lib/simd/src/random.zig
daab053ee43316e1809a84551d573ddd1e5bf3d2
1 const std = @import("std");
2 const crypto = @import("crypto.zig");
3 const rotate = @import("rotate.zig");
4 const shift = @import("shift.zig");
5 const tag = @import("tag.zig");
6
7 const jump_constants = [4]u64{
8 0x180e_c6d3_3cfd_0aba,
9 0xd5a6_1266_f0c9_392c,
10 0xa958_2618_e03f_c9aa,
11 0x39ab_dc45_29b1_661c,
12 };
13
14 const long_jump_constants = [4]u64{
15 0x76e1_5d3e_fefd_cbbf,
16 0xc500_4e44_1c52_2fb3,
17 0x7771_0069_854e_e241,
18 0x3910_9bb0_2acb_e635,
19 };
20
21 const uniform_scale: f64 = 0x1.0p-53;
22
23 pub const SplitMix64 = struct {
24 state: u64,
25
26 pub fn init(state: u64) SplitMix64 {
27 return .{ .state = state };
28 }
29
30 pub fn next(self: *SplitMix64) u64 {
31 self.state +%= 0x9e37_79b9_7f4a_7c15;
32 var value = self.state;
33 value = (value ^ (value >> 30)) *% 0xbf58_476d_1ce4_e5b9;
34 value = (value ^ (value >> 27)) *% 0x94d0_49bb_1331_11eb;
35 return value ^ (value >> 31);
36 }
37 };
38
39 pub const Xoshiro = struct {
40 state: [4]u64,
41
42 pub fn init(seed: u64) Xoshiro {
43 var split = SplitMix64.init(seed);
44 var state: [4]u64 = undefined;
45 for (&state) |*value| value.* = split.next();
46 return .{ .state = state };
47 }
48
49 pub fn initThread(seed: u64, thread_id: u64) Xoshiro {
50 var result = init(seed);
51 var stream: u64 = 0;
52 while (stream < thread_id) : (stream += 1) result.jump();
53 return result;
54 }
55
56 pub fn next(self: *Xoshiro) u64 {
57 const result = std.math.rotl(u64, self.state[0] +% self.state[3], 23) +%
58 self.state[0];
59 const temporary = self.state[1] << 17;
60 self.state[2] ^= self.state[0];
61 self.state[3] ^= self.state[1];
62 self.state[1] ^= self.state[2];
63 self.state[0] ^= self.state[3];
64 self.state[2] ^= temporary;
65 self.state[3] = std.math.rotl(u64, self.state[3], 45);
66 return result;
67 }
68
69 pub fn uniform(self: *Xoshiro) f64 {
70 return @as(f64, @floatFromInt(self.next() >> 11)) * uniform_scale;
71 }
72
73 pub fn getState(self: Xoshiro) [4]u64 {
74 return self.state;
75 }
76
77 pub fn setState(self: *Xoshiro, state: [4]u64) void {
78 self.state = state;
79 }
80
81 pub fn stateSize() usize {
82 return 4;
83 }
84
85 pub fn jump(self: *Xoshiro) void {
86 self.jumpWith(jump_constants);
87 }
88
89 pub fn longJump(self: *Xoshiro) void {
90 self.jumpWith(long_jump_constants);
91 }
92
93 fn jumpWith(self: *Xoshiro, constants: [4]u64) void {
94 var state: [4]u64 = @splat(0);
95 for (constants) |bits| {
96 for (0..64) |bit| {
97 if (bits & (@as(u64, 1) << @intCast(bit)) != 0) {
98 inline for (0..4) |index| state[index] ^= self.state[index];
99 }
100 _ = self.next();
101 }
102 }
103 self.state = state;
104 }
105 };
106
107 pub fn VectorXoshiro(comptime D: type) type {
108 requireU64(D);
109 return struct {
110 state: [4]D.Vector,
111
112 const Self = @This();
113 const DF: type = D.rebind(f64);
114
115 pub fn init(seed: u64) Self {
116 return initThread(seed, 0);
117 }
118
119 pub fn initThread(seed: u64, thread_number: u64) Self {
120 var scalar = Xoshiro.init(seed);
121 var stream: u64 = 0;
122 while (stream < thread_number) : (stream += 1) scalar.longJump();
123 var state: [4]D.Vector = @splat(@splat(0));
124 inline for (0..D.lane_count) |lane| {
125 const scalar_state = scalar.getState();
126 inline for (0..4) |index| state[index][lane] = scalar_state[index];
127 scalar.jump();
128 }
129 return .{ .state = state };
130 }
131
132 pub fn next(self: *Self) D.Vector {
133 return update(&self.state);
134 }
135
136 pub fn fill(self: *Self, output: []u64) void {
137 var index: usize = 0;
138 while (index + D.lane_count <= output.len) : (index += D.lane_count) {
139 const value: [D.lane_count]u64 = self.next();
140 @memcpy(output[index .. index + D.lane_count], &value);
141 }
142 if (index != output.len) {
143 const value: [D.lane_count]u64 = self.next();
144 @memcpy(output[index..], value[0 .. output.len - index]);
145 }
146 }
147
148 pub fn uniform(self: *Self) DF.Vector {
149 return toUniform(self.next());
150 }
151
152 pub fn fillUniform(self: *Self, output: []f64) void {
153 var index: usize = 0;
154 while (index + D.lane_count <= output.len) : (index += D.lane_count) {
155 const value: [D.lane_count]f64 = self.uniform();
156 @memcpy(output[index .. index + D.lane_count], &value);
157 }
158 if (index != output.len) {
159 const value: [D.lane_count]f64 = self.uniform();
160 @memcpy(output[index..], value[0 .. output.len - index]);
161 }
162 }
163
164 pub fn getState(self: Self) [4]D.Vector {
165 return self.state;
166 }
167
168 pub fn setState(self: *Self, state: [4]D.Vector) void {
169 self.state = state;
170 }
171
172 pub fn stateSize() usize {
173 return 4 * D.lane_count;
174 }
175
176 fn update(state: *[4]D.Vector) D.Vector {
177 const result = rotate.rotateRightSame(D, state[0] +% state[3], 41) +% state[0];
178 const temporary = shift.shiftLeft(D, 17, state[1]);
179 state[2] ^= state[0];
180 state[3] ^= state[1];
181 state[1] ^= state[2];
182 state[0] ^= state[3];
183 state[2] ^= temporary;
184 state[3] = rotate.rotateRightSame(D, state[3], 19);
185 return result;
186 }
187
188 fn toUniform(value: D.Vector) DF.Vector {
189 const bits = shift.shiftRight(D, 11, value);
190 var result: DF.Vector = undefined;
191 inline for (0..D.lane_count) |index| {
192 result[index] = @as(f64, @floatFromInt(bits[index])) * uniform_scale;
193 }
194 return result;
195 }
196 };
197 }
198
199 pub fn CachedXoshiro(comptime D: type, comptime cache_size: usize) type {
200 if (cache_size == 0 or !std.math.isPowerOfTwo(cache_size)) {
201 @compileError("cached xoshiro size must be a nonzero power of two");
202 }
203 return struct {
204 generator: Generator,
205 cache: [cache_size]u64,
206 index: usize,
207
208 const Self = @This();
209 const Generator: type = VectorXoshiro(D);
210
211 pub fn init(seed: u64) Self {
212 return initThread(seed, 0);
213 }
214
215 pub fn initThread(seed: u64, thread_number: u64) Self {
216 var result = Self{
217 .generator = Generator.initThread(seed, thread_number),
218 .cache = undefined,
219 .index = 0,
220 };
221 result.generator.fill(&result.cache);
222 return result;
223 }
224
225 pub fn next(self: *Self) u64 {
226 std.debug.assert(self.index <= cache_size);
227 if (self.index == cache_size) {
228 self.generator.fill(&self.cache);
229 self.index = 0;
230 }
231 const result = self.cache[self.index];
232 self.index += 1;
233 return result;
234 }
235
236 pub fn minimum() u64 {
237 return 0;
238 }
239
240 pub fn maximum() u64 {
241 return std.math.maxInt(u64);
242 }
243 };
244 }
245
246 pub fn DefaultCachedXoshiro(comptime D: type) type {
247 return CachedXoshiro(D, 1024);
248 }
249
250 pub const AesCtrEngine = struct {
251 key: [2 * (1 + rounds)]u64,
252
253 const rounds: usize = 5;
254 const D: type = tag.Full128(u8);
255 const D64: type = tag.Full128(u64);
256
257 pub fn initDeterministic() AesCtrEngine {
258 return initKey(.{ 0x243f_6a88_85a3_08d3, 0x1319_8a2e_0370_7344 });
259 }
260
261 pub fn initKey(seed: [2]u64) AesCtrEngine {
262 var result: AesCtrEngine = undefined;
263 result.key[0] = seed[0];
264 result.key[1] = seed[1];
265 inline for (0..rounds) |index| {
266 result.key[2 + 2 * index] = result.key[2 * index + 1] +%
267 0xa409_3822_299f_31d0;
268 result.key[2 + 2 * index + 1] = result.key[2 * index] +%
269 0x082e_fa98_ec4e_6c89;
270 }
271 return result;
272 }
273
274 pub fn initWithEntropy(entropy: anytype) !AesCtrEngine {
275 var bytes: [16]u8 = undefined;
276 try entropy.fill(&bytes);
277 return initKey(@bitCast(bytes));
278 }
279
280 pub fn generate(self: *const AesCtrEngine, stream: u64, counter: u64) u64 {
281 var state: D.Vector = @bitCast(@as(D64.Vector, .{ counter, stream }));
282 state ^= keyVector(self.key[0..2].*);
283 inline for (0..rounds) |index| {
284 state = crypto.aesRound(D, state, keyVector(self.key[2 + 2 * index ..][0..2].*));
285 }
286 const result: D64.Vector = @bitCast(state);
287 return result[0];
288 }
289
290 fn keyVector(pair: [2]u64) D.Vector {
291 return @bitCast(@as(D64.Vector, pair));
292 }
293 };
294
295 pub const RngStream = struct {
296 engine: *const AesCtrEngine,
297 stream: u64,
298 counter: u64 = 0,
299
300 pub fn init(engine: *const AesCtrEngine, stream: u64) RngStream {
301 return .{ .engine = engine, .stream = stream };
302 }
303
304 pub fn next(self: *RngStream) u64 {
305 const result = self.engine.generate(self.stream, self.counter);
306 self.counter +%= 1;
307 return result;
308 }
309
310 pub fn minimum() u64 {
311 return 0;
312 }
313
314 pub fn maximum() u64 {
315 return std.math.maxInt(u64);
316 }
317 };
318
319 pub fn randomNormalizedFloat(stream: *RngStream) f32 {
320 const exponent = @as(u32, @bitCast(@as(f32, 1)));
321 const mantissa = @as(u32, @truncate(stream.next())) & 0x007f_ffff;
322 const value: f32 = @bitCast(exponent | mantissa);
323 std.debug.assert(value >= 1);
324 std.debug.assert(value < 2);
325 const result = (2 * (value - 1)) - 1;
326 std.debug.assert(result >= -1);
327 std.debug.assert(result < 1);
328 return result;
329 }
330
331 pub fn fillRandom(comptime T: type, engine: *const AesCtrEngine, stream_id: u64, output: []T) void {
332 var stream = RngStream.init(engine, stream_id);
333 for (output) |*value| value.* = castRandom(T, stream.next());
334 }
335
336 fn castRandom(comptime T: type, value: u64) T {
337 return switch (@typeInfo(T)) {
338 .int => |info| if (info.signedness == .unsigned)
339 @truncate(value)
340 else blk: {
341 const U = @Int(.unsigned, info.bits);
342 break :blk @bitCast(@as(U, @truncate(value)));
343 },
344 .float => @floatFromInt(value),
345 else => @compileError("fillRandom requires integer or floating-point outputs"),
346 };
347 }
348
349 fn requireU64(comptime D: type) void {
350 if (comptime D.Lane != u64) @compileError("vector xoshiro requires u64 lanes");
351 }
352
353 const TestEntropy = struct {
354 offset: u8,
355
356 fn fill(self: *@This(), output: []u8) !void {
357 for (output, 0..) |*byte, index| byte.* = @intCast(index * 7 + self.offset);
358 }
359 };
360
361 test "Highway scalar xoshiro state replay and stream jumps are deterministic" {
362 var generator = Xoshiro.init(123);
363 const state = generator.getState();
364 const first = generator.next();
365 const second = generator.next();
366 try std.testing.expect(first != second);
367 generator.setState(state);
368 try std.testing.expectEqual(first, generator.next());
369 try std.testing.expectEqual(second, generator.next());
370 var jumped = Xoshiro.initThread(123, 1);
371 var reference = Xoshiro.init(123);
372 reference.jump();
373 try std.testing.expectEqual(reference.getState(), jumped.getState());
374 try std.testing.expectEqual(@as(usize, 4), Xoshiro.stateSize());
375 }
376
377 test "Highway AVX2 scalar and vector xoshiro oracle matches exactly" {
378 var scalar = Xoshiro.init(123);
379 try std.testing.expectEqual(
380 [4]u64{
381 0xb4dc_9bd4_62de_412b,
382 0xfa02_3ce9_f06f_b77c,
383 0xdc12_d311_d371_cbe8,
384 0xafd2_040c_9098_81ff,
385 },
386 scalar.getState(),
387 );
388 for ([_]u64{
389 0xa556_5735_f810_987a,
390 0xd691_4642_e58d_662e,
391 0xaa75_21fe_b709_887f,
392 0x863c_d15c_558d_6bfb,
393 }) |expected| try std.testing.expectEqual(expected, scalar.next());
394 var jumped = Xoshiro.init(123);
395 jumped.jump();
396 try std.testing.expectEqual(
397 [4]u64{
398 0xed6f_8f0b_3989_38af,
399 0x17c5_be01_9e09_5507,
400 0x05a0_a7d3_71bf_d778,
401 0xbc77_948a_a522_4033,
402 },
403 jumped.getState(),
404 );
405 var long_jumped = Xoshiro.init(123);
406 long_jumped.longJump();
407 try std.testing.expectEqual(
408 [4]u64{
409 0xf98f_04cb_080b_8942,
410 0xd884_97e0_e599_e6f1,
411 0x3b4f_07f1_f6d2_6dd4,
412 0x5d3e_2cd5_d040_18c1,
413 },
414 long_jumped.getState(),
415 );
416 const D = tag.FixedTag(u64, 4);
417 const Generator = VectorXoshiro(D);
418 var vector = Generator.init(456);
419 try std.testing.expect(@reduce(.And, vector.next() == @as(D.Vector, .{
420 0xa901_0f72_89ae_bac2,
421 0x6990_d4cc_67a1_47d7,
422 0x76bd_c3fa_db67_4b7c,
423 0x6580_3fd8_b3a1_a35d,
424 })));
425 try std.testing.expect(@reduce(.And, vector.next() == @as(D.Vector, .{
426 0x5fcc_d4ec_4e33_71f4,
427 0x5321_fc38_a086_1200,
428 0x0bce_898c_528d_1264,
429 0xf716_df28_1d99_3979,
430 })));
431 var uniform = Generator.init(456);
432 const real: [4]f64 = uniform.uniform();
433 try std.testing.expectEqual(
434 [4]u64{
435 0x3fe5_2021_ee51_35d7,
436 0x3fda_6435_3319_e850,
437 0x3fdd_af70_feb6_d9d2,
438 0x3fd9_600f_f62c_e868,
439 },
440 @as([4]u64, @bitCast(real)),
441 );
442 }
443
444 test "Highway vector xoshiro interleaves jump-separated scalar streams" {
445 const D = tag.FixedTag(u64, 4);
446 const Generator = VectorXoshiro(D);
447 var generator = Generator.init(456);
448 var streams: [D.lane_count]Xoshiro = undefined;
449 streams[0] = Xoshiro.init(456);
450 for (1..D.lane_count) |index| {
451 streams[index] = streams[index - 1];
452 streams[index].jump();
453 }
454 for (0..7) |_| {
455 const actual: [D.lane_count]u64 = generator.next();
456 for (&streams, 0..) |*stream, lane| {
457 try std.testing.expectEqual(stream.next(), actual[lane]);
458 }
459 }
460 try std.testing.expectEqual(@as(usize, 16), Generator.stateSize());
461 }
462
463 test "Highway vector xoshiro thread streams begin at long-jump partitions" {
464 const D = tag.FixedTag(u64, 4);
465 const Generator = VectorXoshiro(D);
466 var generator = Generator.initThread(654, 3);
467 var reference = Xoshiro.init(654);
468 for (0..3) |_| reference.longJump();
469 const state = generator.getState();
470 inline for (0..D.lane_count) |lane| {
471 const scalar_state = reference.getState();
472 inline for (0..4) |index| {
473 try std.testing.expectEqual(scalar_state[index], state[index][lane]);
474 }
475 reference.jump();
476 }
477 }
478
479 test "Highway vector xoshiro fills awkward tails and uniform doubles" {
480 const D = tag.FixedTag(u64, 4);
481 const Generator = VectorXoshiro(D);
482 var integers = Generator.init(789);
483 var expected = Generator.init(789);
484 var output: [13]u64 = undefined;
485 integers.fill(&output);
486 for (0..4) |block| {
487 const values: [D.lane_count]u64 = expected.next();
488 const begin = block * D.lane_count;
489 const count = @min(D.lane_count, output.len - begin);
490 for (0..count) |lane| try std.testing.expectEqual(values[lane], output[begin + lane]);
491 }
492 var uniforms = Generator.init(789);
493 var real: [13]f64 = undefined;
494 uniforms.fillUniform(&real);
495 for (real) |value| try std.testing.expect(value >= 0 and value < 1);
496 }
497
498 test "Highway cached xoshiro refills without changing stream order" {
499 const D = tag.FixedTag(u64, 4);
500 const Cached = CachedXoshiro(D, 8);
501 const Generator = VectorXoshiro(D);
502 var cached = Cached.init(987);
503 var vector = Generator.init(987);
504 var expected: [24]u64 = undefined;
505 vector.fill(&expected);
506 for (expected) |value| try std.testing.expectEqual(value, cached.next());
507 try std.testing.expectEqual(@as(u64, 0), Cached.minimum());
508 try std.testing.expectEqual(std.math.maxInt(u64), Cached.maximum());
509 }
510
511 test "Highway AES counter streams are deterministic and independently addressed" {
512 const engine = AesCtrEngine.initDeterministic();
513 var first = RngStream.init(&engine, 7);
514 var replay = RngStream.init(&engine, 7);
515 var other = RngStream.init(&engine, 8);
516 for (0..16) |_| {
517 try std.testing.expectEqual(first.next(), replay.next());
518 try std.testing.expect(first.counter == replay.counter);
519 }
520 try std.testing.expect(engine.generate(7, 0) != other.next());
521 }
522
523 test "Highway AVX2 AES counter and normalized float oracle matches exactly" {
524 const engine = AesCtrEngine.initDeterministic();
525 try std.testing.expectEqual(@as(u64, 0xf0ea_ad2f_c4a2_c3e1), engine.generate(7, 0));
526 try std.testing.expectEqual(@as(u64, 0xbf34_0afe_30f5_275d), engine.generate(7, 1));
527 try std.testing.expectEqual(@as(u64, 0x0797_7286_13ce_fc88), engine.generate(8, 0));
528 var stream = RngStream.init(&engine, 0);
529 for ([_]u32{ 0xbed6_6698, 0x3e46_f1b0, 0x3f17_03d4 }) |bits| {
530 try std.testing.expectEqual(bits, @as(u32, @bitCast(randomNormalizedFloat(&stream))));
531 }
532 }
533
534 test "Highway AES counter accepts caller entropy and fills typed buffers" {
535 var entropy = TestEntropy{ .offset = 3 };
536 const engine = try AesCtrEngine.initWithEntropy(&entropy);
537 var other_entropy = TestEntropy{ .offset = 11 };
538 const other_engine = try AesCtrEngine.initWithEntropy(&other_entropy);
539 try std.testing.expect(engine.generate(0, 0) != other_engine.generate(0, 0));
540 var actual: [17]u16 = undefined;
541 fillRandom(u16, &engine, 91, &actual);
542 var stream = RngStream.init(&engine, 91);
543 for (actual) |value| try std.testing.expectEqual(@as(u16, @truncate(stream.next())), value);
544 }
545
546 test "Highway normalized float spans the half-open signed unit interval" {
547 const engine = AesCtrEngine.initDeterministic();
548 var stream = RngStream.init(&engine, 0);
549 var sum: f64 = 0;
550 for (0..20_000) |_| {
551 const value = randomNormalizedFloat(&stream);
552 try std.testing.expect(value >= -1 and value < 1);
553 sum += value;
554 }
555 try std.testing.expect(@abs(sum / 20_000) < 0.02);
556 }
557
558 test "Highway AES counter bit and byte distributions remain balanced" {
559 const engine = AesCtrEngine.initDeterministic();
560 var stream = RngStream.init(&engine, 0);
561 const count: usize = 20_000;
562 var one_bits: u64 = 0;
563 for (0..count) |_| one_bits += @popCount(stream.next());
564 const bit_ratio = @as(f64, @floatFromInt(one_bits)) /
565 @as(f64, @floatFromInt(count * 64));
566 try std.testing.expect(bit_ratio >= 0.49 and bit_ratio <= 0.51);
567 for (0..8) |byte_index| {
568 var counts: [256]u32 = @splat(0);
569 const amount: u6 = @intCast(byte_index * 8);
570 for (0..count) |_| counts[@as(u8, @truncate(stream.next() >> amount))] += 1;
571 const expected = @as(f64, @floatFromInt(count)) / 256;
572 var chi_squared: f64 = 0;
573 for (counts) |observed| {
574 const difference = @as(f64, @floatFromInt(observed)) - expected;
575 chi_squared += difference * difference / expected;
576 }
577 try std.testing.expect(chi_squared >= 170 and chi_squared <= 340);
578 }
579 }