shader: share scalar condition results with direct subgroup votes

Co-authored-by: fxpw <fxpw@outlook.com>

src/graphics/presentation/window/vulkanWindow.cpp

@@ -535,6 +535,7 @@ static vk::Device VulkanCreateDevice(GraphicContext& graphics,
 		subgroup_size_control.pNext = &graphics.mesh_shader_properties;
 	}
 	physical_device.getProperties2(&properties2);
+	EXIT_NOT_IMPLEMENTED(!(properties11.subgroupSupportedOperations & vk::SubgroupFeatureFlagBits::eVote));
 
 	graphics.subgroup_size                 = properties11.subgroupSize;
 	graphics.min_subgroup_size             = subgroup_size_control.minSubgroupSize;

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

@@ -721,17 +721,22 @@ uint32_t EmitDppUpdateU32(ValueEmitContext& ctx, const IR::Inst& inst) {
 
 uint32_t EmitConditionRef(ValueEmitContext& ctx, const IR::Inst& inst) {
 	// Scalar branches include lanes disabled by EXEC, which the branch body may restore.
-	if (ctx.other_half != nullptr && ctx.half != 0) return ctx.other_half->Def(IR::Value(&inst));
+	if (ctx.other_half != nullptr && ctx.half != 0) return ctx.other_half->Result(inst);
 	const auto kind = inst.Flags<CFG::BranchCondition>();
 	if (kind == CFG::BranchCondition::ScalarInstruction) return ctx.Arg(inst, 0);
 	const bool zero = kind == CFG::BranchCondition::ExecZero ||
 	                  kind == CFG::BranchCondition::VccZero || kind == CFG::BranchCondition::SccZero;
-	const auto ballot = ctx.Ballot(inst.Arg(0));
+	// Zero-kind predicates are already inverted by translation.
-	const auto reference = zero ? ctx.Ballot(IR::Value(true))
+	auto predicate = ctx.Arg(inst, 0);
-	                            : ConstantU32CompositeZero(ctx.state, 4);
+	if (ctx.other_half != nullptr) {
-	const auto compare = Binary(ctx.state, zero ? spv::OpIEqual : spv::OpINotEqual,
+		predicate = Binary(ctx.state, zero ? spv::OpLogicalAnd : spv::OpLogicalOr,
-	                            TypeBoolVector(ctx.state, 4), ballot, reference);
+		                   TypeBool(ctx.state), predicate, ctx.HalfArg(inst, 0, 1));
-	return Unary(ctx.state, zero ? spv::OpAll : spv::OpAny, TypeBool(ctx.state), compare);
+	}
+	const auto result = ctx.state.builder.AllocateId();
+	ctx.state.builder.AddFunction(zero ? spv::OpGroupNonUniformAll : spv::OpGroupNonUniformAny,
+	                              TypeBool(ctx.state), result,
+	                              ConstantU32(ctx.state, spv::ScopeSubgroup), predicate);
+	return result;
 }
 
 uint32_t EmitBallot(ValueEmitContext& ctx, IR::Value predicate) {

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

@@ -764,13 +764,17 @@ void DefineModule(EmitterState& state) {
 		state.builder.RequireCapability(spv::CapabilityImageGatherExtended);
 	}
 	if (state.lane_count == 2 || state.program.info.subgroup_barrier ||
-	    state.program.info.subgroup_ballot || state.program.info.subgroup_shuffle ||
+	    state.program.info.subgroup_ballot || state.program.info.subgroup_vote ||
+	    state.program.info.subgroup_shuffle ||
 	    state.program.info.subgroup_arithmetic || state.program.info.subgroup_local_invocation_id) {
 		state.builder.RequireCapability(spv::CapabilityGroupNonUniform);
 	}
 	if (state.lane_count == 2 || state.program.info.subgroup_ballot) {
 		state.builder.RequireCapability(spv::CapabilityGroupNonUniformBallot);
 	}
+	if (state.program.info.subgroup_vote) {
+		state.builder.RequireCapability(spv::CapabilityGroupNonUniformVote);
+	}
 	if (state.program.info.subgroup_shuffle) {
 		state.builder.RequireCapability(spv::CapabilityGroupNonUniformShuffle);
 	}

src/graphics/shader/recompiler/ir/ShaderIR.h

@@ -506,6 +506,7 @@ struct ShaderInfo {
 	bool                            uses_dispatch_threads         = false;
 	bool                            bvh                           = false;
 	bool                            subgroup_ballot               = false;
+	bool                            subgroup_vote                 = false;
 	bool                            subgroup_barrier              = false;
 	bool                            subgroup_shuffle              = false;
 	bool                            subgroup_arithmetic           = false;

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

@@ -510,7 +510,7 @@ void Visit(Program& program, ShaderStageInputInfo input_info, InputUsage& inputs
 		case ValueOpcode::StoreCompletion: info.subgroup_barrier = true; break;
 		case ValueOpcode::BvhIntersect: info.bvh = true; break;
 		case ValueOpcode::ConditionRef:
-			info.subgroup_ballot |=
+			info.subgroup_vote |=
 			    inst.Flags<CFG::BranchCondition>() != CFG::BranchCondition::ScalarInstruction;
 			break;
 		case ValueOpcode::Ballot: info.subgroup_ballot = true; break;
@@ -572,7 +572,8 @@ void CollectShaderInfo(Program& program, ShaderStageInputInfo input_info) {
 	for (auto& buffer: info.buffers) buffer.live = false;
 	info.has_bitwise_xor = info.uses_lds = info.uses_gds = info.uses_flattened_srt =
 	    info.uses_dispatch_threads = info.uses_swizzle = false;
-	info.bvh = info.subgroup_ballot = info.subgroup_barrier = info.subgroup_shuffle = info.subgroup_arithmetic =
+	info.bvh = info.subgroup_ballot = info.subgroup_vote = info.subgroup_barrier =
+	    info.subgroup_shuffle = info.subgroup_arithmetic =
 	    info.subgroup_local_invocation_id = info.compute_derivatives = info.image_gather_extended =
 	    info.function_lds = info.function_scratch = info.pixel_valid_mask = info.buffer_int64_atomics =
 	    info.buffer_u8 = info.buffer_u16 = info.shared_int64_atomics = info.coherent_buffers = info.float64 = false;

tests/ShaderRecompilerComputeTests.cpp

@@ -2017,6 +2017,9 @@ public:
     vk::PhysicalDeviceProperties2 properties{};
     properties.pNext = &subgroup;
     m_physical_device.getProperties2(&properties);
+    Require("VulkanHarness", "subgroup vote support",
+            bool(subgroup.supportedOperations & vk::SubgroupFeatureFlagBits::eVote),
+            "device does not support subgroup All/Any votes");
     return subgroup.subgroupSize;
   }
   // A consumer GPU may lack the production rasterization features the fragment
@@ -30000,10 +30003,13 @@ TestCase BranchVccnzUsesWaveMask() {
   return test;
 }
 
-TestCase ScalarBranchRestoresInactiveLanes(u32 wave_size, u32 threads) {
+TestCase ScalarBranchRestoresInactiveLanes(u32 wave_size, u32 threads, bool upper_nonzero = false) {
   using O = ShaderOpcode;
   TestCase test;
-  if (threads < wave_size) {
+  if (upper_nonzero) {
+    test.name = wave_size == 32 ? "ScalarBranchRestoresLastLaneWave32"
+                               : "ScalarBranchRestoresUpperHalfWave64";
+  } else if (threads < wave_size) {
     test.name = wave_size == 32 ? "ScalarBranchRestoresPartialWave32"
                                : "ScalarBranchRestoresPartialWave64";
   } else {
@@ -30014,10 +30020,16 @@ TestCase ScalarBranchRestoresInactiveLanes(u32 wave_size, u32 threads) {
   code = {EncodeVop1(0x01, 1, InlineU32(32)),
           EncodeVop2(0x25, 1, Vgpr(0), 1),
           EncodeSop1(0x04, 8, 126), // Save all participating lanes.
-          EncodeSMovB32(20, InlineU32(0)),
+          EncodeSMovB32(20, InlineU32(0))};
-          EncodeSMovB32(126, InlineU32(1)), EncodeSMovB32(127, InlineU32(0)),
+  AppendSMovLiteral(&code, 126, upper_nonzero ? (wave_size == 32 ? 0x80000000u : 0u) : 1u);
-          EncodeSopp(0x08, 3), // EXEC is nonzero for the whole wave.
+  AppendSMovLiteral(&code, 127, upper_nonzero && wave_size == 64 ? 0x80000000u : 0u);
-          EncodeSop1(0x04, 126, 8)};
+  if (upper_nonzero) {
+    code.push_back(EncodeSopp(0x09, 1)); // Last lane makes the whole-wave EXEC nonzero.
+    code.push_back(EncodeSopp(0x02, 3));
+  } else {
+    code.push_back(EncodeSopp(0x08, 3)); // EXEC is nonzero for the whole wave.
+  }
+  code.push_back(EncodeSop1(0x04, 126, 8));
   // The branch restores inactive lanes before reading the last lane's value.
   AppendVop3(&code, 0x360, 20, Vgpr(1), InlineU32(threads - 1));
   code.push_back(EncodeSop1(0x04, 126, 8));
@@ -30026,9 +30038,10 @@ TestCase ScalarBranchRestoresInactiveLanes(u32 wave_size, u32 threads) {
   test.initial.assign(threads, 0);
   test.expected.assign(threads, 32 + threads - 1);
   test.opcodes = {O::V_MOV_B32, O::V_ADD_NC_U32, O::S_MOV_B32, O::S_MOV_B64,
-                  O::S_CBRANCH_EXECZ, O::V_READLANE_B32, O::V_LSHLREV_B32,
+                  upper_nonzero ? O::S_CBRANCH_EXECNZ : O::S_CBRANCH_EXECZ,
+                  O::V_READLANE_B32, O::V_LSHLREV_B32,
                   O::BUFFER_STORE_DWORD, O::S_ENDPGM};
-  test.required_spirv = {"OpGroupNonUniformBallot", "OpAll"};
+  test.required_spirv = {upper_nonzero ? "OpGroupNonUniformAny" : "OpGroupNonUniformAll"};
   test.compute_info.wave_size = wave_size;
   test.compute_info.threads_num[0] = threads;
   test.compute_info.thread_ids_num = 1;
@@ -41148,6 +41161,8 @@ std::vector<TestCase> MakeCases() {
   AddCase([] { return ScalarBranchRestoresInactiveLanes(64, 64); });
   AddCase([] { return ScalarBranchRestoresInactiveLanes(32, 4); });
   AddCase([] { return ScalarBranchRestoresInactiveLanes(64, 4); });
+  AddCase([] { return ScalarBranchRestoresInactiveLanes(32, 32, true); });
+  AddCase([] { return ScalarBranchRestoresInactiveLanes(64, 64, true); });
   AddCase(BranchVccnzUsesCarryProducedWaveMask);
   AddCase(ScalarMemRealtimeCapturedPlaceholder);
   AddCase(ScalarMemoryLoadVariants);
@@ -47204,6 +47219,8 @@ int main(int argc, char **argv) {
     RunCase(&vulkan, ScalarBranchRestoresInactiveLanes(64, 64));
     RunCase(&vulkan, ScalarBranchRestoresInactiveLanes(32, 4));
     RunCase(&vulkan, ScalarBranchRestoresInactiveLanes(64, 4));
+    RunCase(&vulkan, ScalarBranchRestoresInactiveLanes(32, 32, true));
+    RunCase(&vulkan, ScalarBranchRestoresInactiveLanes(64, 64, true));
     RunCase(&vulkan, BranchVccnzUsesCarryProducedWaveMask());
     RunCase(&vulkan, ImageSampleAndGather());
     RunCase(&vulkan, ImageGatherExplicitLod());

tests/shaderCfgTests.cpp

@@ -15690,6 +15690,67 @@ void TestSpirvRepeatedEmission() {
         "split-wave re-emission retained native-wave definitions");
 }
 
+void TestSpirvConditionVotes() {
+  using namespace ShaderRecompiler;
+  using K = CFG::BranchCondition;
+  for (const auto [wave, host] : {std::pair{32u, 32u}, std::pair{64u, 64u},
+                                  std::pair{64u, 32u}}) {
+    const uint32_t shader[] = {EncodeSopp(0x01)};
+    ShaderComputeInputInfo compute{};
+    compute.wave_size = wave;
+    compute.host_subgroup_size = host;
+    compute.threads_num[0] = wave;
+    compute.threads_num[1] = compute.threads_num[2] = 1;
+    auto options = MakeCompileOptions(ShaderType::Compute);
+    options.wave_size = wave;
+    options.input_info.compute = &compute;
+    auto compiled = RecompileForTest(shader, options);
+    IR::IREmitter emitter(compiled.program.blocks.front());
+    const auto lane = emitter.Emit(IR::ValueOpcode::LaneId, {});
+    const auto predicate = emitter.Emit(IR::ValueOpcode::IEqual32, {lane, IR::Value(33u)});
+    std::vector<IR::Inst *> references;
+    std::vector<IR::Inst *> conditions;
+    for (const auto kind : {K::ExecZero, K::VccZero, K::SccZero,
+                            K::ExecNonZero, K::VccNonZero, K::SccNonZero}) {
+      const auto condition = emitter.Emit(IR::ValueOpcode::ConditionRef, {predicate}, kind);
+      conditions.push_back(condition.TryInstruction());
+      references.push_back(&compiled.program.blocks.front()->AppendNewInst(
+          IR::ValueOpcode::Reference, {condition}));
+    }
+    const auto scalar = emitter.Emit(IR::ValueOpcode::ConditionRef, {predicate}, K::ScalarInstruction);
+    emitter.Emit(IR::ValueOpcode::Reference, {scalar});
+    compiled.program.shader_info_complete = false;
+    IR::CollectShaderInfo(compiled.program, options.input_info);
+    const auto binary = Spirv::EmitProgram(compiled.program, options.input_info);
+    CheckSpirvBinaryValidates(binary);
+    Check(compiled.program.info.subgroup_vote && !compiled.program.info.subgroup_ballot &&
+              SpirvContainsCapability(binary, spv::CapabilityGroupNonUniformVote) &&
+              SpirvInstructionOpcodeCount(binary, spv::OpGroupNonUniformAll) == 3 &&
+              SpirvInstructionOpcodeCount(binary, spv::OpGroupNonUniformAny) == 3 &&
+              SpirvInstructionOpcodeCount(binary, spv::OpGroupNonUniformBallot) == 0,
+          "condition votes used ballots, duplicated halves, or lost their capability");
+    Check(scalar.TryInstruction()->Definition(0) == predicate.TryInstruction()->Definition(0),
+          "scalar instruction condition acquired a subgroup vote");
+    if (wave != host) {
+      Check(SpirvInstructionOpcodeCount(binary, spv::OpLogicalAnd) == 3 &&
+                SpirvInstructionOpcodeCount(binary, spv::OpLogicalOr) == 3 &&
+                std::ranges::all_of(conditions, [](const auto *inst) {
+                  return inst->Definition(0) == inst->Definition(1);
+                }) && scalar.TryInstruction()->Definition(0) == scalar.TryInstruction()->Definition(1),
+            "condition votes did not combine and reuse both virtual halves");
+    }
+    for (auto *reference : references) reference->Invalidate();
+    IR::EliminateDeadCode(compiled.program.blocks);
+    compiled.program.shader_info_complete = false;
+    IR::CollectShaderInfo(compiled.program, options.input_info);
+    const auto scalar_only = Spirv::EmitProgram(compiled.program, options.input_info);
+    CheckSpirvBinaryValidates(scalar_only);
+    Check(!compiled.program.info.subgroup_vote &&
+              !SpirvContainsCapability(scalar_only, spv::CapabilityGroupNonUniformVote),
+          "recollection retained the removed condition vote capability");
+  }
+}
+
 void TestRepeatedExportsHaveOneInterface() {
   const uint32_t shader[] = {EncodeExp0(0x0c, 0xf), EncodeExp1(0, 1, 2, 3),
                              EncodeExp0(0x0c, 0xf), EncodeExp1(0, 1, 2, 3),
@@ -16027,6 +16088,7 @@ int main() {
   TestDeferredSpirvPhiPatching();
   TestCompilerStageInputOwnership();
   TestSpirvRepeatedEmission();
+  TestSpirvConditionVotes();
   TestRepeatedExportsHaveOneInterface();
   TestNewShaderRecompilerSpirvSizeBaselines();
   TestDemandDrivenSpirvDeclarations();