lib/accy/src/kernel/library/random/philox/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 Philox = model.Philox;
15 const philox_lanes = block.philox_lanes;
16 const philoxWords = block.philoxWords;
17 const outputWord = block.outputWord;
18 const storeGuarded = block.storeGuarded;
19 
20 fn philox_body_each(inner_builder: anytype, index: kernel.Index1D, ctx: anytype) !void {
21     const words = try philoxWords(
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(philox_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 % philox_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 philoxBody(k: anytype, comptime dtype: DType, spec: Philox, args: anytype) !void {
45     _ = try k.forEach1D(spec.count_axis, spec.generators(), .{ .spec = spec, .args = args, .dtype = dtype }, philox_body_each);
46 }
47 fn philoxSpecialization(comptime spec: Philox) entry.Specialization {
48     return .{
49         .dtype = spec.dtype,
50         .operation = .{ .random = .{ .philox = 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 philoxProgram(comptime spec: Philox) type {
58     const Body = struct {
59         fn run(k: anytype, args: anytype) !void {
60             try philoxBody(k, spec.dtype, spec, args);
61         }
62     };
63 
64     return kernel.logical.Program(.{
65         .name = std.fmt.comptimePrint(
66             "accy_kernel_random_philox{}_{}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 philoxEntry(comptime spec: Philox) type {
78     return entry.Entry(philoxProgram(spec), .{
79         .target = std.fmt.comptimePrint(
80             "accy.kernel.random.philox{}_{}r_{}_{s}",
81             .{ spec.count, spec.rounds, spec.threads, spec.dtype.name() },
82         ),
83         .layer = .logical,
84         .category = .random,
85         .specialization = philoxSpecialization(spec),
86     });
87 }
88 pub const Philox8F32 = philoxEntry(.{ .count = 8, .threads = 2 });