shader_recompiler: gl_PrimitiveId scatter for geometry-stage stream compaction (#4330)

Guest shaders that reserve a compacted output slot via subgroup ballot
-> popcount -> atomic add cannot preserve the originating primitives'
relative order once run through the single-invocation subgroup
fallback for stages without subgroup support. gl_PrimitiveId is
already hardware-guaranteed unique and monotonically increasing per
primitive, so CompactionFallbackPass rewrites that pattern into a
gl_PrimitiveId-indexed scatter instead.

Confirmed fixing corrupted/order-scrambled compacted draws in NieR
Automata on a device without geometry-stage subgroup support.

Reviewed-on: https://git.eden-emu.dev/eden-emu/eden/pulls/4330

src/shader_recompiler/CMakeLists.txt

@@ -217,6 +217,7 @@ add_library(shader_recompiler STATIC
     frontend/maxwell/translate_program.h
     host_translate_info.h
     ir_opt/collect_shader_info_pass.cpp
+    ir_opt/compaction_fallback_pass.cpp
     ir_opt/conditional_barrier_pass.cpp
     ir_opt/constant_propagation_pass.cpp
     ir_opt/dead_code_elimination_pass.cpp

src/shader_recompiler/backend/spirv/emit_spirv_warp.cpp

@@ -217,6 +217,10 @@ Id EmitShuffleIndex(EmitContext& ctx, IR::Inst* inst, Id value, Id index, Id cla
 
 Id EmitShuffleUp(EmitContext& ctx, IR::Inst* inst, Id value, Id index, Id clamp,
                  Id segmentation_mask) {
+    if (!StageSupportsSubgroups(ctx)) {
+        SetInBoundsFlag(inst, ctx.false_value);
+        return value;
+    }
     const Id thread_id{EmitLaneId(ctx)};
     const Id max_thread_id{GetMaxThreadId(ctx, thread_id, clamp, segmentation_mask)};
     Id src_thread_id{ctx.OpISub(ctx.U32[1], thread_id, index)};
@@ -232,6 +236,10 @@ Id EmitShuffleUp(EmitContext& ctx, IR::Inst* inst, Id value, Id index, Id clamp,
 
 Id EmitShuffleDown(EmitContext& ctx, IR::Inst* inst, Id value, Id index, Id clamp,
                    Id segmentation_mask) {
+    if (!StageSupportsSubgroups(ctx)) {
+        SetInBoundsFlag(inst, ctx.false_value);
+        return value;
+    }
     const Id thread_id{EmitLaneId(ctx)};
     const Id max_thread_id{GetMaxThreadId(ctx, thread_id, clamp, segmentation_mask)};
     Id src_thread_id{ctx.OpIAdd(ctx.U32[1], thread_id, index)};
@@ -247,6 +255,10 @@ Id EmitShuffleDown(EmitContext& ctx, IR::Inst* inst, Id value, Id index, Id clam
 
 Id EmitShuffleButterfly(EmitContext& ctx, IR::Inst* inst, Id value, Id index, Id clamp,
                         Id segmentation_mask) {
+    if (!StageSupportsSubgroups(ctx)) {
+        SetInBoundsFlag(inst, ctx.false_value);
+        return value;
+    }
     const Id thread_id{EmitLaneId(ctx)};
     const Id max_thread_id{GetMaxThreadId(ctx, thread_id, clamp, segmentation_mask)};
     Id src_thread_id{ctx.OpBitwiseXor(ctx.U32[1], thread_id, index)};

src/shader_recompiler/frontend/maxwell/translate_program.cpp

@@ -299,6 +299,7 @@ IR::Program TranslateProgram(ObjectPool<IR::Inst>& inst_pool, ObjectPool<IR::Blo
     Optimization::PositionPass(env, program);
 
     Optimization::GlobalMemoryToStorageBufferPass(program, normalized_host_info);
+    Optimization::CompactionFallbackPass(program, normalized_host_info);
     Optimization::TexturePass(env, program, normalized_host_info);
 
     if (Settings::values.resolution_info.active || Settings::values.rescale_hack.GetValue()) {

src/shader_recompiler/host_translate_info.h

@@ -38,6 +38,7 @@ struct HostTranslateInfo {
                                                 ///< passthrough shaders
     bool support_conditional_barrier{}; ///< True when the device supports barriers in conditional
                                         ///< control flow
+    bool support_subgroup_in_geometry_stage{};
 
     void ApplyDescriptorLimitPolicy() noexcept {
         if (min_ssbo_alignment == 0) {

src/shader_recompiler/ir_opt/compaction_fallback_pass.cpp

@@ -0,0 +1,69 @@
+// SPDX-FileCopyrightText: Copyright 2026 Eden Emulator Project
+// SPDX-License-Identifier: GPL-3.0-or-later
+
+#include "shader_recompiler/frontend/ir/ir_emitter.h"
+#include "shader_recompiler/host_translate_info.h"
+#include "shader_recompiler/ir_opt/passes.h"
+
+namespace Shader::Optimization {
+namespace {
+
+constexpr int MAX_BALLOT_SEARCH_DEPTH = 8;
+
+bool DependsOnSubgroupBallot(const IR::Value& value, int depth) {
+    if (depth > MAX_BALLOT_SEARCH_DEPTH || value.IsImmediate()) {
+        return false;
+    }
+    IR::Inst* const producer{value.Inst()};
+    if (producer->GetOpcode() == IR::Opcode::SubgroupBallot) {
+        return true;
+    }
+    if (producer->GetOpcode() == IR::Opcode::Phi) {
+        return false;
+    }
+    const size_t num_args{producer->NumArgs()};
+    for (size_t index = 0; index < num_args; ++index) {
+        if (DependsOnSubgroupBallot(producer->Arg(index), depth + 1)) {
+            return true;
+        }
+    }
+    return false;
+}
+
+void RewriteAtomic(IR::Block& block, IR::Inst& inst) {
+    const auto insert_point{IR::Block::InstructionList::s_iterator_to(inst)};
+    IR::IREmitter ir{block, insert_point};
+
+    const IR::U32 amount{inst.Arg(2)};
+    const IR::U32 primitive_id{ir.GetAttributeU32(IR::Attribute::PrimitiveId)};
+    const IR::U32 new_index{ir.IMul(primitive_id, amount)};
+    const IR::U32 new_atomic_value{ir.IAdd(new_index, amount)};
+
+    const IR::Value umax_result{&*block.PrependNewInst(
+        insert_point, IR::Opcode::StorageAtomicUMax32,
+        {inst.Arg(0), inst.Arg(1), new_atomic_value})};
+    static_cast<void>(umax_result);
+
+    inst.ReplaceUsesWith(new_index);
+}
+
+} // Anonymous namespace
+
+void CompactionFallbackPass(IR::Program& program, const HostTranslateInfo& host_info) {
+    if (program.stage != Stage::Geometry || host_info.support_subgroup_in_geometry_stage) {
+        return;
+    }
+    for (IR::Block* const block : program.post_order_blocks) {
+        for (IR::Inst& inst : block->Instructions()) {
+            if (inst.GetOpcode() != IR::Opcode::StorageAtomicIAdd32) {
+                continue;
+            }
+            if (!DependsOnSubgroupBallot(inst.Arg(2), 0)) {
+                continue;
+            }
+            RewriteAtomic(*block, inst);
+        }
+    }
+}
+
+} // namespace Shader::Optimization

src/shader_recompiler/ir_opt/passes.h

@@ -13,6 +13,7 @@ struct HostTranslateInfo;
 namespace Shader::Optimization {
 
 void CollectShaderInfoPass(Environment& env, IR::Program& program);
+void CompactionFallbackPass(IR::Program& program, const HostTranslateInfo& host_info);
 void ConditionalBarrierPass(IR::Program& program);
 void ConstantPropagationPass(Environment& env, IR::Program& program);
 void DeadCodeEliminationPass(IR::Program& program);

src/video_core/renderer_opengl/gl_shader_cache.cpp

@@ -269,6 +269,7 @@ ShaderCache::ShaderCache(Tegra::MaxwellDeviceMemoryManager& device_memory_,
         .support_viewport_index_layer = device.HasVertexViewportLayer(),
         .support_geometry_shader_passthrough = device.HasGeometryShaderPassthrough(),
         .support_conditional_barrier = device.SupportsConditionalBarriers(),
+        .support_subgroup_in_geometry_stage = true,
       } {
     host_info.ApplyDescriptorLimitPolicy();
     if (use_asynchronous_shaders) {

src/video_core/renderer_vulkan/vk_pipeline_cache.cpp

@@ -479,6 +479,7 @@ PipelineCache::PipelineCache(Tegra::MaxwellDeviceMemoryManager& device_memory_,
         .support_viewport_index_layer = device.IsExtShaderViewportIndexLayerSupported(),
         .support_geometry_shader_passthrough = device.IsNvGeometryShaderPassthroughSupported(),
         .support_conditional_barrier = device.SupportsConditionalBarriers(),
+        .support_subgroup_in_geometry_stage = profile.SupportsSubgroupStage(Shader::Stage::Geometry),
     };
     host_info.ApplyDescriptorLimitPolicy();