[Git][ghc/ghc][wip/T23479] 2 commits: Add STG debug from JS Sinker

Serge S. Gulin (@gulin.serge) gitlab at gitlab.haskell.org
Sat Sep 14 20:52:00 UTC 2024



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


Commits:
c878e3d4 by Serge S. Gulin at 2024-09-14T23:51:08+03:00
Add STG debug from JS Sinker

- - - - -
a547d899 by Serge S. Gulin at 2024-09-14T23:51:44+03:00
Add eager Sinker's strings unfloater

- - - - -


4 changed files:

- compiler/GHC/Driver/Flags.hs
- compiler/GHC/Driver/Session.hs
- compiler/GHC/StgToJS/CodeGen.hs
- compiler/GHC/StgToJS/Sinker.hs


Changes:

=====================================
compiler/GHC/Driver/Flags.hs
=====================================
@@ -470,6 +470,7 @@ data DumpFlag
    | Opt_D_dump_stg_cg        -- ^ STG (after stg2stg)
    | Opt_D_dump_stg_tags      -- ^ Result of tag inference analysis.
    | Opt_D_dump_stg_final     -- ^ Final STG (before cmm gen)
+   | Opt_D_dump_stg_from_js_sinker -- ^ STG after JS sinker
    | Opt_D_dump_call_arity
    | Opt_D_dump_exitify
    | Opt_D_dump_dmdanal


=====================================
compiler/GHC/Driver/Session.hs
=====================================
@@ -1515,6 +1515,8 @@ dynamic_flags_deps = [
         "Use `-ddump-stg-from-core` or `-ddump-stg-final` instead"
   , make_ord_flag defGhcFlag "ddump-stg-tags"
         (setDumpFlag Opt_D_dump_stg_tags)
+  , make_ord_flag defGhcFlag "ddump-stg-from-js-sinker"
+        (setDumpFlag Opt_D_dump_stg_tags)
   , make_ord_flag defGhcFlag "ddump-call-arity"
         (setDumpFlag Opt_D_dump_call_arity)
   , make_ord_flag defGhcFlag "ddump-exitify"


=====================================
compiler/GHC/StgToJS/CodeGen.hs
=====================================
@@ -11,7 +11,7 @@ where
 
 import GHC.Prelude
 
-import GHC.Driver.Flags (DumpFlag (Opt_D_dump_js))
+import GHC.Driver.Flags (DumpFlag (Opt_D_dump_js, Opt_D_dump_stg_from_js_sinker))
 
 import GHC.JS.Ppr
 import GHC.JS.JStg.Syntax
@@ -81,7 +81,8 @@ stgToJS logger config stg_binds0 this_mod spt_entries foreign_stubs cccs output_
     -- TODO: avoid top level lifting in core-2-core when the JS backend is
     -- enabled instead of undoing it here
 
-    -- TODO: add dump pass for optimized STG ast for JS
+  putDumpFileMaybe logger Opt_D_dump_stg_from_js_sinker "STG Optimized JS Sinker:" FormatSTG
+    (pprGenStgTopBindings (StgPprOpts False) stg_binds)
 
   (deps,lus) <- runG config this_mod unfloated_binds $ do
     ifProfilingM $ initCostCentres cccs


=====================================
compiler/GHC/StgToJS/Sinker.hs
=====================================
@@ -19,8 +19,8 @@ import GHC.Utils.Misc (partitionWith)
 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 +39,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 +192,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/-/compare/1a145d43264aa9c933fa4bbc1913c0aac28ea174...a547d8990b3ac40ba4a1a187bd7413aa36327d46

-- 
View it on GitLab: https://gitlab.haskell.org/ghc/ghc/-/compare/1a145d43264aa9c933fa4bbc1913c0aac28ea174...a547d8990b3ac40ba4a1a187bd7413aa36327d46
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/3af83fb0/attachment-0001.html>


More information about the ghc-commits mailing list