[Mlir-commits] [mlir] [mlir][linalg] Skip rank reduction for unsigned-cast contractions (PR #225993)
llvmlistbot at llvm.org
llvmlistbot at llvm.org
Thu Oct 1 21:56:20 PDT 2026
https://github.com/speryt updated https://github.com/llvm/llvm-project/pull/225993
>From c7c33f58a3d576eb6cb591a53a6957ffdaa776d0 Mon Sep 17 00:00:00 2001
From: s-peryt <1652033+s-peryt at users.noreply.github.com>
Date: Wed, 23 Sep 2026 18:27:21 -0700
Subject: [PATCH 1/3] [mlir][linalg] Don't rank-reduce contractions whose cast
can't be preserved
Rank reduction could turn a matmul or batch_matmul with
cast_unsigned into vecmat/matvec (or their batch versions). Those ops
have no cast attribute and always sign-extend, so the unsigned
semantics were silently lost.
Only matmul can carry the cast, so skip the rewrite for other targets
when the cast isn't signed.
Fixes #225253
Used Claude to write test.
---
.../Linalg/Transforms/DropUnitDims.cpp | 8 ++
.../Linalg/rank-reduce-contraction-ops.mlir | 93 +++++++++++++++++++
2 files changed, 101 insertions(+)
diff --git a/mlir/lib/Dialect/Linalg/Transforms/DropUnitDims.cpp b/mlir/lib/Dialect/Linalg/Transforms/DropUnitDims.cpp
index dac2891ce65e565..95fbea6197e1db6 100644
--- a/mlir/lib/Dialect/Linalg/Transforms/DropUnitDims.cpp
+++ b/mlir/lib/Dialect/Linalg/Transforms/DropUnitDims.cpp
@@ -963,6 +963,14 @@ struct RankReduceContractionOps : OpRewritePattern<FromOpTy> {
return rewriter.notifyMatchFailure(
contractionOp, "ops with user-defined maps are not supported");
}
+ // linalg.matmul is the only target op with a `cast` attribute.
+ if constexpr (!std::is_same_v<ToOpTy, MatmulOp>) {
+ auto castAttr = dyn_cast_if_present<TypeFnAttr>(
+ contractionOp->getInherentAttr("cast").value_or(Attribute()));
+ if (castAttr && castAttr.getValue() != TypeFn::cast_signed)
+ return rewriter.notifyMatchFailure(
+ contractionOp, "target op cannot represent a non-signed cast");
+ }
auto loc = contractionOp.getLoc();
auto inputs = contractionOp.getDpsInputs();
diff --git a/mlir/test/Dialect/Linalg/rank-reduce-contraction-ops.mlir b/mlir/test/Dialect/Linalg/rank-reduce-contraction-ops.mlir
index b7c095ca96b2d39..cf7ef21f6122528 100644
--- a/mlir/test/Dialect/Linalg/rank-reduce-contraction-ops.mlir
+++ b/mlir/test/Dialect/Linalg/rank-reduce-contraction-ops.mlir
@@ -175,6 +175,99 @@ func.func @batch_matmul_to_vecmat(%arg0: memref<1x1x?xf32>, %arg1: memref<1x?x?x
// -----
+// linalg.vecmat has no cast attribute (it always uses signed casts), so a
+// matmul with an unsigned cast must not be rank-reduced into it.
+func.func @negative_matmul_to_vecmat_unsigned_cast(
+ %arg0: tensor<1x8xi16>, %arg1: tensor<8x6xi64>,
+ %arg2: tensor<1x6xi32>) -> tensor<1x6xi32> {
+ // CHECK-LABEL: @negative_matmul_to_vecmat_unsigned_cast
+ // CHECK-NOT: collapse_shape
+ // CHECK: linalg.matmul {cast = #linalg.type_fn<cast_unsigned>}
+ // CHECK-NOT: linalg.vecmat
+ %0 = linalg.matmul {cast = #linalg.type_fn<cast_unsigned>}
+ ins(%arg0, %arg1 : tensor<1x8xi16>, tensor<8x6xi64>)
+ outs(%arg2 : tensor<1x6xi32>) -> tensor<1x6xi32>
+ return %0 : tensor<1x6xi32>
+}
+
+// -----
+
+func.func @negative_matmul_to_matvec_unsigned_cast(
+ %arg0: tensor<6x8xi16>, %arg1: tensor<8x1xi64>,
+ %arg2: tensor<6x1xi32>) -> tensor<6x1xi32> {
+ // CHECK-LABEL: @negative_matmul_to_matvec_unsigned_cast
+ // CHECK-NOT: collapse_shape
+ // CHECK: linalg.matmul {cast = #linalg.type_fn<cast_unsigned>}
+ // CHECK-NOT: linalg.matvec
+ %0 = linalg.matmul {cast = #linalg.type_fn<cast_unsigned>}
+ ins(%arg0, %arg1 : tensor<6x8xi16>, tensor<8x1xi64>)
+ outs(%arg2 : tensor<6x1xi32>) -> tensor<6x1xi32>
+ return %0 : tensor<6x1xi32>
+}
+
+// -----
+
+func.func @negative_batch_matmul_to_batch_vecmat_unsigned_cast(
+ %arg0: tensor<2x1x8xi16>, %arg1: tensor<2x8x6xi64>,
+ %arg2: tensor<2x1x6xi32>) -> tensor<2x1x6xi32> {
+ // CHECK-LABEL: @negative_batch_matmul_to_batch_vecmat_unsigned_cast
+ // CHECK-NOT: collapse_shape
+ // CHECK: linalg.batch_matmul {cast = #linalg.type_fn<cast_unsigned>}
+ // CHECK-NOT: linalg.batch_vecmat
+ %0 = linalg.batch_matmul {cast = #linalg.type_fn<cast_unsigned>}
+ ins(%arg0, %arg1 : tensor<2x1x8xi16>, tensor<2x8x6xi64>)
+ outs(%arg2 : tensor<2x1x6xi32>) -> tensor<2x1x6xi32>
+ return %0 : tensor<2x1x6xi32>
+}
+
+// -----
+
+func.func @negative_batch_matmul_to_batch_matvec_unsigned_cast(
+ %arg0: tensor<2x6x8xi16>, %arg1: tensor<2x8x1xi64>,
+ %arg2: tensor<2x6x1xi32>) -> tensor<2x6x1xi32> {
+ // CHECK-LABEL: @negative_batch_matmul_to_batch_matvec_unsigned_cast
+ // CHECK-NOT: collapse_shape
+ // CHECK: linalg.batch_matmul {cast = #linalg.type_fn<cast_unsigned>}
+ // CHECK-NOT: linalg.batch_matvec
+ %0 = linalg.batch_matmul {cast = #linalg.type_fn<cast_unsigned>}
+ ins(%arg0, %arg1 : tensor<2x6x8xi16>, tensor<2x8x1xi64>)
+ outs(%arg2 : tensor<2x6x1xi32>) -> tensor<2x6x1xi32>
+ return %0 : tensor<2x6x1xi32>
+}
+
+// -----
+
+// The unit batch dim can still be dropped since linalg.matmul preserves the
+// cast, but the rank reduction must stop there.
+func.func @singleton_batch_matmul_unsigned_cast_to_matmul_only(
+ %arg0: tensor<1x1x8xi16>, %arg1: tensor<1x8x6xi64>,
+ %arg2: tensor<1x1x6xi32>) -> tensor<1x1x6xi32> {
+ // CHECK-LABEL: @singleton_batch_matmul_unsigned_cast_to_matmul_only
+ // CHECK: linalg.matmul {cast = #linalg.type_fn<cast_unsigned>}
+ // CHECK-NOT: linalg.vecmat
+ // CHECK-NOT: linalg.batch_vecmat
+ %0 = linalg.batch_matmul {cast = #linalg.type_fn<cast_unsigned>}
+ ins(%arg0, %arg1 : tensor<1x1x8xi16>, tensor<1x8x6xi64>)
+ outs(%arg2 : tensor<1x1x6xi32>) -> tensor<1x1x6xi32>
+ return %0 : tensor<1x1x6xi32>
+}
+
+// -----
+
+// An explicit signed cast matches the target's semantics and is rank-reduced.
+func.func @matmul_to_vecmat_signed_cast(
+ %arg0: tensor<1x8xi16>, %arg1: tensor<8x6xi64>,
+ %arg2: tensor<1x6xi32>) -> tensor<1x6xi32> {
+ // CHECK-LABEL: @matmul_to_vecmat_signed_cast
+ // CHECK: linalg.vecmat
+ %0 = linalg.matmul {cast = #linalg.type_fn<cast_signed>}
+ ins(%arg0, %arg1 : tensor<1x8xi16>, tensor<8x6xi64>)
+ outs(%arg2 : tensor<1x6xi32>) -> tensor<1x6xi32>
+ return %0 : tensor<1x6xi32>
+}
+
+// -----
+
func.func @matvec_to_dot(%arg0: memref<1x?xf32>, %arg1: memref<?xf32>, %arg2: memref<1xf32>) {
// CHECK-LABEL: @matvec_to_dot
// CHECK-SAME: %[[LHS:[a-zA-Z0-9]+]]: memref<1x?xf32>
>From 924b5631850fcfa90b217a000a4dac94e4114755 Mon Sep 17 00:00:00 2001
From: s-peryt <1652033+s-peryt at users.noreply.github.com>
Date: Sun, 27 Sep 2026 09:12:03 -0700
Subject: [PATCH 2/3] Updated test comment
---
mlir/test/Dialect/Linalg/rank-reduce-contraction-ops.mlir | 4 ++--
1 file changed, 2 insertions(+), 2 deletions(-)
diff --git a/mlir/test/Dialect/Linalg/rank-reduce-contraction-ops.mlir b/mlir/test/Dialect/Linalg/rank-reduce-contraction-ops.mlir
index cf7ef21f6122528..ff5f63f291625c6 100644
--- a/mlir/test/Dialect/Linalg/rank-reduce-contraction-ops.mlir
+++ b/mlir/test/Dialect/Linalg/rank-reduce-contraction-ops.mlir
@@ -175,8 +175,8 @@ func.func @batch_matmul_to_vecmat(%arg0: memref<1x1x?xf32>, %arg1: memref<1x?x?x
// -----
-// linalg.vecmat has no cast attribute (it always uses signed casts), so a
-// matmul with an unsigned cast must not be rank-reduced into it.
+// linalg.vecmat currently has no cast attribute (it always uses signed casts),
+// so a matmul with an unsigned cast must not be rank-reduced into it.
func.func @negative_matmul_to_vecmat_unsigned_cast(
%arg0: tensor<1x8xi16>, %arg1: tensor<8x6xi64>,
%arg2: tensor<1x6xi32>) -> tensor<1x6xi32> {
>From 04ea3f53c6b661484d05fc59eaa913821118ed56 Mon Sep 17 00:00:00 2001
From: s-peryt <1652033+s-peryt at users.noreply.github.com>
Date: Thu, 1 Oct 2026 21:55:49 -0700
Subject: [PATCH 3/3] Added fixes based on PR review
---
.../lib/Dialect/Linalg/Transforms/DropUnitDims.cpp | 8 ++++----
.../Linalg/rank-reduce-contraction-ops.mlir | 14 ++++++++++++++
2 files changed, 18 insertions(+), 4 deletions(-)
diff --git a/mlir/lib/Dialect/Linalg/Transforms/DropUnitDims.cpp b/mlir/lib/Dialect/Linalg/Transforms/DropUnitDims.cpp
index 95fbea6197e1db6..75e6fc9491129ee 100644
--- a/mlir/lib/Dialect/Linalg/Transforms/DropUnitDims.cpp
+++ b/mlir/lib/Dialect/Linalg/Transforms/DropUnitDims.cpp
@@ -964,10 +964,10 @@ struct RankReduceContractionOps : OpRewritePattern<FromOpTy> {
contractionOp, "ops with user-defined maps are not supported");
}
// linalg.matmul is the only target op with a `cast` attribute.
- if constexpr (!std::is_same_v<ToOpTy, MatmulOp>) {
- auto castAttr = dyn_cast_if_present<TypeFnAttr>(
- contractionOp->getInherentAttr("cast").value_or(Attribute()));
- if (castAttr && castAttr.getValue() != TypeFn::cast_signed)
+ if constexpr ((std::is_same_v<FromOpTy, MatmulOp> ||
+ std::is_same_v<FromOpTy, BatchMatmulOp>) &&
+ !std::is_same_v<ToOpTy, MatmulOp>) {
+ if (contractionOp.getCast() != TypeFn::cast_signed)
return rewriter.notifyMatchFailure(
contractionOp, "target op cannot represent a non-signed cast");
}
diff --git a/mlir/test/Dialect/Linalg/rank-reduce-contraction-ops.mlir b/mlir/test/Dialect/Linalg/rank-reduce-contraction-ops.mlir
index ff5f63f291625c6..4f832e1e1dae62c 100644
--- a/mlir/test/Dialect/Linalg/rank-reduce-contraction-ops.mlir
+++ b/mlir/test/Dialect/Linalg/rank-reduce-contraction-ops.mlir
@@ -268,6 +268,20 @@ func.func @matmul_to_vecmat_signed_cast(
// -----
+func.func @batch_matmul_to_batch_vecmat_signed_cast(
+ %arg0: tensor<2x1x8xi16>, %arg1: tensor<2x8x6xi64>,
+ %arg2: tensor<2x1x6xi32>) -> tensor<2x1x6xi32> {
+ // CHECK-LABEL: @batch_matmul_to_batch_vecmat_signed_cast
+ // CHECK: linalg.batch_vecmat
+ // CHECK-NOT: linalg.batch_matmul
+ %0 = linalg.batch_matmul {cast = #linalg.type_fn<cast_signed>}
+ ins(%arg0, %arg1 : tensor<2x1x8xi16>, tensor<2x8x6xi64>)
+ outs(%arg2 : tensor<2x1x6xi32>) -> tensor<2x1x6xi32>
+ return %0 : tensor<2x1x6xi32>
+}
+
+// -----
+
func.func @matvec_to_dot(%arg0: memref<1x?xf32>, %arg1: memref<?xf32>, %arg2: memref<1xf32>) {
// CHECK-LABEL: @matvec_to_dot
// CHECK-SAME: %[[LHS:[a-zA-Z0-9]+]]: memref<1x?xf32>
More information about the Mlir-commits
mailing list