tiny.accy.tensor.interpret

1pub const Graph = struct {2    builder: *trace.Builder,34    pub const Value: type = trace.Value;5    pub const Result: type = program_mod.Program;67    pub fn operation(self: *@This(), step: *step_mod.Step(Value)) !Value {8        var buffer: [program_mod.max_operation_operands]Value = undefined;9        return self.bind(step.op, step_mod.arguments(Value, step.op, step.values, &buffer));10    }1112    pub fn bind(self: *@This(), op: *const program_mod.Operation, args: []const Value) !Value {13        return self.builder.operation(op, args);14    }1516    pub fn finish(self: *@This(), outputs: []const Value) !Result {17        return self.builder.finish(outputs);18    }1920    pub fn builderHandle(self: *@This()) *trace.Builder {21        return self.builder;22    }23};

1pub fn BindContext(comptime BoundValue: type, comptime Handle: type) type {2    return struct {3        next: Handle,4        op: *const program_mod.Operation,5        args: []const BoundValue,67        pub fn default(self: *@This()) !BoundValue {8            return self.next.bind(self.op, self.args);9        }1011        pub fn arg(self: *@This(), index: usize) BoundValue {12            return self.args[index];13        }1415        pub fn builderHandle(self: *@This()) *trace.Builder {16            return self.next.builderHandle();17        }18    };19}

1pub fn FinishContext(comptime BoundValue: type, comptime Handle: type) type {2    return struct {3        next: Handle,45        pub fn default(self: *@This(), outputs: []const BoundValue) !pointerChild(Handle).Result {6            return self.next.finish(outputs);7        }89        pub fn builderHandle(self: *@This()) *trace.Builder {10            return self.next.builderHandle();11        }12    };13}

1pub fn LeafBindContext(comptime BoundValue: type) type {2    return struct {3        op: *const program_mod.Operation,4        args: []const BoundValue,56        pub fn arg(self: *@This(), index: usize) BoundValue {7            return self.args[index];8        }9    };10}

1pub fn result(comptime Initial: type) type {2    if (comptime @hasDecl(Initial, "attach")) return @TypeOf(@as(Initial, undefined).attach(@as(graph_mod.Graph, undefined))).Result;3    return Normalized(Initial).Result;4}

1pub fn run(2    allocator: std.mem.Allocator,3    source: *const program_mod.Program,4    initial: anytype,5) !result(@TypeOf(initial)) {6    var state = normalize(initial);7    const State = @TypeOf(state);8    const Value = State.Value;910    const values = try allocator.alloc(Value, source.valueCount());11    defer allocator.free(values);1213    for (source.operations) |*op| {14        var step = step_mod.Step(Value){15            .source = source,16            .values = values,17            .op = op,18        };19        values[op.id.index] = try state.operation(&step);20    }2122    const outputs = try allocator.alloc(Value, source.outputs.len);23    defer allocator.free(outputs);24    for (source.outputs, 0..) |id, index| {25        outputs[index] = values[id.index];26    }2728    return try state.finish(outputs);29}

1pub fn Semantics(comptime BoundValue: type, comptime Impl: type) type {2    return struct {3        impl: Impl,45        pub const Value: type = BoundValue;6        pub const Result: type = Impl.Result;78        pub fn operation(self: *@This(), step: *step_mod.Step(Value)) !Value {9            var buffer: [program_mod.max_operation_operands]Value = undefined;10            return self.bind(step.op, step_mod.arguments(Value, step.op, step.values, &buffer));11        }1213        pub fn bind(self: *@This(), op: *const program_mod.Operation, args: []const Value) !Value {14            if (comptime @hasDecl(Impl, "bind")) {15                return self.impl.bind(op, args);16            }1718            var ctx = context.LeafBindContext(Value){19                .op = op,20                .args = args,21            };22            if (try dispatch.primitive(Value, &self.impl, &ctx)) |value| return value;23            if (comptime @hasDecl(Impl, "default")) return self.impl.default(&ctx);24            @compileError("tensor leaf semantics without bind must provide primitive handlers and default");25        }2627        pub fn finish(self: *@This(), outputs: []const Value) !Result {28            return self.impl.finish(outputs);29        }3031        pub fn builderHandle(self: *@This()) *trace.Builder {32            return self.impl.builderHandle();33        }34    };35}

1pub fn semantics(comptime BoundValue: type, impl: anytype) Semantics(BoundValue, @TypeOf(impl)) {2    return .{ .impl = impl };3}

1pub fn Layer(comptime BoundValue: type, comptime Next: type, comptime Impl: type) type {2    return dispatch.Layer(BoundValue, Next, Impl);3}

1pub fn Stack(comptime Specs: type) type {2    _ = stackLength(Specs);3    return struct {4        specs: Specs,56        pub fn attach(self: @This(), next: anytype) StackAttach(0, Specs, @TypeOf(next)) {7            return attachFrom(0, self.specs, next);8        }9    };10}

1pub fn With(comptime Impl: type) type {2    return struct {3        impl: Impl,45        pub fn attach(self: @This(), next: anytype) Layer(stateValue(@TypeOf(next)), @TypeOf(next), Impl) {6            return .{7                .next = next,8                .impl = self.impl,9            };10        }11    };12}

