[Mlir-commits] [mlir] [mlir][x86] Fix - replace `vector.load` with `vector.transfer_read` (PR #201503)
Arun Thangamani
llvmlistbot at llvm.org
Thu Jun 4 08:40:19 PDT 2026
https://github.com/arun-thmn updated https://github.com/llvm/llvm-project/pull/201503
>From 93cbada05b66964031536814e0db4f04375f6a06 Mon Sep 17 00:00:00 2001
From: Arun Thangamani <arun.thangamani at intel.com>
Date: Wed, 3 Jun 2026 21:26:36 -0700
Subject: [PATCH 1/5] fix for int8 shuffling (online) - replace vector.load
with vector.transfer_read
---
.../VectorContractToAMXDotProduct.cpp | 29 +++++++++++++++----
.../VectorContractToPackedTypeDotProduct.cpp | 21 +++++++++++---
2 files changed, 40 insertions(+), 10 deletions(-)
diff --git a/mlir/lib/Dialect/X86/Transforms/VectorContractToAMXDotProduct.cpp b/mlir/lib/Dialect/X86/Transforms/VectorContractToAMXDotProduct.cpp
index 1ffc31672296f..36d0af51bfc5d 100644
--- a/mlir/lib/Dialect/X86/Transforms/VectorContractToAMXDotProduct.cpp
+++ b/mlir/lib/Dialect/X86/Transforms/VectorContractToAMXDotProduct.cpp
@@ -264,17 +264,34 @@ static void performShuffle(OpBuilder &rewriter, Location loc, Value matB,
ValueRange iterArgs) {
subviewOffset[subviewOffset.size() - 2] = iv;
- auto vec1 = vector::LoadOp::create(
- rewriter, loc, VectorType::get((16 * offset), ipType), matB,
- ValueRange(subviewOffset));
+ auto flatTy = VectorType::get({2, (16 * (offset / 2))}, ipType);
+ if (ipType.isBF16())
+ flatTy = VectorType::get((16 * offset), ipType);
+
+ int64_t srcRank = (dyn_cast<ShapedType>(matB.getType())).getRank();
+ Value padding = ub::PoisonOp::create(rewriter, loc, ipType);
+ auto map = AffineMap::getMinorIdentityMap(srcRank, flatTy.getRank(),
+ rewriter.getContext());
+ SmallVector<bool> inBounds(flatTy.getRank(), true);
+ Value vec1 = vector::TransferReadOp::create(rewriter, loc, flatTy, matB,
+ ValueRange(subviewOffset),
+ padding, map, inBounds);
+
+ if (!ipType.isBF16())
+ vec1 = vector::ShapeCastOp::create(
+ rewriter, loc, VectorType::get((16 * offset), ipType), vec1);
// Increment the iv by 1 or 2 based on the type to load the next 32/64
// elements
Value incIV = arith::AddIOp::create(rewriter, loc, offsetIndx, iv);
subviewOffset[subviewOffset.size() - 2] = incIV;
- auto vec2 = vector::LoadOp::create(
- rewriter, loc, VectorType::get((16 * offset), ipType), matB,
- ValueRange(subviewOffset));
+
+ Value vec2 = vector::TransferReadOp::create(rewriter, loc, flatTy, matB,
+ ValueRange(subviewOffset),
+ padding, map, inBounds);
+ if (!ipType.isBF16())
+ vec2 = vector::ShapeCastOp::create(
+ rewriter, loc, VectorType::get((16 * offset), ipType), vec2);
vector::ShuffleOp shuffle1;
vector::ShuffleOp shuffle2;
diff --git a/mlir/lib/Dialect/X86/Transforms/VectorContractToPackedTypeDotProduct.cpp b/mlir/lib/Dialect/X86/Transforms/VectorContractToPackedTypeDotProduct.cpp
index be6d8f55e5b4c..d3487369355f0 100644
--- a/mlir/lib/Dialect/X86/Transforms/VectorContractToPackedTypeDotProduct.cpp
+++ b/mlir/lib/Dialect/X86/Transforms/VectorContractToPackedTypeDotProduct.cpp
@@ -82,6 +82,9 @@ static void packNonUnitDimOperandToVNNI(mlir::PatternRewriter &rewriter,
auto elemTy = Ty.getElementType();
auto flatTy = mlir::VectorType::get(nonUnitDimAcc, elemTy);
+ if (elemTy.isSignlessInteger(8))
+ flatTy = mlir::VectorType::get({2, nonUnitDimAcc / 2}, elemTy);
+
Value srcBuff;
SmallVector<Value> indexVals;
@@ -104,8 +107,12 @@ static void packNonUnitDimOperandToVNNI(mlir::PatternRewriter &rewriter,
rewriter.getContext());
SmallVector<bool> inBounds(flatTy.getRank(), true);
- auto vec1 = vector::TransferReadOp::create(rewriter, loc, flatTy, srcBuff,
- indexVals, padding, map, inBounds);
+ Value vec1 = vector::TransferReadOp::create(
+ rewriter, loc, flatTy, srcBuff, indexVals, padding, map, inBounds);
+
+ if (elemTy.isSignlessInteger(8))
+ vec1 = vector::ShapeCastOp::create(
+ rewriter, loc, VectorType::get(nonUnitDimAcc, elemTy), vec1);
unsigned int offset = 1;
if (elemTy.isSignlessInteger(8))
@@ -117,8 +124,14 @@ static void packNonUnitDimOperandToVNNI(mlir::PatternRewriter &rewriter,
indexVals[indexVals.size() - 2]);
indexVals[indexVals.size() - 2] = nextIndx;
- auto vec2 = vector::TransferReadOp::create(rewriter, loc, flatTy, srcBuff,
- indexVals, padding, map, inBounds);
+ Value vec2 = vector::TransferReadOp::create(
+ rewriter, loc, flatTy, srcBuff, indexVals, padding, map, inBounds);
+
+ if (elemTy.isSignlessInteger(8))
+ vec2 = vector::ShapeCastOp::create(
+ rewriter, loc, VectorType::get(nonUnitDimAcc, elemTy), vec2);
+
+ flatTy = mlir::VectorType::get(nonUnitDimAcc, elemTy);
static constexpr int64_t maskLo_bf16[] = {
0, 32, 1, 33, 2, 34, 3, 35, 8, 40, 9, 41, 10, 42, 11, 43,
>From 0bf9a8a1a75fba07173d3d9d7aeaef4bf434a7f6 Mon Sep 17 00:00:00 2001
From: Arun Thangamani <arun.thangamani at intel.com>
Date: Thu, 4 Jun 2026 04:52:48 -0700
Subject: [PATCH 2/5] test-case inclusion + re-name variable
---
.../VectorContractToAMXDotProduct.cpp | 42 +++++++++----------
.../X86/AMX/vector-contract-to-tiled-dp.mlir | 8 ++++
...or-contract-to-packed-type-dotproduct.mlir | 4 ++
3 files changed, 33 insertions(+), 21 deletions(-)
diff --git a/mlir/lib/Dialect/X86/Transforms/VectorContractToAMXDotProduct.cpp b/mlir/lib/Dialect/X86/Transforms/VectorContractToAMXDotProduct.cpp
index 36d0af51bfc5d..30628500e4bad 100644
--- a/mlir/lib/Dialect/X86/Transforms/VectorContractToAMXDotProduct.cpp
+++ b/mlir/lib/Dialect/X86/Transforms/VectorContractToAMXDotProduct.cpp
@@ -264,18 +264,18 @@ static void performShuffle(OpBuilder &rewriter, Location loc, Value matB,
ValueRange iterArgs) {
subviewOffset[subviewOffset.size() - 2] = iv;
- auto flatTy = VectorType::get({2, (16 * (offset / 2))}, ipType);
+ auto ipVectorType = VectorType::get({2, (16 * (offset / 2))}, ipType);
if (ipType.isBF16())
- flatTy = VectorType::get((16 * offset), ipType);
+ ipVectorType = VectorType::get((16 * offset), ipType);
int64_t srcRank = (dyn_cast<ShapedType>(matB.getType())).getRank();
Value padding = ub::PoisonOp::create(rewriter, loc, ipType);
- auto map = AffineMap::getMinorIdentityMap(srcRank, flatTy.getRank(),
- rewriter.getContext());
- SmallVector<bool> inBounds(flatTy.getRank(), true);
- Value vec1 = vector::TransferReadOp::create(rewriter, loc, flatTy, matB,
- ValueRange(subviewOffset),
- padding, map, inBounds);
+ auto map = AffineMap::getMinorIdentityMap(
+ srcRank, ipVectorType.getRank(), rewriter.getContext());
+ SmallVector<bool> inBounds(ipVectorType.getRank(), true);
+ Value vec1 = vector::TransferReadOp::create(
+ rewriter, loc, ipVectorType, matB, ValueRange(subviewOffset),
+ padding, map, inBounds);
if (!ipType.isBF16())
vec1 = vector::ShapeCastOp::create(
@@ -286,9 +286,9 @@ static void performShuffle(OpBuilder &rewriter, Location loc, Value matB,
Value incIV = arith::AddIOp::create(rewriter, loc, offsetIndx, iv);
subviewOffset[subviewOffset.size() - 2] = incIV;
- Value vec2 = vector::TransferReadOp::create(rewriter, loc, flatTy, matB,
- ValueRange(subviewOffset),
- padding, map, inBounds);
+ Value vec2 = vector::TransferReadOp::create(
+ rewriter, loc, ipVectorType, matB, ValueRange(subviewOffset),
+ padding, map, inBounds);
if (!ipType.isBF16())
vec2 = vector::ShapeCastOp::create(
rewriter, loc, VectorType::get((16 * offset), ipType), vec2);
@@ -991,17 +991,17 @@ struct VectorContractToAMXDotProduct
amx::TileStoreOp::create(rewriter, loc, resultBuffer, ValueRange{c0, c0},
dp);
- auto flatTy = mlir::VectorType::get({16, 16}, opType);
+ auto ipVectorType = mlir::VectorType::get({16, 16}, opType);
int64_t srcRank =
(dyn_cast<ShapedType>(resultBuffer.getType())).getRank();
Value padding = ub::PoisonOp::create(rewriter, loc, opType);
- auto map = AffineMap::getMinorIdentityMap(srcRank, flatTy.getRank(),
+ auto map = AffineMap::getMinorIdentityMap(srcRank, ipVectorType.getRank(),
rewriter.getContext());
- SmallVector<bool> inBounds(flatTy.getRank(), true);
+ SmallVector<bool> inBounds(ipVectorType.getRank(), true);
Value vecRow = vector::TransferReadOp::create(
- rewriter, loc, flatTy, resultBuffer, ValueRange{c0, c0}, padding, map,
- inBounds);
+ rewriter, loc, ipVectorType, resultBuffer, ValueRange{c0, c0},
+ padding, map, inBounds);
Value resultOp = contractionUsersAfterYield(contractOp.getResult());
if (auto vecType = llvm::dyn_cast<VectorType>(resultOp.getType()))
@@ -1479,17 +1479,17 @@ struct VectorContractToAMXDotProduct
Value indexOp_i = arith::ConstantIndexOp::create(rewriter, loc, i);
Value indexOp_j = arith::ConstantIndexOp::create(rewriter, loc, j);
- auto flatTy = mlir::VectorType::get({16, 16}, opType);
+ auto ipVectorType = mlir::VectorType::get({16, 16}, opType);
int64_t srcRank =
(dyn_cast<ShapedType>(resultBuffer.getType())).getRank();
Value padding = ub::PoisonOp::create(rewriter, loc, opType);
- auto map = AffineMap::getMinorIdentityMap(srcRank, flatTy.getRank(),
- rewriter.getContext());
- SmallVector<bool> inBounds(flatTy.getRank(), true);
+ auto map = AffineMap::getMinorIdentityMap(
+ srcRank, ipVectorType.getRank(), rewriter.getContext());
+ SmallVector<bool> inBounds(ipVectorType.getRank(), true);
auto vec1 = vector::TransferReadOp::create(
- rewriter, loc, flatTy, resultBuffer,
+ rewriter, loc, ipVectorType, resultBuffer,
ValueRange{indexOp_i, indexOp_j}, padding, map, inBounds);
writeResults.push_back(vec1);
}
diff --git a/mlir/test/Dialect/X86/AMX/vector-contract-to-tiled-dp.mlir b/mlir/test/Dialect/X86/AMX/vector-contract-to-tiled-dp.mlir
index fb2314dfdf506..e9d1032fc9858 100644
--- a/mlir/test/Dialect/X86/AMX/vector-contract-to-tiled-dp.mlir
+++ b/mlir/test/Dialect/X86/AMX/vector-contract-to-tiled-dp.mlir
@@ -1014,6 +1014,10 @@ func.func @online_packing_f8E5M2_matmul_loop(%arg0: memref<64x256xf8E5M2>, %arg1
// CHECK-LABEL: @online_packing_f8E5M2_matmul_loop
// CHECK-COUNT-4: x86.amx.tile_zero : !x86.amx.tile<16x16xf32>
// CHECK: scf.for {{.*}} -> (!x86.amx.tile<16x16xf32>, !x86.amx.tile<16x16xf32>, !x86.amx.tile<16x16xf32>, !x86.amx.tile<16x16xf32>) {
+// CHECK: vector.transfer_read {{.*}} vector<2x32xf8E5M2>
+// CHECK-NEXT: vector.shape_cast {{.*}} : vector<2x32xf8E5M2> to vector<64xf8E5M2>
+// CHECK: vector.transfer_read {{.*}} vector<2x32xf8E5M2>
+// CHECK-NEXT: vector.shape_cast {{.*}} : vector<2x32xf8E5M2> to vector<64xf8E5M2>
// CHECK: vector.shuffle{{.*}}[0, 32, 64, 96, 1, 33, 65, 97, 2, 34, 66, 98, 3, 35, 67, 99, 8, 40, 72, 104, 9, 41, 73, 105, 10, 42, 74, 106, 11, 43, 75, 107, 16, 48, 80, 112, 17, 49, 81, 113, 18, 50, 82, 114, 19, 51, 83, 115, 24, 56, 88, 120, 25, 57, 89, 121, 26, 58, 90, 122, 27, 59, 91, 123] : vector<64xf8E5M2>, vector<64xf8E5M2>
// CHECK-NEXT: vector.shuffle{{.*}}[4, 36, 68, 100, 5, 37, 69, 101, 6, 38, 70, 102, 7, 39, 71, 103, 12, 44, 76, 108, 13, 45, 77, 109, 14, 46, 78, 110, 15, 47, 79, 111, 20, 52, 84, 116, 21, 53, 85, 117, 22, 54, 86, 118, 23, 55, 87, 119, 28, 60, 92, 124, 29, 61, 93, 125, 30, 62, 94, 126, 31, 63, 95, 127] : vector<64xf8E5M2>, vector<64xf8E5M2>
// CHECK: x86.amx.tile_load
@@ -1179,6 +1183,10 @@ func.func @online_packing_int8_matmul_loop(%arg0: memref<64x256xi8>, %arg1: memr
// CHECK-LABEL: @online_packing_int8_matmul_loop
// CHECK-COUNT-4: x86.amx.tile_zero : !x86.amx.tile<16x16xi32>
// CHECK: scf.for {{.*}} -> (!x86.amx.tile<16x16xi32>, !x86.amx.tile<16x16xi32>, !x86.amx.tile<16x16xi32>, !x86.amx.tile<16x16xi32>) {
+// CHECK: vector.transfer_read {{.*}} vector<2x32xi8>
+// CHECK-NEXT: vector.shape_cast {{.*}} : vector<2x32xi8> to vector<64xi8>
+// CHECK: vector.transfer_read {{.*}} vector<2x32xi8>
+// CHECK-NEXT: vector.shape_cast {{.*}} : vector<2x32xi8> to vector<64xi8>
// CHECK: vector.shuffle{{.*}}[0, 32, 64, 96, 1, 33, 65, 97, 2, 34, 66, 98, 3, 35, 67, 99, 8, 40, 72, 104, 9, 41, 73, 105, 10, 42, 74, 106, 11, 43, 75, 107, 16, 48, 80, 112, 17, 49, 81, 113, 18, 50, 82, 114, 19, 51, 83, 115, 24, 56, 88, 120, 25, 57, 89, 121, 26, 58, 90, 122, 27, 59, 91, 123] : vector<64xi8>, vector<64xi8>
// CHECK-NEXT: vector.shuffle{{.*}}[4, 36, 68, 100, 5, 37, 69, 101, 6, 38, 70, 102, 7, 39, 71, 103, 12, 44, 76, 108, 13, 45, 77, 109, 14, 46, 78, 110, 15, 47, 79, 111, 20, 52, 84, 116, 21, 53, 85, 117, 22, 54, 86, 118, 23, 55, 87, 119, 28, 60, 92, 124, 29, 61, 93, 125, 30, 62, 94, 126, 31, 63, 95, 127] : vector<64xi8>, vector<64xi8>
// CHECK: x86.amx.tile_load
diff --git a/mlir/test/Dialect/X86/vector-contract-to-packed-type-dotproduct.mlir b/mlir/test/Dialect/X86/vector-contract-to-packed-type-dotproduct.mlir
index 2efe0f2adc413..e454e88f163d6 100644
--- a/mlir/test/Dialect/X86/vector-contract-to-packed-type-dotproduct.mlir
+++ b/mlir/test/Dialect/X86/vector-contract-to-packed-type-dotproduct.mlir
@@ -866,6 +866,10 @@ func.func @brgemm_int8_flat_avx10(%arg0: memref<16x64x256xi8>, %arg1: memref<16x
// CHECK-LABEL: @brgemm_int8_flat_avx10
// CHECK: vector.shuffle{{.*}}[0, 1, 2, 3, 16, 17, 18, 19, 4, 5, 6, 7, 20, 21, 22, 23] : vector<16xi32>, vector<16xi32>
// CHECK-NEXT: vector.shuffle{{.*}}[8, 9, 10, 11, 24, 25, 26, 27, 12, 13, 14, 15, 28, 29, 30, 31] : vector<16xi32>, vector<16xi32>
+// CHECK: vector.transfer_read {{.*}} vector<2x32xi8>
+// CHECK-NEXT: vector.shape_cast {{.*}} : vector<2x32xi8> to vector<64xi8>
+// CHECK: vector.transfer_read {{.*}} vector<2x32xi8>
+// CHECK-NEXT: vector.shape_cast {{.*}} : vector<2x32xi8> to vector<64xi8>
// CHECK: vector.shuffle{{.*}}[0, 32, 64, 96, 1, 33, 65, 97, 2, 34, 66, 98, 3, 35, 67, 99, 8, 40, 72, 104, 9, 41, 73, 105, 10, 42, 74, 106, 11, 43, 75, 107, 16, 48, 80, 112, 17, 49, 81, 113, 18, 50, 82, 114, 19, 51, 83, 115, 24, 56, 88, 120, 25, 57, 89, 121, 26, 58, 90, 122, 27, 59, 91, 123] : vector<64xi8>, vector<64xi8>
// CHECK-NEXT: vector.shuffle{{.*}}[4, 36, 68, 100, 5, 37, 69, 101, 6, 38, 70, 102, 7, 39, 71, 103, 12, 44, 76, 108, 13, 45, 77, 109, 14, 46, 78, 110, 15, 47, 79, 111, 20, 52, 84, 116, 21, 53, 85, 117, 22, 54, 86, 118, 23, 55, 87, 119, 28, 60, 92, 124, 29, 61, 93, 125, 30, 62, 94, 126, 31, 63, 95, 127] : vector<64xi8>, vector<64xi8>
// CHECK: x86.avx10.dot.i8
>From 0e68ca6548421f9efcef904f72d8f37486966fce Mon Sep 17 00:00:00 2001
From: Arun Thangamani <arun.thangamani at intel.com>
Date: Thu, 4 Jun 2026 05:11:28 -0700
Subject: [PATCH 3/5] test-case inclusion for bf16
---
mlir/test/Dialect/X86/AMX/vector-contract-to-tiled-dp.mlir | 2 ++
.../Dialect/X86/vector-contract-to-packed-type-dotproduct.mlir | 2 ++
2 files changed, 4 insertions(+)
diff --git a/mlir/test/Dialect/X86/AMX/vector-contract-to-tiled-dp.mlir b/mlir/test/Dialect/X86/AMX/vector-contract-to-tiled-dp.mlir
index e9d1032fc9858..fbbd20202620b 100644
--- a/mlir/test/Dialect/X86/AMX/vector-contract-to-tiled-dp.mlir
+++ b/mlir/test/Dialect/X86/AMX/vector-contract-to-tiled-dp.mlir
@@ -826,6 +826,8 @@ func.func @online_packing_bf16_loop(%arg0: memref<16x64x96xbf16>, %arg1: memref<
// CHECK-LABEL: @online_packing_bf16_loop
// CHECK-COUNT-4: x86.amx.tile_zero : !x86.amx.tile<16x16xf32>
// CHECK-COUNT-4: scf.for {{.*}} -> (!x86.amx.tile<16x16xf32>, !x86.amx.tile<16x16xf32>, !x86.amx.tile<16x16xf32>, !x86.amx.tile<16x16xf32>) {
+// CHECK: vector.transfer_read {{.*}} vector<32xbf16>
+// CHECK: vector.transfer_read {{.*}} vector<32xbf16>
// CHECK: vector.shuffle{{.*}}[0, 32, 1, 33, 2, 34, 3, 35, 8, 40, 9, 41, 10, 42, 11, 43, 16, 48, 17, 49, 18, 50, 19, 51, 24, 56, 25, 57, 26, 58, 27, 59] : vector<32xbf16>, vector<32xbf16>
// CHECK-NEXT: vector.shuffle{{.*}}[4, 36, 5, 37, 6, 38, 7, 39, 12, 44, 13, 45, 14, 46, 15, 47, 20, 52, 21, 53, 22, 54, 23, 55, 28, 60, 29, 61, 30, 62, 31, 63] : vector<32xbf16>, vector<32xbf16>
// CHECK: x86.amx.tile_load
diff --git a/mlir/test/Dialect/X86/vector-contract-to-packed-type-dotproduct.mlir b/mlir/test/Dialect/X86/vector-contract-to-packed-type-dotproduct.mlir
index e454e88f163d6..5d32570b21cbe 100644
--- a/mlir/test/Dialect/X86/vector-contract-to-packed-type-dotproduct.mlir
+++ b/mlir/test/Dialect/X86/vector-contract-to-packed-type-dotproduct.mlir
@@ -691,6 +691,8 @@ func.func @brmatmul_bf16dp_flat_layout_loop(%arg0: memref<16x64x32xbf16>, %arg1:
// CHECK-NEXT: vector.shuffle{{.*}}[8, 9, 10, 11, 24, 25, 26, 27, 12, 13, 14, 15, 28, 29, 30, 31] : vector<16xf32>, vector<16xf32>
// CHECK: scf.for
// CHECK: scf.for
+// CHECK: vector.transfer_read {{.*}} vector<32xbf16>
+// CHECK: vector.transfer_read {{.*}} vector<32xbf16>
// CHECK: vector.shuffle{{.*}}[0, 32, 1, 33, 2, 34, 3, 35, 8, 40, 9, 41, 10, 42, 11, 43, 16, 48, 17, 49, 18, 50, 19, 51, 24, 56, 25, 57, 26, 58, 27, 59] : vector<32xbf16>, vector<32xbf16>
// CHECK-NEXT: vector.shuffle{{.*}}[4, 36, 5, 37, 6, 38, 7, 39, 12, 44, 13, 45, 14, 46, 15, 47, 20, 52, 21, 53, 22, 54, 23, 55, 28, 60, 29, 61, 30, 62, 31, 63] : vector<32xbf16>, vector<32xbf16>
// CHECK: x86.avx512.dot
>From 134bc68897a1cbc3008fa072b70bac04e0aa622a Mon Sep 17 00:00:00 2001
From: Arun Thangamani <arun.thangamani at intel.com>
Date: Thu, 4 Jun 2026 08:28:49 -0700
Subject: [PATCH 4/5] added comments for transfer_read
---
.../VectorContractToAMXDotProduct.cpp | 40 ++++++++++---------
1 file changed, 21 insertions(+), 19 deletions(-)
diff --git a/mlir/lib/Dialect/X86/Transforms/VectorContractToAMXDotProduct.cpp b/mlir/lib/Dialect/X86/Transforms/VectorContractToAMXDotProduct.cpp
index 30628500e4bad..af25f1b8a84c9 100644
--- a/mlir/lib/Dialect/X86/Transforms/VectorContractToAMXDotProduct.cpp
+++ b/mlir/lib/Dialect/X86/Transforms/VectorContractToAMXDotProduct.cpp
@@ -264,18 +264,20 @@ static void performShuffle(OpBuilder &rewriter, Location loc, Value matB,
ValueRange iterArgs) {
subviewOffset[subviewOffset.size() - 2] = iv;
- auto ipVectorType = VectorType::get({2, (16 * (offset / 2))}, ipType);
+ // Retrive two rows of vector (32) for int8 and f8 type. For bf16,
+ // retrive one row of vector (32).
+ auto vectorType = VectorType::get({2, (16 * (offset / 2))}, ipType);
if (ipType.isBF16())
- ipVectorType = VectorType::get((16 * offset), ipType);
+ vectorType = VectorType::get((16 * offset), ipType);
int64_t srcRank = (dyn_cast<ShapedType>(matB.getType())).getRank();
Value padding = ub::PoisonOp::create(rewriter, loc, ipType);
- auto map = AffineMap::getMinorIdentityMap(
- srcRank, ipVectorType.getRank(), rewriter.getContext());
- SmallVector<bool> inBounds(ipVectorType.getRank(), true);
+ auto map = AffineMap::getMinorIdentityMap(srcRank, vectorType.getRank(),
+ rewriter.getContext());
+ SmallVector<bool> inBounds(vectorType.getRank(), true);
Value vec1 = vector::TransferReadOp::create(
- rewriter, loc, ipVectorType, matB, ValueRange(subviewOffset),
- padding, map, inBounds);
+ rewriter, loc, vectorType, matB, ValueRange(subviewOffset), padding,
+ map, inBounds);
if (!ipType.isBF16())
vec1 = vector::ShapeCastOp::create(
@@ -287,8 +289,8 @@ static void performShuffle(OpBuilder &rewriter, Location loc, Value matB,
subviewOffset[subviewOffset.size() - 2] = incIV;
Value vec2 = vector::TransferReadOp::create(
- rewriter, loc, ipVectorType, matB, ValueRange(subviewOffset),
- padding, map, inBounds);
+ rewriter, loc, vectorType, matB, ValueRange(subviewOffset), padding,
+ map, inBounds);
if (!ipType.isBF16())
vec2 = vector::ShapeCastOp::create(
rewriter, loc, VectorType::get((16 * offset), ipType), vec2);
@@ -991,17 +993,17 @@ struct VectorContractToAMXDotProduct
amx::TileStoreOp::create(rewriter, loc, resultBuffer, ValueRange{c0, c0},
dp);
- auto ipVectorType = mlir::VectorType::get({16, 16}, opType);
+ auto vectorType = mlir::VectorType::get({16, 16}, opType);
int64_t srcRank =
(dyn_cast<ShapedType>(resultBuffer.getType())).getRank();
Value padding = ub::PoisonOp::create(rewriter, loc, opType);
- auto map = AffineMap::getMinorIdentityMap(srcRank, ipVectorType.getRank(),
+ auto map = AffineMap::getMinorIdentityMap(srcRank, vectorType.getRank(),
rewriter.getContext());
- SmallVector<bool> inBounds(ipVectorType.getRank(), true);
+ SmallVector<bool> inBounds(vectorType.getRank(), true);
Value vecRow = vector::TransferReadOp::create(
- rewriter, loc, ipVectorType, resultBuffer, ValueRange{c0, c0},
- padding, map, inBounds);
+ rewriter, loc, vectorType, resultBuffer, ValueRange{c0, c0}, padding,
+ map, inBounds);
Value resultOp = contractionUsersAfterYield(contractOp.getResult());
if (auto vecType = llvm::dyn_cast<VectorType>(resultOp.getType()))
@@ -1479,17 +1481,17 @@ struct VectorContractToAMXDotProduct
Value indexOp_i = arith::ConstantIndexOp::create(rewriter, loc, i);
Value indexOp_j = arith::ConstantIndexOp::create(rewriter, loc, j);
- auto ipVectorType = mlir::VectorType::get({16, 16}, opType);
+ auto vectorType = mlir::VectorType::get({16, 16}, opType);
int64_t srcRank =
(dyn_cast<ShapedType>(resultBuffer.getType())).getRank();
Value padding = ub::PoisonOp::create(rewriter, loc, opType);
- auto map = AffineMap::getMinorIdentityMap(
- srcRank, ipVectorType.getRank(), rewriter.getContext());
- SmallVector<bool> inBounds(ipVectorType.getRank(), true);
+ auto map = AffineMap::getMinorIdentityMap(srcRank, vectorType.getRank(),
+ rewriter.getContext());
+ SmallVector<bool> inBounds(vectorType.getRank(), true);
auto vec1 = vector::TransferReadOp::create(
- rewriter, loc, ipVectorType, resultBuffer,
+ rewriter, loc, vectorType, resultBuffer,
ValueRange{indexOp_i, indexOp_j}, padding, map, inBounds);
writeResults.push_back(vec1);
}
>From 28d39f96daa876d89e49d1d378f7460ca8841074 Mon Sep 17 00:00:00 2001
From: Arun Thangamani <arun.thangamani at intel.com>
Date: Thu, 4 Jun 2026 08:40:01 -0700
Subject: [PATCH 5/5] fix typo
---
.../Dialect/X86/Transforms/VectorContractToAMXDotProduct.cpp | 4 ++--
1 file changed, 2 insertions(+), 2 deletions(-)
diff --git a/mlir/lib/Dialect/X86/Transforms/VectorContractToAMXDotProduct.cpp b/mlir/lib/Dialect/X86/Transforms/VectorContractToAMXDotProduct.cpp
index af25f1b8a84c9..bfdeb51b65ef1 100644
--- a/mlir/lib/Dialect/X86/Transforms/VectorContractToAMXDotProduct.cpp
+++ b/mlir/lib/Dialect/X86/Transforms/VectorContractToAMXDotProduct.cpp
@@ -264,8 +264,8 @@ static void performShuffle(OpBuilder &rewriter, Location loc, Value matB,
ValueRange iterArgs) {
subviewOffset[subviewOffset.size() - 2] = iv;
- // Retrive two rows of vector (32) for int8 and f8 type. For bf16,
- // retrive one row of vector (32).
+ // Retrieve two rows of vector (32) for int8 and f8 type. For bf16,
+ // retrieve one row of vector (32).
auto vectorType = VectorType::get({2, (16 * (offset / 2))}, ipType);
if (ipType.isBF16())
vectorType = VectorType::get((16 * offset), ipType);
More information about the Mlir-commits
mailing list