shader: complete native task/mesh translation

dm/gfx/tests/test_dmsh_ampere_mesh.py

@@ -28,9 +28,9 @@ def program(*insts): return bytes(0x80) + b''.join(insts) + inst(*EXIT)   # a ze
 # SPIR-V numbers the checks read
 OP_ENTRY_POINT, OP_EXECUTION_MODE, OP_DECORATE, OP_VARIABLE, OP_EMIT_MESH_TASKS, OP_SET_MESH_OUTPUTS = 15, 16, 71, 59, 5294, 5295
 MODEL_TASK, MODEL_MESH = 5364, 5365
-MODE_LOCAL_SIZE, MODE_OUTPUT_VERTICES, MODE_OUTPUT_PRIMITIVES, MODE_OUTPUT_TRIANGLES = 17, 26, 5270, 5298
+MODE_LOCAL_SIZE, MODE_OUTPUT_POINTS, MODE_OUTPUT_VERTICES, MODE_OUTPUT_PRIMITIVES, MODE_OUTPUT_LINES, MODE_OUTPUT_TRIANGLES = 17, 27, 26, 5270, 5269, 5298
 DEC_SPEC_ID, DEC_BUILTIN, DEC_LOCATION, DEC_PER_PRIMITIVE = 1, 11, 30, 5271
-BUILTIN_PRIMITIVE_TRIANGLE_INDICES = 5296
+BUILTIN_POINT_SIZE, BUILTIN_PRIMITIVE_POINT_INDICES, BUILTIN_PRIMITIVE_LINE_INDICES, BUILTIN_PRIMITIVE_TRIANGLE_INDICES = 1, 5294, 5295, 5296
 SC_TASK_PAYLOAD = 5402
 
 
@@ -86,6 +86,17 @@ def test_task(self):
         self.assertEqual(ops[OP_EMIT_MESH_TASKS][0][3], payload[0][1])   # the payload goes with the launch
         self.assertEqual(res['stage'], 'ts')
 
+    def test_task_reads_its_own_payload(self):
+        # 0x927 / 0xb8 writes a record, then 0x923 / 0xb8 reads the same word in another thread. The loaded value becomes the mesh-workgroup count, so it survives optimization.
+        code = program(stage(STAGE_LOAD, HEADER, rd=0), stage(STAGE_STORE, TASK_WRITE, rb=0, imm=0x10),
+                       stage(STAGE_LOAD, TASK_WRITE, rd=1, imm=0x10), stage(STAGE_STORE, TASK_WRITE, rb=1, imm=4))
+        (_, ops), _ = self.translate(code, stage='ts', local=[32, 1, 1], mesh=[0, 0, 64])
+        self.assertEqual([e[0] for e in ops[OP_ENTRY_POINT]], [MODEL_TASK])
+        self.assertEqual(len(ops[OP_EMIT_MESH_TASKS]), 1)
+
+    def test_task_refuses_mesh_topology(self):
+        self.refused(self.TASK, 'task program with mesh output', stage='ts', local=[32, 1, 1], mesh=[0, 0, 64, -5, 1])
+
     # A mesh program: its index (task memory word 0), a payload word by it, the four position planes and the generic planes of vertex 0, one per-primitive value of
     # primitive 0, triangle 0's indices (bytes 4..6) and the primitive count (byte 3)
     MESH = program(stage(STAGE_LOAD, TASK_READ, rd=1), stage(STAGE_LOAD, TASK_READ, rd=2, ra=1, imm=0x10),
@@ -121,6 +132,27 @@ def test_mesh_without_payload(self):   # a mesh program run without a task reads
                                      stage='ms', local=[32, 1, 1], mesh=[32, 16, 128, 0, 1, 2, 3])
         self.assertFalse([v for v in ops[OP_VARIABLE] if v[2] == SC_TASK_PAYLOAD])
 
