lib/accy/src/kernel/library/catalog/match/common.zig

daab053ee43316e1809a84551d573ddd1e5bf3d2

 1 const std = @import("std");
 2 
 3 const library = @import("../../root.zig");
 4 
 5 const entry = library.entry;
 6 
 7 pub fn selectableSpecialization(specialization: entry.Specialization) bool {
 8     return specialization.reductionDependenciesAreValid() and
 9         specialization.reductionReuseScopesAreValid() and
10         specialization.staticParametersAreValid();
11 }
12 pub fn specializationThreadgroup1DMatches(specialization: entry.Specialization, threads: u32) bool {
13     const launch = specialization.launch orelse return false;
14     return launch.threadgroup[0] == threads;
15 }
16 pub fn specializationThreadgroup2DMatches(specialization: entry.Specialization, threads: entry.Threads2D) bool {
17     const launch = specialization.launch orelse return false;
18     return launch.threadgroup[0] == threads.x and
19         launch.threadgroup[1] == threads.y;
20 }
21 pub fn specializationThreadgroup3DMatches(specialization: entry.Specialization, threads: entry.Threads3D) bool {
22     const launch = specialization.launch orelse return false;
23     return launch.threadgroup[0] == threads.x and
24         launch.threadgroup[1] == threads.y and
25         launch.threadgroup[2] == threads.z;
26 }
27 test "kernel library catalog requires valid reduction dependency metadata" {
28     const column = entry.shape1D("col", 4);
29     const rows = entry.shape1D("row", 2);
30     const valid = entry.Specialization{
31         .reductions = &.{
32             entry.reduction("score_dot", .dot_product, column),
33             entry.dependentReduction("score_max", .maximum, column, &.{"score_dot"}),
34         },
35         .reduction_reuse = &.{entry.reductionReuse("score_max", rows)},
36     };
37     const invalid = entry.Specialization{
38         .reductions = &.{
39             entry.dependentReduction("score_max", .maximum, column, &.{"score_dot"}),
40             entry.reduction("score_dot", .dot_product, column),
41         },
42     };
43     const invalid_reuse = entry.Specialization{
44         .reductions = &.{entry.reduction("score_max", .maximum, column)},
45         .reduction_reuse = &.{entry.reductionReuse("score_exp_sum", rows)},
46     };
47 
48     try std.testing.expect(selectableSpecialization(valid));
49     try std.testing.expect(!selectableSpecialization(invalid));
50     try std.testing.expect(!selectableSpecialization(invalid_reuse));
51 }
52 test "kernel library catalog matches specialization threadgroup geometry by rank" {
53     const specialization = entry.Specialization{
54         .launch = entry.launch3D(16, 8, 4, 4, 2, 1),
55     };
56     const no_launch = entry.Specialization{};
57 
58     try std.testing.expect(specializationThreadgroup1DMatches(specialization, 4));
59     try std.testing.expect(specializationThreadgroup2DMatches(specialization, .{ .x = 4, .y = 2 }));
60     try std.testing.expect(specializationThreadgroup3DMatches(specialization, .{ .x = 4, .y = 2, .z = 1 }));
61     try std.testing.expect(!specializationThreadgroup1DMatches(specialization, 2));
62     try std.testing.expect(!specializationThreadgroup2DMatches(specialization, .{ .x = 4, .y = 1 }));
63     try std.testing.expect(!specializationThreadgroup3DMatches(specialization, .{ .x = 4, .y = 2, .z = 2 }));
64     try std.testing.expect(!specializationThreadgroup1DMatches(no_launch, 4));
65     try std.testing.expect(!specializationThreadgroup2DMatches(no_launch, .{ .x = 4, .y = 2 }));
66     try std.testing.expect(!specializationThreadgroup3DMatches(no_launch, .{ .x = 4, .y = 2, .z = 1 }));
67 }