tiny.simd.bfloat

1const std = @import("std");2const tag = @import("tag.zig");34pub const BFloat16 = extern struct {5    pub const is_bfloat16 = true;6    pub const Storage = u16;78    bits: u16,910    pub fn fromBits(bits: u16) BFloat16 {11        return .{ .bits = bits };12    }1314    pub fn fromF32(value: f32) BFloat16 {15        return .{ .bits = bitsFromF32(value) };16    }1718    pub fn fromF64(value: f64) BFloat16 {19        const bits: u64 = @bitCast(value);20        const rounded_bits = (bits & 0xffff_ffc0_0000_0000) |21            ((bits +% 0x0000_003f_ffff_ffff) & 0x0000_0040_0000_0000);22        return fromF32(@floatCast(@as(f64, @bitCast(rounded_bits))));23    }2425    pub fn toF32(self: BFloat16) f32 {26        return f32FromBits(self.bits);27    }2829    pub fn neg(self: BFloat16) BFloat16 {30        return fromBits(self.bits ^ 0x8000);31    }3233    pub fn add(self: BFloat16, other: BFloat16) BFloat16 {34        return fromF32(self.toF32() + other.toF32());35    }3637    pub fn sub(self: BFloat16, other: BFloat16) BFloat16 {38        return fromF32(self.toF32() - other.toF32());39    }4041    pub fn mul(self: BFloat16, other: BFloat16) BFloat16 {42        return fromF32(self.toF32() * other.toF32());43    }4445    pub fn div(self: BFloat16, other: BFloat16) BFloat16 {46        return fromF32(self.toF32() / other.toF32());47    }48};4950pub fn Tag(comptime lanes: usize) type {51    return tag.Descriptor(BFloat16, lanes, 0);52}5354pub fn zero(comptime D: type) D.Vector {55    requireTag(D);56    return @splat(0);57}5859pub fn set(comptime D: type, value: BFloat16) D.Vector {60    requireTag(D);61    return @splat(value.bits);62}6364pub fn load(comptime D: type, input: []const BFloat16) D.Vector {65    requireTag(D);66    std.debug.assert(input.len >= D.lane_count);67    var result: D.Vector = undefined;68    inline for (0..D.lane_count) |index| result[index] = input[index].bits;69    return result;70}7172pub fn store(comptime D: type, value: D.Vector, output: []BFloat16) void {73    requireTag(D);74    std.debug.assert(output.len >= D.lane_count);75    inline for (0..D.lane_count) |index| output[index] = BFloat16.fromBits(value[index]);76}7778pub fn dup128VecFromValues(79    comptime D: type,80    values: [8]BFloat16,81) D.Vector {82    requireTag(D);83    var result: D.Vector = undefined;84    inline for (0..D.lane_count) |index| result[index] = values[index % 8].bits;85    return result;86}8788pub fn neg(comptime D: type, value: D.Vector) D.Vector {89    requireTag(D);90    return value ^ @as(D.Vector, @splat(0x8000));91}9293pub fn demoteF32(comptime D: type, value: anytype) D.Vector {94    requireTag(D);95    validateF32Vector(@TypeOf(value), D.lane_count);96    var result: D.Vector = undefined;97    inline for (0..D.lane_count) |index| result[index] = bitsFromF32(value[index]);98    return result;99}100101pub fn promoteF32(comptime D: type, value: anytype) D.Vector {102    if (comptime D.Lane != f32) @compileError("bfloat16 promotion requires f32 destinations");103    validateBitsVector(@TypeOf(value), D.lane_count);104    var result: D.Vector = undefined;105    inline for (0..D.lane_count) |index| result[index] = f32FromBits(value[index]);106    return result;107}108109pub fn promoteLowerF32(comptime D: type, value: anytype) D.Vector {110    return promoteSelectedF32(D, value, 0);111}112113pub fn promoteUpperF32(comptime D: type, value: anytype) D.Vector {114    return promoteSelectedF32(D, value, D.lane_count);115}116117pub fn promoteEvenF32(comptime D: type, value: anytype) D.Vector {118    return promoteParityF32(D, value, 0);119}120121pub fn promoteOddF32(comptime D: type, value: anytype) D.Vector {122    return promoteParityF32(D, value, 1);123}124125pub fn orderedDemote2F32(comptime D: type, a: anytype, b: @TypeOf(a)) D.Vector {126    requireTag(D);127    if (comptime vectorLanes(@TypeOf(a)) * 2 != D.lane_count) {128        @compileError("ordered bfloat16 demotion requires two half-size f32 vectors");129    }130    validateF32Vector(@TypeOf(a), D.lane_count / 2);131    var result: D.Vector = undefined;132    inline for (0..D.lane_count / 2) |index| {133        result[index] = bitsFromF32(a[index]);134        result[D.lane_count / 2 + index] = bitsFromF32(b[index]);135    }136    return result;137}138139pub fn mulEvenAdd(140    comptime D: type,141    a: anytype,142    b: @TypeOf(a),143    addend: D.Vector,144) D.Vector {145    return mulParityAdd(D, a, b, addend, 0);146}147148pub fn mulOddAdd(149    comptime D: type,150    a: anytype,151    b: @TypeOf(a),152    addend: D.Vector,153) D.Vector {154    return mulParityAdd(D, a, b, addend, 1);155}156157pub fn widenMulPairwiseAdd(comptime D: type, a: anytype, b: @TypeOf(a)) D.Vector {158    return mulOddAdd(D, a, b, mulEvenAdd(D, a, b, @splat(0)));159}160161pub fn maskedWidenMulPairwiseAdd(162    comptime D: type,163    mask: D.Mask,164    a: anytype,165    b: @TypeOf(a),166) D.Vector {167    return @select(f32, mask, widenMulPairwiseAdd(D, a, b), @as(D.Vector, @splat(0)));168}169170pub fn reorderWidenMulAccumulate(171    comptime D: type,172    a: anytype,173    b: @TypeOf(a),174    sum0: D.Vector,175    sum1: *D.Vector,176) D.Vector {177    sum1.* = mulOddAdd(D, a, b, sum1.*);178    return mulEvenAdd(D, a, b, sum0);179}180181pub fn rearrangeToOddPlusEven(comptime D: type, sum0: D.Vector, sum1: D.Vector) D.Vector {182    if (comptime D.Lane != f32) @compileError("bfloat16 accumulation requires f32 sums");183    return sum0 + sum1;184}185186pub fn bitsFromF32(value: f32) u16 {187    const bits: u32 = @bitCast(value);188    const magnitude = bits & 0x7fff_ffff;189    const increment: u32 = if (magnitude < 0x7f80_0000)190        0x7fff + ((bits >> 16) & 1)191    else192        0;193    const quiet_nan: u32 = if (magnitude > 0x7f80_0000) 1 << 6 else 0;194    return @truncate(quiet_nan | ((bits +% increment) >> 16));195}196197pub fn f32FromBits(bits: u16) f32 {198    return @bitCast(@as(u32, bits) << 16);199}200201fn promoteSelectedF32(comptime D: type, value: anytype, comptime offset: usize) D.Vector {202    if (comptime D.Lane != f32) @compileError("bfloat16 promotion requires f32 destinations");203    validateBitsVector(@TypeOf(value), D.lane_count * 2);204    var result: D.Vector = undefined;205    inline for (0..D.lane_count) |index| result[index] = f32FromBits(value[offset + index]);206    return result;207}208209fn promoteParityF32(comptime D: type, value: anytype, comptime parity: usize) D.Vector {210    if (comptime D.Lane != f32) @compileError("bfloat16 promotion requires f32 destinations");211    validateBitsVector(@TypeOf(value), D.lane_count * 2);212    var result: D.Vector = undefined;213    inline for (0..D.lane_count) |index| result[index] = f32FromBits(value[index * 2 + parity]);214    return result;215}216217fn mulParityAdd(218    comptime D: type,219    a: anytype,220    b: @TypeOf(a),221    addend: D.Vector,222    comptime parity: usize,223) D.Vector {224    if (comptime D.Lane != f32) @compileError("bfloat16 multiply-add requires f32 destinations");225    validateBitsVector(@TypeOf(a), D.lane_count * 2);226    var result = addend;227    inline for (0..D.lane_count) |index| {228        result[index] = @mulAdd(229            f32,230            f32FromBits(a[index * 2 + parity]),231            f32FromBits(b[index * 2 + parity]),232            addend[index],233        );234    }235    return result;236}237238fn requireTag(comptime D: type) void {239    if (comptime !@hasDecl(D, "is_bfloat16") or !D.is_bfloat16) {240        @compileError("operation requires a bfloat16 descriptor");241    }242}243244fn validateF32Vector(comptime V: type, comptime lanes: usize) void {245    if (comptime vectorLane(V) != f32 or vectorLanes(V) != lanes) {246        @compileError("operation requires the expected number of f32 lanes");247    }248}249250fn validateBitsVector(comptime V: type, comptime lanes: usize) void {251    if (comptime vectorLane(V) != u16 or vectorLanes(V) != lanes) {252        @compileError("operation requires the expected number of bfloat16 bit lanes");253    }254}255256fn vectorLane(comptime V: type) type {257    return switch (@typeInfo(V)) {258        .vector => |info| info.child,259        else => @compileError("operation requires a vector"),260    };261}262263fn vectorLanes(comptime V: type) usize {264    return switch (@typeInfo(V)) {265        .vector => |info| info.len,266        else => @compileError("operation requires a vector"),267    };268}269270test "Highway bfloat16 scalar conversion rounds to nearest even and preserves classes" {271    try std.testing.expectEqual(@as(u16, 0x3f80), BFloat16.fromF32(1).bits);272    try std.testing.expectEqual(@as(u16, 0xbf80), BFloat16.fromF32(-1).bits);273    try std.testing.expectEqual(@as(u16, 0x3f80), BFloat16.fromF32(1.00390625).bits);274    try std.testing.expectEqual(@as(u16, 0x3f82), BFloat16.fromF32(1.01171875).bits);275    try std.testing.expectEqual(@as(u16, 0x4000), BFloat16.fromF32(1.99609375).bits);276    try std.testing.expectEqual(@as(u16, 0x8000), BFloat16.fromF32(-0.0).bits);277    try std.testing.expectEqual(@as(u16, 0x7f80), BFloat16.fromF32(std.math.inf(f32)).bits);278    try std.testing.expect(std.math.isNan(BFloat16.fromF32(std.math.nan(f32)).toF32()));279}280281test "Highway bfloat16 descriptors preserve scalar type relations" {282    const D = tag.FixedTag(BFloat16, 8);283    try std.testing.expectEqual(BFloat16, tag.TFromD(D));284    try std.testing.expectEqual(Tag(8), D);285    try std.testing.expectEqual(@as(usize, 4), tag.Half(D).lane_count);286    try std.testing.expectEqual(@as(i8, -1), tag.pow2(tag.Half(D)));287    try std.testing.expectEqual(Tag(8), tag.BlockDFromD(Tag(16)));288    try std.testing.expectEqual(tag.FixedTag(i16, 8), tag.RebindToSigned(D));289    try std.testing.expectEqual(tag.FixedTag(u16, 8), tag.RebindToUnsigned(D));290    try std.testing.expectEqual(tag.FixedTag(f32, 4), tag.RepartitionToWide(D));291    try std.testing.expectEqual(u16, tag.MakeUnsigned(BFloat16));292    try std.testing.expectEqual(i16, tag.MakeSigned(BFloat16));293    try std.testing.expectEqual(f32, tag.MakeWide(BFloat16));294    try std.testing.expect(tag.isSpecialFloat(BFloat16));295    try std.testing.expect(!tag.isFloat(BFloat16));296    try std.testing.expect(tag.isSigned(BFloat16));297    try std.testing.expect(!tag.isUnsigned(BFloat16));298}299300test "Highway bfloat16 vector promotion demotion and memory preserve bit patterns" {301    const simd = @import("root.zig");302    const D = Tag(4);303    const F = tag.FixedTag(f32, 4);304    const input: F.Vector = .{ 1, -2, 3.984375, -0.0 };305    const bits = simd.demoteTo(D, input);306    try std.testing.expect(@reduce(.And, simd.promoteTo(F, bits) == @as(F.Vector, .{ 1, -2, 3.984375, -0.0 })));307    var scalar: [4]BFloat16 = undefined;308    store(D, bits, &scalar);309    try std.testing.expect(@reduce(.And, load(D, &scalar) == bits));310    try std.testing.expect(@reduce(.And, neg(D, bits) == (bits ^ @as(D.Vector, @splat(0x8000)))));311}312313test "Highway bfloat16 even odd and pairwise multiply-add widen to f32" {314    const simd = @import("root.zig");315    const D = tag.FixedTag(f32, 4);316    const B = Tag(8);317    const a = demoteF32(B, @as(tag.FixedTag(f32, 8).Vector, .{ 1, 2, 3, 4, 5, 6, 7, 8 }));318    const b = demoteF32(B, @as(tag.FixedTag(f32, 8).Vector, .{ 8, 7, 6, 5, 4, 3, 2, 1 }));319    try std.testing.expect(@reduce(.And, mulEvenAdd(D, a, b, @splat(1)) ==320        @as(D.Vector, .{ 9, 19, 21, 15 })));321    try std.testing.expect(@reduce(.And, mulOddAdd(D, a, b, @splat(1)) ==322        @as(D.Vector, .{ 15, 21, 19, 9 })));323    try std.testing.expect(@reduce(.And, simd.widenMulPairwiseAdd(D, a, b) ==324        @as(D.Vector, .{ 22, 38, 38, 22 })));325    var odd: D.Vector = @splat(0);326    const even = simd.reorderWidenMulAccumulate(D, a, b, @splat(0), &odd);327    try std.testing.expect(@reduce(.And, simd.rearrangeToOddPlusEven(D, even, odd) ==328        @as(D.Vector, .{ 22, 38, 38, 22 })));329}

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