[Git][ghc/ghc][wip/T23479] Add eager Sinker's strings unfloater

Serge S. Gulin (@gulin.serge) gitlab at gitlab.haskell.org
Sat Sep 14 21:45:45 UTC 2024



Serge S. Gulin pushed to branch wip/T23479 at Glasgow Haskell Compiler / GHC


Commits:
2b4abb88 by Serge S. Gulin at 2024-09-15T00:45:29+03:00
Add eager Sinker's strings unfloater

- - - - -


1 changed file:

- compiler/GHC/StgToJS/Sinker.hs


Changes:

=====================================
compiler/GHC/StgToJS/Sinker.hs
=====================================
@@ -21,6 +21,7 @@ import GHC.StgToJS.Utils
 import Data.Char
 import Data.List (partition)
 import Data.Maybe
+import Data.Bifunctor (Bifunctor (..))
 import GHC.Utils.Outputable (showSDocUnsafe, showPprUnsafe)
 import qualified GHC.Utils.Trace as DT
 import Data.ByteString (ByteString)
@@ -39,40 +40,21 @@ sinkPgm :: Module
         -> (UniqFM Id CgStgExpr, [CgStgTopBinding])
 sinkPgm m pgm
   = ( (DT.trace . (++) "<sunk>" . showPprUnsafe) sunk sunk
-    , ((DT.trace . (++) "<pgm''>" . showSDocUnsafe . pprGenStgTopBindings shortStgPprOpts) (map StgTopLifted pgm'') (map StgTopLifted pgm''))
-      ++ ((DT.trace . (++) "<stringLits>" . showSDocUnsafe . pprGenStgTopBindings shortStgPprOpts) stringLits stringLits)
+    , traceDoc "<pgm'''>" (map StgTopLifted pgm''') (map StgTopLifted pgm''')
+      ++ traceDoc "<stringLits>" stringLits stringLits
     )
   where
-    -- selectLifted :: CgStgTopBinding -> Either CgStgBinding (Id, ByteString)
-    selectLifted :: CgStgTopBinding -> Either CgStgBinding CgStgTopBinding
+    traceDoc head = DT.trace . (++) head . showSDocUnsafe . pprGenStgTopBindings shortStgPprOpts
+
+    selectLifted :: CgStgTopBinding -> Either CgStgBinding (Id, ByteString)
     selectLifted (StgTopLifted b)      = Left b
-    -- selectLifted (StgTopStringLit i b) = Right (i, b)
-    selectLifted x = Right x
+    selectLifted (StgTopStringLit i b) = Right (i, b)
 
