Implement RDNA2 per-invocation scratch memory

tests/ShaderRecompilerComputeTests.cpp

@@ -1070,6 +1070,7 @@ CompiledShader CompileCase(const TestCase &test) {
   options.read_memory = ReadTestMemory;
   options.read_memory_data = const_cast<std::vector<u32> *>(&test.initial);
   options.flat_memory_base = test.flat_memory_base;
+  options.scratch_dwords = test.compute_info.scratch_size_dwords;
   if (test.has_compute_info) {
     options.wave_size = test.compute_info.wave_size;
   }
@@ -16335,6 +16336,36 @@ TestCase FlatSegmentIgnoresSaddrAndMasksOffsetMsb() {
   return test;
 }
 
+TestCase ScratchIsPrivatePerInvocation() {
+  using O = ShaderOpcode;
+
+  std::vector<u32> code;
+  code.push_back(EncodeSMovB32(20, InlineU32(0)));
+  AppendVMovU32(&code, 21, 0);
+  code.push_back(EncodeFlat0(0x1c, 1, 4));
+  code.push_back(EncodeFlat1(0, 20, 0, 21));
+  code.push_back(EncodeFlat0(0x0c, 1, 4));
+  code.push_back(EncodeFlat1(5, 20, 0, 21));
+  AppendStoreVgprAtLaneDwordOffset(&code, 5, 0, 0);
+  AppendEnd(&code);
+
+  TestCase test;
+  test.name = "ScratchIsPrivatePerInvocation";
+  test.code = std::move(code);
+  test.expected = {0, 1, 2, 3};
+  test.opcodes = {O::S_MOV_B32,          O::V_MOV_B32,
+                  O::FLAT_STORE_DWORD,   O::FLAT_LOAD_DWORD,
+                  O::BUFFER_STORE_DWORD, O::S_ENDPGM};
+  test.compute_info.threads_num[0] = 4;
+  test.compute_info.threads_num[1] = 1;
+  test.compute_info.threads_num[2] = 1;
+  test.compute_info.thread_ids_num = 1;
+  test.compute_info.workgroup_register = 0;
+  test.compute_info.scratch_size_dwords = 2;
+  test.has_compute_info = true;
+  return test;
+}
+
 TestCase FlatStoreVariants() {
   using O = ShaderOpcode;
 
@@ -18648,6 +18679,7 @@ std::vector<TestCase> MakeCases() {
   AddCase(FlatVirtualAddressRebasesGuestAllocation);
   AddCase(GlobalSignedImmediateRebasesBeforeSaddr);
   AddCase(FlatSegmentIgnoresSaddrAndMasksOffsetMsb);
+  AddCase(ScratchIsPrivatePerInvocation);
   AddCase(FlatStoreVariants);
   AddCase(DsReadWriteVariants);
   AddCase(DsAppendConsumeUsesEncodedLdsSelector);

tests/shaderCfgTests.cpp

@@ -11059,6 +11059,22 @@ void TestComputeLdsAllocationIdentity() {
             ShaderGetIdCS(regs, lds_896, true),
         "compute pipeline identity omitted the LDS allocation");
 
+  mapped.scratch_size_dwords = 7;
+  ShaderMapUserData(regs.cs_regs.data_addr, mapped);
+  regs.cs_regs.lds_size = decode_lds_field(lds_896_rsrc2);
+  ShaderComputeInputInfo scratch_info{};
+  std::span<const uint32_t> scratch_spirv;
+  Check(ShaderCompileInfoCS(regs, sh, scratch_info, scratch_spirv),
+        "compute scratch metadata shader did not compile");
+  Check(scratch_info.scratch_size_dwords == 7,
+        "AGC per-thread scratch size was not propagated");
+  Check(scratch_info.stage.program != nullptr &&
+            scratch_info.stage.program->scratch_dwords == 7,
+        "AGC per-thread scratch size did not reach the compiler program");
+  Check(ShaderGetIdCS(regs, scratch_info, true) !=
+            ShaderGetIdCS(regs, lds_896, true),
+        "compute pipeline identity omitted the scratch allocation");
+
   const uint32_t append_shader[] = {
       EncodeSMovB32(124, 132), // m0 = 4 bytes
       EncodeDs0(0x3e),
@@ -11217,18 +11233,21 @@ void TestNewShaderRecompilerFlatAddressProvenanceBoundaries() {
   options.user_data = user_data;
   options.user_data_count = static_cast<uint32_t>(std::size(user_data));
   options.flat_memory_base = 0;
+  options.scratch_dwords = 1;
   ShaderRecompiler::CompileResult result;
   std::string error;
   const bool compiled =
       ShaderRecompiler::TryRecompile(segmented_shader, options, result, &error);
   Check(compiled, error.c_str());
-  Check(result.program.info.addresses.size() == 2,
-        "segmented address resources were not tracked independently");
-  for (const auto &address : result.program.info.addresses) {
-    Check(
-        address.unbased && address.source == UINT32_MAX,
-        "GLOBAL/SCRATCH null-SADDR incorrectly inherited FLAT VGPR provenance");
-  }
+  Check(result.program.info.addresses.size() == 1 &&
+            result.program.info.addresses[0].kind ==
+                ShaderRecompiler::IR::ResourceKind::Global &&
+            result.program.info.addresses[0].unbased &&
+            result.program.info.addresses[0].source == UINT32_MAX,
+        "GLOBAL null-SADDR did not remain an unbased guest address");
+  Check(Common::ContainsStr(result.ir_dump, "GetScratchResource") &&
+            result.program.scratch_dwords == 1,
+        "SCRATCH incorrectly entered guest address tracking");
 }
 
 void TestNewShaderRecompilerSpirvSizeBaselines() {