lib/accy/src/kernel/library/stencil.zig
daab053ee43316e1809a84551d573ddd1e5bf3d2
1 const std = @import("std");
2 const gpu = @import("gpu");
3 const choir_abi = @import("choir_abi");
4
5 const artifact_product = @import("../../artifact/model/root.zig");
6 const shape = @import("../../choir/shape/root.zig");
7 const entry = @import("entry.zig");
8 const extent_mod = @import("extent.zig");
9 const geometry_mod = @import("geometry.zig");
10 const kernel = @import("../root.zig");
11 const tuning = @import("tuning.zig");
12
13 const DType = choir_abi.DType;
14 const indexExtent = extent_mod.indexExtent;
15 const runtimeExtentArgument = extent_mod.runtimeExtentArgument;
16
17 pub const Window = struct {
18 rows: u64,
19 cols: u64,
20 radius: u32 = 1,
21 dtype: DType = .f32,
22 accumulation_dtype: DType = .f32,
23 threads: entry.Threads2D = .{},
24 row_axis: []const u8 = "r",
25 col_axis: []const u8 = "c",
26 window_axis: []const u8 = "w",
27
28 pub fn side(self: Window) u64 {
29 return 2 * @as(u64, self.radius) + 1;
30 }
31
32 pub fn taps(self: Window) u64 {
33 return self.side() * self.side();
34 }
35
36 pub fn paddedRows(self: Window) u64 {
37 return self.rows + 2 * @as(u64, self.radius);
38 }
39
40 pub fn paddedCols(self: Window) u64 {
41 return self.cols + 2 * @as(u64, self.radius);
42 }
43 };
44
45 pub const window_family_version: u32 = 1;
46 pub const window_radius_max: u32 = 3;
47 const window_thread_caps = geometry_mod.ThreadCaps{
48 .budget = 256,
49 .x_max = 64,
50 .y_max = 16,
51 };
52
53 pub fn windowAccumulationDType(dtype: DType) ?DType {
54 return switch (dtype) {
55 .f32, .f16 => .f32,
56 else => null,
57 };
58 }
59
60 pub fn windowRadiusValid(radius: u32) bool {
61 return radius >= 1 and radius <= window_radius_max;
62 }
63
64 fn windowAccumulationZero(inner: anytype, spec: Window) !kernel.Value {
65 return switch (spec.accumulation_dtype) {
66 .f32 => inner.constantFloat(.f32, 0.0),
67 .f16 => inner.constantFloat(.f16, 0.0),
68 else => error.UnsupportedDType,
69 };
70 }
71
72 fn windowAccumulationValue(inner: anytype, spec: Window, value: anytype) !kernel.Value {
73 return switch (spec.accumulation_dtype) {
74 .f32 => if (comptime @TypeOf(value).scalar_dtype == .f32) value.raw() else (try value.cast(inner, .f32)).raw(),
75 .f16 => if (comptime @TypeOf(value).scalar_dtype == .f16) value.raw() else (try value.cast(inner, .f16)).raw(),
76 else => error.UnsupportedDType,
77 };
78 }
79
80 fn windowOutputValue(inner: anytype, spec: Window, value: kernel.Value) !kernel.Value {
81 if (spec.dtype == spec.accumulation_dtype) return value;
82 return switch (spec.dtype) {
83 .f32 => inner.cast(value, .f32),
84 .f16 => inner.cast(value, .f16),
85 else => error.UnsupportedDType,
86 };
87 }
88
89 pub fn windowCellSum(
90 inner: anytype,
91 spec: Window,
92 src: anytype,
93 weights: anytype,
94 row: kernel.Value,
95 col: kernel.Value,
96 padded_cols: kernel.Value,
97 ) !kernel.Value {
98 var acc = try windowAccumulationZero(inner, spec);
99 const side_extent = spec.side();
100 var dr: u64 = 0;
101 while (dr < side_extent) : (dr += 1) {
102 var dc: u64 = 0;
103 while (dc < side_extent) : (dc += 1) {
104 const dr_value = try inner.constantIndex(try indexExtent(dr));
105 const dc_value = try inner.constantIndex(try indexExtent(dc));
106 const tap_value = try inner.constantIndex(try indexExtent(dr * side_extent + dc));
107 const src_row = try inner.add(row, dr_value);
108 const src_col = try inner.add(col, dc_value);
109 const src_row_offset = try inner.mul(src_row, padded_cols);
110 const src_index = try inner.add(src_row_offset, src_col);
111 const src_value = try src.load(inner, src_index);
112 const weight_value = try weights.load(inner, tap_value);
113 const src_acc = try windowAccumulationValue(inner, spec, src_value);
114 const weight_acc = try windowAccumulationValue(inner, spec, weight_value);
115 const product = try inner.mul(src_acc, weight_acc);
116 acc = try inner.add(acc, product);
117 }
118 }
119 return acc;
120 }
121
122 fn window_body_each(inner: anytype, index: kernel.Index2D, ctx: anytype) !void {
123 const padded_cols = try inner.constantIndex(try indexExtent(ctx.spec.paddedCols()));
124 const cols_stride = try inner.constantIndex(try indexExtent(ctx.spec.cols));
125 const sum = try windowCellSum(
126 inner,
127 ctx.spec,
128 ctx.args.param(.src),
129 ctx.args.param(.weights),
130 index.y.index,
131 index.x.index,
132 padded_cols,
133 );
134 const out_row_offset = try inner.mul(index.y.index, cols_stride);
135 const out_index = try inner.add(out_row_offset, index.x.index);
136 try ctx.args.param(.dst).store(inner, try windowOutputValue(inner, ctx.spec, sum), out_index);
137 }
138
139 fn windowBody(k: anytype, spec: Window, args: anytype) !void {
140 _ = try k.forEach2D(.{
141 .x = kernel.logical.axis(spec.col_axis, spec.cols),
142 .y = kernel.logical.axis(spec.row_axis, spec.rows),
143 }, .{ .spec = spec, .args = args }, window_body_each);
144 }
145
146 fn window_runtime_body_row_active(inner: anytype, ctx: anytype) !void {
147 const col_active = try inner.compare(.lt, ctx.col, ctx.cols_extent);
148 try inner.guardDo(col_active, ctx, window_runtime_body_col_active);
149 }
150
151 fn window_runtime_body_col_active(active_inner: anytype, active_ctx: anytype) !void {
152 const sum = try windowCellSum(
153 active_inner,
154 active_ctx.spec,
155 active_ctx.args.param(.src),
156 active_ctx.args.param(.weights),
157 active_ctx.row,
158 active_ctx.col,
159 active_ctx.padded_cols,
160 );
161 const out_row_offset = try active_inner.mul(active_ctx.row, active_ctx.cols_extent);
162 const out_index = try active_inner.add(out_row_offset, active_ctx.col);
163 try active_ctx.args.param(.dst).store(
164 active_inner,
165 try windowOutputValue(active_inner, active_ctx.spec, sum),
166 out_index,
167 );
168 }
169
170 fn windowRuntimeBody(k: anytype, spec: Window, args: anytype) !void {
171 const row = try k.globalId(.y);
172 const col = try k.globalId(.x);
173 const rows_extent = try k.castIndex(args.param(.rows).raw());
174 const cols_extent = try k.castIndex(args.param(.cols).raw());
175 const halo = try k.constantIndex(try indexExtent(2 * @as(u64, spec.radius)));
176 const padded_cols = try k.add(cols_extent, halo);
177 const row_active = try k.compare(.lt, row, rows_extent);
178 try k.guardDo(row_active, .{
179 .args = args,
180 .spec = spec,
181 .row = row,
182 .col = col,
183 .cols_extent = cols_extent,
184 .padded_cols = padded_cols,
185 }, window_runtime_body_row_active);
186 }
187
188 fn windowFamilySchedule(instance: Window) kernel.logical.schedule.ThreadBlocks {
189 return kernel.logical.schedule.threadBlocks(.{
190 .x = instance.threads.x,
191 .y = instance.threads.y,
192 });
193 }
194
195 fn windowFamily(comptime dtype: DType) type {
196 return kernel.logical.Family(.{
197 .name = std.fmt.comptimePrint("accy_kernel_stencil_window_{s}", .{dtype.name()}),
198 .parameters = .{
199 .dst = kernel.dynamicBuffer(dtype),
200 .src = kernel.dynamicBuffer(dtype),
201 .weights = kernel.dynamicBuffer(dtype),
202 },
203 .Instance = Window,
204 .schedule = windowFamilySchedule,
205 .body = windowBody,
206 });
207 }
208
209 fn windowRuntimeFamily(comptime dtype: DType) type {
210 return kernel.logical.Family(.{
211 .name = std.fmt.comptimePrint("accy_kernel_stencil_window_runtime_{s}", .{dtype.name()}),
212 .parameters = .{
213 .dst = kernel.dynamicBuffer(dtype),
214 .src = kernel.dynamicBuffer(dtype),
215 .weights = kernel.dynamicBuffer(dtype),
216 .rows = kernel.scalar(.i32),
217 .cols = kernel.scalar(.i32),
218 },
219 .Instance = Window,
220 .schedule = windowFamilySchedule,
221 .body = windowRuntimeBody,
222 });
223 }
224
225 pub const WindowFamilyF32 = windowFamily(.f32);
226 pub const WindowFamilyF16 = windowFamily(.f16);
227 pub const WindowRuntimeFamilyF32 = windowRuntimeFamily(.f32);
228 pub const WindowRuntimeFamilyF16 = windowRuntimeFamily(.f16);
229
230 pub fn windowThreadsForExtents(rows: u64, cols: u64) entry.Threads2D {
231 return geometry_mod.threadsForGrid(.{ .rows = rows, .cols = cols }, window_thread_caps);
232 }
233
234 pub fn windowThreadCandidatesForExtents(rows: u64, cols: u64) geometry_mod.ThreadCandidates {
235 return geometry_mod.threadCandidatesForGrid(.{ .rows = rows, .cols = cols }, window_thread_caps);
236 }
237
238 pub fn windowInstanceTarget(allocator: std.mem.Allocator, instance: Window) ![]u8 {
239 return std.fmt.allocPrint(
240 allocator,
241 "accy.kernel.stencil.window{d}x{d}_r{d}_{d}x{d}_{s}",
242 .{ instance.rows, instance.cols, instance.radius, instance.threads.x, instance.threads.y, instance.dtype.name() },
243 );
244 }
245
246 pub fn windowInstanceEntryName(allocator: std.mem.Allocator, instance: Window) ![]u8 {
247 return std.fmt.allocPrint(
248 allocator,
249 "accy_kernel_stencil_window{d}x{d}_r{d}_{d}x{d}_{s}",
250 .{ instance.rows, instance.cols, instance.radius, instance.threads.x, instance.threads.y, instance.dtype.name() },
251 );
252 }
253
254 pub fn windowFamilyTarget(allocator: std.mem.Allocator, instance: Window) ![]u8 {
255 return std.fmt.allocPrint(
256 allocator,
257 "accy.kernel.stencil.window_family_r{d}_{d}x{d}_{s}",
258 .{ instance.radius, instance.threads.x, instance.threads.y, instance.dtype.name() },
259 );
260 }
261
262 pub fn windowFamilyEntryName(allocator: std.mem.Allocator, instance: Window) ![]u8 {
263 return std.fmt.allocPrint(
264 allocator,
265 "accy_kernel_stencil_window_family_r{d}_{d}x{d}_{s}",
266 .{ instance.radius, instance.threads.x, instance.threads.y, instance.dtype.name() },
267 );
268 }
269
270 pub fn windowTuningExtents(instance: Window) [3]u64 {
271 return .{ instance.rows, instance.cols, instance.radius };
272 }
273
274 pub fn windowTuningOperation(instance: Window) entry.Operation {
275 _ = instance;
276 return .{ .stencil = .window };
277 }
278
279 pub fn windowFamilyTuningKey(
280 backing_allocator: std.mem.Allocator,
281 device_fingerprint: u64,
282 instance: Window,
283 ) !tuning.FamilyTuningKey {
284 const family_fingerprint = try windowFamilyFingerprint(backing_allocator, instance);
285 const extents = windowTuningExtents(instance);
286 return tuning.FamilyTuningKey.init(
287 device_fingerprint,
288 family_fingerprint,
289 entry.operationFingerprint(windowTuningOperation(instance)),
290 instance.dtype,
291 window_family_version,
292 extents[0..],
293 ) orelse unreachable;
294 }
295
296 pub fn windowRuntimeArguments(instance: Window) ![2]choir_abi.ScalarArgument {
297 return .{
298 .{ .u32 = try runtimeExtentArgument(instance.rows) },
299 .{ .u32 = try runtimeExtentArgument(instance.cols) },
300 };
301 }
302
303 pub fn windowShapeProfileDimensions(instance: Window) [2]artifact_product.KernelCallShapeProfileDimension {
304 const bounds = windowRuntimeExtentBounds();
305 return .{
306 .{
307 .name = instance.row_axis,
308 .runtime_scalar_argument_index = 0,
309 .bounds = bounds,
310 },
311 .{
312 .name = instance.col_axis,
313 .runtime_scalar_argument_index = 1,
314 .bounds = bounds,
315 },
316 };
317 }
318
319 fn windowRuntimeExtentBounds() shape.Bounds {
320 return .{ .min = 1, .max = extent_mod.runtime_extent_max };
321 }
322
323 fn windowDerivedLaunch(instance: Window) !artifact_product.KernelCallLaunch {
324 if (instance.threads.x == 0 or instance.threads.y == 0) return error.KernelLibraryLaunchThreadgroupMustBeNonzero;
325 return .{ .derived = .{
326 .grid = .{
327 .{ .runtime_u32_ceil_div = .{ .argument_index = 1, .divisor = instance.threads.x } },
328 .{ .runtime_u32_ceil_div = .{ .argument_index = 0, .divisor = instance.threads.y } },
329 .{ .fixed = 1 },
330 },
331 .threadgroup = .{ instance.threads.x, instance.threads.y, 1 },
332 } };
333 }
334
335 pub fn createWindowFamilyArtifact(
336 allocator: std.mem.Allocator,
337 handle: kernel.BackendHandle,
338 instance: Window,
339 options: entry.ArtifactOptions,
340 ) !kernel.OwnedKernelCallArtifact {
341 if (!windowRadiusValid(instance.radius)) return error.StencilRadiusOutOfRange;
342 const target = try windowFamilyTarget(allocator, instance);
343 defer allocator.free(target);
344 const entry_name = try windowFamilyEntryName(allocator, instance);
345 defer allocator.free(entry_name);
346 const family_fingerprint = options.shape_family_fingerprint orelse try windowFamilyFingerprint(allocator, instance);
347 const shape_profile_dimensions = windowShapeProfileDimensions(instance);
348 const shape_profile = options.shape_profile orelse artifact_product.KernelCallShapeProfile{
349 .name = "stencil_window",
350 .fingerprint = family_fingerprint,
351 .dimensions = shape_profile_dimensions[0..],
352 };
353
354 var graph = switch (instance.dtype) {
355 .f32 => try WindowRuntimeFamilyF32.buildNamed(allocator, options.limits, entry_name, instance),
356 .f16 => try WindowRuntimeFamilyF16.buildNamed(allocator, options.limits, entry_name, instance),
357 else => return error.UnsupportedDType,
358 };
359 defer graph.deinit();
360 return kernel.createKernelCallArtifact(allocator, handle, &graph, .{
361 .target = target,
362 .version = window_family_version,
363 .format = options.format,
364 .kernel_plan = options.kernel_plan,
365 .element_count_argument = options.element_count_argument,
366 .shape_family_fingerprint = family_fingerprint,
367 .shape_profile = shape_profile,
368 .launch = options.launch orelse try windowDerivedLaunch(instance),
369 .runtime_scalar_argument_count = if (options.runtime_scalar_argument_count == 0) 2 else options.runtime_scalar_argument_count,
370 .static_arguments = options.static_arguments,
371 });
372 }
373
374 pub fn windowFamilyFingerprint(backing_allocator: std.mem.Allocator, instance: Window) !u64 {
375 var family = try windowShapeFamily(backing_allocator, instance);
376 defer family.deinit();
377 return shape.fingerprint(family);
378 }
379
380 pub fn windowShapeFamily(backing_allocator: std.mem.Allocator, instance: Window) !shape.Family {
381 if (!windowRadiusValid(instance.radius)) return error.StencilRadiusOutOfRange;
382 var builder = try shape.Builder.init(backing_allocator, "stencil_window");
383 errdefer builder.deinit();
384
385 const rows = try builder.symbol(instance.row_axis);
386 const cols = try builder.symbol(instance.col_axis);
387
388 const rows_expr = try builder.symbolExpression(rows);
389 const cols_expr = try builder.symbolExpression(cols);
390 const halo_expr = builder.constantExpression(@intCast(2 * @as(u64, instance.radius)));
391 const padded_rows_expr = try builder.addExpression(rows_expr, halo_expr);
392 const padded_cols_expr = try builder.addExpression(cols_expr, halo_expr);
393 const taps_expr = builder.constantExpression(try indexExtent(instance.taps()));
394
395 _ = try builder.tensor("src", &.{ padded_rows_expr, padded_cols_expr });
396 _ = try builder.tensor("weights", &.{taps_expr});
397 _ = try builder.tensor("out", &.{ rows_expr, cols_expr });
398 try builder.assumeBounds(rows_expr, windowRuntimeExtentBounds());
399 try builder.assumeBounds(cols_expr, windowRuntimeExtentBounds());
400
401 return builder.finish();
402 }
403
404 pub fn windowFamilySpecialization(backing_allocator: std.mem.Allocator, instance: Window) !entry.OwnedSpecialization {
405 if (!windowRadiusValid(instance.radius)) return error.StencilRadiusOutOfRange;
406 var owned = entry.OwnedSpecialization.init(backing_allocator);
407 errdefer owned.deinit();
408 const lifetime_allocator = owned.allocator();
409
410 const inputs = try lifetime_allocator.alloc(entry.Shape, 2);
411 inputs[0] = try entry.runtimeShape2D(
412 lifetime_allocator,
413 instance.row_axis,
414 instance.paddedRows(),
415 instance.col_axis,
416 instance.paddedCols(),
417 );
418 inputs[1] = try entry.runtimeShape1D(lifetime_allocator, instance.window_axis, instance.taps());
419
420 const outputs = try lifetime_allocator.alloc(entry.Shape, 1);
421 outputs[0] = try entry.runtimeShape2D(lifetime_allocator, instance.row_axis, instance.rows, instance.col_axis, instance.cols);
422
423 const reductions = try lifetime_allocator.alloc(entry.Reduction, 1);
424 reductions[0] = try entry.runtimeReduction(
425 lifetime_allocator,
426 "window",
427 .weighted_sum,
428 try entry.runtimeShape1D(lifetime_allocator, instance.window_axis, instance.taps()),
429 );
430
431 owned.value = .{
432 .dtype = instance.dtype,
433 .accumulation_dtype = instance.accumulation_dtype,
434 .operation = .{ .stencil = .window },
435 .inputs = inputs,
436 .outputs = outputs,
437 .reductions = reductions,
438 .schedule = try entry.runtimeThreadBlocks2D(
439 lifetime_allocator,
440 instance.col_axis,
441 instance.cols,
442 instance.row_axis,
443 instance.rows,
444 instance.threads.x,
445 instance.threads.y,
446 ),
447 };
448 owned.value.launch = owned.value.schedule.?.launch();
449 var family = try windowShapeFamily(backing_allocator, instance);
450 errdefer family.deinit();
451 try owned.takeShapeFamily(&family);
452 return owned;
453 }
454
455 pub fn windowInstanceFromSpecialization(specialization: entry.Specialization) ?Window {
456 if (!specialization.scheduleMatchesLaunch()) return null;
457 if (!specialization.operationIs(.{ .stencil = .window })) return null;
458 const dtype = specialization.dtype orelse return null;
459 const accumulation_dtype = specialization.accumulation_dtype orelse return null;
460 if (windowAccumulationDType(dtype) != accumulation_dtype) return null;
461 if (specialization.inputs.len != 2 or specialization.outputs.len != 1 or specialization.reductions.len != 1) return null;
462 const src = specialization.inputs[0];
463 const weights = specialization.inputs[1];
464 const output = specialization.outputs[0];
465 const reduction = specialization.reductions[0];
466 if (src.axes.len != 2 or weights.axes.len != 1 or output.axes.len != 2) return null;
467 if (reduction.shape.axes.len != 1) return null;
468 const rows = output.axes[0].extent;
469 const cols = output.axes[1].extent;
470 if (src.axes[0].extent <= rows or src.axes[1].extent <= cols) return null;
471 const row_halo = src.axes[0].extent - rows;
472 const col_halo = src.axes[1].extent - cols;
473 if (row_halo != col_halo or row_halo % 2 != 0) return null;
474 const radius: u32 = @intCast(row_halo / 2);
475 if (!windowRadiusValid(radius)) return null;
476 const side = 2 * @as(u64, radius) + 1;
477 if (weights.axes[0].extent != side * side) return null;
478 if (!std.mem.eql(u8, src.axes[0].name, output.axes[0].name)) return null;
479 if (!std.mem.eql(u8, src.axes[1].name, output.axes[1].name)) return null;
480 if (!std.mem.eql(u8, reduction.name, "window")) return null;
481 if (reduction.operator != .weighted_sum) return null;
482 if (reduction.shape.axes[0].extent != side * side) return null;
483 if (!std.mem.eql(u8, reduction.shape.axes[0].name, weights.axes[0].name)) return null;
484 const launch = specialization.launch orelse return null;
485 if (launch.threadgroup[0] == 0 or launch.threadgroup[1] == 0) return null;
486 return .{
487 .rows = rows,
488 .cols = cols,
489 .radius = radius,
490 .dtype = dtype,
491 .accumulation_dtype = accumulation_dtype,
492 .threads = .{ .x = launch.threadgroup[0], .y = launch.threadgroup[1] },
493 .row_axis = output.axes[0].name,
494 .col_axis = output.axes[1].name,
495 .window_axis = weights.axes[0].name,
496 };
497 }
498
499 fn windowSpecialization(comptime spec: Window) entry.Specialization {
500 return .{
501 .dtype = spec.dtype,
502 .accumulation_dtype = spec.accumulation_dtype,
503 .operation = .{ .stencil = .window },
504 .inputs = &.{
505 entry.shape2D(spec.row_axis, spec.paddedRows(), spec.col_axis, spec.paddedCols()),
506 entry.shape1D(spec.window_axis, spec.taps()),
507 },
508 .outputs = &.{entry.shape2D(spec.row_axis, spec.rows, spec.col_axis, spec.cols)},
509 .reductions = &.{entry.reduction("window", .weighted_sum, entry.shape1D(spec.window_axis, spec.taps()))},
510 .launch = entry.launch2D(spec.cols, spec.rows, spec.threads.x, spec.threads.y),
511 .schedule = entry.threadBlocks2D(spec.col_axis, spec.cols, spec.row_axis, spec.rows, spec.threads.x, spec.threads.y),
512 };
513 }
514
515 fn windowProgram(comptime spec: Window) type {
516 const Body = struct {
517 fn run(k: anytype, args: anytype) !void {
518 try windowBody(k, spec, args);
519 }
520 };
521
522 return kernel.logical.Program(.{
523 .name = std.fmt.comptimePrint(
524 "accy_kernel_stencil_window{}x{}_r{}_{}x{}_{s}",
525 .{ spec.rows, spec.cols, spec.radius, spec.threads.x, spec.threads.y, spec.dtype.name() },
526 ),
527 .parameters = .{
528 .dst = kernel.dynamicBuffer(spec.dtype),
529 .src = kernel.dynamicBuffer(spec.dtype),
530 .weights = kernel.dynamicBuffer(spec.dtype),
531 },
532 .body = Body.run,
533 }).withSchedule(kernel.logical.schedule.threadBlocks(.{
534 .x = spec.threads.x,
535 .y = spec.threads.y,
536 }));
537 }
538
539 pub fn windowF32(comptime spec: Window) type {
540 if (!windowRadiusValid(spec.radius)) @compileError("kernel library stencil window radius out of range");
541 return entry.Entry(windowProgram(spec), .{
542 .target = std.fmt.comptimePrint(
543 "accy.kernel.stencil.window{}x{}_r{}_{}x{}_{s}",
544 .{ spec.rows, spec.cols, spec.radius, spec.threads.x, spec.threads.y, spec.dtype.name() },
545 ),
546 .layer = .logical,
547 .category = .stencil,
548 .specialization = windowSpecialization(spec),
549 });
550 }
551
552 pub const Window2x3R1F32 = windowF32(.{ .rows = 2, .cols = 3, .radius = 1, .threads = .{ .x = 3, .y = 2 } });
553
554 fn stencilFamilyTuningTestCapabilities(device_id: u32) gpu.BackendCapabilities {
555 return .{ .identity = .{
556 .backend = .cuda,
557 .family = .nvidia_cuda,
558 .name = "stencil-family-tuning-test-device",
559 .vendor_id = 0x10de,
560 .device_id = device_id,
561 } };
562 }
563
564 fn testPaddedInput(comptime count: usize) [count]f32 {
565 var values: [count]f32 = undefined;
566 for (&values, 0..) |*value, index| value.* = @floatFromInt(index);
567 return values;
568 }
569
570 test "stencil window entry runs on CPU" {
571 var src = testPaddedInput(20);
572 var identity_weights = [_]f32{ 0, 0, 0, 0, 1, 0, 0, 0, 0 };
573 var dst = @as([6]f32, @splat(0.0));
574
575 try Window2x3R1F32.runCpu(std.testing.allocator, Window2x3R1F32.Limits.testing, &.{
576 kernel.argumentBuffer(f32, dst[0..]),
577 kernel.argumentBuffer(f32, src[0..]),
578 kernel.argumentBuffer(f32, identity_weights[0..]),
579 });
580 try std.testing.expectEqualSlices(f32, &.{ 6.0, 7.0, 8.0, 11.0, 12.0, 13.0 }, dst[0..]);
581
582 var mixed_weights = [_]f32{ 1, 0, 0, 0, 2, 0, 0, 0, 3 };
583 var mixed_dst = @as([6]f32, @splat(0.0));
584 try Window2x3R1F32.runCpu(std.testing.allocator, Window2x3R1F32.Limits.testing, &.{
585 kernel.argumentBuffer(f32, mixed_dst[0..]),
586 kernel.argumentBuffer(f32, src[0..]),
587 kernel.argumentBuffer(f32, mixed_weights[0..]),
588 });
589 try std.testing.expectEqualSlices(f32, &.{ 48.0, 54.0, 60.0, 78.0, 84.0, 90.0 }, mixed_dst[0..]);
590 }
591
592 test "stencil window runtime family executes explicit runtime extents" {
593 const allocator = std.testing.allocator;
594 const compiled = Window{ .rows = 1, .cols = 1, .radius = 1, .threads = .{ .x = 4, .y = 2 } };
595 const runtime = Window{ .rows = 2, .cols = 3, .radius = 1, .threads = compiled.threads };
596
597 var graph = try WindowRuntimeFamilyF32.build(allocator, WindowRuntimeFamilyF32.Limits.testing, compiled);
598 defer graph.deinit();
599
600 var src = testPaddedInput(20);
601 var weights = [_]f32{ 1, 0, 0, 0, 2, 0, 0, 0, 3 };
602 var dst = @as([6]f32, @splat(0.0));
603
604 const launch_value = try entry.runtimeLaunch2D(runtime.cols, runtime.rows, runtime.threads.x, runtime.threads.y);
605 try graph.runCpuWithLaunch(allocator, &.{
606 kernel.argumentBuffer(f32, dst[0..]),
607 kernel.argumentBuffer(f32, src[0..]),
608 kernel.argumentBuffer(f32, weights[0..]),
609 kernel.argumentI32(@intCast(runtime.rows)),
610 kernel.argumentI32(@intCast(runtime.cols)),
611 }, .{
612 .grid = launch_value.grid,
613 .block = launch_value.threadgroup,
614 });
615 try std.testing.expectEqualSlices(f32, &.{ 48.0, 54.0, 60.0, 78.0, 84.0, 90.0 }, dst[0..]);
616 }
617
618 test "stencil window thread candidates stay legal for output grids" {
619 const candidates = windowThreadCandidatesForExtents(17, 17);
620 try std.testing.expect(candidates.count > 1);
621 for (candidates.slice(), 0..) |candidate, index| {
622 try std.testing.expect(candidate.x != 0);
623 try std.testing.expect(candidate.y != 0);
624 try std.testing.expect(candidate.x * candidate.y <= window_thread_caps.budget);
625 for (candidates.slice()[0..index]) |previous| {
626 try std.testing.expect(!geometry_mod.threadCandidatesEqual(previous, candidate));
627 }
628 }
629 try std.testing.expect(geometry_mod.threadCandidatesEqual(candidates.items[0], windowThreadsForExtents(17, 17)));
630 }
631
632 test "stencil window family instance identity matches fixed entry strings" {
633 const instance = Window{ .rows = 2, .cols = 3, .radius = 1, .threads = .{ .x = 3, .y = 2 } };
634
635 const target = try windowInstanceTarget(std.testing.allocator, instance);
636 defer std.testing.allocator.free(target);
637 try std.testing.expectEqualStrings(Window2x3R1F32.target, target);
638
639 const entry_name = try windowInstanceEntryName(std.testing.allocator, instance);
640 defer std.testing.allocator.free(entry_name);
641 try std.testing.expectEqualStrings(Window2x3R1F32.name, entry_name);
642
643 try std.testing.expectEqual(Window2x3R1F32.version, window_family_version);
644
645 const fresh = Window{ .rows = 64, .cols = 96, .radius = 2, .threads = .{ .x = 8, .y = 4 } };
646 const fresh_target = try windowInstanceTarget(std.testing.allocator, fresh);
647 defer std.testing.allocator.free(fresh_target);
648 try std.testing.expectEqualStrings("accy.kernel.stencil.window64x96_r2_8x4_f32", fresh_target);
649
650 const family_target = try windowFamilyTarget(std.testing.allocator, fresh);
651 defer std.testing.allocator.free(family_target);
652 try std.testing.expectEqualStrings("accy.kernel.stencil.window_family_r2_8x4_f32", family_target);
653
654 const family_entry = try windowFamilyEntryName(std.testing.allocator, fresh);
655 defer std.testing.allocator.free(family_entry);
656 try std.testing.expectEqualStrings("accy_kernel_stencil_window_family_r2_8x4_f32", family_entry);
657
658 const fresh_f16 = Window{ .rows = 64, .cols = 96, .radius = 2, .dtype = .f16, .threads = .{ .x = 8, .y = 4 } };
659 const family_f16_target = try windowFamilyTarget(std.testing.allocator, fresh_f16);
660 defer std.testing.allocator.free(family_f16_target);
661 try std.testing.expectEqualStrings("accy.kernel.stencil.window_family_r2_8x4_f16", family_f16_target);
662 }
663
664 test "stencil window family tuning keys discriminate dtype radius and device" {
665 const allocator = std.testing.allocator;
666 const device = tuning.deviceFingerprint(stencilFamilyTuningTestCapabilities(0x2684));
667
668 const base = try windowFamilyTuningKey(allocator, device, .{ .rows = 8, .cols = 8 });
669 const half = try windowFamilyTuningKey(allocator, device, .{ .rows = 8, .cols = 8, .dtype = .f16, .accumulation_dtype = .f32 });
670 try std.testing.expect(!base.eql(half));
671 try std.testing.expectEqual(base.family_fingerprint, half.family_fingerprint);
672 try std.testing.expectEqual(base.operation_fingerprint, half.operation_fingerprint);
673
674 const wider = try windowFamilyTuningKey(allocator, device, .{ .rows = 8, .cols = 8, .radius = 2 });
675 try std.testing.expect(!base.eql(wider));
676 try std.testing.expect(base.family_fingerprint != wider.family_fingerprint);
677 try std.testing.expectEqual(base.operation_fingerprint, wider.operation_fingerprint);
678
679 const other_device = try windowFamilyTuningKey(
680 allocator,
681 tuning.deviceFingerprint(stencilFamilyTuningTestCapabilities(0x1b80)),
682 .{ .rows = 8, .cols = 8 },
683 );
684 try std.testing.expect(!base.eql(other_device));
685 try std.testing.expectEqual(base.family_fingerprint, other_device.family_fingerprint);
686 try std.testing.expectEqual(base.operation_fingerprint, other_device.operation_fingerprint);
687 }
688
689 test "stencil window family artifact carries runtime launch contract" {
690 const allocator = std.testing.allocator;
691 var state = gpu.recording.BackendState{
692 .allocator = allocator,
693 .kind = .cuda,
694 .format = .cuda_ptx,
695 };
696 const instance = Window{ .rows = 2, .cols = 3, .radius = 1, .threads = .{ .x = 3, .y = 2 } };
697
698 var family_artifact = try createWindowFamilyArtifact(allocator, state.handle(), instance, .{ .limits = .testing });
699 defer family_artifact.deinit();
700 var fixed_artifact = try Window2x3R1F32.createKernelCallArtifact(allocator, state.handle(), .{ .limits = Window2x3R1F32.Limits.testing });
701 defer fixed_artifact.deinit();
702
703 const family_entry = family_artifact.entry();
704 const fixed_entry = fixed_artifact.entry();
705 try std.testing.expect(!std.mem.eql(u8, fixed_entry.target, family_entry.target));
706 try std.testing.expectEqual(fixed_entry.version, family_entry.version);
707 try std.testing.expectEqual(fixed_entry.format, family_entry.format);
708 try std.testing.expectEqualStrings("accy.kernel.stencil.window_family_r1_3x2_f32", family_entry.target);
709 try std.testing.expectEqualStrings("accy_kernel_stencil_window_family_r1_3x2_f32", family_entry.entry_name);
710 try std.testing.expectEqual(@as(u32, 5), family_entry.argument_count);
711 try std.testing.expectEqual(@as(u32, 2), family_entry.runtime_scalar_argument_count);
712 try std.testing.expect(family_entry.required_dtypes.contains(.i32));
713 try std.testing.expect(fixed_entry.shape_family_fingerprint == null);
714 try std.testing.expect(family_entry.shape_family_fingerprint != null);
715 const profile = family_entry.shape_profile orelse return error.TestExpectedShapeProfile;
716 try std.testing.expectEqualStrings("stencil_window", profile.name);
717 try std.testing.expectEqual(family_entry.shape_family_fingerprint.?, profile.fingerprint);
718 try std.testing.expectEqual(@as(usize, 2), profile.dimensions.len);
719 const rows_dimension = profile.runtimeScalarDimension(0) orelse return error.TestExpectedShapeProfile;
720 try std.testing.expectEqualStrings("r", rows_dimension.name);
721 try std.testing.expectEqual(@as(?u64, extent_mod.runtime_extent_max), rows_dimension.bounds.max);
722 switch (family_entry.launch) {
723 .derived => |launch| {
724 try std.testing.expectEqual(@as(u32, 3), launch.threadgroup[0]);
725 try std.testing.expectEqual(@as(u32, 2), launch.threadgroup[1]);
726 switch (launch.grid[0]) {
727 .runtime_u32_ceil_div => |term| {
728 try std.testing.expectEqual(@as(usize, 1), term.argument_index);
729 try std.testing.expectEqual(@as(u32, 3), term.divisor);
730 },
731 else => return error.TestExpectedDerivedLaunch,
732 }
733 switch (launch.grid[1]) {
734 .runtime_u32_ceil_div => |term| {
735 try std.testing.expectEqual(@as(usize, 0), term.argument_index);
736 try std.testing.expectEqual(@as(u32, 2), term.divisor);
737 },
738 else => return error.TestExpectedDerivedLaunch,
739 }
740 },
741 else => return error.TestExpectedDerivedLaunch,
742 }
743 }
744
745 test "stencil window family records fixed-entry specialization metadata" {
746 const instance = Window{ .rows = 2, .cols = 3, .radius = 1, .threads = .{ .x = 3, .y = 2 } };
747 var owned = try windowFamilySpecialization(std.testing.allocator, instance);
748 defer owned.deinit();
749 const specialization = owned.value;
750
751 try std.testing.expect(specialization.operationIs(.{ .stencil = .window }));
752 try std.testing.expectEqual(Window2x3R1F32.specialization.dtype, specialization.dtype);
753 try std.testing.expect(specialization.inputHasExtents(0, &.{ 4, 5 }));
754 try std.testing.expect(specialization.inputHasExtents(1, &.{9}));
755 try std.testing.expect(specialization.outputHasExtents(0, &.{ 2, 3 }));
756 try std.testing.expect(specialization.reductionMatches(0, .{
757 .name = "window",
758 .operator = .weighted_sum,
759 .extents = &.{9},
760 }));
761 try std.testing.expect(specialization.scheduleMatchesLaunch());
762 try std.testing.expect(specialization.shape_family != null);
763 }
764
765 test "stencil window instance round-trips through specialization" {
766 const instance = Window{ .rows = 6, .cols = 9, .radius = 2, .threads = .{ .x = 9, .y = 6 } };
767 var owned = try windowFamilySpecialization(std.testing.allocator, instance);
768 defer owned.deinit();
769
770 const recovered = windowInstanceFromSpecialization(owned.value) orelse return error.TestExpectedWindowInstance;
771 try std.testing.expectEqual(instance.rows, recovered.rows);
772 try std.testing.expectEqual(instance.cols, recovered.cols);
773 try std.testing.expectEqual(instance.radius, recovered.radius);
774 try std.testing.expectEqual(instance.dtype, recovered.dtype);
775 try std.testing.expectEqual(instance.accumulation_dtype, recovered.accumulation_dtype);
776 try std.testing.expectEqual(instance.threads.x, recovered.threads.x);
777 try std.testing.expectEqual(instance.threads.y, recovered.threads.y);
778
779 try std.testing.expectEqual(@as(?Window, null), windowInstanceFromSpecialization(.{}));
780 }
781
782 test "stencil window family rejects out-of-range radii" {
783 const oversized = Window{ .rows = 4, .cols = 4, .radius = window_radius_max + 1, .threads = .{ .x = 4, .y = 4 } };
784 try std.testing.expectError(error.StencilRadiusOutOfRange, windowFamilySpecialization(std.testing.allocator, oversized));
785 try std.testing.expectError(error.StencilRadiusOutOfRange, windowShapeFamily(std.testing.allocator, oversized));
786
787 const zero_radius = Window{ .rows = 4, .cols = 4, .radius = 0, .threads = .{ .x = 4, .y = 4 } };
788 try std.testing.expectError(error.StencilRadiusOutOfRange, windowFamilySpecialization(std.testing.allocator, zero_radius));
789 }