tiny.simd.rotate

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

1const builtin = @import("builtin");2const std = @import("std");34pub fn rotateLeft(5    comptime D: type,6    comptime amount: usize,7    value: D.Vector,8) D.Vector {9    validate(D);10    validateAmount(D, amount);11    const amounts: UnsignedVector(D) = @splat(amount);12    return rotateLeftBits(D, value, amounts);13}1415pub fn rotateRight(16    comptime D: type,17    comptime amount: usize,18    value: D.Vector,19) D.Vector {20    validate(D);21    validateAmount(D, amount);22    const amounts: UnsignedVector(D) = @splat(amount);23    return rotateRightBits(D, value, amounts);24}2526pub fn rotateLeftSame(comptime D: type, value: D.Vector, amount: i32) D.Vector {27    validate(D);28    const normalized: UnsignedLane(D) = @intCast(@mod(amount, @bitSizeOf(D.Lane)));29    return rotateLeftBits(D, value, @splat(normalized));30}3132pub fn rotateRightSame(comptime D: type, value: D.Vector, amount: i32) D.Vector {33    validate(D);34    const normalized: UnsignedLane(D) = @intCast(@mod(amount, @bitSizeOf(D.Lane)));35    return rotateRightBits(D, value, @splat(normalized));36}3738pub fn rol(comptime D: type, value: D.Vector, amounts: D.Vector) D.Vector {39    validate(D);40    const raw: UnsignedVector(D) = @bitCast(amounts);41    return rotateLeftBits(D, value, raw);42}4344pub fn ror(comptime D: type, value: D.Vector, amounts: D.Vector) D.Vector {45    validate(D);46    const raw: UnsignedVector(D) = @bitCast(amounts);47    return rotateRightBits(D, value, raw);48}4950pub fn multiRotateRight(51    comptime D: type,52    value: D.Vector,53    indices: D.repartition(u8).Vector,54) D.Vector {55    validate(D);56    if (@bitSizeOf(D.Lane) != 64) @compileError("multiRotateRight requires 64-bit lanes");57    const bits: @Vector(D.lane_count, u64) = @bitCast(value);58    const index_lanes: [D.lane_count * 8]u8 = indices;59    var result: @Vector(D.lane_count, u64) = @splat(0);60    inline for (0..D.lane_count) |lane_index| {61        var lane: u64 = 0;62        inline for (0..8) |byte_index| {63            const amount = index_lanes[lane_index * 8 + byte_index] & 63;64            const byte: u8 = @truncate(std.math.rotr(u64, bits[lane_index], amount));65            const destination = if (builtin.cpu.arch.endian() == .little)66                byte_index67            else68                byte_index ^ 7;69            lane |= @as(u64, byte) << @intCast(destination * 8);70        }71        result[lane_index] = lane;72    }73    return @bitCast(result);74}7576fn rotateLeftBits(77    comptime D: type,78    value: D.Vector,79    raw_amounts: UnsignedVector(D),80) D.Vector {81    const bits: UnsignedVector(D) = @bitCast(value);82    const zero: UnsignedVector(D) = @splat(0);83    const left: AmountVector(D) = @truncate(raw_amounts);84    const right: AmountVector(D) = @truncate(zero -% raw_amounts);85    return @bitCast((bits << left) | (bits >> right));86}8788fn rotateRightBits(89    comptime D: type,90    value: D.Vector,91    raw_amounts: UnsignedVector(D),92) D.Vector {93    const bits: UnsignedVector(D) = @bitCast(value);94    const zero: UnsignedVector(D) = @splat(0);95    const right: AmountVector(D) = @truncate(raw_amounts);96    const left: AmountVector(D) = @truncate(zero -% raw_amounts);97    return @bitCast((bits >> right) | (bits << left));98}99100fn UnsignedLane(comptime D: type) type {101    return @Int(.unsigned, @bitSizeOf(D.Lane));102}103104fn UnsignedVector(comptime D: type) type {105    return @Vector(D.lane_count, UnsignedLane(D));106}107108fn AmountVector(comptime D: type) type {109    return @Vector(D.lane_count, std.math.Log2Int(UnsignedLane(D)));110}111112fn validate(comptime D: type) void {113    if (@typeInfo(D.Lane) != .int) @compileError("rotations require integer lanes");114}115116fn validateAmount(comptime D: type, comptime amount: usize) void {117    if (amount >= @bitSizeOf(D.Lane)) @compileError("rotation amount exceeds lane width");118}119120fn verifyLaneType(comptime T: type) !void {121    const simd = @import("root.zig");122    const D = simd.FixedTag(T, 4);123    const U = @Int(.unsigned, @bitSizeOf(T));124    const UV = @Vector(4, U);125    const raw: UV = .{ 0, 1, @as(U, 1) << (@bitSizeOf(T) - 1), std.math.maxInt(U) };126    const value: D.Vector = @bitCast(raw);127    try std.testing.expect(@reduce(.And, rotateRight(D, 1, rotateLeft(D, 1, value)) == value));128    try std.testing.expect(@reduce(.And, rotateRightSame(D, rotateLeftSame(D, value, -3), -3) == value));129}130131test "Highway rotations instantiate every integer lane type" {132    inline for (.{ u8, i8, u16, i16, u32, i32, u64, i64 }) |T| {133        try verifyLaneType(T);134    }135}136137test "Highway fixed rotations wrap every shifted bit" {138    const simd = @import("root.zig");139    const D = simd.FixedTag(u8, 4);140    const value: D.Vector = .{ 0x01, 0x80, 0x55, 0x81 };141    try std.testing.expect(@reduce(.And, rotateLeft(D, 1, value) ==142        @as(D.Vector, .{ 0x02, 0x01, 0xaa, 0x03 })));143    try std.testing.expect(@reduce(.And, rotateRight(D, 1, value) ==144        @as(D.Vector, .{ 0x80, 0x40, 0xaa, 0xc0 })));145    try std.testing.expect(@reduce(.And, rotateLeft(D, 0, value) == value));146}147148test "Highway variable rotations mask signed amounts" {149    const simd = @import("root.zig");150    const D = simd.FixedTag(i16, 4);151    const U = simd.FixedTag(u16, 4);152    const value: D.Vector = @bitCast(@as(U.Vector, .{ 1, 0x8000, 0x1234, 0x55aa }));153    const amounts: D.Vector = .{ 0, 1, -4, 17 };154    const round_trip = ror(D, rol(D, value, amounts), amounts);155    try std.testing.expect(@reduce(.And, round_trip == value));156    try std.testing.expect(@reduce(.And, rotateLeftSame(D, value, -1) == rotateRight(D, 1, value)));157}158159test "Highway multi rotate composes bytes within each u64 lane" {160    const simd = @import("root.zig");161    const D = simd.FixedTag(u64, 2);162    const DI = D.repartition(u8);163    const value: D.Vector = .{ 0x0102_0304_0506_0708, 0x1020_3040_5060_7080 };164    const indices: DI.Vector = .{165        0,  8,  16, 24, 32, 40, 48, 56,166        56, 48, 40, 32, 24, 16, 8,  0,167    };168    const expected: D.Vector = if (builtin.cpu.arch.endian() == .little)169        .{ value[0], @byteSwap(value[1]) }170    else171        .{ @byteSwap(value[0]), value[1] };172    try std.testing.expect(@reduce(.And, multiRotateRight(D, value, indices) == expected));173}