lib/accy/src/profiling/sph/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
6 const kernel_mod = @import("kernel.zig");
7
8 const Allocator = std.mem.Allocator;
9 const artifact_mod = accy.artifact;
10 const executable = accy.executable;
11 const kernel = accy.kernel;
12 const spatial = kernel.library.spatial;
13 const sort_library = kernel.library.sort;
14 const scan_library = kernel.library.scan;
15
16 const count: usize = kernel_mod.particle_count;
17 const threads: u32 = kernel_mod.threads_per_block;
18
19 pub const geometry = spatial.GridGeometry{
20 .origin_x = kernel_mod.domain_min,
21 .origin_y = kernel_mod.domain_min,
22 .inv_cell_size = 1.0 / kernel_mod.cell_size,
23 .dims_x = kernel_mod.grid_dims,
24 .dims_y = kernel_mod.grid_dims,
25 };
26
27 const Authored = struct {
28 artifact: kernel.KernelArtifact,
29 loaded: gpu.LoadedArtifact,
30
31 fn init(
32 allocator: Allocator,
33 handle: gpu.BackendHandle,
34 graph: *kernel.Graph,
35 diagnostic_id: []const u8,
36 ) !Authored {
37 var artifact = try kernel.createKernelArtifact(allocator, handle, graph, .{
38 .artifact_format = .cuda_ptx,
39 .authored_kernel_diagnostic_id = diagnostic_id,
40 });
41 errdefer artifact.deinit();
42 const loaded = try handle.loadArtifact(&artifact);
43 return .{ .artifact = artifact, .loaded = loaded };
44 }
45
46 fn deinit(self: *Authored, handle: gpu.BackendHandle) void {
47 handle.destroyObject(self.loaded.id);
48 self.artifact.deinit();
49 }
50 };
51
52 pub const State = struct {
53 allocator: Allocator,
54 handle: gpu.BackendHandle,
55 plan: spatial.GridBuildPlan,
56
57 density_kernel: Authored,
58 force_kernel: Authored,
59 integrate_kernel: Authored,
60 cells_kernel: Authored,
61 count_kernel: Authored,
62
63 sort_artifacts: sort_library.RadixDigitPairsPipelineArtifacts,
64 sort_entries: [5]artifact_mod.KernelCallArtifact,
65 sort_pipeline: artifact_mod.OwnedKernelCallPipeline,
66 sort_pool: []gpu.BufferBinding,
67 sort_artifact_pool: executable.PipelineArtifactPool,
68
69 scan_artifacts: scan_library.DeviceScanPipelineArtifacts,
70 scan_entries: [3]artifact_mod.KernelCallArtifact,
71 scan_pipeline: artifact_mod.OwnedKernelCallPipeline,
72 scan_pool: []gpu.BufferBinding,
73 scan_artifact_pool: executable.PipelineArtifactPool,
74
75 positions_x: [2]gpu.BufferHandle,
76 positions_y: [2]gpu.BufferHandle,
77 velocities_x: [2]gpu.BufferHandle,
78 velocities_y: [2]gpu.BufferHandle,
79 forces_x: gpu.BufferHandle,
80 forces_y: gpu.BufferHandle,
81 density: gpu.BufferHandle,
82 ids: [2]gpu.BufferHandle,
83 payloads: [2]gpu.BufferHandle,
84 counts: gpu.BufferHandle,
85 scanned: gpu.BufferHandle,
86
87 iota: []i32,
88 current: u1 = 0,
89 sorted_slot: usize = 0,
90
91 pub fn deinit(self: *State) void {
92 const handle = self.handle;
93 self.density_kernel.deinit(handle);
94 self.force_kernel.deinit(handle);
95 self.integrate_kernel.deinit(handle);
96 self.cells_kernel.deinit(handle);
97 self.count_kernel.deinit(handle);
98
99 executable.deinitPipelineIntermediates(self.allocator, handle, self.sort_pool);
100 self.sort_artifact_pool.deinit();
101 self.sort_pipeline.deinit();
102 self.sort_artifacts.deinit();
103
104 executable.deinitPipelineIntermediates(self.allocator, handle, self.scan_pool);
105 self.scan_artifact_pool.deinit();
106 self.scan_pipeline.deinit();
107 self.scan_artifacts.deinit();
108
109 for (self.positions_x) |buffer| handle.destroyObject(buffer.id);
110 for (self.positions_y) |buffer| handle.destroyObject(buffer.id);
111 for (self.velocities_x) |buffer| handle.destroyObject(buffer.id);
112 for (self.velocities_y) |buffer| handle.destroyObject(buffer.id);
113 handle.destroyObject(self.forces_x.id);
114 handle.destroyObject(self.forces_y.id);
115 handle.destroyObject(self.density.id);
116 for (self.ids) |buffer| handle.destroyObject(buffer.id);
117 for (self.payloads) |buffer| handle.destroyObject(buffer.id);
118 handle.destroyObject(self.counts.id);
119 handle.destroyObject(self.scanned.id);
120 self.allocator.free(self.iota);
121 }
122
123 pub fn residentBytes(self: *const State) usize {
124 var total: usize = 0;
125 for (self.positions_x) |buffer| addBufferBytes(&total, buffer);
126 for (self.positions_y) |buffer| addBufferBytes(&total, buffer);
127 for (self.velocities_x) |buffer| addBufferBytes(&total, buffer);
128 for (self.velocities_y) |buffer| addBufferBytes(&total, buffer);
129 addBufferBytes(&total, self.forces_x);
130 addBufferBytes(&total, self.forces_y);
131 addBufferBytes(&total, self.density);
132 for (self.ids) |buffer| addBufferBytes(&total, buffer);
133 for (self.payloads) |buffer| addBufferBytes(&total, buffer);
134 addBufferBytes(&total, self.counts);
135 addBufferBytes(&total, self.scanned);
136 for (self.sort_pool) |binding| addByteSize(&total, binding.byte_size);
137 for (self.scan_pool) |binding| addByteSize(&total, binding.byte_size);
138 return total;
139 }
140
141 pub fn neighborEvidence(self: *State, allocator: Allocator) !NeighborEvidence {
142 const instance = spatial.GridNeighborCount{
143 .count = count,
144 .offsets_extent = self.plan.offsets_scan.extent,
145 .threads = threads,
146 };
147 var graph = try spatial.GridNeighborCountRuntimeFamilyF32.build(
148 allocator,
149 spatial.GridNeighborCountRuntimeFamilyF32.Limits.standard,
150 instance,
151 );
152 defer graph.deinit();
153 var authored = try Authored.init(allocator, self.handle, &graph, "profiling/sph-neighbor-evidence");
154 defer authored.deinit(self.handle);
155
156 const neighbors = try allocateBuffer(self.handle, .i32, count);
157 defer self.handle.destroyObject(neighbors.id);
158 const bindings = [_]gpu.BufferBinding{
159 bind(neighbors, .read_write),
160 bind(self.positions_x[self.current], .read_only),
161 bind(self.positions_y[self.current], .read_only),
162 bind(self.payloads[self.sorted_slot], .read_only),
163 bind(self.scanned, .read_only),
164 };
165 const scalars = try spatial.gridNeighborCountRuntimeArguments(instance, geometry, kernel_mod.offsets_stride, kernel_mod.smoothing);
166 const blocks: u32 = @intCast(spatial.gridCellsBlockCount(count, threads));
167 try self.launchAuthored(&authored, bindings[0..], scalars[0..], blocks, threads, null);
168 try self.handle.synchronize(.{ .scope = .device });
169
170 const host = try allocator.alloc(i32, count);
171 defer allocator.free(host);
172 try self.handle.readBuffer(.{ .handle = neighbors, .bytes = std.mem.sliceAsBytes(host) });
173 var total: u64 = 0;
174 for (host) |value| {
175 std.debug.assert(value >= 0);
176 total = std.math.add(u64, total, @intCast(value)) catch unreachable;
177 }
178 return .{
179 .total = total,
180 .fingerprint = std.hash.Wyhash.hash(0, std.mem.sliceAsBytes(host)),
181 };
182 }
183
184 fn bind(buffer: gpu.BufferHandle, access: gpu.BufferAccess) gpu.BufferBinding {
185 return .{
186 .handle = buffer,
187 .access = access,
188 .ownership = buffer.ownership,
189 .byte_size = buffer.byte_size,
190 };
191 }
192
193 fn launchAuthored(
194 self: *State,
195 authored: *Authored,
196 buffers: []const gpu.BufferBinding,
197 scalars: []const choir_abi.ScalarArgument,
198 blocks: u32,
199 launch_threads: u32,
200 stream: ?gpu.StreamHandle,
201 ) !void {
202 try self.handle.launch(.{
203 .artifact = &authored.artifact,
204 .loaded_artifact = authored.loaded,
205 .buffers = buffers,
206 .scalar_arguments = scalars,
207 .geometry = .{
208 .grid = .{ blocks, 1, 1 },
209 .threadgroup = .{ launch_threads, 1, 1 },
210 },
211 .stream = stream,
212 });
213 }
214
215 pub fn buildGrid(self: *State) !void {
216 try self.buildGridOn(null);
217 }
218
219 pub fn buildGridOn(self: *State, stream: ?gpu.StreamHandle) !void {
220 const blocks: u32 = @intCast(spatial.gridCellsBlockCount(count, self.plan.cells_instance.threads));
221 {
222 const bindings = [_]gpu.BufferBinding{
223 bind(self.ids[0], .read_write),
224 bind(self.positions_x[self.current], .read_only),
225 bind(self.positions_y[self.current], .read_only),
226 };
227 const scalars = try spatial.gridCellsRuntimeArguments(self.plan.cells_instance, geometry);
228 try self.launchAuthored(&self.cells_kernel, bindings[0..], scalars[0..], blocks, self.plan.cells_instance.threads, stream);
229 }
230
231 try self.handle.writeBuffer(.{ .handle = self.payloads[0], .bytes = std.mem.sliceAsBytes(self.iota) });
232
233 var source: usize = 0;
234 var pass: u32 = 0;
235 while (pass < self.plan.sort_passes) : (pass += 1) {
236 const shift = pass * sort_library.radix_digit_bits;
237 const scalars = try sort_library.radixDigitHistogramRuntimeArguments(self.plan.sort_instance, shift);
238 const operands = [_]gpu.BufferBinding{
239 bind(self.ids[source], .read_write),
240 bind(self.payloads[source], .read_write),
241 };
242 const results = [_]gpu.BufferBinding{
243 bind(self.ids[1 - source], .read_write),
244 bind(self.payloads[1 - source], .read_write),
245 };
246 try executable.launchPipeline(self.allocator, self.handle, .{
247 .pipeline = self.sort_pipeline.value,
248 .registry = .{ .entries = self.sort_entries[0..] },
249 .format = .cuda_ptx,
250 .operands = operands[0..],
251 .results = results[0..],
252 .runtime_scalar_arguments = scalars[0..],
253 .intermediates = self.sort_pool,
254 .artifacts = &self.sort_artifact_pool,
255 .stream = stream,
256 });
257 source = 1 - source;
258 }
259 self.sorted_slot = source;
260
261 {
262 const bindings = [_]gpu.BufferBinding{
263 bind(self.counts, .read_write),
264 bind(self.ids[self.sorted_slot], .read_only),
265 };
266 const scalars = try spatial.gridCountRuntimeArguments(self.plan.count_instance, self.plan.cells_total);
267 try self.launchAuthored(&self.count_kernel, bindings[0..], scalars[0..], blocks, self.plan.count_instance.threads, stream);
268 }
269
270 {
271 const operands = [_]gpu.BufferBinding{bind(self.counts, .read_only)};
272 const results = [_]gpu.BufferBinding{bind(self.scanned, .read_write)};
273 try executable.launchPipeline(self.allocator, self.handle, .{
274 .pipeline = self.scan_pipeline.value,
275 .registry = .{ .entries = self.scan_entries[0..] },
276 .format = .cuda_ptx,
277 .operands = operands[0..],
278 .results = results[0..],
279 .runtime_scalar_arguments = &.{.{ .u32 = @intCast(self.plan.offsets_scan.extent) }},
280 .intermediates = self.scan_pool,
281 .artifacts = &self.scan_artifact_pool,
282 .stream = stream,
283 });
284 }
285 }
286
287 pub fn gather(self: *State) !void {
288 try self.computeDensityOn(null);
289 try self.computeForceOn(null);
290 }
291
292 pub fn computeDensityOn(self: *State, stream: ?gpu.StreamHandle) !void {
293 const bindings = [_]gpu.BufferBinding{
294 bind(self.density, .read_write),
295 bind(self.positions_x[self.current], .read_only),
296 bind(self.positions_y[self.current], .read_only),
297 bind(self.payloads[self.sorted_slot], .read_only),
298 bind(self.scanned, .read_only),
299 };
300 try self.launchAuthored(&self.density_kernel, bindings[0..], &.{}, kernel_mod.blocks_per_launch, threads, stream);
301 }
302
303 pub fn computeForceOn(self: *State, stream: ?gpu.StreamHandle) !void {
304 const bindings = [_]gpu.BufferBinding{
305 bind(self.forces_x, .read_write),
306 bind(self.forces_y, .read_write),
307 bind(self.positions_x[self.current], .read_only),
308 bind(self.positions_y[self.current], .read_only),
309 bind(self.velocities_x[self.current], .read_only),
310 bind(self.velocities_y[self.current], .read_only),
311 bind(self.density, .read_only),
312 bind(self.payloads[self.sorted_slot], .read_only),
313 bind(self.scanned, .read_only),
314 };
315 try self.launchAuthored(&self.force_kernel, bindings[0..], &.{}, kernel_mod.blocks_per_launch, threads, stream);
316 }
317
318 pub fn integrate(self: *State) !void {
319 try self.integrateOn(null);
320 }
321
322 pub fn integrateOn(self: *State, stream: ?gpu.StreamHandle) !void {
323 const next = self.current ^ 1;
324 const bindings = [_]gpu.BufferBinding{
325 bind(self.positions_x[next], .read_write),
326 bind(self.positions_y[next], .read_write),
327 bind(self.velocities_x[next], .read_write),
328 bind(self.velocities_y[next], .read_write),
329 bind(self.positions_x[self.current], .read_only),
330 bind(self.positions_y[self.current], .read_only),
331 bind(self.velocities_x[self.current], .read_only),
332 bind(self.velocities_y[self.current], .read_only),
333 bind(self.forces_x, .read_only),
334 bind(self.forces_y, .read_only),
335 bind(self.density, .read_only),
336 };
337 try self.launchAuthored(&self.integrate_kernel, bindings[0..], &.{}, kernel_mod.blocks_per_launch, threads, stream);
338 self.current = next;
339 }
340
341 pub fn step(self: *State) !void {
342 try self.buildGrid();
343 try self.gather();
344 try self.integrate();
345 }
346
347 pub fn readBuffer(self: *State, buffer: gpu.BufferHandle, host: []f32) !void {
348 try self.handle.synchronize(.{ .scope = .device });
349 try self.handle.readBuffer(.{ .handle = buffer, .bytes = std.mem.sliceAsBytes(host) });
350 }
351
352 pub fn readState(self: *State, xs: []f32, ys: []f32, vxs: []f32, vys: []f32) !void {
353 try self.readBuffer(self.positions_x[self.current], xs);
354 try self.readBuffer(self.positions_y[self.current], ys);
355 try self.readBuffer(self.velocities_x[self.current], vxs);
356 try self.readBuffer(self.velocities_y[self.current], vys);
357 }
358
359 pub fn readGather(self: *State, densities: []f32, fxs: []f32, fys: []f32) !void {
360 try self.readBuffer(self.density, densities);
361 try self.readBuffer(self.forces_x, fxs);
362 try self.readBuffer(self.forces_y, fys);
363 }
364 };
365
366 pub const NeighborEvidence = struct {
367 total: u64,
368 fingerprint: u64,
369 };
370
371 fn addBufferBytes(total: *usize, buffer: gpu.BufferHandle) void {
372 addByteSize(total, buffer.byte_size);
373 }
374
375 fn addByteSize(total: *usize, byte_size: usize) void {
376 total.* = std.math.add(usize, total.*, byte_size) catch unreachable;
377 }
378
379 fn allocateBuffer(
380 handle: gpu.BackendHandle,
381 dtype: choir_abi.DType,
382 elements: usize,
383 ) !gpu.BufferHandle {
384 return handle.allocateBuffer(.{
385 .byte_size = elements * 4,
386 .alignment = 256,
387 .dtype = dtype,
388 .element_count = elements,
389 });
390 }
391
392 pub fn init(allocator: Allocator, handle: gpu.BackendHandle) !State {
393 const plan = try spatial.gridBuildPlan(count, geometry, threads);
394
395 var density_graph = try kernel_mod.buildDensityGraph(allocator);
396 defer density_graph.deinit();
397 var density_kernel = try Authored.init(allocator, handle, &density_graph, "profiling/sph-density");
398 errdefer density_kernel.deinit(handle);
399
400 var force_graph = try kernel_mod.buildForceGraph(allocator);
401 defer force_graph.deinit();
402 var force_kernel = try Authored.init(allocator, handle, &force_graph, "profiling/sph-force");
403 errdefer force_kernel.deinit(handle);
404
405 var integrate_graph = try kernel_mod.buildIntegrateGraph(allocator);
406 defer integrate_graph.deinit();
407 var integrate_kernel = try Authored.init(allocator, handle, &integrate_graph, "profiling/sph-integrate");
408 errdefer integrate_kernel.deinit(handle);
409
410 const cells_entry = try spatial.gridCellsFamilyEntryName(allocator, plan.cells_instance);
411 defer allocator.free(cells_entry);
412 var cells_graph = try spatial.GridCellsRuntimeFamilyF32.buildNamed(allocator, spatial.GridCellsRuntimeFamilyF32.Limits.standard, cells_entry, plan.cells_instance);
413 defer cells_graph.deinit();
414 var cells_kernel = try Authored.init(allocator, handle, &cells_graph, "profiling/sph-grid-cells");
415 errdefer cells_kernel.deinit(handle);
416
417 const count_entry = try spatial.gridCountFamilyEntryName(allocator, plan.count_instance);
418 defer allocator.free(count_entry);
419 var count_graph = try spatial.GridCountRuntimeFamilyI32.buildNamed(allocator, spatial.GridCountRuntimeFamilyI32.Limits.standard, count_entry, plan.count_instance);
420 defer count_graph.deinit();
421 var count_kernel = try Authored.init(allocator, handle, &count_graph, "profiling/sph-grid-count");
422 errdefer count_kernel.deinit(handle);
423
424 var sort_artifacts = try sort_library.createRadixDigitPairsPipelineArtifacts(allocator, handle, plan.sort_instance, .{ .limits = .standard });
425 errdefer sort_artifacts.deinit();
426 var sort_pipeline = try sort_library.radixDigitPairsPipeline(allocator, plan.sort_instance);
427 errdefer sort_pipeline.deinit();
428 const sort_args = try sort_library.radixDigitHistogramRuntimeArguments(plan.sort_instance, 0);
429 const sort_pool = try executable.allocatePipelineIntermediates(allocator, handle, sort_pipeline.value, sort_args[0..]);
430 errdefer executable.deinitPipelineIntermediates(allocator, handle, sort_pool);
431 const sort_entries = sort_artifacts.entries();
432 var sort_artifact_pool = try executable.loadPipelineArtifacts(
433 allocator,
434 handle,
435 sort_pipeline.value,
436 .{ .entries = sort_entries[0..] },
437 .cuda_ptx,
438 );
439 errdefer sort_artifact_pool.deinit();
440
441 var scan_artifacts = try scan_library.createDeviceScanPipelineArtifacts(allocator, handle, plan.offsets_scan, .{
442 .limits = kernel.Limits.standard,
443 });
444 errdefer scan_artifacts.deinit();
445 var scan_pipeline = try scan_library.deviceScanPipeline(allocator, plan.offsets_scan);
446 errdefer scan_pipeline.deinit();
447 const scan_args = [_]choir_abi.ScalarArgument{.{ .u32 = @intCast(plan.offsets_scan.extent) }};
448 const scan_pool = try executable.allocatePipelineIntermediates(allocator, handle, scan_pipeline.value, scan_args[0..]);
449 errdefer executable.deinitPipelineIntermediates(allocator, handle, scan_pool);
450 const scan_entries = scan_artifacts.entries();
451 var scan_artifact_pool = try executable.loadPipelineArtifacts(
452 allocator,
453 handle,
454 scan_pipeline.value,
455 .{ .entries = scan_entries[0..] },
456 .cuda_ptx,
457 );
458 errdefer scan_artifact_pool.deinit();
459
460 const counts_extent: usize = @intCast(plan.offsets_scan.extent);
461 const iota = try allocator.alloc(i32, count);
462 errdefer allocator.free(iota);
463 for (iota, 0..) |*value, index| value.* = @intCast(index);
464
465 var state = State{
466 .allocator = allocator,
467 .handle = handle,
468 .plan = plan,
469 .density_kernel = density_kernel,
470 .force_kernel = force_kernel,
471 .integrate_kernel = integrate_kernel,
472 .cells_kernel = cells_kernel,
473 .count_kernel = count_kernel,
474 .sort_artifacts = sort_artifacts,
475 .sort_entries = sort_entries,
476 .sort_pipeline = sort_pipeline,
477 .sort_pool = sort_pool,
478 .sort_artifact_pool = sort_artifact_pool,
479 .scan_artifacts = scan_artifacts,
480 .scan_entries = scan_entries,
481 .scan_pipeline = scan_pipeline,
482 .scan_pool = scan_pool,
483 .scan_artifact_pool = scan_artifact_pool,
484 .positions_x = undefined,
485 .positions_y = undefined,
486 .velocities_x = undefined,
487 .velocities_y = undefined,
488 .forces_x = undefined,
489 .forces_y = undefined,
490 .density = undefined,
491 .ids = undefined,
492 .payloads = undefined,
493 .counts = undefined,
494 .scanned = undefined,
495 .iota = iota,
496 };
497
498 for (&state.positions_x, &state.positions_y, &state.velocities_x, &state.velocities_y) |*px, *py, *vx, *vy| {
499 px.* = try allocateBuffer(handle, .f32, count);
500 py.* = try allocateBuffer(handle, .f32, count);
501 vx.* = try allocateBuffer(handle, .f32, count);
502 vy.* = try allocateBuffer(handle, .f32, count);
503 }
504 state.forces_x = try allocateBuffer(handle, .f32, count);
505 state.forces_y = try allocateBuffer(handle, .f32, count);
506 state.density = try allocateBuffer(handle, .f32, count);
507 for (&state.ids, &state.payloads) |*ids, *payload| {
508 ids.* = try allocateBuffer(handle, .i32, count);
509 payload.* = try allocateBuffer(handle, .i32, count);
510 }
511 state.counts = try allocateBuffer(handle, .f32, counts_extent);
512 state.scanned = try allocateBuffer(handle, .f32, counts_extent);
513
514 return state;
515 }
516
517 pub fn seed(state: *State, xs: []const f32, ys: []const f32, vxs: []const f32, vys: []const f32) !void {
518 try state.handle.writeBuffer(.{ .handle = state.positions_x[state.current], .bytes = std.mem.sliceAsBytes(xs) });
519 try state.handle.writeBuffer(.{ .handle = state.positions_y[state.current], .bytes = std.mem.sliceAsBytes(ys) });
520 try state.handle.writeBuffer(.{ .handle = state.velocities_x[state.current], .bytes = std.mem.sliceAsBytes(vxs) });
521 try state.handle.writeBuffer(.{ .handle = state.velocities_y[state.current], .bytes = std.mem.sliceAsBytes(vys) });
522 }