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 faee8351a0e..37550529f0a 100644
--- a/tao_compiler/mlir/disc/tools/disc-transform/TransformOps/TransformOpsExt.cc
+++ b/tao_compiler/mlir/disc/tools/disc-transform/TransformOps/TransformOpsExt.cc
@@ -1002,11 +1002,12 @@ DiagnosedSilenceableFailure LowerMultiLevelPackToLoopOp::applyToOne(
   Value loopInitValue = dst;
   SmallVector<scf::ForOp> forOps;
   SmallVector<Value> srcDimUppers = srcSizes;
+  auto staticSrcDimUppers = llvm::to_vector<>(srcTy.getShape());
   for (int dstIdx = 0; dstIdx < dstRank; ++dstIdx) {
     int logicalIdx = permutationVec[dstIdx];
     if (innerMostDimsInfo.isInnerMostDim(logicalIdx)) continue;
-    Value step =
-        b.create<arith::ConstantIndexOp>(loc, logicalDim2TileSize[logicalIdx]);
+    int64_t staticStep = logicalDim2TileSize[logicalIdx];
+    Value step = b.create<arith::ConstantIndexOp>(loc, staticStep);
     int srcIdx = logicalDim2SrcDim[logicalIdx];
     Value upper = srcDimUppers[srcIdx];
     auto forOp =
@@ -1014,8 +1015,16 @@ DiagnosedSilenceableFailure LowerMultiLevelPackToLoopOp::applyToOne(
     b.setInsertionPoint(forOp.getBody(), forOp.getBody()->begin());
     Value iv = forOp.getInductionVar();
     srcOffsets[srcIdx] = b.create<arith::AddIOp>(loc, srcOffsets[srcIdx], iv);
-    Value remaining = b.create<arith::SubIOp>(loc, upper, iv);
-    srcDimUppers[srcIdx] = buildMin(b, loc, remaining, step);
+    if (staticStep == 1 ||
+        staticSrcDimUppers[srcIdx] != ShapedType::kDynamicSize &&
+            staticSrcDimUppers[srcIdx] % staticStep == 0) {
+      srcDimUppers[srcIdx] = step;
+      staticSrcDimUppers[srcIdx] = staticStep;
+    } else {
+      Value remaining = b.create<arith::SubIOp>(loc, upper, iv);
+      srcDimUppers[srcIdx] = buildMin(b, loc, remaining, step);
+      staticSrcDimUppers[srcIdx] = ShapedType::kDynamicSize;
+    }
     dstOffsets[dstIdx] = b.create<arith::DivSIOp>(loc, iv, step);
     dstSizes[dstIdx] = one;
     forOps.push_back(forOp);
@@ -1044,8 +1053,12 @@ DiagnosedSilenceableFailure LowerMultiLevelPackToLoopOp::applyToOne(
   auto sortedPermutation = innerMostDimsInfo.permutation;
   llvm::sort(sortedPermutation);
   if (sortedPermutation != innerMostDimsInfo.permutation) {
-    Value transposeDst = b.create<tensor::EmptyOp>(
-        loc, innerMostDimsInfo.transposedDimSizes, dstTy.getElementType());
+    auto transposeDstTy = RankedTensorType::get(
+        innerMostDimsInfo.transposedDimSizes, dstTy.getElementType());
+    Value transposeDst =
+        b.create<tensor::ExtractSliceOp>(loc, transposeDstTy, loopInitValue,
+                                         dstOffsets, dstSizes, dstStrides)
+            ->getResult(0);
     srcSlice = makeTransposeOp(b, loc, srcSlice, transposeDst,
                                innerMostDimsInfo.permutation)
                    ->getResult(0);


Test patch:
diff --git a/tao_compiler/mlir/disc/tests/disc-transform/default_schedule_matmul.cc b/tao_compiler/mlir/disc/tests/disc-transform/default_schedule_matmul.cc
index 2a889a77b2b..0d34bf98501 100644
--- a/tao_compiler/mlir/disc/tests/disc-transform/default_schedule_matmul.cc
+++ b/tao_compiler/mlir/disc/tests/disc-transform/default_schedule_matmul.cc
@@ -24,6 +24,7 @@ const std::string c_ft_path =
 
 static bool init_threads = []() {
   setenv("OMP_NUM_THREADS", "1", 1);
+  setenv("DISC_CPU_ENABLE_WEIGHT_PRE_PACKING", "1", 1);
   return true;
 }();
 
diff --git a/tao_compiler/mlir/disc/tests/disc-transform/matmul.cc b/tao_compiler/mlir/disc/tests/disc-transform/matmul.cc
index 1473681733b..96d76716cc0 100644
--- a/tao_compiler/mlir/disc/tests/disc-transform/matmul.cc
+++ b/tao_compiler/mlir/disc/tests/disc-transform/matmul.cc
@@ -24,6 +24,7 @@ const std::string c_ft_path =
 
 static bool init_threads = []() {
   setenv("OMP_NUM_THREADS", "1", 1);
+  setenv("DISC_CPU_ENABLE_WEIGHT_PRE_PACKING", "1", 1);
   return true;
 }();
 
diff --git a/tao_compiler/mlir/disc/tests/disc-transform/matmul_multithread.cc b/tao_compiler/mlir/disc/tests/disc-transform/matmul_multithread.cc
index f8ebcbba0f3..b6a4ad9d148 100644
--- a/tao_compiler/mlir/disc/tests/disc-transform/matmul_multithread.cc
+++ b/tao_compiler/mlir/disc/tests/disc-transform/matmul_multithread.cc
@@ -24,6 +24,7 @@ const std::string c_ft_path =
 
 static bool init_threads = []() {
   setenv("OMP_NUM_THREADS", "8", 1);
+  setenv("DISC_CPU_ENABLE_WEIGHT_PRE_PACKING", "1", 1);
   return true;
 }();
 
diff --git a/tao_compiler/mlir/disc/tests/disc-transform/packed_matmul.cc b/tao_compiler/mlir/disc/tests/disc-transform/packed_matmul.cc
index 5af536ac586..0ff26ad5615 100644
--- a/tao_compiler/mlir/disc/tests/disc-transform/packed_matmul.cc
+++ b/tao_compiler/mlir/disc/tests/disc-transform/packed_matmul.cc
@@ -24,6 +24,7 @@ const std::string c_ft_path =
 
 static bool init_threads = []() {
   setenv("OMP_NUM_THREADS", "1", 1);
+  setenv("DISC_CPU_ENABLE_WEIGHT_PRE_PACKING", "1", 1);
   return true;
 }();
 
diff --git a/tao_compiler/mlir/disc/tools/disc-transform/transforms/tests/lower-multi-level-pack-to-loop.mlir b/tao_compiler/mlir/disc/tools/disc-transform/transforms/tests/lower-multi-level-pack-to-loop.mlir
index fcd652b84a0..20391add4fc 100644
--- a/tao_compiler/mlir/disc/tools/disc-transform/transforms/tests/lower-multi-level-pack-to-loop.mlir
+++ b/tao_compiler/mlir/disc/tools/disc-transform/transforms/tests/lower-multi-level-pack-to-loop.mlir
@@ -56,3 +56,35 @@ transform.structured.canonicalized_sequence failures(propagate) {
   %pack_op = transform.structured.match ops{["disc_linalg_ext.multi_level_pack"]} in %arg1
   transform.disc.lower_multi_level_pack_to_loop %pack_op
 }
+
+// -----
+
+// CHECK-LABEL: @multi_level_pack_with_pad_2
+// CHECK-SAME: (%[[ARG0:.*]]: tensor<?x?xf32>, %[[ARG1:.*]]: index, %[[ARG2:.*]]: index)
+func.func @multi_level_pack_with_pad_2(%arg0: tensor<?x?xf32>, %arg1: index, %arg2: index) -> tensor<?x?x32x1xf32> {
+  %cst = arith.constant 1.0 : f32
+  // CHECK: %[[CST:.*]] = arith.constant 1.000000e+00 : f32
+  // CHECK: %[[INIT:.*]] = tensor.empty(%[[ARG1]], %[[ARG2]]) : tensor<?x?x32x1xf32>
+  // CHECK: %[[DIM0:.*]] = tensor.dim %[[ARG0]], %c0 : tensor<?x?xf32>
+  // CHECK: %[[DIM1:.*]] = tensor.dim %[[ARG0]], %c1 : tensor<?x?xf32>
+  // CHECK: %[[T0:.*]] = scf.for %[[IV0:.*]] = %c0 to %[[DIM0]] step %c32 iter_args(%[[V0:.*]] = %[[INIT]]) -> (tensor<?x?x32x1xf32>) {
+  // CHECK:   %[[T1:.*]] = scf.for %[[IV1:.*]] = %c0 to %[[DIM1]] step %c1 iter_args(%[[V1:.*]] = %[[V0]]) -> (tensor<?x?x32x1xf32>) {
+  // CHECK:     %[[SLICE:.*]] = tensor.extract_slice %[[ARG0]]
+  // CHECK-SAME: %[[IV0]], %[[IV1]]
+  // CHECK:     %[[PAD:.*]] = tensor.pad %[[SLICE]]
+  // CHECK-SAME: low[0, 0]
+  // CHECK-SAME: high[%[[DIM_HIGH:.*]], %c0]
+  // CHECK:     %[[UPDATE:.*]] = tensor.insert_slice %[[PAD]] into %[[V1]]
+  // CHECK:     scf.yield %[[UPDATE]] : tensor<?x?x32x1xf32>
+  // CHECK:   scf.yield %[[T1]] : tensor<?x?x32x1xf32>
+  // CHECK: return %[[T0]]
+  %0 = tensor.empty(%arg1, %arg2) : tensor<?x?x32x1xf32>
+  %1 = disc_linalg_ext.multi_level_pack %arg0 with padding_value(%cst : f32) tile_levels = [1, 1] tile_sizes = [32, 1] permutation = [0, 2, 1, 3] into %0 : (tensor<?x?xf32> tensor<?x?x32x1xf32>) -> tensor<?x?x32x1xf32>
+  return %1 : tensor<?x?x32x1xf32>
+}
+
+transform.structured.canonicalized_sequence failures(propagate) {
+^bb1(%arg1: !pdl.operation):
+  %pack_op = transform.structured.match ops{["disc_linalg_ext.multi_level_pack"]} in %arg1
+  transform.disc.lower_multi_level_pack_to_loop %pack_op
+}
\ No newline at end of file