152 lines · python
1# RUN: %PYTHON %s | FileCheck %s2 3from mlir.ir import *4from mlir.dialects import transform5from mlir.dialects.transform import vector6 7 8def run_apply_patterns(f):9 with Context(), Location.unknown():10 module = Module.create()11 with InsertionPoint(module.body):12 sequence = transform.SequenceOp(13 transform.FailurePropagationMode.Propagate,14 [],15 transform.AnyOpType.get(),16 )17 with InsertionPoint(sequence.body):18 apply = transform.ApplyPatternsOp(sequence.bodyTarget)19 with InsertionPoint(apply.patterns):20 f()21 transform.YieldOp()22 print("\nTEST:", f.__name__)23 print(module)24 return f25 26 27@run_apply_patterns28def non_configurable_patterns():29 # CHECK-LABEL: TEST: non_configurable_patterns30 # CHECK: apply_patterns31 # CHECK: transform.apply_patterns.vector.cast_away_vector_leading_one_dim32 vector.ApplyCastAwayVectorLeadingOneDimPatternsOp()33 # CHECK: transform.apply_patterns.vector.rank_reducing_subview_patterns34 vector.ApplyRankReducingSubviewPatternsOp()35 # CHECK: transform.apply_patterns.vector.transfer_permutation_patterns36 vector.ApplyTransferPermutationPatternsOp()37 # CHECK: transform.apply_patterns.vector.lower_broadcast38 vector.ApplyLowerBroadcastPatternsOp()39 # CHECK: transform.apply_patterns.vector.lower_masks40 vector.ApplyLowerMasksPatternsOp()41 # CHECK: transform.apply_patterns.vector.lower_masked_transfers42 vector.ApplyLowerMaskedTransfersPatternsOp()43 # CHECK: transform.apply_patterns.vector.materialize_masks44 vector.ApplyMaterializeMasksPatternsOp()45 # CHECK: transform.apply_patterns.vector.lower_outerproduct46 vector.ApplyLowerOuterProductPatternsOp()47 # CHECK: transform.apply_patterns.vector.lower_gather48 vector.ApplyLowerGatherPatternsOp()49 # CHECK: transform.apply_patterns.vector.unroll_from_elements50 vector.ApplyUnrollFromElementsPatternsOp()51 # CHECK: transform.apply_patterns.vector.unroll_to_elements52 vector.ApplyUnrollToElementsPatternsOp()53 # CHECK: transform.apply_patterns.vector.lower_scan54 vector.ApplyLowerScanPatternsOp()55 # CHECK: transform.apply_patterns.vector.lower_shape_cast56 vector.ApplyLowerShapeCastPatternsOp()57 58 59@run_apply_patterns60def configurable_patterns():61 # CHECK-LABEL: TEST: configurable_patterns62 # CHECK: apply_patterns63 # CHECK: transform.apply_patterns.vector.lower_transfer64 # CHECK-SAME: max_transfer_rank = 465 vector.ApplyLowerTransferPatternsOp(max_transfer_rank=4)66 # CHECK: transform.apply_patterns.vector.transfer_to_scf67 # CHECK-SAME: max_transfer_rank = 368 # CHECK-SAME: full_unroll = true69 vector.ApplyTransferToScfPatternsOp(max_transfer_rank=3, full_unroll=True)70 71 72@run_apply_patterns73def enum_configurable_patterns():74 # CHECK: transform.apply_patterns.vector.lower_contraction75 vector.ApplyLowerContractionPatternsOp()76 # CHECK: transform.apply_patterns.vector.lower_contraction77 # CHECK-SAME: lowering_strategy = llvmintr78 vector.ApplyLowerContractionPatternsOp(79 lowering_strategy=vector.VectorContractLowering.LLVMIntr80 )81 # CHECK: transform.apply_patterns.vector.lower_contraction82 # CHECK-SAME: lowering_strategy = parallelarith83 vector.ApplyLowerContractionPatternsOp(84 lowering_strategy=vector.VectorContractLowering.ParallelArith85 )86 87 # CHECK: transform.apply_patterns.vector.lower_multi_reduction88 vector.ApplyLowerMultiReductionPatternsOp()89 # CHECK: transform.apply_patterns.vector.lower_multi_reduction90 # This is the default mode, not printed.91 vector.ApplyLowerMultiReductionPatternsOp(92 lowering_strategy=vector.VectorMultiReductionLowering.InnerParallel93 )94 # CHECK: transform.apply_patterns.vector.lower_multi_reduction95 # CHECK-SAME: lowering_strategy = innerreduction96 vector.ApplyLowerMultiReductionPatternsOp(97 lowering_strategy=vector.VectorMultiReductionLowering.InnerReduction98 )99 100 # CHECK: transform.apply_patterns.vector.lower_transpose101 vector.ApplyLowerTransposePatternsOp()102 # CHECK: transform.apply_patterns.vector.lower_transpose103 # This is the default strategy, not printed.104 vector.ApplyLowerTransposePatternsOp(105 lowering_strategy=vector.VectorTransposeLowering.EltWise106 )107 # CHECK: transform.apply_patterns.vector.lower_transpose108 # CHECK-SAME: lowering_strategy = llvmintr109 vector.ApplyLowerTransposePatternsOp(110 lowering_strategy=vector.VectorTransposeLowering.LLVMIntr111 )112 # CHECK: transform.apply_patterns.vector.lower_transpose113 # CHECK-SAME: lowering_strategy = shuffle_1d114 vector.ApplyLowerTransposePatternsOp(115 lowering_strategy=vector.VectorTransposeLowering.Shuffle1D116 )117 # CHECK: transform.apply_patterns.vector.lower_transpose118 # CHECK-SAME: lowering_strategy = shuffle_16x16119 vector.ApplyLowerTransposePatternsOp(120 lowering_strategy=vector.VectorTransposeLowering.Shuffle16x16121 )122 # CHECK: transform.apply_patterns.vector.lower_transpose123 # CHECK-SAME: lowering_strategy = llvmintr124 # CHECK-SAME: avx2_lowering_strategy = true125 vector.ApplyLowerTransposePatternsOp(126 lowering_strategy=vector.VectorTransposeLowering.LLVMIntr,127 avx2_lowering_strategy=True,128 )129 130 # CHECK: transform.apply_patterns.vector.split_transfer_full_partial131 vector.ApplySplitTransferFullPartialPatternsOp()132 # CHECK: transform.apply_patterns.vector.split_transfer_full_partial133 # CHECK-SAME: split_transfer_strategy = none134 vector.ApplySplitTransferFullPartialPatternsOp(135 split_transfer_strategy=vector.VectorTransferSplit.None_136 )137 # CHECK: transform.apply_patterns.vector.split_transfer_full_partial138 # CHECK-SAME: split_transfer_strategy = "vector-transfer"139 vector.ApplySplitTransferFullPartialPatternsOp(140 split_transfer_strategy=vector.VectorTransferSplit.VectorTransfer141 )142 # CHECK: transform.apply_patterns.vector.split_transfer_full_partial143 # This is the default mode, not printed.144 vector.ApplySplitTransferFullPartialPatternsOp(145 split_transfer_strategy=vector.VectorTransferSplit.LinalgCopy146 )147 # CHECK: transform.apply_patterns.vector.split_transfer_full_partial148 # CHECK-SAME: split_transfer_strategy = "force-in-bounds"149 vector.ApplySplitTransferFullPartialPatternsOp(150 split_transfer_strategy=vector.VectorTransferSplit.ForceInBounds151 )152