From 71c72617406cf2eb3db9579307b8bdd7de1c5191 Mon Sep 17 00:00:00 2001 From: nmzik Date: Fri, 9 Oct 2026 15:24:01 +0200 Subject: [PATCH] shader: share indirect texture dispatch with image loads --- .../backend/spirv/spirvEmitterAnalysis.cpp | 9 +- .../backend/spirv/spirvEmitterImage.cpp | 215 +++++++++--------- .../ir/passes/ResourceMaterialization.cpp | 3 +- tests/ShaderRecompilerComputeTests.cpp | 99 +++++--- 4 files changed, 190 insertions(+), 136 deletions(-) diff --git a/src/graphics/shader/recompiler/backend/spirv/spirvEmitterAnalysis.cpp b/src/graphics/shader/recompiler/backend/spirv/spirvEmitterAnalysis.cpp index b3c96be36..776b41498 100644 --- a/src/graphics/shader/recompiler/backend/spirv/spirvEmitterAnalysis.cpp +++ b/src/graphics/shader/recompiler/backend/spirv/spirvEmitterAnalysis.cpp @@ -144,6 +144,12 @@ uint32_t LoadImageDescriptor(EmitterState& state, uint32_t resource, uint32_t mi const auto kind = *IR::DescriptorBindingForImage(state.program.info.images[resource]); state.builder.AddFunction(spv::OpLoad, state.images[IR::ImageBindingIndex(kind)].type, image, pointer); + if (array_index != 0u) { + state.builder.RequireExtension("SPV_EXT_descriptor_indexing"); + state.builder.RequireCapability(spv::CapabilityShaderNonUniform); + state.builder.RequireCapability(spv::CapabilitySampledImageArrayNonUniformIndexing); + state.builder.AddAnnotation(spv::OpDecorate, image, spv::DecorationNonUniform); + } return image; } @@ -174,9 +180,6 @@ uint32_t MakeSampledImage(EmitterState& state, uint32_t resource, uint32_t sampl } state.builder.AddFunction(spv::OpSampledImage, sampled_type, sampled_image, image, sampler_id); if (array_index != 0u) { - state.builder.RequireExtension("SPV_EXT_descriptor_indexing"); - state.builder.RequireCapability(spv::CapabilityShaderNonUniform); - state.builder.RequireCapability(spv::CapabilitySampledImageArrayNonUniformIndexing); state.builder.AddAnnotation(spv::OpDecorate, sampled_image, spv::DecorationNonUniform); } return sampled_image; diff --git a/src/graphics/shader/recompiler/backend/spirv/spirvEmitterImage.cpp b/src/graphics/shader/recompiler/backend/spirv/spirvEmitterImage.cpp index 4d1a1bb53..278995ef4 100644 --- a/src/graphics/shader/recompiler/backend/spirv/spirvEmitterImage.cpp +++ b/src/graphics/shader/recompiler/backend/spirv/spirvEmitterImage.cpp @@ -615,6 +615,90 @@ uint32_t StoreTexel(ValueEmitContext& ctx, const IR::MemoryInfo& mem, uint32_t d return PackImageTexel(ctx, mem, texel); } +template +uint32_t EmitImageAccess(ValueEmitContext& ctx, const IR::Inst& inst, uint32_t result_type, + const Emit& emit) { + auto& state = ctx.state; + const auto& mem = ctx.Memory(inst); + const auto& image = state.program.info.images.at(mem.resource); + if (image.indirect_root != mem.resource) return emit(mem.resource, 0u); + const auto* handle = inst.Arg(0).ResolveInstruction(); + const auto* source = image.source < state.program.descriptor_sources.size() + ? &state.program.descriptor_sources[image.source] + : nullptr; + if (handle == nullptr || source == nullptr || !source->indirect_descriptor.has_value() || + handle->NumArgs() == 0u) { + ctx.Fail(inst, "has invalid indirect image key provenance"); + } + const auto key = ctx.Def(handle->Arg(0)); + if (state.flattened_srt_variable == 0 || + image.indirect_resources.size() < 2u) { + ctx.Fail(inst, "has no indirect image runtime mapping"); + } + const auto selected = EmitIndirectResourceIndex( + state, key, image.indirect_mapping_offset, image.indirect_search_iterations, 0u); + struct ImageRun { + uint32_t first; + uint32_t count; + uint32_t resource; + uint32_t slot_bias; + }; + std::vector runs; + for (uint32_t ordinal = 0; ordinal < image.indirect_resources.size(); ++ordinal) { + const auto resource = image.indirect_resources[ordinal]; + const auto& candidate = state.program.info.images[resource]; + const auto kind = *IR::DescriptorBindingForImage(candidate); + if (!runs.empty()) { + auto& run = runs.back(); + const auto& first = state.program.info.images[run.resource]; + const auto next_slot = run.slot_bias + ordinal; + if (candidate.dimension == first.dimension && candidate.cube == first.cube && + IR::DescriptorBindingForImage(first) == kind && + candidate.mip_count == 1u && first.mip_count == 1u && + (ordinal == 1u || candidate.descriptor_index == next_slot)) { + // Native roots precede appended children; only the root may have a slot gap. + if (ordinal == 1u) + run.slot_bias = candidate.descriptor_index - ordinal; + ++run.count; + continue; + } + } + runs.push_back({ordinal, 1u, resource, + candidate.descriptor_index - ordinal}); + } + const auto EmitRun = [&](uint32_t index) { + const auto& run = runs[index]; + if (run.count == 1u) return emit(run.resource, 0u); + auto slot = Binary(state, spv::OpIAdd, TypeU32(state), selected, + ConstantU32(state, run.slot_bias)); + if (run.first == 0u) { + const auto root_slot = image.descriptor_index; + if (root_slot != run.slot_bias) { + const auto is_root = Binary(state, spv::OpIEqual, TypeBool(state), selected, + ConstantU32(state, 0u)); + const auto root_or_child = state.builder.AllocateId(); + state.builder.AddFunction(spv::OpSelect, TypeU32(state), root_or_child, + is_root, ConstantU32(state, root_slot), slot); + slot = root_or_child; + } + } + return emit(run.resource, slot); + }; + auto run_index = ConstantU32(state, 0u); + for (uint32_t index = 1; index < runs.size(); ++index) { + const auto at_or_after = Binary(state, spv::OpUGreaterThanEqual, TypeBool(state), + selected, ConstantU32(state, runs[index].first)); + const auto next = state.builder.AllocateId(); + state.builder.AddFunction(spv::OpSelect, TypeU32(state), next, at_or_after, + ConstantU32(state, index), run_index); + run_index = next; + } + return runs.size() == 1u + ? EmitRun(0u) + : EmitIndexSwitch(state, run_index, static_cast(runs.size()), + result_type, EmitRun); +} + } // namespace void EmitImage(ValueEmitContext& ctx, const IR::Inst& inst) { @@ -656,29 +740,35 @@ void EmitImage(ValueEmitContext& ctx, const IR::Inst& inst) { return; } if (op == IR::ValueOpcode::ImageRead) { - const auto dimension = image.dimension; - const auto& dimension_info = ImageDimensionInfoFor(dimension); - const auto numeric_class = image.numeric_class; - const auto condition = ctx.Arg(inst, 2); + const auto numeric_class = image.numeric_class; + const auto condition = ctx.Arg(inst, 2); ctx.Define( inst, EmitValueOrDefaultIfCondition( state, condition, TypeU32Vector(state, 4), ConstantU32CompositeZero(state, 4), [&]() { - const auto descriptor = LoadImageDescriptor(state, mem.resource); - const auto color = state.builder.AllocateId(); - const auto coord = CoordU32(ctx, mem, *address, dimension); - if (dimension_info.multisampled != 0u) { - state.builder.AddFunction( - spv::OpImageFetch, ImageVectorType(state, numeric_class, 4), color, - descriptor, coord, spv::ImageOperandsSampleMask, - AddressU32(ctx, mem, *address, dimension_info.coordinate_components)); - } else { - state.builder.AddFunction(spv::OpImageFetch, - ImageVectorType(state, numeric_class, 4), color, - descriptor, coord, spv::ImageOperandsLodMask, - LodU32(ctx, mem, *address, dimension)); - } + const auto color = EmitImageAccess( + ctx, inst, ImageVectorType(state, numeric_class, 4), + [&](uint32_t resource, uint32_t array_index) { + const auto& candidate = state.program.info.images[resource]; + const auto dimension = candidate.dimension; + const auto& dimension_info = ImageDimensionInfoFor(dimension); + const auto descriptor = LoadImageDescriptor(state, resource, 0u, array_index); + const auto fetched = state.builder.AllocateId(); + const auto coord = CoordU32(ctx, mem, *address, dimension); + if (dimension_info.multisampled != 0u) { + state.builder.AddFunction( + spv::OpImageFetch, ImageVectorType(state, candidate.numeric_class, 4), fetched, + descriptor, coord, spv::ImageOperandsSampleMask, + AddressU32(ctx, mem, *address, dimension_info.coordinate_components)); + } else { + state.builder.AddFunction( + spv::OpImageFetch, ImageVectorType(state, candidate.numeric_class, 4), fetched, + descriptor, coord, spv::ImageOperandsLodMask, + LodU32(ctx, mem, *address, mem.image_dimension)); + } + return fetched; + }); return ResultVector(ctx, UnpackImageTexel(ctx, mem, color), numeric_class, false, mem); })); @@ -809,7 +899,7 @@ void EmitImage(ValueEmitContext& ctx, const IR::Inst& inst) { operand_count = 2; } const auto sampler_id = LoadSamplerDescriptor(state, mem.sampler); - const auto EmitSample = [&](uint32_t resource, uint32_t array_index = 0u) { + const auto EmitSample = [&](uint32_t resource, uint32_t array_index) { const auto& candidate = state.program.info.images[resource]; const auto coord = CoordF32(ctx, mem, *address, layout.coord, @@ -822,92 +912,7 @@ void EmitImage(ValueEmitContext& ctx, const IR::Inst& inst) { std::span(operands).first(operand_count)); return sample; }; - if (image.indirect_root != mem.resource) { - const auto sample = EmitSample(mem.resource); - auto result = sample; - if (!dref) { - result = UnpackImageTexel(ctx, mem, sample); - } - ctx.Define(inst, ResultVector(ctx, result, numeric_class, dref, mem)); - return; - } - const auto* handle = image_arg.ResolveInstruction(); - const auto* source = image.source < state.program.descriptor_sources.size() - ? &state.program.descriptor_sources[image.source] - : nullptr; - if (handle == nullptr || source == nullptr || !source->indirect_descriptor.has_value() || - handle->NumArgs() == 0u) { - ctx.Fail(inst, "has invalid indirect image key provenance"); - return; - } - const auto key = ctx.Def(handle->Arg(0)); - if (state.flattened_srt_variable == 0 || - image.indirect_resources.size() < 2u) { - ctx.Fail(inst, "has no indirect image runtime mapping"); - return; - } - const auto selected = EmitIndirectResourceIndex( - state, key, image.indirect_mapping_offset, image.indirect_search_iterations, 0u); - struct SampleRun { - uint32_t first; - uint32_t count; - uint32_t resource; - uint32_t slot_bias; - }; - std::vector runs; - for (uint32_t ordinal = 0; ordinal < image.indirect_resources.size(); ++ordinal) { - const auto resource = image.indirect_resources[ordinal]; - const auto& candidate = state.program.info.images[resource]; - const auto kind = *IR::DescriptorBindingForImage(candidate); - if (!runs.empty()) { - auto& run = runs.back(); - const auto& first = state.program.info.images[run.resource]; - const auto next_slot = run.slot_bias + ordinal; - if (candidate.dimension == first.dimension && candidate.cube == first.cube && - IR::DescriptorBindingForImage(first) == kind && - candidate.mip_count == 1u && first.mip_count == 1u && - (ordinal == 1u || candidate.descriptor_index == next_slot)) { - // Native roots precede appended children; only the root may have a slot gap. - if (ordinal == 1u) - run.slot_bias = candidate.descriptor_index - ordinal; - ++run.count; - continue; - } - } - runs.push_back({ordinal, 1u, resource, - candidate.descriptor_index - ordinal}); - } - const auto EmitRun = [&](uint32_t index) { - const auto& run = runs[index]; - if (run.count == 1u) return EmitSample(run.resource); - auto slot = Binary(state, spv::OpIAdd, TypeU32(state), selected, - ConstantU32(state, run.slot_bias)); - if (run.first == 0u) { - const auto root_slot = image.descriptor_index; - if (root_slot != run.slot_bias) { - const auto is_root = Binary(state, spv::OpIEqual, TypeBool(state), selected, - ConstantU32(state, 0u)); - const auto root_or_child = state.builder.AllocateId(); - state.builder.AddFunction(spv::OpSelect, TypeU32(state), root_or_child, - is_root, ConstantU32(state, root_slot), slot); - slot = root_or_child; - } - } - return EmitSample(run.resource, slot); - }; - auto run_index = ConstantU32(state, 0u); - for (uint32_t index = 1; index < runs.size(); ++index) { - const auto at_or_after = Binary(state, spv::OpUGreaterThanEqual, TypeBool(state), - selected, ConstantU32(state, runs[index].first)); - const auto next = state.builder.AllocateId(); - state.builder.AddFunction(spv::OpSelect, TypeU32(state), next, at_or_after, - ConstantU32(state, index), run_index); - run_index = next; - } - auto result = runs.size() == 1u - ? EmitRun(0u) - : EmitIndexSwitch(state, run_index, static_cast(runs.size()), - result_type, EmitRun); + auto result = EmitImageAccess(ctx, inst, result_type, EmitSample); if (!dref) { result = UnpackImageTexel(ctx, mem, result); } diff --git a/src/graphics/shader/recompiler/ir/passes/ResourceMaterialization.cpp b/src/graphics/shader/recompiler/ir/passes/ResourceMaterialization.cpp index e34c214c9..8a541959a 100644 --- a/src/graphics/shader/recompiler/ir/passes/ResourceMaterialization.cpp +++ b/src/graphics/shader/recompiler/ir/passes/ResourceMaterialization.cpp @@ -1434,7 +1434,8 @@ void ApplyResourceSpecialization(Program& program, const ResourceSpecialization& EXIT_IF(memory.sampler == UINT32_MAX); } EXIT_IF(image.indirect_root == memory.resource && - inst.GetOpcode() != ValueOpcode::ImageSampleRaw); + inst.GetOpcode() != ValueOpcode::ImageSampleRaw && + inst.GetOpcode() != ValueOpcode::ImageRead); } } for (auto& memory: memory_info) { diff --git a/tests/ShaderRecompilerComputeTests.cpp b/tests/ShaderRecompilerComputeTests.cpp index 74a45062b..aaa66741e 100644 --- a/tests/ShaderRecompilerComputeTests.cpp +++ b/tests/ShaderRecompilerComputeTests.cpp @@ -35085,42 +35085,87 @@ void CheckIndirectImageKeySwitch(VulkanHarness &vulkan) { std::vector textures; for (u32 resource = 0; resource < 7u; ++resource) { const auto layers = resource == 4u ? 6u : 1u; - std::vector pixels(layers * 4u, std::bit_cast(float(resource))); + std::vector pixels(layers * 16u, std::bit_cast(float(resource))); + std::vector mip(layers * 4u, + std::bit_cast(resource == 0u ? 0.0f : float(resource + 10u))); textures.push_back(vulkan.CreateImageMips( - test.name, 1, 1, vk::Format::eR32G32B32A32Sfloat, - vk::ImageUsageFlagBits::eSampled, {pixels}, 4u, + test.name, 2, 2, vk::Format::eR32G32B32A32Sfloat, + vk::ImageUsageFlagBits::eSampled, {pixels, mip}, 4u, vk::ImageLayout::eShaderReadOnlyOptimal, vk::ImageType::e2D, resource == 4u ? vk::ImageViewType::e2DArray : vk::ImageViewType::e2D, layers)); } const auto native_sampler = vulkan.CreateSampler(test.name); - for (const u32 wave_size : {32u, 64u}) { - compute.wave_size = wave_size; - compute.host_subgroup_size = vulkan.SubgroupSize(); - compute.threads_num[0] = wave_size; - compute.threads_num[1] = compute.threads_num[2] = 1u; - compiled.program.wave_size = wave_size; - compiled.spirv = ShaderRecompiler::Spirv::EmitProgram(compiled.program, {.compute = &compute}); - compiled.packed_user_data.resize(compiled.program.bindings.ShaderDataDwords()); - ValidateSpirv(test.name, compiled.spirv); - const auto halves = wave_size > compute.host_subgroup_size ? 2u : 1u; - Require(test.name, "native array run samples", tools.Disassemble(compiled.spirv, &text) && - CountText(text, "OpImageSampleExplicitLod") == 3u * halves, - "mixed runs expanded into per-image samples"); - test.initial.assign(wave_size, 0xdeadbeefu); - test.expected.assign(wave_size, 0u); - for (u32 lane = 0; lane < ordinals.size(); ++lane) { - test.expected[lane] = std::bit_cast(float(root.indirect_resources[ordinals[lane]])); + for (const bool fetch : {false, true}) { + if (fetch) { + test.name = "IndirectImageFetchContiguousRuns"; + const auto position = std::ranges::find_if(*block, [&](const Inst &inst) { + return &inst == &sample; + }); + auto &read = *block->PrependNewInst( + position, ValueOpcode::ImageRead, + {Value(&image), Value(&address), Value(true)}, memory_flag_bits); + sample.ReplaceUsesWith(Value(&read), false); + sampler.Invalidate(); + compiled.program.info.samplers.clear(); + compiled.program.info.sampled_pairs.clear(); + compiled.program.info.images[4].cube = false; + compiled.program.memory_info[0].image_sample_flags = 0; + compiled.program.memory_info[0].image_has_mip = true; + for (u32 component = 0; component < 4u; ++component) + address.SetArg(component, Value(0u)); + // Native DIM=2DArray places LOD after the layer, even for a 2D candidate. + address.SetArg(3, Value(1u)); + compiled.program.shader_info_complete = false; + CollectShaderInfo(compiled.program, {.compute = &compute}); } - auto buffer = vulkan.CreateStorageBuffer(test.name, test.initial, test.initial.size()); - vulkan.Dispatch(test, compiled, buffer, nullptr, nullptr, nullptr, nullptr, - native_sampler, textures); - const auto actual = vulkan.ReadBuffer(test.name, buffer, test.expected.size()); - vulkan.DestroyBuffer(&buffer); - CompareWords(test, "root, native slots, cube and mapping bounds", test.expected, actual); + for (const u32 wave_size : {32u, 64u}) { + compute.wave_size = wave_size; + compute.host_subgroup_size = vulkan.SubgroupSize(); + compute.threads_num[0] = wave_size; + compute.threads_num[1] = compute.threads_num[2] = 1u; + compiled.program.wave_size = wave_size; + compiled.spirv = ShaderRecompiler::Spirv::EmitProgram(compiled.program, {.compute = &compute}); + compiled.packed_user_data.resize(compiled.program.bindings.ShaderDataDwords()); + ValidateSpirv(test.name, compiled.spirv); + const auto halves = wave_size > compute.host_subgroup_size ? 2u : 1u; + Require(test.name, "native array image runs", tools.Disassemble(compiled.spirv, &text) && + CountText(text, fetch ? "OpImageFetch" : "OpImageSampleExplicitLod") == 3u * halves, + "mixed runs expanded into per-image accesses"); + if (fetch) { + std::vector decorated; + std::vector fetched_images; + 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(words[0] & 0xffffu); + if (opcode == spv::OpDecorate && words[2] == spv::DecorationNonUniform) + decorated.push_back(words[1]); + if (opcode == spv::OpImageFetch) fetched_images.push_back(words[3]); + offset += words.size(); + } + Require(test.name, "nonuniform fetch operands", + std::ranges::count_if(fetched_images, [&](u32 operand) { + return std::ranges::find(decorated, operand) != decorated.end(); + }) == 2u * halves, + "a dynamically selected fetch image lacks its nonuniform decoration"); + } + test.initial.assign(wave_size, 0xdeadbeefu); + test.expected.assign(wave_size, 0u); + for (u32 lane = 0; lane < ordinals.size(); ++lane) { + const auto resource = root.indirect_resources[ordinals[lane]]; + test.expected[lane] = std::bit_cast(float(resource + (fetch && resource != 0u ? 10u : 0u))); + } + auto buffer = vulkan.CreateStorageBuffer(test.name, test.initial, test.initial.size()); + vulkan.Dispatch(test, compiled, buffer, nullptr, nullptr, nullptr, nullptr, + fetch ? vk::Sampler{} : native_sampler, textures); + const auto actual = vulkan.ReadBuffer(test.name, buffer, test.expected.size()); + vulkan.DestroyBuffer(&buffer); + CompareWords(test, "root, native slots, cube and mapping bounds", test.expected, actual); + } + std::printf("[compute] %-32s ok\n", test.name); } vulkan.Device().destroySampler(native_sampler); for (auto &texture : textures) vulkan.DestroyImage(&texture); - std::printf("[compute] %-32s ok\n", test.name); } TestCase ImageStoreMipSelectsPpsa01340Descriptor() {