tiny.simd.indexed

1const std = @import("std");23pub fn scatterIndex(4    comptime D: type,5    value: D.Vector,6    base: []D.Lane,7    indices: anytype,8) void {9    validateIndices(D, @TypeOf(indices));10    inline for (0..D.lane_count) |lane_index| {11        const destination = checkedIndex(indices[lane_index]);12        std.debug.assert(destination < base.len);13        base[destination] = value[lane_index];14    }15}1617pub fn scatterIndexN(18    comptime D: type,19    value: D.Vector,20    base: []D.Lane,21    indices: anytype,22    count: usize,23) void {24    validateIndices(D, @TypeOf(indices));25    const lanes = @min(count, D.lane_count);26    inline for (0..D.lane_count) |lane_index| {27        if (lane_index < lanes) {28            const destination = checkedIndex(indices[lane_index]);29            std.debug.assert(destination < base.len);30            base[destination] = value[lane_index];31        }32    }33}3435pub fn maskedScatterIndex(36    comptime D: type,37    value: D.Vector,38    mask: D.Mask,39    base: []D.Lane,40    indices: anytype,41) void {42    validateIndices(D, @TypeOf(indices));43    inline for (0..D.lane_count) |lane_index| {44        if (mask[lane_index]) {45            const destination = checkedIndex(indices[lane_index]);46            std.debug.assert(destination < base.len);47            base[destination] = value[lane_index];48        }49    }50}5152pub fn scatterOffset(53    comptime D: type,54    value: D.Vector,55    base: []D.Lane,56    offsets: anytype,57) void {58    validateIndices(D, @TypeOf(offsets));59    inline for (0..D.lane_count) |lane_index| {60        const destination = offsetIndex(D.Lane, offsets[lane_index]);61        std.debug.assert(destination < base.len);62        base[destination] = value[lane_index];63    }64}6566pub fn gatherIndex(comptime D: type, base: []const D.Lane, indices: anytype) D.Vector {67    validateIndices(D, @TypeOf(indices));68    var result: D.Vector = undefined;69    inline for (0..D.lane_count) |lane_index| {70        const source = checkedIndex(indices[lane_index]);71        std.debug.assert(source < base.len);72        result[lane_index] = base[source];73    }74    return result;75}7677pub fn gatherIndexN(78    comptime D: type,79    base: []const D.Lane,80    indices: anytype,81    count: usize,82) D.Vector {83    return gatherIndexNOr(D, @splat(0), base, indices, count);84}8586pub fn gatherIndexNOr(87    comptime D: type,88    inactive: D.Vector,89    base: []const D.Lane,90    indices: anytype,91    count: usize,92) D.Vector {93    const lanes = @min(count, D.lane_count);94    var mask: D.Mask = @splat(false);95    inline for (0..D.lane_count) |index| mask[index] = index < lanes;96    return maskedGatherIndexOr(D, inactive, mask, base, indices);97}9899pub fn maskedGatherIndexOr(100    comptime D: type,101    inactive: D.Vector,102    mask: D.Mask,103    base: []const D.Lane,104    indices: anytype,105) D.Vector {106    validateIndices(D, @TypeOf(indices));107    var result = inactive;108    inline for (0..D.lane_count) |lane_index| {109        if (mask[lane_index]) {110            const source = checkedIndex(indices[lane_index]);111            std.debug.assert(source < base.len);112            result[lane_index] = base[source];113        }114    }115    return result;116}117118pub fn maskedGatherIndex(119    comptime D: type,120    mask: D.Mask,121    base: []const D.Lane,122    indices: anytype,123) D.Vector {124    return maskedGatherIndexOr(D, @splat(0), mask, base, indices);125}126127pub fn gatherOffset(comptime D: type, base: []const D.Lane, offsets: anytype) D.Vector {128    validateIndices(D, @TypeOf(offsets));129    var result: D.Vector = undefined;130    inline for (0..D.lane_count) |lane_index| {131        const source = offsetIndex(D.Lane, offsets[lane_index]);132        std.debug.assert(source < base.len);133        result[lane_index] = base[source];134    }135    return result;136}137138fn checkedIndex(value: anytype) usize {139    const T = @TypeOf(value);140    if (comptime @typeInfo(T) != .int) @compileError("indices must be integers");141    if (comptime @typeInfo(T).int.signedness == .signed) std.debug.assert(value >= 0);142    return @intCast(value);143}144145fn offsetIndex(comptime T: type, offset: anytype) usize {146    const bytes = checkedIndex(offset);147    std.debug.assert(bytes % @sizeOf(T) == 0);148    return bytes / @sizeOf(T);149}150151fn validateIndices(comptime D: type, comptime V: type) void {152    if (comptime @typeInfo(V) != .vector or @typeInfo(V).vector.len != D.lane_count or153        @typeInfo(@typeInfo(V).vector.child) != .int)154    {155        @compileError("indices must be an integer vector with one index per lane");156    }157}158159test "Highway indexed gather and scatter use lane order and clamped counts" {160    const simd = @import("root.zig");161    const D = simd.FixedTag(u32, 4);162    const I = simd.FixedTag(i32, 4);163    const indices: I.Vector = .{ 3, 0, 5, 2 };164    const source = [_]u32{ 10, 11, 12, 13, 14, 15 };165    try std.testing.expect(@reduce(.And, gatherIndex(D, &source, indices) ==166        @as(D.Vector, .{ 13, 10, 15, 12 })));167    try std.testing.expect(@reduce(.And, gatherIndexN(D, &source, indices, 2) ==168        @as(D.Vector, .{ 13, 10, 0, 0 })));169    var output = [_]u32{ 0, 0, 0, 0, 0, 0 };170    scatterIndexN(D, .{ 1, 2, 3, 4 }, &output, indices, 3);171    try std.testing.expectEqualSlices(u32, &.{ 2, 0, 0, 1, 0, 3 }, &output);172}173174test "Highway masked and byte-offset memory avoids inactive indices" {175    const simd = @import("root.zig");176    const D = simd.FixedTag(f64, 2);177    const I = simd.FixedTag(i64, 2);178    const source = [_]f64{ 1, 2, 3, 4 };179    const mask: D.Mask = .{ true, false };180    const indices: I.Vector = .{ 2, -1 };181    try std.testing.expect(@reduce(.And, maskedGatherIndexOr(D, @splat(9), mask, &source, indices) ==182        @as(D.Vector, .{ 3, 9 })));183    const offsets: I.Vector = .{ 24, 0 };184    try std.testing.expect(@reduce(.And, gatherOffset(D, &source, offsets) ==185        @as(D.Vector, .{ 4, 1 })));186    var output = [_]f64{ 0, 0, 0, 0 };187    scatterOffset(D, .{ 7, 8 }, &output, offsets);188    try std.testing.expectEqualSlices(f64, &.{ 8, 0, 0, 7 }, &output);189}

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