Skip to content

Vulkan: classify conv2d method by weight shape in conv2d_local_wg_size #22051

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

Open
msluszniak wants to merge 1 commit into
base: main
Choose a base branch
from
Open

Vulkan: classify conv2d method by weight shape in conv2d_local_wg_size #22051

Changes from all commits
Commits
Show all changes
1 commit
Select commit
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
64 changes: 37 additions & 27 deletions backends/vulkan/runtime/graph/ops/impl/Convolution.cpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -349,6 +349,38 @@ utils::uvec3 create_conv2d_global_wg_size(
}
}

// Determines which convolution method a dispatch uses.
//
// Depthwise and transposed convolutions have shader names of their own, but
// the name alone cannot separate pointwise from sliding window: the sliding
// window shader is itself named "conv2d", and a pointwise convolution also
// takes that name when its weights are prepacked. Those two are therefore
// separated by the weight's spatial extent. Shared by the global and local
// workgroup size functions below so that the two cannot disagree about the
// same dispatch.
Conv2dMethod infer_conv2d_method_from_shader(
ComputeGraph* graph,
const vkapi::ShaderInfo& shader,
const ValueRef weight_data) {
const std::string& kernel_name = shader.kernel_name;
// Checked before the plain "conv2d" test below, which "conv2d_dw" and
// "conv2d_pw" would otherwise match too.
if (kernel_name.find("conv2d_dw") != std::string::npos) {
return Conv2dMethod::Depthwise;
}
if (kernel_name.find("conv2d_pw") != std::string::npos) {
return Conv2dMethod::Pointwise;
}
if (kernel_name.find("conv_transpose2d") != std::string::npos) {
return Conv2dMethod::Transposed;
}
const auto& weight_sizes = graph->get_tref(weight_data)->sizes;
if (weight_sizes.at(2) == 1 && weight_sizes.at(3) == 1) {
return Conv2dMethod::Pointwise;
}
return Conv2dMethod::SlidingWindow;

@SS-JIA SS-JIA Aug 24, 2026

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

also, just to make sure this function covers all possible Conv2dMethod entries, would recommend detecting Conv2dMethod::Depthwise as well.

@msluszniak msluszniak Aug 26, 2026

Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I guess that now every enum entry is reachable.

}

// Custom global workgroup size function for conv2d
utils::uvec3 conv2d_global_wg_size(
ComputeGraph* graph,
Expand All @@ -358,23 +390,8 @@ utils::uvec3 conv2d_global_wg_size(
const ValueRef out = args.at(0).refs.at(0);
const ValueRef weight_data = resize_args.at(0);

// Determine method from shader name
Conv2dMethod method;
if (shader.kernel_name.find("conv2d_pw") != std::string::npos ||
(shader.kernel_name.find("conv2d") != std::string::npos &&
shader.kernel_name.find("conv_transpose2d") == std::string::npos)) {
// Check if it's pointwise by examining weight sizes
const auto& weight_sizes = graph->get_tref(weight_data)->sizes;
if (weight_sizes.at(2) == 1 && weight_sizes.at(3) == 1) {
method = Conv2dMethod::Pointwise;
} else {
method = Conv2dMethod::SlidingWindow;
}
} else if (shader.kernel_name.find("conv_transpose2d") != std::string::npos) {
method = Conv2dMethod::Transposed;
} else {
method = Conv2dMethod::SlidingWindow;
}
const Conv2dMethod method =
infer_conv2d_method_from_shader(graph, shader, weight_data);

// Determine stride_equals_dilation from shader name
bool stride_equals_dilation =
Expand Down Expand Up @@ -402,17 +419,10 @@ utils::uvec3 conv2d_local_wg_size(
const std::vector<ArgGroup>& args,
const std::vector<ValueRef>& resize_args) {
(void)args;
(void)resize_args;

// Determine method from shader name
Conv2dMethod method;
if (shader.kernel_name.find("conv2d_pw") != std::string::npos ||
(shader.kernel_name.find("conv2d") != std::string::npos &&
shader.kernel_name.find("conv_transpose2d") == std::string::npos)) {
method = Conv2dMethod::Pointwise;
} else {
method = Conv2dMethod::SlidingWindow;
}
const ValueRef weight_data = resize_args.at(0);
const Conv2dMethod method =
infer_conv2d_method_from_shader(graph, shader, weight_data);

if (method == Conv2dMethod::Pointwise) {
uint32_t local_wg_size_y = 1;
Expand Down
Loading