lib/accy/src/kernel/library/fused.zig
daab053ee43316e1809a84551d573ddd1e5bf3d2
1 const std = @import("std");
2 const gpu = @import("gpu");
3
4 const activation_mod = @import("../../choir/root.zig").activation;
5 const elementwise = @import("elementwise.zig");
6 const entry = @import("entry.zig");
7 const kernel = @import("../root.zig");
8 const linalg = @import("linalg.zig");
9
10 fn activationName(comptime activation: activation_mod.Kind) []const u8 {
11 return switch (activation) {
12 .gelu => "gelu",
13 .relu => "relu",
14 .silu => "silu",
15 };
16 }
17
18 fn activationValue(inner: anytype, comptime activation: activation_mod.Kind, value: anytype) !@TypeOf(value) {
19 return switch (activation) {
20 .gelu => elementwise.geluTanhApprox(inner, value),
21 .relu => elementwise.relu(inner, value),
22 .silu => elementwise.silu(inner, value),
23 };
24 }
25
26 fn biasActivationSpecialization(comptime spec: entry.Vector1D, comptime activation: activation_mod.Kind) entry.Specialization {
27 return .{
28 .dtype = .f32,
29 .operation = .{ .elementwise = .add },
30 .inputs = &.{
31 entry.shape1D(spec.axis, spec.extent),
32 entry.shape1D(spec.axis, spec.extent),
33 },
34 .outputs = &.{entry.shape1D(spec.axis, spec.extent)},
35 .epilogues = &.{entry.epilogue(.{ .activation = activation })},
36 .launch = entry.launch1D(spec.extent, spec.threads),
37 .schedule = entry.threadBlocks1D(spec.axis, spec.extent, spec.threads),
38 };
39 }
40
41 fn bias_activation_each(inner: anytype, index: kernel.Index1D, ctx: anytype) !void {
42 const value = try ctx.args.param(.src).load(inner, index);
43 const bias = try ctx.args.param(.bias).load(inner, index);
44 const shifted = try value.add(inner, bias);
45 const activated = try activationValue(inner, ctx.activation, shifted);
46 try ctx.args.param(.dst).store(inner, activated, index);
47 }
48
49 fn biasActivationProgram(comptime spec: entry.Vector1D, comptime activation: activation_mod.Kind) type {
50 const Body = struct {
51 fn run(k: anytype, args: anytype) !void {
52 _ = try k.forEach1D(spec.axis, spec.extent, .{ .args = args, .activation = activation }, bias_activation_each);
53 }
54 };
55
56 return kernel.logical.Program(.{
57 .name = std.fmt.comptimePrint("accy_kernel_fused_bias_{s}{}x{}_f32", .{ activationName(activation), spec.extent, spec.threads }),
58 .parameters = .{
59 .dst = kernel.dynamicBuffer(.f32),
60 .src = kernel.dynamicBuffer(.f32),
61 .bias = kernel.dynamicBuffer(.f32),
62 },
63 .body = Body.run,
64 }).withSchedule(kernel.logical.schedule.threadBlocks(.{ .x = spec.threads }));
65 }
66
67 fn biasActivationF32(comptime spec: entry.Vector1D, comptime activation: activation_mod.Kind) type {
68 return entry.Entry(biasActivationProgram(spec, activation), .{
69 .target = std.fmt.comptimePrint("accy.kernel.fused.bias_{s}{}x{}_f32", .{ activationName(activation), spec.extent, spec.threads }),
70 .layer = .logical,
71 .category = .fused,
72 .specialization = biasActivationSpecialization(spec, activation),
73 });
74 }
75
76 pub fn biasGeluF32(comptime spec: entry.Vector1D) type {
77 return biasActivationF32(spec, .gelu);
78 }
79
80 pub fn biasReluF32(comptime spec: entry.Vector1D) type {
81 return biasActivationF32(spec, .relu);
82 }
83
84 pub fn biasSiluF32(comptime spec: entry.Vector1D) type {
85 return biasActivationF32(spec, .silu);
86 }
87
88 pub const BiasGelu8F32 = biasGeluF32(.{ .extent = 8, .threads = 4 });
89 pub const BiasRelu8F32 = biasReluF32(.{ .extent = 8, .threads = 4 });
90 pub const BiasSilu8F32 = biasSiluF32(.{ .extent = 8, .threads = 4 });
91
92 const GatedActivation = activation_mod.Kind;
93
94 fn gatedActivationName(comptime activation: GatedActivation) []const u8 {
95 return switch (activation) {
96 .gelu => "geglu",
97 .relu => "reglu",
98 .silu => "swiglu",
99 };
100 }
101
102 fn gatedActivationValue(inner: anytype, comptime activation: GatedActivation, value: anytype) !@TypeOf(value) {
103 return activationValue(inner, activation, value);
104 }
105
106 fn gatedActivationSpecialization(comptime spec: entry.Vector1D, comptime activation: GatedActivation) entry.Specialization {
107 return .{
108 .dtype = .f32,
109 .operation = .{ .elementwise = .mul },
110 .inputs = &.{
111 entry.shape1D(spec.axis, spec.extent),
112 entry.shape1D(spec.axis, spec.extent),
113 },
114 .outputs = &.{entry.shape1D(spec.axis, spec.extent)},
115 .input_transforms = &.{entry.inputTransform(.{ .activation = activation }, 0, entry.shape1D(spec.axis, spec.extent))},
116 .launch = entry.launch1D(spec.extent, spec.threads),
117 .schedule = entry.threadBlocks1D(spec.axis, spec.extent, spec.threads),
118 };
119 }
120
121 fn gated_activation_each(inner: anytype, index: kernel.Index1D, ctx: anytype) !void {
122 const gate = try ctx.args.param(.gate).load(inner, index);
123 const value = try ctx.args.param(.value).load(inner, index);
124 const activated = try gatedActivationValue(inner, ctx.activation, gate);
125 const output = try activated.mul(inner, value);
126 try ctx.args.param(.dst).store(inner, output, index);
127 }
128
129 fn gatedActivationProgram(comptime spec: entry.Vector1D, comptime activation: GatedActivation) type {
130 const Body = struct {
131 fn run(k: anytype, args: anytype) !void {
132 _ = try k.forEach1D(spec.axis, spec.extent, .{ .args = args, .activation = activation }, gated_activation_each);
133 }
134 };
135
136 return kernel.logical.Program(.{
137 .name = std.fmt.comptimePrint("accy_kernel_fused_{s}{}x{}_f32", .{ gatedActivationName(activation), spec.extent, spec.threads }),
138 .parameters = .{
139 .dst = kernel.dynamicBuffer(.f32),
140 .gate = kernel.dynamicBuffer(.f32),
141 .value = kernel.dynamicBuffer(.f32),
142 },
143 .body = Body.run,
144 }).withSchedule(kernel.logical.schedule.threadBlocks(.{ .x = spec.threads }));
145 }
146
147 fn gatedActivationF32(comptime spec: entry.Vector1D, comptime activation: GatedActivation) type {
148 return entry.Entry(gatedActivationProgram(spec, activation), .{
149 .target = std.fmt.comptimePrint("accy.kernel.fused.{s}{}x{}_f32", .{ gatedActivationName(activation), spec.extent, spec.threads }),
150 .layer = .logical,
151 .category = .fused,
152 .specialization = gatedActivationSpecialization(spec, activation),
153 });
154 }
155
156 pub fn geGluF32(comptime spec: entry.Vector1D) type {
157 return gatedActivationF32(spec, .gelu);
158 }
159
160 pub fn reGluF32(comptime spec: entry.Vector1D) type {
161 return gatedActivationF32(spec, .relu);
162 }
163
164 pub fn swiGluF32(comptime spec: entry.Vector1D) type {
165 return gatedActivationF32(spec, .silu);
166 }
167
168 pub const GeGlu8F32 = geGluF32(.{ .extent = 8, .threads = 4 });
169 pub const ReGlu8F32 = reGluF32(.{ .extent = 8, .threads = 4 });
170 pub const SwiGlu8F32 = swiGluF32(.{ .extent = 8, .threads = 4 });
171
172 fn matrixProductBiasActivationSpecialization(comptime spec: linalg.MatrixProduct, comptime activation: activation_mod.Kind) entry.Specialization {
173 return .{
174 .dtype = .f32,
175 .operation = .{ .linalg = .matrix_product },
176 .equation = "mk,kn,n->mn",
177 .inputs = &.{
178 entry.shape2D(spec.row_axis, spec.m, spec.reduction_axis, spec.k),
179 entry.shape2D(spec.reduction_axis, spec.k, spec.col_axis, spec.n),
180 entry.shape1D(spec.col_axis, spec.n),
181 },
182 .outputs = &.{entry.shape2D(spec.row_axis, spec.m, spec.col_axis, spec.n)},
183 .reductions = &.{entry.reduction("dot", .dot_product, entry.shape1D(spec.reduction_axis, spec.k))},
184 .epilogues = &.{
185 entry.inputEpilogue(.bias_add, 2, entry.shape1D(spec.col_axis, spec.n)),
186 entry.epilogue(.{ .activation = activation }),
187 },
188 .launch = entry.launch2D(spec.n, spec.m, spec.threads.x, spec.threads.y),
189 .schedule = entry.threadBlocks2D(spec.col_axis, spec.n, spec.row_axis, spec.m, spec.threads.x, spec.threads.y),
190 };
191 }
192
193 fn matrix_product_bias_activation_each(inner: anytype, index: kernel.Index2D, ctx: anytype) !void {
194 const sum = try linalg.matrixProductCellSum(inner, ctx.spec, ctx.args.param(.lhs), ctx.args.param(.rhs), index.y.index, index.x.index);
195 const bias = try ctx.args.param(.bias).load(inner, index.x.index);
196 const shifted = try bias.add(inner, sum);
197 const activated = try activationValue(inner, ctx.activation, shifted);
198 const out_index = try linalg.matrixProductOutputIndex(inner, ctx.spec, index.y.index, index.x.index);
199 try ctx.args.param(.dst).store(inner, activated, out_index);
200 }
201
202 fn matrixProductBiasActivationProgram(comptime spec: linalg.MatrixProduct, comptime activation: activation_mod.Kind) type {
203 const Body = struct {
204 fn run(k: anytype, args: anytype) !void {
205 _ = try k.forEach2D(.{
206 .x = kernel.logical.axis(spec.col_axis, spec.n),
207 .y = kernel.logical.axis(spec.row_axis, spec.m),
208 }, .{ .spec = spec, .activation = activation, .args = args }, matrix_product_bias_activation_each);
209 }
210 };
211
212 return kernel.logical.Program(.{
213 .name = std.fmt.comptimePrint(
214 "accy_kernel_fused_matmul_bias_{s}{}x{}x{}_{}x{}_f32",
215 .{ activationName(activation), spec.m, spec.n, spec.k, spec.threads.x, spec.threads.y },
216 ),
217 .parameters = .{
218 .dst = kernel.dynamicBuffer(.f32),
219 .lhs = kernel.dynamicBuffer(.f32),
220 .rhs = kernel.dynamicBuffer(.f32),
221 .bias = kernel.dynamicBuffer(.f32),
222 },
223 .body = Body.run,
224 }).withSchedule(kernel.logical.schedule.threadBlocks(.{
225 .x = spec.threads.x,
226 .y = spec.threads.y,
227 }));
228 }
229
230 fn matrixProductBiasActivationF32(comptime spec: linalg.MatrixProduct, comptime activation: activation_mod.Kind) type {
231 return entry.Entry(matrixProductBiasActivationProgram(spec, activation), .{
232 .target = std.fmt.comptimePrint(
233 "accy.kernel.fused.matmul_bias_{s}{}x{}x{}_{}x{}_f32",
234 .{ activationName(activation), spec.m, spec.n, spec.k, spec.threads.x, spec.threads.y },
235 ),
236 .layer = .logical,
237 .category = .fused,
238 .specialization = matrixProductBiasActivationSpecialization(spec, activation),
239 });
240 }
241
242 pub fn matrixProductBiasGeluF32(comptime spec: linalg.MatrixProduct) type {
243 return matrixProductBiasActivationF32(spec, .gelu);
244 }
245
246 pub fn matrixProductBiasReluF32(comptime spec: linalg.MatrixProduct) type {
247 return matrixProductBiasActivationF32(spec, .relu);
248 }
249
250 pub fn matrixProductBiasSiluF32(comptime spec: linalg.MatrixProduct) type {
251 return matrixProductBiasActivationF32(spec, .silu);
252 }
253
254 pub const MatrixProductBiasGelu2x3x4F32 = matrixProductBiasGeluF32(.{
255 .m = 2,
256 .n = 3,
257 .k = 4,
258 .threads = .{ .x = 2, .y = 2 },
259 });
260 pub const MatrixProductBiasGelu2x3x4ThreadBlocks1x2F32 = matrixProductBiasGeluF32(.{
261 .m = 2,
262 .n = 3,
263 .k = 4,
264 .threads = .{ .x = 1, .y = 2 },
265 });
266 pub const MatrixProductBiasRelu2x3x4F32 = matrixProductBiasReluF32(.{
267 .m = 2,
268 .n = 3,
269 .k = 4,
270 .threads = .{ .x = 2, .y = 2 },
271 });
272 pub const MatrixProductBiasSilu2x3x4F32 = matrixProductBiasSiluF32(.{
273 .m = 2,
274 .n = 3,
275 .k = 4,
276 .threads = .{ .x = 2, .y = 2 },
277 });
278
279 fn matrixVectorProductBiasActivationSpecialization(comptime spec: linalg.MatrixVectorProduct, comptime activation: activation_mod.Kind) entry.Specialization {
280 return .{
281 .dtype = .f32,
282 .operation = .{ .linalg = .matrix_vector_product },
283 .equation = "mk,k,m->m",
284 .inputs = &.{
285 entry.shape2D(spec.row_axis, spec.m, spec.reduction_axis, spec.k),
286 entry.shape1D(spec.reduction_axis, spec.k),
287 entry.shape1D(spec.row_axis, spec.m),
288 },
289 .outputs = &.{entry.shape1D(spec.row_axis, spec.m)},
290 .reductions = &.{entry.reduction("dot", .dot_product, entry.shape1D(spec.reduction_axis, spec.k))},
291 .epilogues = &.{
292 entry.inputEpilogue(.bias_add, 2, entry.shape1D(spec.row_axis, spec.m)),
293 entry.epilogue(.{ .activation = activation }),
294 },
295 .launch = entry.launch1D(spec.m, spec.threads),
296 .schedule = entry.threadBlocks1D(spec.row_axis, spec.m, spec.threads),
297 };
298 }
299
300 fn matrix_vector_product_bias_activation_each(inner: anytype, index: kernel.Index1D, ctx: anytype) !void {
301 const sum = try linalg.matrixVectorProductRowSum(inner, ctx.spec, ctx.args.param(.matrix), ctx.args.param(.vector), index.index);
302 const bias = try ctx.args.param(.bias).load(inner, index);
303 const shifted = try bias.add(inner, sum);
304 const activated = try activationValue(inner, ctx.activation, shifted);
305 try ctx.args.param(.dst).store(inner, activated, index);
306 }
307
308 fn matrixVectorProductBiasActivationProgram(comptime spec: linalg.MatrixVectorProduct, comptime activation: activation_mod.Kind) type {
309 const Body = struct {
310 fn run(k: anytype, args: anytype) !void {
311 _ = try k.forEach1D(spec.row_axis, spec.m, .{ .spec = spec, .activation = activation, .args = args }, matrix_vector_product_bias_activation_each);
312 }
313 };
314
315 return kernel.logical.Program(.{
316 .name = std.fmt.comptimePrint(
317 "accy_kernel_fused_matvec_bias_{s}{}x{}_{}x_f32",
318 .{ activationName(activation), spec.m, spec.k, spec.threads },
319 ),
320 .parameters = .{
321 .dst = kernel.dynamicBuffer(.f32),
322 .matrix = kernel.dynamicBuffer(.f32),
323 .vector = kernel.dynamicBuffer(.f32),
324 .bias = kernel.dynamicBuffer(.f32),
325 },
326 .body = Body.run,
327 }).withSchedule(kernel.logical.schedule.threadBlocks(.{ .x = spec.threads }));
328 }
329
330 fn matrixVectorProductBiasActivationF32(comptime spec: linalg.MatrixVectorProduct, comptime activation: activation_mod.Kind) type {
331 return entry.Entry(matrixVectorProductBiasActivationProgram(spec, activation), .{
332 .target = std.fmt.comptimePrint(
333 "accy.kernel.fused.matvec_bias_{s}{}x{}_{}x_f32",
334 .{ activationName(activation), spec.m, spec.k, spec.threads },
335 ),
336 .layer = .logical,
337 .category = .fused,
338 .specialization = matrixVectorProductBiasActivationSpecialization(spec, activation),
339 });
340 }
341
342 pub fn matrixVectorProductBiasGeluF32(comptime spec: linalg.MatrixVectorProduct) type {
343 return matrixVectorProductBiasActivationF32(spec, .gelu);
344 }
345
346 pub fn matrixVectorProductBiasReluF32(comptime spec: linalg.MatrixVectorProduct) type {
347 return matrixVectorProductBiasActivationF32(spec, .relu);
348 }
349
350 pub fn matrixVectorProductBiasSiluF32(comptime spec: linalg.MatrixVectorProduct) type {
351 return matrixVectorProductBiasActivationF32(spec, .silu);
352 }
353
354 pub const MatrixVectorProductBiasGelu4x8F32 = matrixVectorProductBiasGeluF32(.{
355 .m = 4,
356 .k = 8,
357 .threads = 4,
358 });
359 pub const MatrixVectorProductBiasRelu4x8F32 = matrixVectorProductBiasReluF32(.{
360 .m = 4,
361 .k = 8,
362 .threads = 4,
363 });
364 pub const MatrixVectorProductBiasSilu4x8F32 = matrixVectorProductBiasSiluF32(.{
365 .m = 4,
366 .k = 8,
367 .threads = 4,
368 });
369
370 fn expectedActivation(comptime activation: activation_mod.Kind, value: f32) f32 {
371 return switch (activation) {
372 .gelu => 0.5 * value * (1.0 + std.math.tanh(0.7978845608028654 * (value + 0.044715 * value * value * value))),
373 .relu => if (value > 0.0) value else 0.0,
374 .silu => value / (1.0 + @exp(-value)),
375 };
376 }
377
378 fn expectedBiasActivation(comptime activation: activation_mod.Kind, value: f32, bias: f32) f32 {
379 return expectedActivation(activation, value + bias);
380 }
381
382 fn expectedSwiGlu(gate: f32, value: f32) f32 {
383 return gate / (1.0 + @exp(-gate)) * value;
384 }
385
386 fn expectedGeGlu(gate: f32, value: f32) f32 {
387 return expectedActivation(.gelu, gate) * value;
388 }
389
390 fn expectedReGlu(gate: f32, value: f32) f32 {
391 if (gate > 0.0) return gate * value;
392 return 0.0;
393 }
394
395 test "fused bias gelu entry runs on CPU and records fused operation" {
396 var src = [_]f32{ -3.0, -1.0, -0.5, 0.0, 0.5, 1.0, 2.0, 4.0 };
397 var bias = [_]f32{ 0.25, -0.5, 0.5, 1.0, -0.25, 0.75, -1.0, 0.0 };
398 var dst = @as([8]f32, @splat(0.0));
399
400 try BiasGelu8F32.runCpu(std.testing.allocator, BiasGelu8F32.Limits.testing, &.{
401 kernel.argumentBuffer(f32, dst[0..]),
402 kernel.argumentBuffer(f32, src[0..]),
403 kernel.argumentBuffer(f32, bias[0..]),
404 });
405 for (src, bias, dst) |input, bias_value, actual| {
406 try std.testing.expectApproxEqAbs(expectedBiasActivation(.gelu, input, bias_value), actual, 0.0001);
407 }
408
409 const launch_value = try BiasGelu8F32.launch(std.testing.allocator, BiasGelu8F32.Limits.testing);
410 try std.testing.expectEqual(@as(u32, 2), launch_value.grid[0]);
411 try std.testing.expectEqual(@as(u32, 4), launch_value.block[0]);
412 try std.testing.expectEqual(entry.Category.fused, BiasGelu8F32.category);
413 try std.testing.expect(BiasGelu8F32.specialization.operationIs(.{ .elementwise = .add }));
414 try std.testing.expect(BiasGelu8F32.specialization.epilogueMatches(0, .{ .operator = .{ .activation = .gelu } }));
415 try std.testing.expectEqual(@as(usize, 2), BiasGelu8F32.specialization.inputs.len);
416 try std.testing.expectEqual(@as(u64, 8), BiasGelu8F32.specialization.outputs[0].elementCount().?);
417 try std.testing.expectEqualDeep(BiasGelu8F32.specialization.launch.?, BiasGelu8F32.specialization.schedule.?.launch());
418
419 var snapshot = try BiasGelu8F32.scheduleSnapshot(std.testing.allocator, BiasGelu8F32.Limits.testing);
420 defer snapshot.deinit(std.testing.allocator);
421 try std.testing.expect(BiasGelu8F32.specialization.schedule.?.matchesSnapshot(&snapshot));
422 }
423
424 test "fused bias relu and silu entries run on CPU" {
425 var src = [_]f32{ -3.0, -1.0, -0.5, 0.0, 0.5, 1.0, 2.0, 4.0 };
426 var bias = [_]f32{ 0.25, -0.5, 0.5, 1.0, -0.25, 0.75, -1.0, 0.0 };
427 var relu_dst = @as([8]f32, @splat(0.0));
428 var silu_dst = @as([8]f32, @splat(0.0));
429
430 try BiasRelu8F32.runCpu(std.testing.allocator, BiasRelu8F32.Limits.testing, &.{
431 kernel.argumentBuffer(f32, relu_dst[0..]),
432 kernel.argumentBuffer(f32, src[0..]),
433 kernel.argumentBuffer(f32, bias[0..]),
434 });
435 try BiasSilu8F32.runCpu(std.testing.allocator, BiasSilu8F32.Limits.testing, &.{
436 kernel.argumentBuffer(f32, silu_dst[0..]),
437 kernel.argumentBuffer(f32, src[0..]),
438 kernel.argumentBuffer(f32, bias[0..]),
439 });
440 for (src, bias, relu_dst, silu_dst) |input, bias_value, relu_actual, silu_actual| {
441 try std.testing.expectApproxEqAbs(expectedBiasActivation(.relu, input, bias_value), relu_actual, 0.0001);
442 try std.testing.expectApproxEqAbs(expectedBiasActivation(.silu, input, bias_value), silu_actual, 0.0001);
443 }
444
445 try std.testing.expect(BiasRelu8F32.specialization.epilogueMatches(0, .{ .operator = .{ .activation = .relu } }));
446 try std.testing.expect(BiasSilu8F32.specialization.epilogueMatches(0, .{ .operator = .{ .activation = .silu } }));
447 }
448
449 test "fused swiglu entry runs on CPU and records input transform" {
450 var gate = [_]f32{ -5.0, -2.0, -1.0, 0.0, 0.5, 1.0, 3.0, 6.0 };
451 var value = [_]f32{ 2.0, -1.0, 0.5, 3.0, -2.0, 4.0, 0.25, -0.5 };
452 var dst = @as([8]f32, @splat(0.0));
453
454 try SwiGlu8F32.runCpu(std.testing.allocator, SwiGlu8F32.Limits.testing, &.{
455 kernel.argumentBuffer(f32, dst[0..]),
456 kernel.argumentBuffer(f32, gate[0..]),
457 kernel.argumentBuffer(f32, value[0..]),
458 });
459 for (gate, value, dst) |gate_value, input_value, actual| {
460 try std.testing.expectApproxEqAbs(expectedSwiGlu(gate_value, input_value), actual, 0.0001);
461 }
462
463 const launch_value = try SwiGlu8F32.launch(std.testing.allocator, SwiGlu8F32.Limits.testing);
464 try std.testing.expectEqual(@as(u32, 2), launch_value.grid[0]);
465 try std.testing.expectEqual(@as(u32, 4), launch_value.block[0]);
466 try std.testing.expectEqual(entry.Category.fused, SwiGlu8F32.category);
467 try std.testing.expect(SwiGlu8F32.specialization.operationIs(.{ .elementwise = .mul }));
468 try std.testing.expect(SwiGlu8F32.specialization.inputTransformMatches(0, .{
469 .operator = .{ .activation = .silu },
470 .input_index = 0,
471 .extents = &.{8},
472 }));
473 try std.testing.expectEqual(@as(usize, 2), SwiGlu8F32.specialization.inputs.len);
474 try std.testing.expectEqual(@as(u64, 8), SwiGlu8F32.specialization.outputs[0].elementCount().?);
475 try std.testing.expectEqualDeep(SwiGlu8F32.specialization.launch.?, SwiGlu8F32.specialization.schedule.?.launch());
476
477 var snapshot = try SwiGlu8F32.scheduleSnapshot(std.testing.allocator, SwiGlu8F32.Limits.testing);
478 defer snapshot.deinit(std.testing.allocator);
479 try std.testing.expect(SwiGlu8F32.specialization.schedule.?.matchesSnapshot(&snapshot));
480 }
481
482 test "fused geglu entry runs on CPU and records input transform" {
483 var gate = [_]f32{ -5.0, -2.0, -1.0, 0.0, 0.5, 1.0, 3.0, 6.0 };
484 var value = [_]f32{ 2.0, -1.0, 0.5, 3.0, -2.0, 4.0, 0.25, -0.5 };
485 var dst = @as([8]f32, @splat(0.0));
486
487 try GeGlu8F32.runCpu(std.testing.allocator, GeGlu8F32.Limits.testing, &.{
488 kernel.argumentBuffer(f32, dst[0..]),
489 kernel.argumentBuffer(f32, gate[0..]),
490 kernel.argumentBuffer(f32, value[0..]),
491 });
492 for (gate, value, dst) |gate_value, input_value, actual| {
493 try std.testing.expectApproxEqAbs(expectedGeGlu(gate_value, input_value), actual, 0.0001);
494 }
495
496 const launch_value = try GeGlu8F32.launch(std.testing.allocator, GeGlu8F32.Limits.testing);
497 try std.testing.expectEqual(@as(u32, 2), launch_value.grid[0]);
498 try std.testing.expectEqual(@as(u32, 4), launch_value.block[0]);
499 try std.testing.expectEqual(entry.Category.fused, GeGlu8F32.category);
500 try std.testing.expect(GeGlu8F32.specialization.operationIs(.{ .elementwise = .mul }));
501 try std.testing.expect(GeGlu8F32.specialization.inputTransformMatches(0, .{
502 .operator = .{ .activation = .gelu },
503 .input_index = 0,
504 .extents = &.{8},
505 }));
506 try std.testing.expectEqual(@as(usize, 2), GeGlu8F32.specialization.inputs.len);
507 try std.testing.expectEqual(@as(u64, 8), GeGlu8F32.specialization.outputs[0].elementCount().?);
508 try std.testing.expectEqualDeep(GeGlu8F32.specialization.launch.?, GeGlu8F32.specialization.schedule.?.launch());
509
510 var snapshot = try GeGlu8F32.scheduleSnapshot(std.testing.allocator, GeGlu8F32.Limits.testing);
511 defer snapshot.deinit(std.testing.allocator);
512 try std.testing.expect(GeGlu8F32.specialization.schedule.?.matchesSnapshot(&snapshot));
513 }
514
515 test "fused reglu entry runs on CPU and records input transform" {
516 var gate = [_]f32{ -5.0, -2.0, -1.0, 0.0, 0.5, 1.0, 3.0, 6.0 };
517 var value = [_]f32{ 2.0, -1.0, 0.5, 3.0, -2.0, 4.0, 0.25, -0.5 };
518 var dst = @as([8]f32, @splat(0.0));
519
520 try ReGlu8F32.runCpu(std.testing.allocator, ReGlu8F32.Limits.testing, &.{
521 kernel.argumentBuffer(f32, dst[0..]),
522 kernel.argumentBuffer(f32, gate[0..]),
523 kernel.argumentBuffer(f32, value[0..]),
524 });
525 for (gate, value, dst) |gate_value, input_value, actual| {
526 try std.testing.expectApproxEqAbs(expectedReGlu(gate_value, input_value), actual, 0.0001);
527 }
528
529 const launch_value = try ReGlu8F32.launch(std.testing.allocator, ReGlu8F32.Limits.testing);
530 try std.testing.expectEqual(@as(u32, 2), launch_value.grid[0]);
531 try std.testing.expectEqual(@as(u32, 4), launch_value.block[0]);
532 try std.testing.expectEqual(entry.Category.fused, ReGlu8F32.category);
533 try std.testing.expect(ReGlu8F32.specialization.operationIs(.{ .elementwise = .mul }));
534 try std.testing.expect(ReGlu8F32.specialization.inputTransformMatches(0, .{
535 .operator = .{ .activation = .relu },
536 .input_index = 0,
537 .extents = &.{8},
538 }));
539 try std.testing.expectEqual(@as(usize, 2), ReGlu8F32.specialization.inputs.len);
540 try std.testing.expectEqual(@as(u64, 8), ReGlu8F32.specialization.outputs[0].elementCount().?);
541 try std.testing.expectEqualDeep(ReGlu8F32.specialization.launch.?, ReGlu8F32.specialization.schedule.?.launch());
542
543 var snapshot = try ReGlu8F32.scheduleSnapshot(std.testing.allocator, ReGlu8F32.Limits.testing);
544 defer snapshot.deinit(std.testing.allocator);
545 try std.testing.expect(ReGlu8F32.specialization.schedule.?.matchesSnapshot(&snapshot));
546 }
547
548 test "fused matrix product bias gelu entry runs on CPU and records epilogue operation" {
549 var lhs = [_]f32{
550 1.0, 2.0, 3.0, 4.0,
551 5.0, 6.0, 7.0, 8.0,
552 };
553 var rhs = [_]f32{
554 1.0, 0.0, 2.0,
555 0.0, 1.0, 3.0,
556 1.0, 1.0, 0.0,
557 2.0, 0.0, 1.0,
558 };
559 var bias = [_]f32{ 0.25, -1.0, 0.5 };
560 var dst = @as([6]f32, @splat(0.0));
561
562 try MatrixProductBiasGelu2x3x4F32.runCpu(std.testing.allocator, MatrixProductBiasGelu2x3x4F32.Limits.testing, &.{
563 kernel.argumentBuffer(f32, dst[0..]),
564 kernel.argumentBuffer(f32, lhs[0..]),
565 kernel.argumentBuffer(f32, rhs[0..]),
566 kernel.argumentBuffer(f32, bias[0..]),
567 });
568 const dot = [_]f32{ 12.0, 5.0, 12.0, 28.0, 13.0, 36.0 };
569 for (dot, dst, 0..) |sum, actual, index| {
570 try std.testing.expectApproxEqAbs(expectedBiasActivation(.gelu, sum, bias[index % 3]), actual, 0.0001);
571 }
572
573 const launch_value = try MatrixProductBiasGelu2x3x4F32.launch(std.testing.allocator, MatrixProductBiasGelu2x3x4F32.Limits.testing);
574 try std.testing.expectEqual(@as(u32, 2), launch_value.grid[0]);
575 try std.testing.expectEqual(@as(u32, 1), launch_value.grid[1]);
576 try std.testing.expectEqual(@as(u32, 2), launch_value.block[0]);
577 try std.testing.expectEqual(@as(u32, 2), launch_value.block[1]);
578 try std.testing.expectEqual(entry.Category.fused, MatrixProductBiasGelu2x3x4F32.category);
579 try std.testing.expect(MatrixProductBiasGelu2x3x4F32.specialization.operationIs(.{ .linalg = .matrix_product }));
580 try std.testing.expectEqualStrings("mk,kn,n->mn", MatrixProductBiasGelu2x3x4F32.specialization.equation.?);
581 try std.testing.expectEqual(@as(usize, 3), MatrixProductBiasGelu2x3x4F32.specialization.inputs.len);
582 try std.testing.expectEqual(@as(u64, 3), MatrixProductBiasGelu2x3x4F32.specialization.inputs[2].elementCount().?);
583 try std.testing.expectEqual(@as(u64, 6), MatrixProductBiasGelu2x3x4F32.specialization.outputs[0].elementCount().?);
584 try std.testing.expect(MatrixProductBiasGelu2x3x4F32.specialization.epilogueMatches(0, .{
585 .operator = .bias_add,
586 .input_index = 2,
587 .extents = &.{3},
588 }));
589 try std.testing.expect(MatrixProductBiasGelu2x3x4F32.specialization.epilogueMatches(1, .{ .operator = .{ .activation = .gelu } }));
590 try std.testing.expectEqualDeep(MatrixProductBiasGelu2x3x4F32.specialization.launch.?, MatrixProductBiasGelu2x3x4F32.specialization.schedule.?.launch());
591
592 var snapshot = try MatrixProductBiasGelu2x3x4F32.scheduleSnapshot(std.testing.allocator, MatrixProductBiasGelu2x3x4F32.Limits.testing);
593 defer snapshot.deinit(std.testing.allocator);
594 try std.testing.expect(MatrixProductBiasGelu2x3x4F32.specialization.schedule.?.matchesSnapshot(&snapshot));
595 }
596
597 test "fused matrix product bias relu and silu entries run on CPU" {
598 var lhs = [_]f32{
599 1.0, 2.0, 3.0, 4.0,
600 5.0, 6.0, 7.0, 8.0,
601 };
602 var rhs = [_]f32{
603 1.0, 0.0, 2.0,
604 0.0, 1.0, 3.0,
605 1.0, 1.0, 0.0,
606 2.0, 0.0, 1.0,
607 };
608 var bias = [_]f32{ 0.25, -20.0, 0.5 };
609 var relu_dst = @as([6]f32, @splat(0.0));
610 var silu_dst = @as([6]f32, @splat(0.0));
611
612 try MatrixProductBiasRelu2x3x4F32.runCpu(std.testing.allocator, MatrixProductBiasRelu2x3x4F32.Limits.testing, &.{
613 kernel.argumentBuffer(f32, relu_dst[0..]),
614 kernel.argumentBuffer(f32, lhs[0..]),
615 kernel.argumentBuffer(f32, rhs[0..]),
616 kernel.argumentBuffer(f32, bias[0..]),
617 });
618 try MatrixProductBiasSilu2x3x4F32.runCpu(std.testing.allocator, MatrixProductBiasSilu2x3x4F32.Limits.testing, &.{
619 kernel.argumentBuffer(f32, silu_dst[0..]),
620 kernel.argumentBuffer(f32, lhs[0..]),
621 kernel.argumentBuffer(f32, rhs[0..]),
622 kernel.argumentBuffer(f32, bias[0..]),
623 });
624
625 const dot = [_]f32{ 12.0, 5.0, 12.0, 28.0, 13.0, 36.0 };
626 for (dot, relu_dst, silu_dst, 0..) |sum, relu_actual, silu_actual, index| {
627 try std.testing.expectApproxEqAbs(expectedBiasActivation(.relu, sum, bias[index % 3]), relu_actual, 0.0001);
628 try std.testing.expectApproxEqAbs(expectedBiasActivation(.silu, sum, bias[index % 3]), silu_actual, 0.0001);
629 }
630
631 try std.testing.expect(MatrixProductBiasRelu2x3x4F32.specialization.epilogueMatches(1, .{ .operator = .{ .activation = .relu } }));
632 try std.testing.expect(MatrixProductBiasSilu2x3x4F32.specialization.epilogueMatches(1, .{ .operator = .{ .activation = .silu } }));
633 try std.testing.expectEqualDeep(MatrixProductBiasRelu2x3x4F32.specialization.launch.?, MatrixProductBiasRelu2x3x4F32.specialization.schedule.?.launch());
634 try std.testing.expectEqualDeep(MatrixProductBiasSilu2x3x4F32.specialization.launch.?, MatrixProductBiasSilu2x3x4F32.specialization.schedule.?.launch());
635 }
636
637 test "fused matrix vector product bias gelu entry runs on CPU and records epilogue operation" {
638 var matrix = [_]f32{
639 1.0, 2.0, 3.0, 4.0, 5.0, 6.0, 7.0, 8.0,
640 2.0, 0.0, -2.0, 0.0, 1.0, 0.0, -1.0, 0.0,
641 0.5, 1.0, 1.5, 2.0, 2.5, 3.0, 3.5, 4.0,
642 -1.0, -2.0, 3.0, 4.0, -5.0, 6.0, 7.0, -8.0,
643 };
644 var vector = [_]f32{ 1.0, 0.5, -1.0, 2.0, 0.25, -0.5, 1.5, -2.0 };
645 var bias = [_]f32{ 0.25, -1.0, 0.5, -20.0 };
646 var dst = @as([4]f32, @splat(0.0));
647
648 try MatrixVectorProductBiasGelu4x8F32.runCpu(std.testing.allocator, MatrixVectorProductBiasGelu4x8F32.Limits.testing, &.{
649 kernel.argumentBuffer(f32, dst[0..]),
650 kernel.argumentBuffer(f32, matrix[0..]),
651 kernel.argumentBuffer(f32, vector[0..]),
652 kernel.argumentBuffer(f32, bias[0..]),
653 });
654 const dot = [_]f32{ -0.25, 2.75, -0.125, 25.25 };
655 for (dot, bias, dst) |sum, bias_value, actual| {
656 try std.testing.expectApproxEqAbs(expectedBiasActivation(.gelu, sum, bias_value), actual, 0.0001);
657 }
658
659 const launch_value = try MatrixVectorProductBiasGelu4x8F32.launch(std.testing.allocator, MatrixVectorProductBiasGelu4x8F32.Limits.testing);
660 try std.testing.expectEqual(@as(u32, 1), launch_value.grid[0]);
661 try std.testing.expectEqual(@as(u32, 4), launch_value.block[0]);
662 try std.testing.expectEqual(entry.Category.fused, MatrixVectorProductBiasGelu4x8F32.category);
663 try std.testing.expect(MatrixVectorProductBiasGelu4x8F32.specialization.operationIs(.{ .linalg = .matrix_vector_product }));
664 try std.testing.expectEqualStrings("mk,k,m->m", MatrixVectorProductBiasGelu4x8F32.specialization.equation.?);
665 try std.testing.expectEqual(@as(usize, 3), MatrixVectorProductBiasGelu4x8F32.specialization.inputs.len);
666 try std.testing.expectEqual(@as(u64, 4), MatrixVectorProductBiasGelu4x8F32.specialization.inputs[2].elementCount().?);
667 try std.testing.expectEqual(@as(u64, 4), MatrixVectorProductBiasGelu4x8F32.specialization.outputs[0].elementCount().?);
668 try std.testing.expect(MatrixVectorProductBiasGelu4x8F32.specialization.epilogueMatches(0, .{
669 .operator = .bias_add,
670 .input_index = 2,
671 .extents = &.{4},
672 }));
673 try std.testing.expect(MatrixVectorProductBiasGelu4x8F32.specialization.epilogueMatches(1, .{ .operator = .{ .activation = .gelu } }));
674 try std.testing.expectEqualDeep(MatrixVectorProductBiasGelu4x8F32.specialization.launch.?, MatrixVectorProductBiasGelu4x8F32.specialization.schedule.?.launch());
675
676 var snapshot = try MatrixVectorProductBiasGelu4x8F32.scheduleSnapshot(std.testing.allocator, MatrixVectorProductBiasGelu4x8F32.Limits.testing);
677 defer snapshot.deinit(std.testing.allocator);
678 try std.testing.expect(MatrixVectorProductBiasGelu4x8F32.specialization.schedule.?.matchesSnapshot(&snapshot));
679 }
680
681 test "fused matrix vector product bias relu and silu entries run on CPU" {
682 var matrix = [_]f32{
683 1.0, 2.0, 3.0, 4.0, 5.0, 6.0, 7.0, 8.0,
684 2.0, 0.0, -2.0, 0.0, 1.0, 0.0, -1.0, 0.0,
685 0.5, 1.0, 1.5, 2.0, 2.5, 3.0, 3.5, 4.0,
686 -1.0, -2.0, 3.0, 4.0, -5.0, 6.0, 7.0, -8.0,
687 };
688 var vector = [_]f32{ 1.0, 0.5, -1.0, 2.0, 0.25, -0.5, 1.5, -2.0 };
689 var bias = [_]f32{ 0.25, -1.0, 0.5, -20.0 };
690 var relu_dst = @as([4]f32, @splat(0.0));
691 var silu_dst = @as([4]f32, @splat(0.0));
692
693 try MatrixVectorProductBiasRelu4x8F32.runCpu(std.testing.allocator, MatrixVectorProductBiasRelu4x8F32.Limits.testing, &.{
694 kernel.argumentBuffer(f32, relu_dst[0..]),
695 kernel.argumentBuffer(f32, matrix[0..]),
696 kernel.argumentBuffer(f32, vector[0..]),
697 kernel.argumentBuffer(f32, bias[0..]),
698 });
699 try MatrixVectorProductBiasSilu4x8F32.runCpu(std.testing.allocator, MatrixVectorProductBiasSilu4x8F32.Limits.testing, &.{
700 kernel.argumentBuffer(f32, silu_dst[0..]),
701 kernel.argumentBuffer(f32, matrix[0..]),
702 kernel.argumentBuffer(f32, vector[0..]),
703 kernel.argumentBuffer(f32, bias[0..]),
704 });
705
706 const dot = [_]f32{ -0.25, 2.75, -0.125, 25.25 };
707 for (dot, bias, relu_dst, silu_dst) |sum, bias_value, relu_actual, silu_actual| {
708 try std.testing.expectApproxEqAbs(expectedBiasActivation(.relu, sum, bias_value), relu_actual, 0.0001);
709 try std.testing.expectApproxEqAbs(expectedBiasActivation(.silu, sum, bias_value), silu_actual, 0.0001);
710 }
711
712 try std.testing.expect(MatrixVectorProductBiasRelu4x8F32.specialization.epilogueMatches(1, .{ .operator = .{ .activation = .relu } }));
713 try std.testing.expect(MatrixVectorProductBiasSilu4x8F32.specialization.epilogueMatches(1, .{ .operator = .{ .activation = .silu } }));
714 try std.testing.expectEqualDeep(MatrixVectorProductBiasRelu4x8F32.specialization.launch.?, MatrixVectorProductBiasRelu4x8F32.specialization.schedule.?.launch());
715 try std.testing.expectEqualDeep(MatrixVectorProductBiasSilu4x8F32.specialization.launch.?, MatrixVectorProductBiasSilu4x8F32.specialization.schedule.?.launch());
716 }
717
718 test "fused bias gelu constructor creates independent shape-specialized entries" {
719 const BiasGelu16F32 = biasGeluF32(.{ .extent = 16, .threads = 8 });
720
721 try std.testing.expectEqualStrings("accy.kernel.fused.bias_gelu8x4_f32", BiasGelu8F32.target);
722 try std.testing.expectEqualStrings("accy.kernel.fused.bias_gelu16x8_f32", BiasGelu16F32.target);
723 try std.testing.expectEqual(@as(u64, 8), BiasGelu8F32.specialization.outputs[0].elementCount().?);
724 try std.testing.expectEqual(@as(u64, 16), BiasGelu16F32.specialization.outputs[0].elementCount().?);
725 try std.testing.expect(BiasGelu16F32.specialization.operationIs(.{ .elementwise = .add }));
726 try std.testing.expect(BiasGelu16F32.specialization.epilogueMatches(0, .{ .operator = .{ .activation = .gelu } }));
727 try std.testing.expectEqual(@as(u32, 8), BiasGelu16F32.specialization.launch.?.threadgroup[0]);
728 try std.testing.expectEqual(@as(u32, 2), BiasGelu16F32.specialization.launch.?.grid[0]);
729 }
730
731 test "fused swiglu constructor creates independent shape-specialized entries" {
732 const SwiGlu16F32 = swiGluF32(.{ .extent = 16, .threads = 8 });
733
734 try std.testing.expectEqualStrings("accy.kernel.fused.swiglu8x4_f32", SwiGlu8F32.target);
735 try std.testing.expectEqualStrings("accy.kernel.fused.swiglu16x8_f32", SwiGlu16F32.target);
736 try std.testing.expectEqual(@as(u64, 8), SwiGlu8F32.specialization.outputs[0].elementCount().?);
737 try std.testing.expectEqual(@as(u64, 16), SwiGlu16F32.specialization.outputs[0].elementCount().?);
738 try std.testing.expect(SwiGlu16F32.specialization.operationIs(.{ .elementwise = .mul }));
739 try std.testing.expect(SwiGlu16F32.specialization.inputTransformMatches(0, .{
740 .operator = .{ .activation = .silu },
741 .input_index = 0,
742 .extents = &.{16},
743 }));
744 try std.testing.expectEqual(@as(u32, 8), SwiGlu16F32.specialization.launch.?.threadgroup[0]);
745 try std.testing.expectEqual(@as(u32, 2), SwiGlu16F32.specialization.launch.?.grid[0]);
746 }
747
748 test "fused geglu constructor creates independent shape-specialized entries" {
749 const GeGlu16F32 = geGluF32(.{ .extent = 16, .threads = 8 });
750
751 try std.testing.expectEqualStrings("accy.kernel.fused.geglu8x4_f32", GeGlu8F32.target);
752 try std.testing.expectEqualStrings("accy.kernel.fused.geglu16x8_f32", GeGlu16F32.target);
753 try std.testing.expectEqual(@as(u64, 8), GeGlu8F32.specialization.outputs[0].elementCount().?);
754 try std.testing.expectEqual(@as(u64, 16), GeGlu16F32.specialization.outputs[0].elementCount().?);
755 try std.testing.expect(GeGlu16F32.specialization.operationIs(.{ .elementwise = .mul }));
756 try std.testing.expect(GeGlu16F32.specialization.inputTransformMatches(0, .{
757 .operator = .{ .activation = .gelu },
758 .input_index = 0,
759 .extents = &.{16},
760 }));
761 try std.testing.expectEqual(@as(u32, 8), GeGlu16F32.specialization.launch.?.threadgroup[0]);
762 try std.testing.expectEqual(@as(u32, 2), GeGlu16F32.specialization.launch.?.grid[0]);
763 }
764
765 test "fused reglu constructor creates independent shape-specialized entries" {
766 const ReGlu16F32 = reGluF32(.{ .extent = 16, .threads = 8 });
767
768 try std.testing.expectEqualStrings("accy.kernel.fused.reglu8x4_f32", ReGlu8F32.target);
769 try std.testing.expectEqualStrings("accy.kernel.fused.reglu16x8_f32", ReGlu16F32.target);
770 try std.testing.expectEqual(@as(u64, 8), ReGlu8F32.specialization.outputs[0].elementCount().?);
771 try std.testing.expectEqual(@as(u64, 16), ReGlu16F32.specialization.outputs[0].elementCount().?);
772 try std.testing.expect(ReGlu16F32.specialization.operationIs(.{ .elementwise = .mul }));
773 try std.testing.expect(ReGlu16F32.specialization.inputTransformMatches(0, .{
774 .operator = .{ .activation = .relu },
775 .input_index = 0,
776 .extents = &.{16},
777 }));
778 try std.testing.expectEqual(@as(u32, 8), ReGlu16F32.specialization.launch.?.threadgroup[0]);
779 try std.testing.expectEqual(@as(u32, 2), ReGlu16F32.specialization.launch.?.grid[0]);
780 }
781
782 test "fused matrix product bias gelu constructor creates independent shape-specialized entries" {
783 const MatrixProductBiasGelu4x5x6F32 = matrixProductBiasGeluF32(.{
784 .m = 4,
785 .n = 5,
786 .k = 6,
787 .threads = .{ .x = 4, .y = 2 },
788 });
789
790 try std.testing.expectEqualStrings("accy.kernel.fused.matmul_bias_gelu2x3x4_2x2_f32", MatrixProductBiasGelu2x3x4F32.target);
791 try std.testing.expectEqualStrings("accy.kernel.fused.matmul_bias_gelu2x3x4_1x2_f32", MatrixProductBiasGelu2x3x4ThreadBlocks1x2F32.target);
792 try std.testing.expectEqualStrings("accy.kernel.fused.matmul_bias_gelu4x5x6_4x2_f32", MatrixProductBiasGelu4x5x6F32.target);
793 try std.testing.expect(MatrixProductBiasGelu2x3x4ThreadBlocks1x2F32.specialization.operationIs(.{ .linalg = .matrix_product }));
794 try std.testing.expectEqual(@as(u32, 1), MatrixProductBiasGelu2x3x4ThreadBlocks1x2F32.specialization.launch.?.threadgroup[0]);
795 try std.testing.expectEqual(@as(u32, 2), MatrixProductBiasGelu2x3x4ThreadBlocks1x2F32.specialization.launch.?.threadgroup[1]);
796 try std.testing.expectEqualDeep(MatrixProductBiasGelu2x3x4ThreadBlocks1x2F32.specialization.launch.?, MatrixProductBiasGelu2x3x4ThreadBlocks1x2F32.specialization.schedule.?.launch());
797 try std.testing.expect(MatrixProductBiasGelu4x5x6F32.specialization.operationIs(.{ .linalg = .matrix_product }));
798 try std.testing.expectEqual(@as(usize, 3), MatrixProductBiasGelu4x5x6F32.specialization.inputs.len);
799 try std.testing.expectEqual(@as(u64, 24), MatrixProductBiasGelu4x5x6F32.specialization.inputs[0].elementCount().?);
800 try std.testing.expectEqual(@as(u64, 30), MatrixProductBiasGelu4x5x6F32.specialization.inputs[1].elementCount().?);
801 try std.testing.expectEqual(@as(u64, 5), MatrixProductBiasGelu4x5x6F32.specialization.inputs[2].elementCount().?);
802 try std.testing.expectEqual(@as(u64, 20), MatrixProductBiasGelu4x5x6F32.specialization.outputs[0].elementCount().?);
803 try std.testing.expectEqual(entry.ReductionOperator.dot_product, MatrixProductBiasGelu4x5x6F32.specialization.reductions[0].operator);
804 try std.testing.expectEqual(@as(u64, 6), MatrixProductBiasGelu4x5x6F32.specialization.reductions[0].shape.elementCount().?);
805 try std.testing.expect(MatrixProductBiasGelu4x5x6F32.specialization.epilogueMatches(0, .{
806 .operator = .bias_add,
807 .input_index = 2,
808 .extents = &.{5},
809 }));
810 try std.testing.expect(MatrixProductBiasGelu4x5x6F32.specialization.epilogueMatches(1, .{ .operator = .{ .activation = .gelu } }));
811 try std.testing.expectEqualDeep(MatrixProductBiasGelu4x5x6F32.specialization.launch.?, MatrixProductBiasGelu4x5x6F32.specialization.schedule.?.launch());
812 }
813
814 test "fused matrix vector product bias gelu constructor creates independent shape-specialized entries" {
815 const MatrixVectorProductBiasGelu5x6F32 = matrixVectorProductBiasGeluF32(.{
816 .m = 5,
817 .k = 6,
818 .threads = 4,
819 });
820
821 try std.testing.expectEqualStrings("accy.kernel.fused.matvec_bias_gelu4x8_4x_f32", MatrixVectorProductBiasGelu4x8F32.target);
822 try std.testing.expectEqualStrings("accy.kernel.fused.matvec_bias_gelu5x6_4x_f32", MatrixVectorProductBiasGelu5x6F32.target);
823 try std.testing.expect(MatrixVectorProductBiasGelu5x6F32.specialization.operationIs(.{ .linalg = .matrix_vector_product }));
824 try std.testing.expectEqual(@as(usize, 3), MatrixVectorProductBiasGelu5x6F32.specialization.inputs.len);
825 try std.testing.expectEqual(@as(u64, 30), MatrixVectorProductBiasGelu5x6F32.specialization.inputs[0].elementCount().?);
826 try std.testing.expectEqual(@as(u64, 6), MatrixVectorProductBiasGelu5x6F32.specialization.inputs[1].elementCount().?);
827 try std.testing.expectEqual(@as(u64, 5), MatrixVectorProductBiasGelu5x6F32.specialization.inputs[2].elementCount().?);
828 try std.testing.expectEqual(@as(u64, 5), MatrixVectorProductBiasGelu5x6F32.specialization.outputs[0].elementCount().?);
829 try std.testing.expectEqual(entry.ReductionOperator.dot_product, MatrixVectorProductBiasGelu5x6F32.specialization.reductions[0].operator);
830 try std.testing.expectEqual(@as(u64, 6), MatrixVectorProductBiasGelu5x6F32.specialization.reductions[0].shape.elementCount().?);
831 try std.testing.expect(MatrixVectorProductBiasGelu5x6F32.specialization.epilogueMatches(0, .{
832 .operator = .bias_add,
833 .input_index = 2,
834 .extents = &.{5},
835 }));
836 try std.testing.expect(MatrixVectorProductBiasGelu5x6F32.specialization.epilogueMatches(1, .{ .operator = .{ .activation = .gelu } }));
837 try std.testing.expectEqualDeep(MatrixVectorProductBiasGelu5x6F32.specialization.launch.?, MatrixVectorProductBiasGelu5x6F32.specialization.schedule.?.launch());
838 }
839
840 test "fused bias gelu entry creates registry-ready artifact" {
841 const allocator = std.testing.allocator;
842 var state = gpu.recording.BackendState{
843 .allocator = allocator,
844 .kind = .cuda,
845 .format = .cuda_ptx,
846 };
847
848 var call_artifact = try BiasGelu8F32.createKernelCallArtifact(allocator, state.handle(), .{ .limits = BiasGelu8F32.Limits.testing });
849 defer call_artifact.deinit();
850
851 const artifact = call_artifact.registry().find(BiasGelu8F32.target, BiasGelu8F32.version, .cuda_ptx) orelse {
852 return error.TestExpectedKernelCallArtifact;
853 };
854 try std.testing.expectEqualStrings(BiasGelu8F32.name, artifact.entry_name);
855 try std.testing.expectEqual(@as(u32, 3), artifact.argument_count);
856 try std.testing.expectEqual(gpu.DTypeSet.init(&.{.f32}).bits, artifact.required_dtypes.bits);
857 switch (artifact.launch) {
858 .fixed => |geometry| {
859 try std.testing.expectEqual(BiasGelu8F32.specialization.launch.?.grid[0], geometry.grid[0]);
860 try std.testing.expectEqual(BiasGelu8F32.specialization.launch.?.threadgroup[0], geometry.threadgroup[0]);
861 },
862 else => return error.TestExpectedFixedLaunch,
863 }
864 }
865
866 test "fused swiglu entry creates registry-ready artifact" {
867 const allocator = std.testing.allocator;
868 var state = gpu.recording.BackendState{
869 .allocator = allocator,
870 .kind = .cuda,
871 .format = .cuda_ptx,
872 };
873
874 var call_artifact = try SwiGlu8F32.createKernelCallArtifact(allocator, state.handle(), .{ .limits = SwiGlu8F32.Limits.testing });
875 defer call_artifact.deinit();
876
877 const artifact = call_artifact.registry().find(SwiGlu8F32.target, SwiGlu8F32.version, .cuda_ptx) orelse {
878 return error.TestExpectedKernelCallArtifact;
879 };
880 try std.testing.expectEqualStrings(SwiGlu8F32.name, artifact.entry_name);
881 try std.testing.expectEqual(@as(u32, 3), artifact.argument_count);
882 try std.testing.expectEqual(gpu.DTypeSet.init(&.{.f32}).bits, artifact.required_dtypes.bits);
883 switch (artifact.launch) {
884 .fixed => |geometry| {
885 try std.testing.expectEqual(SwiGlu8F32.specialization.launch.?.grid[0], geometry.grid[0]);
886 try std.testing.expectEqual(SwiGlu8F32.specialization.launch.?.threadgroup[0], geometry.threadgroup[0]);
887 },
888 else => return error.TestExpectedFixedLaunch,
889 }
890 }
891
892 test "fused geglu entry creates registry-ready artifact" {
893 const allocator = std.testing.allocator;
894 var state = gpu.recording.BackendState{
895 .allocator = allocator,
896 .kind = .cuda,
897 .format = .cuda_ptx,
898 };
899
900 var call_artifact = try GeGlu8F32.createKernelCallArtifact(allocator, state.handle(), .{ .limits = GeGlu8F32.Limits.testing });
901 defer call_artifact.deinit();
902
903 const artifact = call_artifact.registry().find(GeGlu8F32.target, GeGlu8F32.version, .cuda_ptx) orelse {
904 return error.TestExpectedKernelCallArtifact;
905 };
906 try std.testing.expectEqualStrings(GeGlu8F32.name, artifact.entry_name);
907 try std.testing.expectEqual(@as(u32, 3), artifact.argument_count);
908 try std.testing.expectEqual(gpu.DTypeSet.init(&.{.f32}).bits, artifact.required_dtypes.bits);
909 switch (artifact.launch) {
910 .fixed => |geometry| {
911 try std.testing.expectEqual(GeGlu8F32.specialization.launch.?.grid[0], geometry.grid[0]);
912 try std.testing.expectEqual(GeGlu8F32.specialization.launch.?.threadgroup[0], geometry.threadgroup[0]);
913 },
914 else => return error.TestExpectedFixedLaunch,
915 }
916 }
917
918 test "fused reglu entry creates registry-ready artifact" {
919 const allocator = std.testing.allocator;
920 var state = gpu.recording.BackendState{
921 .allocator = allocator,
922 .kind = .cuda,
923 .format = .cuda_ptx,
924 };
925
926 var call_artifact = try ReGlu8F32.createKernelCallArtifact(allocator, state.handle(), .{ .limits = ReGlu8F32.Limits.testing });
927 defer call_artifact.deinit();
928
929 const artifact = call_artifact.registry().find(ReGlu8F32.target, ReGlu8F32.version, .cuda_ptx) orelse {
930 return error.TestExpectedKernelCallArtifact;
931 };
932 try std.testing.expectEqualStrings(ReGlu8F32.name, artifact.entry_name);
933 try std.testing.expectEqual(@as(u32, 3), artifact.argument_count);
934 try std.testing.expectEqual(gpu.DTypeSet.init(&.{.f32}).bits, artifact.required_dtypes.bits);
935 switch (artifact.launch) {
936 .fixed => |geometry| {
937 try std.testing.expectEqual(ReGlu8F32.specialization.launch.?.grid[0], geometry.grid[0]);
938 try std.testing.expectEqual(ReGlu8F32.specialization.launch.?.threadgroup[0], geometry.threadgroup[0]);
939 },
940 else => return error.TestExpectedFixedLaunch,
941 }
942 }
943
944 test "fused matrix product bias gelu entry creates registry-ready artifact" {
945 const allocator = std.testing.allocator;
946 var state = gpu.recording.BackendState{
947 .allocator = allocator,
948 .kind = .cuda,
949 .format = .cuda_ptx,
950 };
951
952 var call_artifact = try MatrixProductBiasGelu2x3x4F32.createKernelCallArtifact(allocator, state.handle(), .{ .limits = MatrixProductBiasGelu2x3x4F32.Limits.testing });
953 defer call_artifact.deinit();
954
955 const artifact = call_artifact.registry().find(MatrixProductBiasGelu2x3x4F32.target, MatrixProductBiasGelu2x3x4F32.version, .cuda_ptx) orelse {
956 return error.TestExpectedKernelCallArtifact;
957 };
958 try std.testing.expectEqualStrings(MatrixProductBiasGelu2x3x4F32.name, artifact.entry_name);
959 try std.testing.expectEqual(@as(u32, 4), artifact.argument_count);
960 try std.testing.expectEqual(gpu.DTypeSet.init(&.{.f32}).bits, artifact.required_dtypes.bits);
961 switch (artifact.launch) {
962 .fixed => |geometry| {
963 try std.testing.expectEqual(MatrixProductBiasGelu2x3x4F32.specialization.launch.?.grid[0], geometry.grid[0]);
964 try std.testing.expectEqual(MatrixProductBiasGelu2x3x4F32.specialization.launch.?.grid[1], geometry.grid[1]);
965 try std.testing.expectEqual(MatrixProductBiasGelu2x3x4F32.specialization.launch.?.threadgroup[0], geometry.threadgroup[0]);
966 try std.testing.expectEqual(MatrixProductBiasGelu2x3x4F32.specialization.launch.?.threadgroup[1], geometry.threadgroup[1]);
967 },
968 else => return error.TestExpectedFixedLaunch,
969 }
970 }
971
972 test "fused matrix vector product bias gelu entry creates registry-ready artifact" {
973 const allocator = std.testing.allocator;
974 var state = gpu.recording.BackendState{
975 .allocator = allocator,
976 .kind = .cuda,
977 .format = .cuda_ptx,
978 };
979
980 var call_artifact = try MatrixVectorProductBiasGelu4x8F32.createKernelCallArtifact(allocator, state.handle(), .{ .limits = MatrixVectorProductBiasGelu4x8F32.Limits.testing });
981 defer call_artifact.deinit();
982
983 const artifact = call_artifact.registry().find(MatrixVectorProductBiasGelu4x8F32.target, MatrixVectorProductBiasGelu4x8F32.version, .cuda_ptx) orelse {
984 return error.TestExpectedKernelCallArtifact;
985 };
986 try std.testing.expectEqualStrings(MatrixVectorProductBiasGelu4x8F32.name, artifact.entry_name);
987 try std.testing.expectEqual(@as(u32, 4), artifact.argument_count);
988 try std.testing.expectEqual(gpu.DTypeSet.init(&.{.f32}).bits, artifact.required_dtypes.bits);
989 switch (artifact.launch) {
990 .fixed => |geometry| {
991 try std.testing.expectEqual(MatrixVectorProductBiasGelu4x8F32.specialization.launch.?.grid[0], geometry.grid[0]);
992 try std.testing.expectEqual(MatrixVectorProductBiasGelu4x8F32.specialization.launch.?.threadgroup[0], geometry.threadgroup[0]);
993 },
994 else => return error.TestExpectedFixedLaunch,
995 }
996 }