lib/accy/src/kernel/library/catalog/family/spatial.zig
daab053ee43316e1809a84551d573ddd1e5bf3d2
1 const std = @import("std");
2
3 const descriptor_mod = @import("../root.zig");
4 const geometry = @import("../../root.zig").geometry;
5 const spatial_mod = @import("../../root.zig").spatial;
6 const match_mod = @import("../match/root.zig");
7 const query_mod = descriptor_mod;
8
9 const Descriptor = descriptor_mod.Descriptor;
10 const OwnedDescriptor = descriptor_mod.OwnedDescriptor;
11 const SpatialQuery = query_mod.SpatialQuery;
12
13 pub const SpatialCandidateDescriptors = struct {
14 count: usize = 0,
15 items: [geometry.max_thread_candidates]OwnedDescriptor = undefined,
16
17 pub fn slice(self: *const SpatialCandidateDescriptors) []const OwnedDescriptor {
18 return self.items[0..self.count];
19 }
20
21 pub fn deinit(self: *SpatialCandidateDescriptors) void {
22 for (self.items[0..self.count]) |*descriptor| descriptor.deinit();
23 self.* = undefined;
24 }
25 };
26
27 pub fn selectSpatial(backing_allocator: std.mem.Allocator, query: SpatialQuery) !?OwnedDescriptor {
28 if (!canonicalSpatial(query)) return null;
29 const threads = spatialQueryThreads(query) orelse return null;
30 return try spatialDescriptorForThreads(backing_allocator, query, threads);
31 }
32
33 pub fn selectSpatialCandidates(
34 backing_allocator: std.mem.Allocator,
35 query: SpatialQuery,
36 ) !SpatialCandidateDescriptors {
37 var result = SpatialCandidateDescriptors{};
38 errdefer result.deinit();
39
40 if (!canonicalSpatial(query)) return result;
41 if (query.schedule != null) {
42 if (try selectSpatial(backing_allocator, query)) |descriptor| {
43 result.items[result.count] = descriptor;
44 result.count += 1;
45 }
46 return result;
47 }
48
49 const thread_candidates = spatial_mod.spatialThreadCandidatesForCount(query.count);
50 for (thread_candidates.slice()) |threads| {
51 if (result.count >= result.items.len) break;
52 if (try spatialDescriptorForThreads(backing_allocator, query, threads)) |descriptor| {
53 result.items[result.count] = descriptor;
54 result.count += 1;
55 }
56 }
57 return result;
58 }
59
60 fn canonicalSpatial(query: SpatialQuery) bool {
61 if (query.count == 0) return false;
62 return switch (query.kind) {
63 .grid_cells => query.dtype == .f32,
64 .grid_count => |facts| query.dtype == .i32 and facts.cells != 0,
65 .grid_neighbor_count => |facts| query.dtype == .f32 and facts.cells != 0 and facts.stride != 0,
66 };
67 }
68
69 fn spatialQueryThreads(query: SpatialQuery) ?u32 {
70 if (query.schedule) |requested| {
71 switch (requested) {
72 .thread_blocks => |threads| {
73 if (threads == 0) return null;
74 return threads;
75 },
76 }
77 }
78 return spatial_mod.spatialThreadsForCount(query.count);
79 }
80
81 fn spatialDescriptorForThreads(
82 backing_allocator: std.mem.Allocator,
83 query: SpatialQuery,
84 threads: u32,
85 ) !?OwnedDescriptor {
86 switch (query.kind) {
87 .grid_cells => {
88 const instance = spatial_mod.GridCells{ .count = query.count, .threads = threads };
89 if (!spatial_mod.gridCellsInstanceValid(instance)) return null;
90 var specialization = try spatial_mod.gridCellsFamilySpecialization(backing_allocator, instance);
91 errdefer specialization.deinit();
92 const lifetime_allocator = specialization.allocator();
93 const descriptor = Descriptor{
94 .name = try spatial_mod.gridCellsFamilyEntryName(lifetime_allocator, instance),
95 .metadata = .{
96 .target = try spatial_mod.gridCellsFamilyTarget(lifetime_allocator, instance),
97 .version = spatial_mod.grid_cells_family_version,
98 .layer = .logical,
99 .category = .spatial,
100 .specialization = specialization.value,
101 },
102 };
103 if (!match_mod.spatialDescriptorMatches(descriptor, query)) {
104 specialization.deinit();
105 return null;
106 }
107 return .{ .descriptor = descriptor, .specialization = specialization };
108 },
109 .grid_count => |facts| {
110 const instance = spatial_mod.GridCount{
111 .count = query.count,
112 .cells = facts.cells,
113 .threads = threads,
114 };
115 if (!spatial_mod.gridCountInstanceValid(instance)) return null;
116 var specialization = try spatial_mod.gridCountFamilySpecialization(backing_allocator, instance);
117 errdefer specialization.deinit();
118 const lifetime_allocator = specialization.allocator();
119 const descriptor = Descriptor{
120 .name = try spatial_mod.gridCountFamilyEntryName(lifetime_allocator, instance),
121 .metadata = .{
122 .target = try spatial_mod.gridCountFamilyTarget(lifetime_allocator, instance),
123 .version = spatial_mod.grid_count_family_version,
124 .layer = .logical,
125 .category = .spatial,
126 .specialization = specialization.value,
127 },
128 };
129 if (!match_mod.spatialDescriptorMatches(descriptor, query)) {
130 specialization.deinit();
131 return null;
132 }
133 return .{ .descriptor = descriptor, .specialization = specialization };
134 },
135 .grid_neighbor_count => |facts| {
136 const offsets_extent = @as(u64, facts.cells) * facts.stride;
137 const instance = spatial_mod.GridNeighborCount{
138 .count = query.count,
139 .offsets_extent = offsets_extent,
140 .threads = threads,
141 };
142 if (!spatial_mod.gridNeighborCountInstanceValid(instance)) return null;
143 var specialization = try spatial_mod.gridNeighborCountFamilySpecialization(
144 backing_allocator,
145 instance,
146 );
147 errdefer specialization.deinit();
148 const lifetime_allocator = specialization.allocator();
149 const descriptor = Descriptor{
150 .name = try spatial_mod.gridNeighborCountFamilyEntryName(lifetime_allocator, instance),
151 .metadata = .{
152 .target = try spatial_mod.gridNeighborCountFamilyTarget(lifetime_allocator, instance),
153 .version = spatial_mod.grid_neighbor_count_family_version,
154 .layer = .logical,
155 .category = .spatial,
156 .specialization = specialization.value,
157 },
158 };
159 if (!match_mod.spatialDescriptorMatches(descriptor, query)) {
160 specialization.deinit();
161 return null;
162 }
163 return .{ .descriptor = descriptor, .specialization = specialization };
164 },
165 }
166 }
167
168 test "catalog spatial selection covers all three kernels" {
169 const allocator = std.testing.allocator;
170
171 var cells = (try selectSpatial(allocator, .{
172 .dtype = .f32,
173 .kind = .grid_cells,
174 .count = 5000,
175 })) orelse return error.TestExpectedSpatialDescriptor;
176 defer cells.deinit();
177 try std.testing.expectEqualStrings(
178 "accy.kernel.spatial.grid_cells_family_32_f32",
179 cells.descriptor.metadata.target,
180 );
181 try std.testing.expect(cells.descriptor.metadata.specialization.operationIs(.{ .spatial = .grid_cells }));
182
183 var count = (try selectSpatial(allocator, .{
184 .dtype = .i32,
185 .kind = .{ .grid_count = .{ .cells = 64 } },
186 .count = 5000,
187 .schedule = .{ .thread_blocks = 64 },
188 })) orelse return error.TestExpectedSpatialDescriptor;
189 defer count.deinit();
190 try std.testing.expectEqualStrings(
191 "accy.kernel.spatial.grid_count_family_64_64_i32",
192 count.descriptor.metadata.target,
193 );
194
195 var neighbor = (try selectSpatial(allocator, .{
196 .dtype = .f32,
197 .kind = .{ .grid_neighbor_count = .{ .cells = 64, .stride = 79 } },
198 .count = 5000,
199 .schedule = .{ .thread_blocks = 64 },
200 })) orelse return error.TestExpectedSpatialDescriptor;
201 defer neighbor.deinit();
202 try std.testing.expectEqualStrings(
203 "accy.kernel.spatial.grid_neighbor_count_family_64_f32",
204 neighbor.descriptor.metadata.target,
205 );
206
207 try std.testing.expectEqual(
208 @as(?OwnedDescriptor, null),
209 try selectSpatial(allocator, .{ .dtype = .i32, .kind = .grid_cells, .count = 5000 }),
210 );
211 try std.testing.expectEqual(
212 @as(?OwnedDescriptor, null),
213 try selectSpatial(allocator, .{ .dtype = .f32, .kind = .grid_cells, .count = 0 }),
214 );
215 try std.testing.expectEqual(@as(?OwnedDescriptor, null), try selectSpatial(allocator, .{
216 .dtype = .i32,
217 .kind = .{ .grid_count = .{ .cells = 8192 } },
218 .count = 5000,
219 }));
220
221 var candidates = try selectSpatialCandidates(allocator, .{
222 .dtype = .f32,
223 .kind = .grid_cells,
224 .count = 8192,
225 });
226 defer candidates.deinit();
227 try std.testing.expect(candidates.count >= 2);
228 for (candidates.slice()) |candidate| {
229 try std.testing.expect(candidate.descriptor.metadata.specialization.operationIs(.{ .spatial = .grid_cells }));
230 }
231 }