Add Addr# atomic primops (#17751)

This reuses the codegen used for ByteArray#'s atomic primops.

3 files changed, 130 insertions, 21 deletions

diff --git a/compiler/GHC/Builtin/primops.txt.pp b/compiler/GHC/Builtin/primops.txt.pp
index c292b9ecdc..acee822b0c 100644
--- a/compiler/GHC/Builtin/primops.txt.pp
+++ b/compiler/GHC/Builtin/primops.txt.pp
@@ -1669,7 +1669,7 @@ primop FetchAddByteArrayOp_Int "fetchAddIntArray#" GenPrimOp
 primop FetchSubByteArrayOp_Int "fetchSubIntArray#" GenPrimOp
    MutableByteArray# s -> Int# -> Int# -> State# s -> (# State# s, Int# #)
    {Given an array, and offset in machine words, and a value to subtract,
-    atomically subtract the value to the element. Returns the value of
+    atomically subtract the value from the element. Returns the value of
     the element before the operation. Implies a full memory barrier.}
    with has_side_effects = True
         can_fail = True
@@ -1677,7 +1677,7 @@ primop FetchSubByteArrayOp_Int "fetchSubIntArray#" GenPrimOp
 primop FetchAndByteArrayOp_Int "fetchAndIntArray#" GenPrimOp
    MutableByteArray# s -> Int# -> Int# -> State# s -> (# State# s, Int# #)
    {Given an array, and offset in machine words, and a value to AND,
-    atomically AND the value to the element. Returns the value of the
+    atomically AND the value into the element. Returns the value of the
     element before the operation. Implies a full memory barrier.}
    with has_side_effects = True
         can_fail = True
@@ -1685,7 +1685,7 @@ primop FetchAndByteArrayOp_Int "fetchAndIntArray#" GenPrimOp
 primop FetchNandByteArrayOp_Int "fetchNandIntArray#" GenPrimOp
    MutableByteArray# s -> Int# -> Int# -> State# s -> (# State# s, Int# #)
    {Given an array, and offset in machine words, and a value to NAND,
-    atomically NAND the value to the element. Returns the value of the
+    atomically NAND the value into the element. Returns the value of the
     element before the operation. Implies a full memory barrier.}
    with has_side_effects = True
         can_fail = True
@@ -1693,7 +1693,7 @@ primop FetchNandByteArrayOp_Int "fetchNandIntArray#" GenPrimOp
 primop FetchOrByteArrayOp_Int "fetchOrIntArray#" GenPrimOp
    MutableByteArray# s -> Int# -> Int# -> State# s -> (# State# s, Int# #)
    {Given an array, and offset in machine words, and a value to OR,
-    atomically OR the value to the element. Returns the value of the
+    atomically OR the value into the element. Returns the value of the
     element before the operation. Implies a full memory barrier.}
    with has_side_effects = True
         can_fail = True
@@ -1701,7 +1701,7 @@ primop FetchOrByteArrayOp_Int "fetchOrIntArray#" GenPrimOp
 primop FetchXorByteArrayOp_Int "fetchXorIntArray#" GenPrimOp
    MutableByteArray# s -> Int# -> Int# -> State# s -> (# State# s, Int# #)
    {Given an array, and offset in machine words, and a value to XOR,
-    atomically XOR the value to the element. Returns the value of the
+    atomically XOR the value into the element. Returns the value of the
     element before the operation. Implies a full memory barrier.}
    with has_side_effects = True
         can_fail = True
@@ -2121,6 +2121,67 @@ primop  CasAddrOp_Word "atomicCasWordAddr#" GenPrimOp
    with has_side_effects = True
         can_fail         = True
 
