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)) //