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/TransformOps/TransformOpsExt.cc b/tao_compiler/mlir/disc/tools/disc-transform/TransformOps/TransformOpsExt.cc
index f95658c0eb4..f9ddfeb2c36 100644
--- a/tao_compiler/mlir/disc/tools/disc-transform/TransformOps/TransformOpsExt.cc
+++ b/tao_compiler/mlir/disc/tools/disc-transform/TransformOps/TransformOpsExt.cc
@@ -396,9 +396,10 @@ struct TransferReadOfFillOpPattern
 };
 
 /// convert:
-///   %0 = tensor.extract_slice %arg0[0, 0] [%s0, %s1] [1, 1] : tensor<?x?xf32>
-///   to tensor<?x?xf32> %1 = linalg.fill ins(%cst : f32) outs(%0 :
-///   tensor<?x?xf32>) -> tensor<?x?xf32>
+///   %0 = tensor.extract_slice %arg0[0, 0] [%s0, %s1] [1, 1]
+///       : tensor<?x?xf32> to tensor<?x?xf32>
+///   %1 = linalg.fill ins(%cst : f32) outs(%0
+///       : tensor<?x?xf32>) -> tensor<?x?xf32>
 ///   %2 = vector.transfer_write %result, %1[0, 0]
 ///                     : vector<6x16xf32>, tensor<?x?xf32>
 /// to:
@@ -500,6 +501,80 @@ DiagnosedSilenceableFailure ApplyPatternsOp::applyToOne(
   return DiagnosedSilenceableFailure(success());
 }
 
+//===---------------------------------------------------------------------===//
+// FoldProducerExtractSliceOp
+//===---------------------------------------------------------------------===//
+
+namespace {
+
+LogicalResult tryMergeExtractSliceOp(OpBuilder& b,
+                                     tensor::ExtractSliceOp& sliceOp,
+                                     tensor::ExtractSliceOp producerOp) {
+  if (!strideAllOnes(sliceOp) || !strideAllOnes(producerOp)) return failure();
+  Location loc = sliceOp->getLoc();
+  SmallVector<OpFoldResult> newOffsets;
+  for (const auto& z :
+       llvm::zip(sliceOp.getMixedOffsets(), producerOp.getMixedOffsets())) {
+    Value lhsValue = std::get<0>(z).dyn_cast<Value>();
+    Value rhsValue = std::get<1>(z).dyn_cast<Value>();
+    if (!lhsValue) {
+      lhsValue = b.create<arith::ConstantIndexOp>(
+          loc,
+          std::get<0>(z).dyn_cast<Attribute>().cast<IntegerAttr>().getInt());
+    }
+    if (!rhsValue) {
+      rhsValue = b.create<arith::ConstantIndexOp>(
+          loc,
+          std::get<1>(z).dyn_cast<Attribute>().cast<IntegerAttr>().getInt());
+    }
+    newOffsets.push_back(
+        b.create<arith::AddIOp>(loc, lhsValue, rhsValue).getResult());
+  }
+  sliceOp = b.create<tensor::ExtractSliceOp>(
+      loc, producerOp.getSource(), newOffsets, sliceOp.getMixedSizes(),
+      sliceOp.getMixedStrides());
+  return success();
+}
+
+}  // namespace
+
+void FoldProducerExtractSliceOp::build(OpBuilder& builder,
+                                       OperationState& result, Value target,
+                                       int64_t maxRepeatNum) {
+  MLIRContext* ctx = builder.getContext();
+  result.addOperands(target);
+  result.addAttribute(
+      FoldProducerExtractSliceOp::getMaxRepeatNumAttrName(result.name),
+      builder.getIntegerAttr(builder.getIntegerType(64), maxRepeatNum));
+  result.addTypes({pdl::OperationType::get(ctx)});
+}
+
+DiagnosedSilenceableFailure FoldProducerExtractSliceOp::applyToOne(
+    Operation* target, SmallVectorImpl<Operation*>& results,
+    transform::TransformState& state) {
+  auto sliceOp = dyn_cast<tensor::ExtractSliceOp>(target);
+  if (!sliceOp) {
+    return mlir::emitDefiniteFailure(target,
+                                     "applies only to tensor::ExtractSliceOp");
+  }
+  MLIRContext* ctx = target->getContext();
+
+  OpBuilder b(target);
+  for (int64_t i = 0; i < getMaxRepeatNum(); ++i) {
+    auto producerOp =
+        sliceOp.getSource().getDefiningOp<tensor::ExtractSliceOp>();
+    if (!producerOp || failed(tryMergeExtractSliceOp(b, sliceOp, producerOp)))
+      break;
+  }
+
+  if (sliceOp.getOperation() != target) {
+    target->getResult(0).replaceAllUsesWith(sliceOp.getResult());
+  }
+
+  results.assign({sliceOp.getOperation()});
+  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 ee824fdd512..2af4f1acc9f 100644
--- a/tao_compiler/mlir/disc/tools/disc-transform/TransformOps/TransformOpsExt.td
+++ b/tao_compiler/mlir/disc/tools/disc-transform/TransformOps/TransformOpsExt.td
@@ -101,4 +101,42 @@ def ApplyPatternsOp : Op<Transform_Dialect, "disc.apply_patterns",
   }];
 }
 
