[mesh] Remove link hints from ti.Mesh

python/taichi/lang/ast/ast_transformer.py

@@ -591,6 +591,8 @@ def build_Module(ctx, node):
     @staticmethod
     def build_Attribute(ctx, node):
         build_stmt(ctx, node.value)
+        if isinstance(node.value.ptr, impl.MeshInstance):
+            ctx.meshes.add(node.value.ptr)
         node.ptr = getattr(node.value.ptr, node.attr)
         return node.ptr
 

python/taichi/lang/ast/ast_transformer_utils.py

@@ -148,6 +148,7 @@ def __init__(self,
         self.ast_builder = ast_builder
         self.visited_funcdef = False
         self.is_real_function = is_real_function
+        self.meshes = set()
 
     # e.g.: FunctionDef, Module, Global
     def variable_scope_guard(self):

python/taichi/lang/kernel_impl.py

@@ -491,6 +491,9 @@ def taichi_ast_generator(kernel_cxx):
         taichi_kernel = impl.get_runtime().prog.create_kernel(
             taichi_ast_generator, kernel_name, self.is_grad)
 
+        for mesh in ctx.meshes:
+            mesh.update_relation()
+
         self.kernel_cpp = taichi_kernel
 
         assert key not in self.compiled_functions

python/taichi/lang/mesh.py

@@ -284,6 +284,8 @@ class MeshInstance:
     def __init__(self, _type):
         self._type = _type
         self.mesh_ptr = _ti_core.create_mesh()
+        self.relation_set = set()
+        self.new_relations = set()
 
     def set_owned_offset(self, element_type: MeshElementType,
                          owned_offset: ScalarField):
@@ -310,12 +312,14 @@ def set_patch_max_element_num(self, element_type: MeshElementType,
 
     def set_relation_fixed(self, rel_type: MeshRelationType,
                            value: ScalarField):
+        self.relation_set.add(rel_type)
         _ti_core.set_relation_fixed(self.mesh_ptr, rel_type,
                                     value.vars[0].ptr.snode())
 
     def set_relation_dynamic(self, rel_type: MeshRelationType,
                              value: ScalarField, patch_offset: ScalarField,
                              offset: ScalarField):
+        self.relation_set.add(rel_type)
         _ti_core.set_relation_dynamic(self.mesh_ptr, rel_type,
                                       value.vars[0].ptr.snode(),
                                       patch_offset.vars[0].ptr.snode(),
@@ -325,6 +329,46 @@ def add_mesh_attribute(self, element_type, snode, reorder_type):
         _ti_core.add_mesh_attribute(self.mesh_ptr, element_type, snode,
                                     reorder_type)
 
+    def get_relation_size(self, from_index, to_element_type):
+        from_order = element_order(from_index.element_type)
+        to_order = element_order(to_element_type)
+        self.check_relation(from_order, to_order)
+        return _ti_core.get_relation_size(self.mesh_ptr, from_index.ptr,
+                                          to_element_type)
+
+    def get_relation_access(self, from_index, to_element_type,
+                            neighbor_idx_ptr):
+        from_order = element_order(from_index.element_type)
+        to_order = element_order(to_element_type)
+        self.check_relation(from_order, to_order)
+        return _ti_core.get_relation_access(self.mesh_ptr, from_index.ptr,
+                                            to_element_type, neighbor_idx_ptr)
+
+    def check_relation(self, from_order, to_order):
+        self.new_relations.add((from_order, to_order))
+
+    def update_relation(self):
+        for from_order, to_order in self.new_relations:
+            rel_type = MeshRelationType(
+                relation_by_orders(from_order, to_order))
+            if rel_type not in self.relation_set:
+                meta = self.patcher.get_relation_meta(from_order, to_order)
+                print('new relation')
+
+                def fun(arr, dtype):
+                    field = impl.field(dtype=dtype, shape=arr.shape)
+                    field.from_numpy(arr)
+                    return field
+
+                if from_order <= to_order:
+                    self.set_relation_dynamic(rel_type,
+                                              fun(meta["value"], u16),
+                                              fun(meta["patch_offset"], u32),
+                                              fun(meta["offset"], u16))
+                else:
+                    self.set_relation_fixed(rel_type, fun(meta["value"], u16))
+        self.new_relations.clear()
+
 
 class MeshMetadata:
     def __init__(self, data):
@@ -368,6 +412,8 @@ def __init__(self, data):
                     dtype=u16, shape=len(relation["offset"]))
                 self.relation_fields[rel_type]["patch_offset"] = impl.field(
                     dtype=u32, shape=len(relation["patch_offset"]))
+            self.relation_fields[rel_type]["from_order"] = from_order
+            self.relation_fields[rel_type]["to_order"] = to_order
 
         for element in data["elements"]:
             element_type = MeshElementType(element["order"])
@@ -397,6 +443,10 @@ def __init__(self, data):
 
         self.attrs = {}
         self.attrs["x"] = np.array(data["attrs"]["x"]).reshape(-1, 3)
+        if "patcher" in data:
+            self.patcher = data["patcher"]
+        else:
+            self.patcher = None
 
 
 # Define the Mesh Type, stores the field type info
@@ -434,7 +484,8 @@ def build(self, metadata: MeshMetadata):
 
         instance.set_num_patches(metadata.num_patches)
 
-        for element in self.elements:
+        for element in metadata.element_fields:
+            self.elements.add(element)
             _ti_core.set_num_elements(instance.mesh_ptr, element,
                                       metadata.num_elements[element])
             instance.set_patch_max_element_num(
@@ -459,11 +510,9 @@ def build(self, metadata: MeshMetadata):
             instance.set_index_mapping(element, ConvType.g2r,
                                        metadata.element_fields[element]["g2r"])
 
-        for relation in self.relations:
-            from_order = element_order(relation[0])
-            to_order = element_order(relation[1])
-            rel_type = MeshRelationType(
-                relation_by_orders(from_order, to_order))
+        for rel_type in metadata.relation_fields:
+            from_order = metadata.relation_fields[rel_type]["from_order"]
+            to_order = metadata.relation_fields[rel_type]["to_order"]
             if from_order <= to_order:
                 instance.set_relation_dynamic(
                     rel_type, metadata.relation_fields[rel_type]["value"],
@@ -476,6 +525,8 @@ def build(self, metadata: MeshMetadata):
         if "x" in instance.verts.attr_dict:  # pylint: disable=E1101
             instance.verts.x.from_numpy(metadata.attrs["x"])  # pylint: disable=E1101
 
+        instance.patcher = metadata.patcher
+
         return instance
 
 
@@ -594,15 +645,13 @@ def __init__(self, mesh: MeshInstance, from_index: impl.Expr,
     @property
     def size(self):
         return impl.Expr(
-            _ti_core.get_relation_size(self.mesh.mesh_ptr, self.from_index.ptr,
-                                       self.to_element_type))
+            self.mesh.get_relation_size(self.from_index, self.to_element_type))
 
     def subscript(self, *indices):
         assert len(indices) == 1
-        entry_expr = _ti_core.get_relation_access(self.mesh.mesh_ptr,
-                                                  self.from_index.ptr,
-                                                  self.to_element_type,
-                                                  impl.Expr(indices[0]).ptr)
+        entry_expr = self.mesh.get_relation_access(self.from_index,
+                                                   self.to_element_type,
+                                                   impl.Expr(indices[0]).ptr)
         entry_expr.type_check(impl.get_runtime().prog.config)
         return MeshElementFieldProxy(self.mesh, self.to_element_type,
                                      entry_expr)