Moved add and bench into an examples directory so src is just the library

Otherwise a GpuBuffer of a single f16 is too small (only 2) and a eneven
f16 wouldn't work either

README.md

@@ -0,0 +1,83 @@
+# Minimal Zig WebGPU Compute Library
+
+This is a minimal, self-contained Zig library designed to simplify running compute shaders using WebGPU. It abstracts away much of the boilerplate required for GPU device initialization, memory management, and pipeline execution.
+
+## Core Modules
+
+The library exports five primary components:
+
+* **`GpuDevice`**: Initializes the WebGPU instance, adapter, device, and queue. It is configured to prioritize high performance and automatically requests the `ShaderF16` feature if the adapter supports it. By default, it enforces a 2 GB VRAM limit.
+* **`GpuArena` / `GpuAllocator`**: A memory management layer that tracks allocated VRAM bytes to prevent exceeding the device budget. The arena automatically destroys and releases all tracked WebGPU buffers when deinitialized.
+* **`GpuBuffer`**: Wraps native WebGPU buffers. It automatically aligns buffer sizes forward to a multiple of 4 bytes. It provides a `.load()` method for CPU-to-GPU data transfers (handling both aligned and unaligned lengths smoothly) and a `.read()` method that utilizes a staging buffer to map GPU data back to the CPU.
+* **`GpuProcess`**: Compiles WGSL source code into a compute pipeline. When running a process, it automatically splits the work into manageable chunks (up to 1 GB at a time) and dispatches workgroups of size 256.
+
+## Quick Start Example
+
+Below is a complete, self-contained example demonstrating how to initialize the GPU, load data, run a compute shader, and read the results back to the CPU:
+
+```zig
+
+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});
+}
+```
+
+## Dependencies
+
+* **`wgpu.h`**: The library relies on the WebGPU C API headers to bind to the native system graphics.

build.zig

@@ -1,61 +1,60 @@
-// build.zig
-// zig build run
-//
-// Expects wgpu-native pre-built in libs/wgpu-native/:
-//   include/wgpu.h
-//   lib/libwgpu_native.a   (or .so / .dylib / .dll)
-//
-// Download release: https://github.com/gfx-rs/wgpu-native/releases
-// Pick the archive matching your OS/arch, e.g.:
-//   wgpu-linux-x86_64-release.zip     → libwgpu_native.a  + wgpu.h
-//   wgpu-macos-aarch64-release.zip
-//   wgpu-windows-x86_64-msvc-release.zip
-
 const std = @import("std");
 
-pub fn build(b: *std.Build) void {
+pub fn build(b: *std.Build) !void {
     const target = b.standardTargetOptions(.{});
     const optimize = b.standardOptimizeOption(.{});
 
-    // 1. Define the module so other projects can import it
-    _ = b.addModule("zig-wgpu", .{
+    // Define the module so other projects can import it
+    const mod = b.addModule("zig-wgpu", .{
         .root_source_file = b.path("src/lib.zig"),
-    });
-
-    const exe = b.addExecutable(.{
-        .root_module = b.createModule(.{
-            .root_source_file = b.path("src/main.zig"),
-            .link_libc = true,
         .target = target,
-            .optimize = optimize,
-        }),
-        .name = "gpu_matrix_add",
     });
 
-    // wgpu-native headers + pre-built static library
-    exe.root_module.addIncludePath(b.path("libs/wgpu-native/include"));
-    exe.root_module.addLibraryPath(b.path("libs/wgpu-native/lib"));
-    exe.root_module.addObjectFile(b.path("libs/wgpu-native/lib/libwgpu_native.a"));
+    mod.addIncludePath(b.path("libs/wgpu-native/include"));
+    mod.addLibraryPath(b.path("libs/wgpu-native/lib"));
+    mod.addObjectFile(b.path("libs/wgpu-native/lib/libwgpu_native.a"));
 
     // Platform-specific system frameworks needed by wgpu-native
     const t = target.result;
     if (t.os.tag == .macos) {
-        exe.root_module.linkFramework("Metal", .{});
-        exe.root_module.linkFramework("QuartzCore", .{});
-        exe.root_module.linkFramework("Foundation", .{});
-        exe.root_module.linkFramework("CoreGraphics", .{});
+        mod.linkFramework("Metal", .{});
+        mod.linkFramework("QuartzCore", .{});
+        mod.linkFramework("Foundation", .{});
+        mod.linkFramework("CoreGraphics", .{});
     } else if (t.os.tag == .windows) {
-        exe.root_module.linkSystemLibrary("d3d12", .{});
-        exe.root_module.linkSystemLibrary("dxgi", .{});
-        exe.root_module.linkSystemLibrary("user32", .{});
+        mod.linkSystemLibrary("d3d12", .{});
+        mod.linkSystemLibrary("dxgi", .{});
+        mod.linkSystemLibrary("user32", .{});
     } else {
-        exe.root_module.linkSystemLibrary("vulkan", .{});
-        exe.root_module.linkSystemLibrary("gcc_s", .{});
+        mod.linkSystemLibrary("vulkan", .{});
+        mod.linkSystemLibrary("gcc_s", .{});
     }
 
+    var threaded: std.Io.Threaded = .init_single_threaded;
+    const io = threaded.io();
+
+    var buf: [1024]u8 = undefined;
+    const exemples = try std.Io.Dir.cwd().openDir(io, "examples", .{ .access_sub_paths = false, .iterate = true });
+    var iter = exemples.iterate();
+    while (try iter.next(io)) |entry| {
+        if (entry.kind != .file) continue;
+        if (!std.mem.eql(u8, entry.name[entry.name.len - 4 ..], ".zig")) continue;
+
+        const exe = b.addExecutable(.{
+            .name = entry.name[0 .. entry.name.len - 4],
+            .root_module = b.createModule(.{
+                .root_source_file = b.path(try std.fmt.bufPrint(&buf, "examples/{s}", .{entry.name})),
+                .target = target,
+                .optimize = optimize,
+                .imports = &.{},
+            }),
+        });
+        exe.root_module.addImport("gpu", mod);
+
         b.installArtifact(exe);
 
-    const run = b.addRunArtifact(exe);
-    run.step.dependOn(b.getInstallStep());
-    b.step("run", "Build and run").dependOn(&run.step);
+        const run_step = b.step(entry.name[0 .. entry.name.len - 4], try std.fmt.bufPrint(&buf, "Run {s} demo", .{entry.name}));
+        const run_cmd = b.addRunArtifact(exe);
+        run_step.dependOn(&run_cmd.step);
+    }
 }

build.zig.zon

@@ -1,81 +1,13 @@
 .{
-    // This is the default name used by packages depending on this one. For
-    // example, when a user runs `zig fetch --save <url>`, this field is used
-    // as the key in the `dependencies` table. Although the user can choose a
-    // different name, most users will stick with this provided value.
-    //
-    // It is redundant to include "zig" in this name because it is already
-    // within the Zig package namespace.
     .name = .zig_wgpu,
-    // This is a [Semantic Version](https://semver.org/).
-    // In a future version of Zig it will be used for package deduplication.
-    .version = "0.0.0",
-    // Together with name, this represents a globally unique package
-    // identifier. This field is generated by the Zig toolchain when the
-    // package is first created, and then *never changes*. This allows
-    // unambiguous detection of one package being an updated version of
-    // another.
-    //
-    // When forking a Zig project, this id should be regenerated (delete the
-    // field and run `zig build`) if the upstream project is still maintained.
-    // Otherwise, the fork is *hostile*, attempting to take control over the
-    // original project's identity. Thus it is recommended to leave the comment
-    // on the following line intact, so that it shows up in code reviews that
-    // modify the field.
-    .fingerprint = 0x5d0e853acbc0c2c6, // Changing this has security and trust implications.
-    // Tracks the earliest Zig version that the package considers to be a
-    // supported use case.
+    .version = "0.1.0",
+    .fingerprint = 0x5d0e853acbc0c2c6,
     .minimum_zig_version = "0.16.0",
-    // This field is optional.
-    // Each dependency must either provide a `url` and `hash`, or a `path`.
-    // `zig build --fetch` can be used to fetch all dependencies of a package, recursively.
-    // Once all dependencies are fetched, `zig build` no longer requires
-    // internet connectivity.
-    .dependencies = .{
-        // See `zig fetch --save <url>` for a command-line interface for adding dependencies.
-        //.example = .{
-        //    // When updating this field to a new URL, be sure to delete the corresponding
-        //    // `hash`, otherwise you are communicating that you expect to find the old hash at
-        //    // the new URL. If the contents of a URL change this will result in a hash mismatch
-        //    // which will prevent zig from using it.
-        //    .url = "https://example.com/foo.tar.gz",
-        //
-        //    // This is computed from the file contents of the directory of files that is
-        //    // obtained after fetching `url` and applying the inclusion rules given by
-        //    // `paths`.
-        //    //
-        //    // This field is the source of truth; packages do not come from a `url`; they
-        //    // come from a `hash`. `url` is just one of many possible mirrors for how to
-        //    // obtain a package matching this `hash`.
-        //    //
-        //    // Uses the [multihash](https://multiformats.io/multihash/) format.
-        //    .hash = "...",
-        //
-        //    // When this is provided, the package is found in a directory relative to the
-        //    // build root. In this case the package's hash is irrelevant and therefore not
-        //    // computed. This field and `url` are mutually exclusive.
-        //    .path = "foo",
-        //
-        //    // When this is set to `true`, a package is declared to be lazily
-        //    // fetched. This makes the dependency only get fetched if it is
-        //    // actually used.
-        //    .lazy = false,
-        //},
-    },
-    // Specifies the set of files and directories that are included in this package.
-    // Only files and directories listed here are included in the `hash` that
-    // is computed for this package. Only files listed here will remain on disk
-    // when using the zig package manager. As a rule of thumb, one should list
-    // files required for compilation plus any license(s).
-    // Paths are relative to the build root. Use the empty string (`""`) to refer to
-    // the build root itself.
-    // A directory listed here means that all files within, recursively, are included.
+    .dependencies = .{},
     .paths = .{
         "build.zig",
         "build.zig.zon",
         "src",
-        // For example...
-        //"LICENSE",
-        //"README.md",
+        "libs",
     },
 }

examples/add.zig

@@ -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});
+}

