Skip to content

Commit 0ee1c03

Browse files
matthias-springermatthias-springer
committed
[mlir][bufferize] Remove hoisting functionality from One-Shot Bufferize
The same functionality is already provided by `-buffer-hoisting` and `-buffer-loop-hoisting`. Differential Revision: https://reviews.llvm.org/D126251
1 parent d3187dd commit 0ee1c03

File tree

11 files changed

+35
-128
lines changed

11 files changed

+35
-128
lines changed

‎mlir/include/mlir/Dialect/Bufferization/IR/BufferizableOpInterface.h

Lines changed: 0 additions & 12 deletions
Original file line numberDiff line numberDiff line change
@@ -280,14 +280,6 @@ struct BufferizationOptions {
280280
/// computation. Whether this pays off or not can be very input IR-specific.
281281
bool alwaysAliasingWithDest = true;
282282

283-
/// If set to `true`, try to hoist allocations out of blocks as much as
284-
/// possible. An allocation is not hoisted across allocation hoisting barriers
285-
/// as indicated by `BufferizableOpInterface::isAllocationHoistingBarrier`.
286-
///
287-
/// Examples of allocation hoisting barriers are parallel loops or ops where
288-
/// SSA values cannot be captured from the outside.
289-
bool hoistAllocations = true;
290-
291283
/// Buffer alignment for new memory allocations.
292284
unsigned int bufferAlignment = 128;
293285

@@ -618,10 +610,6 @@ BaseMemRefType
618610
getMemRefTypeWithStaticIdentityLayout(TensorType tensorType,
619611
Attribute memorySpace = {});
620612

621-
/// Try to hoist all new buffer allocations until the next hoisting barrier.
622-
LogicalResult hoistBufferAllocations(Operation *op,
623-
const BufferizationOptions &options);
624-
625613
/// Create alloc/dealloc ops as specified in the bufferization options. If
626614
/// `onlyLeakingAlloc`, only those buffer allocations are processed for which no
627615
/// buffer deallocation can be created. `changed` is set to `true` if the IR was

‎mlir/include/mlir/Dialect/Bufferization/IR/BufferizableOpInterface.td

Lines changed: 0 additions & 16 deletions
Original file line numberDiff line numberDiff line change
@@ -259,22 +259,6 @@ def BufferizableOpInterface : OpInterface<"BufferizableOpInterface"> {
259259
return value.isa<OpResult>();
260260
}]
261261
>,
262-
InterfaceMethod<
263-
/*desc=*/[{
264-
Return `true` if the op is an allocation hoisting barrier. Buffer
265-
allocations will never be beyond such ops. E.g., ops with certain
266-
parallel semantics may be allocation hoisting barriers. The majority
267-
of ops, however, is not a barrier. Therefore, this method returns
268-
`false` by default.
269-
}],
270-
/*retType=*/"bool",
271-
/*methodName=*/"isAllocationHoistingBarrier",
272-
/*args=*/(ins),
273-
/*methodBody=*/"",
274-
/*defaultImplementation=*/[{
275-
return false;
276-
}]
277-
>,
278262
InterfaceMethod<
279263
/*desc=*/[{
280264
Return `true` if the `uRead` and `uWrite` do not constitute a RaW

‎mlir/include/mlir/Dialect/Bufferization/Transforms/Bufferize.h

Lines changed: 2 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -91,9 +91,8 @@ BufferizationOptions getPartialBufferizationOptions();
9191
LogicalResult bufferizeOp(Operation *op,
9292
BufferizationState &bufferizationState);
9393

94-
/// Finalize all buffer allocations.
95-
/// * Hoist buffer allocations as much as possible.
96-
/// * Create alloc/dealloc ops as specified by the bufferization options.
94+
/// Finalize all buffer allocations: Create alloc/dealloc ops as specified by
95+
/// the bufferization options.
9796
LogicalResult finalizeBuffers(Operation *op,
9897
const BufferizationOptions &options);
9998
} // namespace bufferization

‎mlir/lib/Dialect/Bufferization/IR/BufferizableOpInterface.cpp

Lines changed: 0 additions & 57 deletions
Original file line numberDiff line numberDiff line change
@@ -582,63 +582,6 @@ bufferization::createAllocDeallocOps(Operation *op,
582582
return success(!status.wasInterrupted());
583583
}
584584

585-
/// Try to hoist all new buffer allocations until the next hoisting barrier.
586-
// TODO: Consolidate this function with the existing buffer hoisting pass.
587-
LogicalResult
588-
bufferization::hoistBufferAllocations(Operation *op,
589-
const BufferizationOptions &options) {
590-
// Nothing to do if allocation hoisting is deactivated.
591-
if (!options.hoistAllocations)
592-
return success();
593-
594-
// Gather all buffer allocations that were created by the bufferization.
595-
SmallVector<Operation *> allocaOps;
596-
op->walk([&](memref::AllocaOp allocaOp) {
597-
if (allocaOp->hasAttr(kBufferAllocationAttr))
598-
allocaOps.push_back(allocaOp);
599-
});
600-
601-
for (Operation *allocaOp : allocaOps) {
602-
// TODO: Hoisting of allocs with dynamic shape not implemented.
603-
if (!allocaOp->getOpOperands().empty())
604-
continue;
605-
606-
Operation *op = allocaOp->getParentOp();
607-
while (op) {
608-
if (auto bufferizableOp = dyn_cast<BufferizableOpInterface>(op)) {
609-
if (bufferizableOp.isAllocationHoistingBarrier()) {
610-
break;
611-
}
612-
} else {
613-
// Op is not bufferizable: It may not be safe to hoist across this op.
614-
break;
615-
}
616-
op = op->getParentOp();
617-
}
618-
619-
// FuncOp is an allocation hoisting barrier, so this should never happen.
620-
assert(op && "allocation hoisting barrier not found");
621-
622-
// Nothing to do if the insertion point is in the same block.
623-
if (op == allocaOp->getParentOp())
624-
continue;
625-
626-
// `op` may have multiple blocks. Make sure that we insert in the right one.
627-
SmallVector<Block *> blocks;
628-
for (Region &r : op->getRegions())
629-
for (Block &b : r.getBlocks())
630-
blocks.push_back(&b);
631-
auto *insertionBlock = llvm::find_if(
632-
blocks, [&](Block *b) { return b->findAncestorOpInBlock(*allocaOp); });
633-
assert(insertionBlock != blocks.end() && "owning block not found");
634-
635-
// Move to the beginning of the block.
636-
allocaOp->moveBefore(&(*insertionBlock)->front());
637-
}
638-
639-
return success();
640-
}
641-
642585
//===----------------------------------------------------------------------===//
643586
// Bufferization-specific BlockAndValueMapping support with debugging.
644587
//===----------------------------------------------------------------------===//

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

Lines changed: 0 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -295,10 +295,6 @@ static bool hasTensorSemantics(Operation *op) {
295295
LogicalResult
296296
bufferization::finalizeBuffers(Operation *op,
297297
const BufferizationOptions &options) {
298-
// Hoist buffers.
299-
if (failed(hoistBufferAllocations(op, options)))
300-
return failure();
301-
302298
// Create allocation ops for "leaking buffers", i.e., buffer allocations that
303299
// escape block boundaries. If there are no leaking allocs, `hasLeakingAllocs`
304300
// is set to `false`.

‎mlir/lib/Dialect/Bufferization/Transforms/FuncBufferizableOpInterfaceImpl.cpp

Lines changed: 0 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -549,8 +549,6 @@ struct FuncOpInterface
549549
// All function arguments are writable by default.
550550
return true;
551551
}
552-
553-
bool isAllocationHoistingBarrier(Operation *op) const { return true; }
554552
};
555553

556554
} // namespace func_ext

