Gate output-centric image_to_nchw staging to discrete GPUs (#21022)

Differential Revision: D112599710

Pull Request resolved: #21022

Author: SS-JIA

backends/vulkan/runtime/graph/ops/glsl/image_to_nchw.glsl

@@ -19,6 +19,19 @@ ${define_explicit_type_extensions(DTYPE)}
 
 ${define_active_storage_type(STORAGE)}
 
+$if TO_STAGING:
+  #define TO_STAGING
+$if COALESCED_WRITES:
+  #define COALESCED_WRITES
+
+// The coalesced path writes contiguous NCHW staging offsets directly and does
+// not honor buf_meta / buf_layout, so it is only valid for the staging
+// destination. Fail the build if a coalesced non-staging (clone) variant is
+// ever configured.
+#if defined(COALESCED_WRITES) && !defined(TO_STAGING)
+#error "COALESCED_WRITES requires TO_STAGING"
+#endif
+
 layout(std430) buffer;
 
 #include "indexing.glslh"
@@ -39,17 +52,53 @@ $if not TO_STAGING:
   ${layout_declare_spec_const(C, "int", "buf_layout", "CONTIG_LAYOUT_INT")}
 
 void main() {
+#ifdef COALESCED_WRITES
+  // Output-centric dispatch: one thread per staging (NCHW) element. Consecutive
+  // threads write consecutive staging offsets, so writes to the host-visible
+  // staging buffer are fully coalesced. On a discrete GPU the staging buffer is
+  // PCIe-backed, where coalescing the writes matters far more than coalescing
+  // the (VRAM-cached) texture reads -- even though each texel is now fetched up
+  // to 4 times, once per component. On unified-memory (mobile) GPUs the extra
+  // fetches are a net loss, so this variant is gated to discrete GPUs.
+  const int oi = int(gl_GlobalInvocationID.x);
+  const int W = inp.sizes.x;
+  const int H = inp.sizes.y;
+  const int C = inp.sizes.z;
+  const int N = inp.sizes.w;
+  if (oi >= W * H * C * N) {
+    return;
+  }
+
+  TensorIndex4D tidx;
+  tidx.data.x = oi % W;
+  tidx.data.y = (oi / W) % H;
+  tidx.data.z = (oi / (W * H)) % C;
+  tidx.data.w = oi / (W * H * C);
+
+  const TextureElementIndex tex =
+      tensor4d_idx_to_texture_element_idx_simple(inp, tidx, in_layout);
+  #ifdef USING_TEXTURE2D
+  const VEC4_T intex = texelFetch(t_in, tex.pos.xy, 0);
+  #else
+  const VEC4_T intex = texelFetch(t_in, tex.pos, 0);
+  #endif
+  buf_out[oi] = BUF_T(intex[tex.comp]);
+#else
+  // Texel-centric dispatch: one thread per texture texel, writing up to 4 packed
+  // components. Reads are coalesced (one fetch per texel); writes are strided.
+  // Preferred on unified-memory GPUs (mobile) where the staging buffer is not
+  // PCIe-backed, so write coalescing buys nothing and the single fetch wins.
   const ivec3 pos = ivec3(gl_GlobalInvocationID);
   if (out_of_bounds(pos, inp)) {
     return;
   }
 
   TensorIndex4D tidx = texture_pos_to_tensor4d_idx_simple(inp, pos, in_layout);
-#ifdef USING_TEXTURE2D
+  #ifdef USING_TEXTURE2D
   const VEC4_T intex = texelFetch(t_in, pos.xy, 0);
-#else
+  #else
   const VEC4_T intex = texelFetch(t_in, pos, 0);
-#endif
+  #endif
 
   const int packed_dim = get_packed_dim(in_layout);
   int packed_dim_val;
@@ -77,16 +126,17 @@ void main() {
   int limit = min(4, packed_dim_size - packed_dim_val);
 
   for (int comp = 0; comp < limit; comp++) {
-    $if TO_STAGING:
-      // Compute contiguous NCHW index
-      int nchwi = tidx.data.x
-                + tidx.data.y * inp.sizes.x
-                + tidx.data.z * inp.sizes.x * inp.sizes.y
-                + tidx.data.w * inp.sizes.x * inp.sizes.y * inp.sizes.z;
-      buf_out[nchwi] = BUF_T(intex[comp]);
-    $else:
-      int bufi = tensor4d_idx_to_buf_idx(buf_meta, tidx, buf_layout);
-      buf_out[bufi] = BUF_T(intex[comp]);
+#ifdef TO_STAGING
+    // Staging buffer is contiguous NCHW; compute the flat offset directly.
+    int nchwi = tidx.data.x
+              + tidx.data.y * inp.sizes.x
+              + tidx.data.z * inp.sizes.x * inp.sizes.y
+              + tidx.data.w * inp.sizes.x * inp.sizes.y * inp.sizes.z;
+    buf_out[nchwi] = BUF_T(intex[comp]);
+#else
+    int bufi = tensor4d_idx_to_buf_idx(buf_meta, tidx, buf_layout);
+    buf_out[bufi] = BUF_T(intex[comp]);
+#endif
 
     if (packed_dim == 0) {
       tidx.data.x++;
@@ -98,4 +148,5 @@ void main() {
       tidx.data.w++;
     }
   }
+#endif
 }

backends/vulkan/runtime/graph/ops/glsl/image_to_nchw.yaml

@@ -10,6 +10,7 @@ image_to_nchw:
     BUF_DTYPE: float
     STORAGE: texture3d
     TO_STAGING: True
+    COALESCED_WRITES: False
   generate_variant_forall:
     combination:
       parameter_names: [DTYPE, BUF_DTYPE]
@@ -25,5 +26,10 @@ image_to_nchw:
     - NAME: image_to_nchw_texture3d
     - NAME: image_to_nchw_texture2d
       STORAGE: texture2d
+    - NAME: image_to_nchw_coalesced_texture3d
+      COALESCED_WRITES: True
+    - NAME: image_to_nchw_coalesced_texture2d
+      STORAGE: texture2d
+      COALESCED_WRITES: True
     - NAME: clone_image_to_buffer
       TO_STAGING: False

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

@@ -36,6 +36,11 @@ bool is_nchw_to_bitw8_shader(const vkapi::ShaderInfo& shader) {
   return shader_prefix_str == kNchwToBitw8PrefixStr;
 }
 
+bool is_coalesced_image_to_nchw_shader(const vkapi::ShaderInfo& shader) {
+  return shader.kernel_name.find("image_to_nchw_coalesced") !=
+      std::string::npos;
+}
+
 void add_staging_to_tensor_node(
     ComputeGraph& graph,
     const ValueRef in_staging,
@@ -94,19 +99,26 @@ utils::uvec3 tensor_to_staging_global_wg_size(
 
   utils::uvec3 global_wg_size = graph->create_global_wg_size(in_tensor);
 
-  // Normally, the image_to_nchw shader is structured so that each thread reads
-  // one texel from the input texture and writes each component of the texel
-  // into the corresponding location in the output buffer. However, this shader
-  // is structured slightly differently in that each thread writes out a
-  // complete 32 bit integer (containing 4 packed 8-bit integers) into the
-  // output buffer. Therefore, the global work group size for this shader will
-  // be the number of elements in the output buffer divided by 4, as opposed to
-  // the extents of the input texture.
+  // The bitw8 shader writes out a complete 32 bit integer (containing 4 packed
+  // 8-bit integers) per thread, so its global work group size is the number of
+  // elements in the output buffer divided by 4.
   if (is_bitw8_shader(shader)) {
     const uint32_t buffer_len = utils::safe_downcast<uint32_t>(
         graph->get_staging(out_staging)->numel() / 4);
     global_wg_size = {buffer_len, 1, 1};
+  } else if (is_coalesced_image_to_nchw_shader(shader)) {
+    // The coalesced (output-centric) image_to_nchw variant dispatches one
+    // thread per output (staging) element so that consecutive threads write
+    // consecutive NCHW offsets, keeping writes to the PCIe-backed staging
+    // buffer fully coalesced. This mirrors the buffer_to_nchw path, whose
+    // global size is already numel-based via create_global_wg_size.
+    const uint32_t buffer_len = utils::safe_downcast<uint32_t>(
+        graph->get_staging(out_staging)->numel());
+    global_wg_size = {buffer_len, 1, 1};
   }
+  // Otherwise (texel-centric image_to_nchw, used on unified-memory GPUs) keep
+  // the default texel-grid global size from create_global_wg_size: one thread
+  // per texture texel.
 
   return global_wg_size;
 }

backends/vulkan/runtime/graph/ops/utils/StagingUtils.cpp

@@ -80,7 +80,19 @@ vkapi::ShaderInfo get_tensor_to_nchw_shader(
     return VK_KERNEL_FROM_STR(kernel_name);
   }
 
-  kernel_name = "image_to_nchw";
+  // On a discrete GPU the staging buffer is read back over PCIe, so the
+  // output-centric (coalesced-write) variant is a large win. On integrated GPUs
+  // the staging buffer is not PCIe-backed, and the coalesced variant's
+  // redundant texture fetches make it a net loss -- so default to the
+  // texel-centric variant there.
+  //
+  // Gate on device type, not has_unified_memory(): a discrete GPU with
+  // Resizable BAR exposes a DEVICE_LOCAL | HOST_VISIBLE memory type, so
+  // has_unified_memory() returns true for it and would wrongly route it to the
+  // texel-centric path (measured ~10x end-to-end regression on an RTX 4080).
+  const bool coalesced_writes =
+      !graph.context()->adapter_ptr()->is_integrated_gpu();
+  kernel_name = coalesced_writes ? "image_to_nchw_coalesced" : "image_to_nchw";
   add_storage_type_suffix(kernel_name, src_storage_type);
   add_dtype_suffix(kernel_name, src_dtype);
   add_dtype_suffix(kernel_name, staging_dtype);