shader: lower whole-wave OR scans to subgroup reductions

src/graphics/shader/recompiler/backend/spirv/spirvEmitterFlow.cpp

@@ -859,6 +859,10 @@ uint32_t EmitSubgroupUMaxU32(ValueEmitContext& ctx, const IR::Inst& inst) {
 	return EmitSubgroupReduceU32<spv::OpGroupNonUniformUMax, EmitUMax32>(ctx, inst);
 }
 
+uint32_t EmitSubgroupOrU32(ValueEmitContext& ctx, const IR::Inst& inst) {
+	return EmitSubgroupReduceU32<spv::OpGroupNonUniformBitwiseOr, EmitOrU32>(ctx, inst);
+}
+
 uint32_t EmitWriteLane(ValueEmitContext& ctx, const IR::Inst& inst) {
 	auto&      state = ctx.state;
 	const auto hit   = state.builder.AllocateId();

src/graphics/shader/recompiler/backend/spirv/spirvEmitterInstructions.h

@@ -236,6 +236,7 @@ uint32_t              EmitReadFirstLane(ValueEmitContext& ctx, const IR::Inst& i
 uint32_t              EmitReadLane(ValueEmitContext& ctx, const IR::Inst& inst);
 uint32_t              EmitSubgroupUMinU32(ValueEmitContext& ctx, const IR::Inst& inst);
 uint32_t              EmitSubgroupUMaxU32(ValueEmitContext& ctx, const IR::Inst& inst);
+uint32_t              EmitSubgroupOrU32(ValueEmitContext& ctx, const IR::Inst& inst);
 uint32_t              EmitWriteLane(ValueEmitContext& ctx, const IR::Inst& inst);
 uint32_t              EmitPermlane16U32(ValueEmitContext& ctx, const IR::Inst& inst);
 uint32_t              EmitGetAttribute(ValueEmitContext& ctx, const IR::Inst& inst);

src/graphics/shader/recompiler/backend/spirv/spirvEmitterProgram.cpp

@@ -235,7 +235,8 @@ void EmitDirectInstruction(ValueEmitContext& ctx, const IR::Inst& inst) {
 	const auto op = inst.GetOpcode();
 	if (ctx.half != 0 &&
 	    (op == Op::Ballot || op == Op::ReadFirstLane || op == Op::SubgroupUMinU32 ||
-	     op == Op::SubgroupUMaxU32 || op == Op::ReadConstBuffer || op == Op::ReadRealtimeClock ||
+	     op == Op::SubgroupUMaxU32 || op == Op::SubgroupOrU32 || op == Op::ReadConstBuffer ||
+	     op == Op::ReadRealtimeClock ||
 	     (op == Op::LoadAddressU32 && ctx.Memory(inst).kind == IR::ResourceKind::ScalarAddress))) {
 		// Whole-wave results, clocks and SMEM loads (SGPR operands, no EXEC) match in both halves.
 		ctx.Define(inst, ctx.other_half->Result(inst));

src/graphics/shader/recompiler/ir/opcodes/ValueOpcodes.inc

@@ -224,6 +224,7 @@ VALUE_OPCODE(ReadFirstLane, U32, U32, U1)
 VALUE_OPCODE(ReadLane, U32, U32, U32)
 VALUE_OPCODE(SubgroupUMinU32, U32, U32)
 VALUE_OPCODE(SubgroupUMaxU32, U32, U32)
+VALUE_OPCODE(SubgroupOrU32, U32, U32)
 VALUE_OPCODE(WriteLane, U32, U32, U32, U32)
 VALUE_OPCODE(Permlane16U32, U32, U32, U32, U32, U1, U32)
 VALUE_OPCODE(PermuteU32, U32, U32, U32, U1)

src/graphics/shader/recompiler/ir/passes/ReadLaneElimination.cpp

@@ -1,7 +1,7 @@
 #include "graphics/shader/recompiler/ir/passes/ReadLaneElimination.h"
 
 #include <algorithm>
-
+#include <optional>
 #include <queue>
 
 #include <unordered_map>
@@ -186,6 +186,16 @@ private:
 	ValueOpcode  operation;
 };
 
+// The subgroup reduction that replaces a proven whole-wave scan of operation.
+std::optional<ValueOpcode> WaveReduction(ValueOpcode operation) {
+	switch (operation) {
+		case ValueOpcode::UMin32: return ValueOpcode::SubgroupUMinU32;
+		case ValueOpcode::UMax32: return ValueOpcode::SubgroupUMaxU32;
+		case ValueOpcode::BitwiseOr32: return ValueOpcode::SubgroupOrU32;
+		default: return std::nullopt;
+	}
+}
+
 } // namespace
 
 ReadLaneStats EliminateReadLane(Program& program, uint32_t wave_size) {
@@ -197,15 +207,12 @@ ReadLaneStats EliminateReadLane(Program& program, uint32_t wave_size) {
 	for (auto* block: program.blocks) {
 		for (auto& inst: *block) {
 			const auto opcode = inst.GetOpcode();
-			if (opcode == ValueOpcode::UMin32 || opcode == ValueOpcode::UMax32) {
+			if (const auto reduce = WaveReduction(opcode)) {
 				const auto seed = WaveReductionMatcher(block, opcode).Match(inst, wave_size);
 				if (!seed.IsEmpty()) {
 					const auto position = std::find_if(block->begin(), block->end(),
 					                                  [&](const Inst& candidate) { return &candidate == &inst; });
-					const auto reduce    = opcode == ValueOpcode::UMin32
+					auto&      reduction = *block->PrependNewInst(position, *reduce, {seed});
-					                           ? ValueOpcode::SubgroupUMinU32
-					                           : ValueOpcode::SubgroupUMaxU32;
-					auto&      reduction = *block->PrependNewInst(position, reduce, {seed});
 					inst.ReplaceUsesWith(Value(&reduction));
 					stats.rewritten_reductions++;
 				}

src/graphics/shader/recompiler/ir/passes/ShaderInfoCollection.cpp

@@ -524,7 +524,8 @@ void Visit(Program& program, ShaderStageInputInfo input_info, InputUsage& inputs
 		case ValueOpcode::Ballot: info.subgroup_ballot = true; break;
 		case ValueOpcode::LiveMask: info.subgroup_ballot = info.helper_invocation = true; break;
 		case ValueOpcode::SubgroupUMinU32:
-		case ValueOpcode::SubgroupUMaxU32: info.subgroup_arithmetic = true; break;
+		case ValueOpcode::SubgroupUMaxU32:
+		case ValueOpcode::SubgroupOrU32: info.subgroup_arithmetic = true; break;
 		case ValueOpcode::DppMoveU32:
 		case ValueOpcode::ReadFirstLane:
 		case ValueOpcode::ReadLane: {

tests/ScalarProvenanceTests.cpp

@@ -695,7 +695,7 @@ void TestCompleteWaveReduction() {
     Shift, FetchInactive, BoundControl, Dpp8, MoveExec, UpdateExec,
     Previous, DppSource, CrossExec, CrossRow, CrossSelectorLow,
     CrossSelectorHigh, CrossFetch, CrossBound, CrossSource, FinalExec,
-    Endpoint, EndpointSource, CrossBlock, WrongOperation
+    Endpoint, EndpointSource, CrossBlock, WrongOperation, Or
   };
   for (const auto test : {Case::Masked, Case::ScalarExec, Case::Wave32,
        Case::Wave32Rows, Case::Wave32Maximum, Case::RowMask, Case::BankMask,
@@ -704,7 +704,7 @@ void TestCompleteWaveReduction() {
        Case::CrossExec, Case::CrossRow, Case::CrossSelectorLow,
        Case::CrossSelectorHigh, Case::CrossFetch, Case::CrossBound,
        Case::CrossSource, Case::FinalExec, Case::Endpoint, Case::EndpointSource,
-       Case::CrossBlock, Case::WrongOperation}) {
+       Case::CrossBlock, Case::WrongOperation, Case::Or}) {
     Fixture fixture(test == Case::CrossBlock ? 2 : 1);
     const auto source = fixture.Emit(ValueOpcode::LaneId);
     const auto mask = fixture.Emit(ValueOpcode::ULessThan32, {source, Value(19u)});
@@ -720,14 +720,20 @@ void TestCompleteWaveReduction() {
       enabled = fixture.Emit(ValueOpcode::INotEqual32,
           {fixture.Emit(ValueOpcode::BitwiseAnd32, {shifted, Value(1u)}), Value(0u)});
     }
-    // A UMax scan seeds lanes outside EXEC with zero.
+    // UMax and OR scans seed lanes outside EXEC with zero.
     const bool maximum = test == Case::Wave32Maximum;
-    const uint32_t neutral = maximum ? 0u : UINT32_MAX;
+    const bool bitwise_or = test == Case::Or;
+    const uint32_t neutral = maximum || bitwise_or ? 0u : UINT32_MAX;
     const auto seed = fixture.Emit(ValueOpcode::SelectU32, {mask, source, Value(neutral)});
     auto prefix = fixture.Emit(ValueOpcode::SelectU32, {enabled, seed, source});
     Value preserved;
     auto operation = maximum ? ValueOpcode::UMax32 : ValueOpcode::UMin32;
-    if (test == Case::WrongOperation) operation = ValueOpcode::BitwiseOr32;
+    auto reduction_opcode = maximum ? ValueOpcode::SubgroupUMaxU32 : ValueOpcode::SubgroupUMinU32;
+    if (bitwise_or) {
+      operation = ValueOpcode::BitwiseOr32;
+      reduction_opcode = ValueOpcode::SubgroupOrU32;
+    }
+    if (test == Case::WrongOperation) operation = ValueOpcode::IAdd32;
     for (const uint32_t shift : {1u, 2u, 4u, 8u}) {
       DppMoveFlags flags{.control = 0x110u + shift};
       if (shift == 2) {
@@ -787,13 +793,12 @@ void TestCompleteWaveReduction() {
     const uint32_t wave_size = test == Case::Wave32 || rows ? 32u : 64u;
     ConstantPropagationPass(fixture.program.blocks, wave_size);
     const auto stats = EliminateReadLane(fixture.program, wave_size);
-    const bool matches = test == Case::Masked || rows;
+    const bool matches = test == Case::Masked || rows || bitwise_or;
     Check(stats.rewritten_reductions == (matches ? 1u : 0u) && stats.rewritten_reads == 0,
           ("wave reduction proof failed case " + std::to_string(static_cast<unsigned>(test))).c_str());
     if (matches) {
       const auto *reduction = use.Instruction()->Arg(0).ResolveInstruction();
-      Check(reduction->GetOpcode() ==
+      Check(reduction->GetOpcode() == reduction_opcode &&
-                    (maximum ? ValueOpcode::SubgroupUMaxU32 : ValueOpcode::SubgroupUMinU32) &&
                 reduction->Arg(0).Resolve() == seed.Resolve(),
             "wave reduction changed the neutral-masked input");
       RemoveIdentities(fixture.program);

tests/ShaderRecompilerComputeTests.cpp

@@ -45695,60 +45695,75 @@ TestCase ImageAtomicGlc0DoesNotReturnOldValue() {
   return test;
 }
 
-void AppendWave64Minimum(std::vector<u32> &code) {
+// Reduces v13 into v0 with V_MIN_U32/S_MIN_U32 or, for bitwise_or, V_OR_B32/S_OR_B32.
+void AppendWave64Reduction(std::vector<u32> &code, bool bitwise_or) {
+  const u32 vector_op = bitwise_or ? 0x1cu : 0x13u;
   code.insert(code.end(), {EncodeSop1(0x04, 24, 126), // Save contributing lanes.
                           EncodeSop1(0x28, 106, 126)});
-  AppendVop3(&code, 0x101, 13, 193, Vgpr(13), 24); // Inactive lanes contribute UINT_MAX.
+  // Inactive lanes contribute the identity: 0 for OR, UINT_MAX for the minimum.
+  AppendVop3(&code, 0x101, 13, bitwise_or ? 128u : 193u, Vgpr(13), 24);
   for (u32 shift : {1u, 2u, 4u, 8u}) {
-    code.push_back(EncodeVop2(0x13, 13, 250, 13));
+    code.push_back(EncodeVop2(vector_op, 13, 250, 13));
     code.push_back(EncodeVop2Dpp(13, 0x110u + shift));
   }
-  // Exact RDNA2 wave64 reduction from PS 964747887d898821, PCs 0x224..0x244.
+  // RDNA2 wave64 reduction from PS 964747887d898821, PCs 0x224..0x244.
   code.insert(code.end(), {0xd778100cu, 0x0305830du,
-      EncodeVop2(0x13, 13, Vgpr(13), 12), EncodeSop1(0x04, 126, 106),
+      EncodeVop2(vector_op, 13, Vgpr(13), 12), EncodeSop1(0x04, 126, 106),
       0xd7600006u, 0x00013f0du, 0xd7600007u, 0x00017f0du,
-      EncodeSop2(0x07, 8, 6, 7), EncodeVop1(0x01, 0, 8)});
+      EncodeSop2(bitwise_or ? 0x10u : 0x07u, 8, 6, 7), EncodeVop1(0x01, 0, 8)});
 }
 
-GraphicsCase GraphicsSparseWaveMinimum() {
+std::vector<ShaderOpcode> Wave64ReductionOpcodes(bool bitwise_or) {
+  using O = ShaderOpcode;
+  return {bitwise_or ? O::V_OR_B32 : O::V_MIN_U32, O::V_PERMLANEX16_B32,
+          O::V_READLANE_B32, bitwise_or ? O::S_OR_B32 : O::S_MIN_U32};
+}
+
+GraphicsCase GraphicsSparseWaveReduction(bool bitwise_or) {
   GraphicsCase test;
-  test.name = "GraphicsSparseWaveMinimum";
+  test.name = bitwise_or ? "GraphicsSparseWaveOr" : "GraphicsSparseWaveMinimum";
   auto &code = test.fragment_code;
   code = {EncodeVintrp(0x02, 13, 0, 0, 2)};
-  AppendWave64Minimum(code);
+  AppendWave64Reduction(code, bitwise_or);
   code.insert(code.end(), {EncodeExp0(0x00, 0xf), EncodeExp1(0, 0, 0, 0)});
   AppendEnd(&code);
   test.pixel_interpolator_settings = {0x00000400u};
   test.expected_pixel = {0x3e800000u, 0x3e800000u, 0x3e800000u, 0x3e800000u};
-  test.opcodes = {ShaderOpcode::V_MIN_U32, ShaderOpcode::V_PERMLANEX16_B32,
+  test.opcodes = Wave64ReductionOpcodes(bitwise_or);
-                  ShaderOpcode::V_READLANE_B32, ShaderOpcode::S_MIN_U32};
   return test;
 }
 
-TestCase ComputeWave64Minimum(bool masked) {
+TestCase ComputeWave64Reduction(bool masked, bool bitwise_or) {
   TestCase test;
-  test.name = masked ? "ComputeWave64MinimumMasked" : "ComputeWave64Minimum";
   auto &code = test.code;
   code = {EncodeVop2(0x1a, 20, InlineU32(2), 0)};
+  if (bitwise_or) {
+    // Lane i contributes bit i / 4, so each half supplies its own byte of 0xffff.
+    test.name = "ComputeWave64Or";
+    test.expected.assign(64, 0xffffu);
+    code.push_back(EncodeVop2(0x16, 13, InlineU32(2), 0));
+    AppendVop3(&code, 0x11a, 13, Vgpr(13), InlineU32(1));
+  } else {
+    test.name = masked ? "ComputeWave64MinimumMasked" : "ComputeWave64Minimum";
+    test.expected.assign(64, masked ? 38u : 37u);
     AppendVMovU32(&code, 13, 100);
     code.push_back(EncodeVop2(0x26, 13, Vgpr(13), 0)); // 100 - lane: minimum is in the upper half.
+  }
   if (masked) {
     AppendSMovLiteral(&code, 126, 0xfffffffeu);
     AppendSMovLiteral(&code, 127, 0x7fffffffu);
   }
-  AppendWave64Minimum(code);
+  AppendWave64Reduction(code, bitwise_or);
   AppendBufferStoreDword(&code, 0, 20);
   AppendEnd(&code);
   test.initial.assign(64, 0xa5a5a5a5u);
-  test.expected.assign(64, masked ? 38u : 37u);
   if (masked) test.expected.front() = test.expected.back() = 0xa5a5a5a5u;
   test.compute_info.threads_num[0] = 64;
   test.compute_info.threads_num[1] = test.compute_info.threads_num[2] = 1;
   test.compute_info.thread_ids_num = 1;
   test.compute_info.wave_size = 64;
   test.has_compute_info = true;
-  test.opcodes = {ShaderOpcode::V_MIN_U32, ShaderOpcode::V_PERMLANEX16_B32,
+  test.opcodes = Wave64ReductionOpcodes(bitwise_or);
-                  ShaderOpcode::V_READLANE_B32, ShaderOpcode::S_MIN_U32};
   return test;
 }
 
@@ -46627,8 +46642,9 @@ std::vector<TestCase> MakeCases() {
   cases.push_back(VectorDppMirrorAndRotate(64));
   cases.push_back(VectorReadFirstLaneExecMasks(32));
   cases.push_back(VectorReadFirstLaneExecMasks(64));
-  cases.push_back(ComputeWave64Minimum(false));
+  cases.push_back(ComputeWave64Reduction(false, false));
-  cases.push_back(ComputeWave64Minimum(true));
+  cases.push_back(ComputeWave64Reduction(true, false));
+  cases.push_back(ComputeWave64Reduction(false, true));
   AddCase(VectorDppBankMaskPreservesDestination);
   for (const bool row_shift : {false, true})
     for (const bool bound_ctrl : {false, true})
@@ -47072,7 +47088,8 @@ std::vector<TestCase> MakeCases() {
 
 std::vector<GraphicsCase> MakeGraphicsCases() {
   return {
-      GraphicsSparseWaveMinimum(),
+      GraphicsSparseWaveReduction(false),
+      GraphicsSparseWaveReduction(true),
       GraphicsInterpolationExport(),
       GraphicsPositionWExport(),
       GraphicsPackedHalfCentroid(),
@@ -53177,9 +53194,11 @@ int main(int argc, char **argv) {
   }
   if (argc == 2 && std::strcmp(argv[1], "--dpp-only") == 0) {
     VulkanHarness vulkan;
-    RunGraphicsCase(&vulkan, GraphicsSparseWaveMinimum());
+    RunGraphicsCase(&vulkan, GraphicsSparseWaveReduction(false));
-    RunCase(&vulkan, ComputeWave64Minimum(false));
+    RunGraphicsCase(&vulkan, GraphicsSparseWaveReduction(true));
-    RunCase(&vulkan, ComputeWave64Minimum(true));
+    RunCase(&vulkan, ComputeWave64Reduction(false, false));
+    RunCase(&vulkan, ComputeWave64Reduction(true, false));
+    RunCase(&vulkan, ComputeWave64Reduction(false, true));
     RunCase(&vulkan, VectorDpp8Captured(false));
     for (const u32 wave_size : {32u, 64u})
       for (const bool fi : {false, true})