+    def test_point_mesh_with_point_size(self):
+        code = program(*[stage(STAGE_STORE, PLANES, rb=1, imm=0x80 * p) for p in range(5)],
+                       stage(STAGE_STORE, INDICES, rb=RZ, imm=4), stage(STAGE_STORE, COUNTS, rb=1, imm=3))
+        (_, ops), _ = self.translate(code, stage='ms', local=[32, 1, 1], mesh=[32, 32, 64, -5, 1, -4, 0, 1, 2, 3])
+        ep = ops[OP_ENTRY_POINT][0]
+        modes = {m[1]: m[2:] for m in ops[OP_EXECUTION_MODE] if m[0] == ep[1]}
+        self.assertIn(MODE_OUTPUT_POINTS, modes)
+        decos = [d for d in ops[OP_DECORATE] if d[1] == DEC_BUILTIN]
+        self.assertTrue(any(d[2] == BUILTIN_PRIMITIVE_POINT_INDICES for d in decos))
+        self.assertTrue(any(d[2] == DEC_BUILTIN and d[3] == BUILTIN_POINT_SIZE for d in ops.get(72, [])))
+
+    def test_line_mesh(self):
+        code = program(*[stage(STAGE_STORE, PLANES, rb=1, imm=0x80 * p) for p in range(4)],
+                       stage(STAGE_STORE, INDICES, rb=RZ, imm=4), stage(STAGE_STORE, INDICES, rb=1, imm=5),
+                       stage(STAGE_STORE, COUNTS, rb=1, imm=3))
+        (_, ops), _ = self.translate(code, stage='ms', local=[32, 1, 1], mesh=[32, 32, 64, -5, 2, 0, 1, 2, 3])
+        ep = ops[OP_ENTRY_POINT][0]
+        modes = {m[1]: m[2:] for m in ops[OP_EXECUTION_MODE] if m[0] == ep[1]}
+        self.assertIn(MODE_OUTPUT_LINES, modes)
+        self.assertTrue(any(d[1] == DEC_BUILTIN and d[2] == BUILTIN_PRIMITIVE_LINE_INDICES for d in ops[OP_DECORATE]))
+
     def test_refused_mode(self):
         self.refused(program(stage(STAGE_STORE, 0x9800, rb=0)), 'stage memory mode 0x9800', stage='ms', local=[32, 1, 1], mesh=[32, 16, 128, 0, 1, 2, 3])
 

dm/gfx/tests/test_dmsh_raw_reuse.py

