avoid mix reduction type in one fusion block (#1101)

* fix row reduce fusion

* update

Author: Yancey0623

tao_compiler/mlir/disc/transforms/fusion_utils.cc

@@ -1467,6 +1467,23 @@ Value BaseGpuFusionStrategy::getEffectiveShape(FusionPattern& target, Value v) {
   return isa<lmhlo::ReduceOp>(result_op) ? result_op->getOperand(0) : v;
 }
 
+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) {
+    return false;
+  }
+  return BaseFusionStrategy::tryFuse(shapeAnalysis, lhs, rhs, target);
+}
+
 ////////////////////// Stitch-Base CPU FusionStrategy Implemenation /////
 //////////////////////////////////////////////////////////////////
 

tao_compiler/mlir/disc/transforms/fusion_utils.h

@@ -791,7 +791,8 @@ class BaseGpuFusionStrategy : public BaseFusionStrategy {
                       FusionPattern& target) {
     return lhs.isKInputFusion() && rhs.isKInputFusion();
   }
-
+  virtual bool tryFuse(ShapeAnalysis& shapeAnalysis, FusionPattern& lhs,
+                       FusionPattern& rhs, FusionPattern& target) override;
   Value getEffectiveShape(FusionPattern& target, Value v) override;
   virtual StringRef getName() override { return "BaseGpuFusionStrategy"; }
 };
@@ -801,7 +802,8 @@ class StitchGpuFusionStrategy : public FusionStrategy {
   StitchGpuFusionStrategy(const FusionOptions& options)
       : FusionStrategy(options) {}
   virtual bool isFusible(Operation* op) override;
-
+  virtual bool tryFuse(ShapeAnalysis& shapeAnalysis, FusionPattern& lhs,
+                       FusionPattern& rhs, FusionPattern& target) override;
   virtual bool initFusionPattern(ShapeAnalysis& shapeAnalysis,
                                  FusionPattern& fusion_pattern) override;
   virtual StringRef getName() override { return "StitchGpuFusionStrategy"; }

tao_compiler/mlir/disc/transforms/fusion_utils_stitch_gpu.cc

@@ -55,6 +55,23 @@ bool StitchGpuFusionStrategy::isFusible(Operation* op) {
   return true;
 }
 
+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) {
+    return false;
+  }
+  return FusionStrategy::tryFuse(shapeAnalysis, lhs, rhs, target);
+}
+
 Value StitchGpuFusionStrategy::getEffectiveShape(FusionPattern& target,
                                                  Value v) {
   Operation* result_op = target.findLastWriter(v);

tao_compiler/tao

@@ -0,0 +1 @@
+/workspace/BladeDISC/tao