tiny.simd.reduce

1const std = @import("std");23pub fn sum(comptime D: type, value: D.Vector) D.Lane {4    var result: D.Lane = 0;5    inline for (0..D.lane_count) |index| {6        result = if (comptime @typeInfo(D.Lane) == .int)7            result +% value[index]8        else9            result + value[index];10    }11    return result;12}1314pub fn min(comptime D: type, value: D.Vector) D.Lane {15    var result = value[0];16    inline for (1..D.lane_count) |index| {17        result = @min(result, value[index]);18    }19    return result;20}2122pub fn max(comptime D: type, value: D.Vector) D.Lane {23    var result = value[0];24    inline for (1..D.lane_count) |index| {25        result = @max(result, value[index]);26    }27    return result;28}2930pub fn sumOfLanes(comptime D: type, value: D.Vector) D.Vector {31    return @splat(sum(D, value));32}3334pub fn minOfLanes(comptime D: type, value: D.Vector) D.Vector {35    return @splat(min(D, value));36}3738pub fn maxOfLanes(comptime D: type, value: D.Vector) D.Vector {39    return @splat(max(D, value));40}4142pub fn maskedSum(comptime D: type, mask: D.Mask, value: D.Vector) D.Lane {43    var result: D.Lane = 0;44    inline for (0..D.lane_count) |index| {45        if (mask[index]) {46            result = if (comptime @typeInfo(D.Lane) == .int)47                result +% value[index]48            else49                result + value[index];50        }51    }52    return result;53}5455pub fn maskedMin(comptime D: type, mask: D.Mask, value: D.Vector) D.Lane {56    var result: D.Lane = 0;57    var found = false;58    inline for (0..D.lane_count) |index| {59        if (mask[index]) {60            result = if (found) @min(result, value[index]) else value[index];61            found = true;62        }63    }64    return result;65}6667pub fn maskedMax(comptime D: type, mask: D.Mask, value: D.Vector) D.Lane {68    var result: D.Lane = 0;69    var found = false;70    inline for (0..D.lane_count) |index| {71        if (mask[index]) {72            result = if (found) @max(result, value[index]) else value[index];73            found = true;74        }75    }76    return result;77}7879pub fn sumsOf2(comptime D: type, value: D.Vector) D.repartition(wideLane(D.Lane)).Vector {80    const W = wideLane(D.Lane);81    const R = D.repartition(W);82    if (D.lane_count < 2) @compileError("sumsOf2 requires at least two lanes");83    var result: R.Vector = undefined;84    inline for (0..R.lane_count) |index| {85        result[index] = widen(W, value[index * 2]) + widen(W, value[index * 2 + 1]);86    }87    return result;88}8990pub fn sumsOf4(comptime D: type, value: D.Vector) D.repartition(wideLane(wideLane(D.Lane))).Vector {91    const W = wideLane(wideLane(D.Lane));92    const R = D.repartition(W);93    if (D.lane_count < 4) @compileError("sumsOf4 requires at least four lanes");94    var result: R.Vector = undefined;95    inline for (0..R.lane_count) |index| {96        const base = index * 4;97        result[index] = widen(W, value[base]) + widen(W, value[base + 1]) +98            widen(W, value[base + 2]) + widen(W, value[base + 3]);99    }100    return result;101}102103pub fn sumsOf8(comptime D: type, value: D.Vector) D.repartition(wideLane(wideLane(wideLane(D.Lane)))).Vector {104    const W = wideLane(wideLane(wideLane(D.Lane)));105    const R = D.repartition(W);106    if (D.lane_count < 8) @compileError("sumsOf8 requires at least eight lanes");107    var result: R.Vector = undefined;108    inline for (0..R.lane_count) |index| {109        const base = index * 8;110        result[index] = widen(W, value[base]) + widen(W, value[base + 1]) +111            widen(W, value[base + 2]) + widen(W, value[base + 3]) +112            widen(W, value[base + 4]) + widen(W, value[base + 5]) +113            widen(W, value[base + 6]) + widen(W, value[base + 7]);114    }115    return result;116}117118pub fn sumsOf8AbsDiff(119    comptime D: type,120    a: D.Vector,121    b: D.Vector,122) D.repartition(wideLane(wideLane(wideLane(D.Lane)))).Vector {123    requireByteLane(D, "sumsOf8AbsDiff");124    const W = wideLane(wideLane(wideLane(D.Lane)));125    const R = D.repartition(W);126    if (D.lane_count < 8) @compileError("sumsOf8AbsDiff requires at least eight lanes");127    var result: R.Vector = undefined;128    inline for (0..R.lane_count) |index| {129        var total: W = 0;130        inline for (0..8) |offset| {131            total += absoluteDifference(W, a[index * 8 + offset], b[index * 8 + offset]);132        }133        result[index] = total;134    }135    return result;136}137138pub fn sumsOfAdjQuadAbsDiff(139    comptime D: type,140    comptime a_offset: usize,141    comptime b_offset: usize,142    a: D.Vector,143    b: D.Vector,144) D.repartition(wideLane(D.Lane)).Vector {145    requireByteLane(D, "sumsOfAdjQuadAbsDiff");146    if (a_offset > 1 or b_offset > 3) @compileError("quad offsets exceed a 128-bit block");147    const W = wideLane(D.Lane);148    const R = D.repartition(W);149    var result: R.Vector = @splat(0);150    inline for (0..R.lane_count) |index| {151        const a_base = a_offset * 4 + (index / 8) * 16 + (index & 7);152        const b_base = b_offset * 4 + (index / 8) * 16;153        if (a_base + 3 < D.lane_count and b_base + 3 < D.lane_count) {154            var total: W = 0;155            inline for (0..4) |offset| {156                total += absoluteDifference(W, a[a_base + offset], b[b_base + offset]);157            }158            result[index] = total;159        }160    }161    return result;162}163164pub fn sumsOfShuffledQuadAbsDiff(165    comptime D: type,166    comptime index3: usize,167    comptime index2: usize,168    comptime index1: usize,169    comptime index0: usize,170    a: D.Vector,171    b: D.Vector,172) D.repartition(wideLane(D.Lane)).Vector {173    requireByteLane(D, "sumsOfShuffledQuadAbsDiff");174    if (index0 > 3 or index1 > 3 or index2 > 3 or index3 > 3) {175        @compileError("quad shuffle indices must be less than four");176    }177    const W = wideLane(D.Lane);178    const R = D.repartition(W);179    const selectors = [4]usize{ index0, index1, index2, index3 };180    var shuffled: [D.lane_count]D.Lane = @splat(0);181    inline for (0..(D.lane_count + 15) / 16) |block_index| {182        inline for (0..4) |group| {183            inline for (0..4) |byte| {184                const destination = block_index * 16 + group * 4 + byte;185                const source = block_index * 16 + selectors[group] * 4 + byte;186                if (destination < D.lane_count and source < D.lane_count) {187                    shuffled[destination] = a[source];188                }189            }190        }191    }192    var result: R.Vector = @splat(0);193    inline for (0..R.lane_count) |index| {194        const a_base = (index / 4) * 8 + (index & 3);195        const b_base = (index / 2) * 4;196        if (a_base + 3 < D.lane_count and b_base + 3 < D.lane_count) {197            var total: W = 0;198            inline for (0..4) |offset| {199                total += absoluteDifference(W, shuffled[a_base + offset], b[b_base + offset]);200            }201            result[index] = total;202        }203    }204    return result;205}206207fn wideLane(comptime T: type) type {208    return switch (T) {209        i8 => i16,210        u8 => u16,211        i16 => i32,212        u16 => u32,213        i32 => i64,214        u32 => u64,215        f16 => f32,216        f32 => f64,217        else => @compileError("lane type has no Highway wide representation"),218    };219}220221fn widen(comptime T: type, value: anytype) T {222    return if (@typeInfo(T) == .float) @floatCast(value) else @intCast(value);223}224225fn absoluteDifference(comptime T: type, a: anytype, b: @TypeOf(a)) T {226    const a_wide: i128 = a;227    const b_wide: i128 = b;228    const difference = a_wide - b_wide;229    return @intCast(if (difference < 0) -difference else difference);230}231232fn requireByteLane(comptime D: type, comptime operation: []const u8) void {233    if (comptime D.Lane != u8 and D.Lane != i8) {234        @compileError(operation ++ " requires byte integer lanes");235    }236}237238test "reductions match scalar Highway models" {239    const simd = @import("root.zig");240    const D = simd.FixedTag(i16, 8);241    const value: D.Vector = .{ 7, -4, 2, 9, -11, 3, 1, 6 };242    try std.testing.expectEqual(@as(i16, 13), sum(D, value));243    try std.testing.expectEqual(@as(i16, -11), min(D, value));244    try std.testing.expectEqual(@as(i16, 9), max(D, value));245}246247test "integer sum uses Highway wrapping semantics" {248    const simd = @import("root.zig");249    const D = simd.FixedTag(u8, 4);250    try std.testing.expectEqual(@as(u8, 4), sum(D, @as(D.Vector, @splat(129))));251}252253test "Highway broadcast and masked reductions match scalar selection" {254    const simd = @import("root.zig");255    const D = simd.FixedTag(i16, 8);256    const value: D.Vector = .{ 7, -4, 2, 9, -11, 3, 1, 6 };257    const mask: D.Mask = .{ false, true, false, true, true, false, true, false };258    try std.testing.expect(@reduce(.And, sumOfLanes(D, value) == @as(D.Vector, @splat(13))));259    try std.testing.expect(@reduce(.And, minOfLanes(D, value) == @as(D.Vector, @splat(-11))));260    try std.testing.expect(@reduce(.And, maxOfLanes(D, value) == @as(D.Vector, @splat(9))));261    try std.testing.expectEqual(@as(i16, -5), maskedSum(D, mask, value));262    try std.testing.expectEqual(@as(i16, -11), maskedMin(D, mask, value));263    try std.testing.expectEqual(@as(i16, 9), maskedMax(D, mask, value));264}265266test "Highway pairwise sums widen adjacent groups" {267    const simd = @import("root.zig");268    const D8 = simd.FixedTag(i8, 16);269    const value: D8.Vector = .{270        100, 100, -100, -100, 1, 2,  3,  4,271        5,   6,   7,    8,    9, 10, 11, 12,272    };273    const D16 = D8.repartition(i16);274    const D32 = D8.repartition(i32);275    const D64 = D8.repartition(i64);276    try std.testing.expect(@reduce(.And, sumsOf2(D8, value) == @as(D16.Vector, .{277        200, -200, 3, 7, 11, 15, 19, 23,278    })));279    try std.testing.expect(@reduce(.And, sumsOf4(D8, value) == @as(D32.Vector, .{ 0, 10, 26, 42 })));280    try std.testing.expect(@reduce(.And, sumsOf8(D8, value) == @as(D64.Vector, .{ 10, 68 })));281282    const F = simd.FixedTag(f32, 4);283    try std.testing.expect(@reduce(.And, sumsOf2(F, @as(F.Vector, .{ 1.5, 2.5, -4, 1 })) ==284        @as(F.repartition(f64).Vector, .{ 4, -3 })));285}286287fn verifyReductionLaneType(comptime T: type) !void {288    const simd = @import("root.zig");289    const D = simd.FixedTag(T, 4);290    const value: D.Vector = @splat(0);291    const mask: D.Mask = .{ true, false, true, false };292    try std.testing.expectEqual(@as(T, 0), sum(D, value));293    try std.testing.expectEqual(@as(T, 0), min(D, value));294    try std.testing.expectEqual(@as(T, 0), max(D, value));295    try std.testing.expectEqual(@as(T, 0), maskedSum(D, mask, value));296    try std.testing.expectEqual(@as(T, 0), maskedMin(D, mask, value));297    try std.testing.expectEqual(@as(T, 0), maskedMax(D, mask, value));298    try std.testing.expect(@reduce(.And, sumOfLanes(D, value) == value));299    try std.testing.expect(@reduce(.And, minOfLanes(D, value) == value));300    try std.testing.expect(@reduce(.And, maxOfLanes(D, value) == value));301}302303fn verifySumsOf2(comptime T: type) !void {304    const simd = @import("root.zig");305    const D = simd.FixedTag(T, 8);306    const value: D.Vector = @splat(0);307    try std.testing.expect(@reduce(.And, sumsOf2(D, value) == @as(D.repartition(wideLane(T)).Vector, @splat(0))));308}309310fn verifySumsOf4(comptime T: type) !void {311    const simd = @import("root.zig");312    const D = simd.FixedTag(T, 8);313    const W = wideLane(wideLane(T));314    const value: D.Vector = @splat(0);315    try std.testing.expect(@reduce(.And, sumsOf4(D, value) == @as(D.repartition(W).Vector, @splat(0))));316}317318test "Highway reductions instantiate every lane type and supported widening" {319    inline for (.{ u8, i8, u16, i16, u32, i32, u64, i64, f16, f32, f64 }) |T| {320        try verifyReductionLaneType(T);321    }322    inline for (.{ u8, i8, u16, i16, u32, i32, f16, f32 }) |T| try verifySumsOf2(T);323    inline for (.{ u8, i8, u16, i16 }) |T| try verifySumsOf4(T);324    const simd = @import("root.zig");325    inline for (.{ u8, i8 }) |T| {326        const D = simd.FixedTag(T, 8);327        try std.testing.expect(@reduce(.And, sumsOf8(D, @as(D.Vector, @splat(0))) ==328            @as(D.repartition(wideLane(wideLane(wideLane(T)))).Vector, @splat(0))));329    }330}331332test "Highway absolute-difference reductions match scalar block formulas" {333    const simd = @import("root.zig");334    const D = simd.FixedTag(i8, 16);335    const a: D.Vector = .{ 1, 2, 3, 4, 9, 8, 7, 6, -1, -2, -3, -4, 5, 6, 7, 8 };336    const b: D.Vector = .{ 4, 3, 2, 1, 1, 2, 3, 4, 1, 2, 3, 4, 8, 7, 6, 5 };337    try std.testing.expect(@reduce(.And, sumsOf8AbsDiff(D, a, b) ==338        @as(D.repartition(i64).Vector, .{ 28, 28 })));339    try std.testing.expect(@reduce(.And, sumsOfAdjQuadAbsDiff(D, 0, 1, a, b) ==340        @as(D.repartition(i16).Vector, .{ 0, 8, 14, 18, 20, 20, 20, 20 })));341    try std.testing.expect(@reduce(.And, sumsOfShuffledQuadAbsDiff(D, 3, 2, 1, 0, a, b) ==342        @as(D.repartition(i16).Vector, .{ 8, 12, 14, 18, 20, 16, 24, 16 })));343}

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