tiny.accy.tensor.execute

1const std = @import("std");2const gpu = @import("gpu");3const accy = @import("../root.zig");4const lower = @import("lower.zig");5const program_mod = @import("program.zig");6const trace = @import("trace/root.zig");7const types = @import("type/root.zig");89pub const Cpu = struct {10    allocator: std.mem.Allocator,11    state: *gpu.cpu.State,12    fragment: *lower.LoadedFragment,1314    pub fn init(allocator: std.mem.Allocator, program: *const program_mod.Program) !Cpu {15        return initWith(allocator, program, .{ .artifact_format = .cpu_object });16    }1718    pub fn initWith(19        allocator: std.mem.Allocator,20        program: *const program_mod.Program,21        options: lower.FragmentCompilerOptions,22    ) !Cpu {23        const state = try allocator.create(gpu.cpu.State);24        errdefer allocator.destroy(state);25        state.* = gpu.cpu.State.init(allocator);26        errdefer state.deinit();27        const compiled = try lower.compileFragment(allocator, state.handle(), program, options);28        const fragment = try accy.executable.loadFragment(allocator, state.handle(), compiled, options);29        return .{ .allocator = allocator, .state = state, .fragment = fragment };30    }3132    pub fn deinit(self: *Cpu) void {33        self.fragment.deinit();34        self.state.deinit();35        self.allocator.destroy(self.state);36        self.* = undefined;37    }3839    pub fn launch(40        self: *const Cpu,41        scratch: std.mem.Allocator,42        inputs: []const []const u8,43        outputs: []const []u8,44    ) !void {45        try accy.executable.invoke(self.fragment, scratch, scratch, inputs, outputs);46    }47};4849pub fn runCpu(50    allocator: std.mem.Allocator,51    program: *const program_mod.Program,52    inputs: []const []const u8,53    outputs: []const []u8,54) !void {55    var executor = try Cpu.init(allocator, program);56    defer executor.deinit();57    try executor.launch(allocator, inputs, outputs);58}5960fn scaleShiftBody(_: *trace.Builder, args: []const trace.Value) !trace.Value {61    const doubled = try args[0].add(args[0]);62    return doubled.add(args[1]);63}6465test "tensor cpu executor launches one program repeatedly" {66    try @import("../fixture/root.zig").requireNativeCpuArtifacts();67    const allocator = std.testing.allocator;68    var graph = try trace.define(allocator, "execute_scale_shift", &.{69        types.spec(.f32, .{ .lane = 4 }),70        types.spec(.f32, .{ .lane = 4 }),71    }, scaleShiftBody);72    defer graph.deinit();7374    var executor = try Cpu.init(allocator, &graph);75    defer executor.deinit();7677    const base = [_]f32{ 1.0, 2.0, 3.0, 4.0 };78    var shift = [_]f32{ 0.5, 0.5, 0.5, 0.5 };79    var out = @as([4]f32, @splat(0));80    const outputs = [_][]u8{std.mem.sliceAsBytes(out[0..])};8182    try executor.launch(allocator, &.{83        std.mem.sliceAsBytes(base[0..]),84        std.mem.sliceAsBytes(shift[0..]),85    }, outputs[0..]);86    try std.testing.expectEqualSlices(f32, &.{ 2.5, 4.5, 6.5, 8.5 }, out[0..]);8788    shift = .{ -1.0, -1.0, -1.0, -1.0 };89    try executor.launch(allocator, &.{90        std.mem.sliceAsBytes(base[0..]),91        std.mem.sliceAsBytes(shift[0..]),92    }, outputs[0..]);93    try std.testing.expectEqualSlices(f32, &.{ 1.0, 3.0, 5.0, 7.0 }, out[0..]);94}9596test "tensor runCpu runs a program once" {97    try @import("../fixture/root.zig").requireNativeCpuArtifacts();98    const allocator = std.testing.allocator;99    var graph = try trace.define(allocator, "execute_run_once", &.{100        types.spec(.f32, .{ .lane = 3 }),101        types.spec(.f32, .{ .lane = 3 }),102    }, scaleShiftBody);103    defer graph.deinit();104105    const base = [_]f32{ 1.0, -2.0, 0.25 };106    const shift = [_]f32{ 0.0, 1.0, -0.25 };107    var out = @as([3]f32, @splat(0));108    const outputs = [_][]u8{std.mem.sliceAsBytes(out[0..])};109    try runCpu(allocator, &graph, &.{110        std.mem.sliceAsBytes(base[0..]),111        std.mem.sliceAsBytes(shift[0..]),112    }, outputs[0..]);113    try std.testing.expectEqualSlices(f32, &.{ 2.0, -3.0, 0.25 }, out[0..]);114}115116fn constantBody(builder: *trace.Builder, args: []const trace.Value) !trace.Value {117    _ = args;118    return builder.full(.f32, .{ .lane = 4 }, 3.0);119}120121test "tensor runCpu serves constant-only outputs" {122    const allocator = std.testing.allocator;123    var graph = try trace.define(allocator, "execute_constant_output", &.{124        types.spec(.f32, .{ .lane = 3 }),125    }, constantBody);126    defer graph.deinit();127128    const ignored = [_]f32{ 1.0, 2.0, 3.0 };129    var out = @as([4]f32, @splat(0));130    const outputs = [_][]u8{std.mem.sliceAsBytes(out[0..])};131    try runCpu(allocator, &graph, &.{std.mem.sliceAsBytes(ignored[0..])}, outputs[0..]);132    try std.testing.expectEqualSlices(f32, &.{ 3.0, 3.0, 3.0, 3.0 }, out[0..]);133}

1pub const execute = @import("execute.zig");