[Mlir-commits] [mlir] [mlir][linalg] fix specialization of transposed matmul variants (PR #181387)
llvmlistbot at llvm.org
llvmlistbot at llvm.org
Fri Feb 13 09:04:29 PST 2026
llvmbot wrote:
<!--LLVM PR SUMMARY COMMENT-->
@llvm/pr-subscribers-mlir-linalg
Author: None (ziereis)
<details>
<summary>Changes</summary>
Currently all generics that represent `linalg.matmul_transpose_a/b` variants will get specialized to a regular `linalg.matmul` without any special indexing maps. This patch fixes this behavior to correctly emit `linalg.matmul` with the correct indexing maps for the transposed variants.
---
Full diff: https://github.com/llvm/llvm-project/pull/181387.diff
2 Files Affected:
- (modified) mlir/lib/Dialect/Linalg/Transforms/Specialize.cpp (+16-2)
- (modified) mlir/test/Dialect/Linalg/specialize-generic-ops.mlir (+171)
``````````diff
diff --git a/mlir/lib/Dialect/Linalg/Transforms/Specialize.cpp b/mlir/lib/Dialect/Linalg/Transforms/Specialize.cpp
index a71f84dee3bb0..512722196aecc 100644
--- a/mlir/lib/Dialect/Linalg/Transforms/Specialize.cpp
+++ b/mlir/lib/Dialect/Linalg/Transforms/Specialize.cpp
@@ -147,10 +147,10 @@ static LinalgOp replaceWithMatmulVariant(RewriterBase &rewriter, GenericOp op,
castAttrVec = {rewriter.getNamedAttr(
"cast", TypeFnAttr::get(rewriter.getContext(), *castTy))};
- LinalgOp namedOp = rewriter.replaceOpWithNewOp<NamedOpTy>(
+ auto namedOp = rewriter.replaceOpWithNewOp<NamedOpTy>(
op, ValueRange{op.getDpsInputs()[0], op.getDpsInputs()[1]},
ValueRange{op.getDpsInits()[0]}, castAttrVec);
- return namedOp;
+ return cast<LinalgOp>(namedOp.getOperation());
}
// Returns the cast type to use for a matmul-like named op. If the generic
@@ -299,8 +299,22 @@ static FailureOr<LinalgOp> specializeLinalgContractions(RewriterBase &rewriter,
/// Codegen the different matmul variants.
if (numOfBatchDims) {
+ if (a == IndexMatchResult::Transposed)
+ return replaceWithMatmulVariant<BatchMatmulTransposeAOp>(rewriter,
+ genericOp,
+ castTy);
+ if (b == IndexMatchResult::Transposed)
+ return replaceWithMatmulVariant<BatchMatmulTransposeBOp>(rewriter,
+ genericOp,
+ castTy);
return replaceWithMatmulVariant<BatchMatmulOp>(rewriter, genericOp, castTy);
}
+ if (a == IndexMatchResult::Transposed)
+ return replaceWithMatmulVariant<MatmulTransposeAOp>(rewriter, genericOp,
+ castTy);
+ if (b == IndexMatchResult::Transposed)
+ return replaceWithMatmulVariant<MatmulTransposeBOp>(rewriter, genericOp,
+ castTy);
return replaceWithMatmulVariant<MatmulOp>(rewriter, genericOp, castTy);
}
diff --git a/mlir/test/Dialect/Linalg/specialize-generic-ops.mlir b/mlir/test/Dialect/Linalg/specialize-generic-ops.mlir
index 6acf1ca0d4e30..5b2609da662c4 100644
--- a/mlir/test/Dialect/Linalg/specialize-generic-ops.mlir
+++ b/mlir/test/Dialect/Linalg/specialize-generic-ops.mlir
@@ -314,6 +314,73 @@ func.func @negative_op_multi_reduction(%A: tensor<10x20x30xf32>,
// -----
+// Both A and B transposed: no named op variant exists for this combination.
+#mapABt0 = affine_map<(d0, d1, d2) -> (d2, d0)>
+#mapABt1 = affine_map<(d0, d1, d2) -> (d1, d2)>
+#mapABt2 = affine_map<(d0, d1, d2) -> (d0, d1)>
+func.func @negative_matmul_transpose_a_and_b(%A: tensor<?x?xf32>, %B: tensor<?x?xf32>,
+ %Out: tensor<?x?xf32>) -> tensor<?x?xf32> {
+ %0 = linalg.generic
+ {indexing_maps = [#mapABt0, #mapABt1, #mapABt2], iterator_types = ["parallel", "parallel", "reduction"]}
+ ins(%A, %B : tensor<?x?xf32>, tensor<?x?xf32>) outs(%Out : tensor<?x?xf32>) {
+ ^bb0(%in: f32, %in_0: f32, %out: f32):
+ %1 = arith.mulf %in, %in_0 : f32
+ %2 = arith.addf %out, %1 : f32
+ linalg.yield %2 : f32
+ } -> tensor<?x?xf32>
+ return %0 : tensor<?x?xf32>
+}
+
+// CHECK-LABEL: negative_matmul_transpose_a_and_b
+// CHECK: linalg.generic
+
+// -----
+
+// Output transposed: C is accessed as (n, m) instead of (m, n).
+#mapCt0 = affine_map<(d0, d1, d2) -> (d0, d2)>
+#mapCt1 = affine_map<(d0, d1, d2) -> (d2, d1)>
+#mapCt2 = affine_map<(d0, d1, d2) -> (d1, d0)>
+func.func @negative_matmul_transposed_output(%A: tensor<?x?xf32>, %B: tensor<?x?xf32>,
+ %Out: tensor<?x?xf32>) -> tensor<?x?xf32> {
+ %0 = linalg.generic
+ {indexing_maps = [#mapCt0, #mapCt1, #mapCt2], iterator_types = ["parallel", "parallel", "reduction"]}
+ ins(%A, %B : tensor<?x?xf32>, tensor<?x?xf32>) outs(%Out : tensor<?x?xf32>) {
+ ^bb0(%in: f32, %in_0: f32, %out: f32):
+ %1 = arith.mulf %in, %in_0 : f32
+ %2 = arith.addf %out, %1 : f32
+ linalg.yield %2 : f32
+ } -> tensor<?x?xf32>
+ return %0 : tensor<?x?xf32>
+}
+
+// CHECK-LABEL: negative_matmul_transposed_output
+// CHECK: linalg.generic
+
+// -----
+
+// Batch dim not in identity position: batch dim d0 appears at result
+// position 1 in A's map instead of position 0.
+#mapBni0 = affine_map<(d0, d1, d2, d3) -> (d1, d0, d3)>
+#mapBni1 = affine_map<(d0, d1, d2, d3) -> (d0, d3, d2)>
+#mapBni2 = affine_map<(d0, d1, d2, d3) -> (d0, d1, d2)>
+func.func @negative_batch_matmul_non_identity_batch(%A: tensor<4x2x8xf32>, %B: tensor<2x8x16xf32>,
+ %Out: tensor<2x4x16xf32>) -> tensor<2x4x16xf32> {
+ %0 = linalg.generic
+ {indexing_maps = [#mapBni0, #mapBni1, #mapBni2], iterator_types = ["parallel", "parallel", "parallel", "reduction"]}
+ ins(%A, %B : tensor<4x2x8xf32>, tensor<2x8x16xf32>) outs(%Out : tensor<2x4x16xf32>) {
+ ^bb0(%in: f32, %in_0: f32, %out: f32):
+ %1 = arith.mulf %in, %in_0 : f32
+ %2 = arith.addf %out, %1 : f32
+ linalg.yield %2 : f32
+ } -> tensor<2x4x16xf32>
+ return %0 : tensor<2x4x16xf32>
+}
+
+// CHECK-LABEL: negative_batch_matmul_non_identity_batch
+// CHECK: linalg.generic
+
+// -----
+
// TODO: matvec
#map = affine_map<(d0, d1) -> (d0, d1)>
#map1 = affine_map<(d0, d1) -> (d1)>
@@ -331,3 +398,107 @@ func.func @op_matvec(%A: tensor<?x?xf32>, %B: tensor<?xf32>, %Out: tensor<?xf32>
}
// CHECK-LABEL: op_matvec
// CHECK: linalg.generic
+
+// -----
+
+// Matmul transpose A: A is accessed as (k, m) instead of (m, k)
+#mapTA0 = affine_map<(d0, d1, d2) -> (d2, d0)>
+#mapTA1 = affine_map<(d0, d1, d2) -> (d2, d1)>
+#mapTA2 = affine_map<(d0, d1, d2) -> (d0, d1)>
+func.func @op_matmul_transpose_a(%A: tensor<?x?xf32>, %B: tensor<?x?xf32>, %Out: tensor<?x?xf32>) -> tensor<?x?xf32> {
+ %0 = linalg.generic
+ {indexing_maps = [#mapTA0, #mapTA1, #mapTA2], iterator_types = ["parallel", "parallel", "reduction"]}
+ ins(%A, %B : tensor<?x?xf32>, tensor<?x?xf32>) outs(%Out : tensor<?x?xf32>) {
+ ^bb0(%in: f32, %in_0: f32, %out: f32):
+ %1 = arith.mulf %in, %in_0 : f32
+ %2 = arith.addf %out, %1 : f32
+ linalg.yield %2 : f32
+ } -> tensor<?x?xf32>
+ return %0 : tensor<?x?xf32>
+}
+
+// CHECK-DAG: #[[$MAP_TA_A:.+]] = affine_map<(d0, d1, d2) -> (d2, d0)>
+// CHECK-DAG: #[[$MAP_TA_B:.+]] = affine_map<(d0, d1, d2) -> (d2, d1)>
+// CHECK-DAG: #[[$MAP_TA_C:.+]] = affine_map<(d0, d1, d2) -> (d0, d1)>
+// CHECK-LABEL: op_matmul_transpose_a
+// CHECK-SAME: %[[A:.+]]: tensor<?x?xf32>, %[[B:.+]]: tensor<?x?xf32>, %[[Out:.+]]: tensor<?x?xf32>) -> tensor<?x?xf32>
+// CHECK-NOT: linalg.generic
+// CHECK: linalg.matmul indexing_maps = [#[[$MAP_TA_A]], #[[$MAP_TA_B]], #[[$MAP_TA_C]]] ins(%[[A]], %[[B]] : tensor<?x?xf32>, tensor<?x?xf32>) outs(%[[Out]] : tensor<?x?xf32>) -> tensor<?x?xf32>
+
+// -----
+
+// Matmul transpose B: B is accessed as (n, k) instead of (k, n)
+#mapTB0 = affine_map<(d0, d1, d2) -> (d0, d2)>
+#mapTB1 = affine_map<(d0, d1, d2) -> (d1, d2)>
+#mapTB2 = affine_map<(d0, d1, d2) -> (d0, d1)>
+func.func @op_matmul_transpose_b(%A: tensor<?x?xf32>, %B: tensor<?x?xf32>, %Out: tensor<?x?xf32>) -> tensor<?x?xf32> {
+ %0 = linalg.generic
+ {indexing_maps = [#mapTB0, #mapTB1, #mapTB2], iterator_types = ["parallel", "parallel", "reduction"]}
+ ins(%A, %B : tensor<?x?xf32>, tensor<?x?xf32>) outs(%Out : tensor<?x?xf32>) {
+ ^bb0(%in: f32, %in_0: f32, %out: f32):
+ %1 = arith.mulf %in, %in_0 : f32
+ %2 = arith.addf %out, %1 : f32
+ linalg.yield %2 : f32
+ } -> tensor<?x?xf32>
+ return %0 : tensor<?x?xf32>
+}
+
+// CHECK-DAG: #[[$MAP_TB_A:.+]] = affine_map<(d0, d1, d2) -> (d0, d2)>
+// CHECK-DAG: #[[$MAP_TB_B:.+]] = affine_map<(d0, d1, d2) -> (d1, d2)>
+// CHECK-DAG: #[[$MAP_TB_C:.+]] = affine_map<(d0, d1, d2) -> (d0, d1)>
+// CHECK-LABEL: op_matmul_transpose_b
+// CHECK-SAME: %[[A:.+]]: tensor<?x?xf32>, %[[B:.+]]: tensor<?x?xf32>, %[[Out:.+]]: tensor<?x?xf32>) -> tensor<?x?xf32>
+// CHECK-NOT: linalg.generic
+// CHECK: linalg.matmul indexing_maps = [#[[$MAP_TB_A]], #[[$MAP_TB_B]], #[[$MAP_TB_C]]] ins(%[[A]], %[[B]] : tensor<?x?xf32>, tensor<?x?xf32>) outs(%[[Out]] : tensor<?x?xf32>) -> tensor<?x?xf32>
+
+// -----
+
+// Batch matmul transpose A: A is accessed as (b, k, m) instead of (b, m, k)
+#mapBTA0 = affine_map<(d0, d1, d2, d3) -> (d0, d3, d1)>
+#mapBTA1 = affine_map<(d0, d1, d2, d3) -> (d0, d3, d2)>
+#mapBTA2 = affine_map<(d0, d1, d2, d3) -> (d0, d1, d2)>
+func.func @op_batch_matmul_transpose_a(%A: tensor<2x8x4xf32>, %B: tensor<2x8x16xf32>, %Out: tensor<2x4x16xf32>) -> tensor<2x4x16xf32> {
+ %0 = linalg.generic
+ {indexing_maps = [#mapBTA0, #mapBTA1, #mapBTA2], iterator_types = ["parallel", "parallel", "parallel", "reduction"]}
+ ins(%A, %B : tensor<2x8x4xf32>, tensor<2x8x16xf32>) outs(%Out : tensor<2x4x16xf32>) {
+ ^bb0(%in: f32, %in_0: f32, %out: f32):
+ %1 = arith.mulf %in, %in_0 : f32
+ %2 = arith.addf %out, %1 : f32
+ linalg.yield %2 : f32
+ } -> tensor<2x4x16xf32>
+ return %0 : tensor<2x4x16xf32>
+}
+
+// CHECK-DAG: #[[$MAP_BTA_A:.+]] = affine_map<(d0, d1, d2, d3) -> (d0, d3, d1)>
+// CHECK-DAG: #[[$MAP_BTA_B:.+]] = affine_map<(d0, d1, d2, d3) -> (d0, d3, d2)>
+// CHECK-DAG: #[[$MAP_BTA_C:.+]] = affine_map<(d0, d1, d2, d3) -> (d0, d1, d2)>
+// CHECK-LABEL: op_batch_matmul_transpose_a
+// CHECK-SAME: %[[A:.+]]: tensor<2x8x4xf32>, %[[B:.+]]: tensor<2x8x16xf32>, %[[Out:.+]]: tensor<2x4x16xf32>) -> tensor<2x4x16xf32>
+// CHECK-NOT: linalg.generic
+// CHECK: linalg.batch_matmul indexing_maps = [#[[$MAP_BTA_A]], #[[$MAP_BTA_B]], #[[$MAP_BTA_C]]] ins(%[[A]], %[[B]] : tensor<2x8x4xf32>, tensor<2x8x16xf32>) outs(%[[Out]] : tensor<2x4x16xf32>) -> tensor<2x4x16xf32>
+
+// -----
+
+// Batch matmul transpose B: B is accessed as (b, n, k) instead of (b, k, n)
+#mapBTB0 = affine_map<(d0, d1, d2, d3) -> (d0, d1, d3)>
+#mapBTB1 = affine_map<(d0, d1, d2, d3) -> (d0, d2, d3)>
+#mapBTB2 = affine_map<(d0, d1, d2, d3) -> (d0, d1, d2)>
+func.func @op_batch_matmul_transpose_b(%A: tensor<2x4x8xf32>, %B: tensor<2x16x8xf32>, %Out: tensor<2x4x16xf32>) -> tensor<2x4x16xf32> {
+ %0 = linalg.generic
+ {indexing_maps = [#mapBTB0, #mapBTB1, #mapBTB2], iterator_types = ["parallel", "parallel", "parallel", "reduction"]}
+ ins(%A, %B : tensor<2x4x8xf32>, tensor<2x16x8xf32>) outs(%Out : tensor<2x4x16xf32>) {
+ ^bb0(%in: f32, %in_0: f32, %out: f32):
+ %1 = arith.mulf %in, %in_0 : f32
+ %2 = arith.addf %out, %1 : f32
+ linalg.yield %2 : f32
+ } -> tensor<2x4x16xf32>
+ return %0 : tensor<2x4x16xf32>
+}
+
+// CHECK-DAG: #[[$MAP_BTB_A:.+]] = affine_map<(d0, d1, d2, d3) -> (d0, d1, d3)>
+// CHECK-DAG: #[[$MAP_BTB_B:.+]] = affine_map<(d0, d1, d2, d3) -> (d0, d2, d3)>
+// CHECK-DAG: #[[$MAP_BTB_C:.+]] = affine_map<(d0, d1, d2, d3) -> (d0, d1, d2)>
+// CHECK-LABEL: op_batch_matmul_transpose_b
+// CHECK-SAME: %[[A:.+]]: tensor<2x4x8xf32>, %[[B:.+]]: tensor<2x16x8xf32>, %[[Out:.+]]: tensor<2x4x16xf32>) -> tensor<2x4x16xf32>
+// CHECK-NOT: linalg.generic
+// CHECK: linalg.batch_matmul indexing_maps = [#[[$MAP_BTB_A]], #[[$MAP_BTB_B]], #[[$MAP_BTB_C]]] ins(%[[A]], %[[B]] : tensor<2x4x8xf32>, tensor<2x16x8xf32>) outs(%[[Out]] : tensor<2x4x16xf32>) -> tensor<2x4x16xf32>
``````````
</details>
https://github.com/llvm/llvm-project/pull/181387
More information about the Mlir-commits
mailing list