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 27adb6ef701..408d609971e 100644
--- a/tao_compiler/mlir/disc/tools/disc-transform/TransformOps/TransformOpsExt.cc
+++ b/tao_compiler/mlir/disc/tools/disc-transform/TransformOps/TransformOpsExt.cc
@@ -1821,8 +1821,8 @@ DiagnosedSilenceableFailure LinalgEagerlyBackwardInitTensorOp::applyToOne(
     return mlir::emitDefiniteFailure(
         target, "Not support if the result is used by the payload.");
 
-  OpOperand* candidateOperand;
-  OpOperand* rootCandidateOperand;
+  OpOperand* candidateOperand{nullptr};
+  OpOperand* rootCandidateOperand{nullptr};
   for (auto operand : linalgOp.getDpsInputOperands()) {
     if (linalgOp.getMatchingIndexingMap(operand) != affineMap) continue;
     candidateOperand = operand;
@@ -1884,6 +1884,206 @@ DiagnosedSilenceableFailure LinalgEagerlyBackwardInitTensorOp::applyToOne(
   return DiagnosedSilenceableFailure(success());
 }
 
+//===----------------------------------------------------------------------===//
+// DISCFuseIntoContainingOp
+//===----------------------------------------------------------------------===//
+
+namespace {
+
+// First, find the first "scf::ForeachThreadOp" user of `producerOp` and ensure
+/// it is exactly the `containingOp`, otherwise bail.
+/// Then, find the first "extract" user of the tied block argument and tile it
+/// right before its "extract" use. The tiled op is fused under the
+/// `containingOp`.
+/// Return this fused op on success or nullptr if anything fails.
+static Operation* tileAndFuseFirstExtractUseThroughContainingOpBlockArgument(
+    RewriterBase& rewriter, Diagnostic& diag, Operation* producerOp,
+    Operation* containingOp) {
+  LLVM_DEBUG(
+      llvm::dbgs() << "Try to fuse an extract use through block argument\n");
+
+  auto tileableProducer = dyn_cast<TilingInterface>(producerOp);
+  if (!tileableProducer) {
+    diag.attachNote(producerOp->getLoc())
+        << "producer is not a TileableInterface: " << *producerOp;
+    return nullptr;
+  }
+
+  // Search the first use by a "scf::ForOp" user.
+  scf::ForOp forOp;
+  auto itProducerUses =
+      llvm::find_if(tileableProducer->getUses(), [&](OpOperand& use) {
+        forOp = dyn_cast<scf::ForOp>(use.getOwner());
+        return forOp;
+      });
+  // If it's not from the containing op, return.
+  if (!forOp || forOp != containingOp) {
+    diag.attachNote(tileableProducer->getLoc())
+        << "could not find a use by the containing op: " << *tileableProducer;
+    return nullptr;
+  }
+
+  // Search the producer slices accessed within the containing
+  // operation.
+  // TODO: Generalize to more extract/insert/parallel_insert triples.
+  //   Maybe evolve into an interface.
+  OpOperand* pUse = &(*itProducerUses);
+  BlockArgument bbArg = forOp.getBody()->getArgument(
+      pUse->getOperandNumber() - forOp.getNumControlOperands() +
+      forOp.getNumInductionVars());
+
+  // Search the producer slices accessed within the containing operation.
+  // TODO: Generalize to more extract/insert/parallel_insert triples, maybe
+  // evolve into an interface.
+  auto itBBArgUsers = llvm::find_if(bbArg.getUsers(), [&](Operation* user) {
+    auto sliceOp = dyn_cast<tensor::ExtractSliceOp>(user);
+    return sliceOp && containingOp->isProperAncestor(sliceOp);
+  });
+
+  // Find a fusion opportunity.
+  if (itBBArgUsers == bbArg.getUsers().end()) {
+    diag.attachNote(containingOp->getLoc())
+        << "could not find fusion opportunity for bbArg: " << bbArg;
+    return nullptr;
+  }
+  auto sliceOpToTile = cast<tensor::ExtractSliceOp>(*itBBArgUsers);
+
+  // Try to fuse the producer in-place.
+  OpBuilder::InsertionGuard guard(rewriter);
+  rewriter.setInsertionPoint(sliceOpToTile);
+
+  // Replace the use in the tileableProducer before tiling: clone, replace and
+  // then tile.
+  int64_t resultNumber = pUse->get().cast<OpResult>().getResultNumber();
+  LLVM_DEBUG(llvm::dbgs() << "resultNumber: " << resultNumber << "\n");
+
+  // Gather destination tensors.
+  SmallVector<Value> destinationTensors;
+  if (failed(tensor::getOrCreateDestinations(
+          rewriter, tileableProducer->getLoc(), tileableProducer,
+          destinationTensors))) {
+    diag.attachNote(tileableProducer->getLoc())
+        << "failed to get destination tensors for: " << *tileableProducer;
+    return nullptr;
+  }
+
+  BlockAndValueMapping bvm;
+  bvm.map(destinationTensors[resultNumber], bbArg);
+  auto tileableProducerClone =
+      cast<TilingInterface>(rewriter.clone(*tileableProducer, bvm));
+  auto scopeGuard =
+      llvm::make_scope_exit([&]() { rewriter.eraseOp(tileableProducerClone); });
+
+  // Tile the producer.
+  FailureOr<Value> tiledProducer =
+      tileableProducerClone.generateResultTileValue(
+          rewriter, resultNumber, sliceOpToTile.getMixedOffsets(),
+          sliceOpToTile.getMixedSizes());
+  if (failed(tiledProducer)) {
+    diag.attachNote(tileableProducer->getLoc())
+        << "failed to tile producer op: " << *tileableProducer;
+    return nullptr;
+  }
+  LLVM_DEBUG(llvm::dbgs() << "tiledProducer: " << *tiledProducer << "\n");
+
+  // Replace the extract op.
+  Operation* fusedOp = tiledProducer->getDefiningOp();
+  rewriter.replaceOp(sliceOpToTile, fusedOp->getResult(resultNumber));
+
+  // Replace the use in containingOp.
+  rewriter.updateRootInPlace(containingOp, [&]() {
+    containingOp->setOperand(pUse->getOperandNumber(),
+                             destinationTensors[resultNumber]);
+  });
+
+  return fusedOp;
+}
+
+}  // namespace
+
+void DISCFuseIntoContainingOp::getEffects(
+    SmallVectorImpl<MemoryEffects::EffectInstance>& effects) {
+  transform::consumesHandle({this->getOperation()->getOperand(0)}, effects);
+  transform::onlyReadsHandle(this->getOperands().drop_front(), effects);
+  transform::producesHandle(this->getOperation()->getResults(), effects);
+  transform::modifiesPayload(effects);
+}
+
+DiagnosedSilenceableFailure DISCFuseIntoContainingOp::apply(
+    transform::TransformResults& results, transform::TransformState& state) {
+  SmallVector<Operation*> fusedOps;
+  ArrayRef<Operation*> producerOps = state.getPayloadOps(getProducerOp());
+  // If nothing to fuse, propagate success.
+  if (producerOps.empty()) {
+    results.set(getFusedOp().cast<OpResult>(), SmallVector<mlir::Operation*>{});
+    return DiagnosedSilenceableFailure::success();
+  }
+  ArrayRef<Operation*> containingOps = state.getPayloadOps(getContainingOp());
+  if (containingOps.size() != 1) {
+    return emitDefiniteFailure()
+           << "requires exactly one containing_op handle (got "
+           << containingOps.size() << ")";
+  }
+  Operation* containingOp = containingOps.front();
+
+  // Helper function to find the next producer that should be fused. Take any
+  // producer that has a use inside the containing op.
+  SmallVector<Operation*> remainingProducers(producerOps.begin(),
+                                             producerOps.end());
+  auto getNextProducer = [&]() -> FailureOr<Operation*> {
+    for (const auto& it : enumerate(remainingProducers)) {
+      Operation* producerOp = it.value();
+      // The containing op may be a user of producerOp: use isAncestor.
+      int64_t numUsesInContainingOp = llvm::count_if(
+          producerOp->getUsers(),
+          [&](Operation* op) { return containingOp->isAncestor(op); });
+      // TODO: When resolving the TODO below (no duplicate ops), take an op
+      // that has no use among the remaining producers. This is a topological
+      // sorting.
+      if (numUsesInContainingOp > 0) {
+        if (numUsesInContainingOp == 1)
+          remainingProducers.erase(remainingProducers.begin() + it.index());
+        return producerOp;
+      }
+    }
+    return failure();
+  };
+
+  IRRewriter rewriter(getContext());
+  while (!remainingProducers.empty()) {
+    auto nextProducer = getNextProducer();
+    if (failed(nextProducer)) {
+      results.set(getFusedOp().cast<OpResult>(), ArrayRef<Operation*>());
+      Diagnostic diag(containingOp->getLoc(), DiagnosticSeverity::Remark);
+      diag << "could not find next producer to fuse into container";
+      return DiagnosedSilenceableFailure::silenceableFailure(std::move(diag));
+    }
+
+    Operation* producerOp = *nextProducer;
+
+    // Default diagnostic, to be complemented with more failure information.
+    Diagnostic diag(producerOp->getLoc(), DiagnosticSeverity::Remark);
+    diag << "could not fuse " << *producerOp << " into " << *containingOp;
+
+    Operation* tiledContainingOpOperand =
+        tileAndFuseFirstExtractUseThroughContainingOpBlockArgument(
+            rewriter, diag, producerOp, containingOp);
+    if (tiledContainingOpOperand) {
+      LLVM_DEBUG(llvm::dbgs()
+                 << "\nFused an extract use through block argument\n"
+                 << *containingOp);
+      fusedOps.push_back(tiledContainingOpOperand);
+      continue;
+    }
+
+    results.set(getFusedOp().cast<OpResult>(), ArrayRef<Operation*>());
+    return DiagnosedSilenceableFailure::silenceableFailure(std::move(diag));
+  }
+
+  results.set(getFusedOp().cast<OpResult>(), fusedOps);
+  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 c1b50fb8671..8ba3ef5259f 100644
--- a/tao_compiler/mlir/disc/tools/disc-transform/TransformOps/TransformOpsExt.td
+++ b/tao_compiler/mlir/disc/tools/disc-transform/TransformOps/TransformOpsExt.td
@@ -575,4 +575,53 @@ def LinalgEagerlyBackwardInitTensorOp : Op<Transform_Dialect, "disc.linalg.eager
   }];
 }
 
