diff --git a/msl/internal/codegen/functions.go b/msl/internal/codegen/functions.go index 5d76fb8..f6555de 100644 --- a/msl/internal/codegen/functions.go +++ b/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 } diff --git a/snapshot/snapshot_test.go b/snapshot/snapshot_test.go index caca4ef..b81596f 100644 --- a/snapshot/snapshot_test.go +++ b/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. diff --git a/snapshot/testdata/golden/msl/abstract-types-operators.msl b/snapshot/testdata/golden/msl/abstract-types-operators.msl index 6c032d1..7eb0e39 100644 --- a/snapshot/testdata/golden/msl/abstract-types-operators.msl +++ b/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 = {}; } diff --git a/snapshot/testdata/golden/msl/atomicOps-int64.msl b/snapshot/testdata/golden/msl/atomicOps-int64.msl index 1630e26..d2c5ded 100644 --- a/snapshot/testdata/golden/msl/atomicOps-int64.msl +++ b/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++) { diff --git a/snapshot/testdata/golden/msl/atomicOps.msl b/snapshot/testdata/golden/msl/atomicOps.msl index 8aeab45..c319dac 100644 --- a/snapshot/testdata/golden/msl/atomicOps.msl +++ b/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++) { diff --git a/snapshot/testdata/golden/msl/atomics.msl b/snapshot/testdata/golden/msl/atomics.msl index b25ca28..716296b 100644 --- a/snapshot/testdata/golden/msl/atomics.msl +++ b/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); } diff --git a/snapshot/testdata/golden/msl/globals.msl b/snapshot/testdata/golden/msl/globals.msl index 2a8ff6e..bdc0263 100644 --- a/snapshot/testdata/golden/msl/globals.msl +++ b/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); diff --git a/snapshot/testdata/golden/msl/interface.msl b/snapshot/testdata/golden/msl/interface.msl index a886cb0..eb08b98 100644 --- a/snapshot/testdata/golden/msl/interface.msl +++ b/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 = {}; } diff --git a/snapshot/testdata/golden/msl/overrides-atomicCompareExchangeWeak.msl b/snapshot/testdata/golden/msl/overrides-atomicCompareExchangeWeak.msl index e02e134..3361f8c 100644 --- a/snapshot/testdata/golden/msl/overrides-atomicCompareExchangeWeak.msl +++ b/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); } diff --git a/snapshot/testdata/golden/msl/policy-mix.msl b/snapshot/testdata/golden/msl/policy-mix.msl index 4bcb735..c6805a9 100644 --- a/snapshot/testdata/golden/msl/policy-mix.msl +++ b/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 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 = {}; } diff --git a/snapshot/testdata/golden/msl/template-list-trailing-comma.msl b/snapshot/testdata/golden/msl/template-list-trailing-comma.msl index e40f9cd..53038fb 100644 --- a/snapshot/testdata/golden/msl/template-list-trailing-comma.msl +++ b/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 = {}; diff --git a/snapshot/testdata/golden/msl/types_with_comments.msl b/snapshot/testdata/golden/msl/types_with_comments.msl index b3d7b78..f5cf751 100644 --- a/snapshot/testdata/golden/msl/types_with_comments.msl +++ b/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); diff --git a/snapshot/testdata/golden/msl/workgroup-uniform-load.msl b/snapshot/testdata/golden/msl/workgroup-uniform-load.msl index ddc80db..64b1dc7 100644 --- a/snapshot/testdata/golden/msl/workgroup-uniform-load.msl +++ b/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_ = {}; } diff --git a/snapshot/testdata/golden/msl/workgroup-var-init.msl b/snapshot/testdata/golden/msl/workgroup-var-init.msl index 2b5dab9..62a62db 100644 --- a/snapshot/testdata/golden/msl/workgroup-var-init.msl +++ b/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); diff --git a/snapshot/testdata/golden/msl/workgroup_memory.msl b/snapshot/testdata/golden/msl/workgroup_memory.msl index 1669b19..1068835 100644 --- a/snapshot/testdata/golden/msl/workgroup_memory.msl +++ b/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 = {}; }