A task shader's payload past Metal's 16 KB goes through a device buffer: Switch Sports Resort's six refused task / mesh shaders build

dm/gfx/nvn/dmsh_abi.h

@@ -31,6 +31,8 @@
 #define DMSH_C0_SSBO_MS 0x770            // NVN2 mesh stage storage buffer table: its mesh shaders read c0 0x770, 0x780, 0x790 and 0x7b0, and it binds slots from 0
 #define DMSH_MESH_DRAW_SLOT 29           // a task or object function's buffer: the draw's {count, first} workgroups, NVN's DrawMeshTasksIndirect record (dmsh_meshstage.py, nvn_geometry.inc mtk_encode)
 #define DMSH_MESH_GRID 1024              // Metal's limit of threadgroups in a mesh grid: a mesh shader with no task shader is drawn by an object function launching blocks of this many (dmsh_meshstage.py)
+#define DMSH_MESH_PAYLOAD_MAX 16384      // Metal's limit of an object function's payload: a task shader's larger payload goes through a device buffer (dmsh_meshstage.py spill)
+#define DMSH_MESH_SPILL_SLOT 11          // that buffer, at this index of the object and vertex (mesh) tables: word 0 its slots, then one slice of the payload's words rounded to 4 a task workgroup (nvn_geometry.inc mtk_encode)
 #define DMSH_C0_DISPATCH 0x200           // compute: the dispatch's workgroup counts x, y, z (32 bits each; Link's Awakening's light culling)
 #define DMSH_C0_DISPATCH_NVN2 4          // NVN2 compute keeps them here, below its per-shader layout at 0x10 (x and y read by Fire Emblem Fortune's Weave's light culling; z by their order)
 #define DMSH_C0_CS_UBO 0x210             // compute: uniform slot r's address and size at 0x210 + 0x10 r, read through global memory (the build's ubo_ptr bindings); it ends where the storage buffer table starts

dm/gfx/nvn/dmsh_chain.h

@@ -1,3 +1,3 @@
 // Generated file (dm/gfx/tools/dmsh_chain.py --header): the B shader library translation chain version this source tree expects (default options), one for every game;
 // the host loads only <library dir>/b-<V>/. Run dmsh_chain.py --header after changing the chain's tools and commit this file with the change
-#define DMSH_CHAIN "bc8830d4ba5a"
+#define DMSH_CHAIN "eb9a3f4c3394"

dm/gfx/nvn/nvn_geometry.inc

