384 lines · plain
1// RUN: mlir-opt --xegpu-wg-to-sg-distribute -split-input-file %s | FileCheck %s2 3gpu.module @test_1_1_assignment {4 // CHECK-LABEL: create_nd_tdesc5 // CHECK-SAME: %[[ARG_0:.*]]: memref<256x128xf32>6 gpu.func @create_nd_tdesc(%src: memref<256x128xf32>) {7 // CHECK-DAG: %[[SGID:.*]] = gpu.subgroup_id : index8 // CHECK-DAG: %[[REMUX:.*]] = index.remu %[[SGID]], %[[C4:.*]]9 // CHECK-DAG: %[[DIVU:.*]] = index.divu %[[SGID]], %[[C4:.*]]10 // CHECK-DAG: %[[REMUY:.*]] = index.remu %[[DIVU]], %[[C8:.*]]11 // CHECK-DAG: %[[MULY:.*]] = index.mul %[[REMUY]], %[[C32:.*]]12 // CHECK-DAG: %[[MULX:.*]] = index.mul %[[REMUX]], %[[C32:.*]]13 // CHECK-DAG: %[[MODY:.*]] = index.remu %[[MULY]], %[[C256:.*]]14 // CHECK-DAG: %[[MODX:.*]] = index.remu %[[MULX]], %[[C128:.*]]15 // CHECK-DAG: %[[TDESC:.*]] = xegpu.create_nd_tdesc %[[ARG_0]][%[[MODY]], %[[MODX]]] : memref<256x128xf32> -> !xegpu.tensor_desc<32x32xf32, #xegpu.layout<lane_layout = [1, 16], lane_data = [1, 1]>>16 %tdesc = xegpu.create_nd_tdesc %src[0, 0] : memref<256x128xf32>17 -> !xegpu.tensor_desc<256x128xf32, #xegpu.layout<sg_layout = [8, 4], sg_data = [32, 32], lane_layout = [1, 16], lane_data = [1, 1]>>18 gpu.return19 }20 21 // CHECK-LABEL: create_nd_tdesc_from_higher_rank_memref22 // CHECK-SAME: %[[ARG_0:.*]]: memref<3x256x128xf32>23 gpu.func @create_nd_tdesc_from_higher_rank_memref(%src: memref<3x256x128xf32>) {24 // CHECK-DAG: %[[SGID:.*]] = gpu.subgroup_id : index25 // CHECK-DAG: %[[REMUX:.*]] = index.remu %[[SGID]], %[[C4:.*]]26 // CHECK-DAG: %[[DIVU:.*]] = index.divu %[[SGID]], %[[C4:.*]]27 // CHECK-DAG: %[[REMUY:.*]] = index.remu %[[DIVU]], %[[C8:.*]]28 // CHECK-DAG: %[[MULY:.*]] = index.mul %[[REMUY]], %[[C32:.*]]29 // CHECK-DAG: %[[MULX:.*]] = index.mul %[[REMUX]], %[[C32:.*]]30 // CHECK-DAG: %[[MODY:.*]] = index.remu %[[MULY]], %[[C256:.*]]31 // CHECK-DAG: %[[MODX:.*]] = index.remu %[[MULX]], %[[C128:.*]]32 // CHECK-DAG: %[[TDESC:.*]] = xegpu.create_nd_tdesc %[[ARG_0]][1, %[[MODY]], %[[MODX]]] : memref<3x256x128xf32> -> !xegpu.tensor_desc<32x32xf32, #xegpu.layout<lane_layout = [1, 16], lane_data = [1, 1]>>33 %tdesc = xegpu.create_nd_tdesc %src[1, 0, 0] : memref<3x256x128xf32>34 -> !xegpu.tensor_desc<256x128xf32, #xegpu.layout<sg_layout = [8, 4], sg_data = [32, 32], lane_layout = [1, 16], lane_data = [1, 1]>>35 gpu.return36 }37 38 // CHECK-LABEL: load_nd_tdesc39 // CHECK-SAME: %[[ARG_0:.*]]: memref<256x128xf32>40 gpu.func @load_nd_tdesc(%src: memref<256x128xf32>) {41 // CHECK: %[[TDESC:.*]] = xegpu.create_nd_tdesc %[[ARG_0]][{{%.*}}, {{%.*}}] : memref<256x128xf32>42 // CHECK-SAME: -> !xegpu.tensor_desc<32x32xf32, #xegpu.layout<lane_layout = [1, 16], lane_data = [1, 1]>>43 // CHECK: %[[LOAD:.*]] = xegpu.load_nd %[[TDESC]]44 // CHECK-SAME: : !xegpu.tensor_desc<32x32xf32, #xegpu.layout<lane_layout = [1, 16], lane_data = [1, 1]>>45 // CHECK-SAME: -> vector<32x32xf32>46 %tdesc = xegpu.create_nd_tdesc %src[0, 0] : memref<256x128xf32>47 -> !xegpu.tensor_desc<256x128xf32, #xegpu.layout<sg_layout = [8, 4], sg_data = [32, 32], lane_layout = [1, 16], lane_data = [1, 1]>>48 %load = xegpu.load_nd %tdesc49 : !xegpu.tensor_desc<256x128xf32, #xegpu.layout<sg_layout = [8, 4], sg_data = [32, 32], lane_layout = [1, 16], lane_data = [1, 1]>>50 -> vector<256x128xf32>51 gpu.return52 }53 54 // CHECK-LABEL: store_nd55 // CHECK-SAME: %[[ARG_0:.*]]: memref<256x128xf32>56 gpu.func @store_nd(%src: memref<256x128xf32>) {57 // CHECK: %[[TDESC:.*]] = xegpu.create_nd_tdesc %[[ARG_0]][{{%.*}}, {{%.*}}] : memref<256x128xf32>58 // CHECK-SAME: -> !xegpu.tensor_desc<32x32xf32, #xegpu.layout<lane_layout = [1, 16], lane_data = [1, 1]>>59 // CHECK: %[[LOAD:.*]] = xegpu.load_nd %[[TDESC]]60 // CHECK-SAME: : !xegpu.tensor_desc<32x32xf32, #xegpu.layout<lane_layout = [1, 16], lane_data = [1, 1]>>61 // CHECK-SAME: -> vector<32x32xf32>62 // CHECK: xegpu.store_nd %[[LOAD]], %[[TDESC]]63 // CHECK-SAME: : vector<32x32xf32>, !xegpu.tensor_desc<32x32xf32, #xegpu.layout<lane_layout = [1, 16], lane_data = [1, 1]>>64 %tdesc = xegpu.create_nd_tdesc %src[0, 0] : memref<256x128xf32>65 -> !xegpu.tensor_desc<256x128xf32, #xegpu.layout<sg_layout = [8, 4], sg_data = [32, 32], lane_layout = [1, 16], lane_data = [1, 1]>>66 %load = xegpu.load_nd %tdesc67 : !xegpu.tensor_desc<256x128xf32, #xegpu.layout<sg_layout = [8, 4], sg_data = [32, 32], lane_layout = [1, 16], lane_data = [1, 1]>>68 -> vector<256x128xf32>69 xegpu.store_nd %load, %tdesc70 : vector<256x128xf32>, !xegpu.tensor_desc<256x128xf32, #xegpu.layout<sg_layout = [8, 4], sg_data = [32, 32], lane_layout = [1, 16], lane_data = [1, 1]>>71 gpu.return72 }73 74 // CHECK-LABEL: update_nd75 // CHECK-SAME: %[[ARG_0:.*]]: memref<256x128xf32>76 gpu.func @update_nd(%src: memref<256x128xf32>){77 // CHECK: %[[TDESC:.*]] = xegpu.create_nd_tdesc %[[ARG_0]][{{%.*}}, {{%.*}}] : memref<256x128xf32>78 // CHECK-SAME: -> !xegpu.tensor_desc<32x32xf32, #xegpu.layout<lane_layout = [1, 16], lane_data = [1, 1]>>79 // CHECK: %[[UPDATE:.*]] = xegpu.update_nd_offset %[[TDESC]], [0, 16]80 // CHECK-SAME: : !xegpu.tensor_desc<32x32xf32, #xegpu.layout<lane_layout = [1, 16], lane_data = [1, 1]>>81 %tdesc = xegpu.create_nd_tdesc %src[0, 0] : memref<256x128xf32>82 -> !xegpu.tensor_desc<256x128xf32, #xegpu.layout<sg_layout = [8, 4], sg_data = [32, 32], lane_layout = [1, 16], lane_data = [1, 1]>>83 %update = xegpu.update_nd_offset %tdesc, [0, 16]84 : !xegpu.tensor_desc<256x128xf32, #xegpu.layout<sg_layout = [8, 4], sg_data = [32, 32], lane_layout = [1, 16], lane_data = [1, 1]>>85 gpu.return86 }87 88 // CHECK-LABEL: dpas89 gpu.func @dpas(%a: memref<128x128xf16>, %b: memref<128x128xf16>) {90 %tdesc_a = xegpu.create_nd_tdesc %a[0, 0] : memref<128x128xf16>91 -> !xegpu.tensor_desc<128x128xf16, #xegpu.layout<sg_layout = [8, 8], sg_data = [16, 128], lane_layout = [1, 16], lane_data = [1, 1]>>92 %load_a = xegpu.load_nd %tdesc_a93 : !xegpu.tensor_desc<128x128xf16, #xegpu.layout<sg_layout = [8, 8], sg_data = [16, 128], lane_layout = [1, 16], lane_data = [1, 1]>>94 -> vector<128x128xf16>95 %tdesc_b = xegpu.create_nd_tdesc %b[0, 0] : memref<128x128xf16>96 -> !xegpu.tensor_desc<128x128xf16, #xegpu.layout<sg_layout = [8, 8], sg_data = [128, 16], lane_layout = [1, 16], lane_data = [2, 1]>>97 %load_b = xegpu.load_nd %tdesc_b98 : !xegpu.tensor_desc<128x128xf16, #xegpu.layout<sg_layout = [8, 8], sg_data = [128, 16], lane_layout = [1, 16], lane_data = [2, 1]>>99 -> vector<128x128xf16>100 // CHECK: %[[DPAS:.*]] = xegpu.dpas %{{.*}}, %{{.*}} {layout_result_0 = #xegpu.layout<lane_layout = [1, 16], lane_data = [1, 1]>} : vector<16x128xf16>, vector<128x16xf16> -> vector<16x16xf32>101 %dpas = xegpu.dpas %load_a, %load_b102 {layout_result_0 = #xegpu.layout<sg_layout = [8, 8], sg_data = [16, 16], lane_layout = [1, 16], lane_data = [1, 1]>}103 : vector<128x128xf16>, vector<128x128xf16> -> vector<128x128xf32>104 gpu.return105 }106 107 // CHECK-LABEL: dpas_no_sg_data108 gpu.func @dpas_no_sg_data(%a: memref<128x128xf16>, %b: memref<128x128xf16>) {109 %tdesc_a = xegpu.create_nd_tdesc %a[0, 0] : memref<128x128xf16>110 -> !xegpu.tensor_desc<128x128xf16, #xegpu.layout<sg_layout = [8, 8], lane_layout = [1, 16], lane_data = [1, 1],111 order = [1, 0]>>112 %load_a = xegpu.load_nd %tdesc_a113 : !xegpu.tensor_desc<128x128xf16, #xegpu.layout<sg_layout = [8, 8], lane_layout = [1, 16], lane_data = [1, 1],114 order = [1, 0]>>115 -> vector<128x128xf16>116 %tdesc_b = xegpu.create_nd_tdesc %b[0, 0] : memref<128x128xf16>117 -> !xegpu.tensor_desc<128x128xf16, #xegpu.layout<sg_layout = [8, 8], lane_layout = [1, 16], lane_data = [2, 1],118 order = [1, 0]>>119 %load_b = xegpu.load_nd %tdesc_b120 : !xegpu.tensor_desc<128x128xf16, #xegpu.layout<sg_layout = [8, 8], lane_layout = [1, 16], lane_data = [2, 1],121 order = [1, 0]>>122 -> vector<128x128xf16>123 // CHECK: %[[DPAS:.*]] = xegpu.dpas %{{.*}}, %{{.*}} {layout_result_0 = #xegpu.layout<lane_layout = [1, 16], lane_data = [1, 1], order = [1, 0]>} : vector<16x16xf16>, vector<16x16xf16> -> vector<16x16xf32>124 %dpas = xegpu.dpas %load_a, %load_b125 {layout_result_0 = #xegpu.layout<sg_layout = [8, 8], lane_layout = [1, 16], lane_data = [1, 1], order = [1, 0]>}126 : vector<128x128xf16>, vector<128x128xf16> -> vector<128x128xf32>127 gpu.return128 }129 130 // CHECK-LABEL: prefetch_nd_tdesc131 // CHECK-SAME: %[[ARG_0:.*]]: memref<256x128xf32>132 gpu.func @prefetch_nd_tdesc(%src: memref<256x128xf32>) {133 // CHECK: %[[TDESC:.*]] = xegpu.create_nd_tdesc %[[ARG_0]][{{%.*}}, {{%.*}}] : memref<256x128xf32>134 // CHECK-SAME: -> !xegpu.tensor_desc<32x32xf32, #xegpu.layout<lane_layout = [1, 16], lane_data = [1, 1]>>135 // CHECK: xegpu.prefetch_nd %[[TDESC]]136 // CHECK-SAME: : !xegpu.tensor_desc<32x32xf32, #xegpu.layout<lane_layout = [1, 16], lane_data = [1, 1]>>137 %tdesc = xegpu.create_nd_tdesc %src[0, 0] : memref<256x128xf32>138 -> !xegpu.tensor_desc<256x128xf32, #xegpu.layout<sg_layout = [8, 4], sg_data = [32, 32], lane_layout = [1, 16], lane_data = [1, 1]>>139 xegpu.prefetch_nd %tdesc140 : !xegpu.tensor_desc<256x128xf32, #xegpu.layout<sg_layout = [8, 4], sg_data = [32, 32], lane_layout = [1, 16], lane_data = [1, 1]>>141 gpu.return142 }143 144 // CHECK-LABEL: dpas_with_no_create_nd_desc145 gpu.func @dpas_with_no_create_nd_desc(%a: vector<256x128xf32>, %b: vector<128x256xf32>) {146 // CHECK-NOT: vector<32x32xf32>147 %dpas = xegpu.dpas %a, %b148 {layout = #xegpu.layout<sg_layout = [2, 2], sg_data = [12, 12], lane_layout = [2, 2], lane_data = [1, 1]>}149 : vector<256x128xf32>, vector<128x256xf32> -> vector<256x256xf32>150 gpu.return151 }152 153 // CHECK-LABEL: broadcast_dim1154 // CHECK-SAME: %[[ARG_0:.*]]: memref<256x1xf32>155 gpu.func @broadcast_dim1(%src: memref<256x1xf32>) {156 %tdesc = xegpu.create_nd_tdesc %src[0, 0] : memref<256x1xf32>157 -> !xegpu.tensor_desc<256x1xf32, #xegpu.layout<sg_layout = [8, 1], sg_data = [32, 1], lane_layout = [8, 1], lane_data = [1, 1]>>158 %load = xegpu.load_nd %tdesc159 : !xegpu.tensor_desc<256x1xf32, #xegpu.layout<sg_layout = [8, 1], sg_data = [32, 1], lane_layout = [8, 1], lane_data = [1, 1]>>160 -> vector<256x1xf32>161 // CHECK: vector.broadcast {{.*}} {layout_result_0 = #xegpu.layout<lane_layout = [8, 1], lane_data = [1, 1]>}162 // CHECK-SAME: : vector<32x1xf32> to vector<32x32xf32>163 %broadcast = vector.broadcast %load164 {layout_result_0 = #xegpu.layout<sg_layout = [8, 1], sg_data = [32, 32], lane_layout = [8, 1], lane_data = [1, 1]>}165 : vector<256x1xf32> to vector<256x32xf32>166 gpu.return167 }168 169 // CHECK-LABEL: broadcast_dim0170 // CHECK-SAME: %[[ARG_0:.*]]: memref<1x128xf32>171 gpu.func @broadcast_dim0(%src: memref<1x128xf32>) {172 %tdesc = xegpu.create_nd_tdesc %src[0, 0] : memref<1x128xf32>173 -> !xegpu.tensor_desc<1x128xf32, #xegpu.layout<sg_layout = [1, 4], sg_data = [1, 32], lane_layout = [1, 16], lane_data = [1, 1]>>174 %load = xegpu.load_nd %tdesc175 : !xegpu.tensor_desc<1x128xf32, #xegpu.layout<sg_layout = [1, 4], sg_data = [1, 32], lane_layout = [1, 16], lane_data = [1, 1]>>176 -> vector<1x128xf32>177 // CHECK: vector.broadcast {{.*}} {layout_result_0 = #xegpu.layout<lane_layout = [1, 16], lane_data = [1, 1]>}178 // CHECK-SAME: : vector<1x32xf32> to vector<32x32xf32>179 %broadcast = vector.broadcast %load180 {layout_result_0 = #xegpu.layout<sg_layout = [1, 4], sg_data = [32, 32], lane_layout = [1, 16], lane_data = [1, 1]>}181 : vector<1x128xf32> to vector<32x128xf32>182 gpu.return183 }184 185 gpu.func @scf_for(%arg0: memref<1024x1024xf16>, %arg1: memref<1024x1024xf16>, %arg2: memref<1024x1024xf32>) {186 // CHECK-DAG: %[[C0:.*]] = arith.constant 0 : index187 // CHECK-DAG: %[[C128:.*]] = arith.constant 128 : index188 // CHECK-DAG: %[[C1024:.*]] = arith.constant 1024 : index189 %c0 = arith.constant 0 : index190 %c128 = arith.constant 128 : index191 %c1024 = arith.constant 1024 : index192 %block_id_x = gpu.block_id x193 %block_id_y = gpu.block_id y194 %0 = arith.muli %block_id_x, %c128 : index195 %1 = arith.muli %block_id_y, %c128 : index196 %2 = xegpu.create_nd_tdesc %arg2[%0, %1] : memref<1024x1024xf32> -> !xegpu.tensor_desc<128x128xf32, #xegpu.layout<sg_layout = [8, 8], sg_data = [16, 16]>>197 %3 = xegpu.load_nd %2 : !xegpu.tensor_desc<128x128xf32, #xegpu.layout<sg_layout = [8, 8], sg_data = [16, 16]>> -> vector<128x128xf32>198 %4 = xegpu.create_nd_tdesc %arg0[%0, %c0] : memref<1024x1024xf16> -> !xegpu.tensor_desc<128x128xf16, #xegpu.layout<sg_layout = [8, 8], sg_data = [16, 128]>>199 %5 = xegpu.create_nd_tdesc %arg1[%c0, %1] : memref<1024x1024xf16> -> !xegpu.tensor_desc<128x128xf16, #xegpu.layout<sg_layout = [8, 8], sg_data = [128, 16]>>200 201 // CHECK: %[[SCF:.*]]:3 = scf.for %[[ARG3:.*]] = %[[C0]] to %[[C1024]] step %[[C128]]202 // CHECK-SAME: iter_args(%[[ARG4:.*]] = {{.*}}, %[[ARG5:.*]] = {{.*}}, %[[ARG6:.*]] = {{.*}}) ->203 // CHECK-SAME: (!xegpu.tensor_desc<16x128xf16>, !xegpu.tensor_desc<128x16xf16>, vector<16x16xf32>)204 // CHECK: %[[A:.*]] = xegpu.load_nd %[[ARG4]] : !xegpu.tensor_desc<16x128xf16> -> vector<16x128xf16>205 // CHECK: %[[B:.*]] = xegpu.load_nd %[[ARG5]] : !xegpu.tensor_desc<128x16xf16> -> vector<128x16xf16>206 // CHECK: %[[C:.*]] = xegpu.dpas %[[A]], %[[B]], %[[ARG6]] : vector<16x128xf16>, vector<128x16xf16>, vector<16x16xf32> -> vector<16x16xf32>207 // CHECK: %[[AT:.*]] = xegpu.update_nd_offset %[[ARG4]], [%[[C0]], %[[C128]]] : !xegpu.tensor_desc<16x128xf16>208 // CHECK: %[[BT:.*]] = xegpu.update_nd_offset %[[ARG5]], [%[[C128]], %[[C0]]] : !xegpu.tensor_desc<128x16xf16>209 // CHECK: scf.yield %[[AT]], %[[BT]], %[[C]] : !xegpu.tensor_desc<16x128xf16>, !xegpu.tensor_desc<128x16xf16>, vector<16x16xf32>210 %6:3 = scf.for %arg3 = %c0 to %c1024 step %c128 iter_args(%arg4 = %4, %arg5 = %5, %arg6 = %3)211 -> (!xegpu.tensor_desc<128x128xf16, #xegpu.layout<sg_layout = [8, 8], sg_data = [16, 128]>>,212 !xegpu.tensor_desc<128x128xf16, #xegpu.layout<sg_layout = [8, 8], sg_data = [128, 16]>>, vector<128x128xf32>) {213 %8 = xegpu.load_nd %arg4 : !xegpu.tensor_desc<128x128xf16, #xegpu.layout<sg_layout = [8, 8], sg_data = [16, 128]>> -> vector<128x128xf16>214 %9 = xegpu.load_nd %arg5 : !xegpu.tensor_desc<128x128xf16, #xegpu.layout<sg_layout = [8, 8], sg_data = [128, 16]>> -> vector<128x128xf16>215 %10 = xegpu.dpas %8, %9, %arg6 {layout_result_0 = #xegpu.layout<sg_layout = [8, 8], sg_data = [16, 16]>}216 : vector<128x128xf16>, vector<128x128xf16>, vector<128x128xf32> -> vector<128x128xf32>217 %11 = xegpu.update_nd_offset %arg4, [%c0, %c128] : !xegpu.tensor_desc<128x128xf16, #xegpu.layout<sg_layout = [8, 8], sg_data = [16, 128]>>218 %12 = xegpu.update_nd_offset %arg5, [%c128, %c0] : !xegpu.tensor_desc<128x128xf16, #xegpu.layout<sg_layout = [8, 8], sg_data = [128, 16]>>219 scf.yield %11, %12, %10 : !xegpu.tensor_desc<128x128xf16, #xegpu.layout<sg_layout = [8, 8], sg_data = [16, 128]>>,220 !xegpu.tensor_desc<128x128xf16, #xegpu.layout<sg_layout = [8, 8], sg_data = [128, 16]>>, vector<128x128xf32>221 }222 %7 = xegpu.create_nd_tdesc %arg2[%0, %1] : memref<1024x1024xf32>223 -> !xegpu.tensor_desc<128x128xf32, #xegpu.layout<sg_layout = [8, 8], sg_data = [16, 16]>>224 xegpu.store_nd %6#2, %7 : vector<128x128xf32>, !xegpu.tensor_desc<128x128xf32, #xegpu.layout<sg_layout = [8, 8], sg_data = [16, 16]>>225 gpu.return226 }227 228 gpu.func @scf_while_and_condition(%arg0: memref<1024xf32>, %arg1: memref<1024xf32>) {229 %c1_i32 = arith.constant 1 : i32230 %c10_i32 = arith.constant 10 : i32231 %c0_i32 = arith.constant 0 : i32232 %0 = xegpu.create_nd_tdesc %arg0[0] : memref<1024xf32> -> !xegpu.tensor_desc<256xf32, #xegpu.layout<sg_layout = [16], sg_data = [16]>>233 %1 = xegpu.load_nd %0 : !xegpu.tensor_desc<256xf32, #xegpu.layout<sg_layout = [16], sg_data = [16]>> -> vector<256xf32>234 %2 = xegpu.create_nd_tdesc %arg1[0] : memref<1024xf32> -> !xegpu.tensor_desc<256xf32, #xegpu.layout<sg_layout = [16], sg_data = [16]>>235 236 // CHECK: scf.while {{.*}} : (vector<16xf32>, i32) -> (vector<16xf32>, i32)237 %3:2 = scf.while (%arg2 = %1, %arg3 = %c0_i32) : (vector<256xf32>, i32) -> (vector<256xf32>, i32) {238 %4 = arith.cmpi slt, %arg3, %c10_i32 : i32239 // CHECK: scf.condition{{.*}} : vector<16xf32>, i32240 scf.condition(%4) %arg2, %arg3 : vector<256xf32>, i32241 } do {242 // CHECK: (%[[ARG2:.*]]: vector<16xf32>, %[[ARG3:.*]]: i32)243 ^bb0(%arg2: vector<256xf32>, %arg3: i32):244 xegpu.store_nd %arg2, %2 : vector<256xf32>, !xegpu.tensor_desc<256xf32, #xegpu.layout<sg_layout = [16], sg_data = [16]>>245 %4 = arith.addi %arg3, %c1_i32 : i32246 %5 = xegpu.update_nd_offset %0, [256] : !xegpu.tensor_desc<256xf32, #xegpu.layout<sg_layout = [16], sg_data = [16]>>247 %6 = xegpu.load_nd %5 : !xegpu.tensor_desc<256xf32, #xegpu.layout<sg_layout = [16], sg_data = [16]>> -> vector<256xf32>248 scf.yield %6, %4 : vector<256xf32>, i32249 }250 gpu.return251 }252 253 gpu.func @scf_if(%arg0: memref<1024xf32>, %arg1: memref<1024xf32>) {254 %c10 = arith.constant 10 : index255 %id = gpu.subgroup_id : index256 257 %0 = xegpu.create_nd_tdesc %arg0[0] : memref<1024xf32> -> !xegpu.tensor_desc<256xf32, #xegpu.layout<sg_layout = [16], sg_data = [16]>>258 %1 = xegpu.create_nd_tdesc %arg1[0] : memref<1024xf32> -> !xegpu.tensor_desc<256xf32, #xegpu.layout<sg_layout = [16], sg_data = [16]>>259 260 %4 = arith.cmpi eq, %id, %c10 : index261 // CHECK-LABEL: scf.if262 // CHECK-SAME: (vector<16xf32>)263 %5 = scf.if %4 -> (vector<256xf32>) {264 // CHECK-LABEL: xegpu.load_nd265 // CHECK-SAME: !xegpu.tensor_desc<16xf32> -> vector<16xf32>266 %2 = xegpu.load_nd %0 : !xegpu.tensor_desc<256xf32, #xegpu.layout<sg_layout = [16], sg_data = [16]>> -> vector<256xf32>267 // CHECK-LABEL: scf.yield268 // CHECK-SAME: vector<16xf32>269 scf.yield %2 : vector<256xf32>270 } else {271 // CHECK-LABEL: xegpu.load_nd272 // CHECK-SAME: !xegpu.tensor_desc<16xf32> -> vector<16xf32>273 %3 = xegpu.load_nd %1 : !xegpu.tensor_desc<256xf32, #xegpu.layout<sg_layout = [16], sg_data = [16]>> -> vector<256xf32>274 // CHECK-LABEL: scf.yield275 // CHECK-SAME: vector<16xf32>276 scf.yield %3 : vector<256xf32>277 } {layout_result_0 = #xegpu.layout<sg_layout = [16], sg_data = [16]>}278 xegpu.store_nd %5, %0 : vector<256xf32>, !xegpu.tensor_desc<256xf32, #xegpu.layout<sg_layout = [16], sg_data = [16]>>279 gpu.return280 }281 282 gpu.func @scf_if_tensor_desc(%arg0: memref<1024xf32>, %arg1: memref<1024xf32>) {283 %c10 = arith.constant 10 : index284 %id = gpu.subgroup_id : index285 286 %t = xegpu.create_nd_tdesc %arg0[0] : memref<1024xf32> -> !xegpu.tensor_desc<256xf32, #xegpu.layout<sg_layout = [16], sg_data = [16]>>287 %d = xegpu.load_nd %t : !xegpu.tensor_desc<256xf32, #xegpu.layout<sg_layout = [16], sg_data = [16]>> -> vector<256xf32>288 289 %0 = arith.cmpi eq, %id, %c10 : index290 // CHECK-LABEL: scf.if291 // CHECK-SAME: (!xegpu.tensor_desc<16xf32>)292 %1 = scf.if %0 -> (!xegpu.tensor_desc<256xf32, #xegpu.layout<sg_layout = [16], sg_data = [16]>>) {293 // CHECK-LABEL: xegpu.create_nd_tdesc294 // CHECK-SAME: memref<1024xf32> -> !xegpu.tensor_desc<16xf32>295 %2 = xegpu.create_nd_tdesc %arg0[0] : memref<1024xf32> -> !xegpu.tensor_desc<256xf32, #xegpu.layout<sg_layout = [16], sg_data = [16]>>296 // CHECK-LABEL: scf.yield297 // CHECK-SAME: !xegpu.tensor_desc<16xf32>298 scf.yield %2 : !xegpu.tensor_desc<256xf32, #xegpu.layout<sg_layout = [16], sg_data = [16]>>299 } else {300 // CHECK-LABEL: xegpu.create_nd_tdesc301 // CHECK-SAME: memref<1024xf32> -> !xegpu.tensor_desc<16xf32>302 %3 = xegpu.create_nd_tdesc %arg1[0] : memref<1024xf32> -> !xegpu.tensor_desc<256xf32, #xegpu.layout<sg_layout = [16], sg_data = [16]>>303 // CHECK-LABEL: scf.yield304 // CHECK-SAME: !xegpu.tensor_desc<16xf32>305 scf.yield %3 : !xegpu.tensor_desc<256xf32, #xegpu.layout<sg_layout = [16], sg_data = [16]>>306 }307 xegpu.store_nd %d, %1 : vector<256xf32>, !xegpu.tensor_desc<256xf32, #xegpu.layout<sg_layout = [16], sg_data = [16]>>308 gpu.return309 }310 311 // CHECK-LABEL: @subgroup_id_range312 gpu.func @subgroup_id_range(%src: memref<256x128xf32>, %src1: memref<128x256xf32>, %src2: memref<128x64xf32>) {313 %sg_id = gpu.subgroup_id : index314 %c0 = arith.constant 0 : index315 %c1 = arith.constant 1 : index316 %c2 = arith.constant 2 : index317 %c31 = arith.constant 31 : index318 %c3 = arith.constant 3 : index319 %cond1 = arith.cmpi sge, %sg_id, %c0 : index320 %cond2 = arith.cmpi slt, %sg_id, %c1 : index321 %cond = arith.andi %cond1, %cond2 : i1322 scf.if %cond {323 // CHECK-NOT: index.sub324 %tdesc = xegpu.create_nd_tdesc %src[0, 0] : memref<256x128xf32>325 -> !xegpu.tensor_desc<256x128xf32, #xegpu.layout<sg_layout = [8, 4], sg_data = [32, 32], lane_layout = [8, 4], lane_data = [1, 1]>>326 %load = xegpu.load_nd %tdesc327 : !xegpu.tensor_desc<256x128xf32, #xegpu.layout<sg_layout = [8, 4], sg_data = [32, 32], lane_layout = [8, 4], lane_data = [1, 1]>>328 -> vector<256x128xf32>329 } {sg_id_range = #xegpu.range<[0, 32]>}330 %cond3 = arith.cmpi sge, %sg_id, %c2 : index331 %cond4 = arith.cmpi slt, %sg_id, %c31 : index332 %cond5 = arith.andi %cond3, %cond4 : i1333 scf.if %cond5 {334 // CHECK: %[[SGID:.*]] = gpu.subgroup_id : index335 // CHECK: %[[C2:.*]] = arith.constant 2 : index336 // CHECK: %[[SUB:.*]] = index.sub %{{.*}}, %[[C2]]337 %tdesc = xegpu.create_nd_tdesc %src2[0, 0] : memref<128x64xf32>338 -> !xegpu.tensor_desc<128x64xf32, #xegpu.layout<sg_layout = [4, 4], sg_data = [32, 16], lane_layout = [8, 4], lane_data = [1, 1]>>339 %load = xegpu.load_nd %tdesc340 : !xegpu.tensor_desc<128x64xf32, #xegpu.layout<sg_layout = [4, 4], sg_data = [32, 16], lane_layout = [8, 4], lane_data = [1, 1]>>341 -> vector<128x64xf32>342 %exp = math.exp %load {layout_result_0 = #xegpu.layout<sg_layout = [4, 4], sg_data = [32, 16], lane_layout = [8, 4], lane_data = [1, 1]>} : vector<128x64xf32>343 }{sg_id_range = #xegpu.range<[2, 18]>}344 gpu.return345 }346 347 // CHECK-LABEL: @subgroup_id_range_nested_if348 gpu.func @subgroup_id_range_nested_if(%src: memref<256x128xf32>, %src1: memref<128x64xf32>) {349 %sg_id = gpu.subgroup_id : index350 %c1 = arith.constant 1 : i1351 %c3 = arith.constant 3 : index352 %c32 = arith.constant 32 : index353 %tdesc = xegpu.create_nd_tdesc %src[0, 0] : memref<256x128xf32>354 -> !xegpu.tensor_desc<256x128xf32, #xegpu.layout<sg_layout = [8, 4], sg_data = [32, 32], lane_layout = [8, 4], lane_data = [1, 1]>>355 %load = xegpu.load_nd %tdesc356 : !xegpu.tensor_desc<256x128xf32, #xegpu.layout<sg_layout = [8, 4], sg_data = [32, 32], lane_layout = [8, 4], lane_data = [1, 1]>>357 -> vector<256x128xf32>358 %cond1 = arith.cmpi sge, %sg_id, %c3 : index359 %cond2 = arith.cmpi slt, %sg_id, %c32 : index360 %cond = arith.andi %cond1, %cond2 : i1361 scf.if %c1 {362 scf.if %cond {363 // CHECK: %[[SGID:.*]] = gpu.subgroup_id : index364 // CHECK: %[[C3:.*]] = arith.constant 3 : index365 // CHECK: %[[SUB:.*]] = index.sub %{{.*}}, %[[C3]]366 %td = xegpu.create_nd_tdesc %src1[0, 0] : memref<128x64xf32>367 -> !xegpu.tensor_desc<128x64xf32, #xegpu.layout<sg_layout = [4, 4], sg_data = [32, 16], lane_layout = [8, 4], lane_data = [1, 1]>>368 %ld = xegpu.load_nd %td369 : !xegpu.tensor_desc<128x64xf32, #xegpu.layout<sg_layout = [4, 4], sg_data = [32, 16], lane_layout = [8, 4], lane_data = [1, 1]>>370 -> vector<128x64xf32>371 %exp = math.exp %ld {layout_result_0 = #xegpu.layout<sg_layout = [4, 4], sg_data = [32, 16], lane_layout = [8, 4], lane_data = [1, 1]>} : vector<128x64xf32>372 }373 } {sg_id_range = #xegpu.range<[3, 19]>}374 gpu.return375 }376 377 // CHECK-LABEL: distribute_constant378 gpu.func @distribute_constant() {379 // CHECK: arith.constant dense<1.000000e+00> : vector<32x32xf32>380 %cst = arith.constant {layout_result_0 = #xegpu.layout<sg_layout = [8, 4], sg_data = [32, 32]>} dense<1.0> : vector<256x128xf32>381 gpu.return382 }383}384