[bug] MatrixType bug fix: Fix error with quant (#6776)

Issue: #5819 ### Brief Summary Quant types are not primitive types, so we should not make assertions that types only contain primitive types and tensor type. We just need to trigger scalarization when all the operands have `TensorType`. Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com> Co-authored-by: Zhanlue Yang <[email protected]>
taichi-dev · Dec 2, 2022 · a202762 · a202762
1 parent cb53290
commit a202762
Show file tree

Hide file tree

Showing 2 changed files with 58 additions and 42 deletions.
diff --git a/taichi/transforms/scalarize.cpp b/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>());

diff --git a/tests/python/test_quant_atomics.py b/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()