‎mlir/lib/Dialect/Shape/Transforms/BufferizableOpInterfaceImpl.cpp

Lines changed: 0 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -119,11 +119,6 @@ struct AssumingOpInterface
119119
const AnalysisState &state) const {
120120
return BufferRelation::Equivalent;
121121
}
122-
123-
bool isAllocationHoistingBarrier(Operation *op) const {
124-
// Allocations should not be hoisted out of AssumingOps.
125-
return true;
126-
}
127122
};
128123

129124
/// Bufferization of shape.assuming_yield. Bufferized as part of their enclosing

‎mlir/test/Transforms/buffer-hoisting.mlir renamed to ‎mlir/test/Dialect/Bufferization/Transforms/buffer-hoisting.mlir

File renamed without changes.

‎mlir/test/Transforms/buffer-loop-hoisting.mlir renamed to ‎mlir/test/Dialect/Bufferization/Transforms/buffer-loop-hoisting.mlir

File renamed without changes.

‎mlir/test/Dialect/Linalg/one-shot-bufferize.mlir

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -1,4 +1,4 @@
1-
// RUN: mlir-opt %s -one-shot-bufferize="allow-return-allocs bufferize-function-boundaries" -split-input-file | FileCheck %s
1+
// RUN: mlir-opt %s -one-shot-bufferize="allow-return-allocs bufferize-function-boundaries" -buffer-loop-hoisting -split-input-file | FileCheck %s
22

