[Mlir-commits] [mlir] [mlir][linalg] Fix incorrect hoisting of transfers with mismatched masks (PR #232246)
llvmlistbot at llvm.org
llvmlistbot at llvm.org
Thu Oct 15 02:49:45 PDT 2026
llvmorg-github-actions[bot] wrote:
<!--LLVM PR SUMMARY COMMENT-->
@llvm/pr-subscribers-mlir
Author: Pecco (Pecco-314)
<details>
<summary>Changes</summary>
`hoist_redundant_vector_transfers` can change the result of a loop when the read and write have different mask operands.
For example, `%buffer` initially contains `[0, 0]`, `%mask` is `[true, false]`, `%zero` is `0`, and `%ones` is `[1, 1]`:
```mlir
scf.for %i = %c0 to %c2 step %c1 {
%v = vector.transfer_read %buffer[%c0], %zero, %mask
{in_bounds = [true]} : memref<2xi32>, vector<2xi32>
%next = arith.addi %v, %ones : vector<2xi32>
vector.transfer_write %next, %buffer[%c0]
{in_bounds = [true]} : vector<2xi32>, memref<2xi32>
}
```
The masked read returns zero for the second element on every iteration. The two iterations write `[1, 1]` and `[2, 1]`, respectively.
The transformation hoists the read and write out of the loop:
```mlir
%init = vector.transfer_read %buffer[%c0], %zero, %mask
{in_bounds = [true]} : memref<2xi32>, vector<2xi32>
%result = scf.for %i = %c0 to %c2 step %c1
iter_args(%v = %init) -> (vector<2xi32>) {
%next = arith.addi %v, %ones : vector<2xi32>
scf.yield %next : vector<2xi32>
}
vector.transfer_write %result, %buffer[%c0]
{in_bounds = [true]} : vector<2xi32>, memref<2xi32>
```
Here, the mask is applied only before the loop. Both elements are incremented on every iteration, so the final write stores `[2, 2]` instead of `[2, 1]`.
To fix this, hoisting is disallowed when the transfers' mask operands differ, including when only one transfer has a mask. Pairs with identical mask operands or no masks remain eligible for hoisting.
Assisted-by: Codex (OpenAI)
---
Full diff: https://github.com/llvm/llvm-project/pull/232246.diff
2 Files Affected:
- (modified) mlir/lib/Dialect/Linalg/Transforms/Hoisting.cpp (+3-2)
- (modified) mlir/test/Dialect/Linalg/hoist-redundant-vector-transfers.mlir (+206)
``````````diff
diff --git a/mlir/lib/Dialect/Linalg/Transforms/Hoisting.cpp b/mlir/lib/Dialect/Linalg/Transforms/Hoisting.cpp
index 317f36b3569a18..da90044756a50a 100644
--- a/mlir/lib/Dialect/Linalg/Transforms/Hoisting.cpp
+++ b/mlir/lib/Dialect/Linalg/Transforms/Hoisting.cpp
@@ -302,8 +302,8 @@ void mlir::linalg::hoistRedundantVectorTransfers(Operation *root,
<< "\n");
// Approximate aliasing by checking that:
- // 1. indices, vector type and permutation map are the same (i.e., the
- // transfer_read/transfer_write ops are matching),
+ // 1. indices, vector type, permutation map and mask are the same
+ // (i.e., the transfer_read/transfer_write ops are matching),
// 2. source operands for transfer.{read|write} do not originate from
// nor have users that are Ops implementing ViewLikeOpInterface.
// 3. no other operations in the loop access the same memref except
@@ -312,6 +312,7 @@ void mlir::linalg::hoistRedundantVectorTransfers(Operation *root,
// Check 1.
if (transferRead.getIndices() != transferWrite.getIndices() ||
+ transferRead.getMask() != transferWrite.getMask() ||
transferRead.getVectorType() != transferWrite.getVectorType() ||
transferRead.getPermutationMap() != transferWrite.getPermutationMap())
return WalkResult::advance();
diff --git a/mlir/test/Dialect/Linalg/hoist-redundant-vector-transfers.mlir b/mlir/test/Dialect/Linalg/hoist-redundant-vector-transfers.mlir
index b5eb210964b2af..97af6a67dadb25 100644
--- a/mlir/test/Dialect/Linalg/hoist-redundant-vector-transfers.mlir
+++ b/mlir/test/Dialect/Linalg/hoist-redundant-vector-transfers.mlir
@@ -1747,3 +1747,209 @@ module attributes {transform.with_named_sequence} {
transform.yield
}
}
+
+// -----
+
+// Only the read is masked, so hoisting is not safe. The second element must
+// be read as zero on every iteration, leaving 1 in memory after the loop.
+
+// CHECK-LABEL: func.func @masked_read_unmasked_write()
+// CHECK: %[[MASK:.*]] = arith.constant dense<[true, false]>
+// CHECK: scf.for {{.*}} {
+// CHECK-NEXT: %[[READ:.*]] = vector.transfer_read {{.*}}, %[[MASK]]
+// CHECK-NEXT: %[[NEXT:.*]] = arith.addi %[[READ]],
+// CHECK-NEXT: vector.transfer_write %[[NEXT]],
+// CHECK-NEXT: }{{$}}
+// CHECK: memref.load
+func.func @masked_read_unmasked_write() -> i32 {
+ %c0 = arith.constant 0 : index
+ %c1 = arith.constant 1 : index
+ %c2 = arith.constant 2 : index
+ %zero = arith.constant 0 : i32
+ %zeros = arith.constant dense<0> : vector<2xi32>
+ %ones = arith.constant dense<1> : vector<2xi32>
+ %mask = arith.constant dense<[true, false]> : vector<2xi1>
+ %buffer = memref.alloca() : memref<2xi32>
+ vector.transfer_write %zeros, %buffer[%c0]
+ {in_bounds = [true]} : vector<2xi32>, memref<2xi32>
+ scf.for %i = %c0 to %c2 step %c1 {
+ %v = vector.transfer_read %buffer[%c0], %zero, %mask
+ {in_bounds = [true]} : memref<2xi32>, vector<2xi32>
+ %next = arith.addi %v, %ones : vector<2xi32>
+ vector.transfer_write %next, %buffer[%c0]
+ {in_bounds = [true]} : vector<2xi32>, memref<2xi32>
+ }
+ %result = memref.load %buffer[%c1] : memref<2xi32>
+ return %result : i32
+}
+
+module attributes {transform.with_named_sequence} {
+ transform.named_sequence @__transform_main(%root: !transform.any_op {transform.readonly}) {
+ %funcs = transform.structured.match ops{["func.func"]} in %root
+ : (!transform.any_op) -> !transform.any_op
+ transform.structured.hoist_redundant_vector_transfers %funcs verify_non_zero_trip
+ : (!transform.any_op) -> !transform.any_op
+ transform.yield
+ }
+}
+
+// -----
+
+// Only the write is masked, so hoisting is not safe. The next read must see
+// the elements that the masked write leaves unchanged in memory.
+
+// CHECK-LABEL: func.func @unmasked_read_masked_write(
+// CHECK-SAME: %[[MEM:.*]]: memref<2xi32>, %[[MASK:.*]]: vector<2xi1>
+// CHECK: scf.for {{.*}} {
+// CHECK-NEXT: %[[READ:.*]] = vector.transfer_read %[[MEM]]
+// CHECK-NEXT: %[[SWAP:.*]] = vector.shuffle %[[READ]], %{{.*}} [1, 0]
+// CHECK-NEXT: %[[NEXT:.*]] = arith.addi %[[SWAP]],
+// CHECK-NEXT: vector.transfer_write %[[NEXT]], %[[MEM]]{{.*}}, %[[MASK]]
+// CHECK-NEXT: }{{$}}
+func.func @unmasked_read_masked_write(%mem: memref<2xi32>, %mask: vector<2xi1>) {
+ %c0 = arith.constant 0 : index
+ %c1 = arith.constant 1 : index
+ %c2 = arith.constant 2 : index
+ %zero = arith.constant 0 : i32
+ %ones = arith.constant dense<1> : vector<2xi32>
+ scf.for %i = %c0 to %c2 step %c1 {
+ %v = vector.transfer_read %mem[%c0], %zero
+ {in_bounds = [true]} : memref<2xi32>, vector<2xi32>
+ %swap = vector.shuffle %v, %v [1, 0] : vector<2xi32>, vector<2xi32>
+ %next = arith.addi %swap, %ones : vector<2xi32>
+ vector.transfer_write %next, %mem[%c0], %mask
+ {in_bounds = [true]} : vector<2xi32>, memref<2xi32>
+ }
+ return
+}
+
+module attributes {transform.with_named_sequence} {
+ transform.named_sequence @__transform_main(%root: !transform.any_op {transform.readonly}) {
+ %funcs = transform.structured.match ops{["func.func"]} in %root
+ : (!transform.any_op) -> !transform.any_op
+ transform.structured.hoist_redundant_vector_transfers %funcs verify_non_zero_trip
+ : (!transform.any_op) -> !transform.any_op
+ transform.yield
+ }
+}
+
+// -----
+
+// The read and write use separate mask arguments. Hoisting is not safe
+// because the masks may select different elements.
+
+// CHECK-LABEL: func.func @distinct_transfer_masks(
+// CHECK-SAME: %[[MEM:.*]]: memref<2xi32>,
+// CHECK-SAME: %[[MASK:.*]]: vector<2xi1>,
+// CHECK-SAME: %[[READ_MASK:.*]]: vector<2xi1>
+// CHECK: scf.for {{.*}} {
+// CHECK-NEXT: %[[READ:.*]] = vector.transfer_read %[[MEM]]{{.*}}, %[[READ_MASK]]
+// CHECK-NEXT: %[[SWAP:.*]] = vector.shuffle %[[READ]], %{{.*}} [1, 0]
+// CHECK-NEXT: %[[NEXT:.*]] = arith.addi %[[SWAP]],
+// CHECK-NEXT: vector.transfer_write %[[NEXT]], %[[MEM]]{{.*}}, %[[MASK]]
+// CHECK-NEXT: }{{$}}
+func.func @distinct_transfer_masks(
+ %mem: memref<2xi32>, %mask: vector<2xi1>, %read_mask: vector<2xi1>) {
+ %c0 = arith.constant 0 : index
+ %c1 = arith.constant 1 : index
+ %c2 = arith.constant 2 : index
+ %zero = arith.constant 0 : i32
+ %ones = arith.constant dense<1> : vector<2xi32>
+ scf.for %i = %c0 to %c2 step %c1 {
+ %v = vector.transfer_read %mem[%c0], %zero, %read_mask
+ {in_bounds = [true]} : memref<2xi32>, vector<2xi32>
+ %swap = vector.shuffle %v, %v [1, 0] : vector<2xi32>, vector<2xi32>
+ %next = arith.addi %swap, %ones : vector<2xi32>
+ vector.transfer_write %next, %mem[%c0], %mask
+ {in_bounds = [true]} : vector<2xi32>, memref<2xi32>
+ }
+ return
+}
+
+module attributes {transform.with_named_sequence} {
+ transform.named_sequence @__transform_main(%root: !transform.any_op {transform.readonly}) {
+ %funcs = transform.structured.match ops{["func.func"]} in %root
+ : (!transform.any_op) -> !transform.any_op
+ transform.structured.hoist_redundant_vector_transfers %funcs verify_non_zero_trip
+ : (!transform.any_op) -> !transform.any_op
+ transform.yield
+ }
+}
+
+// -----
+
+// Hoisting is safe when both transfers use the same mask and the loop
+// updates each vector element independently.
+
+// CHECK-LABEL: func.func @same_transfer_mask(
+// CHECK-SAME: %[[MEM:.*]]: memref<2xi32>, %[[MASK:.*]]: vector<2xi1>
+// CHECK: %[[READ:.*]] = vector.transfer_read %[[MEM]]{{.*}}, %[[MASK]]
+// CHECK-NEXT: %[[LOOP:.*]] = scf.for {{.*}} iter_args(%[[ARG:.*]] = %[[READ]])
+// CHECK-NEXT: %[[NEXT:.*]] = arith.addi %[[ARG]],
+// CHECK-NEXT: scf.yield %[[NEXT]]
+// CHECK-NEXT: }{{$}}
+// CHECK-NEXT: vector.transfer_write %[[LOOP]], %[[MEM]]{{.*}}, %[[MASK]]
+func.func @same_transfer_mask(%mem: memref<2xi32>, %mask: vector<2xi1>) {
+ %c0 = arith.constant 0 : index
+ %c1 = arith.constant 1 : index
+ %c2 = arith.constant 2 : index
+ %zero = arith.constant 0 : i32
+ %ones = arith.constant dense<1> : vector<2xi32>
+ scf.for %i = %c0 to %c2 step %c1 {
+ %v = vector.transfer_read %mem[%c0], %zero, %mask
+ {in_bounds = [true]} : memref<2xi32>, vector<2xi32>
+ %next = arith.addi %v, %ones : vector<2xi32>
+ vector.transfer_write %next, %mem[%c0], %mask
+ {in_bounds = [true]} : vector<2xi32>, memref<2xi32>
+ }
+ return
+}
+
+module attributes {transform.with_named_sequence} {
+ transform.named_sequence @__transform_main(%root: !transform.any_op {transform.readonly}) {
+ %funcs = transform.structured.match ops{["func.func"]} in %root
+ : (!transform.any_op) -> !transform.any_op
+ transform.structured.hoist_redundant_vector_transfers %funcs verify_non_zero_trip
+ : (!transform.any_op) -> !transform.any_op
+ transform.yield
+ }
+}
+
+// -----
+
+// The write mask changes on every iteration, so the write must stay
+// inside the loop.
+
+// CHECK-LABEL: func.func @loop_variant_write_mask(
+// CHECK: scf.for %[[IV:.*]] = {{.*}} {
+// CHECK-NEXT: %[[MASK:.*]] = vector.create_mask %[[IV]]
+// CHECK-NEXT: %[[READ:.*]] = vector.transfer_read
+// CHECK-NEXT: %[[NEXT:.*]] = arith.addi %[[READ]],
+// CHECK-NEXT: vector.transfer_write %[[NEXT]], {{.*}}, %[[MASK]]
+// CHECK-NEXT: }{{$}}
+func.func @loop_variant_write_mask(%mem: memref<2xi32>) {
+ %c0 = arith.constant 0 : index
+ %c1 = arith.constant 1 : index
+ %c2 = arith.constant 2 : index
+ %zero = arith.constant 0 : i32
+ %ones = arith.constant dense<1> : vector<2xi32>
+ scf.for %i = %c0 to %c2 step %c1 {
+ %mask = vector.create_mask %i : vector<2xi1>
+ %v = vector.transfer_read %mem[%c0], %zero
+ {in_bounds = [true]} : memref<2xi32>, vector<2xi32>
+ %next = arith.addi %v, %ones : vector<2xi32>
+ vector.transfer_write %next, %mem[%c0], %mask
+ {in_bounds = [true]} : vector<2xi32>, memref<2xi32>
+ }
+ return
+}
+
+module attributes {transform.with_named_sequence} {
+ transform.named_sequence @__transform_main(%root: !transform.any_op {transform.readonly}) {
+ %funcs = transform.structured.match ops{["func.func"]} in %root
+ : (!transform.any_op) -> !transform.any_op
+ transform.structured.hoist_redundant_vector_transfers %funcs verify_non_zero_trip
+ : (!transform.any_op) -> !transform.any_op
+ transform.yield
+ }
+}
``````````
</details>
https://github.com/llvm/llvm-project/pull/232246
More information about the Mlir-commits
mailing list