+def FoldProducerExtractSliceOp : Op<Transform_Dialect, "disc.fold_producer_extract_slice",
+    [FunctionalStyleTransformOpTrait,
+     MemoryEffectsOpInterface,
+     TransformEachOpTrait,
+     TransformOpInterface]> {
+  let description = [{
+    Given a tensor.ExtractSliceOp, greedily fold its producers if they are also tensor.ExtractSliceOp.
+
+    Returns the folded new tensor.ExtractSliceOp.
+
+    The following additive attributes can be set:
+      - max_repeat_num: fold at most `max_repeat_num` times.
+
+    Return modes:
+    =============
+    This operation try to fold two tensor.ExtractSliceOp with def-use relationship at most
+    `max_repeat_num` times.
+  }];
+
+  let arguments = (ins PDL_Operation:$target,
+                       I64Attr:$max_repeat_num);
+  let results = (outs PDL_Operation:$result);
+
+  let assemblyFormat = "$target attr-dict";
+  let cppNamespace = "::mlir::disc_ral::transform_dialect";
+
+  let builders = [
+    OpBuilder<(ins "Value":$target, "int64_t":$max_repeat_num)>
+  ];
+
+  let extraClassDeclaration = [{
+    ::mlir::DiagnosedSilenceableFailure applyToOne(
+        ::mlir::Operation *target,
+        ::llvm::SmallVectorImpl<::mlir::Operation *> &results,
+        ::mlir::transform::TransformState &state);
+  }];
+}
+
 #endif // DISC_TRANSFORM_OPS_EXT