examples/bench.zig

@@ -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", .{});
+    }
+}

examples/shaders/add.wgsl

@@ -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
+    }
+}

src/GpuAllocator.zig

@@ -1,140 +1,24 @@
+// GpuAllocator.zig
 const std = @import("std");
-const sh = @import("shaders.zig");
-const c = @import("c.zig").c;
+const GpuDevice = @import("GpuDevice.zig");
+const c = @import("utils.zig").c;
 
 const GpuAllocator = @This();
 
-instance: c.WGPUInstance,
-adapter: c.WGPUAdapter,
-device: c.WGPUDevice,
-queue: c.WGPUQueue,
-
-// Lazily created, cached for lifetime of allocator
-_pip_add: c.WGPUComputePipeline = null,
-_pip_scale: c.WGPUComputePipeline = null,
-
-pub fn init() !GpuAllocator {
-    const instance = c.wgpuCreateInstance(
-        &std.mem.zeroes(c.WGPUInstanceDescriptor),
-    ) orelse return error.NoInstance;
-    errdefer c.wgpuInstanceRelease(instance);
-
-    var ctx = Ctx{};
-    _ = c.wgpuInstanceRequestAdapter(
-        instance,
-        &.{ .powerPreference = c.WGPUPowerPreference_HighPerformance },
-        .{ .callback = onAdapter, .userdata1 = &ctx },
-    );
-    c.wgpuInstanceProcessEvents(instance);
-    const adapter = ctx.adapter orelse return error.NoAdapter;
-    errdefer c.wgpuAdapterRelease(adapter);
-
-    _ = c.wgpuAdapterRequestDevice(
-        adapter,
-        null,
-        .{ .callback = onDevice, .userdata1 = &ctx },
-    );
-    c.wgpuInstanceProcessEvents(instance);
-    const device = ctx.device orelse return error.NoDevice;
-
-    return .{
-        .instance = instance,
-        .adapter = adapter,
-        .device = device,
-        .queue = c.wgpuDeviceGetQueue(device),
-    };
-}
-
-pub fn deinit(self: *GpuAllocator) void {
-    if (self._pip_add) |p| c.wgpuComputePipelineRelease(p);
-    if (self._pip_scale) |p| c.wgpuComputePipelineRelease(p);
-    c.wgpuQueueRelease(self.queue);
-    c.wgpuDeviceRelease(self.device);
-    c.wgpuAdapterRelease(self.adapter);
-    c.wgpuInstanceRelease(self.instance);
-}
-
-// ── Internal ─────────────────────────────────────────────────────────────
-
-pub fn makeBuffer(
-    self: *GpuAllocator,
-    bytes: u64,
-    usage: c.WGPUBufferUsage,
-) !c.WGPUBuffer {
-    return c.wgpuDeviceCreateBuffer(self.device, &.{
-        .usage = usage,
-        .size = bytes,
-    }) orelse error.BufferAlloc;
-}
-
-pub fn pipAdd(self: *GpuAllocator) !c.WGPUComputePipeline {
-    if (self._pip_add == null)
-        self._pip_add = try buildPipeline(self.device, sh.SHADER_ADD);
-    return self._pip_add.?;
-}
-
-pub fn pipScale(self: *GpuAllocator) !c.WGPUComputePipeline {
-    if (self._pip_scale == null)
-        self._pip_scale = try buildPipeline(self.device, sh.SHADER_SCALE);
-    return self._pip_scale.?;
-}
-
-/// Poll until GPU work completes. Use after submit if you need CPU sync.
-pub fn poll(self: *GpuAllocator) void {
-    _ = c.wgpuDevicePoll(self.device, 1, null);
-}
-
-const Ctx = struct {
-    adapter: c.WGPUAdapter = null,
-    device: c.WGPUDevice = null,
+/// The function definitions our underlying implementations must satisfy
+pub const VTable = struct {
+    alloc: *const fn (ctx: *anyopaque, bytes: u64, usage: c.WGPUBufferUsage) anyerror!c.WGPUBuffer,
+    free: *const fn (ctx: *anyopaque, buf_raw: c.WGPUBuffer, size: u64) void,
 };
 
-fn onAdapter(
-    status: c.WGPURequestAdapterStatus,
-    adapter: c.WGPUAdapter,
-    _: c.WGPUStringView,
-    userdata1: ?*anyopaque,
-    _: ?*anyopaque,
-) callconv(.c) void {
-    if (status != c.WGPURequestAdapterStatus_Success) {
-        std.log.err("Adapter request failed (status={d})", .{status});
-        return;
-    }
-    const ctx: *Ctx = @ptrCast(@alignCast(userdata1.?));
-    ctx.adapter = adapter;
+device: GpuDevice,
+ptr: *anyopaque,
+vtable: *const VTable,
+
+pub fn allocBuffer(self: GpuAllocator, bytes: u64, usage: c.WGPUBufferUsage) !c.WGPUBuffer {
+    return self.vtable.alloc(self.ptr, bytes, usage);
 }
 
-fn onDevice(
-    status: c.WGPURequestDeviceStatus,
-    device: c.WGPUDevice,
-    _: c.WGPUStringView,
-    userdata1: ?*anyopaque,
-    _: ?*anyopaque,
-) callconv(.c) void {
-    if (status != c.WGPURequestDeviceStatus_Success) {
-        std.log.err("Device request failed (status={d})", .{status});
-        return;
-    }
-    const ctx: *Ctx = @ptrCast(@alignCast(userdata1.?));
-    ctx.device = device;
-    std.debug.print("{?}", .{device});
-}
-
-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 };
+pub fn freeBuffer(self: GpuAllocator, buf_raw: c.WGPUBuffer, size: u64) void {
+    self.vtable.free(self.ptr, buf_raw, size);
 }

src/GpuArena.zig

@@ -0,0 +1,72 @@
+// GpuArena.zig
+const std = @import("std");
+const GpuDevice = @import("GpuDevice.zig");
+const GpuAllocator = @import("GpuAllocator.zig");
+const c = @import("utils.zig").c;
+
+const GpuArena = @This();
+
+device: GpuDevice,
+tracked_buffers: std.AutoHashMap(c.WGPUBuffer, void),
+allocated_vram_bytes: u64 = 0,
+
+pub fn init(cpu_allocator: std.mem.Allocator, device: GpuDevice) GpuArena {
+    return .{
+        .device = device,
+        .tracked_buffers = .init(cpu_allocator),
+    };
+}
+
+pub fn deinit(self: *GpuArena) void {
+    var it = self.tracked_buffers.keyIterator();
+    while (it.next()) |buf_ptr| {
+        c.wgpuBufferDestroy(buf_ptr.*);
+        c.wgpuBufferRelease(buf_ptr.*);
+    }
+    self.tracked_buffers.deinit();
+}
+
+/// Returns the type-erased immutable interface wrapper
+pub fn gpuAllocator(self: *GpuArena) GpuAllocator {
+    return .{
+        .device = self.device,
+        .ptr = self,
+        .vtable = &.{
+            .alloc = alloc,
+            .free = free,
+        },
+    };
+}
+
+fn alloc(ctx: *anyopaque, bytes: u64, usage: c.WGPUBufferUsage) anyerror!c.WGPUBuffer {
+    const self: *GpuArena = @ptrCast(@alignCast(ctx));
+
+    if (bytes > self.device.limits.maxBufferSize)
+        return error.SingleBufferExceedsLimit;
+
+    if (bytes + self.allocated_vram_bytes > self.device.config.vram_bytes_limit)
+        return error.ExceedsVramBudget;
+
+    const buf = c.wgpuDeviceCreateBuffer(self.device.device, &.{
+        .usage = usage,
+        .size = bytes,
+    }) orelse return error.BufferAlloc;
+    errdefer {
+        c.wgpuBufferDestroy(buf);
+        c.wgpuBufferRelease(buf);
+    }
+
+    try self.tracked_buffers.put(buf, {});
+    self.allocated_vram_bytes += bytes;
+    return buf;
+}
+
+fn free(ctx: *anyopaque, buf_raw: c.WGPUBuffer, size: u64) void {
+    const self: *GpuArena = @ptrCast(@alignCast(ctx));
+
+    if (self.tracked_buffers.remove(buf_raw)) {
+        c.wgpuBufferDestroy(buf_raw);
+        c.wgpuBufferRelease(buf_raw);
+        self.allocated_vram_bytes -= size;
+    }
+}

src/GpuBuffer.zig

@@ -0,0 +1,138 @@
+const std = @import("std");
+const c = @import("utils.zig").c;
+const GpuAllocator = @import("GpuAllocator.zig");
+
+raw: c.WGPUBuffer,
+size: u64, // Now tracks the 4-byte aligned size directly
+usage: c.WGPUBufferUsage,
+gloc: GpuAllocator,
+
+const BufferUsage = enum(u64) {
+    None = 0x0000000000000000,
+    MapRead = 0x0000000000000001,
+    MapWrite = 0x0000000000000002,
+    CopySrc = 0x0000000000000004,
+    CopyDst = 0x0000000000000008,
+    Index = 0x0000000000000010,
+    Vertex = 0x0000000000000020,
+    Uniform = 0x0000000000000040,
+    Storage = 0x0000000000000080,
+    Indirect = 0x0000000000000100,
+    QueryResolve = 0x0000000000000200,
+};
+
+/// Allocates the underlying WebGPU handle and registers it to the parent GpuAllocator
+pub fn init(gloc: GpuAllocator, size: u64, usage: std.EnumSet(BufferUsage)) !@This() {
+    var use: u64 = 0;
+    var iter = usage.iterator();
+    while (iter.next()) |flag| use |= @intFromEnum(flag);
+
+    // Automatically align the buffer size forward to a multiple of 4 bytes under the hood
+    const aligned_size = std.mem.alignForward(u64, size, 4);
+
+    const raw_handle = try gloc.allocBuffer(aligned_size, use);
+    return .{
+        .raw = raw_handle,
+        .size = aligned_size, // Expose the aligned size to the rest of the application
+        .usage = use,
+        .gloc = gloc,
+    };
+}
+
+/// Unregisters from the parent GpuAllocator and cleanly destroys GPU resources
+pub fn deinit(self: @This()) void {
+    self.gloc.freeBuffer(self.raw, self.size);
+}
+
+/// Native getConstMappedRange wrapper
+pub fn getConstMappedRange(self: @This(), offset: u64, size: u64) ?*const anyopaque {
+    return c.wgpuBufferGetConstMappedRange(self.raw, offset, size);
+}
+
+/// Native mapAsync wrapper
+pub fn mapAsync(
+    self: @This(),
+    mode: c.WGPUMapMode,
+    offset: u64,
+    size: u64,
+    callback_info: c.WGPUBufferMapCallbackInfo,
+) void {
+    _ = c.wgpuBufferMapAsync(self.raw, mode, offset, size, callback_info);
+}
+
+/// Native unmap wrapper
+pub fn unmap(self: @This()) void {
+    c.wgpuBufferUnmap(self.raw);
+}
+
+/// CPU to GPU.
+pub fn load(
+    self: @This(),
+    T: type,
+    data: []const T,
+) !void {
+    const bytes = data.len * @sizeOf(T);
+
+    if (bytes == self.size) {
+        // Aligned path: direct download
+        c.wgpuQueueWriteBuffer(self.gloc.device.queue, self.raw, 0, data.ptr, self.size);
+    } else {
+        // Unaligned path: Split the write into an aligned chunk and a padded remainder
+        // to support arbitrary lengths without any allocations or large stack arrays.
+        const aligned_part = (bytes / 4) * 4;
+        if (aligned_part > 0) {
+            c.wgpuQueueWriteBuffer(self.gloc.device.queue, self.raw, 0, data.ptr, aligned_part);
+        }
+
+        var remainder_buf: [4]u8 = .{ 0, 0, 0, 0 };
+        const data_bytes = std.mem.sliceAsBytes(data);
+        @memcpy(remainder_buf[0 .. bytes - aligned_part], data_bytes[aligned_part..bytes]);
+
+        c.wgpuQueueWriteBuffer(self.gloc.device.queue, self.raw, aligned_part, &remainder_buf, 4);
+    }
+}
+
+pub fn read(self: @This(), alloc: std.mem.Allocator, T: type) ![]T {
+    const out = try alloc.alloc(T, @divExact(self.size, @sizeOf(T)));
+
+    const staging = try init(
+        self.gloc,
+        self.size,
+        .initMany(&.{ .MapRead, .CopyDst }),
+    );
+    defer staging.deinit();
+
+    const enc = c.wgpuDeviceCreateCommandEncoder(self.gloc.device.device, null) orelse return error.Encoder;
+    c.wgpuCommandEncoderCopyBufferToBuffer(enc, self.raw, 0, staging.raw, 0, self.size);
+    const cmd = c.wgpuCommandEncoderFinish(enc, null);
+    defer c.wgpuCommandEncoderRelease(enc);
+    defer c.wgpuCommandBufferRelease(cmd);
+    c.wgpuQueueSubmit(self.gloc.device.queue, 1, &cmd);
+
+    var mapped = false;
+    staging.mapAsync(
+        c.WGPUMapMode_Read,
+        0,
+        self.size,
+        .{ .callback = onMapped, .userdata1 = &mapped },
+    );
+    while (!mapped) self.gloc.device.poll();
+
+    const ptr: [*]const T = @ptrCast(@alignCast(
+        staging.getConstMappedRange(0, self.size),
+    ));
+    @memcpy(out[0..out.len], ptr[0..out.len]);
+    staging.unmap();
+
+    return out;
+}
+
+fn onMapped(
+    status: c.WGPUMapAsyncStatus,
+    _: c.WGPUStringView,
+    userdata1: ?*anyopaque,
+    _: ?*anyopaque,
+) callconv(.c) void {
+    const flag: *bool = @ptrCast(@alignCast(userdata1.?));
+    flag.* = (status == c.WGPUMapAsyncStatus_Success);
+}

src/GpuDevice.zig

@@ -0,0 +1,129 @@
+const std = @import("std");
+const c = @import("utils.zig").c;
+const sv = @import("utils.zig").sv;
+
+const Ctx = struct {
+    adapter: c.WGPUAdapter = null,
+    device: c.WGPUDevice = null,
+};
+
+const GpuDeviceConfig = struct {
+    /// VRAM limit. Default 2 GB
+    vram_bytes_limit: u64 = 2 * 1024 * 1024 * 1024,
+    power_preference: enum(c_uint) {
+        Undefined = 0x00000000,
+        LowPower = 0x00000001,
+        HighPerformance = 0x00000002,
+        Force32 = 0x7FFFFFFF,
+    } = .HighPerformance,
+};
+
+instance: c.WGPUInstance,
+adapter: c.WGPUAdapter,
+device: c.WGPUDevice,
+queue: c.WGPUQueue,
+limits: c.WGPULimits,
+
+config: GpuDeviceConfig,
+
+pub fn init(config: GpuDeviceConfig) !@This() {
+    const instance = c.wgpuCreateInstance(
+        &std.mem.zeroes(c.WGPUInstanceDescriptor),
+    ) orelse return error.NoInstance;
+    errdefer c.wgpuInstanceRelease(instance);
+
+    var ctx = Ctx{};
+    _ = c.wgpuInstanceRequestAdapter(
+        instance,
+        &.{ .powerPreference = @intFromEnum(config.power_preference) },
+        .{ .callback = onAdapter, .userdata1 = &ctx },
+    );
+    c.wgpuInstanceProcessEvents(instance);
+    const adapter = ctx.adapter orelse return error.NoAdapter;
+    errdefer c.wgpuAdapterRelease(adapter);
+
+    var supported_features = std.mem.zeroes(c.WGPUSupportedFeatures);
+    c.wgpuAdapterGetFeatures(adapter, &supported_features);
+
+    var supported_limits = std.mem.zeroes(c.WGPULimits);
+    supported_limits.nextInChain = null;
+    if (c.wgpuAdapterGetLimits(adapter, &supported_limits) != 1) return error.FailedToGetAdapterLimits;
+
+    var has_f16 = false;
+    for (0..supported_features.featureCount) |i| {
+        if (supported_features.features[i] == c.WGPUFeatureName_ShaderF16) {
+            has_f16 = true;
+            break;
+        }
+    }
+
+    var feature_buf = [_]c.WGPUFeatureName{c.WGPUFeatureName_ShaderF16};
+    const required_features: []const c.WGPUFeatureName =
+        if (has_f16) feature_buf[0..1] else &.{};
+
+    const device_descriptor = c.WGPUDeviceDescriptor{
+        .nextInChain = null,
+        .label = sv("TensorCompilerDevice"),
+        .requiredFeatureCount = required_features.len,
+        .requiredFeatures = if (required_features.len > 0) required_features.ptr else null,
+        .requiredLimits = &supported_limits,
+    };
+    _ = c.wgpuAdapterRequestDevice(
+        adapter,
+        &device_descriptor,
+        .{ .callback = onDevice, .userdata1 = &ctx },
+    );
+    c.wgpuInstanceProcessEvents(instance);
+    const device = ctx.device orelse return error.NoDevice;
+
+    return .{
+        .instance = instance,
+        .adapter = adapter,
+        .device = device,
+        .queue = c.wgpuDeviceGetQueue(device),
+        .limits = supported_limits,
+        .config = config,
+    };
+}
+
+pub fn deinit(self: @This()) void {
+    c.wgpuQueueRelease(self.queue);
+    c.wgpuDeviceRelease(self.device);
+    c.wgpuAdapterRelease(self.adapter);
+    c.wgpuInstanceRelease(self.instance);
+}
+
+/// Wait for thing to be done
+pub fn poll(self: @This()) void {
+    _ = c.wgpuDevicePoll(self.device, 1, null);
+}
+
+fn onAdapter(
+    status: c.WGPURequestAdapterStatus,
+    adapter: c.WGPUAdapter,
+    _: c.WGPUStringView,
+    userdata1: ?*anyopaque,
+    _: ?*anyopaque,
+) callconv(.c) void {
+    if (status != c.WGPURequestAdapterStatus_Success) {
+        std.log.err("Adapter request failed (status={d})", .{status});
+        return;
+    }
+    const ctx: *Ctx = @ptrCast(@alignCast(userdata1.?));
+    ctx.adapter = adapter;
+}
+
+fn onDevice(
+    status: c.WGPURequestDeviceStatus,
+    device: c.WGPUDevice,
+    _: c.WGPUStringView,
+    userdata1: ?*anyopaque,
+    _: ?*anyopaque,
+) callconv(.c) void {
+    if (status != c.WGPURequestDeviceStatus_Success) {
+        std.log.err("Device request failed (status={d})", .{status});
+        return;
+    }
+    const ctx: *Ctx = @ptrCast(@alignCast(userdata1.?));
+    ctx.device = device;
+}

src/GpuProcess.zig

@@ -0,0 +1,122 @@
+/// GpuProcess is just a pipeline with 2 inpout and 1 output
+/// for now, to see if I make it a bit more generic
+///
+const std = @import("std");
+const c = @import("utils.zig").c;
+const sv = @import("utils.zig").sv;
+const GpuAllocator = @import("GpuAllocator.zig");
+const GpuBuffer = @import("GpuBuffer.zig");
+const GpuDevice = @import("GpuDevice.zig");
+
+pip: 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 .{ .pip = c.wgpuDeviceCreateComputePipeline(device.device, &.{
+        .compute = .{ .module = shader, .entryPoint = sv("main") },
+    }) orelse return error.Pipeline };
+}
+
+pub fn deinit(self: @This()) void {
+    c.wgpuComputePipelineRelease(self.pip);
+}
+
+fn onMapped(
+    status: c.WGPUMapAsyncStatus,
+    _: c.WGPUStringView,
+    userdata1: ?*anyopaque,
+    _: ?*anyopaque,
+) callconv(.c) void {
+    const flag: *bool = @ptrCast(@alignCast(userdata1.?));
+    flag.* = (status == c.WGPUMapAsyncStatus_Success);
+}
+
+// Changed: gloc is passed by value instead of *GpuAllocator
+pub fn run(
+    self: @This(),
+    gloc: GpuAllocator,
+    T: type,
+    buf_a: GpuBuffer,
+    buf_b: GpuBuffer,
+    buf_out: GpuBuffer,
+) !void {
+    const max_chunk_bytes: u64 = 1024 * 1024 * 1024; // 1 GB
+
+    const bytes = buf_a.size;
+    var offset: u64 = 0;
+    while (offset < bytes) {
+        const current_chunk_bytes = @min(max_chunk_bytes, bytes - offset);
+        const current_chunk_elements: u32 = @intCast(current_chunk_bytes / @sizeOf(T));
+
+        const info_buf = try GpuBuffer.init(
+            gloc,
+            @sizeOf(u32),
+            .initMany(&.{ .Uniform, .CopyDst }),
+        );
+        defer info_buf.deinit();
+
+        c.wgpuQueueWriteBuffer(gloc.device.queue, info_buf.raw, 0, &current_chunk_elements, @sizeOf(u32));
+
+        const entries = [_]c.WGPUBindGroupEntry{
+            .{ .binding = 0, .buffer = buf_a.raw, .offset = offset, .size = current_chunk_bytes },
+            .{ .binding = 1, .buffer = buf_b.raw, .offset = offset, .size = current_chunk_bytes },
+            .{ .binding = 2, .buffer = buf_out.raw, .offset = offset, .size = current_chunk_bytes },
+            .{ .binding = 3, .buffer = info_buf.raw, .offset = 0, .size = @sizeOf(u32) },
+        };
+
+        try submitPass(gloc, self.pip, &entries, current_chunk_elements);
+
+        offset += current_chunk_bytes;
+    }
+}
+
+// Changed: gloc is passed by value instead of *GpuAllocator
+fn submitPass(
+    gloc: GpuAllocator,
+    pipeline: c.WGPUComputePipeline,
+    entries: []const c.WGPUBindGroupEntry,
+    n: usize,
+) !void {
+    const bgl = c.wgpuComputePipelineGetBindGroupLayout(pipeline, 0);
+    defer c.wgpuBindGroupLayoutRelease(bgl);
+
+    const bg = c.wgpuDeviceCreateBindGroup(gloc.device.device, &.{
+        .layout = bgl,
+        .entries = entries.ptr,
+        .entryCount = entries.len,
+    }) orelse return error.BindGroup;
+    defer c.wgpuBindGroupRelease(bg);
+
+    const enc = c.wgpuDeviceCreateCommandEncoder(gloc.device.device, null) orelse
+        return error.Encoder;
+    const pass = c.wgpuCommandEncoderBeginComputePass(enc, null);
+    c.wgpuComputePassEncoderSetPipeline(pass, pipeline);
+    c.wgpuComputePassEncoderSetBindGroup(pass, 0, bg, 0, null);
+
+    const WORKGROUP_SIZE = 256;
+    const MAX_WORKGROUPS = 65535;
+
+    const desired_workgroups = ceilDiv(n, WORKGROUP_SIZE);
+    const dispatch_count = @min(desired_workgroups, MAX_WORKGROUPS);
+
+    c.wgpuComputePassEncoderDispatchWorkgroups(pass, @intCast(dispatch_count), 1, 1);
+    c.wgpuComputePassEncoderEnd(pass);
+    c.wgpuComputePassEncoderRelease(pass);
+
+    const cmd = c.wgpuCommandEncoderFinish(enc, null);
+    defer c.wgpuCommandEncoderRelease(enc);
+    defer c.wgpuCommandBufferRelease(cmd);
+    c.wgpuQueueSubmit(gloc.device.queue, 1, &cmd);
+}
+
+fn ceilDiv(n: usize, d: usize) usize {
+    return (n + d - 1) / d;
+}

