[Git][ghc/ghc][wip/romes/12935] fixup! determinism: Sampling uniques in the CG

Rodrigo Mesquita (@alt-romes) gitlab at gitlab.haskell.org
Tue Sep 17 09:05:18 UTC 2024



Rodrigo Mesquita pushed to branch wip/romes/12935 at Glasgow Haskell Compiler / GHC


Commits:
53d54e2b by Rodrigo Mesquita at 2024-09-17T10:04:56+01:00
fixup! determinism: Sampling uniques in the CG

- - - - -


2 changed files:

- compiler/GHC/Driver/GenerateCgIPEStub.hs
- compiler/GHC/Types/Unique/DSM.hs


Changes:

=====================================
compiler/GHC/Driver/GenerateCgIPEStub.hs
=====================================
@@ -37,7 +37,6 @@ import GHC.Unit.Module (moduleNameString)
 import qualified GHC.Utils.Logger as Logger
 import GHC.Utils.Outputable (ppr)
 import GHC.Types.Unique.DSM
-import qualified Control.Monad.Trans.State.Strict as T
 
 {-
 Note [Stacktraces from Info Table Provenance Entries (IPE based stack unwinding)]
@@ -216,15 +215,15 @@ generateCgIPEStub hsc_env this_mod denv (nonCaffySet, moduleLFInfos, infoTablesW
   -- of the code!  Instead, make sure all labels generated for IPE related code
   -- sources uniques from the DUniqSupply gotten from CgStream (see its use in
   -- initInfoTableProv/emitIpeBufferListNode).
-  dus <- liftEff $ UDSMT $ T.get
+  (mIpeStub, ipeCmmGroup) <- liftEff $ UDSMT $ \dus -> do
 
-  -- Yield Cmm for Info Table Provenance Entries (IPEs)
-  let denv' = denv {provInfoTables = Map.mapKeys cit_lbl infoTablesWithTickishes}
-      (((mIpeStub, dus'), ipeCmmGroup), _) =
-          runC (initStgToCmmConfig dflags this_mod) fstate cgState $
-            getCmm (initInfoTableProv initStats (Map.keys infoTablesWithTickishes) denv' dus)
+    -- Yield Cmm for Info Table Provenance Entries (IPEs)
+    let denv' = denv {provInfoTables = Map.mapKeys cit_lbl infoTablesWithTickishes}
+        (((mIpeStub, dus'), ipeCmmGroup), _) =
+            runC (initStgToCmmConfig dflags this_mod) fstate cgState $
+              getCmm (initInfoTableProv initStats (Map.keys infoTablesWithTickishes) denv' dus)
 
-  liftEff $ UDSMT $ T.put dus'
+    return ((mIpeStub, ipeCmmGroup), dus')
 
   (_, ipeCmmGroupSRTs) <- liftEff $ withDUS $ cmmPipeline logger cmm_cfg (emptySRT this_mod) (removeDeterm ipeCmmGroup)
   Stream.yield ipeCmmGroupSRTs


=====================================
compiler/GHC/Types/Unique/DSM.hs
=====================================
@@ -8,7 +8,6 @@ import Control.Monad.Fix
 import GHC.Types.Unique
 import qualified GHC.Utils.Monad.State.Strict as Strict
 import qualified GHC.Types.Unique.Supply as USM
-import qualified Control.Monad.Trans.State.Strict as T
 import Control.Monad.IO.Class
 
 {-
@@ -129,48 +128,68 @@ instance MonadGetUnique USM.UniqSM where
 -- | Transformer version of 'UniqDSM' to use when threading a deterministic
 -- uniq supply over a Monad. Specifically, it is used in the `Stream` of Cmm
 -- decls.
-newtype UniqDSMT m result = UDSMT { unUDSMT :: T.StateT DUniqSupply m (result) }
-  deriving (Functor, Applicative, Monad, MonadIO)
+newtype UniqDSMT m result = UDSMT' { unUDSMT :: DUniqSupply -> m (result, DUniqSupply) }
+  deriving (Functor)
+
+-- Similar to GHC.Utils.Monad.State.Strict, using Note [The one-shot state monad trick]
+
+pattern UDSMT :: (DUniqSupply -> m (result, DUniqSupply)) -> UniqDSMT m result
+pattern UDSMT m <- UDSMT' m
+  where
+    UDSMT m = UDSMT' (oneShot $ \s -> m s)
+{-# COMPLETE UDSMT #-}
+
+instance Monad m => Applicative (UniqDSMT m) where
+  pure x = UDSMT $ \s -> pure (x, s)
+  UDSMT f <*> UDSMT x = UDSMT $ \s0 -> do
+    (f', s1) <- f s0
+    (x', s2) <- x s1
+    pure (f' x', s2)
+
+instance Monad m => Monad (UniqDSMT m) where
+  UDSMT x >>= f = UDSMT $ \s0 -> do
+    (x', s1) <- x s0
+    case f x' of UDSMT y -> y s1
+
+instance MonadIO m => MonadIO (UniqDSMT m) where
+  liftIO x = UDSMT $ \s -> (,s) <$> liftIO x
 
 instance Monad m => MonadGetUnique (UniqDSMT m) where
-  getUniqueM = UDSMT $ do
-    us <- T.get
+  getUniqueM = UDSMT $ \us -> do
     let (u, us') = takeUniqueFromDSupply us
-    T.put us'
-    return u
+    return (u, us')
 
 -- | Set the tag of the running @UniqDSMT@ supply to the given tag and run an action with it.
 -- All uniques produced in the given action will use this tag, until the tag is changed
 -- again.
 setTagUDSMT :: Monad m => Char {-^ Tag -} -> UniqDSMT m a -> UniqDSMT m a
-setTagUDSMT tag act = do
-  origtag <- UDSMT (T.gets getTagDUniqSupply)
-  UDSMT (T.modify (newTagDUniqSupply tag))
-  a <- act
-  UDSMT (T.modify (newTagDUniqSupply origtag)) -- restore original tag
-  return a
+setTagUDSMT tag (UDSMT act) = UDSMT $ \us -> do
+  let origtag = getTagDUniqSupply us
+      new_us  = newTagDUniqSupply tag us
+  (a, us') <- act new_us
+  let us'_origtag = newTagDUniqSupply origtag us'
+      -- restore original tag
+  return (a, us'_origtag)
 
 -- | Like 'runUniqueDSM' but for 'UniqDSMT'
 runUDSMT :: DUniqSupply -> UniqDSMT m a -> m (a, DUniqSupply)
-runUDSMT dus (UDSMT st) = T.runStateT st dus
+runUDSMT dus (UDSMT st) = st dus
 
 -- | Lift an IO action that depends on, and threads through, a unique supply
 -- into UniqDSMT IO.
 withDUS :: (DUniqSupply -> IO (a, DUniqSupply)) -> UniqDSMT IO a
-withDUS f = UDSMT $ do
-  us <- T.get
+withDUS f = UDSMT $ \us -> do
   (a, us') <- liftIO (f us)
-  _  <- T.put us'
-  return a
+  return (a, us')
 
 -- | Change the monad underyling an applied @UniqDSMT@, i.e. transform a
 -- @UniqDSMT m@ into a @UniqDSMT n@ given @m ~> n at .
 hoistUDSMT :: (forall x. m x -> n x) -> UniqDSMT m a -> UniqDSMT n a
-hoistUDSMT nt (UDSMT s) = UDSMT $ T.mapStateT nt s
+hoistUDSMT nt (UDSMT m) = UDSMT $ \s -> nt (m s)
 
 -- | Lift a monadic action @m a@ into an @UniqDSMT m a@
 liftUDSMT :: Functor m => m a -> UniqDSMT m a
-liftUDSMT m = UDSMT $ T.StateT $ \s -> (,s) <$> m
+liftUDSMT m = UDSMT $ \s -> (,s) <$> m
 
 --------------------------------------------------------------------------------
 -- MonadUniqDSM
@@ -184,8 +203,4 @@ instance MonadUniqDSM UniqDSM where
   liftUniqDSM = id
 
 instance Monad m => MonadUniqDSM (UniqDSMT m) where
-  liftUniqDSM act = do
-    us <- UDSMT T.get
-    let (a, us') = runUniqueDSM us act
-    UDSMT (T.put us')
-    return a
+  liftUniqDSM act = UDSMT $ \us -> pure $ runUniqueDSM us act



View it on GitLab: https://gitlab.haskell.org/ghc/ghc/-/commit/53d54e2b2fe19d93b3ac92690cf6f6afb5054996

-- 
View it on GitLab: https://gitlab.haskell.org/ghc/ghc/-/commit/53d54e2b2fe19d93b3ac92690cf6f6afb5054996
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/20240917/5b66fb7f/attachment-0001.html>


More information about the ghc-commits mailing list