Coalesacing

Fix bug IFG adding interference to from 1 reg live range, even if success in clearing reg
Fix bug LRG merging when both masks have already been trimmed.
Missing ARM stack adjust

12 files changed, 199 insertions, 37 deletions

diff --git a/chapter21/src/main/java/com/seaofnodes/simple/codegen/Coalesce.java b/chapter21/src/main/java/com/seaofnodes/simple/codegen/Coalesce.java
new file mode 100644
index 00000000..7496ac3b
--- /dev/null
+++ b/chapter21/src/main/java/com/seaofnodes/simple/codegen/Coalesce.java
@@ -0,0 +1,119 @@
+package com.seaofnodes.simple.codegen;
+
+import com.seaofnodes.simple.Ary;
+import com.seaofnodes.simple.Utils;
+import com.seaofnodes.simple.node.*;
+import java.util.BitSet;
+import java.util.IdentityHashMap;
+
+// Coalesce copies
+abstract public class Coalesce {
+
+    public static boolean coalesce( int round, RegAlloc alloc ) {
+        // Convert the 2-D array of bits (a 1-D array of BitSets) into an
+        // adjacency matrix.
+        int maxlrg = alloc._LRGS.length;
+        for( int i=1; i<maxlrg; i++ ) {
+            BitSet ifg = IFG.IFG.atX(i);
+            if( ifg != null ) {
+                LRG lrg0 = alloc._LRGS[i];
+                for( int lrg = ifg.nextSetBit(0); lrg>=0; lrg=ifg.nextSetBit(lrg+1) ) {
+                    LRG lrg1 = alloc._LRGS[lrg];
+                    lrg0.addNeighbor(lrg1);
+                    lrg1.addNeighbor(lrg0);
+                }
+            }
+        }
+
+        // Walk all the splits, looking for coalesce chances
+        boolean progress = false;
+        for( CFGNode bb : alloc._code._cfg ) {
+            outer:
+            for( int j=0; j<bb.nOuts(); j++ ) {
+                Node n = bb.out(j);
+                if( !(n instanceof SplitNode) ) continue;
+                LRG v1 = alloc.lrg(n);
+                LRG v2 = alloc.lrg(n.in(1));
+                if( v1 != v2 ) {
+                    LRG ov1 = v1, ov2 = v2;
+                    // Get the smaller neighbor count in v1
+                    if( (v1._adj==null ? 0 : v1._adj._len) > (v2._adj==null ? 0 : v2._adj._len) )
+                        { v1 = v2; v2 = alloc.lrg(n); }
+                    int v2len = v2._adj==null ? 0 : v2._adj._len;
+                    // See that they do not conflict (coalescing would make a self-conflict)
+                    if( v1._adj!=null && v2._adj!=null && v1._mask.overlap(v2._mask) ) {
+                        // Walk the smaller neighbors, and add to the larger.
+                        // Check for direct conflicts along the way.
+                        for( LRG v3 : v1._adj ) {
+                            if( v3==v2 ) {
+                                // Direct conflict, skip this LRG pair
+                                v2._adj.setLen(v2len);
+                                continue outer;
+                            }
+                            // Add missing LRGs from V1 to V2
+                            if( v2._adj.find(v3)== -1 )
+                                v2._adj.push(v3);
+                        }
+                    }
+                    // Most constrained mask
+                    RegMask mask = v1._mask==v2._mask ? v1._mask : v1._mask.copy().and(v2._mask);
+                    // Check for capacity
+                    if( (v2._adj==null ? 0 : v2._adj._len) >= mask.size() ) {
+                        // Fails capacity, will not be trivial colorable
+                        v2._adj.setLen(v2len);
+                        continue;
+                    }
+
+                    // Union lrgs.
+                    progress = true;
+
+                    // Since `n` is dying, remove from the LRGS BEFORE doing
+                    // the union; this will let the union code keep the
+                    // matching bits from the other side, preserving some
+                    // info for biased coloring
+                    if( ov1. _machDef==n ) ov1. _machDef = null;
+                    if( ov1._splitDef==n ) ov1._splitDef = null;
+                    if( ov2. _machUse==n ) ov2. _machUse = null;
+                    if( ov2._splitUse==n ) ov2._splitUse = null;
+                    // Keep larger adjacency list.
+                    Ary<LRG> larger = v2._adj;
+                    v2._adj = null;
+
+                    // Union lrgs.  Swaps again, based on small lrg
+                    LRG vWin = v1.union(v2);
+                    vWin._adj = larger;
+                    LRG vLose = v1.leader() ? v2 : v1;
+
+                    // LRGs below v2len were always in v2, but maybe not in v1.
+                    // LRGs above v2len are in v1 but were not in v2.
+
+                    // Walk v2's neighbors, replacing vLose with vWin.
+                    if( larger != null )
+                        for( LRG vn : larger ) {
+                            int idx = vn._adj.find(vLose);
+                            if( idx != -1 )
+                                vn._adj.del(idx);
+                            if( vn._adj.find(vWin) == -1 )
+                                vn._adj.add(vWin);
+                        }
+                } else {
+
+                    // Since `n` is dying, remove from the LRGS
+                    if( v1. _machDef==n ) v1. _machDef=null;
+                    if( v1. _machUse==n ) v1. _machUse=null;
+                    if( v1._splitDef==n ) v1._splitDef=null;
+                    if( v1._splitUse==n ) v1._splitUse=null;
+                }
+
+                // Remove junk split
+                n.removeSplit();
+                j--;
+            }
+        }
+
+        if( progress )
+            alloc.unify();
+
+        return true;
+    }
+}
diff --git a/chapter21/src/main/java/com/seaofnodes/simple/codegen/IFG.java b/chapter21/src/main/java/com/seaofnodes/simple/codegen/IFG.java
index df5976a4..e85bce5d 100644
--- a/chapter21/src/main/java/com/seaofnodes/simple/codegen/IFG.java
+++ b/chapter21/src/main/java/com/seaofnodes/simple/codegen/IFG.java
@@ -19,7 +19,7 @@ abstract public class IFG {
     static final IdentityHashMap<LRG,Node> TMP = new IdentityHashMap<>();
 
     // Inteference Graph: Array of Bitsets
-    private static final Ary<BitSet> IFG = new Ary<>(BitSet.class);
+    static final Ary<BitSet> IFG = new Ary<>(BitSet.class);
     static void resetIFG() {
         for( BitSet bs : IFG )
             if( bs!=null ) bs.clear();
@@ -148,9 +148,10 @@ abstract public class IFG {
                     // last tlrg register at some point, either tlrg or lrg
                     // must fail.  If *n* (a subset of lrg) needs the single
                     // last tlrg register then only tlrg must fail.
-                    if( ((MachNode)n).outregmap().size1() && !tlrg.clr(lrg._mask.firstReg()) )
-                        alloc.fail(tlrg);
-                    else addIFG(lrg,tlrg); // Add interference
+                    if( ((MachNode)n).outregmap().size1() ) {
+                        if( !tlrg.clr(lrg._mask.firstReg()) ) // Clear bit, no interference
+                            alloc.fail(tlrg);                 // Clearing drives mask to empty
+                    } else addIFG(lrg,tlrg); // Add interference
             }
         }
 
@@ -281,22 +282,10 @@ abstract public class IFG {
     // If there's no spare color we'll have to spill this at-risk live range.
 
     public static boolean color(int round, RegAlloc alloc) {
-
-        // Convert the 2-D array of bits (a 1-D array of BitSets) into an
-        // adjacency matrix.
         int maxlrg = alloc._LRGS.length, nlrgs=0;
-        for( int i=1; i<maxlrg; i++ ) {
-            if( alloc._LRGS[i] != null ) nlrgs++;
-            BitSet ifg = IFG.atX(i);
-            if( ifg != null ) {
-                LRG lrg0 = alloc._LRGS[i];
-                for( int lrg = ifg.nextSetBit(0); lrg>=0; lrg=ifg.nextSetBit(lrg+1) ) {
-                    LRG lrg1 = alloc._LRGS[lrg];
-                    lrg0.addNeighbor(lrg1);
-                    lrg1.addNeighbor(lrg0);
-                }
-            }
-        }
+        for( int i=1; i<maxlrg; i++ )
+            if( alloc._LRGS[i] != null )
+                nlrgs++;
 
         // Simplify
 
@@ -429,7 +418,8 @@ abstract public class IFG {
     // Picking a live range that is very close to coloring might allow it to
     // color despite being risky.
     private static int pickRiskyScore( LRG lrg ) {
-        // Pick single-def clonables that are not right next to their single-use
+        // Pick single-def clonables that are not right next to their single-use.
+        // Failing to color these will clone them closer to their uses.
         if( !lrg._multiDef && lrg._machDef.isClone() ) {
             Node def = ((Node)lrg._machDef);
             Node use = ((Node)lrg._machUse);
diff --git a/chapter21/src/main/java/com/seaofnodes/simple/codegen/LRG.java b/chapter21/src/main/java/com/seaofnodes/simple/codegen/LRG.java
index d44dffb8..db5a3242 100644
--- a/chapter21/src/main/java/com/seaofnodes/simple/codegen/LRG.java
+++ b/chapter21/src/main/java/com/seaofnodes/simple/codegen/LRG.java
@@ -2,6 +2,7 @@ package com.seaofnodes.simple.codegen;
 
 import com.seaofnodes.simple.Ary;
 import com.seaofnodes.simple.SB;
+import com.seaofnodes.simple.Utils;
 import com.seaofnodes.simple.node.*;
 
 import java.util.IdentityHashMap;
@@ -151,7 +152,10 @@ public class LRG {
         _splitUse = deepSplit(_splitUse,lrg._splitUse);
 
         // Fold together masks
-        _mask = _mask.and(lrg._mask);
+        RegMask mask = _mask.and(lrg._mask);
+        if( mask==null )
+            mask = _mask.copy().and(lrg._mask);
+        _mask = mask;
         return this;
     }
 
diff --git a/chapter21/src/main/java/com/seaofnodes/simple/codegen/RegAlloc.java b/chapter21/src/main/java/com/seaofnodes/simple/codegen/RegAlloc.java
index ebcf614d..32fca074 100644
--- a/chapter21/src/main/java/com/seaofnodes/simple/codegen/RegAlloc.java
+++ b/chapter21/src/main/java/com/seaofnodes/simple/codegen/RegAlloc.java
@@ -173,6 +173,8 @@ public class RegAlloc {
             BuildLRG.run(round,this) && // if no hard register conflicts
             // Build Interference Graph
             IFG.build(round,this) &&    // If no self conflicts or uncolorable
+            // Conservative coalesce copies
+            Coalesce.coalesce(round,this) &&
             // Color attempt
             IFG.color(round,this);      // If colorable
     }
diff --git a/chapter21/src/main/java/com/seaofnodes/simple/node/cpus/arm/FunARM.java b/chapter21/src/main/java/com/seaofnodes/simple/node/cpus/arm/FunARM.java
index 2e8bf6ad..eaea0ede 100644
--- a/chapter21/src/main/java/com/seaofnodes/simple/node/cpus/arm/FunARM.java
+++ b/chapter21/src/main/java/com/seaofnodes/simple/node/cpus/arm/FunARM.java
@@ -1,13 +1,36 @@
 package com.seaofnodes.simple.node.cpus.arm;
 
 import com.seaofnodes.simple.codegen.*;
+import com.seaofnodes.simple.SB;
 import com.seaofnodes.simple.node.FunNode;
 import com.seaofnodes.simple.node.MachNode;
 
 public class FunARM  extends FunNode implements MachNode {
     FunARM(FunNode fun) { super(fun); }
+    @Override public String op() { return "addi"; }
     @Override public RegMask regmap(int i) { return null; }
     @Override public RegMask outregmap() { return null; }
-    @Override public void encoding( Encoding enc ) { }
     @Override public void postSelect(CodeGen code) { code.link(this);  }
+
+    // Actual stack layout is up to each CPU.
+    // X86, with too many args & spills:
+    // | CALLER |
+    // |  argN  | // slot 1, required by callER
+    // +--------+
+    // |  RPC   | // slot 0, required by callER
+    // | callee | // slot 3, callEE
+    // | callee | // slot 2, callEE
+    // |  PAD16 |
+    // +--------+
+    @Override public void encoding( Encoding enc ) {
+        // Stack frame size: _maxSlot - max(arg), padded to 16.
+        _maxArgSlot = arm.maxSlot(enc._fun.sig());
+        _frameAdjust = (short) (_maxSlot+1 - _maxArgSlot);
+        if( _frameAdjust == 0 ) return; // Skip if no frame adjust
+        enc.add4(arm.imm_inst(0b1001000100, (_frameAdjust*8)&0xFFF, arm.RSP, arm.RSP));
+    }
+
+    @Override public void asm(CodeGen code, SB sb) {
+        sb.p("rsp += #").p(_frameAdjust*8);
+    }
 }
diff --git a/chapter21/src/main/java/com/seaofnodes/simple/node/cpus/arm/RetARM.java b/chapter21/src/main/java/com/seaofnodes/simple/node/cpus/arm/RetARM.java
index 290421aa..b51b4879 100644
--- a/chapter21/src/main/java/com/seaofnodes/simple/node/cpus/arm/RetARM.java
+++ b/chapter21/src/main/java/com/seaofnodes/simple/node/cpus/arm/RetARM.java
@@ -6,7 +6,9 @@ import com.seaofnodes.simple.node.*;
 
 public class RetARM extends ReturnNode implements MachNode {
     RetARM(ReturnNode ret, FunNode fun) { super(ret, fun); fun.setRet(this); }
-    @Override public String op() { return "ret"; }
+    @Override public String op() {
+        return ((FunARM)fun())._frameAdjust > 0 ? "addi" : "ret";
+    }
     // Correct Nodes outside the normal edges
     @Override public void postSelect(CodeGen code) {
         FunNode fun = (FunNode)rpc().in(0);
@@ -16,9 +18,17 @@ public class RetARM extends ReturnNode implements MachNode {
     @Override public RegMask regmap(int i) { return arm.retMask(_fun.sig(),i); }
     @Override public RegMask outregmap() { return null; }
     // RET
-    @Override public void encoding( Encoding enc ) { enc.add4(arm.ret(0b1101011001011111000000)); }
+    @Override public void encoding( Encoding enc ) {
+        int frameAdjust = ((FunARM)fun())._frameAdjust;
+        if( frameAdjust > 0 )
+            enc.add4(arm.imm_inst(0b1001000100, (frameAdjust*-8)&0xFFF, arm.RSP, arm.RSP));
+        enc.add4(arm.ret(0b1101011001011111000000));
+    }
 
     @Override public void asm(CodeGen code, SB sb) {
+        int frameAdjust = ((FunARM)fun())._frameAdjust;
+        if( frameAdjust>0 )
+            sb.p("rsp += #").p(frameAdjust*-8).p("\nret");
         // Post code-gen, just print the "ret"
         if( code._phase.ordinal() <= CodeGen.Phase.RegAlloc.ordinal() )
             // Prints return reg (either X0 or D0), RPC (always R30) and
diff --git a/chapter21/src/main/java/com/seaofnodes/simple/node/cpus/arm/arm.java b/chapter21/src/main/java/com/seaofnodes/simple/node/cpus/arm/arm.java
index 4c3b3fc7..9ac5a718 100644
--- a/chapter21/src/main/java/com/seaofnodes/simple/node/cpus/arm/arm.java
+++ b/chapter21/src/main/java/com/seaofnodes/simple/node/cpus/arm/arm.java
@@ -418,6 +418,20 @@ public class arm extends Machine {
         }
         throw Utils.TODO(); // Pass on stack slot
     }
+    // Return the max stack slot used by this signature, or 0
+    static short maxSlot( TypeFunPtr tfp ) {
+        // Count floats in signature up to index
+        int fcnt=0;
+        for( int i=0; i<tfp.nargs(); i++ )
+            if( tfp.arg(i) instanceof TypeFloat )
+                fcnt++;
+        if( fcnt >= XMMS.length )
+            throw Utils.TODO();
+        RegMask[] cargs = CALLINMASK;
+        if( tfp.nargs()-fcnt >= cargs.length )
+            throw Utils.TODO();
+        return 0;               // No stack args
+    }
 
     static final long CALLEE_SAVE =
             1L<<X19 |
diff --git a/chapter21/src/main/java/com/seaofnodes/simple/node/cpus/riscv/FunRISC.java b/chapter21/src/main/java/com/seaofnodes/simple/node/cpus/riscv/FunRISC.java
index 3bd2570f..2831360d 100644
--- a/chapter21/src/main/java/com/seaofnodes/simple/node/cpus/riscv/FunRISC.java
+++ b/chapter21/src/main/java/com/seaofnodes/simple/node/cpus/riscv/FunRISC.java
@@ -25,7 +25,7 @@ public class FunRISC extends FunNode implements MachNode {
     @Override public void encoding( Encoding enc ) {
         // Stack frame size: _maxSlot - max(arg), padded to 16.
         _maxArgSlot = riscv.maxSlot(enc._fun.sig());
-        _frameAdjust = (short) (_maxSlot - _maxArgSlot);
+        _frameAdjust = (short) (_maxSlot+1 - _maxArgSlot);
         if( _frameAdjust == 0 ) return; // Skip if no frame adjust
         enc.add4(riscv.i_type(riscv.I_TYPE, riscv.SP, 0, riscv.SP, (_frameAdjust*8) & 0xFFF));
     }
diff --git a/chapter21/src/main/java/com/seaofnodes/simple/node/cpus/x86_64_v2/JmpX86.java b/chapter21/src/main/java/com/seaofnodes/simple/node/cpus/x86_64_v2/JmpX86.java
index 10b111e4..a6e280d3 100644
--- a/chapter21/src/main/java/com/seaofnodes/simple/node/cpus/x86_64_v2/JmpX86.java
+++ b/chapter21/src/main/java/com/seaofnodes/simple/node/cpus/x86_64_v2/JmpX86.java
@@ -60,7 +60,7 @@ public class JmpX86 extends IfNode implements MachNode, RIPRelSize {
         String src = code.reg(in(1));
         if( src!="flags" ) sb.p(src).p(" ");
         CFGNode prj = cproj(0); // 0 is True is jump target
-        while( prj.nOuts() == 1 )
+        while( prj.nOuts() == 1 && !(prj instanceof ReturnNode) )
             prj = prj.uctrl(); // Skip empty blocks
         sb.p(label(prj));
     }
diff --git a/chapter21/src/test/java/com/seaofnodes/simple/BrainFuckTest.java b/chapter21/src/test/java/com/seaofnodes/simple/BrainFuckTest.java
index 28111697..eb88ecc1 100644
--- a/chapter21/src/test/java/com/seaofnodes/simple/BrainFuckTest.java
+++ b/chapter21/src/test/java/com/seaofnodes/simple/BrainFuckTest.java
@@ -65,7 +65,7 @@ return output;
 """;
         Chapter21Test.testCPU(src,"x86_64_v2", "SystemV",40,null);
         Chapter21Test.testCPU(src,"riscv"    , "SystemV",30,null);
-        Chapter21Test.testCPU(src,"arm"      , "SystemV",30,null);
+        Chapter21Test.testCPU(src,"arm"      , "SystemV",26,null);
         //assertEquals("Hello World!\n", Eval2.eval(code, 0, 10000));
     }
 }
diff --git a/chapter21/src/test/java/com/seaofnodes/simple/Chapter20Test.java b/chapter21/src/test/java/com/seaofnodes/simple/Chapter20Test.java
index 1a08d6ee..b83db0d1 100644
--- a/chapter21/src/test/java/com/seaofnodes/simple/Chapter20Test.java
+++ b/chapter21/src/test/java/com/seaofnodes/simple/Chapter20Test.java
@@ -40,8 +40,8 @@ public class Chapter20Test {
 
     @Test public void testAlloc0() {
         testCPU("return new u8[arg];","x86_64_v2", "SystemV",3,"return [u8];");
-        testCPU("return new u8[arg];","riscv"    , "SystemV",7,"return [u8];");
-        testCPU("return new u8[arg];","arm"      , "SystemV",7,"return [u8];");
+        testCPU("return new u8[arg];","riscv"    , "SystemV",5,"return [u8];");
+        testCPU("return new u8[arg];","arm"      , "SystemV",5,"return [u8];");
     }
 
     @Test public void testBasic1() {
@@ -66,9 +66,9 @@ val sqrt = { int x ->
 };
 return sqrt(arg) + sqrt(arg+2);
 """;
-        testCPU(src,"x86_64_v2", "SystemV",26,null);
+        testCPU(src,"x86_64_v2", "SystemV",23,null);
         testCPU(src,"riscv"    , "SystemV",19,null);
-        testCPU(src,"arm"      , "SystemV",17,null);
+        testCPU(src,"arm"      , "SystemV",19,null);
     }
 
     @Test
@@ -85,7 +85,7 @@ return sqrt(arg) + sqrt(arg+2);
         String src = "int[] !xs = new int[3]; xs[arg]=1; return xs[arg&1];";
         testCPU(src,"x86_64_v2", "SystemV",3,"return .[];");
         testCPU(src,"riscv"    , "SystemV",6,"return .[];");
-        testCPU(src,"arm"      , "SystemV",6,"return .[];");
+        testCPU(src,"arm"      , "SystemV",4,"return .[];");
     }
 
 
@@ -103,7 +103,7 @@ return sqrt(arg) + sqrt(arg+2);
     public void testString() throws IOException {
         String src = Files.readString(Path.of("src/test/java/com/seaofnodes/simple/progs/stringHash.smp"));
         testCPU(src,"x86_64_v2", "SystemV",18,null);
-        testCPU(src,"riscv"    , "SystemV",17,null);
+        testCPU(src,"riscv"    , "SystemV",20,null);
         testCPU(src,"arm"      , "SystemV",16,null);
     }
 
diff --git a/chapter21/src/test/java/com/seaofnodes/simple/Chapter21Test.java b/chapter21/src/test/java/com/seaofnodes/simple/Chapter21Test.java
index cb55c033..e3b69593 100644
--- a/chapter21/src/test/java/com/seaofnodes/simple/Chapter21Test.java
+++ b/chapter21/src/test/java/com/seaofnodes/simple/Chapter21Test.java
@@ -48,9 +48,9 @@ public class Chapter21Test {
     public void testNewtonFloat() throws IOException {
         String src = Files.readString(Path.of("src/test/java/com/seaofnodes/simple/progs/newtonFloat.smp"))
             + "flt farg = arg; return sqrt(farg) + sqrt(farg+2.0);";
-        testCPU(src,"x86_64_v2", "SystemV",25,null);
-        testCPU(src,"riscv"    , "SystemV",17,null);
-        testCPU(src,"arm"      , "SystemV",18,null);
+        testCPU(src,"x86_64_v2", "SystemV",23,null);
+        testCPU(src,"riscv"    , "SystemV",18,null);
+        testCPU(src,"arm"      , "SystemV",19,null);
     }
 
     @Test
@@ -65,7 +65,7 @@ public class Chapter21Test {
     public void testArray1() throws IOException {
         String src = Files.readString(Path.of("src/test/java/com/seaofnodes/simple/progs/array1.smp"));
         testCPU(src,"x86_64_v2", "SystemV",7,"return .[];");
-        testCPU(src,"riscv"    , "SystemV",9,"return (add,.[],(mul,.[],1000));");
+        testCPU(src,"riscv"    , "SystemV",7,"return (add,.[],(mul,.[],1000));");
         testCPU(src,"arm"      , "SystemV",5,"return (add,.[],(mul,.[],1000));");
     }