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/disc_duplicate_computation_for_fusion.cc b/tao_compiler/mlir/disc/transforms/disc_duplicate_computation_for_fusion.cc
index 74ec6efdafb..4f6d85b1448 100644
--- a/tao_compiler/mlir/disc/transforms/disc_duplicate_computation_for_fusion.cc
+++ b/tao_compiler/mlir/disc/transforms/disc_duplicate_computation_for_fusion.cc
@@ -190,6 +190,20 @@ struct DiscDuplicateComputationForFusionPass
       return;
     }
 
+    if (useTransformSchedule()) {
+      auto strategy =
+          makeNewPlacementAwareFusionStrategy(gpu_enabled_, "transform_based");
+      // const weight can be pre-packed. In some cases, multiple dot ops may
+      // share the same const op while use different packed layouts for the
+      // weight. we duplicate the weight to make sure that each dot general has
+      // its own copy, maximizing the opportunities to do weight pre-packing at
+      // compile time.
+      if (failed(duplicateConstWeightForDotOp(func, *strategy))) {
+        signalPassFailure();
+        return;
+      }
+    }
+
     if (isMemIntensiveOptExperimentalEnabled()) {
       // Populate patterns.
       MLIRContext* ctx = func.getContext();
@@ -201,6 +215,8 @@ struct DiscDuplicateComputationForFusionPass
 
   LogicalResult duplicateBroadcastInDimOp(FuncOp func,
                                           FusionStrategy& strategy);
+  LogicalResult duplicateConstWeightForDotOp(FuncOp func,
+                                             FusionStrategy& strategy);
 };
 
 // Returns true if the value produced by a cheap computation.
@@ -295,6 +311,45 @@ LogicalResult DiscDuplicateComputationForFusionPass::duplicateBroadcastInDimOp(
   return success();
 }
 
