lib/accy/src/kernel/library/random/feistel/runtime.zig

daab053ee43316e1809a84551d573ddd1e5bf3d2

 1 const random = @import("../root.zig");
 2 const model = @import("model.zig");
 3 
 4 const shared = random.base;
 5 const block = random.block;
 6 
 7 const std = shared.std;
 8 const DType = shared.DType;
 9 const kernel = shared.kernel;
10 
11 const Feistel = model.Feistel;
12 const philoxWords = block.philoxWords;
13 const philoxBlock = block.philoxBlock;
14 const wordConstant = block.wordConstant;
15 
16 pub fn feistelPermuteReference(index: u32, rounds: u32, seed: u64, bits: u6) u32 {
17     const half_bits: u5 = @intCast(bits / 2);
18     const mask: u32 = (@as(u32, 1) << half_bits) - 1;
19     var left: u32 = index >> half_bits;
20     var right: u32 = index & mask;
21     const seed_lo: u32 = @truncate(seed);
22     const seed_hi: u32 = @truncate(seed >> 32);
23     var round: u32 = 0;
24     while (round < rounds) : (round += 1) {
25         const mixed = philoxBlock(2, .{ right, round, 0, 0 }, .{ seed_lo, seed_hi })[0] & mask;
26         const next_right = left ^ mixed;
27         left = right;
28         right = next_right;
29     }
30     return (left << half_bits) | right;
31 }
32 
33 fn feistelPermuteValue(k: anytype, spec: Feistel, value: kernel.Value, seed_lo: kernel.Value, seed_hi: kernel.Value) !kernel.Value {
34     const bits = spec.domainBits().?;
35     const half_bits: u5 = @intCast(bits / 2);
36     const mask_value = (@as(u32, 1) << half_bits) - 1;
37     const shift = try k.constantInt(.i32, half_bits);
38     const mask = try wordConstant(k, mask_value);
39     var left = try k.ushr(value, shift);
40     var right = try k.and_(value, mask);
41     var round: u32 = 0;
42     while (round < spec.rounds) : (round += 1) {
43         const round_value = try k.constantInt(.i32, @intCast(round));
44         const words = try philoxWords(k, 2, right, round_value, seed_lo, seed_hi);
45         const mixed = try k.and_(words[0], mask);
46         const next_right = try k.xor(left, mixed);
47         left = right;
48         right = next_right;
49     }
50     return k.or_(try k.shl(left, shift), right);
51 }
52 
53 fn feistel_runtime_body_active(inner_builder: anytype, ctx: anytype) !void {
54     const source = try inner_builder.cast(ctx.element, .i32);
55     const out = try feistelPermuteValue(
56         inner_builder,
57         ctx.spec,
58         source,
59         ctx.args.param(.seed_lo).raw(),
60         ctx.args.param(.seed_hi).raw(),
61     );
62     try ctx.args.param(.dst).store(inner_builder, out, ctx.element);
63 }
64 
65 fn feistelRuntimeBody(k: anytype, comptime dtype: DType, spec: Feistel, args: anytype) !void {
66     if (comptime dtype != .i32) @compileError("feistel permutation supports dtype .i32");
67     const element = try k.globalId(.x);
68     const count = try k.castIndex(args.param(.count).raw());
69     const active = try k.compare(.lt, element, count);
70     try k.guardDo(active, .{ .args = args, .element = element, .spec = spec }, feistel_runtime_body_active);
71 }
72 
73 fn feistelFamilySchedule(instance: Feistel) kernel.logical.schedule.ThreadBlocks {
74     return kernel.logical.schedule.threadBlocks(.{ .x = instance.threads });
75 }
76 
77 fn feistel_runtime_family_body_i32(k: anytype, spec: Feistel, args: anytype) !void {
78     try feistelRuntimeBody(k, .i32, spec, args);
79 }
80 
81 fn feistelRuntimeFamily(comptime dtype: DType) type {
82     return kernel.logical.Family(.{
83         .name = std.fmt.comptimePrint("accy_kernel_random_feistel_runtime_{s}", .{dtype.name()}),
84         .parameters = .{
85             .dst = kernel.dynamicBuffer(dtype),
86             .count = kernel.scalar(.i32),
87             .seed_lo = kernel.scalar(.i32),
88             .seed_hi = kernel.scalar(.i32),
89         },
90         .Instance = Feistel,
91         .schedule = feistelFamilySchedule,
92         .body = switch (dtype) {
93             .i32 => feistel_runtime_family_body_i32,
94             else => @compileError("feistel runtime supports dtype .i32"),
95         },
96     });
97 }
98 
99 pub const FeistelRuntimeFamilyI32 = feistelRuntimeFamily(.i32);