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 37550529f0a..2af04c34949 100644
--- a/tao_compiler/mlir/disc/tools/disc-transform/TransformOps/TransformOpsExt.cc
+++ b/tao_compiler/mlir/disc/tools/disc-transform/TransformOps/TransformOpsExt.cc
@@ -1080,6 +1080,130 @@ DiagnosedSilenceableFailure LowerMultiLevelPackToLoopOp::applyToOne(
   return DiagnosedSilenceableFailure(success());
 }
 
+//===---------------------------------------------------------------------===//
+// InlineReductionInitializer
+//===---------------------------------------------------------------------===//
+
+namespace {
+
+/// Given two MemRefTypes `aT` and `bT`, return a MemRefType to which both can
+/// be cast. If the MemRefTypes don't have the same rank or are not strided,
+/// return null; otherwise:
+///   1. if `aT` and `bT` are cast-compatible, return `aT`.
+///   2. else return a new MemRefType obtained by iterating over the shape and
+///   strides and:
+///     a. keeping the ones that are static and equal across `aT` and `bT`.
+///     b. using a dynamic shape and/or stride for the dimensions that don't
+///        agree.
+MemRefType getCastCompatibleMemRefType(MemRefType aT, MemRefType bT) {
+  if (memref::CastOp::areCastCompatible(aT, bT)) return aT;
+  if (aT.getRank() != bT.getRank()) return MemRefType();
+  int64_t aOffset, bOffset;
+  SmallVector<int64_t, 4> aStrides, bStrides;
+  if (failed(getStridesAndOffset(aT, aStrides, aOffset)) ||
+      failed(getStridesAndOffset(bT, bStrides, bOffset)) ||
+      aStrides.size() != bStrides.size())
+    return MemRefType();
+
+  ArrayRef<int64_t> aShape = aT.getShape(), bShape = bT.getShape();
+  int64_t resOffset;
+  SmallVector<int64_t, 4> resShape(aT.getRank(), 0),
+      resStrides(bT.getRank(), 0);
+  for (int64_t idx = 0, e = aT.getRank(); idx < e; ++idx) {
+    resShape[idx] =
+        (aShape[idx] == bShape[idx]) ? aShape[idx] : ShapedType::kDynamicSize;
+    resStrides[idx] = (aStrides[idx] == bStrides[idx])
+                          ? aStrides[idx]
+                          : ShapedType::kDynamicStrideOrOffset;
+  }
+  resOffset =
+      (aOffset == bOffset) ? aOffset : ShapedType::kDynamicStrideOrOffset;
+  return MemRefType::get(
+      resShape, aT.getElementType(),
+      StridedLayoutAttr::get(aT.getContext(), resOffset, resStrides));
+}
+
+}  // namespace
+
+void InlineReductionInitializerOp::getEffects(
+    SmallVectorImpl<MemoryEffects::EffectInstance>& effects) {
+  transform::consumesHandle({this->getOperation()->getOperand(0)}, effects);
+  transform::onlyReadsHandle({this->getOperation()->getOperand(1)}, effects);
+  transform::consumesHandle({this->getOperation()->getOperand(2)}, effects);
+  transform::producesHandle(this->getOperation()->getResults(), effects);
+  transform::modifiesPayload(effects);
+}
+
+DiagnosedSilenceableFailure InlineReductionInitializerOp::apply(
+    transform::TransformResults& results, transform::TransformState& state) {
+  ArrayRef<Operation*> initOps = state.getPayloadOps(getInitializer());
+  if (initOps.size() != 1)
+    return mlir::emitDefiniteFailure(this->getOperation(),
+                                     "expect only one initializer but got ")
+           << initOps.size();
+  auto fillOp = dyn_cast<linalg::FillOp>(initOps[0]);
+  if (!fillOp)
+    return mlir::emitDefiniteFailure(this->getOperation(),
+                                     "expect the initializer be a fill op");
+
+  ArrayRef<Operation*> loopOps = state.getPayloadOps(getLoop());
+  if (loopOps.size() != 1)
+    return mlir::emitDefiniteFailure(this->getOperation(),
+                                     "expect only one loop but got ")
+           << loopOps.size();
+  auto forOp = dyn_cast<scf::ForOp>(loopOps[0]);
+  if (!forOp)
+    return mlir::emitDefiniteFailure(this->getOperation(),
+                                     "expect the loop to a for op");
+
+  ArrayRef<Operation*> readerOps = state.getPayloadOps(getReader());
+  if (readerOps.size() != 1)
+    return mlir::emitDefiniteFailure(this->getOperation(),
+                                     "expect only one reader but got ")
+           << readerOps.size();
+  auto readerOp = dyn_cast<vector::TransferReadOp>(readerOps[0]);
+  if (!readerOp)
+    return mlir::emitDefiniteFailure(this->getOperation(),
+                                     "expect the reader to a transfer_read op");
+
+  OpBuilder b(fillOp.getOperation());
+  auto vectorTy = readerOp->getResultTypes()[0].cast<VectorType>();
+  auto memrefTy =
+      MemRefType::get(vectorTy.getShape(), vectorTy.getElementType());
+  Value newInitBuffer = b.create<memref::AllocaOp>(
+      fillOp.getLoc(), memrefTy, ValueRange{}, b.getI64IntegerAttr(64));
+  b.create<linalg::FillOp>(fillOp.getLoc(), ValueRange{readerOp.getPadding()},
+                           ValueRange{newInitBuffer});
+
+  b.setInsertionPoint(readerOp.getOperation());
+  Value firstIter =
+      b.create<arith::CmpIOp>(readerOp.getLoc(), arith::CmpIPredicate::eq,
+                              forOp.getLowerBound(), forOp.getInductionVar());
+  auto castTy = getCastCompatibleMemRefType(
+      readerOp.getSource().getType().cast<MemRefType>(), memrefTy);
+  if (!castTy) {
+    return mlir::emitDefiniteFailure(this->getOperation(),
+                                     "failed to find compatible memref type");
+  }
+
+  auto ifOp = b.create<scf::IfOp>(readerOp.getLoc(), TypeRange{castTy},
+                                  firstIter, true);
+  b.setInsertionPointToStart(ifOp.thenBlock());
+  Value thenResult =
+      b.create<memref::CastOp>(readerOp.getLoc(), castTy, newInitBuffer)
+          ->getResult(0);
+  b.create<scf::YieldOp>(readerOp.getLoc(), thenResult);
+  b.setInsertionPointToStart(ifOp.elseBlock());
+  Value elseResult =
+      b.create<memref::CastOp>(readerOp.getLoc(), castTy, readerOp.getSource())
+          ->getResult(0);
+  b.create<scf::YieldOp>(readerOp.getLoc(), elseResult);
+  readerOp->replaceUsesOfWith(readerOp.getSource(), ifOp->getResult(0));
+  fillOp->erase();
+  results.set(getResult().cast<OpResult>(), {readerOp});
+  return DiagnosedSilenceableFailure(success());
+}
+
 }  // namespace transform_dialect
 
 void registerTransformDialectCommonExtension(DialectRegistry& registry) {
diff --git a/tao_compiler/mlir/disc/tools/disc-transform/TransformOps/TransformOpsExt.td b/tao_compiler/mlir/disc/tools/disc-transform/TransformOps/TransformOpsExt.td
index 95de7988c5b..91709c9bb31 100644
--- a/tao_compiler/mlir/disc/tools/disc-transform/TransformOps/TransformOpsExt.td
+++ b/tao_compiler/mlir/disc/tools/disc-transform/TransformOps/TransformOpsExt.td
@@ -263,4 +263,54 @@ def LowerMultiLevelPackToLoopOp : Op<Transform_Dialect, "disc.lower_multi_level_
   }];
 }
 
