Merge pull request #15 from CompilerProgramming/chaitin

Chatin Graph Coloring Register Allocator continued

Author: dibyendumajumdar

optvm/README.md

@@ -4,4 +4,36 @@ This module implements various compiler optimization techniques such as:
 
 * Static Single Assignment
 * Liveness Analysis
-* WIP Graph Coloring Register Allocator (Chaitin)
+* Graph Coloring Register Allocator (Chaitin)
+
+Our goal here is to perform optimizations on the Intermediate Representation targeting an abstract machine, rather than
+a physical machine. Therefore, all our optimization passes will work on the instruction set of this abstract machine.
+
+
+## Guide
+
+* [Register](src/main/java/com/compilerprogramming/ezlang/compiler/Register.java) - implements a virtual register. Virtual registers
+  have a name, type, and id - the id is unique, but name is not. Initially the compiler generates unique registers for every local
+  and temporary, the number of registers grows as we convert to SSA and back out of SSA. Finally, as we run the Chaitin register allocator
+  we shrink down the number of virtual registers to the minimum.
+* [RegisterPool](src/main/java/com/compilerprogramming/ezlang/compiler/RegisterPool.java) - simple pool to allow us to find a register
+  by its id, and to allocate new virtual registers.
+* [CompiledFunction](src/main/java/com/compilerprogramming/ezlang/compiler/CompiledFunction.java) - encapsulates the IR for a single function.
+* [BasicBlock](src/main/java/com/compilerprogramming/ezlang/compiler/BasicBlock.java) - Defines our basic block - which contains instructions
+  that execute sequentially. A basic block ends with a branch. There are two distinguished basic blocks in every function: entry and exit. 
+* [BBHelper](src/main/java/com/compilerprogramming/ezlang/compiler/BBHelper.java) - Some utilities that manipulate basic blocks.
+* [Operand](src/main/java/com/compilerprogramming/ezlang/compiler/Operand.java) - Operands in instructions.
+* [Instruction](src/main/java/com/compilerprogramming/ezlang/compiler/Instruction.java) - Instructions in basic blocks.
+* [DominatorTree](src/main/java/com/compilerprogramming/ezlang/compiler/DominatorTree.java) - Calculates dominator tree and dominance frontiers.
+* [LiveSet](src/main/java/com/compilerprogramming/ezlang/compiler/LiveSet.java) - Bitset used to track liveness of registers.
+* [Liveness](src/main/java/com/compilerprogramming/ezlang/compiler/Liveness.java) - Liveness calculator, works for both SSA and non-SSA forms.
+* [EnterSSA](src/main/java/com/compilerprogramming/ezlang/compiler/EnterSSA.java) - Transforms into SSA, using algorithm by Preston Briggs.
+* [ExitSSA](src/main/java/com/compilerprogramming/ezlang/compiler/ExitSSA.java) - Exits SSA form, using algorithm by Preston Briggs.
+* [LoopFinder](src/main/java/com/compilerprogramming/ezlang/compiler/LoopFinder.java) - Discovers loops.
+* [LoopNest](src/main/java/com/compilerprogramming/ezlang/compiler/LoopNest.java) - Representation of loop nesting.
+* [InterferenceGraph](src/main/java/com/compilerprogramming/ezlang/compiler/InterferenceGraph.java) - Representation of an Interference Graph
+  required by the register allocator.
+* [InterferenceGraphBuilder](src/main/java/com/compilerprogramming/ezlang/compiler/InterferenceGraph.java) - Constructs InteferenceGraph for a set
+  of basic bocks, using liveness information.
+* [ChaitinGraphColoringRegisterAllocator](src/main/java/com/compilerprogramming/ezlang/compiler/ChaitinGraphColoringRegisterAllocator.java) - basic
+  Chaitin Graph Coloring Register Allocator - WIP.

optvm/src/main/java/com/compilerprogramming/ezlang/compiler/BasicBlock.java

@@ -126,6 +126,15 @@ public List<Instruction.Phi> phis() {
         }
         return list;
     }
