[Mlir-commits] [mlir] [mlir][bufferization] Handle disjoint subsets in One-Shot Analysis (PR #224184)
Ziliang Zhang
llvmlistbot at llvm.org
Wed Sep 16 20:29:16 PDT 2026
https://github.com/ziliangzl updated https://github.com/llvm/llvm-project/pull/224184
>From e651cfd72018150360e0492ecf0d30a100f3592d Mon Sep 17 00:00:00 2001
From: Ziliang Zhang <zzl.coding at gmail.com>
Date: Thu, 17 Sep 2026 10:53:05 +0800
Subject: [PATCH 1/2] [mlir][bufferization] Handle disjoint subsets in One-Shot
Analysis
Teach One-Shot Analysis to ignore RaW conflicts between provably
disjoint subset reads and writes. Trace aliasing reads back to subset
extraction ops to handle alias-only operations such as
tensor.extract_slice.
Add tensor and vector tests for disjoint, overlapping, and unknown
subsets.
---
.../Transforms/OneShotAnalysis.cpp | 39 ++++++++++
.../Dialect/Tensor/one-shot-bufferize.mlir | 69 ++++++++++++++++
.../Dialect/Vector/one-shot-bufferize.mlir | 78 +++++++++++++++++++
3 files changed, 186 insertions(+)
diff --git a/mlir/lib/Dialect/Bufferization/Transforms/OneShotAnalysis.cpp b/mlir/lib/Dialect/Bufferization/Transforms/OneShotAnalysis.cpp
index e948a02cc415d..ffdefa79e419f 100644
--- a/mlir/lib/Dialect/Bufferization/Transforms/OneShotAnalysis.cpp
+++ b/mlir/lib/Dialect/Bufferization/Transforms/OneShotAnalysis.cpp
@@ -518,6 +518,45 @@ static bool areNonConflictingSubsets(OpOperand *uRead,
Operation *readingOp = uRead->getOwner();
Operation *conflictingWritingOp = uConflictingWrite->getOwner();
+ auto subsetInsertion =
+ dyn_cast<SubsetInsertionOpInterface>(conflictingWritingOp);
+ if (subsetInsertion &&
+ uConflictingWrite == &subsetInsertion.getDestinationOperand()) {
+ auto writtenSubset = cast<SubsetOpInterface>(conflictingWritingOp);
+ auto isDisjointExtraction = [&](Value value) {
+ auto extraction = value.getDefiningOp<SubsetExtractionOpInterface>();
+ return extraction &&
+ cast<SubsetOpInterface>(extraction.getOperation())
+ .operatesOnDisjointSubset(
+ writtenSubset, [&](Value v1, Value v2) {
+ return state.areEquivalentBufferizedValues(v1, v2);
+ });
+ };
+
+ // 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 a
+ // subset extraction does not itself bufferize to a memory read.
+ if (auto extraction = dyn_cast<SubsetExtractionOpInterface>(readingOp)) {
+ if (uRead == &extraction.getSourceOperand() &&
+ cast<SubsetOpInterface>(readingOp).operatesOnDisjointSubset(
+ writtenSubset, [&](Value v1, Value v2) {
+ return state.areEquivalentBufferizedValues(v1, v2);
+ }))
+ return true;
+ }
+
+ // 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;
+ }
+
// Special rules for matching ExtractSliceOp/InsertSliceOp pairs. If
// uRead is an InsertSliceOp...
if (auto subsetOp = dyn_cast<SubsetInsertionOpInterface>(readingOp)) {
diff --git a/mlir/test/Dialect/Tensor/one-shot-bufferize.mlir b/mlir/test/Dialect/Tensor/one-shot-bufferize.mlir
index bd7115cb3a9c5..5b59176aed05a 100644
--- a/mlir/test/Dialect/Tensor/one-shot-bufferize.mlir
+++ b/mlir/test/Dialect/Tensor/one-shot-bufferize.mlir
@@ -1,4 +1,5 @@
// RUN: mlir-opt %s -one-shot-bufferize="bufferize-function-boundaries" -drop-equivalent-buffer-results -split-input-file | FileCheck %s
+// RUN: mlir-opt %s -one-shot-bufferize="test-analysis-only bufferize-function-boundaries" -split-input-file | FileCheck %s --check-prefix=CHECK-ANALYSIS
// Run fuzzer with different seeds.
// RUN: mlir-opt %s -one-shot-bufferize="test-analysis-only analysis-heuristic=fuzzer analysis-fuzzer-seed=23 bufferize-function-boundaries" -split-input-file -o /dev/null
@@ -8,6 +9,74 @@
// Test bufferization using memref types that have no layout map.
// RUN: mlir-opt %s -one-shot-bufferize="unknown-type-conversion=identity-layout-map bufferize-function-boundaries" -split-input-file -o /dev/null
+// A write to [0, 4) does not conflict with a read from [4, 8), even though
+// tensor.extract_slice itself is alias-only and the actual read is the return.
+
+// CHECK-LABEL: func @disjoint_insert_extract(
+// CHECK-NOT: memref.alloc
+// CHECK: memref.subview
+// CHECK: memref.copy
+// CHECK: memref.subview
+// CHECK-NOT: memref.alloc
+
+// CHECK-ANALYSIS-LABEL: func @disjoint_insert_extract(
+// CHECK-ANALYSIS: tensor.insert_slice
+// CHECK-ANALYSIS-SAME: __inplace_operands_attr__ = ["true", "true"]
+func.func @disjoint_insert_extract(
+ %t: tensor<8xf32> {bufferization.writable = true},
+ %source: tensor<4xf32>) -> (tensor<8xf32>, tensor<4xf32>) {
+ %written = tensor.insert_slice %source into %t[0][4][1]
+ : tensor<4xf32> into tensor<8xf32>
+ %read = tensor.extract_slice %t[4][4][1]
+ : tensor<8xf32> to tensor<4xf32>
+ return %written, %read : tensor<8xf32>, tensor<4xf32>
+}
+
+// -----
+
+// A write to [0, 4) conflicts with a read from [2, 6).
+
+// CHECK-LABEL: func @overlapping_insert_extract(
+// CHECK: %[[ALLOC:.*]] = memref.alloc
+// CHECK: memref.copy %{{.*}}, %[[ALLOC]]
+
+// CHECK-ANALYSIS-LABEL: func @overlapping_insert_extract(
+// CHECK-ANALYSIS: tensor.insert_slice
+// CHECK-ANALYSIS-SAME: __inplace_operands_attr__ = ["true", "false"]
+func.func @overlapping_insert_extract(
+ %t: tensor<8xf32> {bufferization.writable = true},
+ %source: tensor<4xf32>) -> (tensor<8xf32>, tensor<4xf32>) {
+ %written = tensor.insert_slice %source into %t[0][4][1]
+ : tensor<4xf32> into tensor<8xf32>
+ %read = tensor.extract_slice %t[2][4][1]
+ : tensor<8xf32> to tensor<4xf32>
+ return %written, %read : tensor<8xf32>, tensor<4xf32>
+}
+
+// -----
+
+// Dynamic offsets that cannot be proven disjoint remain conservative.
+
+// CHECK-LABEL: func @unknown_insert_extract(
+// CHECK: %[[ALLOC:.*]] = memref.alloc
+// CHECK: memref.copy %{{.*}}, %[[ALLOC]]
+
+// CHECK-ANALYSIS-LABEL: func @unknown_insert_extract(
+// CHECK-ANALYSIS: tensor.insert_slice
+// CHECK-ANALYSIS-SAME: __inplace_operands_attr__ = ["true", "false", "none"]
+func.func @unknown_insert_extract(
+ %t: tensor<8xf32> {bufferization.writable = true},
+ %source: tensor<4xf32>, %write_idx: index,
+ %read_idx: index) -> (tensor<8xf32>, tensor<4xf32>) {
+ %written = tensor.insert_slice %source into %t[%write_idx][4][1]
+ : tensor<4xf32> into tensor<8xf32>
+ %read = tensor.extract_slice %t[%read_idx][4][1]
+ : tensor<8xf32> to tensor<4xf32>
+ return %written, %read : tensor<8xf32>, 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 c2d699b9b013a..004762f939d0d 100644
--- a/mlir/test/Dialect/Vector/one-shot-bufferize.mlir
+++ b/mlir/test/Dialect/Vector/one-shot-bufferize.mlir
@@ -41,3 +41,81 @@ func.func @reading_xfer_write(%t: tensor<5x10xf32>, %v: vector<4x11xf32>) -> ten
%1 = vector.transfer_write %v, %t[%c0, %c0] : vector<4x11xf32>, tensor<5x10xf32>
return %1 : tensor<5x10xf32>
}
+
+// -----
+
+// A write to [0, 4) does not conflict with a read from [4, 8), so the write
+// can bufferize in-place.
+
+// CHECK-LABEL: func @disjoint_transfer_read_write(
+// CHECK-SAME: %[[T:.*]]: memref<8xf32, strided<[?], offset: ?>>
+// CHECK-NOT: memref.alloc
+// CHECK-NOT: memref.copy
+// CHECK: vector.transfer_write %{{.*}}, %[[T]][%{{.*}}]
+// CHECK: vector.transfer_read %[[T]][%{{.*}}]
+// CHECK: return %[[T]],
+
+// CHECK-ANALYSIS-LABEL: func @disjoint_transfer_read_write(
+// CHECK-ANALYSIS: vector.transfer_write
+// CHECK-ANALYSIS-SAME: __inplace_operands_attr__ = ["none", "true", "none"]
+func.func @disjoint_transfer_read_write(
+ %t: tensor<8xf32> {bufferization.writable = true},
+ %v: vector<4xf32>) -> (tensor<8xf32>, vector<4xf32>) {
+ %c0 = arith.constant 0 : index
+ %c4 = arith.constant 4 : index
+ %pad = arith.constant 0.0 : f32
+ %written = vector.transfer_write %v, %t[%c0]
+ : vector<4xf32>, tensor<8xf32>
+ %read = vector.transfer_read %t[%c4], %pad
+ : tensor<8xf32>, vector<4xf32>
+ return %written, %read : tensor<8xf32>, vector<4xf32>
+}
+
+// -----
+
+// A write to [0, 4) conflicts with a read from [2, 6), so the write must
+// bufferize out-of-place.
+
+// CHECK-LABEL: func @overlapping_transfer_read_write(
+// CHECK: %[[ALLOC:.*]] = memref.alloc
+// CHECK: memref.copy %{{.*}}, %[[ALLOC]]
+
+// CHECK-ANALYSIS-LABEL: func @overlapping_transfer_read_write(
+// CHECK-ANALYSIS: vector.transfer_write
+// CHECK-ANALYSIS-SAME: __inplace_operands_attr__ = ["none", "false", "none"]
+func.func @overlapping_transfer_read_write(
+ %t: tensor<8xf32> {bufferization.writable = true},
+ %v: vector<4xf32>) -> (tensor<8xf32>, vector<4xf32>) {
+ %c0 = arith.constant 0 : index
+ %c2 = arith.constant 2 : index
+ %pad = arith.constant 0.0 : f32
+ %written = vector.transfer_write %v, %t[%c0]
+ : vector<4xf32>, tensor<8xf32>
+ %read = vector.transfer_read %t[%c2], %pad
+ : tensor<8xf32>, vector<4xf32>
+ return %written, %read : tensor<8xf32>, vector<4xf32>
+}
+
+// -----
+
+// If the subsets cannot be proven disjoint, the write must bufferize
+// out-of-place.
+
+// CHECK-LABEL: func @unknown_transfer_read_write(
+// CHECK: %[[ALLOC:.*]] = memref.alloc
+// CHECK: memref.copy %{{.*}}, %[[ALLOC]]
+
+// CHECK-ANALYSIS-LABEL: func @unknown_transfer_read_write(
+// CHECK-ANALYSIS: vector.transfer_write
+// CHECK-ANALYSIS-SAME: __inplace_operands_attr__ = ["none", "false", "none"]
+func.func @unknown_transfer_read_write(
+ %t: tensor<8xf32> {bufferization.writable = true},
+ %v: vector<4xf32>, %write_idx: index,
+ %read_idx: index) -> (tensor<8xf32>, vector<4xf32>) {
+ %pad = arith.constant 0.0 : f32
+ %written = vector.transfer_write %v, %t[%write_idx]
+ : vector<4xf32>, tensor<8xf32>
+ %read = vector.transfer_read %t[%read_idx], %pad
+ : tensor<8xf32>, vector<4xf32>
+ return %written, %read : tensor<8xf32>, vector<4xf32>
+}
>From 5d5f83d783456a9fe17403f7c9c863b58051be74 Mon Sep 17 00:00:00 2001
From: Ziliang Zhang <zzl.coding at gmail.com>
Date: Thu, 17 Sep 2026 11:26:22 +0800
Subject: [PATCH 2/2] clang-format
---
mlir/lib/Dialect/Bufferization/Transforms/OneShotAnalysis.cpp | 3 +--
1 file changed, 1 insertion(+), 2 deletions(-)
diff --git a/mlir/lib/Dialect/Bufferization/Transforms/OneShotAnalysis.cpp b/mlir/lib/Dialect/Bufferization/Transforms/OneShotAnalysis.cpp
index ffdefa79e419f..c7e6acc7fa6ba 100644
--- a/mlir/lib/Dialect/Bufferization/Transforms/OneShotAnalysis.cpp
+++ b/mlir/lib/Dialect/Bufferization/Transforms/OneShotAnalysis.cpp
@@ -552,8 +552,7 @@ static bool areNonConflictingSubsets(OpOperand *uRead,
// 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))
+ if (!readOrigins.empty() && llvm::all_of(readOrigins, isDisjointExtraction))
return true;
}
More information about the Mlir-commits
mailing list