tiny.simd.compact

1const std = @import("std");23pub fn compress(comptime D: type, value: D.Vector, mask: D.Mask) D.Vector {4    const lanes: [D.lane_count]D.Lane = @bitCast(value);5    var result: [D.lane_count]D.Lane = undefined;6    var position: usize = 0;7    inline for (0..D.lane_count) |index| {8        if (mask[index]) {9            result[position] = lanes[index];10            position += 1;11        }12    }13    inline for (0..D.lane_count) |index| {14        if (!mask[index]) {15            result[position] = lanes[index];16            position += 1;17        }18    }19    std.debug.assert(position == D.lane_count);20    return @bitCast(result);21}2223pub fn compressIsPartition(comptime D: type) bool {24    _ = D;25    return true;26}2728pub fn compressNot(comptime D: type, value: D.Vector, mask: D.Mask) D.Vector {29    return compress(D, value, !mask);30}3132pub fn compressBlocksNot(comptime D: type, value: D.Vector, mask: D.Mask) D.Vector {33    if (D.Lane != u64) @compileError("compressBlocksNot requires u64 lanes");34    if (D.lane_count < 2 or D.lane_count % 2 != 0) {35        @compileError("compressBlocksNot requires complete 128-bit blocks");36    }37    return compressNot(D, value, mask);38}3940pub fn compressStore(41    comptime D: type,42    value: D.Vector,43    mask: D.Mask,44    output: []D.Lane,45) usize {46    const count = countTrue(D, mask);47    std.debug.assert(output.len >= count);48    const compacted = compress(D, value, mask);49    inline for (0..D.lane_count) |index| {50        if (index < count) output[index] = compacted[index];51    }52    return count;53}5455pub fn compressBlendedStore(56    comptime D: type,57    value: D.Vector,58    mask: D.Mask,59    output: []D.Lane,60) usize {61    return compressStore(D, value, mask, output);62}6364pub fn compressBits(65    comptime D: type,66    value: D.Vector,67    bits: []const u8,68) D.Vector {69    return compress(D, value, @import("compare.zig").loadMaskBits(D, bits));70}7172pub fn compressBitsStore(73    comptime D: type,74    value: D.Vector,75    bits: []const u8,76    output: []D.Lane,77) usize {78    return compressStore(D, value, @import("compare.zig").loadMaskBits(D, bits), output);79}8081pub fn expand(comptime D: type, value: D.Vector, mask: D.Mask) D.Vector {82    const lanes: [D.lane_count]D.Lane = @bitCast(value);83    var result: [D.lane_count]D.Lane = @splat(0);84    var position: usize = 0;85    inline for (0..D.lane_count) |index| {86        if (mask[index]) {87            result[index] = lanes[position];88            position += 1;89        }90    }91    return @bitCast(result);92}9394pub fn loadExpand(95    comptime D: type,96    mask: D.Mask,97    input: []const D.Lane,98) D.Vector {99    const count = countTrue(D, mask);100    std.debug.assert(input.len >= count);101    var result: D.Vector = @splat(0);102    var position: usize = 0;103    inline for (0..D.lane_count) |index| {104        if (mask[index]) {105            result[index] = input[position];106            position += 1;107        }108    }109    return result;110}111112fn countTrue(comptime D: type, mask: D.Mask) usize {113    return @import("compare.zig").countTrue(D, mask);114}115116fn verifyLaneType(comptime T: type) !void {117    const simd = @import("root.zig");118    const D = simd.FixedTag(T, 8);119    const value: D.Vector = @splat(0);120    const mask: D.Mask = .{ true, false, true, false, false, true, false, true };121    const compacted = compress(D, value, mask);122    try std.testing.expect(@reduce(.And, expand(D, compacted, mask) == value));123    var stored: [D.lane_count]T = undefined;124    try std.testing.expectEqual(@as(usize, 4), compressStore(D, value, mask, &stored));125}126127test "Highway compact operations instantiate every lane type" {128    inline for (.{ u8, i8, u16, i16, u32, i32, u64, i64, f16, f32, f64 }) |T| {129        try verifyLaneType(T);130    }131    try std.testing.expect(compressIsPartition(@import("root.zig").FixedTag(u32, 8)));132}133134test "Highway compress is a stable partition and stores its selected prefix" {135    const simd = @import("root.zig");136    const D = simd.FixedTag(i32, 8);137    const value: D.Vector = .{ 10, 11, 12, 13, 14, 15, 16, 17 };138    const mask: D.Mask = .{ false, true, true, false, true, false, false, true };139    const expected: D.Vector = .{ 11, 12, 14, 17, 10, 13, 15, 16 };140    try std.testing.expect(@reduce(.And, compress(D, value, mask) == expected));141    try std.testing.expect(@reduce(.And, compressNot(D, value, !mask) == expected));142143    var output = @as([8]i32, @splat(99));144    try std.testing.expectEqual(@as(usize, 4), compressBlendedStore(D, value, mask, &output));145    try std.testing.expectEqualSlices(i32, &.{ 11, 12, 14, 17, 99, 99, 99, 99 }, &output);146147    var bits = [_]u8{0x96};148    try std.testing.expect(@reduce(.And, compressBits(D, value, &bits) == expected));149    @memset(&output, 99);150    try std.testing.expectEqual(@as(usize, 4), compressBitsStore(D, value, &bits, &output));151    try std.testing.expectEqualSlices(i32, &.{ 11, 12, 14, 17, 99, 99, 99, 99 }, &output);152}153154test "Highway expand scatters consecutive lanes and loadExpand reads only selected input" {155    const simd = @import("root.zig");156    const D = simd.FixedTag(u16, 8);157    const compacted: D.Vector = .{ 20, 30, 50, 80, 91, 92, 93, 94 };158    const mask: D.Mask = .{ false, true, true, false, true, false, false, true };159    const expected: D.Vector = .{ 0, 20, 30, 0, 50, 0, 0, 80 };160    try std.testing.expect(@reduce(.And, expand(D, compacted, mask) == expected));161    try std.testing.expect(@reduce(.And, loadExpand(D, mask, &.{ 20, 30, 50, 80 }) == expected));162}163164test "Highway block compression retains complete u64 pairs" {165    const simd = @import("root.zig");166    const D = simd.FixedTag(u64, 8);167    const value: D.Vector = .{ 0, 1, 2, 3, 4, 5, 6, 7 };168    const not_mask: D.Mask = .{ true, true, false, false, true, true, false, false };169    try std.testing.expect(@reduce(.And, compressBlocksNot(D, value, not_mask) ==170        @as(D.Vector, .{ 2, 3, 6, 7, 0, 1, 4, 5 })));171}

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