src/Mat.zig

@@ -1,204 +0,0 @@
-const std = @import("std");
-const c = @import("c.zig").c;
-const GpuAllocator = @import("GpuAllocator.zig");
-
-const Mat = @This();
-
-buf: c.WGPUBuffer,
-rows: u32,
-cols: u32,
-
-// ── Lifecycle ─────────────────────────────────────────────────────────────
-
-/// Allocate GPU buffer and upload `data`. `data.len` must equal rows*cols.
-pub fn load(
-    gloc: *GpuAllocator,
-    data: []const f32,
-    rows: u32,
-    cols: u32,
-) !Mat {
-    std.debug.assert(data.len == @as(usize, rows) * cols);
-    const bytes = data.len * @sizeOf(f32);
-    const buf = try gloc.makeBuffer(
-        bytes,
-        c.WGPUBufferUsage_Storage |
-            c.WGPUBufferUsage_CopyDst |
-            c.WGPUBufferUsage_CopySrc,
-    );
-    c.wgpuQueueWriteBuffer(gloc.queue, buf, 0, data.ptr, bytes);
-    return .{ .buf = buf, .rows = rows, .cols = cols };
-}
-
-/// Allocate zeroed GPU buffer (no upload).
-pub fn zeros(gloc: *GpuAllocator, rows: u32, cols: u32) !Mat {
-    const bytes: u64 = @as(u64, rows) * cols * @sizeOf(f32);
-    const buf = try gloc.makeBuffer(
-        bytes,
-        c.WGPUBufferUsage_Storage |
-            c.WGPUBufferUsage_CopyDst |
-            c.WGPUBufferUsage_CopySrc,
-    );
-    return .{ .buf = buf, .rows = rows, .cols = cols };
-}
-
-pub fn deinit(self: Mat) void {
-    c.wgpuBufferRelease(self.buf);
-}
-
-pub fn len(self: Mat) u32 {
-    return self.rows * self.cols;
-}
-
-pub fn byteSize(self: Mat) u64 {
-    return @as(u64, self.len()) * @sizeOf(f32);
-}
-
-/// Element-wise add. Shapes must match. Returns new Mat (caller owns).
-pub fn add(self: Mat, gloc: *GpuAllocator, other: Mat) !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();
-
-    const pipeline = try gloc.pipAdd();
-    try dispatch2in1out(gloc, pipeline, self.buf, other.buf, result.buf, self.byteSize(), self.len());
-
-    return result;
-}
-
-/// Element-wise multiply by scalar. Returns new Mat (caller owns).
-pub fn scale(self: Mat, gloc: *GpuAllocator, scalar: f32) !Mat {
-    const result = try Mat.zeros(gloc, self.rows, self.cols);
-    errdefer result.deinit();
-
-    const bytes = self.byteSize();
-    const n = self.len();
-
-    // Upload scalar as uniform buffer
-    const uni_buf = try gloc.makeBuffer(
-        @sizeOf(f32),
-        c.WGPUBufferUsage_Uniform | c.WGPUBufferUsage_CopyDst,
-    );
-    defer c.wgpuBufferRelease(uni_buf);
-    c.wgpuQueueWriteBuffer(gloc.queue, uni_buf, 0, &scalar, @sizeOf(f32));
-
-    const pipeline = try gloc.pipScale();
-    const bgl = c.wgpuComputePipelineGetBindGroupLayout(pipeline, 0);
-    defer c.wgpuBindGroupLayoutRelease(bgl);
-
-    const entries = [_]c.WGPUBindGroupEntry{
-        .{ .binding = 0, .buffer = self.buf, .offset = 0, .size = bytes },
-        .{ .binding = 1, .buffer = result.buf, .offset = 0, .size = bytes },
-        .{ .binding = 2, .buffer = uni_buf, .offset = 0, .size = @sizeOf(f32) },
-    };
-    try submitPass(gloc, pipeline, &entries, n);
-
-    return result;
-}
-
-/// Read GPU buffer back to CPU. `out.len` must be >= rows*cols.
-pub fn read(self: Mat, gloc: *GpuAllocator, out: []f32) !void {
-    std.debug.assert(out.len >= self.len());
-    const bytes = self.byteSize();
-
-    const staging = try gloc.makeBuffer(
-        bytes,
-        c.WGPUBufferUsage_MapRead | c.WGPUBufferUsage_CopyDst,
-    );
-    defer c.wgpuBufferRelease(staging);
-
-    // Copy result → staging
-    const enc = c.wgpuDeviceCreateCommandEncoder(gloc.device, null) orelse
-        return error.Encoder;
-    c.wgpuCommandEncoderCopyBufferToBuffer(enc, self.buf, 0, staging, 0, bytes);
-    const cmd = c.wgpuCommandEncoderFinish(enc, null);
-    defer c.wgpuCommandEncoderRelease(enc);
-    defer c.wgpuCommandBufferRelease(cmd);
-    c.wgpuQueueSubmit(gloc.queue, 1, &cmd);
-
-    // Map and copy to slice
-    var mapped = false;
-    _ = c.wgpuBufferMapAsync(
-        staging,
-        c.WGPUMapMode_Read,
-        0,
-        bytes,
-        .{ .callback = onMapped, .userdata1 = &mapped },
-    );
-    while (!mapped) gloc.poll();
-
-    const ptr: [*]const f32 = @ptrCast(@alignCast(
-        c.wgpuBufferGetConstMappedRange(staging, 0, bytes),
-    ));
-    @memcpy(out[0..self.len()], ptr[0..self.len()]);
-    c.wgpuBufferUnmap(staging);
-}
-
-fn onMapped(
-    status: c.WGPUMapAsyncStatus,
-    _: c.WGPUStringView,
-    userdata1: ?*anyopaque,
-    _: ?*anyopaque,
-) callconv(.c) void {
-    const flag: *bool = @ptrCast(@alignCast(userdata1.?));
-    flag.* = (status == c.WGPUMapAsyncStatus_Success);
-}
-
-// ── Dispatch helpers ──────────────────────────────────────────────────────────
-
-/// Encode + submit a 2-input, 1-output compute pass (used by add).
-fn dispatch2in1out(
-    gloc: *GpuAllocator,
-    pipeline: c.WGPUComputePipeline,
-    buf_a: c.WGPUBuffer,
-    buf_b: c.WGPUBuffer,
-    buf_out: c.WGPUBuffer,
-    bytes: u64,
-    n: u32,
-) !void {
-    const bgl = c.wgpuComputePipelineGetBindGroupLayout(pipeline, 0);
-    defer c.wgpuBindGroupLayoutRelease(bgl);
-
-    const entries = [_]c.WGPUBindGroupEntry{
-        .{ .binding = 0, .buffer = buf_a, .offset = 0, .size = bytes },
-        .{ .binding = 1, .buffer = buf_b, .offset = 0, .size = bytes },
-        .{ .binding = 2, .buffer = buf_out, .offset = 0, .size = bytes },
-    };
-    try submitPass(gloc, pipeline, &entries, n);
-}
-
-/// Create bind group, encode pass, submit. workgroup_size=64.
-fn submitPass(
-    gloc: *GpuAllocator,
-    pipeline: c.WGPUComputePipeline,
-    entries: []const c.WGPUBindGroupEntry,
-    n: u32,
-) !void {
-    const bgl = c.wgpuComputePipelineGetBindGroupLayout(pipeline, 0);
-    defer c.wgpuBindGroupLayoutRelease(bgl);
-
-    const bg = c.wgpuDeviceCreateBindGroup(gloc.device, &.{
-        .layout = bgl,
-        .entries = entries.ptr,
-        .entryCount = entries.len,
-    }) orelse return error.BindGroup;
-    defer c.wgpuBindGroupRelease(bg);
-
-    const enc = c.wgpuDeviceCreateCommandEncoder(gloc.device, null) orelse
-        return error.Encoder;
-    const pass = c.wgpuCommandEncoderBeginComputePass(enc, null);
-    c.wgpuComputePassEncoderSetPipeline(pass, pipeline);
-    c.wgpuComputePassEncoderSetBindGroup(pass, 0, bg, 0, null);
-    c.wgpuComputePassEncoderDispatchWorkgroups(pass, ceilDiv(n, 64), 1, 1);
-    c.wgpuComputePassEncoderEnd(pass);
-    c.wgpuComputePassEncoderRelease(pass);
-
-    const cmd = c.wgpuCommandEncoderFinish(enc, null);
-    defer c.wgpuCommandEncoderRelease(enc);
-    defer c.wgpuCommandBufferRelease(cmd);
-    c.wgpuQueueSubmit(gloc.queue, 1, &cmd);
-}
-
-fn ceilDiv(n: u32, d: u32) u32 {
-    return (n + d - 1) / d;
-}

