lib/accy/src/kernel/library/elementwise.zig
daab053ee43316e1809a84551d573ddd1e5bf3d2
1 const std = @import("std");
2 const gpu = @import("gpu");
3
4 const activation_mod = @import("../../choir/root.zig").activation;
5 const entry = @import("entry.zig");
6 const kernel = @import("../root.zig");
7
8 fn unaryElementwiseSpecialization(comptime spec: entry.Vector1D, comptime operation: entry.Operation) entry.Specialization {
9 return .{
10 .dtype = .f32,
11 .operation = operation,
12 .inputs = &.{entry.shape1D(spec.axis, spec.extent)},
13 .outputs = &.{entry.shape1D(spec.axis, spec.extent)},
14 .launch = entry.launch1D(spec.extent, spec.threads),
15 .schedule = entry.threadBlocks1D(spec.axis, spec.extent, spec.threads),
16 };
17 }
18
19 fn binaryElementwiseSpecialization(comptime spec: entry.Vector1D, comptime operation: entry.Operation) entry.Specialization {
20 return .{
21 .dtype = .f32,
22 .operation = operation,
23 .inputs = &.{
24 entry.shape1D(spec.axis, spec.extent),
25 entry.shape1D(spec.axis, spec.extent),
26 },
27 .outputs = &.{entry.shape1D(spec.axis, spec.extent)},
28 .launch = entry.launch1D(spec.extent, spec.threads),
29 .schedule = entry.threadBlocks1D(spec.axis, spec.extent, spec.threads),
30 };
31 }
32
33 fn vector_add_each(inner: anytype, index: kernel.Index1D, each_args: anytype) !void {
34 const lhs = try each_args.param(.lhs).load(inner, index);
35 const rhs = try each_args.param(.rhs).load(inner, index);
36 const sum = try lhs.add(inner, rhs);
37 try each_args.param(.dst).store(inner, sum, index);
38 }
39
40 fn vectorAddProgram(comptime spec: entry.Vector1D) type {
41 const Body = struct {
42 fn run(k: anytype, args: anytype) !void {
43 _ = try k.forEach1D(spec.axis, spec.extent, args, vector_add_each);
44 }
45 };
46
47 return kernel.logical.Program(.{
48 .name = std.fmt.comptimePrint("accy_kernel_elementwise_add{}x{}_f32", .{ spec.extent, spec.threads }),
49 .parameters = .{
50 .dst = kernel.dynamicBuffer(.f32),
51 .lhs = kernel.dynamicBuffer(.f32),
52 .rhs = kernel.dynamicBuffer(.f32),
53 },
54 .body = Body.run,
55 }).withSchedule(kernel.logical.schedule.threadBlocks(.{ .x = spec.threads }));
56 }
57
58 pub fn vectorAddF32(comptime spec: entry.Vector1D) type {
59 return entry.Entry(vectorAddProgram(spec), .{
60 .target = std.fmt.comptimePrint("accy.kernel.elementwise.add{}x{}_f32", .{ spec.extent, spec.threads }),
61 .layer = .logical,
62 .category = .elementwise,
63 .specialization = binaryElementwiseSpecialization(spec, .{ .elementwise = .add }),
64 });
65 }
66
67 pub const VectorAdd8F32 = vectorAddF32(.{ .extent = 8, .threads = 4 });
68
69 fn axpy_each(inner: anytype, index: kernel.Index1D, each_args: anytype) !void {
70 const x = try each_args.param(.x).load(inner, index);
71 const y = try each_args.param(.y).load(inner, index);
72 const scaled = try x.mul(inner, each_args.param(.alpha));
73 const value = try scaled.add(inner, y);
74 try each_args.param(.dst).store(inner, value, index);
75 }
76
77 fn axpyProgram(comptime spec: entry.Vector1D) type {
78 const Body = struct {
79 fn run(k: anytype, args: anytype) !void {
80 _ = try k.forEach1D(spec.axis, spec.extent, args, axpy_each);
81 }
82 };
83
84 return kernel.logical.Program(.{
85 .name = std.fmt.comptimePrint("accy_kernel_elementwise_axpy{}x{}_f32", .{ spec.extent, spec.threads }),
86 .parameters = .{
87 .dst = kernel.dynamicBuffer(.f32),
88 .x = kernel.dynamicBuffer(.f32),
89 .y = kernel.dynamicBuffer(.f32),
90 .alpha = kernel.scalar(.f32),
91 },
92 .body = Body.run,
93 }).withSchedule(kernel.logical.schedule.threadBlocks(.{ .x = spec.threads }));
94 }
95
96 pub fn axpyF32(comptime spec: entry.Vector1D) type {
97 return entry.Entry(axpyProgram(spec), .{
98 .target = std.fmt.comptimePrint("accy.kernel.elementwise.axpy{}x{}_f32", .{ spec.extent, spec.threads }),
99 .layer = .logical,
100 .category = .elementwise,
101 .specialization = binaryElementwiseSpecialization(spec, .{ .elementwise = .axpy }),
102 });
103 }
104
105 pub const Axpy8F32 = axpyF32(.{ .extent = 8, .threads = 4 });
106
107 pub fn geluTanhApprox(inner: anytype, value: anytype) !@TypeOf(value) {
108 const square = try value.mul(inner, value);
109 const cube = try square.mul(inner, value);
110 const cubic = try cube.mul(inner, 0.044715);
111 const shifted = try value.add(inner, cubic);
112 const scaled = try shifted.mul(inner, 0.7978845608028654);
113 const smooth = try scaled.tanh(inner);
114 const gate = try smooth.add(inner, 1.0);
115 const half_value = try value.mul(inner, 0.5);
116 return half_value.mul(inner, gate);
117 }
118
119 pub fn silu(inner: anytype, value: anytype) !@TypeOf(value) {
120 const negated = try value.mul(inner, -1.0);
121 const exp_value = try negated.exp(inner);
122 const denominator = try exp_value.add(inner, 1.0);
123 return value.div(inner, denominator);
124 }
125
126 pub fn relu(inner: anytype, value: anytype) !@TypeOf(value) {
127 return value.max(inner, 0.0);
128 }
129
130 fn unaryActivationName(comptime activation: activation_mod.Kind) []const u8 {
131 return switch (activation) {
132 .gelu => "gelu",
133 .relu => "relu",
134 .silu => "silu",
135 };
136 }
137
138 fn unaryActivationOperation(comptime activation: activation_mod.Kind) entry.Operation {
139 return .{ .activation = activation };
140 }
141
142 fn unaryActivationValue(inner: anytype, comptime activation: activation_mod.Kind, value: anytype) !@TypeOf(value) {
143 return switch (activation) {
144 .gelu => geluTanhApprox(inner, value),
145 .relu => relu(inner, value),
146 .silu => silu(inner, value),
147 };
148 }
149
150 fn unary_activation_each(inner: anytype, index: kernel.Index1D, ctx: anytype) !void {
151 const value = try ctx.args.param(.src).load(inner, index);
152 const activated = try unaryActivationValue(inner, ctx.activation, value);
153 try ctx.args.param(.dst).store(inner, activated, index);
154 }
155
156 fn unaryActivationProgram(comptime spec: entry.Vector1D, comptime activation: activation_mod.Kind) type {
157 const Body = struct {
158 fn run(k: anytype, args: anytype) !void {
159 _ = try k.forEach1D(spec.axis, spec.extent, .{ .args = args, .activation = activation }, unary_activation_each);
160 }
161 };
162
163 return kernel.logical.Program(.{
164 .name = std.fmt.comptimePrint("accy_kernel_activation_{s}{}x{}_f32", .{ unaryActivationName(activation), spec.extent, spec.threads }),
165 .parameters = .{
166 .dst = kernel.dynamicBuffer(.f32),
167 .src = kernel.dynamicBuffer(.f32),
168 },
169 .body = Body.run,
170 }).withSchedule(kernel.logical.schedule.threadBlocks(.{ .x = spec.threads }));
171 }
172
173 fn unaryActivationF32(comptime spec: entry.Vector1D, comptime activation: activation_mod.Kind) type {
174 return entry.Entry(unaryActivationProgram(spec, activation), .{
175 .target = std.fmt.comptimePrint("accy.kernel.activation.{s}{}x{}_f32", .{ unaryActivationName(activation), spec.extent, spec.threads }),
176 .layer = .logical,
177 .category = .elementwise,
178 .specialization = unaryElementwiseSpecialization(spec, unaryActivationOperation(activation)),
179 });
180 }
181
182 pub fn geluF32(comptime spec: entry.Vector1D) type {
183 return unaryActivationF32(spec, .gelu);
184 }
185
186 pub fn reluF32(comptime spec: entry.Vector1D) type {
187 return unaryActivationF32(spec, .relu);
188 }
189
190 pub fn siluF32(comptime spec: entry.Vector1D) type {
191 return unaryActivationF32(spec, .silu);
192 }
193
194 pub const Gelu8F32 = geluF32(.{ .extent = 8, .threads = 4 });
195 pub const Relu8F32 = reluF32(.{ .extent = 8, .threads = 4 });
196 pub const Silu8F32 = siluF32(.{ .extent = 8, .threads = 4 });
197
198 fn geluApprox(value: f32) f32 {
199 return 0.5 * value * (1.0 + std.math.tanh(0.7978845608028654 * (value + 0.044715 * value * value * value)));
200 }
201
202 fn siluExpected(value: f32) f32 {
203 return value / (1.0 + @exp(-value));
204 }
205
206 fn reluExpected(value: f32) f32 {
207 if (value > 0.0) return value;
208 return 0.0;
209 }
210
211 fn authored_scale_each(inner: anytype, index: kernel.Index1D, each_args: anytype) !void {
212 const value = try each_args.param(.src).load(inner, index);
213 const scaled = try value.mul(inner, each_args.param(.alpha));
214 try each_args.param(.dst).store(inner, scaled, index);
215 }
216
217 fn authoredScaleProgram(comptime spec: entry.Vector1D) type {
218 const Body = struct {
219 fn run(k: anytype, args: anytype) !void {
220 _ = try k.forEach1D(spec.axis, spec.extent, spec.threads, args, authored_scale_each);
221 }
222 };
223
224 return kernel.Program(.{
225 .name = std.fmt.comptimePrint("accy_kernel_authored_elementwise_scale{}x{}_f32", .{ spec.extent, spec.threads }),
226 .parameters = .{
227 .dst = kernel.dynamicBuffer(.f32),
228 .src = kernel.dynamicBuffer(.f32),
229 .alpha = kernel.scalar(.f32),
230 },
231 .body = Body.run,
232 });
233 }
234
235 pub fn authoredScaleF32(comptime spec: entry.Vector1D) type {
236 return entry.Entry(authoredScaleProgram(spec), .{
237 .target = std.fmt.comptimePrint("accy.kernel.authored.elementwise.scale{}x{}_f32", .{ spec.extent, spec.threads }),
238 .layer = .authored,
239 .category = .elementwise,
240 .specialization = unaryElementwiseSpecialization(spec, .{ .elementwise = .scale }),
241 });
242 }
243
244 pub const AuthoredScale8F32 = authoredScaleF32(.{ .extent = 8, .threads = 8 });
245
246 test "elementwise add entry runs on CPU and records schedule" {
247 var lhs = [_]f32{ 1.0, 2.0, -3.0, 4.5, 8.0, -1.0, 0.25, 16.0 };
248 var rhs = [_]f32{ 4.0, -2.0, 6.0, 0.5, -3.0, 5.0, 0.75, -8.0 };
249 var dst = @as([8]f32, @splat(0.0));
250
251 try VectorAdd8F32.runCpu(std.testing.allocator, VectorAdd8F32.Limits.testing, &.{
252 kernel.argumentBuffer(f32, dst[0..]),
253 kernel.argumentBuffer(f32, lhs[0..]),
254 kernel.argumentBuffer(f32, rhs[0..]),
255 });
256 try std.testing.expectEqualSlices(f32, &.{ 5.0, 0.0, 3.0, 5.0, 5.0, 4.0, 1.0, 8.0 }, dst[0..]);
257
258 const launch_value = try VectorAdd8F32.launch(std.testing.allocator, VectorAdd8F32.Limits.testing);
259 try std.testing.expectEqual(@as(u32, 2), launch_value.grid[0]);
260 try std.testing.expectEqual(@as(u32, 4), launch_value.block[0]);
261 }
262
263 test "elementwise constructor creates independent shape-specialized entries" {
264 const VectorAdd16F32 = vectorAddF32(.{ .extent = 16, .threads = 8 });
265
266 try std.testing.expectEqualStrings("accy.kernel.elementwise.add8x4_f32", VectorAdd8F32.target);
267 try std.testing.expectEqualStrings("accy.kernel.elementwise.add16x8_f32", VectorAdd16F32.target);
268 try std.testing.expectEqual(@as(usize, 2), VectorAdd16F32.specialization.inputs.len);
269 try std.testing.expectEqual(@as(u64, 8), VectorAdd8F32.specialization.outputs[0].elementCount().?);
270 try std.testing.expectEqual(@as(u64, 16), VectorAdd16F32.specialization.outputs[0].elementCount().?);
271 try std.testing.expect(VectorAdd16F32.specialization.operationIs(.{ .elementwise = .add }));
272 try std.testing.expectEqual(@as(u32, 8), VectorAdd16F32.specialization.launch.?.threadgroup[0]);
273 try std.testing.expectEqual(@as(u32, 2), VectorAdd16F32.specialization.launch.?.grid[0]);
274 try std.testing.expectEqualDeep(VectorAdd16F32.specialization.launch.?, VectorAdd16F32.specialization.schedule.?.launch());
275
276 var snapshot = try VectorAdd16F32.scheduleSnapshot(std.testing.allocator, VectorAdd16F32.Limits.testing);
277 defer snapshot.deinit(std.testing.allocator);
278 try std.testing.expect(VectorAdd16F32.specialization.schedule.?.matchesSnapshot(&snapshot));
279
280 var lhs: [16]f32 = undefined;
281 var rhs: [16]f32 = undefined;
282 var dst = @as([16]f32, @splat(0.0));
283 for (0..16) |i| {
284 lhs[i] = @floatFromInt(i);
285 rhs[i] = @floatFromInt(16 - i);
286 }
287
288 try VectorAdd16F32.runCpu(std.testing.allocator, VectorAdd16F32.Limits.testing, &.{
289 kernel.argumentBuffer(f32, dst[0..]),
290 kernel.argumentBuffer(f32, lhs[0..]),
291 kernel.argumentBuffer(f32, rhs[0..]),
292 });
293 for (dst) |value| try std.testing.expectEqual(@as(f32, 16.0), value);
294 }
295
296 test "elementwise axpy entry runs on CPU" {
297 var x = [_]f32{ 1.0, 2.0, -3.0, 4.0, 0.5, -1.5, 8.0, 16.0 };
298 var y = [_]f32{ 10.0, -4.0, 1.0, 2.0, 3.0, 6.0, -8.0, 0.0 };
299 var dst = @as([8]f32, @splat(0.0));
300
301 try Axpy8F32.runCpu(std.testing.allocator, Axpy8F32.Limits.testing, &.{
302 kernel.argumentBuffer(f32, dst[0..]),
303 kernel.argumentBuffer(f32, x[0..]),
304 kernel.argumentBuffer(f32, y[0..]),
305 kernel.argumentF32(2.0),
306 });
307 try std.testing.expectEqualSlices(f32, &.{ 12.0, 0.0, -5.0, 10.0, 4.0, 3.0, 8.0, 32.0 }, dst[0..]);
308 }
309
310 test "elementwise gelu entry runs on CPU and records operation" {
311 var src = [_]f32{ -3.0, -1.0, -0.5, 0.0, 0.5, 1.0, 2.0, 4.0 };
312 var dst = @as([8]f32, @splat(0.0));
313
314 try Gelu8F32.runCpu(std.testing.allocator, Gelu8F32.Limits.testing, &.{
315 kernel.argumentBuffer(f32, dst[0..]),
316 kernel.argumentBuffer(f32, src[0..]),
317 });
318 for (src, dst) |input, actual| {
319 try std.testing.expectApproxEqAbs(geluApprox(input), actual, 0.0001);
320 }
321
322 const launch_value = try Gelu8F32.launch(std.testing.allocator, Gelu8F32.Limits.testing);
323 try std.testing.expectEqual(@as(u32, 2), launch_value.grid[0]);
324 try std.testing.expectEqual(@as(u32, 4), launch_value.block[0]);
325 try std.testing.expect(Gelu8F32.specialization.operationIs(.{ .activation = .gelu }));
326 try std.testing.expectEqual(@as(usize, 1), Gelu8F32.specialization.inputs.len);
327 try std.testing.expectEqual(@as(u64, 8), Gelu8F32.specialization.outputs[0].elementCount().?);
328 }
329
330 test "elementwise relu entry runs on CPU and records operation" {
331 var src = [_]f32{ -5.0, -2.0, -1.0, 0.0, 0.5, 1.0, 3.0, 6.0 };
332 var dst = @as([8]f32, @splat(0.0));
333
334 try Relu8F32.runCpu(std.testing.allocator, Relu8F32.Limits.testing, &.{
335 kernel.argumentBuffer(f32, dst[0..]),
336 kernel.argumentBuffer(f32, src[0..]),
337 });
338 for (src, dst) |input, actual| {
339 try std.testing.expectApproxEqAbs(reluExpected(input), actual, 0.0001);
340 }
341
342 const launch_value = try Relu8F32.launch(std.testing.allocator, Relu8F32.Limits.testing);
343 try std.testing.expectEqual(@as(u32, 2), launch_value.grid[0]);
344 try std.testing.expectEqual(@as(u32, 4), launch_value.block[0]);
345 try std.testing.expect(Relu8F32.specialization.operationIs(.{ .activation = .relu }));
346 try std.testing.expectEqual(@as(usize, 1), Relu8F32.specialization.inputs.len);
347 try std.testing.expectEqual(@as(u64, 8), Relu8F32.specialization.outputs[0].elementCount().?);
348 }
349
350 test "elementwise silu entry runs on CPU and records operation" {
351 var src = [_]f32{ -5.0, -2.0, -1.0, 0.0, 0.5, 1.0, 3.0, 6.0 };
352 var dst = @as([8]f32, @splat(0.0));
353
354 try Silu8F32.runCpu(std.testing.allocator, Silu8F32.Limits.testing, &.{
355 kernel.argumentBuffer(f32, dst[0..]),
356 kernel.argumentBuffer(f32, src[0..]),
357 });
358 for (src, dst) |input, actual| {
359 try std.testing.expectApproxEqAbs(siluExpected(input), actual, 0.0001);
360 }
361
362 const launch_value = try Silu8F32.launch(std.testing.allocator, Silu8F32.Limits.testing);
363 try std.testing.expectEqual(@as(u32, 2), launch_value.grid[0]);
364 try std.testing.expectEqual(@as(u32, 4), launch_value.block[0]);
365 try std.testing.expect(Silu8F32.specialization.operationIs(.{ .activation = .silu }));
366 try std.testing.expectEqual(@as(usize, 1), Silu8F32.specialization.inputs.len);
367 try std.testing.expectEqual(@as(u64, 8), Silu8F32.specialization.outputs[0].elementCount().?);
368 }
369
370 test "elementwise entries create registry-ready artifacts" {
371 const allocator = std.testing.allocator;
372 var state = gpu.recording.BackendState{
373 .allocator = allocator,
374 .kind = .cuda,
375 .format = .cuda_ptx,
376 };
377
378 var call_artifact = try VectorAdd8F32.createKernelCallArtifact(allocator, state.handle(), .{ .limits = VectorAdd8F32.Limits.testing });
379 defer call_artifact.deinit();
380
381 const registry = call_artifact.registry();
382 const artifact = registry.find(VectorAdd8F32.target, VectorAdd8F32.version, .cuda_ptx) orelse {
383 return error.TestExpectedKernelCallArtifact;
384 };
385 try std.testing.expectEqualStrings(VectorAdd8F32.name, artifact.entry_name);
386 try std.testing.expectEqual(@as(u32, 3), artifact.argument_count);
387 try std.testing.expectEqual(gpu.DTypeSet.init(&.{.f32}).bits, artifact.required_dtypes.bits);
388 switch (artifact.launch) {
389 .fixed => |geometry| {
390 try std.testing.expectEqual(VectorAdd8F32.specialization.launch.?.grid[0], geometry.grid[0]);
391 try std.testing.expectEqual(VectorAdd8F32.specialization.launch.?.threadgroup[0], geometry.threadgroup[0]);
392 },
393 else => return error.TestExpectedFixedLaunch,
394 }
395 }
396
397 test "authored elementwise entry preserves authored launch" {
398 var src = [_]f32{ 1.0, 2.0, -3.0, 4.0, 0.5, -1.5, 8.0, 16.0 };
399 var dst = @as([8]f32, @splat(0.0));
400
401 try AuthoredScale8F32.runCpu(std.testing.allocator, AuthoredScale8F32.Limits.testing, &.{
402 kernel.argumentBuffer(f32, dst[0..]),
403 kernel.argumentBuffer(f32, src[0..]),
404 kernel.argumentF32(3.0),
405 });
406 try std.testing.expectEqualSlices(f32, &.{ 3.0, 6.0, -9.0, 12.0, 1.5, -4.5, 24.0, 48.0 }, dst[0..]);
407
408 const launch_value = try AuthoredScale8F32.launch(std.testing.allocator, AuthoredScale8F32.Limits.testing);
409 try std.testing.expectEqual(@as(u32, 1), launch_value.grid[0]);
410 try std.testing.expectEqual(@as(u32, 8), launch_value.block[0]);
411 try std.testing.expectEqual(entry.Layer.authored, AuthoredScale8F32.layer);
412 try std.testing.expectEqual(@as(usize, 1), AuthoredScale8F32.specialization.inputs.len);
413 try std.testing.expect(AuthoredScale8F32.specialization.operationIs(.{ .elementwise = .scale }));
414 try std.testing.expectEqual(@as(u32, 8), AuthoredScale8F32.specialization.launch.?.threadgroup[0]);
415 try std.testing.expectEqualDeep(AuthoredScale8F32.specialization.launch.?, AuthoredScale8F32.specialization.schedule.?.launch());
416 }