tiny.accy.tensor.dsl.body

1const std = @import("std");2const tensor = @import("../root.zig");3const trace = tensor.trace;4const parameter = @import("parameter.zig");56pub fn buildInputs(builder: *trace.Builder, comptime parameters: anytype) ![]trace.Value {7    const input_count = comptime parameter.arity(parameters);8    const args = try builder.arena.allocator().alloc(trace.Value, input_count);9    if (comptime parameter.named(parameters)) {10        inline for (@typeInfo(@TypeOf(parameters)).@"struct".field_names, 0..) |field_name, index_value| {11            args[index_value] = try builder.inputSpec(@field(parameters, field_name));12        }13    } else {14        inline for (0..input_count) |index_value| {15            args[index_value] = try builder.inputSpec(parameters[index_value]);16        }17    }18    return args;19}2021pub fn call(comptime definition: anytype, builder: *trace.Builder, args: []const trace.Value) ![]const trace.Value {22    const result = if (comptime parameter.named(definition.parameters))23        try definition.body(builder, parameter.Args(definition.parameters){ .values = args })24    else25        try definition.body(builder, args);26    return collectOutputs(builder, result);27}2829fn collectOutputs(builder: *trace.Builder, result: anytype) ![]const trace.Value {30    const Result = @TypeOf(result);31    if (comptime Result == trace.Value) {32        const outputs = try builder.arena.allocator().alloc(trace.Value, 1);33        outputs[0] = result;34        return outputs;35    }3637    return switch (@typeInfo(Result)) {38        .array => |array| blk: {39            if (array.child != trace.Value) @compileError("tensor Program array outputs must contain tensor Values");40            const outputs = try builder.arena.allocator().alloc(trace.Value, array.len);41            for (result, 0..) |output, index_value| {42                outputs[index_value] = output;43            }44            break :blk outputs;45        },46        .pointer => |pointer| collectPointerOutputs(builder, result, pointer),47        .@"struct" => |info| blk: {48            if (info.field_names.len == 0) @compileError("tensor Program bodies must return at least one output");49            const outputs = try builder.arena.allocator().alloc(trace.Value, info.field_names.len);50            inline for (info.field_names, info.field_types, 0..) |field_name, field_type, index_value| {51                if (field_type != trace.Value) @compileError("tensor Program struct outputs must contain tensor Values");52                outputs[index_value] = @field(result, field_name);53            }54            break :blk outputs;55        },56        else => @compileError("tensor Program bodies must return a Value, array, slice, or struct of Values"),57    };58}5960fn collectPointerOutputs(builder: *trace.Builder, result: anytype, comptime pointer: std.builtin.Type.Pointer) ![]const trace.Value {61    switch (pointer.size) {62        .slice => {63            if (pointer.child != trace.Value) @compileError("tensor Program slice outputs must contain tensor Values");64            const outputs = try builder.arena.allocator().alloc(trace.Value, result.len);65            for (result, 0..) |output, index_value| {66                outputs[index_value] = output;67            }68            return outputs;69        },70        .one => {71            return switch (@typeInfo(pointer.child)) {72                .array => |array| blk: {73                    if (array.child != trace.Value) @compileError("tensor Program array outputs must contain tensor Values");74                    const outputs = try builder.arena.allocator().alloc(trace.Value, array.len);75                    for (result.*, 0..) |output, index_value| {76                        outputs[index_value] = output;77                    }78                    break :blk outputs;79                },80                else => @compileError("tensor Program pointer outputs must point at arrays or slices of tensor Values"),81            };82        },83        else => @compileError("tensor Program pointer outputs must point at arrays or slices of tensor Values"),84    }85}

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