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

daab053ee43316e1809a84551d573ddd1e5bf3d2

 1 const random = @import("../root.zig");
 2 
 3 const shared = random.base;
 4 const block = random.block;
 5 const model = @import("model.zig");
 6 
 7 const std = shared.std;
 8 const DType = shared.DType;
 9 const kernel = shared.kernel;
10 
11 const Philox = model.Philox;
12 const philox_lanes = block.philox_lanes;
13 const philoxWords = block.philoxWords;
14 const outputWord = block.outputWord;
15 const storeGuarded = block.storeGuarded;
16 
17 fn philoxRuntimeBody(k: anytype, comptime dtype: DType, spec: Philox, args: anytype) !void {
18     const generator = try k.globalId(.x);
19     const count = try k.castIndex(args.param(.count).raw());
20     const lanes_extent = try k.constantIndex(philox_lanes);
21     const base = try k.mul(generator, lanes_extent);
22     const words = try philoxWords(
23         k,
24         spec.rounds,
25         try k.cast(generator, .i32),
26         try k.constantInt(.i32, 0),
27         args.param(.seed_lo).raw(),
28         args.param(.seed_hi).raw(),
29     );
30     for (words, 0..) |word, lane| {
31         const element = try k.add(base, try k.constantIndex(@intCast(lane)));
32         const out = try outputWord(k, dtype, word);
33         try storeGuarded(k, args, out, element, count);
34     }
35 }
36 fn philoxFamilySchedule(instance: Philox) kernel.logical.schedule.ThreadBlocks {
37     return kernel.logical.schedule.threadBlocks(.{ .x = instance.threads });
38 }
39 fn philox_runtime_family_body_f32(k: anytype, spec: Philox, args: anytype) !void {
40     try philoxRuntimeBody(k, .f32, spec, args);
41 }
42 fn philox_runtime_family_body_i32(k: anytype, spec: Philox, args: anytype) !void {
43     try philoxRuntimeBody(k, .i32, spec, args);
44 }
45 
46 fn philoxRuntimeFamily(comptime dtype: DType) type {
47     return kernel.logical.Family(.{
48         .name = std.fmt.comptimePrint("accy_kernel_random_philox_runtime_{s}", .{dtype.name()}),
49         .parameters = .{
50             .dst = kernel.dynamicBuffer(dtype),
51             .count = kernel.scalar(.i32),
52             .seed_lo = kernel.scalar(.i32),
53             .seed_hi = kernel.scalar(.i32),
54         },
55         .Instance = Philox,
56         .schedule = philoxFamilySchedule,
57         .body = switch (dtype) {
58             .f32 => philox_runtime_family_body_f32,
59             .i32 => philox_runtime_family_body_i32,
60             else => @compileError("Philox runtime supports dtype .f32 or .i32"),
61         },
62     });
63 }
64 
65 pub const PhiloxRuntimeFamilyF32 = philoxRuntimeFamily(.f32);
66 pub const PhiloxRuntimeFamilyI32 = philoxRuntimeFamily(.i32);