Test patch:
diff --git a/tao_compiler/mlir/disc/tests/disc-transform/data/matmul_nn_d_f32_large_schedule_3.mlir b/tao_compiler/mlir/disc/tests/disc-transform/data/matmul_nn_d_f32_large_schedule_3.mlir
new file mode 100644
index 00000000000..82b37c52104
--- /dev/null
+++ b/tao_compiler/mlir/disc/tests/disc-transform/data/matmul_nn_d_f32_large_schedule_3.mlir
@@ -0,0 +1,57 @@
+transform.structured.canonicalized_sequence failures(propagate) {
+^bb1(%arg1: !pdl.operation):
+  %fill = transform.structured.match ops{["linalg.fill"]} in %arg1
+  %matmul = transform.structured.match ops{["linalg.matmul"]} in %arg1
+
+  %0:2 = transform.structured.tile_to_foreach_thread_op %matmul num_threads [1, 1]
+  transform.structured.fuse_into_containing_op %fill into %0#0
+
+  // first level tile and fuse matmul and fill op.
+  %1:3 = transform.structured.fuse %0#1 {tile_sizes = [288, 256, 0], tile_interchange = [0, 1, 2]}
+  // second level tile and fuse matmul and fill op.
+  %2:3 = transform.structured.fuse %1#0 {tile_sizes = [6, 16, 0], tile_interchange = [0, 1, 2]}
+
+  // gemm reduction axis tiling
+  %3:2 = transform.structured.tile %2#0 [0, 0, 1] {interchange=[0, 1, 2]}
+
+  // clean up
+  %func0 = transform.structured.match ops{["func.func"]} in %arg1
+  transform.disc.apply_patterns %func0 {canonicalization}
+  // fold two extract_slice ops generated by two-level tiling. It's needed to enable following
+  // pad and hosit schedule.
+  %weight_inner_slice = get_producer_of_operand %3#0[1] : (!pdl.operation) -> !pdl.operation
+  transform.disc.fold_producer_extract_slice %weight_inner_slice {max_repeat_num = 1}
+
+  // pad to match the requirement of hardware vector/tensor instruction.
+  %4 = transform.structured.pad %3#0 {padding_values=[0.0 : f32, 0.0 : f32, 0.0 : f32], padding_dimensions=[0, 1, 2], pack_paddings=[1, 1, 0], hoist_paddings=[4, 5, 0], transpose_paddings=[[1, 0], [0, 1], [0, 1]]}
+
+  %func1 = transform.structured.match ops{["func.func"]} in %arg1
+  transform.structured.vectorize %func1 {vectorize_padding}
+
+  %func2 = transform.structured.match ops{["func.func"]} in %arg1
+  transform.disc.apply_patterns %func2 {canonicalization}
+
+  transform.disc.bufferize %arg1
+
+  transform.lower_vectors {
+    contraction_lowering = "outerproduct",
+    multireduction_lowering = "innerparallel",
+    split_transfers = "linalg-copy",
+    // stages = [0, 1, 2, 3, 4, 5, 6, 7],
+    stages = [0, 1, 2, 3],
+    transpose_avx2_lowering = false,
+    transpose_lowering = "eltwise",
+    unroll_vector_transfers = true
+  }
+
+  transform.lower_vectors {
+    contraction_lowering = "outerproduct",
+    multireduction_lowering = "innerparallel",
+    split_transfers = "linalg-copy",
+    // stages = [0, 1, 2, 3, 4, 5, 6, 7],
+    stages = [5, 6, 7],
+    transpose_avx2_lowering = false,
+    transpose_lowering = "eltwise",
+    unroll_vector_transfers = true
+  }
+}
\ No newline at end of file
diff --git a/tao_compiler/mlir/disc/tests/disc-transform/matmul.cc b/tao_compiler/mlir/disc/tests/disc-transform/matmul.cc
index 6a3cc3f3208..536afb7c1a9 100644
--- a/tao_compiler/mlir/disc/tests/disc-transform/matmul.cc
+++ b/tao_compiler/mlir/disc/tests/disc-transform/matmul.cc
@@ -135,4 +135,104 @@ TEST(SimpleTest, MatMulF32_1024x1024x1024_2) {
       /*profiling*/ true));
 }
 
