103 lines · python
1# RUN: %PYTHON %s | FileCheck %s2 3 4from mlir.ir import *5from mlir.dialects import transform6from mlir.dialects.transform import memref7 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 testMemRefAllocaToAllocOpCompact():21 sequence = transform.SequenceOp(22 transform.FailurePropagationMode.Propagate,23 [],24 transform.OperationType.get("memref.alloca"),25 )26 with InsertionPoint(sequence.body):27 memref.MemRefAllocaToGlobalOp(sequence.bodyTarget)28 transform.YieldOp()29 # CHECK-LABEL: TEST: testMemRefAllocaToAllocOpCompact30 # CHECK: = transform.memref.alloca_to_global31 # CHECK-SAME: (!transform.op<"memref.alloca">)32 # CHECK-SAME: -> (!transform.any_op, !transform.any_op)33 34 35@run36def testMemRefAllocaToAllocOpTyped():37 sequence = transform.SequenceOp(38 transform.FailurePropagationMode.Propagate,39 [],40 transform.OperationType.get("memref.alloca"),41 )42 with InsertionPoint(sequence.body):43 memref.MemRefAllocaToGlobalOp(44 transform.OperationType.get("memref.get_global"),45 transform.OperationType.get("memref.global"),46 sequence.bodyTarget,47 )48 transform.YieldOp()49 # CHECK-LABEL: TEST: testMemRefAllocaToAllocOpTyped50 # CHECK: = transform.memref.alloca_to_global51 # CHECK-SAME: -> (!transform.op<"memref.get_global">, !transform.op<"memref.global">)52 53 54@run55def testMemRefMultiBufferOpCompact():56 sequence = transform.SequenceOp(57 transform.FailurePropagationMode.Propagate,58 [],59 transform.OperationType.get("memref.alloc"),60 )61 with InsertionPoint(sequence.body):62 memref.MemRefMultiBufferOp(sequence.bodyTarget, 4)63 transform.YieldOp()64 # CHECK-LABEL: TEST: testMemRefMultiBufferOpCompact65 # CHECK: = transform.memref.multibuffer66 # CHECK-SAME: factor = 4 : i6467 # CHECK-SAME: (!transform.op<"memref.alloc">) -> !transform.any_op68 69 70@run71def testMemRefMultiBufferOpTyped():72 sequence = transform.SequenceOp(73 transform.FailurePropagationMode.Propagate,74 [],75 transform.OperationType.get("memref.alloc"),76 )77 with InsertionPoint(sequence.body):78 memref.MemRefMultiBufferOp(79 transform.OperationType.get("memref.alloc"), sequence.bodyTarget, 480 )81 transform.YieldOp()82 # CHECK-LABEL: TEST: testMemRefMultiBufferOpTyped83 # CHECK: = transform.memref.multibuffer84 # CHECK-SAME: factor = 4 : i6485 # CHECK-SAME: (!transform.op<"memref.alloc">) -> !transform.op<"memref.alloc">86 87 88@run89def testMemRefMultiBufferOpAttributes():90 sequence = transform.SequenceOp(91 transform.FailurePropagationMode.Propagate,92 [],93 transform.OperationType.get("memref.alloc"),94 )95 with InsertionPoint(sequence.body):96 memref.MemRefMultiBufferOp(sequence.bodyTarget, 4, skip_analysis=True)97 transform.YieldOp()98 # CHECK-LABEL: TEST: testMemRefMultiBufferOpAttributes99 # CHECK: = transform.memref.multibuffer100 # CHECK-SAME: factor = 4 : i64101 # CHECK-SAME: skip_analysis102 # CHECK-SAME: (!transform.op<"memref.alloc">) -> !transform.any_op103