lib/accy/src/kernel/library/catalog/family/scan/prefix.zig
daab053ee43316e1809a84551d573ddd1e5bf3d2
1 const std = @import("std");
2
3 const descriptor_mod = @import("../../root.zig");
4 const scan = @import("../../../root.zig").scan;
5 const match_mod = @import("../../match/root.zig");
6 const query_mod = descriptor_mod;
7
8 const Descriptor = descriptor_mod.Descriptor;
9 const OwnedDescriptor = descriptor_mod.OwnedDescriptor;
10 const ScanQuery = query_mod.ScanQuery;
11
12 pub const PrefixSumCandidateDescriptors = struct {
13 count: usize = 0,
14 items: [8]OwnedDescriptor = undefined,
15
16 pub fn slice(self: *const PrefixSumCandidateDescriptors) []const OwnedDescriptor {
17 return self.items[0..self.count];
18 }
19
20 pub fn deinit(self: *PrefixSumCandidateDescriptors) void {
21 for (self.items[0..self.count]) |*descriptor| descriptor.deinit();
22 self.* = undefined;
23 }
24 };
25
26 pub fn selectPrefixSum(backing_allocator: std.mem.Allocator, query: ScanQuery) !?OwnedDescriptor {
27 if (canonicalPrefixSum(query)) {
28 const instance = prefixSumFamilyInstance(query) orelse return null;
29 return try prefixSumDescriptorForInstance(backing_allocator, instance, query);
30 }
31 if (canonicalDeviceScan(query)) {
32 const instance = deviceScanFamilyInstance(query) orelse return null;
33 return try deviceScanDescriptorForInstance(backing_allocator, instance, query);
34 }
35 return null;
36 }
37
38 pub fn selectPrefixSumCandidates(
39 backing_allocator: std.mem.Allocator,
40 query: ScanQuery,
41 ) !PrefixSumCandidateDescriptors {
42 var result = PrefixSumCandidateDescriptors{};
43 errdefer result.deinit();
44
45 if (query.schedule != null) {
46 if (try selectPrefixSum(backing_allocator, query)) |descriptor| {
47 result.items[result.count] = descriptor;
48 result.count += 1;
49 }
50 return result;
51 }
52
53 if (canonicalPrefixSum(query)) {
54 const thread_candidates = scan.prefixSumThreadCandidatesForExtent(query.extent);
55 for (thread_candidates.slice()) |threads| {
56 const instance = scan.PrefixSum{
57 .extent = query.extent,
58 .dtype = query.dtype,
59 .mode = scan.PrefixSumMode.fromOperation(query.kind),
60 .threads = threads,
61 };
62 if (try prefixSumDescriptorForInstance(backing_allocator, instance, query)) |descriptor| {
63 result.items[result.count] = descriptor;
64 result.count += 1;
65 }
66 }
67 return result;
68 }
69
70 if (canonicalDeviceScan(query)) {
71 const thread_candidates = scan.deviceScanThreadCandidatesForExtent(query.extent);
72 for (thread_candidates.slice()) |threads| {
73 const instance = scan.DeviceScan{
74 .extent = query.extent,
75 .dtype = query.dtype,
76 .mode = scan.PrefixSumMode.fromOperation(query.kind),
77 .threads = threads,
78 };
79 if (try deviceScanDescriptorForInstance(backing_allocator, instance, query)) |descriptor| {
80 result.items[result.count] = descriptor;
81 result.count += 1;
82 }
83 }
84 }
85 return result;
86 }
87
88 fn canonicalPrefixSum(query: ScanQuery) bool {
89 if (!scan.prefixSumDTypeSupported(query.dtype)) return false;
90 return query.extent != 0 and query.extent <= scan.prefix_sum_max_threads;
91 }
92
93 fn canonicalDeviceScan(query: ScanQuery) bool {
94 if (!scan.deviceScanDTypeSupported(query.dtype)) return false;
95 if (query.extent <= scan.prefix_sum_max_threads) return false;
96 return scan.deviceScanThreadsForExtent(query.extent) != null;
97 }
98
99 fn deviceScanFamilyInstance(query: ScanQuery) ?scan.DeviceScan {
100 var instance = scan.DeviceScan{
101 .extent = query.extent,
102 .dtype = query.dtype,
103 .mode = scan.PrefixSumMode.fromOperation(query.kind),
104 .threads = scan.deviceScanThreadsForExtent(query.extent) orelse return null,
105 };
106 if (query.schedule) |requested| {
107 switch (requested) {
108 .thread_blocks => |threads| {
109 instance.threads = threads;
110 },
111 }
112 }
113 if (!scan.deviceScanInstanceValid(instance)) return null;
114 return instance;
115 }
116
117 fn deviceScanDescriptorForInstance(
118 backing_allocator: std.mem.Allocator,
119 instance: scan.DeviceScan,
120 query: ScanQuery,
121 ) !?OwnedDescriptor {
122 var specialization = try scan.deviceScanFamilySpecialization(backing_allocator, instance);
123 errdefer specialization.deinit();
124 const lifetime_allocator = specialization.allocator();
125 const descriptor = Descriptor{
126 .name = try scan.deviceScanFamilyTarget(lifetime_allocator, instance),
127 .metadata = .{
128 .target = try scan.deviceScanFamilyTarget(lifetime_allocator, instance),
129 .version = scan.device_scan_family_version,
130 .layer = .logical,
131 .category = .scan,
132 .specialization = specialization.value,
133 },
134 };
135 if (!match_mod.prefixSumDescriptorMatches(descriptor, query)) {
136 specialization.deinit();
137 return null;
138 }
139 return .{ .descriptor = descriptor, .specialization = specialization };
140 }
141
142 fn prefixSumDescriptorForInstance(
143 backing_allocator: std.mem.Allocator,
144 instance: scan.PrefixSum,
145 query: ScanQuery,
146 ) !?OwnedDescriptor {
147 var specialization = try scan.prefixSumFamilySpecialization(backing_allocator, instance);
148 errdefer specialization.deinit();
149 const lifetime_allocator = specialization.allocator();
150 const descriptor = Descriptor{
151 .name = try scan.prefixSumFamilyEntryName(lifetime_allocator, instance),
152 .metadata = .{
153 .target = try scan.prefixSumFamilyTarget(lifetime_allocator, instance),
154 .version = scan.prefix_sum_family_version,
155 .layer = .logical,
156 .category = .scan,
157 .specialization = specialization.value,
158 },
159 };
160 if (!match_mod.prefixSumDescriptorMatches(descriptor, query)) {
161 specialization.deinit();
162 return null;
163 }
164 return .{ .descriptor = descriptor, .specialization = specialization };
165 }
166
167 fn prefixSumFamilyInstance(query: ScanQuery) ?scan.PrefixSum {
168 var instance = scan.PrefixSum{
169 .extent = query.extent,
170 .dtype = query.dtype,
171 .mode = scan.PrefixSumMode.fromOperation(query.kind),
172 .threads = scan.prefixSumThreadsForExtent(query.extent) orelse return null,
173 };
174 if (query.schedule) |requested| {
175 switch (requested) {
176 .thread_blocks => |threads| {
177 instance.threads = threads;
178 },
179 }
180 }
181 if (!scan.prefixSumInstanceValid(instance)) return null;
182 return instance;
183 }