[Git][ghc/ghc][wip/ncg-simd] 2 commits: SIMD: refactor Format datatype

sheaf (@sheaf) gitlab at gitlab.haskell.org
Sat Jun 15 16:35:40 UTC 2024



sheaf pushed to branch wip/ncg-simd at Glasgow Haskell Compiler / GHC


Commits:
84c46f16 by sheaf at 2024-06-15T18:34:58+02:00
SIMD: refactor Format datatype

- - - - -
62168cfc by sheaf at 2024-06-15T18:34:58+02:00
Introduce RegFormat instead of using (Reg, Format)

- - - - -


19 changed files:

- compiler/GHC/CmmToAsm/AArch64/Instr.hs
- compiler/GHC/CmmToAsm/Format.hs
- compiler/GHC/CmmToAsm/Instr.hs
- compiler/GHC/CmmToAsm/PPC.hs
- compiler/GHC/CmmToAsm/PPC/Instr.hs
- compiler/GHC/CmmToAsm/PPC/Regs.hs
- compiler/GHC/CmmToAsm/Reg/Graph.hs
- compiler/GHC/CmmToAsm/Reg/Graph/Coalesce.hs
- compiler/GHC/CmmToAsm/Reg/Graph/Spill.hs
- compiler/GHC/CmmToAsm/Reg/Graph/SpillClean.hs
- compiler/GHC/CmmToAsm/Reg/Graph/SpillCost.hs
- compiler/GHC/CmmToAsm/Reg/Linear.hs
- compiler/GHC/CmmToAsm/Reg/Linear/JoinToTargets.hs
- compiler/GHC/CmmToAsm/Reg/Liveness.hs
- compiler/GHC/CmmToAsm/X86/CodeGen.hs
- compiler/GHC/CmmToAsm/X86/Instr.hs
- compiler/GHC/CmmToAsm/X86/Ppr.hs
- compiler/GHC/Types/Unique/FM.hs
- compiler/GHC/Types/Unique/Set.hs


Changes:

=====================================
compiler/GHC/CmmToAsm/AArch64/Instr.hs
=====================================
@@ -149,8 +149,8 @@ regUsageOfInstr platform instr = case instr of
         -- filtering the usage is necessary, otherwise the register
         -- allocator will try to allocate pre-defined fixed stg
         -- registers as well, as they show up.
-        usage (src, dst) = RU (map (,II64) $ filter (interesting platform) src)
-                              (map (,II64) $ filter (interesting platform) dst)
+        usage (src, dst) = RU (map (\r -> RegFormat r II64) $ filter (interesting platform) src)
+                              (map (\r -> RegFormat r II64) $ filter (interesting platform) dst)
           -- SIMD NCG TODO: the format here is used for register spilling/unspilling.
           -- As the AArch64 NCG does not currently support SIMD registers,
           -- we simply use II64 format for all instructions.