+TEST(SimpleTest, MatMulF32_304x256x256_3) {
+  EnvSetting setting = {
+      {"DISC_TRANSFORM_SCHEDULE_FILE",
+       {c_ft_path + "matmul_nn_d_f32_large_schedule_3.mlir", false}},
+      {"DISC_ENABLE_TRANSFORM_SCHEDULE", {"1", false}},
+      {"DISC_ENABLE_SHAPE_CONSTRAINT_IR", {"1", "false"}},
+      {"DISC_MEM_INTENSIVE_OPT_EXPERIMENTAL", {"0", "false"}}};
+  EnvSettingContext ctx(setting);
+  EXPECT_TRUE(feature_test_main(
+      /*mlir_file_path*/ c_ft_path + "matmul_nn_d_f32.mlir",
+      /*backend_types*/ {BackendType::kAArch64},
+      /*num_inputs*/ 2,
+      /*num_outputs*/ 1,
+      /*input_descriptors*/ {"304x256xf32_X", "256x256xf32_X"},
+      /*output_descriptors*/ {"f32_X"},
+      /*input_vals*/ {},
+      /*expected_output_vals*/ {},
+      /*profiling*/ true));
+}
+
+TEST(SimpleTest, MatMulF32_304x512x256_3) {
+  EnvSetting setting = {
+      {"DISC_TRANSFORM_SCHEDULE_FILE",
+       {c_ft_path + "matmul_nn_d_f32_large_schedule_3.mlir", false}},
+      {"DISC_ENABLE_TRANSFORM_SCHEDULE", {"1", false}},
+      {"DISC_ENABLE_SHAPE_CONSTRAINT_IR", {"1", "false"}},
+      {"DISC_MEM_INTENSIVE_OPT_EXPERIMENTAL", {"0", "false"}}};
+  EnvSettingContext ctx(setting);
+  EXPECT_TRUE(feature_test_main(
+      /*mlir_file_path*/ c_ft_path + "matmul_nn_d_f32.mlir",
+      /*backend_types*/ {BackendType::kAArch64},
+      /*num_inputs*/ 2,
+      /*num_outputs*/ 1,
+      /*input_descriptors*/ {"304x256xf32_X", "256x512xf32_X"},
+      /*output_descriptors*/ {"f32_X"},
+      /*input_vals*/ {},
+      /*expected_output_vals*/ {},
+      /*profiling*/ true));
+}
+
+TEST(SimpleTest, MatMulF32_304x1024x256_3) {
+  EnvSetting setting = {
+      {"DISC_TRANSFORM_SCHEDULE_FILE",
+       {c_ft_path + "matmul_nn_d_f32_large_schedule_3.mlir", false}},
+      {"DISC_ENABLE_TRANSFORM_SCHEDULE", {"1", false}},
+      {"DISC_ENABLE_SHAPE_CONSTRAINT_IR", {"1", "false"}},
+      {"DISC_MEM_INTENSIVE_OPT_EXPERIMENTAL", {"0", "false"}}};
+  EnvSettingContext ctx(setting);
+  EXPECT_TRUE(feature_test_main(
+      /*mlir_file_path*/ c_ft_path + "matmul_nn_d_f32.mlir",
+      /*backend_types*/ {BackendType::kAArch64},
+      /*num_inputs*/ 2,
+      /*num_outputs*/ 1,
+      /*input_descriptors*/ {"304x256xf32_X", "256x1024xf32_X"},
+      /*output_descriptors*/ {"f32_X"},
+      /*input_vals*/ {},
+      /*expected_output_vals*/ {},
+      /*profiling*/ true));
+}
+
+TEST(SimpleTest, MatMulF32_304x1024x512_3) {
+  EnvSetting setting = {
+      {"DISC_TRANSFORM_SCHEDULE_FILE",
+       {c_ft_path + "matmul_nn_d_f32_large_schedule_3.mlir", false}},
+      {"DISC_ENABLE_TRANSFORM_SCHEDULE", {"1", false}},
+      {"DISC_ENABLE_SHAPE_CONSTRAINT_IR", {"1", "false"}},
+      {"DISC_MEM_INTENSIVE_OPT_EXPERIMENTAL", {"0", "false"}}};
+  EnvSettingContext ctx(setting);
+  EXPECT_TRUE(feature_test_main(
+      /*mlir_file_path*/ c_ft_path + "matmul_nn_d_f32.mlir",
+      /*backend_types*/ {BackendType::kAArch64},
+      /*num_inputs*/ 2,
+      /*num_outputs*/ 1,
+      /*input_descriptors*/ {"304x512xf32_X", "512x1024xf32_X"},
+      /*output_descriptors*/ {"f32_X"},
+      /*input_vals*/ {},
+      /*expected_output_vals*/ {},
+      /*profiling*/ true));
+}
+
+TEST(SimpleTest, MatMulF32_1024x1024x1024_3) {
+  EnvSetting setting = {
+      {"DISC_TRANSFORM_SCHEDULE_FILE",
+       {c_ft_path + "matmul_nn_d_f32_large_schedule_3.mlir", false}},
+      {"DISC_ENABLE_TRANSFORM_SCHEDULE", {"1", false}},
+      {"DISC_ENABLE_SHAPE_CONSTRAINT_IR", {"1", "false"}},
+      {"DISC_MEM_INTENSIVE_OPT_EXPERIMENTAL", {"0", "false"}}};
+  EnvSettingContext ctx(setting);
+  EXPECT_TRUE(feature_test_main(
+      /*mlir_file_path*/ c_ft_path + "matmul_nn_d_f32.mlir",
+      /*backend_types*/ {BackendType::kAArch64},
+      /*num_inputs*/ 2,
+      /*num_outputs*/ 1,
+      /*input_descriptors*/ {"1024x1024xf32_X", "1024x1024xf32_X"},
+      /*output_descriptors*/ {"f32_X"},
+      /*input_vals*/ {},
+      /*expected_output_vals*/ {},
+      /*profiling*/ true));
+}
+
 }  // namespace mlir_test
