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 670e5d2d1d1..27adb6ef701 100644
--- a/tao_compiler/mlir/disc/tools/disc-transform/TransformOps/TransformOpsExt.cc
+++ b/tao_compiler/mlir/disc/tools/disc-transform/TransformOps/TransformOpsExt.cc
@@ -1779,6 +1779,111 @@ DiagnosedSilenceableFailure ConvertPaddingPlaceholderToConstOp::applyToOne(
   return DiagnosedSilenceableFailure(success());
 }
 
+//===---------------------------------------------------------------------===//
+// LinalgEagerlyBackwardInitTensorOp
+//===---------------------------------------------------------------------===//
+
+namespace {
+
+OpOperand* backwardToFindValidValueToToReuse(OpOperand* v) {
+  if (!v || !llvm::hasSingleElement(v->get().getUsers())) return nullptr;
+  auto definingOp = v->get().getDefiningOp<linalg::LinalgOp>();
+  if (!definingOp) return nullptr;
+  int64_t resultIdx = v->get().cast<OpResult>().getResultNumber();
+  auto initOperand = definingOp.getDpsInitOperand(resultIdx);
+  if (!definingOp.payloadUsesValueFromOperand(initOperand)) return initOperand;
+  return backwardToFindValidValueToToReuse(initOperand);
+}
+
+}  // namespace
+
+DiagnosedSilenceableFailure LinalgEagerlyBackwardInitTensorOp::applyToOne(
+    Operation* target, SmallVectorImpl<Operation*>& results,
+    transform::TransformState& state) {
+  auto linalgOp = dyn_cast<linalg::LinalgOp>(target);
+  if (!linalgOp)
+    return mlir::emitDefiniteFailure(target, "applies only to linalgOp");
+
+  if (linalgOp.getNumParallelLoops() != linalgOp.getNumLoops())
+    return mlir::emitDefiniteFailure(target,
+                                     "not support reduction loop a.t.m.");
+
+  if (linalgOp.getNumDpsInits() != 1)
+    return mlir::emitDefiniteFailure(target,
+                                     "not support multi outputs a.t.m.");
+
+  auto initOperand = linalgOp.getDpsInitOperand(0);
+  auto affineMap = linalgOp.getMatchingIndexingMap(initOperand);
+  if (!affineMap.isPermutation())
+    return mlir::emitDefiniteFailure(
+        target, "only support permutation mapping for result");
+  if (linalgOp.payloadUsesValueFromOperand(initOperand))
+    return mlir::emitDefiniteFailure(
+        target, "Not support if the result is used by the payload.");
+
+  OpOperand* candidateOperand;
+  OpOperand* rootCandidateOperand;
+  for (auto operand : linalgOp.getDpsInputOperands()) {
+    if (linalgOp.getMatchingIndexingMap(operand) != affineMap) continue;
+    candidateOperand = operand;
+    rootCandidateOperand = backwardToFindValidValueToToReuse(candidateOperand);
+    if (rootCandidateOperand) break;
+  }
+
+  // Early return if no valid candidates.
+  if (!rootCandidateOperand) {
+    results.assign({target});
+    return DiagnosedSilenceableFailure(success());
+  }
+
+  OpBuilder b(target);
+  rootCandidateOperand->getOwner()->replaceUsesOfWith(
+      rootCandidateOperand->get(), initOperand->get());
+  SmallVector<Value> newInputs;
+  SmallVector<AffineMap> newIndexingMaps;
+  SmallVector<Value> newOutputs{candidateOperand->get()};
+  for (auto operand : linalgOp.getDpsInputOperands()) {
+    if (operand == candidateOperand) continue;
+    newInputs.push_back(operand->get());
+    newIndexingMaps.push_back(linalgOp.getMatchingIndexingMap(operand));
+  }
+  // for the new init operand
+  newIndexingMaps.push_back(affineMap);
+  SmallVector<StringRef> iterators = linalgOp.getIteratorTypesArray();
+  SmallVector<Type> resultTypes = linalgOp.hasTensorSemantics()
+                                      ? TypeRange(ValueRange(newOutputs))
+                                      : TypeRange{};
+  SmallVector<NamedAttribute> attrs;
+  for (auto attr : linalgOp->getAttrs()) {
+    if (attr.getName().getValue().starts_with("disc"))
+      attrs.push_back(std::move(attr));
+  }
+
+  auto bodyBuilder = [&](OpBuilder& innerBuilder, Location loc,
+                         ValueRange operands) {
+    BlockAndValueMapping mapping;
+    int64_t operandIdx = candidateOperand->getOperandNumber();
+    if (operandIdx > 0)
+      mapping.map(linalgOp.getBlock()->getArguments().take_front(operandIdx),
+                  operands.take_front(operandIdx));
+    if (operandIdx + 1 < linalgOp.getBlock()->getArguments().size())
+      mapping.map(
+          linalgOp.getBlock()->getArguments().drop_front(operandIdx + 1),
+          operands.drop_front(operandIdx));
+    mapping.map(linalgOp.getBlock()->getArgument(operandIdx), operands.back());
+    for (Operation& op : *linalgOp.getBlock()) {
+      innerBuilder.clone(op, mapping);
+    }
+  };
+
+  auto genericOp = b.create<linalg::GenericOp>(
+      linalgOp.getLoc(), resultTypes, newInputs, newOutputs, newIndexingMaps,
+      iterators, bodyBuilder, attrs);
+  target->replaceAllUsesWith(genericOp->getResults());
+  results.assign({genericOp.getOperation()});
+  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 8cb9eb07ac9..c1b50fb8671 100644
--- a/tao_compiler/mlir/disc/tools/disc-transform/TransformOps/TransformOpsExt.td
+++ b/tao_compiler/mlir/disc/tools/disc-transform/TransformOps/TransformOpsExt.td
@@ -535,4 +535,44 @@ def ConvertPaddingPlaceholderToConstOp : Op<Transform_Dialect, "disc.convert_pad
   }];
 }
 