-    selectLiftedInCase :: [CgStgBinding] -> (Id, ByteString) -> Either (Id, ByteString) (Id, ByteString)
-    selectLiftedInCase bindings (i, _) = undefined
-      where
-        isUsedInCase :: CgStgBinding -> (Id, ByteString) -> Either (Id, ByteString) (Id, ByteString)
-        -- Recursive expressions are not a target
-        isUsedInCase (StgRec _) x = Left x
-        -- Constructors should be left as is
-        isUsedInCase (StgNonRec _ (StgRhsCon {})) x = Left x
-        -- If function has no arguments it is not a target
-        isUsedInCase (StgNonRec i' (StgRhsClosure _ _ _ [] _ _)) x = Left x
-        -- If function has more arguments than 1 it is not a target
-        isUsedInCase (StgNonRec i' (StgRhsClosure _ _ _ (_ : (_ : _)) _ _)) x = Left x
-        -- If function has single argument probably it uses a string
-        isUsedInCase (StgNonRec i' (StgRhsClosure _ _ _ [arg1] _ _)) x
-          | arg1 == i = undefined
-          | otherwise = Left x
-
-    selectLiftedToUnlift :: [CgStgTopBinding] -> ([CgStgTopBinding], [CgStgTopBinding])
-    selectLiftedToUnlift = undefined
-      where
-        xxx = undefined
+    (pgm', stringLits') = partitionWith selectLifted pgm
+    stringLits = uncurry StgTopStringLit <$> stringLits'
 
-    (pgm', stringLits) = partitionWith selectLifted pgm
-    (sunk, pgm'')      = sinkPgm' m $ (DT.trace . (++) "<pgm'>" . showSDocUnsafe . pprGenStgTopBindings shortStgPprOpts) (map StgTopLifted pgm') pgm'
+    (pgm'', _usedStringLitNames) = unfloatStringLits (listToUFM $ first idName <$> stringLits') pgm'
+    (sunk, pgm''') = sinkPgm' m $ (DT.trace . (++) "<pgm'>" . showSDocUnsafe . pprGenStgTopBindings shortStgPprOpts) (map StgTopLifted pgm'') pgm''
 
 sinkPgm'
   :: Module
@@ -211,3 +193,112 @@ topSortDecls _m binds = rest ++ nr'
     nr' | (not . null) [()| CyclicSCC _ <- stronglyConnCompG g]
             = error "topSortDecls: unexpected cycle"
         | otherwise = map node_payload (topologicalSortG g)
+
+-- | We are doing attempts to unfloat string literals back to
+-- the call site. Further special JS optimizations
+-- can generate more performant operations over them.
+unfloatStringLits :: UniqFM Name ByteString -> [CgStgBinding] -> ([CgStgBinding], UniqSet Name)
+unfloatStringLits stringLits allBindings = (binderWithoutChanges ++ binderWithUnfloatedStringLit, actuallyUsedStringLitNames)
+  where
+    (binderWithoutChanges, binderWithUnfloatedStringLitPairs) = partitionWith substituteStringLit allBindings
+
+    binderWithUnfloatedStringLit = fst <$> binderWithUnfloatedStringLitPairs
+    actuallyUsedStringLitNames = unionManyUniqSets (snd <$> binderWithUnfloatedStringLitPairs)
+
+    substituteStringLit :: CgStgBinding -> Either CgStgBinding (CgStgBinding, UniqSet Name)
+    substituteStringLit x@(StgRec bnds)
+      | isEmptyUniqSet names = Left x
+      | otherwise = Right (StgRec bnds', names)
+      where
+        (bnds', names) = extractNames id $ do
+          (i, rhs) <- bnds
+          pure $ case processStgRhs rhs of
+            Nothing -> Left (i, rhs)
+            Just (rhs', names) -> Right ((i, rhs'), names)
+    substituteStringLit x@(StgNonRec binder rhs)
+      = maybe (Left x)
+        (\(body', names) -> Right (StgNonRec binder body', names))
+        (processStgRhs rhs)
+
+    processStgRhs :: CgStgRhs -> Maybe (CgStgRhs, UniqSet Name)
+    processStgRhs (StgRhsCon ccs dataCon mu ticks args typ)
+      | isEmptyUniqSet names = Nothing
+      | otherwise = Just (StgRhsCon ccs dataCon mu ticks unified typ, names)
+      where
+        (unified, names) = substituteArgWithNames args
+    processStgRhs (StgRhsClosure fvs ccs upd bndrs body typ)
+      = (\(body', names) -> (StgRhsClosure fvs ccs upd bndrs body' typ, names)) <$>
+        processStgExpr body
+
+    -- Recursive expressions
+    processStgExpr :: CgStgExpr -> Maybe (CgStgExpr, UniqSet Name)
+    processStgExpr (StgLit _) = Nothing
+    processStgExpr (StgTick _ _) = Nothing
+    processStgExpr (StgLet n b e) =
+      case (substituteStringLit b, processStgExpr e) of
+        (Left _, Nothing) -> Nothing
+        (Right (b', names), Nothing) -> Just (StgLet n b' e, names)
+        (Left _, Just (e', names)) -> Just (StgLet n b e', names)
+        (Right (b', names), Just (e', names')) -> Just (StgLet n b' e', names `unionUniqSets` names')
+    processStgExpr (StgLetNoEscape n b e) =
+      case (substituteStringLit b, processStgExpr e) of
+        (Left _, Nothing) -> Nothing
+        (Right (b', names), Nothing) -> Just (StgLetNoEscape n b' e, names)
+        (Left _, Just (e', names)) -> Just (StgLetNoEscape n b e', names)
+        (Right (b', names), Just (e', names')) -> Just (StgLetNoEscape n b' e', names `unionUniqSets` names')
+    -- We should keep the order: See Note [Case expression invariants]
+    processStgExpr (StgCase e bndr alt_type alts) =
+      case (isEmptyUniqSet names, processStgExpr e) of
+        (True, Nothing) -> Nothing
+        (True, Just (e', names')) -> Just (StgCase e' bndr alt_type alts, names')
+        (False, Nothing) -> Just (StgCase e bndr alt_type unified, names)
+        (False, Just (e', names')) -> Just (StgCase e' bndr alt_type unified, names `unionUniqSets` names')
+      where
+        (unified, names) = extractNames splitAlts alts
+
+        splitAlts :: CgStgAlt -> Either CgStgAlt (CgStgAlt, UniqSet Name)
+        splitAlts alt@(GenStgAlt con bndrs rhs) =
+          case processStgExpr rhs of
+            Nothing -> Left alt
+            Just (alt', names) -> Right (GenStgAlt con bndrs alt', names)
+
+    -- No args
+    processStgExpr (StgApp _ []) = Nothing
+    processStgExpr (StgConApp _ _ [] _) = Nothing
+    processStgExpr (StgOpApp _ [] _) = Nothing
+
+    -- Main targets. Preserving the order of args is important
+    processStgExpr (StgApp fn args@(_:_))
+      | isEmptyUniqSet names = Nothing
+      | otherwise = Just (StgApp fn unified, names)
+      where
+        (unified, names) = substituteArgWithNames args
+    processStgExpr (StgConApp dc n args@(_:_) tys)
+      | isEmptyUniqSet names = Nothing
+      | otherwise = Just (StgConApp dc n unified tys, names)
+      where
+        (unified, names) = substituteArgWithNames args
+    processStgExpr (StgOpApp op args@(_:_) tys)
+      | isEmptyUniqSet names = Nothing
+      | otherwise = Just (StgOpApp op unified tys, names)
+      where
+        (unified, names) = substituteArgWithNames args
+
+    substituteArg :: StgArg -> Either StgArg (StgArg, Name)
+    substituteArg a@(StgLitArg _) = Left a
+    substituteArg a@(StgVarArg i) =
+      let name = idName i
+      in case lookupUFM stringLits name of
+        Nothing -> Left a
+        Just b -> Right (StgLitArg $ LitString b, name)
+
+    substituteArgWithNames = extractNames (second (second unitUniqSet) . substituteArg)
+
+    extractNames :: (a -> Either x (x, UniqSet Name)) -> [a] -> ([x], UniqSet Name)
+    extractNames splitter target =
+      let
+        splitted = splitter <$> target
+        combined = either (, emptyUniqSet) id <$> splitted
+        unified = fst <$> combined
+        names = unionManyUniqSets (snd <$> combined)
+      in (unified, names)



View it on GitLab: https://gitlab.haskell.org/ghc/ghc/-/commit/2b4abb8812cd4301aeb5da3623d3c90d3e2c6fe8

-- 
View it on GitLab: https://gitlab.haskell.org/ghc/ghc/-/commit/2b4abb8812cd4301aeb5da3623d3c90d3e2c6fe8
You're receiving this email because of your account on gitlab.haskell.org.


-------------- next part --------------
An HTML attachment was scrubbed...
URL: <http://mail.haskell.org/pipermail/ghc-commits/attachments/20240914/ece211b9/attachment-0001.html>


More information about the ghc-commits mailing list