[Mlir-commits] [mlir] [mlir] Fix tensor transfer hoisting with mismatched masks or permutation maps (PR #234878)
llvmlistbot at llvm.org
llvmlistbot at llvm.org
Fri Oct 23 01:58:18 PDT 2026
https://github.com/Pecco-314 created https://github.com/llvm/llvm-project/pull/234878
Loop-invariant subset hoisting can pair a tensor transfer read with a write that uses a different mask or permutation map. These operations may access the same region and have the same vector type, but hoisting them can change the result.
For example, suppose `%zeros` is `[0, 0]`, `%mask` is `[true, false]`, `%zero` is `0`, and `%ones` is `[1, 1]`. The following loop runs twice:
```mlir
%result = scf.for %i = %c0 to %c2 step %c1
iter_args(%t = %zeros) -> tensor<2xi32> {
%v = vector.transfer_read %t[%c0], %zero, %mask
: tensor<2xi32>, vector<2xi32>
%next = arith.addi %v, %ones : vector<2xi32>
%out = vector.transfer_write %next, %t[%c0]
: vector<2xi32>, tensor<2xi32>
scf.yield %out : tensor<2xi32>
}
```
The masked read returns zero for the second element on each iteration, so the write always stores one in that position. The final tensor is `[2, 1]`.
Hoisting produces:
```mlir
%init = vector.transfer_read %zeros[%c0], %zero, %mask
: tensor<2xi32>, vector<2xi32>
%values:2 = scf.for %i = %c0 to %c2 step %c1
iter_args(%t = %zeros, %v = %init)
-> (tensor<2xi32>, vector<2xi32>) {
%next = arith.addi %v, %ones : vector<2xi32>
scf.yield %t, %next : tensor<2xi32>, vector<2xi32>
}
%result = vector.transfer_write %values#1, %values#0[%c0]
: vector<2xi32>, tensor<2xi32>
```
The read now happens only once, before the loop. Each iteration adds one to both elements of the previous iteration's vector, producing `[2, 2]`.
Different permutation maps cause a similar problem. With an identity read and a transposed write, each iteration reads back the transposed data. After hoisting, the loop uses the vector directly, and the transpose happens only in the final write.
A new interface method checks that the read and write have matching value types, masks and permutation maps before they can be hoisted. Tensor slices retain the existing type check.
Assisted-by: Codex (OpenAI)
>From 765614ff3792f467f138e6272b85049f29ccb398 Mon Sep 17 00:00:00 2001
From: Yanfeng Gao <gaoyanfeng.gyf at alibaba-inc.com>
Date: Fri, 23 Oct 2026 16:01:22 +0800
Subject: [PATCH] [mlir] Fix tensor transfer hoisting with mismatched masks or
permutation maps
Loop-invariant subset hoisting can pair a tensor transfer read with a
write that uses a different mask or permutation map. These operations
may access the same region and have the same vector type, but hoisting
them can change the result.
For example, suppose `%zeros` is `[0, 0]`, `%mask` is `[true, false]`,
`%zero` is `0`, and `%ones` is `[1, 1]`. The following loop runs twice:
```mlir
%result = scf.for %i = %c0 to %c2 step %c1
iter_args(%t = %zeros) -> tensor<2xi32> {
%v = vector.transfer_read %t[%c0], %zero, %mask
: tensor<2xi32>, vector<2xi32>
%next = arith.addi %v, %ones : vector<2xi32>
%out = vector.transfer_write %next, %t[%c0]
: vector<2xi32>, tensor<2xi32>
scf.yield %out : tensor<2xi32>
}
```
The masked read returns zero for the second element on each iteration,
so the write always stores one in that position. The final tensor is
`[2, 1]`.
Hoisting produces:
```mlir
%init = vector.transfer_read %zeros[%c0], %zero, %mask
: tensor<2xi32>, vector<2xi32>
%values:2 = scf.for %i = %c0 to %c2 step %c1
iter_args(%t = %zeros, %v = %init)
-> (tensor<2xi32>, vector<2xi32>) {
%next = arith.addi %v, %ones : vector<2xi32>
scf.yield %t, %next : tensor<2xi32>, vector<2xi32>
}
%result = vector.transfer_write %values#1, %values#0[%c0]
: vector<2xi32>, tensor<2xi32>
```
The read now happens only once, before the loop. Each iteration adds one
to both elements of the previous iteration's vector, producing `[2, 2]`.
Different permutation maps cause a similar problem. With an identity
read and a transposed write, each iteration reads back the transposed
data. After hoisting, the loop uses the vector directly, and the
transpose happens only in the final write.
A new interface method checks that the read and write have matching
value types, masks and permutation maps before they can be hoisted.
Tensor slices retain the existing type check.
Assisted-by: Codex (OpenAI)
---
.../mlir/Interfaces/SubsetOpInterface.td | 14 +
.../Transforms/SubsetOpInterfaceImpl.cpp | 9 +
.../Utils/LoopInvariantCodeMotionUtils.cpp | 4 +-
.../loop-invariant-subset-hoisting.mlir | 239 ++++++++++++++++++
4 files changed, 263 insertions(+), 3 deletions(-)
diff --git a/mlir/include/mlir/Interfaces/SubsetOpInterface.td b/mlir/include/mlir/Interfaces/SubsetOpInterface.td
index 7000e7dfc89cdb..a63163ff6bc1c9 100644
--- a/mlir/include/mlir/Interfaces/SubsetOpInterface.td
+++ b/mlir/include/mlir/Interfaces/SubsetOpInterface.td
@@ -124,6 +124,20 @@ def SubsetExtractionOpInterface
let cppNamespace = "::mlir";
let methods = [
+ InterfaceMethod<
+ /*desc=*/[{
+ Return "true" if this extraction and the given insertion have compatible
+ value types and op-specific properties.
+ }],
+ /*retType=*/"bool",
+ /*methodName=*/"isCompatibleWithInsertion",
+ /*args=*/(ins "::mlir::SubsetInsertionOpInterface":$insertion),
+ /*methodBody=*/"",
+ /*defaultImplementation=*/[{
+ return $_op->getResult(0).getType() ==
+ insertion.getSourceOperand().get().getType();
+ }]
+ >,
InterfaceMethod<
/*desc=*/[{
Return the source tensor operand.
diff --git a/mlir/lib/Dialect/Vector/Transforms/SubsetOpInterfaceImpl.cpp b/mlir/lib/Dialect/Vector/Transforms/SubsetOpInterfaceImpl.cpp
index 392bbb1c422013..8d371b8aa78531 100644
--- a/mlir/lib/Dialect/Vector/Transforms/SubsetOpInterfaceImpl.cpp
+++ b/mlir/lib/Dialect/Vector/Transforms/SubsetOpInterfaceImpl.cpp
@@ -39,6 +39,15 @@ struct TransferReadOpSubsetExtractionOpInterface
OpOperand &getSourceOperand(Operation *op) const {
return cast<vector::TransferReadOp>(op).getBaseMutable();
}
+
+ bool isCompatibleWithInsertion(Operation *op,
+ SubsetInsertionOpInterface insertion) const {
+ auto read = cast<vector::TransferReadOp>(op);
+ auto write = dyn_cast<vector::TransferWriteOp>(insertion.getOperation());
+ return write && read.getVectorType() == write.getVectorType() &&
+ read.getMask() == write.getMask() &&
+ read.getPermutationMap() == write.getPermutationMap();
+ }
};
struct TransferWriteOpSubsetInsertionOpInterface
diff --git a/mlir/lib/Transforms/Utils/LoopInvariantCodeMotionUtils.cpp b/mlir/lib/Transforms/Utils/LoopInvariantCodeMotionUtils.cpp
index aef81bf2e15a38..87aafdae916de5 100644
--- a/mlir/lib/Transforms/Utils/LoopInvariantCodeMotionUtils.cpp
+++ b/mlir/lib/Transforms/Utils/LoopInvariantCodeMotionUtils.cpp
@@ -143,10 +143,8 @@ class MatchingSubsets {
return llvm::make_filter_range(
llvm::zip(extractions, insertions), [&](auto pair) {
auto [extractionOp, insertionOp] = pair;
- // Hoist only if the extracted and inserted values have the same type.
if (extractionOp && insertionOp &&
- extractionOp->getResult(0).getType() !=
- insertionOp.getSourceOperand().get().getType())
+ !extractionOp.isCompatibleWithInsertion(insertionOp))
return false;
// Hoist only if there are no conflicting subset ops.
return allDisjoint(extractionOp, insertionOp);
diff --git a/mlir/test/Transforms/loop-invariant-subset-hoisting.mlir b/mlir/test/Transforms/loop-invariant-subset-hoisting.mlir
index 57e9f7e772373b..1561030feb1d8c 100644
--- a/mlir/test/Transforms/loop-invariant-subset-hoisting.mlir
+++ b/mlir/test/Transforms/loop-invariant-subset-hoisting.mlir
@@ -645,6 +645,8 @@ func.func @do_not_hoist_multiple_nested_loop_uses(%arg: tensor<?xf32>)
return %0 : tensor<?xf32>
}
+// -----
+
// Ensure that cases with buffer semantics exit gracefully.
// CHECK-LABEL: @hoist_buffer
@@ -660,3 +662,240 @@ func.func @hoist_buffer(%arg0: memref<7x7xf16>) {
}
return
}
+
+// -----
+
+// Only the read is masked, so hoisting is not safe. The second element must
+// be read as zero on every iteration, leaving its final value at 1.
+
+// CHECK-LABEL: func.func @negative_masked_read_unmasked_write(
+// CHECK-NOT: vector.transfer_read
+// CHECK-NOT: vector.transfer_write
+// CHECK: scf.for
+// CHECK: vector.transfer_read
+// CHECK: vector.transfer_write
+// CHECK: scf.yield
+// CHECK-NOT: vector.transfer_read
+// CHECK-NOT: vector.transfer_write
+// CHECK: return
+func.func @negative_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> : tensor<2xi32>
+ %ones = arith.constant dense<1> : vector<2xi32>
+ %mask = arith.constant dense<[true, false]> : vector<2xi1>
+ %result = scf.for %i = %c0 to %c2 step %c1 iter_args(%t = %zeros) -> tensor<2xi32> {
+ %v = vector.transfer_read %t[%c0], %zero, %mask
+ : tensor<2xi32>, vector<2xi32>
+ %next = arith.addi %v, %ones : vector<2xi32>
+ %out = vector.transfer_write %next, %t[%c0]
+ : vector<2xi32>, tensor<2xi32>
+ scf.yield %out : tensor<2xi32>
+ }
+ %value = tensor.extract %result[%c1] : tensor<2xi32>
+ return %value : i32
+}
+
+// -----
+
+// Only the write is masked, so hoisting is not safe. The next read must see
+// the elements that the masked write leaves unchanged in the tensor.
+
+// CHECK-LABEL: func.func @negative_unmasked_read_masked_write(
+// CHECK-NOT: vector.transfer_read
+// CHECK-NOT: vector.transfer_write
+// CHECK: scf.for
+// CHECK: vector.transfer_read
+// CHECK: vector.transfer_write
+// CHECK: scf.yield
+// CHECK-NOT: vector.transfer_read
+// CHECK-NOT: vector.transfer_write
+// CHECK: return
+func.func @negative_unmasked_read_masked_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> : tensor<2xi32>
+ %ones = arith.constant dense<1> : vector<2xi32>
+ %mask = arith.constant dense<[true, false]> : vector<2xi1>
+ %result = scf.for %i = %c0 to %c2 step %c1 iter_args(%t = %zeros) -> tensor<2xi32> {
+ %v = vector.transfer_read %t[%c0], %zero
+ : tensor<2xi32>, vector<2xi32>
+ %swap = vector.shuffle %v, %v [1, 0] : vector<2xi32>, vector<2xi32>
+ %next = arith.addi %swap, %ones : vector<2xi32>
+ %out = vector.transfer_write %next, %t[%c0], %mask
+ : vector<2xi32>, tensor<2xi32>
+ scf.yield %out : tensor<2xi32>
+ }
+ %value = tensor.extract %result[%c0] : tensor<2xi32>
+ return %value : i32
+}
+
+// -----
+
+// The read and write masks select different elements. Hoisting is not safe
+// because the next read replaces the written element with padding.
+
+// CHECK-LABEL: func.func @negative_distinct_transfer_masks(
+// CHECK-NOT: vector.transfer_read
+// CHECK-NOT: vector.transfer_write
+// CHECK: scf.for
+// CHECK: vector.transfer_read
+// CHECK: vector.transfer_write
+// CHECK: scf.yield
+// CHECK-NOT: vector.transfer_read
+// CHECK-NOT: vector.transfer_write
+// CHECK: return
+func.func @negative_distinct_transfer_masks() -> 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> : tensor<2xi32>
+ %ones = arith.constant dense<1> : vector<2xi32>
+ %mask = arith.constant dense<[true, false]> : vector<2xi1>
+ %other_mask = arith.constant dense<[false, true]> : vector<2xi1>
+ %result = scf.for %i = %c0 to %c2 step %c1 iter_args(%t = %zeros) -> tensor<2xi32> {
+ %v = vector.transfer_read %t[%c0], %zero, %mask
+ : tensor<2xi32>, vector<2xi32>
+ %next = arith.addi %v, %ones : vector<2xi32>
+ %out = vector.transfer_write %next, %t[%c0], %other_mask
+ : vector<2xi32>, tensor<2xi32>
+ scf.yield %out : tensor<2xi32>
+ }
+ %value = tensor.extract %result[%c1] : tensor<2xi32>
+ return %value : i32
+}
+
+// -----
+
+// Hoisting is safe when both transfers use the same mask and the loop
+// updates each vector element independently.
+
+// CHECK-LABEL: func.func @same_mask_elementwise(
+// CHECK: vector.transfer_read
+// CHECK: scf.for
+// CHECK-NOT: vector.transfer_read
+// CHECK-NOT: vector.transfer_write
+// CHECK: scf.yield
+// CHECK: vector.transfer_write
+// CHECK: return
+func.func @same_mask_elementwise() -> 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> : tensor<2xi32>
+ %ones = arith.constant dense<1> : vector<2xi32>
+ %mask = arith.constant dense<[true, false]> : vector<2xi1>
+ %result = scf.for %i = %c0 to %c2 step %c1 iter_args(%t = %zeros) -> tensor<2xi32> {
+ %v = vector.transfer_read %t[%c0], %zero, %mask
+ : tensor<2xi32>, vector<2xi32>
+ %next = arith.addi %v, %ones : vector<2xi32>
+ %out = vector.transfer_write %next, %t[%c0], %mask
+ : vector<2xi32>, tensor<2xi32>
+ scf.yield %out : tensor<2xi32>
+ }
+ %value = tensor.extract %result[%c0] : tensor<2xi32>
+ return %value : i32
+}
+
+// -----
+
+// The read and write use different permutation maps, so hoisting is not safe.
+// The write must transpose the vector on every iteration.
+
+// CHECK-LABEL: func.func @negative_mismatched_permutation(
+// CHECK-NOT: vector.transfer_read
+// CHECK-NOT: vector.transfer_write
+// CHECK: scf.for
+// CHECK: vector.transfer_read
+// CHECK: vector.transfer_write
+// CHECK: scf.yield
+// CHECK-NOT: vector.transfer_read
+// CHECK-NOT: vector.transfer_write
+// CHECK: return
+func.func @negative_mismatched_permutation() -> 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> : tensor<2x2xi32>
+ %delta = arith.constant dense<[[0, 1], [0, 0]]> : vector<2x2xi32>
+ %result = scf.for %i = %c0 to %c2 step %c1 iter_args(%t = %zeros) -> tensor<2x2xi32> {
+ %v = vector.transfer_read %t[%c0, %c0], %zero
+ : tensor<2x2xi32>, vector<2x2xi32>
+ %next = arith.addi %v, %delta : vector<2x2xi32>
+ %out = vector.transfer_write %next, %t[%c0, %c0]
+ {permutation_map = affine_map<(d0, d1) -> (d1, d0)>} : vector<2x2xi32>, tensor<2x2xi32>
+ scf.yield %out : tensor<2x2xi32>
+ }
+ %value = tensor.extract %result[%c1, %c0] : tensor<2x2xi32>
+ return %value : i32
+}
+
+// -----
+
+// Hoisting is safe when both transfers use the same permutation map.
+// A transposed write followed by a transposed read preserves the vector.
+
+// CHECK-LABEL: func.func @same_transfer_permutation(
+// CHECK: vector.transfer_read
+// CHECK: scf.for
+// CHECK-NOT: vector.transfer_read
+// CHECK-NOT: vector.transfer_write
+// CHECK: scf.yield
+// CHECK: vector.transfer_write
+// CHECK: return
+func.func @same_transfer_permutation() -> 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> : tensor<2x2xi32>
+ %delta = arith.constant dense<[[0, 1], [0, 0]]> : vector<2x2xi32>
+ %result = scf.for %i = %c0 to %c2 step %c1 iter_args(%t = %zeros) -> tensor<2x2xi32> {
+ %v = vector.transfer_read %t[%c0, %c0], %zero
+ {permutation_map = affine_map<(d0, d1) -> (d1, d0)>} : tensor<2x2xi32>, vector<2x2xi32>
+ %next = arith.addi %v, %delta : vector<2x2xi32>
+ %out = vector.transfer_write %next, %t[%c0, %c0]
+ {permutation_map = affine_map<(d0, d1) -> (d1, d0)>} : vector<2x2xi32>, tensor<2x2xi32>
+ scf.yield %out : tensor<2x2xi32>
+ }
+ %value = tensor.extract %result[%c1, %c0] : tensor<2x2xi32>
+ return %value : i32
+}
+
+// -----
+
+// The slice ops access the same subset, but the extracted and inserted values
+// have different types. They cannot be hoisted as a pair.
+
+// CHECK-LABEL: func.func @negative_mismatched_slice_types(
+// CHECK-NOT: tensor.extract_slice
+// CHECK-NOT: tensor.insert_slice
+// CHECK: scf.for
+// CHECK: tensor.extract_slice
+// CHECK: tensor.expand_shape
+// CHECK: tensor.insert_slice
+// CHECK: scf.yield
+// CHECK-NOT: tensor.extract_slice
+// CHECK-NOT: tensor.insert_slice
+// CHECK: return
+func.func @negative_mismatched_slice_types(%input: tensor<1x2xf32>, %n: index) -> tensor<1x2xf32> {
+ %c0 = arith.constant 0 : index
+ %c1 = arith.constant 1 : index
+ %result = scf.for %i = %c0 to %n step %c1 iter_args(%t = %input) -> tensor<1x2xf32> {
+ %slice = tensor.extract_slice %t[0, 0][1, 2][1, 1]
+ : tensor<1x2xf32> to tensor<2xf32>
+ %expanded = tensor.expand_shape %slice [[0, 1]] output_shape [1, 2]
+ : tensor<2xf32> into tensor<1x2xf32>
+ %out = tensor.insert_slice %expanded into %t[0, 0][1, 2][1, 1]
+ : tensor<1x2xf32> into tensor<1x2xf32>
+ scf.yield %out : tensor<1x2xf32>
+ }
+ return %result : tensor<1x2xf32>
+}
More information about the Mlir-commits
mailing list