content: {-# LANGUAGE TypeFamilies #-}
  {-# LANGUAGE UndecidableInstances #-}
  {-# OPTIONS_GHC -fno-warn-redundant-constraints #-}
  
  -- | A generic transformation for adding memory allocations to a
  -- Futhark program.  Specialised by specific representations in
  -- submodules.
  module Futhark.Pass.ExplicitAllocations
    ( explicitAllocationsGeneric,
      explicitAllocationsInStmsGeneric,
      ExpHint (..),
      defaultExpHints,
      askDefaultSpace,
      Allocable,
      AllocM,
      AllocEnv (..),
      SizeSubst (..),
      allocInStms,
      allocForArray,
      simplifiable,
      mkLetNamesB',
      mkLetNamesB'',
  
      -- * Module re-exports
  
      --
      -- These are highly likely to be needed by any downstream
      -- users.
      module Control.Monad.Reader,
      module Futhark.MonadFreshNames,
      module Futhark.Pass,
      module Futhark.Tools,
    )
  where
  
  import Control.Monad
  import Control.Monad.RWS.Strict
  import Control.Monad.Reader
  import Control.Monad.State
  import Control.Monad.Writer
  import Data.Bifunctor (first)
  import Data.Either (partitionEithers)
  import Data.List (foldl', transpose, zip4)
  import Data.Map.Strict qualified as M
  import Data.Maybe
  import Data.Set qualified as S
  import Futhark.Analysis.SymbolTable (IndexOp)
  import Futhark.Analysis.UsageTable qualified as UT
  import Futhark.IR.Mem
  import Futhark.IR.Mem.LMAD qualified as LMAD
  import Futhark.IR.Prop.Aliases (AliasedOp)
  import Futhark.MonadFreshNames
  import Futhark.Optimise.Simplify.Engine (SimpleOps (..))
  import Futhark.Optimise.Simplify.Engine qualified as Engine
  import Futhark.Optimise.Simplify.Rep (mkWiseBody)
  import Futhark.Pass
  import Futhark.Tools
  import Futhark.Util (maybeNth, splitAt3)
  
  type Allocable fromrep torep inner =
    ( PrettyRep fromrep,
      PrettyRep torep,
      Mem torep inner,
      LetDec torep ~ LetDecMem,
      FParamInfo fromrep ~ DeclType,
      LParamInfo fromrep ~ Type,
      BranchType fromrep ~ ExtType,
      RetType fromrep ~ DeclExtType,
      BodyDec fromrep ~ (),
      BodyDec torep ~ (),
      ExpDec torep ~ (),
      SizeSubst (inner torep),
      BuilderOps torep
    )
  
  data AllocEnv fromrep torep = AllocEnv
    { -- | When allocating memory, put it in this memory space.
      -- This is primarily used to ensure that group-wide
      -- statements store their results in shared memory.
      allocSpace :: Space,
      -- | The set of names that are known to be constants at
      -- kernel compile time.
      envConsts :: S.Set VName,
      allocInOp :: Op fromrep -> AllocM fromrep torep (Op torep),
      envExpHints :: Exp torep -> AllocM fromrep torep [ExpHint]
    }
  
  -- | Monad for adding allocations to an entire program.
  newtype AllocM fromrep torep a
    = AllocM (BuilderT torep (ReaderT (AllocEnv fromrep torep) (State VNameSource)) a)
    deriving
      ( Applicative,
        Functor,
        Monad,
        MonadFreshNames,
        HasScope torep,
        LocalScope torep,
        MonadReader (AllocEnv fromrep torep)
      )
  
  instance (Allocable fromrep torep inner) => MonadBuilder (AllocM fromrep torep) where
    type Rep (AllocM fromrep torep) = torep
  
    mkExpDecM _ _ = pure ()
  
    mkLetNamesM names e = do
      def_space <- askDefaultSpace
      hints <- expHints e
      pat <- patWithAllocations def_space names e hints
      pure $ Let pat (defAux ()) e
  
    mkBodyM stms res = pure $ Body () stms res
  
    addStms = AllocM . addStms
    collectStms (AllocM m) = AllocM $ collectStms m
  
  expHints :: Exp torep -> AllocM fromrep torep [ExpHint]
  expHints e = do
    f <- asks envExpHints
    f e
  
  -- | The space in which we allocate memory if we have no other
  -- preferences or constraints.
  askDefaultSpace :: AllocM fromrep torep Space
  askDefaultSpace = asks allocSpace
  
  runAllocM ::
    (MonadFreshNames m) =>
    Space ->
    (Op fromrep -> AllocM fromrep torep (Op torep)) ->
    (Exp torep -> AllocM fromrep torep [ExpHint]) ->
    AllocM fromrep torep a ->
    m a
  runAllocM space handleOp hints (AllocM m) =
    fmap fst $ modifyNameSource $ runState $ runReaderT (runBuilderT m mempty) env
    where
      env =
        AllocEnv
          { allocSpace = space,
            envConsts = mempty,
            allocInOp = handleOp,
            envExpHints = hints
          }
  
  elemSize :: (Num a) => Type -> a
  elemSize = primByteSize . elemType
  
  arraySizeInBytesExp :: Type -> PrimExp VName
  arraySizeInBytesExp t =
    untyped $ foldl' (*) (elemSize t) $ map pe64 (arrayDims t)
  
  arraySizeInBytes :: (MonadBuilder m) => Type -> m SubExp
  arraySizeInBytes = letSubExp "bytes" <=< toExp . arraySizeInBytesExp
  
  allocForArray' ::
    (MonadBuilder m, Op (Rep m) ~ MemOp inner (Rep m)) =>
    Type ->
    Space ->
    m VName
  allocForArray' t space = do
    size <- arraySizeInBytes t
    letExp "mem" $ Op $ Alloc size space
  
  -- | Allocate memory for a value of the given type.
  allocForArray ::
    (Allocable fromrep torep inner) =>
    Type ->
    Space ->
    AllocM fromrep torep VName
  allocForArray t space = do
    allocForArray' t space
  
  -- | Repair an expression that cannot be assigned an index function.
  -- There is a simple remedy for this: normalise the input arrays and
  -- try again.
  repairExpression ::
    (Allocable fromrep torep inner) =>
    Exp torep ->
    AllocM fromrep torep (Exp torep)
  repairExpression (BasicOp (Reshape v shape)) = do
    v_mem <- fst <$> lookupArraySummary v
    space <- lookupMemSpace v_mem
    v' <- snd <$> ensureDirectArray (Just space) v
    pure $ BasicOp $ Reshape v' shape
  repairExpression e =
    error $ "repairExpression:\n" <> prettyString e
  
  expReturns' ::
    (Allocable fromrep torep inner) =>
    Exp torep ->
    AllocM fromrep torep ([ExpReturns], Exp torep)
  expReturns' e = do
    maybe_rts <- expReturns e
    case maybe_rts of
      Just rts -> pure (rts, e)
      Nothing -> do
        e' <- repairExpression e
        let bad =
              error . unlines $
                [ "expReturns': impossible index transformation",
                  prettyString e,
                  prettyString e'
                ]
        rts <- fromMaybe bad <$> expReturns e'
        pure (rts, e')
  
  allocsForStm ::
    (Allocable fromrep torep inner) =>
    [Ident] ->
    Exp torep ->
    AllocM fromrep torep (Stm torep)
  allocsForStm idents e = do
    def_space <- askDefaultSpace
    hints <- expHints e
    (rts, e') <- expReturns' e
    pes <- allocsForPat def_space idents rts hints
    dec <- mkExpDecM (Pat pes) e'
    pure $ Let (Pat pes) (defAux dec) e'
  
  patWithAllocations ::
    (MonadBuilder m, Mem (Rep m) inner) =>
    Space ->
    [VName] ->
    Exp (Rep m) ->
    [ExpHint] ->
    m (Pat LetDecMem)
  patWithAllocations def_space names e hints = do
    ts' <- instantiateShapes' names <$> expExtType e
    let idents = zipWith Ident names ts'
    rts <- fromMaybe (error "patWithAllocations: ill-typed") <$> expReturns e
    Pat <$> allocsForPat def_space idents rts hints
  
  mkMissingIdents :: (MonadFreshNames m) => [Ident] -> [ExpReturns] -> m [Ident]
  mkMissingIdents idents rts =
    reverse <$> zipWithM f (reverse rts) (map Just (reverse idents) ++ repeat Nothing)
    where
      f _ (Just ident) = pure ident
      f (MemMem space) Nothing = newIdent "ext_mem" $ Mem space
      f _ Nothing = newIdent "ext" $ Prim int64
  
  allocsForPat ::
    (MonadBuilder m, Op (Rep m) ~ MemOp inner (Rep m)) =>
    Space ->
    [Ident] ->
    [ExpReturns] ->
    [ExpHint] ->
    m [PatElem LetDecMem]
  allocsForPat def_space some_idents rts hints = do
    idents <- mkMissingIdents some_idents rts
  
    forM (zip3 idents rts hints) $ \(ident, rt, hint) -> do
      let ident_shape = arrayShape $ identType ident
      case rt of
        MemPrim _ -> do
          summary <- summaryForBindage def_space (identType ident) hint
          pure $ PatElem (identName ident) summary
        MemMem space ->
          pure $ PatElem (identName ident) $ MemMem space
        MemArray bt _ u (Just (ReturnsInBlock mem extlmad)) -> do
          let ixfn = instantiateExtLMAD idents extlmad
          pure . PatElem (identName ident) . MemArray bt ident_shape u $ ArrayIn mem ixfn
        MemArray _ extshape _ Nothing
          | Just _ <- knownShape extshape -> do
              summary <- summaryForBindage def_space (identType ident) hint
              pure $ PatElem (identName ident) summary
        MemArray bt _ u (Just (ReturnsNewBlock _ i extixfn)) -> do
          let ixfn = instantiateExtLMAD idents extixfn
          pure . PatElem (identName ident) . MemArray bt ident_shape u $
            ArrayIn (getIdent idents i) ixfn
        MemAcc acc ispace ts u ->
          pure $ PatElem (identName ident) $ MemAcc acc ispace ts u
        _ -> error "Impossible case reached in allocsForPat!"
    where
      knownShape = mapM known . shapeDims
      known (Free v) = Just v
      known Ext {} = Nothing
  
      getIdent idents i =
        case maybeNth i idents of
          Just ident -> identName ident
          Nothing ->
            error $ "getIdent: Ext " <> show i <> " but pattern has " <> show (length idents) <> " elements: " <> prettyString idents
  
      instantiateExtLMAD idents = fmap $ fmap inst
        where
          inst (Free v) = v
          inst (Ext i) = getIdent idents i
  
  instantiateLMAD :: (Monad m) => ExtLMAD -> m LMAD
  instantiateLMAD = traverse $ traverse inst
    where
      inst Ext {} = error "instantiateLMAD: not yet"
      inst (Free x) = pure x
  
  summaryForBindage ::
    (MonadBuilder m, Op (Rep m) ~ MemOp inner (Rep m)) =>
    Space ->
    Type ->
    ExpHint ->
    m (MemBound NoUniqueness)
  summaryForBindage _ (Prim bt) _ =
    pure $ MemPrim bt
  summaryForBindage _ (Mem space) _ =
    pure $ MemMem space
  summaryForBindage _ (Acc acc ispace ts u) _ =
    pure $ MemAcc acc ispace ts u
  summaryForBindage def_space t@(Array pt shape u) NoHint = do
    m <- allocForArray' t def_space
    pure $ MemArray pt shape u $ ArrayIn m $ LMAD.iota 0 $ map pe64 $ arrayDims t
  summaryForBindage _ t@(Array pt _ _) (Hint lmad space) = do
    bytes <-
      letSubExp "bytes" <=< toExp . untyped $
        primByteSize pt * (1 + LMAD.range lmad)
    m <- letExp "mem" $ Op $ Alloc bytes space
    pure $ MemArray pt (arrayShape t) NoUniqueness $ ArrayIn m lmad
  
  allocInFParams ::
    (Allocable fromrep torep inner) =>
    [(FParam fromrep, Space)] ->
    ([FParam torep] -> AllocM fromrep torep a) ->
    AllocM fromrep torep a
  allocInFParams params m = do
    (valparams, (memparams, ctxparams)) <-
      runWriterT $ mapM (uncurry allocInFParam) params
    let params' = memparams <> ctxparams <> valparams
        summary = scopeOfFParams params'
    localScope summary $ m params'
  
  allocInFParam ::
    (Allocable fromrep torep inner) =>
    FParam fromrep ->
    Space ->
    WriterT
      ([FParam torep], [FParam torep])
      (AllocM fromrep torep)
      (FParam torep)
  allocInFParam param pspace =
    case paramDeclType param of
      Array pt shape u -> do
        let memname = baseName (paramName param) <> "_mem"
            lmad = LMAD.iota 0 $ map pe64 $ shapeDims shape
        mem <- lift $ newVName memname
        tell ([Param (paramAttrs param) mem $ MemMem pspace], [])
        pure param {paramDec = MemArray pt shape u $ ArrayIn mem lmad}
      Prim pt ->
        pure param {paramDec = MemPrim pt}
      Mem space ->
        pure param {paramDec = MemMem space}
      Acc acc ispace ts u ->
        pure param {paramDec = MemAcc acc ispace ts u}
  
  ensureRowMajorArray ::
    (Allocable fromrep torep inner) =>
    Maybe Space ->
    VName ->
    AllocM fromrep torep (VName, VName)
  ensureRowMajorArray space_ok v = do
    (mem, _) <- lookupArraySummary v
    mem_space <- lookupMemSpace mem
    default_space <- askDefaultSpace
    let space = fromMaybe default_space space_ok
    if maybe True (== mem_space) space_ok
      then pure (mem, v)
      else allocLinearArray space (baseName v) v
  
  ensureArrayIn ::
    (Allocable fromrep torep inner) =>
    Space ->
    SubExp ->
    WriterT ([SubExp], [SubExp]) (AllocM fromrep torep) SubExp
  ensureArrayIn _ (Constant v) =
    error $ "ensureArrayIn: " ++ prettyString v ++ " cannot be an array."
  ensureArrayIn space (Var v) = do
    (mem', v') <- lift $ ensureRowMajorArray (Just space) v
    (_, lmad) <- lift $ lookupArraySummary v'
    ctx <- lift $ mapM (letSubExp "lmad_arg" <=< toExp) (LMAD.existentialized lmad)
    tell ([Var mem'], ctx)
    pure $ Var v'
  
  allocInLoopParams ::
    (Allocable fromrep torep inner) =>
    [(FParam fromrep, SubExp)] ->
    ( [(FParam torep, SubExp)] ->
      ([SubExp] -> AllocM fromrep torep ([SubExp], [SubExp])) ->
      AllocM fromrep torep a
    ) ->
    AllocM fromrep torep a
  allocInLoopParams merge m = do
    ((valparams, valargs, handle_loop_subexps), (mem_params, ctx_params)) <-
      runWriterT $ unzip3 <$> mapM allocInLoopParam merge
    let mergeparams' = mem_params <> ctx_params <> valparams
        summary = scopeOfFParams mergeparams'
  
        mk_loop_res ses = do
          (ses', (memargs, ctxargs)) <-
            runWriterT $ zipWithM ($) handle_loop_subexps ses
          pure (memargs <> ctxargs, ses')
  
    (valctx_args, valargs') <- mk_loop_res valargs
    let merge' =
          zip (mem_params <> ctx_params <> valparams) (valctx_args <> valargs')
    localScope summary $ m merge' mk_loop_res
    where
      param_names = namesFromList $ map (paramName . fst) merge
      anyIsLoopParam names = names `namesIntersect` param_names
  
      scalarRes param_t v_mem_space v_lmad (Var res) = do
        -- Try really hard to avoid copying needlessly, but the result
        -- _must_ be in ScalarSpace and have the right index function.
        (res_mem, res_lmad) <- lift $ lookupArraySummary res
        res_mem_space <- lift $ lookupMemSpace res_mem
        (res_mem', res') <-
          if (res_mem_space, res_lmad) == (v_mem_space, v_lmad)
            then pure (res_mem, res)
            else lift $ arrayWithLMAD v_mem_space v_lmad (fromDecl param_t) res
        tell ([Var res_mem'], [])
        pure $ Var res'
      scalarRes _ _ _ se = pure se
  
      allocInLoopParam ::
        (Allocable fromrep torep inner) =>
        (Param DeclType, SubExp) ->
        WriterT
          ([FParam torep], [FParam torep])
          (AllocM fromrep torep)
          ( FParam torep,
            SubExp,
            SubExp -> WriterT ([SubExp], [SubExp]) (AllocM fromrep torep) SubExp
          )
      allocInLoopParam (mergeparam, Var v)
        | param_t@(Array pt shape u) <- paramDeclType mergeparam = do
            (v_mem, v_lmad) <- lift $ lookupArraySummary v
            v_mem_space <- lift $ lookupMemSpace v_mem
  
            -- Loop-invariant array parameters that are in scalar space
            -- are special - we do not wish to existentialise their index
            -- function at all (but the memory block is still existential).
            case v_mem_space of
              ScalarSpace {} ->
                if anyIsLoopParam (freeIn shape)
                  then do
                    -- Arrays with loop-variant shape cannot be in scalar
                    -- space, so copy them elsewhere and try again.
                    space <- lift askDefaultSpace
                    (_, v') <- lift $ allocLinearArray space (baseName v) v
                    allocInLoopParam (mergeparam, Var v')
                  else do
                    p <- newParam "mem_param" $ MemMem v_mem_space
                    tell ([p], [])
  
                    pure
                      ( mergeparam {paramDec = MemArray pt shape u $ ArrayIn (paramName p) v_lmad},
                        Var v,
                        scalarRes param_t v_mem_space v_lmad
                      )
              _ -> do
                (v_mem', v') <- lift $ ensureRowMajorArray Nothing v
                let lmad_ext =
                      LMAD.existentialize 0 $ LMAD.iota 0 $ map pe64 $ shapeDims shape
  
                v_mem_space' <- lift $ lookupMemSpace v_mem'
  
                ctx_params <-
                  replicateM (length (LMAD.existentialized lmad_ext)) $
                    newParam "ctx_param_ext" (MemPrim int64)
  
                param_lmad <-
                  instantiateLMAD $
                    LMAD.substitute
                      ( M.fromList . zip (fmap Ext [0 ..]) $
                          map (le64 . Free . paramName) ctx_params
                      )
                      lmad_ext
  
                mem_param <- newParam "mem_param" $ MemMem v_mem_space'
                tell ([mem_param], ctx_params)
                pure
                  ( mergeparam {paramDec = MemArray pt shape u $ ArrayIn (paramName mem_param) param_lmad},
                    Var v',
                    ensureArrayIn v_mem_space'
                  )
      allocInLoopParam (mergeparam, se) = doDefault mergeparam se =<< lift askDefaultSpace
  
      doDefault mergeparam se space = do
        mergeparam' <- allocInFParam mergeparam space
        pure (mergeparam', se, linearFuncallArg (paramType mergeparam) space)
  
  arrayWithLMAD ::
    (MonadBuilder m, Op (Rep m) ~ MemOp inner (Rep m), LetDec (Rep m) ~ LetDecMem) =>
    Space ->
    LMAD ->
    Type ->
    VName ->
    m (VName, VName)
  arrayWithLMAD space lmad v_t v = do
    let Array pt shape u = v_t
    mem <- allocForArray' v_t space
    v_copy <- newVName $ baseName v <> "_scalcopy"
    let pe = PatElem v_copy $ MemArray pt shape u $ ArrayIn mem lmad
    letBind (Pat [pe]) $ BasicOp $ Replicate mempty $ Var v
    pure (mem, v_copy)
  
  ensureDirectArray ::
    (Allocable fromrep torep inner) =>
    Maybe Space ->
    VName ->
    AllocM fromrep torep (VName, VName)
  ensureDirectArray space_ok v = do
    (mem, lmad) <- lookupArraySummary v
    mem_space <- lookupMemSpace mem
    default_space <- askDefaultSpace
    if LMAD.isDirect lmad && maybe True (== mem_space) space_ok
      then pure (mem, v)
      else needCopy (fromMaybe default_space space_ok)
    where
      needCopy space =
        -- We need to do a new allocation, copy 'v', and make a new
        -- binding for the size of the memory block.
        allocLinearArray space (baseName v) v
  
  allocPermArray ::
    (Allocable fromrep torep inner) =>
    Space ->
    [Int] ->
    Name ->
    VName ->
    AllocM fromrep torep (VName, VName)
  allocPermArray space perm s v = do
    t <- lookupType v
    case t of
      Array pt shape u -> do
        mem <- allocForArray t space
        v' <- newVName $ s <> "_desired_form"
        let info =
              MemArray pt shape u . ArrayIn mem $
                LMAD.permute (LMAD.iota 0 $ map pe64 $ arrayDims t) perm
            pat = Pat [PatElem v' info]
        addStm $ Let pat (defAux ()) $ BasicOp $ Manifest v perm
        pure (mem, v')
      _ ->
        error $ "allocPermArray: " ++ prettyString t
  
  ensurePermArray ::
    (Allocable fromrep torep inner) =>
    Maybe Space ->
    [Int] ->
    VName ->
    AllocM fromrep torep (VName, VName)
  ensurePermArray space_ok perm v = do
    (mem, _) <- lookupArraySummary v
    mem_space <- lookupMemSpace mem
    default_space <- askDefaultSpace
    if maybe True (== mem_space) space_ok
      then pure (mem, v)
      else allocPermArray (fromMaybe default_space space_ok) perm (baseName v) v
  
  allocLinearArray ::
    (Allocable fromrep torep inner) =>
    Space ->
    Name ->
    VName ->
    AllocM fromrep torep (VName, VName)
  allocLinearArray space s v = do
    t <- lookupType v
    let perm = [0 .. arrayRank t - 1]
    allocPermArray space perm s v
  
  funcallArgs ::
    (Allocable fromrep torep inner) =>
    [(SubExp, Diet)] ->
    AllocM fromrep torep [(SubExp, Diet)]
  funcallArgs args = do
    (valargs, (ctx_args, mem_and_size_args)) <- runWriterT $
      forM args $ \(arg, d) -> do
        t <- lift $ subExpType arg
        space <- lift askDefaultSpace
        arg' <- linearFuncallArg t space arg
        pure (arg', d)
    pure $ map (,Observe) (ctx_args <> mem_and_size_args) <> valargs
  
  linearFuncallArg ::
    (Allocable fromrep torep inner) =>
    Type ->
    Space ->
    SubExp ->
    WriterT ([SubExp], [SubExp]) (AllocM fromrep torep) SubExp
  linearFuncallArg Array {} space (Var v) = do
    (mem, arg') <- lift $ ensureDirectArray (Just space) v
    tell ([Var mem], [])
    pure $ Var arg'
  linearFuncallArg _ _ arg =
    pure arg
  
  shiftRetAls :: Int -> Int -> RetAls -> RetAls
  shiftRetAls a b (RetAls is js) =
    RetAls (map (+ a) is) (map (+ b) js)
  
  explicitAllocationsGeneric ::
    (Allocable fromrep torep inner) =>
    Space ->
    (Op fromrep -> AllocM fromrep torep (Op torep)) ->
    (Exp torep -> AllocM fromrep torep [ExpHint]) ->
    Pass fromrep torep
  explicitAllocationsGeneric space handleOp hints =
    Pass "explicit allocations" "Transform program to explicit memory representation" $
      intraproceduralTransformationWithConsts onStms allocInFun
    where
      onStms stms =
        runAllocM space handleOp hints $ collectStms_ $ allocInStms stms $ pure ()
  
      allocInFun consts (FunDef entry attrs fname rettype params fbody) =
        runAllocM space handleOp hints . inScopeOf consts $
          allocInFParams (map (,space) params) $ \params' -> do
            (fbody', mem_rets) <-
              allocInFunBody (map (const $ Just space) rettype) fbody
            let num_extra_params = length params' - length params
                num_extra_rets = length mem_rets
                -- The mem_pals is an over-approximation, like in the case for Apply.
                mem_pals = map fst $ filter (isMem . paramType . snd) $ zip [0 ..] params'
                mem_als = RetAls mem_pals mempty
                rettype' =
                  map (,mem_als) mem_rets
                    ++ zip
                      (memoryInDeclExtType space (length mem_rets) (map fst rettype))
                      (map (shiftRetAls num_extra_params num_extra_rets . snd) rettype)
            pure $ FunDef entry attrs fname rettype' params' fbody'
  
  explicitAllocationsInStmsGeneric ::
    ( MonadFreshNames m,
      HasScope torep m,
      Allocable fromrep torep inner
    ) =>
    Space ->
    (Op fromrep -> AllocM fromrep torep (Op torep)) ->
    (Exp torep -> AllocM fromrep torep [ExpHint]) ->
    Stms fromrep ->
    m (Stms torep)
  explicitAllocationsInStmsGeneric space handleOp hints stms = do
    scope <- askScope
    runAllocM space handleOp hints $
      localScope scope $
        collectStms_ $
          allocInStms stms $
            pure ()
  
  memoryInDeclExtType :: Space -> Int -> [DeclExtType] -> [FunReturns]
  memoryInDeclExtType space k dets = evalState (mapM addMem dets) 0
    where
      addMem (Prim t) = pure $ MemPrim t
      addMem Mem {} = error "memoryInDeclExtType: too much memory"
      addMem (Array pt shape u) = do
        i <- get <* modify (+ 1)
        let shape' = fmap shift shape
        pure . MemArray pt shape' u . ReturnsNewBlock space i $
          LMAD.iota 0 (map convert $ shapeDims shape')
      addMem (Acc acc ispace ts u) = pure $ MemAcc acc ispace ts u
  
      convert (Ext i) = le64 $ Ext i
      convert (Free v) = Free <$> pe64 v
  
      shift (Ext i) = Ext (i + k)
      shift (Free x) = Free x
  
  bodyReturnMemCtx ::
    (Allocable fromrep torep inner) =>
    SubExpRes ->
    AllocM fromrep torep [(SubExpRes, MemInfo ExtSize u MemReturn)]
  bodyReturnMemCtx (SubExpRes _ Constant {}) =
    pure []
  bodyReturnMemCtx (SubExpRes _ (Var v)) = do
    info <- lookupMemInfo v
    case info of
      MemPrim {} -> pure []
      MemAcc {} -> pure []
      MemMem {} -> pure [] -- should not happen
      MemArray _ _ _ (ArrayIn mem _) -> do
        mem_info <- lookupMemInfo mem
        case mem_info of
          MemMem space ->
            pure [(subExpRes $ Var mem, MemMem space)]
          _ -> error $ "bodyReturnMemCtx: not a memory block: " ++ prettyString mem
  
  allocInFunBody ::
    (Allocable fromrep torep inner) =>
    [Maybe Space] ->
    Body fromrep ->
    AllocM fromrep torep (Body torep, [FunReturns])
  allocInFunBody space_oks (Body _ stms res) =
    buildBody . allocInStms stms $ do
      res' <- zipWithM ensureDirect space_oks' res
      (mem_ctx_res, mem_ctx_rets) <- unzip . concat <$> mapM bodyReturnMemCtx res'
      pure (mem_ctx_res <> res', mem_ctx_rets)
    where
      num_vals = length space_oks
      space_oks' = replicate (length res - num_vals) Nothing ++ space_oks
  
  ensureDirect ::
    (Allocable fromrep torep inner) =>
    Maybe Space ->
    SubExpRes ->
    AllocM fromrep torep SubExpRes
  ensureDirect space_ok (SubExpRes cs se) = do
    se_info <- subExpMemInfo se
    SubExpRes cs <$> case (se_info, se) of
      (MemArray {}, Var v) -> do
        (_, v') <- ensureDirectArray space_ok v
        pure $ Var v'
      _ ->
        pure se
  
  allocInStms ::
    (Allocable fromrep torep inner) =>
    Stms fromrep ->
    AllocM fromrep torep a ->
    AllocM fromrep torep a
  allocInStms origstms m = allocInStms' $ stmsToList origstms
    where
      allocInStms' [] = m
      allocInStms' (stm : stms) = do
        allocstms <- collectStms_ $ auxing (stmAux stm) $ allocInStm stm
        addStms allocstms
        let stms_consts = foldMap stmConsts allocstms
            f env = env {envConsts = stms_consts <> envConsts env}
        local f $ allocInStms' stms
  
  allocInStm ::
    (Allocable fromrep torep inner) =>
    Stm fromrep ->
    AllocM fromrep torep ()
  allocInStm (Let (Pat pes) _ e) =
    addStm =<< allocsForStm (map patElemIdent pes) =<< allocInExp e
  
  allocInLambda ::
    (Allocable fromrep torep inner) =>
    [LParam torep] ->
    Body fromrep ->
    AllocM fromrep torep (Lambda torep)
  allocInLambda params body =
    mkLambda params . allocInStms (bodyStms body) $ pure $ bodyResult body
  
  data MemReq
    = MemReq Space
    | NeedsNormalisation Space
    deriving (Eq, Show)
  
  combMemReqs :: MemReq -> MemReq -> MemReq
  combMemReqs x@NeedsNormalisation {} _ = x
  combMemReqs _ y@NeedsNormalisation {} = y
  combMemReqs x@(MemReq x_space) y@MemReq {} =
    if x == y then x else NeedsNormalisation x_space
  
  type MemReqType = MemInfo (Ext SubExp) NoUniqueness MemReq
  
  combMemReqTypes :: MemReqType -> MemReqType -> MemReqType
  combMemReqTypes (MemArray pt shape u x) (MemArray _ _ _ y) =
    MemArray pt shape u $ combMemReqs x y
  combMemReqTypes x _ = x
  
  contextRets :: MemReqType -> [MemInfo d u r]
  contextRets (MemArray _ shape _ (MemReq space)) =
    -- Memory + offset + stride*rank.
    [MemMem space, MemPrim int64]
      ++ replicate (shapeRank shape) (MemPrim int64)
  contextRets (MemArray _ shape _ (NeedsNormalisation space)) =
    -- Memory + offset + stride*rank.
    [MemMem space, MemPrim int64]
      ++ replicate (shapeRank shape) (MemPrim int64)
  contextRets _ = []
  
  -- Add memory information to the body, but do not return memory/lmad
  -- information.  Instead, return restrictions on what the index
  -- function should look like.  We will then (crudely) unify these
  -- restrictions across all bodies.
  allocInMatchBody ::
    (Allocable fromrep torep inner) =>
    [ExtType] ->
    Body fromrep ->
    AllocM fromrep torep (Body torep, [MemReqType])
  allocInMatchBody rets (Body _ stms res) =
    buildBody . allocInStms stms $ do
      restrictions <- zipWithM restriction rets (map resSubExp res)
      pure (res, restrictions)
    where
      restriction t se = do
        v_info <- subExpMemInfo se
        case (t, v_info) of
          (Array pt shape u, MemArray _ _ _ (ArrayIn mem _)) -> do
            space <- lookupMemSpace mem
            pure $ MemArray pt shape u $ MemReq space
          (_, MemMem space) -> pure $ MemMem space
          (_, MemPrim pt) -> pure $ MemPrim pt
          (_, MemAcc acc ispace ts u) -> pure $ MemAcc acc ispace ts u
          _ -> error $ "allocInMatchBody: mismatch: " ++ show (t, v_info)
  
  mkBranchRet :: [MemReqType] -> [BranchTypeMem]
  mkBranchRet reqs =
    let (ctx_rets, res_rets) = foldl helper ([], []) $ zip reqs offsets
     in ctx_rets ++ res_rets
    where
      numCtxNeeded = length . contextRets
  
      offsets = scanl (+) 0 $ map numCtxNeeded reqs
      num_new_ctx = last offsets
  
      helper (ctx_rets_acc, res_rets_acc) (req, ctx_offset) =
        ( ctx_rets_acc ++ contextRets req,
          res_rets_acc ++ [inspect ctx_offset req]
        )
  
      arrayInfo (NeedsNormalisation space) =
        space
      arrayInfo (MemReq space) =
        space
  
      inspect ctx_offset (MemArray pt shape u req) =
        let shape' = fmap (adjustExt num_new_ctx) shape
            space = arrayInfo req
         in MemArray pt shape' u . ReturnsNewBlock space ctx_offset $
              convert
                <$> LMAD.mkExistential (shapeDims shape') (ctx_offset + 1)
      inspect _ (MemAcc acc ispace ts u) = MemAcc acc ispace ts u
      inspect _ (MemPrim pt) = MemPrim pt
      inspect _ (MemMem space) = MemMem space
  
      convert (Ext i) = le64 (Ext i)
      convert (Free v) = Free <$> pe64 v
  
      adjustExt :: Int -> Ext a -> Ext a
      adjustExt _ (Free v) = Free v
      adjustExt k (Ext i) = Ext (k + i)
  
  addCtxToMatchBody ::
    (Allocable fromrep torep inner) =>
    [MemReqType] ->
    Body torep ->
    AllocM fromrep torep (Body torep)
  addCtxToMatchBody reqs body = buildBody_ $ do
    res <- zipWithM normaliseIfNeeded reqs =<< bodyBind body
    ctx <- concat <$> mapM resCtx res
    pure $ ctx ++ res
    where
      normaliseIfNeeded (MemArray _ shape _ (NeedsNormalisation space)) (SubExpRes cs (Var v)) =
        SubExpRes cs . Var . snd
          <$> ensurePermArray (Just space) [0 .. shapeRank shape - 1] v
      normaliseIfNeeded _ res =
        pure res
  
      resCtx (SubExpRes _ Constant {}) =
        pure []
      resCtx (SubExpRes _ (Var v)) = do
        info <- lookupMemInfo v
        case info of
          MemPrim {} -> pure []
          MemAcc {} -> pure []
          MemMem {} -> pure [] -- should not happen
          MemArray _ _ _ (ArrayIn mem lmad) -> do
            lmad_exts <- mapM (letSubExp "lmad_ext" <=< toExp) $ LMAD.existentialized lmad
            pure $ subExpRes (Var mem) : subExpsRes lmad_exts
  
  -- Do a a simple form of invariance analysis to simplify a Match.  It
  -- is unfortunate that we have to do it here, but functions such as
  -- scalarRes will look carefully at the index functions before the
  -- simplifier has a chance to run.  In a perfect world we would
  -- simplify away those copies afterwards. XXX; this should be fixed by
  -- a more general copy-removal pass. See
  -- Futhark.Optimise.EntryPointMem for a very specialised version of
  -- the idea, but which could perhaps be generalised.
  simplifyMatch ::
    (Mem rep inner) =>
    [Case (Body rep)] ->
    Body rep ->
    [BranchTypeMem] ->
    ( [Case (Body rep)],
      Body rep,
      [BranchTypeMem]
    )
  simplifyMatch cases defbody ts =
    let case_reses = map (bodyResult . caseBody) cases
        defbody_res = bodyResult defbody
        (ctx_fixes, variant) =
          partitionEithers . map branchInvariant $
            zip4 [0 ..] (transpose case_reses) defbody_res ts
        (cases_reses, defbody_reses, ts') = unzip3 variant
     in ( zipWith onCase cases (transpose cases_reses),
          onBody defbody defbody_reses,
          foldr (uncurry fixExt) ts' ctx_fixes
        )
    where
      bound_in_branches =
        namesFromList . concatMap (patNames . stmPat) $
          foldMap (bodyStms . caseBody) cases <> bodyStms defbody
  
      onCase c res = fmap (`onBody` res) c
      onBody body res = body {bodyResult = res}
  
      branchInvariant (i, case_reses, defres, t)
        -- If even one branch has a variant result, then we give up.
        | namesIntersect bound_in_branches $ freeIn $ defres : case_reses =
            Right (case_reses, defres, t)
        -- Do all branches return the same value?
        | all ((== resSubExp defres) . resSubExp) case_reses =
            Left (i, resSubExp defres)
        | otherwise =
            Right (case_reses, defres, t)
  
  allocInExp ::
    (Allocable fromrep torep inner) =>
    Exp fromrep ->
    AllocM fromrep torep (Exp torep)
  allocInExp (Loop merge form (Body () bodystms bodyres)) =
    allocInLoopParams merge $ \merge' mk_loop_val -> do
      localScope (scopeOfLoopForm form) $ do
        body' <-
          buildBody_ . allocInStms bodystms $ do
            (valctx, valres') <- mk_loop_val $ map resSubExp bodyres
            pure $ subExpsRes valctx <> zipWith SubExpRes (map resCerts bodyres) valres'
        pure $ Loop merge' form body'
  allocInExp (Apply fname args rettype loc) = do
    args' <- funcallArgs args
    space <- askDefaultSpace
    args_ts <- mapM (subExpType . fst) args'
    -- We assume that every array is going to be in its own memory. Further, we
    -- assume that every result memory block can alias any argument memory block.
    -- This is an overapproximation that can be loosened in the future.
    let mem_als = RetAls (map fst $ filter (isMem . snd) $ zip [0 ..] args_ts) mempty
        mems = replicate num_arrays (MemMem space, mem_als)
        num_extra_args = length args' - length args
        rettype' =
          mems
            ++ zip
              (memoryInDeclExtType space num_arrays (map fst rettype))
              (map (shiftRetAls num_extra_args num_arrays . snd) rettype)
    pure $ Apply fname args' rettype' loc
    where
      num_arrays = length $ filter ((> 0) . arrayRank . declExtTypeOf . fst) rettype
  allocInExp (Match ses cases defbody (MatchDec rets ifsort)) = do
    (defbody', def_reqs) <- allocInMatchBody rets defbody
    (cases', cases_reqs) <- mapAndUnzipM onCase cases
    let reqs = zipWith (foldl combMemReqTypes) def_reqs (transpose cases_reqs)
    defbody'' <- addCtxToMatchBody reqs defbody'
    cases'' <- mapM (traverse $ addCtxToMatchBody reqs) cases'
    let (cases''', defbody''', rets') =
          simplifyMatch cases'' defbody'' $ mkBranchRet reqs
    pure $ Match ses cases''' defbody''' $ MatchDec rets' ifsort
    where
      onCase (Case vs body) = first (Case vs) <$> allocInMatchBody rets body
  allocInExp (WithAcc inputs bodylam) =
    WithAcc <$> mapM onInput inputs <*> onLambda bodylam
    where
      onLambda lam = do
        params <- forM (lambdaParams lam) $ \(Param attrs pv t) ->
          case t of
            Prim Unit -> pure $ Param attrs pv $ MemPrim Unit
            Acc acc ispace ts u -> pure $ Param attrs pv $ MemAcc acc ispace ts u
            _ -> error $ "Unexpected WithAcc lambda param: " ++ prettyString (Param attrs pv t)
        allocInLambda params (lambdaBody lam)
  
      onInput (shape, arrs, op) =
        (shape,arrs,) <$> traverse (onOp shape arrs) op
  
      onOp accshape arrs (lam, nes) = do
        let num_vs = length (lambdaReturnType lam)
            num_is = shapeRank accshape
            (i_params, x_params, y_params) =
              splitAt3 num_is num_vs $ lambdaParams lam
            i_params' = map (\(Param attrs v _) -> Param attrs v $ MemPrim int64) i_params
            is = map (DimFix . Var . paramName) i_params'
        x_params' <- zipWithM (onXParam is) x_params arrs
        y_params' <- zipWithM (onYParam is) y_params arrs
        lam' <-
          allocInLambda
            (i_params' <> x_params' <> y_params')
            (lambdaBody lam)
        pure (lam', nes)
  
      mkP attrs p pt shape u mem lmad is =
        Param attrs p . MemArray pt shape u . ArrayIn mem . LMAD.slice lmad $
          fmap pe64 $
            Slice $
              is ++ map sliceDim (shapeDims shape)
  
      onXParam _ (Param attrs p (Prim t)) _ =
        pure $ Param attrs p (MemPrim t)
      onXParam is (Param attrs p (Array pt shape u)) arr = do
        (mem, lmad) <- lookupArraySummary arr
        pure $ mkP attrs p pt shape u mem lmad is
      onXParam _ p _ =
        error $ "Cannot handle MkAcc param: " ++ prettyString p
  
      onYParam _ (Param attrs p (Prim t)) _ =
        pure $ Param attrs p $ MemPrim t
      onYParam is (Param attrs p (Array pt shape u)) arr = do
        arr_t <- lookupType arr
        space <- askDefaultSpace
        mem <- allocForArray arr_t space
        let base_dims = map pe64 $ arrayDims arr_t
            lmad = LMAD.iota 0 base_dims
        pure $ mkP attrs p pt shape u mem lmad is
      onYParam _ p _ =
        error $ "Cannot handle MkAcc param: " ++ prettyString p
  allocInExp e = mapExpM alloc e
    where
      alloc =
        identityMapper
          { mapOnBody = error "Unhandled Body in ExplicitAllocations",
            mapOnRetType = error "Unhandled RetType in ExplicitAllocations",
            mapOnBranchType = error "Unhandled BranchType in ExplicitAllocations",
            mapOnFParam = error "Unhandled FParam in ExplicitAllocations",
            mapOnLParam = error "Unhandled LParam in ExplicitAllocations",
            mapOnOp = \op -> do
              handle <- asks allocInOp
              handle op
          }
  
  class SizeSubst op where
    opIsConst :: op -> Bool
    opIsConst = const False
  
  instance SizeSubst (NoOp rep)
  
  instance (SizeSubst (op rep)) => SizeSubst (MemOp op rep) where
    opIsConst (Inner op) = opIsConst op
    opIsConst _ = False
  
  stmConsts :: (SizeSubst (Op rep)) => Stm rep -> S.Set VName
  stmConsts (Let pat _ (Op op))
    | opIsConst op = S.fromList $ patNames pat
  stmConsts _ = mempty
  
  mkLetNamesB' ::
    ( LetDec (Rep m) ~ LetDecMem,
      Mem (Rep m) inner,
      MonadBuilder m,
      ExpDec (Rep m) ~ ()
    ) =>
    Space ->
    ExpDec (Rep m) ->
    [VName] ->
    Exp (Rep m) ->
    m (Stm (Rep m))
  mkLetNamesB' space dec names e = do
    pat <- patWithAllocations space names e nohints
    pure $ Let pat (defAux dec) e
    where
      nohints = map (const NoHint) names
  
  mkLetNamesB'' ::
    ( Mem rep inner,
      LetDec rep ~ LetDecMem,
      OpReturns inner,
      ExpDec rep ~ (),
      Rep m ~ Engine.Wise rep,
      HasScope (Engine.Wise rep) m,
      MonadBuilder m,
      AliasedOp inner,
      RephraseOp (MemOp inner),
      Engine.CanBeWise inner,
      ASTConstraints (inner (Engine.Wise rep))
    ) =>
    Space ->
    [VName] ->
    Exp (Engine.Wise rep) ->
    m (Stm (Engine.Wise rep))
  mkLetNamesB'' space names e = do
    pat <- patWithAllocations space names e nohints
    let pat' = Engine.addWisdomToPat pat e
        dec = Engine.mkWiseExpDec pat' () e
    pure $ Let pat' (defAux dec) e
    where
      nohints = map (const NoHint) names
  
  simplifyMemOp ::
    (Engine.SimplifiableRep rep) =>
    ( inner (Engine.Wise rep) ->
      Engine.SimpleM rep (inner (Engine.Wise rep), Stms (Engine.Wise rep))
    ) ->
    MemOp inner (Engine.Wise rep) ->
    Engine.SimpleM rep (MemOp inner (Engine.Wise rep), Stms (Engine.Wise rep))
  simplifyMemOp _ (Alloc size space) =
    (,) <$> (Alloc <$> Engine.simplify size <*> pure space) <*> pure mempty
  simplifyMemOp onInner (Inner k) = do
    (k', hoisted) <- onInner k
    pure (Inner k', hoisted)
  
  simplifiable ::
    ( Engine.SimplifiableRep rep,
      LetDec rep ~ LetDecMem,
      ExpDec rep ~ (),
      BodyDec rep ~ (),
      Mem (Engine.Wise rep) inner,
      Engine.CanBeWise inner,
      RephraseOp inner,
      IsOp inner,
      OpReturns inner,
      AliasedOp inner,
      IndexOp (inner (Engine.Wise rep))
    ) =>
    (inner (Engine.Wise rep) -> UT.UsageTable) ->
    ( inner (Engine.Wise rep) ->
      Engine.SimpleM rep (inner (Engine.Wise rep), Stms (Engine.Wise rep))
    ) ->
    SimpleOps rep
  simplifiable innerUsage simplifyInnerOp =
    SimpleOps mkExpDecS' mkBodyS' protectOp opUsage (simplifyMemOp simplifyInnerOp)
    where
      mkExpDecS' _ pat e =
        pure $ Engine.mkWiseExpDec pat () e
  
      mkBodyS' _ stms res = pure $ mkWiseBody () stms res
  
      protectOp taken pat (Alloc size space) = Just $ do
        tbody <- resultBodyM [size]
        fbody <- resultBodyM [intConst Int64 0]
        size' <-
          letSubExp "hoisted_alloc_size" $
            Match [taken] [Case [Just $ BoolValue True] tbody] fbody $
              MatchDec [MemPrim int64] MatchFallback
        letBind pat $ Op $ Alloc size' space
      protectOp _ _ _ = Nothing
  
      opUsage (Alloc (Var size) _) =
        UT.sizeUsage size
      opUsage (Alloc _ _) =
        mempty
      opUsage (Inner inner) =
        innerUsage inner
  
  data ExpHint
    = NoHint
    | Hint LMAD Space
  
  defaultExpHints :: (ASTRep rep, HasScope rep m) => Exp rep -> m [ExpHint]
  defaultExpHints e = map (const NoHint) <$> expExtType e
  