[Lang] Matrix/Vector refactor: support basic matrix ops

python/taichi/lang/expr.py

@@ -1,8 +1,10 @@
+from typing import Iterable
+
 import numpy as np
 from taichi._lib import core as _ti_core
 from taichi.lang import impl
 from taichi.lang.common_ops import TaichiOperations
-from taichi.lang.exception import TaichiTypeError
+from taichi.lang.exception import TaichiCompilationError, TaichiTypeError
 from taichi.lang.util import is_taichi_class, to_numpy_type
 from taichi.types import primitive_types
 from taichi.types.primitive_types import integer_types, real_types
@@ -39,6 +41,20 @@ def __init__(self, *args, tb=None, dtype=None):
             self.ptr.set_tb(self.tb)
         self.ptr.type_check(impl.get_runtime().prog.config)
 
+    def __getitem__(self, indices):
+        if not isinstance(indices, Iterable):
+            indices = (indices, )
+        return impl.subscript(self, *indices)
+
+    def is_tensor(self):
+        return self.ptr.is_tensor()
+
+    def get_shape(self):
+        if not self.is_tensor():
+            raise TaichiCompilationError(
+                f"Getting shape of non-tensor type: {self.ptr.get_ret_type()}")
+        return self.ptr.get_shape()
+
     def __hash__(self):
         return self.ptr.get_raw_address()
 

taichi/codegen/llvm/codegen_llvm.h

@@ -102,6 +102,22 @@ class TaskCodeGenLLVM : public IRVisitor, public LLVMModuleBuilder {
                              llvm::Value *node_meta,
                              SNode *snode);
 
+  void create_value_cast(
+      UnaryOpStmt *stmt,
+      std::function<llvm::Value *(llvm::Value *, llvm::Type *)> cast_fn,
+      DataType to_ty);
+
+  void create_fp_trunc(
+      UnaryOpStmt *stmt,
+      std::function<llvm::Value *(llvm::Value *, llvm::Type *)> trunc_fn,
+      llvm::Type *to_ty,
+      bool is_tensor,
+      bool trunc_self = false);
+
+  void create_elementwise_binary(
+      BinaryOpStmt *stmt,
+      std::function<llvm::Value *(llvm::Value *lhs, llvm::Value *rhs)> f);
+
   std::unique_ptr<RuntimeObject> emit_struct_meta_object(SNode *snode);
 
   llvm::Value *emit_struct_meta(SNode *snode);

taichi/ir/frontend_ir.cpp

@@ -198,6 +198,32 @@ void UnaryOpExpression::flatten(FlattenContext *ctx) {
   ctx->push_back(std::move(unary));
 }
 
