lib/accy/src/tensor/function.zig
daab053ee43316e1809a84551d573ddd1e5bf3d2
1 const std = @import("std");
2 const batch = @import("batch.zig");
3 const execute = @import("execute.zig");
4 const gradient = @import("grad.zig");
5 const lower = @import("lower.zig");
6 const program_mod = @import("program.zig");
7 const trace = @import("trace/root.zig");
8 const types = @import("type/root.zig");
9
10 pub const Function = struct {
11 allocator: std.mem.Allocator,
12 graph: program_mod.Program,
13 root: ?*Function,
14 derived: std.ArrayListUnmanaged(*Function),
15
16 pub fn define(
17 allocator: std.mem.Allocator,
18 name: []const u8,
19 specs: []const types.Spec,
20 comptime body: anytype,
21 ) !Function {
22 return .{
23 .allocator = allocator,
24 .graph = try trace.define(allocator, name, specs, body),
25 .root = null,
26 .derived = .empty,
27 };
28 }
29
30 pub fn deinit(self: *Function) void {
31 std.debug.assert(self.root == null);
32 for (self.derived.items) |child| {
33 child.graph.deinit();
34 self.allocator.destroy(child);
35 }
36 self.derived.deinit(self.allocator);
37 self.graph.deinit();
38 self.* = undefined;
39 }
40
41 pub fn grad(self: *Function, options: gradient.Options) !*Function {
42 return self.adopt(try gradient.grad(self.allocator, &self.graph, options));
43 }
44
45 pub fn valueAndGrad(self: *Function, options: gradient.Options) !*Function {
46 return self.adopt(try gradient.valueAndGrad(self.allocator, &self.graph, options));
47 }
48
49 pub fn vmap(self: *Function, options: batch.Options) !*Function {
50 return self.adopt(try batch.vmap(self.allocator, &self.graph, options));
51 }
52
53 pub fn compileCpu(self: *const Function) !execute.Cpu {
54 return execute.Cpu.init(self.allocator, &self.graph);
55 }
56
57 pub fn compileFragment(
58 self: *const Function,
59 handle: lower.BackendHandle,
60 options: lower.FragmentCompilerOptions,
61 ) !*lower.CompiledFragment {
62 return lower.compileFragment(self.allocator, handle, &self.graph, options);
63 }
64
65 fn adopt(self: *Function, graph: program_mod.Program) !*Function {
66 var owned = graph;
67 errdefer owned.deinit();
68 const root = self.root orelse self;
69 const child = try root.allocator.create(Function);
70 errdefer root.allocator.destroy(child);
71 child.* = .{
72 .allocator = root.allocator,
73 .graph = owned,
74 .root = root,
75 .derived = .empty,
76 };
77 try root.derived.append(root.allocator, child);
78 return child;
79 }
80 };
81
82 fn productLoss(_: *trace.Builder, args: []const trace.Value) !trace.Value {
83 const product = try args[0].mul(args[1]);
84 return try product.sum(.lane);
85 }
86
87 test "tensor function owns every derived graph through one deinit" {
88 var f = try Function.define(std.testing.allocator, "function_ownership", &.{
89 types.spec(.f32, .{ .lane = 4 }),
90 types.spec(.f32, .{ .lane = 4 }),
91 }, productLoss);
92 defer f.deinit();
93
94 const g = try f.grad(.{ .wrt = &.{ 0, 1 } });
95 const bg = try g.vmap(.{
96 .axis_size = 8,
97 .in_axes = &.{ batch.mapped(0), batch.mapped(0) },
98 });
99 const step = try f.valueAndGrad(.{ .wrt = &.{ 0, 1 } });
100
101 try std.testing.expectEqual(@as(usize, 2), g.graph.outputs.len);
102 try types.expectExtents(&.{ 8, 4 }, bg.graph.typeOf(bg.graph.outputs[0]));
103 try types.expectExtents(&.{ 8, 4 }, bg.graph.typeOf(bg.graph.outputs[1]));
104 try std.testing.expectEqual(@as(usize, 3), step.graph.outputs.len);
105 try types.expectExtents(&.{}, step.graph.typeOf(step.graph.outputs[0]));
106 }
107
108 test "tensor function value and grad matches the free-function path numerically" {
109 try @import("../fixture/root.zig").requireNativeCpuArtifacts();
110 const allocator = std.testing.allocator;
111 var f = try Function.define(allocator, "function_parity", &.{
112 types.spec(.f32, .{ .lane = 4 }),
113 types.spec(.f32, .{ .lane = 4 }),
114 }, productLoss);
115 defer f.deinit();
116
117 const step = try f.valueAndGrad(.{ .wrt = &.{ 0, 1 } });
118 var executor = try step.compileCpu();
119 defer executor.deinit();
120
121 const x = [_]f32{ 1.0, 2.0, 3.0, 4.0 };
122 const y = [_]f32{ 0.5, -1.0, 2.0, 0.25 };
123 const inputs = [_][]const u8{
124 std.mem.sliceAsBytes(x[0..]),
125 std.mem.sliceAsBytes(y[0..]),
126 };
127
128 var loss_value: f32 = 0;
129 var x_grad = @as([4]f32, @splat(0));
130 var y_grad = @as([4]f32, @splat(0));
131 const outputs = [_][]u8{
132 std.mem.asBytes(&loss_value),
133 std.mem.sliceAsBytes(x_grad[0..]),
134 std.mem.sliceAsBytes(y_grad[0..]),
135 };
136 try executor.launch(allocator, inputs[0..], outputs[0..]);
137
138 var source = try trace.define(allocator, "function_parity_reference", &.{
139 types.spec(.f32, .{ .lane = 4 }),
140 types.spec(.f32, .{ .lane = 4 }),
141 }, productLoss);
142 defer source.deinit();
143 var reference = try gradient.valueAndGrad(allocator, &source, .{ .wrt = &.{ 0, 1 } });
144 defer reference.deinit();
145
146 var reference_loss: f32 = 0;
147 var reference_x_grad = @as([4]f32, @splat(0));
148 var reference_y_grad = @as([4]f32, @splat(0));
149 const reference_outputs = [_][]u8{
150 std.mem.asBytes(&reference_loss),
151 std.mem.sliceAsBytes(reference_x_grad[0..]),
152 std.mem.sliceAsBytes(reference_y_grad[0..]),
153 };
154 try execute.runCpu(allocator, &reference, inputs[0..], reference_outputs[0..]);
155
156 try std.testing.expectEqual(reference_loss, loss_value);
157 try std.testing.expectEqualSlices(f32, reference_x_grad[0..], x_grad[0..]);
158 try std.testing.expectEqualSlices(f32, reference_y_grad[0..], y_grad[0..]);
159 try std.testing.expectApproxEqAbs(@as(f32, 5.5), loss_value, 0.000001);
160 }
161
162 test "tensor function composes per-example gradients through vmap of grad" {
163 try @import("../fixture/root.zig").requireNativeCpuArtifacts();
164 const allocator = std.testing.allocator;
165 const rows = 3;
166 const lanes = 4;
167
168 var f = try Function.define(allocator, "function_vmap_grad", &.{
169 types.spec(.f32, .{ .lane = lanes }),
170 types.spec(.f32, .{ .lane = lanes }),
171 }, productLoss);
172 defer f.deinit();
173
174 const per_example = try (try f.grad(.{ .wrt = &.{ 0, 1 } })).vmap(.{
175 .axis_size = rows,
176 .in_axes = &.{ batch.mapped(0), batch.mapped(0) },
177 });
178 var executor = try per_example.compileCpu();
179 defer executor.deinit();
180
181 var x: [rows * lanes]f32 = undefined;
182 var y: [rows * lanes]f32 = undefined;
183 for (0..rows * lanes) |index| {
184 x[index] = @as(f32, @floatFromInt(index)) * 0.5 - 2.0;
185 y[index] = 1.0 - @as(f32, @floatFromInt(index)) * 0.25;
186 }
187
188 var x_grad = @as([(rows * lanes)]f32, @splat(0));
189 var y_grad = @as([(rows * lanes)]f32, @splat(0));
190 const outputs = [_][]u8{
191 std.mem.sliceAsBytes(x_grad[0..]),
192 std.mem.sliceAsBytes(y_grad[0..]),
193 };
194 try executor.launch(allocator, &.{
195 std.mem.sliceAsBytes(x[0..]),
196 std.mem.sliceAsBytes(y[0..]),
197 }, outputs[0..]);
198
199 try std.testing.expectEqualSlices(f32, y[0..], x_grad[0..]);
200 try std.testing.expectEqualSlices(f32, x[0..], y_grad[0..]);
201 }