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/BUILD b/tao_compiler/mlir/disc/tools/disc-transform/BUILD
index 04b7c210eba..65bc1e9cdf3 100644
--- a/tao_compiler/mlir/disc/tools/disc-transform/BUILD
+++ b/tao_compiler/mlir/disc/tools/disc-transform/BUILD
@@ -283,6 +283,7 @@ cc_library(
     srcs = ["transforms/legalize_lmhlo_fusion_to_linalg.cc"],
     deps = [
         ":DISCLinalgExtDialect",
+        ":disc_transform_utils",
         ":pass_details",
         "//tensorflow/compiler/xla/mlir_hlo:lhlo",
         "@llvm-project//llvm:Support",
diff --git a/tao_compiler/mlir/disc/tools/disc-transform/transforms/legalize_lmhlo_fusion_to_linalg.cc b/tao_compiler/mlir/disc/tools/disc-transform/transforms/legalize_lmhlo_fusion_to_linalg.cc
index f1352ddab33..e944c17f798 100644
--- a/tao_compiler/mlir/disc/tools/disc-transform/transforms/legalize_lmhlo_fusion_to_linalg.cc
+++ b/tao_compiler/mlir/disc/tools/disc-transform/transforms/legalize_lmhlo_fusion_to_linalg.cc
@@ -22,6 +22,7 @@
 #include "tensorflow/compiler/mlir/disc/tools/disc-transform/LinalgExt/LinalgExtDialect.h"
 #include "tensorflow/compiler/mlir/disc/tools/disc-transform/LinalgExt/LinalgExtOps.h"
 #include "tensorflow/compiler/mlir/disc/tools/disc-transform/transforms/PassDetail.h"
+#include "tensorflow/compiler/mlir/disc/tools/disc-transform/utils.h"
 
 #define DEBUG_TYPE "disc-legalize-lmhlo-fusion-to-linalg"
 
@@ -82,7 +83,8 @@ LogicalResult emitReturnOp(func::ReturnOp op, OpBuilder& b,
 }
 
 LogicalResult emitDotGeneralOp(lmhlo::DotGeneralOp op, OpBuilder& b,
-                               BlockAndValueMapping& mapping) {
+                               BlockAndValueMapping& mapping,
+                               const std::string& name) {
   Value A = op->getOperand(0);
   Value B = op->getOperand(1);
   Value C = op->getOperand(2);
@@ -96,33 +98,37 @@ LogicalResult emitDotGeneralOp(lmhlo::DotGeneralOp op, OpBuilder& b,
   auto zeroAttr = b.getZeroAttr(ty.getElementType());
   Location loc = op->getLoc();
   Value zero = b.create<arith::ConstantOp>(loc, zeroAttr);
-  Value t0 = b.create<linalg::FillOp>(loc, zero, newC).result();
-  Value t1 =
-      b.create<linalg::MatmulOp>(loc, ValueRange{newA, newB}, ValueRange{t0})
-          .getResult(0);
+  auto fillOp = b.create<linalg::FillOp>(loc, zero, newC);
+  fillOp->setAttr(kDISCLinalgTransformName, b.getStringAttr(name));
+  auto matmulOp = b.create<linalg::MatmulOp>(loc, ValueRange{newA, newB},
+                                             ValueRange{fillOp.getResult(0)});
+  matmulOp->setAttr(kDISCLinalgTransformName, b.getStringAttr(name));
   mapping.erase(C);
-  mapping.map(C, t1);
+  mapping.map(C, matmulOp.getResult(0));
 
   return success();
 }
 
 LogicalResult emitConstOp(lmhlo::ConstantOp op, OpBuilder& b,
-                          BlockAndValueMapping& mapping) {
+                          BlockAndValueMapping& mapping,
+                          const std::string& name) {
   auto resultTy = convertMemRefToTensorType(op->getOperand(0).getType());
   Location loc = op->getLoc();
   auto newOp = b.create<disc_linalg_ext::ConstantWrapperOp>(loc, resultTy,
                                                             op.getValue());
   mapping.erase(op->getOperand(0));
   mapping.map(op->getOperand(0), newOp.getResult());
+  newOp->setAttr(kDISCLinalgTransformName, b.getStringAttr(name));
   return success();
 }
 
 LogicalResult emitLmhloOp(Operation* op, OpBuilder& b,
-                          BlockAndValueMapping& mapping) {
+                          BlockAndValueMapping& mapping,
+                          const std::string& name) {
   if (auto dotGeneralOp = dyn_cast<lmhlo::DotGeneralOp>(op)) {
-    return emitDotGeneralOp(dotGeneralOp, b, mapping);
+    return emitDotGeneralOp(dotGeneralOp, b, mapping, name);
   } else if (auto constOp = dyn_cast<lmhlo::ConstantOp>(op)) {
-    return emitConstOp(constOp, b, mapping);
+    return emitConstOp(constOp, b, mapping, name);
   }
   // TODO(wyzero): support other lmhlo ops.
   return failure();
@@ -130,9 +136,11 @@ LogicalResult emitLmhloOp(Operation* op, OpBuilder& b,
 
 LogicalResult emitLmhloFusionOp(lmhlo::FusionOp op, OpBuilder& b,
                                 BlockAndValueMapping& mapping) {
+  TransformNameAssigner assigner;
   for (Operation& op : op.getRegion().getBlocks().front()) {
     if (isa<lmhlo::TerminatorOp>(&op)) continue;
-    if (failed(emitLmhloOp(&op, b, mapping))) return failure();
+    if (failed(emitLmhloOp(&op, b, mapping, assigner.nameNewOperation(&op))))
+      return failure();
   }
   return success();
 }
diff --git a/tao_compiler/mlir/disc/tools/disc-transform/utils.cc b/tao_compiler/mlir/disc/tools/disc-transform/utils.cc
index 4414e9011ca..2cc76edb787 100644
--- a/tao_compiler/mlir/disc/tools/disc-transform/utils.cc
+++ b/tao_compiler/mlir/disc/tools/disc-transform/utils.cc
@@ -11,6 +11,7 @@
 
 #include "tensorflow/compiler/mlir/disc/tools/disc-transform/utils.h"
 
+#include "llvm/ADT/SmallString.h"
 #include "llvm/Support/Debug.h"
 #include "llvm/Support/SourceMgr.h"
 #include "mlir/Dialect/Linalg/IR/Linalg.h"
@@ -22,6 +23,28 @@
 namespace mlir {
 namespace disc_ral {
 
+const char* kDISCLinalgTransformName = "disc.transform.name";
+
+TransformNameAssigner::TransformNameAssigner(ArrayRef<Operation*> ops) {
+  for (auto op : ops) nameNewOperation(op);
+}
+
+std::string TransformNameAssigner::nameNewOperation(Operation* op) {
+  SmallString<128> fullName;
+  auto opName = op->getName().stripDialect().str();
+  fullName.append(opName);
+  if (auto cnt = nameCounterMap_[opName]++) {
+    fullName.append(("_" + Twine(cnt)).str());
+  }
+
+  return nameMap_[op] = fullName.str();
+}
+
+const std::unordered_map<Operation*, std::string>&
+TransformNameAssigner::getNameMap() {
+  return nameMap_;
+}
+
 /// Create a linalg::GenericOp version of an n-D copy that can further tile,
 /// lower to loops or vectorize, unlike the current implementation of
 /// memref::CopyOp.
diff --git a/tao_compiler/mlir/disc/tools/disc-transform/utils.h b/tao_compiler/mlir/disc/tools/disc-transform/utils.h
index 51eabf183b9..6047f57aac1 100644
--- a/tao_compiler/mlir/disc/tools/disc-transform/utils.h
+++ b/tao_compiler/mlir/disc/tools/disc-transform/utils.h
@@ -16,11 +16,38 @@ limitations under the License.
 #ifndef DISC_TOOLS_DISC_TRANSFORM_UTILS_H_
 #define DISC_TOOLS_DISC_TRANSFORM_UTILS_H_
 
+#include <string>
+#include <unordered_map>
+
 #include "mlir/Dialect/MemRef/IR/MemRef.h"
 
 namespace mlir {
 namespace disc_ral {
 
+// An attribute name used to represent the unique name assigned to each linalg
+// ops. This tag can be used to distinguish different linalg ops when using
+// transform IR.
+extern const char* kDISCLinalgTransformName;
+
+class TransformNameAssigner {
+ public:
+  TransformNameAssigner() = default;
+  TransformNameAssigner(ArrayRef<Operation*> ops);
+
+  // Returns the name assigned to the new `op`. The `op` must not be assigned
+  // name by this assigner before.
+  std::string nameNewOperation(Operation* op);
+
+  // Returns the <op,name> map. The map records the names for all the ops that
+  // have been passed to `nameNewOperation`.
+  const std::unordered_map<Operation*, std::string>& getNameMap();
+
+ private:
+  // Assigned a unique id for the ops having the same op name.
+  std::unordered_map<std::string, int> nameCounterMap_;
+  std::unordered_map<Operation*, std::string> nameMap_;
+};
+
 /// Create a linalg::GenericOp version of an n-D copy that can further tile,
 /// lower to loops or vectorize, unlike the current implementation of
 /// memref::CopyOp.
diff --git a/tao_compiler/mlir/disc/transforms/disc_transform_schedule.cc b/tao_compiler/mlir/disc/transforms/disc_transform_schedule.cc
index ee776bac657..ea11a8fa9f6 100644
--- a/tao_compiler/mlir/disc/transforms/disc_transform_schedule.cc
+++ b/tao_compiler/mlir/disc/transforms/disc_transform_schedule.cc
@@ -87,16 +87,26 @@ bool PatternKindAndStringMapRegistrar = []() {
 using transform::FuseIntoContainingOp;
 using transform::FuseOp;
 using transform::MatchOp;
+using transform::SplitHandlesOp;
 using transform::TileOp;
 using transform::TileToForeachThreadOp;
 using transform::VectorizeOp;
 
 MatchOp buildMatchOp(OpBuilder& b, Location& loc, Value target,
-                     ArrayRef<StringRef> ops) {
+                     ArrayRef<StringRef> ops, StringRef name = {}) {
+  ArrayAttr opNames;
+  if (!ops.empty()) {
+    opNames = b.getStrArrayAttr(ops);
+  }
+  DictionaryAttr attrs;
+  if (!name.empty()) {
+    attrs = b.getDictionaryAttr(
+        b.getNamedAttr(kDISCLinalgTransformName, b.getStringAttr(name)));
+  }
+
   return b.create<MatchOp>(loc, pdl::OperationType::get(b.getContext()), target,
-                           b.getStrArrayAttr(ops),
-                           transform::MatchInterfaceEnumAttr{},
-                           DictionaryAttr{}, TypeAttr{});
+                           opNames, transform::MatchInterfaceEnumAttr{}, attrs,
+                           TypeAttr{});
 }
 
 TileToForeachThreadOp buildTileToForEachThreadOp(OpBuilder& b, Location& loc,
@@ -219,6 +229,13 @@ void buildLowerVectors(OpBuilder& b, Location& loc, ArrayRef<int64_t> stages,
       transposeLowering, transposeAvx2Lowering);
 }
 
+SplitHandlesOp buildSplitHandlesOp(OpBuilder& b, Location& loc, Value target,
+                                   uint64_t num_result_handles) {
+  SmallVector<Type> pdlTypes(num_result_handles,
+                             pdl::OperationType::get(b.getContext()));
+  return b.create<SplitHandlesOp>(loc, pdlTypes, target, num_result_handles);
+}
+
 LogicalResult aarch64GEMMDefaultScheduleFactory(PatternDescription& pd,
                                                 ModuleOp m) {
   OpBuilder b(m);
@@ -235,6 +252,7 @@ LogicalResult aarch64GEMMDefaultScheduleFactory(PatternDescription& pd,
   Value variant = bodyBlock.getArgument(0);
 
   auto& fusionPattern = pd.getFusionPattern();
+  auto nameMap = TransformNameAssigner(fusionPattern.getOpList()).getNameMap();
   auto dotOp =
       dyn_cast_or_null<lmhlo::DotGeneralOp>(fusionPattern.getDominantOp());
   if (!dotOp) {
@@ -266,10 +284,11 @@ LogicalResult aarch64GEMMDefaultScheduleFactory(PatternDescription& pd,
   int64_t N = rhsTranspose ? rhsTy.getShape()[rhsTy.getRank() - 2]
                            : rhsTy.getShape()[rhsTy.getRank() - 1];
 
-  // %fill = transform.structured.match ops{["linalg.fill"]} in %variant
-  Value fill = buildMatchOp(b, loc, variant, {"linalg.fill"});
-  // %matmul = transform.structured.match ops{["linalg.matmul"]} in %arg1
-  Value matmul = buildMatchOp(b, loc, variant, {"linalg.matmul"});
+  // build handle to target dot op.
+  Value fillAndMatmul = buildMatchOp(b, loc, variant, {}, nameMap[dotOp]);
+  auto matmulSplitOp = buildSplitHandlesOp(b, loc, fillAndMatmul, 2);
+  Value fill = matmulSplitOp->getResult(0);
+  Value matmul = matmulSplitOp->getResult(1);
 
   // transform.structured.tile_to_foreach_thread_op %matmul num_threads [1, 1]
   auto forEachThreadOp = buildTileToForEachThreadOp(b, loc, matmul, {1, 1});


Test patch:
diff --git a/tao_compiler/mlir/disc/tools/disc-transform/transforms/tests/legalize-lmhlo-fusion-to-linalg.mlir b/tao_compiler/mlir/disc/tools/disc-transform/transforms/tests/legalize-lmhlo-fusion-to-linalg.mlir
index aa5972bd874..4b53cb65896 100644
--- a/tao_compiler/mlir/disc/tools/disc-transform/transforms/tests/legalize-lmhlo-fusion-to-linalg.mlir
+++ b/tao_compiler/mlir/disc/tools/disc-transform/transforms/tests/legalize-lmhlo-fusion-to-linalg.mlir
@@ -4,8 +4,8 @@
 // CHECK-SAME: (%[[ARG0:.*]]: tensor<?x?xf32>, %[[ARG1:.*]]: tensor<?x?xf32>, %[[ARG2:.*]]: tensor<?x?xf32>)
 func.func @matmul_nn(%arg1: memref<?x?xf32, "cpu">, %arg2: memref<?x?xf32, "cpu">, %arg3: memref<?x?xf32, "cpu">) -> memref<?x?xf32, "cpu"> {
   // CHECK: %[[T0:.*]] = arith.constant 0.000000e+00 : f32
-  // CHECK: %[[T1:.*]] = linalg.fill ins(%[[T0]] : f32) outs(%[[ARG2]] : tensor<?x?xf32>) -> tensor<?x?xf32>
-  // CHECK: %[[T2:.*]] = linalg.matmul ins(%[[ARG0]], %[[ARG1]] : tensor<?x?xf32>, tensor<?x?xf32>) outs(%[[T1]] : tensor<?x?xf32>) -> tensor<?x?xf32>
+  // CHECK: %[[T1:.*]] = linalg.fill {disc.transform.name = "dot_general"} ins(%[[T0]] : f32) outs(%[[ARG2]] : tensor<?x?xf32>) -> tensor<?x?xf32>
+  // CHECK: %[[T2:.*]] = linalg.matmul {disc.transform.name = "dot_general"} ins(%[[ARG0]], %[[ARG1]] : tensor<?x?xf32>, tensor<?x?xf32>) outs(%[[T1]] : tensor<?x?xf32>) -> tensor<?x?xf32>
   // CHECK: return %[[T2]]
   "lmhlo.fusion"() ({
     "lmhlo.dot_general"(%arg1, %arg2, %arg3) {dot_dimension_numbers = #mhlo.dot<lhs_contracting_dimensions = [1], rhs_contracting_dimensions = [0]>} : (memref<?x?xf32, "cpu">, memref<?x?xf32, "cpu">, memref<?x?xf32, "cpu">) -> ()
@@ -19,10 +19,10 @@ func.func @matmul_nn(%arg1: memref<?x?xf32, "cpu">, %arg2: memref<?x?xf32, "cpu"
 // CHECK-LABEL: @packed_matmul_nn
 // CHECK-SAME: (%[[ARG0:.*]]: tensor<?x512xf32>, %[[ARG1:.*]]: tensor<512x1024xf32>, %[[ARG2:.*]]: tensor<?x1024xf32>)
 func.func @packed_matmul_nn(%arg1: memref<?x512xf32, "cpu">, %arg2: memref<512x1024xf32, "cpu">, %arg3: memref<?x1024xf32, "cpu">) -> memref<?x1024xf32, "cpu"> {
-  // CHECK: %[[T0:.*]] = disc_linalg_ext.constant_wrapper dense<-8.000000e-01> : tensor<512x1024xf32>
+  // CHECK: %[[T0:.*]] = disc_linalg_ext.constant_wrapper {disc.transform.name = "constant"} dense<-8.000000e-01> : tensor<512x1024xf32>
   // CHECK: %[[T1:.*]] = arith.constant 0.000000e+00 : f32
-  // CHECK: %[[T2:.*]] = linalg.fill ins(%[[T1]] : f32) outs(%[[ARG2]] : tensor<?x1024xf32>) -> tensor<?x1024xf32>
-  // CHECK: %[[T3:.*]] = linalg.matmul ins(%[[ARG0]], %[[T0]] : tensor<?x512xf32>, tensor<512x1024xf32>) outs(%[[T2]] : tensor<?x1024xf32>) -> tensor<?x1024xf32>
+  // CHECK: %[[T2:.*]] = linalg.fill {disc.transform.name = "dot_general"} ins(%[[T1]] : f32) outs(%[[ARG2]] : tensor<?x1024xf32>) -> tensor<?x1024xf32>
+  // CHECK: %[[T3:.*]] = linalg.matmul {disc.transform.name = "dot_general"} ins(%[[ARG0]], %[[T0]] : tensor<?x512xf32>, tensor<512x1024xf32>) outs(%[[T2]] : tensor<?x1024xf32>) -> tensor<?x1024xf32>
   // CHECK: return %[[T3]]
   "lmhlo.fusion"() ({
     "lmhlo.constant"(%arg2) {disc.device = "cpu", value = dense<-8.000000e-01> : tensor<512x1024xf32>} : (memref<512x1024xf32, "cpu">) -> ()