lib/accy/src/kernel/library/catalog/family/sparse/selection.zig
daab053ee43316e1809a84551d573ddd1e5bf3d2
1 const std = @import("std");
2
3 const descriptor_mod = @import("../../root.zig");
4 const sparse_mod = @import("../../../root.zig").sparse;
5 const tuning_mod = @import("../../../root.zig").tuning;
6 const query_mod = descriptor_mod;
7
8 const candidate_mod = @import("candidate.zig");
9 const descriptor_build = @import("descriptor.zig");
10 const structure_mod = @import("structure/root.zig");
11 const thread_mod = @import("threads.zig");
12
13 const OwnedDescriptor = descriptor_mod.OwnedDescriptor;
14 const SparseQuery = query_mod.SparseQuery;
15 const SparseCandidateDescriptors = candidate_mod.SparseCandidateDescriptors;
16
17 pub fn selectSparse(backing_allocator: std.mem.Allocator, query: SparseQuery) !?OwnedDescriptor {
18 if (!structure_mod.canonicalSparse(query)) return null;
19 const accumulation_dtype = sparse_mod.sparseAccumulationDType(query.dtype) orelse return null;
20 switch (query.kind) {
21 .coo_spmv => |facts| {
22 const structure = structure_mod.sparseResolvedCooStructure(query) orelse return null;
23 const coo_accumulation_dtype = sparse_mod.spmvCooAccumulationDTypeForStructure(structure, query.dtype) orelse return null;
24 const threads = thread_mod.sparseQueryCooThreads(query, facts.rows, facts.nnz, structure) orelse return null;
25 return try descriptor_build.spmvCooDescriptorForInstance(backing_allocator, query, .{
26 .rows = facts.rows,
27 .nnz = facts.nnz,
28 .x_extent = facts.x_extent,
29 .dtype = query.dtype,
30 .accumulation_dtype = coo_accumulation_dtype,
31 .threads = threads,
32 .structure = structure,
33 });
34 },
35 .csr_spmv => |facts| {
36 const structure = structure_mod.sparseResolvedSpmvStructure(query) orelse return null;
37 const threads = thread_mod.sparseQuerySpmvThreads(query, facts.rows, structure) orelse return null;
38 return try descriptor_build.spmvCsrDescriptorForInstance(backing_allocator, query, .{
39 .rows = facts.rows,
40 .nnz = facts.nnz,
41 .x_extent = facts.x_extent,
42 .dtype = query.dtype,
43 .accumulation_dtype = accumulation_dtype,
44 .threads = threads,
45 .structure = structure,
46 });
47 },
48 .ell_spmv => |facts| {
49 const structure = structure_mod.sparseResolvedEllStructure(query) orelse return null;
50 const threads = thread_mod.sparseQueryEllThreads(query, facts.rows) orelse return null;
51 return try descriptor_build.spmvEllDescriptorForInstance(backing_allocator, query, .{
52 .rows = facts.rows,
53 .slots = facts.slots,
54 .x_extent = facts.x_extent,
55 .dtype = query.dtype,
56 .accumulation_dtype = accumulation_dtype,
57 .threads = threads,
58 .structure = structure,
59 });
60 },
61 .sell_spmv => |facts| {
62 const structure = structure_mod.sparseResolvedSellStructure(query) orelse return null;
63 const threads = thread_mod.sparseQuerySellThreads(query, facts.rows) orelse return null;
64 return try descriptor_build.spmvSellDescriptorForInstance(backing_allocator, query, .{
65 .rows = facts.rows,
66 .slice_size = facts.slice_size,
67 .values_size = facts.values_size,
68 .x_extent = facts.x_extent,
69 .dtype = query.dtype,
70 .accumulation_dtype = accumulation_dtype,
71 .threads = threads,
72 .structure = structure,
73 });
74 },
75 .csr_spmm => |facts| {
76 const structure = structure_mod.sparseResolvedSpmmStructure(query) orelse return null;
77 const threads = thread_mod.sparseQuerySpmmThreads(query, facts) orelse return null;
78 return try descriptor_build.spmmCsrDescriptorForInstance(backing_allocator, query, .{
79 .rows = facts.rows,
80 .columns = facts.columns,
81 .nnz = facts.nnz,
82 .x_extent = facts.x_extent,
83 .dtype = query.dtype,
84 .accumulation_dtype = accumulation_dtype,
85 .threads = threads,
86 .structure = structure,
87 });
88 },
89 }
90 }
91
92 pub fn selectSparseCandidates(
93 backing_allocator: std.mem.Allocator,
94 query: SparseQuery,
95 ) !SparseCandidateDescriptors {
96 var result = SparseCandidateDescriptors{};
97 errdefer result.deinit();
98
99 if (!structure_mod.canonicalSparse(query)) return result;
100 switch (query.kind) {
101 .coo_spmv => {
102 if (query.structure != null) {
103 if (try selectSparse(backing_allocator, query)) |descriptor| {
104 result.items[result.count] = descriptor;
105 result.count += 1;
106 }
107 return result;
108 }
109 const structures = [_]query_mod.SparseStructure{ .element_thread, .row_thread };
110 for (structures) |structure| {
111 var candidate = query;
112 candidate.structure = structure;
113 if (try selectSparse(backing_allocator, candidate)) |descriptor| {
114 result.items[result.count] = descriptor;
115 result.count += 1;
116 }
117 }
118 return result;
119 },
120 .csr_spmv => {
121 if (query.structure != null) {
122 if (try selectSparse(backing_allocator, query)) |descriptor| {
123 result.items[result.count] = descriptor;
124 result.count += 1;
125 }
126 return result;
127 }
128 const structures = [_]query_mod.SparseStructure{ .row_thread, .row_warp };
129 for (structures) |structure| {
130 var candidate = query;
131 candidate.structure = structure;
132 if (try selectSparse(backing_allocator, candidate)) |descriptor| {
133 result.items[result.count] = descriptor;
134 result.count += 1;
135 }
136 }
137 return result;
138 },
139 .ell_spmv => |facts| {
140 if (query.structure != null or query.schedule != null) {
141 if (try selectSparse(backing_allocator, query)) |descriptor| {
142 result.items[result.count] = descriptor;
143 result.count += 1;
144 }
145 return result;
146 }
147 if (try selectSparse(backing_allocator, query)) |descriptor| {
148 result.items[result.count] = descriptor;
149 result.count += 1;
150 }
151 const accumulation_dtype = sparse_mod.sparseAccumulationDType(query.dtype) orelse return result;
152 const structure = structure_mod.sparseResolvedEllStructure(query) orelse return result;
153 const thread_candidates = sparse_mod.spmvEllThreadCandidatesForRows(facts.rows);
154 for (thread_candidates.slice()) |threads| {
155 var candidate = sparse_mod.SpmvEll{
156 .rows = facts.rows,
157 .slots = facts.slots,
158 .x_extent = facts.x_extent,
159 .dtype = query.dtype,
160 .accumulation_dtype = accumulation_dtype,
161 .threads = threads,
162 .structure = structure,
163 };
164 candidate.threads = sparse_mod.spmvEllRepresentableThreads(candidate) orelse continue;
165 var scheduled = query;
166 scheduled.schedule = .{ .thread_blocks = candidate.threads };
167 if (try descriptor_build.spmvEllDescriptorForInstance(backing_allocator, scheduled, candidate)) |descriptor_owned| {
168 var descriptor = descriptor_owned;
169 if (candidate_mod.sparseCandidateTargetsContain(result.slice(), descriptor.descriptor.metadata.target)) {
170 descriptor.deinit();
171 continue;
172 }
173 result.items[result.count] = descriptor;
174 result.count += 1;
175 }
176 }
177 return result;
178 },
179 .sell_spmv => |facts| {
180 if (query.structure != null or query.schedule != null) {
181 if (try selectSparse(backing_allocator, query)) |descriptor| {
182 result.items[result.count] = descriptor;
183 result.count += 1;
184 }
185 return result;
186 }
187 if (try selectSparse(backing_allocator, query)) |descriptor| {
188 result.items[result.count] = descriptor;
189 result.count += 1;
190 }
191 const accumulation_dtype = sparse_mod.sparseAccumulationDType(query.dtype) orelse return result;
192 const structure = structure_mod.sparseResolvedSellStructure(query) orelse return result;
193 const thread_candidates = sparse_mod.spmvSellThreadCandidatesForRows(facts.rows);
194 for (thread_candidates.slice()) |threads| {
195 var candidate = sparse_mod.SpmvSell{
196 .rows = facts.rows,
197 .slice_size = facts.slice_size,
198 .values_size = facts.values_size,
199 .x_extent = facts.x_extent,
200 .dtype = query.dtype,
201 .accumulation_dtype = accumulation_dtype,
202 .threads = threads,
203 .structure = structure,
204 };
205 candidate.threads = sparse_mod.spmvSellRepresentableThreads(candidate) orelse continue;
206 var scheduled = query;
207 scheduled.schedule = .{ .thread_blocks = candidate.threads };
208 if (try descriptor_build.spmvSellDescriptorForInstance(backing_allocator, scheduled, candidate)) |descriptor_owned| {
209 var descriptor = descriptor_owned;
210 if (candidate_mod.sparseCandidateTargetsContain(result.slice(), descriptor.descriptor.metadata.target)) {
211 descriptor.deinit();
212 continue;
213 }
214 result.items[result.count] = descriptor;
215 result.count += 1;
216 }
217 }
218 return result;
219 },
220 .csr_spmm => |facts| {
221 if (query.structure != null or query.schedule != null) {
222 if (try selectSparse(backing_allocator, query)) |descriptor| {
223 result.items[result.count] = descriptor;
224 result.count += 1;
225 }
226 return result;
227 }
228 if (try selectSparse(backing_allocator, query)) |descriptor| {
229 result.items[result.count] = descriptor;
230 result.count += 1;
231 }
232 const accumulation_dtype = sparse_mod.sparseAccumulationDType(query.dtype) orelse return result;
233 const structure = structure_mod.sparseResolvedSpmmStructure(query) orelse return result;
234 const thread_candidates = sparse_mod.spmmCsrThreadCandidatesForExtents(facts.rows, facts.columns);
235 for (thread_candidates.slice()) |threads| {
236 var candidate = sparse_mod.SpmmCsr{
237 .rows = facts.rows,
238 .columns = facts.columns,
239 .nnz = facts.nnz,
240 .x_extent = facts.x_extent,
241 .dtype = query.dtype,
242 .accumulation_dtype = accumulation_dtype,
243 .threads = threads,
244 .structure = structure,
245 };
246 candidate.threads = sparse_mod.spmmCsrRepresentableThreads(candidate) orelse continue;
247 var scheduled = query;
248 scheduled.schedule = .{ .thread_blocks_2d = candidate.threads };
249 if (try descriptor_build.spmmCsrDescriptorForInstance(backing_allocator, scheduled, candidate)) |descriptor_owned| {
250 var descriptor = descriptor_owned;
251 if (candidate_mod.sparseCandidateTargetsContain(result.slice(), descriptor.descriptor.metadata.target)) {
252 descriptor.deinit();
253 continue;
254 }
255 result.items[result.count] = descriptor;
256 result.count += 1;
257 }
258 }
259 return result;
260 },
261 }
262 }
263
264 pub fn selectSparseWithTuning(
265 backing_allocator: std.mem.Allocator,
266 query: SparseQuery,
267 reader: tuning_mod.FamilyTuningReader,
268 ) !?OwnedDescriptor {
269 if (!structure_mod.canonicalSparse(query)) return null;
270 if (query.structure != null) return selectSparse(backing_allocator, query);
271 switch (query.kind) {
272 .coo_spmv => |facts| {
273 const structure = structure_mod.sparseResolvedCooStructure(query) orelse return null;
274 const accumulation_dtype = sparse_mod.spmvCooAccumulationDTypeForStructure(structure, query.dtype) orelse return null;
275 const threads = thread_mod.sparseQueryCooThreads(query, facts.rows, facts.nnz, structure) orelse return null;
276 const resolved = try sparse_mod.resolveSpmvCooStructure(backing_allocator, reader, .{
277 .rows = facts.rows,
278 .nnz = facts.nnz,
279 .x_extent = facts.x_extent,
280 .dtype = query.dtype,
281 .accumulation_dtype = accumulation_dtype,
282 .threads = threads,
283 .structure = structure,
284 });
285 var tuned = query;
286 tuned.structure = if (resolved) |resolved_structure| switch (resolved_structure) {
287 .element_thread => .element_thread,
288 .row_thread => .row_thread,
289 } else return selectSparse(backing_allocator, query);
290 return selectSparse(backing_allocator, tuned);
291 },
292 .csr_spmv => |facts| {
293 const accumulation_dtype = sparse_mod.sparseAccumulationDType(query.dtype) orelse return null;
294 const structure = structure_mod.sparseResolvedSpmvStructure(query) orelse return null;
295 const threads = thread_mod.sparseQuerySpmvThreads(query, facts.rows, structure) orelse return null;
296 const resolved = try sparse_mod.resolveSpmvCsrStructure(backing_allocator, reader, .{
297 .rows = facts.rows,
298 .nnz = facts.nnz,
299 .x_extent = facts.x_extent,
300 .dtype = query.dtype,
301 .accumulation_dtype = accumulation_dtype,
302 .threads = threads,
303 .structure = structure,
304 });
305 var tuned = query;
306 tuned.structure = if (resolved) |resolved_structure| switch (resolved_structure) {
307 .row_thread => .row_thread,
308 .row_warp => .row_warp,
309 } else return selectSparse(backing_allocator, query);
310 return selectSparse(backing_allocator, tuned);
311 },
312 .ell_spmv => |facts| {
313 if (query.schedule != null) return selectSparse(backing_allocator, query);
314 const accumulation_dtype = sparse_mod.sparseAccumulationDType(query.dtype) orelse return null;
315 const structure = structure_mod.sparseResolvedEllStructure(query) orelse return null;
316 const threads = thread_mod.sparseQueryEllThreads(query, facts.rows) orelse return null;
317 const resolved = try sparse_mod.resolveSpmvEllThreads(backing_allocator, reader, .{
318 .rows = facts.rows,
319 .slots = facts.slots,
320 .x_extent = facts.x_extent,
321 .dtype = query.dtype,
322 .accumulation_dtype = accumulation_dtype,
323 .threads = threads,
324 .structure = structure,
325 });
326 var tuned = query;
327 tuned.schedule = if (resolved) |resolved_threads|
328 .{ .thread_blocks = resolved_threads }
329 else
330 return selectSparse(backing_allocator, query);
331 return selectSparse(backing_allocator, tuned);
332 },
333 .sell_spmv => |facts| {
334 if (query.schedule != null) return selectSparse(backing_allocator, query);
335 const accumulation_dtype = sparse_mod.sparseAccumulationDType(query.dtype) orelse return null;
336 const structure = structure_mod.sparseResolvedSellStructure(query) orelse return null;
337 const threads = thread_mod.sparseQuerySellThreads(query, facts.rows) orelse return null;
338 const resolved = try sparse_mod.resolveSpmvSellThreads(backing_allocator, reader, .{
339 .rows = facts.rows,
340 .slice_size = facts.slice_size,
341 .values_size = facts.values_size,
342 .x_extent = facts.x_extent,
343 .dtype = query.dtype,
344 .accumulation_dtype = accumulation_dtype,
345 .threads = threads,
346 .structure = structure,
347 });
348 var tuned = query;
349 tuned.schedule = if (resolved) |resolved_threads|
350 .{ .thread_blocks = resolved_threads }
351 else
352 return selectSparse(backing_allocator, query);
353 return selectSparse(backing_allocator, tuned);
354 },
355 .csr_spmm => |facts| {
356 if (query.schedule != null) return selectSparse(backing_allocator, query);
357 const accumulation_dtype = sparse_mod.sparseAccumulationDType(query.dtype) orelse return null;
358 const structure = structure_mod.sparseResolvedSpmmStructure(query) orelse return null;
359 const threads = thread_mod.sparseQuerySpmmThreads(query, facts) orelse return null;
360 const resolved = try sparse_mod.resolveSpmmCsrThreads(backing_allocator, reader, .{
361 .rows = facts.rows,
362 .columns = facts.columns,
363 .nnz = facts.nnz,
364 .x_extent = facts.x_extent,
365 .dtype = query.dtype,
366 .accumulation_dtype = accumulation_dtype,
367 .threads = threads,
368 .structure = structure,
369 });
370 var tuned = query;
371 tuned.schedule = if (resolved) |resolved_threads|
372 .{ .thread_blocks_2d = resolved_threads }
373 else
374 return selectSparse(backing_allocator, query);
375 return selectSparse(backing_allocator, tuned);
376 },
377 }
378 }