From 024a7e5949d2466aff09ae23bc26c98bee71177f Mon Sep 17 00:00:00 2001 From: Stephen Jia Date: Thu, 27 Aug 2026 09:14:50 -0700 Subject: [PATCH] Update [ghstack-poisoned] --- .../graph/ops/glsl/group_norm_reduce_texture.glsl | 9 ++++----- backends/vulkan/runtime/graph/ops/impl/GroupNorm.cpp | 4 +++- 2 files changed, 7 insertions(+), 6 deletions(-) diff --git a/backends/vulkan/runtime/graph/ops/glsl/group_norm_reduce_texture.glsl b/backends/vulkan/runtime/graph/ops/glsl/group_norm_reduce_texture.glsl index a3de804dd0b..51425d47d86 100644 --- a/backends/vulkan/runtime/graph/ops/glsl/group_norm_reduce_texture.glsl +++ b/backends/vulkan/runtime/graph/ops/glsl/group_norm_reduce_texture.glsl @@ -65,17 +65,16 @@ shared float shared_sum_sq[LOCAL_WORK_GROUP_SIZE]; * N is the number of elements in the tensor buffer; each thread computes one * output element. * - * Local work group size: {1, float, 1} - * float should be a power of 2, recommended 64 or 128 threads. This allows - * efficient tree-based reduction in shared memory. Each local group will - * cooperate to compute the output element. + * Local work group size: {1, 1, 64} + * This allows efficient tree-based reduction in shared memory. Each local + * group cooperates to compute one output element. * * Each shader invocation will compute the mean and standard deviation for one * channel group in the input, and write out the corresponding result. */ void group_norm_reduce_C_packed() { const int global_idx = int(gl_GlobalInvocationID.x); - const int local_idx = int(gl_LocalInvocationID.y); + const int local_idx = int(gl_LocalInvocationID.z); // Calculate group dimensions const int D = in_sizes.z / group; // channels per group diff --git a/backends/vulkan/runtime/graph/ops/impl/GroupNorm.cpp b/backends/vulkan/runtime/graph/ops/impl/GroupNorm.cpp index 368b95c9d3b..f18e59229de 100644 --- a/backends/vulkan/runtime/graph/ops/impl/GroupNorm.cpp +++ b/backends/vulkan/runtime/graph/ops/impl/GroupNorm.cpp @@ -29,13 +29,15 @@ utils::uvec3 group_norm_local_wg_size( (void)args; (void)resize_args; - return {1, 64, 1}; + return {1u, 1u, 64u}; } 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);