+def LinalgEagerlyBackwardInitTensorOp : Op<Transform_Dialect, "disc.linalg.eagerly_backward_init_tensor",
+    [FunctionalStyleTransformOpTrait,
+     MemoryEffectsOpInterface,
+     TransformEachOpTrait,
+     TransformOpInterface]> {
+  let description = [{
+    Given a linalg op `A`, eagerly backward its init tensor to its producers. The function of such primitive
+    is to reuse underline buffer.
+
+    Eamples:
+    ```
+      %0 = linalg.generic { indexing_maps = [#map1, #map1, #map1], iterator_types = ["parallel", "parallel"]}
+        ins(%arg0, %arg1) outs(%arg2)
+      %1 = linalg.generic { indexing_maps = [#map1, #map1, #map1], iterator_types = ["parallel", "parallel"]}
+        ins(%0, %0) outs(%arg3)
+    ```
+    After conversion:
+    ```
+      %0 = linalg.generic { indexing_maps = [#map1, #map1], iterator_types = ["parallel", "parallel"]}
+        ins(%arg0, %arg1) outs(%arg3) // Note here use the final output buffer directly
+      %1 = linalg.generic { indexing_maps = [#map1], iterator_types = ["parallel", "parallel"]}
+        ins() outs(%0) // Note here: we promote %0 to be an output operand (thus can be read + write).
+    ```
+
+    Returns the result const op.
+  }];
+
+  let arguments = (ins PDL_Operation:$target);
+  let results = (outs PDL_Operation:$result);
+
+  let assemblyFormat = "$target attr-dict";
+  let cppNamespace = "::mlir::disc_ral::transform_dialect";
+  let extraClassDeclaration = [{
+    ::mlir::DiagnosedSilenceableFailure applyToOne(
+        ::mlir::Operation *target,
+        ::llvm::SmallVectorImpl<::mlir::Operation *> &results,
+        ::mlir::transform::TransformState &state);
+  }];
+}
+
 #endif // DISC_TRANSFORM_OPS_EXT


Test patch:
diff --git a/tao_compiler/mlir/disc/tools/disc-transform/transforms/tests/linalg-eagerly-backward-init-tensor.mlir b/tao_compiler/mlir/disc/tools/disc-transform/transforms/tests/linalg-eagerly-backward-init-tensor.mlir
new file mode 100644
index 00000000000..c4ef155b7ad
--- /dev/null
+++ b/tao_compiler/mlir/disc/tools/disc-transform/transforms/tests/linalg-eagerly-backward-init-tensor.mlir
@@ -0,0 +1,33 @@
+// RUN: disc-opt --disc-transform-dialect-interpreter -split-input-file %s | FileCheck %s
+
+#map = affine_map<(d0, d1) -> (d1)>
+#map1 = affine_map<(d0, d1) -> (d0, d1)>
+
+// CHECK-LABEL: @main
+// CHECK-SAME: (%[[ARG0:.*]]: tensor<?x3072xf32>, %[[ARG4:.*]]: tensor<?x768xf32>, %[[ARG6:.*]]: tensor<?x768xf32>)
+func.func @main(%arg0: tensor<?x3072xf32>, %arg4: tensor<?x768xf32>, %arg6: tensor<?x768xf32>) -> tensor<?x768xf32> {
+  %0 = disc_linalg_ext.constant_wrapper {disc.transform.name = "constant"} dense<-8.000000e-01> : tensor<3072x768xf32>
+  %1 = disc_linalg_ext.constant_wrapper {disc.transform.name = "constant_1"} dense<-1.000000e-01> : tensor<768xf32>
+  %cst = arith.constant 0.000000e+00 : f32
+  // CHECK: %[[T0:.*]] = linalg.fill
+  // CHECK-SAME: outs(%[[ARG6]] : tensor<?x768xf32>)
+  // CHECK-NEXT: %[[T1:.*]] = linalg.matmul
+  // CHECK-SAME: outs(%[[T0]] : tensor<?x768xf32>)
+  // CHECK-NEXT: %[[T2:.*]] = linalg.generic
+  // CHECK-SAME: outs(%[[T1]] : tensor<?x768xf32>)
+  // CHECK: return %[[T2]]
+  %2 = linalg.fill {disc.transform.name = "dot_general"} ins(%cst : f32) outs(%arg4 : tensor<?x768xf32>) -> tensor<?x768xf32>
+  %3 = linalg.matmul {disc.transform.name = "dot_general"} ins(%arg0, %0 : tensor<?x3072xf32>, tensor<3072x768xf32>) outs(%2 : tensor<?x768xf32>) -> tensor<?x768xf32>
+  %4 = linalg.generic {indexing_maps = [#map1, #map, #map1], iterator_types = ["parallel", "parallel"]} ins(%3, %1 : tensor<?x768xf32>, tensor<768xf32>) outs(%arg6 : tensor<?x768xf32>) attrs =  {disc.device = "cpu", disc.transform.name = "subtract"} {
+  ^bb0(%in: f32, %in_0: f32, %out: f32):
+    %6 = arith.subf %in, %in_0 : f32
+    linalg.yield %6 : f32
+  } -> tensor<?x768xf32>
+  return %4 : tensor<?x768xf32>
+}
+
+transform.structured.canonicalized_sequence failures(propagate) {
+^bb0(%arg0: !pdl.operation):
+  %0 = transform.structured.match attributes {disc.transform.name = "subtract"} in %arg0
+  %1 = transform.disc.linalg.eagerly_backward_init_tensor %0
+}