lib/accy/src/choir/record/memory.zig
daab053ee43316e1809a84551d573ddd1e5bf3d2
1 const std = @import("std");
2 const choir_abi = @import("choir_abi");
3 const reference = @import("root.zig").reference;
4
5 pub const LayoutKind = enum {
6 scalar,
7 row_major,
8 dynamic_row_major,
9
10 pub fn name(self: LayoutKind) []const u8 {
11 return switch (self) {
12 .scalar => "scalar",
13 .row_major => "row_major",
14 .dynamic_row_major => "dynamic_row_major",
15 };
16 }
17
18 pub fn hasStaticStrides(self: LayoutKind) bool {
19 return self == .scalar or self == .row_major;
20 }
21 };
22
23 pub const OutputSource = union(enum) {
24 kernel_written: usize,
25 aliased: usize,
26 };
27
28 pub const BoundaryTransfer = enum {
29 none,
30 host_to_device,
31 device_to_host,
32 bidirectional,
33
34 pub fn needsHostInput(self: BoundaryTransfer) bool {
35 return self == .host_to_device or self == .bidirectional;
36 }
37
38 pub fn needsHostOutput(self: BoundaryTransfer) bool {
39 return self == .device_to_host or self == .bidirectional;
40 }
41 };
42
43 pub const MemoryAccess = enum {
44 read_only,
45 write_only,
46 read_write,
47 };
48
49 pub const MemorySpace = enum {
50 host,
51 device_global,
52 device_constant,
53 device_shared,
54 unified,
55
56 pub fn name(self: MemorySpace) []const u8 {
57 return switch (self) {
58 .host => "host",
59 .device_global => "device_global",
60 .device_constant => "device_constant",
61 .device_shared => "device_shared",
62 .unified => "unified",
63 };
64 }
65 };
66
67 pub const BufferRole = struct {
68 input: bool = false,
69 output: bool = false,
70 temporary: bool = false,
71 constant: bool = false,
72
73 pub fn isBoundary(self: BufferRole) bool {
74 return self.input or self.output;
75 }
76 };
77
78 pub const Slot = struct {
79 id: usize,
80 value: reference.Value,
81 producer: ?reference.Operation,
82 function: ?reference.Operation,
83 role: BufferRole,
84 dtype: choir_abi.DType,
85 dims: []const i64,
86 element_count: ?u64,
87 row_major_strides: ?[]const u64,
88 byte_size: ?u64,
89 };
90
91 pub const Elision = struct {
92 value: reference.Value,
93 producer: reference.Operation,
94 root: reference.Operation,
95 cluster_index: usize,
96 };
97
98 pub const Buffers = struct {
99 slots: []const Slot,
100 elisions: []const Elision,
101 input_slot_count: usize,
102 output_slot_count: usize,
103 temporary_slot_count: usize,
104 constant_slot_count: usize,
105 dynamic_slot_count: usize,
106 total_static_bytes: u64,
107 };
108
109 pub const Assignment = struct {
110 slot_id: usize,
111 value: reference.Value,
112 producer: ?reference.Operation,
113 role: BufferRole,
114 space: MemorySpace,
115 access: MemoryAccess,
116 transfer: BoundaryTransfer,
117 byte_size: ?u64,
118 output_source: ?OutputSource,
119 };
120
121 pub const Spaces = struct {
122 assignments: []const Assignment,
123 host_slot_count: usize,
124 device_global_slot_count: usize,
125 device_constant_slot_count: usize,
126 device_shared_slot_count: usize,
127 unified_slot_count: usize,
128 host_input_transfer_count: usize,
129 host_output_transfer_count: usize,
130 dynamic_slot_count: usize,
131 elided_value_count: usize,
132 total_static_bytes: u64,
133 };
134
135 pub const Layout = struct {
136 slot_id: usize,
137 value: reference.Value,
138 producer: ?reference.Operation,
139 role: BufferRole,
140 dtype: choir_abi.DType,
141 memory_space: MemorySpace,
142 kind: LayoutKind,
143 rank: usize,
144 dims: []const i64,
145 element_strides: ?[]const u64,
146 minor_to_major: []const usize,
147 element_count: ?u64,
148 byte_size: ?u64,
149 element_size: u64,
150 alignment: u64,
151 contiguous: bool,
152 static_layout: bool,
153 };
154
155 pub const Layouts = struct {
156 assignments: []const Layout,
157 scalar_layout_count: usize,
158 row_major_layout_count: usize,
159 dynamic_row_major_layout_count: usize,
160 host_slot_count: usize,
161 device_global_slot_count: usize,
162 device_constant_slot_count: usize,
163 device_shared_slot_count: usize,
164 unified_slot_count: usize,
165 dynamic_slot_count: usize,
166 elided_value_count: usize,
167 total_static_bytes: u64,
168 };
169
170 pub const Binding = struct {
171 value: reference.Value,
172 slot_id: usize,
173
174 pub fn lessThan(_: void, left: Binding, right: Binding) bool {
175 const lhs = key(left.value);
176 const rhs = key(right.value);
177 return std.mem.lessThan(u32, &lhs, &rhs);
178 }
179
180 fn key(value: reference.Value) [5]u32 {
181 return switch (value) {
182 .result => |item| .{ item.operation.ordinal, 0, 0, 0, item.position },
183 .argument => |item| .{ item.operation.ordinal, 1, item.region, item.block, item.position },
184 };
185 }
186 };
187
188 pub const Record = struct {
189 buffers: Buffers,
190 bindings: []const Binding,
191 spaces: Spaces,
192 layouts: Layouts,
193 };
194
195 pub fn validate(allocator: std.mem.Allocator, value: Record, dispatch: @import("root.zig").dispatch.Record) !void {
196 try validateBindings(allocator, value);
197 try validateBufferCounts(value.buffers);
198 try validateSpaceCounts(value.spaces, value.buffers.elisions.len);
199 try validateLayoutCounts(value.layouts, value.buffers.elisions.len);
200 const slots = value.buffers.slots.len;
201 if (value.spaces.assignments.len != slots or value.layouts.assignments.len != slots) {
202 return error.InvalidStageRecord;
203 }
204 for (value.buffers.slots, 0..) |slot, index| {
205 if (slot.id != index) return error.InvalidStageRecord;
206 if (slot.row_major_strides) |strides| {
207 if (strides.len != slot.dims.len) return error.InvalidStageRecord;
208 }
209 }
210 for (value.buffers.elisions) |elision| {
211 if (elision.cluster_index >= dispatch.fusion.clusters.len) return error.InvalidStageRecord;
212 }
213 for (value.spaces.assignments, 0..) |assignment, index| {
214 if (assignment.slot_id != index) return error.InvalidStageRecord;
215 if (assignment.output_source) |source| switch (source) {
216 .aliased => |slot| if (slot >= slots) return error.InvalidStageRecord,
217 .kernel_written => |work| if (work >= dispatch.schedule.work_items.len) {
218 return error.InvalidStageRecord;
219 },
220 };
221 }
222 for (value.layouts.assignments, 0..) |layout, index| {
223 if (layout.slot_id != index or layout.rank != layout.dims.len or
224 layout.minor_to_major.len != layout.rank) return error.InvalidStageRecord;
225 if (layout.element_strides) |strides| {
226 if (strides.len != layout.rank) return error.InvalidStageRecord;
227 }
228 for (layout.minor_to_major) |axis| {
229 if (axis >= layout.rank) return error.InvalidStageRecord;
230 }
231 }
232 }
233
234 fn validateBindings(allocator: std.mem.Allocator, value: Record) !void {
235 const primary = try allocator.alloc(bool, value.buffers.slots.len);
236 defer allocator.free(primary);
237 @memset(primary, false);
238 for (value.bindings, 0..) |binding, index| {
239 if (binding.slot_id >= primary.len) return error.InvalidStageRecord;
240 if (index != 0 and !Binding.lessThan({}, value.bindings[index - 1], binding)) {
241 return error.InvalidStageRecord;
242 }
243 if (std.meta.eql(binding.value, value.buffers.slots[binding.slot_id].value)) {
244 primary[binding.slot_id] = true;
245 }
246 }
247 for (primary) |present| if (!present) return error.InvalidStageRecord;
248 }
249
250 const SizeCounts = struct {
251 dynamic: usize = 0,
252 bytes: u64 = 0,
253
254 fn add(self: *SizeCounts, bytes: ?u64) !void {
255 if (bytes) |size| {
256 self.bytes = std.math.add(u64, self.bytes, size) catch
257 return error.InvalidStageRecord;
258 } else self.dynamic += 1;
259 }
260
261 fn matches(self: SizeCounts, value: anytype) bool {
262 return self.dynamic == value.dynamic_slot_count and self.bytes == value.total_static_bytes;
263 }
264 };
265
266 const SpaceCounts = struct {
267 values: [5]usize = @splat(0),
268
269 fn add(self: *SpaceCounts, space: MemorySpace) void {
270 const index: usize = switch (space) {
271 .host => 0,
272 .device_global => 1,
273 .device_constant => 2,
274 .device_shared => 3,
275 .unified => 4,
276 };
277 self.values[index] += 1;
278 }
279
280 fn matches(self: SpaceCounts, value: anytype) bool {
281 return std.mem.eql(usize, &self.values, &.{
282 value.host_slot_count,
283 value.device_global_slot_count,
284 value.device_constant_slot_count,
285 value.device_shared_slot_count,
286 value.unified_slot_count,
287 });
288 }
289 };
290
291 fn validateBufferCounts(value: Buffers) !void {
292 var roles: [4]usize = @splat(0);
293 var sizes: SizeCounts = .{};
294 for (value.slots) |slot| {
295 if (slot.role.input) roles[0] += 1;
296 if (slot.role.output) roles[1] += 1;
297 if (slot.role.temporary) roles[2] += 1;
298 if (slot.role.constant) roles[3] += 1;
299 try sizes.add(slot.byte_size);
300 }
301 if (!sizes.matches(value) or !std.mem.eql(usize, &roles, &.{
302 value.input_slot_count,
303 value.output_slot_count,
304 value.temporary_slot_count,
305 value.constant_slot_count,
306 })) return error.InvalidStageRecord;
307 }
308
309 fn validateSpaceCounts(value: Spaces, elisions: usize) !void {
310 var spaces: SpaceCounts = .{};
311 var sizes: SizeCounts = .{};
312 var inputs: usize = 0;
313 var outputs: usize = 0;
314 for (value.assignments) |assignment| {
315 spaces.add(assignment.space);
316 try sizes.add(assignment.byte_size);
317 if (assignment.transfer.needsHostInput()) inputs += 1;
318 if (assignment.transfer.needsHostOutput()) outputs += 1;
319 }
320 if (!spaces.matches(value) or !sizes.matches(value) or
321 inputs != value.host_input_transfer_count or outputs != value.host_output_transfer_count or
322 elisions != value.elided_value_count) return error.InvalidStageRecord;
323 }
324
325 fn validateLayoutCounts(value: Layouts, elisions: usize) !void {
326 var spaces: SpaceCounts = .{};
327 var sizes: SizeCounts = .{};
328 var kinds: [3]usize = @splat(0);
329 for (value.assignments) |assignment| {
330 spaces.add(assignment.memory_space);
331 try sizes.add(assignment.byte_size);
332 switch (assignment.kind) {
333 .scalar => kinds[0] += 1,
334 .row_major => kinds[1] += 1,
335 .dynamic_row_major => kinds[2] += 1,
336 }
337 }
338 if (!spaces.matches(value) or !sizes.matches(value) or elisions != value.elided_value_count or
339 !std.mem.eql(usize, &kinds, &.{
340 value.scalar_layout_count,
341 value.row_major_layout_count,
342 value.dynamic_row_major_layout_count,
343 })) return error.InvalidStageRecord;
344 }