[CINN] Add type casting for iter_vals (#72204)

ggggxm · web-flow · commit c1c9bfeb5c36 · 2025-04-22T13:48:38.000+08:00
diff --git a/paddle/cinn/ast_gen_ius/ast_gen.cc b/paddle/cinn/ast_gen_ius/ast_gen.cc
@@ -310,6 +310,7 @@ StmtRef AstGen::Build(const ir::Tensor& tensor, TensorGroup* tensor_group) {
       block_vars.push_back(Var(
           Expr(0), shape[i], cinn::UniqName("i" + std::to_string(i)), false));
       optim::ReplaceVarWithExpr(body, axis[i], block_vars[i]);
+      if (shape[i].type() == Int(64)) axis_vars[i]->set_type(Int(64));
       axis_vars[i]->is_reduce_axis = false;
       iter_values.push_back(axis_vars[i]);
     }
diff --git a/paddle/cinn/hlir/pe/elementwise.cc b/paddle/cinn/hlir/pe/elementwise.cc
@@ -338,7 +338,6 @@ ir::Tensor Tril(const ir::Tensor& A,
                               "The Tril op input tensor must have a rank "
                               "greater than or equal to 2."));
         std::vector<Expr> new_indice(indice.end() - 2, indice.end());
-        Expr col_indice = indice.back();
         return ir::Select::Make(new_indice[0] >= new_indice[1] - diagonal,
                                 A(indice),
                                 ir::Zero(A->type()));

Original file line number	Diff line number	Diff line change
`@@ -310,6 +310,7 @@ StmtRef AstGen::Build(const ir::Tensor& tensor, TensorGroup* tensor_group) {`
`310`	`310`	`block_vars.push_back(Var(`
`311`	`311`	`Expr(0), shape[i], cinn::UniqName("i" + std::to_string(i)), false));`
`312`	`312`	`optim::ReplaceVarWithExpr(body, axis[i], block_vars[i]);`
	`313`	`+ if (shape[i].type() == Int(64)) axis_vars[i]->set_type(Int(64));`
`313`	`314`	`axis_vars[i]->is_reduce_axis = false;`
`314`	`315`	`iter_values.push_back(axis_vars[i]);`
`315`	`316`	`}`