[Mlir-commits] [mlir] b0c2017 - [mlir][x86] Fail on missing read source operation (#205077)

[llvmlistbot at llvm.org]

Author: Adam Siemieniuk
Date: 2026-06-22T14:37:20+02:00
New Revision: b0c2017bcff5ae07fe44d3271a269e50e1c98cbd

URL: https://github.com/llvm/llvm-project/commit/b0c2017bcff5ae07fe44d3271a269e50e1c98cbd
DIFF: https://github.com/llvm/llvm-project/commit/b0c2017bcff5ae07fe44d3271a269e50e1c98cbd.diff

LOG: [mlir][x86] Fail on missing read source operation (#205077)

Adds an extra check to AMX lowering to fail gracefully when a source
operation for contraction input data is not found.

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 bfdeb51b65ef1..86291ae03f29e 100644
--- a/mlir/lib/Dialect/X86/Transforms/VectorContractToAMXDotProduct.cpp
+++ b/mlir/lib/Dialect/X86/Transforms/VectorContractToAMXDotProduct.cpp
@@ -1064,6 +1064,10 @@ struct VectorContractToAMXDotProduct
           vectorOpRhs = readOp.getBase().getDefiningOp();
         });
 
+    if (!vectorOpLhs || !vectorOpRhs)
+      return rewriter.notifyMatchFailure(
+          contractOp, "Failed to find LHS or RHS read source operation");
+
     // Retrive all the contaction operation within the loop.
     SmallVector<vector::ContractionOp> ops;
     for (mlir::Operation &op : loopLists[0].getBody()->getOperations()) {

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 fbbd20202620b..71bf62c56a6e5 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
@@ -2009,3 +2009,72 @@ module attributes {transform.with_named_sequence} {
     transform.yield
   }
 }
+
+// -----
+
+!vecA = vector<1x16x32xbf16>
+!vecB = vector<1x32x16xbf16>
+!vecC = vector<16x16xf32>
+
+#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 @negative_no_read_op_source(%arg0: memref<16x64x32xbf16>, %arg1: memref<16x32x128xbf16>, %arg2: memref<64x32xf32>) {
+  %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
+  %c1 = arith.constant 1 : index
+  scf.for %arg3 = %c0 to %c64 step %c32 {
+    %2 = vector.transfer_read %arg2[%arg3, %c0], %0 {in_bounds = [true, true]} : memref<64x32xf32>, !vecC
+    %3 = vector.transfer_read %arg2[%arg3, %c16], %0 {in_bounds = [true, true]} : memref<64x32xf32>, !vecC
+    %4 = arith.addi %arg3, %c16 : index
+    %5 = vector.transfer_read %arg2[%4, %c0], %0 {in_bounds = [true, true]} : memref<64x32xf32>, !vecC
+    %6 = vector.transfer_read %arg2[%4, %c16], %0 {in_bounds = [true, true]} : memref<64x32xf32>, !vecC
+    %7:4 = scf.for %arg4 = %c0 to %c16 step %c1 iter_args(%arg5 = %2, %arg6 = %3, %arg7 = %5, %arg8 = %6) -> (!vecC, !vecC, !vecC, !vecC) {
+      %8 = vector.transfer_read %arg0[%arg4, %arg3, %c0], %1 {in_bounds = [true, true, true]}
+        : memref<16x64x32xbf16>, !vecA
+      %9 = vector.transfer_read %arg0[%arg4, %4, %c0], %1 {in_bounds = [true, true, true]}
+        : memref<16x64x32xbf16>, !vecA
+      %10 = vector.transfer_read %arg1[%arg4, %c0, %c0], %1 {in_bounds = [true, true, true]}
+        : memref<16x32x128xbf16>, !vecB
+      %11 = vector.transfer_read %arg1[%arg4, %c0, %c16], %1 {in_bounds = [true, true, true]}
+        : memref<16x32x128xbf16>, !vecB
+      %12 = vector.contract {indexing_maps = [#map, #map1, #map2], iterator_types =
+              ["reduction", "parallel", "parallel", "reduction"], kind = #vector.kind<add>}
+              %8, %10, %arg5 : !vecA, !vecB into !vecC
+      %13 = vector.contract {indexing_maps = [#map, #map1, #map2], iterator_types =
+              ["reduction", "parallel", "parallel", "reduction"], kind = #vector.kind<add>}
+              %8, %11, %arg6 : !vecA, !vecB into !vecC
+      %14 = vector.contract {indexing_maps = [#map, #map1, #map2], iterator_types =
+              ["reduction", "parallel", "parallel", "reduction"], kind = #vector.kind<add>}
+              %9, %10, %arg7 : !vecA, !vecB into !vecC
+      %15 = vector.contract {indexing_maps = [#map, #map1, #map2], iterator_types =
+              ["reduction", "parallel", "parallel", "reduction"], kind = #vector.kind<add>}
+              %9, %11, %arg8 : !vecA, !vecB into !vecC
+      scf.yield %12, %13, %14, %15 : !vecC, !vecC, !vecC, !vecC
+    }
+    vector.transfer_write %7#3, %arg2[%4, %c16] {in_bounds = [true, true]} : !vecC, memref<64x32xf32>
+    vector.transfer_write %7#2, %arg2[%4, %c0] {in_bounds = [true, true]} : !vecC, memref<64x32xf32>
+    vector.transfer_write %7#1, %arg2[%arg3, %c16] {in_bounds = [true, true]} : !vecC, memref<64x32xf32>
+    vector.transfer_write %7#0, %arg2[%arg3, %c0] {in_bounds = [true, true]} : !vecC, memref<64x32xf32>
+  }
+  return
+}
+
+// CHECK-LABEL: @negative_no_read_op_source
+// CHECK-NOT: x86.amx
+// CHECK: vector.contract
+
+module attributes {transform.with_named_sequence} {
+  transform.named_sequence @__transform_main(%arg0: !transform.any_op {transform.readonly}) {
+    %0 = transform.structured.match ops{["func.func"]} in %arg0 : (!transform.any_op) -> !transform.any_op
+    transform.apply_patterns to %0 {
+      transform.apply_patterns.x86.vector_contract_to_amx_dot_product
+    } : !transform.any_op
+    transform.yield
+  }
+}