Skip to content

unroll 4x4 for gemm and sdpa vulkan, vectorize a and b loading, avoid bank conflict #6524

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
nihui merged 9 commits into from
Jan 28, 2026
Merged

unroll 4x4 for gemm and sdpa vulkan, vectorize a and b loading, avoid bank conflict #6524

Show file tree
Hide file tree
Changes from 2 commits
Commits
Show all changes
9 commits
Select commit Hold shift + click to select a range
b6a54cb
unroll 4x4 for vulkan gemm, vectorize loading a and b
nihui Jan 24, 2026
1c0ae53
unroll 4x4 for vulkan sdpa, vectorize loading a and b, avoid bank con…
nihui Jan 24, 2026
a90cd19
f
nihui Jan 24, 2026
670d121
Update glsl-extension.md
nihui Jan 26, 2026
d6e353a
Update glsl-extension.zh.md
nihui Jan 26, 2026
9635791
gemm vulkan subgroup optimization
nihui Jan 27, 2026
94dd1d8
Merge branch 'opt-vulkan-gemm-2' of github.com:nihui/ncnn into opt-vu…
nihui Jan 27, 2026
a7a9645
apply code-format changes
nihui Jan 27, 2026
c20bff1
enable it
nihui Jan 28, 2026
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
19 changes: 18 additions & 1 deletion src/gpu.cpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -4903,8 +4903,10 @@ int compile_spirv_module(const char* comp_data, int comp_data_size, const Option
custom_defines.append("buffer_sm1(buf,i)", "buf[i]");
custom_defines.append("buffer_sm4(buf,i)", "buf[i]");

custom_defines.append("lfp2afp(v)", "v");
custom_defines.append("lfp2afp(v)", "float(v)");
custom_defines.append("afp2lfp(v)", "bfloat16_t(v)");
custom_defines.append("lfp2afpvec4(v)", "vec4(v)");
custom_defines.append("afp2lfpvec4(v)", "bf16vec4(v)");
}
else if (opt.use_bf16_packed)
{
Expand All @@ -4925,60 +4927,75 @@ int compile_spirv_module(const char* comp_data, int comp_data_size, const Option
if (support_fp16_uniform)
{
custom_defines.append("lfp2afp(v)", "uintBitsToFloat(uint(v)<<16)");
custom_defines.append("afp2lfp(v)", "uint16_t(floatBitsToUint(v)>>16)");
}
else
{
custom_defines.append("lfp2afp(v)", "v");
custom_defines.append("afp2lfp(v)", "v");
}
custom_defines.append("lfp2afpvec4(v)", "vec4(unpackBFloat2x16(v.x),unpackBFloat2x16(v.y))");
custom_defines.append("afp2lfpvec4(v)", "uvec2(packBFloat2x16(v.rg),packBFloat2x16(v.ba))");
}
else if (opt.use_fp16_storage && opt.use_fp16_uniform && opt.use_fp16_arithmetic)
{
custom_defines.append("buffer_sm1(buf,i)", "buf[i]");
custom_defines.append("buffer_sm4(buf,i)", "buf[i]");

custom_defines.append("lfp2afp(v)", "v");
custom_defines.append("afp2lfp(v)", "v");
custom_defines.append("lfp2afpvec4(v)", "v");
custom_defines.append("afp2lfpvec4(v)", "v");
}
else if (opt.use_fp16_storage && opt.use_fp16_arithmetic)
{
custom_defines.append("buffer_sm1(buf,i)", "float(buf[i])");
custom_defines.append("buffer_sm4(buf,i)", "pack64(halfBitsToUint16(buf[i]))");

custom_defines.append("lfp2afp(v)", "float16_t(v)");
custom_defines.append("afp2lfp(v)", "float(v)");
custom_defines.append("lfp2afpvec4(v)", "uint16BitsToHalf(unpack16(v))");
custom_defines.append("afp2lfpvec4(v)", "pack64(halfBitsToUint16(v))");
}
else if (opt.use_fp16_packed && opt.use_fp16_arithmetic)
{
custom_defines.append("buffer_sm1(buf,i)", "unpackHalf2x16(buf[(i)/2])[(i)%2]");
custom_defines.append("buffer_sm4(buf,i)", "buf[i]");

custom_defines.append("lfp2afp(v)", "float16_t(v)");
custom_defines.append("afp2lfp(v)", "float(v)");
custom_defines.append("lfp2afpvec4(v)", "f16vec4(unpackFloat2x16(v.x),unpackFloat2x16(v.y))");
custom_defines.append("afp2lfpvec4(v)", "uvec2(packFloat2x16(v.rg),packFloat2x16(r.ba))");
}
else if (opt.use_fp16_storage)
{
custom_defines.append("buffer_sm1(buf,i)", "float(buf[i])");
custom_defines.append("buffer_sm4(buf,i)", "uvec2(packHalf2x16(vec4(buf[i]).rg),packHalf2x16(vec4(buf[i]).ba))");

custom_defines.append("lfp2afp(v)", "v");
custom_defines.append("afp2lfp(v)", "float(v)");
custom_defines.append("lfp2afpvec4(v)", "vec4(unpackHalf2x16(v.x),unpackHalf2x16(v.y))");
custom_defines.append("afp2lfpvec4(v)", "uvec2(packHalf2x16(v.rg),packHalf2x16(r.ba))");
}
else if (opt.use_fp16_packed)
{
custom_defines.append("buffer_sm1(buf,i)", "unpackHalf2x16(buf[(i)/2])[(i)%2]");
custom_defines.append("buffer_sm4(buf,i)", "buf[i]");

custom_defines.append("lfp2afp(v)", "v");
custom_defines.append("afp2lfp(v)", "v");
custom_defines.append("lfp2afpvec4(v)", "vec4(unpackHalf2x16(v.x),unpackHalf2x16(v.y))");
custom_defines.append("afp2lfpvec4(v)", "uvec2(packHalf2x16(v.x),packHalf2x16(v.y))");
}
else
{
custom_defines.append("buffer_sm1(buf,i)", "buf[i]");
custom_defines.append("buffer_sm4(buf,i)", "buf[i]");

custom_defines.append("lfp2afp(v)", "v");
custom_defines.append("afp2lfp(v)", "v");
custom_defines.append("lfp2afpvec4(v)", "v");
custom_defines.append("afp2lfpvec4(v)", "v");
}

if (opt.use_bf16_storage)
Expand Down
4 changes: 2 additions & 2 deletions src/layer/vulkan/gemm_vulkan.cpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -367,8 +367,8 @@ int Gemm_vulkan::forward(const std::vector<VkMat>& bottom_blobs, std::vector<VkM
const Pipeline* pipeline = pipeline_gemm;

VkMat dispatcher;
dispatcher.w = (N + 1) / 2;
dispatcher.h = (M + 1) / 2;
dispatcher.w = (N + 3) / 4;
dispatcher.h = (M + 3) / 4;
dispatcher.c = 1;
cmd.record_pipeline(pipeline, bindings, constants, dispatcher);
}
Expand Down
8 changes: 4 additions & 4 deletions src/layer/vulkan/sdpa_vulkan.cpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -340,8 +340,8 @@ int SDPA_vulkan::forward(const std::vector<VkMat>& bottom_blobs, std::vector<VkM
else
{
VkMat dispatcher;
dispatcher.w = (N + 1) / 2;
dispatcher.h = (M + 1) / 2;
dispatcher.w = (N + 3) / 4;
dispatcher.h = (M + 3) / 4;
dispatcher.c = B;

cmd.record_pipeline(pipeline_sdpa_qk_cross, bindings, constants, dispatcher);
Expand Down Expand Up @@ -415,8 +415,8 @@ int SDPA_vulkan::forward(const std::vector<VkMat>& bottom_blobs, std::vector<VkM
else
{
VkMat dispatcher;
dispatcher.w = (N + 1) / 2;
dispatcher.h = (M + 1) / 2;
dispatcher.w = (N + 3) / 4;
dispatcher.h = (M + 3) / 4;
dispatcher.c = B;

cmd.record_pipeline(pipeline_sdpa_qkv_cross, bindings, constants, dispatcher);
Expand Down
Loading
Loading