shader: resolve bounded GPU material keys through sampled image arrays

Co-authored-by: TheCruZ <thecruzwtf@gmail.com>
This commit is contained in:
nmzik
2026-10-04 06:05:03 +04:00
parent 139b564c01
commit 843b577880
8 changed files with 266 additions and 43 deletions
@@ -55,6 +55,7 @@ vk::PhysicalDeviceVulkan12Features WindowContext::RequiredVulkan12Features() noe
features.shaderOutputViewportIndex = VK_TRUE;
features.bufferDeviceAddress = VK_TRUE;
features.shaderBufferInt64Atomics = VK_TRUE;
features.shaderSampledImageArrayNonUniformIndexing = VK_TRUE;
return features;
}
@@ -266,6 +267,9 @@ static void VulkanFindPhysicalDevice(vk::Instance instance, vk::SurfaceKHR surfa
required_features12.bufferDeviceAddress);
check_feature(features12.shaderBufferInt64Atomics, "shaderBufferInt64Atomics",
required_features12.shaderBufferInt64Atomics);
check_feature(features12.shaderSampledImageArrayNonUniformIndexing,
"shaderSampledImageArrayNonUniformIndexing",
required_features12.shaderSampledImageArrayNonUniformIndexing);
check_feature(features13.robustImageAccess, "robustImageAccess");
check_feature(features13.dynamicRendering, "dynamicRendering",
required_features13.dynamicRendering);
@@ -83,7 +83,7 @@ uint32_t DescriptorElementPointer(EmitterState& state, uint32_t result_ptr_type,
}
const auto pointer = state.builder.AllocateId();
state.builder.AddFunction(spv::OpAccessChain, result_ptr_type, pointer, variable_id,
ConstantU32(state, array_index));
array_index);
return pointer;
}
@@ -152,8 +152,9 @@ uint32_t ImageViewSizeType(EmitterState& state, ImageDimension dimension) {
}
}
uint32_t LoadImageDescriptor(EmitterState& state, uint32_t resource, uint32_t mip) {
const auto pointer = ImageDescriptorPointer(state, resource, mip);
uint32_t LoadImageDescriptor(EmitterState& state, uint32_t resource, uint32_t mip,
uint32_t array_index) {
const auto pointer = ImageDescriptorPointer(state, resource, mip, array_index);
const auto image = state.builder.AllocateId();
state.builder.AddFunction(spv::OpLoad, ImageType(state, state.program.info.images.at(resource)),
image, pointer);
@@ -167,30 +168,40 @@ uint32_t LoadSamplerDescriptor(EmitterState& state, uint32_t sampler) {
const auto pointer_type =
state.builder.Type(spv::OpTypePointer, spv::StorageClassUniformConstant, sampler_type);
const auto pointer = DescriptorElementPointer(
state, pointer_type, state.sampler_variable, array_index,
state, pointer_type, state.sampler_variable, ConstantU32(state, array_index),
IR::DescriptorBindingKind::Samplers, sampler, "sampler descriptor array was not emitted");
const auto sampler_id = state.builder.AllocateId();
state.builder.AddFunction(spv::OpLoad, sampler_type, sampler_id, pointer);
return sampler_id;
}
uint32_t MakeSampledImage(EmitterState& state, uint32_t resource, uint32_t sampler_id, uint32_t mip) {
uint32_t MakeSampledImage(EmitterState& state, uint32_t resource, uint32_t sampler_id, uint32_t mip,
uint32_t array_index) {
const auto& image_resource = state.program.info.images.at(resource);
EXIT_IF(image_resource.resource_class != IR::ImageResourceClass::Sampled);
const auto image = LoadImageDescriptor(state, resource, mip);
const auto image = LoadImageDescriptor(state, resource, mip, array_index);
const auto sampled_image = state.builder.AllocateId();
const auto sampled_type =
state.builder.Type(spv::OpTypeSampledImage, ImageType(state, image_resource));
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;
}
uint32_t ImageDescriptorPointer(EmitterState& state, uint32_t resource, uint32_t mip) {
uint32_t ImageDescriptorPointer(EmitterState& state, uint32_t resource, uint32_t mip,
uint32_t array_index) {
const auto& image = state.program.info.images.at(resource);
EXIT_IF(mip >= image.mip_count);
const auto kind = IR::DescriptorBindingForImage(image);
EXIT_IF(!kind.has_value());
const auto array_index = ResourceForDescriptor(state, *kind, resource) + mip;
if (array_index == 0u) {
array_index = ConstantU32(state, ResourceForDescriptor(state, *kind, resource) + mip);
}
const auto pointer_type = state.builder.Type(
spv::OpTypePointer, spv::StorageClassUniformConstant, ImageType(state, image));
const auto variable = state.image_variables[IR::ImageBindingIndex(*kind)];
@@ -776,13 +776,13 @@ void EmitImage(ValueEmitContext& ctx, const IR::Inst& inst) {
operands.push_back(AddressF32(ctx, mem, *address, layout.bias));
}
const auto sampler_id = LoadSamplerDescriptor(state, mem.sampler);
const auto EmitSample = [&](uint32_t resource) {
const auto EmitSample = [&](uint32_t resource, uint32_t array_index = 0u) {
const auto& candidate = state.program.info.images[resource];
const auto coord =
CoordF32(ctx, mem, *address, layout.coord,
ImageDimensionInfoFor(candidate.dimension).coordinate_components,
candidate.cube);
const auto sampled = MakeSampledImage(state, resource, sampler_id);
const auto sampled = MakeSampledImage(state, resource, sampler_id, 0u, array_index);
const auto sample = state.builder.AllocateId();
std::vector<uint32_t> sample_operands {result_type, sample, sampled, coord};
if (dref) {
@@ -821,6 +821,41 @@ void EmitImage(ValueEmitContext& ctx, const IR::Inst& inst) {
}
const auto selected = EmitIndirectResourceIndex(
state, key, image.indirect_mapping_offset, image.indirect_search_iterations, 0u);
const auto kind = *IR::DescriptorBindingForImage(image);
const auto& binding = *IR::FindBinding(state.program.bindings, kind);
uint32_t first_child = 0;
bool homogeneous = true;
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];
if (candidate.dimension != image.dimension || candidate.cube != image.cube ||
IR::DescriptorBindingForImage(candidate) != kind ||
candidate.mip_count != 1u) {
homogeneous = false;
break;
}
if (ordinal == 1u) first_child = ResourceForDescriptor(state, kind, resource);
const auto child_slot = first_child + ordinal - 1u;
if (ordinal > 1u && (child_slot >= binding.resources.size() ||
binding.resources[child_slot] != resource)) {
homogeneous = false;
break;
}
}
if (homogeneous) {
const auto child_index = Binary(state, spv::OpIAdd, TypeU32(state), selected,
ConstantU32(state, first_child - 1u));
const auto is_root = Binary(state, spv::OpIEqual, TypeBool(state), selected,
ConstantU32(state, 0u));
const auto array_index = state.builder.AllocateId();
state.builder.AddFunction(spv::OpSelect, TypeU32(state), array_index, is_root,
ConstantU32(state, ResourceForDescriptor(state, kind, mem.resource)),
child_index);
const auto sample = EmitSample(mem.resource, array_index);
const auto result = dref ? sample : UnpackImageTexel(ctx, mem, sample);
ctx.Define(inst, ResultVector(ctx, result, numeric_class, dref, mem));
return;
}
const auto default_label = state.builder.AllocateId();
const auto merge_label = state.builder.AllocateId();
std::vector<uint32_t> labels(image.indirect_resources.size() - 1u);
@@ -316,14 +316,16 @@ uint32_t ImageType(EmitterState& state, const IR::ImageResource& image);
uint32_t ImageViewSizeType(EmitterState& state, ImageDimension dimension);
uint32_t ImageDescriptorPointer(EmitterState& state, uint32_t resource, uint32_t mip = 0);
uint32_t ImageDescriptorPointer(EmitterState& state, uint32_t resource, uint32_t mip = 0,
uint32_t array_index = 0);
uint32_t LoadImageDescriptor(EmitterState& state, uint32_t resource, uint32_t mip = 0);
uint32_t LoadImageDescriptor(EmitterState& state, uint32_t resource, uint32_t mip = 0,
uint32_t array_index = 0);
uint32_t LoadSamplerDescriptor(EmitterState& state, uint32_t sampler);
uint32_t MakeSampledImage(EmitterState& state, uint32_t resource, uint32_t sampler_id,
uint32_t mip = 0);
uint32_t mip = 0, uint32_t array_index = 0);
void EmitStorageImageWrite(EmitterState& state, uint32_t resource, uint32_t mip_lod, uint32_t coord,
uint32_t texel);
@@ -487,6 +487,7 @@ struct DescriptorSource {
uint32_t selector_stride = 0;
uint32_t selector_offset = 0;
uint32_t table_offset = 0;
uint32_t selector_shift = 0;
Value key_count;
Value selector_first;
Value selector_mask;
@@ -285,6 +285,7 @@ bool MaterializeIndirectDescriptor(const ResourcePlan& p
ShaderBufferResource material;
uint32_t first = 0, count = 0;
if (!DecodeBufferDescriptor(material_value, material) || material.Type() != 0u ||
(indirect.selector_shift != 0u && (material.Base48() & 3u) != 0u) ||
material.SwizzleEnabled() ||
material.AddTid() || material.OutOfBounds() != 0u ||
uint64_t {indirect.selector_offset} + 4u > material.Stride() ||
@@ -296,17 +297,20 @@ bool MaterializeIndirectDescriptor(const ResourcePlan& p
indirect.selector_offset + 4u >
uint64_t {UINT32_MAX} + 1u)
return false;
keys.reserve(count);
for (uint64_t index = first; index < uint64_t {first} + count; ++index) {
uint32_t key = 0;
const auto offset = index * material.Stride() + indirect.selector_offset;
if (!ReadScalarTable(material.Base48(), material.GetSize(), offset, runtime,
{&key, 1}))
return false;
keys.push_back(key);
if (material.Stride() == 4u && indirect.selector_offset == 0u) {
keys.resize(count);
if (!ReadScalarTable(material.Base48(), material.GetSize(), uint64_t {first} * 4u,
runtime, keys)) return false;
} else {
keys.reserve(count);
for (uint64_t index = first; index < uint64_t {first} + count; ++index) {
uint32_t key = 0;
const auto offset = index * material.Stride() + indirect.selector_offset;
if (!ReadScalarTable(material.Base48(), material.GetSize(), offset, runtime,
{&key, 1})) return false;
keys.push_back(key);
}
}
std::ranges::sort(keys);
keys.erase(std::unique(keys.begin(), keys.end()), keys.end());
} else if (!indirect.selector_mask.IsEmpty()) {
uint32_t mask = 0;
uint32_t count = 0;
@@ -330,8 +334,6 @@ bool MaterializeIndirectDescriptor(const ResourcePlan& p
keys.push_back(key);
mask &= mask - 1u;
}
std::ranges::sort(keys);
keys.erase(std::unique(keys.begin(), keys.end()), keys.end());
} else {
ShaderBufferResource material;
if (!DecodeBufferDescriptor(material_value, material) || table_value.dword_count != 4u ||
@@ -357,6 +359,12 @@ bool MaterializeIndirectDescriptor(const ResourcePlan& p
}
keys.push_back(key);
}
}
if (indirect.material_source != UINT32_MAX) {
if (indirect.selector_shift != 0u) {
for (auto& key: keys) key >>= indirect.selector_shift;
keys.push_back(0u); // An out-of-range material load returns zero.
}
std::ranges::sort(keys);
keys.erase(std::unique(keys.begin(), keys.end()), keys.end());
}
@@ -426,10 +434,21 @@ struct SamplerPlan {
uint32_t sampler_count = 0;
};
template <typename T, typename Keep>
void CompactImages(std::vector<T>& images, Keep&& keep) {
size_t count = 0;
for (size_t index = 0; index < images.size(); ++index) {
if (!keep(index)) continue;
if (count != index) images[count] = std::move(images[index]);
++count;
}
images.resize(count);
}
struct ImageRemap {
explicit ImageRemap(const ResourceSpecialization& specialization)
: source_count(static_cast<uint32_t>(specialization.images.size())) {
EXIT_IF(specialization.images.size() > indices.size());
: indices(specialization.images.size()),
source_count(static_cast<uint32_t>(specialization.images.size())) {
for (uint32_t index = 0; index < source_count; index++) {
indices[index] = specialization.images[index].fmask ? UINT32_MAX : count++;
}
@@ -446,16 +465,11 @@ struct ImageRemap {
if (count == source_count) {
return;
}
for (uint32_t index = 0; index < source_count; index++) {
if (indices[index] != UINT32_MAX && indices[index] != index) {
images[indices[index]] = std::move(images[index]);
}
}
images.resize(count);
CompactImages(images, [&](size_t index) { return indices[index] != UINT32_MAX; });
}
private:
std::array<uint32_t, ShaderInfo::MaxImages> indices;
std::vector<uint32_t> indices;
uint32_t source_count;
uint32_t count = 0;
};
@@ -602,7 +616,7 @@ static bool BuildResourceSpecialization(const ResourcePlan& program, ResourceSna
}
}
}
ImageRemap(specialization).Apply(snapshot.images);
CompactImages(snapshot.images, [&](size_t index) { return !specialization.images[index].fmask; });
return true;
}
@@ -1146,7 +1160,7 @@ bool MaterializeResources(const ResourcePlan& program, const SrtRuntime& runtime
const auto& indirect = *source->indirect_descriptor;
if (!MaterializeIndirectDescriptor(
program, indirect, i, 8u, observed, clean, snapshot, snapshot.images,
specialization.images, ShaderInfo::MaxImages, [&](DescriptorValue& value) {
specialization.images, UINT32_MAX, [&](DescriptorValue& value) {
if (NullImageDescriptor(value) || !ValidImageDescriptor(value, image.r128))
value.dwords.fill(0u);
return true;
@@ -826,6 +826,7 @@ private:
const auto& b = *descriptor.indirect_descriptor;
if (a.material_source != b.material_source || a.table_source != b.table_source ||
a.selector_stride != b.selector_stride || a.selector_offset != b.selector_offset ||
a.selector_shift != b.selector_shift ||
a.table_offset != b.table_offset || a.sources != b.sources ||
!EquivalentValue(m_program, a.key_count, b.key_count) ||
a.selector_first.IsEmpty() != b.selector_first.IsEmpty() ||
@@ -1549,6 +1550,37 @@ private:
return true;
}
bool MatchUniformizedBufferKey(Value key, const Inst& image,
DescriptorSource::IndirectDescriptor& indirect,
DescriptorSource& material_source) {
const auto* selected = UniformizedMaterialValue(key, image);
if (selected == nullptr) return false;
const auto* shift = selected->Arg(1).Resolve().TryInstruction();
uint32_t amount = 0;
if (shift == nullptr || shift->GetOpcode() != ValueOpcode::ShiftRightLogical32 ||
!ImmediateU32(shift->Arg(1), amount) || amount == 0u || amount >= 32u) return false;
const auto* loaded = shift->Arg(0).Resolve().TryInstruction();
if (loaded == nullptr || loaded->GetOpcode() != ValueOpcode::SelectU32 ||
!EquivalentValue(m_program, selected->Arg(0), loaded->Arg(0))) return false;
const auto* read = loaded->Arg(1).Resolve().TryInstruction();
if (read == nullptr || read->GetOpcode() != ValueOpcode::LoadBufferU32) return false;
const auto& memory = m_program.memory_info[read->Flags<MemoryFlags>().index];
uint32_t offset = 1u, scalar = 1u;
if (memory.kind != ResourceKind::Buffer || memory.typed || memory.formatted ||
!memory.idxen || memory.offen ||
!ImmediateU32(read->Arg(2), offset) || offset != 0u ||
!ImmediateU32(read->Arg(3), scalar) || scalar != 0u ||
!EquivalentValue(m_program, selected->Arg(0), read->Arg(4)) ||
!MakeRuntimeTableSource(*read, material_source)) return false;
// GPU index arithmetic may wrap. Every record in the bounded V# is a candidate.
indirect.material_source = InternSource(material_source);
indirect.selector_first = Value(0u);
indirect.key_count = material_source.dwords[2];
indirect.selector_offset = memory.offset;
indirect.selector_shift = amount;
return true;
}
bool TryMakeIndirectImage(Inst& handle, IndirectDescriptorPlan& plan) {
if (handle.GetOpcode() != ValueOpcode::GetImageResource || handle.NumArgs() != 8u) return false;
DescriptorSource table_source;
@@ -1575,7 +1607,7 @@ private:
}
if ((table_offset & 3u) != 0u ||
(bitscan && table_offset > UINT32_MAX - (32u * 32u - 1u))) return false;
} else {
} else if (!MatchUniformizedBufferKey(key, handle, indirect, material_source)) {
auto* material_read = key.Resolve().TryInstruction();
uint32_t material_memory_index = 0;
const auto* memory = material_read != nullptr