[Mlir-commits] [mlir] [mlir][bufferization] Trace write-side subset provenance (PR #228995)
llvmlistbot at llvm.org
llvmlistbot at llvm.org
Mon Oct 5 00:36:40 PDT 2026
llvmorg-github-actions[bot] wrote:
<!--LLVM PR SUMMARY COMMENT-->
@llvm/pr-subscribers-mlir-vector
Author: Ziliang Zhang (ziliangzl)
<details>
<summary>Changes</summary>
Trace write operands back to enclosing subset extractions to rule out conflicts with disjoint reads. Follow only equivalent alias edges and require every possible read/write origin pair to be disjoint.
Add coverage for fill and vector writes through extracted views, equivalent casts, and multiple write origins.
---
Full diff: https://github.com/llvm/llvm-project/pull/228995.diff
3 Files Affected:
- (modified) mlir/lib/Dialect/Bufferization/Transforms/OneShotAnalysis.cpp (+100-47)
- (modified) mlir/test/Dialect/Tensor/one-shot-bufferize.mlir (+77)
- (modified) mlir/test/Dialect/Vector/one-shot-bufferize.mlir (+127)
``````````diff
diff --git a/mlir/lib/Dialect/Bufferization/Transforms/OneShotAnalysis.cpp b/mlir/lib/Dialect/Bufferization/Transforms/OneShotAnalysis.cpp
index e8b3404843217..bb9e4417c00dc 100644
--- a/mlir/lib/Dialect/Bufferization/Transforms/OneShotAnalysis.cpp
+++ b/mlir/lib/Dialect/Bufferization/Transforms/OneShotAnalysis.cpp
@@ -719,6 +719,99 @@ static bool matchesInsertDestination(const AnalysisState &state,
return llvm::all_of(backwardSlice, matchingSubset);
}
+/// Find the nearest subset extractions from which `opOperand` may originate.
+/// Only equivalent alias edges may be crossed: an unknown alias does not
+/// guarantee that an access through the alias stays within the extracted
+/// subset. Return "false" if any origin is not a subset extraction.
+static bool
+findSubsetExtractionOrigins(const AnalysisState &state, OpOperand *opOperand,
+ SmallVectorImpl<SubsetOpInterface> &subsetOrigins) {
+ auto isSubsetExtraction = [](Value value) {
+ return !!value.getDefiningOp<SubsetExtractionOpInterface>();
+ };
+ TraversalConfig config;
+ config.followEquivalentOnly = true;
+ SetVector<Value> origins = state.findValueInReverseUseDefChain(
+ opOperand, isSubsetExtraction, config);
+ if (origins.empty())
+ return false;
+
+ for (Value origin : origins) {
+ auto extraction = origin.getDefiningOp<SubsetExtractionOpInterface>();
+ if (!extraction)
+ return false;
+ subsetOrigins.push_back(cast<SubsetOpInterface>(extraction.getOperation()));
+ }
+ return true;
+}
+
+/// Return "true" if every read subset is disjoint from every written subset.
+static bool areAllDisjointSubsets(ArrayRef<SubsetOpInterface> readSubsets,
+ ArrayRef<SubsetOpInterface> writtenSubsets,
+ const AnalysisState &state) {
+ return llvm::all_of(readSubsets, [&](SubsetOpInterface readSubset) {
+ return llvm::all_of(writtenSubsets, [&](SubsetOpInterface writtenSubset) {
+ return readSubset.operatesOnDisjointSubset(
+ writtenSubset, [&](Value v1, Value v2) {
+ return state.areEquivalentBufferizedValues(v1, v2);
+ });
+ });
+ });
+}
+
+/// Return "true" if the read and write are confined to provably disjoint
+/// subsets. Prefer the exact subsets of extraction/insertion ops. If those
+/// cannot be compared, use enclosing subset extractions as conservative access
+/// bounds.
+static bool areDisjointSubsetAccesses(OpOperand *uRead,
+ OpOperand *uConflictingWrite,
+ const AnalysisState &state) {
+ Operation *readingOp = uRead->getOwner();
+ Operation *conflictingWritingOp = uConflictingWrite->getOwner();
+
+ SmallVector<SubsetOpInterface, 1> directReadSubsets;
+ if (auto extraction = dyn_cast<SubsetExtractionOpInterface>(readingOp))
+ if (uRead == &extraction.getSourceOperand())
+ directReadSubsets.push_back(cast<SubsetOpInterface>(readingOp));
+
+ SmallVector<SubsetOpInterface> readOrigins;
+ bool readOriginsInitialized = false;
+ bool hasValidReadOrigins = false;
+ auto readsAreDisjoint = [&](ArrayRef<SubsetOpInterface> writtenSubsets) {
+ if (!directReadSubsets.empty() &&
+ areAllDisjointSubsets(directReadSubsets, writtenSubsets, state))
+ return true;
+
+ if (!readOriginsInitialized) {
+ hasValidReadOrigins =
+ findSubsetExtractionOrigins(state, uRead, readOrigins);
+ readOriginsInitialized = true;
+ }
+ return hasValidReadOrigins &&
+ areAllDisjointSubsets(readOrigins, writtenSubsets, state);
+ };
+
+ // First try the exact subset written by an insertion op.
+ if (auto insertion =
+ dyn_cast<SubsetInsertionOpInterface>(conflictingWritingOp)) {
+ if (uConflictingWrite == &insertion.getDestinationOperand()) {
+ SmallVector<SubsetOpInterface, 1> writtenSubset = {
+ cast<SubsetOpInterface>(conflictingWritingOp)};
+ if (readsAreDisjoint(writtenSubset))
+ return true;
+ }
+ }
+
+ // The write may happen through a view that was produced by a subset
+ // extraction. The extraction is a safe (possibly imprecise) bound for every
+ // access through that view. There may be multiple possible origins, e.g.,
+ // after an arith.select; all of them must be disjoint from every read origin.
+ SmallVector<SubsetOpInterface> writtenOrigins;
+ if (!findSubsetExtractionOrigins(state, uConflictingWrite, writtenOrigins))
+ return false;
+ return readsAreDisjoint(writtenOrigins);
+}
+
/// Return "true" if the given "read" and potentially conflicting "write" are
/// not conflicting due to their subset relationship. The comments in this
/// function are expressed in terms of tensor.extract_slice/tensor.insert_slice
@@ -770,56 +863,16 @@ static bool areNonConflictingSubsets(OpOperand *uRead,
return true;
}
+ // A read from a subset does not conflict with a write to a disjoint subset
+ // of an equivalent tensor. The actual read or write may be further down an
+ // aliasing use-def chain. Trace such accesses back to subset extractions,
+ // crossing only equivalent alias edges.
+ if (areDisjointSubsetAccesses(uRead, uConflictingWrite, state))
+ return true;
+
// If uConflictingWrite is an InsertSliceOp...
if (auto subsetOp =
dyn_cast<SubsetInsertionOpInterface>(conflictingWritingOp)) {
- if (uConflictingWrite == &subsetOp.getDestinationOperand()) {
- auto writtenSubset = cast<SubsetOpInterface>(conflictingWritingOp);
- auto isDisjointSubset = [&](SubsetOpInterface readSubset) {
- return readSubset.operatesOnDisjointSubset(
- writtenSubset, [&](Value v1, Value v2) {
- return state.areEquivalentBufferizedValues(v1, v2);
- });
- };
- auto isDisjointExtraction = [&](Value value) {
- auto extraction = value.getDefiningOp<SubsetExtractionOpInterface>();
- return extraction && isDisjointSubset(cast<SubsetOpInterface>(
- extraction.getOperation()));
- };
-
- // Example:
- //
- // %0 = tensor.insert_slice %s into %t[0][4][1]
- // %1 = vector.transfer_read %t[%c4], %cst
- //
- // A read from a subset does not conflict with a write to a disjoint
- // subset of an equivalent tensor. Check the operand roles explicitly
- // because not every subset extraction bufferizes to a memory read.
- if (auto extraction = dyn_cast<SubsetExtractionOpInterface>(readingOp)) {
- if (uRead == &extraction.getSourceOperand() &&
- isDisjointSubset(cast<SubsetOpInterface>(readingOp)))
- return true;
- }
-
- // Example:
- //
- // %0 = tensor.insert_slice %s into %t[0][4][1]
- // %1 = tensor.extract_slice %t[4][4][1]
- // return %0, %1
- //
- // The actual read may be further down the aliasing use-def chain. E.g.,
- // tensor.extract_slice is an alias-only op and the read is attributed to
- // a return or another consumer of its result. Trace such reads back to
- // their subset extractions. Every origin must be a disjoint subset; a
- // non-subset leaf or an extraction that may overlap keeps the analysis
- // conservative.
- SetVector<Value> readOrigins =
- state.findValueInReverseUseDefChain(uRead, isDisjointExtraction);
- if (!readOrigins.empty() &&
- llvm::all_of(readOrigins, isDisjointExtraction))
- return true;
- }
-
// As an example, consider the following IR.
//
// %0 = tensor.extract_slice %t[%a, %b][%c, %d][1, 1] {inplace = [true] }
diff --git a/mlir/test/Dialect/Tensor/one-shot-bufferize.mlir b/mlir/test/Dialect/Tensor/one-shot-bufferize.mlir
index 272e6d12e7e75..fa970d008c109 100644
--- a/mlir/test/Dialect/Tensor/one-shot-bufferize.mlir
+++ b/mlir/test/Dialect/Tensor/one-shot-bufferize.mlir
@@ -119,6 +119,83 @@ func.func @unknown_insert_extract(
// -----
+// A write through an extracted view is confined to that subset. It does not
+// conflict with a read through a disjoint view of the same tensor.
+
+// CHECK-LABEL: func @disjoint_fill_through_extract(
+// CHECK-SAME: %[[T:.*]]: memref<8xf32, strided<[?], offset: ?>>
+// CHECK-NOT: memref.alloc
+// CHECK: %[[LHS:.*]] = memref.subview %[[T]][0] [4] [1]
+// CHECK: %[[RHS:.*]] = memref.subview %[[T]][4] [4] [1]
+// CHECK: linalg.fill {{.*}} outs(%[[LHS]]
+// CHECK: return %[[LHS]], %[[RHS]]
+
+// CHECK-ANALYSIS-LABEL: func @disjoint_fill_through_extract(
+// CHECK-ANALYSIS: tensor.extract_slice
+// CHECK-ANALYSIS-SAME: __inplace_operands_attr__ = ["true"]
+// CHECK-ANALYSIS: tensor.extract_slice
+// CHECK-ANALYSIS-SAME: __inplace_operands_attr__ = ["true"]
+// CHECK-ANALYSIS: linalg.fill
+// CHECK-ANALYSIS-SAME: __inplace_operands_attr__ = ["none", "true"]
+func.func @disjoint_fill_through_extract(
+ %t: tensor<8xf32> {bufferization.writable = true}, %value: f32)
+ -> (tensor<4xf32>, tensor<4xf32>) {
+ %lhs = tensor.extract_slice %t[0][4][1]
+ : tensor<8xf32> to tensor<4xf32>
+ %rhs = tensor.extract_slice %t[4][4][1]
+ : tensor<8xf32> to tensor<4xf32>
+ %written = linalg.fill ins(%value : f32) outs(%lhs : tensor<4xf32>)
+ -> tensor<4xf32>
+ return %written, %rhs : tensor<4xf32>, tensor<4xf32>
+}
+
+// -----
+
+// A write through [0, 4) conflicts with a read through [2, 6).
+
+// CHECK-LABEL: func @overlapping_fill_through_extract(
+// CHECK: %[[ALLOC:.*]] = memref.alloc
+
+// CHECK-ANALYSIS-LABEL: func @overlapping_fill_through_extract(
+// CHECK-ANALYSIS: tensor.extract_slice
+// CHECK-ANALYSIS-SAME: __inplace_operands_attr__ = ["false"]
+func.func @overlapping_fill_through_extract(
+ %t: tensor<8xf32> {bufferization.writable = true}, %value: f32)
+ -> (tensor<4xf32>, tensor<4xf32>) {
+ %lhs = tensor.extract_slice %t[0][4][1]
+ : tensor<8xf32> to tensor<4xf32>
+ %rhs = tensor.extract_slice %t[2][4][1]
+ : tensor<8xf32> to tensor<4xf32>
+ %written = linalg.fill ins(%value : f32) outs(%lhs : tensor<4xf32>)
+ -> tensor<4xf32>
+ return %written, %rhs : tensor<4xf32>, tensor<4xf32>
+}
+
+// -----
+
+// Dynamically positioned views remain conservative.
+
+// CHECK-LABEL: func @unknown_fill_through_extract(
+// CHECK: %[[ALLOC:.*]] = memref.alloc
+
+// CHECK-ANALYSIS-LABEL: func @unknown_fill_through_extract(
+// CHECK-ANALYSIS: tensor.extract_slice
+// CHECK-ANALYSIS-SAME: __inplace_operands_attr__ = ["false", "none"]
+func.func @unknown_fill_through_extract(
+ %t: tensor<8xf32> {bufferization.writable = true}, %value: f32,
+ %write_idx: index, %read_idx: index)
+ -> (tensor<4xf32>, tensor<4xf32>) {
+ %lhs = tensor.extract_slice %t[%write_idx][4][1]
+ : tensor<8xf32> to tensor<4xf32>
+ %rhs = tensor.extract_slice %t[%read_idx][4][1]
+ : tensor<8xf32> to tensor<4xf32>
+ %written = linalg.fill ins(%value : f32) outs(%lhs : tensor<4xf32>)
+ -> tensor<4xf32>
+ return %written, %rhs : tensor<4xf32>, tensor<4xf32>
+}
+
+// -----
+
// CHECK-LABEL: func private @insert_slice_fun
// CHECK-SAME: %[[A0:[a-zA-Z0-9]*]]: memref<?xf32, strided<[?], offset: ?>>,
// CHECK-SAME: %[[A1:[a-zA-Z0-9]*]]: memref<?xf32, strided<[?], offset: ?>>,
diff --git a/mlir/test/Dialect/Vector/one-shot-bufferize.mlir b/mlir/test/Dialect/Vector/one-shot-bufferize.mlir
index a190c94753008..971a2ae20dcd7 100644
--- a/mlir/test/Dialect/Vector/one-shot-bufferize.mlir
+++ b/mlir/test/Dialect/Vector/one-shot-bufferize.mlir
@@ -125,3 +125,130 @@ func.func @unknown_transfer_read_write(
: tensor<8xf32>, vector<4xf32>
return %written, %read : tensor<8xf32>, vector<4xf32>
}
+
+// -----
+
+// A partial vector write through [0, 4) does not conflict with a read through
+// [4, 8). Both the allocation and the copy for the extracted write view can be
+// avoided.
+
+// CHECK-LABEL: func @disjoint_transfer_write_through_extract(
+// CHECK-SAME: %[[T:.*]]: memref<8xf32, strided<[?], offset: ?>>
+// CHECK-NOT: memref.alloc
+// CHECK-NOT: memref.copy
+// CHECK: %[[LHS:.*]] = memref.subview %[[T]][0] [4] [1]
+// CHECK: %[[RHS:.*]] = memref.subview %[[T]][4] [4] [1]
+// CHECK: vector.transfer_write %{{.*}}, %[[LHS]]
+// CHECK: return %[[LHS]], %[[RHS]]
+
+// CHECK-ANALYSIS-LABEL: func @disjoint_transfer_write_through_extract(
+// CHECK-ANALYSIS: tensor.extract_slice
+// CHECK-ANALYSIS-SAME: __inplace_operands_attr__ = ["true"]
+// CHECK-ANALYSIS: tensor.extract_slice
+// CHECK-ANALYSIS-SAME: __inplace_operands_attr__ = ["true"]
+// CHECK-ANALYSIS: vector.transfer_write
+// CHECK-ANALYSIS-SAME: __inplace_operands_attr__ = ["none", "true", "none"]
+func.func @disjoint_transfer_write_through_extract(
+ %t: tensor<8xf32> {bufferization.writable = true},
+ %v: vector<2xf32>) -> (tensor<4xf32>, tensor<4xf32>) {
+ %c0 = arith.constant 0 : index
+ %lhs = tensor.extract_slice %t[0][4][1]
+ : tensor<8xf32> to tensor<4xf32>
+ %rhs = tensor.extract_slice %t[4][4][1]
+ : tensor<8xf32> to tensor<4xf32>
+ %written = vector.transfer_write %v, %lhs[%c0]
+ : vector<2xf32>, tensor<4xf32>
+ return %written, %rhs : tensor<4xf32>, tensor<4xf32>
+}
+
+// -----
+
+// Equivalent aliases between the extraction and the write preserve subset
+// containment.
+
+// CHECK-LABEL: func @disjoint_transfer_write_through_extract_and_cast(
+// CHECK-NOT: memref.alloc
+// CHECK-NOT: memref.copy
+// CHECK: vector.transfer_write
+
+// CHECK-ANALYSIS-LABEL: func @disjoint_transfer_write_through_extract_and_cast(
+// CHECK-ANALYSIS: tensor.extract_slice
+// CHECK-ANALYSIS-SAME: __inplace_operands_attr__ = ["true"]
+// CHECK-ANALYSIS: tensor.cast
+// CHECK-ANALYSIS-SAME: __inplace_operands_attr__ = ["true"]
+// CHECK-ANALYSIS: tensor.cast
+// CHECK-ANALYSIS-SAME: __inplace_operands_attr__ = ["true"]
+// CHECK-ANALYSIS: tensor.extract_slice
+// CHECK-ANALYSIS-SAME: __inplace_operands_attr__ = ["true"]
+func.func @disjoint_transfer_write_through_extract_and_cast(
+ %t: tensor<8xf32> {bufferization.writable = true},
+ %v: vector<2xf32>) -> (tensor<4xf32>, tensor<4xf32>) {
+ %c0 = arith.constant 0 : index
+ %lhs = tensor.extract_slice %t[0][4][1]
+ : tensor<8xf32> to tensor<4xf32>
+ %cast0 = tensor.cast %lhs : tensor<4xf32> to tensor<?xf32>
+ %cast1 = tensor.cast %cast0 : tensor<?xf32> to tensor<4xf32>
+ %rhs = tensor.extract_slice %t[4][4][1]
+ : tensor<8xf32> to tensor<4xf32>
+ %written = vector.transfer_write %v, %cast1[%c0]
+ : vector<2xf32>, tensor<4xf32>
+ return %written, %rhs : tensor<4xf32>, tensor<4xf32>
+}
+
+// -----
+
+// Every possible write origin is disjoint from the read.
+
+// CHECK-LABEL: func @multiple_disjoint_write_origins(
+// CHECK-NOT: memref.alloc
+// CHECK-NOT: memref.copy
+// CHECK: vector.transfer_write
+
+// CHECK-ANALYSIS-LABEL: func @multiple_disjoint_write_origins(
+// CHECK-ANALYSIS: tensor.extract_slice
+// CHECK-ANALYSIS-SAME: __inplace_operands_attr__ = ["true"]
+// CHECK-ANALYSIS: tensor.extract_slice
+// CHECK-ANALYSIS-SAME: __inplace_operands_attr__ = ["true"]
+// CHECK-ANALYSIS: tensor.extract_slice
+// CHECK-ANALYSIS-SAME: __inplace_operands_attr__ = ["true"]
+func.func @multiple_disjoint_write_origins(
+ %t: tensor<12xf32> {bufferization.writable = true},
+ %v: vector<2xf32>, %cond: i1) -> (tensor<4xf32>, tensor<4xf32>) {
+ %c0 = arith.constant 0 : index
+ %lhs0 = tensor.extract_slice %t[0][4][1]
+ : tensor<12xf32> to tensor<4xf32>
+ %lhs1 = tensor.extract_slice %t[4][4][1]
+ : tensor<12xf32> to tensor<4xf32>
+ %rhs = tensor.extract_slice %t[8][4][1]
+ : tensor<12xf32> to tensor<4xf32>
+ %selected = arith.select %cond, %lhs0, %lhs1 : tensor<4xf32>
+ %written = vector.transfer_write %v, %selected[%c0]
+ : vector<2xf32>, tensor<4xf32>
+ return %written, %rhs : tensor<4xf32>, tensor<4xf32>
+}
+
+// -----
+
+// One possible write origin overlaps the read, so the proof must fail.
+
+// CHECK-LABEL: func @one_overlapping_write_origin(
+// CHECK: %[[ALLOC:.*]] = memref.alloc
+
+// CHECK-ANALYSIS-LABEL: func @one_overlapping_write_origin(
+// CHECK-ANALYSIS: tensor.extract_slice
+// CHECK-ANALYSIS-SAME: __inplace_operands_attr__ = ["false"]
+func.func @one_overlapping_write_origin(
+ %t: tensor<12xf32> {bufferization.writable = true},
+ %v: vector<2xf32>, %cond: i1) -> (tensor<4xf32>, tensor<4xf32>) {
+ %c0 = arith.constant 0 : index
+ %lhs0 = tensor.extract_slice %t[0][4][1]
+ : tensor<12xf32> to tensor<4xf32>
+ %lhs1 = tensor.extract_slice %t[8][4][1]
+ : tensor<12xf32> to tensor<4xf32>
+ %rhs = tensor.extract_slice %t[8][4][1]
+ : tensor<12xf32> to tensor<4xf32>
+ %selected = arith.select %cond, %lhs0, %lhs1 : tensor<4xf32>
+ %written = vector.transfer_write %v, %selected[%c0]
+ : vector<2xf32>, tensor<4xf32>
+ return %written, %rhs : tensor<4xf32>, tensor<4xf32>
+}
``````````
</details>
https://github.com/llvm/llvm-project/pull/228995
More information about the Mlir-commits
mailing list