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 a20aa36a96a..c12d1baf3e7 100644
--- a/tao_compiler/mlir/disc/tools/disc-transform/TransformOps/TransformOpsExt.cc
+++ b/tao_compiler/mlir/disc/tools/disc-transform/TransformOps/TransformOpsExt.cc
@@ -23,6 +23,8 @@
 #include "mlir/Dialect/Bufferization/Transforms/AllocTensorElimination.h"
 #include "mlir/Dialect/Bufferization/Transforms/OneShotAnalysis.h"
 #include "mlir/Dialect/Bufferization/Transforms/OneShotModuleBufferize.h"
+#include "mlir/Dialect/Bufferization/Transforms/Passes.h"
+#include "mlir/Dialect/Func/IR/FuncOps.h"
 #include "mlir/Dialect/GPU/IR/GPUDialect.h"
 #include "mlir/Dialect/Linalg/TransformOps/LinalgTransformOps.h"
 #include "mlir/Dialect/Linalg/Transforms/Transforms.h"
@@ -53,6 +55,19 @@ using bufferization::OneShotBufferizationOptions;
 
 namespace {
 
+bool betterToUseAlloca(ShapedType type) {
+  constexpr unsigned kMaximumSizeInBytes = 64 * 1024;  // 64kB
+  constexpr unsigned kBitwidthOfIndexType = 64;
+
+  if (!type || !type.hasStaticShape()) return false;
+
+  // For index types, use the provided size, as the type does not know.
+  unsigned int bitwidth = type.getElementType().isIndex()
+                              ? kBitwidthOfIndexType
+                              : type.getElementTypeBitWidth();
+  return type.getNumElements() * bitwidth <= kMaximumSizeInBytes * 8;
+}
+
 // Allocation callbacks to use with upstream comprehensive bufferization
 static FailureOr<Value> comprehenciveBufferizeAllocationFn(
     OpBuilder& builder, Location loc, MemRefType memRefType,
@@ -66,7 +81,6 @@ static FailureOr<Value> comprehenciveBufferizeAllocationFn(
 LogicalResult comprehensiveBufferizeDeallocationFn(OpBuilder& builder,
                                                    Location loc,
                                                    Value allocation) {
-  builder.create<memref::DeallocOp>(loc, allocation);
   return success();
 }
 
@@ -111,6 +125,8 @@ OneShotBufferizationOptions getBufferizationOptions() {
   options.allocationFn = comprehenciveBufferizeAllocationFn;
   options.deallocationFn = comprehensiveBufferizeDeallocationFn;
   options.memCpyFn = comprehensiveBufferizeCopyFn;
+  options.bufferAlignment = 64;
+  options.createDeallocs = false;
   options.bufferizeFunctionBoundaries = true;
   options.allowReturnAllocs = true;
   options.functionBoundaryTypeConversion =
@@ -191,6 +207,16 @@ DiagnosedSilenceableFailure DISCBufferizeOp::apply(
     return mlir::emitDefiniteFailure(state.getTopLevel(),
                                      "bufferization failed.");
 
+  PassManager pm(getContext());
+  pm.addNestedPass<func::FuncOp>(
+      bufferization::createPromoteBuffersToStackPass([](Value alloc) {
+        return betterToUseAlloca(alloc.getType().cast<ShapedType>());
+      }));
+  pm.addNestedPass<func::FuncOp>(bufferization::createBufferDeallocationPass());
+
+  if (failed(pm.run(state.getTopLevel())))
+    return DiagnosedSilenceableFailure::definiteFailure();
+
   return DiagnosedSilenceableFailure::success();
 }
 


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
index da6f960d807..b6c88e21ca9 100644
--- 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
@@ -9,6 +9,85 @@ func.func @tensor.empty(%arg0 : index) -> tensor<?x?xf32> {
   return %0 : tensor<?x?xf32>
 }
 
+transform.structured.canonicalized_sequence failures(propagate) {
+^bb1(%arg1: !pdl.operation):
+  transform.disc.bufferize %arg1
+}
+
+// -----
+
+// CHECK-LABEL: @use_alloca
+func.func @use_alloca() -> tensor<20x20xf32> {
+  // CHECK: %[[ALLOCA:.*]] = memref.alloca() : memref<10x10xf32>
+  // CHECK: linalg.fill ins({{.*}}) outs(%[[ALLOCA]] : memref<10x10xf32>)
+  // CHECK: %[[RET:.*]] = memref.alloc() {alignment = 64 : i64} : memref<20x20xf32>
+  // CHECK-NOT: dealloc
+  // CHECK: return %[[RET]]
+  %cst0 = arith.constant 0.000000e+00 : f32
+  %cst1 = arith.constant 1.000000e+00 : f32
+  %c0 = arith.constant 0 : index
+  %c10 = arith.constant 10 : index
+  %0 = tensor.empty() : tensor<10x10xf32>
+  %1 = linalg.fill ins(%cst1 : f32) outs(%0 : tensor<10x10xf32>) -> tensor<10x10xf32>
+  %2 = tensor.pad %1 low[%c0, %c0] high[%c10, %c10] {
+  ^bb0(%arg12: index, %arg13: index):
+    tensor.yield %cst0 : f32
+  } : tensor<10x10xf32> to tensor<20x20xf32>
+  return %2 : tensor<20x20xf32>
+}
+
+transform.structured.canonicalized_sequence failures(propagate) {
+^bb1(%arg1: !pdl.operation):
+  transform.disc.bufferize %arg1
+}
+
+// -----
+
+// CHECK-LABEL: @not_use_alloca_due_to_too_large
+func.func @not_use_alloca_due_to_too_large() -> tensor<2000x2000xf32> {
+  // CHECK: %[[ALLOC:.*]] = memref.alloc() {alignment = 64 : i64} : memref<1000x1000xf32>
+  // CHECK: %[[RET:.*]] = memref.alloc() {alignment = 64 : i64} : memref<2000x2000xf32>
+  // CHECK: dealloc %[[ALLOC]]
+  // CHECK: return %[[RET]]
+  %cst0 = arith.constant 0.000000e+00 : f32
+  %cst1 = arith.constant 1.000000e+00 : f32
+  %c0 = arith.constant 0 : index
+  %c1000 = arith.constant 1000 : index
+  %0 = tensor.empty() : tensor<1000x1000xf32>
+  %1 = linalg.fill ins(%cst1 : f32) outs(%0 : tensor<1000x1000xf32>) -> tensor<1000x1000xf32>
+  %2 = tensor.pad %1 low[%c0, %c0] high[%c1000, %c1000] {
+  ^bb0(%arg12: index, %arg13: index):
+    tensor.yield %cst0 : f32
+  } : tensor<1000x1000xf32> to tensor<2000x2000xf32>
+  return %2 : tensor<2000x2000xf32>
+}
+
+transform.structured.canonicalized_sequence failures(propagate) {
+^bb1(%arg1: !pdl.operation):
+  transform.disc.bufferize %arg1
+}
+
+// -----
+
+// CHECK-LABEL: @not_use_alloca_due_to_dynamic_shape
+func.func @not_use_alloca_due_to_dynamic_shape(%arg0: index) -> tensor<?x?xf32> {
+  // CHECK: %[[ALLOC:.*]] = memref.alloc({{.*}}) {alignment = 64 : i64} : memref<?x?xf32>
+  // CHECK: %[[RET:.*]] = memref.alloc({{.*}}) {alignment = 64 : i64} : memref<?x?xf32>
+  // CHECK: dealloc %[[ALLOC]]
+  // CHECK: return %[[RET]]
+  %cst0 = arith.constant 0.000000e+00 : f32
+  %cst1 = arith.constant 1.000000e+00 : f32
+  %c0 = arith.constant 0 : index
+  %c1000 = arith.constant 1000 : index
+  %0 = tensor.empty(%arg0, %arg0) : tensor<?x?xf32>
+  %1 = linalg.fill ins(%cst1 : f32) outs(%0 : tensor<?x?xf32>) -> tensor<?x?xf32>
+  %2 = tensor.pad %1 low[%c0, %c0] high[%c1000, %c1000] {
+  ^bb0(%arg12: index, %arg13: index):
+    tensor.yield %cst0 : f32
+  } : tensor<?x?xf32> to tensor<?x?xf32>
+  return %2 : tensor<?x?xf32>
+}
+
 transform.structured.canonicalized_sequence failures(propagate) {
 ^bb1(%arg1: !pdl.operation):
   transform.disc.bufferize %arg1