lib/accy/src/kernel/library/catalog/match/linalg/matrix.zig

daab053ee43316e1809a84551d573ddd1e5bf3d2

 1 const std = @import("std");
 2 const choir_abi = @import("choir_abi");
 3 
 4 const catalog = @import("../../root.zig");
 5 const library = @import("../../../root.zig");
 6 
 7 const common = @import("../root.zig");
 8 const shapes = @import("../shape/root.zig");
 9 
10 const entry = library.entry;
11 const linalg = library.linalg;
12 const Descriptor = catalog.Descriptor;
13 const MatrixProductSchedule = catalog.MatrixProductSchedule;
14 const MatrixProductShape = shapes.MatrixProductShape;
15 const selectableSpecialization = common.selectableSpecialization;
16 const specializationThreadgroup2DMatches = common.specializationThreadgroup2DMatches;
17 
18 pub fn matrixProductDescriptorMatches(
19     descriptor: Descriptor,
20     shape: MatrixProductShape,
21     schedule: ?MatrixProductSchedule,
22     dtype: choir_abi.DType,
23 ) bool {
24     const metadata = descriptor.metadata;
25     if (metadata.category != .linalg) return false;
26     const specialization = metadata.specialization;
27     if (!selectableSpecialization(specialization)) return false;
28     if (!specialization.scheduleMatchesLaunch()) return false;
29     if (!specialization.operationIs(.{ .linalg = .matrix_product })) return false;
30     if (specialization.dtype != dtype) return false;
31     if (specialization.accumulation_dtype != linalg.matrixProductAccumulationDType(dtype)) return false;
32     const equation = specialization.equation orelse return false;
33     if (!std.mem.eql(u8, equation, "mk,kn->mn")) return false;
34     if (specialization.inputs.len != 2 or specialization.outputs.len != 1) return false;
35     return matrixProductSpecializationMatches(specialization, shape) and
36         matrixProductScheduleMatches(specialization, schedule);
37 }
38 
39 pub fn matrixProductScheduleMatches(specialization: entry.Specialization, schedule: ?MatrixProductSchedule) bool {
40     const requested = schedule orelse return true;
41     return switch (requested) {
42         .thread_blocks => |threads| specializationThreadgroup2DMatches(specialization, threads),
43     };
44 }
45 
46 pub fn matrixProductSpecializationMatches(specialization: entry.Specialization, shape: MatrixProductShape) bool {
47     if (specialization.outputs.len != 1) return false;
48     if (specialization.reductions.len != 1) return false;
49     return specialization.inputHasExtents(0, &.{ shape.m, shape.k }) and
50         specialization.inputHasExtents(1, &.{ shape.k, shape.n }) and
51         specialization.outputHasExtents(0, &.{ shape.m, shape.n }) and
52         specialization.reductionMatches(0, .{
53             .name = "dot",
54             .operator = .dot_product,
55             .extents = &.{shape.k},
56         });
57 }