WorkWrap: Relax "splitFun" warning for join points (#23113)wip/T23113

... and document our ponderings in `Note [Threshold arity for join points]`.
I also introduced a new warning in `finaliseArgBoxities` to see where
we currently are a bit too optimistic wrt. boxity.

Fixes #23113

3 files changed, 78 insertions, 2 deletions

diff --git a/compiler/GHC/Core/Opt/DmdAnal.hs b/compiler/GHC/Core/Opt/DmdAnal.hs
index ece13d894b..cf76563e3a 100644
--- a/compiler/GHC/Core/Opt/DmdAnal.hs
+++ b/compiler/GHC/Core/Opt/DmdAnal.hs
@@ -1940,6 +1940,12 @@ finaliseArgBoxities env fn threshold_arity rhs_dmds div rhs
     --   vcat [text "function:" <+> ppr fn
     --        , text "dmds before:" <+> ppr (map idDemandInfo (filter isId bndrs))
     --        , text "dmds after: " <+>  ppr arg_dmds' ]) $
+    warnPprTrace (isJoinId fn && length rhs_dmds > threshold_arity)
+                 "finaliseArgBoxities: excess rhs_dmds of join point"
+                 (ppr fn <+> ppr threshold_arity <+> ppr rhs_dmds) $
+                  -- It is far from clear that it's OK to ignore excess rhs_dmds
+                  -- here rather than zap all boxity. Hence we warn to collect
+                  -- some examples. See Note [Threshold arity of join points]
     (arg_dmds', set_lam_dmds arg_dmds' rhs)
     -- set_lam_dmds: we must attach the final boxities to the lambda-binders
     -- of the function, both because that's kosher, and because CPR analysis
diff --git a/compiler/GHC/Core/Opt/WorkWrap.hs b/compiler/GHC/Core/Opt/WorkWrap.hs
index 29f1e3973f..cf212a86a0 100644
--- a/compiler/GHC/Core/Opt/WorkWrap.hs
+++ b/compiler/GHC/Core/Opt/WorkWrap.hs
@@ -758,9 +758,11 @@ by LitRubbish (see Note [Drop absent bindings]) so there is no great harm.
 splitFun :: WwOpts -> Id -> CoreExpr -> UniqSM [(Id, CoreExpr)]
 splitFun ww_opts fn_id rhs
   | Just (arg_vars, body) <- collectNValBinders_maybe (length wrap_dmds) rhs
-  = warnPprTrace (not (wrap_dmds `lengthIs` (arityInfo fn_info)))
+  = warnPprTrace (if isJoinId fn_id
+                    then not (arg_vars  `lengthAtMost` idJoinArity fn_id) -- See Note [Threshold arity of join points]
+                    else not (wrap_dmds `lengthIs`     (arityInfo fn_info)))
                  "splitFun"
-                 (ppr fn_id <+> (ppr wrap_dmds $$ ppr cpr)) $
+                 (sep [ ppr fn_id, ppr (arityInfo fn_info), ppr wrap_dmds, ppr cpr]) $
     do { mb_stuff <- mkWwBodies ww_opts fn_id arg_vars (exprType body) wrap_dmds cpr
        ; case mb_stuff of
             Nothing -> -- No useful wrapper; leave the binding alone
diff --git a/compiler/GHC/Types/Demand.hs b/compiler/GHC/Types/Demand.hs
index 09b08b7f36..21c2cdf4e4 100644
--- a/compiler/GHC/Types/Demand.hs
+++ b/compiler/GHC/Types/Demand.hs
@@ -2143,6 +2143,74 @@ type's depth! So in mkDmdSigForArity we make sure to trim the list of
 argument demands to the given threshold arity. Call sites will make sure that
 this corresponds to the arity of the call demand that elicited the wrapped
 demand type. See also Note [What are demand signatures?].
+
+See also Note [Threshold arity of join points] for how the threshold arity of
+join points is special.
+
+Note [Threshold arity of join points]
+~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
+The threshold arity in the demand signature of a join point might be
+
+  * Less than `idArity`:
+      join j x = \pqr. blah in ...(jump j 1)... (jump j 2)...
+    Here idArity is 4, but join-arity is 1.
+  * More than `idArity`:
+      f g = g 42   :: <C(1,L)>
+      h x = f (join j y = (+) y in ... j 13 ...)
+    Note that f's demand on its arg must apply is put on the join expr and hence
+    its RHS.
+    How this is achieved is described in Note [Demand analysis for join points].
+    In this Note, we refer to it as the know-context assumption.
+
+The latter example is interesting, because there analysis ends up with a demand
+/type/ of <1!L><1!L> for the RHS of the `j`, based on the arity 2 signature of
+`(+)`.
+
+Why trim down demand signatures?
+--------------------------------
+When we finalise the demand /signature/ for `j`, we have to trim this signature
+to have a depth (number of arg dmds) of 1.
+
+The reason for that is a tension between the most precise demand signature
+possible (<1L><1L>) and its piggy-backed /boxity/ signature (<!><!>). The latter
+is relevant to WW and if we keep length 2, WW would end up eta-expanding `j` for
+arity 2, destroying its joinpointhood.
+
+So `finaliseArgBoxities` will instead drop one arg, but /keep the boxity info
+on it/. The former is sound because demand analysis knows that `j` will always
+be called with 2 arguments anyway, but the latter may introduce reboxing for `j`
+above if we don't inline `(+)` (which presumably we'll do only for arity 2):
+  join   j y = case y of I# y# -> $wj y#
+       $wj y# = (+) (I# y#)
+  in ...
+Note the reboxing in $wj. On the other hand, all reboxing would
+vanish if we inlined `(+)`, so for now we simply emit a warning in
+GHC.Core.Opt.WorkWrap.splitFun to collect examples when such a potentially
+undesirable split happens.
+
+Undesirable consequence of trimming
+-----------------------------------
+The final demand signature for `j` above is `<1L>` (nevermind boxity here), but
+the *correct* threshold arity is still that of the original demand type,
+namely 2. So we violate the coupling of arg dmds and threshold arity in demand
+signatures we so painfully stated in Note [Understanding DmdType and DmdSig].
+
+It would be unsound to unleash the signature in an arity 1 call context such as
+`Just (j x)`, for example. Nevertheless, because `j` is a join point, all its
+call contexts are fixed and won't change unless the demand on the whole join
+expression changes, so every practical use of the signature is sound even for
+arity 1 (all calls of syntactic arity 1 will ultimately be a call with arity 2).
+
+BUT, certain transformations can destroy the known-context assumption. For
+example, when we demote `j` to a let binding and realise that its RHS does not
+reference `x`, we might be tempted to float it to the top-level. If we do so,
+we should be sure to discard its demand signature, because there is nothing
+preventing to e.g., CSE two calls `j 13` and/or perform it repeatedly.
+
+FloatOut will actually float join points to top-level but is unconcerned about
+this issue. Fortunately, Demand Analysis runs after FloatOut, so this has not
+become an issue in practice; still, it is worth keeping an eye out for and at
+least documenting this potential issue.
 -}
 
 -- | The depth of the wrapped 'DmdType' encodes the arity at which it is safe