Skip to content

[ET-VK][ez][Refactor] Re-order DispatchNode arguments to match shader layout spec #10700

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 1 commit into from
May 5, 2025
Merged

[ET-VK][ez][Refactor] Re-order DispatchNode arguments to match shader layout spec #10700

Show file tree
Hide file tree
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
4 changes: 2 additions & 2 deletions backends/vulkan/runtime/graph/ops/DispatchNode.cpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -21,10 +21,10 @@ DispatchNode::DispatchNode(
const utils::uvec3& local_workgroup_size,
const std::vector<ArgGroup>& args,
const vkapi::ParamsBindList& params,
const std::vector<PushConstantDataInfo>& push_constants,
const vkapi::SpecVarList& spec_vars,
const ResizeFunction& resize_fn,
const std::vector<ValueRef>& resize_args,
const std::vector<PushConstantDataInfo>& push_constants)
const ResizeFunction& resize_fn)
: ExecuteNode(resize_fn, resize_args, args, shader.kernel_name),
shader_(shader),
global_workgroup_size_(global_workgroup_size),
Expand Down
4 changes: 2 additions & 2 deletions backends/vulkan/runtime/graph/ops/DispatchNode.h
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -33,10 +33,10 @@ class DispatchNode final : public ExecuteNode {
const utils::uvec3& local_workgroup_size,
const std::vector<ArgGroup>& args,
const vkapi::ParamsBindList& params,
const std::vector<PushConstantDataInfo>& push_constants = {},
const vkapi::SpecVarList& spec_vars = {},
const ResizeFunction& resize_fn = nullptr,
const std::vector<ValueRef>& resize_args = {},
const std::vector<PushConstantDataInfo>& push_constants = {});
const ResizeFunction& resize_fn = nullptr);

~DispatchNode() override = default;

Expand Down
9 changes: 6 additions & 3 deletions backends/vulkan/runtime/graph/ops/impl/Arange.cpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -94,16 +94,19 @@ void add_arange_node(
graph.create_global_wg_size(out),
graph.create_local_wg_size(out),
// Inputs and Outputs
{{out, vkapi::MemoryAccessType::WRITE}},
{{out, vkapi::kWrite}},
// Shader params buffers
{t_out->sizes_ubo(),
graph.create_params_buffer(start_val),
graph.create_params_buffer(step_val)},
// Push Constants
{},
// Specialization Constants
{},
// Resize Args
{start, end, step},
// Resizing Logic
resize_arange_node,
{start, end, step}));
resize_arange_node));
}

void arange(ComputeGraph& graph, const std::vector<ValueRef>& args) {
Expand Down
15 changes: 11 additions & 4 deletions backends/vulkan/runtime/graph/ops/impl/BatchNorm.cpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -90,12 +90,19 @@ void add_native_batch_norm_node(
VK_KERNEL_FROM_STR(kernel_name),
graph.create_global_wg_size(out_ref),
graph.create_local_wg_size(out_ref),
{{out_ref, vkapi::MemoryAccessType::WRITE},
{{in_ref, arg_weight, arg_bias, arg_mean, arg_var},
vkapi::MemoryAccessType::READ}},
{{out_ref, vkapi::kWrite},
{{in_ref, arg_weight, arg_bias, arg_mean, arg_var}, vkapi::kRead}},
{t_out->logical_limits_ubo(),
graph.create_params_buffer(epsilon),
graph.create_params_buffer(num_texel_per_batch)}));
graph.create_params_buffer(num_texel_per_batch)},
// Push Constants
{},
// Specialization Constants
{},
// Resize Args
{},
// Resizing Logic
nullptr));
}

void native_batch_norm(ComputeGraph& graph, const std::vector<ValueRef>& args) {
Expand Down
38 changes: 20 additions & 18 deletions backends/vulkan/runtime/graph/ops/impl/BinaryOp.cpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -84,19 +84,20 @@ void add_binary_op_texture_node(
graph.create_global_wg_size(out),
graph.create_local_wg_size(out),
// Inputs and Outputs
{{out, vkapi::MemoryAccessType::WRITE},
{{arg1, arg2}, vkapi::MemoryAccessType::READ}},
{{out, vkapi::kWrite}, {{arg1, arg2}, vkapi::kRead}},
// Shader params buffers
{},
// Specialization Constants
{t_out->hashed_layout(), t_in1->hashed_layout(), t_in2->hashed_layout()},
// Resizing Logic
resize_binary_op_node,
{},
// Push Constants
{{graph.sizes_pc_of(out),
graph.sizes_pc_of(arg1),
graph.sizes_pc_of(arg2),
PushConstantDataInfo(&binary_ops_params, sizeof(binary_ops_params))}}));
PushConstantDataInfo(&binary_ops_params, sizeof(binary_ops_params))}},
// Specialization Constants
{t_out->hashed_layout(), t_in1->hashed_layout(), t_in2->hashed_layout()},
// Resize Args
{},
// Resizing Logic
resize_binary_op_node));
}

