lib/accy/src/kernel/library/random/threefry/static.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 entry = shared.entry;
10 const kernel = shared.kernel;
11 const indexExtent = shared.indexExtent;
12 const ceilDivComptime = shared.ceilDivComptime;
13 
14 const Threefry = model.Threefry;
15 const threefry_lanes = block.threefry_lanes;
16 const threefryWords = block.threefryWords;
17 const outputWord = block.outputWord;
18 const storeGuarded = block.storeGuarded;
19 
20 fn threefry_body_each(inner_builder: anytype, index: kernel.Index1D, ctx: anytype) !void {
21     const words = try threefryWords(
22         inner_builder,
23         ctx.spec.rounds,
24         try inner_builder.cast(index.index, .i32),
25         try inner_builder.constantInt(.i32, 0),
26         ctx.args.param(.seed_lo).raw(),
27         ctx.args.param(.seed_hi).raw(),
28     );
29     const lanes_extent = try inner_builder.constantIndex(threefry_lanes);
30     const base = try inner_builder.mul(index.index, lanes_extent);
31     const count = try inner_builder.constantIndex(try indexExtent(ctx.spec.count));
32     const aligned = ctx.spec.count % threefry_lanes == 0;
33     for (words, 0..) |word, lane| {
34         const element = try inner_builder.add(base, try inner_builder.constantIndex(@intCast(lane)));
35         const out = try outputWord(inner_builder, ctx.dtype, word);
36         if (aligned) {
37             try ctx.args.param(.dst).store(inner_builder, out, element);
38         } else {
39             try storeGuarded(inner_builder, ctx.args, out, element, count);
40         }
41     }
42 }
43 
44 pub fn threefryBody(k: anytype, comptime dtype: DType, spec: Threefry, args: anytype) !void {
45     _ = try k.forEach1D(spec.count_axis, spec.generators(), .{ .spec = spec, .args = args, .dtype = dtype }, threefry_body_each);
46 }
47 fn threefrySpecialization(comptime spec: Threefry) entry.Specialization {
48     return .{
49         .dtype = spec.dtype,
50         .operation = .{ .random = .{ .threefry = spec.rounds } },
51         .inputs = &.{},
52         .outputs = &.{entry.shape1D(spec.count_axis, spec.count)},
53         .launch = entry.launch1D(ceilDivComptime(spec.generators(), spec.threads), spec.threads),
54         .schedule = entry.threadBlocks1D("e", spec.generators(), spec.threads),
55     };
56 }
57 fn threefryProgram(comptime spec: Threefry) type {
58     const Body = struct {
59         fn run(k: anytype, args: anytype) !void {
60             try threefryBody(k, spec.dtype, spec, args);
61         }
62     };
63 
64     return kernel.logical.Program(.{
65         .name = std.fmt.comptimePrint(
66             "accy_kernel_random_threefry{}_{}r_{}_{s}",
67             .{ spec.count, spec.rounds, spec.threads, spec.dtype.name() },
68         ),
69         .parameters = .{
70             .dst = kernel.dynamicBuffer(spec.dtype),
71             .seed_lo = kernel.scalar(.i32),
72             .seed_hi = kernel.scalar(.i32),
73         },
74         .body = Body.run,
75     }).withSchedule(kernel.logical.schedule.threadBlocks(.{ .x = spec.threads }));
76 }
77 pub fn threefryEntry(comptime spec: Threefry) type {
78     return entry.Entry(threefryProgram(spec), .{
79         .target = std.fmt.comptimePrint(
80             "accy.kernel.random.threefry{}_{}r_{}_{s}",
81             .{ spec.count, spec.rounds, spec.threads, spec.dtype.name() },
82         ),
83         .layer = .logical,
84         .category = .random,
85         .specialization = threefrySpecialization(spec),
86     });
87 }
88 pub const Threefry8F32 = threefryEntry(.{ .count = 8, .threads = 4 });