510 lines · plain
1// RUN: mlir-opt %s -transform-interpreter -split-input-file -verify-diagnostics | FileCheck %s2 3// CHECK-LABEL: @get_desc_op_a4func.func @get_desc_op_a(%arg0: memref<4096x4096xf16>, %arg1: memref<4096x4096xf16>, %arg2: memref<4096x4096xf16>) {5 %c32 = arith.constant 32 : index6 %c4096 = arith.constant 4096 : index7 %c0 = arith.constant 0 : index8 %0 = xegpu.create_nd_tdesc %arg2 : memref<4096x4096xf16> -> !xegpu.tensor_desc<256x256xf16>9 %1 = xegpu.load_nd %0[%c0, %c0] : !xegpu.tensor_desc<256x256xf16> -> vector<256x256xf16>10 // expected-remark @below {{found desc op}}11 %3 = xegpu.create_nd_tdesc %arg0 : memref<4096x4096xf16> -> !xegpu.tensor_desc<256x32xf16>12 %4 = xegpu.create_nd_tdesc %arg1 : memref<4096x4096xf16> -> !xegpu.tensor_desc<32x256xf16>13 %2 = scf.for %arg3 = %c0 to %c4096 step %c32 iter_args(%arg4 = %1) -> (vector<256x256xf16>) {14 %5 = xegpu.load_nd %3[%c0, %arg3] : !xegpu.tensor_desc<256x32xf16> -> vector<256x32xf16>15 %6 = xegpu.load_nd %4[%arg3, %c0] : !xegpu.tensor_desc<32x256xf16> -> vector<32x256xf16>16 %7 = xegpu.dpas %5, %6, %arg4 : vector<256x32xf16>, vector<32x256xf16>, vector<256x256xf16> -> vector<256x256xf16>17 scf.yield %7 : vector<256x256xf16>18 }19 return20}21 22module attributes {transform.with_named_sequence} {23 transform.named_sequence @__transform_main(%arg1: !transform.any_op {transform.readonly}) {24 %0 = transform.structured.match ops{["xegpu.dpas"]} in %arg1 : (!transform.any_op) -> !transform.any_op25 %1 = transform.get_operand %0[0] : (!transform.any_op) -> !transform.any_value26 %2 = transform.xegpu.get_desc_op %1 : (!transform.any_value) -> !transform.any_op27 transform.debug.emit_remark_at %2, "found desc op" : !transform.any_op28 transform.yield29 }30}31 32// -----33 34// CHECK-LABEL: @get_desc_op_c35func.func @get_desc_op_c(%arg0: memref<4096x4096xf16>, %arg1: memref<4096x4096xf16>, %arg2: memref<4096x4096xf16>) {36 %c32 = arith.constant 32 : index37 %c4096 = arith.constant 4096 : index38 %c0 = arith.constant 0 : index39 // expected-remark @below {{found desc op}}40 %0 = xegpu.create_nd_tdesc %arg2 : memref<4096x4096xf16> -> !xegpu.tensor_desc<256x256xf16>41 %1 = xegpu.load_nd %0[%c0, %c0] : !xegpu.tensor_desc<256x256xf16> -> vector<256x256xf16>42 %3 = xegpu.create_nd_tdesc %arg0 : memref<4096x4096xf16> -> !xegpu.tensor_desc<256x32xf16>43 %4 = xegpu.create_nd_tdesc %arg1 : memref<4096x4096xf16> -> !xegpu.tensor_desc<32x256xf16>44 %2 = scf.for %arg3 = %c0 to %c4096 step %c32 iter_args(%arg4 = %1) -> (vector<256x256xf16>) {45 %5 = xegpu.load_nd %3[%c0, %arg3] : !xegpu.tensor_desc<256x32xf16> -> vector<256x32xf16>46 %6 = xegpu.load_nd %4[%arg3, %c0] : !xegpu.tensor_desc<32x256xf16> -> vector<32x256xf16>47 %7 = xegpu.dpas %5, %6, %arg4 : vector<256x32xf16>, vector<32x256xf16>, vector<256x256xf16> -> vector<256x256xf16>48 scf.yield %7 : vector<256x256xf16>49 }50 return51}52 53module attributes {transform.with_named_sequence} {54 transform.named_sequence @__transform_main(%arg1: !transform.any_op {transform.readonly}) {55 %0 = transform.structured.match ops{["xegpu.dpas"]} in %arg1 : (!transform.any_op) -> !transform.any_op56 %1 = transform.get_operand %0[2] : (!transform.any_op) -> !transform.any_value57 %2 = transform.xegpu.get_desc_op %1 : (!transform.any_value) -> !transform.any_op58 transform.debug.emit_remark_at %2, "found desc op" : !transform.any_op59 transform.yield60 }61}62 63// -----64 65// CHECK-LABEL: @set_desc_layout66func.func @set_desc_layout(%arg0: memref<4096x4096xf16>) {67 // CHECK: %[[V0:.+]] = xegpu.create_nd_tdesc %arg068 // CHECK-SAME: #xegpu.block_tdesc_attr<boundary_check = false>69 // CHECK-SAME: #xegpu.layout<sg_layout = [8, 4], sg_data = [32, 32], inst_data = [8, 16]>>70 %0 = xegpu.create_nd_tdesc %arg0 : memref<4096x4096xf16> -> !xegpu.tensor_desc<256x32xf16, #xegpu.block_tdesc_attr<boundary_check = false>>71 return72}73 74module attributes {transform.with_named_sequence} {75 transform.named_sequence @__transform_main(%arg1: !transform.any_op {transform.readonly}) {76 %0 = transform.structured.match ops{["xegpu.create_nd_tdesc"]} in %arg1 : (!transform.any_op) -> !transform.any_op77 // CHECK: transform.xegpu.set_desc_layout %{{.*}}78 %1 = transform.xegpu.set_desc_layout %0 sg_layout = [8, 4] sg_data = [32, 32] inst_data = [8, 16] : (!transform.any_op) -> !transform.any_op79 transform.yield80 }81}82 83// -----84 85// CHECK-LABEL: @set_desc_layout_minimal86func.func @set_desc_layout_minimal(%arg0: memref<4096x4096xf16>) {87 // CHECK: %[[V0:.+]] = xegpu.create_nd_tdesc %arg088 // CHECK-SAME: #xegpu.layout<sg_layout = [8, 4], sg_data = [32, 32]>>89 %0 = xegpu.create_nd_tdesc %arg0 : memref<4096x4096xf16> -> !xegpu.tensor_desc<256x32xf16>90 return91}92 93module attributes {transform.with_named_sequence} {94 transform.named_sequence @__transform_main(%arg1: !transform.any_op {transform.readonly}) {95 %0 = transform.structured.match ops{["xegpu.create_nd_tdesc"]} in %arg1 : (!transform.any_op) -> !transform.any_op96 // CHECK: transform.xegpu.set_desc_layout %{{.*}}97 %1 = transform.xegpu.set_desc_layout %0 sg_layout = [8, 4] sg_data = [32, 32] : (!transform.any_op) -> !transform.any_op98 transform.yield99 }100}101 102// -----103 104// CHECK-LABEL: @set_desc_layout_param105func.func @set_desc_layout_param(%arg0: memref<4096x4096xf16>) {106 // CHECK: %[[V0:.+]] = xegpu.create_nd_tdesc %arg0107 // CHECK-SAME: #xegpu.layout<sg_layout = [8, 4], sg_data = [32, 32], inst_data = [8, 16]>>108 %0 = xegpu.create_nd_tdesc %arg0 : memref<4096x4096xf16> -> !xegpu.tensor_desc<256x32xf16>109 return110}111 112module attributes {transform.with_named_sequence} {113 transform.named_sequence @__transform_main(%arg1: !transform.any_op {transform.readonly}) {114 %0 = transform.structured.match ops{["xegpu.create_nd_tdesc"]} in %arg1 : (!transform.any_op) -> !transform.any_op115 // CHECK: transform.xegpu.set_desc_layout %{{.*}}116 %layout0 = transform.param.constant 8 : i64 -> !transform.param<i64>117 %1 = transform.xegpu.set_desc_layout %0 sg_layout = [%layout0, 4] sg_data = [32, 32] inst_data = [8, 16] : (!transform.any_op, !transform.param<i64>) -> !transform.any_op118 transform.yield119 }120}121 122// -----123 124// CHECK-LABEL: @set_desc_layout_slice125func.func @set_desc_layout_slice(%arg0: memref<4096xf16>) {126 // CHECK: %[[V0:.+]] = xegpu.create_nd_tdesc %arg0127 // CHECK-SAME: #xegpu.slice<#xegpu.layout<sg_layout = [8, 4], sg_data = [32, 32]>, dims = [0]>128 %0 = xegpu.create_nd_tdesc %arg0 : memref<4096xf16> -> !xegpu.tensor_desc<256xf16>129 return130}131 132module attributes {transform.with_named_sequence} {133 transform.named_sequence @__transform_main(%arg1: !transform.any_op {transform.readonly}) {134 %0 = transform.structured.match ops{["xegpu.create_nd_tdesc"]} in %arg1 : (!transform.any_op) -> !transform.any_op135 // CHECK: transform.xegpu.set_desc_layout %{{.*}}136 %1 = transform.xegpu.set_desc_layout %0 sg_layout = [8, 4] sg_data = [32, 32] slice_dims = [0] : (!transform.any_op) -> !transform.any_op137 transform.yield138 }139}140 141// -----142 143// CHECK-LABEL: @set_op_layout_attr_result_default_index144func.func @set_op_layout_attr_result_default_index(%arg0: memref<4096x4096xf16>, %arg1: memref<4096x4096xf16>, %arg2: memref<4096x4096xf16>) {145 %0 = xegpu.create_nd_tdesc %arg0 : memref<4096x4096xf16> -> !xegpu.tensor_desc<256x32xf16>146 %1 = xegpu.load_nd %0[0, 0] : !xegpu.tensor_desc<256x32xf16> -> vector<256x32xf16>147 %2 = xegpu.create_nd_tdesc %arg1 : memref<4096x4096xf16> -> !xegpu.tensor_desc<32x256xf16>148 %3 = xegpu.load_nd %2[0, 0] : !xegpu.tensor_desc<32x256xf16> -> vector<32x256xf16>149 %4 = xegpu.create_nd_tdesc %arg2 : memref<4096x4096xf16> -> !xegpu.tensor_desc<256x256xf16>150 %5 = xegpu.load_nd %4[0, 0] : !xegpu.tensor_desc<256x256xf16> -> vector<256x256xf16>151 // CHECK: = xegpu.dpas152 // CHECK-SAME: {layout_result_0 = #xegpu.layout<sg_layout = [8, 4], sg_data = [32, 64], inst_data = [8, 16]>}153 %6 = xegpu.dpas %1, %3, %5 : vector<256x32xf16>, vector<32x256xf16>, vector<256x256xf16> -> vector<256x256xf16>154 return155}156 157module attributes {transform.with_named_sequence} {158 transform.named_sequence @__transform_main(%arg1: !transform.any_op {transform.readonly}) {159 %0 = transform.structured.match ops{["xegpu.dpas"]} in %arg1 : (!transform.any_op) -> !transform.any_op160 // CHECK: transform.xegpu.set_op_layout_attr %{{.*}}161 transform.xegpu.set_op_layout_attr %0 result sg_layout = [8, 4] sg_data = [32, 64] inst_data = [8, 16] : !transform.any_op162 transform.yield163 }164}165 166// -----167 168// CHECK-LABEL: @set_op_layout_attr_result_sg_param169func.func @set_op_layout_attr_result_sg_param(%arg0: memref<4096x4096xf16>) {170 %0 = xegpu.create_nd_tdesc %arg0 : memref<4096x4096xf16> -> !xegpu.tensor_desc<256x32xf16>171 %1 = xegpu.load_nd %0[0, 0] : !xegpu.tensor_desc<256x32xf16> -> vector<256x32xf16>172 // CHECK: = arith.extf %1173 // CHECK-SAME: {layout_result_0 = #xegpu.layout<sg_layout = [8, 4], sg_data = [32, 64], inst_data = [8, 16]>}174 %2 = arith.extf %1 : vector<256x32xf16> to vector<256x32xf32>175 return176}177 178module attributes {transform.with_named_sequence} {179 transform.named_sequence @__transform_main(%arg1: !transform.any_op {transform.readonly}) {180 %0 = transform.structured.match ops{["arith.extf"]} in %arg1 : (!transform.any_op) -> !transform.any_op181 // CHECK: transform.xegpu.set_op_layout_attr %{{.*}}182 %layout0 = transform.param.constant 8 : i64 -> !transform.param<i64>183 transform.xegpu.set_op_layout_attr %0 result sg_layout = [%layout0, 4] sg_data = [32, 64] inst_data = [8, 16] : !transform.any_op, !transform.param<i64>184 transform.yield185 }186}187 188// -----189 190// CHECK-LABEL: @set_op_layout_attr_result_sg_param2191func.func @set_op_layout_attr_result_sg_param2(%arg0: memref<4096x4096xf16>) {192 %0 = xegpu.create_nd_tdesc %arg0 : memref<4096x4096xf16> -> !xegpu.tensor_desc<256x32xf16>193 %1 = xegpu.load_nd %0[0, 0] : !xegpu.tensor_desc<256x32xf16> -> vector<256x32xf16>194 // CHECK: = arith.extf %1195 // CHECK-SAME: {layout_result_0 = #xegpu.layout<sg_layout = [8, 4], sg_data = [32, 64], inst_data = [8, 16]>}196 %2 = arith.extf %1 : vector<256x32xf16> to vector<256x32xf32>197 return198}199 200module attributes {transform.with_named_sequence} {201 transform.named_sequence @__transform_main(%arg1: !transform.any_op {transform.readonly}) {202 %0 = transform.structured.match ops{["arith.extf"]} in %arg1 : (!transform.any_op) -> !transform.any_op203 // CHECK: transform.xegpu.set_op_layout_attr %{{.*}}204 %layout0 = transform.param.constant 8 : i64 -> !transform.param<i64>205 %layout1 = transform.param.constant 4 : i64 -> !transform.param<i64>206 transform.xegpu.set_op_layout_attr %0 result sg_layout = [%layout0, %layout1] sg_data = [32, 64] inst_data = [8, 16] : !transform.any_op, !transform.param<i64>, !transform.param<i64>207 transform.yield208 }209}210 211// -----212 213// CHECK-LABEL: @set_op_layout_attr_result0214func.func @set_op_layout_attr_result0(%arg0: memref<4096x4096xf16>) {215 %0 = xegpu.create_nd_tdesc %arg0 : memref<4096x4096xf16> -> !xegpu.tensor_desc<256x32xf16>216 %1 = xegpu.load_nd %0[0, 0] : !xegpu.tensor_desc<256x32xf16> -> vector<256x32xf16>217 // CHECK: = arith.extf %1218 // CHECK-SAME: {layout_result_0 = #xegpu.layout<sg_layout = [8, 4], sg_data = [32, 64], inst_data = [8, 16]>}219 %2 = arith.extf %1 : vector<256x32xf16> to vector<256x32xf32>220 return221}222 223module attributes {transform.with_named_sequence} {224 transform.named_sequence @__transform_main(%arg1: !transform.any_op {transform.readonly}) {225 %0 = transform.structured.match ops{["arith.extf"]} in %arg1 : (!transform.any_op) -> !transform.any_op226 // CHECK: transform.xegpu.set_op_layout_attr %{{.*}}227 transform.xegpu.set_op_layout_attr %0 result index = 0 sg_layout = [8, 4] sg_data = [32, 64] inst_data = [8, 16] : !transform.any_op228 transform.yield229 }230}231 232// -----233 234// CHECK-LABEL: @set_op_layout_attr_result_slice235func.func @set_op_layout_attr_result_slice(%arg0: vector<256xf16>) {236 // CHECK: = arith.extf237 // CHECK-SAME: {layout_result_0 = #xegpu.slice<#xegpu.layout<sg_layout = [8, 4], sg_data = [32, 64], inst_data = [8, 16]>, dims = [0]>}238 %2 = arith.extf %arg0 : vector<256xf16> to vector<256xf32>239 return240}241 242module attributes {transform.with_named_sequence} {243 transform.named_sequence @__transform_main(%arg1: !transform.any_op {transform.readonly}) {244 %0 = transform.structured.match ops{["arith.extf"]} in %arg1 : (!transform.any_op) -> !transform.any_op245 // CHECK: transform.xegpu.set_op_layout_attr %{{.*}}246 transform.xegpu.set_op_layout_attr %0 result index = 0 sg_layout = [8, 4] sg_data = [32, 64] inst_data = [8, 16] slice_dims = [0] : !transform.any_op247 transform.yield248 }249}250 251// -----252 253// CHECK-LABEL: @set_op_layout_attr_operand_minimal254func.func @set_op_layout_attr_operand_minimal(%arg0: memref<4096x4096xf16>) {255 %0 = xegpu.create_nd_tdesc %arg0 : memref<4096x4096xf16> -> !xegpu.tensor_desc<256x32xf16>256 %1 = xegpu.load_nd %0[0, 0] : !xegpu.tensor_desc<256x32xf16> -> vector<256x32xf16>257 // CHECK: = arith.extf %1258 // CHECK-SAME: {layout_operand_0 = #xegpu.layout<sg_layout = [8, 4], sg_data = [32, 64]>}259 %2 = arith.extf %1 : vector<256x32xf16> to vector<256x32xf32>260 return261}262 263module attributes {transform.with_named_sequence} {264 transform.named_sequence @__transform_main(%arg1: !transform.any_op {transform.readonly}) {265 %0 = transform.structured.match ops{["arith.extf"]} in %arg1 : (!transform.any_op) -> !transform.any_op266 // CHECK: transform.xegpu.set_op_layout_attr %{{.*}}267 transform.xegpu.set_op_layout_attr %0 sg_layout = [8, 4] sg_data = [32, 64] : !transform.any_op268 transform.yield269 }270}271 272// -----273 274// CHECK-LABEL: @set_op_layout_attr_operand1275func.func @set_op_layout_attr_operand1(%arg0: memref<4096x4096xf16>, %arg1: memref<4096x4096xf16>) {276 %0 = xegpu.create_nd_tdesc %arg0 : memref<4096x4096xf16> -> !xegpu.tensor_desc<256x32xf16>277 %1 = xegpu.load_nd %0[0, 0] : !xegpu.tensor_desc<256x32xf16> -> vector<256x32xf16>278 %2 = xegpu.create_nd_tdesc %arg1 : memref<4096x4096xf16> -> !xegpu.tensor_desc<256x32xf16>279 %3 = xegpu.load_nd %2[0, 0] : !xegpu.tensor_desc<256x32xf16> -> vector<256x32xf16>280 // CHECK: = arith.addf %1, %3281 // CHECK-SAME: {layout_operand_1 = #xegpu.layout<sg_layout = [8, 4], sg_data = [32, 64], inst_data = [8, 16]>}282 %6 = arith.addf %1, %3 : vector<256x32xf16>283 return284}285 286module attributes {transform.with_named_sequence} {287 transform.named_sequence @__transform_main(%arg1: !transform.any_op {transform.readonly}) {288 %0 = transform.structured.match ops{["arith.addf"]} in %arg1 : (!transform.any_op) -> !transform.any_op289 // CHECK: transform.xegpu.set_op_layout_attr %{{.*}}290 transform.xegpu.set_op_layout_attr %0 index = 1 sg_layout = [8, 4] sg_data = [32, 64] inst_data = [8, 16] : !transform.any_op291 transform.yield292 }293}294 295// -----296 297// CHECK-LABEL: @set_gpu_launch_threads298func.func @set_gpu_launch_threads(%arg0: memref<4096x4096xf16>) {299 // CHECK: %[[C1:.+]] = arith.constant 1 : index300 %c1 = arith.constant 1 : index301 // CHECK: %[[C16:.+]] = arith.constant 16 : index302 %c16 = arith.constant 16 : index303 // CHECK: %[[C8:.+]] = arith.constant 8 : index304 // CHECK: %[[C4:.+]] = arith.constant 4 : index305 // CHECK: %[[C1_0:.+]] = arith.constant 1 : index306 // CHECK: gpu.launch blocks(%{{.*}}, %{{.*}}, %{{.*}}) in (%{{.*}} = %[[C16]], %{{.*}} = %[[C16]], %{{.*}} = %[[C1]])307 // CHECK-SAME: threads(%{{.*}}, %{{.*}}, %{{.*}}) in (%{{.*}} = %[[C8]], %{{.*}} = %[[C4]], %{{.*}} = %[[C1_0]])308 gpu.launch blocks(%arg3, %arg4, %arg5) in (%arg9 = %c16, %arg10 = %c16, %arg11 = %c1) threads(%arg6, %arg7, %arg8) in (%arg12 = %c1, %arg13 = %c1, %arg14 = %c1) {309 gpu.terminator310 }311 return312}313module attributes {transform.with_named_sequence} {314 transform.named_sequence @__transform_main(%arg1: !transform.any_op {transform.readonly}) {315 %0 = transform.structured.match ops{["gpu.launch"]} in %arg1 : (!transform.any_op) -> !transform.any_op316 // CHECK: transform.xegpu.set_gpu_launch_threads %{{.*}}317 transform.xegpu.set_gpu_launch_threads %0 threads = [8, 4, 1] : !transform.any_op318 transform.yield319 }320}321 322// -----323 324// CHECK-LABEL: @set_gpu_launch_threads_param325func.func @set_gpu_launch_threads_param(%arg0: memref<4096x4096xf16>) {326 // CHECK: %[[C1:.+]] = arith.constant 1 : index327 %c1 = arith.constant 1 : index328 // CHECK: %[[C16:.+]] = arith.constant 16 : index329 %c16 = arith.constant 16 : index330 // CHECK: %[[C8:.+]] = arith.constant 8 : index331 // CHECK: %[[C4:.+]] = arith.constant 4 : index332 // CHECK: %[[C1_0:.+]] = arith.constant 1 : index333 // CHECK: gpu.launch blocks(%{{.*}}, %{{.*}}, %{{.*}}) in (%{{.*}} = %[[C16]], %{{.*}} = %[[C16]], %{{.*}} = %[[C1]])334 // CHECK-SAME: threads(%{{.*}}, %{{.*}}, %{{.*}}) in (%{{.*}} = %[[C8]], %{{.*}} = %[[C4]], %{{.*}} = %[[C1_0]])335 gpu.launch blocks(%arg3, %arg4, %arg5) in (%arg9 = %c16, %arg10 = %c16, %arg11 = %c1) threads(%arg6, %arg7, %arg8) in (%arg12 = %c1, %arg13 = %c1, %arg14 = %c1) {336 gpu.terminator337 }338 return339}340module attributes {transform.with_named_sequence} {341 transform.named_sequence @__transform_main(%arg1: !transform.any_op {transform.readonly}) {342 %0 = transform.structured.match ops{["gpu.launch"]} in %arg1 : (!transform.any_op) -> !transform.any_op343 // CHECK: transform.xegpu.set_gpu_launch_threads %{{.*}}344 %th1 = transform.param.constant 4 : i64 -> !transform.param<i64>345 transform.xegpu.set_gpu_launch_threads %0 threads = [8, %th1, 1] : !transform.any_op, !transform.param<i64>346 transform.yield347 }348}349 350// -----351 352// CHECK-LABEL: @insert_prefetch_dpas_a353func.func @insert_prefetch_dpas_a(%arg0: memref<4096x4096xf16>, %arg1: memref<4096x4096xf16>, %arg2: memref<4096x4096xf16>) {354 // CHECK: %[[C32:.+]] = arith.constant 32 : index355 %c32 = arith.constant 32 : index356 %c4096 = arith.constant 4096 : index357 // CHECK: %[[C0:.+]] = arith.constant 0 : index358 %c0 = arith.constant 0 : index359 %0 = xegpu.create_nd_tdesc %arg2 : memref<4096x4096xf16> -> !xegpu.tensor_desc<256x256xf16>360 %1 = xegpu.load_nd %0[%c0, %c0] : !xegpu.tensor_desc<256x256xf16> -> vector<256x256xf16>361 // CHECK: xegpu.create_nd_tdesc %arg0362 // CHECK: xegpu.create_nd_tdesc %arg1363 // CHECK: %[[V0:.+]] = xegpu.create_nd_tdesc %arg0364 // CHECK-SAME: !xegpu.tensor_desc<256x32xf16365 // CHECK: xegpu.prefetch_nd %[[V0]][%[[C0]], %[[C0]]]366 %3 = xegpu.create_nd_tdesc %arg0 : memref<4096x4096xf16> -> !xegpu.tensor_desc<256x32xf16>367 %4 = xegpu.create_nd_tdesc %arg1 : memref<4096x4096xf16> -> !xegpu.tensor_desc<32x256xf16>368 // CHECK: scf.for %[[ARG3:.+]] = %[[C0]]369 %2 = scf.for %arg3 = %c0 to %c4096 step %c32 iter_args(%arg4 = %1) -> (vector<256x256xf16>) {370 // CHECK: %[[ADD:.+]] = arith.addi %[[ARG3]], %[[C32]]371 // CHECK: xegpu.prefetch_nd %[[V0]][%[[C0]], %[[ADD]]]372 %5 = xegpu.load_nd %3[%c0, %arg3] : !xegpu.tensor_desc<256x32xf16> -> vector<256x32xf16>373 %6 = xegpu.load_nd %4[%arg3, %c0] : !xegpu.tensor_desc<32x256xf16> -> vector<32x256xf16>374 %7 = xegpu.dpas %5, %6, %arg4 : vector<256x32xf16>, vector<32x256xf16>, vector<256x256xf16> -> vector<256x256xf16>375 scf.yield %7 : vector<256x256xf16>376 }377 return378}379 380module attributes {transform.with_named_sequence} {381 transform.named_sequence @__transform_main(%arg0: !transform.any_op {transform.readonly}) {382 %func = transform.structured.match ops{["func.func"]} in %arg0 : (!transform.any_op) -> !transform.any_op383 %0 = transform.structured.match ops{["xegpu.dpas"]} in %func : (!transform.any_op) -> !transform.any_op384 %1 = transform.get_operand %0[0] : (!transform.any_op) -> !transform.any_value385 // CHECK: transform.xegpu.insert_prefetch %{{.*}}386 %2 = transform.xegpu.insert_prefetch %1 nb_prefetch = 1 : (!transform.any_value) -> !transform.any_op387 transform.apply_patterns to %func {388 transform.apply_patterns.canonicalization389 } : !transform.any_op390 391 transform.yield392 }393}394 395// -----396 397// CHECK-LABEL: @insert_prefetch_dpas_a_nb_param2398func.func @insert_prefetch_dpas_a_nb_param2(%arg0: memref<4096x4096xf16>, %arg1: memref<4096x4096xf16>, %arg2: memref<4096x4096xf16>) {399 // CHECK: %[[C64:.+]] = arith.constant 64 : index400 // CHECK: %[[C32:.+]] = arith.constant 32 : index401 %c32 = arith.constant 32 : index402 %c4096 = arith.constant 4096 : index403 // CHECK: %[[C0:.+]] = arith.constant 0 : index404 %c0 = arith.constant 0 : index405 %0 = xegpu.create_nd_tdesc %arg2 : memref<4096x4096xf16> -> !xegpu.tensor_desc<256x256xf16>406 %1 = xegpu.load_nd %0[0, 0] : !xegpu.tensor_desc<256x256xf16> -> vector<256x256xf16>407 // CHECK: xegpu.create_nd_tdesc %arg0408 // CHECK: xegpu.create_nd_tdesc %arg1409 // CHECK: %[[V0:.+]] = xegpu.create_nd_tdesc %arg0410 // CHECK-SAME: !xegpu.tensor_desc<256x32xf16411 // CHECK: xegpu.prefetch_nd %[[V0]][0, %[[C0]]]412 // CHECK: xegpu.prefetch_nd %[[V0]][0, %[[C32]]]413 %3 = xegpu.create_nd_tdesc %arg0 : memref<4096x4096xf16> -> !xegpu.tensor_desc<256x32xf16>414 %4 = xegpu.create_nd_tdesc %arg1 : memref<4096x4096xf16> -> !xegpu.tensor_desc<32x256xf16>415 // CHECK: scf.for %[[ARG3:.+]] = %[[C0]]416 %2 = scf.for %arg3 = %c0 to %c4096 step %c32 iter_args(%arg4 = %1) -> (vector<256x256xf16>) {417 // CHECK: %[[ADD:.+]] = arith.addi %[[ARG3]], %[[C64]]418 // CHECK: xegpu.prefetch_nd %[[V0]][0, %[[ADD]]]419 %5 = xegpu.load_nd %3[0, %arg3] : !xegpu.tensor_desc<256x32xf16> -> vector<256x32xf16>420 %6 = xegpu.load_nd %4[%arg3, 0] : !xegpu.tensor_desc<32x256xf16> -> vector<32x256xf16>421 %7 = xegpu.dpas %5, %6, %arg4 : vector<256x32xf16>, vector<32x256xf16>, vector<256x256xf16> -> vector<256x256xf16>422 scf.yield %7 : vector<256x256xf16>423 }424 return425}426 427module attributes {transform.with_named_sequence} {428 transform.named_sequence @__transform_main(%arg0: !transform.any_op {transform.readonly}) {429 %func = transform.structured.match ops{["func.func"]} in %arg0 : (!transform.any_op) -> !transform.any_op430 %0 = transform.structured.match ops{["xegpu.dpas"]} in %func : (!transform.any_op) -> !transform.any_op431 %1 = transform.get_operand %0[0] : (!transform.any_op) -> !transform.any_value432 %nb = transform.param.constant 2 : i64 -> !transform.param<i64>433 // CHECK: transform.xegpu.insert_prefetch %{{.*}}434 %2 = transform.xegpu.insert_prefetch %1 nb_prefetch = %nb : (!transform.any_value, !transform.param<i64>) -> !transform.any_op435 transform.apply_patterns to %func {436 transform.apply_patterns.canonicalization437 } : !transform.any_op438 transform.yield439 }440}441 442// -----443 444// CHECK-LABEL: @convert_layout_a445func.func @convert_layout_a(%arg0: memref<4096x4096xf16>, %arg1: memref<4096x4096xf16>, %arg2: memref<4096x4096xf16>) {446 %c0 = arith.constant 0 : index447 // CHECK: %[[V0:.+]] = xegpu.create_nd_tdesc %arg0448 %0 = xegpu.create_nd_tdesc %arg0 : memref<4096x4096xf16> -> !xegpu.tensor_desc<256x32xf16, #xegpu.layout<sg_layout = [8, 4], sg_data = [32, 32], inst_data = [32, 16]>>449 // CHECK: %[[V1:.+]] = xegpu.load_nd %[[V0]]450 %1 = xegpu.load_nd %0[%c0, %c0] : !xegpu.tensor_desc<256x32xf16, #xegpu.layout<sg_layout = [8, 4], sg_data = [32, 32], inst_data = [32, 16]>> -> vector<256x32xf16>451 // CHECK: %[[V2:.+]] = xegpu.convert_layout %[[V1]]452 // CHECK: input_layout = #xegpu.layout<sg_layout = [8, 4], sg_data = [32, 32], inst_data = [32, 16]>453 // CHECK: target_layout = #xegpu.layout<sg_layout = [8, 4], sg_data = [32, 32], inst_data = [8, 16]>454 %2 = xegpu.create_nd_tdesc %arg1 : memref<4096x4096xf16> -> !xegpu.tensor_desc<32x256xf16>455 %3 = xegpu.load_nd %2[%c0, %c0] : !xegpu.tensor_desc<32x256xf16> -> vector<32x256xf16>456 %4 = xegpu.create_nd_tdesc %arg2 : memref<4096x4096xf16> -> !xegpu.tensor_desc<256x256xf16>457 %5 = xegpu.load_nd %4[%c0, %c0] : !xegpu.tensor_desc<256x256xf16> -> vector<256x256xf16>458 // CHECK: = xegpu.dpas %[[V2]]459 %6 = xegpu.dpas %1, %3, %5 : vector<256x32xf16>, vector<32x256xf16>, vector<256x256xf16> -> vector<256x256xf16>460 return461}462 463module attributes {transform.with_named_sequence} {464 transform.named_sequence @__transform_main(%arg1: !transform.any_op {transform.readonly}) {465 %0 = transform.structured.match ops{["xegpu.dpas"]} in %arg1 : (!transform.any_op) -> !transform.any_op466 %1 = transform.get_operand %0[0] : (!transform.any_op) -> !transform.any_value467 // CHECK: transform.xegpu.convert_layout %{{.*}}468 transform.xegpu.convert_layout %1469 input_sg_layout = [8, 4] input_sg_data = [32, 32] input_inst_data = [32, 16]470 target_sg_layout = [8, 4] target_sg_data = [32, 32] target_inst_data = [8, 16]471 : (!transform.any_value) -> !transform.any_op472 transform.yield473 }474}475 476// -----477 478// CHECK-LABEL: @convert_layout_a_sg_param479func.func @convert_layout_a_sg_param(%arg0: memref<4096x4096xf16>, %arg1: memref<4096x4096xf16>, %arg2: memref<4096x4096xf16>) {480 %c0 = arith.constant 0 : index481 // CHECK: %[[V0:.+]] = xegpu.create_nd_tdesc %arg0482 %0 = xegpu.create_nd_tdesc %arg0 : memref<4096x4096xf16> -> !xegpu.tensor_desc<256x32xf16, #xegpu.layout<sg_layout = [8, 4], sg_data = [32, 32], inst_data = [32, 16]>>483 // CHECK: %[[V1:.+]] = xegpu.load_nd %[[V0]]484 %1 = xegpu.load_nd %0[%c0, %c0] : !xegpu.tensor_desc<256x32xf16, #xegpu.layout<sg_layout = [8, 4], sg_data = [32, 32], inst_data = [32, 16]>> -> vector<256x32xf16>485 // CHECK: %[[V2:.+]] = xegpu.convert_layout %[[V1]]486 // CHECK: input_layout = #xegpu.layout<sg_layout = [8, 4], sg_data = [32, 32], inst_data = [32, 16]>487 // CHECK: target_layout = #xegpu.layout<sg_layout = [8, 4], sg_data = [32, 32], inst_data = [8, 16]>488 %2 = xegpu.create_nd_tdesc %arg1 : memref<4096x4096xf16> -> !xegpu.tensor_desc<32x256xf16>489 %3 = xegpu.load_nd %2[%c0, %c0] : !xegpu.tensor_desc<32x256xf16> -> vector<32x256xf16>490 %4 = xegpu.create_nd_tdesc %arg2 : memref<4096x4096xf16> -> !xegpu.tensor_desc<256x256xf16>491 %5 = xegpu.load_nd %4[%c0, %c0] : !xegpu.tensor_desc<256x256xf16> -> vector<256x256xf16>492 // CHECK: = xegpu.dpas %[[V2]]493 %6 = xegpu.dpas %1, %3, %5 : vector<256x32xf16>, vector<32x256xf16>, vector<256x256xf16> -> vector<256x256xf16>494 return495}496 497module attributes {transform.with_named_sequence} {498 transform.named_sequence @__transform_main(%arg1: !transform.any_op {transform.readonly}) {499 %0 = transform.structured.match ops{["xegpu.dpas"]} in %arg1 : (!transform.any_op) -> !transform.any_op500 %1 = transform.get_operand %0[0] : (!transform.any_op) -> !transform.any_value501 %layout0 = transform.param.constant 8 : i64 -> !transform.param<i64>502 // CHECK: transform.xegpu.convert_layout %{{.*}}503 transform.xegpu.convert_layout %1504 input_sg_layout = [%layout0, 4] input_sg_data = [32, 32] input_inst_data = [32, 16]505 target_sg_layout = [%layout0, 4] target_sg_data = [32, 32] target_inst_data = [8, 16]506 : (!transform.any_value, !transform.param<i64>, !transform.param<i64>) -> !transform.any_op507 transform.yield508 }509}510