tiny.accy.kernel.library.image

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");1011pub const blur_family_version: u32 = 1;12pub const resize_family_version: u32 = 1;13pub const blur_radius_max: u32 = 15;14pub const image_thread_caps = geometry_mod.ThreadCaps{15    .budget = 256,16    .x_max = 64,17    .y_max = 16,18};1920pub const Axis = enum {21    horizontal,22    vertical,2324    pub fn name(self: Axis) []const u8 {25        return switch (self) {26            .horizontal => "h",27            .vertical => "v",28        };29    }30};3132pub const BlurPass = struct {33    radius: u32,34    axis: Axis,35    width: u64 = 1,36    height: u64 = 1,37    threads: entry.Threads2D = .{ .x = 16, .y = 16 },3839    pub fn taps(self: BlurPass) u32 {40        return self.radius * 2 + 1;41    }42};4344pub const Resize = struct {45    dst_width: u64 = 1,46    dst_height: u64 = 1,47    src_width: u64 = 1,48    src_height: u64 = 1,49    threads: entry.Threads2D = .{ .x = 16, .y = 16 },50};5152pub fn imageExtentValid(extent: u64) bool {53    return extent >= 1 and extent <= extent_mod.runtime_extent_max;54}5556pub fn imageThreadsForExtents(width: u64, height: u64) entry.Threads2D {57    return geometry_mod.threadsForGrid(.{ .rows = height, .cols = width }, image_thread_caps);58}5960pub fn imageThreadCandidatesForExtents(width: u64, height: u64) geometry_mod.ThreadCandidates {61    return geometry_mod.threadCandidatesForGrid(.{ .rows = height, .cols = width }, image_thread_caps);62}6364pub fn blurPassInstanceValid(instance: BlurPass) bool {65    if (instance.radius == 0 or instance.radius > blur_radius_max) return false;66    if (!imageExtentValid(instance.width) or !imageExtentValid(instance.height)) return false;67    if (instance.threads.x == 0 or instance.threads.y == 0) return false;68    return true;69}7071pub fn resizeInstanceValid(instance: Resize) bool {72    if (!imageExtentValid(instance.dst_width) or !imageExtentValid(instance.dst_height)) return false;73    if (!imageExtentValid(instance.src_width) or !imageExtentValid(instance.src_height)) return false;74    if (instance.threads.x == 0 or instance.threads.y == 0) return false;75    return true;76}7778pub fn gaussianWeightsAlloc(allocator: std.mem.Allocator, radius: u32, sigma: f32) ![]f32 {79    if (radius == 0 or radius > blur_radius_max) return error.UnsupportedBlurRadius;80    if (!(sigma > 0)) return error.UnsupportedBlurSigma;81    const taps = radius * 2 + 1;82    const weights = try allocator.alloc(f32, taps);83    errdefer allocator.free(weights);84    _ = try gaussianWeights(weights, radius, sigma);85    return weights;86}8788pub fn gaussianWeights(out: []f32, radius: u32, sigma: f32) ![]f32 {89    if (radius == 0 or radius > blur_radius_max) return error.UnsupportedBlurRadius;90    if (!(sigma > 0)) return error.UnsupportedBlurSigma;91    const taps = radius * 2 + 1;92    if (out.len < taps) return error.BufferTooSmall;93    const weights = out[0..taps];94    var total: f32 = 0;95    for (weights, 0..) |*weight, tap| {96        const offset = @as(f32, @floatFromInt(@as(i64, @intCast(tap)) - @as(i64, radius)));97        const value = @exp(-(offset * offset) / (2 * sigma * sigma));98        weight.* = value;99        total += value;100    }101    for (weights) |*weight| weight.* /= total;102    return weights;103}104105const Channels = struct {106    r: kernel.Value,107    g: kernel.Value,108    b: kernel.Value,109    a: kernel.Value,110};111112fn unpackChannels(inner: anytype, pixel: kernel.Value) !Channels {113    const mask = try inner.constantInt(.u32, 0xff);114    return .{115        .r = try inner.cast(try inner.and_(pixel, mask), .f32),116        .g = try inner.cast(try inner.and_(try inner.ushr(pixel, try inner.constantInt(.u32, 8)), mask), .f32),117        .b = try inner.cast(try inner.and_(try inner.ushr(pixel, try inner.constantInt(.u32, 16)), mask), .f32),118        .a = try inner.cast(try inner.ushr(pixel, try inner.constantInt(.u32, 24)), .f32),119    };120}121122fn packChannel(inner: anytype, value: kernel.Value) !kernel.Value {123    const zero = try inner.constantFloat(.f32, 0);124    const limit = try inner.constantFloat(.f32, 255);125    const half = try inner.constantFloat(.f32, 0.5);126    const rounded = try inner.floor(try inner.add(value, half));127    const clamped = try inner.min(try inner.max(rounded, zero), limit);128    return inner.cast(clamped, .u32);129}130131fn packChannels(inner: anytype, channels: Channels) !kernel.Value {132    const r = try packChannel(inner, channels.r);133    const g = try packChannel(inner, channels.g);134    const b = try packChannel(inner, channels.b);135    const a = try packChannel(inner, channels.a);136    return inner.or_(137        try inner.or_(r, try inner.shl(g, try inner.constantInt(.u32, 8))),138        try inner.or_(try inner.shl(b, try inner.constantInt(.u32, 16)), try inner.shl(a, try inner.constantInt(.u32, 24))),139    );140}141142fn blur_pass_body_active(inner: anytype, ctx: anytype) !void {143    const zero_i32 = try inner.constantInt(.i32, 0);144    const axis_extent = switch (ctx.spec.axis) {145        .horizontal => ctx.width,146        .vertical => ctx.height,147    };148    const axis_base = switch (ctx.spec.axis) {149        .horizontal => ctx.col,150        .vertical => ctx.row,151    };152    const axis_last = try inner.sub(try inner.cast(axis_extent, .i32), try inner.constantInt(.i32, 1));153    const base_i32 = try inner.cast(axis_base, .i32);154155    var acc = Channels{156        .r = try inner.constantFloat(.f32, 0),157        .g = try inner.constantFloat(.f32, 0),158        .b = try inner.constantFloat(.f32, 0),159        .a = try inner.constantFloat(.f32, 0),160    };161    const tap_count = ctx.spec.taps();162    var tap: u32 = 0;163    while (tap < tap_count) : (tap += 1) {164        const offset = @as(i64, tap) - @as(i64, ctx.spec.radius);165        const offset_value = try inner.constantInt(.i32, offset);166        const sample_raw = try inner.add(base_i32, offset_value);167        const sample_clamped = try inner.min(try inner.max(sample_raw, zero_i32), axis_last);168        const sample_index = try inner.castIndex(sample_clamped);169        const pixel_index = switch (ctx.spec.axis) {170            .horizontal => try inner.add(try inner.mul(ctx.row, ctx.width), sample_index),171            .vertical => try inner.add(try inner.mul(sample_index, ctx.width), ctx.col),172        };173        const pixel = (try ctx.args.param(.src).load(inner, pixel_index)).raw();174        const weight = (try ctx.args.param(.weights).load(inner, try inner.constantIndex(tap))).raw();175        const channels = try unpackChannels(inner, pixel);176        acc = .{177            .r = try inner.fma(channels.r, weight, acc.r),178            .g = try inner.fma(channels.g, weight, acc.g),179            .b = try inner.fma(channels.b, weight, acc.b),180            .a = try inner.fma(channels.a, weight, acc.a),181        };182    }183    const out_index = try inner.add(try inner.mul(ctx.row, ctx.width), ctx.col);184    try ctx.args.param(.dst).store(inner, try packChannels(inner, acc), out_index);185}186187fn blurPassBody(k: anytype, spec: BlurPass, args: anytype) !void {188    if (!blurPassInstanceValid(spec)) return error.UnsupportedBlurPassInstance;189    const col = try k.globalId(.x);190    const row = try k.globalId(.y);191    const width = try k.castIndex(args.param(.width).raw());192    const height = try k.castIndex(args.param(.height).raw());193    const col_active = try k.compare(.lt, col, width);194    const row_active = try k.compare(.lt, row, height);195    const active = try k.and_(col_active, row_active);196    try k.guardDo(active, .{197        .args = args,198        .spec = spec,199        .col = col,200        .row = row,201        .width = width,202        .height = height,203    }, blur_pass_body_active);204}205206fn blurPassSchedule(instance: BlurPass) kernel.logical.schedule.ThreadBlocks {207    return kernel.logical.schedule.threadBlocks(.{208        .x = instance.threads.x,209        .y = instance.threads.y,210    });211}212213fn blurPassFamily() type {214    return kernel.logical.Family(.{215        .name = "accy_kernel_image_blur_pass_rgba8",216        .parameters = .{217            .dst = kernel.dynamicBuffer(.u32),218            .src = kernel.dynamicBuffer(.u32),219            .weights = kernel.dynamicBuffer(.f32),220            .width = kernel.scalar(.i32),221            .height = kernel.scalar(.i32),222        },223        .Instance = BlurPass,224        .schedule = blurPassSchedule,225        .body = blurPassBody,226    });227}228229pub const BlurPassFamilyRgba8 = blurPassFamily();230231fn resize_body_active(inner: anytype, ctx: anytype) !void {232    const sample_x = try sampleAxis(inner, ctx.col, ctx.args.param(.scale_x).raw(), ctx.args.param(.src_last_x).raw());233    const sample_y = try sampleAxis(inner, ctx.row, ctx.args.param(.scale_y).raw(), ctx.args.param(.src_last_y).raw());234    const src_width_index = try inner.castIndex(ctx.args.param(.src_width).raw());235236    const row0_base = try inner.mul(sample_y.lo, src_width_index);237    const row1_base = try inner.mul(sample_y.hi, src_width_index);238    const p00 = try unpackChannels(inner, (try ctx.args.param(.src).load(inner, try inner.add(row0_base, sample_x.lo))).raw());239    const p01 = try unpackChannels(inner, (try ctx.args.param(.src).load(inner, try inner.add(row0_base, sample_x.hi))).raw());240    const p10 = try unpackChannels(inner, (try ctx.args.param(.src).load(inner, try inner.add(row1_base, sample_x.lo))).raw());241    const p11 = try unpackChannels(inner, (try ctx.args.param(.src).load(inner, try inner.add(row1_base, sample_x.hi))).raw());242243    const top = try lerpChannels(inner, p00, p01, sample_x.frac);244    const bottom = try lerpChannels(inner, p10, p11, sample_x.frac);245    const blended = try lerpChannels(inner, top, bottom, sample_y.frac);246247    const out_index = try inner.add(try inner.mul(ctx.row, ctx.dst_width), ctx.col);248    try ctx.args.param(.dst).store(inner, try packChannels(inner, blended), out_index);249}250251fn resizeBody(k: anytype, spec: Resize, args: anytype) !void {252    if (!resizeInstanceValid(spec)) return error.UnsupportedResizeInstance;253    const col = try k.globalId(.x);254    const row = try k.globalId(.y);255    const dst_width = try k.castIndex(args.param(.dst_width).raw());256    const dst_height = try k.castIndex(args.param(.dst_height).raw());257    const col_active = try k.compare(.lt, col, dst_width);258    const row_active = try k.compare(.lt, row, dst_height);259    const active = try k.and_(col_active, row_active);260    try k.guardDo(active, .{261        .args = args,262        .col = col,263        .row = row,264        .dst_width = dst_width,265    }, resize_body_active);266}267268const AxisSample = struct {269    lo: kernel.Value,270    hi: kernel.Value,271    frac: kernel.Value,272};273274fn sampleAxis(inner: anytype, dst_index: kernel.Value, scale: kernel.Value, last: kernel.Value) !AxisSample {275    const zero = try inner.constantFloat(.f32, 0);276    const half = try inner.constantFloat(.f32, 0.5);277    const one = try inner.constantFloat(.f32, 1);278    const negative_half = try inner.constantFloat(.f32, -0.5);279    const centered = try inner.fma(try inner.add(try inner.cast(try inner.cast(dst_index, .u32), .f32), half), scale, negative_half);280    const clamped = try inner.min(try inner.max(centered, zero), last);281    const lo_f = try inner.floor(clamped);282    const frac = try inner.sub(clamped, lo_f);283    const lo_u32 = try inner.cast(lo_f, .u32);284    const lo = try inner.castIndex(lo_u32);285    const hi_f = try inner.min(try inner.add(lo_f, one), last);286    const hi = try inner.castIndex(try inner.cast(hi_f, .u32));287    return .{ .lo = lo, .hi = hi, .frac = frac };288}289290fn lerpChannels(inner: anytype, from: Channels, to: Channels, t: kernel.Value) !Channels {291    return .{292        .r = try lerp(inner, from.r, to.r, t),293        .g = try lerp(inner, from.g, to.g, t),294        .b = try lerp(inner, from.b, to.b, t),295        .a = try lerp(inner, from.a, to.a, t),296    };297}298299fn lerp(inner: anytype, from: kernel.Value, to: kernel.Value, t: kernel.Value) !kernel.Value {300    return inner.fma(try inner.sub(to, from), t, from);301}302303fn resizeSchedule(instance: Resize) kernel.logical.schedule.ThreadBlocks {304    return kernel.logical.schedule.threadBlocks(.{305        .x = instance.threads.x,306        .y = instance.threads.y,307    });308}309310fn resizeFamily() type {311    return kernel.logical.Family(.{312        .name = "accy_kernel_image_resize_bilinear_rgba8",313        .parameters = .{314            .dst = kernel.dynamicBuffer(.u32),315            .src = kernel.dynamicBuffer(.u32),316            .dst_width = kernel.scalar(.i32),317            .dst_height = kernel.scalar(.i32),318            .src_width = kernel.scalar(.i32),319            .scale_x = kernel.scalar(.f32),320            .scale_y = kernel.scalar(.f32),321            .src_last_x = kernel.scalar(.f32),322            .src_last_y = kernel.scalar(.f32),323        },324        .Instance = Resize,325        .schedule = resizeSchedule,326        .body = resizeBody,327    });328}329330pub const ResizeBilinearFamilyRgba8 = resizeFamily();331332pub fn blurPassFamilyTarget(allocator: std.mem.Allocator, instance: BlurPass) ![]u8 {333    return std.fmt.allocPrint(334        allocator,335        "image_blur_pass_family_r{d}{s}_{d}x{d}_rgba8",336        .{ instance.radius, instance.axis.name(), instance.threads.x, instance.threads.y },337    );338}339340pub fn blurPassFamilyEntryName(allocator: std.mem.Allocator, instance: BlurPass) ![]u8 {341    return std.fmt.allocPrint(342        allocator,343        "accy_image_blur_pass_r{d}{s}_rgba8",344        .{ instance.radius, instance.axis.name() },345    );346}347348pub fn resizeFamilyTarget(allocator: std.mem.Allocator, instance: Resize) ![]u8 {349    return std.fmt.allocPrint(350        allocator,351        "image_resize_bilinear_family_{d}x{d}_rgba8",352        .{ instance.threads.x, instance.threads.y },353    );354}355356pub fn resizeFamilyEntryName(allocator: std.mem.Allocator, instance: Resize) ![]u8 {357    _ = instance;358    return allocator.dupe(u8, "accy_image_resize_bilinear_rgba8");359}360361pub fn blurPassRuntimeArguments(width: u32, height: u32) ![2]choir_abi.ScalarArgument {362    if (width == 0 or height == 0) return error.UnsupportedImageExtent;363    if (width > std.math.maxInt(i32) or height > std.math.maxInt(i32)) return error.UnsupportedImageExtent;364    return .{365        .{ .i32 = @intCast(width) },366        .{ .i32 = @intCast(height) },367    };368}369370pub fn resizeScale(src_extent: u32, dst_extent: u32) f32 {371    return @as(f32, @floatFromInt(src_extent)) / @as(f32, @floatFromInt(dst_extent));372}373374pub fn resizeRuntimeArguments(dst_width: u32, dst_height: u32, src_width: u32, src_height: u32) ![7]choir_abi.ScalarArgument {375    if (dst_width == 0 or dst_height == 0 or src_width == 0 or src_height == 0) return error.UnsupportedImageExtent;376    if (dst_width > std.math.maxInt(i32) or dst_height > std.math.maxInt(i32)) return error.UnsupportedImageExtent;377    if (src_width > std.math.maxInt(i32) or src_height > std.math.maxInt(i32)) return error.UnsupportedImageExtent;378    return .{379        .{ .i32 = @intCast(dst_width) },380        .{ .i32 = @intCast(dst_height) },381        .{ .i32 = @intCast(src_width) },382        .{ .f32 = resizeScale(src_width, dst_width) },383        .{ .f32 = resizeScale(src_height, dst_height) },384        .{ .f32 = @floatFromInt(src_width - 1) },385        .{ .f32 = @floatFromInt(src_height - 1) },386    };387}388389pub fn imageLaunchGeometry(instance_threads: entry.Threads2D, width: u32, height: u32) choir_abi.LaunchGeometry {390    return .{391        .grid = .{392            (width + instance_threads.x - 1) / instance_threads.x,393            (height + instance_threads.y - 1) / instance_threads.y,394            1,395        },396        .threadgroup = .{ instance_threads.x, instance_threads.y, 1 },397    };398}399400pub fn blurPassAxisParameter(axis: Axis) u64 {401    return switch (axis) {402        .horizontal => 0,403        .vertical => 1,404    };405}406407pub fn blurPassShapeProfileDimensions(instance: BlurPass) [2]artifact_product.KernelCallShapeProfileDimension {408    _ = instance;409    const bounds = imageRuntimeExtentBounds();410    return .{411        .{412            .name = "x",413            .runtime_scalar_argument_index = 0,414            .bounds = bounds,415        },416        .{417            .name = "y",418            .runtime_scalar_argument_index = 1,419            .bounds = bounds,420        },421    };422}423424pub fn resizeShapeProfileDimensions(instance: Resize) [3]artifact_product.KernelCallShapeProfileDimension {425    _ = instance;426    const bounds = imageRuntimeExtentBounds();427    return .{428        .{429            .name = "dst_x",430            .runtime_scalar_argument_index = 0,431            .bounds = bounds,432        },433        .{434            .name = "dst_y",435            .runtime_scalar_argument_index = 1,436            .bounds = bounds,437        },438        .{439            .name = "src_x",440            .runtime_scalar_argument_index = 2,441            .bounds = bounds,442        },443    };444}445446fn imageRuntimeExtentBounds() shape.Bounds {447    return .{ .min = 1, .max = extent_mod.runtime_extent_max };448}449450fn imageDerivedLaunch(threads: entry.Threads2D) !artifact_product.KernelCallLaunch {451    if (threads.x == 0 or threads.y == 0) return error.KernelLibraryLaunchThreadgroupMustBeNonzero;452    return .{ .derived = .{453        .grid = .{454            .{ .runtime_u32_ceil_div = .{ .argument_index = 0, .divisor = threads.x } },455            .{ .runtime_u32_ceil_div = .{ .argument_index = 1, .divisor = threads.y } },456            .{ .fixed = 1 },457        },458        .threadgroup = .{ threads.x, threads.y, 1 },459    } };460}461462pub fn blurPassShapeFamily(backing_allocator: std.mem.Allocator, instance: BlurPass) !shape.Family {463    if (!blurPassInstanceValid(instance)) return error.UnsupportedBlurPassInstance;464    var builder = try shape.Builder.init(backing_allocator, "image_blur_pass");465    errdefer builder.deinit();466467    const y = try builder.symbol("y");468    const x = try builder.symbol("x");469    const y_expr = try builder.symbolExpression(y);470    const x_expr = try builder.symbolExpression(x);471    const taps_expr = builder.constantExpression(@intCast(instance.taps()));472473    _ = try builder.tensor("src", &.{ y_expr, x_expr });474    _ = try builder.tensor("weights", &.{taps_expr});475    _ = try builder.tensor("out", &.{ y_expr, x_expr });476    try builder.assumeBounds(y_expr, imageRuntimeExtentBounds());477    try builder.assumeBounds(x_expr, imageRuntimeExtentBounds());478479    return builder.finish();480}481482pub fn resizeShapeFamily(backing_allocator: std.mem.Allocator, instance: Resize) !shape.Family {483    if (!resizeInstanceValid(instance)) return error.UnsupportedResizeInstance;484    var builder = try shape.Builder.init(backing_allocator, "image_resize_bilinear");485    errdefer builder.deinit();486487    const dst_y = try builder.symbol("dst_y");488    const dst_x = try builder.symbol("dst_x");489    const src_y = try builder.symbol("src_y");490    const src_x = try builder.symbol("src_x");491    const dst_y_expr = try builder.symbolExpression(dst_y);492    const dst_x_expr = try builder.symbolExpression(dst_x);493    const src_y_expr = try builder.symbolExpression(src_y);494    const src_x_expr = try builder.symbolExpression(src_x);495496    _ = try builder.tensor("src", &.{ src_y_expr, src_x_expr });497    _ = try builder.tensor("out", &.{ dst_y_expr, dst_x_expr });498    try builder.assumeBounds(dst_y_expr, imageRuntimeExtentBounds());499    try builder.assumeBounds(dst_x_expr, imageRuntimeExtentBounds());500    try builder.assumeBounds(src_y_expr, imageRuntimeExtentBounds());501    try builder.assumeBounds(src_x_expr, imageRuntimeExtentBounds());502503    return builder.finish();504}505506pub fn blurPassFamilyFingerprint(backing_allocator: std.mem.Allocator, instance: BlurPass) !u64 {507    var family = try blurPassShapeFamily(backing_allocator, instance);508    defer family.deinit();509    return shape.fingerprint(family);510}511512pub fn resizeFamilyFingerprint(backing_allocator: std.mem.Allocator, instance: Resize) !u64 {513    var family = try resizeShapeFamily(backing_allocator, instance);514    defer family.deinit();515    return shape.fingerprint(family);516}517518pub fn blurPassFamilySpecialization(backing_allocator: std.mem.Allocator, instance: BlurPass) !entry.OwnedSpecialization {519    if (!blurPassInstanceValid(instance)) return error.UnsupportedBlurPassInstance;520    var owned = entry.OwnedSpecialization.init(backing_allocator);521    errdefer owned.deinit();522    const lifetime_allocator = owned.allocator();523524    const inputs = try lifetime_allocator.alloc(entry.Shape, 2);525    inputs[0] = try entry.runtimeShape2D(lifetime_allocator, "y", instance.height, "x", instance.width);526    inputs[1] = try entry.runtimeShape1D(lifetime_allocator, "tap", instance.taps());527528    const outputs = try lifetime_allocator.alloc(entry.Shape, 1);529    outputs[0] = try entry.runtimeShape2D(lifetime_allocator, "y", instance.height, "x", instance.width);530531    const reductions = try lifetime_allocator.alloc(entry.Reduction, 1);532    reductions[0] = try entry.runtimeReduction(533        lifetime_allocator,534        "blur_tap",535        .weighted_sum,536        try entry.runtimeShape1D(lifetime_allocator, "tap", instance.taps()),537    );538539    const static_parameters = try lifetime_allocator.alloc(entry.StaticParameter, 2);540    static_parameters[0] = try entry.runtimeStaticParameter(lifetime_allocator, "radius", instance.radius);541    static_parameters[1] = try entry.runtimeStaticParameter(lifetime_allocator, "axis", blurPassAxisParameter(instance.axis));542543    owned.value = .{544        .dtype = .u32,545        .accumulation_dtype = .f32,546        .operation = .{ .image = .blur_pass },547        .inputs = inputs,548        .outputs = outputs,549        .reductions = reductions,550        .static_parameters = static_parameters,551        .schedule = try entry.runtimeThreadBlocks2D(552            lifetime_allocator,553            "x",554            instance.width,555            "y",556            instance.height,557            instance.threads.x,558            instance.threads.y,559        ),560    };561    owned.value.launch = owned.value.schedule.?.launch();562    var family = try blurPassShapeFamily(backing_allocator, instance);563    errdefer family.deinit();564    try owned.takeShapeFamily(&family);565    return owned;566}567568pub fn resizeFamilySpecialization(backing_allocator: std.mem.Allocator, instance: Resize) !entry.OwnedSpecialization {569    if (!resizeInstanceValid(instance)) return error.UnsupportedResizeInstance;570    var owned = entry.OwnedSpecialization.init(backing_allocator);571    errdefer owned.deinit();572    const lifetime_allocator = owned.allocator();573574    const inputs = try lifetime_allocator.alloc(entry.Shape, 1);575    inputs[0] = try entry.runtimeShape2D(lifetime_allocator, "src_y", instance.src_height, "src_x", instance.src_width);576577    const outputs = try lifetime_allocator.alloc(entry.Shape, 1);578    outputs[0] = try entry.runtimeShape2D(lifetime_allocator, "dst_y", instance.dst_height, "dst_x", instance.dst_width);579580    owned.value = .{581        .dtype = .u32,582        .accumulation_dtype = .f32,583        .operation = .{ .image = .resize_bilinear },584        .inputs = inputs,585        .outputs = outputs,586        .schedule = try entry.runtimeThreadBlocks2D(587            lifetime_allocator,588            "dst_x",589            instance.dst_width,590            "dst_y",591            instance.dst_height,592            instance.threads.x,593            instance.threads.y,594        ),595    };596    owned.value.launch = owned.value.schedule.?.launch();597    var family = try resizeShapeFamily(backing_allocator, instance);598    errdefer family.deinit();599    try owned.takeShapeFamily(&family);600    return owned;601}602603pub fn blurPassInstanceFromSpecialization(specialization: entry.Specialization) ?BlurPass {604    if (!specialization.scheduleMatchesLaunch()) return null;605    if (!specialization.operationIs(.{ .image = .blur_pass })) return null;606    if (specialization.dtype != .u32 or specialization.accumulation_dtype != .f32) return null;607    if (specialization.inputs.len != 2 or specialization.outputs.len != 1) return null;608    if (specialization.reductions.len != 1 or specialization.static_parameters.len != 2) return null;609610    const src = specialization.inputs[0];611    const weights = specialization.inputs[1];612    const output = specialization.outputs[0];613    const reduction = specialization.reductions[0];614    if (src.axes.len != 2 or weights.axes.len != 1 or output.axes.len != 2) return null;615    if (reduction.shape.axes.len != 1) return null;616    if (!std.mem.eql(u8, src.axes[0].name, output.axes[0].name)) return null;617    if (!std.mem.eql(u8, src.axes[1].name, output.axes[1].name)) return null;618    if (!std.mem.eql(u8, reduction.name, "blur_tap")) return null;619    if (reduction.operator != .weighted_sum) return null;620    if (!std.mem.eql(u8, reduction.shape.axes[0].name, weights.axes[0].name)) return null;621622    const radius_value = specialization.staticParameterValue("radius") orelse return null;623    const axis_value = specialization.staticParameterValue("axis") orelse return null;624    const radius = std.math.cast(u32, radius_value) orelse return null;625    const axis: Axis = switch (axis_value) {626        0 => .horizontal,627        1 => .vertical,628        else => return null,629    };630    const taps = 2 * @as(u64, radius) + 1;631    if (weights.axes[0].extent != taps or reduction.shape.axes[0].extent != taps) return null;632    if (!src.matchesExtents(&.{ output.axes[0].extent, output.axes[1].extent })) return null;633634    const launch = specialization.launch orelse return null;635    const instance = BlurPass{636        .radius = radius,637        .axis = axis,638        .width = output.axes[1].extent,639        .height = output.axes[0].extent,640        .threads = .{ .x = launch.threadgroup[0], .y = launch.threadgroup[1] },641    };642    if (!blurPassInstanceValid(instance)) return null;643    return instance;644}645646pub fn resizeInstanceFromSpecialization(specialization: entry.Specialization) ?Resize {647    if (!specialization.scheduleMatchesLaunch()) return null;648    if (!specialization.operationIs(.{ .image = .resize_bilinear })) return null;649    if (specialization.dtype != .u32 or specialization.accumulation_dtype != .f32) return null;650    if (specialization.inputs.len != 1 or specialization.outputs.len != 1) return null;651    if (specialization.reductions.len != 0 or specialization.static_parameters.len != 0) return null;652653    const src = specialization.inputs[0];654    const output = specialization.outputs[0];655    if (src.axes.len != 2 or output.axes.len != 2) return null;656657    const launch = specialization.launch orelse return null;658    const instance = Resize{659        .dst_width = output.axes[1].extent,660        .dst_height = output.axes[0].extent,661        .src_width = src.axes[1].extent,662        .src_height = src.axes[0].extent,663        .threads = .{ .x = launch.threadgroup[0], .y = launch.threadgroup[1] },664    };665    if (!resizeInstanceValid(instance)) return null;666    return instance;667}668669pub fn createBlurPassFamilyArtifact(670    allocator: std.mem.Allocator,671    handle: kernel.BackendHandle,672    instance: BlurPass,673    options: entry.ArtifactOptions,674) !kernel.OwnedKernelCallArtifact {675    if (!blurPassInstanceValid(instance)) return error.UnsupportedBlurPassInstance;676    const target = try blurPassFamilyTarget(allocator, instance);677    defer allocator.free(target);678    const entry_name = try blurPassFamilyEntryName(allocator, instance);679    defer allocator.free(entry_name);680    const family_fingerprint = options.shape_family_fingerprint orelse try blurPassFamilyFingerprint(allocator, instance);681    const shape_profile_dimensions = blurPassShapeProfileDimensions(instance);682    const shape_profile = options.shape_profile orelse artifact_product.KernelCallShapeProfile{683        .name = "image_blur_pass",684        .fingerprint = family_fingerprint,685        .dimensions = shape_profile_dimensions[0..],686    };687688    var graph = try BlurPassFamilyRgba8.buildNamed(allocator, options.limits, entry_name, instance);689    defer graph.deinit();690    return kernel.createKernelCallArtifact(allocator, handle, &graph, .{691        .target = target,692        .version = blur_family_version,693        .format = options.format,694        .kernel_plan = options.kernel_plan,695        .element_count_argument = options.element_count_argument,696        .shape_family_fingerprint = family_fingerprint,697        .shape_profile = shape_profile,698        .launch = options.launch orelse try imageDerivedLaunch(instance.threads),699        .runtime_scalar_argument_count = if (options.runtime_scalar_argument_count == 0) 2 else options.runtime_scalar_argument_count,700        .static_arguments = options.static_arguments,701    });702}703704pub fn createResizeFamilyArtifact(705    allocator: std.mem.Allocator,706    handle: kernel.BackendHandle,707    instance: Resize,708    options: entry.ArtifactOptions,709) !kernel.OwnedKernelCallArtifact {710    if (!resizeInstanceValid(instance)) return error.UnsupportedResizeInstance;711    const target = try resizeFamilyTarget(allocator, instance);712    defer allocator.free(target);713    const entry_name = try resizeFamilyEntryName(allocator, instance);714    defer allocator.free(entry_name);715    const family_fingerprint = options.shape_family_fingerprint orelse try resizeFamilyFingerprint(allocator, instance);716    const shape_profile_dimensions = resizeShapeProfileDimensions(instance);717    const shape_profile = options.shape_profile orelse artifact_product.KernelCallShapeProfile{718        .name = "image_resize_bilinear",719        .fingerprint = family_fingerprint,720        .dimensions = shape_profile_dimensions[0..],721    };722723    var graph = try ResizeBilinearFamilyRgba8.buildNamed(allocator, options.limits, entry_name, instance);724    defer graph.deinit();725    return kernel.createKernelCallArtifact(allocator, handle, &graph, .{726        .target = target,727        .version = resize_family_version,728        .format = options.format,729        .kernel_plan = options.kernel_plan,730        .element_count_argument = options.element_count_argument,731        .shape_family_fingerprint = family_fingerprint,732        .shape_profile = shape_profile,733        .launch = options.launch orelse try imageDerivedLaunch(instance.threads),734        .runtime_scalar_argument_count = if (options.runtime_scalar_argument_count == 0) 7 else options.runtime_scalar_argument_count,735        .static_arguments = options.static_arguments,736    });737}738739pub fn referenceBlurPass(dst: []u32, src: []const u32, weights: []const f32, width: u32, height: u32, axis: Axis) void {740    const radius: i64 = @intCast((weights.len - 1) / 2);741    var row: u32 = 0;742    while (row < height) : (row += 1) {743        var col: u32 = 0;744        while (col < width) : (col += 1) {745            var acc = [4]f32{ 0, 0, 0, 0 };746            for (weights, 0..) |weight, tap| {747                const offset = @as(i64, @intCast(tap)) - radius;748                const base: i64 = switch (axis) {749                    .horizontal => @as(i64, col),750                    .vertical => @as(i64, row),751                };752                const extent: i64 = switch (axis) {753                    .horizontal => @as(i64, width),754                    .vertical => @as(i64, height),755                };756                const sample = std.math.clamp(base + offset, 0, extent - 1);757                const index: usize = switch (axis) {758                    .horizontal => @as(usize, row) * width + @as(usize, @intCast(sample)),759                    .vertical => @as(usize, @intCast(sample)) * width + col,760                };761                const pixel = src[index];762                acc[0] = @mulAdd(f32, @floatFromInt(pixel & 0xff), weight, acc[0]);763                acc[1] = @mulAdd(f32, @floatFromInt((pixel >> 8) & 0xff), weight, acc[1]);764                acc[2] = @mulAdd(f32, @floatFromInt((pixel >> 16) & 0xff), weight, acc[2]);765                acc[3] = @mulAdd(f32, @floatFromInt(pixel >> 24), weight, acc[3]);766            }767            dst[@as(usize, row) * width + col] = packReferenceChannels(acc);768        }769    }770}771772pub fn referenceResizeBilinear(dst: []u32, src: []const u32, dst_width: u32, dst_height: u32, src_width: u32, src_height: u32) void {773    var row: u32 = 0;774    while (row < dst_height) : (row += 1) {775        var col: u32 = 0;776        while (col < dst_width) : (col += 1) {777            const sx = referenceSampleAxis(col, dst_width, src_width);778            const sy = referenceSampleAxis(row, dst_height, src_height);779            const p00 = referenceUnpack(src[sy.lo * src_width + sx.lo]);780            const p01 = referenceUnpack(src[sy.lo * src_width + sx.hi]);781            const p10 = referenceUnpack(src[sy.hi * src_width + sx.lo]);782            const p11 = referenceUnpack(src[sy.hi * src_width + sx.hi]);783            var blended: [4]f32 = undefined;784            for (0..4) |channel| {785                const top = @mulAdd(f32, p01[channel] - p00[channel], sx.frac, p00[channel]);786                const bottom = @mulAdd(f32, p11[channel] - p10[channel], sx.frac, p10[channel]);787                blended[channel] = @mulAdd(f32, bottom - top, sy.frac, top);788            }789            dst[@as(usize, row) * dst_width + col] = packReferenceChannels(blended);790        }791    }792}793794const ReferenceAxisSample = struct {795    lo: usize,796    hi: usize,797    frac: f32,798};799800fn referenceSampleAxis(dst_index: u32, dst_extent: u32, src_extent: u32) ReferenceAxisSample {801    const scale = resizeScale(src_extent, dst_extent);802    const centered = @mulAdd(f32, @as(f32, @floatFromInt(dst_index)) + 0.5, scale, -0.5);803    const last = @as(f32, @floatFromInt(src_extent)) - 1;804    const clamped = @max(@as(f32, 0), @min(centered, last));805    const lo_f = @floor(clamped);806    const hi_f = @min(lo_f + 1, last);807    return .{808        .lo = @intFromFloat(lo_f),809        .hi = @intFromFloat(hi_f),810        .frac = clamped - lo_f,811    };812}813814fn referenceUnpack(pixel: u32) [4]f32 {815    return .{816        @floatFromInt(pixel & 0xff),817        @floatFromInt((pixel >> 8) & 0xff),818        @floatFromInt((pixel >> 16) & 0xff),819        @floatFromInt(pixel >> 24),820    };821}822823fn packReferenceChannels(channels: [4]f32) u32 {824    var packed_pixel: u32 = 0;825    for (channels, 0..) |value, channel| {826        const rounded = @floor(value + 0.5);827        const clamped = @max(@as(f32, 0), @min(rounded, 255));828        packed_pixel |= @as(u32, @intFromFloat(clamped)) << @intCast(channel * 8);829    }830    return packed_pixel;831}832833const testing = std.testing;834835fn testPixel(seed: usize) u32 {836    var value: u32 = @truncate(seed *% 2654435761);837    value ^= value >> 13;838    value *%= 0x5bd1e995;839    value ^= value >> 15;840    return value;841}842843test "gaussian weights normalize and peak at the center" {844    const weights = try gaussianWeightsAlloc(testing.allocator, 3, 1.4);845    defer testing.allocator.free(weights);846    try testing.expectEqual(@as(usize, 7), weights.len);847    var total: f32 = 0;848    for (weights) |weight| total += weight;849    try testing.expectApproxEqAbs(@as(f32, 1), total, 0.0001);850    for (weights) |weight| try testing.expect(weight <= weights[3]);851}852853test "blur pass family matches the reference on the interpreter" {854    const allocator = testing.allocator;855    const width: u32 = 13;856    const height: u32 = 7;857    const pixel_count = @as(usize, width) * height;858859    const src = try allocator.alloc(u32, pixel_count);860    defer allocator.free(src);861    for (src, 0..) |*pixel, index| pixel.* = testPixel(index);862863    const weights = try gaussianWeightsAlloc(allocator, 2, 1.1);864    defer allocator.free(weights);865866    inline for (.{ Axis.horizontal, Axis.vertical }) |axis| {867        const instance = BlurPass{ .radius = 2, .axis = axis, .threads = .{ .x = 8, .y = 4 } };868        const expected = try allocator.alloc(u32, pixel_count);869        defer allocator.free(expected);870        referenceBlurPass(expected, src, weights, width, height, axis);871872        const actual = try allocator.alloc(u32, pixel_count);873        defer allocator.free(actual);874        @memset(actual, 0);875876        var graph = try BlurPassFamilyRgba8.build(allocator, BlurPassFamilyRgba8.Limits.testing, instance);877        defer graph.deinit();878        const geometry = imageLaunchGeometry(instance.threads, width, height);879        try graph.runCpuWithLaunch(allocator, &.{880            kernel.argumentBuffer(u32, actual),881            kernel.argumentBuffer(u32, @constCast(src)),882            kernel.argumentBuffer(f32, @constCast(weights)),883            kernel.argumentI32(@intCast(width)),884            kernel.argumentI32(@intCast(height)),885        }, .{886            .grid = geometry.grid,887            .block = geometry.threadgroup,888        });889890        try testing.expectEqualSlices(u32, expected, actual);891    }892}893894test "resize bilinear family matches the reference on the interpreter" {895    const allocator = testing.allocator;896    const src_width: u32 = 12;897    const src_height: u32 = 9;898    const dst_width: u32 = 7;899    const dst_height: u32 = 5;900901    const src = try allocator.alloc(u32, @as(usize, src_width) * src_height);902    defer allocator.free(src);903    for (src, 0..) |*pixel, index| pixel.* = testPixel(index +% 17);904905    const expected = try allocator.alloc(u32, @as(usize, dst_width) * dst_height);906    defer allocator.free(expected);907    referenceResizeBilinear(expected, src, dst_width, dst_height, src_width, src_height);908909    const actual = try allocator.alloc(u32, @as(usize, dst_width) * dst_height);910    defer allocator.free(actual);911    @memset(actual, 0);912913    const instance = Resize{ .threads = .{ .x = 8, .y = 4 } };914    var graph = try ResizeBilinearFamilyRgba8.build(allocator, ResizeBilinearFamilyRgba8.Limits.testing, instance);915    defer graph.deinit();916    const geometry = imageLaunchGeometry(instance.threads, dst_width, dst_height);917    try graph.runCpuWithLaunch(allocator, &.{918        kernel.argumentBuffer(u32, actual),919        kernel.argumentBuffer(u32, @constCast(src)),920        kernel.argumentI32(@intCast(dst_width)),921        kernel.argumentI32(@intCast(dst_height)),922        kernel.argumentI32(@intCast(src_width)),923        kernel.argumentF32(resizeScale(src_width, dst_width)),924        kernel.argumentF32(resizeScale(src_height, dst_height)),925        kernel.argumentF32(@floatFromInt(src_width - 1)),926        kernel.argumentF32(@floatFromInt(src_height - 1)),927    }, .{928        .grid = geometry.grid,929        .block = geometry.threadgroup,930    });931932    try testing.expectEqualSlices(u32, expected, actual);933}934935test "upscale resize keeps corner pixels exact" {936    const allocator = testing.allocator;937    const src = [_]u32{ 0xff000011, 0xff000022, 0xff000033, 0xff000044 };938    const dst = try allocator.alloc(u32, 16);939    defer allocator.free(dst);940    referenceResizeBilinear(dst, src[0..], 4, 4, 2, 2);941    try testing.expectEqual(src[0], dst[0]);942    try testing.expectEqual(src[1], dst[3]);943    try testing.expectEqual(src[2], dst[12]);944    try testing.expectEqual(src[3], dst[15]);945}946947test "blur instance validation bounds radius and threads" {948    try testing.expect(blurPassInstanceValid(.{ .radius = 1, .axis = .horizontal }));949    try testing.expect(!blurPassInstanceValid(.{ .radius = 0, .axis = .horizontal }));950    try testing.expect(!blurPassInstanceValid(.{ .radius = blur_radius_max + 1, .axis = .vertical }));951    try testing.expect(!blurPassInstanceValid(.{ .radius = 2, .axis = .vertical, .threads = .{ .x = 0, .y = 4 } }));952}

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