lib/accy/src/kernel/program/execution.zig
daab053ee43316e1809a84551d573ddd1e5bf3d2
1 const std = @import("std");
2 const choir = @import("choir");
3
4 const core = @import("../model/core/root.zig");
5 const model = @import("../model/program/root.zig");
6 const plan = @import("../model/plan/root.zig");
7 const oracle = @import("../oracle/root.zig");
8
9 const Cpu = struct {
10 machine: ?oracle.Machine = null,
11 verified: bool = false,
12
13 fn deinit(self: *Cpu) void {
14 if (self.machine) |*machine| machine.deinit();
15 self.* = undefined;
16 }
17
18 fn run(
19 self: *Cpu,
20 program: *model.Program,
21 allocator: std.mem.Allocator,
22 args: []const oracle.Argument,
23 launch_value: core.schedule.Launch,
24 ) !void {
25 const machine = try self.ensureMachine(allocator);
26 if (!self.verified) {
27 try program.verify();
28 self.verified = true;
29 }
30 try machine.run(program.storage.kernel.func().op, args, launch_value);
31 }
32
33 fn runWithDiagnostic(
34 self: *Cpu,
35 program: *model.Program,
36 allocator: std.mem.Allocator,
37 args: []const oracle.Argument,
38 diagnostic: *oracle.ExecutionDiagnostic,
39 ) !void {
40 diagnostic.len = 0;
41 const machine = try self.ensureMachine(allocator);
42 if (!self.verified) {
43 try program.verifyWithDiagnostic(diagnostic);
44 self.verified = true;
45 }
46 try machine.runWithDiagnostic(program.storage.kernel.func().op, args, try program.launch(), diagnostic);
47 }
48
49 fn ensureMachine(self: *Cpu, allocator: std.mem.Allocator) !*oracle.Machine {
50 if (self.machine) |*machine| {
51 const same_ptr = machine.allocator.ptr == allocator.ptr;
52 const same_vtable = machine.allocator.vtable == allocator.vtable;
53 if (same_ptr and same_vtable) return machine;
54 machine.deinit();
55 self.machine = null;
56 self.verified = false;
57 }
58 self.machine = oracle.Machine.init(allocator);
59 return &self.machine.?;
60 }
61 };
62
63 pub const Executor = struct {
64 program: *model.Program,
65 cpu: Cpu = .{},
66
67 pub fn init(program: *model.Program) Executor {
68 return .{ .program = program };
69 }
70
71 pub fn deinit(self: *Executor) void {
72 self.cpu.deinit();
73 self.* = undefined;
74 }
75
76 pub fn runCpu(self: *Executor, allocator: std.mem.Allocator, args: []const oracle.Argument) !void {
77 try self.runCpuWithLaunch(allocator, args, try self.program.launch());
78 }
79
80 pub fn runCpuWithLaunch(
81 self: *Executor,
82 allocator: std.mem.Allocator,
83 args: []const oracle.Argument,
84 launch_value: core.schedule.Launch,
85 ) !void {
86 try self.cpu.run(self.program, allocator, args, launch_value);
87 }
88
89 pub fn runCpuWithDiagnostic(
90 self: *Executor,
91 allocator: std.mem.Allocator,
92 args: []const oracle.Argument,
93 diagnostic: *oracle.ExecutionDiagnostic,
94 ) !void {
95 try self.cpu.runWithDiagnostic(self.program, allocator, args, diagnostic);
96 }
97 };
98
99 pub const Program = struct {
100 capacity: model.Capacity,
101 model_program: model.Program,
102 cpu: Cpu = .{},
103
104 const Self = @This();
105
106 pub fn init(
107 kernel: core.builder.Kernel,
108 schedule: core.schedule.Schedule,
109 capacity: model.Capacity,
110 ) Self {
111 return .{
112 .capacity = capacity,
113 .model_program = model.Program.init(kernel, schedule, capacity),
114 };
115 }
116
117 pub fn deinit(self: *Self) void {
118 self.cpu.deinit();
119 self.model_program.deinit();
120 self.* = undefined;
121 }
122
123 pub fn verify(self: *Self) !void {
124 try self.model_program.verify();
125 }
126
127 pub fn verifyWithDiagnostic(self: *Self, diagnostic: *core.builder.VerificationDiagnostic) !void {
128 try self.model_program.verifyWithDiagnostic(diagnostic);
129 }
130
131 pub fn launch(self: *const Self) core.schedule.ScheduleError!core.schedule.Launch {
132 return self.model_program.launch();
133 }
134
135 pub fn scheduleSnapshot(self: *const Self, allocator: std.mem.Allocator) core.schedule.ScheduleError!core.schedule.Snapshot {
136 return self.model_program.scheduleSnapshot(allocator);
137 }
138
139 pub fn createPlan(self: *Self, allocator: std.mem.Allocator, options: plan.Options) plan.Error!plan.Plan {
140 return self.model_program.createPlan(allocator, options);
141 }
142
143 pub fn kernelModule(self: *const Self) *choir.ir.Operation {
144 return self.model_program.kernelModule();
145 }
146
147 pub fn params(self: *const Self) []const core.builder.Param {
148 return self.model_program.params();
149 }
150
151 pub fn createCheckedPlan(self: *Self, allocator: std.mem.Allocator, options: plan.Options) !plan.Plan {
152 return self.model_program.createCheckedPlan(allocator, options);
153 }
154
155 pub fn checkPlan(self: *Self, allocator: std.mem.Allocator, authored_plan: *const plan.Plan) !void {
156 try self.model_program.checkPlan(allocator, authored_plan);
157 }
158
159 pub fn runCpu(self: *Self, allocator: std.mem.Allocator, args: []const oracle.Argument) !void {
160 try self.runCpuWithLaunch(allocator, args, try self.launch());
161 }
162
163 pub fn runCpuWithLaunch(
164 self: *Self,
165 allocator: std.mem.Allocator,
166 args: []const oracle.Argument,
167 launch_value: core.schedule.Launch,
168 ) !void {
169 try self.cpu.run(&self.model_program, allocator, args, launch_value);
170 }
171
172 pub fn runCpuWithDiagnostic(
173 self: *Self,
174 allocator: std.mem.Allocator,
175 args: []const oracle.Argument,
176 diagnostic: *oracle.ExecutionDiagnostic,
177 ) !void {
178 try self.cpu.runWithDiagnostic(&self.model_program, allocator, args, diagnostic);
179 }
180
181 pub fn bodyFingerprint(self: *const Self, allocator: std.mem.Allocator) !u64 {
182 return self.model_program.bodyFingerprint(allocator);
183 }
184
185 pub fn scheduleFingerprint(self: *const Self) u64 {
186 return self.model_program.scheduleFingerprint();
187 }
188
189 pub fn fingerprint(self: *const Self, allocator: std.mem.Allocator) !choir.product.incremental.Fingerprint {
190 return self.model_program.fingerprint(allocator);
191 }
192
193 pub fn productStamp(self: *const Self, allocator: std.mem.Allocator) !choir.product.incremental.ProductStamp {
194 return self.model_program.productStamp(allocator);
195 }
196 };
197
198 const surface = model.Surface(Program, .{
199 .registrations = &choir.backends.gpu.dialect_registrations,
200 .preload = &.{"scf"},
201 });
202
203 pub const Index1D = surface.Index1D;
204 pub const VectorIndex1D = surface.VectorIndex1D;
205 pub const DomainAxis = surface.DomainAxis;
206 pub const Domain2D = surface.Domain2D;
207 pub const Domain3D = surface.Domain3D;
208 pub const Index2D = surface.Index2D;
209 pub const Index3D = surface.Index3D;
210 pub const Vec2 = surface.Vec2;
211 pub const Vec3 = surface.Vec3;
212 pub const TypedValue = surface.TypedValue;
213 pub const TypedVec2 = surface.TypedVec2;
214 pub const TypedVec3 = surface.TypedVec3;
215 pub const domainAxis = surface.domainAxis;
216 pub const BufferView = surface.BufferView;
217 pub const Builder = surface.Builder;
218 pub const Limits = surface.Limits;
219 pub const Capacity = surface.Capacity;
220 pub const Guard = surface.Guard;
221 pub const define = surface.define;
222 pub const defineWithDiagnostic = surface.defineWithDiagnostic;
223
224 fn buildCopyProgram(
225 allocator: std.mem.Allocator,
226 name: []const u8,
227 extent: u64,
228 ) !Program {
229 var builder = try Builder.init(allocator, Limits.testing, name, &.{
230 core.builder.dynamicBuffer(.i32),
231 core.builder.dynamicBuffer(.i32),
232 });
233 errdefer builder.deinit();
234
235 const axis = try builder.axis("i", extent);
236 try builder.bind(axis, .thread_x);
237 const src = builder.argument(0);
238 const dst = builder.argument(1);
239 const index = try builder.globalId(.x);
240 const value = try builder.load(src, index);
241 try builder.store(value, dst, index);
242 try builder.return_();
243 return builder.finish();
244 }
245
246 test "Program steady oracle runs make no allocator calls" {
247 var program = try buildCopyProgram(std.testing.allocator, "machine_retained_copy", 4);
248 defer program.deinit();
249
250 var failing = std.testing.FailingAllocator.init(std.testing.allocator, .{});
251 var input = [_]i32{ 3, -1, 7, 12 };
252 var output = [_]i32{ 0, 0, 0, 0 };
253 try program.runCpu(failing.allocator(), &.{
254 oracle.argumentBuffer(i32, input[0..]),
255 oracle.argumentBuffer(i32, output[0..]),
256 });
257 try std.testing.expectEqualSlices(i32, input[0..], output[0..]);
258
259 failing.fail_index = failing.alloc_index;
260 failing.resize_fail_index = failing.resize_index;
261 for (0..8) |_| {
262 @memset(output[0..], 0);
263 try program.runCpu(failing.allocator(), &.{
264 oracle.argumentBuffer(i32, input[0..]),
265 oracle.argumentBuffer(i32, output[0..]),
266 });
267 try std.testing.expectEqualSlices(i32, input[0..], output[0..]);
268 }
269 try std.testing.expect(!failing.has_induced_failure);
270 }
271
272 test "Program reacquires the machine when the run allocator changes" {
273 var program = try buildCopyProgram(std.testing.allocator, "machine_epoch_copy", 4);
274 defer program.deinit();
275
276 var failing = std.testing.FailingAllocator.init(std.testing.allocator, .{});
277 var input = [_]i32{ 5, 6, 7, 8 };
278 var output = [_]i32{ 0, 0, 0, 0 };
279 try program.runCpu(failing.allocator(), &.{
280 oracle.argumentBuffer(i32, input[0..]),
281 oracle.argumentBuffer(i32, output[0..]),
282 });
283 try std.testing.expectEqualSlices(i32, input[0..], output[0..]);
284
285 @memset(output[0..], 0);
286 try program.runCpu(std.testing.allocator, &.{
287 oracle.argumentBuffer(i32, input[0..]),
288 oracle.argumentBuffer(i32, output[0..]),
289 });
290 try std.testing.expectEqualSlices(i32, input[0..], output[0..]);
291 }