tiny.simd.conditional

1const std = @import("std");2const arithmetic = @import("arithmetic.zig");34pub fn ifThenElse(comptime D: type, mask: D.Mask, yes: D.Vector, no: D.Vector) D.Vector {5    return @select(D.Lane, mask, yes, no);6}78pub fn ifThenElseZero(comptime D: type, mask: D.Mask, yes: D.Vector) D.Vector {9    return ifThenElse(D, mask, yes, @as(D.Vector, @splat(0)));10}1112pub fn ifThenZeroElse(comptime D: type, mask: D.Mask, no: D.Vector) D.Vector {13    return ifThenElse(D, mask, @as(D.Vector, @splat(0)), no);14}1516pub fn ifVecThenElse(comptime D: type, mask: D.Vector, yes: D.Vector, no: D.Vector) D.Vector {17    return bitwiseIfThenElse(D, mask, yes, no);18}1920pub fn bitwiseIfThenElse(comptime D: type, mask: D.Vector, yes: D.Vector, no: D.Vector) D.Vector {21    const U = @Int(.unsigned, @bitSizeOf(D.Lane));22    const DU = D.rebind(U);23    const mask_bits: DU.Vector = @bitCast(mask);24    const yes_bits: DU.Vector = @bitCast(yes);25    const no_bits: DU.Vector = @bitCast(no);26    return @bitCast((mask_bits & yes_bits) | (~mask_bits & no_bits));27}2829pub fn zeroIfNegative(comptime D: type, value: D.Vector) D.Vector {30    return ifThenZeroElse(D, negativeMask(D, value), value);31}3233pub fn ifNegativeThenElse(34    comptime D: type,35    sign: D.Vector,36    negative: D.Vector,37    nonnegative: D.Vector,38) D.Vector {39    return ifThenElse(D, negativeMask(D, sign), negative, nonnegative);40}4142pub fn ifNegativeThenElseZero(comptime D: type, sign: D.Vector, negative: D.Vector) D.Vector {43    return ifThenElseZero(D, negativeMask(D, sign), negative);44}4546pub fn ifNegativeThenZeroElse(comptime D: type, sign: D.Vector, nonnegative: D.Vector) D.Vector {47    return ifThenZeroElse(D, negativeMask(D, sign), nonnegative);48}4950pub fn ifNegativeThenNegOrUndefIfZero(51    comptime D: type,52    sign: D.Vector,53    value: D.Vector,54) D.Vector {55    return ifThenElse(D, negativeMask(D, sign), arithmetic.neg(D, value), value);56}5758fn negativeMask(comptime D: type, value: D.Vector) D.Mask {59    return switch (@typeInfo(D.Lane)) {60        .int => |info| if (info.signedness == .signed)61            value < @as(D.Vector, @splat(0))62        else63            @compileError("negative selection requires signed integer or floating-point lanes"),64        .float => blk: {65            const U = @Int(.unsigned, @bitSizeOf(D.Lane));66            const DU = D.rebind(U);67            const bits: DU.Vector = @bitCast(value);68            break :blk bits & @as(DU.Vector, @splat(@as(U, 1) << (@bitSizeOf(U) - 1))) !=69                @as(DU.Vector, @splat(0));70        },71        else => unreachable,72    };73}7475test "Highway conditional selection covers masks vectors and raw bit patterns" {76    const simd = @import("root.zig");77    const D = simd.FixedTag(u32, 4);78    const yes: D.Vector = .{ 1, 2, 3, 4 };79    const no: D.Vector = .{ 10, 20, 30, 40 };80    const mask: D.Mask = .{ true, false, true, false };81    try std.testing.expect(@reduce(.And, ifThenElse(D, mask, yes, no) ==82        @as(D.Vector, .{ 1, 20, 3, 40 })));83    try std.testing.expect(@reduce(.And, ifThenElseZero(D, mask, yes) ==84        @as(D.Vector, .{ 1, 0, 3, 0 })));85    try std.testing.expect(@reduce(.And, ifThenZeroElse(D, mask, no) ==86        @as(D.Vector, .{ 0, 20, 0, 40 })));87    const vector_mask: D.Vector = .{ std.math.maxInt(u32), 0, std.math.maxInt(u32), 0 };88    try std.testing.expect(@reduce(.And, ifVecThenElse(D, vector_mask, yes, no) ==89        @as(D.Vector, .{ 1, 20, 3, 40 })));90    const partial: D.Vector = @splat(0x0f0f_0f0f);91    try std.testing.expect(@reduce(.And, bitwiseIfThenElse(D, partial, @as(D.Vector, @splat(0xaaaa_aaaa)), @as(D.Vector, @splat(0x5555_5555))) ==92        @as(D.Vector, @splat(0x5a5a_5a5a))));93}9495test "Highway sign conditionals use representation sign including negative zero" {96    const simd = @import("root.zig");97    const D = simd.FixedTag(f32, 4);98    const sign: D.Vector = .{ -1, -0.0, 0, 1 };99    const value: D.Vector = .{ 1, 2, 3, 4 };100    try std.testing.expect(@reduce(.And, zeroIfNegative(D, sign) ==101        @as(D.Vector, .{ 0, 0, 0, 1 })));102    try std.testing.expect(@reduce(.And, ifNegativeThenElse(D, sign, value, @as(D.Vector, @splat(9))) ==103        @as(D.Vector, .{ 1, 2, 9, 9 })));104    try std.testing.expect(@reduce(.And, ifNegativeThenElseZero(D, sign, value) ==105        @as(D.Vector, .{ 1, 2, 0, 0 })));106    try std.testing.expect(@reduce(.And, ifNegativeThenZeroElse(D, sign, value) ==107        @as(D.Vector, .{ 0, 0, 3, 4 })));108    try std.testing.expect(@reduce(.And, ifNegativeThenNegOrUndefIfZero(D, sign, value) ==109        @as(D.Vector, .{ -1, -2, 3, 4 })));110}111112test "Highway sign conditionals instantiate signed integer and floating lanes" {113    const simd = @import("root.zig");114    inline for (.{ i8, i16, i32, i64, f16, f32, f64 }) |T| {115        const D = simd.FixedTag(T, 4);116        const value: D.Vector = @splat(1);117        _ = zeroIfNegative(D, value);118        _ = ifNegativeThenElse(D, value, value, value);119        _ = ifNegativeThenNegOrUndefIfZero(D, value, value);120    }121}

1pub const conditional = @import("conditional.zig");