lib/accy/src/kernel/library/random/threefry/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 Threefry = model.Threefry;
12 const threefry_lanes = block.threefry_lanes;
13 const threefryWords = block.threefryWords;
14 const outputWord = block.outputWord;
15 const storeGuarded = block.storeGuarded;
16 
17 fn threefryRuntimeBody(k: anytype, comptime dtype: DType, spec: Threefry, 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(threefry_lanes);
21     const base = try k.mul(generator, lanes_extent);
22     const words = try threefryWords(
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 threefryFamilySchedule(instance: Threefry) kernel.logical.schedule.ThreadBlocks {
37     return kernel.logical.schedule.threadBlocks(.{ .x = instance.threads });
38 }
39 fn threefry_runtime_family_body_f32(k: anytype, spec: Threefry, args: anytype) !void {
40     try threefryRuntimeBody(k, .f32, spec, args);
41 }
42 fn threefry_runtime_family_body_i32(k: anytype, spec: Threefry, args: anytype) !void {
43     try threefryRuntimeBody(k, .i32, spec, args);
44 }
45 
46 fn threefryRuntimeFamily(comptime dtype: DType) type {
47     return kernel.logical.Family(.{
48         .name = std.fmt.comptimePrint("accy_kernel_random_threefry_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 = Threefry,
56         .schedule = threefryFamilySchedule,
57         .body = switch (dtype) {
58             .f32 => threefry_runtime_family_body_f32,
59             .i32 => threefry_runtime_family_body_i32,
60             else => @compileError("Threefry runtime supports dtype .f32 or .i32"),
61         },
62     });
63 }
64 
65 pub const ThreefryRuntimeFamilyF32 = threefryRuntimeFamily(.f32);
66 pub const ThreefryRuntimeFamilyI32 = threefryRuntimeFamily(.i32);