Repository: alibaba/BladeDISC (language: python)

Issue:
Introduce MLIR transform dialect to BladeDISC
We'll start to explore using MLIR [transform dialect ](https://mlir.llvm.org/docs/Dialects/Transform) to do codegen for (fused) compute-intensive pattern. The initial target is to support gemm codegen on ARM platform to address the dynamic shape problem of Arm Compute Library.

The initial plan is:

- [x]  Step 1,  enhance the fusion decision pass. We’ll add a new fusion kind `kTransform` for the transform-based fusion pattern.
- [x]  Step 2,  lower the lmhlo fusion op to linalg on tensor.
- [x]  Step 3,  transform the linalg computation to loops using transform dialect.
- [x]  Step 4,  refined the transformed loop to make it suitable for BladeDISC runtime.
- [x]  Step 5, add a new pass to the disc pass pipeline to drive the above process.
- [x]   Step 6, weight pre-packing support
   - [x]  add `disc_linalg.multi_level_pack` op, used for doing packing.
   - [x]  add  `transform.disc.cache_read` transform op, relying on `disc_linalg.multi_level_pack` op.
   - [x]  add folding support for `disc_linalg.multi_level_pack`.
   - [x]  lower `disc_linalg.multi_level_pack` to loop if it can not be folded.
   - [x]  fuse const weight op into the `kTransform` fusion pattern, lower it to linalg and then schedule it.
- [x]  Step 7,  assign a default schedule for each `kTransform` pattern.  
- [x]  Step 8,  schedule selection logic injection
- [x]  Step 9,  initial model level testing:  bert (albert).
- [x]  Step 10, support nt, tn, tt format GEMM.
- [ ]  Step 11, support batch matmul
- [x]  Step 12, support GEMM epilogue fusion.
- [ ]  Step 13, performance optimization


Discussion / hints:

Fix patch:
diff --git a/tao_compiler/mlir/disc/transforms/fusion_utils.cc b/tao_compiler/mlir/disc/transforms/fusion_utils.cc
index 819ecbf7183..a6709f50843 100644
--- a/tao_compiler/mlir/disc/transforms/fusion_utils.cc
+++ b/tao_compiler/mlir/disc/transforms/fusion_utils.cc
@@ -552,7 +552,7 @@ FusionPatternBase::FusionPatternBase(SmallVectorImpl<Operation*>& op_list)
 // fusion pattern contains.
 int FusionPatternBase::effectiveSize() {
   return llvm::count_if(
-      op_list_, [](Operation* op) { return !matchPattern(op, m_Constant()); });
+      op_list_, [](Operation* op) { return !isa<lmhlo::ConstantOp>(op); });
 }
 
 // Sorts the ops inside the fusion pattern according to the keys provided.
