From 09c993dfe042da0f79a207879505fcce49cb33fb Mon Sep 17 00:00:00 2001 From: nmzik Date: Sun, 18 Oct 2026 21:35:02 +0200 Subject: [PATCH] shader: query dimensions through shared indirect image selection --- .../backend/spirv/spirvEmitterImage.cpp | 11 ++- .../ir/passes/ResourceMaterialization.cpp | 1 + tests/ResourceTrackingTests.cpp | 16 ++-- tests/ShaderRecompilerComputeTests.cpp | 81 ++++++++++++++++++- 4 files changed, 99 insertions(+), 10 deletions(-) diff --git a/src/graphics/shader/recompiler/backend/spirv/spirvEmitterImage.cpp b/src/graphics/shader/recompiler/backend/spirv/spirvEmitterImage.cpp index 2722fbba5..366b7893f 100644 --- a/src/graphics/shader/recompiler/backend/spirv/spirvEmitterImage.cpp +++ b/src/graphics/shader/recompiler/backend/spirv/spirvEmitterImage.cpp @@ -288,10 +288,10 @@ uint32_t ResultVector(ValueEmitContext& ctx, uint32_t value, } uint32_t QueryDimensions(ValueEmitContext& ctx, const IR::MemoryInfo& mem, - const IR::Inst& address) { - const auto dimension = ctx.state.program.info.images.at(mem.resource).dimension; + const IR::Inst& address, uint32_t resource, uint32_t array_index) { + const auto dimension = ctx.state.program.info.images.at(resource).dimension; const auto& info = ImageDimensionInfoFor(dimension); - const auto image = LoadImageDescriptor(ctx.state, mem.resource); + const auto image = LoadImageDescriptor(ctx.state, resource, 0u, array_index); const auto size = ctx.state.builder.AllocateId(); if (info.multisampled != 0u) { ctx.state.builder.AddFunction(spv::OpImageQuerySize, @@ -846,7 +846,10 @@ void EmitImage(ValueEmitContext& ctx, const IR::Inst& inst) { const auto* address = ctx.ImageAddress(inst.Arg(image_info.needs_sampler ? 2 : 1)); if (op == IR::ValueOpcode::ImageQueryDimensions) { state.builder.RequireCapability(spv::CapabilityImageQuery); - ctx.Define(inst, QueryDimensions(ctx, mem, *address)); + ctx.Define(inst, EmitImageAccess(ctx, inst, TypeU32Vector(state, 4), + [&](uint32_t resource, uint32_t array_index) { + return QueryDimensions(ctx, mem, *address, resource, array_index); + })); return; } if (image_info.needs_sampler && diff --git a/src/graphics/shader/recompiler/ir/passes/ResourceMaterialization.cpp b/src/graphics/shader/recompiler/ir/passes/ResourceMaterialization.cpp index b9a3530eb..77be9d668 100644 --- a/src/graphics/shader/recompiler/ir/passes/ResourceMaterialization.cpp +++ b/src/graphics/shader/recompiler/ir/passes/ResourceMaterialization.cpp @@ -1972,6 +1972,7 @@ void ApplyResourceSpecialization(Program& program, const ResourceSpecialization& EXIT_IF(image.indirect_root == memory.resource && inst.GetOpcode() != ValueOpcode::ImageSampleRaw && inst.GetOpcode() != ValueOpcode::ImageGatherRaw && + inst.GetOpcode() != ValueOpcode::ImageQueryDimensions && inst.GetOpcode() != ValueOpcode::ImageRead && inst.GetOpcode() != ValueOpcode::ImageWrite); } diff --git a/tests/ResourceTrackingTests.cpp b/tests/ResourceTrackingTests.cpp index 9077632cd..a7ddaf65e 100644 --- a/tests/ResourceTrackingTests.cpp +++ b/tests/ResourceTrackingTests.cpp @@ -3167,7 +3167,7 @@ void TestUniformizedMaterialImageKeys() { void TestIndexedMaterialImageKeys() { namespace CFG = Libs::Graphics::ShaderRecompiler::CFG; - const auto make_plan = [](ShaderType stage = ShaderType::Compute, bool shared = false) { + const auto make_program = [](ShaderType stage = ShaderType::Compute, bool shared = false) { Fixture f(stage); auto *entry = f.block, *body = f.AddBlock(); auto *later = shared ? f.AddBlock() : nullptr; @@ -3251,9 +3251,10 @@ void TestIndexedMaterialImageKeys() { offset == 12994u * 32u, "inherited indexed selector changed GPU byte offsets into raw material keys"); } - return ExtractResourcePlan(f.program); + return std::move(f.program); }; - auto plan = make_plan(); + auto program = make_program(); + auto plan = ExtractResourcePlan(program); Check(plan.info.images.size() == 3u && plan.selector_keys.size() == 2u, "indexed material fields did not share their query/gather selector domain"); for (const auto &image : plan.info.images) { @@ -3304,6 +3305,9 @@ void TestIndexedMaterialImageKeys() { Check(MaterializeResources(plan, runtime, snapshot, specialization) && memory.selectors == 2u && memory.descriptors == 6u && snapshot.images.size() == 9u, "indexed selector scanned the broad heap or did not share its captured keys"); + ApplyResourceSpecialization(program, specialization); + Check(program.info.images[0].indirect_resources.size() == 3u, + "indirect dimension query lost its specialized candidates"); for (uint32_t root = 0; root < plan.info.images.size(); ++root) { const auto &image = specialization.images[root]; const auto offset = image.indirect_mapping_offset; @@ -3319,7 +3323,8 @@ void TestIndexedMaterialImageKeys() { Check(std::ranges::find(snapshot.specialization_reads, std::pair{base + field, 868u}) != snapshot.specialization_reads.end(), "indexed selector batch escaped writable-alias tracking"); - auto shared = make_plan(ShaderType::Compute, true); + auto shared_program = make_program(ShaderType::Compute, true); + auto shared = ExtractResourcePlan(shared_program); Check(MaterializeResources(shared, runtime, snapshot, specialization) && snapshot.images.size() == 9u, "inherited indexed selector lost its bounded material domain"); memory.words[64u / 4u] = 12996u; @@ -3349,7 +3354,8 @@ void TestIndexedMaterialImageKeys() { "unsupported geometry or writable selector lost the original heap domain"); } user_data = original; - auto graphics = make_plan(ShaderType::Pixel); + auto graphics_program = make_program(ShaderType::Pixel); + auto graphics = ExtractResourcePlan(graphics_program); Check(MaterializeResources(graphics, runtime, snapshot, specialization) && memory.selectors == 0u, "graphics material lookup acquired unproved selector read dependencies"); memory.unreadable = false; diff --git a/tests/ShaderRecompilerComputeTests.cpp b/tests/ShaderRecompilerComputeTests.cpp index 3b2fa77af..a56a63285 100644 --- a/tests/ShaderRecompilerComputeTests.cpp +++ b/tests/ShaderRecompilerComputeTests.cpp @@ -40583,13 +40583,92 @@ void CheckIndirectImageKeySwitch(VulkanHarness &vulkan) { } std::printf("[compute] %-32s ok\n", test.name); } + test.name = "IndirectImageQueryDimensions"; + const auto address_position = std::ranges::find_if(*block, [&](const Inst &inst) { return &inst == &address; }); + auto &query_lod = *block->PrependNewInst(address_position, ValueOpcode::BitwiseAnd32, + {Value(&key), Value(1u)}); + address.SetArg(0, Value(&query_lod)); + const auto access_position = std::ranges::find_if(*block, [&](const Inst &inst) { return &inst == access; }); + auto &query = *block->PrependNewInst(access_position, ValueOpcode::ImageQueryDimensions, + {Value(&image), Value(&address)}, memory_flag_bits); + access->ReplaceUsesWith(Value(&query), false); + store.Invalidate(); + auto &query_offset = block->AppendNewInst(ValueOpcode::IMul32, {Value(&key), Value(16u)}); + auto &query_store = block->AppendNewInst(ValueOpcode::StoreBufferU32x4, + {Value(&output), Value(0u), Value(&query_offset), Value(0u), Value(&query), Value(true)}); + query_store.SetFlags(MemoryFlags{1u, 0u}); + compiled.program.memory_info[1].data_dwords = 4u; + compiled.program.memory_info[0].image_address_components = 1u; + compiled.program.memory_info[0].image_has_mip = false; + compiled.program.info.images[3].numeric_class = Prospero::TextureNumericClass::Uint; + compiled.program.info.images[4].dimension = ShaderRecompiler::Decoder::ImageDimension::Dim2DArray; + constexpr std::array, 7> geometry{{ + {2, 2, 0, 2}, {2, 2, 0, 2}, {4, 2, 0, 3}, {8, 0, 0, 4}, + {2, 2, 6, 2}, {2, 2, 0, 2}, {2, 2, 0, 2}}}; + for (const u32 resource : {2u, 3u, 4u}) { + const auto &size = geometry[resource]; + const u32 height = std::max(size[1], 1u), layers = std::max(size[2], 1u); + std::vector> mips; + for (u32 mip = 0; mip < size[3]; ++mip) + mips.emplace_back(std::max(size[0] >> mip, 1u) * std::max(height >> mip, 1u) * layers * 4u); + vulkan.DestroyImage(&textures[resource]); + textures[resource] = vulkan.CreateImageMips(test.name, size[0], height, + resource == 3u ? vk::Format::eR32G32B32A32Uint : vk::Format::eR32G32B32A32Sfloat, + vk::ImageUsageFlagBits::eSampled, mips, 4u, vk::ImageLayout::eShaderReadOnlyOptimal, + resource == 3u ? vk::ImageType::e1D : vk::ImageType::e2D, + resource == 3u ? vk::ImageViewType::e1D + : resource == 4u ? vk::ImageViewType::e2DArray : vk::ImageViewType::e2D, layers); + } + compiled.program.shader_info_complete = false; + CollectShaderInfo(compiled.program, {.compute = &compute}); + for (const u32 wave_size : {32u, 64u}) { + compute.wave_size = wave_size; + compute.threads_num[0] = wave_size; + 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, "typed query body count", tools.Disassemble(compiled.spirv, &text) && + CountText(text, "OpImageQuerySizeLod") == 3u * halves && + CountText(text, "OpImageQueryLevels") == 3u * halves && + compiled.spirv.size() <= 8192u * halves, + "indirect geometry queries exceeded their per-type body or emitted-size budget"); + std::vector decorated, queried; + for (size_t offset = 5; offset < compiled.spirv.size();) { + const auto *words = compiled.spirv.data() + offset; + const auto op = static_cast(words[0] & 0xffffu); + if (op == spv::OpDecorate && words[2] == spv::DecorationNonUniform) decorated.push_back(words[1]); + if (op == spv::OpImageQuerySizeLod) queried.push_back(words[3]); + offset += words[0] >> 16u; + } + Require(test.name, "nonuniform query image", std::ranges::count_if(queried, [&](u32 operand) { + return std::ranges::find(decorated, operand) != decorated.end(); + }) == halves, "the dynamic query image lacks its nonuniform decoration"); + test.initial.assign(wave_size * 4u, 0xdeadbeefu); + test.expected.clear(); + for (u32 lane = 0; lane < wave_size; ++lane) { + const auto resource = lane < ordinals.size() ? root.indirect_resources[ordinals[lane]] : 0u; + auto size = geometry[resource]; + size[0] = std::max(size[0] >> (lane & 1u), 1u); + if (size[1] != 0u) size[1] = std::max(size[1] >> (lane & 1u), 1u); + test.expected.insert(test.expected.end(), size.begin(), size.end()); + } + auto buffer = vulkan.CreateStorageBuffer(test.name, test.initial, test.initial.size()); + vulkan.Dispatch(test, compiled, buffer, nullptr, nullptr, nullptr, nullptr, {}, textures); + const auto actual = vulkan.ReadBuffer(test.name, buffer, test.expected.size()); + vulkan.DestroyBuffer(&buffer); + CompareWords(test, "selected dimensions, mip levels and unmapped root", test.expected, actual); + } + std::printf("[compute] %-32s ok\n", test.name); + query_store.Invalidate(); + compiled.program.info.images[3].numeric_class = Prospero::TextureNumericClass::Float; vulkan.Device().destroySampler(native_sampler); for (auto &texture : textures) vulkan.DestroyImage(&texture); textures.clear(); // Reuse the GPU key map, including the separated root slot and unmapped suffix, // for storage writes. Every lane has a distinct texel address. - store.Invalidate(); for (auto &inst : *block) if (inst.GetOpcode() == ValueOpcode::ReferenceU32) inst.Invalidate(); compiled.program.info.buffers.clear();