lib/accy/src/kernel/library/catalog/family/linalg/outer.zig

daab053ee43316e1809a84551d573ddd1e5bf3d2

 1 const std = @import("std");
 2 
 3 const descriptor_mod = @import("../../root.zig");
 4 const linalg = @import("../../../root.zig").linalg;
 5 const match_mod = @import("../../match/root.zig");
 6 const query_mod = descriptor_mod;
 7 
 8 const Descriptor = descriptor_mod.Descriptor;
 9 const OwnedDescriptor = descriptor_mod.OwnedDescriptor;
10 const OuterProductQuery = query_mod.OuterProductQuery;
11 const OuterProductSchedule = query_mod.OuterProductSchedule;
12 
13 pub fn selectOuterProduct(backing_allocator: std.mem.Allocator, query: OuterProductQuery) !?OwnedDescriptor {
14     if (query.dtype != .f32) return null;
15     if (!match_mod.canonicalOuterProduct(query)) return null;
16     const shape = match_mod.outerProductShape(query.lhs_dims, query.rhs_dims, query.output_dims) orelse return null;
17     const instance = outerProductFamilyInstance(shape, query.schedule) orelse return null;
18     return try outerProductDescriptorForInstance(backing_allocator, instance, shape, query.schedule);
19 }
20 
21 fn outerProductDescriptorForInstance(
22     backing_allocator: std.mem.Allocator,
23     instance: linalg.OuterProduct,
24     shape: match_mod.OuterProductShape,
25     schedule: ?OuterProductSchedule,
26 ) !?OwnedDescriptor {
27     var specialization = try linalg.outerProductFamilySpecialization(backing_allocator, instance);
28     errdefer specialization.deinit();
29     const lifetime_allocator = specialization.allocator();
30     const descriptor = Descriptor{
31         .name = try linalg.outerProductFamilyEntryName(lifetime_allocator, instance),
32         .metadata = .{
33             .target = try linalg.outerProductFamilyTarget(lifetime_allocator, instance),
34             .version = linalg.outer_product_family_version,
35             .layer = .logical,
36             .category = .linalg,
37             .specialization = specialization.value,
38         },
39     };
40     if (!match_mod.outerProductDescriptorMatches(descriptor, shape, schedule)) {
41         specialization.deinit();
42         return null;
43     }
44     return .{ .descriptor = descriptor, .specialization = specialization };
45 }
46 
47 fn outerProductFamilyInstance(
48     shape: match_mod.OuterProductShape,
49     schedule: ?OuterProductSchedule,
50 ) ?linalg.OuterProduct {
51     var instance = linalg.OuterProduct{
52         .m = shape.m,
53         .n = shape.n,
54         .threads = linalg.outerProductThreadsForExtents(shape.m, shape.n),
55     };
56     if (schedule) |requested| {
57         switch (requested) {
58             .thread_blocks => |threads| {
59                 if (threads.x == 0 or threads.y == 0) return null;
60                 instance.threads = threads;
61             },
62         }
63     }
64     return instance;
65 }