+Expr to_broadcast_tensor(const Expr &elt, const DataType &dt) {
+  TI_ASSERT(dt->is<TensorType>());
+  if (elt->ret_type == dt) {
+    return elt;
+  }
+  auto tensor_type = dt->as<TensorType>();
+  auto elt_type = tensor_type->get_element_type();
+  TI_ASSERT_INFO(elt_type->is<PrimitiveType>(),
+                 "Only primitive types are supported in Tensors, got {}",
+                 elt_type->to_string());
+  std::vector<Expr> broadcast_values(tensor_type->get_num_elements(), elt);
+  return Expr::make<MatrixExpression>(broadcast_values,
+                                      tensor_type->get_shape(), elt->ret_type);
+}
+
+std::tuple<Expr, Expr> unify_binop_operands(const Expr &e1, const Expr &e2) {
+  if (e1->ret_type->is<PrimitiveType>() && e2->ret_type->is<TensorType>()) {
+    return std::tuple(to_broadcast_tensor(e1, e2->ret_type), e2);
+  } else if (e1->ret_type->is<TensorType>() &&
+             e2->ret_type->is<PrimitiveType>()) {
+    return std::tuple(e1, to_broadcast_tensor(e2, e1->ret_type));
+  } else {
+    return std::tuple(e1, e2);
+  }
+}
+
 void BinaryOpExpression::type_check(CompileConfig *config) {
   TI_ASSERT_TYPE_CHECKED(lhs);
   TI_ASSERT_TYPE_CHECKED(rhs);
@@ -217,17 +243,53 @@ void BinaryOpExpression::type_check(CompileConfig *config) {
 
   if ((lhs_type->is<PrimitiveType>() && rhs_type->is<TensorType>()) ||
       (lhs_type->is<TensorType>() && rhs_type->is<PrimitiveType>())) {
-    TI_NOT_IMPLEMENTED;
+    // convert Tensor/Scalar | Scalar/Tensor operations to broadcasting
+    auto [unified_l, unified_r] = unify_binop_operands(lhs, rhs);
+    lhs = unified_l;
+    rhs = unified_r;
+    if (lhs->ret_type == PrimitiveType::unknown)
+      lhs.type_check(config);
+    if (rhs->ret_type == PrimitiveType::unknown)
+      rhs.type_check(config);
+    TI_ASSERT(lhs->ret_type->is<TensorType>());
+    TI_ASSERT(rhs->ret_type->is<TensorType>());
+    lhs_type = lhs->ret_type;
+    rhs_type = rhs->ret_type;
   }
 
+  bool is_tensor_op = false;
+
+  if (lhs_type->is<TensorType>()) {
+    is_tensor_op = true;
+    auto rhs_tensor_type = rhs_type->cast<TensorType>();
+    if (rhs_tensor_type->get_shape() !=
+        lhs_type->cast<TensorType>()->get_shape())
+      // current assume element-wise binary op
+      error();
+  }
+
+  auto make_dt = [&is_tensor_op, this](DataType dt) {
+    if (is_tensor_op) {
+      return TypeFactory::create_tensor_type(
+          this->lhs->ret_type->cast<TensorType>()->get_shape(), dt);
+    } else {
+      return dt;
+    }
+  };
+
   if (binary_is_bitwise(type) &&
-      (!is_integral(lhs_type) || !is_integral(rhs_type)))
+      (!is_integral(lhs_type) || !is_integral(rhs_type)) &&
+      (!is_integral_tensor(lhs_type) || !is_integral_tensor(rhs_type)))
     error();
   if (binary_is_logical(type) &&
-      (lhs_type != PrimitiveType::i32 || rhs_type != PrimitiveType::i32))
+      (lhs_type != PrimitiveType::i32 || rhs_type != PrimitiveType::i32) &&
+      (!is_tensor_op || (lhs_type->cast<TensorType>()->get_element_type() !=
+                             PrimitiveType::i32 ||
+                         rhs_type->cast<TensorType>()->get_element_type() !=
+                             PrimitiveType::i32)))
     error();
   if (is_comparison(type) || binary_is_logical(type)) {
-    ret_type = PrimitiveType::i32;
+    ret_type = make_dt(PrimitiveType::i32);
     return;
   }
   if (is_shift_op(type) ||
@@ -240,20 +302,20 @@ void BinaryOpExpression::type_check(CompileConfig *config) {
   // Try not promoting to fp64 unless necessary
   if (type == BinaryOpType::atan2) {
     if (lhs_type == PrimitiveType::f64 || rhs_type == PrimitiveType::f64) {
-      ret_type = PrimitiveType::f64;
+      ret_type = make_dt(PrimitiveType::f64);
     } else {
-      ret_type = PrimitiveType::f32;
+      ret_type = make_dt(PrimitiveType::f32);
     }
     return;
   }
 
   if (type == BinaryOpType::truediv) {
     auto default_fp = config->default_fp;
-    if (!is_real(lhs_type)) {
-      lhs_type = default_fp;
+    if (!is_real(lhs_type) || !is_real_tensor(lhs_type)) {
+      lhs_type = make_dt(default_fp);
     }
-    if (!is_real(rhs_type)) {
-      rhs_type = default_fp;
+    if (!is_real(rhs_type) || !is_real_tensor(rhs_type)) {
+      rhs_type = make_dt(default_fp);
     }
   }
   ret_type = promoted_type(lhs_type, rhs_type);

taichi/ir/type_factory.cpp

@@ -181,11 +181,25 @@ static DataType to_primitive_type(DataType d) {
 };
 }  // namespace
 
-DataType promoted_type(DataType x, DataType y) {
+DataType promoted_primitive_type(DataType x, DataType y) {
   if (compare_types(to_primitive_type(x), to_primitive_type(y)))
     return x;
   else
     return y;
 }
 
+DataType promoted_type(DataType a, DataType b) {
+  if (a->is<TensorType>() || b->is<TensorType>()) {
+    TI_ASSERT(a->is<TensorType>() && b->is<TensorType>());
+    auto tensor_ty_a = a->cast<TensorType>();
+    auto tensor_ty_b = b->cast<TensorType>();
+    auto promoted_dt = promoted_type(tensor_ty_a->get_element_type(),
+                                     tensor_ty_b->get_element_type());
+    return TypeFactory::create_tensor_type(tensor_ty_a->get_shape(),
+                                           promoted_dt);
+  } else {
+    return promoted_primitive_type(a, b);
+  }
+};
+
 TLANG_NAMESPACE_END

taichi/ir/type_utils.h

@@ -185,6 +185,19 @@ inline TypedConstant get_min_value(DataType dt) {
   }
 }
 
+#define DEFINE_TENSOR_CHECK(func_name)                   \
+  inline bool func_name##_tensor(DataType dt) {          \
+    if (auto tensor_type = dt->cast<TensorType>()) {     \
+      return func_name(tensor_type->get_element_type()); \
+    }                                                    \
+    return false;                                        \
+  }
+
+DEFINE_TENSOR_CHECK(is_quant);
+DEFINE_TENSOR_CHECK(is_real);
+DEFINE_TENSOR_CHECK(is_integral);
+DEFINE_TENSOR_CHECK(is_signed);
+
 class BitStructTypeBuilder {
  public:
   explicit BitStructTypeBuilder(int max_num_bits) {

taichi/python/export_lang.cpp

@@ -777,6 +777,16 @@ void export_lang(py::module &m) {
           },
           py::return_value_policy::reference)
       .def("get_ret_type", &Expr::get_ret_type)
+      .def("is_tensor",
+           [](Expr *expr) { return expr->expr->ret_type->is<TensorType>(); })
+      .def("get_shape",
+           [](Expr *expr) -> std::optional<std::vector<int>> {
+             if (expr->expr->ret_type->is<TensorType>()) {
+               return std::optional<std::vector<int>>(
+                   expr->expr->ret_type->cast<TensorType>()->get_shape());
+             }
+             return std::nullopt;
+           })
       .def("type_check", &Expr::type_check)
       .def("get_expr_name",
            [](Expr *expr) { return expr->cast<FieldExpression>()->name; })

taichi/transforms/type_check.cpp

@@ -286,20 +286,31 @@ class TypeCheck : public IRVisitor {
         stmt->rhs->ret_type->is_primitive(PrimitiveTypeID::unknown))
       error();
     if (stmt->op_type == BinaryOpType::pow &&
-        is_integral(stmt->rhs->ret_type)) {
+        (is_integral(stmt->rhs->ret_type) ||
+         is_integral_tensor(stmt->rhs->ret_type))) {
       stmt->ret_type = stmt->lhs->ret_type;
       return;
     }
 
+    auto make_dt = [stmt](DataType dt) {
+      if (auto tensor_ty = stmt->lhs->ret_type->cast<TensorType>()) {
+        return TypeFactory::create_tensor_type(tensor_ty->get_shape(), dt);
+      } else {
+        return dt;
+      }
+    };
+
     // lower truediv into div
 
     if (stmt->op_type == BinaryOpType::truediv) {
       auto default_fp = config_.default_fp;
-      if (!is_real(stmt->lhs->ret_type)) {
-        cast(stmt->lhs, default_fp);
+      if (!is_real(stmt->lhs->ret_type) ||
+          !is_real_tensor(stmt->lhs->ret_type)) {
+        cast(stmt->lhs, make_dt(default_fp));
       }
-      if (!is_real(stmt->rhs->ret_type)) {
-        cast(stmt->rhs, default_fp);
+      if (!is_real(stmt->rhs->ret_type) ||
+          !is_real_tensor(stmt->rhs->ret_type)) {
+        cast(stmt->rhs, make_dt(default_fp));
       }
       stmt->op_type = BinaryOpType::div;
     }
@@ -309,13 +320,13 @@ class TypeCheck : public IRVisitor {
     if (stmt->op_type == BinaryOpType::atan2) {
       if (stmt->rhs->ret_type == PrimitiveType::f64 ||
           stmt->lhs->ret_type == PrimitiveType::f64) {
-        stmt->ret_type = PrimitiveType::f64;
-        cast(stmt->rhs, PrimitiveType::f64);
-        cast(stmt->lhs, PrimitiveType::f64);
+        stmt->ret_type = make_dt(PrimitiveType::f64);
+        cast(stmt->rhs, make_dt(PrimitiveType::f64));
+        cast(stmt->lhs, make_dt(PrimitiveType::f64));
       } else {
-        stmt->ret_type = PrimitiveType::f32;
-        cast(stmt->rhs, PrimitiveType::f32);
-        cast(stmt->lhs, PrimitiveType::f32);
+        stmt->ret_type = make_dt(PrimitiveType::f32);
+        cast(stmt->rhs, make_dt(PrimitiveType::f32));
+        cast(stmt->lhs, make_dt(PrimitiveType::f32));
       }
     }
 
@@ -359,7 +370,7 @@ class TypeCheck : public IRVisitor {
       error();
     }
     if (is_comparison(stmt->op_type)) {
-      stmt->ret_type = PrimitiveType::i32;
+      stmt->ret_type = make_dt(PrimitiveType::i32);
     } else {
       stmt->ret_type = stmt->lhs->ret_type;
     }

tests/python/test_matrix.py

@@ -817,6 +817,50 @@ def bar():
         bar()
 
 
+@test_utils.test(arch=[ti.cuda, ti.cpu], real_matrix=True, debug=True)
+def test_elementwise_ops():
+    @ti.kernel
+    def test():
+        # TODO: fix parallelization
+        x = ti.Matrix([[1, 2], [3, 4]])
+        # Unify rhs
+        t1 = x + 10
+        ti.loop_config(serialize=True)
+        for i in range(2):
+            for j in range(2):
+                assert t1[i, j] == x[i, j] + 10
+        t2 = x * 2
+        ti.loop_config(serialize=True)
+        for i in range(2):
+            for j in range(2):
+                assert t2[i, j] == x[i, j] * 2
+        # elementwise-add
+        t3 = t1 + t2
+        ti.loop_config(serialize=True)
+        for i in range(2):
+            for j in range(2):
+                assert t3[i, j] == t1[i, j] + t2[i, j]
+        # Unify lhs
+        t4 = 1 / t1
+        # these should be *exactly* equals
+        ti.loop_config(serialize=True)
+        for i in range(2):
+            for j in range(2):
+                assert t4[i, j] == 1 / t1[i, j]
+        t5 = 1 << x
+        ti.loop_config(serialize=True)
+        for i in range(2):
+            for j in range(2):
+                assert t5[i, j] == 1 << x[i, j]
+        t6 = 1 + (x // 2)
+        ti.loop_config(serialize=True)
+        for i in range(2):
+            for j in range(2):
+                assert t6[i, j] == 1 + (x[i, j] // 2)
+
+    test()
+
+
 @test_utils.test()
 def test_vector_vector_t():
     @ti.kernel