content: -- | This module implements a compiler pass for inlining functions,
  -- then removing those that have become dead.
  module Futhark.Optimise.InliningDeadFun
    ( inlineAggressively,
      inlineConservatively,
      removeDeadFunctions,
    )
  where
  
  import Control.Monad
  import Control.Monad.Identity
  import Control.Monad.State
  import Control.Parallel.Strategies
  import Data.Functor (($>))
  import Data.List (partition)
  import Data.Map.Strict qualified as M
  import Data.Maybe
  import Data.Set qualified as S
  import Futhark.Analysis.CallGraph
  import Futhark.Analysis.SymbolTable qualified as ST
  import Futhark.Builder
  import Futhark.IR.SOACS
  import Futhark.IR.SOACS.Simplify
    ( simpleSOACS,
      simplifyConsts,
      simplifyFun,
    )
  import Futhark.Optimise.CSE
  import Futhark.Optimise.Simplify.Rep (addScopeWisdom, informStms)
  import Futhark.Pass
  import Futhark.Transform.CopyPropagate
    ( copyPropagateInFun,
      copyPropagateInProg,
    )
  import Futhark.Transform.Rename
  
  parMapM :: (MonadFreshNames m) => (a -> State VNameSource b) -> [a] -> m [b]
  -- The special-casing of [] is quite important here!  If 'as' is
  -- empty, then we might otherwise create an empty name source below,
  -- which can wreak all kinds of havoc.
  parMapM _ [] = pure []
  parMapM f as =
    modifyNameSource $ \src ->
      let f' a = runState (f a) src
          (bs, srcs) = unzip $ parMap rpar f' as
       in (bs, mconcat srcs)
  
  -- It is more efficient to shrink the program as soon as possible,
  -- rather than wait until it has balooned after full inlining.  This
  -- is the inverse rate at which we perform full simplification after
  -- inlining.  For the other steps we just do copy propagation.  The
  -- simplification rates used have been determined heuristically and
  -- are probably not optimal for any given program.
  inlineFunctions ::
    (MonadFreshNames m) =>
    Int ->
    CallGraph ->
    S.Set Name ->
    Prog SOACS ->
    m (Prog SOACS)
  inlineFunctions simplify_rate cg what_should_be_inlined prog = do
    let consts = progConsts prog
        funs = progFuns prog
        vtable = ST.fromScope (addScopeWisdom (scopeOf consts))
  
    (consts', funs') <- recurse (1, vtable) (consts, funs) what_should_be_inlined
    pure $ prog {progConsts = consts', progFuns = funs'}
    where
      fdmap fds = M.fromList $ zip (map funDefName fds) fds
  
      noCallsTo which from = S.null $ allCalledBy from cg `S.intersection` which
  
      recurse (i, vtable) (consts, funs) to_inline = do
        let (to_inline_now, to_inline_later) =
              S.partition (noCallsTo to_inline) to_inline
            (dont_inline_in, to_inline_in) =
              partition (noCallsTo to_inline_now . funDefName) funs
  
        if null to_inline_now
          then pure (consts, funs)
          else do
            let inlinemap =
                  fdmap $ filter ((`S.member` to_inline_now) . funDefName) dont_inline_in
            (vtable', consts') <-
              if any (`calledByConsts` cg) to_inline_now
                then do
                  consts' <-
                    simplifyConsts . performCSEOnStms
                      =<< inlineInStms inlinemap consts
                  pure (ST.insertStms (informStms consts') mempty, consts')
                else pure (vtable, consts)
  
            let simplifyFun' fd
                  | i `rem` simplify_rate == 0 =
                      copyPropagateInFun simpleSOACS vtable'
                        . performCSEOnFunDef True
                        =<< simplifyFun vtable' fd
                  | otherwise =
                      copyPropagateInFun simpleSOACS vtable' fd
  
                onFun = simplifyFun' <=< inlineInFunDef inlinemap
  
            to_inline_in' <- parMapM onFun to_inline_in
  
            recurse
              (i + 1, vtable')
              (consts', dont_inline_in <> to_inline_in')
              to_inline_later
  
  calledOnce :: CallGraph -> S.Set Name
  calledOnce =
    S.fromList . map fst . filter ((== 1) . snd) . M.toList . numOccurences
  
  inlineBecauseTiny :: Prog SOACS -> S.Set Name
  inlineBecauseTiny = foldMap onFunDef . progFuns
    where
      onFunDef fd
        | (length (bodyStms (funDefBody fd)) <= k)
            || ("inline" `inAttrs` funDefAttrs fd) =
            S.singleton (funDefName fd)
        | otherwise = mempty
        where
          k = length (funDefRetType fd) + length (funDefParams fd)
  
  progStms :: Prog SOACS -> Stms SOACS
  progStms prog =
    progConsts prog <> foldMap (bodyStms . funDefBody) (progFuns prog)
  
  data Used = InSOAC | InAD deriving (Eq, Ord, Show)
  
  directlyCalledInSOACs :: Prog SOACS -> M.Map Name Used
  directlyCalledInSOACs = flip execState mempty . mapM_ (onStm Nothing) . progStms
    where
      onBody :: Maybe Used -> Body SOACS -> State (M.Map Name Used) ()
      onBody u = mapM_ (onStm u) . bodyStms
      onStm u stm = onExp u (stmExp stm) $> stm
      onExp (Just u) (Apply fname _ _ _) = modify $ M.insertWith max fname u
      onExp Nothing Apply {} = pure ()
      onExp u e = walkExpM (walker u) e
      onSOAC u soac = void $ traverseSOACStms (const (traverse (onStm u'))) soac
        where
          u' = max u $ Just $ usage soac
      usage JVP {} = InAD
      usage VJP {} = InAD
      usage _ = InSOAC
      walker u =
        (identityWalker :: Walker SOACS (State (M.Map Name Used)))
          { walkOnBody = const (onBody u),
            walkOnOp = onSOAC u
          }
  
  -- Expand set of function names with all reachable functions.
  withTransitiveCalls :: CallGraph -> M.Map Name Used -> M.Map Name Used
  withTransitiveCalls cg fs
    | fs == fs' = fs
    | otherwise = withTransitiveCalls cg fs'
    where
      look :: (Name, Used) -> M.Map Name Used
      look (f, u) = M.fromList $ map (,u) (S.toList (allCalledBy f cg))
      fs' = foldr (M.unionWith max . look) fs $ M.toList fs
  
  calledInSOACs :: CallGraph -> Prog SOACS -> M.Map Name Used
  calledInSOACs cg prog = withTransitiveCalls cg $ directlyCalledInSOACs prog
  
  -- Inline those functions that are used in SOACs, and which involve
  -- arrays of any kind, as well as any functions used in AD.
  inlineBecauseSOACs :: CallGraph -> Prog SOACS -> S.Set Name
  inlineBecauseSOACs cg prog =
    S.fromList $ mapMaybe onFunDef (progFuns prog)
    where
      called = calledInSOACs cg prog
      isArray = not . primType
      inline _ InAD = True
      inline fd InSOAC =
        any (isArray . paramType) (funDefParams fd)
          || any (isArray . fst) (funDefRetType fd)
          || arrayInBody (funDefBody fd)
      onFunDef fd = do
        guard $ maybe False (inline fd) $ M.lookup (funDefName fd) called
        Just $ funDefName fd
      arrayInBody = any arrayInStm . bodyStms
      arrayInStm stm =
        any isArray (patTypes (stmPat stm)) || arrayInExp (stmExp stm)
      arrayInExp (Match _ cases defbody _) =
        any arrayInBody $ defbody : map caseBody cases
      arrayInExp (Loop _ _ body) =
        arrayInBody body
      arrayInExp _ = False
  
  -- Conservative inlining of functions that are called just once, or
  -- have #[inline] on them.
  consInlineFunctions :: (MonadFreshNames m) => Prog SOACS -> m (Prog SOACS)
  consInlineFunctions prog =
    inlineFunctions 4 cg (calledOnce cg <> inlineBecauseTiny prog) prog
    where
      cg = buildCallGraph prog
  
  -- Inline aggressively; in particular most things called from a SOAC.
  aggInlineFunctions :: (MonadFreshNames m) => Prog SOACS -> m (Prog SOACS)
  aggInlineFunctions prog =
    inlineFunctions 3 cg (inlineBecauseTiny prog <> inlineBecauseSOACs cg prog) prog
    where
      cg = buildCallGraph prog
  
  -- | @inlineInFunDef constf fdmap caller@ inlines in @calleer@ the
  -- functions in @fdmap@ that are called as @constf@. At this point the
  -- preconditions are that if @fdmap@ is not empty, and, more
  -- importantly, the functions in @fdmap@ do not call any other
  -- functions.
  inlineInFunDef ::
    (MonadFreshNames m) =>
    M.Map Name (FunDef SOACS) ->
    FunDef SOACS ->
    m (FunDef SOACS)
  inlineInFunDef fdmap (FunDef entry attrs name rtp args body) =
    FunDef entry attrs name rtp args <$> inlineInBody fdmap body
  
  inlineFunction ::
    (MonadFreshNames m) =>
    Pat Type ->
    StmAux dec ->
    [(SubExp, Diet)] ->
    Safety ->
    Provenance ->
    FunDef SOACS ->
    m (Stms SOACS)
  inlineFunction pat aux args safety p fun = do
    Body _ stms res <-
      renameBody $ mkBody (param_stms <> body_stms) (bodyResult (funDefBody fun))
    pure $ stms <> stmsFromList (zipWith bindSubExpRes (patIdents pat) res)
    where
      param_stms =
        stmsFromList $
          certify (stmAuxCerts aux)
            <$> zipWith bindSubExp (map paramIdent $ funDefParams fun) (map fst args)
  
      body_stms =
        addLocations (stmAuxAttrs aux) safety p $ bodyStms $ funDefBody fun
  
      -- Note that the sizes of arrays may not be correct at this
      -- point - it is crucial that we run copy propagation before
      -- the type checker sees this!
      bindSubExp ident se =
        mkLet [ident] $ BasicOp $ SubExp se
  
      bindSubExpRes ident (SubExpRes cs se) =
        certify cs $ bindSubExp ident se
  
  inlineInStms ::
    (MonadFreshNames m) =>
    M.Map Name (FunDef SOACS) ->
    Stms SOACS ->
    m (Stms SOACS)
  inlineInStms fdmap stms =
    bodyStms <$> inlineInBody fdmap (mkBody stms [])
  
  inlineInBody ::
    (MonadFreshNames m) =>
    M.Map Name (FunDef SOACS) ->
    Body SOACS ->
    m (Body SOACS)
  inlineInBody fdmap = onBody
    where
      inline (Let pat aux (Apply fname args _ safety) : rest)
        | Just fd <- M.lookup fname fdmap,
          not $ "noinline" `inAttrs` funDefAttrs fd,
          not $ "noinline" `inAttrs` stmAuxAttrs aux =
            (<>) <$> inlineFunction pat aux args safety (stmAuxLoc aux) fd <*> inline rest
      inline (stm@(Let _ _ BasicOp {}) : rest) =
        (oneStm stm <>) <$> inline rest
      inline (stm : rest) =
        (<>) <$> (oneStm <$> onStm stm) <*> inline rest
      inline [] =
        pure mempty
  
      onBody (Body dec stms res) =
        Body dec <$> inline (stmsToList stms) <*> pure res
      onStm (Let pat aux e) = Let pat aux <$> mapExpM inliner e
  
      inliner =
        (identityMapper @SOACS)
          { mapOnBody = const onBody,
            mapOnOp = onSOAC
          }
  
      onSOAC =
        mapSOACM identitySOACMapper {mapOnSOACLambda = onLambda}
  
      onLambda (Lambda params ret body) =
        Lambda params ret <$> onBody body
  
  traceLocs :: Provenance -> StmAux () -> StmAux ()
  traceLocs p aux =
    aux {stmAuxLoc = stackProvenance p $ stmAuxLoc aux}
  
  -- Propagate source locations and attributes to the inlined statements.
  -- Attributes are propagated only when applicable (this probably means that
  -- every supported attribute needs to be handled specially here).
  addLocations :: Attrs -> Safety -> Provenance -> Stms SOACS -> Stms SOACS
  addLocations attrs caller_safety p = fmap onStm
    where
      onStm (Let pat aux (BasicOp (Assert cond desc))) =
        Let pat aux' $
          case caller_safety of
            Safe -> BasicOp $ Assert cond desc
            Unsafe -> BasicOp $ SubExp $ Constant UnitValue
        where
          aux' = traceLocs p (withAttrs (attrsForAssert attrs) aux)
      onStm (Let pat aux (Apply fname args t safety)) =
        Let pat aux' $ Apply fname args t $ min caller_safety safety
        where
          aux' = traceLocs p $ aux {stmAuxAttrs = attrs <> stmAuxAttrs aux}
      onStm (Let pat aux (Op soac)) =
        Let pat aux' . Op $
          runIdentity $
            mapSOACM
              identitySOACMapper
                { mapOnSOACLambda = pure . onLambda
                }
              soac
        where
          aux' = traceLocs p $ withAttrs attrs' aux
          attrs' = attrs `withoutAttrs` for_assert
          for_assert = attrsForAssert attrs
          onLambda lam =
            lam {lambdaBody = onBody for_assert $ lambdaBody lam}
      onStm (Let pat aux e) =
        Let pat (traceLocs p aux) $ onExp e
  
      onExp =
        mapExp
          identityMapper
            { mapOnBody = const $ pure . onBody attrs
            }
  
      withAttrs attrs' aux = aux {stmAuxAttrs = attrs' <> stmAuxAttrs aux}
  
      onBody attrs' body =
        body
          { bodyStms = addLocations attrs' caller_safety p $ bodyStms body
          }
  
  -- | Remove functions not ultimately called from an entry point or a
  -- constant.
  removeDeadFunctionsF :: Prog SOACS -> Prog SOACS
  removeDeadFunctionsF prog =
    let cg = buildCallGraph prog
        live_funs = filter ((`isFunInCallGraph` cg) . funDefName) $ progFuns prog
     in prog {progFuns = live_funs}
  
  -- | Inline all functions and remove the resulting dead functions.
  inlineAggressively :: Pass SOACS SOACS
  inlineAggressively =
    Pass
      { passName = "Inline aggressively",
        passDescription = "Aggressively inline and remove resulting dead functions.",
        passFunction =
          copyPropagateInProg simpleSOACS . removeDeadFunctionsF <=< aggInlineFunctions
      }
  
  -- | Inline some functions and remove the resulting dead functions.
  inlineConservatively :: Pass SOACS SOACS
  inlineConservatively =
    Pass
      { passName = "Inline conservatively",
        passDescription = "Conservatively inline and remove resulting dead functions.",
        passFunction =
          copyPropagateInProg simpleSOACS . removeDeadFunctionsF <=< consInlineFunctions
      }
  
  -- | @removeDeadFunctions prog@ removes the functions that are unreachable from
  -- the main function from the program.
  removeDeadFunctions :: Pass SOACS SOACS
  removeDeadFunctions =
    Pass
      { passName = "Remove dead functions",
        passDescription = "Remove the functions that are unreachable from entry points",
        passFunction = pure . removeDeadFunctionsF
      }
  