shader: initialize EXEC from requested dispatch thread extents

Co-authored-by: TheCruZ <thecruzwtf@gmail.com>
Co-authored-by: Mac <me@itsmemac.com>
This commit is contained in:
nmzik
2026-10-04 01:15:39 +04:00
parent f58014ff0e
commit e4aca7a60b
9 changed files with 147 additions and 12 deletions
@@ -302,6 +302,10 @@ void RenderExecutor::DispatchDirect(uint64_t submit_id, CommandBuffer& buffer,
m_context.GetPipelineCache().GetComputePipeline(input_info, compute_program);
auto& bindings = m_compute_bindings;
PrepareBindings(input_info.stage, bindings);
if (program.bindings.dispatch_thread_dword != ShaderRecompiler::IR::PushData::NoStart) {
std::copy(std::begin(input_info.dispatch_threads_num), std::end(input_info.dispatch_threads_num),
bindings.shader_data.begin() + program.bindings.dispatch_thread_dword);
}
FindBuffers(bindings);
if (program.info.uses_dma) {
m_context.PrepareBda();
@@ -85,9 +85,13 @@ void ValidateNativeProgram(const IR::Program& program) {
}
const auto has_shader_data_storage = present[static_cast<size_t>(Kind::ShaderData)];
const auto shader_data_dwords = program.bindings.ShaderDataDwords();
const auto user_data_dwords = program.bindings.user_data_registers.size();
const bool has_dispatch_threads = program.bindings.dispatch_thread_dword != IR::PushData::NoStart;
if ((program.bindings.UsesPushData() &&
!IR::PushData::CanFit(program.bindings.push_data_start_dword, shader_data_dwords)) ||
program.bindings.memory_offset_dword != program.bindings.user_data_registers.size() ||
(has_dispatch_threads && (program.stage != ShaderType::Compute ||
program.bindings.dispatch_thread_dword != user_data_dwords)) ||
program.bindings.memory_offset_dword != user_data_dwords + (has_dispatch_threads ? 3u : 0u) ||
program.bindings.memory_offset_count != buffers.size() ||
has_shader_data_storage != (shader_data_dwords != 0 && !program.bindings.UsesPushData()) ||
!std::is_sorted(program.bindings.user_data_registers.begin(),
@@ -631,6 +631,14 @@ uint32_t EmitGetBuiltin(ValueEmitContext& ctx, IR::Value kind, IR::Value index)
return EmitBuiltinU32(ctx.state, static_cast<IR::StageInputKind>(kind.U32()), index.U32());
}
uint32_t EmitGetDispatchThreadExtent(ValueEmitContext& ctx, const IR::Inst& inst) {
const auto start = ctx.state.program.bindings.dispatch_thread_dword;
if (start == IR::PushData::NoStart || !inst.Arg(0).IsImmediate() || inst.Arg(0).U32() >= 3u) {
ctx.Fail(inst, "invalid dispatch thread extent");
}
return EmitShaderDataDwordLoad(ctx.state, start + inst.Arg(0).U32());
}
uint32_t EmitUndefU1(EmitterState& state, const IR::Inst& inst) {
const auto result = state.builder.AllocateId();
state.builder.AddFunction(spv::OpUndef, TypeId(state, inst.GetType()), result);
@@ -198,6 +198,7 @@ uint32_t EmitMeshDrawParameter(ValueEmitContext& ctx, const IR::Ins
uint32_t EmitGetTessellationAttribute(ValueEmitContext& ctx, const IR::Inst& inst);
void EmitSetTessellationAttribute(ValueEmitContext& ctx, const IR::Inst& inst);
uint32_t EmitGetUserData(EmitterState& state, IR::ScalarReg reg);
uint32_t EmitGetDispatchThreadExtent(ValueEmitContext& ctx, const IR::Inst& inst);
uint32_t EmitGetBuiltin(ValueEmitContext& ctx, IR::Value kind, IR::Value index);
uint32_t EmitUndefU1(EmitterState& state, const IR::Inst& inst);
inline constexpr auto EmitUndefU8 = EmitUndefU1;
@@ -1104,6 +1104,16 @@ IR::Program TranslateProgram(const Decoder::Program& decoded, const CFG::Graph&
IR::U32(IR::Value(total_threads)));
}
}
if (options.stage == ShaderType::Compute &&
options.input_info.compute->dispatch_thread_dimensions) {
for (uint32_t axis = 0; axis < 3u; axis++) {
const auto extent = IR::U32(entry_ir.Emit(
IR::ValueOpcode::GetDispatchThreadExtent, {IR::Value(axis)}));
initial_exec = entry_ir.LogicalAnd(
initial_exec,
entry_ir.ULessThan(builtin(IR::StageInputKind::GlobalInvocationId, axis), extent));
}
}
entry_ir.SetExec(initial_exec);
const auto initial_mask = entry_ir.Emit(IR::ValueOpcode::Ballot, {initial_exec});
entry_ir.SetExecLo(entry_ir.CompositeExtract(initial_mask, 0));
@@ -426,6 +426,7 @@ struct DescriptorBinding {
struct BindingLayout {
uint32_t push_data_start_dword = PushData::NoStart;
uint32_t dispatch_thread_dword = PushData::NoStart;
uint32_t memory_offset_dword = 0;
uint32_t memory_offset_count = 0;
std::vector<uint32_t> user_data_registers;
@@ -8,6 +8,7 @@ VALUE_OPCODE(ConditionRef, U1, U1)
// Translation and register-state pseudo operations.
VALUE_OPCODE(GetUserData, U32, ScalarReg)
VALUE_OPCODE(GetDispatchThreadExtent, U32, U32)
VALUE_OPCODE(GetShaderBase, U64)
VALUE_OPCODE(MeshDrawParameter, U32, U32)
VALUE_OPCODE(MeshAllocate, Void, U32)
@@ -15,11 +15,16 @@ namespace {
std::abort();
}
std::vector<uint32_t> CollectUserData(const Program& program) {
void CollectShaderData(const Program& program, BindingLayout& layout) {
std::array<bool, NumScalarRegs> registers {};
bool uses_dispatch_threads = false;
for (const auto* block: program.blocks) {
for (const auto& inst: *block) {
if (inst.GetOpcode() != ValueOpcode::GetUserData || !inst.HasUses()) {
if (!inst.HasUses()) {
continue;
}
uses_dispatch_threads |= inst.GetOpcode() == ValueOpcode::GetDispatchThreadExtent;
if (inst.GetOpcode() != ValueOpcode::GetUserData) {
continue;
}
if (inst.Arg(0).GetType() != Type::ScalarReg) {
@@ -32,13 +37,16 @@ std::vector<uint32_t> CollectUserData(const Program& program) {
registers[index] = true;
}
}
std::vector<uint32_t> result;
for (uint32_t index = 0; index < registers.size(); index++) {
if (registers[index]) {
result.push_back(index);
layout.user_data_registers.push_back(index);
}
}
return result;
layout.memory_offset_dword = static_cast<uint32_t>(layout.user_data_registers.size());
if (uses_dispatch_threads) {
layout.dispatch_thread_dword = layout.memory_offset_dword;
layout.memory_offset_dword += 3u;
}
}
void AddBinding(BindingLayout& layout, DescriptorBindingKind kind,
@@ -109,8 +117,7 @@ void AllocateBindings(Program& program, uint32_t push_data_start_dword) {
BindingLayout next;
std::vector<uint32_t> buffers;
const bool uses_gds = CollectMemoryResources(program, buffers);
next.user_data_registers = CollectUserData(program);
next.memory_offset_dword = static_cast<uint32_t>(next.user_data_registers.size());
CollectShaderData(program, next);
next.memory_offset_count = static_cast<uint32_t>(buffers.size());
next.push_data_start_dword =
PushData::StartFor(push_data_start_dword, next.ShaderDataDwords());
+103 -4
View File
@@ -14152,6 +14152,13 @@ public:
vk::Sampler sampler = nullptr) {
using Kind = ShaderRecompiler::IR::DescriptorBindingKind;
const auto &layout = compiled.program.bindings;
auto shader_data = compiled.packed_user_data;
if (layout.dispatch_thread_dword != ShaderRecompiler::IR::PushData::NoStart) {
Require(test.name, "dispatch dimensions", test.has_compute_info,
"runtime thread extents are required by the compiled shader");
std::copy_n(test.compute_info.dispatch_threads_num, 3,
shader_data.begin() + layout.dispatch_thread_dword);
}
auto Binding = [&](Kind kind) {
return ShaderRecompiler::IR::FindBinding(layout, kind);
};
@@ -14378,8 +14385,7 @@ public:
}
if (const auto *user = Binding(Kind::ShaderData); user != nullptr) {
user_data_buffer =
CreateStorageBuffer(test.name, compiled.packed_user_data,
compiled.packed_user_data.size());
CreateStorageBuffer(test.name, shader_data, shader_data.size());
user_data_info = {user_data_buffer.buffer, 0, user_data_buffer.size};
vk::WriteDescriptorSet write{};
write.sType = vk::StructureType::eWriteDescriptorSet;
@@ -14574,8 +14580,7 @@ public:
1, &descriptor_set, 0, nullptr);
if (layout.UsesPushData()) {
ShaderRecompiler::IR::PushData push_data;
std::copy(compiled.packed_user_data.begin(),
compiled.packed_user_data.end(),
std::copy(shader_data.begin(), shader_data.end(),
push_data.dwords.begin() + layout.push_data_start_dword);
cmd.pushConstants(pipeline_layout, vk::ShaderStageFlagBits::eCompute, 0,
sizeof(push_data), push_data.dwords.data());
@@ -30601,6 +30606,94 @@ TestCase ImageD16StoreUnpacksHalfPairs() {
return test;
}
void CheckComputeThreadDimensions(VulkanHarness &vulkan) {
using O = ShaderOpcode;
constexpr u32 grid_size = 16 * 8 * 4;
const std::array extents = {std::array{9u, 5u, 3u}, std::array{13u, 7u, 4u},
std::array{3u, 2u, 1u}, std::array{16u, 8u, 4u}};
for (const u32 wave_size : {32u, 64u}) {
for (const bool thread_dimensions : {false, true}) {
TestCase test;
test.name = thread_dimensions ? (wave_size == 32 ? "ThreadDimensions32" : "ThreadDimensions64")
: (wave_size == 32 ? "GroupDimensions32" : "GroupDimensions64");
auto &code = test.code;
code.push_back(EncodeSop1(0x04, 20, 126)); // Save the SPI-provided initial EXEC.
AppendSMovLiteral(&code, 24, grid_size * 4);
for (u32 axis = 0; axis < 3; ++axis) {
code.push_back(EncodeVop1(0x01, 3 + axis, 4 + axis));
code.push_back(EncodeVop2(0x1a, 3 + axis, InlineU32(3 - axis), 3 + axis));
code.push_back(EncodeVop2(0x25, 3 + axis, Vgpr(axis), 3 + axis));
}
code.push_back(EncodeVop2(0x1a, 4, InlineU32(4), 4));
code.push_back(EncodeVop2(0x1a, 5, InlineU32(7), 5));
code.push_back(EncodeVop2(0x25, 3, Vgpr(4), 3));
code.push_back(EncodeVop2(0x25, 3, Vgpr(5), 3));
code.push_back(EncodeVop2(0x25, 20, InlineU32(1), 3));
code.push_back(EncodeVop2(0x1a, 20, InlineU32(2), 20));
AppendVMovU32(&code, 6, 1);
AppendVMovU32(&code, 7, 0);
AppendBufferStoreDword(&code, 6, 20);
AppendBufferStoreOpcode(&code, 0x32, 6, 7);
code.push_back(EncodeSop2(0x0f, 126, 20, InlineU32(1)));
AppendVMovU32(&code, 6, 2); // Divergence changes only lane zero of each wave.
code.push_back(EncodeSop1(0x04, 126, 20));
code.push_back(EncodeVop2(0x25, 20, 24, 20));
AppendBufferStoreDword(&code, 6, 20);
AppendVMovU32(&code, 8, 1);
AppendBufferStoreOpcode(&code, 0x32, 8, 7);
AppendEnd(&code);
test.initial.resize(1 + 2 * grid_size);
test.opcodes = {O::S_MOV_B64, O::S_MOV_B32, O::S_AND_B64, O::V_MOV_B32,
O::V_LSHLREV_B32, O::V_ADD_NC_U32, O::BUFFER_STORE_DWORD,
O::BUFFER_ATOMIC_ADD, O::S_ENDPGM};
test.compute_info.wave_size = wave_size;
test.compute_info.host_subgroup_size = vulkan.SubgroupSize();
test.compute_info.threads_num[0] = 8;
test.compute_info.threads_num[1] = 4;
test.compute_info.threads_num[2] = 2;
std::fill_n(test.compute_info.group_id, 3, true);
test.compute_info.thread_ids_num = 3;
test.compute_info.workgroup_register = 4;
test.compute_info.dispatch_thread_dimensions = thread_dimensions;
test.has_compute_info = true;
test.expand_shader_data_storage = thread_dimensions && wave_size == 32;
const auto compiled = CompileCase(test, vulkan.SubgroupSize());
std::vector<u32> compiled_key;
BuildStageStaticKey(test.compute_info, compiled_key);
for (const auto &extent : extents) {
std::copy(extent.begin(), extent.end(), test.compute_info.dispatch_threads_num);
std::vector<u32> runtime_key;
BuildStageStaticKey(test.compute_info, runtime_key);
Require(test.name, "runtime shader reuse", compiled_key == runtime_key,
"dispatch thread counts changed the compiled shader cache key");
test.dispatch_x = thread_dimensions ? (extent[0] + 7) / 8 : 2;
test.dispatch_y = thread_dimensions ? (extent[1] + 3) / 4 : 2;
test.dispatch_z = thread_dimensions ? (extent[2] + 1) / 2 : 2;
test.expected.assign(test.initial.size(), 0);
for (u32 z = 0; z < 4; ++z) {
for (u32 y = 0; y < 8; ++y) {
for (u32 x = 0; x < 16; ++x) {
if (thread_dimensions && (x >= extent[0] || y >= extent[1] || z >= extent[2])) continue;
const auto index = x + 16 * y + 128 * z;
const auto local_index = x % 8 + 8 * (y % 4) + 32 * (z % 2);
test.expected[0] += 2;
test.expected[1 + index] = 1;
test.expected[1 + grid_size + index] = local_index % wave_size == 0 ? 2 : 1;
}
}
}
auto output = vulkan.CreateStorageBuffer(test.name, test.initial, test.initial.size());
vulkan.Dispatch(test, compiled, output);
const auto actual = vulkan.ReadBuffer(test.name, output, test.expected.size());
vulkan.DestroyBuffer(&output);
CompareWords(test, "initial EXEC, restored EXEC and padded-lane atomics", test.expected, actual);
if (!thread_dimensions) break;
}
std::printf("[compute] %-32s ok\n", test.name);
}
}
}
void CheckIndirectBufferStore(VulkanHarness &vulkan) {
using namespace ShaderRecompiler::IR;
TestCase test;
@@ -37246,6 +37339,11 @@ int main(int argc, char **argv) {
std::setvbuf(stdout, nullptr, _IONBF, 0);
EnsureConfigInitialized();
CheckLeastRecentlyUsedCacheOrdering();
if (argc == 2 && std::strcmp(argv[1], "--thread-dimensions-only") == 0) {
VulkanHarness vulkan;
CheckComputeThreadDimensions(vulkan);
return 0;
}
if (argc == 2 && std::strcmp(argv[1], "--lds-limit-only") == 0) {
VulkanHarness vulkan;
CheckComputeLdsLimit(vulkan);
@@ -38006,6 +38104,7 @@ int main(int argc, char **argv) {
CheckPixelParameterAliases();
CheckRectListShaders();
CheckIndirectBufferStore(vulkan);
CheckComputeThreadDimensions(vulkan);
CheckIndirectImageKeySwitch();
CheckWave64WholeWaveResults();
CheckPs5GameExampleImageClearRuntimeShape();