lib/accy/src/kernel/library/factor.zig
daab053ee43316e1809a84551d573ddd1e5bf3d2
1 const std = @import("std");
2 const choir_abi = @import("choir_abi");
3
4 const artifact_product = @import("../../artifact/model/root.zig");
5 const shape = @import("../../choir/shape/root.zig");
6 const entry = @import("entry.zig");
7 const extent_mod = @import("extent.zig");
8 const kernel = @import("../root.zig");
9
10 const DType = choir_abi.DType;
11 const runtimeExtentArgument = extent_mod.runtimeExtentArgument;
12
13 pub const TileLayout = enum {
14 row_major,
15 interleaved,
16 };
17
18 pub 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 };
25
26 pub const batched_cholesky_family_version: u32 = 1;
27 pub const batched_cholesky_min_n: u32 = 2;
28 pub const batched_cholesky_max_n: u32 = 4;
29 pub const batched_cholesky_max_threads: u32 = 1024;
30 pub const batched_cholesky_max_blocks: u32 = 1024;
31
32 pub fn batchedCholeskyBlockCount(batch: u64, threads: u32) u64 {
33 return (batch + threads - 1) / threads;
34 }
35
36 pub 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 }
42
43 fn 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 }
59
60 fn 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);
64
65 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 }
73
74 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 }
89
90 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 }
99
100 fn 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 }
113
114 fn batchedCholeskyFamilySchedule(instance: BatchedCholesky) kernel.logical.schedule.ThreadBlocks {
115 return kernel.logical.schedule.threadBlocks(.{ .x = instance.threads });
116 }
117
118 fn 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 }
131
132 pub const BatchedCholeskyRuntimeFamilyF32 = batchedCholeskyRuntimeFamily();
133
134 fn layoutTargetSegment(layout: TileLayout) []const u8 {
135 return switch (layout) {
136 .row_major => "",
137 .interleaved => "il_",
138 };
139 }
140
141 pub 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 }
148
149 pub 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 }
156
157 pub fn batchedCholeskyRuntimeArguments(instance: BatchedCholesky) ![1]choir_abi.ScalarArgument {
158 return .{
159 .{ .u32 = try runtimeExtentArgument(instance.batch) },
160 };
161 }
162
163 const testing = std.testing;
164
165 fn 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 }
186
187 fn 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 }
207
208 fn 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);
214
215 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 }
220
221 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 }
226
227 const l = try allocator.alloc(f32, batch * n * n);
228 defer allocator.free(l);
229 @memset(l, -1);
230
231 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 });
241
242 try testing.expectEqualSlices(f32, expected, l);
243 }
244
245 test "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 }
250
251 test "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 };
256
257 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;
261
262 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 });
273
274 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 }
280
281 test "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);
290
291 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 }));
297
298 const args = try batchedCholeskyRuntimeArguments(instance);
299 try testing.expectEqual(@as(u32, 5000), args[0].u32);
300 }
301
302 pub 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 };
309
310 pub const batched_cholesky_solve_family_version: u32 = 1;
311
312 pub 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 }
318
319 fn 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);
325
326 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 }
334
335 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 }
345
346 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 }
355
356 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 }
361
362 fn 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 }
375
376 fn batchedCholeskySolveFamilySchedule(instance: BatchedCholeskySolve) kernel.logical.schedule.ThreadBlocks {
377 return kernel.logical.schedule.threadBlocks(.{ .x = instance.threads });
378 }
379
380 fn 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 }
394
395 pub const BatchedCholeskySolveRuntimeFamilyF32 = batchedCholeskySolveRuntimeFamily();
396
397 pub 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 }
404
405 pub 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 }
412
413 pub fn batchedCholeskySolveRuntimeArguments(instance: BatchedCholeskySolve) ![1]choir_abi.ScalarArgument {
414 return .{
415 .{ .u32 = try runtimeExtentArgument(instance.batch) },
416 };
417 }
418
419 fn 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 }
441
442 fn 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);
448
449 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 }
466
467 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 }
472
473 const x = try allocator.alloc(f32, batch * n);
474 defer allocator.free(x);
475 @memset(x, -1);
476
477 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 });
488
489 try testing.expectEqualSlices(f32, expected, x);
490 }
491
492 test "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 }
497
498 test "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);
504
505 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 }));
509
510 const args = try batchedCholeskySolveRuntimeArguments(instance);
511 try testing.expectEqual(@as(u32, 5000), args[0].u32);
512 }
513
514 pub 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 };
521
522 pub const batched_inverse_family_version: u32 = 1;
523
524 pub 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 }
530
531 fn 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);
535
536 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 }
548
549 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 }
570
571 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 }
578
579 fn 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 }
592
593 fn batchedInverseFamilySchedule(instance: BatchedInverse) kernel.logical.schedule.ThreadBlocks {
594 return kernel.logical.schedule.threadBlocks(.{ .x = instance.threads });
595 }
596
597 fn 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 }
610
611 pub const BatchedInverseRuntimeFamilyF32 = batchedInverseRuntimeFamily();
612
613 pub 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 }
620
621 pub 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 }
628
629 pub fn batchedInverseRuntimeArguments(instance: BatchedInverse) ![1]choir_abi.ScalarArgument {
630 return .{
631 .{ .u32 = try runtimeExtentArgument(instance.batch) },
632 };
633 }
634
635 fn 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);
641
642 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 }
647
648 const inv = try allocator.alloc(f32, batch * n * n);
649 defer allocator.free(inv);
650 @memset(inv, -1);
651
652 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 });
662
663 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 }
678
679 test "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 }
684
685 test "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));
691
692 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 }
697
698 const a_interleaved = try allocator.alloc(f32, batch * n * n);
699 defer allocator.free(a_interleaved);
700 packInterleavedTiles(n, batch, a, a_interleaved);
701
702 const inv_interleaved = try allocator.alloc(f32, batch * n * n);
703 defer allocator.free(inv_interleaved);
704 @memset(inv_interleaved, -1);
705
706 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 });
716
717 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 }
731
732 test "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);
741
742 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 }));
748
749 const args = try batchedInverseRuntimeArguments(instance);
750 try testing.expectEqual(@as(u32, 5000), args[0].u32);
751 }
752
753 pub const batched_cholesky_row_axis = "i";
754 pub const batched_cholesky_col_axis = "j";
755
756 fn factorBatchBounds() shape.Bounds {
757 return .{ .min = 1, .max = extent_mod.runtime_extent_max };
758 }
759
760 fn factorTileBounds() shape.Bounds {
761 return .{ .min = batched_cholesky_min_n, .max = batched_cholesky_max_n };
762 }
763
764 pub 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 }
780
781 pub 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 }
786
787 pub 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();
794
795 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 );
805
806 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 );
816
817 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 }
836
837 pub 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 }
866
867 fn 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 }
872
873 fn 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) and
876 factorAxisMatches(candidate.axes[1], batched_cholesky_row_axis, n) and
877 factorAxisMatches(candidate.axes[2], batched_cholesky_col_axis, n);
878 }
879
880 fn 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) and
883 factorAxisMatches(candidate.axes[1], batched_cholesky_row_axis, n);
884 }
885
886 pub fn batchedCholeskyShapeProfileDimensions(instance: BatchedCholesky) [1]artifact_product.KernelCallShapeProfileDimension {
887 return .{
888 .{ .name = instance.batch_axis, .runtime_scalar_argument_index = 0, .bounds = factorBatchBounds() },
889 };
890 }
891
892 fn 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 }
903
904 pub 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 };
922
923 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 }
938
939 pub 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 }
956
957 pub 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 }
962
963 pub 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();
970
971 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 );
988
989 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 );
997
998 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 }
1017
1018 pub 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 }
1050
1051 pub fn batchedCholeskySolveShapeProfileDimensions(instance: BatchedCholeskySolve) [1]artifact_product.KernelCallShapeProfileDimension {
1052 return .{
1053 .{ .name = instance.batch_axis, .runtime_scalar_argument_index = 0, .bounds = factorBatchBounds() },
1054 };
1055 }
1056
1057 pub 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 };
1075
1076 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 }
1091
1092 pub 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 }
1108
1109 pub 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 }
1114
1115 pub 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();
1122
1123 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 );
1133
1134 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 );
1144
1145 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 }
1164
1165 pub 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 }
1194
1195 pub fn batchedInverseShapeProfileDimensions(instance: BatchedInverse) [1]artifact_product.KernelCallShapeProfileDimension {
1196 return .{
1197 .{ .name = instance.batch_axis, .runtime_scalar_argument_index = 0, .bounds = factorBatchBounds() },
1198 };
1199 }
1200
1201 pub 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 };
1219
1220 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 }
1235
1236 test "linalg batched factor specializations round-trip their instances" {
1237 const allocator = testing.allocator;
1238
1239 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(.{}));
1249
1250 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 );
1266
1267 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 }
1280
1281 test "linalg batched factor specializations reject malformed tensor facts" {
1282 const allocator = testing.allocator;
1283
1284 {
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 }
1291
1292 {
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 }
1299
1300 {
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 }
1310
1311 {
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 }
1321
1322 {
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 }
1330
1331 /// A caller uses this to rearrange a batch of square matrices on the host before a factorization
1332 /// kernel reads them: the function copies `batch` matrices of `n` by `n` values, stored one matrix
1333 /// after another in row order, into an interleaved layout, where the same element of every matrix
1334 /// sits side by side. Element `slot` of matrix `b` lands at index `slot * batch + b`. The caller
1335 /// sizes both slices to hold `n * n * batch` values, and the function neither checks those lengths
1336 /// nor allocates.
1337 pub 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 }
1344
1345 /// A caller uses this to rearrange a batch of vectors, such as right-hand sides for a solve: the
1346 /// function copies `batch` vectors of `n` values, stored one vector after another, into an
1347 /// interleaved layout, where the same element of every vector sits side by side. Element `slot` of
1348 /// 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.
1350 pub 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 }
1357
1358 test "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));
1364
1365 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 }
1375
1376 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);
1382
1383 const l = try allocator.alloc(f32, batch * n * n);
1384 defer allocator.free(l);
1385 @memset(l, -1);
1386
1387 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 });
1397
1398 try testing.expectEqualSlices(f32, expected_interleaved, l);
1399 }
1400
1401 test "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));
1407
1408 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 }
1430
1431 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);
1440
1441 const x = try allocator.alloc(f32, batch * n);
1442 defer allocator.free(x);
1443 @memset(x, -1);
1444
1445 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 });
1456
1457 try testing.expectEqualSlices(f32, expected_interleaved, x);
1458 }
1459
1460 test "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);
1466
1467 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 }