SS-JIA · Nov 2, 2026 · Nov 2, 2026
@@ -12,17 +12,21 @@
 
 ${define_active_storage_type("buffer")}
 
+#extension GL_EXT_control_flow_attributes : require
+
 layout(std430) buffer;
 
 #include "indexing.glslh"
+#include "block_int8x4_utils.glslh"
 
 // Output buffer: packed int8x4 (each int32 = 4 packed int8 along packed_dim)
 ${layout_declare_tensor(B, "w", "t_outp", "int", "buffer")}
 // Input buffer: packed int8x4 (each int32 = 4 packed int8 along packed_dim)
 ${layout_declare_tensor(B, "r", "t_inp", "int", "buffer")}
 
-// Metadata for output and input tensors. Both are int8x4 packed buffers but
-// may use different block layouts (e.g. PACKED_INT8_4W vs PACKED_INT8_4W4C).
+// Metadata for output and input tensors. Both are channels-packed int8x4
+// buffers but may use different block layouts (e.g. PACKED_INT8_4W4C vs
+// PACKED_INT8_4C).
 ${layout_declare_ubo(B, "BufferMetadata", "outp")}
 ${layout_declare_ubo(B, "BufferMetadata", "inp")}
 
@@ -31,131 +35,78 @@ layout(push_constant) uniform restrict Block {
   float inv_scale_out;
   int zp_in;
   int zp_out;
-  int upscale_factor;
-  // Whether we can skip the requantize math and do a pure byte shuffle.
-  // Set to 1 by the host when (scale_in == scale_out) && (zp_in == zp_out).
-  int passthrough;
 };
 
 layout(local_size_x_id = 0, local_size_y_id = 1, local_size_z_id = 2) in;
 
-#include "dispatch.glslh"
-
 ${layout_declare_spec_const(C, "int", "inp_layout", "CONTIG_LAYOUT_INT")}
 ${layout_declare_spec_const(C, "int", "outp_layout", "CONTIG_LAYOUT_INT")}
+// Set by the host when input and output quantization params match, making the
+// shader a pure byte shuffle.
+${layout_declare_spec_const(C, "int", "passthrough", "0")}
 
 /*
- * PixelShuffle map: out[n, c_out, oh, ow] =
- *     in[n, c_out * r * r + (oh % r) * r + (ow % r), oh / r, ow / r].
+ * PixelShuffle with upscale factor 2:
+ *   out[n, c, 2h + dy, 2w + dx] = in[n, 4c + 2dy + dx, h, w]
  *
- * Each thread produces one output int32 word (= 4 consecutive output channels
- * at one (n, oh, ow) spatial position). Channels are the packed dim and
- * packed_dim_block_size is 4, so writing one int word fills 4 channel lanes.
+ * Each thread handles one block of 4 output channels at the 2x2 output pixels
+ * produced by one input pixel. The 4 input words holding input channels
+ * [16 * c4, 16 * c4 + 16) at that pixel form a 4x4 byte matrix where byte
+ * (2dy + dx) of word i is output channel 4 * c4 + i at (2w + dx, 2h + dy). Its
+ * transpose is the 4 output words, so every loaded and stored byte is used.
  *
- * The four channel lanes inside an output int come from four DIFFERENT input
- * words (channels spaced by r*r in the input), so each thread issues 4 input
- * loads. The (oh % r, ow % r) -> input lane mapping is fixed for a given
- * thread because all four output lanes share (oh, ow). Out-of-output-bounds
- * channel lanes (when C_out is not a multiple of 4) are zero-filled.
- *
- * Supported layouts: channels-packed family (PACKED_INT8_4W4C, PACKED_INT8_4C1W,
- * PACKED_INT8_CONV2D). Layout-aware byte indexing is handled by
- * tensor4d_idx_to_buf_idx (which consumes inp_layout / outp_layout).
+ * Global x indexes the channel block, y the input width, and z the input row
+ * and batch.
  */
 
