lib/accy/src/kernel/library/catalog/query.zig
daab053ee43316e1809a84551d573ddd1e5bf3d2
1 const choir_abi = @import("choir_abi");
2 const activation_mod = @import("../../../choir/root.zig").activation;
3 const library = @import("../root.zig");
4
5 const entry = library.entry;
6 const histogram = library.histogram;
7
8 pub const Query = union(enum) {
9 activation: ActivationQuery,
10 attention: AttentionQuery,
11 batched_matrix_product: BatchedMatrixProductQuery,
12 fused: FusedQuery,
13 gather: GatherQuery,
14 layout: LayoutQuery,
15 matrix_product: MatrixProductQuery,
16 matrix_vector_product: MatrixVectorProductQuery,
17 filter: FilterQuery,
18 outer_product: OuterProductQuery,
19 random: RandomQuery,
20 reduction: ReductionQuery,
21 row_normalization: RowNormalizationQuery,
22 row_sparse_cross_entropy: RowSparseCrossEntropyQuery,
23 scan: ScanQuery,
24 sort: SortQuery,
25 sparse: SparseQuery,
26 spatial: SpatialQuery,
27 image: ImageQuery,
28 factor: FactorQuery,
29 scatter: ScatterQuery,
30 scatter_add: ScatterAddQuery,
31 histogram: HistogramQuery,
32 segmented: SegmentedQuery,
33 stencil: StencilQuery,
34 };
35
36 pub const ActivationKind = activation_mod.Kind;
37
38 pub const ActivationQuery = struct {
39 dtype: choir_abi.DType,
40 kind: ActivationKind,
41 extent: u64,
42 };
43
44 pub const AttentionKind = entry.AttentionOperator;
45
46 pub const AttentionSchedule = union(enum) {
47 thread_blocks: entry.Threads3D,
48 };
49
50 pub const AttentionQuery = struct {
51 dtype: choir_abi.DType,
52 kind: AttentionKind,
53 query_indices: []const u8,
54 key_indices: []const u8,
55 value_indices: []const u8,
56 output_indices: []const u8,
57 query_dims: []const i64,
58 key_dims: []const i64,
59 value_dims: []const i64,
60 output_dims: []const i64,
61 schedule: ?AttentionSchedule = null,
62 };
63
64 pub const FusedQuery = union(enum) {
65 vector: FusedVectorQuery,
66 matrix_product: FusedMatrixProductQuery,
67 matrix_vector_product: FusedMatrixVectorProductQuery,
68 row_normalization: FusedRowNormalizationQuery,
69 };
70
71 pub const FusedVectorKind = union(enum) {
72 bias_activation: ActivationKind,
73 gated_activation: ActivationKind,
74 };
75
76 pub const FusedVectorQuery = struct {
77 dtype: choir_abi.DType,
78 kind: FusedVectorKind,
79 extent: u64,
80 };
81
82 pub const FusedLinalgEpilogue = union(enum) {
83 bias_activation: ActivationKind,
84 };
85
86 pub const FusedMatrixProductQuery = struct {
87 dtype: choir_abi.DType,
88 epilogue: FusedLinalgEpilogue,
89 lhs_indices: []const u8,
90 rhs_indices: []const u8,
91 bias_indices: []const u8,
92 output_indices: []const u8,
93 lhs_dims: []const i64,
94 rhs_dims: []const i64,
95 bias_dims: []const i64,
96 output_dims: []const i64,
97 schedule: ?MatrixProductSchedule = null,
98 };
99
100 pub const FusedMatrixVectorProductQuery = struct {
101 dtype: choir_abi.DType,
102 epilogue: FusedLinalgEpilogue,
103 matrix_indices: []const u8,
104 vector_indices: []const u8,
105 bias_indices: []const u8,
106 output_indices: []const u8,
107 matrix_dims: []const i64,
108 vector_dims: []const i64,
109 bias_dims: []const i64,
110 output_dims: []const i64,
111 schedule: ?MatrixVectorProductSchedule = null,
112 };
113
114 pub const MatrixProductSchedule = union(enum) {
115 thread_blocks: entry.Threads2D,
116 };
117
118 pub const BatchedMatrixProductSchedule = union(enum) {
119 thread_blocks: entry.Threads3D,
120 };
121
122 pub const MatrixVectorProductSchedule = union(enum) {
123 thread_blocks: u32,
124 };
125
126 pub const OuterProductSchedule = union(enum) {
127 thread_blocks: entry.Threads2D,
128 };
129
130 pub const FusedRowNormalizationKind = union(enum) {
131 residual: RowNormalizationKind,
132 };
133
134 pub const FusedRowNormalizationQuery = struct {
135 dtype: choir_abi.DType,
136 kind: FusedRowNormalizationKind,
137 rows: u64,
138 cols: u64,
139 schedule: ?RowNormalizationSchedule = null,
140 };
141
142 pub const LayoutKind = entry.LayoutOperator;
143
144 pub const LayoutQuery = struct {
145 dtype: choir_abi.DType,
146 kind: LayoutKind,
147 input_indices: []const u8,
148 output_indices: []const u8,
149 input_dims: []const i64,
150 output_dims: []const i64,
151 };
152
153 pub const TransposeQuery = LayoutQuery;
154
155 pub const MatrixProductQuery = struct {
156 dtype: choir_abi.DType,
157 lhs_indices: []const u8,
158 rhs_indices: []const u8,
159 output_indices: []const u8,
160 lhs_dims: []const i64,
161 rhs_dims: []const i64,
162 output_dims: []const i64,
163 schedule: ?MatrixProductSchedule = null,
164 };
165
166 pub const BatchedMatrixProductQuery = struct {
167 dtype: choir_abi.DType,
168 lhs_indices: []const u8,
169 rhs_indices: []const u8,
170 output_indices: []const u8,
171 lhs_dims: []const i64,
172 rhs_dims: []const i64,
173 output_dims: []const i64,
174 schedule: ?BatchedMatrixProductSchedule = null,
175 };
176
177 pub const MatrixVectorProductQuery = struct {
178 dtype: choir_abi.DType,
179 matrix_indices: []const u8,
180 vector_indices: []const u8,
181 output_indices: []const u8,
182 matrix_dims: []const i64,
183 vector_dims: []const i64,
184 output_dims: []const i64,
185 schedule: ?MatrixVectorProductSchedule = null,
186 };
187
188 pub const OuterProductQuery = struct {
189 dtype: choir_abi.DType,
190 lhs_indices: []const u8,
191 rhs_indices: []const u8,
192 output_indices: []const u8,
193 lhs_dims: []const i64,
194 rhs_dims: []const i64,
195 output_dims: []const i64,
196 schedule: ?OuterProductSchedule = null,
197 };
198
199 pub const ReductionKind = entry.ReductionOperator;
200
201 pub const ReductionOperand = struct {
202 indices: []const u8,
203 dims: []const i64,
204 };
205
206 pub const ReductionQuery = struct {
207 dtype: choir_abi.DType,
208 kind: ReductionKind,
209 inputs: []const ReductionOperand,
210 output_indices: []const u8,
211 output_dims: []const i64,
212 };
213
214 pub const RowNormalizationParameterization = entry.RowNormalizationParameterization;
215 pub const RowNormalizationKind = entry.RowNormalizationOperator;
216
217 pub const RowNormalizationSchedule = union(enum) {
218 thread_blocks: entry.Threads2D,
219 };
220
221 pub const RowNormalizationQuery = struct {
222 dtype: choir_abi.DType,
223 kind: RowNormalizationKind,
224 rows: u64,
225 cols: u64,
226 schedule: ?RowNormalizationSchedule = null,
227 };
228
229 pub const RowSparseCrossEntropySchedule = union(enum) {
230 thread_blocks: u32,
231 };
232
233 pub const RowSparseCrossEntropyQuery = struct {
234 dtype: choir_abi.DType,
235 rows: u64,
236 classes: u64,
237 schedule: ?RowSparseCrossEntropySchedule = null,
238 };
239
240 pub const GatherSchedule = union(enum) {
241 thread_blocks: u32,
242 };
243
244 pub const GatherQuery = struct {
245 dtype: choir_abi.DType,
246 outer: u64 = 1,
247 axis_size: u64,
248 gathered: u64,
249 inner: u64 = 1,
250 schedule: ?GatherSchedule = null,
251 };
252
253 pub const CompactionPredicateKind = enum {
254 nonzero,
255 greater_than,
256 };
257
258 pub const FilterSchedule = union(enum) {
259 thread_blocks: u32,
260 };
261
262 pub const FilterQuery = struct {
263 dtype: choir_abi.DType,
264 predicate: CompactionPredicateKind = .nonzero,
265 extent: u64,
266 schedule: ?FilterSchedule = null,
267 };
268
269 pub const RandomAlgorithm = enum {
270 philox,
271 threefry,
272 squares,
273 };
274
275 pub const RandomSchedule = union(enum) {
276 thread_blocks: u32,
277 };
278
279 pub const RandomQuery = struct {
280 dtype: choir_abi.DType,
281 algorithm: RandomAlgorithm = .philox,
282 count: u64,
283 rounds: u32 = 0,
284 schedule: ?RandomSchedule = null,
285 };
286
287 pub const ScatterSchedule = union(enum) {
288 thread_blocks: u32,
289 };
290
291 pub const ScatterAddSchedule = union(enum) {
292 thread_blocks: u32,
293 shared_bins: u32,
294 };
295
296 pub const HistogramSchedule = union(enum) {
297 thread_blocks: u32,
298 shared_bins: u32,
299 };
300
301 pub const HistogramBinningPolicy = histogram.HistogramBinningPolicy;
302
303 pub const HistogramQuery = struct {
304 dtype: choir_abi.DType,
305 bins: u64,
306 count: u64,
307 binning: HistogramBinningPolicy = .lower_inclusive_upper_exclusive,
308 schedule: ?HistogramSchedule = null,
309 };
310
311 pub const ScatterAddQuery = struct {
312 dtype: choir_abi.DType,
313 outer: u64 = 1,
314 axis_size: u64,
315 updates: u64,
316 inner: u64 = 1,
317 schedule: ?ScatterAddSchedule = null,
318 };
319
320 pub const ScatterQuery = struct {
321 dtype: choir_abi.DType,
322 outer: u64 = 1,
323 axis_size: u64,
324 updates: u64,
325 inner: u64 = 1,
326 schedule: ?ScatterSchedule = null,
327 };
328
329 pub const ScanKind = entry.ScanOperator;
330
331 pub const ScanSchedule = union(enum) {
332 thread_blocks: u32,
333 };
334
335 pub const ScanQuery = struct {
336 dtype: choir_abi.DType,
337 kind: ScanKind,
338 extent: u64,
339 schedule: ?ScanSchedule = null,
340 };
341
342 pub const SortKind = entry.SortOperator;
343
344 pub const SortSchedule = union(enum) {
345 thread_blocks: u32,
346 };
347
348 pub const SortStructure = enum {
349 radix_split,
350 radix_digit,
351 bitonic_block,
352 top_k_block,
353 top_k_block_pairs,
354 };
355
356 /// A caller fills this request to get a sort kernel from the kernel library, the package's
357 /// hand-written kernel families, chosen by a typed request. The request names the element type, the
358 /// sort operation, the number of elements, and `k` for the operations that keep only the first `k`
359 /// results. `structure` picks one way to carry out the sort, such as a radix pass or a bitonic
360 /// block, and `schedule` picks a thread-block size. Setting `structure` or `schedule` changes how
361 /// the sort is built and leaves the sort that is computed the same.
362 pub const SortQuery = struct {
363 dtype: choir_abi.DType,
364 kind: SortKind,
365 extent: u64,
366 k: u64 = 0,
367 structure: ?SortStructure = null,
368 schedule: ?SortSchedule = null,
369 };
370
371 pub const FactorTileLayout = enum {
372 row_major,
373 interleaved,
374 };
375
376 pub const FactorKind = union(enum) {
377 batched_cholesky: struct { n: u32, batch: u64 },
378 batched_cholesky_solve: struct { n: u32, batch: u64 },
379 batched_inverse: struct { n: u32, batch: u64 },
380 };
381
382 pub const FactorSchedule = union(enum) {
383 thread_blocks: u32,
384 };
385
386 /// A caller fills this request to get a batched Cholesky factorization, Cholesky solve or matrix
387 /// inverse from the kernel library. The request names the element type, the operation with its
388 /// matrix size `n` and batch count, the buffer layout, and an optional thread-block count. `layout`
389 /// states how the matrices sit in the device buffer, row by row or interleaved across the batch.
390 /// Because the layout changes where each element is read, it is part of which kernel the request
391 /// names, and it enters the kernel's name.
392 pub const FactorQuery = struct {
393 dtype: choir_abi.DType,
394 kind: FactorKind,
395 layout: FactorTileLayout = .row_major,
396 schedule: ?FactorSchedule = null,
397 };
398
399 pub const SparseStructure = enum {
400 element_thread,
401 row_thread,
402 row_warp,
403 row_column_thread,
404 };
405
406 pub const SparseKind = union(enum) {
407 coo_spmv: struct { rows: u64, nnz: u64, x_extent: u64 },
408 csr_spmv: struct { rows: u64, nnz: u64, x_extent: u64 },
409 csr_spmm: struct { rows: u64, columns: u64, nnz: u64, x_extent: u64 },
410 ell_spmv: struct { rows: u64, slots: u64, x_extent: u64 },
411 sell_spmv: struct { rows: u64, slice_size: u64, values_size: u64, x_extent: u64 },
412 };
413
414 pub const SparseSchedule = union(enum) {
415 thread_blocks: u32,
416 thread_blocks_2d: entry.Threads2D,
417 };
418
419 pub const SparseQuery = struct {
420 dtype: choir_abi.DType,
421 kind: SparseKind,
422 structure: ?SparseStructure = null,
423 schedule: ?SparseSchedule = null,
424 };
425
426 pub const SpatialKind = union(enum) {
427 grid_cells,
428 grid_count: struct { cells: u32 },
429 grid_neighbor_count: struct { cells: u32, stride: u32 },
430 };
431
432 pub const SpatialSchedule = union(enum) {
433 thread_blocks: u32,
434 };
435
436 /// A caller fills this request to get a grid kernel for points in the plane, such as assigning
437 /// points to cells or counting neighbors. The request names the element type, the grid operation,
438 /// the number of points, and an optional thread-block count. Each operation also carries the sizes
439 /// fixed when the kernel is compiled, such as the cell count and, for neighbor counting, the
440 /// stride, beside the point count, which is given at run time.
441 pub const SpatialQuery = struct {
442 dtype: choir_abi.DType,
443 kind: SpatialKind,
444 count: u64,
445 schedule: ?SpatialSchedule = null,
446 };
447
448 pub const ImageKind = union(enum) {
449 blur_pass: struct {
450 radius: u32 = 1,
451 axis: library.image.Axis,
452 },
453 resize_bilinear: struct {
454 src_width: u64,
455 src_height: u64,
456 },
457 };
458
459 pub const ImageSchedule = union(enum) {
460 thread_blocks: entry.Threads2D,
461 };
462
463 pub const ImageQuery = struct {
464 dtype: choir_abi.DType,
465 kind: ImageKind,
466 width: u64,
467 height: u64,
468 schedule: ?ImageSchedule = null,
469 };
470
471 pub const SegmentedKind = entry.SegmentedOperator;
472
473 pub const SegmentedSchedule = union(enum) {
474 thread_blocks: u32,
475 warp_blocks: u32,
476 };
477
478 pub const SegmentedQuery = struct {
479 dtype: choir_abi.DType,
480 kind: SegmentedKind,
481 segments: u64,
482 total: u64,
483 schedule: ?SegmentedSchedule = null,
484 };
485
486 pub const StencilKind = entry.StencilOperator;
487
488 pub const StencilSchedule = union(enum) {
489 thread_blocks: entry.Threads2D,
490 };
491
492 pub const StencilQuery = struct {
493 dtype: choir_abi.DType,
494 kind: StencilKind,
495 rows: u64,
496 cols: u64,
497 radius: u32 = 1,
498 schedule: ?StencilSchedule = null,
499 };
500
501 pub const ArtifactRequest = struct {
502 target: []const u8,
503 version: u32 = 1,
504 options: entry.ArtifactOptions,
505 };