lib/accy/src/kernel/library/catalog/family/histogram.zig
daab053ee43316e1809a84551d573ddd1e5bf3d2
1 const std = @import("std");
2
3 const descriptor_mod = @import("../root.zig");
4 const geometry = @import("../../root.zig").geometry;
5 const histogram_mod = @import("../../root.zig").histogram;
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 HistogramQuery = query_mod.HistogramQuery;
12
13 pub const HistogramCandidateDescriptors = struct {
14 count: usize = 0,
15 items: [2 * geometry.max_thread_candidates]OwnedDescriptor = undefined,
16
17 pub fn slice(self: *const HistogramCandidateDescriptors) []const OwnedDescriptor {
18 return self.items[0..self.count];
19 }
20
21 pub fn deinit(self: *HistogramCandidateDescriptors) void {
22 for (self.items[0..self.count]) |*descriptor| descriptor.deinit();
23 self.* = undefined;
24 }
25 };
26
27 pub fn selectHistogram(backing_allocator: std.mem.Allocator, query: HistogramQuery) !?OwnedDescriptor {
28 if (!canonicalHistogram(query)) return null;
29 const instance = histogramFamilyInstance(query) orelse return null;
30 return try histogramDescriptorForInstance(backing_allocator, instance, query);
31 }
32
33 pub fn selectHistogramCandidates(
34 backing_allocator: std.mem.Allocator,
35 query: HistogramQuery,
36 ) !HistogramCandidateDescriptors {
37 var result = HistogramCandidateDescriptors{};
38 errdefer result.deinit();
39
40 if (!canonicalHistogram(query)) return result;
41 if (query.schedule != null) {
42 if (try selectHistogram(backing_allocator, query)) |descriptor| {
43 result.items[result.count] = descriptor;
44 result.count += 1;
45 }
46 return result;
47 }
48
49 const thread_candidates = histogram_mod.histogramThreadCandidatesForCount(query.count);
50 const variants = [_]histogram_mod.HistogramVariant{ .direct, .shared_bins };
51 for (variants) |variant| {
52 for (thread_candidates.slice()) |threads| {
53 const instance = histogram_mod.Histogram{
54 .bins = query.bins,
55 .count = query.count,
56 .dtype = query.dtype,
57 .binning = query.binning,
58 .variant = variant,
59 .threads = threads,
60 };
61 if (!histogram_mod.histogramInstanceValid(instance)) continue;
62 if (result.count >= result.items.len) break;
63 if (try histogramDescriptorForInstance(backing_allocator, instance, query)) |descriptor| {
64 result.items[result.count] = descriptor;
65 result.count += 1;
66 }
67 }
68 }
69 return result;
70 }
71
72 fn canonicalHistogram(query: HistogramQuery) bool {
73 if (!histogram_mod.histogramDTypeSupported(query.dtype)) return false;
74 return query.bins != 0 and query.count != 0;
75 }
76
77 fn histogramDescriptorForInstance(
78 backing_allocator: std.mem.Allocator,
79 instance: histogram_mod.Histogram,
80 query: HistogramQuery,
81 ) !?OwnedDescriptor {
82 var specialization = try histogram_mod.histogramFamilySpecialization(backing_allocator, instance);
83 errdefer specialization.deinit();
84 const lifetime_allocator = specialization.allocator();
85 const descriptor = Descriptor{
86 .name = try histogram_mod.histogramFamilyEntryName(lifetime_allocator, instance),
87 .metadata = .{
88 .target = try histogram_mod.histogramFamilyTarget(lifetime_allocator, instance),
89 .version = histogram_mod.histogram_family_version,
90 .layer = .logical,
91 .category = .indexing,
92 .specialization = specialization.value,
93 },
94 };
95 if (!match_mod.histogramDescriptorMatches(descriptor, query)) {
96 specialization.deinit();
97 return null;
98 }
99 return .{ .descriptor = descriptor, .specialization = specialization };
100 }
101
102 fn histogramFamilyInstance(query: HistogramQuery) ?histogram_mod.Histogram {
103 var instance = histogram_mod.Histogram{
104 .bins = query.bins,
105 .count = query.count,
106 .dtype = query.dtype,
107 .binning = query.binning,
108 .threads = histogram_mod.histogramThreadsForCount(query.count),
109 };
110 if (query.schedule) |requested| {
111 switch (requested) {
112 .thread_blocks => |threads| {
113 if (threads == 0) return null;
114 if (@as(u64, threads) > query.count) return null;
115 instance.threads = threads;
116 },
117 .shared_bins => |threads| {
118 if (threads == 0) return null;
119 if (@as(u64, threads) > query.count) return null;
120 instance.variant = .shared_bins;
121 instance.threads = threads;
122 },
123 }
124 }
125 if (!histogram_mod.histogramInstanceValid(instance)) return null;
126 return instance;
127 }
128
129 test "catalog histogram candidates enumerate both schedule structures" {
130 const allocator = std.testing.allocator;
131
132 var candidates = try selectHistogramCandidates(allocator, .{
133 .dtype = .f32,
134 .bins = 128,
135 .count = 8192,
136 });
137 defer candidates.deinit();
138
139 var direct_count: usize = 0;
140 var shared_count: usize = 0;
141 for (candidates.slice()) |candidate| {
142 try std.testing.expect(candidate.descriptor.metadata.specialization.operationIs(.{ .indexing = .histogram }));
143 if (candidate.descriptor.metadata.specialization.structureIs("shared_bins")) {
144 shared_count += 1;
145 } else {
146 direct_count += 1;
147 }
148 try std.testing.expect(candidate.descriptor.metadata.specialization.staticParameterMatches(
149 histogram_mod.histogram_binning_policy_parameter,
150 histogram_mod.histogramBinningPolicyCode(.lower_inclusive_upper_exclusive),
151 ));
152 }
153 try std.testing.expect(direct_count >= 2);
154 try std.testing.expectEqual(direct_count, shared_count);
155
156 var pinned = (try selectHistogram(allocator, .{
157 .dtype = .f32,
158 .bins = 128,
159 .count = 8192,
160 .schedule = .{ .shared_bins = 128 },
161 })) orelse return error.TestExpectedHistogramDescriptor;
162 defer pinned.deinit();
163 try std.testing.expectEqualStrings(
164 "accy.kernel.histogram.histogram_family_shared128_128_f32",
165 pinned.descriptor.metadata.target,
166 );
167
168 try std.testing.expectEqual(@as(?OwnedDescriptor, null), try selectHistogram(allocator, .{
169 .dtype = .i32,
170 .bins = 128,
171 .count = 8192,
172 }));
173 }