Compiler Explorer

Source code

const std = @import("std");

fn simulated_prefix_sum_exclusive(chunks: [4]@Vector(16, u8)) [4]@Vector(16, u8) {
    return simulated_prefix_sum_inclusive(.{ std.simd.shiftElementsRight(chunks[3], 1, 0), chunks[0], chunks[1], chunks[2] });
}

fn simulated_prefix_sum_inclusive(chunks: [4]@Vector(16, u8)) [4]@Vector(16, u8) {
    const a, const b, const c, const d = chunks;
    // shift by 1
    const e = a + std.simd.shiftElementsRight(d, 1, 0);
    const f = b + a;
    const g = c + b;
    const h = d + c;

// shift by 2
    const i = e + std.simd.shiftElementsRight(g, 1, 0);
    const j = f + std.simd.shiftElementsRight(h, 1, 0);
    const k = g + e;
    const l = h + f;

var ans = [4]@Vector(16, u8){ i, j, k, l };
    for (&ans) |*m| {
        inline for (0..4) |n| { // emulate shift by 4, 8, 16, and 32
            m.* += std.simd.shiftElementsRight(m.*, 1 << n, 0); 
        }
    }

return ans;
}

export fn prefixSumInterleavedInclusive(a: @Vector(64, u8)) @Vector(64, u8) {
    return @bitCast(simulated_prefix_sum_inclusive(@bitCast(a)));
}