tiny.accy.kernel.library.fused

1const std = @import("std");2const gpu = @import("gpu");34const activation_mod = @import("../../choir/root.zig").activation;5const elementwise = @import("elementwise.zig");6const entry = @import("entry.zig");7const kernel = @import("../root.zig");8const linalg = @import("linalg.zig");910fn activationName(comptime activation: activation_mod.Kind) []const u8 {11    return switch (activation) {12        .gelu => "gelu",13        .relu => "relu",14        .silu => "silu",15    };16}1718fn activationValue(inner: anytype, comptime activation: activation_mod.Kind, value: anytype) !@TypeOf(value) {19    return switch (activation) {20        .gelu => elementwise.geluTanhApprox(inner, value),21        .relu => elementwise.relu(inner, value),22        .silu => elementwise.silu(inner, value),23    };24}2526fn biasActivationSpecialization(comptime spec: entry.Vector1D, comptime activation: activation_mod.Kind) entry.Specialization {27    return .{28        .dtype = .f32,29        .operation = .{ .elementwise = .add },30        .inputs = &.{31            entry.shape1D(spec.axis, spec.extent),32            entry.shape1D(spec.axis, spec.extent),33        },34        .outputs = &.{entry.shape1D(spec.axis, spec.extent)},35        .epilogues = &.{entry.epilogue(.{ .activation = activation })},36        .launch = entry.launch1D(spec.extent, spec.threads),37        .schedule = entry.threadBlocks1D(spec.axis, spec.extent, spec.threads),38    };39}4041fn bias_activation_each(inner: anytype, index: kernel.Index1D, ctx: anytype) !void {42    const value = try ctx.args.param(.src).load(inner, index);43    const bias = try ctx.args.param(.bias).load(inner, index);44    const shifted = try value.add(inner, bias);45    const activated = try activationValue(inner, ctx.activation, shifted);46    try ctx.args.param(.dst).store(inner, activated, index);47}4849fn biasActivationProgram(comptime spec: entry.Vector1D, comptime activation: activation_mod.Kind) type {50    const Body = struct {51        fn run(k: anytype, args: anytype) !void {52            _ = try k.forEach1D(spec.axis, spec.extent, .{ .args = args, .activation = activation }, bias_activation_each);53        }54    };5556    return kernel.logical.Program(.{57        .name = std.fmt.comptimePrint("accy_kernel_fused_bias_{s}{}x{}_f32", .{ activationName(activation), spec.extent, spec.threads }),58        .parameters = .{59            .dst = kernel.dynamicBuffer(.f32),60            .src = kernel.dynamicBuffer(.f32),61            .bias = kernel.dynamicBuffer(.f32),62        },63        .body = Body.run,64    }).withSchedule(kernel.logical.schedule.threadBlocks(.{ .x = spec.threads }));65}6667fn biasActivationF32(comptime spec: entry.Vector1D, comptime activation: activation_mod.Kind) type {68    return entry.Entry(biasActivationProgram(spec, activation), .{69        .target = std.fmt.comptimePrint("accy.kernel.fused.bias_{s}{}x{}_f32", .{ activationName(activation), spec.extent, spec.threads }),70        .layer = .logical,71        .category = .fused,72        .specialization = biasActivationSpecialization(spec, activation),73    });74}7576pub fn biasGeluF32(comptime spec: entry.Vector1D) type {77    return biasActivationF32(spec, .gelu);78}7980pub fn biasReluF32(comptime spec: entry.Vector1D) type {81    return biasActivationF32(spec, .relu);82}8384pub fn biasSiluF32(comptime spec: entry.Vector1D) type {85    return biasActivationF32(spec, .silu);86}8788pub const BiasGelu8F32 = biasGeluF32(.{ .extent = 8, .threads = 4 });89pub const BiasRelu8F32 = biasReluF32(.{ .extent = 8, .threads = 4 });90pub const BiasSilu8F32 = biasSiluF32(.{ .extent = 8, .threads = 4 });9192const GatedActivation = activation_mod.Kind;9394fn gatedActivationName(comptime activation: GatedActivation) []const u8 {95    return switch (activation) {96        .gelu => "geglu",97        .relu => "reglu",98        .silu => "swiglu",99    };100}101102fn gatedActivationValue(inner: anytype, comptime activation: GatedActivation, value: anytype) !@TypeOf(value) {103    return activationValue(inner, activation, value);104}105106fn gatedActivationSpecialization(comptime spec: entry.Vector1D, comptime activation: GatedActivation) entry.Specialization {107    return .{108        .dtype = .f32,109        .operation = .{ .elementwise = .mul },110        .inputs = &.{111            entry.shape1D(spec.axis, spec.extent),112            entry.shape1D(spec.axis, spec.extent),113        },114        .outputs = &.{entry.shape1D(spec.axis, spec.extent)},115        .input_transforms = &.{entry.inputTransform(.{ .activation = activation }, 0, entry.shape1D(spec.axis, spec.extent))},116        .launch = entry.launch1D(spec.extent, spec.threads),117        .schedule = entry.threadBlocks1D(spec.axis, spec.extent, spec.threads),118    };119}120121fn gated_activation_each(inner: anytype, index: kernel.Index1D, ctx: anytype) !void {122    const gate = try ctx.args.param(.gate).load(inner, index);123    const value = try ctx.args.param(.value).load(inner, index);124    const activated = try gatedActivationValue(inner, ctx.activation, gate);125    const output = try activated.mul(inner, value);126    try ctx.args.param(.dst).store(inner, output, index);127}128129fn gatedActivationProgram(comptime spec: entry.Vector1D, comptime activation: GatedActivation) type {130    const Body = struct {131        fn run(k: anytype, args: anytype) !void {132            _ = try k.forEach1D(spec.axis, spec.extent, .{ .args = args, .activation = activation }, gated_activation_each);133        }134    };135136    return kernel.logical.Program(.{137        .name = std.fmt.comptimePrint("accy_kernel_fused_{s}{}x{}_f32", .{ gatedActivationName(activation), spec.extent, spec.threads }),138        .parameters = .{139            .dst = kernel.dynamicBuffer(.f32),140            .gate = kernel.dynamicBuffer(.f32),141            .value = kernel.dynamicBuffer(.f32),142        },143        .body = Body.run,144    }).withSchedule(kernel.logical.schedule.threadBlocks(.{ .x = spec.threads }));145}146147fn gatedActivationF32(comptime spec: entry.Vector1D, comptime activation: GatedActivation) type {148    return entry.Entry(gatedActivationProgram(spec, activation), .{149        .target = std.fmt.comptimePrint("accy.kernel.fused.{s}{}x{}_f32", .{ gatedActivationName(activation), spec.extent, spec.threads }),150        .layer = .logical,151        .category = .fused,152        .specialization = gatedActivationSpecialization(spec, activation),153    });154}155156pub fn geGluF32(comptime spec: entry.Vector1D) type {157    return gatedActivationF32(spec, .gelu);158}159160pub fn reGluF32(comptime spec: entry.Vector1D) type {161    return gatedActivationF32(spec, .relu);162}163164pub fn swiGluF32(comptime spec: entry.Vector1D) type {165    return gatedActivationF32(spec, .silu);166}167168pub const GeGlu8F32 = geGluF32(.{ .extent = 8, .threads = 4 });169pub const ReGlu8F32 = reGluF32(.{ .extent = 8, .threads = 4 });170pub const SwiGlu8F32 = swiGluF32(.{ .extent = 8, .threads = 4 });171172fn matrixProductBiasActivationSpecialization(comptime spec: linalg.MatrixProduct, comptime activation: activation_mod.Kind) entry.Specialization {173    return .{174        .dtype = .f32,175        .operation = .{ .linalg = .matrix_product },176        .equation = "mk,kn,n->mn",177        .inputs = &.{178            entry.shape2D(spec.row_axis, spec.m, spec.reduction_axis, spec.k),179            entry.shape2D(spec.reduction_axis, spec.k, spec.col_axis, spec.n),180            entry.shape1D(spec.col_axis, spec.n),181        },182        .outputs = &.{entry.shape2D(spec.row_axis, spec.m, spec.col_axis, spec.n)},183        .reductions = &.{entry.reduction("dot", .dot_product, entry.shape1D(spec.reduction_axis, spec.k))},184        .epilogues = &.{185            entry.inputEpilogue(.bias_add, 2, entry.shape1D(spec.col_axis, spec.n)),186            entry.epilogue(.{ .activation = activation }),187        },188        .launch = entry.launch2D(spec.n, spec.m, spec.threads.x, spec.threads.y),189        .schedule = entry.threadBlocks2D(spec.col_axis, spec.n, spec.row_axis, spec.m, spec.threads.x, spec.threads.y),190    };191}192193fn matrix_product_bias_activation_each(inner: anytype, index: kernel.Index2D, ctx: anytype) !void {194    const sum = try linalg.matrixProductCellSum(inner, ctx.spec, ctx.args.param(.lhs), ctx.args.param(.rhs), index.y.index, index.x.index);195    const bias = try ctx.args.param(.bias).load(inner, index.x.index);196    const shifted = try bias.add(inner, sum);197    const activated = try activationValue(inner, ctx.activation, shifted);198    const out_index = try linalg.matrixProductOutputIndex(inner, ctx.spec, index.y.index, index.x.index);199    try ctx.args.param(.dst).store(inner, activated, out_index);200}201202fn matrixProductBiasActivationProgram(comptime spec: linalg.MatrixProduct, comptime activation: activation_mod.Kind) type {203    const Body = struct {204        fn run(k: anytype, args: anytype) !void {205            _ = try k.forEach2D(.{206                .x = kernel.logical.axis(spec.col_axis, spec.n),207                .y = kernel.logical.axis(spec.row_axis, spec.m),208            }, .{ .spec = spec, .activation = activation, .args = args }, matrix_product_bias_activation_each);209        }210    };211212    return kernel.logical.Program(.{213        .name = std.fmt.comptimePrint(214            "accy_kernel_fused_matmul_bias_{s}{}x{}x{}_{}x{}_f32",215            .{ activationName(activation), spec.m, spec.n, spec.k, spec.threads.x, spec.threads.y },216        ),217        .parameters = .{218            .dst = kernel.dynamicBuffer(.f32),219            .lhs = kernel.dynamicBuffer(.f32),220            .rhs = kernel.dynamicBuffer(.f32),221            .bias = kernel.dynamicBuffer(.f32),222        },223        .body = Body.run,224    }).withSchedule(kernel.logical.schedule.threadBlocks(.{225        .x = spec.threads.x,226        .y = spec.threads.y,227    }));228}229230fn matrixProductBiasActivationF32(comptime spec: linalg.MatrixProduct, comptime activation: activation_mod.Kind) type {231    return entry.Entry(matrixProductBiasActivationProgram(spec, activation), .{232        .target = std.fmt.comptimePrint(233            "accy.kernel.fused.matmul_bias_{s}{}x{}x{}_{}x{}_f32",234            .{ activationName(activation), spec.m, spec.n, spec.k, spec.threads.x, spec.threads.y },235        ),236        .layer = .logical,237        .category = .fused,238        .specialization = matrixProductBiasActivationSpecialization(spec, activation),239    });240}241242pub fn matrixProductBiasGeluF32(comptime spec: linalg.MatrixProduct) type {243    return matrixProductBiasActivationF32(spec, .gelu);244}245246pub fn matrixProductBiasReluF32(comptime spec: linalg.MatrixProduct) type {247    return matrixProductBiasActivationF32(spec, .relu);248}249250pub fn matrixProductBiasSiluF32(comptime spec: linalg.MatrixProduct) type {251    return matrixProductBiasActivationF32(spec, .silu);252}253254pub const MatrixProductBiasGelu2x3x4F32 = matrixProductBiasGeluF32(.{255    .m = 2,256    .n = 3,257    .k = 4,258    .threads = .{ .x = 2, .y = 2 },259});260pub const MatrixProductBiasGelu2x3x4ThreadBlocks1x2F32 = matrixProductBiasGeluF32(.{261    .m = 2,262    .n = 3,263    .k = 4,264    .threads = .{ .x = 1, .y = 2 },265});266pub const MatrixProductBiasRelu2x3x4F32 = matrixProductBiasReluF32(.{267    .m = 2,268    .n = 3,269    .k = 4,270    .threads = .{ .x = 2, .y = 2 },271});272pub const MatrixProductBiasSilu2x3x4F32 = matrixProductBiasSiluF32(.{273    .m = 2,274    .n = 3,275    .k = 4,276    .threads = .{ .x = 2, .y = 2 },277});278279fn matrixVectorProductBiasActivationSpecialization(comptime spec: linalg.MatrixVectorProduct, comptime activation: activation_mod.Kind) entry.Specialization {280    return .{281        .dtype = .f32,282        .operation = .{ .linalg = .matrix_vector_product },283        .equation = "mk,k,m->m",284        .inputs = &.{285            entry.shape2D(spec.row_axis, spec.m, spec.reduction_axis, spec.k),286            entry.shape1D(spec.reduction_axis, spec.k),287            entry.shape1D(spec.row_axis, spec.m),288        },289        .outputs = &.{entry.shape1D(spec.row_axis, spec.m)},290        .reductions = &.{entry.reduction("dot", .dot_product, entry.shape1D(spec.reduction_axis, spec.k))},291        .epilogues = &.{292            entry.inputEpilogue(.bias_add, 2, entry.shape1D(spec.row_axis, spec.m)),293            entry.epilogue(.{ .activation = activation }),294        },295        .launch = entry.launch1D(spec.m, spec.threads),296        .schedule = entry.threadBlocks1D(spec.row_axis, spec.m, spec.threads),297    };298}299300fn matrix_vector_product_bias_activation_each(inner: anytype, index: kernel.Index1D, ctx: anytype) !void {301    const sum = try linalg.matrixVectorProductRowSum(inner, ctx.spec, ctx.args.param(.matrix), ctx.args.param(.vector), index.index);302    const bias = try ctx.args.param(.bias).load(inner, index);303    const shifted = try bias.add(inner, sum);304    const activated = try activationValue(inner, ctx.activation, shifted);305    try ctx.args.param(.dst).store(inner, activated, index);306}307308fn matrixVectorProductBiasActivationProgram(comptime spec: linalg.MatrixVectorProduct, comptime activation: activation_mod.Kind) type {309    const Body = struct {310        fn run(k: anytype, args: anytype) !void {311            _ = try k.forEach1D(spec.row_axis, spec.m, .{ .spec = spec, .activation = activation, .args = args }, matrix_vector_product_bias_activation_each);312        }313    };314315    return kernel.logical.Program(.{316        .name = std.fmt.comptimePrint(317            "accy_kernel_fused_matvec_bias_{s}{}x{}_{}x_f32",318            .{ activationName(activation), spec.m, spec.k, spec.threads },319        ),320        .parameters = .{321            .dst = kernel.dynamicBuffer(.f32),322            .matrix = kernel.dynamicBuffer(.f32),323            .vector = kernel.dynamicBuffer(.f32),324            .bias = kernel.dynamicBuffer(.f32),325        },326        .body = Body.run,327    }).withSchedule(kernel.logical.schedule.threadBlocks(.{ .x = spec.threads }));328}329330fn matrixVectorProductBiasActivationF32(comptime spec: linalg.MatrixVectorProduct, comptime activation: activation_mod.Kind) type {331    return entry.Entry(matrixVectorProductBiasActivationProgram(spec, activation), .{332        .target = std.fmt.comptimePrint(333            "accy.kernel.fused.matvec_bias_{s}{}x{}_{}x_f32",334            .{ activationName(activation), spec.m, spec.k, spec.threads },335        ),336        .layer = .logical,337        .category = .fused,338        .specialization = matrixVectorProductBiasActivationSpecialization(spec, activation),339    });340}341342pub fn matrixVectorProductBiasGeluF32(comptime spec: linalg.MatrixVectorProduct) type {343    return matrixVectorProductBiasActivationF32(spec, .gelu);344}345346pub fn matrixVectorProductBiasReluF32(comptime spec: linalg.MatrixVectorProduct) type {347    return matrixVectorProductBiasActivationF32(spec, .relu);348}349350pub fn matrixVectorProductBiasSiluF32(comptime spec: linalg.MatrixVectorProduct) type {351    return matrixVectorProductBiasActivationF32(spec, .silu);352}353354pub const MatrixVectorProductBiasGelu4x8F32 = matrixVectorProductBiasGeluF32(.{355    .m = 4,356    .k = 8,357    .threads = 4,358});359pub const MatrixVectorProductBiasRelu4x8F32 = matrixVectorProductBiasReluF32(.{360    .m = 4,361    .k = 8,362    .threads = 4,363});364pub const MatrixVectorProductBiasSilu4x8F32 = matrixVectorProductBiasSiluF32(.{365    .m = 4,366    .k = 8,367    .threads = 4,368});369370fn expectedActivation(comptime activation: activation_mod.Kind, value: f32) f32 {371    return switch (activation) {372        .gelu => 0.5 * value * (1.0 + std.math.tanh(0.7978845608028654 * (value + 0.044715 * value * value * value))),373        .relu => if (value > 0.0) value else 0.0,374        .silu => value / (1.0 + @exp(-value)),375    };376}377378fn expectedBiasActivation(comptime activation: activation_mod.Kind, value: f32, bias: f32) f32 {379    return expectedActivation(activation, value + bias);380}381382fn expectedSwiGlu(gate: f32, value: f32) f32 {383    return gate / (1.0 + @exp(-gate)) * value;384}385386fn expectedGeGlu(gate: f32, value: f32) f32 {387    return expectedActivation(.gelu, gate) * value;388}389390fn expectedReGlu(gate: f32, value: f32) f32 {391    if (gate > 0.0) return gate * value;392    return 0.0;393}394395test "fused bias gelu entry runs on CPU and records fused operation" {396    var src = [_]f32{ -3.0, -1.0, -0.5, 0.0, 0.5, 1.0, 2.0, 4.0 };397    var bias = [_]f32{ 0.25, -0.5, 0.5, 1.0, -0.25, 0.75, -1.0, 0.0 };398    var dst = @as([8]f32, @splat(0.0));399400    try BiasGelu8F32.runCpu(std.testing.allocator, BiasGelu8F32.Limits.testing, &.{401        kernel.argumentBuffer(f32, dst[0..]),402        kernel.argumentBuffer(f32, src[0..]),403        kernel.argumentBuffer(f32, bias[0..]),404    });405    for (src, bias, dst) |input, bias_value, actual| {406        try std.testing.expectApproxEqAbs(expectedBiasActivation(.gelu, input, bias_value), actual, 0.0001);407    }408409    const launch_value = try BiasGelu8F32.launch(std.testing.allocator, BiasGelu8F32.Limits.testing);410    try std.testing.expectEqual(@as(u32, 2), launch_value.grid[0]);411    try std.testing.expectEqual(@as(u32, 4), launch_value.block[0]);412    try std.testing.expectEqual(entry.Category.fused, BiasGelu8F32.category);413    try std.testing.expect(BiasGelu8F32.specialization.operationIs(.{ .elementwise = .add }));414    try std.testing.expect(BiasGelu8F32.specialization.epilogueMatches(0, .{ .operator = .{ .activation = .gelu } }));415    try std.testing.expectEqual(@as(usize, 2), BiasGelu8F32.specialization.inputs.len);416    try std.testing.expectEqual(@as(u64, 8), BiasGelu8F32.specialization.outputs[0].elementCount().?);417    try std.testing.expectEqualDeep(BiasGelu8F32.specialization.launch.?, BiasGelu8F32.specialization.schedule.?.launch());418419    var snapshot = try BiasGelu8F32.scheduleSnapshot(std.testing.allocator, BiasGelu8F32.Limits.testing);420    defer snapshot.deinit(std.testing.allocator);421    try std.testing.expect(BiasGelu8F32.specialization.schedule.?.matchesSnapshot(&snapshot));422}423424test "fused bias relu and silu entries run on CPU" {425    var src = [_]f32{ -3.0, -1.0, -0.5, 0.0, 0.5, 1.0, 2.0, 4.0 };426    var bias = [_]f32{ 0.25, -0.5, 0.5, 1.0, -0.25, 0.75, -1.0, 0.0 };427    var relu_dst = @as([8]f32, @splat(0.0));428    var silu_dst = @as([8]f32, @splat(0.0));429430    try BiasRelu8F32.runCpu(std.testing.allocator, BiasRelu8F32.Limits.testing, &.{431        kernel.argumentBuffer(f32, relu_dst[0..]),432        kernel.argumentBuffer(f32, src[0..]),433        kernel.argumentBuffer(f32, bias[0..]),434    });435    try BiasSilu8F32.runCpu(std.testing.allocator, BiasSilu8F32.Limits.testing, &.{436        kernel.argumentBuffer(f32, silu_dst[0..]),437        kernel.argumentBuffer(f32, src[0..]),438        kernel.argumentBuffer(f32, bias[0..]),439    });440    for (src, bias, relu_dst, silu_dst) |input, bias_value, relu_actual, silu_actual| {441        try std.testing.expectApproxEqAbs(expectedBiasActivation(.relu, input, bias_value), relu_actual, 0.0001);442        try std.testing.expectApproxEqAbs(expectedBiasActivation(.silu, input, bias_value), silu_actual, 0.0001);443    }444445    try std.testing.expect(BiasRelu8F32.specialization.epilogueMatches(0, .{ .operator = .{ .activation = .relu } }));446    try std.testing.expect(BiasSilu8F32.specialization.epilogueMatches(0, .{ .operator = .{ .activation = .silu } }));447}448449test "fused swiglu entry runs on CPU and records input transform" {450    var gate = [_]f32{ -5.0, -2.0, -1.0, 0.0, 0.5, 1.0, 3.0, 6.0 };451    var value = [_]f32{ 2.0, -1.0, 0.5, 3.0, -2.0, 4.0, 0.25, -0.5 };452    var dst = @as([8]f32, @splat(0.0));453454    try SwiGlu8F32.runCpu(std.testing.allocator, SwiGlu8F32.Limits.testing, &.{455        kernel.argumentBuffer(f32, dst[0..]),456        kernel.argumentBuffer(f32, gate[0..]),457        kernel.argumentBuffer(f32, value[0..]),458    });459    for (gate, value, dst) |gate_value, input_value, actual| {460        try std.testing.expectApproxEqAbs(expectedSwiGlu(gate_value, input_value), actual, 0.0001);461    }462463    const launch_value = try SwiGlu8F32.launch(std.testing.allocator, SwiGlu8F32.Limits.testing);464    try std.testing.expectEqual(@as(u32, 2), launch_value.grid[0]);465    try std.testing.expectEqual(@as(u32, 4), launch_value.block[0]);466    try std.testing.expectEqual(entry.Category.fused, SwiGlu8F32.category);467    try std.testing.expect(SwiGlu8F32.specialization.operationIs(.{ .elementwise = .mul }));468    try std.testing.expect(SwiGlu8F32.specialization.inputTransformMatches(0, .{469        .operator = .{ .activation = .silu },470        .input_index = 0,471        .extents = &.{8},472    }));473    try std.testing.expectEqual(@as(usize, 2), SwiGlu8F32.specialization.inputs.len);474    try std.testing.expectEqual(@as(u64, 8), SwiGlu8F32.specialization.outputs[0].elementCount().?);475    try std.testing.expectEqualDeep(SwiGlu8F32.specialization.launch.?, SwiGlu8F32.specialization.schedule.?.launch());476477    var snapshot = try SwiGlu8F32.scheduleSnapshot(std.testing.allocator, SwiGlu8F32.Limits.testing);478    defer snapshot.deinit(std.testing.allocator);479    try std.testing.expect(SwiGlu8F32.specialization.schedule.?.matchesSnapshot(&snapshot));480}481482test "fused geglu entry runs on CPU and records input transform" {483    var gate = [_]f32{ -5.0, -2.0, -1.0, 0.0, 0.5, 1.0, 3.0, 6.0 };484    var value = [_]f32{ 2.0, -1.0, 0.5, 3.0, -2.0, 4.0, 0.25, -0.5 };485    var dst = @as([8]f32, @splat(0.0));486487    try GeGlu8F32.runCpu(std.testing.allocator, GeGlu8F32.Limits.testing, &.{488        kernel.argumentBuffer(f32, dst[0..]),489        kernel.argumentBuffer(f32, gate[0..]),490        kernel.argumentBuffer(f32, value[0..]),491    });492    for (gate, value, dst) |gate_value, input_value, actual| {493        try std.testing.expectApproxEqAbs(expectedGeGlu(gate_value, input_value), actual, 0.0001);494    }495496    const launch_value = try GeGlu8F32.launch(std.testing.allocator, GeGlu8F32.Limits.testing);497    try std.testing.expectEqual(@as(u32, 2), launch_value.grid[0]);498    try std.testing.expectEqual(@as(u32, 4), launch_value.block[0]);499    try std.testing.expectEqual(entry.Category.fused, GeGlu8F32.category);500    try std.testing.expect(GeGlu8F32.specialization.operationIs(.{ .elementwise = .mul }));501    try std.testing.expect(GeGlu8F32.specialization.inputTransformMatches(0, .{502        .operator = .{ .activation = .gelu },503        .input_index = 0,504        .extents = &.{8},505    }));506    try std.testing.expectEqual(@as(usize, 2), GeGlu8F32.specialization.inputs.len);507    try std.testing.expectEqual(@as(u64, 8), GeGlu8F32.specialization.outputs[0].elementCount().?);508    try std.testing.expectEqualDeep(GeGlu8F32.specialization.launch.?, GeGlu8F32.specialization.schedule.?.launch());509510    var snapshot = try GeGlu8F32.scheduleSnapshot(std.testing.allocator, GeGlu8F32.Limits.testing);511    defer snapshot.deinit(std.testing.allocator);512    try std.testing.expect(GeGlu8F32.specialization.schedule.?.matchesSnapshot(&snapshot));513}514515test "fused reglu entry runs on CPU and records input transform" {516    var gate = [_]f32{ -5.0, -2.0, -1.0, 0.0, 0.5, 1.0, 3.0, 6.0 };517    var value = [_]f32{ 2.0, -1.0, 0.5, 3.0, -2.0, 4.0, 0.25, -0.5 };518    var dst = @as([8]f32, @splat(0.0));519520    try ReGlu8F32.runCpu(std.testing.allocator, ReGlu8F32.Limits.testing, &.{521        kernel.argumentBuffer(f32, dst[0..]),522        kernel.argumentBuffer(f32, gate[0..]),523        kernel.argumentBuffer(f32, value[0..]),524    });525    for (gate, value, dst) |gate_value, input_value, actual| {526        try std.testing.expectApproxEqAbs(expectedReGlu(gate_value, input_value), actual, 0.0001);527    }528529    const launch_value = try ReGlu8F32.launch(std.testing.allocator, ReGlu8F32.Limits.testing);530    try std.testing.expectEqual(@as(u32, 2), launch_value.grid[0]);531    try std.testing.expectEqual(@as(u32, 4), launch_value.block[0]);532    try std.testing.expectEqual(entry.Category.fused, ReGlu8F32.category);533    try std.testing.expect(ReGlu8F32.specialization.operationIs(.{ .elementwise = .mul }));534    try std.testing.expect(ReGlu8F32.specialization.inputTransformMatches(0, .{535        .operator = .{ .activation = .relu },536        .input_index = 0,537        .extents = &.{8},538    }));539    try std.testing.expectEqual(@as(usize, 2), ReGlu8F32.specialization.inputs.len);540    try std.testing.expectEqual(@as(u64, 8), ReGlu8F32.specialization.outputs[0].elementCount().?);541    try std.testing.expectEqualDeep(ReGlu8F32.specialization.launch.?, ReGlu8F32.specialization.schedule.?.launch());542543    var snapshot = try ReGlu8F32.scheduleSnapshot(std.testing.allocator, ReGlu8F32.Limits.testing);544    defer snapshot.deinit(std.testing.allocator);545    try std.testing.expect(ReGlu8F32.specialization.schedule.?.matchesSnapshot(&snapshot));546}547548test "fused matrix product bias gelu entry runs on CPU and records epilogue operation" {549    var lhs = [_]f32{550        1.0, 2.0, 3.0, 4.0,551        5.0, 6.0, 7.0, 8.0,552    };553    var rhs = [_]f32{554        1.0, 0.0, 2.0,555        0.0, 1.0, 3.0,556        1.0, 1.0, 0.0,557        2.0, 0.0, 1.0,558    };559    var bias = [_]f32{ 0.25, -1.0, 0.5 };560    var dst = @as([6]f32, @splat(0.0));561562    try MatrixProductBiasGelu2x3x4F32.runCpu(std.testing.allocator, MatrixProductBiasGelu2x3x4F32.Limits.testing, &.{563        kernel.argumentBuffer(f32, dst[0..]),564        kernel.argumentBuffer(f32, lhs[0..]),565        kernel.argumentBuffer(f32, rhs[0..]),566        kernel.argumentBuffer(f32, bias[0..]),567    });568    const dot = [_]f32{ 12.0, 5.0, 12.0, 28.0, 13.0, 36.0 };569    for (dot, dst, 0..) |sum, actual, index| {570        try std.testing.expectApproxEqAbs(expectedBiasActivation(.gelu, sum, bias[index % 3]), actual, 0.0001);571    }572573    const launch_value = try MatrixProductBiasGelu2x3x4F32.launch(std.testing.allocator, MatrixProductBiasGelu2x3x4F32.Limits.testing);574    try std.testing.expectEqual(@as(u32, 2), launch_value.grid[0]);575    try std.testing.expectEqual(@as(u32, 1), launch_value.grid[1]);576    try std.testing.expectEqual(@as(u32, 2), launch_value.block[0]);577    try std.testing.expectEqual(@as(u32, 2), launch_value.block[1]);578    try std.testing.expectEqual(entry.Category.fused, MatrixProductBiasGelu2x3x4F32.category);579    try std.testing.expect(MatrixProductBiasGelu2x3x4F32.specialization.operationIs(.{ .linalg = .matrix_product }));580    try std.testing.expectEqualStrings("mk,kn,n->mn", MatrixProductBiasGelu2x3x4F32.specialization.equation.?);581    try std.testing.expectEqual(@as(usize, 3), MatrixProductBiasGelu2x3x4F32.specialization.inputs.len);582    try std.testing.expectEqual(@as(u64, 3), MatrixProductBiasGelu2x3x4F32.specialization.inputs[2].elementCount().?);583    try std.testing.expectEqual(@as(u64, 6), MatrixProductBiasGelu2x3x4F32.specialization.outputs[0].elementCount().?);584    try std.testing.expect(MatrixProductBiasGelu2x3x4F32.specialization.epilogueMatches(0, .{585        .operator = .bias_add,586        .input_index = 2,587        .extents = &.{3},588    }));589    try std.testing.expect(MatrixProductBiasGelu2x3x4F32.specialization.epilogueMatches(1, .{ .operator = .{ .activation = .gelu } }));590    try std.testing.expectEqualDeep(MatrixProductBiasGelu2x3x4F32.specialization.launch.?, MatrixProductBiasGelu2x3x4F32.specialization.schedule.?.launch());591592    var snapshot = try MatrixProductBiasGelu2x3x4F32.scheduleSnapshot(std.testing.allocator, MatrixProductBiasGelu2x3x4F32.Limits.testing);593    defer snapshot.deinit(std.testing.allocator);594    try std.testing.expect(MatrixProductBiasGelu2x3x4F32.specialization.schedule.?.matchesSnapshot(&snapshot));595}596597test "fused matrix product bias relu and silu entries run on CPU" {598    var lhs = [_]f32{599        1.0, 2.0, 3.0, 4.0,600        5.0, 6.0, 7.0, 8.0,601    };602    var rhs = [_]f32{603        1.0, 0.0, 2.0,604        0.0, 1.0, 3.0,605        1.0, 1.0, 0.0,606        2.0, 0.0, 1.0,607    };608    var bias = [_]f32{ 0.25, -20.0, 0.5 };609    var relu_dst = @as([6]f32, @splat(0.0));610    var silu_dst = @as([6]f32, @splat(0.0));611612    try MatrixProductBiasRelu2x3x4F32.runCpu(std.testing.allocator, MatrixProductBiasRelu2x3x4F32.Limits.testing, &.{613        kernel.argumentBuffer(f32, relu_dst[0..]),614        kernel.argumentBuffer(f32, lhs[0..]),615        kernel.argumentBuffer(f32, rhs[0..]),616        kernel.argumentBuffer(f32, bias[0..]),617    });618    try MatrixProductBiasSilu2x3x4F32.runCpu(std.testing.allocator, MatrixProductBiasSilu2x3x4F32.Limits.testing, &.{619        kernel.argumentBuffer(f32, silu_dst[0..]),620        kernel.argumentBuffer(f32, lhs[0..]),621        kernel.argumentBuffer(f32, rhs[0..]),622        kernel.argumentBuffer(f32, bias[0..]),623    });624625    const dot = [_]f32{ 12.0, 5.0, 12.0, 28.0, 13.0, 36.0 };626    for (dot, relu_dst, silu_dst, 0..) |sum, relu_actual, silu_actual, index| {627        try std.testing.expectApproxEqAbs(expectedBiasActivation(.relu, sum, bias[index % 3]), relu_actual, 0.0001);628        try std.testing.expectApproxEqAbs(expectedBiasActivation(.silu, sum, bias[index % 3]), silu_actual, 0.0001);629    }630631    try std.testing.expect(MatrixProductBiasRelu2x3x4F32.specialization.epilogueMatches(1, .{ .operator = .{ .activation = .relu } }));632    try std.testing.expect(MatrixProductBiasSilu2x3x4F32.specialization.epilogueMatches(1, .{ .operator = .{ .activation = .silu } }));633    try std.testing.expectEqualDeep(MatrixProductBiasRelu2x3x4F32.specialization.launch.?, MatrixProductBiasRelu2x3x4F32.specialization.schedule.?.launch());634    try std.testing.expectEqualDeep(MatrixProductBiasSilu2x3x4F32.specialization.launch.?, MatrixProductBiasSilu2x3x4F32.specialization.schedule.?.launch());635}636637test "fused matrix vector product bias gelu entry runs on CPU and records epilogue operation" {638    var matrix = [_]f32{639        1.0,  2.0,  3.0,  4.0, 5.0,  6.0, 7.0,  8.0,640        2.0,  0.0,  -2.0, 0.0, 1.0,  0.0, -1.0, 0.0,641        0.5,  1.0,  1.5,  2.0, 2.5,  3.0, 3.5,  4.0,642        -1.0, -2.0, 3.0,  4.0, -5.0, 6.0, 7.0,  -8.0,643    };644    var vector = [_]f32{ 1.0, 0.5, -1.0, 2.0, 0.25, -0.5, 1.5, -2.0 };645    var bias = [_]f32{ 0.25, -1.0, 0.5, -20.0 };646    var dst = @as([4]f32, @splat(0.0));647648    try MatrixVectorProductBiasGelu4x8F32.runCpu(std.testing.allocator, MatrixVectorProductBiasGelu4x8F32.Limits.testing, &.{649        kernel.argumentBuffer(f32, dst[0..]),650        kernel.argumentBuffer(f32, matrix[0..]),651        kernel.argumentBuffer(f32, vector[0..]),652        kernel.argumentBuffer(f32, bias[0..]),653    });654    const dot = [_]f32{ -0.25, 2.75, -0.125, 25.25 };655    for (dot, bias, dst) |sum, bias_value, actual| {656        try std.testing.expectApproxEqAbs(expectedBiasActivation(.gelu, sum, bias_value), actual, 0.0001);657    }658659    const launch_value = try MatrixVectorProductBiasGelu4x8F32.launch(std.testing.allocator, MatrixVectorProductBiasGelu4x8F32.Limits.testing);660    try std.testing.expectEqual(@as(u32, 1), launch_value.grid[0]);661    try std.testing.expectEqual(@as(u32, 4), launch_value.block[0]);662    try std.testing.expectEqual(entry.Category.fused, MatrixVectorProductBiasGelu4x8F32.category);663    try std.testing.expect(MatrixVectorProductBiasGelu4x8F32.specialization.operationIs(.{ .linalg = .matrix_vector_product }));664    try std.testing.expectEqualStrings("mk,k,m->m", MatrixVectorProductBiasGelu4x8F32.specialization.equation.?);665    try std.testing.expectEqual(@as(usize, 3), MatrixVectorProductBiasGelu4x8F32.specialization.inputs.len);666    try std.testing.expectEqual(@as(u64, 4), MatrixVectorProductBiasGelu4x8F32.specialization.inputs[2].elementCount().?);667    try std.testing.expectEqual(@as(u64, 4), MatrixVectorProductBiasGelu4x8F32.specialization.outputs[0].elementCount().?);668    try std.testing.expect(MatrixVectorProductBiasGelu4x8F32.specialization.epilogueMatches(0, .{669        .operator = .bias_add,670        .input_index = 2,671        .extents = &.{4},672    }));673    try std.testing.expect(MatrixVectorProductBiasGelu4x8F32.specialization.epilogueMatches(1, .{ .operator = .{ .activation = .gelu } }));674    try std.testing.expectEqualDeep(MatrixVectorProductBiasGelu4x8F32.specialization.launch.?, MatrixVectorProductBiasGelu4x8F32.specialization.schedule.?.launch());675676    var snapshot = try MatrixVectorProductBiasGelu4x8F32.scheduleSnapshot(std.testing.allocator, MatrixVectorProductBiasGelu4x8F32.Limits.testing);677    defer snapshot.deinit(std.testing.allocator);678    try std.testing.expect(MatrixVectorProductBiasGelu4x8F32.specialization.schedule.?.matchesSnapshot(&snapshot));679}680681test "fused matrix vector product bias relu and silu entries run on CPU" {682    var matrix = [_]f32{683        1.0,  2.0,  3.0,  4.0, 5.0,  6.0, 7.0,  8.0,684        2.0,  0.0,  -2.0, 0.0, 1.0,  0.0, -1.0, 0.0,685        0.5,  1.0,  1.5,  2.0, 2.5,  3.0, 3.5,  4.0,686        -1.0, -2.0, 3.0,  4.0, -5.0, 6.0, 7.0,  -8.0,687    };688    var vector = [_]f32{ 1.0, 0.5, -1.0, 2.0, 0.25, -0.5, 1.5, -2.0 };689    var bias = [_]f32{ 0.25, -1.0, 0.5, -20.0 };690    var relu_dst = @as([4]f32, @splat(0.0));691    var silu_dst = @as([4]f32, @splat(0.0));692693    try MatrixVectorProductBiasRelu4x8F32.runCpu(std.testing.allocator, MatrixVectorProductBiasRelu4x8F32.Limits.testing, &.{694        kernel.argumentBuffer(f32, relu_dst[0..]),695        kernel.argumentBuffer(f32, matrix[0..]),696        kernel.argumentBuffer(f32, vector[0..]),697        kernel.argumentBuffer(f32, bias[0..]),698    });699    try MatrixVectorProductBiasSilu4x8F32.runCpu(std.testing.allocator, MatrixVectorProductBiasSilu4x8F32.Limits.testing, &.{700        kernel.argumentBuffer(f32, silu_dst[0..]),701        kernel.argumentBuffer(f32, matrix[0..]),702        kernel.argumentBuffer(f32, vector[0..]),703        kernel.argumentBuffer(f32, bias[0..]),704    });705706    const dot = [_]f32{ -0.25, 2.75, -0.125, 25.25 };707    for (dot, bias, relu_dst, silu_dst) |sum, bias_value, relu_actual, silu_actual| {708        try std.testing.expectApproxEqAbs(expectedBiasActivation(.relu, sum, bias_value), relu_actual, 0.0001);709        try std.testing.expectApproxEqAbs(expectedBiasActivation(.silu, sum, bias_value), silu_actual, 0.0001);710    }711712    try std.testing.expect(MatrixVectorProductBiasRelu4x8F32.specialization.epilogueMatches(1, .{ .operator = .{ .activation = .relu } }));713    try std.testing.expect(MatrixVectorProductBiasSilu4x8F32.specialization.epilogueMatches(1, .{ .operator = .{ .activation = .silu } }));714    try std.testing.expectEqualDeep(MatrixVectorProductBiasRelu4x8F32.specialization.launch.?, MatrixVectorProductBiasRelu4x8F32.specialization.schedule.?.launch());715    try std.testing.expectEqualDeep(MatrixVectorProductBiasSilu4x8F32.specialization.launch.?, MatrixVectorProductBiasSilu4x8F32.specialization.schedule.?.launch());716}717718test "fused bias gelu constructor creates independent shape-specialized entries" {719    const BiasGelu16F32 = biasGeluF32(.{ .extent = 16, .threads = 8 });720721    try std.testing.expectEqualStrings("accy.kernel.fused.bias_gelu8x4_f32", BiasGelu8F32.target);722    try std.testing.expectEqualStrings("accy.kernel.fused.bias_gelu16x8_f32", BiasGelu16F32.target);723    try std.testing.expectEqual(@as(u64, 8), BiasGelu8F32.specialization.outputs[0].elementCount().?);724    try std.testing.expectEqual(@as(u64, 16), BiasGelu16F32.specialization.outputs[0].elementCount().?);725    try std.testing.expect(BiasGelu16F32.specialization.operationIs(.{ .elementwise = .add }));726    try std.testing.expect(BiasGelu16F32.specialization.epilogueMatches(0, .{ .operator = .{ .activation = .gelu } }));727    try std.testing.expectEqual(@as(u32, 8), BiasGelu16F32.specialization.launch.?.threadgroup[0]);728    try std.testing.expectEqual(@as(u32, 2), BiasGelu16F32.specialization.launch.?.grid[0]);729}730731test "fused swiglu constructor creates independent shape-specialized entries" {732    const SwiGlu16F32 = swiGluF32(.{ .extent = 16, .threads = 8 });733734    try std.testing.expectEqualStrings("accy.kernel.fused.swiglu8x4_f32", SwiGlu8F32.target);735    try std.testing.expectEqualStrings("accy.kernel.fused.swiglu16x8_f32", SwiGlu16F32.target);736    try std.testing.expectEqual(@as(u64, 8), SwiGlu8F32.specialization.outputs[0].elementCount().?);737    try std.testing.expectEqual(@as(u64, 16), SwiGlu16F32.specialization.outputs[0].elementCount().?);738    try std.testing.expect(SwiGlu16F32.specialization.operationIs(.{ .elementwise = .mul }));739    try std.testing.expect(SwiGlu16F32.specialization.inputTransformMatches(0, .{740        .operator = .{ .activation = .silu },741        .input_index = 0,742        .extents = &.{16},743    }));744    try std.testing.expectEqual(@as(u32, 8), SwiGlu16F32.specialization.launch.?.threadgroup[0]);745    try std.testing.expectEqual(@as(u32, 2), SwiGlu16F32.specialization.launch.?.grid[0]);746}747748test "fused geglu constructor creates independent shape-specialized entries" {749    const GeGlu16F32 = geGluF32(.{ .extent = 16, .threads = 8 });750751    try std.testing.expectEqualStrings("accy.kernel.fused.geglu8x4_f32", GeGlu8F32.target);752    try std.testing.expectEqualStrings("accy.kernel.fused.geglu16x8_f32", GeGlu16F32.target);753    try std.testing.expectEqual(@as(u64, 8), GeGlu8F32.specialization.outputs[0].elementCount().?);754    try std.testing.expectEqual(@as(u64, 16), GeGlu16F32.specialization.outputs[0].elementCount().?);755    try std.testing.expect(GeGlu16F32.specialization.operationIs(.{ .elementwise = .mul }));756    try std.testing.expect(GeGlu16F32.specialization.inputTransformMatches(0, .{757        .operator = .{ .activation = .gelu },758        .input_index = 0,759        .extents = &.{16},760    }));761    try std.testing.expectEqual(@as(u32, 8), GeGlu16F32.specialization.launch.?.threadgroup[0]);762    try std.testing.expectEqual(@as(u32, 2), GeGlu16F32.specialization.launch.?.grid[0]);763}764765test "fused reglu constructor creates independent shape-specialized entries" {766    const ReGlu16F32 = reGluF32(.{ .extent = 16, .threads = 8 });767768    try std.testing.expectEqualStrings("accy.kernel.fused.reglu8x4_f32", ReGlu8F32.target);769    try std.testing.expectEqualStrings("accy.kernel.fused.reglu16x8_f32", ReGlu16F32.target);770    try std.testing.expectEqual(@as(u64, 8), ReGlu8F32.specialization.outputs[0].elementCount().?);771    try std.testing.expectEqual(@as(u64, 16), ReGlu16F32.specialization.outputs[0].elementCount().?);772    try std.testing.expect(ReGlu16F32.specialization.operationIs(.{ .elementwise = .mul }));773    try std.testing.expect(ReGlu16F32.specialization.inputTransformMatches(0, .{774        .operator = .{ .activation = .relu },775        .input_index = 0,776        .extents = &.{16},777    }));778    try std.testing.expectEqual(@as(u32, 8), ReGlu16F32.specialization.launch.?.threadgroup[0]);779    try std.testing.expectEqual(@as(u32, 2), ReGlu16F32.specialization.launch.?.grid[0]);780}781782test "fused matrix product bias gelu constructor creates independent shape-specialized entries" {783    const MatrixProductBiasGelu4x5x6F32 = matrixProductBiasGeluF32(.{784        .m = 4,785        .n = 5,786        .k = 6,787        .threads = .{ .x = 4, .y = 2 },788    });789790    try std.testing.expectEqualStrings("accy.kernel.fused.matmul_bias_gelu2x3x4_2x2_f32", MatrixProductBiasGelu2x3x4F32.target);791    try std.testing.expectEqualStrings("accy.kernel.fused.matmul_bias_gelu2x3x4_1x2_f32", MatrixProductBiasGelu2x3x4ThreadBlocks1x2F32.target);792    try std.testing.expectEqualStrings("accy.kernel.fused.matmul_bias_gelu4x5x6_4x2_f32", MatrixProductBiasGelu4x5x6F32.target);793    try std.testing.expect(MatrixProductBiasGelu2x3x4ThreadBlocks1x2F32.specialization.operationIs(.{ .linalg = .matrix_product }));794    try std.testing.expectEqual(@as(u32, 1), MatrixProductBiasGelu2x3x4ThreadBlocks1x2F32.specialization.launch.?.threadgroup[0]);795    try std.testing.expectEqual(@as(u32, 2), MatrixProductBiasGelu2x3x4ThreadBlocks1x2F32.specialization.launch.?.threadgroup[1]);796    try std.testing.expectEqualDeep(MatrixProductBiasGelu2x3x4ThreadBlocks1x2F32.specialization.launch.?, MatrixProductBiasGelu2x3x4ThreadBlocks1x2F32.specialization.schedule.?.launch());797    try std.testing.expect(MatrixProductBiasGelu4x5x6F32.specialization.operationIs(.{ .linalg = .matrix_product }));798    try std.testing.expectEqual(@as(usize, 3), MatrixProductBiasGelu4x5x6F32.specialization.inputs.len);799    try std.testing.expectEqual(@as(u64, 24), MatrixProductBiasGelu4x5x6F32.specialization.inputs[0].elementCount().?);800    try std.testing.expectEqual(@as(u64, 30), MatrixProductBiasGelu4x5x6F32.specialization.inputs[1].elementCount().?);801    try std.testing.expectEqual(@as(u64, 5), MatrixProductBiasGelu4x5x6F32.specialization.inputs[2].elementCount().?);802    try std.testing.expectEqual(@as(u64, 20), MatrixProductBiasGelu4x5x6F32.specialization.outputs[0].elementCount().?);803    try std.testing.expectEqual(entry.ReductionOperator.dot_product, MatrixProductBiasGelu4x5x6F32.specialization.reductions[0].operator);804    try std.testing.expectEqual(@as(u64, 6), MatrixProductBiasGelu4x5x6F32.specialization.reductions[0].shape.elementCount().?);805    try std.testing.expect(MatrixProductBiasGelu4x5x6F32.specialization.epilogueMatches(0, .{806        .operator = .bias_add,807        .input_index = 2,808        .extents = &.{5},809    }));810    try std.testing.expect(MatrixProductBiasGelu4x5x6F32.specialization.epilogueMatches(1, .{ .operator = .{ .activation = .gelu } }));811    try std.testing.expectEqualDeep(MatrixProductBiasGelu4x5x6F32.specialization.launch.?, MatrixProductBiasGelu4x5x6F32.specialization.schedule.?.launch());812}813814test "fused matrix vector product bias gelu constructor creates independent shape-specialized entries" {815    const MatrixVectorProductBiasGelu5x6F32 = matrixVectorProductBiasGeluF32(.{816        .m = 5,817        .k = 6,818        .threads = 4,819    });820821    try std.testing.expectEqualStrings("accy.kernel.fused.matvec_bias_gelu4x8_4x_f32", MatrixVectorProductBiasGelu4x8F32.target);822    try std.testing.expectEqualStrings("accy.kernel.fused.matvec_bias_gelu5x6_4x_f32", MatrixVectorProductBiasGelu5x6F32.target);823    try std.testing.expect(MatrixVectorProductBiasGelu5x6F32.specialization.operationIs(.{ .linalg = .matrix_vector_product }));824    try std.testing.expectEqual(@as(usize, 3), MatrixVectorProductBiasGelu5x6F32.specialization.inputs.len);825    try std.testing.expectEqual(@as(u64, 30), MatrixVectorProductBiasGelu5x6F32.specialization.inputs[0].elementCount().?);826    try std.testing.expectEqual(@as(u64, 6), MatrixVectorProductBiasGelu5x6F32.specialization.inputs[1].elementCount().?);827    try std.testing.expectEqual(@as(u64, 5), MatrixVectorProductBiasGelu5x6F32.specialization.inputs[2].elementCount().?);828    try std.testing.expectEqual(@as(u64, 5), MatrixVectorProductBiasGelu5x6F32.specialization.outputs[0].elementCount().?);829    try std.testing.expectEqual(entry.ReductionOperator.dot_product, MatrixVectorProductBiasGelu5x6F32.specialization.reductions[0].operator);830    try std.testing.expectEqual(@as(u64, 6), MatrixVectorProductBiasGelu5x6F32.specialization.reductions[0].shape.elementCount().?);831    try std.testing.expect(MatrixVectorProductBiasGelu5x6F32.specialization.epilogueMatches(0, .{832        .operator = .bias_add,833        .input_index = 2,834        .extents = &.{5},835    }));836    try std.testing.expect(MatrixVectorProductBiasGelu5x6F32.specialization.epilogueMatches(1, .{ .operator = .{ .activation = .gelu } }));837    try std.testing.expectEqualDeep(MatrixVectorProductBiasGelu5x6F32.specialization.launch.?, MatrixVectorProductBiasGelu5x6F32.specialization.schedule.?.launch());838}839840test "fused bias gelu entry creates registry-ready artifact" {841    const allocator = std.testing.allocator;842    var state = gpu.recording.BackendState{843        .allocator = allocator,844        .kind = .cuda,845        .format = .cuda_ptx,846    };847848    var call_artifact = try BiasGelu8F32.createKernelCallArtifact(allocator, state.handle(), .{ .limits = BiasGelu8F32.Limits.testing });849    defer call_artifact.deinit();850851    const artifact = call_artifact.registry().find(BiasGelu8F32.target, BiasGelu8F32.version, .cuda_ptx) orelse {852        return error.TestExpectedKernelCallArtifact;853    };854    try std.testing.expectEqualStrings(BiasGelu8F32.name, artifact.entry_name);855    try std.testing.expectEqual(@as(u32, 3), artifact.argument_count);856    try std.testing.expectEqual(gpu.DTypeSet.init(&.{.f32}).bits, artifact.required_dtypes.bits);857    switch (artifact.launch) {858        .fixed => |geometry| {859            try std.testing.expectEqual(BiasGelu8F32.specialization.launch.?.grid[0], geometry.grid[0]);860            try std.testing.expectEqual(BiasGelu8F32.specialization.launch.?.threadgroup[0], geometry.threadgroup[0]);861        },862        else => return error.TestExpectedFixedLaunch,863    }864}865866test "fused swiglu entry creates registry-ready artifact" {867    const allocator = std.testing.allocator;868    var state = gpu.recording.BackendState{869        .allocator = allocator,870        .kind = .cuda,871        .format = .cuda_ptx,872    };873874    var call_artifact = try SwiGlu8F32.createKernelCallArtifact(allocator, state.handle(), .{ .limits = SwiGlu8F32.Limits.testing });875    defer call_artifact.deinit();876877    const artifact = call_artifact.registry().find(SwiGlu8F32.target, SwiGlu8F32.version, .cuda_ptx) orelse {878        return error.TestExpectedKernelCallArtifact;879    };880    try std.testing.expectEqualStrings(SwiGlu8F32.name, artifact.entry_name);881    try std.testing.expectEqual(@as(u32, 3), artifact.argument_count);882    try std.testing.expectEqual(gpu.DTypeSet.init(&.{.f32}).bits, artifact.required_dtypes.bits);883    switch (artifact.launch) {884        .fixed => |geometry| {885            try std.testing.expectEqual(SwiGlu8F32.specialization.launch.?.grid[0], geometry.grid[0]);886            try std.testing.expectEqual(SwiGlu8F32.specialization.launch.?.threadgroup[0], geometry.threadgroup[0]);887        },888        else => return error.TestExpectedFixedLaunch,889    }890}891892test "fused geglu entry creates registry-ready artifact" {893    const allocator = std.testing.allocator;894    var state = gpu.recording.BackendState{895        .allocator = allocator,896        .kind = .cuda,897        .format = .cuda_ptx,898    };899900    var call_artifact = try GeGlu8F32.createKernelCallArtifact(allocator, state.handle(), .{ .limits = GeGlu8F32.Limits.testing });901    defer call_artifact.deinit();902903    const artifact = call_artifact.registry().find(GeGlu8F32.target, GeGlu8F32.version, .cuda_ptx) orelse {904        return error.TestExpectedKernelCallArtifact;905    };906    try std.testing.expectEqualStrings(GeGlu8F32.name, artifact.entry_name);907    try std.testing.expectEqual(@as(u32, 3), artifact.argument_count);908    try std.testing.expectEqual(gpu.DTypeSet.init(&.{.f32}).bits, artifact.required_dtypes.bits);909    switch (artifact.launch) {910        .fixed => |geometry| {911            try std.testing.expectEqual(GeGlu8F32.specialization.launch.?.grid[0], geometry.grid[0]);912            try std.testing.expectEqual(GeGlu8F32.specialization.launch.?.threadgroup[0], geometry.threadgroup[0]);913        },914        else => return error.TestExpectedFixedLaunch,915    }916}917918test "fused reglu entry creates registry-ready artifact" {919    const allocator = std.testing.allocator;920    var state = gpu.recording.BackendState{921        .allocator = allocator,922        .kind = .cuda,923        .format = .cuda_ptx,924    };925926    var call_artifact = try ReGlu8F32.createKernelCallArtifact(allocator, state.handle(), .{ .limits = ReGlu8F32.Limits.testing });927    defer call_artifact.deinit();928929    const artifact = call_artifact.registry().find(ReGlu8F32.target, ReGlu8F32.version, .cuda_ptx) orelse {930        return error.TestExpectedKernelCallArtifact;931    };932    try std.testing.expectEqualStrings(ReGlu8F32.name, artifact.entry_name);933    try std.testing.expectEqual(@as(u32, 3), artifact.argument_count);934    try std.testing.expectEqual(gpu.DTypeSet.init(&.{.f32}).bits, artifact.required_dtypes.bits);935    switch (artifact.launch) {936        .fixed => |geometry| {937            try std.testing.expectEqual(ReGlu8F32.specialization.launch.?.grid[0], geometry.grid[0]);938            try std.testing.expectEqual(ReGlu8F32.specialization.launch.?.threadgroup[0], geometry.threadgroup[0]);939        },940        else => return error.TestExpectedFixedLaunch,941    }942}943944test "fused matrix product bias gelu entry creates registry-ready artifact" {945    const allocator = std.testing.allocator;946    var state = gpu.recording.BackendState{947        .allocator = allocator,948        .kind = .cuda,949        .format = .cuda_ptx,950    };951952    var call_artifact = try MatrixProductBiasGelu2x3x4F32.createKernelCallArtifact(allocator, state.handle(), .{ .limits = MatrixProductBiasGelu2x3x4F32.Limits.testing });953    defer call_artifact.deinit();954955    const artifact = call_artifact.registry().find(MatrixProductBiasGelu2x3x4F32.target, MatrixProductBiasGelu2x3x4F32.version, .cuda_ptx) orelse {956        return error.TestExpectedKernelCallArtifact;957    };958    try std.testing.expectEqualStrings(MatrixProductBiasGelu2x3x4F32.name, artifact.entry_name);959    try std.testing.expectEqual(@as(u32, 4), artifact.argument_count);960    try std.testing.expectEqual(gpu.DTypeSet.init(&.{.f32}).bits, artifact.required_dtypes.bits);961    switch (artifact.launch) {962        .fixed => |geometry| {963            try std.testing.expectEqual(MatrixProductBiasGelu2x3x4F32.specialization.launch.?.grid[0], geometry.grid[0]);964            try std.testing.expectEqual(MatrixProductBiasGelu2x3x4F32.specialization.launch.?.grid[1], geometry.grid[1]);965            try std.testing.expectEqual(MatrixProductBiasGelu2x3x4F32.specialization.launch.?.threadgroup[0], geometry.threadgroup[0]);966            try std.testing.expectEqual(MatrixProductBiasGelu2x3x4F32.specialization.launch.?.threadgroup[1], geometry.threadgroup[1]);967        },968        else => return error.TestExpectedFixedLaunch,969    }970}971972test "fused matrix vector product bias gelu entry creates registry-ready artifact" {973    const allocator = std.testing.allocator;974    var state = gpu.recording.BackendState{975        .allocator = allocator,976        .kind = .cuda,977        .format = .cuda_ptx,978    };979980    var call_artifact = try MatrixVectorProductBiasGelu4x8F32.createKernelCallArtifact(allocator, state.handle(), .{ .limits = MatrixVectorProductBiasGelu4x8F32.Limits.testing });981    defer call_artifact.deinit();982983    const artifact = call_artifact.registry().find(MatrixVectorProductBiasGelu4x8F32.target, MatrixVectorProductBiasGelu4x8F32.version, .cuda_ptx) orelse {984        return error.TestExpectedKernelCallArtifact;985    };986    try std.testing.expectEqualStrings(MatrixVectorProductBiasGelu4x8F32.name, artifact.entry_name);987    try std.testing.expectEqual(@as(u32, 4), artifact.argument_count);988    try std.testing.expectEqual(gpu.DTypeSet.init(&.{.f32}).bits, artifact.required_dtypes.bits);989    switch (artifact.launch) {990        .fixed => |geometry| {991            try std.testing.expectEqual(MatrixVectorProductBiasGelu4x8F32.specialization.launch.?.grid[0], geometry.grid[0]);992            try std.testing.expectEqual(MatrixVectorProductBiasGelu4x8F32.specialization.launch.?.threadgroup[0], geometry.threadgroup[0]);993        },994        else => return error.TestExpectedFixedLaunch,995    }996}

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