Better block layout

Keeps loops self-contained, handles more variants of infinite loops.
Fold region-to-regions together, in an effort to have fewer junk jumps

8 files changed, 143 insertions, 84 deletions

diff --git a/chapter21/src/main/java/com/seaofnodes/simple/codegen/CodeGen.java b/chapter21/src/main/java/com/seaofnodes/simple/codegen/CodeGen.java
index 2e1a4709..91553f38 100644
--- a/chapter21/src/main/java/com/seaofnodes/simple/codegen/CodeGen.java
+++ b/chapter21/src/main/java/com/seaofnodes/simple/codegen/CodeGen.java
@@ -318,7 +318,7 @@ public class CodeGen {
 
 
     // ---------------------------
-    // Control Flow Graph in RPO order.
+    // Control Flow Graph in Reverse Post Order.
     public int _tGCM;
     public Ary<CFGNode> _cfg = new Ary<>(CFGNode.class);
 
diff --git a/chapter21/src/main/java/com/seaofnodes/simple/codegen/Encoding.java b/chapter21/src/main/java/com/seaofnodes/simple/codegen/Encoding.java
index 5ff2836e..6d6248e4 100644
--- a/chapter21/src/main/java/com/seaofnodes/simple/codegen/Encoding.java
+++ b/chapter21/src/main/java/com/seaofnodes/simple/codegen/Encoding.java
@@ -1,8 +1,10 @@
 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 com.seaofnodes.simple.print.IRPrinter;
 import com.seaofnodes.simple.type.*;
 import java.io.ByteArrayOutputStream;
 import java.util.*;
@@ -152,11 +154,13 @@ public class Encoding {
         _bigCons.put(relo,new Relo(relo,t,(byte)off,(byte)elf));
     }
 
+    // --------------------------------------------------
     void encode() {
         // Basic block layout: negate branches to keep blocks in-order; insert
         // unconditional jumps.  Attempt to keep backwards branches taken,
         // forwards not-taken (this is the default prediction on most
-        // hardware).  Layout is still RPO but with more restrictions.
+        // hardware).  Layout is still Reverse Post Order but with more
+        // restrictions.
         basicBlockLayout();
 
         // Write encoding bits in order into a big byte array.
@@ -171,17 +175,23 @@ public class Encoding {
         patchLocalRelocations();
     }
 
+    // --------------------------------------------------
     // Basic block layout: negate branches to keep blocks in-order; insert
-    // unconditional jumps.  Attempt to keep backwards branches taken,
-    // forwards not-taken (this is the default prediction on most
-    // hardware).  Layout is still RPO but with more restrictions.
+    // unconditional jumps.  Attempt to keep backwards branches taken, forwards
+    // not-taken (this is the default prediction on most hardware).  Layout is
+    // still Reverse Post Order but with more restrictions.
     private void basicBlockLayout() {
+        IdentityHashMap<LoopNode,Ary<CFGNode>> rpos = new IdentityHashMap<>();
         Ary<CFGNode> rpo = new Ary<>(CFGNode.class);
+        rpos.put(_code._start.loop(),rpo);
         BitSet visit = _code.visit();
+        //IdentityHashMap<CFGNode,LoopNode> looptail = new IdentityHashMap<>();
+        rpo.add(_code._stop);
         for( Node n : _code._start._outputs )
             if( n instanceof FunNode fun ) {
                 int x = rpo._len;
-                _rpo_cfg(fun, visit, rpo);
+                //_rpo_cfg2(fun, visit, rpo, looptail);
+                _rpo_cfg(fun, visit, rpos );
                 assert rpo.at(x) instanceof ReturnNode;
             }
         rpo.add(_code._start);
@@ -193,64 +203,45 @@ public class Encoding {
         _code._cfg = rpo;       // Save the new ordering
     }
 
-    // Basic block layout.  Now that RegAlloc is finished, no more spill code
-    // will appear.  We can change our BB layout from RPO to something that
-    // minimizes actual branches, takes advantage of fall-through edges, and
-    // tries to help simple branch predictions: back branches are predicted
-    // taken, forward not-taken.
-    private void _rpo_cfg(CFGNode bb, BitSet visit, Ary<CFGNode> rpo) {
-        if( visit.get(bb._nid) ) return;
+
+    private void _rpo_cfg(CFGNode bb, BitSet visit, IdentityHashMap<LoopNode,Ary<CFGNode>> rpos ) {
+        if( bb==null || visit.get(bb._nid) ) return;
         visit.set(bb._nid);
-        if( bb.nOuts()==0 ) return; // StopNode
-        if( !(bb instanceof IfNode iff) ) {
-            CFGNode next = bb instanceof ReturnNode ? (CFGNode)bb.out(bb.nOuts()-1) : bb.uctrl();
-            // If the *next* BB has already been visited, we may need an
-            // unconditional jump here
-            if( visit.get(next._nid) && !(next instanceof StopNode) ) {
-                boolean needJump = next instanceof LoopNode
-                    // If backwards to a loop, and the block has statements,
-                    // will need a jump.  Empty blocks can just backwards branch.
-                    ? (bb.nOuts()>1)
-                    // Forwards jump.  If all the blocks, in RPO order, to our
-                    // target are empty, we will fall in and not need a jump.
-                    : !isEmptyBackwardsScan(rpo,next);
-                if( needJump ) {
-                    CFGNode jmp = _code._mach.jump();
-                    jmp.setDefX(0,bb);
-                    next.setDef(next._inputs.find(bb),jmp);
-                    rpo.add(jmp);
-                }
-            }
-            _rpo_cfg(next,visit,rpo);
-        } else {
-            boolean invert = false;
+        CFGNode next = bb.uctrl();
+
+        // Loops run an inner "rpo_cfg", then append the entire loop body in
+        // place.  This keeps loop bodies completely contained, although if
+        // Some Day Later we have real profile data we ought to arrange the
+        // branch orderings based on frequency.
+        if( bb instanceof LoopNode loop ) {
+            Ary<CFGNode> body = new Ary<>(CFGNode.class); // Private RPO for the loop
+            rpos.put(loop.loop(),body);                   // Find it via loop tree
+            _rpo_cfg(next,visit,rpos);                    // RPO the loop
+            body.add(loop);                               // Include loop last (first in RPO)
+            Ary<CFGNode> outer = rpos.get(loop.cfg(1).loop());
+            outer.addAll(body); // Append to original CFG
+            return;
+        }
+
+        // IfNodes visit 2 sides, and may choose to reorder them in the RPO
+        if( bb instanceof IfNode iff ) {
             // Pick out the T/F projections
             CProjNode t = iff.cproj(0);
             CProjNode f = iff.cproj(1);
-            int tld = t.loopDepth(), fld = f.loopDepth(), bld = bb.loopDepth();
-            // Decide entering or exiting a loop
-            if( tld==bld ) {
-                // if T is empty, keep
-                if( t.nOuts()>1 &&
-                    // Else swap so T is forward and exits, while F falls into next loop block
-                    ((fld<bld && t.out(0)!=t.loop()) ||
-                     // Jump to an empty block and fall into a busy block
-                     (fld==bld && f.nOuts()==1)) ) {
-                    invert = true;
-                } // Else nothing
-            } else if( tld > bld ) { // True enters a deeper loop
-                throw Utils.TODO();
-            } else if( fld == bld && f.out(0) instanceof LoopNode ) { // Else True exits a loop
-                // if false is the loop backedge, make sure its true/taken
-                invert = true;
-            } // always forward and good
-            // Invert test and Proj fields
-            if( invert ) {
+            int tld = t.loopDepth(), fld = f.loopDepth(), bld = iff.loopDepth();
+            // Invert the branch or not
+            if( tld >= bld &&   // true to exit a loop
+                // false to exit a loop, invert to true
+                (fld < bld ||
+                 // Else not exiting (staying or entering deeper).
+                 // Fall/false into a full block, Jump/true to an empty block.
+                 ( f.nOuts()==1 && (t.nOuts()>1 ||
+                                    // Everything else equal, use pre-order
+                                    t._pre > f._pre) )) ) {
                 iff.negate();
                 t.invert();
                 f.invert();
                 CProjNode tmp=f; f=t; t=tmp; // Swap t/f
-                int d=tld; tld=fld; fld=d;   // Swap depth
             }
 
             // Whichever side is visited first becomes last in the RPO.  With
@@ -258,22 +249,37 @@ public class Encoding {
             // when the False RPO visit returns, the IF is immediately next.
             // When the RPO is reversed, the fall-through path will always be
             // following the IF.
+            _rpo_cfg(t,visit,rpos); // True side first
+            next = f;               // False side last
+        }
 
-            // If loops are involved, attempt to keep them in a line.  Visit
-            // the exits first, so they follow the loop body when the order gets
-            // reversed.
-            if( fld < tld || (fld==tld && f.nOuts()==1) ) {
-                _rpo_cfg(f,visit,rpo);
-                _rpo_cfg(t,visit,rpo);
-            } else {
-                _rpo_cfg(t,visit,rpo);
-                _rpo_cfg(f,visit,rpo);
+        // If the *next* BB has already been visited, and we are not already a
+        // jump, we may need an unconditional forwards jump here
+        Ary<CFGNode> rpo = rpos.get(bb.loop());
+        if( next!=null && visit.get(next._nid) && !(bb instanceof IfNode) ) {
+            boolean needJump = next instanceof LoopNode
+                // Empty blocks from an IF will invert the IF and backwards
+                // branch to the loop head.  If not empty, or not an IF
+                // will need a jump.
+                ? (bb.nOuts()>1 || !(bb.cfg0() instanceof IfNode) )
+                // Forwards jump.  If all the blocks, in RPO order, to our
+                // target are empty, we will fall in and not need a jump.
+                : !isEmptyBackwardsScan(rpo,next);
+            if( needJump ) {
+                CFGNode jmp = _code._mach.jump();
+                jmp._ltree = bb._ltree;
+                jmp.setDefX(0,bb);
+                next.setDef(next._inputs.find(bb),jmp);
+                rpo.add(jmp);
             }
         }
-        assert rpo._len>0 || bb instanceof ReturnNode;
+
+        _rpo_cfg(next,visit,rpos);
         rpo.add(bb);
     }
 
+    // Is the CFG from "next" to the end empty?  This means jumping to "next"
+    // will naturally fall into the end.
     private static boolean isEmptyBackwardsScan(Ary<CFGNode> rpo, CFGNode next) {
         for( int i=rpo._len-1; rpo.at(i)!=next; i-- )
             if( rpo.at(i).nOuts()!=1 )
@@ -281,6 +287,7 @@ public class Encoding {
         return true;
     }
 
+    // --------------------------------------------------
     // Write encoding bits in order into a big byte array.
     // Record opcode start and length.
     public FunNode _fun;        // Currently encoding function
@@ -308,6 +315,7 @@ public class Encoding {
         padN(16,_bits);
     }
 
+    // --------------------------------------------------
     // Short-form RIP-relative support: replace short encodings with long
     // encodings and expand the code, changing all the offsets.
     private void compactShortForm() {
@@ -340,7 +348,7 @@ public class Encoding {
                     if( n instanceof MachNode && !(n instanceof CFGNode) )
                         _opStart[n._nid] += slide;
                 if( bb instanceof RIPRelSize riprel ) {
-                    CFGNode target = ((CFGNode)bb.out(0)).uctrlSkipEmpty();
+                    CFGNode target = (bb instanceof IfNode iff ? iff.cproj(0) : (CFGNode)bb.out(0)).uctrlSkipEmpty();
                     // Delta is from opStart to opStart.  X86 at least counts
                     // the delta from the opEnd, but we don't have the end until
                     // we decide the size - so the encSize has to deal
@@ -375,6 +383,19 @@ public class Encoding {
     }
 
 
+    // --------------------------------------------------
+    // Patch local encodings now
+    void patchLocalRelocations() {
+        // Walk the local code-address relocations
+        for( Node src : _internals.keySet() ) {
+            int start  = _opStart[src._nid];
+            Node dst =  _internals.get(src);
+            int target = _opStart[dst._nid];
+            ((RIPRelSize)src).patch(this, start, _opLen[src._nid], target - start);
+        }
+    }
+
+    // --------------------------------------------------
     // Write the constant pool into the BAOS and optionally patch locally
     void writeConstantPool( BAOS bits, boolean patch ) {
         padN(16,bits);
@@ -409,17 +430,6 @@ public class Encoding {
     }
 
 
-    // Patch local encodings now
-    void patchLocalRelocations() {
-        // Walk the local code-address relocations
-        for( Node src : _internals.keySet() ) {
-            int start  = _opStart[src._nid];
-            Node dst =  _internals.get(src);
-            int target = _opStart[dst._nid];
-            ((RIPRelSize)src).patch(this, start, _opLen[src._nid], target - start);
-        }
-    }
-
     // A series of libc/external calls that Simple can link against in a JIT.
     // Since no runtime in the JVM process, using magic numbers for the CPU
     // emulators to pick up on.
@@ -436,4 +446,12 @@ public class Encoding {
             ((RIPRelSize)src).patch(this, start, _opLen[src._nid], target - start);
         }
     }
+
+    String printCFG() {
+        if( _code._cfg==null ) return "no CFG";
+        SB sb = new SB();
+        for( CFGNode cfg : _code._cfg )
+            IRPrinter.printLine(cfg,sb);
+        return sb.toString();
+    }
 }
diff --git a/chapter21/src/main/java/com/seaofnodes/simple/node/CFGNode.java b/chapter21/src/main/java/com/seaofnodes/simple/node/CFGNode.java
index 0aeee5f3..fd17092f 100644
--- a/chapter21/src/main/java/com/seaofnodes/simple/node/CFGNode.java
+++ b/chapter21/src/main/java/com/seaofnodes/simple/node/CFGNode.java
@@ -106,7 +106,7 @@ public abstract class CFGNode extends Node {
     public LoopNode loop() { return _ltree._head; }
     public int loopDepth() { return _ltree==null ? 0 : _ltree.depth(); }
 
-    LoopTree _ltree;
+    public LoopTree _ltree;
     public int _pre;            // Pre-order numbers for loop tree finding
     static class LoopTree {
         LoopTree _par;
diff --git a/chapter21/src/main/java/com/seaofnodes/simple/node/LoopNode.java b/chapter21/src/main/java/com/seaofnodes/simple/node/LoopNode.java
index cc11d811..1ba182b9 100644
--- a/chapter21/src/main/java/com/seaofnodes/simple/node/LoopNode.java
+++ b/chapter21/src/main/java/com/seaofnodes/simple/node/LoopNode.java
@@ -50,9 +50,9 @@ public class LoopNode extends RegionNode {
         CProjNode t = new CProjNode(iff,0,"True" ).init();
         CProjNode f = new CProjNode(iff,1,"False").init();
         setDef(2,t);            // True continues loop, False (never) exits loop
-        iff._ltree = t._ltree = _ltree;
         ReturnNode ret = fun.ret();
-        f._ltree = ret._ltree;
+        iff._ltree = t._ltree = _ltree;
+        ret._ltree = f._ltree = stop._ltree;
 
         // Now fold control into the exit.  Might have 1 valid exit, or an
         // XCtrl or a bunch of prior NeverNode exits.
diff --git a/chapter21/src/main/java/com/seaofnodes/simple/node/RegionNode.java b/chapter21/src/main/java/com/seaofnodes/simple/node/RegionNode.java
index 2da402b1..583a4088 100644
--- a/chapter21/src/main/java/com/seaofnodes/simple/node/RegionNode.java
+++ b/chapter21/src/main/java/com/seaofnodes/simple/node/RegionNode.java
@@ -1,6 +1,7 @@
 package com.seaofnodes.simple.node;
 
 import com.seaofnodes.simple.*;
+import com.seaofnodes.simple.Parser;
 import com.seaofnodes.simple.codegen.CodeGen;
 import com.seaofnodes.simple.type.Type;
 import java.util.*;
@@ -56,6 +57,29 @@ public class RegionNode extends CFGNode {
             return delDef(2);
         }
 
+        // Flatten stack regions (no loops involved)
+        if( getClass() == RegionNode.class ) {
+            for( int i=1; i<nIns(); i++ )
+                if( cfg(i) instanceof RegionNode region && region.getClass() == RegionNode.class && region.nIns()>2 && !hasMidUser(region) ) {
+                    assert !region.inProgress();
+                    // Fold Phis
+                    for( Node use : _outputs ) {
+                        if( use instanceof PhiNode phi ) {
+                            PhiNode phi2 = phi.in(i) instanceof PhiNode phi2x && phi2x.region()==region ? phi2x : null;
+                            for( int j=1; j<region.nIns(); j++ )
+                                phi.addDef(phi2==null ? phi.in(i) : phi2.in(j));
+                            CodeGen.CODE.add(phi);
+                        }
+                    }
+
+                    // Fold Region
+                    for( int j=1; j<region.nIns(); j++ )
+                        addDef(region.in(j));
+                    setDef(i,Parser.XCTRL);
+                    return this;
+                }
+        }
+
         return null;
     }
 
@@ -98,6 +122,23 @@ public class RegionNode extends CFGNode {
         return false;
     }
 
+    boolean hasMidUser(RegionNode r) {
+        for( Node use : r._outputs ) {
+            if( use==this ) continue;
+            if( use instanceof PhiNode ) {
+                for( Node data : use._outputs ) {
+                    if( !(data instanceof PhiNode phi2) || phi2.region()!=this ) {
+                        use.addDep(this); return true; }
+                    assert phi2.region()==this;
+                }
+            } else {
+                use.addDep(this);
+                return true; } // Control user
+        }
+        return false;
+    }
+
+
     // Immediate dominator of Region is a little more complicated.
     @Override public int idepth() {
         if( _idepth!=0 ) return _idepth;
diff --git a/chapter21/src/main/java/com/seaofnodes/simple/node/cpus/arm/UJmpARM.java b/chapter21/src/main/java/com/seaofnodes/simple/node/cpus/arm/UJmpARM.java
index 4b2d14df..a69c3c7c 100644
--- a/chapter21/src/main/java/com/seaofnodes/simple/node/cpus/arm/UJmpARM.java
+++ b/chapter21/src/main/java/com/seaofnodes/simple/node/cpus/arm/UJmpARM.java
@@ -42,7 +42,7 @@ public class UJmpARM extends CFGNode implements MachNode, RIPRelSize {
 
     @Override public void asm(CodeGen code, SB sb) {
         CFGNode target = uctrl();
-        assert target.nOuts() > 1; // Should optimize jmp to empty targets
+        //assert target.nOuts() > 1; // Should optimize jmp to empty targets
         sb.p(label(target));
     }
 }
diff --git a/chapter21/src/main/java/com/seaofnodes/simple/node/cpus/riscv/UJmpRISC.java b/chapter21/src/main/java/com/seaofnodes/simple/node/cpus/riscv/UJmpRISC.java
index 9f623b0c..6473d81c 100644
--- a/chapter21/src/main/java/com/seaofnodes/simple/node/cpus/riscv/UJmpRISC.java
+++ b/chapter21/src/main/java/com/seaofnodes/simple/node/cpus/riscv/UJmpRISC.java
@@ -40,7 +40,7 @@ public class UJmpRISC extends CFGNode implements MachNode, RIPRelSize {
 
     @Override public void asm(CodeGen code, SB sb) {
         CFGNode target = uctrl();
-        assert target.nOuts() > 1; // Should optimize jmp to empty targets
+        //assert target.nOuts() > 1; // Should optimize jmp to empty targets
         sb.p(label(target));
     }
 }
diff --git a/chapter21/src/main/java/com/seaofnodes/simple/node/cpus/x86_64_v2/UJmpX86.java b/chapter21/src/main/java/com/seaofnodes/simple/node/cpus/x86_64_v2/UJmpX86.java
index c2c804a2..ec4eeb3b 100644
--- a/chapter21/src/main/java/com/seaofnodes/simple/node/cpus/x86_64_v2/UJmpX86.java
+++ b/chapter21/src/main/java/com/seaofnodes/simple/node/cpus/x86_64_v2/UJmpX86.java
@@ -47,7 +47,7 @@ public class UJmpX86 extends CFGNode implements MachNode, RIPRelSize {
 
     @Override public void asm(CodeGen code, SB sb) {
         CFGNode target = uctrl();
-        assert target.nOuts() > 1; // Should optimize jmp to empty targets
+        //assert target.nOuts() > 1; // Should optimize jmp to empty targets
         sb.p(label(target));
     }
 }