tiny.simd.complex

1const std = @import("std");23pub fn complexConj(comptime D: type, value: D.Vector) D.Vector {4    validate(D);5    var result = value;6    inline for (0..D.lane_count / 2) |index| {7        result[index * 2 + 1] = -value[index * 2 + 1];8    }9    return result;10}1112pub fn mulComplex(comptime D: type, a: D.Vector, b: D.Vector) D.Vector {13    return multiply(D, a, b, false, null);14}1516pub fn mulComplexConj(comptime D: type, a: D.Vector, b: D.Vector) D.Vector {17    return multiply(D, a, b, true, null);18}1920pub fn mulComplexAdd(comptime D: type, a: D.Vector, b: D.Vector, c: D.Vector) D.Vector {21    return multiply(D, a, b, false, c);22}2324pub fn mulComplexConjAdd(comptime D: type, a: D.Vector, b: D.Vector, c: D.Vector) D.Vector {25    return multiply(D, a, b, true, c);26}2728pub fn maskedMulComplexConjAdd(29    comptime D: type,30    mask: D.Mask,31    a: D.Vector,32    b: D.Vector,33    c: D.Vector,34) D.Vector {35    return @select(36        D.Lane,37        mask,38        mulComplexConjAdd(D, a, b, c),39        @as(D.Vector, @splat(0)),40    );41}4243pub fn maskedMulComplexConj(44    comptime D: type,45    mask: D.Mask,46    a: D.Vector,47    b: D.Vector,48) D.Vector {49    return @select(50        D.Lane,51        mask,52        mulComplexConj(D, a, b),53        @as(D.Vector, @splat(0)),54    );55}5657pub fn maskedMulComplexOr(58    comptime D: type,59    no: D.Vector,60    mask: D.Mask,61    a: D.Vector,62    b: D.Vector,63) D.Vector {64    return @select(D.Lane, mask, mulComplex(D, a, b), no);65}6667fn multiply(68    comptime D: type,69    a: D.Vector,70    b: D.Vector,71    comptime conjugate_b: bool,72    addend: ?D.Vector,73) D.Vector {74    validate(D);75    var result: D.Vector = undefined;76    inline for (0..D.lane_count / 2) |index| {77        const real = index * 2;78        const imaginary = real + 1;79        const c_real = if (addend) |c| c[real] else 0;80        const c_imaginary = if (addend) |c| c[imaginary] else 0;81        if (conjugate_b) {82            result[real] = @mulAdd(D.Lane, a[imaginary], b[imaginary], @mulAdd(D.Lane, a[real], b[real], c_real));83            result[imaginary] = @mulAdd(D.Lane, a[imaginary], b[real], @mulAdd(D.Lane, -a[real], b[imaginary], c_imaginary));84        } else {85            result[real] = @mulAdd(D.Lane, -a[imaginary], b[imaginary], @mulAdd(D.Lane, a[real], b[real], c_real));86            result[imaginary] = @mulAdd(D.Lane, a[imaginary], b[real], @mulAdd(D.Lane, a[real], b[imaginary], c_imaginary));87        }88    }89    return result;90}9192fn validate(comptime D: type) void {93    if (comptime @typeInfo(D.Lane) != .float or D.lane_count & 1 != 0) {94        @compileError("complex arithmetic requires an even number of floating-point lanes");95    }96}9798test "Highway complex multiply conjugate add and masks preserve lane pairs" {99    const simd = @import("root.zig");100    const D = simd.FixedTag(f32, 8);101    const a: D.Vector = .{ 1, 2, 3, 4, -2, 5, 7, -3 };102    const b: D.Vector = .{ 5, 6, 7, 8, 4, -1, -2, 6 };103    try std.testing.expect(@reduce(.And, complexConj(D, a) ==104        @as(D.Vector, .{ 1, -2, 3, -4, -2, -5, 7, 3 })));105    try std.testing.expect(@reduce(.And, mulComplex(D, a, b) ==106        @as(D.Vector, .{ -7, 16, -11, 52, -3, 22, 4, 48 })));107    try std.testing.expect(@reduce(.And, mulComplexConj(D, a, b) ==108        @as(D.Vector, .{ 17, 4, 53, 4, -13, 18, -32, -36 })));109    try std.testing.expect(@reduce(.And, mulComplexAdd(D, a, b, @as(D.Vector, @splat(1))) ==110        @as(D.Vector, .{ -6, 17, -10, 53, -2, 23, 5, 49 })));111    const mask: D.Mask = .{ true, true, false, false, true, false, false, true };112    try std.testing.expect(@reduce(.And, maskedMulComplexOr(D, @splat(9), mask, a, b) ==113        @as(D.Vector, .{ -7, 16, 9, 9, -3, 9, 9, 48 })));114    _ = maskedMulComplexConj(D, mask, a, b);115    _ = maskedMulComplexConjAdd(D, mask, a, b, @splat(1));116}117118test "Highway complex arithmetic instantiates every floating lane type" {119    const simd = @import("root.zig");120    inline for (.{ f16, f32, f64 }) |T| {121        const D = simd.FixedTag(T, 4);122        const value: D.Vector = @splat(1);123        _ = mulComplex(D, value, value);124        _ = mulComplexConj(D, value, value);125        _ = mulComplexAdd(D, value, value, value);126        _ = mulComplexConjAdd(D, value, value, value);127    }128}

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