refactor(shader): compact whole-quad mask lowering

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

@@ -18,24 +18,37 @@ bool UserDataDwordIndex(const EmitterState& state, IR::Register reg, uint32_t& d
 	return true;
 }
 
-uint32_t EmitWqmLaneU32(EmitterState& state, uint32_t src) {
+uint32_t EmitWqmWordU32(EmitterState& state, uint32_t value) {
-	uint32_t ret = ConstantU32(state, 0);
+	const auto shifted_one = state.builder.AllocateId();
-	for (uint32_t i = 0; i < 8u; i++) {
+	const auto merged_one  = state.builder.AllocateId();
-		const auto mask     = 0x0fu << (i * 4u);
+	const auto shifted_two = state.builder.AllocateId();
-		const auto masked   = state.builder.AllocateId();
+	const auto merged_two  = state.builder.AllocateId();
-		const auto non_zero = state.builder.AllocateId();
+	const auto quad_bits   = state.builder.AllocateId();
-		const auto expanded = state.builder.AllocateId();
+	const auto result      = state.builder.AllocateId();
-		const auto combined = state.builder.AllocateId();
+	state.builder.AddFunction({OpShiftRightLogical, TypeU32(state), shifted_one, value,
-		state.builder.AddFunction(
+	                           ConstantU32(state, 1)});
-		    {OpBitwiseAnd, TypeU32(state), masked, src, ConstantU32(state, mask)});
+	state.builder.AddFunction(
-		state.builder.AddFunction(
+	    {OpBitwiseOr, TypeU32(state), merged_one, value, shifted_one});
-		    {OpINotEqual, TypeBool(state), non_zero, masked, ConstantU32(state, 0)});
+	state.builder.AddFunction({OpShiftRightLogical, TypeU32(state), shifted_two, merged_one,
-		state.builder.AddFunction({OpSelect, TypeU32(state), expanded, non_zero,
+	                           ConstantU32(state, 2)});
-		                           ConstantU32(state, mask), ConstantU32(state, 0)});
+	state.builder.AddFunction(
-		state.builder.AddFunction({OpBitwiseOr, TypeU32(state), combined, ret, expanded});
+	    {OpBitwiseOr, TypeU32(state), merged_two, merged_one, shifted_two});
-		ret = combined;
+	state.builder.AddFunction({OpBitwiseAnd, TypeU32(state), quad_bits, merged_two,
-	}
+	                           ConstantU32(state, 0x11111111u)});
-	return ret;
+	state.builder.AddFunction(
+	    {OpIMul, TypeU32(state), result, quad_bits, ConstantU32(state, 0x0fu)});
+	return result;
+}
+
+uint32_t EmitWqmU64(EmitterState& state, uint32_t value) {
+	const auto low  = state.builder.AllocateId();
+	const auto high = state.builder.AllocateId();
+	state.builder.AddFunction({OpCompositeExtract, TypeU32(state), low, value, 0});
+	state.builder.AddFunction({OpCompositeExtract, TypeU32(state), high, value, 1});
+	const auto result = state.builder.AllocateId();
+	state.builder.AddFunction({OpCompositeConstruct, TypeU32Pair(state), result,
+	                           EmitWqmWordU32(state, low), EmitWqmWordU32(state, high)});
+	return result;
 }
 
 uint32_t BoolAsU32(ValueEmitContext& ctx, uint32_t value) {
@@ -88,24 +101,28 @@ uint32_t EmitWqm(ValueEmitContext& ctx, uint32_t active) {
 	const auto ballot = state.builder.AllocateId();
 	state.builder.AddFunction({OpGroupNonUniformBallot, TypeU32Vector(state, 4), ballot,
 	                           ConstantU32(state, ScopeSubgroup), active});
-	const auto low  = state.builder.AllocateId();
+	const auto low = state.builder.AllocateId();
-	const auto high = state.builder.AllocateId();
 	state.builder.AddFunction({OpCompositeExtract, TypeU32(state), low, ballot, 0});
-	state.builder.AddFunction({OpCompositeExtract, TypeU32(state), high, ballot, 1});
+	const auto wqm_low = EmitWqmWordU32(state, low);
-	const auto wqm_low  = EmitWqmLaneU32(state, low);
-	const auto wqm_high = EmitWqmLaneU32(state, high);
 	const auto lane     = EmitSubgroupLocalInvocationId(state);
-	const auto upper    = state.builder.AllocateId();
+	uint32_t   mask     = wqm_low;
-	const auto mask     = state.builder.AllocateId();
+	uint32_t   bit_lane = lane;
-	const auto bit_lane = state.builder.AllocateId();
+	if (state.wave_size == 64u) {
+		const auto high = state.builder.AllocateId();
+		state.builder.AddFunction({OpCompositeExtract, TypeU32(state), high, ballot, 1});
+		const auto wqm_high = EmitWqmWordU32(state, high);
+		const auto upper    = state.builder.AllocateId();
+		mask                = state.builder.AllocateId();
+		bit_lane            = state.builder.AllocateId();
+		state.builder.AddFunction(
+		    {OpUGreaterThanEqual, TypeBool(state), upper, lane, ConstantU32(state, 32)});
+		state.builder.AddFunction({OpSelect, TypeU32(state), mask, upper, wqm_high, wqm_low});
+		state.builder.AddFunction(
+		    {OpBitwiseAnd, TypeU32(state), bit_lane, lane, ConstantU32(state, 31)});
+	}
 	const auto bit      = state.builder.AllocateId();
 	const auto hit      = state.builder.AllocateId();
 	const auto result   = state.builder.AllocateId();
-	state.builder.AddFunction(
-	    {OpUGreaterThanEqual, TypeBool(state), upper, lane, ConstantU32(state, 32)});
-	state.builder.AddFunction({OpSelect, TypeU32(state), mask, upper, wqm_high, wqm_low});
-	state.builder.AddFunction(
-	    {OpBitwiseAnd, TypeU32(state), bit_lane, lane, ConstantU32(state, 31)});
 	state.builder.AddFunction(
 	    {OpShiftLeftLogical, TypeU32(state), bit, ConstantU32(state, 1), bit_lane});
 	state.builder.AddFunction({OpBitwiseAnd, TypeU32(state), hit, mask, bit});
@@ -404,6 +421,9 @@ bool EmitValueFlow(ValueEmitContext& ctx, const IR::Inst& inst) {
 			ctx.Emit(inst, OpSelect, IR::Type::U32, {write, ctx.Arg(inst, 0), ctx.Arg(inst, 1)});
 			return true;
 		}
+		case IR::ValueOpcode::WqmU64:
+			ctx.Define(inst, EmitWqmU64(ctx.state, ctx.Arg(inst, 0)));
+			return true;
 		case IR::ValueOpcode::WqmMask:
 			ctx.Define(inst, EmitWqm(ctx, ctx.Arg(inst, 0)));
 			return true;

src/graphics/shader/recompiler/frontend/translate/Control.cpp

@@ -170,10 +170,9 @@ bool Translator::TranslateMove(const IR::Instruction& inst) {
 			return true;
 		}
 		case IR::Opcode::WqmB64: {
-			if (!current_per_invocation_masks && inst.dst.kind == IR::OperandKind::Register &&
+			if (!current_per_invocation_masks) {
-			    inst.dst.reg.file == IR::RegisterFile::Scalar) {
+				const auto result = IR::U64(ir.Emit(
-				const auto result =
+				    IR::ValueOpcode::WqmU64, {ReadOperand(inst.src[0], IR::Type::U64)}));
-				    ExpandWholeQuadMask(IR::U64(ReadOperand(inst.src[0], IR::Type::U64)));
 				WriteOperand(inst.dst, result);
 				ir.SetScc(IR::U1(
 				    ir.Emit(IR::ValueOpcode::INotEqual64, {result, IR::Value(uint64_t {0})})));

src/graphics/shader/recompiler/frontend/translate/Translate.cpp

@@ -310,21 +310,6 @@ void Translator::WriteRawU32(const IR::Operand& operand, IR::U32 value) {
 	}
 }
 
-IR::U64 Translator::ExpandWholeQuadMask(IR::U64 value) {
-	const auto shifted1 =
-	    IR::U64(ir.Emit(IR::ValueOpcode::ShiftRightLogical64, {value, IR::Value(1u)}));
-	const auto shifted2 =
-	    IR::U64(ir.Emit(IR::ValueOpcode::ShiftRightLogical64, {value, IR::Value(2u)}));
-	const auto shifted3 =
-	    IR::U64(ir.Emit(IR::ValueOpcode::ShiftRightLogical64, {value, IR::Value(3u)}));
-	const auto merged1 = ir.Emit(IR::ValueOpcode::BitwiseOr64, {value, shifted1});
-	const auto merged2 = ir.Emit(IR::ValueOpcode::BitwiseOr64, {merged1, shifted2});
-	const auto merged3 = ir.Emit(IR::ValueOpcode::BitwiseOr64, {merged2, shifted3});
-	const auto low_bits =
-	    ir.Emit(IR::ValueOpcode::BitwiseAnd64, {merged3, IR::Value(uint64_t {0x1111111111111111})});
-	return IR::U64(ir.Emit(IR::ValueOpcode::IMul64, {low_bits, IR::Value(uint64_t {0xfull})}));
-}
-
 IR::F32 Translator::ApplyF32ResultModifiers(const IR::Operand& operand, IR::F32 value) {
 	if (operand.omod != 0u) {
 		float multiplier = 0.5f;

src/graphics/shader/recompiler/frontend/translate/Translator.h

@@ -34,7 +34,6 @@ private:
 	IR::Value              ReadOperand(const IR::Operand& operand, IR::Type type);
 	IR::U1                 ThreadBit(IR::U32 low, IR::U32 high);
 	void                   WriteRawU32(const IR::Operand& operand, IR::U32 value);
-	IR::U64                ExpandWholeQuadMask(IR::U64 value);
 	IR::F32                ApplyF32ResultModifiers(const IR::Operand& operand, IR::F32 value);
 	void                   WriteOperand(const IR::Operand& operand, IR::Value value);
 	IR::U32                PackHalf2x16(IR::F32 low, IR::F32 high);

src/graphics/shader/recompiler/ir/ValueProgram.cpp

@@ -304,6 +304,13 @@ bool ValidateValueProgram(const ValueProgram& program, bool require_ssa, std::st
 				return Fail(error, fmt::format("untyped opcode {} survived translation",
 				                               ValueOpcodeName(inst.GetOpcode())));
 			}
+			const bool fixed_signature =
+			    inst.GetOpcode() != ValueOpcode::Phi && inst.GetOpcode() != ValueOpcode::Identity;
+			if (fixed_signature && inst.NumArgs() != NumArgsOf(inst.GetOpcode())) {
+				return Fail(error, fmt::format("{} has {} arguments, expected {}",
+				                               ValueOpcodeName(inst.GetOpcode()), inst.NumArgs(),
+				                               NumArgsOf(inst.GetOpcode())));
+			}
 			const auto buffer_components = BufferComponentCount(inst.GetOpcode());
 			if (buffer_components != 0u) {
 				const auto memory_index = inst.Flags<MemoryFlags>().index;
@@ -350,6 +357,12 @@ bool ValidateValueProgram(const ValueProgram& program, bool require_ssa, std::st
 					return Fail(error, fmt::format("{} has an empty argument",
 					                               ValueOpcodeName(inst.GetOpcode())));
 				}
+				if (fixed_signature && arg.GetType() != ArgTypeOf(inst.GetOpcode(), arg_index)) {
+					return Fail(error, fmt::format("{} argument {} has type {}, expected {}",
+					                               ValueOpcodeName(inst.GetOpcode()), arg_index,
+					                               TypeName(arg.GetType()),
+					                               TypeName(ArgTypeOf(inst.GetOpcode(), arg_index))));
+				}
 				if (const auto* definition = arg.TryInstruction();
 				    definition != nullptr && !instructions.contains(definition)) {
 					return Fail(error, "value IR argument has a foreign definition");

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

@@ -79,6 +79,7 @@ VALUE_OPCODE(BitFieldUExtract, U32, U32, U32, U32)
 VALUE_OPCODE(BitFieldSExtract, U32, U32, U32, U32)
 VALUE_OPCODE(DppMoveU32, U32, U32, U1)
 VALUE_OPCODE(DppUpdateU32, U32, U32, U32, U1)
+VALUE_OPCODE(WqmU64, U64, U64)
 VALUE_OPCODE(WqmMask, U1, U1)
 
 // Core value operations used to expand architectural side effects.

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

@@ -357,6 +357,21 @@ void FoldInstruction(Inst& inst) {
 				ReplaceBinaryIdentity(inst, Type::U64, 1u);
 			}
 			return;
+		case ValueOpcode::WqmU64: {
+			const auto value = Arg(inst, 0);
+			if (IsImmediate(value, Type::U64)) {
+				const auto expand = [](uint32_t word) {
+					auto quads = word | (word >> 1u);
+					quads |= quads >> 2u;
+					return (quads & 0x11111111u) * 0x0fu;
+				};
+				const auto low  = expand(static_cast<uint32_t>(value.U64()));
+				const auto high = expand(static_cast<uint32_t>(value.U64() >> 32u));
+				Replace(inst, Value(static_cast<uint64_t>(low) |
+				                    (static_cast<uint64_t>(high) << 32u)));
+			}
+			return;
+		}
 		case ValueOpcode::UDiv32: {
 			const auto rhs = Arg(inst, 1);
 			if (IsImmediate(rhs, Type::U32) && rhs.U32() == 1u) {

tests/ShaderRecompilerComputeTests.cpp

@@ -11599,8 +11599,8 @@ TestCase ScalarShiftAddAndMaskOps() {
   AppendSMovLiteral(&code, 8, 0xffffffffu);
   code.push_back(EncodeSMovB32(9, InlineU32(2)));
   code.push_back(EncodeSop2(0x35, 10, 8, 9));
-  AppendSMovLiteral(&code, 12, 0x0f0f0f0fu);
+  AppendSMovLiteral(&code, 12, 0x80000011u);
-  AppendSMovLiteral(&code, 13, 0x00ff00ffu);
+  AppendSMovLiteral(&code, 13, 0x80000001u);
   code.push_back(EncodeSop1(0x08, 14, 12));
   code.push_back(EncodeSop1(0x0a, 16, 12));
 
@@ -11617,8 +11617,8 @@ TestCase ScalarShiftAddAndMaskOps() {
   return {"ScalarShiftAddAndMaskOps",
           code,
           {},
-          {7, 11, 19, 35, 0xfffffffeu, 1, 0xf0f0f0f0u, 0xff00ff00u, 0x0f0f0f0fu,
+          {7, 11, 19, 35, 0xfffffffeu, 1, 0x7fffffeeu, 0x7ffffffeu, 0xf00000ffu,
-           0x00ff00ffu},
+           0xf000000fu},
           {O::S_MOV_B32, O::S_LSHL1_ADD_U32, O::S_LSHL2_ADD_U32, O::S_LSHL3_ADD_U32,
            O::S_LSHL4_ADD_U32, O::S_ASHR_I32, O::S_MUL_HI_U32, O::S_NOT_B64, O::S_WQM_B64,
            O::V_MOV_B32, O::BUFFER_STORE_DWORD, O::S_ENDPGM}};

tests/shaderCfgTests.cpp

@@ -8503,6 +8503,53 @@ void TestValuePhiValidation() {
   validate_cfg(false);
 }
 
+void TestWqmConstantPropagation() {
+  using namespace ShaderRecompiler::IR;
+
+  ValueProgram program;
+  program.block_storage.push_back(std::make_unique<Block>());
+  program.blocks.push_back(program.block_storage.back().get());
+  program.block_info.emplace_back();
+  IREmitter ir(program.blocks.front());
+
+  constexpr std::array cases = {
+      std::pair{uint64_t{0}, uint64_t{0}},
+      std::pair{uint64_t{1}, uint64_t{0xfull}},
+      std::pair{uint64_t{1} << 31u, uint64_t{0xf} << 28u},
+      std::pair{uint64_t{1} << 32u, uint64_t{0xf} << 32u},
+      std::pair{uint64_t{0x8000001180000001ull},
+                uint64_t{0xf00000fff000000full}},
+  };
+  std::array<Value, cases.size()> results;
+  for (size_t index = 0; index < cases.size(); index++) {
+    results[index] = ir.Emit(ValueOpcode::WqmU64, {Value(cases[index].first)});
+  }
+
+  ConstantPropagationPass(program.blocks);
+  for (size_t index = 0; index < cases.size(); index++) {
+    const auto value = results[index].Resolve();
+    Check(value.IsImmediate() && value.GetType() == Type::U64 &&
+              value.U64() == cases[index].second,
+          "WqmU64 constant propagation crossed a quad or dword boundary");
+  }
+
+  const auto check_invalid_signature = [](ValueOpcode opcode, Value valid, Value invalid) {
+    ValueProgram malformed;
+    malformed.block_storage.push_back(std::make_unique<Block>());
+    malformed.blocks.push_back(malformed.block_storage.back().get());
+    malformed.block_info.emplace_back();
+    IREmitter malformed_ir(malformed.blocks.front());
+    auto value = malformed_ir.Emit(opcode, {valid});
+    value.TryInstruction()->SetArg(0, invalid);
+    std::string error;
+    Check(!ValidateValueProgram(malformed, false, &error) &&
+              Common::ContainsStr(error, "argument 0 has type"),
+          "value IR accepted a WQM operand that violates opcode metadata");
+  };
+  check_invalid_signature(ValueOpcode::WqmU64, Value(uint64_t{1}), Value(true));
+  check_invalid_signature(ValueOpcode::WqmMask, Value(true), Value(uint64_t{1}));
+}
+
 void TestNativeWideValueValidation() {
   using namespace ShaderRecompiler::IR;
 
@@ -8738,6 +8785,55 @@ void TestNewShaderRecompilerPerInvocationMasksWithoutMirrors() {
   Check(Common::ContainsStr(wqm_source, "OpCapability GroupNonUniformBallot") &&
             Common::ContainsStr(wqm_source, "OpGroupNonUniformBallot"),
         "per-invocation scalar WQM omitted its subgroup ballot capability");
+  Check(SpirvInstructionOpcodeCount(result.spirv, 132u) == 2u,
+        "wave64 WQM did not compact exactly two ballot words");
+
+  auto wave32_options = options;
+  wave32_options.wave_size = 32u;
+  Check(ShaderRecompiler::TryRecompile(wqm_shader, wave32_options, result,
+                                       &error),
+        error.c_str());
+  CheckSpirvBinaryValidates(result.spirv);
+  Check(SpirvInstructionOpcodeCount(result.spirv, 132u) == 1u,
+        "wave32 WQM retained the unused high ballot-word expansion");
+
+  const uint32_t native_wqm_shader[] = {
+      EncodeSop1(0x0a, 2, 0), // s_wqm_b64 s[2:3], s[0:1]
+      EncodeVop1(0x01, 0, 2), // v_mov_b32 v0, s2
+      EncodeExp0(0x0c, 0x1), EncodeExp1(0, 0, 0, 0), EncodeSopp(0x01),
+  };
+  auto native_options = options;
+  native_options.lane_mask_mode = ShaderLaneMaskMode::NativeWave;
+  Check(ShaderRecompiler::TryRecompile(native_wqm_shader, native_options,
+                                       result, &error),
+        error.c_str());
+  CheckSpirvBinaryValidates(result.spirv);
+  Check(Common::ContainsStr(result.ir_dump, "WqmU64") &&
+            !Common::ContainsStr(DisassembleSpirvBinary(result.spirv),
+                                 "OpGroupNonUniformBallot"),
+        "native scalar WQM retained frontend expansion or subgroup machinery");
+  Check(SpirvInstructionOpcodeCount(result.spirv, 132u) == 2u &&
+            SpirvInstructionOpcodeCount(result.spirv, 194u) == 4u,
+        "native scalar WQM emitter was folded away or did not use the compact path");
+
+  const uint32_t native_special_wqm_shader[] = {
+      EncodeSop1(0x04, 126, 0),   // s_mov_b64 exec, s[0:1]
+      EncodeSop1(0x0a, 126, 126), // s_wqm_b64 exec, exec
+      EncodeSop1(0x04, 106, 0),   // s_mov_b64 vcc, s[0:1]
+      EncodeSop1(0x0a, 106, 106), // s_wqm_b64 vcc, vcc
+      EncodeVop1(0x01, 0, 126),   // v_mov_b32 v0, exec_lo
+      EncodeVop1(0x01, 1, 106),   // v_mov_b32 v1, vcc_lo
+      EncodeExp0(0x0c, 0x3), EncodeExp1(0, 1, 0, 0), EncodeSopp(0x01),
+  };
+  Check(ShaderRecompiler::TryRecompile(native_special_wqm_shader,
+                                       native_options, result, &error),
+        error.c_str());
+  CheckSpirvBinaryValidates(result.spirv);
+  const auto native_special_source = DisassembleSpirvBinary(result.spirv);
+  Check(CountSourceOccurrences(result.ir_dump, "WqmU64") == 2u &&
+            !Common::ContainsStr(native_special_source,
+                                 "OpGroupNonUniformBallot"),
+        "native EXEC/VCC WQM retained destination-dependent ballot lowering");
 
   const uint32_t cross_lane_shader[] = {
       EncodeSop2(0x25, 126, 132, 128), // s_bfm_b64 exec, 4, 0
@@ -10408,8 +10504,8 @@ void TestNewShaderRecompilerSpirvSizeBaselines() {
   };
   const auto wqm_result =
       compile("wqm", wqm,
-              {.words = 703,
+              {.words = 407,
-               .instructions = 155,
+               .instructions = 98,
                .variables = 4,
                .loads = 3,
                .stores = 1,
@@ -10519,6 +10615,7 @@ int main() {
   TestNewShaderRecompilerSetpcJumpTable();
   TestNewShaderRecompilerSetpcDwordJumpTable();
   TestTypedEntryStateIsMinimal();
+  TestWqmConstantPropagation();
   TestFinalSsaRejectsRegisterStatePseudos();
   TestValuePhiValidation();
   TestNativeWideValueValidation();