lib/accy/src/kernel/library/catalog/registry.zig

daab053ee43316e1809a84551d573ddd1e5bf3d2

 1 const std = @import("std");
 2 
 3 const attention = @import("../root.zig").attention;
 4 const compaction = @import("../root.zig").compaction;
 5 const descriptor_mod = @import("descriptor.zig");
 6 const elementwise = @import("../root.zig").elementwise;
 7 const fused = @import("../root.zig").fused;
 8 const indexing = @import("../root.zig").indexing;
 9 const layout = @import("../root.zig").layout;
10 const linalg = @import("../root.zig").linalg;
11 const normalization = @import("../root.zig").normalization;
12 const random = @import("../root.zig").random;
13 const reduction = @import("../root.zig").reduction;
14 const scan = @import("../root.zig").scan;
15 const segmented = @import("../root.zig").segmented;
16 const stencil = @import("../root.zig").stencil;
17 
18 pub const Descriptor = descriptor_mod.Descriptor;
19 
20 pub const entries = .{
21     attention.ScaledDotProductAttention2x2x3x2x2F32,
22     elementwise.VectorAdd8F32,
23     elementwise.Axpy8F32,
24     elementwise.Gelu8F32,
25     elementwise.Relu8F32,
26     elementwise.Silu8F32,
27     elementwise.AuthoredScale8F32,
28     fused.BiasGelu8F32,
29     fused.BiasRelu8F32,
30     fused.BiasSilu8F32,
31     fused.GeGlu8F32,
32     fused.ReGlu8F32,
33     fused.SwiGlu8F32,
34     fused.MatrixProductBiasGelu2x3x4F32,
35     fused.MatrixProductBiasGelu2x3x4ThreadBlocks1x2F32,
36     fused.MatrixProductBiasRelu2x3x4F32,
37     fused.MatrixProductBiasSilu2x3x4F32,
38     fused.MatrixVectorProductBiasGelu4x8F32,
39     fused.MatrixVectorProductBiasRelu4x8F32,
40     fused.MatrixVectorProductBiasSilu4x8F32,
41     normalization.RowResidualRmsNorm2x4F32,
42     layout.Transpose8x16F32,
43     reduction.Sum8F32,
44     reduction.Dot8F32,
45     normalization.RowSoftmax2x4F32,
46     normalization.RowSoftmax2x4ThreadBlocks2x2F32,
47     normalization.RowLogSoftmax2x4F32,
48     normalization.RowRmsNorm2x4F32,
49     normalization.RowLayerNorm2x4F32,
50     normalization.RowAffineLayerNorm2x4F32,
51     linalg.MatrixProduct4x16x8F32,
52     linalg.MatrixProduct4x16x8ThreadBlocks4x2F32,
53     stencil.Window2x3R1F32,
54     indexing.Gather8F32,
55     indexing.Scatter8F32,
56     compaction.Filter8F32,
57     random.Philox8F32,
58     random.Threefry8F32,
59     segmented.SegmentSum4F32,
60     scan.PrefixSum8F32,
61 };
62 
63 pub const descriptors = buildDescriptors(entries);
64 
65 fn buildDescriptors(comptime entry_types: anytype) [entry_types.len]Descriptor {
66     var result: [entry_types.len]Descriptor = undefined;
67     inline for (entry_types, 0..) |EntryType, index| {
68         result[index] = descriptorFor(EntryType);
69     }
70     return result;
71 }
72 
73 fn descriptorFor(comptime EntryType: type) Descriptor {
74     return .{
75         .name = EntryType.name,
76         .metadata = EntryType.metadata,
77     };
78 }
79 
80 pub fn findEntry(target: []const u8, version: u32) ?Descriptor {
81     for (descriptors) |descriptor| {
82         if (descriptor.metadata.version != version) continue;
83         if (!std.mem.eql(u8, descriptor.metadata.target, target)) continue;
84         return descriptor;
85     }
86     return null;
87 }