+primop FetchAddAddrOp_Word "fetchAddWordAddr#" GenPrimOp
+   Addr# -> Word# -> State# s -> (# State# s, Word# #)
+   {Given an address, and a value to add,
+    atomically add the value to the element. Returns the value of the
+    element before the operation. Implies a full memory barrier.}
+   with has_side_effects = True
+        can_fail = True
+
+primop FetchSubAddrOp_Word "fetchSubWordAddr#" GenPrimOp
+   Addr# -> Word# -> State# s -> (# State# s, Word# #)
+   {Given an address, and a value to subtract,
+    atomically subtract the value from the element. Returns the value of
+    the element before the operation. Implies a full memory barrier.}
+   with has_side_effects = True
+        can_fail = True
+
+primop FetchAndAddrOp_Word "fetchAndWordAddr#" GenPrimOp
+   Addr# -> Word# -> State# s -> (# State# s, Word# #)
+   {Given an address, and a value to AND,
+    atomically AND the value into the element. Returns the value of the
+    element before the operation. Implies a full memory barrier.}
+   with has_side_effects = True
+        can_fail = True
+
+primop FetchNandAddrOp_Word "fetchNandWordAddr#" GenPrimOp
+   Addr# -> Word# -> State# s -> (# State# s, Word# #)
+   {Given an address, and a value to NAND,
+    atomically NAND the value into the element. Returns the value of the
+    element before the operation. Implies a full memory barrier.}
+   with has_side_effects = True
+        can_fail = True
+
+primop FetchOrAddrOp_Word "fetchOrWordAddr#" GenPrimOp
+   Addr# -> Word# -> State# s -> (# State# s, Word# #)
+   {Given an address, and a value to OR,
+    atomically OR the value into the element. Returns the value of the
+    element before the operation. Implies a full memory barrier.}
+   with has_side_effects = True
+        can_fail = True
+
+primop FetchXorAddrOp_Word "fetchXorWordAddr#" GenPrimOp
+   Addr# -> Word# -> State# s -> (# State# s, Word# #)
+   {Given an address, and a value to XOR,
+    atomically XOR the value into the element. Returns the value of the
+    element before the operation. Implies a full memory barrier.}
+   with has_side_effects = True
+        can_fail = True
+
+primop  AtomicReadAddrOp_Word "atomicReadWordAddr#" GenPrimOp
+   Addr# -> State# s -> (# State# s, Word# #)
+   {Given an address, read a machine word.  Implies a full memory barrier.}
+   with has_side_effects = True
+        can_fail = True
+
+primop  AtomicWriteAddrOp_Word "atomicWriteWordAddr#" GenPrimOp
+   Addr# -> Word# -> State# s -> State# s
+   {Given an address, write a machine word. Implies a full memory barrier.}
+   with has_side_effects = True
+        can_fail = True
+
+
 ------------------------------------------------------------------------
 section "Mutable variables"
         {Operations on MutVar\#s.}
diff --git a/compiler/GHC/CmmToAsm/X86/CodeGen.hs b/compiler/GHC/CmmToAsm/X86/CodeGen.hs
index 4480848b90..3883d44717 100644
--- a/compiler/GHC/CmmToAsm/X86/CodeGen.hs
+++ b/compiler/GHC/CmmToAsm/X86/CodeGen.hs
@@ -2121,12 +2121,12 @@ genCCall is32Bit (PrimTarget (MO_AtomicRMW width amop))
         -- final move should go away, because it's the last use of arg
         -- and the first use of dst_r.
         AMO_Add  -> return $ (toOL [ LOCK (XADD format (OpReg arg) (OpAddr amode))
-                                  , MOV format (OpReg arg) (OpReg dst_r)
-                                  ], bid)
+                                   , MOV format (OpReg arg) (OpReg dst_r)
+                                   ], bid)
         AMO_Sub  -> return $ (toOL [ NEGI format (OpReg arg)
-                                  , LOCK (XADD format (OpReg arg) (OpAddr amode))
-                                  , MOV format (OpReg arg) (OpReg dst_r)
-                                  ], bid)
+                                   , LOCK (XADD format (OpReg arg) (OpAddr amode))
+                                   , MOV format (OpReg arg) (OpReg dst_r)
+                                   ], bid)
         -- In these cases we need a new block id, and have to return it so
         -- that later instruction selection can reference it.
         AMO_And  -> cmpxchg_code (\ src dst -> unitOL $ AND format src dst)
