lib/accy/src/kernel/library/sparse.zig
daab053ee43316e1809a84551d573ddd1e5bf3d2
1 const std = @import("std");
2 const gpu = @import("gpu");
3 const choir_abi = @import("choir_abi");
4
5 const artifact_product = @import("../../artifact/model/root.zig");
6 const shape = @import("../../choir/shape/root.zig");
7 const entry = @import("entry.zig");
8 const extent_mod = @import("extent.zig");
9 const geometry_mod = @import("geometry.zig");
10 const kernel = @import("../root.zig");
11 const tuning = @import("tuning.zig");
12
13 const DType = choir_abi.DType;
14 const runtimeExtentArgument = extent_mod.runtimeExtentArgument;
15
16 pub const SpmvCsrStructure = enum {
17 row_thread,
18 row_warp,
19 };
20
21 pub const SpmvCsr = struct {
22 rows: u64,
23 nnz: u64 = 1,
24 x_extent: u64 = 1,
25 dtype: DType = .f32,
26 accumulation_dtype: DType = .f32,
27 threads: u32 = 256,
28 structure: SpmvCsrStructure = .row_warp,
29 row_axis: []const u8 = "r",
30 };
31
32 pub const SpmvCooStructure = enum {
33 element_thread,
34 row_thread,
35 };
36
37 pub const SpmvCoo = struct {
38 rows: u64,
39 nnz: u64 = 1,
40 x_extent: u64 = 1,
41 dtype: DType = .f32,
42 accumulation_dtype: DType = .f32,
43 threads: u32 = 256,
44 structure: SpmvCooStructure = .element_thread,
45 row_axis: []const u8 = "r",
46 nonzero_axis: []const u8 = "n",
47 x_axis: []const u8 = "x",
48 };
49
50 pub const SpmvEllStructure = enum {
51 row_thread,
52 };
53
54 pub const SpmvEll = struct {
55 rows: u64,
56 slots: u64 = 1,
57 x_extent: u64 = 1,
58 dtype: DType = .f32,
59 accumulation_dtype: DType = .f32,
60 threads: u32 = 256,
61 structure: SpmvEllStructure = .row_thread,
62 row_axis: []const u8 = "r",
63 slot_axis: []const u8 = "s",
64 x_axis: []const u8 = "x",
65 };
66
67 pub const SpmvSellStructure = enum {
68 row_thread,
69 };
70
71 pub const SpmvSell = struct {
72 rows: u64,
73 slice_size: u64 = 32,
74 values_size: u64 = 1,
75 x_extent: u64 = 1,
76 dtype: DType = .f32,
77 accumulation_dtype: DType = .f32,
78 threads: u32 = 256,
79 structure: SpmvSellStructure = .row_thread,
80 row_axis: []const u8 = "r",
81 slice_axis: []const u8 = "z",
82 value_axis: []const u8 = "n",
83 x_axis: []const u8 = "x",
84 };
85
86 pub const SpmmCsrStructure = enum {
87 row_column_thread,
88 };
89
90 pub const SpmmCsr = struct {
91 rows: u64,
92 columns: u64,
93 nnz: u64 = 1,
94 x_extent: u64 = 1,
95 dtype: DType = .f32,
96 accumulation_dtype: DType = .f32,
97 threads: entry.Threads2D = .{},
98 structure: SpmmCsrStructure = .row_column_thread,
99 row_axis: []const u8 = "r",
100 column_axis: []const u8 = "c",
101 x_axis: []const u8 = "x",
102 };
103
104 pub const spmv_csr_family_version: u32 = 1;
105 pub const spmv_coo_family_version: u32 = 1;
106 pub const spmv_ell_family_version: u32 = 1;
107 pub const spmv_sell_family_version: u32 = 1;
108 pub const spmm_csr_family_version: u32 = 1;
109 pub const spmv_csr_warp_size: u32 = 32;
110 pub const spmv_csr_max_threads: u32 = 1024;
111 pub const spmv_coo_max_threads: u32 = 1024;
112 pub const spmv_ell_max_threads: u32 = 1024;
113 pub const spmv_sell_max_threads: u32 = 1024;
114 pub const spmm_csr_max_threads: u32 = 1024;
115
116 const spmv_ell_thread_caps = geometry_mod.ThreadCaps1D{ .budget = spmv_ell_max_threads };
117 const spmv_sell_thread_caps = geometry_mod.ThreadCaps1D{ .budget = spmv_sell_max_threads };
118 const spmm_csr_thread_caps = geometry_mod.ThreadCaps{
119 .budget = 256,
120 .x_max = 64,
121 .y_max = 16,
122 };
123
124 pub fn spmvCsrWarpsPerBlock(threads: u32) u32 {
125 return threads / spmv_csr_warp_size;
126 }
127
128 pub fn spmvCsrRowsPerBlock(instance: SpmvCsr) u32 {
129 return switch (instance.structure) {
130 .row_thread => instance.threads,
131 .row_warp => spmvCsrWarpsPerBlock(instance.threads),
132 };
133 }
134
135 pub fn spmvCsrBlockCount(instance: SpmvCsr) u64 {
136 return spmvCsrBlockCountChecked(instance).?;
137 }
138
139 fn sparseFloatAccumulationDType(dtype: DType) ?DType {
140 return switch (dtype) {
141 .f64 => .f64,
142 .f32, .f16 => .f32,
143 else => null,
144 };
145 }
146
147 pub fn sparseAccumulationDType(dtype: DType) ?DType {
148 return sparseFloatAccumulationDType(dtype);
149 }
150
151 pub fn spmvCsrAccumulationDType(dtype: DType) ?DType {
152 return sparseAccumulationDType(dtype);
153 }
154
155 pub fn spmvCooAccumulationDType(dtype: DType) ?DType {
156 const structure = spmvCooDefaultStructure(dtype) orelse return null;
157 return spmvCooAccumulationDTypeForStructure(structure, dtype);
158 }
159
160 pub fn spmvCooDefaultStructure(dtype: DType) ?SpmvCooStructure {
161 return switch (dtype) {
162 .f32 => .element_thread,
163 .f16, .f64 => .row_thread,
164 else => null,
165 };
166 }
167
168 pub fn spmvCooAccumulationDTypeForStructure(structure: SpmvCooStructure, dtype: DType) ?DType {
169 return switch (structure) {
170 .element_thread => switch (dtype) {
171 .f32 => .f32,
172 else => null,
173 },
174 .row_thread => sparseAccumulationDType(dtype),
175 };
176 }
177
178 pub fn spmvEllAccumulationDType(dtype: DType) ?DType {
179 return sparseAccumulationDType(dtype);
180 }
181
182 pub fn spmvSellAccumulationDType(dtype: DType) ?DType {
183 return sparseAccumulationDType(dtype);
184 }
185
186 pub fn spmmCsrAccumulationDType(dtype: DType) ?DType {
187 return sparseAccumulationDType(dtype);
188 }
189
190 fn spmvCsrBlockCountChecked(instance: SpmvCsr) ?u64 {
191 const rows_per_block = spmvCsrRowsPerBlock(instance);
192 if (rows_per_block == 0) return null;
193 const biased = std.math.add(u64, instance.rows, rows_per_block - 1) catch return null;
194 return biased / rows_per_block;
195 }
196
197 pub fn spmvCsrInstanceValid(instance: SpmvCsr) bool {
198 const accumulation_dtype = spmvCsrAccumulationDType(instance.dtype) orelse return false;
199 if (instance.accumulation_dtype != accumulation_dtype) return false;
200 if (instance.rows == 0) return false;
201 if (instance.nnz == 0) return false;
202 if (instance.x_extent == 0) return false;
203 if (instance.threads == 0 or instance.threads > spmv_csr_max_threads) return false;
204 if (spmvCsrLaunchExtentChecked(instance) == null) return false;
205 if (spmvCsrBlockCountChecked(instance) == null) return false;
206 return switch (instance.structure) {
207 .row_thread => true,
208 .row_warp => instance.threads % spmv_csr_warp_size == 0,
209 };
210 }
211
212 fn spmvCooBlockCountChecked(instance: SpmvCoo) ?u64 {
213 if (instance.threads == 0) return null;
214 const extent = spmvCooLaunchExtentChecked(instance) orelse return null;
215 const biased = std.math.add(u64, extent, instance.threads - 1) catch return null;
216 return biased / instance.threads;
217 }
218
219 pub fn spmvCooBlockCount(instance: SpmvCoo) u64 {
220 return spmvCooBlockCountChecked(instance).?;
221 }
222
223 pub fn spmvCooInstanceValid(instance: SpmvCoo) bool {
224 const accumulation_dtype = spmvCooAccumulationDTypeForStructure(instance.structure, instance.dtype) orelse return false;
225 if (instance.accumulation_dtype != accumulation_dtype) return false;
226 if (instance.rows == 0 or instance.nnz == 0 or instance.x_extent == 0) return false;
227 if (instance.threads == 0 or instance.threads > spmv_coo_max_threads) return false;
228 if (spmvCooLaunchExtentChecked(instance) == null) return false;
229 if (spmvCooBlockCountChecked(instance) == null) return false;
230 return switch (instance.structure) {
231 .element_thread => true,
232 .row_thread => true,
233 };
234 }
235
236 fn sparseFloatAccumulationZero(inner_builder: anytype, accumulation_dtype: DType) !kernel.Value {
237 return switch (accumulation_dtype) {
238 .f64 => inner_builder.constantFloat(.f64, 0.0),
239 .f32 => inner_builder.constantFloat(.f32, 0.0),
240 else => error.UnsupportedDType,
241 };
242 }
243
244 fn spmvCsrAccumulationZero(inner_builder: anytype, instance: SpmvCsr) !kernel.Value {
245 return sparseFloatAccumulationZero(inner_builder, instance.accumulation_dtype);
246 }
247
248 fn sparseFloatAccumulationValue(inner_builder: anytype, accumulation_dtype: DType, value: anytype) !kernel.Value {
249 return switch (accumulation_dtype) {
250 .f64 => if (comptime @TypeOf(value).scalar_dtype == .f64) value.raw() else (try value.cast(inner_builder, .f64)).raw(),
251 .f32 => if (comptime @TypeOf(value).scalar_dtype == .f32) value.raw() else (try value.cast(inner_builder, .f32)).raw(),
252 else => error.UnsupportedDType,
253 };
254 }
255
256 fn spmvCsrAccumulationValue(inner_builder: anytype, instance: SpmvCsr, value: anytype) !kernel.Value {
257 return sparseFloatAccumulationValue(inner_builder, instance.accumulation_dtype, value);
258 }
259
260 fn sparseFloatOutputValue(inner_builder: anytype, dtype: DType, accumulation_dtype: DType, value: kernel.Value) !kernel.Value {
261 if (dtype == accumulation_dtype) return value;
262 return switch (dtype) {
263 .f64 => inner_builder.cast(value, .f64),
264 .f32 => inner_builder.cast(value, .f32),
265 .f16 => inner_builder.cast(value, .f16),
266 else => error.UnsupportedDType,
267 };
268 }
269
270 fn spmvCsrOutputValue(inner_builder: anytype, instance: SpmvCsr, value: kernel.Value) !kernel.Value {
271 return sparseFloatOutputValue(inner_builder, instance.dtype, instance.accumulation_dtype, value);
272 }
273
274 fn spmvCooAccumulationValue(inner_builder: anytype, instance: SpmvCoo, value: anytype) !kernel.Value {
275 return sparseFloatAccumulationValue(inner_builder, instance.accumulation_dtype, value);
276 }
277
278 fn spmvCooAccumulationZero(inner_builder: anytype, instance: SpmvCoo) !kernel.Value {
279 return sparseFloatAccumulationZero(inner_builder, instance.accumulation_dtype);
280 }
281
282 fn spmvCooOutputValue(inner_builder: anytype, instance: SpmvCoo, value: kernel.Value) !kernel.Value {
283 return sparseFloatOutputValue(inner_builder, instance.dtype, instance.accumulation_dtype, value);
284 }
285
286 fn spmv_csr_row_value_apply(fold_builder: anytype, element: kernel.Value, current: kernel.Value, ctx: anytype) !kernel.Value {
287 const column_loaded = try ctx.args.param(.cols).load(fold_builder, element);
288 const column_index = try fold_builder.castIndex(column_loaded.raw());
289 const column_lower = try fold_builder.max(column_index, ctx.zero);
290 const x_last = try fold_builder.sub(ctx.x_extent, ctx.one);
291 const column = try fold_builder.min(column_lower, x_last);
292 const x_value = try ctx.args.param(.x).load(fold_builder, column);
293 const matrix_value = try ctx.args.param(.values).load(fold_builder, element);
294 const matrix_acc = try spmvCsrAccumulationValue(fold_builder, ctx.instance, matrix_value);
295 const x_acc = try spmvCsrAccumulationValue(fold_builder, ctx.instance, x_value);
296 const product = try fold_builder.mul(matrix_acc, x_acc);
297 return fold_builder.add(current, product);
298 }
299
300 fn spmvCsrRowValue(
301 inner_builder: anytype,
302 instance: SpmvCsr,
303 args: anytype,
304 row: kernel.Value,
305 lane: kernel.Value,
306 nnz: kernel.Value,
307 x_extent: kernel.Value,
308 ) !kernel.Value {
309 const one = try inner_builder.constantIndex(1);
310 const next = try inner_builder.add(row, one);
311 const begin_loaded = try args.param(.row_ptr).load(inner_builder, row);
312 const end_loaded = try args.param(.row_ptr).load(inner_builder, next);
313 const begin_index = try inner_builder.castIndex(begin_loaded.raw());
314 const end_index = try inner_builder.castIndex(end_loaded.raw());
315 const zero = try inner_builder.constantIndex(0);
316 const end_lower = try inner_builder.max(end_index, zero);
317 const end_clamped = try inner_builder.min(end_lower, nnz);
318 const begin_lower = try inner_builder.max(begin_index, zero);
319 const begin_clamped = try inner_builder.min(begin_lower, end_clamped);
320 const lane_begin = try inner_builder.add(begin_clamped, lane);
321 const stride = try inner_builder.constantIndex(spmv_csr_warp_size);
322
323 const acc_zero = try spmvCsrAccumulationZero(inner_builder, instance);
324 const partial = try inner_builder.fold(lane_begin, end_clamped, stride, acc_zero, .{
325 .instance = instance,
326 .args = args,
327 .x_extent = x_extent,
328 .zero = zero,
329 .one = one,
330 }, spmv_csr_row_value_apply);
331 return inner_builder.warpReduce(.add, partial);
332 }
333
334 fn spmv_csr_row_thread_value_apply(fold_builder: anytype, element: kernel.Value, current: kernel.Value, ctx: anytype) !kernel.Value {
335 const column_loaded = try ctx.args.param(.cols).load(fold_builder, element);
336 const column_index = try fold_builder.castIndex(column_loaded.raw());
337 const column_lower = try fold_builder.max(column_index, ctx.zero);
338 const x_last = try fold_builder.sub(ctx.x_extent, ctx.one);
339 const column = try fold_builder.min(column_lower, x_last);
340 const x_value = try ctx.args.param(.x).load(fold_builder, column);
341 const matrix_value = try ctx.args.param(.values).load(fold_builder, element);
342 const matrix_acc = try spmvCsrAccumulationValue(fold_builder, ctx.instance, matrix_value);
343 const x_acc = try spmvCsrAccumulationValue(fold_builder, ctx.instance, x_value);
344 const product = try fold_builder.mul(matrix_acc, x_acc);
345 return fold_builder.add(current, product);
346 }
347
348 fn spmvCsrRowThreadValue(
349 inner_builder: anytype,
350 instance: SpmvCsr,
351 args: anytype,
352 row: kernel.Value,
353 nnz: kernel.Value,
354 x_extent: kernel.Value,
355 ) !kernel.Value {
356 const one = try inner_builder.constantIndex(1);
357 const next = try inner_builder.add(row, one);
358 const begin_loaded = try args.param(.row_ptr).load(inner_builder, row);
359 const end_loaded = try args.param(.row_ptr).load(inner_builder, next);
360 const begin_index = try inner_builder.castIndex(begin_loaded.raw());
361 const end_index = try inner_builder.castIndex(end_loaded.raw());
362 const zero = try inner_builder.constantIndex(0);
363 const end_lower = try inner_builder.max(end_index, zero);
364 const end_clamped = try inner_builder.min(end_lower, nnz);
365 const begin_lower = try inner_builder.max(begin_index, zero);
366 const begin_clamped = try inner_builder.min(begin_lower, end_clamped);
367
368 const acc_zero = try spmvCsrAccumulationZero(inner_builder, instance);
369 return inner_builder.fold(begin_clamped, end_clamped, one, acc_zero, .{
370 .instance = instance,
371 .args = args,
372 .x_extent = x_extent,
373 .zero = zero,
374 .one = one,
375 }, spmv_csr_row_thread_value_apply);
376 }
377
378 fn spmv_csr_row_thread_runtime_body_active(inner_builder: anytype, ctx: anytype) !void {
379 const sum = try spmvCsrRowThreadValue(inner_builder, ctx.instance, ctx.args, ctx.row, ctx.nnz, ctx.x_extent);
380 const result = try spmvCsrOutputValue(inner_builder, ctx.instance, sum);
381 try ctx.args.param(.y).store(inner_builder, result, ctx.row);
382 }
383
384 fn spmvCsrRowThreadRuntimeBody(k: anytype, spec: SpmvCsr, args: anytype) !void {
385 if (!spmvCsrInstanceValid(spec)) return error.UnsupportedSpmvCsrInstance;
386 const row = try k.globalId(.x);
387 const rows_extent = try k.castIndex(args.param(.rows).raw());
388 const nnz = try k.castIndex(args.param(.nnz).raw());
389 const x_extent = try k.castIndex(args.param(.x_extent).raw());
390 const active = try k.compare(.lt, row, rows_extent);
391 try k.guardDo(active, .{
392 .args = args,
393 .row = row,
394 .nnz = nnz,
395 .x_extent = x_extent,
396 .instance = spec,
397 }, spmv_csr_row_thread_runtime_body_active);
398 }
399
400 fn spmvCsrRuntimeBody(k: anytype, spec: SpmvCsr, args: anytype) !void {
401 return switch (spec.structure) {
402 .row_thread => spmvCsrRowThreadRuntimeBody(k, spec, args),
403 .row_warp => spmvCsrRowWarpRuntimeBody(k, spec, args),
404 };
405 }
406
407 fn spmv_csr_row_warp_runtime_body_active(inner_builder: anytype, ctx: anytype) !void {
408 const sum = try spmvCsrRowValue(inner_builder, ctx.instance, ctx.args, ctx.row, ctx.lane, ctx.nnz, ctx.x_extent);
409 const result = try spmvCsrOutputValue(inner_builder, ctx.instance, sum);
410 const zero = try inner_builder.constantIndex(0);
411 const writer = try inner_builder.compare(.eq, ctx.lane, zero);
412 try inner_builder.guardDo(writer, .{
413 .args = ctx.args,
414 .row = ctx.row,
415 .sum = result,
416 }, spmv_csr_row_warp_runtime_body_writer);
417 }
418
419 fn spmv_csr_row_warp_runtime_body_writer(writer_builder: anytype, writer_ctx: anytype) !void {
420 try writer_ctx.args.param(.y).store(writer_builder, writer_ctx.sum, writer_ctx.row);
421 }
422
423 fn spmvCsrRowWarpRuntimeBody(k: anytype, spec: SpmvCsr, args: anytype) !void {
424 if (!spmvCsrInstanceValid(spec)) return error.UnsupportedSpmvCsrInstance;
425 const element_thread = try k.globalId(.x);
426 const lane = try k.laneId();
427 const warp_size = try k.constantIndex(spmv_csr_warp_size);
428 const row = try k.div(element_thread, warp_size);
429 const rows_extent = try k.castIndex(args.param(.rows).raw());
430 const nnz = try k.castIndex(args.param(.nnz).raw());
431 const x_extent = try k.castIndex(args.param(.x_extent).raw());
432 const active = try k.compare(.lt, row, rows_extent);
433 try k.guardDo(active, .{
434 .args = args,
435 .row = row,
436 .lane = lane,
437 .nnz = nnz,
438 .x_extent = x_extent,
439 .instance = spec,
440 }, spmv_csr_row_warp_runtime_body_active);
441 }
442
443 fn spmvCsrFamilySchedule(instance: SpmvCsr) kernel.logical.schedule.ThreadBlocks {
444 return kernel.logical.schedule.threadBlocks(.{ .x = instance.threads });
445 }
446
447 fn spmvCsrRuntimeFamily(comptime dtype: DType) type {
448 return kernel.logical.Family(.{
449 .name = std.fmt.comptimePrint("accy_kernel_sparse_spmv_csr_runtime_{s}", .{dtype.name()}),
450 .parameters = .{
451 .y = kernel.dynamicBuffer(dtype),
452 .row_ptr = kernel.dynamicBuffer(.i32),
453 .cols = kernel.dynamicBuffer(.i32),
454 .values = kernel.dynamicBuffer(dtype),
455 .x = kernel.dynamicBuffer(dtype),
456 .rows = kernel.scalar(.i32),
457 .nnz = kernel.scalar(.i32),
458 .x_extent = kernel.scalar(.i32),
459 },
460 .Instance = SpmvCsr,
461 .schedule = spmvCsrFamilySchedule,
462 .body = spmvCsrRuntimeBody,
463 });
464 }
465
466 pub const SpmvCsrRuntimeFamilyF64 = spmvCsrRuntimeFamily(.f64);
467 pub const SpmvCsrRuntimeFamilyF32 = spmvCsrRuntimeFamily(.f32);
468 pub const SpmvCsrRuntimeFamilyF16 = spmvCsrRuntimeFamily(.f16);
469
470 pub fn spmvCsrFamilyTarget(allocator: std.mem.Allocator, instance: SpmvCsr) ![]u8 {
471 return std.fmt.allocPrint(
472 allocator,
473 "accy.kernel.sparse.spmv_csr_{s}_family_{d}_{s}",
474 .{ @tagName(instance.structure), instance.threads, instance.dtype.name() },
475 );
476 }
477
478 pub fn spmvCsrFamilyEntryName(allocator: std.mem.Allocator, instance: SpmvCsr) ![]u8 {
479 return std.fmt.allocPrint(
480 allocator,
481 "accy_kernel_sparse_spmv_csr_{s}_family_{d}_{s}",
482 .{ @tagName(instance.structure), instance.threads, instance.dtype.name() },
483 );
484 }
485
486 pub fn spmvCsrRuntimeArguments(
487 instance: SpmvCsr,
488 nnz: u64,
489 x_extent: u64,
490 ) ![3]choir_abi.ScalarArgument {
491 return .{
492 .{ .u32 = try runtimeExtentArgument(instance.rows) },
493 .{ .u32 = try runtimeExtentArgument(nnz) },
494 .{ .u32 = try runtimeExtentArgument(x_extent) },
495 };
496 }
497
498 pub fn spmvCsrMaxRows(instance: SpmvCsr) u64 {
499 _ = instance;
500 return std.math.maxInt(u32);
501 }
502
503 pub fn spmvCsrShapeProfileDimensions(instance: SpmvCsr) [1]artifact_product.KernelCallShapeProfileDimension {
504 return .{
505 .{
506 .name = instance.row_axis,
507 .runtime_scalar_argument_index = 0,
508 .bounds = .{ .min = 1, .max = spmvCsrMaxRows(instance) },
509 },
510 };
511 }
512
513 fn spmvCsrLaunch(instance: SpmvCsr) !artifact_product.KernelCallLaunch {
514 if (!spmvCsrInstanceValid(instance)) return error.UnsupportedSpmvCsrInstance;
515 return .{ .derived = .{
516 .grid = .{
517 .{ .runtime_u32_ceil_div = .{ .argument_index = 0, .divisor = spmvCsrRowsPerBlock(instance) } },
518 .{ .fixed = 1 },
519 .{ .fixed = 1 },
520 },
521 .threadgroup = .{ instance.threads, 1, 1 },
522 } };
523 }
524
525 pub fn spmvCsrShapeFamily(backing_allocator: std.mem.Allocator, instance: SpmvCsr) !shape.Family {
526 var builder = try shape.Builder.init(backing_allocator, "spmv_csr");
527 errdefer builder.deinit();
528 const rows = try builder.symbol(instance.row_axis);
529 const rows_expr = try builder.symbolExpression(rows);
530 _ = try builder.tensor("y", &.{rows_expr});
531 try builder.assumeBounds(rows_expr, .{ .min = 1, .max = spmvCsrMaxRows(instance) });
532 return builder.finish();
533 }
534
535 pub fn spmvCsrFamilyFingerprint(backing_allocator: std.mem.Allocator, instance: SpmvCsr) !u64 {
536 var family = try spmvCsrShapeFamily(backing_allocator, instance);
537 defer family.deinit();
538 return shape.fingerprint(family);
539 }
540
541 pub fn spmvCsrTuningExtents(instance: SpmvCsr) [3]u64 {
542 return .{ instance.rows, instance.nnz, instance.x_extent };
543 }
544
545 pub fn spmvCsrTuningOperation(instance: SpmvCsr) entry.Operation {
546 _ = instance;
547 return .{ .sparse = .csr_spmv };
548 }
549
550 pub fn spmvCsrFamilyTuningKey(
551 backing_allocator: std.mem.Allocator,
552 device_fingerprint: u64,
553 instance: SpmvCsr,
554 ) !tuning.FamilyTuningKey {
555 const family_fingerprint = try spmvCsrFamilyFingerprint(backing_allocator, instance);
556 const extents = spmvCsrTuningExtents(instance);
557 return tuning.FamilyTuningKey.init(
558 device_fingerprint,
559 family_fingerprint,
560 entry.operationFingerprint(spmvCsrTuningOperation(instance)),
561 instance.dtype,
562 spmv_csr_family_version,
563 extents[0..],
564 ) orelse unreachable;
565 }
566
567 pub fn resolveSpmvCsrStructure(
568 backing_allocator: std.mem.Allocator,
569 reader: tuning.FamilyTuningReader,
570 instance: SpmvCsr,
571 ) !?SpmvCsrStructure {
572 const key = try spmvCsrFamilyTuningKey(backing_allocator, reader.device_fingerprint, instance);
573 const record = reader.table.find(key) orelse return null;
574 const structures = [_]SpmvCsrStructure{ .row_thread, .row_warp };
575 for (structures) |structure| {
576 var candidate = instance;
577 candidate.structure = structure;
578 candidate.threads = spmvCsrRepresentableThreads(candidate) orelse continue;
579 if (!spmvCsrInstanceValid(candidate)) continue;
580 const target = try spmvCsrFamilyTarget(backing_allocator, candidate);
581 defer backing_allocator.free(target);
582 if (std.mem.eql(u8, target, record.target)) return structure;
583 }
584 return null;
585 }
586
587 pub fn createSpmvCsrFamilyArtifact(
588 allocator: std.mem.Allocator,
589 handle: kernel.BackendHandle,
590 instance: SpmvCsr,
591 options: entry.ArtifactOptions,
592 ) !kernel.OwnedKernelCallArtifact {
593 if (!spmvCsrInstanceValid(instance)) return error.UnsupportedSpmvCsrInstance;
594 const target = try spmvCsrFamilyTarget(allocator, instance);
595 defer allocator.free(target);
596 const entry_name = try spmvCsrFamilyEntryName(allocator, instance);
597 defer allocator.free(entry_name);
598 const family_fingerprint = options.shape_family_fingerprint orelse try spmvCsrFamilyFingerprint(allocator, instance);
599 const shape_profile_dimensions = spmvCsrShapeProfileDimensions(instance);
600 const shape_profile = options.shape_profile orelse artifact_product.KernelCallShapeProfile{
601 .name = "spmv_csr",
602 .fingerprint = family_fingerprint,
603 .dimensions = shape_profile_dimensions[0..],
604 };
605
606 var graph = switch (instance.dtype) {
607 .f64 => try SpmvCsrRuntimeFamilyF64.buildNamed(allocator, options.limits, entry_name, instance),
608 .f32 => try SpmvCsrRuntimeFamilyF32.buildNamed(allocator, options.limits, entry_name, instance),
609 .f16 => try SpmvCsrRuntimeFamilyF16.buildNamed(allocator, options.limits, entry_name, instance),
610 else => return error.UnsupportedDType,
611 };
612 defer graph.deinit();
613 return kernel.createKernelCallArtifact(allocator, handle, &graph, .{
614 .target = target,
615 .version = spmv_csr_family_version,
616 .format = options.format,
617 .kernel_plan = options.kernel_plan,
618 .element_count_argument = options.element_count_argument,
619 .shape_family_fingerprint = family_fingerprint,
620 .shape_profile = shape_profile,
621 .launch = options.launch orelse try spmvCsrLaunch(instance),
622 .runtime_scalar_argument_count = if (options.runtime_scalar_argument_count == 0) 3 else options.runtime_scalar_argument_count,
623 .static_arguments = options.static_arguments,
624 });
625 }
626
627 fn spmv_coo_element_thread_runtime_body_active(guard_builder: anytype, ctx: anytype) !void {
628 const row_loaded = try ctx.args.param(.row_indices).load(guard_builder, ctx.element);
629 const zero_i32 = try guard_builder.constantInt(.i32, 0);
630 const row_non_negative = try guard_builder.compare(.ge, row_loaded.raw(), zero_i32);
631 try guard_builder.guardDo(row_non_negative, .{
632 .args = ctx.args,
633 .element = ctx.element,
634 .instance = ctx.instance,
635 .row_raw = row_loaded.raw(),
636 }, spmv_coo_element_thread_runtime_body_row_non_negative);
637 }
638
639 fn spmv_coo_element_thread_runtime_body_row_non_negative(row_builder: anytype, row_ctx: anytype) !void {
640 const col_loaded = try row_ctx.args.param(.cols).load(row_builder, row_ctx.element);
641 const zero_i32_inner = try row_builder.constantInt(.i32, 0);
642 const col_non_negative = try row_builder.compare(.ge, col_loaded.raw(), zero_i32_inner);
643 try row_builder.guardDo(col_non_negative, .{
644 .args = row_ctx.args,
645 .element = row_ctx.element,
646 .instance = row_ctx.instance,
647 .row_raw = row_ctx.row_raw,
648 .col_raw = col_loaded.raw(),
649 }, spmv_coo_element_thread_runtime_body_col_non_negative);
650 }
651
652 fn spmv_coo_element_thread_runtime_body_col_non_negative(col_builder: anytype, col_ctx: anytype) !void {
653 const rows_extent = try col_builder.castIndex(col_ctx.args.param(.rows).raw());
654 const x_extent = try col_builder.castIndex(col_ctx.args.param(.x_extent).raw());
655 const row = try col_builder.castIndex(col_ctx.row_raw);
656 const col = try col_builder.castIndex(col_ctx.col_raw);
657 const row_in_range = try col_builder.compare(.lt, row, rows_extent);
658 try col_builder.guardDo(row_in_range, .{
659 .args = col_ctx.args,
660 .element = col_ctx.element,
661 .instance = col_ctx.instance,
662 .row = row,
663 .col = col,
664 .x_extent = x_extent,
665 }, spmv_coo_element_thread_runtime_body_row_in_range);
666 }
667
668 fn spmv_coo_element_thread_runtime_body_row_in_range(row_range_builder: anytype, row_range_ctx: anytype) !void {
669 const col_in_range = try row_range_builder.compare(.lt, row_range_ctx.col, row_range_ctx.x_extent);
670 try row_range_builder.guardDo(col_in_range, .{
671 .args = row_range_ctx.args,
672 .element = row_range_ctx.element,
673 .instance = row_range_ctx.instance,
674 .row = row_range_ctx.row,
675 .col = row_range_ctx.col,
676 }, spmv_coo_element_thread_runtime_body_col_in_range);
677 }
678
679 fn spmv_coo_element_thread_runtime_body_col_in_range(atomic_builder: anytype, atomic_ctx: anytype) !void {
680 const x_value = try atomic_ctx.args.param(.x).load(atomic_builder, atomic_ctx.col);
681 const matrix_value = try atomic_ctx.args.param(.values).load(atomic_builder, atomic_ctx.element);
682 const matrix_acc = try spmvCooAccumulationValue(atomic_builder, atomic_ctx.instance, matrix_value);
683 const x_acc = try spmvCooAccumulationValue(atomic_builder, atomic_ctx.instance, x_value);
684 const product = try atomic_builder.mul(matrix_acc, x_acc);
685 _ = try atomic_ctx.args.param(.y).atomicRmw(atomic_builder, .add, product, atomic_ctx.row);
686 }
687
688 fn spmvCooElementThreadRuntimeBody(k: anytype, spec: SpmvCoo, args: anytype) !void {
689 if (!spmvCooInstanceValid(spec)) return error.UnsupportedSpmvCooInstance;
690 const element = try k.globalId(.x);
691 const nnz = try k.castIndex(args.param(.nnz).raw());
692 const active = try k.compare(.lt, element, nnz);
693 try k.guardDo(active, .{
694 .args = args,
695 .element = element,
696 .instance = spec,
697 }, spmv_coo_element_thread_runtime_body_active);
698 }
699
700 fn spmv_coo_row_thread_value_apply(fold_builder: anytype, element: kernel.Value, current: kernel.Value, ctx: anytype) !kernel.Value {
701 const row_loaded = try ctx.args.param(.row_indices).load(fold_builder, element);
702 const row_i32 = try fold_builder.cast(ctx.row, .i32);
703 const row_matches = try fold_builder.compare(.eq, row_loaded.raw(), row_i32);
704 const col_loaded = try ctx.args.param(.cols).load(fold_builder, element);
705 const zero_i32 = try fold_builder.constantInt(.i32, 0);
706 const col_non_negative = try fold_builder.compare(.ge, col_loaded.raw(), zero_i32);
707 const x_extent_i32 = try fold_builder.cast(ctx.x_extent, .i32);
708 const col_in_range = try fold_builder.compare(.lt, col_loaded.raw(), x_extent_i32);
709 const column_index = try fold_builder.castIndex(col_loaded.raw());
710 const column_lower = try fold_builder.max(column_index, ctx.zero);
711 const x_last = try fold_builder.sub(ctx.x_extent, ctx.one);
712 const column = try fold_builder.min(column_lower, x_last);
713 const x_value = try ctx.args.param(.x).load(fold_builder, column);
714 const matrix_value = try ctx.args.param(.values).load(fold_builder, element);
715 const matrix_acc = try spmvCooAccumulationValue(fold_builder, ctx.instance, matrix_value);
716 const x_acc = try spmvCooAccumulationValue(fold_builder, ctx.instance, x_value);
717 const product = try fold_builder.mul(matrix_acc, x_acc);
718 const row_contribution = try fold_builder.select(row_matches, product, ctx.acc_zero);
719 const non_negative_contribution = try fold_builder.select(col_non_negative, row_contribution, ctx.acc_zero);
720 const contribution = try fold_builder.select(col_in_range, non_negative_contribution, ctx.acc_zero);
721 return fold_builder.add(current, contribution);
722 }
723
724 fn spmvCooRowThreadValue(
725 inner_builder: anytype,
726 instance: SpmvCoo,
727 args: anytype,
728 row: kernel.Value,
729 nnz: kernel.Value,
730 x_extent: kernel.Value,
731 ) !kernel.Value {
732 const zero = try inner_builder.constantIndex(0);
733 const one = try inner_builder.constantIndex(1);
734 const acc_zero = try spmvCooAccumulationZero(inner_builder, instance);
735 return inner_builder.fold(zero, nnz, one, acc_zero, .{
736 .instance = instance,
737 .args = args,
738 .row = row,
739 .x_extent = x_extent,
740 .zero = zero,
741 .one = one,
742 .acc_zero = acc_zero,
743 }, spmv_coo_row_thread_value_apply);
744 }
745
746 fn spmv_coo_row_thread_runtime_body_active(inner_builder: anytype, ctx: anytype) !void {
747 const sum = try spmvCooRowThreadValue(
748 inner_builder,
749 ctx.instance,
750 ctx.args,
751 ctx.row,
752 ctx.nnz,
753 ctx.x_extent,
754 );
755 const initial = try ctx.args.param(.y).load(inner_builder, ctx.row);
756 const initial_acc = try spmvCooAccumulationValue(inner_builder, ctx.instance, initial);
757 const total = try inner_builder.add(initial_acc, sum);
758 const result = try spmvCooOutputValue(inner_builder, ctx.instance, total);
759 try ctx.args.param(.y).store(inner_builder, result, ctx.row);
760 }
761
762 fn spmvCooRowThreadRuntimeBody(k: anytype, spec: SpmvCoo, args: anytype) !void {
763 if (!spmvCooInstanceValid(spec)) return error.UnsupportedSpmvCooInstance;
764 const row = try k.globalId(.x);
765 const rows_extent = try k.castIndex(args.param(.rows).raw());
766 const nnz = try k.castIndex(args.param(.nnz).raw());
767 const x_extent = try k.castIndex(args.param(.x_extent).raw());
768 const active = try k.compare(.lt, row, rows_extent);
769 try k.guardDo(active, .{
770 .args = args,
771 .row = row,
772 .nnz = nnz,
773 .x_extent = x_extent,
774 .instance = spec,
775 }, spmv_coo_row_thread_runtime_body_active);
776 }
777
778 fn spmvCooRuntimeBody(k: anytype, spec: SpmvCoo, args: anytype) !void {
779 return switch (spec.structure) {
780 .element_thread => spmvCooElementThreadRuntimeBody(k, spec, args),
781 .row_thread => spmvCooRowThreadRuntimeBody(k, spec, args),
782 };
783 }
784
785 fn spmvCooFamilySchedule(instance: SpmvCoo) kernel.logical.schedule.ThreadBlocks {
786 return kernel.logical.schedule.threadBlocks(.{ .x = instance.threads });
787 }
788
789 fn spmvCooRuntimeFamily(comptime dtype: DType) type {
790 return kernel.logical.Family(.{
791 .name = std.fmt.comptimePrint("accy_kernel_sparse_spmv_coo_runtime_{s}", .{dtype.name()}),
792 .parameters = .{
793 .y = kernel.dynamicBuffer(dtype),
794 .row_indices = kernel.dynamicBuffer(.i32),
795 .cols = kernel.dynamicBuffer(.i32),
796 .values = kernel.dynamicBuffer(dtype),
797 .x = kernel.dynamicBuffer(dtype),
798 .rows = kernel.scalar(.i32),
799 .nnz = kernel.scalar(.i32),
800 .x_extent = kernel.scalar(.i32),
801 },
802 .Instance = SpmvCoo,
803 .schedule = spmvCooFamilySchedule,
804 .body = spmvCooRuntimeBody,
805 });
806 }
807
808 pub const SpmvCooRuntimeFamilyF64 = spmvCooRuntimeFamily(.f64);
809 pub const SpmvCooRuntimeFamilyF32 = spmvCooRuntimeFamily(.f32);
810 pub const SpmvCooRuntimeFamilyF16 = spmvCooRuntimeFamily(.f16);
811
812 pub fn spmvCooFamilyTarget(allocator: std.mem.Allocator, instance: SpmvCoo) ![]u8 {
813 return std.fmt.allocPrint(
814 allocator,
815 "accy.kernel.sparse.spmv_coo_{s}_family_{d}_{s}",
816 .{ @tagName(instance.structure), instance.threads, instance.dtype.name() },
817 );
818 }
819
820 pub fn spmvCooFamilyEntryName(allocator: std.mem.Allocator, instance: SpmvCoo) ![]u8 {
821 return std.fmt.allocPrint(
822 allocator,
823 "accy_kernel_sparse_spmv_coo_{s}_family_{d}_{s}",
824 .{ @tagName(instance.structure), instance.threads, instance.dtype.name() },
825 );
826 }
827
828 pub fn spmvCooRuntimeArguments(
829 instance: SpmvCoo,
830 nnz: u64,
831 x_extent: u64,
832 ) ![3]choir_abi.ScalarArgument {
833 return .{
834 .{ .u32 = try runtimeExtentArgument(instance.rows) },
835 .{ .u32 = try runtimeExtentArgument(nnz) },
836 .{ .u32 = try runtimeExtentArgument(x_extent) },
837 };
838 }
839
840 pub fn spmvCooMaxRows(instance: SpmvCoo) u64 {
841 _ = instance;
842 return extent_mod.runtime_extent_max;
843 }
844
845 pub fn spmvCooMaxNnz(instance: SpmvCoo) u64 {
846 _ = instance;
847 return extent_mod.runtime_extent_max;
848 }
849
850 pub fn spmvCooMaxXExtent(instance: SpmvCoo) u64 {
851 _ = instance;
852 return extent_mod.runtime_extent_max;
853 }
854
855 pub fn spmvCooShapeProfileDimensions(instance: SpmvCoo) [2]artifact_product.KernelCallShapeProfileDimension {
856 return .{
857 .{
858 .name = instance.row_axis,
859 .runtime_scalar_argument_index = 0,
860 .bounds = .{ .min = 1, .max = spmvCooMaxRows(instance) },
861 },
862 .{
863 .name = instance.nonzero_axis,
864 .runtime_scalar_argument_index = 1,
865 .bounds = .{ .min = 1, .max = spmvCooMaxNnz(instance) },
866 },
867 };
868 }
869
870 fn spmvCooLaunch(instance: SpmvCoo) !artifact_product.KernelCallLaunch {
871 if (!spmvCooInstanceValid(instance)) return error.UnsupportedSpmvCooInstance;
872 return .{ .derived = .{
873 .grid = .{
874 .{ .runtime_u32_ceil_div = .{ .argument_index = spmvCooLaunchArgumentIndex(instance), .divisor = instance.threads } },
875 .{ .fixed = 1 },
876 .{ .fixed = 1 },
877 },
878 .threadgroup = .{ instance.threads, 1, 1 },
879 } };
880 }
881
882 pub fn spmvCooShapeFamily(backing_allocator: std.mem.Allocator, instance: SpmvCoo) !shape.Family {
883 var builder = try shape.Builder.init(backing_allocator, "spmv_coo");
884 errdefer builder.deinit();
885 const rows = try builder.symbol(instance.row_axis);
886 const nnz = try builder.symbol(instance.nonzero_axis);
887 const x_extent = try builder.symbol(instance.x_axis);
888 const rows_expr = try builder.symbolExpression(rows);
889 const nnz_expr = try builder.symbolExpression(nnz);
890 const x_extent_expr = try builder.symbolExpression(x_extent);
891 _ = try builder.tensor("row_indices", &.{nnz_expr});
892 _ = try builder.tensor("cols", &.{nnz_expr});
893 _ = try builder.tensor("values", &.{nnz_expr});
894 _ = try builder.tensor("x", &.{x_extent_expr});
895 _ = try builder.tensor("y", &.{rows_expr});
896 const bounds = shape.Bounds{ .min = 1, .max = extent_mod.runtime_extent_max };
897 try builder.assumeBounds(rows_expr, bounds);
898 try builder.assumeBounds(nnz_expr, bounds);
899 try builder.assumeBounds(x_extent_expr, bounds);
900 return builder.finish();
901 }
902
903 pub fn spmvCooFamilyFingerprint(backing_allocator: std.mem.Allocator, instance: SpmvCoo) !u64 {
904 var family = try spmvCooShapeFamily(backing_allocator, instance);
905 defer family.deinit();
906 return shape.fingerprint(family);
907 }
908
909 pub fn spmvCooTuningExtents(instance: SpmvCoo) [3]u64 {
910 return .{ instance.rows, instance.nnz, instance.x_extent };
911 }
912
913 pub fn spmvCooTuningOperation(instance: SpmvCoo) entry.Operation {
914 _ = instance;
915 return .{ .sparse = .coo_spmv };
916 }
917
918 pub fn spmvCooFamilyTuningKey(
919 backing_allocator: std.mem.Allocator,
920 device_fingerprint: u64,
921 instance: SpmvCoo,
922 ) !tuning.FamilyTuningKey {
923 const family_fingerprint = try spmvCooFamilyFingerprint(backing_allocator, instance);
924 const extents = spmvCooTuningExtents(instance);
925 return tuning.FamilyTuningKey.init(
926 device_fingerprint,
927 family_fingerprint,
928 entry.operationFingerprint(spmvCooTuningOperation(instance)),
929 instance.dtype,
930 spmv_coo_family_version,
931 extents[0..],
932 ) orelse unreachable;
933 }
934
935 pub fn resolveSpmvCooStructure(
936 backing_allocator: std.mem.Allocator,
937 reader: tuning.FamilyTuningReader,
938 instance: SpmvCoo,
939 ) !?SpmvCooStructure {
940 const key = try spmvCooFamilyTuningKey(backing_allocator, reader.device_fingerprint, instance);
941 const record = reader.table.find(key) orelse return null;
942 const structures = [_]SpmvCooStructure{ .element_thread, .row_thread };
943 for (structures) |structure| {
944 var candidate = instance;
945 candidate.structure = structure;
946 candidate.accumulation_dtype = spmvCooAccumulationDTypeForStructure(structure, instance.dtype) orelse continue;
947 candidate.threads = spmvCooRepresentableThreads(candidate) orelse continue;
948 if (!spmvCooInstanceValid(candidate)) continue;
949 const target = try spmvCooFamilyTarget(backing_allocator, candidate);
950 defer backing_allocator.free(target);
951 if (std.mem.eql(u8, target, record.target)) return structure;
952 }
953 return null;
954 }
955
956 pub fn createSpmvCooFamilyArtifact(
957 allocator: std.mem.Allocator,
958 handle: kernel.BackendHandle,
959 instance: SpmvCoo,
960 options: entry.ArtifactOptions,
961 ) !kernel.OwnedKernelCallArtifact {
962 if (!spmvCooInstanceValid(instance)) return error.UnsupportedSpmvCooInstance;
963 const target = try spmvCooFamilyTarget(allocator, instance);
964 defer allocator.free(target);
965 const entry_name = try spmvCooFamilyEntryName(allocator, instance);
966 defer allocator.free(entry_name);
967 const family_fingerprint = options.shape_family_fingerprint orelse try spmvCooFamilyFingerprint(allocator, instance);
968 const shape_profile_dimensions = spmvCooShapeProfileDimensions(instance);
969 const shape_profile = options.shape_profile orelse artifact_product.KernelCallShapeProfile{
970 .name = "spmv_coo",
971 .fingerprint = family_fingerprint,
972 .dimensions = shape_profile_dimensions[0..],
973 };
974
975 var graph = switch (instance.dtype) {
976 .f64 => try SpmvCooRuntimeFamilyF64.buildNamed(allocator, options.limits, entry_name, instance),
977 .f32 => try SpmvCooRuntimeFamilyF32.buildNamed(allocator, options.limits, entry_name, instance),
978 .f16 => try SpmvCooRuntimeFamilyF16.buildNamed(allocator, options.limits, entry_name, instance),
979 else => return error.UnsupportedDType,
980 };
981 defer graph.deinit();
982 return kernel.createKernelCallArtifact(allocator, handle, &graph, .{
983 .target = target,
984 .version = spmv_coo_family_version,
985 .format = options.format,
986 .kernel_plan = options.kernel_plan,
987 .element_count_argument = options.element_count_argument,
988 .shape_family_fingerprint = family_fingerprint,
989 .shape_profile = shape_profile,
990 .launch = options.launch orelse try spmvCooLaunch(instance),
991 .runtime_scalar_argument_count = if (options.runtime_scalar_argument_count == 0) 3 else options.runtime_scalar_argument_count,
992 .static_arguments = options.static_arguments,
993 });
994 }
995
996 fn spmvEllBlockCountChecked(instance: SpmvEll) ?u64 {
997 if (instance.threads == 0) return null;
998 const biased = std.math.add(u64, instance.rows, instance.threads - 1) catch return null;
999 return biased / instance.threads;
1000 }
1001
1002 pub fn spmvEllBlockCount(instance: SpmvEll) u64 {
1003 return spmvEllBlockCountChecked(instance).?;
1004 }
1005
1006 fn spmvEllElementCountChecked(instance: SpmvEll) ?u64 {
1007 return std.math.mul(u64, instance.rows, instance.slots) catch null;
1008 }
1009
1010 pub fn spmvEllInstanceValid(instance: SpmvEll) bool {
1011 const accumulation_dtype = spmvEllAccumulationDType(instance.dtype) orelse return false;
1012 if (instance.accumulation_dtype != accumulation_dtype) return false;
1013 if (instance.rows == 0 or instance.slots == 0 or instance.x_extent == 0) return false;
1014 if (instance.threads == 0 or instance.threads > spmv_ell_max_threads) return false;
1015 if (spmvEllLaunchExtentChecked(instance) == null) return false;
1016 if (spmvEllBlockCountChecked(instance) == null) return false;
1017 if (spmvEllElementCountChecked(instance) == null) return false;
1018 return switch (instance.structure) {
1019 .row_thread => true,
1020 };
1021 }
1022
1023 fn spmvSellBlockCountChecked(instance: SpmvSell) ?u64 {
1024 if (instance.threads == 0) return null;
1025 const biased = std.math.add(u64, instance.rows, instance.threads - 1) catch return null;
1026 return biased / instance.threads;
1027 }
1028
1029 pub fn spmvSellBlockCount(instance: SpmvSell) u64 {
1030 return spmvSellBlockCountChecked(instance).?;
1031 }
1032
1033 fn spmvSellSliceCountChecked(instance: SpmvSell) ?u64 {
1034 if (instance.rows == 0 or instance.slice_size == 0) return null;
1035 const biased = std.math.add(u64, instance.rows, instance.slice_size - 1) catch return null;
1036 return biased / instance.slice_size;
1037 }
1038
1039 pub fn spmvSellSliceCount(instance: SpmvSell) u64 {
1040 return spmvSellSliceCountChecked(instance).?;
1041 }
1042
1043 pub fn spmvSellInstanceValid(instance: SpmvSell) bool {
1044 const accumulation_dtype = spmvSellAccumulationDType(instance.dtype) orelse return false;
1045 if (instance.accumulation_dtype != accumulation_dtype) return false;
1046 if (instance.rows == 0 or instance.slice_size == 0) return false;
1047 if (instance.values_size == 0 or instance.x_extent == 0) return false;
1048 if (instance.threads == 0 or instance.threads > spmv_sell_max_threads) return false;
1049 if (spmvSellLaunchExtentChecked(instance) == null) return false;
1050 if (spmvSellBlockCountChecked(instance) == null) return false;
1051 if (spmvSellSliceCountChecked(instance) == null) return false;
1052 return switch (instance.structure) {
1053 .row_thread => true,
1054 };
1055 }
1056
1057 fn spmvEllAccumulationZero(inner_builder: anytype, instance: SpmvEll) !kernel.Value {
1058 return sparseFloatAccumulationZero(inner_builder, instance.accumulation_dtype);
1059 }
1060
1061 fn spmvEllAccumulationValue(inner_builder: anytype, instance: SpmvEll, value: anytype) !kernel.Value {
1062 return sparseFloatAccumulationValue(inner_builder, instance.accumulation_dtype, value);
1063 }
1064
1065 fn spmvEllOutputValue(inner_builder: anytype, instance: SpmvEll, value: kernel.Value) !kernel.Value {
1066 return sparseFloatOutputValue(inner_builder, instance.dtype, instance.accumulation_dtype, value);
1067 }
1068
1069 fn spmv_ell_row_thread_value_apply(fold_builder: anytype, slot: kernel.Value, current: kernel.Value, ctx: anytype) !kernel.Value {
1070 const slot_offset = try fold_builder.mul(slot, ctx.rows);
1071 const element = try fold_builder.add(slot_offset, ctx.row);
1072 const column_loaded = try ctx.args.param(.cols).load(fold_builder, element);
1073 const active = try fold_builder.compare(.ge, column_loaded.raw(), try fold_builder.constantInt(.i32, 0));
1074 const column_index = try fold_builder.castIndex(column_loaded.raw());
1075 const column_lower = try fold_builder.max(column_index, ctx.zero);
1076 const x_last = try fold_builder.sub(ctx.x_extent, ctx.one);
1077 const column = try fold_builder.min(column_lower, x_last);
1078 const x_value = try ctx.args.param(.x).load(fold_builder, column);
1079 const matrix_value = try ctx.args.param(.values).load(fold_builder, element);
1080 const matrix_acc = try spmvEllAccumulationValue(fold_builder, ctx.instance, matrix_value);
1081 const x_acc = try spmvEllAccumulationValue(fold_builder, ctx.instance, x_value);
1082 const product = try fold_builder.mul(matrix_acc, x_acc);
1083 const next = try fold_builder.add(current, product);
1084 return fold_builder.select(active, next, current);
1085 }
1086
1087 fn spmvEllRowThreadValue(
1088 inner_builder: anytype,
1089 instance: SpmvEll,
1090 args: anytype,
1091 row: kernel.Value,
1092 rows: kernel.Value,
1093 slots: kernel.Value,
1094 x_extent: kernel.Value,
1095 ) !kernel.Value {
1096 const zero = try inner_builder.constantIndex(0);
1097 const one = try inner_builder.constantIndex(1);
1098 const acc_zero = try spmvEllAccumulationZero(inner_builder, instance);
1099 return inner_builder.fold(zero, slots, one, acc_zero, .{
1100 .instance = instance,
1101 .args = args,
1102 .row = row,
1103 .rows = rows,
1104 .x_extent = x_extent,
1105 .zero = zero,
1106 .one = one,
1107 }, spmv_ell_row_thread_value_apply);
1108 }
1109
1110 fn spmv_ell_runtime_body_active(inner_builder: anytype, ctx: anytype) !void {
1111 const sum = try spmvEllRowThreadValue(
1112 inner_builder,
1113 ctx.instance,
1114 ctx.args,
1115 ctx.row,
1116 ctx.rows,
1117 ctx.slots,
1118 ctx.x_extent,
1119 );
1120 const result = try spmvEllOutputValue(inner_builder, ctx.instance, sum);
1121 try ctx.args.param(.y).store(inner_builder, result, ctx.row);
1122 }
1123
1124 fn spmvEllRuntimeBody(k: anytype, spec: SpmvEll, args: anytype) !void {
1125 if (!spmvEllInstanceValid(spec)) return error.UnsupportedSpmvEllInstance;
1126 const row = try k.globalId(.x);
1127 const rows_extent = try k.castIndex(args.param(.rows).raw());
1128 const slots = try k.castIndex(args.param(.slots).raw());
1129 const x_extent = try k.castIndex(args.param(.x_extent).raw());
1130 const active = try k.compare(.lt, row, rows_extent);
1131 try k.guardDo(active, .{
1132 .args = args,
1133 .row = row,
1134 .rows = rows_extent,
1135 .slots = slots,
1136 .x_extent = x_extent,
1137 .instance = spec,
1138 }, spmv_ell_runtime_body_active);
1139 }
1140
1141 fn spmvEllFamilySchedule(instance: SpmvEll) kernel.logical.schedule.ThreadBlocks {
1142 return kernel.logical.schedule.threadBlocks(.{ .x = instance.threads });
1143 }
1144
1145 fn spmvEllRuntimeFamily(comptime dtype: DType) type {
1146 return kernel.logical.Family(.{
1147 .name = std.fmt.comptimePrint("accy_kernel_sparse_spmv_ell_runtime_{s}", .{dtype.name()}),
1148 .parameters = .{
1149 .y = kernel.dynamicBuffer(dtype),
1150 .cols = kernel.dynamicBuffer(.i32),
1151 .values = kernel.dynamicBuffer(dtype),
1152 .x = kernel.dynamicBuffer(dtype),
1153 .rows = kernel.scalar(.i32),
1154 .slots = kernel.scalar(.i32),
1155 .x_extent = kernel.scalar(.i32),
1156 },
1157 .Instance = SpmvEll,
1158 .schedule = spmvEllFamilySchedule,
1159 .body = spmvEllRuntimeBody,
1160 });
1161 }
1162
1163 pub const SpmvEllRuntimeFamilyF64 = spmvEllRuntimeFamily(.f64);
1164 pub const SpmvEllRuntimeFamilyF32 = spmvEllRuntimeFamily(.f32);
1165 pub const SpmvEllRuntimeFamilyF16 = spmvEllRuntimeFamily(.f16);
1166
1167 pub fn spmvEllFamilyTarget(allocator: std.mem.Allocator, instance: SpmvEll) ![]u8 {
1168 return std.fmt.allocPrint(
1169 allocator,
1170 "accy.kernel.sparse.spmv_ell_{s}_family_{d}_{s}",
1171 .{ @tagName(instance.structure), instance.threads, instance.dtype.name() },
1172 );
1173 }
1174
1175 pub fn spmvEllFamilyEntryName(allocator: std.mem.Allocator, instance: SpmvEll) ![]u8 {
1176 return std.fmt.allocPrint(
1177 allocator,
1178 "accy_kernel_sparse_spmv_ell_{s}_family_{d}_{s}",
1179 .{ @tagName(instance.structure), instance.threads, instance.dtype.name() },
1180 );
1181 }
1182
1183 pub fn spmvEllRuntimeArguments(
1184 instance: SpmvEll,
1185 slots: u64,
1186 x_extent: u64,
1187 ) ![3]choir_abi.ScalarArgument {
1188 return .{
1189 .{ .u32 = try runtimeExtentArgument(instance.rows) },
1190 .{ .u32 = try runtimeExtentArgument(slots) },
1191 .{ .u32 = try runtimeExtentArgument(x_extent) },
1192 };
1193 }
1194
1195 pub fn spmvEllMaxRows(instance: SpmvEll) u64 {
1196 _ = instance;
1197 return extent_mod.runtime_extent_max;
1198 }
1199
1200 pub fn spmvEllMaxSlots(instance: SpmvEll) u64 {
1201 _ = instance;
1202 return extent_mod.runtime_extent_max;
1203 }
1204
1205 pub fn spmvEllMaxXExtent(instance: SpmvEll) u64 {
1206 _ = instance;
1207 return extent_mod.runtime_extent_max;
1208 }
1209
1210 pub fn spmvEllShapeProfileDimensions(instance: SpmvEll) [1]artifact_product.KernelCallShapeProfileDimension {
1211 return .{
1212 .{
1213 .name = instance.row_axis,
1214 .runtime_scalar_argument_index = 0,
1215 .bounds = .{ .min = 1, .max = spmvEllMaxRows(instance) },
1216 },
1217 };
1218 }
1219
1220 fn spmvEllLaunch(instance: SpmvEll) !artifact_product.KernelCallLaunch {
1221 if (!spmvEllInstanceValid(instance)) return error.UnsupportedSpmvEllInstance;
1222 return .{ .derived = .{
1223 .grid = .{
1224 .{ .runtime_u32_ceil_div = .{ .argument_index = 0, .divisor = instance.threads } },
1225 .{ .fixed = 1 },
1226 .{ .fixed = 1 },
1227 },
1228 .threadgroup = .{ instance.threads, 1, 1 },
1229 } };
1230 }
1231
1232 pub fn spmvEllShapeFamily(backing_allocator: std.mem.Allocator, instance: SpmvEll) !shape.Family {
1233 var builder = try shape.Builder.init(backing_allocator, "spmv_ell");
1234 errdefer builder.deinit();
1235 const rows = try builder.symbol(instance.row_axis);
1236 const slots = try builder.symbol(instance.slot_axis);
1237 const x_extent = try builder.symbol(instance.x_axis);
1238 const rows_expr = try builder.symbolExpression(rows);
1239 const slots_expr = try builder.symbolExpression(slots);
1240 const x_extent_expr = try builder.symbolExpression(x_extent);
1241 _ = try builder.tensor("cols", &.{ slots_expr, rows_expr });
1242 _ = try builder.tensor("values", &.{ slots_expr, rows_expr });
1243 _ = try builder.tensor("x", &.{x_extent_expr});
1244 _ = try builder.tensor("y", &.{rows_expr});
1245 try builder.assumeBounds(rows_expr, .{ .min = 1, .max = spmvEllMaxRows(instance) });
1246 try builder.assumeBounds(slots_expr, .{ .min = 1, .max = spmvEllMaxSlots(instance) });
1247 try builder.assumeBounds(x_extent_expr, .{ .min = 1, .max = spmvEllMaxXExtent(instance) });
1248 return builder.finish();
1249 }
1250
1251 pub fn spmvEllFamilyFingerprint(backing_allocator: std.mem.Allocator, instance: SpmvEll) !u64 {
1252 var family = try spmvEllShapeFamily(backing_allocator, instance);
1253 defer family.deinit();
1254 return shape.fingerprint(family);
1255 }
1256
1257 pub fn spmvEllTuningExtents(instance: SpmvEll) [3]u64 {
1258 return .{ instance.rows, instance.slots, instance.x_extent };
1259 }
1260
1261 pub fn spmvEllTuningOperation(instance: SpmvEll) entry.Operation {
1262 _ = instance;
1263 return .{ .sparse = .ell_spmv };
1264 }
1265
1266 pub fn spmvEllFamilyTuningKey(
1267 backing_allocator: std.mem.Allocator,
1268 device_fingerprint: u64,
1269 instance: SpmvEll,
1270 ) !tuning.FamilyTuningKey {
1271 const family_fingerprint = try spmvEllFamilyFingerprint(backing_allocator, instance);
1272 const extents = spmvEllTuningExtents(instance);
1273 return tuning.FamilyTuningKey.init(
1274 device_fingerprint,
1275 family_fingerprint,
1276 entry.operationFingerprint(spmvEllTuningOperation(instance)),
1277 instance.dtype,
1278 spmv_ell_family_version,
1279 extents[0..],
1280 ) orelse unreachable;
1281 }
1282
1283 pub fn resolveSpmvEllThreads(
1284 backing_allocator: std.mem.Allocator,
1285 reader: tuning.FamilyTuningReader,
1286 instance: SpmvEll,
1287 ) !?u32 {
1288 const key = try spmvEllFamilyTuningKey(backing_allocator, reader.device_fingerprint, instance);
1289 const record = reader.table.find(key) orelse return null;
1290 var default_candidate = instance;
1291 default_candidate.threads = spmvEllRepresentableThreads(default_candidate) orelse return null;
1292 if (spmvEllInstanceValid(default_candidate)) {
1293 const target = try spmvEllFamilyTarget(backing_allocator, default_candidate);
1294 defer backing_allocator.free(target);
1295 if (std.mem.eql(u8, target, record.target)) return default_candidate.threads;
1296 }
1297 const thread_candidates = spmvEllThreadCandidatesForRows(instance.rows);
1298 for (thread_candidates.slice()) |threads| {
1299 var candidate = instance;
1300 candidate.threads = threads;
1301 candidate.threads = spmvEllRepresentableThreads(candidate) orelse continue;
1302 if (!spmvEllInstanceValid(candidate)) continue;
1303 const target = try spmvEllFamilyTarget(backing_allocator, candidate);
1304 defer backing_allocator.free(target);
1305 if (std.mem.eql(u8, target, record.target)) return candidate.threads;
1306 }
1307 return null;
1308 }
1309
1310 pub fn createSpmvEllFamilyArtifact(
1311 allocator: std.mem.Allocator,
1312 handle: kernel.BackendHandle,
1313 instance: SpmvEll,
1314 options: entry.ArtifactOptions,
1315 ) !kernel.OwnedKernelCallArtifact {
1316 if (!spmvEllInstanceValid(instance)) return error.UnsupportedSpmvEllInstance;
1317 const target = try spmvEllFamilyTarget(allocator, instance);
1318 defer allocator.free(target);
1319 const entry_name = try spmvEllFamilyEntryName(allocator, instance);
1320 defer allocator.free(entry_name);
1321 const family_fingerprint = options.shape_family_fingerprint orelse try spmvEllFamilyFingerprint(allocator, instance);
1322 const shape_profile_dimensions = spmvEllShapeProfileDimensions(instance);
1323 const shape_profile = options.shape_profile orelse artifact_product.KernelCallShapeProfile{
1324 .name = "spmv_ell",
1325 .fingerprint = family_fingerprint,
1326 .dimensions = shape_profile_dimensions[0..],
1327 };
1328
1329 var graph = switch (instance.dtype) {
1330 .f64 => try SpmvEllRuntimeFamilyF64.buildNamed(allocator, options.limits, entry_name, instance),
1331 .f32 => try SpmvEllRuntimeFamilyF32.buildNamed(allocator, options.limits, entry_name, instance),
1332 .f16 => try SpmvEllRuntimeFamilyF16.buildNamed(allocator, options.limits, entry_name, instance),
1333 else => return error.UnsupportedDType,
1334 };
1335 defer graph.deinit();
1336 return kernel.createKernelCallArtifact(allocator, handle, &graph, .{
1337 .target = target,
1338 .version = spmv_ell_family_version,
1339 .format = options.format,
1340 .kernel_plan = options.kernel_plan,
1341 .element_count_argument = options.element_count_argument,
1342 .shape_family_fingerprint = family_fingerprint,
1343 .shape_profile = shape_profile,
1344 .launch = options.launch orelse try spmvEllLaunch(instance),
1345 .runtime_scalar_argument_count = if (options.runtime_scalar_argument_count == 0) 3 else options.runtime_scalar_argument_count,
1346 .static_arguments = options.static_arguments,
1347 });
1348 }
1349
1350 fn spmvSellAccumulationZero(inner_builder: anytype, instance: SpmvSell) !kernel.Value {
1351 return sparseFloatAccumulationZero(inner_builder, instance.accumulation_dtype);
1352 }
1353
1354 fn spmvSellAccumulationValue(inner_builder: anytype, instance: SpmvSell, value: anytype) !kernel.Value {
1355 return sparseFloatAccumulationValue(inner_builder, instance.accumulation_dtype, value);
1356 }
1357
1358 fn spmvSellOutputValue(inner_builder: anytype, instance: SpmvSell, value: kernel.Value) !kernel.Value {
1359 return sparseFloatOutputValue(inner_builder, instance.dtype, instance.accumulation_dtype, value);
1360 }
1361
1362 fn spmv_sell_row_thread_value_apply(fold_builder: anytype, element: kernel.Value, current: kernel.Value, ctx: anytype) !kernel.Value {
1363 const column_loaded = try ctx.args.param(.cols).load(fold_builder, element);
1364 const active = try fold_builder.compare(.ge, column_loaded.raw(), try fold_builder.constantInt(.i32, 0));
1365 const column_index = try fold_builder.castIndex(column_loaded.raw());
1366 const column_lower = try fold_builder.max(column_index, ctx.zero);
1367 const x_last = try fold_builder.sub(ctx.x_extent, ctx.one);
1368 const column = try fold_builder.min(column_lower, x_last);
1369 const x_value = try ctx.args.param(.x).load(fold_builder, column);
1370 const matrix_value = try ctx.args.param(.values).load(fold_builder, element);
1371 const matrix_acc = try spmvSellAccumulationValue(fold_builder, ctx.instance, matrix_value);
1372 const x_acc = try spmvSellAccumulationValue(fold_builder, ctx.instance, x_value);
1373 const product = try fold_builder.mul(matrix_acc, x_acc);
1374 const next = try fold_builder.add(current, product);
1375 return fold_builder.select(active, next, current);
1376 }
1377
1378 fn spmvSellRowThreadValue(
1379 inner_builder: anytype,
1380 instance: SpmvSell,
1381 args: anytype,
1382 row: kernel.Value,
1383 values_size: kernel.Value,
1384 x_extent: kernel.Value,
1385 ) !kernel.Value {
1386 const zero = try inner_builder.constantIndex(0);
1387 const one = try inner_builder.constantIndex(1);
1388 const slice_size = try inner_builder.constantIndex(@intCast(instance.slice_size));
1389 const slice = try inner_builder.div(row, slice_size);
1390 const slice_base = try inner_builder.mul(slice, slice_size);
1391 const local = try inner_builder.sub(row, slice_base);
1392 const next_slice = try inner_builder.add(slice, one);
1393 const begin_loaded = try args.param(.slice_offsets).load(inner_builder, slice);
1394 const end_loaded = try args.param(.slice_offsets).load(inner_builder, next_slice);
1395 const begin_index = try inner_builder.castIndex(begin_loaded.raw());
1396 const end_index = try inner_builder.castIndex(end_loaded.raw());
1397 const end_lower = try inner_builder.max(end_index, zero);
1398 const end_clamped = try inner_builder.min(end_lower, values_size);
1399 const begin_lower = try inner_builder.max(begin_index, zero);
1400 const begin_clamped = try inner_builder.min(begin_lower, end_clamped);
1401 const row_begin_unclamped = try inner_builder.add(begin_clamped, local);
1402 const row_begin = try inner_builder.min(row_begin_unclamped, end_clamped);
1403
1404 const acc_zero = try spmvSellAccumulationZero(inner_builder, instance);
1405 return inner_builder.fold(row_begin, end_clamped, slice_size, acc_zero, .{
1406 .instance = instance,
1407 .args = args,
1408 .x_extent = x_extent,
1409 .zero = zero,
1410 .one = one,
1411 }, spmv_sell_row_thread_value_apply);
1412 }
1413
1414 fn spmv_sell_runtime_body_active(inner_builder: anytype, ctx: anytype) !void {
1415 const one = try inner_builder.constantIndex(1);
1416 const bounded_x = try inner_builder.max(ctx.x_extent, one);
1417 const sum = try spmvSellRowThreadValue(
1418 inner_builder,
1419 ctx.instance,
1420 ctx.args,
1421 ctx.row,
1422 ctx.values_size,
1423 bounded_x,
1424 );
1425 const result = try spmvSellOutputValue(inner_builder, ctx.instance, sum);
1426 try ctx.args.param(.y).store(inner_builder, result, ctx.row);
1427 }
1428
1429 fn spmvSellRuntimeBody(k: anytype, spec: SpmvSell, args: anytype) !void {
1430 if (!spmvSellInstanceValid(spec)) return error.UnsupportedSpmvSellInstance;
1431 const row = try k.globalId(.x);
1432 const rows_extent = try k.castIndex(args.param(.rows).raw());
1433 const values_size = try k.castIndex(args.param(.values_size).raw());
1434 const x_extent = try k.castIndex(args.param(.x_extent).raw());
1435 const active = try k.compare(.lt, row, rows_extent);
1436 try k.guardDo(active, .{
1437 .args = args,
1438 .row = row,
1439 .values_size = values_size,
1440 .x_extent = x_extent,
1441 .instance = spec,
1442 }, spmv_sell_runtime_body_active);
1443 }
1444
1445 fn spmvSellFamilySchedule(instance: SpmvSell) kernel.logical.schedule.ThreadBlocks {
1446 return kernel.logical.schedule.threadBlocks(.{ .x = instance.threads });
1447 }
1448
1449 fn spmvSellRuntimeFamily(comptime dtype: DType) type {
1450 return kernel.logical.Family(.{
1451 .name = std.fmt.comptimePrint("accy_kernel_sparse_spmv_sell_runtime_{s}", .{dtype.name()}),
1452 .parameters = .{
1453 .y = kernel.dynamicBuffer(dtype),
1454 .slice_offsets = kernel.dynamicBuffer(.i32),
1455 .cols = kernel.dynamicBuffer(.i32),
1456 .values = kernel.dynamicBuffer(dtype),
1457 .x = kernel.dynamicBuffer(dtype),
1458 .rows = kernel.scalar(.i32),
1459 .values_size = kernel.scalar(.i32),
1460 .x_extent = kernel.scalar(.i32),
1461 },
1462 .Instance = SpmvSell,
1463 .schedule = spmvSellFamilySchedule,
1464 .body = spmvSellRuntimeBody,
1465 });
1466 }
1467
1468 pub const SpmvSellRuntimeFamilyF64 = spmvSellRuntimeFamily(.f64);
1469 pub const SpmvSellRuntimeFamilyF32 = spmvSellRuntimeFamily(.f32);
1470 pub const SpmvSellRuntimeFamilyF16 = spmvSellRuntimeFamily(.f16);
1471
1472 pub fn spmvSellFamilyTarget(allocator: std.mem.Allocator, instance: SpmvSell) ![]u8 {
1473 return std.fmt.allocPrint(
1474 allocator,
1475 "accy.kernel.sparse.spmv_sell_{s}_slice{d}_family_{d}_{s}",
1476 .{ @tagName(instance.structure), instance.slice_size, instance.threads, instance.dtype.name() },
1477 );
1478 }
1479
1480 pub fn spmvSellFamilyEntryName(allocator: std.mem.Allocator, instance: SpmvSell) ![]u8 {
1481 return std.fmt.allocPrint(
1482 allocator,
1483 "accy_kernel_sparse_spmv_sell_{s}_slice{d}_family_{d}_{s}",
1484 .{ @tagName(instance.structure), instance.slice_size, instance.threads, instance.dtype.name() },
1485 );
1486 }
1487
1488 pub fn spmvSellRuntimeArguments(
1489 instance: SpmvSell,
1490 values_size: u64,
1491 x_extent: u64,
1492 ) ![3]choir_abi.ScalarArgument {
1493 return .{
1494 .{ .u32 = try runtimeExtentArgument(instance.rows) },
1495 .{ .u32 = try runtimeExtentArgument(values_size) },
1496 .{ .u32 = try runtimeExtentArgument(x_extent) },
1497 };
1498 }
1499
1500 pub fn spmvSellMaxRows(instance: SpmvSell) u64 {
1501 _ = instance;
1502 return extent_mod.runtime_extent_max;
1503 }
1504
1505 pub fn spmvSellMaxValuesSize(instance: SpmvSell) u64 {
1506 _ = instance;
1507 return extent_mod.runtime_extent_max;
1508 }
1509
1510 pub fn spmvSellMaxXExtent(instance: SpmvSell) u64 {
1511 _ = instance;
1512 return extent_mod.runtime_extent_max;
1513 }
1514
1515 pub fn spmvSellShapeProfileDimensions(instance: SpmvSell) [1]artifact_product.KernelCallShapeProfileDimension {
1516 return .{
1517 .{
1518 .name = instance.row_axis,
1519 .runtime_scalar_argument_index = 0,
1520 .bounds = .{ .min = 1, .max = spmvSellMaxRows(instance) },
1521 },
1522 };
1523 }
1524
1525 fn spmvSellLaunch(instance: SpmvSell) !artifact_product.KernelCallLaunch {
1526 if (!spmvSellInstanceValid(instance)) return error.UnsupportedSpmvSellInstance;
1527 return .{ .derived = .{
1528 .grid = .{
1529 .{ .runtime_u32_ceil_div = .{ .argument_index = 0, .divisor = instance.threads } },
1530 .{ .fixed = 1 },
1531 .{ .fixed = 1 },
1532 },
1533 .threadgroup = .{ instance.threads, 1, 1 },
1534 } };
1535 }
1536
1537 pub fn spmvSellShapeFamily(backing_allocator: std.mem.Allocator, instance: SpmvSell) !shape.Family {
1538 var builder = try shape.Builder.init(backing_allocator, "spmv_sell");
1539 errdefer builder.deinit();
1540 const rows = try builder.symbol(instance.row_axis);
1541 const slices = try builder.symbol(instance.slice_axis);
1542 const values_size = try builder.symbol(instance.value_axis);
1543 const x_extent = try builder.symbol(instance.x_axis);
1544 const rows_expr = try builder.symbolExpression(rows);
1545 const slices_expr = try builder.symbolExpression(slices);
1546 const slice_offsets_expr = try builder.addExpression(slices_expr, builder.constantExpression(1));
1547 const values_size_expr = try builder.symbolExpression(values_size);
1548 const x_extent_expr = try builder.symbolExpression(x_extent);
1549 _ = try builder.tensor("slice_offsets", &.{slice_offsets_expr});
1550 _ = try builder.tensor("cols", &.{values_size_expr});
1551 _ = try builder.tensor("values", &.{values_size_expr});
1552 _ = try builder.tensor("x", &.{x_extent_expr});
1553 _ = try builder.tensor("y", &.{rows_expr});
1554 const bounds = shape.Bounds{ .min = 1, .max = extent_mod.runtime_extent_max };
1555 try builder.assumeBounds(rows_expr, bounds);
1556 try builder.assumeBounds(slices_expr, bounds);
1557 try builder.assumeBounds(values_size_expr, bounds);
1558 try builder.assumeBounds(x_extent_expr, bounds);
1559 return builder.finish();
1560 }
1561
1562 pub fn spmvSellFamilyFingerprint(backing_allocator: std.mem.Allocator, instance: SpmvSell) !u64 {
1563 var family = try spmvSellShapeFamily(backing_allocator, instance);
1564 defer family.deinit();
1565 return shape.fingerprint(family);
1566 }
1567
1568 pub fn spmvSellTuningExtents(instance: SpmvSell) [4]u64 {
1569 return .{ instance.rows, instance.slice_size, instance.values_size, instance.x_extent };
1570 }
1571
1572 pub fn spmvSellTuningOperation(instance: SpmvSell) entry.Operation {
1573 _ = instance;
1574 return .{ .sparse = .sell_spmv };
1575 }
1576
1577 pub fn spmvSellFamilyTuningKey(
1578 backing_allocator: std.mem.Allocator,
1579 device_fingerprint: u64,
1580 instance: SpmvSell,
1581 ) !tuning.FamilyTuningKey {
1582 const family_fingerprint = try spmvSellFamilyFingerprint(backing_allocator, instance);
1583 const extents = spmvSellTuningExtents(instance);
1584 return tuning.FamilyTuningKey.init(
1585 device_fingerprint,
1586 family_fingerprint,
1587 entry.operationFingerprint(spmvSellTuningOperation(instance)),
1588 instance.dtype,
1589 spmv_sell_family_version,
1590 extents[0..],
1591 ) orelse unreachable;
1592 }
1593
1594 pub fn resolveSpmvSellThreads(
1595 backing_allocator: std.mem.Allocator,
1596 reader: tuning.FamilyTuningReader,
1597 instance: SpmvSell,
1598 ) !?u32 {
1599 const key = try spmvSellFamilyTuningKey(backing_allocator, reader.device_fingerprint, instance);
1600 const record = reader.table.find(key) orelse return null;
1601 var default_candidate = instance;
1602 default_candidate.threads = spmvSellRepresentableThreads(default_candidate) orelse return null;
1603 if (spmvSellInstanceValid(default_candidate)) {
1604 const target = try spmvSellFamilyTarget(backing_allocator, default_candidate);
1605 defer backing_allocator.free(target);
1606 if (std.mem.eql(u8, target, record.target)) return default_candidate.threads;
1607 }
1608 const thread_candidates = spmvSellThreadCandidatesForRows(instance.rows);
1609 for (thread_candidates.slice()) |threads| {
1610 var candidate = instance;
1611 candidate.threads = threads;
1612 candidate.threads = spmvSellRepresentableThreads(candidate) orelse continue;
1613 if (!spmvSellInstanceValid(candidate)) continue;
1614 const target = try spmvSellFamilyTarget(backing_allocator, candidate);
1615 defer backing_allocator.free(target);
1616 if (std.mem.eql(u8, target, record.target)) return candidate.threads;
1617 }
1618 return null;
1619 }
1620
1621 pub fn createSpmvSellFamilyArtifact(
1622 allocator: std.mem.Allocator,
1623 handle: kernel.BackendHandle,
1624 instance: SpmvSell,
1625 options: entry.ArtifactOptions,
1626 ) !kernel.OwnedKernelCallArtifact {
1627 if (!spmvSellInstanceValid(instance)) return error.UnsupportedSpmvSellInstance;
1628 const target = try spmvSellFamilyTarget(allocator, instance);
1629 defer allocator.free(target);
1630 const entry_name = try spmvSellFamilyEntryName(allocator, instance);
1631 defer allocator.free(entry_name);
1632 const family_fingerprint = options.shape_family_fingerprint orelse try spmvSellFamilyFingerprint(allocator, instance);
1633 const shape_profile_dimensions = spmvSellShapeProfileDimensions(instance);
1634 const shape_profile = options.shape_profile orelse artifact_product.KernelCallShapeProfile{
1635 .name = "spmv_sell",
1636 .fingerprint = family_fingerprint,
1637 .dimensions = shape_profile_dimensions[0..],
1638 };
1639
1640 var graph = switch (instance.dtype) {
1641 .f64 => try SpmvSellRuntimeFamilyF64.buildNamed(allocator, options.limits, entry_name, instance),
1642 .f32 => try SpmvSellRuntimeFamilyF32.buildNamed(allocator, options.limits, entry_name, instance),
1643 .f16 => try SpmvSellRuntimeFamilyF16.buildNamed(allocator, options.limits, entry_name, instance),
1644 else => return error.UnsupportedDType,
1645 };
1646 defer graph.deinit();
1647 return kernel.createKernelCallArtifact(allocator, handle, &graph, .{
1648 .target = target,
1649 .version = spmv_sell_family_version,
1650 .format = options.format,
1651 .kernel_plan = options.kernel_plan,
1652 .element_count_argument = options.element_count_argument,
1653 .shape_family_fingerprint = family_fingerprint,
1654 .shape_profile = shape_profile,
1655 .launch = options.launch orelse try spmvSellLaunch(instance),
1656 .runtime_scalar_argument_count = if (options.runtime_scalar_argument_count == 0) 3 else options.runtime_scalar_argument_count,
1657 .static_arguments = options.static_arguments,
1658 });
1659 }
1660
1661 fn spmmCsrThreadCountChecked(instance: SpmmCsr) ?u32 {
1662 if (instance.threads.x == 0 or instance.threads.y == 0) return null;
1663 const count = std.math.mul(u32, instance.threads.x, instance.threads.y) catch return null;
1664 if (count == 0 or count > spmm_csr_max_threads) return null;
1665 return count;
1666 }
1667
1668 fn spmmCsrBlockCountAxisChecked(extent: u64, threads: u32) ?u64 {
1669 if (extent == 0 or threads == 0) return null;
1670 const biased = std.math.add(u64, extent, threads - 1) catch return null;
1671 const blocks = biased / threads;
1672 if (blocks > std.math.maxInt(u32)) return null;
1673 return blocks;
1674 }
1675
1676 pub fn spmmCsrBlockCountX(instance: SpmmCsr) u64 {
1677 return spmmCsrBlockCountXChecked(instance).?;
1678 }
1679
1680 fn spmmCsrBlockCountXChecked(instance: SpmmCsr) ?u64 {
1681 return spmmCsrBlockCountAxisChecked(instance.columns, instance.threads.x);
1682 }
1683
1684 pub fn spmmCsrBlockCountY(instance: SpmmCsr) u64 {
1685 return spmmCsrBlockCountYChecked(instance).?;
1686 }
1687
1688 fn spmmCsrBlockCountYChecked(instance: SpmmCsr) ?u64 {
1689 return spmmCsrBlockCountAxisChecked(instance.rows, instance.threads.y);
1690 }
1691
1692 pub fn spmmCsrInstanceValid(instance: SpmmCsr) bool {
1693 const accumulation_dtype = spmmCsrAccumulationDType(instance.dtype) orelse return false;
1694 if (instance.accumulation_dtype != accumulation_dtype) return false;
1695 if (instance.rows == 0 or instance.columns == 0) return false;
1696 if (instance.nnz == 0 or instance.x_extent == 0) return false;
1697 if (spmmCsrThreadCountChecked(instance) == null) return false;
1698 if (spmmCsrBlockCountXChecked(instance) == null) return false;
1699 if (spmmCsrBlockCountYChecked(instance) == null) return false;
1700 return switch (instance.structure) {
1701 .row_column_thread => true,
1702 };
1703 }
1704
1705 fn spmmCsrAccumulationZero(inner_builder: anytype, instance: SpmmCsr) !kernel.Value {
1706 return sparseFloatAccumulationZero(inner_builder, instance.accumulation_dtype);
1707 }
1708
1709 fn spmmCsrAccumulationValue(inner_builder: anytype, instance: SpmmCsr, value: anytype) !kernel.Value {
1710 return sparseFloatAccumulationValue(inner_builder, instance.accumulation_dtype, value);
1711 }
1712
1713 fn spmmCsrOutputValue(inner_builder: anytype, instance: SpmmCsr, value: kernel.Value) !kernel.Value {
1714 return sparseFloatOutputValue(inner_builder, instance.dtype, instance.accumulation_dtype, value);
1715 }
1716
1717 fn spmmCsrOutputIndex(inner_builder: anytype, row: kernel.Value, column: kernel.Value, columns: kernel.Value) !kernel.Value {
1718 const row_offset = try inner_builder.mul(row, columns);
1719 return inner_builder.add(row_offset, column);
1720 }
1721
1722 fn spmm_csr_cell_value_apply(fold_builder: anytype, element: kernel.Value, current: kernel.Value, ctx: anytype) !kernel.Value {
1723 const column_loaded = try ctx.args.param(.cols).load(fold_builder, element);
1724 const column_index = try fold_builder.castIndex(column_loaded.raw());
1725 const column_lower = try fold_builder.max(column_index, ctx.zero);
1726 const x_last = try fold_builder.sub(ctx.x_extent, ctx.one);
1727 const sparse_column = try fold_builder.min(column_lower, x_last);
1728 const x_row_offset = try fold_builder.mul(sparse_column, ctx.columns);
1729 const x_index = try fold_builder.add(x_row_offset, ctx.column);
1730 const x_value = try ctx.args.param(.x).load(fold_builder, x_index);
1731 const matrix_value = try ctx.args.param(.values).load(fold_builder, element);
1732 const matrix_acc = try spmmCsrAccumulationValue(fold_builder, ctx.instance, matrix_value);
1733 const x_acc = try spmmCsrAccumulationValue(fold_builder, ctx.instance, x_value);
1734 const product = try fold_builder.mul(matrix_acc, x_acc);
1735 return fold_builder.add(current, product);
1736 }
1737
1738 fn spmmCsrCellValue(
1739 inner_builder: anytype,
1740 instance: SpmmCsr,
1741 args: anytype,
1742 row: kernel.Value,
1743 column: kernel.Value,
1744 nnz: kernel.Value,
1745 x_extent: kernel.Value,
1746 columns: kernel.Value,
1747 ) !kernel.Value {
1748 const one = try inner_builder.constantIndex(1);
1749 const next = try inner_builder.add(row, one);
1750 const begin_loaded = try args.param(.row_ptr).load(inner_builder, row);
1751 const end_loaded = try args.param(.row_ptr).load(inner_builder, next);
1752 const begin_index = try inner_builder.castIndex(begin_loaded.raw());
1753 const end_index = try inner_builder.castIndex(end_loaded.raw());
1754 const zero = try inner_builder.constantIndex(0);
1755 const end_lower = try inner_builder.max(end_index, zero);
1756 const end_clamped = try inner_builder.min(end_lower, nnz);
1757 const begin_lower = try inner_builder.max(begin_index, zero);
1758 const begin_clamped = try inner_builder.min(begin_lower, end_clamped);
1759
1760 const acc_zero = try spmmCsrAccumulationZero(inner_builder, instance);
1761 return inner_builder.fold(begin_clamped, end_clamped, one, acc_zero, .{
1762 .instance = instance,
1763 .args = args,
1764 .column = column,
1765 .columns = columns,
1766 .x_extent = x_extent,
1767 .zero = zero,
1768 .one = one,
1769 }, spmm_csr_cell_value_apply);
1770 }
1771
1772 fn spmm_csr_runtime_body_row_active(inner_builder: anytype, ctx: anytype) !void {
1773 const column_active = try inner_builder.compare(.lt, ctx.column, ctx.columns);
1774 try inner_builder.guardDo(column_active, ctx, spmm_csr_runtime_body_column_active);
1775 }
1776
1777 fn spmm_csr_runtime_body_column_active(active_builder: anytype, active_ctx: anytype) !void {
1778 const sum = try spmmCsrCellValue(
1779 active_builder,
1780 active_ctx.instance,
1781 active_ctx.args,
1782 active_ctx.row,
1783 active_ctx.column,
1784 active_ctx.nnz,
1785 active_ctx.x_extent,
1786 active_ctx.columns,
1787 );
1788 const result = try spmmCsrOutputValue(active_builder, active_ctx.instance, sum);
1789 const output_index = try spmmCsrOutputIndex(active_builder, active_ctx.row, active_ctx.column, active_ctx.columns);
1790 try active_ctx.args.param(.y).store(active_builder, result, output_index);
1791 }
1792
1793 fn spmmCsrRuntimeBody(k: anytype, spec: SpmmCsr, args: anytype) !void {
1794 if (!spmmCsrInstanceValid(spec)) return error.UnsupportedSpmmCsrInstance;
1795 const row = try k.globalId(.y);
1796 const column = try k.globalId(.x);
1797 const rows_extent = try k.castIndex(args.param(.rows).raw());
1798 const columns_extent = try k.castIndex(args.param(.columns).raw());
1799 const nnz = try k.castIndex(args.param(.nnz).raw());
1800 const x_extent = try k.castIndex(args.param(.x_extent).raw());
1801 const row_active = try k.compare(.lt, row, rows_extent);
1802 try k.guardDo(row_active, .{
1803 .args = args,
1804 .row = row,
1805 .column = column,
1806 .columns = columns_extent,
1807 .nnz = nnz,
1808 .x_extent = x_extent,
1809 .instance = spec,
1810 }, spmm_csr_runtime_body_row_active);
1811 }
1812
1813 fn spmmCsrFamilySchedule(instance: SpmmCsr) kernel.logical.schedule.ThreadBlocks {
1814 return kernel.logical.schedule.threadBlocks(.{ .x = instance.threads.x, .y = instance.threads.y });
1815 }
1816
1817 fn spmmCsrRuntimeFamily(comptime dtype: DType) type {
1818 return kernel.logical.Family(.{
1819 .name = std.fmt.comptimePrint("accy_kernel_sparse_spmm_csr_runtime_{s}", .{dtype.name()}),
1820 .parameters = .{
1821 .y = kernel.dynamicBuffer(dtype),
1822 .row_ptr = kernel.dynamicBuffer(.i32),
1823 .cols = kernel.dynamicBuffer(.i32),
1824 .values = kernel.dynamicBuffer(dtype),
1825 .x = kernel.dynamicBuffer(dtype),
1826 .rows = kernel.scalar(.i32),
1827 .nnz = kernel.scalar(.i32),
1828 .x_extent = kernel.scalar(.i32),
1829 .columns = kernel.scalar(.i32),
1830 },
1831 .Instance = SpmmCsr,
1832 .schedule = spmmCsrFamilySchedule,
1833 .body = spmmCsrRuntimeBody,
1834 });
1835 }
1836
1837 pub const SpmmCsrRuntimeFamilyF64 = spmmCsrRuntimeFamily(.f64);
1838 pub const SpmmCsrRuntimeFamilyF32 = spmmCsrRuntimeFamily(.f32);
1839 pub const SpmmCsrRuntimeFamilyF16 = spmmCsrRuntimeFamily(.f16);
1840
1841 pub fn spmmCsrFamilyTarget(allocator: std.mem.Allocator, instance: SpmmCsr) ![]u8 {
1842 return std.fmt.allocPrint(
1843 allocator,
1844 "accy.kernel.sparse.spmm_csr_{s}_family_{d}x{d}_{s}",
1845 .{ @tagName(instance.structure), instance.threads.x, instance.threads.y, instance.dtype.name() },
1846 );
1847 }
1848
1849 pub fn spmmCsrFamilyEntryName(allocator: std.mem.Allocator, instance: SpmmCsr) ![]u8 {
1850 return std.fmt.allocPrint(
1851 allocator,
1852 "accy_kernel_sparse_spmm_csr_{s}_family_{d}x{d}_{s}",
1853 .{ @tagName(instance.structure), instance.threads.x, instance.threads.y, instance.dtype.name() },
1854 );
1855 }
1856
1857 pub fn spmmCsrRuntimeArguments(
1858 instance: SpmmCsr,
1859 nnz: u64,
1860 x_extent: u64,
1861 columns: u64,
1862 ) ![4]choir_abi.ScalarArgument {
1863 return .{
1864 .{ .u32 = try runtimeExtentArgument(instance.rows) },
1865 .{ .u32 = try runtimeExtentArgument(nnz) },
1866 .{ .u32 = try runtimeExtentArgument(x_extent) },
1867 .{ .u32 = try runtimeExtentArgument(columns) },
1868 };
1869 }
1870
1871 pub fn spmmCsrMaxRows(instance: SpmmCsr) u64 {
1872 _ = instance;
1873 return extent_mod.runtime_extent_max;
1874 }
1875
1876 pub fn spmmCsrMaxColumns(instance: SpmmCsr) u64 {
1877 _ = instance;
1878 return extent_mod.runtime_extent_max;
1879 }
1880
1881 pub fn spmmCsrShapeProfileDimensions(instance: SpmmCsr) [2]artifact_product.KernelCallShapeProfileDimension {
1882 const bounds = shape.Bounds{ .min = 1, .max = extent_mod.runtime_extent_max };
1883 return .{
1884 .{
1885 .name = instance.row_axis,
1886 .runtime_scalar_argument_index = 0,
1887 .bounds = bounds,
1888 },
1889 .{
1890 .name = instance.column_axis,
1891 .runtime_scalar_argument_index = 3,
1892 .bounds = bounds,
1893 },
1894 };
1895 }
1896
1897 fn spmmCsrLaunch(instance: SpmmCsr) !artifact_product.KernelCallLaunch {
1898 if (!spmmCsrInstanceValid(instance)) return error.UnsupportedSpmmCsrInstance;
1899 return .{ .derived = .{
1900 .grid = .{
1901 .{ .runtime_u32_ceil_div = .{ .argument_index = 3, .divisor = instance.threads.x } },
1902 .{ .runtime_u32_ceil_div = .{ .argument_index = 0, .divisor = instance.threads.y } },
1903 .{ .fixed = 1 },
1904 },
1905 .threadgroup = .{ instance.threads.x, instance.threads.y, 1 },
1906 } };
1907 }
1908
1909 pub fn spmmCsrShapeFamily(backing_allocator: std.mem.Allocator, instance: SpmmCsr) !shape.Family {
1910 var builder = try shape.Builder.init(backing_allocator, "spmm_csr");
1911 errdefer builder.deinit();
1912 const rows = try builder.symbol(instance.row_axis);
1913 const nnz = try builder.symbol(spmv_csr_nonzero_axis);
1914 const x_extent = try builder.symbol(instance.x_axis);
1915 const columns = try builder.symbol(instance.column_axis);
1916 const rows_expr = try builder.symbolExpression(rows);
1917 const row_ptr_expr = try builder.addExpression(rows_expr, builder.constantExpression(1));
1918 const nnz_expr = try builder.symbolExpression(nnz);
1919 const x_extent_expr = try builder.symbolExpression(x_extent);
1920 const columns_expr = try builder.symbolExpression(columns);
1921 _ = try builder.tensor("row_ptr", &.{row_ptr_expr});
1922 _ = try builder.tensor("cols", &.{nnz_expr});
1923 _ = try builder.tensor("values", &.{nnz_expr});
1924 _ = try builder.tensor("x", &.{ x_extent_expr, columns_expr });
1925 _ = try builder.tensor("y", &.{ rows_expr, columns_expr });
1926 const bounds = shape.Bounds{ .min = 1, .max = extent_mod.runtime_extent_max };
1927 try builder.assumeBounds(rows_expr, bounds);
1928 try builder.assumeBounds(nnz_expr, bounds);
1929 try builder.assumeBounds(x_extent_expr, bounds);
1930 try builder.assumeBounds(columns_expr, bounds);
1931 return builder.finish();
1932 }
1933
1934 pub fn spmmCsrFamilyFingerprint(backing_allocator: std.mem.Allocator, instance: SpmmCsr) !u64 {
1935 var family = try spmmCsrShapeFamily(backing_allocator, instance);
1936 defer family.deinit();
1937 return shape.fingerprint(family);
1938 }
1939
1940 pub fn spmmCsrTuningExtents(instance: SpmmCsr) [4]u64 {
1941 return .{ instance.rows, instance.columns, instance.nnz, instance.x_extent };
1942 }
1943
1944 pub fn spmmCsrTuningOperation(instance: SpmmCsr) entry.Operation {
1945 _ = instance;
1946 return .{ .sparse = .csr_spmm };
1947 }
1948
1949 pub fn spmmCsrFamilyTuningKey(
1950 backing_allocator: std.mem.Allocator,
1951 device_fingerprint: u64,
1952 instance: SpmmCsr,
1953 ) !tuning.FamilyTuningKey {
1954 const family_fingerprint = try spmmCsrFamilyFingerprint(backing_allocator, instance);
1955 const extents = spmmCsrTuningExtents(instance);
1956 return tuning.FamilyTuningKey.init(
1957 device_fingerprint,
1958 family_fingerprint,
1959 entry.operationFingerprint(spmmCsrTuningOperation(instance)),
1960 instance.dtype,
1961 spmm_csr_family_version,
1962 extents[0..],
1963 ) orelse unreachable;
1964 }
1965
1966 pub fn resolveSpmmCsrThreads(
1967 backing_allocator: std.mem.Allocator,
1968 reader: tuning.FamilyTuningReader,
1969 instance: SpmmCsr,
1970 ) !?entry.Threads2D {
1971 const key = try spmmCsrFamilyTuningKey(backing_allocator, reader.device_fingerprint, instance);
1972 const record = reader.table.find(key) orelse return null;
1973 var default_candidate = instance;
1974 default_candidate.threads = spmmCsrRepresentableThreads(default_candidate) orelse return null;
1975 {
1976 const target = try spmmCsrFamilyTarget(backing_allocator, default_candidate);
1977 defer backing_allocator.free(target);
1978 if (std.mem.eql(u8, target, record.target)) return default_candidate.threads;
1979 }
1980 const thread_candidates = spmmCsrThreadCandidatesForExtents(instance.rows, instance.columns);
1981 for (thread_candidates.slice()) |threads| {
1982 var candidate = instance;
1983 candidate.threads = threads;
1984 candidate.threads = spmmCsrRepresentableThreads(candidate) orelse continue;
1985 const target = try spmmCsrFamilyTarget(backing_allocator, candidate);
1986 defer backing_allocator.free(target);
1987 if (std.mem.eql(u8, target, record.target)) return candidate.threads;
1988 }
1989 return null;
1990 }
1991
1992 pub fn createSpmmCsrFamilyArtifact(
1993 allocator: std.mem.Allocator,
1994 handle: kernel.BackendHandle,
1995 instance: SpmmCsr,
1996 options: entry.ArtifactOptions,
1997 ) !kernel.OwnedKernelCallArtifact {
1998 if (!spmmCsrInstanceValid(instance)) return error.UnsupportedSpmmCsrInstance;
1999 const target = try spmmCsrFamilyTarget(allocator, instance);
2000 defer allocator.free(target);
2001 const entry_name = try spmmCsrFamilyEntryName(allocator, instance);
2002 defer allocator.free(entry_name);
2003 const family_fingerprint = options.shape_family_fingerprint orelse try spmmCsrFamilyFingerprint(allocator, instance);
2004 const shape_profile_dimensions = spmmCsrShapeProfileDimensions(instance);
2005 const shape_profile = options.shape_profile orelse artifact_product.KernelCallShapeProfile{
2006 .name = "spmm_csr",
2007 .fingerprint = family_fingerprint,
2008 .dimensions = shape_profile_dimensions[0..],
2009 };
2010
2011 var graph = switch (instance.dtype) {
2012 .f64 => try SpmmCsrRuntimeFamilyF64.buildNamed(allocator, options.limits, entry_name, instance),
2013 .f32 => try SpmmCsrRuntimeFamilyF32.buildNamed(allocator, options.limits, entry_name, instance),
2014 .f16 => try SpmmCsrRuntimeFamilyF16.buildNamed(allocator, options.limits, entry_name, instance),
2015 else => return error.UnsupportedDType,
2016 };
2017 defer graph.deinit();
2018 return kernel.createKernelCallArtifact(allocator, handle, &graph, .{
2019 .target = target,
2020 .version = spmm_csr_family_version,
2021 .format = options.format,
2022 .kernel_plan = options.kernel_plan,
2023 .element_count_argument = options.element_count_argument,
2024 .shape_family_fingerprint = family_fingerprint,
2025 .shape_profile = shape_profile,
2026 .launch = options.launch orelse try spmmCsrLaunch(instance),
2027 .runtime_scalar_argument_count = if (options.runtime_scalar_argument_count == 0) 4 else options.runtime_scalar_argument_count,
2028 .static_arguments = options.static_arguments,
2029 });
2030 }
2031
2032 const testing = std.testing;
2033
2034 fn sparseFamilyTuningTestCapabilities(device_id: u32) gpu.BackendCapabilities {
2035 return .{ .identity = .{
2036 .backend = .cuda,
2037 .family = .nvidia_cuda,
2038 .name = "sparse-family-tuning-test-device",
2039 .vendor_id = 0x10de,
2040 .device_id = device_id,
2041 } };
2042 }
2043
2044 fn spmvTestValue(comptime dtype: DType, value: f32) dtype.ZigType() {
2045 return switch (dtype) {
2046 .f64 => @floatCast(value),
2047 .f32 => value,
2048 .f16 => @floatCast(value),
2049 else => @compileError("unsupported SpMV CSR test dtype"),
2050 };
2051 }
2052
2053 fn spmvTestFloat(comptime dtype: DType, value: dtype.ZigType()) f32 {
2054 return switch (dtype) {
2055 .f64 => @floatCast(value),
2056 .f32 => value,
2057 .f16 => @floatCast(value),
2058 else => @compileError("unsupported SpMV CSR test dtype"),
2059 };
2060 }
2061
2062 fn expectSpmvCsrMatchesDense(comptime dtype: DType, comptime structure: SpmvCsrStructure, comptime threads: u32) !void {
2063 const allocator = testing.allocator;
2064 const Scalar = dtype.ZigType();
2065 const rows: usize = 70;
2066 const cols_n: usize = 40;
2067 const instance = SpmvCsr{
2068 .rows = rows,
2069 .dtype = dtype,
2070 .accumulation_dtype = spmvCsrAccumulationDType(dtype).?,
2071 .threads = threads,
2072 .structure = structure,
2073 };
2074 const blocks: u32 = @intCast(spmvCsrBlockCount(instance));
2075 try testing.expect(blocks > 1);
2076
2077 var row_ptr: [rows + 1]i32 = undefined;
2078 var cols_storage: [rows * 8]i32 = undefined;
2079 var values_storage: [rows * 8]Scalar = undefined;
2080 var seed: u32 = 0x2545f491;
2081 var nnz: usize = 0;
2082 row_ptr[0] = 0;
2083 for (0..rows) |row| {
2084 const row_nnz: usize = switch (row % 5) {
2085 0 => 0,
2086 1 => 1,
2087 2 => 3,
2088 3 => 8,
2089 else => 5,
2090 };
2091 var produced: usize = 0;
2092 while (produced < row_nnz) : (produced += 1) {
2093 seed ^= seed << 13;
2094 seed ^= seed >> 17;
2095 seed ^= seed << 5;
2096 cols_storage[nnz] = @intCast(seed % cols_n);
2097 values_storage[nnz] = spmvTestValue(dtype, @floatFromInt((seed >> 8) % 9 + 1));
2098 nnz += 1;
2099 }
2100 row_ptr[row + 1] = @intCast(nnz);
2101 }
2102
2103 var x: [cols_n]Scalar = undefined;
2104 for (&x, 0..) |*value, index| value.* = spmvTestValue(dtype, @floatFromInt((index % 7) + 1));
2105
2106 var y = @as([rows]Scalar, @splat(spmvTestValue(dtype, -1)));
2107 var graph = switch (dtype) {
2108 .f64 => try SpmvCsrRuntimeFamilyF64.build(allocator, SpmvCsrRuntimeFamilyF64.Limits.testing, instance),
2109 .f32 => try SpmvCsrRuntimeFamilyF32.build(allocator, SpmvCsrRuntimeFamilyF32.Limits.testing, instance),
2110 .f16 => try SpmvCsrRuntimeFamilyF16.build(allocator, SpmvCsrRuntimeFamilyF16.Limits.testing, instance),
2111 else => @compileError("unsupported SpMV CSR test dtype"),
2112 };
2113 defer graph.deinit();
2114 try graph.runCpuWithLaunch(allocator, &.{
2115 kernel.argumentBuffer(Scalar, y[0..]),
2116 kernel.argumentBuffer(i32, row_ptr[0..]),
2117 kernel.argumentBuffer(i32, cols_storage[0..nnz]),
2118 kernel.argumentBuffer(Scalar, values_storage[0..nnz]),
2119 kernel.argumentBuffer(Scalar, x[0..]),
2120 kernel.argumentI32(@intCast(rows)),
2121 kernel.argumentI32(@intCast(nnz)),
2122 kernel.argumentI32(@intCast(cols_n)),
2123 }, .{
2124 .grid = .{ blocks, 1, 1 },
2125 .block = .{ instance.threads, 1, 1 },
2126 });
2127
2128 for (0..rows) |row| {
2129 var expected: f32 = 0;
2130 const begin: usize = @intCast(row_ptr[row]);
2131 const end: usize = @intCast(row_ptr[row + 1]);
2132 for (begin..end) |element| {
2133 expected += spmvTestFloat(dtype, values_storage[element]) * spmvTestFloat(dtype, x[@intCast(cols_storage[element])]);
2134 }
2135 const expected_output = spmvTestFloat(dtype, spmvTestValue(dtype, expected));
2136 try testing.expectApproxEqAbs(expected_output, spmvTestFloat(dtype, y[row]), 0.001);
2137 }
2138 }
2139
2140 test "sparse spmv csr matches the dense reference on both structures" {
2141 try expectSpmvCsrMatchesDense(.f64, .row_warp, 64);
2142 try expectSpmvCsrMatchesDense(.f64, .row_thread, 64);
2143 try expectSpmvCsrMatchesDense(.f32, .row_warp, 64);
2144 try expectSpmvCsrMatchesDense(.f32, .row_thread, 64);
2145 try expectSpmvCsrMatchesDense(.f32, .row_thread, 48);
2146 try expectSpmvCsrMatchesDense(.f16, .row_warp, 64);
2147 try expectSpmvCsrMatchesDense(.f16, .row_thread, 64);
2148 }
2149
2150 fn expectSpmvCooMatchesDense(comptime dtype: DType, comptime structure: SpmvCooStructure, threads: u32) !void {
2151 const Scalar = dtype.ZigType();
2152 const allocator = testing.allocator;
2153 const rows: usize = 17;
2154 const nnz: usize = 37;
2155 const x_extent: usize = 13;
2156 const instance = SpmvCoo{
2157 .rows = rows,
2158 .nnz = nnz,
2159 .x_extent = x_extent,
2160 .dtype = dtype,
2161 .accumulation_dtype = spmvCooAccumulationDTypeForStructure(structure, dtype).?,
2162 .threads = threads,
2163 .structure = structure,
2164 };
2165 const blocks: u32 = @intCast(spmvCooBlockCount(instance));
2166 try testing.expect(blocks > 1);
2167
2168 var row_indices: [nnz]i32 = undefined;
2169 var cols_storage: [nnz]i32 = undefined;
2170 var values_storage: [nnz]Scalar = undefined;
2171 for (0..nnz) |element| {
2172 row_indices[element] = @intCast((element * 5 + 2) % rows);
2173 cols_storage[element] = @intCast((element * 7 + 1) % x_extent);
2174 values_storage[element] = spmvTestValue(dtype, @floatFromInt((element % 9) + 1));
2175 }
2176 row_indices[3] = -1;
2177 row_indices[9] = @intCast(rows + 5);
2178 cols_storage[14] = -4;
2179 cols_storage[21] = @intCast(x_extent + 3);
2180 row_indices[24] = 4;
2181 row_indices[25] = 4;
2182 cols_storage[24] = 6;
2183 cols_storage[25] = 6;
2184
2185 var x: [x_extent]Scalar = undefined;
2186 for (&x, 0..) |*value, index| value.* = spmvTestValue(dtype, @floatFromInt((index % 5) + 1));
2187
2188 var y: [rows]Scalar = undefined;
2189 var expected: [rows]f32 = undefined;
2190 for (&y, &expected, 0..) |*actual, *want, row| {
2191 const seed = @as(f32, @floatFromInt(row % 3)) * 0.25;
2192 actual.* = spmvTestValue(dtype, seed);
2193 want.* = seed;
2194 }
2195
2196 var graph = switch (dtype) {
2197 .f64 => try SpmvCooRuntimeFamilyF64.build(allocator, SpmvCooRuntimeFamilyF64.Limits.testing, instance),
2198 .f32 => try SpmvCooRuntimeFamilyF32.build(allocator, SpmvCooRuntimeFamilyF32.Limits.testing, instance),
2199 .f16 => try SpmvCooRuntimeFamilyF16.build(allocator, SpmvCooRuntimeFamilyF16.Limits.testing, instance),
2200 else => @compileError("unsupported SpMV COO test dtype"),
2201 };
2202 defer graph.deinit();
2203 try graph.runCpuWithLaunch(allocator, &.{
2204 kernel.argumentBuffer(Scalar, y[0..]),
2205 kernel.argumentBuffer(i32, row_indices[0..]),
2206 kernel.argumentBuffer(i32, cols_storage[0..]),
2207 kernel.argumentBuffer(Scalar, values_storage[0..]),
2208 kernel.argumentBuffer(Scalar, x[0..]),
2209 kernel.argumentI32(@intCast(rows)),
2210 kernel.argumentI32(@intCast(nnz)),
2211 kernel.argumentI32(@intCast(x_extent)),
2212 }, .{
2213 .grid = .{ blocks, 1, 1 },
2214 .block = .{ instance.threads, 1, 1 },
2215 });
2216
2217 for (0..nnz) |element| {
2218 const row = row_indices[element];
2219 const col = cols_storage[element];
2220 if (row < 0 or row >= @as(i32, @intCast(rows)) or col < 0 or col >= @as(i32, @intCast(x_extent))) continue;
2221 expected[@intCast(row)] += spmvTestFloat(dtype, values_storage[element]) * spmvTestFloat(dtype, x[@intCast(col)]);
2222 }
2223
2224 for (y, expected) |actual, want| {
2225 const expected_output = spmvTestFloat(dtype, spmvTestValue(dtype, want));
2226 try testing.expectApproxEqAbs(expected_output, spmvTestFloat(dtype, actual), 0.001);
2227 }
2228 }
2229
2230 test "sparse spmv coo accumulates matching the dense reference" {
2231 try expectSpmvCooMatchesDense(.f32, .element_thread, 16);
2232 try expectSpmvCooMatchesDense(.f32, .row_thread, 8);
2233 try expectSpmvCooMatchesDense(.f16, .row_thread, 8);
2234 try expectSpmvCooMatchesDense(.f64, .row_thread, 8);
2235 }
2236
2237 fn expectSpmvEllMatchesDense(comptime dtype: DType) !void {
2238 const allocator = testing.allocator;
2239 const Scalar = dtype.ZigType();
2240 const rows: usize = 70;
2241 const slots: usize = 8;
2242 const x_extent: usize = 40;
2243 const instance = SpmvEll{
2244 .rows = rows,
2245 .slots = slots,
2246 .x_extent = x_extent,
2247 .dtype = dtype,
2248 .accumulation_dtype = spmvEllAccumulationDType(dtype).?,
2249 .threads = 48,
2250 };
2251 const blocks: u32 = @intCast(spmvEllBlockCount(instance));
2252 try testing.expect(blocks > 1);
2253
2254 var cols_storage = @as([(rows * slots)]i32, @splat(-1));
2255 var values_storage = @as([(rows * slots)]Scalar, @splat(spmvTestValue(dtype, 97)));
2256 var seed: u32 = 0x6a09e667;
2257 for (0..rows) |row| {
2258 const row_slots: usize = switch (row % 5) {
2259 0 => 0,
2260 1 => 1,
2261 2 => 3,
2262 3 => slots,
2263 else => 5,
2264 };
2265 for (0..row_slots) |slot| {
2266 seed = seed *% 22695477 +% 1;
2267 const element = slot * rows + row;
2268 cols_storage[element] = @intCast(seed % x_extent);
2269 values_storage[element] = spmvTestValue(dtype, @floatFromInt((seed >> 9) % 7 + 1));
2270 }
2271 }
2272
2273 var x: [x_extent]Scalar = undefined;
2274 for (&x, 0..) |*value, index| value.* = spmvTestValue(dtype, @floatFromInt((index % 11) + 1));
2275
2276 var y = @as([rows]Scalar, @splat(spmvTestValue(dtype, -1)));
2277 var graph = switch (dtype) {
2278 .f64 => try SpmvEllRuntimeFamilyF64.build(allocator, SpmvEllRuntimeFamilyF64.Limits.testing, instance),
2279 .f32 => try SpmvEllRuntimeFamilyF32.build(allocator, SpmvEllRuntimeFamilyF32.Limits.testing, instance),
2280 .f16 => try SpmvEllRuntimeFamilyF16.build(allocator, SpmvEllRuntimeFamilyF16.Limits.testing, instance),
2281 else => @compileError("unsupported SpMV ELL test dtype"),
2282 };
2283 defer graph.deinit();
2284 try graph.runCpuWithLaunch(allocator, &.{
2285 kernel.argumentBuffer(Scalar, y[0..]),
2286 kernel.argumentBuffer(i32, cols_storage[0..]),
2287 kernel.argumentBuffer(Scalar, values_storage[0..]),
2288 kernel.argumentBuffer(Scalar, x[0..]),
2289 kernel.argumentI32(@intCast(rows)),
2290 kernel.argumentI32(@intCast(slots)),
2291 kernel.argumentI32(@intCast(x_extent)),
2292 }, .{
2293 .grid = .{ blocks, 1, 1 },
2294 .block = .{ instance.threads, 1, 1 },
2295 });
2296
2297 for (0..rows) |row| {
2298 var expected: f32 = 0;
2299 for (0..slots) |slot| {
2300 const element = slot * rows + row;
2301 const column = cols_storage[element];
2302 if (column >= 0) {
2303 expected += spmvTestFloat(dtype, values_storage[element]) * spmvTestFloat(dtype, x[@intCast(column)]);
2304 }
2305 }
2306 const expected_output = spmvTestFloat(dtype, spmvTestValue(dtype, expected));
2307 try testing.expectApproxEqAbs(expected_output, spmvTestFloat(dtype, y[row]), 0.001);
2308 }
2309 }
2310
2311 test "sparse spmv ell matches the dense reference with padded slots" {
2312 try expectSpmvEllMatchesDense(.f64);
2313 try expectSpmvEllMatchesDense(.f32);
2314 try expectSpmvEllMatchesDense(.f16);
2315 }
2316
2317 fn expectSpmvSellMatchesDense(comptime dtype: DType) !void {
2318 const allocator = testing.allocator;
2319 const Scalar = dtype.ZigType();
2320 const rows: usize = 70;
2321 const slice_size: usize = 8;
2322 const slices: usize = (rows + slice_size - 1) / slice_size;
2323 const x_extent: usize = 40;
2324
2325 var slice_offsets: [slices + 1]i32 = undefined;
2326 var cols_storage = @as([(slices * slice_size * 8)]i32, @splat(-1));
2327 var values_storage = @as([(slices * slice_size * 8)]Scalar, @splat(spmvTestValue(dtype, 97)));
2328 var seed: u32 = 0x243f6a88;
2329 var values_size: usize = 0;
2330 for (0..slices) |slice| {
2331 slice_offsets[slice] = @intCast(values_size);
2332 var slice_slots: usize = 0;
2333 for (0..slice_size) |local| {
2334 const row = slice * slice_size + local;
2335 if (row < rows) slice_slots = @max(slice_slots, sellTestRowSlots(row));
2336 }
2337 for (0..slice_slots) |slot| {
2338 for (0..slice_size) |local| {
2339 const row = slice * slice_size + local;
2340 const element = values_size + slot * slice_size + local;
2341 if (row < rows and slot < sellTestRowSlots(row)) {
2342 seed = seed *% 1103515245 +% 12345;
2343 cols_storage[element] = @intCast(seed % x_extent);
2344 values_storage[element] = spmvTestValue(dtype, @floatFromInt((seed >> 10) % 9 + 1));
2345 }
2346 }
2347 }
2348 values_size += slice_slots * slice_size;
2349 }
2350 slice_offsets[slices] = @intCast(values_size);
2351
2352 const instance = SpmvSell{
2353 .rows = rows,
2354 .slice_size = slice_size,
2355 .values_size = values_size,
2356 .x_extent = x_extent,
2357 .dtype = dtype,
2358 .accumulation_dtype = spmvSellAccumulationDType(dtype).?,
2359 .threads = 48,
2360 };
2361 const blocks: u32 = @intCast(spmvSellBlockCount(instance));
2362 try testing.expect(blocks > 1);
2363 try testing.expectEqual(@as(u64, slices), spmvSellSliceCount(instance));
2364
2365 var x: [x_extent]Scalar = undefined;
2366 for (&x, 0..) |*value, index| value.* = spmvTestValue(dtype, @floatFromInt((index % 11) + 1));
2367
2368 var y = @as([rows]Scalar, @splat(spmvTestValue(dtype, -1)));
2369 var graph = switch (dtype) {
2370 .f64 => try SpmvSellRuntimeFamilyF64.build(allocator, SpmvSellRuntimeFamilyF64.Limits.testing, instance),
2371 .f32 => try SpmvSellRuntimeFamilyF32.build(allocator, SpmvSellRuntimeFamilyF32.Limits.testing, instance),
2372 .f16 => try SpmvSellRuntimeFamilyF16.build(allocator, SpmvSellRuntimeFamilyF16.Limits.testing, instance),
2373 else => @compileError("unsupported SpMV SELL test dtype"),
2374 };
2375 defer graph.deinit();
2376 try graph.runCpuWithLaunch(allocator, &.{
2377 kernel.argumentBuffer(Scalar, y[0..]),
2378 kernel.argumentBuffer(i32, slice_offsets[0..]),
2379 kernel.argumentBuffer(i32, cols_storage[0..values_size]),
2380 kernel.argumentBuffer(Scalar, values_storage[0..values_size]),
2381 kernel.argumentBuffer(Scalar, x[0..]),
2382 kernel.argumentI32(@intCast(rows)),
2383 kernel.argumentI32(@intCast(values_size)),
2384 kernel.argumentI32(@intCast(x_extent)),
2385 }, .{
2386 .grid = .{ blocks, 1, 1 },
2387 .block = .{ instance.threads, 1, 1 },
2388 });
2389
2390 for (0..rows) |row| {
2391 const slice = row / slice_size;
2392 const local = row - slice * slice_size;
2393 const begin: usize = @intCast(slice_offsets[slice]);
2394 const end: usize = @intCast(slice_offsets[slice + 1]);
2395 var expected: f32 = 0;
2396 var element = begin + local;
2397 while (element < end) : (element += slice_size) {
2398 const column = cols_storage[element];
2399 if (column >= 0) {
2400 expected += spmvTestFloat(dtype, values_storage[element]) * spmvTestFloat(dtype, x[@intCast(column)]);
2401 }
2402 }
2403 const expected_output = spmvTestFloat(dtype, spmvTestValue(dtype, expected));
2404 try testing.expectApproxEqAbs(expected_output, spmvTestFloat(dtype, y[row]), 0.001);
2405 }
2406 }
2407
2408 fn sellTestRowSlots(row: usize) usize {
2409 return switch (row % 7) {
2410 0 => 0,
2411 1 => 1,
2412 2 => 2,
2413 3 => 7,
2414 4 => 4,
2415 5 => 6,
2416 else => 3,
2417 };
2418 }
2419
2420 test "sparse spmv sell matches the dense reference with sliced padding" {
2421 try expectSpmvSellMatchesDense(.f64);
2422 try expectSpmvSellMatchesDense(.f32);
2423 try expectSpmvSellMatchesDense(.f16);
2424 }
2425
2426 fn expectSpmmCsrMatchesDense(comptime dtype: DType) !void {
2427 const allocator = testing.allocator;
2428 const Scalar = dtype.ZigType();
2429 const rows: usize = 37;
2430 const columns: usize = 11;
2431 const x_extent: usize = 23;
2432 const instance = SpmmCsr{
2433 .rows = rows,
2434 .columns = columns,
2435 .x_extent = x_extent,
2436 .dtype = dtype,
2437 .accumulation_dtype = spmmCsrAccumulationDType(dtype).?,
2438 .threads = .{ .x = 8, .y = 4 },
2439 };
2440
2441 var row_ptr: [rows + 1]i32 = undefined;
2442 var cols_storage: [rows * 6]i32 = undefined;
2443 var values_storage: [rows * 6]Scalar = undefined;
2444 var seed: u32 = 0x9e3779b9;
2445 var nnz: usize = 0;
2446 row_ptr[0] = 0;
2447 for (0..rows) |row| {
2448 const row_nnz: usize = switch (row % 5) {
2449 0 => 0,
2450 1 => 1,
2451 2 => 4,
2452 3 => 6,
2453 else => 3,
2454 };
2455 var produced: usize = 0;
2456 while (produced < row_nnz) : (produced += 1) {
2457 seed = seed *% 1664525 +% 1013904223;
2458 cols_storage[nnz] = @intCast(seed % x_extent);
2459 values_storage[nnz] = spmvTestValue(dtype, @floatFromInt((seed >> 9) % 7 + 1));
2460 nnz += 1;
2461 }
2462 row_ptr[row + 1] = @intCast(nnz);
2463 }
2464
2465 var x: [x_extent * columns]Scalar = undefined;
2466 for (&x, 0..) |*value, index| value.* = spmvTestValue(dtype, @floatFromInt((index % 13) + 1));
2467
2468 var y = @as([(rows * columns)]Scalar, @splat(spmvTestValue(dtype, -1)));
2469 var graph = switch (dtype) {
2470 .f64 => try SpmmCsrRuntimeFamilyF64.build(allocator, SpmmCsrRuntimeFamilyF64.Limits.testing, instance),
2471 .f32 => try SpmmCsrRuntimeFamilyF32.build(allocator, SpmmCsrRuntimeFamilyF32.Limits.testing, instance),
2472 .f16 => try SpmmCsrRuntimeFamilyF16.build(allocator, SpmmCsrRuntimeFamilyF16.Limits.testing, instance),
2473 else => @compileError("unsupported SpMM CSR test dtype"),
2474 };
2475 defer graph.deinit();
2476 try graph.runCpuWithLaunch(allocator, &.{
2477 kernel.argumentBuffer(Scalar, y[0..]),
2478 kernel.argumentBuffer(i32, row_ptr[0..]),
2479 kernel.argumentBuffer(i32, cols_storage[0..nnz]),
2480 kernel.argumentBuffer(Scalar, values_storage[0..nnz]),
2481 kernel.argumentBuffer(Scalar, x[0..]),
2482 kernel.argumentI32(@intCast(rows)),
2483 kernel.argumentI32(@intCast(nnz)),
2484 kernel.argumentI32(@intCast(x_extent)),
2485 kernel.argumentI32(@intCast(columns)),
2486 }, .{
2487 .grid = .{ @intCast(spmmCsrBlockCountX(instance)), @intCast(spmmCsrBlockCountY(instance)), 1 },
2488 .block = .{ instance.threads.x, instance.threads.y, 1 },
2489 });
2490
2491 for (0..rows) |row| {
2492 for (0..columns) |column| {
2493 var expected: f32 = 0;
2494 const begin: usize = @intCast(row_ptr[row]);
2495 const end: usize = @intCast(row_ptr[row + 1]);
2496 for (begin..end) |element| {
2497 const x_index: usize = @as(usize, @intCast(cols_storage[element])) * columns + column;
2498 expected += spmvTestFloat(dtype, values_storage[element]) * spmvTestFloat(dtype, x[x_index]);
2499 }
2500 const expected_output = spmvTestFloat(dtype, spmvTestValue(dtype, expected));
2501 const output_index = row * columns + column;
2502 try testing.expectApproxEqAbs(expected_output, spmvTestFloat(dtype, y[output_index]), 0.001);
2503 }
2504 }
2505 }
2506
2507 test "sparse spmm csr matches the dense matrix reference" {
2508 try expectSpmmCsrMatchesDense(.f64);
2509 try expectSpmmCsrMatchesDense(.f32);
2510 try expectSpmmCsrMatchesDense(.f16);
2511 }
2512
2513 test "sparse spmv csr identity validity and artifact contract" {
2514 const allocator = testing.allocator;
2515 const instance = SpmvCsr{ .rows = 1000, .threads = 64 };
2516
2517 const target = try spmvCsrFamilyTarget(allocator, instance);
2518 defer allocator.free(target);
2519 try testing.expectEqualStrings("accy.kernel.sparse.spmv_csr_row_warp_family_64_f32", target);
2520 const thread_target = try spmvCsrFamilyTarget(allocator, .{ .rows = 1000, .threads = 64, .structure = .row_thread });
2521 defer allocator.free(thread_target);
2522 try testing.expectEqualStrings("accy.kernel.sparse.spmv_csr_row_thread_family_64_f32", thread_target);
2523 const f16_target = try spmvCsrFamilyTarget(allocator, .{ .rows = 1000, .dtype = .f16, .threads = 64 });
2524 defer allocator.free(f16_target);
2525 try testing.expectEqualStrings("accy.kernel.sparse.spmv_csr_row_warp_family_64_f16", f16_target);
2526 const f64_target = try spmvCsrFamilyTarget(allocator, .{ .rows = 1000, .dtype = .f64, .accumulation_dtype = .f64, .threads = 64 });
2527 defer allocator.free(f64_target);
2528 try testing.expectEqualStrings("accy.kernel.sparse.spmv_csr_row_warp_family_64_f64", f64_target);
2529
2530 try testing.expectEqual(DType.f64, spmvCsrAccumulationDType(.f64).?);
2531 try testing.expectEqual(DType.f32, spmvCsrAccumulationDType(.f16).?);
2532 try testing.expect(spmvCsrInstanceValid(instance));
2533 try testing.expect(spmvCsrInstanceValid(.{ .rows = 10, .dtype = .f64, .accumulation_dtype = .f64, .threads = 64 }));
2534 try testing.expect(!spmvCsrInstanceValid(.{ .rows = 10, .dtype = .f64, .accumulation_dtype = .f32, .threads = 64 }));
2535 try testing.expect(spmvCsrInstanceValid(.{ .rows = 10, .dtype = .f16, .accumulation_dtype = .f32, .threads = 64 }));
2536 try testing.expect(!spmvCsrInstanceValid(.{ .rows = 10, .dtype = .f16, .accumulation_dtype = .f16, .threads = 64 }));
2537 try testing.expect(!spmvCsrInstanceValid(.{ .rows = 0, .threads = 64 }));
2538 try testing.expect(!spmvCsrInstanceValid(.{ .rows = 10, .threads = 48 }));
2539 try testing.expect(spmvCsrInstanceValid(.{ .rows = 10, .threads = 48, .structure = .row_thread }));
2540 try testing.expect(spmvCsrInstanceValid(.{ .rows = 10, .threads = 1, .structure = .row_thread }));
2541 try testing.expect(!spmvCsrInstanceValid(.{ .rows = 10, .threads = 0, .structure = .row_thread }));
2542
2543 var state = gpu.recording.BackendState{
2544 .allocator = allocator,
2545 .kind = .cuda,
2546 .format = .cuda_ptx,
2547 };
2548 var family_artifact = try createSpmvCsrFamilyArtifact(allocator, state.handle(), instance, .{ .limits = .testing });
2549 defer family_artifact.deinit();
2550 const family_entry = family_artifact.entry();
2551 try testing.expectEqualStrings("accy_kernel_sparse_spmv_csr_row_warp_family_64_f32", family_entry.entry_name);
2552 try testing.expectEqual(@as(u32, 8), family_entry.argument_count);
2553 try testing.expectEqual(@as(u32, 3), family_entry.runtime_scalar_argument_count);
2554 switch (family_entry.launch) {
2555 .derived => |launch| {
2556 try testing.expectEqual(@as(u32, 64), launch.threadgroup[0]);
2557 switch (launch.grid[0]) {
2558 .runtime_u32_ceil_div => |axis| try testing.expectEqual(@as(u32, 2), axis.divisor),
2559 else => return error.TestExpectedDerivedGrid,
2560 }
2561 },
2562 else => return error.TestExpectedDerivedLaunch,
2563 }
2564
2565 var f16_family_artifact = try createSpmvCsrFamilyArtifact(
2566 allocator,
2567 state.handle(),
2568 .{ .rows = 1000, .dtype = .f16, .accumulation_dtype = .f32, .threads = 64 },
2569 .{ .limits = .testing },
2570 );
2571 defer f16_family_artifact.deinit();
2572 const f16_family_entry = f16_family_artifact.entry();
2573 try testing.expectEqualStrings("accy_kernel_sparse_spmv_csr_row_warp_family_64_f16", f16_family_entry.entry_name);
2574 try testing.expect(f16_family_entry.required_dtypes.contains(.f16));
2575 try testing.expect(f16_family_entry.required_dtypes.contains(.i32));
2576
2577 var f64_family_artifact = try createSpmvCsrFamilyArtifact(
2578 allocator,
2579 state.handle(),
2580 .{ .rows = 1000, .dtype = .f64, .accumulation_dtype = .f64, .threads = 64 },
2581 .{ .limits = .testing },
2582 );
2583 defer f64_family_artifact.deinit();
2584 const f64_family_entry = f64_family_artifact.entry();
2585 try testing.expectEqualStrings("accy_kernel_sparse_spmv_csr_row_warp_family_64_f64", f64_family_entry.entry_name);
2586 try testing.expect(f64_family_entry.required_dtypes.contains(.f64));
2587 try testing.expect(f64_family_entry.required_dtypes.contains(.i32));
2588 }
2589
2590 test "sparse spmv csr family tuning keys discriminate dtype device and extents" {
2591 const allocator = testing.allocator;
2592 const device = tuning.deviceFingerprint(sparseFamilyTuningTestCapabilities(0x2684));
2593
2594 const row_warp = try spmvCsrFamilyTuningKey(allocator, device, .{
2595 .rows = 70,
2596 .nnz = 560,
2597 .x_extent = 40,
2598 .threads = 256,
2599 .structure = .row_warp,
2600 });
2601 const row_thread = try spmvCsrFamilyTuningKey(allocator, device, .{
2602 .rows = 70,
2603 .nnz = 560,
2604 .x_extent = 40,
2605 .threads = 64,
2606 .structure = .row_thread,
2607 });
2608 try testing.expect(row_warp.eql(row_thread));
2609
2610 const other_nnz = try spmvCsrFamilyTuningKey(allocator, device, .{
2611 .rows = 70,
2612 .nnz = 90,
2613 .x_extent = 40,
2614 });
2615 try testing.expect(!row_warp.eql(other_nnz));
2616
2617 const half = try spmvCsrFamilyTuningKey(allocator, device, .{
2618 .rows = 70,
2619 .nnz = 560,
2620 .x_extent = 40,
2621 .dtype = .f16,
2622 .accumulation_dtype = .f32,
2623 });
2624 try testing.expect(!row_warp.eql(half));
2625
2626 const other_device = try spmvCsrFamilyTuningKey(
2627 allocator,
2628 tuning.deviceFingerprint(sparseFamilyTuningTestCapabilities(0x1b80)),
2629 .{
2630 .rows = 70,
2631 .nnz = 560,
2632 .x_extent = 40,
2633 },
2634 );
2635 try testing.expect(!row_warp.eql(other_device));
2636 }
2637
2638 test "sparse spmv csr family tuning resolves structure" {
2639 const allocator = testing.allocator;
2640 const caps = sparseFamilyTuningTestCapabilities(0x2684);
2641 const device = tuning.deviceFingerprint(caps);
2642 const probe = SpmvCsr{
2643 .rows = 70,
2644 .nnz = 560,
2645 .x_extent = 40,
2646 .threads = 256,
2647 .structure = .row_warp,
2648 };
2649 var winner = probe;
2650 winner.structure = .row_thread;
2651 winner.threads = spmvCsrRepresentableThreads(winner) orelse return error.TestExpectedSparseStructure;
2652 const winner_target = try spmvCsrFamilyTarget(allocator, winner);
2653 defer allocator.free(winner_target);
2654
2655 const records = [_]tuning.FamilyTuningRecord{.{
2656 .key = try spmvCsrFamilyTuningKey(allocator, device, probe),
2657 .target = winner_target,
2658 .winner_median_ns = 800,
2659 .runner_up_median_ns = 1100,
2660 .sample_count = 30,
2661 }};
2662 const reader = tuning.FamilyTuningReader.init(caps, .{ .records = records[0..] });
2663
2664 const resolved = (try resolveSpmvCsrStructure(allocator, reader, probe)) orelse
2665 return error.TestExpectedSparseStructure;
2666 try testing.expectEqual(SpmvCsrStructure.row_thread, resolved);
2667
2668 const miss = try resolveSpmvCsrStructure(allocator, reader, .{
2669 .rows = 70,
2670 .nnz = 90,
2671 .x_extent = 40,
2672 .threads = 256,
2673 .structure = .row_thread,
2674 });
2675 try testing.expectEqual(@as(?SpmvCsrStructure, null), miss);
2676
2677 const stale_records = [_]tuning.FamilyTuningRecord{.{
2678 .key = try spmvCsrFamilyTuningKey(allocator, device, probe),
2679 .target = "accy.kernel.sparse.spmv_csr_row_group_family_256_f32",
2680 .winner_median_ns = 800,
2681 .runner_up_median_ns = 1100,
2682 .sample_count = 30,
2683 }};
2684 const stale_reader = tuning.FamilyTuningReader.init(caps, .{ .records = stale_records[0..] });
2685 try testing.expectEqual(
2686 @as(?SpmvCsrStructure, null),
2687 try resolveSpmvCsrStructure(allocator, stale_reader, probe),
2688 );
2689 }
2690
2691 test "sparse spmv coo family tuning keys discriminate dtype device and extents" {
2692 const allocator = testing.allocator;
2693 const device = tuning.deviceFingerprint(sparseFamilyTuningTestCapabilities(0x2684));
2694
2695 const element_thread = try spmvCooFamilyTuningKey(allocator, device, .{
2696 .rows = 70,
2697 .nnz = 560,
2698 .x_extent = 40,
2699 .threads = 256,
2700 .structure = .element_thread,
2701 });
2702 const row_thread = try spmvCooFamilyTuningKey(allocator, device, .{
2703 .rows = 70,
2704 .nnz = 560,
2705 .x_extent = 40,
2706 .threads = 70,
2707 .structure = .row_thread,
2708 });
2709 try testing.expect(element_thread.eql(row_thread));
2710
2711 const other_nnz = try spmvCooFamilyTuningKey(allocator, device, .{
2712 .rows = 70,
2713 .nnz = 90,
2714 .x_extent = 40,
2715 });
2716 try testing.expect(!element_thread.eql(other_nnz));
2717
2718 const half = try spmvCooFamilyTuningKey(allocator, device, .{
2719 .rows = 70,
2720 .nnz = 560,
2721 .x_extent = 40,
2722 .dtype = .f16,
2723 .accumulation_dtype = .f32,
2724 .threads = 70,
2725 .structure = .row_thread,
2726 });
2727 try testing.expect(!element_thread.eql(half));
2728
2729 const other_device = try spmvCooFamilyTuningKey(
2730 allocator,
2731 tuning.deviceFingerprint(sparseFamilyTuningTestCapabilities(0x1b80)),
2732 .{
2733 .rows = 70,
2734 .nnz = 560,
2735 .x_extent = 40,
2736 },
2737 );
2738 try testing.expect(!element_thread.eql(other_device));
2739 }
2740
2741 test "sparse spmv coo family tuning resolves structure" {
2742 const allocator = testing.allocator;
2743 const caps = sparseFamilyTuningTestCapabilities(0x2684);
2744 const device = tuning.deviceFingerprint(caps);
2745 const probe = SpmvCoo{
2746 .rows = 70,
2747 .nnz = 560,
2748 .x_extent = 40,
2749 .threads = 256,
2750 .structure = .element_thread,
2751 };
2752 var winner = probe;
2753 winner.structure = .row_thread;
2754 winner.threads = spmvCooRepresentableThreads(winner) orelse return error.TestExpectedSparseStructure;
2755 const winner_target = try spmvCooFamilyTarget(allocator, winner);
2756 defer allocator.free(winner_target);
2757
2758 const records = [_]tuning.FamilyTuningRecord{.{
2759 .key = try spmvCooFamilyTuningKey(allocator, device, probe),
2760 .target = winner_target,
2761 .winner_median_ns = 800,
2762 .runner_up_median_ns = 1100,
2763 .sample_count = 30,
2764 }};
2765 const reader = tuning.FamilyTuningReader.init(caps, .{ .records = records[0..] });
2766
2767 const resolved = (try resolveSpmvCooStructure(allocator, reader, probe)) orelse
2768 return error.TestExpectedSparseStructure;
2769 try testing.expectEqual(SpmvCooStructure.row_thread, resolved);
2770
2771 const miss = try resolveSpmvCooStructure(allocator, reader, .{
2772 .rows = 70,
2773 .nnz = 90,
2774 .x_extent = 40,
2775 .threads = 256,
2776 .structure = .element_thread,
2777 });
2778 try testing.expectEqual(@as(?SpmvCooStructure, null), miss);
2779
2780 const stale_records = [_]tuning.FamilyTuningRecord{.{
2781 .key = try spmvCooFamilyTuningKey(allocator, device, probe),
2782 .target = "accy.kernel.sparse.spmv_coo_warp_group_family_256_f32",
2783 .winner_median_ns = 800,
2784 .runner_up_median_ns = 1100,
2785 .sample_count = 30,
2786 }};
2787 const stale_reader = tuning.FamilyTuningReader.init(caps, .{ .records = stale_records[0..] });
2788 try testing.expectEqual(
2789 @as(?SpmvCooStructure, null),
2790 try resolveSpmvCooStructure(allocator, stale_reader, probe),
2791 );
2792 }
2793
2794 test "sparse spmv ell family tuning keys discriminate dtype device and fixed layout" {
2795 const allocator = testing.allocator;
2796 const device = tuning.deviceFingerprint(sparseFamilyTuningTestCapabilities(0x2684));
2797
2798 const block_256 = try spmvEllFamilyTuningKey(allocator, device, .{
2799 .rows = 70,
2800 .slots = 8,
2801 .x_extent = 40,
2802 .threads = 256,
2803 });
2804 const block_64 = try spmvEllFamilyTuningKey(allocator, device, .{
2805 .rows = 70,
2806 .slots = 8,
2807 .x_extent = 40,
2808 .threads = 64,
2809 });
2810 try testing.expect(block_256.eql(block_64));
2811
2812 const other_slots = try spmvEllFamilyTuningKey(allocator, device, .{
2813 .rows = 70,
2814 .slots = 16,
2815 .x_extent = 40,
2816 });
2817 try testing.expect(!block_256.eql(other_slots));
2818
2819 const other_x = try spmvEllFamilyTuningKey(allocator, device, .{
2820 .rows = 70,
2821 .slots = 8,
2822 .x_extent = 41,
2823 });
2824 try testing.expect(!block_256.eql(other_x));
2825
2826 const half = try spmvEllFamilyTuningKey(allocator, device, .{
2827 .rows = 70,
2828 .slots = 8,
2829 .x_extent = 40,
2830 .dtype = .f16,
2831 .accumulation_dtype = .f32,
2832 });
2833 try testing.expect(!block_256.eql(half));
2834
2835 const other_device = try spmvEllFamilyTuningKey(
2836 allocator,
2837 tuning.deviceFingerprint(sparseFamilyTuningTestCapabilities(0x1b80)),
2838 .{
2839 .rows = 70,
2840 .slots = 8,
2841 .x_extent = 40,
2842 },
2843 );
2844 try testing.expect(!block_256.eql(other_device));
2845 }
2846
2847 test "sparse spmv ell family tuning resolves thread blocks" {
2848 const allocator = testing.allocator;
2849 const caps = sparseFamilyTuningTestCapabilities(0x2684);
2850 const device = tuning.deviceFingerprint(caps);
2851 const probe = SpmvEll{
2852 .rows = 70,
2853 .slots = 8,
2854 .x_extent = 40,
2855 .threads = 256,
2856 };
2857 var winner = probe;
2858 winner.threads = 64;
2859 const winner_target = try spmvEllFamilyTarget(allocator, winner);
2860 defer allocator.free(winner_target);
2861
2862 const records = [_]tuning.FamilyTuningRecord{.{
2863 .key = try spmvEllFamilyTuningKey(allocator, device, probe),
2864 .target = winner_target,
2865 .winner_median_ns = 800,
2866 .runner_up_median_ns = 1100,
2867 .sample_count = 30,
2868 }};
2869 const reader = tuning.FamilyTuningReader.init(caps, .{ .records = records[0..] });
2870
2871 const resolved = (try resolveSpmvEllThreads(allocator, reader, probe)) orelse
2872 return error.TestExpectedSparseStructure;
2873 try testing.expectEqual(@as(u32, 64), resolved);
2874
2875 const miss = try resolveSpmvEllThreads(allocator, reader, .{
2876 .rows = 70,
2877 .slots = 9,
2878 .x_extent = 40,
2879 .threads = 256,
2880 });
2881 try testing.expectEqual(@as(?u32, null), miss);
2882
2883 const stale_records = [_]tuning.FamilyTuningRecord{.{
2884 .key = try spmvEllFamilyTuningKey(allocator, device, probe),
2885 .target = "accy.kernel.sparse.spmv_ell_warp_group_family_256_f32",
2886 .winner_median_ns = 800,
2887 .runner_up_median_ns = 1100,
2888 .sample_count = 30,
2889 }};
2890 const stale_reader = tuning.FamilyTuningReader.init(caps, .{ .records = stale_records[0..] });
2891 try testing.expectEqual(
2892 @as(?u32, null),
2893 try resolveSpmvEllThreads(allocator, stale_reader, probe),
2894 );
2895 }
2896
2897 test "sparse spmv sell family tuning keys discriminate dtype device and fixed layout" {
2898 const allocator = testing.allocator;
2899 const device = tuning.deviceFingerprint(sparseFamilyTuningTestCapabilities(0x2684));
2900
2901 const block_256 = try spmvSellFamilyTuningKey(allocator, device, .{
2902 .rows = 70,
2903 .slice_size = 8,
2904 .values_size = 400,
2905 .x_extent = 40,
2906 .threads = 256,
2907 });
2908 const block_64 = try spmvSellFamilyTuningKey(allocator, device, .{
2909 .rows = 70,
2910 .slice_size = 8,
2911 .values_size = 400,
2912 .x_extent = 40,
2913 .threads = 64,
2914 });
2915 try testing.expect(block_256.eql(block_64));
2916
2917 const other_slice = try spmvSellFamilyTuningKey(allocator, device, .{
2918 .rows = 70,
2919 .slice_size = 16,
2920 .values_size = 400,
2921 .x_extent = 40,
2922 });
2923 try testing.expect(!block_256.eql(other_slice));
2924
2925 const other_values = try spmvSellFamilyTuningKey(allocator, device, .{
2926 .rows = 70,
2927 .slice_size = 8,
2928 .values_size = 512,
2929 .x_extent = 40,
2930 });
2931 try testing.expect(!block_256.eql(other_values));
2932
2933 const half = try spmvSellFamilyTuningKey(allocator, device, .{
2934 .rows = 70,
2935 .slice_size = 8,
2936 .values_size = 400,
2937 .x_extent = 40,
2938 .dtype = .f16,
2939 .accumulation_dtype = .f32,
2940 });
2941 try testing.expect(!block_256.eql(half));
2942
2943 const other_device = try spmvSellFamilyTuningKey(
2944 allocator,
2945 tuning.deviceFingerprint(sparseFamilyTuningTestCapabilities(0x1b80)),
2946 .{
2947 .rows = 70,
2948 .slice_size = 8,
2949 .values_size = 400,
2950 .x_extent = 40,
2951 },
2952 );
2953 try testing.expect(!block_256.eql(other_device));
2954 }
2955
2956 test "sparse spmv sell family tuning resolves thread blocks" {
2957 const allocator = testing.allocator;
2958 const caps = sparseFamilyTuningTestCapabilities(0x2684);
2959 const device = tuning.deviceFingerprint(caps);
2960 const probe = SpmvSell{
2961 .rows = 70,
2962 .slice_size = 8,
2963 .values_size = 400,
2964 .x_extent = 40,
2965 .threads = 256,
2966 };
2967 var winner = probe;
2968 winner.threads = 64;
2969 winner.threads = spmvSellRepresentableThreads(winner) orelse return error.TestExpectedSparseStructure;
2970 const winner_target = try spmvSellFamilyTarget(allocator, winner);
2971 defer allocator.free(winner_target);
2972
2973 const records = [_]tuning.FamilyTuningRecord{.{
2974 .key = try spmvSellFamilyTuningKey(allocator, device, probe),
2975 .target = winner_target,
2976 .winner_median_ns = 800,
2977 .runner_up_median_ns = 1100,
2978 .sample_count = 30,
2979 }};
2980 const reader = tuning.FamilyTuningReader.init(caps, .{ .records = records[0..] });
2981
2982 const resolved = (try resolveSpmvSellThreads(allocator, reader, probe)) orelse
2983 return error.TestExpectedSparseStructure;
2984 try testing.expectEqual(@as(u32, 64), resolved);
2985
2986 const miss = try resolveSpmvSellThreads(allocator, reader, .{
2987 .rows = 70,
2988 .slice_size = 16,
2989 .values_size = 400,
2990 .x_extent = 40,
2991 .threads = 256,
2992 });
2993 try testing.expectEqual(@as(?u32, null), miss);
2994
2995 const stale_records = [_]tuning.FamilyTuningRecord{.{
2996 .key = try spmvSellFamilyTuningKey(allocator, device, probe),
2997 .target = "accy.kernel.sparse.spmv_sell_row_thread_slice8_family_96_f32",
2998 .winner_median_ns = 800,
2999 .runner_up_median_ns = 1100,
3000 .sample_count = 30,
3001 }};
3002 const stale_reader = tuning.FamilyTuningReader.init(caps, .{ .records = stale_records[0..] });
3003 try testing.expectEqual(
3004 @as(?u32, null),
3005 try resolveSpmvSellThreads(allocator, stale_reader, probe),
3006 );
3007 }
3008
3009 test "sparse spmm csr family tuning keys discriminate dtype device and extents" {
3010 const allocator = testing.allocator;
3011 const device = tuning.deviceFingerprint(sparseFamilyTuningTestCapabilities(0x2684));
3012
3013 const block_16 = try spmmCsrFamilyTuningKey(allocator, device, .{
3014 .rows = 70,
3015 .columns = 45,
3016 .nnz = 512,
3017 .x_extent = 40,
3018 .threads = .{ .x = 16, .y = 16 },
3019 });
3020 const block_8 = try spmmCsrFamilyTuningKey(allocator, device, .{
3021 .rows = 70,
3022 .columns = 45,
3023 .nnz = 512,
3024 .x_extent = 40,
3025 .threads = .{ .x = 8, .y = 8 },
3026 });
3027 try testing.expect(block_16.eql(block_8));
3028
3029 const other_columns = try spmmCsrFamilyTuningKey(allocator, device, .{
3030 .rows = 70,
3031 .columns = 46,
3032 .nnz = 512,
3033 .x_extent = 40,
3034 });
3035 try testing.expect(!block_16.eql(other_columns));
3036
3037 const other_nnz = try spmmCsrFamilyTuningKey(allocator, device, .{
3038 .rows = 70,
3039 .columns = 45,
3040 .nnz = 768,
3041 .x_extent = 40,
3042 });
3043 try testing.expect(!block_16.eql(other_nnz));
3044
3045 const half = try spmmCsrFamilyTuningKey(allocator, device, .{
3046 .rows = 70,
3047 .columns = 45,
3048 .nnz = 512,
3049 .x_extent = 40,
3050 .dtype = .f16,
3051 .accumulation_dtype = .f32,
3052 });
3053 try testing.expect(!block_16.eql(half));
3054
3055 const other_device = try spmmCsrFamilyTuningKey(
3056 allocator,
3057 tuning.deviceFingerprint(sparseFamilyTuningTestCapabilities(0x1b80)),
3058 .{
3059 .rows = 70,
3060 .columns = 45,
3061 .nnz = 512,
3062 .x_extent = 40,
3063 },
3064 );
3065 try testing.expect(!block_16.eql(other_device));
3066 }
3067
3068 test "sparse spmm csr family tuning resolves thread blocks" {
3069 const allocator = testing.allocator;
3070 const caps = sparseFamilyTuningTestCapabilities(0x2684);
3071 const device = tuning.deviceFingerprint(caps);
3072 const probe = SpmmCsr{
3073 .rows = 70,
3074 .columns = 45,
3075 .nnz = 512,
3076 .x_extent = 40,
3077 .threads = spmmCsrThreadsForExtents(70, 45),
3078 };
3079 const candidates = spmmCsrThreadCandidatesForExtents(probe.rows, probe.columns);
3080 try testing.expect(candidates.count > 1);
3081 var winner = probe;
3082 winner.threads = candidates.items[1];
3083 winner.threads = spmmCsrRepresentableThreads(winner) orelse return error.TestExpectedSparseStructure;
3084 const winner_target = try spmmCsrFamilyTarget(allocator, winner);
3085 defer allocator.free(winner_target);
3086
3087 const records = [_]tuning.FamilyTuningRecord{.{
3088 .key = try spmmCsrFamilyTuningKey(allocator, device, probe),
3089 .target = winner_target,
3090 .winner_median_ns = 800,
3091 .runner_up_median_ns = 1100,
3092 .sample_count = 30,
3093 }};
3094 const reader = tuning.FamilyTuningReader.init(caps, .{ .records = records[0..] });
3095
3096 const resolved = (try resolveSpmmCsrThreads(allocator, reader, probe)) orelse
3097 return error.TestExpectedSparseStructure;
3098 try testing.expectEqual(winner.threads.x, resolved.x);
3099 try testing.expectEqual(winner.threads.y, resolved.y);
3100
3101 const miss = try resolveSpmmCsrThreads(allocator, reader, .{
3102 .rows = 70,
3103 .columns = 46,
3104 .nnz = 512,
3105 .x_extent = 40,
3106 .threads = spmmCsrThreadsForExtents(70, 46),
3107 });
3108 try testing.expectEqual(@as(?entry.Threads2D, null), miss);
3109
3110 const stale_records = [_]tuning.FamilyTuningRecord{.{
3111 .key = try spmmCsrFamilyTuningKey(allocator, device, probe),
3112 .target = "accy.kernel.sparse.spmm_csr_row_column_thread_family_96x4_f32",
3113 .winner_median_ns = 800,
3114 .runner_up_median_ns = 1100,
3115 .sample_count = 30,
3116 }};
3117 const stale_reader = tuning.FamilyTuningReader.init(caps, .{ .records = stale_records[0..] });
3118 try testing.expectEqual(
3119 @as(?entry.Threads2D, null),
3120 try resolveSpmmCsrThreads(allocator, stale_reader, probe),
3121 );
3122 }
3123
3124 test "sparse spmv coo identity validity and artifact contract" {
3125 const allocator = testing.allocator;
3126 const instance = SpmvCoo{ .rows = 1000, .nnz = 4096, .x_extent = 64, .threads = 64 };
3127
3128 const target = try spmvCooFamilyTarget(allocator, instance);
3129 defer allocator.free(target);
3130 try testing.expectEqualStrings("accy.kernel.sparse.spmv_coo_element_thread_family_64_f32", target);
3131
3132 try testing.expectEqual(DType.f32, spmvCooAccumulationDType(.f32).?);
3133 try testing.expectEqual(DType.f64, spmvCooAccumulationDType(.f64).?);
3134 try testing.expectEqual(DType.f32, spmvCooAccumulationDType(.f16).?);
3135 try testing.expectEqual(DType.f32, spmvCooAccumulationDTypeForStructure(.row_thread, .f32).?);
3136 try testing.expectEqual(@as(?DType, null), spmvCooAccumulationDTypeForStructure(.element_thread, .f64));
3137 try testing.expectEqual(@as(?DType, null), spmvCooAccumulationDTypeForStructure(.element_thread, .f16));
3138 try testing.expect(spmvCooInstanceValid(instance));
3139 try testing.expect(!spmvCooInstanceValid(.{ .rows = 0, .nnz = 16, .x_extent = 4, .threads = 64 }));
3140 try testing.expect(!spmvCooInstanceValid(.{ .rows = 10, .nnz = 0, .x_extent = 4, .threads = 64 }));
3141 try testing.expect(!spmvCooInstanceValid(.{ .rows = 10, .nnz = 16, .x_extent = 0, .threads = 64 }));
3142 try testing.expect(!spmvCooInstanceValid(.{ .rows = 10, .nnz = 16, .x_extent = 4, .threads = 0 }));
3143 try testing.expect(!spmvCooInstanceValid(.{ .rows = 10, .nnz = 16, .x_extent = 4, .dtype = .f64, .accumulation_dtype = .f64, .threads = 64 }));
3144 try testing.expect(spmvCooInstanceValid(.{ .rows = 10, .nnz = 16, .x_extent = 4, .dtype = .f64, .accumulation_dtype = .f64, .threads = 64, .structure = .row_thread }));
3145 try testing.expect(spmvCooInstanceValid(.{ .rows = 10, .nnz = 16, .x_extent = 4, .dtype = .f16, .accumulation_dtype = .f32, .threads = 64, .structure = .row_thread }));
3146
3147 var state = gpu.recording.BackendState{
3148 .allocator = allocator,
3149 .kind = .cuda,
3150 .format = .cuda_ptx,
3151 };
3152 var family_artifact = try createSpmvCooFamilyArtifact(allocator, state.handle(), instance, .{ .limits = .testing });
3153 defer family_artifact.deinit();
3154 const family_entry = family_artifact.entry();
3155 try testing.expectEqualStrings("accy_kernel_sparse_spmv_coo_element_thread_family_64_f32", family_entry.entry_name);
3156 try testing.expectEqual(@as(u32, 8), family_entry.argument_count);
3157 try testing.expectEqual(@as(u32, 3), family_entry.runtime_scalar_argument_count);
3158 switch (family_entry.launch) {
3159 .derived => |launch| {
3160 try testing.expectEqual(@as(u32, 64), launch.threadgroup[0]);
3161 switch (launch.grid[0]) {
3162 .runtime_u32_ceil_div => |axis| {
3163 try testing.expectEqual(@as(u32, 1), axis.argument_index);
3164 try testing.expectEqual(@as(u32, 64), axis.divisor);
3165 },
3166 else => return error.TestExpectedDerivedGrid,
3167 }
3168 },
3169 else => return error.TestExpectedDerivedLaunch,
3170 }
3171 try testing.expect(family_entry.required_dtypes.contains(.f32));
3172 try testing.expect(family_entry.required_dtypes.contains(.i32));
3173
3174 var f64_family_artifact = try createSpmvCooFamilyArtifact(
3175 allocator,
3176 state.handle(),
3177 .{ .rows = 1000, .nnz = 4096, .x_extent = 64, .dtype = .f64, .accumulation_dtype = .f64, .threads = 64, .structure = .row_thread },
3178 .{ .limits = .testing },
3179 );
3180 defer f64_family_artifact.deinit();
3181 try testing.expectEqualStrings("accy_kernel_sparse_spmv_coo_row_thread_family_64_f64", f64_family_artifact.entry().entry_name);
3182
3183 var f16_family_artifact = try createSpmvCooFamilyArtifact(
3184 allocator,
3185 state.handle(),
3186 .{ .rows = 1000, .nnz = 4096, .x_extent = 64, .dtype = .f16, .accumulation_dtype = .f32, .threads = 64, .structure = .row_thread },
3187 .{ .limits = .testing },
3188 );
3189 defer f16_family_artifact.deinit();
3190 try testing.expectEqualStrings("accy_kernel_sparse_spmv_coo_row_thread_family_64_f16", f16_family_artifact.entry().entry_name);
3191 }
3192
3193 test "sparse spmv ell identity validity and artifact contract" {
3194 const allocator = testing.allocator;
3195 const instance = SpmvEll{ .rows = 1000, .slots = 8, .x_extent = 64, .threads = 64 };
3196
3197 const target = try spmvEllFamilyTarget(allocator, instance);
3198 defer allocator.free(target);
3199 try testing.expectEqualStrings("accy.kernel.sparse.spmv_ell_row_thread_family_64_f32", target);
3200 const f16_target = try spmvEllFamilyTarget(
3201 allocator,
3202 .{ .rows = 1000, .slots = 8, .x_extent = 64, .dtype = .f16, .accumulation_dtype = .f32, .threads = 64 },
3203 );
3204 defer allocator.free(f16_target);
3205 try testing.expectEqualStrings("accy.kernel.sparse.spmv_ell_row_thread_family_64_f16", f16_target);
3206 const f64_target = try spmvEllFamilyTarget(
3207 allocator,
3208 .{ .rows = 1000, .slots = 8, .x_extent = 64, .dtype = .f64, .accumulation_dtype = .f64, .threads = 64 },
3209 );
3210 defer allocator.free(f64_target);
3211 try testing.expectEqualStrings("accy.kernel.sparse.spmv_ell_row_thread_family_64_f64", f64_target);
3212
3213 try testing.expectEqual(DType.f64, spmvEllAccumulationDType(.f64).?);
3214 try testing.expectEqual(DType.f32, spmvEllAccumulationDType(.f16).?);
3215 try testing.expect(spmvEllInstanceValid(instance));
3216 try testing.expect(spmvEllInstanceValid(.{ .rows = 10, .slots = 3, .x_extent = 4, .dtype = .f64, .accumulation_dtype = .f64, .threads = 64 }));
3217 try testing.expect(!spmvEllInstanceValid(.{ .rows = 10, .slots = 3, .x_extent = 4, .dtype = .f64, .accumulation_dtype = .f32, .threads = 64 }));
3218 try testing.expect(spmvEllInstanceValid(.{ .rows = 10, .slots = 3, .x_extent = 4, .dtype = .f16, .accumulation_dtype = .f32, .threads = 64 }));
3219 try testing.expect(!spmvEllInstanceValid(.{ .rows = 10, .slots = 3, .x_extent = 4, .dtype = .f16, .accumulation_dtype = .f16, .threads = 64 }));
3220 try testing.expect(!spmvEllInstanceValid(.{ .rows = 0, .slots = 3, .x_extent = 4, .threads = 64 }));
3221 try testing.expect(!spmvEllInstanceValid(.{ .rows = 10, .slots = 0, .x_extent = 4, .threads = 64 }));
3222 try testing.expect(!spmvEllInstanceValid(.{ .rows = 10, .slots = 3, .x_extent = 0, .threads = 64 }));
3223 try testing.expect(!spmvEllInstanceValid(.{ .rows = 10, .slots = 3, .x_extent = 4, .threads = 0 }));
3224
3225 var state = gpu.recording.BackendState{
3226 .allocator = allocator,
3227 .kind = .cuda,
3228 .format = .cuda_ptx,
3229 };
3230 var family_artifact = try createSpmvEllFamilyArtifact(allocator, state.handle(), instance, .{ .limits = .testing });
3231 defer family_artifact.deinit();
3232 const family_entry = family_artifact.entry();
3233 try testing.expectEqualStrings("accy_kernel_sparse_spmv_ell_row_thread_family_64_f32", family_entry.entry_name);
3234 try testing.expectEqual(@as(u32, 7), family_entry.argument_count);
3235 try testing.expectEqual(@as(u32, 3), family_entry.runtime_scalar_argument_count);
3236 switch (family_entry.launch) {
3237 .derived => |launch| {
3238 try testing.expectEqual(@as(u32, 64), launch.threadgroup[0]);
3239 switch (launch.grid[0]) {
3240 .runtime_u32_ceil_div => |axis| try testing.expectEqual(@as(u32, 64), axis.divisor),
3241 else => return error.TestExpectedDerivedGrid,
3242 }
3243 },
3244 else => return error.TestExpectedDerivedLaunch,
3245 }
3246
3247 var f16_family_artifact = try createSpmvEllFamilyArtifact(
3248 allocator,
3249 state.handle(),
3250 .{ .rows = 1000, .slots = 8, .x_extent = 64, .dtype = .f16, .accumulation_dtype = .f32, .threads = 64 },
3251 .{ .limits = .testing },
3252 );
3253 defer f16_family_artifact.deinit();
3254 const f16_family_entry = f16_family_artifact.entry();
3255 try testing.expectEqualStrings("accy_kernel_sparse_spmv_ell_row_thread_family_64_f16", f16_family_entry.entry_name);
3256 try testing.expect(f16_family_entry.required_dtypes.contains(.f16));
3257 try testing.expect(f16_family_entry.required_dtypes.contains(.i32));
3258
3259 var f64_family_artifact = try createSpmvEllFamilyArtifact(
3260 allocator,
3261 state.handle(),
3262 .{ .rows = 1000, .slots = 8, .x_extent = 64, .dtype = .f64, .accumulation_dtype = .f64, .threads = 64 },
3263 .{ .limits = .testing },
3264 );
3265 defer f64_family_artifact.deinit();
3266 const f64_family_entry = f64_family_artifact.entry();
3267 try testing.expectEqualStrings("accy_kernel_sparse_spmv_ell_row_thread_family_64_f64", f64_family_entry.entry_name);
3268 try testing.expect(f64_family_entry.required_dtypes.contains(.f64));
3269 try testing.expect(f64_family_entry.required_dtypes.contains(.i32));
3270 }
3271
3272 test "sparse spmv sell identity validity and artifact contract" {
3273 const allocator = testing.allocator;
3274 const instance = SpmvSell{ .rows = 1000, .slice_size = 32, .values_size = 4096, .x_extent = 64, .threads = 64 };
3275
3276 const target = try spmvSellFamilyTarget(allocator, instance);
3277 defer allocator.free(target);
3278 try testing.expectEqualStrings("accy.kernel.sparse.spmv_sell_row_thread_slice32_family_64_f32", target);
3279 const f16_target = try spmvSellFamilyTarget(
3280 allocator,
3281 .{ .rows = 1000, .slice_size = 32, .values_size = 4096, .x_extent = 64, .dtype = .f16, .accumulation_dtype = .f32, .threads = 64 },
3282 );
3283 defer allocator.free(f16_target);
3284 try testing.expectEqualStrings("accy.kernel.sparse.spmv_sell_row_thread_slice32_family_64_f16", f16_target);
3285 const f64_target = try spmvSellFamilyTarget(
3286 allocator,
3287 .{ .rows = 1000, .slice_size = 32, .values_size = 4096, .x_extent = 64, .dtype = .f64, .accumulation_dtype = .f64, .threads = 64 },
3288 );
3289 defer allocator.free(f64_target);
3290 try testing.expectEqualStrings("accy.kernel.sparse.spmv_sell_row_thread_slice32_family_64_f64", f64_target);
3291
3292 try testing.expectEqual(DType.f64, spmvSellAccumulationDType(.f64).?);
3293 try testing.expectEqual(DType.f32, spmvSellAccumulationDType(.f16).?);
3294 try testing.expectEqual(@as(u64, 32), spmvSellSliceCount(instance));
3295 try testing.expect(spmvSellInstanceValid(instance));
3296 try testing.expect(spmvSellInstanceValid(.{ .rows = 10, .slice_size = 4, .values_size = 24, .x_extent = 4, .dtype = .f64, .accumulation_dtype = .f64, .threads = 64 }));
3297 try testing.expect(!spmvSellInstanceValid(.{ .rows = 10, .slice_size = 4, .values_size = 24, .x_extent = 4, .dtype = .f64, .accumulation_dtype = .f32, .threads = 64 }));
3298 try testing.expect(spmvSellInstanceValid(.{ .rows = 10, .slice_size = 4, .values_size = 24, .x_extent = 4, .dtype = .f16, .accumulation_dtype = .f32, .threads = 64 }));
3299 try testing.expect(!spmvSellInstanceValid(.{ .rows = 10, .slice_size = 4, .values_size = 24, .x_extent = 4, .dtype = .f16, .accumulation_dtype = .f16, .threads = 64 }));
3300 try testing.expect(!spmvSellInstanceValid(.{ .rows = 0, .slice_size = 4, .values_size = 24, .x_extent = 4, .threads = 64 }));
3301 try testing.expect(!spmvSellInstanceValid(.{ .rows = 10, .slice_size = 0, .values_size = 24, .x_extent = 4, .threads = 64 }));
3302 try testing.expect(!spmvSellInstanceValid(.{ .rows = 10, .slice_size = 4, .values_size = 0, .x_extent = 4, .threads = 64 }));
3303 try testing.expect(!spmvSellInstanceValid(.{ .rows = 10, .slice_size = 4, .values_size = 24, .x_extent = 0, .threads = 64 }));
3304 try testing.expect(!spmvSellInstanceValid(.{ .rows = 10, .slice_size = 4, .values_size = 24, .x_extent = 4, .threads = 0 }));
3305
3306 var state = gpu.recording.BackendState{
3307 .allocator = allocator,
3308 .kind = .cuda,
3309 .format = .cuda_ptx,
3310 };
3311 var family_artifact = try createSpmvSellFamilyArtifact(allocator, state.handle(), instance, .{ .limits = .testing });
3312 defer family_artifact.deinit();
3313 const family_entry = family_artifact.entry();
3314 try testing.expectEqualStrings("accy_kernel_sparse_spmv_sell_row_thread_slice32_family_64_f32", family_entry.entry_name);
3315 try testing.expectEqual(@as(u32, 8), family_entry.argument_count);
3316 try testing.expectEqual(@as(u32, 3), family_entry.runtime_scalar_argument_count);
3317 switch (family_entry.launch) {
3318 .derived => |launch| {
3319 try testing.expectEqual(@as(u32, 64), launch.threadgroup[0]);
3320 switch (launch.grid[0]) {
3321 .runtime_u32_ceil_div => |axis| try testing.expectEqual(@as(u32, 64), axis.divisor),
3322 else => return error.TestExpectedDerivedGrid,
3323 }
3324 },
3325 else => return error.TestExpectedDerivedLaunch,
3326 }
3327
3328 var f16_family_artifact = try createSpmvSellFamilyArtifact(
3329 allocator,
3330 state.handle(),
3331 .{ .rows = 1000, .slice_size = 32, .values_size = 4096, .x_extent = 64, .dtype = .f16, .accumulation_dtype = .f32, .threads = 64 },
3332 .{ .limits = .testing },
3333 );
3334 defer f16_family_artifact.deinit();
3335 const f16_family_entry = f16_family_artifact.entry();
3336 try testing.expectEqualStrings("accy_kernel_sparse_spmv_sell_row_thread_slice32_family_64_f16", f16_family_entry.entry_name);
3337 try testing.expect(f16_family_entry.required_dtypes.contains(.f16));
3338 try testing.expect(f16_family_entry.required_dtypes.contains(.i32));
3339
3340 var f64_family_artifact = try createSpmvSellFamilyArtifact(
3341 allocator,
3342 state.handle(),
3343 .{ .rows = 1000, .slice_size = 32, .values_size = 4096, .x_extent = 64, .dtype = .f64, .accumulation_dtype = .f64, .threads = 64 },
3344 .{ .limits = .testing },
3345 );
3346 defer f64_family_artifact.deinit();
3347 const f64_family_entry = f64_family_artifact.entry();
3348 try testing.expectEqualStrings("accy_kernel_sparse_spmv_sell_row_thread_slice32_family_64_f64", f64_family_entry.entry_name);
3349 try testing.expect(f64_family_entry.required_dtypes.contains(.f64));
3350 try testing.expect(f64_family_entry.required_dtypes.contains(.i32));
3351 }
3352
3353 test "sparse spmm csr identity validity and artifact contract" {
3354 const allocator = testing.allocator;
3355 const instance = SpmmCsr{
3356 .rows = 1000,
3357 .columns = 37,
3358 .x_extent = 80,
3359 .threads = .{ .x = 8, .y = 4 },
3360 };
3361
3362 const target = try spmmCsrFamilyTarget(allocator, instance);
3363 defer allocator.free(target);
3364 try testing.expectEqualStrings("accy.kernel.sparse.spmm_csr_row_column_thread_family_8x4_f32", target);
3365 const f16_target = try spmmCsrFamilyTarget(allocator, .{
3366 .rows = 1000,
3367 .columns = 37,
3368 .x_extent = 80,
3369 .dtype = .f16,
3370 .accumulation_dtype = .f32,
3371 .threads = .{ .x = 8, .y = 4 },
3372 });
3373 defer allocator.free(f16_target);
3374 try testing.expectEqualStrings("accy.kernel.sparse.spmm_csr_row_column_thread_family_8x4_f16", f16_target);
3375 const f64_target = try spmmCsrFamilyTarget(allocator, .{
3376 .rows = 1000,
3377 .columns = 37,
3378 .x_extent = 80,
3379 .dtype = .f64,
3380 .accumulation_dtype = .f64,
3381 .threads = .{ .x = 8, .y = 4 },
3382 });
3383 defer allocator.free(f64_target);
3384 try testing.expectEqualStrings("accy.kernel.sparse.spmm_csr_row_column_thread_family_8x4_f64", f64_target);
3385
3386 try testing.expectEqual(DType.f64, spmmCsrAccumulationDType(.f64).?);
3387 try testing.expectEqual(DType.f32, spmmCsrAccumulationDType(.f16).?);
3388 try testing.expect(spmmCsrInstanceValid(instance));
3389 try testing.expect(spmmCsrInstanceValid(.{ .rows = 10, .columns = 3, .dtype = .f64, .accumulation_dtype = .f64 }));
3390 try testing.expect(!spmmCsrInstanceValid(.{ .rows = 10, .columns = 3, .dtype = .f64, .accumulation_dtype = .f32 }));
3391 try testing.expect(spmmCsrInstanceValid(.{ .rows = 10, .columns = 3, .dtype = .f16, .accumulation_dtype = .f32 }));
3392 try testing.expect(!spmmCsrInstanceValid(.{ .rows = 10, .columns = 3, .dtype = .f16, .accumulation_dtype = .f16 }));
3393 try testing.expect(!spmmCsrInstanceValid(.{ .rows = 0, .columns = 3 }));
3394 try testing.expect(!spmmCsrInstanceValid(.{ .rows = 10, .columns = 0 }));
3395 try testing.expect(!spmmCsrInstanceValid(.{ .rows = 10, .columns = 3, .threads = .{ .x = 0, .y = 4 } }));
3396 try testing.expect(!spmmCsrInstanceValid(.{ .rows = 10, .columns = 3, .threads = .{ .x = 33, .y = 33 } }));
3397
3398 var state = gpu.recording.BackendState{
3399 .allocator = allocator,
3400 .kind = .cuda,
3401 .format = .cuda_ptx,
3402 };
3403 var family_artifact = try createSpmmCsrFamilyArtifact(allocator, state.handle(), instance, .{ .limits = .testing });
3404 defer family_artifact.deinit();
3405 const family_entry = family_artifact.entry();
3406 try testing.expectEqualStrings("accy_kernel_sparse_spmm_csr_row_column_thread_family_8x4_f32", family_entry.entry_name);
3407 try testing.expectEqual(@as(u32, 9), family_entry.argument_count);
3408 try testing.expectEqual(@as(u32, 4), family_entry.runtime_scalar_argument_count);
3409 switch (family_entry.launch) {
3410 .derived => |launch| {
3411 try testing.expectEqual(@as(u32, 8), launch.threadgroup[0]);
3412 try testing.expectEqual(@as(u32, 4), launch.threadgroup[1]);
3413 switch (launch.grid[0]) {
3414 .runtime_u32_ceil_div => |axis| {
3415 try testing.expectEqual(@as(u32, 3), axis.argument_index);
3416 try testing.expectEqual(@as(u32, 8), axis.divisor);
3417 },
3418 else => return error.TestExpectedDerivedGrid,
3419 }
3420 switch (launch.grid[1]) {
3421 .runtime_u32_ceil_div => |axis| {
3422 try testing.expectEqual(@as(u32, 0), axis.argument_index);
3423 try testing.expectEqual(@as(u32, 4), axis.divisor);
3424 },
3425 else => return error.TestExpectedDerivedGrid,
3426 }
3427 },
3428 else => return error.TestExpectedDerivedLaunch,
3429 }
3430 const profile = family_entry.shape_profile orelse return error.TestExpectedShapeProfile;
3431 try testing.expectEqual(@as(usize, 2), profile.dimensions.len);
3432
3433 var f16_family_artifact = try createSpmmCsrFamilyArtifact(
3434 allocator,
3435 state.handle(),
3436 .{ .rows = 1000, .columns = 37, .x_extent = 80, .dtype = .f16, .accumulation_dtype = .f32, .threads = .{ .x = 8, .y = 4 } },
3437 .{ .limits = .testing },
3438 );
3439 defer f16_family_artifact.deinit();
3440 const f16_family_entry = f16_family_artifact.entry();
3441 try testing.expectEqualStrings("accy_kernel_sparse_spmm_csr_row_column_thread_family_8x4_f16", f16_family_entry.entry_name);
3442 try testing.expect(f16_family_entry.required_dtypes.contains(.f16));
3443 try testing.expect(f16_family_entry.required_dtypes.contains(.i32));
3444
3445 var f64_family_artifact = try createSpmmCsrFamilyArtifact(
3446 allocator,
3447 state.handle(),
3448 .{ .rows = 1000, .columns = 37, .x_extent = 80, .dtype = .f64, .accumulation_dtype = .f64, .threads = .{ .x = 8, .y = 4 } },
3449 .{ .limits = .testing },
3450 );
3451 defer f64_family_artifact.deinit();
3452 const f64_family_entry = f64_family_artifact.entry();
3453 try testing.expectEqualStrings("accy_kernel_sparse_spmm_csr_row_column_thread_family_8x4_f64", f64_family_entry.entry_name);
3454 try testing.expect(f64_family_entry.required_dtypes.contains(.f64));
3455 try testing.expect(f64_family_entry.required_dtypes.contains(.i32));
3456 }
3457
3458 pub const spmv_csr_nonzero_axis = "n";
3459 pub const spmv_csr_column_axis = "x";
3460 pub const spmv_sell_slice_size_parameter = "slice_size";
3461
3462 pub fn spmvCsrLaunchExtent(instance: SpmvCsr) u64 {
3463 return spmvCsrLaunchExtentChecked(instance).?;
3464 }
3465
3466 fn spmvCsrLaunchExtentChecked(instance: SpmvCsr) ?u64 {
3467 return switch (instance.structure) {
3468 .row_thread => instance.rows,
3469 .row_warp => std.math.mul(u64, instance.rows, spmv_csr_warp_size) catch null,
3470 };
3471 }
3472
3473 pub fn spmvCsrRepresentableThreads(instance: SpmvCsr) ?u32 {
3474 const launch_extent = spmvCsrLaunchExtentChecked(instance) orelse return null;
3475 return std.math.cast(u32, @min(@as(u64, instance.threads), launch_extent));
3476 }
3477
3478 pub fn spmvCsrFamilySpecialization(
3479 backing_allocator: std.mem.Allocator,
3480 instance: SpmvCsr,
3481 ) !entry.OwnedSpecialization {
3482 var owned = entry.OwnedSpecialization.init(backing_allocator);
3483 errdefer owned.deinit();
3484 const lifetime_allocator = owned.allocator();
3485
3486 const inputs = try lifetime_allocator.alloc(entry.Shape, 4);
3487 inputs[0] = try entry.runtimeShape1D(lifetime_allocator, instance.row_axis, instance.rows + 1);
3488 inputs[1] = try entry.runtimeShape1D(lifetime_allocator, spmv_csr_nonzero_axis, instance.nnz);
3489 inputs[2] = try entry.runtimeShape1D(lifetime_allocator, spmv_csr_nonzero_axis, instance.nnz);
3490 inputs[3] = try entry.runtimeShape1D(lifetime_allocator, spmv_csr_column_axis, instance.x_extent);
3491
3492 const outputs = try lifetime_allocator.alloc(entry.Shape, 1);
3493 outputs[0] = try entry.runtimeShape1D(lifetime_allocator, instance.row_axis, instance.rows);
3494
3495 owned.value = .{
3496 .dtype = instance.dtype,
3497 .accumulation_dtype = instance.accumulation_dtype,
3498 .operation = .{ .sparse = .csr_spmv },
3499 .inputs = inputs,
3500 .outputs = outputs,
3501 .schedule = try entry.runtimeThreadBlocks1D(
3502 lifetime_allocator,
3503 instance.row_axis,
3504 spmvCsrLaunchExtent(instance),
3505 instance.threads,
3506 ),
3507 .structure = @tagName(instance.structure),
3508 };
3509 owned.value.launch = owned.value.schedule.?.launch();
3510 var family = try spmvCsrShapeFamily(backing_allocator, instance);
3511 errdefer family.deinit();
3512 try owned.takeShapeFamily(&family);
3513 return owned;
3514 }
3515
3516 pub fn spmvCsrInstanceFromSpecialization(specialization: entry.Specialization) ?SpmvCsr {
3517 if (!specialization.scheduleMatchesLaunch()) return null;
3518 const schedule = specialization.schedule orelse return null;
3519 if (!specialization.operationIs(.{ .sparse = .csr_spmv })) return null;
3520 const dtype = specialization.dtype orelse return null;
3521 const accumulation_dtype = spmvCsrAccumulationDType(dtype) orelse return null;
3522 if (specialization.accumulation_dtype != accumulation_dtype) return null;
3523 if (specialization.inputs.len != 4 or specialization.outputs.len != 1) return null;
3524 if (specialization.reductions.len != 0) return null;
3525 const row_ptr = specialization.inputs[0];
3526 const cols = specialization.inputs[1];
3527 const values = specialization.inputs[2];
3528 const x = specialization.inputs[3];
3529 const y = specialization.outputs[0];
3530 if (row_ptr.axes.len != 1 or cols.axes.len != 1 or values.axes.len != 1 or x.axes.len != 1 or y.axes.len != 1) return null;
3531 if (y.axes[0].name.len == 0) return null;
3532 if (cols.axes[0].extent == 0 or x.axes[0].extent == 0) return null;
3533 const row_ptr_extent = std.math.add(u64, y.axes[0].extent, 1) catch return null;
3534 if (!sparseAxisMatches(row_ptr.axes[0], y.axes[0].name, row_ptr_extent)) return null;
3535 if (!sparseAxisMatches(cols.axes[0], spmv_csr_nonzero_axis, cols.axes[0].extent)) return null;
3536 if (!sparseAxisMatches(values.axes[0], spmv_csr_nonzero_axis, cols.axes[0].extent)) return null;
3537 if (!sparseAxisMatches(x.axes[0], spmv_csr_column_axis, x.axes[0].extent)) return null;
3538 const launch = specialization.launch orelse return null;
3539 if (launch.threadgroup[0] == 0) return null;
3540 const structure_name = specialization.structure orelse return null;
3541 const structure = std.meta.stringToEnum(SpmvCsrStructure, structure_name) orelse return null;
3542 const instance = SpmvCsr{
3543 .rows = y.axes[0].extent,
3544 .nnz = cols.axes[0].extent,
3545 .x_extent = x.axes[0].extent,
3546 .dtype = dtype,
3547 .accumulation_dtype = accumulation_dtype,
3548 .threads = launch.threadgroup[0],
3549 .structure = structure,
3550 .row_axis = y.axes[0].name,
3551 };
3552 if (!spmvCsrInstanceValid(instance)) return null;
3553 if (!spmvCsrScheduleMatchesInstance(schedule, instance)) return null;
3554 return instance;
3555 }
3556
3557 fn sparseAxisMatches(axis: entry.Axis, name: []const u8, extent: u64) bool {
3558 if (name.len == 0) return false;
3559 if (!std.mem.eql(u8, axis.name, name)) return false;
3560 return axis.extent == extent;
3561 }
3562
3563 fn sparseScheduleAxisNameMatches(actual: []const u8, row_axis: []const u8, suffix: []const u8) bool {
3564 if (actual.len != row_axis.len + suffix.len) return false;
3565 return std.mem.eql(u8, actual[0..row_axis.len], row_axis) and
3566 std.mem.eql(u8, actual[row_axis.len..], suffix);
3567 }
3568
3569 fn spmvCsrScheduleMatchesInstance(schedule: entry.Schedule, instance: SpmvCsr) bool {
3570 const extent = spmvCsrLaunchExtentChecked(instance) orelse return false;
3571 if (extent <= instance.threads) {
3572 if (schedule.bindings.len != 1) return false;
3573 const binding = schedule.bindings[0];
3574 return sparseAxisMatches(.{ .name = binding.axis, .extent = binding.extent }, instance.row_axis, extent) and
3575 binding.target == .thread_x;
3576 }
3577 if (schedule.bindings.len != 2) return false;
3578 const tile = schedule.bindings[0];
3579 const lane = schedule.bindings[1];
3580 const blocks = spmvCsrBlockCountChecked(instance) orelse return false;
3581 return tile.target == .block_x and
3582 tile.extent == blocks and
3583 sparseScheduleAxisNameMatches(tile.axis, instance.row_axis, "_tile") and
3584 lane.target == .thread_x and
3585 lane.extent == instance.threads and
3586 sparseScheduleAxisNameMatches(lane.axis, instance.row_axis, "_lane");
3587 }
3588
3589 pub fn spmvCooLaunchExtent(instance: SpmvCoo) u64 {
3590 return spmvCooLaunchExtentChecked(instance).?;
3591 }
3592
3593 fn spmvCooLaunchExtentChecked(instance: SpmvCoo) ?u64 {
3594 return switch (instance.structure) {
3595 .element_thread => instance.nnz,
3596 .row_thread => instance.rows,
3597 };
3598 }
3599
3600 pub fn spmvCooRepresentableThreads(instance: SpmvCoo) ?u32 {
3601 const launch_extent = spmvCooLaunchExtentChecked(instance) orelse return null;
3602 return std.math.cast(u32, @min(@as(u64, instance.threads), launch_extent));
3603 }
3604
3605 fn spmvCooLaunchArgumentIndex(instance: SpmvCoo) u32 {
3606 return switch (instance.structure) {
3607 .element_thread => 1,
3608 .row_thread => 0,
3609 };
3610 }
3611
3612 fn spmvCooScheduleAxis(instance: SpmvCoo) []const u8 {
3613 return switch (instance.structure) {
3614 .element_thread => instance.nonzero_axis,
3615 .row_thread => instance.row_axis,
3616 };
3617 }
3618
3619 pub fn spmvCooFamilySpecialization(
3620 backing_allocator: std.mem.Allocator,
3621 instance: SpmvCoo,
3622 ) !entry.OwnedSpecialization {
3623 var owned = entry.OwnedSpecialization.init(backing_allocator);
3624 errdefer owned.deinit();
3625 const lifetime_allocator = owned.allocator();
3626
3627 const inputs = try lifetime_allocator.alloc(entry.Shape, 4);
3628 inputs[0] = try entry.runtimeShape1D(lifetime_allocator, instance.nonzero_axis, instance.nnz);
3629 inputs[1] = try entry.runtimeShape1D(lifetime_allocator, instance.nonzero_axis, instance.nnz);
3630 inputs[2] = try entry.runtimeShape1D(lifetime_allocator, instance.nonzero_axis, instance.nnz);
3631 inputs[3] = try entry.runtimeShape1D(lifetime_allocator, instance.x_axis, instance.x_extent);
3632
3633 const outputs = try lifetime_allocator.alloc(entry.Shape, 1);
3634 outputs[0] = try entry.runtimeShape1D(lifetime_allocator, instance.row_axis, instance.rows);
3635
3636 owned.value = .{
3637 .dtype = instance.dtype,
3638 .accumulation_dtype = instance.accumulation_dtype,
3639 .operation = .{ .sparse = .coo_spmv },
3640 .inputs = inputs,
3641 .outputs = outputs,
3642 .schedule = try entry.runtimeThreadBlocks1D(
3643 lifetime_allocator,
3644 spmvCooScheduleAxis(instance),
3645 spmvCooLaunchExtent(instance),
3646 instance.threads,
3647 ),
3648 .structure = @tagName(instance.structure),
3649 };
3650 owned.value.launch = owned.value.schedule.?.launch();
3651 var family = try spmvCooShapeFamily(backing_allocator, instance);
3652 errdefer family.deinit();
3653 try owned.takeShapeFamily(&family);
3654 return owned;
3655 }
3656
3657 pub fn spmvCooInstanceFromSpecialization(specialization: entry.Specialization) ?SpmvCoo {
3658 if (!specialization.scheduleMatchesLaunch()) return null;
3659 const schedule = specialization.schedule orelse return null;
3660 if (!specialization.operationIs(.{ .sparse = .coo_spmv })) return null;
3661 const dtype = specialization.dtype orelse return null;
3662 if (specialization.inputs.len != 4 or specialization.outputs.len != 1) return null;
3663 if (specialization.reductions.len != 0) return null;
3664 const row_indices = specialization.inputs[0];
3665 const cols = specialization.inputs[1];
3666 const values = specialization.inputs[2];
3667 const x = specialization.inputs[3];
3668 const y = specialization.outputs[0];
3669 if (row_indices.axes.len != 1 or cols.axes.len != 1 or values.axes.len != 1 or x.axes.len != 1 or y.axes.len != 1) return null;
3670 if (row_indices.axes[0].name.len == 0 or x.axes[0].name.len == 0 or y.axes[0].name.len == 0) return null;
3671 if (row_indices.axes[0].extent == 0 or x.axes[0].extent == 0 or y.axes[0].extent == 0) return null;
3672 if (!sparseAxisMatches(cols.axes[0], row_indices.axes[0].name, row_indices.axes[0].extent)) return null;
3673 if (!sparseAxisMatches(values.axes[0], row_indices.axes[0].name, row_indices.axes[0].extent)) return null;
3674 const launch = specialization.launch orelse return null;
3675 if (launch.threadgroup[0] == 0) return null;
3676 const structure_name = specialization.structure orelse return null;
3677 const structure = std.meta.stringToEnum(SpmvCooStructure, structure_name) orelse return null;
3678 const accumulation_dtype = spmvCooAccumulationDTypeForStructure(structure, dtype) orelse return null;
3679 if (specialization.accumulation_dtype != accumulation_dtype) return null;
3680 const instance = SpmvCoo{
3681 .rows = y.axes[0].extent,
3682 .nnz = row_indices.axes[0].extent,
3683 .x_extent = x.axes[0].extent,
3684 .dtype = dtype,
3685 .accumulation_dtype = accumulation_dtype,
3686 .threads = launch.threadgroup[0],
3687 .structure = structure,
3688 .row_axis = y.axes[0].name,
3689 .nonzero_axis = row_indices.axes[0].name,
3690 .x_axis = x.axes[0].name,
3691 };
3692 if (!spmvCooInstanceValid(instance)) return null;
3693 if (!spmvCooScheduleMatchesInstance(schedule, instance)) return null;
3694 return instance;
3695 }
3696
3697 fn spmvCooScheduleMatchesInstance(schedule: entry.Schedule, instance: SpmvCoo) bool {
3698 const extent = spmvCooLaunchExtentChecked(instance) orelse return false;
3699 const axis = spmvCooScheduleAxis(instance);
3700 if (extent <= instance.threads) {
3701 if (schedule.bindings.len != 1) return false;
3702 const binding = schedule.bindings[0];
3703 return sparseAxisMatches(.{ .name = binding.axis, .extent = binding.extent }, axis, extent) and
3704 binding.target == .thread_x;
3705 }
3706 if (schedule.bindings.len != 2) return false;
3707 const tile = schedule.bindings[0];
3708 const lane = schedule.bindings[1];
3709 const blocks = spmvCooBlockCountChecked(instance) orelse return false;
3710 return tile.target == .block_x and
3711 tile.extent == blocks and
3712 sparseScheduleAxisNameMatches(tile.axis, axis, "_tile") and
3713 lane.target == .thread_x and
3714 lane.extent == instance.threads and
3715 sparseScheduleAxisNameMatches(lane.axis, axis, "_lane");
3716 }
3717
3718 pub fn spmvEllLaunchExtent(instance: SpmvEll) u64 {
3719 return spmvEllLaunchExtentChecked(instance).?;
3720 }
3721
3722 fn spmvEllLaunchExtentChecked(instance: SpmvEll) ?u64 {
3723 return switch (instance.structure) {
3724 .row_thread => instance.rows,
3725 };
3726 }
3727
3728 pub fn spmvEllRepresentableThreads(instance: SpmvEll) ?u32 {
3729 const launch_extent = spmvEllLaunchExtentChecked(instance) orelse return null;
3730 if (launch_extent > std.math.maxInt(u32)) return null;
3731 return @min(instance.threads, @as(u32, @intCast(launch_extent)));
3732 }
3733
3734 pub fn spmvEllThreadCandidatesForRows(rows: u64) geometry_mod.Thread1DCandidates {
3735 return geometry_mod.threadCandidatesForExtent(rows, spmv_ell_thread_caps);
3736 }
3737
3738 pub fn spmvEllFamilySpecialization(
3739 backing_allocator: std.mem.Allocator,
3740 instance: SpmvEll,
3741 ) !entry.OwnedSpecialization {
3742 var owned = entry.OwnedSpecialization.init(backing_allocator);
3743 errdefer owned.deinit();
3744 const lifetime_allocator = owned.allocator();
3745
3746 const inputs = try lifetime_allocator.alloc(entry.Shape, 3);
3747 inputs[0] = try entry.runtimeShape2D(
3748 lifetime_allocator,
3749 instance.slot_axis,
3750 instance.slots,
3751 instance.row_axis,
3752 instance.rows,
3753 );
3754 inputs[1] = try entry.runtimeShape2D(
3755 lifetime_allocator,
3756 instance.slot_axis,
3757 instance.slots,
3758 instance.row_axis,
3759 instance.rows,
3760 );
3761 inputs[2] = try entry.runtimeShape1D(lifetime_allocator, instance.x_axis, instance.x_extent);
3762
3763 const outputs = try lifetime_allocator.alloc(entry.Shape, 1);
3764 outputs[0] = try entry.runtimeShape1D(lifetime_allocator, instance.row_axis, instance.rows);
3765
3766 owned.value = .{
3767 .dtype = instance.dtype,
3768 .accumulation_dtype = instance.accumulation_dtype,
3769 .operation = .{ .sparse = .ell_spmv },
3770 .inputs = inputs,
3771 .outputs = outputs,
3772 .schedule = try entry.runtimeThreadBlocks1D(
3773 lifetime_allocator,
3774 instance.row_axis,
3775 spmvEllLaunchExtent(instance),
3776 instance.threads,
3777 ),
3778 .structure = @tagName(instance.structure),
3779 };
3780 owned.value.launch = owned.value.schedule.?.launch();
3781 var family = try spmvEllShapeFamily(backing_allocator, instance);
3782 errdefer family.deinit();
3783 try owned.takeShapeFamily(&family);
3784 return owned;
3785 }
3786
3787 pub fn spmvEllInstanceFromSpecialization(specialization: entry.Specialization) ?SpmvEll {
3788 if (!specialization.scheduleMatchesLaunch()) return null;
3789 const schedule = specialization.schedule orelse return null;
3790 if (!specialization.operationIs(.{ .sparse = .ell_spmv })) return null;
3791 const dtype = specialization.dtype orelse return null;
3792 const accumulation_dtype = spmvEllAccumulationDType(dtype) orelse return null;
3793 if (specialization.accumulation_dtype != accumulation_dtype) return null;
3794 if (specialization.inputs.len != 3 or specialization.outputs.len != 1) return null;
3795 if (specialization.reductions.len != 0) return null;
3796 const cols = specialization.inputs[0];
3797 const values = specialization.inputs[1];
3798 const x = specialization.inputs[2];
3799 const y = specialization.outputs[0];
3800 if (cols.axes.len != 2 or values.axes.len != 2 or x.axes.len != 1 or y.axes.len != 1) return null;
3801 if (cols.axes[0].name.len == 0 or y.axes[0].name.len == 0 or x.axes[0].name.len == 0) return null;
3802 if (cols.axes[0].extent == 0 or x.axes[0].extent == 0) return null;
3803 if (!sparseAxisMatches(cols.axes[1], y.axes[0].name, y.axes[0].extent)) return null;
3804 if (!sparseAxisMatches(values.axes[0], cols.axes[0].name, cols.axes[0].extent)) return null;
3805 if (!sparseAxisMatches(values.axes[1], y.axes[0].name, y.axes[0].extent)) return null;
3806 const launch = specialization.launch orelse return null;
3807 if (launch.threadgroup[0] == 0) return null;
3808 const structure_name = specialization.structure orelse return null;
3809 const structure = std.meta.stringToEnum(SpmvEllStructure, structure_name) orelse return null;
3810 const instance = SpmvEll{
3811 .rows = y.axes[0].extent,
3812 .slots = cols.axes[0].extent,
3813 .x_extent = x.axes[0].extent,
3814 .dtype = dtype,
3815 .accumulation_dtype = accumulation_dtype,
3816 .threads = launch.threadgroup[0],
3817 .structure = structure,
3818 .row_axis = y.axes[0].name,
3819 .slot_axis = cols.axes[0].name,
3820 .x_axis = x.axes[0].name,
3821 };
3822 if (!spmvEllInstanceValid(instance)) return null;
3823 if (!spmvEllScheduleMatchesInstance(schedule, instance)) return null;
3824 return instance;
3825 }
3826
3827 fn spmvEllScheduleMatchesInstance(schedule: entry.Schedule, instance: SpmvEll) bool {
3828 const extent = spmvEllLaunchExtentChecked(instance) orelse return false;
3829 if (extent <= instance.threads) {
3830 if (schedule.bindings.len != 1) return false;
3831 const binding = schedule.bindings[0];
3832 return sparseAxisMatches(.{ .name = binding.axis, .extent = binding.extent }, instance.row_axis, extent) and
3833 binding.target == .thread_x;
3834 }
3835 if (schedule.bindings.len != 2) return false;
3836 const tile = schedule.bindings[0];
3837 const lane = schedule.bindings[1];
3838 const blocks = spmvEllBlockCountChecked(instance) orelse return false;
3839 return tile.target == .block_x and
3840 tile.extent == blocks and
3841 sparseScheduleAxisNameMatches(tile.axis, instance.row_axis, "_tile") and
3842 lane.target == .thread_x and
3843 lane.extent == instance.threads and
3844 sparseScheduleAxisNameMatches(lane.axis, instance.row_axis, "_lane");
3845 }
3846
3847 pub fn spmvSellLaunchExtent(instance: SpmvSell) u64 {
3848 return spmvSellLaunchExtentChecked(instance).?;
3849 }
3850
3851 fn spmvSellLaunchExtentChecked(instance: SpmvSell) ?u64 {
3852 return switch (instance.structure) {
3853 .row_thread => instance.rows,
3854 };
3855 }
3856
3857 pub fn spmvSellRepresentableThreads(instance: SpmvSell) ?u32 {
3858 const launch_extent = spmvSellLaunchExtentChecked(instance) orelse return null;
3859 return std.math.cast(u32, @min(@as(u64, instance.threads), launch_extent));
3860 }
3861
3862 pub fn spmvSellThreadCandidatesForRows(rows: u64) geometry_mod.Thread1DCandidates {
3863 return geometry_mod.threadCandidatesForExtent(rows, spmv_sell_thread_caps);
3864 }
3865
3866 pub fn spmmCsrThreadsForExtents(rows: u64, columns: u64) entry.Threads2D {
3867 return geometry_mod.threadsForGrid(.{ .rows = rows, .cols = columns }, spmm_csr_thread_caps);
3868 }
3869
3870 pub fn spmmCsrThreadCandidatesForExtents(rows: u64, columns: u64) geometry_mod.ThreadCandidates {
3871 return geometry_mod.threadCandidatesForGrid(.{ .rows = rows, .cols = columns }, spmm_csr_thread_caps);
3872 }
3873
3874 pub fn spmmCsrRepresentableThreads(instance: SpmmCsr) ?entry.Threads2D {
3875 if (instance.threads.x == 0 or instance.threads.y == 0) return null;
3876 const x = @min(@as(u64, instance.threads.x), instance.columns);
3877 const y = @min(@as(u64, instance.threads.y), instance.rows);
3878 if (x == 0 or y == 0) return null;
3879 var candidate = instance;
3880 candidate.threads = .{ .x = @intCast(x), .y = @intCast(y) };
3881 if (!spmmCsrInstanceValid(candidate)) return null;
3882 return candidate.threads;
3883 }
3884
3885 pub fn spmvSellFamilySpecialization(
3886 backing_allocator: std.mem.Allocator,
3887 instance: SpmvSell,
3888 ) !entry.OwnedSpecialization {
3889 var owned = entry.OwnedSpecialization.init(backing_allocator);
3890 errdefer owned.deinit();
3891 const lifetime_allocator = owned.allocator();
3892
3893 const inputs = try lifetime_allocator.alloc(entry.Shape, 4);
3894 inputs[0] = try entry.runtimeShape1D(
3895 lifetime_allocator,
3896 instance.slice_axis,
3897 spmvSellSliceCount(instance) + 1,
3898 );
3899 inputs[1] = try entry.runtimeShape1D(lifetime_allocator, instance.value_axis, instance.values_size);
3900 inputs[2] = try entry.runtimeShape1D(lifetime_allocator, instance.value_axis, instance.values_size);
3901 inputs[3] = try entry.runtimeShape1D(lifetime_allocator, instance.x_axis, instance.x_extent);
3902
3903 const outputs = try lifetime_allocator.alloc(entry.Shape, 1);
3904 outputs[0] = try entry.runtimeShape1D(lifetime_allocator, instance.row_axis, instance.rows);
3905
3906 const static_parameters = try lifetime_allocator.alloc(entry.StaticParameter, 1);
3907 static_parameters[0] = try entry.runtimeStaticParameter(
3908 lifetime_allocator,
3909 spmv_sell_slice_size_parameter,
3910 instance.slice_size,
3911 );
3912
3913 owned.value = .{
3914 .dtype = instance.dtype,
3915 .accumulation_dtype = instance.accumulation_dtype,
3916 .operation = .{ .sparse = .sell_spmv },
3917 .inputs = inputs,
3918 .outputs = outputs,
3919 .static_parameters = static_parameters,
3920 .schedule = try entry.runtimeThreadBlocks1D(
3921 lifetime_allocator,
3922 instance.row_axis,
3923 spmvSellLaunchExtent(instance),
3924 instance.threads,
3925 ),
3926 .structure = @tagName(instance.structure),
3927 };
3928 owned.value.launch = owned.value.schedule.?.launch();
3929 var family = try spmvSellShapeFamily(backing_allocator, instance);
3930 errdefer family.deinit();
3931 try owned.takeShapeFamily(&family);
3932 return owned;
3933 }
3934
3935 pub fn spmvSellInstanceFromSpecialization(specialization: entry.Specialization) ?SpmvSell {
3936 if (!specialization.scheduleMatchesLaunch()) return null;
3937 if (!specialization.staticParametersAreValid()) return null;
3938 const schedule = specialization.schedule orelse return null;
3939 if (!specialization.operationIs(.{ .sparse = .sell_spmv })) return null;
3940 const dtype = specialization.dtype orelse return null;
3941 const accumulation_dtype = spmvSellAccumulationDType(dtype) orelse return null;
3942 if (specialization.accumulation_dtype != accumulation_dtype) return null;
3943 if (specialization.inputs.len != 4 or specialization.outputs.len != 1) return null;
3944 if (specialization.reductions.len != 0) return null;
3945 if (specialization.static_parameters.len != 1) return null;
3946 const slice_offsets = specialization.inputs[0];
3947 const cols = specialization.inputs[1];
3948 const values = specialization.inputs[2];
3949 const x = specialization.inputs[3];
3950 const y = specialization.outputs[0];
3951 if (slice_offsets.axes.len != 1 or cols.axes.len != 1 or values.axes.len != 1 or x.axes.len != 1 or y.axes.len != 1) return null;
3952 if (slice_offsets.axes[0].name.len == 0 or cols.axes[0].name.len == 0 or x.axes[0].name.len == 0 or y.axes[0].name.len == 0) return null;
3953 if (slice_offsets.axes[0].extent <= 1 or cols.axes[0].extent == 0 or x.axes[0].extent == 0) return null;
3954 const slice_size = specialization.staticParameterValue(spmv_sell_slice_size_parameter) orelse return null;
3955 if (slice_size == 0 or slice_size > extent_mod.runtime_extent_max) return null;
3956 if (!sparseAxisMatches(values.axes[0], cols.axes[0].name, cols.axes[0].extent)) return null;
3957 const launch = specialization.launch orelse return null;
3958 if (launch.threadgroup[0] == 0) return null;
3959 const structure_name = specialization.structure orelse return null;
3960 const structure = std.meta.stringToEnum(SpmvSellStructure, structure_name) orelse return null;
3961 const instance = SpmvSell{
3962 .rows = y.axes[0].extent,
3963 .slice_size = slice_size,
3964 .values_size = cols.axes[0].extent,
3965 .x_extent = x.axes[0].extent,
3966 .dtype = dtype,
3967 .accumulation_dtype = accumulation_dtype,
3968 .threads = launch.threadgroup[0],
3969 .structure = structure,
3970 .row_axis = y.axes[0].name,
3971 .slice_axis = slice_offsets.axes[0].name,
3972 .value_axis = cols.axes[0].name,
3973 .x_axis = x.axes[0].name,
3974 };
3975 if (!spmvSellInstanceValid(instance)) return null;
3976 const slice_offsets_extent = std.math.add(u64, spmvSellSliceCount(instance), 1) catch return null;
3977 if (!sparseAxisMatches(slice_offsets.axes[0], instance.slice_axis, slice_offsets_extent)) return null;
3978 if (!spmvSellScheduleMatchesInstance(schedule, instance)) return null;
3979 return instance;
3980 }
3981
3982 fn spmvSellScheduleMatchesInstance(schedule: entry.Schedule, instance: SpmvSell) bool {
3983 const extent = spmvSellLaunchExtentChecked(instance) orelse return false;
3984 if (extent <= instance.threads) {
3985 if (schedule.bindings.len != 1) return false;
3986 const binding = schedule.bindings[0];
3987 return sparseAxisMatches(.{ .name = binding.axis, .extent = binding.extent }, instance.row_axis, extent) and
3988 binding.target == .thread_x;
3989 }
3990 if (schedule.bindings.len != 2) return false;
3991 const tile = schedule.bindings[0];
3992 const lane = schedule.bindings[1];
3993 const blocks = spmvSellBlockCountChecked(instance) orelse return false;
3994 return tile.target == .block_x and
3995 tile.extent == blocks and
3996 sparseScheduleAxisNameMatches(tile.axis, instance.row_axis, "_tile") and
3997 lane.target == .thread_x and
3998 lane.extent == instance.threads and
3999 sparseScheduleAxisNameMatches(lane.axis, instance.row_axis, "_lane");
4000 }
4001
4002 pub fn spmmCsrFamilySpecialization(
4003 backing_allocator: std.mem.Allocator,
4004 instance: SpmmCsr,
4005 ) !entry.OwnedSpecialization {
4006 var owned = entry.OwnedSpecialization.init(backing_allocator);
4007 errdefer owned.deinit();
4008 const lifetime_allocator = owned.allocator();
4009
4010 const inputs = try lifetime_allocator.alloc(entry.Shape, 4);
4011 inputs[0] = try entry.runtimeShape1D(lifetime_allocator, instance.row_axis, instance.rows + 1);
4012 inputs[1] = try entry.runtimeShape1D(lifetime_allocator, spmv_csr_nonzero_axis, instance.nnz);
4013 inputs[2] = try entry.runtimeShape1D(lifetime_allocator, spmv_csr_nonzero_axis, instance.nnz);
4014 inputs[3] = try entry.runtimeShape2D(
4015 lifetime_allocator,
4016 instance.x_axis,
4017 instance.x_extent,
4018 instance.column_axis,
4019 instance.columns,
4020 );
4021
4022 const outputs = try lifetime_allocator.alloc(entry.Shape, 1);
4023 outputs[0] = try entry.runtimeShape2D(lifetime_allocator, instance.row_axis, instance.rows, instance.column_axis, instance.columns);
4024
4025 owned.value = .{
4026 .dtype = instance.dtype,
4027 .accumulation_dtype = instance.accumulation_dtype,
4028 .operation = .{ .sparse = .csr_spmm },
4029 .inputs = inputs,
4030 .outputs = outputs,
4031 .schedule = try entry.runtimeThreadBlocks2D(
4032 lifetime_allocator,
4033 instance.column_axis,
4034 instance.columns,
4035 instance.row_axis,
4036 instance.rows,
4037 instance.threads.x,
4038 instance.threads.y,
4039 ),
4040 .structure = @tagName(instance.structure),
4041 };
4042 owned.value.launch = owned.value.schedule.?.launch();
4043 var family = try spmmCsrShapeFamily(backing_allocator, instance);
4044 errdefer family.deinit();
4045 try owned.takeShapeFamily(&family);
4046 return owned;
4047 }
4048
4049 pub fn spmmCsrInstanceFromSpecialization(specialization: entry.Specialization) ?SpmmCsr {
4050 if (!specialization.scheduleMatchesLaunch()) return null;
4051 const schedule = specialization.schedule orelse return null;
4052 if (!specialization.operationIs(.{ .sparse = .csr_spmm })) return null;
4053 const dtype = specialization.dtype orelse return null;
4054 const accumulation_dtype = spmmCsrAccumulationDType(dtype) orelse return null;
4055 if (specialization.accumulation_dtype != accumulation_dtype) return null;
4056 if (specialization.inputs.len != 4 or specialization.outputs.len != 1) return null;
4057 if (specialization.reductions.len != 0) return null;
4058 const row_ptr = specialization.inputs[0];
4059 const cols = specialization.inputs[1];
4060 const values = specialization.inputs[2];
4061 const x = specialization.inputs[3];
4062 const y = specialization.outputs[0];
4063 if (row_ptr.axes.len != 1 or cols.axes.len != 1 or values.axes.len != 1 or x.axes.len != 2 or y.axes.len != 2) return null;
4064 if (y.axes[0].name.len == 0 or y.axes[1].name.len == 0) return null;
4065 if (cols.axes[0].extent == 0 or x.axes[0].extent == 0 or x.axes[1].extent == 0) return null;
4066 const row_ptr_extent = std.math.add(u64, y.axes[0].extent, 1) catch return null;
4067 if (!sparseAxisMatches(row_ptr.axes[0], y.axes[0].name, row_ptr_extent)) return null;
4068 if (!sparseAxisMatches(cols.axes[0], spmv_csr_nonzero_axis, cols.axes[0].extent)) return null;
4069 if (!sparseAxisMatches(values.axes[0], spmv_csr_nonzero_axis, cols.axes[0].extent)) return null;
4070 if (!sparseAxisMatches(x.axes[1], y.axes[1].name, y.axes[1].extent)) return null;
4071 if (x.axes[0].name.len == 0) return null;
4072 const launch = specialization.launch orelse return null;
4073 if (launch.threadgroup[0] == 0 or launch.threadgroup[1] == 0) return null;
4074 const structure_name = specialization.structure orelse return null;
4075 const structure = std.meta.stringToEnum(SpmmCsrStructure, structure_name) orelse return null;
4076 const instance = SpmmCsr{
4077 .rows = y.axes[0].extent,
4078 .columns = y.axes[1].extent,
4079 .nnz = cols.axes[0].extent,
4080 .x_extent = x.axes[0].extent,
4081 .dtype = dtype,
4082 .accumulation_dtype = accumulation_dtype,
4083 .threads = .{ .x = launch.threadgroup[0], .y = launch.threadgroup[1] },
4084 .structure = structure,
4085 .row_axis = y.axes[0].name,
4086 .column_axis = y.axes[1].name,
4087 .x_axis = x.axes[0].name,
4088 };
4089 if (!spmmCsrInstanceValid(instance)) return null;
4090 if (!spmmCsrScheduleMatchesInstance(schedule, instance)) return null;
4091 return instance;
4092 }
4093
4094 fn sparseScheduleBindingMatches(binding: entry.ScheduleBinding, axis: []const u8, extent: u64, target: kernel.BindTarget) bool {
4095 return binding.target == target and sparseAxisMatches(.{ .name = binding.axis, .extent = binding.extent }, axis, extent);
4096 }
4097
4098 fn sparseScheduleTiledAxisMatches(
4099 tile: entry.ScheduleBinding,
4100 lane: entry.ScheduleBinding,
4101 axis: []const u8,
4102 blocks: u64,
4103 threads: u32,
4104 block_target: kernel.BindTarget,
4105 thread_target: kernel.BindTarget,
4106 ) bool {
4107 return tile.target == block_target and
4108 tile.extent == blocks and
4109 sparseScheduleAxisNameMatches(tile.axis, axis, "_tile") and
4110 lane.target == thread_target and
4111 lane.extent == threads and
4112 sparseScheduleAxisNameMatches(lane.axis, axis, "_lane");
4113 }
4114
4115 fn spmmCsrScheduleMatchesInstance(schedule: entry.Schedule, instance: SpmmCsr) bool {
4116 const x_tiled = instance.columns > instance.threads.x;
4117 const y_tiled = instance.rows > instance.threads.y;
4118 const expected_count: usize = @as(usize, if (x_tiled) 2 else 1) + @as(usize, if (y_tiled) 2 else 1);
4119 if (schedule.bindings.len != expected_count) return false;
4120 var index: usize = 0;
4121 if (x_tiled) {
4122 const blocks = spmmCsrBlockCountXChecked(instance) orelse return false;
4123 if (!sparseScheduleTiledAxisMatches(
4124 schedule.bindings[index],
4125 schedule.bindings[index + 1],
4126 instance.column_axis,
4127 blocks,
4128 instance.threads.x,
4129 .block_x,
4130 .thread_x,
4131 )) return false;
4132 index += 2;
4133 } else {
4134 if (!sparseScheduleBindingMatches(schedule.bindings[index], instance.column_axis, instance.columns, .thread_x)) return false;
4135 index += 1;
4136 }
4137 if (y_tiled) {
4138 const blocks = spmmCsrBlockCountYChecked(instance) orelse return false;
4139 if (!sparseScheduleTiledAxisMatches(
4140 schedule.bindings[index],
4141 schedule.bindings[index + 1],
4142 instance.row_axis,
4143 blocks,
4144 instance.threads.y,
4145 .block_y,
4146 .thread_y,
4147 )) return false;
4148 } else {
4149 if (!sparseScheduleBindingMatches(schedule.bindings[index], instance.row_axis, instance.rows, .thread_y)) return false;
4150 }
4151 return true;
4152 }
4153
4154 test "sparse spmv csr specialization round-trips both structures" {
4155 const allocator = testing.allocator;
4156
4157 inline for ([_]SpmvCsrStructure{ .row_thread, .row_warp }) |structure| {
4158 var owned = try spmvCsrFamilySpecialization(
4159 allocator,
4160 .{ .rows = 70, .nnz = 512, .x_extent = 40, .threads = 64, .structure = structure },
4161 );
4162 defer owned.deinit();
4163 try testing.expect(owned.value.structureIs(@tagName(structure)));
4164 try testing.expect(owned.value.operationIs(.{ .sparse = .csr_spmv }));
4165 const recovered = spmvCsrInstanceFromSpecialization(owned.value) orelse {
4166 return error.TestExpectedSpmvCsrInstance;
4167 };
4168 try testing.expectEqual(@as(u64, 70), recovered.rows);
4169 try testing.expectEqual(@as(u64, 512), recovered.nnz);
4170 try testing.expectEqual(@as(u64, 40), recovered.x_extent);
4171 try testing.expectEqual(@as(u32, 64), recovered.threads);
4172 try testing.expectEqual(structure, recovered.structure);
4173 }
4174
4175 var f16_owned = try spmvCsrFamilySpecialization(
4176 allocator,
4177 .{ .rows = 70, .nnz = 512, .x_extent = 40, .dtype = .f16, .accumulation_dtype = .f32, .threads = 64 },
4178 );
4179 defer f16_owned.deinit();
4180 try testing.expectEqual(@as(?DType, .f16), f16_owned.value.dtype);
4181 try testing.expectEqual(@as(?DType, .f32), f16_owned.value.accumulation_dtype);
4182 const f16_recovered = spmvCsrInstanceFromSpecialization(f16_owned.value) orelse {
4183 return error.TestExpectedSpmvCsrInstance;
4184 };
4185 try testing.expectEqual(DType.f16, f16_recovered.dtype);
4186 try testing.expectEqual(DType.f32, f16_recovered.accumulation_dtype);
4187
4188 var f64_owned = try spmvCsrFamilySpecialization(
4189 allocator,
4190 .{ .rows = 70, .nnz = 512, .x_extent = 40, .dtype = .f64, .accumulation_dtype = .f64, .threads = 64 },
4191 );
4192 defer f64_owned.deinit();
4193 try testing.expectEqual(@as(?DType, .f64), f64_owned.value.dtype);
4194 try testing.expectEqual(@as(?DType, .f64), f64_owned.value.accumulation_dtype);
4195 const f64_recovered = spmvCsrInstanceFromSpecialization(f64_owned.value) orelse {
4196 return error.TestExpectedSpmvCsrInstance;
4197 };
4198 try testing.expectEqual(DType.f64, f64_recovered.dtype);
4199 try testing.expectEqual(DType.f64, f64_recovered.accumulation_dtype);
4200 try testing.expectEqual(@as(?SpmvCsr, null), spmvCsrInstanceFromSpecialization(.{}));
4201 }
4202
4203 test "sparse spmv coo specialization round-trips element-thread metadata" {
4204 const allocator = testing.allocator;
4205
4206 var owned = try spmvCooFamilySpecialization(
4207 allocator,
4208 .{ .rows = 70, .nnz = 512, .x_extent = 40, .threads = 64 },
4209 );
4210 defer owned.deinit();
4211 try testing.expect(owned.value.structureIs("element_thread"));
4212 try testing.expect(owned.value.operationIs(.{ .sparse = .coo_spmv }));
4213 const recovered = spmvCooInstanceFromSpecialization(owned.value) orelse {
4214 return error.TestExpectedSpmvCooInstance;
4215 };
4216 try testing.expectEqual(@as(u64, 70), recovered.rows);
4217 try testing.expectEqual(@as(u64, 512), recovered.nnz);
4218 try testing.expectEqual(@as(u64, 40), recovered.x_extent);
4219 try testing.expectEqual(@as(u32, 64), recovered.threads);
4220 try testing.expectEqual(SpmvCooStructure.element_thread, recovered.structure);
4221 try testing.expectEqualStrings("n", recovered.nonzero_axis);
4222 try testing.expectEqualStrings("x", recovered.x_axis);
4223
4224 var f16_owned = try spmvCooFamilySpecialization(
4225 allocator,
4226 .{ .rows = 70, .nnz = 512, .x_extent = 40, .dtype = .f16, .accumulation_dtype = .f32, .threads = 64, .structure = .row_thread },
4227 );
4228 defer f16_owned.deinit();
4229 try testing.expect(f16_owned.value.structureIs("row_thread"));
4230 try testing.expectEqual(@as(?DType, .f16), f16_owned.value.dtype);
4231 try testing.expectEqual(@as(?DType, .f32), f16_owned.value.accumulation_dtype);
4232 const f16_recovered = spmvCooInstanceFromSpecialization(f16_owned.value) orelse {
4233 return error.TestExpectedSpmvCooInstance;
4234 };
4235 try testing.expectEqual(DType.f16, f16_recovered.dtype);
4236 try testing.expectEqual(DType.f32, f16_recovered.accumulation_dtype);
4237 try testing.expectEqual(SpmvCooStructure.row_thread, f16_recovered.structure);
4238
4239 var f64_owned = try spmvCooFamilySpecialization(
4240 allocator,
4241 .{ .rows = 70, .nnz = 512, .x_extent = 40, .dtype = .f64, .accumulation_dtype = .f64, .threads = 64, .structure = .row_thread },
4242 );
4243 defer f64_owned.deinit();
4244 try testing.expect(f64_owned.value.structureIs("row_thread"));
4245 try testing.expectEqual(@as(?DType, .f64), f64_owned.value.dtype);
4246 try testing.expectEqual(@as(?DType, .f64), f64_owned.value.accumulation_dtype);
4247 const f64_recovered = spmvCooInstanceFromSpecialization(f64_owned.value) orelse {
4248 return error.TestExpectedSpmvCooInstance;
4249 };
4250 try testing.expectEqual(DType.f64, f64_recovered.dtype);
4251 try testing.expectEqual(DType.f64, f64_recovered.accumulation_dtype);
4252 try testing.expectEqual(SpmvCooStructure.row_thread, f64_recovered.structure);
4253
4254 try testing.expectEqual(@as(?SpmvCoo, null), spmvCooInstanceFromSpecialization(.{}));
4255 }
4256
4257 test "sparse spmv ell specialization round-trips padded matrix metadata" {
4258 const allocator = testing.allocator;
4259
4260 var owned = try spmvEllFamilySpecialization(
4261 allocator,
4262 .{ .rows = 70, .slots = 8, .x_extent = 40, .threads = 64 },
4263 );
4264 defer owned.deinit();
4265 try testing.expect(owned.value.structureIs("row_thread"));
4266 try testing.expect(owned.value.operationIs(.{ .sparse = .ell_spmv }));
4267 const recovered = spmvEllInstanceFromSpecialization(owned.value) orelse {
4268 return error.TestExpectedSpmvEllInstance;
4269 };
4270 try testing.expectEqual(@as(u64, 70), recovered.rows);
4271 try testing.expectEqual(@as(u64, 8), recovered.slots);
4272 try testing.expectEqual(@as(u64, 40), recovered.x_extent);
4273 try testing.expectEqual(@as(u32, 64), recovered.threads);
4274 try testing.expectEqual(SpmvEllStructure.row_thread, recovered.structure);
4275 try testing.expectEqualStrings("s", recovered.slot_axis);
4276 try testing.expectEqualStrings("x", recovered.x_axis);
4277
4278 var f16_owned = try spmvEllFamilySpecialization(
4279 allocator,
4280 .{ .rows = 70, .slots = 8, .x_extent = 40, .dtype = .f16, .accumulation_dtype = .f32, .threads = 64 },
4281 );
4282 defer f16_owned.deinit();
4283 try testing.expectEqual(@as(?DType, .f16), f16_owned.value.dtype);
4284 try testing.expectEqual(@as(?DType, .f32), f16_owned.value.accumulation_dtype);
4285 const f16_recovered = spmvEllInstanceFromSpecialization(f16_owned.value) orelse {
4286 return error.TestExpectedSpmvEllInstance;
4287 };
4288 try testing.expectEqual(DType.f16, f16_recovered.dtype);
4289 try testing.expectEqual(DType.f32, f16_recovered.accumulation_dtype);
4290
4291 var f64_owned = try spmvEllFamilySpecialization(
4292 allocator,
4293 .{ .rows = 70, .slots = 8, .x_extent = 40, .dtype = .f64, .accumulation_dtype = .f64, .threads = 64 },
4294 );
4295 defer f64_owned.deinit();
4296 try testing.expectEqual(@as(?DType, .f64), f64_owned.value.dtype);
4297 try testing.expectEqual(@as(?DType, .f64), f64_owned.value.accumulation_dtype);
4298 const f64_recovered = spmvEllInstanceFromSpecialization(f64_owned.value) orelse {
4299 return error.TestExpectedSpmvEllInstance;
4300 };
4301 try testing.expectEqual(DType.f64, f64_recovered.dtype);
4302 try testing.expectEqual(DType.f64, f64_recovered.accumulation_dtype);
4303 try testing.expectEqual(@as(?SpmvEll, null), spmvEllInstanceFromSpecialization(.{}));
4304 }
4305
4306 test "sparse spmv sell specialization round-trips sliced matrix metadata" {
4307 const allocator = testing.allocator;
4308
4309 var owned = try spmvSellFamilySpecialization(
4310 allocator,
4311 .{ .rows = 70, .slice_size = 8, .values_size = 400, .x_extent = 40, .threads = 64 },
4312 );
4313 defer owned.deinit();
4314 try testing.expect(owned.value.structureIs("row_thread"));
4315 try testing.expect(owned.value.operationIs(.{ .sparse = .sell_spmv }));
4316 try testing.expect(owned.value.staticParameterMatches(spmv_sell_slice_size_parameter, 8));
4317 const recovered = spmvSellInstanceFromSpecialization(owned.value) orelse {
4318 return error.TestExpectedSpmvSellInstance;
4319 };
4320 try testing.expectEqual(@as(u64, 70), recovered.rows);
4321 try testing.expectEqual(@as(u64, 8), recovered.slice_size);
4322 try testing.expectEqual(@as(u64, 400), recovered.values_size);
4323 try testing.expectEqual(@as(u64, 40), recovered.x_extent);
4324 try testing.expectEqual(@as(u32, 64), recovered.threads);
4325 try testing.expectEqual(SpmvSellStructure.row_thread, recovered.structure);
4326 try testing.expectEqualStrings("z", recovered.slice_axis);
4327 try testing.expectEqualStrings("n", recovered.value_axis);
4328 try testing.expectEqualStrings("x", recovered.x_axis);
4329
4330 var f16_owned = try spmvSellFamilySpecialization(
4331 allocator,
4332 .{ .rows = 70, .slice_size = 8, .values_size = 400, .x_extent = 40, .dtype = .f16, .accumulation_dtype = .f32, .threads = 64 },
4333 );
4334 defer f16_owned.deinit();
4335 try testing.expectEqual(@as(?DType, .f16), f16_owned.value.dtype);
4336 try testing.expectEqual(@as(?DType, .f32), f16_owned.value.accumulation_dtype);
4337 const f16_recovered = spmvSellInstanceFromSpecialization(f16_owned.value) orelse {
4338 return error.TestExpectedSpmvSellInstance;
4339 };
4340 try testing.expectEqual(DType.f16, f16_recovered.dtype);
4341 try testing.expectEqual(DType.f32, f16_recovered.accumulation_dtype);
4342 try testing.expectEqual(@as(u64, 8), f16_recovered.slice_size);
4343
4344 var f64_owned = try spmvSellFamilySpecialization(
4345 allocator,
4346 .{ .rows = 70, .slice_size = 8, .values_size = 400, .x_extent = 40, .dtype = .f64, .accumulation_dtype = .f64, .threads = 64 },
4347 );
4348 defer f64_owned.deinit();
4349 try testing.expectEqual(@as(?DType, .f64), f64_owned.value.dtype);
4350 try testing.expectEqual(@as(?DType, .f64), f64_owned.value.accumulation_dtype);
4351 const f64_recovered = spmvSellInstanceFromSpecialization(f64_owned.value) orelse {
4352 return error.TestExpectedSpmvSellInstance;
4353 };
4354 try testing.expectEqual(DType.f64, f64_recovered.dtype);
4355 try testing.expectEqual(DType.f64, f64_recovered.accumulation_dtype);
4356 try testing.expectEqual(@as(u64, 8), f64_recovered.slice_size);
4357 try testing.expectEqual(@as(?SpmvSell, null), spmvSellInstanceFromSpecialization(.{}));
4358 }
4359
4360 test "sparse spmm csr specialization round-trips matrix metadata" {
4361 const allocator = testing.allocator;
4362
4363 var owned = try spmmCsrFamilySpecialization(
4364 allocator,
4365 .{ .rows = 70, .columns = 45, .nnz = 512, .x_extent = 40, .threads = .{ .x = 8, .y = 4 } },
4366 );
4367 defer owned.deinit();
4368 try testing.expect(owned.value.structureIs("row_column_thread"));
4369 try testing.expect(owned.value.operationIs(.{ .sparse = .csr_spmm }));
4370 const recovered = spmmCsrInstanceFromSpecialization(owned.value) orelse {
4371 return error.TestExpectedSpmmCsrInstance;
4372 };
4373 try testing.expectEqual(@as(u64, 70), recovered.rows);
4374 try testing.expectEqual(@as(u64, 45), recovered.columns);
4375 try testing.expectEqual(@as(u64, 512), recovered.nnz);
4376 try testing.expectEqual(@as(u64, 40), recovered.x_extent);
4377 try testing.expectEqual(@as(u32, 8), recovered.threads.x);
4378 try testing.expectEqual(@as(u32, 4), recovered.threads.y);
4379 try testing.expectEqual(SpmmCsrStructure.row_column_thread, recovered.structure);
4380
4381 var f16_owned = try spmmCsrFamilySpecialization(
4382 allocator,
4383 .{ .rows = 70, .columns = 45, .nnz = 512, .x_extent = 40, .dtype = .f16, .accumulation_dtype = .f32, .threads = .{ .x = 8, .y = 4 } },
4384 );
4385 defer f16_owned.deinit();
4386 try testing.expectEqual(@as(?DType, .f16), f16_owned.value.dtype);
4387 try testing.expectEqual(@as(?DType, .f32), f16_owned.value.accumulation_dtype);
4388 const f16_recovered = spmmCsrInstanceFromSpecialization(f16_owned.value) orelse {
4389 return error.TestExpectedSpmmCsrInstance;
4390 };
4391 try testing.expectEqual(DType.f16, f16_recovered.dtype);
4392 try testing.expectEqual(DType.f32, f16_recovered.accumulation_dtype);
4393
4394 var f64_owned = try spmmCsrFamilySpecialization(
4395 allocator,
4396 .{ .rows = 70, .columns = 45, .nnz = 512, .x_extent = 40, .dtype = .f64, .accumulation_dtype = .f64, .threads = .{ .x = 8, .y = 4 } },
4397 );
4398 defer f64_owned.deinit();
4399 try testing.expectEqual(@as(?DType, .f64), f64_owned.value.dtype);
4400 try testing.expectEqual(@as(?DType, .f64), f64_owned.value.accumulation_dtype);
4401 const f64_recovered = spmmCsrInstanceFromSpecialization(f64_owned.value) orelse {
4402 return error.TestExpectedSpmmCsrInstance;
4403 };
4404 try testing.expectEqual(DType.f64, f64_recovered.dtype);
4405 try testing.expectEqual(DType.f64, f64_recovered.accumulation_dtype);
4406 try testing.expectEqual(@as(?SpmmCsr, null), spmmCsrInstanceFromSpecialization(.{}));
4407 }
4408
4409 test "sparse spmv csr specialization rejects malformed descriptor facts" {
4410 const allocator = testing.allocator;
4411
4412 {
4413 var owned = try spmvCsrFamilySpecialization(
4414 allocator,
4415 .{ .rows = 70, .nnz = 512, .x_extent = 40, .threads = 64, .structure = .row_thread },
4416 );
4417 defer owned.deinit();
4418 const lifetime_allocator = owned.allocator();
4419 const inputs = try lifetime_allocator.alloc(entry.Shape, 4);
4420 inputs[0] = owned.value.inputs[0];
4421 inputs[1] = owned.value.inputs[1];
4422 inputs[2] = owned.value.inputs[2];
4423 inputs[3] = try entry.runtimeShape1D(lifetime_allocator, "bad_x", 40);
4424 owned.value.inputs = inputs;
4425 try testing.expectEqual(@as(?SpmvCsr, null), spmvCsrInstanceFromSpecialization(owned.value));
4426 }
4427
4428 {
4429 var owned = try spmvCsrFamilySpecialization(
4430 allocator,
4431 .{ .rows = 70, .nnz = 512, .x_extent = 40, .threads = 64, .structure = .row_thread },
4432 );
4433 defer owned.deinit();
4434 const lifetime_allocator = owned.allocator();
4435 const inputs = try lifetime_allocator.alloc(entry.Shape, 4);
4436 const zero_x_axes = try lifetime_allocator.alloc(entry.Axis, 1);
4437 zero_x_axes[0] = .{ .name = spmv_csr_column_axis, .extent = 0 };
4438 inputs[0] = owned.value.inputs[0];
4439 inputs[1] = owned.value.inputs[1];
4440 inputs[2] = owned.value.inputs[2];
4441 inputs[3] = .{ .axes = zero_x_axes };
4442 owned.value.inputs = inputs;
4443 try testing.expectEqual(@as(?SpmvCsr, null), spmvCsrInstanceFromSpecialization(owned.value));
4444 }
4445
4446 {
4447 var owned = try spmvCsrFamilySpecialization(
4448 allocator,
4449 .{ .rows = 70, .nnz = 512, .x_extent = 40, .threads = 64, .structure = .row_thread },
4450 );
4451 defer owned.deinit();
4452 const lifetime_allocator = owned.allocator();
4453 const outputs = try lifetime_allocator.alloc(entry.Shape, 1);
4454 const y_axes = try lifetime_allocator.alloc(entry.Axis, 1);
4455 y_axes[0] = .{ .name = "r", .extent = std.math.maxInt(u64) };
4456 outputs[0] = .{ .axes = y_axes };
4457 owned.value.outputs = outputs;
4458 try testing.expectEqual(@as(?SpmvCsr, null), spmvCsrInstanceFromSpecialization(owned.value));
4459 }
4460
4461 {
4462 var owned = try spmvCsrFamilySpecialization(
4463 allocator,
4464 .{ .rows = 70, .nnz = 512, .x_extent = 40, .threads = 64, .structure = .row_warp },
4465 );
4466 defer owned.deinit();
4467 const lifetime_allocator = owned.allocator();
4468 const rows = std.math.maxInt(u64) / @as(u64, spmv_csr_warp_size) + 1;
4469 const inputs = try lifetime_allocator.alloc(entry.Shape, 4);
4470 const row_ptr_axes = try lifetime_allocator.alloc(entry.Axis, 1);
4471 row_ptr_axes[0] = .{ .name = "r", .extent = rows + 1 };
4472 inputs[0] = .{ .axes = row_ptr_axes };
4473 inputs[1] = owned.value.inputs[1];
4474 inputs[2] = owned.value.inputs[2];
4475 inputs[3] = owned.value.inputs[3];
4476 owned.value.inputs = inputs;
4477 const outputs = try lifetime_allocator.alloc(entry.Shape, 1);
4478 const y_axes = try lifetime_allocator.alloc(entry.Axis, 1);
4479 y_axes[0] = .{ .name = "r", .extent = rows };
4480 outputs[0] = .{ .axes = y_axes };
4481 owned.value.outputs = outputs;
4482 try testing.expectEqual(@as(?SpmvCsr, null), spmvCsrInstanceFromSpecialization(owned.value));
4483 }
4484
4485 {
4486 var owned = try spmvCsrFamilySpecialization(
4487 allocator,
4488 .{ .rows = 70, .nnz = 512, .x_extent = 40, .threads = 64, .structure = .row_warp },
4489 );
4490 defer owned.deinit();
4491 const lifetime_allocator = owned.allocator();
4492 owned.value.schedule = try entry.runtimeThreadBlocks1D(lifetime_allocator, "r", 70, 64);
4493 owned.value.launch = owned.value.schedule.?.launch();
4494 try testing.expectEqual(@as(?SpmvCsr, null), spmvCsrInstanceFromSpecialization(owned.value));
4495 }
4496 }
4497
4498 test "sparse spmv coo specialization rejects malformed descriptor facts" {
4499 const allocator = testing.allocator;
4500
4501 {
4502 var owned = try spmvCooFamilySpecialization(
4503 allocator,
4504 .{ .rows = 70, .nnz = 512, .x_extent = 40, .threads = 64 },
4505 );
4506 defer owned.deinit();
4507 const lifetime_allocator = owned.allocator();
4508 const inputs = try lifetime_allocator.alloc(entry.Shape, 4);
4509 inputs[0] = owned.value.inputs[0];
4510 inputs[1] = try entry.runtimeShape1D(lifetime_allocator, "bad_n", 512);
4511 inputs[2] = owned.value.inputs[2];
4512 inputs[3] = owned.value.inputs[3];
4513 owned.value.inputs = inputs;
4514 try testing.expectEqual(@as(?SpmvCoo, null), spmvCooInstanceFromSpecialization(owned.value));
4515 }
4516
4517 {
4518 var owned = try spmvCooFamilySpecialization(
4519 allocator,
4520 .{ .rows = 70, .nnz = 512, .x_extent = 40, .threads = 64 },
4521 );
4522 defer owned.deinit();
4523 const lifetime_allocator = owned.allocator();
4524 const inputs = try lifetime_allocator.alloc(entry.Shape, 4);
4525 const zero_nnz_axes = try lifetime_allocator.alloc(entry.Axis, 1);
4526 zero_nnz_axes[0] = .{ .name = "n", .extent = 0 };
4527 inputs[0] = .{ .axes = zero_nnz_axes };
4528 inputs[1] = owned.value.inputs[1];
4529 inputs[2] = owned.value.inputs[2];
4530 inputs[3] = owned.value.inputs[3];
4531 owned.value.inputs = inputs;
4532 try testing.expectEqual(@as(?SpmvCoo, null), spmvCooInstanceFromSpecialization(owned.value));
4533 }
4534
4535 {
4536 var owned = try spmvCooFamilySpecialization(
4537 allocator,
4538 .{ .rows = 70, .nnz = 512, .x_extent = 40, .threads = 64 },
4539 );
4540 defer owned.deinit();
4541 const lifetime_allocator = owned.allocator();
4542 owned.value.schedule = try entry.runtimeThreadBlocks1D(lifetime_allocator, "r", 70, 64);
4543 owned.value.launch = owned.value.schedule.?.launch();
4544 try testing.expectEqual(@as(?SpmvCoo, null), spmvCooInstanceFromSpecialization(owned.value));
4545 }
4546
4547 {
4548 var owned = try spmvCooFamilySpecialization(
4549 allocator,
4550 .{ .rows = 70, .nnz = 512, .x_extent = 40, .threads = 64 },
4551 );
4552 defer owned.deinit();
4553 owned.value.structure = "row_thread";
4554 try testing.expectEqual(@as(?SpmvCoo, null), spmvCooInstanceFromSpecialization(owned.value));
4555 }
4556 }
4557
4558 test "sparse spmv ell specialization rejects malformed descriptor facts" {
4559 const allocator = testing.allocator;
4560
4561 {
4562 var owned = try spmvEllFamilySpecialization(
4563 allocator,
4564 .{ .rows = 70, .slots = 8, .x_extent = 40, .threads = 64 },
4565 );
4566 defer owned.deinit();
4567 const lifetime_allocator = owned.allocator();
4568 const inputs = try lifetime_allocator.alloc(entry.Shape, 3);
4569 inputs[0] = owned.value.inputs[0];
4570 inputs[1] = try entry.runtimeShape2D(lifetime_allocator, "s", 8, "bad_r", 70);
4571 inputs[2] = owned.value.inputs[2];
4572 owned.value.inputs = inputs;
4573 try testing.expectEqual(@as(?SpmvEll, null), spmvEllInstanceFromSpecialization(owned.value));
4574 }
4575
4576 {
4577 var owned = try spmvEllFamilySpecialization(
4578 allocator,
4579 .{ .rows = 70, .slots = 8, .x_extent = 40, .threads = 64 },
4580 );
4581 defer owned.deinit();
4582 const lifetime_allocator = owned.allocator();
4583 const inputs = try lifetime_allocator.alloc(entry.Shape, 3);
4584 const zero_slot_axes = try lifetime_allocator.alloc(entry.Axis, 2);
4585 zero_slot_axes[0] = .{ .name = "s", .extent = 0 };
4586 zero_slot_axes[1] = .{ .name = "r", .extent = 70 };
4587 inputs[0] = .{ .axes = zero_slot_axes };
4588 inputs[1] = owned.value.inputs[1];
4589 inputs[2] = owned.value.inputs[2];
4590 owned.value.inputs = inputs;
4591 try testing.expectEqual(@as(?SpmvEll, null), spmvEllInstanceFromSpecialization(owned.value));
4592 }
4593
4594 {
4595 var owned = try spmvEllFamilySpecialization(
4596 allocator,
4597 .{ .rows = 70, .slots = 8, .x_extent = 40, .threads = 64 },
4598 );
4599 defer owned.deinit();
4600 const lifetime_allocator = owned.allocator();
4601 const inputs = try lifetime_allocator.alloc(entry.Shape, 3);
4602 const zero_x_axes = try lifetime_allocator.alloc(entry.Axis, 1);
4603 zero_x_axes[0] = .{ .name = "x", .extent = 0 };
4604 inputs[0] = owned.value.inputs[0];
4605 inputs[1] = owned.value.inputs[1];
4606 inputs[2] = .{ .axes = zero_x_axes };
4607 owned.value.inputs = inputs;
4608 try testing.expectEqual(@as(?SpmvEll, null), spmvEllInstanceFromSpecialization(owned.value));
4609 }
4610
4611 {
4612 var owned = try spmvEllFamilySpecialization(
4613 allocator,
4614 .{ .rows = 70, .slots = 8, .x_extent = 40, .threads = 64 },
4615 );
4616 defer owned.deinit();
4617 const lifetime_allocator = owned.allocator();
4618 owned.value.schedule = try entry.runtimeThreadBlocks1D(lifetime_allocator, "bad_r", 70, 64);
4619 owned.value.launch = owned.value.schedule.?.launch();
4620 try testing.expectEqual(@as(?SpmvEll, null), spmvEllInstanceFromSpecialization(owned.value));
4621 }
4622
4623 {
4624 var owned = try spmvEllFamilySpecialization(
4625 allocator,
4626 .{ .rows = 70, .slots = 8, .x_extent = 40, .threads = 64 },
4627 );
4628 defer owned.deinit();
4629 owned.value.structure = "row_warp";
4630 try testing.expectEqual(@as(?SpmvEll, null), spmvEllInstanceFromSpecialization(owned.value));
4631 }
4632 }
4633
4634 test "sparse spmv sell specialization rejects malformed descriptor facts" {
4635 const allocator = testing.allocator;
4636
4637 {
4638 var owned = try spmvSellFamilySpecialization(
4639 allocator,
4640 .{ .rows = 70, .slice_size = 8, .values_size = 400, .x_extent = 40, .threads = 64 },
4641 );
4642 defer owned.deinit();
4643 owned.value.static_parameters = &.{};
4644 try testing.expectEqual(@as(?SpmvSell, null), spmvSellInstanceFromSpecialization(owned.value));
4645 }
4646
4647 {
4648 var owned = try spmvSellFamilySpecialization(
4649 allocator,
4650 .{ .rows = 70, .slice_size = 8, .values_size = 400, .x_extent = 40, .threads = 64 },
4651 );
4652 defer owned.deinit();
4653 const lifetime_allocator = owned.allocator();
4654 const static_parameters = try lifetime_allocator.alloc(entry.StaticParameter, 2);
4655 static_parameters[0] = owned.value.static_parameters[0];
4656 static_parameters[1] = try entry.runtimeStaticParameter(lifetime_allocator, spmv_sell_slice_size_parameter, 8);
4657 owned.value.static_parameters = static_parameters;
4658 try testing.expectEqual(@as(?SpmvSell, null), spmvSellInstanceFromSpecialization(owned.value));
4659 }
4660
4661 {
4662 var owned = try spmvSellFamilySpecialization(
4663 allocator,
4664 .{ .rows = 70, .slice_size = 8, .values_size = 400, .x_extent = 40, .threads = 64 },
4665 );
4666 defer owned.deinit();
4667 const lifetime_allocator = owned.allocator();
4668 const static_parameters = try lifetime_allocator.alloc(entry.StaticParameter, 2);
4669 static_parameters[0] = owned.value.static_parameters[0];
4670 static_parameters[1] = try entry.runtimeStaticParameter(lifetime_allocator, "unused", 1);
4671 owned.value.static_parameters = static_parameters;
4672 try testing.expectEqual(@as(?SpmvSell, null), spmvSellInstanceFromSpecialization(owned.value));
4673 }
4674
4675 {
4676 var owned = try spmvSellFamilySpecialization(
4677 allocator,
4678 .{ .rows = 70, .slice_size = 8, .values_size = 400, .x_extent = 40, .threads = 64 },
4679 );
4680 defer owned.deinit();
4681 const lifetime_allocator = owned.allocator();
4682 const inputs = try lifetime_allocator.alloc(entry.Shape, 4);
4683 inputs[0] = try entry.runtimeShape1D(lifetime_allocator, "z", 9);
4684 inputs[1] = owned.value.inputs[1];
4685 inputs[2] = owned.value.inputs[2];
4686 inputs[3] = owned.value.inputs[3];
4687 owned.value.inputs = inputs;
4688 try testing.expectEqual(@as(?SpmvSell, null), spmvSellInstanceFromSpecialization(owned.value));
4689 }
4690
4691 {
4692 var owned = try spmvSellFamilySpecialization(
4693 allocator,
4694 .{ .rows = 70, .slice_size = 8, .values_size = 400, .x_extent = 40, .threads = 64 },
4695 );
4696 defer owned.deinit();
4697 const lifetime_allocator = owned.allocator();
4698 const inputs = try lifetime_allocator.alloc(entry.Shape, 4);
4699 inputs[0] = owned.value.inputs[0];
4700 inputs[1] = owned.value.inputs[1];
4701 inputs[2] = try entry.runtimeShape1D(lifetime_allocator, "bad_n", 400);
4702 inputs[3] = owned.value.inputs[3];
4703 owned.value.inputs = inputs;
4704 try testing.expectEqual(@as(?SpmvSell, null), spmvSellInstanceFromSpecialization(owned.value));
4705 }
4706
4707 {
4708 var owned = try spmvSellFamilySpecialization(
4709 allocator,
4710 .{ .rows = 70, .slice_size = 8, .values_size = 400, .x_extent = 40, .threads = 64 },
4711 );
4712 defer owned.deinit();
4713 const lifetime_allocator = owned.allocator();
4714 owned.value.schedule = try entry.runtimeThreadBlocks1D(lifetime_allocator, "bad_r", 70, 64);
4715 owned.value.launch = owned.value.schedule.?.launch();
4716 try testing.expectEqual(@as(?SpmvSell, null), spmvSellInstanceFromSpecialization(owned.value));
4717 }
4718
4719 {
4720 var owned = try spmvSellFamilySpecialization(
4721 allocator,
4722 .{ .rows = 70, .slice_size = 8, .values_size = 400, .x_extent = 40, .threads = 64 },
4723 );
4724 defer owned.deinit();
4725 owned.value.structure = "row_warp";
4726 try testing.expectEqual(@as(?SpmvSell, null), spmvSellInstanceFromSpecialization(owned.value));
4727 }
4728 }
4729
4730 test "sparse spmm csr specialization rejects malformed descriptor facts" {
4731 const allocator = testing.allocator;
4732
4733 {
4734 var owned = try spmmCsrFamilySpecialization(
4735 allocator,
4736 .{ .rows = 70, .columns = 45, .nnz = 512, .x_extent = 40, .threads = .{ .x = 8, .y = 4 } },
4737 );
4738 defer owned.deinit();
4739 const lifetime_allocator = owned.allocator();
4740 const inputs = try lifetime_allocator.alloc(entry.Shape, 4);
4741 inputs[0] = owned.value.inputs[0];
4742 inputs[1] = owned.value.inputs[1];
4743 inputs[2] = owned.value.inputs[2];
4744 inputs[3] = try entry.runtimeShape2D(lifetime_allocator, "x", 40, "bad_c", 45);
4745 owned.value.inputs = inputs;
4746 try testing.expectEqual(@as(?SpmmCsr, null), spmmCsrInstanceFromSpecialization(owned.value));
4747 }
4748
4749 {
4750 var owned = try spmmCsrFamilySpecialization(
4751 allocator,
4752 .{ .rows = 70, .columns = 45, .nnz = 512, .x_extent = 40, .threads = .{ .x = 8, .y = 4 } },
4753 );
4754 defer owned.deinit();
4755 const lifetime_allocator = owned.allocator();
4756 const inputs = try lifetime_allocator.alloc(entry.Shape, 4);
4757 const zero_x_axes = try lifetime_allocator.alloc(entry.Axis, 2);
4758 zero_x_axes[0] = .{ .name = "x", .extent = 0 };
4759 zero_x_axes[1] = .{ .name = "c", .extent = 45 };
4760 inputs[0] = owned.value.inputs[0];
4761 inputs[1] = owned.value.inputs[1];
4762 inputs[2] = owned.value.inputs[2];
4763 inputs[3] = .{ .axes = zero_x_axes };
4764 owned.value.inputs = inputs;
4765 try testing.expectEqual(@as(?SpmmCsr, null), spmmCsrInstanceFromSpecialization(owned.value));
4766 }
4767
4768 {
4769 var owned = try spmmCsrFamilySpecialization(
4770 allocator,
4771 .{ .rows = 70, .columns = 45, .nnz = 512, .x_extent = 40, .threads = .{ .x = 8, .y = 4 } },
4772 );
4773 defer owned.deinit();
4774 const lifetime_allocator = owned.allocator();
4775 const inputs = try lifetime_allocator.alloc(entry.Shape, 4);
4776 inputs[0] = try entry.runtimeShape1D(lifetime_allocator, "r", 70);
4777 inputs[1] = owned.value.inputs[1];
4778 inputs[2] = owned.value.inputs[2];
4779 inputs[3] = owned.value.inputs[3];
4780 owned.value.inputs = inputs;
4781 try testing.expectEqual(@as(?SpmmCsr, null), spmmCsrInstanceFromSpecialization(owned.value));
4782 }
4783
4784 {
4785 var owned = try spmmCsrFamilySpecialization(
4786 allocator,
4787 .{ .rows = 70, .columns = 45, .nnz = 512, .x_extent = 40, .threads = .{ .x = 8, .y = 4 } },
4788 );
4789 defer owned.deinit();
4790 const lifetime_allocator = owned.allocator();
4791 owned.value.schedule = try entry.runtimeThreadBlocks2D(lifetime_allocator, "bad_c", 45, "r", 70, 8, 4);
4792 owned.value.launch = owned.value.schedule.?.launch();
4793 try testing.expectEqual(@as(?SpmmCsr, null), spmmCsrInstanceFromSpecialization(owned.value));
4794 }
4795 }