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/LinalgExt/LinalgExtBase.td b/tao_compiler/mlir/disc/tools/disc-transform/LinalgExt/LinalgExtBase.td
index abc021e8c94..dd30b4730c8 100644
--- a/tao_compiler/mlir/disc/tools/disc-transform/LinalgExt/LinalgExtBase.td
+++ b/tao_compiler/mlir/disc/tools/disc-transform/LinalgExt/LinalgExtBase.td
@@ -26,6 +26,7 @@ def DISCLinalgExt_Dialect : Dialect {
     non-structured operations, ie, can not be represented in Linalg operations.
   }];
   let hasCanonicalizer = 1;
+  let hasConstantMaterializer = 1;
   let useDefaultAttributePrinterParser = 0;
 }
 
diff --git a/tao_compiler/mlir/disc/tools/disc-transform/LinalgExt/LinalgExtDialect.cc b/tao_compiler/mlir/disc/tools/disc-transform/LinalgExt/LinalgExtDialect.cc
index bdf92a7bd3f..1b985195eca 100644
--- a/tao_compiler/mlir/disc/tools/disc-transform/LinalgExt/LinalgExtDialect.cc
+++ b/tao_compiler/mlir/disc/tools/disc-transform/LinalgExt/LinalgExtDialect.cc
@@ -14,6 +14,7 @@
 #include "llvm/ADT/SmallVector.h"
 #include "llvm/ADT/TypeSwitch.h"
 #include "llvm/Support/SourceMgr.h"
+#include "mlir/Dialect/Arith/IR/Arith.h"
 #include "mlir/IR/Attributes.h"
 #include "mlir/IR/DialectImplementation.h"
 #include "mlir/IR/OpDefinition.h"
@@ -32,6 +33,13 @@ void DISCLinalgExtDialect::initialize() {
       >();
 }
 
+/// Materialize an integer or floating point constant.
+Operation* DISCLinalgExtDialect::materializeConstant(OpBuilder& builder,
+                                                     Attribute value, Type type,
+                                                     Location loc) {
+  return builder.create<arith::ConstantOp>(loc, value, type);
+}
+
 }  // namespace disc_linalg_ext
 }  // namespace disc_ral
 }  // namespace mlir
\ No newline at end of file
diff --git a/tao_compiler/mlir/disc/tools/disc-transform/LinalgExt/LinalgExtOps.cc b/tao_compiler/mlir/disc/tools/disc-transform/LinalgExt/LinalgExtOps.cc
index 4be60fc98f5..42615ec86af 100644
--- a/tao_compiler/mlir/disc/tools/disc-transform/LinalgExt/LinalgExtOps.cc
+++ b/tao_compiler/mlir/disc/tools/disc-transform/LinalgExt/LinalgExtOps.cc
@@ -280,6 +280,57 @@ LogicalResult MultiLevelPackOp::reifyResultShapes(
   return success();
 }
 