diff --git a/tao_compiler/mlir/disc/transforms/fusion_utils_transform_based.cc b/tao_compiler/mlir/disc/transforms/fusion_utils_transform_based.cc
index deff59cab25..0b3bbcf031d 100644
--- a/tao_compiler/mlir/disc/transforms/fusion_utils_transform_based.cc
+++ b/tao_compiler/mlir/disc/transforms/fusion_utils_transform_based.cc
@@ -18,10 +18,7 @@ namespace disc_ral {
 /////////// Transform based CPU FusionStrategy Implemenation ///////////
 ////////////////////////////////////////////////////////////////////////
 
-bool TransformBasedCpuFusionStrategy::isFusible(Operation* op) {
-  // Only support matmul a.t.m.
-  // TODO(wyzero): support const weight
-  // TODO(wyzero): support gemm + elemwise fusion
+bool isSupportedDot(Operation* op) {
   auto dotOp = dyn_cast<lmhlo::DotGeneralOp>(op);
   if (!dotOp) return false;
 
@@ -45,21 +42,49 @@ bool TransformBasedCpuFusionStrategy::isFusible(Operation* op) {
           rhsCntractingDims.size() == 1 && rhsCntractingDims[0] == 0);
 }
 
+bool TransformBasedCpuFusionStrategy::isFusible(Operation* op) {
+  return isSupportedDot(op) || isa<lmhlo::ConstantOp>(op);
+}
+
 bool TransformBasedCpuFusionStrategy::initFusionPattern(
     ShapeAnalysis& shapeAnalysis, FusionPattern& fusionPattern) {
-  Operation* inferredDominantOp = nullptr;
-  FusionType inferredFusionType = FusionType::kNone;
-
-  // TODO(wyzero): support dot fusion
-  bool allSupported =
-      llvm::all_of(fusionPattern.getOpList(),
-                   [&](Operation* op) { return this->isFusible(op); });
-  if (allSupported && fusionPattern.getOpList().size() == 1) {
-    inferredDominantOp = fusionPattern.getOpList()[0];
-    inferredFusionType = FusionType::kTransform;
+  // firstly init the fusion kind to kNone
+  fusionPattern.setDominantOp(nullptr);
+  fusionPattern.setFusionType(FusionType::kNone);
+
+  // special case for single operation.
+  if (fusionPattern.getOpList().size() == 1) {
+    Operation* op = *fusionPattern.getOpList().begin();
+    if (this->isFusible(op)) {
+      fusionPattern.setDominantOp(op);
+      fusionPattern.setFusionType(FusionType::kTransform);
+    }
+    return true;
   }
-  fusionPattern.setDominantOp(inferredDominantOp);
-  fusionPattern.setFusionType(inferredFusionType);
+
+  DenseSet<Value> dotWeights;
+  DenseSet<Operation*> supportedDotOps;
+  for (Operation* op : fusionPattern.getOpList()) {
+    // early return for the case where there are non supported ops.
+    if (!this->isFusible(op)) return true;
+    if (isSupportedDot(op)) {
+      supportedDotOps.insert(op);
+      dotWeights.insert(op->getOperand(1));
+    }
+  }
+
+  // Only support one gemm a.t.m.
+  if (supportedDotOps.size() != 1) return true;
+
+  // Only support fuse const ops that are used as weights for some dot ops.
+  for (Operation* op : fusionPattern.getOpList()) {
+    if (!isa<lmhlo::ConstantOp>(op)) continue;
+    if (llvm::find(dotWeights, op->getOperand(0)) == dotWeights.end())
+      return true;
+  }
+
+  fusionPattern.setDominantOp(*supportedDotOps.begin());
+  fusionPattern.setFusionType(FusionType::kTransform);
   return true;
 }
 
diff --git a/tao_compiler/mlir/disc/transforms/lhlo_fusion.cc b/tao_compiler/mlir/disc/transforms/lhlo_fusion.cc
index c0865ae178d..682540d9296 100644
--- a/tao_compiler/mlir/disc/transforms/lhlo_fusion.cc
+++ b/tao_compiler/mlir/disc/transforms/lhlo_fusion.cc
@@ -156,9 +156,9 @@ class FusionPlanner {
       FusionPattern& fusion_pattern = cluster->fused_pattern();
       // Make sure the ops in a fusion pattern are in topological ordering.
       fusion_pattern.sortFusionOpListBy(op_to_node_id_);
-      if (!fusion_pattern.isFusible() ||
+      if (!fusion_pattern.isFusible() || fusion_pattern.effectiveSize() < 1 ||
           !fusion_pattern.isTransformBasedFusion() &&
-              fusion_pattern.effectiveSize() <= 1) {
+              fusion_pattern.effectiveSize() == 1) {
         continue;
       }
       plan.emplace_back(fusion_pattern);


Test patch:
diff --git a/tao_compiler/mlir/disc/transforms/tests/cpu-only-lmhlo-fusion.mlir b/tao_compiler/mlir/disc/transforms/tests/cpu-only-lmhlo-fusion.mlir
index db384f52cf8..0a754572fef 100644
--- a/tao_compiler/mlir/disc/transforms/tests/cpu-only-lmhlo-fusion.mlir
+++ b/tao_compiler/mlir/disc/transforms/tests/cpu-only-lmhlo-fusion.mlir
@@ -16,7 +16,7 @@ func.func @custom_call_op(%arg0: memref<?x?xf32, "cpu">, %arg1: memref<?x?xf32,
 
 // TRANSFORM-LABEL: @matmul_nn
 func.func @matmul_nn(%arg0: memref<?x?xf32, "cpu">, %arg1: memref<?x?xf32, "cpu">,
-                          %arg2: memref<?x?xf32, "cpu">, %arg3: memref<?x?xf32, "cpu">) -> memref<?x?xf32, "cpu"> {
+                     %arg2: memref<?x?xf32, "cpu">, %arg3: memref<?x?xf32, "cpu">) -> memref<?x?xf32, "cpu"> {
   // TRANSFORM: "lmhlo.fusion"() ({
   // TRANSFORM-NEXT: lmhlo.dot_general
   // TRANSFORM-NEXT: lmhlo.terminator
@@ -25,4 +25,22 @@ func.func @matmul_nn(%arg0: memref<?x?xf32, "cpu">, %arg1: memref<?x?xf32, "cpu"
   "lmhlo.abs"(%arg0, %arg1) : (memref<?x?xf32, "cpu">, memref<?x?xf32, "cpu">) -> ()
   "lmhlo.dot_general"(%arg1, %arg2, %arg3) {dot_dimension_numbers = #mhlo.dot<lhs_contracting_dimensions = [1], rhs_contracting_dimensions = [0]>} : (memref<?x?xf32, "cpu">, memref<?x?xf32, "cpu">, memref<?x?xf32, "cpu">) -> ()
   return %arg3 : memref<?x?xf32, "cpu">
+}
+
+// -----
+
+// TRANSFORM-LABEL: @matmul_nn
+func.func @matmul_nn_const_weight(%arg0: memref<1024x1024xf32, "cpu">, %arg1: memref<1024x1024xf32, "cpu">,
+                                  %arg2: memref<?x1024xf32, "cpu">, %arg3: memref<?x1024xf32, "cpu">) -> memref<?x1024xf32, "cpu"> {
+  // TRANSFORM: lmhlo.constant
+  // TRANSFORM-NEXT: "lmhlo.fusion"() ({
+  // TRANSFORM-NEXT: lmhlo.constant
+  // TRANSFORM-NEXT: lmhlo.dot_general
+  // TRANSFORM-NEXT: lmhlo.terminator
+  // TRANSFORM-NEXT: })
+  // TRANSFORM-SAME: disc.fusion_type = "kTransform"
+  "lmhlo.constant"(%arg0) {disc.device = "cpu", value = dense<1.0> : tensor<1024x1024xf32>} : (memref<1024x1024xf32, "cpu">) -> ()
+  "lmhlo.constant"(%arg1) {disc.device = "cpu", value = dense<-1.0> : tensor<1024x1024xf32>} : (memref<1024x1024xf32, "cpu">) -> ()
+  "lmhlo.dot_general"(%arg2, %arg1, %arg3) {dot_dimension_numbers = #mhlo.dot<lhs_contracting_dimensions = [1], rhs_contracting_dimensions = [0]>} : (memref<?x1024xf32, "cpu">, memref<1024x1024xf32, "cpu">, memref<?x1024xf32, "cpu">) -> ()
+  return %arg3 : memref<?x1024xf32, "cpu">
 }
\ No newline at end of file