optimizing stack shuffling; architecture description

library/bootstrap/boot-stage1.factor

@@ -113,6 +113,7 @@ sequences io vectors words ;
 
         "/library/bootstrap/image.factor"
 
+        "/library/inference/shuffle.factor"
         "/library/inference/dataflow.factor"
         "/library/inference/inference.factor"
         "/library/inference/branches.factor"
@@ -128,6 +129,7 @@ sequences io vectors words ;
         "/library/inference/call-optimizers.factor"
         "/library/inference/print-dataflow.factor"
         
+        "/library/compiler/architecture.factor"
         "/library/compiler/assembler.factor"
         "/library/compiler/relocate.factor"
         "/library/compiler/xt.factor"

library/bootstrap/boot-stage2.factor

@@ -18,6 +18,7 @@ words ;
 
 cpu "x86" = [
     "/library/compiler/x86/assembler.factor"
+    "/library/compiler/x86/architecture.factor"
     "/library/compiler/x86/generator.factor"
     "/library/compiler/x86/slots.factor"
     "/library/compiler/x86/stack.factor"
@@ -27,6 +28,7 @@ cpu "x86" = [
 
 cpu "ppc" = [
     "/library/compiler/ppc/assembler.factor"
+    "/library/compiler/ppc/architecture.factor"
     "/library/compiler/ppc/generator.factor"
     "/library/compiler/ppc/slots.factor"
     "/library/compiler/ppc/stack.factor"

library/compiler/architecture.factor

@@ -0,0 +1,9 @@
+IN: compiler-frontend
+
+! A few things the front-end needs to know about the back-end.
+
+DEFER: fixnum-imm? ( -- ? )
+#! Can fixnum operations take immediate operands?
+
+DEFER: vregs ( -- n )
+#! Number of vregs

library/compiler/intrinsics.factor

@@ -5,11 +5,6 @@ USING: assembler compiler-backend generic hashtables inference
 kernel kernel-internals lists math math-internals namespaces
 sequences vectors words ;
 
-! Architecture description
-: fixnum-imm?
-    #! Can fixnum operations take immediate operands?
-    cpu "x86" = ;
-
 : node-peek ( node -- value ) node-in-d peek ;
 
 : type-tag ( type -- tag )

library/compiler/ppc/architecture.factor

@@ -0,0 +1,13 @@
+IN: compiler-frontend
+USING: assembler compiler-backend math ;
+
+! Architecture description
+: fixnum-imm? ( -- ? )
+    #! Can fixnum operations take immediate operands?
+    f ;
+
+: vregs ( -- n )
+    #! Number of vregs
+    8 ;
+
+M: vreg v>operand vreg-n 3 + ;

library/compiler/x86/architecture.factor

@@ -0,0 +1,13 @@
+IN: compiler-frontend
+USING: assembler compiler-backend sequences ;
+
+! Architecture description
+: fixnum-imm? ( -- ? )
+    #! Can fixnum operations take immediate operands?
+    t ;
+
+: vregs ( -- n )
+    #! Number of vregs
+    3 ;
+
+M: vreg v>operand vreg-n { EAX ECX EDX } nth ;

library/compiler/x86/generator.factor

@@ -4,8 +4,6 @@ IN: compiler-backend
 USING: alien assembler compiler inference kernel
 kernel-internals lists math memory namespaces sequences words ;
 
-M: vreg v>operand vreg-n { EAX ECX EDX } nth ;
-
 ! Not used on x86
 M: %prologue generate-node drop ;
 

library/inference/dataflow.factor

@@ -45,7 +45,7 @@ C: meet ( values -- value )
 ! representations used by Factor. It annotates concatenative
 ! code with stack flow information and types.
 
-TUPLE: node param in-d out-d in-r out-r
+TUPLE: node param shuffle
        classes literals history
        successor children ;
 
@@ -53,17 +53,29 @@ M: node = eq? ;
 
 : make-node ( param in-d out-d in-r out-r node -- node )
     [
-        >r {{ }} clone {{ }} clone { } clone f f <node> r>
+        >r
+        swapd <shuffle> {{ }} clone {{ }} clone { } clone f f <node>
+        r>
         set-delegate
     ] keep ;
 
+: node-in-d  node-shuffle shuffle-in-d  ;
+: node-in-r  node-shuffle shuffle-in-r  ;
+: node-out-d node-shuffle shuffle-out-d ;
+: node-out-r node-shuffle shuffle-out-r ;
+
+: set-node-in-d  node-shuffle set-shuffle-in-d  ;
+: set-node-in-r  node-shuffle set-shuffle-in-r  ;
+: set-node-out-d node-shuffle set-shuffle-out-d ;
+: set-node-out-r node-shuffle set-shuffle-out-r ;
+
 : empty-node f { } { } { } { } ;
 : param-node ( label) { } { } { } { } ;
 : in-d-node ( inputs) >r f r> { } { } { } ;
 : out-d-node ( outputs) >r f { } r> { } { } ;
 
-: d-tail ( n -- list ) meta-d get tail* >vector ;
+: d-tail ( n -- list ) meta-d get tail* ;
-: r-tail ( n -- list ) meta-r get tail* >vector ;
+: r-tail ( n -- list ) meta-r get tail* ;
 
 : node-child node-children first ;
 
@@ -146,12 +158,6 @@ SYMBOL: current-node
 : with-nesting ( quot -- new-node | quot: -- new-node )
     nest-node 2slip unnest-node ; inline
 
-: copy-effect ( from to -- )
-    over node-in-d over set-node-in-d
-    over node-in-r over set-node-in-r
-    over node-out-d over set-node-out-d
-    swap node-out-r swap set-node-out-r ;
-
 : node-effect ( node -- [[ d-in meta-d ]] )
     dup node-in-d swap node-out-d cons ;
 
@@ -275,11 +281,7 @@ DEFER: subst-value
     ] each-node-with ;
 
 : (clone-node) ( node -- node )
-    clone
+    clone dup node-shuffle clone over set-node-shuffle ;
-    dup node-in-d clone over set-node-in-d
-    dup node-in-r clone over set-node-in-r
-    dup node-out-d clone over set-node-out-d
-    dup node-out-r clone over set-node-out-r ;
 
 : clone-node ( node -- node )
     dup [

library/inference/optimizer.factor

@@ -56,7 +56,14 @@ M: #push optimize-node* ( node -- node/t )
 
 ! #shuffle
 M: #shuffle optimize-node*  ( node -- node/t )
-    [ dup node-in-d empty? swap node-in-r empty? and ] prune-if ;
+    dup node-successor dup #shuffle? [
+        [ >r node-shuffle r> node-shuffle compose-shuffle ] keep
+        [ set-node-shuffle ] keep
+    ] [
+        drop [
+            dup node-in-d empty? swap node-in-r empty? and
+        ] prune-if
+    ] ifte ;
 
 ! #ifte
 : static-branch? ( node -- lit ? )

library/inference/print-dataflow.factor

@@ -17,8 +17,7 @@ M: comment pprint* ( ann -- )
     rot [ <comment> , ] [ 2drop ] ifte ;
 
 : value-str ( prefix values -- str )
-    [ value-uid word-name append ] map-with
+    [ value-uid word-name append ] map-with concat ;
-    " " join ;
 
 : effect-str ( node -- str )
     [
@@ -27,7 +26,7 @@ M: comment pprint* ( ann -- )
         " --" %
         " " over node-out-d value-str %
         " r: " swap node-out-r value-str %
-    ] "" make ;
+    ] "" make 1 swap tail ;
 
 M: #push node>quot ( ? node -- )
     node-out-d [ literal-value literalize ] map % drop ;

library/inference/shuffle.factor

@@ -0,0 +1,59 @@
+IN: inference
+USING: kernel math namespaces sequences ;
+
+TUPLE: shuffle in-d in-r out-d out-r ;
+
+: empty-shuffle { } { } { } { } <shuffle> ;
+
+: cut* ( seq1 seq2 -- seq seq ) [ head* ] 2keep tail* ;
+
+: load-shuffle ( d r shuffle -- )
+    tuck shuffle-in-r [ set ] 2each shuffle-in-d [ set ] 2each ;
+
+: store-shuffle ( shuffle -- d r )
+    dup shuffle-out-d [ get ] map swap shuffle-out-r [ get ] map ;
+
+: shuffle* ( d r shuffle -- d r )
+    [ [ load-shuffle ] keep store-shuffle ] with-scope ;
+
+: split-shuffle ( d r shuffle -- d' r' d r )
+    tuck shuffle-in-r length swap cut*
+    >r >r shuffle-in-d length swap cut*
+    r> swap r> ;
+
+: join-shuffle ( d' r' d r -- d r )
+    swapd append >r append r> ;
+
+: shuffle ( d r shuffle -- d r )
+    #! d and r lengths must be at least the required length for
+    #! the shuffle.
+    [ split-shuffle ] keep shuffle* join-shuffle ;
+
+: fix-compose-d ( s1 s2 -- )
+    over shuffle-out-d over shuffle-in-d length< [
+        over shuffle-out-d length over shuffle-in-d head*
+        [ pick shuffle-in-d append pick set-shuffle-in-d ] keep
+        pick shuffle-out-d append pick set-shuffle-out-d
+    ] when 2drop ;
+
+: fix-compose-r ( s1 s2 -- )
+    over shuffle-out-r over shuffle-in-r length< [
+        over shuffle-out-r length over shuffle-in-r head*
+        [ pick shuffle-in-r append pick set-shuffle-in-r ] keep
+        pick shuffle-out-r append pick set-shuffle-out-r
+    ] when 2drop ;
+
+: compose-shuffle ( s1 s2 -- s1+s2 )
+    #! s1's d and r output lengths must be at least the required
+    #! length for the shuffle. If they are not, a special
+    #! behavior is used which is only valid for the optimizer.
+    >r clone r> clone 2dup fix-compose-d 2dup fix-compose-r
+    >r dup shuffle-out-d over shuffle-out-r r> shuffle
+    >r >r dup shuffle-in-d swap shuffle-in-r r> r> <shuffle> ;
+
+M: shuffle clone ( shuffle -- shuffle )
+    [ shuffle-in-d clone ] keep
+    [ shuffle-in-r clone ] keep
+    [ shuffle-out-d clone ] keep
+    shuffle-out-r clone
+    <shuffle> ;

library/inference/words.factor

@@ -132,7 +132,7 @@ M: symbol apply-object ( word -- )
     dup recursive-label? [
         node,
     ] [
-        node-child splice-node
+        node-child node-successor splice-node
     ] ifte ;
 
 M: compound apply-object ( word -- )

library/test/inference.factor

@@ -2,6 +2,22 @@ IN: temporary
 USING: generic inference kernel lists math math-internals
 namespaces parser sequences test vectors ;
 
+[
+    << shuffle f { "a" } { } { "a" } { "a" } >>
+] [
+    << shuffle f { "a" } { } { "a" "a" } { } >>
+    << shuffle f { "b" } { } { } { "b" } >>
+    compose-shuffle
+] unit-test
+
+[
+    << shuffle f { "b" "a" } { } { "b" "b" } { } >>
+] [
+    << shuffle f { "a" } { } { } { } >>
+    << shuffle f { "b" } { } { "b" "b" } { } >>
+    compose-shuffle
+] unit-test
+
 : simple-effect first2 >r length r> length 2vector ;
 
 [ { 0 2 } ] [ [ 2 "Hello" ] infer simple-effect ] unit-test