tiny.accy.kernel.library.compaction

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 kernel = @import("../root.zig");10const tuning = @import("tuning.zig");1112const DType = choir_abi.DType;13const indexExtent = extent_mod.indexExtent;14const runtimeExtentArgument = extent_mod.runtimeExtentArgument;1516pub const Filter = struct {17    extent: u64,18    dtype: DType = .f32,19    predicate: entry.CompactionPredicate = .nonzero,20    threads: u32 = 256,21    element_axis: []const u8 = "n",22    segment_axis: []const u8 = "b",2324    pub fn segments(self: Filter) u64 {25        return ceilDiv(self.extent, self.threads);26    }2728    pub fn padded(self: Filter) u64 {29        return self.extent + self.segments();30    }31};3233pub fn filterPredicateName(predicate: entry.CompactionPredicate) []const u8 {34    return switch (predicate) {35        .nonzero => "nonzero",36        .greater_than => "greater",37    };38}3940pub const filter_family_version: u32 = 1;41pub const filter_warp_size: u32 = 32;42pub const filter_max_threads: u32 = 1024;4344pub fn filterDTypeSupported(dtype: DType) bool {45    return switch (dtype) {46        .f32, .i32 => true,47        else => false,48    };49}5051pub fn filterInstanceValid(instance: Filter) bool {52    if (!filterDTypeSupported(instance.dtype)) return false;53    if (instance.extent == 0) return false;54    if (instance.threads == 0 or instance.threads > filter_max_threads) return false;55    return instance.threads % filter_warp_size == 0;56}5758fn ceilDiv(numerator: u64, denominator: u64) u64 {59    return numerator / denominator + @intFromBool(numerator % denominator != 0);60}6162pub fn filterBlocksExpectedF32(data: []const f32, threads: u32, dst: []f32) void {63    const segment_count = ceilDiv(data.len, threads);64    for (0..segment_count) |segment| {65        const begin = segment * threads;66        const end = @min(begin + threads, data.len);67        var survivors: usize = 0;68        for (data[begin..end]) |value| {69            if (value != 0.0) {70                dst[begin + survivors] = value;71                survivors += 1;72            }73        }74        dst[data.len + segment] = @floatFromInt(survivors);75    }76}7778pub fn filterBlocksExpectedI32(data: []const i32, threads: u32, dst: []i32) void {79    const segment_count = ceilDiv(data.len, threads);80    for (0..segment_count) |segment| {81        const begin = segment * threads;82        const end = @min(begin + threads, data.len);83        var survivors: usize = 0;84        for (data[begin..end]) |value| {85            if (value != 0) {86                dst[begin + survivors] = value;87                survivors += 1;88            }89        }90        dst[data.len + segment] = @intCast(survivors);91    }92}9394pub fn filterBlocksGreaterExpectedF32(data: []const f32, threads: u32, threshold: f32, dst: []f32) void {95    const segment_count = ceilDiv(data.len, threads);96    for (0..segment_count) |segment| {97        const begin = segment * threads;98        const end = @min(begin + threads, data.len);99        var survivors: usize = 0;100        for (data[begin..end]) |value| {101            if (value > threshold) {102                dst[begin + survivors] = value;103                survivors += 1;104            }105        }106        dst[data.len + segment] = @floatFromInt(survivors);107    }108}109110pub fn filterBlocksGreaterExpectedI32(data: []const i32, threads: u32, threshold: i32, dst: []i32) void {111    const segment_count = ceilDiv(data.len, threads);112    for (0..segment_count) |segment| {113        const begin = segment * threads;114        const end = @min(begin + threads, data.len);115        var survivors: usize = 0;116        for (data[begin..end]) |value| {117            if (value > threshold) {118                dst[begin + survivors] = value;119                survivors += 1;120            }121        }122        dst[data.len + segment] = @intCast(survivors);123    }124}125126fn predicateFlag(127    k: anytype,128    comptime dtype: DType,129    comptime predicate: entry.CompactionPredicate,130    args: anytype,131    element: kernel.Value,132) !kernel.Value {133    const survives = switch (predicate) {134        .nonzero => switch (dtype) {135            .f32 => try k.compare(.ne, element, try k.constantFloat(.f32, 0.0)),136            .i32 => try k.compare(.ne, element, try k.constantInt(.i32, 0)),137            else => @compileError("filter kernels support dtype .f32 or .i32"),138        },139        .greater_than => try k.compare(.gt, element, args.param(.threshold).raw()),140    };141    const one = try k.constantInt(.i32, 1);142    const zero = try k.constantInt(.i32, 0);143    return k.select(survives, one, zero);144}145146fn zeroElement(k: anytype, comptime dtype: DType) !kernel.Value {147    return switch (dtype) {148        .f32 => k.constantFloat(.f32, 0.0),149        .i32 => k.constantInt(.i32, 0),150        else => @compileError("filter kernels support dtype .f32 or .i32"),151    };152}153154fn filter_scan_core_seeds_zero(inner: anytype, ctx: anytype) !void {155    try inner.storeIndex(ctx.zero_flag, ctx.warp_sums, ctx.local);156}157158fn filter_scan_core_is_last_lane(inner: anytype, ctx: anytype) !void {159    try inner.storeIndex(ctx.scanned, ctx.warp_sums, ctx.warp);160}161162fn filter_scan_core_is_first_warp(inner: anytype, ctx: anytype) !void {163    const warp_sum = try inner.loadIndex(ctx.warp_sums, ctx.lane);164    const warp_scan = try inner.warpScan(.add, .inclusive, warp_sum);165    try inner.storeIndex(warp_scan, ctx.warp_sums, ctx.lane);166}167168fn filter_scan_core_survives(inner: anytype, ctx: anytype) !void {169    try ctx.args.param(.dst).store(inner, ctx.element, ctx.destination);170}171172fn filter_scan_core_is_last_thread(inner: anytype, ctx: anytype) !void {173    try ctx.args.param(.dst).store(inner, ctx.count_value, ctx.count_slot);174}175176fn filterScanCore(177    k: anytype,178    comptime dtype: DType,179    comptime predicate: entry.CompactionPredicate,180    args: anytype,181    extent: kernel.Value,182) !void {183    const tid = try k.globalId(.x);184    const local = try k.threadId(.x);185    const block = try k.blockId(.x);186    const block_threads = try k.blockDim(.x);187    const lane = try k.laneId();188    const warp = try k.warpId();189    const zero = try k.constantIndex(0);190    const one = try k.constantIndex(1);191192    const in_range = try k.compare(.lt, tid, extent);193    const extent_minus_one = try k.sub(extent, one);194    const clamped_tid = try k.min(tid, extent_minus_one);195    const loaded = try args.param(.data).load(k, clamped_tid);196    const element = try k.select(in_range, loaded.raw(), try zeroElement(k, dtype));197198    const live = try predicateFlag(k, dtype, predicate, args, element);199    const flag = try k.select(in_range, live, try k.constantInt(.i32, 0));200201    const scanned = try k.warpScan(.add, .inclusive, flag);202203    const warp_sums = try k.sharedBuffer(.i32, filter_warp_size);204    const lane_limit = try k.constantIndex(filter_warp_size - 1);205    const warp_count_value = try k.constantIndex(filter_warp_size);206207    const seeds_zero = try k.compare(.lt, local, warp_count_value);208    try k.guardDo(seeds_zero, .{ .warp_sums = warp_sums, .local = local, .zero_flag = try k.constantInt(.i32, 0) }, filter_scan_core_seeds_zero);209    try k.barrier(.block);210211    const is_last_lane = try k.compare(.eq, lane, lane_limit);212    try k.guardDo(is_last_lane, .{ .warp_sums = warp_sums, .warp = warp, .scanned = scanned }, filter_scan_core_is_last_lane);213    try k.barrier(.block);214215    const is_first_warp = try k.compare(.eq, warp, zero);216    try k.guardDo(is_first_warp, .{ .warp_sums = warp_sums, .lane = lane }, filter_scan_core_is_first_warp);217    try k.barrier(.block);218219    const has_base = try k.compare(.gt, warp, zero);220    const warp_minus_one = try k.sub(warp, one);221    const base_index = try k.select(has_base, warp_minus_one, zero);222    const base_loaded = try k.loadIndex(warp_sums, base_index);223    const base = try k.select(has_base, base_loaded, try k.constantInt(.i32, 0));224    const inclusive = try k.add(scanned, base);225226    const segment_base = try k.mul(block, block_threads);227    const survives = try k.compare(.gt, flag, try k.constantInt(.i32, 0));228    const offset = try k.castIndex(try k.sub(inclusive, try k.constantInt(.i32, 1)));229    const destination = try k.add(segment_base, offset);230    try k.guardDo(survives, .{ .args = args, .element = element, .destination = destination }, filter_scan_core_survives);231232    const block_threads_minus_one = try k.sub(block_threads, one);233    const is_last_thread = try k.compare(.eq, local, block_threads_minus_one);234    const count_value = switch (dtype) {235        .f32 => try k.cast(inclusive, .f32),236        .i32 => inclusive,237        else => @compileError("filter kernels support dtype .f32 or .i32"),238    };239    const count_slot = try k.add(extent, block);240    try k.guardDo(is_last_thread, .{ .args = args, .count_value = count_value, .count_slot = count_slot }, filter_scan_core_is_last_thread);241}242243fn filterBody(244    k: anytype,245    comptime dtype: DType,246    comptime predicate: entry.CompactionPredicate,247    spec: Filter,248    args: anytype,249) !void {250    if (!filterInstanceValid(spec)) return error.UnsupportedFilterInstance;251    const extent = try k.constantIndex(try indexExtent(spec.extent));252    try filterScanCore(k, dtype, predicate, args, extent);253}254255fn filterRuntimeBody(256    k: anytype,257    comptime dtype: DType,258    comptime predicate: entry.CompactionPredicate,259    spec: Filter,260    args: anytype,261) !void {262    if (!filterInstanceValid(spec)) return error.UnsupportedFilterInstance;263    const extent = try k.castIndex(args.param(.extent).raw());264    try filterScanCore(k, dtype, predicate, args, extent);265}266267fn filterFamilySchedule(instance: Filter) kernel.logical.schedule.ThreadBlocks {268    return kernel.logical.schedule.threadBlocks(.{ .x = instance.threads });269}270271fn filter_runtime_family_body_f32(k: anytype, spec: Filter, args: anytype) !void {272    try filterRuntimeBody(k, .f32, .nonzero, spec, args);273}274275fn filter_runtime_family_body_i32(k: anytype, spec: Filter, args: anytype) !void {276    try filterRuntimeBody(k, .i32, .nonzero, spec, args);277}278279fn filterRuntimeFamily(comptime dtype: DType) type {280    return kernel.logical.Family(.{281        .name = std.fmt.comptimePrint("accy_kernel_compaction_filter_runtime_{s}", .{dtype.name()}),282        .parameters = .{283            .dst = kernel.dynamicBuffer(dtype),284            .data = kernel.dynamicBuffer(dtype),285            .extent = kernel.scalar(.i32),286        },287        .Instance = Filter,288        .schedule = filterFamilySchedule,289        .body = switch (dtype) {290            .f32 => filter_runtime_family_body_f32,291            .i32 => filter_runtime_family_body_i32,292            else => @compileError("runtime filter supports dtype .f32 or .i32"),293        },294    });295}296297fn filter_greater_runtime_family_body_f32(k: anytype, spec: Filter, args: anytype) !void {298    try filterRuntimeBody(k, .f32, .greater_than, spec, args);299}300301fn filter_greater_runtime_family_body_i32(k: anytype, spec: Filter, args: anytype) !void {302    try filterRuntimeBody(k, .i32, .greater_than, spec, args);303}304305fn filterGreaterRuntimeFamily(comptime dtype: DType) type {306    return kernel.logical.Family(.{307        .name = std.fmt.comptimePrint("accy_kernel_compaction_filter_greater_runtime_{s}", .{dtype.name()}),308        .parameters = .{309            .dst = kernel.dynamicBuffer(dtype),310            .data = kernel.dynamicBuffer(dtype),311            .extent = kernel.scalar(.i32),312            .threshold = kernel.scalar(dtype),313        },314        .Instance = Filter,315        .schedule = filterFamilySchedule,316        .body = switch (dtype) {317            .f32 => filter_greater_runtime_family_body_f32,318            .i32 => filter_greater_runtime_family_body_i32,319            else => @compileError("runtime greater-than filter supports dtype .f32 or .i32"),320        },321    });322}323324pub const FilterRuntimeFamilyF32 = filterRuntimeFamily(.f32);325pub const FilterRuntimeFamilyI32 = filterRuntimeFamily(.i32);326pub const FilterGreaterRuntimeFamilyF32 = filterGreaterRuntimeFamily(.f32);327pub const FilterGreaterRuntimeFamilyI32 = filterGreaterRuntimeFamily(.i32);328329pub fn filterThreadsForExtent(extent: u64) u32 {330    if (extent >= 256) return 256;331    const wide: u64 = extent + filter_warp_size - 1;332    const rounded: u32 = @intCast((wide / filter_warp_size) * filter_warp_size);333    return @max(rounded, filter_warp_size);334}335336pub const FilterThreadCandidates = struct {337    count: usize = 0,338    items: [6]u32 = @as([6]u32, @splat(0)),339340    pub fn slice(self: *const FilterThreadCandidates) []const u32 {341        return self.items[0..self.count];342    }343};344345pub fn filterThreadCandidatesForExtent(extent: u64) FilterThreadCandidates {346    var result = FilterThreadCandidates{};347    if (extent == 0) return result;348    const base = filterThreadsForExtent(extent);349    result.items[result.count] = base;350    result.count += 1;351    var threads: u32 = filter_warp_size;352    while (threads <= filter_max_threads) : (threads *= 2) {353        if (threads == base) continue;354        if (result.count >= result.items.len) break;355        result.items[result.count] = threads;356        result.count += 1;357    }358    return result;359}360361pub fn filterInstanceTarget(allocator: std.mem.Allocator, instance: Filter) ![]u8 {362    return std.fmt.allocPrint(363        allocator,364        "accy.kernel.compaction.filter{d}_{d}_{s}",365        .{ instance.extent, instance.threads, instance.dtype.name() },366    );367}368369pub fn filterInstanceEntryName(allocator: std.mem.Allocator, instance: Filter) ![]u8 {370    return std.fmt.allocPrint(371        allocator,372        "accy_kernel_compaction_filter{d}_{d}_{s}",373        .{ instance.extent, instance.threads, instance.dtype.name() },374    );375}376377pub fn filterFamilyTarget(allocator: std.mem.Allocator, instance: Filter) ![]u8 {378    return std.fmt.allocPrint(379        allocator,380        "accy.kernel.compaction.filter_family_{s}_{d}_{s}",381        .{ filterPredicateName(instance.predicate), instance.threads, instance.dtype.name() },382    );383}384385pub fn filterFamilyEntryName(allocator: std.mem.Allocator, instance: Filter) ![]u8 {386    return std.fmt.allocPrint(387        allocator,388        "accy_kernel_compaction_filter_family_{s}_{d}_{s}",389        .{ filterPredicateName(instance.predicate), instance.threads, instance.dtype.name() },390    );391}392393pub fn filterTuningExtents(instance: Filter) [1]u64 {394    return .{instance.extent};395}396397pub fn filterTuningOperation(instance: Filter) entry.Operation {398    return .{ .compaction = .{ .blocks = instance.predicate } };399}400401pub fn filterFamilyTuningKey(402    backing_allocator: std.mem.Allocator,403    device_fingerprint: u64,404    instance: Filter,405) !tuning.FamilyTuningKey {406    const family_fingerprint = try filterFamilyFingerprint(backing_allocator, instance);407    const extents = filterTuningExtents(instance);408    return tuning.FamilyTuningKey.init(409        device_fingerprint,410        family_fingerprint,411        entry.operationFingerprint(filterTuningOperation(instance)),412        instance.dtype,413        filter_family_version,414        extents[0..],415    ) orelse unreachable;416}417418pub fn filterRuntimeArguments(instance: Filter) ![1]choir_abi.ScalarArgument {419    return .{420        .{ .u32 = try runtimeExtentArgument(instance.extent) },421    };422}423424pub fn filterGreaterRuntimeArguments(425    instance: Filter,426    threshold: choir_abi.ScalarArgument,427) ![2]choir_abi.ScalarArgument {428    return .{429        .{ .u32 = try runtimeExtentArgument(instance.extent) },430        threshold,431    };432}433434pub fn filterRuntimeScalarArgumentCount(instance: Filter) u32 {435    return switch (instance.predicate) {436        .nonzero => 1,437        .greater_than => 2,438    };439}440441fn filterRuntimeExtentBounds() shape.Bounds {442    return .{ .min = 1, .max = extent_mod.runtime_extent_max };443}444445pub fn filterShapeProfileDimensions(instance: Filter) [1]artifact_product.KernelCallShapeProfileDimension {446    return .{447        .{ .name = instance.element_axis, .runtime_scalar_argument_index = 0, .bounds = filterRuntimeExtentBounds() },448    };449}450451fn filterDerivedLaunch(instance: Filter) !artifact_product.KernelCallLaunch {452    if (instance.threads == 0) return error.KernelLibraryLaunchThreadgroupMustBeNonzero;453    return .{ .derived = .{454        .grid = .{455            .{ .runtime_u32_ceil_div = .{ .argument_index = 0, .divisor = instance.threads } },456            .{ .fixed = 1 },457            .{ .fixed = 1 },458        },459        .threadgroup = .{ instance.threads, 1, 1 },460    } };461}462463pub fn createFilterFamilyArtifact(464    allocator: std.mem.Allocator,465    handle: kernel.BackendHandle,466    instance: Filter,467    options: entry.ArtifactOptions,468) !kernel.OwnedKernelCallArtifact {469    const target = try filterFamilyTarget(allocator, instance);470    defer allocator.free(target);471    const entry_name = try filterFamilyEntryName(allocator, instance);472    defer allocator.free(entry_name);473    const family_fingerprint = options.shape_family_fingerprint orelse try filterFamilyFingerprint(allocator, instance);474    const shape_profile_dimensions = filterShapeProfileDimensions(instance);475    const shape_profile = options.shape_profile orelse artifact_product.KernelCallShapeProfile{476        .name = "filter",477        .fingerprint = family_fingerprint,478        .dimensions = shape_profile_dimensions[0..],479    };480481    var graph = switch (instance.predicate) {482        .nonzero => switch (instance.dtype) {483            .f32 => try FilterRuntimeFamilyF32.buildNamed(allocator, options.limits, entry_name, instance),484            .i32 => try FilterRuntimeFamilyI32.buildNamed(allocator, options.limits, entry_name, instance),485            else => return error.UnsupportedDType,486        },487        .greater_than => switch (instance.dtype) {488            .f32 => try FilterGreaterRuntimeFamilyF32.buildNamed(allocator, options.limits, entry_name, instance),489            .i32 => try FilterGreaterRuntimeFamilyI32.buildNamed(allocator, options.limits, entry_name, instance),490            else => return error.UnsupportedDType,491        },492    };493    defer graph.deinit();494    return kernel.createKernelCallArtifact(allocator, handle, &graph, .{495        .target = target,496        .version = filter_family_version,497        .format = options.format,498        .kernel_plan = options.kernel_plan,499        .element_count_argument = options.element_count_argument,500        .shape_family_fingerprint = family_fingerprint,501        .shape_profile = shape_profile,502        .launch = options.launch orelse try filterDerivedLaunch(instance),503        .runtime_scalar_argument_count = if (options.runtime_scalar_argument_count == 0)504            filterRuntimeScalarArgumentCount(instance)505        else506            options.runtime_scalar_argument_count,507        .static_arguments = options.static_arguments,508    });509}510511pub fn filterFamilyFingerprint(backing_allocator: std.mem.Allocator, instance: Filter) !u64 {512    var family = try filterShapeFamily(backing_allocator, instance);513    defer family.deinit();514    return shape.fingerprint(family);515}516517pub fn filterShapeFamily(backing_allocator: std.mem.Allocator, instance: Filter) !shape.Family {518    var builder = try shape.Builder.init(backing_allocator, "filter");519    errdefer builder.deinit();520521    const element = try builder.symbol(instance.element_axis);522    const segment = try builder.symbol(instance.segment_axis);523    const element_expr = try builder.symbolExpression(element);524    const segment_expr = try builder.symbolExpression(segment);525526    const padded_expr = try builder.addExpression(element_expr, segment_expr);527    _ = try builder.tensor("data", &.{element_expr});528    _ = try builder.tensor("out", &.{padded_expr});529    try builder.assumeBounds(element_expr, filterRuntimeExtentBounds());530    try builder.assumeBounds(segment_expr, filterRuntimeExtentBounds());531532    return builder.finish();533}534535pub fn filterFamilySpecialization(backing_allocator: std.mem.Allocator, instance: Filter) !entry.OwnedSpecialization {536    var owned = entry.OwnedSpecialization.init(backing_allocator);537    errdefer owned.deinit();538    const lifetime_allocator = owned.allocator();539540    const inputs = try lifetime_allocator.alloc(entry.Shape, 1);541    inputs[0] = try entry.runtimeShape1D(lifetime_allocator, instance.element_axis, instance.extent);542543    const outputs = try lifetime_allocator.alloc(entry.Shape, 1);544    outputs[0] = try entry.runtimeShape1D(lifetime_allocator, instance.segment_axis, instance.padded());545546    owned.value = .{547        .dtype = instance.dtype,548        .operation = .{ .compaction = .{ .blocks = instance.predicate } },549        .inputs = inputs,550        .outputs = outputs,551        .schedule = try entry.runtimeThreadBlocks1D(lifetime_allocator, "e", instance.extent, instance.threads),552    };553    owned.value.launch = owned.value.schedule.?.launch();554    var family = try filterShapeFamily(backing_allocator, instance);555    errdefer family.deinit();556    try owned.takeShapeFamily(&family);557    return owned;558}559560pub fn filterInstanceFromSpecialization(specialization: entry.Specialization) ?Filter {561    if (!specialization.scheduleMatchesLaunch()) return null;562    const operation = specialization.operation orelse return null;563    const predicate = switch (operation) {564        .compaction => |compaction_operation| switch (compaction_operation) {565            .blocks => |predicate| predicate,566        },567        else => return null,568    };569    const dtype = specialization.dtype orelse return null;570    if (!filterDTypeSupported(dtype)) return null;571    if (specialization.inputs.len != 1 or specialization.outputs.len != 1) return null;572    if (specialization.reductions.len != 0) return null;573    const data = specialization.inputs[0];574    const packed_output = specialization.outputs[0];575    if (data.axes.len != 1 or packed_output.axes.len != 1) return null;576    const launch = specialization.launch orelse return null;577    if (launch.threadgroup[0] == 0) return null;578    const instance = Filter{579        .extent = data.axes[0].extent,580        .dtype = dtype,581        .predicate = predicate,582        .threads = launch.threadgroup[0],583        .element_axis = data.axes[0].name,584        .segment_axis = packed_output.axes[0].name,585    };586    if (!filterInstanceValid(instance)) return null;587    if (packed_output.axes[0].extent != instance.padded()) return null;588    return instance;589}590591fn ceilDivComptime(comptime numerator: u64, comptime denominator: u32) u32 {592    return @intCast(numerator / denominator + @as(u64, @intFromBool(numerator % denominator != 0)));593}594595fn filterSpecialization(comptime spec: Filter) entry.Specialization {596    return .{597        .dtype = spec.dtype,598        .operation = .{ .compaction = .{ .blocks = spec.predicate } },599        .inputs = &.{entry.shape1D(spec.element_axis, spec.extent)},600        .outputs = &.{entry.shape1D(spec.segment_axis, spec.padded())},601        .launch = entry.launch1D(ceilDivComptime(spec.extent, spec.threads), spec.threads),602        .schedule = entry.threadBlocks1D("e", spec.extent, spec.threads),603    };604}605606fn filterProgram(comptime spec: Filter) type {607    const Body = struct {608        fn run(k: anytype, args: anytype) !void {609            try filterBody(k, spec.dtype, spec.predicate, spec, args);610        }611    };612613    return kernel.logical.Program(.{614        .name = std.fmt.comptimePrint(615            "accy_kernel_compaction_filter{}_{}_{s}",616            .{ spec.extent, spec.threads, spec.dtype.name() },617        ),618        .parameters = .{619            .dst = kernel.dynamicBuffer(spec.dtype),620            .data = kernel.dynamicBuffer(spec.dtype),621        },622        .body = Body.run,623    }).withSchedule(kernel.logical.schedule.threadBlocks(.{ .x = spec.threads }));624}625626pub fn filterEntry(comptime spec: Filter) type {627    return entry.Entry(filterProgram(spec), .{628        .target = std.fmt.comptimePrint(629            "accy.kernel.compaction.filter{}_{}_{s}",630            .{ spec.extent, spec.threads, spec.dtype.name() },631        ),632        .layer = .logical,633        .category = .compaction,634        .specialization = filterSpecialization(spec),635    });636}637638pub const Filter8F32 = filterEntry(.{ .extent = 8, .threads = 32 });639640test "compaction filter entry compacts one block on CPU" {641    const allocator = std.testing.allocator;642    var data = [_]f32{ 0.0, 3.5, 0.0, -1.25, 2.0, 0.0, 0.0, 7.0 };643    var dst = @as([9]f32, @splat(-99.0));644645    const ProgramType = filterProgram(.{ .extent = 8, .threads = 32 });646    var graph = try ProgramType.build(allocator, ProgramType.Limits.testing);647    defer graph.deinit();648    try graph.runCpuWithLaunch(allocator, &.{649        kernel.argumentBuffer(f32, dst[0..]),650        kernel.argumentBuffer(f32, data[0..]),651    }, .{652        .grid = .{ 1, 1, 1 },653        .block = .{ 32, 1, 1 },654    });655656    try std.testing.expectEqual(@as(f32, 4.0), dst[8]);657    try std.testing.expectEqualSlices(f32, &.{ 3.5, -1.25, 2.0, 7.0 }, dst[0..4]);658    for (dst[4..8]) |value| try std.testing.expectEqual(@as(f32, -99.0), value);659}660661test "compaction filter runtime family compacts segments with tail" {662    const allocator = std.testing.allocator;663    const compiled = Filter{ .extent = 1, .threads = 32 };664    const runtime = Filter{ .extent = 70, .threads = 32 };665666    var graph = try FilterRuntimeFamilyF32.build(allocator, FilterRuntimeFamilyF32.Limits.testing, compiled);667    defer graph.deinit();668669    var data: [70]f32 = undefined;670    var seed: u32 = 0x2545f491;671    for (&data, 0..) |*value, index| {672        seed ^= seed << 13;673        seed ^= seed >> 17;674        seed ^= seed << 5;675        value.* = if (seed % 3 == 0) 0.0 else @floatFromInt(index + 1);676    }677678    var dst = @as([73]f32, @splat(-1.0));679680    const launch_value = try entry.runtimeLaunch1D(runtime.extent, runtime.threads);681    try graph.runCpuWithLaunch(allocator, &.{682        kernel.argumentBuffer(f32, dst[0..]),683        kernel.argumentBuffer(f32, data[0..]),684        kernel.argumentI32(@intCast(runtime.extent)),685    }, .{686        .grid = launch_value.grid,687        .block = launch_value.threadgroup,688    });689690    var expected_dst = @as([73]f32, @splat(-1.0));691    filterBlocksExpectedF32(data[0..], runtime.threads, expected_dst[0..]);692693    try std.testing.expectEqualSlices(f32, expected_dst[70..], dst[70..]);694    for (0..3) |segment| {695        const begin = segment * 32;696        const survivors: usize = @intFromFloat(expected_dst[70 + segment]);697        try std.testing.expectEqualSlices(f32, expected_dst[begin .. begin + survivors], dst[begin .. begin + survivors]);698    }699}700701test "compaction filter runtime family compacts i32 data" {702    const allocator = std.testing.allocator;703    const compiled = Filter{ .extent = 1, .dtype = .i32, .threads = 32 };704    const runtime = Filter{ .extent = 40, .dtype = .i32, .threads = 32 };705706    var graph = try FilterRuntimeFamilyI32.build(allocator, FilterRuntimeFamilyI32.Limits.testing, compiled);707    defer graph.deinit();708709    var data: [40]i32 = undefined;710    for (&data, 0..) |*value, index| value.* = if (index % 2 == 0) 0 else @intCast(index);711712    var dst = @as([42]i32, @splat(-1));713714    const launch_value = try entry.runtimeLaunch1D(runtime.extent, runtime.threads);715    try graph.runCpuWithLaunch(allocator, &.{716        kernel.argumentBuffer(i32, dst[0..]),717        kernel.argumentBuffer(i32, data[0..]),718        kernel.argumentI32(@intCast(runtime.extent)),719    }, .{720        .grid = launch_value.grid,721        .block = launch_value.threadgroup,722    });723724    var expected_dst = @as([42]i32, @splat(-1));725    filterBlocksExpectedI32(data[0..], runtime.threads, expected_dst[0..]);726727    try std.testing.expectEqualSlices(i32, expected_dst[40..], dst[40..]);728    for (0..2) |segment| {729        const begin = segment * 32;730        const survivors: usize = @intCast(expected_dst[40 + segment]);731        try std.testing.expectEqualSlices(i32, expected_dst[begin .. begin + survivors], dst[begin .. begin + survivors]);732    }733}734735test "compaction filter greater runtime family keeps survivors above the threshold" {736    const allocator = std.testing.allocator;737    const compiled = Filter{ .extent = 1, .predicate = .greater_than, .threads = 32 };738    const runtime = Filter{ .extent = 70, .predicate = .greater_than, .threads = 32 };739740    var graph = try FilterGreaterRuntimeFamilyF32.build(allocator, FilterGreaterRuntimeFamilyF32.Limits.testing, compiled);741    defer graph.deinit();742743    var data: [70]f32 = undefined;744    var seed: u32 = 0x2545f491;745    for (&data, 0..) |*value, index| {746        seed ^= seed << 13;747        seed ^= seed >> 17;748        seed ^= seed << 5;749        const magnitude: f32 = @floatFromInt(index + 1);750        value.* = if (seed % 2 == 0) -magnitude else magnitude;751    }752753    const threshold: f32 = 20.0;754    var dst = @as([73]f32, @splat(-999.0));755756    const launch_value = try entry.runtimeLaunch1D(runtime.extent, runtime.threads);757    try graph.runCpuWithLaunch(allocator, &.{758        kernel.argumentBuffer(f32, dst[0..]),759        kernel.argumentBuffer(f32, data[0..]),760        kernel.argumentI32(@intCast(runtime.extent)),761        kernel.argumentF32(threshold),762    }, .{763        .grid = launch_value.grid,764        .block = launch_value.threadgroup,765    });766767    var expected_dst = @as([73]f32, @splat(-999.0));768    filterBlocksGreaterExpectedF32(data[0..], runtime.threads, threshold, expected_dst[0..]);769770    try std.testing.expectEqualSlices(f32, expected_dst[70..], dst[70..]);771    for (0..3) |segment| {772        const begin = segment * 32;773        const survivors: usize = @intFromFloat(expected_dst[70 + segment]);774        for (dst[begin .. begin + survivors]) |value| try std.testing.expect(value > threshold);775        try std.testing.expectEqualSlices(f32, expected_dst[begin .. begin + survivors], dst[begin .. begin + survivors]);776    }777}778779test "compaction filter greater runtime family filters i32 data" {780    const allocator = std.testing.allocator;781    const compiled = Filter{ .extent = 1, .dtype = .i32, .predicate = .greater_than, .threads = 32 };782    const runtime = Filter{ .extent = 40, .dtype = .i32, .predicate = .greater_than, .threads = 32 };783784    var graph = try FilterGreaterRuntimeFamilyI32.build(allocator, FilterGreaterRuntimeFamilyI32.Limits.testing, compiled);785    defer graph.deinit();786787    var data: [40]i32 = undefined;788    for (&data, 0..) |*value, index| {789        const magnitude: i32 = @intCast(index);790        value.* = if (index % 3 == 0) -magnitude else magnitude;791    }792793    const threshold: i32 = 11;794    var dst = @as([42]i32, @splat(-999));795796    const launch_value = try entry.runtimeLaunch1D(runtime.extent, runtime.threads);797    try graph.runCpuWithLaunch(allocator, &.{798        kernel.argumentBuffer(i32, dst[0..]),799        kernel.argumentBuffer(i32, data[0..]),800        kernel.argumentI32(@intCast(runtime.extent)),801        kernel.argumentI32(threshold),802    }, .{803        .grid = launch_value.grid,804        .block = launch_value.threadgroup,805    });806807    var expected_dst = @as([42]i32, @splat(-999));808    filterBlocksGreaterExpectedI32(data[0..], runtime.threads, threshold, expected_dst[0..]);809810    try std.testing.expectEqualSlices(i32, expected_dst[40..], dst[40..]);811    for (0..2) |segment| {812        const begin = segment * 32;813        const survivors: usize = @intCast(expected_dst[40 + segment]);814        try std.testing.expectEqualSlices(i32, expected_dst[begin .. begin + survivors], dst[begin .. begin + survivors]);815    }816}817818test "compaction filter family instance identity matches fixed entry strings" {819    const instance = Filter{ .extent = 8, .threads = 32 };820821    const target = try filterInstanceTarget(std.testing.allocator, instance);822    defer std.testing.allocator.free(target);823    try std.testing.expectEqualStrings(Filter8F32.target, target);824825    const entry_name = try filterInstanceEntryName(std.testing.allocator, instance);826    defer std.testing.allocator.free(entry_name);827    try std.testing.expectEqualStrings(Filter8F32.name, entry_name);828829    try std.testing.expectEqual(Filter8F32.version, filter_family_version);830831    const fresh = Filter{ .extent = 1 << 20, .threads = 128, .dtype = .i32 };832    const family_target = try filterFamilyTarget(std.testing.allocator, fresh);833    defer std.testing.allocator.free(family_target);834    try std.testing.expectEqualStrings("accy.kernel.compaction.filter_family_nonzero_128_i32", family_target);835}836837test "compaction filter family tuning keys discriminate predicates" {838    const allocator = std.testing.allocator;839    const device = tuning.deviceFingerprint(.{ .identity = .{840        .backend = .cuda,841        .family = .nvidia_cuda,842        .name = "compaction-family-tuning-test-device",843        .vendor_id = 0x10de,844        .device_id = 0x2684,845    } });846847    const nonzero = Filter{ .extent = 4096, .predicate = .nonzero };848    const greater = Filter{ .extent = 4096, .predicate = .greater_than };849    const nonzero_key = try filterFamilyTuningKey(allocator, device, nonzero);850    const greater_key = try filterFamilyTuningKey(allocator, device, greater);851    try std.testing.expect(!nonzero_key.eql(greater_key));852    try std.testing.expectEqual(nonzero_key.family_fingerprint, greater_key.family_fingerprint);853    try std.testing.expect(nonzero_key.operation_fingerprint != greater_key.operation_fingerprint);854855    const other_dtype = try filterFamilyTuningKey(allocator, device, .{ .extent = 4096, .dtype = .i32 });856    try std.testing.expect(!nonzero_key.eql(other_dtype));857858    const repeat_key = try filterFamilyTuningKey(allocator, device, nonzero);859    try std.testing.expect(nonzero_key.eql(repeat_key));860}861862test "compaction filter family artifact carries runtime launch contract" {863    const allocator = std.testing.allocator;864    var state = gpu.recording.BackendState{865        .allocator = allocator,866        .kind = .cuda,867        .format = .cuda_ptx,868    };869    const instance = Filter{ .extent = 4096, .threads = 128 };870871    var family_artifact = try createFilterFamilyArtifact(allocator, state.handle(), instance, .{ .limits = .testing });872    defer family_artifact.deinit();873874    const family_entry = family_artifact.entry();875    try std.testing.expectEqualStrings("accy.kernel.compaction.filter_family_nonzero_128_f32", family_entry.target);876    try std.testing.expectEqualStrings("accy_kernel_compaction_filter_family_nonzero_128_f32", family_entry.entry_name);877    try std.testing.expectEqual(@as(u32, 3), family_entry.argument_count);878    try std.testing.expectEqual(@as(u32, 1), family_entry.runtime_scalar_argument_count);879    try std.testing.expect(family_entry.required_dtypes.contains(.f32));880    try std.testing.expect(family_entry.shape_family_fingerprint != null);881    const profile = family_entry.shape_profile orelse return error.TestExpectedShapeProfile;882    try std.testing.expectEqualStrings("filter", profile.name);883    switch (family_entry.launch) {884        .derived => |launch| {885            try std.testing.expectEqual(@as(u32, 128), launch.threadgroup[0]);886            switch (launch.grid[0]) {887                .runtime_u32_ceil_div => |term| {888                    try std.testing.expectEqual(@as(usize, 0), term.argument_index);889                    try std.testing.expectEqual(@as(u32, 128), term.divisor);890                },891                else => return error.TestExpectedDerivedLaunch,892            }893        },894        else => return error.TestExpectedDerivedLaunch,895    }896}897898test "compaction filter instance round-trips through specialization" {899    const instance = Filter{ .extent = 1000, .threads = 64, .dtype = .i32 };900    var owned = try filterFamilySpecialization(std.testing.allocator, instance);901    defer owned.deinit();902903    const recovered = filterInstanceFromSpecialization(owned.value) orelse return error.TestExpectedFilterInstance;904    try std.testing.expectEqual(instance.extent, recovered.extent);905    try std.testing.expectEqual(instance.dtype, recovered.dtype);906    try std.testing.expectEqual(instance.predicate, recovered.predicate);907    try std.testing.expectEqual(instance.threads, recovered.threads);908    try std.testing.expectEqual(@as(u64, 16), recovered.segments());909910    try std.testing.expectEqual(@as(?Filter, null), filterInstanceFromSpecialization(.{}));911}912913test "compaction filter greater instance keeps its predicate through identity and specialization" {914    const instance = Filter{ .extent = 4096, .predicate = .greater_than, .threads = 128 };915916    const family_target = try filterFamilyTarget(std.testing.allocator, instance);917    defer std.testing.allocator.free(family_target);918    try std.testing.expectEqualStrings("accy.kernel.compaction.filter_family_greater_128_f32", family_target);919920    const family_entry_name = try filterFamilyEntryName(std.testing.allocator, instance);921    defer std.testing.allocator.free(family_entry_name);922    try std.testing.expectEqualStrings("accy_kernel_compaction_filter_family_greater_128_f32", family_entry_name);923924    try std.testing.expectEqual(@as(u32, 2), filterRuntimeScalarArgumentCount(instance));925    try std.testing.expectEqual(@as(u32, 1), filterRuntimeScalarArgumentCount(.{ .extent = 4096 }));926927    var owned = try filterFamilySpecialization(std.testing.allocator, instance);928    defer owned.deinit();929    const recovered = filterInstanceFromSpecialization(owned.value) orelse return error.TestExpectedFilterInstance;930    try std.testing.expectEqual(entry.CompactionPredicate.greater_than, recovered.predicate);931932    const arguments = try filterGreaterRuntimeArguments(instance, .{ .f32 = 0.5 });933    try std.testing.expectEqual(@as(u32, 4096), arguments[0].u32);934    try std.testing.expectEqual(@as(f32, 0.5), arguments[1].f32);935}936937test "compaction filter greater family artifact requires two runtime scalars" {938    const allocator = std.testing.allocator;939    var state = gpu.recording.BackendState{940        .allocator = allocator,941        .kind = .cuda,942        .format = .cuda_ptx,943    };944    const instance = Filter{ .extent = 4096, .predicate = .greater_than, .threads = 128 };945946    var family_artifact = try createFilterFamilyArtifact(allocator, state.handle(), instance, .{ .limits = .testing });947    defer family_artifact.deinit();948949    const family_entry = family_artifact.entry();950    try std.testing.expectEqualStrings("accy.kernel.compaction.filter_family_greater_128_f32", family_entry.target);951    try std.testing.expectEqual(@as(u32, 4), family_entry.argument_count);952    try std.testing.expectEqual(@as(u32, 2), family_entry.runtime_scalar_argument_count);953}954955test "compaction filter thread candidates stay bounded and lead with the default" {956    const candidates = filterThreadCandidatesForExtent(100_000);957    try std.testing.expect(candidates.count > 2);958    try std.testing.expectEqual(filterThreadsForExtent(100_000), candidates.items[0]);959    for (candidates.slice(), 0..) |candidate, index| {960        try std.testing.expect(candidate != 0);961        try std.testing.expect(candidate % filter_warp_size == 0);962        for (candidates.slice()[0..index]) |previous| try std.testing.expect(previous != candidate);963    }964}

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