@@ -0,0 +1,142 @@
+#!/usr/bin/env python3
+"""Raw-global lookups use the current hint ABI and vectorize wide loads.
+
+The first access tries the tagged 2 MiB-chunk hint before a looped table search.
+Later accesses first check the entry cached for this invocation. LDG.64 and
+LDG.128 resolve once and perform one aligned vector load, including zeroing all
+follower words when the address misses every pool.
+"""
+import json
+import os
+from pathlib import Path
+import re
+import shutil
+import subprocess
+import sys
+import tempfile
+import unittest
+
+sys.path.insert(0, str(Path(__file__).resolve().parent))
+from test_dmsh_ampere_mem import EXIT, Translated, inst  # noqa: E402
+
+TOOLS = Path(__file__).resolve().parents[1] / 'tools'
+STORE = (0x000fe8000000c800, 0x0000080609008388)
+LDG = (0x000ea200081e890f, 0x0000000405067981)
+VEC128 = ((LDG[0] & ~(7 << 9)) | (6 << 9),
+          (LDG[1] & ~0xff0000) | (12 << 16))
+
+
+def waiting(instruction):
+    """Make the instruction wait for every scoreboard dependency."""
+    return (instruction[0] | 0x3f << 52, instruction[1])
+
+
+def s2r_tid(register):
+    """Read SR_TID.X into *register* using the fixture's known SM86 form."""
+    return (0x000e260000002100,
+            0x0000000000007919 | register << 16)
+
+
+@unittest.skipUnless(os.environ.get('DMSH_TEST_XLATE'),
+                     'set DMSH_TEST_XLATE to the offline translator')
+@unittest.skipUnless(shutil.which('spirv-dis'), 'needs spirv-dis')
+class RawReuse(Translated):
+    def spv_dis(self, code):
+        sys.path.insert(0, str(TOOLS))
+        import dmsh_geometry
+
+        with tempfile.TemporaryDirectory(prefix='raw_reuse_') as directory:
+            work = Path(directory)
+            (work / 'k.bin').write_bytes(
+                b''.join(inst(*waiting(op)) for op in code) + inst(*EXIT))
+            job = [dict(key='k', stage='cs', arch='sm86',
+                        code=str(work / 'k.bin'), tex={}, img={}, attr={},
+                        fout={}, local=[32, 1, 1], debug=True,
+                        nvn2_cs=[0, 0, 3, 0x40, 0], shared=0x1000)]
+            (work / 'jobs.json').write_text(json.dumps(job))
+            run = subprocess.run(
+                [os.environ['DMSH_TEST_XLATE'], str(work / 'jobs.json'),
+                 directory], capture_output=True, text=True)
+            self.assertEqual(run.returncode, 0, run.stdout + run.stderr)
+            dmsh_geometry.flatten_sets(work / 'k.spv')
+            if shutil.which('spirv-val'):
+                valid = subprocess.run(
+                    ['spirv-val', '--target-env', 'vulkan1.3',
+                     str(work / 'k.spv')], capture_output=True, text=True)
+                self.assertEqual(valid.returncode, 0,
+                                 valid.stdout + valid.stderr)
+            return subprocess.run(
+                ['spirv-dis', str(work / 'k.spv')], capture_output=True,
+                text=True, check=True).stdout
+
+    @staticmethod
+    def functions(disassembly):
+        return {match.group(1): match.group(0) for match in re.finditer(
+            r'(%\w+) = OpFunction %int None .*?OpFunctionEnd',
+            disassembly, re.S)}
+
+    @staticmethod
+    def main_function(disassembly):
+        return re.search(
+            r'%\w+ = OpFunction %void None .*?OpFunctionEnd',
+            disassembly, re.S).group(0)
+
+    def test_loads_share_one_zero_initialized_last_entry(self):
+        second = (LDG[0], LDG[1] | 4 << 40)
+        disassembly = self.spv_dis(
+            (s2r_tid(9), s2r_tid(5), LDG, STORE, second, STORE))
+        self.assertEqual(len(re.findall(
+            r'OpVariable \S+ Private %\w+', disassembly)), 1)
+        main = self.main_function(disassembly)
+        calls = re.findall(r'OpFunctionCall %int (%\w+)', main)
+        self.assertEqual(len(calls), 2)
+        self.assertEqual(calls[0], calls[1])
+        full = self.functions(disassembly)[calls[0]]
+        self.assertIn('OpLoopMerge', full)
+        self.assertIn('OpConstant %int 8191', disassembly)
+        self.assertRegex(full, r'OpShiftRightLogical %int \S+ %int_21')
+        self.assertRegex(full, r'OpBitwiseAnd %int \S+ %int_8191')
+        self.assertRegex(full, r'OpShiftRightLogical %int \S+ %int_13')
+
+    def test_128_bit_load_uses_one_vector_helper(self):
+        stores = tuple((STORE[0], (STORE[1] & ~0xff000000) |
+                        (register << 24) | (offset << 40))
+                       for register, offset in
+                       ((12, 0), (13, 4), (14, 16), (15, 20)))
+        disassembly = self.spv_dis(
+            (s2r_tid(9), s2r_tid(5), VEC128, *stores))
+        main = self.main_function(disassembly)
+        calls = re.findall(r'OpFunctionCall %int (%\w+)', main)
+        self.assertEqual(len(calls), 1)
+        functions = self.functions(disassembly)
+        helper = functions[calls[0]]
+        self.assertIn('OpLoopMerge', helper)
+        self.assertIn('%int_15', helper,
+                      'the resolver must check the 128-bit load\'s last byte')
+        self.assertRegex(helper, r'OpLoad %v4uint \S+ Aligned 4')
+        self.assertGreaterEqual(
+            len(re.findall(r'OpStore \S+ %uint_0', helper)), 4,
+            'a miss must zero all three follower words and the last-entry cache')
+
+    def test_adjacent_vectors_do_not_share_follower_words(self):
+        first = (VEC128[0] & ~0xf, VEC128[1])
+        second = (first[0],
+                  (first[1] & ~(0xff << 16) & ~(0xffffff << 40)) |
+                  (16 << 16) | (0x40 << 40))
+        keep_first = (STORE[0], (STORE[1] & ~0xff000000) | (12 << 24))
+        keep_second = (STORE[0],
+                       (STORE[1] & ~0xff000000 & ~(0xffffff << 40)) |
+                       (18 << 24) | (4 << 40))
+        disassembly = self.spv_dis(
+            (s2r_tid(9), s2r_tid(5), first, second,
+             keep_first, keep_second))
+        calls = re.findall(
+            r'OpFunctionCall %int (%\w+)', self.main_function(disassembly))
+        self.assertEqual(len(calls), 2)
+        functions = self.functions(disassembly)
+        self.assertEqual(sum('OpLoad %v4uint' in functions[call]
+                             for call in set(calls)), 1)
+
+
+if __name__ == '__main__':
+    unittest.main()

dm/third_party/ryujinx/src/Ryujinx.Graphics.Shader/CodeGen/Glsl/Declarations.cs

