tiny.accy.kernel.library.spatial

1pub const spatial = @import("spatial.zig");

1const std = @import("std");2const choir_abi = @import("choir_abi");34const artifact_product = @import("../../artifact/model/root.zig");5const shape = @import("../../choir/shape/root.zig");6const entry = @import("entry.zig");7const extent_mod = @import("extent.zig");8const geometry_mod = @import("geometry.zig");9const kernel = @import("../root.zig");1011const DType = choir_abi.DType;12const runtimeExtentArgument = extent_mod.runtimeExtentArgument;1314pub const GridCells = struct {15    count: u64,16    threads: u32 = 256,17    point_axis: []const u8 = "p",18};1920pub const grid_cells_family_version: u32 = 1;21pub const grid_cells_warp_size: u32 = 32;22pub const grid_cells_max_threads: u32 = 1024;23pub const grid_cells_max_blocks: u32 = 1024;2425pub const GridGeometry = struct {26    origin_x: f32,27    origin_y: f32,28    inv_cell_size: f32,29    dims_x: u32,30    dims_y: u32,3132    pub fn cellCount(self: GridGeometry) u64 {33        return @as(u64, self.dims_x) * self.dims_y;34    }35};3637pub fn gridCellsBlockCount(count: u64, threads: u32) u64 {38    return gridCellsBlockCountChecked(count, threads).?;39}4041fn 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}4647pub 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}5455pub 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}6061fn 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}7778fn 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}99100fn 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}110111fn gridCellsFamilySchedule(instance: GridCells) kernel.logical.schedule.ThreadBlocks {112    return kernel.logical.schedule.threadBlocks(.{ .x = instance.threads });113}114115fn 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}134135pub const GridCellsRuntimeFamilyF32 = gridCellsRuntimeFamily();136137pub 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}144145pub 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}152153pub 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}167168const testing = std.testing;169170fn 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}177178test "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);191192    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    }215216    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    });234235    for (xs, ys, cells) |x, y, cell| {236        try testing.expectEqual(hostCellId(geometry, x, y), cell);237    }238    _ = runtime_arguments;239}240241test "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);247248    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 }));252253    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 }));256257    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}267268pub const grid_count_shared_cells_cap: u32 = 4096;269270pub const GridCount = struct {271    count: u64,272    cells: u32,273    threads: u32 = 256,274    point_axis: []const u8 = "p",275};276277pub 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}285286fn 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}290291fn 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}302303fn 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}311312fn 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);322323    _ = 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);328329    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);339340    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}349350fn 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}364365fn gridCountFamilySchedule(instance: GridCount) kernel.logical.schedule.ThreadBlocks {366    return kernel.logical.schedule.threadBlocks(.{ .x = instance.threads });367}368369pub const GridCountRuntimeFamilyI32 = gridCountRuntimeFamily();370371pub 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}378379pub 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}386387pub 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}394395test "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);402403    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    }413414    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    });426427    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}435436test "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 };440441    var ids: [count]i32 = undefined;442    for (&ids, 0..) |*id, index| id.* = @intCast(index % 6);443444    var graph = try GridCountRuntimeFamilyI32.build(allocator, GridCountRuntimeFamilyI32.Limits.testing, instance);445    defer graph.deinit();446447    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..]);464465    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..]);476477    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}489490test "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);496497    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 }));501502    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}506507pub const sort_mod = @import("sort.zig");508pub const scan_mod = @import("scan.zig");509510pub 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};518519pub 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;523524    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;538539    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;541542    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}551552test "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));566567    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);578579    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}587588pub const GridNeighborCount = struct {589    count: u64,590    offsets_extent: u64 = 1,591    threads: u32 = 256,592    point_axis: []const u8 = "p",593};594595pub 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}603604fn 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));631632    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);643644        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);650651        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        );662663        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}683684fn 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}701702fn 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}712713fn gridNeighborCountFamilySchedule(instance: GridNeighborCount) kernel.logical.schedule.ThreadBlocks {714    return kernel.logical.schedule.threadBlocks(.{ .x = instance.threads });715}716717fn 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}740741pub const GridNeighborCountRuntimeFamilyF32 = gridNeighborCountRuntimeFamily();742743pub 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}750751pub 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}758759/// Builds the eight scalar arguments for a neighbor-count kernel launch: the function returns, in760/// order, the point count, the grid origin along x and y, the inverse cell size, the cell counts761/// 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 returns763/// `error.UnsupportedGridGeometry` for a grid with zero cells, an inverse cell size of zero or764/// less, or more cells than a signed 32-bit count. The call returns765/// `error.UnsupportedGridNeighborCountInstance` for a zero stride or a radius that is NaN, infinite766/// or negative, and `error.NeighborRadiusExceedsCellSize` for a radius larger than one cell. The767/// call returns `error.ExtentOverflowsIndexRange` when the point count is zero or too large for the768/// kernel's index range.769pub 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}790791test "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));804805    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    }826827    var host_ids: [count]i32 = undefined;828    for (xs, ys, &host_ids) |x, y, *id| id.* = hostCellId(geometry, x, y);829830    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    }842843    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    }852853    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    }864865    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    });886887    try testing.expectEqualSlices(i32, expected[0..], neighbors[0..]);888    try testing.expect(expected[0] >= 1);889}890891test "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);897898    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 }));902903    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);915916    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}921922pub const grid_count_family_version: u32 = 1;923pub const grid_neighbor_count_family_version: u32 = 1;924pub const grid_count_cell_axis = "c";925pub const grid_count_block_axis = "b";926pub const grid_neighbor_offsets_axis = "o";927const spatial_thread_caps = geometry_mod.ThreadCaps1D{};928929pub 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}934935pub 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}942943fn spatialRuntimeExtentBounds() shape.Bounds {944    return .{ .min = 1, .max = extent_mod.runtime_extent_max };945}946947pub 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();950951    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());957958    return builder.finish();959}960961pub 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}966967pub 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();971972    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);975976    const outputs = try lifetime_allocator.alloc(entry.Shape, 1);977    outputs[0] = try entry.runtimeShape1D(lifetime_allocator, instance.point_axis, instance.count);978979    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}992993pub 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}10181019pub fn gridCellsShapeProfileDimensions(instance: GridCells) [1]artifact_product.KernelCallShapeProfileDimension {1020    return .{1021        .{ .name = instance.point_axis, .runtime_scalar_argument_index = 0, .bounds = spatialRuntimeExtentBounds() },1022    };1023}10241025fn 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}10361037pub 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    };10551056    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}10711072pub 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();10751076    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());10871088    return builder.finish();1089}10901091pub 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}10961097pub 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();11011102    const inputs = try lifetime_allocator.alloc(entry.Shape, 1);1103    inputs[0] = try entry.runtimeShape1D(lifetime_allocator, instance.point_axis, instance.count);11041105    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    );11131114    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}11271128pub 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}11551156pub 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}11621163pub 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    };11811182    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}11971198pub 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();12011202    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());12131214    return builder.finish();1215}12161217pub 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}12221223pub 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();12301231    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);12361237    const outputs = try lifetime_allocator.alloc(entry.Shape, 1);1238    outputs[0] = try entry.runtimeShape1D(lifetime_allocator, instance.point_axis, instance.count);12391240    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}12531254pub 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}12851286fn 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}12911292fn 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}12961297fn 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) and1300        std.mem.eql(u8, actual[point_axis.len..], suffix);1301}13021303fn 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 and1309            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 and1315        tile.extent == blocks and1316        spatialScheduleAxisNameMatches(tile.axis, point_axis, "_tile") and1317        lane.target == .thread_x and1318        lane.extent == threads and1319        spatialScheduleAxisNameMatches(lane.axis, point_axis, "_lane");1320}13211322pub fn gridNeighborCountShapeProfileDimensions(instance: GridNeighborCount) [1]artifact_product.KernelCallShapeProfileDimension {1323    return .{1324        .{ .name = instance.point_axis, .runtime_scalar_argument_index = 0, .bounds = spatialRuntimeExtentBounds() },1325    };1326}13271328pub 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    };13461347    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}13621363test "spatial specializations round-trip their instances" {1364    const allocator = std.testing.allocator;13651366    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(.{}));13751376    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));13851386    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}13991400test "spatial specializations reject malformed descriptor shapes" {1401    const allocator = std.testing.allocator;14021403    {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    }14131414    {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    }14221423    {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    }14321433    {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    }14451446    {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}14611462test "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}