1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
|
# 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_return_allocs_from_loops=True,
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_return_allocs_from_loops = true
# 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
|