-// Apply requantize to an int8 lane value.
-int requantize_lane(const int q_in) {
-  if (passthrough != 0) {
-    return q_in;
+// Requantize the 4 int8 lanes of a packed word.
+int requantize_word(const int packed) {
+  int result = 0;
+  [[unroll]] for (int i = 0; i < 4; ++i) {
+    const int q_in = (packed << (24 - 8 * i)) >> 24;
+    const float dq = float(q_in - zp_in) * scale_in;
+    const int q_out =
+        clamp(int(round(dq * inv_scale_out)) + zp_out, -128, 127);
+    result |= (q_out & 0xFF) << (8 * i);
   }
-  // Requantize: round((q_in - zp_in) * scale_in * inv_scale_out) + zp_out,
-  // clamped to int8.
-  float dq = float(q_in - zp_in) * scale_in;
-  int qv = int(round(dq * inv_scale_out)) + zp_out;
-  qv = clamp(qv, -128, 127);
-  return qv;
+  return result;
 }
 
 void main() {
-  // Output sizes (WHCN order via meta.sizes[0])
-  const int W_out = int(safe_idx(outp.sizes[0], 0));
-  const int H_out = int(safe_idx(outp.sizes[0], 1));
-  const int C_out = int(safe_idx(outp.sizes[0], 2));
-  const int N = int(safe_idx(outp.sizes[0], 3));
-
-  // Input sizes
-  const int W_in = int(safe_idx(inp.sizes[0], 0));
-  const int H_in = int(safe_idx(inp.sizes[0], 1));
-  const int C_in = int(safe_idx(inp.sizes[0], 2));
-
-  // One thread per output int32 word: word covers 4 consecutive channels
-  // (along the packed dim) at one (n, oh, ow) spatial position.
-  const int C_words = div_up_4(C_out);
-  const int total_words = N * C_words * H_out * W_out;
-  const int thread_idx = int(linear_idx_from_gid());
-  if (thread_idx >= total_words) {
+  const int c4 = int(gl_GlobalInvocationID.x);
+  const int iw = int(gl_GlobalInvocationID.y);
+  const int H_in = int(inp.sizes[0].y);
+  const int ih = int(gl_GlobalInvocationID.z) % H_in;
+  const int n = int(gl_GlobalInvocationID.z) / H_in;
+
+  const int C_out = int(outp.sizes[0].z);
+  if (4 * c4 >= C_out || iw >= int(inp.sizes[0].x) ||
+      n >= int(inp.sizes[0].w)) {
     return;
   }
 
-  // Decode thread_idx in (W_out, H_out, C_words, N) order.
-  const int ow = thread_idx % W_out;
-  const int oh = (thread_idx / W_out) % H_out;
-  const int c_word = (thread_idx / (W_out * H_out)) % C_words;
-  const int n = thread_idx / (W_out * H_out * C_words);
-  const int c_out_base = c_word * 4;
-
-  const int r = upscale_factor;
-  // (oh % r, ow % r) determines which input channel lane within the input
-  // word group of size r*r — constant for all 4 output channel lanes here.
-  const int offset = (oh % r) * r + (ow % r);
-  const int ih = oh / r;
-  const int iw = ow / r;
-
-  const int c_in_first = c_out_base * r * r + offset;
-
-  // Compute byte_idx for the first lane (i=0) via the layout-aware helper.
-  TensorIndex4D inp_idx;
-  inp_idx.data = ivec4(iw, ih, c_in_first, n);
-  const int byte_idx_first = tensor4d_idx_to_buf_idx(inp, inp_idx, inp_layout);
-
-  // byte_stride between successive c_in advances of r*r = inner_block_size = 4.
-  // Each advance bumps the block-space C coord by 1, so byte_idx grows by
-  // stride[inner_dim] * block_numel. Both factors are layout-only, no second
-  // helper call needed. (Assumes r*r == inner_block_size == 4, enforced by the
-  // C++ dispatch's r==2 and packed_dim_block_size==4 asserts.)
-  const int byte_stride =
-      int(safe_idx(inp.strides[0], get_packed_dim(inp_layout))) * get_block_numel(inp_layout);
-
-  // lane is the byte position within an int32 word, which equals
-  // (intra_block_idx % 4) since block_numel is a multiple of 4. And
-  // intra_block_idx % 4 == inner_offset == c_in_first % 4 == offset.
-  const int lane = offset;
-
-  int packed_out = 0;
+  TensorIndex4D inp_tidx;
+  inp_tidx.data = ivec4(iw, ih, 16 * c4, n);
+  const int inp_base = tensor4d_idx_to_texel_idx(inp, inp_tidx, inp_layout);
+  // Word offset between consecutive channel blocks
+  const int inp_c4_stride =
+      int(inp.strides[0].z) * div_4(get_block_numel(inp_layout));
+
+  ivec4 block = ivec4(0);
   [[unroll]] for (int i = 0; i < 4; ++i) {
-    const int c_out_lane = c_out_base + i;
-    int q_out = 0;
-    if (c_out_lane < C_out) {
-      const int c_in = c_in_first + i * r * r;
-      int q_in;
-      if (iw >= W_in || ih >= H_in || c_in >= C_in) {
-        q_in = zp_in;
-      } else {
-        const int byte_idx = byte_idx_first + i * byte_stride;
-        const int word_idx = div_4(byte_idx);
-        const int packed = t_inp[word_idx];
-        // Sign-extend from 8-bit
-        q_in = ((packed >> (lane * 8)) << 24) >> 24;
+    // Output channels past C_out are padding and stay zero.
+    if (4 * c4 + i < C_out) {
+      block[i] = t_inp[inp_base + i * inp_c4_stride];
+      if (passthrough == 0) {
+        block[i] = requantize_word(block[i]);
       }
-      q_out = requantize_lane(q_in);
     }
-    packed_out |= (q_out & 0xFF) << (i * 8);
   }
+  block = transpose_int8x4_block(block);
 
-  // Store the packed int directly. Output's packed dim is channels with
-  // block size 4, so the byte index for c_out_base aligns to a word boundary.
-  TensorIndex4D outp_idx;
-  outp_idx.data = ivec4(ow, oh, c_out_base, n);
-  const int outp_byte_idx = tensor4d_idx_to_buf_idx(outp, outp_idx, outp_layout);
-  t_outp[div_4(outp_byte_idx)] = packed_out;
+  [[unroll]] for (int i = 0; i < 4; ++i) {
+    TensorIndex4D outp_tidx;
+    outp_tidx.data = ivec4(2 * iw + (i & 1), 2 * ih + (i >> 1), 4 * c4, n);
+    t_outp[tensor4d_idx_to_texel_idx(outp, outp_tidx, outp_layout)] = block[i];
+  }
 }
@@ -39,9 +39,9 @@ void resize_q8ta_pixel_shuffle_node(
   graph->virtual_resize(out, out_sizes);
 }
 
-// Global wg picker: one thread per output int32 word. For a channels-packed
-// int8x4 output with channel block size 4, the number of output int words is
-// N * div_up_4(C_out) * H_out * W_out.
+// Each thread produces one block of 4 output channels at the 2x2 output pixels
+// generated by one input pixel. x indexes the channel block and y the input
+// width.
 GlobalWorkGrid pick_q8ta_pixel_shuffle_gwg(
     ComputeGraph* graph,
     const vkapi::ShaderInfo& shader,
@@ -50,15 +50,13 @@ GlobalWorkGrid pick_q8ta_pixel_shuffle_gwg(
   (void)shader;
   (void)resize_args;
   const ValueRef out = args.at(0).refs.at(0);
-  const std::vector<int64_t>& sizes = graph->sizes_of(out);
-  const int64_t N = utils::val_at(-4, sizes);
-  const int64_t C = utils::val_at(-3, sizes);
-  const int64_t H = utils::val_at(-2, sizes);
-  const int64_t W = utils::val_at(-1, sizes);
-  const int64_t c_words = utils::div_up(C, int64_t(4));
-  const uint32_t total_words =
-      utils::safe_downcast<uint32_t>(N * c_words * H * W);
-  return graph->create_linear_gwg(total_words);
+  const ValueRef in = args.at(1).refs.at(0);
+  const uint32_t C_out = graph->size_at<uint32_t>(-3, out);
+  const uint32_t W_in = graph->size_at<uint32_t>(-1, in);
+  const uint32_t H_in = graph->size_at<uint32_t>(-2, in);
+  const uint32_t N = graph->size_at<uint32_t>(-4, in);
+  return GlobalWorkGrid(
+      {utils::div_up_4(C_out), W_in, H_in * N}, kTiledWorkGrid);
 }
 
 LocalWorkGroup pick_q8ta_pixel_shuffle_lwg(
@@ -67,14 +65,25 @@ LocalWorkGroup pick_q8ta_pixel_shuffle_lwg(
     const GlobalWorkGrid& gwg,
     const std::vector<ArgGroup>& args,
     const std::vector<ValueRef>& resize_args) {
-  (void)graph;
   (void)shader;
-  (void)gwg;
-  (void)args;
   (void)resize_args;
-  // Linear (1D) dispatch: a flat 64-wide workgroup matches the pattern used
-  // by pick_square_lwg_with_block_config in the linear case.
-  return LocalWorkGroup(64u, 1u, 1u);
+  const api::PackedDimInfo& out_info =
+      graph->packed_dim_info_of(args.at(0).refs.at(0));
+  const api::PackedDimInfo& in_info =
+      graph->packed_dim_info_of(args.at(1).refs.at(0));
+  // Span the dims that are contiguous in memory so that each store instruction
+  // covers whole cache lines: channel blocks for PACKED_INT8_4C outputs, width
+  // when both tensors are PACKED_INT8_4C1W, and a mix of both otherwise.
+  LwgShape shape(1u, 2u, 0u);
+  if (!out_info.block_transposed && out_info.outer_packed_dim_block_size == 1) {
+    shape = kLinearLwg;
+  } else if (out_info.block_transposed && in_info.block_transposed) {
+    shape = LwgShape(0u, 1u, 0u);
+  }
+  LocalWorkGroup lwg(
+      shape, graph->context()->adapter_ptr()->recommended_lwg_nthreads());
+  lwg.fit_to_global(gwg);
+  return lwg;
 }
 
 } // namespace
@@ -105,10 +114,8 @@ void add_q8ta_pixel_shuffle_node(
   VK_CHECK_COND(graph.dim_of(packed_int8_output) == 4);
 
   // Both tensors must use a channels-packed int8x4 layout (packed_dim=C with
-  // packed_dim_block_size=4). The supported layouts are PACKED_INT8_4W4C
-  // (outer block on W), PACKED_INT8_4C1W (no outer block), and
-  // PACKED_INT8_CONV2D. Each output thread writes one int word covering 4
-  // consecutive channels at one (n, oh, ow) position.
+  // packed_dim_block_size=4), e.g. PACKED_INT8_4W4C, PACKED_INT8_4C or
+  // PACKED_INT8_4C1W.
   const api::PackedDimInfo& in_info =
       graph.packed_dim_info_of(packed_int8_input);
   const api::PackedDimInfo& out_info =
@@ -119,7 +126,7 @@ void add_q8ta_pixel_shuffle_node(
   VK_CHECK_COND(in_info.packed_dim == WHCN::kChannelsDim);
   VK_CHECK_COND(out_info.packed_dim == WHCN::kChannelsDim);
 
-  // Upscale factor: only r=2 is exercised by the model and tests.
+  // The shader is specialized for r=2.
   const int32_t r = graph.extract_scalar<int32_t>(upscale_factor);
   VK_CHECK_COND(r == 2);
 
@@ -148,29 +155,21 @@ void add_q8ta_pixel_shuffle_node(
   if (scale_diff <= scale_thresh && zp_in == zp_out) {
     passthrough = 1;
   }
-  int32_t r_val = r;
 
   std::vector<PushConstantDataInfo> push_constants = {
       PushConstantDataInfo(&scale_in, sizeof(scale_in)),
       PushConstantDataInfo(&inv_scale_out, sizeof(inv_scale_out)),
       PushConstantDataInfo(&zp_in, sizeof(zp_in)),
       PushConstantDataInfo(&zp_out, sizeof(zp_out)),
-      PushConstantDataInfo(&r_val, sizeof(r_val)),
-      PushConstantDataInfo(&passthrough, sizeof(passthrough)),
   };
 
   // UBOs
   vkapi::ParamsBindList ubos;
   ubos.append(graph.buffer_meta_ubo(packed_int8_output));
   ubos.append(graph.buffer_meta_ubo(packed_int8_input));
 
-  // Each thread writes one int32 output word (= 4 consecutive output channels
-  // at one (n, oh, ow) spatial position). Total threads =
-  //   N * div_up_4(C_out) * H_out * W_out
-  // The custom global wg picker computes this from the output sizes; the
-  // shader internally re-derives the same decomposition from
-  // gl_GlobalInvocationID. The resize-args list still carries the upscale
-  // factor so the resize callback can stamp the output size.
+  // The resize-args list carries the upscale factor so the resize callback can
+  // stamp the output size.
   const ValueRef r_resize_arg = static_cast<ValueRef>(r);
 
   std::string kernel_name = "q8ta_pixel_shuffle";
@@ -188,7 +187,8 @@ void add_q8ta_pixel_shuffle_node(
       push_constants,
       // Specialization Constants
       {graph.hashed_layout_of(packed_int8_input),
-       graph.hashed_layout_of(packed_int8_output)},
+       graph.hashed_layout_of(packed_int8_output),
+       passthrough},
       // Resize args: [upscale_factor]
       {r_resize_arg},
       // Resizing Logic

@@ -39,15 +39,17 @@ TestCase create_test_case_from_config(const PixelShuffleConfig& config) {
   std::string operator_name = "test_etvk." + config.op_name + ".default";
   test_case.set_operator_name(operator_name);
 
-  // FP input: shape [N, C_in, H, W]
+  // FP input: shape [N, C_in, H, W]. With matching qparams, integer inputs and
+  // a unit scale make quantization exact, so the output must match the
+  // reference bit for bit.
   ValueSpec input_tensor(
       config.in_shape,
       vkapi::kFloat,
       utils::kBuffer,
       utils::kWidthPacked,
-      DataGenType::RANDOM);
+      config.same_qparams ? DataGenType::RANDINT8 : DataGenType::RANDOM);
 
-  float input_scale_val = 0.007112f;
+  float input_scale_val = config.same_qparams ? 1.0f : 0.007112f;
   ValueSpec input_scale(input_scale_val);
   int32_t input_zp_val = 0;
   ValueSpec input_zp(input_zp_val);
@@ -85,9 +87,15 @@ TestCase create_test_case_from_config(const PixelShuffleConfig& config) {
   test_case.add_input_spec(out_layout_spec);
   test_case.add_output_spec(output_tensor);
 
-  // Tolerance: ~1 quant step in the bigger of the two scales.
-  float tol = std::max(input_scale_val, output_scale_val) + 1e-4f;
-  test_case.set_abs_tolerance(tol);
+  if (config.same_qparams) {
+    test_case.set_abs_tolerance(0.0f);
+    test_case.set_rel_tolerance(0.0f);
+  } else {
+    // Requantization rounding may differ by ~1 quant step in the bigger of the
+    // two scales.
+    float tol = std::max(input_scale_val, output_scale_val) + 1e-4f;
+    test_case.set_abs_tolerance(tol);
+  }
 
   // Filter shaders that are just measurement overhead (staging copies and the
   // surrounding quantize/dequantize that wrap the operation under test).

@@ -24,10 +24,10 @@ constexpr int64_t kRefDimSizeLimit = 512;
 //                           upscale_factor, in_layout, out_layout) -> fp_out
 // Implementation: fused fast-path kernel. The in_layout / out_layout strings
 // select the channels-packed int8x4 layout used for the temporary quantized
-// tensors. Supported values: "4W4C", "4C1W".
+// tensors. Supported values: "4W4C", "4C", "4C1W".
 // (PACKED_INT8_CONV2D is a Python/serialization-level alias that the runtime
-// resolves to kPackedInt8_4C1W, so it is not exercised separately here -- it
-// would only re-test the same C++ kernel path as "4C1W".)
+// resolves to kPackedInt8_4C or kPackedInt8_4C1W, so it is not exercised
+// separately here.)
 
 struct PixelShuffleConfig {
   std::vector<int64_t> in_shape; // [N, C*r*r, H, W]