diff --git a/compiler/GHC/StgToCmm/Prim.hs b/compiler/GHC/StgToCmm/Prim.hs
index 099a3850dc..625e76f085 100644
--- a/compiler/GHC/StgToCmm/Prim.hs
+++ b/compiler/GHC/StgToCmm/Prim.hs
@@ -850,6 +850,25 @@ emitPrimOp dflags primop = case primop of
     emitPrimCall [res] (MO_Xchg (wordWidth platform)) [src, value]
   InterlockedExchange_Word -> \[src, value] -> opIntoRegs $ \[res] ->
     emitPrimCall [res] (MO_Xchg (wordWidth platform)) [src, value]
+
+  FetchAddAddrOp_Word -> \[addr, n] -> opIntoRegs $ \[res] ->
+    doAtomicAddrRMW res AMO_Add addr (bWord platform) n
+  FetchSubAddrOp_Word -> \[addr, n] -> opIntoRegs $ \[res] ->
+    doAtomicAddrRMW res AMO_Sub addr (bWord platform) n
+  FetchAndAddrOp_Word -> \[addr, n] -> opIntoRegs $ \[res] ->
+    doAtomicAddrRMW res AMO_And addr (bWord platform) n
+  FetchNandAddrOp_Word -> \[addr, n] -> opIntoRegs $ \[res] ->
+    doAtomicAddrRMW res AMO_Nand addr (bWord platform) n
+  FetchOrAddrOp_Word -> \[addr, n] -> opIntoRegs $ \[res] ->
+    doAtomicAddrRMW res AMO_Or addr (bWord platform) n
+  FetchXorAddrOp_Word -> \[addr, n] -> opIntoRegs $ \[res] ->
+    doAtomicAddrRMW res AMO_Xor addr (bWord platform) n
+
+  AtomicReadAddrOp_Word -> \[addr] -> opIntoRegs $ \[res] ->
+    doAtomicReadAddr res addr (bWord platform)
+  AtomicWriteAddrOp_Word -> \[addr, val] -> opIntoRegs $ \[] ->
+    doAtomicWriteAddr addr (bWord platform) val
+
   CasAddrOp_Addr -> \[dst, expected, new] -> opIntoRegs $ \[res] ->
     emitPrimCall [res] (MO_Cmpxchg (wordWidth platform)) [dst, expected, new]
   CasAddrOp_Word -> \[dst, expected, new] -> opIntoRegs $ \[res] ->
@@ -1040,17 +1059,17 @@ emitPrimOp dflags primop = case primop of
 
 -- Atomic read-modify-write
   FetchAddByteArrayOp_Int -> \[mba, ix, n] -> opIntoRegs $ \[res] ->
-    doAtomicRMW res AMO_Add mba ix (bWord platform) n
+    doAtomicByteArrayRMW res AMO_Add mba ix (bWord platform) n
   FetchSubByteArrayOp_Int -> \[mba, ix, n] -> opIntoRegs $ \[res] ->
-    doAtomicRMW res AMO_Sub mba ix (bWord platform) n
+    doAtomicByteArrayRMW res AMO_Sub mba ix (bWord platform) n
   FetchAndByteArrayOp_Int -> \[mba, ix, n] -> opIntoRegs $ \[res] ->
-    doAtomicRMW res AMO_And mba ix (bWord platform) n
+    doAtomicByteArrayRMW res AMO_And mba ix (bWord platform) n
   FetchNandByteArrayOp_Int -> \[mba, ix, n] -> opIntoRegs $ \[res] ->
-    doAtomicRMW res AMO_Nand mba ix (bWord platform) n
+    doAtomicByteArrayRMW res AMO_Nand mba ix (bWord platform) n
   FetchOrByteArrayOp_Int -> \[mba, ix, n] -> opIntoRegs $ \[res] ->
-    doAtomicRMW res AMO_Or mba ix (bWord platform) n
+    doAtomicByteArrayRMW res AMO_Or mba ix (bWord platform) n
   FetchXorByteArrayOp_Int -> \[mba, ix, n] -> opIntoRegs $ \[res] ->
