How can I get rid of `let` in Core?

Question

I have a function that is called frequently in an internal loop. It looks like this:

import qualified Data.Vector.Storable as SV

newtype Timedelta = Timedelta Double

cklsLogDens :: SV.Vector Double -> Timedelta -> Double -> Double -> Double
cklsLogDens p (Timedelta dt) x0 x1 = if si <= 0 then -1e50 else c - 0.5*((x1-mu)/sd)^2 
  where
    al  = p `SV.unsafeIndex` 0
    be  = p `SV.unsafeIndex` 1
    si  = p `SV.unsafeIndex` 2
    xi  = p `SV.unsafeIndex` 3
    sdt = sqrt dt
    mu  = x0 + (al + be*x0)*dt
    sd  = si * (x0 ** xi) * sdt
    c   = sd `seq` -0.5 * log (2*pi*sd^2)

(Data.Vector.Storable is used because this function needs to work on data from a C function later)

GHC has optimized this very nicely (all variables and ops are primitives as far as I can tell), but looking at core, there is one let that is still inside of (what was) the body of the function. I have read here (and somewhere else I don't remember) that 'lets' allocate lazy thunks and can thus be bad for performance in tight loops. Can I get rid of it? If it all possible I would prefer not converting my function into 20 case statements, but if that is too much to ask I'll accept.

Here is the Core:

$wloop_s4Li [Occ=LoopBreaker]
  :: GHC.Prim.Double#
     -> GHC.Prim.Int# -> GHC.Prim.Int# -> GHC.Prim.Double#
[LclId, Arity=3, Str=DmdType LLL]
$wloop_s4Li =
  \ (ww_X4OR :: GHC.Prim.Double#)
    (ww1_X4OW :: GHC.Prim.Int#)
    (ww2_X4P1 :: GHC.Prim.Int#) ->
    case GHC.Prim.<# ww1_X4OW ww2_X4P1 of _ {
      GHC.Types.False -> ww_X4OR;
      GHC.Types.True ->
        case GHC.Prim.<=## x_a4tg 0.0 of _ {
          GHC.Types.False ->
            case GHC.Prim.indexDoubleArray#
                   rb2_a4rT (GHC.Prim.+# rb_a4rR (GHC.Prim.-# ww1_X4OW 1))
            of wild17_X4xM { __DEFAULT ->

            let {
      ----  ^^^^ want to get rid off this! 
      ----
      ----
              ipv1_X2S8 [Dmd=Just L] :: GHC.Prim.Double#
              [LclId, Str=DmdType]
              ipv1_X2S8 =
                GHC.Prim.*##
                  (GHC.Prim.*## x_a4tg (GHC.Prim.**## wild17_X4xM y_a3BN))
                  (GHC.Prim.sqrtDouble# tpl1_B3) } in
            case GHC.Prim.logDouble#
                   (GHC.Prim.*##
                      6.283185307179586 (GHC.Prim.*## ipv1_X2S8 ipv1_X2S8))
            of wild18_X3Gn { __DEFAULT ->
            case GHC.Prim.indexDoubleArray#
                   rb2_a4rT (GHC.Prim.+# rb_a4rR ww1_X4OW)
            of wild19_X4AY { __DEFAULT ->
            case GHC.Prim./##
                   (GHC.Prim.-##
                      wild19_X4AY
                      (GHC.Prim.+##
                         wild17_X4xM
                         (GHC.Prim.*##
                            (GHC.Prim.+##
                               x1_X3GA (GHC.Prim.*## x2_X3cb wild17_X4xM))
                            tpl1_B3)))
                   ipv1_X2S8
            of wild20_X3x8 { __DEFAULT ->
            $wloop_s4Li
              (GHC.Prim.+##
                 ww_X4OR
                 (GHC.Prim.-##
                    (GHC.Prim.negateDouble# (GHC.Prim.*## 0.5 wild18_X3Gn))
                    (GHC.Prim.*##
                       0.5 (GHC.Prim.*## wild20_X3x8 wild20_X3x8))))
              (GHC.Prim.+# ww1_X4OW 1)
              ww2_X4P1
            }
            }
            }
            };
          GHC.Types.True ->
            $wloop_s4Li
              (GHC.Prim.+## ww_X4OR -1.0e50)
              (GHC.Prim.+# ww1_X4OW 1)
              ww2_X4P1
        }
    }; }

(Yes, of course, since you must ask, I am spending waaay too much time on premature optimization...)

Here is the current version with NOINLINE

import qualified Data.Vector.Storable as SV

newtype Timedelta = Timedelta Double

cklsLogDens :: SV.Vector Double -> Timedelta -> Double -> Double -> Double
{-# NOINLINE cklsLogDens #-}
cklsLogDens p (Timedelta dt) x0 x1 = si `seq` (if si <= 0 then -1e50 else (sd `seq` (c - 0.5*((x1-mu)/sd)^2)))
  where
    al  = p `SV.unsafeIndex` 0
    be  = p `SV.unsafeIndex` 1
    si  = p `SV.unsafeIndex` 2
    xi  = p `SV.unsafeIndex` 3
    sdt = sqrt dt
    mu  = x0 + (al + be*x0)*dt
    sd  = si * (x0 ** xi) * sdt
    c   = sd `seq` (-0.5 * log (2*pi*sd^2))

main = putStrLn . show $ cklsLogDens SV.empty (Timedelta 0.1) 0.1 0.15

Corresponding Core snippet:

Main.cklsLogDens [InlPrag=NOINLINE]
  :: Data.Vector.Storable.Vector GHC.Types.Double
     -> Main.Timedelta
     -> GHC.Types.Double
     -> GHC.Types.Double
     -> GHC.Types.Double
[GblId, Arity=4, Caf=NoCafRefs, Str=DmdType U(ALL)LLL]
Main.cklsLogDens =
  \ (p_atw :: Data.Vector.Storable.Vector GHC.Types.Double)
    (ds_dVa :: Main.Timedelta)
    (x0_aty :: GHC.Types.Double)
    (x1_atz :: GHC.Types.Double) ->
    case p_atw
    of _ { Data.Vector.Storable.Vector rb_a2ml rb1_a2mm rb2_a2mn ->
    case GHC.Prim.readDoubleOffAddr#
           @ GHC.Prim.RealWorld rb1_a2mm 2 GHC.Prim.realWorld#
    of _ { (# s2_a2nH, x_a2nI #) ->
    case GHC.Prim.touch#
           @ GHC.ForeignPtr.ForeignPtrContents rb2_a2mn s2_a2nH
    of _ { __DEFAULT ->
    case GHC.Prim.<=## x_a2nI 0.0 of _ {
      GHC.Types.False ->
        case x0_aty of _ { GHC.Types.D# x2_a13d ->
        case GHC.Prim.readDoubleOffAddr#
               @ GHC.Prim.RealWorld rb1_a2mm 3 GHC.Prim.realWorld#
        of _ { (# s1_X2oB, x3_X2oD #) ->
        case GHC.Prim.touch#
               @ GHC.ForeignPtr.ForeignPtrContents rb2_a2mn s1_X2oB
        of _ { __DEFAULT ->
        case ds_dVa
             `cast` (Main.NTCo:Timedelta :: Main.Timedelta ~# GHC.Types.Double)
        of _ { GHC.Types.D# x4_a13m ->
        let {
   --- ^^^^ want to get rid of this!
   ---
          ipv_sYP [Dmd=Just L] :: GHC.Prim.Double#
          [LclId, Str=DmdType]
          ipv_sYP =
            GHC.Prim.*##
              (GHC.Prim.*## x_a2nI (GHC.Prim.**## x2_a13d x3_X2oD))
              (GHC.Prim.sqrtDouble# x4_a13m) } in
        case x1_atz of _ { GHC.Types.D# x5_X14E ->
        case GHC.Prim.readDoubleOffAddr#
               @ GHC.Prim.RealWorld rb1_a2mm 0 GHC.Prim.realWorld#
        of _ { (# s3_X2p2, x6_X2p4 #) ->
        case GHC.Prim.touch#
               @ GHC.ForeignPtr.ForeignPtrContents rb2_a2mn s3_X2p2
        of _ { __DEFAULT ->
        case GHC.Prim.readDoubleOffAddr#
               @ GHC.Prim.RealWorld rb1_a2mm 1 GHC.Prim.realWorld#
        of _ { (# s4_X2pi, x7_X2pk #) ->
        case GHC.Prim.touch#
               @ GHC.ForeignPtr.ForeignPtrContents rb2_a2mn s4_X2pi
        of _ { __DEFAULT ->
        case GHC.Prim.logDouble#
               (GHC.Prim.*## 6.283185307179586 (GHC.Prim.*## ipv_sYP ipv_sYP))
        of wild9_a13D { __DEFAULT ->
        case GHC.Prim./##
               (GHC.Prim.-##
                  x5_X14E
                  (GHC.Prim.+##
                     x2_a13d
                     (GHC.Prim.*##
                        (GHC.Prim.+## x6_X2p4 (GHC.Prim.*## x7_X2pk x2_a13d)) x4_a13m)))
               ipv_sYP
        of wild10_a13O { __DEFAULT ->
        GHC.Types.D#
          (GHC.Prim.-##
             (GHC.Prim.negateDouble# (GHC.Prim.*## 0.5 wild9_a13D))
             (GHC.Prim.*## 0.5 (GHC.Prim.*## wild10_a13O wild10_a13O)))
        }
        }
        }
        }
        }
        }
        }
        }
        }
        }
        };
      GHC.Types.True -> lvl_r2v7
    }
    }
    }
    }

Answer 1

Daniel is right - the let in question does not, in fact, allocate a thunk. That would actually be impossible, as primitive types such as Double# don't have a heap representation. These lets are in fact converted to case expressions before conversion to STG (which is where the "let = allocation" rule actually holds) in the so-called core preparation phase. See the comment on this topic in CorePrep.lhs.

Here's again the Core from before preparation (-ddump-simpl):

    let {
      ipv_sPL [Dmd=Just L] :: GHC.Prim.Double#
      ipv_sPL =
        GHC.Prim.*##
          (GHC.Prim.*## x_a160 (GHC.Prim.**## x1_a11G x2_X17h))
          (GHC.Prim.sqrtDouble# x3_a11P) } in [...]

And here's after (-ddump-prep):

    case GHC.Prim.sqrtDouble# x3_s1aU of sat_s1cB { __DEFAULT ->
    case GHC.Prim.**## x1_s1aQ x2_s1aR of sat_s1cC { __DEFAULT ->
    case GHC.Prim.*## x_s1aC sat_s1cC of sat_s1cD { __DEFAULT ->
    case GHC.Prim.*## sat_s1cD sat_s1cB of ipv_s1aW [Dmd=Just L] { __DEFAULT ->

So there is actually no heap allocation whatsoever.

On the other hand, note that core preparation also explicitly wraps every single application into a let or case statement, producing pretty verbose code. That's why -ddump-simpl is probably considered the default for looking at Core, even though its performance model is actually slightly more surprising.