[Git][ghc/ghc][wip/romes/12935] Major progress in using UniqDSM in CmmToAsm and Ncg backends

Rodrigo Mesquita (@alt-romes) gitlab at gitlab.haskell.org
Wed Jul 3 10:36:13 UTC 2024



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


Commits:
f8ae759b by Rodrigo Mesquita at 2024-07-03T11:34:47+01:00
Major progress in using UniqDSM in CmmToAsm and Ncg backends

- - - - -


28 changed files:

- compiler/GHC/Cmm/BlockId.hs
- compiler/GHC/Cmm/Reducibility.hs
- compiler/GHC/Cmm/UniqueRenamer.hs
- compiler/GHC/CmmToAsm.hs
- compiler/GHC/CmmToAsm/AArch64/CodeGen.hs
- compiler/GHC/CmmToAsm/AArch64/Instr.hs
- compiler/GHC/CmmToAsm/BlockLayout.hs
- compiler/GHC/CmmToAsm/Dwarf.hs
- compiler/GHC/CmmToAsm/Monad.hs
- compiler/GHC/CmmToAsm/PPC/Instr.hs
- compiler/GHC/CmmToAsm/Reg/Graph.hs
- compiler/GHC/CmmToAsm/Reg/Graph/Spill.hs
- compiler/GHC/CmmToAsm/Reg/Linear.hs
- compiler/GHC/CmmToAsm/Reg/Linear/Base.hs
- compiler/GHC/CmmToAsm/Reg/Linear/State.hs
- compiler/GHC/CmmToAsm/Reg/Liveness.hs
- compiler/GHC/CmmToAsm/Wasm.hs
- compiler/GHC/CmmToAsm/Wasm/FromCmm.hs
- compiler/GHC/CmmToAsm/Wasm/Types.hs
- compiler/GHC/CmmToAsm/X86/CodeGen.hs
- compiler/GHC/CmmToAsm/X86/Instr.hs
- compiler/GHC/CmmToLlvm/Base.hs
- compiler/GHC/Data/Graph/Collapse.hs
- compiler/GHC/Driver/CodeOutput.hs
- compiler/GHC/StgToCmm/ExtCode.hs
- + compiler/GHC/Types/Unique/DSM.hs
- compiler/GHC/Wasm/ControlFlow/FromCmm.hs
- compiler/ghc.cabal.in


Changes:

=====================================
compiler/GHC/Cmm/BlockId.hs
=====================================
@@ -16,6 +16,7 @@ import GHC.Types.Id.Info
 import GHC.Types.Name
 import GHC.Types.Unique
 import GHC.Types.Unique.Supply
+import qualified GHC.Types.Unique.DSM as DSM
 
 import GHC.Cmm.Dataflow.Label (Label, mkHooplLabel)
 
@@ -36,8 +37,12 @@ type BlockId = Label
 mkBlockId :: Unique -> BlockId
 mkBlockId unique = mkHooplLabel $ getKey unique
 
-newBlockId :: MonadUnique m => m BlockId
-newBlockId = mkBlockId <$> getUniqueM
+-- If the monad unique instance uses a deterministic unique supply, this will
+-- give you a deterministic unique. Otherwise, it will not. Note that from Cmm
+-- onwards (after deterministic renaming in 'codeGen'), there should only exist
+-- deterministic block labels.
+newBlockId :: DSM.MonadGetUnique m => m BlockId
+newBlockId = mkBlockId <$> DSM.getUniqueM
 
 blockLbl :: BlockId -> CLabel
 blockLbl label = mkLocalBlockLabel (getUnique label)


=====================================
compiler/GHC/Cmm/Reducibility.hs
=====================================
@@ -47,7 +47,7 @@ import GHC.Cmm.Dataflow.Label
 import GHC.Data.Graph.Collapse
 import GHC.Data.Graph.Inductive.Graph
 import GHC.Data.Graph.Inductive.PatriciaTree
-import GHC.Types.Unique.Supply
+import GHC.Types.Unique.DSM
 import GHC.Utils.Panic
 
 -- | Represents the result of a reducibility analysis.
@@ -81,7 +81,7 @@ reducibility gwd =
 -- control-flow graph.
 
 asReducible :: GraphWithDominators CmmNode
-            -> UniqSM (GraphWithDominators CmmNode)
+            -> UniqDSM (GraphWithDominators CmmNode)
 asReducible gwd = case reducibility gwd of
                     Reducible -> return gwd
                     Irreducible -> assertReducible <$> nodeSplit gwd
@@ -97,7 +97,7 @@ assertReducible gwd = case reducibility gwd of
 -- irreducible.
 
 nodeSplit :: GraphWithDominators CmmNode
-          -> UniqSM (GraphWithDominators CmmNode)
+          -> UniqDSM (GraphWithDominators CmmNode)
 nodeSplit gwd =
     graphWithDominators <$> inflate (g_entry g) <$> runNullCollapse collapsed
   where g = gwd_graph gwd
@@ -181,7 +181,7 @@ instance PureSupernode CmmSuper where
   mapLabels = changeLabels
 
 instance Supernode CmmSuper NullCollapseViz where
-  freshen s = liftUniqSM $ relabel s
+  freshen s = liftUniqDSM $ relabel s
 
 
 -- | Return all labels defined within a supernode.
@@ -212,11 +212,11 @@ changeBlockLabels f block = blockJoin entry' middle exit'
 -- | Within the given supernode, replace every defined label (and all
 -- of its uses) with a fresh label.
 
-relabel :: CmmSuper -> UniqSM CmmSuper
+relabel :: CmmSuper -> UniqDSM CmmSuper
 relabel node = do
      finite_map <- foldM addPair mapEmpty $ definedLabels node
      return $ changeLabels (labelChanger finite_map) node
-  where addPair :: LabelMap Label -> Label -> UniqSM (LabelMap Label)
+  where addPair :: LabelMap Label -> Label -> UniqDSM (LabelMap Label)
         addPair map old = do new <- newBlockId
                              return $ mapInsert old new map
         labelChanger :: LabelMap Label -> (Label -> Label)


=====================================
compiler/GHC/Cmm/UniqueRenamer.hs
=====================================
@@ -7,6 +7,8 @@ module GHC.Cmm.UniqueRenamer
 
   -- Careful! Not for general use!
   , DetUniqFM, emptyDetUFM
+
+  , module GHC.Types.Unique.DSM
   )
   where
 
@@ -28,6 +30,7 @@ import GHC.Utils.Outputable as Outputable
 import Data.Tuple (swap)
 import GHC.Types.Id
 import qualified GHC.Types.Unique.Supply as USM
