shader: reuse aligned indirect buffer loads for scalar DWORD reads

Co-authored-by: TheCruZ <thecruzwtf@gmail.com>
Co-authored-by: tototomate123 <rushwars201966@gmail.com>

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

@@ -664,6 +664,7 @@ uint32_t LoadFormattedInBounds(ValueEmitContext& ctx, const IR::MemoryInfo& mem,
 
 uint32_t ConstructU32Composite(EmitterState& state, uint32_t components,
                                const std::array<uint32_t, 4>& values) {
+	if (components == 1u) return values[0];
 	const auto            result = state.builder.AllocateId();
 	std::vector<uint32_t> words {spv::OpCompositeConstruct, TypeU32Composite(state, components),
 	                             result};
@@ -679,7 +680,7 @@ uint32_t FormattedOutOfBoundsValue(ValueEmitContext& ctx, const IR::MemoryInfo&
 		const auto source = ResolveFormattedSource(ctx, mem, plan.info, component);
 		values[component] = FormattedConstant(ctx, plan.info, source.kind);
 	}
-	return components == 1u ? values[0] : ConstructU32Composite(ctx.state, components, values);
+	return ConstructU32Composite(ctx.state, components, values);
 }
 
 uint32_t FormattedLoad(ValueEmitContext& ctx, const IR::Inst& inst, const IR::MemoryInfo& mem) {
@@ -902,15 +903,17 @@ uint32_t LoadIndirectBuffer(ValueEmitContext& ctx, const IR::Inst& inst, uint32_
 	const auto buffer = PrepareIndirectBuffer(ctx, inst);
 	const auto valid_format = Binary(state, spv::OpINotEqual, TypeBool(state), buffer.format,
 	                                  ConstantU32(state, 0));
+	const auto base = Binary(state, spv::OpBitwiseAnd, TypeScalarU64(state), buffer.address,
+	                         ConstantDeviceAddress(state, ~uint64_t {3}));
 	std::array<uint32_t, 4> values {};
 	for (uint32_t component = 0; component < components; ++component) {
-		const auto address = component == 0u ? buffer.address : Binary(
+		const auto address = component == 0u ? base : Binary(
-		    state, spv::OpIAdd, TypeScalarU64(state), buffer.address,
+		    state, spv::OpIAdd, TypeScalarU64(state), base,
 		    ConstantDeviceAddress(state, component * 4u));
-		values[component] = LoadBda(
+		values[component] = EmitValueOrZeroIfCondition(
-		    ctx, address, AndCondition(state, valid_format,
+		    state, AndCondition(state, valid_format,
 		                        IndirectBufferInBounds(state, buffer, component * 4u, 4u, false)),
-		    32u);
+		    [&]() { return LoadBdaDword(ctx, address); });
 	}
 	return ConstructU32Composite(state, components, values);
 }
@@ -960,7 +963,7 @@ uint32_t LoadIndirectFormattedX(ValueEmitContext& ctx, const IR::Inst& inst) {
 	});
 }
 
-uint32_t LoadWideBuffer(ValueEmitContext& ctx, const IR::Inst& inst, uint32_t components) {
+uint32_t LoadBuffer(ValueEmitContext& ctx, const IR::Inst& inst, uint32_t components) {
 	auto& state = ctx.state;
 	return EmitValueOrDefaultIfCondition(
 	    state, ctx.Arg(inst, inst.NumArgs() - 1), TypeU32Composite(state, components),
@@ -1358,8 +1361,9 @@ void EmitLoadMemory(ValueEmitContext& ctx, const IR::Inst& inst) {
 	uint32_t   value;
 	if (mem.kind == IR::ResourceKind::FlatLocal)
 		value = LoadLocalFlat(ctx, inst);
-	else if (buffer_components > 1u)
+	else if (buffer_components > 1u ||
-		value = LoadWideBuffer(ctx, inst, buffer_components);
+	         (buffer_components == 1u && mem.kind == IR::ResourceKind::IndirectBuffer && !mem.formatted))
+		value = LoadBuffer(ctx, inst, buffer_components);
 	else if (shared_components > 1u)
 		value = LoadWideShared(ctx, inst, shared_components);
 	else if (mem.kind == IR::ResourceKind::ScalarAddress)

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

@@ -91,7 +91,8 @@ uint32_t TypeU32Vector(EmitterState& state, uint32_t components) {
 }
 
 uint32_t TypeU32Composite(EmitterState& state, uint32_t components) {
-	EXIT_IF(components < 2u || components > 4u);
+	EXIT_IF(components < 1u || components > 4u);
+	if (components == 1u) return TypeU32(state);
 	return components == 2u ? TypeU32Pair(state) : TypeU32Vector(state, components);
 }
 
@@ -333,8 +334,9 @@ uint32_t ConstantU64(EmitterState& state, uint64_t value) {
 }
 
 uint32_t ConstantU32CompositeZero(EmitterState& state, uint32_t components) {
-	EXIT_IF(components < 2u || components > 4u);
+	EXIT_IF(components < 1u || components > 4u);
 	const auto            zero = ConstantU32(state, 0);
+	if (components == 1u) return zero;
 	std::vector<uint32_t> values(components, zero);
 	return state.builder.Constant(spv::OpConstantComposite, TypeU32Composite(state, components),
 	                              values);

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

@@ -458,7 +458,7 @@ void ValidateProgram(const Program& program, bool require_ssa) {
 				}
 				if (memory.kind == ResourceKind::IndirectBuffer &&
 				    !memory.SupportsIndirectBufferLoad(inst.GetOpcode())) {
-					return Fail("indirect buffer requires a scalar, raw DWORD x2/x3/x4, or formatted X load");
+					return Fail("indirect buffer requires a scalar, raw DWORD x1/x2/x3/x4, or formatted X load");
 				}
 				if (buffer_components > 1u &&
 				    (!vector_buffer || memory.data_bits != 32u ||

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

@@ -76,6 +76,7 @@ struct MemoryInfo {
 		return !typed && data_bits == 32u &&
 		       (formatted ? opcode == ValueOpcode::LoadBufferU32
 		                  : opcode == ValueOpcode::ReadConstBuffer ||
+		                        opcode == ValueOpcode::LoadBufferU32 ||
 		                        opcode == ValueOpcode::LoadBufferU32x2 ||
 		                        opcode == ValueOpcode::LoadBufferU32x3 ||
 		                        opcode == ValueOpcode::LoadBufferU32x4);

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

@@ -2113,7 +2113,7 @@ private:
 				    !memory.SupportsIndirectBufferLoad(op)) {
 					Fail(flags.pc,
 					     "buffer descriptor is not a valid runtime value; GPU-selected access "
-					     "requires a scalar, raw DWORD x2/x3/x4, or formatted X load");
+					     "requires a scalar, raw DWORD x1/x2/x3/x4, or formatted X load");
 				}
 				m_program.memory_info[flags.index].kind = ResourceKind::IndirectBuffer;
 				m_info.uses_dma                         = true;

tests/ShaderRecompilerComputeTests.cpp

@@ -26229,7 +26229,9 @@ TestCase BufferLoadDwordx4ZeroesOnlyOutOfBoundsTail() {
   return test;
 }
 
-TestCase BufferLoadsGpuSelectedDescriptors(bool xyz) {
+TestCase BufferLoadsGpuSelectedDescriptors(u32 width) {
+  const bool scalar = width == 1u;
+  const bool xyz = width == 3u;
   using O = ShaderOpcode;
   constexpr uint64_t GuestBase = 0x0000000110000000ull;
   struct DescriptorCase {
@@ -26253,8 +26255,11 @@ TestCase BufferLoadsGpuSelectedDescriptors(bool xyz) {
       {12, 8, 3, 4, true, true, {3, 4, 5, 0, 0, 0}},
       {12, 4, 3, 4, true, true, {}},
   };
+  // OFFEN byte offset 5 aligns to DWORD; swizzled offsets use eight-index groups.
+  constexpr u32 scalar_expected[] = {2, 2, 2, 2, 2, 0, 3, 0, 10, 10, 0};
   TestCase test;
-  test.name = xyz ? "BufferLoadDwordx3GpuSelectedDescriptors"
+  test.name = scalar ? "BufferLoadDwordGpuSelectedDescriptors"
+              : xyz ? "BufferLoadDwordx3GpuSelectedDescriptors"
                     : "BufferLoadsGpuSelectedDescriptors";
   test.initial.resize(2048);
   for (u32 i = 0; i < std::size(cases); ++i) {
@@ -26280,14 +26285,17 @@ TestCase BufferLoadsGpuSelectedDescriptors(bool xyz) {
     test.code.push_back(EncodeSmem0(0x0a, 8, 0));
     test.code.push_back(EncodeSmem1(520, 20));
     AppendSMovLiteral(&test.code, 22, cases[selected].soffset);
-    AppendVMovU32(&test.code, 21, 1);
+    AppendVMovU32(&test.code, 21, scalar ? 5u : 1u);
-    test.code.push_back(EncodeMubuf0(xyz ? 0x0f : 0x0e, 0, true, false));
+    test.code.push_back(EncodeMubuf0(scalar ? 0x0c : xyz ? 0x0f : 0x0e, 0, !scalar, scalar));
     test.code.push_back(EncodeMubuf1(0, 2, 21, 22));
+    if (!scalar) {
       test.code.push_back(EncodeMubuf0(xyz ? 0x0f : 0x0d, xyz ? 12 : 16, true, false));
       test.code.push_back(EncodeMubuf1(xyz ? 3 : 4, 2, 21, 22));
-    for (u32 component = 0; component < 6; ++component) {
+    }
-      AppendStoreVgpr(&test.code, component, i * 6 + component);
+    const u32 outputs = scalar ? 1u : 6u;
-      const u32 expected = cases[selected].expected[component];
+    for (u32 component = 0; component < outputs; ++component) {
+      AppendStoreVgpr(&test.code, component, i * outputs + component);
+      const u32 expected = scalar ? scalar_expected[selected] : cases[selected].expected[component];
       test.expected.push_back(expected == 0 ? 0 : selected * 100 + expected);
     }
   }
@@ -26296,7 +26304,9 @@ TestCase BufferLoadsGpuSelectedDescriptors(bool xyz) {
   test.opcodes = {O::V_MOV_B32, O::S_MOV_B32, O::BUFFER_LOAD_DWORD,
                   O::V_READFIRSTLANE_B32, O::S_MUL_I32, O::S_BUFFER_LOAD_DWORDX4,
                   O::BUFFER_STORE_DWORD, O::S_ENDPGM};
-  if (xyz) {
+  if (scalar) {
+    test.decoded_counts = {{"BUFFER_LOAD_DWORD ", 2u * std::size(cases)}};
+  } else if (xyz) {
     test.opcodes.push_back(O::BUFFER_LOAD_DWORDX3);
   } else {
     test.opcodes.insert(test.opcodes.end(), {O::BUFFER_LOAD_DWORDX4, O::BUFFER_LOAD_DWORDX2});
@@ -26306,11 +26316,15 @@ TestCase BufferLoadsGpuSelectedDescriptors(bool xyz) {
 }
 
 TestCase BufferLoadsGpuSelectedDescriptors() {
-  return BufferLoadsGpuSelectedDescriptors(false);
+  return BufferLoadsGpuSelectedDescriptors(4u);
 }
 
 TestCase BufferLoadDwordx3GpuSelectedDescriptors() {
-  return BufferLoadsGpuSelectedDescriptors(true);
+  return BufferLoadsGpuSelectedDescriptors(3u);
+}
+
+TestCase BufferLoadDwordGpuSelectedDescriptors() {
+  return BufferLoadsGpuSelectedDescriptors(1u);
 }
 
 TestCase BufferLoadFormatXGpuSelectedDescriptors() {
@@ -32952,6 +32966,7 @@ std::vector<TestCase> MakeCases() {
   AddCase(BufferLoadDwordx4SnapshotsOverlappingAddress);
   AddCase(BufferLoadDwordx4ZeroesOnlyOutOfBoundsTail);
   AddCase(BufferLoadsGpuSelectedDescriptors);
+  AddCase(BufferLoadDwordGpuSelectedDescriptors);
   AddCase(BufferLoadDwordx3GpuSelectedDescriptors);
   AddCase(BufferLoadFormatXGpuSelectedDescriptors);
   AddCase(BufferStoreDwordx4DropsOnlyOutOfBoundsTail);
@@ -38198,6 +38213,7 @@ int main(int argc, char **argv) {
     RunCase(&vulkan, ScalarBufferFromLoopReadlane(64));
     CheckIndirectBufferStore(vulkan);
     RunCase(&vulkan, BufferLoadsGpuSelectedDescriptors());
+    RunCase(&vulkan, BufferLoadDwordGpuSelectedDescriptors());
     RunCase(&vulkan, BufferLoadDwordx3GpuSelectedDescriptors());
     RunCase(&vulkan, BufferLoadFormatXGpuSelectedDescriptors());
     RunCase(&vulkan, BufferLoadFormatXRejectsPartialRecord());