[Mlir-commits] [mlir] 58bc8df - [mlir][x86] AMX memref source check (#206785)
llvmlistbot at llvm.org
llvmlistbot at llvm.org
Mon Jul 6 03:02:58 PDT 2026
Author: Adam Siemieniuk
Date: 2026-07-06T12:02:53+02:00
New Revision: 58bc8dff332e0d4d7609b6dbda62200355bbc75f
URL: https://github.com/llvm/llvm-project/commit/58bc8dff332e0d4d7609b6dbda62200355bbc75f
DIFF: https://github.com/llvm/llvm-project/commit/58bc8dff332e0d4d7609b6dbda62200355bbc75f.diff
LOG: [mlir][x86] AMX memref source check (#206785)
Adds extra checks to AMX lowering to ensure data is read from memrefs as
tensors are not supported by x86 dialect ops.
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 86291ae03f29e..238d03ce4f6af 100644
--- a/mlir/lib/Dialect/X86/Transforms/VectorContractToAMXDotProduct.cpp
+++ b/mlir/lib/Dialect/X86/Transforms/VectorContractToAMXDotProduct.cpp
@@ -68,6 +68,20 @@ static Value collapseInnerDims(OpBuilder &builder, mlir::Location loc,
return memref::CollapseShapeOp::create(builder, loc, input, reassociation);
}
+// Check if a vector.contract operand has a memref read source.
+static bool isReadSrcMemref(Value operand) {
+ Operation *defOp = operand.getDefiningOp();
+ if (!defOp)
+ return false;
+
+ Value srcBuff;
+ llvm::TypeSwitch<Operation *>(operand.getDefiningOp())
+ .Case<TransferReadOp, LoadOp>(
+ [&](auto readOp) { srcBuff = readOp.getOperand(0); });
+
+ return srcBuff && isa<MemRefType>(srcBuff.getType());
+}
+
// Get the MemRef source and offset index for the operands of
// vector.contract.
static FailureOr<std::pair<Value, SmallVector<Value>>>
@@ -86,7 +100,7 @@ getSrcIndxValue(OpBuilder &rewriter, Location loc, Value operand,
srcBuff = readOp.getOperand(0);
});
- if (!srcBuff)
+ if (!srcBuff || !isa<MemRefType>(srcBuff.getType()))
return failure();
if (isNotAcc)
@@ -816,6 +830,11 @@ struct VectorContractToAMXDotProduct
return rewriter.notifyMatchFailure(
contractOp, "The accumulator read is in
diff erent block.");
+ if (!(isReadSrcMemref(contractOp.getLhs()) &&
+ isReadSrcMemref(contractOp.getRhs())))
+ return rewriter.notifyMatchFailure(
+ contractOp, "The LHS or RHS src is not a MemRef type.");
+
unsigned int dimValue = blockingFactor;
if (!isVnni)
dimValue = 16 * blockingFactor;
@@ -825,6 +844,10 @@ struct VectorContractToAMXDotProduct
if (accReadOp->getBlock() == contractOp->getBlock() &&
resultWriteOp->getBlock() == contractOp->getBlock()) {
+ if (!isReadSrcMemref(contractOp.getAcc()))
+ return rewriter.notifyMatchFailure(contractOp,
+ "The ACC src is not a MemRef type.");
+
bool collapse = false;
if (isVnni)
collapse = true;
@@ -844,14 +867,14 @@ struct VectorContractToAMXDotProduct
contractOp.getLhs(), collapse);
if (failed(srcIndxLhs))
return rewriter.notifyMatchFailure(contractOp,
- "The LHS src is not a MemRef type.");
+ "Failed to get the LHS src.");
auto [srcBuffLhs, indicesLhs] = *srcIndxLhs;
auto srcIndxRhs = getSrcIndxValue(rewriter, contractOp.getLoc(),
contractOp.getRhs(), collapse);
if (failed(srcIndxRhs))
return rewriter.notifyMatchFailure(contractOp,
- "The RHS src is not a MemRef type.");
+ "Failed to get the RHS src.");
auto rhsSrc = *srcIndxRhs;
auto srcBuffRhs = rhsSrc.first;
auto indicesRhs = rhsSrc.second;
@@ -860,7 +883,7 @@ struct VectorContractToAMXDotProduct
contractOp.getAcc(), false);
if (failed(srcIndxAcc))
return rewriter.notifyMatchFailure(contractOp,
- "The ACC src is not a MemRef type.");
+ "Failed to get the ACC src.");
auto [srcBuffAcc, indicesAcc] = *srcIndxAcc;
Value c0 = arith::ConstantIndexOp::create(rewriter, loc, 0);
@@ -1043,14 +1066,14 @@ struct VectorContractToAMXDotProduct
contractOp.getLhs(), false);
if (failed(srcIndxLhs))
return rewriter.notifyMatchFailure(contractOp,
- "The LHS src is not a MemRef type.");
+ "Failed to get the LHS src.");
auto [srcBuffLhs, indicesLhs] = *srcIndxLhs;
auto srcIndxRhs = getSrcIndxValue(rewriter, contractOp.getLoc(),
contractOp.getRhs(), false);
if (failed(srcIndxRhs))
return rewriter.notifyMatchFailure(contractOp,
- "The RHS src is not a MemRef type.");
+ "Failed to get the RHS src.");
auto [srcBuffRhs, indicesRhs] = *srcIndxRhs;
Operation *vectorOpLhs;
llvm::TypeSwitch<Operation *>(contractOp.getLhs().getDefiningOp())
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 71bf62c56a6e5..ad983ee5d6ff1 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
@@ -2078,3 +2078,147 @@ module attributes {transform.with_named_sequence} {
transform.yield
}
}
+
+// -----
+
+!vecA = vector<16x16x4xi8>
+!vecB = vector<16x16x4xi8>
+!vecC = vector<16x16xi32>
+!tensorA = tensor<32x16x4xi8>
+!memrefB = memref<16x32x4xi8>
+!memrefC = memref<32x32xi32>
+#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_lhs_tensor(
+ %arg0: !tensorA, %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]} :
+ !tensorA, !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_lhs_tensor
+// 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<16x16x4xi8>
+!vecB = vector<16x16x4xi8>
+!vecC = vector<16x16xi32>
+!memrefA = memref<32x16x4xi8>
+!tensorB = tensor<16x32x4xi8>
+!memrefC = memref<32x32xi32>
+#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_rhs_tensor(
+ %arg0: !memrefA, %arg1: !tensorB, %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]} :
+ !tensorB, !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_rhs_tensor
+// 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<16x16x4xi8>
+!vecB = vector<16x16x4xi8>
+!vecC = vector<16x16xi32>
+!memrefA = memref<32x16x4xi8>
+!memrefB = memref<16x32x4xi8>
+!tensorC = tensor<32x32xi32>
+#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_acc_tensor(
+ %arg0: !memrefA, %arg1: !memrefB, %arg2: !tensorC) -> !tensorC
+{
+ %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]} : !tensorC, !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
+
+ %5 = vector.transfer_write %4, %arg2[%c0, %c0] {in_bounds = [true, true]} : !vecC, !tensorC
+ return %5 : !tensorC
+}
+
+// CHECK-LABEL: @negative_acc_tensor
+// 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
+ }
+}
More information about the Mlir-commits
mailing list