void add_binary_op_buffer_node(
Expand Down Expand Up @@ -127,17 +128,10 @@ void add_binary_op_buffer_node(
graph.create_global_wg_size(out),
graph.create_local_wg_size(out),
// Inputs and Outputs
{{out, vkapi::MemoryAccessType::WRITE},
{{in1, in2}, vkapi::MemoryAccessType::READ}},
{{out, vkapi::kWrite}, {{in1, in2}, vkapi::kRead}},
// Shader params buffers
{},
// Specialization Constants
{graph.packed_dim_of(out),
graph.packed_dim_of(in1),
graph.packed_dim_of(in2)},
// Resizing Logic
resize_binary_op_node,
{},
// Push Constants
{{
graph.sizes_pc_of(in1),
graph.sizes_pc_of(in2),
Expand All @@ -146,7 +140,15 @@ void add_binary_op_buffer_node(
graph.strides_pc_of(in2),
graph.numel_pc_of(out),
PushConstantDataInfo(&alpha_val, sizeof(float)),
}}));
}},
// Specialization Constants
{graph.packed_dim_of(out),
graph.packed_dim_of(in1),
graph.packed_dim_of(in2)},
// Resize Args
{},
// Resizing Logic
resize_binary_op_node));
}

void add_binary_op_node(
Expand Down
12 changes: 12 additions & 0 deletions backends/vulkan/runtime/graph/ops/impl/Clone.cpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -50,8 +50,12 @@ void add_clone_node(
{{out, vkapi::kWrite}, {in, vkapi::kRead}},
// Parameter Buffers
{t_out->logical_limits_ubo()},
// Push Constants
{},
// Specialization Constants
{},
// Resize Args
{},
// Resizing Logic
resize_clone_node));
}
Expand All @@ -74,8 +78,12 @@ void add_image_to_buffer_node(
{{buffer, vkapi::kWrite}, {image, vkapi::kRead}},
// Parameter Buffers
{graph.sizes_ubo(image), graph.strides_ubo(buffer)},
// Push Constants
{},
// Specialization Constants
{graph.hashed_layout_of(image)},
// Resize Args
{},
// Resizing Logic
resize_clone_node));
}
Expand All @@ -98,8 +106,12 @@ void add_buffer_to_image_node(
{{image, vkapi::kWrite}, {buffer, vkapi::kRead}},
// Parameter Buffers
{graph.sizes_ubo(image), graph.strides_ubo(buffer)},
// Push Constants
{},
// Specialization Constants
{graph.hashed_layout_of(image)},
// Resize Args
{},
// Resizing Logic
resize_clone_node));
}
Expand Down
21 changes: 12 additions & 9 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 @@ -444,16 +444,17 @@ void add_conv2d_node(
wg_size,
graph.create_local_wg_size(wg_size),
// Inputs and Outputs
{{out, vkapi::MemoryAccessType::WRITE},
{{in, arg_weight, arg_bias}, vkapi::MemoryAccessType::READ}},
{{out, vkapi::kWrite}, {{in, arg_weight, arg_bias}, vkapi::kRead}},
// Shader params buffers
param_buffers,
// Push Constants
push_constants,
// Specialization Constants
{},
// Resizing Logic
resize_conv2d_node,
// Resize Args
{weight_data, stride, padding, dilation, transposed, output_padding},
push_constants));
// Resizing Logic
resize_conv2d_node));
}

void add_conv1d_node(
Expand Down Expand Up @@ -541,23 +542,25 @@ void add_conv1d_node(
global_size,
local_size,
// Inputs and Outputs
{{out, vkapi::MemoryAccessType::WRITE},
{{in, arg_weight, arg_bias}, vkapi::MemoryAccessType::READ}},
{{out, vkapi::kWrite}, {{in, arg_weight, arg_bias}, vkapi::kRead}},
// Shader params buffers
{
t_out->logical_limits_ubo(),
t_in->sizes_ubo(),
graph.create_params_buffer(kernel_params),
graph.create_params_buffer(out_params),
},
// Push Constants
{},
// Specialization Constants
{t_out->hashed_layout(),
t_in->hashed_layout(),
t_weight->hashed_layout(),
t_bias->hashed_layout()},
// Resize Args
{weight, stride, padding, dilation},
// Resizing Logic
resize_conv1d_node,
{weight, stride, padding, dilation}));
resize_conv1d_node));
}

