297 lines · python
1# RUN: %PYTHON %s | FileCheck %s2 3from mlir.ir import *4from mlir.dialects import transform5from mlir.dialects.transform import xegpu6from mlir.dialects.transform import structured, AnyValueType7 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 getDescOpDefaultIndex():21 sequence = transform.SequenceOp(22 transform.FailurePropagationMode.Propagate,23 [],24 transform.OperationType.get("xegpu.dpas"),25 )26 with InsertionPoint(sequence.body):27 operand = transform.GetOperandOp(AnyValueType.get(), sequence.bodyTarget, [0])28 desc_handle = xegpu.get_desc_op(operand)29 transform.YieldOp()30 # CHECK-LABEL: TEST: getDescOpDefaultIndex31 # CHECK: transform.xegpu.get_desc_op %32 33 34@run35def setDescLayoutMinimal():36 sequence = transform.SequenceOp(37 transform.FailurePropagationMode.Propagate,38 [],39 transform.OperationType.get("xegpu.create_nd_tdesc"),40 )41 with InsertionPoint(sequence.body):42 xegpu.set_desc_layout(sequence.bodyTarget, sg_layout=[6, 4], sg_data=[32, 16])43 transform.YieldOp()44 # CHECK-LABEL: TEST: setDescLayoutMinimal45 # CHECK: %0 = transform.xegpu.set_desc_layout %46 # CHECK: sg_layout = [6, 4]47 # CHECK: sg_data = [32, 16]48 49 50@run51def setDescLayoutInstData():52 sequence = transform.SequenceOp(53 transform.FailurePropagationMode.Propagate,54 [],55 transform.OperationType.get("xegpu.create_nd_tdesc"),56 )57 with InsertionPoint(sequence.body):58 xegpu.set_desc_layout(59 sequence.bodyTarget, sg_layout=[6, 4], sg_data=[32, 16], inst_data=[8, 16]60 )61 transform.YieldOp()62 # CHECK-LABEL: TEST: setDescLayoutInstData63 # CHECK: %0 = transform.xegpu.set_desc_layout %64 # CHECK: sg_layout = [6, 4]65 # CHECK: sg_data = [32, 16]66 # CHECK: inst_data = [8, 16]67 68 69@run70def setDescLayoutSlice():71 sequence = transform.SequenceOp(72 transform.FailurePropagationMode.Propagate,73 [],74 transform.OperationType.get("xegpu.create_nd_tdesc"),75 )76 with InsertionPoint(sequence.body):77 xegpu.set_desc_layout(78 sequence.bodyTarget, sg_layout=[6, 4], sg_data=[32, 16], slice_dims=[0]79 )80 transform.YieldOp()81 # CHECK-LABEL: TEST: setDescLayoutSlice82 # CHECK: %0 = transform.xegpu.set_desc_layout %83 # CHECK: sg_layout = [6, 4]84 # CHECK: sg_data = [32, 16]85 # CHECK: slice_dims = [0]86 87 88@run89def setOpLayoutAttrOperandMinimal():90 sequence = transform.SequenceOp(91 transform.FailurePropagationMode.Propagate,92 [],93 transform.OperationType.get("xegpu.dpas"),94 )95 with InsertionPoint(sequence.body):96 xegpu.set_op_layout_attr(97 sequence.bodyTarget,98 sg_layout=[6, 4],99 sg_data=[32, 16],100 )101 transform.YieldOp()102 # CHECK-LABEL: TEST: setOpLayoutAttr103 # CHECK: transform.xegpu.set_op_layout_attr %104 # NO-CHECK: index = 0105 # NO-CHECK: result106 # CHECK: sg_layout = [6, 4]107 # CHECK: sg_data = [32, 16]108 # NO-CHECK: inst_data109 110 111@run112def setOpLayoutAttrResult():113 sequence = transform.SequenceOp(114 transform.FailurePropagationMode.Propagate,115 [],116 transform.OperationType.get("xegpu.dpas"),117 )118 with InsertionPoint(sequence.body):119 xegpu.set_op_layout_attr(120 sequence.bodyTarget,121 index=0,122 sg_layout=[6, 4],123 sg_data=[32, 16],124 inst_data=[8, 16],125 result=True,126 )127 transform.YieldOp()128 # CHECK-LABEL: TEST: setOpLayoutAttrResult129 # CHECK: transform.xegpu.set_op_layout_attr %130 # NO-CHECK: index = 0131 # CHECK: result132 # CHECK: sg_layout = [6, 4]133 # CHECK: sg_data = [32, 16]134 # CHECK: inst_data = [8, 16]135 136 137@run138def setOpLayoutAttrResultSlice():139 sequence = transform.SequenceOp(140 transform.FailurePropagationMode.Propagate,141 [],142 transform.OperationType.get("xegpu.dpas"),143 )144 with InsertionPoint(sequence.body):145 xegpu.set_op_layout_attr(146 sequence.bodyTarget,147 index=0,148 sg_layout=[6, 4],149 sg_data=[32, 16],150 inst_data=[8, 16],151 slice_dims=[0],152 result=True,153 )154 transform.YieldOp()155 # CHECK-LABEL: TEST: setOpLayoutAttrResultSlice156 # CHECK: transform.xegpu.set_op_layout_attr %157 # NO-CHECK: index = 0158 # CHECK: result159 # CHECK: sg_layout = [6, 4]160 # CHECK: sg_data = [32, 16]161 # CHECK: inst_data = [8, 16]162 # CHECK: slice_dims = [0]163 164 165@run166def setGPULaunchThreadsOp():167 sequence = transform.SequenceOp(168 transform.FailurePropagationMode.Propagate,169 [],170 transform.OperationType.get("gpu.launch"),171 )172 with InsertionPoint(sequence.body):173 xegpu.set_gpu_launch_threads(sequence.bodyTarget, threads=[8, 4, 1])174 transform.YieldOp()175 # CHECK-LABEL: TEST: setGPULaunchThreadsOp176 # CHECK: transform.xegpu.set_gpu_launch_threads177 # CHECK: threads = [8, 4, 1]178 179 180@run181def insertPrefetch0():182 sequence = transform.SequenceOp(183 transform.FailurePropagationMode.Propagate,184 [],185 transform.OperationType.get("xegpu.dpas"),186 )187 with InsertionPoint(sequence.body):188 operand = transform.GetOperandOp(AnyValueType.get(), sequence.bodyTarget, [0])189 xegpu.insert_prefetch(190 operand,191 )192 transform.YieldOp()193 # CHECK-LABEL: TEST: insertPrefetch0194 # CHECK: %[[OPR:.*]] = get_operand195 # CHECK: transform.xegpu.insert_prefetch %[[OPR]]196 197 198@run199def insertPrefetchNbPrefetch():200 sequence = transform.SequenceOp(201 transform.FailurePropagationMode.Propagate,202 [],203 transform.OperationType.get("xegpu.dpas"),204 )205 with InsertionPoint(sequence.body):206 operand = transform.GetOperandOp(AnyValueType.get(), sequence.bodyTarget, [0])207 xegpu.insert_prefetch(208 operand,209 nb_prefetch=2,210 )211 transform.YieldOp()212 # CHECK-LABEL: TEST: insertPrefetchNbPrefetch213 # CHECK: %[[OPR:.*]] = get_operand214 # CHECK: transform.xegpu.insert_prefetch %[[OPR]]215 # CHECK-SAME: nb_prefetch = 2216 217 218@run219def insertPrefetchNbPrefetchParam():220 sequence = transform.SequenceOp(221 transform.FailurePropagationMode.Propagate,222 [],223 transform.OperationType.get("xegpu.dpas"),224 )225 with InsertionPoint(sequence.body):226 operand = transform.GetOperandOp(AnyValueType.get(), sequence.bodyTarget, [0])227 int32_t = IntegerType.get_signless(32)228 param_int32_t = transform.ParamType.get(int32_t)229 nb_param = transform.ParamConstantOp(230 param_int32_t,231 IntegerAttr.get(int32_t, 2),232 )233 xegpu.insert_prefetch(234 operand,235 nb_prefetch=nb_param,236 )237 transform.YieldOp()238 # CHECK-LABEL: TEST: insertPrefetchNbPrefetchParam239 # CHECK: %[[OPR:.*]] = get_operand240 # CHECK: %[[PARAM_OP:.*]] = transform.param.constant 2241 # CHECK: transform.xegpu.insert_prefetch %[[OPR]]242 # CHECK-SAME: nb_prefetch = %[[PARAM_OP]]243 244 245@run246def ConvertLayoutMinimal():247 sequence = transform.SequenceOp(248 transform.FailurePropagationMode.Propagate,249 [],250 transform.OperationType.get("xegpu.dpas"),251 )252 with InsertionPoint(sequence.body):253 operand = transform.GetOperandOp(AnyValueType.get(), sequence.bodyTarget, [0])254 xegpu.convert_layout(255 operand,256 input_sg_layout=[6, 4],257 input_sg_data=[32, 16],258 target_sg_layout=[6, 4],259 target_sg_data=[8, 16],260 )261 transform.YieldOp()262 # CHECK-LABEL: TEST: ConvertLayoutMinimal263 # CHECK: transform.xegpu.convert_layout %264 # CHECK: input_sg_layout = [6, 4]265 # CHECK: input_sg_data = [32, 16]266 # CHECK: target_sg_layout = [6, 4]267 # CHECK: target_sg_data = [8, 16]268 269 270@run271def ConvertLayout():272 sequence = transform.SequenceOp(273 transform.FailurePropagationMode.Propagate,274 [],275 transform.OperationType.get("xegpu.dpas"),276 )277 with InsertionPoint(sequence.body):278 operand = transform.GetOperandOp(AnyValueType.get(), sequence.bodyTarget, [1])279 xegpu.convert_layout(280 operand,281 input_sg_layout=[6, 4],282 input_sg_data=[32, 32],283 input_inst_data=[32, 16],284 target_sg_layout=[6, 4],285 target_sg_data=[32, 32],286 target_inst_data=[8, 16],287 )288 transform.YieldOp()289 # CHECK-LABEL: TEST: ConvertLayout290 # CHECK: transform.xegpu.convert_layout %291 # CHECK: input_sg_layout = [6, 4]292 # CHECK: input_sg_data = [32, 32]293 # CHECK: input_inst_data = [32, 16]294 # CHECK: target_sg_layout = [6, 4]295 # CHECK: target_sg_data = [32, 32]296 # CHECK: target_inst_data = [8, 16]297