Files
KytyPS5/src/graphics/shader/recompiler/ir/passes/ResourceMaterialization.cpp
T

1201 lines
48 KiB
C++

#include "graphics/shader/recompiler/ir/passes/ResourceMaterialization.h"
#include "common/assert.h"
#include "graphics/guest_gpu/gpu_format.h"
#include "graphics/shader/recompiler/ir/ShaderIR.h"
#include "graphics/shader/shaderBindings.h"
#include <algorithm>
#include <array>
#include <bit>
#include <cstdio>
#include <cstring>
#include <fmt/format.h>
#include <functional>
#include <numeric>
#include <unordered_set>
namespace Libs::Graphics::ShaderRecompiler::IR {
namespace {
constexpr uint64_t AddressMask = 0x0000ffffffffffffull;
constexpr uint64_t MaxIndirectImageProbes = 65536u;
bool SpecializationFail(std::string_view message) {
std::fprintf(stderr, "shader resource specialization failed: %.*s\n",
static_cast<int>(message.size()), message.data());
return false;
}
Decoder::ImageDimension DescriptorDimension(const DescriptorValue& descriptor,
Decoder::ImageDimension requested) {
const bool is_array = requested == Decoder::ImageDimension::Dim1DArray ||
requested == Decoder::ImageDimension::Dim2DArray ||
requested == Decoder::ImageDimension::Dim2DMsaaArray;
switch (static_cast<Prospero::ImageType>((descriptor.dwords[3] >> 28u) & 0xfu)) {
case Prospero::ImageType::kColor1D: return Decoder::ImageDimension::Dim1D;
case Prospero::ImageType::kColor1DArray:
if (is_array) {
return Decoder::ImageDimension::Dim1DArray;
}
return Decoder::ImageDimension::Dim1D;
case Prospero::ImageType::kColor3D: return Decoder::ImageDimension::Dim3D;
case Prospero::ImageType::kCube: return Decoder::ImageDimension::Dim2DArray;
case Prospero::ImageType::kColor2DArray:
if (is_array) {
return Decoder::ImageDimension::Dim2DArray;
}
return Decoder::ImageDimension::Dim2D;
case Prospero::ImageType::kColor2DMsaaArray:
if (is_array) {
return Decoder::ImageDimension::Dim2DMsaaArray;
}
return Decoder::ImageDimension::Dim2DMsaa;
case Prospero::ImageType::kColor2D: return Decoder::ImageDimension::Dim2D;
case Prospero::ImageType::kColor2DMsaa: return Decoder::ImageDimension::Dim2DMsaa;
default: return Decoder::ImageDimension::Unknown;
}
}
bool NullImageDescriptor(const DescriptorValue& descriptor) {
return descriptor.dwords[0] == 0 && (descriptor.dwords[1] & 0xffu) == 0;
}
bool ValidImageDescriptor(const DescriptorValue& descriptor, bool r128 = false) {
const auto& words = descriptor.dwords;
// Reject texture descriptors with nonzero reserved bits.
if ((words[1] & 0x20000000u) != 0u || (words[2] & 0x70003000u) != 0u ||
(!r128 && ((words[4] & 0xe000e000u) != 0u || (words[5] & 0xf9000000u) != 0u ||
(words[6] & 0x00007b00u) != 0u))) {
return false;
}
const auto type = static_cast<Prospero::ImageType>((descriptor.dwords[3] >> 28u) & 0xfu);
const auto format = static_cast<Prospero::BufferFormat>((descriptor.dwords[1] >> 20u) & 0x1ffu);
if (type < Prospero::ImageType::kColor1D || format == Prospero::BufferFormat::kInvalid ||
format > Prospero::BufferFormat::kBc7Srgb) {
return false;
}
if (r128 && type != Prospero::ImageType::kColor1D && type != Prospero::ImageType::kColor2D &&
type != Prospero::ImageType::kColor2DMsaa) {
return false;
}
const bool array = type == Prospero::ImageType::kColor1DArray ||
type == Prospero::ImageType::kColor2DArray ||
type == Prospero::ImageType::kColor2DMsaaArray ||
type == Prospero::ImageType::kCube;
if (array && ((words[4] >> 16u) & 0x1fffu) > (words[4] & 0x1fffu)) {
return false;
}
if (type == Prospero::ImageType::kColor2DMsaa ||
type == Prospero::ImageType::kColor2DMsaaArray) {
const auto base_level = (descriptor.dwords[3] >> 12u) & 0xfu;
const auto fragments = (descriptor.dwords[3] >> 16u) & 0xfu;
const auto max_mip = (descriptor.dwords[5] >> 4u) & 0xfu;
return base_level == 0 && fragments >= 1 && fragments <= 3 &&
(r128 || max_mip == fragments);
}
return true;
}
uint32_t DescriptorImageSwizzle(const DescriptorValue& descriptor) {
return descriptor.dwords[3] & 0xfffu;
}
Prospero::BufferFormat ImageConversionFormat(Prospero::BufferFormat format) {
return Prospero::RemapTextureFormat(format) != format ? format
: Prospero::BufferFormat::kInvalid;
}
bool RequiresPointSampler(const ImageResource& image) {
return image.numeric_class == Prospero::TextureNumericClass::Sint ||
image.conversion_format != Prospero::BufferFormat::kInvalid;
}
bool RequiresPointSampler(const ResourceSpecialization::Image& image) {
return image.numeric_class == Prospero::TextureNumericClass::Sint ||
image.conversion_format != Prospero::BufferFormat::kInvalid;
}
bool DescriptorIsCube(const DescriptorValue& descriptor) {
return static_cast<Prospero::ImageType>((descriptor.dwords[3] >> 28u) & 0xfu) ==
Prospero::ImageType::kCube;
}
uint32_t StorageMipCount(const ImageResource& image, const DescriptorValue& descriptor) {
if (image.mip_mode != ImageMipMode::DynamicStorage || NullImageDescriptor(descriptor)) {
return 1;
}
const auto base = (descriptor.dwords[3] >> 12u) & 0xfu;
const auto last = (descriptor.dwords[3] >> 16u) & 0xfu;
return base <= last ? last - base + 1u : 0u;
}
bool DecodeBufferDescriptor(const DescriptorValue& descriptor, ShaderBufferResource& result) {
if (descriptor.dword_count != std::size(result.fields)) {
return false;
}
std::copy_n(descriptor.dwords.begin(), std::size(result.fields), result.fields);
return true;
}
struct ReadCapture {
SrtRuntime source;
std::vector<std::pair<uint64_t, uint64_t>>& ranges;
};
bool CaptureStrictRead(void* userdata, uint64_t address, std::span<uint32_t> values) {
auto& capture = *static_cast<ReadCapture*>(userdata);
if (!capture.source.read_specialization_memory(capture.source.userdata, address, values)) {
return false;
}
capture.ranges.emplace_back(address, values.size_bytes());
return true;
}
bool CaptureOrdinaryRead(void* userdata, uint64_t address, std::span<uint32_t> values) {
auto& capture = *static_cast<ReadCapture*>(userdata);
if (!capture.source.read_memory(capture.source.userdata, address, values)) {
return false;
}
capture.ranges.emplace_back(address, values.size_bytes());
return true;
}
bool WrittenBuffersDisjoint(const ResourcePlan& program, const ResourceSnapshot& snapshot,
std::span<const std::pair<uint64_t, uint64_t>> reads) {
for (uint32_t i = 0; i < program.info.buffers.size(); ++i) {
if (!program.info.buffers[i].written) continue;
ShaderBufferResource buffer;
if (!DecodeBufferDescriptor(snapshot.buffers[i], buffer)) return false;
const auto base = buffer.Base48();
const auto size = buffer.GetSize();
if (size == 0u) continue;
if (size - 1u > AddressMask - base) return false;
for (const auto [address, bytes]: reads) {
if (bytes == 0u) continue;
if (address > AddressMask || bytes - 1u > AddressMask - address ||
(address <= base ? base - address < bytes : address - base < size)) {
return false;
}
}
}
return true;
}
const DescriptorSource* Source(const ResourcePlan& program, uint32_t source) {
if (source >= program.descriptor_sources.size()) {
return nullptr;
}
return &program.descriptor_sources[source];
}
void MarkCleanFlatSlots(const ResourcePlan& program, const DescriptorSource* source,
std::span<uint8_t> slots, Value extra = {}) {
if (source == nullptr && extra.IsEmpty()) {
return;
}
std::vector<Value> pending;
if (source != nullptr) {
pending.assign(source->dwords.begin(), source->dwords.begin() + source->dword_count);
}
if (!extra.IsEmpty()) pending.push_back(extra);
std::vector<const Inst*> visited;
while (!pending.empty()) {
auto value = pending.back().Resolve();
pending.pop_back();
const auto* inst = value.TryInstruction();
if (inst == nullptr || std::ranges::find(visited, inst) != visited.end()) {
continue;
}
visited.push_back(inst);
if (inst->GetOpcode() == ValueOpcode::ReadConst) {
const auto slot = inst->Arg(1).Resolve();
if (slot.IsImmediate() && slot.GetType() == Type::U32 && slot.U32() < slots.size()) {
slots[slot.U32()] = 1u;
pending.push_back(program.srt_reads[slot.U32()].value);
}
continue;
}
for (size_t arg = 0; arg < inst->NumArgs(); arg++) {
pending.push_back(inst->Arg(arg));
}
}
}
bool ReadScalarTable(uint64_t base, uint64_t size, uint32_t dynamic_offset,
const SrtRuntime& runtime, std::span<uint32_t> words) {
const auto offset = static_cast<uint64_t>(dynamic_offset) & ~uint64_t {3};
const auto count = std::min<uint64_t>(words.size(), offset < size ? (size - offset) / 4u : 0u);
std::ranges::fill(words.subspan(count), 0u);
if (count == 0u) {
return true;
}
base &= AddressMask & ~uint64_t {3};
if (offset > AddressMask - base) {
return false;
}
const auto address = base + offset;
const auto prefix = words.first(count);
return prefix.size_bytes() - 1u <= AddressMask - address &&
runtime.read_specialization_memory != nullptr &&
runtime.read_specialization_memory(runtime.userdata, address, prefix);
}
bool MaterializeIndirectImage(const ResourcePlan& program,
const DescriptorSource::IndirectImage& indirect,
const DescriptorValue& material_value,
const DescriptorValue& table_value, uint32_t image_index,
const SrtRuntime& runtime, SrtWalker& clean,
ResourceSnapshot& snapshot,
ResourceSpecialization& specialization) {
uint64_t table_base = 0;
uint64_t table_size = UINT64_MAX; // Scalar addresses have no buffer descriptor bounds.
ShaderBufferResource table;
if (table_value.dword_count == 2u) {
table_base = (static_cast<uint64_t>(table_value.dwords[1]) << 32u) | table_value.dwords[0];
} else if (DecodeBufferDescriptor(table_value, table)) {
table_base = table.Base48();
table_size = table.GetSize();
} else {
return false;
}
auto& keys = program.material_keys;
keys.clear();
if (indirect.material_source == UINT32_MAX) {
uint32_t key_count = 0;
const bool evaluated = clean.Evaluate(indirect.key_count, key_count);
if (std::bit_cast<int32_t>(key_count) <= 0) key_count = 0;
if (table_value.dword_count != 2u || !evaluated ||
key_count > MaxIndirectImageProbes ||
uint64_t {indirect.table_offset} + uint64_t {key_count} * 32u > UINT32_MAX + 1ull) {
return false;
}
keys.resize(key_count);
std::iota(keys.begin(), keys.end(), 0u);
} else if (!indirect.selector_mask.IsEmpty()) {
uint32_t mask = 0;
uint32_t count = 0;
if (material_value.dword_count != 2u || table_value.dword_count != 2u ||
!clean.Evaluate(indirect.selector_mask, mask) ||
!clean.Evaluate(indirect.key_count, count) || count == 0u || count > 32u) {
return false;
}
if (count < 32u) mask &= (1u << count) - 1u;
const auto material_base =
(static_cast<uint64_t>(material_value.dwords[1]) << 32u) | material_value.dwords[0];
keys.reserve(std::popcount(mask));
while (mask != 0u) {
const auto index = std::countr_zero(mask);
const auto offset = static_cast<uint64_t>(indirect.selector_offset) +
static_cast<uint64_t>(index) * indirect.selector_stride;
if (offset > UINT32_MAX) return false;
uint32_t key = 0;
if (!ReadScalarTable(material_base, UINT64_MAX, static_cast<uint32_t>(offset),
runtime, {&key, 1})) return false;
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 ||
material.Stride() != indirect.selector_stride) {
return false;
}
// Enumerate every wrapped scalar-buffer offset that can pass the descriptor bounds.
const auto step = std::gcd<uint64_t>(indirect.selector_stride, uint64_t {1} << 32u);
const auto residue = static_cast<uint64_t>(indirect.selector_offset) % step;
const auto limit = std::min<uint64_t>(UINT32_MAX, material.GetSize() + 3u);
const auto probe_count = residue <= limit ? (limit - residue) / step + 1u : 0u;
if (probe_count > MaxIndirectImageProbes) {
return false;
}
keys.reserve(static_cast<size_t>(probe_count) + 1u);
keys.push_back(0u);
for (uint64_t offset = residue; offset <= limit && probe_count != 0u; offset += step) {
uint32_t key = 0;
if (!ReadScalarTable(material.Base48(), material.GetSize(),
static_cast<uint32_t>(offset), runtime, {&key, 1})) {
return false;
}
keys.push_back(key);
if (limit - offset < step) {
break;
}
}
std::ranges::sort(keys);
keys.erase(std::unique(keys.begin(), keys.end()), keys.end());
}
const auto children_begin = snapshot.images.size();
const auto mapping_offset = snapshot.flattened_srt.size();
const auto root_image = specialization.images[image_index];
snapshot.flattened_srt.resize(mapping_offset + 1u + keys.size() * 2u);
snapshot.flattened_srt[mapping_offset] = static_cast<uint32_t>(keys.size());
for (uint32_t entry = 0; entry < keys.size(); ++entry) {
const auto key = keys[entry];
DescriptorValue candidate;
candidate.dword_count = 8u;
const auto table_offset = (key << 5u) + indirect.table_offset;
if (!ReadScalarTable(table_base, table_size, table_offset, runtime, candidate.dwords)) {
return false;
}
if (NullImageDescriptor(candidate) ||
!ValidImageDescriptor(candidate, program.info.images[image_index].r128)) {
candidate.dwords.fill(0);
}
uint32_t ordinal = 0;
if (entry == 0) {
snapshot.images[image_index] = candidate;
} else if (snapshot.images[image_index] != candidate) {
const auto found = std::find(snapshot.images.begin() + children_begin,
snapshot.images.end(), candidate);
ordinal = static_cast<uint32_t>(found - snapshot.images.begin() - children_begin + 1u);
if (found == snapshot.images.end()) {
if (snapshot.images.size() >= ShaderInfo::MaxImages) {
return false;
}
snapshot.images.push_back(candidate);
auto child = root_image;
child.indirect_root = image_index;
specialization.images.push_back(child);
}
}
snapshot.flattened_srt[mapping_offset + 1u + entry * 2u] = key;
snapshot.flattened_srt[mapping_offset + 2u + entry * 2u] = ordinal;
}
if (snapshot.images.size() == children_begin) {
snapshot.flattened_srt.resize(mapping_offset);
} else {
auto& root = specialization.images[image_index];
root.indirect_root = image_index;
root.indirect_mapping_offset = static_cast<uint32_t>(mapping_offset);
root.indirect_search_iterations = std::bit_width(keys.size());
}
return true;
}
} // namespace
struct SamplerPlan {
std::array<uint32_t, ShaderInfo::MaxSamplers> point_sampler {};
uint32_t sampler_count = 0;
};
struct ImageRemap {
explicit ImageRemap(const ResourceSpecialization& specialization)
: source_count(static_cast<uint32_t>(specialization.images.size())) {
EXIT_IF(specialization.images.size() > indices.size());
for (uint32_t index = 0; index < source_count; index++) {
indices[index] = specialization.images[index].fmask ? UINT32_MAX : count++;
}
}
uint32_t operator[](uint32_t index) const {
EXIT_IF(index >= source_count);
return indices[index];
}
template <typename T>
void Apply(std::vector<T>& images) const {
EXIT_IF(images.size() != source_count);
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);
}
private:
std::array<uint32_t, ShaderInfo::MaxImages> indices;
uint32_t source_count;
uint32_t count = 0;
};
template <typename Images>
bool BuildSamplerPlan(const ShaderInfo& base, const Images& images, SamplerPlan& plan);
static bool BuildResourceSpecialization(const ResourcePlan& program, ResourceSnapshot& snapshot,
ResourceSpecialization& specialization) {
specialization.buffers.clear();
specialization.buffers.reserve(program.info.buffers.size());
for (uint32_t i = 0; i < program.info.buffers.size(); i++) {
auto& descriptor_value = snapshot.buffers[i];
ShaderBufferResource descriptor;
if (!DecodeBufferDescriptor(descriptor_value, descriptor)) {
return SpecializationFail(fmt::format("buffer descriptor {} has invalid width", i));
}
if (descriptor.Type() != 0) {
descriptor_value.dwords.fill(0);
descriptor = {};
}
auto packed_stride = descriptor.PackedStride();
const auto stride = packed_stride & 0x3fffu;
const bool swizzle = stride != 0u && ((packed_stride >> 14u) & 1u) != 0u;
if (stride == 0u) {
packed_stride &= ~((1u << 14u) | (3u << 16u));
} else if (!swizzle) {
packed_stride &= ~(3u << 16u);
}
specialization.buffers.push_back({
.packed_stride = packed_stride,
.descriptor_format = program.info.buffers[i].formatted
? descriptor.Format()
: Prospero::BufferFormat::kInvalid,
.descriptor_swizzle =
program.info.buffers[i].formatted ? descriptor.DstSelXYZW() : DstSel(4, 5, 6, 7),
});
}
for (uint32_t i = 0; i < specialization.images.size(); i++) {
const auto& descriptor = snapshot.images[i];
auto& image = specialization.images[i];
const auto base_index = i < program.info.images.size() ? i : image.indirect_root;
if (base_index >= program.info.images.size()) {
return SpecializationFail(fmt::format("image resource {} has an invalid root", i));
}
const auto& base = program.info.images[base_index];
if (base.resource_class == ImageResourceClass::None ||
(base.atomic && base.resource_class != ImageResourceClass::Storage)) {
return SpecializationFail(fmt::format("image resource {} has an invalid class", i));
}
image.mip_count = StorageMipCount(base, descriptor);
if (image.mip_count == 0u) {
return SpecializationFail(
fmt::format("storage image descriptor {} has an invalid mip range", i));
}
if (NullImageDescriptor(descriptor)) {
image.numeric_class = base.atomic ? Prospero::TextureNumericClass::Uint
: Prospero::TextureNumericClass::Float;
image.dimension = Decoder::ImageDimension::Dim2D;
image.cube = false;
continue;
}
const auto descriptor_dimension = DescriptorDimension(descriptor, base.dimension);
if (descriptor_dimension == Decoder::ImageDimension::Unknown) {
return SpecializationFail(fmt::format(
"image descriptor {} has unsupported type {}: {:08x},{:08x},{:08x},{:08x},"
"{:08x},{:08x},{:08x},{:08x}",
i, (descriptor.dwords[3] >> 28u) & 0xfu, descriptor.dwords[0], descriptor.dwords[1],
descriptor.dwords[2], descriptor.dwords[3], descriptor.dwords[4],
descriptor.dwords[5], descriptor.dwords[6], descriptor.dwords[7]));
}
image.dimension = descriptor_dimension;
image.cube = DescriptorIsCube(descriptor);
const auto format =
static_cast<Prospero::BufferFormat>((descriptor.dwords[1] >> 20u) & 0x1ffu);
if (base.atomic && format != Prospero::BufferFormat::k32UInt) {
return SpecializationFail(
fmt::format("atomic image descriptor {} uses unsupported format {}", i,
static_cast<uint32_t>(format)));
}
const bool storage = base.resource_class == ImageResourceClass::Storage;
image.fmask = Prospero::IsFmaskTextureFormat(format);
if (image.fmask) {
if (storage || base.depth_compare ||
image.indirect_root != ImageResource::NoIndirectImage ||
std::ranges::any_of(program.info.sampled_pairs,
[&](const auto& pair) { return pair.image == i; })) {
return SpecializationFail("FMASK requires a direct image load");
}
}
image.conversion_format = ImageConversionFormat(format);
if (storage || image.conversion_format != Prospero::BufferFormat::kInvalid) {
image.shader_swizzle = DescriptorImageSwizzle(descriptor);
}
const bool raw_sint_storage = storage && format == Prospero::BufferFormat::k32SInt &&
base.written && !base.read && !base.atomic;
image.numeric_class = Prospero::SampledTextureNumericClass(format);
if (storage) {
if ((!raw_sint_storage && image.numeric_class == Prospero::TextureNumericClass::Sint) ||
image.numeric_class == Prospero::TextureNumericClass::Unsupported) {
return SpecializationFail(
fmt::format("storage image descriptor {} uses unsupported format {}", i,
static_cast<uint32_t>(format)));
}
if (raw_sint_storage) {
image.numeric_class = Prospero::TextureNumericClass::Uint;
}
} else if (image.numeric_class == Prospero::TextureNumericClass::Unsupported ||
(base.depth_compare &&
image.numeric_class != Prospero::TextureNumericClass::Float)) {
return SpecializationFail(
fmt::format("sampled image descriptor {} uses unsupported format {}", i,
static_cast<uint32_t>(format)));
}
}
for (uint32_t root_index = 0; root_index < specialization.images.size(); root_index++) {
auto& root = specialization.images[root_index];
if (root.indirect_root != root_index) {
continue;
}
const auto key_count = root.indirect_mapping_offset < snapshot.flattened_srt.size()
? snapshot.flattened_srt[root.indirect_mapping_offset]
: 0u;
if (root.indirect_search_iterations == 0u || key_count < 2u ||
static_cast<size_t>(root.indirect_mapping_offset) + 1u +
static_cast<size_t>(key_count) * 2u >
snapshot.flattened_srt.size()) {
return SpecializationFail("indirect image specialization has an invalid key mapping");
}
uint32_t exemplar = ImageResource::NoIndirectImage;
uint32_t resource_count = 0;
for (uint32_t resource = 0; resource < specialization.images.size(); resource++) {
if (specialization.images[resource].indirect_root != root_index) {
continue;
}
resource_count++;
if (exemplar == ImageResource::NoIndirectImage &&
!NullImageDescriptor(snapshot.images[resource])) {
exemplar = resource;
}
}
if (resource_count < 2u || exemplar == ImageResource::NoIndirectImage) {
return SpecializationFail("indirect image specialization has no typed candidate");
}
const auto& image_class = specialization.images[exemplar];
const auto is_2d = [](Decoder::ImageDimension dimension) {
return dimension == Decoder::ImageDimension::Dim2D ||
dimension == Decoder::ImageDimension::Dim2DArray;
};
for (uint32_t candidate = 0; candidate < specialization.images.size(); candidate++) {
auto& image = specialization.images[candidate];
if (image.indirect_root != root_index) {
continue;
}
if (NullImageDescriptor(snapshot.images[candidate])) {
image.numeric_class = image_class.numeric_class;
image.dimension = image_class.dimension;
image.mip_count = image_class.mip_count;
image.conversion_format = image_class.conversion_format;
image.shader_swizzle = image_class.shader_swizzle;
image.cube = image_class.cube;
}
const bool same_coordinates = image.dimension == image_class.dimension &&
image.cube == image_class.cube;
if (image.numeric_class != image_class.numeric_class ||
(!same_coordinates && !(is_2d(image.dimension) && is_2d(image_class.dimension))) ||
image.mip_count != image_class.mip_count ||
image.conversion_format != image_class.conversion_format ||
image.shader_swizzle != image_class.shader_swizzle) {
return SpecializationFail(
fmt::format("indirect image table at pc 0x{:08x} has incompatible candidates",
program.info.images[root_index].first_use_pc));
}
}
}
SamplerPlan sampler_plan;
if (!BuildSamplerPlan(program.info, specialization.images, sampler_plan)) {
return SpecializationFail("specialized sampler layout exceeds its resource limit");
}
for (uint32_t index = 0; index < program.info.samplers.size(); index++) {
const auto target = sampler_plan.point_sampler[index];
if (target != UINT32_MAX && target >= program.info.samplers.size()) {
snapshot.samplers.push_back(snapshot.samplers[index]);
}
}
ImageRemap(specialization).Apply(snapshot.images);
return true;
}
template <typename Images>
bool BuildSamplerPlan(const ShaderInfo& base, const Images& images, SamplerPlan& plan) {
if (base.samplers.size() > plan.point_sampler.size()) {
return false;
}
std::array<uint8_t, ShaderInfo::MaxSamplers> usage {};
plan.point_sampler.fill(UINT32_MAX);
plan.sampler_count = static_cast<uint32_t>(base.samplers.size());
for (const auto& pair: base.sampled_pairs) {
if (pair.image >= images.size() || pair.sampler >= base.samplers.size()) {
return false;
}
usage[pair.sampler] |= RequiresPointSampler(images[pair.image]) ? 2u : 1u;
}
for (uint32_t index = 0; index < base.samplers.size(); index++) {
if ((usage[index] & 2u) == 0u) {
continue;
}
if ((usage[index] & 1u) == 0u) {
plan.point_sampler[index] = index;
} else {
if (plan.sampler_count >= ShaderInfo::MaxSamplers) {
return false;
}
plan.point_sampler[index] = plan.sampler_count++;
}
}
return true;
}
static std::vector<ResourceBlock> ResourceControlFlow(const Program& program) {
// Any shader write may alias a scalar predicate read, including on a later loop visit.
if (program.blocks.size() != program.block_info.size() || HasShaderMemoryWrites(program)) {
return {};
}
std::unordered_map<uint32_t, uint32_t> indices;
for (uint32_t i = 0; i < program.block_info.size(); i++) {
if (!indices.emplace(program.block_info[i].id, i).second) {
return {};
}
}
std::vector<ResourceBlock> blocks(program.blocks.size());
for (uint32_t i = 0; i < blocks.size(); i++) {
auto& block = blocks[i];
const auto& info = program.block_info[i];
const auto& terminator = info.terminator;
std::vector<uint32_t> successors;
switch (terminator.kind) {
case CFG::TerminatorKind::Branch: successors.push_back(terminator.true_block); break;
case CFG::TerminatorKind::ConditionalBranch:
successors = {terminator.true_block, terminator.false_block};
if (ValidateRuntimeValue(program, info.condition, RuntimeValueType::Integer)) {
block.condition = info.condition;
}
break;
case CFG::TerminatorKind::IndirectBranch:
successors = terminator.indirect_targets;
break;
case CFG::TerminatorKind::Return: break;
default: return {};
}
for (const auto successor: successors) {
const auto found = indices.find(successor);
if (found == indices.end()) {
return {};
}
block.successors.push_back(found->second);
}
for (const auto& inst: *program.blocks[i]) {
const auto op = inst.GetOpcode();
const auto buffer = BufferAccessOf(op);
const auto image = ImageOpcodeInfoOf(op);
if (buffer == BufferAccess::None && image.access == ImageAccess::None) {
continue;
}
const auto& memory = program.memory_info.at(inst.Flags<MemoryFlags>().index);
if (memory.planning_only) {
continue;
}
if (buffer != BufferAccess::None) {
block.sources.push_back(program.info.buffers.at(memory.resource).source);
} else {
block.sources.push_back(program.info.images.at(memory.resource).source);
if (image.needs_sampler) {
block.sources.push_back(program.info.samplers.at(memory.sampler).source);
}
}
}
std::ranges::sort(block.sources);
block.sources.erase(std::unique(block.sources.begin(), block.sources.end()),
block.sources.end());
}
if (std::ranges::none_of(
blocks, [](const ResourceBlock& block) { return !block.condition.IsEmpty(); })) {
return {};
}
return blocks;
}
// Nonnegative affine coefficients for constant, local and workgroup coordinates. Reject modular
// arithmetic that could wrap; runtime coverage also bounds the largest invocation index.
static std::optional<std::array<uint64_t, 3>> FillIndex(Value value, uint32_t axis,
uint32_t depth = 0) {
value = value.Resolve();
if (depth > 32 || value.GetType() != Type::U32) {
return {};
}
if (value.IsImmediate()) {
return std::array<uint64_t, 3> {value.U32(), 0, 0};
}
const auto* inst = value.TryInstruction();
if (inst == nullptr) {
return {};
}
const auto op = inst->GetOpcode();
if (op == ValueOpcode::GetBuiltin && inst->Arg(1) == Value(axis)) {
if (inst->Arg(0) == Value(static_cast<uint32_t>(StageInputKind::LocalInvocationId))) {
return std::array<uint64_t, 3> {0, 1, 0};
}
if (inst->Arg(0) == Value(static_cast<uint32_t>(StageInputKind::WorkgroupId))) {
return std::array<uint64_t, 3> {0, 0, 1};
}
}
if (op != ValueOpcode::IAdd32 && op != ValueOpcode::IMul32 &&
op != ValueOpcode::ShiftLeftLogical32) {
return {};
}
auto left = FillIndex(inst->Arg(0), axis, depth + 1);
auto right = FillIndex(inst->Arg(1), axis, depth + 1);
if (!left || !right) {
return {};
}
if (op == ValueOpcode::IMul32 && ((*right)[1] != 0 || (*right)[2] != 0)) {
std::swap(left, right);
}
if (op != ValueOpcode::IAdd32 && ((*right)[1] != 0 || (*right)[2] != 0)) {
return {};
}
if (op == ValueOpcode::ShiftLeftLogical32) {
if ((*right)[0] >= 32) return {};
(*right)[0] = uint64_t {1} << (*right)[0];
}
for (uint32_t i = 0; i < left->size(); ++i) {
(*left)[i] =
op == ValueOpcode::IAdd32 ? (*left)[i] + (*right)[i] : (*left)[i] * (*right)[0];
if ((*left)[i] > UINT32_MAX) return {};
}
return left;
}
static UniformFillPlan AnalyzeUniformFill(const Program& program) {
if (program.stage != ShaderType::Compute || program.blocks.empty() ||
program.blocks.size() != program.block_info.size() || program.info.uses_dma ||
!program.info.samplers.empty()) {
return {};
}
std::unordered_set<uint32_t> visited;
uint32_t index = 0;
const Inst* store = nullptr;
for (;;) {
if (!visited.insert(index).second) return {};
for (const auto& inst: *program.blocks[index]) {
if (AddressOpcodeInfoOf(inst.GetOpcode()).access != AddressAccess::None) return {};
if (!inst.MayHaveSideEffects()) continue;
if (store != nullptr || (BufferAccessOf(inst.GetOpcode()) != BufferAccess::Write &&
inst.GetOpcode() != ValueOpcode::ImageWrite))
return {};
store = &inst;
}
const auto& term = program.block_info[index].terminator;
if (term.kind == CFG::TerminatorKind::Return) break;
if (term.kind != CFG::TerminatorKind::Branch) return {};
const auto next = std::ranges::find(program.block_info, term.true_block, &BlockInfo::id);
if (next == program.block_info.end()) return {};
index = static_cast<uint32_t>(next - program.block_info.begin());
}
if (store == nullptr || visited.size() != program.blocks.size()) return {};
for (const auto& buffer: program.info.buffers) {
if (buffer.read && (!buffer.scalar || buffer.written)) return {};
}
const auto& memory = program.memory_info.at(store->Flags<MemoryFlags>().index);
UniformFillPlan result;
result.fill.resource = memory.resource;
Value data;
if (store->GetOpcode() == ValueOpcode::ImageWrite) {
if (program.info.images.size() != 1 || memory.dmask != 1 || memory.data_bits != 32 ||
memory.image_has_mip || memory.image_sample_flags != 0 || memory.image_r128 ||
memory.image_dimension != Decoder::ImageDimension::Dim2DArray ||
store->Arg(3).Resolve() != Value(true)) return {};
const auto& image = program.info.images[memory.resource];
if (image.read || image.atomic || image.mip_mode != ImageMipMode::None) return {};
const auto* address = store->Arg(1).ResolveInstruction();
if (address == nullptr || address->GetOpcode() != ValueOpcode::MakeImageAddress) return {};
for (uint32_t axis = 0; axis < 3; ++axis) {
const auto index = FillIndex(address->Arg(axis), axis);
if (!index || (*index)[0] != 0) return {};
if (axis < 2) {
if ((*index)[1] != 1 || (*index)[2] == 0) return {};
} else if ((*index)[1] != 0 || (*index)[2] != 1) {
return {};
}
result.fill.group_stride[axis] = static_cast<uint32_t>((*index)[2]);
}
const auto* values = store->Arg(2).ResolveInstruction();
if (values == nullptr || values->GetOpcode() != ValueOpcode::CompositeConstructU32x4)
return {};
result.fill.kind = UniformFillKind::Image;
result.fill.words = 1;
data = values->Arg(0);
} else {
if (!program.info.images.empty()) return {};
const auto op = store->GetOpcode();
constexpr std::array stores {ValueOpcode::StoreBufferU32, ValueOpcode::StoreBufferU32x2,
ValueOpcode::StoreBufferU32x3, ValueOpcode::StoreBufferU32x4};
const auto store_op = std::ranges::find(stores, op);
if (store_op == stores.end() || store->Arg(2).Resolve() != Value(0u) ||
store->Arg(3).Resolve() != Value(0u) || store->Arg(5).Resolve() != Value(true))
return {};
if (!memory.formatted || memory.typed || !memory.idxen || memory.offen || memory.offset != 0 ||
memory.data_bits != 32 ||
memory.data_dwords != static_cast<uint32_t>(store_op - stores.begin() + 1))
return {};
const auto address = FillIndex(store->Arg(1), 0);
if (!address || (*address)[0] != 0 || (*address)[1] != 1 || (*address)[2] == 0) return {};
result.fill.kind = UniformFillKind::Buffer;
result.fill.group_stride[0] = static_cast<uint32_t>((*address)[2]);
result.fill.words = memory.data_dwords;
data = store->Arg(4);
}
data = data.Resolve();
const auto* vector = data.TryInstruction();
constexpr std::array composites {ValueOpcode::CompositeConstructU32x2,
ValueOpcode::CompositeConstructU32x3,
ValueOpcode::CompositeConstructU32x4};
if (result.fill.words > 1 &&
(vector == nullptr || vector->GetOpcode() != composites[result.fill.words - 2]))
return {};
for (uint32_t i = 0; i < result.fill.words; ++i) {
const auto word = result.fill.words == 1 ? data : vector->Arg(i);
if (word.GetType() != Type::U32 ||
!ValidateRuntimeValue(program, word, RuntimeValueType::Integer))
return {};
result.values[i] = word;
}
return result;
}
ResourcePlan ExtractResourcePlan(const Program& program) {
ResourcePlan plan;
plan.stage = program.stage;
plan.shader_hash = program.shader_hash;
plan.user_data_base = program.user_data_base;
plan.user_data_count = program.user_data_count;
plan.info = program.info;
plan.memory_info = program.memory_info;
plan.srt_plan_complete = program.srt_plan_complete;
plan.resource_tracking_complete = program.resource_tracking_complete;
plan.has_address_writes = program.has_address_writes;
std::unordered_map<const Inst*, Inst*> cloned;
std::function<Value(Value)> Clone = [&](Value value) -> Value {
value = value.Resolve();
const auto* source = value.TryInstruction();
if (source == nullptr) {
return value;
}
if (source->GetOpcode() == ValueOpcode::Phi) {
const auto invariant = ResolveInvariantPhi(program, value);
if (!invariant.IsEmpty() && invariant != value) {
return Clone(invariant);
}
}
if (const auto found = cloned.find(source); found != cloned.end()) {
return Value(found->second);
}
auto& target =
plan.value_storage.emplace_back(source->GetOpcode(), source->Flags<uint64_t>());
cloned.emplace(source, &target);
if (source->GetOpcode() == ValueOpcode::Phi) {
for (size_t index = 0; index < source->NumArgs(); index++) {
target.AddPhiOperand(nullptr, Clone(source->Arg(index)));
}
} else {
for (size_t index = 0; index < source->NumArgs(); index++) {
target.SetArg(index, Clone(source->Arg(index)));
}
}
return Value(&target);
};
plan.descriptor_sources.reserve(program.descriptor_sources.size());
for (const auto& source: program.descriptor_sources) {
auto& target = plan.descriptor_sources.emplace_back();
target.dword_count = source.dword_count;
target.indirect_image = source.indirect_image;
if (target.indirect_image.has_value()) {
target.indirect_image->key_count = Clone(target.indirect_image->key_count);
target.indirect_image->selector_mask = Clone(target.indirect_image->selector_mask);
}
for (uint32_t dword = 0; dword < source.dword_count; dword++) {
target.dwords[dword] = Clone(source.dwords[dword]);
}
}
plan.srt_reads.reserve(program.srt_reads.size());
for (const auto& read: program.srt_reads) {
plan.srt_reads.push_back({Clone(read.value), read.flat_offset});
}
plan.control_flow = ResourceControlFlow(program);
for (auto& block: plan.control_flow) {
block.condition = Clone(block.condition);
}
plan.uniform_fill = AnalyzeUniformFill(program);
for (uint32_t i = 0; i < plan.uniform_fill.fill.words; ++i) {
plan.uniform_fill.values[i] = Clone(plan.uniform_fill.values[i]);
}
plan.clean_flat_slots.resize(plan.srt_reads.size());
for (const auto& image: plan.info.images) {
const auto* source = Source(plan, image.source);
if (source == nullptr || !source->indirect_image.has_value()) {
continue;
}
plan.requires_specialization_memory = true;
MarkCleanFlatSlots(plan, Source(plan, source->indirect_image->material_source),
plan.clean_flat_slots, source->indirect_image->selector_mask);
MarkCleanFlatSlots(plan, Source(plan, source->indirect_image->table_source),
plan.clean_flat_slots);
}
return plan;
}
bool MaterializeResources(const ResourcePlan& program, const SrtRuntime& runtime,
ResourceSnapshot& snapshot, ResourceSpecialization& specialization) {
if (!program.resource_tracking_complete ||
(program.requires_specialization_memory && runtime.read_specialization_memory == nullptr)) {
return false;
}
const bool masked_image = std::ranges::any_of(program.info.images, [&](const auto& image) {
const auto* source = Source(program, image.source);
return source != nullptr && source->indirect_image.has_value() &&
!source->indirect_image->selector_mask.IsEmpty();
});
if (masked_image &&
(program.has_address_writes ||
std::ranges::any_of(program.info.images, &ImageResource::written))) {
return false;
}
const bool capture_reads = masked_image &&
std::ranges::any_of(program.info.buffers, &BufferResource::written);
auto& reads = program.specialization_reads;
ReadCapture capture {runtime, reads};
SrtRuntime observed = runtime;
if (capture_reads) {
reads.clear();
observed.userdata = &capture;
observed.read_specialization_memory = CaptureStrictRead;
if (observed.read_memory != nullptr) observed.read_memory = CaptureOrdinaryRead;
}
SrtWalker clean(program, CleanRuntime(observed));
SrtWalker walker(program, observed, program.clean_flat_slots, &clean);
const auto active = clean.FindActiveSources();
if (!walker.RefreshFlatBuffer(snapshot.flattened_srt)) {
return false;
}
snapshot.uniform_fill = {};
const auto& fill = program.uniform_fill;
const auto words = fill.fill.words;
std::array<uint32_t, 4> stored {};
bool uniform_fill = words != 0;
for (uint32_t i = 0; i < words && uniform_fill; ++i) {
uniform_fill = clean.Evaluate(fill.values[i], stored[i]) && stored[i] == stored[0];
}
if (uniform_fill) {
snapshot.uniform_fill = fill.fill;
snapshot.uniform_fill.value = stored[0];
}
const auto evaluate = [&](uint32_t source, DescriptorValue& value) {
if (source >= program.descriptor_sources.size()) {
return false;
}
if (active.empty() || active[source]) {
return walker.EvaluateDescriptor(source, value);
}
value = {};
value.dword_count = program.descriptor_sources[source].dword_count;
return true;
};
snapshot.buffers.resize(program.info.buffers.size());
for (uint32_t i = 0; i < program.info.buffers.size(); ++i) {
if (!evaluate(program.info.buffers[i].source, snapshot.buffers[i])) {
return false;
}
}
if (capture_reads) {
for (uint32_t i = 0; i < program.info.buffers.size(); ++i) {
const auto& buffer = program.info.buffers[i];
if (!buffer.written || (!active.empty() && !active[buffer.source])) continue;
DescriptorValue strict;
if (!clean.EvaluateDescriptor(buffer.source, strict) ||
strict != snapshot.buffers[i]) return false;
}
}
snapshot.images.resize(program.info.images.size());
specialization.images.resize(program.info.images.size());
for (uint32_t i = 0; i < program.info.images.size(); ++i) {
const auto& image = program.info.images[i];
specialization.images[i] = {
.numeric_class = image.numeric_class,
.dimension = image.dimension,
.mip_count = image.mip_count,
.conversion_format = image.conversion_format,
.shader_swizzle = image.shader_swizzle,
.indirect_root = image.indirect_root,
.indirect_mapping_offset = image.indirect_mapping_offset,
.indirect_search_iterations = image.indirect_search_iterations,
.cube = image.cube,
};
const auto* source = Source(program, image.source);
if (source == nullptr) {
return false;
}
if (source->indirect_image.has_value()) {
snapshot.images[i] = {.dword_count = 8u};
if (!active.empty() && !active[image.source]) {
continue;
}
const auto& indirect = *source->indirect_image;
DescriptorValue material;
DescriptorValue table;
if ((indirect.material_source != UINT32_MAX &&
!clean.EvaluateDescriptor(indirect.material_source, material)) ||
!clean.EvaluateDescriptor(indirect.table_source, table) ||
!MaterializeIndirectImage(program, indirect, material, table, i, observed, clean, snapshot,
specialization)) {
return false;
}
} else {
if (!evaluate(image.source, snapshot.images[i])) {
return false;
}
if (!ValidImageDescriptor(snapshot.images[i], image.r128)) {
snapshot.images[i].dwords.fill(0);
}
}
}
snapshot.samplers.resize(program.info.samplers.size());
for (uint32_t i = 0; i < program.info.samplers.size(); ++i) {
if (!evaluate(program.info.samplers[i].source, snapshot.samplers[i])) {
return false;
}
}
if (capture_reads && !WrittenBuffersDisjoint(program, snapshot, reads)) return false;
snapshot.user_data.assign(runtime.user_data.begin(), runtime.user_data.end());
return BuildResourceSpecialization(program, snapshot, specialization);
}
void ApplyResourceSpecialization(Program& program, const ResourceSpecialization& specialization) {
EXIT_IF(!program.resource_tracking_complete || program.shader_info_complete ||
program.binding_layout_complete);
EXIT_IF(program.info.buffers.size() != specialization.buffers.size() ||
program.info.images.size() > specialization.images.size());
auto buffers = program.info.buffers;
for (size_t index = 0; index < buffers.size(); index++) {
buffers[index].packed_stride = specialization.buffers[index].packed_stride;
buffers[index].descriptor_format = specialization.buffers[index].descriptor_format;
buffers[index].descriptor_swizzle = specialization.buffers[index].descriptor_swizzle;
}
auto images = program.info.images;
images.reserve(specialization.images.size());
for (uint32_t index = 0; index < specialization.images.size(); index++) {
const auto& source = specialization.images[index];
if (index >= images.size()) {
EXIT_IF(source.indirect_root >= program.info.images.size());
images.push_back(program.info.images[source.indirect_root]);
}
auto& image = images[index];
image.numeric_class = source.numeric_class;
image.dimension = source.dimension;
image.mip_count = source.mip_count;
image.conversion_format = source.conversion_format;
image.shader_swizzle = source.shader_swizzle;
image.indirect_root = source.indirect_root;
image.indirect_mapping_offset = source.indirect_mapping_offset;
image.indirect_search_iterations = source.indirect_search_iterations;
image.cube = source.cube;
image.indirect_resources.clear();
}
for (uint32_t index = 0; index < images.size(); index++) {
const auto root = images[index].indirect_root;
if (root != ImageResource::NoIndirectImage) {
EXIT_IF(root >= images.size());
images[root].indirect_resources.push_back(index);
}
}
SamplerPlan sampler_plan;
EXIT_IF(!BuildSamplerPlan(program.info, images, sampler_plan));
auto samplers = program.info.samplers;
auto sampled_pairs = program.info.sampled_pairs;
samplers.reserve(sampler_plan.sampler_count);
for (uint32_t index = 0; index < program.info.samplers.size(); index++) {
const auto target = sampler_plan.point_sampler[index];
if (target == UINT32_MAX) {
continue;
}
if (target == index) {
samplers[index].force_point_filtering = true;
} else {
EXIT_IF(target != samplers.size());
auto sampler = samplers[index];
sampler.force_point_filtering = true;
samplers.push_back(std::move(sampler));
}
}
for (auto& pair: sampled_pairs) {
if (RequiresPointSampler(images[pair.image])) {
EXIT_IF(sampler_plan.point_sampler[pair.sampler] == UINT32_MAX);
pair.sampler = sampler_plan.point_sampler[pair.sampler];
}
samplers[pair.sampler].depth_compare |= images[pair.image].depth_compare;
}
auto memory_info = program.memory_info;
const ImageRemap image_remap(specialization);
for (auto* block: program.blocks) {
for (auto it = block->begin(); it != block->end(); ++it) {
auto& inst = *it;
const auto image_opcode = ImageOpcodeInfoOf(inst.GetOpcode());
if (image_opcode.access == ImageAccess::None) {
continue;
}
const auto index = inst.Flags<MemoryFlags>().index;
EXIT_IF(index >= memory_info.size());
auto& memory = memory_info[index];
EXIT_IF(memory.resource >= images.size());
const auto& image = images[memory.resource];
if (specialization.images[memory.resource].fmask) {
EXIT_IF(inst.GetOpcode() != ValueOpcode::ImageRead || memory.data_bits != 32u);
// Vulkan MSAA stores each sample directly; FMASK's four-bit fragment indices
// therefore map each coverage sample to the same host sample.
constexpr uint32_t indices[] = {0x76543210u, 0xfedcba98u};
std::array<Value, 2> fragments;
for (uint32_t component = 0; component < fragments.size(); component++) {
const auto selected = block->PrependNewInst(
it, ValueOpcode::SelectU32, {inst.Arg(2), Value(indices[component]), Value(0u)});
fragments[component] = Value(&*selected);
}
const auto result = block->PrependNewInst(
it, ValueOpcode::CompositeConstructU32x4,
{fragments[0], fragments[1], Value(0u), Value(0u)});
inst.ReplaceUsesWith(Value(&*result));
continue;
}
if (image_opcode.needs_sampler && RequiresPointSampler(image) &&
memory.sampler < program.info.samplers.size()) {
EXIT_IF(sampler_plan.point_sampler[memory.sampler] == UINT32_MAX);
memory.sampler = sampler_plan.point_sampler[memory.sampler];
}
EXIT_IF(image.indirect_root == memory.resource &&
inst.GetOpcode() != ValueOpcode::ImageSampleRaw);
}
}
for (auto* block: program.blocks) {
for (auto& inst: *block) {
if (inst.GetOpcode() == ValueOpcode::GetImageResource) {
inst.SetFlags(image_remap[inst.Flags<uint32_t>()]);
}
}
}
for (auto& memory: memory_info) {
if (memory.kind == ResourceKind::Image && !memory.planning_only) {
memory.resource = image_remap[memory.resource];
}
}
for (auto& buffer: buffers) {
if (buffer.image_alias != BufferResource::NoImageAlias) {
buffer.image_alias = image_remap[buffer.image_alias];
}
}
for (auto& pair: sampled_pairs) {
pair.image = image_remap[pair.image];
}
for (auto& image: images) {
if (image.indirect_root != ImageResource::NoIndirectImage) {
image.indirect_root = image_remap[image.indirect_root];
}
for (auto& resource: image.indirect_resources) {
resource = image_remap[resource];
}
}
image_remap.Apply(images);
program.info.buffers = std::move(buffers);
program.info.images = std::move(images);
program.info.samplers = std::move(samplers);
program.info.sampled_pairs = std::move(sampled_pairs);
program.memory_info = std::move(memory_info);
}
} // namespace Libs::Graphics::ShaderRecompiler::IR