XeTLA v0.3.6
Intel® Xe Templates for Linear Algebra - API Definition Document
 
Loading...
Searching...
No Matches
store_xe.hpp
Go to the documentation of this file.
1/*******************************************************************************
2* Copyright (c) 2022-2023 Intel Corporation
3*
4* Licensed under the Apache License, Version 2.0 (the "License");
5* you may not use this file except in compliance with the License.
6* You may obtain a copy of the License at
7*
8* http://www.apache.org/licenses/LICENSE-2.0
9*
10* Unless required by applicable law or agreed to in writing, software
11* distributed under the License is distributed on an "AS IS" BASIS,
12* WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
13* See the License for the specific language governing permissions and
14* limitations under the License.
15*******************************************************************************/
16
19
20#pragma once
21
22#include "subgroup/tile/api.hpp"
23#include "subgroup/tile/impl/op_function.hpp"
24#include "subgroup/tile/impl/payload_xe.hpp"
25
26namespace gpu::xetla::subgroup {
27
28namespace detail {
29template <typename tile_t, typename payload_t>
30struct check_store_type {
31 static constexpr bool is_global_2d_xe
32 = (payload_t::memory_space == mem_space::global
33 && (payload_t::message_type == msg_type::block_2d)
34 && (payload_t::arch_tag == gpu_arch::Xe));
35
36 static constexpr bool is_global_block_1d_xe
37 = ((payload_t::memory_space == mem_space::global)
38 && (tile_t::tile_size_y == 1) && (tile_t::block_size_y == 1)
39 && (payload_t::message_type == msg_type::block_1d)
40 && (payload_t::arch_tag == gpu_arch::Xe));
41
42 static constexpr bool is_global_unaligned_2d_xe
43 = (payload_t::memory_space == mem_space::global
44 && (payload_t::message_type == msg_type::unaligned_2d)
45 && (payload_t::arch_tag == gpu_arch::Xe));
46
47 static constexpr bool is_global_atomic_xe
48 = ((payload_t::memory_space == mem_space::global)
49 && (payload_t::message_type == msg_type::atomic_add)
50 && (payload_t::arch_tag == gpu_arch::Xe));
51
52 static constexpr bool is_local_scatter_xe = ((payload_t::memory_space
53 == mem_space::local)
54 && (payload_t::message_type == msg_type::scatter)
55 && (payload_t::arch_tag == gpu_arch::Xe)
56 && (payload_t::tile_desc::register_layout == reg_layout::tiled
57 || payload_t::tile_desc::register_layout
58 == reg_layout::vnni_tiled));
59
60 static constexpr bool is_local_scatter_vnni_col_xe
61 = ((payload_t::memory_space == mem_space::local)
62 && (payload_t::message_type == msg_type::scatter)
63 && (payload_t::arch_tag == gpu_arch::Xe)
64 && (payload_t::tile_desc::register_layout
65 == reg_layout::vnni_tiled_col_major));
66
67 static constexpr bool is_local_block_1d_xe = ((payload_t::memory_space
68 == mem_space::local)
69 && (payload_t::message_type == msg_type::block_1d)
70 && (payload_t::arch_tag == gpu_arch::Xe)
71 && (payload_t::tile_desc::register_layout == reg_layout::tiled));
72};
73
74} // namespace detail
75
86template <cache_hint L1 = cache_hint::write_back,
87 cache_hint L2 = cache_hint::write_back, typename tile_t,
88 typename payload_t>
89__XETLA_API typename std::enable_if_t<
90 detail::check_store_type<tile_t, payload_t>::is_global_2d_xe>
91tile_store(tile_t &tile, payload_t &payload) {
92 using dtype = typename tile_t::dtype;
93 using tile_desc = typename tile_t::tile_desc;
94
95 static constexpr uint32_t tile_size_x = tile_desc::tile_size_x;
96 static constexpr uint32_t block_size_x = tile_desc::block_size_x;
97 static constexpr uint32_t block_size_y = tile_desc::block_size_y;
98 static constexpr uint32_t remained_size_y = tile_desc::remained_size_y;
99
100 static constexpr uint32_t block_elems = tile_desc::block_elems;
101
102 static constexpr uint32_t num_block_x = tile_desc::num_block_x;
103 static constexpr uint32_t num_block_y = tile_desc::num_block_y;
104 static constexpr uint32_t num_block = tile_desc::num_block;
105
106 using load_store_attr = typename arch_attr_t<
107 payload_t::arch_tag>::template load_store_attr<msg_type::block_2d>;
108
109 static constexpr int32_t max_block_width
110 = load_store_attr::max_load_width_in_bytes / sizeof(dtype);
111 static constexpr int32_t max_store_block_height
112 = load_store_attr::max_store_height_in_elem;
113 static_assert((max_block_width % block_size_x) == 0,
114 "max_block_width should be a multiply of block size x.");
115 static constexpr uint32_t elems_per_CL
116 = load_store_attr::cache_line_size_in_bytes / sizeof(dtype);
117 static constexpr uint32_t st_blk_size_y
118 = block_size_y > max_store_block_height ? max_store_block_height
119 : block_size_y;
120 // to make sure full CL store
121 static constexpr uint32_t st_block_x
122 = (((tile_size_x % elems_per_CL) == 0) ? elems_per_CL
123 : block_size_x);
124
125 static constexpr uint8_t arr_len = st_block_x / block_size_x;
126
127 auto payload_2d = payload.payloads.xetla_format<uint32_t, num_block, 16>();
128#pragma unroll
129 for (uint32_t i = 0; i < num_block_y; ++i) {
130 constexpr uint32_t store_block_elems = block_elems * arr_len;
131 auto payload_row = payload_2d.xetla_select<num_block_x, 1, 16, 1>(
132 i * num_block_x, 0);
133 detail::reset_tile_desc_core<num_block_x, block_size_x * arr_len,
134 st_blk_size_y, 1, 1, false>(payload_row);
135#pragma unroll
136 for (uint32_t j = 0; j < num_block_x; j += arr_len) {
137 xetla_tdescriptor tdesc = payload_row.row(j);
138 auto reg_blk = tile.reg.xetla_select<store_block_elems, 1>(
139 (i * num_block_x + j) * block_elems);
140 xetla_vector<dtype, store_block_elems> combine_blk;
141 auto combine_blk_2d = combine_blk.xetla_format<native_type_t<dtype>,
142 block_size_y, block_size_x * arr_len>();
143#pragma unroll
144 for (uint32_t combine_i = 0; combine_i < arr_len; ++combine_i) {
145 combine_blk_2d.xetla_select<block_size_y, 1, block_size_x, 1>(
146 0, combine_i * block_size_x)
147 = reg_blk.xetla_select<block_elems, 1>(
148 combine_i * block_elems);
149 }
150#pragma unroll
151 for (uint32_t ii = 0; ii < block_size_y / st_blk_size_y; ++ii) {
152 constexpr uint32_t store_elems
153 = st_blk_size_y * block_size_x * arr_len;
154 auto st_blk = combine_blk.xetla_select<store_elems, 1>(
155 ii * store_elems);
156 xetla_tstore_global<dtype, store_elems, L1, L2,
157 payload_t::arch_tag>(tdesc, st_blk);
158 xetla_update_tdesc_offsety(
159 tdesc.xetla_format<uint32_t>(), st_blk_size_y);
160 }
161 // exceed hardware limitation
162 if constexpr ((block_size_y % st_blk_size_y) != 0) {
163 constexpr uint32_t blk_remained_start = block_size_y
164 / st_blk_size_y * st_blk_size_y * block_size_x
165 * arr_len;
166 constexpr uint8_t blk_remained_y = block_size_y % st_blk_size_y;
167 constexpr uint8_t blk_remained_elems
168 = blk_remained_y * block_size_x * arr_len;
169 auto st_blk = combine_blk.xetla_select<blk_remained_elems, 1>(
170 blk_remained_start);
171 constexpr uint32_t block_widthx_widthy_arrlen
172 = (block_size_x * arr_len - 1)
173 | ((blk_remained_y - 1) << 8);
174 gpu::xetla::detail::xetla_set_block_widthx_widthy_arrlen(
175 tdesc.xetla_format<uint32_t>(),
176 block_widthx_widthy_arrlen);
177 xetla_tstore_global<dtype, blk_remained_elems, L1, L2,
178 payload_t::arch_tag>(tdesc, st_blk);
179 }
180 }
181 }
182 // process tail
183 if constexpr (remained_size_y > 0) {
184 constexpr uint32_t remained_block_elems
185 = block_size_x * remained_size_y;
186 constexpr uint32_t processed_elems
187 = num_block_y * num_block_x * block_elems;
188 constexpr uint32_t remained_st_blk_size_y
189 = st_blk_size_y > remained_size_y ? remained_size_y
190 : st_blk_size_y;
191 auto payload_row = payload_2d.xetla_select<num_block_x, 1, 16, 1>(
192 num_block_y * num_block_x, 0);
193 detail::reset_tile_desc_core<num_block_x, block_size_x * arr_len,
194 remained_st_blk_size_y, 1, 1, false>(payload_row);
195#pragma unroll
196 for (uint32_t j = 0; j < num_block_x; j += arr_len) {
197 xetla_tdescriptor tdesc = payload_row.row(j);
198 auto reg_blk
199 = tile.reg.xetla_select<remained_block_elems * arr_len, 1>(
200 processed_elems + j * remained_block_elems);
201 // Do combination
202 xetla_vector<dtype, remained_block_elems * arr_len> combine_blk;
203 auto combine_blk_2d = combine_blk.xetla_format<native_type_t<dtype>,
204 remained_size_y, block_size_x * arr_len>();
205#pragma unroll
206 for (uint32_t combine_i = 0; combine_i < arr_len; ++combine_i) {
207 combine_blk_2d
208 .xetla_select<remained_size_y, 1, block_size_x, 1>(
209 0, combine_i * block_size_x)
210 = reg_blk.xetla_select<remained_block_elems, 1>(
211 combine_i * remained_block_elems);
212 }
213#pragma unroll
214 for (uint32_t ii = 0; ii < remained_size_y / remained_st_blk_size_y;
215 ++ii) {
216 constexpr uint32_t store_elems
217 = remained_st_blk_size_y * block_size_x * arr_len;
218 auto st_blk = combine_blk.xetla_select<store_elems, 1>(
219 ii * store_elems);
220 xetla_tstore_global<dtype, store_elems, L1, L2,
221 payload_t::arch_tag>(tdesc, st_blk);
222 xetla_update_tdesc_offsety(
223 tdesc.xetla_format<uint32_t>(), remained_st_blk_size_y);
224 }
225 constexpr uint32_t final_st_blk_size_y
226 = remained_size_y % remained_st_blk_size_y;
227 if constexpr (final_st_blk_size_y != 0) {
228 constexpr uint32_t final_start = remained_size_y
229 / remained_st_blk_size_y * remained_st_blk_size_y
230 * block_size_x * arr_len;
231 constexpr uint32_t final_store_elems
232 = final_st_blk_size_y * block_size_x * arr_len;
233 auto st_blk = combine_blk.xetla_select<final_store_elems, 1>(
234 final_start);
235 constexpr uint32_t block_widthx_widthy_arrlen
236 = (block_size_x * arr_len - 1)
237 | ((final_st_blk_size_y - 1) << 8);
238 gpu::xetla::detail::xetla_set_block_widthx_widthy_arrlen(
239 tdesc.xetla_format<uint32_t>(),
240 block_widthx_widthy_arrlen);
241 xetla_tstore_global<dtype, final_store_elems, L1, L2,
242 payload_t::arch_tag>(tdesc, st_blk);
243 }
244 }
245 }
246}
247
258template <cache_hint L1 = cache_hint::write_back,
259 cache_hint L2 = cache_hint::write_back, typename tile_t,
260 typename payload_t>
261__XETLA_API typename std::enable_if_t<
262 detail::check_store_type<tile_t, payload_t>::is_global_block_1d_xe>
263tile_store(tile_t &tile, payload_t &payload) {
264 using dtype = typename tile_t::dtype;
265 using store_dtype = typename payload_t::mem_dtype;
266
267 static constexpr uint32_t tile_size_x = tile_t::tile_size_x;
268 static constexpr uint32_t scale_factor = payload_t::scale_factor;
269
270 constexpr uint32_t store_len = tile_size_x / scale_factor;
271 if constexpr (store_len >= 64) {
272#pragma unroll
273 for (uint32_t i = 0; i < store_len / 64; i++) {
274 uint32_t offset_x = i * 64 * scale_factor;
275 auto reg_sub
276 = tile.reg.xetla_select<64 * scale_factor, 1>(offset_x);
277 uint32_t address_offset = offset_x * sizeof(dtype);
278
279 xetla_store_global<store_dtype, 64, data_size::default_size, L1,
280 L2>(payload.base_ptr, payload.base_offset + address_offset,
281 reg_sub.xetla_format<store_dtype>());
282 }
283 }
284 constexpr uint32_t tail_len = store_len % 64;
285 uint32_t tail_offset = store_len / 64 * 64 * scale_factor;
286 detail::process_1d_tail<tail_len, 32, detail::process_flag::store, L1, L2>(
287 tile, payload, tail_offset);
288}
289
301template <cache_hint L1 = cache_hint::write_back,
302 cache_hint L3 = cache_hint::write_back, typename tile_t,
303 typename payload_t,
304 typename oob_check_tag = global_atomic_oob_check_on_tag>
305__XETLA_API typename std::enable_if_t<
306 detail::check_store_type<tile_t, payload_t>::is_global_unaligned_2d_xe>
307tile_store(tile_t &tile, payload_t &payload,
308 [[maybe_unused]] oob_check_tag tag = {}) {
309 constexpr bool oob_check = std::is_same<oob_check_tag,
310 global_atomic_oob_check_on_tag>::value;
311 using dtype = typename payload_t::dtype;
312 using tile_desc = typename payload_t::tile_desc;
313 using store_dtype = typename payload_t::mem_dtype;
314 constexpr uint32_t num_channel_y = payload_t::num_channel_y;
315 constexpr uint32_t load_elems = num_channel_y * payload_t::num_channel_x;
316 constexpr uint32_t scale_factor = payload_t::scale_factor;
317
318#pragma unroll
319 for (uint32_t i = 0; i < tile_desc::tile_size_y / tile_desc::block_size_y;
320 i++) {
321 uint32_t offset_y = i * tile_desc::block_size_y;
322#pragma unroll
323 for (uint32_t j = 0; j < tile_desc::num_block_x; j++) {
324 uint32_t offset_x = j * tile_desc::block_size_x;
325 auto reg_sub = tile.reg.xetla_select<tile_desc::block_elems, 1>(
326 (i * tile_desc::num_block_x + j) * tile_desc::block_elems);
327 xetla_mask<load_elems> pred_x = oob_check
328 ? payload.step_x + payload.base_x + offset_x
329 < payload.width_in_elems
330 : 1;
331#pragma unroll
332 for (uint32_t sub_block_y = 0;
333 sub_block_y < tile_desc::block_size_y;
334 sub_block_y += num_channel_y) {
335 xetla_mask<load_elems> pred_y = oob_check
336 ? payload.step_y + payload.base_y + offset_y
337 + sub_block_y
338 < payload.height_in_elems
339 : 1;
340
341 uint32_t address_offset = offset_x * sizeof(dtype)
342 + (offset_y + sub_block_y) * payload.pitch_in_bytes;
343 xetla_store_global<store_dtype, 1, data_size::default_size, L1,
344 L3, load_elems>(payload.base_ptr,
345 (payload.base_offset + address_offset
346 + payload.channel_offset),
347 reg_sub.xetla_select<load_elems * scale_factor, 1>(
348 sub_block_y * tile_desc::block_size_x)
349 .xetla_format<store_dtype>(),
350 (pred_x && pred_y));
351 }
352 }
353 }
354 // process the tail
355 if constexpr ((tile_desc::tile_size_y % tile_desc::block_size_y) != 0) {
356 constexpr uint32_t remained_size_y = tile_desc::remained_size_y;
357 constexpr uint32_t offset_y = tile_desc::tile_size_y - remained_size_y;
358 constexpr uint32_t processed_elems = offset_y * tile_desc::tile_size_x;
359 constexpr uint32_t remain_block_elems
360 = remained_size_y * tile_desc::block_size_x;
361#pragma unroll
362 for (uint32_t j = 0; j < tile_desc::num_block_x; j++) {
363 uint32_t offset_x = j * tile_desc::block_size_x;
364 auto reg_sub = tile.reg.xetla_select<remain_block_elems, 1>(
365 processed_elems + j * remain_block_elems);
366 xetla_mask<load_elems> pred_x = oob_check
367 ? payload.step_x + payload.base_x + offset_x
368 < payload.width_in_elems
369 : 1;
370#pragma unroll
371 for (uint32_t sub_block_y = 0; sub_block_y < remained_size_y;
372 sub_block_y += num_channel_y) {
373 xetla_mask<load_elems> pred_y = oob_check
374 ? payload.step_y + payload.base_y + offset_y
375 + sub_block_y
376 < payload.height_in_elems
377 : 1;
378
379 uint32_t address_offset = offset_x * sizeof(dtype)
380 + (offset_y + sub_block_y) * payload.pitch_in_bytes;
381 xetla_store_global<store_dtype, 1, data_size::default_size, L1,
382 L3, load_elems>(payload.base_ptr,
383 (payload.base_offset + address_offset
384 + payload.channel_offset),
385 reg_sub.xetla_select<load_elems * scale_factor, 1>(
386 sub_block_y * tile_desc::block_size_x)
387 .xetla_format<store_dtype>(),
388 (pred_x && pred_y));
389 }
390 }
391 }
392}
393
405template <cache_hint L1 = cache_hint::uncached,
406 cache_hint L2 = cache_hint::write_back, typename tile_t,
407 typename payload_t,
408 typename oob_check_tag = global_atomic_oob_check_on_tag>
409__XETLA_API typename std::enable_if_t<
410 detail::check_store_type<tile_t, payload_t>::is_global_atomic_xe>
411tile_store(tile_t &tile, payload_t &payload,
412 [[maybe_unused]] oob_check_tag tag = {}) {
413 constexpr bool oob_check = std::is_same<oob_check_tag,
414 global_atomic_oob_check_on_tag>::value;
415 using dtype = typename tile_t::dtype;
416 using tile_desc = typename tile_t::tile_desc;
417
418 static constexpr uint32_t tile_size_x = tile_desc::tile_size_x;
419 static constexpr uint32_t tile_size_y = tile_desc::tile_size_y;
420 static constexpr uint32_t block_size_x = tile_desc::block_size_x;
421 static constexpr uint32_t block_size_y = tile_desc::block_size_y;
422 static constexpr uint32_t block_elems = tile_desc::block_elems;
423 static constexpr uint32_t num_block_x = tile_desc::num_block_x;
424
425 static constexpr atomic_op op_kind
426 = (std::is_same<remove_const_t<dtype>, float>::value
427 || std::is_same<remove_const_t<dtype>, double>::value)
428 ? atomic_op::fadd
429 : atomic_op::iadd;
430
431#pragma unroll
432 for (uint32_t i = 0; i < tile_size_y / block_size_y; i++) {
433 uint32_t offset_y = i * block_size_y;
434#pragma unroll
435 for (uint32_t j = 0; j < num_block_x; j++) {
436 uint32_t offset_x = j * block_size_x;
437 auto reg_sub = tile.reg.xetla_select<block_elems, 1>(
438 (i * num_block_x + j) * block_elems);
439 xetla_mask<payload_t::num_channel> pred_x = oob_check
440 ? (payload.step_x + offset_x + payload.base_x)
441 < payload.width_in_elems
442 : 1;
443#pragma unroll
444 for (uint32_t sub_block_y = 0; sub_block_y < block_size_y;
445 sub_block_y += payload_t::num_channel_y) {
446 xetla_mask<payload_t::num_channel> pred_y = oob_check
447 ? (payload.step_y + offset_y + payload.base_y
448 + sub_block_y)
449 < payload.height_in_elems
450 : 1;
451 uint64_t address_offset = offset_x * sizeof(dtype)
452 + (sub_block_y + offset_y) * payload.pitch_in_bytes;
453
454 xetla_tatomic_store_global<dtype, payload_t::num_channel, L1,
455 L2, op_kind, payload_t::arch_tag>(
456 payload.base_pointer + address_offset,
457 payload.channel_offset,
458 reg_sub.xetla_select<payload_t::store_elems, 1>(
459 sub_block_y * block_size_x),
460 pred_x & pred_y);
461 }
462 }
463 }
464 // process the tail
465 if constexpr ((tile_size_y % block_size_y) != 0) {
466 constexpr uint32_t remained_size_y = tile_desc::remained_size_y;
467 constexpr uint32_t offset_y = tile_size_y - remained_size_y;
468 constexpr uint32_t processed_elems = offset_y * tile_size_x;
469 constexpr uint32_t remain_block_elems = remained_size_y * block_size_x;
470#pragma unroll
471 for (uint32_t j = 0; j < num_block_x; j++) {
472 uint32_t offset_x = j * block_size_x;
473 auto reg_sub = tile.reg.xetla_select<remain_block_elems, 1>(
474 processed_elems + j * remain_block_elems);
475 xetla_mask<payload_t::num_channel> pred_x = oob_check
476 ? (payload.step_x + offset_x + payload.base_x)
477 < payload.width_in_elems
478 : 1;
479#pragma unroll
480 for (uint32_t sub_block_y = 0; sub_block_y < remained_size_y;
481 sub_block_y += payload_t::num_channel_y) {
482 xetla_mask<payload_t::num_channel> pred_y = oob_check
483 ? (payload.step_y + offset_y + payload.base_y
484 + sub_block_y)
485 < payload.height_in_elems
486 : 1;
487 uint64_t address_offset = offset_x * sizeof(dtype)
488 + (sub_block_y + offset_y) * payload.pitch_in_bytes;
489
490 xetla_tatomic_store_global<dtype, payload_t::num_channel, L1,
491 L2, op_kind, payload_t::arch_tag>(
492 (uint64_t)payload.base_pointer + address_offset,
493 payload.channel_offset,
494 reg_sub.xetla_select<payload_t::store_elems, 1>(
495 sub_block_y * block_size_x),
496 pred_x & pred_y);
497 }
498 }
499 }
500}
501
513template <cache_hint L1 = cache_hint::write_back,
514 cache_hint L2 = cache_hint::write_back, typename tile_t,
515 typename payload_t>
516__XETLA_API typename std::enable_if_t<
517 detail::check_store_type<tile_t, payload_t>::is_local_scatter_xe>
518tile_store(tile_t &tile, payload_t &payload) {
519 using dtype = typename tile_t::dtype;
520 using tile_desc = typename tile_t::tile_desc;
521 using store_dtype = typename payload_t::mem_dtype;
522
523 constexpr uint32_t num_channel_y = payload_t::num_channel_y;
524 constexpr uint32_t store_elems = num_channel_y * tile_desc::block_size_x;
525#pragma unroll
526 for (uint32_t i = 0; i < tile_desc::tile_size_y / tile_desc::block_size_y;
527 i++) {
528 uint32_t offset_y = i * tile_desc::block_size_y;
529#pragma unroll
530 for (uint32_t j = 0; j < tile_desc::num_block_x; j++) {
531 uint32_t offset_x = j * tile_desc::block_size_x;
532 auto reg_sub = tile.reg.xetla_select<tile_desc::block_elems, 1>(
533 (i * tile_desc::num_block_x + j) * tile_desc::block_elems);
534#pragma unroll
535 for (uint32_t sub_block_y = 0;
536 sub_block_y < tile_desc::block_size_y;
537 sub_block_y += num_channel_y) {
538 uint32_t address_offset = offset_x * sizeof(dtype)
539 + (sub_block_y + offset_y) * payload.pitch_in_bytes;
540 xetla_store_local<store_dtype>(payload.address + address_offset,
541 reg_sub.xetla_select<store_elems, 1>(
542 sub_block_y * tile_desc::block_size_x)
543 .xetla_format<store_dtype>());
544 }
545 }
546 }
547 // process the tail
548 if constexpr ((tile_desc::tile_size_y % tile_desc::block_size_y) != 0) {
549 constexpr uint32_t remained_size_y = tile_desc::remained_size_y;
550 constexpr uint32_t offset_y = tile_desc::block_size_y - remained_size_y;
551 constexpr uint32_t processed_elems = offset_y * tile_desc::tile_size_x;
552 constexpr uint32_t remain_block_elems
553 = remained_size_y * tile_desc::block_size_x;
554#pragma unroll
555 for (uint32_t j = 0; j < tile_desc::num_block_x; j++) {
556 uint32_t offset_x = j * tile_desc::block_size_x;
557 auto reg_sub = tile.reg.xetla_select<remain_block_elems, 1>(
558 processed_elems + j * remain_block_elems);
559#pragma unroll
560 for (uint32_t sub_block_y = 0; sub_block_y < remained_size_y;
561 sub_block_y += num_channel_y) {
562 uint32_t address_offset = offset_x * sizeof(dtype)
563 + (sub_block_y + offset_y) * payload.pitch_in_bytes;
564 xetla_store_local<store_dtype>(payload.address + address_offset,
565 reg_sub.xetla_select<store_elems, 1>(
566 sub_block_y * tile_desc::block_size_x)
567 .xetla_format<store_dtype>());
568 }
569 }
570 }
571}
572
584template <cache_hint L1 = cache_hint::write_back,
585 cache_hint L2 = cache_hint::write_back, typename tile_t,
586 typename payload_t>
587__XETLA_API typename std::enable_if_t<detail::check_store_type<tile_t,
588 payload_t>::is_local_scatter_vnni_col_xe>
589tile_store(tile_t &tile, payload_t &payload) {
590 using dtype = typename tile_t::dtype;
591 using tile_desc = typename tile_t::tile_desc;
592 using store_dtype = typename payload_t::store_dtype;
593
594 constexpr uint32_t vnni_scale_factor = payload_t::vnni_scale_factor;
595 constexpr uint32_t num_vector_size = payload_t::num_vector_size;
596 constexpr uint32_t num_channel_y = payload_t::num_channel_y;
597 constexpr uint32_t store_elems = payload_t::store_elems;
598#pragma unroll
599 for (uint32_t i = 0; i < tile_desc::tile_size_y / tile_desc::block_size_y;
600 i++) {
601 uint32_t offset_y = i * tile_desc::block_size_y;
602#pragma unroll
603 for (uint32_t j = 0; j < tile_desc::num_block_x; j++) {
604 uint32_t offset_x = j * tile_desc::block_size_x;
605 auto reg_sub = tile.reg.xetla_select<tile_desc::block_elems, 1>(
606 (i * tile_desc::num_block_x + j) * tile_desc::block_elems);
607#pragma unroll
608 for (uint32_t sub_block_y = 0;
609 sub_block_y < tile_desc::block_size_y; sub_block_y
610 += num_channel_y * num_vector_size * vnni_scale_factor) {
611 uint32_t address_offset = offset_x * payload.pitch_in_bytes
612 + (sub_block_y + offset_y) * sizeof(dtype);
613 xetla_store_local<store_dtype, num_vector_size>(
614 payload.address + address_offset,
615 reg_sub.xetla_select<store_elems, 1>(
616 sub_block_y * tile_desc::block_size_x)
617 .xetla_format<store_dtype>());
618 }
619 }
620 }
621 // process the tail
622 if constexpr ((tile_desc::tile_size_y % tile_desc::block_size_y) != 0) {
623 constexpr uint32_t remained_size_y = tile_desc::remained_size_y;
624 constexpr uint32_t offset_y = tile_desc::tile_size_y - remained_size_y;
625 constexpr uint32_t processed_elems = offset_y * tile_desc::tile_size_x;
626 constexpr uint32_t remain_block_elems
627 = remained_size_y * tile_desc::block_size_x;
628#pragma unroll
629 for (uint32_t j = 0; j < tile_desc::num_block_x; j++) {
630 uint32_t offset_x = j * tile_desc::block_size_x;
631 auto reg_sub = tile.reg.xetla_select<remain_block_elems, 1>(
632 processed_elems + j * remain_block_elems);
633#pragma unroll
634 for (uint32_t sub_block_y = 0; sub_block_y < remained_size_y;
635 sub_block_y
636 += num_channel_y * num_vector_size * vnni_scale_factor) {
637 uint32_t address_offset = offset_x * payload.pitch_in_bytes
638 + (sub_block_y + offset_y) * sizeof(dtype);
639 xetla_store_local<store_dtype, num_vector_size>(
640 payload.address + address_offset,
641 reg_sub.xetla_select<store_elems, 1>(
642 sub_block_y * tile_desc::block_size_x)
643 .xetla_format<store_dtype>());
644 }
645 }
646 }
647}
648
660template <cache_hint L1 = cache_hint::write_back,
661 cache_hint L2 = cache_hint::write_back, typename tile_t,
662 typename payload_t>
663__XETLA_API typename std::enable_if_t<
664 detail::check_store_type<tile_t, payload_t>::is_local_block_1d_xe
665 && tile_t::block_size_y != 1>
666tile_store(tile_t &tile, payload_t &payload) {
667 using dtype = typename tile_t::dtype;
668 using tile_desc = typename tile_t::tile_desc;
669 using store_dtype = typename payload_t::mem_dtype;
670
671 constexpr uint32_t vector_size
672 = payload_t::bytes_per_row / sizeof(store_dtype);
673
674#pragma unroll
675 for (uint32_t i = 0; i < tile_desc::tile_size_y / tile_desc::block_size_y;
676 i++) {
677 uint32_t offset_y = i * tile_desc::block_size_y;
678#pragma unroll
679 for (uint32_t j = 0; j < tile_desc::num_block_x; j++) {
680 uint32_t offset_x = j * tile_desc::block_size_x;
681 auto reg_sub = tile.reg.xetla_select<tile_desc::block_elems, 1>(
682 (i * tile_desc::num_block_x + j) * tile_desc::block_elems);
683 auto reg_sub_2d = reg_sub.xetla_format<dtype,
684 tile_desc::block_size_y, tile_desc::block_size_x>();
685 uint32_t address_offset = offset_x * sizeof(dtype)
686 + offset_y * payload.pitch_in_bytes;
687#pragma unroll
688 for (uint32_t row_i = 0; row_i < tile_desc::block_size_y; row_i++) {
689 xetla_store_local<store_dtype, vector_size>(payload.address
690 + address_offset
691 + row_i * payload.pitch_in_bytes,
692 reg_sub_2d.row(row_i).xetla_format<store_dtype>());
693 }
694 }
695 }
696 // process the tail
697 if constexpr ((tile_desc::tile_size_y % tile_desc::block_size_y) != 0) {
698 constexpr uint32_t remained_size_y = tile_desc::remained_size_y;
699 constexpr uint32_t offset_y = tile_desc::tile_size_y - remained_size_y;
700 constexpr uint32_t processed_elems = offset_y * tile_desc::tile_size_x;
701 constexpr uint32_t remain_block_elems
702 = remained_size_y * tile_desc::block_size_x;
703#pragma unroll
704 for (uint32_t j = 0; j < tile_desc::num_block_x; j++) {
705 uint32_t offset_x = j * tile_desc::block_size_x;
706 auto reg_sub = tile.reg.xetla_select<remain_block_elems, 1>(
707 processed_elems + j * remain_block_elems);
708 auto reg_sub_2d = reg_sub.xetla_format<dtype, remained_size_y,
709 tile_desc::block_size_x>();
710 uint32_t address_offset = offset_x * sizeof(dtype)
711 + offset_y * payload.pitch_in_bytes;
712#pragma unroll
713 for (uint32_t row_i = 0; row_i < remained_size_y; row_i++) {
714 xetla_store_local<store_dtype, vector_size>(payload.address
715 + address_offset
716 + row_i * payload.pitch_in_bytes,
717 reg_sub_2d.row(row_i).xetla_format<store_dtype>());
718 }
719 }
720 }
721}
722
734template <cache_hint L1 = cache_hint::write_back,
735 cache_hint L2 = cache_hint::write_back, typename tile_t,
736 typename payload_t>
737__XETLA_API typename std::enable_if_t<
738 detail::check_store_type<tile_t, payload_t>::is_local_block_1d_xe
739 && tile_t::tile_size_y == 1 && tile_t::block_size_y == 1>
740tile_store(tile_t &tile, payload_t &payload) {
741 using dtype = typename tile_t::dtype;
742 using tile_desc = typename payload_t::tile_desc;
743 using store_dtype = typename payload_t::mem_dtype;
744
745 constexpr uint32_t scale_factor = payload_t::scale_factor;
746 constexpr uint32_t store_len = tile_desc::tile_size_x / scale_factor;
747 if constexpr (store_len >= 64) {
748#pragma unroll
749 for (uint32_t j = 0; j < store_len / 64; j++) {
750 uint32_t offset_x = j * 64 * scale_factor;
751 auto reg_sub
752 = tile.reg.xetla_select<64 * scale_factor, 1>(offset_x);
753 uint32_t address_offset = offset_x * sizeof(dtype);
754 xetla_store_local<store_dtype, 64>(payload.address + address_offset,
755 reg_sub.xetla_format<store_dtype>());
756 }
757 }
758 detail::process_1d_tail<store_len % 64, 32, detail::process_flag::store, L1,
759 L2>(tile, payload, store_len / 64 * 64 * scale_factor);
760}
761
762} // namespace gpu::xetla::subgroup
__XETLA_API
#define __XETLA_API
Definition common.hpp:43
gpu::xetla::xetla_tdescriptor
xetla_vector< uint32_t, 16 > xetla_tdescriptor
Description of nd tensor descriptor for load and store.
Definition base_types.hpp:155
gpu::xetla::native_type_t
typename native_type< T >::type native_type_t
Return the native data type of T.
Definition base_types.hpp:106
gpu::xetla::xetla_vector
__ESIMD_NS::simd< native_type_t< Ty >, N > xetla_vector
wrapper for xetla_vector.
Definition base_types.hpp:149
gpu::xetla::xetla_mask
__ESIMD_NS::simd_mask< N > xetla_mask
wrapper for xetla_mask.
Definition base_types.hpp:165
gpu::xetla::xetla_store_global
__XETLA_API void xetla_store_global(Ty *p, xetla_vector< Toffset, N > offsets, xetla_vector< Ty, N *NElts > vals, xetla_mask< N > pred=1)
Stateless scattered store.
Definition memory.hpp:316
gpu::xetla::xetla_tatomic_store_global
__XETLA_API std::enable_if_t< arch_tag==gpu_arch::Xe, void > xetla_tatomic_store_global(uint64_t base_address, xetla_vector< Toffset, N > offset, xetla_vector< Ty, N > data, xetla_mask< N > pred=1)
Tensor atomic store API.
Definition raw_send_load_store.hpp:294
gpu::xetla::xetla_update_tdesc_offsety
__XETLA_API void xetla_update_tdesc_offsety(xetla_tdescriptor_ref tdesc, int32_t doffset_y)
Update the y coordinate in the given tensor descriptor.
Definition raw_send_load_store.hpp:161
gpu::xetla::xetla_tstore_global
__XETLA_API std::enable_if_t< arch_tag==gpu_arch::Xe, void > xetla_tstore_global(xetla_tdescriptor tdesc, xetla_vector< Ty, N > data)
Tensor store API.
Definition raw_send_load_store.hpp:227
gpu::xetla::detail::xetla_set_block_widthx_widthy_arrlen
__XETLA_API void xetla_set_block_widthx_widthy_arrlen(xetla_tdescriptor_ref desc, uint32_t block_widthx_widthy_arrlen)
Definition tensor_descriptor.hpp:79
gpu::xetla::subgroup::detail::process_1d_tail
__XETLA_API std::enable_if_t< base_len==0 > process_1d_tail(tile_t &tile, payload_t &payload, uint32_t offset)
Definition common.hpp:96
gpu::xetla::subgroup
Definition limitation.hpp:457
gpu::xetla::subgroup::tile_store
__XETLA_API std::enable_if_t< detail::check_store_type< tile_t, payload_t >::is_global_2d_xe > tile_store(tile_t &tile, payload_t &payload)
Is the func storing data from register file to global memory.
Definition store_xe.hpp:91
gpu::xetla::cache_hint
cache_hint
L1 or L2 cache hint kinds.
Definition common.hpp:89
gpu::xetla::reg_layout::vnni_tiled_col_major
@ vnni_tiled_col_major
this is vnni tiled format, but for each block, they are stored in col major order
gpu::xetla::fence_scope::tile
@ tile
flush out to the local scope
gpu::xetla::atomic_op
atomic_op
Represents an atomic operation.
Definition common.hpp:142
gpu::xetla::atomic_op::iadd
@ iadd
Atomic signed int add of src1 from memory data and return the old value. see
gpu::xetla::atomic_op::fadd
@ fadd
Atomic float add of src1 from memory data and return the old value. see
op_function.hpp
C++ API.
payload_xe.hpp
C++ API.
gpu::xetla::arch_attr_t
Definition arch_config.hpp:72
gpu::xetla::subgroup::detail::check_store_type
Definition store_xe.hpp:30
gpu::xetla::subgroup::detail::check_store_type::is_global_block_1d_xe
static constexpr bool is_global_block_1d_xe
Definition store_xe.hpp:37
gpu::xetla::subgroup::detail::check_store_type::is_local_block_1d_xe
static constexpr bool is_local_block_1d_xe
Definition store_xe.hpp:67
gpu::xetla::subgroup::detail::check_store_type::is_local_scatter_xe
static constexpr bool is_local_scatter_xe
Definition store_xe.hpp:52
gpu::xetla::subgroup::detail::check_store_type::is_global_atomic_xe
static constexpr bool is_global_atomic_xe
Definition store_xe.hpp:48
gpu::xetla::subgroup::detail::check_store_type::is_global_unaligned_2d_xe
static constexpr bool is_global_unaligned_2d_xe
Definition store_xe.hpp:43
gpu::xetla::subgroup::detail::check_store_type::is_local_scatter_vnni_col_xe
static constexpr bool is_local_scatter_vnni_col_xe
Definition store_xe.hpp:61
gpu::xetla::subgroup::detail::check_store_type::is_global_2d_xe
static constexpr bool is_global_2d_xe
Definition store_xe.hpp:32
gpu::xetla::subgroup::tile_t
Is a struct contains some register file.
Definition api.hpp:99
gpu::xetla::subgroup::tile_t::tile_desc
tile_desc_ tile_desc
Definition api.hpp:101
gpu::xetla::subgroup::tile_t::dtype
dtype_ dtype
Definition api.hpp:100
api.hpp
C++ API.