[Mlir-commits] [mlir] 9a9abb8 - [mlir][x86] Fix VNNI operand load offsets (#230205)
llvmlistbot at llvm.org
llvmlistbot at llvm.org
Fri Oct 9 09:55:45 PDT 2026
Author: Adam Siemieniuk
Date: 2026-10-09T18:55:38+02:00
New Revision: 9a9abb8b32de1b1ac9834d5eb0f9b2664ee95098
URL: https://github.com/llvm/llvm-project/commit/9a9abb8b32de1b1ac9834d5eb0f9b2664ee95098
DIFF: https://github.com/llvm/llvm-project/commit/9a9abb8b32de1b1ac9834d5eb0f9b2664ee95098.diff
LOG: [mlir][x86] Fix VNNI operand load offsets (#230205)
Fixes the tile load offsets of VNNI operands in the AMX contraction
lowering.
An offset into the VNNI packed dims was not scaled by the VNNI factor,
so loads at a non-zero offset read the wrong data. VNNI operands whose
buffer's innermost dim is not the static VNNI factor are now rejected,
as their tiles cannot be loaded directly.
Assisted-by: Claude
Added:
Modified:
mlir/lib/Dialect/X86/Transforms/VectorContractToAMXDotProduct.cpp
mlir/test/Dialect/X86/AMX/vector-contract-to-tiled-dp.mlir
Removed:
################################################################################
diff --git a/mlir/lib/Dialect/X86/Transforms/VectorContractToAMXDotProduct.cpp b/mlir/lib/Dialect/X86/Transforms/VectorContractToAMXDotProduct.cpp
index ec3c14fd3e7081..ef6dc9a00f0d8a 100644
--- a/mlir/lib/Dialect/X86/Transforms/VectorContractToAMXDotProduct.cpp
+++ b/mlir/lib/Dialect/X86/Transforms/VectorContractToAMXDotProduct.cpp
@@ -63,7 +63,7 @@ static Value contractionUsersAfterYield(Value v) {
return contractionUsersAfterYield(parent->getResult(idx));
}
-// Function to collapse the last two dimension (vnni and k) to help the
+// Function to collapse the last two dimension (k or n, and vnni) to help the
// amx.tile_load to correctly load the packed element type.
static Value collapseInnerDims(OpBuilder &builder, Location loc, Value input) {
ShapedType inputType = cast<ShapedType>(input.getType());
@@ -98,6 +98,38 @@ static bool isReadSrcMemref(Value operand) {
return srcBuff && isa<MemRefType>(srcBuff.getType());
}
+// Check if a vector.contract operand reads whole VNNI groups: the memref's
+// static innermost dim is the VNNI factor and the read starts at its origin.
+static bool isVnniOperandRead(Value operand, int64_t vnni) {
+ Operation *defOp = operand.getDefiningOp();
+ if (!defOp)
+ return false;
+
+ Value srcBuff;
+ OpFoldResult vnniIndex;
+ llvm::TypeSwitch<Operation *>(defOp).Case<TransferReadOp, LoadOp>(
+ [&](auto readOp) {
+ srcBuff = readOp.getBase();
+ vnniIndex = readOp.getIndices().back();
+ });
+
+ if (!srcBuff)
+ return false;
+ auto srcType = dyn_cast<MemRefType>(srcBuff.getType());
+ return srcType && srcType.getRank() > 0 &&
+ srcType.getShape().back() == vnni && isZeroInteger(vnniIndex);
+}
+
+// Replaces the indices of the two innermost dims of a VNNI operand by the
+// index of the dim they collapse into.
+static void collapseVnniIndices(OpBuilder &rewriter, Location loc,
+ SmallVectorImpl<Value> &indices, int64_t vnni) {
+ Value inner = indices.pop_back_val();
+ Value cVnni = arith::ConstantIndexOp::create(rewriter, loc, vnni);
+ Value scaled = arith::MulIOp::create(rewriter, loc, indices.back(), cVnni);
+ indices.back() = arith::AddIOp::create(rewriter, loc, scaled, inner);
+}
+
// Get the MemRef source and offset index for the operands of
// vector.contract.
static FailureOr<std::pair<Value, SmallVector<Value>>>
@@ -119,9 +151,6 @@ getSrcIndxValue(OpBuilder &rewriter, Location loc, Value operand,
if (!srcBuff || !isa<MemRefType>(srcBuff.getType()))
return failure();
- if (isNotAcc)
- indexVals.pop_back();
-
SmallVector<Value> indices;
indices.reserve(indexVals.size());
@@ -131,6 +160,8 @@ getSrcIndxValue(OpBuilder &rewriter, Location loc, Value operand,
}
if (isNotAcc) {
+ collapseVnniIndices(rewriter, loc, indices,
+ cast<MemRefType>(srcBuff.getType()).getShape().back());
srcBuff = collapseInnerDims(rewriter, loc, srcBuff);
}
@@ -255,20 +286,13 @@ static unsigned getIndexPosition(Value operand, scf::ForOp loop) {
// Creates amx.tile_loads.
static amx::TileLoadOp createTileLoads(OpBuilder &rewriter, Location loc,
Value operand, Value mat, Type ipType,
- bool rhs, unsigned int offset,
- bool isVnni) {
+ unsigned int offset, bool isVnni) {
auto srcIndx = getSrcIndxValue(rewriter, loc, operand, false);
auto [srcBuff, indices] = *srcIndx;
- if (isVnni) {
- indices.pop_back();
- }
-
- if (rhs && isVnni) {
- auto cOffset = arith::ConstantIndexOp::create(rewriter, loc, offset);
- indices[indices.size() - 1] = arith::MulIOp::create(
- rewriter, loc, indices[indices.size() - 1], cOffset);
- }
+ if (isVnni)
+ collapseVnniIndices(rewriter, loc, indices,
+ cast<MemRefType>(srcBuff.getType()).getShape().back());
amx::TileType tileType = amx::TileType::get({16, (16 * offset)}, ipType);
return amx::TileLoadOp::create(rewriter, loc, tileType, mat, indices);
@@ -462,7 +486,7 @@ createTiledDp(OpBuilder &rewriter, Location loc,
tilesLhs = itLhs->second;
} else {
tilesLhs = createTileLoads(rewriter, loc, ops[i].getLhs(), matA, ipType,
- false, offset, isVnni);
+ offset, isVnni);
readsToTileLoads.try_emplace(readOpLhs, tilesLhs);
}
@@ -473,7 +497,7 @@ createTiledDp(OpBuilder &rewriter, Location loc,
tilesRhs = itRhs->second;
} else {
tilesRhs = createTileLoads(rewriter, loc, ops[i].getRhs(), matB, ipType,
- true, offset, isVnni);
+ offset, isVnni);
readsToTileLoads.try_emplace(readOpRhs, tilesRhs);
}
@@ -865,6 +889,11 @@ struct VectorContractToAMXDotProduct
return rewriter.notifyMatchFailure(
contractOp, "The LHS or RHS src is not a MemRef type.");
+ if (isVnni && !(isVnniOperandRead(contractOp.getLhs(), blockingFactor) &&
+ isVnniOperandRead(contractOp.getRhs(), blockingFactor)))
+ return rewriter.notifyMatchFailure(
+ contractOp, "The LHS or RHS src does not read whole VNNI groups.");
+
unsigned int dimValue = blockingFactor;
if (!isVnni)
dimValue = 16 * blockingFactor;
@@ -1135,6 +1164,12 @@ struct VectorContractToAMXDotProduct
for (Operation &op : loopLists[0].getBody()->getOperations()) {
if (auto contract = dyn_cast<vector::ContractionOp>(op)) {
+ if (isVnni && !(isVnniOperandRead(contract.getLhs(), blockingFactor) &&
+ isVnniOperandRead(contract.getRhs(), blockingFactor)))
+ return rewriter.notifyMatchFailure(
+ contractOp, "The LHS or RHS src of an associated contract "
+ "operation does not read whole VNNI groups.");
+
LogicalResult validate = validateContractOps(
rewriter, contract, dimValue, srcBuffLhs, srcBuffRhs, true);
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 ac31f5f0785c75..c8c17658e3f737 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
@@ -108,6 +108,180 @@ module attributes {transform.with_named_sequence} {
// -----
+// The indices (i, vnni_index) of the two innermost dims of the operands map to
+// i * vnni + vnni_index in the collapsed buffers.
+
+!vecA = vector<16x16x4xi8>
+!vecB = vector<16x16x4xi8>
+!vecC = vector<16x16xi32>
+!memrefA = memref<16x32x4xi8>
+!memrefB = memref<32x32x4xi8>
+!memrefC = memref<16x32xi32>
+#map = affine_map<(d4, d1, d2, d3) -> (d1, d3, d4)>
+#map1 = affine_map<(d4, d1, d2, d3) -> (d3, d2, d4)>
+#map2 = affine_map<(d4, d1, d2, d3) -> (d1, d2)>
+func.func @matmul_int8_vnni_offsets(
+ %arg0: !memrefA, %arg1: !memrefB, %arg2: !memrefC) -> !memrefC
+{
+ %c0 = arith.constant 0 : index
+ %c16 = arith.constant 16 : index
+ %0 = ub.poison : i8
+ %32 = ub.poison : i32
+
+ %1 = vector.transfer_read %arg0[%c0, %c16, %c0], %0 {in_bounds = [true, true, true]} :
+ !memrefA, !vecA
+ %2 = vector.transfer_read %arg1[%c16, %c16, %c0], %0 {in_bounds = [true, true, true]} :
+ !memrefB, !vecB
+
+ %3 = vector.transfer_read %arg2[%c0, %c16], %32 {in_bounds = [true, true]} : !memrefC, !vecC
+
+ %4 = vector.contract {
+ indexing_maps = [#map, #map1, #map2],
+ iterator_types = ["reduction", "parallel", "parallel", "reduction"],
+ kind = #vector.kind<add>}
+ %1, %2, %3 : !vecA, !vecB into !vecC
+
+ vector.transfer_write %4, %arg2[%c0, %c16] {in_bounds = [true, true]} : !vecC, !memrefC
+
+ return %arg2 : !memrefC
+}
+
+// CHECK-LABEL: @matmul_int8_vnni_offsets
+// CHECK-DAG: %[[C0:.+]] = arith.constant 0 : index
+// CHECK-DAG: %[[C16:.+]] = arith.constant 16 : index
+// CHECK-DAG: %[[C64:.+]] = arith.constant 64 : index
+// CHECK-DAG: %[[LHS:.+]] = memref.collapse_shape %arg0 {{\[}}[0], [1, 2]]
+// CHECK-DAG: %[[RHS:.+]] = memref.collapse_shape %arg1 {{\[}}[0], [1, 2]]
+// CHECK: x86.amx.tile_load %[[LHS]][%[[C0]], %[[C64]]] {{.*}} !x86.amx.tile<16x64xi8>
+// CHECK: x86.amx.tile_load %[[RHS]][%[[C16]], %[[C64]]] {{.*}} !x86.amx.tile<16x64xi8>
+// CHECK: x86.amx.tile_load %arg2[%[[C0]], %[[C16]]] {{.*}} !x86.amx.tile<16x16xi32>
+// CHECK: x86.amx.tile_muli
+// CHECK: vector.transfer_write %{{.+}}, %arg2[%[[C0]], %[[C16]]]
+// CHECK-NOT: vector.contract
+
+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_amx_dot_product
+ } : !transform.any_op
+ transform.yield
+ }
+}
+
+// -----
+
+// The innermost dim of the VNNI operand buffers must be the static VNNI factor.
+
+!vecA = vector<16x16x4xi8>
+!vecB = vector<16x16x4xi8>
+!vecC = vector<16x16xi32>
+!memrefA = memref<16x32x?xi8>
+!memrefB = memref<32x32x4xi8>
+!memrefC = memref<16x32xi32>
+#map = affine_map<(d4, d1, d2, d3) -> (d1, d3, d4)>
+#map1 = affine_map<(d4, d1, d2, d3) -> (d3, d2, d4)>
+#map2 = affine_map<(d4, d1, d2, d3) -> (d1, d2)>
+func.func @negative_matmul_int8_dynamic_vnni_dim(
+ %arg0: !memrefA, %arg1: !memrefB, %arg2: !memrefC) -> !memrefC
+{
+ %c0 = arith.constant 0 : index
+ %0 = ub.poison : i8
+ %32 = ub.poison : i32
+
+ %1 = vector.transfer_read %arg0[%c0, %c0, %c0], %0 {in_bounds = [true, true, true]} :
+ !memrefA, !vecA
+ %2 = vector.transfer_read %arg1[%c0, %c0, %c0], %0 {in_bounds = [true, true, true]} :
+ !memrefB, !vecB
+
+ %3 = vector.transfer_read %arg2[%c0, %c0], %32 {in_bounds = [true, true]} : !memrefC, !vecC
+
+ %4 = vector.contract {
+ indexing_maps = [#map, #map1, #map2],
+ iterator_types = ["reduction", "parallel", "parallel", "reduction"],
+ kind = #vector.kind<add>}
+ %1, %2, %3 : !vecA, !vecB into !vecC
+
+ vector.transfer_write %4, %arg2[%c0, %c0] {in_bounds = [true, true]} : !vecC, !memrefC
+
+ return %arg2 : !memrefC
+}
+
+// CHECK-LABEL: @negative_matmul_int8_dynamic_vnni_dim
+// CHECK-NOT: x86.amx
+// CHECK: vector.contract
+
+!memrefBWide = memref<32x32x8xi8>
+func.func @negative_matmul_int8_wide_vnni_dim(
+ %arg0: !memrefB, %arg1: !memrefBWide, %arg2: !memrefC) -> !memrefC
+{
+ %c0 = arith.constant 0 : index
+ %0 = ub.poison : i8
+ %32 = ub.poison : i32
+
+ %1 = vector.transfer_read %arg0[%c0, %c0, %c0], %0 {in_bounds = [true, true, true]} :
+ !memrefB, !vecA
+ %2 = vector.transfer_read %arg1[%c0, %c0, %c0], %0 {in_bounds = [true, true, true]} :
+ !memrefBWide, !vecB
+
+ %3 = vector.transfer_read %arg2[%c0, %c0], %32 {in_bounds = [true, true]} : !memrefC, !vecC
+
+ %4 = vector.contract {
+ indexing_maps = [#map, #map1, #map2],
+ iterator_types = ["reduction", "parallel", "parallel", "reduction"],
+ kind = #vector.kind<add>}
+ %1, %2, %3 : !vecA, !vecB into !vecC
+
+ vector.transfer_write %4, %arg2[%c0, %c0] {in_bounds = [true, true]} : !vecC, !memrefC
+
+ return %arg2 : !memrefC
+}
+
+// CHECK-LABEL: @negative_matmul_int8_wide_vnni_dim
+// CHECK-NOT: x86.amx
+// CHECK: vector.contract
+
+func.func @negative_matmul_int8_vnni_offset(
+ %arg0: !memrefB, %arg1: !memrefB, %arg2: !memrefC) -> !memrefC
+{
+ %c0 = arith.constant 0 : index
+ %c1 = arith.constant 1 : index
+ %0 = ub.poison : i8
+ %32 = ub.poison : i32
+
+ %1 = vector.load %arg0[%c0, %c0, %c1] : !memrefB, !vecA
+ %2 = vector.transfer_read %arg1[%c0, %c0, %c0], %0 {in_bounds = [true, true, true]} :
+ !memrefB, !vecB
+
+ %3 = vector.transfer_read %arg2[%c0, %c0], %32 {in_bounds = [true, true]} : !memrefC, !vecC
+
+ %4 = vector.contract {
+ indexing_maps = [#map, #map1, #map2],
+ iterator_types = ["reduction", "parallel", "parallel", "reduction"],
+ kind = #vector.kind<add>}
+ %1, %2, %3 : !vecA, !vecB into !vecC
+
+ vector.transfer_write %4, %arg2[%c0, %c0] {in_bounds = [true, true]} : !vecC, !memrefC
+
+ return %arg2 : !memrefC
+}
+
+// CHECK-LABEL: @negative_matmul_int8_vnni_offset
+// CHECK-NOT: x86.amx
+// CHECK: vector.contract
+
+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_amx_dot_product
+ } : !transform.any_op
+ transform.yield
+ }
+}
+
+// -----
+
!vecA = vector<1x16x16x2xbf16>
!vecB = vector<1x16x16x2xbf16>
!vecC = vector<16x16xf32>
@@ -693,6 +867,259 @@ module attributes {transform.with_named_sequence} {
// -----
+// The indices (i, vnni_index) of the two innermost dims of the operands map to
+// i * vnni + vnni_index in the collapsed buffers.
+
+!vecAB = vector<16x16x4xi8>
+!vecC = vector<16x16xi32>
+!memrefA = memref<16x32x4xi8, strided<[256, 4, 1], offset: ?>>
+!memrefB = memref<16x32x4xi8, strided<[512, 4, 1], offset: ?>>
+!memrefC = memref<16x32xi32, strided<[128, 1], offset: ?>>
+
+#map = affine_map<(d0, d1, d2, d3) -> (d1, d3, d0)>
+#map1 = affine_map<(d0, d1, d2, d3) -> (d3, d2, d0)>
+#map2 = affine_map<(d0, d1, d2, d3) -> (d1, d2)>
+
+func.func @matmul_int8_loop_vnni_offsets(%arg0: memref<64x64x4xi8>, %arg1: memref<64x128x4xi8>, %arg2: memref<64x128xi32>) {
+ %0 = ub.poison : i32
+ %1 = ub.poison : i8
+ %c0 = arith.constant 0 : index
+ %c64 = arith.constant 64 : index
+ %c128 = arith.constant 128 : index
+ %c16 = arith.constant 16 : index
+ %c32 = arith.constant 32 : index
+ scf.for %arg3 = %c0 to %c64 step %c16 {
+ scf.for %arg4 = %c0 to %c128 step %c32 {
+
+ %subview = memref.subview %arg2[%arg3, %arg4] [16, 32] [1, 1] :
+ memref<64x128xi32> to !memrefC
+ %2 = vector.transfer_read %subview[%c0, %c0], %0 {in_bounds = [true, true]} :
+ !memrefC, !vecC
+ %3 = vector.transfer_read %subview[%c0, %c16], %0 {in_bounds = [true, true]} :
+ !memrefC, !vecC
+
+ %4:2 = scf.for %arg5 = %c0 to %c32 step %c16 iter_args(%arg6 = %2, %arg7 = %3) -> (!vecC, !vecC) {
+
+ %subview_0 = memref.subview %arg0[%arg3, %arg5, 0] [16, 32, 4] [1, 1, 1] :
+ memref<64x64x4xi8> to !memrefA
+ %subview_1 = memref.subview %arg1[%arg5, %arg4, 0] [16, 32, 4] [1, 1, 1] :
+ memref<64x128x4xi8> to !memrefB
+ %5 = vector.transfer_read %subview_0[%c0, %c16, %c0], %1 {in_bounds = [true, true, true]} :
+ !memrefA, !vecAB
+ %6 = vector.transfer_read %subview_1[%c0, %c0, %c0], %1 {in_bounds = [true, true, true]} :
+ !memrefB, !vecAB
+ %7 = vector.transfer_read %subview_1[%c0, %c16, %c0], %1 {in_bounds = [true, true, true]} :
+ !memrefB, !vecAB
+
+ %8 = vector.contract {indexing_maps = [#map, #map1, #map2], iterator_types =
+ ["reduction", "parallel", "parallel", "reduction"], kind = #vector.kind<add>}
+ %5, %6, %arg6 {unroll_shape = array<i64: 4, 16, 16, 16>} : !vecAB, !vecAB into !vecC
+ %9 = vector.contract {indexing_maps = [#map, #map1, #map2], iterator_types =
+ ["reduction", "parallel", "parallel", "reduction"], kind = #vector.kind<add>}
+ %5, %7, %arg7 {unroll_shape = array<i64: 4, 16, 16, 16>} : !vecAB, !vecAB into !vecC
+
+ scf.yield %8, %9 : !vecC, !vecC
+ }
+
+ vector.transfer_write %4#1, %subview[%c0, %c16] {in_bounds = [true, true]} :
+ !vecC, !memrefC
+ vector.transfer_write %4#0, %subview[%c0, %c0] {in_bounds = [true, true]} :
+ !vecC, !memrefC
+ }
+ }
+
+ return
+}
+
+// CHECK-LABEL: @matmul_int8_loop_vnni_offsets
+// CHECK-DAG: %[[C0:.+]] = arith.constant 0 : index
+// CHECK-DAG: %[[C64:.+]] = arith.constant 64 : index
+// CHECK-DAG: %[[LHS:.+]] = memref.collapse_shape %{{.+}} {{\[}}[0], [1, 2]] : memref<16x32x4xi8, strided<[256, 4, 1], offset: ?>> into memref<16x128xi8
+// CHECK-DAG: %[[RHS:.+]] = memref.collapse_shape %{{.+}} {{\[}}[0], [1, 2]] : memref<16x32x4xi8, strided<[512, 4, 1], offset: ?>> into memref<16x128xi8
+// CHECK: x86.amx.tile_load %[[LHS]][%[[C0]], %[[C64]]] {{.*}} !x86.amx.tile<16x64xi8>
+// CHECK: x86.amx.tile_load %[[RHS]][%[[C0]], %[[C0]]] {{.*}} !x86.amx.tile<16x64xi8>
+// CHECK: x86.amx.tile_muli
+// CHECK: x86.amx.tile_load %[[RHS]][%[[C0]], %[[C64]]] {{.*}} !x86.amx.tile<16x64xi8>
+// CHECK: x86.amx.tile_muli
+// CHECK-NOT: vector.contract
+
+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_amx_dot_product
+ } : !transform.any_op
+ transform.yield
+ }
+}
+
+// -----
+
+// All the contractions in the loop must read whole VNNI groups.
+
+!vecAB = vector<16x16x4xi8>
+!vecC = vector<16x16xi32>
+!memrefA = memref<16x32x4xi8, strided<[256, 4, 1], offset: ?>>
+!memrefB = memref<16x32x4xi8, strided<[512, 4, 1], offset: ?>>
+!memrefC = memref<16x32xi32, strided<[128, 1], offset: ?>>
+
+#map = affine_map<(d0, d1, d2, d3) -> (d1, d3, d0)>
+#map1 = affine_map<(d0, d1, d2, d3) -> (d3, d2, d0)>
+#map2 = affine_map<(d0, d1, d2, d3) -> (d1, d2)>
+
+func.func @negative_matmul_int8_loop_vnni_offset(%arg0: memref<64x64x4xi8>, %arg1: memref<64x128x4xi8>, %arg2: memref<64x128xi32>) {
+ %0 = ub.poison : i32
+ %1 = ub.poison : i8
+ %c0 = arith.constant 0 : index
+ %c1 = arith.constant 1 : index
+ %c64 = arith.constant 64 : index
+ %c128 = arith.constant 128 : index
+ %c16 = arith.constant 16 : index
+ %c32 = arith.constant 32 : index
+ scf.for %arg3 = %c0 to %c64 step %c16 {
+ scf.for %arg4 = %c0 to %c128 step %c32 {
+
+ %subview = memref.subview %arg2[%arg3, %arg4] [16, 32] [1, 1] :
+ memref<64x128xi32> to !memrefC
+ %2 = vector.transfer_read %subview[%c0, %c0], %0 {in_bounds = [true, true]} :
+ !memrefC, !vecC
+ %3 = vector.transfer_read %subview[%c0, %c16], %0 {in_bounds = [true, true]} :
+ !memrefC, !vecC
+
+ %4:2 = scf.for %arg5 = %c0 to %c32 step %c16 iter_args(%arg6 = %2, %arg7 = %3) -> (!vecC, !vecC) {
+
+ %subview_0 = memref.subview %arg0[%arg3, %arg5, 0] [16, 32, 4] [1, 1, 1] :
+ memref<64x64x4xi8> to !memrefA
+ %subview_1 = memref.subview %arg1[%arg5, %arg4, 0] [16, 32, 4] [1, 1, 1] :
+ memref<64x128x4xi8> to !memrefB
+ %5 = vector.transfer_read %subview_0[%c0, %c0, %c0], %1 {in_bounds = [true, true, true]} :
+ !memrefA, !vecAB
+ %6 = vector.transfer_read %subview_1[%c0, %c0, %c0], %1 {in_bounds = [true, true, true]} :
+ !memrefB, !vecAB
+ %7 = vector.transfer_read %subview_1[%c0, %c16, %c1], %1 {in_bounds = [true, true, false]} :
+ !memrefB, !vecAB
+
+ %8 = vector.contract {indexing_maps = [#map, #map1, #map2], iterator_types =
+ ["reduction", "parallel", "parallel", "reduction"], kind = #vector.kind<add>}
+ %5, %6, %arg6 : !vecAB, !vecAB into !vecC
+ %9 = vector.contract {indexing_maps = [#map, #map1, #map2], iterator_types =
+ ["reduction", "parallel", "parallel", "reduction"], kind = #vector.kind<add>}
+ %5, %7, %arg7 : !vecAB, !vecAB into !vecC
+
+ scf.yield %8, %9 : !vecC, !vecC
+ }
+
+ vector.transfer_write %4#1, %subview[%c0, %c16] {in_bounds = [true, true]} :
+ !vecC, !memrefC
+ vector.transfer_write %4#0, %subview[%c0, %c0] {in_bounds = [true, true]} :
+ !vecC, !memrefC
+ }
+ }
+
+ return
+}
+
+// CHECK-LABEL: @negative_matmul_int8_loop_vnni_offset
+// CHECK-NOT: x86.amx
+// CHECK: vector.contract
+
+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_amx_dot_product
+ } : !transform.any_op
+ transform.yield
+ }
+}
+
+// -----
+
+!vecAB = vector<16x16x2xbf16>
+!vecC = vector<16x16xf32>
+!memrefA = memref<16x32x2xbf16, strided<[128, 2, 1], offset: ?>>
+!memrefB = memref<16x32x2xbf16, strided<[256, 2, 1], offset: ?>>
+!memrefC = memref<16x32xf32, strided<[128, 1], offset: ?>>
+
+#map = affine_map<(d0, d1, d2, d3) -> (d1, d3, d0)>
+#map1 = affine_map<(d0, d1, d2, d3) -> (d3, d2, d0)>
+#map2 = affine_map<(d0, d1, d2, d3) -> (d1, d2)>
+
+func.func @matmul_bf16_loop_vnni_offsets(%arg0: memref<64x64x2xbf16>, %arg1: memref<64x128x2xbf16>, %arg2: memref<64x128xf32>) {
+ %0 = ub.poison : f32
+ %1 = ub.poison : bf16
+ %c0 = arith.constant 0 : index
+ %c64 = arith.constant 64 : index
+ %c128 = arith.constant 128 : index
+ %c16 = arith.constant 16 : index
+ %c32 = arith.constant 32 : index
+ scf.for %arg3 = %c0 to %c64 step %c16 {
+ scf.for %arg4 = %c0 to %c128 step %c32 {
+
+ %subview = memref.subview %arg2[%arg3, %arg4] [16, 32] [1, 1] :
+ memref<64x128xf32> to !memrefC
+ %2 = vector.transfer_read %subview[%c0, %c0], %0 {in_bounds = [true, true]} :
+ !memrefC, !vecC
+ %3 = vector.transfer_read %subview[%c0, %c16], %0 {in_bounds = [true, true]} :
+ !memrefC, !vecC
+
+ %4:2 = scf.for %arg5 = %c0 to %c32 step %c16 iter_args(%arg6 = %2, %arg7 = %3) -> (!vecC, !vecC) {
+
+ %subview_0 = memref.subview %arg0[%arg3, %arg5, 0] [16, 32, 2] [1, 1, 1] :
+ memref<64x64x2xbf16> to !memrefA
+ %subview_1 = memref.subview %arg1[%arg5, %arg4, 0] [16, 32, 2] [1, 1, 1] :
+ memref<64x128x2xbf16> to !memrefB
+ %5 = vector.transfer_read %subview_0[%c0, %c16, %c0], %1 {in_bounds = [true, true, true]} :
+ !memrefA, !vecAB
+ %6 = vector.transfer_read %subview_1[%c0, %c0, %c0], %1 {in_bounds = [true, true, true]} :
+ !memrefB, !vecAB
+ %7 = vector.transfer_read %subview_1[%c0, %c16, %c0], %1 {in_bounds = [true, true, true]} :
+ !memrefB, !vecAB
+
+ %8 = vector.contract {indexing_maps = [#map, #map1, #map2], iterator_types =
+ ["reduction", "parallel", "parallel", "reduction"], kind = #vector.kind<add>}
+ %5, %6, %arg6 : !vecAB, !vecAB into !vecC
+ %9 = vector.contract {indexing_maps = [#map, #map1, #map2], iterator_types =
+ ["reduction", "parallel", "parallel", "reduction"], kind = #vector.kind<add>}
+ %5, %7, %arg7 : !vecAB, !vecAB into !vecC
+
+ scf.yield %8, %9 : !vecC, !vecC
+ }
+
+ vector.transfer_write %4#1, %subview[%c0, %c16] {in_bounds = [true, true]} :
+ !vecC, !memrefC
+ vector.transfer_write %4#0, %subview[%c0, %c0] {in_bounds = [true, true]} :
+ !vecC, !memrefC
+ }
+ }
+
+ return
+}
+
+// CHECK-LABEL: @matmul_bf16_loop_vnni_offsets
+// CHECK-DAG: %[[C0:.+]] = arith.constant 0 : index
+// CHECK-DAG: %[[C32:.+]] = arith.constant 32 : index
+// CHECK-DAG: %[[LHS:.+]] = memref.collapse_shape %{{.+}} {{\[}}[0], [1, 2]] : memref<16x32x2xbf16, strided<[128, 2, 1], offset: ?>> into memref<16x64xbf16
+// CHECK-DAG: %[[RHS:.+]] = memref.collapse_shape %{{.+}} {{\[}}[0], [1, 2]] : memref<16x32x2xbf16, strided<[256, 2, 1], offset: ?>> into memref<16x64xbf16
+// CHECK: x86.amx.tile_load %[[LHS]][%[[C0]], %[[C32]]] {{.*}} !x86.amx.tile<16x32xbf16>
+// CHECK: x86.amx.tile_load %[[RHS]][%[[C0]], %[[C0]]] {{.*}} !x86.amx.tile<16x32xbf16>
+// CHECK: x86.amx.tile_mulf
+// CHECK: x86.amx.tile_load %[[RHS]][%[[C0]], %[[C32]]] {{.*}} !x86.amx.tile<16x32xbf16>
+// CHECK: x86.amx.tile_mulf
+// CHECK-NOT: vector.contract
+
+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_amx_dot_product
+ } : !transform.any_op
+ transform.yield
+ }
+}
+
+// -----
+
!vecAB = vector<16x16x4xi8>
!vecC = vector<16x16xi32>
!memrefA = memref<16x16x4xi8, strided<[256, 4, 1], offset: ?>>
More information about the Mlir-commits
mailing list