bug fixes for the sinker

compiler/cmm/CmmSink.hs

@@ -10,8 +10,7 @@ import CmmUtils
 import Hoopl
 
 import UniqFM
-import Unique
-import Outputable
+-- import Outputable
 
 import Data.List (partition)
 import qualified Data.Set as Set
@@ -67,6 +66,7 @@ cmmSink graph = ofBlockList (g_entry graph) $ sink mapEmpty $ blocks
   where
 
   liveness = cmmLiveness graph
+  getLive l = mapFindWithDefault Set.empty l liveness
 
   blocks = postorderDfs graph
 
@@ -75,10 +75,11 @@ cmmSink graph = ofBlockList (g_entry graph) $ sink mapEmpty $ blocks
   succ_counts = foldr (\l -> mapInsertWith (+) l 1) mapEmpty all_succs
   join_pts = mapFilter (>1) succ_counts
 
+
   sink :: BlockEnv [Assignment] -> [CmmBlock] -> [CmmBlock]
   sink _ [] = []
   sink sunk (b:bs) =
-    pprTrace "sink" (ppr lbl) $
+    -- pprTrace "sink" (ppr lbl) $
     blockJoin first final_middle last : sink sunk' bs
     where
       lbl = entryLabel b
@@ -86,30 +87,48 @@ cmmSink graph = ofBlockList (g_entry graph) $ sink mapEmpty $ blocks
       (middle', assigs) = walk ann_middles emptyBlock
                                (mapFindWithDefault [] lbl sunk)
 
-      live = Set.unions [ mapFindWithDefault Set.empty l liveness | l <- succs ]
-      live_middle = gen_kill last live
+      succs = successors last
 
+      -- Annotate the middle nodes with the registers live *after*
+      -- the node.  This will help us decide whether we can inline
+      -- an assignment in the current node or not.
+      live = Set.unions (map getLive succs)
+      live_middle = gen_kill last live
       ann_middles = annotate live_middle (blockToList middle)
 
-      getLive l = mapFindWithDefault Set.empty l liveness
-      succs = successors last
+      -- We cannot sink into join points (successors with more than
+      -- one predecessor), so identify the join points and the set
+      -- of registers live in them.
       (joins, nonjoins) = partition (`mapMember` join_pts) succs
-      live_in_nonjoins = concatMap (Set.toList . getLive) nonjoins
-      live_in_joins :: [LocalReg]
-      live_in_joins    = concatMap (Set.toList . getLive) joins
+      live_in_joins = Set.unions (map getLive joins)
+
+      -- We do not want to sink an assignment into multiple branches,
+      -- so identify the set of registers live in multiple successors.
+      -- This is made more complicated because when we sink an assignment
+      -- into one branch, this might change the set of registers that are
+      -- now live in multiple branches.
+      init_live_sets = map getLive nonjoins
+      live_in_multi live_sets r =
+         case filter (Set.member r) live_sets of
+           (_one:_two:_) -> True
+           _ -> False
+
+      -- Now, drop any assignments that we will not sink any further.
+      (dropped_last, assigs') = dropAssignments drop_if init_live_sets assigs
+
+      drop_if a@(r,rhs,_) live_sets = (should_drop, live_sets')
+          where
+            should_drop =  a `conflicts` last
+                        || {- not (isTiny rhs) && -} live_in_multi live_sets r
+                        || r `Set.member` live_in_joins
 
-      -- multilive is a list of registers that are live in more than
-      -- one successor branch, and we should therefore drop them here.
-      multilive = [ r | (r,n) <- ufmToList livemap, n > 1 ]
-         where livemap = foldr (\r m -> addToUFM_C (+) m r (1::Int))
-                            emptyUFM live_in_nonjoins
+            live_sets' | should_drop = live_sets
+                       | otherwise   = map upd live_sets
 
-      (dropped_last, assigs') = dropAssignments drop_if assigs
+            upd set | r `Set.member` set = set `Set.union` live_rhs
+                    | otherwise          = set
 
-      drop_if a@(r,_,_)
-           =  a `conflicts` last
-           || (getUnique r `elem` multilive)
-           || (r `elem` live_in_joins)
+            live_rhs = foldRegsUsed extendRegSet emptyRegSet rhs
 
       final_middle = foldl blockSnoc middle' dropped_last
 
@@ -117,6 +136,11 @@ cmmSink graph = ofBlockList (g_entry graph) $ sink mapEmpty $ blocks
                  mapFromList [ (l, filterAssignments (getLive l) assigs')
                              | l <- succs ]
 
+-- tiny: an expression we don't mind duplicating
+isTiny :: CmmExpr -> Bool
+isTiny (CmmReg _) = True
+isTiny (CmmLit _) = True
+isTiny _other     = False
 
 -- annotate each node with the set of registers live *after* the node
 annotate :: RegSet -> [CmmNode O O] -> [(RegSet, CmmNode O O)]
@@ -143,18 +167,26 @@ walk ((live,node):ns) block as
        where usages :: UniqFM Int
              usages = foldRegsUsed addUsage emptyUFM node
 
-    (dropped, as') = dropAssignments (`conflicts` node1) as1
+    (dropped, as') = dropAssignmentsSimple (`conflicts` node1) as1
     block' = foldl blockSnoc block dropped `blockSnoc` node1
 
 tryToInline :: RegSet -> UniqFM Int -> CmmNode O x -> [Assignment]
             -> (CmmNode O x, [Assignment])
-tryToInline live usages node []
+tryToInline _live _usages node []
   = (node, [])
-tryToInline live usages node ((l,rhs,_) : rest)
-  | not (l `elemRegSet` live),
-    Just 1 <- lookupUFM usages l
-  = tryToInline live' usages' node' rest
+tryToInline live usages node (a@(l,rhs,_) : rest)
+  | occurs_once_in_this_node  = inline_and_discard
+  | False {- isTiny rhs -}    = inline_and_keep
+    -- ^^ seems to make things slightly worse
   where
+        inline_and_discard = tryToInline live' usages' node' rest
+
+        inline_and_keep = (node'', a : rest')
+          where (node'',rest') = inline_and_discard
+
+        occurs_once_in_this_node =
+         not (l `elemRegSet` live) &&  lookupUFM usages l == Just 1
+
         live'   = foldRegsUsed extendRegSet live rhs
         usages' = foldRegsUsed addUsage usages rhs
 
@@ -162,33 +194,40 @@ tryToInline live usages node ((l,rhs,_) : rest)
            where inline (CmmReg    (CmmLocal l'))     | l == l' = rhs
                  inline (CmmRegOff (CmmLocal l') off) | l == l' = cmmOffset rhs off
                  inline other = other
-tryToInline live usages node (assig : rest)
+tryToInline live usages node (assig@(_,rhs,_) : rest)
   = (node', assig : rest')
-  where (node', rest') = tryToInline live usages node rest
+  where (node', rest') = tryToInline live usages' node rest
+        usages' = foldRegsUsed addUsage usages rhs
 
 addUsage :: UniqFM Int -> LocalReg -> UniqFM Int
 addUsage m r = addToUFM_C (+) m r 1
 
 shouldSink :: CmmNode e x -> Maybe Assignment
 shouldSink (CmmAssign (CmmLocal r) e) | no_local_regs = Just (r, e, exprAddr e)
-  where no_local_regs = foldRegsUsed (\_ _ -> False) True e
+  where no_local_regs = True -- foldRegsUsed (\_ _ -> False) True e
 shouldSink _other = Nothing
 
 toNode :: Assignment -> CmmNode O O
 toNode (r,rhs,_) = CmmAssign (CmmLocal r) rhs
 
-dropAssignments :: (Assignment -> Bool) -> [Assignment] -> ([CmmNode O O], [Assignment])
-dropAssignments should_drop assigs
+dropAssignmentsSimple :: (Assignment -> Bool) -> [Assignment] -> ([CmmNode O O], [Assignment])
+dropAssignmentsSimple f = dropAssignments (\a _ -> (f a, ())) ()
+
+dropAssignments :: (Assignment -> s -> (Bool, s)) -> s -> [Assignment] -> ([CmmNode O O], [Assignment])
+dropAssignments should_drop state assigs
  = (dropped, reverse kept)
  where
-   (dropped,kept) = go assigs [] []
+   (dropped,kept) = go state assigs [] []
 
-   go []             dropped kept = (dropped, kept)
-   go (assig : rest) dropped kept
-      | conflict  = go rest (toNode assig : dropped) kept
-      | otherwise = go rest dropped (assig:kept)
+   go _ []             dropped kept = (dropped, kept)
+   go state (assig : rest) dropped kept
+      | conflict  = go state' rest (toNode assig : dropped) kept
+      | otherwise = go state' rest dropped (assig:kept)
       where
-        conflict = should_drop assig || any (assig `conflicts`) dropped
+        (dropit, state') = should_drop assig state
+        conflict = dropit || any (assig `conflicts`) dropped
+
+-- -----------------------------------------------------------------------------
 
 -- | @conflicts (r,e) stmt@ is @False@ if and only if the assignment
 -- @r = e@ can be safely commuted past @stmt@.