Cmm: fix sinking after suspendThread

Suppose a safe call: myCall(x,y,z)

It is lowered into three unsafe calls in Cmm:

  r = suspendThread(...);
  myCall(x,y,z);
  resumeThread(r);

Consider the following situation for myCall arguments:

  x = Sp[..] -- stack
  y = Hp[..] -- heap
  z = R1     -- global register
  r = suspendThread(...);
  myCall(x,y,z);
  resumeThread(r);

The sink pass assumes that unsafe calls clobber memory (heap and stack),
hence x and y assignments are not sunk after `suspendThread`. The sink
pass also correctly handles global register clobbering for all unsafe
calls, except `suspendThread`!

`suspendThread` is special because it releases the capability the thread
is running on. Hence the sink pass must also take into account global
registers that are mapped into memory (in the capability).

In the example above, we could get:

  r = suspendThread(...);
  z = R1
  myCall(x,y,z);
  resumeThread(r);

But this transformation isn't valid if R1 is (BaseReg->rR1) as BaseReg
is invalid between suspendThread and resumeThread. This caused argument
corruption at least with the C backend ("unregisterised") in #19237.

Fix #19237

4 files changed, 31 insertions, 20 deletions

diff --git a/compiler/GHC/Cmm/LayoutStack.hs b/compiler/GHC/Cmm/LayoutStack.hs
index 2fdbb1fe5a..b996427bba 100644
--- a/compiler/GHC/Cmm/LayoutStack.hs
+++ b/compiler/GHC/Cmm/LayoutStack.hs
@@ -12,14 +12,11 @@ import GHC.StgToCmm.Monad      ( newTemp  ) -- XXX layering violation
 import GHC.StgToCmm.Utils      ( callerSaveVolatileRegs  ) -- XXX layering violation
 import GHC.StgToCmm.Foreign    ( saveThreadState, loadThreadState ) -- XXX layering violation
 
-import GHC.Types.Basic
 import GHC.Cmm
 import GHC.Cmm.Info
 import GHC.Cmm.BlockId
-import GHC.Cmm.CLabel
 import GHC.Cmm.Utils
 import GHC.Cmm.Graph
-import GHC.Types.ForeignCall
 import GHC.Cmm.Liveness
 import GHC.Cmm.ProcPoint
 import GHC.Runtime.Heap.Layout
@@ -34,7 +31,6 @@ import GHC.Types.Unique.FM
 import GHC.Utils.Misc
 
 import GHC.Driver.Session
-import GHC.Data.FastString
 import GHC.Utils.Outputable hiding ( isEmpty )
 import GHC.Utils.Panic
 import qualified Data.Set as Set
@@ -1190,21 +1186,14 @@ lowerSafeForeignCall profile block
   | otherwise = return block
 
 
-foreignLbl :: FastString -> CmmExpr
-foreignLbl name = CmmLit (CmmLabel (mkForeignLabel name Nothing ForeignLabelInExternalPackage IsFunction))
-
 callSuspendThread :: Platform -> LocalReg -> Bool -> CmmNode O O
 callSuspendThread platform id intrbl =
-  CmmUnsafeForeignCall
-       (ForeignTarget (foreignLbl (fsLit "suspendThread"))
-        (ForeignConvention CCallConv [AddrHint, NoHint] [AddrHint] CmmMayReturn))
+  CmmUnsafeForeignCall (PrimTarget MO_SuspendThread)
        [id] [baseExpr, mkIntExpr platform (fromEnum intrbl)]
 
 callResumeThread :: LocalReg -> LocalReg -> CmmNode O O
 callResumeThread new_base id =
-  CmmUnsafeForeignCall
-       (ForeignTarget (foreignLbl (fsLit "resumeThread"))
-            (ForeignConvention CCallConv [AddrHint] [AddrHint] CmmMayReturn))
+  CmmUnsafeForeignCall (PrimTarget MO_ResumeThread)
        [new_base] [CmmReg (CmmLocal id)]
 
 -- -----------------------------------------------------------------------------
diff --git a/compiler/GHC/Cmm/MachOp.hs b/compiler/GHC/Cmm/MachOp.hs
index 558cb13a7e..b91263ce47 100644
--- a/compiler/GHC/Cmm/MachOp.hs
+++ b/compiler/GHC/Cmm/MachOp.hs
@@ -638,6 +638,12 @@ data CallishMachOp
   -- Should be an AtomicRMW variant eventually.
   -- Sequential consistent.
   | MO_Xchg Width
