lib/accy/src/tensor/session/session.zig
daab053ee43316e1809a84551d573ddd1e5bf3d2
1 const std = @import("std");
2 const gpu = @import("gpu");
3 const accy = @import("../../root.zig");
4 const tensor = @import("../root.zig");
5
6 const binding = accy.executable.binding;
7 const BackendArtifactPlan = accy.artifact.BackendArtifactPlan;
8 const CompiledFragment = accy.executable.CompiledFragment;
9 const CpuState = gpu.cpu.State;
10 const LoadedFragment = accy.executable.LoadedFragment;
11 const Program = tensor.program.Program;
12
13 /// A caller sizes its bookkeeping by this bound on programs held at once: one session holds at most
14 /// 64 compiled programs at a time. Compiling another program while all 64 are held returns
15 /// `error.TooManyPrograms`.
16 pub const max_programs: u32 = 64;
17
18 /// A caller matches on these to handle session failures: the set lists every failure a session call
19 /// can return. Each failure maps to exactly one stable status number.
20 pub const Error = error{
21 InvalidProgram,
22 UnsupportedVersion,
23 InvalidHandle,
24 InvalidBuffer,
25 TooManyPrograms,
26 OutOfMemory,
27 CompileFailed,
28 LaunchFailed,
29 };
30
31 /// A host sends these numbers to report failures across a language boundary: the codes run from 0
32 /// for success to 8 for a failed launch. Once a number is assigned it keeps its meaning for good.
33 pub const Status = enum(u32) {
34 ok = 0,
35 invalid_program = 1,
36 unsupported_version = 2,
37 invalid_handle = 3,
38 invalid_buffer = 4,
39 too_many_programs = 5,
40 out_of_memory = 6,
41 compile_failed = 7,
42 launch_failed = 8,
43
44 pub fn fromError(err: Error) Status {
45 return switch (err) {
46 error.InvalidProgram => .invalid_program,
47 error.UnsupportedVersion => .unsupported_version,
48 error.InvalidHandle => .invalid_handle,
49 error.InvalidBuffer => .invalid_buffer,
50 error.TooManyPrograms => .too_many_programs,
51 error.OutOfMemory => .out_of_memory,
52 error.CompileFailed => .compile_failed,
53 error.LaunchFailed => .launch_failed,
54 };
55 }
56 };
57
58 /// A caller keeps this handle to name one compiled program in later calls by slot index and
59 /// generation count. `Session.release` changes the slot's generation, so every copy of the old
60 /// handle is refused with `error.InvalidHandle` from then on.
61 pub const Handle = struct {
62 index: u32,
63 generation: u32,
64 };
65
66 /// A caller reads these descriptors to size and type the buffers it passes to a run: the value
67 /// gives the element type, axes, byte length and required alignment of one input or output. `dims`
68 /// points into the compiled program's own types and stays valid until that program is released.
69 pub const Descriptor = struct {
70 dtype: tensor.DType,
71 dims: []const tensor.Dim,
72 byte_size: usize,
73 alignment: u32,
74 };
75
76 const Entry = struct {
77 live: bool = false,
78 generation: u32 = 0,
79 program: Program = undefined,
80 compiled: *const CompiledFragment = undefined,
81 fragment: *LoadedFragment = undefined,
82 descriptors: []Descriptor = &.{},
83 input_count: usize = 0,
84 };
85
86 /// A host creates a session to compile tensor programs for the CPU backend inside this process and
87 /// run them on memory it supplies. The session runs programs forward only, with no transform that
88 /// computes derivatives.
89 pub const Session = struct {
90 allocator: std.mem.Allocator,
91 cpu: *CpuState,
92 entries: [max_programs]Entry = @splat(.{}),
93 /// A caller reads this field to log why the last failing call failed: the field holds the inner
94 /// error behind the most recent failure that went through the session's error mapping, and it
95 /// is for diagnostics only. Failures returned directly, such as a full table, a stale handle, a
96 /// buffer count mismatch or a failed allocation while describing buffers, leave it unchanged.
97 last_cause: ?anyerror = null,
98
99 pub fn init(allocator: std.mem.Allocator) Error!Session {
100 const cpu = allocator.create(CpuState) catch return error.OutOfMemory;
101 cpu.* = CpuState.init(allocator);
102 return .{ .allocator = allocator, .cpu = cpu };
103 }
104
105 pub fn deinit(self: *Session) void {
106 for (&self.entries) |*entry| {
107 if (entry.live) self.releaseEntry(entry);
108 }
109 self.cpu.deinit();
110 self.allocator.destroy(self.cpu);
111 self.* = undefined;
112 }
113
114 /// A caller passes a program's wire bytes, the versioned byte encoding of a tensor program, to
115 /// get a handle it can run. The call decodes the bytes, compiles the program for the CPU, and
116 /// loads the result into a free slot. Bad bytes return `error.InvalidProgram` or
117 /// `error.UnsupportedVersion`, a compile or load failure returns `error.CompileFailed` or
118 /// `error.OutOfMemory`, and a full table returns `error.TooManyPrograms`. The caller may free
119 /// `bytes` once the call returns, because decoding copies names and payloads into the program.
120 pub fn compile(self: *Session, bytes: []const u8) Error!Handle {
121 const index = self.freeIndex() orelse return error.TooManyPrograms;
122 var program = tensor.wire.decode(self.allocator, bytes) catch |err| {
123 return self.fail(decodeError(err), err);
124 };
125 errdefer program.deinit();
126 const cpu = self.cpu.handle();
127 const options = accy.executable.FragmentCompilerOptions{ .artifact_format = .cpu_object };
128 const compiled = tensor.lower.compileFragment(
129 self.allocator,
130 cpu,
131 &program,
132 options,
133 ) catch |err| return self.fail(compileError(err), err);
134 const fragment = accy.executable.loadFragment(
135 self.allocator,
136 cpu,
137 compiled,
138 options,
139 ) catch |err| return self.fail(compileError(err), err);
140 errdefer fragment.deinit();
141 const capabilities = cpu.queryCapabilities() catch |err| {
142 return self.fail(compileError(err), err);
143 };
144 const min_alignment = capabilities.memory.min_buffer_alignment;
145 const artifact_plan = compiled.artifactPlan();
146 const descriptors = try describe(self.allocator, &program, artifact_plan, min_alignment);
147 std.debug.assert(descriptors.len == program.parameters.len + program.outputs.len);
148
149 const entry = &self.entries[index];
150 std.debug.assert(!entry.live);
151 entry.* = .{
152 .live = true,
153 .generation = entry.generation,
154 .program = program,
155 .compiled = compiled,
156 .fragment = fragment,
157 .descriptors = descriptors,
158 .input_count = program.parameters.len,
159 };
160 return .{ .index = index, .generation = entry.generation };
161 }
162
163 pub fn inputDescriptors(self: *const Session, handle: Handle) Error![]const Descriptor {
164 const entry = &self.entries[try self.indexOf(handle)];
165 return entry.descriptors[0..entry.input_count];
166 }
167
168 pub fn outputDescriptors(self: *const Session, handle: Handle) Error![]const Descriptor {
169 const entry = &self.entries[try self.indexOf(handle)];
170 return entry.descriptors[entry.input_count..];
171 }
172
173 /// A caller runs a compiled program on its own input and output buffers. The input and output
174 /// slices follow the program's descriptors in order, and a count that does not match returns
175 /// `error.InvalidBuffer`. The run uses the caller's buffers in place wherever the plans allow,
176 /// and copies any input the program writes so the caller's inputs stay unchanged. Each output
177 /// must share no bytes with any input or with another output, and a buffer that breaks this or
178 /// the descriptor's size or alignment returns `error.InvalidBuffer`. `scratch` backs the
179 /// short-lived state of one launch, and a launch failure returns `error.LaunchFailed`.
180 pub fn invoke(
181 self: *Session,
182 handle: Handle,
183 scratch: std.mem.Allocator,
184 input_bytes: []const []const u8,
185 output_bytes: []const []u8,
186 ) Error!void {
187 const entry = &self.entries[try self.indexOf(handle)];
188 std.debug.assert(entry.live);
189 const output_count = entry.descriptors.len - entry.input_count;
190 if (input_bytes.len != entry.input_count) return error.InvalidBuffer;
191 if (output_bytes.len != output_count) return error.InvalidBuffer;
192 const bindings = binding.prepareBorrowed(
193 self.allocator,
194 self.cpu.handle(),
195 entry.compiled.artifactPlan(),
196 entry.compiled.launchPlan(),
197 input_bytes,
198 output_bytes,
199 ) catch |err| return self.fail(bindError(err), err);
200 defer bindings.deinit();
201 entry.fragment.submitInvocationWithOptions(scratch, bindings, .{}) catch |err| {
202 return self.fail(launchError(err), err);
203 };
204 entry.fragment.completeInvocationWithOptions(.{}) catch |err| {
205 return self.fail(launchError(err), err);
206 };
207 binding.completeBorrowed(bindings, output_bytes) catch |err| {
208 return self.fail(launchError(err), err);
209 };
210 }
211
212 /// Frees the compiled program, its descriptors and its decoded form. After the call the handle
213 /// and every descriptor slice from that program are invalid, and the handle is refused with
214 /// `error.InvalidHandle`.
215 pub fn release(self: *Session, handle: Handle) Error!void {
216 const entry = &self.entries[try self.indexOf(handle)];
217 self.releaseEntry(entry);
218 }
219
220 fn releaseEntry(self: *Session, entry: *Entry) void {
221 std.debug.assert(entry.live);
222 entry.fragment.deinit();
223 self.allocator.free(entry.descriptors);
224 entry.program.deinit();
225 entry.* = .{ .generation = entry.generation +% 1 };
226 }
227
228 fn indexOf(self: *const Session, handle: Handle) Error!u32 {
229 if (handle.index >= max_programs) return error.InvalidHandle;
230 const entry = &self.entries[handle.index];
231 if (!entry.live or entry.generation != handle.generation) return error.InvalidHandle;
232 return handle.index;
233 }
234
235 fn freeIndex(self: *const Session) ?u32 {
236 for (&self.entries, 0..) |*entry, index| {
237 if (!entry.live) return @intCast(index);
238 }
239 return null;
240 }
241
242 fn fail(self: *Session, public: Error, cause: anyerror) Error {
243 self.last_cause = cause;
244 return public;
245 }
246 };
247
248 fn describe(
249 allocator: std.mem.Allocator,
250 program: *const Program,
251 artifact_plan: *const BackendArtifactPlan,
252 min_alignment: u32,
253 ) Error![]Descriptor {
254 const input_count = program.parameters.len;
255 if (artifact_plan.input_slot_ids.len != input_count) return error.CompileFailed;
256 if (artifact_plan.output_slot_ids.len != program.outputs.len) return error.CompileFailed;
257 const count = input_count + program.outputs.len;
258 const descriptors = allocator.alloc(Descriptor, count) catch return error.OutOfMemory;
259 errdefer allocator.free(descriptors);
260
261 const inputs = descriptors[0..input_count];
262 for (program.parameters, artifact_plan.input_slot_ids, inputs) |id, slot_id, *descriptor| {
263 const ty = program.values[id.index];
264 descriptor.* = try describeSlot(ty, artifact_plan, slot_id, min_alignment);
265 }
266 const outputs = descriptors[input_count..];
267 for (program.outputs, artifact_plan.output_slot_ids, outputs) |id, slot_id, *descriptor| {
268 const ty = program.values[id.index];
269 descriptor.* = try describeSlot(ty, artifact_plan, slot_id, min_alignment);
270 }
271 return descriptors;
272 }
273
274 fn describeSlot(
275 ty: tensor.Type,
276 artifact_plan: *const BackendArtifactPlan,
277 slot_id: usize,
278 min_alignment: u32,
279 ) Error!Descriptor {
280 const slot = artifact_plan.slotById(slot_id) orelse return error.CompileFailed;
281 const slot_bytes = slot.byte_size orelse return error.CompileFailed;
282 const byte_size = ty.byteCount() catch return error.CompileFailed;
283 if (slot_bytes != byte_size) return error.CompileFailed;
284 const alignment = @max(slot.alignment, min_alignment, 1);
285 return .{
286 .dtype = ty.dtype,
287 .dims = ty.dims,
288 .byte_size = byte_size,
289 .alignment = std.math.cast(u32, alignment) orelse return error.CompileFailed,
290 };
291 }
292
293 fn decodeError(err: anyerror) Error {
294 return switch (err) {
295 error.OutOfMemory => error.OutOfMemory,
296 error.UnsupportedVersion => error.UnsupportedVersion,
297 else => error.InvalidProgram,
298 };
299 }
300
301 fn compileError(err: anyerror) Error {
302 return switch (err) {
303 error.OutOfMemory => error.OutOfMemory,
304 else => error.CompileFailed,
305 };
306 }
307
308 fn bindError(err: anyerror) Error {
309 return switch (err) {
310 error.OutOfMemory => error.OutOfMemory,
311 error.InvalidBuffer, error.ReadBufferDestinationTooSmall => error.InvalidBuffer,
312 error.CapabilityMismatch => error.InvalidBuffer,
313 else => error.LaunchFailed,
314 };
315 }
316
317 fn launchError(err: anyerror) Error {
318 return switch (err) {
319 error.OutOfMemory => error.OutOfMemory,
320 else => error.LaunchFailed,
321 };
322 }