+static SmallVector<uint64_t> delinearize(uint64_t idx,
+                                         ArrayRef<int64_t> shape) {
+  SmallVector<uint64_t> indices(shape.size());
+  for (int d = static_cast<int>(shape.size()) - 1; d >= 0; --d) {
+    indices[d] = idx % shape[d];
+    idx /= shape[d];
+  }
+  return indices;
+}
+
+LogicalResult MultiLevelPackOp::fold(ArrayRef<Attribute> operands,
+                                     SmallVectorImpl<OpFoldResult>& results) {
+  if (operands.size() < 1 || !operands[0]) return failure();
+  auto srcElemAtts = operands[0].dyn_cast<ElementsAttr>();
+  if (!srcElemAtts) return failure();
+
+  auto srcTy = srcElemAtts.getType();
+  int srcRank = srcTy.getRank();
+  auto dstTy = this->getOutputType();
+  if (!dstTy.hasStaticShape()) return failure();
+  int dstRank = dstTy.getRank();
+
+  auto tileLevelsVec = this->getTileLevelsVec();
+  auto tileSizesVec = this->getTileSizesVec();
+  auto permutationVec = this->getPermutationVec();
+  auto logicalDim2SrcDim =
+      this->getOutputLogicalDimToInputDimMapping(tileLevelsVec, tileSizesVec);
+  auto logicalDim2TileSize =
+      this->getOutputLogicalDimToTileSizeMapping(tileLevelsVec, tileSizesVec);
+
+  SmallVector<Attribute> dstAttrs;
+  for (uint64_t idx = 0; idx < dstTy.getNumElements(); ++idx) {
+    auto indices = delinearize(idx, dstTy.getShape());
+    SmallVector<uint64_t> srcIndices(srcRank, 0);
+    for (int dstIdx = 0; dstIdx < dstRank; ++dstIdx) {
+      int logicalIdx = permutationVec[dstIdx];
+      int srcIdx = logicalDim2SrcDim[logicalIdx];
+      srcIndices[srcIdx] += indices[dstIdx] * logicalDim2TileSize[logicalIdx];
+    }
+    if (srcElemAtts.isValidIndex(srcIndices)) {
+      dstAttrs.push_back(srcElemAtts.getValues<Attribute>()[srcIndices]);
+    } else {
+      if (operands.size() < 3 || !operands[2]) return failure();
+      dstAttrs.push_back(operands[2]);
+    }
+  }
+
+  results.push_back(DenseElementsAttr::get(dstTy, dstAttrs));
+  return success();
+}
+
 namespace {
 /// This is derived from mlir/lib/Dialect/Linalg/IR/LinalgOps.cpp without any
 /// changes.
diff --git a/tao_compiler/mlir/disc/tools/disc-transform/LinalgExt/LinalgExtOps.td b/tao_compiler/mlir/disc/tools/disc-transform/LinalgExt/LinalgExtOps.td
index d0a74972dd9..71cdcd444e0 100644
--- a/tao_compiler/mlir/disc/tools/disc-transform/LinalgExt/LinalgExtOps.td
+++ b/tao_compiler/mlir/disc/tools/disc-transform/LinalgExt/LinalgExtOps.td
@@ -121,6 +121,8 @@ def DISCLinalgExt_MultiLevelPackOp : DISCLinalgExt_Op<"multi_level_pack", [
       )>
   ];
 
+  let hasFolder = 1;
+
   let extraClassDeclaration = extraLinalgExtOpClassDeclaration # [{
 
     // Return the output operand.
@@ -202,6 +204,48 @@ def DISCLinalgExt_MultiLevelPackOp : DISCLinalgExt_Op<"multi_level_pack", [
       return getExpectedResultRank(vs);
     }
 
+    // Returns the map: logical dim of output -> input dim
+    // Here the logical dim is the dim before transpose.
+    DenseMap<int, int> getOutputLogicalDimToInputDimMapping(
+        ArrayRef<int64_t> tileLevelsVec, ArrayRef<int64_t> tileSizesVec) {
+      int logicalIdx = 0;
+      DenseMap<int, int> dst2src;
+      for (int i = 0; i < tileLevelsVec.size(); ++i) {
+        for (int j = 0; j <= tileLevelsVec[i]; ++j) {
+          dst2src[logicalIdx++] = i;
+        }
+      }
+      return dst2src;
+    }
+
+    // Returns the map: logical dim of output -> tile size for each dim
+    // Here the logical dim is the dim before transpose.
+    DenseMap<int, int> getOutputLogicalDimToTileSizeMapping(
+        ArrayRef<int64_t> tileLevelsVec, ArrayRef<int64_t> tileSizesVec) {
+      int logicalIdx = 0;
+      int tileSizeIdx = 0;
+      DenseMap<int, int> logicalDim2TileSize;
+      for (int i = 0; i < tileLevelsVec.size(); ++i) {
+        for (int j = 0; j < tileLevelsVec[i]; ++j) {
+          logicalDim2TileSize[logicalIdx++] = tileSizesVec[tileSizeIdx++];
+        }
+        logicalDim2TileSize[logicalIdx++] = 1;
+      }
+      return logicalDim2TileSize;
+    }
+
+    // Returns the inner most dims
+    DenseSet<int> getOutputInnerMostDims(
+        ArrayRef<int64_t> tileLevelsVec, ArrayRef<int64_t> tileSizesVec) {
+      int logicalIdx = 0;
+      DenseSet<int> innerMostDims;
+      for (int i = 0; i < tileLevelsVec.size(); ++i) {
+        logicalIdx += tileLevelsVec[i];
+        innerMostDims.insert(logicalIdx++);
+      }
+      return innerMostDims;
+    }
+
     // Method to get the shape of the result as `SmallVector<OpFoldResult>`.
     // This is a static method to allow getting the shape of the destination
     // expected while creating a `pack` op.


