[Mlir-commits] [mlir] [mlir][linalg] Preserve unsigned integer widening in named contraction vectorization (PR #216283)
llvmlistbot at llvm.org
llvmlistbot at llvm.org
Fri Aug 14 02:12:56 PDT 2026
https://github.com/Pecco-314 created https://github.com/llvm/llvm-project/pull/216283
This patch fixes a correctness issue in named contraction vectorization of Linalg ops with unsigned integer widening.
When `create_named_contraction` is enabled, an op such as `linalg.matmul {cast = #linalg.type_fn<cast_unsigned>}` was lowered directly to a mixed-width `vector.contract` without materializing the implied `arith.extui`. Since `vector.contract` promotes mixed-width integer operands using sign extension, an `i8 x i8 -> i32` matmul could interpret unsigned inputs as signed values and produce incorrect results.
Materialize `arith.extui` for narrower integer inputs before constructing the `vector.contract`, preserving the required zero-extension semantics.
>From 764e9a52ab63b121401faabb246730da77235a37 Mon Sep 17 00:00:00 2001
From: Yanfeng Gao <gaoyanfeng.gyf at alibaba-inc.com>
Date: Fri, 14 Aug 2026 16:56:55 +0800
Subject: [PATCH] [mlir][linalg] Preserve unsigned integer widening in named
contraction vectorization
This patch fixes a correctness issue in named contraction vectorization of
Linalg ops with unsigned integer widening.
When `create_named_contraction` is enabled, an op such as
`linalg.matmul {cast = #linalg.type_fn<cast_unsigned>}` was lowered directly
to a mixed-width `vector.contract` without materializing the implied
`arith.extui`. Since `vector.contract` promotes mixed-width integer operands
using sign extension, an `i8 x i8 -> i32` matmul could interpret unsigned
inputs as signed values and produce incorrect results.
Materialize `arith.extui` for narrower integer inputs before constructing the
`vector.contract`, preserving the required zero-extension semantics.
---
.../Linalg/Transforms/Vectorization.cpp | 23 ++++
.../vectorization/contraction-interface.mlir | 128 ++++++++++++++++++
2 files changed, 151 insertions(+)
diff --git a/mlir/lib/Dialect/Linalg/Transforms/Vectorization.cpp b/mlir/lib/Dialect/Linalg/Transforms/Vectorization.cpp
index 21ca3108efcd6..c65c078bfa8f5 100644
--- a/mlir/lib/Dialect/Linalg/Transforms/Vectorization.cpp
+++ b/mlir/lib/Dialect/Linalg/Transforms/Vectorization.cpp
@@ -2069,6 +2069,29 @@ vectorizeAsLinalgContraction(RewriterBase &rewriter, VectorizationState &state,
vecOperands.push_back(read);
}
+ bool hasUnsignedCast =
+ TypeSwitch<Operation *, bool>(linalgOp.getOperation())
+ .Case<MatmulOp, BatchMatmulOp, BatchReduceMatmulOp, ContractOp>(
+ [](auto op) { return op.getCast() == TypeFn::cast_unsigned; })
+ .Default(false);
+ if (hasUnsignedCast) {
+ auto accType = dyn_cast<VectorType>(vecOperands[2].getType());
+ auto accElementType =
+ accType ? dyn_cast<IntegerType>(accType.getElementType()) : nullptr;
+ if (accElementType && accElementType.isSignless()) {
+ for (Value &operand : MutableArrayRef(vecOperands).take_front(2)) {
+ auto operandType = cast<VectorType>(operand.getType());
+ auto operandElementType =
+ dyn_cast<IntegerType>(operandType.getElementType());
+ if (!operandElementType || !operandElementType.isSignless() ||
+ operandElementType.getWidth() >= accElementType.getWidth())
+ continue;
+ operand = arith::ExtUIOp::create(
+ rewriter, loc, operandType.clone(accElementType), operand);
+ }
+ }
+ }
+
// Remap iterators from linalg to vector.
SmallVector<Attribute> iterAttrs;
auto iterators = linalgOp.getIteratorTypesArray();
diff --git a/mlir/test/Dialect/Linalg/vectorization/contraction-interface.mlir b/mlir/test/Dialect/Linalg/vectorization/contraction-interface.mlir
index d8f897cca958d..b1e25012c5ff5 100644
--- a/mlir/test/Dialect/Linalg/vectorization/contraction-interface.mlir
+++ b/mlir/test/Dialect/Linalg/vectorization/contraction-interface.mlir
@@ -348,6 +348,96 @@ module attributes {transform.with_named_sequence} {
// -----
+func.func @matmul_mixed_precision_unsigned(
+ %A: tensor<4x16xi8>, %B: tensor<16x4xi8>,
+ %C: tensor<4x4xi32>) -> tensor<4x4xi32> {
+ %0 = linalg.matmul {cast = #linalg.type_fn<cast_unsigned>}
+ ins(%A, %B : tensor<4x16xi8>, tensor<16x4xi8>)
+ outs(%C : tensor<4x4xi32>) -> tensor<4x4xi32>
+ return %0 : tensor<4x4xi32>
+}
+
+// CHECK: #[[$MAP_A:.+]] = affine_map<(d0, d1, d2) -> (d0, d2)>
+// CHECK: #[[$MAP_B:.+]] = affine_map<(d0, d1, d2) -> (d2, d1)>
+// CHECK: #[[$MAP_C:.+]] = affine_map<(d0, d1, d2) -> (d0, d1)>
+// CHECK-LABEL: func.func @matmul_mixed_precision_unsigned(
+// CHECK: %[[LOAD_A:.*]] = vector.transfer_read %{{.*}} : tensor<4x16xi8>, vector<4x16xi8>
+// CHECK: %[[LOAD_B:.*]] = vector.transfer_read %{{.*}} : tensor<16x4xi8>, vector<16x4xi8>
+// CHECK: %[[LOAD_C:.*]] = vector.transfer_read %{{.*}} : tensor<4x4xi32>, vector<4x4xi32>
+// CHECK: %[[EXT_A:.*]] = arith.extui %[[LOAD_A]] : vector<4x16xi8> to vector<4x16xi32>
+// CHECK: %[[EXT_B:.*]] = arith.extui %[[LOAD_B]] : vector<16x4xi8> to vector<16x4xi32>
+// CHECK: %[[CONTRACT:.*]] = vector.contract
+// CHECK-SAME: indexing_maps = [#[[$MAP_A]], #[[$MAP_B]], #[[$MAP_C]]]
+// CHECK-SAME: %[[EXT_A]], %[[EXT_B]], %[[LOAD_C]]
+// CHECK-SAME: : vector<4x16xi32>, vector<16x4xi32> into vector<4x4xi32>
+// CHECK: vector.transfer_write %[[CONTRACT]], %{{.*}} : vector<4x4xi32>, tensor<4x4xi32>
+
+func.func @matmul_mixed_precision_signed(
+ %A: tensor<4x16xi8>, %B: tensor<16x4xi8>,
+ %C: tensor<4x4xi32>) -> tensor<4x4xi32> {
+ %0 = linalg.matmul
+ ins(%A, %B : tensor<4x16xi8>, tensor<16x4xi8>)
+ outs(%C : tensor<4x4xi32>) -> tensor<4x4xi32>
+ return %0 : tensor<4x4xi32>
+}
+
+// CHECK-LABEL: func.func @matmul_mixed_precision_signed(
+// CHECK: %[[SIGNED_A:.*]] = vector.transfer_read %{{.*}} : tensor<4x16xi8>, vector<4x16xi8>
+// CHECK: %[[SIGNED_B:.*]] = vector.transfer_read %{{.*}} : tensor<16x4xi8>, vector<16x4xi8>
+// CHECK: %[[SIGNED_C:.*]] = vector.transfer_read %{{.*}} : tensor<4x4xi32>, vector<4x4xi32>
+// CHECK-NOT: arith.ext
+// CHECK: vector.contract
+// CHECK-SAME: %[[SIGNED_A]], %[[SIGNED_B]], %[[SIGNED_C]]
+// CHECK-SAME: : vector<4x16xi8>, vector<16x4xi8> into vector<4x4xi32>
+
+func.func @matmul_same_precision_unsigned(
+ %A: tensor<4x16xi32>, %B: tensor<16x4xi32>,
+ %C: tensor<4x4xi32>) -> tensor<4x4xi32> {
+ %0 = linalg.matmul {cast = #linalg.type_fn<cast_unsigned>}
+ ins(%A, %B : tensor<4x16xi32>, tensor<16x4xi32>)
+ outs(%C : tensor<4x4xi32>) -> tensor<4x4xi32>
+ return %0 : tensor<4x4xi32>
+}
+
+// CHECK-LABEL: func.func @matmul_same_precision_unsigned(
+// CHECK: %[[SAME_A:.*]] = vector.transfer_read %{{.*}} : tensor<4x16xi32>, vector<4x16xi32>
+// CHECK: %[[SAME_B:.*]] = vector.transfer_read %{{.*}} : tensor<16x4xi32>, vector<16x4xi32>
+// CHECK: %[[SAME_C:.*]] = vector.transfer_read %{{.*}} : tensor<4x4xi32>, vector<4x4xi32>
+// CHECK-NOT: arith.extui
+// CHECK: vector.contract
+// CHECK-SAME: %[[SAME_A]], %[[SAME_B]], %[[SAME_C]]
+
+func.func @matmul_mixed_precision_unsigned_dynamic(
+ %A: tensor<?x?xi8>, %B: tensor<?x?xi8>,
+ %C: tensor<?x?xi32>) -> tensor<?x?xi32> {
+ %0 = linalg.matmul {cast = #linalg.type_fn<cast_unsigned>}
+ ins(%A, %B : tensor<?x?xi8>, tensor<?x?xi8>)
+ outs(%C : tensor<?x?xi32>) -> tensor<?x?xi32>
+ return %0 : tensor<?x?xi32>
+}
+
+// CHECK-LABEL: func.func @matmul_mixed_precision_unsigned_dynamic(
+// CHECK: %[[MASKED_A:.*]] = vector.mask {{.*}} -> vector<4x16xi8>
+// CHECK: %[[MASKED_B:.*]] = vector.mask {{.*}} -> vector<16x4xi8>
+// CHECK: %[[MASKED_C:.*]] = vector.mask {{.*}} -> vector<4x4xi32>
+// CHECK: %[[MASKED_EXT_A:.*]] = arith.extui %[[MASKED_A]] : vector<4x16xi8> to vector<4x16xi32>
+// CHECK: %[[MASKED_EXT_B:.*]] = arith.extui %[[MASKED_B]] : vector<16x4xi8> to vector<16x4xi32>
+// CHECK: vector.mask
+// CHECK-SAME: vector.contract
+// CHECK-SAME: %[[MASKED_EXT_A]], %[[MASKED_EXT_B]], %[[MASKED_C]]
+// CHECK-SAME: : vector<4x4x16xi1> -> vector<4x4xi32>
+
+module attributes {transform.with_named_sequence} {
+ transform.named_sequence @__transform_main(%arg1: !transform.any_op {transform.readonly}) {
+ %0 = transform.structured.match ops{["linalg.matmul"]} in %arg1 : (!transform.any_op) -> !transform.any_op
+ transform.structured.vectorize %0 vector_sizes [4, 4, 16]
+ {create_named_contraction} : !transform.any_op
+ transform.yield
+ }
+}
+
+// -----
+
func.func @batch_matmul(%A: tensor<3x8x4xf32>, %B: tensor<3x4x16xf32>,
%C: tensor<3x8x16xf32>) -> tensor<3x8x16xf32> {
%0 = linalg.batch_matmul
@@ -450,6 +540,44 @@ module attributes {transform.with_named_sequence} {
// -----
+func.func @contract_mixed_precision_unsigned(
+ %A: tensor<4x8x2xi8>, %B: tensor<8x16x2xi8>,
+ %C: tensor<4x16xi32>) -> tensor<4x16xi32> {
+ %0 = linalg.contract
+ indexing_maps = [affine_map<(m, n, k, kk) -> (m, k, kk)>,
+ affine_map<(m, n, k, kk) -> (k, n, kk)>,
+ affine_map<(m, n, k, kk) -> (m, n)>]
+ {cast = #linalg.type_fn<cast_unsigned>}
+ ins(%A, %B : tensor<4x8x2xi8>, tensor<8x16x2xi8>)
+ outs(%C : tensor<4x16xi32>) -> tensor<4x16xi32>
+ return %0 : tensor<4x16xi32>
+}
+
+// CHECK: #[[$MAP_A:.+]] = affine_map<(d0, d1, d2, d3) -> (d0, d2, d3)>
+// CHECK: #[[$MAP_B:.+]] = affine_map<(d0, d1, d2, d3) -> (d2, d1, d3)>
+// CHECK: #[[$MAP_C:.+]] = affine_map<(d0, d1, d2, d3) -> (d0, d1)>
+// CHECK-LABEL: func.func @contract_mixed_precision_unsigned(
+// CHECK: %[[LOAD_A:.*]] = vector.transfer_read %{{.*}} : tensor<4x8x2xi8>, vector<4x8x2xi8>
+// CHECK: %[[LOAD_B:.*]] = vector.transfer_read %{{.*}} : tensor<8x16x2xi8>, vector<8x16x2xi8>
+// CHECK: %[[LOAD_C:.*]] = vector.transfer_read %{{.*}} : tensor<4x16xi32>, vector<4x16xi32>
+// CHECK: %[[EXT_A:.*]] = arith.extui %[[LOAD_A]] : vector<4x8x2xi8> to vector<4x8x2xi32>
+// CHECK: %[[EXT_B:.*]] = arith.extui %[[LOAD_B]] : vector<8x16x2xi8> to vector<8x16x2xi32>
+// CHECK: %[[CONTRACT:.*]] = vector.contract
+// CHECK-SAME: indexing_maps = [#[[$MAP_A]], #[[$MAP_B]], #[[$MAP_C]]]
+// CHECK-SAME: %[[EXT_A]], %[[EXT_B]], %[[LOAD_C]]
+// CHECK-SAME: : vector<4x8x2xi32>, vector<8x16x2xi32> into vector<4x16xi32>
+// CHECK: vector.transfer_write %[[CONTRACT]], %{{.*}} : vector<4x16xi32>, tensor<4x16xi32>
+
+module attributes {transform.with_named_sequence} {
+ transform.named_sequence @__transform_main(%arg1: !transform.any_op {transform.readonly}) {
+ %0 = transform.structured.match ops{["linalg.contract"]} in %arg1 : (!transform.any_op) -> !transform.any_op
+ transform.structured.vectorize %0 {create_named_contraction} : !transform.any_op
+ transform.yield
+ }
+}
+
+// -----
+
/// Generic can represent contractions but it does not implement contraction interface.
/// Thus, direct lowering to vector.contract is not supported.
/// Vectorization still works and applies generic rewrite logic.
More information about the Mlir-commits
mailing list