tiny.accy.choir.einsum

1/// Four figures for one call of `lowerPlan`: bytes of local scratch, structural2/// visits, operations requested from the IR sink, and tensor types requested3/// from it, so a compiler pass can charge one lowering against its work budget4/// before it runs it. The figures assume an equation from `parse` and a plan5/// from `createPlan`, so at most 62 labels, at most 63 inputs, and exactly one6/// step fewer than inputs. Growth of the caller's arena, the internals of the7/// context and of the sink, rewriter queues, parsing and planning are left out,8/// and their owners add their own charges, as the einsum pass does.9pub const LoweringBounds = struct {10    scratch_bytes: u64,11    structural_visits: u64,12    operation_requests: u64,13    tensor_type_requests: u64,1415    /// Bytes charged for one list of `T`: four times its grown capacity for 6216    /// items, alignment slack for each allocation, and one more slice for a17    /// final owned copy, so the bound functions can charge one list of at most18    /// 62 items.19    fn listBytes(comptime T: type) u64 {20        return 4 * std.ArrayList(T).growCapacity(62) * @sizeOf(T) +21            4 * @alignOf(T) + sliceBytes(T);22    }2324    fn sliceBytes(comptime T: type) u64 {25        return 62 * @sizeOf(T) + @alignOf(T) - 1;26    }2728    /// Scratch charged for one contraction step: the reduction of each operand,29    /// the batch and contracted axis lists, the broadcast-and-reduce path, and30    /// the final transpose, so `loweringBounds` can multiply it by the step31    /// count to charge the scratch of a whole plan. The charge is 11 byte32    /// lists, 6 axis lists, 14 slices of dimensions or permutations, and 1 kept33    /// order.34    fn stepScratch() u64 {35        return 11 * listBytes(u8) + 6 * listBytes(i64) +36            14 * sliceBytes(i64) + sliceBytes(u8);37    }38};

