yaml predicated-prefix-sum

yaml_predicated_prefix_sum.zig

const std = @import("std");

export fn columnCountsVec(chunk: @Vector(16, u8)) @TypeOf(chunk) {
    var mask = chunk != @as(@TypeOf(chunk), @splat('\n'));
    var array_count = @select(u8, mask, @as(@TypeOf(chunk), @splat(1)), @as(@TypeOf(chunk), @splat(0)));

    inline for (0..std.math.log2(@sizeOf(@TypeOf(chunk)))) |i| {
        array_count = @select(u8, mask, array_count +% std.simd.shiftElementsRight(array_count, 1 << i, 0), array_count);
        mask = @select(bool, mask, std.simd.shiftElementsRight(mask, 1 << i, false), @as(@TypeOf(mask), @splat(false)));
    }

    return std.simd.shiftElementsRight(array_count, 1, 0);
}

export fn columnCounts(chunk: @Vector(16, u8)) @Vector(16, u8) {
    const newlines: u16 = @bitCast(@as(@Vector(16, u8), chunk) == @as(@Vector(16, u8), @splat('\n')));
    const ones = 0x1111111111111111;
    const ascending_indices = 0xFEDCBA9876543210;
    const restart_nibbles_mask = pdep(newlines, ones ^ 1) *% 0xF;
    const restart_nibbles_indices = pext(ascending_indices, restart_nibbles_mask);
    const prefix_diff = restart_nibbles_indices -% (restart_nibbles_indices << 4);
    const vec: @Vector(8, u8) = @bitCast(ascending_indices -% pdep(prefix_diff, restart_nibbles_mask) *% ones);
    return std.simd.interlace(.{ vec, vec >> @splat(4) }) & @as(@Vector(16, u8), @splat(0xF));
}

//pub fn main() void {
//    const chunk = [16]u8{ 'a', 'b', 'c', '\n', 'e', 'f', '\n', 'h', 'i', '\n', 'k', 'l', 'm', 'n', 'o', 'p' };
//    std.debug.print("{}\n", .{columnCountsVec(chunk)});
//
//    const result = columnCounts(chunk);
//    std.debug.print("{}\n", .{result});
//}

inline fn pdep(src: u64, mask: u64) u64 {
    return asm ("pdep %[mask], %[src], %[ret]"
        : [ret] "=r" (-> u64),
        : [src] "r" (src),
          [mask] "r" (mask),
    );
}

inline fn pext(src: u64, mask: u64) u64 {
    return asm ("pext %[mask], %[src], %[ret]"
        : [ret] "=r" (-> u64),
        : [src] "r" (src),
          [mask] "r" (mask),
    );
}