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 65bc1e9cdf3..8f73306d93d 100644
--- a/tao_compiler/mlir/disc/tools/disc-transform/BUILD
+++ b/tao_compiler/mlir/disc/tools/disc-transform/BUILD
@@ -285,6 +285,7 @@ cc_library(
         ":DISCLinalgExtDialect",
         ":disc_transform_utils",
         ":pass_details",
+        "//tensorflow/compiler/xla/mlir_hlo:legalize_to_linalg_utils",
         "//tensorflow/compiler/xla/mlir_hlo:lhlo",
         "@llvm-project//llvm:Support",
         "@llvm-project//mlir:FuncDialect",
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 e944c17f798..ff287212a39 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
@@ -11,9 +11,11 @@
 
 #include "llvm/Support/Debug.h"
 #include "mlir-hlo/Dialect/lhlo/IR/lhlo_ops.h"
+#include "mlir-hlo/Dialect/mhlo/transforms/legalize_to_linalg_utils.h"
 #include "mlir/Dialect/Arith/IR/Arith.h"
 #include "mlir/Dialect/Func/IR/FuncOps.h"
 #include "mlir/Dialect/Linalg/IR/Linalg.h"
+#include "mlir/Dialect/Linalg/Utils/Utils.h"
 #include "mlir/Dialect/Tensor/IR/Tensor.h"
 #include "mlir/IR/BlockAndValueMapping.h"
 #include "mlir/IR/MLIRContext.h"
@@ -93,18 +95,97 @@ LogicalResult emitDotGeneralOp(lmhlo::DotGeneralOp op, OpBuilder& b,
   Value newB = mapping.lookup(B);
   Value newC = mapping.lookup(C);
 
-  // firstly fill the output buffer using zero.
+  auto lhsTy = A.getType().cast<MemRefType>();
+  auto rhsTy = B.getType().cast<MemRefType>();
+  if (lhsTy.getRank() != 2 || rhsTy.getRank() != 2) {
+    return op->emitError() << "only support rank 2 GEMM a.t.m.\n";
+  }
+
+  auto dimNumbers = op.getDotDimensionNumbers();
+  auto lhsBatchingDims = dimNumbers.getLhsBatchingDimensions();
+  auto rhsBatchingDims = dimNumbers.getRhsBatchingDimensions();
+  auto lhsContractingDims = dimNumbers.getLhsContractingDimensions();
+  auto rhsContractingDims = dimNumbers.getRhsContractingDimensions();
+  if (lhsContractingDims.size() != 1 || rhsContractingDims.size() != 1) {
+    return op->emitError() << "only support exactly 1 contract dim\n";
+  }
+
+  auto numContracting = lhsContractingDims.size();
+  auto outputType = newC.getType().cast<ShapedType>();
+  auto targetRank = outputType.getRank();
+  auto totalLoopCount = numContracting + targetRank;
+  auto lhsRank = lhsTy.getRank();
+  auto lhsExtraDims =
+      lhsRank - lhsBatchingDims.size() - lhsContractingDims.size();
+  auto rhsRank = rhsTy.getRank();
+
+  // 1, fill the output buffer using zero.
   auto ty = newC.getType().cast<ShapedType>();
   auto zeroAttr = b.getZeroAttr(ty.getElementType());
   Location loc = op->getLoc();
   Value zero = b.create<arith::ConstantOp>(loc, zeroAttr);
   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)});
