Skip to content

Commit 855a11e

Browse files
[mlir][bufferize][NFC] Rename allow-return-memref to allow-return-allocs
Also clean up/split test cases. Differential Revision: https://reviews.llvm.org/D121522
1 parent a7c08bc commit 855a11e

18 files changed

+71
-65
lines changed

mlir/include/mlir/Dialect/Bufferization/Transforms/OneShotAnalysis.h

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -46,7 +46,7 @@ struct OneShotBufferizationOptions : public BufferizationOptions {
4646

4747
/// Specifies whether returning newly allocated memrefs should be allowed.
4848
/// Otherwise, a pass failure is triggered.
49-
bool allowReturnMemref = false;
49+
bool allowReturnAllocs = false;
5050
};
5151

5252
/// The BufferizationAliasInfo class maintains a list of buffer aliases and

mlir/include/mlir/Dialect/Bufferization/Transforms/Passes.td

Lines changed: 3 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -173,7 +173,7 @@ def OneShotBufferize : Pass<"one-shot-bufferize", "ModuleOp"> {
173173
One-Shot Bufferize deallocates all buffers that it allocates. Yielding newly
174174
allocated buffers from a block is not supported yet and such IR will be
175175
rejected. For testing purposes and compatibility with partial bufferization,
176-
One-Shot Bufferize can be run with `allow-return-memref=1 create-dealloc=0`
176+
One-Shot Bufferize can be run with `allow-return-allocs=1 create-dealloc=0`
177177
to allow such IR.
178178

179179
One-Shot Bufferize will by default reject IR that contains non-bufferizable
@@ -202,9 +202,9 @@ def OneShotBufferize : Pass<"one-shot-bufferize", "ModuleOp"> {
202202
Bufferize chose to insert a certain buffer copy.
203203
}];
204204
let options = [
205-
Option<"allowReturnMemref", "allow-return-memref", "bool",
205+
Option<"allowReturnAllocs", "allow-return-allocs", "bool",
206206
/*default=*/"false",
207-
"Allows the return of memrefs (for testing purposes only)">,
207+
"Allows the return of new allocations (for testing purposes only)">,
208208
Option<"allowUnknownOps", "allow-unknown-ops", "bool",
209209
/*default=*/"false",
210210
"Allows unknown (not bufferizable) ops in the input IR.">,

mlir/include/mlir/Dialect/Linalg/Passes.td

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -43,9 +43,9 @@ def LinalgComprehensiveModuleBufferize :
4343
Option<"printConflicts", "print-conflicts", "bool",
4444
/*default=*/"false",
4545
"Annotates IR with RaW conflicts. Requires test-analysis-only.">,
46-
Option<"allowReturnMemref", "allow-return-memref", "bool",
46+
Option<"allowReturnAllocs", "allow-return-allocs", "bool",
4747
/*default=*/"false",
48-
"Allows the return of memrefs (for testing purposes only)">,
48+
"Allows the return of new allocations (for testing purposes only)">,
4949
Option<"allowUnknownOps", "allow-unknown-ops", "bool",
5050
/*default=*/"false",
5151
"Allows unknown (not bufferizable) ops in the input IR.">,

mlir/lib/Dialect/Bufferization/Transforms/Bufferize.cpp

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -165,7 +165,7 @@ struct OneShotBufferizePass
165165
if (!options) {
166166
// Make new bufferization options if none were provided when creating the
167167
// pass.
168-
opt.allowReturnMemref = allowReturnMemref;
168+
opt.allowReturnAllocs = allowReturnAllocs;
169169
opt.allowUnknownOps = allowUnknownOps;
170170
opt.analysisFuzzerSeed = analysisFuzzerSeed;
171171
opt.createDeallocs = createDeallocs;

mlir/lib/Dialect/Bufferization/Transforms/OneShotAnalysis.cpp

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -37,7 +37,7 @@
3737
//
3838
// This analysis caters to high-performance codegen where buffer reuse is deemed
3939
// critical: the analysis should fail if the bufferized form of the function
40-
// needs to return a buffer, unless `allowReturnMemref` is enabled.
40+
// needs to return a buffer, unless `allowReturnAllocs` is enabled.
4141

4242
#include "mlir/Dialect/Bufferization/Transforms/OneShotAnalysis.h"
4343

@@ -811,7 +811,7 @@ LogicalResult bufferization::analyzeOp(Operation *op,
811811
}
812812

813813
bool failedAnalysis = false;
814-
if (!options.allowReturnMemref) {
814+
if (!options.allowReturnAllocs) {
815815
SmallVector<Operation *> newOps;
816816
failedAnalysis |=
817817
failed(assertDestinationPassingStyle(op, state, aliasInfo, newOps));

mlir/lib/Dialect/Linalg/ComprehensiveBufferize/ModuleBufferization.cpp

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -1045,7 +1045,7 @@ LogicalResult mlir::linalg::comprehensive_bufferize::runModuleBufferize(
10451045
if (failed(bufferizeFuncOpBoundary(funcOp, rewriter, bufferizationState)))
10461046
return failure();
10471047

1048-
if (!options.allowReturnMemref &&
1048+
if (!options.allowReturnAllocs &&
10491049
llvm::any_of(funcOp.getType().getResults(), [](Type t) {
10501050
return t.isa<MemRefType, UnrankedMemRefType>();
10511051
})) {

mlir/lib/Dialect/Linalg/Transforms/ComprehensiveBufferizePass.cpp

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -89,7 +89,7 @@ void LinalgComprehensiveModuleBufferize::runOnOperation() {
8989
return success();
9090
};
9191
}
92-
opt.allowReturnMemref = allowReturnMemref;
92+
opt.allowReturnAllocs = allowReturnAllocs;
9393
opt.allowUnknownOps = allowUnknownOps;
9494
opt.analysisFuzzerSeed = analysisFuzzerSeed;
9595
opt.createDeallocs = createDeallocs;
Lines changed: 32 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,32 @@
1+
// RUN: mlir-opt %s -one-shot-bufferize="allow-return-allocs allow-unknown-ops" -split-input-file | FileCheck %s
2+
3+
// Run fuzzer with different seeds.
4+
// RUN: mlir-opt %s -one-shot-bufferize="allow-return-allocs test-analysis-only analysis-fuzzer-seed=23" -split-input-file -o /dev/null
5+
// RUN: mlir-opt %s -one-shot-bufferize="allow-return-allocs test-analysis-only analysis-fuzzer-seed=59" -split-input-file -o /dev/null
6+
// RUN: mlir-opt %s -one-shot-bufferize="allow-return-allocs test-analysis-only analysis-fuzzer-seed=91" -split-input-file -o /dev/null
7+
8+
// CHECK-LABEL: func @buffer_not_deallocated(
9+
// CHECK-SAME: %[[t:.*]]: tensor<?xf32>
10+
func @buffer_not_deallocated(%t : tensor<?xf32>, %c : i1) -> tensor<?xf32> {
11+
// CHECK: %[[r:.*]] = scf.if %{{.*}} {
12+
%r = scf.if %c -> tensor<?xf32> {
13+
// CHECK: %[[some_op:.*]] = "test.some_op"
14+
// CHECK: %[[alloc:.*]] = memref.alloc(%[[some_op]])
15+
// CHECK: %[[casted:.*]] = memref.cast %[[alloc]]
16+
// CHECK-NOT: dealloc
17+
// CHECK: scf.yield %[[casted]]
18+
%sz = "test.some_op"() : () -> (index)
19+
%0 = linalg.init_tensor[%sz] : tensor<?xf32>
20+
scf.yield %0 : tensor<?xf32>
21+
} else {
22+
// CHECK: } else {
23+
// CHECK: %[[m:.*]] = bufferization.to_memref %[[t]]
24+
// CHECK: scf.yield %[[m]]
25+
scf.yield %t : tensor<?xf32>
26+
}
27+
// CHECK: }
28+
// CHECK-NOT: dealloc
29+
// CHECK: %[[r_tensor:.*]] = bufferization.to_tensor %[[r]]
30+
// CHECK: return %[[r_tensor]]
31+
return %r : tensor<?xf32>
32+
}

mlir/test/Dialect/Bufferization/Transforms/one-shot-bufferize-compat.mlir

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -1,10 +1,10 @@
11
// RUN: mlir-opt %s \
2-
// RUN: -one-shot-bufferize="allow-return-memref allow-unknown-ops create-deallocs=0" \
2+
// RUN: -one-shot-bufferize="allow-unknown-ops create-deallocs=0" \
33
// RUN: -split-input-file | \
44
// RUN: FileCheck %s --check-prefix=CHECK-NODEALLOC
55

66
// RUN: mlir-opt %s \
7-
// RUN: -one-shot-bufferize="allow-return-memref allow-unknown-ops create-deallocs=0" \
7+
// RUN: -one-shot-bufferize="allow-unknown-ops create-deallocs=0" \
88
// RUN: -buffer-deallocation | \
99
// RUN: FileCheck %s --check-prefix=CHECK-BUFFERDEALLOC
1010

mlir/test/Dialect/Bufferization/Transforms/one-shot-bufferize-partial.mlir

Lines changed: 7 additions & 7 deletions
Original file line numberDiff line numberDiff line change
@@ -1,15 +1,15 @@
1-
// RUN: mlir-opt %s -allow-unregistered-dialect -one-shot-bufferize="allow-return-memref allow-unknown-ops" -split-input-file | FileCheck %s
1+
// RUN: mlir-opt %s -allow-unregistered-dialect -one-shot-bufferize="allow-return-allocs allow-unknown-ops" -split-input-file | FileCheck %s
22

33
// Test bufferization using memref types that have no layout map.
4-
// RUN: mlir-opt %s -allow-unregistered-dialect -one-shot-bufferize="allow-return-memref allow-unknown-ops fully-dynamic-layout-maps=0" -split-input-file | FileCheck %s --check-prefix=CHECK-NO-LAYOUT-MAP
4+
// RUN: mlir-opt %s -allow-unregistered-dialect -one-shot-bufferize="allow-return-allocs allow-unknown-ops fully-dynamic-layout-maps=0" -split-input-file | FileCheck %s --check-prefix=CHECK-NO-LAYOUT-MAP
55

66
// Run fuzzer with different seeds.
7-
// RUN: mlir-opt %s -allow-unregistered-dialect -one-shot-bufferize="allow-return-memref test-analysis-only analysis-fuzzer-seed=23" -split-input-file -o /dev/null
8-
// RUN: mlir-opt %s -allow-unregistered-dialect -one-shot-bufferize="allow-return-memref test-analysis-only analysis-fuzzer-seed=59" -split-input-file -o /dev/null
9-
// RUN: mlir-opt %s -allow-unregistered-dialect -one-shot-bufferize="allow-return-memref test-analysis-only analysis-fuzzer-seed=91" -split-input-file -o /dev/null
7+
// RUN: mlir-opt %s -allow-unregistered-dialect -one-shot-bufferize="allow-return-allocs test-analysis-only analysis-fuzzer-seed=23" -split-input-file -o /dev/null
8+
// RUN: mlir-opt %s -allow-unregistered-dialect -one-shot-bufferize="allow-return-allocs test-analysis-only analysis-fuzzer-seed=59" -split-input-file -o /dev/null
9+
// RUN: mlir-opt %s -allow-unregistered-dialect -one-shot-bufferize="allow-return-allocs test-analysis-only analysis-fuzzer-seed=91" -split-input-file -o /dev/null
1010

11-
// RUN: mlir-opt %s -allow-unregistered-dialect -one-shot-bufferize="dialect-filter=tensor allow-unknown-ops allow-return-memref" -canonicalize -split-input-file | FileCheck %s --check-prefix=CHECK-TENSOR
12-
// RUN: mlir-opt %s -allow-unregistered-dialect -one-shot-bufferize="dialect-filter=scf allow-unknown-ops allow-return-memref" -canonicalize -split-input-file | FileCheck %s --check-prefix=CHECK-SCF
11+
// RUN: mlir-opt %s -allow-unregistered-dialect -one-shot-bufferize="dialect-filter=tensor allow-unknown-ops allow-return-allocs" -canonicalize -split-input-file | FileCheck %s --check-prefix=CHECK-TENSOR
12+
// RUN: mlir-opt %s -allow-unregistered-dialect -one-shot-bufferize="dialect-filter=scf allow-unknown-ops allow-return-allocs" -canonicalize -split-input-file | FileCheck %s --check-prefix=CHECK-SCF
1313

1414
// CHECK: #[[$MAP:.*]] = affine_map<(d0)[s0, s1] -> (d0 * s1 + s0)>
1515

0 commit comments

Comments
 (0)