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 ea35e5f78ca..7c115765ae2 100644
--- a/tao_compiler/mlir/disc/tools/disc-transform/BUILD
+++ b/tao_compiler/mlir/disc/tools/disc-transform/BUILD
@@ -286,6 +286,7 @@ cc_library(
         ":DISCLinalgExtDialect",
         ":disc_transform_utils",
         ":pass_details",
+        "//mlir/disc:disc_lhlo_elemental_utils",
         "@org_tensorflow//tensorflow/compiler/xla/mlir_hlo:legalize_to_linalg_utils",
         "@org_tensorflow//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 75b988d14f6..b8be269d9e2 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
@@ -25,6 +25,7 @@
 #include "mlir/disc/tools/disc-transform/LinalgExt/LinalgExtOps.h"
 #include "mlir/disc/tools/disc-transform/transforms/PassDetail.h"
 #include "mlir/disc/tools/disc-transform/utils.h"
+#include "mlir/disc/transforms/lhlo_elemental_utils.h"
 
 #define DEBUG_TYPE "disc-legalize-lmhlo-fusion-to-linalg"
 
@@ -203,6 +204,128 @@ LogicalResult emitConstOp(lmhlo::ConstantOp op, OpBuilder& b,
   return success();
 }
 
+LogicalResult emitBcastOp(Operation* op, OpBuilder& b,
+                          BlockAndValueMapping& mapping,
+                          const std::string& name) {
+  SmallVector<NamedAttribute> prunedAttrs;
+  if (auto dynBcastInDim = dyn_cast<lmhlo::DynamicBroadcastInDimOp>(op)) {
+    prunedAttrs = linalg::getPrunedAttributeList(dynBcastInDim);
+  } else if (auto bcastInDim = dyn_cast<lmhlo::BroadcastInDimOp>(op)) {
+    prunedAttrs = linalg::getPrunedAttributeList(bcastInDim);
+  } else if (auto bcast = dyn_cast<lmhlo::BroadcastOp>(op)) {
+    prunedAttrs = linalg::getPrunedAttributeList(bcast);
+  } else {
+    return op->emitError() << "not support bcast op\n";
+  }
+  Value in = op->getOperand(0);
+  auto inType = in.getType().cast<MemRefType>();
+  Value out = cast<lmhlo::LmhloOp>(op).getResultBuffer();
+  auto outType = out.getType().cast<MemRefType>();
+  auto dimAttr = op->getAttrOfType<DenseElementsAttr>("broadcast_dimensions");
+  if (!dimAttr)
+    return op->emitError() << "no broadcast_dimensions attr found\n";
+  auto dimensions = dimAttr.getValues<int64_t>();
+
+  // Determine dimension expressions based on whether the dimension is
+  // expanding (0) or non-expanding (identity).
+  // Our fusion pass will make sure there are only supported bcast ops in the
+  // fusion Patterns, thus we can always decide if it's one of the above two
+  // cases in the compile-time.
+  SmallVector<AffineExpr> dimExprs(inType.getRank(), nullptr);
+  if (inType.getRank() != dimensions.size())
+    return op->emitError()
+           << "size of broadcast_dimensions not equal rank of input\n";
+  for (auto [inDimIdx, inDimSize] : llvm::enumerate(inType.getShape())) {
+    int64_t outDimIdx = dimensions[inDimIdx];
+    int64_t outDimSize = outType.getShape()[outDimIdx];
+    bool isExpanding = (inDimSize == 1 && outDimSize != 1);
+    dimExprs[inDimIdx] = isExpanding ? b.getAffineConstantExpr(0)
+                                     : b.getAffineDimExpr(outDimIdx);
+  }
+
+  Location loc = op->getLoc();
+  int64_t nloops = outType.getRank();
+  Value newIn = mapping.lookup(in);
+  Value newOut = mapping.lookup(out);
+  auto newOp = b.create<linalg::GenericOp>(
+      loc, TypeRange{newOut.getType()}, ValueRange{newIn}, ValueRange{newOut},
+      llvm::ArrayRef({AffineMap::get(/*dimCount=*/nloops, /*symbolCount=*/0,
+                                     dimExprs, b.getContext()),
+                      b.getMultiDimIdentityMap(nloops)}),
+      mhlo::getParallelAndReductionIterators(
+          /*nLoops=*/nloops,
+          /*nReduction=*/0),
+      [&](OpBuilder& nestedBuilder, Location /*nested_loc*/, ValueRange args) {
+        nestedBuilder.create<linalg::YieldOp>(loc, *args.begin());
+      },
+      prunedAttrs);
+  newOp->setAttr(kDISCLinalgTransformName, b.getStringAttr(name));
+  mapping.erase(out);
+  mapping.map(out, newOp->getResult(0));
+  return success();
+}
+
+template <typename OpTy>
+LogicalResult emitElemOp(OpTy op, OpBuilder& b, BlockAndValueMapping& mapping,
+                         const std::string& name) {
+  if (!isa<lmhlo::LmhloOp>(op.getOperation()) || op->getNumOperands() < 2)
+    return op->emitError() << "not support lmhlo op\n";
+
+  auto loc = op.getLoc();
+  // Find maximum rank / number of loops.
+  auto getRank = [](Value v) {
+    return v.getType().cast<ShapedType>().getRank();
+  };
+  auto isScalar = [&](Value v) { return getRank(v) == 0; };
+  auto ins = op->getOperands().drop_back();
+  auto it = llvm::find_if_not(ins, isScalar);
+  Value maxRankArg = it != ins.end() ? *it : ins.front();
+  int64_t nloops = getRank(maxRankArg);
+  if (nloops == 0) return op->emitError() << "not support 0d elem ops a.t.m.\n";
+
+  // Apply only if all operands are scalar or have the same rank. Some ops,
+  // like `mhlo.select`, support implicit broadcasting of scalars.
+  if (!llvm::all_of(op->getOperands(), [&](Value v) {
+        int64_t r = getRank(v);
+        return r == 0 || r == nloops;
+      })) {
+    return op->emitError() << "Operands must be os same rank or scalar.\n";
+  }
+
+  Value out = op->getOperands().back();
+  if (getRank(out) != nloops)
+    return op->emitError() << "output of op should have the max rank\n";
+  Value newOut = mapping.lookup(out);
+  SmallVector<Value> newIns;
+  AffineMap scalarMap = AffineMap::get(nloops, 0, b.getContext());
+  AffineMap idMap = b.getMultiDimIdentityMap(nloops);
+  SmallVector<AffineMap, 4> maps;
+  for (Value in : ins) {
+    newIns.push_back(mapping.lookup(in));
+    maps.push_back(isScalar(in) ? scalarMap : idMap);
+  }
+  maps.push_back(b.getMultiDimIdentityMap(nloops));
+
+  auto newOp = b.create<linalg::GenericOp>(
+      loc, TypeRange{newOut.getType()}, newIns, newOut, maps,
+      mhlo::getParallelAndReductionIterators(
+          /*nLoops=*/nloops,
+          /*nReduction=*/0),
+      [&](OpBuilder& nestedBuilder, Location /*nested_loc*/, ValueRange args) {
+        Type innerResultTy = getElementTypeOrSelf(newOut);
+        auto argvec = llvm::to_vector<2>(args.take_front(newIns.size()));
+        auto innerResult = LhloOpToStdScalarOp::map<OpTy>(
+            op, innerResultTy, argvec, &nestedBuilder);
+        nestedBuilder.create<linalg::YieldOp>(loc, innerResult);
+      },
+      linalg::getPrunedAttributeList(op));
+
+  newOp->setAttr(kDISCLinalgTransformName, b.getStringAttr(name));
+  mapping.erase(out);
+  mapping.map(out, newOp->getResult(0));
+  return success();
+}
+
 LogicalResult emitLmhloOp(Operation* op, OpBuilder& b,
                           BlockAndValueMapping& mapping,
                           const std::string& name) {
@@ -210,7 +333,57 @@ LogicalResult emitLmhloOp(Operation* op, OpBuilder& b,
     return emitDotGeneralOp(dotGeneralOp, b, mapping, name);
   } else if (auto constOp = dyn_cast<lmhlo::ConstantOp>(op)) {
     return emitConstOp(constOp, b, mapping, name);
+  } else if (isa<lmhlo::BroadcastInDimOp, lmhlo::BroadcastOp,
+                 lmhlo::DynamicBroadcastInDimOp>(op)) {
+    return emitBcastOp(op, b, mapping, name);
   }
+
+  // clang-format off
+
+  #define ELEM_OP_HANDLER(opType) \
+    else if (auto t##opType = dyn_cast<lmhlo::opType>(op)) { \
+      return emitElemOp(t##opType, b, mapping, name); \
+    }
+  ELEM_OP_HANDLER(AddOp)
+  ELEM_OP_HANDLER(SubtractOp)
+  ELEM_OP_HANDLER(AbsOp)
+  ELEM_OP_HANDLER(CeilOp)
+  ELEM_OP_HANDLER(ConvertOp)
+  ELEM_OP_HANDLER(CopyOp)
+  ELEM_OP_HANDLER(CosineOp)
+  ELEM_OP_HANDLER(ExpOp)
+  ELEM_OP_HANDLER(FloorOp)
+  ELEM_OP_HANDLER(IsFiniteOp)
+  ELEM_OP_HANDLER(LogOp)
+  ELEM_OP_HANDLER(Log1pOp)
+  ELEM_OP_HANDLER(LogisticOp)
+  ELEM_OP_HANDLER(NegOp)
+  ELEM_OP_HANDLER(NotOp)
+  ELEM_OP_HANDLER(RsqrtOp)
+  ELEM_OP_HANDLER(SignOp)
+  ELEM_OP_HANDLER(SineOp)
+  ELEM_OP_HANDLER(SqrtOp)
+  ELEM_OP_HANDLER(TanhOp)
+  ELEM_OP_HANDLER(RoundOp)
+  ELEM_OP_HANDLER(RoundNearestEvenOp)
+  ELEM_OP_HANDLER(AddOp)
+  ELEM_OP_HANDLER(AndOp)
+  ELEM_OP_HANDLER(CompareOp)
+  ELEM_OP_HANDLER(DivOp)
+  ELEM_OP_HANDLER(MaxOp)
+  ELEM_OP_HANDLER(MinOp)
+  ELEM_OP_HANDLER(MulOp)
+  ELEM_OP_HANDLER(OrOp)
+  ELEM_OP_HANDLER(PowOp)
+  ELEM_OP_HANDLER(RemOp)
+  ELEM_OP_HANDLER(SubtractOp)
+  ELEM_OP_HANDLER(SelectOp)
+  ELEM_OP_HANDLER(ClampOp)
+
+  #undef ELEM_OP_HANDLER
+
+  // clang-format on
+
   // TODO(wyzero): support other lmhlo ops.
   return failure();
 }


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 ff1cbaf84e7..dcaf4312ba7 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
@@ -102,4 +102,46 @@ func.func @matmul_tt(%arg1: memref<?x?xf32, "cpu">, %arg2: memref<?x?xf32, "cpu"
     "lmhlo.terminator"() : () -> ()
   }) {disc.device = "cpu", disc.fusion.name = "matmul_nn_kTransform_dot_general__1_1_0", disc.fusion_type = "kTransform"} : () -> ()
   return %arg3 : memref<?x?xf32, "cpu">
-}
\ No newline at end of file
+}
+
+// -----
+
+// CHECK: #map = affine_map<(d0, d1) -> (d1)>
+// CHECK: #map1 = affine_map<(d0, d1) -> (d0, d1)>
+
+// CHECK-LABEL: @matmul_epilogue
+// CHECK-SAME: (%[[ARG0:.*]]: tensor<?x3072xf32>, %[[ARG1:.*]]: tensor<2xindex>, %[[ARG2:.*]]: tensor<3072x768xf32>,
+// CHECK-SAME: %[[ARG3:.*]]: tensor<768xf32>, %[[ARG4:.*]]: tensor<?x768xf32>, %[[ARG5:.*]]: tensor<?x768xf32>,
+// CHECK-SAME: %[[ARG6:.*]]: tensor<?x768xf32>)
+func.func @matmul_epilogue(%arg0: memref<?x3072xf32, "cpu">, %arg1: memref<2xindex, "cpu">, %arg2: memref<3072x768xf32, "cpu">, %arg3: memref<768xf32, "cpu">, %arg4: memref<?x768xf32, "cpu">, %arg5: memref<?x768xf32, "cpu">, %arg6: memref<?x768xf32, "cpu">) -> memref<?x768xf32, "cpu"> {
+  // CHECK: %[[T0:.*]] = disc_linalg_ext.constant_wrapper {disc.transform.name = "constant"} dense<-8.000000e-01> : tensor<3072x768xf32>
+  // CHECK: %[[T1:.*]] = disc_linalg_ext.constant_wrapper {disc.transform.name = "constant_1"} dense<-1.000000e-01> : tensor<768xf32>
+  // CHECK: %[[T2:.*]] = arith.constant 0.000000e+00 : f32
+
+  // CHECK: %[[T3:.*]] = linalg.fill {disc.transform.name = "dot_general"} ins(%[[T2]] : f32) outs(%[[ARG4]] : tensor<?x768xf32>) -> tensor<?x768xf32>
+  // CHECK: %[[T4:.*]] = linalg.matmul {disc.transform.name = "dot_general"} ins(%[[ARG0]], %[[T0]] : tensor<?x3072xf32>, tensor<3072x768xf32>) outs(%[[T3]] : tensor<?x768xf32>) -> tensor<?x768xf32>
+
+  // CHECK: %[[T5:.*]] = linalg.generic
+  // CHECK-SAME: indexing_maps = [#map, #map1]
+  // CHECK-SAME: iterator_types = ["parallel", "parallel"]
+  // CHECK-SAME: ins(%[[T1]] : tensor<768xf32>) outs(%[[ARG5]] : tensor<?x768xf32>)
+  // CHECK-SAME: disc.transform.name = "dynamic_broadcast_in_dim"
+
+  // CHECK: %[[T6:.*]] = linalg.generic
+  // CHECK-SAME: indexing_maps = [#map1, #map1, #map1]
+  // CHECK-SAME: iterator_types = ["parallel", "parallel"]
+  // CHECK-SAME: ins(%[[T4]], %[[T5]] : tensor<?x768xf32>, tensor<?x768xf32>) outs(%[[ARG6]] : tensor<?x768xf32>)
+  // CHECK-SAME: disc.transform.name = "subtract"
+
+  // CHECK: return %[[T6]]
+
+  "lmhlo.fusion"() ({
+    "lmhlo.constant"(%arg2) {disc.device = "cpu", value = dense<-8.000000e-01> : tensor<3072x768xf32>} : (memref<3072x768xf32, "cpu">) -> ()
+    "lmhlo.constant"(%arg3) {disc.device = "cpu", value = dense<-1.000000e-01> : tensor<768xf32>} : (memref<768xf32, "cpu">) -> ()
+    "lmhlo.dot_general"(%arg0, %arg2, %arg4) {disc.device = "cpu", dot_dimension_numbers = #mhlo.dot<lhs_contracting_dimensions = [1], rhs_contracting_dimensions = [0]>} : (memref<?x3072xf32, "cpu">, memref<3072x768xf32, "cpu">, memref<?x768xf32, "cpu">) -> ()
+    "lmhlo.dynamic_broadcast_in_dim"(%arg3, %arg1, %arg5) {broadcast_dimensions = dense<1> : tensor<1xi64>, disc.device = "cpu"} : (memref<768xf32, "cpu">, memref<2xindex, "cpu">, memref<?x768xf32, "cpu">) -> ()
+    "lmhlo.subtract"(%arg4, %arg5, %arg6) {disc.device = "cpu"} : (memref<?x768xf32, "cpu">, memref<?x768xf32, "cpu">, memref<?x768xf32, "cpu">) -> ()
+    "lmhlo.terminator"() : () -> ()
+  }) {disc.device = "cpu", disc.fusion.name = "test", disc.fusion_type = "kTransform"} : () -> ()
+  return %arg6 : memref<?x768xf32, "cpu">
+}