Skip to content

[ET-VK] Modify quantized linear tiling shader to linearly dispatch work to improve thread occupancy and performance. #10508

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
merged 5 commits into from
Apr 28, 2025
Merged

[ET-VK] Modify quantized linear tiling shader to linearly dispatch work to improve thread occupancy and performance. #10508

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
5 commits
Select commit Hold shift + click to select a range
7e168c5
[ET-VK] Modify quantized linear tiling shader to linearly dispatch wo…
trivedivivek Apr 28, 2025
16732cc
Update on "[ET-VK] Modify quantized linear tiling shader to linearly …
trivedivivek Apr 28, 2025
6bac017
Update on "[ET-VK] Modify quantized linear tiling shader to linearly …
trivedivivek Apr 28, 2025
ef1ea85
Update on "[ET-VK] Modify quantized linear tiling shader to linearly …
trivedivivek Apr 28, 2025
ab1de82
Update on "[ET-VK] Modify quantized linear tiling shader to linearly …
trivedivivek Apr 28, 2025
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
9 changes: 6 additions & 3 deletions backends/vulkan/runtime/graph/ops/glsl/q_8w_linear_coop.glsl
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -38,18 +38,21 @@ layout(push_constant) uniform restrict Block {
ivec4 weight_sizes;
};

#include "indexing_utils.h"

layout(local_size_x_id = 0, local_size_y_id = 1, local_size_z_id = 2) in;

shared VEC4_T partial_c[NGROUPS][NWORKERS][TILE_ROWS];

void main() {
const uint out_row = gl_GlobalInvocationID.y * TILE_ROWS;
const uint out_col = gl_GlobalInvocationID.x << 2;
const uint out_width_ntexels = divup4(out_sizes.x);
const uint out_col = (gl_GlobalInvocationID.x % out_width_ntexels) << 2;
const uint out_row = (gl_GlobalInvocationID.x / out_width_ntexels) * TILE_ROWS;

const int gid = int(gl_LocalInvocationID.x); // group id
const int wid = int(gl_LocalInvocationID.z); // worker id

if (out_col >= out_sizes.x || out_row >= out_sizes.y) {
if (out_row >= out_sizes.y) {
return;
}

Expand Down
9 changes: 6 additions & 3 deletions backends/vulkan/runtime/graph/ops/glsl/q_8w_linear_tiled.glsl
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -36,13 +36,16 @@ layout(push_constant) uniform restrict Block {
ivec4 weight_sizes;
};

#include "indexing_utils.h"

layout(local_size_x_id = 0, local_size_y_id = 1, local_size_z_id = 2) in;

void main() {
const uint out_row = gl_GlobalInvocationID.y * TILE_ROWS;
const uint out_col = gl_GlobalInvocationID.x << 2;
const uint out_width_ntexels = divup4(out_sizes.x);
const uint out_col = (gl_GlobalInvocationID.x % out_width_ntexels) << 2;
const uint out_row = (gl_GlobalInvocationID.x / out_width_ntexels) * TILE_ROWS;

if (out_col >= out_sizes.x || out_row >= out_sizes.y) {
if (out_row >= out_sizes.y) {
return;
}

Expand Down
4 changes: 2 additions & 2 deletions backends/vulkan/runtime/graph/ops/glsl/q_8w_linear_tiled.yaml
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -16,10 +16,10 @@ q_8w_linear_tiled:
TILE_ROWS:
- VALUE: 1
SUFFIX: o4x1
- VALUE: 2
SUFFIX: o4x2
- VALUE: 4
SUFFIX: o4x4
- VALUE: 6
SUFFIX: o4x6
shader_variants:
- NAME: q_8w_linear_tiled_texture3d_texture3d_texture2d_texture2d_float
- NAME: q_8w_linear_tiled_buffer_buffer_texture2d_texture2d_float
Expand Down
13 changes: 8 additions & 5 deletions backends/vulkan/runtime/graph/ops/impl/QuantizedLinearInt8.cpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -180,10 +180,10 @@ void add_q_8w_linear_tiled_node(

std::vector<int64_t> mat1_sizes = graph.sizes_of(mat1);
const int64_t M = utils::val_at(-2, mat1_sizes);
int out_tile_nrows = 4;
uint32_t out_tile_nrows = 4;
if (M % 6 == 0) {
kernel_name += "_o4x6";
out_tile_nrows = 6;
kernel_name += "_o4x2";
out_tile_nrows = 2;
} else if (M % 4 == 0) {
kernel_name += "_o4x4";
out_tile_nrows = 4;
Expand All @@ -195,8 +195,11 @@ void add_q_8w_linear_tiled_node(
out_tile_nrows = 4;
}

utils::uvec3 global_wg_size = graph.logical_limits_of(out);
global_wg_size[1] = global_wg_size[1] / out_tile_nrows;
utils::uvec3 out_limits = graph.logical_limits_of(out);
utils::uvec3 global_wg_size = {
out_limits[0] * (utils::div_up(out_limits[1], out_tile_nrows)),
1,
out_limits[2]};

utils::uvec3 local_wg_size{64, 1, 1};
if (use_coop_algorithm) {
Expand Down
Loading