lib/accy/src/kernel/library/random/squares/runtime.zig

daab053ee43316e1809a84551d573ddd1e5bf3d2

 1 const random = @import("../root.zig");
 2 const types = @import("types.zig");
 3 
 4 const base = random.base;
 5 const block = random.block;
 6 
 7 const std = base.std;
 8 const DType = base.DType;
 9 const kernel = base.kernel;
10 
11 const Squares = types.Squares;
12 const zeroExtendWord = block.zeroExtendWord;
13 const squaresWord = block.squaresWord;
14 const squaresKeyValue = block.squaresKeyValue;
15 const outputWord = block.outputWord;
16 
17 fn squares_runtime_body_active(inner_builder: anytype, ctx: anytype) !void {
18     const key = try squaresKeyValue(
19         inner_builder,
20         ctx.args.param(.key_lo).raw(),
21         ctx.args.param(.key_hi).raw(),
22     );
23     const counter = try zeroExtendWord(inner_builder, try inner_builder.cast(ctx.element, .i32));
24     const word = try squaresWord(inner_builder, counter, key);
25     const out = try outputWord(inner_builder, ctx.dtype, word);
26     try ctx.args.param(.dst).store(inner_builder, out, ctx.element);
27 }
28 
29 fn squaresRuntimeBody(k: anytype, comptime dtype: DType, spec: Squares, args: anytype) !void {
30     _ = spec;
31     const element = try k.globalId(.x);
32     const count = try k.castIndex(args.param(.count).raw());
33     const active = try k.compare(.lt, element, count);
34     try k.guardDo(active, .{ .args = args, .element = element, .dtype = dtype }, squares_runtime_body_active);
35 }
36 
37 fn squaresFamilySchedule(instance: Squares) kernel.logical.schedule.ThreadBlocks {
38     return kernel.logical.schedule.threadBlocks(.{ .x = instance.threads });
39 }
40 
41 fn squares_runtime_family_body_f32(k: anytype, spec: Squares, args: anytype) !void {
42     try squaresRuntimeBody(k, .f32, spec, args);
43 }
44 
45 fn squares_runtime_family_body_i32(k: anytype, spec: Squares, args: anytype) !void {
46     try squaresRuntimeBody(k, .i32, spec, args);
47 }
48 
49 fn squaresRuntimeFamily(comptime dtype: DType) type {
50     return kernel.logical.Family(.{
51         .name = std.fmt.comptimePrint("accy_kernel_random_squares_runtime_{s}", .{dtype.name()}),
52         .parameters = .{
53             .dst = kernel.dynamicBuffer(dtype),
54             .count = kernel.scalar(.i32),
55             .key_lo = kernel.scalar(.i32),
56             .key_hi = kernel.scalar(.i32),
57         },
58         .Instance = Squares,
59         .schedule = squaresFamilySchedule,
60         .body = switch (dtype) {
61             .f32 => squares_runtime_family_body_f32,
62             .i32 => squares_runtime_family_body_i32,
63             else => @compileError("Squares runtime supports dtype .f32 or .i32"),
64         },
65     });
66 }
67 
68 pub const SquaresRuntimeFamilyF32 = squaresRuntimeFamily(.f32);
69 pub const SquaresRuntimeFamilyI32 = squaresRuntimeFamily(.i32);