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

daab053ee43316e1809a84551d573ddd1e5bf3d2

  1 const std = @import("std");
  2 const choir_abi = @import("choir_abi");
  3 
  4 const accy = @import("../../../root.zig");
  5 const library = @import("../../library/root.zig");
  6 
  7 const DType = choir_abi.DType;
  8 
  9 pub const RowNormalizationKind = library.RowNormalizationKind;
 10 pub const RowNormalizationParameterization = library.RowNormalizationParameterization;
 11 pub const RowNormalizationSchedule = library.RowNormalizationSchedule;
 12 
 13 pub const RowNormalizationSelectionRequest = struct {
 14     dtype: DType,
 15     kind: RowNormalizationKind,
 16     rows: u64,
 17     cols: u64,
 18     schedule: ?RowNormalizationSchedule = null,
 19 };
 20 
 21 pub const SelectedRowNormalizationKernel = struct {
 22     kind: RowNormalizationKind,
 23     descriptor: library.CatalogDescriptor,
 24 };
 25 
 26 pub fn selectCatalog(request: RowNormalizationSelectionRequest) ?SelectedRowNormalizationKernel {
 27     const descriptor = library.select(.{ .row_normalization = .{
 28         .dtype = request.dtype,
 29         .kind = request.kind,
 30         .rows = request.rows,
 31         .cols = request.cols,
 32         .schedule = request.schedule,
 33     } }) orelse return null;
 34     return .{
 35         .kind = request.kind,
 36         .descriptor = descriptor,
 37     };
 38 }
 39 
 40 test "logical row normalization selection chooses catalog entries" {
 41     const softmax = selectCatalog(.{
 42         .dtype = .f32,
 43         .kind = .softmax,
 44         .rows = 2,
 45         .cols = 4,
 46     }) orelse return error.TestExpectedRowSoftmaxSelection;
 47     const log_softmax = selectCatalog(.{
 48         .dtype = .f32,
 49         .kind = .log_softmax,
 50         .rows = 2,
 51         .cols = 4,
 52     }) orelse return error.TestExpectedRowLogSoftmaxSelection;
 53     const rmsnorm = selectCatalog(.{
 54         .dtype = .f32,
 55         .kind = .{ .rmsnorm = .scale },
 56         .rows = 2,
 57         .cols = 4,
 58     }) orelse return error.TestExpectedRowRmsNormSelection;
 59     const layernorm = selectCatalog(.{
 60         .dtype = .f32,
 61         .kind = .{ .layernorm = .none },
 62         .rows = 2,
 63         .cols = 4,
 64     }) orelse return error.TestExpectedRowLayerNormSelection;
 65     const affine_layernorm = selectCatalog(.{
 66         .dtype = .f32,
 67         .kind = .{ .layernorm = .scale_bias },
 68         .rows = 2,
 69         .cols = 4,
 70     }) orelse return error.TestExpectedRowAffineLayerNormSelection;
 71 
 72     try std.testing.expectEqual(RowNormalizationKind.softmax, softmax.kind);
 73     try std.testing.expectEqualStrings(library.normalization.RowSoftmax2x4F32.target, softmax.descriptor.metadata.target);
 74     try std.testing.expectEqual(RowNormalizationKind.log_softmax, log_softmax.kind);
 75     try std.testing.expectEqualStrings(library.normalization.RowLogSoftmax2x4F32.target, log_softmax.descriptor.metadata.target);
 76     try std.testing.expectEqualDeep(RowNormalizationKind{ .rmsnorm = .scale }, rmsnorm.kind);
 77     try std.testing.expectEqualStrings(library.normalization.RowRmsNorm2x4F32.target, rmsnorm.descriptor.metadata.target);
 78     try std.testing.expectEqualDeep(RowNormalizationKind{ .layernorm = .none }, layernorm.kind);
 79     try std.testing.expectEqualStrings(library.normalization.RowLayerNorm2x4F32.target, layernorm.descriptor.metadata.target);
 80     try std.testing.expectEqualDeep(RowNormalizationKind{ .layernorm = .scale_bias }, affine_layernorm.kind);
 81     try std.testing.expectEqualStrings(library.normalization.RowAffineLayerNorm2x4F32.target, affine_layernorm.descriptor.metadata.target);
 82 }
 83 
 84 test "logical row normalization selection chooses schedule-specialized catalog entries" {
 85     const softmax = selectCatalog(.{
 86         .dtype = .f32,
 87         .kind = .softmax,
 88         .rows = 2,
 89         .cols = 4,
 90         .schedule = .{ .thread_blocks = .{ .x = 2, .y = 2 } },
 91     }) orelse return error.TestExpectedRowSoftmaxScheduleSelection;
 92 
 93     try std.testing.expectEqual(RowNormalizationKind.softmax, softmax.kind);
 94     try std.testing.expectEqualStrings(library.normalization.RowSoftmax2x4ThreadBlocks2x2F32.target, softmax.descriptor.metadata.target);
 95     try std.testing.expectEqual(@as(u32, 2), softmax.descriptor.metadata.specialization.launch.?.threadgroup[0]);
 96     try std.testing.expectEqual(@as(u32, 2), softmax.descriptor.metadata.specialization.launch.?.threadgroup[1]);
 97 }
 98 
 99 test "logical row normalization selection rejects unavailable catalog entry" {
100     try std.testing.expect(selectCatalog(.{
101         .dtype = .f32,
102         .kind = .softmax,
103         .rows = 3,
104         .cols = 4,
105     }) == null);
106     try std.testing.expect(selectCatalog(.{
107         .dtype = .f32,
108         .kind = .log_softmax,
109         .rows = 2,
110         .cols = 5,
111     }) == null);
112     try std.testing.expect(selectCatalog(.{
113         .dtype = .i32,
114         .kind = .{ .rmsnorm = .scale },
115         .rows = 2,
116         .cols = 4,
117     }) == null);
118     try std.testing.expect(selectCatalog(.{
119         .dtype = .f32,
120         .kind = .{ .layernorm = .none },
121         .rows = 2,
122         .cols = 5,
123     }) == null);
124     try std.testing.expect(selectCatalog(.{
125         .dtype = .f32,
126         .kind = .{ .layernorm = .scale },
127         .rows = 2,
128         .cols = 4,
129     }) == null);
130     try std.testing.expect(selectCatalog(.{
131         .dtype = .f32,
132         .kind = .softmax,
133         .rows = 2,
134         .cols = 4,
135         .schedule = .{ .thread_blocks = .{ .x = 8, .y = 2 } },
136     }) == null);
137 }