[lang] Add support for struct on Dynamic SNode

taichi-dev · Nov 2, 2022 · 27e6c19 · 27e6c19
1 parent 2ee18c5
commit 27e6c19
Show file tree

Hide file tree

Showing 7 changed files with 25 additions and 17 deletions.
diff --git a/python/taichi/lang/snode.py b/python/taichi/lang/snode.py
@@ -1,7 +1,8 @@
+import collections.abc
 import numbers
 
 from taichi._lib import core as _ti_core
-from taichi.lang import expr, impl, matrix
+from taichi.lang import expr, impl, matrix, struct
 from taichi.lang.field import BitpackedFields, Field
 
 
@@ -371,11 +372,17 @@ def append(node, indices, val):
         val (:mod:`~taichi.types.primitive_types`): the scalar data to be appended, only i32 value is support for now.
     """
     if isinstance(val, matrix.Matrix):
-        raise ValueError("ti.append only supports appending a scalar value")
+        raise ValueError("ti.append only supports appending a scalar value or a struct")
+    ptrs = []
+    if isinstance(val, struct.Struct):
+        for item in val._members:
+            ptrs.append(expr.Expr(item).ptr)
+    else:
+        ptrs = [expr.Expr(val).ptr]
     a = impl.expr_init(
         _ti_core.expr_snode_append(node._snode.ptr,
                                    expr.make_expr_group(indices),
-                                   expr.Expr(val).ptr))
+                                   ptrs))
     return a
 
 

diff --git a/taichi/analysis/gen_offline_cache_key.cpp b/taichi/analysis/gen_offline_cache_key.cpp
@@ -222,7 +222,7 @@ class ASTSerializer : public IRVisitor, public ExpressionVisitor {
     emit(expr->op_type);
     emit(expr->snode);
     emit(expr->indices.exprs);
-    emit(expr->value);
+    emit(expr->values);
   }
 
   void visit(ConstExpression *expr) override {

diff --git a/taichi/ir/expr.cpp b/taichi/ir/expr.cpp
@@ -83,7 +83,7 @@ Expr expr_rand(DataType dt) {
   return Expr::make<RandExpression>(dt);
 }
 
-Expr snode_append(SNode *snode, const ExprGroup &indices, const Expr &val) {
+Expr snode_append(SNode *snode, const ExprGroup &indices, const std::vector<Expr> &val) {
   return Expr::make<SNodeOpExpression>(snode, SNodeOpType::append, indices,
                                        val);
 }

diff --git a/taichi/ir/expr.h b/taichi/ir/expr.h
@@ -133,7 +133,7 @@ Expr expr_rand() {
   return taichi::lang::expr_rand(get_data_type<T>());
 }
 
-Expr snode_append(SNode *snode, const ExprGroup &indices, const Expr &val);
+Expr snode_append(SNode *snode, const ExprGroup &indices, const std::vector<Expr> &val);
 
 Expr snode_is_active(SNode *snode, const ExprGroup &indices);
 

diff --git a/taichi/ir/expression_printer.h b/taichi/ir/expression_printer.h
@@ -202,9 +202,9 @@ class ExpressionHumanFriendlyPrinter : public ExpressionPrinter {
     emit('(', expr->snode->get_node_type_name_hinted(), ", [");
     emit_vector(expr->indices.exprs);
     emit("]");
-    if (expr->value.expr) {
+    if (!expr->values.empty()) {
       emit(' ');
-      expr->value->accept(this);
+      emit_vector(expr->values);
     }
     emit(')');
   }

diff --git a/taichi/ir/frontend_ir.cpp b/taichi/ir/frontend_ir.cpp
@@ -925,18 +925,19 @@ void SNodeOpExpression::flatten(FlattenContext *ctx) {
   } else if (op_type == SNodeOpType::get_addr) {
     ctx->push_back<SNodeOpStmt>(SNodeOpType::get_addr, snode, ptr, nullptr);
   } else if (op_type == SNodeOpType::append) {
-    flatten_rvalue(value, ctx);
-
+    for (auto &value : values) {
+      flatten_rvalue(value, ctx);
+    }
     auto alloca = ctx->push_back<AllocaStmt>(PrimitiveType::i32);
     auto addr =
         ctx->push_back<SNodeOpStmt>(SNodeOpType::allocate, snode, ptr, alloca);
-    auto ch_addr = ctx->push_back<GetChStmt>(addr, snode, 0);
-    ctx->push_back<GlobalStoreStmt>(ch_addr, value->stmt);
+    for (int i = 0; i < values.size(); i++) {
+      auto ch_addr = ctx->push_back<GetChStmt>(addr, snode, i);
+      ctx->push_back<GlobalStoreStmt>(ch_addr, values[i]->stmt);
+    }
     ctx->push_back<LocalLoadStmt>(alloca);
     TI_ERROR_IF(snode->type != SNodeType::dynamic,
                 "ti.append only works on dynamic nodes.");
-    TI_ERROR_IF(snode->ch.size() != 1,
-                "ti.append only works on single-child dynamic nodes.");
   }
   stmt = ctx->back_stmt();
 }

diff --git a/taichi/ir/frontend_ir.h b/taichi/ir/frontend_ir.h
@@ -730,7 +730,7 @@ class SNodeOpExpression : public Expression {
   SNode *snode;
   SNodeOpType op_type;
   ExprGroup indices;
-  Expr value;
+  std::vector<Expr> values;
 
   SNodeOpExpression(SNode *snode, SNodeOpType op_type, const ExprGroup &indices)
       : snode(snode), op_type(op_type), indices(indices) {
@@ -739,8 +739,8 @@ class SNodeOpExpression : public Expression {
   SNodeOpExpression(SNode *snode,
                     SNodeOpType op_type,
                     const ExprGroup &indices,
-                    const Expr &value)
-      : snode(snode), op_type(op_type), indices(indices), value(value) {
+                    const std::vector<Expr> &values)
+      : snode(snode), op_type(op_type), indices(indices), values(values) {
   }
 
   void type_check(CompileConfig *config) override;