{-# LANGUAGE GeneralizedNewtypeDeriving #-}
{-# LANGUAGE TypeFamilies #-}
{-# LANGUAGE FlexibleContexts #-}
{-# LANGUAGE FlexibleInstances #-}
{-# LANGUAGE MultiParamTypeClasses #-}
{-# LANGUAGE TupleSections #-}
{-# LANGUAGE ConstraintKinds #-}
{-# LANGUAGE UndecidableInstances #-}
{-# LANGUAGE DefaultSignatures #-}
-- | 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

       , Allocable
       , Allocator(..)
       , AllocM
       , AllocEnv(..)
       , SizeSubst(..)
       , allocInStms
       , allocForArray

       , simplifiable
       , arraySizeInBytesExp

       , 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.State
import Control.Monad.Writer
import Control.Monad.Reader
import Control.Monad.RWS.Strict
import qualified Data.Map.Strict as M
import qualified Data.Set as S
import Data.Maybe
import Data.List (zip4, partition, sort)

import Futhark.Optimise.Simplify.Lore
  (mkWiseBody, mkWiseLetStm, removeExpWisdom, removeScopeWisdom)
import Futhark.MonadFreshNames
import Futhark.Representation.Mem
import qualified Futhark.Representation.Mem.IxFun as IxFun
import Futhark.Tools
import qualified Futhark.Analysis.SymbolTable as ST
import Futhark.Optimise.Simplify.Engine (SimpleOps (..))
import qualified Futhark.Optimise.Simplify.Engine as Engine
import Futhark.Pass
import Futhark.Util (splitFromEnd, takeLast)

data AllocStm = SizeComputation VName (PrimExp VName)
              | Allocation VName SubExp Space
              | ArrayCopy VName VName
                    deriving (AllocStm -> AllocStm -> Bool
(AllocStm -> AllocStm -> Bool)
-> (AllocStm -> AllocStm -> Bool) -> Eq AllocStm
forall a. (a -> a -> Bool) -> (a -> a -> Bool) -> Eq a
/= :: AllocStm -> AllocStm -> Bool
$c/= :: AllocStm -> AllocStm -> Bool
== :: AllocStm -> AllocStm -> Bool
$c== :: AllocStm -> AllocStm -> Bool
Eq, Eq AllocStm
Eq AllocStm
-> (AllocStm -> AllocStm -> Ordering)
-> (AllocStm -> AllocStm -> Bool)
-> (AllocStm -> AllocStm -> Bool)
-> (AllocStm -> AllocStm -> Bool)
-> (AllocStm -> AllocStm -> Bool)
-> (AllocStm -> AllocStm -> AllocStm)
-> (AllocStm -> AllocStm -> AllocStm)
-> Ord AllocStm
AllocStm -> AllocStm -> Bool
AllocStm -> AllocStm -> Ordering
AllocStm -> AllocStm -> AllocStm
forall a.
Eq a
-> (a -> a -> Ordering)
-> (a -> a -> Bool)
-> (a -> a -> Bool)
-> (a -> a -> Bool)
-> (a -> a -> Bool)
-> (a -> a -> a)
-> (a -> a -> a)
-> Ord a
min :: AllocStm -> AllocStm -> AllocStm
$cmin :: AllocStm -> AllocStm -> AllocStm
max :: AllocStm -> AllocStm -> AllocStm
$cmax :: AllocStm -> AllocStm -> AllocStm
>= :: AllocStm -> AllocStm -> Bool
$c>= :: AllocStm -> AllocStm -> Bool
> :: AllocStm -> AllocStm -> Bool
$c> :: AllocStm -> AllocStm -> Bool
<= :: AllocStm -> AllocStm -> Bool
$c<= :: AllocStm -> AllocStm -> Bool
< :: AllocStm -> AllocStm -> Bool
$c< :: AllocStm -> AllocStm -> Bool
compare :: AllocStm -> AllocStm -> Ordering
$ccompare :: AllocStm -> AllocStm -> Ordering
$cp1Ord :: Eq AllocStm
Ord, Int -> AllocStm -> ShowS
[AllocStm] -> ShowS
AllocStm -> String
(Int -> AllocStm -> ShowS)
-> (AllocStm -> String) -> ([AllocStm] -> ShowS) -> Show AllocStm
forall a.
(Int -> a -> ShowS) -> (a -> String) -> ([a] -> ShowS) -> Show a
showList :: [AllocStm] -> ShowS
$cshowList :: [AllocStm] -> ShowS
show :: AllocStm -> String
$cshow :: AllocStm -> String
showsPrec :: Int -> AllocStm -> ShowS
$cshowsPrec :: Int -> AllocStm -> ShowS
Show)

bindAllocStm :: (MonadBinder m, Op (Lore m) ~ MemOp inner) =>
                AllocStm -> m ()
bindAllocStm :: AllocStm -> m ()
bindAllocStm (SizeComputation VName
name PrimExp VName
pe) =
  [VName] -> ExpT (Lore m) -> m ()
forall (m :: * -> *).
MonadBinder m =>
[VName] -> Exp (Lore m) -> m ()
letBindNames_ [VName
name] (ExpT (Lore m) -> m ()) -> m (ExpT (Lore m)) -> m ()
forall (m :: * -> *) a b. Monad m => (a -> m b) -> m a -> m b
=<< PrimExp VName -> m (ExpT (Lore m))
forall a (m :: * -> *).
(ToExp a, MonadBinder m) =>
a -> m (Exp (Lore m))
toExp (IntType -> PrimExp VName -> PrimExp VName
forall v. IntType -> PrimExp v -> PrimExp v
coerceIntPrimExp IntType
Int64 PrimExp VName
pe)
bindAllocStm (Allocation VName
name SubExp
size Space
space) =
  [VName] -> ExpT (Lore m) -> m ()
forall (m :: * -> *).
MonadBinder m =>
[VName] -> Exp (Lore m) -> m ()
letBindNames_ [VName
name] (ExpT (Lore m) -> m ()) -> ExpT (Lore m) -> m ()
forall a b. (a -> b) -> a -> b
$ Op (Lore m) -> ExpT (Lore m)
forall lore. Op lore -> ExpT lore
Op (Op (Lore m) -> ExpT (Lore m)) -> Op (Lore m) -> ExpT (Lore m)
forall a b. (a -> b) -> a -> b
$ SubExp -> Space -> MemOp inner
forall inner. SubExp -> Space -> MemOp inner
Alloc SubExp
size Space
space
bindAllocStm (ArrayCopy VName
name VName
src) =
  [VName] -> ExpT (Lore m) -> m ()
forall (m :: * -> *).
MonadBinder m =>
[VName] -> Exp (Lore m) -> m ()
letBindNames_ [VName
name] (ExpT (Lore m) -> m ()) -> ExpT (Lore m) -> m ()
forall a b. (a -> b) -> a -> b
$ BasicOp -> ExpT (Lore m)
forall lore. BasicOp -> ExpT lore
BasicOp (BasicOp -> ExpT (Lore m)) -> BasicOp -> ExpT (Lore m)
forall a b. (a -> b) -> a -> b
$ VName -> BasicOp
Copy VName
src

class (MonadFreshNames m, HasScope lore m, Mem lore) =>
      Allocator lore m where
  addAllocStm :: AllocStm -> m ()
  askDefaultSpace :: m Space

  default addAllocStm :: (Allocable fromlore lore,
                          m ~ AllocM fromlore lore)
                      => AllocStm -> m ()
  addAllocStm (SizeComputation VName
name PrimExp VName
se) =
    [VName] -> Exp (Lore m) -> m ()
forall (m :: * -> *).
MonadBinder m =>
[VName] -> Exp (Lore m) -> m ()
letBindNames_ [VName
name] (ExpT lore -> m ()) -> m (ExpT lore) -> m ()
forall (m :: * -> *) a b. Monad m => (a -> m b) -> m a -> m b
=<< PrimExp VName -> m (Exp (Lore m))
forall a (m :: * -> *).
(ToExp a, MonadBinder m) =>
a -> m (Exp (Lore m))
toExp (IntType -> PrimExp VName -> PrimExp VName
forall v. IntType -> PrimExp v -> PrimExp v
coerceIntPrimExp IntType
Int64 PrimExp VName
se)
  addAllocStm (Allocation VName
name SubExp
size Space
space) =
    [VName] -> Exp (Lore m) -> m ()
forall (m :: * -> *).
MonadBinder m =>
[VName] -> Exp (Lore m) -> m ()
letBindNames_ [VName
name] (Exp (Lore m) -> m ()) -> Exp (Lore m) -> m ()
forall a b. (a -> b) -> a -> b
$ Op lore -> ExpT lore
forall lore. Op lore -> ExpT lore
Op (Op lore -> ExpT lore) -> Op lore -> ExpT lore
forall a b. (a -> b) -> a -> b
$ SubExp -> Space -> Op lore
forall op. AllocOp op => SubExp -> Space -> op
allocOp SubExp
size Space
space
  addAllocStm (ArrayCopy VName
name VName
src) =
    [VName] -> Exp (Lore m) -> m ()
forall (m :: * -> *).
MonadBinder m =>
[VName] -> Exp (Lore m) -> m ()
letBindNames_ [VName
name] (Exp (Lore m) -> m ()) -> Exp (Lore m) -> m ()
forall a b. (a -> b) -> a -> b
$ BasicOp -> ExpT lore
forall lore. BasicOp -> ExpT lore
BasicOp (BasicOp -> ExpT lore) -> BasicOp -> ExpT lore
forall a b. (a -> b) -> a -> b
$ VName -> BasicOp
Copy VName
src

  -- | The subexpression giving the number of elements we should
  -- allocate space for.  See 'ChunkMap' comment.
  dimAllocationSize :: SubExp -> m SubExp

  default dimAllocationSize :: m ~ AllocM fromlore lore
                               => SubExp -> m SubExp
  dimAllocationSize (Var VName
v) =
    -- It is important to recurse here, as the substitution may itself
    -- be a chunk size.
    m SubExp -> (SubExp -> m SubExp) -> Maybe SubExp -> m SubExp
forall b a. b -> (a -> b) -> Maybe a -> b
maybe (SubExp -> m SubExp
forall (m :: * -> *) a. Monad m => a -> m a
return (SubExp -> m SubExp) -> SubExp -> m SubExp
forall a b. (a -> b) -> a -> b
$ VName -> SubExp
Var VName
v) SubExp -> m SubExp
forall lore (m :: * -> *). Allocator lore m => SubExp -> m SubExp
dimAllocationSize (Maybe SubExp -> m SubExp) -> m (Maybe SubExp) -> m SubExp
forall (m :: * -> *) a b. Monad m => (a -> m b) -> m a -> m b
=<< (AllocEnv fromlore lore -> Maybe SubExp) -> m (Maybe SubExp)
forall r (m :: * -> *) a. MonadReader r m => (r -> a) -> m a
asks (VName -> Map VName SubExp -> Maybe SubExp
forall k a. Ord k => k -> Map k a -> Maybe a
M.lookup VName
v (Map VName SubExp -> Maybe SubExp)
-> (AllocEnv fromlore lore -> Map VName SubExp)
-> AllocEnv fromlore lore
-> Maybe SubExp
forall b c a. (b -> c) -> (a -> b) -> a -> c
. AllocEnv fromlore lore -> Map VName SubExp
forall fromlore tolore.
AllocEnv fromlore tolore -> Map VName SubExp
chunkMap)
  dimAllocationSize SubExp
size =
    SubExp -> m SubExp
forall (m :: * -> *) a. Monad m => a -> m a
return SubExp
size

  -- | Get those names that are known to be constants at run-time.
  askConsts :: m (S.Set VName)

  expHints :: Exp lore -> m [ExpHint]
  expHints = ExpT lore -> m [ExpHint]
forall (m :: * -> *) lore.
(Monad m, Attributes lore) =>
Exp lore -> m [ExpHint]
defaultExpHints

allocateMemory :: Allocator lore m =>
                  String -> SubExp -> Space -> m VName
allocateMemory :: String -> SubExp -> Space -> m VName
allocateMemory String
desc SubExp
size Space
space = do
  VName
v <- String -> m VName
forall (m :: * -> *). MonadFreshNames m => String -> m VName
newVName String
desc
  AllocStm -> m ()
forall lore (m :: * -> *). Allocator lore m => AllocStm -> m ()
addAllocStm (AllocStm -> m ()) -> AllocStm -> m ()
forall a b. (a -> b) -> a -> b
$ VName -> SubExp -> Space -> AllocStm
Allocation VName
v SubExp
size Space
space
  VName -> m VName
forall (m :: * -> *) a. Monad m => a -> m a
return VName
v

computeSize :: Allocator lore m =>
               String -> PrimExp VName -> m SubExp
computeSize :: String -> PrimExp VName -> m SubExp
computeSize String
desc PrimExp VName
se = do
  VName
v <- String -> m VName
forall (m :: * -> *). MonadFreshNames m => String -> m VName
newVName String
desc
  AllocStm -> m ()
forall lore (m :: * -> *). Allocator lore m => AllocStm -> m ()
addAllocStm (AllocStm -> m ()) -> AllocStm -> m ()
forall a b. (a -> b) -> a -> b
$ VName -> PrimExp VName -> AllocStm
SizeComputation VName
v PrimExp VName
se
  SubExp -> m SubExp
forall (m :: * -> *) a. Monad m => a -> m a
return (SubExp -> m SubExp) -> SubExp -> m SubExp
forall a b. (a -> b) -> a -> b
$ VName -> SubExp
Var VName
v

type Allocable fromlore tolore =
  (PrettyLore fromlore, PrettyLore tolore,
   Mem tolore,
   FParamAttr fromlore ~ DeclType,
   LParamAttr fromlore ~ Type,
   BranchType fromlore ~ ExtType,
   RetType fromlore ~ DeclExtType,
   BodyAttr fromlore ~ (),
   BodyAttr tolore ~ (),
   ExpAttr tolore ~ (),
   SizeSubst (Op tolore),
   BinderOps tolore)

-- | A mapping from chunk names to their maximum size.  XXX FIXME
-- HACK: This is part of a hack to add loop-invariant allocations to
-- reduce kernels, because memory expansion does not use range
-- analysis yet (it should).
type ChunkMap = M.Map VName SubExp

data AllocEnv fromlore tolore  =
  AllocEnv { AllocEnv fromlore tolore -> Map VName SubExp
chunkMap :: ChunkMap
           , AllocEnv fromlore tolore -> Bool
aggressiveReuse :: Bool
             -- ^ Aggressively try to reuse memory in do-loops -
             -- should be True inside kernels, False outside.
           , AllocEnv fromlore tolore -> Space
allocSpace :: Space
             -- ^ When allocating memory, put it in this memory space.
             -- This is primarily used to ensure that group-wide
             -- statements store their results in local memory.
           , AllocEnv fromlore tolore -> Set VName
envConsts :: S.Set VName
             -- ^ The set of names that are known to be constants at
             -- kernel compile time.
           , AllocEnv fromlore tolore
-> Op fromlore -> AllocM fromlore tolore (Op tolore)
allocInOp :: Op fromlore -> AllocM fromlore tolore (Op tolore)
           , AllocEnv fromlore tolore
-> Exp tolore -> AllocM fromlore tolore [ExpHint]
envExpHints :: Exp tolore -> AllocM fromlore tolore [ExpHint]
           }

-- | Monad for adding allocations to an entire program.
newtype AllocM fromlore tolore a =
  AllocM (BinderT tolore (ReaderT (AllocEnv fromlore tolore) (State VNameSource)) a)
  deriving (Functor (AllocM fromlore tolore)
a -> AllocM fromlore tolore a
Functor (AllocM fromlore tolore)
-> (forall a. a -> AllocM fromlore tolore a)
-> (forall a b.
    AllocM fromlore tolore (a -> b)
    -> AllocM fromlore tolore a -> AllocM fromlore tolore b)
-> (forall a b c.
    (a -> b -> c)
    -> AllocM fromlore tolore a
    -> AllocM fromlore tolore b
    -> AllocM fromlore tolore c)
-> (forall a b.
    AllocM fromlore tolore a
    -> AllocM fromlore tolore b -> AllocM fromlore tolore b)
-> (forall a b.
    AllocM fromlore tolore a
    -> AllocM fromlore tolore b -> AllocM fromlore tolore a)
-> Applicative (AllocM fromlore tolore)
AllocM fromlore tolore a
-> AllocM fromlore tolore b -> AllocM fromlore tolore b
AllocM fromlore tolore a
-> AllocM fromlore tolore b -> AllocM fromlore tolore a
AllocM fromlore tolore (a -> b)
-> AllocM fromlore tolore a -> AllocM fromlore tolore b
(a -> b -> c)
-> AllocM fromlore tolore a
-> AllocM fromlore tolore b
-> AllocM fromlore tolore c
forall a. a -> AllocM fromlore tolore a
forall fromlore tolore. Functor (AllocM fromlore tolore)
forall a b.
AllocM fromlore tolore a
-> AllocM fromlore tolore b -> AllocM fromlore tolore a
forall a b.
AllocM fromlore tolore a
-> AllocM fromlore tolore b -> AllocM fromlore tolore b
forall a b.
AllocM fromlore tolore (a -> b)
-> AllocM fromlore tolore a -> AllocM fromlore tolore b
forall fromlore tolore a. a -> AllocM fromlore tolore a
forall a b c.
(a -> b -> c)
-> AllocM fromlore tolore a
-> AllocM fromlore tolore b
-> AllocM fromlore tolore c
forall fromlore tolore a b.
AllocM fromlore tolore a
-> AllocM fromlore tolore b -> AllocM fromlore tolore a
forall fromlore tolore a b.
AllocM fromlore tolore a
-> AllocM fromlore tolore b -> AllocM fromlore tolore b
forall fromlore tolore a b.
AllocM fromlore tolore (a -> b)
-> AllocM fromlore tolore a -> AllocM fromlore tolore b
forall fromlore tolore a b c.
(a -> b -> c)
-> AllocM fromlore tolore a
-> AllocM fromlore tolore b
-> AllocM fromlore tolore c
forall (f :: * -> *).
Functor f
-> (forall a. a -> f a)
-> (forall a b. f (a -> b) -> f a -> f b)
-> (forall a b c. (a -> b -> c) -> f a -> f b -> f c)
-> (forall a b. f a -> f b -> f b)
-> (forall a b. f a -> f b -> f a)
-> Applicative f
<* :: AllocM fromlore tolore a
-> AllocM fromlore tolore b -> AllocM fromlore tolore a
$c<* :: forall fromlore tolore a b.
AllocM fromlore tolore a
-> AllocM fromlore tolore b -> AllocM fromlore tolore a
*> :: AllocM fromlore tolore a
-> AllocM fromlore tolore b -> AllocM fromlore tolore b
$c*> :: forall fromlore tolore a b.
AllocM fromlore tolore a
-> AllocM fromlore tolore b -> AllocM fromlore tolore b
liftA2 :: (a -> b -> c)
-> AllocM fromlore tolore a
-> AllocM fromlore tolore b
-> AllocM fromlore tolore c
$cliftA2 :: forall fromlore tolore a b c.
(a -> b -> c)
-> AllocM fromlore tolore a
-> AllocM fromlore tolore b
-> AllocM fromlore tolore c
<*> :: AllocM fromlore tolore (a -> b)
-> AllocM fromlore tolore a -> AllocM fromlore tolore b
$c<*> :: forall fromlore tolore a b.
AllocM fromlore tolore (a -> b)
-> AllocM fromlore tolore a -> AllocM fromlore tolore b
pure :: a -> AllocM fromlore tolore a
$cpure :: forall fromlore tolore a. a -> AllocM fromlore tolore a
$cp1Applicative :: forall fromlore tolore. Functor (AllocM fromlore tolore)
Applicative, a -> AllocM fromlore tolore b -> AllocM fromlore tolore a
(a -> b) -> AllocM fromlore tolore a -> AllocM fromlore tolore b
(forall a b.
 (a -> b) -> AllocM fromlore tolore a -> AllocM fromlore tolore b)
-> (forall a b.
    a -> AllocM fromlore tolore b -> AllocM fromlore tolore a)
-> Functor (AllocM fromlore tolore)
forall a b.
a -> AllocM fromlore tolore b -> AllocM fromlore tolore a
forall a b.
(a -> b) -> AllocM fromlore tolore a -> AllocM fromlore tolore b
forall fromlore tolore a b.
a -> AllocM fromlore tolore b -> AllocM fromlore tolore a
forall fromlore tolore a b.
(a -> b) -> AllocM fromlore tolore a -> AllocM fromlore tolore b
forall (f :: * -> *).
(forall a b. (a -> b) -> f a -> f b)
-> (forall a b. a -> f b -> f a) -> Functor f
<$ :: a -> AllocM fromlore tolore b -> AllocM fromlore tolore a
$c<$ :: forall fromlore tolore a b.
a -> AllocM fromlore tolore b -> AllocM fromlore tolore a
fmap :: (a -> b) -> AllocM fromlore tolore a -> AllocM fromlore tolore b
$cfmap :: forall fromlore tolore a b.
(a -> b) -> AllocM fromlore tolore a -> AllocM fromlore tolore b
Functor, Applicative (AllocM fromlore tolore)
a -> AllocM fromlore tolore a
Applicative (AllocM fromlore tolore)
-> (forall a b.
    AllocM fromlore tolore a
    -> (a -> AllocM fromlore tolore b) -> AllocM fromlore tolore b)
-> (forall a b.
    AllocM fromlore tolore a
    -> AllocM fromlore tolore b -> AllocM fromlore tolore b)
-> (forall a. a -> AllocM fromlore tolore a)
-> Monad (AllocM fromlore tolore)
AllocM fromlore tolore a
-> (a -> AllocM fromlore tolore b) -> AllocM fromlore tolore b
AllocM fromlore tolore a
-> AllocM fromlore tolore b -> AllocM fromlore tolore b
forall a. a -> AllocM fromlore tolore a
forall fromlore tolore. Applicative (AllocM fromlore tolore)
forall a b.
AllocM fromlore tolore a
-> AllocM fromlore tolore b -> AllocM fromlore tolore b
forall a b.
AllocM fromlore tolore a
-> (a -> AllocM fromlore tolore b) -> AllocM fromlore tolore b
forall fromlore tolore a. a -> AllocM fromlore tolore a
forall fromlore tolore a b.
AllocM fromlore tolore a
-> AllocM fromlore tolore b -> AllocM fromlore tolore b
forall fromlore tolore a b.
AllocM fromlore tolore a
-> (a -> AllocM fromlore tolore b) -> AllocM fromlore tolore b
forall (m :: * -> *).
Applicative m
-> (forall a b. m a -> (a -> m b) -> m b)
-> (forall a b. m a -> m b -> m b)
-> (forall a. a -> m a)
-> Monad m
return :: a -> AllocM fromlore tolore a
$creturn :: forall fromlore tolore a. a -> AllocM fromlore tolore a
>> :: AllocM fromlore tolore a
-> AllocM fromlore tolore b -> AllocM fromlore tolore b
$c>> :: forall fromlore tolore a b.
AllocM fromlore tolore a
-> AllocM fromlore tolore b -> AllocM fromlore tolore b
>>= :: AllocM fromlore tolore a
-> (a -> AllocM fromlore tolore b) -> AllocM fromlore tolore b
$c>>= :: forall fromlore tolore a b.
AllocM fromlore tolore a
-> (a -> AllocM fromlore tolore b) -> AllocM fromlore tolore b
$cp1Monad :: forall fromlore tolore. Applicative (AllocM fromlore tolore)
Monad,
             Monad (AllocM fromlore tolore)
Applicative (AllocM fromlore tolore)
AllocM fromlore tolore VNameSource
Applicative (AllocM fromlore tolore)
-> Monad (AllocM fromlore tolore)
-> AllocM fromlore tolore VNameSource
-> (VNameSource -> AllocM fromlore tolore ())
-> MonadFreshNames (AllocM fromlore tolore)
VNameSource -> AllocM fromlore tolore ()
forall fromlore tolore. Monad (AllocM fromlore tolore)
forall fromlore tolore. Applicative (AllocM fromlore tolore)
forall fromlore tolore. AllocM fromlore tolore VNameSource
forall fromlore tolore. VNameSource -> AllocM fromlore tolore ()
forall (m :: * -> *).
Applicative m
-> Monad m
-> m VNameSource
-> (VNameSource -> m ())
-> MonadFreshNames m
putNameSource :: VNameSource -> AllocM fromlore tolore ()
$cputNameSource :: forall fromlore tolore. VNameSource -> AllocM fromlore tolore ()
getNameSource :: AllocM fromlore tolore VNameSource
$cgetNameSource :: forall fromlore tolore. AllocM fromlore tolore VNameSource
$cp2MonadFreshNames :: forall fromlore tolore. Monad (AllocM fromlore tolore)
$cp1MonadFreshNames :: forall fromlore tolore. Applicative (AllocM fromlore tolore)
MonadFreshNames,
             HasScope tolore,
             LocalScope tolore,
             MonadReader (AllocEnv fromlore tolore))

instance (Allocable fromlore tolore, Allocator tolore (AllocM fromlore tolore)) =>
         MonadBinder (AllocM fromlore tolore) where
  type Lore (AllocM fromlore tolore) = tolore

  mkExpAttrM :: Pattern (Lore (AllocM fromlore tolore))
-> Exp (Lore (AllocM fromlore tolore))
-> AllocM fromlore tolore (ExpAttr (Lore (AllocM fromlore tolore)))
mkExpAttrM Pattern (Lore (AllocM fromlore tolore))
_ Exp (Lore (AllocM fromlore tolore))
_ = () -> AllocM fromlore tolore ()
forall (m :: * -> *) a. Monad m => a -> m a
return ()

  mkLetNamesM :: [VName]
-> Exp (Lore (AllocM fromlore tolore))
-> AllocM fromlore tolore (Stm (Lore (AllocM fromlore tolore)))
mkLetNamesM [VName]
names Exp (Lore (AllocM fromlore tolore))
e = do
    PatternT (LetAttr tolore)
pat <- [VName]
-> Exp tolore -> AllocM fromlore tolore (PatternT (LetAttr tolore))
forall lore (m :: * -> *).
(Allocator lore m, ExpAttr lore ~ ()) =>
[VName] -> Exp lore -> m (Pattern lore)
patternWithAllocations [VName]
names Exp tolore
Exp (Lore (AllocM fromlore tolore))
e
    Stm tolore -> AllocM fromlore tolore (Stm tolore)
forall (m :: * -> *) a. Monad m => a -> m a
return (Stm tolore -> AllocM fromlore tolore (Stm tolore))
-> Stm tolore -> AllocM fromlore tolore (Stm tolore)
forall a b. (a -> b) -> a -> b
$ PatternT (LetAttr tolore)
-> StmAux (ExpAttr tolore) -> Exp tolore -> Stm tolore
forall lore.
Pattern lore -> StmAux (ExpAttr lore) -> Exp lore -> Stm lore
Let PatternT (LetAttr tolore)
pat (() -> StmAux ()
forall attr. attr -> StmAux attr
defAux ()) Exp tolore
Exp (Lore (AllocM fromlore tolore))
e

  mkBodyM :: Stms (Lore (AllocM fromlore tolore))
-> Result
-> AllocM fromlore tolore (Body (Lore (AllocM fromlore tolore)))
mkBodyM Stms (Lore (AllocM fromlore tolore))
bnds Result
res = BodyT tolore -> AllocM fromlore tolore (BodyT tolore)
forall (m :: * -> *) a. Monad m => a -> m a
return (BodyT tolore -> AllocM fromlore tolore (BodyT tolore))
-> BodyT tolore -> AllocM fromlore tolore (BodyT tolore)
forall a b. (a -> b) -> a -> b
$ BodyAttr tolore -> Stms tolore -> Result -> BodyT tolore
forall lore. BodyAttr lore -> Stms lore -> Result -> BodyT lore
Body () Stms tolore
Stms (Lore (AllocM fromlore tolore))
bnds Result
res

  addStms :: Stms (Lore (AllocM fromlore tolore)) -> AllocM fromlore tolore ()
addStms Stms (Lore (AllocM fromlore tolore))
binding = BinderT
  tolore (ReaderT (AllocEnv fromlore tolore) (State VNameSource)) ()
-> AllocM fromlore tolore ()
forall fromlore tolore a.
BinderT
  tolore (ReaderT (AllocEnv fromlore tolore) (State VNameSource)) a
-> AllocM fromlore tolore a
AllocM (BinderT
   tolore (ReaderT (AllocEnv fromlore tolore) (State VNameSource)) ()
 -> AllocM fromlore tolore ())
-> BinderT
     tolore (ReaderT (AllocEnv fromlore tolore) (State VNameSource)) ()
-> AllocM fromlore tolore ()
forall a b. (a -> b) -> a -> b
$ Stms tolore
-> BinderT
     tolore (ReaderT (AllocEnv fromlore tolore) (State VNameSource)) ()
forall (m :: * -> *) lore.
Monad m =>
Stms lore -> BinderT lore m ()
addBinderStms Stms tolore
Stms (Lore (AllocM fromlore tolore))
binding
  collectStms :: AllocM fromlore tolore a
-> AllocM fromlore tolore (a, Stms (Lore (AllocM fromlore tolore)))
collectStms (AllocM BinderT
  tolore (ReaderT (AllocEnv fromlore tolore) (State VNameSource)) a
m) = BinderT
  tolore
  (ReaderT (AllocEnv fromlore tolore) (State VNameSource))
  (a, Stms tolore)
-> AllocM fromlore tolore (a, Stms tolore)
forall fromlore tolore a.
BinderT
  tolore (ReaderT (AllocEnv fromlore tolore) (State VNameSource)) a
-> AllocM fromlore tolore a
AllocM (BinderT
   tolore
   (ReaderT (AllocEnv fromlore tolore) (State VNameSource))
   (a, Stms tolore)
 -> AllocM fromlore tolore (a, Stms tolore))
-> BinderT
     tolore
     (ReaderT (AllocEnv fromlore tolore) (State VNameSource))
     (a, Stms tolore)
-> AllocM fromlore tolore (a, Stms tolore)
forall a b. (a -> b) -> a -> b
$ BinderT
  tolore (ReaderT (AllocEnv fromlore tolore) (State VNameSource)) a
-> BinderT
     tolore
     (ReaderT (AllocEnv fromlore tolore) (State VNameSource))
     (a, Stms tolore)
forall (m :: * -> *) lore a.
Monad m =>
BinderT lore m a -> BinderT lore m (a, Stms lore)
collectBinderStms BinderT
  tolore (ReaderT (AllocEnv fromlore tolore) (State VNameSource)) a
m
  certifying :: Certificates
-> AllocM fromlore tolore a -> AllocM fromlore tolore a
certifying Certificates
cs (AllocM BinderT
  tolore (ReaderT (AllocEnv fromlore tolore) (State VNameSource)) a
m) = BinderT
  tolore (ReaderT (AllocEnv fromlore tolore) (State VNameSource)) a
-> AllocM fromlore tolore a
forall fromlore tolore a.
BinderT
  tolore (ReaderT (AllocEnv fromlore tolore) (State VNameSource)) a
-> AllocM fromlore tolore a
AllocM (BinderT
   tolore (ReaderT (AllocEnv fromlore tolore) (State VNameSource)) a
 -> AllocM fromlore tolore a)
-> BinderT
     tolore (ReaderT (AllocEnv fromlore tolore) (State VNameSource)) a
-> AllocM fromlore tolore a
forall a b. (a -> b) -> a -> b
$ Certificates
-> BinderT
     tolore (ReaderT (AllocEnv fromlore tolore) (State VNameSource)) a
-> BinderT
     tolore (ReaderT (AllocEnv fromlore tolore) (State VNameSource)) a
forall (m :: * -> *) lore a.
(MonadFreshNames m, BinderOps lore) =>
Certificates -> BinderT lore m a -> BinderT lore m a
certifyingBinder Certificates
cs BinderT
  tolore (ReaderT (AllocEnv fromlore tolore) (State VNameSource)) a
m

instance (Allocable fromlore tolore) =>
         Allocator tolore (AllocM fromlore tolore) where
  expHints :: Exp tolore -> AllocM fromlore tolore [ExpHint]
expHints Exp tolore
e = do
    Exp tolore -> AllocM fromlore tolore [ExpHint]
f <- (AllocEnv fromlore tolore
 -> Exp tolore -> AllocM fromlore tolore [ExpHint])
-> AllocM
     fromlore tolore (Exp tolore -> AllocM fromlore tolore [ExpHint])
forall r (m :: * -> *) a. MonadReader r m => (r -> a) -> m a
asks AllocEnv fromlore tolore
-> Exp tolore -> AllocM fromlore tolore [ExpHint]
forall fromlore tolore.
AllocEnv fromlore tolore
-> Exp tolore -> AllocM fromlore tolore [ExpHint]
envExpHints
    Exp tolore -> AllocM fromlore tolore [ExpHint]
f Exp tolore
e
  askDefaultSpace :: AllocM fromlore tolore Space
askDefaultSpace = (AllocEnv fromlore tolore -> Space) -> AllocM fromlore tolore Space
forall r (m :: * -> *) a. MonadReader r m => (r -> a) -> m a
asks AllocEnv fromlore tolore -> Space
forall fromlore tolore. AllocEnv fromlore tolore -> Space
allocSpace

  askConsts :: AllocM fromlore tolore (Set VName)
askConsts = (AllocEnv fromlore tolore -> Set VName)
-> AllocM fromlore tolore (Set VName)
forall r (m :: * -> *) a. MonadReader r m => (r -> a) -> m a
asks AllocEnv fromlore tolore -> Set VName
forall fromlore tolore. AllocEnv fromlore tolore -> Set VName
envConsts

runAllocM :: MonadFreshNames m =>
             (Op fromlore -> AllocM fromlore tolore (Op tolore))
          -> (Exp tolore -> AllocM fromlore tolore [ExpHint])
          -> AllocM fromlore tolore a -> m a
runAllocM :: (Op fromlore -> AllocM fromlore tolore (Op tolore))
-> (Exp tolore -> AllocM fromlore tolore [ExpHint])
-> AllocM fromlore tolore a
-> m a
runAllocM Op fromlore -> AllocM fromlore tolore (Op tolore)
handleOp Exp tolore -> AllocM fromlore tolore [ExpHint]
hints (AllocM BinderT
  tolore (ReaderT (AllocEnv fromlore tolore) (State VNameSource)) a
m) =
  ((a, Stms tolore) -> a) -> m (a, Stms tolore) -> m a
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap (a, Stms tolore) -> a
forall a b. (a, b) -> a
fst (m (a, Stms tolore) -> m a) -> m (a, Stms tolore) -> m a
forall a b. (a -> b) -> a -> b
$ (VNameSource -> ((a, Stms tolore), VNameSource))
-> m (a, Stms tolore)
forall (m :: * -> *) a.
MonadFreshNames m =>
(VNameSource -> (a, VNameSource)) -> m a
modifyNameSource ((VNameSource -> ((a, Stms tolore), VNameSource))
 -> m (a, Stms tolore))
-> (VNameSource -> ((a, Stms tolore), VNameSource))
-> m (a, Stms tolore)
forall a b. (a -> b) -> a -> b
$ State VNameSource (a, Stms tolore)
-> VNameSource -> ((a, Stms tolore), VNameSource)
forall s a. State s a -> s -> (a, s)
runState (State VNameSource (a, Stms tolore)
 -> VNameSource -> ((a, Stms tolore), VNameSource))
-> State VNameSource (a, Stms tolore)
-> VNameSource
-> ((a, Stms tolore), VNameSource)
forall a b. (a -> b) -> a -> b
$ ReaderT
  (AllocEnv fromlore tolore) (State VNameSource) (a, Stms tolore)
-> AllocEnv fromlore tolore -> State VNameSource (a, Stms tolore)
forall r (m :: * -> *) a. ReaderT r m a -> r -> m a
runReaderT (BinderT
  tolore (ReaderT (AllocEnv fromlore tolore) (State VNameSource)) a
-> Scope tolore
-> ReaderT
     (AllocEnv fromlore tolore) (State VNameSource) (a, Stms tolore)
forall (m :: * -> *) lore a.
MonadFreshNames m =>
BinderT lore m a -> Scope lore -> m (a, Stms lore)
runBinderT BinderT
  tolore (ReaderT (AllocEnv fromlore tolore) (State VNameSource)) a
m Scope tolore
forall a. Monoid a => a
mempty) AllocEnv fromlore tolore
env
  where env :: AllocEnv fromlore tolore
env = AllocEnv :: forall fromlore tolore.
Map VName SubExp
-> Bool
-> Space
-> Set VName
-> (Op fromlore -> AllocM fromlore tolore (Op tolore))
-> (Exp tolore -> AllocM fromlore tolore [ExpHint])
-> AllocEnv fromlore tolore
AllocEnv { chunkMap :: Map VName SubExp
chunkMap = Map VName SubExp
forall a. Monoid a => a
mempty
                       , aggressiveReuse :: Bool
aggressiveReuse = Bool
False
                       , allocSpace :: Space
allocSpace = Space
DefaultSpace
                       , envConsts :: Set VName
envConsts = Set VName
forall a. Monoid a => a
mempty
                       , allocInOp :: Op fromlore -> AllocM fromlore tolore (Op tolore)
allocInOp = Op fromlore -> AllocM fromlore tolore (Op tolore)
handleOp
                       , envExpHints :: Exp tolore -> AllocM fromlore tolore [ExpHint]
envExpHints = Exp tolore -> AllocM fromlore tolore [ExpHint]
hints
                       }

-- | Monad for adding allocations to a single pattern.
newtype PatAllocM lore a = PatAllocM (RWS
                                      (Scope lore)
                                      [AllocStm]
                                      VNameSource
                                      a)
                    deriving (Functor (PatAllocM lore)
a -> PatAllocM lore a
Functor (PatAllocM lore)
-> (forall a. a -> PatAllocM lore a)
-> (forall a b.
    PatAllocM lore (a -> b) -> PatAllocM lore a -> PatAllocM lore b)
-> (forall a b c.
    (a -> b -> c)
    -> PatAllocM lore a -> PatAllocM lore b -> PatAllocM lore c)
-> (forall a b.
    PatAllocM lore a -> PatAllocM lore b -> PatAllocM lore b)
-> (forall a b.
    PatAllocM lore a -> PatAllocM lore b -> PatAllocM lore a)
-> Applicative (PatAllocM lore)
PatAllocM lore a -> PatAllocM lore b -> PatAllocM lore b
PatAllocM lore a -> PatAllocM lore b -> PatAllocM lore a
PatAllocM lore (a -> b) -> PatAllocM lore a -> PatAllocM lore b
(a -> b -> c)
-> PatAllocM lore a -> PatAllocM lore b -> PatAllocM lore c
forall lore. Functor (PatAllocM lore)
forall a. a -> PatAllocM lore a
forall lore a. a -> PatAllocM lore a
forall a b.
PatAllocM lore a -> PatAllocM lore b -> PatAllocM lore a
forall a b.
PatAllocM lore a -> PatAllocM lore b -> PatAllocM lore b
forall a b.
PatAllocM lore (a -> b) -> PatAllocM lore a -> PatAllocM lore b
forall lore a b.
PatAllocM lore a -> PatAllocM lore b -> PatAllocM lore a
forall lore a b.
PatAllocM lore a -> PatAllocM lore b -> PatAllocM lore b
forall lore a b.
PatAllocM lore (a -> b) -> PatAllocM lore a -> PatAllocM lore b
forall a b c.
(a -> b -> c)
-> PatAllocM lore a -> PatAllocM lore b -> PatAllocM lore c
forall lore a b c.
(a -> b -> c)
-> PatAllocM lore a -> PatAllocM lore b -> PatAllocM lore c
forall (f :: * -> *).
Functor f
-> (forall a. a -> f a)
-> (forall a b. f (a -> b) -> f a -> f b)
-> (forall a b c. (a -> b -> c) -> f a -> f b -> f c)
-> (forall a b. f a -> f b -> f b)
-> (forall a b. f a -> f b -> f a)
-> Applicative f
<* :: PatAllocM lore a -> PatAllocM lore b -> PatAllocM lore a
$c<* :: forall lore a b.
PatAllocM lore a -> PatAllocM lore b -> PatAllocM lore a
*> :: PatAllocM lore a -> PatAllocM lore b -> PatAllocM lore b
$c*> :: forall lore a b.
PatAllocM lore a -> PatAllocM lore b -> PatAllocM lore b
liftA2 :: (a -> b -> c)
-> PatAllocM lore a -> PatAllocM lore b -> PatAllocM lore c
$cliftA2 :: forall lore a b c.
(a -> b -> c)
-> PatAllocM lore a -> PatAllocM lore b -> PatAllocM lore c
<*> :: PatAllocM lore (a -> b) -> PatAllocM lore a -> PatAllocM lore b
$c<*> :: forall lore a b.
PatAllocM lore (a -> b) -> PatAllocM lore a -> PatAllocM lore b
pure :: a -> PatAllocM lore a
$cpure :: forall lore a. a -> PatAllocM lore a
$cp1Applicative :: forall lore. Functor (PatAllocM lore)
Applicative, a -> PatAllocM lore b -> PatAllocM lore a
(a -> b) -> PatAllocM lore a -> PatAllocM lore b
(forall a b. (a -> b) -> PatAllocM lore a -> PatAllocM lore b)
-> (forall a b. a -> PatAllocM lore b -> PatAllocM lore a)
-> Functor (PatAllocM lore)
forall a b. a -> PatAllocM lore b -> PatAllocM lore a
forall a b. (a -> b) -> PatAllocM lore a -> PatAllocM lore b
forall lore a b. a -> PatAllocM lore b -> PatAllocM lore a
forall lore a b. (a -> b) -> PatAllocM lore a -> PatAllocM lore b
forall (f :: * -> *).
(forall a b. (a -> b) -> f a -> f b)
-> (forall a b. a -> f b -> f a) -> Functor f
<$ :: a -> PatAllocM lore b -> PatAllocM lore a
$c<$ :: forall lore a b. a -> PatAllocM lore b -> PatAllocM lore a
fmap :: (a -> b) -> PatAllocM lore a -> PatAllocM lore b
$cfmap :: forall lore a b. (a -> b) -> PatAllocM lore a -> PatAllocM lore b
Functor, Applicative (PatAllocM lore)
a -> PatAllocM lore a
Applicative (PatAllocM lore)
-> (forall a b.
    PatAllocM lore a -> (a -> PatAllocM lore b) -> PatAllocM lore b)
-> (forall a b.
    PatAllocM lore a -> PatAllocM lore b -> PatAllocM lore b)
-> (forall a. a -> PatAllocM lore a)
-> Monad (PatAllocM lore)
PatAllocM lore a -> (a -> PatAllocM lore b) -> PatAllocM lore b
PatAllocM lore a -> PatAllocM lore b -> PatAllocM lore b
forall lore. Applicative (PatAllocM lore)
forall a. a -> PatAllocM lore a
forall lore a. a -> PatAllocM lore a
forall a b.
PatAllocM lore a -> PatAllocM lore b -> PatAllocM lore b
forall a b.
PatAllocM lore a -> (a -> PatAllocM lore b) -> PatAllocM lore b
forall lore a b.
PatAllocM lore a -> PatAllocM lore b -> PatAllocM lore b
forall lore a b.
PatAllocM lore a -> (a -> PatAllocM lore b) -> PatAllocM lore b
forall (m :: * -> *).
Applicative m
-> (forall a b. m a -> (a -> m b) -> m b)
-> (forall a b. m a -> m b -> m b)
-> (forall a. a -> m a)
-> Monad m
return :: a -> PatAllocM lore a
$creturn :: forall lore a. a -> PatAllocM lore a
>> :: PatAllocM lore a -> PatAllocM lore b -> PatAllocM lore b
$c>> :: forall lore a b.
PatAllocM lore a -> PatAllocM lore b -> PatAllocM lore b
>>= :: PatAllocM lore a -> (a -> PatAllocM lore b) -> PatAllocM lore b
$c>>= :: forall lore a b.
PatAllocM lore a -> (a -> PatAllocM lore b) -> PatAllocM lore b
$cp1Monad :: forall lore. Applicative (PatAllocM lore)
Monad,
                              HasScope lore,
                              MonadWriter [AllocStm],
                              Monad (PatAllocM lore)
Applicative (PatAllocM lore)
PatAllocM lore VNameSource
Applicative (PatAllocM lore)
-> Monad (PatAllocM lore)
-> PatAllocM lore VNameSource
-> (VNameSource -> PatAllocM lore ())
-> MonadFreshNames (PatAllocM lore)
VNameSource -> PatAllocM lore ()
forall lore. Monad (PatAllocM lore)
forall lore. Applicative (PatAllocM lore)
forall lore. PatAllocM lore VNameSource
forall lore. VNameSource -> PatAllocM lore ()
forall (m :: * -> *).
Applicative m
-> Monad m
-> m VNameSource
-> (VNameSource -> m ())
-> MonadFreshNames m
putNameSource :: VNameSource -> PatAllocM lore ()
$cputNameSource :: forall lore. VNameSource -> PatAllocM lore ()
getNameSource :: PatAllocM lore VNameSource
$cgetNameSource :: forall lore. PatAllocM lore VNameSource
$cp2MonadFreshNames :: forall lore. Monad (PatAllocM lore)
$cp1MonadFreshNames :: forall lore. Applicative (PatAllocM lore)
MonadFreshNames)

instance Mem lore => Allocator lore (PatAllocM lore) where
  addAllocStm :: AllocStm -> PatAllocM lore ()
addAllocStm = [AllocStm] -> PatAllocM lore ()
forall w (m :: * -> *). MonadWriter w m => w -> m ()
tell ([AllocStm] -> PatAllocM lore ())
-> (AllocStm -> [AllocStm]) -> AllocStm -> PatAllocM lore ()
forall b c a. (b -> c) -> (a -> b) -> a -> c
. AllocStm -> [AllocStm]
forall (f :: * -> *) a. Applicative f => a -> f a
pure
  dimAllocationSize :: SubExp -> PatAllocM lore SubExp
dimAllocationSize = SubExp -> PatAllocM lore SubExp
forall (m :: * -> *) a. Monad m => a -> m a
return
  askDefaultSpace :: PatAllocM lore Space
askDefaultSpace = Space -> PatAllocM lore Space
forall (m :: * -> *) a. Monad m => a -> m a
return Space
DefaultSpace
  askConsts :: PatAllocM lore (Set VName)
askConsts = Set VName -> PatAllocM lore (Set VName)
forall (f :: * -> *) a. Applicative f => a -> f a
pure Set VName
forall a. Monoid a => a
mempty

runPatAllocM :: MonadFreshNames m =>
                PatAllocM lore a -> Scope lore
             -> m (a, [AllocStm])
runPatAllocM :: PatAllocM lore a -> Scope lore -> m (a, [AllocStm])
runPatAllocM (PatAllocM RWS (Scope lore) [AllocStm] VNameSource a
m) Scope lore
mems =
  (VNameSource -> ((a, [AllocStm]), VNameSource))
-> m (a, [AllocStm])
forall (m :: * -> *) a.
MonadFreshNames m =>
(VNameSource -> (a, VNameSource)) -> m a
modifyNameSource ((VNameSource -> ((a, [AllocStm]), VNameSource))
 -> m (a, [AllocStm]))
-> (VNameSource -> ((a, [AllocStm]), VNameSource))
-> m (a, [AllocStm])
forall a b. (a -> b) -> a -> b
$ (a, VNameSource, [AllocStm]) -> ((a, [AllocStm]), VNameSource)
forall a b b. (a, b, b) -> ((a, b), b)
frob ((a, VNameSource, [AllocStm]) -> ((a, [AllocStm]), VNameSource))
-> (VNameSource -> (a, VNameSource, [AllocStm]))
-> VNameSource
-> ((a, [AllocStm]), VNameSource)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. RWS (Scope lore) [AllocStm] VNameSource a
-> Scope lore -> VNameSource -> (a, VNameSource, [AllocStm])
forall r w s a. RWS r w s a -> r -> s -> (a, s, w)
runRWS RWS (Scope lore) [AllocStm] VNameSource a
m Scope lore
mems
  where frob :: (a, b, b) -> ((a, b), b)
frob (a
a,b
s,b
w) = ((a
a,b
w),b
s)

arraySizeInBytesExp :: Type -> PrimExp VName
arraySizeInBytesExp :: Type -> PrimExp VName
arraySizeInBytesExp Type
t =
  [PrimExp VName] -> PrimExp VName
forall (t :: * -> *) a. (Foldable t, Num a) => t a -> a
product
    [ PrimExp VName -> PrimExp VName
forall v. PrimExp v -> PrimExp v
toInt64 (PrimExp VName -> PrimExp VName) -> PrimExp VName -> PrimExp VName
forall a b. (a -> b) -> a -> b
$ [PrimExp VName] -> PrimExp VName
forall (t :: * -> *) a. (Foldable t, Num a) => t a -> a
product ([PrimExp VName] -> PrimExp VName)
-> [PrimExp VName] -> PrimExp VName
forall a b. (a -> b) -> a -> b
$ (SubExp -> PrimExp VName) -> Result -> [PrimExp VName]
forall a b. (a -> b) -> [a] -> [b]
map (PrimType -> SubExp -> PrimExp VName
primExpFromSubExp PrimType
int32) (Type -> Result
forall u. TypeBase Shape u -> Result
arrayDims Type
t)
    , PrimValue -> PrimExp VName
forall v. PrimValue -> PrimExp v
ValueExp (PrimValue -> PrimExp VName) -> PrimValue -> PrimExp VName
forall a b. (a -> b) -> a -> b
$ IntValue -> PrimValue
IntValue (IntValue -> PrimValue) -> IntValue -> PrimValue
forall a b. (a -> b) -> a -> b
$ Int64 -> IntValue
Int64Value (Int64 -> IntValue) -> Int64 -> IntValue
forall a b. (a -> b) -> a -> b
$ PrimType -> Int64
forall a. Num a => PrimType -> a
primByteSize (PrimType -> Int64) -> PrimType -> Int64
forall a b. (a -> b) -> a -> b
$ Type -> PrimType
forall shape u. TypeBase shape u -> PrimType
elemType Type
t ]
  where toInt64 :: PrimExp v -> PrimExp v
toInt64 = ConvOp -> PrimExp v -> PrimExp v
forall v. ConvOp -> PrimExp v -> PrimExp v
ConvOpExp (ConvOp -> PrimExp v -> PrimExp v)
-> ConvOp -> PrimExp v -> PrimExp v
forall a b. (a -> b) -> a -> b
$ IntType -> IntType -> ConvOp
SExt IntType
Int32 IntType
Int64

arraySizeInBytesExpM :: Allocator lore m => Type -> m (PrimExp VName)
arraySizeInBytesExpM :: Type -> m (PrimExp VName)
arraySizeInBytesExpM Type
t = do
  Result
dims <- (SubExp -> m SubExp) -> Result -> m Result
forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
(a -> m b) -> t a -> m (t b)
mapM SubExp -> m SubExp
forall lore (m :: * -> *). Allocator lore m => SubExp -> m SubExp
dimAllocationSize (Type -> Result
forall u. TypeBase Shape u -> Result
arrayDims Type
t)
  let dim_prod_i32 :: PrimExp VName
dim_prod_i32 = [PrimExp VName] -> PrimExp VName
forall (t :: * -> *) a. (Foldable t, Num a) => t a -> a
product ([PrimExp VName] -> PrimExp VName)
-> [PrimExp VName] -> PrimExp VName
forall a b. (a -> b) -> a -> b
$ (SubExp -> PrimExp VName) -> Result -> [PrimExp VName]
forall a b. (a -> b) -> [a] -> [b]
map (PrimExp VName -> PrimExp VName
forall v. PrimExp v -> PrimExp v
toInt64 (PrimExp VName -> PrimExp VName)
-> (SubExp -> PrimExp VName) -> SubExp -> PrimExp VName
forall b c a. (b -> c) -> (a -> b) -> a -> c
. PrimType -> SubExp -> PrimExp VName
primExpFromSubExp PrimType
int32) Result
dims
  let elm_size_i64 :: PrimExp VName
elm_size_i64 = PrimValue -> PrimExp VName
forall v. PrimValue -> PrimExp v
ValueExp (PrimValue -> PrimExp VName) -> PrimValue -> PrimExp VName
forall a b. (a -> b) -> a -> b
$ IntValue -> PrimValue
IntValue (IntValue -> PrimValue) -> IntValue -> PrimValue
forall a b. (a -> b) -> a -> b
$ Int64 -> IntValue
Int64Value (Int64 -> IntValue) -> Int64 -> IntValue
forall a b. (a -> b) -> a -> b
$ PrimType -> Int64
forall a. Num a => PrimType -> a
primByteSize (PrimType -> Int64) -> PrimType -> Int64
forall a b. (a -> b) -> a -> b
$ Type -> PrimType
forall shape u. TypeBase shape u -> PrimType
elemType Type
t
  PrimExp VName -> m (PrimExp VName)
forall (m :: * -> *) a. Monad m => a -> m a
return (PrimExp VName -> m (PrimExp VName))
-> PrimExp VName -> m (PrimExp VName)
forall a b. (a -> b) -> a -> b
$ [PrimExp VName] -> PrimExp VName
forall (t :: * -> *) a. (Foldable t, Num a) => t a -> a
product [ PrimExp VName
dim_prod_i32, PrimExp VName
elm_size_i64 ]
  where toInt64 :: PrimExp v -> PrimExp v
toInt64 = ConvOp -> PrimExp v -> PrimExp v
forall v. ConvOp -> PrimExp v -> PrimExp v
ConvOpExp (ConvOp -> PrimExp v -> PrimExp v)
-> ConvOp -> PrimExp v -> PrimExp v
forall a b. (a -> b) -> a -> b
$ IntType -> IntType -> ConvOp
SExt IntType
Int32 IntType
Int64

arraySizeInBytes :: Allocator lore m => Type -> m SubExp
arraySizeInBytes :: Type -> m SubExp
arraySizeInBytes = String -> PrimExp VName -> m SubExp
forall lore (m :: * -> *).
Allocator lore m =>
String -> PrimExp VName -> m SubExp
computeSize String
"bytes" (PrimExp VName -> m SubExp)
-> (Type -> m (PrimExp VName)) -> Type -> m SubExp
forall (m :: * -> *) b c a.
Monad m =>
(b -> m c) -> (a -> m b) -> a -> m c
<=< Type -> m (PrimExp VName)
forall lore (m :: * -> *).
Allocator lore m =>
Type -> m (PrimExp VName)
arraySizeInBytesExpM

-- | Allocate memory for a value of the given type.
allocForArray :: Allocator lore m =>
                 Type -> Space -> m VName
allocForArray :: Type -> Space -> m VName
allocForArray Type
t Space
space = do
  SubExp
size <- Type -> m SubExp
forall lore (m :: * -> *). Allocator lore m => Type -> m SubExp
arraySizeInBytes Type
t
  String -> SubExp -> Space -> m VName
forall lore (m :: * -> *).
Allocator lore m =>
String -> SubExp -> Space -> m VName
allocateMemory String
"mem" SubExp
size Space
space

allocsForStm :: (Allocator lore m, ExpAttr lore ~ ()) =>
                [Ident] -> [Ident] -> Exp lore
             -> m (Stm lore)
allocsForStm :: [Ident] -> [Ident] -> Exp lore -> m (Stm lore)
allocsForStm [Ident]
sizeidents [Ident]
validents Exp lore
e = do
  [ExpReturns]
rts <- Exp lore -> m [ExpReturns]
forall (m :: * -> *) lore.
(Monad m, HasScope lore m, Mem lore) =>
Exp lore -> m [ExpReturns]
expReturns Exp lore
e
  [ExpHint]
hints <- Exp lore -> m [ExpHint]
forall lore (m :: * -> *).
Allocator lore m =>
ExpT lore -> m [ExpHint]
expHints Exp lore
e
  ([PatElemT LetAttrMem]
ctxElems, [PatElemT LetAttrMem]
valElems) <- [Ident]
-> [Ident]
-> [ExpReturns]
-> [ExpHint]
-> m ([PatElem lore], [PatElem lore])
forall lore (m :: * -> *).
Allocator lore m =>
[Ident]
-> [Ident]
-> [ExpReturns]
-> [ExpHint]
-> m ([PatElem lore], [PatElem lore])
allocsForPattern [Ident]
sizeidents [Ident]
validents [ExpReturns]
rts [ExpHint]
hints
  Stm lore -> m (Stm lore)
forall (m :: * -> *) a. Monad m => a -> m a
return (Stm lore -> m (Stm lore)) -> Stm lore -> m (Stm lore)
forall a b. (a -> b) -> a -> b
$ Pattern lore -> StmAux (ExpAttr lore) -> Exp lore -> Stm lore
forall lore.
Pattern lore -> StmAux (ExpAttr lore) -> Exp lore -> Stm lore
Let ([PatElemT LetAttrMem]
-> [PatElemT LetAttrMem] -> PatternT LetAttrMem
forall attr. [PatElemT attr] -> [PatElemT attr] -> PatternT attr
Pattern [PatElemT LetAttrMem]
ctxElems [PatElemT LetAttrMem]
valElems) (() -> StmAux ()
forall attr. attr -> StmAux attr
defAux ()) Exp lore
e

patternWithAllocations :: (Allocator lore m, ExpAttr lore ~ ()) =>
                          [VName]
                       -> Exp lore
                       -> m (Pattern lore)
patternWithAllocations :: [VName] -> Exp lore -> m (Pattern lore)
patternWithAllocations [VName]
names Exp lore
e = do
  ([Type]
ts',[Ident]
sizes) <- [TypeBase ExtShape NoUniqueness] -> m ([Type], [Ident])
forall (m :: * -> *) u.
MonadFreshNames m =>
[TypeBase ExtShape u] -> m ([TypeBase Shape u], [Ident])
instantiateShapes' ([TypeBase ExtShape NoUniqueness] -> m ([Type], [Ident]))
-> m [TypeBase ExtShape NoUniqueness] -> m ([Type], [Ident])
forall (m :: * -> *) a b. Monad m => (a -> m b) -> m a -> m b
=<< Exp lore -> m [TypeBase ExtShape NoUniqueness]
forall lore (m :: * -> *).
(HasScope lore m, TypedOp (Op lore)) =>
Exp lore -> m [TypeBase ExtShape NoUniqueness]
expExtType Exp lore
e
  let identForBindage :: VName -> Type -> f Ident
identForBindage VName
name Type
t =
        Ident -> f Ident
forall (f :: * -> *) a. Applicative f => a -> f a
pure (Ident -> f Ident) -> Ident -> f Ident
forall a b. (a -> b) -> a -> b
$ VName -> Type -> Ident
Ident VName
name Type
t
  [Ident]
vals <- [m Ident] -> m [Ident]
forall (t :: * -> *) (m :: * -> *) a.
(Traversable t, Monad m) =>
t (m a) -> m (t a)
sequence [ VName -> Type -> m Ident
forall (f :: * -> *). Applicative f => VName -> Type -> f Ident
identForBindage VName
name Type
t | (VName
name, Type
t) <- [VName] -> [Type] -> [(VName, Type)]
forall a b. [a] -> [b] -> [(a, b)]
zip [VName]
names [Type]
ts' ]
  Stm lore -> PatternT LetAttrMem
forall lore. Stm lore -> Pattern lore
stmPattern (Stm lore -> PatternT LetAttrMem)
-> m (Stm lore) -> m (PatternT LetAttrMem)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> [Ident] -> [Ident] -> Exp lore -> m (Stm lore)
forall lore (m :: * -> *).
(Allocator lore m, ExpAttr lore ~ ()) =>
[Ident] -> [Ident] -> Exp lore -> m (Stm lore)
allocsForStm [Ident]
sizes [Ident]
vals Exp lore
e

allocsForPattern :: Allocator lore m =>
                    [Ident] -> [Ident] -> [ExpReturns] -> [ExpHint]
                 -> m ([PatElem lore],
                       [PatElem lore])
allocsForPattern :: [Ident]
-> [Ident]
-> [ExpReturns]
-> [ExpHint]
-> m ([PatElem lore], [PatElem lore])
allocsForPattern [Ident]
sizeidents [Ident]
validents [ExpReturns]
rts [ExpHint]
hints = do
  let sizes' :: [PatElemT LetAttrMem]
sizes' = [ VName -> LetAttrMem -> PatElemT LetAttrMem
forall attr. VName -> attr -> PatElemT attr
PatElem VName
size (LetAttrMem -> PatElemT LetAttrMem)
-> LetAttrMem -> PatElemT LetAttrMem
forall a b. (a -> b) -> a -> b
$ PrimType -> LetAttrMem
forall d u ret. PrimType -> MemInfo d u ret
MemPrim PrimType
int32 | VName
size <- (Ident -> VName) -> [Ident] -> [VName]
forall a b. (a -> b) -> [a] -> [b]
map Ident -> VName
identName [Ident]
sizeidents ]
  ([PatElemT LetAttrMem]
vals, ([PatElemT LetAttrMem]
exts, [PatElemT LetAttrMem]
mems)) <-
    WriterT
  ([PatElemT LetAttrMem], [PatElemT LetAttrMem])
  m
  [PatElemT LetAttrMem]
-> m ([PatElemT LetAttrMem],
      ([PatElemT LetAttrMem], [PatElemT LetAttrMem]))
forall w (m :: * -> *) a. WriterT w m a -> m (a, w)
runWriterT (WriterT
   ([PatElemT LetAttrMem], [PatElemT LetAttrMem])
   m
   [PatElemT LetAttrMem]
 -> m ([PatElemT LetAttrMem],
       ([PatElemT LetAttrMem], [PatElemT LetAttrMem])))
-> WriterT
     ([PatElemT LetAttrMem], [PatElemT LetAttrMem])
     m
     [PatElemT LetAttrMem]
-> m ([PatElemT LetAttrMem],
      ([PatElemT LetAttrMem], [PatElemT LetAttrMem]))
forall a b. (a -> b) -> a -> b
$ [(Ident, ExpReturns, ExpHint)]
-> ((Ident, ExpReturns, ExpHint)
    -> WriterT
         ([PatElemT LetAttrMem], [PatElemT LetAttrMem])
         m
         (PatElemT LetAttrMem))
-> WriterT
     ([PatElemT LetAttrMem], [PatElemT LetAttrMem])
     m
     [PatElemT LetAttrMem]
forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
t a -> (a -> m b) -> m (t b)
forM ([Ident]
-> [ExpReturns] -> [ExpHint] -> [(Ident, ExpReturns, ExpHint)]
forall a b c. [a] -> [b] -> [c] -> [(a, b, c)]
zip3 [Ident]
validents [ExpReturns]
rts [ExpHint]
hints) (((Ident, ExpReturns, ExpHint)
  -> WriterT
       ([PatElemT LetAttrMem], [PatElemT LetAttrMem])
       m
       (PatElemT LetAttrMem))
 -> WriterT
      ([PatElemT LetAttrMem], [PatElemT LetAttrMem])
      m
      [PatElemT LetAttrMem])
-> ((Ident, ExpReturns, ExpHint)
    -> WriterT
         ([PatElemT LetAttrMem], [PatElemT LetAttrMem])
         m
         (PatElemT LetAttrMem))
-> WriterT
     ([PatElemT LetAttrMem], [PatElemT LetAttrMem])
     m
     [PatElemT LetAttrMem]
forall a b. (a -> b) -> a -> b
$ \(Ident
ident, ExpReturns
rt, ExpHint
hint) -> do
      let shape :: Shape
shape = Type -> Shape
forall shape u. ArrayShape shape => TypeBase shape u -> shape
arrayShape (Type -> Shape) -> Type -> Shape
forall a b. (a -> b) -> a -> b
$ Ident -> Type
identType Ident
ident
      case ExpReturns
rt of
        MemPrim PrimType
_ -> do
          LetAttrMem
summary <- m LetAttrMem
-> WriterT
     ([PatElemT LetAttrMem], [PatElemT LetAttrMem]) m LetAttrMem
forall (t :: (* -> *) -> * -> *) (m :: * -> *) a.
(MonadTrans t, Monad m) =>
m a -> t m a
lift (m LetAttrMem
 -> WriterT
      ([PatElemT LetAttrMem], [PatElemT LetAttrMem]) m LetAttrMem)
-> m LetAttrMem
-> WriterT
     ([PatElemT LetAttrMem], [PatElemT LetAttrMem]) m LetAttrMem
forall a b. (a -> b) -> a -> b
$ Type -> ExpHint -> m LetAttrMem
forall lore (m :: * -> *).
Allocator lore m =>
Type -> ExpHint -> m LetAttrMem
summaryForBindage (Ident -> Type
identType Ident
ident) ExpHint
hint
          PatElemT LetAttrMem
-> WriterT
     ([PatElemT LetAttrMem], [PatElemT LetAttrMem])
     m
     (PatElemT LetAttrMem)
forall (m :: * -> *) a. Monad m => a -> m a
return (PatElemT LetAttrMem
 -> WriterT
      ([PatElemT LetAttrMem], [PatElemT LetAttrMem])
      m
      (PatElemT LetAttrMem))
-> PatElemT LetAttrMem
-> WriterT
     ([PatElemT LetAttrMem], [PatElemT LetAttrMem])
     m
     (PatElemT LetAttrMem)
forall a b. (a -> b) -> a -> b
$ VName -> LetAttrMem -> PatElemT LetAttrMem
forall attr. VName -> attr -> PatElemT attr
PatElem (Ident -> VName
identName Ident
ident) LetAttrMem
summary

        MemMem Space
space ->
          PatElemT LetAttrMem
-> WriterT
     ([PatElemT LetAttrMem], [PatElemT LetAttrMem])
     m
     (PatElemT LetAttrMem)
forall (m :: * -> *) a. Monad m => a -> m a
return (PatElemT LetAttrMem
 -> WriterT
      ([PatElemT LetAttrMem], [PatElemT LetAttrMem])
      m
      (PatElemT LetAttrMem))
-> PatElemT LetAttrMem
-> WriterT
     ([PatElemT LetAttrMem], [PatElemT LetAttrMem])
     m
     (PatElemT LetAttrMem)
forall a b. (a -> b) -> a -> b
$ VName -> LetAttrMem -> PatElemT LetAttrMem
forall attr. VName -> attr -> PatElemT attr
PatElem (Ident -> VName
identName Ident
ident) (LetAttrMem -> PatElemT LetAttrMem)
-> LetAttrMem -> PatElemT LetAttrMem
forall a b. (a -> b) -> a -> b
$
          Space -> LetAttrMem
forall d u ret. Space -> MemInfo d u ret
MemMem Space
space

        MemArray PrimType
bt ExtShape
_ NoUniqueness
u (Just (ReturnsInBlock VName
mem ExtIxFun
extixfun)) -> do
          ([PatElemT LetAttrMem]
patels, IxFun
ixfn) <- Ident
-> ExtIxFun
-> WriterT
     ([PatElemT LetAttrMem], [PatElemT LetAttrMem])
     m
     ([PatElemT LetAttrMem], IxFun)
forall (m :: * -> *) d u ret.
MonadFreshNames m =>
Ident -> ExtIxFun -> m ([PatElemT (MemInfo d u ret)], IxFun)
instantiateExtIxFun Ident
ident ExtIxFun
extixfun
          ([PatElemT LetAttrMem], [PatElemT LetAttrMem])
-> WriterT ([PatElemT LetAttrMem], [PatElemT LetAttrMem]) m ()
forall w (m :: * -> *). MonadWriter w m => w -> m ()
tell ([PatElemT LetAttrMem]
patels, [])

          PatElemT LetAttrMem
-> WriterT
     ([PatElemT LetAttrMem], [PatElemT LetAttrMem])
     m
     (PatElemT LetAttrMem)
forall (m :: * -> *) a. Monad m => a -> m a
return (PatElemT LetAttrMem
 -> WriterT
      ([PatElemT LetAttrMem], [PatElemT LetAttrMem])
      m
      (PatElemT LetAttrMem))
-> PatElemT LetAttrMem
-> WriterT
     ([PatElemT LetAttrMem], [PatElemT LetAttrMem])
     m
     (PatElemT LetAttrMem)
forall a b. (a -> b) -> a -> b
$ VName -> LetAttrMem -> PatElemT LetAttrMem
forall attr. VName -> attr -> PatElemT attr
PatElem (Ident -> VName
identName Ident
ident) (LetAttrMem -> PatElemT LetAttrMem)
-> LetAttrMem -> PatElemT LetAttrMem
forall a b. (a -> b) -> a -> b
$
            PrimType -> Shape -> NoUniqueness -> MemBind -> LetAttrMem
forall d u ret.
PrimType -> ShapeBase d -> u -> ret -> MemInfo d u ret
MemArray PrimType
bt Shape
shape NoUniqueness
u (MemBind -> LetAttrMem) -> MemBind -> LetAttrMem
forall a b. (a -> b) -> a -> b
$
            VName -> IxFun -> MemBind
ArrayIn VName
mem IxFun
ixfn

        MemArray PrimType
_ ExtShape
extshape NoUniqueness
_ Maybe MemReturn
Nothing
          | Just Result
_ <- ExtShape -> Maybe Result
forall b. ShapeBase (Ext b) -> Maybe [b]
knownShape ExtShape
extshape -> do
            LetAttrMem
summary <- m LetAttrMem
-> WriterT
     ([PatElemT LetAttrMem], [PatElemT LetAttrMem]) m LetAttrMem
forall (t :: (* -> *) -> * -> *) (m :: * -> *) a.
(MonadTrans t, Monad m) =>
m a -> t m a
lift (m LetAttrMem
 -> WriterT
      ([PatElemT LetAttrMem], [PatElemT LetAttrMem]) m LetAttrMem)
-> m LetAttrMem
-> WriterT
     ([PatElemT LetAttrMem], [PatElemT LetAttrMem]) m LetAttrMem
forall a b. (a -> b) -> a -> b
$ Type -> ExpHint -> m LetAttrMem
forall lore (m :: * -> *).
Allocator lore m =>
Type -> ExpHint -> m LetAttrMem
summaryForBindage (Ident -> Type
identType Ident
ident) ExpHint
hint
            PatElemT LetAttrMem
-> WriterT
     ([PatElemT LetAttrMem], [PatElemT LetAttrMem])
     m
     (PatElemT LetAttrMem)
forall (m :: * -> *) a. Monad m => a -> m a
return (PatElemT LetAttrMem
 -> WriterT
      ([PatElemT LetAttrMem], [PatElemT LetAttrMem])
      m
      (PatElemT LetAttrMem))
-> PatElemT LetAttrMem
-> WriterT
     ([PatElemT LetAttrMem], [PatElemT LetAttrMem])
     m
     (PatElemT LetAttrMem)
forall a b. (a -> b) -> a -> b
$ VName -> LetAttrMem -> PatElemT LetAttrMem
forall attr. VName -> attr -> PatElemT attr
PatElem (Ident -> VName
identName Ident
ident) LetAttrMem
summary

        MemArray PrimType
bt ExtShape
_ NoUniqueness
u (Just (ReturnsNewBlock Space
space Int
_ ExtIxFun
extixfn)) -> do
          -- treat existential index function first
          ([PatElemT LetAttrMem]
patels, IxFun
ixfn) <- Ident
-> ExtIxFun
-> WriterT
     ([PatElemT LetAttrMem], [PatElemT LetAttrMem])
     m
     ([PatElemT LetAttrMem], IxFun)
forall (m :: * -> *) d u ret.
MonadFreshNames m =>
Ident -> ExtIxFun -> m ([PatElemT (MemInfo d u ret)], IxFun)
instantiateExtIxFun Ident
ident ExtIxFun
extixfn
          ([PatElemT LetAttrMem], [PatElemT LetAttrMem])
-> WriterT ([PatElemT LetAttrMem], [PatElemT LetAttrMem]) m ()
forall w (m :: * -> *). MonadWriter w m => w -> m ()
tell ([PatElemT LetAttrMem]
patels, [])

          Ident
memid <- m Ident
-> WriterT ([PatElemT LetAttrMem], [PatElemT LetAttrMem]) m Ident
forall (t :: (* -> *) -> * -> *) (m :: * -> *) a.
(MonadTrans t, Monad m) =>
m a -> t m a
lift (m Ident
 -> WriterT ([PatElemT LetAttrMem], [PatElemT LetAttrMem]) m Ident)
-> m Ident
-> WriterT ([PatElemT LetAttrMem], [PatElemT LetAttrMem]) m Ident
forall a b. (a -> b) -> a -> b
$ Ident -> Space -> m Ident
forall (m :: * -> *).
MonadFreshNames m =>
Ident -> Space -> m Ident
mkMemIdent Ident
ident Space
space
          ([PatElemT LetAttrMem], [PatElemT LetAttrMem])
-> WriterT ([PatElemT LetAttrMem], [PatElemT LetAttrMem]) m ()
forall w (m :: * -> *). MonadWriter w m => w -> m ()
tell ([], [VName -> LetAttrMem -> PatElemT LetAttrMem
forall attr. VName -> attr -> PatElemT attr
PatElem (Ident -> VName
identName Ident
memid) (LetAttrMem -> PatElemT LetAttrMem)
-> LetAttrMem -> PatElemT LetAttrMem
forall a b. (a -> b) -> a -> b
$ Space -> LetAttrMem
forall d u ret. Space -> MemInfo d u ret
MemMem Space
space])
          PatElemT LetAttrMem
-> WriterT
     ([PatElemT LetAttrMem], [PatElemT LetAttrMem])
     m
     (PatElemT LetAttrMem)
forall (m :: * -> *) a. Monad m => a -> m a
return (PatElemT LetAttrMem
 -> WriterT
      ([PatElemT LetAttrMem], [PatElemT LetAttrMem])
      m
      (PatElemT LetAttrMem))
-> PatElemT LetAttrMem
-> WriterT
     ([PatElemT LetAttrMem], [PatElemT LetAttrMem])
     m
     (PatElemT LetAttrMem)
forall a b. (a -> b) -> a -> b
$ VName -> LetAttrMem -> PatElemT LetAttrMem
forall attr. VName -> attr -> PatElemT attr
PatElem (Ident -> VName
identName Ident
ident) (LetAttrMem -> PatElemT LetAttrMem)
-> LetAttrMem -> PatElemT LetAttrMem
forall a b. (a -> b) -> a -> b
$ PrimType -> Shape -> NoUniqueness -> MemBind -> LetAttrMem
forall d u ret.
PrimType -> ShapeBase d -> u -> ret -> MemInfo d u ret
MemArray PrimType
bt Shape
shape NoUniqueness
u (MemBind -> LetAttrMem) -> MemBind -> LetAttrMem
forall a b. (a -> b) -> a -> b
$
            VName -> IxFun -> MemBind
ArrayIn (Ident -> VName
identName Ident
memid) IxFun
ixfn

        ExpReturns
_ -> String
-> WriterT
     ([PatElemT LetAttrMem], [PatElemT LetAttrMem])
     m
     (PatElemT LetAttrMem)
forall a. HasCallStack => String -> a
error String
"Impossible case reached in allocsForPattern!"

  ([PatElemT LetAttrMem], [PatElemT LetAttrMem])
-> m ([PatElemT LetAttrMem], [PatElemT LetAttrMem])
forall (m :: * -> *) a. Monad m => a -> m a
return ([PatElemT LetAttrMem]
sizes' [PatElemT LetAttrMem]
-> [PatElemT LetAttrMem] -> [PatElemT LetAttrMem]
forall a. Semigroup a => a -> a -> a
<> [PatElemT LetAttrMem]
exts [PatElemT LetAttrMem]
-> [PatElemT LetAttrMem] -> [PatElemT LetAttrMem]
forall a. Semigroup a => a -> a -> a
<> [PatElemT LetAttrMem]
mems,
          [PatElemT LetAttrMem]
vals)
  where knownShape :: ShapeBase (Ext b) -> Maybe [b]
knownShape = (Ext b -> Maybe b) -> [Ext b] -> Maybe [b]
forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
(a -> m b) -> t a -> m (t b)
mapM Ext b -> Maybe b
forall a. Ext a -> Maybe a
known ([Ext b] -> Maybe [b])
-> (ShapeBase (Ext b) -> [Ext b]) -> ShapeBase (Ext b) -> Maybe [b]
forall b c a. (b -> c) -> (a -> b) -> a -> c
. ShapeBase (Ext b) -> [Ext b]
forall d. ShapeBase d -> [d]
shapeDims
        known :: Ext a -> Maybe a
known (Free a
v) = a -> Maybe a
forall a. a -> Maybe a
Just a
v
        known Ext{} = Maybe a
forall a. Maybe a
Nothing

        mkMemIdent :: (MonadFreshNames m) => Ident -> Space -> m Ident
        mkMemIdent :: Ident -> Space -> m Ident
mkMemIdent Ident
ident Space
space = do
          let memname :: String
memname = VName -> String
baseString (Ident -> VName
identName Ident
ident) String -> ShowS
forall a. Semigroup a => a -> a -> a
<> String
"_mem"
          String -> Type -> m Ident
forall (m :: * -> *).
MonadFreshNames m =>
String -> Type -> m Ident
newIdent String
memname (Type -> m Ident) -> Type -> m Ident
forall a b. (a -> b) -> a -> b
$ Space -> Type
forall shape u. Space -> TypeBase shape u
Mem Space
space

        instantiateExtIxFun :: MonadFreshNames m =>
                               Ident -> ExtIxFun ->
                               m ([PatElemT (MemInfo d u ret)], IxFun)
        instantiateExtIxFun :: Ident -> ExtIxFun -> m ([PatElemT (MemInfo d u ret)], IxFun)
instantiateExtIxFun Ident
idd ExtIxFun
ext_ixfn = do
          let isAndPtps :: [(Int, PrimType)]
isAndPtps = Set (Int, PrimType) -> [(Int, PrimType)]
forall a. Set a -> [a]
S.toList (Set (Int, PrimType) -> [(Int, PrimType)])
-> Set (Int, PrimType) -> [(Int, PrimType)]
forall a b. (a -> b) -> a -> b
$
                          ((Ext VName, PrimType) -> Set (Int, PrimType))
-> Set (Ext VName, PrimType) -> Set (Int, PrimType)
forall (t :: * -> *) m a.
(Foldable t, Monoid m) =>
(a -> m) -> t a -> m
foldMap (Ext VName, PrimType) -> Set (Int, PrimType)
forall a. (Ext a, PrimType) -> Set (Int, PrimType)
onlyExts (Set (Ext VName, PrimType) -> Set (Int, PrimType))
-> Set (Ext VName, PrimType) -> Set (Int, PrimType)
forall a b. (a -> b) -> a -> b
$
                          (PrimExp (Ext VName) -> Set (Ext VName, PrimType))
-> ExtIxFun -> Set (Ext VName, PrimType)
forall (t :: * -> *) m a.
(Foldable t, Monoid m) =>
(a -> m) -> t a -> m
foldMap PrimExp (Ext VName) -> Set (Ext VName, PrimType)
forall a. Ord a => PrimExp a -> Set (a, PrimType)
leafExpTypes ExtIxFun
ext_ixfn

          -- Find the existentials that reuse the sizeidents, and
          -- those that need new pattern elements.  Assumes that the
          -- Exts form a contiguous interval of integers.
          let ([(Int, PrimType)]
size_exts, [(Int, PrimType)]
new_exts) =
                ((Int, PrimType) -> Bool)
-> [(Int, PrimType)] -> ([(Int, PrimType)], [(Int, PrimType)])
forall a. (a -> Bool) -> [a] -> ([a], [a])
span ((Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
<[Ident] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length [Ident]
sizeidents) (Int -> Bool)
-> ((Int, PrimType) -> Int) -> (Int, PrimType) -> Bool
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (Int, PrimType) -> Int
forall a b. (a, b) -> a
fst) ([(Int, PrimType)] -> ([(Int, PrimType)], [(Int, PrimType)]))
-> [(Int, PrimType)] -> ([(Int, PrimType)], [(Int, PrimType)])
forall a b. (a -> b) -> a -> b
$ [(Int, PrimType)] -> [(Int, PrimType)]
forall a. Ord a => [a] -> [a]
sort [(Int, PrimType)]
isAndPtps
          ([(Ext VName, PrimExp (Ext VName))]
new_substs, [PatElemT (MemInfo d u ret)]
patels) <-
            ([((Ext VName, PrimExp (Ext VName)), PatElemT (MemInfo d u ret))]
 -> ([(Ext VName, PrimExp (Ext VName))],
     [PatElemT (MemInfo d u ret)]))
-> m [((Ext VName, PrimExp (Ext VName)),
       PatElemT (MemInfo d u ret))]
-> m ([(Ext VName, PrimExp (Ext VName))],
      [PatElemT (MemInfo d u ret)])
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap [((Ext VName, PrimExp (Ext VName)), PatElemT (MemInfo d u ret))]
-> ([(Ext VName, PrimExp (Ext VName))],
    [PatElemT (MemInfo d u ret)])
forall a b. [(a, b)] -> ([a], [b])
unzip (m [((Ext VName, PrimExp (Ext VName)), PatElemT (MemInfo d u ret))]
 -> m ([(Ext VName, PrimExp (Ext VName))],
       [PatElemT (MemInfo d u ret)]))
-> m [((Ext VName, PrimExp (Ext VName)),
       PatElemT (MemInfo d u ret))]
-> m ([(Ext VName, PrimExp (Ext VName))],
      [PatElemT (MemInfo d u ret)])
forall a b. (a -> b) -> a -> b
$ [(Int, PrimType)]
-> ((Int, PrimType)
    -> m ((Ext VName, PrimExp (Ext VName)),
          PatElemT (MemInfo d u ret)))
-> m [((Ext VName, PrimExp (Ext VName)),
       PatElemT (MemInfo d u ret))]
forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
t a -> (a -> m b) -> m (t b)
forM [(Int, PrimType)]
new_exts (((Int, PrimType)
  -> m ((Ext VName, PrimExp (Ext VName)),
        PatElemT (MemInfo d u ret)))
 -> m [((Ext VName, PrimExp (Ext VName)),
        PatElemT (MemInfo d u ret))])
-> ((Int, PrimType)
    -> m ((Ext VName, PrimExp (Ext VName)),
          PatElemT (MemInfo d u ret)))
-> m [((Ext VName, PrimExp (Ext VName)),
       PatElemT (MemInfo d u ret))]
forall a b. (a -> b) -> a -> b
$ \(Int
i, PrimType
t) -> do
            VName
v <- String -> m VName
forall (m :: * -> *). MonadFreshNames m => String -> m VName
newVName (String -> m VName) -> String -> m VName
forall a b. (a -> b) -> a -> b
$ VName -> String
baseString (Ident -> VName
identName Ident
idd) String -> ShowS
forall a. Semigroup a => a -> a -> a
<> String
"_ixfn"
            ((Ext VName, PrimExp (Ext VName)), PatElemT (MemInfo d u ret))
-> m ((Ext VName, PrimExp (Ext VName)), PatElemT (MemInfo d u ret))
forall (m :: * -> *) a. Monad m => a -> m a
return ((Int -> Ext VName
forall a. Int -> Ext a
Ext Int
i, Ext VName -> PrimType -> PrimExp (Ext VName)
forall v. v -> PrimType -> PrimExp v
LeafExp (VName -> Ext VName
forall a. a -> Ext a
Free VName
v) PrimType
t),
                    VName -> MemInfo d u ret -> PatElemT (MemInfo d u ret)
forall attr. VName -> attr -> PatElemT attr
PatElem VName
v (MemInfo d u ret -> PatElemT (MemInfo d u ret))
-> MemInfo d u ret -> PatElemT (MemInfo d u ret)
forall a b. (a -> b) -> a -> b
$ PrimType -> MemInfo d u ret
forall d u ret. PrimType -> MemInfo d u ret
MemPrim PrimType
t)
          let size_substs :: [(Ext VName, PrimExp (Ext VName))]
size_substs = ((Int, PrimType) -> Ident -> (Ext VName, PrimExp (Ext VName)))
-> [(Int, PrimType)]
-> [Ident]
-> [(Ext VName, PrimExp (Ext VName))]
forall a b c. (a -> b -> c) -> [a] -> [b] -> [c]
zipWith (\(Int
i, PrimType
t) Ident
ident ->
                                    (Int -> Ext VName
forall a. Int -> Ext a
Ext Int
i, Ext VName -> PrimType -> PrimExp (Ext VName)
forall v. v -> PrimType -> PrimExp v
LeafExp (VName -> Ext VName
forall a. a -> Ext a
Free (Ident -> VName
identName Ident
ident)) PrimType
t))
                            [(Int, PrimType)]
size_exts [Ident]
sizeidents
              substs :: Map (Ext VName) (PrimExp (Ext VName))
substs = [(Ext VName, PrimExp (Ext VName))]
-> Map (Ext VName) (PrimExp (Ext VName))
forall k a. Ord k => [(k, a)] -> Map k a
M.fromList ([(Ext VName, PrimExp (Ext VName))]
 -> Map (Ext VName) (PrimExp (Ext VName)))
-> [(Ext VName, PrimExp (Ext VName))]
-> Map (Ext VName) (PrimExp (Ext VName))
forall a b. (a -> b) -> a -> b
$ [(Ext VName, PrimExp (Ext VName))]
new_substs [(Ext VName, PrimExp (Ext VName))]
-> [(Ext VName, PrimExp (Ext VName))]
-> [(Ext VName, PrimExp (Ext VName))]
forall a. Semigroup a => a -> a -> a
<> [(Ext VName, PrimExp (Ext VName))]
size_substs
          IxFun
ixfn <- ExtIxFun -> m IxFun
forall (m :: * -> *). Monad m => ExtIxFun -> m IxFun
instantiateIxFun (ExtIxFun -> m IxFun) -> ExtIxFun -> m IxFun
forall a b. (a -> b) -> a -> b
$ Map (Ext VName) (PrimExp (Ext VName)) -> ExtIxFun -> ExtIxFun
forall a.
Ord a =>
Map a (PrimExp a) -> IxFun (PrimExp a) -> IxFun (PrimExp a)
IxFun.substituteInIxFun Map (Ext VName) (PrimExp (Ext VName))
substs ExtIxFun
ext_ixfn

          ([PatElemT (MemInfo d u ret)], IxFun)
-> m ([PatElemT (MemInfo d u ret)], IxFun)
forall (m :: * -> *) a. Monad m => a -> m a
return ([PatElemT (MemInfo d u ret)]
patels, IxFun
ixfn)

onlyExts :: (Ext a, PrimType) -> S.Set (Int, PrimType)
onlyExts :: (Ext a, PrimType) -> Set (Int, PrimType)
onlyExts (Free a
_, PrimType
_) = Set (Int, PrimType)
forall a. Set a
S.empty
onlyExts (Ext Int
i, PrimType
t) = (Int, PrimType) -> Set (Int, PrimType)
forall a. a -> Set a
S.singleton (Int
i, PrimType
t)


instantiateIxFun :: Monad m => ExtIxFun -> m IxFun
instantiateIxFun :: ExtIxFun -> m IxFun
instantiateIxFun = (PrimExp (Ext VName) -> m (PrimExp VName)) -> ExtIxFun -> m IxFun
forall (t :: * -> *) (f :: * -> *) a b.
(Traversable t, Applicative f) =>
(a -> f b) -> t a -> f (t b)
traverse ((PrimExp (Ext VName) -> m (PrimExp VName)) -> ExtIxFun -> m IxFun)
-> (PrimExp (Ext VName) -> m (PrimExp VName))
-> ExtIxFun
-> m IxFun
forall a b. (a -> b) -> a -> b
$ (Ext VName -> m VName) -> PrimExp (Ext VName) -> m (PrimExp VName)
forall (t :: * -> *) (f :: * -> *) a b.
(Traversable t, Applicative f) =>
(a -> f b) -> t a -> f (t b)
traverse Ext VName -> m VName
forall (m :: * -> *) a. Monad m => Ext a -> m a
inst
  where inst :: Ext a -> m a
inst Ext{} = String -> m a
forall a. HasCallStack => String -> a
error String
"instantiateIxFun: not yet"
        inst (Free a
x) = a -> m a
forall (m :: * -> *) a. Monad m => a -> m a
return a
x

summaryForBindage :: Allocator lore m =>
                     Type -> ExpHint
                  -> m (MemBound NoUniqueness)
summaryForBindage :: Type -> ExpHint -> m LetAttrMem
summaryForBindage (Prim PrimType
bt) ExpHint
_ =
  LetAttrMem -> m LetAttrMem
forall (m :: * -> *) a. Monad m => a -> m a
return (LetAttrMem -> m LetAttrMem) -> LetAttrMem -> m LetAttrMem
forall a b. (a -> b) -> a -> b
$ PrimType -> LetAttrMem
forall d u ret. PrimType -> MemInfo d u ret
MemPrim PrimType
bt
summaryForBindage (Mem Space
space) ExpHint
_ =
  LetAttrMem -> m LetAttrMem
forall (m :: * -> *) a. Monad m => a -> m a
return (LetAttrMem -> m LetAttrMem) -> LetAttrMem -> m LetAttrMem
forall a b. (a -> b) -> a -> b
$ Space -> LetAttrMem
forall d u ret. Space -> MemInfo d u ret
MemMem Space
space
summaryForBindage t :: Type
t@(Array PrimType
bt Shape
shape NoUniqueness
u) ExpHint
NoHint = do
  VName
m <- Type -> Space -> m VName
forall lore (m :: * -> *).
Allocator lore m =>
Type -> Space -> m VName
allocForArray Type
t (Space -> m VName) -> m Space -> m VName
forall (m :: * -> *) a b. Monad m => (a -> m b) -> m a -> m b
=<< m Space
forall lore (m :: * -> *). Allocator lore m => m Space
askDefaultSpace
  LetAttrMem -> m LetAttrMem
forall (m :: * -> *) a. Monad m => a -> m a
return (LetAttrMem -> m LetAttrMem) -> LetAttrMem -> m LetAttrMem
forall a b. (a -> b) -> a -> b
$ PrimType -> Shape -> NoUniqueness -> VName -> Type -> LetAttrMem
forall u. PrimType -> Shape -> u -> VName -> Type -> MemBound u
directIxFun PrimType
bt Shape
shape NoUniqueness
u VName
m Type
t
summaryForBindage Type
t (Hint IxFun
ixfun Space
space) = do
  let bt :: PrimType
bt = Type -> PrimType
forall shape u. TypeBase shape u -> PrimType
elemType Type
t
  SubExp
bytes <- String -> PrimExp VName -> m SubExp
forall lore (m :: * -> *).
Allocator lore m =>
String -> PrimExp VName -> m SubExp
computeSize String
"bytes" (PrimExp VName -> m SubExp) -> PrimExp VName -> m SubExp
forall a b. (a -> b) -> a -> b
$
           [PrimExp VName] -> PrimExp VName
forall (t :: * -> *) a. (Foldable t, Num a) => t a -> a
product [ConvOp -> PrimExp VName -> PrimExp VName
forall v. ConvOp -> PrimExp v -> PrimExp v
ConvOpExp (IntType -> IntType -> ConvOp
SExt IntType
Int32 IntType
Int64) ([PrimExp VName] -> PrimExp VName
forall (t :: * -> *) a. (Foldable t, Num a) => t a -> a
product (IxFun -> [PrimExp VName]
forall num. IxFun num -> Shape num
IxFun.base IxFun
ixfun)),
                    Int64 -> PrimExp VName
forall a b. (Integral a, Num b) => a -> b
fromIntegral (PrimType -> Int64
forall a. Num a => PrimType -> a
primByteSize (Type -> PrimType
forall shape u. TypeBase shape u -> PrimType
elemType Type
t)::Int64)]
  VName
m <- String -> SubExp -> Space -> m VName
forall lore (m :: * -> *).
Allocator lore m =>
String -> SubExp -> Space -> m VName
allocateMemory String
"mem" SubExp
bytes Space
space
  LetAttrMem -> m LetAttrMem
forall (m :: * -> *) a. Monad m => a -> m a
return (LetAttrMem -> m LetAttrMem) -> LetAttrMem -> m LetAttrMem
forall a b. (a -> b) -> a -> b
$ PrimType -> Shape -> NoUniqueness -> MemBind -> LetAttrMem
forall d u ret.
PrimType -> ShapeBase d -> u -> ret -> MemInfo d u ret
MemArray PrimType
bt (Type -> Shape
forall shape u. ArrayShape shape => TypeBase shape u -> shape
arrayShape Type
t) NoUniqueness
NoUniqueness (MemBind -> LetAttrMem) -> MemBind -> LetAttrMem
forall a b. (a -> b) -> a -> b
$ VName -> IxFun -> MemBind
ArrayIn VName
m IxFun
ixfun

lookupMemSpace :: (HasScope lore m, Monad m) => VName -> m Space
lookupMemSpace :: VName -> m Space
lookupMemSpace VName
v = do
  Type
t <- VName -> m Type
forall lore (m :: * -> *). HasScope lore m => VName -> m Type
lookupType VName
v
  case Type
t of
    Mem Space
space -> Space -> m Space
forall (m :: * -> *) a. Monad m => a -> m a
return Space
space
    Type
_ -> String -> m Space
forall a. HasCallStack => String -> a
error (String -> m Space) -> String -> m Space
forall a b. (a -> b) -> a -> b
$ String
"lookupMemSpace: " String -> ShowS
forall a. [a] -> [a] -> [a]
++ VName -> String
forall a. Pretty a => a -> String
pretty VName
v String -> ShowS
forall a. [a] -> [a] -> [a]
++ String
" is not a memory block."

directIxFun :: PrimType -> Shape -> u -> VName -> Type -> MemBound u
directIxFun :: PrimType -> Shape -> u -> VName -> Type -> MemBound u
directIxFun PrimType
bt Shape
shape u
u VName
mem Type
t =
  PrimType -> Shape -> u -> MemBind -> MemBound u
forall d u ret.
PrimType -> ShapeBase d -> u -> ret -> MemInfo d u ret
MemArray PrimType
bt Shape
shape u
u (MemBind -> MemBound u) -> MemBind -> MemBound u
forall a b. (a -> b) -> a -> b
$ VName -> IxFun -> MemBind
ArrayIn VName
mem (IxFun -> MemBind) -> IxFun -> MemBind
forall a b. (a -> b) -> a -> b
$
  [PrimExp VName] -> IxFun
forall num. IntegralExp num => Shape num -> IxFun num
IxFun.iota ([PrimExp VName] -> IxFun) -> [PrimExp VName] -> IxFun
forall a b. (a -> b) -> a -> b
$ (SubExp -> PrimExp VName) -> Result -> [PrimExp VName]
forall a b. (a -> b) -> [a] -> [b]
map (PrimType -> SubExp -> PrimExp VName
primExpFromSubExp PrimType
int32) (Result -> [PrimExp VName]) -> Result -> [PrimExp VName]
forall a b. (a -> b) -> a -> b
$ Type -> Result
forall u. TypeBase Shape u -> Result
arrayDims Type
t

allocInFParams :: (Allocable fromlore tolore) =>
                  [(FParam fromlore, Space)] ->
                  ([FParam tolore] -> AllocM fromlore tolore a)
               -> AllocM fromlore tolore a
allocInFParams :: [(FParam fromlore, Space)]
-> ([FParam tolore] -> AllocM fromlore tolore a)
-> AllocM fromlore tolore a
allocInFParams [(FParam fromlore, Space)]
params [FParam tolore] -> AllocM fromlore tolore a
m = do
  ([Param FParamMem]
valparams, [Param FParamMem]
memparams) <-
    WriterT
  [Param FParamMem] (AllocM fromlore tolore) [Param FParamMem]
-> AllocM fromlore tolore ([Param FParamMem], [Param FParamMem])
forall w (m :: * -> *) a. WriterT w m a -> m (a, w)
runWriterT (WriterT
   [Param FParamMem] (AllocM fromlore tolore) [Param FParamMem]
 -> AllocM fromlore tolore ([Param FParamMem], [Param FParamMem]))
-> WriterT
     [Param FParamMem] (AllocM fromlore tolore) [Param FParamMem]
-> AllocM fromlore tolore ([Param FParamMem], [Param FParamMem])
forall a b. (a -> b) -> a -> b
$ ((Param DeclType, Space)
 -> WriterT
      [Param FParamMem] (AllocM fromlore tolore) (Param FParamMem))
-> [(Param DeclType, Space)]
-> WriterT
     [Param FParamMem] (AllocM fromlore tolore) [Param FParamMem]
forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
(a -> m b) -> t a -> m (t b)
mapM ((Param DeclType
 -> Space
 -> WriterT
      [Param FParamMem] (AllocM fromlore tolore) (Param FParamMem))
-> (Param DeclType, Space)
-> WriterT
     [Param FParamMem] (AllocM fromlore tolore) (Param FParamMem)
forall a b c. (a -> b -> c) -> (a, b) -> c
uncurry Param DeclType
-> Space
-> WriterT
     [Param FParamMem] (AllocM fromlore tolore) (Param FParamMem)
forall fromlore tolore.
Allocable fromlore tolore =>
FParam fromlore
-> Space
-> WriterT [FParam tolore] (AllocM fromlore tolore) (FParam tolore)
allocInFParam) [(Param DeclType, Space)]
[(FParam fromlore, Space)]
params
  let params' :: [Param FParamMem]
params' = [Param FParamMem]
memparams [Param FParamMem] -> [Param FParamMem] -> [Param FParamMem]
forall a. Semigroup a => a -> a -> a
<> [Param FParamMem]
valparams
      summary :: Scope tolore
summary = [Param FParamMem] -> Scope tolore
forall lore attr.
(FParamAttr lore ~ attr) =>
[Param attr] -> Scope lore
scopeOfFParams [Param FParamMem]
params'
  Scope tolore
-> AllocM fromlore tolore a -> AllocM fromlore tolore a
forall lore (m :: * -> *) a.
LocalScope lore m =>
Scope lore -> m a -> m a
localScope Scope tolore
summary (AllocM fromlore tolore a -> AllocM fromlore tolore a)
-> AllocM fromlore tolore a -> AllocM fromlore tolore a
forall a b. (a -> b) -> a -> b
$ [FParam tolore] -> AllocM fromlore tolore a
m [FParam tolore]
[Param FParamMem]
params'

allocInFParam :: (Allocable fromlore tolore) =>
                 FParam fromlore
              -> Space
              -> WriterT [FParam tolore]
                 (AllocM fromlore tolore) (FParam tolore)
allocInFParam :: FParam fromlore
-> Space
-> WriterT [FParam tolore] (AllocM fromlore tolore) (FParam tolore)
allocInFParam FParam fromlore
param Space
pspace =
  case Param DeclType -> DeclType
forall attr. DeclTyped attr => Param attr -> DeclType
paramDeclType Param DeclType
FParam fromlore
param of
    Array PrimType
bt Shape
shape Uniqueness
u -> do
      let memname :: String
memname = VName -> String
baseString (Param DeclType -> VName
forall attr. Param attr -> VName
paramName Param DeclType
FParam fromlore
param) String -> ShowS
forall a. Semigroup a => a -> a -> a
<> String
"_mem"
          ixfun :: IxFun
ixfun = [PrimExp VName] -> IxFun
forall num. IntegralExp num => Shape num -> IxFun num
IxFun.iota ([PrimExp VName] -> IxFun) -> [PrimExp VName] -> IxFun
forall a b. (a -> b) -> a -> b
$ (SubExp -> PrimExp VName) -> Result -> [PrimExp VName]
forall a b. (a -> b) -> [a] -> [b]
map (PrimType -> SubExp -> PrimExp VName
primExpFromSubExp PrimType
int32) (Result -> [PrimExp VName]) -> Result -> [PrimExp VName]
forall a b. (a -> b) -> a -> b
$ Shape -> Result
forall d. ShapeBase d -> [d]
shapeDims Shape
shape
      VName
mem <- AllocM fromlore tolore VName
-> WriterT [Param FParamMem] (AllocM fromlore tolore) VName
forall (t :: (* -> *) -> * -> *) (m :: * -> *) a.
(MonadTrans t, Monad m) =>
m a -> t m a
lift (AllocM fromlore tolore VName
 -> WriterT [Param FParamMem] (AllocM fromlore tolore) VName)
-> AllocM fromlore tolore VName
-> WriterT [Param FParamMem] (AllocM fromlore tolore) VName
forall a b. (a -> b) -> a -> b
$ String -> AllocM fromlore tolore VName
forall (m :: * -> *). MonadFreshNames m => String -> m VName
newVName String
memname
      [Param FParamMem]
-> WriterT [Param FParamMem] (AllocM fromlore tolore) ()
forall w (m :: * -> *). MonadWriter w m => w -> m ()
tell [VName -> FParamMem -> Param FParamMem
forall attr. VName -> attr -> Param attr
Param VName
mem (FParamMem -> Param FParamMem) -> FParamMem -> Param FParamMem
forall a b. (a -> b) -> a -> b
$ Space -> FParamMem
forall d u ret. Space -> MemInfo d u ret
MemMem Space
pspace]
      Param FParamMem
-> WriterT
     [Param FParamMem] (AllocM fromlore tolore) (Param FParamMem)
forall (m :: * -> *) a. Monad m => a -> m a
return Param DeclType
FParam fromlore
param { paramAttr :: FParamMem
paramAttr =  PrimType -> Shape -> Uniqueness -> MemBind -> FParamMem
forall d u ret.
PrimType -> ShapeBase d -> u -> ret -> MemInfo d u ret
MemArray PrimType
bt Shape
shape Uniqueness
u (MemBind -> FParamMem) -> MemBind -> FParamMem
forall a b. (a -> b) -> a -> b
$ VName -> IxFun -> MemBind
ArrayIn VName
mem IxFun
ixfun }
    Prim PrimType
bt ->
      Param FParamMem
-> WriterT
     [Param FParamMem] (AllocM fromlore tolore) (Param FParamMem)
forall (m :: * -> *) a. Monad m => a -> m a
return Param DeclType
FParam fromlore
param { paramAttr :: FParamMem
paramAttr = PrimType -> FParamMem
forall d u ret. PrimType -> MemInfo d u ret
MemPrim PrimType
bt }
    Mem Space
space ->
      Param FParamMem
-> WriterT
     [Param FParamMem] (AllocM fromlore tolore) (Param FParamMem)
forall (m :: * -> *) a. Monad m => a -> m a
return Param DeclType
FParam fromlore
param { paramAttr :: FParamMem
paramAttr = Space -> FParamMem
forall d u ret. Space -> MemInfo d u ret
MemMem Space
space }

allocInMergeParams :: (Allocable fromlore tolore,
                       Allocator tolore (AllocM fromlore tolore)) =>
                      [VName]
                   -> [(FParam fromlore,SubExp)]
                   -> ([FParam tolore]
                       -> [FParam tolore]
                       -> ([SubExp] -> AllocM fromlore tolore ([SubExp], [SubExp]))
                       -> AllocM fromlore tolore a)
                   -> AllocM fromlore tolore a
allocInMergeParams :: [VName]
-> [(FParam fromlore, SubExp)]
-> ([FParam tolore]
    -> [FParam tolore]
    -> (Result -> AllocM fromlore tolore (Result, Result))
    -> AllocM fromlore tolore a)
-> AllocM fromlore tolore a
allocInMergeParams [VName]
variant [(FParam fromlore, SubExp)]
merge [FParam tolore]
-> [FParam tolore]
-> (Result -> AllocM fromlore tolore (Result, Result))
-> AllocM fromlore tolore a
m = do
  (([Param FParamMem]
valparams, [SubExp -> WriterT Result (AllocM fromlore tolore) SubExp]
handle_loop_subexps), [Param FParamMem]
mem_params) <-
    WriterT
  [Param FParamMem]
  (AllocM fromlore tolore)
  ([Param FParamMem],
   [SubExp -> WriterT Result (AllocM fromlore tolore) SubExp])
-> AllocM
     fromlore
     tolore
     (([Param FParamMem],
       [SubExp -> WriterT Result (AllocM fromlore tolore) SubExp]),
      [Param FParamMem])
forall w (m :: * -> *) a. WriterT w m a -> m (a, w)
runWriterT (WriterT
   [Param FParamMem]
   (AllocM fromlore tolore)
   ([Param FParamMem],
    [SubExp -> WriterT Result (AllocM fromlore tolore) SubExp])
 -> AllocM
      fromlore
      tolore
      (([Param FParamMem],
        [SubExp -> WriterT Result (AllocM fromlore tolore) SubExp]),
       [Param FParamMem]))
-> WriterT
     [Param FParamMem]
     (AllocM fromlore tolore)
     ([Param FParamMem],
      [SubExp -> WriterT Result (AllocM fromlore tolore) SubExp])
-> AllocM
     fromlore
     tolore
     (([Param FParamMem],
       [SubExp -> WriterT Result (AllocM fromlore tolore) SubExp]),
      [Param FParamMem])
forall a b. (a -> b) -> a -> b
$ [(Param FParamMem,
  SubExp -> WriterT Result (AllocM fromlore tolore) SubExp)]
-> ([Param FParamMem],
    [SubExp -> WriterT Result (AllocM fromlore tolore) SubExp])
forall a b. [(a, b)] -> ([a], [b])
unzip ([(Param FParamMem,
   SubExp -> WriterT Result (AllocM fromlore tolore) SubExp)]
 -> ([Param FParamMem],
     [SubExp -> WriterT Result (AllocM fromlore tolore) SubExp]))
-> WriterT
     [Param FParamMem]
     (AllocM fromlore tolore)
     [(Param FParamMem,
       SubExp -> WriterT Result (AllocM fromlore tolore) SubExp)]
-> WriterT
     [Param FParamMem]
     (AllocM fromlore tolore)
     ([Param FParamMem],
      [SubExp -> WriterT Result (AllocM fromlore tolore) SubExp])
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> ((Param DeclType, SubExp)
 -> WriterT
      [Param FParamMem]
      (AllocM fromlore tolore)
      (Param FParamMem,
       SubExp -> WriterT Result (AllocM fromlore tolore) SubExp))
-> [(Param DeclType, SubExp)]
-> WriterT
     [Param FParamMem]
     (AllocM fromlore tolore)
     [(Param FParamMem,
       SubExp -> WriterT Result (AllocM fromlore tolore) SubExp)]
forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
(a -> m b) -> t a -> m (t b)
mapM (Param DeclType, SubExp)
-> WriterT
     [Param FParamMem]
     (AllocM fromlore tolore)
     (Param FParamMem,
      SubExp -> WriterT Result (AllocM fromlore tolore) SubExp)
allocInMergeParam [(Param DeclType, SubExp)]
[(FParam fromlore, SubExp)]
merge
  let mergeparams' :: [Param FParamMem]
mergeparams' = [Param FParamMem]
mem_params [Param FParamMem] -> [Param FParamMem] -> [Param FParamMem]
forall a. Semigroup a => a -> a -> a
<> [Param FParamMem]
valparams
      summary :: Scope tolore
summary = [Param FParamMem] -> Scope tolore
forall lore attr.
(FParamAttr lore ~ attr) =>
[Param attr] -> Scope lore
scopeOfFParams [Param FParamMem]
mergeparams'

      mk_loop_res :: Result -> AllocM fromlore tolore (Result, Result)
mk_loop_res Result
ses = do
        (Result
valargs, Result
memargs) <-
          WriterT Result (AllocM fromlore tolore) Result
-> AllocM fromlore tolore (Result, Result)
forall w (m :: * -> *) a. WriterT w m a -> m (a, w)
runWriterT (WriterT Result (AllocM fromlore tolore) Result
 -> AllocM fromlore tolore (Result, Result))
-> WriterT Result (AllocM fromlore tolore) Result
-> AllocM fromlore tolore (Result, Result)
forall a b. (a -> b) -> a -> b
$ ((SubExp -> WriterT Result (AllocM fromlore tolore) SubExp)
 -> SubExp -> WriterT Result (AllocM fromlore tolore) SubExp)
-> [SubExp -> WriterT Result (AllocM fromlore tolore) SubExp]
-> Result
-> WriterT Result (AllocM fromlore tolore) Result
forall (m :: * -> *) a b c.
Applicative m =>
(a -> b -> m c) -> [a] -> [b] -> m [c]
zipWithM (SubExp -> WriterT Result (AllocM fromlore tolore) SubExp)
-> SubExp -> WriterT Result (AllocM fromlore tolore) SubExp
forall a b. (a -> b) -> a -> b
($) [SubExp -> WriterT Result (AllocM fromlore tolore) SubExp]
handle_loop_subexps Result
ses
        (Result, Result) -> AllocM fromlore tolore (Result, Result)
forall (m :: * -> *) a. Monad m => a -> m a
return (Result
memargs, Result
valargs)

  Scope tolore
-> AllocM fromlore tolore a -> AllocM fromlore tolore a
forall lore (m :: * -> *) a.
LocalScope lore m =>
Scope lore -> m a -> m a
localScope Scope tolore
summary (AllocM fromlore tolore a -> AllocM fromlore tolore a)
-> AllocM fromlore tolore a -> AllocM fromlore tolore a
forall a b. (a -> b) -> a -> b
$ [FParam tolore]
-> [FParam tolore]
-> (Result -> AllocM fromlore tolore (Result, Result))
-> AllocM fromlore tolore a
m [FParam tolore]
[Param FParamMem]
mem_params [FParam tolore]
[Param FParamMem]
valparams Result -> AllocM fromlore tolore (Result, Result)
mk_loop_res
  where allocInMergeParam :: (Param DeclType, SubExp)
-> WriterT
     [Param FParamMem]
     (AllocM fromlore tolore)
     (Param FParamMem,
      SubExp -> WriterT Result (AllocM fromlore tolore) SubExp)
allocInMergeParam (Param DeclType
mergeparam, Var VName
v)
          | Array PrimType
bt Shape
shape Uniqueness
u <- Param DeclType -> DeclType
forall attr. DeclTyped attr => Param attr -> DeclType
paramDeclType Param DeclType
mergeparam = do
              (VName
mem, IxFun
ixfun) <- AllocM fromlore tolore (VName, IxFun)
-> WriterT
     [Param FParamMem] (AllocM fromlore tolore) (VName, IxFun)
forall (t :: (* -> *) -> * -> *) (m :: * -> *) a.
(MonadTrans t, Monad m) =>
m a -> t m a
lift (AllocM fromlore tolore (VName, IxFun)
 -> WriterT
      [Param FParamMem] (AllocM fromlore tolore) (VName, IxFun))
-> AllocM fromlore tolore (VName, IxFun)
-> WriterT
     [Param FParamMem] (AllocM fromlore tolore) (VName, IxFun)
forall a b. (a -> b) -> a -> b
$ VName -> AllocM fromlore tolore (VName, IxFun)
forall lore (m :: * -> *).
(Mem lore, HasScope lore m, Monad m) =>
VName -> m (VName, IxFun)
lookupArraySummary VName
v
              Space
space <- AllocM fromlore tolore Space
-> WriterT [Param FParamMem] (AllocM fromlore tolore) Space
forall (t :: (* -> *) -> * -> *) (m :: * -> *) a.
(MonadTrans t, Monad m) =>
m a -> t m a
lift (AllocM fromlore tolore Space
 -> WriterT [Param FParamMem] (AllocM fromlore tolore) Space)
-> AllocM fromlore tolore Space
-> WriterT [Param FParamMem] (AllocM fromlore tolore) Space
forall a b. (a -> b) -> a -> b
$ VName -> AllocM fromlore tolore Space
forall lore (m :: * -> *).
(HasScope lore m, Monad m) =>
VName -> m Space
lookupMemSpace VName
mem
              Bool
reuse <- (AllocEnv fromlore tolore -> Bool)
-> WriterT [Param FParamMem] (AllocM fromlore tolore) Bool
forall r (m :: * -> *) a. MonadReader r m => (r -> a) -> m a
asks AllocEnv fromlore tolore -> Bool
forall fromlore tolore. AllocEnv fromlore tolore -> Bool
aggressiveReuse
              if Space
space Space -> Space -> Bool
forall a. Eq a => a -> a -> Bool
/= String -> Space
Space String
"local" Bool -> Bool -> Bool
&&
                 Bool
reuse Bool -> Bool -> Bool
&&
                 Uniqueness
u Uniqueness -> Uniqueness -> Bool
forall a. Eq a => a -> a -> Bool
== Uniqueness
Unique Bool -> Bool -> Bool
&&
                 Param DeclType -> Bool
loopInvariantShape Param DeclType
mergeparam
                then (Param FParamMem,
 SubExp -> WriterT Result (AllocM fromlore tolore) SubExp)
-> WriterT
     [Param FParamMem]
     (AllocM fromlore tolore)
     (Param FParamMem,
      SubExp -> WriterT Result (AllocM fromlore tolore) SubExp)
forall (m :: * -> *) a. Monad m => a -> m a
return (Param DeclType
mergeparam { paramAttr :: FParamMem
paramAttr = PrimType -> Shape -> Uniqueness -> MemBind -> FParamMem
forall d u ret.
PrimType -> ShapeBase d -> u -> ret -> MemInfo d u ret
MemArray PrimType
bt Shape
shape Uniqueness
Unique (MemBind -> FParamMem) -> MemBind -> FParamMem
forall a b. (a -> b) -> a -> b
$ VName -> IxFun -> MemBind
ArrayIn VName
mem IxFun
ixfun },
                             AllocM fromlore tolore SubExp
-> WriterT Result (AllocM fromlore tolore) SubExp
forall (t :: (* -> *) -> * -> *) (m :: * -> *) a.
(MonadTrans t, Monad m) =>
m a -> t m a
lift (AllocM fromlore tolore SubExp
 -> WriterT Result (AllocM fromlore tolore) SubExp)
-> (SubExp -> AllocM fromlore tolore SubExp)
-> SubExp
-> WriterT Result (AllocM fromlore tolore) SubExp
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Type -> VName -> IxFun -> SubExp -> AllocM fromlore tolore SubExp
forall fromlore tolore.
(Allocable fromlore tolore,
 Allocator tolore (AllocM fromlore tolore)) =>
Type -> VName -> IxFun -> SubExp -> AllocM fromlore tolore SubExp
ensureArrayIn (Param DeclType -> Type
forall attr. Typed attr => Param attr -> Type
paramType Param DeclType
mergeparam) VName
mem IxFun
ixfun)
                else do Space
def_space <- (AllocEnv fromlore tolore -> Space)
-> WriterT [Param FParamMem] (AllocM fromlore tolore) Space
forall r (m :: * -> *) a. MonadReader r m => (r -> a) -> m a
asks AllocEnv fromlore tolore -> Space
forall fromlore tolore. AllocEnv fromlore tolore -> Space
allocSpace
                        Param DeclType
-> Space
-> WriterT
     [FParam tolore]
     (AllocM fromlore tolore)
     (FParam tolore,
      SubExp -> WriterT Result (AllocM fromlore tolore) SubExp)
forall tolore fromlore tolore fromlore.
(PrettyLore fromlore, PrettyLore fromlore, AllocOp (Op tolore),
 AllocOp (Op tolore), Checkable tolore, Checkable tolore,
 OpReturns tolore, OpReturns tolore, SizeSubst (Op tolore),
 SizeSubst (Op tolore), BinderOps tolore, BinderOps tolore,
 FParamAttr tolore ~ FParamMem, BranchType tolore ~ BranchTypeMem,
 RetType fromlore ~ DeclExtType, BodyAttr fromlore ~ (),
 FParamAttr fromlore ~ DeclType, LParamAttr tolore ~ LetAttrMem,
 ExpAttr tolore ~ (), RetType tolore ~ RetTypeMem,
 BranchType fromlore ~ TypeBase ExtShape NoUniqueness,
 BodyAttr tolore ~ (), LParamAttr fromlore ~ Type,
 LetAttr tolore ~ LetAttrMem, BodyAttr tolore ~ (),
 LetAttr tolore ~ LetAttrMem, LParamAttr fromlore ~ Type,
 ExpAttr tolore ~ (),
 BranchType fromlore ~ TypeBase ExtShape NoUniqueness,
 RetType tolore ~ RetTypeMem, BodyAttr fromlore ~ (),
 LParamAttr tolore ~ LetAttrMem, FParamAttr fromlore ~ DeclType,
 FParamAttr tolore ~ FParamMem, RetType fromlore ~ DeclExtType,
 BranchType tolore ~ BranchTypeMem) =>
Param DeclType
-> Space
-> WriterT
     [Param (FParamAttr tolore)]
     (AllocM fromlore tolore)
     (Param (FParamAttr tolore),
      SubExp -> WriterT Result (AllocM fromlore tolore) SubExp)
doDefault Param DeclType
mergeparam Space
def_space

        allocInMergeParam (Param DeclType
mergeparam, SubExp
_) = Param DeclType
-> Space
-> WriterT
     [FParam tolore]
     (AllocM fromlore tolore)
     (FParam tolore,
      SubExp -> WriterT Result (AllocM fromlore tolore) SubExp)
forall tolore fromlore tolore fromlore.
(PrettyLore fromlore, PrettyLore fromlore, AllocOp (Op tolore),
 AllocOp (Op tolore), Checkable tolore, Checkable tolore,
 OpReturns tolore, OpReturns tolore, SizeSubst (Op tolore),
 SizeSubst (Op tolore), BinderOps tolore, BinderOps tolore,
 FParamAttr tolore ~ FParamMem, BranchType tolore ~ BranchTypeMem,
 RetType fromlore ~ DeclExtType, BodyAttr fromlore ~ (),
 FParamAttr fromlore ~ DeclType, LParamAttr tolore ~ LetAttrMem,
 ExpAttr tolore ~ (), RetType tolore ~ RetTypeMem,
 BranchType fromlore ~ TypeBase ExtShape NoUniqueness,
 BodyAttr tolore ~ (), LParamAttr fromlore ~ Type,
 LetAttr tolore ~ LetAttrMem, BodyAttr tolore ~ (),
 LetAttr tolore ~ LetAttrMem, LParamAttr fromlore ~ Type,
 ExpAttr tolore ~ (),
 BranchType fromlore ~ TypeBase ExtShape NoUniqueness,
 RetType tolore ~ RetTypeMem, BodyAttr fromlore ~ (),
 LParamAttr tolore ~ LetAttrMem, FParamAttr fromlore ~ DeclType,
 FParamAttr tolore ~ FParamMem, RetType fromlore ~ DeclExtType,
 BranchType tolore ~ BranchTypeMem) =>
Param DeclType
-> Space
-> WriterT
     [Param (FParamAttr tolore)]
     (AllocM fromlore tolore)
     (Param (FParamAttr tolore),
      SubExp -> WriterT Result (AllocM fromlore tolore) SubExp)
doDefault Param DeclType
mergeparam (Space
 -> WriterT
      [Param FParamMem]
      (AllocM fromlore tolore)
      (Param FParamMem,
       SubExp -> WriterT Result (AllocM fromlore tolore) SubExp))
-> WriterT [Param FParamMem] (AllocM fromlore tolore) Space
-> WriterT
     [Param FParamMem]
     (AllocM fromlore tolore)
     (Param FParamMem,
      SubExp -> WriterT Result (AllocM fromlore tolore) SubExp)
forall (m :: * -> *) a b. Monad m => (a -> m b) -> m a -> m b
=<< AllocM fromlore tolore Space
-> WriterT [Param FParamMem] (AllocM fromlore tolore) Space
forall (t :: (* -> *) -> * -> *) (m :: * -> *) a.
(MonadTrans t, Monad m) =>
m a -> t m a
lift AllocM fromlore tolore Space
forall lore (m :: * -> *). Allocator lore m => m Space
askDefaultSpace

        doDefault :: Param DeclType
-> Space
-> WriterT
     [Param (FParamAttr tolore)]
     (AllocM fromlore tolore)
     (Param (FParamAttr tolore),
      SubExp -> WriterT Result (AllocM fromlore tolore) SubExp)
doDefault Param DeclType
mergeparam Space
space = do
          Param (FParamAttr tolore)
mergeparam' <- FParam fromlore
-> Space
-> WriterT
     [Param (FParamAttr tolore)]
     (AllocM fromlore tolore)
     (Param (FParamAttr tolore))
forall fromlore tolore.
Allocable fromlore tolore =>
FParam fromlore
-> Space
-> WriterT [FParam tolore] (AllocM fromlore tolore) (FParam tolore)
allocInFParam Param DeclType
FParam fromlore
mergeparam Space
space
          (Param (FParamAttr tolore),
 SubExp -> WriterT Result (AllocM fromlore tolore) SubExp)
-> WriterT
     [Param (FParamAttr tolore)]
     (AllocM fromlore tolore)
     (Param (FParamAttr tolore),
      SubExp -> WriterT Result (AllocM fromlore tolore) SubExp)
forall (m :: * -> *) a. Monad m => a -> m a
return (Param (FParamAttr tolore)
mergeparam', Type
-> Space
-> SubExp
-> WriterT Result (AllocM fromlore tolore) SubExp
forall fromlore tolore.
(Allocable fromlore tolore,
 Allocator tolore (AllocM fromlore tolore)) =>
Type
-> Space
-> SubExp
-> WriterT Result (AllocM fromlore tolore) SubExp
linearFuncallArg (Param DeclType -> Type
forall attr. Typed attr => Param attr -> Type
paramType Param DeclType
mergeparam) Space
space)

        variant_names :: [VName]
variant_names = [VName]
variant [VName] -> [VName] -> [VName]
forall a. [a] -> [a] -> [a]
++ ((Param DeclType, SubExp) -> VName)
-> [(Param DeclType, SubExp)] -> [VName]
forall a b. (a -> b) -> [a] -> [b]
map (Param DeclType -> VName
forall attr. Param attr -> VName
paramName (Param DeclType -> VName)
-> ((Param DeclType, SubExp) -> Param DeclType)
-> (Param DeclType, SubExp)
-> VName
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (Param DeclType, SubExp) -> Param DeclType
forall a b. (a, b) -> a
fst) [(Param DeclType, SubExp)]
[(FParam fromlore, SubExp)]
merge
        loopInvariantShape :: Param DeclType -> Bool
loopInvariantShape =
          Bool -> Bool
not (Bool -> Bool)
-> (Param DeclType -> Bool) -> Param DeclType -> Bool
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (VName -> Bool) -> [VName] -> Bool
forall (t :: * -> *) a. Foldable t => (a -> Bool) -> t a -> Bool
any (VName -> [VName] -> Bool
forall (t :: * -> *) a. (Foldable t, Eq a) => a -> t a -> Bool
`elem` [VName]
variant_names) ([VName] -> Bool)
-> (Param DeclType -> [VName]) -> Param DeclType -> Bool
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Result -> [VName]
subExpVars (Result -> [VName])
-> (Param DeclType -> Result) -> Param DeclType -> [VName]
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Type -> Result
forall u. TypeBase Shape u -> Result
arrayDims (Type -> Result)
-> (Param DeclType -> Type) -> Param DeclType -> Result
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Param DeclType -> Type
forall attr. Typed attr => Param attr -> Type
paramType

ensureArrayIn :: (Allocable fromlore tolore,
                  Allocator tolore (AllocM fromlore tolore)) =>
                 Type -> VName -> IxFun -> SubExp
              -> AllocM fromlore tolore SubExp
ensureArrayIn :: Type -> VName -> IxFun -> SubExp -> AllocM fromlore tolore SubExp
ensureArrayIn Type
_ VName
_ IxFun
_ (Constant PrimValue
v) =
  String -> AllocM fromlore tolore SubExp
forall a. HasCallStack => String -> a
error (String -> AllocM fromlore tolore SubExp)
-> String -> AllocM fromlore tolore SubExp
forall a b. (a -> b) -> a -> b
$ String
"ensureArrayIn: " String -> ShowS
forall a. [a] -> [a] -> [a]
++ PrimValue -> String
forall a. Pretty a => a -> String
pretty PrimValue
v String -> ShowS
forall a. [a] -> [a] -> [a]
++ String
" cannot be an array."
ensureArrayIn Type
t VName
mem IxFun
ixfun (Var VName
v) = do
  (VName
src_mem, IxFun
src_ixfun) <- VName -> AllocM fromlore tolore (VName, IxFun)
forall lore (m :: * -> *).
(Mem lore, HasScope lore m, Monad m) =>
VName -> m (VName, IxFun)
lookupArraySummary VName
v
  if VName
src_mem VName -> VName -> Bool
forall a. Eq a => a -> a -> Bool
== VName
mem Bool -> Bool -> Bool
&& IxFun
src_ixfun IxFun -> IxFun -> Bool
forall a. Eq a => a -> a -> Bool
== IxFun
ixfun
    then SubExp -> AllocM fromlore tolore SubExp
forall (m :: * -> *) a. Monad m => a -> m a
return (SubExp -> AllocM fromlore tolore SubExp)
-> SubExp -> AllocM fromlore tolore SubExp
forall a b. (a -> b) -> a -> b
$ VName -> SubExp
Var VName
v
    else do Ident
copy <- String -> Type -> AllocM fromlore tolore Ident
forall (m :: * -> *).
MonadFreshNames m =>
String -> Type -> m Ident
newIdent (VName -> String
baseString VName
v String -> ShowS
forall a. [a] -> [a] -> [a]
++ String
"_ensure_copy") Type
t
            let summary :: LetAttrMem
summary = PrimType -> Shape -> NoUniqueness -> MemBind -> LetAttrMem
forall d u ret.
PrimType -> ShapeBase d -> u -> ret -> MemInfo d u ret
MemArray (Type -> PrimType
forall shape u. TypeBase shape u -> PrimType
elemType Type
t) (Type -> Shape
forall shape u. ArrayShape shape => TypeBase shape u -> shape
arrayShape Type
t) NoUniqueness
NoUniqueness (MemBind -> LetAttrMem) -> MemBind -> LetAttrMem
forall a b. (a -> b) -> a -> b
$
                          VName -> IxFun -> MemBind
ArrayIn VName
mem IxFun
ixfun
                pat :: PatternT LetAttrMem
pat = [PatElemT LetAttrMem]
-> [PatElemT LetAttrMem] -> PatternT LetAttrMem
forall attr. [PatElemT attr] -> [PatElemT attr] -> PatternT attr
Pattern [] [VName -> LetAttrMem -> PatElemT LetAttrMem
forall attr. VName -> attr -> PatElemT attr
PatElem (Ident -> VName
identName Ident
copy) LetAttrMem
summary]
            Pattern (Lore (AllocM fromlore tolore))
-> Exp (Lore (AllocM fromlore tolore)) -> AllocM fromlore tolore ()
forall (m :: * -> *).
MonadBinder m =>
Pattern (Lore m) -> Exp (Lore m) -> m ()
letBind_ Pattern (Lore (AllocM fromlore tolore))
PatternT LetAttrMem
pat (Exp (Lore (AllocM fromlore tolore)) -> AllocM fromlore tolore ())
-> Exp (Lore (AllocM fromlore tolore)) -> AllocM fromlore tolore ()
forall a b. (a -> b) -> a -> b
$ BasicOp -> ExpT tolore
forall lore. BasicOp -> ExpT lore
BasicOp (BasicOp -> ExpT tolore) -> BasicOp -> ExpT tolore
forall a b. (a -> b) -> a -> b
$ VName -> BasicOp
Copy VName
v
            SubExp -> AllocM fromlore tolore SubExp
forall (m :: * -> *) a. Monad m => a -> m a
return (SubExp -> AllocM fromlore tolore SubExp)
-> SubExp -> AllocM fromlore tolore SubExp
forall a b. (a -> b) -> a -> b
$ VName -> SubExp
Var (VName -> SubExp) -> VName -> SubExp
forall a b. (a -> b) -> a -> b
$ Ident -> VName
identName Ident
copy

ensureDirectArray :: (Allocable fromlore tolore,
                      Allocator tolore (AllocM fromlore tolore)) =>
                     Maybe Space -> VName -> AllocM fromlore tolore (VName, SubExp)
ensureDirectArray :: Maybe Space -> VName -> AllocM fromlore tolore (VName, SubExp)
ensureDirectArray Maybe Space
space_ok VName
v = do
  (VName
mem, IxFun
ixfun) <- VName -> AllocM fromlore tolore (VName, IxFun)
forall lore (m :: * -> *).
(Mem lore, HasScope lore m, Monad m) =>
VName -> m (VName, IxFun)
lookupArraySummary VName
v
  Space
mem_space <- VName -> AllocM fromlore tolore Space
forall lore (m :: * -> *).
(HasScope lore m, Monad m) =>
VName -> m Space
lookupMemSpace VName
mem
  Space
default_space <- AllocM fromlore tolore Space
forall lore (m :: * -> *). Allocator lore m => m Space
askDefaultSpace
  if IxFun -> Bool
forall num. (Eq num, IntegralExp num) => IxFun num -> Bool
IxFun.isDirect IxFun
ixfun Bool -> Bool -> Bool
&& Bool -> (Space -> Bool) -> Maybe Space -> Bool
forall b a. b -> (a -> b) -> Maybe a -> b
maybe Bool
True (Space -> Space -> Bool
forall a. Eq a => a -> a -> Bool
==Space
mem_space) Maybe Space
space_ok
    then (VName, SubExp) -> AllocM fromlore tolore (VName, SubExp)
forall (m :: * -> *) a. Monad m => a -> m a
return (VName
mem, VName -> SubExp
Var VName
v)
    else Space -> AllocM fromlore tolore (VName, SubExp)
needCopy (Space -> Maybe Space -> Space
forall a. a -> Maybe a -> a
fromMaybe Space
default_space Maybe Space
space_ok)
  where needCopy :: Space -> AllocM fromlore tolore (VName, SubExp)
needCopy Space
space =
          -- We need to do a new allocation, copy 'v', and make a new
          -- binding for the size of the memory block.
          Space -> String -> VName -> AllocM fromlore tolore (VName, SubExp)
forall fromlore tolore.
(Allocable fromlore tolore,
 Allocator tolore (AllocM fromlore tolore)) =>
Space -> String -> VName -> AllocM fromlore tolore (VName, SubExp)
allocLinearArray Space
space (VName -> String
baseString VName
v) VName
v

allocLinearArray :: (Allocable fromlore tolore, Allocator tolore (AllocM fromlore tolore)) =>
                    Space -> String -> VName
                 -> AllocM fromlore tolore (VName, SubExp)
allocLinearArray :: Space -> String -> VName -> AllocM fromlore tolore (VName, SubExp)
allocLinearArray Space
space String
s VName
v = do
  Type
t <- VName -> AllocM fromlore tolore Type
forall lore (m :: * -> *). HasScope lore m => VName -> m Type
lookupType VName
v
  VName
mem <- Type -> Space -> AllocM fromlore tolore VName
forall lore (m :: * -> *).
Allocator lore m =>
Type -> Space -> m VName
allocForArray Type
t Space
space
  Ident
v' <- String -> Type -> AllocM fromlore tolore Ident
forall (m :: * -> *).
MonadFreshNames m =>
String -> Type -> m Ident
newIdent (String
s String -> ShowS
forall a. [a] -> [a] -> [a]
++ String
"_linear") Type
t
  let pat :: PatternT LetAttrMem
pat = [PatElemT LetAttrMem]
-> [PatElemT LetAttrMem] -> PatternT LetAttrMem
forall attr. [PatElemT attr] -> [PatElemT attr] -> PatternT attr
Pattern [] [VName -> LetAttrMem -> PatElemT LetAttrMem
forall attr. VName -> attr -> PatElemT attr
PatElem (Ident -> VName
identName Ident
v') (LetAttrMem -> PatElemT LetAttrMem)
-> LetAttrMem -> PatElemT LetAttrMem
forall a b. (a -> b) -> a -> b
$
                        PrimType -> Shape -> NoUniqueness -> VName -> Type -> LetAttrMem
forall u. PrimType -> Shape -> u -> VName -> Type -> MemBound u
directIxFun (Type -> PrimType
forall shape u. TypeBase shape u -> PrimType
elemType Type
t) (Type -> Shape
forall shape u. ArrayShape shape => TypeBase shape u -> shape
arrayShape Type
t)
                        NoUniqueness
NoUniqueness VName
mem Type
t]
  Stm (Lore (AllocM fromlore tolore)) -> AllocM fromlore tolore ()
forall (m :: * -> *). MonadBinder m => Stm (Lore m) -> m ()
addStm (Stm (Lore (AllocM fromlore tolore)) -> AllocM fromlore tolore ())
-> Stm (Lore (AllocM fromlore tolore)) -> AllocM fromlore tolore ()
forall a b. (a -> b) -> a -> b
$ Pattern tolore
-> StmAux (ExpAttr tolore) -> Exp tolore -> Stm tolore
forall lore.
Pattern lore -> StmAux (ExpAttr lore) -> Exp lore -> Stm lore
Let Pattern tolore
PatternT LetAttrMem
pat (() -> StmAux ()
forall attr. attr -> StmAux attr
defAux ()) (Exp tolore -> Stm tolore) -> Exp tolore -> Stm tolore
forall a b. (a -> b) -> a -> b
$ BasicOp -> Exp tolore
forall lore. BasicOp -> ExpT lore
BasicOp (BasicOp -> Exp tolore) -> BasicOp -> Exp tolore
forall a b. (a -> b) -> a -> b
$ VName -> BasicOp
Copy VName
v
  (VName, SubExp) -> AllocM fromlore tolore (VName, SubExp)
forall (m :: * -> *) a. Monad m => a -> m a
return (VName
mem, VName -> SubExp
Var (VName -> SubExp) -> VName -> SubExp
forall a b. (a -> b) -> a -> b
$ Ident -> VName
identName Ident
v')

funcallArgs :: (Allocable fromlore tolore,
                Allocator tolore (AllocM fromlore tolore)) =>
               [(SubExp,Diet)] -> AllocM fromlore tolore [(SubExp,Diet)]
funcallArgs :: [(SubExp, Diet)] -> AllocM fromlore tolore [(SubExp, Diet)]
funcallArgs [(SubExp, Diet)]
args = do
  ([(SubExp, Diet)]
valargs, Result
mem_and_size_args) <- WriterT Result (AllocM fromlore tolore) [(SubExp, Diet)]
-> AllocM fromlore tolore ([(SubExp, Diet)], Result)
forall w (m :: * -> *) a. WriterT w m a -> m (a, w)
runWriterT (WriterT Result (AllocM fromlore tolore) [(SubExp, Diet)]
 -> AllocM fromlore tolore ([(SubExp, Diet)], Result))
-> WriterT Result (AllocM fromlore tolore) [(SubExp, Diet)]
-> AllocM fromlore tolore ([(SubExp, Diet)], Result)
forall a b. (a -> b) -> a -> b
$ [(SubExp, Diet)]
-> ((SubExp, Diet)
    -> WriterT Result (AllocM fromlore tolore) (SubExp, Diet))
-> WriterT Result (AllocM fromlore tolore) [(SubExp, Diet)]
forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
t a -> (a -> m b) -> m (t b)
forM [(SubExp, Diet)]
args (((SubExp, Diet)
  -> WriterT Result (AllocM fromlore tolore) (SubExp, Diet))
 -> WriterT Result (AllocM fromlore tolore) [(SubExp, Diet)])
-> ((SubExp, Diet)
    -> WriterT Result (AllocM fromlore tolore) (SubExp, Diet))
-> WriterT Result (AllocM fromlore tolore) [(SubExp, Diet)]
forall a b. (a -> b) -> a -> b
$ \(SubExp
arg,Diet
d) -> do
    Type
t <- AllocM fromlore tolore Type
-> WriterT Result (AllocM fromlore tolore) Type
forall (t :: (* -> *) -> * -> *) (m :: * -> *) a.
(MonadTrans t, Monad m) =>
m a -> t m a
lift (AllocM fromlore tolore Type
 -> WriterT Result (AllocM fromlore tolore) Type)
-> AllocM fromlore tolore Type
-> WriterT Result (AllocM fromlore tolore) Type
forall a b. (a -> b) -> a -> b
$ SubExp -> AllocM fromlore tolore Type
forall t (m :: * -> *). HasScope t m => SubExp -> m Type
subExpType SubExp
arg
    Space
space <- AllocM fromlore tolore Space
-> WriterT Result (AllocM fromlore tolore) Space
forall (t :: (* -> *) -> * -> *) (m :: * -> *) a.
(MonadTrans t, Monad m) =>
m a -> t m a
lift AllocM fromlore tolore Space
forall lore (m :: * -> *). Allocator lore m => m Space
askDefaultSpace
    SubExp
arg' <- Type
-> Space
-> SubExp
-> WriterT Result (AllocM fromlore tolore) SubExp
forall fromlore tolore.
(Allocable fromlore tolore,
 Allocator tolore (AllocM fromlore tolore)) =>
Type
-> Space
-> SubExp
-> WriterT Result (AllocM fromlore tolore) SubExp
linearFuncallArg Type
t Space
space SubExp
arg
    (SubExp, Diet)
-> WriterT Result (AllocM fromlore tolore) (SubExp, Diet)
forall (m :: * -> *) a. Monad m => a -> m a
return (SubExp
arg', Diet
d)
  [(SubExp, Diet)] -> AllocM fromlore tolore [(SubExp, Diet)]
forall (m :: * -> *) a. Monad m => a -> m a
return ([(SubExp, Diet)] -> AllocM fromlore tolore [(SubExp, Diet)])
-> [(SubExp, Diet)] -> AllocM fromlore tolore [(SubExp, Diet)]
forall a b. (a -> b) -> a -> b
$ (SubExp -> (SubExp, Diet)) -> Result -> [(SubExp, Diet)]
forall a b. (a -> b) -> [a] -> [b]
map (,Diet
Observe) Result
mem_and_size_args [(SubExp, Diet)] -> [(SubExp, Diet)] -> [(SubExp, Diet)]
forall a. Semigroup a => a -> a -> a
<> [(SubExp, Diet)]
valargs

linearFuncallArg :: (Allocable fromlore tolore,
                     Allocator tolore (AllocM fromlore tolore)) =>
                    Type -> Space -> SubExp
                 -> WriterT [SubExp] (AllocM fromlore tolore) SubExp
linearFuncallArg :: Type
-> Space
-> SubExp
-> WriterT Result (AllocM fromlore tolore) SubExp
linearFuncallArg Array{} Space
space (Var VName
v) = do
  (VName
mem, SubExp
arg') <- AllocM fromlore tolore (VName, SubExp)
-> WriterT Result (AllocM fromlore tolore) (VName, SubExp)
forall (t :: (* -> *) -> * -> *) (m :: * -> *) a.
(MonadTrans t, Monad m) =>
m a -> t m a
lift (AllocM fromlore tolore (VName, SubExp)
 -> WriterT Result (AllocM fromlore tolore) (VName, SubExp))
-> AllocM fromlore tolore (VName, SubExp)
-> WriterT Result (AllocM fromlore tolore) (VName, SubExp)
forall a b. (a -> b) -> a -> b
$ Maybe Space -> VName -> AllocM fromlore tolore (VName, SubExp)
forall fromlore tolore.
(Allocable fromlore tolore,
 Allocator tolore (AllocM fromlore tolore)) =>
Maybe Space -> VName -> AllocM fromlore tolore (VName, SubExp)
ensureDirectArray (Space -> Maybe Space
forall a. a -> Maybe a
Just Space
space) VName
v
  Result -> WriterT Result (AllocM fromlore tolore) ()
forall w (m :: * -> *). MonadWriter w m => w -> m ()
tell [VName -> SubExp
Var VName
mem]
  SubExp -> WriterT Result (AllocM fromlore tolore) SubExp
forall (m :: * -> *) a. Monad m => a -> m a
return SubExp
arg'
linearFuncallArg Type
_ Space
_ SubExp
arg =
  SubExp -> WriterT Result (AllocM fromlore tolore) SubExp
forall (m :: * -> *) a. Monad m => a -> m a
return SubExp
arg

explicitAllocationsGeneric :: (Allocable fromlore tolore,
                               Allocator tolore (AllocM fromlore tolore)) =>
                              (Op fromlore -> AllocM fromlore tolore (Op tolore))
                           -> (Exp tolore -> AllocM fromlore tolore [ExpHint])
                           -> Pass fromlore tolore
explicitAllocationsGeneric :: (Op fromlore -> AllocM fromlore tolore (Op tolore))
-> (Exp tolore -> AllocM fromlore tolore [ExpHint])
-> Pass fromlore tolore
explicitAllocationsGeneric Op fromlore -> AllocM fromlore tolore (Op tolore)
handleOp Exp tolore -> AllocM fromlore tolore [ExpHint]
hints =
  String
-> String
-> (Prog fromlore -> PassM (Prog tolore))
-> Pass fromlore tolore
forall fromlore tolore.
String
-> String
-> (Prog fromlore -> PassM (Prog tolore))
-> Pass fromlore tolore
Pass String
"explicit allocations" String
"Transform program to explicit memory representation" ((Prog fromlore -> PassM (Prog tolore)) -> Pass fromlore tolore)
-> (Prog fromlore -> PassM (Prog tolore)) -> Pass fromlore tolore
forall a b. (a -> b) -> a -> b
$
  (Stms fromlore -> PassM (Stms tolore))
-> (Stms tolore -> FunDef fromlore -> PassM (FunDef tolore))
-> Prog fromlore
-> PassM (Prog tolore)
forall fromlore tolore.
(Stms fromlore -> PassM (Stms tolore))
-> (Stms tolore -> FunDef fromlore -> PassM (FunDef tolore))
-> Prog fromlore
-> PassM (Prog tolore)
intraproceduralTransformationWithConsts Stms fromlore -> PassM (Stms tolore)
onStms Stms tolore -> FunDef fromlore -> PassM (FunDef tolore)
allocInFun
  where onStms :: Stms fromlore -> PassM (Stms tolore)
onStms Stms fromlore
stms = (Op fromlore -> AllocM fromlore tolore (Op tolore))
-> (Exp tolore -> AllocM fromlore tolore [ExpHint])
-> AllocM fromlore tolore (Stms tolore)
-> PassM (Stms tolore)
forall (m :: * -> *) fromlore tolore a.
MonadFreshNames m =>
(Op fromlore -> AllocM fromlore tolore (Op tolore))
-> (Exp tolore -> AllocM fromlore tolore [ExpHint])
-> AllocM fromlore tolore a
-> m a
runAllocM Op fromlore -> AllocM fromlore tolore (Op tolore)
handleOp Exp tolore -> AllocM fromlore tolore [ExpHint]
hints (AllocM fromlore tolore (Stms tolore) -> PassM (Stms tolore))
-> AllocM fromlore tolore (Stms tolore) -> PassM (Stms tolore)
forall a b. (a -> b) -> a -> b
$ Stms fromlore
-> (Stms tolore -> AllocM fromlore tolore (Stms tolore))
-> AllocM fromlore tolore (Stms tolore)
forall fromlore tolore a.
Allocable fromlore tolore =>
Stms fromlore
-> (Stms tolore -> AllocM fromlore tolore a)
-> AllocM fromlore tolore a
allocInStms Stms fromlore
stms Stms tolore -> AllocM fromlore tolore (Stms tolore)
forall (f :: * -> *) a. Applicative f => a -> f a
pure

        allocInFun :: Stms tolore -> FunDef fromlore -> PassM (FunDef tolore)
allocInFun Stms tolore
consts (FunDef Maybe EntryPoint
entry Name
fname [RetType fromlore]
rettype [FParam fromlore]
params BodyT fromlore
fbody) =
          (Op fromlore -> AllocM fromlore tolore (Op tolore))
-> (Exp tolore -> AllocM fromlore tolore [ExpHint])
-> AllocM fromlore tolore (FunDef tolore)
-> PassM (FunDef tolore)
forall (m :: * -> *) fromlore tolore a.
MonadFreshNames m =>
(Op fromlore -> AllocM fromlore tolore (Op tolore))
-> (Exp tolore -> AllocM fromlore tolore [ExpHint])
-> AllocM fromlore tolore a
-> m a
runAllocM Op fromlore -> AllocM fromlore tolore (Op tolore)
handleOp Exp tolore -> AllocM fromlore tolore [ExpHint]
hints (AllocM fromlore tolore (FunDef tolore) -> PassM (FunDef tolore))
-> AllocM fromlore tolore (FunDef tolore) -> PassM (FunDef tolore)
forall a b. (a -> b) -> a -> b
$ Stms tolore
-> AllocM fromlore tolore (FunDef tolore)
-> AllocM fromlore tolore (FunDef tolore)
forall lore a (m :: * -> *) b.
(Scoped lore a, LocalScope lore m) =>
a -> m b -> m b
inScopeOf Stms tolore
consts (AllocM fromlore tolore (FunDef tolore)
 -> AllocM fromlore tolore (FunDef tolore))
-> AllocM fromlore tolore (FunDef tolore)
-> AllocM fromlore tolore (FunDef tolore)
forall a b. (a -> b) -> a -> b
$
          [(FParam fromlore, Space)]
-> ([FParam tolore] -> AllocM fromlore tolore (FunDef tolore))
-> AllocM fromlore tolore (FunDef tolore)
forall fromlore tolore a.
Allocable fromlore tolore =>
[(FParam fromlore, Space)]
-> ([FParam tolore] -> AllocM fromlore tolore a)
-> AllocM fromlore tolore a
allocInFParams ([Param DeclType] -> [Space] -> [(Param DeclType, Space)]
forall a b. [a] -> [b] -> [(a, b)]
zip [Param DeclType]
[FParam fromlore]
params ([Space] -> [(Param DeclType, Space)])
-> [Space] -> [(Param DeclType, Space)]
forall a b. (a -> b) -> a -> b
$ Space -> [Space]
forall a. a -> [a]
repeat Space
DefaultSpace) (([FParam tolore] -> AllocM fromlore tolore (FunDef tolore))
 -> AllocM fromlore tolore (FunDef tolore))
-> ([FParam tolore] -> AllocM fromlore tolore (FunDef tolore))
-> AllocM fromlore tolore (FunDef tolore)
forall a b. (a -> b) -> a -> b
$ \[FParam tolore]
params' -> do
          BodyT tolore
fbody' <- AllocM fromlore tolore (Body (Lore (AllocM fromlore tolore)))
-> AllocM fromlore tolore (Body (Lore (AllocM fromlore tolore)))
forall (m :: * -> *).
MonadBinder m =>
m (Body (Lore m)) -> m (Body (Lore m))
insertStmsM (AllocM fromlore tolore (Body (Lore (AllocM fromlore tolore)))
 -> AllocM fromlore tolore (Body (Lore (AllocM fromlore tolore))))
-> AllocM fromlore tolore (Body (Lore (AllocM fromlore tolore)))
-> AllocM fromlore tolore (Body (Lore (AllocM fromlore tolore)))
forall a b. (a -> b) -> a -> b
$ [Maybe Space]
-> BodyT fromlore -> AllocM fromlore tolore (BodyT tolore)
forall fromlore tolore.
(Allocable fromlore tolore,
 Allocator tolore (AllocM fromlore tolore)) =>
[Maybe Space]
-> Body fromlore -> AllocM fromlore tolore (Body tolore)
allocInFunBody
                    ((DeclExtType -> Maybe Space) -> [DeclExtType] -> [Maybe Space]
forall a b. (a -> b) -> [a] -> [b]
map (Maybe Space -> DeclExtType -> Maybe Space
forall a b. a -> b -> a
const (Maybe Space -> DeclExtType -> Maybe Space)
-> Maybe Space -> DeclExtType -> Maybe Space
forall a b. (a -> b) -> a -> b
$ Space -> Maybe Space
forall a. a -> Maybe a
Just Space
DefaultSpace) [DeclExtType]
[RetType fromlore]
rettype) BodyT fromlore
fbody
          FunDef tolore -> AllocM fromlore tolore (FunDef tolore)
forall (m :: * -> *) a. Monad m => a -> m a
return (FunDef tolore -> AllocM fromlore tolore (FunDef tolore))
-> FunDef tolore -> AllocM fromlore tolore (FunDef tolore)
forall a b. (a -> b) -> a -> b
$ Maybe EntryPoint
-> Name
-> [RetType tolore]
-> [FParam tolore]
-> BodyT tolore
-> FunDef tolore
forall lore.
Maybe EntryPoint
-> Name
-> [RetType lore]
-> [FParam lore]
-> BodyT lore
-> FunDef lore
FunDef Maybe EntryPoint
entry Name
fname ([DeclExtType] -> [RetTypeMem]
memoryInDeclExtType [DeclExtType]
[RetType fromlore]
rettype) [FParam tolore]
params' BodyT tolore
fbody'

explicitAllocationsInStmsGeneric :: (MonadFreshNames m, HasScope tolore m,
                                     Allocable fromlore tolore) =>
                                    (Op fromlore -> AllocM fromlore tolore (Op tolore))
                                 -> (Exp tolore -> AllocM fromlore tolore [ExpHint])
                                 -> Stms fromlore -> m (Stms tolore)
explicitAllocationsInStmsGeneric :: (Op fromlore -> AllocM fromlore tolore (Op tolore))
-> (Exp tolore -> AllocM fromlore tolore [ExpHint])
-> Stms fromlore
-> m (Stms tolore)
explicitAllocationsInStmsGeneric Op fromlore -> AllocM fromlore tolore (Op tolore)
handleOp Exp tolore -> AllocM fromlore tolore [ExpHint]
hints Stms fromlore
stms = do
  Scope tolore
scope <- m (Scope tolore)
forall lore (m :: * -> *). HasScope lore m => m (Scope lore)
askScope
  (Op fromlore -> AllocM fromlore tolore (Op tolore))
-> (Exp tolore -> AllocM fromlore tolore [ExpHint])
-> AllocM fromlore tolore (Stms tolore)
-> m (Stms tolore)
forall (m :: * -> *) fromlore tolore a.
MonadFreshNames m =>
(Op fromlore -> AllocM fromlore tolore (Op tolore))
-> (Exp tolore -> AllocM fromlore tolore [ExpHint])
-> AllocM fromlore tolore a
-> m a
runAllocM Op fromlore -> AllocM fromlore tolore (Op tolore)
handleOp Exp tolore -> AllocM fromlore tolore [ExpHint]
hints (AllocM fromlore tolore (Stms tolore) -> m (Stms tolore))
-> AllocM fromlore tolore (Stms tolore) -> m (Stms tolore)
forall a b. (a -> b) -> a -> b
$ Scope tolore
-> AllocM fromlore tolore (Stms tolore)
-> AllocM fromlore tolore (Stms tolore)
forall lore (m :: * -> *) a.
LocalScope lore m =>
Scope lore -> m a -> m a
localScope Scope tolore
scope (AllocM fromlore tolore (Stms tolore)
 -> AllocM fromlore tolore (Stms tolore))
-> AllocM fromlore tolore (Stms tolore)
-> AllocM fromlore tolore (Stms tolore)
forall a b. (a -> b) -> a -> b
$ Stms fromlore
-> (Stms tolore -> AllocM fromlore tolore (Stms tolore))
-> AllocM fromlore tolore (Stms tolore)
forall fromlore tolore a.
Allocable fromlore tolore =>
Stms fromlore
-> (Stms tolore -> AllocM fromlore tolore a)
-> AllocM fromlore tolore a
allocInStms Stms fromlore
stms Stms tolore -> AllocM fromlore tolore (Stms tolore)
forall (m :: * -> *) a. Monad m => a -> m a
return

memoryInDeclExtType :: [DeclExtType] -> [FunReturns]
memoryInDeclExtType :: [DeclExtType] -> [RetTypeMem]
memoryInDeclExtType [DeclExtType]
ts = State Int [RetTypeMem] -> Int -> [RetTypeMem]
forall s a. State s a -> s -> a
evalState ((DeclExtType -> StateT Int Identity RetTypeMem)
-> [DeclExtType] -> State Int [RetTypeMem]
forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
(a -> m b) -> t a -> m (t b)
mapM DeclExtType -> StateT Int Identity RetTypeMem
forall (m :: * -> *) u.
MonadState Int m =>
TypeBase ExtShape u -> m (MemInfo (Ext SubExp) u MemReturn)
addAttr [DeclExtType]
ts) (Int -> [RetTypeMem]) -> Int -> [RetTypeMem]
forall a b. (a -> b) -> a -> b
$ [DeclExtType] -> Int
forall u. [TypeBase ExtShape u] -> Int
startOfFreeIDRange [DeclExtType]
ts
  where addAttr :: TypeBase ExtShape u -> m (MemInfo (Ext SubExp) u MemReturn)
addAttr (Prim PrimType
t) = MemInfo (Ext SubExp) u MemReturn
-> m (MemInfo (Ext SubExp) u MemReturn)
forall (m :: * -> *) a. Monad m => a -> m a
return (MemInfo (Ext SubExp) u MemReturn
 -> m (MemInfo (Ext SubExp) u MemReturn))
-> MemInfo (Ext SubExp) u MemReturn
-> m (MemInfo (Ext SubExp) u MemReturn)
forall a b. (a -> b) -> a -> b
$ PrimType -> MemInfo (Ext SubExp) u MemReturn
forall d u ret. PrimType -> MemInfo d u ret
MemPrim PrimType
t
        addAttr Mem{} = String -> m (MemInfo (Ext SubExp) u MemReturn)
forall a. HasCallStack => String -> a
error String
"memoryInDeclExtType: too much memory"
        addAttr (Array PrimType
bt ExtShape
shape u
u) = do
          Int
i <- m Int
forall s (m :: * -> *). MonadState s m => m s
get m Int -> m () -> m Int
forall (f :: * -> *) a b. Applicative f => f a -> f b -> f a
<* (Int -> Int) -> m ()
forall s (m :: * -> *). MonadState s m => (s -> s) -> m ()
modify (Int -> Int -> Int
forall a. Num a => a -> a -> a
+Int
1)
          MemInfo (Ext SubExp) u MemReturn
-> m (MemInfo (Ext SubExp) u MemReturn)
forall (m :: * -> *) a. Monad m => a -> m a
return (MemInfo (Ext SubExp) u MemReturn
 -> m (MemInfo (Ext SubExp) u MemReturn))
-> MemInfo (Ext SubExp) u MemReturn
-> m (MemInfo (Ext SubExp) u MemReturn)
forall a b. (a -> b) -> a -> b
$ PrimType
-> ExtShape -> u -> MemReturn -> MemInfo (Ext SubExp) u MemReturn
forall d u ret.
PrimType -> ShapeBase d -> u -> ret -> MemInfo d u ret
MemArray PrimType
bt ExtShape
shape u
u (MemReturn -> MemInfo (Ext SubExp) u MemReturn)
-> MemReturn -> MemInfo (Ext SubExp) u MemReturn
forall a b. (a -> b) -> a -> b
$ Space -> Int -> ExtIxFun -> MemReturn
ReturnsNewBlock Space
DefaultSpace Int
i (ExtIxFun -> MemReturn) -> ExtIxFun -> MemReturn
forall a b. (a -> b) -> a -> b
$
            Shape (PrimExp (Ext VName)) -> ExtIxFun
forall num. IntegralExp num => Shape num -> IxFun num
IxFun.iota (Shape (PrimExp (Ext VName)) -> ExtIxFun)
-> Shape (PrimExp (Ext VName)) -> ExtIxFun
forall a b. (a -> b) -> a -> b
$ (Ext SubExp -> PrimExp (Ext VName))
-> [Ext SubExp] -> Shape (PrimExp (Ext VName))
forall a b. (a -> b) -> [a] -> [b]
map Ext SubExp -> PrimExp (Ext VName)
convert ([Ext SubExp] -> Shape (PrimExp (Ext VName)))
-> [Ext SubExp] -> Shape (PrimExp (Ext VName))
forall a b. (a -> b) -> a -> b
$ ExtShape -> [Ext SubExp]
forall d. ShapeBase d -> [d]
shapeDims ExtShape
shape

        convert :: Ext SubExp -> PrimExp (Ext VName)
convert (Ext Int
i) = Ext VName -> PrimType -> PrimExp (Ext VName)
forall v. v -> PrimType -> PrimExp v
LeafExp (Int -> Ext VName
forall a. Int -> Ext a
Ext Int
i) PrimType
int32
        convert (Free SubExp
v) = VName -> Ext VName
forall a. a -> Ext a
Free (VName -> Ext VName) -> PrimExp VName -> PrimExp (Ext VName)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> PrimType -> SubExp -> PrimExp VName
primExpFromSubExp PrimType
int32 SubExp
v

startOfFreeIDRange :: [TypeBase ExtShape u] -> Int
startOfFreeIDRange :: [TypeBase ExtShape u] -> Int
startOfFreeIDRange = Set Int -> Int
forall a. Set a -> Int
S.size (Set Int -> Int)
-> ([TypeBase ExtShape u] -> Set Int)
-> [TypeBase ExtShape u]
-> Int
forall b c a. (b -> c) -> (a -> b) -> a -> c
. [TypeBase ExtShape u] -> Set Int
forall u. [TypeBase ExtShape u] -> Set Int
shapeContext

bodyReturnMemCtx :: (Allocable fromlore tolore, Allocator tolore (AllocM fromlore tolore)) =>
                    SubExp -> AllocM fromlore tolore [SubExp]
bodyReturnMemCtx :: SubExp -> AllocM fromlore tolore Result
bodyReturnMemCtx Constant{} =
  Result -> AllocM fromlore tolore Result
forall (m :: * -> *) a. Monad m => a -> m a
return []
bodyReturnMemCtx (Var VName
v) = do
  LetAttrMem
info <- VName -> AllocM fromlore tolore LetAttrMem
forall lore (m :: * -> *).
(HasScope lore m, Mem lore) =>
VName -> m LetAttrMem
lookupMemInfo VName
v
  case LetAttrMem
info of
    MemPrim{} -> Result -> AllocM fromlore tolore Result
forall (m :: * -> *) a. Monad m => a -> m a
return []
    MemMem{} -> Result -> AllocM fromlore tolore Result
forall (m :: * -> *) a. Monad m => a -> m a
return [] -- should not happen
    MemArray PrimType
_ Shape
_ NoUniqueness
_ (ArrayIn VName
mem IxFun
_) -> Result -> AllocM fromlore tolore Result
forall (m :: * -> *) a. Monad m => a -> m a
return [VName -> SubExp
Var VName
mem]

allocInFunBody :: (Allocable fromlore tolore, Allocator tolore (AllocM fromlore tolore)) =>
                  [Maybe Space] -> Body fromlore -> AllocM fromlore tolore (Body tolore)
allocInFunBody :: [Maybe Space]
-> Body fromlore -> AllocM fromlore tolore (Body tolore)
allocInFunBody [Maybe Space]
space_oks (Body BodyAttr fromlore
_ Stms fromlore
bnds Result
res) =
  Stms fromlore
-> (Stms tolore -> AllocM fromlore tolore (Body tolore))
-> AllocM fromlore tolore (Body tolore)
forall fromlore tolore a.
Allocable fromlore tolore =>
Stms fromlore
-> (Stms tolore -> AllocM fromlore tolore a)
-> AllocM fromlore tolore a
allocInStms Stms fromlore
bnds ((Stms tolore -> AllocM fromlore tolore (Body tolore))
 -> AllocM fromlore tolore (Body tolore))
-> (Stms tolore -> AllocM fromlore tolore (Body tolore))
-> AllocM fromlore tolore (Body tolore)
forall a b. (a -> b) -> a -> b
$ \Stms tolore
bnds' -> do
    (Result
res'', Stms tolore
allocs) <- AllocM fromlore tolore Result
-> AllocM
     fromlore tolore (Result, Stms (Lore (AllocM fromlore tolore)))
forall (m :: * -> *) a.
MonadBinder m =>
m a -> m (a, Stms (Lore m))
collectStms (AllocM fromlore tolore Result
 -> AllocM
      fromlore tolore (Result, Stms (Lore (AllocM fromlore tolore))))
-> AllocM fromlore tolore Result
-> AllocM
     fromlore tolore (Result, Stms (Lore (AllocM fromlore tolore)))
forall a b. (a -> b) -> a -> b
$ do
      Result
res' <- (Maybe Space -> SubExp -> AllocM fromlore tolore SubExp)
-> [Maybe Space] -> Result -> AllocM fromlore tolore Result
forall (m :: * -> *) a b c.
Applicative m =>
(a -> b -> m c) -> [a] -> [b] -> m [c]
zipWithM Maybe Space -> SubExp -> AllocM fromlore tolore SubExp
forall fromlore tolore.
(Allocable fromlore tolore,
 Allocator tolore (AllocM fromlore tolore)) =>
Maybe Space -> SubExp -> AllocM fromlore tolore SubExp
ensureDirect [Maybe Space]
space_oks' Result
res
      let (Result
ctx_res, Result
val_res) = Int -> Result -> (Result, Result)
forall a. Int -> [a] -> ([a], [a])
splitFromEnd Int
num_vals Result
res'
      Result
mem_ctx_res <- [Result] -> Result
forall (t :: * -> *) a. Foldable t => t [a] -> [a]
concat ([Result] -> Result)
-> AllocM fromlore tolore [Result] -> AllocM fromlore tolore Result
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> (SubExp -> AllocM fromlore tolore Result)
-> Result -> AllocM fromlore tolore [Result]
forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
(a -> m b) -> t a -> m (t b)
mapM SubExp -> AllocM fromlore tolore Result
forall fromlore tolore.
(Allocable fromlore tolore,
 Allocator tolore (AllocM fromlore tolore)) =>
SubExp -> AllocM fromlore tolore Result
bodyReturnMemCtx Result
val_res
      Result -> AllocM fromlore tolore Result
forall (m :: * -> *) a. Monad m => a -> m a
return (Result -> AllocM fromlore tolore Result)
-> Result -> AllocM fromlore tolore Result
forall a b. (a -> b) -> a -> b
$ Result
ctx_res Result -> Result -> Result
forall a. Semigroup a => a -> a -> a
<> Result
mem_ctx_res Result -> Result -> Result
forall a. Semigroup a => a -> a -> a
<> Result
val_res
    Body tolore -> AllocM fromlore tolore (Body tolore)
forall (m :: * -> *) a. Monad m => a -> m a
return (Body tolore -> AllocM fromlore tolore (Body tolore))
-> Body tolore -> AllocM fromlore tolore (Body tolore)
forall a b. (a -> b) -> a -> b
$ BodyAttr tolore -> Stms tolore -> Result -> Body tolore
forall lore. BodyAttr lore -> Stms lore -> Result -> BodyT lore
Body () (Stms tolore
bnds'Stms tolore -> Stms tolore -> Stms tolore
forall a. Semigroup a => a -> a -> a
<>Stms tolore
allocs) Result
res''
  where num_vals :: Int
num_vals = [Maybe Space] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length [Maybe Space]
space_oks
        space_oks' :: [Maybe Space]
space_oks' = Int -> Maybe Space -> [Maybe Space]
forall a. Int -> a -> [a]
replicate (Result -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length Result
res Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
num_vals) Maybe Space
forall a. Maybe a
Nothing [Maybe Space] -> [Maybe Space] -> [Maybe Space]
forall a. [a] -> [a] -> [a]
++ [Maybe Space]
space_oks

ensureDirect :: (Allocable fromlore tolore, Allocator tolore (AllocM fromlore tolore)) =>
                Maybe Space -> SubExp -> AllocM fromlore tolore SubExp
ensureDirect :: Maybe Space -> SubExp -> AllocM fromlore tolore SubExp
ensureDirect Maybe Space
_ se :: SubExp
se@Constant{} = SubExp -> AllocM fromlore tolore SubExp
forall (m :: * -> *) a. Monad m => a -> m a
return SubExp
se
ensureDirect Maybe Space
space_ok (Var VName
v) = do
  Bool
bt <- Type -> Bool
forall shape u. TypeBase shape u -> Bool
primType (Type -> Bool)
-> AllocM fromlore tolore Type -> AllocM fromlore tolore Bool
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> VName -> AllocM fromlore tolore Type
forall lore (m :: * -> *). HasScope lore m => VName -> m Type
lookupType VName
v
  if Bool
bt
    then SubExp -> AllocM fromlore tolore SubExp
forall (m :: * -> *) a. Monad m => a -> m a
return (SubExp -> AllocM fromlore tolore SubExp)
-> SubExp -> AllocM fromlore tolore SubExp
forall a b. (a -> b) -> a -> b
$ VName -> SubExp
Var VName
v
    else do (VName
_, SubExp
v') <- Maybe Space -> VName -> AllocM fromlore tolore (VName, SubExp)
forall fromlore tolore.
(Allocable fromlore tolore,
 Allocator tolore (AllocM fromlore tolore)) =>
Maybe Space -> VName -> AllocM fromlore tolore (VName, SubExp)
ensureDirectArray Maybe Space
space_ok VName
v
            SubExp -> AllocM fromlore tolore SubExp
forall (m :: * -> *) a. Monad m => a -> m a
return SubExp
v'

allocInStms :: (Allocable fromlore tolore) =>
               Stms fromlore -> (Stms tolore -> AllocM fromlore tolore a)
            -> AllocM fromlore tolore a
allocInStms :: Stms fromlore
-> (Stms tolore -> AllocM fromlore tolore a)
-> AllocM fromlore tolore a
allocInStms Stms fromlore
origstms Stms tolore -> AllocM fromlore tolore a
m = [Stm fromlore] -> Stms tolore -> AllocM fromlore tolore a
allocInStms' (Stms fromlore -> [Stm fromlore]
forall lore. Stms lore -> [Stm lore]
stmsToList Stms fromlore
origstms) Stms tolore
forall a. Monoid a => a
mempty
  where allocInStms' :: [Stm fromlore] -> Stms tolore -> AllocM fromlore tolore a
allocInStms' [] Stms tolore
stms' =
          Stms tolore -> AllocM fromlore tolore a
m Stms tolore
stms'
        allocInStms' (Stm fromlore
x:[Stm fromlore]
xs) Stms tolore
stms' = do
          Stms tolore
allocstms <- Stm fromlore -> AllocM fromlore tolore (Stms tolore)
forall tolore fromlore.
(PrettyLore fromlore, AllocOp (Op tolore), Checkable tolore,
 OpReturns tolore, SizeSubst (Op tolore), BinderOps tolore,
 FParamAttr tolore ~ FParamMem, FParamAttr fromlore ~ DeclType,
 LParamAttr tolore ~ LetAttrMem, ExpAttr tolore ~ (),
 LParamAttr fromlore ~ Type, LetAttr tolore ~ LetAttrMem,
 BodyAttr tolore ~ (),
 BranchType fromlore ~ TypeBase ExtShape NoUniqueness,
 RetType tolore ~ RetTypeMem, BodyAttr fromlore ~ (),
 RetType fromlore ~ DeclExtType,
 BranchType tolore ~ BranchTypeMem) =>
Stm fromlore -> AllocM fromlore tolore (Stms tolore)
allocInStm' Stm fromlore
x
          Scope tolore
-> AllocM fromlore tolore a -> AllocM fromlore tolore a
forall lore (m :: * -> *) a.
LocalScope lore m =>
Scope lore -> m a -> m a
localScope (Stms tolore -> Scope tolore
forall lore a. Scoped lore a => a -> Scope lore
scopeOf Stms tolore
allocstms) (AllocM fromlore tolore a -> AllocM fromlore tolore a)
-> AllocM fromlore tolore a -> AllocM fromlore tolore a
forall a b. (a -> b) -> a -> b
$ do
            let stms_substs :: Map VName SubExp
stms_substs = (Stm tolore -> Map VName SubExp) -> Stms tolore -> Map VName SubExp
forall (t :: * -> *) m a.
(Foldable t, Monoid m) =>
(a -> m) -> t a -> m
foldMap Stm tolore -> Map VName SubExp
forall lore. SizeSubst (Op lore) => Stm lore -> Map VName SubExp
sizeSubst Stms tolore
allocstms
                stms_consts :: Set VName
stms_consts = (Stm tolore -> Set VName) -> Stms tolore -> Set VName
forall (t :: * -> *) m a.
(Foldable t, Monoid m) =>
(a -> m) -> t a -> m
foldMap Stm tolore -> Set VName
forall lore. SizeSubst (Op lore) => Stm lore -> Set VName
stmConsts Stms tolore
allocstms
                f :: AllocEnv fromlore tolore -> AllocEnv fromlore tolore
f AllocEnv fromlore tolore
env = AllocEnv fromlore tolore
env { chunkMap :: Map VName SubExp
chunkMap = Map VName SubExp
stms_substs Map VName SubExp -> Map VName SubExp -> Map VName SubExp
forall a. Semigroup a => a -> a -> a
<> AllocEnv fromlore tolore -> Map VName SubExp
forall fromlore tolore.
AllocEnv fromlore tolore -> Map VName SubExp
chunkMap AllocEnv fromlore tolore
env
                            , envConsts :: Set VName
envConsts = Set VName
stms_consts Set VName -> Set VName -> Set VName
forall a. Semigroup a => a -> a -> a
<> AllocEnv fromlore tolore -> Set VName
forall fromlore tolore. AllocEnv fromlore tolore -> Set VName
envConsts AllocEnv fromlore tolore
env
                            }
            (AllocEnv fromlore tolore -> AllocEnv fromlore tolore)
-> AllocM fromlore tolore a -> AllocM fromlore tolore a
forall r (m :: * -> *) a. MonadReader r m => (r -> r) -> m a -> m a
local AllocEnv fromlore tolore -> AllocEnv fromlore tolore
f (AllocM fromlore tolore a -> AllocM fromlore tolore a)
-> AllocM fromlore tolore a -> AllocM fromlore tolore a
forall a b. (a -> b) -> a -> b
$ [Stm fromlore] -> Stms tolore -> AllocM fromlore tolore a
allocInStms' [Stm fromlore]
xs (Stms tolore
stms'Stms tolore -> Stms tolore -> Stms tolore
forall a. Semigroup a => a -> a -> a
<>Stms tolore
allocstms)
        allocInStm' :: Stm fromlore -> AllocM fromlore tolore (Stms tolore)
allocInStm' Stm fromlore
bnd = do
          ((),Stms tolore
stms') <- AllocM fromlore tolore ()
-> AllocM
     fromlore tolore ((), Stms (Lore (AllocM fromlore tolore)))
forall (m :: * -> *) a.
MonadBinder m =>
m a -> m (a, Stms (Lore m))
collectStms (AllocM fromlore tolore ()
 -> AllocM
      fromlore tolore ((), Stms (Lore (AllocM fromlore tolore))))
-> AllocM fromlore tolore ()
-> AllocM
     fromlore tolore ((), Stms (Lore (AllocM fromlore tolore)))
forall a b. (a -> b) -> a -> b
$ Certificates
-> AllocM fromlore tolore () -> AllocM fromlore tolore ()
forall (m :: * -> *) a. MonadBinder m => Certificates -> m a -> m a
certifying (Stm fromlore -> Certificates
forall lore. Stm lore -> Certificates
stmCerts Stm fromlore
bnd) (AllocM fromlore tolore () -> AllocM fromlore tolore ())
-> AllocM fromlore tolore () -> AllocM fromlore tolore ()
forall a b. (a -> b) -> a -> b
$ Stm fromlore -> AllocM fromlore tolore ()
forall fromlore tolore.
(Allocable fromlore tolore,
 Allocator tolore (AllocM fromlore tolore)) =>
Stm fromlore -> AllocM fromlore tolore ()
allocInStm Stm fromlore
bnd
          Stms tolore -> AllocM fromlore tolore (Stms tolore)
forall (m :: * -> *) a. Monad m => a -> m a
return Stms tolore
stms'

allocInStm :: (Allocable fromlore tolore, Allocator tolore (AllocM fromlore tolore)) =>
              Stm fromlore -> AllocM fromlore tolore ()
allocInStm :: Stm fromlore -> AllocM fromlore tolore ()
allocInStm (Let (Pattern [PatElemT (LetAttr fromlore)]
sizeElems [PatElemT (LetAttr fromlore)]
valElems) StmAux (ExpAttr fromlore)
_ Exp fromlore
e) = do
  Exp tolore
e' <- Exp fromlore -> AllocM fromlore tolore (Exp tolore)
forall fromlore tolore.
(Allocable fromlore tolore,
 Allocator tolore (AllocM fromlore tolore)) =>
Exp fromlore -> AllocM fromlore tolore (Exp tolore)
allocInExp Exp fromlore
e
  let sizeidents :: [Ident]
sizeidents = (PatElemT (LetAttr fromlore) -> Ident)
-> [PatElemT (LetAttr fromlore)] -> [Ident]
forall a b. (a -> b) -> [a] -> [b]
map PatElemT (LetAttr fromlore) -> Ident
forall attr. Typed attr => PatElemT attr -> Ident
patElemIdent [PatElemT (LetAttr fromlore)]
sizeElems
      validents :: [Ident]
validents = (PatElemT (LetAttr fromlore) -> Ident)
-> [PatElemT (LetAttr fromlore)] -> [Ident]
forall a b. (a -> b) -> [a] -> [b]
map PatElemT (LetAttr fromlore) -> Ident
forall attr. Typed attr => PatElemT attr -> Ident
patElemIdent [PatElemT (LetAttr fromlore)]
valElems
  Stm tolore
bnd <- [Ident]
-> [Ident] -> Exp tolore -> AllocM fromlore tolore (Stm tolore)
forall lore (m :: * -> *).
(Allocator lore m, ExpAttr lore ~ ()) =>
[Ident] -> [Ident] -> Exp lore -> m (Stm lore)
allocsForStm [Ident]
sizeidents [Ident]
validents Exp tolore
e'
  Stm (Lore (AllocM fromlore tolore)) -> AllocM fromlore tolore ()
forall (m :: * -> *). MonadBinder m => Stm (Lore m) -> m ()
addStm Stm tolore
Stm (Lore (AllocM fromlore tolore))
bnd

allocInExp :: (Allocable fromlore tolore, Allocator tolore (AllocM fromlore tolore)) =>
              Exp fromlore -> AllocM fromlore tolore (Exp tolore)
allocInExp :: Exp fromlore -> AllocM fromlore tolore (Exp tolore)
allocInExp (DoLoop [(FParam fromlore, SubExp)]
ctx [(FParam fromlore, SubExp)]
val LoopForm fromlore
form (Body () Stms fromlore
bodybnds Result
bodyres)) =
  [VName]
-> [(FParam fromlore, SubExp)]
-> ([FParam tolore]
    -> [FParam tolore]
    -> (Result -> AllocM fromlore tolore (Result, Result))
    -> AllocM fromlore tolore (Exp tolore))
-> AllocM fromlore tolore (Exp tolore)
forall fromlore tolore a.
(Allocable fromlore tolore,
 Allocator tolore (AllocM fromlore tolore)) =>
[VName]
-> [(FParam fromlore, SubExp)]
-> ([FParam tolore]
    -> [FParam tolore]
    -> (Result -> AllocM fromlore tolore (Result, Result))
    -> AllocM fromlore tolore a)
-> AllocM fromlore tolore a
allocInMergeParams [VName]
forall a. Monoid a => a
mempty [(FParam fromlore, SubExp)]
ctx (([FParam tolore]
  -> [FParam tolore]
  -> (Result -> AllocM fromlore tolore (Result, Result))
  -> AllocM fromlore tolore (Exp tolore))
 -> AllocM fromlore tolore (Exp tolore))
-> ([FParam tolore]
    -> [FParam tolore]
    -> (Result -> AllocM fromlore tolore (Result, Result))
    -> AllocM fromlore tolore (Exp tolore))
-> AllocM fromlore tolore (Exp tolore)
forall a b. (a -> b) -> a -> b
$ \[FParam tolore]
_ [FParam tolore]
ctxparams' Result -> AllocM fromlore tolore (Result, Result)
_ ->
  [VName]
-> [(FParam fromlore, SubExp)]
-> ([FParam tolore]
    -> [FParam tolore]
    -> (Result -> AllocM fromlore tolore (Result, Result))
    -> AllocM fromlore tolore (Exp tolore))
-> AllocM fromlore tolore (Exp tolore)
forall fromlore tolore a.
(Allocable fromlore tolore,
 Allocator tolore (AllocM fromlore tolore)) =>
[VName]
-> [(FParam fromlore, SubExp)]
-> ([FParam tolore]
    -> [FParam tolore]
    -> (Result -> AllocM fromlore tolore (Result, Result))
    -> AllocM fromlore tolore a)
-> AllocM fromlore tolore a
allocInMergeParams ((Param FParamMem -> VName) -> [Param FParamMem] -> [VName]
forall a b. (a -> b) -> [a] -> [b]
map Param FParamMem -> VName
forall attr. Param attr -> VName
paramName [FParam tolore]
[Param FParamMem]
ctxparams') [(FParam fromlore, SubExp)]
val (([FParam tolore]
  -> [FParam tolore]
  -> (Result -> AllocM fromlore tolore (Result, Result))
  -> AllocM fromlore tolore (Exp tolore))
 -> AllocM fromlore tolore (Exp tolore))
-> ([FParam tolore]
    -> [FParam tolore]
    -> (Result -> AllocM fromlore tolore (Result, Result))
    -> AllocM fromlore tolore (Exp tolore))
-> AllocM fromlore tolore (Exp tolore)
forall a b. (a -> b) -> a -> b
$
  \[FParam tolore]
new_ctx_params [FParam tolore]
valparams' Result -> AllocM fromlore tolore (Result, Result)
mk_loop_val -> do
  LoopForm tolore
form' <- LoopForm fromlore -> AllocM fromlore tolore (LoopForm tolore)
forall fromlore tolore.
(Allocable fromlore tolore,
 Allocator tolore (AllocM fromlore tolore)) =>
LoopForm fromlore -> AllocM fromlore tolore (LoopForm tolore)
allocInLoopForm LoopForm fromlore
form
  Scope tolore
-> AllocM fromlore tolore (Exp tolore)
-> AllocM fromlore tolore (Exp tolore)
forall lore (m :: * -> *) a.
LocalScope lore m =>
Scope lore -> m a -> m a
localScope (LoopForm tolore -> Scope tolore
forall lore a. Scoped lore a => a -> Scope lore
scopeOf LoopForm tolore
form') (AllocM fromlore tolore (Exp tolore)
 -> AllocM fromlore tolore (Exp tolore))
-> AllocM fromlore tolore (Exp tolore)
-> AllocM fromlore tolore (Exp tolore)
forall a b. (a -> b) -> a -> b
$ do
    (Result
valinit_ctx, Result
valinit') <- Result -> AllocM fromlore tolore (Result, Result)
mk_loop_val Result
valinit
    BodyT tolore
body' <- AllocM fromlore tolore (Body (Lore (AllocM fromlore tolore)))
-> AllocM fromlore tolore (Body (Lore (AllocM fromlore tolore)))
forall (m :: * -> *).
MonadBinder m =>
m (Body (Lore m)) -> m (Body (Lore m))
insertStmsM (AllocM fromlore tolore (Body (Lore (AllocM fromlore tolore)))
 -> AllocM fromlore tolore (Body (Lore (AllocM fromlore tolore))))
-> AllocM fromlore tolore (Body (Lore (AllocM fromlore tolore)))
-> AllocM fromlore tolore (Body (Lore (AllocM fromlore tolore)))
forall a b. (a -> b) -> a -> b
$ Stms fromlore
-> (Stms tolore -> AllocM fromlore tolore (BodyT tolore))
-> AllocM fromlore tolore (BodyT tolore)
forall fromlore tolore a.
Allocable fromlore tolore =>
Stms fromlore
-> (Stms tolore -> AllocM fromlore tolore a)
-> AllocM fromlore tolore a
allocInStms Stms fromlore
bodybnds ((Stms tolore -> AllocM fromlore tolore (BodyT tolore))
 -> AllocM fromlore tolore (BodyT tolore))
-> (Stms tolore -> AllocM fromlore tolore (BodyT tolore))
-> AllocM fromlore tolore (BodyT tolore)
forall a b. (a -> b) -> a -> b
$ \Stms tolore
bodybnds' -> do
      ((Result
val_ses,Result
valres'),Stms tolore
val_retbnds) <- AllocM fromlore tolore (Result, Result)
-> AllocM
     fromlore
     tolore
     ((Result, Result), Stms (Lore (AllocM fromlore tolore)))
forall (m :: * -> *) a.
MonadBinder m =>
m a -> m (a, Stms (Lore m))
collectStms (AllocM fromlore tolore (Result, Result)
 -> AllocM
      fromlore
      tolore
      ((Result, Result), Stms (Lore (AllocM fromlore tolore))))
-> AllocM fromlore tolore (Result, Result)
-> AllocM
     fromlore
     tolore
     ((Result, Result), Stms (Lore (AllocM fromlore tolore)))
forall a b. (a -> b) -> a -> b
$ Result -> AllocM fromlore tolore (Result, Result)
mk_loop_val Result
valres
      BodyT tolore -> AllocM fromlore tolore (BodyT tolore)
forall (m :: * -> *) a. Monad m => a -> m a
return (BodyT tolore -> AllocM fromlore tolore (BodyT tolore))
-> BodyT tolore -> AllocM fromlore tolore (BodyT tolore)
forall a b. (a -> b) -> a -> b
$ BodyAttr tolore -> Stms tolore -> Result -> BodyT tolore
forall lore. BodyAttr lore -> Stms lore -> Result -> BodyT lore
Body () (Stms tolore
bodybnds'Stms tolore -> Stms tolore -> Stms tolore
forall a. Semigroup a => a -> a -> a
<>Stms tolore
val_retbnds) (Result
ctxresResult -> Result -> Result
forall a. [a] -> [a] -> [a]
++Result
val_sesResult -> Result -> Result
forall a. [a] -> [a] -> [a]
++Result
valres')
    Exp tolore -> AllocM fromlore tolore (Exp tolore)
forall (m :: * -> *) a. Monad m => a -> m a
return (Exp tolore -> AllocM fromlore tolore (Exp tolore))
-> Exp tolore -> AllocM fromlore tolore (Exp tolore)
forall a b. (a -> b) -> a -> b
$
      [(FParam tolore, SubExp)]
-> [(FParam tolore, SubExp)]
-> LoopForm tolore
-> BodyT tolore
-> Exp tolore
forall lore.
[(FParam lore, SubExp)]
-> [(FParam lore, SubExp)]
-> LoopForm lore
-> BodyT lore
-> ExpT lore
DoLoop
      ([Param FParamMem] -> Result -> [(Param FParamMem, SubExp)]
forall a b. [a] -> [b] -> [(a, b)]
zip ([FParam tolore]
[Param FParamMem]
ctxparams'[Param FParamMem] -> [Param FParamMem] -> [Param FParamMem]
forall a. [a] -> [a] -> [a]
++[FParam tolore]
[Param FParamMem]
new_ctx_params) (Result
ctxinitResult -> Result -> Result
forall a. [a] -> [a] -> [a]
++Result
valinit_ctx))
      ([Param FParamMem] -> Result -> [(Param FParamMem, SubExp)]
forall a b. [a] -> [b] -> [(a, b)]
zip [FParam tolore]
[Param FParamMem]
valparams' Result
valinit')
      LoopForm tolore
form' BodyT tolore
body'
  where ([Param DeclType]
_ctxparams, Result
ctxinit) = [(Param DeclType, SubExp)] -> ([Param DeclType], Result)
forall a b. [(a, b)] -> ([a], [b])
unzip [(Param DeclType, SubExp)]
[(FParam fromlore, SubExp)]
ctx
        ([Param DeclType]
_valparams, Result
valinit) = [(Param DeclType, SubExp)] -> ([Param DeclType], Result)
forall a b. [(a, b)] -> ([a], [b])
unzip [(Param DeclType, SubExp)]
[(FParam fromlore, SubExp)]
val
        (Result
ctxres, Result
valres) = Int -> Result -> (Result, Result)
forall a. Int -> [a] -> ([a], [a])
splitAt ([(Param DeclType, SubExp)] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length [(Param DeclType, SubExp)]
[(FParam fromlore, SubExp)]
ctx) Result
bodyres
allocInExp (Apply Name
fname [(SubExp, Diet)]
args [RetType fromlore]
rettype (Safety, SrcLoc, [SrcLoc])
loc) = do
  [(SubExp, Diet)]
args' <- [(SubExp, Diet)] -> AllocM fromlore tolore [(SubExp, Diet)]
forall fromlore tolore.
(Allocable fromlore tolore,
 Allocator tolore (AllocM fromlore tolore)) =>
[(SubExp, Diet)] -> AllocM fromlore tolore [(SubExp, Diet)]
funcallArgs [(SubExp, Diet)]
args
  Exp tolore -> AllocM fromlore tolore (Exp tolore)
forall (m :: * -> *) a. Monad m => a -> m a
return (Exp tolore -> AllocM fromlore tolore (Exp tolore))
-> Exp tolore -> AllocM fromlore tolore (Exp tolore)
forall a b. (a -> b) -> a -> b
$ Name
-> [(SubExp, Diet)]
-> [RetType tolore]
-> (Safety, SrcLoc, [SrcLoc])
-> Exp tolore
forall lore.
Name
-> [(SubExp, Diet)]
-> [RetType lore]
-> (Safety, SrcLoc, [SrcLoc])
-> ExpT lore
Apply Name
fname [(SubExp, Diet)]
args' ([DeclExtType] -> [RetTypeMem]
memoryInDeclExtType [DeclExtType]
[RetType fromlore]
rettype) (Safety, SrcLoc, [SrcLoc])
loc
allocInExp (If SubExp
cond BodyT fromlore
tbranch0 BodyT fromlore
fbranch0 (IfAttr [BranchType fromlore]
rets IfSort
ifsort)) = do
  let num_rets :: Int
num_rets = [TypeBase ExtShape NoUniqueness] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length [TypeBase ExtShape NoUniqueness]
[BranchType fromlore]
rets
  -- switch to the explicit-mem rep, but do nothing about results
  (BodyT tolore
tbranch, [Maybe MemBind]
tm_ixfs) <- Int
-> BodyT fromlore
-> AllocM fromlore tolore (BodyT tolore, [Maybe MemBind])
forall fromlore tolore.
(Allocable fromlore tolore,
 Allocator tolore (AllocM fromlore tolore)) =>
Int
-> Body fromlore
-> AllocM fromlore tolore (Body tolore, [Maybe MemBind])
allocInIfBody Int
num_rets BodyT fromlore
tbranch0
  (BodyT tolore
fbranch, [Maybe MemBind]
fm_ixfs) <- Int
-> BodyT fromlore
-> AllocM fromlore tolore (BodyT tolore, [Maybe MemBind])
forall fromlore tolore.
(Allocable fromlore tolore,
 Allocator tolore (AllocM fromlore tolore)) =>
Int
-> Body fromlore
-> AllocM fromlore tolore (Body tolore, [Maybe MemBind])
allocInIfBody Int
num_rets BodyT fromlore
fbranch0
  [Maybe Space]
tspaces <- Int -> BodyT tolore -> AllocM fromlore tolore [Maybe Space]
forall tolore (m :: * -> *).
(Mem tolore, LocalScope tolore m) =>
Int -> Body tolore -> m [Maybe Space]
mkSpaceOks Int
num_rets BodyT tolore
tbranch
  [Maybe Space]
fspaces <- Int -> BodyT tolore -> AllocM fromlore tolore [Maybe Space]
forall tolore (m :: * -> *).
(Mem tolore, LocalScope tolore m) =>
Int -> Body tolore -> m [Maybe Space]
mkSpaceOks Int
num_rets BodyT tolore
fbranch
  -- try to generalize (antiunify) the index functions of the then and else bodies
  let sp_substs :: [(Maybe Space, Maybe (ExtIxFun, [(PrimExp VName, PrimExp VName)]))]
sp_substs = ((Maybe Space, Maybe MemBind)
 -> (Maybe Space, Maybe MemBind)
 -> (Maybe Space,
     Maybe (ExtIxFun, [(PrimExp VName, PrimExp VName)])))
-> [(Maybe Space, Maybe MemBind)]
-> [(Maybe Space, Maybe MemBind)]
-> [(Maybe Space,
     Maybe (ExtIxFun, [(PrimExp VName, PrimExp VName)]))]
forall a b c. (a -> b -> c) -> [a] -> [b] -> [c]
zipWith (Maybe Space, Maybe MemBind)
-> (Maybe Space, Maybe MemBind)
-> (Maybe Space,
    Maybe (ExtIxFun, [(PrimExp VName, PrimExp VName)]))
generalize ([Maybe Space] -> [Maybe MemBind] -> [(Maybe Space, Maybe MemBind)]
forall a b. [a] -> [b] -> [(a, b)]
zip [Maybe Space]
tspaces [Maybe MemBind]
tm_ixfs) ([Maybe Space] -> [Maybe MemBind] -> [(Maybe Space, Maybe MemBind)]
forall a b. [a] -> [b] -> [(a, b)]
zip [Maybe Space]
fspaces [Maybe MemBind]
fm_ixfs)
      ([Maybe Space]
spaces, [Maybe (ExtIxFun, [(PrimExp VName, PrimExp VName)])]
subs) = [(Maybe Space, Maybe (ExtIxFun, [(PrimExp VName, PrimExp VName)]))]
-> ([Maybe Space],
    [Maybe (ExtIxFun, [(PrimExp VName, PrimExp VName)])])
forall a b. [(a, b)] -> ([a], [b])
unzip [(Maybe Space, Maybe (ExtIxFun, [(PrimExp VName, PrimExp VName)]))]
sp_substs
      tsubs :: [Maybe (ExtIxFun, [PrimExp VName])]
tsubs = (Maybe (ExtIxFun, [(PrimExp VName, PrimExp VName)])
 -> Maybe (ExtIxFun, [PrimExp VName]))
-> [Maybe (ExtIxFun, [(PrimExp VName, PrimExp VName)])]
-> [Maybe (ExtIxFun, [PrimExp VName])]
forall a b. (a -> b) -> [a] -> [b]
map (((PrimExp VName, PrimExp VName) -> PrimExp VName)
-> Maybe (ExtIxFun, [(PrimExp VName, PrimExp VName)])
-> Maybe (ExtIxFun, [PrimExp VName])
forall a.
((a, a) -> a)
-> Maybe (ExtIxFun, [(a, a)]) -> Maybe (ExtIxFun, [a])
selectSub (PrimExp VName, PrimExp VName) -> PrimExp VName
forall a b. (a, b) -> a
fst) [Maybe (ExtIxFun, [(PrimExp VName, PrimExp VName)])]
subs
      fsubs :: [Maybe (ExtIxFun, [PrimExp VName])]
fsubs = (Maybe (ExtIxFun, [(PrimExp VName, PrimExp VName)])
 -> Maybe (ExtIxFun, [PrimExp VName]))
-> [Maybe (ExtIxFun, [(PrimExp VName, PrimExp VName)])]
-> [Maybe (ExtIxFun, [PrimExp VName])]
forall a b. (a -> b) -> [a] -> [b]
map (((PrimExp VName, PrimExp VName) -> PrimExp VName)
-> Maybe (ExtIxFun, [(PrimExp VName, PrimExp VName)])
-> Maybe (ExtIxFun, [PrimExp VName])
forall a.
((a, a) -> a)
-> Maybe (ExtIxFun, [(a, a)]) -> Maybe (ExtIxFun, [a])
selectSub (PrimExp VName, PrimExp VName) -> PrimExp VName
forall a b. (a, b) -> b
snd) [Maybe (ExtIxFun, [(PrimExp VName, PrimExp VName)])]
subs
  (BodyT tolore
tbranch', [BranchTypeMem]
trets) <- [TypeBase ExtShape NoUniqueness]
-> BodyT tolore
-> [Maybe Space]
-> [Maybe (ExtIxFun, [PrimExp VName])]
-> AllocM fromlore tolore (BodyT tolore, [BranchTypeMem])
forall fromlore tolore.
(Allocable fromlore tolore,
 Allocator tolore (AllocM fromlore tolore)) =>
[TypeBase ExtShape NoUniqueness]
-> Body tolore
-> [Maybe Space]
-> [Maybe (ExtIxFun, [PrimExp VName])]
-> AllocM fromlore tolore (Body tolore, [BranchTypeMem])
addResCtxInIfBody [TypeBase ExtShape NoUniqueness]
[BranchType fromlore]
rets BodyT tolore
tbranch [Maybe Space]
spaces [Maybe (ExtIxFun, [PrimExp VName])]
tsubs
  (BodyT tolore
fbranch', [BranchTypeMem]
frets) <- [TypeBase ExtShape NoUniqueness]
-> BodyT tolore
-> [Maybe Space]
-> [Maybe (ExtIxFun, [PrimExp VName])]
-> AllocM fromlore tolore (BodyT tolore, [BranchTypeMem])
forall fromlore tolore.
(Allocable fromlore tolore,
 Allocator tolore (AllocM fromlore tolore)) =>
[TypeBase ExtShape NoUniqueness]
-> Body tolore
-> [Maybe Space]
-> [Maybe (ExtIxFun, [PrimExp VName])]
-> AllocM fromlore tolore (Body tolore, [BranchTypeMem])
addResCtxInIfBody [TypeBase ExtShape NoUniqueness]
[BranchType fromlore]
rets BodyT tolore
fbranch [Maybe Space]
spaces [Maybe (ExtIxFun, [PrimExp VName])]
fsubs
  if [BranchTypeMem]
frets [BranchTypeMem] -> [BranchTypeMem] -> Bool
forall a. Eq a => a -> a -> Bool
/= [BranchTypeMem]
trets then String -> AllocM fromlore tolore (Exp tolore)
forall a. HasCallStack => String -> a
error String
"In allocInExp, IF case: antiunification of then/else produce different ExtInFn!"
    else do -- above is a sanity check; implementation continues on else branch
    let res_then :: Result
res_then = BodyT tolore -> Result
forall lore. BodyT lore -> Result
bodyResult BodyT tolore
tbranch'
        res_else :: Result
res_else = BodyT tolore -> Result
forall lore. BodyT lore -> Result
bodyResult BodyT tolore
fbranch'
        size_ext :: Int
size_ext = Result -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length Result
res_then Int -> Int -> Int
forall a. Num a => a -> a -> a
- [BranchTypeMem] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length [BranchTypeMem]
trets
        ([(SubExp, SubExp, Int)]
ind_ses0, [(SubExp, SubExp, Int)]
r_then_else) =
            ((SubExp, SubExp, Int) -> Bool)
-> [(SubExp, SubExp, Int)]
-> ([(SubExp, SubExp, Int)], [(SubExp, SubExp, Int)])
forall a. (a -> Bool) -> [a] -> ([a], [a])
partition (\(SubExp
r_then, SubExp
r_else, Int
_) -> SubExp
r_then SubExp -> SubExp -> Bool
forall a. Eq a => a -> a -> Bool
== SubExp
r_else) ([(SubExp, SubExp, Int)]
 -> ([(SubExp, SubExp, Int)], [(SubExp, SubExp, Int)]))
-> [(SubExp, SubExp, Int)]
-> ([(SubExp, SubExp, Int)], [(SubExp, SubExp, Int)])
forall a b. (a -> b) -> a -> b
$
            Result -> Result -> [Int] -> [(SubExp, SubExp, Int)]
forall a b c. [a] -> [b] -> [c] -> [(a, b, c)]
zip3 Result
res_then Result
res_else [Int
0 .. Int
size_ext Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
1]
        (Result
r_then_ext, Result
r_else_ext, [Int]
_) = [(SubExp, SubExp, Int)] -> (Result, Result, [Int])
forall a b c. [(a, b, c)] -> ([a], [b], [c])
unzip3 [(SubExp, SubExp, Int)]
r_then_else
        ind_ses :: [(Int, SubExp)]
ind_ses = ((SubExp, SubExp, Int) -> Int -> (Int, SubExp))
-> [(SubExp, SubExp, Int)] -> [Int] -> [(Int, SubExp)]
forall a b c. (a -> b -> c) -> [a] -> [b] -> [c]
zipWith (\(SubExp
se, SubExp
_, Int
i) Int
k -> (Int
iInt -> Int -> Int
forall a. Num a => a -> a -> a
-Int
k, SubExp
se)) [(SubExp, SubExp, Int)]
ind_ses0
                  [Int
0 .. [(SubExp, SubExp, Int)] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length [(SubExp, SubExp, Int)]
ind_ses0 Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
1]
        rets'' :: [BranchTypeMem]
rets'' = ([BranchTypeMem] -> (Int, SubExp) -> [BranchTypeMem])
-> [BranchTypeMem] -> [(Int, SubExp)] -> [BranchTypeMem]
forall (t :: * -> *) b a.
Foldable t =>
(b -> a -> b) -> b -> t a -> b
foldl (\[BranchTypeMem]
acc (Int
i, SubExp
se) -> Int -> SubExp -> [BranchTypeMem] -> [BranchTypeMem]
forall t. FixExt t => Int -> SubExp -> t -> t
fixExt Int
i SubExp
se [BranchTypeMem]
acc) [BranchTypeMem]
trets [(Int, SubExp)]
ind_ses
        tbranch'' :: BodyT tolore
tbranch'' = BodyT tolore
tbranch' { bodyResult :: Result
bodyResult = Result
r_then_ext Result -> Result -> Result
forall a. [a] -> [a] -> [a]
++ Int -> Result -> Result
forall a. Int -> [a] -> [a]
drop Int
size_ext Result
res_then }
        fbranch'' :: BodyT tolore
fbranch'' = BodyT tolore
fbranch' { bodyResult :: Result
bodyResult = Result
r_else_ext Result -> Result -> Result
forall a. [a] -> [a] -> [a]
++ Int -> Result -> Result
forall a. Int -> [a] -> [a]
drop Int
size_ext Result
res_else }
        res_if_expr :: Exp tolore
res_if_expr = SubExp
-> BodyT tolore
-> BodyT tolore
-> IfAttr (BranchType tolore)
-> Exp tolore
forall lore.
SubExp
-> BodyT lore
-> BodyT lore
-> IfAttr (BranchType lore)
-> ExpT lore
If SubExp
cond BodyT tolore
tbranch'' BodyT tolore
fbranch'' (IfAttr (BranchType tolore) -> Exp tolore)
-> IfAttr (BranchType tolore) -> Exp tolore
forall a b. (a -> b) -> a -> b
$ [BranchTypeMem] -> IfSort -> IfAttr BranchTypeMem
forall rt. [rt] -> IfSort -> IfAttr rt
IfAttr [BranchTypeMem]
rets'' IfSort
ifsort
    Exp tolore -> AllocM fromlore tolore (Exp tolore)
forall (m :: * -> *) a. Monad m => a -> m a
return Exp tolore
res_if_expr
      where generalize :: (Maybe Space, Maybe MemBind) -> (Maybe Space, Maybe MemBind)
                       -> (Maybe Space, Maybe (ExtIxFun, [(PrimExp VName, PrimExp VName)]))
            generalize :: (Maybe Space, Maybe MemBind)
-> (Maybe Space, Maybe MemBind)
-> (Maybe Space,
    Maybe (ExtIxFun, [(PrimExp VName, PrimExp VName)]))
generalize (Just Space
sp1, Just (ArrayIn VName
_ IxFun
ixf1)) (Just Space
sp2, Just (ArrayIn VName
_ IxFun
ixf2)) =
              if Space
sp1 Space -> Space -> Bool
forall a. Eq a => a -> a -> Bool
/= Space
sp2 then (Space -> Maybe Space
forall a. a -> Maybe a
Just Space
sp1, Maybe (ExtIxFun, [(PrimExp VName, PrimExp VName)])
forall a. Maybe a
Nothing)
              else case IxFun
-> IxFun -> Maybe (ExtIxFun, [(PrimExp VName, PrimExp VName)])
forall v.
Eq v =>
IxFun (PrimExp v)
-> IxFun (PrimExp v)
-> Maybe (IxFun (PrimExp (Ext v)), [(PrimExp v, PrimExp v)])
IxFun.leastGeneralGeneralization IxFun
ixf1 IxFun
ixf2 of
                Just (ExtIxFun
ixf, [(PrimExp VName, PrimExp VName)]
m) -> (Space -> Maybe Space
forall a. a -> Maybe a
Just Space
sp1, (ExtIxFun, [(PrimExp VName, PrimExp VName)])
-> Maybe (ExtIxFun, [(PrimExp VName, PrimExp VName)])
forall a. a -> Maybe a
Just (ExtIxFun
ixf, [(PrimExp VName, PrimExp VName)]
m))
                Maybe (ExtIxFun, [(PrimExp VName, PrimExp VName)])
Nothing -> (Space -> Maybe Space
forall a. a -> Maybe a
Just Space
sp1, Maybe (ExtIxFun, [(PrimExp VName, PrimExp VName)])
forall a. Maybe a
Nothing)
            generalize (Maybe Space
mbsp1, Maybe MemBind
_) (Maybe Space, Maybe MemBind)
_ = (Maybe Space
mbsp1, Maybe (ExtIxFun, [(PrimExp VName, PrimExp VName)])
forall a. Maybe a
Nothing)

            selectSub :: ((a, a) -> a) -> Maybe (ExtIxFun, [(a, a)]) ->
                         Maybe (ExtIxFun, [a])
            selectSub :: ((a, a) -> a)
-> Maybe (ExtIxFun, [(a, a)]) -> Maybe (ExtIxFun, [a])
selectSub (a, a) -> a
f (Just (ExtIxFun
ixfn, [(a, a)]
m)) = (ExtIxFun, [a]) -> Maybe (ExtIxFun, [a])
forall a. a -> Maybe a
Just (ExtIxFun
ixfn, ((a, a) -> a) -> [(a, a)] -> [a]
forall a b. (a -> b) -> [a] -> [b]
map (a, a) -> a
f [(a, a)]
m)
            selectSub (a, a) -> a
_ Maybe (ExtIxFun, [(a, a)])
Nothing = Maybe (ExtIxFun, [a])
forall a. Maybe a
Nothing

            -- | Just introduces the new representation (index functions); but
            -- does not unify (e.g., does not ensures direct); implementation
            -- extends `allocInBodyNoDirect`, but also return `MemBind`
            allocInIfBody :: (Allocable fromlore tolore, Allocator tolore (AllocM fromlore tolore)) =>
                             Int -> Body fromlore -> AllocM fromlore tolore (Body tolore, [Maybe MemBind])
            allocInIfBody :: Int
-> Body fromlore
-> AllocM fromlore tolore (Body tolore, [Maybe MemBind])
allocInIfBody Int
num_vals (Body BodyAttr fromlore
_ Stms fromlore
bnds Result
res) =
              Stms fromlore
-> (Stms tolore
    -> AllocM fromlore tolore (Body tolore, [Maybe MemBind]))
-> AllocM fromlore tolore (Body tolore, [Maybe MemBind])
forall fromlore tolore a.
Allocable fromlore tolore =>
Stms fromlore
-> (Stms tolore -> AllocM fromlore tolore a)
-> AllocM fromlore tolore a
allocInStms Stms fromlore
bnds ((Stms tolore
  -> AllocM fromlore tolore (Body tolore, [Maybe MemBind]))
 -> AllocM fromlore tolore (Body tolore, [Maybe MemBind]))
-> (Stms tolore
    -> AllocM fromlore tolore (Body tolore, [Maybe MemBind]))
-> AllocM fromlore tolore (Body tolore, [Maybe MemBind])
forall a b. (a -> b) -> a -> b
$ \Stms tolore
bnds' -> do
                let (Result
_, Result
val_res) = Int -> Result -> (Result, Result)
forall a. Int -> [a] -> ([a], [a])
splitFromEnd Int
num_vals Result
res
                [Maybe MemBind]
mem_ixfs <- (SubExp -> AllocM fromlore tolore (Maybe MemBind))
-> Result -> AllocM fromlore tolore [Maybe MemBind]
forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
(a -> m b) -> t a -> m (t b)
mapM SubExp -> AllocM fromlore tolore (Maybe MemBind)
forall lore (m :: * -> *).
(HasScope lore m, Monad m, AllocOp (Op lore), Checkable lore,
 OpReturns lore, FParamAttr lore ~ FParamMem,
 LParamAttr lore ~ LetAttrMem, LetAttr lore ~ LetAttrMem,
 RetType lore ~ RetTypeMem, BranchType lore ~ BranchTypeMem) =>
SubExp -> m (Maybe MemBind)
bodyReturnMIxf Result
val_res
                (Body tolore, [Maybe MemBind])
-> AllocM fromlore tolore (Body tolore, [Maybe MemBind])
forall (m :: * -> *) a. Monad m => a -> m a
return (BodyAttr tolore -> Stms tolore -> Result -> Body tolore
forall lore. BodyAttr lore -> Stms lore -> Result -> BodyT lore
Body () Stms tolore
bnds' Result
res, [Maybe MemBind]
mem_ixfs)
                  where
                    bodyReturnMIxf :: SubExp -> m (Maybe MemBind)
bodyReturnMIxf Constant{} = Maybe MemBind -> m (Maybe MemBind)
forall (m :: * -> *) a. Monad m => a -> m a
return Maybe MemBind
forall a. Maybe a
Nothing
                    bodyReturnMIxf (Var VName
v) = do
                      LetAttrMem
info <- VName -> m LetAttrMem
forall lore (m :: * -> *).
(HasScope lore m, Mem lore) =>
VName -> m LetAttrMem
lookupMemInfo VName
v
                      case LetAttrMem
info of
                        MemArray PrimType
_ptp Shape
_shp NoUniqueness
_u MemBind
mem_ixf -> Maybe MemBind -> m (Maybe MemBind)
forall (m :: * -> *) a. Monad m => a -> m a
return (Maybe MemBind -> m (Maybe MemBind))
-> Maybe MemBind -> m (Maybe MemBind)
forall a b. (a -> b) -> a -> b
$ MemBind -> Maybe MemBind
forall a. a -> Maybe a
Just MemBind
mem_ixf
                        LetAttrMem
_ -> Maybe MemBind -> m (Maybe MemBind)
forall (m :: * -> *) a. Monad m => a -> m a
return Maybe MemBind
forall a. Maybe a
Nothing
allocInExp Exp fromlore
e = Mapper fromlore tolore (AllocM fromlore tolore)
-> Exp fromlore -> AllocM fromlore tolore (Exp tolore)
forall (m :: * -> *) flore tlore.
(Applicative m, Monad m) =>
Mapper flore tlore m -> Exp flore -> m (Exp tlore)
mapExpM Mapper fromlore tolore (AllocM fromlore tolore)
alloc Exp fromlore
e
  where alloc :: Mapper fromlore tolore (AllocM fromlore tolore)
alloc =
          Mapper Any Any (AllocM fromlore tolore)
forall (m :: * -> *) lore. Monad m => Mapper lore lore m
identityMapper { mapOnBody :: Scope tolore
-> BodyT fromlore -> AllocM fromlore tolore (BodyT tolore)
mapOnBody = String
-> Scope tolore
-> BodyT fromlore
-> AllocM fromlore tolore (BodyT tolore)
forall a. HasCallStack => String -> a
error String
"Unhandled Body in ExplicitAllocations"
                         , mapOnRetType :: RetType fromlore -> AllocM fromlore tolore (RetType tolore)
mapOnRetType = String
-> RetType fromlore -> AllocM fromlore tolore (RetType tolore)
forall a. HasCallStack => String -> a
error String
"Unhandled RetType in ExplicitAllocations"
                         , mapOnBranchType :: BranchType fromlore -> AllocM fromlore tolore (BranchType tolore)
mapOnBranchType = String
-> BranchType fromlore
-> AllocM fromlore tolore (BranchType tolore)
forall a. HasCallStack => String -> a
error String
"Unhandled BranchType in ExplicitAllocations"
                         , mapOnFParam :: FParam fromlore -> AllocM fromlore tolore (FParam tolore)
mapOnFParam = String -> FParam fromlore -> AllocM fromlore tolore (FParam tolore)
forall a. HasCallStack => String -> a
error String
"Unhandled FParam in ExplicitAllocations"
                         , mapOnLParam :: LParam fromlore -> AllocM fromlore tolore (LParam tolore)
mapOnLParam = String -> LParam fromlore -> AllocM fromlore tolore (LParam tolore)
forall a. HasCallStack => String -> a
error String
"Unhandled LParam in ExplicitAllocations"
                         , mapOnOp :: Op fromlore -> AllocM fromlore tolore (Op tolore)
mapOnOp = \Op fromlore
op -> do Op fromlore -> AllocM fromlore tolore (Op tolore)
handle <- (AllocEnv fromlore tolore
 -> Op fromlore -> AllocM fromlore tolore (Op tolore))
-> AllocM
     fromlore tolore (Op fromlore -> AllocM fromlore tolore (Op tolore))
forall r (m :: * -> *) a. MonadReader r m => (r -> a) -> m a
asks AllocEnv fromlore tolore
-> Op fromlore -> AllocM fromlore tolore (Op tolore)
forall fromlore tolore.
AllocEnv fromlore tolore
-> Op fromlore -> AllocM fromlore tolore (Op tolore)
allocInOp
                                               Op fromlore -> AllocM fromlore tolore (Op tolore)
handle Op fromlore
op
                         }

addResCtxInIfBody :: (Allocable fromlore tolore, Allocator tolore (AllocM fromlore tolore)) =>
                     [ExtType] -> Body tolore -> [Maybe Space] ->
                     [Maybe (ExtIxFun, [PrimExp VName])] ->
                     AllocM fromlore tolore (Body tolore, [BodyReturns])
addResCtxInIfBody :: [TypeBase ExtShape NoUniqueness]
-> Body tolore
-> [Maybe Space]
-> [Maybe (ExtIxFun, [PrimExp VName])]
-> AllocM fromlore tolore (Body tolore, [BranchTypeMem])
addResCtxInIfBody [TypeBase ExtShape NoUniqueness]
ifrets (Body BodyAttr tolore
_ Stms tolore
bnds Result
res) [Maybe Space]
spaces [Maybe (ExtIxFun, [PrimExp VName])]
substs = do
  let num_vals :: Int
num_vals = [TypeBase ExtShape NoUniqueness] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length [TypeBase ExtShape NoUniqueness]
ifrets
      (Result
ctx_res, Result
val_res) = Int -> Result -> (Result, Result)
forall a. Int -> [a] -> ([a], [a])
splitFromEnd Int
num_vals Result
res
  ((Result
res', [BranchTypeMem]
bodyrets'), Stms tolore
all_body_stms) <- AllocM fromlore tolore (Result, [BranchTypeMem])
-> AllocM
     fromlore
     tolore
     ((Result, [BranchTypeMem]), Stms (Lore (AllocM fromlore tolore)))
forall (m :: * -> *) a.
MonadBinder m =>
m a -> m (a, Stms (Lore m))
collectStms (AllocM fromlore tolore (Result, [BranchTypeMem])
 -> AllocM
      fromlore
      tolore
      ((Result, [BranchTypeMem]), Stms (Lore (AllocM fromlore tolore))))
-> AllocM fromlore tolore (Result, [BranchTypeMem])
-> AllocM
     fromlore
     tolore
     ((Result, [BranchTypeMem]), Stms (Lore (AllocM fromlore tolore)))
forall a b. (a -> b) -> a -> b
$ do
    (Stm tolore -> AllocM fromlore tolore ())
-> Stms tolore -> AllocM fromlore tolore ()
forall (t :: * -> *) (m :: * -> *) a b.
(Foldable t, Monad m) =>
(a -> m b) -> t a -> m ()
mapM_ Stm tolore -> AllocM fromlore tolore ()
forall (m :: * -> *). MonadBinder m => Stm (Lore m) -> m ()
addStm Stms tolore
bnds
    (Result
val_res', Result
ext_ses_res, Result
mem_ctx_res, [BranchTypeMem]
bodyrets, Int
total_existentials) <-
      ((Result, Result, Result, [BranchTypeMem], Int)
 -> (TypeBase ExtShape NoUniqueness, SubExp,
     Maybe (ExtIxFun, [PrimExp VName]), Maybe Space)
 -> AllocM
      fromlore tolore (Result, Result, Result, [BranchTypeMem], Int))
-> (Result, Result, Result, [BranchTypeMem], Int)
-> [(TypeBase ExtShape NoUniqueness, SubExp,
     Maybe (ExtIxFun, [PrimExp VName]), Maybe Space)]
-> AllocM
     fromlore tolore (Result, Result, Result, [BranchTypeMem], Int)
forall (t :: * -> *) (m :: * -> *) b a.
(Foldable t, Monad m) =>
(b -> a -> m b) -> b -> t a -> m b
foldM (Result, Result, Result, [BranchTypeMem], Int)
-> (TypeBase ExtShape NoUniqueness, SubExp,
    Maybe (ExtIxFun, [PrimExp VName]), Maybe Space)
-> AllocM
     fromlore tolore (Result, Result, Result, [BranchTypeMem], Int)
forall tolore fromlore u.
(PrettyLore fromlore, AllocOp (Op tolore), Checkable tolore,
 OpReturns tolore, SizeSubst (Op tolore), BinderOps tolore,
 FParamAttr tolore ~ FParamMem, FParamAttr fromlore ~ DeclType,
 LParamAttr tolore ~ LetAttrMem, ExpAttr tolore ~ (),
 LParamAttr fromlore ~ Type, LetAttr tolore ~ LetAttrMem,
 BodyAttr tolore ~ (),
 BranchType fromlore ~ TypeBase ExtShape NoUniqueness,
 RetType tolore ~ RetTypeMem, BodyAttr fromlore ~ (),
 RetType fromlore ~ DeclExtType,
 BranchType tolore ~ BranchTypeMem) =>
(Result, Result, Result, [MemInfo (Ext SubExp) u MemReturn], Int)
-> (TypeBase ExtShape u, SubExp, Maybe (ExtIxFun, [PrimExp VName]),
    Maybe Space)
-> AllocM
     fromlore
     tolore
     (Result, Result, Result, [MemInfo (Ext SubExp) u MemReturn], Int)
helper ([], [], [], [], Result -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length Result
ctx_res) ([TypeBase ExtShape NoUniqueness]
-> Result
-> [Maybe (ExtIxFun, [PrimExp VName])]
-> [Maybe Space]
-> [(TypeBase ExtShape NoUniqueness, SubExp,
     Maybe (ExtIxFun, [PrimExp VName]), Maybe Space)]
forall a b c d. [a] -> [b] -> [c] -> [d] -> [(a, b, c, d)]
zip4 [TypeBase ExtShape NoUniqueness]
ifrets Result
val_res [Maybe (ExtIxFun, [PrimExp VName])]
substs [Maybe Space]
spaces)
    (Result, [BranchTypeMem])
-> AllocM fromlore tolore (Result, [BranchTypeMem])
forall (m :: * -> *) a. Monad m => a -> m a
return (Result
ctx_res Result -> Result -> Result
forall a. Semigroup a => a -> a -> a
<> Result
ext_ses_res Result -> Result -> Result
forall a. Semigroup a => a -> a -> a
<> Result
mem_ctx_res Result -> Result -> Result
forall a. Semigroup a => a -> a -> a
<> Result
val_res',
             -- We need to adjust the ReturnsNewBlock existentials, because they
             -- should always be numbered _after_ all other existentials in the
             -- return values.
            [BranchTypeMem] -> [BranchTypeMem]
forall a. [a] -> [a]
reverse ([BranchTypeMem] -> [BranchTypeMem])
-> [BranchTypeMem] -> [BranchTypeMem]
forall a b. (a -> b) -> a -> b
$ ([BranchTypeMem], Int) -> [BranchTypeMem]
forall a b. (a, b) -> a
fst (([BranchTypeMem], Int) -> [BranchTypeMem])
-> ([BranchTypeMem], Int) -> [BranchTypeMem]
forall a b. (a -> b) -> a -> b
$ (([BranchTypeMem], Int) -> BranchTypeMem -> ([BranchTypeMem], Int))
-> ([BranchTypeMem], Int)
-> [BranchTypeMem]
-> ([BranchTypeMem], Int)
forall (t :: * -> *) b a.
Foldable t =>
(b -> a -> b) -> b -> t a -> b
foldl ([BranchTypeMem], Int) -> BranchTypeMem -> ([BranchTypeMem], Int)
adjustNewBlockExistential ([], Int
total_existentials) [BranchTypeMem]
bodyrets)
  Body tolore
body' <- Stms (Lore (AllocM fromlore tolore))
-> Result
-> AllocM fromlore tolore (Body (Lore (AllocM fromlore tolore)))
forall (m :: * -> *).
MonadBinder m =>
Stms (Lore m) -> Result -> m (Body (Lore m))
mkBodyM Stms tolore
Stms (Lore (AllocM fromlore tolore))
all_body_stms Result
res'
  (Body tolore, [BranchTypeMem])
-> AllocM fromlore tolore (Body tolore, [BranchTypeMem])
forall (m :: * -> *) a. Monad m => a -> m a
return (Body tolore
body', [BranchTypeMem]
bodyrets')
    where
      helper :: (Result, Result, Result, [MemInfo (Ext SubExp) u MemReturn], Int)
-> (TypeBase ExtShape u, SubExp, Maybe (ExtIxFun, [PrimExp VName]),
    Maybe Space)
-> AllocM
     fromlore
     tolore
     (Result, Result, Result, [MemInfo (Ext SubExp) u MemReturn], Int)
helper (Result
res_acc, Result
ext_acc, Result
ctx_acc, [MemInfo (Ext SubExp) u MemReturn]
br_acc, Int
k) (TypeBase ExtShape u
ifr, SubExp
r, Maybe (ExtIxFun, [PrimExp VName])
mbixfsub, Maybe Space
sp) =
        case Maybe (ExtIxFun, [PrimExp VName])
mbixfsub of
          Maybe (ExtIxFun, [PrimExp VName])
Nothing -> do
            -- does NOT generalize/antiunify; ensure direct
            SubExp
r' <- Maybe Space -> SubExp -> AllocM fromlore tolore SubExp
forall fromlore tolore.
(Allocable fromlore tolore,
 Allocator tolore (AllocM fromlore tolore)) =>
Maybe Space -> SubExp -> AllocM fromlore tolore SubExp
ensureDirect Maybe Space
sp SubExp
r
            Result
mem_ctx_r <- SubExp -> AllocM fromlore tolore Result
forall fromlore tolore.
(Allocable fromlore tolore,
 Allocator tolore (AllocM fromlore tolore)) =>
SubExp -> AllocM fromlore tolore Result
bodyReturnMemCtx SubExp
r'
            let body_ret :: MemInfo (Ext SubExp) u MemReturn
body_ret = TypeBase ExtShape u
-> Maybe Space -> MemInfo (Ext SubExp) u MemReturn
forall u.
TypeBase ExtShape u
-> Maybe Space -> MemInfo (Ext SubExp) u MemReturn
inspect TypeBase ExtShape u
ifr Maybe Space
sp
            (Result, Result, Result, [MemInfo (Ext SubExp) u MemReturn], Int)
-> AllocM
     fromlore
     tolore
     (Result, Result, Result, [MemInfo (Ext SubExp) u MemReturn], Int)
forall (m :: * -> *) a. Monad m => a -> m a
return (Result
res_acc Result -> Result -> Result
forall a. [a] -> [a] -> [a]
++ [SubExp
r'],
                    Result
ext_acc,
                    Result
ctx_acc Result -> Result -> Result
forall a. [a] -> [a] -> [a]
++ Result
mem_ctx_r,
                    [MemInfo (Ext SubExp) u MemReturn]
br_acc [MemInfo (Ext SubExp) u MemReturn]
-> [MemInfo (Ext SubExp) u MemReturn]
-> [MemInfo (Ext SubExp) u MemReturn]
forall a. [a] -> [a] -> [a]
++ [MemInfo (Ext SubExp) u MemReturn
body_ret],
                    Int
k)
          Just (ExtIxFun
ixfn, [PrimExp VName]
m) -> do -- generalizes
            let i :: Int
i = [PrimExp VName] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length [PrimExp VName]
m
            Result
ext_ses <- (PrimExp VName -> AllocM fromlore tolore SubExp)
-> [PrimExp VName] -> AllocM fromlore tolore Result
forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
(a -> m b) -> t a -> m (t b)
mapM (String
-> (VName
    -> AllocM fromlore tolore (Exp (Lore (AllocM fromlore tolore))))
-> PrimExp VName
-> AllocM fromlore tolore SubExp
forall (m :: * -> *) v.
MonadBinder m =>
String -> (v -> m (Exp (Lore m))) -> PrimExp v -> m SubExp
primExpToSubExp String
"ixfn_exist"
                             (ExpT tolore -> AllocM fromlore tolore (ExpT tolore)
forall (m :: * -> *) a. Monad m => a -> m a
return (ExpT tolore -> AllocM fromlore tolore (ExpT tolore))
-> (VName -> ExpT tolore)
-> VName
-> AllocM fromlore tolore (ExpT tolore)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. BasicOp -> ExpT tolore
forall lore. BasicOp -> ExpT lore
BasicOp (BasicOp -> ExpT tolore)
-> (VName -> BasicOp) -> VName -> ExpT tolore
forall b c a. (b -> c) -> (a -> b) -> a -> c
. SubExp -> BasicOp
SubExp (SubExp -> BasicOp) -> (VName -> SubExp) -> VName -> BasicOp
forall b c a. (b -> c) -> (a -> b) -> a -> c
. VName -> SubExp
Var))
                       [PrimExp VName]
m
            Result
mem_ctx_r <- SubExp -> AllocM fromlore tolore Result
forall fromlore tolore.
(Allocable fromlore tolore,
 Allocator tolore (AllocM fromlore tolore)) =>
SubExp -> AllocM fromlore tolore Result
bodyReturnMemCtx SubExp
r
            let sp' :: Space
sp' = Space -> Maybe Space -> Space
forall a. a -> Maybe a -> a
fromMaybe Space
DefaultSpace Maybe Space
sp
                ixfn' :: ExtIxFun
ixfn' = (PrimExp (Ext VName) -> PrimExp (Ext VName))
-> ExtIxFun -> ExtIxFun
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap (Int -> PrimExp (Ext VName) -> PrimExp (Ext VName)
adjustExtPE Int
k) ExtIxFun
ixfn
                exttp :: MemInfo (Ext SubExp) u MemReturn
exttp = case TypeBase ExtShape u
ifr of
                          Array PrimType
pt ExtShape
shp' u
u ->
                            PrimType
-> ExtShape -> u -> MemReturn -> MemInfo (Ext SubExp) u MemReturn
forall d u ret.
PrimType -> ShapeBase d -> u -> ret -> MemInfo d u ret
MemArray PrimType
pt ExtShape
shp' u
u (MemReturn -> MemInfo (Ext SubExp) u MemReturn)
-> MemReturn -> MemInfo (Ext SubExp) u MemReturn
forall a b. (a -> b) -> a -> b
$
                            Space -> Int -> ExtIxFun -> MemReturn
ReturnsNewBlock Space
sp' Int
0 ExtIxFun
ixfn'
                          TypeBase ExtShape u
_ -> String -> MemInfo (Ext SubExp) u MemReturn
forall a. HasCallStack => String -> a
error String
"Impossible case reached in addResCtxInIfBody"
            (Result, Result, Result, [MemInfo (Ext SubExp) u MemReturn], Int)
-> AllocM
     fromlore
     tolore
     (Result, Result, Result, [MemInfo (Ext SubExp) u MemReturn], Int)
forall (m :: * -> *) a. Monad m => a -> m a
return (Result
res_acc Result -> Result -> Result
forall a. [a] -> [a] -> [a]
++ [SubExp
r],
                    Result
ext_acc Result -> Result -> Result
forall a. [a] -> [a] -> [a]
++ Result
ext_ses,
                    Result
ctx_acc Result -> Result -> Result
forall a. [a] -> [a] -> [a]
++ Result
mem_ctx_r,
                    [MemInfo (Ext SubExp) u MemReturn]
br_acc [MemInfo (Ext SubExp) u MemReturn]
-> [MemInfo (Ext SubExp) u MemReturn]
-> [MemInfo (Ext SubExp) u MemReturn]
forall a. [a] -> [a] -> [a]
++ [MemInfo (Ext SubExp) u MemReturn
exttp],
                    Int
k Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
i)

      adjustNewBlockExistential :: ([BodyReturns], Int) -> BodyReturns -> ([BodyReturns], Int)
      adjustNewBlockExistential :: ([BranchTypeMem], Int) -> BranchTypeMem -> ([BranchTypeMem], Int)
adjustNewBlockExistential ([BranchTypeMem]
acc, Int
k) (MemArray PrimType
pt ExtShape
shp NoUniqueness
u (ReturnsNewBlock Space
space Int
_ ExtIxFun
ixfun)) =
        (PrimType -> ExtShape -> NoUniqueness -> MemReturn -> BranchTypeMem
forall d u ret.
PrimType -> ShapeBase d -> u -> ret -> MemInfo d u ret
MemArray PrimType
pt ExtShape
shp NoUniqueness
u (Space -> Int -> ExtIxFun -> MemReturn
ReturnsNewBlock Space
space Int
k ExtIxFun
ixfun) BranchTypeMem -> [BranchTypeMem] -> [BranchTypeMem]
forall a. a -> [a] -> [a]
: [BranchTypeMem]
acc, Int
k Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
1)
      adjustNewBlockExistential ([BranchTypeMem]
acc, Int
k) BranchTypeMem
x = (BranchTypeMem
x BranchTypeMem -> [BranchTypeMem] -> [BranchTypeMem]
forall a. a -> [a] -> [a]
: [BranchTypeMem]
acc, Int
k)

      inspect :: TypeBase ExtShape u
-> Maybe Space -> MemInfo (Ext SubExp) u MemReturn
inspect (Array PrimType
pt ExtShape
shape u
u) Maybe Space
space =
        let space' :: Space
space' = Space -> Maybe Space -> Space
forall a. a -> Maybe a -> a
fromMaybe Space
DefaultSpace Maybe Space
space
            bodyret :: MemInfo (Ext SubExp) u MemReturn
bodyret = PrimType
-> ExtShape -> u -> MemReturn -> MemInfo (Ext SubExp) u MemReturn
forall d u ret.
PrimType -> ShapeBase d -> u -> ret -> MemInfo d u ret
MemArray PrimType
pt ExtShape
shape u
u (MemReturn -> MemInfo (Ext SubExp) u MemReturn)
-> MemReturn -> MemInfo (Ext SubExp) u MemReturn
forall a b. (a -> b) -> a -> b
$ Space -> Int -> ExtIxFun -> MemReturn
ReturnsNewBlock Space
space' Int
0 (ExtIxFun -> MemReturn) -> ExtIxFun -> MemReturn
forall a b. (a -> b) -> a -> b
$
              Shape (PrimExp (Ext VName)) -> ExtIxFun
forall num. IntegralExp num => Shape num -> IxFun num
IxFun.iota (Shape (PrimExp (Ext VName)) -> ExtIxFun)
-> Shape (PrimExp (Ext VName)) -> ExtIxFun
forall a b. (a -> b) -> a -> b
$ (Ext SubExp -> PrimExp (Ext VName))
-> [Ext SubExp] -> Shape (PrimExp (Ext VName))
forall a b. (a -> b) -> [a] -> [b]
map Ext SubExp -> PrimExp (Ext VName)
convert ([Ext SubExp] -> Shape (PrimExp (Ext VName)))
-> [Ext SubExp] -> Shape (PrimExp (Ext VName))
forall a b. (a -> b) -> a -> b
$ ExtShape -> [Ext SubExp]
forall d. ShapeBase d -> [d]
shapeDims ExtShape
shape
        in MemInfo (Ext SubExp) u MemReturn
bodyret
      inspect (Prim PrimType
pt) Maybe Space
_ = PrimType -> MemInfo (Ext SubExp) u MemReturn
forall d u ret. PrimType -> MemInfo d u ret
MemPrim PrimType
pt
      inspect (Mem Space
space) Maybe Space
_ = Space -> MemInfo (Ext SubExp) u MemReturn
forall d u ret. Space -> MemInfo d u ret
MemMem Space
space

      convert :: Ext SubExp -> PrimExp (Ext VName)
convert (Ext Int
i) = Ext VName -> PrimType -> PrimExp (Ext VName)
forall v. v -> PrimType -> PrimExp v
LeafExp (Int -> Ext VName
forall a. Int -> Ext a
Ext Int
i) PrimType
int32
      convert (Free SubExp
v) = VName -> Ext VName
forall a. a -> Ext a
Free (VName -> Ext VName) -> PrimExp VName -> PrimExp (Ext VName)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> PrimType -> SubExp -> PrimExp VName
primExpFromSubExp PrimType
int32 SubExp
v

      adjustExtV :: Int -> Ext VName -> Ext VName
      adjustExtV :: Int -> Ext VName -> Ext VName
adjustExtV Int
_ (Free VName
v) = VName -> Ext VName
forall a. a -> Ext a
Free VName
v
      adjustExtV Int
k (Ext Int
i) = Int -> Ext VName
forall a. Int -> Ext a
Ext (Int
k Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
i)

      adjustExtPE :: Int -> PrimExp (Ext VName) -> PrimExp (Ext VName)
      adjustExtPE :: Int -> PrimExp (Ext VName) -> PrimExp (Ext VName)
adjustExtPE Int
k = (Ext VName -> Ext VName)
-> PrimExp (Ext VName) -> PrimExp (Ext VName)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap (Int -> Ext VName -> Ext VName
adjustExtV Int
k)

mkSpaceOks :: (Mem tolore, LocalScope tolore m) =>
              Int -> Body tolore -> m [Maybe Space]
mkSpaceOks :: Int -> Body tolore -> m [Maybe Space]
mkSpaceOks Int
num_vals (Body BodyAttr tolore
_ Stms tolore
stms Result
res) =
  Stms tolore -> m [Maybe Space] -> m [Maybe Space]
forall lore a (m :: * -> *) b.
(Scoped lore a, LocalScope lore m) =>
a -> m b -> m b
inScopeOf Stms tolore
stms (m [Maybe Space] -> m [Maybe Space])
-> m [Maybe Space] -> m [Maybe Space]
forall a b. (a -> b) -> a -> b
$
  (SubExp -> m (Maybe Space)) -> Result -> m [Maybe Space]
forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
(a -> m b) -> t a -> m (t b)
mapM SubExp -> m (Maybe Space)
forall lore (m :: * -> *).
(HasScope lore m, AllocOp (Op lore), Monad m, Checkable lore,
 OpReturns lore, FParamAttr lore ~ FParamMem,
 LetAttr lore ~ LetAttrMem, LParamAttr lore ~ LetAttrMem,
 BranchType lore ~ BranchTypeMem, RetType lore ~ RetTypeMem) =>
SubExp -> m (Maybe Space)
mkSpaceOK (Result -> m [Maybe Space]) -> Result -> m [Maybe Space]
forall a b. (a -> b) -> a -> b
$ Int -> Result -> Result
forall a. Int -> [a] -> [a]
takeLast Int
num_vals Result
res
  where mkSpaceOK :: SubExp -> m (Maybe Space)
mkSpaceOK (Var VName
v) = do
          LetAttrMem
v_info <- VName -> m LetAttrMem
forall lore (m :: * -> *).
(HasScope lore m, Mem lore) =>
VName -> m LetAttrMem
lookupMemInfo VName
v
          case LetAttrMem
v_info of MemArray PrimType
_ Shape
_ NoUniqueness
_ (ArrayIn VName
mem IxFun
_) -> do
                           LetAttrMem
mem_info <- VName -> m LetAttrMem
forall lore (m :: * -> *).
(HasScope lore m, Mem lore) =>
VName -> m LetAttrMem
lookupMemInfo VName
mem
                           case LetAttrMem
mem_info of MemMem Space
space -> Maybe Space -> m (Maybe Space)
forall (m :: * -> *) a. Monad m => a -> m a
return (Maybe Space -> m (Maybe Space)) -> Maybe Space -> m (Maybe Space)
forall a b. (a -> b) -> a -> b
$ Space -> Maybe Space
forall a. a -> Maybe a
Just Space
space
                                            LetAttrMem
_ -> Maybe Space -> m (Maybe Space)
forall (m :: * -> *) a. Monad m => a -> m a
return Maybe Space
forall a. Maybe a
Nothing
                         LetAttrMem
_ -> Maybe Space -> m (Maybe Space)
forall (m :: * -> *) a. Monad m => a -> m a
return Maybe Space
forall a. Maybe a
Nothing
        mkSpaceOK SubExp
_ = Maybe Space -> m (Maybe Space)
forall (m :: * -> *) a. Monad m => a -> m a
return Maybe Space
forall a. Maybe a
Nothing

allocInLoopForm :: (Allocable fromlore tolore,
                    Allocator tolore (AllocM fromlore tolore)) =>
                   LoopForm fromlore -> AllocM fromlore tolore (LoopForm tolore)
allocInLoopForm :: LoopForm fromlore -> AllocM fromlore tolore (LoopForm tolore)
allocInLoopForm (WhileLoop VName
v) = LoopForm tolore -> AllocM fromlore tolore (LoopForm tolore)
forall (m :: * -> *) a. Monad m => a -> m a
return (LoopForm tolore -> AllocM fromlore tolore (LoopForm tolore))
-> LoopForm tolore -> AllocM fromlore tolore (LoopForm tolore)
forall a b. (a -> b) -> a -> b
$ VName -> LoopForm tolore
forall lore. VName -> LoopForm lore
WhileLoop VName
v
allocInLoopForm (ForLoop VName
i IntType
it SubExp
n [(LParam fromlore, VName)]
loopvars) =
  VName
-> IntType -> SubExp -> [(LParam tolore, VName)] -> LoopForm tolore
forall lore.
VName
-> IntType -> SubExp -> [(LParam lore, VName)] -> LoopForm lore
ForLoop VName
i IntType
it SubExp
n ([(Param LetAttrMem, VName)] -> LoopForm tolore)
-> AllocM fromlore tolore [(Param LetAttrMem, VName)]
-> AllocM fromlore tolore (LoopForm tolore)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> ((Param Type, VName)
 -> AllocM fromlore tolore (Param LetAttrMem, VName))
-> [(Param Type, VName)]
-> AllocM fromlore tolore [(Param LetAttrMem, VName)]
forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
(a -> m b) -> t a -> m (t b)
mapM (Param Type, VName)
-> AllocM fromlore tolore (Param LetAttrMem, VName)
allocInLoopVar [(Param Type, VName)]
[(LParam fromlore, VName)]
loopvars
  where allocInLoopVar :: (Param Type, VName)
-> AllocM fromlore tolore (Param LetAttrMem, VName)
allocInLoopVar (Param Type
p,VName
a) = do
          (VName
mem, IxFun
ixfun) <- VName -> AllocM fromlore tolore (VName, IxFun)
forall lore (m :: * -> *).
(Mem lore, HasScope lore m, Monad m) =>
VName -> m (VName, IxFun)
lookupArraySummary VName
a
          case Param Type -> Type
forall attr. Typed attr => Param attr -> Type
paramType Param Type
p of
            Array PrimType
bt Shape
shape NoUniqueness
u -> do
              [PrimExp VName]
dims <- (SubExp -> PrimExp VName) -> Result -> [PrimExp VName]
forall a b. (a -> b) -> [a] -> [b]
map (PrimType -> SubExp -> PrimExp VName
primExpFromSubExp PrimType
int32) (Result -> [PrimExp VName])
-> (Type -> Result) -> Type -> [PrimExp VName]
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Type -> Result
forall u. TypeBase Shape u -> Result
arrayDims (Type -> [PrimExp VName])
-> AllocM fromlore tolore Type
-> AllocM fromlore tolore [PrimExp VName]
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> VName -> AllocM fromlore tolore Type
forall lore (m :: * -> *). HasScope lore m => VName -> m Type
lookupType VName
a
              let ixfun' :: IxFun
ixfun' = IxFun -> Slice (PrimExp VName) -> IxFun
forall num.
(Eq num, IntegralExp num) =>
IxFun num -> Slice num -> IxFun num
IxFun.slice IxFun
ixfun (Slice (PrimExp VName) -> IxFun) -> Slice (PrimExp VName) -> IxFun
forall a b. (a -> b) -> a -> b
$
                           [PrimExp VName] -> Slice (PrimExp VName) -> Slice (PrimExp VName)
forall d. Num d => [d] -> [DimIndex d] -> [DimIndex d]
fullSliceNum [PrimExp VName]
dims [PrimExp VName -> DimIndex (PrimExp VName)
forall d. d -> DimIndex d
DimFix (PrimExp VName -> DimIndex (PrimExp VName))
-> PrimExp VName -> DimIndex (PrimExp VName)
forall a b. (a -> b) -> a -> b
$ VName -> PrimType -> PrimExp VName
forall v. v -> PrimType -> PrimExp v
LeafExp VName
i PrimType
int32]
              (Param LetAttrMem, VName)
-> AllocM fromlore tolore (Param LetAttrMem, VName)
forall (m :: * -> *) a. Monad m => a -> m a
return (Param Type
p { paramAttr :: LetAttrMem
paramAttr = PrimType -> Shape -> NoUniqueness -> MemBind -> LetAttrMem
forall d u ret.
PrimType -> ShapeBase d -> u -> ret -> MemInfo d u ret
MemArray PrimType
bt Shape
shape NoUniqueness
u (MemBind -> LetAttrMem) -> MemBind -> LetAttrMem
forall a b. (a -> b) -> a -> b
$ VName -> IxFun -> MemBind
ArrayIn VName
mem IxFun
ixfun' }, VName
a)
            Prim PrimType
bt ->
              (Param LetAttrMem, VName)
-> AllocM fromlore tolore (Param LetAttrMem, VName)
forall (m :: * -> *) a. Monad m => a -> m a
return (Param Type
p { paramAttr :: LetAttrMem
paramAttr = PrimType -> LetAttrMem
forall d u ret. PrimType -> MemInfo d u ret
MemPrim PrimType
bt }, VName
a)
            Mem Space
space ->
              (Param LetAttrMem, VName)
-> AllocM fromlore tolore (Param LetAttrMem, VName)
forall (m :: * -> *) a. Monad m => a -> m a
return (Param Type
p { paramAttr :: LetAttrMem
paramAttr = Space -> LetAttrMem
forall d u ret. Space -> MemInfo d u ret
MemMem Space
space }, VName
a)

class SizeSubst op where
  opSizeSubst :: PatternT attr -> op -> ChunkMap
  opIsConst :: op -> Bool
  opIsConst = Bool -> op -> Bool
forall a b. a -> b -> a
const Bool
False

instance SizeSubst () where
  opSizeSubst :: PatternT attr -> () -> Map VName SubExp
opSizeSubst PatternT attr
_ ()
_ = Map VName SubExp
forall a. Monoid a => a
mempty

instance SizeSubst op => SizeSubst (MemOp op) where
  opSizeSubst :: PatternT attr -> MemOp op -> Map VName SubExp
opSizeSubst PatternT attr
pat (Inner op
op) = PatternT attr -> op -> Map VName SubExp
forall op attr.
SizeSubst op =>
PatternT attr -> op -> Map VName SubExp
opSizeSubst PatternT attr
pat op
op
  opSizeSubst PatternT attr
_ MemOp op
_ = Map VName SubExp
forall a. Monoid a => a
mempty

  opIsConst :: MemOp op -> Bool
opIsConst (Inner op
op) = op -> Bool
forall op. SizeSubst op => op -> Bool
opIsConst op
op
  opIsConst MemOp op
_ = Bool
False

sizeSubst :: SizeSubst (Op lore) => Stm lore -> ChunkMap
sizeSubst :: Stm lore -> Map VName SubExp
sizeSubst (Let Pattern lore
pat StmAux (ExpAttr lore)
_ (Op Op lore
op)) = Pattern lore -> Op lore -> Map VName SubExp
forall op attr.
SizeSubst op =>
PatternT attr -> op -> Map VName SubExp
opSizeSubst Pattern lore
pat Op lore
op
sizeSubst Stm lore
_ = Map VName SubExp
forall a. Monoid a => a
mempty

stmConsts :: SizeSubst (Op lore) => Stm lore -> S.Set VName
stmConsts :: Stm lore -> Set VName
stmConsts (Let Pattern lore
pat StmAux (ExpAttr lore)
_ (Op Op lore
op))
  | Op lore -> Bool
forall op. SizeSubst op => op -> Bool
opIsConst Op lore
op = [VName] -> Set VName
forall a. Ord a => [a] -> Set a
S.fromList ([VName] -> Set VName) -> [VName] -> Set VName
forall a b. (a -> b) -> a -> b
$ Pattern lore -> [VName]
forall attr. PatternT attr -> [VName]
patternNames Pattern lore
pat
stmConsts Stm lore
_ = Set VName
forall a. Monoid a => a
mempty

mkLetNamesB' :: (Op (Lore m) ~ MemOp inner,
                 MonadBinder m, ExpAttr (Lore m) ~ (),
                 Allocator (Lore m) (PatAllocM (Lore m))) =>
                ExpAttr (Lore m) -> [VName] -> Exp (Lore m) -> m (Stm (Lore m))
mkLetNamesB' :: ExpAttr (Lore m) -> [VName] -> Exp (Lore m) -> m (Stm (Lore m))
mkLetNamesB' ExpAttr (Lore m)
attr [VName]
names Exp (Lore m)
e = do
  Scope (Lore m)
scope <- m (Scope (Lore m))
forall lore (m :: * -> *). HasScope lore m => m (Scope lore)
askScope
  PatternT (LetAttr (Lore m))
pat <- Scope (Lore m)
-> [VName] -> Exp (Lore m) -> m (PatternT (LetAttr (Lore m)))
forall (m :: * -> *) lore inner.
(MonadBinder m, ExpAttr lore ~ (), Op (Lore m) ~ MemOp inner,
 Allocator lore (PatAllocM lore)) =>
Scope lore -> [VName] -> Exp lore -> m (Pattern lore)
bindPatternWithAllocations Scope (Lore m)
scope [VName]
names Exp (Lore m)
e
  Stm (Lore m) -> m (Stm (Lore m))
forall (m :: * -> *) a. Monad m => a -> m a
return (Stm (Lore m) -> m (Stm (Lore m)))
-> Stm (Lore m) -> m (Stm (Lore m))
forall a b. (a -> b) -> a -> b
$ PatternT (LetAttr (Lore m))
-> StmAux (ExpAttr (Lore m)) -> Exp (Lore m) -> Stm (Lore m)
forall lore.
Pattern lore -> StmAux (ExpAttr lore) -> Exp lore -> Stm lore
Let PatternT (LetAttr (Lore m))
pat (() -> StmAux ()
forall attr. attr -> StmAux attr
defAux ()
ExpAttr (Lore m)
attr) Exp (Lore m)
e

mkLetNamesB'' :: (Op (Lore m) ~ MemOp inner, ExpAttr lore ~ (),
                   HasScope (Engine.Wise lore) m, Allocator lore (PatAllocM lore),
                   MonadBinder m, Engine.CanBeWise (Op lore)) =>
                 [VName] -> Exp (Engine.Wise lore)
              -> m (Stm (Engine.Wise lore))
mkLetNamesB'' :: [VName] -> Exp (Wise lore) -> m (Stm (Wise lore))
mkLetNamesB'' [VName]
names Exp (Wise lore)
e = do
  Scope lore
scope <- Scope (Wise lore) -> Scope lore
forall lore. Scope (Wise lore) -> Scope lore
Engine.removeScopeWisdom (Scope (Wise lore) -> Scope lore)
-> m (Scope (Wise lore)) -> m (Scope lore)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> m (Scope (Wise lore))
forall lore (m :: * -> *). HasScope lore m => m (Scope lore)
askScope
  (PatternT LetAttrMem
pat, [AllocStm]
prestms) <- PatAllocM lore (PatternT LetAttrMem)
-> Scope lore -> m (PatternT LetAttrMem, [AllocStm])
forall (m :: * -> *) lore a.
MonadFreshNames m =>
PatAllocM lore a -> Scope lore -> m (a, [AllocStm])
runPatAllocM ([VName] -> Exp lore -> PatAllocM lore (Pattern lore)
forall lore (m :: * -> *).
(Allocator lore m, ExpAttr lore ~ ()) =>
[VName] -> Exp lore -> m (Pattern lore)
patternWithAllocations [VName]
names (Exp lore -> PatAllocM lore (Pattern lore))
-> Exp lore -> PatAllocM lore (Pattern lore)
forall a b. (a -> b) -> a -> b
$ Exp (Wise lore) -> Exp lore
forall lore. CanBeWise (Op lore) => Exp (Wise lore) -> Exp lore
Engine.removeExpWisdom Exp (Wise lore)
e) Scope lore
scope
  (AllocStm -> m ()) -> [AllocStm] -> m ()
forall (t :: * -> *) (m :: * -> *) a b.
(Foldable t, Monad m) =>
(a -> m b) -> t a -> m ()
mapM_ AllocStm -> m ()
forall (m :: * -> *) inner.
(MonadBinder m, Op (Lore m) ~ MemOp inner) =>
AllocStm -> m ()
bindAllocStm [AllocStm]
prestms
  let pat' :: Pattern (Wise lore)
pat' = Pattern lore -> Exp (Wise lore) -> Pattern (Wise lore)
forall lore.
(Attributes lore, CanBeWise (Op lore)) =>
Pattern lore -> Exp (Wise lore) -> Pattern (Wise lore)
Engine.addWisdomToPattern Pattern lore
PatternT LetAttrMem
pat Exp (Wise lore)
e
      attr :: ExpAttr (Wise lore)
attr = Pattern (Wise lore)
-> ExpAttr lore -> Exp (Wise lore) -> ExpAttr (Wise lore)
forall lore.
(Attributes lore, CanBeWise (Op lore)) =>
Pattern (Wise lore)
-> ExpAttr lore -> Exp (Wise lore) -> ExpAttr (Wise lore)
Engine.mkWiseExpAttr Pattern (Wise lore)
pat' () Exp (Wise lore)
e
  Stm (Wise lore) -> m (Stm (Wise lore))
forall (m :: * -> *) a. Monad m => a -> m a
return (Stm (Wise lore) -> m (Stm (Wise lore)))
-> Stm (Wise lore) -> m (Stm (Wise lore))
forall a b. (a -> b) -> a -> b
$ Pattern (Wise lore)
-> StmAux (ExpAttr (Wise lore))
-> Exp (Wise lore)
-> Stm (Wise lore)
forall lore.
Pattern lore -> StmAux (ExpAttr lore) -> Exp lore -> Stm lore
Let Pattern (Wise lore)
pat' ((ExpWisdom, ()) -> StmAux (ExpWisdom, ())
forall attr. attr -> StmAux attr
defAux (ExpWisdom, ())
ExpAttr (Wise lore)
attr) Exp (Wise lore)
e

simplifiable :: (Engine.SimplifiableLore lore,
                 ExpAttr lore ~ (),
                 BodyAttr lore ~ (),
                 Op lore ~ MemOp inner,
                 Allocator lore (PatAllocM lore)) =>
                (inner -> Engine.SimpleM lore (Engine.OpWithWisdom inner, Stms (Engine.Wise lore)))
             -> SimpleOps lore
simplifiable :: (inner -> SimpleM lore (OpWithWisdom inner, Stms (Wise lore)))
-> SimpleOps lore
simplifiable inner -> SimpleM lore (OpWithWisdom inner, Stms (Wise lore))
simplifyInnerOp =
  (SymbolTable (Wise lore)
 -> Pattern (Wise lore)
 -> Exp (Wise lore)
 -> SimpleM lore (ExpAttr (Wise lore)))
-> (SymbolTable (Wise lore)
    -> Stms (Wise lore) -> Result -> SimpleM lore (Body (Wise lore)))
-> (SymbolTable (Wise lore)
    -> [VName]
    -> Exp (Wise lore)
    -> SimpleM lore (Stm (Wise lore), Stms (Wise lore)))
-> Protect (Binder (Wise lore))
-> SimplifyOp lore (Op lore)
-> SimpleOps lore
forall lore.
(SymbolTable (Wise lore)
 -> Pattern (Wise lore)
 -> Exp (Wise lore)
 -> SimpleM lore (ExpAttr (Wise lore)))
-> (SymbolTable (Wise lore)
    -> Stms (Wise lore) -> Result -> SimpleM lore (Body (Wise lore)))
-> (SymbolTable (Wise lore)
    -> [VName]
    -> Exp (Wise lore)
    -> SimpleM lore (Stm (Wise lore), Stms (Wise lore)))
-> Protect (Binder (Wise lore))
-> SimplifyOp lore (Op lore)
-> SimpleOps lore
SimpleOps SymbolTable (Wise lore)
-> Pattern (Wise lore)
-> Exp (Wise lore)
-> SimpleM lore (ExpAttr (Wise lore))
forall (m :: * -> *) lore p.
(Monad m, Attributes lore, CanBeWise (Op lore),
 ExpAttr lore ~ ()) =>
p
-> PatternT (VarWisdom, LetAttr lore)
-> Exp (Wise lore)
-> m (ExpWisdom, ExpAttr lore)
mkExpAttrS' SymbolTable (Wise lore)
-> Stms (Wise lore) -> Result -> SimpleM lore (Body (Wise lore))
forall (m :: * -> *) lore p.
(Monad m, Attributes lore, CanBeWise (Op lore),
 BodyAttr lore ~ ()) =>
p -> Stms (Wise lore) -> Result -> m (Body (Wise lore))
mkBodyS' SymbolTable (Wise lore)
-> [VName]
-> Exp (Wise lore)
-> SimpleM lore (Stm (Wise lore), Stms (Wise lore))
forall lore (m :: * -> *) somelore lore inner.
(BinderOps lore, MonadFreshNames m, HasScope somelore m,
 AllocOp (Op lore), Checkable lore, OpReturns lore,
 CanBeWise (Op lore), FParamAttr lore ~ FParamMem,
 LetAttr somelore ~ LetAttr lore, RetType lore ~ RetTypeMem,
 Op lore ~ MemOp inner, LetAttr lore ~ LetAttrMem,
 FParamAttr somelore ~ FParamAttr lore,
 LParamAttr somelore ~ LParamAttr lore,
 LParamAttr lore ~ LetAttrMem, ExpAttr lore ~ (),
 BranchType lore ~ BranchTypeMem) =>
SymbolTable (Wise lore)
-> [VName] -> Exp (Wise lore) -> m (Stm (Wise lore), Stms lore)
mkLetNamesS' Protect (Binder (Wise lore))
forall (m :: * -> *) d u ret inner inner.
(MonadBinder m, BranchType (Lore m) ~ MemInfo d u ret,
 Op (Lore m) ~ MemOp inner) =>
SubExp
-> PatternT (LetAttr (Lore m)) -> MemOp inner -> Maybe (m ())
protectOp SimplifyOp lore (Op lore)
MemOp inner
-> SimpleM lore (MemOp (OpWithWisdom inner), Stms (Wise lore))
simplifyOp
  where mkExpAttrS' :: p
-> PatternT (VarWisdom, LetAttr lore)
-> Exp (Wise lore)
-> m (ExpWisdom, ExpAttr lore)
mkExpAttrS' p
_ PatternT (VarWisdom, LetAttr lore)
pat Exp (Wise lore)
e =
          (ExpWisdom, ExpAttr lore) -> m (ExpWisdom, ExpAttr lore)
forall (m :: * -> *) a. Monad m => a -> m a
return ((ExpWisdom, ExpAttr lore) -> m (ExpWisdom, ExpAttr lore))
-> (ExpWisdom, ExpAttr lore) -> m (ExpWisdom, ExpAttr lore)
forall a b. (a -> b) -> a -> b
$ Pattern (Wise lore)
-> ExpAttr lore -> Exp (Wise lore) -> ExpAttr (Wise lore)
forall lore.
(Attributes lore, CanBeWise (Op lore)) =>
Pattern (Wise lore)
-> ExpAttr lore -> Exp (Wise lore) -> ExpAttr (Wise lore)
Engine.mkWiseExpAttr PatternT (VarWisdom, LetAttr lore)
Pattern (Wise lore)
pat () Exp (Wise lore)
e

        mkBodyS' :: p -> Stms (Wise lore) -> Result -> m (Body (Wise lore))
mkBodyS' p
_ Stms (Wise lore)
bnds Result
res = Body (Wise lore) -> m (Body (Wise lore))
forall (m :: * -> *) a. Monad m => a -> m a
return (Body (Wise lore) -> m (Body (Wise lore)))
-> Body (Wise lore) -> m (Body (Wise lore))
forall a b. (a -> b) -> a -> b
$ BodyAttr lore -> Stms (Wise lore) -> Result -> Body (Wise lore)
forall lore.
(Attributes lore, CanBeWise (Op lore)) =>
BodyAttr lore -> Stms (Wise lore) -> Result -> Body (Wise lore)
mkWiseBody () Stms (Wise lore)
bnds Result
res

        mkLetNamesS' :: SymbolTable (Wise lore)
-> [VName] -> Exp (Wise lore) -> m (Stm (Wise lore), Stms lore)
mkLetNamesS' SymbolTable (Wise lore)
vtable [VName]
names Exp (Wise lore)
e = do
          (PatternT (LetAttr lore)
pat', Stms lore
stms) <- Binder lore (PatternT (LetAttr lore))
-> m (PatternT (LetAttr lore), Stms lore)
forall (m :: * -> *) somelore lore a.
(MonadFreshNames m, HasScope somelore m,
 SameScope somelore lore) =>
Binder lore a -> m (a, Stms lore)
runBinder (Binder lore (PatternT (LetAttr lore))
 -> m (PatternT (LetAttr lore), Stms lore))
-> Binder lore (PatternT (LetAttr lore))
-> m (PatternT (LetAttr lore), Stms lore)
forall a b. (a -> b) -> a -> b
$ Scope lore
-> [VName] -> Exp lore -> Binder lore (PatternT (LetAttr lore))
forall (m :: * -> *) lore inner.
(MonadBinder m, ExpAttr lore ~ (), Op (Lore m) ~ MemOp inner,
 Allocator lore (PatAllocM lore)) =>
Scope lore -> [VName] -> Exp lore -> m (Pattern lore)
bindPatternWithAllocations Scope lore
env [VName]
names (Exp lore -> Binder lore (PatternT (LetAttr lore)))
-> Exp lore -> Binder lore (PatternT (LetAttr lore))
forall a b. (a -> b) -> a -> b
$
                          Exp (Wise lore) -> Exp lore
forall lore. CanBeWise (Op lore) => Exp (Wise lore) -> Exp lore
removeExpWisdom Exp (Wise lore)
e
          (Stm (Wise lore), Stms lore) -> m (Stm (Wise lore), Stms lore)
forall (m :: * -> *) a. Monad m => a -> m a
return (PatternT (LetAttr lore)
-> StmAux (ExpAttr lore) -> Exp (Wise lore) -> Stm (Wise lore)
forall lore.
(Attributes lore, CanBeWise (Op lore)) =>
Pattern lore
-> StmAux (ExpAttr lore) -> Exp (Wise lore) -> Stm (Wise lore)
mkWiseLetStm PatternT (LetAttr lore)
pat' (() -> StmAux ()
forall attr. attr -> StmAux attr
defAux ()) Exp (Wise lore)
e, Stms lore
stms)
          where env :: Scope lore
env = Scope (Wise lore) -> Scope lore
forall lore. Scope (Wise lore) -> Scope lore
removeScopeWisdom (Scope (Wise lore) -> Scope lore)
-> Scope (Wise lore) -> Scope lore
forall a b. (a -> b) -> a -> b
$ SymbolTable (Wise lore) -> Scope (Wise lore)
forall lore. SymbolTable lore -> Scope lore
ST.toScope SymbolTable (Wise lore)
vtable

        protectOp :: SubExp
-> PatternT (LetAttr (Lore m)) -> MemOp inner -> Maybe (m ())
protectOp SubExp
taken PatternT (LetAttr (Lore m))
pat (Alloc SubExp
size Space
space) = m () -> Maybe (m ())
forall a. a -> Maybe a
Just (m () -> Maybe (m ())) -> m () -> Maybe (m ())
forall a b. (a -> b) -> a -> b
$ do
          BodyT (Lore m)
tbody <- Result -> m (BodyT (Lore m))
forall (m :: * -> *). MonadBinder m => Result -> m (Body (Lore m))
resultBodyM [SubExp
size]
          BodyT (Lore m)
fbody <- Result -> m (BodyT (Lore m))
forall (m :: * -> *). MonadBinder m => Result -> m (Body (Lore m))
resultBodyM [IntType -> Integer -> SubExp
intConst IntType
Int64 Integer
0]
          SubExp
size' <- String -> Exp (Lore m) -> m SubExp
forall (m :: * -> *).
MonadBinder m =>
String -> Exp (Lore m) -> m SubExp
letSubExp String
"hoisted_alloc_size" (Exp (Lore m) -> m SubExp) -> Exp (Lore m) -> m SubExp
forall a b. (a -> b) -> a -> b
$
                   SubExp
-> BodyT (Lore m)
-> BodyT (Lore m)
-> IfAttr (BranchType (Lore m))
-> Exp (Lore m)
forall lore.
SubExp
-> BodyT lore
-> BodyT lore
-> IfAttr (BranchType lore)
-> ExpT lore
If SubExp
taken BodyT (Lore m)
tbody BodyT (Lore m)
fbody (IfAttr (BranchType (Lore m)) -> Exp (Lore m))
-> IfAttr (BranchType (Lore m)) -> Exp (Lore m)
forall a b. (a -> b) -> a -> b
$ [MemInfo d u ret] -> IfSort -> IfAttr (MemInfo d u ret)
forall rt. [rt] -> IfSort -> IfAttr rt
IfAttr [PrimType -> MemInfo d u ret
forall d u ret. PrimType -> MemInfo d u ret
MemPrim PrimType
int64] IfSort
IfFallback
          PatternT (LetAttr (Lore m)) -> Exp (Lore m) -> m ()
forall (m :: * -> *).
MonadBinder m =>
Pattern (Lore m) -> Exp (Lore m) -> m ()
letBind_ PatternT (LetAttr (Lore m))
pat (Exp (Lore m) -> m ()) -> Exp (Lore m) -> m ()
forall a b. (a -> b) -> a -> b
$ Op (Lore m) -> Exp (Lore m)
forall lore. Op lore -> ExpT lore
Op (Op (Lore m) -> Exp (Lore m)) -> Op (Lore m) -> Exp (Lore m)
forall a b. (a -> b) -> a -> b
$ SubExp -> Space -> MemOp inner
forall inner. SubExp -> Space -> MemOp inner
Alloc SubExp
size' Space
space
        protectOp SubExp
_ PatternT (LetAttr (Lore m))
_ MemOp inner
_ = Maybe (m ())
forall a. Maybe a
Nothing

        simplifyOp :: MemOp inner
-> SimpleM lore (MemOp (OpWithWisdom inner), Stms (Wise lore))
simplifyOp (Alloc SubExp
size Space
space) =
          (,) (MemOp (OpWithWisdom inner)
 -> Stms (Wise lore)
 -> (MemOp (OpWithWisdom inner), Stms (Wise lore)))
-> SimpleM lore (MemOp (OpWithWisdom inner))
-> SimpleM
     lore
     (Stms (Wise lore)
      -> (MemOp (OpWithWisdom inner), Stms (Wise lore)))
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> (SubExp -> Space -> MemOp (OpWithWisdom inner)
forall inner. SubExp -> Space -> MemOp inner
Alloc (SubExp -> Space -> MemOp (OpWithWisdom inner))
-> SimpleM lore SubExp
-> SimpleM lore (Space -> MemOp (OpWithWisdom inner))
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> SubExp -> SimpleM lore SubExp
forall e lore.
(Simplifiable e, SimplifiableLore lore) =>
e -> SimpleM lore e
Engine.simplify SubExp
size SimpleM lore (Space -> MemOp (OpWithWisdom inner))
-> SimpleM lore Space -> SimpleM lore (MemOp (OpWithWisdom inner))
forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> Space -> SimpleM lore Space
forall (f :: * -> *) a. Applicative f => a -> f a
pure Space
space) SimpleM
  lore
  (Stms (Wise lore)
   -> (MemOp (OpWithWisdom inner), Stms (Wise lore)))
-> SimpleM lore (Stms (Wise lore))
-> SimpleM lore (MemOp (OpWithWisdom inner), Stms (Wise lore))
forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> Stms (Wise lore) -> SimpleM lore (Stms (Wise lore))
forall (f :: * -> *) a. Applicative f => a -> f a
pure Stms (Wise lore)
forall a. Monoid a => a
mempty
        simplifyOp (Inner inner
k) = do (OpWithWisdom inner
k', Stms (Wise lore)
hoisted) <- inner -> SimpleM lore (OpWithWisdom inner, Stms (Wise lore))
simplifyInnerOp inner
k
                                  (MemOp (OpWithWisdom inner), Stms (Wise lore))
-> SimpleM lore (MemOp (OpWithWisdom inner), Stms (Wise lore))
forall (m :: * -> *) a. Monad m => a -> m a
return (OpWithWisdom inner -> MemOp (OpWithWisdom inner)
forall inner. inner -> MemOp inner
Inner OpWithWisdom inner
k', Stms (Wise lore)
hoisted)

bindPatternWithAllocations :: (MonadBinder m,
                               ExpAttr lore ~ (),
                               Op (Lore m) ~ MemOp inner,
                               Allocator lore (PatAllocM lore)) =>
                              Scope lore -> [VName] -> Exp lore
                           -> m (Pattern lore)
bindPatternWithAllocations :: Scope lore -> [VName] -> Exp lore -> m (Pattern lore)
bindPatternWithAllocations Scope lore
types [VName]
names Exp lore
e = do
  (PatternT LetAttrMem
pat,[AllocStm]
prebnds) <- PatAllocM lore (PatternT LetAttrMem)
-> Scope lore -> m (PatternT LetAttrMem, [AllocStm])
forall (m :: * -> *) lore a.
MonadFreshNames m =>
PatAllocM lore a -> Scope lore -> m (a, [AllocStm])
runPatAllocM ([VName] -> Exp lore -> PatAllocM lore (Pattern lore)
forall lore (m :: * -> *).
(Allocator lore m, ExpAttr lore ~ ()) =>
[VName] -> Exp lore -> m (Pattern lore)
patternWithAllocations [VName]
names Exp lore
e) Scope lore
types
  (AllocStm -> m ()) -> [AllocStm] -> m ()
forall (t :: * -> *) (m :: * -> *) a b.
(Foldable t, Monad m) =>
(a -> m b) -> t a -> m ()
mapM_ AllocStm -> m ()
forall (m :: * -> *) inner.
(MonadBinder m, Op (Lore m) ~ MemOp inner) =>
AllocStm -> m ()
bindAllocStm [AllocStm]
prebnds
  PatternT LetAttrMem -> m (PatternT LetAttrMem)
forall (m :: * -> *) a. Monad m => a -> m a
return PatternT LetAttrMem
pat

data ExpHint = NoHint
             | Hint IxFun Space

defaultExpHints :: (Monad m, Attributes lore) => Exp lore -> m [ExpHint]
defaultExpHints :: Exp lore -> m [ExpHint]
defaultExpHints Exp lore
e = [ExpHint] -> m [ExpHint]
forall (m :: * -> *) a. Monad m => a -> m a
return ([ExpHint] -> m [ExpHint]) -> [ExpHint] -> m [ExpHint]
forall a b. (a -> b) -> a -> b
$ Int -> ExpHint -> [ExpHint]
forall a. Int -> a -> [a]
replicate (Exp lore -> Int
forall lore.
(Annotations lore, TypedOp (Op lore)) =>
Exp lore -> Int
expExtTypeSize Exp lore
e) ExpHint
NoHint