gpu: jit: conv: fix SLM message check

oneapi-src · Nov 2, 2022 · 365ac20 · 365ac20
1 parent 9d4bf94
commit 365ac20
Showing 1 changed file with 2 additions and 1 deletion.
diff --git a/src/gpu/jit/conv/config.hpp b/src/gpu/jit/conv/config.hpp
@@ -1021,7 +1021,8 @@ class conv_config_t : public conv_problem_t {
             //Check that SLM can be stored with oword messages.
             int tg_size = tg_grid_dim[0] * tg_grid_dim[1] * tg_grid_dim[2];
             int bytes_per_tg = (m_tg_blk * k_blk * a_data_type_size);
-            bool can_split_a = bytes_per_tg % 16 == 0
+            int align = is_bwd_w ? 32 : 16;
+            bool can_split_a = bytes_per_tg % align == 0
                     && bytes_per_tg / tg_size >= k_blk && k_blk % 2 == 0;
             use_a_slm = (tg_grid_dim[0] > 1) && can_split_a;
         }