shader: keep per-invocation mesh exports out of shared memory

This commit is contained in:
nmzik
2026-09-05 15:48:29 +04:00
parent 49ff962f69
commit 9c00c7077b
3 changed files with 103 additions and 17 deletions
@@ -59,6 +59,7 @@ enum : uint32_t {
StorageClassInput = 1,
StorageClassOutput = 3,
StorageClassWorkgroup = 4,
StorageClassPrivate = 6,
StorageClassFunction = 7,
StorageClassPushConstant = 9,
StorageClassImage = 11,
@@ -16,8 +16,9 @@ uint32_t MeshElement(EmitterState& state, uint32_t variable, uint32_t storage, u
return pointer;
}
uint32_t MeshLoad(EmitterState& state, uint32_t variable, uint32_t type, uint32_t index) {
const auto pointer = MeshElement(state, variable, StorageClassWorkgroup, type, index);
uint32_t MeshLoad(EmitterState& state, uint32_t variable, uint32_t storage, uint32_t type,
uint32_t index) {
const auto pointer = MeshElement(state, variable, storage, type, index);
const auto value = state.builder.AllocateId();
state.builder.AddFunction({OpLoad, type, value, pointer});
return value;
@@ -41,8 +42,11 @@ void DefineMeshOutputs(EmitterState& state) {
output.variable_id = MeshArray(
state, StorageClassOutput, type,
output.kind == IR::StageOutputKind::Layer ? mesh.max_primitives : mesh.max_vertices);
output.mesh_data_variable =
MeshArray(state, StorageClassWorkgroup, type, mesh.max_vertices);
// Only Layer is read by another invocation, through the primitive's provoking vertex.
const bool shared = output.kind == IR::StageOutputKind::Layer;
output.mesh_data_variable = MeshArray(
state, shared ? StorageClassWorkgroup : StorageClassPrivate, type,
shared ? mesh.max_vertices : state.lane_count);
state.interface_variables.push_back(output.variable_id);
state.builder.AddName(output.variable_id, output.debug_name.c_str());
if (output.kind == IR::StageOutputKind::Parameter) {
@@ -59,7 +63,7 @@ void DefineMeshOutputs(EmitterState& state) {
}
state.mesh_allocation = MeshArray(state, StorageClassWorkgroup, TypeU32(state), 2);
state.mesh_primitive_data =
MeshArray(state, StorageClassWorkgroup, TypeU32(state), mesh.max_primitives);
MeshArray(state, StorageClassPrivate, TypeU32(state), state.lane_count);
state.mesh_primitives =
MeshArray(state, StorageClassOutput, TypeU32Vector(state, 3), mesh.max_primitives);
state.mesh_cull = MeshArray(state, StorageClassOutput, TypeBool(state), mesh.max_primitives);
@@ -80,13 +84,17 @@ uint32_t MeshOutputPointer(EmitterState& state, IR::StageOutputKind kind, uint32
EXIT("mesh export has no output binding: kind=%u index=%u\n", static_cast<uint32_t>(kind),
index);
}
return MeshElement(state, output->mesh_data_variable, StorageClassWorkgroup,
MeshOutputType(state, kind), EmitLocalInvocationIndex(state));
const bool shared = kind == IR::StageOutputKind::Layer;
return MeshElement(state, output->mesh_data_variable,
shared ? StorageClassWorkgroup : StorageClassPrivate,
MeshOutputType(state, kind),
shared ? EmitLocalInvocationIndex(state)
: ConstantU32(state, state.lane_half));
}
uint32_t MeshPrimitivePointer(EmitterState& state) {
return MeshElement(state, state.mesh_primitive_data, StorageClassWorkgroup, TypeU32(state),
EmitLocalInvocationIndex(state));
return MeshElement(state, state.mesh_primitive_data, StorageClassPrivate, TypeU32(state),
ConstantU32(state, state.lane_half));
}
void EmitMeshAllocate(ValueEmitContext& ctx, const IR::Inst& inst) {
@@ -118,10 +126,10 @@ void EmitMeshEntryPoint(EmitterState& state) {
state.builder.AddFunction(
{OpControlBarrier, ConstantU32(state, ScopeWorkgroup), ConstantU32(state, ScopeWorkgroup),
ConstantU32(state, MemorySemanticsAcquireRelease | MemorySemanticsWorkgroupMemory)});
const auto vertices =
MeshLoad(state, state.mesh_allocation, TypeU32(state), ConstantU32(state, 0));
const auto primitives =
MeshLoad(state, state.mesh_allocation, TypeU32(state), ConstantU32(state, 1));
const auto vertices = MeshLoad(state, state.mesh_allocation, StorageClassWorkgroup,
TypeU32(state), ConstantU32(state, 0));
const auto primitives = MeshLoad(state, state.mesh_allocation, StorageClassWorkgroup,
TypeU32(state), ConstantU32(state, 1));
state.builder.AddFunction({5295u, vertices, primitives}); // OpSetMeshOutputsEXT
for (uint32_t half = 0; half < state.lane_count; half++) {
state.lane_half = half;
@@ -134,7 +142,8 @@ void EmitMeshEntryPoint(EmitterState& state) {
continue;
}
const auto type = MeshOutputType(state, output.kind);
const auto value = MeshLoad(state, output.mesh_data_variable, type, index);
const auto value = MeshLoad(state, output.mesh_data_variable, StorageClassPrivate,
type, ConstantU32(state, half));
const auto pointer =
MeshElement(state, output.variable_id, StorageClassOutput, type, index);
state.builder.AddFunction({OpStore, pointer, value});
@@ -143,7 +152,8 @@ void EmitMeshEntryPoint(EmitterState& state) {
const auto is_primitive = state.builder.AllocateId();
state.builder.AddFunction({OpULessThan, TypeBool(state), is_primitive, index, primitives});
EmitIfCondition(state, is_primitive, [&] {
const auto packed = MeshLoad(state, state.mesh_primitive_data, TypeU32(state), index);
const auto packed = MeshLoad(state, state.mesh_primitive_data, StorageClassPrivate,
TypeU32(state), ConstantU32(state, half));
uint32_t vertex[3] {};
for (uint32_t component = 0; component < 3; component++) {
vertex[component] = state.builder.AllocateId();
@@ -170,8 +180,9 @@ void EmitMeshEntryPoint(EmitterState& state) {
if (output.kind != IR::StageOutputKind::Layer) {
continue;
}
const auto layer = MeshLoad(state, output.mesh_data_variable, TypeU32(state),
vertex[state.input_info.vertex->mesh.provoking_vertex]);
const auto layer = MeshLoad(
state, output.mesh_data_variable, StorageClassWorkgroup, TypeU32(state),
vertex[state.input_info.vertex->mesh.provoking_vertex]);
const auto pointer = MeshElement(state, output.variable_id, StorageClassOutput,
TypeU32(state), index);
state.builder.AddFunction({OpStore, pointer, layer});
+74
View File
@@ -9054,6 +9054,79 @@ void TestFusedShaderHandoffPreservesRegisters() {
Check(allocations == 1u, "fused shader omitted the back shader allocation");
}
void TestMeshExportStorage() {
const uint32_t front[] = {
EncodeSMovB32(12, 255), 0x1003u,
EncodeSop1(0x20, 0, 6), // merged-stage handoff
};
const uint32_t back[] = {
EncodeSMovB32(124, 12), EncodeSopp(0x10, 9), // GS allocation
EncodeVop1(0x01, 0, 128),
EncodeDs0(0x0d), EncodeDs1(0, 0, 0), // retain guest LDS
EncodeExp0(0x0c, 0xf, false), EncodeExp1(0, 0, 0, 0),
EncodeExp0(0x20, 0xf, false), EncodeExp1(0, 0, 0, 0),
EncodeExp0(0x21, 0xf, false), EncodeExp1(0, 0, 0, 0),
EncodeExp0(0x22, 0xf, false), EncodeExp1(0, 0, 0, 0),
EncodeExp0(0x0d, 0x4, false), EncodeExp1(0, 0, 0, 0), // Layer
EncodeExp0(0x14, 0x1), EncodeExp1(0, 0, 0, 0), // primitive
EncodeSopp(0x01),
};
ShaderVertexInputInfo input{};
input.pa_cl_vs_out_cntl = (1u << 21u) | (1u << 18u);
auto &mesh = input.mesh;
mesh.threads_num[0] = 192;
mesh.threads_num[1] = mesh.threads_num[2] = 1;
mesh.lds_size_dwords = 3840;
mesh.primitives_per_group = 62;
mesh.vertices_per_group = 64;
mesh.max_vertices = 192;
mesh.max_primitives = 176;
ShaderRecompiler::CompileOptions options{};
options.stage = ShaderType::Mesh;
options.input_info.vertex = &input;
options.back_code = back;
for (const uint32_t subgroup_size : {32u, 64u}) {
mesh.host_subgroup_size = subgroup_size;
const auto result = RecompileForTest(front, options);
CheckSpirvBinaryValidates(result.spirv);
const auto &binary = result.spirv;
std::vector<uint32_t> sizes(binary[3]), constants(binary[3]);
uint32_t shared_bytes = 0, private_bytes = 0;
for (size_t i = 5; i < binary.size(); i += binary[i] >> 16u) {
switch (binary[i] & 0xffffu) {
case 21u: // OpTypeInt
case 22u: // OpTypeFloat
sizes[binary[i + 1]] = binary[i + 2] / 8u;
break;
case 23u: // OpTypeVector
sizes[binary[i + 1]] = sizes[binary[i + 2]] * binary[i + 3];
break;
case 28u: // OpTypeArray
sizes[binary[i + 1]] = sizes[binary[i + 2]] * constants[binary[i + 3]];
break;
case 32u: // OpTypePointer
sizes[binary[i + 1]] = sizes[binary[i + 3]];
break;
case 43u: // OpConstant
constants[binary[i + 2]] = binary[i + 3];
break;
case 59u: // OpVariable
if (binary[i + 3] == 4u || binary[i + 3] == 6u) {
Check(sizes[binary[i + 1]] != 0, "unmeasured mesh staging type");
(binary[i + 3] == 4u ? shared_bytes : private_bytes) += sizes[binary[i + 1]];
}
break;
}
}
// The Pathless shader 35869c17ce783c88 exceeds the host's 28 KiB budget
// when its four vertex exports and primitive exports are shared arrays.
Check(shared_bytes == 3840u * 4u + 192u * 4u + 8u && shared_bytes <= 28672u,
"mesh staging must retain guest LDS, shared Layer and allocation within the host budget");
Check(private_bytes == (4u * 16u + 4u) * (64u / subgroup_size),
"mesh vertex and primitive exports lost their separate logical-lane storage");
}
}
void TestMergedShaderUserDataSnapshot() {
using namespace ShaderRecompiler;
const uint32_t front[] = {
@@ -12482,6 +12555,7 @@ int main() {
TestNewShaderRecompilerBranchConditionForms();
TestNewShaderRecompilerSetpcBranch();
TestFusedShaderHandoffPreservesRegisters();
TestMeshExportStorage();
TestMergedShaderUserDataSnapshot();
TestMeshInputAssembly();
TestNewShaderRecompilerSetpcJumpTable();