lib/accy/src/kernel/library/catalog/family/scatter/scatter.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 ScatterQuery = query_mod.ScatterQuery;
 12 
 13 pub const ScatterCandidateDescriptors = struct {
 14     count: usize = 0,
 15     items: [geometry.max_thread_candidates]OwnedDescriptor = undefined,
 16 
 17     pub fn slice(self: *const ScatterCandidateDescriptors) []const OwnedDescriptor {
 18         return self.items[0..self.count];
 19     }
 20 
 21     pub fn deinit(self: *ScatterCandidateDescriptors) void {
 22         for (self.items[0..self.count]) |*descriptor| descriptor.deinit();
 23         self.* = undefined;
 24     }
 25 };
 26 
 27 pub fn selectScatter(backing_allocator: std.mem.Allocator, query: ScatterQuery) !?OwnedDescriptor {
 28     if (!canonicalScatter(query)) return null;
 29     const instance = scatterFamilyInstance(query) orelse return null;
 30     return try scatterDescriptorForInstance(backing_allocator, instance, query);
 31 }
 32 
 33 pub fn selectScatterCandidates(
 34     backing_allocator: std.mem.Allocator,
 35     query: ScatterQuery,
 36 ) !ScatterCandidateDescriptors {
 37     var result = ScatterCandidateDescriptors{};
 38     errdefer result.deinit();
 39 
 40     if (!canonicalScatter(query)) return result;
 41     if (query.schedule != null) {
 42         if (try selectScatter(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.axis_size * query.inner;
 50     const thread_candidates = indexing.scatterThreadCandidatesForTotal(total);
 51     for (thread_candidates.slice()) |threads| {
 52         const instance = indexing.Scatter{
 53             .outer = query.outer,
 54             .axis_size = query.axis_size,
 55             .updates = query.updates,
 56             .inner = query.inner,
 57             .dtype = query.dtype,
 58             .threads = threads,
 59         };
 60         if (try scatterDescriptorForInstance(backing_allocator, instance, query)) |descriptor| {
 61             result.items[result.count] = descriptor;
 62             result.count += 1;
 63         }
 64     }
 65     return result;
 66 }
 67 
 68 fn canonicalScatter(query: ScatterQuery) bool {
 69     if (!indexing.scatterDTypeSupported(query.dtype)) return false;
 70     return query.outer != 0 and query.axis_size != 0 and query.updates != 0 and query.inner != 0;
 71 }
 72 
 73 fn scatterDescriptorForInstance(
 74     backing_allocator: std.mem.Allocator,
 75     instance: indexing.Scatter,
 76     query: ScatterQuery,
 77 ) !?OwnedDescriptor {
 78     var specialization = try indexing.scatterFamilySpecialization(backing_allocator, instance);
 79     errdefer specialization.deinit();
 80     const lifetime_allocator = specialization.allocator();
 81     const descriptor = Descriptor{
 82         .name = try indexing.scatterFamilyEntryName(lifetime_allocator, instance),
 83         .metadata = .{
 84             .target = try indexing.scatterFamilyTarget(lifetime_allocator, instance),
 85             .version = indexing.scatter_family_version,
 86             .layer = .logical,
 87             .category = .indexing,
 88             .specialization = specialization.value,
 89         },
 90     };
 91     if (!match_mod.scatterDescriptorMatches(descriptor, query)) {
 92         specialization.deinit();
 93         return null;
 94     }
 95     return .{ .descriptor = descriptor, .specialization = specialization };
 96 }
 97 
 98 fn scatterFamilyInstance(query: ScatterQuery) ?indexing.Scatter {
 99     const total = query.outer * query.axis_size * query.inner;
100     var instance = indexing.Scatter{
101         .outer = query.outer,
102         .axis_size = query.axis_size,
103         .updates = query.updates,
104         .inner = query.inner,
105         .dtype = query.dtype,
106         .threads = indexing.scatterThreadsForTotal(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 }