src/c.zig

@@ -1 +0,0 @@
-pub const c = @cImport(@cInclude("wgpu.h"));

src/lib.zig

@@ -1 +1,5 @@
 pub const GpuAllocator = @import("GpuAllocator.zig");
+pub const GpuArena = @import("GpuArena.zig");
+pub const GpuBuffer = @import("GpuBuffer.zig");
+pub const GpuDevice = @import("GpuDevice.zig");
+pub const GpuProcess = @import("GpuProcess.zig");

src/main.zig

@@ -1,50 +0,0 @@
-const std = @import("std");
-const GpuAllocator = @import("GpuAllocator.zig");
-const Mat = @import("Mat.zig");
-
-pub fn main() !void {
-    var gloc = try GpuAllocator.init();
-    defer gloc.deinit();
-
-    // Input data: a[i] = i, b[i] = 15 - i  →  add should give all 15s
-    var data_a: [16]f32 = undefined;
-    var data_b: [16]f32 = undefined;
-    for (0..16) |i| {
-        data_a[i] = @floatFromInt(i);
-        data_b[i] = @floatFromInt(15 - i);
-    }
-
-    const a = try Mat.load(&gloc, &data_a, 4, 4);
-    defer a.deinit();
-    const b = try Mat.load(&gloc, &data_b, 4, 4);
-    defer b.deinit();
-
-    // a + b
-    const sum = try a.add(&gloc, b);
-    defer sum.deinit();
-
-    // sum * 2
-    const scaled = try sum.scale(&gloc, 2.0);
-    defer scaled.deinit();
-
-    // Read back
-    var out_sum: [16]f32 = undefined;
-    var out_scaled: [16]f32 = undefined;
-    try sum.read(&gloc, &out_sum);
-    try scaled.read(&gloc, &out_scaled);
-
-    // Print
-    std.debug.print("\na + b  (expect all 15):\n", .{});
-    printMat(&out_sum, 4, 4);
-
-    std.debug.print("\n(a + b) * 2  (expect all 30):\n", .{});
-    printMat(&out_scaled, 4, 4);
-}
-
-fn printMat(data: []const f32, rows: u32, cols: u32) void {
-    for (0..rows) |r| {
-        for (0..cols) |col|
-            std.debug.print("{d:6.0}", .{data[r * cols + col]});
-        std.debug.print("\n", .{});
-    }
-}

