tiny.accy.tensor.random

1const std = @import("std");2const choir_abi = @import("choir_abi");3const accy = @import("../../root.zig");4const tensor = @import("../root.zig");56const DType = tensor.DType;7const Dim = tensor.Dim;8const Type = tensor.Type;9const Value = tensor.Value;10const Builder = tensor.Builder;1112pub const Algorithm = enum {13    philox,14};1516pub const KeySpec = struct {17    algorithm: Algorithm = .philox,18    rounds: u32 = accy.kernel.library.random.philox_default_rounds,19    threads: u32 = 256,20};2122pub const Key = struct {23    value: Value,24    spec: KeySpec,2526    pub fn split(self: Key, dims_struct: anytype) !Key {27        var buffer: [tensor.types.dimCount(@TypeOf(dims_struct))]Dim = undefined;28        tensor.types.fillDims(dims_struct, &buffer);29        if (buffer.len != 1) return error.RankMismatch;30        return splitKey(self, buffer[0]);31    }3233    pub fn splitDim(self: Key, dim: Dim) !Key {34        return splitKey(self, dim);35    }3637    pub fn uniform(self: Key, dtype: DType) !Value {38        return uniformKey(self, dtype);39    }4041    pub fn counterUniform(self: Key, counter: Value, dims_struct: anytype, dtype: DType) !Value {42        var buffer: [tensor.types.dimCount(@TypeOf(dims_struct))]Dim = undefined;43        tensor.types.fillDims(dims_struct, &buffer);44        if (buffer.len != 1) return error.RankMismatch;45        return uniformKeyCounter(self, counter, buffer[0], dtype);46    }4748    pub fn counterUniformDim(self: Key, counter: Value, dim: Dim, dtype: DType) !Value {49        return uniformKeyCounter(self, counter, dim, dtype);50    }51};5253pub fn input(builder: *Builder, dims: []const Dim, spec_value: KeySpec) !Key {54    return .{55        .value = try builder.inputDims(.key, dims),56        .spec = spec_value,57    };58}5960pub fn seed(builder: *Builder, seed_value: u64, spec_value: KeySpec) !Key {61    var payload = [_]choir_abi.Key{choir_abi.Key.init(seed_value)};62    return .{63        .value = try builder.constantBytes(Type.scalar(.key), std.mem.sliceAsBytes(payload[0..])),64        .spec = spec_value,65    };66}6768pub fn bind(value: Value, spec_value: KeySpec) !Key {69    if (value.ty.dtype != .key) return error.DTypeMismatch;70    return .{ .value = value, .spec = spec_value };71}7273fn splitKey(source: Key, dim: Dim) !Key {74    if (dim.extent <= 0) return error.InvalidDimension;75    try tensor.types.validateAuthoredDims(&.{dim});76    if (source.value.ty.rank() != 0) return error.RankMismatch;77    const builder = source.value.builder;78    const count: u64 = @intCast(dim.extent);79    const ty = try Type.init(builder.arena.allocator(), .key, &.{dim});80    const target = switch (source.spec.algorithm) {81        .philox => try accy.kernel.library.random.philoxKeySplitFamilyTarget(82            builder.arena.allocator(),83            .{84                .count = count,85                .rounds = source.spec.rounds,86                .threads = source.spec.threads,87            },88        ),89    };90    return .{91        .value = try builder.customCall(target, philoxSplitVersion(source.spec), &.{source.value}, ty),92        .spec = source.spec,93    };94}9596pub fn uniform(keys: Key, dtype: DType) !Value {97    return uniformKey(keys, dtype);98}99100fn uniformKey(keys: Key, dtype: DType) !Value {101    if (!accy.kernel.library.random.randomDTypeSupported(dtype)) return error.UnsupportedDType;102    if (keys.value.ty.rank() == 0) return error.RankMismatch;103    const count = try keys.value.ty.elementCount();104    const builder = keys.value.builder;105    const ty = try Type.init(builder.arena.allocator(), dtype, keys.value.ty.dims);106    const target = switch (keys.spec.algorithm) {107        .philox => try accy.kernel.library.random.philoxKeyUniformFamilyTarget(108            builder.arena.allocator(),109            .{110                .count = count,111                .rounds = keys.spec.rounds,112                .dtype = dtype,113                .threads = keys.spec.threads,114            },115        ),116    };117    return builder.customCall(target, philoxUniformVersion(keys.spec), &.{keys.value}, ty);118}119120fn uniformKeyCounter(key: Key, counter: Value, dim: Dim, dtype: DType) !Value {121    if (!accy.kernel.library.random.randomDTypeSupported(dtype)) return error.UnsupportedDType;122    if (dim.extent <= 0) return error.InvalidDimension;123    try tensor.types.validateAuthoredDims(&.{dim});124    if (key.value.ty.rank() != 0 or counter.ty.rank() != 0) return error.RankMismatch;125    if (counter.ty.dtype != .i32) return error.DTypeMismatch;126    const builder = key.value.builder;127    const count: u64 = @intCast(dim.extent);128    const ty = try Type.init(builder.arena.allocator(), dtype, &.{dim});129    const target = switch (key.spec.algorithm) {130        .philox => try accy.kernel.library.random.philoxKeyCounterUniformFamilyTarget(131            builder.arena.allocator(),132            .{133                .count = count,134                .rounds = key.spec.rounds,135                .dtype = dtype,136                .threads = key.spec.threads,137            },138        ),139    };140    return builder.customCall(target, philoxCounterUniformVersion(key.spec), &.{ key.value, counter }, ty);141}142143fn philoxSplitVersion(spec_value: KeySpec) u32 {144    return switch (spec_value.algorithm) {145        .philox => accy.kernel.library.random.philox_key_split_family_version,146    };147}148149fn philoxUniformVersion(spec_value: KeySpec) u32 {150    return switch (spec_value.algorithm) {151        .philox => accy.kernel.library.random.philox_key_uniform_family_version,152    };153}154155fn philoxCounterUniformVersion(spec_value: KeySpec) u32 {156    return switch (spec_value.algorithm) {157        .philox => accy.kernel.library.random.philox_key_counter_uniform_family_version,158    };159}

1pub const random = @import("random/root.zig");