lib/accy/src/choir/einsum/planner.zig

daab053ee43316e1809a84551d573ddd1e5bf3d2

   1 const std = @import("std");
   2 const alloc_arena = @import("alloc_arena");
   3 const builtin = @import("builtin");
   4 const testing = std.testing;
   5 
   6 const spec = @import("spec.zig");
   7 
   8 const Equation = spec.Equation;
   9 const IndexSet = spec.IndexSet;
  10 
  11 pub const NodeId = u32;
  12 
  13 pub const Strategy = enum {
  14     auto,
  15     left_to_right,
  16     greedy,
  17     beam,
  18     anytime,
  19     optimal,
  20 };
  21 
  22 pub const default_exact_state_limit: usize = 262_144;
  23 pub const default_auto_beam_width: usize = 256;
  24 
  25 pub const Options = struct {
  26     strategy: Strategy = .auto,
  27     max_exact_inputs: u8 = 18,
  28     exact_state_limit: usize = default_exact_state_limit,
  29     beam_width: usize = 64,
  30     auto_beam_width: usize = default_auto_beam_width,
  31 
  32     /// Returns the structural visits charged for one `createPlan` run over
  33     /// `input_count` inputs under these options, whatever pruning or cache hits
  34     /// happen during the run. The einsum pass calls this with the operand count
  35     /// to charge planning against its work budget. The charge covers the
  36     /// dimensions the cost model reads, scans of the kept sets, completion
  37     /// searches, ordering and the hash maps, and it leaves out parsing,
  38     /// lowering, allocator internals and tensor execution. One cost evaluation
  39     /// is charged as eight passes over 256 labels, plus scans of the inputs and
  40     /// fixed metadata. Partition work is charged by the number of ways to
  41     /// assign each input outside, left or right, and the greedy strategy adds
  42     /// its lookahead. Map work is charged for four maps with lookups, inserts,
  43     /// growth and collisions, with eight times as many buckets as entries. The
  44     /// automatic strategy is charged for the path it will take: exact search
  45     /// when the state count fits, widening plus beam above the exact input
  46     /// limit, and beam otherwise. The errors are `error.EmptyEquation`,
  47     /// `error.TooManyInputs` above 63, `error.ExactInputLimitExceeded` for
  48     /// exact search past its limit, and `error.WorkOverflow`.
  49     pub fn workBound(self: Options, input_count: usize) !u64 {
  50         if (input_count == 0) return error.EmptyEquation;
  51         if (input_count > 63) return error.TooManyInputs;
  52         var work: PlannerWork = .{ .inputs = input_count };
  53         if (precomputeStateCount(input_count)) |states| try work.add(16 * states);
  54         switch (self.strategy) {
  55             .left_to_right => try work.linear(false),
  56             .greedy => try work.linear(true),
  57             .optimal => try work.optimal(self.max_exact_inputs),
  58             .beam => try work.beam(@max(self.beam_width, 1), false),
  59             .anytime => try work.widening(@max(self.beam_width, 1)),
  60             .auto => {
  61                 if (canUseExactPlan(input_count, self)) {
  62                     try work.optimal(self.max_exact_inputs);
  63                 } else if (input_count > self.max_exact_inputs) {
  64                     const width = @max(self.auto_beam_width, 1);
  65                     try work.widening(width);
  66                     try work.beam(width, true);
  67                 } else {
  68                     try work.beam(@max(self.beam_width, 1), false);
  69                 }
  70             },
  71         }
  72         return work.visits;
  73     }
  74 
  75     /// Returns the bytes of allocator storage one `createPlan` run may request,
  76     /// counting list growth that is never reclaimed and the arenas of beam
  77     /// search. The einsum pass calls this with the operand count to charge
  78     /// planning's memory against its scratch budget. This bound leaves out
  79     /// parsing, lowering, the overhead of the caller's arena and stack storage.
  80     /// The figure follows the same strategy choice `createPlan` makes, and for
  81     /// a widening search it adds up the storage of every width tried. A test
  82     /// runs every strategy in a fixed buffer of this size that never frees.
  83     pub fn storageBound(self: Options, input_count: usize) !u64 {
  84         if (input_count == 0) return error.EmptyEquation;
  85         if (input_count > 63) return error.TooManyInputs;
  86         var storage: PlannerStorage = .{};
  87         if (precomputeStateCount(input_count)) |states| {
  88             try storage.slice(IndexSet, 2 * states);
  89         }
  90         switch (self.strategy) {
  91             .left_to_right, .greedy => try storage.linear(input_count),
  92             .optimal => try storage.optimal(input_count, self.max_exact_inputs),
  93             .beam => try storage.beam(input_count, @max(self.beam_width, 1), false),
  94             .anytime => try storage.widening(input_count, @max(self.beam_width, 1)),
  95             .auto => {
  96                 if (canUseExactPlan(input_count, self)) {
  97                     try storage.optimal(input_count, self.max_exact_inputs);
  98                 } else if (input_count > self.max_exact_inputs) {
  99                     const width = @max(self.auto_beam_width, 1);
 100                     try storage.widening(input_count, width);
 101                     try storage.beam(input_count, width, true);
 102                 } else {
 103                     try storage.beam(input_count, @max(self.beam_width, 1), false);
 104                 }
 105             },
 106         }
 107         return storage.bytes;
 108     }
 109 };
 110 
 111 const StorageError = error{WorkOverflow};
 112 
 113 const PlannerWork = struct {
 114     inputs: usize,
 115     visits: u64 = 0,
 116 
 117     fn sum(a: u64, b: u64) StorageError!u64 {
 118         return std.math.add(u64, a, b) catch return error.WorkOverflow;
 119     }
 120 
 121     fn product(a: u64, b: u64) StorageError!u64 {
 122         return std.math.mul(u64, a, b) catch return error.WorkOverflow;
 123     }
 124 
 125     fn add(self: *PlannerWork, visits: u64) StorageError!void {
 126         self.visits = try sum(self.visits, visits);
 127     }
 128 
 129     fn costs(self: *PlannerWork, count: u64) StorageError!void {
 130         try self.add(try product(count, 64 + 8 * (256 + self.inputs)));
 131     }
 132 
 133     fn greedyEvaluations(count: usize) u64 {
 134         std.debug.assert(count >= 2);
 135         std.debug.assert(count <= 63);
 136         const pairs = candidatePairCount(count);
 137         return if (count == 2) pairs else pairs * (2 + candidatePairCount(count - 1));
 138     }
 139 
 140     fn linear(self: *PlannerWork, lookahead: bool) StorageError!void {
 141         try self.costs(4 * self.inputs);
 142         var count = self.inputs;
 143         while (count > 1) : (count -= 1) {
 144             try self.costs(if (lookahead) greedyEvaluations(count) else 1);
 145             try self.add(16 * count);
 146         }
 147     }
 148 
 149     fn partitions(count: usize) StorageError!u64 {
 150         std.debug.assert(count <= 20);
 151         var result: u64 = 1;
 152         for (0..count) |_| result = try product(result, 3);
 153         return result;
 154     }
 155 
 156     fn optimal(self: *PlannerWork, max_inputs: u8) !void {
 157         if (self.inputs > max_inputs or self.inputs > 20) {
 158             return error.ExactInputLimitExceeded;
 159         }
 160         try self.costs(try sum(try partitions(self.inputs), 4 * self.inputs));
 161     }
 162 
 163     fn completion(count: usize, exact_limit: usize) StorageError!u64 {
 164         if (count <= 1) return 0;
 165         if (count <= exact_limit) {
 166             return sum(try partitions(count), exact_completion_state_count);
 167         }
 168         var result: u64 = 4 * count;
 169         var remaining = count;
 170         while (remaining > 1) : (remaining -= 1) {
 171             result = try sum(result, candidatePairCount(remaining));
 172         }
 173         return result;
 174     }
 175 
 176     fn beam(self: *PlannerWork, width: usize, exploratory: bool) StorageError!void {
 177         std.debug.assert(width >= 1);
 178         if (self.inputs == 1) return self.costs(4);
 179         try self.linear(true);
 180         const refinement = @max(
 181             try product(width, default_beam_refinement_factor),
 182             beam_refinement_min,
 183         );
 184         const exact_limit: usize = if (exploratory)
 185             auto_exploration_completion_exact_limit
 186         else
 187             default_completion_exact_limit;
 188         const score_limit: usize = if (exploratory)
 189             auto_exploration_completion_score_limit
 190         else
 191             default_completion_score_limit;
 192         var frontier: u64 = 1;
 193         var total: u64 = 0;
 194         var count = self.inputs;
 195         while (count > 1) : (count -= 1) {
 196             const candidates = try product(frontier, candidatePairCount(count));
 197             total = try sum(total, candidates);
 198             try self.costs(try product(candidates, 2 + candidatePairCount(count - 1)));
 199             try self.add(try product(candidates, 16 * count));
 200             const ordering = try sum(candidates, 1);
 201             try self.add(try product(32, try product(ordering, ordering)));
 202             const refined = @min(candidates, refinement);
 203             if (count - 1 <= score_limit) {
 204                 try self.costs(try product(refined, try completion(count - 1, exact_limit)));
 205             }
 206             frontier = @min(width, candidates);
 207         }
 208         try self.maps(total);
 209     }
 210 
 211     fn maps(self: *PlannerWork, entries: u64) StorageError!void {
 212         const padded = try sum(entries, 8);
 213         const probes = try product(512, try product(padded, padded));
 214         try self.add(try product(probes, try product(self.inputs + 1, self.inputs + 1)));
 215     }
 216 
 217     fn widening(self: *PlannerWork, final_width: usize) StorageError!void {
 218         var width: usize = 1;
 219         for (0..@bitSizeOf(usize) + 1) |_| {
 220             try self.beam(width, false);
 221             if (width == final_width) return;
 222             width = nextBeamWidth(width, final_width);
 223         }
 224         unreachable;
 225     }
 226 };
 227 
 228 const PlannerStorage = struct {
 229     bytes: u64 = 0,
 230 
 231     fn add(self: *PlannerStorage, bytes: u64) StorageError!void {
 232         self.bytes = std.math.add(u64, self.bytes, bytes) catch return error.WorkOverflow;
 233     }
 234 
 235     fn product(a: u64, b: u64) StorageError!u64 {
 236         return std.math.mul(u64, a, b) catch return error.WorkOverflow;
 237     }
 238 
 239     fn slice(self: *PlannerStorage, comptime T: type, count: u64) StorageError!void {
 240         if (count == 0) return;
 241         try self.add(try product(count, @sizeOf(T)));
 242         try self.add(128);
 243     }
 244 
 245     fn list(self: *PlannerStorage, comptime T: type, count: u64) StorageError!void {
 246         if (count == 0) return;
 247         const minimum = std.math.cast(usize, count) orelse return error.WorkOverflow;
 248         const grown: u64 = std.ArrayList(T).growCapacity(minimum);
 249         try self.slice(T, try product(4, grown));
 250     }
 251 
 252     fn map(
 253         self: *PlannerStorage,
 254         comptime K: type,
 255         comptime V: type,
 256         count: u64,
 257     ) StorageError!void {
 258         if (count == 0) return;
 259         const padded = std.math.add(u64, count, 8) catch return error.WorkOverflow;
 260         const buckets = try product(8, padded);
 261         try self.add(try product(buckets, @sizeOf(K) + @sizeOf(V) + 65));
 262     }
 263 
 264     fn beam(
 265         self: *PlannerStorage,
 266         inputs: usize,
 267         width: usize,
 268         exploratory: bool,
 269     ) StorageError!void {
 270         if (inputs == 1) return;
 271         const widened = try product(width, default_beam_refinement_factor);
 272         const refinement = @max(widened, beam_refinement_min);
 273         var scratch: PlannerStorage = .{};
 274         try scratch.slice(Active, inputs);
 275         try scratch.list(BeamState, 1);
 276         var frontier: u64 = 1;
 277         var steps: u64 = 0;
 278         var immediate: u64 = 0;
 279         var exact: u64 = 0;
 280         var greedy: u64 = 0;
 281         var count = inputs;
 282         while (count > 1) : (count -= 1) {
 283             const candidates = try product(frontier, candidatePairCount(count));
 284             steps = std.math.add(u64, steps, candidates) catch return error.WorkOverflow;
 285             try scratch.add(try product(candidates, (count - 1) * @sizeOf(Active) + 128));
 286             try scratch.list(BeamState, candidates);
 287             const refined = @min(candidates, refinement);
 288             const factor = default_beam_refinement_factor;
 289             const deduplicate = switch (inputs > beam_immediate_cache_min) {
 290                 true => shouldDeduplicateBeamCandidates(true, count, factor),
 291                 false => shouldDeduplicateBeamCandidates(false, count, factor),
 292             };
 293             if (deduplicate) {
 294                 try scratch.map(GreedyCompletionKey, usize, refined);
 295             }
 296             if (inputs > beam_immediate_cache_min and count - 1 >= beam_immediate_cache_min) {
 297                 immediate = std.math.add(u64, immediate, candidates) catch
 298                     return error.WorkOverflow;
 299             }
 300             const exact_limit: usize = switch (exploratory) {
 301                 true => auto_exploration_completion_exact_limit,
 302                 false => default_completion_exact_limit,
 303             };
 304             const score_limit: usize = switch (exploratory) {
 305                 true => auto_exploration_completion_score_limit,
 306                 false => default_completion_score_limit,
 307             };
 308             if (count - 1 <= score_limit) {
 309                 const entries = if (count - 1 <= exact_limit) &exact else &greedy;
 310                 entries.* = std.math.add(u64, entries.*, refined) catch return error.WorkOverflow;
 311             }
 312             frontier = @min(width, candidates);
 313         }
 314         try scratch.list(BeamStep, steps);
 315         try scratch.map(GreedyCompletionKey, u128, immediate);
 316         try scratch.map(ExactCompletionKey, u128, exact);
 317         try scratch.map(GreedyCompletionKey, u128, greedy);
 318         try self.add(try product(8, scratch.bytes));
 319         try self.slice(Step, inputs - 1);
 320     }
 321 
 322     fn widening(self: *PlannerStorage, inputs: usize, final_width: usize) StorageError!void {
 323         var width: usize = 1;
 324         for (0..@bitSizeOf(usize) + 1) |_| {
 325             try self.beam(inputs, width, false);
 326             if (width == final_width) return;
 327             width = nextBeamWidth(width, final_width);
 328         }
 329         unreachable;
 330     }
 331 
 332     fn optimal(self: *PlannerStorage, inputs: usize, max_inputs: u8) !void {
 333         if (inputs > max_inputs or inputs > 20) return error.ExactInputLimitExceeded;
 334         if (inputs == 1) return;
 335         try self.slice(DpState, @as(u64, 1) << @intCast(inputs));
 336         try self.list(Step, inputs - 1);
 337         try self.slice(Step, inputs - 1);
 338     }
 339 
 340     fn linear(self: *PlannerStorage, inputs: usize) StorageError!void {
 341         try self.slice(Active, inputs);
 342         try self.list(Step, inputs - 1);
 343         try self.slice(Step, inputs - 1);
 344     }
 345 };
 346 
 347 const mask_precompute_limit: usize = 262_144;
 348 const mask_precompute_min_inputs: usize = 11;
 349 
 350 const PlanningContext = struct {
 351     equation: *const Equation,
 352     keep_sets: []IndexSet = &.{},
 353 
 354     fn init(allocator: std.mem.Allocator, equation: *const Equation) PlanError!PlanningContext {
 355         var context = PlanningContext{ .equation = equation };
 356         const state_count = precomputeStateCount(equation.inputs.len) orelse return context;
 357         const mask_sets = try allocator.alloc(IndexSet, state_count);
 358         defer allocator.free(mask_sets);
 359         context.keep_sets = try allocator.alloc(IndexSet, state_count);
 360         errdefer allocator.free(context.keep_sets);
 361 
 362         mask_sets[0] = .{};
 363         var mask: usize = 1;
 364         while (mask < state_count) : (mask += 1) {
 365             const previous = mask & (mask - 1);
 366             const input_index: usize = @intCast(@ctz(mask));
 367             mask_sets[mask] = mask_sets[previous].unioned(equation.inputs[input_index].index_set);
 368         }
 369 
 370         const all_mask = state_count - 1;
 371         mask = 0;
 372         while (mask < state_count) : (mask += 1) {
 373             context.keep_sets[mask] = equation.output_set.unioned(mask_sets[all_mask ^ mask]);
 374         }
 375 
 376         return context;
 377     }
 378 
 379     fn deinit(self: *PlanningContext, allocator: std.mem.Allocator) void {
 380         if (self.keep_sets.len != 0) allocator.free(self.keep_sets);
 381         self.* = undefined;
 382     }
 383 
 384     inline fn resultIndices(self: *const PlanningContext, combined_mask: u64, pair_indices: IndexSet) IndexSet {
 385         return pair_indices.intersected(self.keepSet(combined_mask));
 386     }
 387 
 388     inline fn keepSet(self: *const PlanningContext, mask: u64) IndexSet {
 389         if (self.keep_sets.len != 0) return self.keep_sets[@intCast(mask)];
 390         var keep = self.equation.output_set;
 391         for (self.equation.inputs, 0..) |input, index| {
 392             if ((mask & (@as(u64, 1) << @intCast(index))) != 0) continue;
 393             keep = keep.unioned(input.index_set);
 394         }
 395         return keep;
 396     }
 397 
 398     inline fn elementCount(self: *const PlanningContext, indices: IndexSet) u128 {
 399         return self.equation.elementCount(indices);
 400     }
 401 };
 402 
 403 pub const Step = struct {
 404     lhs: NodeId,
 405     rhs: NodeId,
 406     result: NodeId,
 407     lhs_inputs: u64,
 408     rhs_inputs: u64,
 409     result_inputs: u64,
 410     lhs_indices: IndexSet,
 411     rhs_indices: IndexSet,
 412     result_indices: IndexSet,
 413     summed_indices: IndexSet,
 414     scalar_cost: u128,
 415     flop_cost: u128,
 416     result_elements: u128,
 417     work_elements: u128,
 418 };
 419 
 420 pub const SearchStats = struct {
 421     candidate_contractions: u64 = 0,
 422     states_created: u64 = 0,
 423     states_kept: u64 = 0,
 424     states_pruned: u64 = 0,
 425     frontier_peak: u64 = 0,
 426 };
 427 
 428 pub const Plan = struct {
 429     allocator: std.mem.Allocator,
 430     strategy: Strategy,
 431     steps: []Step,
 432     total_scalar_cost: u128,
 433     total_flop_cost: u128,
 434     peak_intermediate_elements: u128,
 435     peak_workspace_elements: u128,
 436     output_elements: u128,
 437     selected_beam_width: usize,
 438     search: SearchStats,
 439 
 440     pub fn deinit(self: *Plan) void {
 441         self.allocator.free(self.steps);
 442         self.* = undefined;
 443     }
 444 };
 445 
 446 pub const PlanError = error{
 447     EmptyEquation,
 448     TooManyInputs,
 449     ExactInputLimitExceeded,
 450     InvalidState,
 451 } || std.mem.Allocator.Error;
 452 
 453 const Active = struct {
 454     node: NodeId,
 455     input_mask: u64,
 456     indices: IndexSet,
 457 };
 458 
 459 const PairChoice = struct {
 460     lhs_index: usize,
 461     rhs_index: usize,
 462     cost: u128,
 463     flop_cost: u128,
 464     score: u128,
 465     result_elements: u128,
 466     work_elements: u128,
 467 };
 468 
 469 const StepCost = struct {
 470     scalar: u128,
 471     flop: u128,
 472     work: u128,
 473 };
 474 
 475 const DpState = struct {
 476     valid: bool = false,
 477     scalar_cost: u128 = 0,
 478     flop_cost: u128 = 0,
 479     peak: u128 = 0,
 480     indices: IndexSet = .{},
 481     left: u64 = 0,
 482     right: u64 = 0,
 483 };
 484 
 485 const BeamState = struct {
 486     active: []Active,
 487     active_count: usize,
 488     step_ref: ?usize = null,
 489     step_count: usize = 0,
 490     scalar_cost: u128 = 0,
 491     flop_cost: u128 = 0,
 492     score: u128 = 0,
 493     peak: u128 = 0,
 494     preserve: bool = false,
 495 };
 496 
 497 const BeamStep = struct {
 498     previous: ?usize,
 499     step: Step,
 500 };
 501 
 502 const BeamCandidateEstimate = struct {
 503     combined_mask: u64,
 504     pair_indices: IndexSet,
 505     result_indices: IndexSet,
 506     cost: StepCost,
 507     prefix_flop_cost: u128,
 508 };
 509 
 510 const CompletionState = struct {
 511     valid: bool = false,
 512     flop_cost: u128 = 0,
 513     indices: IndexSet = .{},
 514     input_mask: u64 = 0,
 515 };
 516 
 517 const ExactCompletionKey = struct {
 518     count: u8 = 0,
 519     masks: [beam_exact_completion_max_inputs]u64 = @as([beam_exact_completion_max_inputs]u64, @splat(0)),
 520 };
 521 
 522 const GreedyCompletionKey = struct {
 523     count: u8 = 0,
 524     masks: [max_completion_key_inputs]u64 = @as([max_completion_key_inputs]u64, @splat(0)),
 525 };
 526 
 527 const GreedyCompletionKeyContext = struct {
 528     pub fn hash(_: GreedyCompletionKeyContext, key: GreedyCompletionKey) u64 {
 529         var hasher = std.hash.Wyhash.init(0);
 530         hasher.update(std.mem.asBytes(&key.count));
 531         const count: usize = key.count;
 532         for (key.masks[0..count]) |mask| {
 533             hasher.update(std.mem.asBytes(&mask));
 534         }
 535         return hasher.final();
 536     }
 537 
 538     pub fn eql(_: GreedyCompletionKeyContext, lhs: GreedyCompletionKey, rhs: GreedyCompletionKey) bool {
 539         if (lhs.count != rhs.count) return false;
 540         const count: usize = lhs.count;
 541         return std.mem.eql(u64, lhs.masks[0..count], rhs.masks[0..count]);
 542     }
 543 };
 544 
 545 const GreedyCompletionMap = std.HashMap(GreedyCompletionKey, u128, GreedyCompletionKeyContext, std.hash_map.default_max_load_percentage);
 546 const BeamPartitionMap = std.HashMap(GreedyCompletionKey, usize, GreedyCompletionKeyContext, std.hash_map.default_max_load_percentage);
 547 
 548 const CompletionCache = struct {
 549     exact: std.AutoHashMap(ExactCompletionKey, u128),
 550     greedy: GreedyCompletionMap,
 551 
 552     fn init(allocator: std.mem.Allocator) CompletionCache {
 553         return .{
 554             .exact = std.AutoHashMap(ExactCompletionKey, u128).init(allocator),
 555             .greedy = GreedyCompletionMap.init(allocator),
 556         };
 557     }
 558 
 559     fn deinit(self: *CompletionCache) void {
 560         self.exact.deinit();
 561         self.greedy.deinit();
 562     }
 563 };
 564 
 565 const ImmediateCompletionCache = GreedyCompletionMap;
 566 
 567 const beam_exact_completion_max_inputs = 9;
 568 const exact_completion_state_count = @as(usize, 1) << beam_exact_completion_max_inputs;
 569 const beam_immediate_cache_min = 10;
 570 const default_beam_refinement_factor = 32;
 571 const beam_refinement_min = 512;
 572 const beam_partition_dedup_max_inputs = 10;
 573 const beam_refined_partition_dedup_max_inputs = 16;
 574 const default_completion_exact_limit = 8;
 575 const default_completion_score_limit = 15;
 576 const auto_exploration_completion_exact_limit = 9;
 577 const auto_exploration_completion_score_limit = 12;
 578 const max_completion_key_inputs = 63;
 579 
 580 pub fn createPlan(allocator: std.mem.Allocator, equation: *const Equation, options: Options) PlanError!Plan {
 581     if (equation.inputs.len == 0) return error.EmptyEquation;
 582     if (equation.inputs.len > 63) return error.TooManyInputs;
 583     var context = try PlanningContext.init(allocator, equation);
 584     defer context.deinit(allocator);
 585     return switch (options.strategy) {
 586         .auto => createAutoPlan(allocator, &context, options),
 587         .left_to_right => createLinearPlan(allocator, &context, .left_to_right),
 588         .greedy => createLinearPlan(allocator, &context, .greedy),
 589         .beam => createBeamPlan(allocator, &context, options.beam_width),
 590         .anytime => createAnytimePlan(allocator, &context, options.beam_width),
 591         .optimal => createOptimalPlan(allocator, &context, options.max_exact_inputs),
 592     };
 593 }
 594 
 595 fn createAutoPlan(allocator: std.mem.Allocator, context: *const PlanningContext, options: Options) PlanError!Plan {
 596     const equation = context.equation;
 597     if (canUseExactPlan(equation.inputs.len, options)) {
 598         return createOptimalPlan(allocator, context, options.max_exact_inputs);
 599     }
 600     if (equation.inputs.len > options.max_exact_inputs) {
 601         return createAutoBeamPlan(allocator, context, options.auto_beam_width);
 602     }
 603     return createBeamPlan(allocator, context, options.beam_width);
 604 }
 605 
 606 fn canUseExactPlan(input_count: usize, options: Options) bool {
 607     if (input_count > options.max_exact_inputs) return false;
 608     if (input_count > 20) return false;
 609     const state_count = exactStateCount(input_count) orelse return false;
 610     return state_count <= options.exact_state_limit;
 611 }
 612 
 613 fn exactStateCount(input_count: usize) ?usize {
 614     if (input_count >= @bitSizeOf(usize)) return null;
 615     return @as(usize, 1) << @intCast(input_count);
 616 }
 617 
 618 fn precomputeStateCount(input_count: usize) ?usize {
 619     if (input_count < mask_precompute_min_inputs) return null;
 620     const state_count = exactStateCount(input_count) orelse return null;
 621     if (state_count > mask_precompute_limit) return null;
 622     return state_count;
 623 }
 624 
 625 fn createLinearPlan(allocator: std.mem.Allocator, context: *const PlanningContext, strategy: Strategy) PlanError!Plan {
 626     const equation = context.equation;
 627     var steps = std.ArrayListUnmanaged(Step).empty;
 628     errdefer steps.deinit(allocator);
 629     var search: SearchStats = .{
 630         .states_created = @intCast(equation.inputs.len),
 631         .states_kept = @intCast(equation.inputs.len),
 632         .frontier_peak = @intCast(equation.inputs.len),
 633     };
 634 
 635     var active = try allocator.alloc(Active, equation.inputs.len);
 636     defer allocator.free(active);
 637     for (equation.inputs, 0..) |input, index| {
 638         active[index] = .{
 639             .node = @intCast(index),
 640             .input_mask = @as(u64, 1) << @intCast(index),
 641             .indices = input.index_set,
 642         };
 643     }
 644 
 645     var active_count = equation.inputs.len;
 646     while (active_count > 1) {
 647         const choice = switch (strategy) {
 648             .left_to_right => blk: {
 649                 search.candidate_contractions += 1;
 650                 break :blk pairChoice(context, active[0], active[1], 0, 1);
 651             },
 652             .greedy => bestGreedyPair(context, active[0..active_count], &search),
 653             .auto, .beam, .anytime, .optimal => unreachable,
 654         };
 655         const lhs = active[choice.lhs_index];
 656         const rhs = active[choice.rhs_index];
 657         const combined_mask = lhs.input_mask | rhs.input_mask;
 658         const pair_indices = lhs.indices.unioned(rhs.indices);
 659         const result_indices = resultIndices(context, combined_mask, pair_indices);
 660         const result_node: NodeId = @intCast(equation.inputs.len + steps.items.len);
 661         const step: Step = .{
 662             .lhs = lhs.node,
 663             .rhs = rhs.node,
 664             .result = result_node,
 665             .lhs_inputs = lhs.input_mask,
 666             .rhs_inputs = rhs.input_mask,
 667             .result_inputs = combined_mask,
 668             .lhs_indices = lhs.indices,
 669             .rhs_indices = rhs.indices,
 670             .result_indices = result_indices,
 671             .summed_indices = pair_indices.without(result_indices),
 672             .scalar_cost = choice.cost,
 673             .flop_cost = choice.flop_cost,
 674             .result_elements = choice.result_elements,
 675             .work_elements = choice.work_elements,
 676         };
 677         try steps.append(allocator, step);
 678         active[choice.lhs_index] = .{
 679             .node = result_node,
 680             .input_mask = combined_mask,
 681             .indices = result_indices,
 682         };
 683         removeActive(active[0..active_count], choice.rhs_index);
 684         active_count -= 1;
 685         search.states_created += 1;
 686         search.states_kept += 1;
 687     }
 688 
 689     const owned_steps = try steps.toOwnedSlice(allocator);
 690     return finishPlan(allocator, equation, strategy, owned_steps, search);
 691 }
 692 
 693 fn createBeamPlan(allocator: std.mem.Allocator, context: *const PlanningContext, beam_width: usize) PlanError!Plan {
 694     return createBeamPlanWithRefinement(allocator, context, beam_width, default_beam_refinement_factor);
 695 }
 696 
 697 fn createBeamPlanWithRefinement(
 698     allocator: std.mem.Allocator,
 699     context: *const PlanningContext,
 700     beam_width: usize,
 701     refinement_factor: usize,
 702 ) PlanError!Plan {
 703     return createBeamPlanWithRefinementAndScoreLimit(allocator, context, beam_width, refinement_factor, default_completion_exact_limit, default_completion_score_limit);
 704 }
 705 
 706 fn createBeamPlanWithRefinementAndScoreLimit(
 707     allocator: std.mem.Allocator,
 708     context: *const PlanningContext,
 709     beam_width: usize,
 710     refinement_factor: usize,
 711     completion_exact_limit: usize,
 712     completion_score_limit: usize,
 713 ) PlanError!Plan {
 714     if (context.equation.inputs.len > beam_immediate_cache_min) {
 715         return createBeamPlanMode(true, allocator, context, beam_width, refinement_factor, completion_exact_limit, completion_score_limit);
 716     }
 717     return createBeamPlanMode(false, allocator, context, beam_width, refinement_factor, completion_exact_limit, completion_score_limit);
 718 }
 719 
 720 fn createBeamPlanMode(
 721     comptime use_immediate_cache: bool,
 722     allocator: std.mem.Allocator,
 723     context: *const PlanningContext,
 724     beam_width: usize,
 725     refinement_factor: usize,
 726     completion_exact_limit: usize,
 727     completion_score_limit: usize,
 728 ) PlanError!Plan {
 729     const equation = context.equation;
 730     const width = @max(beam_width, 1);
 731     if (equation.inputs.len == 1) {
 732         const steps = try allocator.alloc(Step, 0);
 733         var plan = finishPlan(allocator, equation, .beam, steps, .{
 734             .states_created = 1,
 735             .states_kept = 1,
 736             .frontier_peak = 1,
 737         });
 738         plan.selected_beam_width = width;
 739         return plan;
 740     }
 741 
 742     var scratch_state = alloc_arena.Arena.init(allocator);
 743     defer scratch_state.deinit();
 744     const scratch = scratch_state.allocator();
 745 
 746     const initial_active = try scratch.alloc(Active, equation.inputs.len);
 747     for (equation.inputs, 0..) |input, index| {
 748         initial_active[index] = .{
 749             .node = @intCast(index),
 750             .input_mask = @as(u64, 1) << @intCast(index),
 751             .indices = input.index_set,
 752         };
 753     }
 754 
 755     var beam = std.ArrayListUnmanaged(BeamState).empty;
 756     defer beam.deinit(scratch);
 757     try beam.append(scratch, .{
 758         .active = initial_active,
 759         .active_count = equation.inputs.len,
 760         .preserve = true,
 761     });
 762     var beam_steps = std.ArrayListUnmanaged(BeamStep).empty;
 763     defer beam_steps.deinit(scratch);
 764     var completion_cache = CompletionCache.init(scratch);
 765     defer completion_cache.deinit();
 766     var partition_cache = BeamPartitionMap.init(scratch);
 767     defer partition_cache.deinit();
 768     var immediate_cache: ImmediateCompletionCache = undefined;
 769     if (use_immediate_cache) immediate_cache = ImmediateCompletionCache.init(scratch);
 770     defer if (use_immediate_cache) immediate_cache.deinit();
 771     var search: SearchStats = .{
 772         .states_created = 1,
 773         .states_kept = 1,
 774         .frontier_peak = 1,
 775     };
 776     const incumbent_flop_bound = greedyUpperFlopCost(context);
 777 
 778     while (beam.items[0].active_count > 1) {
 779         var candidates = std.ArrayListUnmanaged(BeamState).empty;
 780         errdefer candidates.deinit(scratch);
 781         for (beam.items) |state| {
 782             for (0..state.active_count) |lhs_index| {
 783                 for (lhs_index + 1..state.active_count) |rhs_index| {
 784                     const estimate = beamCandidateEstimate(context, state, lhs_index, rhs_index);
 785                     search.candidate_contractions += 1;
 786                     if (state.active_count > 2 and estimate.prefix_flop_cost >= incumbent_flop_bound) {
 787                         search.states_pruned += 1;
 788                         continue;
 789                     }
 790                     try candidates.append(scratch, try beamCandidate(
 791                         use_immediate_cache,
 792                         scratch,
 793                         context,
 794                         state,
 795                         lhs_index,
 796                         rhs_index,
 797                         estimate,
 798                         &beam_steps,
 799                         if (use_immediate_cache) &immediate_cache else null,
 800                     ));
 801                     search.states_created += 1;
 802                 }
 803             }
 804         }
 805         if (candidates.items.len == 0) return error.InvalidState;
 806         if (use_immediate_cache) {
 807             retainBeamRefinementCandidates(&candidates, beamRefinementWidth(width, refinement_factor));
 808         } else {
 809             std.mem.sort(BeamState, candidates.items, {}, beamStateLessThan);
 810             retainSortedBeamCandidates(&candidates, beamRefinementWidth(width, refinement_factor));
 811         }
 812         if (shouldDeduplicateBeamCandidates(use_immediate_cache, beam.items[0].active_count, refinement_factor)) {
 813             try deduplicateBeamCandidates(&partition_cache, &candidates);
 814         }
 815         try scoreBeamCandidates(context, candidates.items, &completion_cache, completion_exact_limit, completion_score_limit);
 816         if (use_immediate_cache) {
 817             retainBeamCandidates(&candidates, width);
 818         } else {
 819             std.mem.sort(BeamState, candidates.items, {}, beamStateLessThan);
 820             retainSortedBeamCandidates(&candidates, width);
 821             markPreservedBeamPath(candidates.items);
 822         }
 823         search.states_kept += @intCast(candidates.items.len);
 824         search.frontier_peak = @max(search.frontier_peak, @as(u64, @intCast(candidates.items.len)));
 825         beam.deinit(scratch);
 826         beam = candidates;
 827     }
 828 
 829     const best = beam.items[0];
 830     const steps = try allocator.alloc(Step, best.step_count);
 831     writeBeamSteps(steps, beam_steps.items, best.step_ref);
 832     var plan = finishPlan(allocator, equation, .beam, steps, search);
 833     plan.selected_beam_width = width;
 834     return plan;
 835 }
 836 
 837 fn createAnytimePlan(allocator: std.mem.Allocator, context: *const PlanningContext, max_beam_width: usize) PlanError!Plan {
 838     return createAnytimePlanWithRefinement(allocator, context, max_beam_width, default_beam_refinement_factor);
 839 }
 840 
 841 fn createAutoBeamPlan(allocator: std.mem.Allocator, context: *const PlanningContext, max_beam_width: usize) PlanError!Plan {
 842     var best = try createAnytimePlanWithRefinement(allocator, context, max_beam_width, default_beam_refinement_factor);
 843     errdefer best.deinit();
 844     var search = best.search;
 845 
 846     var exploratory = try createBeamPlanWithRefinementAndScoreLimit(
 847         allocator,
 848         context,
 849         max_beam_width,
 850         default_beam_refinement_factor,
 851         auto_exploration_completion_exact_limit,
 852         auto_exploration_completion_score_limit,
 853     );
 854     search = mergedSearchStats(search, exploratory.search);
 855     if (planBeats(exploratory, best)) {
 856         best.deinit();
 857         best = exploratory;
 858     } else {
 859         exploratory.deinit();
 860     }
 861 
 862     best.search = search;
 863     return best;
 864 }
 865 
 866 fn createAnytimePlanWithRefinement(
 867     allocator: std.mem.Allocator,
 868     context: *const PlanningContext,
 869     max_beam_width: usize,
 870     refinement_factor: usize,
 871 ) PlanError!Plan {
 872     const final_width = @max(max_beam_width, 1);
 873     var width: usize = 1;
 874     var best = try createBeamPlanWithRefinement(allocator, context, width, refinement_factor);
 875     errdefer best.deinit();
 876     var search = best.search;
 877 
 878     while (width != final_width) {
 879         width = nextBeamWidth(width, final_width);
 880         var candidate = try createBeamPlanWithRefinement(allocator, context, width, refinement_factor);
 881         search = mergedSearchStats(search, candidate.search);
 882         if (planBeats(candidate, best)) {
 883             best.deinit();
 884             best = candidate;
 885         } else {
 886             candidate.deinit();
 887         }
 888     }
 889 
 890     best.strategy = .anytime;
 891     best.search = search;
 892     return best;
 893 }
 894 
 895 fn nextBeamWidth(width: usize, final_width: usize) usize {
 896     if (width >= final_width) return final_width;
 897     if (width > final_width / 2) return final_width;
 898     return @max(width * 2, width + 1);
 899 }
 900 
 901 fn planBeats(candidate: Plan, incumbent: Plan) bool {
 902     if (candidate.total_flop_cost != incumbent.total_flop_cost) return candidate.total_flop_cost < incumbent.total_flop_cost;
 903     if (candidate.peak_intermediate_elements != incumbent.peak_intermediate_elements) return candidate.peak_intermediate_elements < incumbent.peak_intermediate_elements;
 904     if (candidate.peak_workspace_elements != incumbent.peak_workspace_elements) return candidate.peak_workspace_elements < incumbent.peak_workspace_elements;
 905     return candidate.total_scalar_cost < incumbent.total_scalar_cost;
 906 }
 907 
 908 fn mergedSearchStats(lhs: SearchStats, rhs: SearchStats) SearchStats {
 909     return .{
 910         .candidate_contractions = saturatingAddU64(lhs.candidate_contractions, rhs.candidate_contractions),
 911         .states_created = saturatingAddU64(lhs.states_created, rhs.states_created),
 912         .states_kept = saturatingAddU64(lhs.states_kept, rhs.states_kept),
 913         .states_pruned = saturatingAddU64(lhs.states_pruned, rhs.states_pruned),
 914         .frontier_peak = @max(lhs.frontier_peak, rhs.frontier_peak),
 915     };
 916 }
 917 
 918 fn beamCandidateEstimate(
 919     context: *const PlanningContext,
 920     state: BeamState,
 921     lhs_index: usize,
 922     rhs_index: usize,
 923 ) BeamCandidateEstimate {
 924     const lhs = state.active[lhs_index];
 925     const rhs = state.active[rhs_index];
 926     const combined_mask = lhs.input_mask | rhs.input_mask;
 927     const pair_indices = lhs.indices.unioned(rhs.indices);
 928     const result_indices = resultIndices(context, combined_mask, pair_indices);
 929     const cost = stepCost(context, lhs.indices, rhs.indices, result_indices);
 930     return .{
 931         .combined_mask = combined_mask,
 932         .pair_indices = pair_indices,
 933         .result_indices = result_indices,
 934         .cost = cost,
 935         .prefix_flop_cost = saturatingAdd(state.flop_cost, cost.flop),
 936     };
 937 }
 938 
 939 fn beamCandidate(
 940     comptime use_immediate_cache: bool,
 941     allocator: std.mem.Allocator,
 942     context: *const PlanningContext,
 943     state: BeamState,
 944     lhs_index: usize,
 945     rhs_index: usize,
 946     estimate: BeamCandidateEstimate,
 947     beam_steps: *std.ArrayListUnmanaged(BeamStep),
 948     immediate_cache: ?*ImmediateCompletionCache,
 949 ) PlanError!BeamState {
 950     const equation = context.equation;
 951     const lhs = state.active[lhs_index];
 952     const rhs = state.active[rhs_index];
 953     const result_node: NodeId = @intCast(equation.inputs.len + state.step_count);
 954     const step: Step = .{
 955         .lhs = lhs.node,
 956         .rhs = rhs.node,
 957         .result = result_node,
 958         .lhs_inputs = lhs.input_mask,
 959         .rhs_inputs = rhs.input_mask,
 960         .result_inputs = estimate.combined_mask,
 961         .lhs_indices = lhs.indices,
 962         .rhs_indices = rhs.indices,
 963         .result_indices = estimate.result_indices,
 964         .summed_indices = estimate.pair_indices.without(estimate.result_indices),
 965         .scalar_cost = estimate.cost.scalar,
 966         .flop_cost = estimate.cost.flop,
 967         .result_elements = context.elementCount(estimate.result_indices),
 968         .work_elements = estimate.cost.work,
 969     };
 970 
 971     try beam_steps.append(allocator, .{
 972         .previous = state.step_ref,
 973         .step = step,
 974     });
 975     const step_ref = beam_steps.items.len - 1;
 976 
 977     const active = try allocator.alloc(Active, state.active_count - 1);
 978     writeBeamCandidateActive(
 979         active,
 980         state.active[0..state.active_count],
 981         lhs_index,
 982         rhs_index,
 983         .{
 984             .node = result_node,
 985             .input_mask = estimate.combined_mask,
 986             .indices = estimate.result_indices,
 987         },
 988     );
 989 
 990     const all_mask = allInputsMask(equation.inputs.len);
 991     const scalar_cost = saturatingAdd(state.scalar_cost, step.scalar_cost);
 992     const peak = @max(state.peak, stepPeakContribution(step, all_mask));
 993     const completion_score = if (use_immediate_cache and active.len >= beam_immediate_cache_min)
 994         try immediateCompletionFlopCostCached(context, active, immediate_cache.?)
 995     else
 996         immediateCompletionFlopCost(context, active);
 997     return .{
 998         .active = active,
 999         .active_count = active.len,
1000         .step_ref = step_ref,
1001         .step_count = state.step_count + 1,
1002         .scalar_cost = scalar_cost,
1003         .flop_cost = estimate.prefix_flop_cost,
1004         .score = saturatingAdd(estimate.prefix_flop_cost, completion_score),
1005         .peak = peak,
1006         .preserve = state.preserve,
1007     };
1008 }
1009 
1010 fn beamRefinementWidth(width: usize, refinement_factor: usize) usize {
1011     return @max(width * refinement_factor, beam_refinement_min);
1012 }
1013 
1014 fn writeBeamCandidateActive(active: []Active, source: []const Active, lhs_index: usize, rhs_index: usize, result: Active) void {
1015     var write: usize = 0;
1016     var inserted = false;
1017     for (source, 0..) |entry, index| {
1018         if (index == lhs_index or index == rhs_index) continue;
1019         if (!inserted and result.input_mask < entry.input_mask) {
1020             active[write] = result;
1021             write += 1;
1022             inserted = true;
1023         }
1024         active[write] = entry;
1025         write += 1;
1026     }
1027     if (!inserted) {
1028         active[write] = result;
1029     }
1030 }
1031 
1032 fn shouldDeduplicateBeamCandidates(comptime use_immediate_cache: bool, active_count: usize, refinement_factor: usize) bool {
1033     return (use_immediate_cache and refinement_factor == default_beam_refinement_factor and active_count <= beam_refined_partition_dedup_max_inputs) or
1034         active_count <= beam_partition_dedup_max_inputs;
1035 }
1036 
1037 fn retainBeamRefinementCandidates(candidates: *std.ArrayListUnmanaged(BeamState), width: usize) void {
1038     retainBestBeamCandidates(candidates, width, false);
1039 }
1040 
1041 fn scoreBeamCandidates(context: *const PlanningContext, candidates: []BeamState, cache: *CompletionCache, completion_exact_limit: usize, completion_score_limit: usize) PlanError!void {
1042     for (candidates) |*candidate| {
1043         if (candidate.active_count > completion_score_limit) continue;
1044         candidate.score = saturatingAdd(candidate.flop_cost, try completionFlopCost(context, candidate.active[0..candidate.active_count], cache, completion_exact_limit));
1045     }
1046 }
1047 
1048 fn deduplicateBeamCandidates(cache: *BeamPartitionMap, candidates: *std.ArrayListUnmanaged(BeamState)) PlanError!void {
1049     cache.clearRetainingCapacity();
1050     var write: usize = 0;
1051     for (candidates.items) |candidate| {
1052         const key = greedyCompletionKey(candidate.active[0..candidate.active_count]);
1053         if (cache.get(key)) |index| {
1054             const preserve = candidates.items[index].preserve or candidate.preserve;
1055             if (beamStateSelectionLessThan(candidate, candidates.items[index])) {
1056                 candidates.items[index] = candidate;
1057             }
1058             candidates.items[index].preserve = preserve;
1059         } else {
1060             candidates.items[write] = candidate;
1061             try cache.put(key, write);
1062             write += 1;
1063         }
1064     }
1065     candidates.items.len = write;
1066 }
1067 
1068 fn retainBeamCandidates(candidates: *std.ArrayListUnmanaged(BeamState), width: usize) void {
1069     retainBestBeamCandidates(candidates, width, true);
1070     markPreservedBeamPath(candidates.items);
1071 }
1072 
1073 fn retainSortedBeamCandidates(candidates: *std.ArrayListUnmanaged(BeamState), width: usize) void {
1074     const retained = @min(width, candidates.items.len);
1075     var preserved_index: ?usize = null;
1076     for (candidates.items, 0..) |candidate, index| {
1077         if (!candidate.preserve) continue;
1078         preserved_index = index;
1079         break;
1080     }
1081     if (preserved_index) |index| {
1082         if (index >= retained and retained > 0) {
1083             candidates.items[retained - 1] = candidates.items[index];
1084             std.mem.sort(BeamState, candidates.items[0..retained], {}, beamStateLessThan);
1085         }
1086     }
1087     candidates.items.len = retained;
1088 }
1089 
1090 fn retainBestBeamCandidates(candidates: *std.ArrayListUnmanaged(BeamState), width: usize, sort_retained: bool) void {
1091     const retained = @min(width, candidates.items.len);
1092     if (retained == 0) {
1093         candidates.items.len = 0;
1094         return;
1095     }
1096     if (candidates.items.len > retained) {
1097         const preserved = bestPreservedBeamCandidate(candidates.items);
1098         selectBestBeamPrefix(candidates.items, retained);
1099         candidates.items.len = retained;
1100         if (preserved) |candidate| {
1101             if (!beamCandidateRetained(candidates.items, candidate)) {
1102                 candidates.items[retained - 1] = candidate;
1103             }
1104         }
1105     }
1106     candidates.items.len = retained;
1107     if (sort_retained) std.mem.sort(BeamState, candidates.items, {}, beamStateLessThan);
1108 }
1109 
1110 fn bestPreservedBeamCandidate(candidates: []const BeamState) ?BeamState {
1111     var best: ?BeamState = null;
1112     for (candidates) |candidate| {
1113         if (!candidate.preserve) continue;
1114         if (best == null or beamStateSelectionLessThan(candidate, best.?)) {
1115             best = candidate;
1116         }
1117     }
1118     return best;
1119 }
1120 
1121 fn beamCandidateRetained(candidates: []const BeamState, needle: BeamState) bool {
1122     for (candidates) |candidate| {
1123         if (sameBeamCandidate(candidate, needle)) return true;
1124     }
1125     return false;
1126 }
1127 
1128 fn sameBeamCandidate(lhs: BeamState, rhs: BeamState) bool {
1129     return lhs.step_ref == rhs.step_ref and lhs.step_count == rhs.step_count;
1130 }
1131 
1132 fn selectBestBeamPrefix(candidates: []BeamState, retained: usize) void {
1133     const target = retained - 1;
1134     var left: usize = 0;
1135     var right = candidates.len - 1;
1136     while (left < right) {
1137         const pivot = partitionBeamCandidates(candidates, left, right, left + (right - left) / 2);
1138         if (target == pivot) return;
1139         if (target < pivot) {
1140             if (pivot == 0) return;
1141             right = pivot - 1;
1142         } else {
1143             left = pivot + 1;
1144         }
1145     }
1146 }
1147 
1148 fn partitionBeamCandidates(candidates: []BeamState, left: usize, right: usize, pivot_index: usize) usize {
1149     const pivot = candidates[pivot_index];
1150     std.mem.swap(BeamState, &candidates[pivot_index], &candidates[right]);
1151     var store = left;
1152     for (left..right) |index| {
1153         if (beamStateSelectionLessThan(candidates[index], pivot)) {
1154             std.mem.swap(BeamState, &candidates[store], &candidates[index]);
1155             store += 1;
1156         }
1157     }
1158     std.mem.swap(BeamState, &candidates[store], &candidates[right]);
1159     return store;
1160 }
1161 
1162 fn markPreservedBeamPath(states: []BeamState) void {
1163     var preserved = false;
1164     for (states) |*state| {
1165         if (state.preserve and !preserved) {
1166             preserved = true;
1167         } else {
1168             state.preserve = false;
1169         }
1170     }
1171 }
1172 
1173 fn writeBeamSteps(steps: []Step, beam_steps: []const BeamStep, step_ref: ?usize) void {
1174     var cursor = step_ref;
1175     var write = steps.len;
1176     while (cursor) |index| {
1177         write -= 1;
1178         const beam_step = beam_steps[index];
1179         steps[write] = beam_step.step;
1180         cursor = beam_step.previous;
1181     }
1182 }
1183 
1184 fn beamStateLessThan(_: void, lhs: BeamState, rhs: BeamState) bool {
1185     if (lhs.score != rhs.score) return lhs.score < rhs.score;
1186     if (lhs.flop_cost != rhs.flop_cost) return lhs.flop_cost < rhs.flop_cost;
1187     if (lhs.peak != rhs.peak) return lhs.peak < rhs.peak;
1188     if (lhs.scalar_cost != rhs.scalar_cost) return lhs.scalar_cost < rhs.scalar_cost;
1189     return lhs.step_count < rhs.step_count;
1190 }
1191 
1192 fn beamStateSelectionLessThan(lhs: BeamState, rhs: BeamState) bool {
1193     if (beamStateLessThan({}, lhs, rhs)) return true;
1194     if (beamStateLessThan({}, rhs, lhs)) return false;
1195     const lhs_step_ref = lhs.step_ref orelse std.math.maxInt(usize);
1196     const rhs_step_ref = rhs.step_ref orelse std.math.maxInt(usize);
1197     return lhs_step_ref < rhs_step_ref;
1198 }
1199 
1200 fn createOptimalPlan(allocator: std.mem.Allocator, context: *const PlanningContext, max_exact_inputs: u8) PlanError!Plan {
1201     const equation = context.equation;
1202     const input_count = equation.inputs.len;
1203     if (input_count > max_exact_inputs) return error.ExactInputLimitExceeded;
1204     if (input_count > 20) return error.ExactInputLimitExceeded;
1205     if (input_count == 1) {
1206         const steps = try allocator.alloc(Step, 0);
1207         return finishPlan(allocator, equation, .optimal, steps, .{
1208             .states_created = 1,
1209             .states_kept = 1,
1210             .frontier_peak = 1,
1211         });
1212     }
1213 
1214     const state_count: usize = @as(usize, 1) << @intCast(input_count);
1215     var states = try allocator.alloc(DpState, state_count);
1216     defer allocator.free(states);
1217     @memset(states, .{});
1218 
1219     for (equation.inputs, 0..) |input, index| {
1220         const mask: u64 = @as(u64, 1) << @intCast(index);
1221         states[@intCast(mask)] = .{
1222             .valid = true,
1223             .scalar_cost = 0,
1224             .flop_cost = 0,
1225             .peak = 0,
1226             .indices = input.index_set,
1227         };
1228     }
1229     var search: SearchStats = .{
1230         .states_created = @intCast(input_count),
1231         .states_kept = @intCast(input_count),
1232         .frontier_peak = @intCast(input_count),
1233     };
1234 
1235     var mask: u64 = 1;
1236     const all_mask: u64 = (@as(u64, 1) << @intCast(input_count)) - 1;
1237     while (mask <= all_mask) : (mask += 1) {
1238         if (isSingleton(mask)) continue;
1239         var mask_had_state = false;
1240         var sub = (mask - 1) & mask;
1241         while (sub != 0) : (sub = (sub - 1) & mask) {
1242             const other = mask ^ sub;
1243             if (other == 0 or sub > other) continue;
1244             const left = states[@intCast(sub)];
1245             const right = states[@intCast(other)];
1246             if (!left.valid or !right.valid) continue;
1247             search.candidate_contractions += 1;
1248             const pair_indices = left.indices.unioned(right.indices);
1249             const result = resultIndices(context, mask, pair_indices);
1250             const cost = stepCost(context, left.indices, right.indices, result);
1251             const result_elements = context.elementCount(result);
1252             const scalar_cost = saturatingAdd(saturatingAdd(left.scalar_cost, right.scalar_cost), cost.scalar);
1253             const flop_cost = saturatingAdd(saturatingAdd(left.flop_cost, right.flop_cost), cost.flop);
1254             const peak = @max(@max(left.peak, right.peak), stepPeakContributionValues(result_elements, cost.work, mask, all_mask));
1255             if (!states[@intCast(mask)].valid or flop_cost < states[@intCast(mask)].flop_cost or
1256                 (flop_cost == states[@intCast(mask)].flop_cost and peak < states[@intCast(mask)].peak) or
1257                 (flop_cost == states[@intCast(mask)].flop_cost and peak == states[@intCast(mask)].peak and scalar_cost < states[@intCast(mask)].scalar_cost))
1258             {
1259                 if (!mask_had_state) {
1260                     search.states_created += 1;
1261                     mask_had_state = true;
1262                 }
1263                 states[@intCast(mask)] = .{
1264                     .valid = true,
1265                     .scalar_cost = scalar_cost,
1266                     .flop_cost = flop_cost,
1267                     .peak = peak,
1268                     .indices = result,
1269                     .left = sub,
1270                     .right = other,
1271                 };
1272             }
1273         }
1274         if (mask_had_state) search.states_kept += 1;
1275     }
1276     if (!states[@intCast(all_mask)].valid) return error.InvalidState;
1277     search.frontier_peak = @max(search.frontier_peak, search.states_kept);
1278 
1279     var steps = std.ArrayListUnmanaged(Step).empty;
1280     errdefer steps.deinit(allocator);
1281     _ = try appendDpSteps(allocator, context, states, all_mask, &steps);
1282     const owned_steps = try steps.toOwnedSlice(allocator);
1283     return finishPlan(allocator, equation, .optimal, owned_steps, search);
1284 }
1285 
1286 fn appendDpSteps(
1287     allocator: std.mem.Allocator,
1288     context: *const PlanningContext,
1289     states: []const DpState,
1290     mask: u64,
1291     steps: *std.ArrayListUnmanaged(Step),
1292 ) PlanError!NodeId {
1293     const equation = context.equation;
1294     if (isSingleton(mask)) return @intCast(@ctz(mask));
1295     const state = states[@intCast(mask)];
1296     const lhs = try appendDpSteps(allocator, context, states, state.left, steps);
1297     const rhs = try appendDpSteps(allocator, context, states, state.right, steps);
1298     const lhs_state = states[@intCast(state.left)];
1299     const rhs_state = states[@intCast(state.right)];
1300     const pair_indices = lhs_state.indices.unioned(rhs_state.indices);
1301     const result = resultIndices(context, mask, pair_indices);
1302     const cost = stepCost(context, lhs_state.indices, rhs_state.indices, result);
1303     const result_node: NodeId = @intCast(equation.inputs.len + steps.items.len);
1304     try steps.append(allocator, .{
1305         .lhs = lhs,
1306         .rhs = rhs,
1307         .result = result_node,
1308         .lhs_inputs = state.left,
1309         .rhs_inputs = state.right,
1310         .result_inputs = mask,
1311         .lhs_indices = lhs_state.indices,
1312         .rhs_indices = rhs_state.indices,
1313         .result_indices = result,
1314         .summed_indices = pair_indices.without(result),
1315         .scalar_cost = cost.scalar,
1316         .flop_cost = cost.flop,
1317         .result_elements = context.elementCount(result),
1318         .work_elements = cost.work,
1319     });
1320     return result_node;
1321 }
1322 
1323 fn finishPlan(allocator: std.mem.Allocator, equation: *const Equation, strategy: Strategy, steps: []Step, search: SearchStats) Plan {
1324     var total: u128 = 0;
1325     var flops: u128 = 0;
1326     var peak: u128 = 0;
1327     var workspace: u128 = 0;
1328     const all_mask = allInputsMask(equation.inputs.len);
1329     for (steps) |step| {
1330         total = saturatingAdd(total, step.scalar_cost);
1331         flops = saturatingAdd(flops, step.flop_cost);
1332         peak = @max(peak, stepPeakContribution(step, all_mask));
1333         workspace = @max(workspace, step.work_elements);
1334     }
1335     if (steps.len == 0 and equation.inputs.len == 1) {
1336         const input_indices = equation.inputs[0].index_set;
1337         const summed = input_indices.without(equation.output_set);
1338         if (!summed.isEmpty()) {
1339             const reduction_cost = equation.elementCount(input_indices);
1340             total = saturatingAdd(total, reduction_cost);
1341             flops = saturatingAdd(flops, reduction_cost);
1342         }
1343     }
1344     return .{
1345         .allocator = allocator,
1346         .strategy = strategy,
1347         .steps = steps,
1348         .total_scalar_cost = total,
1349         .total_flop_cost = flops,
1350         .peak_intermediate_elements = peak,
1351         .peak_workspace_elements = workspace,
1352         .output_elements = equation.elementCount(equation.output_set),
1353         .selected_beam_width = 0,
1354         .search = search,
1355     };
1356 }
1357 
1358 fn bestGreedyPair(context: *const PlanningContext, active: []const Active, search: *SearchStats) PairChoice {
1359     search.candidate_contractions += candidatePairCount(active.len);
1360     return bestGreedyPairEstimate(context, active);
1361 }
1362 
1363 fn bestGreedyPairEstimate(context: *const PlanningContext, active: []const Active) PairChoice {
1364     var best = pairChoice(context, active[0], active[1], 0, 1);
1365     best.score = greedyScore(context, active, best);
1366     for (0..active.len) |lhs_index| {
1367         for (lhs_index + 1..active.len) |rhs_index| {
1368             if (lhs_index == 0 and rhs_index == 1) continue;
1369             var candidate = pairChoice(context, active[lhs_index], active[rhs_index], lhs_index, rhs_index);
1370             candidate.score = greedyScore(context, active, candidate);
1371             if (candidate.score < best.score or
1372                 (candidate.score == best.score and candidate.flop_cost < best.flop_cost) or
1373                 (candidate.score == best.score and candidate.flop_cost == best.flop_cost and choicePeak(candidate) < choicePeak(best)))
1374             {
1375                 best = candidate;
1376             }
1377         }
1378     }
1379     return best;
1380 }
1381 
1382 fn greedyUpperFlopCost(context: *const PlanningContext) u128 {
1383     const equation = context.equation;
1384     if (equation.inputs.len <= 1) {
1385         const input_indices = equation.inputs[0].index_set;
1386         const summed = input_indices.without(equation.output_set);
1387         if (summed.isEmpty()) return 0;
1388         return context.elementCount(input_indices);
1389     }
1390 
1391     var active: [63]Active = undefined;
1392     for (equation.inputs, 0..) |input, index| {
1393         active[index] = .{
1394             .node = 0,
1395             .input_mask = @as(u64, 1) << @intCast(index),
1396             .indices = input.index_set,
1397         };
1398     }
1399     var active_count = equation.inputs.len;
1400     var total: u128 = 0;
1401     while (active_count > 1) {
1402         const choice = bestGreedyPairEstimate(context, active[0..active_count]);
1403         const lhs = active[choice.lhs_index];
1404         const rhs = active[choice.rhs_index];
1405         const combined_mask = lhs.input_mask | rhs.input_mask;
1406         const pair_indices = lhs.indices.unioned(rhs.indices);
1407         const result = resultIndices(context, combined_mask, pair_indices);
1408         total = saturatingAdd(total, choice.flop_cost);
1409         active[choice.lhs_index] = .{
1410             .node = 0,
1411             .input_mask = combined_mask,
1412             .indices = result,
1413         };
1414         removeActive(active[0..active_count], choice.rhs_index);
1415         active_count -= 1;
1416     }
1417     return total;
1418 }
1419 
1420 fn greedyCompletionFlopCost(context: *const PlanningContext, active: []const Active) u128 {
1421     if (active.len <= 1) return 0;
1422 
1423     var scratch: [63]Active = undefined;
1424     @memcpy(scratch[0..active.len], active);
1425     var active_count = active.len;
1426     var total: u128 = 0;
1427     while (active_count > 1) {
1428         const choice = bestImmediatePairEstimate(context, scratch[0..active_count]);
1429         const lhs = scratch[choice.lhs_index];
1430         const rhs = scratch[choice.rhs_index];
1431         const combined_mask = lhs.input_mask | rhs.input_mask;
1432         const pair_indices = lhs.indices.unioned(rhs.indices);
1433         const result = resultIndices(context, combined_mask, pair_indices);
1434         total = saturatingAdd(total, choice.flop_cost);
1435         scratch[choice.lhs_index] = .{
1436             .node = 0,
1437             .input_mask = combined_mask,
1438             .indices = result,
1439         };
1440         removeActive(scratch[0..active_count], choice.rhs_index);
1441         active_count -= 1;
1442     }
1443     return total;
1444 }
1445 
1446 fn immediateCompletionFlopCost(context: *const PlanningContext, active: []const Active) u128 {
1447     if (active.len <= 1) return 0;
1448     return bestImmediatePairEstimate(context, active).flop_cost;
1449 }
1450 
1451 fn immediateCompletionFlopCostCached(
1452     context: *const PlanningContext,
1453     active: []const Active,
1454     cache: *ImmediateCompletionCache,
1455 ) PlanError!u128 {
1456     std.debug.assert(active.len >= beam_immediate_cache_min);
1457     const key = greedyCompletionKey(active);
1458     if (cache.get(key)) |cost| return cost;
1459     const cost = immediateCompletionFlopCost(context, active);
1460     try cache.put(key, cost);
1461     return cost;
1462 }
1463 
1464 fn completionFlopCost(context: *const PlanningContext, active: []const Active, cache: *CompletionCache, completion_exact_limit: usize) PlanError!u128 {
1465     std.debug.assert(completion_exact_limit <= beam_exact_completion_max_inputs);
1466     if (active.len <= completion_exact_limit) {
1467         const key = exactCompletionKey(active);
1468         if (cache.exact.get(key)) |cost| return cost;
1469         const cost = exactCompletionFlopCost(context, active);
1470         try cache.exact.put(key, cost);
1471         return cost;
1472     }
1473     const key = greedyCompletionKey(active);
1474     if (cache.greedy.get(key)) |cost| return cost;
1475     const cost = greedyCompletionFlopCost(context, active);
1476     try cache.greedy.put(key, cost);
1477     return cost;
1478 }
1479 
1480 fn exactCompletionKey(active: []const Active) ExactCompletionKey {
1481     std.debug.assert(active.len <= beam_exact_completion_max_inputs);
1482     var key = ExactCompletionKey{ .count = @intCast(active.len) };
1483     writeCompletionMasks(key.masks[0..], active);
1484     return key;
1485 }
1486 
1487 fn greedyCompletionKey(active: []const Active) GreedyCompletionKey {
1488     std.debug.assert(active.len <= max_completion_key_inputs);
1489     var key = GreedyCompletionKey{ .count = @intCast(active.len) };
1490     writeCompletionMasks(key.masks[0..], active);
1491     return key;
1492 }
1493 
1494 fn writeCompletionMasks(masks: []u64, active: []const Active) void {
1495     for (active, 0..) |entry, index| {
1496         if (index > 0) std.debug.assert(active[index - 1].input_mask < entry.input_mask);
1497         masks[index] = entry.input_mask;
1498     }
1499 }
1500 
1501 fn exactCompletionFlopCost(context: *const PlanningContext, active: []const Active) u128 {
1502     if (active.len <= 1) return 0;
1503     std.debug.assert(active.len <= beam_exact_completion_max_inputs);
1504 
1505     var states: [exact_completion_state_count]CompletionState = undefined;
1506     @memset(&states, .{});
1507     for (active, 0..) |entry, index| {
1508         const mask: u64 = @as(u64, 1) << @intCast(index);
1509         states[@intCast(mask)] = .{
1510             .valid = true,
1511             .flop_cost = 0,
1512             .indices = entry.indices,
1513             .input_mask = entry.input_mask,
1514         };
1515     }
1516 
1517     const all_mask = (@as(u64, 1) << @intCast(active.len)) - 1;
1518     var mask: u64 = 1;
1519     while (mask <= all_mask) : (mask += 1) {
1520         if (isSingleton(mask)) continue;
1521         var sub = (mask - 1) & mask;
1522         while (sub != 0) : (sub = (sub - 1) & mask) {
1523             const other = mask ^ sub;
1524             if (other == 0 or sub > other) continue;
1525             const left = states[@intCast(sub)];
1526             const right = states[@intCast(other)];
1527             if (!left.valid or !right.valid) continue;
1528             const pair_indices = left.indices.unioned(right.indices);
1529             const input_mask = left.input_mask | right.input_mask;
1530             const result = resultIndices(context, input_mask, pair_indices);
1531             const cost = stepCost(context, left.indices, right.indices, result);
1532             const flop_cost = saturatingAdd(saturatingAdd(left.flop_cost, right.flop_cost), cost.flop);
1533             if (!states[@intCast(mask)].valid or flop_cost < states[@intCast(mask)].flop_cost) {
1534                 states[@intCast(mask)] = .{
1535                     .valid = true,
1536                     .flop_cost = flop_cost,
1537                     .indices = result,
1538                     .input_mask = input_mask,
1539                 };
1540             }
1541         }
1542     }
1543     return states[@intCast(all_mask)].flop_cost;
1544 }
1545 
1546 fn bestImmediatePairEstimate(context: *const PlanningContext, active: []const Active) PairChoice {
1547     var best = pairChoice(context, active[0], active[1], 0, 1);
1548     for (0..active.len) |lhs_index| {
1549         for (lhs_index + 1..active.len) |rhs_index| {
1550             if (lhs_index == 0 and rhs_index == 1) continue;
1551             const candidate = pairChoice(context, active[lhs_index], active[rhs_index], lhs_index, rhs_index);
1552             if (candidate.flop_cost < best.flop_cost or
1553                 (candidate.flop_cost == best.flop_cost and choicePeak(candidate) < choicePeak(best)))
1554             {
1555                 best = candidate;
1556             }
1557         }
1558     }
1559     return best;
1560 }
1561 
1562 fn candidatePairCount(active_count: usize) u64 {
1563     if (active_count < 2) return 0;
1564     return @intCast((active_count * (active_count - 1)) / 2);
1565 }
1566 
1567 inline fn pairChoice(context: *const PlanningContext, lhs: Active, rhs: Active, lhs_index: usize, rhs_index: usize) PairChoice {
1568     const pair_indices = lhs.indices.unioned(rhs.indices);
1569     const combined_mask = lhs.input_mask | rhs.input_mask;
1570     const result = resultIndices(context, combined_mask, pair_indices);
1571     const cost = stepCost(context, lhs.indices, rhs.indices, result);
1572     return .{
1573         .lhs_index = lhs_index,
1574         .rhs_index = rhs_index,
1575         .cost = cost.scalar,
1576         .flop_cost = cost.flop,
1577         .score = cost.flop,
1578         .result_elements = context.elementCount(result),
1579         .work_elements = cost.work,
1580     };
1581 }
1582 
1583 inline fn stepCost(context: *const PlanningContext, lhs: IndexSet, rhs: IndexSet, result: IndexSet) StepCost {
1584     const pair_indices = lhs.unioned(rhs);
1585     const summed = pair_indices.without(result);
1586     const lhs_local = summed.intersected(lhs).without(rhs);
1587     const rhs_local = summed.intersected(rhs).without(lhs);
1588     const shared_summed = summed.intersected(lhs.intersected(rhs));
1589     const lhs_after = lhs.without(lhs_local);
1590     const rhs_after = rhs.without(rhs_local);
1591     const dot_indices = pair_indices.without(lhs_local).without(rhs_local);
1592     const dot_scalar = context.elementCount(dot_indices);
1593     var scalar = dot_scalar;
1594     var flops = flopCost(dot_scalar, shared_summed);
1595     var work: u128 = 0;
1596     if (!lhs_local.isEmpty()) {
1597         const reduction = context.elementCount(lhs);
1598         scalar = saturatingAdd(scalar, reduction);
1599         flops = saturatingAdd(flops, reduction);
1600         work = @max(work, context.elementCount(lhs_after));
1601     }
1602     if (!rhs_local.isEmpty()) {
1603         const reduction = context.elementCount(rhs);
1604         scalar = saturatingAdd(scalar, reduction);
1605         flops = saturatingAdd(flops, reduction);
1606         work = @max(work, context.elementCount(rhs_after));
1607     }
1608     if (!shared_summed.isEmpty() and !usesDotGeneralIndices(lhs_after, rhs_after, result, shared_summed)) {
1609         work = @max(work, dot_scalar);
1610     }
1611     return .{
1612         .scalar = scalar,
1613         .flop = flops,
1614         .work = work,
1615     };
1616 }
1617 
1618 fn usesDotGeneralIndices(lhs: IndexSet, rhs: IndexSet, result: IndexSet, shared_summed: IndexSet) bool {
1619     const shared_result = lhs.intersected(rhs).intersected(result);
1620     return shared_result.isEmpty() and
1621         lhs.count() == 2 and
1622         rhs.count() == 2 and
1623         result.count() == 2 and
1624         shared_summed.count() == 1;
1625 }
1626 
1627 fn stepPeakContribution(step: Step, all_mask: u64) u128 {
1628     return stepPeakContributionValues(step.result_elements, step.work_elements, step.result_inputs, all_mask);
1629 }
1630 
1631 fn stepPeakContributionValues(result_elements: u128, work_elements: u128, result_inputs: u64, all_mask: u64) u128 {
1632     const result_peak = if (result_inputs == all_mask) 0 else result_elements;
1633     return @max(result_peak, work_elements);
1634 }
1635 
1636 fn choicePeak(choice: PairChoice) u128 {
1637     return @max(choice.result_elements, choice.work_elements);
1638 }
1639 
1640 inline fn immediateFlopCost(context: *const PlanningContext, lhs: Active, rhs: Active) u128 {
1641     const pair_indices = lhs.indices.unioned(rhs.indices);
1642     const combined_mask = lhs.input_mask | rhs.input_mask;
1643     const result = resultIndices(context, combined_mask, pair_indices);
1644     return stepCost(context, lhs.indices, rhs.indices, result).flop;
1645 }
1646 
1647 inline fn resultIndices(context: *const PlanningContext, combined_mask: u64, pair_indices: IndexSet) IndexSet {
1648     return context.resultIndices(combined_mask, pair_indices);
1649 }
1650 
1651 fn removeActive(active: []Active, index: usize) void {
1652     var cursor = index;
1653     while (cursor + 1 < active.len) : (cursor += 1) {
1654         active[cursor] = active[cursor + 1];
1655     }
1656 }
1657 
1658 fn allInputsMask(input_count: usize) u64 {
1659     return (@as(u64, 1) << @intCast(input_count)) - 1;
1660 }
1661 
1662 fn isSingleton(mask: u64) bool {
1663     return mask != 0 and (mask & (mask - 1)) == 0;
1664 }
1665 
1666 fn saturatingAdd(lhs: u128, rhs: u128) u128 {
1667     const max = std.math.maxInt(u128);
1668     if (lhs > max - rhs) return max;
1669     return lhs + rhs;
1670 }
1671 
1672 fn saturatingAddU64(lhs: u64, rhs: u64) u64 {
1673     const max = std.math.maxInt(u64);
1674     if (lhs > max - rhs) return max;
1675     return lhs + rhs;
1676 }
1677 
1678 fn saturatingMul(lhs: u128, rhs: u128) u128 {
1679     if (lhs == 0 or rhs == 0) return 0;
1680     const max = std.math.maxInt(u128);
1681     if (lhs > max / rhs) return max;
1682     return lhs * rhs;
1683 }
1684 
1685 fn flopCost(scalar_cost: u128, summed_indices: IndexSet) u128 {
1686     if (summed_indices.isEmpty()) return scalar_cost;
1687     return saturatingMul(scalar_cost, 2);
1688 }
1689 
1690 fn expectSetContains(set: IndexSet, indices: []const u8) !void {
1691     for (indices) |index| try testing.expect(set.contains(index));
1692     try testing.expectEqual(@as(u32, @intCast(indices.len)), set.count());
1693 }
1694 
1695 fn greedyScore(context: *const PlanningContext, active: []const Active, choice: PairChoice) u128 {
1696     if (active.len <= 2) return choice.flop_cost;
1697 
1698     var next: [63]Active = undefined;
1699     const lhs = active[choice.lhs_index];
1700     const rhs = active[choice.rhs_index];
1701     const combined_mask = lhs.input_mask | rhs.input_mask;
1702     const pair_indices = lhs.indices.unioned(rhs.indices);
1703     const result = resultIndices(context, combined_mask, pair_indices);
1704     var next_count: usize = 0;
1705     for (active, 0..) |entry, index| {
1706         if (index == choice.rhs_index) continue;
1707         if (index == choice.lhs_index) {
1708             next[next_count] = .{
1709                 .node = 0,
1710                 .input_mask = combined_mask,
1711                 .indices = result,
1712             };
1713         } else {
1714             next[next_count] = entry;
1715         }
1716         next_count += 1;
1717     }
1718 
1719     var best_next = immediateFlopCost(context, next[0], next[1]);
1720     for (0..next_count) |lhs_index| {
1721         for (lhs_index + 1..next_count) |rhs_index| {
1722             best_next = @min(best_next, immediateFlopCost(context, next[lhs_index], next[rhs_index]));
1723         }
1724     }
1725     return saturatingAdd(choice.flop_cost, best_next);
1726 }
1727 
1728 test "einsum planner computes matrix product cost" {
1729     const shapes = [_][]const u64{ &.{ 2, 3 }, &.{ 3, 5 } };
1730     var equation = try spec.parse(testing.allocator, "ik,kj->ij", &shapes);
1731     defer equation.deinit();
1732 
1733     var plan = try createPlan(testing.allocator, &equation, .{ .strategy = .optimal });
1734     defer plan.deinit();
1735 
1736     try testing.expectEqual(Strategy.optimal, plan.strategy);
1737     try testing.expectEqual(@as(usize, 1), plan.steps.len);
1738     try testing.expectEqual(@as(u128, 30), plan.total_scalar_cost);
1739     try testing.expectEqual(@as(u128, 60), plan.total_flop_cost);
1740     try testing.expectEqual(@as(u128, 0), plan.peak_intermediate_elements);
1741     try testing.expectEqual(@as(u128, 0), plan.peak_workspace_elements);
1742     try testing.expectEqual(@as(u128, 10), plan.output_elements);
1743     try testing.expectEqual(@as(usize, 0), plan.selected_beam_width);
1744     try testing.expectEqual(@as(u64, 1), plan.search.candidate_contractions);
1745     try testing.expectEqual(@as(u64, 3), plan.search.states_created);
1746     try testing.expectEqual(@as(u64, 3), plan.search.states_kept);
1747     try testing.expectEqual(@as(u64, 0), plan.search.states_pruned);
1748     try testing.expectEqual(@as(u64, 3), plan.search.frontier_peak);
1749     try expectSetContains(plan.steps[0].summed_indices, "k");
1750     try expectSetContains(plan.steps[0].result_indices, "ij");
1751 }
1752 
1753 test "einsum planner handles scalar contraction output" {
1754     const shapes = [_][]const u64{ &.{7}, &.{7} };
1755     var equation = try spec.parse(testing.allocator, "i,i->", &shapes);
1756     defer equation.deinit();
1757 
1758     var plan = try createPlan(testing.allocator, &equation, .{ .strategy = .optimal });
1759     defer plan.deinit();
1760 
1761     try testing.expectEqual(@as(usize, 1), plan.steps.len);
1762     try testing.expectEqual(@as(u128, 7), plan.total_scalar_cost);
1763     try testing.expectEqual(@as(u128, 14), plan.total_flop_cost);
1764     try testing.expectEqual(@as(u128, 7), plan.peak_intermediate_elements);
1765     try testing.expectEqual(@as(u128, 7), plan.peak_workspace_elements);
1766     try testing.expectEqual(@as(u128, 1), plan.output_elements);
1767     try testing.expect(plan.steps[0].result_indices.isEmpty());
1768 }
1769 
1770 test "einsum planner accounts for single input reductions" {
1771     const shapes = [_][]const u64{&.{7}};
1772     var equation = try spec.parse(testing.allocator, "i->", &shapes);
1773     defer equation.deinit();
1774 
1775     var plan = try createPlan(testing.allocator, &equation, .{ .strategy = .optimal });
1776     defer plan.deinit();
1777 
1778     try testing.expectEqual(@as(usize, 0), plan.steps.len);
1779     try testing.expectEqual(@as(u128, 7), plan.total_scalar_cost);
1780     try testing.expectEqual(@as(u128, 7), plan.total_flop_cost);
1781     try testing.expectEqual(@as(u128, 0), plan.peak_intermediate_elements);
1782     try testing.expectEqual(@as(u128, 0), plan.peak_workspace_elements);
1783     try testing.expectEqual(@as(u128, 1), plan.output_elements);
1784 }
1785 
1786 test "einsum planner distinguishes matrix chain orderings" {
1787     const shapes = [_][]const u64{ &.{ 1000, 2 }, &.{ 2, 100 }, &.{ 100, 10 } };
1788     var equation = try spec.parse(testing.allocator, "ik,kl,lj->ij", &shapes);
1789     defer equation.deinit();
1790 
1791     var left = try createPlan(testing.allocator, &equation, .{ .strategy = .left_to_right });
1792     defer left.deinit();
1793     var greedy = try createPlan(testing.allocator, &equation, .{ .strategy = .greedy });
1794     defer greedy.deinit();
1795     var beam = try createPlan(testing.allocator, &equation, .{ .strategy = .beam });
1796     defer beam.deinit();
1797     var optimal = try createPlan(testing.allocator, &equation, .{ .strategy = .optimal });
1798     defer optimal.deinit();
1799 
1800     try testing.expectEqual(@as(u128, 1_200_000), left.total_scalar_cost);
1801     try testing.expectEqual(@as(u128, 22_000), greedy.total_scalar_cost);
1802     try testing.expectEqual(@as(u128, 22_000), beam.total_scalar_cost);
1803     try testing.expectEqual(@as(u128, 22_000), optimal.total_scalar_cost);
1804     try testing.expectEqual(Strategy.beam, beam.strategy);
1805     try testing.expectEqual(@as(u128, 44_000), optimal.total_flop_cost);
1806     try testing.expectEqual(@as(u64, 0b110), optimal.steps[0].result_inputs);
1807     try testing.expectEqual(@as(u128, 20), optimal.peak_intermediate_elements);
1808 }
1809 
1810 test "einsum exact planner optimizes four operand chain" {
1811     const shapes = [_][]const u64{ &.{ 20, 3 }, &.{ 3, 40 }, &.{ 40, 4 }, &.{ 4, 30 } };
1812     var equation = try spec.parse(testing.allocator, "ab,bc,cd,de->ae", &shapes);
1813     defer equation.deinit();
1814 
1815     var left = try createPlan(testing.allocator, &equation, .{ .strategy = .left_to_right });
1816     defer left.deinit();
1817     var greedy = try createPlan(testing.allocator, &equation, .{ .strategy = .greedy });
1818     defer greedy.deinit();
1819     var beam = try createPlan(testing.allocator, &equation, .{ .strategy = .beam, .beam_width = 3 });
1820     defer beam.deinit();
1821     var optimal = try createPlan(testing.allocator, &equation, .{ .strategy = .optimal });
1822     defer optimal.deinit();
1823 
1824     try testing.expect(left.total_scalar_cost > optimal.total_scalar_cost);
1825     try testing.expectEqual(optimal.total_flop_cost, greedy.total_flop_cost);
1826     try testing.expectEqual(optimal.total_flop_cost, beam.total_flop_cost);
1827     try testing.expectEqual(@as(usize, 3), optimal.steps.len);
1828     try testing.expectEqual(@as(u128, 2_640), optimal.total_scalar_cost);
1829     try testing.expectEqual(@as(u128, 5_280), optimal.total_flop_cost);
1830     try testing.expectEqual(@as(u64, 25), optimal.search.candidate_contractions);
1831     try testing.expectEqual(@as(u64, 15), optimal.search.states_created);
1832     try testing.expectEqual(@as(u64, 15), optimal.search.states_kept);
1833     try testing.expectEqual(@as(u64, 15), optimal.search.frontier_peak);
1834 }
1835 
1836 test "einsum exact planner optimizes reported flop cost" {
1837     const shapes = [_][]const u64{ &.{ 5, 11 }, &.{ 5, 5 }, &.{ 11, 5, 19 }, &.{ 19, 5 } };
1838     var equation = try spec.parse(testing.allocator, "cd,ca,dae,ec->", &shapes);
1839     defer equation.deinit();
1840 
1841     var beam = try createPlan(testing.allocator, &equation, .{ .strategy = .beam });
1842     defer beam.deinit();
1843     var optimal = try createPlan(testing.allocator, &equation, .{ .strategy = .optimal });
1844     defer optimal.deinit();
1845 
1846     try testing.expectEqual(@as(u128, 5_595), optimal.total_scalar_cost);
1847     try testing.expectEqual(@as(u128, 10_915), optimal.total_flop_cost);
1848     try testing.expectEqual(optimal.total_flop_cost, beam.total_flop_cost);
1849 }
1850 
1851 test "einsum flop cost does not double pure outer-product steps" {
1852     const shapes = [_][]const u64{ &.{ 32, 4 }, &.{ 4, 64 }, &.{ 4, 3 }, &.{ 4, 16 } };
1853     var equation = try spec.parse(testing.allocator, "ab,bc,bd,be->acde", &shapes);
1854     defer equation.deinit();
1855 
1856     var greedy = try createPlan(testing.allocator, &equation, .{ .strategy = .greedy });
1857     defer greedy.deinit();
1858     var beam = try createPlan(testing.allocator, &equation, .{ .strategy = .beam });
1859     defer beam.deinit();
1860     var optimal = try createPlan(testing.allocator, &equation, .{ .strategy = .optimal });
1861     defer optimal.deinit();
1862 
1863     try testing.expectEqual(@as(u128, 396_032), optimal.total_scalar_cost);
1864     try testing.expectEqual(@as(u128, 789_248), optimal.total_flop_cost);
1865     try testing.expectEqual(optimal.total_flop_cost, greedy.total_flop_cost);
1866     try testing.expectEqual(optimal.total_flop_cost, beam.total_flop_cost);
1867 }
1868 
1869 test "einsum planner prices operand-local reductions before product" {
1870     const shapes = [_][]const u64{ &.{ 2, 3 }, &.{ 5, 7 } };
1871     var equation = try spec.parse(testing.allocator, "ab,cd->ac", &shapes);
1872     defer equation.deinit();
1873 
1874     var plan = try createPlan(testing.allocator, &equation, .{ .strategy = .optimal });
1875     defer plan.deinit();
1876 
1877     try testing.expectEqual(@as(usize, 1), plan.steps.len);
1878     try testing.expectEqual(@as(u128, 51), plan.steps[0].scalar_cost);
1879     try testing.expectEqual(@as(u128, 51), plan.steps[0].flop_cost);
1880     try testing.expectEqual(@as(u128, 5), plan.steps[0].work_elements);
1881     try testing.expectEqual(@as(u128, 51), plan.total_scalar_cost);
1882     try testing.expectEqual(@as(u128, 51), plan.total_flop_cost);
1883     try testing.expectEqual(@as(u128, 5), plan.peak_intermediate_elements);
1884     try testing.expectEqual(@as(u128, 5), plan.peak_workspace_elements);
1885 }
1886 
1887 test "einsum planner accounts for non-matmul contraction workspace" {
1888     const shapes = [_][]const u64{ &.{ 2, 3, 5 }, &.{ 2, 7, 5 } };
1889     var equation = try spec.parse(testing.allocator, "bqd,bkd->bqk", &shapes);
1890     defer equation.deinit();
1891 
1892     var plan = try createPlan(testing.allocator, &equation, .{ .strategy = .optimal });
1893     defer plan.deinit();
1894 
1895     try testing.expectEqual(@as(usize, 1), plan.steps.len);
1896     try testing.expectEqual(@as(u128, 210), plan.steps[0].scalar_cost);
1897     try testing.expectEqual(@as(u128, 420), plan.steps[0].flop_cost);
1898     try testing.expectEqual(@as(u128, 210), plan.steps[0].work_elements);
1899     try testing.expectEqual(@as(u128, 210), plan.peak_intermediate_elements);
1900     try testing.expectEqual(@as(u128, 210), plan.peak_workspace_elements);
1901 }
1902 
1903 test "einsum beam width controls retained frontier" {
1904     const shapes = [_][]const u64{ &.{ 32, 4 }, &.{ 4, 64 }, &.{ 4, 3 }, &.{ 4, 16 } };
1905     var equation = try spec.parse(testing.allocator, "ab,bc,bd,be->acde", &shapes);
1906     defer equation.deinit();
1907 
1908     var narrow = try createPlan(testing.allocator, &equation, .{ .strategy = .beam, .beam_width = 2 });
1909     defer narrow.deinit();
1910     var wide = try createPlan(testing.allocator, &equation, .{ .strategy = .beam });
1911     defer wide.deinit();
1912 
1913     try testing.expectEqual(@as(u128, 789_248), narrow.total_flop_cost);
1914     try testing.expectEqual(@as(u128, 789_248), wide.total_flop_cost);
1915     try testing.expectEqual(@as(usize, 2), narrow.selected_beam_width);
1916     try testing.expectEqual(@as(usize, 64), wide.selected_beam_width);
1917     try testing.expectEqual(@as(u64, 14), narrow.search.candidate_contractions);
1918     try testing.expectEqual(@as(u64, 15), narrow.search.states_created);
1919     try testing.expectEqual(@as(u64, 6), narrow.search.states_kept);
1920     try testing.expectEqual(@as(u64, 0), narrow.search.states_pruned);
1921     try testing.expectEqual(@as(u64, 2), narrow.search.frontier_peak);
1922 }
1923 
1924 test "einsum beam ranking uses greedy completion estimate" {
1925     const shapes = [_][]const u64{ &.{ 32, 4 }, &.{ 4, 64 }, &.{ 4, 3 }, &.{ 4, 16 }, &.{ 4, 5 }, &.{ 4, 8 } };
1926     var equation = try spec.parse(testing.allocator, "ab,bc,bd,be,bf,bg->acdefg", &shapes);
1927     defer equation.deinit();
1928 
1929     var narrow = try createPlan(testing.allocator, &equation, .{ .strategy = .beam, .beam_width = 2 });
1930     defer narrow.deinit();
1931     var optimal = try createPlan(testing.allocator, &equation, .{ .strategy = .optimal });
1932     defer optimal.deinit();
1933 
1934     try testing.expectEqual(@as(u128, 31_473_504), optimal.total_flop_cost);
1935     try testing.expectEqual(optimal.total_flop_cost, narrow.total_flop_cost);
1936     try testing.expectEqual(@as(u64, 55), narrow.search.candidate_contractions);
1937     try testing.expectEqual(@as(u64, 2), narrow.search.frontier_peak);
1938     try testing.expectEqual(@as(u64, 0), narrow.search.states_pruned);
1939 }
1940 
1941 test "einsum beam prunes states above greedy incumbent bound" {
1942     const shapes = [_][]const u64{
1943         &.{ 32, 32, 7 },
1944         &.{ 11, 16 },
1945         &.{ 64, 32, 16, 32 },
1946         &.{ 7, 32, 64, 11 },
1947         &.{ 16, 32, 4 },
1948         &.{ 4, 16 },
1949     };
1950     var equation = try spec.parse(testing.allocator, "dfh,ae,cgeb,hfca,efi,ie->d", &shapes);
1951     defer equation.deinit();
1952 
1953     var plan = try createPlan(testing.allocator, &equation, .{ .strategy = .beam });
1954     defer plan.deinit();
1955 
1956     try testing.expectEqual(@as(u128, 2_108_928), plan.total_flop_cost);
1957     try testing.expect(plan.search.states_pruned > 0);
1958     try testing.expect(plan.search.candidate_contractions > plan.search.states_created);
1959 }
1960 
1961 test "einsum anytime planner keeps best widening result" {
1962     const shapes = [_][]const u64{ &.{ 32, 4 }, &.{ 4, 64 }, &.{ 4, 3 }, &.{ 4, 16 }, &.{ 4, 5 }, &.{ 4, 8 } };
1963     var equation = try spec.parse(testing.allocator, "ab,bc,bd,be,bf,bg->acdefg", &shapes);
1964     defer equation.deinit();
1965 
1966     var beam = try createPlan(testing.allocator, &equation, .{ .strategy = .beam, .beam_width = 8 });
1967     defer beam.deinit();
1968     var anytime = try createPlan(testing.allocator, &equation, .{ .strategy = .anytime, .beam_width = 8 });
1969     defer anytime.deinit();
1970 
1971     try testing.expectEqual(Strategy.anytime, anytime.strategy);
1972     try testing.expect(anytime.selected_beam_width > 0);
1973     try testing.expect(anytime.selected_beam_width <= 8);
1974     try testing.expect(anytime.total_flop_cost <= beam.total_flop_cost);
1975     try testing.expect(anytime.search.candidate_contractions > beam.search.candidate_contractions);
1976     try testing.expectEqual(@as(u64, 8), anytime.search.frontier_peak);
1977 }
1978 
1979 test "einsum auto planner selects exact plan within state budget" {
1980     const shapes = [_][]const u64{
1981         &.{ 7, 5, 7 },
1982         &.{ 11, 4, 7 },
1983         &.{ 7, 5, 11 },
1984         &.{ 8, 3, 7, 5 },
1985         &.{ 7, 8, 7 },
1986         &.{ 5, 7 },
1987         &.{ 7, 16, 7, 5 },
1988         &.{ 5, 16, 7, 11 },
1989         &.{ 8, 5, 11, 5 },
1990         &.{ 8, 3, 7, 5 },
1991     };
1992     var equation = try spec.parse(testing.allocator, "hik,jak,hfj,edhg,hek,bh,kchb,gckj,egjf,edhf->acgi", &shapes);
1993     defer equation.deinit();
1994 
1995     var plan = try createPlan(testing.allocator, &equation, .{ .strategy = .auto });
1996     defer plan.deinit();
1997 
1998     try testing.expectEqual(Strategy.optimal, plan.strategy);
1999     try testing.expectEqual(@as(usize, 0), plan.selected_beam_width);
2000     try testing.expectEqual(@as(u128, 622_020), plan.total_flop_cost);
2001     try testing.expectEqual(@as(u128, 172_480), plan.peak_workspace_elements);
2002 }
2003 
2004 test "einsum auto planner falls back to beam outside exact state budget" {
2005     const shapes = [_][]const u64{
2006         &.{ 7, 5, 7 },
2007         &.{ 11, 4, 7 },
2008         &.{ 7, 5, 11 },
2009         &.{ 8, 3, 7, 5 },
2010         &.{ 7, 8, 7 },
2011         &.{ 5, 7 },
2012         &.{ 7, 16, 7, 5 },
2013         &.{ 5, 16, 7, 11 },
2014         &.{ 8, 5, 11, 5 },
2015         &.{ 8, 3, 7, 5 },
2016     };
2017     var equation = try spec.parse(testing.allocator, "hik,jak,hfj,edhg,hek,bh,kchb,gckj,egjf,edhf->acgi", &shapes);
2018     defer equation.deinit();
2019 
2020     var plan = try createPlan(testing.allocator, &equation, .{ .strategy = .auto, .exact_state_limit = 1 });
2021     defer plan.deinit();
2022 
2023     try testing.expectEqual(Strategy.beam, plan.strategy);
2024     try testing.expectEqual(@as(usize, 64), plan.selected_beam_width);
2025     try testing.expectEqual(@as(u128, 622_020), plan.total_flop_cost);
2026     try testing.expectEqual(@as(u128, 172_480), plan.peak_workspace_elements);
2027 }
2028 
2029 test "einsum auto planner uses anytime beyond exact input boundary" {
2030     const shapes = [_][]const u64{
2031         &.{ 2, 2 },
2032         &.{ 2, 2 },
2033         &.{ 2, 2 },
2034         &.{ 2, 2 },
2035         &.{ 2, 2 },
2036         &.{ 2, 2 },
2037         &.{ 2, 2 },
2038         &.{ 2, 2 },
2039         &.{ 2, 2 },
2040         &.{ 2, 2 },
2041         &.{ 2, 2 },
2042         &.{ 2, 2 },
2043         &.{ 2, 2 },
2044         &.{ 2, 2 },
2045         &.{ 2, 2 },
2046         &.{ 2, 2 },
2047         &.{ 2, 2 },
2048         &.{ 2, 2 },
2049         &.{ 2, 2 },
2050     };
2051     var equation = try spec.parse(testing.allocator, "ab,bc,cd,de,ef,fg,gh,hi,ij,jk,kl,lm,mn,no,op,pq,qr,rs,st->at", &shapes);
2052     defer equation.deinit();
2053 
2054     var plan = try createPlan(testing.allocator, &equation, .{ .strategy = .auto, .auto_beam_width = 2 });
2055     defer plan.deinit();
2056 
2057     try testing.expectEqual(Strategy.anytime, plan.strategy);
2058     try testing.expect(plan.selected_beam_width > 0);
2059     try testing.expect(plan.selected_beam_width <= 2);
2060     try testing.expectEqual(@as(usize, 18), plan.steps.len);
2061 }
2062 
2063 test "einsum auto planner optimizes seventeen input tail fixture" {
2064     if (builtin.mode != .fast) return error.SkipZigTest;
2065 
2066     const shapes = [_][]const u64{
2067         &.{ 3, 4, 8, 8 },
2068         &.{ 7, 8, 5 },
2069         &.{ 3, 8, 2, 5 },
2070         &.{ 7, 8, 3, 8 },
2071         &.{ 7, 8, 4 },
2072         &.{ 8, 2, 8, 4 },
2073         &.{ 8, 3, 7, 8 },
2074         &.{ 2, 11, 3, 5 },
2075         &.{ 2, 5, 7, 2 },
2076         &.{ 5, 8, 8, 2 },
2077         &.{ 5, 4, 8, 8 },
2078         &.{ 11, 3, 8 },
2079         &.{ 11, 4, 3, 3 },
2080         &.{ 4, 11, 7, 2 },
2081         &.{ 8, 11, 8, 8 },
2082         &.{ 8, 11, 7, 8 },
2083         &.{ 3, 8, 7, 8 },
2084     };
2085     var equation = try spec.parse(testing.allocator, "Abot,ktz,Vcjz,IJVp,ILs,MOcs,MSho,OPSg,Ogkw,Hapw,Hbpu,deu,PQUe,QXhj,FXac,FPhn,Ackn->HUX", &shapes);
2086     defer equation.deinit();
2087 
2088     var plan = try createPlan(testing.allocator, &equation, .{ .strategy = .auto });
2089     defer plan.deinit();
2090 
2091     try testing.expectEqual(Strategy.optimal, plan.strategy);
2092     try testing.expectEqual(@as(u128, 142_374_072), plan.total_flop_cost);
2093 }
2094 
2095 test "einsum auto planner uses exact boundary for fifteen inputs" {
2096     if (builtin.mode != .fast) return error.SkipZigTest;
2097 
2098     const shapes = [_][]const u64{
2099         &.{ 8, 11, 13 },
2100         &.{ 8, 4 },
2101         &.{ 11, 11 },
2102         &.{ 11, 13, 13, 11 },
2103         &.{ 13, 11, 11, 4 },
2104         &.{ 7, 7, 4 },
2105         &.{ 7, 11, 3 },
2106         &.{ 4, 2, 13 },
2107         &.{ 13, 4, 11, 3 },
2108         &.{ 7, 4, 3 },
2109         &.{ 3, 11 },
2110         &.{ 4, 13, 7, 4 },
2111         &.{ 8, 13 },
2112         &.{ 4, 13, 7, 3 },
2113         &.{ 3, 2, 3, 7 },
2114     };
2115     var equation = try spec.parse(testing.allocator, "bam,bj,ka,alik,moaf,pdj,dkc,jhe,ijkc,djc,ca,jidf,bm,jeng,ghcn->nlio", &shapes);
2116     defer equation.deinit();
2117 
2118     var beam = try createPlan(testing.allocator, &equation, .{ .strategy = .beam });
2119     defer beam.deinit();
2120     var auto = try createPlan(testing.allocator, &equation, .{ .strategy = .auto });
2121     defer auto.deinit();
2122     var capped = try createPlan(testing.allocator, &equation, .{ .strategy = .auto, .exact_state_limit = 16_384 });
2123     defer capped.deinit();
2124 
2125     try testing.expectEqual(Strategy.optimal, auto.strategy);
2126     try testing.expectEqual(Strategy.beam, capped.strategy);
2127     try testing.expectEqual(@as(u128, 4_434_221), auto.total_flop_cost);
2128     try testing.expectEqual(beam.total_flop_cost, capped.total_flop_cost);
2129     try testing.expect(beam.total_flop_cost > auto.total_flop_cost);
2130 }
2131 
2132 test "einsum exact state default covers eighteen inputs" {
2133     try testing.expect(canUseExactPlan(16, .{}));
2134     try testing.expect(canUseExactPlan(17, .{}));
2135     try testing.expect(canUseExactPlan(18, .{}));
2136     try testing.expect(!canUseExactPlan(19, .{}));
2137     try testing.expect(!canUseExactPlan(18, .{ .exact_state_limit = 131_072 }));
2138 }
2139 
2140 test "einsum auto planner uses exact boundary for sixteen inputs" {
2141     if (builtin.mode != .fast) return error.SkipZigTest;
2142 
2143     const shapes = [_][]const u64{
2144         &.{ 11, 7, 3, 8 },
2145         &.{ 3, 7, 3, 11 },
2146         &.{ 13, 3, 2 },
2147         &.{ 8, 5, 2, 13 },
2148         &.{ 8, 5, 13, 8 },
2149         &.{ 3, 13, 11, 2 },
2150         &.{ 13, 11, 3 },
2151         &.{ 11, 8, 13, 4 },
2152         &.{ 11, 8, 13, 4 },
2153         &.{ 13, 11, 7, 8 },
2154         &.{ 13, 8, 13, 8 },
2155         &.{ 8, 3, 8 },
2156         &.{ 3, 11, 7 },
2157         &.{ 13, 11, 8, 11 },
2158         &.{ 13, 5, 13, 2 },
2159         &.{ 5, 11, 7, 8 },
2160     };
2161     var equation = try spec.parse(testing.allocator, "XZab,YZas,CYn,bhnu,fhkz,cksx,JPc,Pbiw,Dfuw,ADZp,AGip,GMl,MSU,QSfs,ALQx,LXZz->QXw", &shapes);
2162     defer equation.deinit();
2163 
2164     var auto = try createPlan(testing.allocator, &equation, .{ .strategy = .auto });
2165     defer auto.deinit();
2166     var capped = try createPlan(testing.allocator, &equation, .{ .strategy = .auto, .exact_state_limit = 32_768 });
2167     defer capped.deinit();
2168 
2169     try testing.expectEqual(Strategy.optimal, auto.strategy);
2170     try testing.expectEqual(Strategy.beam, capped.strategy);
2171     try testing.expectEqual(@as(u128, 754_408_598), auto.total_flop_cost);
2172     try testing.expectEqual(@as(u128, 2_835_181_970), capped.total_flop_cost);
2173 }
2174 
2175 test "einsum planner work bound includes every greedy lookahead pair" {
2176     for (2..17) |count| {
2177         var evaluations: u64 = 0;
2178         for (0..count) |left| {
2179             for (left + 1..count) |_| {
2180                 evaluations += 1;
2181                 if (count == 2) continue;
2182                 evaluations += 1;
2183                 for (0..count - 1) |next_left| {
2184                     for (next_left + 1..count - 1) |_| evaluations += 1;
2185                 }
2186             }
2187         }
2188         try testing.expectEqual(evaluations, PlannerWork.greedyEvaluations(count));
2189         const linear = try (Options{ .strategy = .left_to_right }).workBound(count);
2190         const greedy = try (Options{ .strategy = .greedy }).workBound(count);
2191         try testing.expect(greedy >= linear);
2192         if (count > 2) try testing.expect(greedy > linear);
2193     }
2194 }
2195 
2196 test "einsum planner work bound covers finite partition enumeration" {
2197     for (0..11) |count| {
2198         const states = @as(u64, 1) << @intCast(count);
2199         var assignments: u64 = 0;
2200         for (0..states) |mask| {
2201             var subset = mask;
2202             for (0..states) |_| {
2203                 assignments += 1;
2204                 if (subset == 0) break;
2205                 subset = (subset - 1) & mask;
2206             } else unreachable;
2207         }
2208         try testing.expectEqual(assignments, try PlannerWork.partitions(count));
2209         if (count > 1) {
2210             try testing.expect(try PlannerWork.completion(count, count) >= assignments);
2211         }
2212     }
2213 }
2214 
2215 test "einsum planner work bound follows exact automatic and widening choices" {
2216     const shapes = [_][]const u64{ &.{ 32, 4 }, &.{ 4, 64 }, &.{ 4, 3 }, &.{ 4, 16 } };
2217     var equation = try spec.parse(testing.allocator, "ab,bc,bd,be->acde", &shapes);
2218     defer equation.deinit();
2219     for (comptime std.meta.tags(Strategy)) |strategy| {
2220         const options: Options = .{ .strategy = strategy, .beam_width = 3 };
2221         var plan = try createPlan(testing.allocator, &equation, options);
2222         defer plan.deinit();
2223         try testing.expect(try options.workBound(4) > plan.search.candidate_contractions);
2224         try testing.expect(try options.workBound(1) > 0);
2225     }
2226     const exact = try (Options{ .strategy = .optimal }).workBound(4);
2227     try testing.expectEqual(exact, try (Options{}).workBound(4));
2228     try testing.expectEqual(exact, try (Options{
2229         .strategy = .optimal,
2230         .exact_state_limit = 0,
2231     }).workBound(4));
2232     const beam = try (Options{ .strategy = .beam, .beam_width = 0 }).workBound(4);
2233     try testing.expectEqual(beam, try (Options{ .exact_state_limit = 0, .beam_width = 0 })
2234         .workBound(4));
2235     var widening: u64 = 0;
2236     for ([_]usize{ 1, 2, 3 }) |width| {
2237         widening += try (Options{ .strategy = .beam, .beam_width = width }).workBound(4);
2238     }
2239     try testing.expectEqual(widening, try (Options{ .strategy = .anytime, .beam_width = 3 })
2240         .workBound(4));
2241     try testing.expect(try (Options{ .max_exact_inputs = 2, .auto_beam_width = 3 })
2242         .workBound(4) > widening);
2243 }
2244 
2245 test "einsum planner work bound handles table boundaries and checked overflow" {
2246     const linear: Options = .{ .strategy = .left_to_right };
2247     for ([_]usize{ 10, 11, 18, 19, 63 }) |count| {
2248         try testing.expect(try linear.workBound(count) > count);
2249     }
2250     try testing.expect(try linear.workBound(18) > try linear.workBound(19));
2251     try testing.expectError(error.EmptyEquation, linear.workBound(0));
2252     try testing.expectError(error.TooManyInputs, linear.workBound(64));
2253     const exact: Options = .{ .strategy = .optimal, .max_exact_inputs = 63 };
2254     try testing.expectError(error.ExactInputLimitExceeded, exact.workBound(21));
2255     const enormous: Options = .{ .strategy = .beam, .beam_width = std.math.maxInt(usize) };
2256     try testing.expectError(error.WorkOverflow, enormous.workBound(11));
2257     try testing.expectError(error.WorkOverflow, PlannerWork.product(std.math.maxInt(u64), 2));
2258     try testing.expectError(error.WorkOverflow, PlannerWork.sum(std.math.maxInt(u64), 1));
2259 }
2260 
2261 fn storagePlanWitness(text: []const u8, shapes: []const []const u64, options: Options) !void {
2262     var equation = try spec.parse(testing.allocator, text, shapes);
2263     defer equation.deinit();
2264     var expected = try createPlan(testing.allocator, &equation, options);
2265     defer expected.deinit();
2266     try testing.expect(try options.workBound(shapes.len) > expected.search.candidate_contractions);
2267     const bound = try options.storageBound(shapes.len);
2268     const bytes = try testing.allocator.alloc(u8, @intCast(bound));
2269     defer testing.allocator.free(bytes);
2270     var buffer = std.heap.FixedBufferAllocator.init(bytes);
2271     const base = buffer.allocator();
2272     const vtable: std.mem.Allocator.VTable = .{
2273         .alloc = base.vtable.alloc,
2274         .resize = std.mem.Allocator.noResize,
2275         .remap = std.mem.Allocator.noRemap,
2276         .free = std.mem.Allocator.noFree,
2277     };
2278     const allocator: std.mem.Allocator = .{ .ptr = base.ptr, .vtable = &vtable };
2279     var plan = try createPlan(allocator, &equation, options);
2280     defer plan.deinit();
2281     try testing.expectEqual(expected.strategy, plan.strategy);
2282     try testing.expectEqual(expected.selected_beam_width, plan.selected_beam_width);
2283     try testing.expectEqualDeep(expected.steps, plan.steps);
2284     try testing.expectEqualDeep(expected.search, plan.search);
2285     try testing.expectEqual(expected.total_flop_cost, plan.total_flop_cost);
2286     try testing.expect(buffer.end_index <= bound);
2287 }
2288 
2289 test "einsum planner storage bound covers every strategy and automatic branch" {
2290     const shapes = [_][]const u64{ &.{ 32, 4 }, &.{ 4, 64 }, &.{ 4, 3 }, &.{ 4, 16 } };
2291     for (comptime std.meta.tags(Strategy)) |strategy| {
2292         try storagePlanWitness("ab,bc,bd,be->acde", &shapes, .{
2293             .strategy = strategy,
2294             .beam_width = 3,
2295         });
2296         try storagePlanWitness("ab->a", &.{&.{ 2, 3 }}, .{ .strategy = strategy });
2297     }
2298     try storagePlanWitness("ab,bc,bd,be->acde", &shapes, .{
2299         .exact_state_limit = 0,
2300         .beam_width = 0,
2301     });
2302     try storagePlanWitness("ab,bc,bd,be->acde", &shapes, .{
2303         .max_exact_inputs = 2,
2304         .auto_beam_width = 3,
2305     });
2306     const exact = try (Options{ .strategy = .optimal }).storageBound(4);
2307     try testing.expectEqual(exact, try (Options{}).storageBound(4));
2308     try testing.expectEqual(exact, try (Options{
2309         .strategy = .optimal,
2310         .exact_state_limit = 0,
2311     }).storageBound(4));
2312 }
2313 
2314 test "einsum planner storage bound covers completion caches and widening" {
2315     const shapes: [19][]const u64 = @splat(&.{ 2, 2 });
2316     const short = "ab,bc,cd,de,ef,fg,gh,hi,ij,jk,kl->al";
2317     for ([_]Strategy{ .beam, .anytime, .optimal }) |strategy| {
2318         try storagePlanWitness(short, shapes[0..11], .{
2319             .strategy = strategy,
2320             .beam_width = 3,
2321         });
2322     }
2323     try storagePlanWitness(
2324         "ab,bc,cd,de,ef,fg,gh,hi,ij,jk,kl,lm,mn,no,op,pq,qr,rs,st->at",
2325         &shapes,
2326         .{ .auto_beam_width = 2 },
2327     );
2328     const beam = try (Options{ .strategy = .beam, .beam_width = 3 }).storageBound(11);
2329     const anytime = try (Options{ .strategy = .anytime, .beam_width = 3 }).storageBound(11);
2330     try testing.expect(anytime > beam);
2331 }
2332 
2333 test "einsum planner storage bound covers precomputation and input boundaries" {
2334     const shapes: [63][]const u64 = @splat(&.{2});
2335     var text: [129]u8 = undefined;
2336     for (0..63) |index| {
2337         text[index * 2] = 'i';
2338         text[index * 2 + 1] = ',';
2339     }
2340     for ([_]usize{ 10, 11, 18, 19, 63 }) |count| {
2341         var equation = text;
2342         @memcpy(equation[count * 2 - 1 ..][0..3], "->i");
2343         try storagePlanWitness(equation[0 .. count * 2 + 2], shapes[0..count], .{
2344             .strategy = .left_to_right,
2345         });
2346     }
2347     const linear: Options = .{ .strategy = .left_to_right };
2348     const with_tables = try linear.storageBound(18);
2349     try testing.expect(with_tables >= 2 * (1 << 18) * @sizeOf(IndexSet));
2350     try testing.expect(try linear.storageBound(19) < with_tables);
2351     try testing.expectError(error.EmptyEquation, linear.storageBound(0));
2352     try testing.expectError(error.TooManyInputs, linear.storageBound(64));
2353     try testing.expectError(error.ExactInputLimitExceeded, (Options{
2354         .strategy = .optimal,
2355         .max_exact_inputs = 63,
2356     }).storageBound(21));
2357     try testing.expectError(error.WorkOverflow, (Options{
2358         .strategy = .beam,
2359         .beam_width = std.math.maxInt(usize),
2360     }).storageBound(2));
2361 }