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:
e2e model test on: Bert Base (TF) and Albert (PyTorch),  on g6r, using single thread. Note that we only have one default schedule for all shape and the schedule is known to be less performant when n or k is large, thus the initial performance is supposed to be improved when we support schedule selection logic.

Bert Base (TF) 

input | TF 2.8(s) | DISC-ACL(s) | DISC-Transform(s) | speedup (DISC-transform / DISC-ACL)
-- | -- | -- | -- | --
(1, 128) | 0.742 | 0.638 | 0.656 | 97.3%
(2, 128) | 1.41 | 1.24 | 1.27 | 97.6%
(4, 128) | 2.85 | 2.36 | 2.55 | 92.5%
(8, 128) | 5.84 | 4.68 | 5.07 | 92.3%
(16, 128) | 11.9 | 9.55 | 10.2 | 93.6%
----


Albert (PyTorch)

input | TorchScript | OnnxRuntime | DISC-ACL | DISC-Transform
-- | -- | -- | -- | --
(2, 12) | 0.197 | 0.140 | 0.117 | 0.139



Fix patch:
diff --git a/tao_compiler/mlir/disc/tools/disc-transform/TransformOps/TransformOpsExt.cc b/tao_compiler/mlir/disc/tools/disc-transform/TransformOps/TransformOpsExt.cc
index 4f55b0de279..b105d70bb50 100644
--- a/tao_compiler/mlir/disc/tools/disc-transform/TransformOps/TransformOpsExt.cc
+++ b/tao_compiler/mlir/disc/tools/disc-transform/TransformOps/TransformOpsExt.cc
@@ -2374,6 +2374,93 @@ DiagnosedSilenceableFailure ReductionOutputFuseOp::apply(
   return DiagnosedSilenceableFailure(success());
 }
 
+//===---------------------------------------------------------------------===//
+// ReductionInputFuseOp
+//===---------------------------------------------------------------------===//
+
+void ReductionInputFuseOp::getEffects(
+    SmallVectorImpl<MemoryEffects::EffectInstance>& effects) {
+  transform::consumesHandle({this->getOperation()->getOperand(0)}, effects);
+  transform::consumesHandle({this->getOperation()->getOperand(1)}, effects);
+  transform::producesHandle(this->getOperation()->getResults(), effects);
+  transform::modifiesPayload(effects);
+}
+
+DiagnosedSilenceableFailure ReductionInputFuseOp::apply(
+    transform::TransformResults& results, transform::TransformState& state) {
+  ArrayRef<Operation*> targetOps = state.getPayloadOps(getTarget());
+  if (targetOps.size() != 1)
+    return mlir::emitDefiniteFailure(this->getOperation(),
+                                     "expect only one target but got ")
+           << targetOps.size();
+  Operation* target = targetOps[0];
+  auto linalgOp = dyn_cast<linalg::LinalgOp>(target);
+  if (!linalgOp)
+    return mlir::emitDefiniteFailure(this->getOperation(),
+                                     "expect the target is a linalg op");
+
+  ArrayRef<Operation*> loopOps = state.getPayloadOps(getLoop());
+  if (loopOps.size() != 1)
+    return mlir::emitDefiniteFailure(this->getOperation(),
+                                     "expect only one loop but got ")
+           << loopOps.size();
+  auto forOp = dyn_cast<scf::ForOp>(loopOps[0]);
+  if (!forOp)
+    return mlir::emitDefiniteFailure(this->getOperation(),
+                                     "expect the loop to a for op");
+
+  if (linalgOp->getNumResults() != 1)
+    return mlir::emitDefiniteFailure(this->getOperation(),
+                                     "expect the target has one result");
+
+  auto iterOperands = forOp.getIterOperands();
+  auto argIt = llvm::find(iterOperands, linalgOp->getResult(0));
+  if (argIt == iterOperands.end())
+    return mlir::emitDefiniteFailure(
+        this->getOperation(),
+        "expect the result of target is the iter arg of the loop.");
+  Value iterArg = forOp.getRegionIterArg(argIt - iterOperands.begin());
+  auto iterArgUsers = iterArg.getUsers();
+  auto userIt = llvm::find_if(iterArgUsers, [&](Operation* user) {
+    auto sliceOp = dyn_cast<tensor::ExtractSliceOp>(user);
+    return sliceOp && forOp->isProperAncestor(sliceOp);
+  });
+  if (userIt == iterArgUsers.end())
+    return mlir::emitDefiniteFailure(
+        this->getOperation(),
+        "could not find candidate extract_slice of target in the loop.");
+  auto sliceOp = cast<tensor::ExtractSliceOp>(*userIt);
+
+  OpBuilder b(sliceOp);
+  forOp->setOperand(
+      forOp.getNumControlOperands() + (argIt - iterOperands.begin()),
+      linalgOp.getDpsInitOperands().front()->get());
+  auto clonedTarget = b.clone(*target);
+  clonedTarget->replaceUsesOfWith(linalgOp.getDpsInitOperands().front()->get(),
+                                  iterArg);
+  auto tileableTarget = cast<TilingInterface>(clonedTarget);
+  FailureOr<Value> tiledTarget = tileableTarget.generateResultTileValue(
+      b, 0, sliceOp.getMixedOffsets(), sliceOp.getMixedSizes());
+  if (failed(tiledTarget)) {
+    return mlir::emitDefiniteFailure(this->getOperation(),
+                                     "failed to tile target op\n");
+  }
+  Operation* tiledOp = tiledTarget->getDefiningOp();
+  Value pred =
+      b.create<arith::CmpIOp>(tiledOp->getLoc(), arith::CmpIPredicate::eq,
+                              forOp.getInductionVar(), forOp.getLowerBound());
+  auto conditionalGenericOp = convertGenericOpToConditionalGenericOp(
+      b, pred, cast<linalg::LinalgOp>(tiledOp));
+  if (!conditionalGenericOp) {
+    return mlir::emitDefiniteFailure(
+        this->getOperation(), "failed to build conditional_generic op\n");
+  }
+  sliceOp->replaceAllUsesWith(conditionalGenericOp->getResults());
+  results.set(getTiledTarget().cast<OpResult>(), {conditionalGenericOp});
+  results.set(getFusedLoop().cast<OpResult>(), {forOp});
+  return DiagnosedSilenceableFailure(success());
+}
+
 }  // namespace transform_dialect
 
 void registerTransformDialectCommonExtension(DialectRegistry& registry) {
diff --git a/tao_compiler/mlir/disc/tools/disc-transform/TransformOps/TransformOpsExt.td b/tao_compiler/mlir/disc/tools/disc-transform/TransformOps/TransformOpsExt.td
index 865848c17de..872ea1e6158 100644
--- a/tao_compiler/mlir/disc/tools/disc-transform/TransformOps/TransformOpsExt.td
+++ b/tao_compiler/mlir/disc/tools/disc-transform/TransformOps/TransformOpsExt.td
@@ -683,4 +683,50 @@ def ReductionOutputFuseOp : Op<Transform_Dialect, "disc.reduction_output_fuse",
   let cppNamespace = "::mlir::disc_ral::transform_dialect";
 }
 
