lib/choir/src/abi/dtype.zig
daab053ee43316e1809a84551d573ddd1e5bf3d2
1 const std = @import("std");
2
3 pub const Bf16 = packed struct(u16) {
4 bits: u16,
5
6 pub fn fromF32(x: f32) Bf16 {
7 const f32_bits: u32 = @bitCast(x);
8 return .{ .bits = @truncate(f32_bits >> 16) };
9 }
10
11 pub fn toF32(self: Bf16) f32 {
12 const widened: u32 = @as(u32, self.bits) << 16;
13 return @bitCast(widened);
14 }
15 };
16
17 pub const Key = packed struct(u64) {
18 bits: u64,
19
20 pub fn init(bits: u64) Key {
21 return .{ .bits = bits };
22 }
23
24 pub fn fromWords(low: u32, high: u32) Key {
25 return .{ .bits = @as(u64, low) | (@as(u64, high) << 32) };
26 }
27
28 pub fn lo(self: Key) u32 {
29 return @truncate(self.bits);
30 }
31
32 pub fn hi(self: Key) u32 {
33 return @truncate(self.bits >> 32);
34 }
35 };
36
37 pub const DType = enum(u8) {
38 i1 = 0,
39 i8 = 1,
40 i16 = 2,
41 i32 = 3,
42 i64 = 4,
43 u8 = 5,
44 u16 = 6,
45 u32 = 7,
46 u64 = 8,
47 f16 = 9,
48 bf16 = 10,
49 f32 = 11,
50 f64 = 12,
51 key = 13,
52
53 pub fn ZigType(comptime self: DType) type {
54 return switch (self) {
55 .i1 => bool,
56 .i8 => i8,
57 .i16 => i16,
58 .i32 => i32,
59 .i64 => i64,
60 .u8 => u8,
61 .u16 => u16,
62 .u32 => u32,
63 .u64 => u64,
64 .f16 => f16,
65 .bf16 => Bf16,
66 .f32 => f32,
67 .f64 => f64,
68 .key => Key,
69 };
70 }
71
72 pub fn sizeOf(self: DType) u8 {
73 return switch (self) {
74 inline else => |tag| @sizeOf(tag.ZigType()),
75 };
76 }
77
78 pub fn alignOf(self: DType) u8 {
79 return switch (self) {
80 inline else => |tag| @alignOf(tag.ZigType()),
81 };
82 }
83
84 pub fn name(self: DType) []const u8 {
85 return @tagName(self);
86 }
87
88 pub fn isFloat(self: DType) bool {
89 return switch (self) {
90 .f16, .bf16, .f32, .f64 => true,
91 else => false,
92 };
93 }
94
95 pub fn isSignedInt(self: DType) bool {
96 return switch (self) {
97 .i8, .i16, .i32, .i64 => true,
98 else => false,
99 };
100 }
101
102 pub fn isUnsignedInt(self: DType) bool {
103 return switch (self) {
104 .u8, .u16, .u32, .u64 => true,
105 else => false,
106 };
107 }
108
109 pub fn isBool(self: DType) bool {
110 return self == .i1;
111 }
112
113 pub fn isKey(self: DType) bool {
114 return self == .key;
115 }
116
117 pub fn isNumeric(self: DType) bool {
118 return self.isFloat() or self.isSignedInt() or self.isUnsignedInt();
119 }
120
121 pub fn fromName(s: []const u8) ?DType {
122 if (std.mem.eql(u8, s, "bool")) return .i1;
123 inline for (
124 @typeInfo(DType).@"enum".field_names,
125 @typeInfo(DType).@"enum".field_values,
126 ) |field_name, field_name_value| {
127 const field = .{ .name = field_name, .value = field_name_value };
128 if (std.mem.eql(u8, s, field.name)) return @fromBackingInt(@intCast(field.value));
129 }
130 return null;
131 }
132
133 pub fn fromZigType(comptime T: type) ?DType {
134 inline for (
135 @typeInfo(DType).@"enum".field_names,
136 @typeInfo(DType).@"enum".field_values,
137 ) |field_name, field_name_value| {
138 const field = .{ .name = field_name, .value = field_name_value };
139 const tag: DType = comptime @fromBackingInt(@intCast(field.value));
140 if (T == comptime tag.ZigType()) return tag;
141 }
142 return null;
143 }
144 };
145
146 test "name returns canonical tag spelling" {
147 try std.testing.expectEqualStrings("i1", DType.i1.name());
148 try std.testing.expectEqualStrings("i32", DType.i32.name());
149 try std.testing.expectEqualStrings("u64", DType.u64.name());
150 try std.testing.expectEqualStrings("f16", DType.f16.name());
151 try std.testing.expectEqualStrings("bf16", DType.bf16.name());
152 try std.testing.expectEqualStrings("f64", DType.f64.name());
153 try std.testing.expectEqualStrings("key", DType.key.name());
154 }
155
156 test "fromName round-trips name for every variant" {
157 inline for (
158 @typeInfo(DType).@"enum".field_names,
159 @typeInfo(DType).@"enum".field_values,
160 ) |f_name, f_name_value| {
161 const f = .{ .name = f_name, .value = f_name_value };
162 const tag: DType = @fromBackingInt(@intCast(f.value));
163 try std.testing.expectEqual(@as(?DType, tag), DType.fromName(tag.name()));
164 }
165 }
166
167 test "fromName accepts the bool alias for i1" {
168 try std.testing.expectEqual(@as(?DType, .i1), DType.fromName("bool"));
169 }
170
171 test "fromName rejects unknown spellings" {
172 try std.testing.expectEqual(@as(?DType, null), DType.fromName(""));
173 try std.testing.expectEqual(@as(?DType, null), DType.fromName("complex64"));
174 try std.testing.expectEqual(@as(?DType, null), DType.fromName("F32"));
175 }
176
177 test "sizeOf widths match the IEEE / two's-complement layout" {
178 try std.testing.expectEqual(@as(u8, 1), DType.i1.sizeOf());
179 try std.testing.expectEqual(@as(u8, 1), DType.i8.sizeOf());
180 try std.testing.expectEqual(@as(u8, 2), DType.i16.sizeOf());
181 try std.testing.expectEqual(@as(u8, 4), DType.i32.sizeOf());
182 try std.testing.expectEqual(@as(u8, 8), DType.i64.sizeOf());
183 try std.testing.expectEqual(@as(u8, 1), DType.u8.sizeOf());
184 try std.testing.expectEqual(@as(u8, 2), DType.u16.sizeOf());
185 try std.testing.expectEqual(@as(u8, 4), DType.u32.sizeOf());
186 try std.testing.expectEqual(@as(u8, 8), DType.u64.sizeOf());
187 try std.testing.expectEqual(@as(u8, 2), DType.f16.sizeOf());
188 try std.testing.expectEqual(@as(u8, 2), DType.bf16.sizeOf());
189 try std.testing.expectEqual(@as(u8, 4), DType.f32.sizeOf());
190 try std.testing.expectEqual(@as(u8, 8), DType.f64.sizeOf());
191 try std.testing.expectEqual(@as(u8, 8), DType.key.sizeOf());
192 }
193
194 test "sizeOf and alignOf agree with ZigType for every variant" {
195 inline for (
196 @typeInfo(DType).@"enum".field_names,
197 @typeInfo(DType).@"enum".field_values,
198 ) |f_name, f_name_value| {
199 const f = .{ .name = f_name, .value = f_name_value };
200 const tag: DType = comptime @fromBackingInt(@intCast(f.value));
201 const T = tag.ZigType();
202 try std.testing.expectEqual(@as(u8, @sizeOf(T)), tag.sizeOf());
203 try std.testing.expectEqual(@as(u8, @alignOf(T)), tag.alignOf());
204 }
205 }
206
207 test "ZigType maps i1 to bool" {
208 try std.testing.expect(DType.i1.ZigType() == bool);
209 }
210
211 test "ZigType maps integer dtypes to native Zig integers" {
212 try std.testing.expect(DType.i8.ZigType() == i8);
213 try std.testing.expect(DType.i16.ZigType() == i16);
214 try std.testing.expect(DType.i32.ZigType() == i32);
215 try std.testing.expect(DType.i64.ZigType() == i64);
216 try std.testing.expect(DType.u8.ZigType() == u8);
217 try std.testing.expect(DType.u16.ZigType() == u16);
218 try std.testing.expect(DType.u32.ZigType() == u32);
219 try std.testing.expect(DType.u64.ZigType() == u64);
220 }
221
222 test "ZigType maps float dtypes to native Zig floats and Bf16" {
223 try std.testing.expect(DType.f16.ZigType() == f16);
224 try std.testing.expect(DType.bf16.ZigType() == Bf16);
225 try std.testing.expect(DType.f32.ZigType() == f32);
226 try std.testing.expect(DType.f64.ZigType() == f64);
227 }
228
229 test "ZigType maps key dtype to opaque key payload" {
230 try std.testing.expect(DType.key.ZigType() == Key);
231 }
232
233 test "fromZigType is the inverse of ZigType" {
234 inline for (
235 @typeInfo(DType).@"enum".field_names,
236 @typeInfo(DType).@"enum".field_values,
237 ) |f_name, f_name_value| {
238 const f = .{ .name = f_name, .value = f_name_value };
239 const tag: DType = comptime @fromBackingInt(@intCast(f.value));
240 try std.testing.expectEqual(@as(?DType, tag), comptime DType.fromZigType(tag.ZigType()));
241 }
242 }
243
244 test "fromZigType rejects types outside the dtype set" {
245 try std.testing.expectEqual(@as(?DType, null), DType.fromZigType(u128));
246 try std.testing.expectEqual(@as(?DType, null), DType.fromZigType(usize));
247 try std.testing.expectEqual(@as(?DType, null), DType.fromZigType(f80));
248 }
249
250 test "isFloat / isSignedInt / isUnsignedInt / isBool / isKey partition the enum" {
251 inline for (
252 @typeInfo(DType).@"enum".field_names,
253 @typeInfo(DType).@"enum".field_values,
254 ) |f_name, f_name_value| {
255 const f = .{ .name = f_name, .value = f_name_value };
256 const tag: DType = @fromBackingInt(@intCast(f.value));
257 const flags = [_]bool{ tag.isFloat(), tag.isSignedInt(), tag.isUnsignedInt(), tag.isBool(), tag.isKey() };
258 var trues: u32 = 0;
259 for (flags) |b| if (b) {
260 trues += 1;
261 };
262 try std.testing.expectEqual(@as(u32, 1), trues);
263 }
264 }
265
266 test "isFloat names the float dtypes" {
267 try std.testing.expect(DType.f16.isFloat());
268 try std.testing.expect(DType.bf16.isFloat());
269 try std.testing.expect(DType.f32.isFloat());
270 try std.testing.expect(DType.f64.isFloat());
271 try std.testing.expect(!DType.i32.isFloat());
272 try std.testing.expect(!DType.i1.isFloat());
273 }
274
275 test "key payload exposes stable counter words" {
276 const key = Key.fromWords(0x89abcdef, 0x01234567);
277 try std.testing.expectEqual(@as(u64, 0x01234567_89abcdef), key.bits);
278 try std.testing.expectEqual(@as(u32, 0x89abcdef), key.lo());
279 try std.testing.expectEqual(@as(u32, 0x01234567), key.hi());
280 try std.testing.expect(!DType.key.isNumeric());
281 }
282
283 test "Bf16 fromF32 / toF32 round-trips simple values exactly" {
284 const cases = [_]f32{ 0.0, 1.0, -1.0, 2.0, 0.5, -0.5, 4.0 };
285 for (cases) |v| {
286 const x = Bf16.fromF32(v);
287 try std.testing.expectEqual(v, x.toF32());
288 }
289 }
290
291 test "Bf16 truncates the low 16 bits of f32" {
292 const x = Bf16.fromF32(1.0);
293 const expected_bits: u32 = @bitCast(@as(f32, 1.0));
294 try std.testing.expectEqual(@as(u16, @truncate(expected_bits >> 16)), x.bits);
295 }