-- Copyright 2020 Google LLC
--
-- Use of this source code is governed by a BSD-style
-- license that can be found in the LICENSE file or at
-- https://developers.google.com/open-source/licenses/bsd

module Parallelize (parallelizeModule, dceModule) where

import Control.Monad.State.Strict
import Control.Monad.Reader
import Data.Foldable

import Optimize
import Syntax
import Builder
import Cat
import Subst
import Type
import Util (for)

-- TODO: extractParallelism can benefit a lot from horizontal fusion (can happen be after)
-- TODO: Parallelism extraction can emit some really cheap (but not trivial)
--       fors and we should inline them.
parallelizeModule :: Module -> Module
parallelizeModule = dceModule . extractParallelism

-- === Flattening parallelism ===

-- It is more convenient to treat the block result as "cheap" and treat
-- all potentially expensive operations as appearing in decls for uniformity.
data ABlock = ABlock (Nest Decl) Atom

asABlock :: Block -> ABlock
asABlock block = ABlock decls result
  where
    scope = freeVars block
    ((result, decls), _) = runBuilder scope mempty $ scopedDecls $ emitBlock block

data LoopSubst = LoopEnv
  { loopBinders :: [Var]              -- In scope of the original program
  , delayedApps :: Subst (Atom, [Atom]) -- (n @> (arr, bs)), n and bs in scope of the original program
                                      --   arr in scope of the newly constructed program!
  }
data AccSubst = AccEnv { activeAccs :: Env (Var, (Var, BaseMonoid)) } -- (original region, reference, its base monoid)

type TLParallelM = SubstBuilderT (State AccSubst)   -- Top-level non-parallel statements
type LoopM       = ReaderT LoopSubst TLParallelM  -- Generation of (parallel) loop nests

runLoopM :: LoopM a -> TLParallelM a
runLoopM m = runReaderT m $ LoopSubst mempty mempty

extractParallelism :: Module -> Module
extractParallelism = transformModuleAsBlock go
  where go block = fst $ evalState (runSubstBuilderT mempty mempty
                     (traverseBlock parallelTrav block)) $ AccSubst mempty

parallelTrav :: TraversalDef TLParallelM
parallelTrav = ( traverseDecl parallelTrav
               , parallelTraverseExpr
               , traverseAtom parallelTrav
               )

