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

daab053ee43316e1809a84551d573ddd1e5bf3d2

  1 const std = @import("std");
  2 
  3 const descriptor_mod = @import("../root.zig");
  4 const entry = @import("../../root.zig").entry;
  5 const geometry = @import("../../root.zig").geometry;
  6 const image = @import("../../root.zig").image;
  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 ImageQuery = query_mod.ImageQuery;
 13 
 14 pub const ImageCandidateDescriptors = struct {
 15     count: usize = 0,
 16     items: [geometry.max_thread_candidates]OwnedDescriptor = undefined,
 17 
 18     pub fn slice(self: *const ImageCandidateDescriptors) []const OwnedDescriptor {
 19         return self.items[0..self.count];
 20     }
 21 
 22     pub fn deinit(self: *ImageCandidateDescriptors) void {
 23         for (self.items[0..self.count]) |*descriptor| descriptor.deinit();
 24         self.* = undefined;
 25     }
 26 };
 27 
 28 pub fn selectImage(backing_allocator: std.mem.Allocator, query: ImageQuery) !?OwnedDescriptor {
 29     if (!canonicalImage(query)) return null;
 30     const threads = imageQueryThreads(query) orelse return null;
 31     return try imageDescriptorForThreads(backing_allocator, query, threads);
 32 }
 33 
 34 pub fn selectImageCandidates(
 35     backing_allocator: std.mem.Allocator,
 36     query: ImageQuery,
 37 ) !ImageCandidateDescriptors {
 38     var result = ImageCandidateDescriptors{};
 39     errdefer result.deinit();
 40 
 41     if (!canonicalImage(query)) return result;
 42     if (query.schedule != null) {
 43         if (try selectImage(backing_allocator, query)) |descriptor| {
 44             result.items[result.count] = descriptor;
 45             result.count += 1;
 46         }
 47         return result;
 48     }
 49 
 50     const thread_candidates = image.imageThreadCandidatesForExtents(query.width, query.height);
 51     for (thread_candidates.slice()) |threads| {
 52         if (result.count >= result.items.len) break;
 53         if (try imageDescriptorForThreads(backing_allocator, query, threads)) |descriptor| {
 54             result.items[result.count] = descriptor;
 55             result.count += 1;
 56         }
 57     }
 58     return result;
 59 }
 60 
 61 fn canonicalImage(query: ImageQuery) bool {
 62     if (query.dtype != .u32) return false;
 63     if (!image.imageExtentValid(query.width) or !image.imageExtentValid(query.height)) return false;
 64     return switch (query.kind) {
 65         .blur_pass => |facts| image.blurPassInstanceValid(.{
 66             .radius = facts.radius,
 67             .axis = facts.axis,
 68             .width = query.width,
 69             .height = query.height,
 70         }),
 71         .resize_bilinear => |facts| image.resizeInstanceValid(.{
 72             .dst_width = query.width,
 73             .dst_height = query.height,
 74             .src_width = facts.src_width,
 75             .src_height = facts.src_height,
 76         }),
 77     };
 78 }
 79 
 80 fn imageQueryThreads(query: ImageQuery) ?entry.Threads2D {
 81     if (query.schedule) |requested| {
 82         switch (requested) {
 83             .thread_blocks => |threads| {
 84                 if (threads.x == 0 or threads.y == 0) return null;
 85                 return threads;
 86             },
 87         }
 88     }
 89     return image.imageThreadsForExtents(query.width, query.height);
 90 }
 91 
 92 fn imageDescriptorForThreads(
 93     backing_allocator: std.mem.Allocator,
 94     query: ImageQuery,
 95     threads: entry.Threads2D,
 96 ) !?OwnedDescriptor {
 97     switch (query.kind) {
 98         .blur_pass => |facts| {
 99             const instance = image.BlurPass{
100                 .radius = facts.radius,
101                 .axis = facts.axis,
102                 .width = query.width,
103                 .height = query.height,
104                 .threads = threads,
105             };
106             if (!image.blurPassInstanceValid(instance)) return null;
107             var specialization = try image.blurPassFamilySpecialization(backing_allocator, instance);
108             errdefer specialization.deinit();
109             const lifetime_allocator = specialization.allocator();
110             const descriptor = Descriptor{
111                 .name = try image.blurPassFamilyEntryName(lifetime_allocator, instance),
112                 .metadata = .{
113                     .target = try image.blurPassFamilyTarget(lifetime_allocator, instance),
114                     .version = image.blur_family_version,
115                     .layer = .logical,
116                     .category = .image,
117                     .specialization = specialization.value,
118                 },
119             };
120             if (!match_mod.imageDescriptorMatches(descriptor, query)) {
121                 specialization.deinit();
122                 return null;
123             }
124             return .{ .descriptor = descriptor, .specialization = specialization };
125         },
126         .resize_bilinear => |facts| {
127             const instance = image.Resize{
128                 .dst_width = query.width,
129                 .dst_height = query.height,
130                 .src_width = facts.src_width,
131                 .src_height = facts.src_height,
132                 .threads = threads,
133             };
134             if (!image.resizeInstanceValid(instance)) return null;
135             var specialization = try image.resizeFamilySpecialization(backing_allocator, instance);
136             errdefer specialization.deinit();
137             const lifetime_allocator = specialization.allocator();
138             const descriptor = Descriptor{
139                 .name = try image.resizeFamilyEntryName(lifetime_allocator, instance),
140                 .metadata = .{
141                     .target = try image.resizeFamilyTarget(lifetime_allocator, instance),
142                     .version = image.resize_family_version,
143                     .layer = .logical,
144                     .category = .image,
145                     .specialization = specialization.value,
146                 },
147             };
148             if (!match_mod.imageDescriptorMatches(descriptor, query)) {
149                 specialization.deinit();
150                 return null;
151             }
152             return .{ .descriptor = descriptor, .specialization = specialization };
153         },
154     }
155 }