lib/accy/src/kernel/library/catalog/family/linalg/matrix.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 linalg = @import("../../../root.zig").linalg;
  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 MatrixProductQuery = query_mod.MatrixProductQuery;
 13 const MatrixProductSchedule = query_mod.MatrixProductSchedule;
 14 
 15 pub const MatrixProductCandidateDescriptors = struct {
 16     count: usize = 0,
 17     items: [geometry.max_thread_candidates]OwnedDescriptor = undefined,
 18 
 19     pub fn slice(self: *const MatrixProductCandidateDescriptors) []const OwnedDescriptor {
 20         return self.items[0..self.count];
 21     }
 22 
 23     pub fn deinit(self: *MatrixProductCandidateDescriptors) void {
 24         for (self.items[0..self.count]) |*descriptor| descriptor.deinit();
 25         self.* = undefined;
 26     }
 27 };
 28 
 29 pub fn selectMatrixProduct(backing_allocator: std.mem.Allocator, query: MatrixProductQuery) !?OwnedDescriptor {
 30     const accumulation_dtype = linalg.matrixProductAccumulationDType(query.dtype) orelse return null;
 31     if (!match_mod.canonicalMatrixProduct(query)) return null;
 32     const shape = match_mod.matrixProductShape(query.lhs_dims, query.rhs_dims, query.output_dims) orelse return null;
 33     const instance = matrixProductFamilyInstance(shape, query.schedule, query.dtype, accumulation_dtype) orelse return null;
 34     return try matrixProductDescriptorForInstance(backing_allocator, instance, shape, query.schedule, query.dtype);
 35 }
 36 
 37 pub fn selectMatrixProductCandidates(
 38     backing_allocator: std.mem.Allocator,
 39     query: MatrixProductQuery,
 40 ) !MatrixProductCandidateDescriptors {
 41     var result = MatrixProductCandidateDescriptors{};
 42     errdefer result.deinit();
 43 
 44     const accumulation_dtype = linalg.matrixProductAccumulationDType(query.dtype) orelse return result;
 45     if (!match_mod.canonicalMatrixProduct(query)) return result;
 46     const shape = match_mod.matrixProductShape(query.lhs_dims, query.rhs_dims, query.output_dims) orelse return result;
 47     if (query.schedule != null) {
 48         if (try selectMatrixProduct(backing_allocator, query)) |descriptor| {
 49             result.items[result.count] = descriptor;
 50             result.count += 1;
 51         }
 52         return result;
 53     }
 54 
 55     const thread_candidates = linalg.matrixProductThreadCandidatesForExtents(shape.m, shape.n);
 56     for (thread_candidates.slice()) |threads| {
 57         const instance = linalg.MatrixProduct{
 58             .m = shape.m,
 59             .n = shape.n,
 60             .k = shape.k,
 61             .dtype = query.dtype,
 62             .accumulation_dtype = accumulation_dtype,
 63             .threads = threads,
 64         };
 65         if (try matrixProductDescriptorForInstance(backing_allocator, instance, shape, null, query.dtype)) |descriptor| {
 66             result.items[result.count] = descriptor;
 67             result.count += 1;
 68         }
 69     }
 70     return result;
 71 }
 72 
 73 fn matrixProductDescriptorForInstance(
 74     backing_allocator: std.mem.Allocator,
 75     instance: linalg.MatrixProduct,
 76     shape: match_mod.MatrixProductShape,
 77     schedule: ?MatrixProductSchedule,
 78     dtype: choir_abi.DType,
 79 ) !?OwnedDescriptor {
 80     var specialization = try linalg.matrixProductFamilySpecialization(backing_allocator, instance);
 81     errdefer specialization.deinit();
 82     const lifetime_allocator = specialization.allocator();
 83     const descriptor = Descriptor{
 84         .name = try linalg.matrixProductFamilyEntryName(lifetime_allocator, instance),
 85         .metadata = .{
 86             .target = try linalg.matrixProductFamilyTarget(lifetime_allocator, instance),
 87             .version = linalg.matrix_product_family_version,
 88             .layer = .logical,
 89             .category = .linalg,
 90             .specialization = specialization.value,
 91         },
 92     };
 93     if (!match_mod.matrixProductDescriptorMatches(descriptor, shape, schedule, dtype)) {
 94         specialization.deinit();
 95         return null;
 96     }
 97     return .{ .descriptor = descriptor, .specialization = specialization };
 98 }
 99 
100 fn matrixProductFamilyInstance(
101     shape: match_mod.MatrixProductShape,
102     schedule: ?MatrixProductSchedule,
103     dtype: choir_abi.DType,
104     accumulation_dtype: choir_abi.DType,
105 ) ?linalg.MatrixProduct {
106     var instance = linalg.MatrixProduct{
107         .m = shape.m,
108         .n = shape.n,
109         .k = shape.k,
110         .dtype = dtype,
111         .accumulation_dtype = accumulation_dtype,
112         .threads = linalg.matrixProductThreadsForExtents(shape.m, shape.n),
113     };
114     if (schedule) |requested| {
115         switch (requested) {
116             .thread_blocks => |threads| {
117                 if (threads.x == 0 or threads.y == 0) return null;
118                 instance.threads = threads;
119             },
120         }
121     }
122     return instance;
123 }