lib/accy/src/kernel/library/random/fold/squares/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 SquaresFold = types.SquaresFold;
14 const zeroExtendWord = block.zeroExtendWord;
15 const squaresWord = block.squaresWord;
16 const squaresKeyValue = block.squaresKeyValue;
17 const foldInitial = common.foldInitial;
18 const foldCombine = common.foldCombine;
19 
20 fn squares_fold_runtime_body_active(inner_builder: anytype, ctx: anytype) !void {
21     const zero = try inner_builder.constantIndex(0);
22     const one = try inner_builder.constantIndex(1);
23     const initial = try foldInitial(inner_builder, ctx.dtype);
24     const key = try squaresKeyValue(
25         inner_builder,
26         ctx.args.param(.key_lo).raw(),
27         ctx.args.param(.key_hi).raw(),
28     );
29     const counter_low = try zeroExtendWord(inner_builder, try inner_builder.cast(ctx.output, .i32));
30     const folded = try inner_builder.fold(zero, ctx.samples, one, initial, .{
31         .args = ctx.args,
32         .counter_low = counter_low,
33         .key = key,
34         .dtype = ctx.dtype,
35     }, squares_fold_runtime_body_apply);
36     try ctx.args.param(.dst).store(inner_builder, folded, ctx.output);
37 }
38 
39 fn squares_fold_runtime_body_apply(fold_builder: anytype, sample: kernel.Value, current: kernel.Value, fold_ctx: anytype) !kernel.Value {
40     const half = try fold_builder.constantInt(.i64, 32);
41     const sample_wide = try zeroExtendWord(fold_builder, try fold_builder.cast(sample, .i32));
42     const counter = try fold_builder.or_(fold_ctx.counter_low, try fold_builder.shl(sample_wide, half));
43     const word = try squaresWord(fold_builder, counter, fold_ctx.key);
44     return foldCombine(fold_builder, fold_ctx.dtype, current, word);
45 }
46 
47 fn squaresFoldRuntimeBody(k: anytype, comptime dtype: DType, spec: SquaresFold, args: anytype) !void {
48     _ = spec;
49     const output = try k.globalId(.x);
50     const count = try k.castIndex(args.param(.count).raw());
51     const samples = try k.castIndex(args.param(.samples).raw());
52     const active = try k.compare(.lt, output, count);
53     try k.guardDo(active, .{ .args = args, .output = output, .samples = samples, .dtype = dtype }, squares_fold_runtime_body_active);
54 }
55 
56 fn squaresFoldFamilySchedule(instance: SquaresFold) kernel.logical.schedule.ThreadBlocks {
57     return kernel.logical.schedule.threadBlocks(.{ .x = instance.threads });
58 }
59 
60 fn squares_fold_runtime_family_body_f32(k: anytype, spec: SquaresFold, args: anytype) !void {
61     try squaresFoldRuntimeBody(k, .f32, spec, args);
62 }
63 
64 fn squares_fold_runtime_family_body_i32(k: anytype, spec: SquaresFold, args: anytype) !void {
65     try squaresFoldRuntimeBody(k, .i32, spec, args);
66 }
67 
68 fn squaresFoldRuntimeFamily(comptime dtype: DType) type {
69     return kernel.logical.Family(.{
70         .name = std.fmt.comptimePrint("accy_kernel_random_squares_fold_runtime_{s}", .{dtype.name()}),
71         .parameters = .{
72             .dst = kernel.dynamicBuffer(dtype),
73             .count = kernel.scalar(.i32),
74             .samples = kernel.scalar(.i32),
75             .key_lo = kernel.scalar(.i32),
76             .key_hi = kernel.scalar(.i32),
77         },
78         .Instance = SquaresFold,
79         .schedule = squaresFoldFamilySchedule,
80         .body = switch (dtype) {
81             .f32 => squares_fold_runtime_family_body_f32,
82             .i32 => squares_fold_runtime_family_body_i32,
83             else => @compileError("Squares fold runtime supports dtype .f32 or .i32"),
84         },
85     });
86 }
87 
88 pub const SquaresFoldRuntimeFamilyF32 = squaresFoldRuntimeFamily(.f32);
89 pub const SquaresFoldRuntimeFamilyI32 = squaresFoldRuntimeFamily(.i32);