lib/accy/src/kernel/library/catalog/match/segmented.zig

daab053ee43316e1809a84551d573ddd1e5bf3d2

 1 const catalog = @import("../root.zig");
 2 const library = @import("../../root.zig");
 3 
 4 const common = @import("common.zig");
 5 
 6 const entry = library.entry;
 7 const segmented = library.segmented;
 8 const Descriptor = catalog.Descriptor;
 9 const SegmentedQuery = catalog.SegmentedQuery;
10 
11 const selectableSpecialization = common.selectableSpecialization;
12 
13 pub fn segmentSumDescriptorMatches(descriptor: Descriptor, query: SegmentedQuery) bool {
14     const metadata = descriptor.metadata;
15     if (metadata.category != .segmented) return false;
16     const specialization = metadata.specialization;
17     if (!selectableSpecialization(specialization)) return false;
18     if (!specialization.scheduleMatchesLaunch()) return false;
19     if (!specialization.operationIs(.{ .segmented = .segment_sum })) return false;
20     if (specialization.dtype != query.dtype) return false;
21     if (specialization.inputs.len != 2 or specialization.outputs.len != 1) return false;
22     if (specialization.reductions.len != 0) return false;
23     return specialization.inputHasExtents(0, &.{query.total}) and
24         specialization.inputHasExtents(1, &.{query.segments + 1}) and
25         specialization.outputHasExtents(0, &.{query.segments}) and
26         segmentedScheduleMatches(specialization, query);
27 }
28 fn segmentedScheduleMatches(specialization: entry.Specialization, query: SegmentedQuery) bool {
29     const requested = query.schedule orelse return true;
30     const launch = specialization.launch orelse return false;
31     const granularity = segmented.segmentSumGranularityFromLaunch(launch, query.segments) orelse return false;
32     return switch (requested) {
33         .thread_blocks => |threads| granularity == .thread and launch.threadgroup[0] == threads,
34         .warp_blocks => |threads| granularity == .warp and launch.threadgroup[0] == threads,
35     };
36 }