[CINN] Fix reshape op fusion with axis of length one (#69250)

huangjiyi · web-flow · commit 88c922a54eb4 · 2024-11-11T11:25:52.000+08:00
diff --git a/paddle/cinn/operator_fusion/pir_graph_analyzing/shardable_axes_base.cc b/paddle/cinn/operator_fusion/pir_graph_analyzing/shardable_axes_base.cc
@@ -336,6 +336,13 @@ ShardableAxesSignature CreateSignatureForReshape(
     return shape_analysis->IsProductEqual(
         op->operand_source(0), 0, lhs_end, op->result(0), 0, rhs_end);
   };
+  const auto axis_equal = [&](int input_axis, int output_axis) {
+    const auto& input_sym =
+        shape_analysis->GetProductDimExpr(op->operand_source(0), {input_axis});
+    const auto& output_sym =
+        shape_analysis->GetProductDimExpr(op->result(0), {output_axis});
+    return shape_analysis->IsEqual(input_sym, output_sym);
+  };
   const auto axis_equal_one = [&shape_analysis](pir::Value v, int axis) {
     const auto& sym = shape_analysis->GetProductDimExpr(v, {axis});
     return shape_analysis->IsEqual(sym, symbol::DimExpr(1));
@@ -360,19 +367,19 @@ ShardableAxesSignature CreateSignatureForReshape(
                  ::common::errors::InvalidArgument(
                      "Shape product should be equal for reshape op."));
 
-  std::vector<std::pair<int, int>> partion_indices = {{-1, -1}};
+  std::vector<std::pair<int, int>> partion_indices = {{0, 0}};
   for (int i = 1, j = 1; i <= input_rank && j <= output_rank;) {
     if (shape_product_equal(i, j)) {
-      partion_indices.emplace_back((i++ - 1), (j++ - 1));
+      partion_indices.emplace_back(i++, j++);
       if (i > input_rank || j > output_rank) {
-        partion_indices.back().first = input_rank - 1;
-        partion_indices.back().second = output_rank - 1;
+        partion_indices.back().first = input_rank;
+        partion_indices.back().second = output_rank;
       }
     } else if (j < output_rank) {
       j++;
     } else if (i < input_rank) {
       i++;
-      j = partion_indices.back().second + 2;
+      j = partion_indices.back().second + 1;
     } else {
       PADDLE_THROW(::common::errors::InvalidArgument(
           "Shape product should be equal with whole shape."));
@@ -381,22 +388,27 @@ ShardableAxesSignature CreateSignatureForReshape(
 
   std::vector<std::string> output_axes;
   for (int i = 1; i < partion_indices.size(); ++i) {
-    const auto& in_start = partion_indices[i - 1].first + 1;
+    const auto& in_start = partion_indices[i - 1].first;
     const auto& in_end = partion_indices[i].first;
-    const auto& out_start = partion_indices[i - 1].second + 1;
+    const auto& out_start = partion_indices[i - 1].second;
     const auto& out_end = partion_indices[i].second;
-    if (in_end == in_start && out_end == out_start) {
-      output_axes.emplace_back(input_axes[in_end]);
+    if (in_end == in_start + 1 && out_end == out_start + 1) {
+      output_axes.emplace_back(input_axes[in_start]);
     } else {
-      for (int i = out_start; i <= out_end; ++i) {
+      for (int i = out_start; i < out_end; ++i) {
         output_axes.emplace_back(ShardableAxesInfoManager::GetUniqueName());
-        for (int j = in_start; j <= in_end; ++j) {
-          if (!axis_equal_one(op->operand_source(0), j) &&
-              !axis_equal_one(op->result(0), i)) {
+        if (axis_equal_one(op->result(0), i)) {
+          continue;
+        }
+        for (int j = in_start; j < in_end; ++j) {
+          if (!axis_equal_one(op->operand_source(0), j) && !axis_equal(j, i)) {
             axes_manager->related_axes_map()[input_axes[j]].insert(
                 output_axes.back());
-            VLOG(4) << "Relate " << input_axes[j] << " to "
-                    << output_axes.back();
+            VLOG(4) << "Relate input axis[" << j << "]: " << input_axes[j]
+                    << " to output axis[" << i << "]: " << output_axes.back();
+          } else if (axis_equal(j, i)) {
+            output_axes.back() = input_axes[j];
+            break;
           }
         }
       }
diff --git a/paddle/cinn/operator_fusion/policy/iters_fusion_policy.cc b/paddle/cinn/operator_fusion/policy/iters_fusion_policy.cc
@@ -36,7 +36,8 @@ bool ItersFusionPolicy::CheckItersRelation(const PatternNodePtr& source,
     for (const auto& related_iter : source_related_iters) {
       if (iters_manager_->CanFindRelatedIters(related_iter,
                                               target_unique_iters)) {
-        VLOG(4) << "Can not fuse because find related iters";
+        VLOG(4) << "Can not fuse from " << source->fusion_iters().DebugStr()
+                << " to " << target->fusion_iters().DebugStr();
         return false;
       }
     }

Original file line number	Diff line number	Diff line change
`@@ -36,7 +36,8 @@ bool ItersFusionPolicy::CheckItersRelation(const PatternNodePtr& source,`
`36`	`36`	`for (const auto& related_iter : source_related_iters) {`
`37`	`37`	`if (iters_manager_->CanFindRelatedIters(related_iter,`
`38`	`38`	`target_unique_iters)) {`
`39`		`- VLOG(4) << "Can not fuse because find related iters";`
	`39`	`+ VLOG(4) << "Can not fuse from " << source->fusion_iters().DebugStr()`
	`40`	`+ << " to " << target->fusion_iters().DebugStr();`
`40`	`41`	`return false;`
`41`	`42`	`}`
`42`	`43`	`}`