Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

[lang] Store relations with 16-bit type #4779

Merged
merged 4 commits into from
Apr 13, 2022
Merged

[lang] Store relations with 16-bit type #4779

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
4 commits
Select commit Hold shift + click to select a range
e4a921f
offset/value from int32 -> uint16, add auxiliary data 'patch_offset'
g1n0st Apr 13, 2022
f202948
update mesh test
g1n0st Apr 13, 2022
6e81e7f
remove i16
g1n0st Apr 13, 2022
95e6d58
[pre-commit.ci] auto fixes from pre-commit.com hooks
pre-commit-ci[bot] Apr 13, 2022
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
15 changes: 11 additions & 4 deletions python/taichi/lang/mesh.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -10,7 +10,7 @@
_MatrixFieldElement)
from taichi.lang.struct import StructField
from taichi.lang.util import python_scope
from taichi.types import i32
from taichi.types import i32, u16, u32
from taichi.types.compound_types import CompoundType

from taichi import lang
Expand Down Expand Up @@ -286,9 +286,11 @@ def set_relation_fixed(self, rel_type: MeshRelationType,
value.vars[0].ptr.snode())

def set_relation_dynamic(self, rel_type: MeshRelationType,
value: ScalarField, offset: ScalarField):
value: ScalarField, patch_offset: ScalarField,
offset: ScalarField):
_ti_core.set_relation_dynamic(self.mesh_ptr, rel_type,
value.vars[0].ptr.snode(),
patch_offset.vars[0].ptr.snode(),
offset.vars[0].ptr.snode())

def add_mesh_attribute(self, element_type, snode, reorder_type):
Expand Down Expand Up @@ -332,10 +334,12 @@ def __init__(self, data):
relation_by_orders(from_order, to_order))
self.relation_fields[rel_type] = {}
self.relation_fields[rel_type]["value"] = impl.field(
dtype=i32, shape=len(relation["value"]))
dtype=u16, shape=len(relation["value"]))
if from_order <= to_order:
self.relation_fields[rel_type]["offset"] = impl.field(
dtype=i32, shape=len(relation["offset"]))
dtype=u16, shape=len(relation["offset"]))
self.relation_fields[rel_type]["patch_offset"] = impl.field(
dtype=u32, shape=len(relation["patch_offset"]))

for element in data["elements"]:
element_type = MeshElementType(element["order"])
Expand All @@ -358,6 +362,8 @@ def __init__(self, data):
self.relation_fields[rel_type]["value"].from_numpy(
np.array(relation["value"]))
if from_order <= to_order:
self.relation_fields[rel_type]["patch_offset"].from_numpy(
np.array(relation["patch_offset"]))
self.relation_fields[rel_type]["offset"].from_numpy(
np.array(relation["offset"]))

Expand Down Expand Up @@ -422,6 +428,7 @@ def build(self, metadata: MeshMetadata):
if from_order <= to_order:
instance.set_relation_dynamic(
rel_type, metadata.relation_fields[rel_type]["value"],
metadata.relation_fields[rel_type]["patch_offset"],
metadata.relation_fields[rel_type]["offset"])
else:
instance.set_relation_fixed(
Expand Down
5 changes: 3 additions & 2 deletions taichi/ir/mesh.h
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -53,8 +53,8 @@ MeshRelationType relation_by_orders(int from_order, int to_order);
MeshRelationType inverse_relation(MeshRelationType rel);

struct MeshLocalRelation {
MeshLocalRelation(SNode *value_, SNode *offset_)
: value(value_), offset(offset_) {
MeshLocalRelation(SNode *value_, SNode *patch_offset_, SNode *offset_)
: value(value_), patch_offset(patch_offset_), offset(offset_) {
fixed = false;
}

Expand All @@ -64,6 +64,7 @@ struct MeshLocalRelation {

bool fixed;
SNode *value{nullptr};
SNode *patch_offset{nullptr};
SNode *offset{nullptr};
};

Expand Down
4 changes: 2 additions & 2 deletions taichi/ir/statements.h
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -1629,7 +1629,7 @@ class MeshRelationAccessStmt : public Stmt {
mesh_idx(mesh_idx),
to_type(to_type),
neighbor_idx(neighbor_idx) {
this->ret_type = PrimitiveType::i32;
this->ret_type = PrimitiveType::u16;
TI_STMT_REG_FIELDS;
}

Expand All @@ -1640,7 +1640,7 @@ class MeshRelationAccessStmt : public Stmt {
mesh_idx(mesh_idx),
to_type(to_type),
neighbor_idx(nullptr) {
this->ret_type = PrimitiveType::i32;
this->ret_type = PrimitiveType::u16;
TI_STMT_REG_FIELDS;
}

Expand Down
6 changes: 3 additions & 3 deletions taichi/python/export_lang.cpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -1066,9 +1066,9 @@ void export_lang(py::module &m) {

m.def("set_relation_dynamic",
[](mesh::MeshPtr &mesh_ptr, mesh::MeshRelationType type, SNode *value,
SNode *offset) {
mesh_ptr.ptr->relations.insert(
std::pair(type, mesh::MeshLocalRelation(value, offset)));
SNode *patch_offset, SNode *offset) {
mesh_ptr.ptr->relations.insert(std::pair(
type, mesh::MeshLocalRelation(value, patch_offset, offset)));
});
}

Expand Down
9 changes: 7 additions & 2 deletions taichi/transforms/demote_mesh_statements.cpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -103,10 +103,13 @@ void demote_mesh_statements_offload(OffloadedStmt *offload,
stmt->replace_with(std::move(block));
}
} else { // low-to-high or same-order
SNode *rel_offset = stmt->mesh->relations.find(rel_type)->second.offset;
const auto &rel = stmt->mesh->relations.find(rel_type)->second;
SNode *rel_offset = rel.offset;
SNode *rel_patch_offset = rel.patch_offset;
VecStatement block;
Stmt *patch_idx = block.push_back<MeshPatchIndexStmt>();
Stmt *owned_offset = offload->owned_offset_local.find(from_type)->second;
Stmt *patch_offset = get_load(rel_patch_offset, patch_idx, block);
Stmt *index_offset = block.push_back<BinaryOpStmt>(
BinaryOpType::add, patch_idx, owned_offset);
Stmt *index = block.push_back<BinaryOpStmt>(BinaryOpType::add,
Expand All @@ -121,8 +124,10 @@ void demote_mesh_statements_offload(OffloadedStmt *offload,
block.push_back<BinaryOpStmt>(BinaryOpType::sub, offset_1, offset);
} else {
SNode *rel_value = stmt->mesh->relations.find(rel_type)->second.value;
Stmt *val_index = block.push_back<BinaryOpStmt>(
Stmt *val_local_index = block.push_back<BinaryOpStmt>(
BinaryOpType::add, offset, stmt->neighbor_idx);
Stmt *val_index = block.push_back<BinaryOpStmt>(
BinaryOpType::add, val_local_index, patch_offset);
[[maybe_unused]] Stmt *val = get_load(rel_value, val_index, block);
}
stmt->replace_with(std::move(block));
Expand Down
Loading