1pub const Options = struct {2    strategy: Strategy = .auto,3    max_exact_inputs: u8 = 18,4    exact_state_limit: usize = default_exact_state_limit,5    beam_width: usize = 64,6    auto_beam_width: usize = default_auto_beam_width,78    /// Returns the structural visits charged for one `createPlan` run over9    /// `input_count` inputs under these options, whatever pruning or cache hits10    /// happen during the run. The einsum pass calls this with the operand count11    /// to charge planning against its work budget. The charge covers the12    /// dimensions the cost model reads, scans of the kept sets, completion13    /// searches, ordering and the hash maps, and it leaves out parsing,14    /// lowering, allocator internals and tensor execution. One cost evaluation15    /// is charged as eight passes over 256 labels, plus scans of the inputs and16    /// fixed metadata. Partition work is charged by the number of ways to17    /// assign each input outside, left or right, and the greedy strategy adds18    /// its lookahead. Map work is charged for four maps with lookups, inserts,19    /// growth and collisions, with eight times as many buckets as entries. The20    /// automatic strategy is charged for the path it will take: exact search21    /// when the state count fits, widening plus beam above the exact input22    /// limit, and beam otherwise. The errors are `error.EmptyEquation`,23    /// `error.TooManyInputs` above 63, `error.ExactInputLimitExceeded` for24    /// exact search past its limit, and `error.WorkOverflow`.25    pub fn workBound(self: Options, input_count: usize) !u64 {26        if (input_count == 0) return error.EmptyEquation;27        if (input_count > 63) return error.TooManyInputs;28        var work: PlannerWork = .{ .inputs = input_count };29        if (precomputeStateCount(input_count)) |states| try work.add(16 * states);30        switch (self.strategy) {31            .left_to_right => try work.linear(false),32            .greedy => try work.linear(true),33            .optimal => try work.optimal(self.max_exact_inputs),34            .beam => try work.beam(@max(self.beam_width, 1), false),35            .anytime => try work.widening(@max(self.beam_width, 1)),36            .auto => {37                if (canUseExactPlan(input_count, self)) {38                    try work.optimal(self.max_exact_inputs);39                } else if (input_count > self.max_exact_inputs) {40                    const width = @max(self.auto_beam_width, 1);41                    try work.widening(width);42                    try work.beam(width, true);43                } else {44                    try work.beam(@max(self.beam_width, 1), false);45                }46            },47        }48        return work.visits;49    }5051    /// Returns the bytes of allocator storage one `createPlan` run may request,52    /// counting list growth that is never reclaimed and the arenas of beam53    /// search. The einsum pass calls this with the operand count to charge54    /// planning's memory against its scratch budget. This bound leaves out55    /// parsing, lowering, the overhead of the caller's arena and stack storage.56    /// The figure follows the same strategy choice `createPlan` makes, and for57    /// a widening search it adds up the storage of every width tried. A test58    /// runs every strategy in a fixed buffer of this size that never frees.59    pub fn storageBound(self: Options, input_count: usize) !u64 {60        if (input_count == 0) return error.EmptyEquation;61        if (input_count > 63) return error.TooManyInputs;62        var storage: PlannerStorage = .{};63        if (precomputeStateCount(input_count)) |states| {64            try storage.slice(IndexSet, 2 * states);65        }66        switch (self.strategy) {67            .left_to_right, .greedy => try storage.linear(input_count),68            .optimal => try storage.optimal(input_count, self.max_exact_inputs),69            .beam => try storage.beam(input_count, @max(self.beam_width, 1), false),70            .anytime => try storage.widening(input_count, @max(self.beam_width, 1)),71            .auto => {72                if (canUseExactPlan(input_count, self)) {73                    try storage.optimal(input_count, self.max_exact_inputs);74                } else if (input_count > self.max_exact_inputs) {75                    const width = @max(self.auto_beam_width, 1);76                    try storage.widening(input_count, width);77                    try storage.beam(input_count, width, true);78                } else {79                    try storage.beam(input_count, @max(self.beam_width, 1), false);80                }81            },82        }83        return storage.bytes;84    }85};

1pub const Plan = struct {2    allocator: std.mem.Allocator,3    strategy: Strategy,4    steps: []Step,5    total_scalar_cost: u128,6    total_flop_cost: u128,7    peak_intermediate_elements: u128,8    peak_workspace_elements: u128,9    output_elements: u128,10    selected_beam_width: usize,11    search: SearchStats,1213    pub fn deinit(self: *Plan) void {14        self.allocator.free(self.steps);15        self.* = undefined;16    }17};

1pub const Step = struct {2    lhs: NodeId,3    rhs: NodeId,4    result: NodeId,5    lhs_inputs: u64,6    rhs_inputs: u64,7    result_inputs: u64,8    lhs_indices: IndexSet,9    rhs_indices: IndexSet,10    result_indices: IndexSet,11    summed_indices: IndexSet,12    scalar_cost: u128,13    flop_cost: u128,14    result_elements: u128,15    work_elements: u128,16};

1pub const Strategy = enum {2    auto,3    left_to_right,4    greedy,5    beam,6    anytime,7    optimal,8};

1pub const Equation = struct {2    allocator: std.mem.Allocator,3    inputs: []Input,4    output: []const u8,5    output_set: IndexSet,6    index_set: IndexSet,7    dimensions: [256]u64,89    pub fn deinit(self: *Equation) void {10        for (self.inputs) |input| {11            self.allocator.free(input.indices);12            self.allocator.free(input.dims);13        }14        self.allocator.free(self.inputs);15        self.allocator.free(self.output);16        self.* = undefined;17    }1819    pub fn dimension(self: *const Equation, index: u8) u64 {20        std.debug.assert(self.index_set.contains(index));21        return self.dimensions[index];22    }2324    pub fn elementCount(self: *const Equation, indices: IndexSet) u128 {25        var product: u128 = 1;26        for (indices.words, 0..) |initial_word, word_index| {27            var word = initial_word;28            while (word != 0) {29                const bit: usize = @intCast(@ctz(word));30                const index: u8 = @intCast(word_index * 64 + bit);31                product = saturatingMul(product, self.dimension(index));32                word &= word - 1;33            }34        }35        return product;36    }3738    pub fn inputSet(self: *const Equation, input_mask: u64) IndexSet {39        var set: IndexSet = .{};40        for (self.inputs, 0..) |input, index| {41            if ((input_mask & (@as(u64, 1) << @intCast(index))) == 0) continue;42            set = set.unioned(input.index_set);43        }44        return set;45    }46};

1pub const IndexSet = struct {2    words: [4]u64 = .{ 0, 0, 0, 0 },34    pub fn add(self: *IndexSet, index: u8) void {5        self.words[wordIndex(index)] |= bitMask(index);6    }78    pub fn contains(self: IndexSet, index: u8) bool {9        return (self.words[wordIndex(index)] & bitMask(index)) != 0;10    }1112    pub fn unioned(self: IndexSet, other: IndexSet) IndexSet {13        return .{ .words = .{14            self.words[0] | other.words[0],15            self.words[1] | other.words[1],16            self.words[2] | other.words[2],17            self.words[3] | other.words[3],18        } };19    }2021    pub fn intersected(self: IndexSet, other: IndexSet) IndexSet {22        return .{ .words = .{23            self.words[0] & other.words[0],24            self.words[1] & other.words[1],25            self.words[2] & other.words[2],26            self.words[3] & other.words[3],27        } };28    }2930    pub fn without(self: IndexSet, other: IndexSet) IndexSet {31        return .{ .words = .{32            self.words[0] & ~other.words[0],33            self.words[1] & ~other.words[1],34            self.words[2] & ~other.words[2],35            self.words[3] & ~other.words[3],36        } };37    }3839    pub fn eql(self: IndexSet, other: IndexSet) bool {40        return std.mem.eql(u64, &self.words, &other.words);41    }4243    pub fn isEmpty(self: IndexSet) bool {44        return self.words[0] == 0 and self.words[1] == 0 and self.words[2] == 0 and self.words[3] == 0;45    }4647    pub fn count(self: IndexSet) u32 {48        return @popCount(self.words[0]) + @popCount(self.words[1]) + @popCount(self.words[2]) + @popCount(self.words[3]);49    }5051    fn wordIndex(index: u8) usize {52        return @intCast(index >> 6);53    }5455    fn bitMask(index: u8) u64 {56        const shift: u6 = @intCast(index & 63);57        return @as(u64, 1) << shift;58    }59};

1pub const Input = struct {2    indices: []const u8,3    index_set: IndexSet,4    dims: []const u64,5};

1/// Two figures for one parse: the bytes the parser asks its allocator for,2/// counting list growth it cannot give back, and the structural visits. The3/// einsum pass reads these two figures to charge a parse before it runs it. The4/// figures leave out stack storage, the internals of the allocator or arena,5/// and the borrowed equation text and shapes. The parser reads an input's6/// dimensions only after that input's rank matches its labels, which are unique7/// letters and digits.8pub const ParseBounds = struct {9    allocation_capacity: u64 = 0,10    structural_visits: u64 = 0,1112    const Error = error{WorkOverflow};1314    fn add(a: u64, b: u64) Error!u64 {15        return std.math.add(u64, a, b) catch return error.WorkOverflow;16    }1718    fn multiply(a: u64, b: u64) Error!u64 {19        return std.math.mul(u64, a, b) catch return error.WorkOverflow;20    }2122    fn slice(comptime T: type, count: usize) Error!u64 {23        if (count == 0) return 0;24        return add(try multiply(@sizeOf(T), count), @alignOf(T) - 1);25    }2627    /// Charges the bytes for one growing list of `count` items of `T` for the28    /// parse bound: every capacity the list grows through, then one more slice29    /// for the final `toOwnedSlice` copy that happens when shrinking cannot30    /// resize or remap in place.31    fn list(comptime T: type, count: usize) Error!u64 {32        std.debug.assert(count <= 63);33        var capacity: usize = 0;34        var bytes: u64 = 0;35        while (capacity < count) {36            capacity = std.ArrayList(T).growCapacity(capacity + 1);37            bytes = try add(bytes, try slice(T, capacity));38        }39        return add(bytes, try slice(T, count));40    }41};

1pub const LowerError = error{2    InputCountMismatch,3    InvalidPlan,4    InvalidDimension,5} || std.mem.Allocator.Error;

1pub fn lowerPlan(2    allocator: std.mem.Allocator,3    fb: *choir_root.semantic.FunctionBuilder,4    equation: *const Equation,5    plan: *const Plan,6    inputs: []const *ir.Value,7    dtype: DType,8) !*ir.Value {9    var sink = FunctionSink{ .fb = fb };10    return try lowerPlanWithSink(allocator, &sink, equation, plan, inputs, dtype);11}

1pub fn lowerPlanWithRewriter(2    allocator: std.mem.Allocator,3    rewriter: *rewrite.PatternRewriter,4    equation: *const Equation,5    plan: *const Plan,6    inputs: []const *ir.Value,7    dtype: DType,8) !*ir.Value {9    var sink = RewriterSink{ .rewriter = rewriter };10    return try lowerPlanWithSink(allocator, &sink, equation, plan, inputs, dtype);11}

1/// Returns the figures for an equation with `input_count` inputs, or2/// `error.EmptyEquation` for none and `error.TooManyInputs` above 63, and that3/// check keeps every product in range. The einsum pass calls this with the4/// operand count to add the lowering's share to its work estimate. Each step is5/// charged for both operand reductions and for the general broadcast-and-reduce6/// path, even when the matrix-product path skips them. One reduction may7/// transpose, make a zero, reshape, reduce, and reshape back. The operation8/// figure is 19 per step and the tensor-type figure 17 per step, and a single9/// input is charged 3 of each for its own reduction and the final transpose.10/// Visits count searches over labels and axes, which grow with the square of11/// the rank, together with local copies and cleanup. Visits measure neither12/// machine instructions nor tensor execution.13pub fn loweringBounds(input_count: usize) !LoweringBounds {14    if (input_count == 0) return error.EmptyEquation;15    if (input_count > 63) return error.TooManyInputs;16    const steps: u64 = input_count - 1;17    const states: u64 = input_count + steps;18    const scratch = states * @sizeOf(State) + @alignOf(State) - 1 + if (steps == 0)19        LoweringBounds.listBytes(u8) + LoweringBounds.listBytes(i64) +20            3 * LoweringBounds.sliceBytes(i64)21    else22        steps * LoweringBounds.stepScratch();23    return .{24        .scratch_bytes = scratch,25        .structural_visits = 8 * scratch + 64 * 63 * 63 * @max(steps, 1) + 8 * states,26        .operation_requests = if (steps == 0) 3 else 19 * steps,27        .tensor_type_requests = if (steps == 0) 3 else 17 * steps,28    };29}

1pub const NodeId = u32;

1pub const PlanError = error{2    EmptyEquation,3    TooManyInputs,4    ExactInputLimitExceeded,5    InvalidState,6} || std.mem.Allocator.Error;

1pub fn createPlan(allocator: std.mem.Allocator, equation: *const Equation, options: Options) PlanError!Plan {2    if (equation.inputs.len == 0) return error.EmptyEquation;3    if (equation.inputs.len > 63) return error.TooManyInputs;4    var context = try PlanningContext.init(allocator, equation);5    defer context.deinit(allocator);6    return switch (options.strategy) {7        .auto => createAutoPlan(allocator, &context, options),8        .left_to_right => createLinearPlan(allocator, &context, .left_to_right),9        .greedy => createLinearPlan(allocator, &context, .greedy),10        .beam => createBeamPlan(allocator, &context, options.beam_width),11        .anytime => createAnytimePlan(allocator, &context, options.beam_width),12        .optimal => createOptimalPlan(allocator, &context, options.max_exact_inputs),13    };14}

1pub const default_auto_beam_width: usize = 256;

1pub const default_exact_state_limit: usize = 262_144;

1pub const ParseError = error{2    MissingArrow,3    MultipleArrows,4    EmptyInputList,5    InvalidIndex,6    RepeatedInputIndex,7    RepeatedOutputIndex,8    OutputIndexMissing,9    ArityMismatch,10    RankMismatch,11    DimensionMismatch,12    TooManyInputs,13} || std.mem.Allocator.Error;

1pub fn parse(allocator: std.mem.Allocator, equation: []const u8, input_shapes: []const []const u64) ParseError!Equation {2    const arrow = findArrow(equation) orelse return error.MissingArrow;3    if (findArrow(equation[arrow + 2 ..]) != null) return error.MultipleArrows;4    if (input_shapes.len > 63) return error.TooManyInputs;56    const lhs = equation[0..arrow];7    const rhs = equation[arrow + 2 ..];8    var dimensions: [256]u64 = undefined;9    var global_indices: IndexSet = .{};10    var inputs = std.ArrayListUnmanaged(Input).empty;11    errdefer inputs.deinit(allocator);12    errdefer freeInputs(allocator, inputs.items);1314    var input_index: usize = 0;15    var term_iter = std.mem.splitScalar(u8, lhs, ',');16    while (term_iter.next()) |term| {17        if (input_index >= input_shapes.len) return error.ArityMismatch;18        const input = try parseInput(allocator, term, input_shapes[input_index], &dimensions, &global_indices);19        var input_owned = true;20        errdefer if (input_owned) freeInput(allocator, input);21        try inputs.append(allocator, input);22        input_owned = false;23        input_index += 1;24    }25    if (input_index == 0) return error.EmptyInputList;26    if (input_index != input_shapes.len) return error.ArityMismatch;2728    const output = try parseOutput(allocator, rhs, global_indices);29    errdefer allocator.free(output.indices);3031    return .{32        .allocator = allocator,33        .inputs = try inputs.toOwnedSlice(allocator),34        .output = output.indices,35        .output_set = output.index_set,36        .index_set = global_indices,37        .dimensions = dimensions,38    };39}

1/// Returns the figures for an equation of `equation_bytes` bytes and2/// `input_count` inputs, computed from those two numbers alone with no parsing3/// and no allocation, so the einsum pass can call this with the equation length4/// and operand count before parsing. The figures cover a successful parse and5/// every error the parse can stop at. The parser looks for the arrow and for a6/// second arrow before it checks the input limit, so above 63 inputs the7/// figures charge only the text scan and no storage. Label lists stop at 628/// distinct letters and digits, so extra whitespace lengthens the scan but adds9/// no storage. Structural visits cover copies, checks of dimensions and label10/// sets, transfers, and cleanup after a failed prefix. A test parses inside a11/// fixed buffer of the charged size that never frees.12pub fn parseBounds(equation_bytes: usize, input_count: usize) ParseBounds.Error!ParseBounds {13    var bounds: ParseBounds = .{14        .structural_visits = try ParseBounds.add(1, try ParseBounds.multiply(16, equation_bytes)),15    };16    if (input_count > 63) return bounds;17    const labels = @min(equation_bytes, 62);18    const label_storage = try ParseBounds.list(u8, labels);19    const input_storage = try ParseBounds.add(label_storage, try ParseBounds.slice(u64, labels));20    bounds.allocation_capacity = try ParseBounds.add(21        try ParseBounds.list(Input, input_count),22        try ParseBounds.add(try ParseBounds.multiply(input_count, input_storage), label_storage),23    );24    bounds.structural_visits = try ParseBounds.add(bounds.structural_visits, try ParseBounds.add(25        try ParseBounds.multiply(4, bounds.allocation_capacity),26        16 * (input_count + 1),27    ));28    return bounds;29}

1const spec = @import("spec.zig");2const planner = @import("planner.zig");3const lowering = @import("lowering.zig");45pub const Equation = spec.Equation;6pub const Input = spec.Input;7pub const IndexSet = spec.IndexSet;8pub const ParseError = spec.ParseError;9pub const ParseBounds = spec.ParseBounds;1011pub const NodeId = planner.NodeId;12pub const Options = planner.Options;13pub const Plan = planner.Plan;14pub const PlanError = planner.PlanError;15pub const Step = planner.Step;16pub const Strategy = planner.Strategy;17pub const default_exact_state_limit = planner.default_exact_state_limit;18pub const default_auto_beam_width = planner.default_auto_beam_width;19pub const LowerError = lowering.LowerError;20pub const LoweringBounds = lowering.LoweringBounds;2122pub const parse = spec.parse;23pub const parseBounds = spec.parseBounds;24pub const createPlan = planner.createPlan;25pub const loweringBounds = lowering.loweringBounds;26pub const lowerPlan = lowering.lowerPlan;27pub const lowerPlanWithRewriter = lowering.lowerPlanWithRewriter;

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