Fixed minor issue

Author: jim19930609

c_api/src/taichi_core_impl.cpp

@@ -558,8 +558,12 @@ void ti_launch_compute_graph(TiRuntime runtime,
           }
         }
 
-        ndarrays.emplace_back(
-            taichi::lang::Ndarray(devalloc, *prim_ty, shape, elem_shape));
+        taichi::lang::DataType dtype = *prim_ty;
+        if (elem_shape.size() > 0) {
+          dtype = taichi::lang::TypeFactory::get_instance().get_tensor_type(
+              elem_shape, dtype);
+        }
+        ndarrays.emplace_back(taichi::lang::Ndarray(devalloc, dtype, shape));
         arg_map.emplace(std::make_pair(
             arg.name, taichi::lang::aot::IValue::create(ndarrays.back())));
         break;

python/taichi/lang/matrix.py

@@ -1900,10 +1900,10 @@ def __init__(self, n, m, dtype, shape, layout):
 
         self.layout = layout
         self.shape = tuple(shape)
-        self.element_type = TensorType((self.n, self.m), self.dtype)
+        self.element_type = TensorType((self.n, self.m), dtype)
         # TODO: we should pass in element_type, shape, layout instead.
         self.arr = impl.get_runtime().prog.create_ndarray(
-            self.element_type, shape, layout)
+            cook_dtype(self.element_type.ptr), shape, layout)
 
     @property
     def element_shape(self):
@@ -1915,7 +1915,7 @@ def element_shape(self):
             >>> arr.element_shape
             (2, 2)
         """
-        return tuple(self.arr.element_shape)
+        return tuple(self.arr.element_shape())
 
     @python_scope
     def __setitem__(self, key, value):
@@ -1999,12 +1999,12 @@ def __init__(self, n, dtype, shape, layout):
         super().__init__()
         # TODO(zhanlue): remove self.dtype and migrate its usages to element_type
         self.dtype = cook_dtype(dtype)
+
         self.layout = layout
         self.shape = tuple(shape)
-        self.element_type = TensorType((n, ), self.dtype)
-        # TODO: pass in element_type, shape, layout directly
+        self.element_type = TensorType((n, ), dtype)
         self.arr = impl.get_runtime().prog.create_ndarray(
-            self.element_type, shape, layout)
+            cook_dtype(self.element_type.ptr), shape, layout)
 
     @property
     def element_shape(self):
@@ -2016,7 +2016,7 @@ def element_shape(self):
             >>> a.element_shape
             (3,)
         """
-        return tuple(self.arr.element_shape)
+        return tuple(self.arr.element_shape())
 
     @python_scope
     def __setitem__(self, key, value):

python/taichi/types/compound_types.py

@@ -1,14 +1,27 @@
+from taichi._lib.utils import ti_python_core as _ti_python_core
+
 import taichi
 
+_type_factory = _ti_python_core.get_type_factory_instance()
+
 
 class CompoundType:
     pass
 
 
 class TensorType(CompoundType):
     def __init__(self, shape, dtype):
-        self.dtype = dtype
-        self.shape = shape
+        if isinstance(dtype, _ti_python_core.DataType):
+            dtype = dtype.get_ptr()
+        self.ptr = _type_factory.get_tensor_type(shape, dtype)
+        self.shape = self.get_shape()
+        self.dtype = self.get_element_type()
+
+    def get_shape(self):
+        return self.ptr.get_element_shape()
+
+    def get_element_type(self):
+        return self.ptr.get_element_type()
 
 
 # TODO: maybe move MatrixType, StructType here to avoid the circular import?

taichi/aot/graph_data.cpp

