shader: bind dispatch spans while keeping GPU record counts native

Co-authored-by: TheCruZ <thecruzwtf@gmail.com>
This commit is contained in:
nmzik
2026-10-04 10:13:35 +04:00
parent cd9fc6ff2d
commit 8b05b6730e
11 changed files with 207 additions and 16 deletions
@@ -641,6 +641,7 @@ CompileResult CompileProgram(TranslateResult translated, const CompileOptions& o
for (auto& inst: *block) {
const auto op = inst.GetOpcode();
uint32_t first = 0;
bool native_buffer_count = false;
if (op == IR::ValueOpcode::GetBufferResource) {
if (std::ranges::any_of(inst.Uses(), [&](const IR::Use& use) {
return ir.memory_info[use.user->Flags<IR::MemoryFlags>().index].kind ==
@@ -649,6 +650,8 @@ CompileResult CompileProgram(TranslateResult translated, const CompileOptions& o
continue;
}
const auto resource = inst.Flags<uint32_t>();
native_buffer_count = resource < ir.info.buffers.size() &&
ir.info.buffers[resource].dispatch_stride != 0u;
first = resource < ir.info.buffers.size() &&
ir.info.buffers[resource].indirect_root == resource
? 1u
@@ -661,6 +664,7 @@ CompileResult CompileProgram(TranslateResult translated, const CompileOptions& o
continue;
}
for (size_t index = first; index < inst.NumArgs(); index++) {
if (native_buffer_count && index == 2u) continue;
inst.SetArg(index, IR::Value(0u));
}
}
@@ -506,7 +506,15 @@ void StoreLocalFlat(ValueEmitContext& ctx, const IR::Inst& inst) {
void StoreWordPrepared(ValueEmitContext& ctx, const IR::Inst& inst, const IR::MemoryInfo& mem,
const MemoryResourceAccess& resource, uint32_t data) {
const auto index = EmitMemoryElementIndex(ctx.state, resource, DwordIndex(ctx, inst, mem));
EmitIfCondition(ctx.state, EmitMemoryElementInBounds(ctx.state, resource, index), [&]() {
auto in_bounds = EmitMemoryElementInBounds(ctx.state, resource, index);
if (mem.kind == IR::ResourceKind::Buffer &&
ctx.state.program.info.buffers[mem.resource].dispatch_stride != 0u) {
const auto& handle = *inst.Arg(0).ResolveInstruction();
in_bounds = AndCondition(
ctx.state, in_bounds,
Binary(ctx.state, spv::OpULessThan, TypeBool(ctx.state), ctx.Arg(inst, 1), ctx.Arg(handle, 2)));
}
EmitIfCondition(ctx.state, in_bounds, [&]() {
StoreWordInBounds(ctx, resource, index, data);
});
}
@@ -106,6 +106,9 @@ struct BufferResource {
uint32_t source = 0;
uint32_t first_use_pc = 0;
uint32_t max_byte_extent = 0;
// Nonzero proves index = local_x + workgroup_x * dispatch_stride. NUM_RECORDS
// stays on the GPU; host binding capacity comes from the known dispatch extent.
uint32_t dispatch_stride = 0;
uint32_t packed_stride = 0;
Prospero::BufferFormat descriptor_format = Prospero::BufferFormat::kInvalid;
uint32_t descriptor_swizzle = DstSel(4, 5, 6, 7);
@@ -776,8 +776,8 @@ static Value FillValue(Value value, Value guard) {
// 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, Value guard,
uint32_t depth = 0) {
std::optional<std::array<uint64_t, 3>> AffineIndex(Value value, uint32_t axis, Value guard,
uint32_t depth) {
value = FillValue(value, guard);
if (depth > 32 || value.GetType() != Type::U32) {
return {};
@@ -802,8 +802,8 @@ static std::optional<std::array<uint64_t, 3>> FillIndex(Value value, uint32_t ax
op != ValueOpcode::ShiftLeftLogical32) {
return {};
}
auto left = FillIndex(inst->Arg(0), axis, guard, depth + 1);
auto right = FillIndex(inst->Arg(1), axis, guard, depth + 1);
auto left = AffineIndex(inst->Arg(0), axis, guard, depth + 1);
auto right = AffineIndex(inst->Arg(1), axis, guard, depth + 1);
if (!left || !right) {
return {};
}
@@ -870,7 +870,7 @@ static UniformFillPlan AnalyzeUniformFill(const Program& program) {
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, guard);
const auto index = AffineIndex(address->Arg(axis), axis, guard);
if (!index || (*index)[0] != 0) return {};
if (axis < 2) {
if ((*index)[1] != 1 || (*index)[2] == 0) return {};
@@ -900,7 +900,7 @@ static UniformFillPlan AnalyzeUniformFill(const Program& program) {
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, guard);
const auto address = AffineIndex(store->Arg(1), 0, guard);
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]);
@@ -1138,6 +1138,12 @@ bool MaterializeResources(const ResourcePlan& program, const SrtRuntime& runtime
auto packed_stride = descriptor.PackedStride();
const auto stride = packed_stride & 0x3fffu;
const bool swizzle = stride != 0u && ((packed_stride >> 14u) & 1u) != 0u;
if (base.dispatch_stride != 0u && descriptor.Base48() != 0u &&
(descriptor.OutOfBounds() != 0u || stride == 0u || swizzle ||
(packed_stride & (1u << 20u)) != 0u || base.max_byte_extent > stride)) {
return SpecializationFail(fmt::format(
"dispatch-sized buffer {} requires an unswizzled structured DWORD layout", i));
}
if (stride == 0u) {
packed_stride &= ~((1u << 14u) | (3u << 16u));
} else if (!swizzle) {
@@ -43,6 +43,10 @@ struct ResourceSpecialization {
// its values and is independent of the translated shader CFG.
ResourcePlan ExtractResourcePlan(const Program& program);
// Nonnegative constant/local/workgroup coefficients under the instruction's active guard.
std::optional<std::array<uint64_t, 3>> AffineIndex(Value value, uint32_t axis, Value guard,
uint32_t depth = 0);
// Refreshes cached resources and specialization in place. A failed refresh must not be used.
bool MaterializeResources(const ResourcePlan& program, const SrtRuntime& runtime,
ResourceSnapshot& snapshot, ResourceSpecialization& specialization);
@@ -3,6 +3,7 @@
#include "common/assert.h"
#include "graphics/shader/recompiler/ir/ShaderIR.h"
#include "graphics/shader/recompiler/ir/passes/DeadCodeElimination.h"
#include "graphics/shader/recompiler/ir/passes/ResourceMaterialization.h"
#include "graphics/shader/recompiler/ir/passes/SrtWalker.h"
#include <algorithm>
@@ -1955,9 +1956,32 @@ private:
}
}
uint32_t AddBuffer(uint32_t source, const MemoryInfo& memory, ValueOpcode op, uint32_t pc) {
bool TryMakeDispatchBuffer(const Inst& inst, const Inst& handle, uint32_t& source,
uint32_t& dispatch_stride) {
const auto flags = inst.Flags<MemoryFlags>();
const auto& memory = m_program.memory_info[flags.index];
if (m_program.stage != ShaderType::Compute || inst.GetOpcode() != ValueOpcode::StoreBufferU32 ||
memory.kind != ResourceKind::Buffer || memory.typed || memory.formatted ||
!memory.idxen || memory.offen || memory.data_bits != 32u || memory.data_dwords != 1u ||
inst.Arg(2).Resolve() != Value(0u) || inst.Arg(3).Resolve() != Value(0u)) return false;
const auto index = AffineIndex(inst.Arg(1), 0, inst.Arg(5).Resolve());
if (!index || (*index)[0] != 0u || (*index)[1] != 1u || (*index)[2] == 0u) return false;
DescriptorSource descriptor;
MakeSource(handle, 4u, false, false, memory.resource * 4u, descriptor, flags.pc);
// Binding capacity is derived from the dispatch. The native handle keeps its count.
descriptor.dwords[2] = Value(0u);
uint32_t bad_dword = 0;
if (!ValidateSource(descriptor, bad_dword)) return false;
source = InternSource(descriptor);
dispatch_stride = static_cast<uint32_t>((*index)[2]);
return true;
}
uint32_t AddBuffer(uint32_t source, const MemoryInfo& memory, ValueOpcode op, uint32_t pc,
uint32_t dispatch_stride) {
for (uint32_t i = 0; i < m_info.buffers.size(); i++) {
if (m_info.buffers[i].source == source) {
if (m_info.buffers[i].source == source &&
m_info.buffers[i].dispatch_stride == dispatch_stride) {
Merge(m_info.buffers[i], memory, op, pc);
return i;
}
@@ -1968,6 +1992,7 @@ private:
BufferResource resource;
resource.source = source;
resource.first_use_pc = pc;
resource.dispatch_stride = dispatch_stride;
Merge(resource, memory, op, pc);
m_info.buffers.push_back(resource);
return static_cast<uint32_t>(m_info.buffers.size() - 1);
@@ -2123,12 +2148,14 @@ private:
uint32_t resource = 0;
if (buffer != BufferAccess::None) {
uint32_t dispatch_stride = 0;
handle = inst.Arg(0).Resolve().TryInstruction();
const auto* indirect = handle == nullptr ? nullptr : FindIndirectDescriptor(*handle);
if (indirect != nullptr) {
source = indirect->source;
} else if (!GetHandle(inst.Arg(0), ValueOpcode::GetBufferResource, 4, flags.pc,
memory.resource * 4u, handle, source)) {
memory.resource * 4u, handle, source) &&
!TryMakeDispatchBuffer(inst, *handle, source, dispatch_stride)) {
if (memory.kind != (op == ValueOpcode::ReadConstBuffer ? ResourceKind::ScalarBuffer
: ResourceKind::Buffer) ||
!memory.SupportsIndirectBufferLoad(op)) {
@@ -2140,7 +2167,7 @@ private:
m_info.uses_dma = true;
return;
}
resource = AddBuffer(source, memory, op, flags.pc);
resource = AddBuffer(source, memory, op, flags.pc, dispatch_stride);
if (resource == UINT32_MAX) {
Fail(flags.pc, "buffer resource limit exceeded");
}