diff --git a/src/GpuAllocator.zig b/src/GpuAllocator.zig index 9600710..efc005d 100644 --- a/src/GpuAllocator.zig +++ b/src/GpuAllocator.zig @@ -1,5 +1,4 @@ const std = @import("std"); -const sh = @import("shaders.zig"); const GpuDevice = @import("GpuDevice.zig"); const c = @import("c.zig").c; @@ -8,25 +7,16 @@ const GpuAllocator = @This(); device: GpuDevice, cpu_allocator: std.mem.Allocator, tracked_buffers: std.AutoHashMap(c.WGPUBuffer, void), -pipelines: struct { - add: c.WGPUComputePipeline, -}, pub fn init(cpu_allocator: std.mem.Allocator, device: GpuDevice) !GpuAllocator { return .{ .device = device, .cpu_allocator = cpu_allocator, .tracked_buffers = .init(cpu_allocator), - .pipelines = .{ - .add = try buildPipeline(device.device, sh.SHADER_ADD), - }, }; } pub fn deinit(self: *GpuAllocator) void { - inline for (@typeInfo(@TypeOf(self.pipelines)).@"struct".fields) |field| - c.wgpuComputePipelineRelease(@field(self.pipelines, field.name)); - var it = self.tracked_buffers.keyIterator(); while (it.next()) |buf_ptr| { const buf = buf_ptr.*; @@ -56,22 +46,3 @@ pub fn unregisterAndDestroyBuffer(self: *GpuAllocator, buf: c.WGPUBuffer) void { c.wgpuBufferRelease(buf); } } - -fn buildPipeline(device: c.WGPUDevice, wgsl: []const u8) !c.WGPUComputePipeline { - var wgsl_src = c.WGPUShaderSourceWGSL{ - .chain = .{ .sType = c.WGPUSType_ShaderSourceWGSL }, - .code = sv(wgsl), - }; - const shader = c.wgpuDeviceCreateShaderModule(device, &.{ - .nextInChain = @ptrCast(&wgsl_src), - }) orelse return error.Shader; - defer c.wgpuShaderModuleRelease(shader); - - return c.wgpuDeviceCreateComputePipeline(device, &.{ - .compute = .{ .module = shader, .entryPoint = sv("main") }, - }) orelse error.Pipeline; -} - -fn sv(s: []const u8) c.WGPUStringView { - return .{ .data = s.ptr, .length = s.len }; -} diff --git a/src/GpuDevice.zig b/src/GpuDevice.zig index be4026c..601cd02 100644 --- a/src/GpuDevice.zig +++ b/src/GpuDevice.zig @@ -1,5 +1,4 @@ const std = @import("std"); -const sh = @import("shaders.zig"); const c = @import("c.zig").c; const Ctx = struct { diff --git a/src/GpuPipeline.zig b/src/GpuPipeline.zig new file mode 100644 index 0000000..68b3540 --- /dev/null +++ b/src/GpuPipeline.zig @@ -0,0 +1,28 @@ +const std = @import("std"); +const GpuDevice = @import("GpuDevice.zig"); +const c = @import("c.zig").c; + +raw: c.WGPUComputePipeline, + +pub fn init(device: GpuDevice, wgsl: []const u8) !@This() { + var wgsl_src = c.WGPUShaderSourceWGSL{ + .chain = .{ .sType = c.WGPUSType_ShaderSourceWGSL }, + .code = sv(wgsl), + }; + const shader = c.wgpuDeviceCreateShaderModule(device.device, &.{ + .nextInChain = @ptrCast(&wgsl_src), + }) orelse return error.Shader; + defer c.wgpuShaderModuleRelease(shader); + + return .{ .raw = c.wgpuDeviceCreateComputePipeline(device.device, &.{ + .compute = .{ .module = shader, .entryPoint = sv("main") }, + }) orelse return error.Pipeline }; +} + +pub fn deinit(self: @This()) void { + c.wgpuComputePipelineRelease(self.raw); +} + +fn sv(s: []const u8) c.WGPUStringView { + return .{ .data = s.ptr, .length = s.len }; +} diff --git a/src/Mat.zig b/src/Mat.zig index 144132a..06fcfbc 100644 --- a/src/Mat.zig +++ b/src/Mat.zig @@ -3,6 +3,7 @@ const std = @import("std"); const c = @import("c.zig").c; const GpuAllocator = @import("GpuAllocator.zig"); const GpuBuffer = @import("GpuBuffer.zig"); +const GpuPipeline = @import("GpuPipeline.zig"); const Mat = @This(); @@ -52,13 +53,13 @@ pub fn byteSize(self: Mat) u64 { return @as(u64, self.len()) * @sizeOf(f32); } -pub fn add(self: Mat, gloc: *GpuAllocator, other: Mat) !Mat { +pub fn run(self: Mat, gloc: *GpuAllocator, other: Mat, pip: GpuPipeline) !Mat { std.debug.assert(self.rows == other.rows and self.cols == other.cols); const result = try Mat.zeros(gloc, self.rows, self.cols); errdefer result.deinit(); - try dispatch2in1out(gloc, gloc.pipelines.add, self.buf, other.buf, result.buf, self.byteSize()); + try dispatch2in1out(gloc, pip.raw, self.buf, other.buf, result.buf, self.byteSize()); return result; } diff --git a/src/main.zig b/src/main.zig index 8f7748a..673ce9c 100644 --- a/src/main.zig +++ b/src/main.zig @@ -1,6 +1,7 @@ const std = @import("std"); const GpuDevice = @import("GpuDevice.zig"); const GpuAllocator = @import("GpuAllocator.zig"); +const GpuPipeline = @import("GpuPipeline.zig"); const Mat = @import("Mat.zig"); pub fn main(init: std.process.Init) !void { @@ -10,6 +11,9 @@ pub fn main(init: std.process.Init) !void { var gloc = try GpuAllocator.init(init.gpa, device); defer gloc.deinit(); + const add_pip = try GpuPipeline.init(device, @embedFile("shaders/add.wgsl")); + defer add_pip.deinit(); + // Define the sizes you want to benchmark const sizes = [_]usize{ 1, @@ -54,7 +58,7 @@ pub fn main(init: std.process.Init) !void { defer b.deinit(); // a + b - const sum = try a.add(&gloc, b); + const sum = try a.run(&gloc, b, add_pip); defer sum.deinit(); const out = try sum.read(&gloc, allocator); diff --git a/src/shaders.zig b/src/shaders.zig deleted file mode 100644 index 26f2330..0000000 --- a/src/shaders.zig +++ /dev/null @@ -1,3 +0,0 @@ -const std = @import("std"); - -pub const SHADER_ADD = @embedFile("shaders/add.wgsl");