From 1b31697489ee82da161c385f81ae0d9ae3164253 Mon Sep 17 00:00:00 2001 From: Stephen Jia Date: Mon, 2 Nov 2026 08:44:26 -0800 Subject: [PATCH] [ET-VK][pixel_shuffle] Rewrite q8ta_pixel_shuffle to write whole output words ## Ulterior Motive Reduce the inference latency of int8 segmentation CNNs that upsample with pixel shuffle on mobile GPUs. Per-op profiling of such a 512x512 model found `q8ta_pixel_shuffle` to be its most expensive operator, running far below memory bandwidth. ## Rationale **What**: Rewrites `q8ta_pixel_shuffle` so that each invocation turns the 4 input words holding 16 channels at one input pixel into 4 full output words with a 4x4 byte transpose. The global work grid and local workgroup size change to match. The test sets add the `PACKED_INT8_4C` layout and every pair of channels-packed layouts, and check equal-qparams cases bit-exactly. **Why**: The old shader ran one invocation per output word, with output width as the fastest-varying index and the channel block as the slowest. Channels-packed int8x4 layouts store channels innermost, so neighbouring invocations wrote words 64 B apart for a `PACKED_INT8_4C` output, or 16 B apart for `PACKED_INT8_4W4C`. Each 64 B output line was therefore completed over 16 (4C) or 4 (4W4C) separate passes over the whole tensor. Once the output no longer fits in cache, every pass writes partial lines back to memory. In addition, the 4 lanes of an output word come from 4 different input words, so each invocation issued 4 loads to produce 4 bytes. On Adreno 830, the old shader took the following for [1,256,256,256] -> [1,64,512,512]: | Layouts (in -> out) | Time | | --- | --- | | 4W4C -> 4C | 10.65 ms | | 4W4C -> 4W4C | 2.93 ms | | 4C1W -> 4C1W (width innermost, which matches the old thread order) | 0.80 ms | Its 4W4C -> 4C time grew faster than the data size: 0.26 / 2.0 / 10.7 ms at 64^2 / 128^2 / 256^2 input planes. The new shader takes 0.55 ms for that 4W4C -> 4C shape (19x faster) and is 2.8-4.1x faster on smaller shapes. ## Details ### Shader With upscale factor 2, `out[n, c, 2h + dy, 2w + dx] = in[n, 4c + 2dy + dx, h, w]`. Each invocation handles output channel block `k` at one input pixel `(h, w)`. Input word `i` holds input channels `16k + 4i` to `16k + 4i + 3`, which are output channel `4k + i` at the 4 output pixels of the 2x2 block. The 4 input words form a 4x4 byte matrix whose transpose is the 4 output words: ``` input word i = [ out ch 4k+i at (dy,dx) = (0,0) | (0,1) | (1,0) | (1,1) ] | transpose_int8x4_block v output word j = [ out ch 4k | 4k+1 | 4k+2 | 4k+3 ] at pixel (2h + j/2, 2w + j%2) ``` Every loaded and stored byte is used exactly once. The requantize formula is unchanged and is applied per word before the transpose. `passthrough` (input and output qparams match) is now a specialization constant, so the requantize path compiles out in that case. The `upscale_factor` push constant is removed; the dispatch already rejects any factor other than 2. ### Dispatch The global work grid is (output channel blocks, input width, input height x batch). The local workgroup shape depends on the layouts, so that one workgroup's stores cover contiguous memory. Values below are for 64 output channels and were chosen from sweeps on Adreno 830: - `PACKED_INT8_4C` output: channel blocks first (16x4x1), so 16 invocations write one 64 B line. - `PACKED_INT8_4C1W` on both sides: width (1x64x1). - Other layout pairs: 4x16x1. ### Tests and benchmark sets - The `PACKED_INT8_4C` layout is added. Correctness covers all 9 pairs of `PACKED_INT8_4W4C`, `PACKED_INT8_4C` and `PACKED_INT8_4C1W`. - New correctness shapes: C_out = 1 (three padding lanes) and batch 2. - Equal-qparams cases use integer-valued inputs with scale 1 and zero tolerance, so they check bit-exactness. Requantize cases keep a one-step tolerance: the GPU multiplies by the reciprocal scale and GLSL `round()` breaks ties in an implementation-defined direction, while the CPU reference divides. - The performance set has 81 cases (9 layout pairs x 9 shapes), including decoder-style [1,256,H,W] inputs with H = W = 16 to 256. Suggested review order: `q8ta_pixel_shuffle.glsl`, `Q8taPixelShuffle.cpp`, then the test files. Authored with AI assistance (Claude Code). Differential Revision: [D129089325](https://our.internmc.facebook.com/intern/diff/D129089325/) **NOTE FOR REVIEWERS**: This PR has internal Meta-specific changes or comments, please review them on [Phabricator](https://our.internmc.facebook.com/intern/diff/D129089325/)! ghstack-source-id: 462609407 Pull-Request: https://github.com/pytorch/executorch/pull/24440 --- .../graph/ops/glsl/q8ta_pixel_shuffle.glsl | 161 ++++++------------ .../graph/ops/impl/Q8taPixelShuffle.cpp | 68 ++++---- .../q8ta/pixel_shuffle/pixel_shuffle.cpp | 20 ++- .../cases/q8ta/pixel_shuffle/pixel_shuffle.h | 6 +- .../test_sets/pixel_shuffle_sets.cpp | 41 ++--- .../test_ops/impl/TestQ8taPixelShuffle.cpp | 13 +- 6 files changed, 136 insertions(+), 173 deletions(-) diff --git a/backends/vulkan/runtime/graph/ops/glsl/q8ta_pixel_shuffle.glsl b/backends/vulkan/runtime/graph/ops/glsl/q8ta_pixel_shuffle.glsl index 69ce99ec839..638799e1a3c 100644 --- a/backends/vulkan/runtime/graph/ops/glsl/q8ta_pixel_shuffle.glsl +++ b/backends/vulkan/runtime/graph/ops/glsl/q8ta_pixel_shuffle.glsl @@ -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]; + } } diff --git a/backends/vulkan/runtime/graph/ops/impl/Q8taPixelShuffle.cpp b/backends/vulkan/runtime/graph/ops/impl/Q8taPixelShuffle.cpp index 45748ece4a7..2bf14fdf152 100644 --- a/backends/vulkan/runtime/graph/ops/impl/Q8taPixelShuffle.cpp +++ b/backends/vulkan/runtime/graph/ops/impl/Q8taPixelShuffle.cpp @@ -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& 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(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(-3, out); + const uint32_t W_in = graph->size_at(-1, in); + const uint32_t H_in = graph->size_at(-2, in); + const uint32_t N = graph->size_at(-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& args, const std::vector& 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(upscale_factor); VK_CHECK_COND(r == 2); @@ -148,15 +155,12 @@ void add_q8ta_pixel_shuffle_node( if (scale_diff <= scale_thresh && zp_in == zp_out) { passthrough = 1; } - int32_t r_val = r; std::vector 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 @@ -164,13 +168,8 @@ void add_q8ta_pixel_shuffle_node( 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(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 diff --git a/backends/vulkan/test/ops/benchmark/cases/q8ta/pixel_shuffle/pixel_shuffle.cpp b/backends/vulkan/test/ops/benchmark/cases/q8ta/pixel_shuffle/pixel_shuffle.cpp index 5f805171fac..e97b5f564c8 100644 --- a/backends/vulkan/test/ops/benchmark/cases/q8ta/pixel_shuffle/pixel_shuffle.cpp +++ b/backends/vulkan/test/ops/benchmark/cases/q8ta/pixel_shuffle/pixel_shuffle.cpp @@ -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). diff --git a/backends/vulkan/test/ops/benchmark/cases/q8ta/pixel_shuffle/pixel_shuffle.h b/backends/vulkan/test/ops/benchmark/cases/q8ta/pixel_shuffle/pixel_shuffle.h index 42dcc329359..c85491fe62a 100644 --- a/backends/vulkan/test/ops/benchmark/cases/q8ta/pixel_shuffle/pixel_shuffle.h +++ b/backends/vulkan/test/ops/benchmark/cases/q8ta/pixel_shuffle/pixel_shuffle.h @@ -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 in_shape; // [N, C*r*r, H, W] diff --git a/backends/vulkan/test/ops/benchmark/cases/q8ta/pixel_shuffle/test_sets/pixel_shuffle_sets.cpp b/backends/vulkan/test/ops/benchmark/cases/q8ta/pixel_shuffle/test_sets/pixel_shuffle_sets.cpp index 13605714009..4f4cea07ffd 100644 --- a/backends/vulkan/test/ops/benchmark/cases/q8ta/pixel_shuffle/test_sets/pixel_shuffle_sets.cpp +++ b/backends/vulkan/test/ops/benchmark/cases/q8ta/pixel_shuffle/test_sets/pixel_shuffle_sets.cpp @@ -7,7 +7,6 @@ #include #include -#include #include namespace executorch { @@ -17,15 +16,9 @@ namespace q8ta_pixel_shuffle { namespace { -// All (in_layout, out_layout) pairs across the channels-packed int8 family. -// CONV2D is a Python-level alias that resolves to 4C1W at the C++ runtime, so -// it is not listed separately -- it would just re-run the 4C1W kernel path. -const std::vector> kLayoutPairs = { - {"4W4C", "4W4C"}, - {"4W4C", "4C1W"}, - {"4C1W", "4W4C"}, - {"4C1W", "4C1W"}, -}; +// Channels-packed int8 layouts. CONV2D is a Python-level alias that resolves to +// 4C or 4C1W at the C++ runtime, so it is not listed separately. +const std::vector kLayouts = {"4W4C", "4C", "4C1W"}; std::vector generate_test_cases( const std::vector>& in_shapes, @@ -36,16 +29,18 @@ std::vector generate_test_cases( for (const auto& shape : in_shapes) { for (bool same_qp : same_qparams) { - for (const auto& layouts : kLayoutPairs) { - PixelShuffleConfig cfg; - cfg.in_shape = shape; - cfg.upscale_factor = upscale_factor; - cfg.same_qparams = same_qp; - cfg.in_layout = layouts.first; - cfg.out_layout = layouts.second; - cfg.test_case_name = test_case_name; + for (const std::string& in_layout : kLayouts) { + for (const std::string& out_layout : kLayouts) { + PixelShuffleConfig cfg; + cfg.in_shape = shape; + cfg.upscale_factor = upscale_factor; + cfg.same_qparams = same_qp; + cfg.in_layout = in_layout; + cfg.out_layout = out_layout; + cfg.test_case_name = test_case_name; - test_cases.push_back(create_test_case_from_config(cfg)); + test_cases.push_back(create_test_case_from_config(cfg)); + } } } } @@ -64,6 +59,8 @@ REGISTER_TEST_CASE_SET("q8ta/pixel_shuffle", "correctness") { {1, 24, 8, 4}, // out: [1, 6, 16, 8] {1, 32, 12, 8}, // out: [1, 8, 24, 16] {1, 96, 16, 9}, // out: [1, 24, 32, 18] - first model shape + {1, 4, 5, 7}, // out: [1, 1, 10, 14] - 3 padding channel lanes + {2, 32, 6, 5}, // out: [2, 8, 12, 10] - batch > 1 }; const std::vector same_qparams = {true, false}; return { @@ -82,6 +79,12 @@ REGISTER_TEST_CASE_SET("q8ta/pixel_shuffle", "performance") { {1, 96, 32, 18}, {1, 96, 64, 36}, {1, 96, 128, 72}, + // Upsampling decoder with 64 output channels. + {1, 256, 16, 16}, + {1, 256, 32, 32}, + {1, 256, 64, 64}, + {1, 256, 128, 128}, + {1, 256, 256, 256}, }; const std::vector same_qparams = {true}; // residual-style: scales match return { diff --git a/backends/vulkan/test/ops/benchmark/test_ops/impl/TestQ8taPixelShuffle.cpp b/backends/vulkan/test/ops/benchmark/test_ops/impl/TestQ8taPixelShuffle.cpp index 70dd8c9ca39..c33c9696a33 100644 --- a/backends/vulkan/test/ops/benchmark/test_ops/impl/TestQ8taPixelShuffle.cpp +++ b/backends/vulkan/test/ops/benchmark/test_ops/impl/TestQ8taPixelShuffle.cpp @@ -19,13 +19,14 @@ namespace { // channels-packed int8x4 GPUMemoryLayout enum. // // Note: PACKED_INT8_CONV2D is a Python/serialization-level alias that the -// runtime resolves to kPackedInt8_4C1W (see VulkanBackend.cpp). At the C++ -// runtime layer there is no distinct kPackedInt8_CONV2D enum, so testing -// "CONV2D" here would just be duplicate runtime work over "4C1W". We -// therefore only expose the two real enum values to the test driver. +// runtime resolves to kPackedInt8_4C (im2col convolution inputs) or +// kPackedInt8_4C1W. At the C++ runtime layer there is no distinct +// kPackedInt8_CONV2D enum, so only the resolved layouts are exposed here. utils::GPUMemoryLayout layout_from_string(const std::string& s) { if (s == "4W4C") { return utils::kPackedInt8_4W4C; + } else if (s == "4C") { + return utils::kPackedInt8_4C; } else if (s == "4C1W") { return utils::kPackedInt8_4C1W; } @@ -46,8 +47,8 @@ utils::GPUMemoryLayout layout_from_string(const std::string& s) { // float output_scale, // int output_zp, // int upscale_factor, -// str in_layout, // "4W4C" | "4C1W" -// str out_layout, // "4W4C" | "4C1W" +// str in_layout, // "4W4C" | "4C" | "4C1W" +// str out_layout, // "4W4C" | "4C" | "4C1W" // ) -> Tensor (float, output shape = pixel_shuffle(in, r)) //