33
// Run fuzzer with different seeds.
44
// RUN: mlir-opt %s -one-shot-bufferize="allow-return-allocs test-analysis-only analysis-fuzzer-seed=23 bufferize-function-boundaries" -split-input-file -o /dev/null

‎mlir/test/Dialect/SCF/one-shot-bufferize.mlir

Lines changed: 32 additions & 28 deletions
Original file line numberDiff line numberDiff line change
@@ -367,18 +367,20 @@ func.func @scf_while_non_equiv_condition(%arg0: tensor<5xi1>,
367367
%idx: index)
368368
-> (tensor<5xi1>, tensor<5xi1>)
369369
{
370-
// These allocation used to be inside the scf.while loop, but they were
371-
// hoisted.
372-
// CHECK: %[[a0:.*]] = memref.alloc() {{.*}} : memref<5xi1>
373-
// CHECK: %[[a1:.*]] = memref.alloc() {{.*}} : memref<5xi1>
374-
// CHECK: %[[loop:.*]]:2 = scf.while (%[[w0:.*]] = %[[arg0]], %[[w1:.*]] = %[[arg1]]) {{.*}} {
370+
// CHECK: %[[clone1:.*]] = bufferization.clone %[[arg1]]
371+
// CHECK: %[[clone0:.*]] = bufferization.clone %[[arg0]]
372+
// CHECK: %[[loop:.*]]:2 = scf.while (%[[w0:.*]] = %[[clone0]], %[[w1:.*]] = %[[clone1]]) {{.*}} {
375373
%r0, %r1 = scf.while (%w0 = %arg0, %w1 = %arg1)
376374
: (tensor<5xi1>, tensor<5xi1>) -> (tensor<5xi1>, tensor<5xi1>) {
377375
// CHECK: %[[condition:.*]] = memref.load %[[w0]]
376+
// CHECK: %[[a1:.*]] = memref.alloc() {{.*}} : memref<5xi1>
378377
// CHECK: memref.copy %[[w1]], %[[a1]]
379-
// CHECK: %[[casted1:.*]] = memref.cast %[[a1]]
378+
// CHECK: memref.dealloc %[[w1]]
379+
// CHECK: %[[a0:.*]] = memref.alloc() {{.*}} : memref<5xi1>
380380
// CHECK: memref.copy %[[w0]], %[[a0]]
381+
// CHECK: memref.dealloc %[[w0]]
381382
// CHECK: %[[casted0:.*]] = memref.cast %[[a0]]
383+
// CHECK: %[[casted1:.*]] = memref.cast %[[a1]]
382384
// CHECK: scf.condition(%[[condition]]) %[[casted1]], %[[casted0]]
383385
%condition = tensor.extract %w0[%idx] : tensor<5xi1>
384386
scf.condition(%condition) %w1, %w0 : tensor<5xi1>, tensor<5xi1>
@@ -410,42 +412,43 @@ func.func @scf_while_non_equiv_condition_and_body(%arg0: tensor<5xi1>,
410412
%idx: index)
411413
-> (tensor<5xi1>, tensor<5xi1>)
412414
{
413-
// These allocation used to be inside the scf.while loop, but they were
414-
// hoisted.
415-
// CHECK: %[[a0:.*]] = memref.alloc() {{.*}} : memref<5xi1>
416-
// CHECK: %[[a1:.*]] = memref.alloc() {{.*}} : memref<5xi1>
417-
// CHECK: %[[a2:.*]] = memref.alloc() {{.*}} : memref<5xi1>
418-
// CHECK: %[[a3:.*]] = memref.alloc() {{.*}} : memref<5xi1>
419-
// CHECK: %[[loop:.*]]:2 = scf.while (%[[w0:.*]] = %[[arg0]], %[[w1:.*]] = %[[arg1]]) {{.*}} {
415+
// CHECK: %[[clone1:.*]] = bufferization.clone %[[arg1]]
416+
// CHECK: %[[clone0:.*]] = bufferization.clone %[[arg0]]
417+
// CHECK: %[[loop:.*]]:2 = scf.while (%[[w0:.*]] = %[[clone0]], %[[w1:.*]] = %[[clone1]]) {{.*}} {
420418
%r0, %r1 = scf.while (%w0 = %arg0, %w1 = %arg1)
421419
: (tensor<5xi1>, tensor<5xi1>) -> (tensor<5xi1>, tensor<5xi1>) {
422420
// CHECK: %[[condition:.*]] = memref.load %[[w0]]
423-
// CHECK: memref.copy %[[w1]], %[[a3]]
424-
// CHECK: %[[casted3:.*]] = memref.cast %[[a3]]
425-
// CHECK: memref.copy %[[w0]], %[[a2]]
426-
// CHECK: %[[casted2:.*]] = memref.cast %[[a2]]
427-
// CHECK: scf.condition(%[[condition]]) %[[casted3]], %[[casted2]]
421+
// CHECK: %[[a1:.*]] = memref.alloc() {{.*}} : memref<5xi1>
422+
// CHECK: memref.copy %[[w1]], %[[a1]]
423+
// CHECK: memref.dealloc %[[w1]]
424+
// CHECK: %[[a0:.*]] = memref.alloc() {{.*}} : memref<5xi1>
425+
// CHECK: memref.copy %[[w0]], %[[a0]]
426+
// CHECK: memref.dealloc %[[w0]]
427+
// CHECK: %[[casted0:.*]] = memref.cast %[[a0]]
428+
// CHECK: %[[casted1:.*]] = memref.cast %[[a1]]
429+
// CHECK: scf.condition(%[[condition]]) %[[casted1]], %[[casted0]]
428430
%condition = tensor.extract %w0[%idx] : tensor<5xi1>
429431
scf.condition(%condition) %w1, %w0 : tensor<5xi1>, tensor<5xi1>
430432
} do {
431433
^bb0(%b0: tensor<5xi1>, %b1: tensor<5xi1>):
432434
// CHECK: } do {
433435
// CHECK: ^bb0(%[[b0:.*]]: memref<5xi1, #{{.*}}>, %[[b1:.*]]: memref<5xi1, #{{.*}}):
434436
// CHECK: memref.store %{{.*}}, %[[b0]]
435-
// CHECK: memref.copy %[[b1]], %[[a1]]
436-
// CHECK: %[[casted1:.*]] = memref.cast %[[a1]]
437-
// CHECK: memref.copy %[[b0]], %[[a0]]
438-
// CHECK: %[[casted0:.*]] = memref.cast %[[a0]]
439-
// CHECK: scf.yield %[[casted1]], %[[casted0]]
437+
// CHECK: %[[a3:.*]] = memref.alloc() {{.*}} : memref<5xi1>
438+
// CHECK: memref.copy %[[b1]], %[[a3]]
439+
// CHECK: memref.dealloc %[[b1]]
440+
// CHECK: %[[a2:.*]] = memref.alloc() {{.*}} : memref<5xi1>
441+
// CHECK: memref.copy %[[b0]], %[[a2]]
442+
// CHECK: %[[casted2:.*]] = memref.cast %[[a2]]
443+
// CHECK: %[[casted3:.*]] = memref.cast %[[a3]]
444+
// CHECK: scf.yield %[[casted3]], %[[casted2]]
440445
// CHECK: }
441446
%pos = "dummy.some_op"() : () -> (index)
442447
%val = "dummy.another_op"() : () -> (i1)
443448
%1 = tensor.insert %val into %b0[%pos] : tensor<5xi1>
444449
scf.yield %b1, %1 : tensor<5xi1>, tensor<5xi1>
445450
}
446451

447-
// CHECK-DAG: memref.dealloc %[[a0]]
448-
// CHECK-DAG: memref.dealloc %[[a1]]
449452
// CHECK: return %[[loop]]#0, %[[loop]]#1
450453
return %r0, %r1 : tensor<5xi1>, tensor<5xi1>
451454
}
@@ -454,19 +457,20 @@ func.func @scf_while_non_equiv_condition_and_body(%arg0: tensor<5xi1>,
454457

455458
// CHECK-LABEL: func @scf_while_iter_arg_result_mismatch(
456459
// CHECK-SAME: %[[arg0:.*]]: memref<5xi1, #{{.*}}>, %[[arg1:.*]]: memref<5xi1, #{{.*}}>
457-
// CHECK: %[[alloc1:.*]] = memref.alloc() {{.*}} : memref<5xi1>
458460
// CHECK: %[[alloc2:.*]] = memref.alloc() {{.*}} : memref<5xi1>
459-
// CHECK: scf.while (%[[arg3:.*]] = %[[arg1]]) : (memref<5xi1, #{{.*}}) -> () {
461+
// CHECK: %[[clone:.*]] = bufferization.clone %[[arg1]]
462+
// CHECK: scf.while (%[[arg3:.*]] = %[[clone]]) : (memref<5xi1, #{{.*}}) -> () {
463+
// CHECK: memref.dealloc %[[arg3]]
460464
// CHECK: %[[load:.*]] = memref.load %[[arg0]]
461465
// CHECK: scf.condition(%[[load]])
462466
// CHECK: } do {
463467
// CHECK: memref.copy %[[arg0]], %[[alloc2]]
464468
// CHECK: memref.store %{{.*}}, %[[alloc2]]
469+
// CHECK: %[[alloc1:.*]] = memref.alloc() {{.*}} : memref<5xi1>
465470
// CHECK: memref.copy %[[alloc2]], %[[alloc1]]
466471
// CHECK: %[[casted:.*]] = memref.cast %[[alloc1]] : memref<5xi1> to memref<5xi1, #{{.*}}>
467472
// CHECK: scf.yield %[[casted]]
468473
// CHECK: }
469-
// CHECK-DAG: memref.dealloc %[[alloc1]]
470474
// CHECK-DAG: memref.dealloc %[[alloc2]]
471475
func.func @scf_while_iter_arg_result_mismatch(%arg0: tensor<5xi1>,
472476
%arg1: tensor<5xi1>,

0 commit comments

Comments
 (0)