Moved add and bench into an examples directory so src is just the library
This commit is contained in:
adrien
@@ -0,0 +1,59 @@
|
||||
const std = @import("std");
|
||||
const gpu = @import("gpu");
|
||||
const GpuDevice = gpu.GpuDevice;
|
||||
const GpuArena = gpu.GpuArena;
|
||||
const GpuBuffer = gpu.GpuBuffer;
|
||||
const GpuProcess = gpu.GpuProcess;
|
||||
|
||||
pub fn main(init: std.process.Init) !void {
|
||||
const allocator = init.gpa;
|
||||
|
||||
// 1. Open GPU Device
|
||||
const device = try GpuDevice.init(.{});
|
||||
defer device.deinit();
|
||||
|
||||
// 2. Create a GPU Arena to manage VRAM
|
||||
var grena = GpuArena.init(allocator, device);
|
||||
defer grena.deinit();
|
||||
const gloc = grena.gpuAllocator();
|
||||
|
||||
// 3. Load the WGSL compute pipeline
|
||||
const add_process = try GpuProcess.init(device, @embedFile("shaders/add.wgsl"));
|
||||
defer add_process.deinit();
|
||||
|
||||
// 4. Setup CPU data
|
||||
const len: usize = 16;
|
||||
const data_a = try allocator.alloc(f16, len);
|
||||
defer allocator.free(data_a);
|
||||
const data_b = try allocator.alloc(f16, len);
|
||||
defer allocator.free(data_b);
|
||||
|
||||
for (0..len) |i| {
|
||||
data_a[i] = @floatFromInt(i);
|
||||
data_b[i] = @floatFromInt(len - 1 - i);
|
||||
}
|
||||
|
||||
// 5. Initialize raw GPU Buffers
|
||||
// We pass the EnumSet inline using `.initMany` since the Enum itself isn't exported
|
||||
const byte_size = len * @sizeOf(f16);
|
||||
const buf_a = try GpuBuffer.init(gloc, byte_size, .initMany(&.{ .Storage, .CopyDst, .CopySrc }));
|
||||
const buf_b = try GpuBuffer.init(gloc, byte_size, .initMany(&.{ .Storage, .CopyDst, .CopySrc }));
|
||||
const buf_out = try GpuBuffer.init(gloc, byte_size, .initMany(&.{ .Storage, .CopyDst, .CopySrc }));
|
||||
|
||||
// Note: The buffers are safely tied to the GpuArena which will automatically
|
||||
// release them at the end. You can also manually call buf_x.deinit() if desired.
|
||||
|
||||
// 6. Transfer data from CPU slices to GPU Buffers
|
||||
try buf_a.load(f16, data_a);
|
||||
try buf_b.load(f16, data_b);
|
||||
|
||||
// 7. Dispatch the Compute Process
|
||||
// We pass the data type (f16) to allow GpuProcess to calculate chunks correctly
|
||||
try add_process.run(gloc, f16, buf_a, buf_b, buf_out);
|
||||
|
||||
// 8. Map and copy the resulting buffer back to the CPU
|
||||
const out = try buf_out.read(allocator, f16);
|
||||
defer allocator.free(out);
|
||||
|
||||
std.debug.print("Result: {any}\n", .{out});
|
||||
}
|
||||
@@ -0,0 +1,180 @@
|
||||
const std = @import("std");
|
||||
const gpu = @import("gpu");
|
||||
const GpuDevice = gpu.GpuDevice;
|
||||
const GpuArena = gpu.GpuArena;
|
||||
const GpuAllocator = gpu.GpuAllocator;
|
||||
const GpuBuffer = gpu.GpuBuffer;
|
||||
const GpuProcess = gpu.GpuProcess;
|
||||
|
||||
/// Minimal implementation of a f16 Vector
|
||||
const Vec = struct {
|
||||
buf: GpuBuffer,
|
||||
len: usize,
|
||||
|
||||
// Changed: gloc is passed by value (const)
|
||||
pub fn initZero(gloc: GpuAllocator, len: usize) !Vec {
|
||||
return .{
|
||||
.buf = try GpuBuffer.init(
|
||||
gloc,
|
||||
len * @sizeOf(f16),
|
||||
.initMany(&.{ .Storage, .CopyDst, .CopySrc }),
|
||||
),
|
||||
.len = len,
|
||||
};
|
||||
}
|
||||
|
||||
// Changed: gloc is passed by value
|
||||
pub fn initLoad(gloc: GpuAllocator, data: []const f16) !Vec {
|
||||
var self = try initZero(gloc, data.len);
|
||||
try self.load(data); // Direct access via the interface copy
|
||||
return self;
|
||||
}
|
||||
|
||||
pub fn deinit(self: Vec) void {
|
||||
self.buf.deinit();
|
||||
}
|
||||
|
||||
/// CPU to GPU.
|
||||
pub fn load(self: Vec, data: []const f16) !void {
|
||||
try self.buf.load(f16, data);
|
||||
}
|
||||
|
||||
// Changed: gloc is passed by value instead of *GpuAllocator
|
||||
pub fn run(self: Vec, gloc: GpuAllocator, other: Vec, process: GpuProcess) !Vec {
|
||||
std.debug.assert(self.len == other.len);
|
||||
|
||||
const result = try Vec.initZero(gloc, self.len);
|
||||
errdefer result.deinit();
|
||||
|
||||
try process.run(gloc, f16, self.buf, other.buf, result.buf);
|
||||
return result;
|
||||
}
|
||||
|
||||
// Changed: gloc is passed by value instead of *GpuAllocator
|
||||
pub fn read(self: Vec, alloc: std.mem.Allocator) ![]f16 {
|
||||
return self.buf.read(alloc, f16);
|
||||
}
|
||||
};
|
||||
|
||||
pub fn main(init: std.process.Init) !void {
|
||||
const device = try GpuDevice.init(.{ .vram_bytes_limit = 4 * 1024 * 1024 * 1024 });
|
||||
defer device.deinit();
|
||||
|
||||
var grena = GpuArena.init(init.gpa, device);
|
||||
defer grena.deinit();
|
||||
|
||||
const gloc = grena.gpuAllocator();
|
||||
|
||||
const add_pip = try GpuProcess.init(device, @embedFile("shaders/add.wgsl"));
|
||||
defer add_pip.deinit();
|
||||
|
||||
const allocator = init.gpa;
|
||||
|
||||
const sizes = [_]usize{
|
||||
1,
|
||||
256,
|
||||
1024,
|
||||
4 * 1024,
|
||||
4 * 4 * 1024,
|
||||
4 * 4 * 4 * 1024,
|
||||
4 * 4 * 4 * 4 * 1024,
|
||||
1024 * 1024,
|
||||
4 * 1024 * 1024,
|
||||
4 * 4 * 1024 * 1024,
|
||||
4 * 4 * 4 * 1024 * 1024,
|
||||
// 4 * 4 * 4 * 4 * 1024 * 1024,
|
||||
// 4 * 4 * 4 * 4 * 4 * 1024 * 1024,
|
||||
};
|
||||
|
||||
const iterations = 10;
|
||||
|
||||
// Updated headers to include VRAM footprint info
|
||||
std.debug.print("\n| Size (MB) | Phase | Time (ms) | GB/s | VRAM Peak |\n", .{});
|
||||
std.debug.print("|----------:|:------------------|-----------:|---------:|----------:|\n", .{});
|
||||
|
||||
for (sizes) |size| {
|
||||
// --- Phase 1: Host Init/Alloc (Outside the iteration loop for pure host prep) ---
|
||||
const data_a = try allocator.alloc(f16, size);
|
||||
defer allocator.free(data_a);
|
||||
const data_b = try allocator.alloc(f16, size);
|
||||
defer allocator.free(data_b);
|
||||
|
||||
for (0..size) |i| {
|
||||
data_a[i] = @floatFromInt(i);
|
||||
data_b[i] = @floatFromInt(size - 1 - i);
|
||||
}
|
||||
|
||||
// Track best times across iterations
|
||||
var min_alloc_ns: u64 = std.math.maxInt(u64);
|
||||
var min_transfer_ns: u64 = std.math.maxInt(u64);
|
||||
var min_compute_ns: u64 = std.math.maxInt(u64);
|
||||
|
||||
// Track peak VRAM usage observed during the iterations
|
||||
var peak_vram_bytes: u64 = 0;
|
||||
|
||||
for (0..iterations) |_| {
|
||||
// --- 1. GPU ALLOCATION PHASE ---
|
||||
const alloc_start = std.Io.Clock.awake.now(init.io);
|
||||
|
||||
const a = try Vec.initLoad(gloc, data_a);
|
||||
defer a.deinit();
|
||||
const b = try Vec.initLoad(gloc, data_b);
|
||||
defer b.deinit();
|
||||
|
||||
const alloc_duration = alloc_start.durationTo(std.Io.Clock.awake.now(init.io));
|
||||
const alloc_ns = @as(u64, @intCast(alloc_duration.toNanoseconds()));
|
||||
if (alloc_ns < min_alloc_ns) min_alloc_ns = alloc_ns;
|
||||
|
||||
// --- 2. COMPUTE PHASE ---
|
||||
const compute_start = std.Io.Clock.awake.now(init.io);
|
||||
|
||||
const sum = try a.run(gloc, b, add_pip);
|
||||
defer sum.deinit();
|
||||
|
||||
// All 3 buffers (a, b, sum) are currently resident in VRAM here.
|
||||
// Querying now catches the true peak allocation step.
|
||||
if (grena.allocated_vram_bytes > peak_vram_bytes)
|
||||
peak_vram_bytes = grena.allocated_vram_bytes;
|
||||
|
||||
device.poll();
|
||||
|
||||
const compute_duration = compute_start.durationTo(std.Io.Clock.awake.now(init.io));
|
||||
const compute_ns = @as(u64, @intCast(compute_duration.toNanoseconds()));
|
||||
if (compute_ns < min_compute_ns) min_compute_ns = compute_ns;
|
||||
|
||||
// --- 3. TRANSFER PHASE (Device -> Host) ---
|
||||
const transfer_start = std.Io.Clock.awake.now(init.io);
|
||||
|
||||
const out = try sum.read(allocator);
|
||||
defer allocator.free(out);
|
||||
|
||||
const transfer_duration = transfer_start.durationTo(std.Io.Clock.awake.now(init.io));
|
||||
const transfer_ns = @as(u64, @intCast(transfer_duration.toNanoseconds()));
|
||||
if (transfer_ns < min_transfer_ns) min_transfer_ns = transfer_ns;
|
||||
}
|
||||
|
||||
// --- Metrics Calculations ---
|
||||
const f_size = @as(f64, @floatFromInt(size));
|
||||
const element_bytes = f_size * @as(f64, @floatFromInt(@sizeOf(f16)));
|
||||
const mb = element_bytes / (1024.0 * 1024.0);
|
||||
|
||||
// Individual Phase Timings (ms)
|
||||
const alloc_ms = @as(f64, @floatFromInt(min_alloc_ns)) / 1_000_000.0;
|
||||
const compute_ms = @as(f64, @floatFromInt(min_compute_ns)) / 1_000_000.0;
|
||||
const transfer_ms = @as(f64, @floatFromInt(min_transfer_ns)) / 1_000_000.0;
|
||||
|
||||
// Bandwidth Calculations
|
||||
const alloc_gb_s = (element_bytes * 2.0 / 1_000_000_000.0) / (@as(f64, @floatFromInt(min_alloc_ns)) / 1_000_000_000.0);
|
||||
const compute_gb_s = (element_bytes * 3.0 / 1_000_000_000.0) / (@as(f64, @floatFromInt(min_compute_ns)) / 1_000_000_000.0);
|
||||
const transfer_gb_s = (element_bytes * 1.0 / 1_000_000_000.0) / (@as(f64, @floatFromInt(min_transfer_ns)) / 1_000_000_000.0);
|
||||
|
||||
// Convert Peak VRAM bytes to Megabytes for clean display
|
||||
const peak_vram_mb = @as(f64, @floatFromInt(peak_vram_bytes)) / (1024.0 * 1024.0);
|
||||
|
||||
// Print Results per Size Block with VRAM column aligned
|
||||
std.debug.print("| {d:9.2} | 1. GPU Alloc/Load | {d:10.3} | {d:8.2} | |\n", .{ mb, alloc_ms, alloc_gb_s });
|
||||
std.debug.print("| | 2. Compute | {d:10.3} | {d:8.2} | {d:7.2} MB|\n", .{ compute_ms, compute_gb_s, peak_vram_mb });
|
||||
std.debug.print("| | 3. Transfer (D->H)| {d:10.3} | {d:8.2} | |\n", .{ transfer_ms, transfer_gb_s });
|
||||
std.debug.print("|-----------|-------------------|------------|---------:|----------:|\n", .{});
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,24 @@
|
||||
enable f16;
|
||||
|
||||
@group(0) @binding(0) var<storage, read> A: array<f16>;
|
||||
@group(0) @binding(1) var<storage, read> B: array<f16>;
|
||||
@group(0) @binding(2) var<storage, read_write> C: array<f16>;
|
||||
@group(0) @binding(3) var<uniform> size: u32;
|
||||
|
||||
@compute @workgroup_size(256)
|
||||
fn main(
|
||||
@builtin(global_invocation_id) global_id : vec3<u32>,
|
||||
@builtin(num_workgroups) num_workgroups: vec3<u32>
|
||||
) {
|
||||
// 1. Calculate the total number of threads across the entire grid
|
||||
let total_threads = num_workgroups.x * 256u;
|
||||
|
||||
// 2. Start at this thread's unique global ID
|
||||
var index = global_id.x;
|
||||
|
||||
// 3. Stride through the tensor elements
|
||||
while (index < size) {
|
||||
C[index] = A[index] + B[index];
|
||||
index += total_threads; // Jump forward by the total thread count
|
||||
}
|
||||
}
|
||||
Reference in New Issue
Block a user