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
+31 −27
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
58 changes: 31 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,32 +349,43 @@ utils::uvec3 create_conv2d_global_wg_size(
}
}

// Custom global workgroup size function for conv2d
utils::uvec3 conv2d_global_wg_size(
// Determines which convolution method a dispatch uses.
//
// The shader name alone is not enough: the sliding window and depthwise
// shaders are also named "conv2d", so a name test that accepts "conv2d"
// matches all of them. The pointwise case is therefore confirmed against 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 conv2d_method_from_dispatch(
ComputeGraph* graph,
const vkapi::ShaderInfo& shader,
const std::vector<ArgGroup>& args,
const std::vector<ValueRef>& resize_args) {
const ValueRef out = args.at(0).refs.at(0);
const ValueRef weight_data = resize_args.at(0);

// Determine method from shader name
Conv2dMethod method;
const ValueRef weight_data) {
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;
return Conv2dMethod::Pointwise;
}
} else if (shader.kernel_name.find("conv_transpose2d") != std::string::npos) {
method = Conv2dMethod::Transposed;
} else {
method = Conv2dMethod::SlidingWindow;
return Conv2dMethod::SlidingWindow;
}
if (shader.kernel_name.find("conv_transpose2d") != std::string::npos) {
return Conv2dMethod::Transposed;
}
return Conv2dMethod::SlidingWindow;
}

// Custom global workgroup size function for conv2d
utils::uvec3 conv2d_global_wg_size(
ComputeGraph* graph,
const vkapi::ShaderInfo& shader,
const std::vector<ArgGroup>& args,
const std::vector<ValueRef>& resize_args) {
const ValueRef out = args.at(0).refs.at(0);
const ValueRef weight_data = resize_args.at(0);

const Conv2dMethod method =
conv2d_method_from_dispatch(graph, shader, weight_data);

// Determine stride_equals_dilation from shader name
bool stride_equals_dilation =
Expand Down Expand Up @@ -402,17 +413,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 =
conv2d_method_from_dispatch(graph, shader, weight_data);

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