@@ -25,6 +25,7 @@ void CompiledGraph::run(
       const aot::IValue &ival = found->second;
       if (ival.tag == aot::ArgKind::kNdarray) {
         Ndarray *arr = reinterpret_cast<Ndarray *>(ival.val);
+
         TI_ERROR_IF(arr->get_element_shape() != symbolic_arg.element_shape,
                     "Mismatched shape information for argument {}",
                     symbolic_arg.name);
@@ -33,12 +34,24 @@ void CompiledGraph::run(
                     "field_dim={} but got an ndarray with field_dim={}",
                     symbolic_arg.name, symbolic_arg.field_dim,
                     arr->shape.size());
-        TI_ERROR_IF(arr->dtype != symbolic_arg.dtype(),
+
+        DataType symbolic_arg_primitive_dtype = symbolic_arg.dtype();
+        if (symbolic_arg.dtype()->is<TensorType>()) {
+          symbolic_arg_primitive_dtype =
+              symbolic_arg.dtype()->cast<TensorType>()->get_element_type();
+        }
+
+        DataType arr_primitive_dtype = arr->dtype;
+        if (arr->dtype->is<TensorType>()) {
+          arr_primitive_dtype =
+              arr->dtype->cast<TensorType>()->get_element_type();
+        }
+
+        TI_ERROR_IF(arr_primitive_dtype != symbolic_arg_primitive_dtype,
                     "Dispatch node is compiled for argument {} with "
                     "dtype={} but got an ndarray with dtype={}",
-                    symbolic_arg.name, symbolic_arg.dtype().to_string(),
-                    arr->dtype.to_string());
-
+                    symbolic_arg.name, symbolic_arg_primitive_dtype.to_string(),
+                    arr_primitive_dtype.to_string());
         ctx.set_arg_ndarray(i, arr->get_device_allocation_ptr_as_int(),
                             arr->shape);
       } else if (ival.tag == aot::ArgKind::kScalar) {

taichi/ir/type.h

@@ -42,6 +42,14 @@ class TI_DLL_EXPORT Type {
 
   bool is_primitive(PrimitiveTypeID type) const;
 
+  virtual std::vector<int> get_shape() const {
+    return {};
+  }
+
+  virtual Type *get_element_type() const {
+    return nullptr;
+  }
+
   virtual Type *get_compute_type() {
     TI_NOT_IMPLEMENTED;
   }
@@ -160,7 +168,7 @@ class TensorType : public Type {
       : shape_(std::move(shape)), element_(element) {
   }
 
-  Type *get_element_type() const {
+  Type *get_element_type() const override {
     return element_;
   }
 
@@ -171,7 +179,7 @@ class TensorType : public Type {
     return num_elements;
   }
 
-  std::vector<int> get_shape() const {
+  std::vector<int> get_shape() const override {
     return shape_;
   }
 
@@ -339,7 +347,7 @@ class QuantArrayType : public Type {
     return physical_type_;
   }
 
-  Type *get_element_type() const {
+  Type *get_element_type() const override {
     return element_type_;
   }
 

taichi/program/ndarray.cpp

@@ -91,10 +91,11 @@ Ndarray::Ndarray(DeviceAllocation &devalloc,
                  const DataType type,
                  const std::vector<int> &shape,
                  const std::vector<int> &element_shape,
-                 ExternalArrayLayout layout) {
-  TI_ASSERT(type->is<PrimitiveType>());
-  auto tensor_type = TypeFactory::create_tensor_type(element_shape, type);
-  Ndarray(devalloc, tensor_type, shape, layout);
+                 ExternalArrayLayout layout)
+    : Ndarray(devalloc,
+              TypeFactory::create_tensor_type(element_shape, type),
+              shape,
+              layout) {
 }
 
 Ndarray::~Ndarray() {

taichi/python/export_lang.cpp

@@ -1090,7 +1090,10 @@ void export_lang(py::module &m) {
 
   // Type system
 
-  py::class_<Type>(m, "Type").def("to_string", &Type::to_string);
+  py::class_<Type>(m, "Type")
+      .def("to_string", &Type::to_string)
+      .def("get_element_shape", &Type::get_shape)
+      .def("get_element_type", &Type::get_element_type);
 
   // Note that it is important to specify py::return_value_policy::reference for
   // the factory methods, otherwise pybind11 will delete the Types owned by
@@ -1104,7 +1107,9 @@ void export_lang(py::module &m) {
            py::return_value_policy::reference)
       .def("get_quant_float_type", &TypeFactory::get_quant_float_type,
            py::arg("digits_type"), py::arg("exponent_type"),
-           py::arg("compute_type"), py::return_value_policy::reference);
+           py::arg("compute_type"), py::return_value_policy::reference)
+      .def("get_tensor_type", &TypeFactory::get_tensor_type, py::arg("shape"),
+           py::arg("element_type"), py::return_value_policy::reference);
 
   m.def("get_type_factory_instance", TypeFactory::get_instance,
         py::return_value_policy::reference);