+// mainly copied from mlir community, adding support to fuse through
+// iter_args for normal scf.for ops.
+
+def DISCFuseIntoContainingOp :
+    Op<Transform_Dialect, "disc.fuse_into_containing_op",
+      [DeclareOpInterfaceMethods<MemoryEffectsOpInterface>,
+       DeclareOpInterfaceMethods<TransformOpInterface>]> {
+  let description = [{Fuse a producer into a containing operation.}];
+
+  let summary = [{
+    Fuses the `producer_op` into the `containing_op`.
+    Returns a handle to the fused ops.
+
+    The producer is typically a slice of a tileable op (i.e., implements
+    TilingInterface). In that case, this transform computes the accessed
+    producer slice inside of the containing op ("tile and fuse"). Otherwise,
+    the entire producer is cloned inside the containing op ("clone and fuse").
+
+    The containing op handle must be associated with exactly one payload op. The
+    producer op handle may be associated with multiple payload ops. This
+    transform fuses producers one-by-one, always picking an unspecified producer
+    that has at least one use inside the containing op among the
+    producers.
+
+    Note: If a producer has multiple uses inside the containing op, it is
+    currently tiled and/or cloned multiple times into the containing op.
+    TODO: Reuse already fused OpResults instead of tiling/cloning a second time
+    when possible. Fuse producers according to a topological sorting to achieve
+    the largest amount of reuse.
+
+    #### Return modes
+
+    If at least one producer could not be fused, this operation fails silently.
+    This is the case when tiling fails or when no producer op could be found
+    among the remaining producers that has at least one use within the
+    containing op. I.e., "producers" that are not consumed within the containing
+    op are rejected by this operation.
+
+    This operation reads and frees the producer handle.
+    This operation reads the containing op handle.
+  }];
+
+  let arguments = (ins PDL_Operation:$producer_op,
+                       PDL_Operation:$containing_op);
+  let results = (outs PDL_Operation:$fused_op);
+  let assemblyFormat = "$producer_op `into` $containing_op attr-dict";
+  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/fuse-into-containing-op.mlir b/tao_compiler/mlir/disc/tools/disc-transform/transforms/tests/fuse-into-containing-op.mlir
new file mode 100644
index 00000000000..465f793a59c
--- /dev/null
+++ b/tao_compiler/mlir/disc/tools/disc-transform/transforms/tests/fuse-into-containing-op.mlir
@@ -0,0 +1,48 @@
+// RUN: disc-opt --disc-transform-dialect-interpreter -split-input-file %s | FileCheck %s
+
+#map = affine_map<(d0)[s0] -> (-d0 + s0, 288)>
+#map1 = affine_map<(d0, d1) -> (d1)>
+#map2 = affine_map<(d0, d1) -> (d0, d1)>
+
+// CHECK: @main
+// CHECK-SAME: (%[[ARG0:.*]]: tensor<?x3072xf32>, %[[ARG1:.*]]: tensor<?x768xf32>)
+func.func @main(%arg0: tensor<?x3072xf32>, %arg1: tensor<?x768xf32>) -> tensor<?x768xf32> {
+  %cst = arith.constant 0.000000e+00 : f32
+  %c0 = arith.constant 0 : index
+  %c288 = arith.constant 288 : index
+  %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>
+  %2 = linalg.fill {disc.transform.name = "dot_general"} ins(%cst : f32) outs(%arg1 : 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>
+  %dim = tensor.dim %arg0, %c0 : tensor<?x3072xf32>
+  // CHECK: %[[T0:.*]] = scf.for
+  // CHECK-SAME: iter_args(%[[T1:.*]] = %[[ARG1]])
+  %4 = scf.for %arg2 = %c0 to %dim step %c288 iter_args(%arg3 = %3) -> (tensor<?x768xf32>) {
+    // CHECK: %[[T2:.*]] = tensor.extract_slice %[[T1]]
+    // CHECK: %[[T3:.*]] = linalg.fill
+    // CHECK-SAME: outs(%[[T2]] : tensor<?x768xf32>)
+    // CHECK-NEXT: %[[T4:.*]] = linalg.matmul
+    // CHECK-SAME: outs(%[[T3]] : tensor<?x768xf32>)
+    // CHECK-NEXT: %[[T5:.*]] = linalg.generic
+    // CHECK-SAME: outs(%[[T4]] : tensor<?x768xf32>)
+    %5 = affine.min #map(%arg2)[%dim]
+    %extracted_slice = tensor.extract_slice %arg3[%arg2, 0] [%5, 768] [1, 1] : tensor<?x768xf32> to tensor<?x768xf32>
+    %6 = linalg.generic {indexing_maps = [#map1, #map2], iterator_types = ["parallel", "parallel"]} ins(%1 : tensor<768xf32>) outs(%extracted_slice : tensor<?x768xf32>) attrs =  {disc.device = "cpu", disc.transform.name = "subtract"} {
+    ^bb0(%in: f32, %out: f32):
+      %7 = arith.subf %out, %in : f32
+      linalg.yield %7 : f32
+    } -> tensor<?x768xf32>
+    %inserted_slice = tensor.insert_slice %6 into %arg3[%arg2, 0] [%5, 768] [1, 1] : tensor<?x768xf32> into tensor<?x768xf32>
+    scf.yield %inserted_slice : tensor<?x768xf32>
+  }
+  return %4 : tensor<?x768xf32>
+}
+transform.structured.canonicalized_sequence failures(propagate) {
+^bb0(%arg0: !pdl.operation):
+  %0 = transform.structured.match ops{["scf.for"]} in %arg0
+  %1 = transform.structured.match ops{["linalg.matmul"]}  in %arg0
+  %2 = transform.structured.match ops{["linalg.fill"]}  in %arg0
+  %3 = transform.disc.fuse_into_containing_op %1 into %0
+  %4 = transform.disc.fuse_into_containing_op %2 into %0
+  %5 = transform.disc.apply_patterns %arg0 {canonicalization}
+}
\ No newline at end of file