tiny.accy.kernel.library.stencil

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

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