Popular use grouping for RegAlloc

Fix arm BL encodings; ADD/SUB with reversed sign.

8 files changed, 210 insertions, 66 deletions

diff --git a/chapter25/src/main/java/com/seaofnodes/simple/codegen/RegAlloc.java b/chapter25/src/main/java/com/seaofnodes/simple/codegen/RegAlloc.java
index 4a880a19..196a1142 100644
--- a/chapter25/src/main/java/com/seaofnodes/simple/codegen/RegAlloc.java
+++ b/chapter25/src/main/java/com/seaofnodes/simple/codegen/RegAlloc.java
@@ -344,56 +344,40 @@ public class RegAlloc {
         return true;
     }
 
-    // Single-def live range with an empty mask.  There are many single-reg
-    // uses.  Theory is there's many repeats of the same reg amongst the uses.
-    // In of splitting once per use, start by splitting into groups based on
-    // required input register.
+    // Popular single-def value: share a split among uses with compatible masks.
     boolean splitEmptyMaskByUse( byte round, LRG lrg ) {
         Node def = (Node)lrg._machDef;
+        // Snapshot distinct users: rewriting an input changes def's output list,
+        // and a call can use the same value in more than one argument.
+        Ary<Node> uses = new Ary<>(Node.class);
+        for( Node use : def._outputs )
+            if( uses.find(use)==-1 ) uses.push(use);
 
-        // Look at each use, and break into non-overlapping register classes.
         Ary<RegMask> rclass = new Ary<>(RegMask.class);
-        boolean done=false;
         int ncalls=0;
-        while( !done ) {
-            done = true;
-            for( Node use : def._outputs )
-                if( use instanceof MachNode mach ) {
-                    if( mach instanceof CallNode ) ncalls++;
-                    for( int i=1; i<use.nIns(); i++ )
-                        if( use.in(i)==def )
-                            done = putIntoRegClass( rclass, mach.regmap(i) );
-                }
-        }
-
-        // See how many register classes we split into.  Generally not
-        // productive to split like this across calls, which are going to kill
-        // all registers anyways.
-        if( rclass._len <= 1 || ncalls > 1 ) return false;
+        for( Node use : uses )
+            if( use instanceof MachNode mach ) {
+                // Sharing copies across several calls usually loses to their kills.
+                if( use instanceof CallNode && ++ncalls > 1 ) return false;
+                for( int i=1; i<use.nIns(); i++ )
+                    if( use.in(i)==def && mach.regmap(i)!=null )
+                        putIntoRegClass(rclass,mach.regmap(i));
+            }
+        if( rclass._len <= 1 ) return false;
 
-        // Split by class
-        Ary<Node> ns = new Ary<>(Node.class);
         for( RegMask rmask : rclass ) {
-            ns.addAll(def._outputs);
             Node split = makeSplit(def,"popular",round,lrg,rmask);
-            split.insertAfter( def );
+            split.insertAfter(def);
             if( split.nIns()>1 ) split.setDef(1,def);
-            // all uses by class to split
-            for( int j=0; j < def._outputs._len; j++ ) {
-                Node use = def._outputs.at(j);
-                if( use instanceof MachNode mach && use!=split ) {
-                    // Check all use inputs for n, in case there's several
-                    for( int i = 1; i < use.nIns(); i++ )
-                        // Find a def input, and check register class
-                        if( use.in( i ) == def && mach.regmap( i ).overlap( rmask ) ) {
-                            RegMask m = mach.regmap( i );
-                            if( m!=null && mach.regmap( i ).overlap( rmask ) )
-                                // Modify use to use the split version specialized to this rclass
-                                { use.setDefOrdered( i, split ); j--; break; }
+            for( Node use : uses )
+                if( use instanceof MachNode mach )
+                    for( int i=1; i<use.nIns(); i++ ) {
+                        RegMask mask = mach.regmap(i);
+                        if( use.in(i)==def && mask!=null && mask.overlap(rmask) ) {
+                            assert rmask.and(mask)==rmask;
+                            use.setDefOrdered(i,split);
                         }
-                }
-            }
-            ns.clear();
+                    }
         }
         return true;
     }
@@ -419,22 +403,20 @@ public class RegAlloc {
         return progress;
     }
 
-    // Put use into a register class, perhaps adding a class or perhaps
-    // narrowing a class (and causing a repeat)
-    private static boolean putIntoRegClass( Ary<RegMask> rclass, RegMask rmask ) {
+    // Narrowing preserves all earlier users of a class; new classes are disjoint.
+    // Each use contains its first overlapping class, so one pass suffices.
+    private static void putIntoRegClass( Ary<RegMask> rclass, RegMask rmask ) {
         for( int i=0; i<rclass._len; i++ ) {
             RegMask omask = rclass.at(i);
-            if( omask.and(rmask) == omask ) return true; // Within the same register class
-            if( omask.overlap(rmask) ) {
+            if( !omask.overlap(rmask) ) continue;
+            if( omask.and(rmask)!=omask )
                 rclass.set(i,new RegMask(omask.copy().and(rmask)));
-                return false;   // Need go again
-            }
+            return;
         }
-        // Add a new class, no need to go again
         rclass.push(rmask);
-        return true;
     }
 
+
     // Self conflicts require Phis (or two-address).
     // Insert a split after every def.
     boolean splitSelfConflict( byte round, LRG lrg ) {
diff --git a/chapter25/src/main/java/com/seaofnodes/simple/node/cpus/arm/FunARM.java b/chapter25/src/main/java/com/seaofnodes/simple/node/cpus/arm/FunARM.java
index 44a417f9..b9e85202 100644
--- a/chapter25/src/main/java/com/seaofnodes/simple/node/cpus/arm/FunARM.java
+++ b/chapter25/src/main/java/com/seaofnodes/simple/node/cpus/arm/FunARM.java
@@ -10,13 +10,12 @@ public class FunARM  extends FunNode implements MachNode {
     FunARM(FunNode fun) { super(fun); }
     @Override public void computeFrameAdjust(CodeGen code, int maxReg) {
         super.computeFrameAdjust(code,maxReg);
-        if( _hasCalls )         // If non-leaf, pad to 16b
-            _frameAdjust = ((_frameAdjust+8) & -16);
+        _frameAdjust = (_frameAdjust+15) & -16; // Keep SP aligned, including leaf spills.
     }
     @Override public void encoding( Encoding enc ) {
         int sz = _frameAdjust;
         if( sz == 0 ) return;   // Skip if no frame adjust
         if( sz >= 1L<<12 ) throw Utils.TODO();
-        enc.add4(arm.imm_inst(arm.OPI_ADD, -sz&0xFFF, arm.RSP, arm.RSP));
+        enc.add4(arm.imm_inst(arm.OPI_SUB, sz, arm.RSP, arm.RSP));
     }
 }
diff --git a/chapter25/src/main/java/com/seaofnodes/simple/node/cpus/arm/RetARM.java b/chapter25/src/main/java/com/seaofnodes/simple/node/cpus/arm/RetARM.java
index 2136ebf5..4597666b 100644
--- a/chapter25/src/main/java/com/seaofnodes/simple/node/cpus/arm/RetARM.java
+++ b/chapter25/src/main/java/com/seaofnodes/simple/node/cpus/arm/RetARM.java
@@ -9,7 +9,7 @@ public class RetARM extends ReturnNode implements MachNode {
     @Override public void encoding( Encoding enc ) {
         int frameAdjust = ((FunARM)fun())._frameAdjust;
         if( frameAdjust > 0 )
-            enc.add4(arm.imm_inst(arm.OPI_ADD, (frameAdjust*-8)&0xFFF, arm.RSP, arm.RSP));
+            enc.add4(arm.imm_inst(arm.OPI_ADD, frameAdjust, arm.RSP, arm.RSP));
         enc.add4(arm.ret(arm.OP_RET));
     }
 }
diff --git a/chapter25/src/main/java/com/seaofnodes/simple/node/cpus/arm/arm.java b/chapter25/src/main/java/com/seaofnodes/simple/node/cpus/arm/arm.java
index 20343fb4..510d5e78 100644
--- a/chapter25/src/main/java/com/seaofnodes/simple/node/cpus/arm/arm.java
+++ b/chapter25/src/main/java/com/seaofnodes/simple/node/cpus/arm/arm.java
@@ -319,6 +319,11 @@ public class arm extends Machine {
         short self = enc.reg(n);
         short reg1 = enc.reg(n2);
 
+        // ADD/SUB encode an unsigned magnitude; switch operation for a negative addend.
+        if( imm12<0 && (opcode==OPI_ADD || opcode==OPI_SUB) ) {
+            opcode ^= OPI_ADD ^ OPI_SUB;
+            imm12 = -imm12;
+        }
         int body = imm_inst(opcode, imm12&0xFFF, reg1, self);
         enc.add4(body);
     }
@@ -429,7 +434,7 @@ public class arm extends Machine {
     }
 
     public static int ret(int opcode) {
-        return (opcode << 10);
+        return (opcode << 10) | (X30 << 5);
     }
 
     // FMOV (scalar, immediate)
diff --git a/chapter25/src/test/java/com/seaofnodes/simple/Chapter21Test.java b/chapter25/src/test/java/com/seaofnodes/simple/Chapter21Test.java
index e194829a..10d52ea7 100644
--- a/chapter25/src/test/java/com/seaofnodes/simple/Chapter21Test.java
+++ b/chapter25/src/test/java/com/seaofnodes/simple/Chapter21Test.java
@@ -10,6 +10,119 @@ import org.junit.Test;
 import static org.junit.Assert.*;
 
 public class Chapter21Test {
+    @Test public void testArmCallInstructions() {
+        EvalArm64 cpu = new EvalArm64(new byte[512],256);
+        for( int delta : new int[]{4,32,-4} ) {
+            cpu._pc = 128;
+            for( int i=0; i<32; i++ ) cpu.regs[i]=1000+i;
+            cpu.st4(128,0x94000000 | ((delta>>2)&0x03FFFFFF)); // BL
+            assertEquals(0,cpu.step(1));
+            assertEquals(128+delta,cpu._pc);
+            for( int i=0; i<32; i++ ) assertEquals(i==30 ? 132 : 1000+i,cpu.regs[i]);
+        }
+        for( int rn : new int[]{9,30} )
+            for( boolean call : new boolean[]{false,true} ) {
+                cpu._pc=4;
+                cpu.regs[rn]=64;
+                cpu.st4(4,(call ? 0xD63F0000 : 0xD65F0000) | (rn<<5)); // BLR / RET
+                assertEquals(0,cpu.step(1));
+                assertEquals(64,cpu._pc);
+                if( call ) assertEquals(8,cpu.regs[30]);
+            }
+    }
+
+    @Test public void testArmImmediateArithmetic() {
+        EvalArm64 cpu = new EvalArm64(new byte[16],16);
+        // Unsigned imm12, including bit 11 and the optional 12-bit shift.
+        for( int[] test : new int[][]{{0x91200020,2148},{0xD1200020,-1948},
+                                     {0x91400420,4196},{0xD1400420,-3996}} ) {
+            cpu._pc=0; cpu.regs[1]=100;
+            cpu.N=true; cpu.Z=false; cpu.C=true; cpu.V=true;
+            cpu.st4(0,test[0]);
+            assertEquals(0,cpu.step(1));
+            assertEquals(test[1],cpu.regs[0]);
+            assertTrue(cpu.N); assertFalse(cpu.Z); assertTrue(cpu.C); assertTrue(cpu.V);
+        }
+    }
+
+    @Test public void testArmCallsAndFrames() throws IOException {
+        String src = """
+            val sum = { int n ->
+                if( n<=0 ) return 1;
+                return sum(n-1)+n;
+            };
+            val run = { int n -> return sum(n)+sum(n+1); };
+            """;
+        assertTrue("Ordinary calls must survive optimization",checkArmCalls(src,"run",4,27)>0);
+    }
+
+    @Test public void testArmVectorGrowth() throws IOException {
+        String src = """
+            struct _Vec {
+                u32 !len;
+                int[] !buf;
+                new _Vec = { int[] initial -> buf=initial; };
+                val grow = { int sz ->
+                    var buf2 = new int[sz];
+                    for( int i=0; i<len; i++ ) buf2[i]=buf[i];
+                    buf=buf2;
+                };
+                val add = { int e ->
+                    if( len>=buf# ) grow(buf#*2);
+                    buf[len++]=e;
+                    return self;
+                };
+            };
+            val run = { int arg ->
+                val v = new _Vec(new int[4]).add(2).add(3).add(5).add(7).add(11);
+                return v.len*100+v.buf[arg];
+            };
+            """;
+        checkArmCalls(src,"run",4,511);
+    }
+
+    private static int checkArmCalls(String src, String entryName, int arg, int result) throws IOException {
+        CodeGen code = new CodeGen(src).driver("arm","SystemV",true,false);
+        byte[] image = new byte[1<<20];
+        System.arraycopy(code._encoding.bits(),0,image,0,code._encoding.bits().length);
+        EvalArm64 cpu = new EvalArm64(image,1<<16);
+        boolean entry=false;
+        int frames=0, calls=0;
+        for( var bb : code._cfg ) {
+            if( bb instanceof com.seaofnodes.simple.node.FunNode fun ) {
+                int off = code._encoding.opStart(fun);
+                if( entryName.equals(fun._name) ) { cpu._pc=off; entry=true; }
+                int frame=fun._frameAdjust;
+                assertEquals(0,frame&15);
+                if( frame>0 ) {
+                    frames++;
+                    assertEquals(0xD10003FF | (frame<<10),cpu.ld4s(off)); // SUB SP,SP,#bytes
+                }
+            }
+            if( bb instanceof com.seaofnodes.simple.node.ReturnNode ret ) {
+                int off=code._encoding.opStart(ret);
+                int frame=ret.fun()._frameAdjust;
+                if( frame>0 ) {
+                    assertEquals(0x910003FF | (frame<<10),cpu.ld4s(off)); // ADD SP,SP,#bytes
+                    off+=4;
+                }
+                assertEquals(0xD65F03C0,cpu.ld4s(off)); // RET X30
+            }
+            if( bb instanceof com.seaofnodes.simple.node.cpus.arm.CallARM ) calls++;
+        }
+        assertTrue("Must execute the program entry",entry);
+        assertTrue("Exercise stack frames",frames>0);
+        for( int i=19; i<=29; i++ ) cpu.regs[i]=1000+i;
+        cpu.regs[0]=arg;
+        assertEquals(0,cpu.step(10000));
+        assertEquals("Must return to the harness",0,cpu._pc);
+        assertEquals(result,cpu.regs[0]);
+        assertEquals(1<<16,cpu.regs[31]);
+        for( int i=19; i<=29; i++ ) assertEquals(1000+i,cpu.regs[i]);
+        return calls;
+    }
+
+
     @Test public void testCoalescing() { com.seaofnodes.simple.codegen.RegAllocTestSupport.coalescing(); }
 
     @Test public void testArmSubtractRegisters() {
diff --git a/chapter25/src/test/java/com/seaofnodes/simple/Chapter23Test.java b/chapter25/src/test/java/com/seaofnodes/simple/Chapter23Test.java
index 923395bf..2f0490a5 100644
--- a/chapter25/src/test/java/com/seaofnodes/simple/Chapter23Test.java
+++ b/chapter25/src/test/java/com/seaofnodes/simple/Chapter23Test.java
@@ -15,6 +15,8 @@ import org.junit.Test;
 import static org.junit.Assert.*;
 
 public class Chapter23Test {
+    @Test public void testPopularUses() { com.seaofnodes.simple.codegen.RegAllocTestSupport.popularUses(); }
+
     @Test public void testDiagnosticTypeAccessors() throws Exception {
         new CodeGen("return 0;").parse();
         var tfp = TypeFunPtr.TEST;
diff --git a/chapter25/src/test/java/com/seaofnodes/simple/EvalArm64.java b/chapter25/src/test/java/com/seaofnodes/simple/EvalArm64.java
index 89e09f69..aa128b72 100644
--- a/chapter25/src/test/java/com/seaofnodes/simple/EvalArm64.java
+++ b/chapter25/src/test/java/com/seaofnodes/simple/EvalArm64.java
@@ -307,11 +307,11 @@ public class EvalArm64 {
                 break;
             }
 
-            case 0x91: {  // add(immediate)
+            case 0x91, 0xD1: { // ADD/SUB (immediate), unsigned imm12, optional LSL #12
                 int rn = (ir >> 5) & 0x1F;
-                int immediate = ir << 10 >> 20;
-                // only hit it twice
-                rval = regs[rn] + immediate;
+                long immediate = (ir >>> 10) & 0xFFF;
+                if( (ir & (1<<22))!=0 ) immediate <<= 12;
+                rval = opcode1==0x91 ? regs[rn] + immediate : regs[rn] - immediate;
                 break;
             }
             case 0x92: {        // and(immediate)
@@ -330,6 +330,7 @@ public class EvalArm64 {
             }
 
             case 0x94, 0x95, 0x96, 0x97: { // bl
+                rdid = -1; // BL writes only X30, not a displacement-selected register.
                 rval = pc + 4;
                 regs[arm.X30] = rval;
                 int imm26 = (ir & 0x03FFFFFF);
@@ -575,15 +576,15 @@ public class EvalArm64 {
                 break;
             }
             case 0xD6: {
-                if( (ir & 0xFFFF) != 0 ) throw Utils.TODO();
-                if( ((ir>>16) & 0xFF) != 0x5F ) throw Utils.TODO(); // 0x3F is register call
+                int op = ir & 0xFFFFFC1F;
+                int rn = (ir >>> 5) & 31;
+                int target = (int)regs[rn]; // Read before BLR overwrites X30.
                 rdid = -1;
-                int rn = ir << 22 >> 27;
-                assert rn == 0;
-                pc = (int)regs[30];
-                // Return to zero breaks simulation
-                if( pc==0 ) break outer;
-                pc -= 4;        // Undo post-increment
+                if( op==0xD63F0000 ) regs[arm.X30] = pc+4; // BLR
+                else if( op!=0xD65F0000 ) { trap = 3; break; } // RET
+                pc = target;
+                if( op==0xD65F0000 && pc==0 ) break outer;
+                pc -= 4;
                 break;
             }
 
diff --git a/chapter25/src/test/java/com/seaofnodes/simple/codegen/RegAllocTestSupport.java b/chapter25/src/test/java/com/seaofnodes/simple/codegen/RegAllocTestSupport.java
index 5e1b276e..a92859d4 100644
--- a/chapter25/src/test/java/com/seaofnodes/simple/codegen/RegAllocTestSupport.java
+++ b/chapter25/src/test/java/com/seaofnodes/simple/codegen/RegAllocTestSupport.java
@@ -105,6 +105,48 @@ public class RegAllocTestSupport {
         assertTrue(BuildLRG.run(0,new RegAlloc(code)));
     }
 
+    private static class CallOp extends CallNode implements MachNode {
+        CallOp(Node ctrl, Node def) { super(null,ctrl,null,def,def); }
+        @Override public RegMask regmap(int i) { return i==2 ? A : i==3 ? B : null; }
+        @Override public RegMask outregmap() { return null; }
+        @Override public void encoding(Encoding enc) { throw new AssertionError(); }
+    }
+
+    public static void popularUses() {
+        for( int calls=0; calls<=2; calls++ ) {
+            CodeGen code = graph();
+            RegAlloc alloc = new RegAlloc(code);
+            RegionNode block = new RegionNode(null,code._start,code._start);
+            Op def = new Op(new RegMask(7L),null,null,false,block);
+            LRG lrg = alloc.newLRG(def,null);
+            lrg._machDef = def;
+            lrg._mask = new RegMask(0L);
+            Op ab = new Op(null,new RegMask(3L),null,false,block,def);
+            Op bc = new Op(null,new RegMask(6L),null,false,block,def);
+            Op a = new Op(null,A,null,false,block,def,def);
+            Op b = new Op(null,B,null,false,block,def);
+            Op dep = new Op(null,null,null,false,block,def);
+            CallOp call = calls>0 ? new CallOp(block,def) : null;
+            if( calls==2 ) new CallOp(block,def);
+            boolean split = alloc.splitEmptyMaskByUse((byte)0,lrg);
+            assertEquals(calls<2,split);
+            if( !split ) { assertSame(def,a.in(1)); assertSame(def,b.in(1)); continue; }
+            // Two shared copies, even with overlapping masks and repeated inputs.
+            assertTrue(a.in(1) instanceof SplitNode);
+            assertTrue(b.in(1) instanceof SplitNode);
+            assertNotSame(a.in(1),b.in(1));
+            assertSame(a.in(1),a.in(2));
+            assertSame(b.in(1),ab.in(1));
+            assertSame(b.in(1),bc.in(1));
+            assertSame(def,dep.in(1)); // A scheduling dependency has no register.
+            assertEquals(3,def.nOuts()); // Two copies and the dependency.
+            if( call!=null ) {
+                assertSame(a.in(1),call.in(2));
+                assertSame(b.in(1),call.in(3));
+            }
+        }
+    }
+
     // Check the scheduled graph before encoding adds untyped branches or rewrites tail calls.
     public static class CheckedCodeGen extends CodeGen {
         public CheckedCodeGen(String src) { super(src); }