=====================================
compiler/GHC/CmmToAsm/Format.hs
=====================================
@@ -1,7 +1,13 @@
+{-# LANGUAGE DataKinds #-}
 {-# LANGUAGE LambdaCase #-}
+{-# LANGUAGE MagicHash #-}
 {-# LANGUAGE PatternSynonyms #-}
+{-# LANGUAGE UnboxedTuples #-}
 {-# LANGUAGE ViewPatterns #-}
 
+{-# OPTIONS_GHC -Wno-duplicate-exports #-}
+  -- Allow bundling II8, II16... with both Format and ScalarFormat
+
 -- | Formats on this architecture
 --      A Format is a combination of width and class
 --
@@ -13,19 +19,23 @@
 --              properly. eg SPARC doesn't care about FF80.
 --
 module GHC.CmmToAsm.Format (
-    Format(.., IntegerFormat),
+    Format(Format, VecFormat, II8, II16, II32, II64, FF32, FF64, ..),
     ScalarFormat(..),
     intFormat,
     floatFormat,
     isIntFormat,
+    isIntScalarFormat,
     isFloatFormat,
+    isFloatScalarFormat,
     vecFormat,
     isVecFormat,
     cmmTypeFormat,
     formatToWidth,
     formatInBytes,
-    scalarWidth,
-    isIntScalarFormat,
+    RegFormat(..),
+    takeVirtualRegs,
+    takeRealRegs,
+    mapRegFormatSet,
 )
 
 where
@@ -33,9 +43,15 @@ where
 import GHC.Prelude
 
 import GHC.Cmm
+import GHC.Platform.Reg ( Reg(..), RealReg, VirtualReg )
+import GHC.Types.Unique ( Uniquable(..) )
+import GHC.Types.Unique.Set
 import GHC.Utils.Outputable
 import GHC.Utils.Panic
 
+import Data.Coerce
+import Data.Word (Word8)
+
 {- Note [GHC's data format representations]
 ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
 GHC has severals types that represent various aspects of data format.
@@ -71,54 +87,57 @@ These include:
 --        here.  I've removed them from the x86 version, we'll see what happens --SDM
 
 -- ToDo: quite a few occurrences of Format could usefully be replaced by Width
-
-data Format
-        = II8
-        | II16
-        | II32
-        | II64
-        | FF32
-        | FF64
-        | VecFormat !Length       -- ^ number of elements
-                    !ScalarFormat -- ^ format of each element
-        deriving (Show, Eq, Ord)
-
-pattern IntegerFormat :: Format
-pattern IntegerFormat <- ( isIntegerFormat -> True )
-{-# COMPLETE IntegerFormat, FF32, FF64, VecFormat #-}
-
-isIntegerFormat :: Format -> Bool
-isIntegerFormat = \case
-  II8  -> True
-  II16 -> True
-  II32 -> True
-  II64 -> True
-  _    -> False
-
-
+newtype Format = MkFormat { formatBits :: Word8 }
+  deriving (Eq, Ord)
+newtype ScalarFormat = ScalarFormat { scalarFormatBits :: Word8 }
+  deriving (Eq, Ord)
+
+pattern II8, II16, II32, II64, FF32, FF64 :: Coercible a Word8 => a
+pattern II8  <- ( coerce -> ( 0b0_000 :: Word8 ) ) where { II8  = coerce (0b0_000 :: Word8) }
+pattern II16 <- ( coerce -> ( 0b0_001 :: Word8 ) ) where { II16 = coerce (0b0_001 :: Word8) }
+pattern II32 <- ( coerce -> ( 0b0_010 :: Word8 ) ) where { II32 = coerce (0b0_010 :: Word8) }
+pattern II64 <- ( coerce -> ( 0b0_011 :: Word8 ) ) where { II64 = coerce (0b0_011 :: Word8) }
+pattern FF32 <- ( coerce -> ( 0b1_010 :: Word8 ) ) where { FF32 = coerce (0b1_010 :: Word8) }
+pattern FF64 <- ( coerce -> ( 0b1_011 :: Word8 ) ) where { FF64 = coerce (0b1_011 :: Word8) }
+
+pattern Format :: Length -> ScalarFormat -> Format
+pattern Format lg b <- ( getFormat -> (# _, lg, b #) )
+  where
+    Format lg b = MkFormat $ scalarFormatBits b .|. ( fromIntegral ( finiteBitSize lg - 1 - countLeadingZeros lg ) `shiftL` 4 )
+pattern VecFormat :: Length -> ScalarFormat -> Format
+pattern VecFormat lg b <- ( getFormat -> (# True, lg, b #) )
+  where
+    VecFormat lg b = Format lg b
+
+{-# COMPLETE Format :: Format #-}
+{-# COMPLETE II8, II16, II32, II64, FF32, FF64, VecFormat :: Format #-}
+{-# COMPLETE II8, II16, II32, II64, FF32, FF64 :: ScalarFormat #-}
+getFormat :: Format -> (# Bool, Length, ScalarFormat #)
+getFormat ( MkFormat b ) = (# lg > 1, lg, ScalarFormat (b .&. 0b0000_1111) #)
+  where
+    lg = bit ( fromIntegral b `shiftR` 4 )
+
+instance Show ScalarFormat where
+  show = \case
+    II8  -> "II8"
+    II16 -> "II16"
+    II32 -> "II32"
+    II64 -> "II64"
+    FF32 -> "FF32"
+    FF64 -> "FF64"
+instance Show Format where
+  show (Format l f)
+    | l == 1
+    = show f
+    | otherwise
+    = "V" ++ show l ++ show f
 instance Outputable Format where
   ppr fmt = text (show fmt)
 
-data ScalarFormat
-  = FmtInt8
-  | FmtInt16
-  | FmtInt32
-  | FmtInt64
-  | FmtFloat
-  | FmtDouble
-  deriving (Show, Eq, Ord)
-
-isIntScalarFormat :: ScalarFormat -> Bool
-isIntScalarFormat FmtInt8 = True
-isIntScalarFormat FmtInt16 = True
-isIntScalarFormat FmtInt32 = True
-isIntScalarFormat FmtInt64 = True
-isIntScalarFormat _ = False
-
 -- | Get the integer format of this width.
 intFormat :: Width -> Format
 intFormat width
- = case width of
+  = case width of
         W8      -> II8
         W16     -> II16
         W32     -> II32
@@ -130,36 +149,37 @@ intFormat width
 -- | Check if a format represents a vector
 isVecFormat :: Format -> Bool
 isVecFormat (VecFormat {}) = True
-isVecFormat _              = False
+isVecFormat _ = False
 
 -- | Get the float format of this width.
 floatFormat :: Width -> Format
 floatFormat width
  = case width of
-        W32     -> FF32
-        W64     -> FF64
+        W32   -> FF32
+        W64   -> FF64
+        other -> pprPanic "Format.floatFormat" (ppr other)
 
-        other   -> pprPanic "Format.floatFormat" (ppr other)
-
--- | Check if a format represent an integer value.
+-- | Check if a format represents a scalar integer value.
 isIntFormat :: Format -> Bool
-isIntFormat = not . isFloatFormat
+isIntFormat(Format l f)
+  = l == 1 && isIntScalarFormat f
+isIntScalarFormat :: ScalarFormat -> Bool
+isIntScalarFormat = not . isFloatScalarFormat
 
--- | Check if a format represents a floating point value.
+-- | Check if a format represents a scalar floating point value.
 isFloatFormat :: Format -> Bool
-isFloatFormat format
- = case format of
-        FF32    -> True
-        FF64    -> True
-        _       -> False
+isFloatFormat (Format l f)
+  = l == 1 && isFloatScalarFormat f
 
+isFloatScalarFormat :: ScalarFormat -> Bool
+isFloatScalarFormat (ScalarFormat b) = testBit b 3
 
 -- | Convert a Cmm type to a Format.
 cmmTypeFormat :: CmmType -> Format
 cmmTypeFormat ty
-        | isFloatType ty        = floatFormat (typeWidth ty)
-        | isVecType ty          = vecFormat ty
-        | otherwise             = intFormat (typeWidth ty)
+        | isFloatType ty = floatFormat (typeWidth ty)
+        | isVecType ty   = vecFormat ty
+        | otherwise      = intFormat (typeWidth ty)
 
 vecFormat :: CmmType -> Format
 vecFormat ty =
@@ -167,37 +187,65 @@ vecFormat ty =
       elemTy = vecElemType ty
    in if isFloatType elemTy
       then case typeWidth elemTy of
-             W32 -> VecFormat l FmtFloat
-             W64 -> VecFormat l FmtDouble
+             W32 -> Format l FF32
+             W64 -> Format l FF64
              _   -> pprPanic "Incorrect vector element width" (ppr elemTy)
       else case typeWidth elemTy of
-             W8  -> VecFormat l FmtInt8
-             W16 -> VecFormat l FmtInt16
-             W32 -> VecFormat l FmtInt32
-             W64 -> VecFormat l FmtInt64
+             W8  -> Format l II8
+             W16 -> Format l II16
+             W32 -> Format l II32
+             W64 -> Format l II64
              _   -> pprPanic "Incorrect vector element width" (ppr elemTy)
 
 -- | Get the Width of a Format.
 formatToWidth :: Format -> Width
-formatToWidth format
- = case format of
-        II8  -> W8
-        II16 -> W16
-        II32 -> W32
-        II64 -> W64
-        FF32 -> W32
-        FF64 -> W64
-        VecFormat l s ->
-          widthFromBytes (l * widthInBytes (scalarWidth s))
-
-scalarWidth :: ScalarFormat -> Width
-scalarWidth = \case
-  FmtInt8   -> W8
-  FmtInt16  -> W16
-  FmtInt32  -> W32
-  FmtInt64  -> W64
-  FmtFloat  -> W32
-  FmtDouble -> W64
+formatToWidth (Format l f)
+  | l == 1
+  = go f
+  | otherwise
+  = widthFromBytes (l * widthInBytes (go f))
+  where
+    go = \case
+      II8  -> W8
+      II16 -> W16
+      II32 -> W32
+      II64 -> W64
+      FF32 -> W32
+      FF64 -> W64
 
 formatInBytes :: Format -> Int
 formatInBytes = widthInBytes . formatToWidth
+
+--------------------------------------------------------------------------------
+
+-- | A typed register: a register, together with the specific format we
+-- are using it at.
+data RegFormat
+    = RegFormat
+    { regFormatReg :: {-# UNPACK #-} !Reg
+    , regFormatFormat :: !Format
+    }
+
+instance Show RegFormat where
+  show (RegFormat reg fmt) = show reg ++ "::" ++ show fmt
+
+instance Uniquable RegFormat where
+  getUnique = getUnique . regFormatReg
+
+instance Outputable RegFormat where
+  ppr (RegFormat reg fmt) = ppr reg <+> dcolon <+> ppr fmt
+
+-- | Take all the virtual registers from this set.
+takeVirtualRegs :: UniqSet RegFormat -> UniqSet VirtualReg
+takeVirtualRegs = mapMaybeUniqSet_sameUnique $
+  \ case { RegFormat { regFormatReg = RegVirtual vr } -> Just vr; _ -> Nothing }
+  -- See Note [Unique Determinism and code generation]
+
+-- | Take all the real registers from this set.
+takeRealRegs :: UniqSet RegFormat -> UniqSet RealReg
+takeRealRegs = mapMaybeUniqSet_sameUnique $
+  \ case { RegFormat { regFormatReg = RegReal rr } -> Just rr; _ -> Nothing }
+  -- See Note [Unique Determinism and code generation]
+
+mapRegFormatSet :: (Reg -> Reg) -> UniqSet RegFormat -> UniqSet RegFormat
+mapRegFormatSet f = mapUniqSet ( \ ( RegFormat r fmt ) -> RegFormat ( f r ) fmt )


=====================================
compiler/GHC/CmmToAsm/Instr.hs
=====================================
@@ -30,8 +30,8 @@ import GHC.CmmToAsm.Format
 --
 data RegUsage
         = RU    {
-                reads :: [(Reg, Format)],
-                writes :: [(Reg, Format)]
+                reads :: [RegFormat],
+                writes :: [RegFormat]
                 }
         deriving Show
 


=====================================
compiler/GHC/CmmToAsm/PPC.hs
=====================================
@@ -48,8 +48,8 @@ instance Instruction PPC.Instr where
    jumpDestsOfInstr    = PPC.jumpDestsOfInstr
    canFallthroughTo    = PPC.canFallthroughTo
    patchJumpInstr      = PPC.patchJumpInstr
-   mkSpillInstr cfg reg _ i j       = PPC.mkSpillInstr cfg reg i j
-   mkLoadInstr cfg reg _ i j = PPC.mkLoadInstr cfg reg i j
+   mkSpillInstr        = PPC.mkSpillInstr
+   mkLoadInstr         = PPC.mkLoadInstr
    takeDeltaInstr      = PPC.takeDeltaInstr
    isMetaInstr         = PPC.isMetaInstr
    mkRegRegMoveInstr _ = PPC.mkRegRegMoveInstr


=====================================
compiler/GHC/CmmToAsm/PPC/Instr.hs
=====================================
@@ -41,9 +41,7 @@ import GHC.CmmToAsm.PPC.Cond
 import GHC.CmmToAsm.Types
 import GHC.CmmToAsm.Instr (RegUsage(..), noUsage)
 import GHC.CmmToAsm.Format
-import GHC.CmmToAsm.Reg.Target
 import GHC.CmmToAsm.Config
-import GHC.Platform.Reg.Class
 import GHC.Platform.Reg
 
 import GHC.Platform.Regs
@@ -68,8 +66,8 @@ import Data.Maybe (fromMaybe)
 --
 archWordFormat :: Bool -> Format
 archWordFormat is32Bit
- | is32Bit   = II32
- | otherwise = II64
+  | is32Bit   = II32
+  | otherwise = II64
 
 
 mkStackAllocInstr :: Platform -> Int -> [Instr]
@@ -391,8 +389,8 @@ regUsageOfInstr platform instr
     FMADD _ _ rt ra rc rb   -> usage ([ra, rc, rb], [rt])
     _                       -> noUsage
   where
-    usage (src, dst) = RU (map (,II64) $ filter (interesting platform) src)
-                          (map (,II64) $ filter (interesting platform) dst)
+    usage (src, dst) = RU (map (\ r -> RegFormat r II64) $ filter (interesting platform) src)
+                          (map (\ r -> RegFormat r II64) $ filter (interesting platform) dst)
       -- SIMD NCG TODO: the format here is used for register spilling/unspilling.
       -- As the PowerPC NCG does not currently support SIMD registers,
       -- we simply use II64 format for all instructions.
@@ -551,21 +549,16 @@ patchJumpInstr insn patchF
 mkSpillInstr
    :: NCGConfig
    -> Reg       -- register to spill
+   -> Format
    -> Int       -- current stack delta
    -> Int       -- spill slot to use
    -> [Instr]
 
-mkSpillInstr config reg delta slot
+mkSpillInstr config reg fmt delta slot
   = let platform = ncgPlatform config
         off      = spillSlotToOffset platform slot
-        arch     = platformArch platform
     in
-    let fmt = case targetClassOfReg platform reg of
-                RcInteger -> case arch of
-                                ArchPPC -> II32
-                                _       -> II64
-                RcFloatOrVector  -> FF64
-        instr = case makeImmediate W32 True (off-delta) of
+    let instr = case makeImmediate W32 True (off-delta) of
                 Just _  -> ST
                 Nothing -> STFAR -- pseudo instruction: 32 bit offsets
 
@@ -575,21 +568,16 @@ mkSpillInstr config reg delta slot
 mkLoadInstr
    :: NCGConfig
    -> Reg       -- register to load
+   -> Format
    -> Int       -- current stack delta
    -> Int       -- spill slot to use
    -> [Instr]
 
-mkLoadInstr config reg delta slot
+mkLoadInstr config reg fmt delta slot
   = let platform = ncgPlatform config
         off      = spillSlotToOffset platform slot
-        arch     = platformArch platform
     in
-    let fmt = case targetClassOfReg platform reg of
-                RcInteger ->  case arch of
-                                 ArchPPC -> II32
-                                 _       -> II64
-                RcFloatOrVector  -> FF64
-        instr = case makeImmediate W32 True (off-delta) of
+    let instr = case makeImmediate W32 True (off-delta) of
                 Just _  -> LD
                 Nothing -> LDFAR -- pseudo instruction: 32 bit offsets
 


=====================================
compiler/GHC/CmmToAsm/PPC/Regs.hs
=====================================
@@ -108,13 +108,14 @@ realRegSqueeze cls rr
 
 
 mkVirtualReg :: Unique -> Format -> VirtualReg
-mkVirtualReg u format
-   | not (isFloatFormat format) = VirtualRegI u
-   | otherwise
-   = case format of
-        FF32    -> VirtualRegD u
-        FF64    -> VirtualRegD u
-        _       -> panic "mkVirtualReg"
+mkVirtualReg u fmt =
+  case fmt of
+    VecFormat {} -> panic "mkVirtualReg: vector register"
+    _ | not (isFloatFormat fmt)
+      -> VirtualRegI u
+    FF32    -> VirtualRegD u
+    FF64    -> VirtualRegD u
+    _       -> panic "mkVirtualReg"
 
 regDotColor :: RealReg -> SDoc
 regDotColor reg


=====================================
compiler/GHC/CmmToAsm/Reg/Graph.hs
=====================================
@@ -21,6 +21,7 @@ import GHC.CmmToAsm.Reg.Graph.TrivColorable
 import GHC.CmmToAsm.Instr
 import GHC.CmmToAsm.Reg.Target
 import GHC.CmmToAsm.Config
+import GHC.CmmToAsm.Format
 import GHC.CmmToAsm.Types
 import GHC.Platform.Reg.Class
 import GHC.Platform.Reg
@@ -336,21 +337,21 @@ buildGraph platform code
 -- | Add some conflict edges to the graph.
 --   Conflicts between virtual and real regs are recorded as exclusions.
 graphAddConflictSet
-        :: RegMap (Reg, fmt)
+        :: UniqSet RegFormat
         -> Color.Graph VirtualReg RegClass RealReg
         -> Color.Graph VirtualReg RegClass RealReg
 
 graphAddConflictSet regs graph
- = let  virtuals        = mkUniqSet
-                        [ vr | (RegVirtual vr, _) <- nonDetEltsUFM regs ]
+ = let  virtuals = takeVirtualRegs regs
+        reals    = takeRealRegs regs
 
         graph1  = Color.addConflicts virtuals classOfVirtualReg graph
 
         graph2  = foldr (\(r1, r2) -> Color.addExclusion r1 classOfVirtualReg r2)
                         graph1
                         [ (vr, rr)
-                                | (RegVirtual vr, _) <- nonDetEltsUFM regs
-                                , (RegReal    rr, _) <- nonDetEltsUFM regs]
+                        | vr <- nonDetEltsUniqSet virtuals
+                        , rr <- nonDetEltsUniqSet reals ]
                           -- See Note [Unique Determinism and code generation]
 
    in   graph2


=====================================
compiler/GHC/CmmToAsm/Reg/Graph/Coalesce.hs
=====================================
@@ -15,6 +15,8 @@ import GHC.Data.Graph.Directed
 import GHC.Types.Unique.FM
 import GHC.Types.Unique.Supply
 import GHC.Platform (Platform)
+import GHC.Types.Unique.Set
+import GHC.Types.Unique (getUnique)
 
 
 -- | Do register coalescing on this top level thing
@@ -87,8 +89,8 @@ slurpJoinMovs platform live
         slurpLI    rs (LiveInstr _      Nothing)    = rs
         slurpLI    rs (LiveInstr instr (Just live))
                 | Just (r1, r2) <- takeRegRegMoveInstr platform instr
-                , elemUFM r1 $ liveDieRead live
-                , elemUFM r2 $ liveBorn live
+                , elemUniqSet_Directly (getUnique r1) $ liveDieRead live
+                , elemUniqSet_Directly (getUnique r2) $ liveBorn live
 
                 -- only coalesce movs between two virtuals for now,
                 -- else we end up with allocatable regs in the live


=====================================
compiler/GHC/CmmToAsm/Reg/Graph/Spill.hs
=====================================
@@ -17,6 +17,8 @@ import GHC.Cmm hiding (RegSet)
 import GHC.Cmm.BlockId
 import GHC.Cmm.Dataflow.Label
 
+import GHC.Data.List.SetOps (nubOrdBy)
+
 import GHC.Utils.Monad
 import GHC.Utils.Monad.State.Strict
 import GHC.Types.Unique
@@ -27,11 +29,14 @@ import GHC.Utils.Outputable
 import GHC.Utils.Panic
 import GHC.Platform
 
-import Data.List (nub, (\\), intersect)
+import Data.Function ( on )
+import Data.List (intersectBy)
 import Data.Maybe
 import Data.IntSet              (IntSet)
 import qualified Data.IntSet    as IntSet
-import GHC.CmmToAsm.Format
+import GHC.CmmToAsm.Format ( RegFormat(RegFormat, regFormatReg) )
+
+
 
 
 -- | Spill all these virtual regs to stack slots.
@@ -139,7 +144,7 @@ regSpill_top platform regSlotMap cmm
         -- then record the fact that these slots are now live in those blocks
         -- in the given slotmap.
         patchLiveSlot
-                :: BlockMap IntSet -> BlockId -> RegMap (Reg, Format) -> BlockMap IntSet
+                :: BlockMap IntSet -> BlockId -> UniqSet RegFormat-> BlockMap IntSet
 
         patchLiveSlot slotMap blockId regsLive
          = let
@@ -148,9 +153,8 @@ regSpill_top platform regSlotMap cmm
                                 $ mapLookup blockId slotMap
 
                 moreSlotsLive   = IntSet.fromList
-                                $ mapMaybe (lookupUFM regSlotMap)
-                                $ map fst
-                                $ nonDetEltsUFM regsLive
+                                $ mapMaybe (lookupUFM regSlotMap . regFormatReg)
+                                $ nonDetEltsUniqSet regsLive
                     -- See Note [Unique Determinism and code generation]
 
                 slotMap'
@@ -190,18 +194,20 @@ regSpill_instr platform regSlotMap (LiveInstr instr (Just _)) = do
 
   -- sometimes a register is listed as being read more than once,
   --      nub this so we don't end up inserting two lots of spill code.
-  let rsRead_             = nub rlRead
-  let rsWritten_          = nub rlWritten
+  let rsRead_             = nubOrdBy (nonDetCmpUnique `on` getUnique) rlRead
+      rsWritten_          = nubOrdBy (nonDetCmpUnique `on` getUnique) rlWritten
 
   -- if a reg is modified, it appears in both lists, want to undo this..
-  let rsRead              = rsRead_    \\ rsWritten_
-  let rsWritten           = rsWritten_ \\ rsRead_
-  let rsModify            = intersect rsRead_ rsWritten_
+  let rsModify            = intersectBy ((==) `on` getUnique) rsRead_ rsWritten_
+      modified            = mkUniqSet rsModify
+      rsRead              = filter (\ r -> not $ elementOfUniqSet r modified) rsRead_
+      rsWritten           = filter (\ r -> not $ elementOfUniqSet r modified) rsWritten_
+
 
   -- work out if any of the regs being used are currently being spilled.
-  let rsSpillRead         = filter (\(r,_) -> elemUFM r regSlotMap) rsRead
-  let rsSpillWritten      = filter (\(r,_) -> elemUFM r regSlotMap) rsWritten
-  let rsSpillModify       = filter (\(r,_) -> elemUFM r regSlotMap) rsModify
+  let rsSpillRead         = filter (\r -> elemUFM (regFormatReg r) regSlotMap) rsRead
+  let rsSpillWritten      = filter (\r -> elemUFM (regFormatReg r) regSlotMap) rsWritten
+  let rsSpillModify       = filter (\r -> elemUFM (regFormatReg r) regSlotMap) rsModify
 
   -- rewrite the instr and work out spill code.
   (instr1, prepost1)      <- mapAccumLM (spillRead   regSlotMap) instr  rsSpillRead
@@ -226,10 +232,10 @@ spillRead
         :: Instruction instr
         => UniqFM Reg Int
         -> instr
-        -> (Reg, Format)
+        -> RegFormat
         -> SpillM (instr, ([LiveInstr instr'], [LiveInstr instr']))
 
-spillRead regSlotMap instr (reg, fmt)
+spillRead regSlotMap instr (RegFormat reg fmt)
  | Just slot     <- lookupUFM regSlotMap reg
  = do    (instr', nReg)  <- patchInstr reg instr
 
@@ -249,10 +255,10 @@ spillWrite
         :: Instruction instr
         => UniqFM Reg Int
         -> instr
-        -> (Reg, Format)
+        -> RegFormat
         -> SpillM (instr, ([LiveInstr instr'], [LiveInstr instr']))
 
-spillWrite regSlotMap instr (reg, fmt)
+spillWrite regSlotMap instr (RegFormat reg fmt)
  | Just slot     <- lookupUFM regSlotMap reg
  = do    (instr', nReg)  <- patchInstr reg instr
 
@@ -272,10 +278,10 @@ spillModify
         :: Instruction instr
         => UniqFM Reg Int
         -> instr
-        -> (Reg, Format)
+        -> RegFormat
         -> SpillM (instr, ([LiveInstr instr'], [LiveInstr instr']))
 
-spillModify regSlotMap instr (reg, fmt)
+spillModify regSlotMap instr (RegFormat reg fmt)
  | Just slot     <- lookupUFM regSlotMap reg
  = do    (instr', nReg)  <- patchInstr reg instr
 


=====================================
compiler/GHC/CmmToAsm/Reg/Graph/SpillClean.hs
=====================================
@@ -36,6 +36,7 @@ module GHC.CmmToAsm.Reg.Graph.SpillClean (
 import GHC.Prelude
 
 import GHC.CmmToAsm.Reg.Liveness
+import GHC.CmmToAsm.Format
 import GHC.CmmToAsm.Instr
 import GHC.Platform.Reg
 
@@ -215,7 +216,7 @@ cleanForward platform blockId assoc acc (li : instrs)
         -- Writing to a reg changes its value.
         | LiveInstr instr _     <- li
         , RU _ written          <- regUsageOfInstr platform instr
-        = let assoc'    = foldr delAssoc assoc (map SReg $ nub $ map fst written)
+        = let assoc'    = foldr delAssoc assoc (map SReg $ nub $ map regFormatReg written)
           in  cleanForward platform blockId assoc' (li : acc) instrs
 
 


=====================================
compiler/GHC/CmmToAsm/Reg/Graph/SpillCost.hs
=====================================
@@ -1,4 +1,5 @@
 {-# LANGUAGE GADTs               #-}
+{-# LANGUAGE LambdaCase          #-}
 {-# LANGUAGE ScopedTypeVariables #-}
 
 module GHC.CmmToAsm.Reg.Graph.SpillCost (
@@ -34,6 +35,7 @@ import GHC.Utils.Panic
 import GHC.Platform
 import GHC.Utils.Monad.State.Strict
 import GHC.CmmToAsm.CFG
+import GHC.CmmToAsm.Format
 
 import Data.List        (nub, minimumBy)
 import Data.Maybe
@@ -99,7 +101,7 @@ slurpSpillCostInfo platform cfg cmm
         countBlock info freqMap (BasicBlock blockId instrs)
                 | LiveInfo _ _ blockLive _ <- info
                 , Just rsLiveEntry  <- mapLookup blockId blockLive
-                , rsLiveEntry_virt  <- takeVirtuals rsLiveEntry
+                , rsLiveEntry_virt  <- takeVirtualRegs rsLiveEntry
                 = countLIs (ceiling $ blockFreq freqMap blockId) rsLiveEntry_virt instrs
 
                 | otherwise
@@ -129,13 +131,13 @@ slurpSpillCostInfo platform cfg cmm
 
                 -- Increment counts for what regs were read/written from.
                 let (RU read written)   = regUsageOfInstr platform instr
-                mapM_ (incUses scale) $ mapMaybe takeVirtualReg $ nub $ map fst read
-                mapM_ (incDefs scale) $ mapMaybe takeVirtualReg $ nub $ map fst written
+                mapM_ (incUses scale) $ nub $ mapMaybe (takeVirtualReg . regFormatReg) read
+                mapM_ (incDefs scale) $ nub $ mapMaybe (takeVirtualReg . regFormatReg) written
 
                 -- Compute liveness for entry to next instruction.
-                let liveDieRead_virt    = takeVirtuals (liveDieRead  live)
-                let liveDieWrite_virt   = takeVirtuals (liveDieWrite live)
-                let liveBorn_virt       = takeVirtuals (liveBorn     live)
+                let liveDieRead_virt    = takeVirtualRegs (liveDieRead  live)
+                let liveDieWrite_virt   = takeVirtualRegs (liveDieWrite live)
+                let liveBorn_virt       = takeVirtualRegs (liveBorn     live)
 
                 let rsLiveAcross
                         = rsLiveEntry `minusUniqSet` liveDieRead_virt
@@ -157,13 +159,6 @@ slurpSpillCostInfo platform cfg cmm
           | otherwise
           = 1.0 -- Only if no cfg given
 
--- | Take all the virtual registers from this set.
-takeVirtuals :: RegMap (Reg, fmt) -> UniqSet VirtualReg
-takeVirtuals m = mkUniqSet
-  [ vr | (RegVirtual vr, _) <- nonDetEltsUFM m ]
-  -- See Note [Unique Determinism and code generation]
-
-
 -- | Choose a node to spill from this graph
 chooseSpill
         :: SpillCostInfo


=====================================
compiler/GHC/CmmToAsm/Reg/Linear.hs
=====================================
@@ -117,6 +117,7 @@ import GHC.CmmToAsm.Reg.Liveness
 import GHC.CmmToAsm.Reg.Utils
 import GHC.CmmToAsm.Instr
 import GHC.CmmToAsm.Config
+import GHC.CmmToAsm.Format
 import GHC.CmmToAsm.Types
 import GHC.Platform.Reg
 import GHC.Platform.Reg.Class (RegClass(..))
@@ -128,6 +129,7 @@ import GHC.Cmm hiding (RegSet)
 import GHC.Data.Graph.Directed
 import GHC.Types.Unique
 import GHC.Types.Unique.FM
+import GHC.Types.Unique.Set
 import GHC.Types.Unique.Supply
 import GHC.Utils.Outputable
 import GHC.Utils.Panic
@@ -136,7 +138,6 @@ import GHC.Platform
 import Data.Maybe
 import Data.List (partition, nub)
 import Control.Monad
-import GHC.CmmToAsm.Format
 
 -- -----------------------------------------------------------------------------
 -- Top level of the register allocator
@@ -203,7 +204,7 @@ linearRegAlloc
         :: forall instr. (Instruction instr)
         => NCGConfig
         -> [BlockId] -- ^ entry points
-        -> BlockMap (UniqFM Reg (Reg, Format))
+        -> BlockMap (UniqSet RegFormat)
               -- ^ live regs on entry to each basic block
         -> [SCC (LiveBasicBlock instr)]
               -- ^ instructions annotated with "deaths"
@@ -242,7 +243,7 @@ linearRegAlloc'
         => NCGConfig
         -> freeRegs
         -> [BlockId]                    -- ^ entry points
-        -> BlockMap (UniqFM Reg (Reg, Format))              -- ^ live regs on entry to each basic block
+        -> BlockMap (UniqSet RegFormat)              -- ^ live regs on entry to each basic block
         -> [SCC (LiveBasicBlock instr)] -- ^ instructions annotated with "deaths"
         -> UniqSM ([NatBasicBlock instr], RegAllocStats, Int)
 
@@ -256,7 +257,7 @@ linearRegAlloc' config initFreeRegs entry_ids block_live sccs
 
 linearRA_SCCs :: OutputableRegConstraint freeRegs instr
               => [BlockId]
-              -> BlockMap (UniqFM Reg (Reg, Format))
+              -> BlockMap (UniqSet RegFormat)
               -> [NatBasicBlock instr]
               -> [SCC (LiveBasicBlock instr)]
               -> RegM freeRegs [NatBasicBlock instr]
@@ -291,7 +292,7 @@ linearRA_SCCs entry_ids block_live blocksAcc (CyclicSCC blocks : sccs)
 
 process :: forall freeRegs instr. (OutputableRegConstraint freeRegs instr)
         => [BlockId]
-        -> BlockMap (UniqFM Reg (Reg, Format))
+        -> BlockMap (UniqSet RegFormat)
         -> [GenBasicBlock (LiveInstr instr)]
         -> RegM freeRegs [[NatBasicBlock instr]]
 process entry_ids block_live =
@@ -330,7 +331,7 @@ process entry_ids block_live =
 --
 processBlock
         :: OutputableRegConstraint freeRegs instr
-        => BlockMap (UniqFM Reg (Reg, Format))              -- ^ live regs on entry to each basic block
+        => BlockMap (UniqSet RegFormat)              -- ^ live regs on entry to each basic block
         -> LiveBasicBlock instr         -- ^ block to do register allocation on
         -> RegM freeRegs [NatBasicBlock instr]   -- ^ block with registers allocated
 
@@ -347,7 +348,7 @@ processBlock block_live (BasicBlock id instrs)
 -- | Load the freeregs and current reg assignment into the RegM state
 --      for the basic block with this BlockId.
 initBlock :: FR freeRegs
-          => BlockId -> BlockMap (UniqFM Reg (Reg, Format)) -> RegM freeRegs ()
+          => BlockId -> BlockMap (UniqSet RegFormat) -> RegM freeRegs ()
 initBlock id block_live
  = do   platform    <- getPlatform
         block_assig <- getBlockAssigR
@@ -364,7 +365,7 @@ initBlock id block_live
                             setFreeRegsR    (frInitFreeRegs platform)
                           Just live ->
                             setFreeRegsR $ foldl' (flip $ frAllocateReg platform) (frInitFreeRegs platform)
-                                                  [ r | ( RegReal r, _ ) <- nonDetEltsUFM live ]
+                                                  (nonDetEltsUniqSet $ takeRealRegs live)
                             -- See Note [Unique Determinism and code generation]
                         setAssigR       emptyRegMap
 
@@ -377,7 +378,7 @@ initBlock id block_live
 -- | Do allocation for a sequence of instructions.
 linearRA
         :: forall freeRegs instr. (OutputableRegConstraint freeRegs instr)
-        => BlockMap (UniqFM Reg (Reg, Format))                      -- ^ map of what vregs are live on entry to each block.
+        => BlockMap (UniqSet RegFormat)                      -- ^ map of what vregs are live on entry to each block.
         -> BlockId                              -- ^ id of the current block, for debugging.
         -> [LiveInstr instr]                    -- ^ liveness annotated instructions in this block.
         -> RegM freeRegs
@@ -402,7 +403,7 @@ linearRA block_live block_id = go [] []
 -- | Do allocation for a single instruction.
 raInsn
         :: OutputableRegConstraint freeRegs instr
-        => BlockMap (UniqFM Reg (Reg, Format))                      -- ^ map of what vregs are love on entry to each block.
+        => BlockMap (UniqSet RegFormat)                      -- ^ map of what vregs are love on entry to each block.
         -> [instr]                              -- ^ accumulator for instructions already processed.
         -> BlockId                              -- ^ the id of the current block, for debugging
         -> LiveInstr instr                      -- ^ the instr to have its regs allocated, with liveness info.
@@ -433,7 +434,7 @@ raInsn block_live new_instrs id (LiveInstr (Instr instr) (Just live))
     -- (we can't eliminate it if the source register is on the stack, because
     --  we do not want to use one spill slot for different virtual registers)
     case takeRegRegMoveInstr platform instr of
-        Just (src,dst)  | Just (_, fmt) <- lookupUFM (liveDieRead live) src,
+        Just (src,dst)  | Just (RegFormat _ fmt) <- lookupUniqSet_Directly (liveDieRead live) (getUnique src),
                           isVirtualReg dst,
                           not (dst `elemUFM` assig),
                           isRealReg src || isInReg src assig -> do
@@ -457,8 +458,8 @@ raInsn block_live new_instrs id (LiveInstr (Instr instr) (Just live))
            return (new_instrs, [])
 
         _ -> genRaInsn block_live new_instrs id instr
-                        (map fst $ nonDetEltsUFM $ liveDieRead live)
-                        (map fst $ nonDetEltsUFM $ liveDieWrite live)
+                        (map regFormatReg $ nonDetEltsUniqSet $ liveDieRead live)
+                        (map regFormatReg $ nonDetEltsUniqSet $ liveDieWrite live)
                         -- See Note [Unique Determinism and code generation]
 
 raInsn _ _ _ instr
@@ -487,7 +488,7 @@ isInReg src assig | Just (InReg _) <- lookupUFM assig src = True
 
 genRaInsn :: forall freeRegs instr.
              (OutputableRegConstraint freeRegs instr)
-          => BlockMap (UniqFM Reg (Reg, Format))
+          => BlockMap (UniqSet RegFormat)
           -> [instr]
           -> BlockId
           -> instr
@@ -500,13 +501,13 @@ genRaInsn block_live new_instrs block_id instr r_dying w_dying = do
   platform <- getPlatform
   case regUsageOfInstr platform instr of { RU read written ->
     do
-    let real_written    = [ rr  | (RegReal     rr, _) <- written ] :: [RealReg]
-    let virt_written    = [ vr  | (RegVirtual  vr, _) <- written ]
+    let real_written    = [ rr  | RegFormat { regFormatReg = RegReal     rr } <- written ] :: [RealReg]
+    let virt_written    = [ vr  | RegFormat { regFormatReg = RegVirtual  vr } <- written ]
 
     -- we don't need to do anything with real registers that are
     -- only read by this instr.  (the list is typically ~2 elements,
     -- so using nub isn't a problem).
-    let virt_read       = nub [ vr      | (RegVirtual vr, _) <- read ] :: [VirtualReg]
+    let virt_read       = nub [ vr | RegFormat { regFormatReg = RegVirtual vr }<- read ] :: [VirtualReg]
 
 --     do
 --         let real_read       = nub [ rr      | (RegReal rr) <- read]
@@ -872,7 +873,7 @@ allocRegsAndSpill_spill reading keep spills alloc r rs assig spill_loc
         let regclass = classOfVirtualReg r
             freeRegs_thisClass  = frGetFreeRegs platform regclass freeRegs :: [RealReg]
             vr_fmt = case r of
-                VirtualRegV128 {} -> VecFormat 2 FmtDouble
+                VirtualRegV128 {} -> VecFormat 2 FF64
                 -- It doesn't really matter whether we use e.g. v2f64 or v4f32
                 -- or v4i32 etc here. This is perhaps a sign that 'Format'
                 -- is not the right type to use here, but that is a battle


=====================================
compiler/GHC/CmmToAsm/Reg/Linear/JoinToTargets.hs
=====================================
@@ -32,13 +32,14 @@ import GHC.Types.Unique.FM
 
 import GHC.Utils.Outputable
 import GHC.CmmToAsm.Format
+import GHC.Types.Unique.Set
 
 -- | For a jump instruction at the end of a block, generate fixup code so its
 --      vregs are in the correct regs for its destination.
 --
 joinToTargets
         :: (FR freeRegs, Instruction instr)
-        => BlockMap (RegMap (Reg, Format)) -- ^ maps the unique of the blockid to the set of vregs
+        => BlockMap (UniqSet RegFormat) -- ^ maps the unique of the blockid to the set of vregs
                                         --      that are known to be live on the entry to each block.
 
         -> BlockId                      -- ^ id of the current block
@@ -62,7 +63,7 @@ joinToTargets block_live id instr
 -----
 joinToTargets'
         :: (FR freeRegs, Instruction instr)
-        => BlockMap (RegMap (Reg, Format))               -- ^ maps the unique of the blockid to the set of vregs
+        => BlockMap (UniqSet RegFormat) -- ^ maps the unique of the blockid to the set of vregs
                                         --      that are known to be live on the entry to each block.
 
         -> [NatBasicBlock instr]        -- ^ acc blocks of fixup code.
@@ -90,7 +91,7 @@ joinToTargets' block_live new_blocks block_id instr (dest:dests)
         -- adjust the current assignment to remove any vregs that are not live
         -- on entry to the destination block.
         let Just live_set       = mapLookup dest block_live
-        let still_live uniq _   = uniq `elemUFM_Directly` live_set
+        let still_live uniq _   = uniq `elemUniqSet_Directly` live_set
         let adjusted_assig      = filterUFM_Directly still_live assig
 
         -- and free up those registers which are now free.
@@ -99,7 +100,7 @@ joinToTargets' block_live new_blocks block_id instr (dest:dests)
                         -- This is non-deterministic but we do not
                         -- currently support deterministic code-generation.
                         -- See Note [Unique Determinism and code generation]
-                        , not (elemUFM_Directly reg live_set)
+                        , not (elemUniqSet_Directly reg live_set)
                         , r          <- regsOfLoc loc ]
 
         case lookupBlockAssignment  dest block_assig of
@@ -116,7 +117,7 @@ joinToTargets' block_live new_blocks block_id instr (dest:dests)
 
 -- this is the first time we jumped to this block.
 joinToTargets_first :: (FR freeRegs, Instruction instr)
-                    => BlockMap (UniqFM Reg (Reg, Format))
+                    => BlockMap (UniqSet RegFormat)
                     -> [NatBasicBlock instr]
                     -> BlockId
                     -> instr
@@ -145,7 +146,7 @@ joinToTargets_first block_live new_blocks block_id instr dest dests
 
 -- we've jumped to this block before
 joinToTargets_again :: (Instruction instr, FR freeRegs)
-                    => BlockMap (UniqFM Reg (Reg, Format))
+                    => BlockMap (UniqSet RegFormat)
                     -> [NatBasicBlock instr]
                     -> BlockId
                     -> instr


=====================================
compiler/GHC/CmmToAsm/Reg/Liveness.hs
=====================================
@@ -64,6 +64,7 @@ import Data.List (mapAccumL, partition)
 import Data.Maybe
 import Data.IntSet              (IntSet)
 import GHC.CmmToAsm.Format
+import GHC.Types.Unique (Uniquable(..))
 
 -----------------------------------------------------------------------------
 type RegSet = UniqSet Reg
@@ -110,8 +111,8 @@ instance Instruction instr => Instruction (InstrSR instr) where
         regUsageOfInstr platform i
          = case i of
                 Instr  instr    -> regUsageOfInstr platform instr
-                SPILL  reg fmt _    -> RU [(reg, fmt)] []
-                RELOAD _ reg fmt    -> RU [] [(reg, fmt)]
+                SPILL  reg fmt _    -> RU [RegFormat reg fmt] []
+                RELOAD _ reg fmt    -> RU [] [RegFormat reg fmt]
 
         patchRegsOfInstr i f
          = case i of
@@ -187,9 +188,9 @@ data LiveInstr instr
 
 data Liveness
         = Liveness
-        { liveBorn      :: RegMap (Reg, Format)       -- ^ registers born in this instruction (written to for first time).
-        , liveDieRead   :: RegMap (Reg, Format)       -- ^ registers that died because they were read for the last time.
-        , liveDieWrite  :: RegMap (Reg, Format) }     -- ^ registers that died because they were clobbered by something.
+        { liveBorn      :: UniqSet RegFormat      -- ^ registers born in this instruction (written to for first time).
+        , liveDieRead   :: UniqSet RegFormat      -- ^ registers that died because they were read for the last time.
+        , liveDieWrite  :: UniqSet RegFormat}     -- ^ registers that died because they were clobbered by something.
 
 
 -- | Stash regs live on entry to each basic block in the info part of the cmm code.
@@ -198,7 +199,7 @@ data LiveInfo
                 (LabelMap RawCmmStatics)  -- cmm info table static stuff
                 [BlockId]                 -- entry points (first one is the
                                           -- entry point for the proc).
-                (BlockMap (UniqFM Reg (Reg, Format)))         -- argument locals live on entry to this block
+                (BlockMap (UniqSet RegFormat))         -- argument locals live on entry to this block
                 (BlockMap IntSet)         -- stack slots live on entry to this block
 
 
@@ -244,11 +245,11 @@ instance Outputable instr
                         , pprRegs (text "# w_dying: ") (liveDieWrite live) ]
                     $+$ space)
 
-         where  pprRegs :: Outputable a => SDoc -> RegMap a -> SDoc
+         where  pprRegs :: SDoc -> UniqSet RegFormat -> SDoc
                 pprRegs name regs
-                 | isNullUFM regs  = empty
+                 | isEmptyUniqSet regs  = empty
                  | otherwise            = name <>
-                     (pprUFM regs (hcat . punctuate space . map ppr))
+                     (pprUFM (getUniqSet regs) (hcat . punctuate space . map ppr))
 
 instance OutputableP env instr => OutputableP env (LiveInstr instr) where
    pdoc env i = ppr (fmap (pdoc env) i)
@@ -328,7 +329,7 @@ slurpConflicts
         :: Instruction instr
         => Platform
         -> LiveCmmDecl statics instr
-        -> (Bag (UniqFM Reg (Reg, Format)), Bag (Reg, Reg))
+        -> (Bag (UniqSet RegFormat), Bag (Reg, Reg))
 
 slurpConflicts platform live
         = slurpCmm (emptyBag, emptyBag) live
@@ -362,23 +363,23 @@ slurpConflicts platform live
          = let
                 -- regs that die because they are read for the last time at the start of an instruction
                 --      are not live across it.
-                rsLiveAcross    = rsLiveEntry `minusUFM` (liveDieRead live)
+                rsLiveAcross    = rsLiveEntry `minusUniqSet` (liveDieRead live)
 
                 -- regs live on entry to the next instruction.
                 --      be careful of orphans, make sure to delete dying regs _after_ unioning
                 --      in the ones that are born here.
-                rsLiveNext      = (rsLiveAcross `plusUFM` (liveBorn     live))
-                                                `minusUFM`  (liveDieWrite live)
+                rsLiveNext      = (rsLiveAcross `unionUniqSets` (liveBorn     live))
+                                                `minusUniqSet`  (liveDieWrite live)
 
                 -- orphan vregs are the ones that die in the same instruction they are born in.
                 --      these are likely to be results that are never used, but we still
                 --      need to assign a hreg to them..
-                rsOrphans       = intersectUFM
+                rsOrphans       = intersectUniqSets
                                         (liveBorn live)
-                                        (plusUFM (liveDieWrite live) (liveDieRead live))
+                                        (unionUniqSets (liveDieWrite live) (liveDieRead live))
 
                 --
-                rsConflicts     = plusUFM rsLiveNext rsOrphans
+                rsConflicts     = unionUniqSets rsLiveNext rsOrphans
 
           in    case takeRegRegMoveInstr platform instr of
                  Just rr        -> slurpLIs rsLiveNext
@@ -622,7 +623,7 @@ patchEraseLive platform patchF cmm
          | LiveInfo static id blockMap mLiveSlots <- info
          = let
                   -- See Note [Unique Determinism and code generation]
-                blockMap'       = mapMap (mapKeysUFM patchF) blockMap
+                blockMap'       = mapMap (mapRegFormatSet patchF) blockMap
 
                 info'           = LiveInfo static id blockMap' mLiveSlots
            in   CmmProc info' label live $ map patchSCC sccs
@@ -651,8 +652,8 @@ patchEraseLive platform patchF cmm
                 | r1 == r2      = True
 
                 -- destination reg is never used
-                | elemUFM r2 (liveBorn live)
-                , elemUFM r2 (liveDieRead live) || elemUFM r2 (liveDieWrite live)
+                | elemUniqSet_Directly (getUnique r2) (liveBorn live)
+                , elemUniqSet_Directly (getUnique r2) (liveDieRead live) || elemUniqSet_Directly (getUnique r2) (liveDieWrite live)
                 = True
 
                 | otherwise     = False
@@ -675,9 +676,9 @@ patchRegsLiveInstr patchF li
                 (patchRegsOfInstr instr patchF)
                 (Just live
                         { -- WARNING: have to go via lists here because patchF changes the uniq in the Reg
-                          liveBorn      = mapKeysUFM patchF $ liveBorn live
-                        , liveDieRead   = mapKeysUFM patchF $ liveDieRead live
-                        , liveDieWrite  = mapKeysUFM patchF $ liveDieWrite live })
+                          liveBorn      = mapRegFormatSet patchF $ liveBorn live
+                        , liveDieRead   = mapRegFormatSet patchF $ liveDieRead live
+                        , liveDieWrite  = mapRegFormatSet patchF $ liveDieWrite live })
                           -- See Note [Unique Determinism and code generation]
 
 --------------------------------------------------------------------------------
@@ -867,7 +868,7 @@ computeLiveness
         -> [SCC (LiveBasicBlock instr)]
         -> ([SCC (LiveBasicBlock instr)],       -- instructions annotated with list of registers
                                                 -- which are "dead after this instruction".
-               BlockMap (UniqFM Reg (Reg, Format)))                 -- blocks annotated with set of live registers
+               BlockMap (UniqSet RegFormat))                 -- blocks annotated with set of live registers
                                                 -- on entry to the block.
 
 computeLiveness platform sccs
@@ -882,11 +883,11 @@ computeLiveness platform sccs
 livenessSCCs
        :: Instruction instr
        => Platform
-       -> BlockMap (UniqFM Reg (Reg, Format))
+       -> BlockMap (UniqSet RegFormat)
        -> [SCC (LiveBasicBlock instr)]          -- accum
        -> [SCC (LiveBasicBlock instr)]
        -> ( [SCC (LiveBasicBlock instr)]
-          , BlockMap (UniqFM Reg (Reg, Format)))
+          , BlockMap (UniqSet RegFormat))
 
 livenessSCCs _ blockmap done []
         = (done, blockmap)
@@ -915,8 +916,8 @@ livenessSCCs platform blockmap done
 
             linearLiveness
                 :: Instruction instr
-                => BlockMap (UniqFM Reg (Reg, Format)) -> [LiveBasicBlock instr]
-                -> (BlockMap (UniqFM Reg (Reg, Format)), [LiveBasicBlock instr])
+                => BlockMap (UniqSet RegFormat) -> [LiveBasicBlock instr]
+                -> (BlockMap (UniqSet RegFormat), [LiveBasicBlock instr])
 
             linearLiveness = mapAccumL (livenessBlock platform)
 
@@ -935,14 +936,14 @@ livenessSCCs platform blockmap done
 livenessBlock
         :: Instruction instr
         => Platform
-        -> BlockMap (UniqFM Reg (Reg, Format))
+        -> BlockMap (UniqSet RegFormat)
         -> LiveBasicBlock instr
-        -> (BlockMap (UniqFM Reg (Reg, Format)), LiveBasicBlock instr)
+        -> (BlockMap (UniqSet RegFormat), LiveBasicBlock instr)
 
 livenessBlock platform blockmap (BasicBlock block_id instrs)
  = let
         (regsLiveOnEntry, instrs1)
-            = livenessBack platform emptyUFM blockmap [] (reverse instrs)
+            = livenessBack platform emptyUniqSet blockmap [] (reverse instrs)
         blockmap'       = mapInsert block_id regsLiveOnEntry blockmap
 
         instrs2         = livenessForward platform regsLiveOnEntry instrs1
@@ -957,7 +958,7 @@ livenessBlock platform blockmap (BasicBlock block_id instrs)
 livenessForward
         :: Instruction instr
         => Platform
-        -> UniqFM Reg (Reg, Format) -- regs live on this instr
+        -> UniqSet RegFormat -- regs live on this instr
         -> [LiveInstr instr] -> [LiveInstr instr]
 
 livenessForward _        _           []  = []
@@ -967,14 +968,13 @@ livenessForward platform rsLiveEntry (li@(LiveInstr instr mLive) : lis)
                 RU _ written  = regUsageOfInstr platform instr
                 -- Regs that are written to but weren't live on entry to this instruction
                 --      are recorded as being born here.
-                rsBorn          = listToUFM
-                                $ map ( \ ( r, fmt ) -> ( r, ( r, fmt ) ) )
-                                $ filter (\( r, _) -> not $ elemUFM r rsLiveEntry)
+                rsBorn          = mkUniqSet
+                                $ filter (\ r -> not $ elemUniqSet_Directly (getUnique r) rsLiveEntry)
                                 $ written
 
-                rsLiveNext      = (rsLiveEntry `plusUFM` rsBorn)
-                                        `minusUFM` (liveDieRead live)
-                                        `minusUFM` (liveDieWrite live)
+                rsLiveNext      = (rsLiveEntry `unionUniqSets` rsBorn)
+                                        `minusUniqSet` (liveDieRead live)
+                                        `minusUniqSet` (liveDieWrite live)
 
         in LiveInstr instr (Just live { liveBorn = rsBorn })
                 : livenessForward platform rsLiveNext lis
@@ -989,11 +989,11 @@ livenessForward platform rsLiveEntry (li@(LiveInstr instr mLive) : lis)
 livenessBack
         :: Instruction instr
         => Platform
-        -> UniqFM Reg (Reg, Format)            -- regs live on this instr
-        -> BlockMap (UniqFM Reg (Reg, Format)) -- regs live on entry to other BBs
+        -> UniqSet RegFormat            -- regs live on this instr
+        -> BlockMap (UniqSet RegFormat) -- regs live on entry to other BBs
         -> [LiveInstr instr]            -- instructions (accum)
         -> [LiveInstr instr]            -- instructions
-        -> (UniqFM Reg (Reg, Format), [LiveInstr instr])
+        -> (UniqSet RegFormat, [LiveInstr instr])
 
 livenessBack _        liveregs _        done []  = (liveregs, done)
 
@@ -1006,10 +1006,10 @@ livenessBack platform liveregs blockmap acc (instr : instrs)
 liveness1
         :: Instruction instr
         => Platform
-        -> UniqFM Reg (Reg, Format)
-        -> BlockMap (UniqFM Reg (Reg, Format))
+        -> UniqSet RegFormat
+        -> BlockMap (UniqSet RegFormat)
         -> LiveInstr instr
-        -> (UniqFM Reg (Reg, Format), LiveInstr instr)
+        -> (UniqSet RegFormat, LiveInstr instr)
 
 liveness1 _ liveregs _ (LiveInstr instr _)
         | isMetaInstr instr
@@ -1020,14 +1020,14 @@ liveness1 platform liveregs blockmap (LiveInstr instr _)
         | not_a_branch
         = (liveregs1, LiveInstr instr
                         (Just $ Liveness
-                        { liveBorn      = emptyUFM
+                        { liveBorn      = emptyUniqSet
                         , liveDieRead   = r_dying
                         , liveDieWrite  = w_dying }))
 
         | otherwise
         = (liveregs_br, LiveInstr instr
                         (Just $ Liveness
-                        { liveBorn      = emptyUFM
+                        { liveBorn      = emptyUniqSet
                         , liveDieRead   = r_dying_br
                         , liveDieWrite  = w_dying }))
 
@@ -1036,18 +1036,21 @@ liveness1 platform liveregs blockmap (LiveInstr instr _)
 
             -- registers that were written here are dead going backwards.
             -- registers that were read here are live going backwards.
-            liveregs1   = (liveregs `delListFromUFM` (map fst written))
-                                    `addListToUFM` (map (\(r, fmt) -> (r, (r,fmt))) read)
+            liveregs1   = (liveregs `delListFromUniqSet` written)
+                                    `addListToUniqSet` read
 
             -- registers that are not live beyond this point, are recorded
             --  as dying here.
-            r_dying     = listToUFM
-                          [ (reg, (reg, fmt)) | (reg, fmt) <- read, reg `notElem` map fst written,
-                              not (elemUFM reg liveregs) ]
+            r_dying     = mkUniqSet
+                          [ reg
+                          | reg@(RegFormat r _) <- read
+                          , not $ any (\ w -> getUnique w == getUnique r) written
+                          , not (elementOfUniqSet reg liveregs) ]
 
-            w_dying     = listToUFM
-                          [ (reg, (reg, fmt)) | (reg, fmt) <- written,
-                             not (elemUFM reg liveregs) ]
+            w_dying     = mkUniqSet
+                          [ reg
+                          | reg <- written
+                          , not (elementOfUniqSet reg liveregs) ]
 
             -- union in the live regs from all the jump destinations of this
             -- instruction.
@@ -1057,14 +1060,14 @@ liveness1 platform liveregs blockmap (LiveInstr instr _)
             targetLiveRegs target
                   = case mapLookup target blockmap of
                                 Just ra -> ra
-                                Nothing -> emptyUFM
+                                Nothing -> emptyUniqSet
 
-            live_from_branch = plusUFMList (map targetLiveRegs targets)
+            live_from_branch = unionManyUniqSets (map targetLiveRegs targets)
 
-            liveregs_br = liveregs1 `plusUFM` live_from_branch
+            liveregs_br = liveregs1 `unionUniqSets` live_from_branch
 
             -- registers that are live only in the branch targets should
             -- be listed as dying here.
-            live_branch_only = live_from_branch `minusUFM` liveregs
-            r_dying_br  = r_dying `plusUFM` live_branch_only
+            live_branch_only = live_from_branch `minusUniqSet` liveregs
+            r_dying_br  = (r_dying `unionUniqSets` live_branch_only)
                           -- See Note [Unique Determinism and code generation]


=====================================
compiler/GHC/CmmToAsm/X86/CodeGen.hs
=====================================
@@ -1220,12 +1220,12 @@ getRegister' platform is32Bit (CmmMachOp mop [x]) = do -- unary MachOps
 
         vector_float_negate_avx :: Length -> Width -> CmmExpr -> NatM Register
         vector_float_negate_avx l w expr = do
-          tmp                  <- getNewRegNat (VecFormat l FmtFloat)
+          tmp                  <- getNewRegNat (VecFormat l FF32)
           (reg, exp)           <- getSomeReg expr
           Amode addr addr_code <- memConstant (mkAlignment $ widthInBytes W32) (CmmFloat 0.0 W32)
           let format   = case w of
-                           W32 -> VecFormat l FmtFloat
-                           W64 -> VecFormat l FmtDouble
+                           W32 -> VecFormat l FF32
+                           W64 -> VecFormat l FF64
                            _ -> pprPanic "Cannot negate vector of width" (ppr w)
               code dst = case w of
                            W32 -> exp `appOL` addr_code `snocOL`
@@ -1240,11 +1240,11 @@ getRegister' platform is32Bit (CmmMachOp mop [x]) = do -- unary MachOps
 
         vector_float_negate_sse :: Length -> Width -> CmmExpr -> NatM Register
         vector_float_negate_sse l w expr = do
-          tmp                  <- getNewRegNat (VecFormat l FmtFloat)
+          tmp                  <- getNewRegNat (VecFormat l FF32)
           (reg, exp)           <- getSomeReg expr
           let format   = case w of
-                           W32 -> VecFormat l FmtFloat
-                           W64 -> VecFormat l FmtDouble
+                           W32 -> VecFormat l FF32
+                           W64 -> VecFormat l FF64
                            _ -> pprPanic "Cannot negate vector of width" (ppr w)
               code dst = exp `snocOL`
                          (XOR format (OpReg tmp) (OpReg tmp)) `snocOL`
@@ -1260,7 +1260,7 @@ getRegister' platform is32Bit (CmmMachOp mop [x]) = do -- unary MachOps
         vector_float_broadcast_avx len W32 expr
           = do
           (reg, exp) <- getSomeReg expr
-          let f    = VecFormat len FmtFloat
+          let f    = VecFormat len FF32
               addr = spRel platform 0
            in return $ Any f (\dst -> exp    `snocOL`
                                     (MOVU f (OpReg reg) (OpAddr addr)) `snocOL`
@@ -1268,7 +1268,7 @@ getRegister' platform is32Bit (CmmMachOp mop [x]) = do -- unary MachOps
         vector_float_broadcast_avx len W64 expr
           = do
           (reg, exp) <- getSomeReg expr
-          let f    = VecFormat len FmtDouble
+          let f    = VecFormat len FF64
               addr = spRel platform 0
            in return $ Any f (\dst -> exp `snocOL`
                                     (MOVU f (OpReg reg) (OpAddr addr)) `snocOL`
@@ -1284,7 +1284,7 @@ getRegister' platform is32Bit (CmmMachOp mop [x]) = do -- unary MachOps
         vector_float_broadcast_sse len W32 expr
           = do
           (reg, exp) <- getSomeReg expr
-          let f        = VecFormat len FmtFloat
+          let f        = VecFormat len FF32
               addr     = spRel platform 0
               code dst = exp `snocOL`
                          (MOVU f (OpReg reg) (OpAddr addr)) `snocOL`
@@ -1307,7 +1307,7 @@ getRegister' platform is32Bit (CmmMachOp mop [x]) = do -- unary MachOps
         vector_int_broadcast len W64 expr
           = do
           (reg, exp) <- getSomeReg expr
-          let fmt = VecFormat len FmtInt64
+          let fmt = VecFormat len II64
           return $ Any fmt (\dst -> exp `snocOL`
                                     (MOV II64 (OpReg reg) (OpReg dst)) `snocOL`
                                     (PUNPCKLQDQ fmt (OpReg dst) dst) `snocOL`
@@ -1652,8 +1652,8 @@ getRegister' platform is32Bit (CmmMachOp mop [x, y]) = do -- dyadic MachOps
       (reg1, exp1) <- getSomeReg expr1
       (reg2, exp2) <- getSomeReg expr2
       let format   = case w of
-                       W32 -> VecFormat l FmtFloat
-                       W64 -> VecFormat l FmtDouble
+                       W32 -> VecFormat l FF32
+                       W64 -> VecFormat l FF64
                        _ -> pprPanic "Operation not supported for width " (ppr w)
           code dst = case op of
             VA_Add -> arithInstr VADD
@@ -1676,8 +1676,8 @@ getRegister' platform is32Bit (CmmMachOp mop [x, y]) = do -- dyadic MachOps
       (reg1, exp1) <- getSomeReg expr1
       (reg2, exp2) <- getSomeReg expr2
       let format   = case w of
-                       W32 -> VecFormat l FmtFloat
-                       W64 -> VecFormat l FmtDouble
+                       W32 -> VecFormat l FF32
+                       W64 -> VecFormat l FF64
                        _ -> pprPanic "Operation not supported for width " (ppr w)
           code dst = case op of
             VA_Add -> arithInstr ADD
@@ -1700,7 +1700,7 @@ getRegister' platform is32Bit (CmmMachOp mop [x, y]) = do -- dyadic MachOps
     vector_float_unpack l W32 expr (CmmLit lit)
       = do
       (r, exp) <- getSomeReg expr
-      let format   = VecFormat l FmtFloat
+      let format   = VecFormat l FF32
           imm      = litToImm lit
           code dst
             = case lit of
@@ -1711,7 +1711,7 @@ getRegister' platform is32Bit (CmmMachOp mop [x, y]) = do -- dyadic MachOps
     vector_float_unpack l W64 expr (CmmLit lit)
       = do
       (r, exp) <- getSomeReg expr
-      let format   = VecFormat l FmtDouble
+      let format   = VecFormat l FF64
           code dst
             = case lit of
                 CmmInt 0 _ -> exp `snocOL`
@@ -1732,7 +1732,7 @@ getRegister' platform is32Bit (CmmMachOp mop [x, y]) = do -- dyadic MachOps
     vector_float_unpack_sse l W32 expr (CmmLit lit)
       = do
       (r,exp) <- getSomeReg expr
-      let format   = VecFormat l FmtFloat
+      let format   = VecFormat l FF32
           imm      = litToImm lit
           code dst
             = case lit of
@@ -1752,7 +1752,7 @@ getRegister' platform is32Bit (CmmMachOp mop [x, y]) = do -- dyadic MachOps
     vector_int_unpack_sse l at 2 W64 expr (CmmLit lit)
       = do
       (r, exp) <- getSomeReg expr
-      let fmt = VecFormat l FmtInt64
+      let fmt = VecFormat l II64
       tmp <- getNewRegNat fmt
       let code dst
             = case lit of
@@ -1770,7 +1770,7 @@ getRegister' platform is32Bit (CmmMachOp mop [x, y]) = do -- dyadic MachOps
     vector_shuffle_float l w v1 v2 is = do
       (r1, exp1) <- getSomeReg v1
       (r2, exp2) <- getSomeReg v2
-      let fmt = VecFormat l (if w == W32 then FmtFloat else FmtDouble)
+      let fmt = VecFormat l (if w == W32 then FF32 else FF64)
           code dst
             = exp1 `appOL` (exp2 `appOL` shuffleInstructions fmt r1 r2 is dst)
       return (Any fmt code)
@@ -1778,7 +1778,7 @@ getRegister' platform is32Bit (CmmMachOp mop [x, y]) = do -- dyadic MachOps
     shuffleInstructions :: Format -> Reg -> Reg -> [Int] -> Reg -> OrdList Instr
     shuffleInstructions fmt v1 v2 is dst =
       case fmt of
-        VecFormat 2 FmtDouble ->
+        VecFormat 2 FF64 ->
           case is of
             [i1, i2] -> case (i1, i2) of
               (0,0) -> unitOL (VSHUFPD fmt (ImmInt 0b00) (OpReg v1) v1 dst)
@@ -1799,7 +1799,7 @@ getRegister' platform is32Bit (CmmMachOp mop [x, y]) = do -- dyadic MachOps
               (3,1) -> unitOL (VSHUFPD fmt (ImmInt 0b11) (OpReg v1) v2 dst)
               _ -> pprPanic "vector shuffle: indices out of bounds 0 <= i <= 3" (ppr is)
             _ -> pprPanic "vector shuffle: wrong number of indices (expected 2)" (ppr is)
-        VecFormat 4 FmtFloat
+        VecFormat 4 FF32
           -- indices 0 <= i <= 7
           | all ( (>= 0) <&&> (<= 7) ) is ->
           case is of
@@ -1885,7 +1885,7 @@ getRegister' platform _is32Bit (CmmMachOp mop [x, y, z]) = do -- ternary MachOps
       = do
       fn          <- getAnyReg vecExpr
       (r, exp)    <- getSomeReg valExpr
-      let fmt      = VecFormat len FmtFloat
+      let fmt      = VecFormat len FF32
           imm      = litToImm (CmmInt (offset `shiftL` 4) W32)
           code dst = exp `appOL`
                      (fn dst) `snocOL`
@@ -1896,7 +1896,7 @@ getRegister' platform _is32Bit (CmmMachOp mop [x, y, z]) = do -- ternary MachOps
       = do
         (valReg, valExp) <- getSomeReg valExpr
         (vecReg, vecExp) <- getSomeReg vecExpr
-        let fmt = VecFormat len FmtDouble
+        let fmt = VecFormat len FF64
             code dst
               = case offset of
                   CmmInt 0 _ -> valExp `appOL`
@@ -1934,7 +1934,7 @@ getRegister' platform _is32Bit (CmmMachOp mop [x, y, z]) = do -- ternary MachOps
         pprTraceM "vecExpr:" (pdoc platform vecExpr)
         (valReg, valExp) <- getSomeReg valExpr
         (vecReg, vecExp) <- getSomeReg vecExpr
-        let fmt = VecFormat len FmtInt64
+        let fmt = VecFormat len II64
         tmp <- getNewRegNat fmt
         pprTraceM "tmp:" (ppr tmp)
         let code dst
@@ -2382,7 +2382,7 @@ addAlignmentCheck align reg =
   where
     check :: Format -> Reg -> InstrBlock
     check fmt reg =
-        assert (not $ isFloatFormat fmt) $
+        assert (isIntFormat fmt) $
         toOL [ TEST fmt (OpImm $ ImmInt $ align-1) (OpReg reg)
              , JXX_GBL NE $ ImmCLbl mkBadAlignmentLabel
              ]


=====================================
compiler/GHC/CmmToAsm/X86/Instr.hs
=====================================
@@ -629,11 +629,10 @@ regUsageOfInstr platform instr
               use_index EAIndexNone   tl = tl
               use_index (EAIndex i _) tl = i : tl
 
-    mkRUR fmt src = src' `seq` RU (map (,fmt) src') []
+    mkRUR fmt src = src' `seq` RU (map (\ r -> RegFormat r fmt) src') []
         where src' = filter (interesting platform) src
 
-
-    mkRU fmt src dst = src' `seq` dst' `seq` RU (map (,fmt) src') (map (,fmt) dst')
+    mkRU fmt src dst = src' `seq` dst' `seq` RU (map (\ r -> RegFormat r fmt) src') (map (\ r -> RegFormat r fmt) dst')
         where src' = filter (interesting platform) src
               dst' = filter (interesting platform) dst
 


=====================================
compiler/GHC/CmmToAsm/X86/Ppr.hs
=====================================
@@ -431,13 +431,13 @@ pprFormat x = case x of
   II64  -> text "q"
   FF32  -> text "ss"      -- "scalar single-precision float" (SSE2)
   FF64  -> text "sd"      -- "scalar double-precision float" (SSE2)
-  VecFormat _ FmtFloat  -> text "ps"
-  VecFormat _ FmtDouble -> text "pd"
+  VecFormat _ FF32  -> text "ps"
+  VecFormat _ FF64  -> text "pd"
   -- TODO: this is shady because it only works for certain instructions
-  VecFormat _ FmtInt8   -> text "b"
-  VecFormat _ FmtInt16  -> text "w"
-  VecFormat _ FmtInt32  -> text "l"
-  VecFormat _ FmtInt64  -> text "q"
+  VecFormat _ II8   -> text "b"
+  VecFormat _ II16  -> text "w"
+  VecFormat _ II32  -> text "l"
+  VecFormat _ II64  -> text "q"
 
 pprFormat_x87 :: IsLine doc => Format -> doc
 pprFormat_x87 x = case x of
@@ -781,9 +781,9 @@ pprInstr platform i = case i of
    BT format imm src
       -> pprFormatImmOp (text "bt") format imm src
 
-   CMP format src dst
-     | isFloatFormat format -> pprFormatOpOp (text "ucomi") format src dst -- SSE2
-     | otherwise            -> pprFormatOpOp (text "cmp")   format src dst
+   CMP fmt@(Format _ s) src dst
+     | isFloatScalarFormat s -> pprFormatOpOp (text "ucomi") fmt src dst -- SSE2
+     | otherwise             -> pprFormatOpOp (text "cmp")  fmt src dst
 
    TEST format src dst
       -> pprFormatOpOp (text "test") format' src dst
@@ -1051,14 +1051,13 @@ pprInstr platform i = case i of
       char '\t' <> name <> pprBroadcastFormat format <> space
 
    pprBroadcastFormat :: Format -> Line doc
-   pprBroadcastFormat (VecFormat _ f)
-     = case f of
-         FmtFloat  -> text "ss"
-         FmtDouble -> text "sd"
-         FmtInt8   -> text "b"
-         FmtInt16  -> text "w"
-         FmtInt32  -> text "d"
-         FmtInt64  -> text "q"
+   pprBroadcastFormat (VecFormat _ f) = case f of
+         FF32  -> text "ss"
+         FF64  -> text "sd"
+         II8   -> text "b"
+         II16  -> text "w"
+         II32  -> text "d"
+         II64  -> text "q"
    pprBroadcastFormat _ = panic "Scalar Format invading vector operation"
 
    pprFormatImmOp :: Line doc -> Format -> Imm -> Operand -> doc


=====================================
compiler/GHC/Types/Unique/FM.hs
=====================================
@@ -71,8 +71,7 @@ module GHC.Types.Unique.FM (
         nonDetStrictFoldUFM_Directly,
         anyUFM, allUFM, seqEltsUFM,
         mapUFM, mapUFM_Directly, strictMapUFM,
-        mapKeysUFM,
-        mapMaybeUFM, mapMaybeWithKeyUFM,
+        mapMaybeUFM, mapMaybeUFM_sameUnique, mapMaybeWithKeyUFM,
         elemUFM, elemUFM_Directly,
         filterUFM, filterUFM_Directly, partitionUFM,
         sizeUFM,
@@ -390,7 +389,12 @@ mapUFM :: (elt1 -> elt2) -> UniqFM key elt1 -> UniqFM key elt2
 mapUFM f (UFM m) = UFM (M.map f m)
 
 mapMaybeUFM :: (elt1 -> Maybe elt2) -> UniqFM key elt1 -> UniqFM key elt2
-mapMaybeUFM f (UFM m) = UFM (M.mapMaybe f m)
+mapMaybeUFM = mapMaybeUFM_sameUnique
+
+-- | Like 'Data.Map.mapMaybe', but you must ensure the passed-in function does
+-- not modify the unique.
+mapMaybeUFM_sameUnique :: (elt1 -> Maybe elt2) -> UniqFM key1 elt1 -> UniqFM key2 elt2
+mapMaybeUFM_sameUnique f (UFM m) = UFM (M.mapMaybe f m)
 
 mapMaybeWithKeyUFM :: (Unique -> elt1 -> Maybe elt2) -> UniqFM key elt1 -> UniqFM key elt2
 mapMaybeWithKeyUFM f (UFM m) = UFM (M.mapMaybeWithKey (f . mkUniqueGrimily) m)
@@ -398,10 +402,6 @@ mapMaybeWithKeyUFM f (UFM m) = UFM (M.mapMaybeWithKey (f . mkUniqueGrimily) m)
 mapUFM_Directly :: (Unique -> elt1 -> elt2) -> UniqFM key elt1 -> UniqFM key elt2
 mapUFM_Directly f (UFM m) = UFM (M.mapWithKey (f . mkUniqueGrimily) m)
 
--- | Map over the keys in a 'UniqFM'.
-mapKeysUFM :: Uniquable key' => (key -> key') -> UniqFM key (key, b) -> UniqFM key' (key', b)
-mapKeysUFM f m = listToUFM $ map ( \ (r, fmt) -> let r' = f r in (r', (r', fmt)) ) $ nonDetEltsUFM m
-
 strictMapUFM :: (a -> b) -> UniqFM k a -> UniqFM k b
 strictMapUFM f (UFM a) = UFM $ MS.map f a
 


=====================================
compiler/GHC/Types/Unique/Set.hs
=====================================
@@ -44,6 +44,7 @@ module GHC.Types.Unique.Set (
         nonDetEltsUniqSet,
         nonDetKeysUniqSet,
         nonDetStrictFoldUniqSet,
+        mapMaybeUniqSet_sameUnique,
 
         -- UniqueSet
         UniqueSet(..),
@@ -205,6 +206,11 @@ nonDetStrictFoldUniqSet c n (UniqSet s) = nonDetStrictFoldUFM c n s
 mapUniqSet :: Uniquable b => (a -> b) -> UniqSet a -> UniqSet b
 mapUniqSet f = mkUniqSet . map f . nonDetEltsUniqSet
 
+-- | Like 'Data.Set.mapMaybe', but you must ensure the passed in function
+-- does not change the 'Unique'.
+mapMaybeUniqSet_sameUnique :: (a -> Maybe b) -> UniqSet a -> UniqSet b
+mapMaybeUniqSet_sameUnique f (UniqSet a) = UniqSet $ mapMaybeUFM_sameUnique f a
+
 -- Two 'UniqSet's are considered equal if they contain the same
 -- uniques.
 instance Eq (UniqSet a) where



View it on GitLab: https://gitlab.haskell.org/ghc/ghc/-/compare/baa104c6ba96284ecc96eb82eb0f8581abb65b84...62168cfcccf795af856a9a096af78c599502477c

-- 
View it on GitLab: https://gitlab.haskell.org/ghc/ghc/-/compare/baa104c6ba96284ecc96eb82eb0f8581abb65b84...62168cfcccf795af856a9a096af78c599502477c
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/20240615/3cab07e7/attachment-0001.html>


More information about the ghc-commits mailing list