[Mlir-commits] [mlir] [mlir][x86] Fix result tracing through loops (PR #230229)
llvmlistbot at llvm.org
llvmlistbot at llvm.org
Thu Oct 8 10:55:09 PDT 2026
llvmorg-github-actions[bot] wrote:
<!--LLVM PR SUMMARY COMMENT-->
@llvm/pr-subscribers-mlir
Author: Adam Siemieniuk (adam-smnk)
<details>
<summary>Changes</summary>
Fixes the search for the write of a contraction result in the AMX lowering when the result is passed through loop args.
The value was followed to the wrong loop result, which either found the wrong write or crashed when the loop has few results.
Assisted-by: Claude
---
Full diff: https://github.com/llvm/llvm-project/pull/230229.diff
2 Files Affected:
- (modified) mlir/lib/Dialect/X86/Utils/X86Utils.cpp (+4-2)
- (modified) mlir/test/Dialect/X86/AMX/vector-contract-to-tiled-dp.mlir (+192)
``````````diff
diff --git a/mlir/lib/Dialect/X86/Utils/X86Utils.cpp b/mlir/lib/Dialect/X86/Utils/X86Utils.cpp
index 554e8d20a7a627..b5bfdaf9efaad2 100644
--- a/mlir/lib/Dialect/X86/Utils/X86Utils.cpp
+++ b/mlir/lib/Dialect/X86/Utils/X86Utils.cpp
@@ -257,8 +257,10 @@ Operation *traceToVectorWriteLikeUserOperation(Value v) {
// --- SCF FOR ---
if (auto forOp = dyn_cast<scf::ForOp>(user)) {
- unsigned idx = use.getOperandNumber();
- if (auto *res = traceToVectorWriteLikeUserOperation(forOp.getResult(idx)))
+ OpResult loopResult = forOp.getTiedLoopResult(&use);
+ if (!loopResult)
+ continue;
+ if (auto *res = traceToVectorWriteLikeUserOperation(loopResult))
return res;
continue;
}
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 3e6d53fd4f1457..ac31f5f0785c75 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
@@ -778,6 +778,198 @@ module attributes {transform.with_named_sequence} {
// -----
+// The contraction results are carried through another loop before being
+// written.
+
+!vecAB = vector<16x16x4xi8>
+!vecC = vector<16x16xi32>
+!memrefA = memref<16x16x4xi8, 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_result_through_loop(
+ %arg0: memref<64x64x4xi8>, %arg1: memref<64x128x4xi8>, %arg2: memref<64x128xi32>,
+ %bias: !vecC) {
+ %cst = arith.constant dense<0> : vector<16x16xi32>
+ %1 = ub.poison : i8
+ %c0 = arith.constant 0 : index
+ %c1 = arith.constant 1 : index
+ %c2 = arith.constant 2 : 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
+
+ %4:2 = scf.for %arg5 = %c0 to %c64 step %c16 iter_args(%arg6 = %cst, %arg7 = %cst) -> (!vecC, !vecC) {
+
+ %subview_0 = memref.subview %arg0[%arg3, %arg5, 0] [16, 16, 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, %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
+ }
+
+ %10:2 = scf.for %arg5 = %c0 to %c2 step %c1 iter_args(%arg6 = %4#0, %arg7 = %4#1) -> (!vecC, !vecC) {
+ %11 = arith.addi %arg6, %bias : !vecC
+ %12 = arith.addi %arg7, %bias : !vecC
+ scf.yield %11, %12 : !vecC, !vecC
+ }
+
+ vector.transfer_write %10#1, %subview[%c0, %c16] {in_bounds = [true, true]} :
+ !vecC, !memrefC
+ vector.transfer_write %10#0, %subview[%c0, %c0] {in_bounds = [true, true]} :
+ !vecC, !memrefC
+ }
+ }
+
+ return
+}
+
+// CHECK-LABEL: @matmul_int8_loop_result_through_loop
+// CHECK-COUNT-2: x86.amx.tile_zero : !x86.amx.tile<16x16xi32>
+// CHECK: %[[TILES:.+]]:2 = scf.for {{.*}} -> (!x86.amx.tile<16x16xi32>, !x86.amx.tile<16x16xi32>) {
+// CHECK: x86.amx.tile_store %[[BUF:.+]][%{{.+}}, %{{.+}}], %[[TILES]]#0
+// CHECK: x86.amx.tile_store %[[BUF]][%{{.+}}, %{{.+}}], %[[TILES]]#1
+// CHECK: %[[RES0:.+]] = vector.transfer_read %[[BUF]]
+// CHECK: %[[RES1:.+]] = vector.transfer_read %[[BUF]]
+// CHECK: scf.for {{.*}} iter_args(%{{.+}} = %[[RES0]], %{{.+}} = %[[RES1]])
+// 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 loop carrying the contraction results also carries other values, which
+// are written to a different buffer.
+
+!vecAB = vector<16x16x4xi8>
+!vecC = vector<16x16xi32>
+!memrefA = memref<16x16x4xi8, 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_result_through_loop_extra_args(
+ %arg0: memref<64x64x4xi8>, %arg1: memref<64x128x4xi8>, %arg2: memref<64x128xi32>,
+ %bias: !vecC, %out: memref<16x16xi32>) {
+ %cst = arith.constant dense<0> : vector<16x16xi32>
+ %1 = ub.poison : i8
+ %c0 = arith.constant 0 : index
+ %c1 = arith.constant 1 : index
+ %c2 = arith.constant 2 : 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
+
+ %4:2 = scf.for %arg5 = %c0 to %c64 step %c16 iter_args(%arg6 = %cst, %arg7 = %cst) -> (!vecC, !vecC) {
+
+ %subview_0 = memref.subview %arg0[%arg3, %arg5, 0] [16, 16, 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, %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
+ }
+
+ %10:5 = scf.for %arg5 = %c0 to %c2 step %c1
+ iter_args(%arg6 = %bias, %arg7 = %4#1, %arg8 = %bias, %arg9 = %4#0, %arg10 = %bias)
+ -> (!vecC, !vecC, !vecC, !vecC, !vecC) {
+ %11 = arith.addi %arg9, %arg6 : !vecC
+ %12 = arith.addi %arg7, %arg6 : !vecC
+ scf.yield %arg6, %12, %arg8, %11, %arg10 : !vecC, !vecC, !vecC, !vecC, !vecC
+ }
+
+ vector.transfer_write %10#1, %subview[%c0, %c16] {in_bounds = [true, true]} :
+ !vecC, !memrefC
+ vector.transfer_write %10#3, %subview[%c0, %c0] {in_bounds = [true, true]} :
+ !vecC, !memrefC
+ vector.transfer_write %10#0, %out[%c0, %c0] {in_bounds = [true, true]} :
+ !vecC, memref<16x16xi32>
+ vector.transfer_write %10#2, %out[%c0, %c0] {in_bounds = [true, true]} :
+ !vecC, memref<16x16xi32>
+ vector.transfer_write %10#4, %out[%c0, %c0] {in_bounds = [true, true]} :
+ !vecC, memref<16x16xi32>
+ }
+ }
+
+ return
+}
+
+// CHECK-LABEL: @matmul_int8_loop_result_through_loop_extra_args
+// CHECK: %[[TILES:.+]]:2 = scf.for {{.*}} -> (!x86.amx.tile<16x16xi32>, !x86.amx.tile<16x16xi32>) {
+// CHECK: %[[BUF:.+]] = memref.alloca() : memref<16x32xi32>
+// CHECK: x86.amx.tile_store %[[BUF]][%{{.+}}, %{{.+}}], %[[TILES]]#0
+// CHECK: x86.amx.tile_store %[[BUF]][%{{.+}}, %{{.+}}], %[[TILES]]#1
+// CHECK: %[[RES0:.+]] = vector.transfer_read %[[BUF]]
+// CHECK: %[[RES1:.+]] = vector.transfer_read %[[BUF]]
+// CHECK: scf.for {{.*}} iter_args(%{{.+}} = %{{.+}}, %{{.+}} = %[[RES1]], %{{.+}} = %{{.+}}, %{{.+}} = %[[RES0]], %{{.+}} = %{{.+}})
+// 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
+ }
+}
+
+// -----
+
!vecA = vector<1x16x32xbf16>
!vecB = vector<1x32x16xbf16>
!vecC = vector<16x16xf32>
``````````
</details>
https://github.com/llvm/llvm-project/pull/230229
More information about the Mlir-commits
mailing list