fix reg_blk and reg_tmp size misaligned issue (#559)

Co-authored-by: Jacky, Deng <[email protected]>
intel · Jun 14, 2023 · 76c593f · 76c593f
1 parent 971a2ee
commit 76c593f
Show file tree

Hide file tree

Showing 2 changed files with 19 additions and 5 deletions.
diff --git a/include/subgroup/tile/impl/load_xe.hpp b/include/subgroup/tile/impl/load_xe.hpp
@@ -247,7 +247,9 @@ tile_load(tile_t &tile, payload_t &payload) {
             constexpr uint32_t ld_blk_height = reg_transpose
                     ? detail::getNextPowerOf2<ld_blk_size_y>()
                     : ld_blk_size_y;
-            xetla_vector<dtype, ld_blk_height * block_size_x * arr_len> reg_tmp;
+            constexpr uint32_t tmp_size
+                    = ld_blk_height * block_size_x * arr_len;
+            xetla_vector<dtype, tmp_size> reg_tmp;
 #pragma unroll
             for (int ii = 0; ii < block_size_y / ld_blk_size_y; ++ii) {
                 constexpr uint32_t load_elems
@@ -267,7 +269,7 @@ tile_load(tile_t &tile, payload_t &payload) {
                                       .xetla_select<block_size_x, 1,
                                               ld_blk_size_y, 1>(0, 0);
                 } else {
-                    reg_blk = reg_tmp;
+                    reg_blk.xetla_select<tmp_size, 1>(ii * tmp_size) = reg_tmp;
                 }
 
                 if constexpr (mem_transpose) {
@@ -332,7 +334,9 @@ tile_load(tile_t &tile, payload_t &payload) {
             constexpr uint32_t ld_blk_height = reg_transpose
                     ? detail::getNextPowerOf2<remained_ld_blk_size_y>()
                     : remained_ld_blk_size_y;
-            xetla_vector<dtype, ld_blk_height * block_size_x * arr_len> reg_tmp;
+            constexpr uint32_t tmp_size
+                    = ld_blk_height * block_size_x * arr_len;
+            xetla_vector<dtype, tmp_size> reg_tmp;
 #pragma unroll
             for (int ii = 0; ii < remained_size_y / remained_ld_blk_size_y;
                     ++ii) {
@@ -353,7 +357,7 @@ tile_load(tile_t &tile, payload_t &payload) {
                                       .xetla_select<block_size_x, 1,
                                               remained_ld_blk_size_y, 1>(0, 0);
                 } else {
-                    reg_blk = reg_tmp;
+                    reg_blk.xetla_select<tmp_size, 1>(ii * tmp_size) = reg_tmp;
                 }
                 if constexpr (mem_transpose) {
                     xetla_update_tdesc_offsetx(tdesc.xetla_format<uint32_t>(),

diff --git a/tests/unit/tile_load_store/main.cpp b/tests/unit/tile_load_store/main.cpp
@@ -165,7 +165,7 @@ TEST(tile_load_store, esimd) {
             nd_range, result_validate);
 }
 
-TEST(tile_load_transpose_store, esimd) {
+TEST(tile_load_transpose_store_1, esimd) {
     cl::sycl::nd_range<1> nd_range({1}, {1});
     auto result_validate
             = std::bind(tile_load_store_result_validate<int, false, true>, _1,
@@ -175,6 +175,16 @@ TEST(tile_load_transpose_store, esimd) {
                     64>>(nd_range, result_validate);
 }
 
+TEST(tile_load_transpose_store_2, esimd) {
+    cl::sycl::nd_range<1> nd_range({1}, {1});
+    auto result_validate
+            = std::bind(tile_load_store_result_validate<int, false, true>, _1,
+                    _2, _3, 128, 64, 32, 32);
+    kernel_run<int,
+            tile_load_store_func<int, 128, 64, 128, 32, 32, 8, 16, false, true,
+                    64>>(nd_range, result_validate);
+}
+
 TEST(tile_load_transform_store, esimd) {
     cl::sycl::nd_range<1> nd_range({1}, {1});
     auto result_validate