lib/accy/src/kernel/library/spatial.zig
daab053ee43316e1809a84551d573ddd1e5bf3d2
1 const std = @import("std");
2 const choir_abi = @import("choir_abi");
3
4 const artifact_product = @import("../../artifact/model/root.zig");
5 const shape = @import("../../choir/shape/root.zig");
6 const entry = @import("entry.zig");
7 const extent_mod = @import("extent.zig");
8 const geometry_mod = @import("geometry.zig");
9 const kernel = @import("../root.zig");
10
11 const DType = choir_abi.DType;
12 const runtimeExtentArgument = extent_mod.runtimeExtentArgument;
13
14 pub const GridCells = struct {
15 count: u64,
16 threads: u32 = 256,
17 point_axis: []const u8 = "p",
18 };
19
20 pub const grid_cells_family_version: u32 = 1;
21 pub const grid_cells_warp_size: u32 = 32;
22 pub const grid_cells_max_threads: u32 = 1024;
23 pub const grid_cells_max_blocks: u32 = 1024;
24
25 pub const GridGeometry = struct {
26 origin_x: f32,
27 origin_y: f32,
28 inv_cell_size: f32,
29 dims_x: u32,
30 dims_y: u32,
31
32 pub fn cellCount(self: GridGeometry) u64 {
33 return @as(u64, self.dims_x) * self.dims_y;
34 }
35 };
36
37 pub fn gridCellsBlockCount(count: u64, threads: u32) u64 {
38 return gridCellsBlockCountChecked(count, threads).?;
39 }
40
41 fn gridCellsBlockCountChecked(count: u64, threads: u32) ?u64 {
42 if (threads == 0) return null;
43 const biased = std.math.add(u64, count, threads - 1) catch return null;
44 return biased / threads;
45 }
46
47 pub fn gridCellsInstanceValid(instance: GridCells) bool {
48 if (instance.count == 0) return false;
49 if (instance.point_axis.len == 0) return false;
50 if (instance.threads == 0 or instance.threads > grid_cells_max_threads) return false;
51 if (instance.threads % grid_cells_warp_size != 0) return false;
52 return extent_mod.blockCountWithinLimit(instance.count, instance.threads, grid_cells_max_blocks);
53 }
54
55 pub fn gridGeometryValid(geometry: GridGeometry) bool {
56 if (geometry.dims_x == 0 or geometry.dims_y == 0) return false;
57 if (!(geometry.inv_cell_size > 0)) return false;
58 return geometry.cellCount() <= std.math.maxInt(i32);
59 }
60
61 fn gridCellAxisIndex(
62 b: anytype,
63 coordinate: kernel.Value,
64 origin: kernel.Value,
65 inv_cell_size: kernel.Value,
66 dim: kernel.Value,
67 ) !kernel.Value {
68 const offset = try b.sub(coordinate, origin);
69 const scaled = try b.mul(offset, inv_cell_size);
70 const raw_index = try b.cast(scaled, .i32);
71 const zero = try b.constantInt(.i32, 0);
72 const one = try b.constantInt(.i32, 1);
73 const last = try b.sub(dim, one);
74 const non_negative = try b.max(raw_index, zero);
75 return b.min(non_negative, last);
76 }
77
78 fn grid_cells_body_active(inner: anytype, ctx: anytype) !void {
79 const x = try ctx.args.param(.x).load(inner, ctx.point);
80 const y = try ctx.args.param(.y).load(inner, ctx.point);
81 const cx = try gridCellAxisIndex(
82 inner,
83 x.raw(),
84 ctx.args.param(.origin_x).raw(),
85 ctx.args.param(.inv_cell_size).raw(),
86 ctx.args.param(.dims_x).raw(),
87 );
88 const cy = try gridCellAxisIndex(
89 inner,
90 y.raw(),
91 ctx.args.param(.origin_y).raw(),
92 ctx.args.param(.inv_cell_size).raw(),
93 ctx.args.param(.dims_y).raw(),
94 );
95 const row = try inner.mul(cy, ctx.args.param(.dims_x).raw());
96 const cell = try inner.add(row, cx);
97 try ctx.args.param(.cells).store(inner, cell, ctx.point);
98 }
99
100 fn gridCellsBody(k: anytype, spec: GridCells, args: anytype) !void {
101 if (!gridCellsInstanceValid(spec)) return error.UnsupportedGridCellsInstance;
102 const point = try k.globalId(.x);
103 const count = try k.castIndex(args.param(.count).raw());
104 const active = try k.compare(.lt, point, count);
105 try k.guardDo(active, .{
106 .args = args,
107 .point = point,
108 }, grid_cells_body_active);
109 }
110
111 fn gridCellsFamilySchedule(instance: GridCells) kernel.logical.schedule.ThreadBlocks {
112 return kernel.logical.schedule.threadBlocks(.{ .x = instance.threads });
113 }
114
115 fn gridCellsRuntimeFamily() type {
116 return kernel.logical.Family(.{
117 .name = "accy_kernel_spatial_grid_cells_runtime_f32",
118 .parameters = .{
119 .cells = kernel.dynamicBuffer(.i32),
120 .x = kernel.dynamicBuffer(.f32),
121 .y = kernel.dynamicBuffer(.f32),
122 .count = kernel.scalar(.i32),
123 .origin_x = kernel.scalar(.f32),
124 .origin_y = kernel.scalar(.f32),
125 .inv_cell_size = kernel.scalar(.f32),
126 .dims_x = kernel.scalar(.i32),
127 .dims_y = kernel.scalar(.i32),
128 },
129 .Instance = GridCells,
130 .schedule = gridCellsFamilySchedule,
131 .body = gridCellsBody,
132 });
133 }
134
135 pub const GridCellsRuntimeFamilyF32 = gridCellsRuntimeFamily();
136
137 pub fn gridCellsFamilyTarget(allocator: std.mem.Allocator, instance: GridCells) ![]u8 {
138 return std.fmt.allocPrint(
139 allocator,
140 "accy.kernel.spatial.grid_cells_family_{d}_f32",
141 .{instance.threads},
142 );
143 }
144
145 pub fn gridCellsFamilyEntryName(allocator: std.mem.Allocator, instance: GridCells) ![]u8 {
146 return std.fmt.allocPrint(
147 allocator,
148 "accy_kernel_spatial_grid_cells_family_{d}_f32",
149 .{instance.threads},
150 );
151 }
152
153 pub fn gridCellsRuntimeArguments(
154 instance: GridCells,
155 geometry: GridGeometry,
156 ) ![6]choir_abi.ScalarArgument {
157 if (!gridGeometryValid(geometry)) return error.UnsupportedGridGeometry;
158 return .{
159 .{ .u32 = try runtimeExtentArgument(instance.count) },
160 .{ .f32 = geometry.origin_x },
161 .{ .f32 = geometry.origin_y },
162 .{ .f32 = geometry.inv_cell_size },
163 .{ .u32 = geometry.dims_x },
164 .{ .u32 = geometry.dims_y },
165 };
166 }
167
168 const testing = std.testing;
169
170 fn hostCellId(geometry: GridGeometry, x: f32, y: f32) i32 {
171 const fx = (x - geometry.origin_x) * geometry.inv_cell_size;
172 const fy = (y - geometry.origin_y) * geometry.inv_cell_size;
173 const cx = std.math.clamp(@as(i32, @intFromFloat(fx)), 0, @as(i32, @intCast(geometry.dims_x - 1)));
174 const cy = std.math.clamp(@as(i32, @intFromFloat(fy)), 0, @as(i32, @intCast(geometry.dims_y - 1)));
175 return cy * @as(i32, @intCast(geometry.dims_x)) + cx;
176 }
177
178 test "spatial grid cells matches the host reference with boundary clamping" {
179 const allocator = testing.allocator;
180 const count: usize = 70;
181 const instance = GridCells{ .count = count, .threads = 32 };
182 const geometry = GridGeometry{
183 .origin_x = -1.0,
184 .origin_y = -1.0,
185 .inv_cell_size = 4.0,
186 .dims_x = 8,
187 .dims_y = 8,
188 };
189 const blocks: u32 = @intCast(gridCellsBlockCount(instance.count, instance.threads));
190 try testing.expect(blocks > 1);
191
192 var xs: [count]f32 = undefined;
193 var ys: [count]f32 = undefined;
194 var seed: u32 = 0x2545f491;
195 for (&xs, &ys, 0..) |*x, *y, index| {
196 seed ^= seed << 13;
197 seed ^= seed >> 17;
198 seed ^= seed << 5;
199 const fx = @as(f32, @floatFromInt(seed % 1000)) / 250.0 - 2.0;
200 seed ^= seed << 13;
201 seed ^= seed >> 17;
202 seed ^= seed << 5;
203 const fy = @as(f32, @floatFromInt(seed % 1000)) / 250.0 - 2.0;
204 x.* = fx;
205 y.* = fy;
206 if (index == 0) {
207 x.* = -1.0;
208 y.* = -1.0;
209 }
210 if (index == 1) {
211 x.* = 5.0;
212 y.* = -9.0;
213 }
214 }
215
216 var cells = @as([count]i32, @splat(-1));
217 var graph = try GridCellsRuntimeFamilyF32.build(allocator, GridCellsRuntimeFamilyF32.Limits.testing, instance);
218 defer graph.deinit();
219 const runtime_arguments = try gridCellsRuntimeArguments(instance, geometry);
220 try graph.runCpuWithLaunch(allocator, &.{
221 kernel.argumentBuffer(i32, cells[0..]),
222 kernel.argumentBuffer(f32, xs[0..]),
223 kernel.argumentBuffer(f32, ys[0..]),
224 kernel.argumentI32(@intCast(count)),
225 kernel.argumentF32(geometry.origin_x),
226 kernel.argumentF32(geometry.origin_y),
227 kernel.argumentF32(geometry.inv_cell_size),
228 kernel.argumentI32(@intCast(geometry.dims_x)),
229 kernel.argumentI32(@intCast(geometry.dims_y)),
230 }, .{
231 .grid = .{ blocks, 1, 1 },
232 .block = .{ instance.threads, 1, 1 },
233 });
234
235 for (xs, ys, cells) |x, y, cell| {
236 try testing.expectEqual(hostCellId(geometry, x, y), cell);
237 }
238 _ = runtime_arguments;
239 }
240
241 test "spatial grid cells identity and validity" {
242 const allocator = testing.allocator;
243 const instance = GridCells{ .count = 5000, .threads = 64 };
244 const target = try gridCellsFamilyTarget(allocator, instance);
245 defer allocator.free(target);
246 try testing.expectEqualStrings("accy.kernel.spatial.grid_cells_family_64_f32", target);
247
248 try testing.expect(gridCellsInstanceValid(instance));
249 try testing.expect(!gridCellsInstanceValid(.{ .count = 0, .threads = 64 }));
250 try testing.expect(!gridCellsInstanceValid(.{ .count = 10, .threads = 48 }));
251 try testing.expect(!gridCellsInstanceValid(.{ .count = std.math.maxInt(u64), .threads = 32 }));
252
253 try testing.expect(gridGeometryValid(.{ .origin_x = 0, .origin_y = 0, .inv_cell_size = 1, .dims_x = 4, .dims_y = 4 }));
254 try testing.expect(!gridGeometryValid(.{ .origin_x = 0, .origin_y = 0, .inv_cell_size = 0, .dims_x = 4, .dims_y = 4 }));
255 try testing.expect(!gridGeometryValid(.{ .origin_x = 0, .origin_y = 0, .inv_cell_size = 1, .dims_x = 0, .dims_y = 4 }));
256
257 const args = try gridCellsRuntimeArguments(instance, .{
258 .origin_x = -1,
259 .origin_y = -1,
260 .inv_cell_size = 4,
261 .dims_x = 8,
262 .dims_y = 8,
263 });
264 try testing.expectEqual(@as(u32, 5000), args[0].u32);
265 try testing.expectEqual(@as(u32, 8), args[4].u32);
266 }
267
268 pub const grid_count_shared_cells_cap: u32 = 4096;
269
270 pub const GridCount = struct {
271 count: u64,
272 cells: u32,
273 threads: u32 = 256,
274 point_axis: []const u8 = "p",
275 };
276
277 pub fn gridCountInstanceValid(instance: GridCount) bool {
278 if (instance.count == 0) return false;
279 if (instance.point_axis.len == 0) return false;
280 if (instance.cells == 0 or instance.cells > grid_count_shared_cells_cap) return false;
281 if (instance.threads == 0 or instance.threads > grid_cells_max_threads) return false;
282 if (instance.threads % grid_cells_warp_size != 0) return false;
283 return extent_mod.blockCountWithinLimit(instance.count, instance.threads, grid_cells_max_blocks);
284 }
285
286 fn grid_count_body_zero_bin(loop_builder: anytype, bin: kernel.Value, acc: kernel.Value, ctx: anytype) !kernel.Value {
287 try loop_builder.storeIndex(ctx.zero_count, ctx.shared_bins, bin);
288 return acc;
289 }
290
291 fn grid_count_body_active(inner: anytype, ctx: anytype) !void {
292 const id_loaded = try ctx.args.param(.ids).load(inner, ctx.point);
293 const zero = try inner.constantInt(.i32, 0);
294 const one_i32 = try inner.constantInt(.i32, 1);
295 const last_raw = try inner.sub(ctx.capped, one_i32);
296 const last = try inner.max(last_raw, zero);
297 const non_negative = try inner.max(id_loaded.raw(), zero);
298 const clamped = try inner.min(non_negative, last);
299 const bin = try inner.castIndex(clamped);
300 _ = try inner.atomicRmwIndex(.add, one_i32, ctx.shared_bins, bin);
301 }
302
303 fn grid_count_body_grid(loop_builder: anytype, bin: kernel.Value, acc: kernel.Value, ctx: anytype) !kernel.Value {
304 const partial = try loop_builder.loadIndex(ctx.shared_bins, bin);
305 const partial_value = try loop_builder.cast(partial, .f32);
306 const column = try loop_builder.mul(bin, ctx.grid);
307 const cell = try loop_builder.add(column, ctx.block);
308 try loop_builder.storeIndex(partial_value, ctx.args.param(.counts).raw(), cell);
309 return acc;
310 }
311
312 fn gridCountBody(k: anytype, spec: GridCount, args: anytype) !void {
313 if (!gridCountInstanceValid(spec)) return error.UnsupportedGridCountInstance;
314 const shared_bins = try k.sharedBuffer(.i32, spec.cells);
315 const zero_count = try k.constantInt(.i32, 0);
316 const thread = try k.castIndex(try k.threadId(.x));
317 const stride = try k.castIndex(try k.blockDim(.x));
318 const cells_cap = try k.constantInt(.i32, @intCast(spec.cells));
319 const total_non_negative = try k.max(args.param(.cells_total).raw(), zero_count);
320 const capped = try k.min(total_non_negative, cells_cap);
321 const bins = try k.castIndex(capped);
322
323 _ = try k.fold(thread, bins, stride, zero_count, .{
324 .shared_bins = shared_bins,
325 .zero_count = zero_count,
326 }, grid_count_body_zero_bin);
327 try k.barrier(.block);
328
329 const point = try k.globalId(.x);
330 const count = try k.castIndex(args.param(.count).raw());
331 const active = try k.compare(.lt, point, count);
332 try k.guardDo(active, .{
333 .args = args,
334 .point = point,
335 .shared_bins = shared_bins,
336 .capped = capped,
337 }, grid_count_body_active);
338 try k.barrier(.block);
339
340 const block = try k.blockId(.x);
341 const grid = try k.gridDim(.x);
342 _ = try k.fold(thread, bins, stride, zero_count, .{
343 .args = args,
344 .shared_bins = shared_bins,
345 .block = block,
346 .grid = grid,
347 }, grid_count_body_grid);
348 }
349
350 fn gridCountRuntimeFamily() type {
351 return kernel.logical.Family(.{
352 .name = "accy_kernel_spatial_grid_count_runtime_i32",
353 .parameters = .{
354 .counts = kernel.dynamicBuffer(.f32),
355 .ids = kernel.dynamicBuffer(.i32),
356 .count = kernel.scalar(.i32),
357 .cells_total = kernel.scalar(.i32),
358 },
359 .Instance = GridCount,
360 .schedule = gridCountFamilySchedule,
361 .body = gridCountBody,
362 });
363 }
364
365 fn gridCountFamilySchedule(instance: GridCount) kernel.logical.schedule.ThreadBlocks {
366 return kernel.logical.schedule.threadBlocks(.{ .x = instance.threads });
367 }
368
369 pub const GridCountRuntimeFamilyI32 = gridCountRuntimeFamily();
370
371 pub fn gridCountFamilyTarget(allocator: std.mem.Allocator, instance: GridCount) ![]u8 {
372 return std.fmt.allocPrint(
373 allocator,
374 "accy.kernel.spatial.grid_count_family_{d}_{d}_i32",
375 .{ instance.cells, instance.threads },
376 );
377 }
378
379 pub fn gridCountFamilyEntryName(allocator: std.mem.Allocator, instance: GridCount) ![]u8 {
380 return std.fmt.allocPrint(
381 allocator,
382 "accy_kernel_spatial_grid_count_family_{d}_{d}_i32",
383 .{ instance.cells, instance.threads },
384 );
385 }
386
387 pub fn gridCountRuntimeArguments(instance: GridCount, cells_total: u32) ![2]choir_abi.ScalarArgument {
388 if (cells_total == 0 or cells_total > instance.cells) return error.UnsupportedGridCountInstance;
389 return .{
390 .{ .u32 = try runtimeExtentArgument(instance.count) },
391 .{ .u32 = cells_total },
392 };
393 }
394
395 test "spatial grid count tallies precomputed ids per block in column-major order" {
396 const allocator = testing.allocator;
397 const count: usize = 90;
398 const cells_total: u32 = 12;
399 const instance = GridCount{ .count = count, .cells = 64, .threads = 32 };
400 const blocks: u32 = @intCast(gridCellsBlockCount(instance.count, instance.threads));
401 try testing.expectEqual(@as(u32, 3), blocks);
402
403 var ids: [count]i32 = undefined;
404 var seed: u32 = 0x2545f491;
405 for (&ids, 0..) |*id, index| {
406 seed ^= seed << 13;
407 seed ^= seed >> 17;
408 seed ^= seed << 5;
409 id.* = @intCast(seed % cells_total);
410 if (index == 7) id.* = -3;
411 if (index == 13) id.* = 200;
412 }
413
414 var counts = @as([(12 * 3)]f32, @splat(-1));
415 var graph = try GridCountRuntimeFamilyI32.build(allocator, GridCountRuntimeFamilyI32.Limits.testing, instance);
416 defer graph.deinit();
417 try graph.runCpuWithLaunch(allocator, &.{
418 kernel.argumentBuffer(f32, counts[0..]),
419 kernel.argumentBuffer(i32, ids[0..]),
420 kernel.argumentI32(@intCast(count)),
421 kernel.argumentI32(@intCast(cells_total)),
422 }, .{
423 .grid = .{ blocks, 1, 1 },
424 .block = .{ instance.threads, 1, 1 },
425 });
426
427 var expected = @as([(12 * 3)]f32, @splat(0));
428 for (ids, 0..) |id, index| {
429 const clamped: usize = @intCast(std.math.clamp(id, 0, @as(i32, @intCast(cells_total - 1))));
430 const block = index / instance.threads;
431 expected[clamped * 3 + block] += 1;
432 }
433 try testing.expectEqualSlices(f32, expected[0..], counts[0..]);
434 }
435
436 test "spatial grid count clamps runtime cell totals to the compiled shared cap" {
437 const allocator = testing.allocator;
438 const count: usize = 32;
439 const instance = GridCount{ .count = count, .cells = 4, .threads = 32 };
440
441 var ids: [count]i32 = undefined;
442 for (&ids, 0..) |*id, index| id.* = @intCast(index % 6);
443
444 var graph = try GridCountRuntimeFamilyI32.build(allocator, GridCountRuntimeFamilyI32.Limits.testing, instance);
445 defer graph.deinit();
446
447 var over_counts = @as([5]f32, @splat(-1));
448 try graph.runCpuWithLaunch(allocator, &.{
449 kernel.argumentBuffer(f32, over_counts[0..]),
450 kernel.argumentBuffer(i32, ids[0..]),
451 kernel.argumentI32(@intCast(count)),
452 kernel.argumentI32(5),
453 }, .{
454 .grid = .{ 1, 1, 1 },
455 .block = .{ instance.threads, 1, 1 },
456 });
457 var expected = @as([5]f32, @splat(0));
458 for (ids) |id| {
459 const clamped: usize = @intCast(std.math.clamp(id, 0, 3));
460 expected[clamped] += 1;
461 }
462 expected[4] = -1;
463 try testing.expectEqualSlices(f32, expected[0..], over_counts[0..]);
464
465 var zero_counts = @as([4]f32, @splat(-1));
466 try graph.runCpuWithLaunch(allocator, &.{
467 kernel.argumentBuffer(f32, zero_counts[0..]),
468 kernel.argumentBuffer(i32, ids[0..]),
469 kernel.argumentI32(@intCast(count)),
470 kernel.argumentI32(0),
471 }, .{
472 .grid = .{ 1, 1, 1 },
473 .block = .{ instance.threads, 1, 1 },
474 });
475 try testing.expectEqualSlices(f32, &[_]f32{ -1, -1, -1, -1 }, zero_counts[0..]);
476
477 var negative_counts = @as([4]f32, @splat(-1));
478 try graph.runCpuWithLaunch(allocator, &.{
479 kernel.argumentBuffer(f32, negative_counts[0..]),
480 kernel.argumentBuffer(i32, ids[0..]),
481 kernel.argumentI32(@intCast(count)),
482 kernel.argumentI32(-7),
483 }, .{
484 .grid = .{ 1, 1, 1 },
485 .block = .{ instance.threads, 1, 1 },
486 });
487 try testing.expectEqualSlices(f32, &[_]f32{ -1, -1, -1, -1 }, negative_counts[0..]);
488 }
489
490 test "spatial grid count identity and validity" {
491 const allocator = testing.allocator;
492 const instance = GridCount{ .count = 5000, .cells = 64, .threads = 64 };
493 const target = try gridCountFamilyTarget(allocator, instance);
494 defer allocator.free(target);
495 try testing.expectEqualStrings("accy.kernel.spatial.grid_count_family_64_64_i32", target);
496
497 try testing.expect(gridCountInstanceValid(instance));
498 try testing.expect(!gridCountInstanceValid(.{ .count = 5000, .cells = 0, .threads = 64 }));
499 try testing.expect(!gridCountInstanceValid(.{ .count = 5000, .cells = 8192, .threads = 64 }));
500 try testing.expect(!gridCountInstanceValid(.{ .count = std.math.maxInt(u64), .cells = 64, .threads = 32 }));
501
502 try testing.expectError(error.UnsupportedGridCountInstance, gridCountRuntimeArguments(instance, 65));
503 const args = try gridCountRuntimeArguments(instance, 48);
504 try testing.expectEqual(@as(u32, 48), args[1].u32);
505 }
506
507 pub const sort_mod = @import("sort.zig");
508 pub const scan_mod = @import("scan.zig");
509
510 pub const GridBuildPlan = struct {
511 cells_instance: GridCells,
512 sort_instance: sort_mod.RadixSplit,
513 count_instance: GridCount,
514 offsets_scan: scan_mod.DeviceScan,
515 sort_passes: u32,
516 cells_total: u32,
517 };
518
519 pub fn gridBuildPlan(count: u64, geometry: GridGeometry, threads: u32) !GridBuildPlan {
520 if (!gridGeometryValid(geometry)) return error.UnsupportedGridGeometry;
521 const cells_total: u32 = @intCast(geometry.cellCount());
522 if (cells_total > grid_count_shared_cells_cap) return error.UnsupportedGridGeometry;
523
524 const cells_instance = GridCells{ .count = count, .threads = threads };
525 if (!gridCellsInstanceValid(cells_instance)) return error.UnsupportedGridCellsInstance;
526 const sort_instance = sort_mod.RadixSplit{ .extent = count, .threads = threads };
527 if (!sort_mod.radixSplitInstanceValid(sort_instance)) return error.UnsupportedGridCellsInstance;
528 const count_instance = GridCount{ .count = count, .cells = cells_total, .threads = threads };
529 if (!gridCountInstanceValid(count_instance)) return error.UnsupportedGridCountInstance;
530 const count_blocks = gridCellsBlockCount(count, threads);
531 const counts_extent = std.math.mul(u64, cells_total, count_blocks) catch return error.UnsupportedGridGeometry;
532 const offsets_scan = scan_mod.DeviceScan{
533 .extent = counts_extent,
534 .mode = .exclusive,
535 .threads = scan_mod.deviceScanThreadsForExtent(counts_extent) orelse return error.UnsupportedGridGeometry,
536 };
537 if (!scan_mod.deviceScanInstanceValid(offsets_scan)) return error.UnsupportedGridGeometry;
538
539 const id_bits: u32 = 32 - @clz(@max(cells_total - 1, 1));
540 const digit_passes = (id_bits + sort_mod.radix_digit_bits - 1) / sort_mod.radix_digit_bits;
541
542 return .{
543 .cells_instance = cells_instance,
544 .sort_instance = sort_instance,
545 .count_instance = count_instance,
546 .offsets_scan = offsets_scan,
547 .sort_passes = digit_passes,
548 .cells_total = cells_total,
549 };
550 }
551
552 test "spatial grid build plan derives pass counts and stage instances" {
553 const geometry = GridGeometry{
554 .origin_x = -1.0,
555 .origin_y = -1.0,
556 .inv_cell_size = 4.0,
557 .dims_x = 8,
558 .dims_y = 8,
559 };
560 const plan = try gridBuildPlan(5000, geometry, 64);
561 try testing.expectEqual(@as(u32, 64), plan.cells_total);
562 try testing.expectEqual(@as(u32, 2), plan.sort_passes);
563 try testing.expectEqual(@as(u64, 64 * 79), plan.offsets_scan.extent);
564 try testing.expectEqual(scan_mod.PrefixSumMode.exclusive, plan.offsets_scan.mode);
565 try testing.expectEqual(plan.offsets_scan.extent, @as(u64, plan.cells_total) * gridCellsBlockCount(5000, plan.count_instance.threads));
566
567 const wide = GridGeometry{
568 .origin_x = 0,
569 .origin_y = 0,
570 .inv_cell_size = 1.0,
571 .dims_x = 64,
572 .dims_y = 64,
573 };
574 const wide_plan = try gridBuildPlan(5000, wide, 64);
575 try testing.expectEqual(@as(u32, 4096), wide_plan.cells_total);
576 try testing.expectEqual(@as(u32, 3), wide_plan.sort_passes);
577 try testing.expectEqual(@as(u64, 4096 * 79), wide_plan.offsets_scan.extent);
578
579 try testing.expectError(error.UnsupportedGridGeometry, gridBuildPlan(5000, .{
580 .origin_x = 0,
581 .origin_y = 0,
582 .inv_cell_size = 1.0,
583 .dims_x = 128,
584 .dims_y = 64,
585 }, 64));
586 }
587
588 pub const GridNeighborCount = struct {
589 count: u64,
590 offsets_extent: u64 = 1,
591 threads: u32 = 256,
592 point_axis: []const u8 = "p",
593 };
594
595 pub fn gridNeighborCountInstanceValid(instance: GridNeighborCount) bool {
596 if (instance.count == 0) return false;
597 if (instance.offsets_extent == 0) return false;
598 if (instance.point_axis.len == 0) return false;
599 if (instance.threads == 0 or instance.threads > grid_cells_max_threads) return false;
600 if (instance.threads % grid_cells_warp_size != 0) return false;
601 return extent_mod.blockCountWithinLimit(instance.count, instance.threads, grid_cells_max_blocks);
602 }
603
604 fn grid_neighbor_count_body_active(inner: anytype, ctx: anytype) !void {
605 const px = try ctx.args.param(.x).load(inner, ctx.point);
606 const py = try ctx.args.param(.y).load(inner, ctx.point);
607 const dims_x = ctx.args.param(.dims_x).raw();
608 const dims_y = ctx.args.param(.dims_y).raw();
609 const cx = try gridCellAxisIndex(
610 inner,
611 px.raw(),
612 ctx.args.param(.origin_x).raw(),
613 ctx.args.param(.inv_cell_size).raw(),
614 dims_x,
615 );
616 const cy = try gridCellAxisIndex(
617 inner,
618 py.raw(),
619 ctx.args.param(.origin_y).raw(),
620 ctx.args.param(.inv_cell_size).raw(),
621 dims_y,
622 );
623 const zero = try inner.constantInt(.i32, 0);
624 const one = try inner.constantInt(.i32, 1);
625 const step = try inner.constantIndex(1);
626 const cells_total = try inner.mul(dims_x, dims_y);
627 const stride = ctx.args.param(.stride).raw();
628 const count_scalar = ctx.args.param(.count).raw();
629 const col_lo = try inner.max(try inner.sub(cx, one), zero);
630 const col_hi = try inner.min(try inner.add(cx, one), try inner.sub(dims_x, one));
631
632 var total = zero;
633 inline for ([_]i64{ -1, 0, 1 }) |row_offset| {
634 const offset_const = try inner.constantInt(.i32, row_offset);
635 const ny = try inner.add(cy, offset_const);
636 const row_low = try inner.compare(.ge, ny, zero);
637 const row_high = try inner.compare(.lt, ny, dims_y);
638 const row_valid = try inner.and_(row_low, row_high);
639 const row_base = try inner.mul(ny, dims_x);
640 const first_cell = try inner.add(row_base, col_lo);
641 const next_cell = try inner.add(try inner.add(row_base, col_hi), one);
642 const has_next = try inner.compare(.lt, next_cell, cells_total);
643
644 const first_safe = try inner.select(row_valid, first_cell, zero);
645 const start_loaded = try ctx.args.param(.offsets).load(
646 inner,
647 try inner.castIndex(try inner.mul(first_safe, stride)),
648 );
649 const start_value = try inner.cast(start_loaded.raw(), .i32);
650
651 const next_guard = try inner.and_(row_valid, has_next);
652 const next_safe = try inner.select(next_guard, next_cell, zero);
653 const end_loaded = try ctx.args.param(.offsets).load(
654 inner,
655 try inner.castIndex(try inner.mul(next_safe, stride)),
656 );
657 const end_value = try inner.select(
658 has_next,
659 try inner.cast(end_loaded.raw(), .i32),
660 count_scalar,
661 );
662
663 const start_position = try inner.select(row_valid, start_value, zero);
664 const end_position = try inner.select(row_valid, end_value, zero);
665 total = try inner.fold(
666 try inner.castIndex(start_position),
667 try inner.castIndex(end_position),
668 step,
669 total,
670 .{
671 .args = ctx.args,
672 .point = ctx.point,
673 .px = px,
674 .py = py,
675 .one = one,
676 .zero = zero,
677 },
678 grid_neighbor_count_body_visit,
679 );
680 }
681 try ctx.args.param(.neighbors).store(inner, total, ctx.point);
682 }
683
684 fn grid_neighbor_count_body_visit(loop_builder: anytype, position: kernel.Value, acc: kernel.Value, fold_ctx: anytype) !kernel.Value {
685 const candidate = try fold_ctx.args.param(.sorted_indices).load(loop_builder, position);
686 const candidate_index = try loop_builder.castIndex(candidate.raw());
687 const qx = try fold_ctx.args.param(.x).load(loop_builder, candidate_index);
688 const qy = try fold_ctx.args.param(.y).load(loop_builder, candidate_index);
689 const dx = try loop_builder.sub(fold_ctx.px.raw(), qx.raw());
690 const dy = try loop_builder.sub(fold_ctx.py.raw(), qy.raw());
691 const dist2 = try loop_builder.add(
692 try loop_builder.mul(dx, dx),
693 try loop_builder.mul(dy, dy),
694 );
695 const within = try loop_builder.compare(.le, dist2, fold_ctx.args.param(.radius2).raw());
696 const not_self = try loop_builder.compare(.ne, candidate_index, fold_ctx.point);
697 const hit = try loop_builder.and_(within, not_self);
698 const contribution = try loop_builder.select(hit, fold_ctx.one, fold_ctx.zero);
699 return loop_builder.add(acc, contribution);
700 }
701
702 fn gridNeighborCountBody(k: anytype, spec: GridNeighborCount, args: anytype) !void {
703 if (!gridNeighborCountInstanceValid(spec)) return error.UnsupportedGridNeighborCountInstance;
704 const point = try k.globalId(.x);
705 const count = try k.castIndex(args.param(.count).raw());
706 const active = try k.compare(.lt, point, count);
707 try k.guardDo(active, .{
708 .args = args,
709 .point = point,
710 }, grid_neighbor_count_body_active);
711 }
712
713 fn gridNeighborCountFamilySchedule(instance: GridNeighborCount) kernel.logical.schedule.ThreadBlocks {
714 return kernel.logical.schedule.threadBlocks(.{ .x = instance.threads });
715 }
716
717 fn gridNeighborCountRuntimeFamily() type {
718 return kernel.logical.Family(.{
719 .name = "accy_kernel_spatial_grid_neighbor_count_runtime_f32",
720 .parameters = .{
721 .neighbors = kernel.dynamicBuffer(.i32),
722 .x = kernel.dynamicBuffer(.f32),
723 .y = kernel.dynamicBuffer(.f32),
724 .sorted_indices = kernel.dynamicBuffer(.i32),
725 .offsets = kernel.dynamicBuffer(.f32),
726 .count = kernel.scalar(.i32),
727 .origin_x = kernel.scalar(.f32),
728 .origin_y = kernel.scalar(.f32),
729 .inv_cell_size = kernel.scalar(.f32),
730 .dims_x = kernel.scalar(.i32),
731 .dims_y = kernel.scalar(.i32),
732 .stride = kernel.scalar(.i32),
733 .radius2 = kernel.scalar(.f32),
734 },
735 .Instance = GridNeighborCount,
736 .schedule = gridNeighborCountFamilySchedule,
737 .body = gridNeighborCountBody,
738 });
739 }
740
741 pub const GridNeighborCountRuntimeFamilyF32 = gridNeighborCountRuntimeFamily();
742
743 pub fn gridNeighborCountFamilyTarget(allocator: std.mem.Allocator, instance: GridNeighborCount) ![]u8 {
744 return std.fmt.allocPrint(
745 allocator,
746 "accy.kernel.spatial.grid_neighbor_count_family_{d}_f32",
747 .{instance.threads},
748 );
749 }
750
751 pub fn gridNeighborCountFamilyEntryName(allocator: std.mem.Allocator, instance: GridNeighborCount) ![]u8 {
752 return std.fmt.allocPrint(
753 allocator,
754 "accy_kernel_spatial_grid_neighbor_count_family_{d}_f32",
755 .{instance.threads},
756 );
757 }
758
759 /// Builds the eight scalar arguments for a neighbor-count kernel launch: the function returns, in
760 /// order, the point count, the grid origin along x and y, the inverse cell size, the cell counts
761 /// along x and y, the stride, and the squared radius. The radius must be finite, zero or greater,
762 /// and at most one grid cell, so the search reaches only the neighboring cells. The call returns
763 /// `error.UnsupportedGridGeometry` for a grid with zero cells, an inverse cell size of zero or
764 /// less, or more cells than a signed 32-bit count. The call returns
765 /// `error.UnsupportedGridNeighborCountInstance` for a zero stride or a radius that is NaN, infinite
766 /// or negative, and `error.NeighborRadiusExceedsCellSize` for a radius larger than one cell. The
767 /// call returns `error.ExtentOverflowsIndexRange` when the point count is zero or too large for the
768 /// kernel's index range.
769 pub fn gridNeighborCountRuntimeArguments(
770 instance: GridNeighborCount,
771 geometry: GridGeometry,
772 stride: u32,
773 radius: f32,
774 ) ![8]choir_abi.ScalarArgument {
775 if (!gridGeometryValid(geometry)) return error.UnsupportedGridGeometry;
776 if (stride == 0) return error.UnsupportedGridNeighborCountInstance;
777 if (!std.math.isFinite(radius) or radius < 0) return error.UnsupportedGridNeighborCountInstance;
778 if (radius * geometry.inv_cell_size > 1.0) return error.NeighborRadiusExceedsCellSize;
779 return .{
780 .{ .u32 = try runtimeExtentArgument(instance.count) },
781 .{ .f32 = geometry.origin_x },
782 .{ .f32 = geometry.origin_y },
783 .{ .f32 = geometry.inv_cell_size },
784 .{ .u32 = geometry.dims_x },
785 .{ .u32 = geometry.dims_y },
786 .{ .u32 = stride },
787 .{ .f32 = radius * radius },
788 };
789 }
790
791 test "spatial grid neighbor count matches the quadratic host reference" {
792 const allocator = testing.allocator;
793 const count: usize = 48;
794 const geometry = GridGeometry{
795 .origin_x = 0,
796 .origin_y = 0,
797 .inv_cell_size = 1.0,
798 .dims_x = 4,
799 .dims_y = 4,
800 };
801 const radius: f32 = 0.75;
802 const instance = GridNeighborCount{ .count = count, .threads = 32 };
803 const blocks: u32 = @intCast(gridCellsBlockCount(instance.count, instance.threads));
804
805 var xs: [count]f32 = undefined;
806 var ys: [count]f32 = undefined;
807 var seed: u32 = 0x9e3779b9;
808 for (&xs, &ys, 0..) |*x, *y, index| {
809 seed ^= seed << 13;
810 seed ^= seed >> 17;
811 seed ^= seed << 5;
812 x.* = @as(f32, @floatFromInt(seed % 1000)) / 200.0 - 0.5;
813 seed ^= seed << 13;
814 seed ^= seed >> 17;
815 seed ^= seed << 5;
816 y.* = @as(f32, @floatFromInt(seed % 1000)) / 200.0 - 0.5;
817 if (index == 0) {
818 x.* = -2.0;
819 y.* = -2.0;
820 }
821 if (index == 1) {
822 x.* = -1.6;
823 y.* = -2.0;
824 }
825 }
826
827 var host_ids: [count]i32 = undefined;
828 for (xs, ys, &host_ids) |x, y, *id| id.* = hostCellId(geometry, x, y);
829
830 var order: [count]i32 = undefined;
831 for (&order, 0..) |*value, index| value.* = @intCast(index);
832 var sort_index: usize = 1;
833 while (sort_index < count) : (sort_index += 1) {
834 const key = order[sort_index];
835 const key_cell = host_ids[@intCast(key)];
836 var slot = sort_index;
837 while (slot > 0 and host_ids[@intCast(order[slot - 1])] > key_cell) : (slot -= 1) {
838 order[slot] = order[slot - 1];
839 }
840 order[slot] = key;
841 }
842
843 const cells_total: usize = 16;
844 var counts_by_cell = @as([cells_total]u32, @splat(0));
845 for (host_ids) |id| counts_by_cell[@intCast(id)] += 1;
846 var offsets: [cells_total]f32 = undefined;
847 var prefix: u32 = 0;
848 for (&offsets, counts_by_cell) |*offset, cell_count| {
849 offset.* = @floatFromInt(prefix);
850 prefix += cell_count;
851 }
852
853 var expected: [count]i32 = undefined;
854 for (0..count) |a| {
855 var total: i32 = 0;
856 for (0..count) |b| {
857 if (a == b) continue;
858 const dx = xs[a] - xs[b];
859 const dy = ys[a] - ys[b];
860 if (dx * dx + dy * dy <= radius * radius) total += 1;
861 }
862 expected[a] = total;
863 }
864
865 var neighbors = @as([count]i32, @splat(-1));
866 var graph = try GridNeighborCountRuntimeFamilyF32.build(allocator, GridNeighborCountRuntimeFamilyF32.Limits.testing, instance);
867 defer graph.deinit();
868 try graph.runCpuWithLaunch(allocator, &.{
869 kernel.argumentBuffer(i32, neighbors[0..]),
870 kernel.argumentBuffer(f32, xs[0..]),
871 kernel.argumentBuffer(f32, ys[0..]),
872 kernel.argumentBuffer(i32, order[0..]),
873 kernel.argumentBuffer(f32, offsets[0..]),
874 kernel.argumentI32(@intCast(count)),
875 kernel.argumentF32(geometry.origin_x),
876 kernel.argumentF32(geometry.origin_y),
877 kernel.argumentF32(geometry.inv_cell_size),
878 kernel.argumentI32(@intCast(geometry.dims_x)),
879 kernel.argumentI32(@intCast(geometry.dims_y)),
880 kernel.argumentI32(1),
881 kernel.argumentF32(radius * radius),
882 }, .{
883 .grid = .{ blocks, 1, 1 },
884 .block = .{ instance.threads, 1, 1 },
885 });
886
887 try testing.expectEqualSlices(i32, expected[0..], neighbors[0..]);
888 try testing.expect(expected[0] >= 1);
889 }
890
891 test "spatial grid neighbor count identity and validity" {
892 const allocator = testing.allocator;
893 const instance = GridNeighborCount{ .count = 5000, .threads = 64 };
894 const target = try gridNeighborCountFamilyTarget(allocator, instance);
895 defer allocator.free(target);
896 try testing.expectEqualStrings("accy.kernel.spatial.grid_neighbor_count_family_64_f32", target);
897
898 try testing.expect(gridNeighborCountInstanceValid(instance));
899 try testing.expect(!gridNeighborCountInstanceValid(.{ .count = 0, .threads = 64 }));
900 try testing.expect(!gridNeighborCountInstanceValid(.{ .count = 10, .threads = 48 }));
901 try testing.expect(!gridNeighborCountInstanceValid(.{ .count = std.math.maxInt(u64), .threads = 32 }));
902
903 const geometry = GridGeometry{
904 .origin_x = -1,
905 .origin_y = -1,
906 .inv_cell_size = 4,
907 .dims_x = 8,
908 .dims_y = 8,
909 };
910 const args = try gridNeighborCountRuntimeArguments(instance, geometry, 3, 0.2);
911 try testing.expectEqual(@as(u32, 5000), args[0].u32);
912 try testing.expectEqual(@as(u32, 3), args[6].u32);
913 const radius_runtime: f32 = 0.2;
914 try testing.expectEqual(radius_runtime * radius_runtime, args[7].f32);
915
916 try testing.expectError(error.UnsupportedGridNeighborCountInstance, gridNeighborCountRuntimeArguments(instance, geometry, 0, 0.2));
917 try testing.expectError(error.UnsupportedGridNeighborCountInstance, gridNeighborCountRuntimeArguments(instance, geometry, 3, -0.5));
918 try testing.expectError(error.UnsupportedGridNeighborCountInstance, gridNeighborCountRuntimeArguments(instance, geometry, 3, std.math.nan(f32)));
919 try testing.expectError(error.NeighborRadiusExceedsCellSize, gridNeighborCountRuntimeArguments(instance, geometry, 3, 0.3));
920 }
921
922 pub const grid_count_family_version: u32 = 1;
923 pub const grid_neighbor_count_family_version: u32 = 1;
924 pub const grid_count_cell_axis = "c";
925 pub const grid_count_block_axis = "b";
926 pub const grid_neighbor_offsets_axis = "o";
927 const spatial_thread_caps = geometry_mod.ThreadCaps1D{};
928
929 pub fn spatialThreadsForCount(count: u64) u32 {
930 const raw = geometry_mod.threadsForExtent(count, spatial_thread_caps);
931 if (raw % grid_cells_warp_size != 0) return grid_cells_warp_size;
932 return raw;
933 }
934
935 pub fn spatialThreadCandidatesForCount(count: u64) geometry_mod.Thread1DCandidates {
936 var candidates = geometry_mod.threadCandidatesForExtent(count, spatial_thread_caps);
937 for (candidates.items[0..candidates.count]) |*threads| {
938 if (threads.* % grid_cells_warp_size != 0) threads.* = grid_cells_warp_size;
939 }
940 return candidates;
941 }
942
943 fn spatialRuntimeExtentBounds() shape.Bounds {
944 return .{ .min = 1, .max = extent_mod.runtime_extent_max };
945 }
946
947 pub fn gridCellsShapeFamily(backing_allocator: std.mem.Allocator, instance: GridCells) !shape.Family {
948 var builder = try shape.Builder.init(backing_allocator, "grid_cells");
949 errdefer builder.deinit();
950
951 const point = try builder.symbol(instance.point_axis);
952 const point_expr = try builder.symbolExpression(point);
953 _ = try builder.tensor("cells", &.{point_expr});
954 _ = try builder.tensor("x", &.{point_expr});
955 _ = try builder.tensor("y", &.{point_expr});
956 try builder.assumeBounds(point_expr, spatialRuntimeExtentBounds());
957
958 return builder.finish();
959 }
960
961 pub fn gridCellsFamilyFingerprint(backing_allocator: std.mem.Allocator, instance: GridCells) !u64 {
962 var family = try gridCellsShapeFamily(backing_allocator, instance);
963 defer family.deinit();
964 return shape.fingerprint(family);
965 }
966
967 pub fn gridCellsFamilySpecialization(backing_allocator: std.mem.Allocator, instance: GridCells) !entry.OwnedSpecialization {
968 var owned = entry.OwnedSpecialization.init(backing_allocator);
969 errdefer owned.deinit();
970 const lifetime_allocator = owned.allocator();
971
972 const inputs = try lifetime_allocator.alloc(entry.Shape, 2);
973 inputs[0] = try entry.runtimeShape1D(lifetime_allocator, instance.point_axis, instance.count);
974 inputs[1] = try entry.runtimeShape1D(lifetime_allocator, instance.point_axis, instance.count);
975
976 const outputs = try lifetime_allocator.alloc(entry.Shape, 1);
977 outputs[0] = try entry.runtimeShape1D(lifetime_allocator, instance.point_axis, instance.count);
978
979 owned.value = .{
980 .dtype = .f32,
981 .operation = .{ .spatial = .grid_cells },
982 .inputs = inputs,
983 .outputs = outputs,
984 .schedule = try entry.runtimeThreadBlocks1D(lifetime_allocator, instance.point_axis, instance.count, instance.threads),
985 };
986 owned.value.launch = owned.value.schedule.?.launch();
987 var family = try gridCellsShapeFamily(backing_allocator, instance);
988 errdefer family.deinit();
989 try owned.takeShapeFamily(&family);
990 return owned;
991 }
992
993 pub fn gridCellsInstanceFromSpecialization(specialization: entry.Specialization) ?GridCells {
994 if (!specialization.scheduleMatchesLaunch()) return null;
995 const schedule = specialization.schedule orelse return null;
996 if (!specialization.operationIs(.{ .spatial = .grid_cells })) return null;
997 const dtype = specialization.dtype orelse return null;
998 if (dtype != .f32) return null;
999 if (specialization.inputs.len != 2 or specialization.outputs.len != 1) return null;
1000 if (specialization.reductions.len != 0) return null;
1001 const cells = specialization.outputs[0];
1002 if (cells.axes.len != 1) return null;
1003 const point_axis = cells.axes[0].name;
1004 const count = cells.axes[0].extent;
1005 if (!spatialShape1DMatches(specialization.inputs[0], point_axis, count)) return null;
1006 if (!spatialShape1DMatches(specialization.inputs[1], point_axis, count)) return null;
1007 const launch = specialization.launch orelse return null;
1008 if (launch.threadgroup[0] == 0) return null;
1009 const instance = GridCells{
1010 .count = count,
1011 .threads = launch.threadgroup[0],
1012 .point_axis = point_axis,
1013 };
1014 if (!gridCellsInstanceValid(instance)) return null;
1015 if (!spatialThreadBlocksMatch(schedule, instance.point_axis, instance.count, instance.threads)) return null;
1016 return instance;
1017 }
1018
1019 pub fn gridCellsShapeProfileDimensions(instance: GridCells) [1]artifact_product.KernelCallShapeProfileDimension {
1020 return .{
1021 .{ .name = instance.point_axis, .runtime_scalar_argument_index = 0, .bounds = spatialRuntimeExtentBounds() },
1022 };
1023 }
1024
1025 fn spatialDerivedLaunch(threads: u32) !artifact_product.KernelCallLaunch {
1026 if (threads == 0) return error.KernelLibraryLaunchThreadgroupMustBeNonzero;
1027 return .{ .derived = .{
1028 .grid = .{
1029 .{ .runtime_u32_ceil_div = .{ .argument_index = 0, .divisor = threads } },
1030 .{ .fixed = 1 },
1031 .{ .fixed = 1 },
1032 },
1033 .threadgroup = .{ threads, 1, 1 },
1034 } };
1035 }
1036
1037 pub fn createGridCellsFamilyArtifact(
1038 allocator: std.mem.Allocator,
1039 handle: kernel.BackendHandle,
1040 instance: GridCells,
1041 options: entry.ArtifactOptions,
1042 ) !kernel.OwnedKernelCallArtifact {
1043 if (!gridCellsInstanceValid(instance)) return error.InvalidKernelLibraryEntry;
1044 const target = try gridCellsFamilyTarget(allocator, instance);
1045 defer allocator.free(target);
1046 const entry_name = try gridCellsFamilyEntryName(allocator, instance);
1047 defer allocator.free(entry_name);
1048 const family_fingerprint = options.shape_family_fingerprint orelse try gridCellsFamilyFingerprint(allocator, instance);
1049 const shape_profile_dimensions = gridCellsShapeProfileDimensions(instance);
1050 const shape_profile = options.shape_profile orelse artifact_product.KernelCallShapeProfile{
1051 .name = "grid_cells",
1052 .fingerprint = family_fingerprint,
1053 .dimensions = shape_profile_dimensions[0..],
1054 };
1055
1056 var graph = try GridCellsRuntimeFamilyF32.buildNamed(allocator, options.limits, entry_name, instance);
1057 defer graph.deinit();
1058 return kernel.createKernelCallArtifact(allocator, handle, &graph, .{
1059 .target = target,
1060 .version = grid_cells_family_version,
1061 .format = options.format,
1062 .kernel_plan = options.kernel_plan,
1063 .element_count_argument = options.element_count_argument,
1064 .shape_family_fingerprint = family_fingerprint,
1065 .shape_profile = shape_profile,
1066 .launch = options.launch orelse try spatialDerivedLaunch(instance.threads),
1067 .runtime_scalar_argument_count = if (options.runtime_scalar_argument_count == 0) 6 else options.runtime_scalar_argument_count,
1068 .static_arguments = options.static_arguments,
1069 });
1070 }
1071
1072 pub fn gridCountShapeFamily(backing_allocator: std.mem.Allocator, instance: GridCount) !shape.Family {
1073 var builder = try shape.Builder.init(backing_allocator, "grid_count");
1074 errdefer builder.deinit();
1075
1076 const point = try builder.symbol(instance.point_axis);
1077 const cell = try builder.symbol(grid_count_cell_axis);
1078 const block = try builder.symbol(grid_count_block_axis);
1079 const point_expr = try builder.symbolExpression(point);
1080 const cell_expr = try builder.symbolExpression(cell);
1081 const block_expr = try builder.symbolExpression(block);
1082 _ = try builder.tensor("counts", &.{ cell_expr, block_expr });
1083 _ = try builder.tensor("ids", &.{point_expr});
1084 try builder.assumeBounds(point_expr, spatialRuntimeExtentBounds());
1085 try builder.assumeBounds(cell_expr, spatialRuntimeExtentBounds());
1086 try builder.assumeBounds(block_expr, spatialRuntimeExtentBounds());
1087
1088 return builder.finish();
1089 }
1090
1091 pub fn gridCountFamilyFingerprint(backing_allocator: std.mem.Allocator, instance: GridCount) !u64 {
1092 var family = try gridCountShapeFamily(backing_allocator, instance);
1093 defer family.deinit();
1094 return shape.fingerprint(family);
1095 }
1096
1097 pub fn gridCountFamilySpecialization(backing_allocator: std.mem.Allocator, instance: GridCount) !entry.OwnedSpecialization {
1098 var owned = entry.OwnedSpecialization.init(backing_allocator);
1099 errdefer owned.deinit();
1100 const lifetime_allocator = owned.allocator();
1101
1102 const inputs = try lifetime_allocator.alloc(entry.Shape, 1);
1103 inputs[0] = try entry.runtimeShape1D(lifetime_allocator, instance.point_axis, instance.count);
1104
1105 const outputs = try lifetime_allocator.alloc(entry.Shape, 1);
1106 outputs[0] = try entry.runtimeShape2D(
1107 lifetime_allocator,
1108 grid_count_cell_axis,
1109 instance.cells,
1110 grid_count_block_axis,
1111 gridCellsBlockCount(instance.count, instance.threads),
1112 );
1113
1114 owned.value = .{
1115 .dtype = .i32,
1116 .operation = .{ .spatial = .grid_count },
1117 .inputs = inputs,
1118 .outputs = outputs,
1119 .schedule = try entry.runtimeThreadBlocks1D(lifetime_allocator, instance.point_axis, instance.count, instance.threads),
1120 };
1121 owned.value.launch = owned.value.schedule.?.launch();
1122 var family = try gridCountShapeFamily(backing_allocator, instance);
1123 errdefer family.deinit();
1124 try owned.takeShapeFamily(&family);
1125 return owned;
1126 }
1127
1128 pub fn gridCountInstanceFromSpecialization(specialization: entry.Specialization) ?GridCount {
1129 if (!specialization.scheduleMatchesLaunch()) return null;
1130 const schedule = specialization.schedule orelse return null;
1131 if (!specialization.operationIs(.{ .spatial = .grid_count })) return null;
1132 const dtype = specialization.dtype orelse return null;
1133 if (dtype != .i32) return null;
1134 if (specialization.inputs.len != 1 or specialization.outputs.len != 1) return null;
1135 if (specialization.reductions.len != 0) return null;
1136 const ids = specialization.inputs[0];
1137 const counts = specialization.outputs[0];
1138 if (ids.axes.len != 1 or counts.axes.len != 2) return null;
1139 const launch = specialization.launch orelse return null;
1140 if (launch.threadgroup[0] == 0) return null;
1141 const cells = std.math.cast(u32, counts.axes[0].extent) orelse return null;
1142 if (!spatialAxisMatches(counts.axes[0], grid_count_cell_axis, cells)) return null;
1143 const instance = GridCount{
1144 .count = ids.axes[0].extent,
1145 .cells = cells,
1146 .threads = launch.threadgroup[0],
1147 .point_axis = ids.axes[0].name,
1148 };
1149 if (!gridCountInstanceValid(instance)) return null;
1150 const blocks = gridCellsBlockCountChecked(instance.count, instance.threads) orelse return null;
1151 if (!spatialAxisMatches(counts.axes[1], grid_count_block_axis, blocks)) return null;
1152 if (!spatialThreadBlocksMatch(schedule, instance.point_axis, instance.count, instance.threads)) return null;
1153 return instance;
1154 }
1155
1156 pub fn gridCountShapeProfileDimensions(instance: GridCount) [2]artifact_product.KernelCallShapeProfileDimension {
1157 return .{
1158 .{ .name = instance.point_axis, .runtime_scalar_argument_index = 0, .bounds = spatialRuntimeExtentBounds() },
1159 .{ .name = grid_count_cell_axis, .runtime_scalar_argument_index = 1, .bounds = spatialRuntimeExtentBounds() },
1160 };
1161 }
1162
1163 pub fn createGridCountFamilyArtifact(
1164 allocator: std.mem.Allocator,
1165 handle: kernel.BackendHandle,
1166 instance: GridCount,
1167 options: entry.ArtifactOptions,
1168 ) !kernel.OwnedKernelCallArtifact {
1169 if (!gridCountInstanceValid(instance)) return error.InvalidKernelLibraryEntry;
1170 const target = try gridCountFamilyTarget(allocator, instance);
1171 defer allocator.free(target);
1172 const entry_name = try gridCountFamilyEntryName(allocator, instance);
1173 defer allocator.free(entry_name);
1174 const family_fingerprint = options.shape_family_fingerprint orelse try gridCountFamilyFingerprint(allocator, instance);
1175 const shape_profile_dimensions = gridCountShapeProfileDimensions(instance);
1176 const shape_profile = options.shape_profile orelse artifact_product.KernelCallShapeProfile{
1177 .name = "grid_count",
1178 .fingerprint = family_fingerprint,
1179 .dimensions = shape_profile_dimensions[0..],
1180 };
1181
1182 var graph = try GridCountRuntimeFamilyI32.buildNamed(allocator, options.limits, entry_name, instance);
1183 defer graph.deinit();
1184 return kernel.createKernelCallArtifact(allocator, handle, &graph, .{
1185 .target = target,
1186 .version = grid_count_family_version,
1187 .format = options.format,
1188 .kernel_plan = options.kernel_plan,
1189 .element_count_argument = options.element_count_argument,
1190 .shape_family_fingerprint = family_fingerprint,
1191 .shape_profile = shape_profile,
1192 .launch = options.launch orelse try spatialDerivedLaunch(instance.threads),
1193 .runtime_scalar_argument_count = if (options.runtime_scalar_argument_count == 0) 2 else options.runtime_scalar_argument_count,
1194 .static_arguments = options.static_arguments,
1195 });
1196 }
1197
1198 pub fn gridNeighborCountShapeFamily(backing_allocator: std.mem.Allocator, instance: GridNeighborCount) !shape.Family {
1199 var builder = try shape.Builder.init(backing_allocator, "grid_neighbor_count");
1200 errdefer builder.deinit();
1201
1202 const point = try builder.symbol(instance.point_axis);
1203 const offsets = try builder.symbol(grid_neighbor_offsets_axis);
1204 const point_expr = try builder.symbolExpression(point);
1205 const offsets_expr = try builder.symbolExpression(offsets);
1206 _ = try builder.tensor("neighbors", &.{point_expr});
1207 _ = try builder.tensor("x", &.{point_expr});
1208 _ = try builder.tensor("y", &.{point_expr});
1209 _ = try builder.tensor("sorted_indices", &.{point_expr});
1210 _ = try builder.tensor("offsets", &.{offsets_expr});
1211 try builder.assumeBounds(point_expr, spatialRuntimeExtentBounds());
1212 try builder.assumeBounds(offsets_expr, spatialRuntimeExtentBounds());
1213
1214 return builder.finish();
1215 }
1216
1217 pub fn gridNeighborCountFamilyFingerprint(backing_allocator: std.mem.Allocator, instance: GridNeighborCount) !u64 {
1218 var family = try gridNeighborCountShapeFamily(backing_allocator, instance);
1219 defer family.deinit();
1220 return shape.fingerprint(family);
1221 }
1222
1223 pub fn gridNeighborCountFamilySpecialization(
1224 backing_allocator: std.mem.Allocator,
1225 instance: GridNeighborCount,
1226 ) !entry.OwnedSpecialization {
1227 var owned = entry.OwnedSpecialization.init(backing_allocator);
1228 errdefer owned.deinit();
1229 const lifetime_allocator = owned.allocator();
1230
1231 const inputs = try lifetime_allocator.alloc(entry.Shape, 4);
1232 inputs[0] = try entry.runtimeShape1D(lifetime_allocator, instance.point_axis, instance.count);
1233 inputs[1] = try entry.runtimeShape1D(lifetime_allocator, instance.point_axis, instance.count);
1234 inputs[2] = try entry.runtimeShape1D(lifetime_allocator, instance.point_axis, instance.count);
1235 inputs[3] = try entry.runtimeShape1D(lifetime_allocator, grid_neighbor_offsets_axis, instance.offsets_extent);
1236
1237 const outputs = try lifetime_allocator.alloc(entry.Shape, 1);
1238 outputs[0] = try entry.runtimeShape1D(lifetime_allocator, instance.point_axis, instance.count);
1239
1240 owned.value = .{
1241 .dtype = .f32,
1242 .operation = .{ .spatial = .grid_neighbor_count },
1243 .inputs = inputs,
1244 .outputs = outputs,
1245 .schedule = try entry.runtimeThreadBlocks1D(lifetime_allocator, instance.point_axis, instance.count, instance.threads),
1246 };
1247 owned.value.launch = owned.value.schedule.?.launch();
1248 var family = try gridNeighborCountShapeFamily(backing_allocator, instance);
1249 errdefer family.deinit();
1250 try owned.takeShapeFamily(&family);
1251 return owned;
1252 }
1253
1254 pub fn gridNeighborCountInstanceFromSpecialization(specialization: entry.Specialization) ?GridNeighborCount {
1255 if (!specialization.scheduleMatchesLaunch()) return null;
1256 const schedule = specialization.schedule orelse return null;
1257 if (!specialization.operationIs(.{ .spatial = .grid_neighbor_count })) return null;
1258 const dtype = specialization.dtype orelse return null;
1259 if (dtype != .f32) return null;
1260 if (specialization.inputs.len != 4 or specialization.outputs.len != 1) return null;
1261 if (specialization.reductions.len != 0) return null;
1262 const points = specialization.outputs[0];
1263 if (points.axes.len != 1) return null;
1264 const point_axis = points.axes[0].name;
1265 const count = points.axes[0].extent;
1266 if (!spatialShape1DMatches(specialization.inputs[0], point_axis, count)) return null;
1267 if (!spatialShape1DMatches(specialization.inputs[1], point_axis, count)) return null;
1268 if (!spatialShape1DMatches(specialization.inputs[2], point_axis, count)) return null;
1269 const offsets = specialization.inputs[3];
1270 if (offsets.axes.len != 1) return null;
1271 const offsets_extent = offsets.axes[0].extent;
1272 if (!spatialAxisMatches(offsets.axes[0], grid_neighbor_offsets_axis, offsets_extent)) return null;
1273 const launch = specialization.launch orelse return null;
1274 if (launch.threadgroup[0] == 0) return null;
1275 const instance = GridNeighborCount{
1276 .count = count,
1277 .offsets_extent = offsets_extent,
1278 .threads = launch.threadgroup[0],
1279 .point_axis = point_axis,
1280 };
1281 if (!gridNeighborCountInstanceValid(instance)) return null;
1282 if (!spatialThreadBlocksMatch(schedule, instance.point_axis, instance.count, instance.threads)) return null;
1283 return instance;
1284 }
1285
1286 fn spatialAxisMatches(axis: entry.Axis, name: []const u8, extent: u64) bool {
1287 if (name.len == 0) return false;
1288 if (!std.mem.eql(u8, axis.name, name)) return false;
1289 return axis.extent == extent;
1290 }
1291
1292 fn spatialShape1DMatches(candidate: entry.Shape, name: []const u8, extent: u64) bool {
1293 if (candidate.axes.len != 1) return false;
1294 return spatialAxisMatches(candidate.axes[0], name, extent);
1295 }
1296
1297 fn spatialScheduleAxisNameMatches(actual: []const u8, point_axis: []const u8, suffix: []const u8) bool {
1298 if (actual.len != point_axis.len + suffix.len) return false;
1299 return std.mem.eql(u8, actual[0..point_axis.len], point_axis) and
1300 std.mem.eql(u8, actual[point_axis.len..], suffix);
1301 }
1302
1303 fn spatialThreadBlocksMatch(schedule: entry.Schedule, point_axis: []const u8, count: u64, threads: u32) bool {
1304 const blocks = gridCellsBlockCountChecked(count, threads) orelse return false;
1305 if (count <= threads) {
1306 if (schedule.bindings.len != 1) return false;
1307 const binding = schedule.bindings[0];
1308 return binding.target == .thread_x and
1309 spatialAxisMatches(.{ .name = binding.axis, .extent = binding.extent }, point_axis, count);
1310 }
1311 if (schedule.bindings.len != 2) return false;
1312 const tile = schedule.bindings[0];
1313 const lane = schedule.bindings[1];
1314 return tile.target == .block_x and
1315 tile.extent == blocks and
1316 spatialScheduleAxisNameMatches(tile.axis, point_axis, "_tile") and
1317 lane.target == .thread_x and
1318 lane.extent == threads and
1319 spatialScheduleAxisNameMatches(lane.axis, point_axis, "_lane");
1320 }
1321
1322 pub fn gridNeighborCountShapeProfileDimensions(instance: GridNeighborCount) [1]artifact_product.KernelCallShapeProfileDimension {
1323 return .{
1324 .{ .name = instance.point_axis, .runtime_scalar_argument_index = 0, .bounds = spatialRuntimeExtentBounds() },
1325 };
1326 }
1327
1328 pub fn createGridNeighborCountFamilyArtifact(
1329 allocator: std.mem.Allocator,
1330 handle: kernel.BackendHandle,
1331 instance: GridNeighborCount,
1332 options: entry.ArtifactOptions,
1333 ) !kernel.OwnedKernelCallArtifact {
1334 if (!gridNeighborCountInstanceValid(instance)) return error.InvalidKernelLibraryEntry;
1335 const target = try gridNeighborCountFamilyTarget(allocator, instance);
1336 defer allocator.free(target);
1337 const entry_name = try gridNeighborCountFamilyEntryName(allocator, instance);
1338 defer allocator.free(entry_name);
1339 const family_fingerprint = options.shape_family_fingerprint orelse try gridNeighborCountFamilyFingerprint(allocator, instance);
1340 const shape_profile_dimensions = gridNeighborCountShapeProfileDimensions(instance);
1341 const shape_profile = options.shape_profile orelse artifact_product.KernelCallShapeProfile{
1342 .name = "grid_neighbor_count",
1343 .fingerprint = family_fingerprint,
1344 .dimensions = shape_profile_dimensions[0..],
1345 };
1346
1347 var graph = try GridNeighborCountRuntimeFamilyF32.buildNamed(allocator, options.limits, entry_name, instance);
1348 defer graph.deinit();
1349 return kernel.createKernelCallArtifact(allocator, handle, &graph, .{
1350 .target = target,
1351 .version = grid_neighbor_count_family_version,
1352 .format = options.format,
1353 .kernel_plan = options.kernel_plan,
1354 .element_count_argument = options.element_count_argument,
1355 .shape_family_fingerprint = family_fingerprint,
1356 .shape_profile = shape_profile,
1357 .launch = options.launch orelse try spatialDerivedLaunch(instance.threads),
1358 .runtime_scalar_argument_count = if (options.runtime_scalar_argument_count == 0) 8 else options.runtime_scalar_argument_count,
1359 .static_arguments = options.static_arguments,
1360 });
1361 }
1362
1363 test "spatial specializations round-trip their instances" {
1364 const allocator = std.testing.allocator;
1365
1366 var cells_owned = try gridCellsFamilySpecialization(allocator, .{ .count = 5000, .threads = 64 });
1367 defer cells_owned.deinit();
1368 try std.testing.expect(cells_owned.value.operationIs(.{ .spatial = .grid_cells }));
1369 const cells_recovered = gridCellsInstanceFromSpecialization(cells_owned.value) orelse {
1370 return error.TestExpectedGridCellsInstance;
1371 };
1372 try std.testing.expectEqual(@as(u64, 5000), cells_recovered.count);
1373 try std.testing.expectEqual(@as(u32, 64), cells_recovered.threads);
1374 try std.testing.expectEqual(@as(?GridCells, null), gridCellsInstanceFromSpecialization(.{}));
1375
1376 var count_owned = try gridCountFamilySpecialization(allocator, .{ .count = 5000, .cells = 64, .threads = 64 });
1377 defer count_owned.deinit();
1378 const count_recovered = gridCountInstanceFromSpecialization(count_owned.value) orelse {
1379 return error.TestExpectedGridCountInstance;
1380 };
1381 try std.testing.expectEqual(@as(u64, 5000), count_recovered.count);
1382 try std.testing.expectEqual(@as(u32, 64), count_recovered.cells);
1383 try std.testing.expectEqual(@as(u32, 64), count_recovered.threads);
1384 try std.testing.expectEqual(@as(?GridCount, null), gridCountInstanceFromSpecialization(cells_owned.value));
1385
1386 var neighbor_owned = try gridNeighborCountFamilySpecialization(allocator, .{ .count = 5000, .offsets_extent = 64 * 79, .threads = 64 });
1387 defer neighbor_owned.deinit();
1388 const neighbor_recovered = gridNeighborCountInstanceFromSpecialization(neighbor_owned.value) orelse {
1389 return error.TestExpectedGridNeighborCountInstance;
1390 };
1391 try std.testing.expectEqual(@as(u64, 5000), neighbor_recovered.count);
1392 try std.testing.expectEqual(@as(u64, 64 * 79), neighbor_recovered.offsets_extent);
1393 try std.testing.expectEqual(@as(u32, 64), neighbor_recovered.threads);
1394 try std.testing.expectEqual(
1395 @as(?GridNeighborCount, null),
1396 gridNeighborCountInstanceFromSpecialization(count_owned.value),
1397 );
1398 }
1399
1400 test "spatial specializations reject malformed descriptor shapes" {
1401 const allocator = std.testing.allocator;
1402
1403 {
1404 var owned = try gridCellsFamilySpecialization(allocator, .{ .count = 5000, .threads = 64 });
1405 defer owned.deinit();
1406 const lifetime_allocator = owned.allocator();
1407 const inputs = try lifetime_allocator.alloc(entry.Shape, 2);
1408 inputs[0] = try entry.runtimeShape1D(lifetime_allocator, "bad_p", 5000);
1409 inputs[1] = owned.value.inputs[1];
1410 owned.value.inputs = inputs;
1411 try std.testing.expectEqual(@as(?GridCells, null), gridCellsInstanceFromSpecialization(owned.value));
1412 }
1413
1414 {
1415 var owned = try gridCellsFamilySpecialization(allocator, .{ .count = 5000, .threads = 64 });
1416 defer owned.deinit();
1417 const lifetime_allocator = owned.allocator();
1418 owned.value.schedule = try entry.runtimeThreadBlocks1D(lifetime_allocator, "q", 5000, 64);
1419 owned.value.launch = owned.value.schedule.?.launch();
1420 try std.testing.expectEqual(@as(?GridCells, null), gridCellsInstanceFromSpecialization(owned.value));
1421 }
1422
1423 {
1424 var owned = try gridCountFamilySpecialization(allocator, .{ .count = 5000, .cells = 64, .threads = 64 });
1425 defer owned.deinit();
1426 const lifetime_allocator = owned.allocator();
1427 const outputs = try lifetime_allocator.alloc(entry.Shape, 1);
1428 outputs[0] = try entry.runtimeShape2D(lifetime_allocator, grid_count_cell_axis, 64, "bad_b", gridCellsBlockCount(5000, 64));
1429 owned.value.outputs = outputs;
1430 try std.testing.expectEqual(@as(?GridCount, null), gridCountInstanceFromSpecialization(owned.value));
1431 }
1432
1433 {
1434 var owned = try gridNeighborCountFamilySpecialization(allocator, .{ .count = 5000, .offsets_extent = 64 * 79, .threads = 64 });
1435 defer owned.deinit();
1436 const lifetime_allocator = owned.allocator();
1437 const inputs = try lifetime_allocator.alloc(entry.Shape, 4);
1438 inputs[0] = owned.value.inputs[0];
1439 inputs[1] = owned.value.inputs[1];
1440 inputs[2] = try entry.runtimeShape1D(lifetime_allocator, "p", 4999);
1441 inputs[3] = owned.value.inputs[3];
1442 owned.value.inputs = inputs;
1443 try std.testing.expectEqual(@as(?GridNeighborCount, null), gridNeighborCountInstanceFromSpecialization(owned.value));
1444 }
1445
1446 {
1447 var owned = try gridNeighborCountFamilySpecialization(allocator, .{ .count = 5000, .offsets_extent = 64 * 79, .threads = 64 });
1448 defer owned.deinit();
1449 const lifetime_allocator = owned.allocator();
1450 const inputs = try lifetime_allocator.alloc(entry.Shape, 4);
1451 const zero_offsets_axes = try lifetime_allocator.alloc(entry.Axis, 1);
1452 zero_offsets_axes[0] = .{ .name = grid_neighbor_offsets_axis, .extent = 0 };
1453 inputs[0] = owned.value.inputs[0];
1454 inputs[1] = owned.value.inputs[1];
1455 inputs[2] = owned.value.inputs[2];
1456 inputs[3] = .{ .axes = zero_offsets_axes };
1457 owned.value.inputs = inputs;
1458 try std.testing.expectEqual(@as(?GridNeighborCount, null), gridNeighborCountInstanceFromSpecialization(owned.value));
1459 }
1460 }
1461
1462 test "spatial thread selection rounds to warps" {
1463 try std.testing.expectEqual(@as(u32, 32), spatialThreadsForCount(5));
1464 try std.testing.expect(spatialThreadsForCount(8192) % grid_cells_warp_size == 0);
1465 const candidates = spatialThreadCandidatesForCount(8192);
1466 try std.testing.expect(candidates.count >= 2);
1467 for (candidates.slice()) |threads| {
1468 try std.testing.expect(threads % grid_cells_warp_size == 0);
1469 }
1470 }