Implement graph coloring register allocator

[jserv]

Description

Replace the current simple register allocator with a graph coloring allocator based on the Chaitin-Briggs algorithm. This classical approach models register allocation as a graph coloring problem where variables are nodes and interference represents edges.

Current Problem

The existing allocator lacks global view of register allocation, leading to:

• Suboptimal register assignments
• Excessive spilling in loops
• Poor handling of high register pressure
• No consideration of variable importance

Graph Coloring Algorithm

1. Core Data Structures

typedef struct interference_node {
    var_t *var;
    bitset_t *neighbors;      // Variables that interfere
    int degree;               // Number of neighbors
    int color;                // Assigned register (-1 if spilled)
    float spill_cost;
    bool on_stack;
    bool precolored;          // Fixed register (e.g., function args)
    struct interference_node *next;
} ig_node_t;

typedef struct {
    ig_node_t **nodes;        // Array of nodes
    int node_count;
    int num_colors;           // Available registers
    stack_t *simplify_stack;
    list_t *spill_worklist;
    list_t *freeze_worklist;
    list_t *simplify_worklist;
} ig_graph_t;

2. Building Interference Graph

void build_interference_graph(func_t *func, ig_graph_t *graph) {
    // Step 1: Compute live ranges
    compute_live_ranges(func);
    
    // Step 2: Create nodes for each variable
    for (var in all_variables) {
        ig_node_t *node = create_node(var);
        add_to_graph(graph, node);
    }
    
    // Step 3: Add interference edges
    for (bb in func->bbs) {
        live_set_t *live = get_live_out(bb);
        
        for (insn in reverse(bb->insns)) {
            // Definition kills, use makes live
            if (insn->rd) {
                // rd interferes with all live variables
                for (var in live) {
                    if (var != insn->rd) {
                        add_interference(graph, insn->rd, var);
                    }
                }
                remove_from_set(live, insn->rd);
            }
            
            // Add uses to live set
            if (insn->rs1) add_to_set(live, insn->rs1);
            if (insn->rs2) add_to_set(live, insn->rs2);
        }
    }
}

3. Simplification Phase

void simplify(ig_graph_t *graph) {
    while (!worklist_empty(graph->simplify_worklist)) {
        ig_node_t *node = select_node(graph->simplify_worklist);
        
        if (node->degree < graph->num_colors) {
            // Can definitely color this node
            push(graph->simplify_stack, node);
            remove_from_graph(node);
            
            // Update neighbors' degrees
            for (neighbor in node->neighbors) {
                neighbor->degree--;
                if (neighbor->degree < graph->num_colors) {
                    move_to_simplify_worklist(neighbor);
                }
            }
        }
    }
}

4. Spill Selection

ig_node_t *select_spill_node(ig_graph_t *graph) {
    ig_node_t *best = NULL;
    float min_cost = INFINITY;
    
    for (node in graph->spill_worklist) {
        // Spill cost heuristic: (uses + defs) / degree
        float cost = node->spill_cost / node->degree;
        
        // Prefer spilling outside loops
        if (node->var->loop_depth > 0) {
            cost *= pow(10, node->var->loop_depth);
        }
        
        if (cost < min_cost) {
            min_cost = cost;
            best = node;
        }
    }
    
    return best;
}

void potential_spill(ig_graph_t *graph) {
    while (!worklist_empty(graph->spill_worklist)) {
        ig_node_t *node = select_spill_node(graph);
        push(graph->simplify_stack, node);
        remove_from_graph(node);
        // Mark as potential spill
        node->may_spill = true;
    }
}

5. Coloring Phase

void assign_colors(ig_graph_t *graph) {
    while (!stack_empty(graph->simplify_stack)) {
        ig_node_t *node = pop(graph->simplify_stack);
        
        // Find available colors
        bitset_t *used_colors = create_bitset(graph->num_colors);
        
        for (neighbor in node->neighbors) {
            if (neighbor->color >= 0) {
                set_bit(used_colors, neighbor->color);
            }
        }
        
        // Try to assign a color
        int color = find_first_zero_bit(used_colors);
        
        if (color < graph->num_colors) {
            node->color = color;
        } else {
            // Must spill this variable
            add_to_spill_list(node->var);
        }
    }
}

6. Main Allocation Loop

void allocate_registers(func_t *func) {
    bool done = false;
    
    while (!done) {
        ig_graph_t *graph = create_graph();
        
        // Build interference graph
        build_interference_graph(func, graph);
        
        // Coalesce moves (optional optimization)
        coalesce_moves(graph);
        
        // Simplify graph
        make_worklist(graph);
        simplify(graph);
        
        // Handle remaining high-degree nodes
        if (!worklist_empty(graph->spill_worklist)) {
            potential_spill(graph);
        }
        
        // Assign colors
        assign_colors(graph);
        
        // Check if we need to spill
        if (has_spills(graph)) {
            // Insert spill code and retry
            rewrite_program(func, graph);
        } else {
            done = true;
            apply_allocation(func, graph);
        }
        
        destroy_graph(graph);
    }
}

7. Spill Code Generation

void insert_spill_code(func_t *func, var_t *spilled) {
    int stack_slot = allocate_stack_slot(spilled);
    
    // Insert store after definition
    for (insn in all_instructions) {
        if (insn->rd == spilled) {
            insert_after(insn, create_store(spilled, stack_slot));
        }
    }
    
    // Insert load before use
    for (insn in all_instructions) {
        if (uses(insn, spilled)) {
            var_t *temp = create_temp_var();
            insert_before(insn, create_load(temp, stack_slot));
            replace_use(insn, spilled, temp);
        }
    }
}

Test Cases

// High interference - forces graph coloring
void test_high_interference() {
    int a=1, b=2, c=3, d=4, e=5, f=6, g=7, h=8;
    int i=9, j=10, k=11, l=12, m=13, n=14, o=15, p=16;
    
    // All variables interfere
    return a+b+c+d+e+f+g+h+i+j+k+l+m+n+o+p;
}

// Disjoint live ranges - should reuse registers
void test_disjoint_ranges() {
    {
        int a = 1, b = 2, c = 3;
        use(a, b, c);
    }  // a,b,c dead
    {
        int d = 4, e = 5, f = 6;  // Should reuse a,b,c's registers
        use(d, e, f);
    }
}

// Loop with interference
void test_loop_interference(int n) {
    int sum = 0;
    for (int i = 0; i < n; i++) {
        int t1 = i * 2;
        int t2 = i * 3;
        int t3 = i * 4;
        sum += t1 + t2 + t3;  // All interfere here
    }
    return sum;
}

Benefits

• Optimal register allocation for many programs
• Global view of register allocation
• Handles complex interference patterns
• Well-studied algorithm with proven properties

Success Metrics

• 30-50% reduction in spill code
• Successful compilation of complex programs
• Performance improvement on benchmarks
• Correct handling of all test cases

References

• Chaitin et al. "Register allocation via coloring" (1981)
• Briggs et al. "Improvements to graph coloring" (1989)
• Appel & George "Optimal spilling for CISC machines" (2001)