Refine colreduction fusion strategy of kStitch. (alibaba#1257)

JamesTheZ · Oct 9, 2023 · c46084d · c46084d
1 parent da990f4
commit c46084d
Show file tree

Hide file tree

Showing 5 changed files with 44 additions and 76 deletions.
diff --git a/.github/workflows/pytorch_pre_cpu.yml b/.github/workflows/pytorch_pre_cpu.yml
diff --git a/.github/workflows/pytorch_pre_gpu.yml b/.github/workflows/pytorch_pre_gpu.yml
diff --git a/tao_compiler/mlir/disc/transforms/fusion_utils.cc b/tao_compiler/mlir/disc/transforms/fusion_utils.cc
@@ -1475,16 +1475,28 @@ bool BaseGpuFusionStrategy::tryFuse(ShapeAnalysis& shapeAnalysis,
                                     FusionPattern& lhs, FusionPattern& rhs,
                                     FusionPattern& target) {
   // TODO(Yancey): support fusion with different reduction type
-  bool has_row_reduction = llvm::any_of(target.getOpList(), [](Operation* op) {
-    return isRank2RowReduction(op);
-  });
-  bool has_col_reduciton = llvm::any_of(target.getOpList(), [](Operation* op) {
-    return isRank2ColReduction(op);
-  });
-
-  if (has_row_reduction && has_col_reduciton) {
+  bool has_rank2_row_reduction =
+      llvm::any_of(target.getOpList(),
+                   [](Operation* op) { return isRank2RowReduction(op); });
+  bool has_rank2_col_reduction =
+      llvm::any_of(target.getOpList(),
+                   [](Operation* op) { return isRank2ColReduction(op); });
+
+  if (has_rank2_row_reduction && has_rank2_col_reduction) {
     return false;
   }
+
+  if (has_rank2_col_reduction) {
+    const auto& results = target.getResults();
+    auto ref_shape = getEffectiveShape(target, results[0]);
+    if (llvm::any_of(results, [&](Value result) {
+          auto op = target.findLastWriter(result);
+          return isa<lmhlo::TransposeOp>(op);
+        })) {
+      return false;
+    }
+  }
+
   return BaseFusionStrategy::tryFuse(shapeAnalysis, lhs, rhs, target);
 }
 

diff --git a/tao_compiler/mlir/disc/transforms/fusion_utils_stitch_gpu.cc b/tao_compiler/mlir/disc/transforms/fusion_utils_stitch_gpu.cc
@@ -59,16 +59,28 @@ bool StitchGpuFusionStrategy::tryFuse(ShapeAnalysis& shapeAnalysis,
                                       FusionPattern& lhs, FusionPattern& rhs,
                                       FusionPattern& target) {
   // TODO(Yancey): support fusion with different reduction type
-  bool has_row_reduction = llvm::any_of(target.getOpList(), [](Operation* op) {
-    return isRank2RowReduction(op);
-  });
-  bool has_col_reduciton = llvm::any_of(target.getOpList(), [](Operation* op) {
-    return isRank2ColReduction(op);
-  });
-
-  if (has_row_reduction && has_col_reduciton) {
+  bool has_rank2_row_reduction =
+      llvm::any_of(target.getOpList(),
+                   [](Operation* op) { return isRank2RowReduction(op); });
+  bool has_rank2_col_reduction =
+      llvm::any_of(target.getOpList(),
+                   [](Operation* op) { return isRank2ColReduction(op); });
+
+  if (has_rank2_row_reduction && has_rank2_col_reduction) {
     return false;
   }
+
+  if (has_rank2_col_reduction) {
+    const auto& results = target.getResults();
+    auto ref_shape = getEffectiveShape(target, results[0]);
+    if (llvm::any_of(results, [&](Value result) {
+          auto op = target.findLastWriter(result);
+          return isa<lmhlo::TransposeOp>(op);
+        })) {
+      return false;
+    }
+  }
+
   return FusionStrategy::tryFuse(shapeAnalysis, lhs, rhs, target);
 }
 
@@ -428,9 +440,8 @@ bool StitchGpuFusionStrategy::findFusionPatternTypeAndSubroot(
               return true;
             }
             Value shape = getEffectiveShape(fusion_pattern, result);
-            return isRank2ColReduction(op)
-                       ? shapeAnalysis.isShapeEqual(ref_shape, shape)
-                       : shapeAnalysis.isSameNumElements(ref_shape, shape);
+            return isRank2ColReduction(op) &&
+                   shapeAnalysis.isShapeEqual(ref_shape, shape);
           })) {
         return false;
       }

diff --git a/tao_compiler/mlir/disc/transforms/lhlo_legalize_roots_to_loops.cc b/tao_compiler/mlir/disc/transforms/lhlo_legalize_roots_to_loops.cc
@@ -650,7 +650,7 @@ LogicalResult lowerWithScheduleColReduction(
 
   SmallVector<Value, 4> yield_values_for_if;
 
-  ValueRange load_index({row_index, col_index});
+  SmallVector<Value, 2> load_index({row_index, col_index});
   b.setInsertionPointToStart(&if_row_valid_op.getThenRegion().front());
   int col_red_root_op_idx = 0;
   for (auto* root_op : root_ops) {
@@ -924,7 +924,7 @@ LogicalResult lowerWithScheduleColReductionTileH(
 
   SmallVector<Value, 4> yield_values_for_if;
 
-  ValueRange load_index({row_index, col_index});
+  SmallVector<Value, 2> load_index({row_index, col_index});
   b.setInsertionPointToStart(&if_row_valid_op.getThenRegion().front());
   int col_red_root_op_idx = 0;
   for (auto* root_op : root_ops) {