Vulkan: take im2col when the gather is cheap, not only when c_out is large

**Updated twice; this is the third rule.** Both earlier versions regressed a model, so the history is worth stating:

- A `c_out` gate regressed fast_sam on Mali by 18% (GPU 345 -> 405 ms). Its 80x80 convolutions have `c_out` 64, so the gate pushed them onto the direct shader.
- Bounding the gather size regressed fcn_resnet50 by 19% (2938 -> 3505 ms). Its head is a 2048->512 3x3 at 65x65, which wants a 77.9M element gather but reads it back 512 times, so it is exactly the case im2col exists for.

Neither variable works alone because they are **independent reasons** to take im2col. A large gather is fine when `c_out` is large, since the GEMM amortizes it `c_out` times. When `c_out` is small it is only worth paying if the matrix is small outright, which is where the direct shader on Mali loses badly.

So this keeps upstream's `c_out` rule and adds a cheap-gather disjunct, scoped to Mali so no other device changes behaviour.

Mali-G76 (Galaxy S10+), fp16, interleaved GPU timestamps, median. "before" is upstream today:

| model | before | c_out gate | gather bound | **this** |
| --- | --- | --- | --- | --- |
| style_transfer_candy @640 | 1912.15 | 1436.31 | 1410.94 | **1411.67** |
| fast_sam s @640 | 345.09 | 404.54 | n/a | **341.61** |
| fcn_resnet50 @520 | 2938.13 | 2993.47 | 3505.06 | **2948.10** |

fcn gives up 0.3%, at the edge of the run-to-run range. The other two are wins, and neither earlier rule was a win on all three.

Regression check (these runners also carry pytorch#22784, which is why the mobilenet models move):

| model | before | this |
| --- | --- | --- |
| lraspp @520 | 71.28 | 67.78 |
| deeplabv3-mnv3 @520 | 165.27 | 161.76 |
| selfie @256 | 6.49 | 6.18 |
| ssdlite320 | 37.79 | 35.80 |

32M is calibrated on this device, not derived. The cases it separates are far apart (~4M and ~30M on one side, ~100M and up on the other), so the exact value is not delicate, but a maintainer with other Mali parts may want to revisit it. Adreno is deliberately untouched.


cc @SS-JIA @manuelcandales @digantdesai @cbilgin

Author: msluszniak

backends/vulkan/runtime/graph/ops/impl/Convolution.cpp

@@ -312,15 +312,24 @@ Conv2dMethod get_conv2d_method(
 //   - groups == 1
 //   - dilation == 1 (all dims)
 //
-// Selection rule: use im2col on Mali universally, or once the output channel
-// count is large enough to amortize the fixed ~N*K_total im2col gather cost.
+// Selection rule: use im2col once the output channel count is large enough to
+// amortize the fixed ~N*K_total im2col gather cost.
 constexpr int64_t kIm2colMinCOut = 128;
 
+// A cheap gather is a second, independent reason to take im2col. The gather
+// materializes an N x K_total matrix before the GEMM runs, and a large gather
+// is fine when c_out is large because the GEMM reads it back c_out times. When
+// c_out is small it is only worth paying if the matrix is small outright, which
+// is where the direct shader on Mali loses badly: an 80x80 3x3 conv wants ~4M
+// elements, a 640x640 9x9 conv wants ~100M.
+constexpr int64_t kIm2colMaxCheapGatherElements = 32 * 1024 * 1024;
+
 bool should_use_conv2d_im2col(
     ComputeGraph& graph,
     const ValueRef weight_data,
     const int64_t groups_val,
-    const Kernel2dParams& kernel_params) {
+    const Kernel2dParams& kernel_params,
+    const ValueRef out) {
   if (groups_val != 1) {
     return false;
   }
@@ -329,7 +338,14 @@ bool should_use_conv2d_im2col(
   }
   const auto weight_sizes = graph.sizes_of(weight_data);
   const int64_t c_out = weight_sizes.at(0);
-  return graph.device_is_mali() || c_out >= kIm2colMinCOut;
+  const int64_t k_total =
+      weight_sizes.at(1) * weight_sizes.at(2) * weight_sizes.at(3);
+
+  const auto out_sizes = graph.sizes_of(out);
+  const size_t ndim = out_sizes.size();
+  const int64_t n = out_sizes.at(ndim - 1) * out_sizes.at(ndim - 2);
+  return c_out >= kIm2colMinCOut ||
+      (graph.device_is_mali() && n * k_total <= kIm2colMaxCheapGatherElements);
 }
 
 GlobalWorkGrid create_conv2d_gwg(
@@ -539,7 +555,8 @@ void add_conv2d_node(
   // device); the default (false) reproduces the production routing exactly.
   const bool use_im2col = !force_direct &&
       method == Conv2dMethod::SlidingWindow &&
-      should_use_conv2d_im2col(graph, weight_data, groups_val, kernel_params);
+      should_use_conv2d_im2col(
+          graph, weight_data, groups_val, kernel_params, out);
   if (use_im2col) {
     return conv2d_gemm_impl(
         graph,