refactor(shader): compact whole-quad mask lowering

This commit is contained in:
nmzik
2026-08-19 08:21:18 +02:00
parent 5383972e96
commit c9e942a972
9 changed files with 186 additions and 57 deletions
@@ -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 ret = ConstantU32(state, 0);
for (uint32_t i = 0; i < 8u; i++) {
const auto mask = 0x0fu << (i * 4u);
const auto masked = state.builder.AllocateId();
const auto non_zero = state.builder.AllocateId();
const auto expanded = state.builder.AllocateId();
const auto combined = state.builder.AllocateId();
state.builder.AddFunction(
{OpBitwiseAnd, TypeU32(state), masked, src, ConstantU32(state, mask)});
state.builder.AddFunction(
{OpINotEqual, TypeBool(state), non_zero, masked, ConstantU32(state, 0)});
state.builder.AddFunction({OpSelect, TypeU32(state), expanded, non_zero,
ConstantU32(state, mask), ConstantU32(state, 0)});
state.builder.AddFunction({OpBitwiseOr, TypeU32(state), combined, ret, expanded});
ret = combined;
}
return ret;
uint32_t EmitWqmWordU32(EmitterState& state, uint32_t value) {
const auto shifted_one = state.builder.AllocateId();
const auto merged_one = state.builder.AllocateId();
const auto shifted_two = state.builder.AllocateId();
const auto merged_two = state.builder.AllocateId();
const auto quad_bits = state.builder.AllocateId();
const auto result = state.builder.AllocateId();
state.builder.AddFunction({OpShiftRightLogical, TypeU32(state), shifted_one, value,
ConstantU32(state, 1)});
state.builder.AddFunction(
{OpBitwiseOr, TypeU32(state), merged_one, value, shifted_one});
state.builder.AddFunction({OpShiftRightLogical, TypeU32(state), shifted_two, merged_one,
ConstantU32(state, 2)});
state.builder.AddFunction(
{OpBitwiseOr, TypeU32(state), merged_two, merged_one, shifted_two});
state.builder.AddFunction({OpBitwiseAnd, TypeU32(state), quad_bits, merged_two,
ConstantU32(state, 0x11111111u)});
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 high = state.builder.AllocateId();
const auto low = 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 = EmitWqmLaneU32(state, low);
const auto wqm_high = EmitWqmLaneU32(state, high);
const auto wqm_low = EmitWqmWordU32(state, low);
const auto lane = EmitSubgroupLocalInvocationId(state);
const auto upper = state.builder.AllocateId();
const auto mask = state.builder.AllocateId();
const auto bit_lane = state.builder.AllocateId();
uint32_t mask = wqm_low;
uint32_t bit_lane = lane;
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;
@@ -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 &&
inst.dst.reg.file == IR::RegisterFile::Scalar) {
const auto result =
ExpandWholeQuadMask(IR::U64(ReadOperand(inst.src[0], IR::Type::U64)));
if (!current_per_invocation_masks) {
const auto result = IR::U64(ir.Emit(
IR::ValueOpcode::WqmU64, {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})})));
@@ -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;
@@ -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);
@@ -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");
@@ -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.
@@ -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) {
+4 -4
View File
@@ -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, 13, 0x00ff00ffu);
AppendSMovLiteral(&code, 12, 0x80000011u);
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,
0x00ff00ffu},
{7, 11, 19, 35, 0xfffffffeu, 1, 0x7fffffeeu, 0x7ffffffeu, 0xf00000ffu,
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}};
+99 -2
View File
@@ -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,
.instructions = 155,
{.words = 407,
.instructions = 98,
.variables = 4,
.loads = 3,
.stores = 1,
@@ -10519,6 +10615,7 @@ int main() {
TestNewShaderRecompilerSetpcJumpTable();
TestNewShaderRecompilerSetpcDwordJumpTable();
TestTypedEntryStateIsMinimal();
TestWqmConstantPropagation();
TestFinalSsaRejectsRegisterStatePseudos();
TestValuePhiValidation();
TestNativeWideValueValidation();