From 226d219813c4e073382b0fb49231b9f1745d40af Mon Sep 17 00:00:00 2001 From: Vivek Trivedi <5340687+trivedivivek@users.noreply.github.com> Date: Mon, 27 Jan 2025 09:21:35 -0800 Subject: [PATCH 1/3] [ET-VK] Using push constants for conv2d pw. Pull Request resolved: https://github.com/pytorch/executorch/pull/7814 This diff is related to the use of push constants for convolutional pw (pointwise) in Executorch's Vulkan backend. This optimization improves performance and memory usage. ghstack-source-id: 263238730 @exported-using-ghexport Differential Revision: [D68400677](https://our.internmc.facebook.com/intern/diff/D68400677/) --- .../runtime/graph/ops/glsl/conv2d_pw.glsl | 23 +++-- .../runtime/graph/ops/impl/Convolution.cpp | 83 ++++++++++++++----- 2 files changed, 78 insertions(+), 28 deletions(-) diff --git a/backends/vulkan/runtime/graph/ops/glsl/conv2d_pw.glsl b/backends/vulkan/runtime/graph/ops/glsl/conv2d_pw.glsl index f72c487fa78..0413eb7b7aa 100644 --- a/backends/vulkan/runtime/graph/ops/glsl/conv2d_pw.glsl +++ b/backends/vulkan/runtime/graph/ops/glsl/conv2d_pw.glsl @@ -24,11 +24,20 @@ ${layout_declare_tensor(0, "w", "t_out", DTYPE, "texture3d")} ${layout_declare_tensor(1, "r", "t_in", DTYPE, "texture3d")} ${layout_declare_tensor(2, "r", "t_kernel", DTYPE, "texture2d")} ${layout_declare_tensor(3, "r", "t_bias", DTYPE, "texture2d")} -${layout_declare_ubo(4, "ivec3", "out_limits")} -${layout_declare_ubo(5, "ivec4", "in_sizes")} -${layout_declare_ubo(6, "ivec2", "kernel_size", "ivec2", "stride", "ivec2", "padding", "ivec2", "dilation")} -${layout_declare_ubo(7, "ivec2", "overlay_region", "int", "in_group_size")} -${layout_declare_ubo(8, "float", "out_min", "float", "out_max")} + +layout(push_constant) uniform restrict Block { + ivec4 out_limits; + ivec4 in_sizes; + ivec2 kernel_size; + ivec2 stride; + ivec2 padding; + ivec2 dilation; + ivec2 overlay_region; + int in_group_size; + int dummy_padding; + float out_min; + float out_max; +}; layout(local_size_x_id = 0, local_size_y_id = 1, local_size_z_id = 2) in; @@ -70,7 +79,7 @@ void main() { // If the top left position is out of bounds, then this invocation will have // no work to do. - if (any(greaterThanEqual(ivec3(pos[0], gpos.z), out_limits))) { + if (any(greaterThanEqual(ivec3(pos[0], gpos.z), out_limits.xyz))) { return; } @@ -144,7 +153,7 @@ void main() { for (int i = 0; i < TILE_SIZE * TILE_SIZE; ++i) { const ivec2 pos = pos_shared[(shared_mem_stride * i) + gl_LocalInvocationIndex]; - if (all(lessThan(ivec3(pos, gpos.z), out_limits))) { + if (all(lessThan(ivec3(pos, gpos.z), out_limits.xyz))) { imageStore(t_out, ivec3(pos, gpos.z), op(sum[i], out_min, out_max)); } } diff --git a/backends/vulkan/runtime/graph/ops/impl/Convolution.cpp b/backends/vulkan/runtime/graph/ops/impl/Convolution.cpp index 8c369914c1b..3c367f334d9 100644 --- a/backends/vulkan/runtime/graph/ops/impl/Convolution.cpp +++ b/backends/vulkan/runtime/graph/ops/impl/Convolution.cpp @@ -407,27 +407,68 @@ void add_conv2d_node( wg_size = {wg_size[0] * wg_size[1] * wg_size[2], 1, 1}; } - graph.execute_nodes().emplace_back(new DispatchNode( - graph, - shader, - wg_size, - graph.create_local_wg_size(wg_size), - // Inputs and Outputs - {{out, vkapi::MemoryAccessType::WRITE}, - {{in, arg_weight, arg_bias}, vkapi::MemoryAccessType::READ}}, - // Shader params buffers - { - t_out->logical_limits_ubo(), - t_in->sizes_ubo(), - graph.create_params_buffer(kernel_params), - graph.create_params_buffer(extra_params), - graph.create_params_buffer(out_params), - }, - // Specialization Constants - {}, - // Resizing Logic - resize_conv2d_node, - {weight_data, stride, padding, dilation, transposed, output_padding})); + if (method == Conv2dMethod::Pointwise) { + const utils::ivec4 kernel_param_size_stride = { + kernel_params.kernel_size[0], + kernel_params.kernel_size[1], + kernel_params.stride[0], + kernel_params.stride[1]}; + + const utils::ivec4 kernel_param_pad_dial = { + kernel_params.padding[0], + kernel_params.padding[1], + kernel_params.dilation[0], + kernel_params.dilation[1]}; + + graph.execute_nodes().emplace_back(new DispatchNode( + graph, + shader, + wg_size, + graph.create_local_wg_size(wg_size), + // Inputs and Outputs + {{out, vkapi::MemoryAccessType::WRITE}, + {{in, arg_weight, arg_bias}, vkapi::MemoryAccessType::READ}}, + // Shader params buffers + {}, + // Specialization Constants + {}, + // Resizing Logic + resize_conv2d_node, + {weight_data, stride, padding, dilation, transposed, output_padding}, + { + graph.logical_limits_pc_of(out), + graph.sizes_pc_of(in), + PushConstantDataInfo( + &kernel_param_size_stride, sizeof(kernel_param_size_stride)), + PushConstantDataInfo( + &kernel_param_pad_dial, sizeof(kernel_param_pad_dial)), + PushConstantDataInfo( + &extra_params, sizeof(extra_params), sizeof(utils::ivec4)), + PushConstantDataInfo(&out_params, sizeof(out_params)), + })); + } else { + graph.execute_nodes().emplace_back(new DispatchNode( + graph, + shader, + wg_size, + graph.create_local_wg_size(wg_size), + // Inputs and Outputs + {{out, vkapi::MemoryAccessType::WRITE}, + {{in, arg_weight, arg_bias}, vkapi::MemoryAccessType::READ}}, + // Shader params buffers + { + t_out->logical_limits_ubo(), + t_in->sizes_ubo(), + graph.create_params_buffer(kernel_params), + graph.create_params_buffer(extra_params), + graph.create_params_buffer(out_params), + }, + // Specialization Constants + {}, + // Resizing Logic + resize_conv2d_node, + {weight_data, stride, padding, dilation, transposed, output_padding})); + } } void add_conv1d_node( From 680916535673bd225035c6826a3cde26fa298f70 Mon Sep 17 00:00:00 2001 From: Vivek Trivedi <5340687+trivedivivek@users.noreply.github.com> Date: Mon, 27 Jan 2025 09:21:36 -0800 Subject: [PATCH 2/3] [ET-VK] Minor improvements to conv2d pw and dw bounds check. Pull Request resolved: https://github.com/pytorch/executorch/pull/7815 This diff contains minor improvements to the conv2d pw and dw bounds check in the Vulkan backend for Executorch. ghstack-source-id: 263238731 @exported-using-ghexport Differential Revision: [D68400689](https://our.internmc.facebook.com/intern/diff/D68400689/) --- backends/vulkan/runtime/graph/ops/glsl/conv2d_dw.glsl | 2 +- .../vulkan/runtime/graph/ops/glsl/conv2d_dw_output_tile.glsl | 2 +- .../runtime/graph/ops/glsl/conv2d_dw_sned_output_tile.glsl | 2 +- backends/vulkan/runtime/graph/ops/glsl/conv2d_pw.glsl | 2 +- 4 files changed, 4 insertions(+), 4 deletions(-) diff --git a/backends/vulkan/runtime/graph/ops/glsl/conv2d_dw.glsl b/backends/vulkan/runtime/graph/ops/glsl/conv2d_dw.glsl index 103f3cfdd76..8a845b6a8a6 100644 --- a/backends/vulkan/runtime/graph/ops/glsl/conv2d_dw.glsl +++ b/backends/vulkan/runtime/graph/ops/glsl/conv2d_dw.glsl @@ -41,7 +41,7 @@ void main() { div_by_x % out_limits.y, div_by_x / out_limits.y); - if (any(greaterThanEqual(pos, out_limits))) { + if (pos.z >= out_limits.z) { return; } diff --git a/backends/vulkan/runtime/graph/ops/glsl/conv2d_dw_output_tile.glsl b/backends/vulkan/runtime/graph/ops/glsl/conv2d_dw_output_tile.glsl index 9e69fdd1fe7..c05c7e4450d 100644 --- a/backends/vulkan/runtime/graph/ops/glsl/conv2d_dw_output_tile.glsl +++ b/backends/vulkan/runtime/graph/ops/glsl/conv2d_dw_output_tile.glsl @@ -59,7 +59,7 @@ void main() { pos.y *= BATCH_SIZE_Y; // do not process if top pixel does not fit within the output range - if (any(greaterThanEqual(pos, out_limits))) { + if (pos.z >= out_limits.z) { return; } diff --git a/backends/vulkan/runtime/graph/ops/glsl/conv2d_dw_sned_output_tile.glsl b/backends/vulkan/runtime/graph/ops/glsl/conv2d_dw_sned_output_tile.glsl index d0fc6707bff..bb70ee1aabb 100644 --- a/backends/vulkan/runtime/graph/ops/glsl/conv2d_dw_sned_output_tile.glsl +++ b/backends/vulkan/runtime/graph/ops/glsl/conv2d_dw_sned_output_tile.glsl @@ -44,7 +44,7 @@ void main() { div_by_x % out_limits.y, div_by_x / out_limits.y); - if (any(greaterThanEqual(pos, out_limits))) { + if (pos.z >= out_limits.z) { return; } diff --git a/backends/vulkan/runtime/graph/ops/glsl/conv2d_pw.glsl b/backends/vulkan/runtime/graph/ops/glsl/conv2d_pw.glsl index 0413eb7b7aa..051d9d69837 100644 --- a/backends/vulkan/runtime/graph/ops/glsl/conv2d_pw.glsl +++ b/backends/vulkan/runtime/graph/ops/glsl/conv2d_pw.glsl @@ -79,7 +79,7 @@ void main() { // If the top left position is out of bounds, then this invocation will have // no work to do. - if (any(greaterThanEqual(ivec3(pos[0], gpos.z), out_limits.xyz))) { + if (gpos.z >= out_limits.z) { return; } From fdb4f3f282e53c0546f5e85a7377e4da2f19ded7 Mon Sep 17 00:00:00 2001 From: Vivek Trivedi <5340687+trivedivivek@users.noreply.github.com> Date: Mon, 27 Jan 2025 09:21:38 -0800 Subject: [PATCH 3/3] [ET-VK] Splitting TILE_SIZE to TILE_SIZE_X and TILE_SIZE_Y in conv2d pw. Pull Request resolved: https://github.com/pytorch/executorch/pull/7816 This diff splits the `TILE_SIZE` variable in the `conv2d_pw` GLSL code into `TILE_SIZE_X` and `TILE_SIZE_Y`. This change is made so tile size in different dimensions can be tuned separately. ghstack-source-id: 263238734 @exported-using-ghexport Differential Revision: [D68400783](https://our.internmc.facebook.com/intern/diff/D68400783/) --- .../runtime/graph/ops/glsl/conv2d_pw.glsl | 30 +++++++++---------- .../runtime/graph/ops/glsl/conv2d_pw.yaml | 3 +- 2 files changed, 17 insertions(+), 16 deletions(-) diff --git a/backends/vulkan/runtime/graph/ops/glsl/conv2d_pw.glsl b/backends/vulkan/runtime/graph/ops/glsl/conv2d_pw.glsl index 051d9d69837..59f9f3880f8 100644 --- a/backends/vulkan/runtime/graph/ops/glsl/conv2d_pw.glsl +++ b/backends/vulkan/runtime/graph/ops/glsl/conv2d_pw.glsl @@ -12,7 +12,8 @@ #define VEC4_T ${texel_type(DTYPE)} -#define TILE_SIZE ${TILE_SIZE} +#define TILE_SIZE_X ${TILE_SIZE_X} +#define TILE_SIZE_Y ${TILE_SIZE_Y} #define op(X, A, B) ${OPERATOR} @@ -43,7 +44,7 @@ layout(local_size_x_id = 0, local_size_y_id = 1, local_size_z_id = 2) in; // shared memory to hold calculated positions, this would reduce register usage thus improving performance. // 64 is the number of threads in the local wg -$num_shared = 64 * TILE_SIZE * TILE_SIZE +$num_shared = 64 * TILE_SIZE_X * TILE_SIZE_Y shared ivec2 pos_shared[${num_shared}]; /* @@ -52,8 +53,8 @@ shared ivec2 pos_shared[${num_shared}]; * size is only 1x1, making it easier to re-use loaded texels from t_kernel. */ void main() { - const ivec2 out_limits_scaled = (out_limits.xy + TILE_SIZE - 1) / TILE_SIZE; - const uint shared_mem_stride = gl_WorkGroupSize.x * gl_WorkGroupSize.y * gl_WorkGroupSize.z; + const ivec2 out_limits_scaled = (out_limits.xy + ivec2(TILE_SIZE_X - 1, TILE_SIZE_Y - 1)) / ivec2(TILE_SIZE_X, TILE_SIZE_Y); + const uint shared_mem_stride = 64; const uint div_by_x = gl_GlobalInvocationID.x / out_limits_scaled.x; const ivec3 gpos = ivec3( @@ -67,11 +68,10 @@ void main() { // +--------+--------+ // | pos[2] | pos[3] | // +--------+--------+ - ivec2 pos[TILE_SIZE * TILE_SIZE]; - for (int y = 0, i = 0; y < TILE_SIZE; ++y) { - for (int x = 0; x < TILE_SIZE; ++x) { - pos[i] = ivec2( - gpos.x * TILE_SIZE + x, gpos.y * TILE_SIZE + y); + ivec2 pos[TILE_SIZE_X * TILE_SIZE_Y]; + for (int y = 0, i = 0; y < TILE_SIZE_Y; ++y) { + for (int x = 0; x < TILE_SIZE_X; ++x) { + pos[i] = ivec2(gpos.x * TILE_SIZE_X + x, gpos.y * TILE_SIZE_Y + y); pos_shared[(shared_mem_stride * i) + gl_LocalInvocationIndex] = pos[i]; i++; } @@ -86,14 +86,14 @@ void main() { // Compute the index of the input texture that needs to be loaded for each // output position. Note that negative indices can be produced indicating that // the top-left element is in a region added by padding. - ivec2 ipos[TILE_SIZE * TILE_SIZE]; - for (int i = 0; i < TILE_SIZE * TILE_SIZE; ++i) { + ivec2 ipos[TILE_SIZE_X * TILE_SIZE_Y]; + for (int i = 0; i < TILE_SIZE_X * TILE_SIZE_Y; ++i) { ipos[i] = pos[i] * stride - padding; } - vec4 sum[TILE_SIZE * TILE_SIZE]; + vec4 sum[TILE_SIZE_X * TILE_SIZE_Y]; sum[0] = texelFetch(t_bias, ivec2(gpos.z, 0), 0); - for (int i = 1; i < TILE_SIZE * TILE_SIZE; ++i) { + for (int i = 1; i < TILE_SIZE_X * TILE_SIZE_Y; ++i) { sum[i] = sum[0]; } @@ -109,7 +109,7 @@ void main() { const vec4 ktex_3 = texelFetchOffset(t_kernel, ivec2(z, gpos.z), 0, ivec2(3, 0)); #pragma unroll - for (int i = 0; i < TILE_SIZE * TILE_SIZE; ++i) { + for (int i = 0; i < TILE_SIZE_X * TILE_SIZE_Y; ++i) { const vec4 in_tex = texelFetch(t_in, ivec3(ipos[i], z4), 0); // For 2x2 tile size algorithm works as follows. // To explain the calculations below, the contents of one in_tex and the @@ -151,7 +151,7 @@ void main() { } } - for (int i = 0; i < TILE_SIZE * TILE_SIZE; ++i) { + for (int i = 0; i < TILE_SIZE_X * TILE_SIZE_Y; ++i) { const ivec2 pos = pos_shared[(shared_mem_stride * i) + gl_LocalInvocationIndex]; if (all(lessThan(ivec3(pos, gpos.z), out_limits.xyz))) { imageStore(t_out, ivec3(pos, gpos.z), op(sum[i], out_min, out_max)); diff --git a/backends/vulkan/runtime/graph/ops/glsl/conv2d_pw.yaml b/backends/vulkan/runtime/graph/ops/glsl/conv2d_pw.yaml index e77bc4f350a..1f0e8fb71be 100644 --- a/backends/vulkan/runtime/graph/ops/glsl/conv2d_pw.yaml +++ b/backends/vulkan/runtime/graph/ops/glsl/conv2d_pw.yaml @@ -9,7 +9,8 @@ conv2d_pw: OPERATOR: X NDIM: 3 DTYPE: float - TILE_SIZE: 2 + TILE_SIZE_X: 2 + TILE_SIZE_Y: 2 generate_variant_forall: DTYPE: - VALUE: half