603 lines · plain
1// RUN: mlir-opt --xegpu-wg-to-sg-distribute -split-input-file %s | FileCheck %s2 3gpu.module @test_distribution {4 // CHECK-LABEL: create_nd_tdesc_no_offset5 // CHECK-SAME: %[[ARG_0:.*]]: memref<256x128xf32>6 gpu.func @create_nd_tdesc_no_offset(%src: memref<256x128xf32>) {7 // CHECK: xegpu.create_nd_tdesc %[[ARG_0]] : memref<256x128xf32>8 // CHECK-SAME: -> !xegpu.tensor_desc<32x32xf32, #xegpu.layout<lane_layout = [1, 16], lane_data = [1, 1]>>9 %tdesc = xegpu.create_nd_tdesc %src : memref<256x128xf32>10 -> !xegpu.tensor_desc<256x128xf32, #xegpu.layout<sg_layout = [8, 4], sg_data = [32, 32], lane_layout = [1, 16], lane_data = [1, 1]>>11 gpu.return12 }13 14 // CHECK-LABEL: create_nd_tdesc_with_ptr15 // CHECK-SAME: %[[ARG_0:.*]]: ui6416 gpu.func @create_nd_tdesc_with_ptr(%src: ui64, %w : index, %h : index, %x : index, %y : index) {17 // CHECK: xegpu.create_nd_tdesc %[[ARG_0]], shape : [{{.*}}, {{.*}}], strides : [{{.*}}, {{.*}}] : ui6418 // CHECK-SAME: -> !xegpu.tensor_desc<32x32xf32, #xegpu.layout<lane_layout = [1, 16], lane_data = [1, 1]>>19 %c1 = arith.constant 1 : index20 %tdesc = xegpu.create_nd_tdesc %src, shape:[%h, %w], strides: [%w, %c1] : ui6421 -> !xegpu.tensor_desc<256x128xf32, #xegpu.layout<sg_layout = [8, 4], sg_data = [32, 32], lane_layout = [1, 16], lane_data = [1, 1]>>22 gpu.return23 }24 25 // CHECK-LABEL: load_nd_tdesc_with_offset26 gpu.func @load_nd_tdesc_with_offset(%src: memref<256x128xf32>) {27 //CHECK: %[[TDESC:.*]] = xegpu.create_nd_tdesc %{{.*}} : memref<256x128xf32> -> !xegpu.tensor_desc<32x32xf32, #xegpu.layout<lane_layout = [1, 16], lane_data = [1, 1]>>28 //CHECK-DAG: %[[SGID:.*]] = gpu.subgroup_id : index29 //CHECK-DAG: %[[C4:.*]] = arith.constant 4 : index30 //CHECK-DAG: %[[SGIDX:.*]] = index.remu %[[SGID]], %[[C4]]31 //CHECK-DAG: %[[SGIDY_TMP:.*]] = index.divu %[[SGID]], %[[C4]]32 //CHECK-DAG: %[[C8:.*]] = arith.constant 8 : index33 //CHECK-DAG: %[[SGIDY:.*]] = index.remu %[[SGIDY_TMP]], %[[C8]]34 //CHECK-DAG: %[[C32:.*]] = arith.constant 32 : index35 //CHECK-DAG: %[[L_OFF_Y:.*]] = index.mul %[[SGIDY]], %[[C32]]36 //CHECK-DAG: %[[L_OFF_X:.*]] = index.mul %[[SGIDX]], %[[C32]]37 //CHECK-DAG: %[[C256:.*]] = arith.constant 256 : index38 //CHECK-DAG: %[[OFF_Y:.*]] = index.remu %[[L_OFF_Y]], %[[C256]]39 //CHECK-DAG: %[[C128:.*]] = arith.constant 128 : index40 //CHECK-DAG: %[[OFF_X:.*]] = index.remu %[[L_OFF_X]], %[[C128]]41 //CHECK-DAG: %[[LOAD:.*]] = xegpu.load_nd %[[TDESC]][{{%.*}}, {{%.*}}] : !xegpu.tensor_desc<32x32xf32, #xegpu.layout<lane_layout = [1, 16], lane_data = [1, 1]>> -> vector<32x32xf32>42 %tdesc = xegpu.create_nd_tdesc %src : memref<256x128xf32>43 -> !xegpu.tensor_desc<256x128xf32, #xegpu.layout<sg_layout = [8, 4], sg_data = [32, 32], lane_layout = [1, 16], lane_data = [1, 1]>>44 %load = xegpu.load_nd %tdesc[0, 0]45 : !xegpu.tensor_desc<256x128xf32, #xegpu.layout<sg_layout = [8, 4], sg_data = [32, 32], lane_layout = [1, 16], lane_data = [1, 1]>>46 -> vector<256x128xf32>47 gpu.return48 }49 50 // CHECK-LABEL: store_nd_with_offsets51 // CHECK-SAME: %[[ARG_0:.*]]: memref<256x128xf32>52 gpu.func @store_nd_with_offsets(%src: memref<256x128xf32>) {53 //CHECK: xegpu.store_nd %{{.*}}, {{%.*}}[{{%.*}}, {{%.*}}] : vector<32x32xf32>, !xegpu.tensor_desc<32x32xf32, #xegpu.layout<lane_layout = [1, 16], lane_data = [1, 1]>>54 %tdesc = xegpu.create_nd_tdesc %src: memref<256x128xf32>55 -> !xegpu.tensor_desc<256x128xf32, #xegpu.layout<sg_layout = [8, 4], sg_data = [32, 32], lane_layout = [1, 16], lane_data = [1, 1]>>56 %load = xegpu.load_nd %tdesc[0, 0]57 : !xegpu.tensor_desc<256x128xf32, #xegpu.layout<sg_layout = [8, 4], sg_data = [32, 32], lane_layout = [1, 16], lane_data = [1, 1]>>58 -> vector<256x128xf32>59 xegpu.store_nd %load, %tdesc[0, 0]60 : vector<256x128xf32>, !xegpu.tensor_desc<256x128xf32, #xegpu.layout<sg_layout = [8, 4], sg_data = [32, 32], lane_layout = [1, 16], lane_data = [1, 1]>>61 gpu.return62}63 64 // CHECK-LABEL: prefetch_nd_tdesc_with_offset65 // CHECK-SAME: %[[ARG_0:.*]]: memref<256x128xf32>66 gpu.func @prefetch_nd_tdesc_with_offset(%src: memref<256x128xf32>) {67 //CHECK: xegpu.prefetch_nd %{{.*}}[{{%.*}}, {{%.*}}] : !xegpu.tensor_desc<32x32xf32, #xegpu.layout<lane_layout = [1, 16], lane_data = [1, 1]>>68 %cst0 = arith.constant 0 : index69 %tdesc = xegpu.create_nd_tdesc %src : memref<256x128xf32>70 -> !xegpu.tensor_desc<256x128xf32, #xegpu.layout<sg_layout = [8, 4], sg_data = [32, 32], lane_layout = [1, 16], lane_data = [1, 1]>>71 xegpu.prefetch_nd %tdesc[%cst0, %cst0]72 : !xegpu.tensor_desc<256x128xf32, #xegpu.layout<sg_layout = [8, 4], sg_data = [32, 32], lane_layout = [1, 16], lane_data = [1, 1]>>73 gpu.return74 }75 76 // CHECK-LABEL: dpas77 gpu.func @dpas(%a: memref<128x128xf16>, %b: memref<128x128xf16>) {78 // CHECK: %[[DPAS:.*]] = xegpu.dpas %{{.*}}, %{{.*}} {layout_result_0 = #xegpu.layout<lane_layout = [1, 16], lane_data = [1, 1]>} : vector<16x128xf16>, vector<128x16xf16> -> vector<16x16xf32>79 %tdesc_a = xegpu.create_nd_tdesc %a : memref<128x128xf16>80 -> !xegpu.tensor_desc<128x128xf16, #xegpu.layout<sg_layout = [8, 8], sg_data = [16, 128], lane_layout = [1, 16], lane_data = [1, 1]>>81 %load_a = xegpu.load_nd %tdesc_a[0, 0]82 : !xegpu.tensor_desc<128x128xf16, #xegpu.layout<sg_layout = [8, 8], sg_data = [16, 128], lane_layout = [1, 16], lane_data = [1, 1]>>83 -> vector<128x128xf16>84 %tdesc_b = xegpu.create_nd_tdesc %b : memref<128x128xf16>85 -> !xegpu.tensor_desc<128x128xf16, #xegpu.layout<sg_layout = [8, 8], sg_data = [128, 16], lane_layout = [1, 16], lane_data = [2, 1]>>86 %load_b = xegpu.load_nd %tdesc_b[0, 0]87 : !xegpu.tensor_desc<128x128xf16, #xegpu.layout<sg_layout = [8, 8], sg_data = [128, 16], lane_layout = [1, 16], lane_data = [2, 1]>>88 -> vector<128x128xf16>89 %dpas = xegpu.dpas %load_a, %load_b90 {layout_result_0 = #xegpu.layout<sg_layout = [8, 8], sg_data = [16, 16], lane_layout = [1, 16], lane_data = [1, 1]>}91 : vector<128x128xf16>, vector<128x128xf16> -> vector<128x128xf32>92 gpu.return93 }94 95 // CHECK-LABEL: dpas_no_sg_data96 gpu.func @dpas_no_sg_data(%a: memref<128x128xf16>, %b: memref<128x128xf16>) {97 // 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>98 %tdesc_a = xegpu.create_nd_tdesc %a : memref<128x128xf16>99 -> !xegpu.tensor_desc<128x128xf16, #xegpu.layout<sg_layout = [8, 8], lane_layout = [1, 16], lane_data = [1, 1],100 order = [1, 0]>>101 %load_a = xegpu.load_nd %tdesc_a[0, 0]102 : !xegpu.tensor_desc<128x128xf16, #xegpu.layout<sg_layout = [8, 8], lane_layout = [1, 16], lane_data = [1, 1],103 order = [1, 0]>>104 -> vector<128x128xf16>105 %tdesc_b = xegpu.create_nd_tdesc %b : memref<128x128xf16>106 -> !xegpu.tensor_desc<128x128xf16, #xegpu.layout<sg_layout = [8, 8], lane_layout = [1, 16], lane_data = [2, 1],107 order = [1, 0]>>108 %load_b = xegpu.load_nd %tdesc_b[0, 0]109 : !xegpu.tensor_desc<128x128xf16, #xegpu.layout<sg_layout = [8, 8], lane_layout = [1, 16], lane_data = [2, 1],110 order = [1, 0]>>111 -> vector<128x128xf16>112 %dpas = xegpu.dpas %load_a, %load_b113 {layout_result_0 = #xegpu.layout<sg_layout = [8, 8], lane_layout = [1, 16], lane_data = [1, 1], order = [1, 0]>}114 : vector<128x128xf16>, vector<128x128xf16> -> vector<128x128xf32>115 gpu.return116 }117 118 // CHECK-LABEL: dpas_with_no_create_nd_desc119 gpu.func @dpas_with_no_create_nd_desc(%a: vector<256x128xf32>, %b: vector<128x256xf32>) {120 // CHECK-NOT: vector<32x32xf32>121 %dpas = xegpu.dpas %a, %b122 {layout = #xegpu.layout<sg_layout = [8, 8], sg_data = [16, 16], lane_layout = [1, 16], lane_data = [1, 1]>}123 : vector<256x128xf32>, vector<128x256xf32> -> vector<256x256xf32>124 gpu.return125 }126 127 // CHECK-LABEL: broadcast_dim1128 // CHECK-SAME: %[[ARG_0:.*]]: memref<256x1xf32>129 gpu.func @broadcast_dim1(%src: memref<256x1xf32>) {130 %tdesc = xegpu.create_nd_tdesc %src : memref<256x1xf32>131 -> !xegpu.tensor_desc<256x1xf32, #xegpu.layout<sg_layout = [8, 1], sg_data = [32, 1], lane_layout = [8, 1], lane_data = [1, 1]>>132 %load = xegpu.load_nd %tdesc[0, 0]133 : !xegpu.tensor_desc<256x1xf32, #xegpu.layout<sg_layout = [8, 1], sg_data = [32, 1], lane_layout = [8, 1], lane_data = [1, 1]>>134 -> vector<256x1xf32>135 // CHECK: vector.broadcast {{.*}} {layout_result_0 = #xegpu.layout<lane_layout = [8, 1], lane_data = [1, 1]>}136 // CHECK-SAME: : vector<32x1xf32> to vector<32x32xf32>137 %broadcast = vector.broadcast %load138 {layout_result_0 = #xegpu.layout<sg_layout = [8, 1], sg_data = [32, 32], lane_layout = [8, 1], lane_data = [1, 1]>}139 : vector<256x1xf32> to vector<256x32xf32>140 gpu.return141 }142 143 // CHECK-LABEL: broadcast_dim0144 // CHECK-SAME: %[[ARG_0:.*]]: memref<1x128xf32>145 gpu.func @broadcast_dim0(%src: memref<1x128xf32>) {146 %tdesc = xegpu.create_nd_tdesc %src : memref<1x128xf32>147 -> !xegpu.tensor_desc<1x128xf32, #xegpu.layout<sg_layout = [1, 4], sg_data = [1, 32], lane_layout = [1, 16], lane_data = [1, 1]>>148 %load = xegpu.load_nd %tdesc[0, 0]149 : !xegpu.tensor_desc<1x128xf32, #xegpu.layout<sg_layout = [1, 4], sg_data = [1, 32], lane_layout = [1, 16], lane_data = [1, 1]>>150 -> vector<1x128xf32>151 // CHECK: vector.broadcast {{.*}} {layout_result_0 = #xegpu.layout<lane_layout = [1, 16], lane_data = [1, 1]>}152 // CHECK-SAME: : vector<1x32xf32> to vector<32x32xf32>153 %broadcast = vector.broadcast %load154 {layout_result_0 = #xegpu.layout<sg_layout = [1, 4], sg_data = [32, 32], lane_layout = [1, 16], lane_data = [1, 1]>}155 : vector<1x128xf32> to vector<32x128xf32>156 gpu.return157 }158 159 // CHECK-LABEL: gemm_with_load_store_offset160 // CHECK-SAME: %[[ARG_0:.*]]: memref<1024x1024xf16>, %[[ARG_1:.*]]: memref<1024x1024xf16>, %[[ARG_2:.*]]: memref<1024x1024xf32>161 gpu.func @gemm_with_load_store_offset(%arg0: memref<1024x1024xf16>, %arg1: memref<1024x1024xf16>, %arg2: memref<1024x1024xf32>) {162 //CHECK: [[c0:%.+]] = arith.constant 0 : index163 //CHECK: [[c128:%.+]] = arith.constant 128 : index164 //CHECK: [[c1024:%.+]] = arith.constant 1024 : index165 %c0 = arith.constant 0 : index166 %c128 = arith.constant 128 : index167 %c1024 = arith.constant 1024 : index168 %block_id_x = gpu.block_id x169 %block_id_y = gpu.block_id y170 %0 = arith.muli %block_id_x, %c128 : index171 %1 = arith.muli %block_id_y, %c128 : index172 %2 = xegpu.create_nd_tdesc %arg2 : memref<1024x1024xf32> -> !xegpu.tensor_desc<128x128xf32, #xegpu.layout<sg_layout = [8, 8], sg_data = [16, 16]>>173 // CHECK: [[DESC_A:%.+]] = xegpu.create_nd_tdesc %[[ARG_0]] : memref<1024x1024xf16> -> !xegpu.tensor_desc<16x128xf16>174 // CHECK: [[DESC_B:%.+]] = xegpu.create_nd_tdesc %[[ARG_1]] : memref<1024x1024xf16> -> !xegpu.tensor_desc<128x16xf16>175 %3 = xegpu.create_nd_tdesc %arg0 : memref<1024x1024xf16> -> !xegpu.tensor_desc<128x128xf16, #xegpu.layout<sg_layout = [8, 8], sg_data = [16, 128]>>176 %4 = xegpu.create_nd_tdesc %arg1 : memref<1024x1024xf16> -> !xegpu.tensor_desc<128x128xf16, #xegpu.layout<sg_layout = [8, 8], sg_data = [128, 16]>>177 // load_nd with offset178 %5 = xegpu.load_nd %2[%0, %1] : !xegpu.tensor_desc<128x128xf32, #xegpu.layout<sg_layout = [8, 8], sg_data = [16, 16]>> -> vector<128x128xf32>179 %6 = xegpu.load_nd %3[%0, %c0] : !xegpu.tensor_desc<128x128xf16, #xegpu.layout<sg_layout = [8, 8], sg_data = [16, 128]>> -> vector<128x128xf16>180 %7 = xegpu.load_nd %4[%c0, %1] : !xegpu.tensor_desc<128x128xf16, #xegpu.layout<sg_layout = [8, 8], sg_data = [128, 16]>> -> vector<128x128xf16>181 // scf.for loop182 // CHECK: [[scf:%.+]]:3 = scf.for [[arg3:%.+]] = [[c0]] to [[c1024]] step [[c128]]183 // CHECK-SAME: iter_args([[arg4:%.+]] = {{.*}}, [[arg5:%.+]] = {{.*}}, [[arg6:%.+]] = {{.*}}) ->184 // CHECK-SAME: (vector<16x128xf16>, vector<128x16xf16>, vector<16x16xf32>)185 // CHECK: [[c:%.+]] = xegpu.dpas [[arg4]], [[arg5]], [[arg6]] : vector<16x128xf16>, vector<128x16xf16>, vector<16x16xf32> -> vector<16x16xf32>186 // CHECK: [[a:%.+]] = xegpu.load_nd [[DESC_A]][{{%.*}}, {{%.*}}] : !xegpu.tensor_desc<16x128xf16> -> vector<16x128xf16>187 // CHECK: [[b:%.+]] = xegpu.load_nd [[DESC_B]][{{%.*}}, {{%.*}}] : !xegpu.tensor_desc<128x16xf16> -> vector<128x16xf16>188 // CHECK: scf.yield [[a]], [[b]], [[c]] : vector<16x128xf16>, vector<128x16xf16>, vector<16x16xf32>189 %8:3 = scf.for %arg3 = %c0 to %c1024 step %c128 iter_args(%arg4 = %6, %arg5 = %7, %arg6 = %5)190 -> (vector<128x128xf16>, vector<128x128xf16>, vector<128x128xf32>) {191 // load_nd with offset inside loop192 %9 = xegpu.dpas %arg4, %arg5, %arg6 {layout_result_0 = #xegpu.layout<sg_layout = [8, 8], sg_data = [16, 16]>}193 : vector<128x128xf16>, vector<128x128xf16>, vector<128x128xf32> -> vector<128x128xf32>194 %10 = xegpu.load_nd %3[%arg3, %c0] : !xegpu.tensor_desc<128x128xf16, #xegpu.layout<sg_layout = [8, 8], sg_data = [16, 128]>> -> vector<128x128xf16>195 %11 = xegpu.load_nd %4[%c0, %arg3] : !xegpu.tensor_desc<128x128xf16, #xegpu.layout<sg_layout = [8, 8], sg_data = [128, 16]>> -> vector<128x128xf16>196 scf.yield %10, %11, %9 : vector<128x128xf16>, vector<128x128xf16>, vector<128x128xf32>197 }198 // store_nd with offset199 xegpu.store_nd %8#2, %2[%0, %1] : vector<128x128xf32>, !xegpu.tensor_desc<128x128xf32, #xegpu.layout<sg_layout = [8, 8], sg_data = [16, 16]>>200 gpu.return201 }202 203 // CHECK-LABEL: @subgroup_id_range204 gpu.func @subgroup_id_range(%src: memref<256x128xf32>, %src1: memref<128x256xf32>, %src2: memref<128x64xf32>) {205 %sg_id = gpu.subgroup_id : index206 %c0 = arith.constant 0 : index207 %c1 = arith.constant 1 : index208 %c2 = arith.constant 2 : index209 %c31 = arith.constant 31 : index210 %c3 = arith.constant 3 : index211 %cond1 = arith.cmpi sge, %sg_id, %c0 : index212 %cond2 = arith.cmpi slt, %sg_id, %c1 : index213 %cond = arith.andi %cond1, %cond2 : i1214 scf.if %cond {215 // CHECK-NOT: index.sub216 %tdesc = xegpu.create_nd_tdesc %src : memref<256x128xf32>217 -> !xegpu.tensor_desc<256x128xf32, #xegpu.layout<sg_layout = [8, 4], sg_data = [32, 32], lane_layout = [8, 4], lane_data = [1, 1]>>218 %load = xegpu.load_nd %tdesc[0, 0]219 : !xegpu.tensor_desc<256x128xf32, #xegpu.layout<sg_layout = [8, 4], sg_data = [32, 32], lane_layout = [8, 4], lane_data = [1, 1]>>220 -> vector<256x128xf32>221 } {sg_id_range = #xegpu.range<[0, 32]>}222 %cond3 = arith.cmpi sge, %sg_id, %c2 : index223 %cond4 = arith.cmpi slt, %sg_id, %c31 : index224 %cond5 = arith.andi %cond3, %cond4 : i1225 scf.if %cond5 {226 // CHECK: %[[SGID:.*]] = gpu.subgroup_id : index227 // CHECK: %[[C2:.*]] = arith.constant 2 : index228 // CHECK: %[[SUB:.*]] = index.sub %{{.*}}, %[[C2]]229 %tdesc = xegpu.create_nd_tdesc %src2 : memref<128x64xf32>230 -> !xegpu.tensor_desc<128x64xf32, #xegpu.layout<sg_layout = [4, 4], sg_data = [32, 16], lane_layout = [8, 4], lane_data = [1, 1]>>231 %load = xegpu.load_nd %tdesc[0, 0]232 : !xegpu.tensor_desc<128x64xf32, #xegpu.layout<sg_layout = [4, 4], sg_data = [32, 16], lane_layout = [8, 4], lane_data = [1, 1]>>233 -> vector<128x64xf32>234 %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>235 }{sg_id_range = #xegpu.range<[2, 18]>}236 gpu.return237 }238 239 // CHECK-LABEL: @subgroup_id_range_nested_if240 gpu.func @subgroup_id_range_nested_if(%src: memref<256x128xf32>, %src1: memref<128x64xf32>) {241 %sg_id = gpu.subgroup_id : index242 %c1 = arith.constant 1 : i1243 %c3 = arith.constant 3 : index244 %c32 = arith.constant 32 : index245 %tdesc = xegpu.create_nd_tdesc %src : memref<256x128xf32>246 -> !xegpu.tensor_desc<256x128xf32, #xegpu.layout<sg_layout = [8, 4], sg_data = [32, 32], lane_layout = [8, 4], lane_data = [1, 1]>>247 %load = xegpu.load_nd %tdesc[0, 0]248 : !xegpu.tensor_desc<256x128xf32, #xegpu.layout<sg_layout = [8, 4], sg_data = [32, 32], lane_layout = [8, 4], lane_data = [1, 1]>>249 -> vector<256x128xf32>250 %cond1 = arith.cmpi sge, %sg_id, %c3 : index251 %cond2 = arith.cmpi slt, %sg_id, %c32 : index252 %cond = arith.andi %cond1, %cond2 : i1253 scf.if %c1 {254 scf.if %cond {255 // CHECK: %[[SGID:.*]] = gpu.subgroup_id : index256 // CHECK: %[[C3:.*]] = arith.constant 3 : index257 // CHECK: %[[SUB:.*]] = index.sub %{{.*}}, %[[C3]]258 %td = xegpu.create_nd_tdesc %src1 : memref<128x64xf32>259 -> !xegpu.tensor_desc<128x64xf32, #xegpu.layout<sg_layout = [4, 4], sg_data = [32, 16], lane_layout = [8, 4], lane_data = [1, 1]>>260 %ld = xegpu.load_nd %td[0, 0]261 : !xegpu.tensor_desc<128x64xf32, #xegpu.layout<sg_layout = [4, 4], sg_data = [32, 16], lane_layout = [8, 4], lane_data = [1, 1]>>262 -> vector<128x64xf32>263 %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>264 }265 } {sg_id_range = #xegpu.range<[3, 19]>}266 gpu.return267 }268 269 // CHECK-LABEL: @load_gather270 // CHECK-SAME: %[[ARG0:.*]]: memref<?xf16>271 gpu.func @load_gather(%src : memref<?xf16>) {272 // CHECK: %[[CST:.*]] = arith.constant dense<0> : vector<32x4xindex>273 // CHECK: %[[MASK:.*]] = arith.constant dense<true> : vector<32x4xi1>274 // CHECK: %[[LOAD:.*]] = xegpu.load %[[ARG0]][%[[CST]]], %[[MASK]] <{chunk_size = 1 : i64, l1_hint = #xegpu.cache_hint<cached>}>275 // CHECK-SAME: : memref<?xf16>, vector<32x4xindex>, vector<32x4xi1> -> vector<32x4xf16>276 %offset = arith.constant {layout_result_0 = #xegpu.layout<sg_layout = [8, 4], sg_data = [32, 4]>} dense<0> : vector<256x16xindex>277 %mask = arith.constant {layout_result_0 = #xegpu.layout<sg_layout = [8, 4], sg_data = [32, 4]>} dense<1> : vector<256x16xi1>278 %load = xegpu.load %src[%offset], %mask {chunk_size = 1, layout_result_0 = #xegpu.layout<sg_layout = [8, 4], sg_data = [32, 4]>, l1_hint = #xegpu.cache_hint<cached>}279 : memref<?xf16>, vector<256x16xindex>, vector<256x16xi1> -> vector<256x16xf16>280 gpu.return281 }282 283 // CHECK-LABEL: @store_scatter284 // CHECK-SAME: %[[ARG0:.*]]: memref<256xf16>285 gpu.func @store_scatter(%dest : memref<256xf16>) {286 // CHECK: %[[VAL:.*]] = arith.constant {layout_result_0 = #xegpu.layout<inst_data = [8]>} dense<2.550000e+01> : vector<8xf16>287 // CHECK: %[[CST:.*]] = arith.constant {layout_result_0 = #xegpu.layout<inst_data = [8]>} dense<0> : vector<8xindex>288 // CHECK: %[[MASK:.*]] = arith.constant {layout_result_0 = #xegpu.layout<inst_data = [8]>} dense<true> : vector<8xi1>289 // CHECK: xegpu.store %[[VAL]], %[[ARG0]][%[[CST]]], %[[MASK]] <{chunk_size = 1 : i64, l1_hint = #xegpu.cache_hint<cached>, layout = #xegpu.layout<inst_data = [8]>}>290 // CHECK-SAME: {layout_operand_0 = #xegpu.layout<inst_data = [8]>, layout_operand_2 = #xegpu.layout<inst_data = [8]>,291 // CHECK-SAME: layout_operand_3 = #xegpu.layout<inst_data = [8]>}292 // CHECK-SAME: : vector<8xf16>, memref<256xf16>, vector<8xindex>, vector<8xi1>293 %val = arith.constant {layout_result_0 = #xegpu.layout<sg_layout = [32], sg_data = [8], inst_data = [8]>} dense<25.5> : vector<256xf16>294 %offset = arith.constant {layout_result_0 = #xegpu.layout<sg_layout = [32], sg_data = [8], inst_data = [8]>} dense<0> : vector<256xindex>295 %mask = arith.constant {layout_result_0 = #xegpu.layout<sg_layout = [32], sg_data = [8], inst_data = [8]>} dense<1> : vector<256xi1>296 xegpu.store %val, %dest[%offset], %mask {chunk_size = 1, layout_operand_0 = #xegpu.layout<sg_layout = [32], sg_data = [8], inst_data = [8]>, 297 layout_operand_2 = #xegpu.layout<sg_layout = [32], sg_data = [8], inst_data = [8]>,298 layout_operand_3 = #xegpu.layout<sg_layout = [32], sg_data = [8], inst_data = [8]>,299 l1_hint = #xegpu.cache_hint<cached>}300 : vector<256xf16>, memref<256xf16>, vector<256xindex>, vector<256xi1>301 gpu.return302 }303 304 // CHECK-LABEL: @load_with_non_unit_chunk_size305 // CHECK-SAME: %[[ARG0:.*]]: memref<?xf16>306 gpu.func @load_with_non_unit_chunk_size(%src : memref<?xf16>) {307 // CHECK: %[[CST:.*]] = arith.constant dense<0> : vector<8xindex>308 // CHECK: %[[MASK:.*]] = arith.constant dense<true> : vector<8xi1>309 // CHECK: %[[LOAD:.*]] = xegpu.load %[[ARG0]][%[[CST]]], %[[MASK]] <{chunk_size = 4 : i64, l1_hint = #xegpu.cache_hint<cached>}>310 // CHECK-SAME: : memref<?xf16>, vector<8xindex>, vector<8xi1> -> vector<8x4xf16>311 %offset = arith.constant {layout_result_0 = #xegpu.layout<sg_layout = [32], sg_data = [8]>} dense<0> : vector<256xindex>312 %mask = arith.constant {layout_result_0 = #xegpu.layout<sg_layout = [32], sg_data = [8]>} dense<1> : vector<256xi1>313 %load = xegpu.load %src[%offset], %mask {chunk_size = 4, layout_result_0 = #xegpu.layout<sg_layout = [32, 1], sg_data = [8, 4]>, l1_hint = #xegpu.cache_hint<cached>}314 : memref<?xf16>, vector<256xindex>, vector<256xi1> -> vector<256x4xf16>315 gpu.return316 }317 318 // CHECK-LABEL: distribute_load_matrix319 // CHECK-SAME: [[arg0:%.+]]: memref<32768xi8, 3>320 gpu.func @distribute_load_matrix(%arg0: memref<32768xi8, 3>) {321 //CHECK: [[mdesc:%.+]] = xegpu.create_mem_desc [[arg0]] : memref<32768xi8, 3> -> !xegpu.mem_desc<64x128xf32>322 //CHECK: [[sgid:%.+]] = gpu.subgroup_id : index323 //CHECK: [[c4:%.+]] = arith.constant 4 : index324 //CHECK: [[sgidx:%.+]] = index.remu [[sgid]], [[c4]]325 //CHECK: [[sgidy_tmp:%.+]] = index.divu [[sgid]], [[c4]]326 //CHECK: [[c2:%.+]] = arith.constant 2 : index327 //CHECK: [[sgidy:%.+]] = index.remu [[sgidy_tmp]], [[c2]]328 //CHECK: [[c32:%.+]] = arith.constant 32 : index329 //CHECK: [[l_off_y:%.+]] = index.mul [[sgidy]], [[c32]]330 //CHECK: [[c32_0:%.+]] = arith.constant 32 : index331 //CHECK: [[l_off_x:%.+]] = index.mul [[sgidx]], [[c32_0]]332 //CHECK: [[c64:%.+]] = arith.constant 64 : index333 //CHECK: [[off_y:%.+]] = index.remu [[l_off_y]], [[c64]]334 //CHECK: [[c128:%.+]] = arith.constant 128 : index335 //CHECK: [[off_x:%.+]] = index.remu [[l_off_x]], [[c128]]336 //CHECK: xegpu.load_matrix [[mdesc]][[[off_y]], [[off_x]]] <{layout = #xegpu.layout<lane_layout = [2, 8], lane_data = [1, 1]>}>: !xegpu.mem_desc<64x128xf32>, index, index -> vector<32x32xf32>337 %0 = xegpu.create_mem_desc %arg0 : memref<32768xi8, 3> -> !xegpu.mem_desc<64x128xf32>338 %1 = xegpu.load_matrix %0[0, 0] <{layout = #xegpu.layout<sg_layout = [2, 4], sg_data = [32, 32], lane_layout = [2, 8], lane_data = [1, 1]>}>: !xegpu.mem_desc<64x128xf32> -> vector<64x128xf32>339 gpu.return340 }341 342 //CHECK-LABEL: distribute_store_matrix343 //CHECK-SAME: [[arg0:%.+]]: memref<32768xi8, 3>344 gpu.func @distribute_store_matrix(%arg0 : memref<32768xi8, 3>) {345 //CHECK: [[cst:%.+]] = arith.constant dense<1.000000e+00> : vector<32x32xf32>346 //CHECK: [[mdesc:%.+]] = xegpu.create_mem_desc [[arg0]] : memref<32768xi8, 3> -> !xegpu.mem_desc<64x128xf32>347 //CHECK: [[sgid:%.+]] = gpu.subgroup_id : index348 //CHECK: [[c4:%.+]] = arith.constant 4 : index349 //CHECK: [[sgidx:%.+]] = index.remu [[sgid]], [[c4]]350 //CHECK: [[sgidy_tmp:%.+]] = index.divu [[sgid]], [[c4]]351 //CHECK: [[c2:%.+]] = arith.constant 2 : index352 //CHECK: [[sgidy:%.+]] = index.remu [[sgidy_tmp]], [[c2]]353 //CHECK: [[c32:%.+]] = arith.constant 32 : index354 //CHECK: [[l_off_y:%.+]] = index.mul [[sgidy]], [[c32]]355 //CHECK: [[c32_0:%.+]] = arith.constant 32 : index356 //CHECK: [[l_off_x:%.+]] = index.mul [[sgidx]], [[c32_0]]357 //CHECK: [[c64:%.+]] = arith.constant 64 : index358 //CHECK: [[off_y:%.+]] = index.remu [[l_off_y]], [[c64]]359 //CHECK: [[c128:%.+]] = arith.constant 128 : index360 //CHECK: [[off_x:%.+]] = index.remu [[l_off_x]], [[c128]]361 //CHECK: xegpu.store_matrix [[cst]], [[mdesc]][[[off_y]], [[off_x]]] : vector<32x32xf32>, !xegpu.mem_desc<64x128xf32>, index, index362 %cst = arith.constant {layout_result_0 = #xegpu.layout<sg_layout = [2, 4], sg_data = [32, 32]>} dense<1.0> : vector<64x128xf32>363 %mdesc = xegpu.create_mem_desc %arg0 : memref<32768xi8, 3> -> !xegpu.mem_desc<64x128xf32>364 xegpu.store_matrix %cst, %mdesc[0, 0] {layout = #xegpu.layout<sg_layout = [2, 4], sg_data = [32, 32]>} : vector<64x128xf32>, !xegpu.mem_desc<64x128xf32>365 gpu.return366 }367 368 // CHECK-LABEL: @vector_reduce_dim_0369 gpu.func @vector_reduce_dim_0(%src: memref<4x128xf32>) {370 %cst = arith.constant {layout_result_0 = #xegpu.slice<#xegpu.layout<sg_layout = [1, 32], sg_data = [4, 4]>, dims = [0]>} dense<1.0> : vector<128xf32>371 %tdesc = xegpu.create_nd_tdesc %src : memref<4x128xf32>372 -> !xegpu.tensor_desc<4x128xf32, #xegpu.layout<sg_layout = [1, 32], sg_data = [4, 4]>>373 %load = xegpu.load_nd %tdesc[0, 0]374 : !xegpu.tensor_desc<4x128xf32, #xegpu.layout<sg_layout = [1, 32], sg_data = [4, 4]>>375 -> vector<4x128xf32>376 // CHECK: vector.multi_reduction <add>, {{.*}}, {{.*}} [0] : vector<4x4xf32> to vector<4xf32>377 %reduce = vector.multi_reduction <add>, %load, %cst {layout_result_0 = #xegpu.slice<#xegpu.layout<sg_layout = [1, 32], sg_data = [4, 4]>, dims = [0]>} [0]378 : vector<4x128xf32> to vector<128xf32>379 gpu.return380 }381 382 // CHECK-LABEL: @vector_reduce_dim_1383 gpu.func @vector_reduce_dim_1(%src: memref<256x64xf32>) {384 %cst = arith.constant {layout_result_0 = #xegpu.slice<#xegpu.layout<sg_layout = [16, 1], sg_data = [16, 64]>, dims = [1]>} dense<1.0> : vector<256xf32>385 %tdesc = xegpu.create_nd_tdesc %src : memref<256x64xf32>386 -> !xegpu.tensor_desc<256x64xf32, #xegpu.layout<sg_layout = [16, 1], sg_data = [16, 64]>>387 %load = xegpu.load_nd %tdesc[0, 0]388 : !xegpu.tensor_desc<256x64xf32, #xegpu.layout<sg_layout = [16, 1], sg_data = [16, 64]>>389 -> vector<256x64xf32>390 // CHECK: vector.multi_reduction <add>, {{.*}}, {{.*}} [1] : vector<16x64xf32> to vector<16xf32>391 %reduce = vector.multi_reduction <add>, %load, %cst {layout_result_0 = #xegpu.slice<#xegpu.layout<sg_layout = [16, 1], sg_data = [16, 64]>, dims = [1]>} [1]392 : vector<256x64xf32> to vector<256xf32>393 gpu.return394 }395 396 // CHECK-LABEL: @vector_reduce_4D397 gpu.func @vector_reduce_4D(%src: ui64) {398 %cst_acc = arith.constant {layout_result_0 = #xegpu.slice<#xegpu.layout<sg_layout = [4, 2, 6, 1], sg_data = [1, 1, 1, 32]>, dims = [3]>} dense<0.0> : vector<4x2x6xf16>399 %offset = arith.constant {layout_result_0 = #xegpu.layout<sg_layout = [4, 2, 6, 1], sg_data = [1, 1, 1, 32]>} dense<0> : vector<4x2x6x32xindex>400 %mask = arith.constant {layout_result_0 = #xegpu.layout<sg_layout = [4, 2, 6, 1], sg_data = [1, 1, 1, 32]>} dense<true> : vector<4x2x6x32xi1>401 %load = xegpu.load %src[%offset], %mask {layout_result_0 = #xegpu.layout<sg_layout = [4, 2, 6, 1], sg_data = [1, 1, 1, 32]>} : ui64, vector<4x2x6x32xindex>, vector<4x2x6x32xi1> -> vector<4x2x6x32xf16>402 // CHECK: vector.multi_reduction <add>, {{.*}}, {{.*}} [3] : vector<1x1x1x32xf16> to vector<1x1x1xf16>403 %reduce = vector.multi_reduction <add>, %load, %cst_acc {layout_result_0 = #xegpu.slice<#xegpu.layout<sg_layout = [4, 2, 6, 1], sg_data = [1, 1, 1, 32]>, dims = [3]>} [3]404 : vector<4x2x6x32xf16> to vector<4x2x6xf16>405 gpu.return406 }407 408 // CHECK-LABEL: vector_step_op409 gpu.func @vector_step_op_slice_attr() {410 //CHECK: [[sgId:%.+]] = gpu.subgroup_id : index411 //CHECK: [[c8:%.+]] = arith.constant 8 : index412 //CHECK: [[sgidx:%.+]] = index.remu [[sgId]], [[c8]]413 //CHECK: [[sgidy_tmp:%.+]] = index.divu [[sgId]], [[c8]]414 //CHECK: [[c4:%.+]] = arith.constant 4 : index415 //CHECK: [[sgidy:%.+]] = index.remu [[sgidy_tmp]], [[c4]]416 //CHECK: [[c32:%.+]] = arith.constant 32 : index417 //CHECK: [[LY:%.+]] = index.mul [[sgidy]], [[c32]]418 //CHECK: [[c128:%.+]] = arith.constant 128 : index419 //CHECK: [[MODY:%.+]] = index.remu [[LY]], [[c128]]420 //CHECK: [[BASE:%.+]] = vector.step : vector<32xindex>421 //CHECK: [[CAST:%.+]] = vector.broadcast [[MODY]] : index to vector<32xindex>422 //CHECK: [[ADD:%.+]] = arith.addi [[BASE]], [[CAST]] : vector<32xindex>423 %step = vector.step {layout_result_0 = #xegpu.slice<#xegpu.layout<sg_layout = [4, 8], sg_data = [32, 32]>, dims = [1]>}: vector<128xindex>424 gpu.return425 }426 427 gpu.func @vector_step_op_layout_attr() {428 //CHECK: [[sgId:%.+]] = gpu.subgroup_id : index429 //CHECK: [[c16:%.+]] = arith.constant 16 : index430 //CHECK: [[sgidx:%.+]] = index.remu [[sgId]], [[c16]]431 //CHECK: [[c8:%.+]] = arith.constant 8 : index432 //CHECK: [[LOCALY:%.+]] = index.mul [[sgidx]], [[c8]]433 //CHECK: [[c128:%.+]] = arith.constant 128 : index434 //CHECK: [[MODY:%.+]] = index.remu [[LOCALY]], [[c128]]435 //CHECK: [[BASE:%.+]] = vector.step : vector<8xindex>436 //CHECK: [[CAST:%.+]] = vector.broadcast [[MODY]] : index to vector<8xindex>437 //CHECK: [[ADD:%.+]] = arith.addi [[BASE]], [[CAST]] : vector<8xindex>438 %step = vector.step {layout_result_0 = #xegpu.layout<sg_layout = [16], sg_data = [8]>}: vector<128xindex>439 gpu.return440 }441 442 // CHECK-LABEL: constant_with_slice_attr443 gpu.func @constant_with_slice_attr() {444 //CHECK: [[cst:%.+]] = arith.constant dense<10> : vector<1xindex>445 %cst = arith.constant {layout_result_0 = #xegpu.slice<#xegpu.layout<sg_layout = [4, 2, 6, 1], sg_data = [1, 1, 1, 1]>, dims = [1, 2, 3]>} dense<10> : vector<4xindex>446 gpu.return447 }448 449 // CHECK-LABEL: vector_shape_cast450 gpu.func @vector_shape_cast() {451 %cst = arith.constant {layout_result_0 = #xegpu.slice<#xegpu.layout<sg_layout = [8, 1, 1, 4], sg_data = [1, 1, 1, 32]>, dims = [0, 1, 2]>} dense<10> : vector<128xindex>452 %step = vector.step {layout_result_0 = #xegpu.slice<#xegpu.layout<sg_layout = [8, 1, 1, 4], sg_data = [1, 1, 1, 32]>, dims = [0, 1, 2]>} : vector<128xindex>453 %muli = arith.muli %cst, %step {layout_result_0 = #xegpu.slice<#xegpu.layout<sg_layout = [8, 1, 1, 4], sg_data = [1, 1, 1, 32]>, dims = [0, 1, 2]>} : vector<128xindex>454 //CHECK: vector.shape_cast {{.*}} : vector<32xindex> to vector<1x1x1x32xindex>455 %shape_cast = vector.shape_cast %muli {layout_result_0 = #xegpu.layout<sg_layout = [8, 1, 1, 4], sg_data = [1, 1, 1, 32]>} : vector<128xindex> to vector<1x1x1x128xindex>456 gpu.return457 }458 459 // CHECK-LABEL: vector_broadcast460 gpu.func @vector_broadcast(%arg0: index, %arg1: index) {461 %muli = arith.muli %arg0, %arg1 : index462 // CHECK: vector.broadcast {{.*}} : index to vector<1x1x1x32xindex>463 %broadcast = vector.broadcast %muli {layout_result_0 = #xegpu.layout<sg_layout = [4, 2, 6, 1], sg_data = [1, 1, 1, 32]>} : index to vector<4x2x6x32xindex>464 gpu.return465 }466 467 // CHECK-LABEL: vector_transpose468 gpu.func @vector_transpose(%src: memref<256x32xf32>) {469 %tdesc = xegpu.create_nd_tdesc %src : memref<256x32xf32>470 -> !xegpu.tensor_desc<256x32xf32, #xegpu.layout<sg_layout = [4, 8], sg_data = [64, 32], lane_layout = [16, 1], lane_data = [1, 1], order =[0, 1]>>471 %load = xegpu.load_nd %tdesc[0, 0]472 : !xegpu.tensor_desc<256x32xf32, #xegpu.layout<sg_layout = [4, 8], sg_data = [64, 32], lane_layout = [16, 1], lane_data = [1, 1], order =[0, 1]>>473 -> vector<256x32xf32>474 //CHECK: vector.transpose {{.*}}, [1, 0] {layout_result_0 = #xegpu.layout<lane_layout = [1, 16], lane_data = [1, 1], order = [1, 0]>} : vector<64x32xf32> to vector<32x64xf32>475 %trans = vector.transpose %load, [1, 0] {layout_result_0 = #xegpu.layout<sg_layout = [8, 4], sg_data = [32, 64], lane_layout = [1, 16], lane_data = [1, 1], order =[1, 0]>} : vector<256x32xf32> to vector<32x256xf32>476 gpu.return477 }478 479 // CHECK-LABEL: non_splat_constant_2D480 gpu.func @non_splat_constant_2D() {481 // CHECK-DAG: %[[CST:.*]] = arith.constant dense<0> : vector<1x1xindex>482 // CHECK-DAG: %[[SGID:.*]] = gpu.subgroup_id : index483 // CHECK-DAG: %[[SGIDX:.*]] = index.remu %[[SGID]], %{{.*}}484 // CHECK-DAG: %[[SGIDY_TMP:.*]] = index.divu %[[SGID]], %{{.*}}485 // CHECK-DAG: %[[SGIDY:.*]] = index.remu %[[SGIDY_TMP]], %{{.*}}486 // CHECK-DAG: %[[IDY:.*]] = index.remu %[[SGIDY]], %{{.*}}487 // CHECK-DAG: %[[IDX:.*]] = index.remu %[[SGIDX]], %{{.*}}488 // CHECK-DAG: %[[STRIDECOL:.*]] = arith.muli %[[IDY]], %[[C16:.*]] : index489 // CHECK-DAG: %[[ADD:.*]] = arith.addi %[[C0:.*]], %[[STRIDECOL]] : index490 // CHECK-DAG: %[[STRIDEROW:.*]] = arith.muli %[[IDX]], %[[C0:.*]] : index491 // CHECK-DAG: %[[ADDSTRIDES:.*]] = arith.addi %[[ADD]], %[[STRIDEROW]] : index492 // CHECK-DAG: %[[BCAST:.*]] = vector.broadcast %[[ADDSTRIDES]] : index to vector<1x1xindex>493 // CHECK-DAG: arith.addi %[[CST]], %[[BCAST]] : vector<1x1xindex>494 %cst = arith.constant {layout_result_0 = #xegpu.layout<sg_layout = [32, 1], sg_data = [1, 1]>} dense<[[0], [16], [32], [48], [64], [80], [96], [112], [128], [144], [160], [176], [192], [208], [224], [240], [256], [272], [288], [304], [320], [336], [352], [368], [384], [400], [416], [432], [448], [464], [480], [496]]> : vector<32x1xindex>495 gpu.return496 }497 498 // CHECK-LABEL: non_splat_constant_2D_non_unit_dim499 gpu.func @non_splat_constant_2D_non_unit_dim() {500 // CHECK-DAG: %[[BASECST:.*]] = arith.constant dense<{{\[}}{{\[}}0, 16{{\]}}, {{\[}}8, 24{{\]}}{{\]}}> : vector<2x2xindex>501 // CHECK-DAG: %[[SGID:.*]] = gpu.subgroup_id : index502 // CHECK-DAG: %[[SGIDX:.*]] = index.remu %[[SGID]], %{{.*}}503 // CHECK-DAG: %[[SGIDY_TMP:.*]] = index.divu %[[SGID]], %{{.*}}504 // CHECK-DAG: %[[SGIDY:.*]] = index.remu %[[SGIDY_TMP]], %{{.*}}505 // CHECK-DAG: %[[MULY:.*]] = index.mul %[[SGIDY]], %[[C2:.*]]506 // CHECK-DAG: %[[MULX:.*]] = index.mul %[[SGIDX]], %{{.*}}507 // CHECK-DAG: %[[REMU_Y:.*]] = index.remu %[[MULY]], %[[C8:.*]]508 // CHECK-DAG: %[[REMU_X:.*]] = index.remu %[[MULX]], %{{.*}}509 // CHECK-DAG: %[[MUL5:.*]] = arith.muli %[[REMU_Y]], %{{.*}} : index510 // CHECK-DAG: %[[ADD:.*]] = arith.addi %[[C0:.*]], %[[MUL5]] : index511 // CHECK-DAG: %[[MUL6:.*]] = arith.muli %[[REMU_X]], %[[C16:.*]] : index512 // CHECK-DAG: %[[ADDSTRIDES:.*]] = arith.addi %[[ADD]], %[[MUL6]] : index513 // CHECK-DAG: %[[BCAST:.*]] = vector.broadcast %[[ADDSTRIDES]] : index to vector<2x2xindex>514 // CHECK-DAG: %[[ADDCST:.*]] = arith.addi %[[BASECST]], %[[BCAST]] : vector<2x2xindex>515 %cst_8x8 = arith.constant {layout_result_0 = #xegpu.layout<sg_layout = [4, 4], sg_data = [2, 2]>} dense<[516 [0, 16, 32, 48, 64, 80, 96, 112],517 [8, 24, 40, 56, 72, 88, 104, 120],518 [16, 32, 48, 64, 80, 96, 112, 128],519 [24, 40, 56, 72, 88, 104, 120, 136],520 [32, 48, 64, 80, 96, 112, 128, 144],521 [40, 56, 72, 88, 104, 120, 136, 152],522 [48, 64, 80, 96, 112, 128, 144, 160],523 [56, 72, 88, 104, 120, 136, 152, 168]524 ]> : vector<8x8xindex>525 gpu.return526 }527 528 // CHECK-LABEL: non_splat_constant529 gpu.func @non_splat_constant() {530 // CHECK-DAG: %[[CST:.*]] = arith.constant dense<0> : vector<1xindex>531 // CHECK-DAG: %[[SGID:.*]] = gpu.subgroup_id : index532 // CHECK-DAG: %[[REMU:.*]] = index.remu %[[SGID]], %{{.*}}533 // CHECK-DAG: %[[REMU2:.*]] = index.remu %[[REMU]], %{{.*}}534 // CHECK-DAG: %[[MUL:.*]] = arith.muli %[[REMU2]], %[[C16:.*]] : index535 // CHECK-DAG: %[[ADDSTRIDES:.*]] = arith.addi %[[C0:.*]], %[[MUL]] : index536 // CHECK-DAG: %[[BCAST:.*]] = vector.broadcast %[[ADDSTRIDES]] : index to vector<1xindex>537 // CHECK-DAG: %[[ADD:.*]] = arith.addi %[[CST]], %[[BCAST]] : vector<1xindex>538 %cst = arith.constant {layout_result_0 = #xegpu.layout<sg_layout = [32], sg_data = [1]>} dense<[0, 16, 32, 48, 64, 80, 96, 112, 128, 144, 160, 176, 192, 208, 224, 240, 256, 272, 288, 304, 320, 336, 352, 368, 384, 400, 416, 432, 448, 464, 480, 496]> : vector<32xindex>539 // CHECK: arith.constant dense<{{\[}}{{\[}}0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15{{\]}}{{\]}}> : vector<1x16xindex>540 %cst_1 = arith.constant {layout_result_0 = #xegpu.layout<sg_layout = [32, 1], sg_data = [1, 16]>} dense<[[0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15]]> : vector<1x16xindex>541 gpu.return542 }543 544 // CHECK-LABEL: scalar_broadcast545 gpu.func @scalar_broadcast(%arg0: index) {546 // CHECK: vector.broadcast {{.*}} : index to vector<1x1x1xindex>547 %broadcast = vector.broadcast %arg0 {layout_result_0 = #xegpu.layout<sg_layout = [4, 8, 1], sg_data = [1, 1, 1]>} : index to vector<4x1x1xindex>548 gpu.return549 }550 551 // CHECK-LABEL: vector_mask_1D552 gpu.func @vector_mask_1D() {553 // CHECK-DAG: %[[SGID:.*]] = gpu.subgroup_id : index554 // CHECK-DAG: %[[REMU:.*]] = index.remu %[[SGID]], %[[C2:.*]]555 // CHECK-DAG: %[[MUL:.*]] = index.mul %[[REMU]], %[[C16:.*]]556 // CHECK-DAG: %[[REMU2:.*]] = index.remu %[[MUL]], %[[C32:.*]]557 // CHECK-DAG: %[[SUB:.*]] = arith.subi %[[C8:.*]], %[[REMU2]] : index558 // CHECK-DAG: %[[MAX:.*]] = arith.maxsi %[[SUB]], %[[C0:.*]] : index559 // CHECK-DAG: %[[MIN:.*]] = arith.minsi %[[MAX]], %[[C16:.*]] : index560 // CHECK-DAG: %[[MASK:.*]] = vector.create_mask %[[MIN]] : vector<16xi1>561 %constant_mask = vector.constant_mask [8] {layout_result_0 = #xegpu.layout<sg_layout = [2], sg_data = [16]>} : vector<32xi1>562 gpu.return563 }564 565 // CHECK-LABEL: vector_mask_2D566 gpu.func @vector_mask_2D() {567 // CHECK-DAG: %[[SGID:.*]] = gpu.subgroup_id : index568 // CHECK-DAG: %[[SGIDX:.*]] = index.remu %[[SGID]], %[[C4:.*]]569 // CHECK-DAG: %[[SGIDY_TMP:.*]] = index.divu %[[SGID]], %[[C4:.*]]570 // CHECK-DAG: %[[SGIDY:.*]] = index.remu %[[SGIDY_TMP]], %[[C8:.*]]571 // CHECK-DAG: %[[ROW:.*]] = index.mul %[[SGIDY]], %[[C32:.*]]572 // CHECK-DAG: %[[COL:.*]] = index.mul %[[SGIDX]], %[[C32:.*]]573 // CHECK-DAG: %[[MODROW:.*]] = index.remu %[[ROW]], %[[C256:.*]]574 // CHECK-DAG: %[[MODCOL:.*]] = index.remu %[[COL]], %[[C128:.*]]575 // CHECK-DAG: %[[SUBROW:.*]] = arith.subi %[[C16:.*]], %[[MODROW]] : index576 // CHECK-DAG: %[[MAXROW:.*]] = arith.maxsi %[[SUBROW]], %[[C4:.*]] : index577 // CHECK-DAG: %[[MINROW:.*]] = arith.minsi %[[MAXROW]], %[[C32:.*]] : index578 // CHECK-DAG: %[[SUBCOL:.*]] = arith.subi %[[C16:.*]], %[[MODCOL]] : index579 // CHECK-DAG: %[[MAXCOL:.*]] = arith.maxsi %[[SUBCOL]], %[[C7:.*]] : index580 // CHECK-DAG: %[[MINCOL:.*]] = arith.minsi %[[MAXCOL]], %[[C32:.*]] : index581 // CHECK-DAG: %[[MASK:.*]] = vector.create_mask %[[MINROW]], %[[MINCOL]] : vector<32x32xi1>582 %constant_mask = vector.constant_mask [16, 16] {layout_result_0 = #xegpu.layout<sg_layout = [8, 4], sg_data = [32, 32]>} : vector<256x128xi1>583 gpu.return584 }585 586 // CHECK-LABEL: distribute_load_slice_attr587 gpu.func @distribute_load_slice_attr() {588 %2 = memref.alloca() {alignment = 1024} : memref<4096xf32>589 %offset = arith.constant {layout_result_0 = #xegpu.layout<sg_layout = [8], sg_data = [32], inst_data = [16]> } dense<0> : vector<256xindex>590 %mask = arith.constant {layout_result_0 = #xegpu.layout<sg_layout = [8], sg_data = [32], inst_data = [16]> } dense<1> : vector<256xi1>591 592 // CHECK: %[[LOAD:.*]] = xegpu.load {{.*}} <{chunk_size = 1 : i64, layout = #xegpu.slice<#xegpu.layout<inst_data = [8, 16]>, dims = [0]>}>593 // CHECK-SAME: {layout_result_0 = #xegpu.slice<#xegpu.layout<inst_data = [8, 16]>, dims = [0]>} :594 // CHECK-SAME: memref<4096xf32>, vector<32xindex>, vector<32xi1> -> vector<32xf32>595 %3 = xegpu.load %2[%offset], %mask {chunk_size = 1, layout_result_0 = #xegpu.slice<#xegpu.layout<sg_layout = [8, 8], sg_data = [32, 32], inst_data = [8, 16]>, dims = [0]> } : memref<4096xf32>, vector<256xindex>, vector<256xi1> -> vector<256xf32>596 597 // CHECK: %[[BROADCAST:.*]] = vector.broadcast %[[LOAD]] {layout_result_0 = #xegpu.layout<inst_data = [8, 16]>} : vector<32xf32> to vector<32x32xf32>598 %4 = vector.broadcast %3 {layout_result_0 =599 #xegpu.layout<sg_layout = [8, 8], sg_data = [32, 32], inst_data = [8, 16]>} : vector<256xf32> to vector<256x256xf32>600 gpu.return601 }602}603