[ir] [refactor] Let the type of Alloca be pointer

ghstack-source-id: 663f7b3c9e383c290b47ddd4183b64d540b35c41
Pull Request resolved: #8007

taichi/codegen/cuda/codegen_cuda.cpp

@@ -172,8 +172,8 @@ class TaskCodeGenCUDA : public TaskCodeGenLLVM {
 
   void visit(AllocaStmt *stmt) override {
     // Override shared memory codegen logic for large shared memory
-    if (stmt->ret_type->is<TensorType>() && stmt->is_shared) {
-      auto tensor_type = stmt->ret_type->cast<TensorType>();
+    auto tensor_type = stmt->ret_type.ptr_removed()->cast<TensorType>();
+    if (tensor_type && stmt->is_shared) {
       size_t shared_array_bytes =
           tensor_type->get_num_elements() *
           data_type_size(tensor_type->get_element_type());

taichi/codegen/llvm/codegen_llvm.cpp

@@ -126,8 +126,10 @@ void TaskCodeGenLLVM::visit(Block *stmt_list) {
 }
 
 void TaskCodeGenLLVM::visit(AllocaStmt *stmt) {
-  if (stmt->ret_type->is<TensorType>()) {
-    auto tensor_type = stmt->ret_type->cast<TensorType>();
+  TI_ASSERT(stmt->ret_type.is_pointer());
+  auto alloca_type = stmt->ret_type.ptr_removed();
+  if (alloca_type->is<TensorType>()) {
+    auto tensor_type = alloca_type->cast<TensorType>();
     auto type = tlctx->get_data_type(tensor_type);
     if (stmt->is_shared) {
       auto base = new llvm::GlobalVariable(
@@ -141,12 +143,10 @@ void TaskCodeGenLLVM::visit(AllocaStmt *stmt) {
       llvm_val[stmt] = create_entry_block_alloca(type);
     }
   } else {
-    llvm_val[stmt] =
-        create_entry_block_alloca(stmt->ret_type, stmt->ret_type.is_pointer());
+    llvm_val[stmt] = create_entry_block_alloca(alloca_type);
     // initialize as zero if element is not a pointer
-    if (!stmt->ret_type.is_pointer())
-      builder->CreateStore(tlctx->get_constant(stmt->ret_type, 0),
-                           llvm_val[stmt]);
+    if (!alloca_type->is<PointerType>())
+      builder->CreateStore(tlctx->get_constant(alloca_type, 0), llvm_val[stmt]);
   }
 }
 

taichi/codegen/llvm/llvm_codegen_utils.h

@@ -83,8 +83,6 @@ class LLVMModuleBuilder {
 
   llvm::Value *create_entry_block_alloca(DataType dt, bool is_pointer = false) {
     auto type = tlctx->get_data_type(dt);
-    if (is_pointer)
-      type = llvm::PointerType::get(type, 0);
     return create_entry_block_alloca(type);
   }
 

taichi/codegen/spirv/spirv_codegen.cpp

@@ -274,8 +274,8 @@ class TaskCodegen : public IRVisitor {
 
   void visit(AllocaStmt *alloca) override {
     spirv::Value ptr_val;
-    if (alloca->ret_type->is<TensorType>()) {
-      auto tensor_type = alloca->ret_type->cast<TensorType>();
+    auto alloca_type = alloca->ret_type.ptr_removed();
+    if (auto tensor_type = alloca_type->cast<TensorType>()) {
       auto elem_num = tensor_type->get_num_elements();
       spirv::SType elem_type =
           ir_->get_primitive_type(tensor_type->get_element_type());
@@ -288,7 +288,7 @@ class TaskCodegen : public IRVisitor {
       }
     } else {
       // Alloca for a single variable
-      spirv::SType src_type = ir_->get_primitive_type(alloca->element_type());
+      spirv::SType src_type = ir_->get_primitive_type(alloca_type);
       ptr_val = ir_->alloca_variable(src_type);
       ir_->store_variable(ptr_val, ir_->get_zero(src_type));
     }

taichi/ir/control_flow_graph.cpp

@@ -247,7 +247,7 @@ Stmt *CFGNode::get_store_forwarding_data(Stmt *var, int position) const {
     // result: the value to store
     Stmt *result = irpass::analysis::get_store_data(
         block->statements[last_def_position].get());
-    bool is_tensor_involved = var->ret_type->is<TensorType>();
+    bool is_tensor_involved = var->ret_type.ptr_removed()->is<TensorType>();
     if (!(var->is<AllocaStmt>() && !is_tensor_involved)) {
       // In between the store stmt and current stmt,
       // if there's a third-stmt that "may" have stored a "different value" to
@@ -355,7 +355,7 @@ Stmt *CFGNode::get_store_forwarding_data(Stmt *var, int position) const {
 
   // Check for aliased address
   // There's a store to the same dest_addr before this stmt
-  bool is_tensor_involved = var->ret_type->is<TensorType>();
+  bool is_tensor_involved = var->ret_type.ptr_removed()->is<TensorType>();
   if (!(var->is<AllocaStmt>() && !is_tensor_involved)) {
     // In between the store stmt and current stmt,
     // if there's a third-stmt that "may" have stored a "different value" to
@@ -443,7 +443,8 @@ bool CFGNode::store_to_load_forwarding(bool after_lower_access,
           continue;
 
         // special case of alloca (initialized to 0)
-        auto zero = Stmt::make<ConstStmt>(TypedConstant(result->ret_type, 0));
+        auto zero = Stmt::make<ConstStmt>(
+            TypedConstant(result->ret_type.ptr_removed(), 0));
         replace_with(i, std::move(zero), true);
       } else {
         if (result->ret_type.ptr_removed()->is<TensorType>() &&

taichi/ir/frontend_ir.cpp

@@ -37,7 +37,8 @@ FrontendAssignStmt::FrontendAssignStmt(const Expr &lhs, const Expr &rhs)
     : lhs(lhs), rhs(rhs) {
   TI_ASSERT(lhs->is_lvalue());
   if (lhs.is<IdExpression>() && lhs->ret_type == PrimitiveType::unknown) {
-    lhs.expr->ret_type = rhs.get_rvalue_type();
+    lhs.expr->ret_type =
+        TypeFactory::get_instance().get_pointer_type(rhs.get_rvalue_type());
   }
 }
 
@@ -127,7 +128,8 @@ void FrontendForStmt::init_loop_vars(const ExprGroup &loop_vars) {
 
 void FrontendForStmt::add_loop_var(const Expr &loop_var) {
   loop_var_ids.push_back(loop_var.cast<IdExpression>()->id);
-  loop_var.expr->ret_type = PrimitiveType::i32;
+  loop_var.expr->ret_type =
+      TypeFactory::get_instance().get_pointer_type(PrimitiveType::i32);
 }
 
 FrontendFuncDefStmt::FrontendFuncDefStmt(const FrontendFuncDefStmt &o)
@@ -883,7 +885,7 @@ void IndexExpression::type_check(const CompileConfig *) {
         "Invalid IndexExpression: the source is not among field, ndarray or "
         "local tensor");
   }
-
+  ret_type = TypeFactory::get_instance().get_pointer_type(ret_type);
   for (auto &indices : indices_group) {
     for (int i = 0; i < indices.exprs.size(); i++) {
       auto &expr = indices.exprs[i];
@@ -960,7 +962,7 @@ void LoopUniqueExpression::flatten(FlattenContext *ctx) {
 
 void IdExpression::flatten(FlattenContext *ctx) {
   stmt = ctx->current_block->lookup_var(id);
-  if (!ret_type->is_primitive(PrimitiveTypeID::unknown)) {
+  if (stmt->ret_type->is_primitive(PrimitiveTypeID::unknown)) {
     stmt->ret_type = ret_type;
   }
 }
@@ -1514,7 +1516,7 @@ Expr ASTBuilder::expr_subscript(const Expr &expr,
                                 std::string tb) {
   TI_ASSERT(expr.is<FieldExpression>() || expr.is<MatrixFieldExpression>() ||
             expr.is<ExternalTensorExpression>() ||
-            is_tensor(expr.expr->ret_type));
+            is_tensor(expr.expr->ret_type.ptr_removed()));
 
   // IndexExpression without ret_shape is used for matrix indexing,
   // where each entry of ExprGroup is interpreted as indexing into a specific
@@ -1677,7 +1679,7 @@ std::vector<Expr> ASTBuilder::expand_exprs(const std::vector<Expr> &exprs) {
         elem.expr->ret_type = struct_type->get_element_type(indices);
         expanded_exprs.push_back(elem);
       }
-    } else if (!expr->ret_type->is<TensorType>()) {
+    } else if (!expr->ret_type.ptr_removed()->is<TensorType>()) {
       expanded_exprs.push_back(expr);
     } else {
       // Expand TensorType expr
@@ -1695,7 +1697,7 @@ std::vector<Expr> ASTBuilder::expand_exprs(const std::vector<Expr> &exprs) {
           return {ind0, ind1, ind2, ind3}
 
       */
-      auto tensor_type = expr->ret_type->cast<TensorType>();
+      auto tensor_type = expr->ret_type.ptr_removed()->cast<TensorType>();
 
       Expr id_expr;
       if (expr.is<IdExpression>()) {
@@ -1708,7 +1710,7 @@ std::vector<Expr> ASTBuilder::expand_exprs(const std::vector<Expr> &exprs) {
         for (int i = 0; i < shape[0]; i++) {
           auto ind = Expr(std::make_shared<IndexExpression>(
               id_expr, ExprGroup(Expr(i)), expr->tb));
-          ind.expr->ret_type = tensor_type->get_element_type();
+          ind->type_check(nullptr);
           expanded_exprs.push_back(ind);
         }
       } else {
@@ -1717,7 +1719,7 @@ std::vector<Expr> ASTBuilder::expand_exprs(const std::vector<Expr> &exprs) {
           for (int j = 0; j < shape[1]; j++) {
             auto ind = Expr(std::make_shared<IndexExpression>(
                 id_expr, ExprGroup(Expr(i), Expr(j)), expr->tb));
-            ind.expr->ret_type = tensor_type->get_element_type();
+            ind->type_check(nullptr);
             expanded_exprs.push_back(ind);
           }
         }

taichi/ir/frontend_ir.h

@@ -86,7 +86,8 @@ class FrontendAllocaStmt : public Stmt {
                      DataType element,
                      bool is_shared = false)
       : ident(lhs), is_shared(is_shared) {
-    ret_type = DataType(TypeFactory::create_tensor_type(shape, element));
+    ret_type = TypeFactory::get_instance().get_pointer_type(
+        DataType(TypeFactory::create_tensor_type(shape, element)));
   }
 
   bool is_shared;
@@ -500,6 +501,7 @@ class ExternalTensorExpression : public Expression {
 
   void type_check(const CompileConfig *config) override {
     ret_type = dt;
+    ret_type.set_is_pointer(true);
     config_ = config;
   }
 
@@ -585,7 +587,7 @@ class MatrixExpression : public Expression {
                    std::vector<int> shape,
                    DataType element_type)
       : elements(elements) {
-    this->dt = DataType(TypeFactory::create_tensor_type(shape, element_type));
+    dt = TypeFactory::create_tensor_type(shape, element_type);
   }
 
   void type_check(const CompileConfig *config) override;

taichi/ir/statements.cpp

@@ -86,6 +86,7 @@ MatrixOfMatrixPtrStmt::MatrixOfMatrixPtrStmt(const std::vector<Stmt *> &stmts,
                                              DataType dt)
     : stmts(stmts) {
   ret_type = dt;
+  ret_type.set_is_pointer(true);
   TI_STMT_REG_FIELDS;
 }
 

taichi/ir/statements.h

@@ -19,15 +19,20 @@ class Function;
 class AllocaStmt : public Stmt, public ir_traits::Store {
  public:
   explicit AllocaStmt(DataType type) : is_shared(false) {
-    ret_type = type;
+    if (type->is_primitive(PrimitiveTypeID::unknown)) {
+      ret_type = type;
+    } else {
+      ret_type = TypeFactory::get_instance().get_pointer_type(type);
+    }
     TI_STMT_REG_FIELDS;
   }
 
   AllocaStmt(const std::vector<int> &shape,
              DataType type,
              bool is_shared = false)
       : is_shared(is_shared) {
-    ret_type = TypeFactory::create_tensor_type(shape, type);
+    ret_type = TypeFactory::get_instance().get_pointer_type(
+        TypeFactory::create_tensor_type(shape, type));
     TI_STMT_REG_FIELDS;
   }
 

taichi/ir/type_factory.cpp

@@ -222,7 +222,7 @@ static bool compare_types(DataType x, DataType y) {
 static DataType to_primitive_type(DataType d) {
   if (d->is<PointerType>()) {
     d = d->as<PointerType>()->get_pointee_type();
-    TI_WARN("promoted_type got a pointer input.");
+    TI_ERROR("promoted_type got a pointer input.");
   }
 
   if (d->is<TensorType>()) {