Make keepAlive# out-of-line

This is a naive approach to fixing the unsoundness noticed in #21708.
Specifically, we remove the lowering of `keepAlive#` via CorePrep and
instead turn it into an out-of-line primop.
This is simple, inefficient (since the continuation must now be heap
allocated), but good enough for 9.4.1. We will revisit this
(particiularly via #16098) in a future release.

Metric Increase:
    T4978
    T7257
    T9203
    T13701
    T14697

(cherry picked from commit d75c540d439510491b45f64c1113762dcb251ae1)

7 files changed, 31 insertions, 36 deletions

diff --git a/compiler/GHC/Builtin/primops.txt.pp b/compiler/GHC/Builtin/primops.txt.pp
index cf132d0a7f..4db381f691 100644
--- a/compiler/GHC/Builtin/primops.txt.pp
+++ b/compiler/GHC/Builtin/primops.txt.pp
@@ -3283,6 +3283,7 @@ primop KeepAliveOp "keepAlive#" GenPrimOp
    { \tt{keepAlive# x s k} keeps the value \tt{x} alive during the execution
      of the computation \tt{k}. }
    with
+   out_of_line = True
    strictness = { \ _arity -> mkClosedDmdSig [topDmd, topDmd, strictOnceApply1Dmd] topDiv }
 
 
diff --git a/compiler/GHC/CoreToStg/Prep.hs b/compiler/GHC/CoreToStg/Prep.hs
index 80d64f4fb1..b1b4399f58 100644
--- a/compiler/GHC/CoreToStg/Prep.hs
+++ b/compiler/GHC/CoreToStg/Prep.hs
@@ -29,10 +29,7 @@ import GHC.Tc.Utils.Env
 import GHC.Unit
 
 import GHC.Builtin.Names
-import GHC.Builtin.PrimOps
-import GHC.Builtin.PrimOps.Ids (primOpId)
 import GHC.Builtin.Types
-import GHC.Builtin.Types.Prim ( realWorldStatePrimTy )
 
 import GHC.Core.Utils
 import GHC.Core.Opt.Arity
@@ -1050,36 +1047,6 @@ cpeApp top_env expr
         = let (terminal, args') = collect_args arg
           in cpe_app env terminal (args' ++ args)
 
-    -- See Note [keepAlive# magic].
-    cpe_app env
-            (Var f)
-            args
-        | Just KeepAliveOp <- isPrimOpId_maybe f
-        , CpeApp (Type arg_lev)
-          : CpeApp (Type _result_rep)
-          : CpeApp (Type arg_ty)
-          : CpeApp (Type result_ty)
-          : CpeApp arg
-          : CpeApp s0
-          : CpeApp k
-          : rest <- args
-        = do { y  <- newVar (cpSubstTy env result_ty)
-             ; s2 <- newVar realWorldStatePrimTy
-             ; -- beta reduce if possible
-             ; (floats, k') <- case k of
-                  Lam s body -> cpe_app (extendCorePrepEnvExpr env s s0) body rest
-                  _          -> cpe_app env k (CpeApp s0 : rest)
-             ; let touchId = primOpId TouchOp
-                   expr = Case k' y result_ty [Alt DEFAULT [] rhs]
-                   rhs = let scrut = mkApps (Var touchId) [Type arg_lev, Type arg_ty, arg, Var realWorldPrimId]
-                         in Case scrut s2 result_ty [Alt DEFAULT [] (Var y)]
-             ; (floats', expr') <- cpeBody env expr
-             ; return (floats `appendFloats` floats', expr')
-             }
-        | Just KeepAliveOp <- isPrimOpId_maybe f
-        = pprPanic "invalid keepAlive# application" $
-            vcat [ text "args:" <+> ppr args ]
-
     -- runRW# magic
     cpe_app env (Var f) (CpeApp _runtimeRep@Type{} : CpeApp _type@Type{} : CpeApp arg : rest)
         | f `hasKey` runRWKey
diff --git a/compiler/GHC/StgToCmm/Prim.hs b/compiler/GHC/StgToCmm/Prim.hs
index b49ad24edd..a0e6ecf871 100644
--- a/compiler/GHC/StgToCmm/Prim.hs
+++ b/compiler/GHC/StgToCmm/Prim.hs
@@ -1640,9 +1640,7 @@ emitPrimOp cfg primop =
   TraceEventBinaryOp -> alwaysExternal
   TraceMarkerOp -> alwaysExternal
   SetThreadAllocationCounter -> alwaysExternal
-
-  -- See Note [keepAlive# magic] in GHC.CoreToStg.Prep.
-  KeepAliveOp -> panic "keepAlive# should have been eliminated in CorePrep"
+  KeepAliveOp -> alwaysExternal
 
  where
   profile  = stgToCmmProfile  cfg
diff --git a/rts/PrimOps.cmm b/rts/PrimOps.cmm
index 84c5850f97..2db2703afc 100644
--- a/rts/PrimOps.cmm
+++ b/rts/PrimOps.cmm
@@ -2918,3 +2918,23 @@ stg_setThreadAllocationCounterzh ( I64 counter )
     StgTSO_alloc_limit(CurrentTSO) = counter + TO_I64(offset);
     return ();
 }
+
+
+#define KEEP_ALIVE_FRAME_FIELDS(w_,p_,info_ptr,p1,p2,c)   \
+  w_ info_ptr,                                            \
+  PROF_HDR_FIELDS(w_,p1,p2)                               \
+  p_ c
+
+stg_keepAlivezh ( P_ c,      /* :: v */
+                  P_ io      /* :: IO p */ )
+{
+    STK_CHK_GEN();
+    jump stg_ap_v_fast
+        (KEEP_ALIVE_FRAME_FIELDS(,,stg_keepAlive_frame_info, CCCS, 0, c))(io);
+}
+
+INFO_TABLE_RET(stg_keepAlive_frame, RET_SMALL, KEEP_ALIVE_FRAME_FIELDS(W_,P_, info_ptr, p1, p2, c))
+    return (P_ ret)
+{
+    return (ret);
+}
diff --git a/rts/RtsSymbols.c b/rts/RtsSymbols.c
index 2818df6ff3..3285aa4ff4 100644
--- a/rts/RtsSymbols.c
+++ b/rts/RtsSymbols.c
@@ -785,6 +785,7 @@ extern char **environ;
       SymI_HasDataProto(stg_raiseUnderflowzh)                               \
       SymI_HasDataProto(stg_raiseOverflowzh)                                \
       SymI_HasDataProto(stg_raiseIOzh)                                      \
+      SymI_HasDataProto(stg_keepAlivezh)                                    \
       SymI_HasDataProto(stg_paniczh)                                        \
       SymI_HasDataProto(stg_absentErrorzh)                                  \
       SymI_HasDataProto(stg_readTVarzh)                                     \
diff --git a/rts/include/rts/storage/Closures.h b/rts/include/rts/storage/Closures.h
index 41861abac9..ce1809b930 100644
--- a/rts/include/rts/storage/Closures.h
+++ b/rts/include/rts/storage/Closures.h
@@ -253,6 +253,12 @@ typedef struct _StgUpdateFrame {
 } StgUpdateFrame;
 
 
+// Closure types: RET_SMALL
+typedef struct {
+    StgHeader  header;
+    StgClosure *c;
+} StgKeepAliveFrame;
+
 // Stack frame, when we call catch one of these will be put on the stack so we
 // know to handle exceptions with the supplied handler
 //
diff --git a/rts/include/stg/MiscClosures.h b/rts/include/stg/MiscClosures.h
index e87eba0931..02212d99f8 100644
--- a/rts/include/stg/MiscClosures.h
+++ b/rts/include/stg/MiscClosures.h
@@ -61,6 +61,7 @@ RTS_RET(stg_unmaskAsyncExceptionszh_ret);
 RTS_RET(stg_maskUninterruptiblezh_ret);
 RTS_RET(stg_maskAsyncExceptionszh_ret);
 RTS_RET(stg_stack_underflow_frame);
+RTS_RET(stg_keepAlive_frame);
 RTS_RET(stg_restore_cccs);
 RTS_RET(stg_restore_cccs_eval);
 
@@ -497,6 +498,7 @@ RTS_FUN_DECL(stg_raiseUnderflowzh);
 RTS_FUN_DECL(stg_raiseOverflowzh);
 RTS_FUN_DECL(stg_raiseIOzh);
 RTS_FUN_DECL(stg_paniczh);
+RTS_FUN_DECL(stg_keepAlivezh);
 RTS_FUN_DECL(stg_absentErrorzh);
 
 RTS_FUN_DECL(stg_makeStableNamezh);