xe: conv: jit: rework walk order heuristic

Author: rjoursler

src/gpu/intel/conv/jit/config.cpp

@@ -1644,31 +1644,62 @@ walk_order_t compute_walk_order(const config_t &cfg) {
         size_t ab_bytes = get_memory_footprint(cfg, inner, outer);
         if (ab_bytes <= l3_size) grid_inner = std::move(outer);
     }
+
+    auto &w_inner = grid_inner[pvars::ow];
+    auto &h_inner = grid_inner[pvars::oh];
+
+    // Prefer square spatial dimensions to increase cache reuse due to iteration
+    // over kernel spatial dimensions.
+    auto rebalance_hw = [&]() {
+        if (!cfg.prb().is_fwd) return false;
+        if (grid_tile[pvars::oh] % (h_inner * 2)) return false;
+        if (w_inner % 2) return false;
+        if (w_inner < h_inner * 4) return false;
+        return true;
+    };
+
+    while (rebalance_hw()) {
+        w_inner /= 2;
+        h_inner *= 2;
+    }
+
     // Add the blocks in this order:
     // - Step 1. Add grid_inner blocks (fitting L3 cache)
     // - Step 2. Add the remaining M/N blocks
     // - Step 3. Add the remaining B/K blocks
     // Within a step follow the default walk order between dimensions.
     walk_order_t walk_order;
     for (int step = 0; step < 3; step++) {
-        for (auto &b : default_walk_order.blocks()) {
-            switch (step) {
-                case 0:
-                    if (grid_inner.has(b.dim)) {
-                        walk_order.add(b.dim, grid_inner[b.dim], 0);
-                    }
-                    break;
-                case 1:
-                case 2:
-                    dim_t rem = utils::div_up(
-                            grid_tile[b.dim], grid_inner.get(b.dim, 1));
-                    if (rem == 1) continue;
-                    auto bmnk = to_gemm(b.dim, prb);
-                    bool is_bk = utils::one_of(bmnk, pvars::b, pvars::k);
-                    if ((step == 2) != is_bk) continue;
-                    walk_order.add(b.dim, rem, 0);
-                    break;
+        if (step == 0) {
+            // Transpose spatial for better reuse
+            auto blocks = default_walk_order.blocks();
+            for (size_t i = 0; i < blocks.size() - 1; i++) {
+                if (cfg.prb().is_fwd && blocks[i].dim == pvars::ow
+                        && blocks[i + 1].dim == pvars::oh) {
+                    std::swap(blocks[i], blocks[i + 1]);
+                }
+            }
+            for (auto &b : blocks) {
+                if (grid_inner.has(b.dim)) {
+                    walk_order.add(b.dim, grid_inner[b.dim], 0);
+                    printf("step %d: %s: %ld\n", step, b.dim.str().c_str(),
+                            grid_inner[b.dim]);
+                }
             }
+            continue;
+        }
+
+        for (auto &b : default_walk_order.blocks()) {
+            dim_t rem
+                    = utils::div_up(grid_tile[b.dim], grid_inner.get(b.dim, 1));
+            if (rem == 1) continue;
+            auto bmnk = to_gemm(b.dim, prb);
+            bool is_bk = utils::one_of(bmnk, pvars::b, pvars::k);
+            if ((step == 2) != is_bk) continue;
+            walk_order.add(b.dim, rem, 0);
+            printf("step %d: %s: %ld\n", step, b.dim.str().c_str(),
+                    grid_inner[b.dim]);
+            break;
         }
     }
     walk_order.finalize(grid_tile);