parallelTraverseExpr :: Expr -> TLParallelM Expr
parallelTraverseExpr expr = case expr of
  Hof (For ParallelFor _) -> traverseExpr substTraversalDef expr
  Hof (For (RegularFor _) fbody@(LamVal b body)) -> do
    -- TODO: functionEffs is an overapproximation of the effects that really appear inside
    refs <- gets activeAccs
    let allowedRegions = foldMap (\(varType . fst . snd -> RefTy (Var reg) _) -> reg @> ()) refs
    (EffectRow bodyEffs t) <- substBuilderR $ functionEffs fbody
    let onlyAllowedEffects = all (parallelizableEffect allowedRegions) $ toList bodyEffs
    case t == Nothing && onlyAllowedEffects of
      True -> do
        b' <- substBuilderR b
        liftM Atom $ runLoopM $ withLoopBinder b' $ buildParallelBlock $ asABlock body
      False -> nothingSpecial
  Hof (RunWriter bm (BinaryFunVal h@(Bind hName) b _ body)) -> do
    ~(RefTy _ accTy) <- traverseAtom substTraversalDef $ binderType b
    liftM Atom $ emitRunWriter (binderNameHint b) accTy bm \ref@(Var refVar) -> do
      let RefTy h' _ = varType refVar
      modify \accSubst -> accEnv { activeAccs = activeAccs accEnv <> b @> (hName, (refVar, bm)) }
      res <- extendR (h @> SubstVal h' <> b @> SubstVal ref) $ evalBlockE parallelTrav body
      modify \accSubst -> accEnv { activeAccs = activeAccs accEnv `envDiff` (b @> ()) }
      return res
  -- TODO: Do some alias analysis. This is not fundamentally hard, but it is a little annoying.
  --       We would have to track not only the base references, but also all the aliases, along
  --       with their relationships. Then, when we emit local effects in emitLoops, we would have
  --       to recreate all the aliases. We could do that by carrying around a block and using
  --       SubstBuilder to take care of renaming for us.
  Op (IndexRef ref _) -> disallowRef ref >> nothingSpecial
  Op (ProjRef  _ ref) -> disallowRef ref >> nothingSpecial
  _ -> nothingSpecial
  where
    nothingSpecial = traverseExpr parallelTrav expr
    disallowRef ~(Var refVar) =
      modify \accSubst -> accEnv { activeAccs = activeAccs accEnv `envDiff` (refVar @> ()) }

parallelizableEffect :: Subst () -> Effect -> Bool
parallelizableEffect allowedRegions effect = case effect of
  RWSEffect Writer h | h `isin` allowedRegions -> True
  -- TODO: we should be able to parallelize the exception effect too
  _ -> False

-- Precondition: This is never called with no binders in the loop env
buildParallelBlock :: ABlock -> LoopM Atom
buildParallelBlock ablock@(ABlock decls result) = do
  decision <- analyzeDecls decls
  lbs <- asks loopBinders
  let loopVars = fmap Var lbs
  case decision of
    -- The only thing an empty ABlock can do is read indices and some local variables
    -- that have been gathered into arrays. Those things are really cheap though, and
    -- emitting a for in that case would only gather them into another table.
    Emit -> case decls of
      Empty -> reduce =<< unflattenIndexBundle lbs =<< emitLoops (flip buildLam TabArrow) ablock
        where reduce = lift . traverseAtom appReduceTraversalDef
      _     -> unflattenIndexBundle lbs =<< emitLoops (buildForAnn ParallelFor) ablock
    Split prologue (arrb, loop@(Abs i lbody)) epilogue -> do
      prologueApps <- case prologue of
            Empty -> return mempty
            _ -> do
              -- TODO: This can break miserably with dependent values!
              let freeContVars = freeVars loop <> freeVars epilogue
              let prologueCtxVars = bindingsAsVars $ boundVars prologue `envIntersect` freeContVars
              let (prologueResult, prologueBundleDesc) = mkBundle $ fmap Var prologueCtxVars
              let prologueBlock = ABlock prologue prologueResult
              prologueCtxAtom <- emitLoops (buildForAnn ParallelFor) prologueBlock
              prologueCtxArrs <- mapM (unflattenIndexBundle lbs) =<<
                                    unpackBundleTab prologueCtxAtom prologueBundleDesc
              return $ foldMap (\(v, arr) -> v @> (arr, loopVars)) $ zip prologueCtxVars prologueCtxArrs
      delayApps prologueApps $ do
        i' <- lift $ substBuilderR i
        loopAtom <- withLoopBinder i' $ buildParallelBlock $ asABlock lbody
        delayApps (arrb @> (loopAtom, loopVars)) $
          buildParallelBlock $ ABlock epilogue result

unflattenIndexBundle :: MonadBuilder m => [Var] -> Atom -> m Atom
unflattenIndexBundle []  arr = return arr
unflattenIndexBundle [_] arr = return arr
unflattenIndexBundle ivs arr = buildNaryLam TabArrow (toNest $ fmap Bind ivs) $ app arr . fst . mkBundle

type Loop = Abs Binder Block
data NestDecision = Emit | Split (Nest Decl) (Binder, Loop) (Nest Decl)

pattern LoopHof :: Binder -> Block -> Expr
pattern LoopHof i body <- Hof (For _ (Lam (Abs i (PureArrow, body))))

-- TODO: Implement something less aggressive, maybe choose between different tactics?
--       After all, we might have some interesting shape information available.
--       We might also want to limit ourselves to flattening only perfectly nested loops.
analyzeDecls :: Nest Decl -> LoopM NestDecision
analyzeDecls declsNest = splitAggressively
  where
    splitAggressively = do
      let decls = toList declsNest
      let (prologue, rest) = break (\case (Let _ _ (LoopHof _ _)) -> True; _ -> False) decls
      -- TODO: Don't split if the domain is known to be small (e.g. a 3D vector)
      -- TODO: If the prologue is cheap (and probably isn't used in the epilogue), then
      --       we could decide to inline it inside the loop and replicate the compute!
      --       For example, this will unnecessarily get split into three loops:
      --       for i.
      --         xi = x i
      --         for j.
      --           xij = xi j
      --           for z.
      --             xij z
      -- TODO: Check that the loop binder is in scope given only the other binders
      --       (i.e. it's not dependent on a loop-local variable.
      --       BTW we might be able to handle the loop-local cases too, right?
      --       Because the array of locals will be emitted anyway.
      --       We might just get a length-indexed array (not that we support them yet...)!
      return $ case rest of
        [] -> Emit
        ~(Let _ b (LoopHof i body):epilogue) -> Split (toNest prologue) (b, Abs i body) (toNest epilogue)

withLoopBinder :: Binder -> LoopM a -> LoopM a
withLoopBinder b m = do
  v <- case b of
    Bind v    -> return v
    Ignore ty -> withNameHint ("i"::String) $ freshVarE UnknownBinder ty
  local (\env -> env { loopBinders = loopBinders env <> [v]}) m

delayApps :: Subst (Atom, [Atom]) -> LoopM a -> LoopM a
delayApps apps = local (\env -> env { delayedApps = apps <> delayedApps env })

emitLoops :: (Binder -> (Atom -> TLParallelM Atom) -> TLParallelM Atom)
          -> ABlock -> LoopM Atom
emitLoops buildPureLoop (ABlock decls result) = do
  -- TODO: Deal with dependent binders properly
  lbs <- asks loopBinders
  dapps <- asks delayedApps
  refs  <- gets activeAccs
  -- TODO: Filter out the refs that are unused in the body.
  let oldRefNames = envNames refs
  let (oldRegionNames, newRefs) = unzip $ toList refs
  let (iterTy, iterBundleDesc) = mkBundleTy $ fmap varType lbs
  let buildBody pari = do
        is <- unpackBundle pari iterBundleDesc
        extendR (newSubst lbs $ map SubstVal is) $ do
          ctxSubst <- flip traverseNames dapps \_ (arr, idx) ->
            -- XXX: arr is namespaced in the new program
            SubstVal <$> (foldM appTryReduce arr =<< substBuilderR idx)
          extendR ctxSubst $ evalBlockE appReduceTraversalDef $ Block decls $ Atom result
  lift $ case null refs of
    True -> buildPureLoop (Bind $ "pari" :> iterTy) buildBody
    False -> do
      body <- do
        buildLam (Bind $ "gtid" :> IdxRepTy) PureArrow \gtid -> do
          buildLam (Bind $ "nthr" :> IdxRepTy) PureArrow \nthr -> do
            let threadRange = TC $ ParIndexRange iterTy gtid nthr
            let writerSpecs = for newRefs \(ref, bm) -> (varName ref, derefType (varType ref), bm)
            emitRunWriters writerSpecs $ \localRefs -> do
              buildFor Fwd (Bind $ "tidx" :> threadRange) \tidx -> do
                pari <- emitOp $ Sink tidx
                let regionSubst = newEnv oldRegionNames $ for localRefs \(getType -> RefTy h _) -> SubstVal h
                extendR (regionSubst <> newEnv oldRefNames (map SubstVal localRefs)) $ buildBody pari
      (ans, updateList) <- fromPair =<< (emit $ Hof $ PTileReduce (fmap snd newRefs) iterTy body)
      updates <- unpackRightLeaningConsList updateList
      forM_ (zip newRefs updates) $ \((ref, bm), update) -> do
        updater <- mextendForRef (Var ref) bm update
        emitOp $ PrimEffect (Var ref) $ MExtend updater
      return ans
  where
    derefType ~(RefTy _ accTy) = accTy