void conv(ComputeGraph& graph, const std::vector<ValueRef>& args) {
Expand Down
54 changes: 31 additions & 23 deletions backends/vulkan/runtime/graph/ops/impl/Copy.cpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -50,19 +50,22 @@ void add_copy_offset_node(
},
// Parameter buffers
{},
// Push Constants
{
PushConstantDataInfo(&range, sizeof(range), sizeof(ivec4)),
PushConstantDataInfo(&src_offset, sizeof(src_offset), sizeof(ivec4)),
PushConstantDataInfo(&dst_offset, sizeof(dst_offset), sizeof(ivec4)),
},
// Specialization Constants
{graph.hashed_layout_of(out),
graph.hashed_layout_of(in),
(calc_out_pos_using_src_chnl ? 1
: calc_in_pos_using_dst_chnl ? 2
: 0)},
nullptr,
// Resize Args
{},
{
PushConstantDataInfo(&range, sizeof(range), sizeof(ivec4)),
PushConstantDataInfo(&src_offset, sizeof(src_offset), sizeof(ivec4)),
PushConstantDataInfo(&dst_offset, sizeof(dst_offset), sizeof(ivec4)),
}));
// Resizing Logic
nullptr));
}

void add_copy_packed_dim_offset_node(
Expand Down Expand Up @@ -138,22 +141,25 @@ void add_copy_packed_dim_offset_node(
graph.create_local_wg_size(global_wg_size),
// Inputs and Outputs
{
{out, vkapi::MemoryAccessType::WRITE},
{out, vkapi::MemoryAccessType::READ},
{in, vkapi::MemoryAccessType::READ},
{out, vkapi::kWrite},
{out, vkapi::kRead},
{in, vkapi::kRead},
},
// Parameter buffers
{},
// Specialization Constants
{graph.hashed_layout_of(out), graph.hashed_layout_of(in)},
nullptr,
{},
// Push Constants
{
PushConstantDataInfo(
&final_range, sizeof(final_range), sizeof(ivec4)),
PushConstantDataInfo(&src_offset, sizeof(src_offset), sizeof(ivec4)),
PushConstantDataInfo(&dst_offset, sizeof(dst_offset), sizeof(ivec4)),
}));
},
// Specialization Constants
{graph.hashed_layout_of(out), graph.hashed_layout_of(in)},
// Resize Args
{},
// Resizing Logic
nullptr));
}

void add_copy_channel_offset_node(
Expand Down Expand Up @@ -248,22 +254,24 @@ void add_copy_channel_offset_node(
local_size,
// Inputs and Outputs
{
{out, vkapi::MemoryAccessType::WRITE},
{out, vkapi::MemoryAccessType::READ},
{in, vkapi::MemoryAccessType::READ},
{out, vkapi::kWrite},
{out, vkapi::kRead},
{in, vkapi::kRead},
},
// Parameter buffers
{},
// Specialization Constants
{graph.hashed_layout_of(out), graph.hashed_layout_of(in)},
nullptr,
{},
// Push Constants
{graph.sizes_pc_of(out),
graph.sizes_pc_of(in),
PushConstantDataInfo(&range_params, sizeof(range_params)),
PushConstantDataInfo(&offset_params, sizeof(offset_params)),
PushConstantDataInfo(
&src_channel_offset, sizeof(src_channel_offset))}));
PushConstantDataInfo(&src_channel_offset, sizeof(src_channel_offset))},
// Specialization Constants
{graph.hashed_layout_of(out), graph.hashed_layout_of(in)},
// Resize Args
{},
// Resizing Logic
nullptr));
}
}

Expand Down
9 changes: 8 additions & 1 deletion backends/vulkan/runtime/graph/ops/impl/Embedding.cpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -58,9 +58,16 @@ void add_embedding_node(
{
t_out->sizes_ubo(),
},
// Push Constants
{},
// Specialization Constants
{t_out->hashed_layout(),
t_in->hashed_layout(),
t_weight->hashed_layout()}));
t_weight->hashed_layout()},
// Resize Args
{},
// Resizing Logic
nullptr));
}

void embedding(ComputeGraph& graph, const std::vector<ValueRef>& args) {
Expand Down
4 changes: 4 additions & 0 deletions backends/vulkan/runtime/graph/ops/impl/Flip.cpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -74,8 +74,12 @@ void add_flip_node(
graph.sizes_ubo(out),
graph.create_params_buffer(dim_bitmap),
},
// Push Constants
{},
// Specialization Constants
{},
// Resize Args
{},
// Resizing Logic
resize_flip_node));
}
Expand Down
9 changes: 6 additions & 3 deletions backends/vulkan/runtime/graph/ops/impl/Full.cpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -50,14 +50,17 @@ void add_full_node(
graph.create_global_wg_size(out),
graph.create_local_wg_size(out),
// Inputs and Outputs
{{out, vkapi::MemoryAccessType::WRITE}},
{{out, vkapi::kWrite}},
// Shader params buffers
{t_out->sizes_ubo(), graph.create_params_buffer(fill_value_val)},
// Push Constants
{},
// Specialization Constants
{SV(t_out->packed_dim())},
// Resize Args
{size_or_in},
// Resizing Logic
resize_full_node,
{size_or_in}));
resize_full_node));
}

void full(ComputeGraph& graph, const std::vector<ValueRef>& args) {
Expand Down
Loading
Loading