+def InlineReductionInitializerOp : Op<Transform_Dialect, "disc.inline_reduction_initializer",
+    [DeclareOpInterfaceMethods<MemoryEffectsOpInterface>,
+     DeclareOpInterfaceMethods<TransformOpInterface>]> {
+  let description = [{
+    Given a buffer popluated by a reduction loop named `l0`, inline the initializer of the buffer
+    into the loop `l0`.
+
+    Returns the new xfer_read op.
+
+    Example:
+     convert from:
+     ```
+      linalg.fill ins(%cst : f32) outs(%0 : memref<?x?xf32>)
+      scf.for %iv = %init to %stop step %step {
+        %1 = memref.subview %0[...] : memref<?x?xf32> to memref<?x?xf32, strided<[?, 1], offset: 0>>
+        %2 = vector.transfer_read %1[%c0, %c0], %cst : memref<?x?xf32, strided<[?, 1], offset: 0>>, vector<1x16xf32>
+        use(%2)
+      }
+     ```
+     to:
+     ```
+      %0 = memref.alloca() : memref<1x16xf32>
+      linalg.fill ins(%cst : f32) outs(%0 : memref<1x16xf32>)
+      scf.for %iv = %init to %stop step %step {
+        %1 = memref.subview %0[...] : memref<?x?xf32> to memref<?x?xf32, strided<[?, 1], offset: 0>>
+        %first_step = arith.cmpi eq %iv, %init : i1
+        %2 = scf.if %first_step -> memref<?x?xf32, strided<[?, 1], offset: ?>> {
+          %cast = memref.cast %0 :  memref<1x16xf32> to memref<?x?xf32, strided<[?, 1], offset: ?>>
+          scf.yield %cast : memref<?x?xf32, strided<[?, 1], offset: ?>>
+        } else {
+          %cast = memref.cast %1 :  memref<?x?xf32, strided<[?, 1], offset: 0>> to memref<?x?xf32, strided<[?, 1], offset: ?>>
+          scf.yield %cast : memref<?x?xf32, strided<[?, 1], offset: ?>>
+        }
+        %3 = vector.transfer_read %2[%c0, %c0], %cst : memref<?x?xf32, strided<[?, 1], offset: 0>>, vector<1x16xf32>
+        use(%3)
+      }
+     ```
+  }];
+
+  let arguments = (ins PDL_Operation:$initializer,
+                       PDL_Operation:$loop,
+                       PDL_Operation:$reader);
+  let results = (outs PDL_Operation:$result);
+
+  let assemblyFormat = [{
+    attr-dict $initializer `for` `reader` $reader  `into` `loop` $loop
+  }];
+  let cppNamespace = "::mlir::disc_ral::transform_dialect";
+}
+
 #endif // DISC_TRANSFORM_OPS_EXT


