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) {