+
+  // 2, build the matmul op.
+  Operation* matmulOp;
+  if (lhsContractingDims[0] == 1 && rhsContractingDims[0] == 0) {
+    // nn format gemm
+    matmulOp = b.create<linalg::MatmulOp>(loc, ValueRange{newA, newB},
+                                          ValueRange{fillOp.getResult(0)});
+  } else {
+    // nt, tn, tt format gemm
+    SmallVector<AffineMap, 3> indexingMaps;
+
+    auto getMap = [&](int64_t rank, ArrayRef<int64_t> batchingDims,
+                      ArrayRef<int64_t> contractingDims, size_t extraDims) {
+      llvm::SmallVector<AffineExpr> indices(rank);
+      for (const auto& i : llvm::enumerate(batchingDims)) {
+        indices[i.value()] = b.getAffineDimExpr(i.index());
+      }
+      for (const auto& i : llvm::enumerate(contractingDims)) {
+        indices[i.value()] = b.getAffineDimExpr(i.index() + targetRank);
+      }
+      for (int i = 0; i < rank; ++i) {
+        if (!indices[i]) {
+          indices[i] = b.getAffineDimExpr(extraDims++);
+        }
+      }
+      indexingMaps.push_back(AffineMap::get(/*dimCount=*/totalLoopCount,
+                                            /*symbolCount=*/0, indices,
+                                            op->getContext()));
+    };
+    getMap(lhsRank, lhsBatchingDims, lhsContractingDims,
+           lhsBatchingDims.size());
+    getMap(rhsRank, rhsBatchingDims, rhsContractingDims,
+           rhsBatchingDims.size() + lhsExtraDims);
+
+    {
+      SmallVector<AffineExpr, 4> dimExprs;
+      dimExprs.reserve(targetRank);
+      for (unsigned i = 0; i < targetRank; ++i)
+        dimExprs.push_back(b.getAffineDimExpr(i));
+      indexingMaps.push_back(AffineMap::get(/*dimCount=*/totalLoopCount,
+                                            /*symbolCount=*/0, dimExprs,
+                                            op.getContext()));
+    }
+
+    matmulOp = b.create<linalg::GenericOp>(
+        loc, /*resultTensorTypes=*/TypeRange{outputType},
+        /*inputs=*/ValueRange{newA, newB},
+        /*outputBuffers=*/ValueRange{fillOp.getResult(0)}, indexingMaps,
+        mhlo::getParallelAndReductionIterators(
+            /*nLoops=*/totalLoopCount,
+            /*nReduction=*/numContracting),
+        [](OpBuilder& b, Location loc, ValueRange) {
+          ImplicitLocOpBuilder builder(loc, b);
+          linalg::MatmulOp::regionBuilder(builder, *b.getInsertionBlock(), {});
+        },
+        linalg::getPrunedAttributeList(op));
+  }
   matmulOp->setAttr(kDISCLinalgTransformName, b.getStringAttr(name));
   mapping.erase(C);
-  mapping.map(C, matmulOp.getResult(0));
+  mapping.map(C, matmulOp->getResult(0));
 
   return success();
 }
diff --git a/tao_compiler/mlir/disc/transforms/disc_transform_schedule.cc b/tao_compiler/mlir/disc/transforms/disc_transform_schedule.cc
index ea11a8fa9f6..07db691c052 100644
--- a/tao_compiler/mlir/disc/transforms/disc_transform_schedule.cc
+++ b/tao_compiler/mlir/disc/transforms/disc_transform_schedule.cc
@@ -274,9 +274,6 @@ LogicalResult aarch64GEMMDefaultScheduleFactory(PatternDescription& pd,
   }
   bool lhsTranspose = (lhsCntractingDims[0] == lhsTy.getRank() - 2);
   bool rhsTranspose = (rhsCntractingDims[0] == rhsTy.getRank() - 1);
-  if (lhsTranspose || rhsTranspose) {
-    return m->emitError() << "not support transposed gemm a.t.m.\n";
-  }
   int64_t M = lhsTranspose ? lhsTy.getShape()[lhsTy.getRank() - 1]
                            : lhsTy.getShape()[lhsTy.getRank() - 2];
   int64_t K = lhsTranspose ? lhsTy.getShape()[lhsTy.getRank() - 2]
@@ -374,8 +371,17 @@ LogicalResult aarch64GEMMDefaultScheduleFactory(PatternDescription& pd,
     }
     auto loopN0 = buildGetParentForOp(b, loc, padForInput, loopLevel);
     bool inputIsPadded = mIsPadded || kIsPadded;
