gpu: jit: conv: add restriction for a/b subtiles

bviyer · Dec 1, 2022 · b36c74f · b36c74f
1 parent c5bfd3a
commit b36c74f
Show file tree

Hide file tree

Showing 2 changed files with 12 additions and 0 deletions.
diff --git a/src/gpu/jit/conv/config.cpp b/src/gpu/jit/conv/config.cpp
@@ -1954,6 +1954,7 @@ bool try_reduce_grf_usage(conv_config_t &cfg) {
         // XXX: avoid layout mismatch for B loads
         if (cfg.hw() >= ngen::HW::XeHPC && prb.is_bwd_w)
             max_b_subtiles = std::min(2, max_b_subtiles);
+        if (h.prb_iter_ndims('n') > 1) max_b_subtiles = 1;
         while (cfg.subtiles().b() < max_b_subtiles) {
             cfg.subtiles().set_b(cfg.subtiles().b() * 2);
             int est_regs = estimate_register_count(cfg);
@@ -1964,6 +1965,7 @@ bool try_reduce_grf_usage(conv_config_t &cfg) {
         int m_iter_blk = h.iter_dim('m');
         int max_a_subtiles = std::min((cfg.slm().a() ? 4 : 2), m_iter_blk / 8);
         if (cfg.subtiles().b() > 1) max_a_subtiles = 1;
+        if (h.prb_iter_ndims('m') > 1) max_a_subtiles = 1;
         while (cfg.subtiles().a() < max_a_subtiles) {
             cfg.subtiles().set_a(cfg.subtiles().a() * 2);
             int est_regs = estimate_register_count(cfg);

diff --git a/src/gpu/jit/conv/config.hpp b/src/gpu/jit/conv/config.hpp
@@ -1258,6 +1258,16 @@ class bmnk_dim_helper_t {
         return ret;
     }
 
+    int prb_iter_ndims(char bmnk) const {
+        int ret = 0;
+        for (auto &kv : cfg_.iter_dims().get()) {
+            if (to_bmnk(kv.first) != bmnk) continue;
+            if (kv.second == 1) continue;
+            ret++;
+        }
+        return ret;
+    }
+
     int thread_group_dim(char bmnk) const {
         int ret = 1;
         for (auto &kv : cfg_.thread_group_dims().get()) {