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/tools/disc-transform/TransformOps/TransformOpsExt.cc b/tao_compiler/mlir/disc/tools/disc-transform/TransformOps/TransformOpsExt.cc
index aba9ce65fb3..a20aa36a96a 100644
--- a/tao_compiler/mlir/disc/tools/disc-transform/TransformOps/TransformOpsExt.cc
+++ b/tao_compiler/mlir/disc/tools/disc-transform/TransformOps/TransformOpsExt.cc
@@ -128,6 +128,46 @@ OneShotBufferizationOptions getBufferizationOptions() {
   return options;
 }
 
+/// Pattern to convert tensor.empty to bufferizaiton::AllocTensorOp.
+struct EmptyTensorLoweringPattern : public OpRewritePattern<tensor::EmptyOp> {
+  using OpRewritePattern<tensor::EmptyOp>::OpRewritePattern;
+
+  LogicalResult matchAndRewrite(tensor::EmptyOp op,
+                                PatternRewriter& rewriter) const override {
+    rewriter.replaceOpWithNewOp<bufferization::AllocTensorOp>(
+        op, op.getType(), op.getDynamicSizes());
+    return success();
+  }
+};
+
+LogicalResult bufferizeTensorEmptyOps(
+    transform::TransformState& state, ModuleOp moduleOp,
+    const OneShotBufferizationOptions& options) {
+  /// 1, first convert tensor.emtpy -> TensorAlloc
+  RewritePatternSet patterns(moduleOp->getContext());
+  patterns.add<EmptyTensorLoweringPattern>(patterns.getContext());
+  TrackingListener listener(state);
+  GreedyRewriteConfig config;
+  LogicalResult result = applyPatternsAndFoldGreedily(
+      state.getTopLevel(), std::move(patterns), config, &listener);
+  LogicalResult listenerResult = listener.checkErrorState();
+  if (failed(result) || failed(listenerResult)) {
+    return moduleOp->emitError() << "failed to bufferize tensor.empty\n";
+  }
+  /// 2, some TensorAlloc can be folded. Try to do some optimizations in such
+  /// case.
+  IRRewriter rewriter(moduleOp->getContext());
+  OneShotAnalysisState oneShotState(moduleOp, options);
+  if (failed(bufferization::analyzeOp(moduleOp, oneShotState)) ||
+      failed(bufferization::insertSliceAnchoredAllocTensorEliminationStep(
+          rewriter, moduleOp, oneShotState))) {
+    return moduleOp->emitError()
+           << "failed to analyze the module for bufferization\n";
+  }
+
+  return success();
+}
+
 }  // namespace
 
 DiagnosedSilenceableFailure DISCBufferizeOp::apply(
@@ -140,6 +180,13 @@ DiagnosedSilenceableFailure DISCBufferizeOp::apply(
 
   auto moduleOp = cast<ModuleOp>(payload.front());
   auto options = getBufferizationOptions();
+
+  // Bufferize tensor.empty
+  if (failed(bufferizeTensorEmptyOps(state, moduleOp, options))) {
+    return mlir::emitDefiniteFailure(state.getTopLevel(),
+                                     "failed to bufferize tensor.empty.");
+  }
+
   if (failed(bufferization::runOneShotModuleBufferize(moduleOp, options)))
     return mlir::emitDefiniteFailure(state.getTopLevel(),
                                      "bufferization failed.");


Test patch:
diff --git a/tao_compiler/mlir/disc/tools/disc-transform/transforms/tests/disc-bufferize.mlir b/tao_compiler/mlir/disc/tools/disc-transform/transforms/tests/disc-bufferize.mlir
new file mode 100644
index 00000000000..da6f960d807
--- /dev/null
+++ b/tao_compiler/mlir/disc/tools/disc-transform/transforms/tests/disc-bufferize.mlir
@@ -0,0 +1,15 @@
+// RUN: disc-opt --disc-transform-dialect-interpreter -split-input-file %s | FileCheck %s
+
+// CHECK-LABEL: @tensor.empty
+func.func @tensor.empty(%arg0 : index) -> tensor<?x?xf32> {
+  %cst = arith.constant 0.000000e+00 : f32
+  // CHECK: %[[T0:.*]] = memref.alloc({{.*}})
+  // CHECK-NEXT: return %[[T0]]
+  %0 = tensor.empty(%arg0, %arg0) : tensor<?x?xf32>
+  return %0 : tensor<?x?xf32>
+}
+
+transform.structured.canonicalized_sequence failures(propagate) {
+^bb1(%arg1: !pdl.operation):
+  transform.disc.bufferize %arg1
+}
\ No newline at end of file