-    buildCacheRead(b, loc, padForInput, loopN0, {1, 1}, {6, 1}, inputIsPadded,
-                   {0, 2, 3, 1});
+    SmallVector<int64_t> tileSizes;
+    SmallVector<int64_t> permutation;
+    if (lhsTranspose) {
+      tileSizes = {K0, M1};
+      permutation = {2, 0, 1, 3};
+    } else {
+      tileSizes = {M1, K0};
+      permutation = {0, 2, 3, 1};
+    }
+    buildCacheRead(b, loc, padForInput, loopN0, {1, 1}, tileSizes,
+                   inputIsPadded, permutation);
   }
 
   // Check if we need to pack the weight, one of the following conditions:
@@ -386,8 +392,17 @@ LogicalResult aarch64GEMMDefaultScheduleFactory(PatternDescription& pd,
   if (packWeight) {
     bool weightIsPadded = nIsPadded || kIsPadded;
     forEachThreadLoop = buildMatchOp(b, loc, variant, {"scf.foreach_thread"});
-    buildCacheRead(b, loc, padForWeight, forEachThreadLoop, {1, 1}, {1, 16},
-                   weightIsPadded, {2, 0, 1, 3});
+    SmallVector<int64_t> tileSizes;
+    SmallVector<int64_t> permutation;
+    if (rhsTranspose) {
+      tileSizes = {N1, K0};
+      permutation = {0, 2, 3, 1};
+    } else {
+      tileSizes = {K0, N1};
+      permutation = {2, 0, 1, 3};
+    }
+    buildCacheRead(b, loc, padForWeight, forEachThreadLoop, {1, 1}, tileSizes,
+                   weightIsPadded, permutation);
   }
 
   variant = buildRunCanonicalizer(b, loc, variant);
diff --git a/tao_compiler/mlir/disc/transforms/fusion_utils_transform_based.cc b/tao_compiler/mlir/disc/transforms/fusion_utils_transform_based.cc
index 18ce52804a6..01a0f83387f 100644
--- a/tao_compiler/mlir/disc/transforms/fusion_utils_transform_based.cc
+++ b/tao_compiler/mlir/disc/transforms/fusion_utils_transform_based.cc
@@ -35,11 +35,9 @@ bool isSupportedDot(Operation* op) {
       dimNumbers.getRhsBatchingDimensions().size() != 0)
     return false;
 
-  // TODO(wyzero): support dot_general with fused transpose
   auto lhsCntractingDims = dimNumbers.getLhsContractingDimensions();
   auto rhsCntractingDims = dimNumbers.getRhsContractingDimensions();
-  return (lhsCntractingDims.size() == 1 && lhsCntractingDims[0] == 1 &&
-          rhsCntractingDims.size() == 1 && rhsCntractingDims[0] == 0);
+  return (lhsCntractingDims.size() == 1 && rhsCntractingDims.size() == 1);
 }
 
 bool TransformBasedCpuFusionStrategy::isFusible(Operation* op) {


Test patch:
diff --git a/tao_compiler/mlir/disc/tests/disc-transform/data/default_schedule_matmul_nt_d_f32.mlir b/tao_compiler/mlir/disc/tests/disc-transform/data/default_schedule_matmul_nt_d_f32.mlir
new file mode 100644
index 00000000000..2532d6c9258
--- /dev/null
+++ b/tao_compiler/mlir/disc/tests/disc-transform/data/default_schedule_matmul_nt_d_f32.mlir
@@ -0,0 +1,9 @@
+module attributes {tf.versions = {bad_consumers = [], min_consumer = 0 : i32, producer = 0 : i32}} {
+  func.func @main(%arg0: tensor<?x?xf32>, %arg1: tensor<?x?xf32>) -> (tensor<?x?xf32>) attributes {tf.entry_function = {inputs = "{{INPUTS}}", outputs = "{{OUTPUTS}}", input_placements="{{INPUT_PLACEMENTS}}", output_placements="{{OUTPUT_PLACEMENTS}}"}} {
+    %graph = tf_executor.graph {
+      %0:2 = tf_executor.island wraps "tf.MatMul"(%arg0, %arg1) {transpose_a = false, transpose_b = true} : (tensor<?x?xf32>, tensor<?x?xf32>) -> (tensor<?x?xf32>)
+      tf_executor.fetch %0 : tensor<?x?xf32>
+    }
+    return %graph : tensor<?x?xf32>
+  }
+}
\ No newline at end of file
diff --git a/tao_compiler/mlir/disc/tests/disc-transform/data/default_schedule_matmul_tn_d_f32.mlir b/tao_compiler/mlir/disc/tests/disc-transform/data/default_schedule_matmul_tn_d_f32.mlir
new file mode 100644
index 00000000000..206c1f3d293
--- /dev/null
+++ b/tao_compiler/mlir/disc/tests/disc-transform/data/default_schedule_matmul_tn_d_f32.mlir
@@ -0,0 +1,9 @@
+module attributes {tf.versions = {bad_consumers = [], min_consumer = 0 : i32, producer = 0 : i32}} {
+  func.func @main(%arg0: tensor<?x?xf32>, %arg1: tensor<?x?xf32>) -> (tensor<?x?xf32>) attributes {tf.entry_function = {inputs = "{{INPUTS}}", outputs = "{{OUTPUTS}}", input_placements="{{INPUT_PLACEMENTS}}", output_placements="{{OUTPUT_PLACEMENTS}}"}} {
+    %graph = tf_executor.graph {
+      %0:2 = tf_executor.island wraps "tf.MatMul"(%arg0, %arg1) {transpose_a = true, transpose_b = false} : (tensor<?x?xf32>, tensor<?x?xf32>) -> (tensor<?x?xf32>)
+      tf_executor.fetch %0 : tensor<?x?xf32>
+    }
+    return %graph : tensor<?x?xf32>
+  }
+}
\ No newline at end of file
diff --git a/tao_compiler/mlir/disc/tests/disc-transform/data/default_schedule_matmul_tt_d_f32.mlir b/tao_compiler/mlir/disc/tests/disc-transform/data/default_schedule_matmul_tt_d_f32.mlir
new file mode 100644
index 00000000000..27086621569
--- /dev/null
+++ b/tao_compiler/mlir/disc/tests/disc-transform/data/default_schedule_matmul_tt_d_f32.mlir
@@ -0,0 +1,9 @@
+module attributes {tf.versions = {bad_consumers = [], min_consumer = 0 : i32, producer = 0 : i32}} {
+  func.func @main(%arg0: tensor<?x?xf32>, %arg1: tensor<?x?xf32>) -> (tensor<?x?xf32>) attributes {tf.entry_function = {inputs = "{{INPUTS}}", outputs = "{{OUTPUTS}}", input_placements="{{INPUT_PLACEMENTS}}", output_placements="{{OUTPUT_PLACEMENTS}}"}} {
+    %graph = tf_executor.graph {
+      %0:2 = tf_executor.island wraps "tf.MatMul"(%arg0, %arg1) {transpose_a = true, transpose_b = true} : (tensor<?x?xf32>, tensor<?x?xf32>) -> (tensor<?x?xf32>)
+      tf_executor.fetch %0 : tensor<?x?xf32>
+    }
+    return %graph : tensor<?x?xf32>
+  }
+}
\ No newline at end of file
diff --git a/tao_compiler/mlir/disc/tests/disc-transform/default_schedule_matmul.cc b/tao_compiler/mlir/disc/tests/disc-transform/default_schedule_matmul.cc
index c9606e82409..2a889a77b2b 100644
--- a/tao_compiler/mlir/disc/tests/disc-transform/default_schedule_matmul.cc
+++ b/tao_compiler/mlir/disc/tests/disc-transform/default_schedule_matmul.cc
@@ -95,6 +95,57 @@ TEST(DynamicShapeMatmul, F32_24x768x3072_Using_Default_Schedule) {
       /*profiling*/ true));
 }
 
