tiny.accy.kernel.library.factor

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 kernel = @import("../root.zig");910const DType = choir_abi.DType;11const runtimeExtentArgument = extent_mod.runtimeExtentArgument;1213pub const TileLayout = enum {14    row_major,15    interleaved,16};1718pub const BatchedCholesky = struct {19    batch: u64,20    n: u32 = 3,21    threads: u32 = 256,22    layout: TileLayout = .row_major,23    batch_axis: []const u8 = "b",24};2526pub const batched_cholesky_family_version: u32 = 1;27pub const batched_cholesky_min_n: u32 = 2;28pub const batched_cholesky_max_n: u32 = 4;29pub const batched_cholesky_max_threads: u32 = 1024;30pub const batched_cholesky_max_blocks: u32 = 1024;3132pub fn batchedCholeskyBlockCount(batch: u64, threads: u32) u64 {33    return (batch + threads - 1) / threads;34}3536pub fn batchedCholeskyInstanceValid(instance: BatchedCholesky) bool {37    if (instance.batch == 0) return false;38    if (instance.n < batched_cholesky_min_n or instance.n > batched_cholesky_max_n) return false;39    if (instance.threads == 0 or instance.threads > batched_cholesky_max_threads) return false;40    return extent_mod.blockCountWithinLimit(instance.batch, instance.threads, batched_cholesky_max_blocks);41}4243fn factorElementIndex(44    inner: anytype,45    layout: TileLayout,46    base: kernel.Value,47    batch_index: kernel.Value,48    system: kernel.Value,49    slot: usize,50) !kernel.Value {51    return switch (layout) {52        .row_major => inner.add(base, try inner.constantIndex(@intCast(slot))),53        .interleaved => inner.add(54            try inner.mul(try inner.constantIndex(@intCast(slot)), batch_index),55            system,56        ),57    };58}5960fn batched_cholesky_body_active(inner: anytype, ctx: anytype) !void {61    const n: usize = @intCast(ctx.n);62    const tile = try inner.constantIndex(@intCast(n * n));63    const base = try inner.mul(ctx.system, tile);6465    var a_values: [batched_cholesky_max_n][batched_cholesky_max_n]kernel.Value = undefined;66    for (0..n) |i| {67        for (0..i + 1) |j| {68            const index = try factorElementIndex(inner, ctx.layout, base, ctx.batch_index, ctx.system, i * n + j);69            const loaded = try ctx.args.param(.a).load(inner, index);70            a_values[i][j] = loaded.raw();71        }72    }7374    var l_values: [batched_cholesky_max_n][batched_cholesky_max_n]kernel.Value = undefined;75    for (0..n) |j| {76        var diagonal = a_values[j][j];77        for (0..j) |c| {78            diagonal = try inner.sub(diagonal, try inner.mul(l_values[j][c], l_values[j][c]));79        }80        l_values[j][j] = try inner.sqrt(diagonal);81        for (j + 1..n) |i| {82            var sum = a_values[i][j];83            for (0..j) |c| {84                sum = try inner.sub(sum, try inner.mul(l_values[i][c], l_values[j][c]));85            }86            l_values[i][j] = try inner.div(sum, l_values[j][j]);87        }88    }8990    const zero = try inner.constantFloat(.f32, 0);91    for (0..n) |i| {92        for (0..n) |j| {93            const value = if (j <= i) l_values[i][j] else zero;94            const index = try factorElementIndex(inner, ctx.layout, base, ctx.batch_index, ctx.system, i * n + j);95            try ctx.args.param(.l).store(inner, value, index);96        }97    }98}99100fn batchedCholeskyBody(k: anytype, spec: BatchedCholesky, args: anytype) !void {101    if (!batchedCholeskyInstanceValid(spec)) return error.UnsupportedBatchedCholeskyInstance;102    const system = try k.globalId(.x);103    const batch = try k.castIndex(args.param(.batch).raw());104    const active = try k.compare(.lt, system, batch);105    try k.guardDo(active, .{106        .args = args,107        .system = system,108        .n = spec.n,109        .layout = spec.layout,110        .batch_index = batch,111    }, batched_cholesky_body_active);112}113114fn batchedCholeskyFamilySchedule(instance: BatchedCholesky) kernel.logical.schedule.ThreadBlocks {115    return kernel.logical.schedule.threadBlocks(.{ .x = instance.threads });116}117118fn batchedCholeskyRuntimeFamily() type {119    return kernel.logical.Family(.{120        .name = "accy_kernel_linalg_batched_cholesky_runtime_f32",121        .parameters = .{122            .l = kernel.dynamicBuffer(.f32),123            .a = kernel.dynamicBuffer(.f32),124            .batch = kernel.scalar(.i32),125        },126        .Instance = BatchedCholesky,127        .schedule = batchedCholeskyFamilySchedule,128        .body = batchedCholeskyBody,129    });130}131132pub const BatchedCholeskyRuntimeFamilyF32 = batchedCholeskyRuntimeFamily();133134fn layoutTargetSegment(layout: TileLayout) []const u8 {135    return switch (layout) {136        .row_major => "",137        .interleaved => "il_",138    };139}140141pub fn batchedCholeskyFamilyTarget(allocator: std.mem.Allocator, instance: BatchedCholesky) ![]u8 {142    return std.fmt.allocPrint(143        allocator,144        "accy.kernel.linalg.batched_cholesky_family_{d}_{d}_{s}f32",145        .{ instance.n, instance.threads, layoutTargetSegment(instance.layout) },146    );147}148149pub fn batchedCholeskyFamilyEntryName(allocator: std.mem.Allocator, instance: BatchedCholesky) ![]u8 {150    return std.fmt.allocPrint(151        allocator,152        "accy_kernel_linalg_batched_cholesky_family_{d}_{d}_{s}f32",153        .{ instance.n, instance.threads, layoutTargetSegment(instance.layout) },154    );155}156157pub fn batchedCholeskyRuntimeArguments(instance: BatchedCholesky) ![1]choir_abi.ScalarArgument {158    return .{159        .{ .u32 = try runtimeExtentArgument(instance.batch) },160    };161}162163const testing = std.testing;164165fn hostCholesky(n: usize, a: []const f32, l: []f32) void {166    for (0..n) |i| {167        for (0..n) |j| {168            l[i * n + j] = 0;169        }170    }171    for (0..n) |j| {172        var diagonal = a[j * n + j];173        for (0..j) |c| {174            diagonal -= l[j * n + c] * l[j * n + c];175        }176        l[j * n + j] = @sqrt(diagonal);177        for (j + 1..n) |i| {178            var sum = a[i * n + j];179            for (0..j) |c| {180                sum -= l[i * n + c] * l[j * n + c];181            }182            l[i * n + j] = sum / l[j * n + j];183        }184    }185}186187fn fillSpdTile(n: usize, seed_base: u32, a: []f32) void {188    var m: [batched_cholesky_max_n * batched_cholesky_max_n]f32 = undefined;189    var seed: u32 = seed_base | 1;190    for (m[0 .. n * n]) |*value| {191        seed ^= seed << 13;192        seed ^= seed >> 17;193        seed ^= seed << 5;194        value.* = @as(f32, @floatFromInt(seed % 1000)) / 1000.0;195    }196    for (0..n) |i| {197        for (0..n) |j| {198            var sum: f32 = 0;199            for (0..n) |c| {200                sum += m[i * n + c] * m[j * n + c];201            }202            if (i == j) sum += @floatFromInt(n);203            a[i * n + j] = sum;204        }205    }206}207208fn expectBatchedCholeskyMatchesHost(comptime n: usize, threads: u32) !void {209    const allocator = testing.allocator;210    const batch: usize = 40;211    const instance = BatchedCholesky{ .batch = batch, .n = n, .threads = threads };212    const blocks: u32 = @intCast(batchedCholeskyBlockCount(instance.batch, instance.threads));213    try testing.expect(blocks > 1);214215    const a = try allocator.alloc(f32, batch * n * n);216    defer allocator.free(a);217    for (0..batch) |b| {218        fillSpdTile(n, @intCast(b * 977 + 13), a[b * n * n ..][0 .. n * n]);219    }220221    const expected = try allocator.alloc(f32, batch * n * n);222    defer allocator.free(expected);223    for (0..batch) |b| {224        hostCholesky(n, a[b * n * n ..][0 .. n * n], expected[b * n * n ..][0 .. n * n]);225    }226227    const l = try allocator.alloc(f32, batch * n * n);228    defer allocator.free(l);229    @memset(l, -1);230231    var graph = try BatchedCholeskyRuntimeFamilyF32.build(allocator, BatchedCholeskyRuntimeFamilyF32.Limits.testing, instance);232    defer graph.deinit();233    try graph.runCpuWithLaunch(allocator, &.{234        kernel.argumentBuffer(f32, l),235        kernel.argumentBuffer(f32, a),236        kernel.argumentI32(@intCast(batch)),237    }, .{238        .grid = .{ blocks, 1, 1 },239        .block = .{ instance.threads, 1, 1 },240    });241242    try testing.expectEqualSlices(f32, expected, l);243}244245test "linalg batched cholesky matches the host reference at every supported size" {246    try expectBatchedCholeskyMatchesHost(2, 32);247    try expectBatchedCholeskyMatchesHost(3, 32);248    try expectBatchedCholeskyMatchesHost(4, 32);249}250251test "linalg batched cholesky propagates NaN for non-SPD tiles in band" {252    const allocator = testing.allocator;253    const n: usize = 3;254    const batch: usize = 2;255    const instance = BatchedCholesky{ .batch = batch, .n = n, .threads = 32 };256257    var a = @as([(batch * n * n)]f32, @splat(0));258    fillSpdTile(n, 17, a[0 .. n * n]);259    fillSpdTile(n, 18, a[n * n ..][0 .. n * n]);260    a[n * n] = -4.0;261262    var l = @as([(batch * n * n)]f32, @splat(-1));263    var graph = try BatchedCholeskyRuntimeFamilyF32.build(allocator, BatchedCholeskyRuntimeFamilyF32.Limits.testing, instance);264    defer graph.deinit();265    try graph.runCpuWithLaunch(allocator, &.{266        kernel.argumentBuffer(f32, l[0..]),267        kernel.argumentBuffer(f32, a[0..]),268        kernel.argumentI32(@intCast(batch)),269    }, .{270        .grid = .{ 1, 1, 1 },271        .block = .{ instance.threads, 1, 1 },272    });273274    var expected_good: [n * n]f32 = undefined;275    hostCholesky(n, a[0 .. n * n], expected_good[0..]);276    try testing.expectEqualSlices(f32, expected_good[0..], l[0 .. n * n]);277    try testing.expect(std.math.isNan(l[n * n]));278    try testing.expect(std.math.isNan(l[n * n + n + 1]));279}280281test "linalg batched cholesky identity and validity" {282    const allocator = testing.allocator;283    const instance = BatchedCholesky{ .batch = 5000, .n = 3, .threads = 64 };284    const target = try batchedCholeskyFamilyTarget(allocator, instance);285    defer allocator.free(target);286    try testing.expectEqualStrings("accy.kernel.linalg.batched_cholesky_family_3_64_f32", target);287    const entry_name = try batchedCholeskyFamilyEntryName(allocator, instance);288    defer allocator.free(entry_name);289    try testing.expectEqualStrings("accy_kernel_linalg_batched_cholesky_family_3_64_f32", entry_name);290291    try testing.expect(batchedCholeskyInstanceValid(instance));292    try testing.expect(!batchedCholeskyInstanceValid(.{ .batch = 0, .n = 3, .threads = 64 }));293    try testing.expect(!batchedCholeskyInstanceValid(.{ .batch = 100, .n = 1, .threads = 64 }));294    try testing.expect(!batchedCholeskyInstanceValid(.{ .batch = 100, .n = 5, .threads = 64 }));295    try testing.expect(!batchedCholeskyInstanceValid(.{ .batch = 100, .n = 3, .threads = 0 }));296    try testing.expect(!batchedCholeskyInstanceValid(.{ .batch = std.math.maxInt(u64), .n = 3, .threads = 32 }));297298    const args = try batchedCholeskyRuntimeArguments(instance);299    try testing.expectEqual(@as(u32, 5000), args[0].u32);300}301302pub const BatchedCholeskySolve = struct {303    batch: u64,304    n: u32 = 3,305    threads: u32 = 256,306    layout: TileLayout = .row_major,307    batch_axis: []const u8 = "b",308};309310pub const batched_cholesky_solve_family_version: u32 = 1;311312pub fn batchedCholeskySolveInstanceValid(instance: BatchedCholeskySolve) bool {313    if (instance.batch == 0) return false;314    if (instance.n < batched_cholesky_min_n or instance.n > batched_cholesky_max_n) return false;315    if (instance.threads == 0 or instance.threads > batched_cholesky_max_threads) return false;316    return extent_mod.blockCountWithinLimit(instance.batch, instance.threads, batched_cholesky_max_blocks);317}318319fn batched_cholesky_solve_body_active(inner: anytype, ctx: anytype) !void {320    const n: usize = @intCast(ctx.n);321    const tile = try inner.constantIndex(@intCast(n * n));322    const tile_base = try inner.mul(ctx.system, tile);323    const width = try inner.constantIndex(@intCast(n));324    const vector_base = try inner.mul(ctx.system, width);325326    var l_values: [batched_cholesky_max_n][batched_cholesky_max_n]kernel.Value = undefined;327    for (0..n) |i| {328        for (0..i + 1) |j| {329            const index = try factorElementIndex(inner, ctx.layout, tile_base, ctx.batch_index, ctx.system, i * n + j);330            const loaded = try ctx.args.param(.l).load(inner, index);331            l_values[i][j] = loaded.raw();332        }333    }334335    var solution: [batched_cholesky_max_n]kernel.Value = undefined;336    for (0..n) |i| {337        const index = try factorElementIndex(inner, ctx.layout, vector_base, ctx.batch_index, ctx.system, i);338        const loaded = try ctx.args.param(.rhs).load(inner, index);339        var sum = loaded.raw();340        for (0..i) |j| {341            sum = try inner.sub(sum, try inner.mul(l_values[i][j], solution[j]));342        }343        solution[i] = try inner.div(sum, l_values[i][i]);344    }345346    var row = n;347    while (row > 0) {348        row -= 1;349        var sum = solution[row];350        for (row + 1..n) |j| {351            sum = try inner.sub(sum, try inner.mul(l_values[j][row], solution[j]));352        }353        solution[row] = try inner.div(sum, l_values[row][row]);354    }355356    for (0..n) |i| {357        const index = try factorElementIndex(inner, ctx.layout, vector_base, ctx.batch_index, ctx.system, i);358        try ctx.args.param(.x).store(inner, solution[i], index);359    }360}361362fn batchedCholeskySolveBody(k: anytype, spec: BatchedCholeskySolve, args: anytype) !void {363    if (!batchedCholeskySolveInstanceValid(spec)) return error.UnsupportedBatchedCholeskySolveInstance;364    const system = try k.globalId(.x);365    const batch = try k.castIndex(args.param(.batch).raw());366    const active = try k.compare(.lt, system, batch);367    try k.guardDo(active, .{368        .args = args,369        .system = system,370        .n = spec.n,371        .layout = spec.layout,372        .batch_index = batch,373    }, batched_cholesky_solve_body_active);374}375376fn batchedCholeskySolveFamilySchedule(instance: BatchedCholeskySolve) kernel.logical.schedule.ThreadBlocks {377    return kernel.logical.schedule.threadBlocks(.{ .x = instance.threads });378}379380fn batchedCholeskySolveRuntimeFamily() type {381    return kernel.logical.Family(.{382        .name = "accy_kernel_linalg_batched_cholesky_solve_runtime_f32",383        .parameters = .{384            .x = kernel.dynamicBuffer(.f32),385            .l = kernel.dynamicBuffer(.f32),386            .rhs = kernel.dynamicBuffer(.f32),387            .batch = kernel.scalar(.i32),388        },389        .Instance = BatchedCholeskySolve,390        .schedule = batchedCholeskySolveFamilySchedule,391        .body = batchedCholeskySolveBody,392    });393}394395pub const BatchedCholeskySolveRuntimeFamilyF32 = batchedCholeskySolveRuntimeFamily();396397pub fn batchedCholeskySolveFamilyTarget(allocator: std.mem.Allocator, instance: BatchedCholeskySolve) ![]u8 {398    return std.fmt.allocPrint(399        allocator,400        "accy.kernel.linalg.batched_cholesky_solve_family_{d}_{d}_{s}f32",401        .{ instance.n, instance.threads, layoutTargetSegment(instance.layout) },402    );403}404405pub fn batchedCholeskySolveFamilyEntryName(allocator: std.mem.Allocator, instance: BatchedCholeskySolve) ![]u8 {406    return std.fmt.allocPrint(407        allocator,408        "accy_kernel_linalg_batched_cholesky_solve_family_{d}_{d}_{s}f32",409        .{ instance.n, instance.threads, layoutTargetSegment(instance.layout) },410    );411}412413pub fn batchedCholeskySolveRuntimeArguments(instance: BatchedCholeskySolve) ![1]choir_abi.ScalarArgument {414    return .{415        .{ .u32 = try runtimeExtentArgument(instance.batch) },416    };417}418419fn hostCholeskySolve(n: usize, l: []const f32, rhs: []const f32, x: []f32) void {420    var solution: [batched_cholesky_max_n]f32 = undefined;421    for (0..n) |i| {422        var sum = rhs[i];423        for (0..i) |j| {424            sum -= l[i * n + j] * solution[j];425        }426        solution[i] = sum / l[i * n + i];427    }428    var row = n;429    while (row > 0) {430        row -= 1;431        var sum = solution[row];432        for (row + 1..n) |j| {433            sum -= l[j * n + row] * solution[j];434        }435        solution[row] = sum / l[row * n + row];436    }437    for (0..n) |i| {438        x[i] = solution[i];439    }440}441442fn expectBatchedCholeskySolveMatchesHost(comptime n: usize, threads: u32) !void {443    const allocator = testing.allocator;444    const batch: usize = 40;445    const instance = BatchedCholeskySolve{ .batch = batch, .n = n, .threads = threads };446    const blocks: u32 = @intCast(batchedCholeskyBlockCount(instance.batch, instance.threads));447    try testing.expect(blocks > 1);448449    const a = try allocator.alloc(f32, batch * n * n);450    defer allocator.free(a);451    const l = try allocator.alloc(f32, batch * n * n);452    defer allocator.free(l);453    const rhs = try allocator.alloc(f32, batch * n);454    defer allocator.free(rhs);455    var seed: u32 = 0x51f15eed;456    for (0..batch) |b| {457        fillSpdTile(n, @intCast(b * 977 + 13), a[b * n * n ..][0 .. n * n]);458        hostCholesky(n, a[b * n * n ..][0 .. n * n], l[b * n * n ..][0 .. n * n]);459        for (rhs[b * n ..][0..n]) |*value| {460            seed ^= seed << 13;461            seed ^= seed >> 17;462            seed ^= seed << 5;463            value.* = @as(f32, @floatFromInt(seed % 1000)) / 500.0 - 1.0;464        }465    }466467    const expected = try allocator.alloc(f32, batch * n);468    defer allocator.free(expected);469    for (0..batch) |b| {470        hostCholeskySolve(n, l[b * n * n ..][0 .. n * n], rhs[b * n ..][0..n], expected[b * n ..][0..n]);471    }472473    const x = try allocator.alloc(f32, batch * n);474    defer allocator.free(x);475    @memset(x, -1);476477    var graph = try BatchedCholeskySolveRuntimeFamilyF32.build(allocator, BatchedCholeskySolveRuntimeFamilyF32.Limits.testing, instance);478    defer graph.deinit();479    try graph.runCpuWithLaunch(allocator, &.{480        kernel.argumentBuffer(f32, x),481        kernel.argumentBuffer(f32, l),482        kernel.argumentBuffer(f32, rhs),483        kernel.argumentI32(@intCast(batch)),484    }, .{485        .grid = .{ blocks, 1, 1 },486        .block = .{ instance.threads, 1, 1 },487    });488489    try testing.expectEqualSlices(f32, expected, x);490}491492test "linalg batched cholesky solve matches the host reference at every supported size" {493    try expectBatchedCholeskySolveMatchesHost(2, 32);494    try expectBatchedCholeskySolveMatchesHost(3, 32);495    try expectBatchedCholeskySolveMatchesHost(4, 32);496}497498test "linalg batched cholesky solve identity and validity" {499    const allocator = testing.allocator;500    const instance = BatchedCholeskySolve{ .batch = 5000, .n = 3, .threads = 64 };501    const target = try batchedCholeskySolveFamilyTarget(allocator, instance);502    defer allocator.free(target);503    try testing.expectEqualStrings("accy.kernel.linalg.batched_cholesky_solve_family_3_64_f32", target);504505    try testing.expect(batchedCholeskySolveInstanceValid(instance));506    try testing.expect(!batchedCholeskySolveInstanceValid(.{ .batch = 0, .n = 3, .threads = 64 }));507    try testing.expect(!batchedCholeskySolveInstanceValid(.{ .batch = 100, .n = 5, .threads = 64 }));508    try testing.expect(!batchedCholeskySolveInstanceValid(.{ .batch = std.math.maxInt(u64), .n = 3, .threads = 32 }));509510    const args = try batchedCholeskySolveRuntimeArguments(instance);511    try testing.expectEqual(@as(u32, 5000), args[0].u32);512}513514pub const BatchedInverse = struct {515    batch: u64,516    n: u32 = 3,517    threads: u32 = 256,518    layout: TileLayout = .row_major,519    batch_axis: []const u8 = "b",520};521522pub const batched_inverse_family_version: u32 = 1;523524pub fn batchedInverseInstanceValid(instance: BatchedInverse) bool {525    if (instance.batch == 0) return false;526    if (instance.n < batched_cholesky_min_n or instance.n > batched_cholesky_max_n) return false;527    if (instance.threads == 0 or instance.threads > batched_cholesky_max_threads) return false;528    return extent_mod.blockCountWithinLimit(instance.batch, instance.threads, batched_cholesky_max_blocks);529}530531fn batched_inverse_body_active(inner: anytype, ctx: anytype) !void {532    const n: usize = @intCast(ctx.n);533    const tile = try inner.constantIndex(@intCast(n * n));534    const base = try inner.mul(ctx.system, tile);535536    var a_values: [batched_cholesky_max_n][batched_cholesky_max_n]kernel.Value = undefined;537    var inv_values: [batched_cholesky_max_n][batched_cholesky_max_n]kernel.Value = undefined;538    const zero = try inner.constantFloat(.f32, 0);539    const one = try inner.constantFloat(.f32, 1);540    for (0..n) |i| {541        for (0..n) |j| {542            const index = try factorElementIndex(inner, ctx.layout, base, ctx.batch_index, ctx.system, i * n + j);543            const loaded = try ctx.args.param(.a).load(inner, index);544            a_values[i][j] = loaded.raw();545            inv_values[i][j] = if (i == j) one else zero;546        }547    }548549    for (0..n) |pivot_index| {550        const pivot = a_values[pivot_index][pivot_index];551        for (0..n) |j| {552            a_values[pivot_index][j] = try inner.div(a_values[pivot_index][j], pivot);553            inv_values[pivot_index][j] = try inner.div(inv_values[pivot_index][j], pivot);554        }555        for (0..n) |row| {556            if (row == pivot_index) continue;557            const factor = a_values[row][pivot_index];558            for (0..n) |j| {559                a_values[row][j] = try inner.sub(560                    a_values[row][j],561                    try inner.mul(factor, a_values[pivot_index][j]),562                );563                inv_values[row][j] = try inner.sub(564                    inv_values[row][j],565                    try inner.mul(factor, inv_values[pivot_index][j]),566                );567            }568        }569    }570571    for (0..n) |i| {572        for (0..n) |j| {573            const index = try factorElementIndex(inner, ctx.layout, base, ctx.batch_index, ctx.system, i * n + j);574            try ctx.args.param(.inv).store(inner, inv_values[i][j], index);575        }576    }577}578579fn batchedInverseBody(k: anytype, spec: BatchedInverse, args: anytype) !void {580    if (!batchedInverseInstanceValid(spec)) return error.UnsupportedBatchedInverseInstance;581    const system = try k.globalId(.x);582    const batch = try k.castIndex(args.param(.batch).raw());583    const active = try k.compare(.lt, system, batch);584    try k.guardDo(active, .{585        .args = args,586        .system = system,587        .n = spec.n,588        .layout = spec.layout,589        .batch_index = batch,590    }, batched_inverse_body_active);591}592593fn batchedInverseFamilySchedule(instance: BatchedInverse) kernel.logical.schedule.ThreadBlocks {594    return kernel.logical.schedule.threadBlocks(.{ .x = instance.threads });595}596597fn batchedInverseRuntimeFamily() type {598    return kernel.logical.Family(.{599        .name = "accy_kernel_linalg_batched_inverse_runtime_f32",600        .parameters = .{601            .inv = kernel.dynamicBuffer(.f32),602            .a = kernel.dynamicBuffer(.f32),603            .batch = kernel.scalar(.i32),604        },605        .Instance = BatchedInverse,606        .schedule = batchedInverseFamilySchedule,607        .body = batchedInverseBody,608    });609}610611pub const BatchedInverseRuntimeFamilyF32 = batchedInverseRuntimeFamily();612613pub fn batchedInverseFamilyTarget(allocator: std.mem.Allocator, instance: BatchedInverse) ![]u8 {614    return std.fmt.allocPrint(615        allocator,616        "accy.kernel.linalg.batched_inverse_family_{d}_{d}_{s}f32",617        .{ instance.n, instance.threads, layoutTargetSegment(instance.layout) },618    );619}620621pub fn batchedInverseFamilyEntryName(allocator: std.mem.Allocator, instance: BatchedInverse) ![]u8 {622    return std.fmt.allocPrint(623        allocator,624        "accy_kernel_linalg_batched_inverse_family_{d}_{d}_{s}f32",625        .{ instance.n, instance.threads, layoutTargetSegment(instance.layout) },626    );627}628629pub fn batchedInverseRuntimeArguments(instance: BatchedInverse) ![1]choir_abi.ScalarArgument {630    return .{631        .{ .u32 = try runtimeExtentArgument(instance.batch) },632    };633}634635fn expectBatchedInverseResidual(comptime n: usize, threads: u32) !void {636    const allocator = testing.allocator;637    const batch: usize = 40;638    const instance = BatchedInverse{ .batch = batch, .n = n, .threads = threads };639    const blocks: u32 = @intCast(batchedCholeskyBlockCount(instance.batch, instance.threads));640    try testing.expect(blocks > 1);641642    const a = try allocator.alloc(f32, batch * n * n);643    defer allocator.free(a);644    for (0..batch) |b| {645        fillSpdTile(n, @intCast(b * 977 + 13), a[b * n * n ..][0 .. n * n]);646    }647648    const inv = try allocator.alloc(f32, batch * n * n);649    defer allocator.free(inv);650    @memset(inv, -1);651652    var graph = try BatchedInverseRuntimeFamilyF32.build(allocator, BatchedInverseRuntimeFamilyF32.Limits.testing, instance);653    defer graph.deinit();654    try graph.runCpuWithLaunch(allocator, &.{655        kernel.argumentBuffer(f32, inv),656        kernel.argumentBuffer(f32, a),657        kernel.argumentI32(@intCast(batch)),658    }, .{659        .grid = .{ blocks, 1, 1 },660        .block = .{ instance.threads, 1, 1 },661    });662663    for (0..batch) |b| {664        const a_tile = a[b * n * n ..][0 .. n * n];665        const inv_tile = inv[b * n * n ..][0 .. n * n];666        for (0..n) |i| {667            for (0..n) |j| {668                var product: f32 = 0;669                for (0..n) |c| {670                    product += a_tile[i * n + c] * inv_tile[c * n + j];671                }672                const expected: f32 = if (i == j) 1 else 0;673                try testing.expectApproxEqAbs(expected, product, 0.001);674            }675        }676    }677}678679test "linalg batched inverse residual matches identity at every supported size" {680    try expectBatchedInverseResidual(2, 32);681    try expectBatchedInverseResidual(3, 32);682    try expectBatchedInverseResidual(4, 32);683}684685test "linalg batched inverse interleaved layout matches residual contract" {686    const allocator = testing.allocator;687    const n: usize = 3;688    const batch: usize = 40;689    const instance = BatchedInverse{ .batch = batch, .n = n, .threads = 32, .layout = .interleaved };690    const blocks: u32 = @intCast(batchedCholeskyBlockCount(batch, instance.threads));691692    const a = try allocator.alloc(f32, batch * n * n);693    defer allocator.free(a);694    for (0..batch) |b| {695        fillSpdTile(n, @intCast(b * 977 + 13), a[b * n * n ..][0 .. n * n]);696    }697698    const a_interleaved = try allocator.alloc(f32, batch * n * n);699    defer allocator.free(a_interleaved);700    packInterleavedTiles(n, batch, a, a_interleaved);701702    const inv_interleaved = try allocator.alloc(f32, batch * n * n);703    defer allocator.free(inv_interleaved);704    @memset(inv_interleaved, -1);705706    var graph = try BatchedInverseRuntimeFamilyF32.build(allocator, BatchedInverseRuntimeFamilyF32.Limits.testing, instance);707    defer graph.deinit();708    try graph.runCpuWithLaunch(allocator, &.{709        kernel.argumentBuffer(f32, inv_interleaved),710        kernel.argumentBuffer(f32, a_interleaved),711        kernel.argumentI32(@intCast(batch)),712    }, .{713        .grid = .{ blocks, 1, 1 },714        .block = .{ instance.threads, 1, 1 },715    });716717    for (0..batch) |b| {718        const a_tile = a[b * n * n ..][0 .. n * n];719        for (0..n) |i| {720            for (0..n) |j| {721                var product: f32 = 0;722                for (0..n) |c| {723                    product += a_tile[i * n + c] * inv_interleaved[(c * n + j) * batch + b];724                }725                const expected: f32 = if (i == j) 1 else 0;726                try testing.expectApproxEqAbs(expected, product, 0.001);727            }728        }729    }730}731732test "linalg batched inverse identity and validity" {733    const allocator = testing.allocator;734    const instance = BatchedInverse{ .batch = 5000, .n = 3, .threads = 64 };735    const target = try batchedInverseFamilyTarget(allocator, instance);736    defer allocator.free(target);737    try testing.expectEqualStrings("accy.kernel.linalg.batched_inverse_family_3_64_f32", target);738    const entry_name = try batchedInverseFamilyEntryName(allocator, instance);739    defer allocator.free(entry_name);740    try testing.expectEqualStrings("accy_kernel_linalg_batched_inverse_family_3_64_f32", entry_name);741742    try testing.expect(batchedInverseInstanceValid(instance));743    try testing.expect(!batchedInverseInstanceValid(.{ .batch = 0, .n = 3, .threads = 64 }));744    try testing.expect(!batchedInverseInstanceValid(.{ .batch = 100, .n = 1, .threads = 64 }));745    try testing.expect(!batchedInverseInstanceValid(.{ .batch = 100, .n = 5, .threads = 64 }));746    try testing.expect(!batchedInverseInstanceValid(.{ .batch = 100, .n = 3, .threads = 0 }));747    try testing.expect(!batchedInverseInstanceValid(.{ .batch = std.math.maxInt(u64), .n = 3, .threads = 32 }));748749    const args = try batchedInverseRuntimeArguments(instance);750    try testing.expectEqual(@as(u32, 5000), args[0].u32);751}752753pub const batched_cholesky_row_axis = "i";754pub const batched_cholesky_col_axis = "j";755756fn factorBatchBounds() shape.Bounds {757    return .{ .min = 1, .max = extent_mod.runtime_extent_max };758}759760fn factorTileBounds() shape.Bounds {761    return .{ .min = batched_cholesky_min_n, .max = batched_cholesky_max_n };762}763764pub fn batchedCholeskyShapeFamily(backing_allocator: std.mem.Allocator, instance: BatchedCholesky) !shape.Family {765    var builder = try shape.Builder.init(backing_allocator, "batched_cholesky");766    errdefer builder.deinit();767    const batch_symbol = try builder.symbol(instance.batch_axis);768    const row = try builder.symbol(batched_cholesky_row_axis);769    const col = try builder.symbol(batched_cholesky_col_axis);770    const batch_expr = try builder.symbolExpression(batch_symbol);771    const row_expr = try builder.symbolExpression(row);772    const col_expr = try builder.symbolExpression(col);773    _ = try builder.tensor("l", &.{ batch_expr, row_expr, col_expr });774    _ = try builder.tensor("a", &.{ batch_expr, row_expr, col_expr });775    try builder.assumeBounds(batch_expr, factorBatchBounds());776    try builder.assumeBounds(row_expr, factorTileBounds());777    try builder.assumeBounds(col_expr, factorTileBounds());778    return builder.finish();779}780781pub fn batchedCholeskyFamilyFingerprint(backing_allocator: std.mem.Allocator, instance: BatchedCholesky) !u64 {782    var family = try batchedCholeskyShapeFamily(backing_allocator, instance);783    defer family.deinit();784    return shape.fingerprint(family);785}786787pub fn batchedCholeskyFamilySpecialization(788    backing_allocator: std.mem.Allocator,789    instance: BatchedCholesky,790) !entry.OwnedSpecialization {791    var owned = entry.OwnedSpecialization.init(backing_allocator);792    errdefer owned.deinit();793    const lifetime_allocator = owned.allocator();794795    const inputs = try lifetime_allocator.alloc(entry.Shape, 1);796    inputs[0] = try entry.runtimeShape3D(797        lifetime_allocator,798        instance.batch_axis,799        instance.batch,800        batched_cholesky_row_axis,801        instance.n,802        batched_cholesky_col_axis,803        instance.n,804    );805806    const outputs = try lifetime_allocator.alloc(entry.Shape, 1);807    outputs[0] = try entry.runtimeShape3D(808        lifetime_allocator,809        instance.batch_axis,810        instance.batch,811        batched_cholesky_row_axis,812        instance.n,813        batched_cholesky_col_axis,814        instance.n,815    );816817    owned.value = .{818        .dtype = .f32,819        .operation = .{ .linalg = .batched_cholesky },820        .inputs = inputs,821        .outputs = outputs,822        .schedule = try entry.runtimeThreadBlocks1D(823            lifetime_allocator,824            instance.batch_axis,825            instance.batch,826            instance.threads,827        ),828        .layout = @tagName(instance.layout),829    };830    owned.value.launch = owned.value.schedule.?.launch();831    var family = try batchedCholeskyShapeFamily(backing_allocator, instance);832    errdefer family.deinit();833    try owned.takeShapeFamily(&family);834    return owned;835}836837pub fn batchedCholeskyInstanceFromSpecialization(specialization: entry.Specialization) ?BatchedCholesky {838    if (!specialization.scheduleMatchesLaunch()) return null;839    if (!specialization.operationIs(.{ .linalg = .batched_cholesky })) return null;840    const dtype = specialization.dtype orelse return null;841    if (dtype != .f32) return null;842    if (specialization.inputs.len != 1 or specialization.outputs.len != 1) return null;843    if (specialization.reductions.len != 0) return null;844    const tiles = specialization.outputs[0];845    if (tiles.axes.len != 3) return null;846    const batch_axis = tiles.axes[0].name;847    const batch = tiles.axes[0].extent;848    const tile_extent = tiles.axes[1].extent;849    if (!factorTileShapeMatches(tiles, batch_axis, batch, tile_extent)) return null;850    if (!factorTileShapeMatches(specialization.inputs[0], batch_axis, batch, tile_extent)) return null;851    const launch = specialization.launch orelse return null;852    if (launch.threadgroup[0] == 0) return null;853    const n = std.math.cast(u32, tile_extent) orelse return null;854    const layout_name = specialization.layout orelse return null;855    const layout = std.meta.stringToEnum(TileLayout, layout_name) orelse return null;856    const instance = BatchedCholesky{857        .batch = batch,858        .n = n,859        .threads = launch.threadgroup[0],860        .layout = layout,861        .batch_axis = batch_axis,862    };863    if (!batchedCholeskyInstanceValid(instance)) return null;864    return instance;865}866867fn factorAxisMatches(axis: entry.Axis, name: []const u8, extent: u64) bool {868    if (name.len == 0) return false;869    if (!std.mem.eql(u8, axis.name, name)) return false;870    return axis.extent == extent;871}872873fn factorTileShapeMatches(candidate: entry.Shape, batch_axis: []const u8, batch: u64, n: u64) bool {874    if (candidate.axes.len != 3) return false;875    return factorAxisMatches(candidate.axes[0], batch_axis, batch) and876        factorAxisMatches(candidate.axes[1], batched_cholesky_row_axis, n) and877        factorAxisMatches(candidate.axes[2], batched_cholesky_col_axis, n);878}879880fn factorVectorShapeMatches(candidate: entry.Shape, batch_axis: []const u8, batch: u64, n: u64) bool {881    if (candidate.axes.len != 2) return false;882    return factorAxisMatches(candidate.axes[0], batch_axis, batch) and883        factorAxisMatches(candidate.axes[1], batched_cholesky_row_axis, n);884}885886pub fn batchedCholeskyShapeProfileDimensions(instance: BatchedCholesky) [1]artifact_product.KernelCallShapeProfileDimension {887    return .{888        .{ .name = instance.batch_axis, .runtime_scalar_argument_index = 0, .bounds = factorBatchBounds() },889    };890}891892fn factorDerivedLaunch(threads: u32) !artifact_product.KernelCallLaunch {893    if (threads == 0) return error.KernelLibraryLaunchThreadgroupMustBeNonzero;894    return .{ .derived = .{895        .grid = .{896            .{ .runtime_u32_ceil_div = .{ .argument_index = 0, .divisor = threads } },897            .{ .fixed = 1 },898            .{ .fixed = 1 },899        },900        .threadgroup = .{ threads, 1, 1 },901    } };902}903904pub fn createBatchedCholeskyFamilyArtifact(905    allocator: std.mem.Allocator,906    handle: kernel.BackendHandle,907    instance: BatchedCholesky,908    options: entry.ArtifactOptions,909) !kernel.OwnedKernelCallArtifact {910    if (!batchedCholeskyInstanceValid(instance)) return error.InvalidKernelLibraryEntry;911    const target = try batchedCholeskyFamilyTarget(allocator, instance);912    defer allocator.free(target);913    const entry_name = try batchedCholeskyFamilyEntryName(allocator, instance);914    defer allocator.free(entry_name);915    const family_fingerprint = options.shape_family_fingerprint orelse try batchedCholeskyFamilyFingerprint(allocator, instance);916    const shape_profile_dimensions = batchedCholeskyShapeProfileDimensions(instance);917    const shape_profile = options.shape_profile orelse artifact_product.KernelCallShapeProfile{918        .name = "batched_cholesky",919        .fingerprint = family_fingerprint,920        .dimensions = shape_profile_dimensions[0..],921    };922923    var graph = try BatchedCholeskyRuntimeFamilyF32.buildNamed(allocator, options.limits, entry_name, instance);924    defer graph.deinit();925    return kernel.createKernelCallArtifact(allocator, handle, &graph, .{926        .target = target,927        .version = batched_cholesky_family_version,928        .format = options.format,929        .kernel_plan = options.kernel_plan,930        .element_count_argument = options.element_count_argument,931        .shape_family_fingerprint = family_fingerprint,932        .shape_profile = shape_profile,933        .launch = options.launch orelse try factorDerivedLaunch(instance.threads),934        .runtime_scalar_argument_count = if (options.runtime_scalar_argument_count == 0) 1 else options.runtime_scalar_argument_count,935        .static_arguments = options.static_arguments,936    });937}938939pub fn batchedCholeskySolveShapeFamily(backing_allocator: std.mem.Allocator, instance: BatchedCholeskySolve) !shape.Family {940    var builder = try shape.Builder.init(backing_allocator, "batched_cholesky_solve");941    errdefer builder.deinit();942    const batch_symbol = try builder.symbol(instance.batch_axis);943    const row = try builder.symbol(batched_cholesky_row_axis);944    const col = try builder.symbol(batched_cholesky_col_axis);945    const batch_expr = try builder.symbolExpression(batch_symbol);946    const row_expr = try builder.symbolExpression(row);947    const col_expr = try builder.symbolExpression(col);948    _ = try builder.tensor("x", &.{ batch_expr, row_expr });949    _ = try builder.tensor("l", &.{ batch_expr, row_expr, col_expr });950    _ = try builder.tensor("rhs", &.{ batch_expr, row_expr });951    try builder.assumeBounds(batch_expr, factorBatchBounds());952    try builder.assumeBounds(row_expr, factorTileBounds());953    try builder.assumeBounds(col_expr, factorTileBounds());954    return builder.finish();955}956957pub fn batchedCholeskySolveFamilyFingerprint(backing_allocator: std.mem.Allocator, instance: BatchedCholeskySolve) !u64 {958    var family = try batchedCholeskySolveShapeFamily(backing_allocator, instance);959    defer family.deinit();960    return shape.fingerprint(family);961}962963pub fn batchedCholeskySolveFamilySpecialization(964    backing_allocator: std.mem.Allocator,965    instance: BatchedCholeskySolve,966) !entry.OwnedSpecialization {967    var owned = entry.OwnedSpecialization.init(backing_allocator);968    errdefer owned.deinit();969    const lifetime_allocator = owned.allocator();970971    const inputs = try lifetime_allocator.alloc(entry.Shape, 2);972    inputs[0] = try entry.runtimeShape3D(973        lifetime_allocator,974        instance.batch_axis,975        instance.batch,976        batched_cholesky_row_axis,977        instance.n,978        batched_cholesky_col_axis,979        instance.n,980    );981    inputs[1] = try entry.runtimeShape2D(982        lifetime_allocator,983        instance.batch_axis,984        instance.batch,985        batched_cholesky_row_axis,986        instance.n,987    );988989    const outputs = try lifetime_allocator.alloc(entry.Shape, 1);990    outputs[0] = try entry.runtimeShape2D(991        lifetime_allocator,992        instance.batch_axis,993        instance.batch,994        batched_cholesky_row_axis,995        instance.n,996    );997998    owned.value = .{999        .dtype = .f32,1000        .operation = .{ .linalg = .batched_cholesky_solve },1001        .inputs = inputs,1002        .outputs = outputs,1003        .schedule = try entry.runtimeThreadBlocks1D(1004            lifetime_allocator,1005            instance.batch_axis,1006            instance.batch,1007            instance.threads,1008        ),1009        .layout = @tagName(instance.layout),1010    };1011    owned.value.launch = owned.value.schedule.?.launch();1012    var family = try batchedCholeskySolveShapeFamily(backing_allocator, instance);1013    errdefer family.deinit();1014    try owned.takeShapeFamily(&family);1015    return owned;1016}10171018pub fn batchedCholeskySolveInstanceFromSpecialization(specialization: entry.Specialization) ?BatchedCholeskySolve {1019    if (!specialization.scheduleMatchesLaunch()) return null;1020    if (!specialization.operationIs(.{ .linalg = .batched_cholesky_solve })) return null;1021    const dtype = specialization.dtype orelse return null;1022    if (dtype != .f32) return null;1023    if (specialization.inputs.len != 2 or specialization.outputs.len != 1) return null;1024    if (specialization.reductions.len != 0) return null;1025    const tiles = specialization.inputs[0];1026    const rhs = specialization.inputs[1];1027    const solution = specialization.outputs[0];1028    if (solution.axes.len != 2) return null;1029    const batch_axis = solution.axes[0].name;1030    const batch = solution.axes[0].extent;1031    const tile_extent = solution.axes[1].extent;1032    if (!factorVectorShapeMatches(solution, batch_axis, batch, tile_extent)) return null;1033    if (!factorTileShapeMatches(tiles, batch_axis, batch, tile_extent)) return null;1034    if (!factorVectorShapeMatches(rhs, batch_axis, batch, tile_extent)) return null;1035    const launch = specialization.launch orelse return null;1036    if (launch.threadgroup[0] == 0) return null;1037    const n = std.math.cast(u32, tile_extent) orelse return null;1038    const layout_name = specialization.layout orelse return null;1039    const layout = std.meta.stringToEnum(TileLayout, layout_name) orelse return null;1040    const instance = BatchedCholeskySolve{1041        .batch = batch,1042        .n = n,1043        .threads = launch.threadgroup[0],1044        .layout = layout,1045        .batch_axis = batch_axis,1046    };1047    if (!batchedCholeskySolveInstanceValid(instance)) return null;1048    return instance;1049}10501051pub fn batchedCholeskySolveShapeProfileDimensions(instance: BatchedCholeskySolve) [1]artifact_product.KernelCallShapeProfileDimension {1052    return .{1053        .{ .name = instance.batch_axis, .runtime_scalar_argument_index = 0, .bounds = factorBatchBounds() },1054    };1055}10561057pub fn createBatchedCholeskySolveFamilyArtifact(1058    allocator: std.mem.Allocator,1059    handle: kernel.BackendHandle,1060    instance: BatchedCholeskySolve,1061    options: entry.ArtifactOptions,1062) !kernel.OwnedKernelCallArtifact {1063    if (!batchedCholeskySolveInstanceValid(instance)) return error.InvalidKernelLibraryEntry;1064    const target = try batchedCholeskySolveFamilyTarget(allocator, instance);1065    defer allocator.free(target);1066    const entry_name = try batchedCholeskySolveFamilyEntryName(allocator, instance);1067    defer allocator.free(entry_name);1068    const family_fingerprint = options.shape_family_fingerprint orelse try batchedCholeskySolveFamilyFingerprint(allocator, instance);1069    const shape_profile_dimensions = batchedCholeskySolveShapeProfileDimensions(instance);1070    const shape_profile = options.shape_profile orelse artifact_product.KernelCallShapeProfile{1071        .name = "batched_cholesky_solve",1072        .fingerprint = family_fingerprint,1073        .dimensions = shape_profile_dimensions[0..],1074    };10751076    var graph = try BatchedCholeskySolveRuntimeFamilyF32.buildNamed(allocator, options.limits, entry_name, instance);1077    defer graph.deinit();1078    return kernel.createKernelCallArtifact(allocator, handle, &graph, .{1079        .target = target,1080        .version = batched_cholesky_solve_family_version,1081        .format = options.format,1082        .kernel_plan = options.kernel_plan,1083        .element_count_argument = options.element_count_argument,1084        .shape_family_fingerprint = family_fingerprint,1085        .shape_profile = shape_profile,1086        .launch = options.launch orelse try factorDerivedLaunch(instance.threads),1087        .runtime_scalar_argument_count = if (options.runtime_scalar_argument_count == 0) 1 else options.runtime_scalar_argument_count,1088        .static_arguments = options.static_arguments,1089    });1090}10911092pub fn batchedInverseShapeFamily(backing_allocator: std.mem.Allocator, instance: BatchedInverse) !shape.Family {1093    var builder = try shape.Builder.init(backing_allocator, "batched_inverse");1094    errdefer builder.deinit();1095    const batch_symbol = try builder.symbol(instance.batch_axis);1096    const row = try builder.symbol(batched_cholesky_row_axis);1097    const col = try builder.symbol(batched_cholesky_col_axis);1098    const batch_expr = try builder.symbolExpression(batch_symbol);1099    const row_expr = try builder.symbolExpression(row);1100    const col_expr = try builder.symbolExpression(col);1101    _ = try builder.tensor("inv", &.{ batch_expr, row_expr, col_expr });1102    _ = try builder.tensor("a", &.{ batch_expr, row_expr, col_expr });1103    try builder.assumeBounds(batch_expr, factorBatchBounds());1104    try builder.assumeBounds(row_expr, factorTileBounds());1105    try builder.assumeBounds(col_expr, factorTileBounds());1106    return builder.finish();1107}11081109pub fn batchedInverseFamilyFingerprint(backing_allocator: std.mem.Allocator, instance: BatchedInverse) !u64 {1110    var family = try batchedInverseShapeFamily(backing_allocator, instance);1111    defer family.deinit();1112    return shape.fingerprint(family);1113}11141115pub fn batchedInverseFamilySpecialization(1116    backing_allocator: std.mem.Allocator,1117    instance: BatchedInverse,1118) !entry.OwnedSpecialization {1119    var owned = entry.OwnedSpecialization.init(backing_allocator);1120    errdefer owned.deinit();1121    const lifetime_allocator = owned.allocator();11221123    const inputs = try lifetime_allocator.alloc(entry.Shape, 1);1124    inputs[0] = try entry.runtimeShape3D(1125        lifetime_allocator,1126        instance.batch_axis,1127        instance.batch,1128        batched_cholesky_row_axis,1129        instance.n,1130        batched_cholesky_col_axis,1131        instance.n,1132    );11331134    const outputs = try lifetime_allocator.alloc(entry.Shape, 1);1135    outputs[0] = try entry.runtimeShape3D(1136        lifetime_allocator,1137        instance.batch_axis,1138        instance.batch,1139        batched_cholesky_row_axis,1140        instance.n,1141        batched_cholesky_col_axis,1142        instance.n,1143    );11441145    owned.value = .{1146        .dtype = .f32,1147        .operation = .{ .linalg = .batched_inverse },1148        .inputs = inputs,1149        .outputs = outputs,1150        .schedule = try entry.runtimeThreadBlocks1D(1151            lifetime_allocator,1152            instance.batch_axis,1153            instance.batch,1154            instance.threads,1155        ),1156        .layout = @tagName(instance.layout),1157    };1158    owned.value.launch = owned.value.schedule.?.launch();1159    var family = try batchedInverseShapeFamily(backing_allocator, instance);1160    errdefer family.deinit();1161    try owned.takeShapeFamily(&family);1162    return owned;1163}11641165pub fn batchedInverseInstanceFromSpecialization(specialization: entry.Specialization) ?BatchedInverse {1166    if (!specialization.scheduleMatchesLaunch()) return null;1167    if (!specialization.operationIs(.{ .linalg = .batched_inverse })) return null;1168    const dtype = specialization.dtype orelse return null;1169    if (dtype != .f32) return null;1170    if (specialization.inputs.len != 1 or specialization.outputs.len != 1) return null;1171    if (specialization.reductions.len != 0) return null;1172    const tiles = specialization.outputs[0];1173    if (tiles.axes.len != 3) return null;1174    const batch_axis = tiles.axes[0].name;1175    const batch = tiles.axes[0].extent;1176    const tile_extent = tiles.axes[1].extent;1177    if (!factorTileShapeMatches(tiles, batch_axis, batch, tile_extent)) return null;1178    if (!factorTileShapeMatches(specialization.inputs[0], batch_axis, batch, tile_extent)) return null;1179    const launch = specialization.launch orelse return null;1180    if (launch.threadgroup[0] == 0) return null;1181    const n = std.math.cast(u32, tile_extent) orelse return null;1182    const layout_name = specialization.layout orelse return null;1183    const layout = std.meta.stringToEnum(TileLayout, layout_name) orelse return null;1184    const instance = BatchedInverse{1185        .batch = batch,1186        .n = n,1187        .threads = launch.threadgroup[0],1188        .layout = layout,1189        .batch_axis = batch_axis,1190    };1191    if (!batchedInverseInstanceValid(instance)) return null;1192    return instance;1193}11941195pub fn batchedInverseShapeProfileDimensions(instance: BatchedInverse) [1]artifact_product.KernelCallShapeProfileDimension {1196    return .{1197        .{ .name = instance.batch_axis, .runtime_scalar_argument_index = 0, .bounds = factorBatchBounds() },1198    };1199}12001201pub fn createBatchedInverseFamilyArtifact(1202    allocator: std.mem.Allocator,1203    handle: kernel.BackendHandle,1204    instance: BatchedInverse,1205    options: entry.ArtifactOptions,1206) !kernel.OwnedKernelCallArtifact {1207    if (!batchedInverseInstanceValid(instance)) return error.InvalidKernelLibraryEntry;1208    const target = try batchedInverseFamilyTarget(allocator, instance);1209    defer allocator.free(target);1210    const entry_name = try batchedInverseFamilyEntryName(allocator, instance);1211    defer allocator.free(entry_name);1212    const family_fingerprint = options.shape_family_fingerprint orelse try batchedInverseFamilyFingerprint(allocator, instance);1213    const shape_profile_dimensions = batchedInverseShapeProfileDimensions(instance);1214    const shape_profile = options.shape_profile orelse artifact_product.KernelCallShapeProfile{1215        .name = "batched_inverse",1216        .fingerprint = family_fingerprint,1217        .dimensions = shape_profile_dimensions[0..],1218    };12191220    var graph = try BatchedInverseRuntimeFamilyF32.buildNamed(allocator, options.limits, entry_name, instance);1221    defer graph.deinit();1222    return kernel.createKernelCallArtifact(allocator, handle, &graph, .{1223        .target = target,1224        .version = batched_inverse_family_version,1225        .format = options.format,1226        .kernel_plan = options.kernel_plan,1227        .element_count_argument = options.element_count_argument,1228        .shape_family_fingerprint = family_fingerprint,1229        .shape_profile = shape_profile,1230        .launch = options.launch orelse try factorDerivedLaunch(instance.threads),1231        .runtime_scalar_argument_count = if (options.runtime_scalar_argument_count == 0) 1 else options.runtime_scalar_argument_count,1232        .static_arguments = options.static_arguments,1233    });1234}12351236test "linalg batched factor specializations round-trip their instances" {1237    const allocator = testing.allocator;12381239    var cholesky_owned = try batchedCholeskyFamilySpecialization(allocator, .{ .batch = 5000, .n = 3, .threads = 64 });1240    defer cholesky_owned.deinit();1241    try testing.expect(cholesky_owned.value.operationIs(.{ .linalg = .batched_cholesky }));1242    const cholesky_recovered = batchedCholeskyInstanceFromSpecialization(cholesky_owned.value) orelse {1243        return error.TestExpectedBatchedCholeskyInstance;1244    };1245    try testing.expectEqual(@as(u64, 5000), cholesky_recovered.batch);1246    try testing.expectEqual(@as(u32, 3), cholesky_recovered.n);1247    try testing.expectEqual(@as(u32, 64), cholesky_recovered.threads);1248    try testing.expectEqual(@as(?BatchedCholesky, null), batchedCholeskyInstanceFromSpecialization(.{}));12491250    var solve_owned = try batchedCholeskySolveFamilySpecialization(allocator, .{ .batch = 5000, .n = 4, .threads = 32 });1251    defer solve_owned.deinit();1252    const solve_recovered = batchedCholeskySolveInstanceFromSpecialization(solve_owned.value) orelse {1253        return error.TestExpectedBatchedCholeskySolveInstance;1254    };1255    try testing.expectEqual(@as(u64, 5000), solve_recovered.batch);1256    try testing.expectEqual(@as(u32, 4), solve_recovered.n);1257    try testing.expectEqual(@as(u32, 32), solve_recovered.threads);1258    try testing.expectEqual(1259        @as(?BatchedCholeskySolve, null),1260        batchedCholeskySolveInstanceFromSpecialization(cholesky_owned.value),1261    );1262    try testing.expectEqual(1263        @as(?BatchedCholesky, null),1264        batchedCholeskyInstanceFromSpecialization(solve_owned.value),1265    );12661267    var inverse_owned = try batchedInverseFamilySpecialization(allocator, .{ .batch = 4096, .n = 2, .threads = 128 });1268    defer inverse_owned.deinit();1269    const inverse_recovered = batchedInverseInstanceFromSpecialization(inverse_owned.value) orelse {1270        return error.TestExpectedBatchedInverseInstance;1271    };1272    try testing.expectEqual(@as(u64, 4096), inverse_recovered.batch);1273    try testing.expectEqual(@as(u32, 2), inverse_recovered.n);1274    try testing.expectEqual(@as(u32, 128), inverse_recovered.threads);1275    try testing.expectEqual(1276        @as(?BatchedInverse, null),1277        batchedInverseInstanceFromSpecialization(cholesky_owned.value),1278    );1279}12801281test "linalg batched factor specializations reject malformed tensor facts" {1282    const allocator = testing.allocator;12831284    {1285        var owned = try batchedCholeskyFamilySpecialization(allocator, .{ .batch = 5000, .n = 3, .threads = 64 });1286        defer owned.deinit();1287        const lifetime_allocator = owned.allocator();1288        owned.value.inputs = &.{try entry.runtimeShape3D(lifetime_allocator, "b", 4096, "i", 3, "j", 3)};1289        try testing.expectEqual(@as(?BatchedCholesky, null), batchedCholeskyInstanceFromSpecialization(owned.value));1290    }12911292    {1293        var owned = try batchedCholeskyFamilySpecialization(allocator, .{ .batch = 5000, .n = 3, .threads = 64 });1294        defer owned.deinit();1295        const lifetime_allocator = owned.allocator();1296        owned.value.outputs = &.{try entry.runtimeShape3D(lifetime_allocator, "b", 5000, "row", 3, "j", 3)};1297        try testing.expectEqual(@as(?BatchedCholesky, null), batchedCholeskyInstanceFromSpecialization(owned.value));1298    }12991300    {1301        var owned = try batchedCholeskySolveFamilySpecialization(allocator, .{ .batch = 5000, .n = 3, .threads = 64 });1302        defer owned.deinit();1303        const lifetime_allocator = owned.allocator();1304        const inputs = try lifetime_allocator.alloc(entry.Shape, 2);1305        inputs[0] = owned.value.inputs[0];1306        inputs[1] = try entry.runtimeShape2D(lifetime_allocator, "b", 4096, "i", 3);1307        owned.value.inputs = inputs;1308        try testing.expectEqual(@as(?BatchedCholeskySolve, null), batchedCholeskySolveInstanceFromSpecialization(owned.value));1309    }13101311    {1312        var owned = try batchedCholeskySolveFamilySpecialization(allocator, .{ .batch = 5000, .n = 3, .threads = 64 });1313        defer owned.deinit();1314        const lifetime_allocator = owned.allocator();1315        const inputs = try lifetime_allocator.alloc(entry.Shape, 2);1316        inputs[0] = owned.value.inputs[0];1317        inputs[1] = try entry.runtimeShape2D(lifetime_allocator, "b", 5000, "i", 4);1318        owned.value.inputs = inputs;1319        try testing.expectEqual(@as(?BatchedCholeskySolve, null), batchedCholeskySolveInstanceFromSpecialization(owned.value));1320    }13211322    {1323        var owned = try batchedInverseFamilySpecialization(allocator, .{ .batch = 5000, .n = 3, .threads = 64 });1324        defer owned.deinit();1325        const lifetime_allocator = owned.allocator();1326        owned.value.outputs = &.{try entry.runtimeShape3D(lifetime_allocator, "b", 5000, "i", 3, "k", 3)};1327        try testing.expectEqual(@as(?BatchedInverse, null), batchedInverseInstanceFromSpecialization(owned.value));1328    }1329}13301331/// A caller uses this to rearrange a batch of square matrices on the host before a factorization1332/// kernel reads them: the function copies `batch` matrices of `n` by `n` values, stored one matrix1333/// after another in row order, into an interleaved layout, where the same element of every matrix1334/// sits side by side. Element `slot` of matrix `b` lands at index `slot * batch + b`. The caller1335/// sizes both slices to hold `n * n * batch` values, and the function neither checks those lengths1336/// nor allocates.1337pub fn packInterleavedTiles(n: usize, batch: usize, row_major: []const f32, interleaved: []f32) void {1338    for (0..batch) |b| {1339        for (0..n * n) |slot| {1340            interleaved[slot * batch + b] = row_major[b * n * n + slot];1341        }1342    }1343}13441345/// A caller uses this to rearrange a batch of vectors, such as right-hand sides for a solve: the1346/// function copies `batch` vectors of `n` values, stored one vector after another, into an1347/// interleaved layout, where the same element of every vector sits side by side. Element `slot` of1348/// vector `b` lands at index `slot * batch + b`. The caller sizes both slices to hold `n * batch`1349/// values, and the function neither checks those lengths nor allocates.1350pub fn packInterleavedVectors(n: usize, batch: usize, row_major: []const f32, interleaved: []f32) void {1351    for (0..batch) |b| {1352        for (0..n) |slot| {1353            interleaved[slot * batch + b] = row_major[b * n + slot];1354        }1355    }1356}13571358test "linalg batched cholesky interleaved layout matches the row-major results" {1359    const allocator = testing.allocator;1360    const n: usize = 3;1361    const batch: usize = 40;1362    const instance = BatchedCholesky{ .batch = batch, .n = n, .threads = 32, .layout = .interleaved };1363    const blocks: u32 = @intCast(batchedCholeskyBlockCount(batch, instance.threads));13641365    const a = try allocator.alloc(f32, batch * n * n);1366    defer allocator.free(a);1367    for (0..batch) |b| {1368        fillSpdTile(n, @intCast(b * 977 + 13), a[b * n * n ..][0 .. n * n]);1369    }1370    const expected = try allocator.alloc(f32, batch * n * n);1371    defer allocator.free(expected);1372    for (0..batch) |b| {1373        hostCholesky(n, a[b * n * n ..][0 .. n * n], expected[b * n * n ..][0 .. n * n]);1374    }13751376    const a_interleaved = try allocator.alloc(f32, batch * n * n);1377    defer allocator.free(a_interleaved);1378    packInterleavedTiles(n, batch, a, a_interleaved);1379    const expected_interleaved = try allocator.alloc(f32, batch * n * n);1380    defer allocator.free(expected_interleaved);1381    packInterleavedTiles(n, batch, expected, expected_interleaved);13821383    const l = try allocator.alloc(f32, batch * n * n);1384    defer allocator.free(l);1385    @memset(l, -1);13861387    var graph = try BatchedCholeskyRuntimeFamilyF32.build(allocator, BatchedCholeskyRuntimeFamilyF32.Limits.testing, instance);1388    defer graph.deinit();1389    try graph.runCpuWithLaunch(allocator, &.{1390        kernel.argumentBuffer(f32, l),1391        kernel.argumentBuffer(f32, a_interleaved),1392        kernel.argumentI32(@intCast(batch)),1393    }, .{1394        .grid = .{ blocks, 1, 1 },1395        .block = .{ instance.threads, 1, 1 },1396    });13971398    try testing.expectEqualSlices(f32, expected_interleaved, l);1399}14001401test "linalg batched cholesky solve interleaved layout matches the row-major results" {1402    const allocator = testing.allocator;1403    const n: usize = 3;1404    const batch: usize = 40;1405    const instance = BatchedCholeskySolve{ .batch = batch, .n = n, .threads = 32, .layout = .interleaved };1406    const blocks: u32 = @intCast(batchedCholeskyBlockCount(batch, instance.threads));14071408    const a = try allocator.alloc(f32, batch * n * n);1409    defer allocator.free(a);1410    const l_row = try allocator.alloc(f32, batch * n * n);1411    defer allocator.free(l_row);1412    const rhs_row = try allocator.alloc(f32, batch * n);1413    defer allocator.free(rhs_row);1414    var seed: u32 = 0x1a40a75d;1415    for (0..batch) |b| {1416        fillSpdTile(n, @intCast(b * 977 + 13), a[b * n * n ..][0 .. n * n]);1417        hostCholesky(n, a[b * n * n ..][0 .. n * n], l_row[b * n * n ..][0 .. n * n]);1418        for (rhs_row[b * n ..][0..n]) |*value| {1419            seed ^= seed << 13;1420            seed ^= seed >> 17;1421            seed ^= seed << 5;1422            value.* = @as(f32, @floatFromInt(seed % 1000)) / 500.0 - 1.0;1423        }1424    }1425    const expected_row = try allocator.alloc(f32, batch * n);1426    defer allocator.free(expected_row);1427    for (0..batch) |b| {1428        hostCholeskySolve(n, l_row[b * n * n ..][0 .. n * n], rhs_row[b * n ..][0..n], expected_row[b * n ..][0..n]);1429    }14301431    const l_interleaved = try allocator.alloc(f32, batch * n * n);1432    defer allocator.free(l_interleaved);1433    packInterleavedTiles(n, batch, l_row, l_interleaved);1434    const rhs_interleaved = try allocator.alloc(f32, batch * n);1435    defer allocator.free(rhs_interleaved);1436    packInterleavedVectors(n, batch, rhs_row, rhs_interleaved);1437    const expected_interleaved = try allocator.alloc(f32, batch * n);1438    defer allocator.free(expected_interleaved);1439    packInterleavedVectors(n, batch, expected_row, expected_interleaved);14401441    const x = try allocator.alloc(f32, batch * n);1442    defer allocator.free(x);1443    @memset(x, -1);14441445    var graph = try BatchedCholeskySolveRuntimeFamilyF32.build(allocator, BatchedCholeskySolveRuntimeFamilyF32.Limits.testing, instance);1446    defer graph.deinit();1447    try graph.runCpuWithLaunch(allocator, &.{1448        kernel.argumentBuffer(f32, x),1449        kernel.argumentBuffer(f32, l_interleaved),1450        kernel.argumentBuffer(f32, rhs_interleaved),1451        kernel.argumentI32(@intCast(batch)),1452    }, .{1453        .grid = .{ blocks, 1, 1 },1454        .block = .{ instance.threads, 1, 1 },1455    });14561457    try testing.expectEqualSlices(f32, expected_interleaved, x);1458}14591460test "linalg batched factor layout rides identity and round-trips" {1461    const allocator = testing.allocator;1462    const instance = BatchedCholesky{ .batch = 5000, .n = 3, .threads = 64, .layout = .interleaved };1463    const target = try batchedCholeskyFamilyTarget(allocator, instance);1464    defer allocator.free(target);1465    try testing.expectEqualStrings("accy.kernel.linalg.batched_cholesky_family_3_64_il_f32", target);14661467    var owned = try batchedCholeskyFamilySpecialization(allocator, instance);1468    defer owned.deinit();1469    try testing.expect(owned.value.layoutIs("interleaved"));1470    const recovered = batchedCholeskyInstanceFromSpecialization(owned.value) orelse {1471        return error.TestExpectedBatchedCholeskyInstance;1472    };1473    try testing.expectEqual(TileLayout.interleaved, recovered.layout);1474}