xref: /llvm-project/mlir/test/python/dialects/transform_bufferization_ext.py (revision f0549017533bd073173b7fba4c5295c7deb455e9) 
1# RUN: %PYTHON %s | FileCheck %s
2
3from mlir.ir import *
4from mlir.dialects import transform
5from mlir.dialects.transform import bufferization
6
7
8def run(f):
9    with Context(), Location.unknown():
10        module = Module.create()
11        with InsertionPoint(module.body):
12            print("\nTEST:", f.__name__)
13            f()
14        print(module)
15    return f
16
17
18@run
19def testEmptyTensorToAllocTensorOpCompact():
20    sequence = transform.SequenceOp(
21        transform.FailurePropagationMode.PROPAGATE,
22        [],
23        transform.OperationType.get("tensor.empty"),
24    )
25    with InsertionPoint(sequence.body):
26        bufferization.EmptyTensorToAllocTensorOp(sequence.bodyTarget)
27        transform.YieldOp()
28    # CHECK-LABEL: TEST: testEmptyTensorToAllocTensorOpCompact
29    # CHECK: = transform.bufferization.empty_tensor_to_alloc_tensor
30    # CHECK-SAME: (!transform.op<"tensor.empty">) -> !transform.op<"bufferization.alloc_tensor">
31
32
33@run
34def testEmptyTensorToAllocTensorOpTyped():
35    sequence = transform.SequenceOp(
36        transform.FailurePropagationMode.PROPAGATE,
37        [],
38        transform.OperationType.get("tensor.empty"),
39    )
40    with InsertionPoint(sequence.body):
41        bufferization.EmptyTensorToAllocTensorOp(
42            transform.OperationType.get("bufferization.alloc_tensor"),
43            sequence.bodyTarget,
44        )
45        transform.YieldOp()
46    # CHECK-LABEL: TEST: testEmptyTensorToAllocTensorOpTyped
47    # CHECK: = transform.bufferization.empty_tensor_to_alloc_tensor
48    # CHECK-SAME: (!transform.op<"tensor.empty">) -> !transform.op<"bufferization.alloc_tensor">
49
50
51@run
52def testOneShotBufferizeOpCompact():
53    sequence = transform.SequenceOp(
54        transform.FailurePropagationMode.PROPAGATE, [], transform.AnyOpType.get()
55    )
56    with InsertionPoint(sequence.body):
57        bufferization.OneShotBufferizeOp(sequence.bodyTarget)
58        transform.YieldOp()
59    # CHECK-LABEL: TEST: testOneShotBufferizeOpCompact
60    # CHECK: = transform.bufferization.one_shot_bufferize
61    # CHECK-SAME: (!transform.any_op) -> !transform.any_op
62
63
64@run
65def testOneShotBufferizeOpTyped():
66    sequence = transform.SequenceOp(
67        transform.FailurePropagationMode.PROPAGATE, [], transform.AnyOpType.get()
68    )
69    with InsertionPoint(sequence.body):
70        bufferization.OneShotBufferizeOp(
71            transform.OperationType.get("test.dummy"),
72            sequence.bodyTarget,
73        )
74        transform.YieldOp()
75    # CHECK-LABEL: TEST: testOneShotBufferizeOpTyped
76    # CHECK: = transform.bufferization.one_shot_bufferize
77    # CHECK-SAME: (!transform.any_op) -> !transform.op<"test.dummy">
78
79
80@run
81def testOneShotBufferizeOpAttributes():
82    sequence = transform.SequenceOp(
83        transform.FailurePropagationMode.PROPAGATE, [], transform.AnyOpType.get()
84    )
85    with InsertionPoint(sequence.body):
86        bufferization.OneShotBufferizeOp(
87            sequence.bodyTarget,
88            allow_return_allocs=True,
89            allow_unknown_ops=True,
90            bufferize_function_boundaries=True,
91            create_deallocs=False,
92            function_boundary_type_conversion=bufferization.LayoutMapOption.IDENTITY_LAYOUT_MAP,
93            memcpy_op="linalg.copy",
94            print_conflicts=True,
95            test_analysis_only=True,
96        )
97        transform.YieldOp()
98    # CHECK-LABEL: TEST: testOneShotBufferizeOpAttributes
99    # CHECK: = transform.bufferization.one_shot_bufferize
100    # CHECK-SAME: layout{IdentityLayoutMap}
101    # CHECK-SAME: allow_return_allocs = true
102    # CHECK-SAME: allow_unknown_ops = true
103    # CHECK-SAME: bufferize_function_boundaries = true
104    # CHECK-SAME: create_deallocs = false
105    # CHECK-SAME: memcpy_op = "linalg.copy"
106    # CHECK-SAME: print_conflicts = true
107    # CHECK-SAME: test_analysis_only = true
108    # CHECK-SAME: (!transform.any_op) -> !transform.any_op
109