lib/accy/src/tensor/execute.zig

daab053ee43316e1809a84551d573ddd1e5bf3d2

  1 const std = @import("std");
  2 const gpu = @import("gpu");
  3 const accy = @import("../root.zig");
  4 const lower = @import("lower.zig");
  5 const program_mod = @import("program.zig");
  6 const trace = @import("trace/root.zig");
  7 const types = @import("type/root.zig");
  8 
  9 pub const Cpu = struct {
 10     allocator: std.mem.Allocator,
 11     state: *gpu.cpu.State,
 12     fragment: *lower.LoadedFragment,
 13 
 14     pub fn init(allocator: std.mem.Allocator, program: *const program_mod.Program) !Cpu {
 15         return initWith(allocator, program, .{ .artifact_format = .cpu_object });
 16     }
 17 
 18     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     }
 31 
 32     pub fn deinit(self: *Cpu) void {
 33         self.fragment.deinit();
 34         self.state.deinit();
 35         self.allocator.destroy(self.state);
 36         self.* = undefined;
 37     }
 38 
 39     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 };
 48 
 49 pub 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 }
 59 
 60 fn 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 }
 64 
 65 test "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();
 73 
 74     var executor = try Cpu.init(allocator, &graph);
 75     defer executor.deinit();
 76 
 77     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..])};
 81 
 82     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..]);
 87 
 88     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 }
 95 
 96 test "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();
104 
105     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 }
115 
116 fn constantBody(builder: *trace.Builder, args: []const trace.Value) !trace.Value {
117     _ = args;
118     return builder.full(.f32, .{ .lane = 4 }, 3.0);
119 }
120 
121 test "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();
127 
128     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 }