src/reference.zig

@@ -1,239 +0,0 @@
-// Minimal WebGPU compute in Zig: element-wise matrix addition
-// Uses wgpu-native C bindings.
-// Build: see ../build.zig
-//
-// Data flow:
-//   CPU (mat_a, mat_b) → GPU storage buffers → compute shader → GPU buf_c
-//   → staging buffer (mapped) → CPU read → print
-
-const std = @import("std");
-const c = @cImport(@cInclude("wgpu.h"));
-
-// ── Config ────────────────────────────────────────────────────────────────────
-const ROWS: u32 = 4;
-const COLS: u32 = 4;
-const N = ROWS * COLS; // 16 elements
-const BUF_BYTES = N * @sizeOf(f32);
-
-// ── WGSL Compute Shader ───────────────────────────────────────────────────────
-//   workgroup_size(4,4) matches one full 4×4 matrix → dispatch(1,1,1)
-const SHADER =
-    \\@group(0) @binding(0) var<storage, read>       mat_a : array<f32>;
-    \\@group(0) @binding(1) var<storage, read>       mat_b : array<f32>;
-    \\@group(0) @binding(2) var<storage, read_write> mat_c : array<f32>;
-    \\
-    \\@compute @workgroup_size(4, 4)
-    \\fn main(@builtin(global_invocation_id) gid : vec3<u32>) {
-    \\    let idx = gid.y * 4u + gid.x;
-    \\    if (idx < arrayLength(&mat_c)) {
-    \\        mat_c[idx] = mat_a[idx] + mat_b[idx];
-    \\    }
-    \\}
-;
-
-// ── Callback state ────────────────────────────────────────────────────────────
-const Ctx = struct {
-    adapter: c.WGPUAdapter = null,
-    device: c.WGPUDevice = null,
-};
-
-fn onAdapter(
-    status: c.WGPURequestAdapterStatus,
-    adapter: c.WGPUAdapter,
-    _: c.WGPUStringView,
-    userdata1: ?*anyopaque,
-    _: ?*anyopaque,
-) callconv(.c) void {
-    if (status != c.WGPURequestAdapterStatus_Success) {
-        std.log.err("Adapter request failed (status={d})", .{status});
-        return;
-    }
-    const ctx: *Ctx = @ptrCast(@alignCast(userdata1.?));
-    ctx.adapter = adapter;
-}
-
-fn onDevice(
-    status: c.WGPURequestDeviceStatus,
-    device: c.WGPUDevice,
-    _: c.WGPUStringView,
-    userdata1: ?*anyopaque,
-    _: ?*anyopaque,
-) callconv(.c) void {
-    if (status != c.WGPURequestDeviceStatus_Success) {
-        std.log.err("Device request failed (status={d})", .{status});
-        return;
-    }
-    const ctx: *Ctx = @ptrCast(@alignCast(userdata1.?));
-    ctx.device = device;
-    std.debug.print("{?}", .{device});
-}
-
-fn onMapped(
-    status: c.WGPUMapAsyncStatus,
-    _: c.WGPUStringView,
-    userdata1: ?*anyopaque,
-    _: ?*anyopaque,
-) callconv(.c) void {
-    const flag: *bool = @ptrCast(@alignCast(userdata1.?));
-    flag.* = (status == c.WGPUMapAsyncStatus_Success);
-}
-
-fn sv(s: []const u8) c.WGPUStringView {
-    return .{ .data = s.ptr, .length = s.len };
-}
-
-// ── Main ──────────────────────────────────────────────────────────────────────
-pub fn main() !void {
-
-    // 1. Instance ──────────────────────────────────────────────────────────────
-    const instance = c.wgpuCreateInstance(&std.mem.zeroes(c.WGPUInstanceDescriptor)) orelse
-        return error.NoInstance;
-    defer c.wgpuInstanceRelease(instance);
-
-    // 2. Adapter (async → poll) ────────────────────────────────────────────────
-    var ctx = Ctx{};
-    _ = c.wgpuInstanceRequestAdapter(
-        instance,
-        &.{ .powerPreference = c.WGPUPowerPreference_HighPerformance },
-        .{ .callback = onAdapter, .userdata1 = &ctx },
-    );
-    c.wgpuInstanceProcessEvents(instance); // drive callbacks
-    const adapter = ctx.adapter orelse return error.NoAdapter;
-    defer c.wgpuAdapterRelease(adapter);
-
-    // 3. Device ────────────────────────────────────────────────────────────────
-    _ = c.wgpuAdapterRequestDevice(adapter, null, .{ .callback = onDevice, .userdata1 = &ctx });
-    c.wgpuInstanceProcessEvents(instance);
-    const device = ctx.device orelse return error.NoDevice;
-    defer c.wgpuDeviceRelease(device);
-
-    const queue = c.wgpuDeviceGetQueue(device);
-    defer c.wgpuQueueRelease(queue);
-
-    // 4. Input data ────────────────────────────────────────────────────────────
-    //   mat_a[i] = i        (0 … 15)
-    //   mat_b[i] = 15 − i   → every element of mat_c should equal 15
-    var mat_a: [N]f32 = undefined;
-    var mat_b: [N]f32 = undefined;
-    for (0..N) |i| {
-        mat_a[i] = @floatFromInt(i);
-        mat_b[i] = @floatFromInt(N - 1 - i);
-    }
-
-    // 5. GPU Buffers ───────────────────────────────────────────────────────────
-    const buf_a = c.wgpuDeviceCreateBuffer(device, &.{
-        .usage = c.WGPUBufferUsage_Storage | c.WGPUBufferUsage_CopyDst,
-        .size = BUF_BYTES,
-    }) orelse return error.BufferA;
-
-    const buf_b = c.wgpuDeviceCreateBuffer(device, &.{
-        .usage = c.WGPUBufferUsage_Storage | c.WGPUBufferUsage_CopyDst,
-        .size = BUF_BYTES,
-    }) orelse return error.BufferB;
-
-    // buf_c: GPU-only result; staging: CPU-readable copy
-    const buf_c = c.wgpuDeviceCreateBuffer(device, &.{
-        .usage = c.WGPUBufferUsage_Storage | c.WGPUBufferUsage_CopySrc,
-        .size = BUF_BYTES,
-    }) orelse return error.BufferC;
-
-    const buf_staging = c.wgpuDeviceCreateBuffer(device, &.{
-        .usage = c.WGPUBufferUsage_MapRead | c.WGPUBufferUsage_CopyDst,
-        .size = BUF_BYTES,
-    }) orelse return error.BufferStaging;
-
-    defer c.wgpuBufferRelease(buf_a);
-    defer c.wgpuBufferRelease(buf_b);
-    defer c.wgpuBufferRelease(buf_c);
-    defer c.wgpuBufferRelease(buf_staging);
-
-    // Upload inputs
-    c.wgpuQueueWriteBuffer(queue, buf_a, 0, &mat_a, BUF_BYTES);
-    c.wgpuQueueWriteBuffer(queue, buf_b, 0, &mat_b, BUF_BYTES);
-
-    // 6. Shader module ─────────────────────────────────────────────────────────
-    // ✅ New API (0.20+)
-    var wgsl_src = c.WGPUShaderSourceWGSL{
-        .chain = .{ .sType = c.WGPUSType_ShaderSourceWGSL },
-        .code = sv(SHADER),
-    };
-    const shader = c.wgpuDeviceCreateShaderModule(device, &.{
-        .nextInChain = @ptrCast(&wgsl_src),
-    }) orelse return error.Shader;
-
-    // 7. Compute pipeline (layout auto-inferred from shader) ───────────────────
-    // ✅
-    const pipeline = c.wgpuDeviceCreateComputePipeline(device, &.{
-        .compute = .{
-            .module = shader,
-            .entryPoint = sv("main"),
-        },
-    }) orelse return error.Pipeline;
-    defer c.wgpuComputePipelineRelease(pipeline);
-
-    // 8. Bind group ────────────────────────────────────────────────────────────
-    const bgl = c.wgpuComputePipelineGetBindGroupLayout(pipeline, 0);
-    defer c.wgpuBindGroupLayoutRelease(bgl);
-
-    const entries = [_]c.WGPUBindGroupEntry{
-        .{ .binding = 0, .buffer = buf_a, .offset = 0, .size = BUF_BYTES },
-        .{ .binding = 1, .buffer = buf_b, .offset = 0, .size = BUF_BYTES },
-        .{ .binding = 2, .buffer = buf_c, .offset = 0, .size = BUF_BYTES },
-    };
-    const bind_group = c.wgpuDeviceCreateBindGroup(device, &.{
-        .layout = bgl,
-        .entries = &entries,
-        .entryCount = entries.len,
-    }) orelse return error.BindGroup;
-    defer c.wgpuBindGroupRelease(bind_group);
-
-    // 9. Encode compute pass + buffer copy ────────────────────────────────────
-    const encoder = c.wgpuDeviceCreateCommandEncoder(device, null) orelse
-        return error.Encoder;
-
-    const pass = c.wgpuCommandEncoderBeginComputePass(encoder, null);
-    c.wgpuComputePassEncoderSetPipeline(pass, pipeline);
-    c.wgpuComputePassEncoderSetBindGroup(pass, 0, bind_group, 0, null);
-    // dispatch(1,1,1): one workgroup of size (4,4) covers the whole 4×4 matrix
-    c.wgpuComputePassEncoderDispatchWorkgroups(pass, 1, 1, 1);
-    c.wgpuComputePassEncoderEnd(pass);
-    c.wgpuComputePassEncoderRelease(pass);
-
-    // Copy result buffer → CPU-readable staging buffer
-    c.wgpuCommandEncoderCopyBufferToBuffer(encoder, buf_c, 0, buf_staging, 0, BUF_BYTES);
-
-    const cmdbuf = c.wgpuCommandEncoderFinish(encoder, null);
-    defer c.wgpuCommandEncoderRelease(encoder);
-    defer c.wgpuCommandBufferRelease(cmdbuf);
-
-    // 10. Submit ───────────────────────────────────────────────────────────────
-    c.wgpuQueueSubmit(queue, 1, &cmdbuf);
-
-    // 11. Map staging buffer back to CPU ──────────────────────────────────────
-    var mapped = false;
-    _ = c.wgpuBufferMapAsync(
-        buf_staging,
-        c.WGPUMapMode_Read,
-        0,
-        BUF_BYTES,
-        .{ .callback = onMapped, .userdata1 = &mapped },
-    );
-    // Poll the device until the async map completes
-    while (!mapped) _ = c.wgpuDevicePoll(device, 1, null);
-
-    const ptr: [*]const f32 = @ptrCast(@alignCast(
-        c.wgpuBufferGetConstMappedRange(buf_staging, 0, BUF_BYTES),
-    ));
-    const result = ptr[0..N];
-
-    // 12. Print ────────────────────────────────────────────────────────────────
-    std.debug.print("\nmat_a + mat_b  ({d}×{d}):\n", .{ ROWS, COLS });
-    for (0..ROWS) |r| {
-        for (0..COLS) |col|
-            std.debug.print("{d:6.0}", .{result[r * COLS + col]});
-        std.debug.print("\n", .{});
-    }
-    // Expected output: every cell = 15.0
-
-    c.wgpuBufferUnmap(buf_staging);
-}