@@ -556,6 +556,11 @@ static void geo_mesh_encode(id<MTLRenderCommandEncoder> en_) {
 // equivalent fallback: the demanded pipeline must finish compiling before its first draw, so each one built at a draw is recorded as an M row (pso_user.tsv) and the prewarm builds
 // the rows before the game reaches them (docs/reference/shaders.md, Native mesh pipelines compiled at their first draw)
 static struct { int on, miss; NvShader *ts, *ms; uint32_t first, count; uint64_t ind, ind_rec; } g_mtk;   // ind: an indirect draw's threadgroups (mdx_mask), ind_rec: its {count, first} record
+// A task shader's payload past Metal's object payload (DMSH_MESH_PAYLOAD_MAX) goes through a buffer at DMSH_MESH_SPILL_SLOT of both stages (dmsh_meshstage.py spill): the Metal payload is
+// the slot's number only. The buffer has a slice for each task workgroup of the draw, up to MTK_SPILL_MAX; an indirect draw's count is written on the GPU, so it gets MTK_SPILL_MAX
+enum { MTK_SPILL_MAX = 1024 };
+static int mtk_spill(const NvShader *ts) { return ts && ts->mesh_payload > DMSH_MESH_PAYLOAD_MAX; }
+static NSUInteger mtk_payload_len(const NvShader *ts) { return !ts ? 16 : mtk_spill(ts) ? 16 : (ts->mesh_payload + 15u) & ~15u; }
 static void draw_(const p_draw *dw, const p_draw_idx *di);
 static uint64_t g_mtk_n, g_mtk_wait;
 static nvn_hmap g_pso_mtk = {.klen = sizeof(pso_key)};   // only the submit thread uses it: prewarm workers build pipelines (mtk_build) and the submit thread claims them (warm_claim)
@@ -574,7 +579,7 @@ static MTL4MeshRenderPipelineDescriptor *mtk_desc4(MTLRenderPipelineDescriptor *
     if (!mf || !of || !f4) return nil;
     MTL4MeshRenderPipelineDescriptor *m = [MTL4MeshRenderPipelineDescriptor new];
     m.meshFunctionDescriptor = mf; m.fragmentFunctionDescriptor = f4; m.maxTotalThreadsPerMeshThreadgroup = ms->mesh_threads;
-    m.objectFunctionDescriptor = of; m.payloadMemoryLength = ts ? (ts->mesh_payload + 15u) & ~15u : 16; m.maxTotalThreadsPerObjectThreadgroup = ts ? ts->mesh_threads : 1;   // with no task: a block's first workgroup, one thread a block
+    m.objectFunctionDescriptor = of; m.payloadMemoryLength = mtk_payload_len(ts); m.maxTotalThreadsPerObjectThreadgroup = ts ? ts->mesh_threads : 1;   // with no task: a block's first workgroup, one thread a block
     m.rasterSampleCount = d.rasterSampleCount; m.alphaToCoverageState = d.alphaToCoverageEnabled ? MTL4AlphaToCoverageStateEnabled : MTL4AlphaToCoverageStateDisabled;
     for (int i = 0; i < 8; i++) pso_col4(m.colorAttachments[i], d.colorAttachments[i]);
     return m;
@@ -604,7 +609,7 @@ static id<MTLRenderPipelineState> mtk_build(const pso_key *k, NvShader *ts, NvSh
     }
     MTLMeshRenderPipelineDescriptor *m = [MTLMeshRenderPipelineDescriptor new];
     m.meshFunction = fn_dev(ms->fn); m.fragmentFunction = fn_dev(ff); m.maxTotalThreadsPerMeshThreadgroup = ms->mesh_threads;
-    m.objectFunction = fn_dev(ts ? ts->fn : ms->mesh_obj); m.payloadMemoryLength = ts ? (ts->mesh_payload + 15u) & ~15u : 16; m.maxTotalThreadsPerObjectThreadgroup = ts ? ts->mesh_threads : 1;
+    m.objectFunction = fn_dev(ts ? ts->fn : ms->mesh_obj); m.payloadMemoryLength = mtk_payload_len(ts); m.maxTotalThreadsPerObjectThreadgroup = ts ? ts->mesh_threads : 1;
     m.rasterSampleCount = d.rasterSampleCount; m.alphaToCoverageEnabled = d.alphaToCoverageEnabled;
     for (int i = 0; i < 8; i++) m.colorAttachments[i] = d.colorAttachments[i];
     m.depthAttachmentPixelFormat = d.depthAttachmentPixelFormat; m.stencilAttachmentPixelFormat = d.stencilAttachmentPixelFormat;
@@ -659,6 +664,15 @@ static void mtk_encode(id<MTLRenderCommandEncoder> en_) {
     if (rec_on()) rec_table_op(RO_T_ADDR, g_ot, draw_gpu, DMSH_MESH_DRAW_SLOT); else [g_ot setAddress:draw_gpu atIndex:DMSH_MESH_DRAW_SLOT];
     g_args_dirty[4] = 1; encode_args(en, ms, 4, 1);
     // A task shader's threadgroups are the draw's workgroups; with no task, the object function's are blocks of DMSH_MESH_GRID of them (mdx_mask writes their count at +16)
+    if (mtk_spill(ts)) {   // the task's payload slices: word 0 their count, the slices from word 4
+        uint32_t words = ((ts->mesh_payload + 15u) & ~15u) >> 2, cap = g_mtk.ind ? MTK_SPILL_MAX : MIN(g_mtk.count, (uint32_t)MTK_SPILL_MAX);
+        if (!g_mtk.ind && g_mtk.count > MTK_SPILL_MAX) { static int n; if (n++ < 4) MLOG("[RED] task shader %s: %u workgroups, payload slices for %u: the rest launch no mesh workgroups", ts->key, g_mtk.count, cap); }
+        id<MTLBuffer> sb; size_t so = arena_alloc(16 + (size_t)MAX(cap, 1u) * words * 4, &sb); *(uint32_t *)((uint8_t *)sb.contents + so) = cap;
+        uint64_t sg = sb.gpuAddress + so;
+        if (rec_on()) { rec_table_op(RO_T_ADDR, g_ot, sg, DMSH_MESH_SPILL_SLOT); rec_table_op(RO_T_ADDR, g_vt, sg, DMSH_MESH_SPILL_SLOT); }
+        else { [g_ot setAddress:sg atIndex:DMSH_MESH_SPILL_SLOT]; [g_vt setAddress:sg atIndex:DMSH_MESH_SPILL_SLOT]; }
+        static int said; if (!said++) MLOG("task shader %s: its %u-byte payload goes through a buffer (%u workgroups, %u-byte slices)", ts->key, ts->mesh_payload, g_mtk.ind ? 0 : g_mtk.count, words * 4);
+    }
     MTLSize to = MTLSizeMake(ts ? ts->mesh_threads : 1, 1, 1), tm = MTLSizeMake(ms->mesh_threads, 1, 1);
     if (g_mtk.ind) [en drawMeshThreadgroupsWithIndirectBuffer:g_mtk.ind + (ts ? 0 : 16) threadsPerObjectThreadgroup:to threadsPerMeshThreadgroup:tm];
     else [en drawMeshThreadgroups:MTLSizeMake(ts ? g_mtk.count : (g_mtk.count + DMSH_MESH_GRID - 1) / DMSH_MESH_GRID, 1, 1) threadsPerObjectThreadgroup:to threadsPerMeshThreadgroup:tm];

dm/gfx/tests/test_dmsh_meshstage.py

@@ -33,6 +33,16 @@
     _23[%d] = (*spvDescriptorSet0.cp_c3)._m0[0];
 }
 """
+# A spill mesh shader's kernel: the translator's dm_stage_in follows the stage memory, and the body reads it (ISBERD 0x38)
+KERNEL_IN = HEAD + """
+kernel void main0(constant spvDescriptorSetBuffer0& spvDescriptorSet0 [[buffer(0)]], uint3 gl_LocalInvocationID [[thread_position_in_threadgroup]])
+{
+    threadgroup spvUnsafeArray<uint, %d> _23;
+    threadgroup spvUnsafeArray<uint, %d> _24;
+    if (gl_LocalInvocationID.x != 0u) return;
+    _23[%d] = _24[int(gl_LocalInvocationID.x) + 1] + _24[0];
+}
+"""
 
 
 def sph(pos=0xf0, generics=b''):
@@ -61,9 +71,18 @@ def test_per_primitive_sph(self):   # after the instructions, at control +0x64 w
         self.assertEqual((i['kind'], i['threads'], i['payload'], i['task'], i['max_v'], i['max_p'], i['primitive']), ('ms', 8, 1408, True, 32, 16, [(0, 0)]))
         self.assertEqual(ms.info(control(7, max_v=128, max_p=64), area)['primitive'], [])
 
-    def test_task_memory(self):   # the payload's bytes hold the mesh count at +4, as the game's stores do
+    def test_task_memory(self):   # the payload's bytes hold the mesh count at +4, as the game's stores do. A task reads no input: its outputs follow word 0's block
         lay, size = ms.memory(ms.info(control(6), bytes(0x80)))
-        self.assertEqual(lay, [0, 0x10, 0x110, 0x210]); self.assertEqual(size, 0x210)
+        self.assertEqual(lay, [0, 0x10, 0x10, 0x110]); self.assertEqual(size, 0x110)
+
+    def test_spill_memory(self):   # Nintendo Switch Sports Resort's 07abb052a670b8ba / d60092d7424cb5b4: 27,904 payload bytes, past Metal's 16 KB
+        t = ms.info(control(6, payload=27904), bytes(0x80))
+        self.assertTrue(ms.spill(t)); self.assertFalse(ms.spill(ms.info(control(6, payload=abi.MESH_PAYLOAD_MAX), bytes(0x80))))
+        self.assertEqual(ms.memory(t), ([0, 0x10, 0x10, 0x10 + 27904], 0x10 + 27904))
+        m = ms.info(control(7, payload=27904, task=1, max_v=64, max_p=126), sph(0xf0, b'\x30'))
+        lay, size = ms.memory(m)
+        self.assertEqual(lay[2], 0x10); self.assertEqual(lay[4], 27904 // 4); self.assertLessEqual(size, ms.METAL_MAX_SHARED)   # the input is not in the stage memory
+        self.assertEqual(ms.spill_words(m), 6976)
 
     def test_mesh_memory(self):   # 128 vertices: four blocks of 32. Six components a block, 0x80 bytes each; the primitive bytes follow, 3 a triangle after the count's word
         i = ms.info(control(7, payload=128, max_v=128, max_p=64), sph(0xf0, b'\x30'))
@@ -93,10 +112,16 @@ def test_own_shared_memory(self):
 
 
 class Stage(unittest.TestCase):
-    def task(self):
-        i = ms.info(control(6), bytes(0x80))
+    def task(self, payload=256):
+        i = ms.info(control(6, payload=payload), bytes(0x80))
         return ms.to_stage(KERNEL % (ms.memory(i)[1] // 4, ms.memory(i)[0][2] // 4 + 1), i), i
 
+    def spill_mesh(self, with_input=True):
+        i = ms.info(control(7, payload=27904, task=1, max_v=64, max_p=126), sph(0xf0, b'\x30'))
+        lay, size = ms.memory(i)
+        src = KERNEL_IN % (size // 4, lay[4], lay[3] // 4) if with_input else KERNEL % (size // 4, lay[3] // 4)
+        return ms.to_stage(src, i), i
+
     def mesh(self, task=False, primitive=False, vertices=128, primitives=64, pos=0xf0, verts=0, reads=None):
         area = sph(pos, b'\x30') + (sph(0, b'\x01') if primitive else b'')
         i = ms.info(control(7, payload=128, task=int(task), max_v=vertices, max_p=primitives, prim_at=0x80 if primitive else 0, verts=verts), area)
@@ -135,6 +160,27 @@ def test_mesh_count_and_indices(self):
         self.assertIn(f'uint np = min(_23[{prim >> 2}] >> 24, 64u);', src)
         self.assertIn(f'uint at = {prim + 4}u + n;', src)
 
+    def test_spill_task_stores_its_payload_in_its_slice(self):
+        src, i = self.task(payload=27904); out = ms.memory(i)[0][2] // 4
+        self.assertIn(f'device uint* dm_spill [[buffer({abi.MESH_SPILL_SLOT})]]', src)
+        self.assertIn('struct dm_payload_t { uint slot; };', src)
+        self.assertIn('if (dm_tg >= dm_spill[0])', src)   # past the host's slices: no launch
+        self.assertIn('device uint* dm_slice = dm_spill + 4u + dm_tg * 6976u;', src)
+        self.assertIn(f'dm_slice[i] = _23[{out}u + i];', src)
+        self.assertIn('dm_payload.slot = dm_tg;', src)
+        self.assertNotIn('dm_payload.w', src)
+
+    def test_spill_mesh_reads_its_input_from_the_slice(self):
+        src, i = self.spill_mesh()
+        self.assertNotIn('threadgroup spvUnsafeArray<uint, 6976> _24;', src)   # the array left the body
+        self.assertIn('dm_in_t _24)', src)
+        self.assertIn('dm_in_t dm_in = { dm_spill + 4u + dm_payload.slot * 6976u, dm_tg };', src)   # word 0: the workgroup's number among its task's launches
+        self.assertIn(', _23, dm_in);', src)
+        self.assertIn(f'const device uint* dm_spill [[buffer({abi.MESH_SPILL_SLOT})]]', src)
+
+    def test_spill_mesh_without_its_input_refused(self):
+        with self.assertRaises(ValueError): self.spill_mesh(with_input=False)
+
     def test_mesh_with_a_task_reads_its_number_from_the_payload(self):
         src = self.mesh(task=True)[0]
         self.assertIn('_23[4] = dm_tg;', src)
@@ -217,7 +263,7 @@ def test_position_is_invariant(self):   # two mesh shaders of one geometry must
 
     @unittest.skipUnless(shutil.which('xcrun') and subprocess.run(['xcrun', '-sdk', 'macosx', '-f', 'metal'], capture_output=True).returncode == 0, 'no Metal compiler')
     def test_compiles(self):
-        cases = [self.task()[0], self.mesh()[0], self.mesh(task=True)[0], self.mesh(primitive=True)[0], self.mesh(vertices=32, primitives=32, pos=0xf8, verts=1, reads={5: 0x7})[0], self.mesh(primitive=True, reads={0: 1, 1: 0xf, 5: 0xf})[0], self.local_word()]
+        cases = [self.task()[0], self.task(payload=27904)[0], self.spill_mesh()[0], self.mesh()[0], self.mesh(task=True)[0], self.mesh(primitive=True)[0], self.mesh(vertices=32, primitives=32, pos=0xf8, verts=1, reads={5: 0x7})[0], self.mesh(primitive=True, reads={0: 1, 1: 0xf, 5: 0xf})[0], self.local_word()]
         with tempfile.TemporaryDirectory() as d:
             for n, src in enumerate(cases):
                 f = Path(d, f's{n}.metal'); f.write_text(src)

dm/gfx/tools/dmsh_chain.lock

@@ -1,5 +1,5 @@
 # Translation chain tool fingerprints (written by tool_lock.py --update-lock; regenerate after changing tools and commit it)
-dmsh_xlate	0406b50b2856b1b4
+dmsh_xlate	cebc4fc4564b4a4f
 xlate_source	657232fd201f16a1
 spirv-cross	e37945d8e441d34d
 metal_macosx	Apple metal version 32023.922 (metalfe-32023.922.4)