lib/accy/src/tensor/random/root.zig
daab053ee43316e1809a84551d573ddd1e5bf3d2
1 const std = @import("std");
2 const choir_abi = @import("choir_abi");
3 const accy = @import("../../root.zig");
4 const tensor = @import("../root.zig");
5
6 const DType = tensor.DType;
7 const Dim = tensor.Dim;
8 const Type = tensor.Type;
9 const Value = tensor.Value;
10 const Builder = tensor.Builder;
11
12 pub const Algorithm = enum {
13 philox,
14 };
15
16 pub const KeySpec = struct {
17 algorithm: Algorithm = .philox,
18 rounds: u32 = accy.kernel.library.random.philox_default_rounds,
19 threads: u32 = 256,
20 };
21
22 pub const Key = struct {
23 value: Value,
24 spec: KeySpec,
25
26 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 }
32
33 pub fn splitDim(self: Key, dim: Dim) !Key {
34 return splitKey(self, dim);
35 }
36
37 pub fn uniform(self: Key, dtype: DType) !Value {
38 return uniformKey(self, dtype);
39 }
40
41 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 }
47
48 pub fn counterUniformDim(self: Key, counter: Value, dim: Dim, dtype: DType) !Value {
49 return uniformKeyCounter(self, counter, dim, dtype);
50 }
51 };
52
53 pub 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 }
59
60 pub 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 }
67
68 pub 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 }
72
73 fn 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 }
95
96 pub fn uniform(keys: Key, dtype: DType) !Value {
97 return uniformKey(keys, dtype);
98 }
99
100 fn 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 }
119
120 fn 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 }
142
143 fn philoxSplitVersion(spec_value: KeySpec) u32 {
144 return switch (spec_value.algorithm) {
145 .philox => accy.kernel.library.random.philox_key_split_family_version,
146 };
147 }
148
149 fn philoxUniformVersion(spec_value: KeySpec) u32 {
150 return switch (spec_value.algorithm) {
151 .philox => accy.kernel.library.random.philox_key_uniform_family_version,
152 };
153 }
154
155 fn 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 }