lib/accy/src/kernel/library/catalog/family/linalg/batch.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 BatchedMatrixProductQuery = query_mod.BatchedMatrixProductQuery;
11 const BatchedMatrixProductSchedule = query_mod.BatchedMatrixProductSchedule;
12 
13 pub fn selectBatchedMatrixProduct(backing_allocator: std.mem.Allocator, query: BatchedMatrixProductQuery) !?OwnedDescriptor {
14     if (query.dtype != .f32) return null;
15     if (!match_mod.canonicalBatchedMatrixProduct(query)) return null;
16     const shape = match_mod.batchedMatrixProductShape(query.lhs_dims, query.rhs_dims, query.output_dims) orelse return null;
17     const instance = batchedMatrixProductFamilyInstance(shape, query.schedule) orelse return null;
18     return try batchedMatrixProductDescriptorForInstance(backing_allocator, instance, shape, query.schedule);
19 }
20 
21 fn batchedMatrixProductDescriptorForInstance(
22     backing_allocator: std.mem.Allocator,
23     instance: linalg.BatchedMatrixProduct,
24     shape: match_mod.BatchedMatrixProductShape,
25     schedule: ?BatchedMatrixProductSchedule,
26 ) !?OwnedDescriptor {
27     var specialization = try linalg.batchedMatrixProductFamilySpecialization(backing_allocator, instance);
28     errdefer specialization.deinit();
29     const lifetime_allocator = specialization.allocator();
30     const descriptor = Descriptor{
31         .name = try linalg.batchedMatrixProductFamilyEntryName(lifetime_allocator, instance),
32         .metadata = .{
33             .target = try linalg.batchedMatrixProductFamilyTarget(lifetime_allocator, instance),
34             .version = linalg.batched_matrix_product_family_version,
35             .layer = .logical,
36             .category = .linalg,
37             .specialization = specialization.value,
38         },
39     };
40     if (!match_mod.batchedMatrixProductDescriptorMatches(descriptor, shape, schedule)) {
41         specialization.deinit();
42         return null;
43     }
44     return .{ .descriptor = descriptor, .specialization = specialization };
45 }
46 
47 fn batchedMatrixProductFamilyInstance(
48     shape: match_mod.BatchedMatrixProductShape,
49     schedule: ?BatchedMatrixProductSchedule,
50 ) ?linalg.BatchedMatrixProduct {
51     var instance = linalg.BatchedMatrixProduct{
52         .batch = shape.batch,
53         .m = shape.m,
54         .n = shape.n,
55         .k = shape.k,
56         .threads = linalg.batchedMatrixProductThreadsForExtents(shape.batch, shape.m, shape.n),
57     };
58     if (schedule) |requested| {
59         switch (requested) {
60             .thread_blocks => |threads| {
61                 if (threads.x == 0 or threads.y == 0 or threads.z == 0) return null;
62                 instance.threads = threads;
63             },
64         }
65     }
66     return instance;
67 }