+def ReductionInputFuseOp : Op<Transform_Dialect, "disc.reduction_input_fuse",
+    [DeclareOpInterfaceMethods<MemoryEffectsOpInterface>,
+     DeclareOpInterfaceMethods<TransformOpInterface>]> {
+  let description = [{
+    Given a linalg op, input fuse it into the reduction loop.
+
+    Returns the fused linalg op and the new loop op.
+
+    Example:
+     convert from:
+     ```
+      %0 = linalg.fill(%init)
+      %1 = scf.for %arg1 = %c0 to %c768 step %c512 iter_args(%arg0 = %0) -> (tensor<?x3072xf32>) {
+        ...
+        %t0 = tensor.extract_slice %arg0[0, 0] [...] [1, 1] : tensor<?x3072xf32> to tensor<?x3072xf32>
+        %t1 = linalg.matmul {disc.transform.name = "dot_general"} ins(...) outs(%t0 : tensor<?x3072xf32>) -> tensor<?x3072xf32>
+        %t2 = tensor.insert_slice %t1 into %arg0[0, 0] [...] [1, 1] : tensor<?x3072xf32> into tensor<?x3072xf32>
+        scf.yield %t2 : tensor<?x3072xf32>
+      }
+      use(%0)
+     ```
+     to:
+     ```
+      %1 = scf.for %arg1 = %c0 to %c768 step %c512 iter_args(%arg0 = %0) -> (tensor<?x3072xf32>) {
+        ...
+        %pred = ... // is first iteration?
+        %t0 = disc_linalg_ext.conditional_generic ins(%pred) outs(%arg0)
+        %t1 = tensor.extract_slice %t0[0, 0] [...] [1, 1] : tensor<?x3072xf32> to tensor<?x3072xf32>
+        %t2 = linalg.matmul {disc.transform.name = "dot_general"} ins(...) outs(%t1 : tensor<?x3072xf32>) -> tensor<?x3072xf32>
+        %t3 = tensor.insert_slice %t2 into %arg0[0, 0] [...] [1, 1] : tensor<?x3072xf32> into tensor<?x3072xf32>
+        scf.yield %t3 : tensor<?x3072xf32>
+      }
+      use(%0)
+     ```
+  }];
+
+  let arguments = (ins PDL_Operation:$target,
+                       PDL_Operation:$loop);
+  let results = (outs PDL_Operation:$tiled_target, PDL_Operation:$fused_loop);
+
+  let assemblyFormat = [{
+    attr-dict $target `into` $loop
+  }];
+  let cppNamespace = "::mlir::disc_ral::transform_dialect";
+}
+
 #endif // DISC_TRANSFORM_OPS_EXT


