[Lang] MatrixNdarray refactor part12: Add scalarization for AtomicOpStmt

taichi/ir/frontend_ir.cpp

@@ -751,7 +751,7 @@ void IdExpression::flatten(FlattenContext *ctx) {
   }
 }
 
-void AtomicOpExpression::type_check(CompileConfig *) {
+void AtomicOpExpression::type_check(CompileConfig *config) {
   TI_ASSERT_TYPE_CHECKED(dest);
   TI_ASSERT_TYPE_CHECKED(val);
   auto error = [&]() {
@@ -760,11 +760,34 @@ void AtomicOpExpression::type_check(CompileConfig *) {
         atomic_op_type_name(op_type), dest->ret_type->to_string(),
         val->ret_type->to_string()));
   };
-  if (!val->ret_type->is<PrimitiveType>())
+
+  // Broadcast val to dest if neccessary
+  auto val_dtype = val->ret_type;
+  auto dest_dtype = dest->ret_type.ptr_removed();
+  if (dest_dtype->is<PrimitiveType>() && val_dtype->is<TensorType>()) {
+    error();
+  }
+
+  if (val_dtype->is<PrimitiveType>() && dest_dtype->is<TensorType>()) {
+    auto broadcasted_expr = to_broadcast_tensor(val, dest_dtype);
+    val = std::move(broadcasted_expr);
+    val.type_check(config);
+  }
+
+  // Validate dtype
+  auto dtype = val->ret_type;
+  if (dtype->is<TensorType>()) {
+    dtype = dtype.get_element_type();
+  }
+
+  if (!dtype->is<PrimitiveType>()) {
     error();
+  }
+
   if (is_quant(dest->ret_type)) {
     ret_type = dest->ret_type->get_compute_type();
-  } else if (dest->ret_type->is<PrimitiveType>()) {
+  } else if (dest->ret_type->is<PrimitiveType>() ||
+             dest->ret_type->is<TensorType>()) {
     ret_type = dest->ret_type;
   } else {
     error();

taichi/transforms/scalarize.cpp

@@ -258,6 +258,92 @@ class Scalarize : public BasicStmtVisitor {
     }
   }
 
+  /*
+    Before:
+      TensorType<4 x i32> val = AtomicStmt(TensorType<4 x i32>* dest,
+                                           TensorType<4 x i32>  val)
+
+    After:
+      i32* dest_ptr_0 = MatrixPtrStmt(dest, 0)
+      i32* dest_ptr_1 = MatrixPtrStmt(dest, 1)
+      i32* dest_ptr_2 = MatrixPtrStmt(dest, 2)
+      i32* dest_ptr_3 = MatrixPtrStmt(dest, 3)
+
+      i32 dest_val0 = AtomicStmt(dest_ptr_0,
+                                 val->cast<MatrixInitStmt>()->val[0])
+      i32 dest_val1 = AtomicStmt(dest_ptr_1,
+                                 val->cast<MatrixInitStmt>()->val[1])
+      i32 dest_val2 = AtomicStmt(dest_ptr_2,
+                                 val->cast<MatrixInitStmt>()->val[2])
+      i32 dest_val3 = AtomicStmt(dest_ptr_3,
+                                 val->cast<MatrixInitStmt>()->val[3])
+
+      tmp = MatrixInitStmt(dest_val0, dest_val1,
+                           dest_val2, dest_val3)
+
+      stmt->replace_all_usages_with(tmp)
+  */
+  void visit(AtomicOpStmt *stmt) override {
+    auto dest_dtype = stmt->dest->ret_type.ptr_removed();
+    auto val_dtype = stmt->val->ret_type;
+
+    if (dest_dtype->is<PrimitiveType>() && val_dtype->is<PrimitiveType>()) {
+      return;
+    }
+
+    // AtomicOpExpression::type_check() have taken care of the broadcasting,
+    // but the type conversions are delayed until irpass::type_check().
+    // So we only check for the shape here.
+    TI_ASSERT(dest_dtype->is<TensorType>() && val_dtype->is<TensorType>());
+    TI_ASSERT(dest_dtype->cast<TensorType>()->get_shape() ==
+              val_dtype->cast<TensorType>()->get_shape());
+
+    if (dest_dtype->is<TensorType>() && val_dtype->is<TensorType>()) {
+      // Scalarization for LoadStmt should have already replaced val operand
+      // to MatrixInitStmt
+      TI_ASSERT(stmt->val->is<MatrixInitStmt>());
+
+      auto val_matrix_init_stmt = stmt->val->cast<MatrixInitStmt>();
+      std::vector<Stmt *> val_values = val_matrix_init_stmt->values;
+
+      size_t num_elements = val_values.size();
+      auto primitive_type = stmt->ret_type.get_element_type();
+
+      // Scalarize dest & val
+      std::vector<Stmt *> matrix_init_values;
+      for (size_t i = 0; i < num_elements; i++) {
+        // scalarize to dest_i
+        auto const_stmt = std::make_unique<ConstStmt>(
+            TypedConstant(get_data_type<int32>(), i));
+        auto matrix_ptr_stmt =
+            std::make_unique<MatrixPtrStmt>(stmt->dest, const_stmt.get());
+
+        // scalarize to val_i
+        auto val_stmt = val_values[i];
+
+        // assemble to scalarized atomic_op
+        auto atomic_stmt = std::make_unique<AtomicOpStmt>(
+            stmt->op_type, matrix_ptr_stmt.get(), val_stmt);
+        atomic_stmt->ret_type = primitive_type;
+
+        matrix_init_values.push_back(atomic_stmt.get());
+
+        modifier_.insert_before(stmt, std::move(const_stmt));
+        modifier_.insert_before(stmt, std::move(matrix_ptr_stmt));
+        modifier_.insert_before(stmt, std::move(atomic_stmt));
+      }
+
+      auto matrix_init_stmt =
+          std::make_unique<MatrixInitStmt>(matrix_init_values);
+      matrix_init_stmt->ret_type = stmt->ret_type;
+
+      stmt->replace_usages_with(matrix_init_stmt.get());
+      modifier_.insert_before(stmt, std::move(matrix_init_stmt));
+
+      modifier_.erase(stmt);
+    }
+  }
+
   void visit(GlobalStoreStmt *stmt) override {
     scalarize_store_stmt<GlobalStoreStmt>(stmt);
   }

tests/python/test_matrix.py

@@ -1044,3 +1044,33 @@ def verify(x):
     field = ti.Matrix.field(2, 2, ti.f32, shape=5)
     ndarray = ti.Matrix.ndarray(2, 2, ti.f32, shape=5)
     _test_field_and_ndarray(field, ndarray, func, verify)
+
+
+@test_utils.test(arch=[ti.cuda, ti.cpu],
+                 real_matrix=True,
+                 real_matrix_scalarize=True,
+                 debug=True)
+def test_atomic_op_scalarize():
+    @ti.func
+    def func(x: ti.template()):
+        x[0] = [1., 2., 3.]
+        tmp = ti.Vector([3., 2., 1.])
+        z = ti.atomic_add(x[0], tmp)
+        assert z[0] == 1.
+        assert z[1] == 2.
+        assert z[2] == 3.
+
+        # Broadcasting
+        x[1] = [1., 1., 1.]
+        g = ti.atomic_add(x[1], 2)
+        assert g[0] == 1.
+        assert g[1] == 1.
+        assert g[2] == 1.
+
+    def verify(x):
+        assert (x[0] == [4., 4., 4.]).all()
+        assert (x[1] == [3., 3., 3.]).all()
+
+    field = ti.Vector.field(n=3, dtype=ti.f32, shape=10)
+    ndarray = ti.Vector.ndarray(n=3, dtype=ti.f32, shape=(10))
+    _test_field_and_ndarray(field, ndarray, func, verify)