Skip to content

[ET-VK] Efficient tiled int8 matmul #9804

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 2 commits into from
Apr 1, 2025
Merged

[ET-VK] Efficient tiled int8 matmul #9804

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
2 commits
Select commit Hold shift + click to select a range
de2aab6
[ET-VK] Store weights transposed for int8 linear
SS-JIA Mar 31, 2025
1907ae2
[ET-VK] Efficient tiled int8 matmul
SS-JIA Mar 31, 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
21 changes: 19 additions & 2 deletions backends/vulkan/runtime/graph/ops/glsl/nchw_to_bitw8_image_nobitw8buffer.glsl
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -27,6 +27,8 @@ ${layout_declare_ubo(B, "ivec4", "sizes")}
layout(local_size_x_id = 0, local_size_y_id = 1, local_size_z_id = 2) in;

${layout_declare_spec_const(C, "int", "t_layout", "DEFAULT_LAYOUT")}
${layout_declare_spec_const(C, "int", "transpose_hw", "0")}

const lowp ivec4 axis_map = unhash_axis_map(t_layout);
const lowp int packed_dim = unhash_packed_dim(t_layout);

Expand All @@ -41,8 +43,23 @@ int extend_sign(int x) {
}

ivec4 read_texel(ivec4 tidx) {
ivec4 tidx_to_use = tidx;
ivec4 sizes_to_use = sizes;
int packed_dim_to_use = packed_dim;
if (transpose_hw == 1) {
sizes_to_use.xy = sizes_to_use.yx;
tidx_to_use.xy = tidx.yx;

if (packed_dim == 1) {
packed_dim_to_use = 0;
}
if (packed_dim == 0) {
packed_dim_to_use = 1;
}
}

const ivec4 buf_indices = tidx_to_nchwi(
tidx, sizes, packed_dim);
tidx_to_use, sizes_to_use, packed_dim_to_use);

