6a91dfedeb
This is the first commit in a series with the goal to rework the BufferDeallocation pass. Currently, this pass heavily relies on copies to perform correct deallocations, which leads to very slow code and potentially high memory usage. Additionally, there are unsupported cases such as returning memrefs which this series of commits aims to add support for as well. This first commit removes the deallocation capabilities of one-shot-bufferization.One-shot-bufferization should never deallocate any memrefs as this should be entirely handled by the buffer-deallocation pass going forward. This means the allow-return-allocs pass option will default to true now, create-deallocs defaults to false and they, as well as the escape attribute indicating whether a memref escapes the current region, will be removed. The documentation should w.r.t. these pass option changes should also be updated in this commit. Reviewed By: springerm Differential Revision: https://reviews.llvm.org/D156662
106 lines
3.6 KiB
Python
106 lines
3.6 KiB
Python
# RUN: %PYTHON %s | FileCheck %s
|
|
|
|
from mlir.ir import *
|
|
from mlir.dialects import transform
|
|
from mlir.dialects.transform import bufferization
|
|
from mlir.dialects.bufferization import LayoutMapOption
|
|
|
|
|
|
def run(f):
|
|
with Context(), Location.unknown():
|
|
module = Module.create()
|
|
with InsertionPoint(module.body):
|
|
print("\nTEST:", f.__name__)
|
|
f()
|
|
print(module)
|
|
return f
|
|
|
|
|
|
@run
|
|
def testEmptyTensorToAllocTensorOpCompact():
|
|
sequence = transform.SequenceOp(
|
|
transform.FailurePropagationMode.Propagate,
|
|
[],
|
|
transform.OperationType.get("tensor.empty"),
|
|
)
|
|
with InsertionPoint(sequence.body):
|
|
bufferization.EmptyTensorToAllocTensorOp(sequence.bodyTarget)
|
|
transform.YieldOp()
|
|
# CHECK-LABEL: TEST: testEmptyTensorToAllocTensorOpCompact
|
|
# CHECK: = transform.bufferization.empty_tensor_to_alloc_tensor
|
|
# CHECK-SAME: (!transform.op<"tensor.empty">) -> !transform.op<"bufferization.alloc_tensor">
|
|
|
|
|
|
@run
|
|
def testEmptyTensorToAllocTensorOpTyped():
|
|
sequence = transform.SequenceOp(
|
|
transform.FailurePropagationMode.Propagate,
|
|
[],
|
|
transform.OperationType.get("tensor.empty"),
|
|
)
|
|
with InsertionPoint(sequence.body):
|
|
bufferization.EmptyTensorToAllocTensorOp(
|
|
transform.OperationType.get("bufferization.alloc_tensor"),
|
|
sequence.bodyTarget,
|
|
)
|
|
transform.YieldOp()
|
|
# CHECK-LABEL: TEST: testEmptyTensorToAllocTensorOpTyped
|
|
# CHECK: = transform.bufferization.empty_tensor_to_alloc_tensor
|
|
# CHECK-SAME: (!transform.op<"tensor.empty">) -> !transform.op<"bufferization.alloc_tensor">
|
|
|
|
|
|
@run
|
|
def testOneShotBufferizeOpCompact():
|
|
sequence = transform.SequenceOp(
|
|
transform.FailurePropagationMode.Propagate, [], transform.AnyOpType.get()
|
|
)
|
|
with InsertionPoint(sequence.body):
|
|
bufferization.OneShotBufferizeOp(sequence.bodyTarget)
|
|
transform.YieldOp()
|
|
# CHECK-LABEL: TEST: testOneShotBufferizeOpCompact
|
|
# CHECK: = transform.bufferization.one_shot_bufferize
|
|
# CHECK-SAME: (!transform.any_op) -> !transform.any_op
|
|
|
|
|
|
@run
|
|
def testOneShotBufferizeOpTyped():
|
|
sequence = transform.SequenceOp(
|
|
transform.FailurePropagationMode.Propagate, [], transform.AnyOpType.get()
|
|
)
|
|
with InsertionPoint(sequence.body):
|
|
bufferization.OneShotBufferizeOp(
|
|
transform.OperationType.get("test.dummy"),
|
|
sequence.bodyTarget,
|
|
)
|
|
transform.YieldOp()
|
|
# CHECK-LABEL: TEST: testOneShotBufferizeOpTyped
|
|
# CHECK: = transform.bufferization.one_shot_bufferize
|
|
# CHECK-SAME: (!transform.any_op) -> !transform.op<"test.dummy">
|
|
|
|
|
|
@run
|
|
def testOneShotBufferizeOpAttributes():
|
|
sequence = transform.SequenceOp(
|
|
transform.FailurePropagationMode.Propagate, [], transform.AnyOpType.get()
|
|
)
|
|
with InsertionPoint(sequence.body):
|
|
bufferization.OneShotBufferizeOp(
|
|
sequence.bodyTarget,
|
|
allow_unknown_ops=True,
|
|
bufferize_function_boundaries=True,
|
|
function_boundary_type_conversion=LayoutMapOption.IdentityLayoutMap,
|
|
memcpy_op="linalg.copy",
|
|
print_conflicts=True,
|
|
test_analysis_only=True,
|
|
)
|
|
transform.YieldOp()
|
|
# CHECK-LABEL: TEST: testOneShotBufferizeOpAttributes
|
|
# CHECK: = transform.bufferization.one_shot_bufferize
|
|
# CHECK-SAME: layout{IdentityLayoutMap}
|
|
# CHECK-SAME: allow_unknown_ops = true
|
|
# CHECK-SAME: bufferize_function_boundaries = true
|
|
# CHECK-SAME: memcpy_op = "linalg.copy"
|
|
# CHECK-SAME: print_conflicts = true
|
|
# CHECK-SAME: test_analysis_only = true
|
|
# CHECK-SAME: (!transform.any_op) -> !transform.any_op
|