[Lang] MatrixNdarray refactor part13: Add scalarization for TernaryOpStmt

taichi/ir/frontend_ir.cpp

@@ -391,24 +391,63 @@ void make_ifte(Expression::FlattenContext *ctx,
   return;
 }
 
-void TernaryOpExpression::type_check(CompileConfig *) {
+void TernaryOpExpression::type_check(CompileConfig *config) {
   TI_ASSERT_TYPE_CHECKED(op1);
   TI_ASSERT_TYPE_CHECKED(op2);
   TI_ASSERT_TYPE_CHECKED(op3);
   auto op1_type = op1->ret_type;
   auto op2_type = op2->ret_type;
   auto op3_type = op3->ret_type;
+
   auto error = [&]() {
     throw TaichiTypeError(
         fmt::format("unsupported operand type(s) for '{}': '{}', '{}' and '{}'",
                     ternary_type_name(type), op1->ret_type->to_string(),
                     op2->ret_type->to_string(), op3->ret_type->to_string()));
   };
-  if (op1_type != PrimitiveType::i32)
-    error();
-  if (!op2_type->is<PrimitiveType>() || !op3_type->is<PrimitiveType>())
+
+  bool is_valid = true;
+  bool is_tensor = false;
+  if (op1_type->is<TensorType>() && op2_type->is<TensorType>() &&
+      op3_type->is<TensorType>()) {
+    // valid
+    is_tensor = true;
+    if (op1_type->cast<TensorType>()->get_shape() !=
+        op2_type->cast<TensorType>()->get_shape()) {
+      is_valid = false;
+    }
+    if (op2_type->cast<TensorType>()->get_shape() !=
+        op3_type->cast<TensorType>()->get_shape()) {
+      is_valid = false;
+    }
+    op1_type = op1_type->cast<TensorType>()->get_element_type();
+    op2_type = op2_type->cast<TensorType>()->get_element_type();
+    op3_type = op3_type->cast<TensorType>()->get_element_type();
+
+  } else if (op1_type->is<PrimitiveType>() && op2_type->is<PrimitiveType>() &&
+             op3_type->is<PrimitiveType>()) {
+    // valid
+  } else {
+    is_valid = false;
+  }
+
+  if (op1_type != PrimitiveType::i32) {
+    is_valid = false;
+  }
+  if (!op2_type->is<PrimitiveType>() || !op3_type->is<PrimitiveType>()) {
+    is_valid = false;
+  }
+
+  if (!is_valid)
     error();
-  ret_type = promoted_type(op2_type, op3_type);
+
+  if (is_tensor) {
+    auto primitive_dtype = promoted_type(op2_type, op3_type);
+    ret_type = TypeFactory::create_tensor_type(
+        op2->ret_type->cast<TensorType>()->get_shape(), primitive_dtype);
+  } else {
+    ret_type = promoted_type(op2_type, op3_type);
+  }
 }
 
 void TernaryOpExpression::flatten(FlattenContext *ctx) {
@@ -425,6 +464,7 @@ void TernaryOpExpression::flatten(FlattenContext *ctx) {
   }
   stmt = ctx->back_stmt();
   stmt->tb = tb;
+  stmt->ret_type = ret_type;
 }
 
 void InternalFuncCallExpression::type_check(CompileConfig *) {

taichi/transforms/scalarize.cpp

@@ -258,6 +258,90 @@ class Scalarize : public BasicStmtVisitor {
     }
   }
 
+  /*
+    Before:
+      TensorType<4 x i32> val = TernaryStmt(TensorType<4 x i32> cond,
+                                            TensorType<4 x i32> lhs,
+                                            TensorType<4 x i32> rhs)
+
+    After:
+      i32 val0 = TernaryStmt(cond->cast<MatrixInitStmt>()->val[0],
+                             lhs->cast<MatrixInitStmt>()->val[0],
+                             rhs->cast<MatrixInitStmt>()->val[0])
+
+      i32 val1 = TernaryStmt(cond->cast<MatrixInitStmt>()->val[1],
+                             lhs->cast<MatrixInitStmt>()->val[1],
+                             rhs->cast<MatrixInitStmt>()->val[1])
+
+      i32 val2 = TernaryStmt(cond->cast<MatrixInitStmt>()->val[2],
+                             lhs->cast<MatrixInitStmt>()->val[2],
+                             rhs->cast<MatrixInitStmt>()->val[2])
+
+      i32 val3 = TernaryStmt(cond->cast<MatrixInitStmt>()->val[3],
+                             lhs->cast<MatrixInitStmt>()->val[3],
+                             rhs->cast<MatrixInitStmt>()->val[3])
+
+      tmp = MatrixInitStmt(val0, val1, val2, val3)
+
+      stmt->replace_all_usages_with(tmp)
+  */
+  void visit(TernaryOpStmt *stmt) override {
+    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>() &&
+        op3_dtype->is<TensorType>()) {
+      TI_ASSERT(stmt->op1->is<MatrixInitStmt>());
+      TI_ASSERT(stmt->op2->is<MatrixInitStmt>());
+      TI_ASSERT(stmt->op3->is<MatrixInitStmt>());
+
+      auto cond_matrix_init_stmt = stmt->op1->cast<MatrixInitStmt>();
+      std::vector<Stmt *> cond_vals = cond_matrix_init_stmt->values;
+
+      auto op2_matrix_init_stmt = stmt->op2->cast<MatrixInitStmt>();
+      std::vector<Stmt *> op2_vals = op2_matrix_init_stmt->values;
+
+      auto op3_matrix_init_stmt = stmt->op3->cast<MatrixInitStmt>();
+      std::vector<Stmt *> op3_vals = op3_matrix_init_stmt->values;
+
+      TI_ASSERT(cond_vals.size() == op2_vals.size());
+      TI_ASSERT(op2_vals.size() == op3_vals.size());
+
+      size_t num_elements = cond_vals.size();
+      auto primitive_type = stmt->ret_type.get_element_type();
+      std::vector<Stmt *> matrix_init_values;
+      for (size_t i = 0; i < num_elements; i++) {
+        auto ternary_stmt = std::make_unique<TernaryOpStmt>(
+            stmt->op_type, cond_vals[i], op2_vals[i], op3_vals[i]);
+        matrix_init_values.push_back(ternary_stmt.get());
+        ternary_stmt->ret_type = primitive_type;
+
+        modifier_.insert_before(stmt, std::move(ternary_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,23 @@ 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_ternary_op_scalarize():
+    @ti.kernel
+    def test():
+        cond = ti.Vector([1, 0, 1])
+        x = ti.Vector([3, 3, 3])
+        y = ti.Vector([5, 5, 5])
+
+        z = ti.select(cond, x, y)
+
+        assert z[0] == 3
+        assert z[1] == 5
+        assert z[2] == 3
+
+    test()