gpu: jit: conv: support OIHW/OHWI weights via reorder

oneapi-src · Jan 11, 2023 · 5bd5d52 · 5bd5d52
1 parent 3d6412a
commit 5bd5d52
Showing 1 changed file with 6 additions and 4 deletions.
diff --git a/src/gpu/jit/conv/config.cpp b/src/gpu/jit/conv/config.cpp
@@ -781,10 +781,12 @@ status_t init_tensor_layouts(conv_config_t &cfg, convolution_pd_t *pd) {
         if (!matches_tag(dst_md, dst_tag)) user_dst_tag = "axb";
     }
 
-    // Allow internal reorder from oihw/ohwi to more optimal weights layout.
-    if (allow_wei_reorder) {
-        if (matches_tag(wei_md, "abx")) user_wei_tag = "abx";
-        if (matches_tag(wei_md, "axb")) user_wei_tag = "axb";
+    // Allow internal reorder for plain weights.
+    for (auto *t : {"abx", "axb"}) {
+        if (matches_tag(wei_md, t)) {
+            user_wei_tag = t;
+            break;
+        }
     }
 
     if (user_src_tag.empty()) user_src_tag = src_tag;