Replace atomicModifyMutVar#

Reviewers: simonmar, hvr, bgamari, erikd, fryguybob, rrnewton

Reviewed By: simonmar

Subscribers: fryguybob, rwbarton, thomie, carter

GHC Trac Issues: #15364

Differential Revision: https://phabricator.haskell.org/D4884

12 files changed, 255 insertions, 60 deletions

diff --git a/compiler/prelude/primops.txt.pp b/compiler/prelude/primops.txt.pp
index 468299f5d2..b8c6811668 100644
--- a/compiler/prelude/primops.txt.pp
+++ b/compiler/prelude/primops.txt.pp
@@ -2224,25 +2224,37 @@ primop  WriteMutVarOp "writeMutVar#"  GenPrimOp
 primop  SameMutVarOp "sameMutVar#" GenPrimOp
    MutVar# s a -> MutVar# s a -> Int#
 
--- Note [Why not an unboxed tuple in atomicModifyMutVar#?]
+-- Note [Why not an unboxed tuple in atomicModifyMutVar2#?]
 -- ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
 --
--- Looking at the type of atomicModifyMutVar#, one might wonder why
+-- Looking at the type of atomicModifyMutVar2#, one might wonder why
 -- it doesn't return an unboxed tuple. e.g.,
 --
---   MutVar# s a -> (a -> (# a, b #)) -> State# s -> (# State# s, b #)
+--   MutVar# s a -> (a -> (# a, b #)) -> State# s -> (# State# s, a, (# a, b #) #)
 --
--- The reason is that atomicModifyMutVar# relies on laziness for its atomicity.
--- Given a MutVar# containing x, atomicModifyMutVar# merely replaces the
+-- The reason is that atomicModifyMutVar2# relies on laziness for its atomicity.
+-- Given a MutVar# containing x, atomicModifyMutVar2# merely replaces
 -- its contents with a thunk of the form (fst (f x)). This can be done using an
 -- atomic compare-and-swap as it is merely replacing a pointer.
 
-primop  AtomicModifyMutVarOp "atomicModifyMutVar#" GenPrimOp
-   MutVar# s a -> (a -> b) -> State# s -> (# State# s, c #)
-   { Modify the contents of a {\tt MutVar\#}. Note that this isn't strictly
-     speaking the correct type for this function, it should really be
-     {\tt MutVar# s a -> (a -> (a,b)) -> State# s -> (# State# s, b #)}, however
-     we don't know about pairs here. }
+primop  AtomicModifyMutVar2Op "atomicModifyMutVar2#" GenPrimOp
+   MutVar# s a -> (a -> c) -> State# s -> (# State# s, a, c #)
+   { Modify the contents of a {\tt MutVar\#}, returning the previous
+     contents and the result of applying the given function to the
+     previous contents. Note that this isn't strictly
+     speaking the correct type for this function; it should really be
+     {\tt MutVar# s a -> (a -> (a,b)) -> State# s -> (# State# s, a, (a, b) #)},
+     but we don't know about pairs here. }
+   with
+   out_of_line = True
+   has_side_effects = True
+   can_fail         = True
+
+primop  AtomicModifyMutVar_Op "atomicModifyMutVar_#" GenPrimOp
+   MutVar# s a -> (a -> a) -> State# s -> (# State# s, a, a #)
+   { Modify the contents of a {\tt MutVar\#}, returning the previous
+     contents and the result of applying the given function to the
+     previous contents. }
    with
    out_of_line = True
    has_side_effects = True
diff --git a/includes/stg/MiscClosures.h b/includes/stg/MiscClosures.h
index 98363b342b..6ae3df1390 100644
--- a/includes/stg/MiscClosures.h
+++ b/includes/stg/MiscClosures.h
@@ -390,7 +390,8 @@ RTS_FUN_DECL(stg_copySmallMutableArrayzh);
 RTS_FUN_DECL(stg_casSmallArrayzh);
 
 RTS_FUN_DECL(stg_newMutVarzh);
-RTS_FUN_DECL(stg_atomicModifyMutVarzh);
+RTS_FUN_DECL(stg_atomicModifyMutVar2zh);
+RTS_FUN_DECL(stg_atomicModifyMutVarzuzh);
 RTS_FUN_DECL(stg_casMutVarzh);
 
 RTS_FUN_DECL(stg_isEmptyMVarzh);
diff --git a/libraries/base/Data/IORef.hs b/libraries/base/Data/IORef.hs
index c6275f5433..6f07a097f7 100644
--- a/libraries/base/Data/IORef.hs
+++ b/libraries/base/Data/IORef.hs
@@ -1,5 +1,6 @@
 {-# LANGUAGE Trustworthy #-}
 {-# LANGUAGE NoImplicitPrelude, MagicHash, UnboxedTuples #-}
+{-# LANGUAGE BangPatterns #-}
 
 -----------------------------------------------------------------------------
 -- |
@@ -36,8 +37,7 @@ module Data.IORef
 
 import GHC.Base
 import GHC.STRef
-import GHC.IORef hiding (atomicModifyIORef)
-import qualified GHC.IORef
+import GHC.IORef
 import GHC.Weak
 
 -- |Make a 'Weak' pointer to an 'IORef', using the second argument as a finalizer
@@ -91,18 +91,9 @@ modifyIORef' ref f = do
 -- Use 'atomicModifyIORef'' or 'atomicWriteIORef' to avoid this problem.
 --
 atomicModifyIORef :: IORef a -> (a -> (a,b)) -> IO b
-atomicModifyIORef = GHC.IORef.atomicModifyIORef
-
--- | Strict version of 'atomicModifyIORef'.  This forces both the value stored
--- in the 'IORef' as well as the value returned.
---
--- @since 4.6.0.0
-atomicModifyIORef' :: IORef a -> (a -> (a,b)) -> IO b
-atomicModifyIORef' ref f = do
-    b <- atomicModifyIORef ref $ \a ->
-            case f a of
-                v@(a',_) -> a' `seq` v
-    b `seq` return b
+atomicModifyIORef ref f = do
+  (_old, ~(_new, res)) <- atomicModifyIORef2 ref f
+  pure res
 
 -- | Variant of 'writeIORef' with the \"barrier to reordering\" property that
 -- 'atomicModifyIORef' has.
@@ -110,8 +101,8 @@ atomicModifyIORef' ref f = do
 -- @since 4.6.0.0
 atomicWriteIORef :: IORef a -> a -> IO ()
 atomicWriteIORef ref a = do
-    x <- atomicModifyIORef ref (\_ -> (a, ()))
-    x `seq` return ()
+  _ <- atomicSwapIORef ref a
+  pure ()
 
 {- $memmodel
 
diff --git a/libraries/base/GHC/Conc/Windows.hs b/libraries/base/GHC/Conc/Windows.hs
index 6c4136428a..ed5e0452a0 100644
--- a/libraries/base/GHC/Conc/Windows.hs
+++ b/libraries/base/GHC/Conc/Windows.hs
@@ -131,7 +131,7 @@ waitForDelayEvent :: Int -> IO ()
 waitForDelayEvent usecs = do
   m <- newEmptyMVar
   target <- calculateTarget usecs
-  atomicModifyIORef pendingDelays (\xs -> (Delay target m : xs, ()))
+  _ <- atomicModifyIORef'_ pendingDelays (\xs -> Delay target m : xs)
   prodServiceThread
   takeMVar m
 
@@ -140,7 +140,7 @@ waitForDelayEventSTM :: Int -> IO (TVar Bool)
 waitForDelayEventSTM usecs = do
    t <- atomically $ newTVar False
    target <- calculateTarget usecs
-   atomicModifyIORef pendingDelays (\xs -> (DelaySTM target t : xs, ()))
+   _ <- atomicModifyIORef'_ pendingDelays (\xs -> DelaySTM target t : xs)
    prodServiceThread
    return t
 
@@ -219,10 +219,10 @@ foreign import ccall unsafe "getOrSetGHCConcWindowsProddingStore"
 
 prodServiceThread :: IO ()
 prodServiceThread = do
-  -- NB. use atomicModifyIORef here, otherwise there are race
+  -- NB. use atomicSwapIORef here, otherwise there are race
   -- conditions in which prodding is left at True but the server is
   -- blocked in select().
-  was_set <- atomicModifyIORef prodding $ \b -> (True,b)
+  was_set <- atomicSwapIORef prodding True
   when (not was_set) wakeupIOManager
 
 -- ----------------------------------------------------------------------------
@@ -239,7 +239,7 @@ service_loop :: HANDLE          -- read end of pipe
 
 service_loop wakeup old_delays = do
   -- pick up new delay requests
-  new_delays <- atomicModifyIORef pendingDelays (\a -> ([],a))
+  new_delays <- atomicSwapIORef pendingDelays []
   let  delays = foldr insertDelay old_delays new_delays
 
   now <- getMonotonicUSec
@@ -262,8 +262,7 @@ service_loop wakeup old_delays = do
 
 service_cont :: HANDLE -> [DelayReq] -> IO ()
 service_cont wakeup delays = do
-  r <- atomicModifyIORef prodding (\_ -> (False,False))
-  r `seq` return () -- avoid space leak
+  _ <- atomicSwapIORef prodding False
   service_loop wakeup delays
 
 -- must agree with rts/win32/ThrIOManager.c
diff --git a/libraries/base/GHC/Event/Control.hs b/libraries/base/GHC/Event/Control.hs
index 5b4a81b38b..779d60d5d7 100644
--- a/libraries/base/GHC/Event/Control.hs
+++ b/libraries/base/GHC/Event/Control.hs
@@ -126,7 +126,7 @@ newControl shouldRegister = allocaArray 2 $ \fds -> do
 -- file after it has been closed.
 closeControl :: Control -> IO ()
 closeControl w = do
-  atomicModifyIORef (controlIsDead w) (\_ -> (True, ()))
+  _ <- atomicSwapIORef (controlIsDead w) True
   _ <- c_close . fromIntegral . controlReadFd $ w
   _ <- c_close . fromIntegral . controlWriteFd $ w
   when (didRegisterWakeupFd w) $ c_setIOManagerWakeupFd (-1)
diff --git a/libraries/base/GHC/Exts.hs b/libraries/base/GHC/Exts.hs
index 11d329ac19..3f5b630ab9 100755
--- a/libraries/base/GHC/Exts.hs
+++ b/libraries/base/GHC/Exts.hs
@@ -34,6 +34,9 @@ module GHC.Exts
         uncheckedIShiftL64#, uncheckedIShiftRA64#,
         isTrue#,
 
+        -- * Compat wrapper
+        atomicModifyMutVar#,
+
         -- * Fusion
         build, augment,
 
@@ -219,3 +222,27 @@ instance IsList CallStack where
   type (Item CallStack) = (String, SrcLoc)
   fromList = fromCallSiteList
   toList   = getCallStack
+
+-- | An implementation of the old @atomicModifyMutVar#@ primop in
+-- terms of the new 'atomicModifyMutVar2#' primop, for backwards
+-- compatibility. The type of this function is a bit bogus. It's
+-- best to think of it as having type
+--
+-- @
+-- atomicModifyMutVar#
+--   :: MutVar# s a
+--   -> (a -> (a, b))
+--   -> State# s
+--   -> (# State# s, b #)
+-- @
+--
+-- but there may be code that uses this with other two-field record
+-- types.
+atomicModifyMutVar#
+  :: MutVar# s a
+  -> (a -> b)
+  -> State# s
+  -> (# State# s, c #)
+atomicModifyMutVar# mv f s =
+  case unsafeCoerce# (atomicModifyMutVar2# mv f s) of
+    (# s', _, ~(_, res) #) -> (# s', res #)
diff --git a/libraries/base/GHC/ForeignPtr.hs b/libraries/base/GHC/ForeignPtr.hs
index 043de1f94b..d80412aac1 100644
--- a/libraries/base/GHC/ForeignPtr.hs
+++ b/libraries/base/GHC/ForeignPtr.hs
@@ -321,7 +321,7 @@ addForeignPtrConcFinalizer_ _ _ =
 
 insertHaskellFinalizer :: IORef Finalizers -> IO () -> IO Bool
 insertHaskellFinalizer r f = do
-  !wasEmpty <- atomicModifyIORef r $ \finalizers -> case finalizers of
+  !wasEmpty <- atomicModifyIORefP r $ \finalizers -> case finalizers of
       NoFinalizers -> (HaskellFinalizers [f], True)
       HaskellFinalizers fs -> (HaskellFinalizers (f:fs), False)
       _ -> noMixingError
@@ -352,8 +352,8 @@ ensureCFinalizerWeak ref@(IORef (STRef r#)) value = do
       NoFinalizers -> IO $ \s ->
           case mkWeakNoFinalizer# r# (unsafeCoerce# value) s of { (# s1, w #) ->
              -- See Note [MallocPtr finalizers] (#10904)
-          case atomicModifyMutVar# r# (update w) s1 of
-              { (# s2, (weak, needKill ) #) ->
+          case atomicModifyMutVar2# r# (update w) s1 of
+              { (# s2, _, (_, (weak, needKill )) #) ->
           if needKill
             then case finalizeWeak# w s2 of { (# s3, _, _ #) ->
               (# s3, weak #) }
@@ -370,7 +370,8 @@ noMixingError = errorWithoutStackTrace $
 
 foreignPtrFinalizer :: IORef Finalizers -> IO ()
 foreignPtrFinalizer r = do
-  fs <- atomicModifyIORef r $ \fs -> (NoFinalizers, fs) -- atomic, see #7170
+  fs <- atomicSwapIORef r NoFinalizers
+             -- atomic, see #7170
   case fs of
     NoFinalizers -> return ()
     CFinalizers w -> IO $ \s -> case finalizeWeak# w s of
diff --git a/libraries/base/GHC/IORef.hs b/libraries/base/GHC/IORef.hs
index 7377690f0f..d04ae728fd 100644
--- a/libraries/base/GHC/IORef.hs
+++ b/libraries/base/GHC/IORef.hs
@@ -1,6 +1,7 @@
 {-# LANGUAGE Unsafe #-}
 {-# LANGUAGE NoImplicitPrelude, MagicHash #-}
-{-# OPTIONS_GHC -funbox-strict-fields #-}
+{-# LANGUAGE UnboxedTuples #-}
+{-# LANGUAGE BangPatterns #-}
 {-# OPTIONS_HADDOCK hide #-}
 
 -----------------------------------------------------------------------------
@@ -19,7 +20,9 @@
 
 module GHC.IORef (
         IORef(..),
-        newIORef, readIORef, writeIORef, atomicModifyIORef
+        newIORef, readIORef, writeIORef, atomicModifyIORef2Lazy,
+        atomicModifyIORef2, atomicModifyIORefLazy_, atomicModifyIORef'_,
+        atomicModifyIORefP, atomicSwapIORef, atomicModifyIORef'
     ) where
 
 import GHC.Base
@@ -48,6 +51,120 @@ readIORef  (IORef var) = stToIO (readSTRef var)
 writeIORef  :: IORef a -> a -> IO ()
 writeIORef (IORef var) v = stToIO (writeSTRef var v)
 
-atomicModifyIORef :: IORef a -> (a -> (a,b)) -> IO b
-atomicModifyIORef (IORef (STRef r#)) f = IO $ \s -> atomicModifyMutVar# r# f s
+-- Atomically apply a function to the contents of an 'IORef',
+-- installing its first component in the 'IORef' and returning
+-- the old contents and the result of applying the function.
+-- The result of the function application (the pair) is not forced.
+-- As a result, this can lead to memory leaks. It is generally better
+-- to use 'atomicModifyIORef2'.
+atomicModifyIORef2Lazy :: IORef a -> (a -> (a,b)) -> IO (a, (a, b))
+atomicModifyIORef2Lazy (IORef (STRef r#)) f =
+  IO (\s -> case atomicModifyMutVar2# r# f s of
+              (# s', old, res #) -> (# s', (old, res) #))
 
+-- Atomically apply a function to the contents of an 'IORef',
+-- installing its first component in the 'IORef' and returning
+-- the old contents and the result of applying the function.
+-- The result of the function application (the pair) is forced,
+-- but neither of its components is.
+atomicModifyIORef2 :: IORef a -> (a -> (a,b)) -> IO (a, (a, b))
+atomicModifyIORef2 ref f = do
+  r@(_old, (_new, _res)) <- atomicModifyIORef2Lazy ref f
+  return r
+
+-- | A version of 'Data.IORef.atomicModifyIORef' that forces
+-- the (pair) result of the function.
+atomicModifyIORefP :: IORef a -> (a -> (a,b)) -> IO b
+atomicModifyIORefP ref f = do
+  (_old, (_,r)) <- atomicModifyIORef2 ref f
+  pure r
+
+-- | Atomically apply a function to the contents of an
+-- 'IORef' and return the old and new values. The result
+-- of the function is not forced. As this can lead to a
+-- memory leak, it is usually better to use `atomicModifyIORef'_`.
+atomicModifyIORefLazy_ :: IORef a -> (a -> a) -> IO (a, a)
+atomicModifyIORefLazy_ (IORef (STRef ref)) f = IO $ \s ->
+  case atomicModifyMutVar_# ref f s of
+    (# s', old, new #) -> (# s', (old, new) #)
+
+-- | Atomically apply a function to the contents of an
+-- 'IORef' and return the old and new values. The result
+-- of the function is forced.
+atomicModifyIORef'_ :: IORef a -> (a -> a) -> IO (a, a)
+atomicModifyIORef'_ ref f = do
+  (old, !new) <- atomicModifyIORefLazy_ ref f
+  return (old, new)
+
+-- | Atomically replace the contents of an 'IORef', returning
+-- the old contents.
+atomicSwapIORef :: IORef a -> a -> IO a
+-- Bad implementation! This will be a primop shortly.
+atomicSwapIORef (IORef (STRef ref)) new = IO $ \s ->
+  case atomicModifyMutVar2# ref (\_old -> Box new) s of
+    (# s', old, Box _new #) -> (# s', old #)
+
+data Box a = Box a
+
+-- | Strict version of 'Data.IORef.atomicModifyIORef'. This forces both
+-- the value stored in the 'IORef' and the value returned. The new value
+-- is installed in the 'IORef' before the returned value is forced.
+-- So
+--
+-- @atomicModifyIORef' ref (\x -> (x+1, undefined))@
+--
+-- will increment the 'IORef' and then throw an exception in the calling
+-- thread.
+--
+-- @since 4.6.0.0
+atomicModifyIORef' :: IORef a -> (a -> (a,b)) -> IO b
+-- See Note [atomicModifyIORef' definition]
+atomicModifyIORef' ref f = do
+  (_old, (_new, !res)) <- atomicModifyIORef2 ref $
+    \old -> case f old of
+       r@(!_new, _res) -> r
+  pure res
+
+-- Note [atomicModifyIORef' definition]
+-- ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
+--
+-- atomicModifyIORef' was historically defined
+--
+--    atomicModifyIORef' ref f = do
+--        b <- atomicModifyIORef ref $ \a ->
+--                case f a of
+--                    v@(a',_) -> a' `seq` v
+--        b `seq` return b
+--
+-- The most obvious definition, now that we have atomicModifyMutVar2#,
+-- would be
+--
+--    atomicModifyIORef' ref f = do
+--      (_old, (!_new, !res)) <- atomicModifyIORef2 ref f
+--      pure res
+--
+-- Why do we force the new value on the "inside" instead of afterwards?
+-- I initially thought the latter would be okay, but then I realized
+-- that if we write
+--
+--   atomicModifyIORef' ref $ \x -> (x + 5, x - 5)
+--
+-- then we'll end up building a pair of thunks to calculate x + 5
+-- and x - 5. That's no good! With the more complicated definition,
+-- we avoid this problem; the result pair is strict in the new IORef
+-- contents. Of course, if the function passed to atomicModifyIORef'
+-- doesn't inline, we'll build a closure for it. But that was already
+-- true for the historical definition of atomicModifyIORef' (in terms
+-- of atomicModifyIORef), so we shouldn't lose anything. Note that
+-- in keeping with the historical behavior, we *don't* propagate the
+-- strict demand on the result inwards. In particular,
+--
+--   atomicModifyIORef' ref (\x -> (x + 1, undefined))
+--
+-- will increment the IORef and throw an exception; it will not
+-- install an undefined value in the IORef.
+--
+-- A clearer version, in my opinion (but one quite incompatible with
+-- the traditional one) would only force the new IORef value and not
+-- the result. This version would have been relatively inefficient
+-- to implement using atomicModifyMutVar#, but is just fine now.
diff --git a/libraries/base/System/IO.hs b/libraries/base/System/IO.hs
index e02c30d63c..4a9b8837ab 100644
--- a/libraries/base/System/IO.hs
+++ b/libraries/base/System/IO.hs
@@ -583,7 +583,7 @@ tempCounter = unsafePerformIO $ newIORef 0
 rand_string :: IO String
 rand_string = do
   r1 <- c_getpid
-  r2 <- atomicModifyIORef tempCounter (\n -> (n+1, n))
+  (r2, _) <- atomicModifyIORef'_ tempCounter (+1)
   return $ show r1 ++ "-" ++ show r2
 
 data OpenNewFileResult
diff --git a/rts/PrimOps.cmm b/rts/PrimOps.cmm
index 6081fabe93..4e4c6a6947 100644
--- a/rts/PrimOps.cmm
+++ b/rts/PrimOps.cmm
@@ -564,9 +564,9 @@ stg_casMutVarzh ( gcptr mv, gcptr old, gcptr new )
 #endif
 }
 
-stg_atomicModifyMutVarzh ( gcptr mv, gcptr f )
+stg_atomicModifyMutVar2zh ( gcptr mv, gcptr f )
 {
-    W_ z, x, y, r, h;
+    W_ z, x, y, h;
 
     /* If x is the current contents of the MutVar#, then
        We want to make the new contents point to
@@ -575,13 +575,12 @@ stg_atomicModifyMutVarzh ( gcptr mv, gcptr f )
 
        and the return value is
 
-         (sel_1 (f x))
+         (# x, (f x) #)
 
         obviously we can share (f x).
 
          z = [stg_ap_2 f x]  (max (HS + 2) MIN_UPD_SIZE)
          y = [stg_sel_0 z]   (max (HS + 1) MIN_UPD_SIZE)
-         r = [stg_sel_1 z]   (max (HS + 1) MIN_UPD_SIZE)
     */
 
 #if defined(MIN_UPD_SIZE) && MIN_UPD_SIZE > 1
@@ -600,7 +599,7 @@ stg_atomicModifyMutVarzh ( gcptr mv, gcptr f )
 #define TICK_ALLOC_THUNK_2() TICK_ALLOC_UP_THK(WDS(2),0)
 #endif
 
-#define SIZE (THUNK_2_SIZE + THUNK_1_SIZE + THUNK_1_SIZE)
+#define SIZE (THUNK_2_SIZE + THUNK_1_SIZE)
 
     HP_CHK_GEN_TICKY(SIZE);
 
@@ -618,13 +617,6 @@ stg_atomicModifyMutVarzh ( gcptr mv, gcptr f )
     LDV_RECORD_CREATE(y);
     StgThunk_payload(y,0) = z;
 
-    TICK_ALLOC_THUNK_1();
-    CCCS_ALLOC(THUNK_1_SIZE);
-    r = y - THUNK_1_SIZE;
-    SET_HDR(r, stg_sel_1_upd_info, CCCS);
-    LDV_RECORD_CREATE(r);
-    StgThunk_payload(r,0) = z;
-
   retry:
     x = StgMutVar_var(mv);
     StgThunk_payload(z,1) = x;
@@ -639,9 +631,62 @@ stg_atomicModifyMutVarzh ( gcptr mv, gcptr f )
         ccall dirty_MUT_VAR(BaseReg "ptr", mv "ptr");
     }
 
-    return (r);
+    return (x,z);
+}
+
+stg_atomicModifyMutVarzuzh ( gcptr mv, gcptr f )
+{
+    W_ z, x, h;
+
+    /* If x is the current contents of the MutVar#, then
+       We want to make the new contents point to
+
+         (f x)
+
+       and the return value is
+
+         (# x, (f x) #)
+
+        obviously we can share (f x).
+
+         z = [stg_ap_2 f x]  (max (HS + 2) MIN_UPD_SIZE)
+    */
+
+#if defined(MIN_UPD_SIZE) && MIN_UPD_SIZE > 2
+#define THUNK_SIZE (SIZEOF_StgThunkHeader + WDS(MIN_UPD_SIZE))
+#define TICK_ALLOC_THUNK() TICK_ALLOC_UP_THK(WDS(2),WDS(MIN_UPD_SIZE-2))
+#else
+#define THUNK_SIZE (SIZEOF_StgThunkHeader + WDS(2))
+#define TICK_ALLOC_THUNK() TICK_ALLOC_UP_THK(WDS(2),0)
+#endif
+
+    HP_CHK_GEN_TICKY(THUNK_SIZE);
+
+    TICK_ALLOC_THUNK();
+    CCCS_ALLOC(THUNK_SIZE);
+    z = Hp - THUNK_SIZE + WDS(1);
+    SET_HDR(z, stg_ap_2_upd_info, CCCS);
+    LDV_RECORD_CREATE(z);
+    StgThunk_payload(z,0) = f;
+
+  retry:
+    x = StgMutVar_var(mv);
+    StgThunk_payload(z,1) = x;
+#if defined(THREADED_RTS)
+    (h) = prim %cmpxchgW(mv + SIZEOF_StgHeader + OFFSET_StgMutVar_var, x, z);
+    if (h != x) { goto retry; }
+#else
+    StgMutVar_var(mv) = z;
+#endif
+
+    if (GET_INFO(mv) == stg_MUT_VAR_CLEAN_info) {
+        ccall dirty_MUT_VAR(BaseReg "ptr", mv "ptr");
+    }
+
+    return (x,z);
 }
 
+
 /* -----------------------------------------------------------------------------
    Weak Pointer Primitives
    -------------------------------------------------------------------------- */
diff --git a/rts/RtsSymbols.c b/rts/RtsSymbols.c
index 87fa98dd4f..1543a9df5f 100644
--- a/rts/RtsSymbols.c
+++ b/rts/RtsSymbols.c
@@ -669,7 +669,8 @@
       SymI_HasProto(stg_newMutVarzh)                                    \
       SymI_HasProto(stg_newTVarzh)                                      \
       SymI_HasProto(stg_noDuplicatezh)                                  \
-      SymI_HasProto(stg_atomicModifyMutVarzh)                           \
+      SymI_HasProto(stg_atomicModifyMutVar2zh)                          \
+      SymI_HasProto(stg_atomicModifyMutVarzuzh)                         \
       SymI_HasProto(stg_casMutVarzh)                                    \
       SymI_HasProto(stg_newPinnedByteArrayzh)                           \
       SymI_HasProto(stg_newAlignedPinnedByteArrayzh)                    \
diff --git a/testsuite/tests/perf/should_run/all.T b/testsuite/tests/perf/should_run/all.T
index a8107c487a..7a52492f45 100644
--- a/testsuite/tests/perf/should_run/all.T
+++ b/testsuite/tests/perf/should_run/all.T
@@ -317,10 +317,11 @@ test('T7257',
                         # expected value: 1246287228 (i386/Linux)
                         # 2016-04-06: 989850664 (i386/Linux) no idea what happened
                         # 2017-03-25: 869850704 (x86/Linux, 64-bit machine) probably sizeExpr fix
-                       (wordsize(64), 1414893248, 5)]),
+                       (wordsize(64), 1297293264, 5)]),
                         # 2012-09-21: 1774893760 (amd64/Linux)
                         # 2015-11-03: 1654893248 (amd64/Linux)
                         # 2016-06-22: 1414893248 (amd64/Linux, sizeExpr fix)
+                        # 2018-06-22: 1297293264 (amd64/Linux, atomicModifyMutVar# replacement)
       stats_num_field('peak_megabytes_allocated',
                       [(wordsize(32), 217, 5),
                         # 2012-10-08: 217 (x86/Linux)