lib/accy/src/kernel/logical/selection/activation.zig

daab053ee43316e1809a84551d573ddd1e5bf3d2

 1 const std = @import("std");
 2 const choir_abi = @import("choir_abi");
 3 
 4 const accy = @import("../../../root.zig");
 5 const activation = @import("../../../choir/root.zig").activation;
 6 const library = @import("../../library/root.zig");
 7 
 8 const DType = choir_abi.DType;
 9 
10 pub const ActivationKind = activation.Kind;
11 
12 pub const ActivationSelectionRequest = struct {
13     dtype: DType,
14     kind: ActivationKind,
15     extent: u64,
16 };
17 
18 pub const SelectedActivationKernel = struct {
19     kind: ActivationKind,
20     descriptor: library.CatalogDescriptor,
21 };
22 
23 pub fn selectCatalog(request: ActivationSelectionRequest) ?SelectedActivationKernel {
24     const descriptor = library.select(.{ .activation = .{
25         .dtype = request.dtype,
26         .kind = request.kind,
27         .extent = request.extent,
28     } }) orelse return null;
29     return .{
30         .kind = request.kind,
31         .descriptor = descriptor,
32     };
33 }
34 
35 test "logical activation selection chooses catalog entries" {
36     const gelu = selectCatalog(.{
37         .dtype = .f32,
38         .kind = .gelu,
39         .extent = 8,
40     }) orelse return error.TestExpectedGeluSelection;
41     const relu = selectCatalog(.{
42         .dtype = .f32,
43         .kind = .relu,
44         .extent = 8,
45     }) orelse return error.TestExpectedReluSelection;
46     const silu = selectCatalog(.{
47         .dtype = .f32,
48         .kind = .silu,
49         .extent = 8,
50     }) orelse return error.TestExpectedSiluSelection;
51 
52     try std.testing.expectEqual(ActivationKind.gelu, gelu.kind);
53     try std.testing.expectEqualStrings(library.elementwise.Gelu8F32.target, gelu.descriptor.metadata.target);
54     try std.testing.expectEqual(ActivationKind.relu, relu.kind);
55     try std.testing.expectEqualStrings(library.elementwise.Relu8F32.target, relu.descriptor.metadata.target);
56     try std.testing.expectEqual(ActivationKind.silu, silu.kind);
57     try std.testing.expectEqualStrings(library.elementwise.Silu8F32.target, silu.descriptor.metadata.target);
58 }
59 
60 test "logical activation selection rejects unavailable catalog entry" {
61     try std.testing.expect(selectCatalog(.{
62         .dtype = .i32,
63         .kind = .silu,
64         .extent = 8,
65     }) == null);
66     try std.testing.expect(selectCatalog(.{
67         .dtype = .f32,
68         .kind = .gelu,
69         .extent = 16,
70     }) == null);
71     try std.testing.expect(selectCatalog(.{
72         .dtype = .i32,
73         .kind = .relu,
74         .extent = 8,
75     }) == null);
76     try std.testing.expect(selectCatalog(.{
77         .dtype = .f32,
78         .kind = .relu,
79         .extent = 0,
80     }) == null);
81     try std.testing.expect(selectCatalog(.{
82         .dtype = .f32,
83         .kind = .silu,
84         .extent = 0,
85     }) == null);
86 }