+TEST(DynamicShapeMatmul_nt, F32_304x1024x512_Using_Default_Schedule) {
+  EnvSetting setting = {{"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 + "default_schedule_matmul_nt_d_f32.mlir",
+      /*backend_types*/ {BackendType::kAArch64},
+      /*num_inputs*/ 2,
+      /*num_outputs*/ 1,
+      /*input_descriptors*/ {"24x512xf32_X", "1024x512xf32_X"},
+      /*output_descriptors*/ {"f32_X"},
+      /*input_vals*/ {},
+      /*expected_output_vals*/ {},
+      /*profiling*/ true));
+}
+
+TEST(DynamicShapeMatmul_tn, F32_304x1024x512_Using_Default_Schedule) {
+  EnvSetting setting = {{"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 + "default_schedule_matmul_tn_d_f32.mlir",
+      /*backend_types*/ {BackendType::kAArch64},
+      /*num_inputs*/ 2,
+      /*num_outputs*/ 1,
+      /*input_descriptors*/ {"512x24xf32_X", "512x1024xf32_X"},
+      /*output_descriptors*/ {"f32_X"},
+      /*input_vals*/ {},
+      /*expected_output_vals*/ {},
+      /*profiling*/ true));
+}
+
+TEST(DynamicShapeMatmul_tt, F32_304x1024x512_Using_Default_Schedule) {
+  EnvSetting setting = {{"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 + "default_schedule_matmul_tt_d_f32.mlir",
+      /*backend_types*/ {BackendType::kAArch64},
+      /*num_inputs*/ 2,
+      /*num_outputs*/ 1,
+      /*input_descriptors*/ {"512x24xf32_X", "1024x512xf32_X"},
+      /*output_descriptors*/ {"f32_X"},
+      /*input_vals*/ {},
+      /*expected_output_vals*/ {},
+      /*profiling*/ true));
+}
+
 TEST(StaticShapeMatmul, F32_24x24x64_Using_Default_Schedule) {
   EnvSetting setting = {{"DISC_ENABLE_TRANSFORM_SCHEDULE", {"1", false}},
                         {"DISC_ENABLE_SHAPE_CONSTRAINT_IR", {"1", false}},
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 4b53cb65896..ff1cbaf84e7 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
@@ -30,4 +30,76 @@ func.func @packed_matmul_nn(%arg1: memref<?x512xf32, "cpu">, %arg2: memref<512x1
     "lmhlo.terminator"() : () -> ()
   }) {disc.device = "cpu", disc.fusion.name = "matmul_nn_kTransform_dot_general__1_1_0", disc.fusion_type = "kTransform"} : () -> ()
   return %arg3 : memref<?x1024xf32, "cpu">
+}
+
+// -----
+
+// CHECK: #map = affine_map<(d0, d1, d2) -> (d2, d0)>
+// CHECK: #map1 = affine_map<(d0, d1, d2) -> (d2, d1)>
+// CHECK: #map2 = affine_map<(d0, d1, d2) -> (d0, d1)>
+
+// CHECK-LABEL: @matmul_tn
+// CHECK-SAME: (%[[ARG0:.*]]: tensor<?x?xf32>, %[[ARG1:.*]]: tensor<?x?xf32>, %[[ARG2:.*]]: tensor<?x?xf32>)
+func.func @matmul_tn(%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 {disc.transform.name = "dot_general"} ins(%[[T0]] : f32) outs(%[[ARG2]] : tensor<?x?xf32>) -> tensor<?x?xf32>
+  // CHECK: %[[T2:.*]] = linalg.generic
+  // CHECK-SAME: indexing_maps = [#map, #map1, #map2]
+  // CHECK-SAME: iterator_types = ["parallel", "parallel", "reduction"]
+  // CHECK-SAME: ins(%[[ARG0]], %[[ARG1]] : tensor<?x?xf32>, tensor<?x?xf32>) outs(%[[T1]]
+  // CHECK-SAME: {disc.transform.name = "dot_general"}
+  // CHECK: return %[[T2]]
+  "lmhlo.fusion"() ({
+    "lmhlo.dot_general"(%arg1, %arg2, %arg3) {dot_dimension_numbers = #mhlo.dot<lhs_contracting_dimensions = [0], rhs_contracting_dimensions = [0]>} : (memref<?x?xf32, "cpu">, memref<?x?xf32, "cpu">, 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">
+}
+
+// -----
+
+// CHECK: #map = affine_map<(d0, d1, d2) -> (d0, d2)>
+// CHECK: #map1 = affine_map<(d0, d1, d2) -> (d1, d2)>
+// CHECK: #map2 = affine_map<(d0, d1, d2) -> (d0, d1)>
+
+// CHECK-LABEL: @matmul_nt
+// CHECK-SAME: (%[[ARG0:.*]]: tensor<?x?xf32>, %[[ARG1:.*]]: tensor<?x?xf32>, %[[ARG2:.*]]: tensor<?x?xf32>)
+func.func @matmul_nt(%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 {disc.transform.name = "dot_general"} ins(%[[T0]] : f32) outs(%[[ARG2]] : tensor<?x?xf32>) -> tensor<?x?xf32>
+  // CHECK: %[[T2:.*]] = linalg.generic
+  // CHECK-SAME: indexing_maps = [#map, #map1, #map2]
+  // CHECK-SAME: iterator_types = ["parallel", "parallel", "reduction"]
+  // CHECK-SAME: ins(%[[ARG0]], %[[ARG1]] : tensor<?x?xf32>, tensor<?x?xf32>) outs(%[[T1]]
+  // CHECK-SAME: {disc.transform.name = "dot_general"}
+  // CHECK: return %[[T2]]
+  "lmhlo.fusion"() ({
+    "lmhlo.dot_general"(%arg1, %arg2, %arg3) {dot_dimension_numbers = #mhlo.dot<lhs_contracting_dimensions = [1], rhs_contracting_dimensions = [1]>} : (memref<?x?xf32, "cpu">, memref<?x?xf32, "cpu">, 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">
+}
+
+// -----
+
+// CHECK: #map = affine_map<(d0, d1, d2) -> (d2, d0)>
+// CHECK: #map1 = affine_map<(d0, d1, d2) -> (d1, d2)>
+// CHECK: #map2 = affine_map<(d0, d1, d2) -> (d0, d1)>
+
+// CHECK-LABEL: @matmul_tt
+// CHECK-SAME: (%[[ARG0:.*]]: tensor<?x?xf32>, %[[ARG1:.*]]: tensor<?x?xf32>, %[[ARG2:.*]]: tensor<?x?xf32>)
+func.func @matmul_tt(%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 {disc.transform.name = "dot_general"} ins(%[[T0]] : f32) outs(%[[ARG2]] : tensor<?x?xf32>) -> tensor<?x?xf32>
+  // CHECK: %[[T2:.*]] = linalg.generic
+  // CHECK-SAME: indexing_maps = [#map, #map1, #map2]
+  // CHECK-SAME: iterator_types = ["parallel", "parallel", "reduction"]
+  // CHECK-SAME: ins(%[[ARG0]], %[[ARG1]] : tensor<?x?xf32>, tensor<?x?xf32>) outs(%[[T1]]
+  // CHECK-SAME: {disc.transform.name = "dot_general"}
+  // CHECK: return %[[T2]]
+  "lmhlo.fusion"() ({
+    "lmhlo.dot_general"(%arg1, %arg2, %arg3) {dot_dimension_numbers = #mhlo.dot<lhs_contracting_dimensions = [0], rhs_contracting_dimensions = [1]>} : (memref<?x?xf32, "cpu">, memref<?x?xf32, "cpu">, 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