fix(msl): declare workgroup vars at function-body scope in kernel entry points

msl/internal/codegen/functions.go

@@ -573,12 +573,26 @@ func (w *Writer) writeEntryPoint(epIdx int, ep *ir.EntryPoint) error {
 	// Collect all parameters, then format them
 	paramCount := 0
 
+	// Build set of globals actually referenced by this entry point (direct + transitive).
+	// Rust naga only emits resources that the entry point actually uses, not ALL globals.
+	epUsedGlobals := make(map[uint32]struct{})
+	if globals, ok := w.funcPassThroughGlobals[epFuncHandle(epIdx)]; ok {
+		for _, h := range globals {
+			epUsedGlobals[h] = struct{}{}
+		}
+	}
+
 	// Check if we need workgroup zero-initialization for this entry point.
-	// This requires: compute shader + ZeroInitializeWorkgroupMemory + workgroup vars.
+	// This requires: compute shader + ZeroInitializeWorkgroupMemory + workgroup vars
+	// actually used by this entry point (matching Rust naga, which filters by
+	// !fun_info[handle].is_empty()).
 	needWorkgroupInit := false
 	localInvocationIDName := ""
 	if w.options.ZeroInitializeWorkgroupMemory && ep.Stage == ir.StageCompute {
-		for _, global := range w.module.GlobalVariables {
+		for i, global := range w.module.GlobalVariables {
+			if _, used := epUsedGlobals[uint32(i)]; !used {
+				continue
+			}
 			if global.Space == ir.SpaceWorkGroup {
 				needWorkgroupInit = true
 				break
@@ -684,15 +698,6 @@ func (w *Writer) writeEntryPoint(epIdx int, ep *ir.EntryPoint) error {
 	// preventing collisions when multiple groups share binding numbers.
 	w.computeResourceMap(ep.Name)
 
-	// Build set of globals actually referenced by this entry point (direct + transitive).
-	// Rust naga only emits resources that the entry point actually uses, not ALL globals.
-	epUsedGlobals := make(map[uint32]struct{})
-	if globals, ok := w.funcPassThroughGlobals[epFuncHandle(epIdx)]; ok {
-		for _, h := range globals {
-			epUsedGlobals[h] = struct{}{}
-		}
-	}
-
 	// Global variable parameters — emitted in declaration order (matching Rust naga).
 	// This includes both resource bindings (device/constant with [[buffer]]/[[texture]])
 	// and workgroup variables (threadgroup without binding attributes).
@@ -717,11 +722,12 @@ func (w *Writer) writeEntryPoint(epIdx int, ep *ir.EntryPoint) error {
 			w.writeEntryPointParam(paramCount, paramStr)
 			paramCount++
 		} else if global.Space == ir.SpaceWorkGroup {
-			// Workgroup variable — threadgroup parameter without binding attribute
-			name := w.getName(nameKey{kind: nameKeyGlobalVariable, handle1: uint32(i)})
-			typeName := w.writeTypeName(global.Type, StorageAccess(0))
-			w.writeEntryPointParam(paramCount, fmt.Sprintf("threadgroup %s& %s", typeName, name))
-			paramCount++
+			// Workgroup variables are NOT emitted as entry-point parameters.
+			// Metal requires the host to call setThreadgroupMemoryLength:atIndex:
+			// for threadgroup entry-point parameters; instead they are declared
+			// at function-body scope inside the kernel (see below), which needs
+			// no host-side setup.
+			continue
 		} else if global.Space == ir.SpaceImmediate {
 			// Immediate data variable — constant buffer parameter.
 			// Resolve binding slot from per-entry-point ImmediatesBuffer config.
@@ -858,11 +864,30 @@ func (w *Writer) writeEntryPoint(epIdx int, ep *ir.EntryPoint) error {
 	// For simple results, inline aggregate initialization is used.
 	_ = outputStructName
 
+	// Workgroup variable declarations — function-body scope, threadgroup
+	// address space. Declared inside the kernel (legal MSL) rather than as
+	// entry-point parameters, because threadgroup *parameters* require the
+	// host to call setThreadgroupMemoryLength:atIndex: which the pure-Go
+	// Metal HAL does not do. Names are identical to the former parameters,
+	// so all body references and helper-function call sites resolve unchanged.
+	// Must come BEFORE the zero-init prologue, which references these names.
+	for i, global := range w.module.GlobalVariables {
+		if _, used := epUsedGlobals[uint32(i)]; !used {
+			continue
+		}
+		if global.Space != ir.SpaceWorkGroup {
+			continue
+		}
+		name := w.getName(nameKey{kind: nameKeyGlobalVariable, handle1: uint32(i)})
+		typeName := w.writeTypeName(global.Type, StorageAccess(0))
+		w.WriteLine("threadgroup %s %s;", typeName, name)
+	}
+
 	// Workgroup zero-initialization prologue.
 	// Matches Rust naga: zero all workgroup vars if __local_invocation_id == uint3(0).
 	// Must come BEFORE local variables and private var locals.
 	if needWorkgroupInit {
-		if err := w.writeWorkgroupZeroInit(localInvocationIDName); err != nil {
+		if err := w.writeWorkgroupZeroInit(localInvocationIDName, epUsedGlobals); err != nil {
 			return err
 		}
 	}
@@ -1843,12 +1868,16 @@ func resolveInterpolationString(interp *ir.Interpolation) string {
 }
 
 // writeWorkgroupZeroInit writes the zero-initialization prologue for workgroup variables.
-// Matches Rust naga: check __local_invocation_id == uint3(0), then zero-init all workgroup vars.
-func (w *Writer) writeWorkgroupZeroInit(localInvIDName string) error {
+// Matches Rust naga: check __local_invocation_id == uint3(0), then zero-init the workgroup
+// vars used by this entry point (Rust filters by !fun_info[handle].is_empty()).
+func (w *Writer) writeWorkgroupZeroInit(localInvIDName string, usedGlobals map[uint32]struct{}) error {
 	w.WriteLine("if (%sall(%s == %suint3(0u))) {", Namespace, localInvIDName, Namespace)
 	w.PushIndent()
 
 	for i, global := range w.module.GlobalVariables {
+		if _, used := usedGlobals[uint32(i)]; !used {
+			continue
+		}
 		if global.Space != ir.SpaceWorkGroup {
 			continue
 		}

snapshot/snapshot_test.go

@@ -102,6 +102,13 @@ func TestSnapshots(t *testing.T) {
 //   - "extra-opdecorate": we emit an extra OpDecorate (e.g. NonWritable) that Rust omits.
 //   - "no-compact-pass": shader has no entry points — Rust compact pass removes all dead code,
 //     we emit full output. Rust reference is empty/minimal, comparison is meaningless.
+//   - "msl-threadgroup-function-scope": we declare workgroup variables at function-body
+//     scope inside the kernel ("threadgroup T name;") instead of as threadgroup
+//     entry-point parameters like Rust naga. Rust's approach requires the host to call
+//     setThreadgroupMemoryLength:atIndex: (Rust wgpu-hal does this in its Metal encoder);
+//     the pure-Go wgpu Metal HAL does not, so parameter-based threadgroup memory is
+//     unsized and reads/writes silently no-op. Function-scope declarations are legal MSL
+//     and need no host-side setup. Generated code is otherwise identical.
 var referenceAllowList = map[string]string{
 	"atomicOps":                 "workgroup-layout-free",
 	"atomicOps-int64":           "workgroup-layout-free",
@@ -110,6 +117,13 @@ var referenceAllowList = map[string]string{
 	"bits":                      "missing-int8-capability",
 	"binding-buffer-arrays":     "extra-opdecorate",
 	"ptr-deref-test":            "no-compact-pass",
+
+	"abstract-types-operators":            "msl-threadgroup-function-scope",
+	"globals":                             "msl-threadgroup-function-scope",
+	"interface":                           "msl-threadgroup-function-scope",
+	"overrides-atomicCompareExchangeWeak": "msl-threadgroup-function-scope",
+	"policy-mix":                          "msl-threadgroup-function-scope",
+	"workgroup-uniform-load":              "msl-threadgroup-function-scope",
 }
 
 // TestRustReference compares our compiled output against Rust naga reference outputs.

snapshot/testdata/golden/msl/abstract-types-operators.msl

@@ -115,8 +115,8 @@ void wgpu_4435_(
 
 kernel void main_(
   metal::uint3 __local_invocation_id [[thread_position_in_threadgroup]]
-, threadgroup type_3& a
 ) {
+    threadgroup type_3 a;
     if (metal::all(__local_invocation_id == metal::uint3(0u))) {
         a = {};
     }

snapshot/testdata/golden/msl/atomicOps-int64.msl

@@ -85,10 +85,10 @@ kernel void cs_main(
 , device metal::atomic_ulong& storage_atomic_scalar [[user(fake0)]]
 , device type_2& storage_atomic_arr [[user(fake0)]]
 , device Struct& storage_struct [[user(fake0)]]
-, threadgroup metal::atomic_ulong& workgroup_atomic_scalar
-, threadgroup type_2& workgroup_atomic_arr
-, threadgroup Struct& workgroup_struct
 ) {
+    threadgroup metal::atomic_ulong workgroup_atomic_scalar;
+    threadgroup type_2 workgroup_atomic_arr;
+    threadgroup Struct workgroup_struct;
     if (metal::all(id == metal::uint3(0u))) {
         metal::atomic_store_explicit(&workgroup_atomic_scalar, 0, metal::memory_order_relaxed);
         for (int __i0 = 0; __i0 < 2; __i0++) {

snapshot/testdata/golden/msl/atomicOps.msl

@@ -85,10 +85,10 @@ kernel void cs_main(
 , device metal::atomic_uint& storage_atomic_scalar [[user(fake0)]]
 , device type_2& storage_atomic_arr [[user(fake0)]]
 , device Struct& storage_struct [[user(fake0)]]
-, threadgroup metal::atomic_uint& workgroup_atomic_scalar
-, threadgroup type_2& workgroup_atomic_arr
-, threadgroup Struct& workgroup_struct
 ) {
+    threadgroup metal::atomic_uint workgroup_atomic_scalar;
+    threadgroup type_2 workgroup_atomic_arr;
+    threadgroup Struct workgroup_struct;
     if (metal::all(id == metal::uint3(0u))) {
         metal::atomic_store_explicit(&workgroup_atomic_scalar, 0, metal::memory_order_relaxed);
         for (int __i0 = 0; __i0 < 2; __i0++) {

snapshot/testdata/golden/msl/atomics.msl

@@ -21,10 +21,10 @@ struct main_Input {
 kernel void main_(
   uint lid [[thread_index_in_threadgroup]]
 , metal::uint3 __local_invocation_id [[thread_position_in_threadgroup]]
-, threadgroup metal::atomic_uint& shared_counter
 , device type_2& result [[user(fake0)]]
 , constant _mslBufferSizes& _buffer_sizes [[user(fake0)]]
 ) {
+    threadgroup metal::atomic_uint shared_counter;
     if (metal::all(__local_invocation_id == metal::uint3(0u))) {
         metal::atomic_store_explicit(&shared_counter, 0, metal::memory_order_relaxed);
     }

snapshot/testdata/golden/msl/globals.msl

@@ -69,8 +69,6 @@ void test_msl_packed_vec3_(
 
 kernel void main_(
   metal::uint3 __local_invocation_id [[thread_position_in_threadgroup]]
-, threadgroup type_2& wg
-, threadgroup metal::atomic_uint& at_1
 , device FooStruct& alignment [[user(fake0)]]
 , device type_6 const& dummy [[user(fake0)]]
 , constant type_8& float_vecs [[user(fake0)]]
@@ -80,6 +78,8 @@ kernel void main_(
 , constant type_15& global_nested_arrays_of_matrices_4x2_ [[user(fake0)]]
 , constant _mslBufferSizes& _buffer_sizes [[user(fake0)]]
 ) {
+    threadgroup type_2 wg;
+    threadgroup metal::atomic_uint at_1;
     if (metal::all(__local_invocation_id == metal::uint3(0u))) {
         wg = {};
         metal::atomic_store_explicit(&at_1, 0, metal::memory_order_relaxed);

snapshot/testdata/golden/msl/interface.msl

@@ -80,8 +80,8 @@ kernel void compute(
 , uint local_index [[thread_index_in_threadgroup]]
 , metal::uint3 wg_id [[threadgroup_position_in_grid]]
 , metal::uint3 num_wgs [[threadgroups_per_grid]]
-, threadgroup type_4& output
 ) {
+    threadgroup type_4 output;
     if (metal::all(local_id == metal::uint3(0u))) {
         output = {};
     }

snapshot/testdata/golden/msl/overrides-atomicCompareExchangeWeak.msl

@@ -44,8 +44,8 @@ constant int o = int {};
 
 kernel void f(
   metal::uint3 __local_invocation_id [[thread_position_in_threadgroup]]
-, threadgroup metal::atomic_uint& a
 ) {
+    threadgroup metal::atomic_uint a;
     if (metal::all(__local_invocation_id == metal::uint3(0u))) {
         metal::atomic_store_explicit(&a, 0, metal::memory_order_relaxed);
     }

snapshot/testdata/golden/msl/policy-mix.msl

@@ -57,8 +57,8 @@ kernel void main_(
 , device InStorage const& in_storage [[user(fake0)]]
 , constant InUniform& in_uniform [[user(fake0)]]
 , metal::texture2d_array<float, metal::access::sample> image_2d_array [[user(fake0)]]
-, threadgroup type_5& in_workgroup
 ) {
+    threadgroup type_5 in_workgroup;
     if (metal::all(__local_invocation_id == metal::uint3(0u))) {
         in_workgroup = {};
     }

snapshot/testdata/golden/msl/template-list-trailing-comma.msl

@@ -22,12 +22,12 @@ typedef uint type_2[1];
 
 kernel void main_(
   metal::uint3 __local_invocation_id [[thread_position_in_threadgroup]]
-, threadgroup type_1& sized_comma
-, threadgroup type_1& sized_no_comma
 , device type_2 const& unsized_comma [[user(fake0)]]
 , device type_2& unsized_no_comma [[user(fake0)]]
 , constant _mslBufferSizes& _buffer_sizes [[user(fake0)]]
 ) {
+    threadgroup type_1 sized_comma;
+    threadgroup type_1 sized_no_comma;
     if (metal::all(__local_invocation_id == metal::uint3(0u))) {
         sized_comma = {};
         sized_no_comma = {};

snapshot/testdata/golden/msl/types_with_comments.msl

@@ -30,10 +30,9 @@ void test_g(
 
 kernel void test_ep(
   metal::uint3 __local_invocation_id [[thread_position_in_threadgroup]]
-, threadgroup metal::float2x2& w_mem2_
 ) {
+    threadgroup metal::float2x2 w_mem2_;
     if (metal::all(__local_invocation_id == metal::uint3(0u))) {
-        w_mem = {};
         w_mem2_ = {};
     }
     metal::threadgroup_barrier(metal::mem_flags::mem_threadgroup);

snapshot/testdata/golden/msl/workgroup-uniform-load.msl

@@ -20,8 +20,8 @@ struct test_workgroupUniformLoadInput {
 kernel void test_workgroupUniformLoad(
   metal::uint3 workgroup_id [[threadgroup_position_in_grid]]
 , metal::uint3 __local_invocation_id [[thread_position_in_threadgroup]]
-, threadgroup type_2& arr_i32_
 ) {
+    threadgroup type_2 arr_i32_;
     if (metal::all(__local_invocation_id == metal::uint3(0u))) {
         arr_i32_ = {};
     }

snapshot/testdata/golden/msl/workgroup-var-init.msl

@@ -27,9 +27,9 @@ struct WStruct {
 
 kernel void main_(
   metal::uint3 __local_invocation_id [[thread_position_in_threadgroup]]
-, threadgroup WStruct& w_mem
 , device type_1& output [[user(fake0)]]
 ) {
+    threadgroup WStruct w_mem;
     if (metal::all(__local_invocation_id == metal::uint3(0u))) {
         w_mem.arr = {};
         metal::atomic_store_explicit(&w_mem.atom, 0, metal::memory_order_relaxed);

snapshot/testdata/golden/msl/workgroup_memory.msl

@@ -32,10 +32,10 @@ kernel void main_(
   uint lid [[thread_index_in_threadgroup]]
 , metal::uint3 gid [[thread_position_in_grid]]
 , metal::uint3 __local_invocation_id [[thread_position_in_threadgroup]]
-, threadgroup type_1& shared_data
 , device type_3& output [[user(fake0)]]
 , constant _mslBufferSizes& _buffer_sizes [[user(fake0)]]
 ) {
+    threadgroup type_1 shared_data;
     if (metal::all(__local_invocation_id == metal::uint3(0u))) {
         shared_data = {};
     }