@@ -331,7 +331,7 @@ private static void DeclareMemories(CodeGenContext context, IEnumerable<MemoryDe
                 }
                 else
                 {
-                    context.AppendLine($"{prefix}{typeName} {memory.Name};");
+                    context.AppendLine($"{prefix}{typeName} {memory.Name}{(memory.ZeroInit ? $" = {typeName}(0)" : string.Empty)};");
                 }
             }
         }

dm/third_party/ryujinx/src/Ryujinx.Graphics.Shader/CodeGen/Spirv/Declarations.cs

@@ -116,8 +116,9 @@ private static void DeclareMemories(
                 }
 
                 StorageClass memoryStorage = SharedStorageClass(storage, memory);
-                SpvInstruction pointerType = context.TypePointer(memoryStorage, context.GetType(memory.Type, memory.ArrayLength));
-                SpvInstruction variable = context.Variable(pointerType, memoryStorage);
+                SpvInstruction type = context.GetType(memory.Type, memory.ArrayLength);
+                SpvInstruction pointerType = context.TypePointer(memoryStorage, type);
+                SpvInstruction variable = context.Variable(pointerType, memoryStorage, memory.ZeroInit ? context.ConstantNull(type) : null);
 
                 context.AddGlobalVariable(variable);
 

dm/third_party/ryujinx/src/Ryujinx.Graphics.Shader/CodeGen/Spirv/Instructions.cs

@@ -1800,6 +1800,11 @@ private static OperationResult GenerateAtomicMemoryCas(CodeGenContext context, A
 
         private static OperationResult GenerateLoadOrStore(CodeGenContext context, AstOperation operation, bool isStore)
         {
+            if (!isStore && operation.StorageKind == StorageKind.GlobalMemoryVector)
+            {
+                return GenerateRawVectorLoad(context, operation);
+            }
+
             // nxnative (after azmasm/switch2mac, MIT): a Switch 2 mesh program's store into its per-vertex / per-primitive planes goes straight into the mesh outputs
             // (CodeGenContext.MeshPlaneStore, written by SpirvGenerator.GenerateMeshEntry), so the values are not held twice
             if (isStore && operation.StorageKind == StorageKind.SharedMemory && context.Properties.Mesh is { IsMesh: true } mesh && operation.SourcesCount == 3 &&
@@ -1840,6 +1845,48 @@ private static OperationResult GenerateLoadOrStore(CodeGenContext context, AstOp
             }
         }
 
+        // Load an aligned 64- or 128-bit raw value through one physical
+        // pointer. Word zero is returned; the rest go to the private array
+        // consumed by the original follower loads.
+        private static OperationResult GenerateRawVectorLoad(CodeGenContext context, AstOperation operation)
+        {
+            if (operation.GetSource(2) is not AstOperand { Type: OperandType.Constant } countOperand ||
+                operation.GetSource(3) is not AstOperand { Type: OperandType.Constant } memoryOperand)
+            {
+                throw new InvalidOperationException(
+                    $"The word count and local memory of a {operation.StorageKind} load must be constants.");
+            }
+
+            int count = countOperand.Value;
+            SpvInstruction vectorType = context.TypeVector(context.TypeU32(), count);
+            SpvInstruction addressWords = context.CompositeConstruct(
+                context.TypeVector(context.TypeU32(), 2),
+                context.Get(AggregateType.U32, operation.GetSource(0)),
+                context.Get(AggregateType.U32, operation.GetSource(1)));
+            SpvInstruction pointer = context.ConvertUToPtr(
+                context.TypePointer(StorageClass.PhysicalStorageBuffer, vectorType),
+                context.Bitcast(context.TypeInt(64, 0), addressWords));
+
+            SpvInstruction value = context.Load(vectorType, pointer, MemoryAccessMask.Aligned);
+            value.AddOperand((SpvLiteralInteger)4);
+            SpvInstruction wordPointerType = context.TypePointer(StorageClass.Private, context.TypeU32());
+
+            for (int word = 1; word < count; word++)
+            {
+                SpvInstruction element = context.AccessChain(
+                    wordPointerType,
+                    context.LocalMemories[memoryOperand.Value],
+                    context.Constant(context.TypeS32(), word - 1));
+                context.Store(
+                    element,
+                    context.CompositeExtract(context.TypeU32(), value, (SpvLiteralInteger)word));
+            }
+
+            return new OperationResult(
+                AggregateType.U32,
+                context.CompositeExtract(context.TypeU32(), value, (SpvLiteralInteger)0));
+        }
+
         private static SpvInstruction GetStoragePointer(CodeGenContext context, AstOperation operation, out AggregateType varType)
         {
             StorageKind storageKind = operation.StorageKind;