Test patch:
diff --git a/tao_compiler/mlir/disc/tools/disc-transform/LinalgExt/tests/canonicalize.mlir b/tao_compiler/mlir/disc/tools/disc-transform/LinalgExt/tests/canonicalize.mlir
new file mode 100644
index 00000000000..e47974d4a33
--- /dev/null
+++ b/tao_compiler/mlir/disc/tools/disc-transform/LinalgExt/tests/canonicalize.mlir
@@ -0,0 +1,40 @@
+// RUN: disc-opt --canonicalize -split-input-file %s | FileCheck %s
+
+// CHECK-LABEL: @multi_level_pack_fold
+func.func @multi_level_pack_fold() -> tensor<1x2x2x2xf32> {
+  // CHECK: %[[CST:.*]] =  arith.constant
+  // CHECK-SAME{LITERAL}: dense<[[[[0.000000e+00, 1.000000e+00], [4.000000e+00, 5.000000e+00]], [[2.000000e+00, 3.000000e+00], [6.000000e+00, 7.000000e+00]]]]> : tensor<1x2x2x2xf32>
+  // CHECK: return %[[CST]]
+  %cst = arith.constant dense<[[0.0, 1.0, 2.0, 3.0], [4.0, 5.0, 6.0, 7.0]]> : tensor<2x4xf32>
+  %0 = tensor.empty() : tensor<1x2x2x2xf32>
+  %1 = disc_linalg_ext.multi_level_pack %cst with tile_levels = [1, 1] tile_sizes = [2, 2] permutation = [0, 2, 1, 3] into %0 : (tensor<2x4xf32> tensor<1x2x2x2xf32>) -> tensor<1x2x2x2xf32>
+  return %1 : tensor<1x2x2x2xf32>
+}
+
+// -----
+
+// CHECK-LABEL: @multi_level_pack_fold_2
+func.func @multi_level_pack_fold_2() -> tensor<1x2x2x2xf32> {
+  // CHECK: %[[CST:.*]] =  arith.constant
+  // CHECK-SAME{LITERAL}: dense<[[[[0.000000e+00, 4.000000e+00], [1.000000e+00, 5.000000e+00]], [[2.000000e+00, 6.000000e+00], [3.000000e+00, 7.000000e+00]]]]> : tensor<1x2x2x2xf32>
+  // CHECK: return %[[CST]]
+  %cst = arith.constant dense<[[0.0, 1.0, 2.0, 3.0], [4.0, 5.0, 6.0, 7.0]]> : tensor<2x4xf32>
+  %0 = tensor.empty() : tensor<1x2x2x2xf32>
+  %1 = disc_linalg_ext.multi_level_pack %cst with tile_levels = [1, 1] tile_sizes = [2, 2] permutation = [0, 2, 3, 1] into %0 : (tensor<2x4xf32> tensor<1x2x2x2xf32>) -> tensor<1x2x2x2xf32>
+  return %1 : tensor<1x2x2x2xf32>
+}
+
+// -----
+
+// CHECK-LABEL: @multi_level_pack_fold_with_pad
+func.func @multi_level_pack_fold_with_pad() -> tensor<1x3x2x2xf32> {
+  // CHECK: %[[CST:.*]] =  arith.constant
+  // CHECK-SAME{LITERAL}: dense<[[[[0.000000e+00, 1.000000e+00], [5.000000e+00, 6.000000e+00]], [[2.000000e+00, 3.000000e+00], [7.000000e+00, 8.000000e+00]], [[4.000000e+00, 1.280000e+02], [9.000000e+00, 1.280000e+02]]]]>
+  // CHECK: return %[[CST]]
+  %cst = arith.constant dense<[[0.0, 1.0, 2.0, 3.0, 4.0], [5.0, 6.0, 7.0, 8.0, 9.0]]> : tensor<2x5xf32>
+  %init = arith.constant 128.0 : f32
+  %0 = tensor.empty() : tensor<1x3x2x2xf32>
+  %1 = disc_linalg_ext.multi_level_pack %cst with padding_value(%init: f32) tile_levels = [1, 1] tile_sizes = [2, 2] permutation = [0, 2, 1, 3] into %0 : (tensor<2x5xf32> tensor<1x3x2x2xf32>) -> tensor<1x3x2x2xf32>
+  return %1 : tensor<1x3x2x2xf32>
+}
+