233 lines · plain
1// RUN: mlir-opt --xegpu-wg-to-sg-distribute -split-input-file %s | FileCheck %s2 3gpu.module @test_round_robin_assignment {4 // CHECK-LABEL: create_nd_tdesc5 // CHECK-SAME: %[[ARG_0:.*]]: memref<256x128xf32>6 gpu.func @create_nd_tdesc(%src: memref<256x128xf32>) {7 // CHECK-COUNT-4: xegpu.create_nd_tdesc %[[ARG_0]][%{{.*}}, %{{.*}}] : memref<256x128xf32> -> !xegpu.tensor_desc<16x16xf32, #xegpu.layout<lane_layout = [1, 16], lane_data = [1, 1]>>8 // CHECK-NOT: xegpu.create_nd_tdesc9 %tdesc = xegpu.create_nd_tdesc %src[0, 0] : memref<256x128xf32>10 -> !xegpu.tensor_desc<256x128xf32, #xegpu.layout<sg_layout = [8, 4], sg_data = [16, 16], lane_layout = [1, 16], lane_data = [1, 1]>>11 gpu.return12 }13 14 // CHECK-LABEL: create_nd_tdesc_with_shared_data15 // CHECK-SAME: %[[ARG_0:.*]]: memref<256x128xf32>16 gpu.func @create_nd_tdesc_with_shared_data(%src: memref<256x128xf32>) {17 // CHECK: %[[SGID:.*]] = gpu.subgroup_id : index18 // CHECK: %[[C4:.*]] = arith.constant 4 : index19 // CHECK: %[[IDX:.*]] = index.remu %[[SGID]], %[[C4]]20 // CHECK: %[[IDY_DIV:.*]] = index.divu %[[SGID]], %[[C4]]21 // CHECK: %[[C8:.*]] = arith.constant 8 : index22 // CHECK: %[[IDY:.*]] = index.remu %[[IDY_DIV]], %[[C8]]23 // CHECK: %[[C16:.*]] = arith.constant 16 : index24 // CHECK: %[[LY:.*]] = index.mul %[[IDY]], %[[C16]]25 // CHECK: %[[C64:.*]] = arith.constant 64 : index26 // CHECK: %[[LX:.*]] = index.mul %[[IDX]], %[[C64]]27 // CHECK: %[[C128:.*]] = arith.constant 128 : index28 // CHECK: %[[OFFY:.*]] = index.remu %[[LY]], %[[C128]]29 // CHECK: %[[C64_1:.*]] = arith.constant 64 : index30 // CHECK: %[[OFFX:.*]] = index.remu %[[LX]], %[[C64_1]]31 // CHECK: xegpu.create_nd_tdesc %[[ARG_0]][%[[OFFY]], %[[OFFX]]] : memref<256x128xf32> -> !xegpu.tensor_desc<16x64xf32>32 %tdesc = xegpu.create_nd_tdesc %src[0, 0] : memref<256x128xf32>33 -> !xegpu.tensor_desc<128x64xf32, #xegpu.layout<sg_layout = [8, 4], sg_data = [16, 64]>>34 gpu.return35 }36 37 // CHECK-LABEL: load_nd_tdesc38 // CHECK-SAME: %[[ARG_0:.*]]: memref<256x128xf32>39 gpu.func @load_nd_tdesc(%src: memref<256x128xf32>) {40 %tdesc = xegpu.create_nd_tdesc %src[0, 0] : memref<256x128xf32>41 -> !xegpu.tensor_desc<256x128xf32, #xegpu.layout<sg_layout = [8, 4], sg_data = [16, 16], lane_layout = [1, 16], lane_data = [1, 1]>>42 // CHECK-COUNT-4: xegpu.load_nd %{{.*}} : !xegpu.tensor_desc<16x16xf32, #xegpu.layout<lane_layout = [1, 16], lane_data = [1, 1]>> -> vector<16x16xf32>43 // CHECK-NOT: xegpu.load_nd44 %load = xegpu.load_nd %tdesc45 : !xegpu.tensor_desc<256x128xf32, #xegpu.layout<sg_layout = [8, 4], sg_data = [16, 16], lane_layout = [1, 16], lane_data = [1, 1]>>46 -> vector<256x128xf32>47 gpu.return48 }49 50 // CHECK-LABEL: store_nd51 // CHECK-SAME: %[[ARG_0:.*]]: memref<256x128xf32>52 gpu.func @store_nd(%src: memref<256x128xf32>) {53 %tdesc = xegpu.create_nd_tdesc %src[0, 0] : memref<256x128xf32>54 -> !xegpu.tensor_desc<256x128xf32, #xegpu.layout<sg_layout = [8, 4], sg_data = [16, 16], lane_layout = [1, 16], lane_data = [1, 1]>>55 // CHECK-COUNT-4: xegpu.store_nd %{{.*}}, %{{.*}} : vector<16x16xf32>, !xegpu.tensor_desc<16x16xf32, #xegpu.layout<lane_layout = [1, 16], lane_data = [1, 1]>>56 // CHECK-NOT: xegpu.store_nd57 %load = xegpu.load_nd %tdesc58 : !xegpu.tensor_desc<256x128xf32, #xegpu.layout<sg_layout = [8, 4], sg_data = [16, 16], lane_layout = [1, 16], lane_data = [1, 1]>>59 -> vector<256x128xf32>60 xegpu.store_nd %load, %tdesc61 : vector<256x128xf32>, !xegpu.tensor_desc<256x128xf32, #xegpu.layout<sg_layout = [8, 4], sg_data = [16, 16], lane_layout = [1, 16], lane_data = [1, 1]>>62 gpu.return63 }64 65 // CHECK-LABEL: update_nd66 // CHECK-SAME: %[[ARG_0:.*]]: memref<256x128xf32>67 gpu.func @update_nd(%src: memref<256x128xf32>){68 %tdesc = xegpu.create_nd_tdesc %src[0, 0] : memref<256x128xf32>69 -> !xegpu.tensor_desc<256x128xf32, #xegpu.layout<sg_layout = [8, 4], sg_data = [16, 16], lane_layout = [1, 16], lane_data = [1, 1]>>70 // CHECK-COUNT-4: xegpu.update_nd_offset %{{.*}}, [0, 16] : !xegpu.tensor_desc<16x16xf32, #xegpu.layout<lane_layout = [1, 16], lane_data = [1, 1]>>71 // CHECK-NOT: xegpu.update_nd_offset72 %update = xegpu.update_nd_offset %tdesc, [0, 16]73 : !xegpu.tensor_desc<256x128xf32, #xegpu.layout<sg_layout = [8, 4], sg_data = [16, 16], lane_layout = [1, 16], lane_data = [1, 1]>>74 gpu.return75 }76 77 // CHECK-LABEL: dpas78 // CHECK-SAME: (%[[ARG_0:.*]]: memref<256x128xf16>, %[[ARG_1:.*]]: memref<128x256xf16>)79 gpu.func @dpas(%a: memref<256x128xf16>, %b: memref<128x256xf16>) {80 // CHECK-COUNT-4: xegpu.create_nd_tdesc %[[ARG_0]][%{{.*}}, %{{.*}}] : memref<256x128xf16> -> !xegpu.tensor_desc<16x16xf16, #xegpu.layout<lane_layout = [1, 16], lane_data = [1, 1]>>81 // CHECK-COUNT-4: xegpu.create_nd_tdesc %[[ARG_1]][%{{.*}}, %{{.*}}] : memref<128x256xf16> -> !xegpu.tensor_desc<16x16xf16, #xegpu.layout<lane_layout = [1, 16], lane_data = [2, 1]>>82 // CHECK-COUNT-16: xegpu.dpas %{{.*}}, %{{.*}} {layout_result_0 = #xegpu.layout<lane_layout = [1, 16], lane_data = [1, 1]>} : vector<16x16xf16>, vector<16x16xf16> -> vector<16x16xf32>83 // CHECK-NOT: xegpu.dpas84 %tdesc_a = xegpu.create_nd_tdesc %a[0, 0] : memref<256x128xf16>85 -> !xegpu.tensor_desc<256x128xf16, #xegpu.layout<sg_layout = [8, 4], sg_data = [16, 16], lane_layout = [1, 16], lane_data = [1, 1]>>86 %load_a = xegpu.load_nd %tdesc_a87 : !xegpu.tensor_desc<256x128xf16, #xegpu.layout<sg_layout = [8, 4], sg_data = [16, 16], lane_layout = [1, 16], lane_data = [1, 1]>>88 -> vector<256x128xf16>89 %tdesc_b = xegpu.create_nd_tdesc %b[0, 0] : memref<128x256xf16>90 -> !xegpu.tensor_desc<128x256xf16, #xegpu.layout<sg_layout = [4, 8], sg_data = [16, 16], lane_layout = [1, 16], lane_data = [2, 1]>>91 %load_b = xegpu.load_nd %tdesc_b92 : !xegpu.tensor_desc<128x256xf16, #xegpu.layout<sg_layout = [4, 8], sg_data = [16, 16], lane_layout = [1, 16], lane_data = [2, 1]>>93 -> vector<128x256xf16>94 %dpas = xegpu.dpas %load_a, %load_b95 {layout_result_0 = #xegpu.layout<sg_layout = [8, 8], sg_data = [16, 16], lane_layout = [1, 16], lane_data = [1, 1]>}96 : vector<256x128xf16>, vector<128x256xf16> -> vector<256x256xf32>97 gpu.return98 }99 100 // CHECK-LABEL: prefetch_nd_tdesc101 // CHECK-SAME: %[[ARG_0:.*]]: memref<256x128xf32>102 gpu.func @prefetch_nd_tdesc(%src: memref<256x128xf32>) {103 // CHECK-COUNT-4: xegpu.prefetch_nd %{{.*}} : !xegpu.tensor_desc<16x16xf32, #xegpu.layout<lane_layout = [1, 16], lane_data = [1, 1]>>104 // CHECK-NOT: xegpu.prefetch_nd105 %tdesc = xegpu.create_nd_tdesc %src[0, 0] : memref<256x128xf32>106 -> !xegpu.tensor_desc<256x128xf32, #xegpu.layout<sg_layout = [8, 4], sg_data = [16, 16], lane_layout = [1, 16], lane_data = [1, 1]>>107 xegpu.prefetch_nd %tdesc108 : !xegpu.tensor_desc<256x128xf32, #xegpu.layout<sg_layout = [8, 4], sg_data = [16, 16], lane_layout = [1, 16], lane_data = [1, 1]>>109 gpu.return110 }111 112 // CHECK-LABEL: broadcast113 // CHECK-SAME: %[[ARG_0:.*]]: memref<128x1xf32>114 gpu.func @broadcast(%src: memref<128x1xf32>) {115 %tdesc = xegpu.create_nd_tdesc %src[0, 0] : memref<128x1xf32>116 -> !xegpu.tensor_desc<128x1xf32, #xegpu.layout<sg_layout = [4, 1], sg_data = [16, 1], lane_layout = [8, 1], lane_data = [1, 1]>>117 %load = xegpu.load_nd %tdesc118 : !xegpu.tensor_desc<128x1xf32, #xegpu.layout<sg_layout = [4, 1], sg_data = [16, 1], lane_layout = [8, 1], lane_data = [1, 1]>>119 -> vector<128x1xf32>120 // CHECK-COUNT-2: vector.broadcast {{.*}} {layout_result_0 = #xegpu.layout<lane_layout = [8, 1], lane_data = [1, 1]>} : vector<16x1xf32> to vector<16x32xf32>121 // CHECK-NOT: vector.broadcast122 %broadcast = vector.broadcast %load123 {layout_result_0 = #xegpu.layout<sg_layout = [4, 1], sg_data = [16, 32], lane_layout = [8, 1], lane_data = [1, 1]>}124 : vector<128x1xf32> to vector<128x64xf32>125 gpu.return126 }127 128 gpu.func @scf_for(%arg0: memref<1024xf32>, %arg1: memref<1024xf32>) {129 %c1 = arith.constant 1 : index130 %c10 = arith.constant 10 : index131 %c0 = arith.constant 0 : index132 %c256 = arith.constant 256 : index133 %c1024 = arith.constant 1024 : index134 %0 = xegpu.create_nd_tdesc %arg0[0] : memref<1024xf32> -> !xegpu.tensor_desc<256xf32, #xegpu.layout<sg_layout = [8], sg_data = [16]>>135 %1 = xegpu.create_nd_tdesc %arg1[0] : memref<1024xf32> -> !xegpu.tensor_desc<256xf32, #xegpu.layout<sg_layout = [8], sg_data = [16]>>136 // CHECK-LABEL: scf.for137 // CHECK-SAME: (!xegpu.tensor_desc<16xf32>, !xegpu.tensor_desc<16xf32>, !xegpu.tensor_desc<16xf32>, !xegpu.tensor_desc<16xf32>)138 %2:2 = scf.for %arg2 = %c0 to %c1024 step %c256 iter_args(%arg3 = %0, %arg4 = %1)139 -> (!xegpu.tensor_desc<256xf32, #xegpu.layout<sg_layout = [8], sg_data = [16]>>, !xegpu.tensor_desc<256xf32, #xegpu.layout<sg_layout = [8], sg_data = [16]>>) {140 %3 = xegpu.load_nd %0 : !xegpu.tensor_desc<256xf32, #xegpu.layout<sg_layout = [8], sg_data = [16]>> -> vector<256xf32>141 xegpu.store_nd %3, %arg3 : vector<256xf32>, !xegpu.tensor_desc<256xf32, #xegpu.layout<sg_layout = [8], sg_data = [16]>>142 %4 = xegpu.update_nd_offset %arg3, [256] : !xegpu.tensor_desc<256xf32, #xegpu.layout<sg_layout = [8], sg_data = [16]>>143 %5 = xegpu.update_nd_offset %arg4, [256] : !xegpu.tensor_desc<256xf32, #xegpu.layout<sg_layout = [8], sg_data = [16]>>144 // CHECK-LABEL: scf.yield145 // CHECK-SAME: !xegpu.tensor_desc<16xf32>, !xegpu.tensor_desc<16xf32>, !xegpu.tensor_desc<16xf32>, !xegpu.tensor_desc<16xf32>146 scf.yield %4, %5 : !xegpu.tensor_desc<256xf32, #xegpu.layout<sg_layout = [8], sg_data = [16]>>, !xegpu.tensor_desc<256xf32, #xegpu.layout<sg_layout = [8], sg_data = [16]>>147 }148 gpu.return149 }150 151 gpu.func @scf_while_and_condition(%arg0: memref<1024xf32>, %arg1: memref<1024xf32>) {152 %c1_i32 = arith.constant 1 : i32153 %c10_i32 = arith.constant 10 : i32154 %c0_i32 = arith.constant 0 : i32155 %0 = xegpu.create_nd_tdesc %arg0[0] : memref<1024xf32> -> !xegpu.tensor_desc<256xf32, #xegpu.layout<sg_layout = [8], sg_data = [16]>>156 %1 = xegpu.load_nd %0 : !xegpu.tensor_desc<256xf32, #xegpu.layout<sg_layout = [8], sg_data = [16]>> -> vector<256xf32>157 %2 = xegpu.create_nd_tdesc %arg1[0] : memref<1024xf32> -> !xegpu.tensor_desc<256xf32, #xegpu.layout<sg_layout = [8], sg_data = [16]>>158 // CHECK: scf.while ({{.*}}) : (vector<16xf32>, vector<16xf32>, i32) -> (vector<16xf32>, vector<16xf32>, i32)159 %3:2 = scf.while (%arg2 = %1, %arg3 = %c0_i32) : (vector<256xf32>, i32) -> (vector<256xf32>, i32) {160 %4 = arith.cmpi slt, %arg3, %c10_i32 : i32161 // CHECK: scf.condition{{.*}} : vector<16xf32>, vector<16xf32>, i32162 scf.condition(%4) %arg2, %arg3 : vector<256xf32>, i32163 } do {164 // CHECK: ([[arg2:%.+]]: vector<16xf32>, [[arg3:%.+]]: vector<16xf32>, [[arg4:%.+]]: i32)165 ^bb0(%arg2: vector<256xf32>, %arg3: i32):166 xegpu.store_nd %arg2, %2 : vector<256xf32>, !xegpu.tensor_desc<256xf32, #xegpu.layout<sg_layout = [8], sg_data = [16]>>167 %4 = arith.addi %arg3, %c1_i32 : i32168 %5 = xegpu.update_nd_offset %0, [256] : !xegpu.tensor_desc<256xf32, #xegpu.layout<sg_layout = [8], sg_data = [16]>>169 %6 = xegpu.load_nd %5 : !xegpu.tensor_desc<256xf32, #xegpu.layout<sg_layout = [8], sg_data = [16]>> -> vector<256xf32>170 scf.yield %6, %4 : vector<256xf32>, i32171 }172 gpu.return173 }174 175 gpu.func @scf_if(%arg0: memref<1024xf32>, %arg1: memref<1024xf32>) {176 %c10 = arith.constant 10 : index177 %0 = gpu.subgroup_id : index178 %1 = xegpu.create_nd_tdesc %arg0[0] : memref<1024xf32> -> !xegpu.tensor_desc<256xf32, #xegpu.layout<sg_layout = [8], sg_data = [16]>>179 %2 = xegpu.create_nd_tdesc %arg1[0] : memref<1024xf32> -> !xegpu.tensor_desc<256xf32, #xegpu.layout<sg_layout = [8], sg_data = [16]>>180 %3 = arith.cmpi eq, %0, %c10 : index181 // CHECK-LABEL: scf.if182 // CHECK-SAME: (vector<16xf32>, vector<16xf32>)183 %4 = scf.if %3 -> (vector<256xf32>) {184 %5 = xegpu.load_nd %1 : !xegpu.tensor_desc<256xf32, #xegpu.layout<sg_layout = [8], sg_data = [16]>> -> vector<256xf32>185 // CHECK-LABEL: scf.yield186 // CHECK-SAME: vector<16xf32>, vector<16xf32>187 scf.yield %5 : vector<256xf32>188 } else {189 %5 = xegpu.load_nd %2 : !xegpu.tensor_desc<256xf32, #xegpu.layout<sg_layout = [8], sg_data = [16]>> -> vector<256xf32>190 // CHECK-LABEL: scf.yield191 // CHECK-SAME: vector<16xf32>, vector<16xf32>192 scf.yield %5 : vector<256xf32>193 } {layout_result_0 = #xegpu.layout<sg_layout = [8], sg_data = [16]>}194 xegpu.store_nd %4, %1 : vector<256xf32>, !xegpu.tensor_desc<256xf32, #xegpu.layout<sg_layout = [8], sg_data = [16]>>195 gpu.return196 }197 198 gpu.func @scf_if_tensor_desc(%arg0: memref<1024xf32>, %arg1: memref<1024xf32>) {199 %c10 = arith.constant 10 : index200 %id = gpu.subgroup_id : index201 202 %t = xegpu.create_nd_tdesc %arg0[0] : memref<1024xf32> -> !xegpu.tensor_desc<256xf32, #xegpu.layout<sg_layout = [8], sg_data = [16]>>203 %d = xegpu.load_nd %t : !xegpu.tensor_desc<256xf32, #xegpu.layout<sg_layout = [8], sg_data = [16]>> -> vector<256xf32>204 205 %0 = arith.cmpi eq, %id, %c10 : index206 // CHECK-LABEL: scf.if207 // CHECK-SAME: (!xegpu.tensor_desc<16xf32>, !xegpu.tensor_desc<16xf32>)208 %1 = scf.if %0 -> (!xegpu.tensor_desc<256xf32, #xegpu.layout<sg_layout = [8], sg_data = [16]>>) {209 %2 = xegpu.create_nd_tdesc %arg0[0] : memref<1024xf32> -> !xegpu.tensor_desc<256xf32, #xegpu.layout<sg_layout = [8], sg_data = [16]>>210 // CHECK-LABEL: scf.yield211 // CHECK-SAME: !xegpu.tensor_desc<16xf32>, !xegpu.tensor_desc<16xf32>212 scf.yield %2 : !xegpu.tensor_desc<256xf32, #xegpu.layout<sg_layout = [8], sg_data = [16]>>213 } else {214 %3 = xegpu.create_nd_tdesc %arg1[0] : memref<1024xf32> -> !xegpu.tensor_desc<256xf32, #xegpu.layout<sg_layout = [8], sg_data = [16]>>215 // CHECK-LABEL: scf.yield216 // CHECK-SAME: !xegpu.tensor_desc<16xf32>, !xegpu.tensor_desc<16xf32>217 scf.yield %3 : !xegpu.tensor_desc<256xf32, #xegpu.layout<sg_layout = [8], sg_data = [16]>>218 }219 xegpu.store_nd %d, %1 : vector<256xf32>, !xegpu.tensor_desc<256xf32, #xegpu.layout<sg_layout = [8], sg_data = [16]>>220 gpu.return221 }222 223 gpu.func @convert_layout_optimal(%arg0: memref<32x64xf32>) {224 %0 = xegpu.create_nd_tdesc %arg0[0, 0] : memref<32x64xf32> -> !xegpu.tensor_desc<32x64xf32, #xegpu.layout<sg_layout = [2, 2], sg_data = [16, 16], inst_data = [16, 16]>>225 // CHECK-COUNT-2: xegpu.load_nd {{.*}} : !xegpu.tensor_desc<16x16xf32, #xegpu.layout<inst_data = [16, 16]>> -> vector<16x16xf32>226 // CHECK-COUNT-2: xegpu.convert_layout {{.*}} <{input_layout = #xegpu.layout<inst_data = [16, 16]>, target_layout = #xegpu.layout<inst_data = [8, 16]>}> : vector<16x16xf32>227 %1 = xegpu.load_nd %0 : !xegpu.tensor_desc<32x64xf32, #xegpu.layout<sg_layout = [2, 2], sg_data = [16, 16], inst_data = [16, 16]>> -> vector<32x64xf32>228 %2 = xegpu.convert_layout %1 <{input_layout = #xegpu.layout<sg_layout = [2, 2], sg_data = [16, 16], inst_data = [16, 16]>,229 target_layout = #xegpu.layout<sg_layout = [2, 2], sg_data = [16, 16], inst_data = [8, 16]>}> : vector<32x64xf32>230 gpu.return231 }232}233