Test patch:
diff --git a/tao_compiler/mlir/disc/tools/disc-transform/transforms/tests/inline-reduction-initializer.mlir b/tao_compiler/mlir/disc/tools/disc-transform/transforms/tests/inline-reduction-initializer.mlir
new file mode 100644
index 00000000000..b5381259943
--- /dev/null
+++ b/tao_compiler/mlir/disc/tools/disc-transform/transforms/tests/inline-reduction-initializer.mlir
@@ -0,0 +1,51 @@
+// RUN: disc-opt --disc-transform-dialect-interpreter -split-input-file %s | FileCheck %s
+
+
+// CHECK-LABEL: @inline_reduction_loop_initializer
+// CHECK-SAME: (%[[ARG0:.*]]: memref<?x?xf32>, %[[ARG1:.*]]: memref<?xf32>)
+func.func @inline_reduction_loop_initializer(%arg0: memref<?x?xf32>, %arg1: memref<?xf32>) -> memref<?x?xf32> {
+  %cst = arith.constant 0.000000e+00 : f32
+  %c0 = arith.constant 0 : index
+  %c1 = arith.constant 1 : index
+  %c16 = arith.constant 16 : index
+  %c372 = arith.constant 372 : index
+  %d0 = memref.dim %arg0, %c0 : memref<?x?xf32>
+  %d1 = memref.dim %arg0, %c1 : memref<?x?xf32>
+  // CHECK: %[[NEW_INIT_BUFFER:.*]] = memref.alloca() {alignment = 64 : i64} : memref<16xf32>
+  // CHECK-NEXT: linalg.fill ins(%cst : f32) outs(%[[NEW_INIT_BUFFER]] : memref<16xf32>)
+  // CHECK-NOT: linalg.fill
+  // CHECK: scf.for %[[IV0:.*]] = %c0 to %[[DIM0:.*]] step %c372
+  // CHECK-NEXT:  %[[FIRST_ITER:.*]] = arith.cmpi eq, %[[IV0]], %c0 : index
+  // CHECK-NEXT:  scf.for %[[IV1:.*]] = %c0 to %[[DIM1:.*]] step %c16
+  // CHECK-NEXT:    %[[T0:.*]] = scf.if %[[FIRST_ITER]] -> (memref<?xf32>) {
+  // CHECK-NEXT:      %[[CAST0:.*]] = memref.cast %[[NEW_INIT_BUFFER]] : memref<16xf32> to memref<?xf32>
+  // CHECK-NEXT:      scf.yield %[[CAST0]] : memref<?xf32>
+  // CHECK-NEXT:    } else {
+  // CHECK-NEXT:      scf.yield %[[ARG1]] : memref<?xf32>
+  // CHECK-NEXT:    }
+  // CHECK-NEXT:    %[[T1:.*]] = vector.transfer_read %[[T0]]
+  linalg.fill ins(%cst : f32) outs(%arg1 : memref<?xf32>)
+  scf.for %iv0 = %c0 to %d0 step %c372 {
+    scf.for %iv1 = %c0 to %d1 step %c16 {
+      %0 = vector.transfer_read %arg1[%iv1], %cst : memref<?xf32>, vector<16xf32>
+      %1 = scf.for %iv2 = %c0 to %c372 step %c1 iter_args(%acc = %0) -> (vector<16xf32>) {
+        %iv3 = arith.addi %iv0, %iv1 : index
+        %1 = vector.transfer_read %arg0[%iv3, %iv1], %cst : memref<?x?xf32>, vector<16xf32>
+        %2 = arith.addf %1, %acc : vector<16xf32>
+        scf.yield %2 : vector<16xf32>
+      }
+      vector.transfer_write %1, %arg1[%iv1] : vector<16xf32>, memref<?xf32>
+    }
+  }
+
+  return %arg0: memref<?x?xf32>
+}
+
+transform.structured.canonicalized_sequence failures(propagate) {
+^bb0(%arg0: !pdl.operation):
+  %fill = transform.structured.match ops{["linalg.fill"]} in %arg0
+  %readers = transform.structured.match ops{["vector.transfer_read"]} in %arg0
+  %reader_for_output, %reader_for_input = split_handles %readers in [2] : (!pdl.operation) -> (!pdl.operation, !pdl.operation)
+  %loop = transform.loop.get_parent_for %reader_for_output {num_loops = 2 : i64} : (!pdl.operation) -> !pdl.operation
+  transform.disc.inline_reduction_initializer %fill for reader %reader_for_output into loop %loop
+}