shader: share indirect buffer array access after an empty root

Co-authored-by: fxpw <fxpw@outlook.com>

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

@@ -400,7 +400,7 @@ uint32_t StorageBufferPackedStride(const EmitterState& state, const IR::MemoryIn
 
 Prospero::BufferFormat StorageBufferFormat(const EmitterState& state, const IR::MemoryInfo& mem);
 
-bool UsesSharedIndirectBufferAccess(const EmitterState& state, uint32_t resource);
+uint32_t SharedIndirectBufferResource(const EmitterState& state, uint32_t resource);
 
 void EmitMemoryOffsets(EmitterState& state);
 

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

@@ -309,20 +309,24 @@ auto EmitBufferAccess(ValueEmitContext& ctx, const IR::Inst& inst, IR::MemoryInf
 	    state.program.info.buffers[mem.resource].indirect_root != mem.resource) {
 		return access(mem.resource);
 	}
-	const auto& buffer = state.program.info.buffers[mem.resource];
+	const auto root_resource = mem.resource;
+	const auto& buffer = state.program.info.buffers[root_resource];
 	const auto& handle = *inst.Arg(0).ResolveInstruction();
 	const auto selected = ctx.Arg(handle, 0);
-	const auto shared = UsesSharedIndirectBufferAccess(state, mem.resource);
+	const auto representative = SharedIndirectBufferResource(state, root_resource);
+	const auto shared = representative != UINT32_MAX;
+	const auto empty_root = shared && buffer.packed_stride == 0u;
 	const auto selected_access = [&]() {
 		if (shared) {
 			const auto root_slot = buffer.descriptor_index;
 			const auto child_slot = state.program.info.buffers[buffer.indirect_resources[1]].descriptor_index;
 			auto slot = EmitAddU32(state, selected, ConstantU32(state, child_slot - 1u));
-			if (child_slot != root_slot + 1u) {
+			if (!empty_root && child_slot != root_slot + 1u) {
 				slot = Select(state, TypeU32(state),
 				    Binary(state, spv::OpIEqual, TypeBool(state), selected, ConstantU32(state, 0u)),
 				    ConstantU32(state, root_slot), slot);
 			}
+			mem.resource = representative;
 			return emit(mem, PrepareMemoryResourceAccess(state, mem, slot));
 		}
 		return EmitIndexSwitch(state, selected, buffer.indirect_resources.size(),
@@ -332,12 +336,19 @@ auto EmitBufferAccess(ValueEmitContext& ctx, const IR::Inst& inst, IR::MemoryInf
 	if constexpr (!has_result) {
 		if (!shared) return selected_access();
 	}
-	const auto valid = Binary(state, spv::OpULessThan, TypeBool(state), selected,
+	auto valid = Binary(state, spv::OpULessThan, TypeBool(state), selected,
 	                          ConstantU32(state, buffer.indirect_resources.size()));
 	if constexpr (has_result) {
 		return EmitValueOrDefaultIfCondition(state, valid, TypeU32Composite(state, components),
-		                                    ConstantU32CompositeZero(state, components), selected_access);
+		    ConstantU32CompositeZero(state, components), [&]() {
+			    if (!empty_root) return selected_access();
+			    return EmitValueOrDefaultIfCondition(state,
+			        Binary(state, spv::OpINotEqual, TypeBool(state), selected, ConstantU32(state, 0u)),
+			        TypeU32Composite(state, components), access(root_resource), selected_access);
+		    });
 	} else {
+		if (empty_root) valid = AndCondition(state, valid,
+		    Binary(state, spv::OpINotEqual, TypeBool(state), selected, ConstantU32(state, 0u)));
 		EmitIfCondition(state, valid, selected_access);
 	}
 }

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

@@ -108,24 +108,28 @@ Prospero::BufferFormat StorageBufferFormat(const EmitterState& state, const IR::
 	return state.program.info.buffers[mem.resource].descriptor_format;
 }
 
-bool UsesSharedIndirectBufferAccess(const EmitterState& state, uint32_t resource) {
+uint32_t SharedIndirectBufferResource(const EmitterState& state, uint32_t resource) {
 	const auto& buffers = state.program.info.buffers;
 	const auto& root = buffers[resource];
 	if (root.indirect_root != resource || root.descriptor_index == UINT32_MAX ||
-	    root.packed_stride == 0u || root.indirect_resources.size() < 2u) {
+	    root.indirect_resources.size() < 2u) {
-		return false;
+		return UINT32_MAX;
 	}
+	const auto first = root.packed_stride == 0u ? 1u : 0u;
+	const auto representative = root.indirect_resources[first];
+	const auto& shared = buffers[representative];
+	if (shared.packed_stride == 0u) return UINT32_MAX;
 	uint32_t next_child_slot = UINT32_MAX;
-	for (uint32_t ordinal = 0; ordinal < root.indirect_resources.size(); ++ordinal) {
+	for (uint32_t ordinal = first; ordinal < root.indirect_resources.size(); ++ordinal) {
 		const auto& candidate = buffers[root.indirect_resources[ordinal]];
-		if (candidate.packed_stride != root.packed_stride ||
+		if (candidate.packed_stride != shared.packed_stride ||
-		    candidate.descriptor_format != root.descriptor_format ||
+		    candidate.descriptor_format != shared.descriptor_format ||
-		    candidate.descriptor_swizzle != root.descriptor_swizzle ||
+		    candidate.descriptor_swizzle != shared.descriptor_swizzle ||
 		    candidate.descriptor_index == UINT32_MAX ||
-		    (ordinal > 1u && candidate.descriptor_index != next_child_slot)) return false;
+		    (ordinal > 1u && candidate.descriptor_index != next_child_slot)) return UINT32_MAX;
 		if (ordinal != 0u) next_child_slot = candidate.descriptor_index + 1u;
 	}
-	return true;
+	return representative;
 }
 
 void EmitMemoryOffsets(EmitterState& state) {
@@ -133,7 +137,7 @@ void EmitMemoryOffsets(EmitterState& state) {
 	state.memory_byte_offsets.resize(count);
 	std::vector<bool> dynamic_offsets(count);
 	for (uint32_t resource = 0; resource < state.program.info.buffers.size(); ++resource) {
-		if (!UsesSharedIndirectBufferAccess(state, resource)) continue;
+		if (SharedIndirectBufferResource(state, resource) == UINT32_MAX) continue;
 		for (const auto candidate: state.program.info.buffers[resource].indirect_resources) {
 			dynamic_offsets[state.program.info.buffers[candidate].descriptor_index] = true;
 		}

tests/ShaderRecompilerComputeTests.cpp

@@ -37012,7 +37012,7 @@ void CheckBoundedRecordBufferAccess(VulkanHarness &vulkan) {
                             .limits.minStorageBufferOffsetAlignment;
   const u32 slot = (4096u + alignment - 1u) / alignment * alignment;
   for (const u32 wave : {32u, 64u}) {
-    for (const u32 count : {0u, 2u, 6u}) {
+    for (const u32 count : {0u, 2u, 3u, 6u}) {
       TestCase test;
       test.name = wave == 32 ? "BoundedRecordBufferAccess32" : "BoundedRecordBufferAccess64";
       constexpr u32 table = 256u;
@@ -37063,9 +37063,10 @@ void CheckBoundedRecordBufferAccess(VulkanHarness &vulkan) {
         const u32 dst = slot * (row + 7u) + row + 1u;
         const auto format = count == 6u && row == 0u ? F::k32_32_32_32Float :
                             count == 6u && row == 1u ? F::k8_8_8_8UNorm : F::k16_16_16_16Float;
-        const u32 swizzle = row == 2u ? DstSel(1, 4, 1, 4) : DstSel(4, 5, 6, 7);
+        const bool empty_one = row == (count == 3u ? 0u : 2u);
-        std::array<u32, 4> input{row == 2u ? 0u : src, 16u << 16u,
+        const u32 swizzle = empty_one ? DstSel(1, 4, 1, 4) : DstSel(4, 5, 6, 7);
-                                row == 2u ? 0u : 2u, static_cast<u32>(format) << 12u | swizzle};
+        std::array<u32, 4> input{empty_one ? 0u : src, 16u << 16u,
+                                empty_one ? 0u : 2u, static_cast<u32>(format) << 12u | swizzle};
         if (row == 3u) input.fill(0u);
         const std::array<u32, 4> destination{dst, 3u << 16u, 16u, 0x5204u};
         std::copy(input.begin(), input.end(), test.initial.begin() + table / 4u + row * 48u + 40u);
@@ -37087,7 +37088,7 @@ void CheckBoundedRecordBufferAccess(VulkanHarness &vulkan) {
       for (u32 row = 0; row < count; ++row) {
         for (u32 lane = 0; lane < wave; ++lane) {
           std::array<u32, 2> value{};
-          if (row == 2u) value = {0x3c00u, 0x3c00u};
+          if (row == (count == 3u ? 0u : 2u)) value = {0x3c00u, 0x3c00u};
           else if (row != 3u && lane < 2u && !(row == 5u && lane == 1u))
             value = count == 6u && row == 1u ? std::array<u32, 2>{0x3c003c00u, 0x3c003c00u}
                                             : std::array<u32, 2>{0x38003c00u, 0xc0000000u};
@@ -37097,7 +37098,7 @@ void CheckBoundedRecordBufferAccess(VulkanHarness &vulkan) {
             expected_bytes[slot * (row + 7u) + row + 1u + lane * 3u] = std::byte{0xc3};
         }
       }
-      if (count == 2u) {
+      if (count == 2u || count == 3u) {
         test.required_spirv = {"StorageBufferArrayNonUniformIndexing"};
         test.forbidden_spirv = {"OpSwitch", "PhysicalStorageBuffer", "convert_indirect_buffer_format"};
       } else if (count == 6u) test.required_spirv = {"OpSwitch"};
@@ -37347,8 +37348,24 @@ void CheckGpuSelectedFormattedRecordStore(VulkanHarness &vulkan) {
     test.decoded_counts = {{"V_LSHLREV_B32", wide ? 1u : 2u},
         {wide ? "BUFFER_STORE_FORMAT_XYZW" : "BUFFER_STORE_FORMAT_X", 1u}};
     test.required_spirv = {"get_bda_pointer", "convert_indirect_buffer_format"};
+    if (!wide) test.required_spirv.push_back("StorageBufferArrayNonUniformIndexing");
     test.forbidden_spirv = {"OpAtomicCompareExchange"};
     const auto compiled = CompileCase(test, vulkan.SubgroupSize());
+    if (!wide) {
+      std::set<u32> nonuniform;
+      u32 shared_stores = 0;
+      for (size_t offset = 5; offset < compiled.spirv.size();) {
+        const auto words = std::span(compiled.spirv).subspan(offset, compiled.spirv[offset] >> 16u);
+        const auto opcode = static_cast<spv::Op>(words[0] & 0xffffu);
+        if (opcode == spv::OpDecorate && words[2] == spv::DecorationNonUniform)
+          nonuniform.insert(words[1]);
+        if (opcode == spv::OpStore && nonuniform.contains(words[1])) ++shared_stores;
+        offset += words.size();
+      }
+      Require(test.name, "one shared store after the null candidate",
+              shared_stores == (wave > vulkan.SubgroupSize() ? 2u : 1u),
+              "homogeneous children duplicated the native store per descriptor");
+    }
     u32 bda_loads = 0, native_stores = 0;
     for (const auto *block : compiled.program.blocks)
       for (const auto &inst : *block) {