src/shaders.zig

@@ -1,28 +0,0 @@
-pub const SHADER_ADD =
-    \\@group(0) @binding(0) var<storage, read>       a : array<f32>;
-    \\@group(0) @binding(1) var<storage, read>       b : array<f32>;
-    \\@group(0) @binding(2) var<storage, read_write> out : array<f32>;
-    \\
-    \\@compute @workgroup_size(64)
-    \\fn main(@builtin(global_invocation_id) gid : vec3<u32>) {
-    \\    let i = gid.x;
-    \\    if (i < arrayLength(&out)) {
-    \\        out[i] = a[i] + b[i];
-    \\    }
-    \\}
-;
-
-pub const SHADER_SCALE =
-    \\struct Uniforms { scalar : f32 }
-    \\@group(0) @binding(0) var<storage, read>       a   : array<f32>;
-    \\@group(0) @binding(1) var<storage, read_write> out : array<f32>;
-    \\@group(0) @binding(2) var<uniform>             u   : Uniforms;
-    \\
-    \\@compute @workgroup_size(64)
-    \\fn main(@builtin(global_invocation_id) gid : vec3<u32>) {
-    \\    let i = gid.x;
-    \\    if (i < arrayLength(&out)) {
-    \\        out[i] = a[i] * u.scalar;
-    \\    }
-    \\}
-;

src/utils.zig

@@ -0,0 +1,5 @@
+pub const c = @cImport(@cInclude("wgpu.h"));
+
+pub fn sv(s: []const u8) c.WGPUStringView {
+    return .{ .data = s.ptr, .length = s.len };
+}