tiny.accy.tensor.types.dim

1const std = @import("std");2const choir_abi = @import("choir_abi");3const accy = @import("../../root.zig");45pub const DType = choir_abi.DType;67pub const generated_prefix: u8 = '#';89pub const Error = error{10    InvalidDimension,11    ShapeElementOverflow,12    ShapeByteOverflow,13    ShapeMismatch,14    DTypeMismatch,15    RankMismatch,16    AxisOutOfRange,17    AxisNotFound,18    AxisNameEmpty,19    AxisNameReserved,20    AxisExtentMismatch,21    DuplicateAxis,22    InvalidPermutation,23    ReshapeElementMismatch,24    BroadcastRankMismatch,25    BroadcastDimensionMismatch,26    ContractDTypeMismatch,27    SplitExtentMismatch,28    MergeExtentOverflow,29    ReduceInitNotScalar,30    PayloadLengthMismatch,31};3233pub const Dim = struct {34    name: []const u8,35    extent: i64,36};3738pub const Spec = struct {39    dtype: DType,40    dims: []const Dim,4142    pub fn init(dtype: DType, dims: []const Dim) Spec {43        return .{ .dtype = dtype, .dims = dims };44    }45};4647pub const Type = struct {48    dtype: DType,49    dims: []const Dim,5051    pub fn init(allocator: std.mem.Allocator, dtype: DType, dims: []const Dim) !Type {52        return .{53            .dtype = dtype,54            .dims = try copyDims(allocator, dims),55        };56    }5758    pub fn fromSpec(allocator: std.mem.Allocator, value: Spec) !Type {59        return init(allocator, value.dtype, value.dims);60    }6162    pub fn scalar(dtype: DType) Type {63        return .{ .dtype = dtype, .dims = &.{} };64    }6566    pub fn eql(self: Type, other: Type) bool {67        return self.dtype == other.dtype and sameDims(self.dims, other.dims);68    }6970    pub fn rank(self: Type) usize {71        return self.dims.len;72    }7374    pub fn extent(self: Type, index: usize) i64 {75        return self.dims[index].extent;76    }7778    pub fn find(self: Type, name: []const u8) ?usize {79        return findDim(self.dims, name);80    }8182    pub fn elementCount(self: Type) !usize {83        return dimsElementCount(self.dims);84    }8586    pub fn byteCount(self: Type) !usize {87        return std.math.mul(usize, try self.elementCount(), self.dtype.sizeOf()) catch error.ShapeByteOverflow;88    }89};9091pub fn spec(comptime dtype: DType, comptime dims_struct: anytype) Spec {92    const dims_array = comptime blk: {93        var array: [dimCount(@TypeOf(dims_struct))]Dim = undefined;94        fillDims(dims_struct, &array);95        break :blk array;96    };97    return .{ .dtype = dtype, .dims = &dims_array };98}99100pub fn specDims(dtype: DType, dims: []const Dim) Spec {101    return Spec.init(dtype, dims);102}103104pub fn dimCount(comptime Dims: type) usize {105    const info = @typeInfo(Dims);106    if (info != .@"struct") {107        @compileError("tensor dims must be a struct literal mapping axis names to extents");108    }109    if (info.@"struct".field_names.len == 0) return 0;110    if (info.@"struct".is_tuple) {111        @compileError("tensor dims must be a struct literal mapping axis names to extents");112    }113    return info.@"struct".field_names.len;114}115116pub fn fillDims(dims_struct: anytype, buffer: *[dimCount(@TypeOf(dims_struct))]Dim) void {117    const field_names = @typeInfo(@TypeOf(dims_struct)).@"struct".field_names;118    inline for (field_names, 0..) |field_name, index| {119        buffer[index] = .{120            .name = field_name,121            .extent = @intCast(@field(dims_struct, field_name)),122        };123    }124}125126pub fn axisCount(comptime axes: anytype) usize {127    const Axes = @TypeOf(axes);128    return switch (@typeInfo(Axes)) {129        .enum_literal => 1,130        .@"struct" => |info| blk: {131            if (!info.is_tuple) {132                @compileError("tensor axes must be an enum literal or a tuple of enum literals");133            }134            break :blk info.field_names.len;135        },136        else => @compileError("tensor axes must be an enum literal or a tuple of enum literals"),137    };138}139140pub fn axisNames(comptime axes: anytype) *const [axisCount(axes)][]const u8 {141    comptime var names: [axisCount(axes)][]const u8 = undefined;142    switch (@typeInfo(@TypeOf(axes))) {143        .enum_literal => names[0] = @tagName(axes),144        .@"struct" => |info| {145            inline for (info.field_names, 0..) |field_name, index| {146                const value = @field(axes, field_name);147                if (@typeInfo(@TypeOf(value)) != .enum_literal) {148                    @compileError("tensor axes tuple entries must be enum literals");149                }150                names[index] = @tagName(value);151            }152        },153        else => unreachable,154    }155    const final = names;156    return &final;157}158159pub fn isNameSlice(comptime Axes: type) bool {160    return Axes == []const []const u8 or Axes == [][]const u8;161}162163pub fn validateName(name: []const u8) !void {164    if (name.len == 0) return error.AxisNameEmpty;165}166167pub fn validateAuthoredName(name: []const u8) !void {168    try validateName(name);169    if (name[0] == generated_prefix) return error.AxisNameReserved;170}171172pub fn validateDims(dims: []const Dim) !void {173    for (dims, 0..) |dim, index| {174        try validateName(dim.name);175        if (dim.extent < 0) return error.InvalidDimension;176        for (dims[0..index]) |seen| {177            if (std.mem.eql(u8, seen.name, dim.name)) return error.DuplicateAxis;178        }179    }180}181182pub fn validateAuthoredDims(dims: []const Dim) !void {183    try validateDims(dims);184    for (dims) |dim| {185        try validateAuthoredName(dim.name);186    }187}188189pub fn copyDims(allocator: std.mem.Allocator, dims: []const Dim) ![]const Dim {190    try validateDims(dims);191    const owned = try allocator.alloc(Dim, dims.len);192    for (dims, owned) |dim, *slot| {193        slot.* = .{194            .name = try allocator.dupe(u8, dim.name),195            .extent = dim.extent,196        };197    }198    return owned;199}200201pub fn findDim(dims: []const Dim, name: []const u8) ?usize {202    for (dims, 0..) |dim, index| {203        if (std.mem.eql(u8, dim.name, name)) return index;204    }205    return null;206}207208pub fn sameDims(lhs: []const Dim, rhs: []const Dim) bool {209    if (lhs.len != rhs.len) return false;210    for (lhs, rhs) |left, right| {211        if (left.extent != right.extent) return false;212        if (!std.mem.eql(u8, left.name, right.name)) return false;213    }214    return true;215}216217pub fn dimsElementCount(dims: []const Dim) !usize {218    var count: usize = 1;219    for (dims) |dim| {220        if (dim.extent < 0) return error.InvalidDimension;221        count = std.math.mul(usize, count, @intCast(dim.extent)) catch return error.ShapeElementOverflow;222    }223    return count;224}225226pub fn extents(allocator: std.mem.Allocator, dims: []const Dim) ![]const i64 {227    const owned = try allocator.alloc(i64, dims.len);228    for (dims, owned) |dim, *slot| {229        slot.* = dim.extent;230    }231    return owned;232}233234pub fn sameType(lhs: Type, rhs: Type) !void {235    if (lhs.dtype != rhs.dtype) return error.DTypeMismatch;236    if (!sameDims(lhs.dims, rhs.dims)) return error.ShapeMismatch;237}238239test "dims literal parsing preserves declaration order" {240    const parsed = spec(.f32, .{ .point = 16, .sample = 8 });241    try std.testing.expectEqual(@as(usize, 2), parsed.dims.len);242    try std.testing.expectEqualStrings("point", parsed.dims[0].name);243    try std.testing.expectEqual(@as(i64, 16), parsed.dims[0].extent);244    try std.testing.expectEqualStrings("sample", parsed.dims[1].name);245    try std.testing.expectEqual(@as(i64, 8), parsed.dims[1].extent);246}247248test "axis names parse from literals and tuples" {249    const single = axisNames(.sample);250    try std.testing.expectEqualStrings("sample", single[0]);251252    const pair = axisNames(.{ .point, .sample });253    try std.testing.expectEqualStrings("point", pair[0]);254    try std.testing.expectEqualStrings("sample", pair[1]);255}256257test "type owns validated dims and computes byte count" {258    const ty = try Type.init(std.testing.allocator, .f32, &.{259        .{ .name = "a", .extent = 2 },260        .{ .name = "b", .extent = 3 },261        .{ .name = "c", .extent = 4 },262    });263    defer freeTestType(std.testing.allocator, ty);264265    try std.testing.expectEqual(@as(usize, 3), ty.rank());266    try std.testing.expectEqual(@as(usize, 24), try ty.elementCount());267    try std.testing.expectEqual(@as(usize, 96), try ty.byteCount());268    try std.testing.expectEqual(@as(?usize, 1), ty.find("b"));269    try std.testing.expectEqual(@as(?usize, null), ty.find("missing"));270}271272test "dims validation rejects bad names and extents" {273    try std.testing.expectError(error.InvalidDimension, Type.init(274        std.testing.allocator,275        .f32,276        &.{.{ .name = "a", .extent = -1 }},277    ));278    try std.testing.expectError(error.DuplicateAxis, Type.init(279        std.testing.allocator,280        .f32,281        &.{ .{ .name = "a", .extent = 2 }, .{ .name = "a", .extent = 3 } },282    ));283    try std.testing.expectError(error.AxisNameEmpty, Type.init(284        std.testing.allocator,285        .f32,286        &.{.{ .name = "", .extent = 2 }},287    ));288    try std.testing.expectError(error.AxisNameReserved, validateAuthoredDims(&.{289        .{ .name = "#batch", .extent = 2 },290    }));291    try validateDims(&.{.{ .name = "#batch", .extent = 2 }});292}293294pub fn freeTestType(allocator: std.mem.Allocator, ty: Type) void {295    for (ty.dims) |dim| allocator.free(dim.name);296    allocator.free(ty.dims);297}298299pub fn expectExtents(expected: []const i64, ty: Type) !void {300    try std.testing.expectEqual(expected.len, ty.dims.len);301    for (expected, ty.dims) |extent, dim| {302        try std.testing.expectEqual(extent, dim.extent);303    }304}

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