-    doAtomicRMW res AMO_Xor mba ix (bWord platform) n
+    doAtomicByteArrayRMW res AMO_Xor mba ix (bWord platform) n
   AtomicReadByteArrayOp_Int -> \[mba, ix] -> opIntoRegs $ \[res] ->
     doAtomicReadByteArray res mba ix (bWord platform)
   AtomicWriteByteArrayOp_Int -> \[mba, ix, val] -> opIntoRegs $ \[] ->
@@ -2855,22 +2874,33 @@ doWriteSmallPtrArrayOp addr idx val = do
 -- | Emit an atomic modification to a byte array element. The result
 -- reg contains that previous value of the element. Implies a full
 -- memory barrier.
-doAtomicRMW :: LocalReg      -- ^ Result reg
+doAtomicByteArrayRMW
+            :: LocalReg      -- ^ Result reg
             -> AtomicMachOp  -- ^ Atomic op (e.g. add)
             -> CmmExpr       -- ^ MutableByteArray#
             -> CmmExpr       -- ^ Index
             -> CmmType       -- ^ Type of element by which we are indexing
             -> CmmExpr       -- ^ Op argument (e.g. amount to add)
             -> FCode ()
-doAtomicRMW res amop mba idx idx_ty n = do
+doAtomicByteArrayRMW res amop mba idx idx_ty n = do
     profile <- getProfile
     platform <- getPlatform
     let width = typeWidth idx_ty
         addr  = cmmIndexOffExpr platform (arrWordsHdrSize profile)
                 width mba idx
+    doAtomicAddrRMW res amop addr idx_ty n
+
+doAtomicAddrRMW
+            :: LocalReg      -- ^ Result reg
+            -> AtomicMachOp  -- ^ Atomic op (e.g. add)
+            -> CmmExpr       -- ^ Addr#
+            -> CmmType       -- ^ Pointed value type
+            -> CmmExpr       -- ^ Op argument (e.g. amount to add)
+            -> FCode ()
+doAtomicAddrRMW res amop addr ty n = do
     emitPrimCall
         [ res ]
-        (MO_AtomicRMW width amop)
+        (MO_AtomicRMW (typeWidth ty) amop)
         [ addr, n ]
 
 -- | Emit an atomic read to a byte array that acts as a memory barrier.
@@ -2886,9 +2916,18 @@ doAtomicReadByteArray res mba idx idx_ty = do
     let width = typeWidth idx_ty
         addr  = cmmIndexOffExpr platform (arrWordsHdrSize profile)
                 width mba idx
+    doAtomicReadAddr res addr idx_ty
+
+-- | Emit an atomic read to an address that acts as a memory barrier.
+doAtomicReadAddr
+    :: LocalReg  -- ^ Result reg
+    -> CmmExpr   -- ^ Addr#
+    -> CmmType   -- ^ Type of element by which we are indexing
+    -> FCode ()
+doAtomicReadAddr res addr ty = do
     emitPrimCall
         [ res ]
-        (MO_AtomicRead width)
+        (MO_AtomicRead (typeWidth ty))
         [ addr ]
 
 -- | Emit an atomic write to a byte array that acts as a memory barrier.
@@ -2904,9 +2943,18 @@ doAtomicWriteByteArray mba idx idx_ty val = do
     let width = typeWidth idx_ty
         addr  = cmmIndexOffExpr platform (arrWordsHdrSize profile)
                 width mba idx
+    doAtomicWriteAddr addr idx_ty val
+
+-- | Emit an atomic write to an address that acts as a memory barrier.
+doAtomicWriteAddr
+    :: CmmExpr   -- ^ Addr#
+    -> CmmType   -- ^ Type of element by which we are indexing
+    -> CmmExpr   -- ^ Value to write
+    -> FCode ()
+doAtomicWriteAddr addr ty val = do
     emitPrimCall
         [ {- no results -} ]
-        (MO_AtomicWrite width)
+        (MO_AtomicWrite (typeWidth ty))
         [ addr, val ]
 
 doCasByteArray