+
+  -- These rts provided functions are special: suspendThread releases the
+  -- capability, hence we mustn't sink any use of data stored in the capability
+  -- after this instruction.
+  | MO_SuspendThread
+  | MO_ResumeThread
   deriving (Eq, Show)
 
 -- | The operation to perform atomically.
@@ -653,13 +659,16 @@ data AtomicMachOp =
 pprCallishMachOp :: CallishMachOp -> SDoc
 pprCallishMachOp mo = text (show mo)
 
+-- | Return (results_hints,args_hints)
 callishMachOpHints :: CallishMachOp -> ([ForeignHint], [ForeignHint])
 callishMachOpHints op = case op of
-  MO_Memcpy _  -> ([], [AddrHint,AddrHint,NoHint])
-  MO_Memset _  -> ([], [AddrHint,NoHint,NoHint])
-  MO_Memmove _ -> ([], [AddrHint,AddrHint,NoHint])
-  MO_Memcmp _  -> ([], [AddrHint, AddrHint, NoHint])
-  _            -> ([],[])
+  MO_Memcpy _      -> ([], [AddrHint,AddrHint,NoHint])
+  MO_Memset _      -> ([], [AddrHint,NoHint,NoHint])
+  MO_Memmove _     -> ([], [AddrHint,AddrHint,NoHint])
+  MO_Memcmp _      -> ([], [AddrHint, AddrHint, NoHint])
+  MO_SuspendThread -> ([AddrHint], [AddrHint,NoHint])
+  MO_ResumeThread  -> ([AddrHint], [AddrHint])
+  _                -> ([],[])
   -- empty lists indicate NoHint
 
 -- | The alignment of a 'memcpy'-ish operation.
diff --git a/compiler/GHC/Cmm/Parser.y b/compiler/GHC/Cmm/Parser.y
index 4576eb9b38..a83feff8cf 100644
--- a/compiler/GHC/Cmm/Parser.y
+++ b/compiler/GHC/Cmm/Parser.y
@@ -1027,6 +1027,9 @@ callishMachOps platform = listToUFM $
         ( "memmove", memcpyLikeTweakArgs MO_Memmove ),
         ( "memcmp", memcpyLikeTweakArgs MO_Memcmp ),
 
+        ( "suspendThread", (MO_SuspendThread,) ),
+        ( "resumeThread",  (MO_ResumeThread,) ),
+
         ("prefetch0", (MO_Prefetch_Data 0,)),
         ("prefetch1", (MO_Prefetch_Data 1,)),
         ("prefetch2", (MO_Prefetch_Data 2,)),
diff --git a/compiler/GHC/Cmm/Sink.hs b/compiler/GHC/Cmm/Sink.hs
index cd13d6b655..3ef58b3648 100644
--- a/compiler/GHC/Cmm/Sink.hs
+++ b/compiler/GHC/Cmm/Sink.hs
@@ -622,10 +622,16 @@ conflicts platform (r, rhs, addr) node
   -- (5) foreign calls clobber heap: see Note [Foreign calls clobber heap]
   | CmmUnsafeForeignCall{} <- node, memConflicts addr AnyMem      = True
 
-  -- (6) native calls clobber any memory
+  -- (6) suspendThread clobbers every global register not backed by a real
+  -- register. It also clobbers heap and stack but this is handled by (5)
+  | CmmUnsafeForeignCall (PrimTarget MO_SuspendThread) _ _ <- node
+  , foldRegsUsed platform (\b g -> globalRegMaybe platform g == Nothing || b) False rhs
+  = True
+
+  -- (7) native calls clobber any memory
   | CmmCall{} <- node, memConflicts addr AnyMem                   = True
 
-  -- (7) otherwise, no conflict
+  -- (8) otherwise, no conflict
   | otherwise = False
 
 {- Note [Inlining foldRegsDefd]
@@ -759,6 +765,10 @@ data AbsMem
 -- Some CallishMachOp imply a memory barrier e.g. AtomicRMW and
 -- therefore we should never float any memory operations across one of
 -- these calls.
+--
+-- `suspendThread` releases the capability used by the thread, hence we mustn't
+-- float accesses to heap, stack or virtual global registers stored in the
+-- capability (e.g. with unregisterised build, see #19237).
 
 
 bothMems :: AbsMem -> AbsMem -> AbsMem