lib/accy/src/kernel/library/normalization.zig
daab053ee43316e1809a84551d573ddd1e5bf3d2
1 const std = @import("std");
2 const gpu = @import("gpu");
3
4 const entry = @import("entry.zig");
5 const kernel = @import("../root.zig");
6
7 const RowNormalizationParameterization = entry.RowNormalizationParameterization;
8
9 pub const Threads2D = struct {
10 x: u32 = 4,
11 y: u32 = 2,
12 };
13
14 pub const Row = struct {
15 rows: u64,
16 cols: u64,
17 threads: Threads2D = .{},
18 row_axis: []const u8 = "row",
19 col_axis: []const u8 = "col",
20 };
21
22 const RowDistribution = enum {
23 softmax,
24 log_softmax,
25 };
26
27 fn rowDistributionName(comptime kind: RowDistribution) []const u8 {
28 return switch (kind) {
29 .softmax => "softmax",
30 .log_softmax => "log_softmax",
31 };
32 }
33
34 fn rowDistributionOperator(comptime kind: RowDistribution) entry.RowNormalizationOperator {
35 return switch (kind) {
36 .softmax => .softmax,
37 .log_softmax => .log_softmax,
38 };
39 }
40
41 fn indexUpper(comptime extent: u64) i64 {
42 if (extent > @as(u64, @intCast(std.math.maxInt(i64)))) {
43 @compileError("kernel library normalization extent overflows index range");
44 }
45 return @intCast(extent);
46 }
47
48 fn floatExtent(comptime extent: u64) f64 {
49 return @floatFromInt(extent);
50 }
51
52 fn rowDomain(comptime spec: Row) kernel.logical.Domain2D {
53 return .{
54 .x = kernel.logical.axis(spec.col_axis, spec.cols),
55 .y = kernel.logical.axis(spec.row_axis, spec.rows),
56 };
57 }
58
59 fn rowMatrixShape(comptime spec: Row) entry.Shape {
60 return entry.shape2D(spec.row_axis, spec.rows, spec.col_axis, spec.cols);
61 }
62
63 fn rowShape(comptime spec: Row) entry.Shape {
64 return entry.shape1D(spec.row_axis, spec.rows);
65 }
66
67 fn columnShape(comptime spec: Row) entry.Shape {
68 return entry.shape1D(spec.col_axis, spec.cols);
69 }
70
71 fn rowLaunch(comptime spec: Row) entry.Launch {
72 return entry.launch2D(spec.cols, spec.rows, spec.threads.x, spec.threads.y);
73 }
74
75 fn rowOffset(k: anytype, index: kernel.Index2D, comptime spec: Row) !kernel.Value {
76 const cols_stride = try k.constantIndex(indexUpper(spec.cols));
77 return k.mul(index.y.index, cols_stride);
78 }
79
80 fn rowItemIndex(k: anytype, index: kernel.Index2D, offset: kernel.Value) !kernel.Value {
81 return k.add(offset, index.x.index);
82 }
83
84 fn foldColumnsFrom(
85 k: anytype,
86 comptime spec: Row,
87 comptime lower: i64,
88 initial: anytype,
89 context: anytype,
90 comptime body: anytype,
91 ) !@TypeOf(initial) {
92 return k.foldRange(lower, indexUpper(spec.cols), 1, initial, context, body);
93 }
94
95 fn foldColumns(k: anytype, comptime spec: Row, initial: anytype, context: anytype, comptime body: anytype) !@TypeOf(initial) {
96 return foldColumnsFrom(k, spec, 0, initial, context, body);
97 }
98
99 fn rowDistributionSpecialization(comptime spec: Row, comptime kind: RowDistribution) entry.Specialization {
100 return .{
101 .dtype = .f32,
102 .operation = .{ .row_normalization = rowDistributionOperator(kind) },
103 .inputs = &.{rowMatrixShape(spec)},
104 .outputs = &.{rowMatrixShape(spec)},
105 .reductions = &.{
106 entry.reduction("row_max", .maximum, columnShape(spec)),
107 entry.dependentReduction("row_exp_sum", .sum_exp_shifted, columnShape(spec), &.{"row_max"}),
108 },
109 .reduction_reuse = &.{
110 entry.reductionReuse("row_max", rowShape(spec)),
111 entry.reductionReuse("row_exp_sum", rowShape(spec)),
112 },
113 .launch = rowLaunch(spec),
114 .schedule = entry.threadBlocks2D(spec.col_axis, spec.cols, spec.row_axis, spec.rows, spec.threads.x, spec.threads.y),
115 };
116 }
117
118 fn rowWeightedSpecialization(comptime spec: Row) entry.Specialization {
119 return .{
120 .dtype = .f32,
121 .operation = .{ .row_normalization = .{ .rmsnorm = .scale } },
122 .inputs = &.{
123 rowMatrixShape(spec),
124 columnShape(spec),
125 },
126 .outputs = &.{rowMatrixShape(spec)},
127 .reductions = &.{entry.reduction("row_sum_squares", .sum_squares, columnShape(spec))},
128 .reduction_reuse = &.{entry.reductionReuse("row_sum_squares", rowShape(spec))},
129 .launch = rowLaunch(spec),
130 .schedule = entry.threadBlocks2D(spec.col_axis, spec.cols, spec.row_axis, spec.rows, spec.threads.x, spec.threads.y),
131 };
132 }
133
134 fn rowResidualRmsNormSpecialization(comptime spec: Row) entry.Specialization {
135 return .{
136 .dtype = .f32,
137 .operation = .{ .row_normalization = .{ .rmsnorm = .scale } },
138 .inputs = &.{
139 rowMatrixShape(spec),
140 rowMatrixShape(spec),
141 columnShape(spec),
142 },
143 .outputs = &.{rowMatrixShape(spec)},
144 .reductions = &.{entry.reduction("row_sum_squares", .sum_squares, columnShape(spec))},
145 .reduction_reuse = &.{entry.reductionReuse("row_sum_squares", rowShape(spec))},
146 .input_transforms = &.{entry.inputTransform(.residual_add, 1, rowMatrixShape(spec))},
147 .launch = rowLaunch(spec),
148 .schedule = entry.threadBlocks2D(spec.col_axis, spec.cols, spec.row_axis, spec.rows, spec.threads.x, spec.threads.y),
149 };
150 }
151
152 fn rowLayerNormName(comptime parameterization: RowNormalizationParameterization) []const u8 {
153 return switch (parameterization) {
154 .none => "layernorm",
155 .scale => "layernorm_scale",
156 .scale_bias => "layernorm_affine",
157 };
158 }
159
160 fn rowLayerNormInputs(comptime spec: Row, comptime parameterization: RowNormalizationParameterization) []const entry.Shape {
161 return switch (parameterization) {
162 .none => &.{rowMatrixShape(spec)},
163 .scale => &.{
164 rowMatrixShape(spec),
165 columnShape(spec),
166 },
167 .scale_bias => &.{
168 rowMatrixShape(spec),
169 columnShape(spec),
170 columnShape(spec),
171 },
172 };
173 }
174
175 fn rowLayerNormSpecialization(comptime spec: Row, comptime parameterization: RowNormalizationParameterization) entry.Specialization {
176 return .{
177 .dtype = .f32,
178 .operation = .{ .row_normalization = .{ .layernorm = parameterization } },
179 .inputs = rowLayerNormInputs(spec, parameterization),
180 .outputs = &.{rowMatrixShape(spec)},
181 .reductions = &.{
182 entry.reduction("row_sum", .sum, columnShape(spec)),
183 entry.dependentReduction("row_variance_sum", .sum_squared_difference, columnShape(spec), &.{"row_sum"}),
184 },
185 .reduction_reuse = &.{
186 entry.reductionReuse("row_sum", rowShape(spec)),
187 entry.reductionReuse("row_variance_sum", rowShape(spec)),
188 },
189 .launch = rowLaunch(spec),
190 .schedule = entry.threadBlocks2D(spec.col_axis, spec.cols, spec.row_axis, spec.rows, spec.threads.x, spec.threads.y),
191 };
192 }
193
194 fn rowDistributionValue(
195 inner: anytype,
196 comptime kind: RowDistribution,
197 shifted: kernel.Value,
198 denominator: kernel.Value,
199 ) !kernel.Value {
200 return switch (kind) {
201 .softmax => blk: {
202 const numerator = try inner.exp(shifted);
203 break :blk try inner.div(numerator, denominator);
204 },
205 .log_softmax => blk: {
206 const log_denominator = try inner.log(denominator);
207 break :blk try inner.sub(shifted, log_denominator);
208 },
209 };
210 }
211
212 fn row_distribution_each(inner: anytype, index: kernel.Index2D, ctx: anytype) !void {
213 const row_offset = try rowOffset(inner, index, ctx.spec);
214 const item_index = try rowItemIndex(inner, index, row_offset);
215 const each_args = ctx.args;
216 const src = each_args.param(.src);
217 const dst = each_args.param(.dst);
218 const first = try src.load(inner, row_offset);
219 const row_max = try foldColumnsFrom(inner, ctx.spec, 1, first.raw(), .{
220 .src = src,
221 .row_offset = row_offset,
222 }, row_distribution_row_max);
223 const current = try src.load(inner, item_index);
224 const shifted = try inner.sub(current.raw(), row_max);
225 const zero = try inner.constantFloat(.f32, 0.0);
226 const denominator = try foldColumns(inner, ctx.spec, zero, .{
227 .src = src,
228 .row_offset = row_offset,
229 .row_max = row_max,
230 }, row_distribution_denominator);
231 const normalized = try rowDistributionValue(inner, ctx.kind, shifted, denominator);
232 try dst.store(inner, normalized, item_index);
233 }
234
235 fn row_distribution_row_max(fold_inner: anytype, offset: kernel.Value, acc: kernel.Value, ctx: anytype) !kernel.Value {
236 const reduce_index = try fold_inner.add(ctx.row_offset, offset);
237 const value = try ctx.src.load(fold_inner, reduce_index);
238 return fold_inner.max(acc, value.raw());
239 }
240
241 fn row_distribution_denominator(fold_inner: anytype, offset: kernel.Value, acc: kernel.Value, ctx: anytype) !kernel.Value {
242 const reduce_index = try fold_inner.add(ctx.row_offset, offset);
243 const value = try ctx.src.load(fold_inner, reduce_index);
244 const reduce_shifted = try fold_inner.sub(value.raw(), ctx.row_max);
245 const exp_value = try fold_inner.exp(reduce_shifted);
246 return fold_inner.add(acc, exp_value);
247 }
248
249 fn rowDistributionProgram(comptime spec: Row, comptime kind: RowDistribution) type {
250 const Body = struct {
251 fn run(k: anytype, args: anytype) !void {
252 _ = try k.forEach2D(rowDomain(spec), .{ .spec = spec, .kind = kind, .args = args }, row_distribution_each);
253 }
254 };
255
256 return kernel.logical.Program(.{
257 .name = std.fmt.comptimePrint(
258 "accy_kernel_normalization_{s}{}x{}_{}x{}_f32",
259 .{ rowDistributionName(kind), spec.rows, spec.cols, spec.threads.x, spec.threads.y },
260 ),
261 .parameters = .{
262 .dst = kernel.dynamicBuffer(.f32),
263 .src = kernel.dynamicBuffer(.f32),
264 },
265 .body = Body.run,
266 }).withSchedule(kernel.logical.schedule.threadBlocks(.{
267 .x = spec.threads.x,
268 .y = spec.threads.y,
269 }));
270 }
271
272 fn rowDistributionF32(comptime spec: Row, comptime kind: RowDistribution) type {
273 return entry.Entry(rowDistributionProgram(spec, kind), .{
274 .target = std.fmt.comptimePrint(
275 "accy.kernel.normalization.{s}{}x{}_{}x{}_f32",
276 .{ rowDistributionName(kind), spec.rows, spec.cols, spec.threads.x, spec.threads.y },
277 ),
278 .layer = .logical,
279 .category = .normalization,
280 .specialization = rowDistributionSpecialization(spec, kind),
281 });
282 }
283
284 pub fn rowSoftmaxF32(comptime spec: Row) type {
285 return rowDistributionF32(spec, .softmax);
286 }
287
288 pub fn rowLogSoftmaxF32(comptime spec: Row) type {
289 return rowDistributionF32(spec, .log_softmax);
290 }
291
292 pub const RowSoftmax2x4F32 = rowSoftmaxF32(.{
293 .rows = 2,
294 .cols = 4,
295 .threads = .{ .x = 4, .y = 2 },
296 });
297 pub const RowSoftmax2x4ThreadBlocks2x2F32 = rowSoftmaxF32(.{
298 .rows = 2,
299 .cols = 4,
300 .threads = .{ .x = 2, .y = 2 },
301 });
302 pub const RowLogSoftmax2x4F32 = rowLogSoftmaxF32(.{
303 .rows = 2,
304 .cols = 4,
305 .threads = .{ .x = 4, .y = 2 },
306 });
307
308 fn row_rms_norm_each(inner: anytype, index: kernel.Index2D, ctx: anytype) !void {
309 const row_offset = try rowOffset(inner, index, ctx.spec);
310 const item_index = try rowItemIndex(inner, index, row_offset);
311 const each_args = ctx.args;
312 const src = each_args.param(.src);
313 const scale = each_args.param(.scale);
314 const zero = try inner.constantFloat(.f32, 0.0);
315 const sum_squares = try foldColumns(inner, ctx.spec, zero, .{
316 .src = src,
317 .row_offset = row_offset,
318 }, row_rms_norm_sum_squares);
319 const cols_count = try inner.constantFloat(.f32, floatExtent(ctx.spec.cols));
320 const mean_square = try inner.div(sum_squares, cols_count);
321 const stabilized = try inner.add(mean_square, each_args.param(.epsilon).raw());
322 const root = try inner.sqrt(stabilized);
323 const value = try src.load(inner, item_index);
324 const normalized = try inner.div(value.raw(), root);
325 const weight = try scale.load(inner, index.x.index);
326 const weighted = try inner.mul(normalized, weight.raw());
327 try each_args.param(.dst).store(inner, weighted, item_index);
328 }
329
330 fn row_rms_norm_sum_squares(fold_inner: anytype, offset: kernel.Value, acc: kernel.Value, ctx: anytype) !kernel.Value {
331 const reduce_index = try fold_inner.add(ctx.row_offset, offset);
332 const value = try ctx.src.load(fold_inner, reduce_index);
333 const square = try fold_inner.mul(value.raw(), value.raw());
334 return fold_inner.add(acc, square);
335 }
336
337 fn rowRmsNormProgram(comptime spec: Row) type {
338 const Body = struct {
339 fn run(k: anytype, args: anytype) !void {
340 _ = try k.forEach2D(rowDomain(spec), .{ .spec = spec, .args = args }, row_rms_norm_each);
341 }
342 };
343
344 return kernel.logical.Program(.{
345 .name = std.fmt.comptimePrint(
346 "accy_kernel_normalization_rmsnorm{}x{}_{}x{}_f32",
347 .{ spec.rows, spec.cols, spec.threads.x, spec.threads.y },
348 ),
349 .parameters = .{
350 .dst = kernel.dynamicBuffer(.f32),
351 .src = kernel.dynamicBuffer(.f32),
352 .scale = kernel.dynamicBuffer(.f32),
353 .epsilon = kernel.scalar(.f32),
354 },
355 .body = Body.run,
356 }).withSchedule(kernel.logical.schedule.threadBlocks(.{
357 .x = spec.threads.x,
358 .y = spec.threads.y,
359 }));
360 }
361
362 pub fn rowRmsNormF32(comptime spec: Row) type {
363 return entry.Entry(rowRmsNormProgram(spec), .{
364 .target = std.fmt.comptimePrint(
365 "accy.kernel.normalization.rmsnorm{}x{}_{}x{}_f32",
366 .{ spec.rows, spec.cols, spec.threads.x, spec.threads.y },
367 ),
368 .layer = .logical,
369 .category = .normalization,
370 .specialization = rowWeightedSpecialization(spec),
371 });
372 }
373
374 pub const RowRmsNorm2x4F32 = rowRmsNormF32(.{
375 .rows = 2,
376 .cols = 4,
377 .threads = .{ .x = 4, .y = 2 },
378 });
379
380 fn row_residual_rms_norm_each(inner: anytype, index: kernel.Index2D, ctx: anytype) !void {
381 const row_offset = try rowOffset(inner, index, ctx.spec);
382 const item_index = try rowItemIndex(inner, index, row_offset);
383 const each_args = ctx.args;
384 const src = each_args.param(.src);
385 const residual = each_args.param(.residual);
386 const scale = each_args.param(.scale);
387 const zero = try inner.constantFloat(.f32, 0.0);
388 const sum_squares = try foldColumns(inner, ctx.spec, zero, .{
389 .src = src,
390 .residual = residual,
391 .row_offset = row_offset,
392 }, row_residual_rms_norm_sum_squares);
393 const cols_count = try inner.constantFloat(.f32, floatExtent(ctx.spec.cols));
394 const mean_square = try inner.div(sum_squares, cols_count);
395 const stabilized = try inner.add(mean_square, each_args.param(.epsilon).raw());
396 const root = try inner.sqrt(stabilized);
397 const value = try src.load(inner, item_index);
398 const residual_value = try residual.load(inner, item_index);
399 const combined = try inner.add(value.raw(), residual_value.raw());
400 const normalized = try inner.div(combined, root);
401 const weight = try scale.load(inner, index.x.index);
402 const weighted = try inner.mul(normalized, weight.raw());
403 try each_args.param(.dst).store(inner, weighted, item_index);
404 }
405
406 fn row_residual_rms_norm_sum_squares(fold_inner: anytype, offset: kernel.Value, acc: kernel.Value, ctx: anytype) !kernel.Value {
407 const reduce_index = try fold_inner.add(ctx.row_offset, offset);
408 const value = try ctx.src.load(fold_inner, reduce_index);
409 const residual_value = try ctx.residual.load(fold_inner, reduce_index);
410 const combined = try fold_inner.add(value.raw(), residual_value.raw());
411 const square = try fold_inner.mul(combined, combined);
412 return fold_inner.add(acc, square);
413 }
414
415 fn rowResidualRmsNormProgram(comptime spec: Row) type {
416 const Body = struct {
417 fn run(k: anytype, args: anytype) !void {
418 _ = try k.forEach2D(rowDomain(spec), .{ .spec = spec, .args = args }, row_residual_rms_norm_each);
419 }
420 };
421
422 return kernel.logical.Program(.{
423 .name = std.fmt.comptimePrint(
424 "accy_kernel_fused_row_residual_rmsnorm{}x{}_{}x{}_f32",
425 .{ spec.rows, spec.cols, spec.threads.x, spec.threads.y },
426 ),
427 .parameters = .{
428 .dst = kernel.dynamicBuffer(.f32),
429 .src = kernel.dynamicBuffer(.f32),
430 .residual = kernel.dynamicBuffer(.f32),
431 .scale = kernel.dynamicBuffer(.f32),
432 .epsilon = kernel.scalar(.f32),
433 },
434 .body = Body.run,
435 }).withSchedule(kernel.logical.schedule.threadBlocks(.{
436 .x = spec.threads.x,
437 .y = spec.threads.y,
438 }));
439 }
440
441 pub fn rowResidualRmsNormF32(comptime spec: Row) type {
442 return entry.Entry(rowResidualRmsNormProgram(spec), .{
443 .target = std.fmt.comptimePrint(
444 "accy.kernel.fused.row_residual_rmsnorm{}x{}_{}x{}_f32",
445 .{ spec.rows, spec.cols, spec.threads.x, spec.threads.y },
446 ),
447 .layer = .logical,
448 .category = .fused,
449 .specialization = rowResidualRmsNormSpecialization(spec),
450 });
451 }
452
453 pub const RowResidualRmsNorm2x4F32 = rowResidualRmsNormF32(.{
454 .rows = 2,
455 .cols = 4,
456 .threads = .{ .x = 4, .y = 2 },
457 });
458
459 fn rowLayerNormOutput(
460 inner: anytype,
461 index: kernel.Index2D,
462 normalized: kernel.Value,
463 each_args: anytype,
464 comptime parameterization: RowNormalizationParameterization,
465 ) !kernel.Value {
466 return switch (parameterization) {
467 .none => normalized,
468 .scale => blk: {
469 const scale = each_args.param(.scale);
470 const weight = try scale.load(inner, index.x.index);
471 break :blk try inner.mul(normalized, weight.raw());
472 },
473 .scale_bias => blk: {
474 const scale = each_args.param(.scale);
475 const bias = each_args.param(.bias);
476 const weight = try scale.load(inner, index.x.index);
477 const scaled = try inner.mul(normalized, weight.raw());
478 const shift = try bias.load(inner, index.x.index);
479 break :blk try inner.add(scaled, shift.raw());
480 },
481 };
482 }
483
484 fn row_layer_norm_each(inner: anytype, index: kernel.Index2D, ctx: anytype) !void {
485 const row_offset = try rowOffset(inner, index, ctx.spec);
486 const item_index = try rowItemIndex(inner, index, row_offset);
487 const each_args = ctx.args;
488 const src = each_args.param(.src);
489 const zero = try inner.constantFloat(.f32, 0.0);
490 const sum = try foldColumns(inner, ctx.spec, zero, .{
491 .src = src,
492 .row_offset = row_offset,
493 }, row_layer_norm_sum);
494 const cols_count = try inner.constantFloat(.f32, floatExtent(ctx.spec.cols));
495 const mean = try inner.div(sum, cols_count);
496 const variance_sum = try foldColumns(inner, ctx.spec, zero, .{
497 .src = src,
498 .row_offset = row_offset,
499 .mean = mean,
500 }, row_layer_norm_variance_sum);
501 const variance = try inner.div(variance_sum, cols_count);
502 const stabilized = try inner.add(variance, each_args.param(.epsilon).raw());
503 const root = try inner.sqrt(stabilized);
504 const value = try src.load(inner, item_index);
505 const centered = try inner.sub(value.raw(), mean);
506 const normalized = try inner.div(centered, root);
507 const output = try rowLayerNormOutput(inner, index, normalized, each_args, ctx.parameterization);
508 try each_args.param(.dst).store(inner, output, item_index);
509 }
510
511 fn row_layer_norm_sum(fold_inner: anytype, offset: kernel.Value, acc: kernel.Value, ctx: anytype) !kernel.Value {
512 const reduce_index = try fold_inner.add(ctx.row_offset, offset);
513 const value = try ctx.src.load(fold_inner, reduce_index);
514 return fold_inner.add(acc, value.raw());
515 }
516
517 fn row_layer_norm_variance_sum(fold_inner: anytype, offset: kernel.Value, acc: kernel.Value, ctx: anytype) !kernel.Value {
518 const reduce_index = try fold_inner.add(ctx.row_offset, offset);
519 const value = try ctx.src.load(fold_inner, reduce_index);
520 const centered = try fold_inner.sub(value.raw(), ctx.mean);
521 const square = try fold_inner.mul(centered, centered);
522 return fold_inner.add(acc, square);
523 }
524
525 fn rowLayerNormProgram(comptime spec: Row, comptime parameterization: RowNormalizationParameterization) type {
526 const Body = struct {
527 fn run(k: anytype, args: anytype) !void {
528 _ = try k.forEach2D(rowDomain(spec), .{ .spec = spec, .parameterization = parameterization, .args = args }, row_layer_norm_each);
529 }
530 };
531
532 return switch (parameterization) {
533 .none => kernel.logical.Program(.{
534 .name = std.fmt.comptimePrint(
535 "accy_kernel_normalization_{s}{}x{}_{}x{}_f32",
536 .{ rowLayerNormName(parameterization), spec.rows, spec.cols, spec.threads.x, spec.threads.y },
537 ),
538 .parameters = .{
539 .dst = kernel.dynamicBuffer(.f32),
540 .src = kernel.dynamicBuffer(.f32),
541 .epsilon = kernel.scalar(.f32),
542 },
543 .body = Body.run,
544 }).withSchedule(kernel.logical.schedule.threadBlocks(.{
545 .x = spec.threads.x,
546 .y = spec.threads.y,
547 })),
548 .scale => kernel.logical.Program(.{
549 .name = std.fmt.comptimePrint(
550 "accy_kernel_normalization_{s}{}x{}_{}x{}_f32",
551 .{ rowLayerNormName(parameterization), spec.rows, spec.cols, spec.threads.x, spec.threads.y },
552 ),
553 .parameters = .{
554 .dst = kernel.dynamicBuffer(.f32),
555 .src = kernel.dynamicBuffer(.f32),
556 .scale = kernel.dynamicBuffer(.f32),
557 .epsilon = kernel.scalar(.f32),
558 },
559 .body = Body.run,
560 }).withSchedule(kernel.logical.schedule.threadBlocks(.{
561 .x = spec.threads.x,
562 .y = spec.threads.y,
563 })),
564 .scale_bias => kernel.logical.Program(.{
565 .name = std.fmt.comptimePrint(
566 "accy_kernel_normalization_{s}{}x{}_{}x{}_f32",
567 .{ rowLayerNormName(parameterization), spec.rows, spec.cols, spec.threads.x, spec.threads.y },
568 ),
569 .parameters = .{
570 .dst = kernel.dynamicBuffer(.f32),
571 .src = kernel.dynamicBuffer(.f32),
572 .scale = kernel.dynamicBuffer(.f32),
573 .bias = kernel.dynamicBuffer(.f32),
574 .epsilon = kernel.scalar(.f32),
575 },
576 .body = Body.run,
577 }).withSchedule(kernel.logical.schedule.threadBlocks(.{
578 .x = spec.threads.x,
579 .y = spec.threads.y,
580 })),
581 };
582 }
583
584 fn rowParameterizedLayerNormF32(comptime spec: Row, comptime parameterization: RowNormalizationParameterization) type {
585 return entry.Entry(rowLayerNormProgram(spec, parameterization), .{
586 .target = std.fmt.comptimePrint(
587 "accy.kernel.normalization.{s}{}x{}_{}x{}_f32",
588 .{ rowLayerNormName(parameterization), spec.rows, spec.cols, spec.threads.x, spec.threads.y },
589 ),
590 .layer = .logical,
591 .category = .normalization,
592 .specialization = rowLayerNormSpecialization(spec, parameterization),
593 });
594 }
595
596 pub fn rowLayerNormF32(comptime spec: Row) type {
597 return rowParameterizedLayerNormF32(spec, .none);
598 }
599
600 pub fn rowAffineLayerNormF32(comptime spec: Row) type {
601 return rowParameterizedLayerNormF32(spec, .scale_bias);
602 }
603
604 pub const RowLayerNorm2x4F32 = rowLayerNormF32(.{
605 .rows = 2,
606 .cols = 4,
607 .threads = .{ .x = 4, .y = 2 },
608 });
609 pub const RowAffineLayerNorm2x4F32 = rowAffineLayerNormF32(.{
610 .rows = 2,
611 .cols = 4,
612 .threads = .{ .x = 4, .y = 2 },
613 });
614
615 test "normalization row softmax entry runs on CPU" {
616 var src = [_]f32{
617 1.0, 2.0, 3.0, 4.0,
618 1.0, 1.0, 1.0, 1.0,
619 };
620 var dst = @as([8]f32, @splat(0.0));
621
622 try RowSoftmax2x4F32.runCpu(std.testing.allocator, RowSoftmax2x4F32.Limits.testing, &.{
623 kernel.argumentBuffer(f32, dst[0..]),
624 kernel.argumentBuffer(f32, src[0..]),
625 });
626
627 const expected = [_]f32{
628 0.032058604, 0.08714432, 0.23688282, 0.6439143,
629 0.25, 0.25, 0.25, 0.25,
630 };
631 for (expected, dst) |expected_value, actual| {
632 try std.testing.expectApproxEqAbs(expected_value, actual, 0.0001);
633 }
634
635 const launch_value = try RowSoftmax2x4F32.launch(std.testing.allocator, RowSoftmax2x4F32.Limits.testing);
636 try std.testing.expectEqual(@as(u32, 1), launch_value.grid[0]);
637 try std.testing.expectEqual(@as(u32, 1), launch_value.grid[1]);
638 try std.testing.expectEqual(@as(u32, 4), launch_value.block[0]);
639 try std.testing.expectEqual(@as(u32, 2), launch_value.block[1]);
640 }
641
642 test "normalization row log softmax entry runs on CPU" {
643 var src = [_]f32{
644 1.0, 2.0, 3.0, 4.0,
645 1.0, 1.0, 1.0, 1.0,
646 };
647 var dst = @as([8]f32, @splat(0.0));
648
649 try RowLogSoftmax2x4F32.runCpu(std.testing.allocator, RowLogSoftmax2x4F32.Limits.testing, &.{
650 kernel.argumentBuffer(f32, dst[0..]),
651 kernel.argumentBuffer(f32, src[0..]),
652 });
653
654 const expected = [_]f32{
655 -3.4401898, -2.4401898, -1.4401897, -0.4401897,
656 -1.3862944, -1.3862944, -1.3862944, -1.3862944,
657 };
658 for (expected, dst) |expected_value, actual| {
659 try std.testing.expectApproxEqAbs(expected_value, actual, 0.0001);
660 }
661
662 const launch_value = try RowLogSoftmax2x4F32.launch(std.testing.allocator, RowLogSoftmax2x4F32.Limits.testing);
663 try std.testing.expectEqual(@as(u32, 1), launch_value.grid[0]);
664 try std.testing.expectEqual(@as(u32, 1), launch_value.grid[1]);
665 try std.testing.expectEqual(@as(u32, 4), launch_value.block[0]);
666 try std.testing.expectEqual(@as(u32, 2), launch_value.block[1]);
667 }
668
669 test "normalization row rmsnorm entry runs on CPU" {
670 var src = [_]f32{
671 1.0, 2.0, 3.0, 4.0,
672 2.0, 0.0, -2.0, 0.0,
673 };
674 var scale = [_]f32{ 1.0, 0.5, 2.0, -1.0 };
675 var dst = @as([8]f32, @splat(0.0));
676
677 try RowRmsNorm2x4F32.runCpu(std.testing.allocator, RowRmsNorm2x4F32.Limits.testing, &.{
678 kernel.argumentBuffer(f32, dst[0..]),
679 kernel.argumentBuffer(f32, src[0..]),
680 kernel.argumentBuffer(f32, scale[0..]),
681 kernel.argumentF32(0.00001),
682 });
683
684 const expected = [_]f32{
685 0.36514813, 0.36514813, 2.1908886, -1.4605925,
686 1.4142101, 0.0, -2.8284202, -0.0,
687 };
688 for (expected, dst) |expected_value, actual| {
689 try std.testing.expectApproxEqAbs(expected_value, actual, 0.0001);
690 }
691
692 const launch_value = try RowRmsNorm2x4F32.launch(std.testing.allocator, RowRmsNorm2x4F32.Limits.testing);
693 try std.testing.expectEqual(@as(u32, 1), launch_value.grid[0]);
694 try std.testing.expectEqual(@as(u32, 1), launch_value.grid[1]);
695 try std.testing.expectEqual(@as(u32, 4), launch_value.block[0]);
696 try std.testing.expectEqual(@as(u32, 2), launch_value.block[1]);
697 }
698
699 test "fused row residual rmsnorm entry runs on CPU" {
700 var src = [_]f32{
701 1.0, 2.0, 3.0, 4.0,
702 2.0, 0.0, -2.0, 0.0,
703 };
704 var residual = [_]f32{
705 0.5, -0.5, 1.0, -1.0,
706 -1.0, 1.0, 0.5, -0.5,
707 };
708 var scale = [_]f32{ 1.0, 0.5, 2.0, -1.0 };
709 var dst = @as([8]f32, @splat(0.0));
710
711 try RowResidualRmsNorm2x4F32.runCpu(std.testing.allocator, RowResidualRmsNorm2x4F32.Limits.testing, &.{
712 kernel.argumentBuffer(f32, dst[0..]),
713 kernel.argumentBuffer(f32, src[0..]),
714 kernel.argumentBuffer(f32, residual[0..]),
715 kernel.argumentBuffer(f32, scale[0..]),
716 kernel.argumentF32(0.00001),
717 });
718
719 const expected = [_]f32{
720 0.5523444, 0.2761722, 2.9458368, -1.1046888,
721 0.9428049, 0.4714024, -2.8284146, 0.4714024,
722 };
723 for (expected, dst) |expected_value, actual| {
724 try std.testing.expectApproxEqAbs(expected_value, actual, 0.0001);
725 }
726
727 const launch_value = try RowResidualRmsNorm2x4F32.launch(std.testing.allocator, RowResidualRmsNorm2x4F32.Limits.testing);
728 try std.testing.expectEqual(@as(u32, 1), launch_value.grid[0]);
729 try std.testing.expectEqual(@as(u32, 1), launch_value.grid[1]);
730 try std.testing.expectEqual(@as(u32, 4), launch_value.block[0]);
731 try std.testing.expectEqual(@as(u32, 2), launch_value.block[1]);
732 try std.testing.expectEqual(entry.Category.fused, RowResidualRmsNorm2x4F32.category);
733 try std.testing.expect(RowResidualRmsNorm2x4F32.specialization.operationIs(.{ .row_normalization = .{ .rmsnorm = .scale } }));
734 try std.testing.expect(RowResidualRmsNorm2x4F32.specialization.inputTransformMatches(0, .{
735 .operator = .residual_add,
736 .input_index = 1,
737 .extents = &.{ 2, 4 },
738 }));
739 try std.testing.expect(RowResidualRmsNorm2x4F32.specialization.reductionMatches(0, .{
740 .name = "row_sum_squares",
741 .operator = .sum_squares,
742 .extents = &.{4},
743 }));
744 try std.testing.expectEqualDeep(RowResidualRmsNorm2x4F32.specialization.launch.?, RowResidualRmsNorm2x4F32.specialization.schedule.?.launch());
745 }
746
747 test "normalization row layernorm entry runs on CPU" {
748 var src = [_]f32{
749 1.0, 2.0, 3.0, 4.0,
750 2.0, 0.0, -2.0, 0.0,
751 };
752 var dst = @as([8]f32, @splat(0.0));
753
754 try RowLayerNorm2x4F32.runCpu(std.testing.allocator, RowLayerNorm2x4F32.Limits.testing, &.{
755 kernel.argumentBuffer(f32, dst[0..]),
756 kernel.argumentBuffer(f32, src[0..]),
757 kernel.argumentF32(0.00001),
758 });
759
760 const expected = [_]f32{
761 -1.3416355, -0.44721183, 0.44721183, 1.3416355,
762 1.4142101, 0.0, -1.4142101, 0.0,
763 };
764 for (expected, dst) |expected_value, actual| {
765 try std.testing.expectApproxEqAbs(expected_value, actual, 0.0001);
766 }
767
768 const launch_value = try RowLayerNorm2x4F32.launch(std.testing.allocator, RowLayerNorm2x4F32.Limits.testing);
769 try std.testing.expectEqual(@as(u32, 1), launch_value.grid[0]);
770 try std.testing.expectEqual(@as(u32, 1), launch_value.grid[1]);
771 try std.testing.expectEqual(@as(u32, 4), launch_value.block[0]);
772 try std.testing.expectEqual(@as(u32, 2), launch_value.block[1]);
773 }
774
775 test "normalization row affine layernorm entry runs on CPU" {
776 var src = [_]f32{
777 1.0, 2.0, 3.0, 4.0,
778 2.0, 0.0, -2.0, 0.0,
779 };
780 var scale = [_]f32{ 1.0, 0.5, 2.0, -1.0 };
781 var bias = [_]f32{ 0.1, -0.2, 0.3, 0.4 };
782 var dst = @as([8]f32, @splat(0.0));
783
784 try RowAffineLayerNorm2x4F32.runCpu(std.testing.allocator, RowAffineLayerNorm2x4F32.Limits.testing, &.{
785 kernel.argumentBuffer(f32, dst[0..]),
786 kernel.argumentBuffer(f32, src[0..]),
787 kernel.argumentBuffer(f32, scale[0..]),
788 kernel.argumentBuffer(f32, bias[0..]),
789 kernel.argumentF32(0.00001),
790 });
791
792 const expected = [_]f32{
793 -1.2416355, -0.4236059, 1.1944237, -0.9416355,
794 1.5142101, -0.2, -2.5284202, 0.4,
795 };
796 for (expected, dst) |expected_value, actual| {
797 try std.testing.expectApproxEqAbs(expected_value, actual, 0.0001);
798 }
799
800 const launch_value = try RowAffineLayerNorm2x4F32.launch(std.testing.allocator, RowAffineLayerNorm2x4F32.Limits.testing);
801 try std.testing.expectEqual(@as(u32, 1), launch_value.grid[0]);
802 try std.testing.expectEqual(@as(u32, 1), launch_value.grid[1]);
803 try std.testing.expectEqual(@as(u32, 4), launch_value.block[0]);
804 try std.testing.expectEqual(@as(u32, 2), launch_value.block[1]);
805 }
806
807 test "normalization constructor creates independent shape-specialized entries" {
808 const RowSoftmax3x5F32 = rowSoftmaxF32(.{
809 .rows = 3,
810 .cols = 5,
811 .threads = .{ .x = 5, .y = 1 },
812 });
813 const RowLogSoftmax3x5F32 = rowLogSoftmaxF32(.{
814 .rows = 3,
815 .cols = 5,
816 .threads = .{ .x = 5, .y = 1 },
817 });
818 const RowRmsNorm3x5F32 = rowRmsNormF32(.{
819 .rows = 3,
820 .cols = 5,
821 .threads = .{ .x = 5, .y = 1 },
822 });
823 const RowResidualRmsNorm3x5F32 = rowResidualRmsNormF32(.{
824 .rows = 3,
825 .cols = 5,
826 .threads = .{ .x = 5, .y = 1 },
827 });
828 const RowLayerNorm3x5F32 = rowLayerNormF32(.{
829 .rows = 3,
830 .cols = 5,
831 .threads = .{ .x = 5, .y = 1 },
832 });
833 const RowAffineLayerNorm3x5F32 = rowAffineLayerNormF32(.{
834 .rows = 3,
835 .cols = 5,
836 .threads = .{ .x = 5, .y = 1 },
837 });
838
839 try std.testing.expectEqualStrings("accy.kernel.normalization.softmax2x4_4x2_f32", RowSoftmax2x4F32.target);
840 try std.testing.expectEqualStrings("accy.kernel.normalization.softmax2x4_2x2_f32", RowSoftmax2x4ThreadBlocks2x2F32.target);
841 try std.testing.expectEqualStrings("accy.kernel.normalization.softmax3x5_5x1_f32", RowSoftmax3x5F32.target);
842 try std.testing.expectEqualStrings("accy.kernel.normalization.log_softmax2x4_4x2_f32", RowLogSoftmax2x4F32.target);
843 try std.testing.expectEqualStrings("accy.kernel.normalization.log_softmax3x5_5x1_f32", RowLogSoftmax3x5F32.target);
844 try std.testing.expectEqualStrings("accy.kernel.normalization.rmsnorm2x4_4x2_f32", RowRmsNorm2x4F32.target);
845 try std.testing.expectEqualStrings("accy.kernel.normalization.rmsnorm3x5_5x1_f32", RowRmsNorm3x5F32.target);
846 try std.testing.expectEqualStrings("accy.kernel.fused.row_residual_rmsnorm2x4_4x2_f32", RowResidualRmsNorm2x4F32.target);
847 try std.testing.expectEqualStrings("accy.kernel.fused.row_residual_rmsnorm3x5_5x1_f32", RowResidualRmsNorm3x5F32.target);
848 try std.testing.expectEqualStrings("accy.kernel.normalization.layernorm2x4_4x2_f32", RowLayerNorm2x4F32.target);
849 try std.testing.expectEqualStrings("accy.kernel.normalization.layernorm3x5_5x1_f32", RowLayerNorm3x5F32.target);
850 try std.testing.expectEqualStrings("accy.kernel.normalization.layernorm_affine2x4_4x2_f32", RowAffineLayerNorm2x4F32.target);
851 try std.testing.expectEqualStrings("accy.kernel.normalization.layernorm_affine3x5_5x1_f32", RowAffineLayerNorm3x5F32.target);
852 try std.testing.expectEqual(entry.Category.normalization, RowSoftmax2x4F32.category);
853 try std.testing.expectEqual(entry.Category.normalization, RowSoftmax2x4ThreadBlocks2x2F32.category);
854 try std.testing.expectEqual(entry.Category.normalization, RowLogSoftmax2x4F32.category);
855 try std.testing.expectEqual(entry.Category.normalization, RowRmsNorm2x4F32.category);
856 try std.testing.expectEqual(entry.Category.fused, RowResidualRmsNorm2x4F32.category);
857 try std.testing.expectEqual(entry.Category.normalization, RowLayerNorm2x4F32.category);
858 try std.testing.expectEqual(entry.Category.normalization, RowAffineLayerNorm2x4F32.category);
859 try std.testing.expect(RowSoftmax3x5F32.specialization.operationIs(.{ .row_normalization = .softmax }));
860 try std.testing.expect(RowSoftmax2x4ThreadBlocks2x2F32.specialization.operationIs(.{ .row_normalization = .softmax }));
861 try std.testing.expect(RowLogSoftmax3x5F32.specialization.operationIs(.{ .row_normalization = .log_softmax }));
862 try std.testing.expectEqual(@as(usize, 1), RowSoftmax3x5F32.specialization.inputs.len);
863 try std.testing.expectEqual(@as(u64, 15), RowSoftmax3x5F32.specialization.outputs[0].elementCount().?);
864 try std.testing.expectEqual(@as(usize, 2), RowSoftmax3x5F32.specialization.reductions.len);
865 try std.testing.expectEqualStrings("row_max", RowSoftmax3x5F32.specialization.reductions[0].name);
866 try std.testing.expectEqual(entry.ReductionOperator.maximum, RowSoftmax3x5F32.specialization.reductions[0].operator);
867 try std.testing.expectEqual(@as(u64, 5), RowSoftmax3x5F32.specialization.reductions[0].shape.elementCount().?);
868 try std.testing.expectEqualStrings("row_exp_sum", RowSoftmax3x5F32.specialization.reductions[1].name);
869 try std.testing.expectEqual(entry.ReductionOperator.sum_exp_shifted, RowSoftmax3x5F32.specialization.reductions[1].operator);
870 try std.testing.expectEqual(@as(usize, 1), RowSoftmax3x5F32.specialization.reductions[1].dependencies.len);
871 try std.testing.expectEqualStrings("row_max", RowSoftmax3x5F32.specialization.reductions[1].dependencies[0]);
872 try std.testing.expect(RowSoftmax3x5F32.specialization.reductionReuseScopesAreValid());
873 try std.testing.expectEqual(@as(usize, 2), RowSoftmax3x5F32.specialization.reduction_reuse.len);
874 try std.testing.expect(RowSoftmax3x5F32.specialization.reductionReuseMatches(0, .{ .reduction = "row_max", .extents = &.{3} }));
875 try std.testing.expect(RowSoftmax3x5F32.specialization.reductionReuseMatches(1, .{ .reduction = "row_exp_sum", .extents = &.{3} }));
876 try std.testing.expectEqual(@as(u64, 30), RowSoftmax3x5F32.specialization.estimatedElementOps().?);
877 try std.testing.expectEqual(@as(u32, 1), RowSoftmax3x5F32.specialization.launch.?.grid[0]);
878 try std.testing.expectEqual(@as(u32, 3), RowSoftmax3x5F32.specialization.launch.?.grid[1]);
879 try std.testing.expectEqualDeep(RowSoftmax3x5F32.specialization.launch.?, RowSoftmax3x5F32.specialization.schedule.?.launch());
880 try std.testing.expectEqual(@as(usize, 3), RowSoftmax3x5F32.specialization.schedule.?.bindings.len);
881 try std.testing.expectEqualStrings("col", RowSoftmax3x5F32.specialization.schedule.?.bindings[0].axis);
882 try std.testing.expectEqual(kernel.BindTarget.thread_x, RowSoftmax3x5F32.specialization.schedule.?.bindings[0].target);
883 try std.testing.expectEqualStrings("row_tile", RowSoftmax3x5F32.specialization.schedule.?.bindings[1].axis);
884 try std.testing.expectEqual(kernel.BindTarget.block_y, RowSoftmax3x5F32.specialization.schedule.?.bindings[1].target);
885 try std.testing.expectEqualStrings("row_lane", RowSoftmax3x5F32.specialization.schedule.?.bindings[2].axis);
886 try std.testing.expectEqual(kernel.BindTarget.thread_y, RowSoftmax3x5F32.specialization.schedule.?.bindings[2].target);
887 var softmax_snapshot = try RowSoftmax3x5F32.scheduleSnapshot(std.testing.allocator, RowSoftmax3x5F32.Limits.testing);
888 defer softmax_snapshot.deinit(std.testing.allocator);
889 try std.testing.expect(RowSoftmax3x5F32.specialization.schedule.?.matchesSnapshot(&softmax_snapshot));
890
891 try std.testing.expectEqual(@as(usize, 1), RowLogSoftmax3x5F32.specialization.inputs.len);
892 try std.testing.expectEqual(@as(u64, 15), RowLogSoftmax3x5F32.specialization.outputs[0].elementCount().?);
893 try std.testing.expectEqual(@as(usize, 2), RowLogSoftmax3x5F32.specialization.reductions.len);
894 try std.testing.expectEqualStrings("row_max", RowLogSoftmax3x5F32.specialization.reductions[0].name);
895 try std.testing.expectEqual(entry.ReductionOperator.maximum, RowLogSoftmax3x5F32.specialization.reductions[0].operator);
896 try std.testing.expectEqualStrings("row_exp_sum", RowLogSoftmax3x5F32.specialization.reductions[1].name);
897 try std.testing.expectEqual(entry.ReductionOperator.sum_exp_shifted, RowLogSoftmax3x5F32.specialization.reductions[1].operator);
898 try std.testing.expectEqual(@as(usize, 1), RowLogSoftmax3x5F32.specialization.reductions[1].dependencies.len);
899 try std.testing.expectEqualStrings("row_max", RowLogSoftmax3x5F32.specialization.reductions[1].dependencies[0]);
900 try std.testing.expect(RowLogSoftmax3x5F32.specialization.reductionReuseScopesAreValid());
901 try std.testing.expectEqual(@as(usize, 2), RowLogSoftmax3x5F32.specialization.reduction_reuse.len);
902 try std.testing.expect(RowLogSoftmax3x5F32.specialization.reductionReuseMatches(0, .{ .reduction = "row_max", .extents = &.{3} }));
903 try std.testing.expect(RowLogSoftmax3x5F32.specialization.reductionReuseMatches(1, .{ .reduction = "row_exp_sum", .extents = &.{3} }));
904 try std.testing.expectEqual(@as(u64, 30), RowLogSoftmax3x5F32.specialization.estimatedElementOps().?);
905 try std.testing.expectEqualDeep(RowLogSoftmax3x5F32.specialization.launch.?, RowLogSoftmax3x5F32.specialization.schedule.?.launch());
906
907 try std.testing.expect(RowRmsNorm3x5F32.specialization.operationIs(.{ .row_normalization = .{ .rmsnorm = .scale } }));
908 try std.testing.expectEqual(@as(usize, 2), RowRmsNorm3x5F32.specialization.inputs.len);
909 try std.testing.expectEqual(@as(u64, 15), RowRmsNorm3x5F32.specialization.inputs[0].elementCount().?);
910 try std.testing.expectEqual(@as(u64, 5), RowRmsNorm3x5F32.specialization.inputs[1].elementCount().?);
911 try std.testing.expectEqual(@as(u64, 15), RowRmsNorm3x5F32.specialization.outputs[0].elementCount().?);
912 try std.testing.expectEqualStrings("row_sum_squares", RowRmsNorm3x5F32.specialization.reductions[0].name);
913 try std.testing.expectEqual(entry.ReductionOperator.sum_squares, RowRmsNorm3x5F32.specialization.reductions[0].operator);
914 try std.testing.expectEqual(@as(u64, 5), RowRmsNorm3x5F32.specialization.reductions[0].shape.elementCount().?);
915 try std.testing.expect(RowRmsNorm3x5F32.specialization.reductionReuseScopesAreValid());
916 try std.testing.expectEqual(@as(usize, 1), RowRmsNorm3x5F32.specialization.reduction_reuse.len);
917 try std.testing.expect(RowRmsNorm3x5F32.specialization.reductionReuseMatches(0, .{ .reduction = "row_sum_squares", .extents = &.{3} }));
918 try std.testing.expectEqual(@as(u64, 15), RowRmsNorm3x5F32.specialization.estimatedElementOps().?);
919 try std.testing.expectEqualDeep(RowRmsNorm3x5F32.specialization.launch.?, RowRmsNorm3x5F32.specialization.schedule.?.launch());
920
921 try std.testing.expect(RowResidualRmsNorm3x5F32.specialization.operationIs(.{ .row_normalization = .{ .rmsnorm = .scale } }));
922 try std.testing.expectEqual(@as(usize, 3), RowResidualRmsNorm3x5F32.specialization.inputs.len);
923 try std.testing.expectEqual(@as(u64, 15), RowResidualRmsNorm3x5F32.specialization.inputs[0].elementCount().?);
924 try std.testing.expectEqual(@as(u64, 15), RowResidualRmsNorm3x5F32.specialization.inputs[1].elementCount().?);
925 try std.testing.expectEqual(@as(u64, 5), RowResidualRmsNorm3x5F32.specialization.inputs[2].elementCount().?);
926 try std.testing.expectEqual(@as(u64, 15), RowResidualRmsNorm3x5F32.specialization.outputs[0].elementCount().?);
927 try std.testing.expectEqual(@as(usize, 1), RowResidualRmsNorm3x5F32.specialization.input_transforms.len);
928 try std.testing.expect(RowResidualRmsNorm3x5F32.specialization.inputTransformMatches(0, .{
929 .operator = .residual_add,
930 .input_index = 1,
931 .extents = &.{ 3, 5 },
932 }));
933 try std.testing.expectEqualStrings("row_sum_squares", RowResidualRmsNorm3x5F32.specialization.reductions[0].name);
934 try std.testing.expectEqual(entry.ReductionOperator.sum_squares, RowResidualRmsNorm3x5F32.specialization.reductions[0].operator);
935 try std.testing.expectEqual(@as(u64, 5), RowResidualRmsNorm3x5F32.specialization.reductions[0].shape.elementCount().?);
936 try std.testing.expect(RowResidualRmsNorm3x5F32.specialization.reductionReuseScopesAreValid());
937 try std.testing.expectEqual(@as(usize, 1), RowResidualRmsNorm3x5F32.specialization.reduction_reuse.len);
938 try std.testing.expect(RowResidualRmsNorm3x5F32.specialization.reductionReuseMatches(0, .{ .reduction = "row_sum_squares", .extents = &.{3} }));
939 try std.testing.expectEqual(@as(u64, 15), RowResidualRmsNorm3x5F32.specialization.estimatedElementOps().?);
940 try std.testing.expectEqualDeep(RowResidualRmsNorm3x5F32.specialization.launch.?, RowResidualRmsNorm3x5F32.specialization.schedule.?.launch());
941
942 try std.testing.expect(RowLayerNorm3x5F32.specialization.operationIs(.{ .row_normalization = .{ .layernorm = .none } }));
943 try std.testing.expectEqual(@as(usize, 1), RowLayerNorm3x5F32.specialization.inputs.len);
944 try std.testing.expectEqual(@as(u64, 15), RowLayerNorm3x5F32.specialization.inputs[0].elementCount().?);
945 try std.testing.expectEqual(@as(u64, 15), RowLayerNorm3x5F32.specialization.outputs[0].elementCount().?);
946 try std.testing.expectEqual(@as(usize, 2), RowLayerNorm3x5F32.specialization.reductions.len);
947 try std.testing.expectEqualStrings("row_sum", RowLayerNorm3x5F32.specialization.reductions[0].name);
948 try std.testing.expectEqual(entry.ReductionOperator.sum, RowLayerNorm3x5F32.specialization.reductions[0].operator);
949 try std.testing.expectEqual(@as(u64, 5), RowLayerNorm3x5F32.specialization.reductions[0].shape.elementCount().?);
950 try std.testing.expectEqualStrings("row_variance_sum", RowLayerNorm3x5F32.specialization.reductions[1].name);
951 try std.testing.expectEqual(entry.ReductionOperator.sum_squared_difference, RowLayerNorm3x5F32.specialization.reductions[1].operator);
952 try std.testing.expectEqual(@as(usize, 1), RowLayerNorm3x5F32.specialization.reductions[1].dependencies.len);
953 try std.testing.expectEqualStrings("row_sum", RowLayerNorm3x5F32.specialization.reductions[1].dependencies[0]);
954 try std.testing.expect(RowLayerNorm3x5F32.specialization.reductionReuseScopesAreValid());
955 try std.testing.expectEqual(@as(usize, 2), RowLayerNorm3x5F32.specialization.reduction_reuse.len);
956 try std.testing.expect(RowLayerNorm3x5F32.specialization.reductionReuseMatches(0, .{ .reduction = "row_sum", .extents = &.{3} }));
957 try std.testing.expect(RowLayerNorm3x5F32.specialization.reductionReuseMatches(1, .{ .reduction = "row_variance_sum", .extents = &.{3} }));
958 try std.testing.expectEqual(@as(u64, 30), RowLayerNorm3x5F32.specialization.estimatedElementOps().?);
959 try std.testing.expectEqualDeep(RowLayerNorm3x5F32.specialization.launch.?, RowLayerNorm3x5F32.specialization.schedule.?.launch());
960
961 try std.testing.expect(RowAffineLayerNorm3x5F32.specialization.operationIs(.{ .row_normalization = .{ .layernorm = .scale_bias } }));
962 try std.testing.expectEqual(@as(usize, 3), RowAffineLayerNorm3x5F32.specialization.inputs.len);
963 try std.testing.expectEqual(@as(u64, 15), RowAffineLayerNorm3x5F32.specialization.inputs[0].elementCount().?);
964 try std.testing.expectEqual(@as(u64, 5), RowAffineLayerNorm3x5F32.specialization.inputs[1].elementCount().?);
965 try std.testing.expectEqual(@as(u64, 5), RowAffineLayerNorm3x5F32.specialization.inputs[2].elementCount().?);
966 try std.testing.expectEqual(@as(u64, 15), RowAffineLayerNorm3x5F32.specialization.outputs[0].elementCount().?);
967 try std.testing.expectEqual(@as(usize, 2), RowAffineLayerNorm3x5F32.specialization.reductions.len);
968 try std.testing.expectEqualStrings("row_sum", RowAffineLayerNorm3x5F32.specialization.reductions[0].name);
969 try std.testing.expectEqual(entry.ReductionOperator.sum, RowAffineLayerNorm3x5F32.specialization.reductions[0].operator);
970 try std.testing.expectEqualStrings("row_variance_sum", RowAffineLayerNorm3x5F32.specialization.reductions[1].name);
971 try std.testing.expectEqual(entry.ReductionOperator.sum_squared_difference, RowAffineLayerNorm3x5F32.specialization.reductions[1].operator);
972 try std.testing.expectEqual(@as(usize, 1), RowAffineLayerNorm3x5F32.specialization.reductions[1].dependencies.len);
973 try std.testing.expectEqualStrings("row_sum", RowAffineLayerNorm3x5F32.specialization.reductions[1].dependencies[0]);
974 try std.testing.expect(RowAffineLayerNorm3x5F32.specialization.reductionReuseScopesAreValid());
975 try std.testing.expectEqual(@as(usize, 2), RowAffineLayerNorm3x5F32.specialization.reduction_reuse.len);
976 try std.testing.expect(RowAffineLayerNorm3x5F32.specialization.reductionReuseMatches(0, .{ .reduction = "row_sum", .extents = &.{3} }));
977 try std.testing.expect(RowAffineLayerNorm3x5F32.specialization.reductionReuseMatches(1, .{ .reduction = "row_variance_sum", .extents = &.{3} }));
978 try std.testing.expectEqual(@as(u64, 30), RowAffineLayerNorm3x5F32.specialization.estimatedElementOps().?);
979 try std.testing.expectEqualDeep(RowAffineLayerNorm3x5F32.specialization.launch.?, RowAffineLayerNorm3x5F32.specialization.schedule.?.launch());
980 }
981
982 test "normalization row softmax entry creates registry-ready artifact" {
983 const allocator = std.testing.allocator;
984 var state = gpu.recording.BackendState{
985 .allocator = allocator,
986 .kind = .cuda,
987 .format = .cuda_ptx,
988 };
989
990 var call_artifact = try RowSoftmax2x4F32.createKernelCallArtifact(allocator, state.handle(), .{ .limits = RowSoftmax2x4F32.Limits.testing });
991 defer call_artifact.deinit();
992
993 const artifact = call_artifact.registry().find(RowSoftmax2x4F32.target, RowSoftmax2x4F32.version, .cuda_ptx) orelse {
994 return error.TestExpectedKernelCallArtifact;
995 };
996 try std.testing.expectEqualStrings(RowSoftmax2x4F32.name, artifact.entry_name);
997 try std.testing.expectEqual(@as(u32, 2), artifact.argument_count);
998 try std.testing.expectEqual(gpu.DTypeSet.init(&.{.f32}).bits, artifact.required_dtypes.bits);
999 switch (artifact.launch) {
1000 .fixed => |geometry| {
1001 try std.testing.expectEqual(RowSoftmax2x4F32.specialization.launch.?.grid[0], geometry.grid[0]);
1002 try std.testing.expectEqual(RowSoftmax2x4F32.specialization.launch.?.grid[1], geometry.grid[1]);
1003 try std.testing.expectEqual(RowSoftmax2x4F32.specialization.launch.?.threadgroup[0], geometry.threadgroup[0]);
1004 try std.testing.expectEqual(RowSoftmax2x4F32.specialization.launch.?.threadgroup[1], geometry.threadgroup[1]);
1005 },
1006 else => return error.TestExpectedFixedLaunch,
1007 }
1008 }
1009
1010 test "normalization row log softmax entry creates registry-ready artifact" {
1011 const allocator = std.testing.allocator;
1012 var state = gpu.recording.BackendState{
1013 .allocator = allocator,
1014 .kind = .cuda,
1015 .format = .cuda_ptx,
1016 };
1017
1018 var call_artifact = try RowLogSoftmax2x4F32.createKernelCallArtifact(allocator, state.handle(), .{ .limits = RowLogSoftmax2x4F32.Limits.testing });
1019 defer call_artifact.deinit();
1020
1021 const artifact = call_artifact.registry().find(RowLogSoftmax2x4F32.target, RowLogSoftmax2x4F32.version, .cuda_ptx) orelse {
1022 return error.TestExpectedKernelCallArtifact;
1023 };
1024 try std.testing.expectEqualStrings(RowLogSoftmax2x4F32.name, artifact.entry_name);
1025 try std.testing.expectEqual(@as(u32, 2), artifact.argument_count);
1026 try std.testing.expectEqual(gpu.DTypeSet.init(&.{.f32}).bits, artifact.required_dtypes.bits);
1027 switch (artifact.launch) {
1028 .fixed => |geometry| {
1029 try std.testing.expectEqual(RowLogSoftmax2x4F32.specialization.launch.?.grid[0], geometry.grid[0]);
1030 try std.testing.expectEqual(RowLogSoftmax2x4F32.specialization.launch.?.grid[1], geometry.grid[1]);
1031 try std.testing.expectEqual(RowLogSoftmax2x4F32.specialization.launch.?.threadgroup[0], geometry.threadgroup[0]);
1032 try std.testing.expectEqual(RowLogSoftmax2x4F32.specialization.launch.?.threadgroup[1], geometry.threadgroup[1]);
1033 },
1034 else => return error.TestExpectedFixedLaunch,
1035 }
1036 }
1037
1038 test "normalization row rmsnorm entry creates registry-ready artifact" {
1039 const allocator = std.testing.allocator;
1040 var state = gpu.recording.BackendState{
1041 .allocator = allocator,
1042 .kind = .cuda,
1043 .format = .cuda_ptx,
1044 };
1045
1046 var call_artifact = try RowRmsNorm2x4F32.createKernelCallArtifact(allocator, state.handle(), .{ .limits = RowRmsNorm2x4F32.Limits.testing });
1047 defer call_artifact.deinit();
1048
1049 const artifact = call_artifact.registry().find(RowRmsNorm2x4F32.target, RowRmsNorm2x4F32.version, .cuda_ptx) orelse {
1050 return error.TestExpectedKernelCallArtifact;
1051 };
1052 try std.testing.expectEqualStrings(RowRmsNorm2x4F32.name, artifact.entry_name);
1053 try std.testing.expectEqual(@as(u32, 4), artifact.argument_count);
1054 try std.testing.expectEqual(gpu.DTypeSet.init(&.{.f32}).bits, artifact.required_dtypes.bits);
1055 switch (artifact.launch) {
1056 .fixed => |geometry| {
1057 try std.testing.expectEqual(RowRmsNorm2x4F32.specialization.launch.?.grid[0], geometry.grid[0]);
1058 try std.testing.expectEqual(RowRmsNorm2x4F32.specialization.launch.?.grid[1], geometry.grid[1]);
1059 try std.testing.expectEqual(RowRmsNorm2x4F32.specialization.launch.?.threadgroup[0], geometry.threadgroup[0]);
1060 try std.testing.expectEqual(RowRmsNorm2x4F32.specialization.launch.?.threadgroup[1], geometry.threadgroup[1]);
1061 },
1062 else => return error.TestExpectedFixedLaunch,
1063 }
1064 }
1065
1066 test "fused row residual rmsnorm entry creates registry-ready artifact" {
1067 const allocator = std.testing.allocator;
1068 var state = gpu.recording.BackendState{
1069 .allocator = allocator,
1070 .kind = .cuda,
1071 .format = .cuda_ptx,
1072 };
1073
1074 var call_artifact = try RowResidualRmsNorm2x4F32.createKernelCallArtifact(allocator, state.handle(), .{ .limits = RowResidualRmsNorm2x4F32.Limits.testing });
1075 defer call_artifact.deinit();
1076
1077 const artifact = call_artifact.registry().find(RowResidualRmsNorm2x4F32.target, RowResidualRmsNorm2x4F32.version, .cuda_ptx) orelse {
1078 return error.TestExpectedKernelCallArtifact;
1079 };
1080 try std.testing.expectEqualStrings(RowResidualRmsNorm2x4F32.name, artifact.entry_name);
1081 try std.testing.expectEqual(@as(u32, 5), artifact.argument_count);
1082 try std.testing.expectEqual(gpu.DTypeSet.init(&.{.f32}).bits, artifact.required_dtypes.bits);
1083 switch (artifact.launch) {
1084 .fixed => |geometry| {
1085 try std.testing.expectEqual(RowResidualRmsNorm2x4F32.specialization.launch.?.grid[0], geometry.grid[0]);
1086 try std.testing.expectEqual(RowResidualRmsNorm2x4F32.specialization.launch.?.grid[1], geometry.grid[1]);
1087 try std.testing.expectEqual(RowResidualRmsNorm2x4F32.specialization.launch.?.threadgroup[0], geometry.threadgroup[0]);
1088 try std.testing.expectEqual(RowResidualRmsNorm2x4F32.specialization.launch.?.threadgroup[1], geometry.threadgroup[1]);
1089 },
1090 else => return error.TestExpectedFixedLaunch,
1091 }
1092 }
1093
1094 test "normalization row layernorm entry creates registry-ready artifact" {
1095 const allocator = std.testing.allocator;
1096 var state = gpu.recording.BackendState{
1097 .allocator = allocator,
1098 .kind = .cuda,
1099 .format = .cuda_ptx,
1100 };
1101
1102 var call_artifact = try RowLayerNorm2x4F32.createKernelCallArtifact(allocator, state.handle(), .{ .limits = RowLayerNorm2x4F32.Limits.testing });
1103 defer call_artifact.deinit();
1104
1105 const artifact = call_artifact.registry().find(RowLayerNorm2x4F32.target, RowLayerNorm2x4F32.version, .cuda_ptx) orelse {
1106 return error.TestExpectedKernelCallArtifact;
1107 };
1108 try std.testing.expectEqualStrings(RowLayerNorm2x4F32.name, artifact.entry_name);
1109 try std.testing.expectEqual(@as(u32, 3), artifact.argument_count);
1110 try std.testing.expectEqual(gpu.DTypeSet.init(&.{.f32}).bits, artifact.required_dtypes.bits);
1111 switch (artifact.launch) {
1112 .fixed => |geometry| {
1113 try std.testing.expectEqual(RowLayerNorm2x4F32.specialization.launch.?.grid[0], geometry.grid[0]);
1114 try std.testing.expectEqual(RowLayerNorm2x4F32.specialization.launch.?.grid[1], geometry.grid[1]);
1115 try std.testing.expectEqual(RowLayerNorm2x4F32.specialization.launch.?.threadgroup[0], geometry.threadgroup[0]);
1116 try std.testing.expectEqual(RowLayerNorm2x4F32.specialization.launch.?.threadgroup[1], geometry.threadgroup[1]);
1117 },
1118 else => return error.TestExpectedFixedLaunch,
1119 }
1120 }
1121
1122 test "normalization row affine layernorm entry creates registry-ready artifact" {
1123 const allocator = std.testing.allocator;
1124 var state = gpu.recording.BackendState{
1125 .allocator = allocator,
1126 .kind = .cuda,
1127 .format = .cuda_ptx,
1128 };
1129
1130 var call_artifact = try RowAffineLayerNorm2x4F32.createKernelCallArtifact(allocator, state.handle(), .{ .limits = RowAffineLayerNorm2x4F32.Limits.testing });
1131 defer call_artifact.deinit();
1132
1133 const artifact = call_artifact.registry().find(RowAffineLayerNorm2x4F32.target, RowAffineLayerNorm2x4F32.version, .cuda_ptx) orelse {
1134 return error.TestExpectedKernelCallArtifact;
1135 };
1136 try std.testing.expectEqualStrings(RowAffineLayerNorm2x4F32.name, artifact.entry_name);
1137 try std.testing.expectEqual(@as(u32, 5), artifact.argument_count);
1138 try std.testing.expectEqual(gpu.DTypeSet.init(&.{.f32}).bits, artifact.required_dtypes.bits);
1139 switch (artifact.launch) {
1140 .fixed => |geometry| {
1141 try std.testing.expectEqual(RowAffineLayerNorm2x4F32.specialization.launch.?.grid[0], geometry.grid[0]);
1142 try std.testing.expectEqual(RowAffineLayerNorm2x4F32.specialization.launch.?.grid[1], geometry.grid[1]);
1143 try std.testing.expectEqual(RowAffineLayerNorm2x4F32.specialization.launch.?.threadgroup[0], geometry.threadgroup[0]);
1144 try std.testing.expectEqual(RowAffineLayerNorm2x4F32.specialization.launch.?.threadgroup[1], geometry.threadgroup[1]);
1145 },
1146 else => return error.TestExpectedFixedLaunch,
1147 }
1148 }