+    public int whichPred(BasicBlock s) {
+        int i = 0;
+        for (BasicBlock p: s.predecessors) {
+            if (p == this)
+                return i;
+            i++;
+        }
+        throw new IllegalStateException();
+    }
     public static StringBuilder toStr(StringBuilder sb, BasicBlock bb, BitSet visited, boolean dumpLiveness)
     {
         if (visited.get(bb.bid))

optvm/src/main/java/com/compilerprogramming/ezlang/compiler/ChaitinGraphColoringRegisterAllocator.java

@@ -1,57 +1,156 @@
 package com.compilerprogramming.ezlang.compiler;
 
-import java.util.ArrayList;
-import java.util.List;
+import java.util.*;
+import java.util.stream.IntStream;
 
+/**
+ * Implement the original graph coloring algorithm described by Chaitin.
+ *
+ * TODO spilling
+ */
 public class ChaitinGraphColoringRegisterAllocator {
 
-    public ChaitinGraphColoringRegisterAllocator(CompiledFunction function) {
-        coalesce(function);
+    public ChaitinGraphColoringRegisterAllocator() {
     }
 
-    private void coalesce(CompiledFunction function) {
+    public Map<Integer, Integer> assignRegisters(CompiledFunction function, int numRegisters) {
+        if (function.isSSA) throw new IllegalStateException("Register allocation should be done after exiting SSA");
+        var g = coalesce(function);
+        var registers = registersInIR(function);
+        var colors = IntStream.range(0, numRegisters).boxed().toList();
+        // TODO pre-assign regs to args
+        // TODO spilling
+        var assignments = colorGraph(g, registers, new HashSet<>(colors));
+        return assignments;
+    }
+
+    /**
+     * Chaitin: coalesce_nodes - coalesce away copy operations
+     */
+    public InterferenceGraph coalesce(CompiledFunction function) {
         boolean changed = true;
+        InterferenceGraph igraph = null;
         while (changed) {
-            var igraph = new InterferenceGraphBuilder().build(function);
-            changed = coalesceRegisters(function, igraph);
+            igraph = new InterferenceGraphBuilder().build(function);
+            changed = coalesceCopyOperations(function, igraph);
         }
+        return igraph;
     }
 
-    private boolean coalesceRegisters(CompiledFunction function, InterferenceGraph igraph) {
+    /**
+     * Chaitin: coalesce_nodes - coalesce away copy operations
+     */
+    private boolean coalesceCopyOperations(CompiledFunction function, InterferenceGraph igraph) {
         boolean changed = false;
         for (var block: function.getBlocks()) {
-            List<Integer> instructionsToRemove = new ArrayList<>();
-            for (int j = 0; j < block.instructions.size(); j++) {
-                Instruction i = block.instructions.get(j);
-                if (i instanceof Instruction.Move move
-                    && move.from() instanceof Operand.RegisterOperand targetOperand) {
+            Iterator<Instruction> iter = block.instructions.iterator();
+            while (iter.hasNext()) {
+                Instruction instruction = iter.next();
+                if (instruction instanceof Instruction.Move move
+                    && move.from() instanceof Operand.RegisterOperand registerTarget) {
                     Register source = move.def();
-                    Register target = targetOperand.reg;
+                    Register target = registerTarget.reg;
                     if (source.id != target.id &&
                         !igraph.interfere(target.id, source.id)) {
                         igraph.rename(source.id, target.id);
-                        rewriteInstructions(function, i, source, target);
-                        instructionsToRemove.add(j);
+                        rewriteInstructions(function, instruction, source, target);
+                        iter.remove();
                         changed = true;
                     }
                 }
             }
-            for (var j: instructionsToRemove) {
-                block.instructions.set(j, new Instruction.NoOp());
-            }
         }
         return changed;
     }
 
-    private void rewriteInstructions(CompiledFunction function, Instruction notNeeded, Register source, Register target) {
+    /**
+     * Chaitin: rewrite_il
+     */
+    private void rewriteInstructions(CompiledFunction function, Instruction deadInstruction, Register source, Register target) {
         for (var block: function.getBlocks()) {
             for (Instruction i: block.instructions) {
-                if (i == notNeeded)
+                if (i == deadInstruction)
                     continue;
                 if (i.definesVar() && source.id == i.def().id)
                     i.replaceDef(target);
                 i.replaceUse(source, target);
             }
         }
     }
+
+    /**
+     * Get the list of registers in use in the Intermediate Code
+     * Chaitin: registers_in_il()
+     */
+    private Set<Integer> registersInIR(CompiledFunction function) {
+        Set<Integer> registers = new HashSet<>();
+        for (var block: function.getBlocks()) {
+            Iterator<Instruction> iter = block.instructions.iterator();
+            while (iter.hasNext()) {
+                Instruction instruction = iter.next();
+                if (instruction.definesVar())
+                    registers.add(instruction.def().id);
+                for (Register use: instruction.uses())
+                    registers.add(use.id);
+            }
+        }
+        return registers;
+    }
+
+    /**
+     * Chaitin: color_graph line 2-3
+     */
+    private Integer findNodeWithNeighborCountLessThan(InterferenceGraph g, Set<Integer> nodes, int numColors) {
+        for (var node: nodes) {
+            if (g.neighbors(node).size() < numColors) {
+                return node;
+            }
+        }
+        return null;
+    }
+
+    private Set<Integer> getNeighborColors(InterferenceGraph g, Integer node, Map<Integer,Integer> assignedColors) {
+        Set<Integer> colors = new HashSet<>();
+        for (var neighbour: g.neighbors(node)) {
+            var c = assignedColors.get(neighbour);
+            if (c != null) {
+                colors.add(c);
+            }
+        }
+        return colors;
+    }
+
+    private Integer chooseSomeColorNotAssignedToNeighbors(Set<Integer> colors, Set<Integer> neighborColors) {
+        // Create new color set that removes the colors assigned to neighbors
+        var set = new HashSet<>(colors);
+        set.removeAll(neighborColors);
+        // pick a random color (we pick the first)
+        return set.stream().findAny().orElseThrow();
+    }
+
+    private static HashSet<Integer> subtract(Set<Integer> originalSet, Integer node) {
+        var reducedSet = new HashSet<>(originalSet);
+        reducedSet.remove(node);
+        return reducedSet;
+    }
+
+    /**
+     * Chaitin: color_graph
+     */
+    private Map<Integer, Integer> colorGraph(InterferenceGraph g, Set<Integer> nodes, Set<Integer> colors) {
+        if (nodes.size() == 0)
+            return new HashMap<>();
+        var numColors = colors.size();
+        var node = findNodeWithNeighborCountLessThan(g, nodes, numColors);
+        if (node == null)
+            return null;
+        var coloring = colorGraph(g.dup().subtract(node), subtract(nodes, node), colors);
+        if (coloring == null)
+            return null;
+        var neighbourColors = getNeighborColors(g, node, coloring);
+        var color = chooseSomeColorNotAssignedToNeighbors(colors, neighbourColors);
+        coloring.put(node, color);
+        return coloring;
+    }
+
 }

optvm/src/main/java/com/compilerprogramming/ezlang/compiler/SSATransform.java renamed to optvm/src/main/java/com/compilerprogramming/ezlang/compiler/EnterSSA.java

@@ -8,7 +8,7 @@
  * 'Practical Improvements to the Construction and Destruction
  * of Single Static Assigment Form' by Preston Briggs.
  */
-public class SSATransform {
+public class EnterSSA {
 
     CompiledFunction function;
     DominatorTree domTree;
@@ -22,12 +22,12 @@ public class SSATransform {
     int[] counters;
     VersionStack[] stacks;
 
-    public SSATransform(CompiledFunction bytecodeFunction) {
+    public EnterSSA(CompiledFunction bytecodeFunction) {
         this.function = bytecodeFunction;
         setupGlobals();
         computeDomTreeAndDominanceFrontiers();
         this.blocks = domTree.blocks;
-        findGlobalVars();
+        findNonLocalNames();
         insertPhis();
         renameVars();
         bytecodeFunction.isSSA = true;
@@ -47,7 +47,7 @@ private void setupGlobals() {
      * Compute set of registers that are live across multiple blocks
      * i.e. are not exclusively used in a single block.
      */
-    private void findGlobalVars() {
+    private void findNonLocalNames() {
         for (BasicBlock block : blocks) {
             var varKill = new HashSet<Integer>();
             for (Instruction instruction: block.instructions) {
@@ -142,7 +142,7 @@ void search(BasicBlock block) {
         }
         // Update phis in successor blocks
         for (BasicBlock s: block.successors) {
-            int j = whichPred(s,block);
+            int j = block.whichPred(s);
             for (Instruction.Phi phi: s.phis()) {
                 Register oldReg = phi.input(j);
                 phi.replaceInput(j, stacks[oldReg.nonSSAId()].top());
@@ -161,16 +161,6 @@ void search(BasicBlock block) {
         }
     }
 
-    public static int whichPred(BasicBlock s, BasicBlock block) {
-        int i = 0;
-        for (BasicBlock p: s.predecessors) {
-            if (p == block)
-                return i;
-            i++;
-        }
-        throw new IllegalStateException();
-    }
-
     private void initVersionCounters() {
         counters = new int[nonLocalNames.length];
         stacks = new VersionStack[nonLocalNames.length];

optvm/src/main/java/com/compilerprogramming/ezlang/compiler/ExitSSA.java

@@ -82,7 +82,7 @@ private void scheduleCopies(BasicBlock block, List<Integer> pushed) {
         Map<Integer, Register> map = new HashMap<>();
         BitSet usedByAnother = new BitSet(function.registerPool.numRegisters()*2);
         for (BasicBlock s: block.successors) {
-            int j = SSATransform.whichPred(s, block);
+            int j = block.whichPred(s);
             for (Instruction.Phi phi: s.phis()) {
                 Register dst = phi.value();
                 Register src = phi.input(j);   // jth operand of phi node

optvm/src/main/java/com/compilerprogramming/ezlang/compiler/InterferenceGraph.java

@@ -3,7 +3,7 @@
 import java.util.*;
 
 public class InterferenceGraph {
-    Map<Integer, Set<Integer>> edges = new HashMap<>();
+    private Map<Integer, Set<Integer>> edges = new HashMap<>();
 
     private Set<Integer> addNode(Integer node) {
         var set = edges.get(node);
@@ -24,6 +24,32 @@ public void addEdge(Integer from, Integer to) {
         set2.add(from);
     }
 
+    /**
+     * Remove a node from the interference graph
+     * deleting it from all adjacency lists
+     */
+    public InterferenceGraph subtract(Integer node) {
+        edges.remove(node);
+        for (var key : edges.keySet()) {
+            var neighbours = edges.get(key);
+            neighbours.remove(key);
+        }
+        return this;
+    }
+
+    /**
+     * Duplicate an interference graph
+     */
+    public InterferenceGraph dup() {
+        var igraph = new InterferenceGraph();
+        igraph.edges = new HashMap<>();
+        for (var key : edges.keySet()) {
+            var neighbours = edges.get(key);
+            igraph.edges.put(key, new HashSet<>(neighbours));
+        }
+        return igraph;
+    }
+
     public boolean interfere(Integer from, Integer to) {
         var set = edges.get(from);
         return set != null && set.contains(to);
@@ -50,8 +76,15 @@ public void rename(Integer source, Integer target) {
         }
     }
 
-    public Set<Integer> adjacents(Integer node) {
-        return edges.get(node);
+    /**
+     * Get neighbours of the node
+     * Chaitin: neighbors()
+     */
+    public Set<Integer> neighbors(Integer node) {
+        var adjacents = edges.get(node);
+        if (adjacents == null)
+            adjacents = Collections.emptySet();
+        return adjacents;
     }
 
     public static final class Edge {

optvm/src/main/java/com/compilerprogramming/ezlang/compiler/LiveSet.java

@@ -32,7 +32,7 @@ public void remove(List<Register> regs) {
             remove(r);
         }
     }
-    public boolean isMember(Register r) {
+    public boolean contains(Register r) {
         return get(r.id);
     }
     public LiveSet intersect(LiveSet other) {
@@ -43,7 +43,10 @@ public LiveSet union(LiveSet other) {
         or(other);
         return this;
     }
-    public LiveSet intersectNot(LiveSet other) {
+    /**
+     * Computes this - other.
+     */
+    public LiveSet subtract(LiveSet other) {
         andNot(other);
         return this;
     }