lib/accy/src/kernel/library/random/key/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 const PhiloxKeySplit = types.PhiloxKeySplit;
11 const PhiloxKeyUniform = types.PhiloxKeyUniform;
12 const PhiloxKeyCounterUniform = types.PhiloxKeyCounterUniform;
13 const philoxWords = block.philoxWords;
14 const outputWord = block.outputWord;
15 const storeGuarded = block.storeGuarded;
16
17 fn philox_key_split_body_active(inner_builder: anytype, ctx: anytype) !void {
18 try ctx.args.param(.dst).store(inner_builder, ctx.out, ctx.element);
19 }
20
21 fn philoxKeySplitBody(k: anytype, spec: PhiloxKeySplit, args: anytype) !void {
22 const element = try k.globalId(.x);
23 const count = try k.castIndex(args.param(.count).raw());
24 const seed = try args.param(.key).load(k, try k.constantIndex(0));
25 const words = try philoxWords(
26 k,
27 spec.rounds,
28 try k.cast(element, .i32),
29 try k.constantInt(.i32, 0),
30 seed.lo,
31 seed.hi,
32 );
33 const out = kernel.KeyValue{ .lo = words[0], .hi = words[1] };
34 const active = try k.compare(.lt, element, count);
35 try k.guardDo(active, .{ .args = args, .out = out, .element = element }, philox_key_split_body_active);
36 }
37
38 fn philoxKeyUniformBody(k: anytype, comptime dtype: DType, spec: PhiloxKeyUniform, args: anytype) !void {
39 const element = try k.globalId(.x);
40 const count = try k.castIndex(args.param(.count).raw());
41 const key = try args.param(.keys).load(k, element);
42 const words = try philoxWords(
43 k,
44 spec.rounds,
45 try k.constantInt(.i32, 0),
46 try k.constantInt(.i32, 0),
47 key.lo,
48 key.hi,
49 );
50 const out = try outputWord(k, dtype, words[0]);
51 try storeGuarded(k, args, out, element, count);
52 }
53
54 fn philoxKeyCounterUniformBody(k: anytype, comptime dtype: DType, spec: PhiloxKeyCounterUniform, args: anytype) !void {
55 const element = try k.globalId(.x);
56 const count = try k.castIndex(args.param(.count).raw());
57 const seed = try args.param(.key).load(k, try k.constantIndex(0));
58 const counter_value = try args.param(.counter).load(k, try k.constantIndex(0));
59 const words = try philoxWords(
60 k,
61 spec.rounds,
62 try k.cast(element, .i32),
63 counter_value.raw(),
64 seed.lo,
65 seed.hi,
66 );
67 const out = try outputWord(k, dtype, words[0]);
68 try storeGuarded(k, args, out, element, count);
69 }
70
71 fn philoxKeySplitSchedule(instance: PhiloxKeySplit) kernel.logical.schedule.ThreadBlocks {
72 return kernel.logical.schedule.threadBlocks(.{ .x = instance.threads });
73 }
74
75 fn philoxKeyUniformSchedule(instance: PhiloxKeyUniform) kernel.logical.schedule.ThreadBlocks {
76 return kernel.logical.schedule.threadBlocks(.{ .x = instance.threads });
77 }
78
79 fn philoxKeyCounterUniformSchedule(instance: PhiloxKeyCounterUniform) kernel.logical.schedule.ThreadBlocks {
80 return kernel.logical.schedule.threadBlocks(.{ .x = instance.threads });
81 }
82
83 fn philox_key_split_runtime_family_body(k: anytype, spec: PhiloxKeySplit, args: anytype) !void {
84 try philoxKeySplitBody(k, spec, args);
85 }
86
87 pub const PhiloxKeySplitRuntimeFamily = kernel.logical.Family(.{
88 .name = "accy_kernel_random_philox_key_split_runtime_key",
89 .parameters = .{
90 .dst = kernel.dynamicBuffer(.key),
91 .key = kernel.dynamicBuffer(.key),
92 .count = kernel.scalar(.i32),
93 },
94 .Instance = PhiloxKeySplit,
95 .schedule = philoxKeySplitSchedule,
96 .body = philox_key_split_runtime_family_body,
97 });
98
99 fn philox_key_uniform_runtime_family_body_f32(k: anytype, spec: PhiloxKeyUniform, args: anytype) !void {
100 try philoxKeyUniformBody(k, .f32, spec, args);
101 }
102
103 fn philox_key_uniform_runtime_family_body_i32(k: anytype, spec: PhiloxKeyUniform, args: anytype) !void {
104 try philoxKeyUniformBody(k, .i32, spec, args);
105 }
106
107 fn philoxKeyUniformRuntimeFamily(comptime dtype: DType) type {
108 return kernel.logical.Family(.{
109 .name = std.fmt.comptimePrint("accy_kernel_random_philox_key_uniform_runtime_{s}", .{dtype.name()}),
110 .parameters = .{
111 .dst = kernel.dynamicBuffer(dtype),
112 .keys = kernel.dynamicBuffer(.key),
113 .count = kernel.scalar(.i32),
114 },
115 .Instance = PhiloxKeyUniform,
116 .schedule = philoxKeyUniformSchedule,
117 .body = switch (dtype) {
118 .f32 => philox_key_uniform_runtime_family_body_f32,
119 .i32 => philox_key_uniform_runtime_family_body_i32,
120 else => @compileError("Philox key uniform runtime supports dtype .f32 or .i32"),
121 },
122 });
123 }
124
125 pub const PhiloxKeyUniformRuntimeFamilyF32 = philoxKeyUniformRuntimeFamily(.f32);
126 pub const PhiloxKeyUniformRuntimeFamilyI32 = philoxKeyUniformRuntimeFamily(.i32);
127
128 fn philox_key_counter_uniform_runtime_family_body_f32(k: anytype, spec: PhiloxKeyCounterUniform, args: anytype) !void {
129 try philoxKeyCounterUniformBody(k, .f32, spec, args);
130 }
131
132 fn philox_key_counter_uniform_runtime_family_body_i32(k: anytype, spec: PhiloxKeyCounterUniform, args: anytype) !void {
133 try philoxKeyCounterUniformBody(k, .i32, spec, args);
134 }
135
136 fn philoxKeyCounterUniformRuntimeFamily(comptime dtype: DType) type {
137 return kernel.logical.Family(.{
138 .name = std.fmt.comptimePrint("accy_kernel_random_philox_key_counter_uniform_runtime_{s}", .{dtype.name()}),
139 .parameters = .{
140 .dst = kernel.dynamicBuffer(dtype),
141 .key = kernel.dynamicBuffer(.key),
142 .counter = kernel.dynamicBuffer(.i32),
143 .count = kernel.scalar(.i32),
144 },
145 .Instance = PhiloxKeyCounterUniform,
146 .schedule = philoxKeyCounterUniformSchedule,
147 .body = switch (dtype) {
148 .f32 => philox_key_counter_uniform_runtime_family_body_f32,
149 .i32 => philox_key_counter_uniform_runtime_family_body_i32,
150 else => @compileError("Philox key counter runtime supports dtype .f32 or .i32"),
151 },
152 });
153 }
154
155 pub const PhiloxKeyCounterUniformRuntimeFamilyF32 = philoxKeyCounterUniformRuntimeFamily(.f32);
156 pub const PhiloxKeyCounterUniformRuntimeFamilyI32 = philoxKeyCounterUniformRuntimeFamily(.i32);