lib/accy/src/kernel/library/random/fold/threefry/runtime.zig

daab053ee43316e1809a84551d573ddd1e5bf3d2

 1 const random = @import("../../root.zig");
 2 
 3 const base = random.base;
 4 const block = random.block;
 5 const fold = @import("../root.zig");
 6 const common = fold.common;
 7 const types = @import("types.zig");
 8 
 9 const std = base.std;
10 const DType = base.DType;
11 const kernel = base.kernel;
12 
13 const ThreefryFold = types.ThreefryFold;
14 const threefryWords = block.threefryWords;
15 const foldInitial = common.foldInitial;
16 const foldCombine = common.foldCombine;
17 
18 fn threefry_fold_runtime_body_active(inner_builder: anytype, ctx: anytype) !void {
19     const zero = try inner_builder.constantIndex(0);
20     const one = try inner_builder.constantIndex(1);
21     const initial = try foldInitial(inner_builder, ctx.dtype);
22     const counter = try inner_builder.cast(ctx.output, .i32);
23     const folded = try inner_builder.fold(zero, ctx.samples, one, initial, .{
24         .spec = ctx.spec,
25         .args = ctx.args,
26         .counter = counter,
27         .dtype = ctx.dtype,
28     }, threefry_fold_runtime_body_apply);
29     try ctx.args.param(.dst).store(inner_builder, folded, ctx.output);
30 }
31 
32 fn threefry_fold_runtime_body_apply(fold_builder: anytype, sample: kernel.Value, current: kernel.Value, fold_ctx: anytype) !kernel.Value {
33     const words = try threefryWords(
34         fold_builder,
35         fold_ctx.spec.rounds,
36         fold_ctx.counter,
37         try fold_builder.cast(sample, .i32),
38         fold_ctx.args.param(.seed_lo).raw(),
39         fold_ctx.args.param(.seed_hi).raw(),
40     );
41     var accumulator = current;
42     for (words) |word| accumulator = try foldCombine(fold_builder, fold_ctx.dtype, accumulator, word);
43     return accumulator;
44 }
45 
46 fn threefryFoldRuntimeBody(k: anytype, comptime dtype: DType, spec: ThreefryFold, args: anytype) !void {
47     const output = try k.globalId(.x);
48     const count = try k.castIndex(args.param(.count).raw());
49     const samples = try k.castIndex(args.param(.samples).raw());
50     const active = try k.compare(.lt, output, count);
51     try k.guardDo(active, .{ .spec = spec, .args = args, .output = output, .samples = samples, .dtype = dtype }, threefry_fold_runtime_body_active);
52 }
53 
54 fn threefryFoldFamilySchedule(instance: ThreefryFold) kernel.logical.schedule.ThreadBlocks {
55     return kernel.logical.schedule.threadBlocks(.{ .x = instance.threads });
56 }
57 
58 fn threefry_fold_runtime_family_body_f32(k: anytype, spec: ThreefryFold, args: anytype) !void {
59     try threefryFoldRuntimeBody(k, .f32, spec, args);
60 }
61 
62 fn threefry_fold_runtime_family_body_i32(k: anytype, spec: ThreefryFold, args: anytype) !void {
63     try threefryFoldRuntimeBody(k, .i32, spec, args);
64 }
65 
66 fn threefryFoldRuntimeFamily(comptime dtype: DType) type {
67     return kernel.logical.Family(.{
68         .name = std.fmt.comptimePrint("accy_kernel_random_threefry_fold_runtime_{s}", .{dtype.name()}),
69         .parameters = .{
70             .dst = kernel.dynamicBuffer(dtype),
71             .count = kernel.scalar(.i32),
72             .samples = kernel.scalar(.i32),
73             .seed_lo = kernel.scalar(.i32),
74             .seed_hi = kernel.scalar(.i32),
75         },
76         .Instance = ThreefryFold,
77         .schedule = threefryFoldFamilySchedule,
78         .body = switch (dtype) {
79             .f32 => threefry_fold_runtime_family_body_f32,
80             .i32 => threefry_fold_runtime_family_body_i32,
81             else => @compileError("Threefry fold runtime supports dtype .f32 or .i32"),
82         },
83     });
84 }
85 
86 pub const ThreefryFoldRuntimeFamilyF32 = threefryFoldRuntimeFamily(.f32);
87 pub const ThreefryFoldRuntimeFamilyI32 = threefryFoldRuntimeFamily(.i32);