Test patch:
diff --git a/tao_compiler/mlir/disc/tools/disc-transform/transforms/tests/reduction-input-fuse.mlir b/tao_compiler/mlir/disc/tools/disc-transform/transforms/tests/reduction-input-fuse.mlir
new file mode 100644
index 00000000000..d7a5172b8e8
--- /dev/null
+++ b/tao_compiler/mlir/disc/tools/disc-transform/transforms/tests/reduction-input-fuse.mlir
@@ -0,0 +1,47 @@
+// RUN: disc-opt --disc-transform-dialect-interpreter -split-input-file %s | FileCheck %s
+
+#map = affine_map<(d0, d1) -> (d0, d1)>
+#map1 = affine_map<(d0, d1) -> (512, d0 - d1)>
+
+// CHECK-LABEL: @test_reduction_input_fuse
+// CHECK-SAME: (%[[ARG0:.*]]: tensor<?x?xf32>, %[[ARG1:.*]]: tensor<?x?xf32>, %[[ARG2:.*]]: tensor<?x?xf32>)
+func.func @test_reduction_input_fuse(%arg0 : tensor<?x?xf32>, %arg1 : tensor<?x?xf32>, %arg2 : tensor<?x?xf32>) -> tensor<?x?xf32> {
+  %c0 = arith.constant 0 : index
+  %c1 = arith.constant 1 : index
+  %c512 = arith.constant 512 : index
+  // CHECK: %[[M:.*]] = tensor.dim %[[ARG0]], %c0 : tensor<?x?xf32>
+  // CHECK: %[[K:.*]] = tensor.dim %[[ARG0]], %c1 : tensor<?x?xf32>
+  // CHECK: %[[N:.*]] = tensor.dim %[[ARG1]], %c1 : tensor<?x?xf32>
+  %m = tensor.dim %arg0, %c0 : tensor<?x?xf32>
+  %k = tensor.dim %arg0, %c1 : tensor<?x?xf32>
+  %n = tensor.dim %arg1, %c1 : tensor<?x?xf32>
+  %cst = arith.constant 0.000000e+00 : f32
+  %t0 = linalg.fill {disc.transform.name = "fill"} ins(%cst : f32) outs(%arg0 : tensor<?x?xf32>) -> tensor<?x?xf32>
+  // CHECK: scf.for %[[IV:.*]] = %c0 to %[[K]] step %c512 iter_args(%[[V0:.*]] = %[[ARG0]])
+  %t1 = scf.for %iv = %c0 to %k step %c512 iter_args(%v0 = %t0) -> (tensor<?x?xf32>) {
+    %r0 = affine.min #map1(%k, %iv)
+    // CHECK: %[[R1:.*]] = tensor.extract_slice %[[ARG0]]
+    // CHECK: %[[R2:.*]] = tensor.extract_slice %[[ARG1]]
+    // CHECK: %[[R3:.*]] = tensor.extract_slice %[[V0]]
+    // CHECK: %[[R4:.*]] = disc_linalg_ext.conditional_generic
+    // CHECK-SAME: outs(%[[R3]] : tensor<?x?xf32>)
+    // CHECK: %[[R5:.*]] = linalg.matmul
+    // CHECK-SAME: ins(%[[R1]], %[[R2]] : tensor<?x?xf32>, tensor<?x?xf32>)
+    // CHECK-SAME: outs(%[[R4]] : tensor<?x?xf32>)
+    %r1 = tensor.extract_slice %arg0[0, %iv] [%m, %r0] [1, 1] : tensor<?x?xf32> to tensor<?x?xf32>
+    %r2 = tensor.extract_slice %arg1[%iv, 0] [%r0, %n] [1, 1] : tensor<?x?xf32> to tensor<?x?xf32>
+    %r3 = tensor.extract_slice %v0[0, 0] [%m, %n] [1, 1] : tensor<?x?xf32> to tensor<?x?xf32>
+    %r4 = linalg.matmul {disc.transform.name = "dot_general"} ins(%r1, %r2 : tensor<?x?xf32>, tensor<?x?xf32>) outs(%r3 : tensor<?x?xf32>) -> tensor<?x?xf32>
+    %r5 = tensor.insert_slice %r4 into %v0[0, 0] [%m, %n] [1, 1] : tensor<?x?xf32> into tensor<?x?xf32>
+    scf.yield %r5 : tensor<?x?xf32>
+  }
+  return %t1 : tensor<?x?xf32>
+}
+
+
+transform.structured.canonicalized_sequence failures(propagate) {
+^bb0(%arg0: !pdl.operation):
+  %0 = transform.structured.match attributes {disc.transform.name = "fill"} in %arg0
+  %1 = transform.structured.match ops{["scf.for"]} in %arg0
+  %2, %3 = transform.disc.reduction_input_fuse %0 into %1
+}
\ No newline at end of file