diff --git a/tao_compiler/mlir/disc/tools/disc-transform/transforms/tests/fold-producer-extract-slice.mlir b/tao_compiler/mlir/disc/tools/disc-transform/transforms/tests/fold-producer-extract-slice.mlir
new file mode 100644
index 00000000000..7c488645912
--- /dev/null
+++ b/tao_compiler/mlir/disc/tools/disc-transform/transforms/tests/fold-producer-extract-slice.mlir
@@ -0,0 +1,57 @@
+// RUN: disc-opt --disc-transform-dialect-interpreter -split-input-file %s | FileCheck %s
+
+// CHECK-LABEL: @fold_extracted_slice
+// CHECK-SAME: (%[[ARG0:.*]]: tensor<?x?xf32>, %[[ARG1:.*]]: index, %[[ARG2:.*]]: index, %[[ARG3:.*]]: index, %[[ARG4:.*]]: index)
+func.func @fold_extracted_slice(%arg0: tensor<?x?xf32>, %arg1: index, %arg2: index, %arg3: index, %arg4: index) -> tensor<?x?xf32> {
+  // CHECK: %[[T0:.*]] = arith.addi %[[ARG2]], %[[ARG1]]
+  // CHECK: %[[T1:.*]] = tensor.extract_slice %[[ARG0]]
+  // CHECK-SAME: %[[T0]], %[[T0]]
+  // CHECK-SAME: %[[ARG4]], %[[ARG4]]
+  // CHECK-SAME: 1, 1
+  // CHECK-NOT: tensor.extract_slice
+  // CHECK: %[[T2:.*]] = linalg.fill
+  // CHECK-SAME: %[[T1]]
+  // CHECK: return %[[T2]]
+  %0 = tensor.extract_slice %arg0[%arg1, %arg1][%arg3, %arg3][1, 1] : tensor<?x?xf32> to tensor<?x?xf32>
+  %1 = tensor.extract_slice %0[%arg2, %arg2][%arg4, %arg4][1, 1] : tensor<?x?xf32> to tensor<?x?xf32>
+  %cst = arith.constant 0.000000e+00 : f32
+  %2 = linalg.fill ins(%cst : f32) outs(%1 : tensor<?x?xf32>) -> tensor<?x?xf32>
+  return %2 : tensor<?x?xf32>
+}
+
+transform.structured.canonicalized_sequence failures(propagate) {
+^bb1(%arg1: !pdl.operation):
+  %fill = transform.structured.match ops{["linalg.fill"]} in %arg1
+  %slice = get_producer_of_operand %fill[1] : (!pdl.operation) -> !pdl.operation
+  transform.disc.fold_producer_extract_slice %slice {max_repeat_num = 1}
+}
+
+// -----
+
+// CHECK-LABEL: @fold_two_extracted_slice
+// CHECK-SAME: (%[[ARG0:.*]]: tensor<?x?xf32>, %[[ARG1:.*]]: index, %[[ARG2:.*]]: index, %[[ARG3:.*]]: index, %[[ARG4:.*]]: index)
+func.func @fold_two_extracted_slice(%arg0: tensor<?x?xf32>, %arg1: index, %arg2: index, %arg3: index, %arg4: index) -> tensor<?x?xf32> {
+  // CHECK: %[[T0:.*]] = arith.addi %[[ARG1]], %[[ARG2]]
+  // CHECK: %[[T1:.*]] = arith.addi %[[T0]], %[[ARG1]]
+  // CHECK: %[[T2:.*]] = tensor.extract_slice %[[ARG0]]
+  // CHECK-SAME: %[[T1]], %[[T1]]
+  // CHECK-SAME: %[[ARG4]], %[[ARG4]]
+  // CHECK-SAME: 1, 1
+  // CHECK-NOT: tensor.extract_slice
+  // CHECK: %[[T3:.*]] = linalg.fill
+  // CHECK-SAME: %[[T2]]
+  // CHECK: return %[[T3]]
+  %0 = tensor.extract_slice %arg0[%arg1, %arg1][%arg3, %arg3][1, 1] : tensor<?x?xf32> to tensor<?x?xf32>
+  %1 = tensor.extract_slice %0[%arg2, %arg2][%arg4, %arg4][1, 1] : tensor<?x?xf32> to tensor<?x?xf32>
+  %2 = tensor.extract_slice %1[%arg1, %arg1][%arg4, %arg4][1, 1] : tensor<?x?xf32> to tensor<?x?xf32>
+  %cst = arith.constant 0.000000e+00 : f32
+  %3 = linalg.fill ins(%cst : f32) outs(%2 : tensor<?x?xf32>) -> tensor<?x?xf32>
+  return %3 : tensor<?x?xf32>
+}
+
+transform.structured.canonicalized_sequence failures(propagate) {
+^bb1(%arg1: !pdl.operation):
+  %fill = transform.structured.match ops{["linalg.fill"]} in %arg1
+  %slice = get_producer_of_operand %fill[1] : (!pdl.operation) -> !pdl.operation
+  transform.disc.fold_producer_extract_slice %slice {max_repeat_num = 2}
+}