1pub fn bind(impl: anytype) With(@TypeOf(impl)) {2    return with(impl);3}

1pub fn layer(comptime BoundValue: type, next: anytype, impl: anytype) Layer(BoundValue, @TypeOf(next), @TypeOf(impl)) {2    return .{3        .next = next,4        .impl = impl,5    };6}

1pub fn stack(specs: anytype) Stack(@TypeOf(specs)) {2    return .{ .specs = specs };3}

1pub fn with(impl: anytype) With(@TypeOf(impl)) {2    return .{ .impl = impl };3}

1pub fn Step(comptime Value: type) type {2    return struct {3        source: *const program_mod.Program,4        values: []const Value,5        op: *const program_mod.Operation,67        pub fn read(self: @This(), id: program_mod.Id) Value {8            return self.values[id.index];9        }1011        pub fn typeOf(self: @This(), id: program_mod.Id) program_mod.Type {12            return self.source.typeOf(id);13        }14    };15}

1pub fn arguments(comptime Value: type, op: *const program_mod.Operation, values: []const Value, buffer: *[program_mod.max_operation_operands]Value) []const Value {2    return switch (op.kind) {3        .parameter, .constant, .iota => buffer[0..0],4        .unary => |unary| blk: {5            buffer[0] = values[unary.input.index];6            break :blk buffer[0..1];7        },8        .binary => |binary| blk: {9            buffer[0] = values[binary.lhs.index];10            buffer[1] = values[binary.rhs.index];11            break :blk buffer[0..2];12        },13        .broadcast => |broadcast| blk: {14            buffer[0] = values[broadcast.input.index];15            break :blk buffer[0..1];16        },17        .broadcast_in_dim => |broadcast| blk: {18            buffer[0] = values[broadcast.input.index];19            break :blk buffer[0..1];20        },21        .reshape => |reshape| blk: {22            buffer[0] = values[reshape.input.index];23            break :blk buffer[0..1];24        },25        .transpose => |transpose| blk: {26            buffer[0] = values[transpose.input.index];27            break :blk buffer[0..1];28        },29        .reduce => |reduce| blk: {30            buffer[0] = values[reduce.input.index];31            buffer[1] = values[reduce.init.index];32            break :blk buffer[0..2];33        },34        .gather => |gather| blk: {35            buffer[0] = values[gather.input.index];36            buffer[1] = values[gather.indices.index];37            break :blk buffer[0..2];38        },39        .scatter_add => |scatter_add| blk: {40            buffer[0] = values[scatter_add.input.index];41            buffer[1] = values[scatter_add.indices.index];42            buffer[2] = values[scatter_add.updates.index];43            break :blk buffer[0..3];44        },45        .sparse_cross_entropy => |sparse_cross_entropy| blk: {46            buffer[0] = values[sparse_cross_entropy.logits.index];47            buffer[1] = values[sparse_cross_entropy.targets.index];48            break :blk buffer[0..2];49        },50        .dot_general => |dot| blk: {51            buffer[0] = values[dot.lhs.index];52            buffer[1] = values[dot.rhs.index];53            break :blk buffer[0..2];54        },55        .compare => |compare| blk: {56            buffer[0] = values[compare.lhs.index];57            buffer[1] = values[compare.rhs.index];58            break :blk buffer[0..2];59        },60        .select => |select| blk: {61            buffer[0] = values[select.pred.index];62            buffer[1] = values[select.on_true.index];63            buffer[2] = values[select.on_false.index];64            break :blk buffer[0..3];65        },66        .custom_call => |custom| blk: {67            for (custom.operands, 0..) |operand, index| buffer[index] = values[operand.index];68            break :blk buffer[0..custom.operands.len];69        },70        .scan => |scan| blk: {71            for (scan.inits, 0..) |init, index| buffer[index] = values[init.index];72            break :blk buffer[0..scan.inits.len];73        },74        .projection => |projection| blk: {75            buffer[0] = values[projection.source.index];76            break :blk buffer[0..1];77        },78    };79}

1const context_mod = @import("context.zig");2const execute = @import("execute.zig");3const graph_mod = @import("graph.zig");4const leaf = @import("leaf.zig");5const spec = @import("spec.zig");6const step_mod = @import("step.zig");78pub const Step = step_mod.Step;9pub const arguments = step_mod.arguments;10pub const Graph = graph_mod.Graph;11pub const Semantics = leaf.Semantics;12pub const BindContext = context_mod.BindContext;13pub const FinishContext = context_mod.FinishContext;14pub const LeafBindContext = context_mod.LeafBindContext;15pub const With = spec.With;16pub const Stack = spec.Stack;17pub const Layer = spec.Layer;18pub const semantics = leaf.semantics;19pub const with = spec.with;20pub const bind = spec.bind;21pub const stack = spec.stack;22pub const layer = spec.layer;23pub const run = execute.run;24pub const result = execute.result;

1pub const interpret = @import("interpret/root.zig");