tiny.accy.tensor.function

1const std = @import("std");2const batch = @import("batch.zig");3const execute = @import("execute.zig");4const gradient = @import("grad.zig");5const lower = @import("lower.zig");6const program_mod = @import("program.zig");7const trace = @import("trace/root.zig");8const types = @import("type/root.zig");910pub const Function = struct {11    allocator: std.mem.Allocator,12    graph: program_mod.Program,13    root: ?*Function,14    derived: std.ArrayListUnmanaged(*Function),1516    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    }2930    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    }4041    pub fn grad(self: *Function, options: gradient.Options) !*Function {42        return self.adopt(try gradient.grad(self.allocator, &self.graph, options));43    }4445    pub fn valueAndGrad(self: *Function, options: gradient.Options) !*Function {46        return self.adopt(try gradient.valueAndGrad(self.allocator, &self.graph, options));47    }4849    pub fn vmap(self: *Function, options: batch.Options) !*Function {50        return self.adopt(try batch.vmap(self.allocator, &self.graph, options));51    }5253    pub fn compileCpu(self: *const Function) !execute.Cpu {54        return execute.Cpu.init(self.allocator, &self.graph);55    }5657    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    }6465    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};8182fn 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}8687test "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();9394    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 } });100101    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}107108test "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();116117    const step = try f.valueAndGrad(.{ .wrt = &.{ 0, 1 } });118    var executor = try step.compileCpu();119    defer executor.deinit();120121    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    };127128    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..]);137138    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();145146    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..]);155156    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}161162test "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;167168    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();173174    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();180181    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    }187188    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..]);198199    try std.testing.expectEqualSlices(f32, y[0..], x_grad[0..]);200    try std.testing.expectEqualSlices(f32, x[0..], y_grad[0..]);201}

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