+import GHC.Types.Unique.DSM
 
 {-
 --------------------------------------------------------------------------------
@@ -276,93 +279,3 @@ instance (UniqRenamable a) => UniqRenamable (Maybe a) where
 panicMapKeysNotInjective :: a -> b -> c
 panicMapKeysNotInjective _ _ = error "this should be impossible because the function which maps keys should be injective"
 
---------------------------------------------------------------------------------
--- UniqDSM (ToDo: For this to make sense in this module, rename the module to
--- something like GHC.Cmm.UniqueDeterminism). Write notes....
-
--- todo: Do I need to use the one-shot state monad trick? Probably yes.
-
--- check: UniqSM is only used before Cmm (grep for it), afterwards only UniqDSM is used.
-
--- todo: use UniqSM for UniqRenamable? We've basically re-implemented this logic
--- there, but without the unboxing it feels? Maybe not, since we carry the
--- mappings too.
-
-newtype DUniqSupply = DUS Word64 -- supply uniques iteratively
-type DUniqResult result = (# result, DUniqSupply #)
-
-pattern DUniqResult :: a -> b -> (# a, b #)
-pattern DUniqResult x y = (# x, y #)
-{-# COMPLETE DUniqResult #-}
-
--- | A monad which just gives the ability to obtain 'Unique's deterministically.
--- There's no splitting.
-newtype UniqDSM result = UDSM { unUDSM :: DUniqSupply -> DUniqResult result }
-  deriving Functor
-
-instance Monad UniqDSM where
-  (>>=) (UDSM f) cont = UDSM $ \us0 -> case f us0 of
-    DUniqResult result us1 -> unUDSM (cont result) us1
-  (>>)  = (*>)
-  {-# INLINE (>>=) #-}
-  {-# INLINE (>>) #-}
-
-instance Applicative UniqDSM where
-  pure result = UDSM (DUniqResult result)
-  (UDSM f) <*> (UDSM x) = UDSM $ \us0 -> case f us0 of
-    DUniqResult ff us1 -> case x us1 of
-      DUniqResult xx us2 -> DUniqResult (ff xx) us2
-  (*>) (UDSM expr) (UDSM cont) = UDSM $ \us0 -> case expr us0 of
-    DUniqResult _ us1 -> cont us1
-  {-# INLINE pure #-}
-  {-# INLINE (*>) #-}
-
-instance MonadFix UniqDSM where
-  mfix m = UDSM (\us0 -> let (r,us1) = runUniqueDSM us0 (m r) in DUniqResult r us1)
-
-getUniqueDSM :: UniqDSM Unique
-getUniqueDSM = UDSM (\(DUS us0) -> DUniqResult (mkUniqueGrimily us0) (DUS $ us0+1))
-
-takeUniqueFromDSupply :: DUniqSupply -> (Unique, DUniqSupply)
-takeUniqueFromDSupply d =
-  case unUDSM getUniqueDSM d of
-    DUniqResult x y -> (x, y)
-
--- Write Note about the importance of locality in uniques that are deterministic
---
--- If you use a tag which collides with other names, you'll get a uniques
--- deterministically colliding with existing symbols.
---
--- (e.g. easy to observe if you do this wrong)
---
--- Ideally, we'd thread the same deterministic unique supply all the way
--- throughout the Cmm pipeline, starting off from hte deterministic rename
--- pass.
-initDUniqSupply :: Char -> Word64 -> DUniqSupply
-initDUniqSupply c firstUniq =
-  let !tag = mkTag c
-  in DUS (tag .|. firstUniq)
-
-newTagDUniqSupply :: Char -> DUniqSupply -> DUniqSupply
-newTagDUniqSupply c (DUS w) = DUS $ getKey $ newTagUnique (mkUniqueGrimily w) c
-
-runUniqueDSM :: DUniqSupply -> UniqDSM a -> (a, DUniqSupply)
-runUniqueDSM ds (UDSM f) =
-  case f ds of
-    DUniqResult uq us -> (uq, us)
-
-class Monad m => MonadGetUnique m where
-  getUniqueM :: m Unique
-
-instance MonadGetUnique UniqDSM where
-  getUniqueM = getUniqueDSM
-
-instance MonadGetUnique USM.UniqSM where
-  getUniqueM = USM.getUniqueM
-
-{-
-Note [Cmm Local Deterministic Uniques]
-~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
-TODO!!!!!
-TODO!!!!!
--}


=====================================
compiler/GHC/CmmToAsm.hs
=====================================
@@ -100,7 +100,7 @@ import GHC.Cmm.GenericOpt
 import GHC.Cmm.CLabel
 
 import GHC.Types.Unique.FM
-import GHC.Types.Unique.Supply
+import GHC.Types.Unique.DSM
 import GHC.Driver.DynFlags
 import GHC.Driver.Ppr
 import GHC.Utils.Misc
@@ -129,7 +129,7 @@ import System.IO
 import System.Directory ( getCurrentDirectory )
 
 --------------------
-nativeCodeGen :: forall a . Logger -> ToolSettings -> NCGConfig -> ModLocation -> Handle -> UniqSupply
+nativeCodeGen :: forall a . Logger -> ToolSettings -> NCGConfig -> ModLocation -> Handle -> DUniqSupply
               -> Stream IO RawCmmGroup a
               -> IO a
 nativeCodeGen logger ts config modLoc h us cmms
@@ -203,7 +203,7 @@ nativeCodeGen' :: (OutputableP Platform statics, Outputable jumpDest, Instructio
                -> ModLocation
                -> NcgImpl statics instr jumpDest
                -> Handle
-               -> UniqSupply
+               -> DUniqSupply
                -> Stream IO RawCmmGroup a
                -> IO a
 nativeCodeGen' logger config modLoc ncgImpl h us cmms
@@ -223,9 +223,9 @@ finishNativeGen :: Instruction instr
                 -> NCGConfig
                 -> ModLocation
                 -> BufHandle
-                -> UniqSupply
+                -> DUniqSupply
                 -> NativeGenAcc statics instr
-                -> IO UniqSupply
+                -> IO DUniqSupply
 finishNativeGen logger config modLoc bufh us ngs
  = withTimingSilent logger (text "NCG") (`seq` ()) $ do
         -- Write debug data and finish
@@ -284,19 +284,19 @@ cmmNativeGenStream :: forall statics jumpDest instr a . (OutputableP Platform st
               -> ModLocation
               -> NcgImpl statics instr jumpDest
               -> BufHandle
-              -> UniqSupply
+              -> DUniqSupply
               -> Stream.Stream IO RawCmmGroup a
               -> NativeGenAcc statics instr
-              -> IO (NativeGenAcc statics instr, UniqSupply, a)
+              -> IO (NativeGenAcc statics instr, DUniqSupply, a)
 
 cmmNativeGenStream logger config modLoc ncgImpl h us cmm_stream ngs
  = loop us (Stream.runStream cmm_stream) ngs
   where
     ncglabel = text "NCG"
-    loop :: UniqSupply
+    loop :: DUniqSupply
               -> Stream.StreamS IO RawCmmGroup a
               -> NativeGenAcc statics instr
-              -> IO (NativeGenAcc statics instr, UniqSupply, a)
+              -> IO (NativeGenAcc statics instr, DUniqSupply, a)
     loop us s ngs =
       case s of
         Stream.Done a ->
@@ -345,17 +345,17 @@ cmmNativeGens :: forall statics instr jumpDest.
               -> NcgImpl statics instr jumpDest
               -> BufHandle
               -> LabelMap DebugBlock
-              -> UniqSupply
+              -> DUniqSupply
               -> [RawCmmDecl]
               -> NativeGenAcc statics instr
               -> Int
-              -> IO (NativeGenAcc statics instr, UniqSupply)
+              -> IO (NativeGenAcc statics instr, DUniqSupply)
 
 cmmNativeGens logger config ncgImpl h dbgMap = go
   where
-    go :: UniqSupply -> [RawCmmDecl]
+    go :: DUniqSupply -> [RawCmmDecl]
        -> NativeGenAcc statics instr -> Int
-       -> IO (NativeGenAcc statics instr, UniqSupply)
+       -> IO (NativeGenAcc statics instr, DUniqSupply)
 
     go us [] ngs !_ =
         return (ngs, us)
@@ -424,12 +424,12 @@ cmmNativeGen
     :: forall statics instr jumpDest. (Instruction instr, OutputableP Platform statics, Outputable jumpDest)
     => Logger
     -> NcgImpl statics instr jumpDest
-        -> UniqSupply
+        -> DUniqSupply
         -> DwarfFiles
         -> LabelMap DebugBlock
         -> RawCmmDecl                                   -- ^ the cmm to generate code for
         -> Int                                          -- ^ sequence number of this top thing
-        -> IO   ( UniqSupply
+        -> IO   ( DUniqSupply
                 , DwarfFiles
                 , [NatCmmDecl statics instr]                -- native code
                 , [CLabel]                                  -- things imported by this cmm
@@ -468,7 +468,7 @@ cmmNativeGen logger ncgImpl us fileIds dbgMap cmm count
         -- generate native code from cmm
         let ((native, lastMinuteImports, fileIds', nativeCfgWeights), usGen) =
                 {-# SCC "genMachCode" #-}
-                initUs us $ genMachCode config
+                runUniqueDSM us $ genMachCode config
                                         (cmmTopCodeGen ncgImpl)
                                         fileIds dbgMap opt_cmm cmmCfg
 
@@ -486,7 +486,7 @@ cmmNativeGen logger ncgImpl us fileIds dbgMap cmm count
                                 else Nothing
         let (withLiveness, usLive) =
                 {-# SCC "regLiveness" #-}
-                initUs usGen
+                runUniqueDSM usGen
                         $ mapM (cmmTopLiveness livenessCfg platform) native
 
         putDumpFileMaybe logger
@@ -494,8 +494,6 @@ cmmNativeGen logger ncgImpl us fileIds dbgMap cmm count
                 FormatCMM
                 (vcat $ map (pprLiveCmmDecl platform) withLiveness)
 
-        -- ROMES:TODO: RENAME VIRTUAL REGISTERS DETERMINISTICALLY
-
         -- allocate registers
         (alloced, usAlloc, ppr_raStatsColor, ppr_raStatsLinear, raStats, stack_updt_blks) <-
          if ( ncgRegsGraph config || ncgRegsIterative config )
@@ -510,7 +508,7 @@ cmmNativeGen logger ncgImpl us fileIds dbgMap cmm count
                 -- do the graph coloring register allocation
                 let ((alloced, maybe_more_stack, regAllocStats), usAlloc)
                         = {-# SCC "RegAlloc-color" #-}
-                          initUs usLive
+                          runUniqueDSM usLive
                           $ Color.regAlloc
                                 config
                                 alloc_regs
@@ -520,13 +518,13 @@ cmmNativeGen logger ncgImpl us fileIds dbgMap cmm count
                                 livenessCfg
 
                 let ((alloced', stack_updt_blks), usAlloc')
-                        = initUs usAlloc $
-                                case maybe_more_stack of
-                                Nothing     -> return (alloced, [])
-                                Just amount -> do
-                                    (alloced',stack_updt_blks) <- unzip <$>
-                                                (mapM ((ncgAllocMoreStack ncgImpl) amount) alloced)
-                                    return (alloced', concat stack_updt_blks )
+                        = runUniqueDSM usAlloc $
+                            case maybe_more_stack of
+                            Nothing     -> return (alloced, [])
+                            Just amount -> do
+                                (alloced',stack_updt_blks) <- unzip <$>
+                                            (mapM ((ncgAllocMoreStack ncgImpl) amount) alloced)
+                                return (alloced', concat stack_updt_blks )
 
 
                 -- dump out what happened during register allocation
@@ -571,7 +569,7 @@ cmmNativeGen logger ncgImpl us fileIds dbgMap cmm count
 
                 let ((alloced, regAllocStats, stack_updt_blks), usAlloc)
                         = {-# SCC "RegAlloc-linear" #-}
-                          initUs usLive
+                          runUniqueDSM usLive
                           $ liftM unzip3
                           $ mapM reg_alloc withLiveness
 
@@ -643,7 +641,7 @@ cmmNativeGen logger ncgImpl us fileIds dbgMap cmm count
         -- sequenced :: [NatCmmDecl statics instr]
         let (sequenced, us_seq) =
                         {-# SCC "sequenceBlocks" #-}
-                        initUs usAlloc $ mapM (BlockLayout.sequenceTop
+                        runUniqueDSM usAlloc $ mapM (BlockLayout.sequenceTop
                                 ncgImpl optimizedCFG)
                             shorted
 
@@ -913,7 +911,7 @@ genMachCode
         -> LabelMap DebugBlock
         -> RawCmmDecl
         -> CFG
-        -> UniqSM
+        -> UniqDSM
                 ( [NatCmmDecl statics instr]
                 , [CLabel]
                 , DwarfFiles
@@ -921,15 +919,16 @@ genMachCode
                 )
 
 genMachCode config cmmTopCodeGen fileIds dbgMap cmm_top cmm_cfg
-  = do  { initial_us <- getUniqueSupplyM
-        ; let initial_st           = mkNatM_State initial_us 0 config
-                                                  fileIds dbgMap cmm_cfg
-              (new_tops, final_st) = initNat initial_st (cmmTopCodeGen cmm_top)
-              final_delta          = natm_delta final_st
-              final_imports        = natm_imports final_st
-              final_cfg            = natm_cfg final_st
-        ; if   final_delta == 0
-          then return (new_tops, final_imports
-                      , natm_fileid final_st, final_cfg)
-          else pprPanic "genMachCode: nonzero final delta" (int final_delta)
-    }
+  = UDSM $ \initial_us -> do
+      { let initial_st           = mkNatM_State initial_us 0 config
+                                                fileIds dbgMap cmm_cfg
+            (new_tops, final_st) = initNat initial_st (cmmTopCodeGen cmm_top)
+            final_delta          = natm_delta final_st
+            final_imports        = natm_imports final_st
+            final_cfg            = natm_cfg final_st
+      ; if   final_delta == 0
+        then DUniqResult
+              (new_tops, final_imports
+                , natm_fileid final_st, final_cfg) (natm_us final_st)
+        else DUniqResult (pprPanic "genMachCode: nonzero final delta" (int final_delta)) undefined
+      }


=====================================
compiler/GHC/CmmToAsm/AArch64/CodeGen.hs
=====================================
@@ -1,4 +1,4 @@
-{-# language GADTs #-}
+{-# language GADTs, LambdaCase #-}
 {-# LANGUAGE OverloadedStrings #-}
 module GHC.CmmToAsm.AArch64.CodeGen (
       cmmTopCodeGen
@@ -44,7 +44,7 @@ import GHC.Cmm.Dataflow.Label
 import GHC.Cmm.Dataflow.Graph
 import GHC.Types.Tickish ( GenTickish(..) )
 import GHC.Types.SrcLoc  ( srcSpanFile, srcSpanStartLine, srcSpanStartCol )
-import GHC.Types.Unique.Supply
+import GHC.Types.Unique.DSM
 
 -- The rest:
 import GHC.Data.OrdList
@@ -1440,7 +1440,7 @@ genCondJump bid expr = do
       _ -> pprPanic "AArch64.genCondJump: " (text $ show expr)
 
 -- A conditional jump with at least +/-128M jump range
-genCondFarJump :: MonadUnique m => Cond -> Target -> m InstrBlock
+genCondFarJump :: MonadGetUnique m => Cond -> Target -> m InstrBlock
 genCondFarJump cond far_target = do
   skip_lbl_id <- newBlockId
   jmp_lbl_id <- newBlockId
@@ -2272,7 +2272,7 @@ data BlockInRange = InRange | NotInRange Target
 
 -- See Note [AArch64 far jumps]
 makeFarBranches :: Platform -> LabelMap RawCmmStatics -> [NatBasicBlock Instr]
-                -> UniqSM [NatBasicBlock Instr]
+                -> UniqDSM [NatBasicBlock Instr]
 makeFarBranches {- only used when debugging -} _platform statics basic_blocks = do
   -- All offsets/positions are counted in multiples of 4 bytes (the size of AArch64 instructions)
   -- That is an offset of 1 represents a 4-byte/one instruction offset.
@@ -2293,7 +2293,7 @@ makeFarBranches {- only used when debugging -} _platform statics basic_blocks =
     long_bz_jump_size =  4 :: Int
 
     -- Replace out of range conditional jumps with unconditional jumps.
-    replace_blk :: LabelMap Int -> Int -> GenBasicBlock Instr -> UniqSM (Int, [GenBasicBlock Instr])
+    replace_blk :: LabelMap Int -> Int -> GenBasicBlock Instr -> UniqDSM (Int, [GenBasicBlock Instr])
     replace_blk !m !pos (BasicBlock lbl instrs) = do
       -- Account for a potential info table before the label.
       let !block_pos = pos + infoTblSize_maybe lbl
@@ -2307,12 +2307,14 @@ makeFarBranches {- only used when debugging -} _platform statics basic_blocks =
       let final_blocks = BasicBlock lbl top : split_blocks
       pure (pos', final_blocks)
 
-    replace_jump :: LabelMap Int -> Int -> Instr -> UniqSM (Int, [Instr])
+    replace_jump :: LabelMap Int -> Int -> Instr -> UniqDSM (Int, [Instr])
     replace_jump !m !pos instr = do
       case instr of
         ANN ann instr -> do
-          (idx,instr':instrs') <- replace_jump m pos instr
-          pure (idx, ANN ann instr':instrs')
+          replace_jump m pos instr >>= \case
+            (idx,instr':instrs') ->
+              pure (idx, ANN ann instr':instrs')
+            (idx,[]) -> pprPanic "replace_jump" (text "empty return list for " <+> ppr idx)
         BCOND cond t
           -> case target_in_range m t pos of
               InRange -> pure (pos+long_bc_jump_size,[instr])


=====================================
compiler/GHC/CmmToAsm/AArch64/Instr.hs
=====================================
@@ -23,7 +23,7 @@ import GHC.Cmm
 import GHC.Cmm.CLabel
 import GHC.Utils.Outputable
 import GHC.Platform
-import GHC.Types.Unique.Supply
+import GHC.Types.Unique.DSM
 
 import GHC.Utils.Panic
 
@@ -473,13 +473,13 @@ allocMoreStack
   :: Platform
   -> Int
   -> NatCmmDecl statics GHC.CmmToAsm.AArch64.Instr.Instr
-  -> UniqSM (NatCmmDecl statics GHC.CmmToAsm.AArch64.Instr.Instr, [(BlockId,BlockId)])
+  -> UniqDSM (NatCmmDecl statics GHC.CmmToAsm.AArch64.Instr.Instr, [(BlockId,BlockId)])
 
 allocMoreStack _ _ top@(CmmData _ _) = return (top,[])
 allocMoreStack platform slots proc@(CmmProc info lbl live (ListGraph code)) = do
     let entries = entryBlocks proc
 
-    uniqs <- getUniquesM
+    retargetList <- mapM (\e -> (e,) . mkBlockId <$> getUniqueM) entries
 
     let
       delta = ((x + stackAlign - 1) `quot` stackAlign) * stackAlign -- round up
@@ -488,8 +488,6 @@ allocMoreStack platform slots proc@(CmmProc info lbl live (ListGraph code)) = do
       alloc   = mkStackAllocInstr   platform delta
       dealloc = mkStackDeallocInstr platform delta
 
-      retargetList = (zip entries (map mkBlockId uniqs))
-
       new_blockmap :: LabelMap BlockId
       new_blockmap = mapFromList retargetList
 


=====================================
compiler/GHC/CmmToAsm/BlockLayout.hs
=====================================
@@ -48,7 +48,7 @@ import Data.STRef
 import Control.Monad.ST.Strict
 import Control.Monad (foldM, unless)
 import GHC.Data.UnionFind
-import GHC.Types.Unique.Supply (UniqSM)
+import GHC.Types.Unique.DSM (UniqDSM)
 
 {-
   Note [CFG based code layout]
@@ -793,7 +793,7 @@ sequenceTop
     => NcgImpl statics instr jumpDest
     -> Maybe CFG -- ^ CFG if we have one.
     -> NatCmmDecl statics instr -- ^ Function to serialize
-    -> UniqSM (NatCmmDecl statics instr)
+    -> UniqDSM (NatCmmDecl statics instr)
 
 sequenceTop _       _           top@(CmmData _ _) = pure top
 sequenceTop ncgImpl edgeWeights (CmmProc info lbl live (ListGraph blocks)) = do


=====================================
compiler/GHC/CmmToAsm/Dwarf.hs
=====================================
@@ -14,7 +14,7 @@ import GHC.Unit.Module
 import GHC.Utils.Outputable
 import GHC.Platform
 import GHC.Types.Unique
-import GHC.Types.Unique.Supply
+import GHC.Types.Unique.DSM
 
 import GHC.CmmToAsm.Dwarf.Constants
 import GHC.CmmToAsm.Dwarf.Types
@@ -31,8 +31,7 @@ import System.FilePath
 import qualified GHC.Cmm.Dataflow.Label as H
 
 -- | Generate DWARF/debug information
-dwarfGen :: IsDoc doc => String -> NCGConfig -> ModLocation -> UniqSupply -> [DebugBlock]
-            -> (doc, UniqSupply)
+dwarfGen :: IsDoc doc => String -> NCGConfig -> ModLocation -> DUniqSupply -> [DebugBlock] -> (doc, DUniqSupply)
 dwarfGen _        _      _      us []     = (empty, us)
 dwarfGen compPath config modLoc us blocks =
   let platform = ncgPlatform config
@@ -65,7 +64,7 @@ dwarfGen compPath config modLoc us blocks =
 
   -- .debug_info section: Information records on procedures and blocks
       -- unique to identify start and end compilation unit .debug_inf
-      (unitU, us') = takeUniqFromSupply us
+      (unitU, us') = takeUniqueFromDSupply us
       infoSct = vcat [ line (dwarfInfoLabel <> colon)
                      , dwarfInfoSection platform
                      , compileUnitHeader platform unitU
@@ -79,7 +78,7 @@ dwarfGen compPath config modLoc us blocks =
                 line (dwarfLineLabel <> colon)
 
   -- .debug_frame section: Information about the layout of the GHC stack
-      (framesU, us'') = takeUniqFromSupply us'
+      (framesU, us'') = takeUniqueFromDSupply us'
       frameSct = dwarfFrameSection platform $$
                  line (dwarfFrameLabel <> colon) $$
                  pprDwarfFrame platform (debugFrame platform framesU procs)
@@ -90,8 +89,8 @@ dwarfGen compPath config modLoc us blocks =
       aranges = dwarfARangesSection platform $$ pprDwarfARanges platform aranges' unitU
 
   in (infoSct $$ abbrevSct $$ lineSct $$ frameSct $$ aranges, us'')
-{-# SPECIALIZE dwarfGen :: String -> NCGConfig -> ModLocation -> UniqSupply -> [DebugBlock] -> (SDoc, UniqSupply) #-}
-{-# SPECIALIZE dwarfGen :: String -> NCGConfig -> ModLocation -> UniqSupply -> [DebugBlock] -> (HDoc, UniqSupply) #-} -- see Note [SPECIALIZE to HDoc] in GHC.Utils.Outputable
+{-# SPECIALIZE dwarfGen :: String -> NCGConfig -> ModLocation -> DUniqSupply -> [DebugBlock] -> (SDoc, DUniqSupply) #-}
+{-# SPECIALIZE dwarfGen :: String -> NCGConfig -> ModLocation -> DUniqSupply -> [DebugBlock] -> (HDoc, DUniqSupply) #-} -- see Note [SPECIALIZE to HDoc] in GHC.Utils.Outputable
 
 -- | Build an address range entry for one proc.
 -- With split sections, each proc needs its own entry, since they may get


=====================================
compiler/GHC/CmmToAsm/Monad.hs
=====================================
@@ -62,7 +62,7 @@ import GHC.Cmm.Expr             (LocalReg (..), isWord64)
 
 import GHC.Data.FastString      ( FastString )
 import GHC.Types.Unique.FM
-import GHC.Types.Unique.Supply
+import GHC.Types.Unique.DSM
 import GHC.Types.Unique         ( Unique )
 import GHC.Unit.Module
 
@@ -109,11 +109,11 @@ data NcgImpl statics instr jumpDest = NcgImpl {
     maxSpillSlots             :: Int,
     allocatableRegs           :: [RealReg],
     ncgAllocMoreStack         :: Int -> NatCmmDecl statics instr
-                              -> UniqSM (NatCmmDecl statics instr, [(BlockId,BlockId)]),
+                              -> UniqDSM (NatCmmDecl statics instr, [(BlockId,BlockId)]),
     -- ^ The list of block ids records the redirected jumps to allow us to update
     -- the CFG.
     ncgMakeFarBranches        :: Platform -> LabelMap RawCmmStatics -> [NatBasicBlock instr]
-                              -> UniqSM [NatBasicBlock instr],
+                              -> UniqDSM [NatBasicBlock instr],
     extractUnwindPoints       :: [instr] -> [UnwindPoint],
     -- ^ given the instruction sequence of a block, produce a list of
     -- the block's 'UnwindPoint's
@@ -178,7 +178,7 @@ mistake would readily show up in performance tests). -}
 
 data NatM_State
         = NatM_State {
-                natm_us          :: UniqSupply,
+                natm_us          :: DUniqSupply,
                 natm_delta       :: Int, -- ^ Stack offset for unwinding information
                 natm_imports     :: [(CLabel)],
                 natm_pic         :: Maybe Reg,
@@ -205,7 +205,7 @@ pattern NatM f <- NatM' (runState -> f)
 unNat :: NatM a -> NatM_State -> (a, NatM_State)
 unNat (NatM a) = a
 
-mkNatM_State :: UniqSupply -> Int -> NCGConfig ->
+mkNatM_State :: DUniqSupply -> Int -> NCGConfig ->
                 DwarfFiles -> LabelMap DebugBlock -> CFG -> NatM_State
 mkNatM_State us delta config
         = \dwf dbg cfg ->
@@ -223,19 +223,13 @@ mkNatM_State us delta config
 initNat :: NatM_State -> NatM a -> (a, NatM_State)
 initNat = flip unNat
 
-instance MonadUnique NatM where
-  getUniqueSupplyM = NatM $ \st ->
-      case splitUniqSupply (natm_us st) of
-          (us1, us2) -> (us1, st {natm_us = us2})
-
+instance MonadGetUnique NatM where
   getUniqueM = NatM $ \st ->
-      case takeUniqFromSupply (natm_us st) of
-          (uniq, us') -> (uniq, st {natm_us = us'})
+      case takeUniqueFromDSupply (natm_us st) of
+        (uniq, us') -> (uniq, st {natm_us = us'})
 
 getUniqueNat :: NatM Unique
-getUniqueNat = NatM $ \ st ->
-    case takeUniqFromSupply $ natm_us st of
-    (uniq, us') -> (uniq, st {natm_us = us'})
+getUniqueNat = getUniqueM
 
 getDeltaNat :: NatM Int
 getDeltaNat = NatM $ \ st -> (natm_delta st, st)


=====================================
compiler/GHC/CmmToAsm/PPC/Instr.hs
=====================================
@@ -55,7 +55,7 @@ import GHC.Cmm.CLabel
 import GHC.Utils.Panic
 import GHC.Platform
 import GHC.Types.Unique.FM (listToUFM, lookupUFM)
-import GHC.Types.Unique.Supply
+import GHC.Types.Unique.DSM
 
 import Data.Foldable (toList)
 import qualified Data.List.NonEmpty as NE
@@ -105,7 +105,7 @@ allocMoreStack
   :: Platform
   -> Int
   -> NatCmmDecl statics GHC.CmmToAsm.PPC.Instr.Instr
-  -> UniqSM (NatCmmDecl statics GHC.CmmToAsm.PPC.Instr.Instr, [(BlockId,BlockId)])
+  -> UniqDSM (NatCmmDecl statics GHC.CmmToAsm.PPC.Instr.Instr, [(BlockId,BlockId)])
 
 allocMoreStack _ _ top@(CmmData _ _) = return (top,[])
 allocMoreStack platform slots (CmmProc info lbl live (ListGraph code)) = do
@@ -117,7 +117,7 @@ allocMoreStack platform slots (CmmProc info lbl live (ListGraph code)) = do
                         | entry `elem` infos -> infos
                         | otherwise          -> entry : infos
 
-    uniqs <- getUniquesM
+    retargetList <- mapM (\e -> (e,) . mkBlockId <$> getUniqueM) entries
 
     let
         delta = ((x + stackAlign - 1) `quot` stackAlign) * stackAlign -- round up
@@ -126,8 +126,6 @@ allocMoreStack platform slots (CmmProc info lbl live (ListGraph code)) = do
         alloc   = mkStackAllocInstr   platform delta
         dealloc = mkStackDeallocInstr platform delta
 
-        retargetList = (zip entries (map mkBlockId uniqs))
-
         new_blockmap :: LabelMap BlockId
         new_blockmap = mapFromList retargetList
 
@@ -698,7 +696,7 @@ makeFarBranches
         :: Platform
         -> LabelMap RawCmmStatics
         -> [NatBasicBlock Instr]
-        -> UniqSM [NatBasicBlock Instr]
+        -> UniqDSM [NatBasicBlock Instr]
 makeFarBranches _platform info_env blocks
     | NE.last blockAddresses < nearLimit = return blocks
     | otherwise = return $ zipWith handleBlock blockAddressList blocks


=====================================
compiler/GHC/CmmToAsm/Reg/Graph.hs
=====================================
@@ -31,7 +31,7 @@ import GHC.Utils.Panic
 import GHC.Platform
 import GHC.Types.Unique.FM
 import GHC.Types.Unique.Set
-import GHC.Types.Unique.Supply
+import GHC.Types.Unique.DSM
 import GHC.Utils.Misc (seqList)
 import GHC.CmmToAsm.CFG
 
@@ -57,8 +57,8 @@ regAlloc
         -> Int                          -- ^ current number of spill slots
         -> [LiveCmmDecl statics instr]  -- ^ code annotated with liveness information.
         -> Maybe CFG                    -- ^ CFG of basic blocks if available
-        -> UniqSM ( [NatCmmDecl statics instr]
-                  , Maybe Int, [RegAllocStats statics instr] )
+        -> UniqDSM ( [NatCmmDecl statics instr]
+                   , Maybe Int, [RegAllocStats statics instr] )
            -- ^ code with registers allocated, additional stacks required
            -- and stats for each stage of allocation
 
@@ -107,7 +107,7 @@ regAlloc_spin
         -> [RegAllocStats statics instr] -- ^ Current regalloc stats to add to.
         -> [LiveCmmDecl statics instr]   -- ^ Liveness annotated code to allocate.
         -> Maybe CFG
-        -> UniqSM ( [NatCmmDecl statics instr]
+        -> UniqDSM ( [NatCmmDecl statics instr]
                   , [RegAllocStats statics instr]
                   , Int                  -- Slots in use
                   , Color.Graph VirtualReg RegClass RealReg)
@@ -305,7 +305,7 @@ regAlloc_spin config spinCount triv regsFree slotsFree slotsCount debug_codeGrap
 buildGraph
         :: Instruction instr
         => [LiveCmmDecl statics instr]
-        -> UniqSM (Color.Graph VirtualReg RegClass RealReg)
+        -> UniqDSM (Color.Graph VirtualReg RegClass RealReg)
 
 buildGraph code
  = do


=====================================
compiler/GHC/CmmToAsm/Reg/Graph/Spill.hs
=====================================
@@ -22,7 +22,7 @@ import GHC.Utils.Monad.State.Strict
 import GHC.Types.Unique
 import GHC.Types.Unique.FM
 import GHC.Types.Unique.Set
-import GHC.Types.Unique.Supply
+import GHC.Types.Unique.DSM
 import GHC.Utils.Outputable
 import GHC.Utils.Panic
 import GHC.Platform
@@ -52,7 +52,7 @@ regSpill
         -> UniqSet Int                  -- ^ available stack slots
         -> Int                          -- ^ current number of spill slots.
         -> UniqSet VirtualReg           -- ^ the regs to spill
-        -> UniqSM
+        -> UniqDSM
             ([LiveCmmDecl statics instr]
                  -- code with SPILL and RELOAD meta instructions added.
             , UniqSet Int               -- left over slots
@@ -81,17 +81,20 @@ regSpill platform code slotsFree slotCount regs
                     -- See Note [Unique Determinism and code generation]
 
                 -- Grab the unique supply from the monad.
-                us      <- getUniqueSupplyM
+                UDSM $ \us ->
 
-                -- Run the spiller on all the blocks.
-                let (code', state')     =
-                        runState (mapM (regSpill_top platform regSlotMap) code)
-                                 (initSpillS us)
+                  -- Run the spiller on all the blocks.
+                  let (code', state')     =
+                          runState (mapM (regSpill_top platform regSlotMap) code)
+                                   (initSpillS us)
 
-                return  ( code'
+                   in DUniqResult
+                        ( code'
                         , minusUniqSet slotsFree (mkUniqSet slots)
                         , slotCount
                         , makeSpillStats state')
+                        ( stateUS state' )
+
 
 
 -- | Spill some registers to stack slots in a top-level thing.
@@ -323,21 +326,28 @@ patchReg1 old new instr
 
 -- Spiller monad --------------------------------------------------------------
 -- | State monad for the spill code generator.
-type SpillM a
-        = State SpillS a
+type SpillM = State SpillS
 
 -- | Spill code generator state.
 data SpillS
         = SpillS
         { -- | Unique supply for generating fresh vregs.
-          stateUS       :: UniqSupply
+          stateUS       :: DUniqSupply
 
           -- | Spilled vreg vs the number of times it was loaded, stored.
         , stateSpillSL  :: UniqFM Reg (Reg, Int, Int) }
 
+instance MonadGetUnique SpillM where
+  getUniqueM = do
+    us <- gets stateUS
+    case takeUniqueFromDSupply us of
+     (uniq, us')
+      -> do modify $ \s -> s { stateUS = us' }
+            return uniq
+
 
 -- | Create a new spiller state.
-initSpillS :: UniqSupply -> SpillS
+initSpillS :: DUniqSupply -> SpillS
 initSpillS uniqueSupply
         = SpillS
         { stateUS       = uniqueSupply
@@ -346,12 +356,7 @@ initSpillS uniqueSupply
 
 -- | Allocate a new unique in the spiller monad.
 newUnique :: SpillM Unique
-newUnique
- = do   us      <- gets stateUS
-        case takeUniqFromSupply us of
-         (uniq, us')
-          -> do modify $ \s -> s { stateUS = us' }
-                return uniq
+newUnique = getUniqueM
 
 
 -- | Add a spill/reload count to a stats record for a register.


=====================================
compiler/GHC/CmmToAsm/Reg/Linear.hs
=====================================
@@ -129,7 +129,7 @@ import GHC.Data.Graph.Directed
 import GHC.Types.Unique
 import GHC.Types.Unique.Set
 import GHC.Types.Unique.FM
-import GHC.Types.Unique.Supply
+import GHC.Types.Unique.DSM
 import GHC.Utils.Outputable
 import GHC.Utils.Panic
 import GHC.Platform
@@ -146,11 +146,11 @@ regAlloc
         :: Instruction instr
         => NCGConfig
         -> LiveCmmDecl statics instr
-        -> UniqSM ( NatCmmDecl statics instr
-                  , Maybe Int  -- number of extra stack slots required,
-                               -- beyond maxSpillSlots
-                  , Maybe RegAllocStats
-                  )
+        -> UniqDSM ( NatCmmDecl statics instr
+                   , Maybe Int  -- number of extra stack slots required,
+                                -- beyond maxSpillSlots
+                   , Maybe RegAllocStats
+                   )
 
 regAlloc _ (CmmData sec d)
         = return
@@ -207,7 +207,7 @@ linearRegAlloc
               -- ^ live regs on entry to each basic block
         -> [SCC (LiveBasicBlock instr)]
               -- ^ instructions annotated with "deaths"
-        -> UniqSM ([NatBasicBlock instr], RegAllocStats, Int)
+        -> UniqDSM ([NatBasicBlock instr], RegAllocStats, Int)
 
 linearRegAlloc config entry_ids block_live sccs
  = case platformArch platform of
@@ -228,7 +228,7 @@ linearRegAlloc config entry_ids block_live sccs
       ArchUnknown    -> panic "linearRegAlloc ArchUnknown"
  where
   go :: (FR regs, Outputable regs)
-     => regs -> UniqSM ([NatBasicBlock instr], RegAllocStats, Int)
+     => regs -> UniqDSM ([NatBasicBlock instr], RegAllocStats, Int)
   go f = linearRegAlloc' config f entry_ids block_live sccs
   platform = ncgPlatform config
 
@@ -244,14 +244,14 @@ linearRegAlloc'
         -> [BlockId]                    -- ^ entry points
         -> BlockMap RegSet              -- ^ live regs on entry to each basic block
         -> [SCC (LiveBasicBlock instr)] -- ^ instructions annotated with "deaths"
-        -> UniqSM ([NatBasicBlock instr], RegAllocStats, Int)
+        -> UniqDSM ([NatBasicBlock instr], RegAllocStats, Int)
 
 linearRegAlloc' config initFreeRegs entry_ids block_live sccs
- = do   us      <- getUniqueSupplyM
-        let !(_, !stack, !stats, !blocks) =
-                runR config emptyBlockAssignment initFreeRegs emptyRegMap emptyStackMap us
-                    $ linearRA_SCCs entry_ids block_live [] sccs
-        return  (blocks, stats, getStackUse stack)
+ = UDSM $ \us -> do
+    let !(_, !stack, !stats, !blocks, us') =
+            runR config emptyBlockAssignment initFreeRegs emptyRegMap emptyStackMap us
+                $ linearRA_SCCs entry_ids block_live [] sccs
+     in DUniqResult (blocks, stats, getStackUse stack) us'
 
 
 linearRA_SCCs :: OutputableRegConstraint freeRegs instr


=====================================
compiler/GHC/CmmToAsm/Reg/Linear/Base.hs
=====================================
@@ -32,7 +32,7 @@ import GHC.Platform.Reg
 import GHC.Utils.Outputable
 import GHC.Types.Unique
 import GHC.Types.Unique.FM
-import GHC.Types.Unique.Supply
+import GHC.Types.Unique.DSM
 import GHC.Cmm.BlockId
 import GHC.Cmm.Dataflow.Label
 import GHC.CmmToAsm.Reg.Utils
@@ -170,7 +170,7 @@ data RA_State freeRegs
         , ra_stack      :: StackMap
 
         -- | unique supply for generating names for join point fixup blocks.
-        , ra_us         :: UniqSupply
+        , ra_us         :: DUniqSupply
 
         -- | Record why things were spilled, for -ddrop-asm-stats.
         --      Just keep a list here instead of a map of regs -> reasons.


=====================================
compiler/GHC/CmmToAsm/Reg/Linear/State.hs
=====================================
@@ -49,7 +49,7 @@ import GHC.Cmm.BlockId
 
 import GHC.Platform
 import GHC.Types.Unique
-import GHC.Types.Unique.Supply
+import GHC.Types.Unique.DSM
 import GHC.Exts (oneShot)
 
 import Control.Monad (ap)
@@ -91,9 +91,9 @@ runR    :: NCGConfig
         -> freeRegs
         -> RegMap Loc
         -> StackMap
-        -> UniqSupply
+        -> DUniqSupply
         -> RegM freeRegs a
-        -> (BlockAssignment freeRegs, StackMap, RegAllocStats, a)
+        -> (BlockAssignment freeRegs, StackMap, RegAllocStats, a, DUniqSupply)
 
 runR config block_assig freeregs assig stack us thing =
   case unReg thing
@@ -109,7 +109,7 @@ runR config block_assig freeregs assig stack us thing =
                 , ra_fixups     = [] })
    of
         RA_Result state returned_thing
-         ->     (ra_blockassig state, ra_stack state, makeRAStats state, returned_thing)
+         ->  (ra_blockassig state, ra_stack state, makeRAStats state, returned_thing, ra_us state)
 
 
 -- | Make register allocator stats from its final state.
@@ -169,7 +169,7 @@ getDeltaR = mkRegM $ \s -> RA_Result s (ra_delta s)
 
 getUniqueR :: RegM freeRegs Unique
 getUniqueR = mkRegM $ \s ->
-  case takeUniqFromSupply (ra_us s) of
+  case takeUniqueFromDSupply (ra_us s) of
     (uniq, us) -> RA_Result s{ra_us = us} uniq
 
 


=====================================
compiler/GHC/CmmToAsm/Reg/Liveness.hs
=====================================
@@ -56,7 +56,7 @@ import GHC.Utils.Panic
 import GHC.Platform
 import GHC.Types.Unique.Set
 import GHC.Types.Unique.FM
-import GHC.Types.Unique.Supply
+import GHC.Types.Unique.DSM
 import GHC.Data.Bag
 import GHC.Utils.Monad.State.Strict
 
@@ -690,7 +690,7 @@ cmmTopLiveness
         => Maybe CFG
         -> Platform
         -> NatCmmDecl statics instr
-        -> UniqSM (LiveCmmDecl statics instr)
+        -> UniqDSM (LiveCmmDecl statics instr)
 cmmTopLiveness cfg platform cmm
         = regLiveness platform $ natCmmTopToLive cfg cmm
 
@@ -784,7 +784,7 @@ regLiveness
         :: Instruction instr
         => Platform
         -> LiveCmmDecl statics instr
-        -> UniqSM (LiveCmmDecl statics instr)
+        -> UniqDSM (LiveCmmDecl statics instr)
 
 regLiveness _ (CmmData i d)
         = return $ CmmData i d


=====================================
compiler/GHC/CmmToAsm/Wasm.hs
=====================================
@@ -21,7 +21,7 @@ import GHC.Driver.DynFlags
 import GHC.Platform
 import GHC.Prelude
 import GHC.Settings
-import GHC.Types.Unique.Supply
+import GHC.Types.Unique.DSM
 import GHC.Unit
 import GHC.Utils.Logger
 import GHC.Utils.Outputable (text)
@@ -32,7 +32,7 @@ ncgWasm ::
   Logger ->
   Platform ->
   ToolSettings ->
-  UniqSupply ->
+  DUniqSupply ->
   ModLocation ->
   Handle ->
   Stream IO RawCmmGroup a ->
@@ -58,7 +58,7 @@ ncgWasm ncg_config logger platform ts us loc h cmms = do
 streamCmmGroups ::
   NCGConfig ->
   Platform ->
-  UniqSupply ->
+  DUniqSupply ->
   Stream IO RawCmmGroup a ->
   IO (a, WasmCodeGenState 'I32)
 streamCmmGroups ncg_config platform us cmms =


=====================================
compiler/GHC/CmmToAsm/Wasm/FromCmm.hs
=====================================
@@ -48,7 +48,7 @@ import GHC.Types.Unique
 import GHC.Types.Unique.FM
 import GHC.Types.Unique.Map
 import GHC.Types.Unique.Set
-import GHC.Types.Unique.Supply
+import GHC.Types.Unique.DSM
 import GHC.Utils.Outputable hiding ((<>))
 import GHC.Utils.Panic
 import GHC.Wasm.ControlFlow.FromCmm
@@ -1572,11 +1572,9 @@ lower_CmmGraph :: CLabel -> CmmGraph -> WasmCodeGenM w (FuncBody w)
 lower_CmmGraph lbl g = do
   ty_word <- wasmWordTypeM
   platform <- wasmPlatformM
-  us <- getUniqueSupplyM
   body <-
     structuredControl
       platform
-      us
       (\_ -> lower_CmmExpr_Typed lbl ty_word)
       (lower_CmmActions lbl)
       g


=====================================
compiler/GHC/CmmToAsm/Wasm/Types.hs
=====================================
@@ -45,6 +45,7 @@ module GHC.CmmToAsm.Wasm.Types
     wasmStateM,
     wasmModifyM,
     wasmExecM,
+    wasmRunM
   )
 where
 
@@ -66,9 +67,10 @@ import GHC.Types.Unique
 import GHC.Types.Unique.FM
 import GHC.Types.Unique.Map
 import GHC.Types.Unique.Set
-import GHC.Types.Unique.Supply
+import GHC.Types.Unique.DSM
 import GHC.Utils.Monad.State.Strict
 import GHC.Utils.Outputable hiding ((<>))
+import GHC.Data.Graph.Collapse (MonadUniqDSM(..))
 import Unsafe.Coerce
 
 -- | WebAssembly type of a WebAssembly value that WebAssembly code
@@ -419,10 +421,10 @@ data WasmCodeGenState w = WasmCodeGenState
       UniqFM LocalReg LocalInfo,
     localRegsCount ::
       Int,
-    wasmUniqSupply :: UniqSupply
+    wasmDUniqSupply :: DUniqSupply
   }
 
-initialWasmCodeGenState :: Platform -> UniqSupply -> WasmCodeGenState w
+initialWasmCodeGenState :: Platform -> DUniqSupply -> WasmCodeGenState w
 initialWasmCodeGenState platform us =
   WasmCodeGenState
     { wasmPlatform =
@@ -436,12 +438,17 @@ initialWasmCodeGenState platform us =
         [],
       localRegs = emptyUFM,
       localRegsCount = 0,
-      wasmUniqSupply = us
+      wasmDUniqSupply = us
     }
 
 newtype WasmCodeGenM w a = WasmCodeGenM (State (WasmCodeGenState w) a)
   deriving newtype (Functor, Applicative, Monad)
 
+instance MonadUniqDSM (WasmCodeGenM w) where
+  liftUniqDSM (UDSM m) = wasmStateM $ \st ->
+    let DUniqResult a us' = m (wasmDUniqSupply st)
+     in (# a, st{wasmDUniqSupply=us'} #)
+
 wasmGetsM :: (WasmCodeGenState w -> a) -> WasmCodeGenM w a
 wasmGetsM = coerce . gets
 
@@ -465,18 +472,13 @@ wasmStateM = coerce . State
 wasmModifyM :: (WasmCodeGenState w -> WasmCodeGenState w) -> WasmCodeGenM w ()
 wasmModifyM = coerce . modify
 
-wasmEvalM :: WasmCodeGenM w a -> WasmCodeGenState w -> a
-wasmEvalM (WasmCodeGenM s) = evalState s
-
 wasmExecM :: WasmCodeGenM w a -> WasmCodeGenState w -> WasmCodeGenState w
 wasmExecM (WasmCodeGenM s) = execState s
 
-instance MonadUnique (WasmCodeGenM w) where
-  getUniqueSupplyM = wasmGetsM wasmUniqSupply
+wasmRunM :: WasmCodeGenM w a -> WasmCodeGenState w -> (a, WasmCodeGenState w)
+wasmRunM (WasmCodeGenM s) = runState s
+
+instance MonadGetUnique (WasmCodeGenM w) where
   getUniqueM = wasmStateM $
-    \s at WasmCodeGenState {..} -> case takeUniqFromSupply wasmUniqSupply of
-      (u, us) -> (# u, s {wasmUniqSupply = us} #)
-  getUniquesM = do
-    u <- getUniqueM
-    s <- WasmCodeGenM get
-    pure $ u:(wasmEvalM getUniquesM s)
+    \s at WasmCodeGenState {..} -> case takeUniqueFromDSupply wasmDUniqSupply of
+      (u, us) -> (# u, s {wasmDUniqSupply = us} #)


=====================================
compiler/GHC/CmmToAsm/X86/CodeGen.hs
=====================================
@@ -79,7 +79,7 @@ import GHC.Utils.Constants (debugIsOn)
 import GHC.Utils.Panic
 import GHC.Data.FastString
 import GHC.Utils.Misc
-import GHC.Types.Unique.Supply ( getUniqueM )
+import GHC.Types.Unique.DSM ( getUniqueM )
 
 import Control.Monad
 import Data.Foldable (fold)


=====================================
compiler/GHC/CmmToAsm/X86/Instr.hs
=====================================
@@ -64,7 +64,7 @@ import GHC.Platform
 import GHC.Cmm.CLabel
 import GHC.Types.Unique.Set
 import GHC.Types.Unique
-import GHC.Types.Unique.Supply
+import GHC.Types.Unique.DSM
 import GHC.Types.Basic (Alignment)
 import GHC.Cmm.DebugBlock (UnwindTable)
 
@@ -987,13 +987,13 @@ allocMoreStack
   :: Platform
   -> Int
   -> NatCmmDecl statics GHC.CmmToAsm.X86.Instr.Instr
-  -> UniqSM (NatCmmDecl statics GHC.CmmToAsm.X86.Instr.Instr, [(BlockId,BlockId)])
+  -> UniqDSM (NatCmmDecl statics GHC.CmmToAsm.X86.Instr.Instr, [(BlockId,BlockId)])
 
 allocMoreStack _ _ top@(CmmData _ _) = return (top,[])
 allocMoreStack platform slots proc@(CmmProc info lbl live (ListGraph code)) = do
     let entries = entryBlocks proc
 
-    uniqs <- getUniquesM
+    retargetList <- mapM (\e -> (e,) . mkBlockId <$> getUniqueM) entries
 
     let
       delta = ((x + stackAlign - 1) `quot` stackAlign) * stackAlign -- round up
@@ -1002,8 +1002,6 @@ allocMoreStack platform slots proc@(CmmProc info lbl live (ListGraph code)) = do
       alloc   = mkStackAllocInstr   platform delta
       dealloc = mkStackDeallocInstr platform delta
 
-      retargetList = (zip entries (map mkBlockId uniqs))
-
       new_blockmap :: LabelMap BlockId
       new_blockmap = mapFromList retargetList
 


=====================================
compiler/GHC/CmmToLlvm/Base.hs
=====================================
@@ -56,6 +56,7 @@ import GHC.Types.Unique
 import GHC.Utils.BufHandle   ( BufHandle )
 import GHC.Types.Unique.Set
 import GHC.Types.Unique.Supply
+import qualified GHC.Types.Unique.DSM as DSM
 import GHC.Utils.Logger
 
 import Data.Maybe (fromJust)
@@ -300,6 +301,13 @@ instance MonadUnique LlvmM where
         tag <- getEnv envTag
         liftIO $! uniqFromTag tag
 
+-- TODO: If you want Llvm code to be deterministic, this should use a
+-- deterministic unique supply to get the Id.
+instance DSM.MonadGetUnique LlvmM where
+  getUniqueM = do
+    tag <- getEnv envTag
+    liftIO $! uniqFromTag tag
+
 -- | Lifting of IO actions. Not exported, as we want to encapsulate IO.
 liftIO :: IO a -> LlvmM a
 liftIO m = LlvmM $ \env -> do x <- m


=====================================
compiler/GHC/Data/Graph/Collapse.hs
=====================================
@@ -11,7 +11,7 @@ module GHC.Data.Graph.Collapse
   , VizCollapseMonad(..)
   , NullCollapseViz(..)
   , runNullCollapse
-  , MonadUniqSM(..)
+  , MonadUniqDSM(..)
   )
 where
 
@@ -24,7 +24,7 @@ import Data.Semigroup
 
 import GHC.Cmm.Dataflow.Label
 import GHC.Data.Graph.Inductive.Graph
-import GHC.Types.Unique.Supply
+import GHC.Types.Unique.DSM
 import GHC.Utils.Panic hiding (assert)
 
 
@@ -59,23 +59,21 @@ Functional Graph Library (Hackage package `fgl`, modules
 -- care about visualization, you would use the `NullCollapseViz`
 -- monad, in which these operations are no-ops.
 
-class (Monad m) => MonadUniqSM m where
-  liftUniqSM :: UniqSM a -> m a
+class (Monad m) => MonadUniqDSM m where
+  liftUniqDSM :: UniqDSM a -> m a
 
-class (MonadUniqSM m, Graph gr, Supernode s m) => VizCollapseMonad m gr s where
+class (MonadUniqDSM m, Graph gr, Supernode s m) => VizCollapseMonad m gr s where
   consumeByInGraph :: Node -> Node -> gr s () -> m ()
   splitGraphAt :: gr s () -> LNode s -> m ()
   finalGraph :: gr s () -> m ()
 
-
-
 -- | The identity monad as a `VizCollapseMonad`.  Use this monad when
 -- you want efficiency in graph collapse.
-newtype NullCollapseViz a = NullCollapseViz { unNCV :: UniqSM a }
-  deriving (Functor, Applicative, Monad, MonadUnique)
+newtype NullCollapseViz a = NullCollapseViz { unNCV :: UniqDSM a }
+  deriving (Functor, Applicative, Monad, MonadGetUnique)
 
-instance MonadUniqSM NullCollapseViz where
-  liftUniqSM = NullCollapseViz
+instance MonadUniqDSM NullCollapseViz where
+  liftUniqDSM = NullCollapseViz
 
 instance (Graph gr, Supernode s NullCollapseViz) =>
     VizCollapseMonad NullCollapseViz gr s where
@@ -83,7 +81,7 @@ instance (Graph gr, Supernode s NullCollapseViz) =>
   splitGraphAt _ _ = return ()
   finalGraph _ = return ()
 
-runNullCollapse :: NullCollapseViz a -> UniqSM a
+runNullCollapse :: NullCollapseViz a -> UniqDSM a
 runNullCollapse = unNCV
 
 
@@ -158,7 +156,7 @@ class (Semigroup node) => PureSupernode node where
   superLabel :: node -> Label
   mapLabels :: (Label -> Label) -> (node -> node)
 
-class (MonadUnique m, PureSupernode node) => Supernode node m where
+class (MonadGetUnique m, PureSupernode node) => Supernode node m where
   freshen :: node -> m node
 
   -- ghost method


=====================================
compiler/GHC/Driver/CodeOutput.hs
=====================================
@@ -56,7 +56,7 @@ import GHC.Unit.Finder      ( mkStubPaths )
 import GHC.Types.SrcLoc
 import GHC.Types.CostCentre
 import GHC.Types.ForeignStubs
-import GHC.Types.Unique.Supply ( mkSplitUniqSupply )
+import GHC.Types.Unique.DSM
 
 import System.Directory
 import System.FilePath
@@ -198,7 +198,7 @@ outputAsm :: Logger
           -> Stream IO RawCmmGroup a
           -> IO a
 outputAsm logger dflags this_mod location filenm cmm_stream = do
-  ncg_uniqs <- mkSplitUniqSupply 'n'
+  let ncg_uniqs = initDUniqSupply 'n' 0 {- See Note [Cmm Local Deterministic Uniques], or should we receive it as input?-}
   debugTraceMsg logger 4 (text "Outputing asm to" <+> text filenm)
   let ncg_config = initNCGConfig dflags this_mod
   {-# SCC "OutputAsm" #-} doOutput filenm $


=====================================
compiler/GHC/StgToCmm/ExtCode.hs
=====================================
@@ -57,6 +57,7 @@ import GHC.Unit.Module
 import GHC.Types.Unique.FM
 import GHC.Types.Unique
 import GHC.Types.Unique.Supply
+import qualified GHC.Types.Unique.DSM as DSM
 
 import Control.Monad (ap)
 import GHC.Utils.Outputable (SDocContext)
@@ -102,6 +103,9 @@ instance MonadUnique CmmParse where
     u <- getUniqueM
     return (decls, u)
 
+instance DSM.MonadGetUnique CmmParse where
+  getUniqueM = GHC.Types.Unique.Supply.getUniqueM
+
 getProfile :: CmmParse Profile
 getProfile = EC (\_ _ d -> (d,) <$> F.getProfile)
 


=====================================
compiler/GHC/Types/Unique/DSM.hs
=====================================
@@ -0,0 +1,97 @@
+{-# LANGUAGE UnboxedTuples, PatternSynonyms #-}
+module GHC.Types.Unique.DSM where
+
+import GHC.Prelude
+import GHC.Word
+import Control.Monad.Fix
+import GHC.Types.Unique
+import qualified GHC.Types.Unique.Supply as USM
+
+-- todo: Do I need to use the one-shot state monad trick? Probably yes.
+
+-- check: UniqSM is only used before Cmm (grep for it), afterwards only UniqDSM is used.
+
+-- todo: use UniqSM for UniqRenamable? We've basically re-implemented this logic
+-- there, but without the unboxing it feels? Maybe not, since we carry the
+-- mappings too.
+
+newtype DUniqSupply = DUS Word64 -- supply uniques iteratively
+type DUniqResult result = (# result, DUniqSupply #)
+
+pattern DUniqResult :: a -> b -> (# a, b #)
+pattern DUniqResult x y = (# x, y #)
+{-# COMPLETE DUniqResult #-}
+
+-- | A monad which just gives the ability to obtain 'Unique's deterministically.
+-- There's no splitting.
+newtype UniqDSM result = UDSM { unUDSM :: DUniqSupply -> DUniqResult result }
+  deriving Functor
+
+instance Monad UniqDSM where
+  (>>=) (UDSM f) cont = UDSM $ \us0 -> case f us0 of
+    DUniqResult result us1 -> unUDSM (cont result) us1
+  (>>)  = (*>)
+  {-# INLINE (>>=) #-}
+  {-# INLINE (>>) #-}
+
+instance Applicative UniqDSM where
+  pure result = UDSM (DUniqResult result)
+  (UDSM f) <*> (UDSM x) = UDSM $ \us0 -> case f us0 of
+    DUniqResult ff us1 -> case x us1 of
+      DUniqResult xx us2 -> DUniqResult (ff xx) us2
+  (*>) (UDSM expr) (UDSM cont) = UDSM $ \us0 -> case expr us0 of
+    DUniqResult _ us1 -> cont us1
+  {-# INLINE pure #-}
+  {-# INLINE (*>) #-}
+
+instance MonadFix UniqDSM where
+  mfix m = UDSM (\us0 -> let (r,us1) = runUniqueDSM us0 (m r) in DUniqResult r us1)
+
+getUniqueDSM :: UniqDSM Unique
+getUniqueDSM = UDSM (\(DUS us0) -> DUniqResult (mkUniqueGrimily us0) (DUS $ us0+1))
+
+takeUniqueFromDSupply :: DUniqSupply -> (Unique, DUniqSupply)
+takeUniqueFromDSupply d =
+  case unUDSM getUniqueDSM d of
+    DUniqResult x y -> (x, y)
+
+-- Write Note about the importance of locality in uniques that are deterministic
+--
+-- If you use a tag which collides with other names, you'll get a uniques
+-- deterministically colliding with existing symbols.
+--
+-- (e.g. easy to observe if you do this wrong)
+--
+-- Ideally, we'd thread the same deterministic unique supply all the way
+-- throughout the Cmm pipeline, starting off from hte deterministic rename
+-- pass.
+initDUniqSupply :: Char -> Word64 -> DUniqSupply
+initDUniqSupply c firstUniq =
+  let !tag = mkTag c
+  in DUS (tag .|. firstUniq)
+
+newTagDUniqSupply :: Char -> DUniqSupply -> DUniqSupply
+newTagDUniqSupply c (DUS w) = DUS $ getKey $ newTagUnique (mkUniqueGrimily w) c
+
+runUniqueDSM :: DUniqSupply -> UniqDSM a -> (a, DUniqSupply)
+runUniqueDSM ds (UDSM f) =
+  case f ds of
+    DUniqResult uq us -> (uq, us)
+
+-- Add explanation on how this gives you a deterministic way of getting uniques
+-- if the instance uses a deterministic unique supply.
+class Monad m => MonadGetUnique m where
+  getUniqueM :: m Unique
+
+instance MonadGetUnique UniqDSM where
+  getUniqueM = getUniqueDSM
+
+instance MonadGetUnique USM.UniqSM where
+  getUniqueM = USM.getUniqueM
+
+{-
+Note [Cmm Local Deterministic Uniques]
+~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
+TODO!!!!!
+TODO!!!!!
+-}


=====================================
compiler/GHC/Wasm/ControlFlow/FromCmm.hs
=====================================
@@ -21,10 +21,11 @@ import GHC.Cmm.Dataflow.Label
 import GHC.Cmm.Reducibility
 import GHC.Cmm.Switch
 
+import GHC.Data.Graph.Collapse (MonadUniqDSM (liftUniqDSM))
 import GHC.CmmToAsm.Wasm.Types
 
 import GHC.Platform
-import GHC.Types.Unique.Supply
+import GHC.Types.Unique.DSM
 import GHC.Utils.Misc
 import GHC.Utils.Panic
 import GHC.Utils.Outputable ( Outputable, text, (<+>), ppr
@@ -138,22 +139,20 @@ emptyPost _ = False
 -- | Convert a Cmm CFG to WebAssembly's structured control flow.
 
 structuredControl :: forall expr stmt m .
-                     Applicative m
+                     MonadUniqDSM m
                   => Platform  -- ^ needed for offset calculation
-                  -> UniqSupply
                   -> (Label -> CmmExpr -> m expr) -- ^ translator for expressions
                   -> (Label -> CmmActions -> m stmt) -- ^ translator for straight-line code
                   -> CmmGraph -- ^ CFG to be translated
                   -> m (WasmControl stmt expr '[] '[ 'I32])
-structuredControl platform us txExpr txBlock g' =
-   doTree returns dominatorTree emptyContext
- where
+structuredControl platform txExpr txBlock g' = do
+  gwd :: GraphWithDominators CmmNode <-
+    liftUniqDSM $ asReducible $ graphWithDominators g'
+
+  let
    g :: CmmGraph
    g = gwd_graph gwd
 
-   gwd :: GraphWithDominators CmmNode
-   gwd = initUs_ us $ asReducible $ graphWithDominators g'
-
    dominatorTree :: Tree.Tree CmmBlock-- Dominator tree in which children are sorted
                                        -- with highest reverse-postorder number first
    dominatorTree = fmap blockLabeled $ sortTree $ gwdDominatorTree gwd
@@ -313,7 +312,7 @@ structuredControl platform us txExpr txBlock g' =
    dominates lbl blockname =
        lbl == blockname || dominatorsMember lbl (gwdDominatorsOf gwd blockname)
 
-
+  doTree returns dominatorTree emptyContext
 
 nodeBody :: CmmBlock -> CmmActions
 nodeBody (BlockCC _first middle _last) = middle


=====================================
compiler/ghc.cabal.in
=====================================
@@ -894,6 +894,7 @@ Library
         GHC.Types.Unique
         GHC.Types.Unique.DFM
         GHC.Types.Unique.DSet
+        GHC.Types.Unique.DSM
         GHC.Types.Unique.FM
         GHC.Types.Unique.Map
         GHC.Types.Unique.MemoFun



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

-- 
This project does not include diff previews in email notifications.
View it on GitLab: https://gitlab.haskell.org/ghc/ghc/-/commit/f8ae759bae9cb4d417ecacf590d4a721df2711d8
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/20240703/6d278c9d/attachment-0001.html>


More information about the ghc-commits mailing list