DOCSEARCHbyAlgolia
Submit Docs

regalloc2

repository·main·Indexed 42 days ago

https://github.com/bytecodealliance/regalloc2

A Rust implementation of a backtracking register allocator inspired by IonMonkey. It features a fast compile-time allocator called Fastalloc, which implements a reverse linear scan register allocation/SSRA algorithm with Selection, Assignment, Eviction, and Edit Insertion phases. Version 0.15.2 includes extensive testing infrastructure, fuzzing harnesses, and checkers.

Tokens
19.3K
Snippets
31
Records
80
Agent score
38%

What's inside regalloc2

  1. Overview of regalloc2

    main

    regalloc2 is a register allocator implemented in Rust. It originated as a port of IonMonkey's backtracking register allocator but has since been generalized, optimized, and improved. It includes extensive testing infrastructure, such as fuzzing harnesses and checkers, which were not present in the original IonMonkey allocator.

    For a deep dive into the allocator's mechanics, refer to the design overview.

  2. How the allocator probes physical registers

    main

    If a bundle requires a register, the allocator probes physical registers in a specific order to minimize contention:

    1. Hint Register: If the spillset provides a hint (from the last allocation in that spillset), it is probed first.
    2. Preferred Registers: All registers marked as preferred.
    3. Non-preferred Registers: All remaining registers in the required class.

    Probing Mechanism: To distribute contention, the allocator uses an offset for probing. Instead of always starting at register 0, it starts at an index determined by a heuristic (the sum of the bundle index and the instruction index of the first range's start). It then marches through the sequence, wrapping around if necessary.

    For each register, the allocator checks:

    • If the entire range is free.
    • If not, the vector of conflicting bundles and the highest weight among them.
    • If not, the first conflict point.

  3. Core allocation data structures: Ranges, Bundles, and VRegs

    main

    The allocator manages state through several hierarchical structures:

    • ProgPoint: A precise location in the program consisting of an instruction index and a Before or After suffix (packed into the LSB of a u32).
    • Live-range: A contiguous, half-open range of ProgPoints [from, to) where a VReg is live. Each range contains a vector of Uses (packed into three u32s containing the Operand, ProgPoint, slot, and weight).
    • Bundles: A collection of non-overlapping live-ranges. Bundles carry an allocation (starting as none), a spillset, and metadata like spill weight, priority, and property flags (minimal, contains fixed constraints, contains stack constraints).
    • VRegs: Virtual registers, which contain vectors of their constituent live-ranges.
    • Physical Registers: Managed via an allocation map (a BTree) indexed by ProgPoint ranges. The map uses a custom comparison operator where overlapping ranges are considered equal, allowing for single-probe checks to see if a range is free.

  4. How VReg spillslots are managed

    main

    When a Virtual Register (VReg) requires a spillslot, a dedicated slot is allocated and recorded in vreg_spillslots. If a VReg never needs to spill, no slot is allocated.

    To maintain consistency, the allocator ensures a VReg is in its spillslot when expected. During the processing of a def operand, before it is deallocated, the allocator inserts an edit to move the value from its current allocation (as tracked in vreg_allocs) to its dedicated spillslot (if one exists in vreg_spillslots).

  5. How Operand Allocation is performed

    main

    When allocating an operand, Fastalloc checks if the VReg's current allocation (vreg_allocs) already satisfies the operand's constraints.

    Case 1: Current allocation satisfies constraints

    • Assignment: Set the operand's final allocation to the current one.
    • Selection: Mark the PReg as unavailable in the appropriate available_pregs set.
    • LRU Update: Update the LRU for the register class to mark this PReg as the Most Recently Used.
    • Note: No eviction or edit insertion is required.

    Case 2: Current allocation does NOT satisfy constraints

    • Selection: Create a 'draw-from' set from available_pregs:
      • Early uses / Late defs: Use the early or late set respectively.
      • Late uses / Early defs: Use the intersection of early and late sets.
      • Traverse the LRU from the end (least recently used) to find the first PReg in the draw-from set.
      • Mark the PReg as MRU, unavailable in available_pregs, and evict the previous occupant.
    • Assignment: Set the operand's final allocation to the selected PReg.
    • Lifecycle Management:
      • If this is the first time the VReg is encountered, insert it into live_vregs and update vreg_in_preg.
      • If the operand is a def, insert an edit after the instruction to move the value from the new allocation to the expected post-instruction allocation.

  6. How debug information is constructed

    main

    Debug locations are built after all blocks have completed processing, using liverange information tracked during allocation.

    Tracking Liveranges

    • First Allocation: When a VReg is first allocated, its liverange end is saved in the vreg_to_live_inst_range vector.
    • Definition: When a VReg's definition is encountered, its liverange beginning is saved.
    • Location: The allocation it will occupy throughout that range is also recorded.

    Determining Allocation for a Liverange

    To determine the correct allocation for a VReg throughout its liverange, the allocator checks its spillslot status:

    1. If the VReg has no dedicated spillslot, it implies the VReg was never evicted or spilled. Its current entry in vreg_allocs is used as the location for the entire liverange.
    2. If the VReg has a dedicated spillslot, it implies the VReg was evicted, was a livein of a predecessor, or was a block parameter. In this case, the dedicated spillslot is recorded as the allocation for the vreg_to_live_inst_range info.

  7. Deciding between splitting and evicting bundles

    main

    When a register probe results in a conflict, the allocator must choose between evicting existing bundles or splitting the current bundle.

    Eviction

    Eviction is chosen if the cost (the maximum weight of the conflicting bundles) is lower than the cost of splitting.

    Splitting

    A split is chosen if the bundle is not already minimal AND:

    • The allocator has already attempted to evict once without success,
    • OR the weight of the current bundle is less than the eviction cost.

    Split Cost Calculation: A split's cost is defined as the cost of the move (the weight of a normal def operand at the split point) plus the cost of all bundles beyond the split point that remain in conflict.

    Important Constraints:

    • If a conflict involves a fixed allocation, eviction is not an option; the allocator must attempt a split.
    • The allocator never evicts if the current bundle is less important (lower weight) than the bundles being evicted, even if splitting would be more expensive. This ensures forward progress.

  8. How reuse operands are handled in Fastalloc

    main

    Fastalloc manages operand reuse through specific constraints to minimize register pressure:

    • Reuse def operands: A new operand is created that is identical to the reuse def, but its constraints are updated to match the constraints of the reused input.
    • Reused inputs: A new operand is created with a fixed register constraint, forcing it to use the same register that was assigned to the original reuse def.

    To prevent conflicts, the allocator marks registers as unavailable in the available_pregs sets during selection: if an operand is a reused input, its PReg is marked unavailable in the available_pregs late set; if it is a reuse def, the selected register is marked unavailable in the available_pregs early set.

  9. How Fastalloc works: Design Overview

    main

    Fastalloc is a register allocator optimized for fast compile times. It implements a reverse linear scan register allocation/SSRA algorithm.

    Key Characteristics:

    • Direction: Allocation proceeds in reverse order: from the last block to the first block, and within each block, from the last instruction to the first instruction.
    • Phases: For every operand in an instruction, the allocator performs four distinct phases: Selection, Assignment, Eviction, and Edit Insertion.
    • Operand Lifecycle: A virtual register (VReg) typically changes its allocation state a few times: from none $\rightarrow$ a physical register (PReg) $\rightarrow$ a spillslot (if evicted or crossing block boundaries) $\rightarrow$ none (when its definition is reached).

  10. How bundle merging works in Ion

    main

    After live-range vectors are built, the allocator attempts to merge them into larger units called bundles. This is the only stage where bundles increase in size.

    Merging Logic:

    • Initial State: Each non-pinned vreg starts with one bundle.
    • Target Pairs: The allocator attempts to merge two types of pairs:
      1. Reused-input to its corresponding output.
      2. Across blockparam assignments.
    • Conflict Checking:
      • A coarse check is performed using 1-bit flags for fixed-reg or stack-only constraints.
      • If the coarse check passes, a detailed check computes merged requirements to ensure no Conflict (lattice bottom value) occurs.
    • Implementation Detail: To avoid the complexity of an in-place mergesort, the allocator appends one vector to the other and uses Rust's built-in sort to merge the sorted liverange vectors.

    Once merging is complete, bundle properties (priorities and weights) are computed and the bundles are enqueued in a priority queue for the allocation phase.

  11. How program-derived alloc-invariant data works

    main

    Before the core allocation begins, certain data structures are computed deterministically from the input program and remain read-only during allocation:

    • Livein/Liveout Bitsets: Stored on the Env, these bitsets (liveins and liveouts) record which VRegs are live entering and leaving each basic block. They are computed via backward iterative dataflow analysis and are exact (no over-approximation).
    • Blockparam Vectors: Two vectors, blockparam_ins and blockparam_outs, represent blockparam flows. They use a 'build lazily then sort' pattern to provide dense storage and efficient scanning.
      • blockparam_ins contains (to-vreg, to-block, from-block) tuples.
      • blockparam_outs contains (from-vreg, from-block, to-block, to-vreg) tuples.
      • There is exactly one 'out' tuple for every 'in' tuple, allowing the allocator to scan both to generate moves.

  12. How bundle splitting and post-processing works

    main

    The Split Procedure

    When a bundle is split at a specific point:

    1. A new bundle is created for the second half ("rest").
    2. Liveranges are distributed based on the split point. If a split point falls in the middle of a liverange, that liverange is also split.
    3. New liveranges are appended to the corresponding virtual register (vreg) liverange vectors.
    4. A "minimal split" may be used to simply peel off the first use to ensure forward progress.

    Post-split Cleanup (Trimming)

    To improve results, the allocator performs post-processing to move "loose ends" to a spill bundle:

    • If the last liverange in the first-half bundle extends beyond its last use, the excess is moved to a spill bundle as an empty liverange.
    • If the first liverange in the second-half bundle starts before its first use, the preceding part is moved to a spill bundle.

    This reduces register contention and ensures that "flow-through" parts of a vreg's liveness are kept in a single, non-split bundle.