[CINN]Fix infer shape bug (#62867)

phlrain · web-flow · commit a415e9b068b7 · 2024-03-21T14:06:34.000+08:00
* update

* udpate

* fix bug
diff --git a/paddle/cinn/hlir/framework/pir/group.h b/paddle/cinn/hlir/framework/pir/group.h
@@ -63,6 +63,10 @@ struct Group {
                                ::pir::IrMapping& ir_mapping,
                                const Options& option = Options()) const;
 
+  bool HasShapeOrDataExprs(const ::pir::Value& value) const {
+    return value_to_shape_or_data_exprs_.count(value);
+  }
+
   const symbol::ShapeOrDataDimExprs& GetShapeOrDataExprs(
       const ::pir::Value& value) const {
     CHECK(value_to_shape_or_data_exprs_.count(value))
diff --git a/paddle/cinn/hlir/framework/pir/op_lowering_impl.cc b/paddle/cinn/hlir/framework/pir/op_lowering_impl.cc
@@ -227,20 +227,22 @@ BucketLoweredFuncsWrapper OpLowererImpl::BucketLower(const GroupPtr& group,
   }
   std::vector<ir::Tensor> group_func_arg_tensors_copy = group_func_arg_tensors;
   std::vector<ir::Argument> group_func_args;
+  std::vector<ir::Tensor> infer_shape_tensor_args;
   std::vector<ir::LoweredFunc> funcs = PostProcess(group,
                                                    tensor_map,
                                                    apply_group_schedule,
                                                    {scheduled_func_bodies},
                                                    &group_func_arg_tensors_copy,
-                                                   &group_func_args);
+                                                   &group_func_args,
+                                                   &infer_shape_tensor_args);
   CHECK_EQ(funcs.size(), cond2func_bodies.size());
   BucketLoweredFuncsWrapper funcs_wrapper;
   for (int i = 0; i < funcs.size(); ++i) {
     funcs_wrapper.predicate2funcs.emplace_back(cond2func_bodies[i].first,
                                                funcs[i]);
   }
-  funcs_wrapper.infer_shape_func = GenerateInferShapeFunc(
-      group, group_func_arg_tensors_copy, group_func_args);
+  funcs_wrapper.infer_shape_func =
+      GenerateInferShapeFunc(group, infer_shape_tensor_args, group_func_args);
 
   return funcs_wrapper;
 }
@@ -363,12 +365,14 @@ std::vector<ir::LoweredFunc> OpLowererImpl::LowerMapExpr(
   // including preparing function args and temporary variables,
   // applying low-level optimization passes, etc.
   std::vector<ir::Argument> group_func_args;
+  std::vector<ir::Tensor> infer_shape_tensor_args;
   return PostProcess(group,
                      *tensor_map,
                      apply_op_schedule,
                      {ir_sch.GetModule().GetExprs()[0]},
                      group_func_arg_tensors,
-                     &group_func_args);
+                     &group_func_args,
+                     &infer_shape_tensor_args);
 }
 
 std::vector<ir::LoweredFunc> OpLowererImpl::LowerGroup(
@@ -439,12 +443,14 @@ std::vector<ir::LoweredFunc> OpLowererImpl::LowerGroup(
   // including preparing function args and temporary variables,
   // applying low-level optimization passes, etc.
   std::vector<ir::Argument> group_func_args;
+  std::vector<ir::Tensor> infer_shape_args;
   return PostProcess(group,
                      tensor_map,
                      do_op_schedule,
                      {ir_sch->GetModule().GetExprs().at(0)},
                      &group_func_arg_tensors,
-                     &group_func_args);
+                     &group_func_args,
+                     &infer_shape_args);
 }
 
 void OpLowererImpl::BuildBroadcastInfo(const GroupPtr& group,
@@ -652,7 +658,8 @@ std::vector<ir::LoweredFunc> OpLowererImpl::PostProcess(
     bool done_op_schedule,
     std::vector<ir::Expr> func_bodies,
     std::vector<ir::Tensor>* group_func_arg_tensors,
-    std::vector<ir::Argument>* group_func_args) {
+    std::vector<ir::Argument>* group_func_args,
+    std::vector<ir::Tensor>* infer_shape_arg_tensor) {
   // 1.Prepare function args
   group->input_names.clear();
   std::unordered_set<std::string> arg_name_set;
@@ -673,6 +680,17 @@ std::vector<ir::LoweredFunc> OpLowererImpl::PostProcess(
       continue;
     }
     auto tensor = tensor_map.at(op_result);
+    if (group->HasShapeOrDataExprs(op_result)) {
+      tensor->shape.clear();
+      for (size_t i = 0;
+           i < group->GetShapeOrDataExprs(op_result).shape().size();
+           ++i) {
+        ir::Dim t(tensor->name,
+                  group->GetShapeOrDataExprs(op_result).shape()[i]);
+        tensor->shape.push_back(t->dim_expr);
+      }
+    }
+    infer_shape_arg_tensor->push_back(tensor);
     if ((op_result.defining_op()->name() == "cinn_op.reshape") &&
         erase_reshape.count(op_result.defining_op())) {
       tensor = tensor_map.at(op_result.defining_op()->operand_source(0));
@@ -1172,9 +1190,6 @@ ir::LoweredFunc OpLowererImpl::GenerateInferShapeFunc(
   int output_tensor_idx = 0;
   for (int tensor_arg_idx = 0; tensor_arg_idx < group_func_arg_tensors.size();
        ++tensor_arg_idx) {
-    if (group_func_args[tensor_arg_idx].is_input()) {
-      continue;
-    }
     auto tensor_dim = group_func_arg_tensors[tensor_arg_idx]->sym_shape;
     int tensor_dim_size = tensor_dim.size();
     auto tensor_shape = group_func_arg_tensors[tensor_arg_idx]->shape;
diff --git a/paddle/cinn/hlir/framework/pir/op_lowering_impl.h b/paddle/cinn/hlir/framework/pir/op_lowering_impl.h
@@ -131,7 +131,8 @@ class OpLowererImpl : public OpLowererImplBase<GroupPtr> {
       bool done_op_schedule,
       std::vector<ir::Expr> func_bodies,
       std::vector<ir::Tensor>* group_func_arg_tensors,
-      std::vector<ir::Argument>* group_func_args);
+      std::vector<ir::Argument>* group_func_args,
+      std::vector<ir::Tensor>* infer_shape_arg_tensor);
 
   /**
    * @brief Lower an Op set to CINN IR.
