From f805e9494a0d5af6d34b34595b301c06a1822867 Mon Sep 17 00:00:00 2001 From: Vivek Trivedi <5340687+trivedivivek@users.noreply.github.com> Date: Wed, 8 Jan 2025 21:00:59 -0800 Subject: [PATCH 1/2] [ET-VK] Adding a common utility function to calculate 3d output position based on unique index. Pull Request resolved: https://github.com/pytorch/executorch/pull/7522 This diff adds an indexing utils header file used in Vulkan backend of Executorch. The header file includes functions for converting a global index to u16 indices based on input sizes. ghstack-source-id: 260707858 @exported-using-ghexport Differential Revision: [D67821941](https://our.internmc.facebook.com/intern/diff/D67821941/) --- .../runtime/graph/ops/glsl/conv2d_dw.glsl | 7 ++----- .../graph/ops/glsl/conv2d_dw_output_tile.glsl | 9 +++------ .../runtime/graph/ops/glsl/conv2d_pw.glsl | 10 +++------- .../graph/ops/glsl/indexing_utils_u16.h | 19 +++++++++++++++++++ 4 files changed, 27 insertions(+), 18 deletions(-) create mode 100644 backends/vulkan/runtime/graph/ops/glsl/indexing_utils_u16.h diff --git a/backends/vulkan/runtime/graph/ops/glsl/conv2d_dw.glsl b/backends/vulkan/runtime/graph/ops/glsl/conv2d_dw.glsl index 43a4f7c8dc7..5d7c69ab654 100644 --- a/backends/vulkan/runtime/graph/ops/glsl/conv2d_dw.glsl +++ b/backends/vulkan/runtime/graph/ops/glsl/conv2d_dw.glsl @@ -14,7 +14,7 @@ #define op(X, A, B) ${OPERATOR} -#include "indexing_utils.h" +#include "indexing_utils_u16.h" layout(std430) buffer; @@ -35,10 +35,7 @@ layout(local_size_x_id = 0, local_size_y_id = 1, local_size_z_id = 2) in; * output at a single output location. */ void main() { - const ivec3 pos = ivec3( - gl_GlobalInvocationID.x % out_limits.x, - (gl_GlobalInvocationID.x / out_limits.x) % out_limits.y, - gl_GlobalInvocationID.x / (out_limits.x * out_limits.y)); + const ivec3 pos = idx_to_u16pos_x_wise(gl_GlobalInvocationID.x, out_limits.x, out_limits.y); if (any(greaterThanEqual(pos, out_limits))) { 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 b2ae4953a78..20fb9374bec 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 @@ -18,7 +18,7 @@ #define op(X, A, B) ${OPERATOR} -#include "indexing_utils.h" +#include "indexing_utils_u16.h" layout(std430) buffer; @@ -43,12 +43,9 @@ layout(local_size_x_id = 0, local_size_y_id = 1, local_size_z_id = 2) in; void main() { // y divided up by batch size is used to determine 3d position // since work size is calculated by x * ((y + B_Y - 1) / B_Y) * z - const uint out_limits_y_scaled = (out_limits.y + BATCH_SIZE_Y - 1) / BATCH_SIZE_Y; + const int out_limits_y_scaled = (out_limits.y + BATCH_SIZE_Y - 1) / BATCH_SIZE_Y; - u16vec3 pos = u16vec3( - gl_GlobalInvocationID.x % out_limits.x, - ((gl_GlobalInvocationID.x / out_limits.x) % out_limits_y_scaled), - gl_GlobalInvocationID.x / (out_limits.x * out_limits_y_scaled)); + u16vec3 pos = idx_to_u16pos_x_wise(gl_GlobalInvocationID.x, out_limits.x, out_limits_y_scaled); // scale pos.y by batch size, because that's the top pixel to be processed pos.y *= uint16_t(BATCH_SIZE_Y); diff --git a/backends/vulkan/runtime/graph/ops/glsl/conv2d_pw.glsl b/backends/vulkan/runtime/graph/ops/glsl/conv2d_pw.glsl index 23ad912c11a..ad5d4adb134 100644 --- a/backends/vulkan/runtime/graph/ops/glsl/conv2d_pw.glsl +++ b/backends/vulkan/runtime/graph/ops/glsl/conv2d_pw.glsl @@ -16,7 +16,7 @@ #define op(X, A, B) ${OPERATOR} -#include "indexing_utils.h" +#include "indexing_utils_u16.h" layout(std430) buffer; @@ -43,13 +43,10 @@ shared u16vec2 pos_shared[gl_WorkGroupSize.x * gl_WorkGroupSize.y * gl_WorkGroup * size is only 1x1, making it easier to re-use loaded texels from t_kernel. */ void main() { - const uvec2 out_limits_scaled = (out_limits.xy + TILE_SIZE - 1) / TILE_SIZE; + 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 u16vec3 gpos = u16vec3( - gl_GlobalInvocationID.x % out_limits_scaled.x, - (gl_GlobalInvocationID.x / out_limits_scaled.x) % out_limits_scaled.y, - gl_GlobalInvocationID.x / (out_limits_scaled.x * out_limits_scaled.y)); + const u16vec3 gpos = idx_to_u16pos_x_wise(gl_GlobalInvocationID.x, out_limits_scaled.x, out_limits_scaled.y); // Output position for TILE_SIZE = 2 // +--------+--------+ @@ -98,7 +95,6 @@ void main() { const vec4 ktex_2 = texelFetchOffset(t_kernel, u16vec2(z, gpos.z), 0, u16vec2(2, 0)); const vec4 ktex_3 = texelFetchOffset(t_kernel, u16vec2(z, gpos.z), 0, u16vec2(3, 0)); - #pragma unroll for (int i = 0; i < TILE_SIZE * TILE_SIZE; ++i) { const vec4 in_tex = texelFetch(t_in, u16vec3(ipos[i], z4), 0); diff --git a/backends/vulkan/runtime/graph/ops/glsl/indexing_utils_u16.h b/backends/vulkan/runtime/graph/ops/glsl/indexing_utils_u16.h new file mode 100644 index 00000000000..6dc59b63039 --- /dev/null +++ b/backends/vulkan/runtime/graph/ops/glsl/indexing_utils_u16.h @@ -0,0 +1,19 @@ +/* + * Copyright (c) Meta Platforms, Inc. and affiliates. + * All rights reserved. + * + * This source code is licensed under the BSD-style license found in the + * LICENSE file in the root directory of this source tree. + */ + +#ifndef INDEXING_UTILS_U16_H +#define INDEXING_UTILS_U16_H + +#extension GL_EXT_shader_explicit_arithmetic_types_int16 : require + +u16vec3 idx_to_u16pos_x_wise(uint idx, int size_x, int size_y) { + const uint div_by_x = idx / size_x; + return u16vec3(idx % size_x, div_by_x % size_y, div_by_x / size_y); +} + +#endif // INDEXING_UTILS_U16_H From 4a4dab5c3f7b4f14c00fe751962ed264b2bf97a5 Mon Sep 17 00:00:00 2001 From: Vivek Trivedi <5340687+trivedivivek@users.noreply.github.com> Date: Thu, 9 Jan 2025 08:26:54 -0800 Subject: [PATCH 2/2] [ET-VK] Adding batch processing in x axis to conv2d dw shader by caching input texel for reuse. Pull Request resolved: https://github.com/pytorch/executorch/pull/7526 This diff adds batch processing in the x axis to the conv2d dw shader by reusing input texel overlapping between consecutive tiles. The changes include modifying the glsl code for the conv2d dw output tile, adding a new parameter to the yaml file, and modifying the Convolution.cpp file to use the new parameter. ghstack-source-id: 260707856 Differential Revision: [D67868671](https://our.internmc.facebook.com/intern/diff/D67868671/) --- .../graph/ops/glsl/conv2d_dw_output_tile.glsl | 71 +++++++++++-------- .../graph/ops/glsl/conv2d_dw_output_tile.yaml | 1 + .../runtime/graph/ops/glsl/indexing_utils.h | 5 ++ .../runtime/graph/ops/impl/Convolution.cpp | 2 +- 4 files changed, 48 insertions(+), 31 deletions(-) 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 20fb9374bec..32d0229d96d 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 @@ -14,11 +14,13 @@ #define TILE_SIZE ${TILE_SIZE} +#define BATCH_SIZE_X ${BATCH_SIZE_X} + #define BATCH_SIZE_Y ${BATCH_SIZE_Y} #define op(X, A, B) ${OPERATOR} -#include "indexing_utils_u16.h" +#include "indexing_utils.h" layout(std430) buffer; @@ -41,70 +43,79 @@ layout(local_size_x_id = 0, local_size_y_id = 1, local_size_z_id = 2) in; * output at a single output location. */ void main() { - // y divided up by batch size is used to determine 3d position + // x and y are divided by batch size to determine 3d position // since work size is calculated by x * ((y + B_Y - 1) / B_Y) * z - const int out_limits_y_scaled = (out_limits.y + BATCH_SIZE_Y - 1) / BATCH_SIZE_Y; + const ivec2 out_limits_xy_scaled = (out_limits.xy + ivec2(BATCH_SIZE_X, BATCH_SIZE_Y) - 1) / ivec2(BATCH_SIZE_X, BATCH_SIZE_Y); - u16vec3 pos = idx_to_u16pos_x_wise(gl_GlobalInvocationID.x, out_limits.x, out_limits_y_scaled); + ivec3 pos = idx_to_ipos_x_wise(gl_GlobalInvocationID.x, out_limits_xy_scaled.x, out_limits_xy_scaled.y); - // scale pos.y by batch size, because that's the top pixel to be processed - pos.y *= uint16_t(BATCH_SIZE_Y); + // scale pos.xy by batch sizes, because that's the top pixel to be processed + pos.x *= BATCH_SIZE_X; + pos.y *= BATCH_SIZE_Y; // do not process if top pixel does not fit within the output range - if (any(greaterThanEqual(u16vec3(pos.x, pos.y, pos.z), out_limits))) { + if (any(greaterThanEqual(pos, out_limits))) { return; } // Compute the index of the top-left element of the overlay region. Negative // indices indicate that the top-left element is in a region added by padding. - const u16vec2 ipos = pos.xy * u16vec2(stride) - u16vec2(padding); + const ivec2 ipos = pos.xy * stride - padding; // Compute the start and end of the input indices to load. Padding is assumed // to be constant 0 padding, so any reads from the padding region is skipped. - const u16vec2 start = ipos; - const u16vec2 end = ipos + u16vec2(overlay_region.xy); + const ivec2 start = ipos; + const ivec2 end = ipos + overlay_region.xy; // sum outputs - VEC4_T sum[BATCH_SIZE_Y]; + VEC4_T sum[BATCH_SIZE_Y][BATCH_SIZE_X]; - sum[0] = texelFetch(t_bias, u16vec2(pos.z, 0), 0); - for (int i = 1; i < BATCH_SIZE_Y; i++) { - sum[i] = sum[0]; + sum[0][0] = texelFetch(t_bias, ivec2(pos.z, 0), 0); + for (int y = 0; y < BATCH_SIZE_Y; y++) { + for (int x = 0; x < BATCH_SIZE_X; x++) { + sum[y][x] = sum[0][0]; + } } // array to store input texels - VEC4_T in_texels[TILE_SIZE]; + VEC4_T in_texels[TILE_SIZE + BATCH_SIZE_X - 1]; // array to store kernel data of previous y VEC4_T prev_kernel_line[TILE_SIZE]; - uint16_t kx = uint16_t(0); - for (uint16_t y = start.y, i = uint16_t(0); i < uint16_t(TILE_SIZE + BATCH_SIZE_Y - 1); y += uint16_t(dilation.y), i++) { - for (uint16_t x = start.x, j = uint16_t(0); j < uint16_t(TILE_SIZE); x += uint16_t(dilation.x), j++) { - in_texels[int(j)] = texelFetch(t_in, u16vec3(x, y, pos.z), 0); + int kx = 0; + for (int y = start.y, i = 0; i < TILE_SIZE + BATCH_SIZE_Y - 1; y += dilation.y, i++) { + for (int x = start.x, j = 0; j < TILE_SIZE + BATCH_SIZE_X - 1; x += dilation.x, j++) { + in_texels[j] = texelFetch(t_in, ivec3(x, y, pos.z), 0); } // from 2nd iteration onwards accumulate dot product in 2nd sum // based on kernel line data fetched in previous iteration and input texel from this iteration - if (i > uint16_t(0)) { - for (uint16_t j = uint16_t(0); j < uint16_t(TILE_SIZE); j++) { - sum[1] = fma(in_texels[int(j)], prev_kernel_line[int(j)], sum[1]); + if (i > 0) { + for (int j = 0; j < TILE_SIZE; j++) { + for (int s = 0; s < BATCH_SIZE_X; s++) { + sum[1][s] = fma(in_texels[j + s], prev_kernel_line[j], sum[1][s]); + } } } // accumulate dot product in 1st sum only until tile size - if (i < uint16_t(TILE_SIZE)) { - for (uint16_t j = uint16_t(0); j < uint16_t(TILE_SIZE); j++, kx++) { - prev_kernel_line[int(j)] = texelFetch(t_kernel, u16vec2(kx, pos.z), 0); - sum[0] = fma(in_texels[int(j)], prev_kernel_line[int(j)], sum[0]); + if (i < TILE_SIZE) { + for (int j = 0; j < TILE_SIZE; j++, kx++) { + prev_kernel_line[j] = texelFetch(t_kernel, ivec2(kx, pos.z), 0); + for (int s = 0; s < BATCH_SIZE_X; s++) { + sum[0][s] = fma(in_texels[j + s], prev_kernel_line[j], sum[0][s]); + } } } } - for (int i = 0; i < BATCH_SIZE_Y; i++) { - if (any(greaterThanEqual(u16vec3(pos.x, pos.y + i, pos.z), out_limits))) { - continue; + for (int y = 0; y < BATCH_SIZE_Y; y++) { + for (int x = 0; x < BATCH_SIZE_X; x++) { + if (any(greaterThanEqual(ivec3(pos.x + x, pos.y + y, pos.z), out_limits))) { + continue; + } + imageStore(t_out, ivec3(pos.x + x, pos.y + y, pos.z), op(sum[y][x], out_min, out_max)); } - imageStore(t_out, u16vec3(pos.x, pos.y + i, pos.z), op(sum[i], out_min, out_max)); } } diff --git a/backends/vulkan/runtime/graph/ops/glsl/conv2d_dw_output_tile.yaml b/backends/vulkan/runtime/graph/ops/glsl/conv2d_dw_output_tile.yaml index bb197c2c187..9cf6c22c6ca 100644 --- a/backends/vulkan/runtime/graph/ops/glsl/conv2d_dw_output_tile.yaml +++ b/backends/vulkan/runtime/graph/ops/glsl/conv2d_dw_output_tile.yaml @@ -10,6 +10,7 @@ conv2d_dw_output_tile: NDIM: 3 DTYPE: float TILE_SIZE: 3 + BATCH_SIZE_X: 4 BATCH_SIZE_Y: 2 generate_variant_forall: DTYPE: diff --git a/backends/vulkan/runtime/graph/ops/glsl/indexing_utils.h b/backends/vulkan/runtime/graph/ops/glsl/indexing_utils.h index 0b372ab70a4..1d3a60cb293 100644 --- a/backends/vulkan/runtime/graph/ops/glsl/indexing_utils.h +++ b/backends/vulkan/runtime/graph/ops/glsl/indexing_utils.h @@ -223,6 +223,11 @@ ivec3 lpos_to_pos(const ivec3 lpos, const ivec4 axis_map) { return pos; } +ivec3 idx_to_ipos_x_wise(uint idx, int size_x, int size_y) { + const uint div_by_x = idx / size_x; + return ivec3(idx % size_x, div_by_x % size_y, div_by_x / size_y); +} + #ifdef USING_BUFFER #define load_texel(buf, idx) buf[idx] #elif defined(USING_TEXTURE2D) diff --git a/backends/vulkan/runtime/graph/ops/impl/Convolution.cpp b/backends/vulkan/runtime/graph/ops/impl/Convolution.cpp index 3519635ac7e..64c145fb7e7 100644 --- a/backends/vulkan/runtime/graph/ops/impl/Convolution.cpp +++ b/backends/vulkan/runtime/graph/ops/impl/Convolution.cpp @@ -299,7 +299,7 @@ utils::uvec3 create_conv2d_global_wg_size( } else if (method == Conv2dMethod::Depthwise) { const utils::uvec3 image_extents = graph.logical_limits_of(out); return { - utils::div_up(image_extents[0u], 1u), + utils::div_up(image_extents[0u], 4u), utils::div_up(image_extents[1u], 2u), image_extents[2u]}; } else {