Implement RDNA2 per-invocation scratch memory

src/graphics/shader/recompiler/ShaderRecompiler.cpp

@@ -46,7 +46,8 @@ std::string MakeIrDump(const CFG::Graph& cfg, const IR::Program& ir) {
 	std::string dump = "CFG:\n";
 	dump += CFG::GraphToString(cfg);
 	dump += "\nIR:\n";
-	dump += fmt::format("mode={}\n", ir.dispatcher_fallback ? "dispatcher" : "structured");
+	dump += fmt::format("mode={} scratch_dwords={}\n",
+	                    ir.dispatcher_fallback ? "dispatcher" : "structured", ir.scratch_dwords);
 	dump += ir.values != nullptr ? IR::ValueProgramToString(*ir.values) : IR::ProgramToString(ir);
 	return dump;
 }
@@ -686,6 +687,7 @@ bool TryRecompile(std::span<const uint32_t> code, const CompileOptions& options,
 		output.shader_hash     = options.shader_hash;
 		output.user_data_base  = options.user_data_base;
 		output.user_data_count = options.user_data_count;
+		output.scratch_dwords  = options.scratch_dwords;
 		return true;
 	};
 	if (!lower_program(ir)) {

src/graphics/shader/recompiler/ShaderRecompiler.h

@@ -18,6 +18,7 @@ struct CompileOptions {
 	uint32_t                    wave_size       = 64;
 	uint32_t                    user_data_base  = 0;
 	uint32_t                    user_data_count = 64;
+	uint32_t                    scratch_dwords  = 0;
 	uint64_t                    shader_hash     = 0;
 	uint64_t                    shader_base     = 0;
 	std::optional<uint64_t>     flat_memory_base;

src/graphics/shader/recompiler/backend/spirv/SpirvEmitter.cpp

@@ -158,6 +158,11 @@ bool ValidateNativeProgram(const IR::Program& program, std::string* error) {
 						return Fail(error, "typed address handle has an invalid dense resource");
 					}
 					break;
+				case IR::ValueOpcode::GetScratchResource:
+					if (inst.NumArgs() != 0 || program.scratch_dwords == 0) {
+						return Fail(error, "typed scratch handle has invalid shader metadata");
+					}
+					break;
 				case IR::ValueOpcode::GetImageResource:
 					if (dense >= program.info.images.size()) {
 						return Fail(error, "typed image handle has an invalid dense resource");
@@ -198,6 +203,18 @@ bool AnalyzeProgramRequirements(IR::Program& program, std::string* error) {
 	};
 	for (const auto* block: program.values->blocks) {
 		for (const auto& inst: *block) {
+			if (IR::AddressOpcodeInfoOf(inst.GetOpcode()).access != IR::AddressAccess::None) {
+				const auto memory_index = inst.Flags<IR::MemoryFlags>().index;
+				if (memory_index >= program.values->memory_info.size()) {
+					return Fail(error, "address operation has invalid memory metadata");
+				}
+				if (program.values->memory_info[memory_index].kind == IR::ResourceKind::Scratch) {
+					if (program.scratch_dwords == 0) {
+						return Fail(error, "scratch operation has no per-thread storage");
+					}
+					requirements.function_scratch = true;
+				}
+			}
 			if (IR::BufferAccessOf(inst.GetOpcode()) != IR::BufferAccess::None) {
 				const auto memory_index = inst.Flags<IR::MemoryFlags>().index;
 				if (memory_index >= program.values->memory_info.size()) {

src/graphics/shader/recompiler/backend/spirv/spirvEmitterFlow.cpp

@@ -535,6 +535,7 @@ bool EmitValueFlow(ValueEmitContext& ctx, const IR::Inst& inst) {
 		case IR::ValueOpcode::GetSrtResource:
 		case IR::ValueOpcode::GetBufferResource:
 		case IR::ValueOpcode::GetAddressResource:
+		case IR::ValueOpcode::GetScratchResource:
 		case IR::ValueOpcode::GetImageResource:
 		case IR::ValueOpcode::GetSamplerResource:
 		case IR::ValueOpcode::MakeImageAddress: return true;

src/graphics/shader/recompiler/backend/spirv/spirvEmitterInternal.h

@@ -321,6 +321,7 @@ struct EmitterState {
 	uint32_t                                         vsharp_storage_variable = 0;
 	uint32_t                                         flattened_srt_variable  = 0;
 	uint32_t                                         lds_variable            = 0;
+	uint32_t                                         scratch_variable        = 0;
 	std::array<uint32_t, 14>                         sampled_image_variables {};
 	std::array<uint32_t, 10>                         storage_image_variables {};
 	uint32_t                                         sampler_variable                      = 0;
@@ -356,8 +357,8 @@ uint32_t TypeFunction(EmitterState& state);
 uint32_t TypeStorageBufferPointer(EmitterState& state);
 uint32_t TypeStorageBufferElementPointer(EmitterState& state);
 uint32_t TypePushConstantElementPointer(EmitterState& state);
-uint32_t TypeLdsArrayPointer(EmitterState& state, uint32_t storage_class);
+uint32_t TypeU32ArrayPointer(EmitterState& state, uint32_t storage_class, uint32_t dwords);
-uint32_t TypeLdsElementPointer(EmitterState& state, uint32_t storage_class);
+uint32_t TypeU32ElementPointer(EmitterState& state, uint32_t storage_class);
 
 inline void EmitLabel(EmitterState& state, uint32_t label) {
 	state.current_label = label;

src/graphics/shader/recompiler/backend/spirv/spirvEmitterMemory.cpp

@@ -105,12 +105,12 @@ uint32_t AddressByteAddress(ValueEmitContext& ctx, const IR::Inst& inst, const I
 	}
 	const auto immediate_low  = ConstantU32(state, static_cast<uint32_t>(immediate));
 	const auto immediate_high = ConstantU32(state, immediate < 0 ? UINT32_MAX : 0u);
-	const auto base           = state.program.info.addresses[mem.resource].specialized_base;
 	if (mem.kind == IR::ResourceKind::Scratch) {
 		low              = AddU64Low(state, low, high, immediate_low, immediate_high, high);
 		const auto valid = Binary(state, OpIEqual, TypeBool(state), high, ConstantU32(state, 0));
 		return Select(state, TypeU32(state), valid, low, ConstantU32(state, UINT32_MAX));
 	}
+	const auto base = state.program.info.addresses[mem.resource].specialized_base;
 	if (mem.kind == IR::ResourceKind::Flat ||
 	    (mem.kind == IR::ResourceKind::Global &&
 	     state.program.info.addresses[mem.resource].unbased)) {

src/graphics/shader/recompiler/backend/spirv/spirvEmitterMemoryHelpers.cpp

@@ -75,7 +75,8 @@ static void EnsureLdsStorage(EmitterState& state) {
 		EXIT("function LDS was not prepared before SPIR-V function emission\n");
 	}
 	state.lds_variable = state.builder.DefineGlobalVariable(
-	    TypeLdsArrayPointer(state, StorageClassWorkgroup), StorageClassWorkgroup);
+	    TypeU32ArrayPointer(state, StorageClassWorkgroup, LdsDwordCount(state)),
+	    StorageClassWorkgroup);
 	state.builder.AddName(state.lds_variable, "lds_dwords");
 }
 
@@ -98,10 +99,16 @@ MemoryResourceAccess PrepareMemoryResourceAccess(EmitterState& state, const IR::
 			}
 			access.length = state.gds_length;
 			return access;
+		case IR::ResourceKind::Scratch:
+			if (state.scratch_variable == 0) {
+				EXIT("scratch storage was not prepared before SPIR-V function emission\n");
+			}
+			access.object_pointer = state.scratch_variable;
+			access.length         = ConstantU32(state, state.program.scratch_dwords);
+			return access;
 		case IR::ResourceKind::ScalarAddress:
 		case IR::ResourceKind::Flat:
-		case IR::ResourceKind::Global:
+		case IR::ResourceKind::Global: {
-		case IR::ResourceKind::Scratch: {
 			if (state.address_memory_variable == 0) {
 				ExitDescriptorBindingFailure(state, IR::DescriptorBindingKind::AddressMemory,
 				                             mem.resource,
@@ -155,11 +162,11 @@ uint32_t EmitMemoryElementInBounds(EmitterState& state, const MemoryResourceAcce
 uint32_t EmitMemoryElementPointer(EmitterState& state, const MemoryResourceAccess& access,
                                   uint32_t index) {
 	const auto pointer = state.builder.AllocateId();
-	if (access.kind == IR::ResourceKind::Lds) {
+	if (access.kind == IR::ResourceKind::Lds || access.kind == IR::ResourceKind::Scratch) {
-		state.builder.AddFunction({OpAccessChain,
+		const auto storage_class = access.kind == IR::ResourceKind::Scratch ? StorageClassFunction
-		                           TypeLdsElementPointer(state, state.stage == ShaderType::Compute
+		                           : state.stage == ShaderType::Compute     ? StorageClassWorkgroup
-		                                                            ? StorageClassWorkgroup
+		                                                                    : StorageClassFunction;
-		                                                            : StorageClassFunction),
+		state.builder.AddFunction({OpAccessChain, TypeU32ElementPointer(state, storage_class),
 		                           pointer, access.object_pointer, index});
 	} else {
 		state.builder.AddFunction({OpAccessChain, TypeStorageBufferElementPointer(state), pointer,

src/graphics/shader/recompiler/backend/spirv/spirvEmitterModule.cpp

@@ -94,13 +94,13 @@ uint32_t TypePushConstantElementPointer(EmitterState& state) {
 	return TypePointer(state, StorageClassPushConstant, TypeU32(state));
 }
 
-uint32_t TypeLdsArrayPointer(EmitterState& state, uint32_t storage_class) {
+uint32_t TypeU32ArrayPointer(EmitterState& state, uint32_t storage_class, uint32_t dwords) {
-	const auto count = ConstantU32(state, std::max(LdsDwordCount(state), 1u));
+	const auto count = ConstantU32(state, std::max(dwords, 1u));
 	const auto array = state.builder.Type(OpTypeArray, {TypeU32(state), count});
 	return TypePointer(state, storage_class, array);
 }
 
-uint32_t TypeLdsElementPointer(EmitterState& state, uint32_t storage_class) {
+uint32_t TypeU32ElementPointer(EmitterState& state, uint32_t storage_class) {
 	return TypePointer(state, storage_class, TypeU32(state));
 }
 
@@ -601,6 +601,9 @@ void DefineModule(EmitterState& state) {
 	if (state.requirements.function_lds) {
 		state.lds_variable = state.builder.AllocateId();
 	}
+	if (state.requirements.function_scratch) {
+		state.scratch_variable = state.builder.AllocateId();
+	}
 	state.main_func   = state.builder.AllocateId();
 	state.entry_label = state.builder.AllocateId();
 
@@ -659,6 +662,9 @@ void DefineModule(EmitterState& state) {
 	if (state.requirements.function_lds) {
 		state.builder.AddName(state.lds_variable, "lds_dwords");
 	}
+	if (state.requirements.function_scratch) {
+		state.builder.AddName(state.scratch_variable, "scratch_dwords");
+	}
 	AddInputAnnotationsAndNames(state);
 	AddOutputAnnotationsAndNames(state);
 	AddDescriptorAnnotationsAndNames(state);

src/graphics/shader/recompiler/backend/spirv/spirvEmitterProgram.cpp

@@ -580,9 +580,16 @@ bool EmitValueProgram(EmitterState& state, const IR::ValueProgram& program, std:
 	    {OpFunction, TypeVoid(state), state.main_func, FunctionControlNone, TypeFunction(state)});
 	EmitLabel(state, state.entry_label);
 	if (state.requirements.function_lds) {
-		state.builder.AddFunction({OpVariable, TypeLdsArrayPointer(state, StorageClassFunction),
+		state.builder.AddFunction(
+		    {OpVariable, TypeU32ArrayPointer(state, StorageClassFunction, LdsDwordCount(state)),
 		     state.lds_variable, StorageClassFunction});
 	}
+	if (state.requirements.function_scratch) {
+		state.builder.AddFunction(
+		    {OpVariable,
+		     TypeU32ArrayPointer(state, StorageClassFunction, state.program.scratch_dwords),
+		     state.scratch_variable, StorageClassFunction});
+	}
 	if (state.pixel_valid_mask_variable != 0) {
 		state.builder.AddFunction({OpVariable,
 		                           TypePointer(state, StorageClassFunction, TypeU32(state)),

src/graphics/shader/recompiler/frontend/translate/Memory.cpp

@@ -34,7 +34,7 @@ Translator::AddressOperands Translator::ReadAddressOperands(const IR::Instructio
 		                            high_or_base.reg.file != IR::RegisterFile::Vector
 		                        ? ReadU32(high_or_base)
 		                        : low;
-		return {GetAddressResource(offset, IR::Value(0u)), offset, IR::Value(0u)};
+		return {ir.Emit(IR::ValueOpcode::GetScratchResource), offset, IR::Value(0u)};
 	}
 	if (inst.memory.kind == IR::ResourceKind::Global &&
 	    high_or_base.kind == IR::OperandKind::Register &&

src/graphics/shader/recompiler/ir/ShaderIR.cpp

@@ -222,9 +222,9 @@ MemoryInfo MemoryInfoFromDecoded(const Decoder::Instruction& decoded, ResourceKi
 		mem.image_nsa_addr[i] = decoded.image_nsa_addr[i];
 	}
 	mem.memory_segment = decoded.memory_segment;
-	mem.address_is_full =
+	mem.address_is_full = kind == ResourceKind::Flat ||
-	    kind == ResourceKind::Flat || kind == ResourceKind::Scratch ||
+	                      (kind == ResourceKind::Global &&
-	    (kind == ResourceKind::Global && decoded.src1.kind == Decoder::OperandKind::Vgpr);
+	                       decoded.src1.kind == Decoder::OperandKind::Vgpr);
 	mem.data_signed   = decoded.data_signed;
 	mem.typed         = decoded.typed;
 	mem.formatted     = decoded.formatted;
@@ -236,7 +236,7 @@ MemoryInfo MemoryInfoFromDecoded(const Decoder::Instruction& decoded, ResourceKi
 	mem.offen         = decoded.offen;
 	mem.resource      = ResourceIndexFromOperand(decoded.src1);
 	mem.sampler       = ResourceIndexFromOperand(decoded.src2);
-	if (kind == ResourceKind::Lds || kind == ResourceKind::Gds) {
+	if (kind == ResourceKind::Lds || kind == ResourceKind::Gds || kind == ResourceKind::Scratch) {
 		mem.resource = 0;
 		mem.sampler  = 0;
 	}

src/graphics/shader/recompiler/ir/ShaderIR.h

@@ -430,6 +430,7 @@ struct SpirvRequirements {
 	bool compute_derivatives          = false;
 	bool image_gather_extended        = false;
 	bool function_lds                 = false;
+	bool function_scratch             = false;
 	bool pixel_valid_mask             = false;
 };
 
@@ -440,6 +441,7 @@ struct Program {
 	uint32_t                      wave_size           = 64;
 	uint32_t                      user_data_base      = 0;
 	uint32_t                      user_data_count     = 64;
+	uint32_t                      scratch_dwords      = 0;
 	bool                          dispatcher_fallback = false;
 	CFG::FailureKind              cfg_failure_kind    = CFG::FailureKind::None;
 	std::string                   fallback_reason;

src/graphics/shader/recompiler/ir/opcodes/ValueOpcodes.inc

@@ -192,6 +192,7 @@ VALUE_OPCODE(Permlane16U32, U32, U32, U32, U32, U1)
 VALUE_OPCODE(GetSrtResource, SrtResource)
 VALUE_OPCODE(GetBufferResource, BufferResource, U32, U32, U32, U32)
 VALUE_OPCODE(GetAddressResource, AddressResource, U32, U32)
+VALUE_OPCODE(GetScratchResource, AddressResource)
 VALUE_OPCODE(GetImageResource, ImageResource, U32, U32, U32, U32, U32, U32, U32, U32)
 VALUE_OPCODE(GetSamplerResource, SamplerResource, U32, U32, U32, U32)
 VALUE_OPCODE(MakeImageAddress, ImageAddress, U32, U32, U32, U32, U32, U32, U32, U32, U32, U32, U32,

src/graphics/shader/recompiler/ir/passes/ResourceMaterialization.cpp

@@ -679,10 +679,9 @@ bool MaterializeResources(const Program& program, const SrtRuntime& runtime,
 		} else {
 			if (!runtime.flat_memory_base.has_value()) {
 				if (error != nullptr) {
-					*error =
+					*error = fmt::format("unbased {} address at pc 0x{:08x} requires runtime "
-					    fmt::format("unbased {} address at pc 0x{:08x} requires runtime "
 					                     "guest-address translation",
-					                address.kind == ResourceKind::Flat ? "FLAT" : "global/scratch",
+					                     address.kind == ResourceKind::Flat ? "FLAT" : "global",
 					                     address.first_use_pc);
 				}
 				return false;

src/graphics/shader/recompiler/ir/passes/ResourceTracking.cpp

@@ -764,6 +764,18 @@ private:
 			if (!IsAddressResourceKind(memory.kind)) {
 				return Fail(flags.pc, error, "address operation has invalid resource kind");
 			}
+			if (memory.kind == ResourceKind::Scratch) {
+				handle = inst.Arg(0).Resolve().TryInstruction();
+				if (handle == nullptr || handle->GetOpcode() != ValueOpcode::GetScratchResource ||
+				    handle->NumArgs() != 0) {
+					return Fail(flags.pc, error, "scratch operation requires GetScratchResource");
+				}
+				if (m_program.scratch_dwords == 0) {
+					return Fail(flags.pc, error,
+					            "scratch operation requires a nonzero AGC per-thread size");
+				}
+				return true;
+			}
 			bool unbased = false;
 			if (!GetAddressHandle(inst, inst.Arg(0), flags.pc, handle, source, unbased, error)) {
 				return false;

src/graphics/shader/shader.cpp

@@ -794,6 +794,7 @@ static bool ShaderGetStaticInputInfoVS(const HW::VertexShaderInfo& regs,
 		     shader_addr);
 		return false;
 	}
+	info.scratch_size_dwords = data.scratch_size_dwords;
 
 	if (data.user_data == nullptr) {
 		LOGF("ShaderGetInputInfoVS(): no AGC user data for shader=0x%016" PRIx64 " es=0x%016" PRIx64
@@ -842,6 +843,12 @@ static void ShaderGetStaticInputInfoPS(
 	KYTY_PROFILER_FUNCTION();
 
 	ps_info = {};
+	ShaderMappedData data;
+	if (!ShaderGetMappedData(regs.ps_regs.data_addr, data)) {
+		EXIT("ShaderGetInputInfoPS(): shader=0x%016" PRIx64 " is missing from ShaderMap\n",
+		     regs.ps_regs.data_addr);
+	}
+	ps_info.scratch_size_dwords = data.scratch_size_dwords;
 
 	// SPI_PS_IN_CONTROL.NUM_INTERP occupies bits 5:0. Keep the remaining control
 	// flags in the hardware state and extract only the input count here.
@@ -891,11 +898,17 @@ static void ShaderGetStaticInputInfoCS(const HW::ComputeShaderInfo& regs,
                                        const HW::ShaderRegisters& /*sh*/,
                                        ShaderComputeInputInfo& info) {
 	info = {};
+	ShaderMappedData data;
+	if (!ShaderGetMappedData(regs.cs_regs.data_addr, data)) {
+		EXIT("ShaderGetInputInfoCS(): shader=0x%016" PRIx64 " is missing from ShaderMap\n",
+		     regs.cs_regs.data_addr);
+	}
 
 	info.threads_num[0]      = regs.cs_regs.num_thread_x;
 	info.threads_num[1]      = regs.cs_regs.num_thread_y;
 	info.threads_num[2]      = regs.cs_regs.num_thread_z;
 	info.lds_size_dwords     = static_cast<uint32_t>(regs.cs_regs.lds_size) * 128u;
+	info.scratch_size_dwords = data.scratch_size_dwords;
 	info.group_id[0]         = regs.cs_regs.tgid_x_en != 0;
 	info.group_id[1]         = regs.cs_regs.tgid_y_en != 0;
 	info.group_id[2]         = regs.cs_regs.tgid_z_en != 0;
@@ -1450,6 +1463,7 @@ bool ShaderCompileSpirvVS(const HW::VertexShaderInfo& regs, const HW::ShaderRegi
 	options.shader_base                = shader_addr;
 	options.user_data_base             = 8;
 	options.user_data_count            = regs.gs_regs.rsrc2.user_sgpr;
+	options.scratch_dwords             = input_info.scratch_size_dwords;
 	options.user_data                  = regs.gs_user_sgpr.value;
 	options.read_specialization_memory = ReadShaderGuestMemory;
 	options.descriptor_set             = 0;
@@ -1504,6 +1518,7 @@ bool ShaderCompileSpirvPS(const HW::PixelShaderInfo& regs, const HW::ShaderRegis
 	options.shader_hash                = shader_hash;
 	options.shader_base                = shader_addr;
 	options.user_data_count            = regs.ps_regs.rsrc2.user_sgpr;
+	options.scratch_dwords             = input_info.scratch_size_dwords;
 	options.user_data                  = regs.ps_user_sgpr.value;
 	options.read_specialization_memory = ReadShaderGuestMemory;
 	options.descriptor_set             = input_info.descriptor_set;
@@ -1554,6 +1569,7 @@ bool ShaderCompileSpirvCS(const HW::ComputeShaderInfo& regs, const HW::ShaderReg
 	options.shader_hash                = shader_addr;
 	options.shader_base                = shader_addr;
 	options.user_data_count            = regs.cs_regs.user_sgpr;
+	options.scratch_dwords             = input_info.scratch_size_dwords;
 	options.user_data                  = regs.cs_user_sgpr.value;
 	options.read_specialization_memory = ReadShaderGuestMemory;
 	options.descriptor_set             = 0;
@@ -1610,6 +1626,7 @@ ShaderId ShaderGetIdVS(const HW::VertexShaderInfo& regs, const ShaderVertexInput
 	ret.ids.push_back(static_cast<uint32_t>(input_info.fetch_embedded));
 	ret.ids.push_back(input_info.resources_num);
 	ret.ids.push_back(input_info.export_count);
+	ret.ids.push_back(input_info.scratch_size_dwords);
 
 	for (int i = 0; i < input_info.resources_num; i++) {
 		const auto& r  = input_info.resources[i];
@@ -1663,6 +1680,7 @@ ShaderId ShaderGetIdPS(const HW::PixelShaderInfo& regs, const ShaderPixelInputIn
 
 	ret.ids.push_back(input_info.descriptor_set);
 	ret.ids.push_back(input_info.push_constant_offset);
+	ret.ids.push_back(input_info.scratch_size_dwords);
 	ret.ids.push_back(input_info.input_num);
 	ret.ids.push_back(input_info.ps_system_input_base);
 	ret.ids.push_back(static_cast<uint32_t>(input_info.ps_pos_x));
@@ -1726,6 +1744,7 @@ ShaderId ShaderGetIdCS(const HW::ComputeShaderInfo& regs, const ShaderComputeInp
 	ret.ids.push_back(input_info.wave_size);
 	ret.ids.push_back(input_info.thread_ids_num);
 	ret.ids.push_back(input_info.lds_size_dwords);
+	ret.ids.push_back(input_info.scratch_size_dwords);
 
 	for (int i = 0; i < 3; i++) {
 		ret.ids.push_back(input_info.threads_num[i]);

src/graphics/shader/shader.h

@@ -95,6 +95,7 @@ struct ShaderVertexInputInfo {
 	int                     fetch_buffer_reg    = 0;
 	int                     buffers_num         = 0;
 	int                     export_count        = 0;
+	uint32_t                scratch_size_dwords = 0;
 	uint32_t                param_export_mask   = 0;
 	bool                    fetch_external      = false;
 	bool                    fetch_embedded      = false;
@@ -104,6 +105,7 @@ struct ShaderComputeInputInfo {
 	uint32_t           threads_num[3]             = {0, 0, 0};
 	uint32_t           dispatch_threads_num[3]    = {0, 0, 0};
 	uint32_t           lds_size_dwords            = 0;
+	uint32_t           scratch_size_dwords        = 0;
 	bool               group_id[3]                = {false, false, false};
 	bool               dispatch_thread_dimensions = false;
 	uint32_t           wave_size                  = 64;
@@ -122,6 +124,7 @@ struct ShaderPixelInputInfo {
 	uint32_t                                       mrt_output_mask              = 0;
 	uint32_t                                       descriptor_set               = 0;
 	uint32_t                                       push_constant_offset         = 0;
+	uint32_t                                       scratch_size_dwords          = 0;
 	bool                                           ps_pos_x                     = false;
 	bool                                           ps_pos_y                     = false;
 	bool                                           ps_pos_xy                    = false;
@@ -241,6 +244,7 @@ struct ShaderMappedData {
 	ShaderSemantic* input_semantics     = nullptr;
 	uint32_t        num_input_semantics = 0;
 	uint32_t        code_size_bytes     = 0;
+	uint32_t        scratch_size_dwords = 0;
 };
 
 void ShaderInit();

src/libs/agc.cpp

@@ -663,6 +663,7 @@ int KYTY_SYSV_ABI AgcCreateShader(Shader** dst, void* header, const volatile voi
 	map.input_semantics     = h->input_semantics;
 	map.num_input_semantics = h->num_input_semantics;
 	map.code_size_bytes     = h->shader_size;
+	map.scratch_size_dwords = h->scratch_size_dw_per_thread;
 
 	ShaderMapUserData(base, map);
 

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,
+  Check(result.program.info.addresses.size() == 1 &&
-        "segmented address resources were not tracked independently");
+            result.program.info.addresses[0].kind ==
-  for (const auto &address : result.program.info.addresses) {
+                ShaderRecompiler::IR::ResourceKind::Global &&
-    Check(
+            result.program.info.addresses[0].unbased &&
-        address.unbased && address.source == UINT32_MAX,
+            result.program.info.addresses[0].source == UINT32_MAX,
-        "GLOBAL/SCRATCH null-SADDR incorrectly inherited FLAT VGPR provenance");
+        "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() {