From 2cf24e468f513d98e152624139d42885c3f0811e Mon Sep 17 00:00:00 2001 From: Stephen Jia Date: Fri, 28 Aug 2026 07:54:39 -0700 Subject: [PATCH] [ET-VK][ops] Add batch support to q8ta convolutions Add batch-aware direct, pointwise, depthwise, and im2col q8ta convolution dispatch. Batched im2col uses an NCHW scratch tensor and batch-strided shader indexing, with conservative groups=1 high-K/small-spatial routing bounded by a 32 MiB scratch cap. Grouped and large-spatial convolutions stay on the direct path. Authored with Codex. Differential Revision: [D117869783](https://our.internmc.facebook.com/intern/diff/D117869783/) ghstack-source-id: 421281491 Pull-Request: https://github.com/pytorch/executorch/pull/22254 --- .../runtime/graph/ops/glsl/q8ta_conv2d.glsl | 17 +++- .../graph/ops/glsl/q8ta_conv2d_dw.glsl | 17 +++- .../graph/ops/glsl/q8ta_conv2d_pw.glsl | 24 ++++-- .../runtime/graph/ops/glsl/q8ta_im2col.glsl | 17 ++-- .../runtime/graph/ops/impl/Q8taConv2d.cpp | 83 +++++++++++++++---- .../runtime/graph/ops/impl/Q8taConv2dDW.cpp | 5 +- .../graph/ops/impl/Q8taConv2dIm2Col.cpp | 19 +++-- .../runtime/graph/ops/impl/Q8taConv2dPW.cpp | 7 +- .../vulkan/test/custom_ops/conv2d_utils.h | 1 + .../test/custom_ops/test_q8ta_conv2d.cpp | 77 +++++++++++++++-- .../test/custom_ops/test_q8ta_conv2d_dw.cpp | 54 ++++++++++-- .../test/custom_ops/test_q8ta_conv2d_pw.cpp | 54 ++++++++++-- 12 files changed, 311 insertions(+), 64 deletions(-) diff --git a/backends/vulkan/runtime/graph/ops/glsl/q8ta_conv2d.glsl b/backends/vulkan/runtime/graph/ops/glsl/q8ta_conv2d.glsl index 821f7f79b0e..644ef466bb6 100644 --- a/backends/vulkan/runtime/graph/ops/glsl/q8ta_conv2d.glsl +++ b/backends/vulkan/runtime/graph/ops/glsl/q8ta_conv2d.glsl @@ -77,14 +77,18 @@ void main() { // Thread mapping int oc4 = int(gl_GlobalInvocationID.z); int w4 = int(gl_GlobalInvocationID.x); + const int H = int(outp.sizes[0][1]); + const int hn = int(gl_GlobalInvocationID.y); + const int n = hn / H; + const int h = hn % H; // Initialize output tensor index (WHCN order) // Each thread handles 4 adjacent widths starting at base_out_w TensorIndex4D outp_tidx; outp_tidx.data[0] = w4 * 4; - outp_tidx.data[1] = int(gl_GlobalInvocationID.y); + outp_tidx.data[1] = h; outp_tidx.data[2] = oc4 * 4; - outp_tidx.data[3] = 0; + outp_tidx.data[3] = n; const int W = int(outp.sizes[0][0]); const int OC = int(outp.sizes[0][2]); @@ -113,6 +117,7 @@ void main() { const int inp_w_stride = int(inp.strides[0][0]); const int inp_h_stride = int(inp.strides[0][1]); const int inp_c_stride = int(inp.strides[0][2]); + const int inp_n_stride = int(inp.strides[0][3]); const int w_texel_step = conv2d_params.dilation.x * inp_w_stride; const int h_texel_step = conv2d_params.dilation.y * inp_h_stride; const int subtile_w_step = conv2d_params.stride.x * inp_w_stride; @@ -122,7 +127,7 @@ void main() { inp_tidx.data[0] = outp_tidx.data[0] * conv2d_params.stride.x - conv2d_params.padding.x; inp_tidx.data[1] = outp_tidx.data[1] * conv2d_params.stride.y - conv2d_params.padding.y; inp_tidx.data[2] = ic_group_start; - inp_tidx.data[3] = 0; + inp_tidx.data[3] = n; int base_inp_texel_idx; if (get_outer_packed_dim_block_size(inp_layout) == 1) { @@ -172,7 +177,11 @@ void main() { // inp_texel_idx = tensor4d_idx_to_texel_idx(inp, inp_tidx, inp_layout); const int w4 = div_4(inp_tidx.data[0]); const int inp_c4 = div_4(inp_tidx.data[2]); - inp_texel_idx = (inp_tidx.data[1] * inp_h_stride + w4 * inp_w_stride + inp_c4) * 4 + mod_4(inp_tidx.data[0]); + inp_texel_idx = + (n * inp_n_stride + inp_tidx.data[1] * inp_h_stride + + w4 * inp_w_stride + inp_c4) * + 4 + + mod_4(inp_tidx.data[0]); } packed_input = t_packed_int8_input[inp_texel_idx]; } diff --git a/backends/vulkan/runtime/graph/ops/glsl/q8ta_conv2d_dw.glsl b/backends/vulkan/runtime/graph/ops/glsl/q8ta_conv2d_dw.glsl index 7f4d03887df..673cfd68952 100644 --- a/backends/vulkan/runtime/graph/ops/glsl/q8ta_conv2d_dw.glsl +++ b/backends/vulkan/runtime/graph/ops/glsl/q8ta_conv2d_dw.glsl @@ -71,14 +71,18 @@ ivec4 quantize(const vec4 texel, const float inv_scale, const int zp) { void main() { const int c4 = int(gl_GlobalInvocationID.z); + const int H = int(outp.sizes[0][1]); + const int hn = int(gl_GlobalInvocationID.y); + const int n = hn / H; + const int h = hn % H; // Initialize output tensor index (WHCN order) // Each thread handles 4 adjacent widths starting at base_out_w TensorIndex4D outp_tidx; outp_tidx.data[0] = int(gl_GlobalInvocationID.x) * 4; - outp_tidx.data[1] = int(gl_GlobalInvocationID.y); + outp_tidx.data[1] = h; outp_tidx.data[2] = c4 * 4; - outp_tidx.data[3] = 0; + outp_tidx.data[3] = n; const int W = int(outp.sizes[0][0]); const int C4 = int(div_up_4(outp.sizes[0][2])); @@ -94,6 +98,7 @@ void main() { // Get strides for width and height dimensions (in texel space) const int w_stride = int(inp.strides[0][0]); const int h_stride = int(inp.strides[0][1]); + const int n_stride = int(inp.strides[0][3]); // Pre-compute step sizes for efficient indexing const int w_texel_step = conv2d_params.dilation.x * w_stride; @@ -106,7 +111,7 @@ void main() { inp_tidx.data[0] = outp_tidx.data[0] * conv2d_params.stride.x - conv2d_params.padding.x; inp_tidx.data[1] = outp_tidx.data[1] * conv2d_params.stride.y - conv2d_params.padding.y; inp_tidx.data[2] = outp_tidx.data[2]; - inp_tidx.data[3] = 0; // batch = 0 since N == 1 + inp_tidx.data[3] = n; int base_inp_texel_idx; if (get_outer_packed_dim_block_size(inp_layout) == 1) { @@ -152,7 +157,11 @@ void main() { // inp_texel_idx = base_inp_texel_idx + div_4(w_offset) * w_stride + mod_4(w_offset); // inp_texel_idx = tensor4d_idx_to_texel_idx(inp, inp_tidx, inp_layout); const int w4 = div_4(inp_tidx.data[0]); - inp_texel_idx = (inp_tidx.data[1] * h_stride + w4 * w_stride + c4) * 4 + mod_4(inp_tidx.data[0]); + inp_texel_idx = + (n * n_stride + inp_tidx.data[1] * h_stride + + w4 * w_stride + c4) * + 4 + + mod_4(inp_tidx.data[0]); } const int packed_input = t_packed_int8_input[inp_texel_idx]; input_4c = unpack_int8x4(packed_input); diff --git a/backends/vulkan/runtime/graph/ops/glsl/q8ta_conv2d_pw.glsl b/backends/vulkan/runtime/graph/ops/glsl/q8ta_conv2d_pw.glsl index 9f60bea9948..aeb98f7a41b 100644 --- a/backends/vulkan/runtime/graph/ops/glsl/q8ta_conv2d_pw.glsl +++ b/backends/vulkan/runtime/graph/ops/glsl/q8ta_conv2d_pw.glsl @@ -73,14 +73,17 @@ ${layout_declare_spec_const(C, "int", "inp_layout", "CONTIG_LAYOUT_INT")} int compute_outp_buffer_idx( const int w_block_idx, const int h_idx, - const int c_block_idx) { + const int c_block_idx, + const int n_idx) { if (get_outer_packed_dim_block_size(outp_layout) == 1) { - return h_idx * int(outp.strides[0][1]) + return n_idx * int(outp.strides[0][3]) + + h_idx * int(outp.strides[0][1]) + mul_4(w_block_idx) * int(outp.strides[0][0]) + c_block_idx * int(outp.strides[0][2]); } else { return mul_4( - h_idx * int(outp.strides[0][1]) + n_idx * int(outp.strides[0][3]) + + h_idx * int(outp.strides[0][1]) + w_block_idx * int(outp.strides[0][0]) + c_block_idx * int(outp.strides[0][2])); } @@ -91,20 +94,22 @@ void main() { // Thread mapping: each thread handles TILE_M widths x TILE_N output channels. // gl_GlobalInvocationID.x -> output channel blocks. // gl_GlobalInvocationID.y -> width blocks. - // gl_GlobalInvocationID.z → batch (or height * batch combined) + // gl_GlobalInvocationID.z -> height * batch. const int oc_block_idx = int(gl_GlobalInvocationID.x) * TILE_N4; const int ow_block_idx = int(gl_GlobalInvocationID.y) * TILE_M4; - const int oh = int(gl_GlobalInvocationID.z); // Get output extents in block space (div_up_4 for packed dimensions) const int W = int(outp.sizes[0][0]); const int W4 = div_up_4(int(outp.sizes[0][0])); const int H = int(outp.sizes[0][1]); const int OC4 = div_up_4(int(outp.sizes[0][2])); + const int hn = int(gl_GlobalInvocationID.z); + const int n = hn / H; + const int oh = hn % H; // Bounds check in block space if (ow_block_idx >= W4 || - oh >= H || + n >= int(outp.sizes[0][3]) || oc_block_idx >= OC4) { return; } @@ -118,6 +123,7 @@ void main() { const int inp_w_stride = int(inp.strides[0][0]); const int inp_h_stride = int(inp.strides[0][1]); const int inp_c_stride = int(inp.strides[0][2]); + const int inp_n_stride = int(inp.strides[0][3]); // Initialize int32 accumulator ivec4 out_accum[TILE_M][TILE_N4]; @@ -133,7 +139,8 @@ void main() { // Compute initial input tile index with group offset // For grouped im2col, each group's K range starts at group_idx * K4_per_group // For non-grouped (groups=1), group_idx is always 0 so offset is 0 - int input_idx = oh * inp_h_stride + int input_idx = n * inp_n_stride + + oh * inp_h_stride + ow_block_idx * inp_w_stride + group_idx * K4_per_group; @@ -256,7 +263,8 @@ void main() { const int base_outp_buffer_idx = compute_outp_buffer_idx( ow_block_idx + m4, oh, - oc_block_idx + n4); + oc_block_idx + n4, + n); if (oc_block_idx + n4 < OC4) { // Store individual ints from the ivec4 const int subtile_w_limit = min(4, W - mul_4(ow_block_idx + m4)); diff --git a/backends/vulkan/runtime/graph/ops/glsl/q8ta_im2col.glsl b/backends/vulkan/runtime/graph/ops/glsl/q8ta_im2col.glsl index 314788cb857..b0cc4866a03 100644 --- a/backends/vulkan/runtime/graph/ops/glsl/q8ta_im2col.glsl +++ b/backends/vulkan/runtime/graph/ops/glsl/q8ta_im2col.glsl @@ -50,15 +50,19 @@ void main() { const int im2col_W4 = div_up_4(im2col_sizes.x); const int im2col_H = im2col_sizes.y; const int im2col_Z4 = div_up_4(im2col_sizes.z); + const int im2col_N = im2col_sizes.w; // im2col block index from linear output buffer index const int c4_idx = out_buf_idx % im2col_Z4; const int row = out_buf_idx / im2col_Z4; const int w4_idx = row % im2col_W4; - const int h_idx = row / im2col_W4; + const int hn_idx = row / im2col_W4; + const int h_idx = hn_idx % im2col_H; + const int n_idx = hn_idx / im2col_H; // out of bounds check - if (w4_idx >= im2col_W4 || h_idx >= im2col_H || c4_idx >= im2col_Z4) { + if (w4_idx >= im2col_W4 || h_idx >= im2col_H || + c4_idx >= im2col_Z4 || n_idx >= im2col_N) { return; } @@ -108,12 +112,14 @@ void main() { const int x_mod = mod_4(x); int scalar_idx; if (get_outer_packed_dim_block_size(inp_layout) == 1) { - scalar_idx = input_y * int(inp.strides[0][1]) + scalar_idx = n_idx * int(inp.strides[0][3]) + + input_y * int(inp.strides[0][1]) + x * int(inp.strides[0][0]) + z4 * int(inp.strides[0][2]); } else { scalar_idx = mul_4( - input_y * int(inp.strides[0][1]) + n_idx * int(inp.strides[0][3]) + + input_y * int(inp.strides[0][1]) + x4 * int(inp.strides[0][0]) + z4) + x_mod; } @@ -122,7 +128,8 @@ void main() { } // store_packed_int8_output_tile (with TILE_M4=1, TILE_N4=1) - const int buffer_idx = h_idx * int(im2col_outp.strides[0][1]) + const int buffer_idx = n_idx * int(im2col_outp.strides[0][3]) + + h_idx * int(im2col_outp.strides[0][1]) + w4_idx * int(im2col_outp.strides[0][0]) + c4_idx; diff --git a/backends/vulkan/runtime/graph/ops/impl/Q8taConv2d.cpp b/backends/vulkan/runtime/graph/ops/impl/Q8taConv2d.cpp index 0f00b5e80f0..70171a04820 100644 --- a/backends/vulkan/runtime/graph/ops/impl/Q8taConv2d.cpp +++ b/backends/vulkan/runtime/graph/ops/impl/Q8taConv2d.cpp @@ -42,6 +42,52 @@ bool q8ta_conv2d_check_4w4c_packed_dim_info(const api::PackedDimInfo& info) { info.outer_packed_dim_block_size == 4; } +namespace { + +uint64_t q8ta_conv2d_im2col_scratch_limit(ComputeGraph& graph) { + constexpr uint64_t kMaxBatchedIm2ColScratchBytes = 32ULL * 1024ULL * 1024ULL; + const uint64_t device_scratch_limit = + graph.context()->adapter_ptr()->max_buffer_numel(); + return device_scratch_limit < kMaxBatchedIm2ColScratchBytes + ? device_scratch_limit + : kMaxBatchedIm2ColScratchBytes; +} + +bool should_use_q8ta_conv2d_im2col( + ComputeGraph& graph, + const int64_t batch, + const int64_t groups, + const int64_t in_channels_per_group, + const int64_t flattened_kernel_size, + const int64_t out_height, + const int64_t out_width) { + const bool im2col_eligible = in_channels_per_group % 4 == 0; + if (!im2col_eligible) { + return false; + } + + const int64_t spatial_out = out_height * out_width; + if (batch > 1) { + constexpr int64_t kMinFlattenedKernelSize = 1024; + constexpr int64_t kMaxSpatialOutput = 64; + const uint64_t scratch_bytes = static_cast(batch) * + static_cast(flattened_kernel_size) * + static_cast(out_height) * + static_cast(utils::align_up_4(out_width)); + return groups == 1 && flattened_kernel_size >= kMinFlattenedKernelSize && + spatial_out <= kMaxSpatialOutput && + scratch_bytes <= q8ta_conv2d_im2col_scratch_limit(graph); + } + + if (graph.device_is_mali()) { + return true; + } + + return groups == 1 && (in_channels_per_group >= 32 || spatial_out <= 4096); +} + +} // namespace + // // Workgroup size selection functions // @@ -70,13 +116,16 @@ GlobalWorkGrid pick_q8ta_conv2d_gwg( const uint32_t W = graph->size_at(-1, output); const uint32_t H = graph->size_at(-2, output); const uint32_t C = graph->size_at(-3, output); + const uint32_t N = graph->size_at(-4, output); // Each thread processes 4 adjacent width positions and 4 channels (4Wx4C // tile) const uint32_t W4 = utils::div_up_4(W); const uint32_t C4 = utils::div_up_4(C); - return GlobalWorkGrid({W4, H, C4}, kTiledWorkGrid); + return GlobalWorkGrid( + {W4, utils::safe_downcast(static_cast(H) * N), C4}, + kTiledWorkGrid); } /** @@ -466,33 +515,33 @@ void q8ta_conv2d_general( void q8ta_conv2d(ComputeGraph& graph, const std::vector& args) { const ValueRef input = args.at(0); + const ValueRef kernel_size_ref = args.at(9); const ValueRef groups_ref = args.at(13); const ValueRef output = args.at(15); const int64_t groups = graph.extract_scalar(groups_ref); const int64_t in_channels = graph.size_at(-3, input); const int64_t in_channels_per_group = in_channels / groups; + const int64_t batch = graph.size_at(-4, input); const int64_t H_out = graph.size_at(-2, output); const int64_t W_out = graph.size_at(-1, output); - const int64_t spatial_out = H_out * W_out; - - // Im2col requires input channels per group to be a multiple of 4 - const bool im2col_eligible = in_channels_per_group % 4 == 0; - - bool use_im2col = false; - if (graph.device_is_mali()) { - // On Mali, im2col is faster than the general shader across the board. - use_im2col = im2col_eligible; - } else { - // Default: on Adreno and unknown GPU architectures, im2col is only - // beneficial for ungrouped convolutions with sufficient channel depth or - // small spatial output. For grouped convolutions, the general shader is - // more efficient (0.7-0.95x regression measured on Adreno). - use_im2col = im2col_eligible && groups == 1 && - (in_channels_per_group >= 32 || spatial_out <= 4096); + int64_t flattened_kernel_size; + { + const auto kernel_size = graph.get_int_list(kernel_size_ref); + flattened_kernel_size = utils::align_up_4( + in_channels_per_group * kernel_size->at(0) * kernel_size->at(1)); } + const bool use_im2col = should_use_q8ta_conv2d_im2col( + graph, + batch, + groups, + in_channels_per_group, + flattened_kernel_size, + H_out, + W_out); + if (use_im2col) { q8ta_conv2d_im2col(graph, args); } else { diff --git a/backends/vulkan/runtime/graph/ops/impl/Q8taConv2dDW.cpp b/backends/vulkan/runtime/graph/ops/impl/Q8taConv2dDW.cpp index 505e07bd9e5..182e8d684d2 100644 --- a/backends/vulkan/runtime/graph/ops/impl/Q8taConv2dDW.cpp +++ b/backends/vulkan/runtime/graph/ops/impl/Q8taConv2dDW.cpp @@ -34,13 +34,16 @@ GlobalWorkGrid pick_q8ta_conv2d_dw_gwg( const uint32_t W = graph->size_at(-1, output); const uint32_t H = graph->size_at(-2, output); const uint32_t C = graph->size_at(-3, output); + const uint32_t N = graph->size_at(-4, output); // Each thread processes 4 adjacent width positions and 4 channels (4Wx4C // tile) const uint32_t W4 = utils::div_up_4(W); const uint32_t C4 = utils::div_up_4(C); - return GlobalWorkGrid({W4, H, C4}, kTiledWorkGrid); + return GlobalWorkGrid( + {W4, utils::safe_downcast(static_cast(H) * N), C4}, + kTiledWorkGrid); } LocalWorkGroup pick_q8ta_conv2d_dw_lwg( diff --git a/backends/vulkan/runtime/graph/ops/impl/Q8taConv2dIm2Col.cpp b/backends/vulkan/runtime/graph/ops/impl/Q8taConv2dIm2Col.cpp index f8d630c2250..e93723c5125 100644 --- a/backends/vulkan/runtime/graph/ops/impl/Q8taConv2dIm2Col.cpp +++ b/backends/vulkan/runtime/graph/ops/impl/Q8taConv2dIm2Col.cpp @@ -32,16 +32,17 @@ GlobalWorkGrid pick_q8ta_im2col_gwg( const ValueRef im2col_output = args.at(0).refs.at(0); - std::vector im2col_sizes = graph->sizes_of(im2col_output); - const uint32_t K = utils::safe_downcast(im2col_sizes[0]); - const uint32_t H = utils::safe_downcast(im2col_sizes[1]); - const uint32_t W = utils::safe_downcast(im2col_sizes[2]); + const uint32_t N = graph->size_at(-4, im2col_output); + const uint32_t K = graph->size_at(-3, im2col_output); + const uint32_t H = graph->size_at(-2, im2col_output); + const uint32_t W = graph->size_at(-1, im2col_output); const uint32_t K4 = utils::div_up_4(K); const uint32_t W4 = utils::div_up_4(W); // Each thread handles one 4x4 block in the output - return graph->create_linear_gwg(K4 * W4 * H); + return graph->create_linear_gwg( + utils::safe_downcast(static_cast(K4) * W4 * H * N)); } LocalWorkGroup pick_q8ta_im2col_lwg( @@ -70,6 +71,7 @@ std::vector calculate_q8ta_im2col_sizes( const ValueRef& kernel_size, const ValueRef& groups) { std::vector in_sizes = graph->sizes_of(input); + const int64_t batch = utils::val_at(-4, in_sizes); const int64_t in_channels = utils::val_at(-3, in_sizes); std::vector out_sizes = graph->sizes_of(output); @@ -93,7 +95,7 @@ std::vector calculate_q8ta_im2col_sizes( const int64_t W = utils::align_up_4(out_width); const int64_t H = out_height; - return {K, H, W}; + return {batch, K, H, W}; } // @@ -102,7 +104,7 @@ std::vector calculate_q8ta_im2col_sizes( // resize_args = { input, kernel_size, stride, padding, dilation, groups } // -// The im2col scratch tensor is [K, H_out, align_up_4(W_out)] where K (the +// The im2col scratch tensor is [N, K, H_out, align_up_4(W_out)] where K (the // flattened conv window, channel/kernel-derived) is shape-independent and // H_out/W_out are the conv output spatial dims. The downstream PW GEMM that // consumes this scratch is resized separately (it preserves H/W). Without this, @@ -122,6 +124,7 @@ void resize_q8ta_im2col_node( const ValueRef groups = resize_args.at(5); const std::vector in_sizes = graph->sizes_of(in); + const int64_t batch = utils::val_at(-4, in_sizes); // Conv output H/W from the current input. const std::vector out_hw = calc_out_sizes_hw( @@ -146,7 +149,7 @@ void resize_q8ta_im2col_node( const int64_t K = flattened_kernel_len * groups_val; const int64_t W = utils::align_up_4(out_width); - graph->virtual_resize(im2col_out, {K, out_height, W}); + graph->virtual_resize(im2col_out, {batch, K, out_height, W}); } // diff --git a/backends/vulkan/runtime/graph/ops/impl/Q8taConv2dPW.cpp b/backends/vulkan/runtime/graph/ops/impl/Q8taConv2dPW.cpp index 99ddd74117b..ee234319e8c 100644 --- a/backends/vulkan/runtime/graph/ops/impl/Q8taConv2dPW.cpp +++ b/backends/vulkan/runtime/graph/ops/impl/Q8taConv2dPW.cpp @@ -33,6 +33,7 @@ GlobalWorkGrid pick_q8ta_conv2d_pw_gwg( const uint32_t W = graph->size_at(-1, output); const uint32_t H = graph->size_at(-2, output); const uint32_t C = graph->size_at(-3, output); + const uint32_t N = graph->size_at(-4, output); // Each thread covers a 4-width x 4-channel output block. // Tile constants must match TILE_M4 / TILE_N4 in q8ta_conv2d_pw.glsl. @@ -45,9 +46,11 @@ GlobalWorkGrid pick_q8ta_conv2d_pw_gwg( // Global workgroup size: // x = output channels / (TILE_N4 * 4) = C4 / TILE_N4 = C4 // y = width / (TILE_M4 * 4) = W4 / TILE_M4 = W4 - // z = height + // z = height * batch return GlobalWorkGrid( - {utils::div_up(C4, TILE_N4), utils::div_up(W4, TILE_M4), H}, + {utils::div_up(C4, TILE_N4), + utils::div_up(W4, TILE_M4), + utils::safe_downcast(static_cast(H) * N)}, kTiledWorkGrid); } diff --git a/backends/vulkan/test/custom_ops/conv2d_utils.h b/backends/vulkan/test/custom_ops/conv2d_utils.h index 416f6c50061..34853ec383a 100644 --- a/backends/vulkan/test/custom_ops/conv2d_utils.h +++ b/backends/vulkan/test/custom_ops/conv2d_utils.h @@ -68,6 +68,7 @@ struct Conv2dConfig { Padding padding; Dilation dilation; int32_t groups; // Number of groups for grouped convolution + int32_t batch = 1; std::string test_case_name = "placeholder"; std::string op_name = "conv2d"; diff --git a/backends/vulkan/test/custom_ops/test_q8ta_conv2d.cpp b/backends/vulkan/test/custom_ops/test_q8ta_conv2d.cpp index 5b1de615d29..b30212feb72 100644 --- a/backends/vulkan/test/custom_ops/test_q8ta_conv2d.cpp +++ b/backends/vulkan/test/custom_ops/test_q8ta_conv2d.cpp @@ -36,9 +36,12 @@ static TestCase create_test_case_from_config( int64_t H_out = config.get_output_height(); int64_t W_out = config.get_output_width(); - // Input tensor (float/half) - [1, C_in, H_in, W_in] (batch size always 1) + // Input tensor (float/half) - [N, C_in, H_in, W_in] std::vector input_size = { - 1, config.channels.in, config.input_size.h, config.input_size.w}; + config.batch, + config.channels.in, + config.input_size.h, + config.input_size.w}; utils::GPUMemoryLayout fp_memory_layout = fp_storage_type == utils::kBuffer ? utils::kWidthPacked @@ -47,7 +50,8 @@ static TestCase create_test_case_from_config( // Create test case name std::string prefix = config.test_case_name.substr(0, 4); // "ACCU" or "PERF" std::string dtype_str = dtype_short(input_dtype); - std::string in_shape = "[1," + std::to_string(config.channels.in) + "," + + std::string in_shape = "[" + std::to_string(config.batch) + "," + + std::to_string(config.channels.in) + "," + std::to_string(config.input_size.h) + "," + std::to_string(config.input_size.w) + "]"; std::string weight_shape = "[" + std::to_string(config.channels.out) + "," + @@ -159,9 +163,9 @@ static TestCase create_test_case_from_config( // Kernel size parameters ValueSpec kernel_size({config.kernel.h, config.kernel.w}); - // Output tensor (float/half) - [1, C_out, H_out, W_out] (batch size always 1) + // Output tensor (float/half) - [N, C_out, H_out, W_out] ValueSpec output( - {1, config.channels.out, H_out, W_out}, + {config.batch, config.channels.out, H_out, W_out}, input_dtype, fp_storage_type, fp_memory_layout, @@ -475,6 +479,69 @@ static std::vector generate_quantized_conv2d_test_cases() { } } + std::vector batch_configs = { + {OutInChannels(16, 32), + InputSize2D(7, 7), + KernelSize(3, 3), + Stride(1, 1), + Padding(1, 1), + Dilation(1, 1), + 1, + 2}, + {OutInChannels(32, 3), + InputSize2D(256, 256), + KernelSize(3, 3), + Stride(2, 2), + Padding(1, 1), + Dilation(1, 1), + 1, + 1}, + {OutInChannels(32, 3), + InputSize2D(256, 256), + KernelSize(3, 3), + Stride(2, 2), + Padding(1, 1), + Dilation(1, 1), + 1, + 60}, + {OutInChannels(512, 256), + InputSize2D(10, 13), + KernelSize(3, 3), + Stride(2, 2), + Padding(1, 1), + Dilation(1, 1), + 1, + 60}}; + + for (auto& config : batch_configs) { + const bool is_performance = config.batch > kRefDimSizeLimit || + config.channels.out > kRefDimSizeLimit || + config.channels.in > kRefDimSizeLimit || + config.input_size.h > kRefDimSizeLimit || + config.input_size.w > kRefDimSizeLimit; + config.op_name = "conv2d_q8ta_q8csw_q8to"; + config.test_case_name = make_test_case_name( + config, is_performance, utils::kTexture3D, utils::kBuffer); + test_cases.push_back(create_test_case_from_config( + config, vkapi::kFloat, utils::kTexture3D, utils::kPackedInt8_4C1W)); + if (config.batch == 2) { + test_cases.push_back(create_test_case_from_config( + config, + vkapi::kFloat, + utils::kTexture3D, + utils::kPackedInt8_4C1W, + /*impl_selector=*/"im2col")); + test_cases.push_back(create_test_case_from_config( + config, vkapi::kFloat, utils::kTexture3D, utils::kPackedInt8_4W4C)); + test_cases.push_back(create_test_case_from_config( + config, + vkapi::kFloat, + utils::kTexture3D, + utils::kPackedInt8_4W4C, + /*impl_selector=*/"im2col")); + } + } + return test_cases; } diff --git a/backends/vulkan/test/custom_ops/test_q8ta_conv2d_dw.cpp b/backends/vulkan/test/custom_ops/test_q8ta_conv2d_dw.cpp index 9813eeaa9d6..2dbb4909adb 100644 --- a/backends/vulkan/test/custom_ops/test_q8ta_conv2d_dw.cpp +++ b/backends/vulkan/test/custom_ops/test_q8ta_conv2d_dw.cpp @@ -37,9 +37,12 @@ TestCase create_test_case_from_config( int64_t H_out = config.get_output_height(); int64_t W_out = config.get_output_width(); - // Input tensor (float/half) - [1, C_in, H_in, W_in] (batch size always 1) + // Input tensor (float/half) - [N, C_in, H_in, W_in] std::vector input_size = { - 1, config.channels.in, config.input_size.h, config.input_size.w}; + config.batch, + config.channels.in, + config.input_size.h, + config.input_size.w}; utils::GPUMemoryLayout fp_memory_layout = fp_storage_type == utils::kBuffer ? utils::kWidthPacked @@ -48,7 +51,8 @@ TestCase create_test_case_from_config( // Create test case name std::string prefix = config.test_case_name.substr(0, 4); // "ACCU" or "PERF" std::string dtype_str = dtype_short(input_dtype); - std::string in_shape = "[1," + std::to_string(config.channels.in) + "," + + std::string in_shape = "[" + std::to_string(config.batch) + "," + + std::to_string(config.channels.in) + "," + std::to_string(config.input_size.h) + "," + std::to_string(config.input_size.w) + "]"; // depthwise: weight is [C_out, 1, K_h, K_w] @@ -168,9 +172,9 @@ TestCase create_test_case_from_config( // Kernel size parameters ValueSpec kernel_size({config.kernel.h, config.kernel.w}); - // Output tensor (float/half) - [1, C_out, H_out, W_out] (batch size always 1) + // Output tensor (float/half) - [N, C_out, H_out, W_out] ValueSpec output( - {1, config.channels.out, H_out, W_out}, + {config.batch, config.channels.out, H_out, W_out}, input_dtype, fp_storage_type, fp_memory_layout, @@ -395,6 +399,46 @@ std::vector generate_quantized_conv2d_dw_test_cases() { } } + std::vector batch_configs = { + {OutInChannels(8, 8), + InputSize2D(8, 8), + KernelSize(3, 3), + Stride(1, 1), + Padding(1, 1), + Dilation(1, 1), + 8, + 2}, + {OutInChannels(128, 128), + InputSize2D(64, 64), + KernelSize(5, 5), + Stride(2, 2), + Padding(2, 2), + Dilation(1, 1), + 128, + 1}, + {OutInChannels(128, 128), + InputSize2D(64, 64), + KernelSize(5, 5), + Stride(2, 2), + Padding(2, 2), + Dilation(1, 1), + 128, + 60}}; + + for (auto& config : batch_configs) { + const bool is_performance = config.channels.out > kRefDimSizeLimit || + config.channels.in > kRefDimSizeLimit; + config.op_name = "conv2d_q8ta_q8csw_q8to"; + config.test_case_name = make_test_case_name( + config, is_performance, utils::kTexture3D, utils::kBuffer); + test_cases.push_back(create_test_case_from_config( + config, vkapi::kFloat, utils::kTexture3D, utils::kPackedInt8_4C1W)); + if (config.batch == 2) { + test_cases.push_back(create_test_case_from_config( + config, vkapi::kFloat, utils::kTexture3D, utils::kPackedInt8_4W4C)); + } + } + return test_cases; } diff --git a/backends/vulkan/test/custom_ops/test_q8ta_conv2d_pw.cpp b/backends/vulkan/test/custom_ops/test_q8ta_conv2d_pw.cpp index d2b84101940..ee7d8c9e5bf 100644 --- a/backends/vulkan/test/custom_ops/test_q8ta_conv2d_pw.cpp +++ b/backends/vulkan/test/custom_ops/test_q8ta_conv2d_pw.cpp @@ -36,9 +36,12 @@ static TestCase create_test_case_from_config( int64_t H_out = config.get_output_height(); int64_t W_out = config.get_output_width(); - // Input tensor (float/half) - [1, C_in, H_in, W_in] (batch size always 1) + // Input tensor (float/half) - [N, C_in, H_in, W_in] std::vector input_size = { - 1, config.channels.in, config.input_size.h, config.input_size.w}; + config.batch, + config.channels.in, + config.input_size.h, + config.input_size.w}; utils::GPUMemoryLayout fp_memory_layout = fp_storage_type == utils::kBuffer ? utils::kWidthPacked @@ -47,7 +50,8 @@ static TestCase create_test_case_from_config( // Create test case name std::string prefix = config.test_case_name.substr(0, 4); // "ACCU" or "PERF" std::string dtype_str = dtype_short(input_dtype); - std::string in_shape = "[1," + std::to_string(config.channels.in) + "," + + std::string in_shape = "[" + std::to_string(config.batch) + "," + + std::to_string(config.channels.in) + "," + std::to_string(config.input_size.h) + "," + std::to_string(config.input_size.w) + "]"; std::string weight_shape = "[" + std::to_string(config.channels.out) + "," + @@ -160,9 +164,9 @@ static TestCase create_test_case_from_config( // Kernel size parameters ValueSpec kernel_size({config.kernel.h, config.kernel.w}); - // Output tensor (float/half) - [1, C_out, H_out, W_out] (batch size always 1) + // Output tensor (float/half) - [N, C_out, H_out, W_out] ValueSpec output( - {1, config.channels.out, H_out, W_out}, + {config.batch, config.channels.out, H_out, W_out}, input_dtype, fp_storage_type, fp_memory_layout, @@ -353,6 +357,46 @@ static std::vector generate_quantized_conv2d_pw_test_cases() { } } + std::vector batch_configs = { + {OutInChannels(8, 8), + InputSize2D(8, 8), + KernelSize(1, 1), + Stride(1, 1), + Padding(0, 0), + Dilation(1, 1), + 1, + 2}, + {OutInChannels(64, 32), + InputSize2D(128, 128), + KernelSize(1, 1), + Stride(1, 1), + Padding(0, 0), + Dilation(1, 1), + 1, + 1}, + {OutInChannels(64, 32), + InputSize2D(128, 128), + KernelSize(1, 1), + Stride(1, 1), + Padding(0, 0), + Dilation(1, 1), + 1, + 60}}; + + for (auto& config : batch_configs) { + const bool is_performance = config.input_size.h > kRefDimSizeLimit || + config.input_size.w > kRefDimSizeLimit; + config.op_name = "conv2d_q8ta_q8csw_q8to"; + config.test_case_name = make_test_case_name( + config, is_performance, utils::kTexture3D, utils::kBuffer); + test_cases.push_back(create_test_case_from_config( + config, vkapi::kFloat, utils::kTexture3D, utils::kPackedInt8_4C1W)); + if (config.batch == 2) { + test_cases.push_back(create_test_case_from_config( + config, vkapi::kFloat, utils::kTexture3D, utils::kPackedInt8_4W4C)); + } + } + return test_cases; }