lib/accy/src/profiling/reaction/device.zig
daab053ee43316e1809a84551d573ddd1e5bf3d2
1 const std = @import("std");
2 const gpu = @import("gpu");
3 const choir_abi = @import("choir_abi");
4 const accy = @import("accy");
5 const sys = @import("sys");
6
7 const kernel_mod = @import("kernel.zig");
8 const options_mod = @import("options.zig");
9
10 const Allocator = std.mem.Allocator;
11
12 pub const State = struct {
13 handle: gpu.BackendHandle,
14 artifact: accy.kernel.KernelArtifact,
15 loaded: gpu.LoadedArtifact,
16 fields: [2][2]gpu.BufferBinding,
17 geometry: choir_abi.LaunchGeometry,
18 compile_ns: u64,
19 load_ns: u64,
20 current: u1 = 0,
21
22 pub fn deinit(self: *State) void {
23 for (self.fields) |pair| {
24 for (pair) |binding| self.handle.destroyObject(binding.handle.id);
25 }
26 self.handle.destroyObject(self.loaded.id);
27 self.artifact.deinit();
28 }
29
30 pub fn ptxBytes(self: *const State) usize {
31 return switch (self.artifact.payload) {
32 .text => |text| text.len,
33 else => 0,
34 };
35 }
36
37 pub fn step(self: *State) !void {
38 const source = self.fields[self.current];
39 const target = self.fields[self.current ^ 1];
40 try self.handle.launch(.{
41 .artifact = &self.artifact,
42 .loaded_artifact = self.loaded,
43 .buffers = &.{ source[0], source[1], target[0], target[1] },
44 .geometry = self.geometry,
45 });
46 self.current ^= 1;
47 }
48
49 pub fn readU(self: *State, host: []f32) !void {
50 try self.readField(0, host);
51 }
52
53 pub fn readV(self: *State, host: []f32) !void {
54 try self.readField(1, host);
55 }
56
57 fn readField(self: *State, slot: u1, host: []f32) !void {
58 try self.handle.synchronize(.{ .scope = .device });
59 try self.handle.readBuffer(.{
60 .handle = self.fields[self.current][slot].handle,
61 .bytes = std.mem.sliceAsBytes(host),
62 });
63 }
64 };
65
66 pub fn init(backing: Allocator, handle: gpu.BackendHandle, side: u32, u: []const f32, v: []const f32) !State {
67 const compile_start = sys.time.nanoTimestamp();
68 var graph = try kernel_mod.buildStepGraph(backing, side);
69 defer graph.deinit();
70 var artifact = try accy.kernel.createKernelArtifact(backing, handle, &graph, .{
71 .authored_kernel_diagnostic_id = "profiling/reaction-step",
72 });
73 errdefer artifact.deinit();
74 const compile_ns: u64 = @intCast(sys.time.nanoTimestamp() - compile_start);
75
76 const load_start = sys.time.nanoTimestamp();
77 const loaded = try handle.loadArtifact(&artifact);
78 errdefer handle.destroyObject(loaded.id);
79 const load_ns: u64 = @intCast(sys.time.nanoTimestamp() - load_start);
80
81 const cells = @as(usize, side) * side;
82 var fields: [2][2]gpu.BufferBinding = undefined;
83 var created: usize = 0;
84 errdefer {
85 var index: usize = 0;
86 while (index < created) : (index += 1) {
87 handle.destroyObject(fields[index / 2][index % 2].handle.id);
88 }
89 }
90 for (0..2) |pair| {
91 for (0..2) |slot| {
92 fields[pair][slot] = try allocateFieldBuffer(handle, cells);
93 created += 1;
94 }
95 }
96 try handle.writeBuffer(.{ .handle = fields[0][0].handle, .bytes = std.mem.sliceAsBytes(u) });
97 try handle.writeBuffer(.{ .handle = fields[0][1].handle, .bytes = std.mem.sliceAsBytes(v) });
98
99 const blocks: u32 = @intCast(cells / options_mod.threads_per_block);
100 return .{
101 .handle = handle,
102 .artifact = artifact,
103 .loaded = loaded,
104 .fields = fields,
105 .geometry = .{
106 .grid = .{ blocks, 1, 1 },
107 .threadgroup = .{ options_mod.threads_per_block, 1, 1 },
108 },
109 .compile_ns = compile_ns,
110 .load_ns = load_ns,
111 };
112 }
113
114 pub fn oracleMaxAbsError(device_u: []const f32, device_v: []const f32, u: []const f32, v: []const f32, scratch: [][]f32, side: u32) f32 {
115 kernel_mod.referenceStep(u, v, scratch[0], scratch[1], side);
116 var max_error: f32 = 0.0;
117 for (device_u, scratch[0]) |device_value, host_value| {
118 max_error = @max(max_error, @abs(device_value - host_value));
119 }
120 for (device_v, scratch[1]) |device_value, host_value| {
121 max_error = @max(max_error, @abs(device_value - host_value));
122 }
123 return max_error;
124 }
125
126 fn allocateFieldBuffer(handle: gpu.BackendHandle, element_count: usize) !gpu.BufferBinding {
127 const device_buffer = try handle.allocateBuffer(.{
128 .byte_size = element_count * @sizeOf(f32),
129 .alignment = 256,
130 .dtype = .f32,
131 .element_count = element_count,
132 });
133 return .{
134 .handle = device_buffer,
135 .access = .read_write,
136 .ownership = device_buffer.ownership,
137 .byte_size = device_buffer.byte_size,
138 };
139 }
140
141 test "oracleMaxAbsError is zero against the reference itself and catches drift in either field" {
142 const side: u32 = 32;
143 const cells = @as(usize, side) * side;
144 const u = try std.testing.allocator.alloc(f32, cells);
145 defer std.testing.allocator.free(u);
146 const v = try std.testing.allocator.alloc(f32, cells);
147 defer std.testing.allocator.free(v);
148 kernel_mod.seedState(u, v, side);
149
150 var scratch: [2][]f32 = undefined;
151 for (&scratch) |*buffer| buffer.* = try std.testing.allocator.alloc(f32, cells);
152 defer for (scratch) |buffer| std.testing.allocator.free(buffer);
153 const mirror_u = try std.testing.allocator.alloc(f32, cells);
154 defer std.testing.allocator.free(mirror_u);
155 const mirror_v = try std.testing.allocator.alloc(f32, cells);
156 defer std.testing.allocator.free(mirror_v);
157
158 kernel_mod.referenceStep(u, v, mirror_u, mirror_v, side);
159 const stepped_u = try std.testing.allocator.dupe(f32, mirror_u);
160 defer std.testing.allocator.free(stepped_u);
161 const stepped_v = try std.testing.allocator.dupe(f32, mirror_v);
162 defer std.testing.allocator.free(stepped_v);
163
164 try std.testing.expectApproxEqAbs(
165 @as(f32, 0.0),
166 oracleMaxAbsError(stepped_u, stepped_v, u, v, scratch[0..], side),
167 0.0000001,
168 );
169
170 stepped_u[cells / 2] += 0.25;
171 try std.testing.expect(oracleMaxAbsError(stepped_u, stepped_v, u, v, scratch[0..], side) > 0.2);
172 stepped_u[cells / 2] -= 0.25;
173 stepped_v[cells / 3] += 0.25;
174 try std.testing.expect(oracleMaxAbsError(stepped_u, stepped_v, u, v, scratch[0..], side) > 0.2);
175 }