[ir] Let the type of ExternalTensorExpression be an ndarray struct

ghstack-source-id: 429c5e9f00736ace86cdd11d65c861da2119550f Pull Request resolved: #8191
taichi-dev · Jun 15, 2023 · 1f4da93 · 1f4da93
1 parent eec09b3
commit 1f4da93
Show file tree

Hide file tree

Showing 2 changed files with 3 additions and 1 deletion.
diff --git a/taichi/ir/frontend_ir.cpp b/taichi/ir/frontend_ir.cpp
@@ -608,6 +608,7 @@ void InternalFuncCallExpression::flatten(FlattenContext *ctx) {
 void ExternalTensorExpression::flatten(FlattenContext *ctx) {
   auto type =
       TypeFactory::get_instance().get_ndarray_struct_type(dt, ndim, needs_grad);
+  type = TypeFactory::get_instance().get_pointer_type((Type *)type);
 
   auto ptr = Stmt::make<ArgLoadStmt>(arg_id, type, /*is_ptr=*/true,
                                      /*create_load=*/false);

diff --git a/taichi/ir/frontend_ir.h b/taichi/ir/frontend_ir.h
@@ -502,7 +502,8 @@ class ExternalTensorExpression : public Expression {
   }
 
   void type_check(const CompileConfig *config) override {
-    ret_type = dt;
+    ret_type = TypeFactory::get_instance().get_ndarray_struct_type(dt, ndim,
+                                                                   needs_grad);
     ret_type.set_is_pointer(true);
     config_ = config;
   }