int shift = (1 << 8) - 1;
ivec4 masks;
Expand Down Expand Up @@ -70,7 +87,7 @@ ivec4 read_texel(ivec4 tidx) {

void main() {
const ivec3 lpos = ivec3(gl_GlobalInvocationID);
const ivec4 tidx = lpos_to_tidx(lpos, sizes, axis_map.w, packed_dim);
ivec4 tidx = lpos_to_tidx(lpos, sizes, axis_map.w, packed_dim);

if (any(greaterThanEqual(tidx, sizes))) {
return;
Expand Down
9 changes: 8 additions & 1 deletion backends/vulkan/runtime/graph/ops/glsl/nchw_to_buffer.glsl
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -21,6 +21,7 @@ layout(local_size_x_id = 0, local_size_y_id = 1, local_size_z_id = 2) in;
// This constant is unused in this shader but is kept so that the signature is
// consistent with nchw_to_image.
${layout_declare_spec_const(C, "int", "UNUSED_layout", "0")}
${layout_declare_spec_const(C, "int", "transpose_hw", "0")}

void main() {
int out_bufi = int(gl_GlobalInvocationID.x);
Expand All @@ -29,7 +30,13 @@ void main() {
}

ivec4 out_tidx = bufi_to_tidx(out_bufi, out_strides);
const int in_nchwi = tidx_to_nchwi(out_tidx, out_sizes);

ivec4 sizes = out_sizes;
if (transpose_hw == 1) {
sizes.xy = sizes.yx;
out_tidx.xy = out_tidx.yx;
}
const int in_nchwi = tidx_to_nchwi(out_tidx, sizes);

t_out[out_bufi] = nchw_in[in_nchwi];
}
21 changes: 19 additions & 2 deletions backends/vulkan/runtime/graph/ops/glsl/nchw_to_image.glsl
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -30,14 +30,31 @@ $if not FROM_STAGING:
layout(local_size_x_id = 0, local_size_y_id = 1, local_size_z_id = 2) in;

${layout_declare_spec_const(C, "int", "t_layout", "DEFAULT_LAYOUT")}
${layout_declare_spec_const(C, "int", "transpose_hw", "0")}

const lowp ivec4 axis_map = unhash_axis_map(t_layout);
const lowp int packed_dim = unhash_packed_dim(t_layout);

VEC4_T read_texel(ivec4 tidx) {
ivec4 tidx_to_use = tidx;
ivec4 sizes_to_use = sizes;
int packed_dim_to_use = packed_dim;
if (transpose_hw == 1) {
sizes_to_use.xy = sizes_to_use.yx;
tidx_to_use.xy = tidx.yx;

if (packed_dim == 1) {
packed_dim_to_use = 0;
}
if (packed_dim == 0) {
packed_dim_to_use = 1;
}
}

$if FROM_STAGING:
const ivec4 buf_indices = tidx_to_nchwi(tidx, sizes, packed_dim);
const ivec4 buf_indices = tidx_to_nchwi(tidx_to_use, sizes_to_use, packed_dim_to_use);
$else:
const ivec4 buf_indices = tidx_to_4bufi(tidx, buf_strides, packed_dim);
const ivec4 buf_indices = tidx_to_4bufi(tidx_to_use, buf_strides, packed_dim_to_use);

VEC4_T texel = VEC4_T(0);
if (tidx[packed_dim] < sizes[packed_dim]) {
Expand Down
31 changes: 15 additions & 16 deletions backends/vulkan/runtime/graph/ops/glsl/q_8w_linear.glsl
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -64,24 +64,21 @@ void main() {

FLOAT_T outval = FLOAT_T(0.0);

// Initial mat1 tensor idx will be (0, out_tidx.y, out_tidx.z, 0)
int mat1_offset = out_tidx.y * mat1_strides.y + out_tidx.z * qmat2_strides.z;
// Initial qmat2 tensor idx wil be (0, out_tidx.x, 0, 0); note that the qmat2
// tensor is transposed
int qmat2_offset = out_tidx.x * qmat2_strides.y;
int qmat2_offset = out_tidx.x;

// TODO(ssjia): optimize memory access pattern by traversing mat1 x in inner loop
for (int i = 0; i < mat1_sizes.x; i++) {
const FLOAT_T mat1_val = t_mat1[mat1_offset];
const FLOAT_T mat2_val = t_qmat2[qmat2_offset] * scale;
const FLOAT_T mat2_val = FLOAT_T(t_qmat2[qmat2_offset]);

outval += mat1_val * mat2_val;

mat1_offset++;
qmat2_offset++;
qmat2_offset += qmat2_strides.y;
}

t_out[out_bufi] = outval;
t_out[out_bufi] = outval * scale;
}

#else // USING_TEXTURE
Expand All @@ -97,25 +94,27 @@ void main() {
return;
}

const uint16_t qmat2_pos_y = out_pos.x * uint16_t(4);
const uint16_t qmat2_pos_x = out_pos.x;

VEC4_T outtex = VEC4_T(0);

const VEC4_T scales = load_texel(t_scales, u16vec3(out_pos.x, 0, 0));

VEC4_T mat1_tex;
VEC4_T mat2_tex[4];
for (
uint16_t i = uint16_t(0), x = uint16_t(0);
i < uint16_t(mat1_sizes.x);
i += uint16_t(4), x++)
{
const VEC4_T mat1_tex = load_texel(t_mat1, u16vec3(x, out_pos.y, 0));
const VEC4_T sums = VEC4_T(
dot(mat1_tex, load_texel(t_qmat2, u16vec3(x, qmat2_pos_y, 0))),
dot(mat1_tex, load_texel(t_qmat2, u16vec3(x, qmat2_pos_y + uint16_t(1), 0))),
dot(mat1_tex, load_texel(t_qmat2, u16vec3(x, qmat2_pos_y + uint16_t(2), 0))),
dot(mat1_tex, load_texel(t_qmat2, u16vec3(x, qmat2_pos_y + uint16_t(3), 0))));

outtex += sums;
mat1_tex = load_texel(t_mat1, u16vec3(x, out_pos.y, 0));

mat2_tex[0] = load_texel(t_qmat2, u16vec3(out_pos.x, i, 0));
mat2_tex[1] = load_texel(t_qmat2, u16vec3(out_pos.x, i + uint16_t(1), 0));
mat2_tex[2] = load_texel(t_qmat2, u16vec3(out_pos.x, i + uint16_t(2), 0));
mat2_tex[3] = load_texel(t_qmat2, u16vec3(out_pos.x, i + uint16_t(3), 0));

outtex += mat1_tex.x * mat2_tex[0] + mat1_tex.y * mat2_tex[1] + mat1_tex.z * mat2_tex[2] + mat1_tex.w * mat2_tex[3];
}

outtex *= scales;
Expand Down
212 changes: 0 additions & 212 deletions backends/vulkan/runtime/graph/ops/glsl/q_8w_linear_optimized.glsl
View file
Open in desktop

This file was deleted.

Loading
Loading