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

daab053ee43316e1809a84551d573ddd1e5bf3d2

  1 const std = @import("std");
  2 
  3 const descriptor_mod = @import("../root.zig");
  4 const geometry = @import("../../root.zig").geometry;
  5 const indexing = @import("../../root.zig").indexing;
  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 GatherQuery = query_mod.GatherQuery;
 12 
 13 pub const GatherCandidateDescriptors = struct {
 14     count: usize = 0,
 15     items: [geometry.max_thread_candidates]OwnedDescriptor = undefined,
 16 
 17     pub fn slice(self: *const GatherCandidateDescriptors) []const OwnedDescriptor {
 18         return self.items[0..self.count];
 19     }
 20 
 21     pub fn deinit(self: *GatherCandidateDescriptors) void {
 22         for (self.items[0..self.count]) |*descriptor| descriptor.deinit();
 23         self.* = undefined;
 24     }
 25 };
 26 
 27 pub fn selectGather(backing_allocator: std.mem.Allocator, query: GatherQuery) !?OwnedDescriptor {
 28     if (!canonicalGather(query)) return null;
 29     const instance = gatherFamilyInstance(query) orelse return null;
 30     return try gatherDescriptorForInstance(backing_allocator, instance, query);
 31 }
 32 
 33 pub fn selectGatherCandidates(
 34     backing_allocator: std.mem.Allocator,
 35     query: GatherQuery,
 36 ) !GatherCandidateDescriptors {
 37     var result = GatherCandidateDescriptors{};
 38     errdefer result.deinit();
 39 
 40     if (!canonicalGather(query)) return result;
 41     if (query.schedule != null) {
 42         if (try selectGather(backing_allocator, query)) |descriptor| {
 43             result.items[result.count] = descriptor;
 44             result.count += 1;
 45         }
 46         return result;
 47     }
 48 
 49     const total = query.outer * query.gathered * query.inner;
 50     const thread_candidates = indexing.gatherThreadCandidatesForTotal(total);
 51     for (thread_candidates.slice()) |threads| {
 52         const instance = indexing.Gather{
 53             .outer = query.outer,
 54             .axis_size = query.axis_size,
 55             .gathered = query.gathered,
 56             .inner = query.inner,
 57             .dtype = query.dtype,
 58             .threads = threads,
 59         };
 60         if (try gatherDescriptorForInstance(backing_allocator, instance, query)) |descriptor| {
 61             result.items[result.count] = descriptor;
 62             result.count += 1;
 63         }
 64     }
 65     return result;
 66 }
 67 
 68 fn canonicalGather(query: GatherQuery) bool {
 69     if (!indexing.gatherDTypeSupported(query.dtype)) return false;
 70     return query.outer != 0 and query.axis_size != 0 and query.gathered != 0 and query.inner != 0;
 71 }
 72 
 73 fn gatherDescriptorForInstance(
 74     backing_allocator: std.mem.Allocator,
 75     instance: indexing.Gather,
 76     query: GatherQuery,
 77 ) !?OwnedDescriptor {
 78     var specialization = try indexing.gatherFamilySpecialization(backing_allocator, instance);
 79     errdefer specialization.deinit();
 80     const lifetime_allocator = specialization.allocator();
 81     const descriptor = Descriptor{
 82         .name = try indexing.gatherFamilyEntryName(lifetime_allocator, instance),
 83         .metadata = .{
 84             .target = try indexing.gatherFamilyTarget(lifetime_allocator, instance),
 85             .version = indexing.gather_family_version,
 86             .layer = .logical,
 87             .category = .indexing,
 88             .specialization = specialization.value,
 89         },
 90     };
 91     if (!match_mod.gatherDescriptorMatches(descriptor, query)) {
 92         specialization.deinit();
 93         return null;
 94     }
 95     return .{ .descriptor = descriptor, .specialization = specialization };
 96 }
 97 
 98 fn gatherFamilyInstance(query: GatherQuery) ?indexing.Gather {
 99     const total = query.outer * query.gathered * query.inner;
100     var instance = indexing.Gather{
101         .outer = query.outer,
102         .axis_size = query.axis_size,
103         .gathered = query.gathered,
104         .inner = query.inner,
105         .dtype = query.dtype,
106         .threads = indexing.gatherThreadsForTotal(total),
107     };
108     if (query.schedule) |requested| {
109         switch (requested) {
110             .thread_blocks => |threads| {
111                 if (threads == 0) return null;
112                 if (@as(u64, threads) > total) return null;
113                 instance.threads = threads;
114             },
115         }
116     }
117     return instance;
118 }