tiny.simd.shift

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

1const std = @import("std");23pub fn shiftLeft(4    comptime D: type,5    comptime amount: usize,6    value: D.Vector,7) D.Vector {8    validate(D);9    validateAmount(D, amount);10    const bits: UnsignedVector(D) = @bitCast(value);11    return @bitCast(bits << sameAmount(D, amount));12}1314pub fn shiftRight(15    comptime D: type,16    comptime amount: usize,17    value: D.Vector,18) D.Vector {19    validate(D);20    validateAmount(D, amount);21    return value >> sameAmount(D, amount);22}2324pub fn shiftLeftSame(comptime D: type, value: D.Vector, amount: usize) D.Vector {25    validate(D);26    std.debug.assert(amount < @bitSizeOf(D.Lane));27    const bits: UnsignedVector(D) = @bitCast(value);28    return @bitCast(bits << sameAmount(D, amount));29}3031pub fn shiftRightSame(comptime D: type, value: D.Vector, amount: usize) D.Vector {32    validate(D);33    std.debug.assert(amount < @bitSizeOf(D.Lane));34    return value >> sameAmount(D, amount);35}3637pub fn shl(comptime D: type, value: D.Vector, amounts: D.Vector) D.Vector {38    validate(D);39    const bits: UnsignedVector(D) = @bitCast(value);40    return @bitCast(bits << laneAmounts(D, amounts));41}4243pub fn shr(comptime D: type, value: D.Vector, amounts: D.Vector) D.Vector {44    validate(D);45    return value >> laneAmounts(D, amounts);46}4748pub fn averageRound(comptime D: type, a: D.Vector, b: D.Vector) D.Vector {49    validate(D);50    return (a | b) -% ((a ^ b) >> sameAmount(D, 1));51}5253pub fn roundingShiftRight(54    comptime D: type,55    comptime amount: usize,56    value: D.Vector,57) D.Vector {58    validateAmount(D, amount);59    if (amount == 0) return value;60    return averageRound(D, shiftRight(D, amount - 1, value), @splat(0));61}6263pub fn roundingShiftRightSame(64    comptime D: type,65    value: D.Vector,66    amount: usize,67) D.Vector {68    std.debug.assert(amount < @bitSizeOf(D.Lane));69    if (amount == 0) return value;70    return averageRound(D, shiftRightSame(D, value, amount - 1), @splat(0));71}7273pub fn roundingShr(comptime D: type, value: D.Vector, amounts: D.Vector) D.Vector {74    validate(D);75    const U = UnsignedVector(D);76    const raw: U = @bitCast(amounts);77    const zero: U = @splat(0);78    const one: U = @splat(1);79    const scaled_amounts: U = @select(UnsignedLane(D), raw == zero, zero, raw -% one);80    const scaled = shr(D, value, @bitCast(scaled_amounts));81    const other = @select(D.Lane, raw == zero, scaled, @as(D.Vector, @splat(0)));82    return averageRound(D, scaled, other);83}8485pub fn maskedShiftLeftOr(86    comptime D: type,87    comptime amount: usize,88    no: D.Vector,89    mask: D.Mask,90    value: D.Vector,91) D.Vector {92    return @select(D.Lane, mask, shiftLeft(D, amount, value), no);93}9495pub fn maskedShiftLeft(96    comptime D: type,97    comptime amount: usize,98    mask: D.Mask,99    value: D.Vector,100) D.Vector {101    return maskedShiftLeftOr(D, amount, @splat(0), mask, value);102}103104pub fn maskedShiftRightOr(105    comptime D: type,106    comptime amount: usize,107    no: D.Vector,108    mask: D.Mask,109    value: D.Vector,110) D.Vector {111    return @select(D.Lane, mask, shiftRight(D, amount, value), no);112}113114pub fn maskedShiftRight(115    comptime D: type,116    comptime amount: usize,117    mask: D.Mask,118    value: D.Vector,119) D.Vector {120    return maskedShiftRightOr(D, amount, @splat(0), mask, value);121}122123pub fn maskedShlOr(124    comptime D: type,125    no: D.Vector,126    mask: D.Mask,127    value: D.Vector,128    amounts: D.Vector,129) D.Vector {130    return @select(D.Lane, mask, shl(D, value, amounts), no);131}132133pub fn maskedShl(134    comptime D: type,135    mask: D.Mask,136    value: D.Vector,137    amounts: D.Vector,138) D.Vector {139    return maskedShlOr(D, @splat(0), mask, value, amounts);140}141142pub fn maskedShrOr(143    comptime D: type,144    no: D.Vector,145    mask: D.Mask,146    value: D.Vector,147    amounts: D.Vector,148) D.Vector {149    return @select(D.Lane, mask, shr(D, value, amounts), no);150}151152pub fn maskedShr(153    comptime D: type,154    mask: D.Mask,155    value: D.Vector,156    amounts: D.Vector,157) D.Vector {158    return maskedShrOr(D, @splat(0), mask, value, amounts);159}160161fn UnsignedLane(comptime D: type) type {162    return @Int(.unsigned, @bitSizeOf(D.Lane));163}164165fn UnsignedVector(comptime D: type) type {166    return @Vector(D.lane_count, UnsignedLane(D));167}168169fn AmountVector(comptime D: type) type {170    return @Vector(D.lane_count, std.math.Log2Int(UnsignedLane(D)));171}172173fn sameAmount(comptime D: type, amount: usize) AmountVector(D) {174    return @splat(@intCast(amount));175}176177fn laneAmounts(comptime D: type, amounts: D.Vector) AmountVector(D) {178    const raw: UnsignedVector(D) = @bitCast(amounts);179    return @truncate(raw);180}181182fn validate(comptime D: type) void {183    if (@typeInfo(D.Lane) != .int) @compileError("shifts require integer lanes");184}185186fn validateAmount(comptime D: type, comptime amount: usize) void {187    validate(D);188    if (amount >= @bitSizeOf(D.Lane)) @compileError("shift amount exceeds lane width");189}190191fn verifyLaneType(comptime T: type) !void {192    const simd = @import("root.zig");193    const D = simd.FixedTag(T, 4);194    const U = @Int(.unsigned, @bitSizeOf(T));195    const UV = @Vector(4, U);196    const raw: UV = .{ 0, 1, @as(U, 1) << (@bitSizeOf(T) - 1), std.math.maxInt(U) };197    const value: D.Vector = @bitCast(raw);198    const amounts: D.Vector = @splat(@as(T, 1));199    try std.testing.expect(@reduce(.And, shiftLeft(D, 1, value) == shl(D, value, amounts)));200    try std.testing.expect(@reduce(.And, shiftRight(D, 1, value) == shr(D, value, amounts)));201    try std.testing.expect(@reduce(.And, shiftLeftSame(D, value, 1) == shiftLeft(D, 1, value)));202    try std.testing.expect(@reduce(.And, shiftRightSame(D, value, 1) == shiftRight(D, 1, value)));203    try std.testing.expect(@reduce(.And, roundingShiftRight(D, 0, value) == value));204}205206test "Highway shifts instantiate every integer lane type" {207    inline for (.{ u8, i8, u16, i16, u32, i32, u64, i64 }) |T| {208        try verifyLaneType(T);209    }210}211212test "Highway fixed and variable shifts preserve signed semantics" {213    const simd = @import("root.zig");214    const D = simd.FixedTag(i8, 4);215    const value: D.Vector = .{ -128, -3, 2, 127 };216    const amounts: D.Vector = .{ 0, 1, 2, 7 };217    try std.testing.expect(@reduce(.And, shiftLeft(D, 1, value) ==218        @as(D.Vector, .{ 0, -6, 4, -2 })));219    try std.testing.expect(@reduce(.And, shiftRight(D, 1, value) ==220        @as(D.Vector, .{ -64, -2, 1, 63 })));221    try std.testing.expect(@reduce(.And, shl(D, value, amounts) ==222        @as(D.Vector, .{ -128, -6, 8, -128 })));223    try std.testing.expect(@reduce(.And, shr(D, value, amounts) ==224        @as(D.Vector, .{ -128, -2, 0, 0 })));225}226227test "Highway rounding shifts round upward at the half bit" {228    const simd = @import("root.zig");229    const DU = simd.FixedTag(u8, 4);230    const DI = simd.FixedTag(i8, 4);231    try std.testing.expect(@reduce(.And, roundingShiftRight(DU, 1, @as(DU.Vector, .{ 0, 1, 2, 255 })) ==232        @as(DU.Vector, .{ 0, 1, 1, 128 })));233    try std.testing.expect(@reduce(.And, roundingShiftRight(DI, 1, @as(DI.Vector, .{ -128, -3, 2, 127 })) ==234        @as(DI.Vector, .{ -64, -1, 1, 64 })));235    const value: DU.Vector = .{ 0, 3, 7, 255 };236    const amounts: DU.Vector = .{ 0, 1, 2, 7 };237    try std.testing.expect(@reduce(.And, roundingShiftRightSame(DU, value, 2) ==238        @as(DU.Vector, .{ 0, 1, 2, 64 })));239    try std.testing.expect(@reduce(.And, roundingShr(DU, value, amounts) ==240        @as(DU.Vector, .{ 0, 2, 2, 2 })));241}242243test "Highway masked shifts preserve inactive lanes" {244    const simd = @import("root.zig");245    const D = simd.FixedTag(u16, 4);246    const value: D.Vector = .{ 1, 2, 3, 4 };247    const mask: D.Mask = .{ true, false, true, false };248    const no: D.Vector = @splat(9);249    try std.testing.expect(@reduce(.And, maskedShiftLeftOr(D, 2, no, mask, value) ==250        @as(D.Vector, .{ 4, 9, 12, 9 })));251    try std.testing.expect(@reduce(.And, maskedShiftRight(D, 1, mask, value) ==252        @as(D.Vector, .{ 0, 0, 1, 0 })));253    try std.testing.expect(@reduce(.And, maskedShiftLeft(D, 2, mask, value) ==254        @as(D.Vector, .{ 4, 0, 12, 0 })));255    const amounts: D.Vector = .{ 0, 1, 2, 3 };256    try std.testing.expect(@reduce(.And, maskedShl(D, mask, value, amounts) ==257        @as(D.Vector, .{ 1, 0, 12, 0 })));258    try std.testing.expect(@reduce(.And, maskedShrOr(D, no, mask, value, amounts) ==259        @as(D.Vector, .{ 1, 9, 0, 9 })));260    try std.testing.expect(@reduce(.And, maskedShr(D, mask, value, amounts) ==261        @as(D.Vector, .{ 1, 0, 0, 0 })));262}