[Mlir-commits] [mlir] [mlir][linalg] Preserve unsigned integer widening in named contraction vectorization (PR #216283)
llvmlistbot at llvm.org
llvmlistbot at llvm.org
Wed Aug 26 01:09:03 PDT 2026
https://github.com/Pecco-314 updated https://github.com/llvm/llvm-project/pull/216283
>From cb527c12a1f467091a1ceb1355958301133bc88d 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 1/3] [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.
Assisted-by: Codex (OpenAI)
---
.../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.
>From b96c95dc1147a28c34df234b4fdd1dcb077f36e9 Mon Sep 17 00:00:00 2001
From: Yanfeng Gao <gaoyanfeng.gyf at alibaba-inc.com>
Date: Tue, 25 Aug 2026 20:43:21 +0800
Subject: [PATCH 2/3] [mlir][linalg] Address contraction vectorization review
Query the common cast attribute without enumerating concrete contraction ops, document the integer widening behavior, and split the new tests into independent inputs.
Assisted-by: Codex (OpenAI)
---
.../Linalg/Transforms/Vectorization.cpp | 9 ++---
.../vectorization/contraction-interface.mlir | 35 +++++++++++++++++++
2 files changed, 40 insertions(+), 4 deletions(-)
diff --git a/mlir/lib/Dialect/Linalg/Transforms/Vectorization.cpp b/mlir/lib/Dialect/Linalg/Transforms/Vectorization.cpp
index c65c078bfa8f5..aeedec87769e5 100644
--- a/mlir/lib/Dialect/Linalg/Transforms/Vectorization.cpp
+++ b/mlir/lib/Dialect/Linalg/Transforms/Vectorization.cpp
@@ -2069,11 +2069,12 @@ vectorizeAsLinalgContraction(RewriterBase &rewriter, VectorizationState &state,
vecOperands.push_back(read);
}
+ // When integer operands are narrower than the accumulator, vector.contract
+ // promotes them by sign extension. This matches cast_signed. To preserve
+ // cast_unsigned, zero-extend the operands explicitly before the contraction.
+ auto castAttr = linalgOp->getAttrOfType<TypeFnAttr>("cast");
bool hasUnsignedCast =
- TypeSwitch<Operation *, bool>(linalgOp.getOperation())
- .Case<MatmulOp, BatchMatmulOp, BatchReduceMatmulOp, ContractOp>(
- [](auto op) { return op.getCast() == TypeFn::cast_unsigned; })
- .Default(false);
+ castAttr && castAttr.getValue() == TypeFn::cast_unsigned;
if (hasUnsignedCast) {
auto accType = dyn_cast<VectorType>(vecOperands[2].getType());
auto accElementType =
diff --git a/mlir/test/Dialect/Linalg/vectorization/contraction-interface.mlir b/mlir/test/Dialect/Linalg/vectorization/contraction-interface.mlir
index b1e25012c5ff5..e2884ad562e8e 100644
--- a/mlir/test/Dialect/Linalg/vectorization/contraction-interface.mlir
+++ b/mlir/test/Dialect/Linalg/vectorization/contraction-interface.mlir
@@ -372,6 +372,17 @@ func.func @matmul_mixed_precision_unsigned(
// CHECK-SAME: : vector<4x16xi32>, vector<16x4xi32> into vector<4x4xi32>
// CHECK: vector.transfer_write %[[CONTRACT]], %{{.*}} : vector<4x4xi32>, tensor<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 @matmul_mixed_precision_signed(
%A: tensor<4x16xi8>, %B: tensor<16x4xi8>,
%C: tensor<4x4xi32>) -> tensor<4x4xi32> {
@@ -381,6 +392,8 @@ func.func @matmul_mixed_precision_signed(
return %0 : tensor<4x4xi32>
}
+// vector.contract sign-extends mixed-width integer operands by default, so no
+// explicit extension is needed for cast_signed.
// 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>
@@ -390,6 +403,17 @@ func.func @matmul_mixed_precision_signed(
// CHECK-SAME: %[[SIGNED_A]], %[[SIGNED_B]], %[[SIGNED_C]]
// CHECK-SAME: : vector<4x16xi8>, vector<16x4xi8> into 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 @matmul_same_precision_unsigned(
%A: tensor<4x16xi32>, %B: tensor<16x4xi32>,
%C: tensor<4x4xi32>) -> tensor<4x4xi32> {
@@ -407,6 +431,17 @@ func.func @matmul_same_precision_unsigned(
// CHECK: vector.contract
// CHECK-SAME: %[[SAME_A]], %[[SAME_B]], %[[SAME_C]]
+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 @matmul_mixed_precision_unsigned_dynamic(
%A: tensor<?x?xi8>, %B: tensor<?x?xi8>,
%C: tensor<?x?xi32>) -> tensor<?x?xi32> {
>From c08b8475c11674581309a3750485ab2d8a7bbb83 Mon Sep 17 00:00:00 2001
From: Yanfeng Gao <gaoyanfeng.gyf at alibaba-inc.com>
Date: Wed, 26 Aug 2026 16:06:07 +0800
Subject: [PATCH 3/3] [mlir][linalg] Preserve float-to-integer contraction
casts
Named contraction vectorization constructs vector.contract directly. Materialize fptosi or fptoui for floating-point operands with an integer accumulator, while retaining vector.contract's implicit signed integer promotion and the explicit extension needed for unsigned widening.
Assisted-by: Codex (OpenAI)
---
.../Linalg/Transforms/Vectorization.cpp | 42 ++++++++-----
.../vectorization/contraction-interface.mlir | 60 +++++++++++++++++++
2 files changed, 86 insertions(+), 16 deletions(-)
diff --git a/mlir/lib/Dialect/Linalg/Transforms/Vectorization.cpp b/mlir/lib/Dialect/Linalg/Transforms/Vectorization.cpp
index aeedec87769e5..2289fc565e173 100644
--- a/mlir/lib/Dialect/Linalg/Transforms/Vectorization.cpp
+++ b/mlir/lib/Dialect/Linalg/Transforms/Vectorization.cpp
@@ -2069,26 +2069,36 @@ vectorizeAsLinalgContraction(RewriterBase &rewriter, VectorizationState &state,
vecOperands.push_back(read);
}
- // When integer operands are narrower than the accumulator, vector.contract
- // promotes them by sign extension. This matches cast_signed. To preserve
- // cast_unsigned, zero-extend the operands explicitly before the contraction.
+ // Preserve the contraction's cast semantics when converting operands to the
+ // integer accumulator type. vector.contract provides an implicit signed
+ // integer promotion; the cases below materialize explicit casts as needed.
auto castAttr = linalgOp->getAttrOfType<TypeFnAttr>("cast");
bool hasUnsignedCast =
castAttr && castAttr.getValue() == TypeFn::cast_unsigned;
- 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())
+ 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());
+ Type operandElementType = operandType.getElementType();
+ VectorType castType = operandType.clone(accElementType);
+ if (isa<FloatType>(operandElementType)) {
+ // Floating-point operands require an explicit signed or unsigned
+ // conversion to the integer accumulator type.
+ if (hasUnsignedCast)
+ operand = arith::FPToUIOp::create(rewriter, loc, castType, operand);
+ else
+ operand = arith::FPToSIOp::create(rewriter, loc, castType, operand);
+ } else {
+ auto operandIntegerType = dyn_cast<IntegerType>(operandElementType);
+ if (!operandIntegerType || !operandIntegerType.isSignless() ||
+ operandIntegerType.getWidth() >= accElementType.getWidth())
continue;
- operand = arith::ExtUIOp::create(
- rewriter, loc, operandType.clone(accElementType), operand);
+ // Integer vector.contract implicitly sign-extends the operands.
+ // Unsigned promotion requires an explicit zero extension.
+ if (hasUnsignedCast)
+ operand = arith::ExtUIOp::create(rewriter, loc, castType, operand);
}
}
}
diff --git a/mlir/test/Dialect/Linalg/vectorization/contraction-interface.mlir b/mlir/test/Dialect/Linalg/vectorization/contraction-interface.mlir
index e2884ad562e8e..b2fbadeeba3c6 100644
--- a/mlir/test/Dialect/Linalg/vectorization/contraction-interface.mlir
+++ b/mlir/test/Dialect/Linalg/vectorization/contraction-interface.mlir
@@ -383,6 +383,66 @@ module attributes {transform.with_named_sequence} {
// -----
+func.func @matmul_float_to_unsigned_integer(
+ %A: tensor<4x16xf16>, %B: tensor<16x4xf32>,
+ %C: tensor<4x4xi32>) -> tensor<4x4xi32> {
+ %0 = linalg.matmul {cast = #linalg.type_fn<cast_unsigned>}
+ ins(%A, %B : tensor<4x16xf16>, tensor<16x4xf32>)
+ outs(%C : tensor<4x4xi32>) -> tensor<4x4xi32>
+ return %0 : tensor<4x4xi32>
+}
+
+// CHECK-LABEL: func.func @matmul_float_to_unsigned_integer(
+// CHECK: %[[LOAD_A:.*]] = vector.transfer_read %{{.*}} : tensor<4x16xf16>, vector<4x16xf16>
+// CHECK: %[[LOAD_B:.*]] = vector.transfer_read %{{.*}} : tensor<16x4xf32>, vector<16x4xf32>
+// CHECK: %[[LOAD_C:.*]] = vector.transfer_read %{{.*}} : tensor<4x4xi32>, vector<4x4xi32>
+// CHECK: %[[CAST_A:.*]] = arith.fptoui %[[LOAD_A]] : vector<4x16xf16> to vector<4x16xi32>
+// CHECK: %[[CAST_B:.*]] = arith.fptoui %[[LOAD_B]] : vector<16x4xf32> to vector<16x4xi32>
+// CHECK: %[[CONTRACT:.*]] = vector.contract
+// CHECK-SAME: %[[CAST_A]], %[[CAST_B]], %[[LOAD_C]]
+// CHECK-SAME: : vector<4x16xi32>, vector<16x4xi32> into 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 @matmul_float_to_signed_integer(
+ %A: tensor<4x16xf16>, %B: tensor<16x4xf32>,
+ %C: tensor<4x4xi32>) -> tensor<4x4xi32> {
+ %0 = linalg.matmul
+ ins(%A, %B : tensor<4x16xf16>, tensor<16x4xf32>)
+ outs(%C : tensor<4x4xi32>) -> tensor<4x4xi32>
+ return %0 : tensor<4x4xi32>
+}
+
+// CHECK-LABEL: func.func @matmul_float_to_signed_integer(
+// CHECK: %[[LOAD_A:.*]] = vector.transfer_read %{{.*}} : tensor<4x16xf16>, vector<4x16xf16>
+// CHECK: %[[LOAD_B:.*]] = vector.transfer_read %{{.*}} : tensor<16x4xf32>, vector<16x4xf32>
+// CHECK: %[[LOAD_C:.*]] = vector.transfer_read %{{.*}} : tensor<4x4xi32>, vector<4x4xi32>
+// CHECK: %[[CAST_A:.*]] = arith.fptosi %[[LOAD_A]] : vector<4x16xf16> to vector<4x16xi32>
+// CHECK: %[[CAST_B:.*]] = arith.fptosi %[[LOAD_B]] : vector<16x4xf32> to vector<16x4xi32>
+// CHECK: %[[CONTRACT:.*]] = vector.contract
+// CHECK-SAME: %[[CAST_A]], %[[CAST_B]], %[[LOAD_C]]
+// CHECK-SAME: : vector<4x16xi32>, vector<16x4xi32> into 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 @matmul_mixed_precision_signed(
%A: tensor<4x16xi8>, %B: tensor<16x4xi8>,
%C: tensor<4x4xi32>) -> tensor<4x4xi32> {
More information about the Mlir-commits
mailing list