+LogicalResult
+DiscDuplicateComputationForFusionPass::duplicateConstWeightForDotOp(
+    FuncOp func, FusionStrategy& strategy) {
+  SmallVector<lmhlo::DotGeneralOp> dotOps;
+  func->walk([&](lmhlo::DotGeneralOp op) {
+    if (strategy.isFusible(op.getOperation())) dotOps.emplace_back(op);
+  });
+  for (lmhlo::DotGeneralOp dotOp : dotOps) {
+    Value weight = dotOp.getRhs();
+    Value rootWeightMemref = getRootMemRef(weight);
+    bool hasOtherUsers = false;
+    Operation* constOp = nullptr;
+    for (Operation* user : getValueUsers(rootWeightMemref)) {
+      if (user == dotOp.getOperation()) continue;
+      if (isa<lmhlo::ConstantOp>(user)) {
+        if (constOp)
+          return dotOp->emitError()
+                 << "weight buffer consumsed by mulitple const ops\n";
+        constOp = user;
+        continue;
+      }
+      hasOtherUsers = true;
+    }
+    if (!constOp || !hasOtherUsers) continue;
+    auto weightTy = weight.getType().cast<MemRefType>();
+    auto constTy = constOp->getOperand(0).getType().cast<MemRefType>();
+    // TODO(wyzero): support the case where there are cast ops between const op
+    // and dot op.
+    if (weightTy != constTy) continue;
+    OpBuilder b(dotOp);
+    Location loc = dotOp->getLoc();
+    Value newWeight = b.create<memref::AllocOp>(loc, weightTy);
+    Operation* clonedConstOp = b.clone(*constOp);
+    clonedConstOp->replaceUsesOfWith(constOp->getOperand(0), newWeight);
+    dotOp->replaceUsesOfWith(weight, newWeight);
+  }
+  return success();
+}
+
 }  // namespace
 
 std::unique_ptr<OperationPass<func::FuncOp>>
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 0b3bbcf031d..18ce52804a6 100644
--- a/tao_compiler/mlir/disc/transforms/fusion_utils_transform_based.cc
+++ b/tao_compiler/mlir/disc/transforms/fusion_utils_transform_based.cc
@@ -76,10 +76,13 @@ bool TransformBasedCpuFusionStrategy::initFusionPattern(
   // 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.
+  // Only support fuse const ops that are used as weights for some dot ops and
+  // not consumed by ops outside the fusion pattern.
   for (Operation* op : fusionPattern.getOpList()) {
     if (!isa<lmhlo::ConstantOp>(op)) continue;
-    if (llvm::find(dotWeights, op->getOperand(0)) == dotWeights.end())
+    if (llvm::find(dotWeights, op->getOperand(0)) == dotWeights.end() ||
+        llvm::find(fusionPattern.getRootOps(), op) !=
+            fusionPattern.getRootOps().end())
       return true;
   }
 


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 0a754572fef..3f3f0cd6484 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
@@ -1,5 +1,5 @@
 // RUN: disc-opt -pass-pipeline='func.func(disc-fusion{gpu-enabled=false fusion-strategy=base})' -split-input-file %s -o - | FileCheck %s --check-prefix=BASE
-// RUN: DISC_ENABLE_TRANSFORM_SCHEDULE=1 disc-opt -pass-pipeline='func.func(disc-fusion{gpu-enabled=false fusion-strategy=stitch})' -split-input-file %s -o - | FileCheck %s --check-prefix=TRANSFORM
+// RUN: DISC_ENABLE_TRANSFORM_SCHEDULE=1 disc-opt -split-input-file -pass-pipeline='func.func(disc-fusion{gpu-enabled=false fusion-strategy=stitch})' %s -o - | FileCheck %s --check-prefix=TRANSFORM
 
 // BASE-LABEL: @custom_call_op
 // BASE-SAME: (%[[ARG0:.*]]: memref<?x?xf32, "cpu">, %[[ARG1:.*]]: memref<?x?xf32, "cpu">, %[[ARG2:.*]]: memref<?x?xf32, "cpu">, %[[ARG3:.*]]: memref<?x?xf32, "cpu">) -> memref<?x?xf32, "cpu">
@@ -29,7 +29,7 @@ func.func @matmul_nn(%arg0: memref<?x?xf32, "cpu">, %arg1: memref<?x?xf32, "cpu"
 
 // -----
 
-// TRANSFORM-LABEL: @matmul_nn
+// TRANSFORM-LABEL: @matmul_nn_const_weight
 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
@@ -43,4 +43,21 @@ func.func @matmul_nn_const_weight(%arg0: memref<1024x1024xf32, "cpu">, %arg1: me
   "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">
+}
+
+// -----
+
+// TRANSFORM-LABEL: @matmul_nn_const_weight_with_external_user
+func.func @matmul_nn_const_weight_with_external_user(%arg1: memref<1024x1024xf32, "cpu">,
+                                                     %arg2: memref<?x1024xf32, "cpu">, %arg3: memref<?x1024xf32, "cpu">) -> (memref<1024x1024xf32, "cpu">, memref<?x1024xf32, "cpu">) {
+  // TRANSFORM: lmhlo.constant
+  // TRANSFORM-NEXT: "lmhlo.fusion"() ({
+  // TRANSFORM-NEXT: lmhlo.dot_general
+  // TRANSFORM-NEXT: lmhlo.terminator
+  // TRANSFORM-NEXT: })
+  // TRANSFORM-SAME: disc.fusion_type = "kTransform"
+  // TRANSFORM-NEXT: return
+  "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 %arg1, %arg3 : memref<1024x1024xf32, "cpu">, memref<?x1024xf32, "cpu">
 }
\ No newline at end of file
diff --git a/tao_compiler/mlir/disc/transforms/tests/disc-duplicate-computation-for-fusion.mlir b/tao_compiler/mlir/disc/transforms/tests/disc-duplicate-computation-for-fusion.mlir
index ef041a56f82..5921e734a13 100644
--- a/tao_compiler/mlir/disc/transforms/tests/disc-duplicate-computation-for-fusion.mlir
+++ b/tao_compiler/mlir/disc/transforms/tests/disc-duplicate-computation-for-fusion.mlir
@@ -1,10 +1,11 @@
-// RUN: disc-opt --disc-duplicate-computation-for-fusion -split-input-file %s | FileCheck %s
+// RUN: disc-opt -split-input-file --disc-duplicate-computation-for-fusion %s | FileCheck %s --check-prefix=BASE
+// RUN: DISC_ENABLE_TRANSFORM_SCHEDULE=1 disc-opt -split-input-file --disc-duplicate-computation-for-fusion=gpu-enabled=false %s -o - | FileCheck %s --check-prefix=TRANSFORM
 
-// CHECK-LABEL: @main
+// BASE-LABEL: @main
 func.func @main(%arg0 : memref<?x?xf32, "gpu">, %arg1 : memref<f32, "gpu">, %arg2 : memref<2xi32, "cpu">) -> memref<?x?xf32, "gpu"> {
   // make sure dynamic_broadcast_in_dim is duplicated
-  // CHECK: lmhlo.dynamic_broadcast_in_dim
-  // CHECK: lmhlo.dynamic_broadcast_in_dim
+  // BASE: lmhlo.dynamic_broadcast_in_dim
+  // BASE: lmhlo.dynamic_broadcast_in_dim
   %c0 = arith.constant 0 : index
   %c1 = arith.constant 1 : index
   %0 = memref.dim %arg0, %c0 : memref<?x?xf32, "gpu">
@@ -24,11 +25,11 @@ func.func @main(%arg0 : memref<?x?xf32, "gpu">, %arg1 : memref<f32, "gpu">, %arg
 
 // -----
 
-// CHECK-LABEL: @main
+// BASE-LABEL: @main
 func.func @main(%arg0 : memref<?x128xf32, "gpu">, %arg1 : memref<128xf32, "gpu">, %arg2 : memref<1xi32, "cpu">) -> memref<?x128xf32, "gpu"> {
   // make sure dynamic_broadcast_in_dim is duplicated
-  // CHECK: lmhlo.dynamic_broadcast_in_dim
-  // CHECK: lmhlo.dynamic_broadcast_in_dim
+  // BASE: lmhlo.dynamic_broadcast_in_dim
+  // BASE: lmhlo.dynamic_broadcast_in_dim
   %c0 = arith.constant 0 : index
   "lmhlo.constant"(%arg1) {value = dense<0.1> : tensor<128xf32>} : (memref<128xf32, "gpu">) -> ()
   %0 = memref.dim %arg0, %c0 : memref<?x128xf32, "gpu">
@@ -43,4 +44,20 @@ func.func @main(%arg0 : memref<?x128xf32, "gpu">, %arg1 : memref<128xf32, "gpu">
   %6 = memref.alloc(%0) : memref<?x128xf32, "gpu">
   "lmhlo.add"(%2, %5, %6) : (memref<?x128xf32, "gpu">, memref<?x128xf32, "gpu">, memref<?x128xf32, "gpu">) -> ()
   return %6 :  memref<?x128xf32, "gpu">
+}
+
+// -----
+
+// TRANSFORM-LABEL: @duplicate_dot_weight
+func.func @duplicate_dot_weight(%arg0 : memref<128x128xf32, "cpu">, %arg1 : memref<?x128xf32, "cpu">,
+                                %arg2 : memref<?x128xf32, "cpu">, %arg3 : memref<?x128xf32, "cpu">, %arg4 : memref<?x128xf32, "cpu">) -> (memref<?x128xf32, "cpu">, memref<?x128xf32, "cpu">) {
+  // TRANSFORM: lmhlo.constant
+  // TRANSFORM-NEXT: memref.alloc
+  // TRANSFORM-NEXT: lmhlo.constant
+  // TRANSFORM-NEXT: lmhlo.dot_general
+  // TRANSFORM-NEXT: lmhlo.dot_general
+  "lmhlo.constant"(%arg0) {value = dense<0.1> : tensor<128x128xf32>} : (memref<128x128xf32, "cpu">) -> ()
+  "lmhlo.dot_general"(%arg1, %arg0, %arg2) {dot_dimension_numbers = #mhlo.dot<lhs_contracting_dimensions = [1], rhs_contracting_dimensions = [0]>} : (memref<?x128xf32, "cpu">, memref<128x128xf32, "cpu">, memref<?x128xf32, "cpu">) -> ()
+  "lmhlo.dot_general"(%arg3, %arg0, %arg4) {dot_dimension_numbers = #mhlo.dot<lhs_contracting_dimensions = [1], rhs_contracting_dimensions = [0]>} : (memref<?x128xf32, "cpu">, memref<128x128xf32, "cpu">, memref<?x128xf32, "cpu">) -> ()
+  return %arg2, %arg4 :  memref<?x128xf32, "cpu">, memref<?x128xf32, "cpu">
 }
\ No newline at end of file