diff --git a/backends/vulkan/runtime/api/Context.cpp b/backends/vulkan/runtime/api/Context.cpp index 41c9d39ddcd..05c7e5a635c 100644 --- a/backends/vulkan/runtime/api/Context.cpp +++ b/backends/vulkan/runtime/api/Context.cpp @@ -73,16 +73,11 @@ void Context::cmd_reset_querypool() { void Context::report_shader_dispatch_start( const std::string& shader_name, - const utils::uvec3& global_wg_size, - const LocalWorkGroup& local_wg_size, + const GlobalWorkGrid& gwg, + const LocalWorkGroup& lwg, const uint32_t dispatch_id) { if (querypool_) { - querypool_.shader_profile_begin( - cmd_, - dispatch_id, - shader_name, - vkapi::create_extent3d(global_wg_size), - vkapi::create_extent3d((utils::uvec3)local_wg_size)); + querypool_.shader_profile_begin(cmd_, dispatch_id, shader_name, gwg, lwg); } } @@ -133,7 +128,7 @@ void Context::check_device_capabilities(const vkapi::ShaderInfo& shader) { vkapi::DescriptorSet Context::get_descriptor_set( const vkapi::ShaderInfo& shader_descriptor, - const LocalWorkGroup& local_workgroup_size, + const LocalWorkGroup& lwg, const vkapi::SpecVarList& additional_constants, const uint32_t push_constants_size) { VkDescriptorSetLayout shader_layout = @@ -142,10 +137,7 @@ vkapi::DescriptorSet Context::get_descriptor_set( VkPipelineLayout pipeline_layout = pipeline_layout_cache().retrieve(shader_layout, push_constants_size); - vkapi::SpecVarList spec_constants = { - SV(local_workgroup_size[0u]), - SV(local_workgroup_size[1u]), - SV(local_workgroup_size[2u])}; + vkapi::SpecVarList spec_constants = {SV(lwg[0u]), SV(lwg[1u]), SV(lwg[2u])}; spec_constants.append(additional_constants); @@ -158,7 +150,7 @@ vkapi::DescriptorSet Context::get_descriptor_set( spec_constants, resolved_required_subgroup_size}); - cmd_.bind_pipeline(pipeline, pipeline_layout, local_workgroup_size); + cmd_.bind_pipeline(pipeline, pipeline_layout, lwg); return descriptor_pool().get_descriptor_set( shader_layout, shader_descriptor.kernel_layout); @@ -168,30 +160,35 @@ void Context::register_shader_dispatch( const vkapi::DescriptorSet& descriptors, vkapi::PipelineBarrier& pipeline_barrier, const vkapi::ShaderInfo& shader_descriptor, - const utils::uvec3& global_workgroup_size, + const GlobalWorkGrid& gwg, + const LocalWorkGroup& lwg, const void* push_constants_data, const uint32_t push_constants_size) { + vkapi::Adapter* const adapter = adapter_ptr(); + const LocalWorkGroup& dispatch_lwg = + gwg.required_lwg_size().is_valid() ? gwg.required_lwg_size() : lwg; + dispatch_lwg.validate( + adapter->max_compute_workgroup_size(), + adapter->max_compute_workgroup_invocations()); + gwg.validate( + dispatch_lwg, + adapter->max_compute_workgroup_count(), + shader_descriptor.out_tile_size); + // Adjust the global workgroup size based on the output tile size - uint32_t global_wg_w = utils::div_up( - global_workgroup_size[0u], shader_descriptor.out_tile_size[0u]); - uint32_t global_wg_h = utils::div_up( - global_workgroup_size[1u], shader_descriptor.out_tile_size[1u]); - uint32_t global_wg_d = utils::div_up( - global_workgroup_size[2u], shader_descriptor.out_tile_size[2u]); + uint32_t gwg_w = utils::div_up(gwg[0u], shader_descriptor.out_tile_size[0u]); + uint32_t gwg_h = utils::div_up(gwg[1u], shader_descriptor.out_tile_size[1u]); + uint32_t gwg_d = utils::div_up(gwg[2u], shader_descriptor.out_tile_size[2u]); // Submitting a global work group size of 0 is undefined behaviour. If this is // detected then submit a single workgroup instead. - if (global_wg_w == 0u || global_wg_h == 0u || global_wg_d == 0u) { - global_wg_w = 1u; - global_wg_h = 1u; - global_wg_d = 1u; + if (gwg_w == 0u || gwg_h == 0u || gwg_d == 0u) { + gwg_w = 1u; + gwg_h = 1u; + gwg_d = 1u; } - const utils::uvec3 effective_global_wg = { - global_wg_w, - global_wg_h, - global_wg_d, - }; + const GlobalWorkGrid effective_gwg({gwg_w, gwg_h, gwg_d}, kExplicitWorkGrid); cmd_.bind_descriptors(descriptors.get_bind_handle()); cmd_.insert_barrier(pipeline_barrier); @@ -205,7 +202,7 @@ void Context::register_shader_dispatch( pipeline_layout, push_constants_data, push_constants_size); } - cmd_.dispatch(effective_global_wg); + cmd_.dispatch(effective_gwg, dispatch_lwg); } void Context::register_barrier(vkapi::PipelineBarrier& pipeline_barrier) { @@ -324,11 +321,8 @@ VkPipeline Context::get_shader_pipeline( VkPipelineLayout pipeline_layout = pipeline_layout_cache().retrieve(shader_layout, push_constants_size); - const LocalWorkGroup local_workgroup_size(4u, 4u, 1u); - vkapi::SpecVarList spec_constants = { - SV(local_workgroup_size[0u]), - SV(local_workgroup_size[1u]), - SV(local_workgroup_size[2u])}; + const LocalWorkGroup lwg(4u, 4u, 1u); + vkapi::SpecVarList spec_constants = {SV(lwg[0u]), SV(lwg[1u]), SV(lwg[2u])}; spec_constants.append(additional_constants); diff --git a/backends/vulkan/runtime/api/Context.h b/backends/vulkan/runtime/api/Context.h index d54cd2cce48..22abede2a5b 100644 --- a/backends/vulkan/runtime/api/Context.h +++ b/backends/vulkan/runtime/api/Context.h @@ -151,8 +151,8 @@ class Context final { */ void report_shader_dispatch_start( const std::string& shader_name, - const utils::uvec3& global_wg_size, - const LocalWorkGroup& local_wg_size, + const GlobalWorkGrid& gwg, + const LocalWorkGroup& lwg, const uint32_t dispatch_id = UINT32_MAX); /* @@ -197,15 +197,16 @@ class Context final { inline vkapi::DescriptorSet get_descriptor_set( const vkapi::ShaderInfo& shader_descriptor, - const LocalWorkGroup& local_work_group_size) { - return get_descriptor_set(shader_descriptor, local_work_group_size, {}, 0u); + const LocalWorkGroup& lwg) { + return get_descriptor_set(shader_descriptor, lwg, {}, 0u); } void register_shader_dispatch( const vkapi::DescriptorSet&, vkapi::PipelineBarrier&, const vkapi::ShaderInfo&, - const utils::uvec3&, + const GlobalWorkGrid&, + const LocalWorkGroup&, const void* = nullptr, const uint32_t = 0); @@ -220,8 +221,8 @@ class Context final { bool submit_compute_job( const vkapi::ShaderInfo&, vkapi::PipelineBarrier&, - const utils::uvec3&, - const utils::uvec3&, + const GlobalWorkGrid&, + const LocalWorkGroup&, const vkapi::SpecVarList&, VkFence fence_handle, const uint32_t dispatch_id, @@ -334,8 +335,8 @@ template inline bool Context::submit_compute_job( const vkapi::ShaderInfo& shader, vkapi::PipelineBarrier& pipeline_barrier, - const utils::uvec3& global_work_group, - const utils::uvec3& local_work_group_size, + const GlobalWorkGrid& gwg, + const LocalWorkGroup& lwg, const vkapi::SpecVarList& specialization_constants, VkFence fence_handle, const uint32_t dispatch_id, @@ -368,20 +369,17 @@ inline bool Context::submit_compute_job( set_cmd(); + const LocalWorkGroup& dispatch_lwg = + gwg.required_lwg_size().is_valid() ? gwg.required_lwg_size() : lwg; + report_shader_dispatch_start( - shader.kernel_name, - global_work_group, - LocalWorkGroup(local_work_group_size), - dispatch_id); + shader.kernel_name, gwg, dispatch_lwg, dispatch_id); // Factor out template parameter independent code to minimize code bloat. // Note that push constants are not exposed yet via this API, therefore the // push constants size is assumed to be 0. - vkapi::DescriptorSet descriptor_set = get_descriptor_set( - shader, - LocalWorkGroup(local_work_group_size), - specialization_constants, - 0u); + vkapi::DescriptorSet descriptor_set = + get_descriptor_set(shader, dispatch_lwg, specialization_constants, 0u); detail::bind( descriptor_set, @@ -390,7 +388,7 @@ inline bool Context::submit_compute_job( // Factor out template parameter independent code to minimize code bloat. register_shader_dispatch( - descriptor_set, pipeline_barrier, shader, global_work_group); + descriptor_set, pipeline_barrier, shader, gwg, dispatch_lwg); report_shader_dispatch_end(); diff --git a/backends/vulkan/runtime/graph/ComputeGraph.cpp b/backends/vulkan/runtime/graph/ComputeGraph.cpp index 6a49c8b00f3..f23d1f19c66 100644 --- a/backends/vulkan/runtime/graph/ComputeGraph.cpp +++ b/backends/vulkan/runtime/graph/ComputeGraph.cpp @@ -845,7 +845,7 @@ void ComputeGraph::update_descriptor_counts( void ComputeGraph::register_pipeline_to_create( const vkapi::ShaderInfo& shader_info, - const LocalWorkGroup& local_workgroup_size, + const LocalWorkGroup& lwg, const vkapi::SpecVarList& spec_vars, const std::vector& push_constants) { VkDescriptorSetLayout shader_layout = @@ -857,10 +857,7 @@ void ComputeGraph::register_pipeline_to_create( pc_offset += pc.write(pc_data.data(), pc_offset, kMaxPushConstantSize); } - vkapi::SpecVarList spec_constants = { - SV(local_workgroup_size[0u]), - SV(local_workgroup_size[1u]), - SV(local_workgroup_size[2u])}; + vkapi::SpecVarList spec_constants = {SV(lwg[0u]), SV(lwg[1u]), SV(lwg[2u])}; spec_constants.append(spec_vars); @@ -889,64 +886,46 @@ void ComputeGraph::register_pipeline_to_create( pipeline_descriptors_.insert(desc); } -utils::uvec3 ComputeGraph::create_global_wg_size(const ValueRef idx) { +GlobalWorkGrid ComputeGraph::create_gwg(const ValueRef idx) { if (is_buffer_storage(idx)) { - return {uint32_t(numel_of(idx)), 1u, 1u}; + return create_linear_gwg(utils::safe_downcast(numel_of(idx))); } - return logical_limits_of(idx); -} -utils::uvec3 ComputeGraph::create_local_wg_size( - const utils::uvec3 global_wg_size) { - if (config_.enable_local_wg_size_override) { - return config_.local_wg_size_override; - } + return GlobalWorkGrid( + utils::make_uvec3(logical_limits_of(idx)), kTextureExtentsWorkGrid); +} - // array containing axis index and global workgroup size - std::pair global_wg_size_desc[] = { - {0u, global_wg_size[0]}, - {1u, global_wg_size[1]}, - {2u, global_wg_size[2]}}; +GlobalWorkGrid ComputeGraph::create_linear_gwg(const uint64_t numel) { + vkapi::Adapter* const adapter = context()->adapter_ptr(); + GlobalWorkGrid gwg( + {utils::safe_downcast(numel), 1u, 1u}, kLinearWorkGrid); + gwg.wrap_linear_dispatch( + adapter->max_compute_workgroup_count(), + adapter->recommended_lwg_nthreads()); + return gwg; +} - // sort the global workgroup size in descending order - if (global_wg_size_desc[0].second < global_wg_size_desc[1].second) { - std::swap(global_wg_size_desc[0], global_wg_size_desc[1]); - } - if (global_wg_size_desc[1].second < global_wg_size_desc[2].second) { - std::swap(global_wg_size_desc[1], global_wg_size_desc[2]); - } - if (global_wg_size_desc[0].second < global_wg_size_desc[1].second) { - std::swap(global_wg_size_desc[0], global_wg_size_desc[1]); +LocalWorkGroup ComputeGraph::create_lwg(const GlobalWorkGrid& gwg) { + if (gwg.required_lwg_size().is_valid()) { + return gwg.required_lwg_size(); } - utils::uvec3 local_group_size = { - 8, - global_wg_size_desc[1].second >= 4u ? 4u - : global_wg_size_desc[1].second >= 2u ? 2u - : 1u, - global_wg_size_desc[2].second >= 2u ? 2u : 1u}; + vkapi::Adapter* const adapter = context()->adapter_ptr(); - if (global_wg_size_desc[2u].second == 1) { - if (global_wg_size_desc[1u].second == 1) { - local_group_size[0u] = 64; - local_group_size[1u] = 1; - } else if (global_wg_size_desc[1u].second % 4 == 0) { - local_group_size[0u] = 16; - local_group_size[1u] = 4; - } else { - local_group_size[0u] = 32; - local_group_size[1u] = 2; - } + utils::uvec3 shape_weights{ + gwg[0] > 1u ? 1u : 0u, gwg[1] > 1u ? 1u : 0u, gwg[2] > 1u ? 1u : 0u}; + if (shape_weights == utils::uvec3{0u, 0u, 0u}) { + shape_weights[0] = 1u; } - return { - local_group_size[global_wg_size_desc[0].first], - local_group_size[global_wg_size_desc[1].first], - local_group_size[global_wg_size_desc[2].first]}; + LocalWorkGroup lwg( + LwgShape(shape_weights), adapter->recommended_lwg_nthreads()); + lwg.fit_to_global(gwg); + return lwg; } -utils::uvec3 ComputeGraph::create_local_wg_size(const ValueRef idx) { - return create_local_wg_size(create_global_wg_size(idx)); +LocalWorkGroup ComputeGraph::create_lwg(const ValueRef idx) { + return create_lwg(create_gwg(idx)); } void ComputeGraph::bind_tensor_to_descriptor_set( diff --git a/backends/vulkan/runtime/graph/ComputeGraph.h b/backends/vulkan/runtime/graph/ComputeGraph.h index 192e20f99b3..1de890efb38 100644 --- a/backends/vulkan/runtime/graph/ComputeGraph.h +++ b/backends/vulkan/runtime/graph/ComputeGraph.h @@ -1012,7 +1012,7 @@ class ComputeGraph final { void register_pipeline_to_create( const vkapi::ShaderInfo& shader_info, - const LocalWorkGroup& local_workgroup_size, + const LocalWorkGroup& lwg, const vkapi::SpecVarList& spec_vars, const std::vector& push_constants); @@ -1025,37 +1025,36 @@ class ComputeGraph final { // /* - * Create a global workgroup size for a given `api::vTensor` value assuming - * that every shader invocation calculates one texel element of the output - * tensor. + * Create a global invocation size for a given `api::vTensor` value assuming + * that every shader invocation calculates one element of the output tensor. * * For tensors that use texture storage, the image extents of the * `api::vTensor` will be used to set the global workgroup size. * - * For tensor that use buffer storage, the number of texels in the texel - * buffer will be used to set the x component of the global workgroup size. - * All other components will be set to 1 (i.e. {ntexels, 1, 1} will be - * returned). + * Buffer tensors use linear dispatch intent. Oversized X dimensions are + * wrapped across X and Y according to device workgroup-count limits. */ - utils::uvec3 create_global_wg_size(const ValueRef idx); + GlobalWorkGrid create_gwg(const ValueRef idx); + + GlobalWorkGrid create_linear_gwg(const uint64_t numel); /* - * Suggest a local workgroup size for a given global workgroup size. + * Suggest a local workgroup size for a given global invocation size. * * The local workgroup size will be formed to try and minimize the number of * inactive invocations. * - * Currently, the local workgroup size is hard-coded to contain a total of 64 - * shader invocations. In the future, this value can be configured. + * Linear dispatches return their binding hint. Other dispatches use a shape + * heuristic targeting the adapter's recommended thread count. */ - utils::uvec3 create_local_wg_size(const utils::uvec3 global_wg_size); + LocalWorkGroup create_lwg(const GlobalWorkGrid& gwg); /* * Convenience function to suggest a local workgroup size for a given * `api::vTensor` value, assuming that every shader invocation calculates one * texel element of the output tensor. */ - utils::uvec3 create_local_wg_size(const ValueRef idx); + LocalWorkGroup create_lwg(const ValueRef idx); void bind_tensor_to_descriptor_set( const ValueRef ref, diff --git a/backends/vulkan/runtime/graph/GraphConfig.cpp b/backends/vulkan/runtime/graph/GraphConfig.cpp index 9a919a42573..6372cfb41c5 100644 --- a/backends/vulkan/runtime/graph/GraphConfig.cpp +++ b/backends/vulkan/runtime/graph/GraphConfig.cpp @@ -61,9 +61,6 @@ GraphConfig::GraphConfig() { // dispatches. By default, this functionality is disabled. enable_querypool = false; - enable_local_wg_size_override = false; - local_wg_size_override = {}; - has_data_dependent_shapes = false; expect_dynamic_shapes = false; force_resize = false; @@ -82,10 +79,4 @@ void GraphConfig::set_memory_layout_override( memory_layout_override = memory_layout; } -void GraphConfig::set_local_wg_size_override( - const utils::uvec3& local_wg_size) { - enable_local_wg_size_override = true; - local_wg_size_override = local_wg_size; -} - } // namespace vkcompute diff --git a/backends/vulkan/runtime/graph/GraphConfig.h b/backends/vulkan/runtime/graph/GraphConfig.h index 20d01362ef1..8d3ca3d611d 100644 --- a/backends/vulkan/runtime/graph/GraphConfig.h +++ b/backends/vulkan/runtime/graph/GraphConfig.h @@ -30,9 +30,6 @@ struct GraphConfig final { bool enable_querypool; - bool enable_local_wg_size_override; - utils::uvec3 local_wg_size_override; - // If true, then resize functions should always be called even if input shapes // have not changed. bool has_data_dependent_shapes = false; @@ -82,7 +79,6 @@ struct GraphConfig final { void set_storage_type_override(utils::StorageType storage_type); void set_memory_layout_override(utils::GPUMemoryLayout memory_layout); - void set_local_wg_size_override(const utils::uvec3& local_wg_size); }; } // namespace vkcompute diff --git a/backends/vulkan/runtime/graph/ops/BlitNode.cpp b/backends/vulkan/runtime/graph/ops/BlitNode.cpp index 6c780321641..8de8f7f9304 100644 --- a/backends/vulkan/runtime/graph/ops/BlitNode.cpp +++ b/backends/vulkan/runtime/graph/ops/BlitNode.cpp @@ -44,7 +44,10 @@ void BlitNode::encode(ComputeGraph* graph) { kernel_name += vkapi::to_string(graph->dtype_of(dst_)); context->report_shader_dispatch_start( - kernel_name, utils::uvec3(), LocalWorkGroup(), node_id_); + kernel_name, + GlobalWorkGrid({0u, 0u, 0u}, kExplicitWorkGrid), + LocalWorkGroup(), + node_id_); context->register_blit( pipeline_barrier, diff --git a/backends/vulkan/runtime/graph/ops/DispatchNode.cpp b/backends/vulkan/runtime/graph/ops/DispatchNode.cpp index ab48ec3f4c3..af4d5250dcd 100644 --- a/backends/vulkan/runtime/graph/ops/DispatchNode.cpp +++ b/backends/vulkan/runtime/graph/ops/DispatchNode.cpp @@ -17,8 +17,8 @@ namespace vkcompute { DispatchNode::DispatchNode( ComputeGraph& graph, const vkapi::ShaderInfo& shader, - const utils::uvec3& global_workgroup_size, - const utils::uvec3& local_workgroup_size, + const GlobalWorkGrid& gwg, + const LocalWorkGroup& lwg, const std::vector& args, const vkapi::ParamsBindList& params, const std::vector& push_constants, @@ -27,8 +27,8 @@ DispatchNode::DispatchNode( const ResizeFunction& resize_fn) : ExecuteNode(resize_fn, resize_args, args, shader.kernel_name), shader_(shader), - global_workgroup_size_(global_workgroup_size), - local_workgroup_size_(local_workgroup_size), + gwg_(gwg), + lwg_(gwg.required_lwg_size().is_valid() ? gwg.required_lwg_size() : lwg), params_(params), spec_vars_(spec_vars), push_constants_(push_constants) { @@ -37,7 +37,7 @@ DispatchNode::DispatchNode( void DispatchNode::prepare_pipelines(ComputeGraph* graph) { graph->register_pipeline_to_create( - shader_, local_workgroup_size_, spec_vars_, push_constants_); + shader_, lwg_, spec_vars_, push_constants_); } void DispatchNode::encode(ComputeGraph* graph) { @@ -46,8 +46,7 @@ void DispatchNode::encode(ComputeGraph* graph) { } // If any global wg size element is 0, then skip encoding this shader - if (global_workgroup_size_[0] == 0 || global_workgroup_size_[1] == 0 || - global_workgroup_size_[2] == 0) { + if (gwg_[0] == 0 || gwg_[1] == 0 || gwg_[2] == 0) { return; } @@ -75,12 +74,12 @@ void DispatchNode::encode(ComputeGraph* graph) { #else shader_.kernel_name, #endif - global_workgroup_size_, - local_workgroup_size_, + gwg_, + lwg_, node_id_); vkapi::DescriptorSet descriptor_set = context->get_descriptor_set( - shader_, local_workgroup_size_, spec_vars_, push_constants_offset_); + shader_, lwg_, spec_vars_, push_constants_offset_); uint32_t idx = 0; idx = bind_values_to_descriptor_set( @@ -92,7 +91,8 @@ void DispatchNode::encode(ComputeGraph* graph) { descriptor_set, pipeline_barrier, shader_, - global_workgroup_size_, + gwg_, + lwg_, push_constants_data_.data(), push_constants_offset_); diff --git a/backends/vulkan/runtime/graph/ops/DispatchNode.h b/backends/vulkan/runtime/graph/ops/DispatchNode.h index f0c4245695d..1b39608eff4 100644 --- a/backends/vulkan/runtime/graph/ops/DispatchNode.h +++ b/backends/vulkan/runtime/graph/ops/DispatchNode.h @@ -31,8 +31,8 @@ class DispatchNode : public ExecuteNode { explicit DispatchNode( ComputeGraph& graph, const vkapi::ShaderInfo& shader, - const utils::uvec3& global_workgroup_size, - const utils::uvec3& local_workgroup_size, + const GlobalWorkGrid& gwg, + const LocalWorkGroup& lwg, const std::vector& args, const vkapi::ParamsBindList& params, const std::vector& push_constants = {}, @@ -50,8 +50,8 @@ class DispatchNode : public ExecuteNode { protected: vkapi::ShaderInfo shader_; - utils::uvec3 global_workgroup_size_; - LocalWorkGroup local_workgroup_size_; + GlobalWorkGrid gwg_; + LocalWorkGroup lwg_; const vkapi::ParamsBindList params_; const vkapi::SpecVarList spec_vars_; const std::vector push_constants_; diff --git a/backends/vulkan/runtime/graph/ops/DynamicDispatchNode.cpp b/backends/vulkan/runtime/graph/ops/DynamicDispatchNode.cpp index 4861afc43ad..afdb5d571ca 100644 --- a/backends/vulkan/runtime/graph/ops/DynamicDispatchNode.cpp +++ b/backends/vulkan/runtime/graph/ops/DynamicDispatchNode.cpp @@ -15,8 +15,8 @@ namespace vkcompute { DynamicDispatchNode::DynamicDispatchNode( ComputeGraph& graph, const PickShaderFn& pick_shader_fn, - const PickGlobalFn& pick_global_wg_fn, - const PickLocalFn& pick_local_wg_fn, + const PickGwgFn& pick_gwg_fn, + const PickLwgFn& pick_lwg_fn, const std::vector& args, const vkapi::ParamsBindList& params, const std::vector& push_constants, @@ -26,8 +26,8 @@ DynamicDispatchNode::DynamicDispatchNode( : DispatchNode( graph, pick_shader_fn(&graph, args, resize_args), - {1u, 1u, 1u}, - {8u, 8u, 1u}, + GlobalWorkGrid({1u, 1u, 1u}, kExplicitWorkGrid), + LocalWorkGroup(8u, 8u, 1u), args, params, push_constants, @@ -35,25 +35,26 @@ DynamicDispatchNode::DynamicDispatchNode( resize_args, resize_fn), pick_shader_fn_(pick_shader_fn), - pick_global_wg_fn_(pick_global_wg_fn), - pick_local_wg_fn_(pick_local_wg_fn) { - global_workgroup_size_ = - pick_global_wg_fn(&graph, shader_, args, resize_args); - local_workgroup_size_ = LocalWorkGroup(pick_local_wg_fn( - &graph, shader_, global_workgroup_size_, args, resize_args)); + pick_gwg_fn_(pick_gwg_fn), + pick_lwg_fn_(pick_lwg_fn) { + gwg_ = pick_gwg_fn(&graph, shader_, args, resize_args); + lwg_ = pick_lwg_fn(&graph, shader_, gwg_, args, resize_args); + if (gwg_.required_lwg_size().is_valid()) { + lwg_ = gwg_.required_lwg_size(); + } // Calculate dispatch grid similar to Context.cpp register_shader_dispatch wg_dispatch_grid_ = { - utils::div_up(global_workgroup_size_[0], local_workgroup_size_[0]), - utils::div_up(global_workgroup_size_[1], local_workgroup_size_[1]), - utils::div_up(global_workgroup_size_[2], local_workgroup_size_[2])}; + utils::div_up(gwg_[0], lwg_[0]), + utils::div_up(gwg_[1], lwg_[1]), + utils::div_up(gwg_[2], lwg_[2])}; } DynamicDispatchNode::DynamicDispatchNode( ComputeGraph& graph, const vkapi::ShaderInfo& shader, - const PickGlobalFn& pick_global_wg_fn, - const PickLocalFn& pick_local_wg_fn, + const PickGwgFn& pick_gwg_fn, + const PickLwgFn& pick_lwg_fn, const std::vector& args, const vkapi::ParamsBindList& params, const std::vector& push_constants, @@ -63,8 +64,8 @@ DynamicDispatchNode::DynamicDispatchNode( : DispatchNode( graph, shader, - {1u, 1u, 1u}, - {8u, 8u, 1u}, + GlobalWorkGrid({1u, 1u, 1u}, kExplicitWorkGrid), + LocalWorkGroup(8u, 8u, 1u), args, params, push_constants, @@ -72,17 +73,18 @@ DynamicDispatchNode::DynamicDispatchNode( resize_args, resize_fn), pick_shader_fn_{nullptr}, - pick_global_wg_fn_(pick_global_wg_fn), - pick_local_wg_fn_(pick_local_wg_fn) { - global_workgroup_size_ = - pick_global_wg_fn(&graph, shader_, args, resize_args); - local_workgroup_size_ = LocalWorkGroup(pick_local_wg_fn( - &graph, shader_, global_workgroup_size_, args, resize_args)); + pick_gwg_fn_(pick_gwg_fn), + pick_lwg_fn_(pick_lwg_fn) { + gwg_ = pick_gwg_fn(&graph, shader_, args, resize_args); + lwg_ = pick_lwg_fn(&graph, shader_, gwg_, args, resize_args); + if (gwg_.required_lwg_size().is_valid()) { + lwg_ = gwg_.required_lwg_size(); + } // Calculate the work group grid that will be dispatched wg_dispatch_grid_ = { - utils::div_up(global_workgroup_size_[0], local_workgroup_size_[0]), - utils::div_up(global_workgroup_size_[1], local_workgroup_size_[1]), - utils::div_up(global_workgroup_size_[2], local_workgroup_size_[2])}; + utils::div_up(gwg_[0], lwg_[0]), + utils::div_up(gwg_[1], lwg_[1]), + utils::div_up(gwg_[2], lwg_[2])}; } bool DynamicDispatchNode::trigger_resize(ComputeGraph* graph) { @@ -108,29 +110,30 @@ bool DynamicDispatchNode::trigger_resize(ComputeGraph* graph) { dispatch_params_changed = true; } } - if (pick_global_wg_fn_) { + if (pick_gwg_fn_) { // Note that if global workgroup size changes, then the dispatch params // may not actually be different. The actual value to check is the // work group grid size that will be dispatched, which is calculated // below. - global_workgroup_size_ = - pick_global_wg_fn_(graph, shader_, args_, resize_args_); + gwg_ = pick_gwg_fn_(graph, shader_, args_, resize_args_); } - if (pick_local_wg_fn_) { - utils::uvec3 new_local_wg_uvec3 = pick_local_wg_fn_( - graph, shader_, global_workgroup_size_, args_, resize_args_); - LocalWorkGroup new_local_wg = LocalWorkGroup(new_local_wg_uvec3); - if (local_workgroup_size_ != new_local_wg) { - local_workgroup_size_ = new_local_wg; + if (pick_lwg_fn_) { + LocalWorkGroup new_lwg = + pick_lwg_fn_(graph, shader_, gwg_, args_, resize_args_); + if (gwg_.required_lwg_size().is_valid()) { + new_lwg = gwg_.required_lwg_size(); + } + if (lwg_ != new_lwg) { + lwg_ = new_lwg; dispatch_params_changed = true; } } // Always recompute the new dispatch grid and check if it's different utils::uvec3 new_wg_dispatch_grid = { - utils::div_up(global_workgroup_size_[0], local_workgroup_size_[0]), - utils::div_up(global_workgroup_size_[1], local_workgroup_size_[1]), - utils::div_up(global_workgroup_size_[2], local_workgroup_size_[2])}; + utils::div_up(gwg_[0], lwg_[0]), + utils::div_up(gwg_[1], lwg_[1]), + utils::div_up(gwg_[2], lwg_[2])}; // Check if the new dispatch grid is different from the old one if (wg_dispatch_grid_ != new_wg_dispatch_grid) { diff --git a/backends/vulkan/runtime/graph/ops/DynamicDispatchNode.h b/backends/vulkan/runtime/graph/ops/DynamicDispatchNode.h index d3b82968eb2..7902c3d0d34 100644 --- a/backends/vulkan/runtime/graph/ops/DynamicDispatchNode.h +++ b/backends/vulkan/runtime/graph/ops/DynamicDispatchNode.h @@ -30,23 +30,23 @@ class DynamicDispatchNode final : public DispatchNode { ComputeGraph*, const std::vector&, const std::vector&)>; - using PickGlobalFn = const std::function&, const std::vector&)>; - using PickLocalFn = const std::function&, const std::vector&)>; explicit DynamicDispatchNode( ComputeGraph& graph, const PickShaderFn& pick_shader_fn, - const PickGlobalFn& pick_global_wg_fn, - const PickLocalFn& pick_local_wg_fn, + const PickGwgFn& pick_gwg_fn, + const PickLwgFn& pick_lwg_fn, const std::vector& args, const vkapi::ParamsBindList& params, const std::vector& push_constants, @@ -57,8 +57,8 @@ class DynamicDispatchNode final : public DispatchNode { explicit DynamicDispatchNode( ComputeGraph& graph, const vkapi::ShaderInfo& shader, - const PickGlobalFn& pick_global_wg_fn, - const PickLocalFn& pick_local_wg_fn, + const PickGwgFn& pick_gwg_fn, + const PickLwgFn& pick_lwg_fn, const std::vector& args, const vkapi::ParamsBindList& params, const std::vector& push_constants, @@ -72,8 +72,8 @@ class DynamicDispatchNode final : public DispatchNode { protected: const PickShaderFn pick_shader_fn_; - const PickGlobalFn pick_global_wg_fn_; - const PickLocalFn pick_local_wg_fn_; + const PickGwgFn pick_gwg_fn_; + const PickLwgFn pick_lwg_fn_; utils::uvec3 wg_dispatch_grid_{1u, 1u, 1u}; diff --git a/backends/vulkan/runtime/graph/ops/PrepackNode.cpp b/backends/vulkan/runtime/graph/ops/PrepackNode.cpp index f771870c566..40bd9bcfbca 100644 --- a/backends/vulkan/runtime/graph/ops/PrepackNode.cpp +++ b/backends/vulkan/runtime/graph/ops/PrepackNode.cpp @@ -18,16 +18,16 @@ namespace vkcompute { PrepackNode::PrepackNode( ComputeGraph& graph, const vkapi::ShaderInfo& shader, - const utils::uvec3& global_workgroup_size, - const utils::uvec3& local_workgroup_size, + const GlobalWorkGrid& gwg, + const LocalWorkGroup& lwg, const ValueRef tref, const ValueRef packed, const vkapi::ParamsBindList& params, const vkapi::SpecVarList& spec_vars, const std::vector& push_constants) : shader_(shader), - global_workgroup_size_(global_workgroup_size), - local_workgroup_size_(local_workgroup_size), + gwg_(gwg), + lwg_(gwg.required_lwg_size().is_valid() ? gwg.required_lwg_size() : lwg), tref_(tref), packed_(packed), params_(params), @@ -102,7 +102,7 @@ api::StagingBuffer PrepackNode::create_staging_buffer(ComputeGraph* graph) { void PrepackNode::prepare_pipelines(ComputeGraph* graph) { graph->register_pipeline_to_create( - shader_, local_workgroup_size_, spec_vars_, push_constants_); + shader_, lwg_, spec_vars_, push_constants_); } void PrepackNode::encode(ComputeGraph* graph) { @@ -131,7 +131,7 @@ void PrepackNode::encode(ComputeGraph* graph) { vkapi::PipelineBarrier pipeline_barrier{}; vkapi::DescriptorSet descriptor_set = context->get_descriptor_set( - shader_, local_workgroup_size_, spec_vars_, push_constants_offset); + shader_, lwg_, spec_vars_, push_constants_offset); uint32_t idx = 0; graph->bind_tensor_to_descriptor_set( @@ -147,7 +147,8 @@ void PrepackNode::encode(ComputeGraph* graph) { descriptor_set, pipeline_barrier, shader_, - global_workgroup_size_, + gwg_, + lwg_, push_constants_data.data(), push_constants_offset); } diff --git a/backends/vulkan/runtime/graph/ops/PrepackNode.h b/backends/vulkan/runtime/graph/ops/PrepackNode.h index 73e70fc78d5..8a301ef1e0a 100644 --- a/backends/vulkan/runtime/graph/ops/PrepackNode.h +++ b/backends/vulkan/runtime/graph/ops/PrepackNode.h @@ -32,8 +32,8 @@ class PrepackNode final { PrepackNode( ComputeGraph& graph, const vkapi::ShaderInfo& shader, - const utils::uvec3& global_workgroup_size, - const utils::uvec3& local_workgroup_size, + const GlobalWorkGrid& gwg, + const LocalWorkGroup& lwg, const ValueRef tref, const ValueRef packed, const vkapi::ParamsBindList& params, @@ -53,8 +53,8 @@ class PrepackNode final { protected: uint32_t node_id_; const vkapi::ShaderInfo shader_; - const utils::uvec3 global_workgroup_size_; - const LocalWorkGroup local_workgroup_size_; + GlobalWorkGrid gwg_; + const LocalWorkGroup lwg_; const ValueRef tref_; const ValueRef packed_; const vkapi::ParamsBindList params_; diff --git a/backends/vulkan/runtime/graph/ops/glsl/q8ta_conv2d.glsl b/backends/vulkan/runtime/graph/ops/glsl/q8ta_conv2d.glsl index 821f7f79b0e..644ef466bb6 100644 --- a/backends/vulkan/runtime/graph/ops/glsl/q8ta_conv2d.glsl +++ b/backends/vulkan/runtime/graph/ops/glsl/q8ta_conv2d.glsl @@ -77,14 +77,18 @@ void main() { // Thread mapping int oc4 = int(gl_GlobalInvocationID.z); int w4 = int(gl_GlobalInvocationID.x); + const int H = int(outp.sizes[0][1]); + const int hn = int(gl_GlobalInvocationID.y); + const int n = hn / H; + const int h = hn % H; // Initialize output tensor index (WHCN order) // Each thread handles 4 adjacent widths starting at base_out_w TensorIndex4D outp_tidx; outp_tidx.data[0] = w4 * 4; - outp_tidx.data[1] = int(gl_GlobalInvocationID.y); + outp_tidx.data[1] = h; outp_tidx.data[2] = oc4 * 4; - outp_tidx.data[3] = 0; + outp_tidx.data[3] = n; const int W = int(outp.sizes[0][0]); const int OC = int(outp.sizes[0][2]); @@ -113,6 +117,7 @@ void main() { const int inp_w_stride = int(inp.strides[0][0]); const int inp_h_stride = int(inp.strides[0][1]); const int inp_c_stride = int(inp.strides[0][2]); + const int inp_n_stride = int(inp.strides[0][3]); const int w_texel_step = conv2d_params.dilation.x * inp_w_stride; const int h_texel_step = conv2d_params.dilation.y * inp_h_stride; const int subtile_w_step = conv2d_params.stride.x * inp_w_stride; @@ -122,7 +127,7 @@ void main() { inp_tidx.data[0] = outp_tidx.data[0] * conv2d_params.stride.x - conv2d_params.padding.x; inp_tidx.data[1] = outp_tidx.data[1] * conv2d_params.stride.y - conv2d_params.padding.y; inp_tidx.data[2] = ic_group_start; - inp_tidx.data[3] = 0; + inp_tidx.data[3] = n; int base_inp_texel_idx; if (get_outer_packed_dim_block_size(inp_layout) == 1) { @@ -172,7 +177,11 @@ void main() { // inp_texel_idx = tensor4d_idx_to_texel_idx(inp, inp_tidx, inp_layout); const int w4 = div_4(inp_tidx.data[0]); const int inp_c4 = div_4(inp_tidx.data[2]); - inp_texel_idx = (inp_tidx.data[1] * inp_h_stride + w4 * inp_w_stride + inp_c4) * 4 + mod_4(inp_tidx.data[0]); + inp_texel_idx = + (n * inp_n_stride + inp_tidx.data[1] * inp_h_stride + + w4 * inp_w_stride + inp_c4) * + 4 + + mod_4(inp_tidx.data[0]); } packed_input = t_packed_int8_input[inp_texel_idx]; } diff --git a/backends/vulkan/runtime/graph/ops/glsl/q8ta_conv2d_dw.glsl b/backends/vulkan/runtime/graph/ops/glsl/q8ta_conv2d_dw.glsl index 7f4d03887df..673cfd68952 100644 --- a/backends/vulkan/runtime/graph/ops/glsl/q8ta_conv2d_dw.glsl +++ b/backends/vulkan/runtime/graph/ops/glsl/q8ta_conv2d_dw.glsl @@ -71,14 +71,18 @@ ivec4 quantize(const vec4 texel, const float inv_scale, const int zp) { void main() { const int c4 = int(gl_GlobalInvocationID.z); + const int H = int(outp.sizes[0][1]); + const int hn = int(gl_GlobalInvocationID.y); + const int n = hn / H; + const int h = hn % H; // Initialize output tensor index (WHCN order) // Each thread handles 4 adjacent widths starting at base_out_w TensorIndex4D outp_tidx; outp_tidx.data[0] = int(gl_GlobalInvocationID.x) * 4; - outp_tidx.data[1] = int(gl_GlobalInvocationID.y); + outp_tidx.data[1] = h; outp_tidx.data[2] = c4 * 4; - outp_tidx.data[3] = 0; + outp_tidx.data[3] = n; const int W = int(outp.sizes[0][0]); const int C4 = int(div_up_4(outp.sizes[0][2])); @@ -94,6 +98,7 @@ void main() { // Get strides for width and height dimensions (in texel space) const int w_stride = int(inp.strides[0][0]); const int h_stride = int(inp.strides[0][1]); + const int n_stride = int(inp.strides[0][3]); // Pre-compute step sizes for efficient indexing const int w_texel_step = conv2d_params.dilation.x * w_stride; @@ -106,7 +111,7 @@ void main() { inp_tidx.data[0] = outp_tidx.data[0] * conv2d_params.stride.x - conv2d_params.padding.x; inp_tidx.data[1] = outp_tidx.data[1] * conv2d_params.stride.y - conv2d_params.padding.y; inp_tidx.data[2] = outp_tidx.data[2]; - inp_tidx.data[3] = 0; // batch = 0 since N == 1 + inp_tidx.data[3] = n; int base_inp_texel_idx; if (get_outer_packed_dim_block_size(inp_layout) == 1) { @@ -152,7 +157,11 @@ void main() { // inp_texel_idx = base_inp_texel_idx + div_4(w_offset) * w_stride + mod_4(w_offset); // inp_texel_idx = tensor4d_idx_to_texel_idx(inp, inp_tidx, inp_layout); const int w4 = div_4(inp_tidx.data[0]); - inp_texel_idx = (inp_tidx.data[1] * h_stride + w4 * w_stride + c4) * 4 + mod_4(inp_tidx.data[0]); + inp_texel_idx = + (n * n_stride + inp_tidx.data[1] * h_stride + + w4 * w_stride + c4) * + 4 + + mod_4(inp_tidx.data[0]); } const int packed_input = t_packed_int8_input[inp_texel_idx]; input_4c = unpack_int8x4(packed_input); diff --git a/backends/vulkan/runtime/graph/ops/glsl/q8ta_conv2d_pw.glsl b/backends/vulkan/runtime/graph/ops/glsl/q8ta_conv2d_pw.glsl index 9f60bea9948..aeb98f7a41b 100644 --- a/backends/vulkan/runtime/graph/ops/glsl/q8ta_conv2d_pw.glsl +++ b/backends/vulkan/runtime/graph/ops/glsl/q8ta_conv2d_pw.glsl @@ -73,14 +73,17 @@ ${layout_declare_spec_const(C, "int", "inp_layout", "CONTIG_LAYOUT_INT")} int compute_outp_buffer_idx( const int w_block_idx, const int h_idx, - const int c_block_idx) { + const int c_block_idx, + const int n_idx) { if (get_outer_packed_dim_block_size(outp_layout) == 1) { - return h_idx * int(outp.strides[0][1]) + return n_idx * int(outp.strides[0][3]) + + h_idx * int(outp.strides[0][1]) + mul_4(w_block_idx) * int(outp.strides[0][0]) + c_block_idx * int(outp.strides[0][2]); } else { return mul_4( - h_idx * int(outp.strides[0][1]) + n_idx * int(outp.strides[0][3]) + + h_idx * int(outp.strides[0][1]) + w_block_idx * int(outp.strides[0][0]) + c_block_idx * int(outp.strides[0][2])); } @@ -91,20 +94,22 @@ void main() { // Thread mapping: each thread handles TILE_M widths x TILE_N output channels. // gl_GlobalInvocationID.x -> output channel blocks. // gl_GlobalInvocationID.y -> width blocks. - // gl_GlobalInvocationID.z → batch (or height * batch combined) + // gl_GlobalInvocationID.z -> height * batch. const int oc_block_idx = int(gl_GlobalInvocationID.x) * TILE_N4; const int ow_block_idx = int(gl_GlobalInvocationID.y) * TILE_M4; - const int oh = int(gl_GlobalInvocationID.z); // Get output extents in block space (div_up_4 for packed dimensions) const int W = int(outp.sizes[0][0]); const int W4 = div_up_4(int(outp.sizes[0][0])); const int H = int(outp.sizes[0][1]); const int OC4 = div_up_4(int(outp.sizes[0][2])); + const int hn = int(gl_GlobalInvocationID.z); + const int n = hn / H; + const int oh = hn % H; // Bounds check in block space if (ow_block_idx >= W4 || - oh >= H || + n >= int(outp.sizes[0][3]) || oc_block_idx >= OC4) { return; } @@ -118,6 +123,7 @@ void main() { const int inp_w_stride = int(inp.strides[0][0]); const int inp_h_stride = int(inp.strides[0][1]); const int inp_c_stride = int(inp.strides[0][2]); + const int inp_n_stride = int(inp.strides[0][3]); // Initialize int32 accumulator ivec4 out_accum[TILE_M][TILE_N4]; @@ -133,7 +139,8 @@ void main() { // Compute initial input tile index with group offset // For grouped im2col, each group's K range starts at group_idx * K4_per_group // For non-grouped (groups=1), group_idx is always 0 so offset is 0 - int input_idx = oh * inp_h_stride + int input_idx = n * inp_n_stride + + oh * inp_h_stride + ow_block_idx * inp_w_stride + group_idx * K4_per_group; @@ -256,7 +263,8 @@ void main() { const int base_outp_buffer_idx = compute_outp_buffer_idx( ow_block_idx + m4, oh, - oc_block_idx + n4); + oc_block_idx + n4, + n); if (oc_block_idx + n4 < OC4) { // Store individual ints from the ivec4 const int subtile_w_limit = min(4, W - mul_4(ow_block_idx + m4)); diff --git a/backends/vulkan/runtime/graph/ops/glsl/q8ta_im2col.glsl b/backends/vulkan/runtime/graph/ops/glsl/q8ta_im2col.glsl index 314788cb857..b0cc4866a03 100644 --- a/backends/vulkan/runtime/graph/ops/glsl/q8ta_im2col.glsl +++ b/backends/vulkan/runtime/graph/ops/glsl/q8ta_im2col.glsl @@ -50,15 +50,19 @@ void main() { const int im2col_W4 = div_up_4(im2col_sizes.x); const int im2col_H = im2col_sizes.y; const int im2col_Z4 = div_up_4(im2col_sizes.z); + const int im2col_N = im2col_sizes.w; // im2col block index from linear output buffer index const int c4_idx = out_buf_idx % im2col_Z4; const int row = out_buf_idx / im2col_Z4; const int w4_idx = row % im2col_W4; - const int h_idx = row / im2col_W4; + const int hn_idx = row / im2col_W4; + const int h_idx = hn_idx % im2col_H; + const int n_idx = hn_idx / im2col_H; // out of bounds check - if (w4_idx >= im2col_W4 || h_idx >= im2col_H || c4_idx >= im2col_Z4) { + if (w4_idx >= im2col_W4 || h_idx >= im2col_H || + c4_idx >= im2col_Z4 || n_idx >= im2col_N) { return; } @@ -108,12 +112,14 @@ void main() { const int x_mod = mod_4(x); int scalar_idx; if (get_outer_packed_dim_block_size(inp_layout) == 1) { - scalar_idx = input_y * int(inp.strides[0][1]) + scalar_idx = n_idx * int(inp.strides[0][3]) + + input_y * int(inp.strides[0][1]) + x * int(inp.strides[0][0]) + z4 * int(inp.strides[0][2]); } else { scalar_idx = mul_4( - input_y * int(inp.strides[0][1]) + n_idx * int(inp.strides[0][3]) + + input_y * int(inp.strides[0][1]) + x4 * int(inp.strides[0][0]) + z4) + x_mod; } @@ -122,7 +128,8 @@ void main() { } // store_packed_int8_output_tile (with TILE_M4=1, TILE_N4=1) - const int buffer_idx = h_idx * int(im2col_outp.strides[0][1]) + const int buffer_idx = n_idx * int(im2col_outp.strides[0][3]) + + h_idx * int(im2col_outp.strides[0][1]) + w4_idx * int(im2col_outp.strides[0][0]) + c4_idx; diff --git a/backends/vulkan/runtime/graph/ops/impl/AdamwStep.cpp b/backends/vulkan/runtime/graph/ops/impl/AdamwStep.cpp index 489e95a773d..4d881567e7b 100644 --- a/backends/vulkan/runtime/graph/ops/impl/AdamwStep.cpp +++ b/backends/vulkan/runtime/graph/ops/impl/AdamwStep.cpp @@ -70,8 +70,8 @@ void add_adamw_step_node( graph.execute_nodes().emplace_back(new DynamicDispatchNode( graph, VK_KERNEL_FROM_STR(kernel_name), - default_pick_global_wg_size, - default_pick_local_wg_size, + default_pick_gwg, + default_pick_lwg, // Inputs and Outputs {{param, vkapi::kReadWrite}, {m, vkapi::kReadWrite}, diff --git a/backends/vulkan/runtime/graph/ops/impl/Arange.cpp b/backends/vulkan/runtime/graph/ops/impl/Arange.cpp index bf6345c0f16..f635c9282f2 100644 --- a/backends/vulkan/runtime/graph/ops/impl/Arange.cpp +++ b/backends/vulkan/runtime/graph/ops/impl/Arange.cpp @@ -91,8 +91,8 @@ void add_arange_node( graph.execute_nodes().emplace_back(new DynamicDispatchNode( graph, VK_KERNEL_FROM_STR(kernel_name), - default_pick_global_wg_size, - default_pick_local_wg_size, + default_pick_gwg, + default_pick_lwg, // Inputs and Outputs {{out, vkapi::kWrite}}, // Shader params buffers diff --git a/backends/vulkan/runtime/graph/ops/impl/BatchNorm.cpp b/backends/vulkan/runtime/graph/ops/impl/BatchNorm.cpp index a6dd8f07f53..23ef5520079 100644 --- a/backends/vulkan/runtime/graph/ops/impl/BatchNorm.cpp +++ b/backends/vulkan/runtime/graph/ops/impl/BatchNorm.cpp @@ -99,8 +99,8 @@ void add_native_batch_norm_node( graph.execute_nodes().emplace_back(new DynamicDispatchNode( graph, VK_KERNEL_FROM_STR(kernel_name), - default_pick_global_wg_size, - default_pick_local_wg_size, + default_pick_gwg, + default_pick_lwg, {{out_ref, vkapi::kWrite}, {{in_ref, arg_weight, arg_bias, arg_mean, arg_var}, vkapi::kRead}}, {graph.logical_limits_ubo(out_ref), diff --git a/backends/vulkan/runtime/graph/ops/impl/BinaryOp.cpp b/backends/vulkan/runtime/graph/ops/impl/BinaryOp.cpp index 9e696a008fe..d78f4d6a3ff 100644 --- a/backends/vulkan/runtime/graph/ops/impl/BinaryOp.cpp +++ b/backends/vulkan/runtime/graph/ops/impl/BinaryOp.cpp @@ -94,8 +94,8 @@ void add_binary_op_node( graph.execute_nodes().emplace_back(new DynamicDispatchNode( graph, VK_KERNEL_FROM_STR(kernel_name), - default_pick_global_wg_size, - default_pick_local_wg_size, + default_pick_gwg, + default_pick_lwg, // Inputs and Outputs {{out, vkapi::kWrite}, {{arg1, arg2}, vkapi::kRead}}, // Shader params buffers diff --git a/backends/vulkan/runtime/graph/ops/impl/BinaryScalarOp.cpp b/backends/vulkan/runtime/graph/ops/impl/BinaryScalarOp.cpp index 68f22403d6b..57bfff019ab 100644 --- a/backends/vulkan/runtime/graph/ops/impl/BinaryScalarOp.cpp +++ b/backends/vulkan/runtime/graph/ops/impl/BinaryScalarOp.cpp @@ -103,8 +103,8 @@ void add_binary_scalar_op_node( graph.execute_nodes().emplace_back(new DynamicDispatchNode( graph, VK_KERNEL_FROM_STR(kernel_name), - default_pick_global_wg_size, - default_pick_local_wg_size, + default_pick_gwg, + default_pick_lwg, // Inputs and Outputs {{out, vkapi::kWrite}, {arg, vkapi::kRead}}, // Shader params buffers diff --git a/backends/vulkan/runtime/graph/ops/impl/ChooseQParams.cpp b/backends/vulkan/runtime/graph/ops/impl/ChooseQParams.cpp index fdacce0236c..cd1f9510bad 100644 --- a/backends/vulkan/runtime/graph/ops/impl/ChooseQParams.cpp +++ b/backends/vulkan/runtime/graph/ops/impl/ChooseQParams.cpp @@ -51,7 +51,7 @@ vkapi::ShaderInfo pick_choose_qparams_per_row_shader( return VK_KERNEL_FROM_STR(kernel_name); } -utils::uvec3 pick_choose_qparams_per_row_global_wg_size( +GlobalWorkGrid pick_choose_qparams_per_row_gwg( ComputeGraph* graph, const vkapi::ShaderInfo& shader, const std::vector& args, @@ -61,23 +61,10 @@ utils::uvec3 pick_choose_qparams_per_row_global_wg_size( const ValueRef input = args.at(1).refs.at(0); const uint32_t height = graph->size_at(-2, input); - return {1u, utils::div_up_4(height), 1u}; -} - -utils::uvec3 pick_choose_qparams_per_row_local_wg_size( - ComputeGraph* graph, - const vkapi::ShaderInfo& shader, - const utils::uvec3& global_workgroup_size, - const std::vector& args, - const std::vector& resize_args) { - (void)global_workgroup_size; - (void)args; - (void)resize_args; - - uint32_t outputs_per_wg = 1u; - uint32_t workers_per_output = 64u; - - return {workers_per_output, outputs_per_wg, 1u}; + return GlobalWorkGrid( + {1u, utils::div_up_4(height), 1u}, + kTiledWorkGrid, + LocalWorkGroup(64u, 1u, 1u)); } void add_choose_qparams_per_row_node( @@ -115,8 +102,8 @@ void add_choose_qparams_per_row_node( graph.execute_nodes().emplace_back(new DynamicDispatchNode( graph, pick_choose_qparams_per_row_shader, - pick_choose_qparams_per_row_global_wg_size, - pick_choose_qparams_per_row_local_wg_size, + pick_choose_qparams_per_row_gwg, + pick_required_lwg, // Inputs and Outputs {{{input_scales, input_zps}, vkapi::kWrite}, {input, vkapi::kRead}}, // Shader param buffers diff --git a/backends/vulkan/runtime/graph/ops/impl/Clone.cpp b/backends/vulkan/runtime/graph/ops/impl/Clone.cpp index d80d1e72cb3..df075c6244d 100644 --- a/backends/vulkan/runtime/graph/ops/impl/Clone.cpp +++ b/backends/vulkan/runtime/graph/ops/impl/Clone.cpp @@ -44,8 +44,8 @@ void add_clone_node( graph.execute_nodes().emplace_back(new DynamicDispatchNode( graph, VK_KERNEL_FROM_STR(kernel_name), - default_pick_global_wg_size, - default_pick_local_wg_size, + default_pick_gwg, + default_pick_lwg, // Inputs and Outputs {{out, vkapi::kWrite}, {in, vkapi::kRead}}, // Parameter Buffers @@ -60,7 +60,7 @@ void add_clone_node( resize_clone_node)); } -utils::uvec3 clone_image_to_buffer_global_wg_size( +GlobalWorkGrid clone_image_to_buffer_gwg( ComputeGraph* graph, const vkapi::ShaderInfo& shader, const std::vector& args, @@ -68,7 +68,7 @@ utils::uvec3 clone_image_to_buffer_global_wg_size( (void)shader; (void)resize_args; const ValueRef image = args.at(1).refs.at(0); - return graph->create_global_wg_size(image); + return graph->create_gwg(image); } void add_image_to_buffer_node( @@ -83,8 +83,8 @@ void add_image_to_buffer_node( graph.execute_nodes().emplace_back(new DynamicDispatchNode( graph, shader, - clone_image_to_buffer_global_wg_size, - default_pick_local_wg_size, + clone_image_to_buffer_gwg, + default_pick_lwg, // Input and Outputs {{buffer, vkapi::kWrite}, {image, vkapi::kRead}}, // Parameter Buffers: TextureMetadata for image, BufferMetadata for buffer @@ -111,8 +111,8 @@ void add_buffer_to_image_node( graph.execute_nodes().emplace_back(new DynamicDispatchNode( graph, shader, - default_pick_global_wg_size, - default_pick_local_wg_size, + default_pick_gwg, + default_pick_lwg, // Input and Outputs {{image, vkapi::kWrite}, {buffer, vkapi::kRead}}, // Parameter Buffers: TextureMetadata for image, BufferMetadata for buffer diff --git a/backends/vulkan/runtime/graph/ops/impl/Common.cpp b/backends/vulkan/runtime/graph/ops/impl/Common.cpp index 40ca0510383..a63b58b15b0 100644 --- a/backends/vulkan/runtime/graph/ops/impl/Common.cpp +++ b/backends/vulkan/runtime/graph/ops/impl/Common.cpp @@ -93,7 +93,7 @@ int32_t BlockConfig::outer_dim_from_packed_int(int32_t packed_int) { // Default workgroup size functions // -utils::uvec3 default_pick_global_wg_size( +GlobalWorkGrid default_pick_gwg( ComputeGraph* graph, const vkapi::ShaderInfo& shader, const std::vector& args, @@ -101,65 +101,82 @@ utils::uvec3 default_pick_global_wg_size( (void)shader; (void)resize_args; const ValueRef out = args.at(0).refs.at(0); - return graph->create_global_wg_size(out); + return graph->create_gwg(out); } -utils::uvec3 default_pick_local_wg_size( +LocalWorkGroup default_pick_lwg( ComputeGraph* graph, const vkapi::ShaderInfo& shader, - const utils::uvec3& global_workgroup_size, + const GlobalWorkGrid& gwg, const std::vector& args, const std::vector& resize_args) { (void)shader; (void)args; (void)resize_args; - return graph->create_local_wg_size(global_workgroup_size); + return graph->create_lwg(gwg); } -utils::uvec3 pick_hw_square_wg_size( +LocalWorkGroup pick_required_lwg( ComputeGraph* graph, const vkapi::ShaderInfo& shader, - const utils::uvec3& global_workgroup_size, + const GlobalWorkGrid& gwg, const std::vector& args, const std::vector& resize_args) { (void)graph; (void)shader; (void)args; (void)resize_args; - // Some inactive invocations are okay; set 6 as the threshold to use the - // a square wg size. - if (global_workgroup_size[0u] >= 6 && global_workgroup_size[1u] >= 6) { - return {8u, 8u, 1u}; - } - // If width dim is sufficiently small, then bias towards height dim to reduce - // the number of inactive invocations. - if (global_workgroup_size[0u] < 6u) { - return {4u, 16u, 1u}; - } - return {16u, 4u, 1u}; + VK_CHECK_COND( + gwg.required_lwg_size().is_valid(), + "Dispatch requires a valid local workgroup"); + return gwg.required_lwg_size(); } -utils::uvec3 pick_wc_square_wg_size( +LocalWorkGroup pick_xy_square_lwg( ComputeGraph* graph, const vkapi::ShaderInfo& shader, - const utils::uvec3& global_workgroup_size, + const GlobalWorkGrid& gwg, const std::vector& args, const std::vector& resize_args) { - (void)graph; (void)shader; (void)args; (void)resize_args; - // Some inactive invocations are okay; set 6 as the threshold to use the - // a square wg size. - if (global_workgroup_size[0u] >= 6 && global_workgroup_size[2u] >= 6) { - return {8u, 1u, 8u}; - } - // If channels dim is sufficiently small, then bias towards width dim to - // reduce the number of inactive invocations. - if (global_workgroup_size[2u] < 2u) { - return {64u, 1u, 1u}; - } - return {16u, 1u, 4u}; + LocalWorkGroup lwg( + kSquareLwg, graph->context()->adapter_ptr()->recommended_lwg_nthreads()); + lwg.fit_to_global(gwg); + return lwg; +} + +LocalWorkGroup pick_xz_square_lwg( + ComputeGraph* graph, + const vkapi::ShaderInfo& shader, + const GlobalWorkGrid& gwg, + const std::vector& args, + const std::vector& resize_args) { + (void)shader; + (void)args; + (void)resize_args; + LocalWorkGroup lwg( + LwgShape(1u, 0u, 1u), + graph->context()->adapter_ptr()->recommended_lwg_nthreads()); + lwg.fit_to_global(gwg); + return lwg; +} + +LocalWorkGroup pick_120shape_lwg( + ComputeGraph* graph, + const vkapi::ShaderInfo& shader, + const GlobalWorkGrid& gwg, + const std::vector& args, + const std::vector& resize_args) { + (void)shader; + (void)args; + (void)resize_args; + LocalWorkGroup lwg( + LwgShape(1u, 2u, 0u), + graph->context()->adapter_ptr()->recommended_lwg_nthreads()); + lwg.fit_to_global(gwg); + return lwg; } BlockConfig create_block_config_from_io_packed_dims( @@ -213,7 +230,7 @@ BlockConfig create_block_config_from_other( other.inner_dim_block_size}; } -utils::uvec3 pick_linear_global_wg_with_block_config( +GlobalWorkGrid pick_linear_gwg_with_block_config( ComputeGraph* graph, const vkapi::ShaderInfo& shader, const std::vector& args, @@ -260,10 +277,10 @@ utils::uvec3 pick_linear_global_wg_with_block_config( // Return linear workgroup size: {total_blocks, 1u, 1u} const uint32_t total_blocks = num_inner_blocks * num_outer_blocks * num_planes; - return {total_blocks, 1u, 1u}; + return graph->create_linear_gwg(total_blocks); } -utils::uvec3 pick_extents_global_wg_with_block_config( +GlobalWorkGrid pick_extents_gwg_with_block_config( ComputeGraph* graph, const vkapi::ShaderInfo& shader, const std::vector& args, @@ -328,22 +345,22 @@ utils::uvec3 pick_extents_global_wg_with_block_config( } const uint32_t z_threads = utils::safe_downcast(C_for_z * N); - return {x_threads, y_threads, z_threads}; + return GlobalWorkGrid({x_threads, y_threads, z_threads}, kTiledWorkGrid); } -utils::uvec3 pick_square_local_wg_with_block_config( +LocalWorkGroup pick_square_lwg_with_block_config( ComputeGraph* graph, const vkapi::ShaderInfo& shader, - const utils::uvec3& global_workgroup_size, + const GlobalWorkGrid& gwg, const std::vector& args, const std::vector& extra_args) { (void)graph; (void)shader; (void)args; - // Detect linear dispatch pattern: global_wg = {total_blocks, 1, 1} - if (global_workgroup_size[1u] == 1u && global_workgroup_size[2u] == 1u) { - return {64u, 1u, 1u}; + // Detect linear dispatch pattern: gwg = {total_blocks, 1, 1} + if (gwg[1u] == 1u && gwg[2u] == 1u) { + return LocalWorkGroup(64u, 1u, 1u); } // Extents dispatch: use 8x8 square on inner_dim and outer_dim axes @@ -366,7 +383,7 @@ utils::uvec3 pick_square_local_wg_with_block_config( uint32_t local_y = (inner_dim == 1 || outer_dim == 1) ? 8u : 1u; uint32_t local_z = (inner_dim == 2 || outer_dim == 2) ? 8u : 1u; - return {local_x, local_y, local_z}; + return LocalWorkGroup(local_x, local_y, local_z); } } // namespace vkcompute diff --git a/backends/vulkan/runtime/graph/ops/impl/Common.h b/backends/vulkan/runtime/graph/ops/impl/Common.h index 84cacc8e4f7..fb37f93ba4f 100644 --- a/backends/vulkan/runtime/graph/ops/impl/Common.h +++ b/backends/vulkan/runtime/graph/ops/impl/Common.h @@ -83,9 +83,9 @@ struct BlockConfig { /** * Creates a global workgroup size based on the first output tensor in the args. * This is a utility function that extracts the output tensor from - * args.at(0).refs.at(0) and calls graph->create_global_wg_size(out) on it. + * args.at(0).refs.at(0) and calls graph->create_gwg(out) on it. */ -utils::uvec3 default_pick_global_wg_size( +GlobalWorkGrid default_pick_gwg( ComputeGraph* graph, const vkapi::ShaderInfo& shader, const std::vector& args, @@ -94,12 +94,19 @@ utils::uvec3 default_pick_global_wg_size( /** * Creates a local workgroup size based on the first output tensor in the args. * This is a utility function that extracts the output tensor from - * args.at(0).refs.at(0) and calls graph->create_local_wg_size(out) on it. + * args.at(0).refs.at(0) and calls graph->create_lwg(out) on it. */ -utils::uvec3 default_pick_local_wg_size( +LocalWorkGroup default_pick_lwg( ComputeGraph* graph, const vkapi::ShaderInfo& shader, - const utils::uvec3& global_workgroup_size, + const GlobalWorkGrid& gwg, + const std::vector& args, + const std::vector& resize_args); + +LocalWorkGroup pick_required_lwg( + ComputeGraph* graph, + const vkapi::ShaderInfo& shader, + const GlobalWorkGrid& gwg, const std::vector& args, const std::vector& resize_args); @@ -114,17 +121,24 @@ utils::uvec3 default_pick_local_wg_size( * = W * K elements from the weight tensor, resulting in (W + H) * K unique * elements in total. */ -utils::uvec3 pick_hw_square_wg_size( +LocalWorkGroup pick_xy_square_lwg( + ComputeGraph* graph, + const vkapi::ShaderInfo& shader, + const GlobalWorkGrid& gwg, + const std::vector& args, + const std::vector& resize_args); + +LocalWorkGroup pick_xz_square_lwg( ComputeGraph* graph, const vkapi::ShaderInfo& shader, - const utils::uvec3& global_workgroup_size, + const GlobalWorkGrid& gwg, const std::vector& args, const std::vector& resize_args); -utils::uvec3 pick_wc_square_wg_size( +LocalWorkGroup pick_120shape_lwg( ComputeGraph* graph, const vkapi::ShaderInfo& shader, - const utils::uvec3& global_workgroup_size, + const GlobalWorkGrid& gwg, const std::vector& args, const std::vector& resize_args); @@ -202,7 +216,7 @@ BlockConfig create_block_config_from_other( * * @return Global workgroup size as {total_blocks, 1, 1} */ -utils::uvec3 pick_linear_global_wg_with_block_config( +GlobalWorkGrid pick_linear_gwg_with_block_config( ComputeGraph* graph, const vkapi::ShaderInfo& shader, const std::vector& args, @@ -224,7 +238,7 @@ utils::uvec3 pick_linear_global_wg_with_block_config( * * @return Global workgroup size as {x_threads, y_threads, z_threads} */ -utils::uvec3 pick_extents_global_wg_with_block_config( +GlobalWorkGrid pick_extents_gwg_with_block_config( ComputeGraph* graph, const vkapi::ShaderInfo& shader, const std::vector& args, @@ -238,10 +252,10 @@ utils::uvec3 pick_extents_global_wg_with_block_config( * - extra_args.at(0): Packed int32_t block configuration cast to ValueRef * (created via static_cast(BlockConfig::as_packed_int())) * - * For linear dispatch (buffer storage, global_wg = {total_blocks, 1, 1}): + * For linear dispatch (buffer storage, gwg = {total_blocks, 1, 1}): * - Returns {64, 1, 1} * - * For extents dispatch (texture storage, global_wg = {x, y, z}): + * For extents dispatch (texture storage, gwg = {x, y, z}): * - Returns an 8x8 square configuration where: * - Axes corresponding to inner_dim and outer_dim are set to 8 * - The remaining axis is set to 1 @@ -250,10 +264,10 @@ utils::uvec3 pick_extents_global_wg_with_block_config( * * @return Local workgroup size optimized for the dispatch pattern */ -utils::uvec3 pick_square_local_wg_with_block_config( +LocalWorkGroup pick_square_lwg_with_block_config( ComputeGraph* graph, const vkapi::ShaderInfo& shader, - const utils::uvec3& global_workgroup_size, + const GlobalWorkGrid& gwg, const std::vector& args, const std::vector& extra_args); diff --git a/backends/vulkan/runtime/graph/ops/impl/Concat.cpp b/backends/vulkan/runtime/graph/ops/impl/Concat.cpp index ea123de592b..0f64d0ef4bd 100644 --- a/backends/vulkan/runtime/graph/ops/impl/Concat.cpp +++ b/backends/vulkan/runtime/graph/ops/impl/Concat.cpp @@ -57,7 +57,7 @@ void resize_concat_node( graph->virtual_resize(out, new_out_sizes); } -utils::uvec3 concat_pick_global_wg_size( +GlobalWorkGrid concat_pick_gwg( ComputeGraph* graph, const vkapi::ShaderInfo& shader, const std::vector& args, @@ -91,7 +91,7 @@ utils::uvec3 concat_pick_global_wg_size( } if (graph->is_buffer_storage(out)) { - return {total_input_numel, 1, 1}; + return graph->create_linear_gwg(total_input_numel); } // The texture implementation is similar, except each invocation writes out 4 @@ -124,7 +124,7 @@ utils::uvec3 concat_pick_global_wg_size( const uint32_t inp_volume_texel_numel = utils::multiply_integers(inp_volume_texel_sizes); - return {inp_volume_texel_numel, 1, 1}; + return graph->create_linear_gwg(inp_volume_texel_numel); } void add_concat_node( @@ -166,8 +166,8 @@ void add_concat_node( graph.execute_nodes().emplace_back(new DispatchNode( graph, VK_KERNEL_FROM_STR(kernel_name), - {1, 1, 1}, - {1, 1, 1}, + GlobalWorkGrid({1u, 1u, 1u}, kExplicitWorkGrid), + LocalWorkGroup(1u, 1u, 1u), // Inputs and Outputs {{concat_offset, vkapi::kWrite}}, // Parameter buffers @@ -230,8 +230,8 @@ void add_concat_node( graph.execute_nodes().emplace_back(new DynamicDispatchNode( graph, VK_KERNEL_FROM_STR(kernel_name), - concat_pick_global_wg_size, - default_pick_local_wg_size, + concat_pick_gwg, + default_pick_lwg, // Inputs and Outputs {{out, vkapi::kReadWrite}, {batch_inputs, vkapi::kRead}, @@ -271,8 +271,8 @@ void add_concat_node( graph.execute_nodes().emplace_back(new DispatchNode( graph, VK_KERNEL_FROM_STR(kernel_name), - {1u, 1u, 1u}, - {1u, 1u, 1u}, + GlobalWorkGrid({1u, 1u, 1u}, kExplicitWorkGrid), + LocalWorkGroup(1u, 1u, 1u), // Inputs and Outputs {{concat_offset, vkapi::kReadWrite}}, // Parameter buffers diff --git a/backends/vulkan/runtime/graph/ops/impl/Conv1dDW.cpp b/backends/vulkan/runtime/graph/ops/impl/Conv1dDW.cpp index 88d421e6994..ba59090c124 100644 --- a/backends/vulkan/runtime/graph/ops/impl/Conv1dDW.cpp +++ b/backends/vulkan/runtime/graph/ops/impl/Conv1dDW.cpp @@ -55,7 +55,7 @@ struct Conv1dDWClampParams final { float output_max; }; -utils::uvec3 pick_conv1d_dw_global_wg_size( +GlobalWorkGrid pick_conv1d_dw_gwg( ComputeGraph* graph, const vkapi::ShaderInfo& shader, const std::vector& args, @@ -70,7 +70,7 @@ utils::uvec3 pick_conv1d_dw_global_wg_size( const uint32_t N = graph->dim_of(out) >= 3 ? graph->size_at(-3, out) : 1; - return {utils::div_up_4(C), L_out, N}; + return GlobalWorkGrid({utils::div_up_4(C), L_out, N}, kTiledWorkGrid); } void add_conv1d_dw_node( @@ -130,8 +130,8 @@ void add_conv1d_dw_node( graph.execute_nodes().emplace_back(new DynamicDispatchNode( graph, VK_KERNEL_FROM_STR(kernel_name), - pick_conv1d_dw_global_wg_size, - default_pick_local_wg_size, + pick_conv1d_dw_gwg, + pick_120shape_lwg, // Inputs and Outputs {{out, vkapi::kWrite}, {read_inputs, vkapi::kRead}}, // Shader params buffers diff --git a/backends/vulkan/runtime/graph/ops/impl/Conv1dPW.cpp b/backends/vulkan/runtime/graph/ops/impl/Conv1dPW.cpp index 90dada6b58e..de2fcdbb806 100644 --- a/backends/vulkan/runtime/graph/ops/impl/Conv1dPW.cpp +++ b/backends/vulkan/runtime/graph/ops/impl/Conv1dPW.cpp @@ -54,10 +54,11 @@ static ValueRef prepack_conv1d_pw_weight( weight_storage, utils::kWidthPacked); - utils::uvec3 global_wg_size = { - utils::safe_downcast(N4), - utils::safe_downcast(K4), - 1u}; + const GlobalWorkGrid gwg( + {utils::safe_downcast(N4), + utils::safe_downcast(K4), + 1u}, + kTiledWorkGrid); struct PackParams { int32_t N; @@ -76,8 +77,8 @@ static ValueRef prepack_conv1d_pw_weight( graph.prepack_nodes().emplace_back(new PrepackNode( graph, VK_KERNEL_FROM_STR(kernel_name), - global_wg_size, - graph.create_local_wg_size(global_wg_size), + gwg, + graph.create_lwg(gwg), weight_data, packed_weight, {}, @@ -125,7 +126,7 @@ vkapi::ShaderInfo pick_conv1d_pw_shader( return VK_KERNEL_FROM_STR(kernel_name); } -utils::uvec3 pick_conv1d_pw_global_wg_size( +GlobalWorkGrid pick_conv1d_pw_gwg( ComputeGraph* graph, const vkapi::ShaderInfo& shader, const std::vector& args, @@ -141,7 +142,9 @@ utils::uvec3 pick_conv1d_pw_global_wg_size( graph->dim_of(out) >= 3 ? graph->size_at(-3, out) : 1; // X=OC4 (div_up_4(C_out)), Y=L/tile_m, Z=N_batch - return {utils::div_up_4(C_out), utils::div_up(L, kTileM), N_batch}; + return GlobalWorkGrid( + {utils::div_up_4(C_out), utils::div_up(L, kTileM), N_batch}, + kTiledWorkGrid); } void add_conv1d_pw_node( @@ -189,8 +192,8 @@ void add_conv1d_pw_node( graph.execute_nodes().emplace_back(new DynamicDispatchNode( graph, pick_conv1d_pw_shader, - pick_conv1d_pw_global_wg_size, - pick_hw_square_wg_size, + pick_conv1d_pw_gwg, + pick_xy_square_lwg, // Inputs and Outputs {{out, vkapi::kWrite}, {read_inputs, vkapi::kRead}}, // Shader params buffers diff --git a/backends/vulkan/runtime/graph/ops/impl/Conv2dDW.cpp b/backends/vulkan/runtime/graph/ops/impl/Conv2dDW.cpp index a9d8483b2e2..3cfcf1eebce 100644 --- a/backends/vulkan/runtime/graph/ops/impl/Conv2dDW.cpp +++ b/backends/vulkan/runtime/graph/ops/impl/Conv2dDW.cpp @@ -47,8 +47,8 @@ ValueRef prepack_dw_weights(ComputeGraph& graph, const ValueRef vref) { graph.prepack_nodes().emplace_back(new PrepackNode( graph, VK_KERNEL_FROM_STR(kernel_name), - graph.create_global_wg_size(v), - graph.create_local_wg_size(v), + graph.create_gwg(v), + graph.create_lwg(v), vref, v, {}, @@ -102,7 +102,7 @@ std::string pick_conv2d_dw_shader( // Workgroup size // -utils::uvec3 conv2d_dw_global_wg_size( +GlobalWorkGrid conv2d_dw_gwg( ComputeGraph* graph, const vkapi::ShaderInfo& shader, const std::vector& args, @@ -116,11 +116,13 @@ utils::uvec3 conv2d_dw_global_wg_size( if (uses_output_tile) { const bool is_sned = shader.kernel_name.find("_sned") != std::string::npos; - const utils::uvec3 image_extents = graph->create_global_wg_size(out); + const utils::uvec3 image_extents = graph->create_gwg(out).extents(); if (is_sned) { // sned output_tile shaders: no batch division, just flatten W*H - return {image_extents[0] * image_extents[1], image_extents[2], 1}; + return GlobalWorkGrid( + {image_extents[0] * image_extents[1], image_extents[2], 1u}, + kTiledWorkGrid); } // stride==dilation output_tile shaders: apply batch division @@ -133,26 +135,28 @@ utils::uvec3 conv2d_dw_global_wg_size( uint32_t scaled_x = utils::div_up(image_extents[0], batch_x); uint32_t scaled_y = utils::div_up(image_extents[1], batch_y); - return {scaled_x * scaled_y, image_extents[2], 1}; + return GlobalWorkGrid( + {scaled_x * scaled_y, image_extents[2], 1u}, kTiledWorkGrid); } // Base conv2d_dw shader: fully linearized dispatch - const utils::uvec3 base_extents = graph->create_global_wg_size(out); - return {base_extents[0] * base_extents[1] * base_extents[2], 1, 1}; + const utils::uvec3 base_extents = graph->create_gwg(out).extents(); + return graph->create_linear_gwg( + base_extents[0] * base_extents[1] * base_extents[2]); } -utils::uvec3 conv2d_dw_local_wg_size( +LocalWorkGroup conv2d_dw_lwg( ComputeGraph* graph, const vkapi::ShaderInfo& shader, - const utils::uvec3& global_workgroup_size, + const GlobalWorkGrid& gwg, const std::vector& args, const std::vector& resize_args) { (void)graph; (void)shader; - (void)global_workgroup_size; + (void)gwg; (void)args; (void)resize_args; - return {64, 1, 1}; + return LocalWorkGroup(64u, 1u, 1u); } // @@ -232,8 +236,8 @@ void add_conv2d_dw_node( graph.execute_nodes().emplace_back(new DynamicDispatchNode( graph, shader, - conv2d_dw_global_wg_size, - conv2d_dw_local_wg_size, + conv2d_dw_gwg, + conv2d_dw_lwg, // Inputs and Outputs {{out, vkapi::kWrite}, {{in, arg_weight, arg_bias}, vkapi::kRead}}, // Shader params buffers diff --git a/backends/vulkan/runtime/graph/ops/impl/Conv2dGemm.cpp b/backends/vulkan/runtime/graph/ops/impl/Conv2dGemm.cpp index 923c7167876..6285dc92d4a 100644 --- a/backends/vulkan/runtime/graph/ops/impl/Conv2dGemm.cpp +++ b/backends/vulkan/runtime/graph/ops/impl/Conv2dGemm.cpp @@ -98,10 +98,11 @@ ValueRef prepack_conv2d_gemm_weight( weight_storage, utils::kWidthPacked); - const utils::uvec3 global_wg_size = { - utils::safe_downcast(N4), - utils::safe_downcast(K4), - 1u}; + const GlobalWorkGrid gwg( + {utils::safe_downcast(N4), + utils::safe_downcast(K4), + 1u}, + kTiledWorkGrid); // Push constants must be uploaded in <= 16-byte (one ivec4) chunks; the // shader's Block reads them back as dims0 / dims1. Layout must match @@ -125,8 +126,8 @@ ValueRef prepack_conv2d_gemm_weight( graph.prepack_nodes().emplace_back(new PrepackNode( graph, VK_KERNEL_FROM_STR(kernel_name), - global_wg_size, - graph.create_local_wg_size(global_wg_size), + gwg, + graph.create_lwg(gwg), weight_data, packed_weight, {}, @@ -164,7 +165,7 @@ vkapi::ShaderInfo pick_conv2d_gemm_shader( // handles): both are build-time constants, so packing the ints directly into // the slots avoids materializing graph Values for them. Read with static_cast, // never get_int. -utils::uvec3 pick_conv2d_gemm_global_wg_size( +GlobalWorkGrid pick_conv2d_gemm_gwg( ComputeGraph* graph, const vkapi::ShaderInfo& shader, const std::vector& args, @@ -185,14 +186,14 @@ utils::uvec3 pick_conv2d_gemm_global_wg_size( // skip tracks the dynamic shape. (The num_tiles count is still fixed at build // time; this only zeroes the work of tiles that fall off the live region.) if (oh_offset >= static_cast(H_out)) { - return {0u, 0u, 0u}; + return GlobalWorkGrid({0u, 0u, 0u}, kTiledWorkGrid); } // Every live tile dispatches oh_tile output-height rows (oh_tile * W per-tile // M); trailing threads past the real H_out no-op in the shader. const uint32_t M_tile = static_cast(oh_tile) * W; const uint32_t N4 = utils::div_up_4(C_out); // TILE_N4=1, TILE_M=4 - return {N4, utils::div_up(M_tile, 4u), 1}; + return GlobalWorkGrid({N4, utils::div_up(M_tile, 4u), 1u}, kTiledWorkGrid); } // Recompute the conv output sizes from the current input shape and resize the @@ -292,8 +293,8 @@ void add_conv2d_gemm_node( graph.execute_nodes().emplace_back(new DynamicDispatchNode( graph, pick_conv2d_gemm_shader, - pick_conv2d_gemm_global_wg_size, - pick_hw_square_wg_size, + pick_conv2d_gemm_gwg, + pick_xy_square_lwg, // Inputs and Outputs {{out, vkapi::kWrite}, {{im2col_in, packed_weight, packed_bias}, vkapi::kRead}}, diff --git a/backends/vulkan/runtime/graph/ops/impl/Conv2dIm2Col.cpp b/backends/vulkan/runtime/graph/ops/impl/Conv2dIm2Col.cpp index 062c85c9edd..e1bbded799e 100644 --- a/backends/vulkan/runtime/graph/ops/impl/Conv2dIm2Col.cpp +++ b/backends/vulkan/runtime/graph/ops/impl/Conv2dIm2Col.cpp @@ -49,7 +49,7 @@ Im2colExtents im2col_extents_of(ComputeGraph* graph, const ValueRef im2col) { return {m, k_total / 4u}; } -utils::uvec3 pick_conv2d_im2col_global_wg_size( +GlobalWorkGrid pick_conv2d_im2col_gwg( ComputeGraph* graph, const vkapi::ShaderInfo& shader, const std::vector& args, @@ -59,23 +59,23 @@ utils::uvec3 pick_conv2d_im2col_global_wg_size( const ValueRef im2col_out = args.at(0).refs.at(0); const Im2colExtents ext = im2col_extents_of(graph, im2col_out); // Global wg: one thread per (k4, m) vec4 in the output. - return {ext.k4_total, ext.m, 1u}; + return GlobalWorkGrid({ext.k4_total, ext.m, 1u}, kTiledWorkGrid); } -utils::uvec3 pick_conv2d_im2col_local_wg_size( +LocalWorkGroup pick_conv2d_im2col_lwg( ComputeGraph* graph, const vkapi::ShaderInfo& shader, - const utils::uvec3& global_workgroup_size, + const GlobalWorkGrid& gwg, const std::vector& args, const std::vector& resize_args) { (void)graph; (void)shader; - (void)global_workgroup_size; + (void)gwg; (void)args; (void)resize_args; // Fixed {16, 4, 1} mirrors the original static dispatch — one thread per // (k4, m) vec4 with 16 K-tiles × 4 M positions per workgroup. - return {16u, 4u, 1u}; + return LocalWorkGroup(16u, 4u, 1u); } // Recompute the im2col scratch extents from the current input shape and @@ -197,8 +197,8 @@ void add_conv2d_im2col_node( graph.execute_nodes().emplace_back(new DynamicDispatchNode( graph, shader, - pick_conv2d_im2col_global_wg_size, - pick_conv2d_im2col_local_wg_size, + pick_conv2d_im2col_gwg, + pick_conv2d_im2col_lwg, // Inputs and Outputs {{im2col_out, vkapi::kWrite}, {in, vkapi::kRead}}, // UBOs diff --git a/backends/vulkan/runtime/graph/ops/impl/Conv2dPW.cpp b/backends/vulkan/runtime/graph/ops/impl/Conv2dPW.cpp index 43657e7f4ec..7c5f9f0bc0a 100644 --- a/backends/vulkan/runtime/graph/ops/impl/Conv2dPW.cpp +++ b/backends/vulkan/runtime/graph/ops/impl/Conv2dPW.cpp @@ -66,7 +66,7 @@ vkapi::ShaderInfo pick_conv2d_pw_tiled_shader( return VK_KERNEL_FROM_STR(kernel_name); } -utils::uvec3 pick_conv2d_pw_tiled_global_wg_size( +GlobalWorkGrid pick_conv2d_pw_tiled_gwg( ComputeGraph* graph, const vkapi::ShaderInfo& shader, const std::vector& args, @@ -80,7 +80,7 @@ utils::uvec3 pick_conv2d_pw_tiled_global_wg_size( uint32_t M = H * W; uint32_t N4 = utils::div_up_4(C_out); // TILE_N4=1, TILE_M=4 - return {N4, utils::div_up(M, 4u), 1}; + return GlobalWorkGrid({N4, utils::div_up(M, 4u), 1u}, kTiledWorkGrid); } // @@ -119,10 +119,11 @@ ValueRef prepack_conv2d_pw_weight( weight_storage, utils::kWidthPacked); - utils::uvec3 global_wg_size = { - utils::safe_downcast(N4), - utils::safe_downcast(K4), - 1u}; + const GlobalWorkGrid gwg( + {utils::safe_downcast(N4), + utils::safe_downcast(K4), + 1u}, + kTiledWorkGrid); PackParams pack_params{ utils::safe_downcast(N), utils::safe_downcast(K), 1, 1}; @@ -135,8 +136,8 @@ ValueRef prepack_conv2d_pw_weight( graph.prepack_nodes().emplace_back(new PrepackNode( graph, VK_KERNEL_FROM_STR(pack_kernel_name), - global_wg_size, - graph.create_local_wg_size(global_wg_size), + gwg, + graph.create_lwg(gwg), weight_data, packed_weight, {}, @@ -188,8 +189,8 @@ void add_conv2d_pw_tiled_node( graph.execute_nodes().emplace_back(new DynamicDispatchNode( graph, pick_conv2d_pw_tiled_shader, - pick_conv2d_pw_tiled_global_wg_size, - pick_hw_square_wg_size, + pick_conv2d_pw_tiled_gwg, + pick_xy_square_lwg, // Inputs and Outputs {{out, vkapi::kWrite}, {{in, packed_weight, packed_bias}, vkapi::kRead}}, // Shader params buffers diff --git a/backends/vulkan/runtime/graph/ops/impl/Convolution.cpp b/backends/vulkan/runtime/graph/ops/impl/Convolution.cpp index 5df73556ab6..4e8ff50fdac 100644 --- a/backends/vulkan/runtime/graph/ops/impl/Convolution.cpp +++ b/backends/vulkan/runtime/graph/ops/impl/Convolution.cpp @@ -126,8 +126,8 @@ ValueRef prepack_biases( graph.prepack_nodes().emplace_back(new PrepackNode( graph, shader, - graph.create_global_wg_size(v), - graph.create_local_wg_size(v), + graph.create_gwg(v), + graph.create_lwg(v), vref, v, param_buffers, @@ -224,8 +224,8 @@ ValueRef prepack_weights( graph.prepack_nodes().emplace_back(new PrepackNode( graph, shader, - graph.create_global_wg_size(v), - graph.create_local_wg_size(v), + graph.create_gwg(v), + graph.create_lwg(v), vref, v, {}, @@ -332,7 +332,7 @@ bool should_use_conv2d_im2col( return graph.device_is_mali() || c_out >= kIm2colMinCOut; } -utils::uvec3 create_conv2d_global_wg_size( +GlobalWorkGrid create_conv2d_gwg( ComputeGraph& graph, const Conv2dMethod method, const ValueRef out, @@ -340,17 +340,18 @@ utils::uvec3 create_conv2d_global_wg_size( const bool stride_equals_dilation) { if (method == Conv2dMethod::Pointwise) { const utils::uvec3 image_extents = graph.logical_limits_of(out); - return { - utils::div_up(image_extents[0u], 1u), - utils::div_up(image_extents[1u], 4u), - image_extents[2u]}; + return GlobalWorkGrid( + {utils::div_up(image_extents[0u], 1u), + utils::div_up(image_extents[1u], 4u), + image_extents[2u]}, + kTiledWorkGrid); } else { - return graph.create_global_wg_size(out); + return graph.create_gwg(out); } } // Custom global workgroup size function for conv2d -utils::uvec3 conv2d_global_wg_size( +GlobalWorkGrid conv2d_gwg( ComputeGraph* graph, const vkapi::ShaderInfo& shader, const std::vector& args, @@ -380,25 +381,26 @@ utils::uvec3 conv2d_global_wg_size( bool stride_equals_dilation = shader.kernel_name.find("_sned") == std::string::npos; - utils::uvec3 wg_size = create_conv2d_global_wg_size( + const GlobalWorkGrid wg_size = create_conv2d_gwg( *graph, method, out, weight_data, stride_equals_dilation); if (method == Conv2dMethod::Pointwise) { - wg_size = {wg_size[0] * wg_size[1], wg_size[2], 1}; + utils::uvec3 pointwise_wg_size = {wg_size[0] * wg_size[1], wg_size[2], 1u}; if (shader.kernel_name.find("s1p0") != std::string::npos) { - wg_size[0] *= 4; + pointwise_wg_size[0] *= 4; } + return GlobalWorkGrid(pointwise_wg_size, kTiledWorkGrid); } return wg_size; } // Custom local workgroup size function for conv2d -utils::uvec3 conv2d_local_wg_size( +LocalWorkGroup conv2d_lwg( ComputeGraph* graph, const vkapi::ShaderInfo& shader, - const utils::uvec3& global_workgroup_size, + const GlobalWorkGrid& gwg, const std::vector& args, const std::vector& resize_args) { (void)args; @@ -415,22 +417,22 @@ utils::uvec3 conv2d_local_wg_size( } if (method == Conv2dMethod::Pointwise) { - uint32_t local_wg_size_y = 1; - if (global_workgroup_size[1] % 8 == 0) { - local_wg_size_y = 8; - } else if (global_workgroup_size[1] % 4 == 0) { - local_wg_size_y = 4; - } else if (global_workgroup_size[1] % 2 == 0) { - local_wg_size_y = 2; + uint32_t lwg_y = 1; + if (gwg[1] % 8 == 0) { + lwg_y = 8; + } else if (gwg[1] % 4 == 0) { + lwg_y = 4; + } else if (gwg[1] % 2 == 0) { + lwg_y = 2; } - return {64 / local_wg_size_y, local_wg_size_y, 1}; + return LocalWorkGroup(64u / lwg_y, lwg_y, 1u); } else { - return graph->create_local_wg_size(global_workgroup_size); + return graph->create_lwg(gwg); } } // Custom global workgroup size function for conv1d -utils::uvec3 conv1d_global_wg_size( +GlobalWorkGrid conv1d_gwg( ComputeGraph* graph, const vkapi::ShaderInfo& shader, const std::vector& args, @@ -439,12 +441,14 @@ utils::uvec3 conv1d_global_wg_size( (void)resize_args; const ValueRef out = args.at(0).refs.at(0); - return {// out length - graph->size_at(-1, out), - // out channels - static_cast(graph->size_at(-2, out)), - // out batches - utils::div_up_4(graph->size_at(-3, out))}; + return GlobalWorkGrid( + {// out length + graph->size_at(-1, out), + // out channels + static_cast(graph->size_at(-2, out)), + // out batches + utils::div_up_4(graph->size_at(-3, out))}, + kTiledWorkGrid); } void add_conv2d_node( @@ -593,8 +597,8 @@ void add_conv2d_node( graph.execute_nodes().emplace_back(new DynamicDispatchNode( graph, shader, - conv2d_global_wg_size, - conv2d_local_wg_size, + conv2d_gwg, + conv2d_lwg, // Inputs and Outputs {{out, vkapi::kWrite}, {{in, arg_weight, arg_bias}, vkapi::kRead}}, // Shader params buffers @@ -684,8 +688,8 @@ void add_conv1d_node( graph.execute_nodes().emplace_back(new DynamicDispatchNode( graph, VK_KERNEL_FROM_STR(kernel_name), - conv1d_global_wg_size, - default_pick_local_wg_size, + conv1d_gwg, + default_pick_lwg, // Inputs and Outputs {{out, vkapi::kWrite}, {{in, arg_weight, arg_bias}, vkapi::kRead}}, // Shader params buffers diff --git a/backends/vulkan/runtime/graph/ops/impl/Embedding.cpp b/backends/vulkan/runtime/graph/ops/impl/Embedding.cpp index c8c5ae847b0..61ba9349b45 100644 --- a/backends/vulkan/runtime/graph/ops/impl/Embedding.cpp +++ b/backends/vulkan/runtime/graph/ops/impl/Embedding.cpp @@ -69,8 +69,8 @@ void add_embedding_node( graph.execute_nodes().emplace_back(new DynamicDispatchNode( graph, VK_KERNEL_FROM_STR(kernel_name), - default_pick_global_wg_size, - default_pick_local_wg_size, + default_pick_gwg, + default_pick_lwg, // Inputs and Outputs {{out, vkapi::kWrite}, {{indices, weight}, vkapi::kRead}}, // Shader params buffers @@ -99,8 +99,8 @@ void add_embedding_legacy_node( graph.execute_nodes().emplace_back(new DynamicDispatchNode( graph, VK_KERNEL_FROM_STR(kernel_name), - default_pick_global_wg_size, - default_pick_local_wg_size, + default_pick_gwg, + default_pick_lwg, {{out, vkapi::kWrite}, {{in, weight}, vkapi::kRead}}, { graph.sizes_ubo(out), diff --git a/backends/vulkan/runtime/graph/ops/impl/EmbeddingQ4gsw.cpp b/backends/vulkan/runtime/graph/ops/impl/EmbeddingQ4gsw.cpp index 46a65c9284c..ad4786d1c60 100644 --- a/backends/vulkan/runtime/graph/ops/impl/EmbeddingQ4gsw.cpp +++ b/backends/vulkan/runtime/graph/ops/impl/EmbeddingQ4gsw.cpp @@ -22,7 +22,7 @@ namespace vkcompute { -utils::uvec3 pick_embedding_q4gsw_global_wg_size( +GlobalWorkGrid pick_embedding_q4gsw_gwg( ComputeGraph* graph, const vkapi::ShaderInfo& shader, const std::vector& args, @@ -38,7 +38,7 @@ utils::uvec3 pick_embedding_q4gsw_global_wg_size( depth *= static_cast(sizes[i]); } - return {blocks_per_row, height, depth}; + return GlobalWorkGrid({blocks_per_row, height, depth}, kTiledWorkGrid); } void resize_embedding_q4gsw_node( @@ -99,8 +99,8 @@ void add_embedding_q4gsw_node( graph.execute_nodes().emplace_back(new DynamicDispatchNode( graph, VK_KERNEL_FROM_STR(kernel_name), - pick_embedding_q4gsw_global_wg_size, - default_pick_local_wg_size, + pick_embedding_q4gsw_gwg, + default_pick_lwg, {{out, vkapi::kWrite}, {{indices, weight, weight_scales}, vkapi::kRead}}, param_ubos, push_constants, diff --git a/backends/vulkan/runtime/graph/ops/impl/Expand.cpp b/backends/vulkan/runtime/graph/ops/impl/Expand.cpp index 6308d333eaf..5629aa6d2f9 100644 --- a/backends/vulkan/runtime/graph/ops/impl/Expand.cpp +++ b/backends/vulkan/runtime/graph/ops/impl/Expand.cpp @@ -66,8 +66,8 @@ void add_expand_node( graph.execute_nodes().emplace_back(new DynamicDispatchNode( graph, VK_KERNEL_FROM_STR(kernel_name), - default_pick_global_wg_size, - default_pick_local_wg_size, + default_pick_gwg, + default_pick_lwg, {{out, vkapi::kWrite}, {in, vkapi::kRead}}, // Parameter buffers param_buffers, diff --git a/backends/vulkan/runtime/graph/ops/impl/Flip.cpp b/backends/vulkan/runtime/graph/ops/impl/Flip.cpp index 52288734704..0db7b7a8979 100644 --- a/backends/vulkan/runtime/graph/ops/impl/Flip.cpp +++ b/backends/vulkan/runtime/graph/ops/impl/Flip.cpp @@ -17,7 +17,7 @@ namespace vkcompute { // Custom global workgroup size function for flip -utils::uvec3 flip_global_wg_size( +GlobalWorkGrid flip_gwg( ComputeGraph* graph, const vkapi::ShaderInfo& shader, const std::vector& args, @@ -25,7 +25,7 @@ utils::uvec3 flip_global_wg_size( (void)shader; (void)resize_args; const ValueRef out = args.at(0).refs.at(0); - return graph->create_global_wg_size(out); + return graph->create_gwg(out); } void check_flip_args( @@ -75,8 +75,8 @@ void add_flip_node( graph.execute_nodes().emplace_back(new DynamicDispatchNode( graph, VK_KERNEL_FROM_STR(kernel_name), - flip_global_wg_size, - default_pick_local_wg_size, + flip_gwg, + default_pick_lwg, // Inputs and Outputs { {out, vkapi::kWrite}, diff --git a/backends/vulkan/runtime/graph/ops/impl/Full.cpp b/backends/vulkan/runtime/graph/ops/impl/Full.cpp index 5458fdce7df..e5e577da7ef 100644 --- a/backends/vulkan/runtime/graph/ops/impl/Full.cpp +++ b/backends/vulkan/runtime/graph/ops/impl/Full.cpp @@ -47,8 +47,8 @@ void add_full_node( graph.execute_nodes().emplace_back(new DynamicDispatchNode( graph, VK_KERNEL_FROM_STR(kernel_name), - default_pick_global_wg_size, - default_pick_local_wg_size, + default_pick_gwg, + default_pick_lwg, // Inputs and Outputs {{out, vkapi::kWrite}}, // Shader params buffers diff --git a/backends/vulkan/runtime/graph/ops/impl/FusedCe.cpp b/backends/vulkan/runtime/graph/ops/impl/FusedCe.cpp index f66e0e4fc46..25adf34b7a0 100644 --- a/backends/vulkan/runtime/graph/ops/impl/FusedCe.cpp +++ b/backends/vulkan/runtime/graph/ops/impl/FusedCe.cpp @@ -20,7 +20,7 @@ namespace vkcompute { using namespace utils; -utils::uvec3 fused_ce_global_wg_size( +GlobalWorkGrid fused_ce_gwg( ComputeGraph* graph, const vkapi::ShaderInfo& shader, const std::vector& args, @@ -28,48 +28,23 @@ utils::uvec3 fused_ce_global_wg_size( (void)shader; (void)resize_args; const ValueRef loss_partial = args.at(0).refs.at(1); - return { - 1u, utils::safe_downcast(graph->numel_of(loss_partial)), 1u}; + return GlobalWorkGrid( + {1u, utils::safe_downcast(graph->numel_of(loss_partial)), 1u}, + kTiledWorkGrid, + LocalWorkGroup(64u, 1u, 1u)); } -utils::uvec3 fused_ce_local_wg_size( +GlobalWorkGrid fused_ce_sum_gwg( ComputeGraph* graph, const vkapi::ShaderInfo& shader, - const utils::uvec3& global_workgroup_size, const std::vector& args, const std::vector& resize_args) { (void)graph; (void)shader; - (void)global_workgroup_size; (void)args; (void)resize_args; - return {64u, 1u, 1u}; -} - -utils::uvec3 fused_ce_sum_global_wg_size( - ComputeGraph* graph, - const vkapi::ShaderInfo& shader, - const std::vector& args, - const std::vector& resize_args) { - (void)graph; - (void)shader; - (void)args; - (void)resize_args; - return {1u, 1u, 1u}; -} - -utils::uvec3 fused_ce_sum_local_wg_size( - ComputeGraph* graph, - const vkapi::ShaderInfo& shader, - const utils::uvec3& global_workgroup_size, - const std::vector& args, - const std::vector& resize_args) { - (void)graph; - (void)shader; - (void)global_workgroup_size; - (void)args; - (void)resize_args; - return {64u, 1u, 1u}; + return GlobalWorkGrid( + {1u, 1u, 1u}, kTiledWorkGrid, LocalWorkGroup(64u, 1u, 1u)); } void resize_fused_ce_node( @@ -146,8 +121,8 @@ void fused_ce(ComputeGraph& graph, const std::vector& args) { graph.execute_nodes().emplace_back(new DynamicDispatchNode( graph, VK_KERNEL_FROM_STR(kernel_name), - fused_ce_global_wg_size, - fused_ce_local_wg_size, + fused_ce_gwg, + pick_required_lwg, // Inputs and Outputs {{{dlogits, loss_partial}, vkapi::kWrite}, {{logits, labels}, vkapi::kRead}}, @@ -174,8 +149,8 @@ void fused_ce(ComputeGraph& graph, const std::vector& args) { graph.execute_nodes().emplace_back(new DynamicDispatchNode( graph, VK_KERNEL_FROM_STR(sum_kernel_name), - fused_ce_sum_global_wg_size, - fused_ce_sum_local_wg_size, + fused_ce_sum_gwg, + pick_required_lwg, // Inputs and Outputs {{loss, vkapi::kWrite}, {loss_partial, vkapi::kRead}}, // Shader params buffers diff --git a/backends/vulkan/runtime/graph/ops/impl/Gather.cpp b/backends/vulkan/runtime/graph/ops/impl/Gather.cpp index 1d8173619bd..cea4a2d5959 100644 --- a/backends/vulkan/runtime/graph/ops/impl/Gather.cpp +++ b/backends/vulkan/runtime/graph/ops/impl/Gather.cpp @@ -52,8 +52,8 @@ void add_gather_node( graph.execute_nodes().emplace_back(new DynamicDispatchNode( graph, VK_KERNEL_FROM_STR(kernel_name), - default_pick_global_wg_size, - default_pick_local_wg_size, + default_pick_gwg, + default_pick_lwg, // Inputs and Outputs {{out, vkapi::kWrite}, {{input, index}, vkapi::kRead}}, // Shader params buffers diff --git a/backends/vulkan/runtime/graph/ops/impl/GemmCommon.cpp b/backends/vulkan/runtime/graph/ops/impl/GemmCommon.cpp index ae32581d5a3..b763dcaeb38 100644 --- a/backends/vulkan/runtime/graph/ops/impl/GemmCommon.cpp +++ b/backends/vulkan/runtime/graph/ops/impl/GemmCommon.cpp @@ -61,10 +61,11 @@ ValueRef prepack_fp_linear_weight( weight_storage, utils::kWidthPacked); - utils::uvec3 global_wg_size = { - utils::safe_downcast(N4), - utils::safe_downcast(K4), - utils::safe_downcast(B)}; + const GlobalWorkGrid gwg( + {utils::safe_downcast(N4), + utils::safe_downcast(K4), + utils::safe_downcast(B)}, + kTiledWorkGrid); struct PackParams { int32_t N; @@ -86,8 +87,8 @@ ValueRef prepack_fp_linear_weight( graph.prepack_nodes().emplace_back(new PrepackNode( graph, VK_KERNEL_FROM_STR(kernel_name), - global_wg_size, - graph.create_local_wg_size(global_wg_size), + gwg, + graph.create_lwg(gwg), weight_data, packed_weight, {}, diff --git a/backends/vulkan/runtime/graph/ops/impl/GemmCoopmat.cpp b/backends/vulkan/runtime/graph/ops/impl/GemmCoopmat.cpp index d5aff62ac62..6b62fa6a9ba 100644 --- a/backends/vulkan/runtime/graph/ops/impl/GemmCoopmat.cpp +++ b/backends/vulkan/runtime/graph/ops/impl/GemmCoopmat.cpp @@ -31,7 +31,7 @@ static vkapi::ShaderInfo pick_linear_coopmat_shader( return VK_KERNEL_FROM_STR(kernel_name); } -static utils::uvec3 pick_linear_coopmat_global_wg_size( +static GlobalWorkGrid pick_linear_coopmat_gwg( ComputeGraph* graph, const vkapi::ShaderInfo& shader, const std::vector& args, @@ -49,25 +49,14 @@ static utils::uvec3 pick_linear_coopmat_global_wg_size( // to launch exactly num_tiles_n x num_tiles_m workgroups. // // The framework computes the group count as - // group_count = div_up(global_wg_size, local_wg_size) - // (see Context.cpp + Command.cpp). With local_wg = (kCoopmatInvocations, + // group_count = div_up(gwg, lwg) + // (see Context.cpp + Command.cpp). With lwg = (kCoopmatInvocations, // 1, 1), multiplying num_tiles_n by kCoopmatInvocations cancels the // div, yielding group_count.x = num_tiles_n. - return {num_tiles_n * kCoopmatInvocations, num_tiles_m, 1}; -} - -static utils::uvec3 pick_linear_coopmat_local_wg_size( - ComputeGraph* graph, - const vkapi::ShaderInfo& shader, - const utils::uvec3& global_workgroup_size, - const std::vector& args, - const std::vector& resize_args) { - (void)graph; - (void)shader; - (void)global_workgroup_size; - (void)args; - (void)resize_args; - return {kCoopmatInvocations, 1, 1}; + return GlobalWorkGrid( + {num_tiles_n * kCoopmatInvocations, num_tiles_m, 1u}, + kTiledWorkGrid, + LocalWorkGroup(kCoopmatInvocations, 1u, 1u)); } void add_linear_coopmat_node( @@ -104,8 +93,8 @@ void add_linear_coopmat_node( graph.execute_nodes().emplace_back(new DynamicDispatchNode( graph, pick_linear_coopmat_shader, - pick_linear_coopmat_global_wg_size, - pick_linear_coopmat_local_wg_size, + pick_linear_coopmat_gwg, + pick_required_lwg, // Inputs and Outputs {{out, vkapi::kWrite}, {read_inputs, vkapi::kRead}}, // Shader params buffers @@ -134,7 +123,7 @@ static vkapi::ShaderInfo pick_matmul_coopmat_shader( return VK_KERNEL_FROM_STR(kernel_name); } -static utils::uvec3 pick_matmul_coopmat_global_wg_size( +static GlobalWorkGrid pick_matmul_coopmat_gwg( ComputeGraph* graph, const vkapi::ShaderInfo& shader, const std::vector& args, @@ -152,25 +141,14 @@ static utils::uvec3 pick_matmul_coopmat_global_wg_size( // to launch exactly num_tiles_n x num_tiles_m workgroups. // // The framework computes the group count as - // group_count = div_up(global_wg_size, local_wg_size) - // (see Context.cpp + Command.cpp). With local_wg = (kCoopmatInvocations, + // group_count = div_up(gwg, lwg) + // (see Context.cpp + Command.cpp). With lwg = (kCoopmatInvocations, // 1, 1), multiplying num_tiles_n by kCoopmatInvocations cancels the // div, yielding group_count.x = num_tiles_n. - return {num_tiles_n * kCoopmatInvocations, num_tiles_m, 1}; -} - -static utils::uvec3 pick_matmul_coopmat_local_wg_size( - ComputeGraph* graph, - const vkapi::ShaderInfo& shader, - const utils::uvec3& global_workgroup_size, - const std::vector& args, - const std::vector& resize_args) { - (void)graph; - (void)shader; - (void)global_workgroup_size; - (void)args; - (void)resize_args; - return {kCoopmatInvocations, 1, 1}; + return GlobalWorkGrid( + {num_tiles_n * kCoopmatInvocations, num_tiles_m, 1u}, + kTiledWorkGrid, + LocalWorkGroup(kCoopmatInvocations, 1u, 1u)); } void add_matmul_coopmat_node( @@ -190,8 +168,8 @@ void add_matmul_coopmat_node( graph.execute_nodes().emplace_back(new DynamicDispatchNode( graph, pick_matmul_coopmat_shader, - pick_matmul_coopmat_global_wg_size, - pick_matmul_coopmat_local_wg_size, + pick_matmul_coopmat_gwg, + pick_required_lwg, // Inputs and Outputs — same binding order as matmul_vec {{out, vkapi::kWrite}, {{mat1, mat2}, vkapi::kRead}}, // Shader params buffers — same UBOs as matmul_vec diff --git a/backends/vulkan/runtime/graph/ops/impl/GridPriors.cpp b/backends/vulkan/runtime/graph/ops/impl/GridPriors.cpp index 5f39c16d405..784b421fada 100644 --- a/backends/vulkan/runtime/graph/ops/impl/GridPriors.cpp +++ b/backends/vulkan/runtime/graph/ops/impl/GridPriors.cpp @@ -50,8 +50,8 @@ void add_grid_priors_node( graph.execute_nodes().emplace_back(new DynamicDispatchNode( graph, VK_KERNEL_FROM_STR(kernel_name), - default_pick_global_wg_size, - default_pick_local_wg_size, + default_pick_gwg, + default_pick_lwg, // Inputs and Outputs { {out, vkapi::kWrite}, diff --git a/backends/vulkan/runtime/graph/ops/impl/GridSampler2d.cpp b/backends/vulkan/runtime/graph/ops/impl/GridSampler2d.cpp index f5b10ad6576..0f5515233dc 100644 --- a/backends/vulkan/runtime/graph/ops/impl/GridSampler2d.cpp +++ b/backends/vulkan/runtime/graph/ops/impl/GridSampler2d.cpp @@ -92,8 +92,8 @@ void add_grid_sampler_2d_node( graph.execute_nodes().emplace_back(new DynamicDispatchNode( graph, VK_KERNEL_FROM_STR(kernel_name), - default_pick_global_wg_size, - default_pick_local_wg_size, + default_pick_gwg, + default_pick_lwg, // Inputs and Outputs {{out, vkapi::kWrite}, {{in, grid}, vkapi::kRead}}, // Shader params buffers. `meta_ubo` packs sizes, limits, axis_map, and diff --git a/backends/vulkan/runtime/graph/ops/impl/GroupNorm.cpp b/backends/vulkan/runtime/graph/ops/impl/GroupNorm.cpp index f18e59229de..cd26cc4c521 100644 --- a/backends/vulkan/runtime/graph/ops/impl/GroupNorm.cpp +++ b/backends/vulkan/runtime/graph/ops/impl/GroupNorm.cpp @@ -17,27 +17,28 @@ namespace vkcompute { -utils::uvec3 group_norm_local_wg_size( +GlobalWorkGrid group_norm_gwg( ComputeGraph* graph, const vkapi::ShaderInfo& shader, - const utils::uvec3& global_workgroup_size, const std::vector& args, const std::vector& resize_args) { - (void)graph; (void)shader; - (void)global_workgroup_size; - (void)args; (void)resize_args; - return {1u, 1u, 64u}; + const ValueRef mean = args.at(0).refs.at(0); + const LocalWorkGroup lwg(1u, 1u, 64u); + GlobalWorkGrid gwg( + {utils::safe_downcast(graph->numel_of(mean)), 1u, 1u}, + kLinearWorkGrid); + gwg.wrap_linear_dispatch( + graph->context()->adapter_ptr()->max_compute_workgroup_count(), lwg); + return gwg; } void resize_group_norm_texture_node( ComputeGraph* graph, const std::vector& args, const std::vector& resize_args) { - VK_CHECK_COND(graph != nullptr); - // Extract tensor references from args const ValueRef out = args.at(0).refs.at(0); const ValueRef in = args.at(1).refs.at(0); @@ -105,8 +106,8 @@ void add_native_group_norm_node( graph.execute_nodes().emplace_back(new DynamicDispatchNode( graph, VK_KERNEL_FROM_STR(kernel_name), - default_pick_global_wg_size, - group_norm_local_wg_size, + group_norm_gwg, + pick_required_lwg, // Inputs and Outputs {{{mean, rstd}, vkapi::kWrite}, {in, vkapi::kRead}}, // Shader params buffers @@ -138,8 +139,8 @@ void add_native_group_norm_node( graph.execute_nodes().emplace_back(new DynamicDispatchNode( graph, VK_KERNEL_FROM_STR(norm_kernel_name), - default_pick_global_wg_size, - default_pick_local_wg_size, + default_pick_gwg, + default_pick_lwg, // Inputs and Outputs {{out, vkapi::kWrite}, {{in, arg_weight, arg_bias, mean, rstd}, vkapi::kRead}}, diff --git a/backends/vulkan/runtime/graph/ops/impl/IndexSelect.cpp b/backends/vulkan/runtime/graph/ops/impl/IndexSelect.cpp index 576711a86f1..20ab76813a7 100644 --- a/backends/vulkan/runtime/graph/ops/impl/IndexSelect.cpp +++ b/backends/vulkan/runtime/graph/ops/impl/IndexSelect.cpp @@ -42,8 +42,8 @@ void add_index_select_channel_node( graph.execute_nodes().emplace_back(new DynamicDispatchNode( graph, VK_KERNEL_FROM_STR(kernel_name), - default_pick_global_wg_size, - default_pick_local_wg_size, + default_pick_gwg, + default_pick_lwg, {{out, vkapi::kWrite}, {{in, idx}, vkapi::kRead}}, {graph.sizes_ubo(out), graph.sizes_ubo(in)}, // Push Constants @@ -96,8 +96,8 @@ void add_index_select_node( graph.execute_nodes().emplace_back(new DynamicDispatchNode( graph, VK_KERNEL_FROM_STR(kernel_name), - default_pick_global_wg_size, - default_pick_local_wg_size, + default_pick_gwg, + default_pick_lwg, {{out, vkapi::kWrite}, {{in, idx}, vkapi::kRead}}, {graph.sizes_ubo(out), graph.create_params_buffer(params)}, // Push Constants diff --git a/backends/vulkan/runtime/graph/ops/impl/IndexTensor.cpp b/backends/vulkan/runtime/graph/ops/impl/IndexTensor.cpp index afce5ef3812..ddd8e8994b1 100644 --- a/backends/vulkan/runtime/graph/ops/impl/IndexTensor.cpp +++ b/backends/vulkan/runtime/graph/ops/impl/IndexTensor.cpp @@ -51,8 +51,8 @@ void add_index_tensor_node( graph.execute_nodes().emplace_back(new DynamicDispatchNode( graph, VK_KERNEL_FROM_STR(kernel_name), - default_pick_global_wg_size, - default_pick_local_wg_size, + default_pick_gwg, + default_pick_lwg, // Inputs and Outputs {{out, vkapi::kWrite}, {{self, index}, vkapi::kRead}}, // Shader params buffers diff --git a/backends/vulkan/runtime/graph/ops/impl/Int8x4Staging.cpp b/backends/vulkan/runtime/graph/ops/impl/Int8x4Staging.cpp index eb1d9965f30..b6d15e23e88 100644 --- a/backends/vulkan/runtime/graph/ops/impl/Int8x4Staging.cpp +++ b/backends/vulkan/runtime/graph/ops/impl/Int8x4Staging.cpp @@ -32,14 +32,14 @@ void add_prepack_int8x4_buffer_node( // (e.g., kPackedInt8_4C with C=3 pads to C=4). uint32_t num_texels = utils::safe_downcast(graph.padded_numel_of(tensor) / 4); - utils::uvec3 global_wg_size = {num_texels, 1, 1}; - utils::uvec3 local_wg_size = graph.create_local_wg_size(global_wg_size); + const GlobalWorkGrid gwg({num_texels, 1u, 1u}, kTiledWorkGrid); + const LocalWorkGroup lwg = graph.create_lwg(gwg); graph.prepack_nodes().emplace_back(new PrepackNode( graph, VK_KERNEL_FROM_STR(kernel_name), - global_wg_size, - local_wg_size, + gwg, + lwg, // Input and Output tensor_data, tensor, @@ -49,7 +49,7 @@ void add_prepack_int8x4_buffer_node( {graph.hashed_layout_of(tensor)})); } -static utils::uvec3 staging_to_int8x4_buffer_global_wg_size( +static GlobalWorkGrid staging_to_int8x4_buffer_gwg( ComputeGraph* graph, const vkapi::ShaderInfo& shader, const std::vector& args, @@ -59,7 +59,7 @@ static utils::uvec3 staging_to_int8x4_buffer_global_wg_size( const ValueRef out_tensor = args.at(0).refs.at(0); const uint32_t num_texels = utils::safe_downcast(graph->padded_numel_of(out_tensor) / 4); - return {num_texels, 1, 1}; + return graph->create_linear_gwg(num_texels); } void add_staging_to_int8x4_buffer_node( @@ -76,8 +76,8 @@ void add_staging_to_int8x4_buffer_node( graph.execute_nodes().emplace_back(new DynamicDispatchNode( graph, VK_KERNEL_FROM_STR("nchw_to_int8x4_buffer"), - staging_to_int8x4_buffer_global_wg_size, - default_pick_local_wg_size, + staging_to_int8x4_buffer_gwg, + default_pick_lwg, // Input and Output {{tensor, vkapi::kWrite}, {in_staging, vkapi::kRead}}, // Parameter Buffers @@ -92,7 +92,7 @@ void add_staging_to_int8x4_buffer_node( nullptr)); } -static utils::uvec3 int8x4_buffer_to_staging_global_wg_size( +static GlobalWorkGrid int8x4_buffer_to_staging_gwg( ComputeGraph* graph, const vkapi::ShaderInfo& shader, const std::vector& args, @@ -104,7 +104,7 @@ static utils::uvec3 int8x4_buffer_to_staging_global_wg_size( const int32_t numel = graph->numel_of(in_tensor); const uint32_t num_out_int32s = utils::safe_downcast((numel + 3) / 4); - return {num_out_int32s, 1, 1}; + return graph->create_linear_gwg(num_out_int32s); } void add_int8x4_buffer_to_staging_node( @@ -121,8 +121,8 @@ void add_int8x4_buffer_to_staging_node( graph.execute_nodes().emplace_back(new DynamicDispatchNode( graph, VK_KERNEL_FROM_STR("int8x4_buffer_to_nchw"), - int8x4_buffer_to_staging_global_wg_size, - default_pick_local_wg_size, + int8x4_buffer_to_staging_gwg, + default_pick_lwg, // Input and Output {{staging_data, vkapi::kWrite}, {tensor, vkapi::kRead}}, // Parameter Buffers diff --git a/backends/vulkan/runtime/graph/ops/impl/Linear.cpp b/backends/vulkan/runtime/graph/ops/impl/Linear.cpp index 13415f33919..86f27b63586 100644 --- a/backends/vulkan/runtime/graph/ops/impl/Linear.cpp +++ b/backends/vulkan/runtime/graph/ops/impl/Linear.cpp @@ -65,7 +65,7 @@ vkapi::ShaderInfo pick_linear_shader( return VK_KERNEL_FROM_STR(kernel_name); } -utils::uvec3 pick_linear_global_wg_size( +GlobalWorkGrid pick_linear_gwg( ComputeGraph* graph, const vkapi::ShaderInfo& shader, const std::vector& args, @@ -77,7 +77,8 @@ utils::uvec3 pick_linear_global_wg_size( uint32_t M = graph->size_at(-2, out); uint32_t B = graph->dim_of(out) >= 3 ? graph->size_at(-3, out) : 1; uint32_t tile_m = pick_matmul_tile_m(graph, out); - return {utils::div_up_4(N), utils::div_up(M, tile_m), B}; + return GlobalWorkGrid( + {utils::div_up_4(N), utils::div_up(M, tile_m), B}, kTiledWorkGrid); } void add_linear_tiled_node( @@ -122,8 +123,8 @@ void add_linear_tiled_node( graph.execute_nodes().emplace_back(new DynamicDispatchNode( graph, pick_linear_shader, - pick_linear_global_wg_size, - pick_hw_square_wg_size, + pick_linear_gwg, + pick_xy_square_lwg, // Inputs and Outputs {{out, vkapi::kWrite}, {read_inputs, vkapi::kRead}}, // Shader params buffers diff --git a/backends/vulkan/runtime/graph/ops/impl/LinearDw.cpp b/backends/vulkan/runtime/graph/ops/impl/LinearDw.cpp index 1ab5e6976ca..9adff9dcef4 100644 --- a/backends/vulkan/runtime/graph/ops/impl/LinearDw.cpp +++ b/backends/vulkan/runtime/graph/ops/impl/LinearDw.cpp @@ -27,7 +27,7 @@ void resize_linear_dW_node( graph->virtual_resize(dW, {N, K}); } -utils::uvec3 linear_dW_global_wg_size( +GlobalWorkGrid linear_dW_gwg( ComputeGraph* graph, const vkapi::ShaderInfo& shader, const std::vector& args, @@ -38,21 +38,21 @@ utils::uvec3 linear_dW_global_wg_size( const uint32_t N = graph->size_at(-2, dW); const uint32_t K = graph->size_at(-1, dW); const uint32_t tiles = utils::div_up_4(N) * utils::div_up_4(K); - return {tiles, 1u, 1u}; + return graph->create_linear_gwg(tiles); } -utils::uvec3 linear_dW_local_wg_size( +LocalWorkGroup linear_dW_lwg( ComputeGraph* graph, const vkapi::ShaderInfo& shader, - const utils::uvec3& global_workgroup_size, + const GlobalWorkGrid& gwg, const std::vector& args, const std::vector& resize_args) { (void)graph; (void)shader; - (void)global_workgroup_size; + (void)gwg; (void)args; (void)resize_args; - return {64u, 1u, 1u}; + return LocalWorkGroup(64u, 1u, 1u); } void linear_dW(ComputeGraph& graph, const std::vector& args) { @@ -89,8 +89,8 @@ void linear_dW(ComputeGraph& graph, const std::vector& args) { graph.execute_nodes().emplace_back(new DynamicDispatchNode( graph, VK_KERNEL_FROM_STR(kernel_name), - linear_dW_global_wg_size, - linear_dW_local_wg_size, + linear_dW_gwg, + linear_dW_lwg, // Inputs and Outputs {{dW, vkapi::kWrite}, {{d_out, x}, vkapi::kRead}}, // Shader params buffers diff --git a/backends/vulkan/runtime/graph/ops/impl/Matmul.cpp b/backends/vulkan/runtime/graph/ops/impl/Matmul.cpp index ed6825f6126..26ec08863cc 100644 --- a/backends/vulkan/runtime/graph/ops/impl/Matmul.cpp +++ b/backends/vulkan/runtime/graph/ops/impl/Matmul.cpp @@ -78,7 +78,7 @@ vkapi::ShaderInfo pick_matmul_tiled_shader( return VK_KERNEL_FROM_STR(kernel_name); } -utils::uvec3 pick_matmul_tiled_global_wg_size( +GlobalWorkGrid pick_matmul_tiled_gwg( ComputeGraph* graph, const vkapi::ShaderInfo& shader, const std::vector& args, @@ -90,7 +90,8 @@ utils::uvec3 pick_matmul_tiled_global_wg_size( uint32_t M = graph->size_at(-2, out); uint32_t B = graph->dim_of(out) >= 3 ? graph->size_at(-3, out) : 1; uint32_t tile_m = pick_matmul_tile_m(graph, out); - return {utils::div_up_4(N), utils::div_up(M, tile_m), B}; + return GlobalWorkGrid( + {utils::div_up_4(N), utils::div_up(M, tile_m), B}, kTiledWorkGrid); } void add_matmul_tiled_node( @@ -106,8 +107,8 @@ void add_matmul_tiled_node( graph.execute_nodes().emplace_back(new DynamicDispatchNode( graph, pick_matmul_tiled_shader, - pick_matmul_tiled_global_wg_size, - pick_hw_square_wg_size, + pick_matmul_tiled_gwg, + pick_xy_square_lwg, // Inputs and Outputs {{out, vkapi::kWrite}, {{mat1, mat2}, vkapi::kRead}}, // Shader params buffers @@ -147,8 +148,8 @@ void add_addmm_tiled_node( graph.execute_nodes().emplace_back(new DynamicDispatchNode( graph, pick_matmul_tiled_shader, - pick_matmul_tiled_global_wg_size, - pick_hw_square_wg_size, + pick_matmul_tiled_gwg, + pick_xy_square_lwg, // Inputs and Outputs {{out, vkapi::kWrite}, {{mat1, mat2, bias}, vkapi::kRead}}, // Shader params buffers diff --git a/backends/vulkan/runtime/graph/ops/impl/NativeLayerNorm.cpp b/backends/vulkan/runtime/graph/ops/impl/NativeLayerNorm.cpp index 73d1ea908e9..b1d8965a4d1 100644 --- a/backends/vulkan/runtime/graph/ops/impl/NativeLayerNorm.cpp +++ b/backends/vulkan/runtime/graph/ops/impl/NativeLayerNorm.cpp @@ -50,7 +50,7 @@ void resize_native_layer_norm_node( graph->virtual_resize(rstd, mean_size); } -utils::uvec3 layer_norm_buffer_global_wg_size( +GlobalWorkGrid layer_norm_buffer_gwg( ComputeGraph* graph, const vkapi::ShaderInfo& shader, const std::vector& args, @@ -60,21 +60,8 @@ utils::uvec3 layer_norm_buffer_global_wg_size( const ValueRef mean_tensor = args.at(0).refs.at(1); const uint32_t num_rows = utils::safe_downcast(graph->numel_of(mean_tensor)); - return {1u, num_rows, 1u}; -} - -utils::uvec3 layer_norm_buffer_local_wg_size( - ComputeGraph* graph, - const vkapi::ShaderInfo& shader, - const utils::uvec3& global_workgroup_size, - const std::vector& args, - const std::vector& resize_args) { - (void)graph; - (void)shader; - (void)global_workgroup_size; - (void)args; - (void)resize_args; - return {64u, 1u, 1u}; + return GlobalWorkGrid( + {1u, num_rows, 1u}, kTiledWorkGrid, LocalWorkGroup(64u, 1u, 1u)); } void add_native_layer_norm_node( @@ -131,9 +118,8 @@ void add_native_layer_norm_node( graph.execute_nodes().emplace_back(new DynamicDispatchNode( graph, VK_KERNEL_FROM_STR(kernel_name), - is_buffer ? layer_norm_buffer_global_wg_size - : default_pick_global_wg_size, - is_buffer ? layer_norm_buffer_local_wg_size : default_pick_local_wg_size, + is_buffer ? layer_norm_buffer_gwg : default_pick_gwg, + is_buffer ? pick_required_lwg : default_pick_lwg, // Inputs and Outputs {{{out_tensor, mean_tensor, rstd_tensor}, vkapi::kWrite}, {{in, arg_weight, arg_bias}, vkapi::kRead}}, diff --git a/backends/vulkan/runtime/graph/ops/impl/Pad.cpp b/backends/vulkan/runtime/graph/ops/impl/Pad.cpp index 7f872512c05..d2593171c56 100644 --- a/backends/vulkan/runtime/graph/ops/impl/Pad.cpp +++ b/backends/vulkan/runtime/graph/ops/impl/Pad.cpp @@ -71,8 +71,8 @@ void add_constant_pad_nd_node( graph.execute_nodes().emplace_back(new DynamicDispatchNode( graph, VK_KERNEL_FROM_STR(kernel_name), - default_pick_global_wg_size, - default_pick_local_wg_size, + default_pick_gwg, + default_pick_lwg, {{out, vkapi::kWrite}, {in, vkapi::kRead}}, // Parameter buffers {graph.meta_ubo(out), diff --git a/backends/vulkan/runtime/graph/ops/impl/Permute.cpp b/backends/vulkan/runtime/graph/ops/impl/Permute.cpp index 8081424cfb7..00e9924b544 100644 --- a/backends/vulkan/runtime/graph/ops/impl/Permute.cpp +++ b/backends/vulkan/runtime/graph/ops/impl/Permute.cpp @@ -170,8 +170,8 @@ void add_permute_node( graph.execute_nodes().emplace_back(new DynamicDispatchNode( graph, VK_KERNEL_FROM_STR(kernel_name), - default_pick_global_wg_size, - default_pick_local_wg_size, + default_pick_gwg, + default_pick_lwg, {{out, vkapi::kWrite}, {in, vkapi::kRead}}, param_ubos, push_constants, diff --git a/backends/vulkan/runtime/graph/ops/impl/PixelShuffle.cpp b/backends/vulkan/runtime/graph/ops/impl/PixelShuffle.cpp index 24c00b9d7af..306c3906a9d 100644 --- a/backends/vulkan/runtime/graph/ops/impl/PixelShuffle.cpp +++ b/backends/vulkan/runtime/graph/ops/impl/PixelShuffle.cpp @@ -64,8 +64,8 @@ void add_pixel_shuffle_node( graph.execute_nodes().emplace_back(new DynamicDispatchNode( graph, VK_KERNEL_FROM_STR(kernel_name), - default_pick_global_wg_size, - default_pick_local_wg_size, + default_pick_gwg, + default_pick_lwg, // Inputs and Outputs {{out, vkapi::kWrite}, {in, vkapi::kRead}}, // Shader params buffers diff --git a/backends/vulkan/runtime/graph/ops/impl/Pool.cpp b/backends/vulkan/runtime/graph/ops/impl/Pool.cpp index d405825fad1..a85b96ea157 100644 --- a/backends/vulkan/runtime/graph/ops/impl/Pool.cpp +++ b/backends/vulkan/runtime/graph/ops/impl/Pool.cpp @@ -110,8 +110,8 @@ void add_max_pool2d_node( graph.execute_nodes().emplace_back(new DynamicDispatchNode( graph, VK_KERNEL_FROM_STR(kernel_name), - default_pick_global_wg_size, - default_pick_local_wg_size, + default_pick_gwg, + default_pick_lwg, // Inputs and Outputs {{{out_tensor, indices_tensor}, vkapi::kWrite}, {in, vkapi::kRead}}, // Shader params buffers @@ -179,8 +179,8 @@ void add_avg_pool2d_node( graph.execute_nodes().emplace_back(new DynamicDispatchNode( graph, VK_KERNEL_FROM_STR(kernel_name), - default_pick_global_wg_size, - default_pick_local_wg_size, + default_pick_gwg, + default_pick_lwg, // Inputs and Outputs {{out, vkapi::kWrite}, {in, vkapi::kRead}}, // Shader params buffers diff --git a/backends/vulkan/runtime/graph/ops/impl/Preprocess.cpp b/backends/vulkan/runtime/graph/ops/impl/Preprocess.cpp index e8bfb97a4be..32ea9677347 100644 --- a/backends/vulkan/runtime/graph/ops/impl/Preprocess.cpp +++ b/backends/vulkan/runtime/graph/ops/impl/Preprocess.cpp @@ -20,9 +20,9 @@ namespace vkcompute { // // M and K are read from fp_input's live sizes (resize_args[0]) so that // virtual_resize updates flow through. When M == 1 the transpose is a no-op -// (the downstream GEMV path reads fp_input directly) and global_wg returns +// (the downstream GEMV path reads fp_input directly) and gwg returns // {0,0,0} to make DispatchNode::encode() skip the recording entirely. -static utils::uvec3 transpose_cast_global_wg_size( +static GlobalWorkGrid transpose_cast_gwg( ComputeGraph* graph, const vkapi::ShaderInfo& shader, const std::vector& args, @@ -35,29 +35,31 @@ static utils::uvec3 transpose_cast_global_wg_size( const uint32_t K = static_cast(utils::val_at(-1, in_sizes)); if (M == 1u) { - return {0u, 0u, 0u}; + return GlobalWorkGrid({0u, 0u, 0u}, kTiledWorkGrid); } bool is_4x4 = shader.kernel_name.find("4x4") != std::string::npos; if (is_4x4) { - return {utils::div_up(K, 4u), utils::div_up(M, 4u), 1u}; + return GlobalWorkGrid( + {utils::div_up(K, 4u), utils::div_up(M, 4u), 1u}, kTiledWorkGrid); } - return {K, utils::div_up(M, 4u), 1u}; + return GlobalWorkGrid({K, utils::div_up(M, 4u), 1u}, kTiledWorkGrid); } -static utils::uvec3 transpose_cast_local_wg_size( +static LocalWorkGroup transpose_cast_lwg( ComputeGraph* graph, const vkapi::ShaderInfo& shader, - const utils::uvec3& global_workgroup_size, + const GlobalWorkGrid& gwg, const std::vector& args, const std::vector& resize_args) { (void)graph; - (void)global_workgroup_size; + (void)gwg; (void)args; (void)resize_args; bool is_4x4 = shader.kernel_name.find("4x4") != std::string::npos; - return is_4x4 ? utils::uvec3{2u, 16u, 1u} : utils::uvec3{8u, 8u, 1u}; + return LocalWorkGroup( + is_4x4 ? utils::uvec3{2u, 16u, 1u} : utils::uvec3{8u, 8u, 1u}); } // Resize the transposed output tensor to match current fp_input dimensions. @@ -102,8 +104,8 @@ void add_transpose_cast_contig_to_vectorized_node( graph.execute_nodes().emplace_back(new DynamicDispatchNode( graph, VK_KERNEL_FROM_STR(kernel_name), - transpose_cast_global_wg_size, - transpose_cast_local_wg_size, + transpose_cast_gwg, + transpose_cast_lwg, {{output, vkapi::kWrite}, {fp_input, vkapi::kRead}}, {graph.sizes_ubo(fp_input)}, {}, diff --git a/backends/vulkan/runtime/graph/ops/impl/Q4gswLinear.cpp b/backends/vulkan/runtime/graph/ops/impl/Q4gswLinear.cpp index 62322602ac3..cf8afc8c8c1 100644 --- a/backends/vulkan/runtime/graph/ops/impl/Q4gswLinear.cpp +++ b/backends/vulkan/runtime/graph/ops/impl/Q4gswLinear.cpp @@ -56,7 +56,7 @@ namespace { // 2. An adaptive nc-coop GEMV DynamicDispatchNode whose global WG // self-gates to {0,0,0} at M!=1 — handles decode (M==1) only. // -// The framework re-invokes pick_shader_fn / pick_global_wg / pick_local_wg +// The framework re-invokes pick_shader_fn / pick_gwg / pick_lwg // on every trigger_resize(), so M transitions across `virtual_resize` are // routed to the correct node without re-encode beyond what the changed WG // shape requires. @@ -168,7 +168,7 @@ vkapi::ShaderInfo pick_q4gsw_nc_coop_shader( // Global WG for the nc-coop GEMV. Self-gates to {0,0,0} when M != 1 so the // node is a no-op on prefill (the parallel GEMM dispatch handles M>1). -utils::uvec3 pick_q4gsw_nc_coop_global_wg( +GlobalWorkGrid pick_q4gsw_nc_coop_gwg( ComputeGraph* graph, const vkapi::ShaderInfo& shader, const std::vector& args, @@ -179,38 +179,26 @@ utils::uvec3 pick_q4gsw_nc_coop_global_wg( const std::vector out_sizes = graph->sizes_of(out); const uint32_t M = utils::safe_downcast(utils::val_at(-2, out_sizes)); - if (M != 1u) { - return {0u, 0u, 0u}; - } const uint32_t N = utils::safe_downcast(utils::val_at(-1, out_sizes)); - const uint32_t N8 = (N + 7u) / 8u; const CoopVariant v = pick_coop_variant_for_N(N); + const LocalWorkGroup lwg(1u, v.num_groups, v.workers_per_group); + if (M != 1u) { + return GlobalWorkGrid({0u, 0u, 0u}, kTiledWorkGrid, lwg); + } + const uint32_t N8 = (N + 7u) / 8u; const uint32_t wgs_along_x = utils::div_up(N8, v.num_groups); - return {wgs_along_x, v.num_groups, v.workers_per_group}; -} - -utils::uvec3 pick_q4gsw_nc_coop_local_wg( - ComputeGraph* graph, - const vkapi::ShaderInfo& shader, - const utils::uvec3& global_workgroup_size, - const std::vector& args, - const std::vector& resize_args) { - (void)shader; - (void)global_workgroup_size; - (void)resize_args; - const ValueRef out = args.at(0).refs.at(0); - const uint32_t N = - utils::safe_downcast(utils::val_at(-1, graph->sizes_of(out))); - const CoopVariant v = pick_coop_variant_for_N(N); - return {1u, v.num_groups, v.workers_per_group}; + GlobalWorkGrid gwg({wgs_along_x * v.num_groups, 1u, 1u}, kLinearWorkGrid); + gwg.wrap_linear_dispatch( + graph->context()->adapter_ptr()->max_compute_workgroup_count(), lwg); + return gwg; } } // namespace // Global WG picker for the fp32 GEMM path. Exposed so the forced-shader test // selectors (GEMM_W_4X8) can dispatch the same kernel with arbitrary M. -utils::uvec3 pick_q4gsw_linear_gemm_global_wg( +GlobalWorkGrid pick_q4gsw_linear_gemm_gwg( ComputeGraph* graph, const vkapi::ShaderInfo& shader, const std::vector& args, @@ -224,26 +212,28 @@ utils::uvec3 pick_q4gsw_linear_gemm_global_wg( const uint32_t M = utils::safe_downcast(utils::val_at(-2, out_sizes)); // fp32 GEMM: 4M x 8N per-thread tile. - return {utils::div_up(N, kGemmTileN), utils::div_up(M, kGemmTileM), 1u}; + return GlobalWorkGrid( + {utils::div_up(N, kGemmTileN), utils::div_up(M, kGemmTileM), 1u}, + kTiledWorkGrid); } // Local WG picker for the fp32 GEMM path. -utils::uvec3 pick_q4gsw_linear_gemm_local_wg( +LocalWorkGroup pick_q4gsw_linear_gemm_lwg( ComputeGraph* graph, const vkapi::ShaderInfo& shader, - const utils::uvec3& global_workgroup_size, + const GlobalWorkGrid& gwg, const std::vector& args, const std::vector& resize_args) { (void)graph; (void)shader; - (void)global_workgroup_size; + (void)gwg; (void)args; (void)resize_args; - return {8u, 8u, 1u}; + return LocalWorkGroup(8u, 8u, 1u); } // Global WG picker for the fp16 tin GEMM path. -utils::uvec3 pick_q4gsw_linear_tin_gemm_global_wg( +GlobalWorkGrid pick_q4gsw_linear_tin_gemm_gwg( ComputeGraph* graph, const vkapi::ShaderInfo& shader, const std::vector& args, @@ -258,22 +248,24 @@ utils::uvec3 pick_q4gsw_linear_tin_gemm_global_wg( utils::safe_downcast(utils::val_at(-2, out_sizes)); // fp16 tin GEMM: 8M x 4N per-thread tile. Shader x/y are swapped relative // to the fp32 GEMM — x = M tiles, y = N tiles. - return {utils::div_up(M, kTinGemmTileM), utils::div_up(N, kTinGemmTileN), 1u}; + return GlobalWorkGrid( + {utils::div_up(M, kTinGemmTileM), utils::div_up(N, kTinGemmTileN), 1u}, + kTiledWorkGrid); } // Local WG picker for the fp16 tin GEMM path. -utils::uvec3 pick_q4gsw_linear_tin_gemm_local_wg( +LocalWorkGroup pick_q4gsw_linear_tin_gemm_lwg( ComputeGraph* graph, const vkapi::ShaderInfo& shader, - const utils::uvec3& global_workgroup_size, + const GlobalWorkGrid& gwg, const std::vector& args, const std::vector& resize_args) { (void)graph; (void)shader; - (void)global_workgroup_size; + (void)gwg; (void)args; (void)resize_args; - return {1u, 128u, 1u}; + return LocalWorkGroup(1u, 128u, 1u); } namespace { @@ -283,7 +275,7 @@ namespace { // M==1 decode; this gate prevents the GEMM shader from running at M==1 and // overwriting the nc-coop output. The ungated pickers remain available for // forced-shader test selectors that need to dispatch GEMM at arbitrary M. -utils::uvec3 pick_q4gsw_linear_gemm_gated_global_wg( +GlobalWorkGrid pick_q4gsw_linear_gemm_gated_gwg( ComputeGraph* graph, const vkapi::ShaderInfo& shader, const std::vector& args, @@ -292,12 +284,12 @@ utils::uvec3 pick_q4gsw_linear_gemm_gated_global_wg( const uint32_t M = utils::safe_downcast(utils::val_at(-2, graph->sizes_of(out))); if (M == 1u) { - return {0u, 0u, 0u}; + return GlobalWorkGrid({0u, 0u, 0u}, kTiledWorkGrid); } - return pick_q4gsw_linear_gemm_global_wg(graph, shader, args, resize_args); + return pick_q4gsw_linear_gemm_gwg(graph, shader, args, resize_args); } -utils::uvec3 pick_q4gsw_linear_tin_gemm_gated_global_wg( +GlobalWorkGrid pick_q4gsw_linear_tin_gemm_gated_gwg( ComputeGraph* graph, const vkapi::ShaderInfo& shader, const std::vector& args, @@ -306,9 +298,9 @@ utils::uvec3 pick_q4gsw_linear_tin_gemm_gated_global_wg( const uint32_t M = utils::safe_downcast(utils::val_at(-2, graph->sizes_of(out))); if (M == 1u) { - return {0u, 0u, 0u}; + return GlobalWorkGrid({0u, 0u, 0u}, kTiledWorkGrid); } - return pick_q4gsw_linear_tin_gemm_global_wg(graph, shader, args, resize_args); + return pick_q4gsw_linear_tin_gemm_gwg(graph, shader, args, resize_args); } } // namespace @@ -365,16 +357,17 @@ ValueRef prepack_q4_w_4x8_nc_buffer( // identical layout. const int32_t n4_pitch = utils::safe_downcast(N4_padded); - utils::uvec3 global_wg = { - utils::safe_downcast(K4), - utils::safe_downcast(N8), - 1u}; + const GlobalWorkGrid gwg( + {utils::safe_downcast(K4), + utils::safe_downcast(N8), + 1u}, + kTiledWorkGrid); graph.prepack_nodes().emplace_back(new PrepackNode( graph, VK_KERNEL_FROM_STR("pack_q4_linear_weight__w_4x8_nc_buffer"), - global_wg, - graph.create_local_wg_size(global_wg), + gwg, + graph.create_lwg(gwg), weight_data, packed_weight, {}, @@ -409,7 +402,7 @@ ValueRef prepack_q4_scales( // 1. GEMM DynamicDispatchNode — self-gates to {0,0,0} when M==1. // 2. nc-coop GEMV DynamicDispatchNode — self-gates to {0,0,0} when M!=1. // Together they cover decode (M==1) and prefill (M>1) without re-encode cost, -// since the framework re-runs pick_shader_fn + pick_global_wg on every +// since the framework re-runs pick_shader_fn + pick_gwg on every // trigger_resize() and re-encodes only when the chosen kernel changes. // // The fp16 path additionally requires a transpose preprocess dispatch @@ -423,7 +416,7 @@ ValueRef prepack_q4_scales( // (output, fp_input, transposed_input, q4_weights, scales, bias), where // `transposed_input` is a 0-element dummy (nc-coop never reads it). // -// Self-gates to {0,0,0} when M != 1 via pick_q4gsw_nc_coop_global_wg, so the +// Self-gates to {0,0,0} when M != 1 via pick_q4gsw_nc_coop_gwg, so the // node is a no-op at prefill. At decode, pick_q4gsw_nc_coop_shader selects // the nc-buffer coop variant whose (NUM_GROUPS, WORKERS_PER_GROUP) decomp is // best for the current N. The nc-buffer payload is byte-identical to the @@ -449,8 +442,8 @@ void add_q4gsw_linear_nc_coop_gemv_node( graph.execute_nodes().emplace_back(new DynamicDispatchNode( graph, pick_q4gsw_nc_coop_shader, - pick_q4gsw_nc_coop_global_wg, - pick_q4gsw_nc_coop_local_wg, + pick_q4gsw_nc_coop_gwg, + pick_required_lwg, {{output, vkapi::kWrite}, {{fp_input, dummy_transposed_input.vref, @@ -513,8 +506,8 @@ void add_q4gsw_linear_w_4x8_node( graph.execute_nodes().emplace_back(new DynamicDispatchNode( graph, pick_q4gsw_linear_w_4x8_shader, - pick_q4gsw_linear_gemm_gated_global_wg, - pick_q4gsw_linear_gemm_local_wg, + pick_q4gsw_linear_gemm_gated_gwg, + pick_q4gsw_linear_gemm_lwg, {{output, vkapi::kWrite}, {{fp_input, dummy_transposed_input.vref, @@ -608,8 +601,8 @@ void add_q4gsw_linear_tin_w_4x8_node( graph.execute_nodes().emplace_back(new DynamicDispatchNode( graph, pick_q4gsw_linear_tin_w_4x8_shader, - pick_q4gsw_linear_tin_gemm_gated_global_wg, - pick_q4gsw_linear_tin_gemm_local_wg, + pick_q4gsw_linear_tin_gemm_gated_gwg, + pick_q4gsw_linear_tin_gemm_lwg, {{output, vkapi::kWrite}, {{fp_input, transposed_input.vref, diff --git a/backends/vulkan/runtime/graph/ops/impl/Q4gswLinear.h b/backends/vulkan/runtime/graph/ops/impl/Q4gswLinear.h index d3268b4ec7c..6f707a52eee 100644 --- a/backends/vulkan/runtime/graph/ops/impl/Q4gswLinear.h +++ b/backends/vulkan/runtime/graph/ops/impl/Q4gswLinear.h @@ -52,31 +52,31 @@ ValueRef prepack_q4_scales( vkapi::ScalarType dtype); // Global/local workgroup pickers for the fp32 GEMM path. -utils::uvec3 pick_q4gsw_linear_gemm_global_wg( +GlobalWorkGrid pick_q4gsw_linear_gemm_gwg( ComputeGraph* graph, const vkapi::ShaderInfo& shader, const std::vector& args, const std::vector& resize_args); -utils::uvec3 pick_q4gsw_linear_gemm_local_wg( +LocalWorkGroup pick_q4gsw_linear_gemm_lwg( ComputeGraph* graph, const vkapi::ShaderInfo& shader, - const utils::uvec3& global_workgroup_size, + const GlobalWorkGrid& gwg, const std::vector& args, const std::vector& resize_args); // Global/local workgroup pickers for the fp16 tin GEMM path — // {ceil(M/8), ceil(N/4), 1} global, {1, 128, 1} local. -utils::uvec3 pick_q4gsw_linear_tin_gemm_global_wg( +GlobalWorkGrid pick_q4gsw_linear_tin_gemm_gwg( ComputeGraph* graph, const vkapi::ShaderInfo& shader, const std::vector& args, const std::vector& resize_args); -utils::uvec3 pick_q4gsw_linear_tin_gemm_local_wg( +LocalWorkGroup pick_q4gsw_linear_tin_gemm_lwg( ComputeGraph* graph, const vkapi::ShaderInfo& shader, - const utils::uvec3& global_workgroup_size, + const GlobalWorkGrid& gwg, const std::vector& args, const std::vector& resize_args); diff --git a/backends/vulkan/runtime/graph/ops/impl/Q8taBinary.cpp b/backends/vulkan/runtime/graph/ops/impl/Q8taBinary.cpp index 44fbc4bc8f6..8241533078c 100644 --- a/backends/vulkan/runtime/graph/ops/impl/Q8taBinary.cpp +++ b/backends/vulkan/runtime/graph/ops/impl/Q8taBinary.cpp @@ -113,15 +113,15 @@ void add_q8ta_binary_node( const BlockConfig block_config = create_block_config_for_tensor(graph, packed_int8_output); - // Cast block config to ValueRef for pick_linear_global_wg_with_block_config + // Cast block config to ValueRef for pick_linear_gwg_with_block_config const ValueRef block_config_ref = static_cast(block_config.as_packed_int()); graph.execute_nodes().emplace_back(new DynamicDispatchNode( graph, VK_KERNEL_FROM_STR(kernel_name), - pick_linear_global_wg_with_block_config, - pick_square_local_wg_with_block_config, + pick_linear_gwg_with_block_config, + pick_square_lwg_with_block_config, // Inputs and Outputs {{packed_int8_output, vkapi::kWrite}, {{packed_int8_input_a, packed_int8_input_b}, vkapi::kRead}}, diff --git a/backends/vulkan/runtime/graph/ops/impl/Q8taClone.cpp b/backends/vulkan/runtime/graph/ops/impl/Q8taClone.cpp index 6c688af802d..a1ed2371ec7 100644 --- a/backends/vulkan/runtime/graph/ops/impl/Q8taClone.cpp +++ b/backends/vulkan/runtime/graph/ops/impl/Q8taClone.cpp @@ -39,7 +39,7 @@ void add_q8ta_clone_node( const BlockConfig inp_block_config = create_block_config_from_other( graph, packed_int8_input, outp_block_config); - // Cast block config to ValueRef for pick_*_global_wg_with_block_config + // Cast block config to ValueRef for pick_*_gwg_with_block_config // Use inp_block_config since shader uses inp_block_config for indexing const ValueRef block_config_ref = static_cast(inp_block_config.as_packed_int()); @@ -48,8 +48,8 @@ void add_q8ta_clone_node( graph.execute_nodes().emplace_back(new DynamicDispatchNode( graph, VK_KERNEL_FROM_STR(kernel_name), - pick_linear_global_wg_with_block_config, - pick_square_local_wg_with_block_config, + pick_linear_gwg_with_block_config, + pick_square_lwg_with_block_config, // Inputs and Outputs {{packed_int8_output, vkapi::kWrite}, {packed_int8_input, vkapi::kRead}}, // Shader params buffers diff --git a/backends/vulkan/runtime/graph/ops/impl/Q8taConv2d.cpp b/backends/vulkan/runtime/graph/ops/impl/Q8taConv2d.cpp index b9f17021ea0..70171a04820 100644 --- a/backends/vulkan/runtime/graph/ops/impl/Q8taConv2d.cpp +++ b/backends/vulkan/runtime/graph/ops/impl/Q8taConv2d.cpp @@ -42,6 +42,52 @@ bool q8ta_conv2d_check_4w4c_packed_dim_info(const api::PackedDimInfo& info) { info.outer_packed_dim_block_size == 4; } +namespace { + +uint64_t q8ta_conv2d_im2col_scratch_limit(ComputeGraph& graph) { + constexpr uint64_t kMaxBatchedIm2ColScratchBytes = 32ULL * 1024ULL * 1024ULL; + const uint64_t device_scratch_limit = + graph.context()->adapter_ptr()->max_buffer_numel(); + return device_scratch_limit < kMaxBatchedIm2ColScratchBytes + ? device_scratch_limit + : kMaxBatchedIm2ColScratchBytes; +} + +bool should_use_q8ta_conv2d_im2col( + ComputeGraph& graph, + const int64_t batch, + const int64_t groups, + const int64_t in_channels_per_group, + const int64_t flattened_kernel_size, + const int64_t out_height, + const int64_t out_width) { + const bool im2col_eligible = in_channels_per_group % 4 == 0; + if (!im2col_eligible) { + return false; + } + + const int64_t spatial_out = out_height * out_width; + if (batch > 1) { + constexpr int64_t kMinFlattenedKernelSize = 1024; + constexpr int64_t kMaxSpatialOutput = 64; + const uint64_t scratch_bytes = static_cast(batch) * + static_cast(flattened_kernel_size) * + static_cast(out_height) * + static_cast(utils::align_up_4(out_width)); + return groups == 1 && flattened_kernel_size >= kMinFlattenedKernelSize && + spatial_out <= kMaxSpatialOutput && + scratch_bytes <= q8ta_conv2d_im2col_scratch_limit(graph); + } + + if (graph.device_is_mali()) { + return true; + } + + return groups == 1 && (in_channels_per_group >= 32 || spatial_out <= 4096); +} + +} // namespace + // // Workgroup size selection functions // @@ -57,7 +103,7 @@ bool q8ta_conv2d_check_4w4c_packed_dim_info(const api::PackedDimInfo& info) { * * Each thread processes a 4Wx4C tile of output elements. */ -utils::uvec3 pick_q8ta_conv2d_global_wg_size( +GlobalWorkGrid pick_q8ta_conv2d_gwg( ComputeGraph* graph, const vkapi::ShaderInfo& shader, const std::vector& args, @@ -70,13 +116,16 @@ utils::uvec3 pick_q8ta_conv2d_global_wg_size( const uint32_t W = graph->size_at(-1, output); const uint32_t H = graph->size_at(-2, output); const uint32_t C = graph->size_at(-3, output); + const uint32_t N = graph->size_at(-4, output); // Each thread processes 4 adjacent width positions and 4 channels (4Wx4C // tile) const uint32_t W4 = utils::div_up_4(W); const uint32_t C4 = utils::div_up_4(C); - return {W4, H, C4}; + return GlobalWorkGrid( + {W4, utils::safe_downcast(static_cast(H) * N), C4}, + kTiledWorkGrid); } /** @@ -86,10 +135,10 @@ utils::uvec3 pick_q8ta_conv2d_global_wg_size( * - {8, 1, 8} for very large tensors: best baseline performance * - {64, 1, 1} for narrow channel dimensions: minimize inactive invocations */ -utils::uvec3 pick_q8ta_conv2d_local_wg_size( +LocalWorkGroup pick_q8ta_conv2d_lwg( ComputeGraph* graph, const vkapi::ShaderInfo& shader, - const utils::uvec3& global_workgroup_size, + const GlobalWorkGrid& gwg, const std::vector& args, const std::vector& resize_args) { (void)shader; @@ -102,34 +151,32 @@ utils::uvec3 pick_q8ta_conv2d_local_wg_size( // For very large tensors (H >= 100 and large x/z), use {8, 1, 8} // This configuration performed best for 128x128 tensors in experiments - if (H >= 100 && global_workgroup_size[0u] >= 24 && - global_workgroup_size[2u] >= 24) { - return {8u, 1u, 8u}; + if (H >= 100 && gwg[0u] >= 24 && gwg[2u] >= 24) { + return LocalWorkGroup(8u, 1u, 8u); } // For medium-sized tensors, use {4, 2, 8} for better height parallelism // This configuration showed +57% improvement on 81x81 tensors - if (global_workgroup_size[0u] >= 4 && global_workgroup_size[1u] >= 2 && - global_workgroup_size[2u] >= 8) { - return {4u, 2u, 8u}; + if (gwg[0u] >= 4 && gwg[1u] >= 2 && gwg[2u] >= 8) { + return LocalWorkGroup(4u, 2u, 8u); } // For tensors with sufficient x and z dimensions, use square configuration - if (global_workgroup_size[0u] >= 6 && global_workgroup_size[2u] >= 6) { - return {8u, 1u, 8u}; + if (gwg[0u] >= 6 && gwg[2u] >= 6) { + return LocalWorkGroup(8u, 1u, 8u); } // If x dimension is very small, bias towards z dimension - if (global_workgroup_size[0u] < 2u) { - return {1u, 1u, 64u}; + if (gwg[0u] < 2u) { + return LocalWorkGroup(1u, 1u, 64u); } // If z dimension is very small, bias towards x dimension - if (global_workgroup_size[2u] < 2u) { - return {64u, 1u, 1u}; + if (gwg[2u] < 2u) { + return LocalWorkGroup(64u, 1u, 1u); } - return {16u, 1u, 4u}; + return LocalWorkGroup(16u, 1u, 4u); } // @@ -192,10 +239,11 @@ ValueRef prepack_quantized_conv2d_weight( storage_type, utils::kWidthPacked); - utils::uvec3 global_wg_size = { - utils::safe_downcast(num_blocks_x), - utils::safe_downcast(num_blocks_y), - 1u}; + const GlobalWorkGrid gwg( + {utils::safe_downcast(num_blocks_x), + utils::safe_downcast(num_blocks_y), + 1u}, + kTiledWorkGrid); std::string kernel_name = "pack_q8_conv2d_weights"; add_storage_type_suffix(kernel_name, storage_type); @@ -203,8 +251,8 @@ ValueRef prepack_quantized_conv2d_weight( graph.prepack_nodes().emplace_back(new PrepackNode( graph, VK_KERNEL_FROM_STR(kernel_name), - global_wg_size, - graph.create_local_wg_size(global_wg_size), + gwg, + graph.create_lwg(gwg), // Inputs and Outputs weight_data, packed_weight, @@ -357,8 +405,8 @@ void add_q8ta_conv2d_node( graph.execute_nodes().emplace_back(new DynamicDispatchNode( graph, VK_KERNEL_FROM_STR(kernel_name), - pick_q8ta_conv2d_global_wg_size, - pick_q8ta_conv2d_local_wg_size, + pick_q8ta_conv2d_gwg, + pick_q8ta_conv2d_lwg, // Inputs and Outputs {{packed_int8_output, vkapi::kWrite}, {{packed_int8_input, @@ -467,33 +515,33 @@ void q8ta_conv2d_general( void q8ta_conv2d(ComputeGraph& graph, const std::vector& args) { const ValueRef input = args.at(0); + const ValueRef kernel_size_ref = args.at(9); const ValueRef groups_ref = args.at(13); const ValueRef output = args.at(15); const int64_t groups = graph.extract_scalar(groups_ref); const int64_t in_channels = graph.size_at(-3, input); const int64_t in_channels_per_group = in_channels / groups; + const int64_t batch = graph.size_at(-4, input); const int64_t H_out = graph.size_at(-2, output); const int64_t W_out = graph.size_at(-1, output); - const int64_t spatial_out = H_out * W_out; - - // Im2col requires input channels per group to be a multiple of 4 - const bool im2col_eligible = in_channels_per_group % 4 == 0; - - bool use_im2col = false; - if (graph.device_is_mali()) { - // On Mali, im2col is faster than the general shader across the board. - use_im2col = im2col_eligible; - } else { - // Default: on Adreno and unknown GPU architectures, im2col is only - // beneficial for ungrouped convolutions with sufficient channel depth or - // small spatial output. For grouped convolutions, the general shader is - // more efficient (0.7-0.95x regression measured on Adreno). - use_im2col = im2col_eligible && groups == 1 && - (in_channels_per_group >= 32 || spatial_out <= 4096); + int64_t flattened_kernel_size; + { + const auto kernel_size = graph.get_int_list(kernel_size_ref); + flattened_kernel_size = utils::align_up_4( + in_channels_per_group * kernel_size->at(0) * kernel_size->at(1)); } + const bool use_im2col = should_use_q8ta_conv2d_im2col( + graph, + batch, + groups, + in_channels_per_group, + flattened_kernel_size, + H_out, + W_out); + if (use_im2col) { q8ta_conv2d_im2col(graph, args); } else { diff --git a/backends/vulkan/runtime/graph/ops/impl/Q8taConv2dDW.cpp b/backends/vulkan/runtime/graph/ops/impl/Q8taConv2dDW.cpp index 914ca1a23ef..182e8d684d2 100644 --- a/backends/vulkan/runtime/graph/ops/impl/Q8taConv2dDW.cpp +++ b/backends/vulkan/runtime/graph/ops/impl/Q8taConv2dDW.cpp @@ -21,7 +21,7 @@ namespace vkcompute { // Shader dispatch utilities // -utils::uvec3 pick_q8ta_conv2d_dw_global_wg_size( +GlobalWorkGrid pick_q8ta_conv2d_dw_gwg( ComputeGraph* graph, const vkapi::ShaderInfo& shader, const std::vector& args, @@ -34,19 +34,22 @@ utils::uvec3 pick_q8ta_conv2d_dw_global_wg_size( const uint32_t W = graph->size_at(-1, output); const uint32_t H = graph->size_at(-2, output); const uint32_t C = graph->size_at(-3, output); + const uint32_t N = graph->size_at(-4, output); // Each thread processes 4 adjacent width positions and 4 channels (4Wx4C // tile) const uint32_t W4 = utils::div_up_4(W); const uint32_t C4 = utils::div_up_4(C); - return {W4, H, C4}; + return GlobalWorkGrid( + {W4, utils::safe_downcast(static_cast(H) * N), C4}, + kTiledWorkGrid); } -utils::uvec3 pick_q8ta_conv2d_dw_local_wg_size( +LocalWorkGroup pick_q8ta_conv2d_dw_lwg( ComputeGraph* graph, const vkapi::ShaderInfo& shader, - const utils::uvec3& global_workgroup_size, + const GlobalWorkGrid& gwg, const std::vector& args, const std::vector& resize_args) { (void)graph; @@ -56,18 +59,18 @@ utils::uvec3 pick_q8ta_conv2d_dw_local_wg_size( // Some inactive invocations are okay; set 6 as the threshold to use the // a square wg size. - if (global_workgroup_size[0u] >= 6 && global_workgroup_size[2u] >= 6) { - return {8u, 1u, 8u}; + if (gwg[0u] >= 6 && gwg[2u] >= 6) { + return LocalWorkGroup(8u, 1u, 8u); } // If channels dim is sufficiently small, then bias towards width dim to // reduce the number of inactive invocations. - if (global_workgroup_size[2u] < 2u) { - return {64u, 1u, 1u}; + if (gwg[2u] < 2u) { + return LocalWorkGroup(64u, 1u, 1u); } - return {16u, 1u, 4u}; + return LocalWorkGroup(16u, 1u, 4u); } -utils::uvec3 int8_conv2d_dw_global_wg_size( +GlobalWorkGrid int8_conv2d_dw_gwg( ComputeGraph* graph, const vkapi::ShaderInfo& shader, const std::vector& args, @@ -81,7 +84,7 @@ utils::uvec3 int8_conv2d_dw_global_wg_size( const uint32_t W4 = utils::div_up_4(W); const uint32_t C4 = utils::div_up_4(C); - return {C4 * W4 * H, 1, 1}; + return graph->create_linear_gwg(C4 * W4 * H); } // @@ -146,10 +149,11 @@ ValueRef prepack_quantized_conv2d_dw_weight( storage_type, utils::kWidthPacked); - utils::uvec3 global_wg_size = { - utils::safe_downcast(num_blocks_OC), - utils::safe_downcast(num_blocks_K), - 1u}; + const GlobalWorkGrid gwg( + {utils::safe_downcast(num_blocks_OC), + utils::safe_downcast(num_blocks_K), + 1u}, + kTiledWorkGrid); std::string kernel_name = "pack_q8_conv2d_dw_weights"; add_storage_type_suffix(kernel_name, storage_type); @@ -157,8 +161,8 @@ ValueRef prepack_quantized_conv2d_dw_weight( graph.prepack_nodes().emplace_back(new PrepackNode( graph, VK_KERNEL_FROM_STR(kernel_name), - global_wg_size, - graph.create_local_wg_size(global_wg_size), + gwg, + graph.create_lwg(gwg), // Inputs and Outputs weight_data, packed_weight, @@ -282,8 +286,8 @@ void add_conv2d_dw_q8ta_q8csw_q8to_4w4c_node( graph.execute_nodes().emplace_back(new DynamicDispatchNode( graph, VK_KERNEL_FROM_STR(kernel_name), - int8_conv2d_dw_global_wg_size, - default_pick_local_wg_size, + int8_conv2d_dw_gwg, + default_pick_lwg, // Inputs and Outputs {{packed_int8_output, vkapi::kWrite}, {{packed_int8_input, @@ -387,8 +391,8 @@ void add_q8ta_conv2d_dw_node( graph.execute_nodes().emplace_back(new DynamicDispatchNode( graph, VK_KERNEL_FROM_STR(kernel_name), - pick_q8ta_conv2d_dw_global_wg_size, - pick_q8ta_conv2d_dw_local_wg_size, + pick_q8ta_conv2d_dw_gwg, + pick_q8ta_conv2d_dw_lwg, // Inputs and Outputs {{packed_int8_output, vkapi::kWrite}, {{packed_int8_input, diff --git a/backends/vulkan/runtime/graph/ops/impl/Q8taConv2dIm2Col.cpp b/backends/vulkan/runtime/graph/ops/impl/Q8taConv2dIm2Col.cpp index 9aa6e7b05d1..e93723c5125 100644 --- a/backends/vulkan/runtime/graph/ops/impl/Q8taConv2dIm2Col.cpp +++ b/backends/vulkan/runtime/graph/ops/impl/Q8taConv2dIm2Col.cpp @@ -22,7 +22,7 @@ namespace vkcompute { // Shader dispatch utilities // -utils::uvec3 pick_q8ta_im2col_global_wg_size( +GlobalWorkGrid pick_q8ta_im2col_gwg( ComputeGraph* graph, const vkapi::ShaderInfo& shader, const std::vector& args, @@ -32,31 +32,32 @@ utils::uvec3 pick_q8ta_im2col_global_wg_size( const ValueRef im2col_output = args.at(0).refs.at(0); - std::vector im2col_sizes = graph->sizes_of(im2col_output); - const uint32_t K = utils::safe_downcast(im2col_sizes[0]); - const uint32_t H = utils::safe_downcast(im2col_sizes[1]); - const uint32_t W = utils::safe_downcast(im2col_sizes[2]); + const uint32_t N = graph->size_at(-4, im2col_output); + const uint32_t K = graph->size_at(-3, im2col_output); + const uint32_t H = graph->size_at(-2, im2col_output); + const uint32_t W = graph->size_at(-1, im2col_output); const uint32_t K4 = utils::div_up_4(K); const uint32_t W4 = utils::div_up_4(W); // Each thread handles one 4x4 block in the output - return {K4 * W4 * H, 1, 1}; + return graph->create_linear_gwg( + utils::safe_downcast(static_cast(K4) * W4 * H * N)); } -utils::uvec3 pick_q8ta_im2col_local_wg_size( +LocalWorkGroup pick_q8ta_im2col_lwg( ComputeGraph* graph, const vkapi::ShaderInfo& shader, - const utils::uvec3& global_workgroup_size, + const GlobalWorkGrid& gwg, const std::vector& args, const std::vector& resize_args) { (void)graph; (void)shader; (void)args; (void)resize_args; - (void)global_workgroup_size; + (void)gwg; - return {64, 1, 1}; + return LocalWorkGroup(64u, 1u, 1u); } // @@ -70,6 +71,7 @@ std::vector calculate_q8ta_im2col_sizes( const ValueRef& kernel_size, const ValueRef& groups) { std::vector in_sizes = graph->sizes_of(input); + const int64_t batch = utils::val_at(-4, in_sizes); const int64_t in_channels = utils::val_at(-3, in_sizes); std::vector out_sizes = graph->sizes_of(output); @@ -93,7 +95,7 @@ std::vector calculate_q8ta_im2col_sizes( const int64_t W = utils::align_up_4(out_width); const int64_t H = out_height; - return {K, H, W}; + return {batch, K, H, W}; } // @@ -102,7 +104,7 @@ std::vector calculate_q8ta_im2col_sizes( // resize_args = { input, kernel_size, stride, padding, dilation, groups } // -// The im2col scratch tensor is [K, H_out, align_up_4(W_out)] where K (the +// The im2col scratch tensor is [N, K, H_out, align_up_4(W_out)] where K (the // flattened conv window, channel/kernel-derived) is shape-independent and // H_out/W_out are the conv output spatial dims. The downstream PW GEMM that // consumes this scratch is resized separately (it preserves H/W). Without this, @@ -122,6 +124,7 @@ void resize_q8ta_im2col_node( const ValueRef groups = resize_args.at(5); const std::vector in_sizes = graph->sizes_of(in); + const int64_t batch = utils::val_at(-4, in_sizes); // Conv output H/W from the current input. const std::vector out_hw = calc_out_sizes_hw( @@ -146,7 +149,7 @@ void resize_q8ta_im2col_node( const int64_t K = flattened_kernel_len * groups_val; const int64_t W = utils::align_up_4(out_width); - graph->virtual_resize(im2col_out, {K, out_height, W}); + graph->virtual_resize(im2col_out, {batch, K, out_height, W}); } // @@ -212,8 +215,8 @@ void add_q8ta_im2col_node( graph.execute_nodes().emplace_back(new DynamicDispatchNode( graph, VK_KERNEL_FROM_STR(kernel_name), - pick_q8ta_im2col_global_wg_size, - pick_q8ta_im2col_local_wg_size, + pick_q8ta_im2col_gwg, + pick_q8ta_im2col_lwg, // Inputs and Outputs {{packed_int8_im2col, vkapi::kWrite}, {packed_int8_input, vkapi::kRead}}, // Shader params buffers diff --git a/backends/vulkan/runtime/graph/ops/impl/Q8taConv2dPW.cpp b/backends/vulkan/runtime/graph/ops/impl/Q8taConv2dPW.cpp index 4fb7f0fa775..ee234319e8c 100644 --- a/backends/vulkan/runtime/graph/ops/impl/Q8taConv2dPW.cpp +++ b/backends/vulkan/runtime/graph/ops/impl/Q8taConv2dPW.cpp @@ -20,7 +20,7 @@ namespace vkcompute { // Shader dispatch utilities // -utils::uvec3 pick_q8ta_conv2d_pw_global_wg_size( +GlobalWorkGrid pick_q8ta_conv2d_pw_gwg( ComputeGraph* graph, const vkapi::ShaderInfo& shader, const std::vector& args, @@ -33,6 +33,7 @@ utils::uvec3 pick_q8ta_conv2d_pw_global_wg_size( const uint32_t W = graph->size_at(-1, output); const uint32_t H = graph->size_at(-2, output); const uint32_t C = graph->size_at(-3, output); + const uint32_t N = graph->size_at(-4, output); // Each thread covers a 4-width x 4-channel output block. // Tile constants must match TILE_M4 / TILE_N4 in q8ta_conv2d_pw.glsl. @@ -45,25 +46,28 @@ utils::uvec3 pick_q8ta_conv2d_pw_global_wg_size( // Global workgroup size: // x = output channels / (TILE_N4 * 4) = C4 / TILE_N4 = C4 // y = width / (TILE_M4 * 4) = W4 / TILE_M4 = W4 - // z = height - return {utils::div_up(C4, TILE_N4), utils::div_up(W4, TILE_M4), H}; + // z = height * batch + return GlobalWorkGrid( + {utils::div_up(C4, TILE_N4), + utils::div_up(W4, TILE_M4), + utils::safe_downcast(static_cast(H) * N)}, + kTiledWorkGrid); } -utils::uvec3 pick_q8ta_conv2d_pw_local_wg_size( +LocalWorkGroup pick_q8ta_conv2d_pw_lwg( ComputeGraph* graph, const vkapi::ShaderInfo& shader, - const utils::uvec3& global_workgroup_size, + const GlobalWorkGrid& gwg, const std::vector& args, const std::vector& resize_args) { - return pick_hw_square_wg_size( - graph, shader, global_workgroup_size, args, resize_args); + return pick_xy_square_lwg(graph, shader, gwg, args, resize_args); } // // 4W4C shader dispatch utilities // -utils::uvec3 pick_q8ta_conv2d_pw_4w4c_global_wg_size( +GlobalWorkGrid pick_q8ta_conv2d_pw_4w4c_gwg( ComputeGraph* graph, const vkapi::ShaderInfo& shader, const std::vector& args, @@ -91,17 +95,18 @@ utils::uvec3 pick_q8ta_conv2d_pw_4w4c_global_wg_size( // x = output channels / (TILE_N4 * 4) = C4 / TILE_N4 // y = width / (TILE_M4 * 4) = W4 / TILE_M4 // z = height - return {utils::div_up(C4, TILE_N4), utils::div_up(W4, TILE_M4), H}; + return GlobalWorkGrid( + {utils::div_up(C4, TILE_N4), utils::div_up(W4, TILE_M4), H}, + kTiledWorkGrid); } -utils::uvec3 pick_q8ta_conv2d_pw_4w4c_local_wg_size( +LocalWorkGroup pick_q8ta_conv2d_pw_4w4c_lwg( ComputeGraph* graph, const vkapi::ShaderInfo& shader, - const utils::uvec3& global_workgroup_size, + const GlobalWorkGrid& gwg, const std::vector& args, const std::vector& resize_args) { - return pick_hw_square_wg_size( - graph, shader, global_workgroup_size, args, resize_args); + return pick_xy_square_lwg(graph, shader, gwg, args, resize_args); } // @@ -155,10 +160,11 @@ ValueRef prepack_quantized_conv2d_pw_weight( storage_type, utils::kWidthPacked); - utils::uvec3 global_wg_size = { - utils::safe_downcast(num_blocks_x), - utils::safe_downcast(num_blocks_y), - 1u}; + const GlobalWorkGrid gwg( + {utils::safe_downcast(num_blocks_x), + utils::safe_downcast(num_blocks_y), + 1u}, + kTiledWorkGrid); std::string kernel_name = "pack_q8_conv2d_weights"; add_storage_type_suffix(kernel_name, storage_type); @@ -166,8 +172,8 @@ ValueRef prepack_quantized_conv2d_pw_weight( graph.prepack_nodes().emplace_back(new PrepackNode( graph, VK_KERNEL_FROM_STR(kernel_name), - global_wg_size, - graph.create_local_wg_size(global_wg_size), + gwg, + graph.create_lwg(gwg), // Inputs and Outputs weight_data, packed_weight, @@ -338,8 +344,8 @@ void add_q8ta_conv2d_pw_node( graph.execute_nodes().emplace_back(new DynamicDispatchNode( graph, VK_KERNEL_FROM_STR(kernel_name), - pick_q8ta_conv2d_pw_global_wg_size, - pick_q8ta_conv2d_pw_local_wg_size, + pick_q8ta_conv2d_pw_gwg, + pick_q8ta_conv2d_pw_lwg, {{packed_int8_output, vkapi::kWrite}, {{packed_int8_input, packed_weight, diff --git a/backends/vulkan/runtime/graph/ops/impl/Q8taConv2dTransposed.cpp b/backends/vulkan/runtime/graph/ops/impl/Q8taConv2dTransposed.cpp index 7e3c4166e3c..91d745fb1fa 100644 --- a/backends/vulkan/runtime/graph/ops/impl/Q8taConv2dTransposed.cpp +++ b/backends/vulkan/runtime/graph/ops/impl/Q8taConv2dTransposed.cpp @@ -63,7 +63,7 @@ void resize_q8ta_conv2d_transposed_node( // stride-alignment patterns). Keeping local_y=1 ensures all threads in a // workgroup process the same height row, maximizing branch coherence. -utils::uvec3 pick_q8ta_conv2d_transposed_global_wg_size( +GlobalWorkGrid pick_q8ta_conv2d_transposed_gwg( ComputeGraph* graph, const vkapi::ShaderInfo& shader, const std::vector& args, @@ -80,13 +80,13 @@ utils::uvec3 pick_q8ta_conv2d_transposed_global_wg_size( const uint32_t W4 = utils::div_up_4(W); const uint32_t C4 = utils::div_up_4(C); - return {W4, H, C4}; + return GlobalWorkGrid({W4, H, C4}, kTiledWorkGrid); } -utils::uvec3 pick_q8ta_conv2d_transposed_local_wg_size( +LocalWorkGroup pick_q8ta_conv2d_transposed_lwg( ComputeGraph* graph, const vkapi::ShaderInfo& shader, - const utils::uvec3& global_workgroup_size, + const GlobalWorkGrid& gwg, const std::vector& args, const std::vector& resize_args) { (void)shader; @@ -95,16 +95,16 @@ utils::uvec3 pick_q8ta_conv2d_transposed_local_wg_size( (void)args; // Always keep local_y=1 to avoid branch divergence between height rows. - if (global_workgroup_size[0u] >= 6 && global_workgroup_size[2u] >= 6) { - return {8u, 1u, 8u}; + if (gwg[0u] >= 6 && gwg[2u] >= 6) { + return LocalWorkGroup(8u, 1u, 8u); } - if (global_workgroup_size[0u] < 2u) { - return {1u, 1u, 64u}; + if (gwg[0u] < 2u) { + return LocalWorkGroup(1u, 1u, 64u); } - if (global_workgroup_size[2u] < 2u) { - return {64u, 1u, 1u}; + if (gwg[2u] < 2u) { + return LocalWorkGroup(64u, 1u, 1u); } - return {16u, 1u, 4u}; + return LocalWorkGroup(16u, 1u, 4u); } void add_q8ta_conv2d_transposed_node( @@ -200,8 +200,8 @@ void add_q8ta_conv2d_transposed_node( graph.execute_nodes().emplace_back(new DynamicDispatchNode( graph, VK_KERNEL_FROM_STR(kernel_name), - pick_q8ta_conv2d_transposed_global_wg_size, - pick_q8ta_conv2d_transposed_local_wg_size, + pick_q8ta_conv2d_transposed_gwg, + pick_q8ta_conv2d_transposed_lwg, // Inputs and Outputs {{packed_int8_output, vkapi::kWrite}, {{packed_int8_input, diff --git a/backends/vulkan/runtime/graph/ops/impl/Q8taLinear.cpp b/backends/vulkan/runtime/graph/ops/impl/Q8taLinear.cpp index 92daf9d8ac5..683b7b5ce07 100644 --- a/backends/vulkan/runtime/graph/ops/impl/Q8taLinear.cpp +++ b/backends/vulkan/runtime/graph/ops/impl/Q8taLinear.cpp @@ -29,7 +29,7 @@ bool q8ta_linear_check_packed_dim_info(const api::PackedDimInfo& info) { // Workgroup size selection // -utils::uvec3 q8ta_linear_global_wg_size( +GlobalWorkGrid q8ta_linear_gwg( ComputeGraph* graph, const vkapi::ShaderInfo& shader, const std::vector& args, @@ -50,17 +50,16 @@ utils::uvec3 q8ta_linear_global_wg_size( const uint32_t num_N_tiles = utils::div_up(N, N_per_tile); const uint32_t num_M_tiles = utils::div_up(M, M_per_tile); - return {num_N_tiles, num_M_tiles, 1}; + return GlobalWorkGrid({num_N_tiles, num_M_tiles, 1u}, kTiledWorkGrid); } -utils::uvec3 q8ta_linear_local_wg_size( +LocalWorkGroup q8ta_linear_lwg( ComputeGraph* graph, const vkapi::ShaderInfo& shader, - const utils::uvec3& global_workgroup_size, + const GlobalWorkGrid& gwg, const std::vector& args, const std::vector& resize_args) { - return pick_hw_square_wg_size( - graph, shader, global_workgroup_size, args, resize_args); + return pick_xy_square_lwg(graph, shader, gwg, args, resize_args); } // @@ -144,8 +143,8 @@ void add_q8ta_linear_node( graph.execute_nodes().emplace_back(new DynamicDispatchNode( graph, VK_KERNEL_FROM_STR(kernel_name), - q8ta_linear_global_wg_size, - q8ta_linear_local_wg_size, + q8ta_linear_gwg, + q8ta_linear_lwg, // Inputs and Outputs {{packed_int8_output, vkapi::kWrite}, {{packed_int8_input, diff --git a/backends/vulkan/runtime/graph/ops/impl/Q8taLinearGemv.cpp b/backends/vulkan/runtime/graph/ops/impl/Q8taLinearGemv.cpp index 2885ad86f35..7a63d33d546 100644 --- a/backends/vulkan/runtime/graph/ops/impl/Q8taLinearGemv.cpp +++ b/backends/vulkan/runtime/graph/ops/impl/Q8taLinearGemv.cpp @@ -30,7 +30,7 @@ static bool q8ta_linear_gemv_check_packed_dim_info( // Workgroup size selection // -utils::uvec3 q8ta_linear_gemv_global_wg_size( +GlobalWorkGrid q8ta_linear_gemv_gwg( ComputeGraph* graph, const vkapi::ShaderInfo& shader, const std::vector& args, @@ -47,23 +47,11 @@ utils::uvec3 q8ta_linear_gemv_global_wg_size( const uint32_t N_per_tile = 8; const uint32_t num_N_tiles = utils::div_up(N, N_per_tile); - return {num_N_tiles, 1, 1}; -} - -utils::uvec3 q8ta_linear_gemv_local_wg_size( - ComputeGraph* graph, - const vkapi::ShaderInfo& shader, - const utils::uvec3& global_workgroup_size, - const std::vector& args, - const std::vector& resize_args) { - (void)graph; - (void)shader; - (void)global_workgroup_size; - (void)args; - (void)resize_args; - - // Cooperative algorithm: 64 threads share the K reduction - return {1, 1, 64}; + const LocalWorkGroup lwg(1u, 1u, 64u); + GlobalWorkGrid gwg({num_N_tiles, 1u, 1u}, kLinearWorkGrid); + gwg.wrap_linear_dispatch( + graph->context()->adapter_ptr()->max_compute_workgroup_count(), lwg); + return gwg; } // @@ -120,8 +108,8 @@ void add_q8ta_linear_gemv_node( graph.execute_nodes().emplace_back(new DynamicDispatchNode( graph, VK_KERNEL_FROM_STR(kernel_name), - q8ta_linear_gemv_global_wg_size, - q8ta_linear_gemv_local_wg_size, + q8ta_linear_gemv_gwg, + pick_required_lwg, // Inputs and Outputs {{packed_int8_output, vkapi::kWrite}, {{packed_int8_input, diff --git a/backends/vulkan/runtime/graph/ops/impl/Q8taPixelShuffle.cpp b/backends/vulkan/runtime/graph/ops/impl/Q8taPixelShuffle.cpp index 74712654fd4..45748ece4a7 100644 --- a/backends/vulkan/runtime/graph/ops/impl/Q8taPixelShuffle.cpp +++ b/backends/vulkan/runtime/graph/ops/impl/Q8taPixelShuffle.cpp @@ -42,7 +42,7 @@ void resize_q8ta_pixel_shuffle_node( // Global wg picker: one thread per output int32 word. For a channels-packed // int8x4 output with channel block size 4, the number of output int words is // N * div_up_4(C_out) * H_out * W_out. -utils::uvec3 pick_q8ta_pixel_shuffle_global_wg( +GlobalWorkGrid pick_q8ta_pixel_shuffle_gwg( ComputeGraph* graph, const vkapi::ShaderInfo& shader, const std::vector& args, @@ -58,23 +58,23 @@ utils::uvec3 pick_q8ta_pixel_shuffle_global_wg( const int64_t c_words = utils::div_up(C, int64_t(4)); const uint32_t total_words = utils::safe_downcast(N * c_words * H * W); - return {total_words, 1u, 1u}; + return graph->create_linear_gwg(total_words); } -utils::uvec3 pick_q8ta_pixel_shuffle_local_wg( +LocalWorkGroup pick_q8ta_pixel_shuffle_lwg( ComputeGraph* graph, const vkapi::ShaderInfo& shader, - const utils::uvec3& global_workgroup_size, + const GlobalWorkGrid& gwg, const std::vector& args, const std::vector& resize_args) { (void)graph; (void)shader; - (void)global_workgroup_size; + (void)gwg; (void)args; (void)resize_args; // Linear (1D) dispatch: a flat 64-wide workgroup matches the pattern used - // by pick_square_local_wg_with_block_config in the linear case. - return {64u, 1u, 1u}; + // by pick_square_lwg_with_block_config in the linear case. + return LocalWorkGroup(64u, 1u, 1u); } } // namespace @@ -178,8 +178,8 @@ void add_q8ta_pixel_shuffle_node( graph.execute_nodes().emplace_back(new DynamicDispatchNode( graph, VK_KERNEL_FROM_STR(kernel_name), - pick_q8ta_pixel_shuffle_global_wg, - pick_q8ta_pixel_shuffle_local_wg, + pick_q8ta_pixel_shuffle_gwg, + pick_q8ta_pixel_shuffle_lwg, // Inputs and Outputs {{packed_int8_output, vkapi::kWrite}, {packed_int8_input, vkapi::kRead}}, // Shader params buffers diff --git a/backends/vulkan/runtime/graph/ops/impl/Q8taQuantizeDequantize.cpp b/backends/vulkan/runtime/graph/ops/impl/Q8taQuantizeDequantize.cpp index fb0ffcab14c..cbf9d750a5c 100644 --- a/backends/vulkan/runtime/graph/ops/impl/Q8taQuantizeDequantize.cpp +++ b/backends/vulkan/runtime/graph/ops/impl/Q8taQuantizeDequantize.cpp @@ -66,7 +66,7 @@ void add_q8ta_quantize_node( const BlockConfig inp_block_config = create_block_config_from_other(graph, fp_input, outp_block_config); - // Cast block config to ValueRef for pick_*_global_wg_with_block_config + // Cast block config to ValueRef for pick_*_gwg_with_block_config // Use inp_block_config since shader uses inp_block_config for indexing const ValueRef block_config_ref = static_cast(inp_block_config.as_packed_int()); @@ -74,15 +74,15 @@ void add_q8ta_quantize_node( // Choose dispatch function based on FP input storage type: // - Buffer: use linear dispatch (better performance) // - Texture: use extents-style 3D dispatch (better performance) - auto pick_global_wg_size = (inp_storage == utils::kBuffer) - ? pick_linear_global_wg_with_block_config - : pick_extents_global_wg_with_block_config; + auto pick_gwg = (inp_storage == utils::kBuffer) + ? pick_linear_gwg_with_block_config + : pick_extents_gwg_with_block_config; graph.execute_nodes().emplace_back(new DynamicDispatchNode( graph, VK_KERNEL_FROM_STR(kernel_name), - pick_global_wg_size, - pick_square_local_wg_with_block_config, + pick_gwg, + pick_square_lwg_with_block_config, // Inputs and Outputs {{packed_int8_output, vkapi::kWrite}, {fp_input, vkapi::kRead}}, // Shader params buffers @@ -138,7 +138,7 @@ void add_q8ta_dequantize_node( const BlockConfig inp_block_config = create_block_config_from_other( graph, packed_int8_input, outp_block_config); - // Cast block config to ValueRef for pick_*_global_wg_with_block_config + // Cast block config to ValueRef for pick_*_gwg_with_block_config // Use inp_block_config since shader uses inp_block_config for indexing const ValueRef block_config_ref = static_cast(inp_block_config.as_packed_int()); @@ -146,15 +146,15 @@ void add_q8ta_dequantize_node( // Choose dispatch function based on FP output storage type: // - Buffer: use linear dispatch (better performance) // - Texture: use extents-style 3D dispatch (better performance) - auto pick_global_wg_size = (outp_storage == utils::kBuffer) - ? pick_linear_global_wg_with_block_config - : pick_extents_global_wg_with_block_config; + auto pick_gwg = (outp_storage == utils::kBuffer) + ? pick_linear_gwg_with_block_config + : pick_extents_gwg_with_block_config; graph.execute_nodes().emplace_back(new DynamicDispatchNode( graph, VK_KERNEL_FROM_STR(kernel_name), - pick_global_wg_size, - pick_square_local_wg_with_block_config, + pick_gwg, + pick_square_lwg_with_block_config, // Inputs and Outputs {{fp_output, vkapi::kWrite}, {packed_int8_input, vkapi::kRead}}, // Shader params buffers diff --git a/backends/vulkan/runtime/graph/ops/impl/Q8taUnary.cpp b/backends/vulkan/runtime/graph/ops/impl/Q8taUnary.cpp index f8b606f3dfa..f4d6c15ef60 100644 --- a/backends/vulkan/runtime/graph/ops/impl/Q8taUnary.cpp +++ b/backends/vulkan/runtime/graph/ops/impl/Q8taUnary.cpp @@ -75,8 +75,8 @@ void add_q8ta_unary_node( graph.execute_nodes().emplace_back(new DynamicDispatchNode( graph, VK_KERNEL_FROM_STR(kernel_name), - pick_linear_global_wg_with_block_config, - pick_square_local_wg_with_block_config, + pick_linear_gwg_with_block_config, + pick_square_lwg_with_block_config, // Inputs and Outputs {{packed_int8_output, vkapi::kWrite}, {packed_int8_input, vkapi::kRead}}, // Shader params buffers diff --git a/backends/vulkan/runtime/graph/ops/impl/QuantizeDequantize.cpp b/backends/vulkan/runtime/graph/ops/impl/QuantizeDequantize.cpp index 98f97eab572..97c939dcabf 100644 --- a/backends/vulkan/runtime/graph/ops/impl/QuantizeDequantize.cpp +++ b/backends/vulkan/runtime/graph/ops/impl/QuantizeDequantize.cpp @@ -44,7 +44,7 @@ std::tuple get_quantized_input_num_blocks( return std::make_tuple(num_blocks_M, num_blocks_K); } -utils::uvec3 quantize_and_pack_4h4w_global_wg_size( +GlobalWorkGrid quantize_and_pack_4h4w_gwg( ComputeGraph* graph, const vkapi::ShaderInfo& shader, const std::vector& args, @@ -54,10 +54,11 @@ utils::uvec3 quantize_and_pack_4h4w_global_wg_size( std::tie(num_blocks_M, num_blocks_K) = get_quantized_input_num_blocks(*graph, input); - return { - utils::safe_downcast(num_blocks_K), - utils::safe_downcast(num_blocks_M), - 1u}; + return GlobalWorkGrid( + {utils::safe_downcast(num_blocks_K), + utils::safe_downcast(num_blocks_M), + 1u}, + kTiledWorkGrid); } vkapi::ShaderInfo pick_quantize_and_pack_4h4w_with_group_sums_shader( @@ -87,7 +88,7 @@ vkapi::ShaderInfo pick_quantize_and_pack_4h4w_with_group_sums_shader( return VK_KERNEL_FROM_STR(shader_name); } -utils::uvec3 pick_quantize_and_pack_4h4w_with_group_sums_global_wg_size( +GlobalWorkGrid pick_quantize_and_pack_4h4w_with_group_sums_gwg( ComputeGraph* graph, const vkapi::ShaderInfo& shader, const std::vector& args, @@ -98,8 +99,12 @@ utils::uvec3 pick_quantize_and_pack_4h4w_with_group_sums_global_wg_size( // rationale for this is that gemv is a memory bound operation and may not // necessarily benefit from quantizing the input and computing with integer // accumulation. + const LocalWorkGroup required_lwg = + shader.kernel_name.find("o4w16") != std::string::npos + ? LocalWorkGroup(4u, 1u, 16u) + : LocalWorkGroup(2u, 1u, 32u); if (is_gemv(graph, fp_input)) { - return {0u, 0u, 0u}; + return GlobalWorkGrid({0u, 0u, 0u}, kTiledWorkGrid, required_lwg); } const ValueRef group_size = resize_args.at(0); @@ -112,37 +117,12 @@ utils::uvec3 pick_quantize_and_pack_4h4w_with_group_sums_global_wg_size( const int64_t num_groups = num_blocks_K / blocks_per_group; - return { - utils::safe_downcast(num_groups), - utils::safe_downcast(num_blocks_M), - 1u}; -} - -utils::uvec3 pick_quantize_and_pack_4h4w_with_group_sums_local_wg_size( - ComputeGraph* graph, - const vkapi::ShaderInfo& shader, - const utils::uvec3& global_workgroup_size, - const std::vector& args, - const std::vector& resize_args) { - (void)shader; - (void)resize_args; - - const ValueRef fp_input = args.at(1).refs.at(0); - // For gemv, skip the quantize input step since the quantized linear is - // computed as a weight only quantized linear operation. - if (is_gemv(graph, fp_input)) { - return {1u, 1u, 1u}; - } - - uint32_t groups_per_wg = 2u; - uint32_t workers_per_group = 32u; - - if (shader.kernel_name.find("o4w16") != std::string::npos) { - groups_per_wg = 4u; - workers_per_group = 16u; - } - - return {groups_per_wg, 1u, workers_per_group}; + return GlobalWorkGrid( + {utils::safe_downcast(num_groups), + utils::safe_downcast(num_blocks_M), + 1u}, + kTiledWorkGrid, + required_lwg); } // @@ -184,8 +164,8 @@ void add_quantize_and_pack_4h4w_node( graph.execute_nodes().emplace_back(new DynamicDispatchNode( graph, VK_KERNEL_FROM_STR(shader_name), - quantize_and_pack_4h4w_global_wg_size, - default_pick_local_wg_size, + quantize_and_pack_4h4w_gwg, + default_pick_lwg, // Inputs and Outputs {{packed_int_input, vkapi::kWrite}, {fp_input, vkapi::kRead}}, // Shader params buffers @@ -222,8 +202,8 @@ void add_quantize_and_pack_4h4w_with_group_sums_node( graph.execute_nodes().emplace_back(new DynamicDispatchNode( graph, pick_quantize_and_pack_4h4w_with_group_sums_shader, - pick_quantize_and_pack_4h4w_with_group_sums_global_wg_size, - pick_quantize_and_pack_4h4w_with_group_sums_local_wg_size, + pick_quantize_and_pack_4h4w_with_group_sums_gwg, + pick_required_lwg, // Inputs and Outputs {{{packed_int_input, int_input_sums}, vkapi::kWrite}, {{fp_input, packed_input_scales, packed_input_zps}, vkapi::kRead}}, @@ -241,7 +221,7 @@ void add_quantize_and_pack_4h4w_with_group_sums_node( // Dispatch utilities (Conv2d) // -utils::uvec3 pick_quantize_and_pack_4w4c_global_wg_size( +GlobalWorkGrid pick_quantize_and_pack_4w4c_gwg( ComputeGraph* graph, const vkapi::ShaderInfo& shader, const std::vector& args, @@ -255,10 +235,10 @@ utils::uvec3 pick_quantize_and_pack_4w4c_global_wg_size( const uint32_t W4 = utils::div_up_4(W); const uint32_t C4 = utils::div_up_4(C); - return {W4, H, C4}; + return GlobalWorkGrid({W4, H, C4}, kTiledWorkGrid); } -utils::uvec3 pick_unpack_4w4c_and_dequantize_global_wg_size( +GlobalWorkGrid pick_unpack_4w4c_and_dequantize_gwg( ComputeGraph* graph, const vkapi::ShaderInfo& shader, const std::vector& args, @@ -272,7 +252,7 @@ utils::uvec3 pick_unpack_4w4c_and_dequantize_global_wg_size( const uint32_t W4 = utils::div_up_4(W); const uint32_t C4 = utils::div_up_4(C); - return {W4, H, C4}; + return GlobalWorkGrid({W4, H, C4}, kTiledWorkGrid); } // @@ -307,8 +287,8 @@ void add_quantize_and_pack_4w4c_node( graph.execute_nodes().emplace_back(new DynamicDispatchNode( graph, VK_KERNEL_FROM_STR(kernel_name), - pick_quantize_and_pack_4w4c_global_wg_size, - pick_wc_square_wg_size, + pick_quantize_and_pack_4w4c_gwg, + pick_xz_square_lwg, // Inputs and Outputs {{packed_int8_input, vkapi::kWrite}, {fp_input, vkapi::kRead}}, // Shader params buffers @@ -351,8 +331,8 @@ void add_unpack_4w4c_and_dequantize_node( graph.execute_nodes().emplace_back(new DynamicDispatchNode( graph, VK_KERNEL_FROM_STR(kernel_name), - pick_unpack_4w4c_and_dequantize_global_wg_size, - default_pick_local_wg_size, + pick_unpack_4w4c_and_dequantize_gwg, + default_pick_lwg, // Inputs and Outputs {{fp_output, vkapi::kWrite}, {packed_int8_output, vkapi::kRead}}, // Shader params buffers diff --git a/backends/vulkan/runtime/graph/ops/impl/QuantizedConvolution.cpp b/backends/vulkan/runtime/graph/ops/impl/QuantizedConvolution.cpp index ebc276ee347..30415cd4f08 100644 --- a/backends/vulkan/runtime/graph/ops/impl/QuantizedConvolution.cpp +++ b/backends/vulkan/runtime/graph/ops/impl/QuantizedConvolution.cpp @@ -262,7 +262,7 @@ std::vector calculate_output_im2col_sizes( // Shader dispatch utilities // -utils::uvec3 im2col_global_wg_size( +GlobalWorkGrid im2col_gwg( ComputeGraph* graph, const vkapi::ShaderInfo& shader, const std::vector& args, @@ -281,10 +281,10 @@ utils::uvec3 im2col_global_wg_size( const uint32_t K4 = utils::div_up(K, 4u); const uint32_t M4 = utils::div_up(M, 4u); - return {K4, M4, 1}; + return GlobalWorkGrid({K4, M4, 1u}, kTiledWorkGrid); } -utils::uvec3 im2col_packed_int8_global_wg_size( +GlobalWorkGrid im2col_packed_int8_gwg( ComputeGraph* graph, const vkapi::ShaderInfo& shader, const std::vector& args, @@ -299,19 +299,24 @@ utils::uvec3 im2col_packed_int8_global_wg_size( const uint32_t K4 = utils::div_up(K, 4u); const uint32_t W4 = utils::div_up(W, 4u); - return {K4 * W4 * H, 1, 1}; + return graph->create_linear_gwg(K4 * W4 * H); } -utils::uvec3 im2col_packed_int8_local_wg_size( +LocalWorkGroup im2col_packed_int8_lwg( ComputeGraph* graph, const vkapi::ShaderInfo& shader, - const utils::uvec3& global_workgroup_size, + const GlobalWorkGrid& gwg, const std::vector& args, const std::vector& resize_args) { - return {64, 1, 1}; + (void)graph; + (void)shader; + (void)gwg; + (void)args; + (void)resize_args; + return LocalWorkGroup(64u, 1u, 1u); } -utils::uvec3 col2im_global_wg_size( +GlobalWorkGrid col2im_gwg( ComputeGraph* graph, const vkapi::ShaderInfo& shader, const std::vector& args, @@ -327,10 +332,10 @@ utils::uvec3 col2im_global_wg_size( const uint32_t N4 = utils::div_up(N, 4u); const uint32_t M4 = utils::div_up(M, 4u); - return {N4, M4, 1}; + return GlobalWorkGrid({N4, M4, 1u}, kTiledWorkGrid); } -utils::uvec3 pick_static_quantized_conv2d_global_wg_size( +GlobalWorkGrid pick_static_quantized_conv2d_gwg( ComputeGraph* graph, const vkapi::ShaderInfo& shader, const std::vector& args, @@ -351,17 +356,16 @@ utils::uvec3 pick_static_quantized_conv2d_global_wg_size( const uint32_t num_W_tiles = utils::div_up(W, W_per_tile); const uint32_t num_C_tiles = utils::div_up(C, C_per_tile); - return {num_C_tiles, num_W_tiles, H}; + return GlobalWorkGrid({num_C_tiles, num_W_tiles, H}, kTiledWorkGrid); } -utils::uvec3 pick_static_quantized_conv2d_local_wg_size( +LocalWorkGroup pick_static_quantized_conv2d_lwg( ComputeGraph* graph, const vkapi::ShaderInfo& shader, - const utils::uvec3& global_workgroup_size, + const GlobalWorkGrid& gwg, const std::vector& args, const std::vector& resize_args) { - return pick_hw_square_wg_size( - graph, shader, global_workgroup_size, args, resize_args); + return pick_xy_square_lwg(graph, shader, gwg, args, resize_args); } // @@ -447,8 +451,8 @@ void add_input_im2col_node( graph.execute_nodes().emplace_back(new DynamicDispatchNode( graph, VK_KERNEL_FROM_STR(kernel_name), - im2col_global_wg_size, - default_pick_local_wg_size, + im2col_gwg, + default_pick_lwg, // Inputs and Outputs {{input_im2col, vkapi::kWrite}, {input_image, vkapi::kRead}}, // Shader params buffers @@ -500,8 +504,8 @@ void add_input_im2col_packed_int8_node( graph.execute_nodes().emplace_back(new DynamicDispatchNode( graph, VK_KERNEL_FROM_STR(kernel_name), - im2col_packed_int8_global_wg_size, - im2col_packed_int8_local_wg_size, + im2col_packed_int8_gwg, + im2col_packed_int8_lwg, // Inputs and Outputs {{input_im2col, vkapi::kWrite}, {input, vkapi::kRead}}, // Shader params buffers @@ -565,8 +569,8 @@ void add_quantize_and_pack_im2col_node( graph.execute_nodes().emplace_back(new DynamicDispatchNode( graph, VK_KERNEL_FROM_STR(kernel_name), - im2col_global_wg_size, - default_pick_local_wg_size, + im2col_gwg, + default_pick_lwg, // Inputs and Outputs {{input_int_im2col, vkapi::kWrite}, {input_image, vkapi::kRead}}, // Shader params buffers @@ -633,8 +637,8 @@ void add_conv2d_q8csw_linear_node( graph.execute_nodes().emplace_back(new DynamicDispatchNode( graph, VK_KERNEL_FROM_STR(kernel_name), - col2im_global_wg_size, - quantized_linear_local_wg_size, + col2im_gwg, + quantized_linear_lwg, // Inputs and Outputs {{output_image, vkapi::kWrite}, {{input_im2col, packed_weight, packed_weight_scales, packed_bias}, @@ -715,8 +719,8 @@ void add_conv2d_q8ta_q8csw_linear_node( graph.execute_nodes().emplace_back(new DynamicDispatchNode( graph, VK_KERNEL_FROM_STR(kernel_name), - col2im_global_wg_size, - quantized_linear_local_wg_size, + col2im_gwg, + quantized_linear_lwg, // Inputs and Outputs {{output_image, vkapi::kWrite}, {{input_int_im2col, @@ -804,8 +808,8 @@ void add_conv2d_q8ta_q8csw_q8to_4w4c_node( graph.execute_nodes().emplace_back(new DynamicDispatchNode( graph, VK_KERNEL_FROM_STR(kernel_name), - pick_static_quantized_conv2d_global_wg_size, - pick_static_quantized_conv2d_local_wg_size, + pick_static_quantized_conv2d_gwg, + pick_static_quantized_conv2d_lwg, // Inputs and Outputs {{packed_int8_output, vkapi::kWrite}, {{packed_int8_input_im2col, diff --git a/backends/vulkan/runtime/graph/ops/impl/QuantizedLinear.cpp b/backends/vulkan/runtime/graph/ops/impl/QuantizedLinear.cpp index 45588e7e2e5..25a8d0b89ef 100644 --- a/backends/vulkan/runtime/graph/ops/impl/QuantizedLinear.cpp +++ b/backends/vulkan/runtime/graph/ops/impl/QuantizedLinear.cpp @@ -83,7 +83,7 @@ static CoopmatTileDims coopmat_tile_dims(const std::string& kernel_name) { return {kCoopmatTileM, kCoopmatTileN, kCoopmatTileK, kCoopmatInvocations}; } -utils::uvec3 quantized_linear_global_wg_size( +GlobalWorkGrid quantized_linear_gwg( ComputeGraph* graph, const vkapi::ShaderInfo& shader, const std::vector& args, @@ -97,14 +97,17 @@ utils::uvec3 quantized_linear_global_wg_size( const uint32_t M = utils::val_at(-2, out_sizes); // Coopmat variants dispatch a 256-thread WG per 64x64 output tile. Mirrors - // GemmCoopmat.cpp's pick_linear_coopmat_global_wg_size — the multiplication + // GemmCoopmat.cpp's pick_linear_coopmat_gwg — the multiplication // by kCoopmatInvocations cancels the framework's div_up, since - // local_wg = {256, 1, 1}. + // lwg = {256, 1, 1}. if (shader.kernel_name.find("_coopmat") != std::string::npos) { const CoopmatTileDims dims = coopmat_tile_dims(shader.kernel_name); const uint32_t num_tiles_n = utils::div_up(N, dims.n); const uint32_t num_tiles_m = utils::div_up(M, dims.m); - return {num_tiles_n * dims.wg_size, num_tiles_m, 1}; + return GlobalWorkGrid( + {num_tiles_n * dims.wg_size, num_tiles_m, 1u}, + kTiledWorkGrid, + LocalWorkGroup(dims.wg_size, 1u, 1u)); } uint32_t N_per_tile = 4; @@ -126,30 +129,25 @@ utils::uvec3 quantized_linear_global_wg_size( const uint32_t num_M_tiles = utils::div_up(M, M_per_tile); // Otherwise, each output tile contains 4 columns and 4 rows - return {num_N_tiles, num_M_tiles, 1}; + if (shader.kernel_name.find("_coop") != std::string::npos) { + return GlobalWorkGrid( + {num_N_tiles, num_M_tiles, 1u}, + kTiledWorkGrid, + LocalWorkGroup(1u, 1u, 64u)); + } + return GlobalWorkGrid({num_N_tiles, num_M_tiles, 1u}, kTiledWorkGrid); } -utils::uvec3 quantized_linear_local_wg_size( +LocalWorkGroup quantized_linear_lwg( ComputeGraph* graph, const vkapi::ShaderInfo& shader, - const utils::uvec3& global_workgroup_size, + const GlobalWorkGrid& gwg, const std::vector& args, const std::vector& resize_args) { - // Coopmat variants use a per-shader workgroup size (q4gsw/q8csw = 128, - // dq8ca = 256) — must match the WG_SIZE the shader yaml resolves to. - if (shader.kernel_name.find("_coopmat") != std::string::npos) { - return {coopmat_tile_dims(shader.kernel_name).wg_size, 1, 1}; - } - - const bool use_coop_algorithm = - shader.kernel_name.find("_coop") != std::string::npos; - - if (use_coop_algorithm) { - return {1, 1, 64}; - } else { - return pick_hw_square_wg_size( - graph, shader, global_workgroup_size, args, resize_args); + if (gwg.required_lwg_size().is_valid()) { + return pick_required_lwg(graph, shader, gwg, args, resize_args); } + return pick_xy_square_lwg(graph, shader, gwg, args, resize_args); } // Returns true when the q4gsw coopmat shader can be dispatched for this @@ -429,25 +427,26 @@ ValueRef prepack_quantized_linear_weight( ValueRef qmat2 = graph.add_tensor( qmat2_sizes, vkcompute::vkapi::kInt, storage_type, utils::kWidthPacked); - utils::uvec3 global_wg_size; + utils::uvec3 global_extents; if (weight_quant_config.nbits == 4) { // For 4-bit quantization, each thread writes out two adjacent blocks - global_wg_size = { + global_extents = { utils::safe_downcast(utils::div_up(num_blocks_K, int64_t(2))), utils::safe_downcast(num_blocks_N), 1u}; } else { - global_wg_size = { + global_extents = { utils::safe_downcast(num_blocks_N), utils::safe_downcast(num_blocks_K), 1u}; } + const GlobalWorkGrid gwg(global_extents, kTiledWorkGrid); graph.prepack_nodes().emplace_back(new PrepackNode( graph, VK_KERNEL_FROM_STR(kernel_name), - global_wg_size, - graph.create_local_wg_size(global_wg_size), + gwg, + graph.create_lwg(gwg), // Inputs and Outputs qmat2_data, qmat2, @@ -514,8 +513,8 @@ void add_linear_qw_node( graph.execute_nodes().emplace_back(new DynamicDispatchNode( graph, pick_linear_qw_shader, - quantized_linear_global_wg_size, - quantized_linear_local_wg_size, + quantized_linear_gwg, + quantized_linear_lwg, // Inputs and Outputs {{output, vkapi::kWrite}, {{fp_input, packed_weight, packed_weight_scales, packed_bias}, @@ -589,8 +588,8 @@ void add_linear_qa_qw_node( graph.execute_nodes().emplace_back(new DynamicDispatchNode( graph, VK_KERNEL_FROM_STR(kernel_name), - quantized_linear_global_wg_size, - quantized_linear_local_wg_size, + quantized_linear_gwg, + quantized_linear_lwg, // Inputs and Outputs {{output, vkapi::kWrite}, {{packed_int_input, @@ -661,8 +660,8 @@ void add_linear_dqa_qw_node( graph.execute_nodes().emplace_back(new DynamicDispatchNode( graph, pick_linear_dqa_qw_shader, - quantized_linear_global_wg_size, - quantized_linear_local_wg_size, + quantized_linear_gwg, + quantized_linear_lwg, // Inputs and Outputs {{output, vkapi::kWrite}, {{fp_input, diff --git a/backends/vulkan/runtime/graph/ops/impl/QuantizedLinear.h b/backends/vulkan/runtime/graph/ops/impl/QuantizedLinear.h index 7b62c98390d..7cb0e172c4a 100644 --- a/backends/vulkan/runtime/graph/ops/impl/QuantizedLinear.h +++ b/backends/vulkan/runtime/graph/ops/impl/QuantizedLinear.h @@ -14,10 +14,10 @@ namespace vkcompute { -utils::uvec3 quantized_linear_local_wg_size( +LocalWorkGroup quantized_linear_lwg( ComputeGraph* graph, const vkapi::ShaderInfo& shader, - const utils::uvec3& global_workgroup_size, + const GlobalWorkGrid& gwg, const std::vector& args, const std::vector& resize_args); diff --git a/backends/vulkan/runtime/graph/ops/impl/QuantizedLinearBackward.cpp b/backends/vulkan/runtime/graph/ops/impl/QuantizedLinearBackward.cpp index 493daa47126..8bfe1681bbf 100644 --- a/backends/vulkan/runtime/graph/ops/impl/QuantizedLinearBackward.cpp +++ b/backends/vulkan/runtime/graph/ops/impl/QuantizedLinearBackward.cpp @@ -32,7 +32,7 @@ void resize_linear_q4gsw_backward_node( graph->virtual_resize(d_x, new_sizes); } -utils::uvec3 linear_q4gsw_backward_global_wg_size( +GlobalWorkGrid linear_q4gsw_backward_gwg( ComputeGraph* graph, const vkapi::ShaderInfo& shader, const std::vector& args, @@ -43,21 +43,21 @@ utils::uvec3 linear_q4gsw_backward_global_wg_size( const uint32_t K = graph->size_at(-1, d_x); const uint32_t M = utils::safe_downcast(graph->numel_of(d_x) / K); const uint32_t tiles = utils::div_up_4(M) * utils::div_up_4(K); - return {tiles, 1u, 1u}; + return graph->create_linear_gwg(tiles); } -utils::uvec3 linear_q4gsw_backward_local_wg_size( +LocalWorkGroup linear_q4gsw_backward_lwg( ComputeGraph* graph, const vkapi::ShaderInfo& shader, - const utils::uvec3& global_workgroup_size, + const GlobalWorkGrid& gwg, const std::vector& args, const std::vector& resize_args) { (void)graph; (void)shader; - (void)global_workgroup_size; + (void)gwg; (void)args; (void)resize_args; - return {64u, 1u, 1u}; + return LocalWorkGroup(64u, 1u, 1u); } void linear_q4gsw_backward( @@ -106,8 +106,8 @@ void linear_q4gsw_backward( graph.execute_nodes().emplace_back(new DynamicDispatchNode( graph, VK_KERNEL_FROM_STR(kernel_name), - linear_q4gsw_backward_global_wg_size, - linear_q4gsw_backward_local_wg_size, + linear_q4gsw_backward_gwg, + linear_q4gsw_backward_lwg, // Inputs and Outputs {{d_x, vkapi::kWrite}, {{d_out, packed_weight, packed_scales}, vkapi::kRead}}, diff --git a/backends/vulkan/runtime/graph/ops/impl/QuantizedLinearQCSNW.cpp b/backends/vulkan/runtime/graph/ops/impl/QuantizedLinearQCSNW.cpp index 18958ccc3ce..5d9311e9761 100644 --- a/backends/vulkan/runtime/graph/ops/impl/QuantizedLinearQCSNW.cpp +++ b/backends/vulkan/runtime/graph/ops/impl/QuantizedLinearQCSNW.cpp @@ -17,7 +17,7 @@ namespace vkcompute { // Custom global workgroup size function for linear_qcs8w -utils::uvec3 linear_qcs8w_global_wg_size( +GlobalWorkGrid linear_qcs8w_gwg( ComputeGraph* graph, const vkapi::ShaderInfo& shader, const std::vector& args, @@ -25,26 +25,27 @@ utils::uvec3 linear_qcs8w_global_wg_size( (void)shader; (void)resize_args; const ValueRef out = args.at(0).refs.at(0); - return {static_cast(graph->numel_of(out)), 1, 1}; + return graph->create_linear_gwg( + utils::safe_downcast(graph->numel_of(out))); } // Custom local workgroup size function for linear_qcs8w -utils::uvec3 linear_qcs8w_local_wg_size( +LocalWorkGroup linear_qcs8w_lwg( ComputeGraph* graph, const vkapi::ShaderInfo& shader, - const utils::uvec3& global_workgroup_size, + const GlobalWorkGrid& gwg, const std::vector& args, const std::vector& resize_args) { (void)graph; (void)shader; - (void)global_workgroup_size; + (void)gwg; (void)args; (void)resize_args; - return {64, 1, 1}; + return LocalWorkGroup(64u, 1u, 1u); } // Custom global workgroup size function for linear_qcsnw_tiled -utils::uvec3 linear_qcsnw_tiled_global_wg_size( +GlobalWorkGrid linear_qcsnw_tiled_gwg( ComputeGraph* graph, const vkapi::ShaderInfo& shader, const std::vector& args, @@ -79,34 +80,33 @@ utils::uvec3 linear_qcsnw_tiled_global_wg_size( } utils::uvec3 out_limits = graph->logical_limits_of(out); - uint32_t global_wg_x = utils::div_up(out_limits[0], out_tile_ntxcols); - return { - global_wg_x * (utils::div_up(out_limits[1], out_tile_nrows)), - 1, + uint32_t gwg_x = utils::div_up(out_limits[0], out_tile_ntxcols); + const utils::uvec3 extents{ + gwg_x * (utils::div_up(out_limits[1], out_tile_nrows)), + 1u, out_limits[2]}; + if (shader.kernel_name.find("_coop") != std::string::npos) { + return GlobalWorkGrid(extents, kTiledWorkGrid, LocalWorkGroup(8u, 1u, 8u)); + } + return GlobalWorkGrid(extents, kTiledWorkGrid); } // Custom local workgroup size function for linear_qcsnw_tiled -utils::uvec3 linear_qcsnw_tiled_local_wg_size( +LocalWorkGroup linear_qcsnw_tiled_lwg( ComputeGraph* graph, const vkapi::ShaderInfo& shader, - const utils::uvec3& global_workgroup_size, + const GlobalWorkGrid& gwg, const std::vector& args, const std::vector& resize_args) { (void)graph; - (void)global_workgroup_size; + (void)gwg; (void)args; (void)resize_args; - // Check if using cooperative algorithm from shader name - bool use_coop_algorithm = - shader.kernel_name.find("_coop") != std::string::npos; - - if (use_coop_algorithm) { - return {8, 1, 8}; - } else { - return {64, 1, 1}; + if (gwg.required_lwg_size().is_valid()) { + return pick_required_lwg(graph, shader, gwg, args, resize_args); } + return LocalWorkGroup(64u, 1u, 1u); } void check_linear_qcsnw_args( @@ -229,15 +229,11 @@ void add_linear_qcs8w_node( graph.sizes_pc_of(q_mat2)}; } - const utils::uvec3 global_wg = { - static_cast(graph.numel_of(out_W_packed)), 1, 1}; - const utils::uvec3 local_wg{64, 1, 1}; - graph.execute_nodes().emplace_back(new DynamicDispatchNode( graph, VK_KERNEL_FROM_STR(kernel_name), - linear_qcs8w_global_wg_size, - linear_qcs8w_local_wg_size, + linear_qcs8w_gwg, + linear_qcs8w_lwg, // Inputs and Outputs {{out_W_packed, vkapi::MemoryAccessType::WRITE}, {{mat1_W_packed, q_mat2, scales}, vkapi::MemoryAccessType::READ}}, @@ -309,44 +305,21 @@ void add_linear_qcsnw_tiled_node( std::vector mat1_sizes = graph.sizes_of(mat1); const int64_t M = utils::val_at(-2, mat1_sizes); - uint32_t out_tile_nrows = 1; if (M % 3 == 0) { kernel_name += "_o4x3"; - out_tile_nrows = 3; } else if (M % 4 == 0) { kernel_name += "_o4x4"; - out_tile_nrows = 4; } else if (M % 2 == 0) { kernel_name += "_o4x2"; - out_tile_nrows = 2; } else { kernel_name += "_o4x1"; - out_tile_nrows = 1; - } - - // Number of output texels in the output tile - uint32_t out_tile_ntxcols = 1; - if (quant_nbits == 4) { - out_tile_ntxcols = 2; - } - - utils::uvec3 out_limits = graph.logical_limits_of(out); - uint32_t global_wg_x = utils::div_up(out_limits[0], out_tile_ntxcols); - utils::uvec3 global_wg_size = { - global_wg_x * (utils::div_up(out_limits[1], out_tile_nrows)), - 1, - out_limits[2]}; - - utils::uvec3 local_wg_size{64, 1, 1}; - if (use_coop_algorithm) { - local_wg_size = {8, 1, 8}; } graph.execute_nodes().emplace_back(new DynamicDispatchNode( graph, VK_KERNEL_FROM_STR(kernel_name), - linear_qcsnw_tiled_global_wg_size, - linear_qcsnw_tiled_local_wg_size, + linear_qcsnw_tiled_gwg, + linear_qcsnw_tiled_lwg, // Inputs and Outputs {{out, vkapi::kWrite}, {{mat1, q_mat2, scales}, vkapi::kRead}}, // Shader params buffers diff --git a/backends/vulkan/runtime/graph/ops/impl/QuantizedLinearRequant.cpp b/backends/vulkan/runtime/graph/ops/impl/QuantizedLinearRequant.cpp index ab7336fcc33..34acfeded68 100644 --- a/backends/vulkan/runtime/graph/ops/impl/QuantizedLinearRequant.cpp +++ b/backends/vulkan/runtime/graph/ops/impl/QuantizedLinearRequant.cpp @@ -34,7 +34,7 @@ void resize_q4gsw_requant_node( graph->virtual_resize(packed, {K4 * N4_padded * 2}); } -utils::uvec3 q4gsw_requant_global_wg_size( +GlobalWorkGrid q4gsw_requant_gwg( ComputeGraph* graph, const vkapi::ShaderInfo& shader, const std::vector& args, @@ -48,21 +48,21 @@ utils::uvec3 q4gsw_requant_global_wg_size( const uint32_t K4 = K / 4u; const uint32_t N4 = (N + 3u) / 4u; const uint32_t N8 = (N4 + 1u) / 2u; - return {K4, N8, 1u}; + return GlobalWorkGrid({K4, N8, 1u}, kTiledWorkGrid); } -utils::uvec3 q4gsw_requant_local_wg_size( +LocalWorkGroup q4gsw_requant_lwg( ComputeGraph* graph, const vkapi::ShaderInfo& shader, - const utils::uvec3& global_workgroup_size, + const GlobalWorkGrid& gwg, const std::vector& args, const std::vector& resize_args) { (void)graph; (void)shader; - (void)global_workgroup_size; + (void)gwg; (void)args; (void)resize_args; - return {8u, 8u, 1u}; + return LocalWorkGroup(8u, 8u, 1u); } void q4gsw_requant(ComputeGraph& graph, const std::vector& args) { @@ -108,8 +108,8 @@ void q4gsw_requant(ComputeGraph& graph, const std::vector& args) { graph.execute_nodes().emplace_back(new DynamicDispatchNode( graph, VK_KERNEL_FROM_STR(kernel_name), - q4gsw_requant_global_wg_size, - q4gsw_requant_local_wg_size, + q4gsw_requant_gwg, + q4gsw_requant_lwg, // Inputs and Outputs {{packed, vkapi::kWrite}, {{latent, packed_scales}, vkapi::kRead}}, // Shader params buffers diff --git a/backends/vulkan/runtime/graph/ops/impl/Reduce.cpp b/backends/vulkan/runtime/graph/ops/impl/Reduce.cpp index 856783ce219..f684906cf57 100644 --- a/backends/vulkan/runtime/graph/ops/impl/Reduce.cpp +++ b/backends/vulkan/runtime/graph/ops/impl/Reduce.cpp @@ -72,48 +72,28 @@ void resize_reduce_per_row_node( graph->virtual_resize(out, new_sizes); } -utils::uvec3 reduce_global_wg_size( +GlobalWorkGrid reduce_gwg( ComputeGraph* graph, const vkapi::ShaderInfo& shader, const std::vector& args, const std::vector& resize_args) { (void)shader; const ValueRef out = args.at(0).refs.at(0); - const int32_t reduce_dim_whcn = - graph->extract_scalar(resize_args.at(1)); - - utils::uvec3 global_wg_size = graph->logical_limits_of(out); - global_wg_size[reduce_dim_whcn] = 1; - return global_wg_size; -} - -utils::uvec3 reduce_local_wg_size( - ComputeGraph* graph, - const vkapi::ShaderInfo& shader, - const utils::uvec3& global_workgroup_size, - const std::vector& args, - const std::vector& resize_args) { - (void)shader; - (void)args; - (void)global_workgroup_size; - const int32_t reduce_dim_whcn = graph->extract_scalar(resize_args.at(1)); const int64_t group_dim_whcn = graph->extract_scalar(resize_args.at(2)); - // This should match the value of MAX_NTHREADS in the reduce shader. - constexpr uint32_t max_nthreads = 16; - - const uint32_t nworkers_per_group = 4; - const uint32_t ngroups = 4; + utils::uvec3 extents = graph->logical_limits_of(out); + extents[reduce_dim_whcn] = 1; + constexpr uint32_t max_nthreads = 16u; + constexpr uint32_t nworkers_per_group = 4u; + constexpr uint32_t ngroups = 4u; VK_CHECK_COND(nworkers_per_group * ngroups <= max_nthreads); - - utils::uvec3 local_wg_size{1, 1, 1}; - local_wg_size[reduce_dim_whcn] = nworkers_per_group; - local_wg_size[group_dim_whcn] = ngroups; - - return local_wg_size; + utils::uvec3 lwg_extents{1u, 1u, 1u}; + lwg_extents[reduce_dim_whcn] = nworkers_per_group; + lwg_extents[group_dim_whcn] = ngroups; + return GlobalWorkGrid(extents, kTiledWorkGrid, LocalWorkGroup(lwg_extents)); } void add_reduce_node( @@ -161,8 +141,8 @@ void add_reduce_node( graph.execute_nodes().emplace_back(new DynamicDispatchNode( graph, VK_KERNEL_FROM_STR(kernel_name), - reduce_global_wg_size, - reduce_local_wg_size, + reduce_gwg, + pick_required_lwg, // Inputs and Outputs {{out, vkapi::kWrite}, {in, vkapi::kRead}}, // Shader params buffers @@ -238,8 +218,8 @@ void add_reduce2d_node( graph.execute_nodes().emplace_back(new DynamicDispatchNode( graph, VK_KERNEL_FROM_STR(kernel_name), - reduce_global_wg_size, - reduce_local_wg_size, + reduce_gwg, + pick_required_lwg, // Inputs and Outputs {{out, vkapi::kWrite}, {in, vkapi::kRead}}, // Shader params buffers @@ -257,7 +237,7 @@ void add_reduce2d_node( resize_reduce2d_node)); } -utils::uvec3 reduce_per_row_global_wg_size( +GlobalWorkGrid reduce_per_row_gwg( ComputeGraph* graph, const vkapi::ShaderInfo& shader, const std::vector& args, @@ -266,23 +246,10 @@ utils::uvec3 reduce_per_row_global_wg_size( (void)resize_args; const ValueRef out = args.at(0).refs.at(0); - return {1u, utils::safe_downcast(graph->numel_of(out)), 1u}; -} - -utils::uvec3 reduce_per_row_local_wg_size( - ComputeGraph* graph, - const vkapi::ShaderInfo& shader, - const utils::uvec3& global_workgroup_size, - const std::vector& args, - const std::vector& resize_args) { - (void)global_workgroup_size; - (void)args; - (void)resize_args; - - uint32_t outputs_per_wg = 1u; - uint32_t workers_per_output = 64u; - - return {workers_per_output, outputs_per_wg, 1u}; + return GlobalWorkGrid( + {1u, utils::safe_downcast(graph->numel_of(out)), 1u}, + kTiledWorkGrid, + LocalWorkGroup(64u, 1u, 1u)); } void add_reduce_per_row_node( @@ -305,9 +272,9 @@ void add_reduce_per_row_node( graph, VK_KERNEL_FROM_STR(kernel_name), // Global workgroup size function - reduce_per_row_global_wg_size, + reduce_per_row_gwg, // Local workgroup size function - reduce_per_row_local_wg_size, + pick_required_lwg, // Inputs and Outputs {{output, vkapi::kWrite}, {input, vkapi::kRead}}, // Shader param buffers diff --git a/backends/vulkan/runtime/graph/ops/impl/Repeat.cpp b/backends/vulkan/runtime/graph/ops/impl/Repeat.cpp index 48b990f4622..99b12abc487 100644 --- a/backends/vulkan/runtime/graph/ops/impl/Repeat.cpp +++ b/backends/vulkan/runtime/graph/ops/impl/Repeat.cpp @@ -53,8 +53,8 @@ void add_repeat_node( graph.execute_nodes().emplace_back(new DynamicDispatchNode( graph, VK_KERNEL_FROM_STR(kernel_name), - default_pick_global_wg_size, - default_pick_local_wg_size, + default_pick_gwg, + default_pick_lwg, {{out, vkapi::kWrite}, {in, vkapi::kRead}}, {graph.meta_ubo(out), graph.meta_ubo(in)}, {}, diff --git a/backends/vulkan/runtime/graph/ops/impl/RepeatInterleave.cpp b/backends/vulkan/runtime/graph/ops/impl/RepeatInterleave.cpp index 221d0d23f51..e2cae010360 100644 --- a/backends/vulkan/runtime/graph/ops/impl/RepeatInterleave.cpp +++ b/backends/vulkan/runtime/graph/ops/impl/RepeatInterleave.cpp @@ -53,8 +53,8 @@ void add_repeat_interleave_node( graph.execute_nodes().emplace_back(new DynamicDispatchNode( graph, VK_KERNEL_FROM_STR(kernel_name), - default_pick_global_wg_size, - default_pick_local_wg_size, + default_pick_gwg, + default_pick_lwg, // Inputs and Outputs {{out, vkapi::MemoryAccessType::WRITE}, {in, vkapi::MemoryAccessType::READ}}, diff --git a/backends/vulkan/runtime/graph/ops/impl/RmsNorm.cpp b/backends/vulkan/runtime/graph/ops/impl/RmsNorm.cpp index 1eb267f9794..5238bf0a74e 100644 --- a/backends/vulkan/runtime/graph/ops/impl/RmsNorm.cpp +++ b/backends/vulkan/runtime/graph/ops/impl/RmsNorm.cpp @@ -28,7 +28,7 @@ void resize_rms_norm_node( graph->virtual_resize(out, graph->sizes_of(in)); } -utils::uvec3 rms_norm_global_wg_size( +GlobalWorkGrid rms_norm_gwg( ComputeGraph* graph, const vkapi::ShaderInfo& shader, const std::vector& args, @@ -40,21 +40,8 @@ utils::uvec3 rms_norm_global_wg_size( const int64_t hidden = sizes.back(); const int64_t numel = graph->numel_of(in); const uint32_t num_rows = utils::safe_downcast(numel / hidden); - return {1u, num_rows, 1u}; -} - -utils::uvec3 rms_norm_local_wg_size( - ComputeGraph* graph, - const vkapi::ShaderInfo& shader, - const utils::uvec3& global_workgroup_size, - const std::vector& args, - const std::vector& resize_args) { - (void)graph; - (void)shader; - (void)global_workgroup_size; - (void)args; - (void)resize_args; - return {64u, 1u, 1u}; + return GlobalWorkGrid( + {1u, num_rows, 1u}, kTiledWorkGrid, LocalWorkGroup(64u, 1u, 1u)); } void add_rms_norm_node( @@ -87,8 +74,8 @@ void add_rms_norm_node( graph.execute_nodes().emplace_back(new DynamicDispatchNode( graph, VK_KERNEL_FROM_STR(kernel_name), - rms_norm_global_wg_size, - rms_norm_local_wg_size, + rms_norm_gwg, + pick_required_lwg, // Inputs and Outputs {{out, vkapi::kWrite}, {{in, arg_weight}, vkapi::kRead}}, // Shader params buffers diff --git a/backends/vulkan/runtime/graph/ops/impl/RotaryEmbedding.cpp b/backends/vulkan/runtime/graph/ops/impl/RotaryEmbedding.cpp index d1e70cc8c41..2333c718b7e 100644 --- a/backends/vulkan/runtime/graph/ops/impl/RotaryEmbedding.cpp +++ b/backends/vulkan/runtime/graph/ops/impl/RotaryEmbedding.cpp @@ -33,7 +33,7 @@ void resize_rotary_embedding_node( graph->virtual_resize(xk_out, xk_sizes); } -utils::uvec3 rotary_embedding_global_wg_size( +GlobalWorkGrid rotary_embedding_gwg( ComputeGraph* graph, const vkapi::ShaderInfo& shader, const std::vector& args, @@ -53,7 +53,7 @@ utils::uvec3 rotary_embedding_global_wg_size( // Input tokens sequence length const uint32_t S = graph->size_at(-3, xq_out); - return {D8, QH, S}; + return GlobalWorkGrid({D8, QH, S}, kTiledWorkGrid); } void add_rotary_embedding_node( @@ -91,8 +91,8 @@ void add_rotary_embedding_node( graph.execute_nodes().emplace_back(new DynamicDispatchNode( graph, VK_KERNEL_FROM_STR(kernel_name), - rotary_embedding_global_wg_size, - default_pick_local_wg_size, + rotary_embedding_gwg, + default_pick_lwg, // Inputs and Outputs {{{xq_out, xk_out}, vkapi::kWrite}, {{xq, xk, freqs_cos, freqs_sin}, vkapi::kRead}}, @@ -121,7 +121,7 @@ void apply_rotary_emb(ComputeGraph& graph, const std::vector& args) { // HuggingFace RoPE variant // -utils::uvec3 rotary_embedding_hf_global_wg_size( +GlobalWorkGrid rotary_embedding_hf_gwg( ComputeGraph* graph, const vkapi::ShaderInfo& shader, const std::vector& args, @@ -139,7 +139,7 @@ utils::uvec3 rotary_embedding_hf_global_wg_size( const uint32_t QH = graph->size_at(-2, xq_out); const uint32_t S = graph->size_at(-3, xq_out); - return {D4, QH, S}; + return GlobalWorkGrid({D4, QH, S}, kTiledWorkGrid); } void add_rotary_embedding_hf_node( @@ -189,8 +189,8 @@ void add_rotary_embedding_hf_node( graph.execute_nodes().emplace_back(new DynamicDispatchNode( graph, VK_KERNEL_FROM_STR(kernel_name), - rotary_embedding_hf_global_wg_size, - default_pick_local_wg_size, + rotary_embedding_hf_gwg, + default_pick_lwg, // Inputs and Outputs {{{xq_out, xk_out}, vkapi::kWrite}, {{xq, xk, freqs_cos, freqs_sin}, vkapi::kRead}}, @@ -243,7 +243,7 @@ void resize_rotary_embedding_interleaved_node( graph->virtual_resize(out, graph->sizes_of(in)); } -utils::uvec3 rotary_embedding_interleaved_global_wg_size( +GlobalWorkGrid rotary_embedding_interleaved_gwg( ComputeGraph* graph, const vkapi::ShaderInfo& shader, const std::vector& args, @@ -261,7 +261,7 @@ utils::uvec3 rotary_embedding_interleaved_global_wg_size( const uint32_t C = static_cast(out_sizes.at(2)); // One thread per output texel of 4 elements along C. - return {utils::div_up_4(C), N, B}; + return GlobalWorkGrid({utils::div_up_4(C), N, B}, kTiledWorkGrid); } void add_rotary_embedding_interleaved_node( @@ -303,8 +303,8 @@ void add_rotary_embedding_interleaved_node( graph.execute_nodes().emplace_back(new DynamicDispatchNode( graph, VK_KERNEL_FROM_STR(kernel_name), - rotary_embedding_interleaved_global_wg_size, - default_pick_local_wg_size, + rotary_embedding_interleaved_gwg, + default_pick_lwg, // Inputs and Outputs {{out, vkapi::kWrite}, {{x, freqs_cis}, vkapi::kRead}}, // Parameter buffers diff --git a/backends/vulkan/runtime/graph/ops/impl/SDPA.cpp b/backends/vulkan/runtime/graph/ops/impl/SDPA.cpp index c4259d4c530..39854b23d02 100644 --- a/backends/vulkan/runtime/graph/ops/impl/SDPA.cpp +++ b/backends/vulkan/runtime/graph/ops/impl/SDPA.cpp @@ -143,7 +143,7 @@ void resize_sdpa_out_node( // Shader dispatch pick functions // -utils::uvec3 kv_cache_update_global_wg_size( +GlobalWorkGrid kv_cache_update_gwg( ComputeGraph* graph, const vkapi::ShaderInfo& shader, const std::vector& args, @@ -157,7 +157,8 @@ utils::uvec3 kv_cache_update_global_wg_size( const uint32_t num_heads = graph->size_at(-2, projected); const uint32_t seq_len = graph->size_at(-3, projected); - return {utils::div_up_4(head_dim_size), seq_len, num_heads}; + return GlobalWorkGrid( + {utils::div_up_4(head_dim_size), seq_len, num_heads}, kTiledWorkGrid); } // resize_args layout for SDPA dispatch pickers mirrors the node creation @@ -264,12 +265,11 @@ vkapi::ShaderInfo pick_sdpa_qk_shader( } } -utils::uvec3 pick_sdpa_qk_global_wg_size( +GlobalWorkGrid pick_sdpa_qk_gwg( ComputeGraph* graph, const vkapi::ShaderInfo& shader, const std::vector& args, const std::vector& resize_args) { - (void)shader; (void)args; const SDPAMode mode = mode_of(resize_args); const ValueRef q = resize_args.at(0); @@ -280,30 +280,30 @@ utils::uvec3 pick_sdpa_qk_global_wg_size( // Dispatch grid: (context_len tiles, S tiles, H * B). const uint32_t N4 = utils::div_up_4(static_cast(d.context_len)); const uint32_t M4 = utils::div_up_4(static_cast(d.S)); - return {N4, M4, static_cast(d.H * d.B)}; + const utils::uvec3 extents{N4, M4, static_cast(d.H * d.B)}; + if (shader.kernel_name.find("_coop") != std::string::npos) { + return GlobalWorkGrid(extents, kTiledWorkGrid, LocalWorkGroup(1u, 64u, 1u)); + } + return GlobalWorkGrid(extents, kTiledWorkGrid); } -utils::uvec3 pick_sdpa_qk_local_wg_size( +LocalWorkGroup pick_sdpa_qk_lwg( ComputeGraph* graph, const vkapi::ShaderInfo& shader, - const utils::uvec3& global_workgroup_size, + const GlobalWorkGrid& gwg, const std::vector& args, const std::vector& resize_args) { + if (gwg.required_lwg_size().is_valid()) { + return pick_required_lwg(graph, shader, gwg, args, resize_args); + } const SDPAMode mode = mode_of(resize_args); if (mode == SDPAMode::LLM) { - const bool use_coop_algorithm = - shader.kernel_name.find("_coop") != std::string::npos; - if (use_coop_algorithm) { - return {1, 64, 1}; - } - return pick_hw_square_wg_size( - graph, shader, global_workgroup_size, args, resize_args); + return pick_xy_square_lwg(graph, shader, gwg, args, resize_args); } - return default_pick_local_wg_size( - graph, shader, global_workgroup_size, args, resize_args); + return default_pick_lwg(graph, shader, gwg, args, resize_args); } -utils::uvec3 pick_sdpa_softmax_global_wg_size( +GlobalWorkGrid pick_sdpa_softmax_gwg( ComputeGraph* graph, const vkapi::ShaderInfo& shader, const std::vector& args, @@ -321,24 +321,12 @@ utils::uvec3 pick_sdpa_softmax_global_wg_size( : graph->size_at(-2, q); const int64_t B = (mode == SDPAMode::LLM) ? 1 : graph->size_at(-4, q); - return { - 1, - static_cast(seq_len), - static_cast(num_q_heads * B)}; -} - -utils::uvec3 pick_sdpa_softmax_local_wg_size( - ComputeGraph* graph, - const vkapi::ShaderInfo& shader, - const utils::uvec3& global_workgroup_size, - const std::vector& args, - const std::vector& resize_args) { - (void)graph; - (void)shader; - (void)global_workgroup_size; - (void)args; - (void)resize_args; - return {64, 1, 1}; + return GlobalWorkGrid( + {1u, + static_cast(seq_len), + static_cast(num_q_heads * B)}, + kTiledWorkGrid, + LocalWorkGroup(64u, 1u, 1u)); } vkapi::ShaderInfo pick_sdpa_av_shader( @@ -371,7 +359,7 @@ vkapi::ShaderInfo pick_sdpa_av_shader( // consistent win on Adreno (AV ~1.14-1.63x) but a regression on Mali at // common decode contexts (~0.67-0.86x, interleaved median), so it is // selected vendor-adaptively (tests can pin it via shader_override; see - // resolve_use_tile2). pick_sdpa_av_global_wg_size keys the x-dim + // resolve_use_tile2). pick_sdpa_av_gwg keys the x-dim // collapse off this same _tile2 suffix. if (resolve_use_tile2(graph, shader_override)) { shader_name += "_tile2"; @@ -397,7 +385,7 @@ vkapi::ShaderInfo pick_sdpa_av_shader( } } -utils::uvec3 pick_sdpa_av_global_wg_size( +GlobalWorkGrid pick_sdpa_av_gwg( ComputeGraph* graph, const vkapi::ShaderInfo& shader, const std::vector& args, @@ -425,30 +413,33 @@ utils::uvec3 pick_sdpa_av_global_wg_size( shader.kernel_name.find("_tile2") != std::string::npos ? utils::div_up(N4, 2u) : N4; - return {x_dim, M4, static_cast(num_kv_heads * d.B)}; + return GlobalWorkGrid( + {x_dim, M4, static_cast(num_kv_heads * d.B)}, + kTiledWorkGrid, + LocalWorkGroup(1u, 64u, 1u)); } - return {N4, M4, static_cast(d.H * d.B)}; + const utils::uvec3 extents{N4, M4, static_cast(d.H * d.B)}; + if (shader.kernel_name.find("_coop") != std::string::npos) { + return GlobalWorkGrid(extents, kTiledWorkGrid, LocalWorkGroup(1u, 64u, 1u)); + } + return GlobalWorkGrid(extents, kTiledWorkGrid); } -utils::uvec3 pick_sdpa_av_local_wg_size( +LocalWorkGroup pick_sdpa_av_lwg( ComputeGraph* graph, const vkapi::ShaderInfo& shader, - const utils::uvec3& global_workgroup_size, + const GlobalWorkGrid& gwg, const std::vector& args, const std::vector& resize_args) { + if (gwg.required_lwg_size().is_valid()) { + return pick_required_lwg(graph, shader, gwg, args, resize_args); + } const SDPAMode mode = mode_of(resize_args); if (mode == SDPAMode::LLM) { - const bool use_coop_algorithm = - shader.kernel_name.find("_coop") != std::string::npos; - if (use_coop_algorithm) { - return {1, 64, 1}; - } - return pick_hw_square_wg_size( - graph, shader, global_workgroup_size, args, resize_args); + return pick_xy_square_lwg(graph, shader, gwg, args, resize_args); } - return default_pick_local_wg_size( - graph, shader, global_workgroup_size, args, resize_args); + return default_pick_lwg(graph, shader, gwg, args, resize_args); } // @@ -473,8 +464,8 @@ void add_sdpa_kv_cache_update_node( graph.execute_nodes().emplace_back(new DynamicDispatchNode( graph, VK_KERNEL_FROM_STR(kernel_name), - kv_cache_update_global_wg_size, - default_pick_local_wg_size, + kv_cache_update_gwg, + default_pick_lwg, // Inputs and Outputs {{cache, vkapi::kWrite}, {projected, vkapi::kRead}}, // Shader param buffers @@ -522,8 +513,8 @@ void add_sdpa_compute_attn_weights_node( graph.execute_nodes().emplace_back(new DynamicDispatchNode( graph, pick_sdpa_qk_shader, - pick_sdpa_qk_global_wg_size, - pick_sdpa_qk_local_wg_size, + pick_sdpa_qk_gwg, + pick_sdpa_qk_lwg, // Inputs and Outputs {{attn_weights, vkapi::kWrite}, {read_inputs, vkapi::kRead}}, // Shader param buffers @@ -574,8 +565,8 @@ void add_sdpa_attn_weights_softmax_node( graph.execute_nodes().emplace_back(new DynamicDispatchNode( graph, VK_KERNEL_FROM_STR(shader_name), - pick_sdpa_softmax_global_wg_size, - pick_sdpa_softmax_local_wg_size, + pick_sdpa_softmax_gwg, + pick_required_lwg, // Inputs and Outputs {{attn_weights_softmax, vkapi::kWrite}, {attn_weights, vkapi::kRead}}, // Shader param buffers @@ -633,8 +624,8 @@ void add_sdpa_compute_out_node( graph.execute_nodes().emplace_back(new DynamicDispatchNode( graph, pick_sdpa_av_shader, - pick_sdpa_av_global_wg_size, - pick_sdpa_av_local_wg_size, + pick_sdpa_av_gwg, + pick_sdpa_av_lwg, // Inputs and Outputs {{out, vkapi::kWrite}, {{attn_weights_softmax, v}, vkapi::kRead}}, // Shader param buffers diff --git a/backends/vulkan/runtime/graph/ops/impl/ScalarTensor.cpp b/backends/vulkan/runtime/graph/ops/impl/ScalarTensor.cpp index 82fc5c977d3..ca2fe79b7c4 100644 --- a/backends/vulkan/runtime/graph/ops/impl/ScalarTensor.cpp +++ b/backends/vulkan/runtime/graph/ops/impl/ScalarTensor.cpp @@ -31,8 +31,8 @@ void scalar_tensor(ComputeGraph& graph, const std::vector& args) { graph.execute_nodes().emplace_back(new DispatchNode( graph, VK_KERNEL_FROM_STR(kernel_name), - graph.create_global_wg_size(out), - graph.create_local_wg_size(out), + graph.create_gwg(out), + graph.create_lwg(out), // Inputs and Outputs {{out, vkapi::kWrite}}, // Shader params buffers diff --git a/backends/vulkan/runtime/graph/ops/impl/Softmax.cpp b/backends/vulkan/runtime/graph/ops/impl/Softmax.cpp index 55ecf466ab6..ab69b421bcc 100644 --- a/backends/vulkan/runtime/graph/ops/impl/Softmax.cpp +++ b/backends/vulkan/runtime/graph/ops/impl/Softmax.cpp @@ -18,7 +18,7 @@ namespace vkcompute { using namespace utils; -utils::uvec3 pick_softmax_global_wg_size( +GlobalWorkGrid pick_softmax_gwg( ComputeGraph* graph, const vkapi::ShaderInfo& shader, const std::vector& args, @@ -33,53 +33,23 @@ utils::uvec3 pick_softmax_global_wg_size( int32_t reduce_dim = normalize(dim, ndim); reduce_dim = nchw_dim_to_whcn_dim(reduce_dim, ndim); + utils::uvec3 lwg_extents{1u, 1u, 1u}; + lwg_extents[reduce_dim] = 4u; if (graph->is_buffer_storage(out)) { - utils::uvec3 global_size = { + utils::uvec3 extents = { graph->size_at(-1, out), graph->size_at(-2, out), graph->size_at(-3, out) * graph->size_at(-4, out)}; - global_size[reduce_dim] = 1; - return global_size; - } - - utils::uvec3 global_size = graph->logical_limits_of(out); - global_size[reduce_dim] = 1; - return global_size; -} - -utils::uvec3 pick_softmax_local_wg_size( - ComputeGraph* graph, - const vkapi::ShaderInfo& shader, - const utils::uvec3& global_workgroup_size, - const std::vector& args, - const std::vector& resize_args) { - (void)shader; - (void)global_workgroup_size; - - const ValueRef out = args.at(0).refs.at(0); - const ValueRef in = args.at(1).refs.at(0); - const int dim = resize_args.at(0); - - const int64_t ndim = graph->dim_of(in); - int32_t reduce_dim = normalize(dim, ndim); - reduce_dim = nchw_dim_to_whcn_dim(reduce_dim, ndim); - - const uint32_t nworkers_per_group = 4; - - if (graph->is_buffer_storage(out)) { - utils::uvec3 local_wg_size{1, 1, 1}; - local_wg_size[reduce_dim] = nworkers_per_group; - return local_wg_size; + extents[reduce_dim] = 1; + return GlobalWorkGrid(extents, kTiledWorkGrid, LocalWorkGroup(lwg_extents)); } + utils::uvec3 extents = graph->logical_limits_of(out); + extents[reduce_dim] = 1; const int64_t group_dim_xyz = graph->extract_scalar(resize_args.at(1)); - const uint32_t ngroups = 4; - - utils::uvec3 local_wg_size{1, 1, 1}; - local_wg_size[reduce_dim] = nworkers_per_group; - local_wg_size[group_dim_xyz] = ngroups; - return local_wg_size; + lwg_extents[group_dim_xyz] = 4u; + return GlobalWorkGrid(extents, kTiledWorkGrid, LocalWorkGroup(lwg_extents)); } void resize_softmax_node( @@ -140,8 +110,8 @@ void add_softmax_node( const int other_dim_1 = (reduce_dim_xyz + 1) % 3; const int other_dim_2 = (reduce_dim_xyz + 2) % 3; int32_t group_dim; - utils::uvec3 global_wg_size = graph.logical_limits_of(out); - if (global_wg_size[other_dim_1] > global_wg_size[other_dim_2]) { + const utils::uvec3 extents = graph.logical_limits_of(out); + if (extents[other_dim_1] > extents[other_dim_2]) { group_dim = other_dim_1; } else { group_dim = other_dim_2; @@ -157,8 +127,8 @@ void add_softmax_node( graph.execute_nodes().emplace_back(new DynamicDispatchNode( graph, VK_KERNEL_FROM_STR(kernel_name), - pick_softmax_global_wg_size, - pick_softmax_local_wg_size, + pick_softmax_gwg, + pick_required_lwg, // Inputs and Outputs {{out, vkapi::kWrite}, {in, vkapi::kRead}}, // Shader params buffers diff --git a/backends/vulkan/runtime/graph/ops/impl/Split.cpp b/backends/vulkan/runtime/graph/ops/impl/Split.cpp index c38189d99f6..2a2176a602d 100644 --- a/backends/vulkan/runtime/graph/ops/impl/Split.cpp +++ b/backends/vulkan/runtime/graph/ops/impl/Split.cpp @@ -42,8 +42,8 @@ void add_split_node( graph.execute_nodes().emplace_back(new DynamicDispatchNode( graph, VK_KERNEL_FROM_STR(kernel_name), - default_pick_global_wg_size, - default_pick_local_wg_size, + default_pick_gwg, + default_pick_lwg, // Inputs and Outputs {{out, vkapi::kWrite}, {input, vkapi::kRead}}, // Shader params buffers diff --git a/backends/vulkan/runtime/graph/ops/impl/Staging.cpp b/backends/vulkan/runtime/graph/ops/impl/Staging.cpp index 0b2c1799283..37c8b4d4f3f 100644 --- a/backends/vulkan/runtime/graph/ops/impl/Staging.cpp +++ b/backends/vulkan/runtime/graph/ops/impl/Staging.cpp @@ -72,8 +72,8 @@ void add_staging_to_tensor_node( graph.execute_nodes().emplace_back(new DynamicDispatchNode( graph, shader, - default_pick_global_wg_size, - default_pick_local_wg_size, + default_pick_gwg, + default_pick_lwg, // Input and Outputs {{out_tensor, vkapi::kWrite}, {in_staging, vkapi::kRead}}, // Parameter Buffers @@ -88,7 +88,7 @@ void add_staging_to_tensor_node( nullptr)); } -utils::uvec3 tensor_to_staging_global_wg_size( +GlobalWorkGrid tensor_to_staging_gwg( ComputeGraph* graph, const vkapi::ShaderInfo& shader, const std::vector& args, @@ -97,7 +97,7 @@ utils::uvec3 tensor_to_staging_global_wg_size( const ValueRef in_tensor = args.at(1).refs.at(0); const ValueRef out_staging = args.at(0).refs.at(0); - utils::uvec3 global_wg_size = graph->create_global_wg_size(in_tensor); + GlobalWorkGrid gwg = graph->create_gwg(in_tensor); // The bitw8 shader writes out a complete 32 bit integer (containing 4 packed // 8-bit integers) per thread, so its global work group size is the number of @@ -105,22 +105,22 @@ utils::uvec3 tensor_to_staging_global_wg_size( if (is_bitw8_shader(shader)) { const uint32_t buffer_len = utils::safe_downcast( graph->get_staging(out_staging)->numel() / 4); - global_wg_size = {buffer_len, 1, 1}; + gwg = graph->create_linear_gwg(buffer_len); } else if (is_coalesced_image_to_nchw_shader(shader)) { // The coalesced (output-centric) image_to_nchw variant dispatches one // thread per output (staging) element so that consecutive threads write // consecutive NCHW offsets, keeping writes to the PCIe-backed staging // buffer fully coalesced. This mirrors the buffer_to_nchw path, whose - // global size is already numel-based via create_global_wg_size. + // global size is already numel-based via create_gwg. const uint32_t buffer_len = utils::safe_downcast( graph->get_staging(out_staging)->numel()); - global_wg_size = {buffer_len, 1, 1}; + gwg = graph->create_linear_gwg(buffer_len); } // Otherwise (texel-centric image_to_nchw, used on unified-memory GPUs) keep - // the default texel-grid global size from create_global_wg_size: one thread + // the default texel-grid global size from create_gwg: one thread // per texture texel. - return global_wg_size; + return gwg; } void add_tensor_to_staging_node( @@ -155,8 +155,8 @@ void add_tensor_to_staging_node( graph.execute_nodes().emplace_back(new DynamicDispatchNode( graph, shader, - tensor_to_staging_global_wg_size, - default_pick_local_wg_size, + tensor_to_staging_gwg, + default_pick_lwg, // Input and Outputs {{out_staging, vkapi::kWrite}, {in_tensor, vkapi::kRead}}, // Parameter Buffers @@ -204,8 +204,8 @@ void add_prepack_standard_node( graph.prepack_nodes().emplace_back(new PrepackNode( graph, shader, - graph.create_global_wg_size(tensor), - graph.create_local_wg_size(tensor), + graph.create_gwg(tensor), + graph.create_lwg(tensor), // Input and Outputs tensor_data, tensor, @@ -288,8 +288,8 @@ void add_prepack_direct_copy_buffer_node( graph.prepack_nodes().emplace_back(new PrepackNode( graph, shader, - graph.create_global_wg_size(tensor), - graph.create_local_wg_size(tensor), + graph.create_gwg(tensor), + graph.create_lwg(tensor), // Input and Outputs tensor_data, tensor, @@ -329,9 +329,9 @@ ValueRef prepack_int4_linear_weight_transposed_interleaved( ValueRef qmat2 = graph.add_tensor( qmat2_sizes, vkcompute::vkapi::kByte, storage_type, utils::kWidthPacked); - utils::uvec3 global_wg_size; - global_wg_size = graph.logical_limits_of(qmat2); - global_wg_size[1] = utils::div_up(global_wg_size[1], uint32_t(2)); + utils::uvec3 global_extents = graph.logical_limits_of(qmat2); + global_extents[1] = utils::div_up(global_extents[1], uint32_t(2)); + const GlobalWorkGrid gwg(global_extents, kTiledWorkGrid); std::string kernel_name = graph.context()->adapter_ptr()->has_full_int8_buffers_support() @@ -342,8 +342,8 @@ ValueRef prepack_int4_linear_weight_transposed_interleaved( graph.prepack_nodes().emplace_back(new PrepackNode( graph, VK_KERNEL_FROM_STR(kernel_name), - global_wg_size, - graph.create_local_wg_size(global_wg_size), + gwg, + graph.create_lwg(gwg), // Inputs and Outputs qmat2_data, qmat2, diff --git a/backends/vulkan/runtime/graph/ops/impl/Tan.cpp b/backends/vulkan/runtime/graph/ops/impl/Tan.cpp index 223f082d6a6..7f9db0e2ef4 100644 --- a/backends/vulkan/runtime/graph/ops/impl/Tan.cpp +++ b/backends/vulkan/runtime/graph/ops/impl/Tan.cpp @@ -43,8 +43,8 @@ void add_tan_node(ComputeGraph& graph, const ValueRef in, const ValueRef out) { graph.execute_nodes().emplace_back(new DynamicDispatchNode( graph, VK_KERNEL_FROM_STR(kernel_name), - default_pick_global_wg_size, - default_pick_local_wg_size, + default_pick_gwg, + default_pick_lwg, // Inputs and Outputs {{out, vkapi::kWrite}, {in, vkapi::kRead}}, // Shader params buffers diff --git a/backends/vulkan/runtime/graph/ops/impl/Transfer.cpp b/backends/vulkan/runtime/graph/ops/impl/Transfer.cpp index 12ced29bba9..d51070bf5f5 100644 --- a/backends/vulkan/runtime/graph/ops/impl/Transfer.cpp +++ b/backends/vulkan/runtime/graph/ops/impl/Transfer.cpp @@ -94,8 +94,8 @@ void add_transfer_copy_node( graph.execute_nodes().emplace_back(new DynamicDispatchNode( graph, VK_KERNEL_FROM_STR(kernel_name), - default_pick_global_wg_size, - default_pick_local_wg_size, + default_pick_gwg, + default_pick_lwg, // Inputs and Outputs {{out, vkapi::kWrite}, {in, vkapi::kRead}}, // Parameter buffers diff --git a/backends/vulkan/runtime/graph/ops/impl/UnaryOp.cpp b/backends/vulkan/runtime/graph/ops/impl/UnaryOp.cpp index 79d60b3f870..6a50cb2f6a9 100644 --- a/backends/vulkan/runtime/graph/ops/impl/UnaryOp.cpp +++ b/backends/vulkan/runtime/graph/ops/impl/UnaryOp.cpp @@ -48,8 +48,8 @@ void add_unary_op_node( graph.execute_nodes().emplace_back(new DynamicDispatchNode( graph, VK_KERNEL_FROM_STR(kernel_name), - default_pick_global_wg_size, - default_pick_local_wg_size, + default_pick_gwg, + default_pick_lwg, // Inputs and Outputs {{out, vkapi::kWrite}, {in, vkapi::kRead}}, // Shader params buffers diff --git a/backends/vulkan/runtime/graph/ops/impl/Unfold.cpp b/backends/vulkan/runtime/graph/ops/impl/Unfold.cpp index bcab825198e..80e094ddbfb 100644 --- a/backends/vulkan/runtime/graph/ops/impl/Unfold.cpp +++ b/backends/vulkan/runtime/graph/ops/impl/Unfold.cpp @@ -84,8 +84,8 @@ void add_unfold_copy_node( graph.execute_nodes().emplace_back(new DynamicDispatchNode( graph, VK_KERNEL_FROM_STR(kernel_name), - default_pick_global_wg_size, - default_pick_local_wg_size, + default_pick_gwg, + default_pick_lwg, {{output, vkapi::kWrite}, {input, vkapi::kRead}}, {graph.meta_ubo(output), graph.meta_ubo(input)}, {}, diff --git a/backends/vulkan/runtime/graph/ops/impl/Upsample.cpp b/backends/vulkan/runtime/graph/ops/impl/Upsample.cpp index 6662ae367c5..6403ae80ed0 100644 --- a/backends/vulkan/runtime/graph/ops/impl/Upsample.cpp +++ b/backends/vulkan/runtime/graph/ops/impl/Upsample.cpp @@ -118,8 +118,8 @@ void add_upsample_nearest2d_node( graph.execute_nodes().emplace_back(new DynamicDispatchNode( graph, VK_KERNEL_FROM_STR(kernel_name), - default_pick_global_wg_size, - default_pick_local_wg_size, + default_pick_gwg, + default_pick_lwg, // Inputs and Outputs {{out, vkapi::MemoryAccessType::WRITE}, {in, vkapi::MemoryAccessType::READ}}, diff --git a/backends/vulkan/runtime/graph/ops/impl/Var.cpp b/backends/vulkan/runtime/graph/ops/impl/Var.cpp index d8fd367f18a..8747fc9ac58 100644 --- a/backends/vulkan/runtime/graph/ops/impl/Var.cpp +++ b/backends/vulkan/runtime/graph/ops/impl/Var.cpp @@ -16,29 +16,13 @@ namespace vkcompute { using namespace utils; // Custom global workgroup size function for var_buffer -utils::uvec3 var_buffer_global_wg_size( +GlobalWorkGrid var_buffer_gwg( ComputeGraph* graph, const vkapi::ShaderInfo& shader, const std::vector& args, const std::vector& resize_args) { (void)shader; - (void)resize_args; const ValueRef out = args.at(0).refs.at(0); - return { - graph->size_at(-1, out), - graph->size_at(-2, out), - graph->size_at(-3, out) * graph->size_at(-4, out)}; -} - -// Custom local workgroup size function for var_buffer -utils::uvec3 var_buffer_local_wg_size( - ComputeGraph* graph, - const vkapi::ShaderInfo& shader, - const utils::uvec3& global_workgroup_size, - const std::vector& args, - const std::vector& resize_args) { - (void)shader; - (void)global_workgroup_size; const ValueRef in = args.at(1).refs.at(0); const int dim = resize_args.at(0); @@ -46,14 +30,18 @@ utils::uvec3 var_buffer_local_wg_size( int32_t reduce_dim = normalize(dim, ndim); reduce_dim = nchw_dim_to_whcn_dim(reduce_dim, ndim); - const uint32_t nworkers_per_group = 4; - utils::uvec3 local_wg_size{1, 1, 1}; - local_wg_size[reduce_dim] = nworkers_per_group; - return local_wg_size; + utils::uvec3 lwg_extents{1u, 1u, 1u}; + lwg_extents[reduce_dim] = 4u; + return GlobalWorkGrid( + {graph->size_at(-1, out), + graph->size_at(-2, out), + graph->size_at(-3, out) * graph->size_at(-4, out)}, + kTiledWorkGrid, + LocalWorkGroup(lwg_extents)); } // Custom global workgroup size function for var_texture -utils::uvec3 var_texture_global_wg_size( +GlobalWorkGrid var_texture_gwg( ComputeGraph* graph, const vkapi::ShaderInfo& shader, const std::vector& args, @@ -67,39 +55,18 @@ utils::uvec3 var_texture_global_wg_size( int32_t reduce_dim = normalize(dim, ndim); reduce_dim = nchw_dim_to_whcn_dim(reduce_dim, ndim); - utils::uvec3 global_wg_size = graph->logical_limits_of(out); - global_wg_size[reduce_dim] = 1; - return global_wg_size; -} - -// Custom local workgroup size function for var_texture -utils::uvec3 var_texture_local_wg_size( - ComputeGraph* graph, - const vkapi::ShaderInfo& shader, - const utils::uvec3& global_workgroup_size, - const std::vector& args, - const std::vector& resize_args) { - (void)shader; - const ValueRef in = args.at(1).refs.at(0); - const int dim = resize_args.at(0); - - const int64_t ndim = graph->dim_of(in); - int32_t reduce_dim = normalize(dim, ndim); - reduce_dim = nchw_dim_to_whcn_dim(reduce_dim, ndim); - - const uint32_t nworkers_per_group = 4; - const uint32_t ngroups = 4; - - utils::uvec3 local_wg_size{1, 1, 1}; - local_wg_size[reduce_dim] = nworkers_per_group; + utils::uvec3 extents = graph->logical_limits_of(out); + extents[reduce_dim] = 1; + utils::uvec3 lwg_extents{1u, 1u, 1u}; + lwg_extents[reduce_dim] = 4u; const int other_dim_1 = (reduce_dim + 1) % 3; const int other_dim_2 = (reduce_dim + 2) % 3; - if (global_workgroup_size[other_dim_1] > global_workgroup_size[other_dim_2]) { - local_wg_size[other_dim_1] = ngroups; + if (extents[other_dim_1] > extents[other_dim_2]) { + lwg_extents[other_dim_1] = 4u; } else { - local_wg_size[other_dim_2] = ngroups; + lwg_extents[other_dim_2] = 4u; } - return local_wg_size; + return GlobalWorkGrid(extents, kTiledWorkGrid, LocalWorkGroup(lwg_extents)); } void resize_var_node( @@ -142,16 +109,6 @@ void add_var_buffer_node( add_storage_type_suffix(kernel_name, graph.storage_type_of(out)); add_dtype_suffix(kernel_name, graph.dtype_of(out)); - const uint32_t nworkers_per_group = 4; - - utils::uvec3 global_wg_size = { - graph.size_at(-1, out), - graph.size_at(-2, out), - graph.size_at(-3, out) * graph.size_at(-4, out)}; - - utils::uvec3 local_wg_size{1, 1, 1}; - local_wg_size[reduce_dim] = nworkers_per_group; - std::vector push_constants; int32_t unbiased_int = static_cast(unbiased); push_constants.emplace_back(&unbiased_int, sizeof(unbiased_int)); @@ -159,8 +116,8 @@ void add_var_buffer_node( graph.execute_nodes().emplace_back(new DynamicDispatchNode( graph, VK_KERNEL_FROM_STR(kernel_name), - var_buffer_global_wg_size, - var_buffer_local_wg_size, + var_buffer_gwg, + pick_required_lwg, // Inputs and Outputs {{out, vkapi::kWrite}, {in, vkapi::kRead}}, // Shader params buffers @@ -211,19 +168,14 @@ void add_var_texture_node( const uint32_t ngroups = 4; VK_CHECK_COND(nworkers_per_group * ngroups <= max_nthreads); - utils::uvec3 global_wg_size = graph.logical_limits_of(out); - global_wg_size[reduce_dim] = 1; - - utils::uvec3 local_wg_size{1, 1, 1}; - local_wg_size[reduce_dim] = nworkers_per_group; + utils::uvec3 extents = graph.logical_limits_of(out); + extents[reduce_dim] = 1; const int other_dim_1 = (reduce_dim + 1) % 3; const int other_dim_2 = (reduce_dim + 2) % 3; int32_t group_dim; - if (global_wg_size[other_dim_1] > global_wg_size[other_dim_2]) { - local_wg_size[other_dim_1] = ngroups; + if (extents[other_dim_1] > extents[other_dim_2]) { group_dim = other_dim_1; } else { - local_wg_size[other_dim_2] = ngroups; group_dim = other_dim_2; } @@ -234,8 +186,8 @@ void add_var_texture_node( graph.execute_nodes().emplace_back(new DynamicDispatchNode( graph, VK_KERNEL_FROM_STR(kernel_name), - var_texture_global_wg_size, - var_texture_local_wg_size, + var_texture_gwg, + pick_required_lwg, // Inputs and Outputs {{out, vkapi::kWrite}, {in, vkapi::kRead}}, // Shader params buffers diff --git a/backends/vulkan/runtime/graph/ops/impl/View.cpp b/backends/vulkan/runtime/graph/ops/impl/View.cpp index 6dea5efa9ea..4ff4a4b5fd9 100644 --- a/backends/vulkan/runtime/graph/ops/impl/View.cpp +++ b/backends/vulkan/runtime/graph/ops/impl/View.cpp @@ -93,8 +93,8 @@ void add_view_copy_node( graph.execute_nodes().emplace_back(new DynamicDispatchNode( graph, VK_KERNEL_FROM_STR(kernel_name), - default_pick_global_wg_size, - default_pick_local_wg_size, + default_pick_gwg, + default_pick_lwg, // Inputs and Outputs {{out, vkapi::kWrite}, {in, vkapi::kRead}}, // Parameter Buffers @@ -135,8 +135,8 @@ void add_view_copy_convert_node( graph.execute_nodes().emplace_back(new DynamicDispatchNode( graph, VK_KERNEL_FROM_STR(kernel_name), - default_pick_global_wg_size, - default_pick_local_wg_size, + default_pick_gwg, + default_pick_lwg, // Inputs and Outputs {{out, vkapi::kWrite}, {in, vkapi::kRead}}, // Parameter Buffers diff --git a/backends/vulkan/runtime/graph/ops/impl/Where.cpp b/backends/vulkan/runtime/graph/ops/impl/Where.cpp index ebd12eabf2e..dd95ffca057 100644 --- a/backends/vulkan/runtime/graph/ops/impl/Where.cpp +++ b/backends/vulkan/runtime/graph/ops/impl/Where.cpp @@ -59,8 +59,8 @@ void add_where_node( graph.execute_nodes().emplace_back(new DynamicDispatchNode( graph, VK_KERNEL_FROM_STR(kernel_name), - default_pick_global_wg_size, - default_pick_local_wg_size, + default_pick_gwg, + default_pick_lwg, // Inputs and Outputs {{out, vkapi::kWrite}, {{cond, self, other}, vkapi::kRead}}, // Parameter buffers diff --git a/backends/vulkan/runtime/graph/ops/impl/utils/TensorUtils.cpp b/backends/vulkan/runtime/graph/ops/impl/utils/TensorUtils.cpp index a52572289a4..df83bdec7e0 100644 --- a/backends/vulkan/runtime/graph/ops/impl/utils/TensorUtils.cpp +++ b/backends/vulkan/runtime/graph/ops/impl/utils/TensorUtils.cpp @@ -77,20 +77,20 @@ utils::ivec2 create_broadcast_params( // Work group size calculation functions // -utils::uvec3 adaptive_work_group_size(const utils::uvec3& global_work_group) { - utils::uvec3 local_group_size = {4, 4, 4}; - if (global_work_group[2u] == 1) { - if (global_work_group[1u] < 8) { - local_group_size[0u] = 16; - local_group_size[1u] = 4; - local_group_size[2u] = 1; +LocalWorkGroup adaptive_lwg(const GlobalWorkGrid& gwg) { + utils::uvec3 extents = {4u, 4u, 4u}; + if (gwg[2u] == 1) { + if (gwg[1u] < 8) { + extents[0u] = 16u; + extents[1u] = 4u; + extents[2u] = 1u; } else { - local_group_size[0u] = 8; - local_group_size[1u] = 8; - local_group_size[2u] = 1; + extents[0u] = 8u; + extents[1u] = 8u; + extents[2u] = 1u; } } - return local_group_size; + return LocalWorkGroup(extents); } } // namespace vkcompute diff --git a/backends/vulkan/runtime/graph/ops/impl/utils/TensorUtils.h b/backends/vulkan/runtime/graph/ops/impl/utils/TensorUtils.h index 05234c7790f..f8787e61afd 100644 --- a/backends/vulkan/runtime/graph/ops/impl/utils/TensorUtils.h +++ b/backends/vulkan/runtime/graph/ops/impl/utils/TensorUtils.h @@ -48,7 +48,7 @@ utils::ivec2 create_broadcast_params( // Work group size calculation functions // -utils::uvec3 adaptive_work_group_size(const utils::uvec3& global_work_group); +LocalWorkGroup adaptive_lwg(const GlobalWorkGrid& gwg); // // Tensor dim utilities diff --git a/backends/vulkan/runtime/vk_api/Adapter.h b/backends/vulkan/runtime/vk_api/Adapter.h index a1b7f2962ec..6eae09e8eb6 100644 --- a/backends/vulkan/runtime/vk_api/Adapter.h +++ b/backends/vulkan/runtime/vk_api/Adapter.h @@ -427,6 +427,30 @@ class Adapter final { return physical_device_.properties.limits.maxStorageBufferRange; } + inline utils::uvec3 max_compute_workgroup_count() const { + const auto& limits = physical_device_.properties.limits; + return { + limits.maxComputeWorkGroupCount[0], + limits.maxComputeWorkGroupCount[1], + limits.maxComputeWorkGroupCount[2]}; + } + + inline utils::uvec3 max_compute_workgroup_size() const { + const auto& limits = physical_device_.properties.limits; + return { + limits.maxComputeWorkGroupSize[0], + limits.maxComputeWorkGroupSize[1], + limits.maxComputeWorkGroupSize[2]}; + } + + inline uint32_t max_compute_workgroup_invocations() const { + return physical_device_.properties.limits.maxComputeWorkGroupInvocations; + } + + inline uint32_t recommended_lwg_nthreads() const { + return 64u; + } + // Command Buffer Submission void submit_cmd( diff --git a/backends/vulkan/runtime/vk_api/Command.cpp b/backends/vulkan/runtime/vk_api/Command.cpp index c4e1cb20cf4..d3a11a373ea 100644 --- a/backends/vulkan/runtime/vk_api/Command.cpp +++ b/backends/vulkan/runtime/vk_api/Command.cpp @@ -81,7 +81,7 @@ void CommandBuffer::end() { void CommandBuffer::bind_pipeline( VkPipeline pipeline, VkPipelineLayout pipeline_layout, - const LocalWorkGroup& local_workgroup_size) { + const LocalWorkGroup& lwg) { VK_CHECK_COND( state_ == CommandBuffer::State::RECORDING, "Vulkan CommandBuffer: called bind_pipeline() on a command buffer whose state " @@ -94,7 +94,7 @@ void CommandBuffer::bind_pipeline( } bound_.pipeline_layout = pipeline_layout; - bound_.local_workgroup_size = local_workgroup_size; + bound_.lwg = lwg; state_ = CommandBuffer::State::PIPELINE_BOUND; } @@ -191,18 +191,23 @@ void CommandBuffer::insert_barrier_only(PipelineBarrier& pipeline_barrier) { record_barrier(pipeline_barrier); } -void CommandBuffer::dispatch(const utils::uvec3& global_workgroup_size) { +void CommandBuffer::dispatch( + const GlobalWorkGrid& gwg, + const LocalWorkGroup& lwg) { VK_CHECK_COND( state_ == CommandBuffer::State::BARRIERS_INSERTED, "Vulkan CommandBuffer: called dispatch() on a command buffer whose state " "is not BARRIERS_INSERTED."); + VK_CHECK_COND( + bound_.lwg == lwg, + "Vulkan CommandBuffer: dispatch local workgroup does not match bound " + "pipeline."); vkCmdDispatch( handle_, - utils::div_up(global_workgroup_size[0u], bound_.local_workgroup_size[0u]), - utils::div_up(global_workgroup_size[1u], bound_.local_workgroup_size[1u]), - utils::div_up( - global_workgroup_size[2u], bound_.local_workgroup_size[2u])); + utils::div_up(gwg[0u], lwg[0u]), + utils::div_up(gwg[1u], lwg[1u]), + utils::div_up(gwg[2u], lwg[2u])); state_ = CommandBuffer::State::RECORDING; } diff --git a/backends/vulkan/runtime/vk_api/Command.h b/backends/vulkan/runtime/vk_api/Command.h index 4797f788c5a..7c2922fb546 100644 --- a/backends/vulkan/runtime/vk_api/Command.h +++ b/backends/vulkan/runtime/vk_api/Command.h @@ -51,19 +51,19 @@ class CommandBuffer final { struct Bound { VkPipeline pipeline; VkPipelineLayout pipeline_layout; - LocalWorkGroup local_workgroup_size; + LocalWorkGroup lwg; VkDescriptorSet descriptors; explicit Bound() : pipeline{VK_NULL_HANDLE}, pipeline_layout{VK_NULL_HANDLE}, - local_workgroup_size{0u, 0u, 0u}, + lwg{0u, 0u, 0u}, descriptors{VK_NULL_HANDLE} {} inline void reset() { pipeline = VK_NULL_HANDLE; pipeline_layout = VK_NULL_HANDLE; - local_workgroup_size = LocalWorkGroup{0u, 0u, 0u}; + lwg = LocalWorkGroup{0u, 0u, 0u}; descriptors = VK_NULL_HANDLE; } }; @@ -95,7 +95,7 @@ class CommandBuffer final { void insert_barrier(PipelineBarrier& pipeline_barrier); void insert_barrier_only(PipelineBarrier& pipeline_barrier); - void dispatch(const utils::uvec3&); + void dispatch(const GlobalWorkGrid&, const LocalWorkGroup&); void blit(vkapi::VulkanImage& src, vkapi::VulkanImage& dst); void write_timestamp(VkQueryPool, const uint32_t) const; diff --git a/backends/vulkan/runtime/vk_api/DispatchGrid.cpp b/backends/vulkan/runtime/vk_api/DispatchGrid.cpp index 6f805789e07..a95704e188f 100644 --- a/backends/vulkan/runtime/vk_api/DispatchGrid.cpp +++ b/backends/vulkan/runtime/vk_api/DispatchGrid.cpp @@ -308,6 +308,10 @@ const utils::uvec3& GlobalWorkGrid::extents() const { return extents_; } +uint32_t GlobalWorkGrid::operator[](const int idx) const { + return extents_[idx]; +} + const LocalWorkGroup& GlobalWorkGrid::required_lwg_size() const { return required_lwg_; } @@ -323,7 +327,20 @@ bool GlobalWorkGrid::is_linear() const { void GlobalWorkGrid::wrap_linear_dispatch( const utils::uvec3& max_wg_count, const uint32_t target_total_nthreads) { - if (!is_linear() || required_lwg_.is_valid()) { + wrap_linear_dispatch( + max_wg_count, LocalWorkGroup(kLinearLwg, target_total_nthreads)); +} + +void GlobalWorkGrid::wrap_linear_dispatch( + const utils::uvec3& max_wg_count, + const LocalWorkGroup& required_lwg) { + if (!is_linear()) { + return; + } + if (required_lwg_.is_valid()) { + VK_CHECK_COND( + required_lwg_ == required_lwg, + "Linear dispatch local workgroup size must match its wrapping hint"); return; } @@ -331,33 +348,70 @@ void GlobalWorkGrid::wrap_linear_dispatch( extents_[1] == 1u && extents_[2] == 1u, "Linear dispatch wrapping requires one-dimensional input extents"); VK_CHECK_COND( - max_wg_count[0] > 0u && max_wg_count[1] > 0u, - "Linear dispatch requires nonzero X and Y workgroup limits"); + required_lwg.is_valid(), + "Linear dispatch requires nonzero local workgroup dimensions"); + VK_CHECK_COND( + max_wg_count[0] > 0u && max_wg_count[1] > 0u && max_wg_count[2] > 0u, + "Linear dispatch requires nonzero workgroup limits"); - const LocalWorkGroup required_lwg(kLinearLwg, target_total_nthreads); const uint64_t lwg_x = required_lwg.x(); - const uint64_t required_workgroups = - utils::div_up(uint64_t(extents_[0]), lwg_x); - if (required_workgroups <= max_wg_count[0]) { + const uint64_t lwg_y = required_lwg.y(); + const uint64_t lwg_plane = lwg_x * lwg_y; + const uint64_t required_wgs = utils::div_up(uint64_t(extents_[0]), lwg_plane); + if (required_wgs <= max_wg_count[0]) { + extents_ = { + utils::safe_downcast( + utils::div_up(uint64_t(extents_[0]), lwg_y)), + utils::safe_downcast(lwg_y), + 1u}; required_lwg_ = required_lwg; return; } const uint64_t square_width = static_cast( - std::ceil(std::sqrt(static_cast(required_workgroups)))); + std::ceil(std::sqrt(static_cast(required_wgs)))); + const uint64_t min_width = + utils::div_up(required_wgs, uint64_t(max_wg_count[1])); const uint64_t workgroups_x = - std::min(square_width, max_wg_count[0]); - const uint64_t workgroups_y = - utils::div_up(required_workgroups, workgroups_x); + std::min(std::max(square_width, min_width), max_wg_count[0]); + const uint64_t workgroups_y = utils::div_up(required_wgs, workgroups_x); VK_CHECK_COND( workgroups_y <= max_wg_count[1], "Linear dispatch exceeds two-dimensional workgroup limits"); extents_ = { utils::safe_downcast(workgroups_x * lwg_x), - utils::safe_downcast(workgroups_y), + utils::safe_downcast(workgroups_y * lwg_y), 1u}; required_lwg_ = required_lwg; } +void GlobalWorkGrid::validate( + const LocalWorkGroup& lwg, + const utils::uvec3& max_wg_count, + const LocalWorkGroup& output_tile) const { + VK_CHECK_COND(output_tile.is_valid(), "Dispatch dimensions must be nonzero"); + VK_CHECK_COND( + !required_lwg_.is_valid() || required_lwg_ == lwg, + "Linear dispatch local workgroup size does not match wrapping hint"); + + utils::uvec3 effective_extents = { + utils::div_up(extents_[0], output_tile[0]), + utils::div_up(extents_[1], output_tile[1]), + utils::div_up(extents_[2], output_tile[2])}; + if (effective_extents[0] == 0u || effective_extents[1] == 0u || + effective_extents[2] == 0u) { + effective_extents = {1u, 1u, 1u}; + } + + const utils::uvec3 wg_count = { + utils::div_up(effective_extents[0], lwg[0]), + utils::div_up(effective_extents[1], lwg[1]), + utils::div_up(effective_extents[2], lwg[2])}; + VK_CHECK_COND( + wg_count[0] <= max_wg_count[0] && wg_count[1] <= max_wg_count[1] && + wg_count[2] <= max_wg_count[2], + "Shader dispatch exceeds device workgroup count limits"); +} + } // namespace vkcompute diff --git a/backends/vulkan/runtime/vk_api/DispatchGrid.h b/backends/vulkan/runtime/vk_api/DispatchGrid.h index e7e9f75d09d..5bb085a7aea 100644 --- a/backends/vulkan/runtime/vk_api/DispatchGrid.h +++ b/backends/vulkan/runtime/vk_api/DispatchGrid.h @@ -108,6 +108,7 @@ class GlobalWorkGrid final { bool operator!=(const GlobalWorkGrid& other) const; const utils::uvec3& extents() const; + uint32_t operator[](int idx) const; const LocalWorkGroup& required_lwg_size() const; DispatchGridIntent intent() const; bool is_linear() const; @@ -115,6 +116,14 @@ class GlobalWorkGrid final { void wrap_linear_dispatch( const utils::uvec3& max_wg_count, uint32_t target_total_nthreads = 64u); + void wrap_linear_dispatch( + const utils::uvec3& max_wg_count, + const LocalWorkGroup& required_lwg); + + void validate( + const LocalWorkGroup& lwg, + const utils::uvec3& max_wg_count, + const LocalWorkGroup& output_tile = LocalWorkGroup(1u, 1u, 1u)) const; }; } // namespace vkcompute diff --git a/backends/vulkan/runtime/vk_api/QueryPool.cpp b/backends/vulkan/runtime/vk_api/QueryPool.cpp index e8b3ca55206..51e121fc716 100644 --- a/backends/vulkan/runtime/vk_api/QueryPool.cpp +++ b/backends/vulkan/runtime/vk_api/QueryPool.cpp @@ -40,7 +40,7 @@ QueryPool::QueryPool(const QueryPoolConfig& config, const Adapter* adapter_p) device_(VK_NULL_HANDLE), querypool_(VK_NULL_HANDLE), num_queries_(0u), - shader_durations_(0), + shader_durations_{}, mutex_{} { initialize(adapter_p); } @@ -104,8 +104,8 @@ void QueryPool::shader_profile_begin( const CommandBuffer& cmd, const uint32_t dispatch_id, const std::string& kernel_name, - const VkExtent3D global_workgroup_size, - const VkExtent3D local_workgroup_size) { + const GlobalWorkGrid& gwg, + const LocalWorkGroup& lwg) { EARLY_RETURN_IF_UNINITIALIZED(); std::lock_guard lock(mutex_); @@ -116,8 +116,8 @@ void QueryPool::shader_profile_begin( // Execution Properties dispatch_id, kernel_name, - global_workgroup_size, - local_workgroup_size, + gwg, + lwg, // Query indexes query_idx, // start query idx UINT32_MAX, // end query idx @@ -164,16 +164,14 @@ void QueryPool::extract_results() { } } -std::ostream& operator<<(std::ostream& os, const VkExtent3D& extents) { - os << "{" << extents.width << ", " << extents.height << ", " << extents.depth - << "}"; +std::ostream& operator<<(std::ostream& os, const utils::uvec3& extents) { + os << "{" << extents[0] << ", " << extents[1] << ", " << extents[2] << "}"; return os; } -std::string stringize(const VkExtent3D& extents) { +std::string stringize(const utils::uvec3& extents) { std::stringstream ss; - ss << "{" << extents.width << ", " << extents.height << ", " << extents.depth - << "}"; + ss << extents; return ss.str(); } @@ -191,14 +189,10 @@ std::vector QueryPool::get_shader_timestamp_data() { /* .end_time_ns = */ entry.end_time_ns, /* .metadata = */ ShaderMetadata{ - /* .global_workgroup_size = */ - {entry.global_workgroup_size.width, - entry.global_workgroup_size.height, - entry.global_workgroup_size.depth}, - /* .local_workgroup_size = */ - {entry.local_workgroup_size.width, - entry.local_workgroup_size.height, - entry.local_workgroup_size.depth}, + /* .gwg = */ + {entry.gwg[0u], entry.gwg[1u], entry.gwg[2u]}, + /* .lwg = */ + {entry.lwg.x(), entry.lwg.y(), entry.lwg.z()}, }}); } return shader_result; @@ -234,8 +228,9 @@ std::string QueryPool::generate_string_report() { ss << std::left; ss << std::setw(kernel_name_w) << entry.kernel_name; - ss << std::setw(global_size_w) << stringize(entry.global_workgroup_size); - ss << std::setw(local_size_w) << stringize(entry.local_workgroup_size); + ss << std::setw(global_size_w) << stringize(entry.gwg.extents()); + ss << std::setw(local_size_w) + << stringize(static_cast(entry.lwg)); ss << std::right << std::setw(duration_w) << exec_duration_ns.count(); ss << std::endl; } @@ -265,8 +260,8 @@ std::string QueryPool::generate_tsv_string_report() { entry.execution_duration_ns); ss << entry.kernel_name << "\t"; - ss << stringize(entry.global_workgroup_size) << "\t"; - ss << stringize(entry.local_workgroup_size) << "\t"; + ss << stringize(entry.gwg.extents()) << "\t"; + ss << stringize(static_cast(entry.lwg)) << "\t"; ss << exec_duration_ns.count() << "\t"; ss << std::endl; } diff --git a/backends/vulkan/runtime/vk_api/QueryPool.h b/backends/vulkan/runtime/vk_api/QueryPool.h index 94bd99584eb..8357a06521e 100644 --- a/backends/vulkan/runtime/vk_api/QueryPool.h +++ b/backends/vulkan/runtime/vk_api/QueryPool.h @@ -27,8 +27,8 @@ namespace vkcompute { namespace vkapi { struct ShaderMetadata final { - const uint32_t global_workgroup_size[3]; - const uint32_t local_workgroup_size[3]; + const uint32_t gwg[3]; + const uint32_t lwg[3]; }; struct ShaderResult final { @@ -50,8 +50,8 @@ struct ShaderDuration final { // Execution Properties uint32_t dispatch_id; std::string kernel_name; - VkExtent3D global_workgroup_size; - VkExtent3D local_workgroup_size; + GlobalWorkGrid gwg; + LocalWorkGroup lwg; // Query indexes uint32_t start_query_idx; @@ -103,8 +103,8 @@ class QueryPool final { const CommandBuffer&, const uint32_t, const std::string&, - const VkExtent3D, - const VkExtent3D); + const GlobalWorkGrid&, + const LocalWorkGroup&); void shader_profile_end(const CommandBuffer&); diff --git a/backends/vulkan/test/custom_ops/conv2d_utils.h b/backends/vulkan/test/custom_ops/conv2d_utils.h index 416f6c50061..34853ec383a 100644 --- a/backends/vulkan/test/custom_ops/conv2d_utils.h +++ b/backends/vulkan/test/custom_ops/conv2d_utils.h @@ -68,6 +68,7 @@ struct Conv2dConfig { Padding padding; Dilation dilation; int32_t groups; // Number of groups for grouped convolution + int32_t batch = 1; std::string test_case_name = "placeholder"; std::string op_name = "conv2d"; diff --git a/backends/vulkan/test/custom_ops/impl/AddPrototype.cpp b/backends/vulkan/test/custom_ops/impl/AddPrototype.cpp index dc35153baf0..a59cd1ba197 100644 --- a/backends/vulkan/test/custom_ops/impl/AddPrototype.cpp +++ b/backends/vulkan/test/custom_ops/impl/AddPrototype.cpp @@ -31,23 +31,22 @@ vkapi::ShaderInfo pick_add_shader( } // Global workgroup size function for add operations -utils::uvec3 add_global_wg_size( +GlobalWorkGrid add_gwg( ComputeGraph* graph, const vkapi::ShaderInfo& shader, const std::vector& args, const std::vector& resize_args) { - return default_pick_global_wg_size(graph, shader, args, resize_args); + return default_pick_gwg(graph, shader, args, resize_args); } // Local workgroup size function for add operations -utils::uvec3 add_local_wg_size( +LocalWorkGroup add_lwg( ComputeGraph* graph, const vkapi::ShaderInfo& shader, - const utils::uvec3& global_workgroup_size, + const GlobalWorkGrid& gwg, const std::vector& args, const std::vector& resize_args) { - return default_pick_local_wg_size( - graph, shader, global_workgroup_size, args, resize_args); + return default_pick_lwg(graph, shader, gwg, args, resize_args); } void add_prototype(ComputeGraph& graph, const std::vector& args) { @@ -86,8 +85,8 @@ void add_prototype(ComputeGraph& graph, const std::vector& args) { graph.execute_nodes().emplace_back(new DynamicDispatchNode( graph, pick_add_shader, - add_global_wg_size, - add_local_wg_size, + add_gwg, + add_lwg, // Inputs and Outputs {{output, vkapi::kWrite}, {{input_a, input_b}, vkapi::kRead}}, // Shader params buffers diff --git a/backends/vulkan/test/custom_ops/impl/TestConv2dDw.cpp b/backends/vulkan/test/custom_ops/impl/TestConv2dDw.cpp index 8709e4bdc2c..85c2308b11f 100644 --- a/backends/vulkan/test/custom_ops/impl/TestConv2dDw.cpp +++ b/backends/vulkan/test/custom_ops/impl/TestConv2dDw.cpp @@ -81,7 +81,7 @@ static std::string pick_conv2d_dw_shader_with_selector( return kernel_name; } -static utils::uvec3 conv2d_dw_global_wg_size_fn( +static GlobalWorkGrid conv2d_dw_gwg_fn( ComputeGraph* graph, const vkapi::ShaderInfo& shader, const std::vector& args, @@ -94,10 +94,12 @@ static utils::uvec3 conv2d_dw_global_wg_size_fn( if (uses_output_tile) { const bool is_sned = shader.kernel_name.find("_sned") != std::string::npos; - const utils::uvec3 image_extents = graph->create_global_wg_size(out); + const utils::uvec3 image_extents = graph->create_gwg(out).extents(); if (is_sned) { - return {image_extents[0] * image_extents[1], image_extents[2], 1}; + return GlobalWorkGrid( + {image_extents[0] * image_extents[1], image_extents[2], 1u}, + kTiledWorkGrid); } uint32_t batch_x = 4u; @@ -109,25 +111,27 @@ static utils::uvec3 conv2d_dw_global_wg_size_fn( uint32_t scaled_x = utils::div_up(image_extents[0], batch_x); uint32_t scaled_y = utils::div_up(image_extents[1], batch_y); - return {scaled_x * scaled_y, image_extents[2], 1}; + return GlobalWorkGrid( + {scaled_x * scaled_y, image_extents[2], 1u}, kTiledWorkGrid); } - const utils::uvec3 base_extents = graph->create_global_wg_size(out); - return {base_extents[0] * base_extents[1] * base_extents[2], 1, 1}; + const utils::uvec3 base_extents = graph->create_gwg(out).extents(); + return graph->create_linear_gwg( + base_extents[0] * base_extents[1] * base_extents[2]); } -static utils::uvec3 conv2d_dw_local_wg_size_fn( +static LocalWorkGroup conv2d_dw_lwg_fn( ComputeGraph* graph, const vkapi::ShaderInfo& shader, - const utils::uvec3& global_workgroup_size, + const GlobalWorkGrid& gwg, const std::vector& args, const std::vector& resize_args) { (void)graph; (void)shader; - (void)global_workgroup_size; + (void)gwg; (void)args; (void)resize_args; - return {64, 1, 1}; + return LocalWorkGroup(64u, 1u, 1u); } static ValueRef prepack_dw_weights(ComputeGraph& graph, const ValueRef vref) { @@ -156,8 +160,8 @@ static ValueRef prepack_dw_weights(ComputeGraph& graph, const ValueRef vref) { graph.prepack_nodes().emplace_back(new PrepackNode( graph, VK_KERNEL_FROM_STR(kernel_name), - graph.create_global_wg_size(v), - graph.create_local_wg_size(v), + graph.create_gwg(v), + graph.create_lwg(v), vref, v, {}, @@ -271,8 +275,8 @@ static void conv2d_dw_with_selector( graph.execute_nodes().emplace_back(new DynamicDispatchNode( graph, shader, - conv2d_dw_global_wg_size_fn, - conv2d_dw_local_wg_size_fn, + conv2d_dw_gwg_fn, + conv2d_dw_lwg_fn, {{out, vkapi::kWrite}, {{in, arg_weight, arg_bias}, vkapi::kRead}}, param_buffers, push_constants, diff --git a/backends/vulkan/test/custom_ops/impl/TestFpaQ4gswLinear.cpp b/backends/vulkan/test/custom_ops/impl/TestFpaQ4gswLinear.cpp index acdbd1de307..a9e5ad96e95 100644 --- a/backends/vulkan/test/custom_ops/impl/TestFpaQ4gswLinear.cpp +++ b/backends/vulkan/test/custom_ops/impl/TestFpaQ4gswLinear.cpp @@ -175,7 +175,7 @@ vkapi::ShaderInfo pick_forced_shader_coop_kc( // NUM_GROUPS). The framework computes num_WGs = div_up(global, local), so the // global x-axis is set to that count directly (with local.x == 1). template -utils::uvec3 pick_q4gsw_coop_global_wg( +GlobalWorkGrid pick_q4gsw_coop_gwg( ComputeGraph* graph, const vkapi::ShaderInfo& shader, const std::vector& args, @@ -188,23 +188,27 @@ utils::uvec3 pick_q4gsw_coop_global_wg( utils::safe_downcast(utils::val_at(-1, out_sizes)); const uint32_t N8 = (N + 7u) / 8u; const uint32_t wgs_along_x = utils::div_up(N8, NUM_GROUPS); - return {wgs_along_x, NUM_GROUPS, WORKERS_PER_GROUP}; + const LocalWorkGroup lwg(1u, NUM_GROUPS, WORKERS_PER_GROUP); + GlobalWorkGrid gwg({wgs_along_x * NUM_GROUPS, 1u, 1u}, kLinearWorkGrid); + gwg.wrap_linear_dispatch( + graph->context()->adapter_ptr()->max_compute_workgroup_count(), lwg); + return gwg; } // Local WG picker for the coop GEMV — LWG=(1, NUM_GROUPS, WORKERS_PER_GROUP). template -utils::uvec3 pick_q4gsw_coop_local_wg( +LocalWorkGroup pick_q4gsw_coop_lwg( ComputeGraph* graph, const vkapi::ShaderInfo& shader, - const utils::uvec3& global_workgroup_size, + const GlobalWorkGrid& gwg, const std::vector& args, const std::vector& resize_args) { (void)graph; (void)shader; - (void)global_workgroup_size; + (void)gwg; (void)args; (void)resize_args; - return {1u, NUM_GROUPS, WORKERS_PER_GROUP}; + return LocalWorkGroup(1u, NUM_GROUPS, WORKERS_PER_GROUP); } // Spec-constant LWG values for the q4gsw_linear_gemv__w_4x8_nc[_nosg] shaders @@ -218,7 +222,7 @@ constexpr uint32_t kGemvNumSubgroups = 4u; // WG pickers for the legacy sg/nosg GEMV shaders. Used only by test selectors // 1 (GEMV_W_4X8) and 2 (GEMV_W_4X8_NOSG); the production dispatcher never // references these shaders. -utils::uvec3 pick_q4gsw_legacy_gemv_global_wg( +GlobalWorkGrid pick_q4gsw_legacy_gemv_gwg( ComputeGraph* graph, const vkapi::ShaderInfo& shader, const std::vector& args, @@ -229,21 +233,21 @@ utils::uvec3 pick_q4gsw_legacy_gemv_global_wg( const uint32_t N = utils::safe_downcast(utils::val_at(-1, graph->sizes_of(out))); // Each thread owns one row-pair along x; y-dim splits K-blocks across waves. - return {N / 2u, kGemvNumSubgroups, 1u}; + return GlobalWorkGrid({N / 2u, kGemvNumSubgroups, 1u}, kTiledWorkGrid); } -utils::uvec3 pick_q4gsw_legacy_gemv_local_wg( +LocalWorkGroup pick_q4gsw_legacy_gemv_lwg( ComputeGraph* graph, const vkapi::ShaderInfo& shader, - const utils::uvec3& global_workgroup_size, + const GlobalWorkGrid& gwg, const std::vector& args, const std::vector& resize_args) { (void)graph; (void)shader; - (void)global_workgroup_size; + (void)gwg; (void)args; (void)resize_args; - return {kGemvSubgroupSize, kGemvNumSubgroups, 1u}; + return LocalWorkGroup(kGemvSubgroupSize, kGemvNumSubgroups, 1u); } // @@ -285,7 +289,7 @@ void legacy_q4gsw_resize_linear_node( graph->virtual_resize(output, new_out_sizes); } -utils::uvec3 legacy_q4gsw_global_wg_size( +GlobalWorkGrid legacy_q4gsw_gwg( ComputeGraph* graph, const vkapi::ShaderInfo& shader, const std::vector& args, @@ -311,23 +315,22 @@ utils::uvec3 legacy_q4gsw_global_wg_size( const uint32_t num_N_tiles = utils::div_up(N, N_per_tile); const uint32_t num_M_tiles = utils::div_up(M, M_per_tile); - return {num_N_tiles, num_M_tiles, 1}; + return GlobalWorkGrid({num_N_tiles, num_M_tiles, 1u}, kTiledWorkGrid); } -utils::uvec3 legacy_q4gsw_local_wg_size( +LocalWorkGroup legacy_q4gsw_lwg( ComputeGraph* graph, const vkapi::ShaderInfo& shader, - const utils::uvec3& global_workgroup_size, + const GlobalWorkGrid& gwg, const std::vector& args, const std::vector& resize_args) { const bool use_coop_algorithm = shader.kernel_name.find("_coop") != std::string::npos; if (use_coop_algorithm) { - return {1, 1, 64}; + return LocalWorkGroup(1u, 1u, 64u); } - return pick_hw_square_wg_size( - graph, shader, global_workgroup_size, args, resize_args); + return pick_xy_square_lwg(graph, shader, gwg, args, resize_args); } vkapi::ShaderInfo legacy_q4gsw_pick_shader( @@ -407,16 +410,17 @@ ValueRef legacy_prepack_q4gsw_weight( qmat2_sizes, vkapi::kInt, storage_type, utils::kWidthPacked); // 4-bit prepack: each thread writes two adjacent blocks along K. - utils::uvec3 global_wg_size = { - utils::safe_downcast(utils::div_up(num_blocks_K, int64_t(2))), - utils::safe_downcast(num_blocks_N), - 1u}; + const GlobalWorkGrid gwg( + {utils::safe_downcast(utils::div_up(num_blocks_K, int64_t(2))), + utils::safe_downcast(num_blocks_N), + 1u}, + kTiledWorkGrid); graph.prepack_nodes().emplace_back(new PrepackNode( graph, VK_KERNEL_FROM_STR(kernel_name), - global_wg_size, - graph.create_local_wg_size(global_wg_size), + gwg, + graph.create_lwg(gwg), qmat2_data, qmat2, // UBOs @@ -472,8 +476,8 @@ void add_legacy_q4gsw_linear_node( graph.execute_nodes().emplace_back(new DynamicDispatchNode( graph, legacy_q4gsw_pick_shader, - legacy_q4gsw_global_wg_size, - legacy_q4gsw_local_wg_size, + legacy_q4gsw_gwg, + legacy_q4gsw_lwg, // Inputs and Outputs (legacy 5-binding layout) {{output, vkapi::kWrite}, {{fp_input, packed_weight, packed_weight_scales, packed_bias}, @@ -542,21 +546,21 @@ void add_q4gsw_linear_coop_kc_forced_node( ComputeGraph*, const std::vector&, const std::vector&); - using PickWgFn = utils::uvec3 (*)( + using PickGwgFn = GlobalWorkGrid (*)( ComputeGraph*, const vkapi::ShaderInfo&, const std::vector&, const std::vector&); - using PickLocalWgFn = utils::uvec3 (*)( + using PickLwgFn = LocalWorkGroup (*)( ComputeGraph*, const vkapi::ShaderInfo&, - const utils::uvec3&, + const GlobalWorkGrid&, const std::vector&, const std::vector&); PickShaderFn pick_shader = nullptr; - PickWgFn pick_global = nullptr; - PickLocalWgFn pick_local = nullptr; + PickGwgFn pick_gwg = nullptr; + PickLwgFn pick_lwg = nullptr; // NOLINTNEXTLINE(clang-diagnostic-switch-enum) switch (kind) { @@ -564,20 +568,20 @@ void add_q4gsw_linear_coop_kc_forced_node( case TestKernelKind::GEMV_COOP_W_4X8_NC_BUFFER_G1W64: pick_shader = pick_forced_shader_coop_kc; - pick_global = pick_q4gsw_coop_global_wg<1u, 64u>; - pick_local = pick_q4gsw_coop_local_wg<1u, 64u>; + pick_gwg = pick_q4gsw_coop_gwg<1u, 64u>; + pick_lwg = pick_q4gsw_coop_lwg<1u, 64u>; break; case TestKernelKind::GEMV_COOP_W_4X8_NC_BUFFER_G4W16: pick_shader = pick_forced_shader_coop_kc< TestKernelKind::GEMV_COOP_W_4X8_NC_BUFFER_G4W16>; - pick_global = pick_q4gsw_coop_global_wg<4u, 16u>; - pick_local = pick_q4gsw_coop_local_wg<4u, 16u>; + pick_gwg = pick_q4gsw_coop_gwg<4u, 16u>; + pick_lwg = pick_q4gsw_coop_lwg<4u, 16u>; break; case TestKernelKind::GEMV_COOP_W_4X8_NC_BUFFER_G8W8: pick_shader = pick_forced_shader_coop_kc< TestKernelKind::GEMV_COOP_W_4X8_NC_BUFFER_G8W8>; - pick_global = pick_q4gsw_coop_global_wg<8u, 8u>; - pick_local = pick_q4gsw_coop_local_wg<8u, 8u>; + pick_gwg = pick_q4gsw_coop_gwg<8u, 8u>; + pick_lwg = pick_q4gsw_coop_lwg<8u, 8u>; break; default: VK_THROW("add_q4gsw_linear_coop_kc_forced_node: non-coop kind"); @@ -586,8 +590,8 @@ void add_q4gsw_linear_coop_kc_forced_node( graph.execute_nodes().emplace_back(new DynamicDispatchNode( graph, pick_shader, - pick_global, - pick_local, + pick_gwg, + pick_lwg, {{output, vkapi::kWrite}, {{fp_input, dummy_transposed_input.vref, @@ -666,43 +670,43 @@ void add_q4gsw_linear_forced_node( ComputeGraph*, const std::vector&, const std::vector&); - using PickWgFn = utils::uvec3 (*)( + using PickGwgFn = GlobalWorkGrid (*)( ComputeGraph*, const vkapi::ShaderInfo&, const std::vector&, const std::vector&); - using PickLocalWgFn = utils::uvec3 (*)( + using PickLwgFn = LocalWorkGroup (*)( ComputeGraph*, const vkapi::ShaderInfo&, - const utils::uvec3&, + const GlobalWorkGrid&, const std::vector&, const std::vector&); PickShaderFn pick_shader = nullptr; - PickWgFn pick_global = nullptr; - PickLocalWgFn pick_local = nullptr; + PickGwgFn pick_gwg = nullptr; + PickLwgFn pick_lwg = nullptr; // NOLINTNEXTLINE(clang-diagnostic-switch-enum) switch (kind) { case TestKernelKind::GEMM_W_4X8: pick_shader = pick_forced_shader; - pick_global = pick_q4gsw_linear_gemm_global_wg; - pick_local = pick_q4gsw_linear_gemm_local_wg; + pick_gwg = pick_q4gsw_linear_gemm_gwg; + pick_lwg = pick_q4gsw_linear_gemm_lwg; break; case TestKernelKind::GEMV_W_4X8: pick_shader = pick_forced_shader; - pick_global = pick_q4gsw_legacy_gemv_global_wg; - pick_local = pick_q4gsw_legacy_gemv_local_wg; + pick_gwg = pick_q4gsw_legacy_gemv_gwg; + pick_lwg = pick_q4gsw_legacy_gemv_lwg; break; case TestKernelKind::GEMM_TIN_W_4X8: pick_shader = pick_forced_shader; - pick_global = pick_q4gsw_linear_tin_gemm_global_wg; - pick_local = pick_q4gsw_linear_tin_gemm_local_wg; + pick_gwg = pick_q4gsw_linear_tin_gemm_gwg; + pick_lwg = pick_q4gsw_linear_tin_gemm_lwg; break; case TestKernelKind::GEMV_W_4X8_NOSG: pick_shader = pick_forced_shader; - pick_global = pick_q4gsw_legacy_gemv_global_wg; - pick_local = pick_q4gsw_legacy_gemv_local_wg; + pick_gwg = pick_q4gsw_legacy_gemv_gwg; + pick_lwg = pick_q4gsw_legacy_gemv_lwg; break; case TestKernelKind::PROD: default: @@ -712,8 +716,8 @@ void add_q4gsw_linear_forced_node( graph.execute_nodes().emplace_back(new DynamicDispatchNode( graph, pick_shader, - pick_global, - pick_local, + pick_gwg, + pick_lwg, {{output, vkapi::kWrite}, {{fp_input, transposed_input_ref, diff --git a/backends/vulkan/test/custom_ops/test_q8ta_conv2d.cpp b/backends/vulkan/test/custom_ops/test_q8ta_conv2d.cpp index 5b1de615d29..b30212feb72 100644 --- a/backends/vulkan/test/custom_ops/test_q8ta_conv2d.cpp +++ b/backends/vulkan/test/custom_ops/test_q8ta_conv2d.cpp @@ -36,9 +36,12 @@ static TestCase create_test_case_from_config( int64_t H_out = config.get_output_height(); int64_t W_out = config.get_output_width(); - // Input tensor (float/half) - [1, C_in, H_in, W_in] (batch size always 1) + // Input tensor (float/half) - [N, C_in, H_in, W_in] std::vector input_size = { - 1, config.channels.in, config.input_size.h, config.input_size.w}; + config.batch, + config.channels.in, + config.input_size.h, + config.input_size.w}; utils::GPUMemoryLayout fp_memory_layout = fp_storage_type == utils::kBuffer ? utils::kWidthPacked @@ -47,7 +50,8 @@ static TestCase create_test_case_from_config( // Create test case name std::string prefix = config.test_case_name.substr(0, 4); // "ACCU" or "PERF" std::string dtype_str = dtype_short(input_dtype); - std::string in_shape = "[1," + std::to_string(config.channels.in) + "," + + std::string in_shape = "[" + std::to_string(config.batch) + "," + + std::to_string(config.channels.in) + "," + std::to_string(config.input_size.h) + "," + std::to_string(config.input_size.w) + "]"; std::string weight_shape = "[" + std::to_string(config.channels.out) + "," + @@ -159,9 +163,9 @@ static TestCase create_test_case_from_config( // Kernel size parameters ValueSpec kernel_size({config.kernel.h, config.kernel.w}); - // Output tensor (float/half) - [1, C_out, H_out, W_out] (batch size always 1) + // Output tensor (float/half) - [N, C_out, H_out, W_out] ValueSpec output( - {1, config.channels.out, H_out, W_out}, + {config.batch, config.channels.out, H_out, W_out}, input_dtype, fp_storage_type, fp_memory_layout, @@ -475,6 +479,69 @@ static std::vector generate_quantized_conv2d_test_cases() { } } + std::vector batch_configs = { + {OutInChannels(16, 32), + InputSize2D(7, 7), + KernelSize(3, 3), + Stride(1, 1), + Padding(1, 1), + Dilation(1, 1), + 1, + 2}, + {OutInChannels(32, 3), + InputSize2D(256, 256), + KernelSize(3, 3), + Stride(2, 2), + Padding(1, 1), + Dilation(1, 1), + 1, + 1}, + {OutInChannels(32, 3), + InputSize2D(256, 256), + KernelSize(3, 3), + Stride(2, 2), + Padding(1, 1), + Dilation(1, 1), + 1, + 60}, + {OutInChannels(512, 256), + InputSize2D(10, 13), + KernelSize(3, 3), + Stride(2, 2), + Padding(1, 1), + Dilation(1, 1), + 1, + 60}}; + + for (auto& config : batch_configs) { + const bool is_performance = config.batch > kRefDimSizeLimit || + config.channels.out > kRefDimSizeLimit || + config.channels.in > kRefDimSizeLimit || + config.input_size.h > kRefDimSizeLimit || + config.input_size.w > kRefDimSizeLimit; + config.op_name = "conv2d_q8ta_q8csw_q8to"; + config.test_case_name = make_test_case_name( + config, is_performance, utils::kTexture3D, utils::kBuffer); + test_cases.push_back(create_test_case_from_config( + config, vkapi::kFloat, utils::kTexture3D, utils::kPackedInt8_4C1W)); + if (config.batch == 2) { + test_cases.push_back(create_test_case_from_config( + config, + vkapi::kFloat, + utils::kTexture3D, + utils::kPackedInt8_4C1W, + /*impl_selector=*/"im2col")); + test_cases.push_back(create_test_case_from_config( + config, vkapi::kFloat, utils::kTexture3D, utils::kPackedInt8_4W4C)); + test_cases.push_back(create_test_case_from_config( + config, + vkapi::kFloat, + utils::kTexture3D, + utils::kPackedInt8_4W4C, + /*impl_selector=*/"im2col")); + } + } + return test_cases; } diff --git a/backends/vulkan/test/custom_ops/test_q8ta_conv2d_dw.cpp b/backends/vulkan/test/custom_ops/test_q8ta_conv2d_dw.cpp index 9813eeaa9d6..2dbb4909adb 100644 --- a/backends/vulkan/test/custom_ops/test_q8ta_conv2d_dw.cpp +++ b/backends/vulkan/test/custom_ops/test_q8ta_conv2d_dw.cpp @@ -37,9 +37,12 @@ TestCase create_test_case_from_config( int64_t H_out = config.get_output_height(); int64_t W_out = config.get_output_width(); - // Input tensor (float/half) - [1, C_in, H_in, W_in] (batch size always 1) + // Input tensor (float/half) - [N, C_in, H_in, W_in] std::vector input_size = { - 1, config.channels.in, config.input_size.h, config.input_size.w}; + config.batch, + config.channels.in, + config.input_size.h, + config.input_size.w}; utils::GPUMemoryLayout fp_memory_layout = fp_storage_type == utils::kBuffer ? utils::kWidthPacked @@ -48,7 +51,8 @@ TestCase create_test_case_from_config( // Create test case name std::string prefix = config.test_case_name.substr(0, 4); // "ACCU" or "PERF" std::string dtype_str = dtype_short(input_dtype); - std::string in_shape = "[1," + std::to_string(config.channels.in) + "," + + std::string in_shape = "[" + std::to_string(config.batch) + "," + + std::to_string(config.channels.in) + "," + std::to_string(config.input_size.h) + "," + std::to_string(config.input_size.w) + "]"; // depthwise: weight is [C_out, 1, K_h, K_w] @@ -168,9 +172,9 @@ TestCase create_test_case_from_config( // Kernel size parameters ValueSpec kernel_size({config.kernel.h, config.kernel.w}); - // Output tensor (float/half) - [1, C_out, H_out, W_out] (batch size always 1) + // Output tensor (float/half) - [N, C_out, H_out, W_out] ValueSpec output( - {1, config.channels.out, H_out, W_out}, + {config.batch, config.channels.out, H_out, W_out}, input_dtype, fp_storage_type, fp_memory_layout, @@ -395,6 +399,46 @@ std::vector generate_quantized_conv2d_dw_test_cases() { } } + std::vector batch_configs = { + {OutInChannels(8, 8), + InputSize2D(8, 8), + KernelSize(3, 3), + Stride(1, 1), + Padding(1, 1), + Dilation(1, 1), + 8, + 2}, + {OutInChannels(128, 128), + InputSize2D(64, 64), + KernelSize(5, 5), + Stride(2, 2), + Padding(2, 2), + Dilation(1, 1), + 128, + 1}, + {OutInChannels(128, 128), + InputSize2D(64, 64), + KernelSize(5, 5), + Stride(2, 2), + Padding(2, 2), + Dilation(1, 1), + 128, + 60}}; + + for (auto& config : batch_configs) { + const bool is_performance = config.channels.out > kRefDimSizeLimit || + config.channels.in > kRefDimSizeLimit; + config.op_name = "conv2d_q8ta_q8csw_q8to"; + config.test_case_name = make_test_case_name( + config, is_performance, utils::kTexture3D, utils::kBuffer); + test_cases.push_back(create_test_case_from_config( + config, vkapi::kFloat, utils::kTexture3D, utils::kPackedInt8_4C1W)); + if (config.batch == 2) { + test_cases.push_back(create_test_case_from_config( + config, vkapi::kFloat, utils::kTexture3D, utils::kPackedInt8_4W4C)); + } + } + return test_cases; } diff --git a/backends/vulkan/test/custom_ops/test_q8ta_conv2d_pw.cpp b/backends/vulkan/test/custom_ops/test_q8ta_conv2d_pw.cpp index d2b84101940..ee7d8c9e5bf 100644 --- a/backends/vulkan/test/custom_ops/test_q8ta_conv2d_pw.cpp +++ b/backends/vulkan/test/custom_ops/test_q8ta_conv2d_pw.cpp @@ -36,9 +36,12 @@ static TestCase create_test_case_from_config( int64_t H_out = config.get_output_height(); int64_t W_out = config.get_output_width(); - // Input tensor (float/half) - [1, C_in, H_in, W_in] (batch size always 1) + // Input tensor (float/half) - [N, C_in, H_in, W_in] std::vector input_size = { - 1, config.channels.in, config.input_size.h, config.input_size.w}; + config.batch, + config.channels.in, + config.input_size.h, + config.input_size.w}; utils::GPUMemoryLayout fp_memory_layout = fp_storage_type == utils::kBuffer ? utils::kWidthPacked @@ -47,7 +50,8 @@ static TestCase create_test_case_from_config( // Create test case name std::string prefix = config.test_case_name.substr(0, 4); // "ACCU" or "PERF" std::string dtype_str = dtype_short(input_dtype); - std::string in_shape = "[1," + std::to_string(config.channels.in) + "," + + std::string in_shape = "[" + std::to_string(config.batch) + "," + + std::to_string(config.channels.in) + "," + std::to_string(config.input_size.h) + "," + std::to_string(config.input_size.w) + "]"; std::string weight_shape = "[" + std::to_string(config.channels.out) + "," + @@ -160,9 +164,9 @@ static TestCase create_test_case_from_config( // Kernel size parameters ValueSpec kernel_size({config.kernel.h, config.kernel.w}); - // Output tensor (float/half) - [1, C_out, H_out, W_out] (batch size always 1) + // Output tensor (float/half) - [N, C_out, H_out, W_out] ValueSpec output( - {1, config.channels.out, H_out, W_out}, + {config.batch, config.channels.out, H_out, W_out}, input_dtype, fp_storage_type, fp_memory_layout, @@ -353,6 +357,46 @@ static std::vector generate_quantized_conv2d_pw_test_cases() { } } + std::vector batch_configs = { + {OutInChannels(8, 8), + InputSize2D(8, 8), + KernelSize(1, 1), + Stride(1, 1), + Padding(0, 0), + Dilation(1, 1), + 1, + 2}, + {OutInChannels(64, 32), + InputSize2D(128, 128), + KernelSize(1, 1), + Stride(1, 1), + Padding(0, 0), + Dilation(1, 1), + 1, + 1}, + {OutInChannels(64, 32), + InputSize2D(128, 128), + KernelSize(1, 1), + Stride(1, 1), + Padding(0, 0), + Dilation(1, 1), + 1, + 60}}; + + for (auto& config : batch_configs) { + const bool is_performance = config.input_size.h > kRefDimSizeLimit || + config.input_size.w > kRefDimSizeLimit; + config.op_name = "conv2d_q8ta_q8csw_q8to"; + config.test_case_name = make_test_case_name( + config, is_performance, utils::kTexture3D, utils::kBuffer); + test_cases.push_back(create_test_case_from_config( + config, vkapi::kFloat, utils::kTexture3D, utils::kPackedInt8_4C1W)); + if (config.batch == 2) { + test_cases.push_back(create_test_case_from_config( + config, vkapi::kFloat, utils::kTexture3D, utils::kPackedInt8_4W4C)); + } + } + return test_cases; } diff --git a/backends/vulkan/test/custom_ops/utils.cpp b/backends/vulkan/test/custom_ops/utils.cpp index 2b8e7208f5e..60c3ddee30c 100644 --- a/backends/vulkan/test/custom_ops/utils.cpp +++ b/backends/vulkan/test/custom_ops/utils.cpp @@ -954,8 +954,8 @@ void BenchmarkResult::add_iter_timing(float time_us) { void BenchmarkResult::add_shader_timing( const std::string& shader_name, float time_us, - const uint32_t global_wg[3], - const uint32_t local_wg[3]) { + const uint32_t gwg[3], + const uint32_t lwg[3]) { // Find existing shader timing or create new one for (auto& st : shader_timings_) { if (st.shader_name == shader_name) { @@ -968,12 +968,12 @@ void BenchmarkResult::add_shader_timing( ShaderTiming new_timing; new_timing.shader_name = shader_name; new_timing.iter_timings_us.push_back(time_us); - new_timing.global_wg_size[0] = global_wg[0]; - new_timing.global_wg_size[1] = global_wg[1]; - new_timing.global_wg_size[2] = global_wg[2]; - new_timing.local_wg_size[0] = local_wg[0]; - new_timing.local_wg_size[1] = local_wg[1]; - new_timing.local_wg_size[2] = local_wg[2]; + new_timing.gwg[0] = gwg[0]; + new_timing.gwg[1] = gwg[1]; + new_timing.gwg[2] = gwg[2]; + new_timing.lwg[0] = lwg[0]; + new_timing.lwg[1] = lwg[1]; + new_timing.lwg[2] = lwg[2]; shader_timings_.push_back(std::move(new_timing)); } @@ -1068,12 +1068,11 @@ void BenchmarkResult::print_summary( const auto& st = shader_timings_[0]; std::cout << std::left << std::setw(OPERATOR_NAME_WIDTH) << truncate_shader_name(st.shader_name) << " " << std::left - << std::setw(GLOBAL_WG_WIDTH) - << format_wg_size(st.global_wg_size) << std::left - << std::setw(LOCAL_WG_WIDTH) << format_wg_size(st.local_wg_size) - << std::left << std::setw(KERNEL_NAME_WIDTH) - << get_kernel_name() << std::right << " " - << std::setw(SIZE_INFO_WIDTH) << size_info + << std::setw(GLOBAL_WG_WIDTH) << format_wg_size(st.gwg) + << std::left << std::setw(LOCAL_WG_WIDTH) + << format_wg_size(st.lwg) << std::left + << std::setw(KERNEL_NAME_WIDTH) << get_kernel_name() + << std::right << " " << std::setw(SIZE_INFO_WIDTH) << size_info << std::setw(TIMING_WIDTH) << std::fixed << std::setprecision(3) << get_avg_time_us() << " μs " << std::setw(GFLOPS_WIDTH) << std::fixed << std::setprecision(3) << total_gflops @@ -1088,10 +1087,9 @@ void BenchmarkResult::print_summary( // Shader lines don't show test case info std::cout << std::left << std::setw(OPERATOR_NAME_WIDTH) << truncate_shader_name(st.shader_name) << " " << std::left - << std::setw(GLOBAL_WG_WIDTH) - << format_wg_size(st.global_wg_size) << std::left - << std::setw(LOCAL_WG_WIDTH) - << format_wg_size(st.local_wg_size) << std::left + << std::setw(GLOBAL_WG_WIDTH) << format_wg_size(st.gwg) + << std::left << std::setw(LOCAL_WG_WIDTH) + << format_wg_size(st.lwg) << std::left << std::setw(KERNEL_NAME_WIDTH) << "" << std::right << " " << std::setw(SIZE_INFO_WIDTH) << "" << std::setw(TIMING_WIDTH) << std::fixed << std::setprecision(3) << shader_avg_time @@ -1626,8 +1624,8 @@ BenchmarkResult execute_test_case( result.add_shader_timing( shader_result.kernel_name, duration_us, - shader_result.metadata.global_workgroup_size, - shader_result.metadata.local_workgroup_size); + shader_result.metadata.gwg, + shader_result.metadata.lwg); } } // gpu_time_us aggregates chained_dispatches worth of shader runs; divide @@ -2160,11 +2158,9 @@ ValueRef quantized_weights_canvas( const ValueRef weight_ref) { const auto original_sizes = graph.sizes_of(weight_ref); - // Get the 2 highest values of original_sizes std::vector sorted_sizes = original_sizes; std::sort(sorted_sizes.begin(), sorted_sizes.end(), std::greater()); int64_t largest1 = sorted_sizes.size() > 0 ? sorted_sizes[0] : 0; - int64_t largest2 = sorted_sizes.size() > 1 ? sorted_sizes[1] : 0; std::vector final_sizes = {1, largest1, largest1}; @@ -2190,19 +2186,14 @@ ValueRef quantized_weights_canvas( ValueRef packed_weight = graph.add_tensor( final_sizes, vkapi::kInt, utils::kTexture3D, utils::kWidthPacked); - utils::uvec3 global_wg_size{ - utils::div_up(utils::safe_downcast(largest1), uint32_t(4)), - utils::safe_downcast(largest2), - utils::safe_downcast(std::min(largest1, int64_t(2048)))}; - std::string kernel_name = "packed_int32_canvas"; add_storage_type_suffix(kernel_name, graph.storage_type_of(packed_weight)); graph.prepack_nodes().emplace_back(new PrepackNode( graph, VK_KERNEL_FROM_STR(kernel_name), - graph.create_global_wg_size(packed_weight), - graph.create_local_wg_size(packed_weight), + graph.create_gwg(packed_weight), + graph.create_lwg(packed_weight), weight_ref, packed_weight, // UBOs @@ -2218,11 +2209,9 @@ ValueRef quantized_weights_canvas( ValueRef float_tensor_canvas(ComputeGraph& graph, const ValueRef weight_ref) { const auto original_sizes = graph.sizes_of(weight_ref); - // Get the 2 highest values of original_sizes std::vector sorted_sizes = original_sizes; std::sort(sorted_sizes.begin(), sorted_sizes.end(), std::greater()); int64_t largest1 = sorted_sizes.size() > 0 ? sorted_sizes[0] : 0; - int64_t largest2 = sorted_sizes.size() > 1 ? sorted_sizes[1] : 0; std::vector final_sizes = {1, largest1, largest1}; @@ -2248,16 +2237,11 @@ ValueRef float_tensor_canvas(ComputeGraph& graph, const ValueRef weight_ref) { ValueRef packed_weight = graph.add_tensor( final_sizes, vkapi::kFloat, utils::kTexture3D, utils::kWidthPacked); - utils::uvec3 global_wg_size{ - utils::div_up(utils::safe_downcast(largest1), uint32_t(4)), - utils::safe_downcast(largest2), - utils::safe_downcast(std::min(largest1, int64_t(2048)))}; - graph.prepack_nodes().emplace_back(new PrepackNode( graph, VK_KERNEL_FROM_STR("float_canvas"), - graph.create_global_wg_size(packed_weight), - graph.create_local_wg_size(packed_weight), + graph.create_gwg(packed_weight), + graph.create_lwg(packed_weight), weight_ref, packed_weight, // UBOs diff --git a/backends/vulkan/test/custom_ops/utils.h b/backends/vulkan/test/custom_ops/utils.h index 81bad5e9df0..2174ceb5618 100644 --- a/backends/vulkan/test/custom_ops/utils.h +++ b/backends/vulkan/test/custom_ops/utils.h @@ -694,8 +694,8 @@ enum class CorrectnessStatus { struct ShaderTiming { std::string shader_name; std::vector iter_timings_us; // Individual iteration timings - uint32_t global_wg_size[3] = {0, 0, 0}; - uint32_t local_wg_size[3] = {0, 0, 0}; + uint32_t gwg[3] = {0, 0, 0}; + uint32_t lwg[3] = {0, 0, 0}; float get_avg_time_us() const { if (iter_timings_us.empty()) { @@ -730,8 +730,8 @@ class BenchmarkResult { void add_shader_timing( const std::string& shader_name, float time_us, - const uint32_t global_wg[3], - const uint32_t local_wg[3]); + const uint32_t gwg[3], + const uint32_t lwg[3]); // Get per-shader timing data const std::vector& get_shader_timings() const { diff --git a/backends/vulkan/test/op_tests/cases.py b/backends/vulkan/test/op_tests/cases.py index 6299a56a0ae..4225619b330 100644 --- a/backends/vulkan/test/op_tests/cases.py +++ b/backends/vulkan/test/op_tests/cases.py @@ -71,13 +71,23 @@ def get_binary_elementwise_inputs(): ] highdim_test_suite.test_name_suffix = "highdim" + large_buffer_test_suite = VkTestSuite( + [ + ((5000000,), (5000000,)), + ] + ) + large_buffer_test_suite.storage_types = ["utils::kBuffer"] + large_buffer_test_suite.layouts = ["utils::kWidthPacked"] + large_buffer_test_suite.data_range = (1, 2) + large_buffer_test_suite.test_name_suffix = "large_buffer" + for suite in [test_suite, highdim_test_suite]: suite.layouts = [ "utils::kWidthPacked", "utils::kChannelsPacked", ] - return [test_suite, highdim_test_suite] + return [test_suite, highdim_test_suite, large_buffer_test_suite] # Eq requires a different test generator so it was split from the other test case. @@ -999,14 +1009,22 @@ def get_view_inputs(): highdim_test_suite.test_name_suffix = "highdim" highdim_test_suite.data_gen = "make_seq_tensor" - for suite in [test_suite, highdim_test_suite]: + large_buffer_test_suite = VkTestSuite( + [ + ((30, 3, 256, 256), (30, 3, 65536)), + ] + ) + large_buffer_test_suite.storage_types = ["utils::kBuffer"] + large_buffer_test_suite.test_name_suffix = "large_buffer" + + for suite in [test_suite, highdim_test_suite, large_buffer_test_suite]: suite.layouts = [ # "utils::kWidthPacked", "utils::kHeightPacked", "utils::kChannelsPacked", ] - return [test_suite, highdim_test_suite] + return [test_suite, highdim_test_suite, large_buffer_test_suite] @register_test_suite("aten.slice_copy.Tensor") diff --git a/backends/vulkan/test/utils/test_utils.cpp b/backends/vulkan/test/utils/test_utils.cpp index 097e3036b29..fdd45baa6b7 100644 --- a/backends/vulkan/test/utils/test_utils.cpp +++ b/backends/vulkan/test/utils/test_utils.cpp @@ -18,6 +18,21 @@ using namespace vkcompute; +namespace { + +GlobalWorkGrid make_linear_dispatch( + api::Context* const context, + const uint64_t numel, + const LocalWorkGroup& lwg) { + vkapi::Adapter* const adapter = context->adapter_ptr(); + GlobalWorkGrid gwg( + {utils::safe_downcast(numel), 1u, 1u}, kLinearWorkGrid); + gwg.wrap_linear_dispatch(adapter->max_compute_workgroup_count(), lwg); + return gwg; +} + +} // namespace + bool is_bitw8(vkapi::ScalarType dtype) { return dtype == vkapi::kByte || dtype == vkapi::kChar || dtype == vkapi::kQInt8 || dtype == vkapi::kQUInt8; @@ -92,12 +107,13 @@ void record_nchw_to_buffer_op( api::vTensor& v_dst) { vkapi::PipelineBarrier pipeline_barrier{}; vkapi::SpecVarList specialization_constants = {v_dst.hashed_layout()}; + const LocalWorkGroup lwg(64u, 1u, 1u); context->submit_compute_job( get_nchw_to_tensor_shader(v_dst, true), pipeline_barrier, - {uint32_t(v_dst.numel()), 1, 1}, - {64, 1, 1}, + make_linear_dispatch(context, v_dst.numel(), lwg), + lwg, specialization_constants, VK_NULL_HANDLE, 0, @@ -114,11 +130,12 @@ void record_buffer_to_nchw_op( api::vTensor& v_src, vkapi::VulkanBuffer& dst_buffer) { vkapi::PipelineBarrier pipeline_barrier{}; + const LocalWorkGroup lwg(64u, 1u, 1u); context->submit_compute_job( get_tensor_to_nchw_shader(v_src, true), pipeline_barrier, - {uint32_t(v_src.numel()), 1, 1}, - {64, 1, 1}, + make_linear_dispatch(context, v_src.numel(), lwg), + lwg, {}, VK_NULL_HANDLE, 0, @@ -137,6 +154,9 @@ void record_nchw_to_image_op( bool int8_buffer_enabled = context->adapter_ptr()->has_full_int8_buffers_support(); auto shader = get_nchw_to_tensor_shader(v_dst, int8_buffer_enabled); + const GlobalWorkGrid gwg( + utils::make_uvec3(v_dst.logical_limits()), kTextureExtentsWorkGrid); + const LocalWorkGroup lwg = adaptive_lwg(gwg); // bitw8 _no_pc shaders expect ivec4 sizes UBO; regular shaders expect // TextureMetadata UBO. The bitw8 path is only used when int8 buffers are @@ -147,8 +167,8 @@ void record_nchw_to_image_op( context->submit_compute_job( shader, pipeline_barrier, - v_dst.logical_limits(), - adaptive_work_group_size(v_dst.logical_limits()), + gwg, + lwg, specialization_constants, VK_NULL_HANDLE, 0, @@ -162,8 +182,8 @@ void record_nchw_to_image_op( context->submit_compute_job( shader, pipeline_barrier, - v_dst.logical_limits(), - adaptive_work_group_size(v_dst.logical_limits()), + gwg, + lwg, specialization_constants, VK_NULL_HANDLE, 0, @@ -186,6 +206,9 @@ void record_image_to_nchw_op( bool int8_buffer_enabled = context->adapter_ptr()->has_full_int8_buffers_support(); auto shader = get_tensor_to_nchw_shader(v_src, int8_buffer_enabled); + const GlobalWorkGrid gwg( + utils::make_uvec3(v_src.logical_limits()), kTextureExtentsWorkGrid); + const LocalWorkGroup lwg = adaptive_lwg(gwg); // bitw8 _no_pc shaders expect ivec4 sizes UBO; regular shaders expect // TextureMetadata UBO. The bitw8 path is only used when int8 buffers are @@ -196,8 +219,8 @@ void record_image_to_nchw_op( context->submit_compute_job( shader, pipeline_barrier, - v_src.logical_limits(), - adaptive_work_group_size(v_src.logical_limits()), + gwg, + lwg, specialization_constants, VK_NULL_HANDLE, 0, @@ -208,8 +231,8 @@ void record_image_to_nchw_op( context->submit_compute_job( shader, pipeline_barrier, - v_src.logical_limits(), - adaptive_work_group_size(v_src.logical_limits()), + gwg, + lwg, specialization_constants, VK_NULL_HANDLE, 0, @@ -225,7 +248,8 @@ void record_bitw8_image_to_nchw_nobitw8buffer_op( api::StagingBuffer& dst_buffer) { vkapi::PipelineBarrier pipeline_barrier{}; uint32_t buffer_len = utils::safe_downcast(dst_buffer.numel() / 4); - utils::uvec3 global_wg_size = {buffer_len, 1, 1}; + const LocalWorkGroup lwg(64u, 1u, 1u); + const GlobalWorkGrid gwg = make_linear_dispatch(context, buffer_len, lwg); std::string kernel_name = "bitw8_image_to_nchw_nobitw8buffer_no_pc"; add_storage_type_suffix(kernel_name, v_src.storage_type()); @@ -234,8 +258,8 @@ void record_bitw8_image_to_nchw_nobitw8buffer_op( context->submit_compute_job( VK_KERNEL_FROM_STR(kernel_name), pipeline_barrier, - global_wg_size, - adaptive_work_group_size(global_wg_size), + gwg, + lwg, {v_src.hashed_layout()}, VK_NULL_HANDLE, 0, @@ -256,11 +280,14 @@ void record_binary_op( vkapi::PipelineBarrier pipeline_barrier{}; vkapi::SpecVarList specialization_constants = {}; + const GlobalWorkGrid gwg( + utils::make_uvec3(v_dst.logical_limits()), kTextureExtentsWorkGrid); + const LocalWorkGroup lwg = adaptive_lwg(gwg); context->submit_compute_job( VK_KERNEL_FROM_STR(kernel_name), pipeline_barrier, - v_dst.logical_limits(), - adaptive_work_group_size(v_dst.logical_limits()), + gwg, + lwg, specialization_constants, VK_NULL_HANDLE, 0, @@ -320,11 +347,12 @@ void record_index_fill_buffer(api::Context* context, api::vTensor& v_ten) { { vkapi::PipelineBarrier pipeline_barrier{}; vkapi::SpecVarList specialization_constants = {}; + const LocalWorkGroup lwg(64u, 1u, 1u); api::context()->submit_compute_job( VK_KERNEL_FROM_STR(kernel_name), pipeline_barrier, - {uint32_t(v_ten.numel()), 1, 1}, - {64, 1, 1}, + make_linear_dispatch(context, v_ten.numel(), lwg), + lwg, specialization_constants, VK_NULL_HANDLE, 0, @@ -344,11 +372,12 @@ void record_scalar_add_buffer( vkapi::SpecVarList specialization_constants = {SV(offset)}; std::string kernel = "scalar_add_buffer"; add_dtype_suffix(kernel, v_ten.dtype()); + const LocalWorkGroup lwg(64u, 1u, 1u); api::context()->submit_compute_job( VK_KERNEL_FROM_STR(kernel), pipeline_barrier, - {uint32_t(v_ten.numel()), 1, 1}, - {64, 1, 1}, + make_linear_dispatch(context, v_ten.numel(), lwg), + lwg, specialization_constants, VK_NULL_HANDLE, 0, @@ -368,8 +397,9 @@ void record_reference_matmul( api::context()->submit_compute_job( VK_KERNEL(reference_matmul), pipeline_barrier, - {uint32_t(out.size(1)), uint32_t(out.size(0)), 1}, - {64, 1, 1}, + GlobalWorkGrid( + {uint32_t(out.size(1)), uint32_t(out.size(0)), 1u}, kTiledWorkGrid), + LocalWorkGroup(64u, 1u, 1u), {}, VK_NULL_HANDLE, 0, @@ -399,7 +429,8 @@ void record_matmul_texture3d( add_storage_type_suffix(kernel_name, out.storage_type()); add_dtype_suffix(kernel_name, out.dtype()); - utils::uvec3 global_wg_size = out.logical_limits(); + const GlobalWorkGrid gwg( + utils::make_uvec3(out.logical_limits()), kTextureExtentsWorkGrid); struct PushConstants { utils::ivec4 out_sizes; @@ -435,11 +466,11 @@ void record_matmul_texture3d( vkapi::SpecVarList specialization_constants = { out.hashed_layout(), mat1.hashed_layout(), mat2.hashed_layout()}; - utils::uvec3 local_wg_size = {8, 8, 1}; + const LocalWorkGroup lwg(8u, 8u, 1u); vkapi::DescriptorSet descriptor_set = api::context()->get_descriptor_set( VK_KERNEL_FROM_STR(kernel_name), - LocalWorkGroup(local_wg_size), + lwg, specialization_constants, sizeof(push_constants)); @@ -458,7 +489,8 @@ void record_matmul_texture3d( descriptor_set, pipeline_barrier, VK_KERNEL_FROM_STR(kernel_name), - global_wg_size, + gwg, + lwg, &push_constants, sizeof(push_constants)); } diff --git a/backends/vulkan/test/vulkan_compute_api_test.cpp b/backends/vulkan/test/vulkan_compute_api_test.cpp index 2d8694bf33f..95776e42304 100644 --- a/backends/vulkan/test/vulkan_compute_api_test.cpp +++ b/backends/vulkan/test/vulkan_compute_api_test.cpp @@ -19,6 +19,8 @@ #include +#include + #include #include @@ -1173,8 +1175,8 @@ TEST_F(VulkanComputeAPITest, spec_var_shader_test) { context()->submit_compute_job( VK_KERNEL(fill_buffer), pipeline_barrier, - {64, 1, 1}, - {len_div4, 1, 1}, + GlobalWorkGrid({64u, 1u, 1u}, kExplicitWorkGrid), + LocalWorkGroup(len_div4, 1u, 1u), {SV(scale), SV(offset)}, VK_NULL_HANDLE, 0, @@ -1220,8 +1222,8 @@ TEST_F(VulkanComputeAPITest, update_params_between_submit) { context()->submit_compute_job( VK_KERNEL_FROM_STR(kernel_name), pipeline_barrier, - {4, 4, 4}, - {4, 4, 4}, + GlobalWorkGrid({4u, 4u, 4u}, kExplicitWorkGrid), + LocalWorkGroup(4u, 4u, 4u), specialization_constants, VK_NULL_HANDLE, 0, @@ -1287,8 +1289,8 @@ void test_storage_buffer_type(const size_t len) { context()->submit_compute_job( VK_KERNEL_FROM_STR(kernel_name), pipeline_barrier, - {64, 1, 1}, - {len_div4, 1, 1}, + GlobalWorkGrid({64u, 1u, 1u}, kExplicitWorkGrid), + LocalWorkGroup(len_div4, 1u, 1u), specialization_constants, VK_NULL_HANDLE, 0, @@ -1690,8 +1692,8 @@ TEST_F(VulkanComputeAPITest, print_object_sizes) { EXPECT_TRUE(sizeof(StagingBuffer) < 500); // Current known size on 64 bit system: 608 B EXPECT_TRUE(sizeof(ComputeGraph) < 700); - // Current known size on 64 bit system: 248 B - EXPECT_TRUE(sizeof(DispatchNode) < 500); + // Current known size on 64 bit system: 528 B + EXPECT_TRUE(sizeof(DispatchNode) < 600); } TEST_F(VulkanComputeAPITest, test_tensor_creation_from_vulkan_image) { @@ -2105,8 +2107,8 @@ TEST(VulkanComputeGraphTest, test_simple_graph_with_symint) { graph.execute_nodes().emplace_back(new DispatchNode( graph, VK_KERNEL_FROM_STR("scalar_add_texture"), - graph.create_global_wg_size(a.value), - graph.create_local_wg_size(a.value), + graph.create_gwg(a.value), + graph.create_lwg(a.value), // Inputs and Outputs {{out.value, vkapi::MemoryAccessType::WRITE}}, // Shader params buffers @@ -2689,8 +2691,9 @@ void run_from_gpu_test( context()->submit_compute_job( VK_KERNEL_FROM_STR(kernel_name), pipeline_barrier, - vten.logical_limits(), - {4, 4, 4}, + GlobalWorkGrid( + utils::make_uvec3(vten.logical_limits()), kTextureExtentsWorkGrid), + LocalWorkGroup(4u, 4u, 4u), {vten.packed_dim(), offset}, VK_NULL_HANDLE, 0, @@ -3322,26 +3325,28 @@ vkapi::ShaderInfo pick_dynamic_dispatch_shader( return VK_KERNEL_FROM_STR(kernel_name); } -utils::uvec3 pick_dynamic_dispatch_global_wg_size( +GlobalWorkGrid pick_dynamic_dispatch_gwg( ComputeGraph* graph, const vkapi::ShaderInfo& shader, const std::vector& args, const std::vector& resize_args) { (void)shader; const ValueRef out = args[0].refs[0]; - return graph->logical_limits_of(out); + return GlobalWorkGrid( + utils::make_uvec3(graph->logical_limits_of(out)), + kTextureExtentsWorkGrid); } -utils::uvec3 pick_dynamic_dispatch_local_wg_size( +LocalWorkGroup pick_dynamic_dispatch_lwg( ComputeGraph* graph, const vkapi::ShaderInfo& shader, - const utils::uvec3& global_workgroup_size, + const GlobalWorkGrid& gwg, const std::vector& args, const std::vector& resize_args) { (void)graph; (void)shader; - (void)global_workgroup_size; - return {64, 1, 1}; + (void)gwg; + return LocalWorkGroup(64u, 1u, 1u); } void resize_dynamic_dispatch_node( @@ -3365,8 +3370,8 @@ void add_dynamic_dispatch_test_node( graph.execute_nodes().emplace_back(new DynamicDispatchNode( graph, pick_dynamic_dispatch_shader, - pick_dynamic_dispatch_global_wg_size, - pick_dynamic_dispatch_local_wg_size, + pick_dynamic_dispatch_gwg, + pick_dynamic_dispatch_lwg, // Inputs and Outputs {{out, vkapi::kWrite}, {{mat1, mat2}, vkapi::kRead}}, // Shader params buffers @@ -3529,7 +3534,7 @@ TEST(VulkanComputeGraphTest, test_int8x4_staging_round_trip) { } } -TEST(VulkanWorkGroupSizeTest, local_workgroup_size) { +TEST(VulkanWorkGroupSizeTest, lwg) { const LocalWorkGroup lwg(64u, 2u, 1u); EXPECT_EQ(lwg.x(), 64u); @@ -3552,7 +3557,7 @@ TEST(VulkanWorkGroupSizeTest, local_workgroup_size) { EXPECT_THROW(LocalWorkGroup(3u, 2u, 1u), vkapi::Error); } -TEST(VulkanWorkGroupSizeTest, local_workgroup_size_validation) { +TEST(VulkanWorkGroupSizeTest, lwg_validation) { const utils::uvec3 max_lwg{1024u, 1024u, 64u}; const LocalWorkGroup lwg(8u, 8u, 1u); @@ -3666,6 +3671,26 @@ TEST(VulkanWorkGroupSizeTest, linear_gwg_wraps_across_xy) { EXPECT_TRUE(gwg.is_linear()); } +TEST(VulkanWorkGroupSizeTest, linear_gwg_allows_cooperative_z_lanes) { + const LocalWorkGroup lwg(1u, 1u, 64u); + const utils::uvec3 max_wg_count{65536u, 65536u, 65536u}; + GlobalWorkGrid gwg({65537u, 1u, 1u}, kLinearWorkGrid); + gwg.wrap_linear_dispatch(max_wg_count, lwg); + + EXPECT_EQ(gwg.extents(), utils::uvec3({257u, 256u, 1u})); + EXPECT_EQ(gwg.required_lwg_size(), lwg); +} + +TEST(VulkanWorkGroupSizeTest, linear_gwg_allows_xy_cooperation) { + const LocalWorkGroup lwg(1u, 8u, 8u); + const utils::uvec3 max_wg_count{65536u, 65536u, 65536u}; + GlobalWorkGrid gwg({65537u * 8u, 1u, 1u}, kLinearWorkGrid); + gwg.wrap_linear_dispatch(max_wg_count, lwg); + + EXPECT_EQ(gwg.extents(), utils::uvec3({257u, 2048u, 1u})); + EXPECT_EQ(gwg.required_lwg_size(), lwg); +} + TEST(VulkanWorkGroupSizeTest, linear_gwg_rejects_insufficient_y) { const utils::uvec3 max_wg_count{2u, 1u, 1u}; GlobalWorkGrid gwg({3u * 64u, 1u, 1u}, kLinearWorkGrid); @@ -3673,6 +3698,15 @@ TEST(VulkanWorkGroupSizeTest, linear_gwg_rejects_insufficient_y) { EXPECT_THROW(gwg.wrap_linear_dispatch(max_wg_count, 64u), vkapi::Error); } +TEST(VulkanWorkGroupSizeTest, linear_gwg_uses_available_x_capacity) { + const LocalWorkGroup lwg(64u, 1u, 1u); + const utils::uvec3 max_wg_count{4u, 2u, 1u}; + GlobalWorkGrid gwg({7u * 64u, 1u, 1u}, kLinearWorkGrid); + gwg.wrap_linear_dispatch(max_wg_count, lwg); + + EXPECT_EQ(gwg.extents(), utils::uvec3({4u * 64u, 2u, 1u})); +} + TEST(VulkanWorkGroupSizeTest, explicit_gwg_preserves_extents) { const GlobalWorkGrid gwg({31u, 17u, 5u}, kExplicitWorkGrid); @@ -3695,3 +3729,75 @@ TEST(VulkanWorkGroupSizeTest, gwg_intents) { EXPECT_EQ(texture_grid.extents(), extents); EXPECT_FALSE(texture_grid.required_lwg_size().is_valid()); } + +TEST(VulkanWorkGroupSizeTest, required_lwg_picker) { + const LocalWorkGroup required_lwg(1u, 64u, 1u); + const GlobalWorkGrid gwg({32u, 16u, 1u}, kTiledWorkGrid, required_lwg); + + EXPECT_EQ(gwg.required_lwg_size(), required_lwg); + EXPECT_EQ( + pick_required_lwg(nullptr, vkapi::ShaderInfo{}, gwg, {}, {}), + required_lwg); +} + +TEST(VulkanWorkGroupSizeTest, required_lwg_picker_rejects_missing_lwg) { + const GlobalWorkGrid gwg({32u, 16u, 1u}, kTiledWorkGrid); + + EXPECT_THROW( + pick_required_lwg(nullptr, vkapi::ShaderInfo{}, gwg, {}, {}), + vkapi::Error); +} + +TEST(VulkanWorkGroupSizeTest, explicit_gwg_rejects_excess_workgroups) { + const LocalWorkGroup lwg(2u, 1u, 1u); + const GlobalWorkGrid gwg({5u, 1u, 1u}, kExplicitWorkGrid); + + EXPECT_THROW(gwg.validate(lwg, {2u, 1u, 1u}), vkapi::Error); +} + +TEST(VulkanWorkGroupSizeTest, linear_gwg_requires_its_local_hint) { + const LocalWorkGroup lwg(64u, 1u, 1u); + GlobalWorkGrid gwg({100u, 1u, 1u}, kLinearWorkGrid); + gwg.wrap_linear_dispatch({65536u, 65536u, 65536u}, lwg); + + EXPECT_THROW( + gwg.validate(LocalWorkGroup(32u, 1u, 1u), {65536u, 65536u, 65536u}), + vkapi::Error); +} + +TEST(VulkanWorkGroupSizeTest, adapter_dispatch_recommendations) { + if (!api::available()) { + GTEST_SKIP(); + } + api::Context* const context = api::context(); + const auto* const adapter = context->adapter_ptr(); + const utils::uvec3 max_wg_count = adapter->max_compute_workgroup_count(); + const utils::uvec3 max_lwg = adapter->max_compute_workgroup_size(); + + EXPECT_GT(max_wg_count[0], 0u); + EXPECT_GT(max_wg_count[1], 0u); + EXPECT_GT(max_wg_count[2], 0u); + EXPECT_GT(max_lwg[0], 0u); + EXPECT_GT(adapter->max_compute_workgroup_invocations(), 0u); + EXPECT_EQ(adapter->recommended_lwg_nthreads(), 64u); +} + +TEST(VulkanWorkGroupSizeTest, compute_graph_preserves_dispatch_intent) { + if (!api::available()) { + GTEST_SKIP(); + } + GraphConfig config; + ComputeGraph graph(config); + + const ValueRef buffer = graph.add_tensor( + {257}, vkapi::kFloat, utils::kBuffer, utils::kWidthPacked); + const ValueRef texture = graph.add_tensor( + {1, 2, 3, 4}, vkapi::kFloat, utils::kTexture3D, utils::kChannelsPacked); + + const GlobalWorkGrid buffer_gwg = graph.create_gwg(buffer); + const GlobalWorkGrid texture_gwg = graph.create_gwg(texture); + + EXPECT_TRUE(buffer_gwg.is_linear()); + EXPECT_EQ(graph.create_lwg(buffer_gwg), buffer_gwg.required_lwg_size()); + EXPECT_FALSE(texture_gwg.is_linear()); +}