[bug] MatrixType bug fix: Fix error with quant

taichi/transforms/scalarize.cpp

@@ -226,20 +226,19 @@ class Scalarize : public BasicStmtVisitor {
   void visit(BinaryOpStmt *stmt) override {
     auto lhs_dtype = stmt->lhs->ret_type;
     auto rhs_dtype = stmt->rhs->ret_type;
-
-    if (lhs_dtype->is<PrimitiveType>() && rhs_dtype->is<PrimitiveType>()) {
-      return;
-    }
-
-    if (lhs_dtype->is<TensorType>() && rhs_dtype->is<TensorType>()) {
+    if (lhs_dtype->is<TensorType>() || rhs_dtype->is<TensorType>()) {
+      // Make sure broadcasting has been correctly applied by
+      // BinaryOpExpression::type_check().
+      TI_ASSERT(lhs_dtype->is<TensorType>() && rhs_dtype->is<TensorType>());
+      // However, since the type conversions are delayed until
+      // irpass::type_check(), we only check for the shape here.
+      TI_ASSERT(lhs_dtype->cast<TensorType>()->get_shape() ==
+                rhs_dtype->cast<TensorType>()->get_shape());
       // Scalarization for LoadStmt should have already replaced both operands
-      // to MatrixInitStmt
+      // to MatrixInitStmt.
       TI_ASSERT(stmt->lhs->is<MatrixInitStmt>());
       TI_ASSERT(stmt->rhs->is<MatrixInitStmt>());
 
-      TI_ASSERT(lhs_dtype->cast<TensorType>()->get_shape() ==
-                rhs_dtype->cast<TensorType>()->get_shape());
-
       auto lhs_matrix_init_stmt = stmt->lhs->cast<MatrixInitStmt>();
       std::vector<Stmt *> lhs_vals = lhs_matrix_init_stmt->values;
 
@@ -322,21 +321,16 @@ class Scalarize : public BasicStmtVisitor {
   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>()) {
+    if (dest_dtype->is<TensorType>() || val_dtype->is<TensorType>()) {
+      // Make sure broadcasting has been correctly applied by
+      // AtomicOpExpression::type_check().
+      TI_ASSERT(dest_dtype->is<TensorType>() && val_dtype->is<TensorType>());
+      // However, since the type conversions are delayed until
+      // irpass::type_check(), we only check for the shape here.
+      TI_ASSERT(dest_dtype->cast<TensorType>()->get_shape() ==
+                val_dtype->cast<TensorType>()->get_shape());
       // Scalarization for LoadStmt should have already replaced val operand
-      // to MatrixInitStmt
+      // to MatrixInitStmt.
       TI_ASSERT(stmt->val->is<MatrixInitStmt>());
 
       auto val_matrix_init_stmt = stmt->val->cast<MatrixInitStmt>();
@@ -411,20 +405,18 @@ class Scalarize : public BasicStmtVisitor {
     auto cond_dtype = stmt->op1->ret_type;
     auto op2_dtype = stmt->op2->ret_type;
     auto op3_dtype = stmt->op3->ret_type;
-
-    if (cond_dtype->is<PrimitiveType>() && op2_dtype->is<PrimitiveType>() &&
-        op3_dtype->is<PrimitiveType>()) {
-      return;
-    }
-
-    // TernaryOpExpression::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(cond_dtype.get_shape() == op2_dtype.get_shape());
-    TI_ASSERT(op2_dtype.get_shape() == op3_dtype.get_shape());
-
-    if (cond_dtype->is<TensorType>() && op2_dtype->is<TensorType>() &&
+    if (cond_dtype->is<TensorType>() || op2_dtype->is<TensorType>() ||
         op3_dtype->is<TensorType>()) {
+      // Make sure broadcasting has been correctly applied by
+      // TernaryOpExpression::type_check().
+      TI_ASSERT(cond_dtype->is<TensorType>() && op2_dtype->is<TensorType>() &&
+                op3_dtype->is<TensorType>());
+      // However, since the type conversions are delayed until
+      // irpass::type_check(), we only check for the shape here.
+      TI_ASSERT(cond_dtype.get_shape() == op2_dtype.get_shape());
+      TI_ASSERT(op2_dtype.get_shape() == op3_dtype.get_shape());
+      // Scalarization for LoadStmt should have already replaced all operands
+      // to MatrixInitStmt.
       TI_ASSERT(stmt->op1->is<MatrixInitStmt>());
       TI_ASSERT(stmt->op2->is<MatrixInitStmt>());
       TI_ASSERT(stmt->op3->is<MatrixInitStmt>());

tests/python/test_quant_atomics.py

@@ -43,9 +43,7 @@ def foo():
     assert z[None] == 3
 
 
-@test_utils.test(require=[ti.extension.quant_basic, ti.extension.data64],
-                 debug=True)
-def test_quant_int_atomics_b64():
+def _test_quant_int_atomics_b64():
     qi13 = ti.types.quant.int(13, True)
 
     x = ti.field(dtype=qi13)
@@ -68,8 +66,21 @@ def foo():
     assert x[2] == 315
 
 
-@test_utils.test(require=ti.extension.quant_basic, debug=True)
-def test_quant_fixed_atomics():
+@test_utils.test(require=[ti.extension.quant_basic, ti.extension.data64],
+                 debug=True)
+def test_quant_int_atomics_b64():
+    _test_quant_int_atomics_b64()
+
+
+@test_utils.test(require=[ti.extension.quant_basic, ti.extension.data64],
+                 debug=True,
+                 real_matrix=True,
+                 real_matrix_scalarize=True)
+def test_quant_int_atomics_b64_real_matrix_scalarize():
+    _test_quant_int_atomics_b64()
+
+
+def _test_quant_fixed_atomics():
     qfxt13 = ti.types.quant.fixed(bits=13, signed=True, scale=0.1)
     qfxt19 = ti.types.quant.fixed(bits=19, signed=False, scale=0.1)
 
@@ -91,3 +102,16 @@ def foo():
     foo()
     assert x[None] == approx(-3.3)
     assert y[None] == approx(1124.4)
+
+
+@test_utils.test(require=ti.extension.quant_basic, debug=True)
+def test_quant_fixed_atomics():
+    _test_quant_fixed_atomics()
+
+
+@test_utils.test(require=ti.extension.quant_basic,
+                 debug=True,
+                 real_matrix=True,
+                 real_matrix_scalarize=True)
+def test_quant_fixed_atomics_real_matrix_scalarize():
+    _test_quant_fixed_atomics()