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

daab053ee43316e1809a84551d573ddd1e5bf3d2

  1 const std = @import("std");
  2 
  3 const descriptor_mod = @import("../root.zig");
  4 const geometry = @import("../../root.zig").geometry;
  5 const segmented = @import("../../root.zig").segmented;
  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 SegmentedQuery = query_mod.SegmentedQuery;
 12 
 13 pub const SegmentSumCandidateDescriptors = struct {
 14     count: usize = 0,
 15     items: [2 * geometry.max_thread_candidates]OwnedDescriptor = undefined,
 16 
 17     pub fn slice(self: *const SegmentSumCandidateDescriptors) []const OwnedDescriptor {
 18         return self.items[0..self.count];
 19     }
 20 
 21     pub fn deinit(self: *SegmentSumCandidateDescriptors) void {
 22         for (self.items[0..self.count]) |*descriptor| descriptor.deinit();
 23         self.* = undefined;
 24     }
 25 };
 26 
 27 pub fn selectSegmentSum(backing_allocator: std.mem.Allocator, query: SegmentedQuery) !?OwnedDescriptor {
 28     if (!canonicalSegmentSum(query)) return null;
 29     const instance = segmentSumFamilyInstance(query) orelse return null;
 30     return try segmentSumDescriptorForInstance(backing_allocator, instance, query);
 31 }
 32 
 33 pub fn selectSegmentSumCandidates(
 34     backing_allocator: std.mem.Allocator,
 35     query: SegmentedQuery,
 36 ) !SegmentSumCandidateDescriptors {
 37     var result = SegmentSumCandidateDescriptors{};
 38     errdefer result.deinit();
 39 
 40     if (!canonicalSegmentSum(query)) return result;
 41     if (query.schedule != null) {
 42         if (try selectSegmentSum(backing_allocator, query)) |descriptor| {
 43             result.items[result.count] = descriptor;
 44             result.count += 1;
 45         }
 46         return result;
 47     }
 48 
 49     const thread_candidates = segmented.segmentSumThreadCandidatesForSegments(query.segments);
 50     for (thread_candidates.slice()) |threads| {
 51         const instance = segmented.SegmentSum{
 52             .segments = query.segments,
 53             .total = query.total,
 54             .dtype = query.dtype,
 55             .threads = threads,
 56         };
 57         if (try segmentSumDescriptorForInstance(backing_allocator, instance, query)) |descriptor| {
 58             result.items[result.count] = descriptor;
 59             result.count += 1;
 60         }
 61     }
 62 
 63     const warp_extent = query.segments * segmented.segment_sum_warp_size;
 64     const warp_candidates = segmented.segmentSumThreadCandidatesForSegments(warp_extent);
 65     for (warp_candidates.slice()) |threads| {
 66         if (!segmented.segmentSumGranularityValid(.warp, threads)) continue;
 67         const instance = segmented.SegmentSum{
 68             .segments = query.segments,
 69             .total = query.total,
 70             .dtype = query.dtype,
 71             .granularity = .warp,
 72             .threads = threads,
 73         };
 74         if (try segmentSumDescriptorForInstance(backing_allocator, instance, query)) |descriptor| {
 75             result.items[result.count] = descriptor;
 76             result.count += 1;
 77         }
 78     }
 79     return result;
 80 }
 81 
 82 fn canonicalSegmentSum(query: SegmentedQuery) bool {
 83     if (query.kind != .segment_sum) return false;
 84     if (!segmented.segmentSumDTypeSupported(query.dtype)) return false;
 85     return query.segments != 0 and query.total != 0;
 86 }
 87 
 88 fn segmentSumDescriptorForInstance(
 89     backing_allocator: std.mem.Allocator,
 90     instance: segmented.SegmentSum,
 91     query: SegmentedQuery,
 92 ) !?OwnedDescriptor {
 93     var specialization = try segmented.segmentSumFamilySpecialization(backing_allocator, instance);
 94     errdefer specialization.deinit();
 95     const lifetime_allocator = specialization.allocator();
 96     const descriptor = Descriptor{
 97         .name = try segmented.segmentSumFamilyEntryName(lifetime_allocator, instance),
 98         .metadata = .{
 99             .target = try segmented.segmentSumFamilyTarget(lifetime_allocator, instance),
100             .version = segmented.segment_sum_family_version,
101             .layer = .logical,
102             .category = .segmented,
103             .specialization = specialization.value,
104         },
105     };
106     if (!match_mod.segmentSumDescriptorMatches(descriptor, query)) {
107         specialization.deinit();
108         return null;
109     }
110     return .{ .descriptor = descriptor, .specialization = specialization };
111 }
112 
113 fn segmentSumFamilyInstance(query: SegmentedQuery) ?segmented.SegmentSum {
114     var instance = segmented.SegmentSum{
115         .segments = query.segments,
116         .total = query.total,
117         .dtype = query.dtype,
118         .threads = segmented.segmentSumThreadsForSegments(query.segments),
119     };
120     if (query.schedule) |requested| {
121         switch (requested) {
122             .thread_blocks => |threads| {
123                 if (threads == 0) return null;
124                 if (@as(u64, threads) > query.segments) return null;
125                 instance.threads = threads;
126             },
127             .warp_blocks => |threads| {
128                 if (!segmented.segmentSumGranularityValid(.warp, threads)) return null;
129                 if (@as(u64, threads) > query.segments * segmented.segment_sum_warp_size) return null;
130                 instance.granularity = .warp;
131                 instance.threads = threads;
132             },
133         }
134     }
135     return instance;
136 }