tiny.accy.tensor.wire.format

1//! Layout, version 1, of a tensor program written as bytes, shared by the encoder and the decoder.2//! A program written by one process has to be read back by another with every type intact, and a3//! reader has to reject bytes it cannot trust, because a reader that took result types from the4//! bytes would trust whatever type the sender wrote, correct or not.5//!6//! A program's bytes (its *wire bytes*) hold a magic number, a version, the program name and its7//! root graph. A graph is written as an operation count, the operations, and the ids of its8//! outputs. Each operation starts with a number that says which kind of operation it is (its *kind9//! code*), then only the fields that `Builder.operation` needs for that kind, so the decoder works10//! out every other result type itself. An operation can repeat a body graph a set number of times,11//! carrying values from one step to the next (a *scan*), and such an operation also writes its12//! length and the ids of its initial values, with its body graph after it. Integers are13//! little-endian, and every list, whether of bytes, ids, dimensions or integers, starts with its14//! length as a 32-bit count.15const std = @import("std");16const tensor = @import("../root.zig");1718const program = tensor.program;1920/// A reader checks these four bytes first to tell wire bytes, the versioned byte encoding of a21/// tensor program, from other data: the value is the four bytes "ACTP" read as a little-endian22/// 32-bit number.23pub const magic: u32 = std.mem.readInt(u32, "ACTP", .little);24/// A writer and a reader agree on the layout through this number: the current layout is version 1.25/// The version is raised before any change to how this format writes anything.26pub const version: u32 = 1;27/// A caller relies on this bound when building programs for transfer with nested scans that repeat28/// a body graph: a program may nest scans inside scan bodies at most eight levels deep. Encoding a29/// deeper program returns `error.ScanTooDeep`.30pub const max_scan_depth: usize = 8;3132/// A caller matches on these to tell why a program failed to encode or decode, so the set lists the33/// wire failures other than bytes that end early or run past the program: a wrong magic number,34/// another version, an unknown code, a reference to a missing value, too many operands and scans35/// nested too deep. Bytes that end early or run past the program return `error.InvalidArtifact`.36pub const Error = error{37    BadMagic,38    UnsupportedVersion,39    UnknownTag,40    InvalidReference,41    TooManyOperands,42    ScanTooDeep,43};4445/// The function assigns each operation kind, element type and option a fixed number on the wire: it46/// gives each value of `E` the number of its position in `order`. A compile-time check fails the47/// build when `order` leaves out a value of `E` or lists one twice. Decoding a number past the end48/// of `order` returns `error.UnknownTag`.49pub fn Codes(comptime E: type, comptime order: []const E) type {50    comptime {51        if (order.len != @typeInfo(E).@"enum".field_names.len) {52            @compileError("wire codes must cover every value of " ++ @typeName(E));53        }54        for (order, 0..) |value, index| {55            for (order[0..index]) |previous| {56                if (previous == value) @compileError("duplicate wire code in " ++ @typeName(E));57            }58        }59    }60    return struct {61        pub fn code(value: E) u32 {62            for (order, 0..) |candidate, index| {63                if (candidate == value) return @intCast(index);64            }65            unreachable;66        }6768        pub fn decode(value_code: u32) Error!E {69            if (value_code >= order.len) return error.UnknownTag;70            const value = order[value_code];71            std.debug.assert(code(value) == value_code);72            return value;73        }74    };75}7677pub const Tag = std.meta.Tag(program.Kind);7879pub const kind_codes = Codes(Tag, &.{80    .parameter,81    .constant,82    .unary,83    .binary,84    .iota,85    .broadcast,86    .broadcast_in_dim,87    .reshape,88    .transpose,89    .reduce,90    .gather,91    .scatter_add,92    .sparse_cross_entropy,93    .dot_general,94    .compare,95    .select,96    .custom_call,97    .scan,98    .projection,99});100101pub const dtype_codes = Codes(tensor.DType, &.{102    .i1, .i8, .i16, .i32, .i64, .u8, .u16, .u32, .u64, .f16, .bf16, .f32, .f64, .key,103});104105pub const unary_codes = Codes(program.Unary, &.{106    .neg, .abs, .exp, .log, .sqrt, .tanh, .sin, .cos, .tan,107});108109pub const binary_codes = Codes(program.Binary, &.{ .add, .sub, .mul, .div, .max, .min, .pow });110111pub const reducer_codes = Codes(program.Reducer, &.{ .sum, .max, .min });112113pub const compare_codes = Codes(program.CompareDirection, &.{ .lt, .le, .gt, .ge, .eq, .ne });114115test "tensor wire codes pin enum positions and reject unknown codes" {116    try std.testing.expectEqual(@as(u32, 0), kind_codes.code(.parameter));117    try std.testing.expectEqual(@as(u32, 18), kind_codes.code(.projection));118    try std.testing.expectEqual(Tag.scan, try kind_codes.decode(17));119    try std.testing.expectEqual(program.Unary.tan, try unary_codes.decode(8));120    try std.testing.expectEqual(tensor.DType.key, try dtype_codes.decode(13));121    try std.testing.expectError(error.UnknownTag, unary_codes.decode(9));122    try std.testing.expectError(error.UnknownTag, dtype_codes.decode(14));123    try std.testing.expectError(error.UnknownTag, kind_codes.decode(std.math.maxInt(u32)));124}

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