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

daab053ee43316e1809a84551d573ddd1e5bf3d2

  1 const std = @import("std");
  2 const choir_abi = @import("choir_abi");
  3 
  4 const descriptor_mod = @import("../root.zig");
  5 const geometry = @import("../../root.zig").geometry;
  6 const stencil = @import("../../root.zig").stencil;
  7 const match_mod = @import("../match/root.zig");
  8 const query_mod = descriptor_mod;
  9 
 10 const Descriptor = descriptor_mod.Descriptor;
 11 const OwnedDescriptor = descriptor_mod.OwnedDescriptor;
 12 const StencilQuery = query_mod.StencilQuery;
 13 
 14 pub const StencilWindowCandidateDescriptors = struct {
 15     count: usize = 0,
 16     items: [geometry.max_thread_candidates]OwnedDescriptor = undefined,
 17 
 18     pub fn slice(self: *const StencilWindowCandidateDescriptors) []const OwnedDescriptor {
 19         return self.items[0..self.count];
 20     }
 21 
 22     pub fn deinit(self: *StencilWindowCandidateDescriptors) void {
 23         for (self.items[0..self.count]) |*descriptor| descriptor.deinit();
 24         self.* = undefined;
 25     }
 26 };
 27 
 28 pub fn selectStencilWindow(backing_allocator: std.mem.Allocator, query: StencilQuery) !?OwnedDescriptor {
 29     const accumulation_dtype = stencil.windowAccumulationDType(query.dtype) orelse return null;
 30     if (!canonicalStencilWindow(query)) return null;
 31     const instance = stencilWindowInstance(query, accumulation_dtype) orelse return null;
 32     return try stencilWindowDescriptorForInstance(backing_allocator, instance, query);
 33 }
 34 
 35 pub fn selectStencilWindowCandidates(
 36     backing_allocator: std.mem.Allocator,
 37     query: StencilQuery,
 38 ) !StencilWindowCandidateDescriptors {
 39     var result = StencilWindowCandidateDescriptors{};
 40     errdefer result.deinit();
 41 
 42     const accumulation_dtype = stencil.windowAccumulationDType(query.dtype) orelse return result;
 43     if (!canonicalStencilWindow(query)) return result;
 44     if (query.schedule != null) {
 45         if (try selectStencilWindow(backing_allocator, query)) |descriptor| {
 46             result.items[result.count] = descriptor;
 47             result.count += 1;
 48         }
 49         return result;
 50     }
 51 
 52     const thread_candidates = stencil.windowThreadCandidatesForExtents(query.rows, query.cols);
 53     for (thread_candidates.slice()) |threads| {
 54         const instance = stencil.Window{
 55             .rows = query.rows,
 56             .cols = query.cols,
 57             .radius = query.radius,
 58             .dtype = query.dtype,
 59             .accumulation_dtype = accumulation_dtype,
 60             .threads = threads,
 61         };
 62         if (try stencilWindowDescriptorForInstance(backing_allocator, instance, query)) |descriptor| {
 63             result.items[result.count] = descriptor;
 64             result.count += 1;
 65         }
 66     }
 67     return result;
 68 }
 69 
 70 fn canonicalStencilWindow(query: StencilQuery) bool {
 71     if (query.kind != .window) return false;
 72     if (query.rows == 0 or query.cols == 0) return false;
 73     return stencil.windowRadiusValid(query.radius);
 74 }
 75 
 76 fn stencilWindowDescriptorForInstance(
 77     backing_allocator: std.mem.Allocator,
 78     instance: stencil.Window,
 79     query: StencilQuery,
 80 ) !?OwnedDescriptor {
 81     var specialization = try stencil.windowFamilySpecialization(backing_allocator, instance);
 82     errdefer specialization.deinit();
 83     const lifetime_allocator = specialization.allocator();
 84     const descriptor = Descriptor{
 85         .name = try stencil.windowFamilyEntryName(lifetime_allocator, instance),
 86         .metadata = .{
 87             .target = try stencil.windowFamilyTarget(lifetime_allocator, instance),
 88             .version = stencil.window_family_version,
 89             .layer = .logical,
 90             .category = .stencil,
 91             .specialization = specialization.value,
 92         },
 93     };
 94     if (!match_mod.stencilWindowDescriptorMatches(descriptor, query)) {
 95         specialization.deinit();
 96         return null;
 97     }
 98     return .{ .descriptor = descriptor, .specialization = specialization };
 99 }
100 
101 fn stencilWindowInstance(query: StencilQuery, accumulation_dtype: choir_abi.DType) ?stencil.Window {
102     var instance = stencil.Window{
103         .rows = query.rows,
104         .cols = query.cols,
105         .radius = query.radius,
106         .dtype = query.dtype,
107         .accumulation_dtype = accumulation_dtype,
108         .threads = stencil.windowThreadsForExtents(query.rows, query.cols),
109     };
110     if (query.schedule) |requested| {
111         switch (requested) {
112             .thread_blocks => |threads| {
113                 if (threads.x == 0 or threads.y == 0) return null;
114                 instance.threads = threads;
115             },
116         }
117     }
118     return instance;
119 }