476 lines · plain
1// RUN: mlir-opt --test-xegpu-unrolling-patterns -split-input-file %s | FileCheck %s2 3gpu.module @test {4 5 // CHECK-LABEL: create_nd_tdesc6 // CHECK-SAME: [[arg0:%.+]]: memref<24x32xf32>7 // CHECK-COUNT-6: [[tdesc:%.+]] = xegpu.create_nd_tdesc [[arg0]][{{.*}}] : memref<24x32xf32> -> !xegpu.tensor_desc<8x16xf32>8 // CHECK: [[cast:%.+]] = builtin.unrealized_conversion_cast9 // CHECK-SAME: !xegpu.tensor_desc<8x16xf32>, !xegpu.tensor_desc<8x16xf32>,10 // CHECK-SAME: !xegpu.tensor_desc<8x16xf32>, !xegpu.tensor_desc<8x16xf32>,11 // CHECK-SAME: !xegpu.tensor_desc<8x16xf32>, !xegpu.tensor_desc<8x16xf32>12 // CHECK-SAME: to !xegpu.tensor_desc<24x32xf32, #xegpu.layout<inst_data = [8, 16]>> {__xegpu_blocking_tile_shape__ = array<i64: 8, 16>, __xegpu_blocking_unpack__}13 gpu.func @create_nd_tdesc(%src: memref<24x32xf32>) -> !xegpu.tensor_desc<24x32xf32, #xegpu.layout<inst_data = [8, 16]>> {14 %tdesc = xegpu.create_nd_tdesc %src[0, 0] : memref<24x32xf32> -> !xegpu.tensor_desc<24x32xf32, #xegpu.layout<inst_data = [8, 16]>>15 gpu.return %tdesc : !xegpu.tensor_desc<24x32xf32, #xegpu.layout<inst_data = [8, 16]>>16 }17 18 //-----19 20 // CHECK-LABEL: create_nd_tdesc_1d21 // CHECK-SAME: [[arg0:%.+]]: memref<64xf32>22 // CHECK-COUNT-2: [[tdesc:%.+]] = xegpu.create_nd_tdesc [[arg0]][{{.*}}] : memref<64xf32> -> !xegpu.tensor_desc<16xf32>23 // CHECK: [[cast:%.+]] = builtin.unrealized_conversion_cast24 // CHECK-SAME: !xegpu.tensor_desc<16xf32>, !xegpu.tensor_desc<16xf32>25 // CHECK-SAME: to !xegpu.tensor_desc<32xf32, #xegpu.layout<inst_data = [16]>> {__xegpu_blocking_tile_shape__ = array<i64: 16>, __xegpu_blocking_unpack__}26 gpu.func @create_nd_tdesc_1d(%src: memref<64xf32>) -> !xegpu.tensor_desc<32xf32, #xegpu.layout<inst_data = [16]>> {27 %tdesc = xegpu.create_nd_tdesc %src[0] : memref<64xf32> -> !xegpu.tensor_desc<32xf32, #xegpu.layout<inst_data = [16]>>28 gpu.return %tdesc : !xegpu.tensor_desc<32xf32, #xegpu.layout<inst_data = [16]>>29 }30 31 //-----32 33 // CHECK-LABEL: update_nd_tdesc34 // CHECK-SAME: [[arg0:%.+]]: memref<24x32xf32>35 // CHECK-COUNT-6: [[tdesc:%.+]] = xegpu.create_nd_tdesc [[arg0]][{{.*}}] : memref<24x32xf32> -> !xegpu.tensor_desc<8x16xf32>36 // CHECK-COUNT-6: [[update:%.+]] = xegpu.update_nd_offset {{.*}} : !xegpu.tensor_desc<8x16xf32>37 gpu.func @update_nd_tdesc(%src: memref<24x32xf32>) -> !xegpu.tensor_desc<24x32xf32, #xegpu.layout<inst_data = [8, 16]>> {38 %tdesc = xegpu.create_nd_tdesc %src[0, 0] : memref<24x32xf32> -> !xegpu.tensor_desc<24x32xf32, #xegpu.layout<inst_data = [8, 16]>>39 %update = xegpu.update_nd_offset %tdesc, [0, 16] : !xegpu.tensor_desc<24x32xf32, #xegpu.layout<inst_data = [8, 16]>>40 gpu.return %update : !xegpu.tensor_desc<24x32xf32, #xegpu.layout<inst_data = [8, 16]>>41 }42 43 //-----44 45 // CHECK-LABEL: update_nd_tdesc_1d46 // CHECK-SAME: [[arg0:%.+]]: memref<64xf32>47 // CHECK-COUNT-2: [[tdesc:%.+]] = xegpu.create_nd_tdesc [[arg0]][{{.*}}] : memref<64xf32> -> !xegpu.tensor_desc<16xf32>48 // CHECK-COUNT-2: [[update:%.+]] = xegpu.update_nd_offset {{.*}} : !xegpu.tensor_desc<16xf32>49 gpu.func @update_nd_tdesc_1d(%src: memref<64xf32>) -> !xegpu.tensor_desc<32xf32, #xegpu.layout<inst_data = [16]>> {50 %tdesc = xegpu.create_nd_tdesc %src[0] : memref<64xf32> -> !xegpu.tensor_desc<32xf32, #xegpu.layout<inst_data = [16]>>51 %update = xegpu.update_nd_offset %tdesc, [32] : !xegpu.tensor_desc<32xf32, #xegpu.layout<inst_data = [16]>>52 gpu.return %update : !xegpu.tensor_desc<32xf32, #xegpu.layout<inst_data = [16]>>53 }54 55 //-----56 57 // CHECK-LABEL: prefetch_nd_tdesc58 // CHECK-SAME: [[arg0:%.+]]: memref<24x32xf32>59 // CHECK-COUNT-6: [[tdesc:%.+]] = xegpu.create_nd_tdesc [[arg0]][{{.*}}] : memref<24x32xf32> -> !xegpu.tensor_desc<8x16xf32>60 // CHECK-COUNT-6: xegpu.prefetch_nd {{.*}} : !xegpu.tensor_desc<8x16xf32>61 gpu.func @prefetch_nd_tdesc(%src: memref<24x32xf32>) {62 %tdesc = xegpu.create_nd_tdesc %src[0, 0] : memref<24x32xf32> -> !xegpu.tensor_desc<24x32xf32, #xegpu.layout<inst_data = [8, 16]>>63 xegpu.prefetch_nd %tdesc : !xegpu.tensor_desc<24x32xf32, #xegpu.layout<inst_data = [8, 16]>>64 gpu.return65 }66 67 //-----68 69 // CHECK-LABEL: prefetch_nd_tdesc_1d70 // CHECK-SAME: [[arg0:%.+]]: memref<64xf32>71 // CHECK-COUNT-4: [[tdesc:%.+]] = xegpu.create_nd_tdesc [[arg0]][{{.*}}] : memref<64xf32> -> !xegpu.tensor_desc<16xf32>72 // CHECK-COUNT-4: xegpu.prefetch_nd {{.*}} : !xegpu.tensor_desc<16xf32>73 gpu.func @prefetch_nd_tdesc_1d(%src: memref<64xf32>) {74 %tdesc = xegpu.create_nd_tdesc %src[0] : memref<64xf32> -> !xegpu.tensor_desc<64xf32, #xegpu.layout<inst_data = [16]>>75 xegpu.prefetch_nd %tdesc : !xegpu.tensor_desc<64xf32, #xegpu.layout<inst_data = [16]>>76 gpu.return77 }78 79 //-----80 // CHECK-LABEL: load_nd81 // CHECK-SAME: [[arg0:%.+]]: memref<24x32xf32>82 // CHECK-COUNT-6: [[tdesc:%.+]] = xegpu.create_nd_tdesc [[arg0]][{{.*}}] : memref<24x32xf32> -> !xegpu.tensor_desc<8x16xf32>83 // CHECK-COUNT-6: [[ld:%.+]] = xegpu.load_nd {{.*}} : !xegpu.tensor_desc<8x16xf32> -> vector<8x16xf32>84 // CHECK-COUNT-6: [[insert:%.+]] = vector.insert_strided_slice {{.*}} : vector<8x16xf32> into vector<24x32xf32>85 gpu.func @load_nd(%src: memref<24x32xf32>) -> vector<24x32xf32> {86 %tdesc = xegpu.create_nd_tdesc %src[0, 0] : memref<24x32xf32> -> !xegpu.tensor_desc<24x32xf32, #xegpu.layout<inst_data = [8, 16]>>87 %ld = xegpu.load_nd %tdesc: !xegpu.tensor_desc<24x32xf32, #xegpu.layout<inst_data = [8, 16]>> -> vector<24x32xf32>88 gpu.return %ld : vector<24x32xf32>89 }90 91 //-----92 93 // CHECK-LABEL: load_nd_1d94 // CHECK-SAME: [[arg0:%.+]]: memref<64xf32>95 // CHECK-COUNT-4: [[tdesc:%.+]] = xegpu.create_nd_tdesc [[arg0]][{{.*}}] : memref<64xf32> -> !xegpu.tensor_desc<16xf32>96 // CHECK-COUNT-4: [[ld:%.+]] = xegpu.load_nd {{.*}} : !xegpu.tensor_desc<16xf32> -> vector<16xf32>97 // CHECK-COUNT-4: [[insert:%.+]] = vector.insert_strided_slice {{.*}} : vector<16xf32> into vector<64xf32>98 gpu.func @load_nd_1d(%src: memref<64xf32>) -> vector<64xf32> {99 %tdesc = xegpu.create_nd_tdesc %src[0] : memref<64xf32> -> !xegpu.tensor_desc<64xf32, #xegpu.layout<inst_data = [16]>>100 %data = xegpu.load_nd %tdesc: !xegpu.tensor_desc<64xf32, #xegpu.layout<inst_data = [16]>> -> vector<64xf32>101 gpu.return %data : vector<64xf32>102 }103 104 //-----105 106 // CHECK-LABEL: store_nd107 // CHECK-SAME: [[arg0:%.+]]: memref<24x32xf32>108 // CHECK-COUNT-6: [[tdesc:%.+]] = xegpu.create_nd_tdesc [[arg0]][{{.*}}] : memref<24x32xf32> -> !xegpu.tensor_desc<8x16xf32>109 // CHECK-COUNT-6: xegpu.store_nd {{.*}} : vector<8x16xf32>, !xegpu.tensor_desc<8x16xf32>110 gpu.func @store_nd(%src: memref<24x32xf32>) {111 %tdesc = xegpu.create_nd_tdesc %src[0, 0] : memref<24x32xf32> -> !xegpu.tensor_desc<24x32xf32, #xegpu.layout<inst_data = [8, 16]>>112 %data = arith.constant dense<9.0> : vector<24x32xf32>113 xegpu.store_nd %data, %tdesc: vector<24x32xf32>, !xegpu.tensor_desc<24x32xf32, #xegpu.layout<inst_data = [8, 16]>>114 gpu.return115 }116 117 //-----118 119 // CHECK-LABEL: store_nd_1d120 // CHECK-SAME: [[arg0:%.+]]: memref<64xf32>121 // CHECK-COUNT-4: [[tdesc:%.+]] = xegpu.create_nd_tdesc [[arg0]][{{.*}}] : memref<64xf32> -> !xegpu.tensor_desc<16xf32>122 // CHECK-COUNT-4: xegpu.store_nd {{.*}} : vector<16xf32>, !xegpu.tensor_desc<16xf32>123 gpu.func @store_nd_1d(%src: memref<64xf32>) {124 %tdesc = xegpu.create_nd_tdesc %src[0] : memref<64xf32> -> !xegpu.tensor_desc<64xf32, #xegpu.layout<inst_data = [16]>>125 %data = arith.constant dense<9.0> : vector<64xf32>126 xegpu.store_nd %data, %tdesc: vector<64xf32>, !xegpu.tensor_desc<64xf32, #xegpu.layout<inst_data = [16]>>127 gpu.return128 }129 130 //-----131 132 // CHECK-LABEL: createNd_loadNd_storeNd133 // CHECK-SAME: [[arg0:%.+]]: memref<24x32xf32>134 //CHECK-COUNT-6: [[tdesc:%.+]] = xegpu.create_nd_tdesc [[arg0]][{{.*}}] : memref<24x32xf32> -> !xegpu.tensor_desc<8x16xf32>135 //CHECK-COUNT-6: [[data:%.+]] = xegpu.load_nd {{.*}} : !xegpu.tensor_desc<8x16xf32> -> vector<8x16xf32>136 //CHECK-COUNT-6: [[insert:%.+]] = vector.insert_strided_slice {{.*}} : vector<8x16xf32> into vector<24x32xf32>137 //CHECK: [[add:%.+]] = arith.addf {{.*}} : vector<24x32xf32>138 //CHECK-COUNT-6: [[extract:%.+]] = vector.extract_strided_slice {{.*}} : vector<24x32xf32> to vector<8x16xf32>139 //CHECK-COUNT-6: xegpu.store_nd {{.*}} : vector<8x16xf32>, !xegpu.tensor_desc<8x16xf32>140 gpu.func @createNd_loadNd_storeNd(%src: memref<24x32xf32>) {141 %tdesc = xegpu.create_nd_tdesc %src[0, 0] : memref<24x32xf32> -> !xegpu.tensor_desc<24x32xf32, #xegpu.layout<inst_data = [8, 16]>>142 %data = arith.constant dense<9.0> : vector<24x32xf32>143 %ld = xegpu.load_nd %tdesc: !xegpu.tensor_desc<24x32xf32, #xegpu.layout<inst_data = [8, 16]>> -> vector<24x32xf32>144 %add = arith.addf %data, %ld : vector<24x32xf32>145 xegpu.store_nd %add, %tdesc: vector<24x32xf32>, !xegpu.tensor_desc<24x32xf32, #xegpu.layout<inst_data = [8, 16]>>146 gpu.return147 }148 149 //-----150 151 // CHECK-LABEL: dpas152 // CHECK-SAME: [[arg0:%.+]]: vector<32x32xf16>, [[arg1:%.+]]: vector<32x32xf16>153 //CHECK-COUNT-8: [[extract1:%.+]] = vector.extract_strided_slice [[arg0]] {{.*}} : vector<32x32xf16> to vector<8x16xf16>154 //CHECK-COUNT-4: [[extract2:%.+]] = vector.extract_strided_slice [[arg1]] {{.*}} : vector<32x32xf16> to vector<16x16xf16>155 //CHECK-COUNT-16: [[dpas:%.+]] = xegpu.dpas {{.*}} -> vector<8x16xf32>156 //CHECK-COUNT-8: [[insert:%.+]] = vector.insert_strided_slice {{.*}} : vector<8x16xf32> into vector<32x32xf32>157 gpu.func @dpas(%a: vector<32x32xf16>, %b: vector<32x32xf16>) -> vector<32x32xf32> {158 %c = xegpu.dpas %a, %b : vector<32x32xf16>, vector<32x32xf16> -> vector<32x32xf32>159 gpu.return %c : vector<32x32xf32>160 }161 162//-----163 164 // CHECK-LABEL: create_tdesc_vec165 // CHECK-SAME: [[arg0:%.+]]: ui64166 // CHECK-COUNT-2: xegpu.create_tdesc [[arg0]], {{.*}} : ui64, vector<16xindex> -> !xegpu.tensor_desc<16xf32, #xegpu.scatter_tdesc_attr<>>167 gpu.func @create_tdesc_vec(%src: ui64) -> !xegpu.tensor_desc<32xf32, #xegpu.scatter_tdesc_attr<>, #xegpu.layout<inst_data = [16]>> {168 %cst = arith.constant dense<[169 0, 8, 16, 24, 32, 40, 48, 56,170 64, 72, 80, 88, 96, 104, 112, 120,171 128, 136, 144, 152, 160, 168, 176, 184,172 192, 200, 208, 216, 224, 232, 240, 248173 ]> : vector<32xindex>174 %tdesc = xegpu.create_tdesc %src, %cst : ui64, vector<32xindex> -> !xegpu.tensor_desc<32xf32, #xegpu.scatter_tdesc_attr<>, #xegpu.layout<inst_data = [16]>>175 gpu.return %tdesc : !xegpu.tensor_desc<32xf32, #xegpu.scatter_tdesc_attr<>, #xegpu.layout<inst_data = [16]>>176 }177 178//-----179 180 // CHECK-LABEL: create_tdesc_step181 // CHECK-SAME: [[arg0:%.+]]: ui64182 // CHECK-COUNT-2: xegpu.create_tdesc [[arg0]], {{.*}} : ui64, vector<16xindex> -> !xegpu.tensor_desc<16xf32, #xegpu.scatter_tdesc_attr<>>183 gpu.func @create_tdesc_step(%src: ui64) -> !xegpu.tensor_desc<32xf32, #xegpu.scatter_tdesc_attr<>, #xegpu.layout<inst_data = [16]>> {184 %step = arith.constant dense<8> : vector<32xindex>185 %seq = vector.step : vector<32xindex>186 %cst = arith.muli %seq, %step : vector<32xindex>187 %tdesc = xegpu.create_tdesc %src, %cst : ui64, vector<32xindex> -> !xegpu.tensor_desc<32xf32, #xegpu.scatter_tdesc_attr<>, #xegpu.layout<inst_data = [16]>>188 gpu.return %tdesc : !xegpu.tensor_desc<32xf32, #xegpu.scatter_tdesc_attr<>, #xegpu.layout<inst_data = [16]>>189 }190 191//-----192 193 // CHECK-LABEL: load194 // CHECK-SAME: [[arg0:%.+]]: ui64195 // CHECK-COUNT-2: xegpu.create_tdesc [[arg0]], {{.*}} : ui64, vector<16xindex> -> !xegpu.tensor_desc<16xf32, #xegpu.scatter_tdesc_attr<>>196 // CHECK-COUNT-2: xegpu.load {{.*}} : !xegpu.tensor_desc<16xf32, #xegpu.scatter_tdesc_attr<>>, vector<16xi1> -> vector<16xf32>197 gpu.func @load(%src: ui64) -> vector<32xf32> {198 %cst = arith.constant dense<[199 0, 8, 16, 24, 32, 40, 48, 56,200 64, 72, 80, 88, 96, 104, 112, 120,201 128, 136, 144, 152, 160, 168, 176, 184,202 192, 200, 208, 216, 224, 232, 240, 248203 ]> : vector<32xindex>204 205 %c17 = arith.constant 17: index206 %mask = vector.create_mask %c17: vector<32xi1>207 %tdesc = xegpu.create_tdesc %src, %cst : ui64, vector<32xindex> -> !xegpu.tensor_desc<32xf32, #xegpu.scatter_tdesc_attr<>, #xegpu.layout<inst_data = [16]>>208 %ld = xegpu.load %tdesc, %mask: !xegpu.tensor_desc<32xf32, #xegpu.scatter_tdesc_attr<>, #xegpu.layout<inst_data = [16]>>, vector<32xi1> -> vector<32xf32>209 210 gpu.return %ld : vector<32xf32>211 }212 213//-----214 215 216 // CHECK-LABEL: load_with_offsets217 // CHECK-SAME: [[arg0:%.+]]: ui64218 // CHECK-COUNT-2: xegpu.load {{.*}}[{{.*}}], {{.*}} <{chunk_size = 1 : i64, l1_hint = #xegpu.cache_hint<cached>}> : ui64, vector<16xindex>, vector<16xi1> -> vector<16xf32>219 gpu.func @load_with_offsets(%src: ui64) -> vector<32xf32> {220 %cst = arith.constant dense<[221 0, 8, 16, 24, 32, 40, 48, 56,222 64, 72, 80, 88, 96, 104, 112, 120,223 128, 136, 144, 152, 160, 168, 176, 184,224 192, 200, 208, 216, 224, 232, 240, 248225 ]> : vector<32xindex>226 227 %c17 = arith.constant 17: index228 %mask = vector.create_mask %c17: vector<32xi1>229 %ld = xegpu.load %src[%cst], %mask {chunk_size = 1, layout_result_0 = #xegpu.layout<inst_data = [16]>, l1_hint = #xegpu.cache_hint<cached>} : ui64, vector<32xindex>, vector<32xi1> -> vector<32xf32>230 231 gpu.return %ld : vector<32xf32>232 }233 234//-----235 236 // CHECK-LABEL: prefetch237 // CHECK-SAME: [[arg0:%.+]]: ui64238 // CHECK-COUNT-2: xegpu.create_tdesc [[arg0]], {{.*}} : ui64, vector<16xindex> -> !xegpu.tensor_desc<16xf32, #xegpu.scatter_tdesc_attr<>>239 // CHECK-COUNT-2: xegpu.prefetch {{.*}} : !xegpu.tensor_desc<16xf32, #xegpu.scatter_tdesc_attr<>>240 gpu.func @prefetch(%src: ui64) {241 242 %cst = arith.constant dense<[243 0, 8, 16, 24, 32, 40, 48, 56,244 64, 72, 80, 88, 96, 104, 112, 120,245 128, 136, 144, 152, 160, 168, 176, 184,246 192, 200, 208, 216, 224, 232, 240, 248247 ]> : vector<32xindex>248 249 %tdesc = xegpu.create_tdesc %src, %cst : ui64, vector<32xindex> -> !xegpu.tensor_desc<32xf32, #xegpu.scatter_tdesc_attr<>, #xegpu.layout<inst_data = [16]>>250 251 xegpu.prefetch %tdesc: !xegpu.tensor_desc<32xf32, #xegpu.scatter_tdesc_attr<>, #xegpu.layout<inst_data = [16]>>252 gpu.return253 }254 255//-----256 257 // CHECK-LABEL: store258 // CHECK-SAME: [[arg0:%.+]]: ui64259 // CHECK-COUNT-2: xegpu.create_tdesc [[arg0]], {{.*}} : ui64, vector<16xindex> -> !xegpu.tensor_desc<16xf32, #xegpu.scatter_tdesc_attr<>>260 // CHECK-COUNT-2: xegpu.store {{.*}} : vector<16xf32>, !xegpu.tensor_desc<16xf32, #xegpu.scatter_tdesc_attr<>>, vector<16xi1>261 gpu.func @store(%src: ui64) {262 %cst = arith.constant dense<[263 0, 8, 16, 24, 32, 40, 48, 56,264 64, 72, 80, 88, 96, 104, 112, 120,265 128, 136, 144, 152, 160, 168, 176, 184,266 192, 200, 208, 216, 224, 232, 240, 248267 ]> : vector<32xindex>268 269 %c17 = arith.constant 17: index270 %mask = vector.create_mask %c17: vector<32xi1>271 272 %st_vec = arith.constant dense<1023.0>: vector<32xf32>273 %tdesc = xegpu.create_tdesc %src, %cst : ui64, vector<32xindex> -> !xegpu.tensor_desc<32xf32, #xegpu.scatter_tdesc_attr<>, #xegpu.layout<inst_data = [16]>>274 xegpu.store %st_vec, %tdesc, %mask: vector<32xf32>, !xegpu.tensor_desc<32xf32, #xegpu.scatter_tdesc_attr<>, #xegpu.layout<inst_data = [16]>>, vector<32xi1>275 276 gpu.return277 }278 279 //-----280 281 // CHECK-LABEL: store_with_offsets282 // CHECK-SAME: [[arg0:%.+]]: ui64283 // CHECK-COUNT-2: xegpu.store {{.*}}[{{.*}}], {{.*}} <{chunk_size = 1 : i64, l1_hint = #xegpu.cache_hint<cached>}> : vector<16xf32>, ui64, vector<16xindex>, vector<16xi1>284 gpu.func @store_with_offsets(%src: ui64) {285 %cst = arith.constant dense<[286 0, 8, 16, 24, 32, 40, 48, 56,287 64, 72, 80, 88, 96, 104, 112, 120,288 128, 136, 144, 152, 160, 168, 176, 184,289 192, 200, 208, 216, 224, 232, 240, 248290 ]> : vector<32xindex>291 292 %c17 = arith.constant 17: index293 %mask = vector.create_mask %c17: vector<32xi1>294 295 %st_vec = arith.constant dense<1023.0>: vector<32xf32>296 xegpu.store %st_vec, %src[%cst], %mask {chunk_size = 1, layout = #xegpu.layout<inst_data = [16]>, l1_hint = #xegpu.cache_hint<cached>} : vector<32xf32>, ui64, vector<32xindex>, vector<32xi1>297 298 gpu.return299 }300 301//-----302 // CHECK-LABEL: create_tdesc_step_chunk303 // CHECK-SAME: [[arg0:%.+]]: ui64304 // CHECK-COUNT-2: xegpu.create_tdesc [[arg0]], {{.*}} : ui64, vector<16xindex> -> !xegpu.tensor_desc<16x4xf32, #xegpu.scatter_tdesc_attr<chunk_size = 4 : i64>>305 gpu.func @create_tdesc_step_chunk(%src: ui64) -> !xegpu.tensor_desc<32x4xf32, #xegpu.scatter_tdesc_attr<chunk_size=4>, #xegpu.layout<inst_data = [16, 4]>> {306 %step = arith.constant dense<8> : vector<32xindex>307 %seq = vector.step : vector<32xindex>308 %cst = arith.muli %seq, %step : vector<32xindex>309 %tdesc = xegpu.create_tdesc %src, %cst : ui64, vector<32xindex> -> !xegpu.tensor_desc<32x4xf32, #xegpu.scatter_tdesc_attr<chunk_size=4>, #xegpu.layout<inst_data = [16, 4]>>310 gpu.return %tdesc : !xegpu.tensor_desc<32x4xf32, #xegpu.scatter_tdesc_attr<chunk_size=4>, #xegpu.layout<inst_data = [16, 4]>>311 }312 313//-----314 // CHECK-LABEL: create_tdesc_step_chunk2315 // CHECK-SAME: [[arg0:%.+]]: ui64316 // CHECK-COUNT-4: xegpu.create_tdesc [[arg0]], {{.*}} : ui64, vector<16xindex> -> !xegpu.tensor_desc<16x2xf32, #xegpu.scatter_tdesc_attr<chunk_size = 2 : i64>>317 gpu.func @create_tdesc_step_chunk2(%src: ui64) -> !xegpu.tensor_desc<32x4xf32, #xegpu.scatter_tdesc_attr<chunk_size=4>, #xegpu.layout<inst_data = [16, 2]>> {318 %step = arith.constant dense<8> : vector<32xindex>319 %seq = vector.step : vector<32xindex>320 %cst = arith.muli %seq, %step : vector<32xindex>321 %tdesc = xegpu.create_tdesc %src, %cst : ui64, vector<32xindex> -> !xegpu.tensor_desc<32x4xf32, #xegpu.scatter_tdesc_attr<chunk_size=4>, #xegpu.layout<inst_data = [16, 2]>>322 gpu.return %tdesc : !xegpu.tensor_desc<32x4xf32, #xegpu.scatter_tdesc_attr<chunk_size=4>, #xegpu.layout<inst_data = [16, 2]>>323 }324 325// CHECK-LABEL: create_tdesc_step_chunk3326 // CHECK-SAME: [[arg0:%.+]]: ui64327 // CHECK: xegpu.create_tdesc [[arg0]], {{.*}} : ui64, vector<16xindex> -> !xegpu.tensor_desc<16x2xf32, #xegpu.scatter_tdesc_attr<chunk_size = 2 : i64>>328 // CHECK: arith.addi %{{.*}}, %{{.*}} : vector<16xindex>329 // CHECK: xegpu.create_tdesc [[arg0]], {{.*}} : ui64, vector<16xindex> -> !xegpu.tensor_desc<16x2xf32, #xegpu.scatter_tdesc_attr<chunk_size = 2 : i64>>330 // CHECK: arith.addi %{{.*}}, %{{.*}} : vector<16xindex>331 // CHECK: xegpu.create_tdesc [[arg0]], {{.*}} : ui64, vector<16xindex> -> !xegpu.tensor_desc<16x2xf32, #xegpu.scatter_tdesc_attr<chunk_size = 2 : i64>>332 // CHECK: arith.addi %{{.*}}, %{{.*}} : vector<16xindex>333 // CHECK: xegpu.create_tdesc [[arg0]], {{.*}} : ui64, vector<16xindex> -> !xegpu.tensor_desc<16x2xf32, #xegpu.scatter_tdesc_attr<chunk_size = 2 : i64>>334 gpu.func @create_tdesc_step_chunk3(%src: ui64) -> !xegpu.tensor_desc<16x8xf32, #xegpu.scatter_tdesc_attr<chunk_size=8>, #xegpu.layout<inst_data = [16, 2]>> {335 %step = arith.constant dense<8> : vector<16xindex>336 %seq = vector.step : vector<16xindex>337 %cst = arith.muli %seq, %step : vector<16xindex>338 %tdesc = xegpu.create_tdesc %src, %cst : ui64, vector<16xindex> -> !xegpu.tensor_desc<16x8xf32, #xegpu.scatter_tdesc_attr<chunk_size=8>, #xegpu.layout<inst_data = [16, 2]>>339 gpu.return %tdesc : !xegpu.tensor_desc<16x8xf32, #xegpu.scatter_tdesc_attr<chunk_size=8>, #xegpu.layout<inst_data = [16, 2]>>340 }341 342//-----343 // CHECK-LABEL: load_chunk344 // CHECK-SAME: [[arg0:%.+]]: ui64345 // CHECK-COUNT-4: xegpu.create_tdesc [[arg0]], {{.*}} : ui64, vector<16xindex> -> !xegpu.tensor_desc<16x2xf32, #xegpu.scatter_tdesc_attr<chunk_size = 2 : i64>>346 // CHECK-COUNT-4: xegpu.load {{.*}} <{l1_hint = #xegpu.cache_hint<cached>, l2_hint = #xegpu.cache_hint<uncached>}> : !xegpu.tensor_desc<16x2xf32, #xegpu.scatter_tdesc_attr<chunk_size = 2 : i64>>, vector<16xi1> -> vector<16x2xf32>347 348 gpu.func @load_chunk(%src: ui64) -> vector<32x4xf32> {349 %cst = arith.constant dense<[350 0, 8, 16, 24, 32, 40, 48, 56,351 64, 72, 80, 88, 96, 104, 112, 120,352 128, 136, 144, 152, 160, 168, 176, 184,353 192, 200, 208, 216, 224, 232, 240, 248354 ]> : vector<32xindex>355 356 %c17 = arith.constant 17: index357 %mask = vector.create_mask %c17: vector<32xi1>358 359 %tdesc = xegpu.create_tdesc %src, %cst : ui64, vector<32xindex> -> !xegpu.tensor_desc<32x4xf32, #xegpu.scatter_tdesc_attr<chunk_size=4>, #xegpu.layout<inst_data = [16, 2]>>360 %ld = xegpu.load %tdesc, %mask <{l1_hint = #xegpu.cache_hint<cached>, l2_hint = #xegpu.cache_hint<uncached>}>: !xegpu.tensor_desc<32x4xf32, #xegpu.scatter_tdesc_attr<chunk_size=4>, #xegpu.layout<inst_data = [16, 2]>>, vector<32xi1> -> vector<32x4xf32>361 362 gpu.return %ld : vector<32x4xf32>363 }364 365//-----366 // CHECK-LABEL: load_with_offsets_chunk367 // CHECK-SAME: [[arg0:%.+]]: ui64368 // CHECK: [[cst:%.+]] = arith.constant dense<0.000000e+00> : vector<32x4xf32>369 // CHECK: [[cst0:%.+]] = arith.constant dense<[130, 138, 146, 154, 162, 170, 178, 186, 194, 202, 210, 218, 226, 234, 242, 250]> : vector<16xindex>370 // CHECK: [[cst1:%.+]] = arith.constant dense<[2, 10, 18, 26, 34, 42, 50, 58, 66, 74, 82, 90, 98, 106, 114, 122]> : vector<16xindex>371 // CHECK: [[cst2:%.+]] = arith.constant dense<[128, 136, 144, 152, 160, 168, 176, 184, 192, 200, 208, 216, 224, 232, 240, 248]> : vector<16xindex>372 // CHECK: [[cst3:%.+]] = arith.constant dense<[0, 8, 16, 24, 32, 40, 48, 56, 64, 72, 80, 88, 96, 104, 112, 120]> : vector<16xindex>373 // CHECK-COUNT-4: xegpu.load {{.*}}[{{.*}}], {{.*}} <{chunk_size = 2 : i64, l1_hint = #xegpu.cache_hint<cached>}> : ui64, vector<16xindex>, vector<16xi1> -> vector<16x2xf32>374 gpu.func @load_with_offsets_chunk(%src: ui64) -> vector<32x4xf32> {375 %cst = arith.constant dense<[376 0, 8, 16, 24, 32, 40, 48, 56,377 64, 72, 80, 88, 96, 104, 112, 120,378 128, 136, 144, 152, 160, 168, 176, 184,379 192, 200, 208, 216, 224, 232, 240, 248380 ]> : vector<32xindex>381 382 %c17 = arith.constant 17: index383 %mask = vector.create_mask %c17: vector<32xi1>384 %ld = xegpu.load %src[%cst], %mask {chunk_size = 4, layout_result_0 = #xegpu.layout<inst_data = [16, 2]>, l1_hint = #xegpu.cache_hint<cached>} : ui64, vector<32xindex>, vector<32xi1> -> vector<32x4xf32>385 gpu.return %ld : vector<32x4xf32>386 }387 388//-----389 // CHECK-LABEL: store_chunk390 // CHECK-SAME: [[arg0:%.+]]: ui64391 // CHECK-COUNT-4: xegpu.create_tdesc [[arg0]], {{.*}} : ui64, vector<16xindex> -> !xegpu.tensor_desc<16x2xf32, #xegpu.scatter_tdesc_attr<chunk_size = 2 : i64>>392 // CHECK-COUNT-4: xegpu.store {{.*}} <{l1_hint = #xegpu.cache_hint<cached>, l2_hint = #xegpu.cache_hint<uncached>}> : vector<16x2xf32>, !xegpu.tensor_desc<16x2xf32, #xegpu.scatter_tdesc_attr<chunk_size = 2 : i64>>, vector<16xi1>393 gpu.func @store_chunk(%src: ui64) {394 %cst = arith.constant dense<[395 0, 8, 16, 24, 32, 40, 48, 56,396 64, 72, 80, 88, 96, 104, 112, 120,397 128, 136, 144, 152, 160, 168, 176, 184,398 192, 200, 208, 216, 224, 232, 240, 248399 ]> : vector<32xindex>400 401 %c17 = arith.constant 17: index402 %mask = vector.create_mask %c17: vector<32xi1>403 404 %st_vec = arith.constant dense<1023.>: vector<32x4xf32>405 %tdesc = xegpu.create_tdesc %src, %cst : ui64, vector<32xindex> -> !xegpu.tensor_desc<32x4xf32, #xegpu.scatter_tdesc_attr<chunk_size=4>, #xegpu.layout<inst_data = [16, 2]>>406 xegpu.store %st_vec, %tdesc, %mask <{l1_hint = #xegpu.cache_hint<cached>, l2_hint = #xegpu.cache_hint<uncached>}>: vector<32x4xf32>, !xegpu.tensor_desc<32x4xf32, #xegpu.scatter_tdesc_attr<chunk_size=4>, #xegpu.layout<inst_data = [16,2]>>, vector<32xi1>407 408 gpu.return409 }410 411//-----412 // CHECK-LABEL: store_with_offsets_chunk413 // CHECK-SAME: [[arg0:%.+]]: ui64414 // CHECK: [[cst:%.+]] = arith.constant dense<1.023000e+03> : vector<16x2xf32415 // CHECK: [[cst0:%.+]] = arith.constant dense<[130, 138, 146, 154, 162, 170, 178, 186, 194, 202, 210, 218, 226, 234, 242, 250]> : vector<16xindex>416 // CHECK: [[cst1:%.+]] = arith.constant dense<[2, 10, 18, 26, 34, 42, 50, 58, 66, 74, 82, 90, 98, 106, 114, 122]> : vector<16xindex>417 // CHECK: [[cst2:%.+]] = arith.constant dense<[128, 136, 144, 152, 160, 168, 176, 184, 192, 200, 208, 216, 224, 232, 240, 248]> : vector<16xindex>418 // CHECK: [[cst3:%.+]] = arith.constant dense<[0, 8, 16, 24, 32, 40, 48, 56, 64, 72, 80, 88, 96, 104, 112, 120]> : vector<16xindex>419 // CHECK-COUNT-4: xegpu.store {{.*}}[{{.*}}], {{.*}} <{chunk_size = 2 : i64, l1_hint = #xegpu.cache_hint<cached>}> : vector<16x2xf32>, ui64, vector<16xindex>, vector<16xi1>420 gpu.func @store_with_offsets_chunk(%src: ui64) {421 %cst = arith.constant dense<[422 0, 8, 16, 24, 32, 40, 48, 56,423 64, 72, 80, 88, 96, 104, 112, 120,424 128, 136, 144, 152, 160, 168, 176, 184,425 192, 200, 208, 216, 224, 232, 240, 248426 ]> : vector<32xindex>427 428 %c17 = arith.constant 17: index429 %mask = vector.create_mask %c17: vector<32xi1>430 431 %st_vec = arith.constant dense<1023.>: vector<32x4xf32>432 xegpu.store %st_vec, %src[%cst], %mask {chunk_size = 4, layout = #xegpu.layout<inst_data = [16, 2]>, l1_hint = #xegpu.cache_hint<cached>} : vector<32x4xf32>, ui64, vector<32xindex>, vector<32xi1>433 gpu.return434 }435 436//-----437 // CHECK-LABEL: prefetch_chunk438 // CHECK-SAME: [[arg0:%.+]]: ui64439 // CHECK-COUNT-2: xegpu.create_tdesc [[arg0]], {{.*}} : ui64, vector<16xindex> -> !xegpu.tensor_desc<16x2xf32, #xegpu.scatter_tdesc_attr<chunk_size = 2 : i64>>440 // CHECK-COUNT-2: xegpu.prefetch {{.*}} : !xegpu.tensor_desc<16x2xf32, #xegpu.scatter_tdesc_attr<chunk_size = 2 : i64>>441 gpu.func @prefetch_chunk(%src: ui64) {442 %cst = arith.constant dense<[443 0, 8, 16, 24, 32, 40, 48, 56,444 64, 72, 80, 88, 96, 104, 112, 120,445 128, 136, 144, 152, 160, 168, 176, 184,446 192, 200, 208, 216, 224, 232, 240, 248447 ]> : vector<32xindex>448 %tdesc = xegpu.create_tdesc %src, %cst : ui64, vector<32xindex> -> !xegpu.tensor_desc<32x4xf32, #xegpu.scatter_tdesc_attr<chunk_size=4>, #xegpu.layout<inst_data = [16, 2]>>449 xegpu.prefetch %tdesc: !xegpu.tensor_desc<32x4xf32, #xegpu.scatter_tdesc_attr<chunk_size=4>, #xegpu.layout<inst_data = [16, 2]>>450 451 gpu.return452 }453 454//-----455 // CHECK-LABEL: update_chunk456 // CHECK-SAME: [[arg0:%.+]]: ui64457 // CHECK-COUNT-4: xegpu.create_tdesc [[arg0]], {{.*}} : ui64, vector<16xindex> -> !xegpu.tensor_desc<16x2xf32, #xegpu.scatter_tdesc_attr<chunk_size = 2 : i64>>458 // CHECK-COUNT-4: xegpu.update_offset {{.*}} : !xegpu.tensor_desc<16x2xf32, #xegpu.scatter_tdesc_attr<chunk_size = 2 : i64>>, vector<16xindex>459 gpu.func @update_chunk(%src: ui64) -> !xegpu.tensor_desc<32x4xf32, #xegpu.scatter_tdesc_attr<chunk_size=4>, #xegpu.layout<inst_data = [16, 2]>> {460 %cst = arith.constant dense<[461 0, 8, 16, 24, 32, 40, 48, 56,462 64, 72, 80, 88, 96, 104, 112, 120,463 128, 136, 144, 152, 160, 168, 176, 184,464 192, 200, 208, 216, 224, 232, 240, 248465 ]> : vector<32xindex>466 %delta = arith.constant dense<32>: vector<32xindex>467 %tdesc = xegpu.create_tdesc %src, %cst : ui64, vector<32xindex> -> !xegpu.tensor_desc<32x4xf32, #xegpu.scatter_tdesc_attr<chunk_size=4>, #xegpu.layout<inst_data = [16, 2]>>468 469 %new_tdesc = xegpu.update_offset %tdesc, %delta470 : !xegpu.tensor_desc<32x4xf32, #xegpu.scatter_tdesc_attr<chunk_size=4>, #xegpu.layout<inst_data = [16, 2]>>, vector<32xindex>471 472 gpu.return %new_tdesc : !xegpu.tensor_desc<32x4xf32, #xegpu.scatter_tdesc_attr<chunk_size=4>, #xegpu.layout<inst_data = [16, 2]>>473 }474}475 476