108 lines · python
1# RUN: %PYTHON %s | FileCheck %s2 3from mlir.ir import *4from mlir.dialects import transform5from mlir.dialects.transform import bufferization6from mlir.dialects.bufferization import LayoutMapOption7 8 9def run(f):10 with Context(), Location.unknown():11 module = Module.create()12 with InsertionPoint(module.body):13 print("\nTEST:", f.__name__)14 f()15 print(module)16 return f17 18 19@run20def testEmptyTensorToAllocTensorOpCompact():21 sequence = transform.SequenceOp(22 transform.FailurePropagationMode.Propagate,23 [],24 transform.OperationType.get("tensor.empty"),25 )26 with InsertionPoint(sequence.body):27 bufferization.EmptyTensorToAllocTensorOp(sequence.bodyTarget)28 transform.YieldOp()29 # CHECK-LABEL: TEST: testEmptyTensorToAllocTensorOpCompact30 # CHECK: = transform.bufferization.empty_tensor_to_alloc_tensor31 # CHECK-SAME: (!transform.op<"tensor.empty">) -> !transform.op<"bufferization.alloc_tensor">32 33 34@run35def testEmptyTensorToAllocTensorOpTyped():36 sequence = transform.SequenceOp(37 transform.FailurePropagationMode.Propagate,38 [],39 transform.OperationType.get("tensor.empty"),40 )41 with InsertionPoint(sequence.body):42 bufferization.EmptyTensorToAllocTensorOp(43 transform.OperationType.get("bufferization.alloc_tensor"),44 sequence.bodyTarget,45 )46 transform.YieldOp()47 # CHECK-LABEL: TEST: testEmptyTensorToAllocTensorOpTyped48 # CHECK: = transform.bufferization.empty_tensor_to_alloc_tensor49 # CHECK-SAME: (!transform.op<"tensor.empty">) -> !transform.op<"bufferization.alloc_tensor">50 51 52@run53def testOneShotBufferizeOpCompact():54 sequence = transform.SequenceOp(55 transform.FailurePropagationMode.Propagate, [], transform.AnyOpType.get()56 )57 with InsertionPoint(sequence.body):58 bufferization.OneShotBufferizeOp(sequence.bodyTarget)59 transform.YieldOp()60 # CHECK-LABEL: TEST: testOneShotBufferizeOpCompact61 # CHECK: = transform.bufferization.one_shot_bufferize62 # CHECK-SAME: (!transform.any_op) -> !transform.any_op63 64 65@run66def testOneShotBufferizeOpTyped():67 sequence = transform.SequenceOp(68 transform.FailurePropagationMode.Propagate, [], transform.AnyOpType.get()69 )70 with InsertionPoint(sequence.body):71 bufferization.OneShotBufferizeOp(72 transform.OperationType.get("test.dummy"),73 sequence.bodyTarget,74 )75 transform.YieldOp()76 # CHECK-LABEL: TEST: testOneShotBufferizeOpTyped77 # CHECK: = transform.bufferization.one_shot_bufferize78 # CHECK-SAME: (!transform.any_op) -> !transform.op<"test.dummy">79 80 81@run82def testOneShotBufferizeOpAttributes():83 sequence = transform.SequenceOp(84 transform.FailurePropagationMode.Propagate, [], transform.AnyOpType.get()85 )86 with InsertionPoint(sequence.body):87 bufferization.OneShotBufferizeOp(88 sequence.bodyTarget,89 allow_return_allocs_from_loops=True,90 allow_unknown_ops=True,91 bufferize_function_boundaries=True,92 function_boundary_type_conversion=LayoutMapOption.IdentityLayoutMap,93 memcpy_op="linalg.copy",94 print_conflicts=True,95 test_analysis_only=True,96 )97 transform.YieldOp()98 # CHECK-LABEL: TEST: testOneShotBufferizeOpAttributes99 # CHECK: = transform.bufferization.one_shot_bufferize100 # CHECK-SAME: layout{IdentityLayoutMap}101 # CHECK-SAME: allow_return_allocs_from_loops = true102 # CHECK-SAME: allow_unknown_ops = true103 # CHECK-SAME: bufferize_function_boundaries = true104 # CHECK-SAME: memcpy_op = "linalg.copy"105 # CHECK-SAME: print_conflicts = true106 # CHECK-SAME: test_analysis_only = true107 # CHECK-SAME: (!transform.any_op) -> !transform.any_op108