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

ghstack-source-id: 8b5cea534c215bcacff943ff89294700ffa63ce3 Pull Request resolved: taichi-dev#8191
PGZXB · Jul 13, 2023 · 871a493 · 871a493
1 parent 6cf2d46
commit 871a493
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;
   }