[Mlir-commits] [mlir] [mlir][x86] Fix - Replace `load` with `transfer_read` to support on tensor type. (PR #194543)
Arun Thangamani
llvmlistbot at llvm.org
Mon Apr 27 23:05:04 PDT 2026
https://github.com/arun-thmn updated https://github.com/llvm/llvm-project/pull/194543
>From e9b4c6edf26d1b4ab29328befa27c5816b378364 Mon Sep 17 00:00:00 2001
From: Arun Thangamani <arun.thangamani at intel.com>
Date: Mon, 27 Apr 2026 22:22:49 -0700
Subject: [PATCH 1/2] replace load with transfer_read to support on tensor
types as well
---
.../VectorContractToPackedTypeDotProduct.cpp | 95 +++++++++++--------
1 file changed, 58 insertions(+), 37 deletions(-)
diff --git a/mlir/lib/Dialect/X86/Transforms/VectorContractToPackedTypeDotProduct.cpp b/mlir/lib/Dialect/X86/Transforms/VectorContractToPackedTypeDotProduct.cpp
index 4d0c82085f8a8..6154d85d58a08 100644
--- a/mlir/lib/Dialect/X86/Transforms/VectorContractToPackedTypeDotProduct.cpp
+++ b/mlir/lib/Dialect/X86/Transforms/VectorContractToPackedTypeDotProduct.cpp
@@ -59,13 +59,12 @@ static void rewriteUses(mlir::Value oldVal, mlir::Value newVal,
// Function to convert the flat layout A or B matrix vector<32xbf16>
// into VNNI packed layout using the vpunpack operations
-static void packNonUnitDimOperandToVNNI(mlir::PatternRewriter &rewriter,
- mlir::Operation *opA,
- mlir::Operation *opB,
- mlir::vector::ContractionOp contractA,
- mlir::vector::ContractionOp contractB,
- int64_t nonUnitDimAcc,
- mlir::VectorType Ty) {
+static LogicalResult
+packNonUnitDimOperandToVNNI(mlir::PatternRewriter &rewriter,
+ mlir::Operation *opA, mlir::Operation *opB,
+ mlir::vector::ContractionOp contractA,
+ mlir::vector::ContractionOp contractB,
+ int64_t nonUnitDimAcc, mlir::VectorType Ty) {
bool opABeforeopB = opA->isBeforeInBlock(opB);
@@ -98,7 +97,22 @@ static void packNonUnitDimOperandToVNNI(mlir::PatternRewriter &rewriter,
});
});
- auto vec1 = vector::LoadOp::create(rewriter, loc, flatTy, srcBuff, indexVals);
+ int64_t srcRank;
+ if (auto memrefTy = dyn_cast<MemRefType>(srcBuff.getType())) {
+ srcRank = memrefTy.getRank();
+ }
+
+ if (auto tensorTy = dyn_cast<RankedTensorType>(srcBuff.getType())) {
+ srcRank = tensorTy.getRank();
+ }
+
+ Value padding = ub::PoisonOp::create(rewriter, loc, elemTy);
+ auto map = AffineMap::getMinorIdentityMap(srcRank, flatTy.getRank(),
+ rewriter.getContext());
+ SmallVector<bool> inBounds(flatTy.getRank(), true);
+
+ auto vec1 = vector::TransferReadOp::create(rewriter, loc, flatTy, srcBuff,
+ indexVals, padding, map, inBounds);
unsigned int offset = 1;
if (elemTy.isSignlessInteger(8))
@@ -110,7 +124,8 @@ static void packNonUnitDimOperandToVNNI(mlir::PatternRewriter &rewriter,
indexVals[indexVals.size() - 2]);
indexVals[indexVals.size() - 2] = nextIndx;
- auto vec2 = vector::LoadOp::create(rewriter, loc, flatTy, srcBuff, indexVals);
+ auto vec2 = vector::TransferReadOp::create(rewriter, loc, flatTy, srcBuff,
+ indexVals, padding, map, inBounds);
static constexpr int64_t maskLo_bf16[] = {
0, 32, 1, 33, 2, 34, 3, 35, 8, 40, 9, 41, 10, 42, 11, 43,
@@ -160,6 +175,8 @@ static void packNonUnitDimOperandToVNNI(mlir::PatternRewriter &rewriter,
rewriteUses(opA->getResult(0), newA.getResult(), contractA, rewriter);
rewriteUses(opB->getResult(0), newB.getResult(), contractB, rewriter);
+
+ return success();
}
// Implements packed type outer product contraction as a sequence
@@ -313,6 +330,38 @@ struct VectorContractToPackedTypeDotProduct
return rewriter.notifyMatchFailure(contractOp,
"Could not find a contract pair");
+ if (!isNonUnitDimOperandShuffled(nonUnitDimOperand)) {
+ Value nonUnitDimOperandPairContract = rhsHasMultipleNonUnitDims
+ ? pairContractOp.getRhs()
+ : pairContractOp.getLhs();
+
+ // Get the non-packed A or B matrix's vector<32xbf16> elements.
+ Operation *nonUnitDimReadOp =
+ traceToVectorReadLikeParentOperation(nonUnitDimOperand);
+ Operation *nonUnitDimReadOpPairContract =
+ traceToVectorReadLikeParentOperation(nonUnitDimOperandPairContract);
+
+ if (!nonUnitDimReadOp || !nonUnitDimReadOpPairContract)
+ return rewriter.notifyMatchFailure(
+ contractOp, "Could not find a valid contract pair");
+
+ VectorType nonUnitDimTy = rhsHasMultipleNonUnitDims
+ ? contractOp.getRhsType()
+ : contractOp.getLhsType();
+
+ LogicalResult vnniPacked = packNonUnitDimOperandToVNNI(
+ rewriter, nonUnitDimReadOp, nonUnitDimReadOpPairContract,
+ contractOp, pairContractOp, blockingFactor * nonUnitDimValue,
+ nonUnitDimTy);
+
+ if (failed(vnniPacked))
+ return rewriter.notifyMatchFailure(
+ contractOp, "The input source must be MemRef type.");
+
+ nonUnitDimOperand = rhsHasMultipleNonUnitDims ? contractOp.getRhs()
+ : contractOp.getLhs();
+ }
+
// Validate and shuffle the accumulator
if (accRead) {
// Trace back to the load or transfer_read operations of the contract
@@ -370,34 +419,6 @@ struct VectorContractToPackedTypeDotProduct
contractOp,
"Write to accumulator is not by transfer_write or store");
}
-
- if (!isNonUnitDimOperandShuffled(nonUnitDimOperand)) {
- Value nonUnitDimOperandPairContract = rhsHasMultipleNonUnitDims
- ? pairContractOp.getRhs()
- : pairContractOp.getLhs();
-
- // Get the non-packed A or B matrix's vector<32xbf16> elements.
- Operation *nonUnitDimReadOp =
- traceToVectorReadLikeParentOperation(nonUnitDimOperand);
- Operation *nonUnitDimReadOpPairContract =
- traceToVectorReadLikeParentOperation(nonUnitDimOperandPairContract);
-
- if (!nonUnitDimReadOp || !nonUnitDimReadOpPairContract)
- return rewriter.notifyMatchFailure(
- contractOp, "Could not find a valid contract pair");
-
- VectorType nonUnitDimTy = rhsHasMultipleNonUnitDims
- ? contractOp.getRhsType()
- : contractOp.getLhsType();
-
- packNonUnitDimOperandToVNNI(
- rewriter, nonUnitDimReadOp, nonUnitDimReadOpPairContract,
- contractOp, pairContractOp, blockingFactor * nonUnitDimValue,
- nonUnitDimTy);
-
- nonUnitDimOperand = rhsHasMultipleNonUnitDims ? contractOp.getRhs()
- : contractOp.getLhs();
- }
}
rewriter.setInsertionPoint(contractOp);
>From a1bdae260f7ae6b3a31cb29401f865699eab2df8 Mon Sep 17 00:00:00 2001
From: Arun Thangamani <arun.thangamani at intel.com>
Date: Mon, 27 Apr 2026 23:04:50 -0700
Subject: [PATCH 2/2] unit-test for tensor tyope.
---
.../VectorContractToPackedTypeDotProduct.cpp | 75 ++++++-------
...or-contract-to-packed-type-dotproduct.mlir | 105 ++++++++++++++++++
2 files changed, 140 insertions(+), 40 deletions(-)
diff --git a/mlir/lib/Dialect/X86/Transforms/VectorContractToPackedTypeDotProduct.cpp b/mlir/lib/Dialect/X86/Transforms/VectorContractToPackedTypeDotProduct.cpp
index 6154d85d58a08..dc2a764c39a98 100644
--- a/mlir/lib/Dialect/X86/Transforms/VectorContractToPackedTypeDotProduct.cpp
+++ b/mlir/lib/Dialect/X86/Transforms/VectorContractToPackedTypeDotProduct.cpp
@@ -59,12 +59,13 @@ static void rewriteUses(mlir::Value oldVal, mlir::Value newVal,
// Function to convert the flat layout A or B matrix vector<32xbf16>
// into VNNI packed layout using the vpunpack operations
-static LogicalResult
-packNonUnitDimOperandToVNNI(mlir::PatternRewriter &rewriter,
- mlir::Operation *opA, mlir::Operation *opB,
- mlir::vector::ContractionOp contractA,
- mlir::vector::ContractionOp contractB,
- int64_t nonUnitDimAcc, mlir::VectorType Ty) {
+static void packNonUnitDimOperandToVNNI(mlir::PatternRewriter &rewriter,
+ mlir::Operation *opA,
+ mlir::Operation *opB,
+ mlir::vector::ContractionOp contractA,
+ mlir::vector::ContractionOp contractB,
+ int64_t nonUnitDimAcc,
+ mlir::VectorType Ty) {
bool opABeforeopB = opA->isBeforeInBlock(opB);
@@ -175,8 +176,6 @@ packNonUnitDimOperandToVNNI(mlir::PatternRewriter &rewriter,
rewriteUses(opA->getResult(0), newA.getResult(), contractA, rewriter);
rewriteUses(opB->getResult(0), newB.getResult(), contractB, rewriter);
-
- return success();
}
// Implements packed type outer product contraction as a sequence
@@ -330,38 +329,6 @@ struct VectorContractToPackedTypeDotProduct
return rewriter.notifyMatchFailure(contractOp,
"Could not find a contract pair");
- if (!isNonUnitDimOperandShuffled(nonUnitDimOperand)) {
- Value nonUnitDimOperandPairContract = rhsHasMultipleNonUnitDims
- ? pairContractOp.getRhs()
- : pairContractOp.getLhs();
-
- // Get the non-packed A or B matrix's vector<32xbf16> elements.
- Operation *nonUnitDimReadOp =
- traceToVectorReadLikeParentOperation(nonUnitDimOperand);
- Operation *nonUnitDimReadOpPairContract =
- traceToVectorReadLikeParentOperation(nonUnitDimOperandPairContract);
-
- if (!nonUnitDimReadOp || !nonUnitDimReadOpPairContract)
- return rewriter.notifyMatchFailure(
- contractOp, "Could not find a valid contract pair");
-
- VectorType nonUnitDimTy = rhsHasMultipleNonUnitDims
- ? contractOp.getRhsType()
- : contractOp.getLhsType();
-
- LogicalResult vnniPacked = packNonUnitDimOperandToVNNI(
- rewriter, nonUnitDimReadOp, nonUnitDimReadOpPairContract,
- contractOp, pairContractOp, blockingFactor * nonUnitDimValue,
- nonUnitDimTy);
-
- if (failed(vnniPacked))
- return rewriter.notifyMatchFailure(
- contractOp, "The input source must be MemRef type.");
-
- nonUnitDimOperand = rhsHasMultipleNonUnitDims ? contractOp.getRhs()
- : contractOp.getLhs();
- }
-
// Validate and shuffle the accumulator
if (accRead) {
// Trace back to the load or transfer_read operations of the contract
@@ -419,6 +386,34 @@ struct VectorContractToPackedTypeDotProduct
contractOp,
"Write to accumulator is not by transfer_write or store");
}
+
+ if (!isNonUnitDimOperandShuffled(nonUnitDimOperand)) {
+ Value nonUnitDimOperandPairContract = rhsHasMultipleNonUnitDims
+ ? pairContractOp.getRhs()
+ : pairContractOp.getLhs();
+
+ // Get the non-packed A or B matrix's vector<32xbf16> elements.
+ Operation *nonUnitDimReadOp =
+ traceToVectorReadLikeParentOperation(nonUnitDimOperand);
+ Operation *nonUnitDimReadOpPairContract =
+ traceToVectorReadLikeParentOperation(nonUnitDimOperandPairContract);
+
+ if (!nonUnitDimReadOp || !nonUnitDimReadOpPairContract)
+ return rewriter.notifyMatchFailure(
+ contractOp, "Could not find a valid contract pair");
+
+ VectorType nonUnitDimTy = rhsHasMultipleNonUnitDims
+ ? contractOp.getRhsType()
+ : contractOp.getLhsType();
+
+ packNonUnitDimOperandToVNNI(
+ rewriter, nonUnitDimReadOp, nonUnitDimReadOpPairContract,
+ contractOp, pairContractOp, blockingFactor * nonUnitDimValue,
+ nonUnitDimTy);
+
+ nonUnitDimOperand = rhsHasMultipleNonUnitDims ? contractOp.getRhs()
+ : contractOp.getLhs();
+ }
}
rewriter.setInsertionPoint(contractOp);
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 f861d357739a3..2efe0f2adc413 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
@@ -874,6 +874,111 @@ func.func @brgemm_int8_flat_avx10(%arg0: memref<16x64x256xi8>, %arg1: memref<16x
// CHECK-NEXT: vector.shuffle{{.*}}[8, 9, 10, 11, 24, 25, 26, 27, 12, 13, 14, 15, 28, 29, 30, 31] : vector<16xi32>, vector<16xi32>
+module attributes {transform.with_named_sequence} {
+ transform.named_sequence @__transform_main(%arg1: !transform.any_op {transform.readonly}) {
+ %func = transform.structured.match ops{["func.func"]} in %arg1 : (!transform.any_op) -> !transform.any_op
+ transform.apply_patterns to %func {
+ transform.apply_patterns.x86.vector_contract_to_packed_type_dot_product
+ } : !transform.any_op
+ transform.yield
+ }
+}
+
+// -----
+
+!vecA = vector<1x1x2xbf16>
+!vecB = vector<1x2x16xbf16>
+!vecC = vector<1x16xf32>
+
+#map = affine_map<(d0, d1, d2, d3) -> (d0, d1, d3)>
+#map1 = affine_map<(d0, d1, d2, d3) -> (d0, d3, d2)>
+#map2 = affine_map<(d0, d1, d2, d3) -> (d1, d2)>
+
+func.func @brmatmul_bf16dp_flat_layout_loop_tensor_type(%arg0: tensor<16x64x32xbf16>, %arg1: tensor<16x32x64xbf16>, %arg2: tensor<64x64xf32>) -> tensor<64x64xf32> {
+ %0 = ub.poison : f32
+ %1 = ub.poison : bf16
+ %c0 = arith.constant 0 : index
+ %c64 = arith.constant 64 : index
+ %c16 = arith.constant 16 : index
+ %c32 = arith.constant 32 : index
+ %c2 = arith.constant 2 : index
+ %c1 = arith.constant 1 : index
+ %2 = scf.for %arg3 = %c0 to %c64 step %c2 iter_args(%arg4 = %arg2) -> (tensor<64x64xf32>) {
+ %3 = scf.for %arg5 = %c0 to %c64 step %c32 iter_args(%arg6 = %arg4) -> (tensor<64x64xf32>) {
+
+ %extracted_slice = tensor.extract_slice %arg6[%arg3, %arg5] [2, 32] [1, 1]
+ : tensor<64x64xf32> to tensor<2x32xf32>
+ %4 = vector.transfer_read %extracted_slice[%c0, %c0], %0 {in_bounds = [true, true]}
+ : tensor<2x32xf32>, !vecC
+ %5 = vector.transfer_read %extracted_slice[%c0, %c16], %0 {in_bounds = [true, true]}
+ : tensor<2x32xf32>, !vecC
+ %6 = vector.transfer_read %extracted_slice[%c1, %c0], %0 {in_bounds = [true, true]}
+ : tensor<2x32xf32>, !vecC
+ %7 = vector.transfer_read %extracted_slice[%c1, %c16], %0 {in_bounds = [true, true]}
+ : tensor<2x32xf32>, !vecC
+ %8:4 = scf.for %arg7 = %c0 to %c16 step %c1 iter_args(%arg8 = %4, %arg9 = %5, %arg10 = %6, %arg11 = %7) -> (!vecC, !vecC, !vecC, !vecC) {
+ %13:4 = scf.for %arg12 = %c0 to %c32 step %c2 iter_args(%arg13 = %arg8, %arg14 = %arg9, %arg15 = %arg10, %arg16 = %arg11) -> (!vecC, !vecC, !vecC, !vecC) {
+
+ %extracted_slice_0 = tensor.extract_slice %arg0[%arg7, %arg3, %arg12] [1, 2, 2] [1, 1, 1]
+ : tensor<16x64x32xbf16> to tensor<1x2x2xbf16>
+ %extracted_slice_1 = tensor.extract_slice %arg1[%arg7, %arg12, %arg5] [1, 2, 32] [1, 1, 1]
+ : tensor<16x32x64xbf16> to tensor<1x2x32xbf16>
+ %14 = vector.transfer_read %extracted_slice_0[%c0, %c0, %c0], %1 {in_bounds = [true, true, true]}
+ : tensor<1x2x2xbf16>, !vecA
+ %15 = vector.transfer_read %extracted_slice_0[%c0, %c1, %c0], %1 {in_bounds = [true, true, true]}
+ : tensor<1x2x2xbf16>, !vecA
+ %16 = vector.transfer_read %extracted_slice_1[%c0, %c0, %c0], %1 {in_bounds = [true, true, true]}
+ : tensor<1x2x32xbf16>, !vecB
+ %17 = vector.transfer_read %extracted_slice_1[%c0, %c0, %c16], %1 {in_bounds = [true, true, true]}
+ : tensor<1x2x32xbf16>, !vecB
+ %18 = vector.contract {indexing_maps = [#map, #map1, #map2], iterator_types =
+ ["reduction", "parallel", "parallel", "reduction"], kind = #vector.kind<add>} %14, %16, %arg13
+ {unroll_shape = array<i64: 1, 1, 16, 2>} : !vecA, !vecB into !vecC
+ %19 = vector.contract {indexing_maps = [#map, #map1, #map2], iterator_types =
+ ["reduction", "parallel", "parallel", "reduction"], kind = #vector.kind<add>} %14, %17, %arg14
+ {unroll_shape = array<i64: 1, 1, 16, 2>} : !vecA, !vecB into !vecC
+ %20 = vector.contract {indexing_maps = [#map, #map1, #map2], iterator_types =
+ ["reduction", "parallel", "parallel", "reduction"], kind = #vector.kind<add>} %15, %16, %arg15
+ {unroll_shape = array<i64: 1, 1, 16, 2>} : !vecA, !vecB into !vecC
+ %21 = vector.contract {indexing_maps = [#map, #map1, #map2], iterator_types =
+ ["reduction", "parallel", "parallel", "reduction"], kind = #vector.kind<add>} %15, %17, %arg16
+ {unroll_shape = array<i64: 1, 1, 16, 2>} : !vecA, !vecB into !vecC
+ scf.yield %18, %19, %20, %21 : !vecC, !vecC, !vecC, !vecC
+ }
+ scf.yield %13#0, %13#1, %13#2, %13#3 : !vecC, !vecC, !vecC, !vecC
+ }
+
+ %9 = vector.transfer_write %8#3, %extracted_slice[%c1, %c16] {in_bounds = [true, true]}
+ : !vecC, tensor<2x32xf32>
+ %10 = vector.transfer_write %8#2, %9[%c1, %c0] {in_bounds = [true, true]}
+ : !vecC, tensor<2x32xf32>
+ %11 = vector.transfer_write %8#1, %10[%c0, %c16] {in_bounds = [true, true]}
+ : !vecC, tensor<2x32xf32>
+ %12 = vector.transfer_write %8#0, %11[%c0, %c0] {in_bounds = [true, true]}
+ : !vecC, tensor<2x32xf32>
+ %inserted_slice = tensor.insert_slice %12 into %arg6[%arg3, %arg5] [2, 32] [1, 1]
+ : tensor<2x32xf32> into tensor<64x64xf32>
+ scf.yield %inserted_slice : tensor<64x64xf32>
+ }
+ scf.yield %3 : tensor<64x64xf32>
+ }
+ return %2 : tensor<64x64xf32>
+}
+
+// CHECK-LABEL: @brmatmul_bf16dp_flat_layout_loop_tensor_type
+// CHECK: vector.shuffle{{.*}}[0, 1, 2, 3, 16, 17, 18, 19, 4, 5, 6, 7, 20, 21, 22, 23] : vector<16xf32>, vector<16xf32>
+// 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.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
+// CHECK: x86.avx512.dot
+// CHECK: scf.yield
+// CHECK: scf.yield
+// CHECK: vector.shuffle{{.*}}[0, 1, 2, 3, 16, 17, 18, 19, 4, 5, 6, 7, 20, 21, 22, 23] : vector<16xf32>, vector<16xf32>
+// CHECK-NEXT: vector.shuffle{{.*}}[8, 9, 10, 11, 24, 25, 26, 27, 12, 13, 14, 15, 28, 29, 30, 31] : vector<16xf32>, vector<16xf32>
+
module attributes {transform.with_named_sequence} {
transform.named_sequence @__transform_main(%arg1: !transform.any_op {transform.readonly}) {
%func = transform.structured.match ops{["func.func"]} in %arg1 : (!transform.any_op) -> !transform.any_op
More information about the Mlir-commits
mailing list