{-# LANGUAGE TypeFamilies #-}
{-# LANGUAGE UndecidableInstances #-}

-- | Building blocks for defining representations where every array
-- is given information about which memory block is it based in, and
-- how array elements map to memory block offsets.
--
-- There are two primary concepts you will need to understand:
--
--  1. Memory blocks, which are Futhark values of type v'Mem'
--     (parametrized with their size).  These correspond to arbitrary
--     blocks of memory, and are created using the 'Alloc' operation.
--
--  2. Index functions, which describe a mapping from the index space
--     of an array (eg. a two-dimensional space for an array of type
--     @[[int]]@) to a one-dimensional offset into a memory block.
--     Thus, index functions describe how arbitrary-dimensional arrays
--     are mapped to the single-dimensional world of memory.
--
-- At a conceptual level, imagine that we have a two-dimensional array
-- @a@ of 32-bit integers, consisting of @n@ rows of @m@ elements
-- each.  This array could be represented in classic row-major format
-- with an index function like the following:
--
-- @
--   f(i,j) = i * m + j
-- @
--
-- When we want to know the location of element @a[2,3]@, we simply
-- call the index function as @f(2,3)@ and obtain @2*m+3@.  We could
-- also have chosen another index function, one that represents the
-- array in column-major (or "transposed") format:
--
-- @
--   f(i,j) = j * n + i
-- @
--
-- Index functions are not Futhark-level functions, but a special
-- construct that the final code generator will eventually use to
-- generate concrete access code.  By modifying the index functions we
-- can change how an array is represented in memory, which can permit
-- memory access pattern optimisations.
--
-- Every time we bind an array, whether in a @let@-binding, @loop@
-- merge parameter, or @lambda@ parameter, we have an annotation
-- specifying a memory block and an index function.  In some cases,
-- such as @let@-bindings for many expressions, we are free to specify
-- an arbitrary index function and memory block - for example, we get
-- to decide where 'Copy' stores its result - but in other cases the
-- type rules of the expression chooses for us.  For example, 'Index'
-- always produces an array in the same memory block as its input, and
-- with the same index function, except with some indices fixed.
module Futhark.IR.Mem
  ( LetDecMem,
    FParamMem,
    LParamMem,
    RetTypeMem,
    BranchTypeMem,
    MemOp (..),
    traverseMemOpStms,
    MemInfo (..),
    MemBound,
    MemBind (..),
    MemReturn (..),
    LMAD,
    ExtLMAD,
    isStaticLMAD,
    ExpReturns,
    BodyReturns,
    FunReturns,
    noModeReturns,
    bodyReturnsToExpReturns,
    Mem,
    HasLetDecMem (..),
    OpReturns (..),
    varReturns,
    expReturns,
    extReturns,
    nameInfoToMemInfo,
    lookupMemInfo,
    subExpMemInfo,
    lookupArraySummary,
    lookupMemSpace,
    existentialiseLMAD,

    -- * Type checking parts
    matchBranchReturnType,
    matchPatToExp,
    matchFunctionReturnType,
    matchLoopResultMem,
    bodyReturnsFromPat,
    checkMemInfo,

    -- * Module re-exports
    module Futhark.IR.Prop,
    module Futhark.IR.Traversals,
    module Futhark.IR.Pretty,
    module Futhark.IR.Syntax,
    module Futhark.Analysis.PrimExp.Convert,
  )
where

import Control.Category
import Control.Monad
import Control.Monad.Except
import Control.Monad.Reader
import Control.Monad.State
import Data.Foldable (traverse_)
import Data.Function ((&))
import Data.Kind qualified
import Data.List (elemIndex, find)
import Data.Map.Strict qualified as M
import Data.Maybe
import Data.Text qualified as T
import Futhark.Analysis.Metrics
import Futhark.Analysis.PrimExp.Convert
import Futhark.Analysis.PrimExp.Simplify
import Futhark.Analysis.SymbolTable qualified as ST
import Futhark.IR.Aliases
  ( Aliases,
    CanBeAliased (..),
    removeExpAliases,
    removePatAliases,
    removeScopeAliases,
  )
import Futhark.IR.Mem.LMAD qualified as LMAD
import Futhark.IR.Pretty
import Futhark.IR.Prop
import Futhark.IR.Prop.Aliases
import Futhark.IR.Syntax
import Futhark.IR.Traversals
import Futhark.IR.TypeCheck qualified as TC
import Futhark.Optimise.Simplify.Engine qualified as Engine
import Futhark.Optimise.Simplify.Rep
import Futhark.Transform.Rename
import Futhark.Transform.Substitute
import Futhark.Util
import Futhark.Util.Pretty (docText, indent, ppTupleLines', pretty, (<+>), (</>))
import Futhark.Util.Pretty qualified as PP
import Prelude hiding (id, (.))

type LetDecMem = MemInfo SubExp NoMode MemBind

type FParamMem = MemInfo SubExp Diet MemBind

type LParamMem = MemInfo SubExp NoMode MemBind

type RetTypeMem = FunReturns

type BranchTypeMem = BodyReturns

-- | The class of pattern element decorators that contain memory
-- information.
class HasLetDecMem t where
  letDecMem :: t -> LetDecMem

instance HasLetDecMem LetDecMem where
  letDecMem :: LetDecMem -> LetDecMem
letDecMem = LetDecMem -> LetDecMem
forall a. a -> a
forall {k} (cat :: k -> k -> *) (a :: k). Category cat => cat a a
id

instance (HasLetDecMem b) => HasLetDecMem (a, b) where
  letDecMem :: (a, b) -> LetDecMem
letDecMem = b -> LetDecMem
forall t. HasLetDecMem t => t -> LetDecMem
letDecMem (b -> LetDecMem) -> ((a, b) -> b) -> (a, b) -> LetDecMem
forall b c a. (b -> c) -> (a -> b) -> a -> c
forall {k} (cat :: k -> k -> *) (b :: k) (c :: k) (a :: k).
Category cat =>
cat b c -> cat a b -> cat a c
. (a, b) -> b
forall a b. (a, b) -> b
snd

type Mem rep inner =
  ( FParamInfo rep ~ FParamMem,
    LParamInfo rep ~ LParamMem,
    HasLetDecMem (LetDec rep),
    RetType rep ~ RetTypeMem,
    BranchType rep ~ BranchTypeMem,
    ASTRep rep,
    OpReturns inner,
    RephraseOp inner,
    ASTConstraints (inner rep),
    FreeIn (inner rep),
    OpC rep ~ MemOp inner
  )

instance IsRetType FunReturns where
  primRetType :: PrimType -> FunReturns
primRetType = PrimType -> FunReturns
forall d o ret. PrimType -> MemInfo d o ret
MemPrim
  applyRetType :: forall dec.
Typed dec =>
[FunReturns]
-> [Param dec] -> [(SubExp, Type)] -> Maybe [FunReturns]
applyRetType = [FunReturns]
-> [Param dec] -> [(SubExp, Type)] -> Maybe [FunReturns]
forall dec.
Typed dec =>
[FunReturns]
-> [Param dec] -> [(SubExp, Type)] -> Maybe [FunReturns]
applyFunReturns

instance IsBodyType BodyReturns where
  primBodyType :: PrimType -> FunReturns
primBodyType = PrimType -> FunReturns
forall d o ret. PrimType -> MemInfo d o ret
MemPrim

data MemOp (inner :: Data.Kind.Type -> Data.Kind.Type) (rep :: Data.Kind.Type)
  = -- | Allocate a memory block.
    Alloc SubExp Space
  | -- | If the given array has a direct index function (as 'LMAD.isDirect'),
    -- returns the array unchanged. Otherwise, make a copy of the array. The
    -- choice is made dynamically, and the goal is to handle cases where we must
    -- return something that is direct (mostly entry point return types), but we
    -- cannot statically be sure of the layout of some array.
    EnsureDirect VName
  | Inner (inner rep)
  deriving (MemOp inner rep -> MemOp inner rep -> Bool
(MemOp inner rep -> MemOp inner rep -> Bool)
-> (MemOp inner rep -> MemOp inner rep -> Bool)
-> Eq (MemOp inner rep)
forall a. (a -> a -> Bool) -> (a -> a -> Bool) -> Eq a
forall (inner :: * -> *) rep.
Eq (inner rep) =>
MemOp inner rep -> MemOp inner rep -> Bool
$c== :: forall (inner :: * -> *) rep.
Eq (inner rep) =>
MemOp inner rep -> MemOp inner rep -> Bool
== :: MemOp inner rep -> MemOp inner rep -> Bool
$c/= :: forall (inner :: * -> *) rep.
Eq (inner rep) =>
MemOp inner rep -> MemOp inner rep -> Bool
/= :: MemOp inner rep -> MemOp inner rep -> Bool
Eq, Eq (MemOp inner rep)
Eq (MemOp inner rep) =>
(MemOp inner rep -> MemOp inner rep -> Ordering)
-> (MemOp inner rep -> MemOp inner rep -> Bool)
-> (MemOp inner rep -> MemOp inner rep -> Bool)
-> (MemOp inner rep -> MemOp inner rep -> Bool)
-> (MemOp inner rep -> MemOp inner rep -> Bool)
-> (MemOp inner rep -> MemOp inner rep -> MemOp inner rep)
-> (MemOp inner rep -> MemOp inner rep -> MemOp inner rep)
-> Ord (MemOp inner rep)
MemOp inner rep -> MemOp inner rep -> Bool
MemOp inner rep -> MemOp inner rep -> Ordering
MemOp inner rep -> MemOp inner rep -> MemOp inner rep
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
forall (inner :: * -> *) rep.
Ord (inner rep) =>
Eq (MemOp inner rep)
forall (inner :: * -> *) rep.
Ord (inner rep) =>
MemOp inner rep -> MemOp inner rep -> Bool
forall (inner :: * -> *) rep.
Ord (inner rep) =>
MemOp inner rep -> MemOp inner rep -> Ordering
forall (inner :: * -> *) rep.
Ord (inner rep) =>
MemOp inner rep -> MemOp inner rep -> MemOp inner rep
$ccompare :: forall (inner :: * -> *) rep.
Ord (inner rep) =>
MemOp inner rep -> MemOp inner rep -> Ordering
compare :: MemOp inner rep -> MemOp inner rep -> Ordering
$c< :: forall (inner :: * -> *) rep.
Ord (inner rep) =>
MemOp inner rep -> MemOp inner rep -> Bool
< :: MemOp inner rep -> MemOp inner rep -> Bool
$c<= :: forall (inner :: * -> *) rep.
Ord (inner rep) =>
MemOp inner rep -> MemOp inner rep -> Bool
<= :: MemOp inner rep -> MemOp inner rep -> Bool
$c> :: forall (inner :: * -> *) rep.
Ord (inner rep) =>
MemOp inner rep -> MemOp inner rep -> Bool
> :: MemOp inner rep -> MemOp inner rep -> Bool
$c>= :: forall (inner :: * -> *) rep.
Ord (inner rep) =>
MemOp inner rep -> MemOp inner rep -> Bool
>= :: MemOp inner rep -> MemOp inner rep -> Bool
$cmax :: forall (inner :: * -> *) rep.
Ord (inner rep) =>
MemOp inner rep -> MemOp inner rep -> MemOp inner rep
max :: MemOp inner rep -> MemOp inner rep -> MemOp inner rep
$cmin :: forall (inner :: * -> *) rep.
Ord (inner rep) =>
MemOp inner rep -> MemOp inner rep -> MemOp inner rep
min :: MemOp inner rep -> MemOp inner rep -> MemOp inner rep
Ord, Int -> MemOp inner rep -> ShowS
[MemOp inner rep] -> ShowS
MemOp inner rep -> String
(Int -> MemOp inner rep -> ShowS)
-> (MemOp inner rep -> String)
-> ([MemOp inner rep] -> ShowS)
-> Show (MemOp inner rep)
forall a.
(Int -> a -> ShowS) -> (a -> String) -> ([a] -> ShowS) -> Show a
forall (inner :: * -> *) rep.
Show (inner rep) =>
Int -> MemOp inner rep -> ShowS
forall (inner :: * -> *) rep.
Show (inner rep) =>
[MemOp inner rep] -> ShowS
forall (inner :: * -> *) rep.
Show (inner rep) =>
MemOp inner rep -> String
$cshowsPrec :: forall (inner :: * -> *) rep.
Show (inner rep) =>
Int -> MemOp inner rep -> ShowS
showsPrec :: Int -> MemOp inner rep -> ShowS
$cshow :: forall (inner :: * -> *) rep.
Show (inner rep) =>
MemOp inner rep -> String
show :: MemOp inner rep -> String
$cshowList :: forall (inner :: * -> *) rep.
Show (inner rep) =>
[MemOp inner rep] -> ShowS
showList :: [MemOp inner rep] -> ShowS
Show)

-- | A helper for defining 'TraverseOpStms'.
traverseMemOpStms ::
  (Monad m) =>
  OpStmsTraverser m (inner rep) rep ->
  OpStmsTraverser m (MemOp inner rep) rep
traverseMemOpStms :: forall (m :: * -> *) (inner :: * -> *) rep.
Monad m =>
OpStmsTraverser m (inner rep) rep
-> OpStmsTraverser m (MemOp inner rep) rep
traverseMemOpStms OpStmsTraverser m (inner rep) rep
_ Scope rep -> Stms rep -> m (Stms rep)
_ op :: MemOp inner rep
op@Alloc {} = MemOp inner rep -> m (MemOp inner rep)
forall a. a -> m a
forall (f :: * -> *) a. Applicative f => a -> f a
pure MemOp inner rep
op
traverseMemOpStms OpStmsTraverser m (inner rep) rep
_ Scope rep -> Stms rep -> m (Stms rep)
_ op :: MemOp inner rep
op@EnsureDirect {} = MemOp inner rep -> m (MemOp inner rep)
forall a. a -> m a
forall (f :: * -> *) a. Applicative f => a -> f a
pure MemOp inner rep
op
traverseMemOpStms OpStmsTraverser m (inner rep) rep
onInner Scope rep -> Stms rep -> m (Stms rep)
f (Inner inner rep
inner) = inner rep -> MemOp inner rep
forall (inner :: * -> *) rep. inner rep -> MemOp inner rep
Inner (inner rep -> MemOp inner rep)
-> m (inner rep) -> m (MemOp inner rep)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> OpStmsTraverser m (inner rep) rep
onInner Scope rep -> Stms rep -> m (Stms rep)
f inner rep
inner

instance (RephraseOp inner) => RephraseOp (MemOp inner) where
  rephraseInOp :: forall (m :: * -> *) from to.
Monad m =>
Rephraser m from to -> MemOp inner from -> m (MemOp inner to)
rephraseInOp Rephraser m from to
_ (Alloc SubExp
e Space
space) = MemOp inner to -> m (MemOp inner to)
forall a. a -> m a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (SubExp -> Space -> MemOp inner to
forall (inner :: * -> *) rep. SubExp -> Space -> MemOp inner rep
Alloc SubExp
e Space
space)
  rephraseInOp Rephraser m from to
_ (EnsureDirect VName
v) = MemOp inner to -> m (MemOp inner to)
forall a. a -> m a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (VName -> MemOp inner to
forall (inner :: * -> *) rep. VName -> MemOp inner rep
EnsureDirect VName
v)
  rephraseInOp Rephraser m from to
r (Inner inner from
x) = inner to -> MemOp inner to
forall (inner :: * -> *) rep. inner rep -> MemOp inner rep
Inner (inner to -> MemOp inner to) -> m (inner to) -> m (MemOp inner to)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Rephraser m from to -> inner from -> m (inner to)
forall (op :: * -> *) (m :: * -> *) from to.
(RephraseOp op, Monad m) =>
Rephraser m from to -> op from -> m (op to)
forall (m :: * -> *) from to.
Monad m =>
Rephraser m from to -> inner from -> m (inner to)
rephraseInOp Rephraser m from to
r inner from
x

instance (FreeIn (inner rep)) => FreeIn (MemOp inner rep) where
  freeIn' :: MemOp inner rep -> FV
freeIn' (Alloc SubExp
size Space
_) = SubExp -> FV
forall a. FreeIn a => a -> FV
freeIn' SubExp
size
  freeIn' (EnsureDirect VName
v) = VName -> FV
forall a. FreeIn a => a -> FV
freeIn' VName
v
  freeIn' (Inner inner rep
k) = inner rep -> FV
forall a. FreeIn a => a -> FV
freeIn' inner rep
k

instance (TypedOp inner) => TypedOp (MemOp inner) where
  opType :: forall rep (m :: * -> *).
HasScope rep m =>
MemOp inner rep -> m [ExtType]
opType (Alloc SubExp
_ Space
space) = [ExtType] -> m [ExtType]
forall a. a -> m a
forall (f :: * -> *) a. Applicative f => a -> f a
pure [Space -> ExtType
forall shape o. Space -> TypeBase shape o
Mem Space
space]
  opType (EnsureDirect VName
v) = Type -> [ExtType]
f (Type -> [ExtType]) -> m Type -> m [ExtType]
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> VName -> m Type
forall rep (m :: * -> *). HasScope rep m => VName -> m Type
lookupType VName
v
    where
      f :: Type -> [ExtType]
f (Array PrimType
pt ShapeBase SubExp
shape NoMode
_) =
        [Space -> ExtType
forall shape o. Space -> TypeBase shape o
Mem Space
DefaultSpace, PrimType -> ShapeBase (Ext SubExp) -> NoMode -> ExtType
forall shape o. PrimType -> shape -> o -> TypeBase shape o
Array PrimType
pt ((SubExp -> Ext SubExp)
-> ShapeBase SubExp -> ShapeBase (Ext SubExp)
forall a b. (a -> b) -> ShapeBase a -> ShapeBase b
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap SubExp -> Ext SubExp
forall a. a -> Ext a
Free ShapeBase SubExp
shape) NoMode
NoMode]
      f Type
_ = String -> [ExtType]
forall a. HasCallStack => String -> a
error (String -> [ExtType]) -> String -> [ExtType]
forall a b. (a -> b) -> a -> b
$ String
"EnsureDirect applied to non-array: " String -> ShowS
forall a. [a] -> [a] -> [a]
++ VName -> String
forall a. Show a => a -> String
show VName
v
  opType (Inner inner rep
k) = inner rep -> m [ExtType]
forall rep (m :: * -> *).
HasScope rep m =>
inner rep -> m [ExtType]
forall (op :: * -> *) rep (m :: * -> *).
(TypedOp op, HasScope rep m) =>
op rep -> m [ExtType]
opType inner rep
k

instance (AliasedOp inner) => AliasedOp (MemOp inner) where
  opAliases :: forall rep. Aliased rep => MemOp inner rep -> [Names]
opAliases Alloc {} = [Names
forall a. Monoid a => a
mempty]
  -- XXX - why not mem aliases for EnsureDirect?
  opAliases (EnsureDirect VName
v) = [Names
forall a. Monoid a => a
mempty, VName -> Names
oneName VName
v]
  opAliases (Inner inner rep
k) = inner rep -> [Names]
forall rep. Aliased rep => inner rep -> [Names]
forall (op :: * -> *) rep.
(AliasedOp op, Aliased rep) =>
op rep -> [Names]
opAliases inner rep
k

  consumedInOp :: forall rep. Aliased rep => MemOp inner rep -> Names
consumedInOp Alloc {} = Names
forall a. Monoid a => a
mempty
  consumedInOp EnsureDirect {} = Names
forall a. Monoid a => a
mempty
  consumedInOp (Inner inner rep
k) = inner rep -> Names
forall rep. Aliased rep => inner rep -> Names
forall (op :: * -> *) rep.
(AliasedOp op, Aliased rep) =>
op rep -> Names
consumedInOp inner rep
k

instance (CanBeAliased inner) => CanBeAliased (MemOp inner) where
  addOpAliases :: forall rep.
AliasableRep rep =>
AliasTable -> MemOp inner rep -> MemOp inner (Aliases rep)
addOpAliases AliasTable
_ (Alloc SubExp
se Space
space) = SubExp -> Space -> MemOp inner (Aliases rep)
forall (inner :: * -> *) rep. SubExp -> Space -> MemOp inner rep
Alloc SubExp
se Space
space
  addOpAliases AliasTable
_ (EnsureDirect VName
v) = VName -> MemOp inner (Aliases rep)
forall (inner :: * -> *) rep. VName -> MemOp inner rep
EnsureDirect VName
v
  addOpAliases AliasTable
aliases (Inner inner rep
k) = inner (Aliases rep) -> MemOp inner (Aliases rep)
forall (inner :: * -> *) rep. inner rep -> MemOp inner rep
Inner (inner (Aliases rep) -> MemOp inner (Aliases rep))
-> inner (Aliases rep) -> MemOp inner (Aliases rep)
forall a b. (a -> b) -> a -> b
$ AliasTable -> inner rep -> inner (Aliases rep)
forall rep.
AliasableRep rep =>
AliasTable -> inner rep -> inner (Aliases rep)
forall (op :: * -> *) rep.
(CanBeAliased op, AliasableRep rep) =>
AliasTable -> op rep -> op (Aliases rep)
addOpAliases AliasTable
aliases inner rep
k

instance (Rename (inner rep)) => Rename (MemOp inner rep) where
  rename :: MemOp inner rep -> RenameM (MemOp inner rep)
rename (Alloc SubExp
size Space
space) = SubExp -> Space -> MemOp inner rep
forall (inner :: * -> *) rep. SubExp -> Space -> MemOp inner rep
Alloc (SubExp -> Space -> MemOp inner rep)
-> RenameM SubExp -> RenameM (Space -> MemOp inner rep)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> SubExp -> RenameM SubExp
forall a. Rename a => a -> RenameM a
rename SubExp
size RenameM (Space -> MemOp inner rep)
-> RenameM Space -> RenameM (MemOp inner rep)
forall a b. RenameM (a -> b) -> RenameM a -> RenameM b
forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> Space -> RenameM Space
forall a. a -> RenameM a
forall (f :: * -> *) a. Applicative f => a -> f a
pure Space
space
  rename (EnsureDirect VName
v) = VName -> MemOp inner rep
forall (inner :: * -> *) rep. VName -> MemOp inner rep
EnsureDirect (VName -> MemOp inner rep)
-> RenameM VName -> RenameM (MemOp inner rep)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> VName -> RenameM VName
forall a. Rename a => a -> RenameM a
rename VName
v
  rename (Inner inner rep
k) = inner rep -> MemOp inner rep
forall (inner :: * -> *) rep. inner rep -> MemOp inner rep
Inner (inner rep -> MemOp inner rep)
-> RenameM (inner rep) -> RenameM (MemOp inner rep)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> inner rep -> RenameM (inner rep)
forall a. Rename a => a -> RenameM a
rename inner rep
k

instance (Substitute (inner rep)) => Substitute (MemOp inner rep) where
  substituteNames :: Map VName VName -> MemOp inner rep -> MemOp inner rep
substituteNames Map VName VName
subst (Alloc SubExp
size Space
space) = SubExp -> Space -> MemOp inner rep
forall (inner :: * -> *) rep. SubExp -> Space -> MemOp inner rep
Alloc (Map VName VName -> SubExp -> SubExp
forall a. Substitute a => Map VName VName -> a -> a
substituteNames Map VName VName
subst SubExp
size) Space
space
  substituteNames Map VName VName
subst (EnsureDirect VName
v) = VName -> MemOp inner rep
forall (inner :: * -> *) rep. VName -> MemOp inner rep
EnsureDirect (Map VName VName -> VName -> VName
forall a. Substitute a => Map VName VName -> a -> a
substituteNames Map VName VName
subst VName
v)
  substituteNames Map VName VName
subst (Inner inner rep
k) = inner rep -> MemOp inner rep
forall (inner :: * -> *) rep. inner rep -> MemOp inner rep
Inner (inner rep -> MemOp inner rep) -> inner rep -> MemOp inner rep
forall a b. (a -> b) -> a -> b
$ Map VName VName -> inner rep -> inner rep
forall a. Substitute a => Map VName VName -> a -> a
substituteNames Map VName VName
subst inner rep
k

instance (PP.Pretty (inner rep)) => PP.Pretty (MemOp inner rep) where
  pretty :: forall ann. MemOp inner rep -> Doc ann
pretty (Alloc SubExp
e Space
DefaultSpace) = Doc ann
"alloc" Doc ann -> Doc ann -> Doc ann
forall a. Semigroup a => a -> a -> a
<> [Doc ann] -> Doc ann
forall a. [Doc a] -> Doc a
PP.apply [SubExp -> Doc ann
forall a ann. Pretty a => a -> Doc ann
forall ann. SubExp -> Doc ann
PP.pretty SubExp
e]
  pretty (Alloc SubExp
e Space
s) = Doc ann
"alloc" Doc ann -> Doc ann -> Doc ann
forall a. Semigroup a => a -> a -> a
<> [Doc ann] -> Doc ann
forall a. [Doc a] -> Doc a
PP.apply [SubExp -> Doc ann
forall a ann. Pretty a => a -> Doc ann
forall ann. SubExp -> Doc ann
PP.pretty SubExp
e, Space -> Doc ann
forall a ann. Pretty a => a -> Doc ann
forall ann. Space -> Doc ann
PP.pretty Space
s]
  pretty (EnsureDirect VName
v) = Doc ann
"ensure_direct" Doc ann -> Doc ann -> Doc ann
forall a. Semigroup a => a -> a -> a
<> [Doc ann] -> Doc ann
forall a. [Doc a] -> Doc a
PP.apply [VName -> Doc ann
forall a ann. Pretty a => a -> Doc ann
forall ann. VName -> Doc ann
PP.pretty VName
v]
  pretty (Inner inner rep
k) = inner rep -> Doc ann
forall ann. inner rep -> Doc ann
forall a ann. Pretty a => a -> Doc ann
PP.pretty inner rep
k

instance (OpMetrics (inner rep)) => OpMetrics (MemOp inner rep) where
  opMetrics :: MemOp inner rep -> MetricsM ()
opMetrics Alloc {} = Text -> MetricsM ()
seen Text
"Alloc"
  opMetrics EnsureDirect {} = Text -> MetricsM ()
seen Text
"EnsureDirect"
  opMetrics (Inner inner rep
k) = inner rep -> MetricsM ()
forall op. OpMetrics op => op -> MetricsM ()
opMetrics inner rep
k

instance (IsOp inner) => IsOp (MemOp inner) where
  safeOp :: forall rep. ASTRep rep => MemOp inner rep -> Bool
safeOp (Alloc (Constant (IntValue (Int64Value Int64
k))) Space
_) = Int64
k Int64 -> Int64 -> Bool
forall a. Ord a => a -> a -> Bool
>= Int64
0
  safeOp Alloc {} = Bool
False
  safeOp EnsureDirect {} = Bool
True
  safeOp (Inner inner rep
k) = inner rep -> Bool
forall rep. ASTRep rep => inner rep -> Bool
forall (op :: * -> *) rep. (IsOp op, ASTRep rep) => op rep -> Bool
safeOp inner rep
k
  cheapOp :: forall rep. ASTRep rep => MemOp inner rep -> Bool
cheapOp (Inner inner rep
k) = inner rep -> Bool
forall rep. ASTRep rep => inner rep -> Bool
forall (op :: * -> *) rep. (IsOp op, ASTRep rep) => op rep -> Bool
cheapOp inner rep
k
  cheapOp Alloc {} = Bool
True
  cheapOp EnsureDirect {} = Bool
False
  opDependencies :: forall rep. ASTRep rep => MemOp inner rep -> [Names]
opDependencies (Alloc SubExp
_ Space
e) = [Space -> Names
forall a. FreeIn a => a -> Names
freeIn Space
e]
  opDependencies (EnsureDirect VName
v) = [Names
forall a. Monoid a => a
mempty, VName -> Names
forall a. FreeIn a => a -> Names
freeIn VName
v]
  opDependencies (Inner inner rep
op) = inner rep -> [Names]
forall rep. ASTRep rep => inner rep -> [Names]
forall (op :: * -> *) rep.
(IsOp op, ASTRep rep) =>
op rep -> [Names]
opDependencies inner rep
op

instance (CanBeWise inner) => CanBeWise (MemOp inner) where
  addOpWisdom :: forall rep.
Informing rep =>
MemOp inner rep -> MemOp inner (Wise rep)
addOpWisdom (Alloc SubExp
size Space
space) = SubExp -> Space -> MemOp inner (Wise rep)
forall (inner :: * -> *) rep. SubExp -> Space -> MemOp inner rep
Alloc SubExp
size Space
space
  addOpWisdom (EnsureDirect VName
v) = VName -> MemOp inner (Wise rep)
forall (inner :: * -> *) rep. VName -> MemOp inner rep
EnsureDirect VName
v
  addOpWisdom (Inner inner rep
k) = inner (Wise rep) -> MemOp inner (Wise rep)
forall (inner :: * -> *) rep. inner rep -> MemOp inner rep
Inner (inner (Wise rep) -> MemOp inner (Wise rep))
-> inner (Wise rep) -> MemOp inner (Wise rep)
forall a b. (a -> b) -> a -> b
$ inner rep -> inner (Wise rep)
forall rep. Informing rep => inner rep -> inner (Wise rep)
forall (op :: * -> *) rep.
(CanBeWise op, Informing rep) =>
op rep -> op (Wise rep)
addOpWisdom inner rep
k

instance (ST.IndexOp (inner rep)) => ST.IndexOp (MemOp inner rep) where
  indexOp :: forall rep.
(ASTRep rep, IndexOp (Op rep)) =>
SymbolTable rep
-> Int
-> MemOp inner rep
-> [TPrimExp Int64 VName]
-> Maybe Indexed
indexOp SymbolTable rep
vtable Int
k (Inner inner rep
op) [TPrimExp Int64 VName]
is = SymbolTable rep
-> Int -> inner rep -> [TPrimExp Int64 VName] -> Maybe Indexed
forall rep.
(ASTRep rep, IndexOp (Op rep)) =>
SymbolTable rep
-> Int -> inner rep -> [TPrimExp Int64 VName] -> Maybe Indexed
forall op rep.
(IndexOp op, ASTRep rep, IndexOp (Op rep)) =>
SymbolTable rep
-> Int -> op -> [TPrimExp Int64 VName] -> Maybe Indexed
ST.indexOp SymbolTable rep
vtable Int
k inner rep
op [TPrimExp Int64 VName]
is
  indexOp SymbolTable rep
_ Int
_ MemOp inner rep
_ [TPrimExp Int64 VName]
_ = Maybe Indexed
forall a. Maybe a
Nothing

-- | The LMAD representation used for memory annotations.
type LMAD = LMAD.LMAD (TPrimExp Int64 VName)

-- | An index function that may contain existential variables.
type ExtLMAD = LMAD.LMAD (TPrimExp Int64 (Ext VName))

-- | A summary of the memory information for every let-bound identifier,
-- function parameter, and return value. Parameterisered over mode, dimension,
-- and auxiliary array information.
data MemInfo d o ret
  = -- | A primitive value.
    MemPrim PrimType
  | -- | A memory block.
    MemMem Space
  | -- | The array is stored in the named memory block, and with the
    -- given index function.  The index function maps indices in the
    -- array to /element/ offset, /not/ byte offsets!  To translate to
    -- byte offsets, multiply the offset with the size of the array
    -- element type.
    MemArray PrimType (ShapeBase d) o ret
  | -- | An accumulator, which is not stored anywhere.
    MemAcc VName Shape [Type]
  deriving (MemInfo d o ret -> MemInfo d o ret -> Bool
(MemInfo d o ret -> MemInfo d o ret -> Bool)
-> (MemInfo d o ret -> MemInfo d o ret -> Bool)
-> Eq (MemInfo d o ret)
forall a. (a -> a -> Bool) -> (a -> a -> Bool) -> Eq a
forall d o ret.
(Eq d, Eq o, Eq ret) =>
MemInfo d o ret -> MemInfo d o ret -> Bool
$c== :: forall d o ret.
(Eq d, Eq o, Eq ret) =>
MemInfo d o ret -> MemInfo d o ret -> Bool
== :: MemInfo d o ret -> MemInfo d o ret -> Bool
$c/= :: forall d o ret.
(Eq d, Eq o, Eq ret) =>
MemInfo d o ret -> MemInfo d o ret -> Bool
/= :: MemInfo d o ret -> MemInfo d o ret -> Bool
Eq, Int -> MemInfo d o ret -> ShowS
[MemInfo d o ret] -> ShowS
MemInfo d o ret -> String
(Int -> MemInfo d o ret -> ShowS)
-> (MemInfo d o ret -> String)
-> ([MemInfo d o ret] -> ShowS)
-> Show (MemInfo d o ret)
forall a.
(Int -> a -> ShowS) -> (a -> String) -> ([a] -> ShowS) -> Show a
forall d o ret.
(Show d, Show o, Show ret) =>
Int -> MemInfo d o ret -> ShowS
forall d o ret.
(Show d, Show o, Show ret) =>
[MemInfo d o ret] -> ShowS
forall d o ret.
(Show d, Show o, Show ret) =>
MemInfo d o ret -> String
$cshowsPrec :: forall d o ret.
(Show d, Show o, Show ret) =>
Int -> MemInfo d o ret -> ShowS
showsPrec :: Int -> MemInfo d o ret -> ShowS
$cshow :: forall d o ret.
(Show d, Show o, Show ret) =>
MemInfo d o ret -> String
show :: MemInfo d o ret -> String
$cshowList :: forall d o ret.
(Show d, Show o, Show ret) =>
[MemInfo d o ret] -> ShowS
showList :: [MemInfo d o ret] -> ShowS
Show, Eq (MemInfo d o ret)
Eq (MemInfo d o ret) =>
(MemInfo d o ret -> MemInfo d o ret -> Ordering)
-> (MemInfo d o ret -> MemInfo d o ret -> Bool)
-> (MemInfo d o ret -> MemInfo d o ret -> Bool)
-> (MemInfo d o ret -> MemInfo d o ret -> Bool)
-> (MemInfo d o ret -> MemInfo d o ret -> Bool)
-> (MemInfo d o ret -> MemInfo d o ret -> MemInfo d o ret)
-> (MemInfo d o ret -> MemInfo d o ret -> MemInfo d o ret)
-> Ord (MemInfo d o ret)
MemInfo d o ret -> MemInfo d o ret -> Bool
MemInfo d o ret -> MemInfo d o ret -> Ordering
MemInfo d o ret -> MemInfo d o ret -> MemInfo d o ret
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
forall d o ret. (Ord d, Ord o, Ord ret) => Eq (MemInfo d o ret)
forall d o ret.
(Ord d, Ord o, Ord ret) =>
MemInfo d o ret -> MemInfo d o ret -> Bool
forall d o ret.
(Ord d, Ord o, Ord ret) =>
MemInfo d o ret -> MemInfo d o ret -> Ordering
forall d o ret.
(Ord d, Ord o, Ord ret) =>
MemInfo d o ret -> MemInfo d o ret -> MemInfo d o ret
$ccompare :: forall d o ret.
(Ord d, Ord o, Ord ret) =>
MemInfo d o ret -> MemInfo d o ret -> Ordering
compare :: MemInfo d o ret -> MemInfo d o ret -> Ordering
$c< :: forall d o ret.
(Ord d, Ord o, Ord ret) =>
MemInfo d o ret -> MemInfo d o ret -> Bool
< :: MemInfo d o ret -> MemInfo d o ret -> Bool
$c<= :: forall d o ret.
(Ord d, Ord o, Ord ret) =>
MemInfo d o ret -> MemInfo d o ret -> Bool
<= :: MemInfo d o ret -> MemInfo d o ret -> Bool
$c> :: forall d o ret.
(Ord d, Ord o, Ord ret) =>
MemInfo d o ret -> MemInfo d o ret -> Bool
> :: MemInfo d o ret -> MemInfo d o ret -> Bool
$c>= :: forall d o ret.
(Ord d, Ord o, Ord ret) =>
MemInfo d o ret -> MemInfo d o ret -> Bool
>= :: MemInfo d o ret -> MemInfo d o ret -> Bool
$cmax :: forall d o ret.
(Ord d, Ord o, Ord ret) =>
MemInfo d o ret -> MemInfo d o ret -> MemInfo d o ret
max :: MemInfo d o ret -> MemInfo d o ret -> MemInfo d o ret
$cmin :: forall d o ret.
(Ord d, Ord o, Ord ret) =>
MemInfo d o ret -> MemInfo d o ret -> MemInfo d o ret
min :: MemInfo d o ret -> MemInfo d o ret -> MemInfo d o ret
Ord) --- XXX Ord?

type MemBound o = MemInfo SubExp o MemBind

instance (FixExt ret) => ExtTyped (MemInfo ExtSize NoMode ret) where
  extTypeOf :: MemInfo (Ext SubExp) NoMode ret -> ExtType
extTypeOf (MemPrim PrimType
pt) = PrimType -> ExtType
forall shape o. PrimType -> TypeBase shape o
Prim PrimType
pt
  extTypeOf (MemMem Space
space) = Space -> ExtType
forall shape o. Space -> TypeBase shape o
Mem Space
space
  extTypeOf (MemArray PrimType
pt ShapeBase (Ext SubExp)
shape NoMode
o ret
_) = PrimType -> ShapeBase (Ext SubExp) -> NoMode -> ExtType
forall shape o. PrimType -> shape -> o -> TypeBase shape o
Array PrimType
pt ShapeBase (Ext SubExp)
shape NoMode
o
  extTypeOf (MemAcc VName
acc ShapeBase SubExp
ispace [Type]
ts) = VName -> ShapeBase SubExp -> [Type] -> ExtType
forall shape o.
VName -> ShapeBase SubExp -> [Type] -> TypeBase shape o
Acc VName
acc ShapeBase SubExp
ispace [Type]
ts

instance (FixExt ret) => FixExt (MemInfo ExtSize o ret) where
  fixExt :: Int
-> SubExp
-> MemInfo (Ext SubExp) o ret
-> MemInfo (Ext SubExp) o ret
fixExt Int
_ SubExp
_ (MemPrim PrimType
pt) = PrimType -> MemInfo (Ext SubExp) o ret
forall d o ret. PrimType -> MemInfo d o ret
MemPrim PrimType
pt
  fixExt Int
_ SubExp
_ (MemMem Space
space) = Space -> MemInfo (Ext SubExp) o ret
forall d o ret. Space -> MemInfo d o ret
MemMem Space
space
  fixExt Int
_ SubExp
_ (MemAcc VName
acc ShapeBase SubExp
ispace [Type]
ts) = VName -> ShapeBase SubExp -> [Type] -> MemInfo (Ext SubExp) o ret
forall d o ret.
VName -> ShapeBase SubExp -> [Type] -> MemInfo d o ret
MemAcc VName
acc ShapeBase SubExp
ispace [Type]
ts
  fixExt Int
i SubExp
se (MemArray PrimType
pt ShapeBase (Ext SubExp)
shape o
o ret
ret) =
    PrimType
-> ShapeBase (Ext SubExp) -> o -> ret -> MemInfo (Ext SubExp) o ret
forall d o ret.
PrimType -> ShapeBase d -> o -> ret -> MemInfo d o ret
MemArray PrimType
pt (Int -> SubExp -> ShapeBase (Ext SubExp) -> ShapeBase (Ext SubExp)
forall t. FixExt t => Int -> SubExp -> t -> t
fixExt Int
i SubExp
se ShapeBase (Ext SubExp)
shape) o
o (Int -> SubExp -> ret -> ret
forall t. FixExt t => Int -> SubExp -> t -> t
fixExt Int
i SubExp
se ret
ret)

  mapExt :: (Int -> Int)
-> MemInfo (Ext SubExp) o ret -> MemInfo (Ext SubExp) o ret
mapExt Int -> Int
_ (MemPrim PrimType
pt) = PrimType -> MemInfo (Ext SubExp) o ret
forall d o ret. PrimType -> MemInfo d o ret
MemPrim PrimType
pt
  mapExt Int -> Int
_ (MemMem Space
space) = Space -> MemInfo (Ext SubExp) o ret
forall d o ret. Space -> MemInfo d o ret
MemMem Space
space
  mapExt Int -> Int
_ (MemAcc VName
acc ShapeBase SubExp
ispace [Type]
ts) = VName -> ShapeBase SubExp -> [Type] -> MemInfo (Ext SubExp) o ret
forall d o ret.
VName -> ShapeBase SubExp -> [Type] -> MemInfo d o ret
MemAcc VName
acc ShapeBase SubExp
ispace [Type]
ts
  mapExt Int -> Int
f (MemArray PrimType
pt ShapeBase (Ext SubExp)
shape o
o ret
ret) =
    PrimType
-> ShapeBase (Ext SubExp) -> o -> ret -> MemInfo (Ext SubExp) o ret
forall d o ret.
PrimType -> ShapeBase d -> o -> ret -> MemInfo d o ret
MemArray PrimType
pt ((Int -> Int) -> ShapeBase (Ext SubExp) -> ShapeBase (Ext SubExp)
forall t. FixExt t => (Int -> Int) -> t -> t
mapExt Int -> Int
f ShapeBase (Ext SubExp)
shape) o
o ((Int -> Int) -> ret -> ret
forall t. FixExt t => (Int -> Int) -> t -> t
mapExt Int -> Int
f ret
ret)

instance Typed (MemInfo SubExp Diet ret) where
  typeOf :: MemInfo SubExp Diet ret -> Type
typeOf = TypeBase (ShapeBase SubExp) Diet -> Type
forall shape o. TypeBase shape o -> TypeBase shape NoMode
fromDecl (TypeBase (ShapeBase SubExp) Diet -> Type)
-> (MemInfo SubExp Diet ret -> TypeBase (ShapeBase SubExp) Diet)
-> MemInfo SubExp Diet ret
-> Type
forall b c a. (b -> c) -> (a -> b) -> a -> c
forall {k} (cat :: k -> k -> *) (b :: k) (c :: k) (a :: k).
Category cat =>
cat b c -> cat a b -> cat a c
. MemInfo SubExp Diet ret -> TypeBase (ShapeBase SubExp) Diet
forall t. DeclTyped t => t -> TypeBase (ShapeBase SubExp) Diet
declTypeOf

instance Typed (MemInfo SubExp NoMode ret) where
  typeOf :: MemInfo SubExp NoMode ret -> Type
typeOf (MemPrim PrimType
pt) = PrimType -> Type
forall shape o. PrimType -> TypeBase shape o
Prim PrimType
pt
  typeOf (MemMem Space
space) = Space -> Type
forall shape o. Space -> TypeBase shape o
Mem Space
space
  typeOf (MemArray PrimType
bt ShapeBase SubExp
shape NoMode
o ret
_) = PrimType -> ShapeBase SubExp -> NoMode -> Type
forall shape o. PrimType -> shape -> o -> TypeBase shape o
Array PrimType
bt ShapeBase SubExp
shape NoMode
o
  typeOf (MemAcc VName
acc ShapeBase SubExp
ispace [Type]
ts) = VName -> ShapeBase SubExp -> [Type] -> Type
forall shape o.
VName -> ShapeBase SubExp -> [Type] -> TypeBase shape o
Acc VName
acc ShapeBase SubExp
ispace [Type]
ts

instance DeclTyped (MemInfo SubExp Diet ret) where
  declTypeOf :: MemInfo SubExp Diet ret -> TypeBase (ShapeBase SubExp) Diet
declTypeOf (MemPrim PrimType
bt) = PrimType -> TypeBase (ShapeBase SubExp) Diet
forall shape o. PrimType -> TypeBase shape o
Prim PrimType
bt
  declTypeOf (MemMem Space
space) = Space -> TypeBase (ShapeBase SubExp) Diet
forall shape o. Space -> TypeBase shape o
Mem Space
space
  declTypeOf (MemArray PrimType
bt ShapeBase SubExp
shape Diet
o ret
_) = PrimType
-> ShapeBase SubExp -> Diet -> TypeBase (ShapeBase SubExp) Diet
forall shape o. PrimType -> shape -> o -> TypeBase shape o
Array PrimType
bt ShapeBase SubExp
shape Diet
o
  declTypeOf (MemAcc VName
acc ShapeBase SubExp
ispace [Type]
ts) = VName
-> ShapeBase SubExp -> [Type] -> TypeBase (ShapeBase SubExp) Diet
forall shape o.
VName -> ShapeBase SubExp -> [Type] -> TypeBase shape o
Acc VName
acc ShapeBase SubExp
ispace [Type]
ts

instance (FreeIn d, FreeIn ret) => FreeIn (MemInfo d o ret) where
  freeIn' :: MemInfo d o ret -> FV
freeIn' (MemArray PrimType
_ ShapeBase d
shape o
_ ret
ret) = ShapeBase d -> FV
forall a. FreeIn a => a -> FV
freeIn' ShapeBase d
shape FV -> FV -> FV
forall a. Semigroup a => a -> a -> a
<> ret -> FV
forall a. FreeIn a => a -> FV
freeIn' ret
ret
  freeIn' (MemMem Space
s) = Space -> FV
forall a. FreeIn a => a -> FV
freeIn' Space
s
  freeIn' MemPrim {} = FV
forall a. Monoid a => a
mempty
  freeIn' (MemAcc VName
acc ShapeBase SubExp
ispace [Type]
ts) = (VName, ShapeBase SubExp, [Type]) -> FV
forall a. FreeIn a => a -> FV
freeIn' (VName
acc, ShapeBase SubExp
ispace, [Type]
ts)

instance (Substitute d, Substitute ret) => Substitute (MemInfo d o ret) where
  substituteNames :: Map VName VName -> MemInfo d o ret -> MemInfo d o ret
substituteNames Map VName VName
subst (MemArray PrimType
bt ShapeBase d
shape o
o ret
ret) =
    PrimType -> ShapeBase d -> o -> ret -> MemInfo d o ret
forall d o ret.
PrimType -> ShapeBase d -> o -> ret -> MemInfo d o ret
MemArray
      PrimType
bt
      (Map VName VName -> ShapeBase d -> ShapeBase d
forall a. Substitute a => Map VName VName -> a -> a
substituteNames Map VName VName
subst ShapeBase d
shape)
      o
o
      (Map VName VName -> ret -> ret
forall a. Substitute a => Map VName VName -> a -> a
substituteNames Map VName VName
subst ret
ret)
  substituteNames Map VName VName
substs (MemAcc VName
acc ShapeBase SubExp
ispace [Type]
ts) =
    VName -> ShapeBase SubExp -> [Type] -> MemInfo d o ret
forall d o ret.
VName -> ShapeBase SubExp -> [Type] -> MemInfo d o ret
MemAcc
      (Map VName VName -> VName -> VName
forall a. Substitute a => Map VName VName -> a -> a
substituteNames Map VName VName
substs VName
acc)
      (Map VName VName -> ShapeBase SubExp -> ShapeBase SubExp
forall a. Substitute a => Map VName VName -> a -> a
substituteNames Map VName VName
substs ShapeBase SubExp
ispace)
      (Map VName VName -> [Type] -> [Type]
forall a. Substitute a => Map VName VName -> a -> a
substituteNames Map VName VName
substs [Type]
ts)
  substituteNames Map VName VName
_ (MemMem Space
space) =
    Space -> MemInfo d o ret
forall d o ret. Space -> MemInfo d o ret
MemMem Space
space
  substituteNames Map VName VName
_ (MemPrim PrimType
bt) =
    PrimType -> MemInfo d o ret
forall d o ret. PrimType -> MemInfo d o ret
MemPrim PrimType
bt

instance (Substitute d, Substitute ret) => Rename (MemInfo d o ret) where
  rename :: MemInfo d o ret -> RenameM (MemInfo d o ret)
rename = MemInfo d o ret -> RenameM (MemInfo d o ret)
forall a. Substitute a => a -> RenameM a
substituteRename

simplifyLMAD ::
  (Engine.SimplifiableRep rep) =>
  LMAD ->
  Engine.SimpleM rep LMAD
simplifyLMAD :: forall rep. SimplifiableRep rep => LMAD -> SimpleM rep LMAD
simplifyLMAD = (TPrimExp Int64 VName -> SimpleM rep (TPrimExp Int64 VName))
-> LMAD -> SimpleM rep LMAD
forall (t :: * -> *) (f :: * -> *) a b.
(Traversable t, Applicative f) =>
(a -> f b) -> t a -> f (t b)
forall (f :: * -> *) a b.
Applicative f =>
(a -> f b) -> LMAD a -> f (LMAD b)
traverse ((TPrimExp Int64 VName -> SimpleM rep (TPrimExp Int64 VName))
 -> LMAD -> SimpleM rep LMAD)
-> (TPrimExp Int64 VName -> SimpleM rep (TPrimExp Int64 VName))
-> LMAD
-> SimpleM rep LMAD
forall a b. (a -> b) -> a -> b
$ (PrimExp VName -> TPrimExp Int64 VName)
-> SimpleM rep (PrimExp VName)
-> SimpleM rep (TPrimExp Int64 VName)
forall a b. (a -> b) -> SimpleM rep a -> SimpleM rep b
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap PrimExp VName -> TPrimExp Int64 VName
forall v. PrimExp v -> TPrimExp Int64 v
isInt64 (SimpleM rep (PrimExp VName) -> SimpleM rep (TPrimExp Int64 VName))
-> (TPrimExp Int64 VName -> SimpleM rep (PrimExp VName))
-> TPrimExp Int64 VName
-> SimpleM rep (TPrimExp Int64 VName)
forall b c a. (b -> c) -> (a -> b) -> a -> c
forall {k} (cat :: k -> k -> *) (b :: k) (c :: k) (a :: k).
Category cat =>
cat b c -> cat a b -> cat a c
. PrimExp VName -> SimpleM rep (PrimExp VName)
forall rep.
SimplifiableRep rep =>
PrimExp VName -> SimpleM rep (PrimExp VName)
simplifyPrimExp (PrimExp VName -> SimpleM rep (PrimExp VName))
-> (TPrimExp Int64 VName -> PrimExp VName)
-> TPrimExp Int64 VName
-> SimpleM rep (PrimExp VName)
forall b c a. (b -> c) -> (a -> b) -> a -> c
forall {k} (cat :: k -> k -> *) (b :: k) (c :: k) (a :: k).
Category cat =>
cat b c -> cat a b -> cat a c
. TPrimExp Int64 VName -> PrimExp VName
forall {k} (t :: k) v. TPrimExp t v -> PrimExp v
untyped

simplifyExtLMAD ::
  (Engine.SimplifiableRep rep) =>
  ExtLMAD ->
  Engine.SimpleM rep ExtLMAD
simplifyExtLMAD :: forall rep. SimplifiableRep rep => ExtLMAD -> SimpleM rep ExtLMAD
simplifyExtLMAD = (TPrimExp Int64 (Ext VName)
 -> SimpleM rep (TPrimExp Int64 (Ext VName)))
-> ExtLMAD -> SimpleM rep ExtLMAD
forall (t :: * -> *) (f :: * -> *) a b.
(Traversable t, Applicative f) =>
(a -> f b) -> t a -> f (t b)
forall (f :: * -> *) a b.
Applicative f =>
(a -> f b) -> LMAD a -> f (LMAD b)
traverse ((TPrimExp Int64 (Ext VName)
  -> SimpleM rep (TPrimExp Int64 (Ext VName)))
 -> ExtLMAD -> SimpleM rep ExtLMAD)
-> (TPrimExp Int64 (Ext VName)
    -> SimpleM rep (TPrimExp Int64 (Ext VName)))
-> ExtLMAD
-> SimpleM rep ExtLMAD
forall a b. (a -> b) -> a -> b
$ (PrimExp (Ext VName) -> TPrimExp Int64 (Ext VName))
-> SimpleM rep (PrimExp (Ext VName))
-> SimpleM rep (TPrimExp Int64 (Ext VName))
forall a b. (a -> b) -> SimpleM rep a -> SimpleM rep b
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap PrimExp (Ext VName) -> TPrimExp Int64 (Ext VName)
forall v. PrimExp v -> TPrimExp Int64 v
isInt64 (SimpleM rep (PrimExp (Ext VName))
 -> SimpleM rep (TPrimExp Int64 (Ext VName)))
-> (TPrimExp Int64 (Ext VName)
    -> SimpleM rep (PrimExp (Ext VName)))
-> TPrimExp Int64 (Ext VName)
-> SimpleM rep (TPrimExp Int64 (Ext VName))
forall b c a. (b -> c) -> (a -> b) -> a -> c
forall {k} (cat :: k -> k -> *) (b :: k) (c :: k) (a :: k).
Category cat =>
cat b c -> cat a b -> cat a c
. PrimExp (Ext VName) -> SimpleM rep (PrimExp (Ext VName))
forall rep.
SimplifiableRep rep =>
PrimExp (Ext VName) -> SimpleM rep (PrimExp (Ext VName))
simplifyExtPrimExp (PrimExp (Ext VName) -> SimpleM rep (PrimExp (Ext VName)))
-> (TPrimExp Int64 (Ext VName) -> PrimExp (Ext VName))
-> TPrimExp Int64 (Ext VName)
-> SimpleM rep (PrimExp (Ext VName))
forall b c a. (b -> c) -> (a -> b) -> a -> c
forall {k} (cat :: k -> k -> *) (b :: k) (c :: k) (a :: k).
Category cat =>
cat b c -> cat a b -> cat a c
. TPrimExp Int64 (Ext VName) -> PrimExp (Ext VName)
forall {k} (t :: k) v. TPrimExp t v -> PrimExp v
untyped

isStaticLMAD :: ExtLMAD -> Maybe LMAD
isStaticLMAD :: ExtLMAD -> Maybe LMAD
isStaticLMAD = (TPrimExp Int64 (Ext VName) -> Maybe (TPrimExp Int64 VName))
-> ExtLMAD -> Maybe LMAD
forall (t :: * -> *) (f :: * -> *) a b.
(Traversable t, Applicative f) =>
(a -> f b) -> t a -> f (t b)
forall (f :: * -> *) a b.
Applicative f =>
(a -> f b) -> LMAD a -> f (LMAD b)
traverse ((TPrimExp Int64 (Ext VName) -> Maybe (TPrimExp Int64 VName))
 -> ExtLMAD -> Maybe LMAD)
-> (TPrimExp Int64 (Ext VName) -> Maybe (TPrimExp Int64 VName))
-> ExtLMAD
-> Maybe LMAD
forall a b. (a -> b) -> a -> b
$ (Ext VName -> Maybe VName)
-> TPrimExp Int64 (Ext VName) -> Maybe (TPrimExp Int64 VName)
forall (t :: * -> *) (f :: * -> *) a b.
(Traversable t, Applicative f) =>
(a -> f b) -> t a -> f (t b)
forall (f :: * -> *) a b.
Applicative f =>
(a -> f b) -> TPrimExp Int64 a -> f (TPrimExp Int64 b)
traverse Ext VName -> Maybe VName
forall {a}. Ext a -> Maybe a
inst
  where
    inst :: Ext a -> Maybe a
inst Ext {} = Maybe a
forall a. Maybe a
Nothing
    inst (Free a
x) = a -> Maybe a
forall a. a -> Maybe a
Just a
x

instance
  (Engine.Simplifiable d, Engine.Simplifiable ret) =>
  Engine.Simplifiable (MemInfo d o ret)
  where
  simplify :: forall rep.
SimplifiableRep rep =>
MemInfo d o ret -> SimpleM rep (MemInfo d o ret)
simplify (MemPrim PrimType
bt) =
    MemInfo d o ret -> SimpleM rep (MemInfo d o ret)
forall a. a -> SimpleM rep a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (MemInfo d o ret -> SimpleM rep (MemInfo d o ret))
-> MemInfo d o ret -> SimpleM rep (MemInfo d o ret)
forall a b. (a -> b) -> a -> b
$ PrimType -> MemInfo d o ret
forall d o ret. PrimType -> MemInfo d o ret
MemPrim PrimType
bt
  simplify (MemMem Space
space) =
    MemInfo d o ret -> SimpleM rep (MemInfo d o ret)
forall a. a -> SimpleM rep a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (MemInfo d o ret -> SimpleM rep (MemInfo d o ret))
-> MemInfo d o ret -> SimpleM rep (MemInfo d o ret)
forall a b. (a -> b) -> a -> b
$ Space -> MemInfo d o ret
forall d o ret. Space -> MemInfo d o ret
MemMem Space
space
  simplify (MemArray PrimType
bt ShapeBase d
shape o
o ret
ret) =
    PrimType -> ShapeBase d -> o -> ret -> MemInfo d o ret
forall d o ret.
PrimType -> ShapeBase d -> o -> ret -> MemInfo d o ret
MemArray PrimType
bt (ShapeBase d -> o -> ret -> MemInfo d o ret)
-> SimpleM rep (ShapeBase d)
-> SimpleM rep (o -> ret -> MemInfo d o ret)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> ShapeBase d -> SimpleM rep (ShapeBase d)
forall rep.
SimplifiableRep rep =>
ShapeBase d -> SimpleM rep (ShapeBase d)
forall e rep.
(Simplifiable e, SimplifiableRep rep) =>
e -> SimpleM rep e
Engine.simplify ShapeBase d
shape SimpleM rep (o -> ret -> MemInfo d o ret)
-> SimpleM rep o -> SimpleM rep (ret -> MemInfo d o ret)
forall a b. SimpleM rep (a -> b) -> SimpleM rep a -> SimpleM rep b
forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> o -> SimpleM rep o
forall a. a -> SimpleM rep a
forall (f :: * -> *) a. Applicative f => a -> f a
pure o
o SimpleM rep (ret -> MemInfo d o ret)
-> SimpleM rep ret -> SimpleM rep (MemInfo d o ret)
forall a b. SimpleM rep (a -> b) -> SimpleM rep a -> SimpleM rep b
forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> ret -> SimpleM rep ret
forall rep. SimplifiableRep rep => ret -> SimpleM rep ret
forall e rep.
(Simplifiable e, SimplifiableRep rep) =>
e -> SimpleM rep e
Engine.simplify ret
ret
  simplify (MemAcc VName
acc ShapeBase SubExp
ispace [Type]
ts) =
    VName -> ShapeBase SubExp -> [Type] -> MemInfo d o ret
forall d o ret.
VName -> ShapeBase SubExp -> [Type] -> MemInfo d o ret
MemAcc (VName -> ShapeBase SubExp -> [Type] -> MemInfo d o ret)
-> SimpleM rep VName
-> SimpleM rep (ShapeBase SubExp -> [Type] -> MemInfo d o ret)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> VName -> SimpleM rep VName
forall rep. SimplifiableRep rep => VName -> SimpleM rep VName
forall e rep.
(Simplifiable e, SimplifiableRep rep) =>
e -> SimpleM rep e
Engine.simplify VName
acc SimpleM rep (ShapeBase SubExp -> [Type] -> MemInfo d o ret)
-> SimpleM rep (ShapeBase SubExp)
-> SimpleM rep ([Type] -> MemInfo d o ret)
forall a b. SimpleM rep (a -> b) -> SimpleM rep a -> SimpleM rep b
forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> ShapeBase SubExp -> SimpleM rep (ShapeBase SubExp)
forall rep.
SimplifiableRep rep =>
ShapeBase SubExp -> SimpleM rep (ShapeBase SubExp)
forall e rep.
(Simplifiable e, SimplifiableRep rep) =>
e -> SimpleM rep e
Engine.simplify ShapeBase SubExp
ispace SimpleM rep ([Type] -> MemInfo d o ret)
-> SimpleM rep [Type] -> SimpleM rep (MemInfo d o ret)
forall a b. SimpleM rep (a -> b) -> SimpleM rep a -> SimpleM rep b
forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> [Type] -> SimpleM rep [Type]
forall rep. SimplifiableRep rep => [Type] -> SimpleM rep [Type]
forall e rep.
(Simplifiable e, SimplifiableRep rep) =>
e -> SimpleM rep e
Engine.simplify [Type]
ts

instance
  ( PP.Pretty (ShapeBase d),
    PP.Pretty (TypeBase (ShapeBase d) o),
    PP.Pretty d,
    PP.Pretty o,
    PP.Pretty ret
  ) =>
  PP.Pretty (MemInfo d o ret)
  where
  pretty :: forall ann. MemInfo d o ret -> Doc ann
pretty (MemPrim PrimType
bt) = PrimType -> Doc ann
forall a ann. Pretty a => a -> Doc ann
forall ann. PrimType -> Doc ann
PP.pretty PrimType
bt
  pretty (MemMem Space
DefaultSpace) = Doc ann
"mem"
  pretty (MemMem Space
s) = Doc ann
"mem" Doc ann -> Doc ann -> Doc ann
forall a. Semigroup a => a -> a -> a
<> Space -> Doc ann
forall a ann. Pretty a => a -> Doc ann
forall ann. Space -> Doc ann
PP.pretty Space
s
  pretty (MemArray PrimType
bt ShapeBase d
shape o
o ret
ret) =
    TypeBase (ShapeBase d) o -> Doc ann
forall a ann. Pretty a => a -> Doc ann
forall ann. TypeBase (ShapeBase d) o -> Doc ann
PP.pretty (PrimType -> ShapeBase d -> o -> TypeBase (ShapeBase d) o
forall shape o. PrimType -> shape -> o -> TypeBase shape o
Array PrimType
bt ShapeBase d
shape o
o) Doc ann -> Doc ann -> Doc ann
forall ann. Doc ann -> Doc ann -> Doc ann
<+> Doc ann
"@" Doc ann -> Doc ann -> Doc ann
forall ann. Doc ann -> Doc ann -> Doc ann
<+> ret -> Doc ann
forall ann. ret -> Doc ann
forall a ann. Pretty a => a -> Doc ann
PP.pretty ret
ret
  pretty (MemAcc VName
acc ShapeBase SubExp
ispace [Type]
ts) =
    Type -> Doc ann
forall a ann. Pretty a => a -> Doc ann
forall ann. Type -> Doc ann
PP.pretty (VName -> ShapeBase SubExp -> [Type] -> Type
forall shape o.
VName -> ShapeBase SubExp -> [Type] -> TypeBase shape o
Acc VName
acc ShapeBase SubExp
ispace [Type]
ts :: Type)

-- | Memory information for an array bound somewhere in the program.
data MemBind
  = -- | Located in this memory block with this index
    -- function.
    ArrayIn VName LMAD
  deriving (Int -> MemBind -> ShowS
[MemBind] -> ShowS
MemBind -> String
(Int -> MemBind -> ShowS)
-> (MemBind -> String) -> ([MemBind] -> ShowS) -> Show MemBind
forall a.
(Int -> a -> ShowS) -> (a -> String) -> ([a] -> ShowS) -> Show a
$cshowsPrec :: Int -> MemBind -> ShowS
showsPrec :: Int -> MemBind -> ShowS
$cshow :: MemBind -> String
show :: MemBind -> String
$cshowList :: [MemBind] -> ShowS
showList :: [MemBind] -> ShowS
Show)

instance Eq MemBind where
  MemBind
_ == :: MemBind -> MemBind -> Bool
== MemBind
_ = Bool
True

instance Ord MemBind where
  MemBind
_ compare :: MemBind -> MemBind -> Ordering
`compare` MemBind
_ = Ordering
EQ

instance Rename MemBind where
  rename :: MemBind -> RenameM MemBind
rename = MemBind -> RenameM MemBind
forall a. Substitute a => a -> RenameM a
substituteRename

instance Substitute MemBind where
  substituteNames :: Map VName VName -> MemBind -> MemBind
substituteNames Map VName VName
substs (ArrayIn VName
ident LMAD
lmad) =
    VName -> LMAD -> MemBind
ArrayIn (Map VName VName -> VName -> VName
forall a. Substitute a => Map VName VName -> a -> a
substituteNames Map VName VName
substs VName
ident) (Map VName VName -> LMAD -> LMAD
forall a. Substitute a => Map VName VName -> a -> a
substituteNames Map VName VName
substs LMAD
lmad)

instance PP.Pretty MemBind where
  pretty :: forall ann. MemBind -> Doc ann
pretty (ArrayIn VName
mem LMAD
lmad) =
    VName -> Doc ann
forall a ann. Pretty a => a -> Doc ann
forall ann. VName -> Doc ann
PP.pretty VName
mem Doc ann -> Doc ann -> Doc ann
forall ann. Doc ann -> Doc ann -> Doc ann
<+> Doc ann
"->" Doc ann -> Doc ann -> Doc ann
forall ann. Doc ann -> Doc ann -> Doc ann
PP.</> LMAD -> Doc ann
forall a ann. Pretty a => a -> Doc ann
forall ann. LMAD -> Doc ann
PP.pretty LMAD
lmad

instance FreeIn MemBind where
  freeIn' :: MemBind -> FV
freeIn' (ArrayIn VName
mem LMAD
lmad) = VName -> FV
forall a. FreeIn a => a -> FV
freeIn' VName
mem FV -> FV -> FV
forall a. Semigroup a => a -> a -> a
<> LMAD -> FV
forall a. FreeIn a => a -> FV
freeIn' LMAD
lmad

-- | A description of the memory properties of an array being returned
-- by an operation.
data MemReturn
  = -- | The array is located in a memory block that is
    -- already in scope.
    ReturnsInBlock VName ExtLMAD
  | -- | The operation returns a new (existential) memory
    -- block.
    ReturnsNewBlock Space Int ExtLMAD
  deriving (Int -> MemReturn -> ShowS
[MemReturn] -> ShowS
MemReturn -> String
(Int -> MemReturn -> ShowS)
-> (MemReturn -> String)
-> ([MemReturn] -> ShowS)
-> Show MemReturn
forall a.
(Int -> a -> ShowS) -> (a -> String) -> ([a] -> ShowS) -> Show a
$cshowsPrec :: Int -> MemReturn -> ShowS
showsPrec :: Int -> MemReturn -> ShowS
$cshow :: MemReturn -> String
show :: MemReturn -> String
$cshowList :: [MemReturn] -> ShowS
showList :: [MemReturn] -> ShowS
Show)

instance Eq MemReturn where
  MemReturn
_ == :: MemReturn -> MemReturn -> Bool
== MemReturn
_ = Bool
True

instance Ord MemReturn where
  MemReturn
_ compare :: MemReturn -> MemReturn -> Ordering
`compare` MemReturn
_ = Ordering
EQ

instance Rename MemReturn where
  rename :: MemReturn -> RenameM MemReturn
rename = MemReturn -> RenameM MemReturn
forall a. Substitute a => a -> RenameM a
substituteRename

instance Substitute MemReturn where
  substituteNames :: Map VName VName -> MemReturn -> MemReturn
substituteNames Map VName VName
substs (ReturnsInBlock VName
ident ExtLMAD
lmad) =
    VName -> ExtLMAD -> MemReturn
ReturnsInBlock (Map VName VName -> VName -> VName
forall a. Substitute a => Map VName VName -> a -> a
substituteNames Map VName VName
substs VName
ident) (Map VName VName -> ExtLMAD -> ExtLMAD
forall a. Substitute a => Map VName VName -> a -> a
substituteNames Map VName VName
substs ExtLMAD
lmad)
  substituteNames Map VName VName
substs (ReturnsNewBlock Space
space Int
i ExtLMAD
lmad) =
    Space -> Int -> ExtLMAD -> MemReturn
ReturnsNewBlock Space
space Int
i (Map VName VName -> ExtLMAD -> ExtLMAD
forall a. Substitute a => Map VName VName -> a -> a
substituteNames Map VName VName
substs ExtLMAD
lmad)

instance FixExt MemReturn where
  fixExt :: Int -> SubExp -> MemReturn -> MemReturn
fixExt Int
i (Var VName
v) (ReturnsNewBlock Space
_ Int
j ExtLMAD
lmad)
    | Int
j Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
== Int
i =
        VName -> ExtLMAD -> MemReturn
ReturnsInBlock VName
v (ExtLMAD -> MemReturn) -> ExtLMAD -> MemReturn
forall a b. (a -> b) -> a -> b
$
          Int -> PrimExp VName -> ExtLMAD -> ExtLMAD
fixExtLMAD
            Int
i
            (PrimType -> SubExp -> PrimExp VName
primExpFromSubExp PrimType
int64 (VName -> SubExp
Var VName
v))
            ExtLMAD
lmad
  fixExt Int
i SubExp
se (ReturnsNewBlock Space
space Int
j ExtLMAD
lmad) =
    Space -> Int -> ExtLMAD -> MemReturn
ReturnsNewBlock
      Space
space
      Int
j'
      (Int -> PrimExp VName -> ExtLMAD -> ExtLMAD
fixExtLMAD Int
i (PrimType -> SubExp -> PrimExp VName
primExpFromSubExp PrimType
int64 SubExp
se) ExtLMAD
lmad)
    where
      j' :: Int
j'
        | Int
i Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
< Int
j = Int
j Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
1
        | Bool
otherwise = Int
j
  fixExt Int
i SubExp
se (ReturnsInBlock VName
mem ExtLMAD
lmad) =
    VName -> ExtLMAD -> MemReturn
ReturnsInBlock VName
mem (Int -> PrimExp VName -> ExtLMAD -> ExtLMAD
fixExtLMAD Int
i (PrimType -> SubExp -> PrimExp VName
primExpFromSubExp PrimType
int64 SubExp
se) ExtLMAD
lmad)

  mapExt :: (Int -> Int) -> MemReturn -> MemReturn
mapExt Int -> Int
f (ReturnsNewBlock Space
space Int
i ExtLMAD
lmad) =
    Space -> Int -> ExtLMAD -> MemReturn
ReturnsNewBlock Space
space (Int -> Int
f Int
i) ExtLMAD
lmad
  mapExt Int -> Int
f (ReturnsInBlock VName
mem ExtLMAD
lmad) =
    VName -> ExtLMAD -> MemReturn
ReturnsInBlock VName
mem ((TPrimExp Int64 (Ext VName) -> TPrimExp Int64 (Ext VName))
-> ExtLMAD -> ExtLMAD
forall a b. (a -> b) -> LMAD a -> LMAD b
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap ((Ext VName -> Ext VName)
-> TPrimExp Int64 (Ext VName) -> TPrimExp Int64 (Ext VName)
forall a b. (a -> b) -> TPrimExp Int64 a -> TPrimExp Int64 b
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap Ext VName -> Ext VName
f') ExtLMAD
lmad)
    where
      f' :: Ext VName -> Ext VName
f' (Ext Int
i) = Int -> Ext VName
forall a. Int -> Ext a
Ext (Int -> Ext VName) -> Int -> Ext VName
forall a b. (a -> b) -> a -> b
$ Int -> Int
f Int
i
      f' Ext VName
v = Ext VName
v

fixExtLMAD :: Int -> PrimExp VName -> ExtLMAD -> ExtLMAD
fixExtLMAD :: Int -> PrimExp VName -> ExtLMAD -> ExtLMAD
fixExtLMAD Int
i PrimExp VName
e = (TPrimExp Int64 (Ext VName) -> TPrimExp Int64 (Ext VName))
-> ExtLMAD -> ExtLMAD
forall a b. (a -> b) -> LMAD a -> LMAD b
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap ((TPrimExp Int64 (Ext VName) -> TPrimExp Int64 (Ext VName))
 -> ExtLMAD -> ExtLMAD)
-> (TPrimExp Int64 (Ext VName) -> TPrimExp Int64 (Ext VName))
-> ExtLMAD
-> ExtLMAD
forall a b. (a -> b) -> a -> b
$ PrimExp (Ext VName) -> TPrimExp Int64 (Ext VName)
forall v. PrimExp v -> TPrimExp Int64 v
isInt64 (PrimExp (Ext VName) -> TPrimExp Int64 (Ext VName))
-> (TPrimExp Int64 (Ext VName) -> PrimExp (Ext VName))
-> TPrimExp Int64 (Ext VName)
-> TPrimExp Int64 (Ext VName)
forall b c a. (b -> c) -> (a -> b) -> a -> c
forall {k} (cat :: k -> k -> *) (b :: k) (c :: k) (a :: k).
Category cat =>
cat b c -> cat a b -> cat a c
. (Ext VName -> PrimType -> PrimExp (Ext VName))
-> PrimExp (Ext VName) -> PrimExp (Ext VName)
forall a b. (a -> PrimType -> PrimExp b) -> PrimExp a -> PrimExp b
replaceInPrimExp Ext VName -> PrimType -> PrimExp (Ext VName)
update (PrimExp (Ext VName) -> PrimExp (Ext VName))
-> (TPrimExp Int64 (Ext VName) -> PrimExp (Ext VName))
-> TPrimExp Int64 (Ext VName)
-> PrimExp (Ext VName)
forall b c a. (b -> c) -> (a -> b) -> a -> c
forall {k} (cat :: k -> k -> *) (b :: k) (c :: k) (a :: k).
Category cat =>
cat b c -> cat a b -> cat a c
. TPrimExp Int64 (Ext VName) -> PrimExp (Ext VName)
forall {k} (t :: k) v. TPrimExp t v -> PrimExp v
untyped
  where
    update :: Ext VName -> PrimType -> PrimExp (Ext VName)
update (Ext Int
j) PrimType
t
      | Int
j Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
> 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 -> Ext VName) -> Int -> Ext VName
forall a b. (a -> b) -> a -> b
$ Int
j Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
1) PrimType
t
      | Int
j Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
== Int
i = (VName -> Ext VName) -> PrimExp VName -> PrimExp (Ext VName)
forall a b. (a -> b) -> PrimExp a -> PrimExp b
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap VName -> Ext VName
forall a. a -> Ext a
Free PrimExp VName
e
      | Bool
otherwise = Ext VName -> PrimType -> PrimExp (Ext VName)
forall v. v -> PrimType -> PrimExp v
LeafExp (Int -> Ext VName
forall a. Int -> Ext a
Ext Int
j) PrimType
t
    update (Free VName
x) PrimType
t = Ext VName -> PrimType -> PrimExp (Ext VName)
forall v. v -> PrimType -> PrimExp v
LeafExp (VName -> Ext VName
forall a. a -> Ext a
Free VName
x) PrimType
t

leafExp :: Int -> TPrimExp Int64 (Ext a)
leafExp :: forall a. Int -> TPrimExp Int64 (Ext a)
leafExp Int
i = PrimExp (Ext a) -> TPrimExp Int64 (Ext a)
forall v. PrimExp v -> TPrimExp Int64 v
isInt64 (PrimExp (Ext a) -> TPrimExp Int64 (Ext a))
-> PrimExp (Ext a) -> TPrimExp Int64 (Ext a)
forall a b. (a -> b) -> a -> b
$ Ext a -> PrimType -> PrimExp (Ext a)
forall v. v -> PrimType -> PrimExp v
LeafExp (Int -> Ext a
forall a. Int -> Ext a
Ext Int
i) PrimType
int64

existentialiseLMAD :: [VName] -> LMAD -> ExtLMAD
existentialiseLMAD :: [VName] -> LMAD -> ExtLMAD
existentialiseLMAD [VName]
ctx = Map (Ext VName) (TPrimExp Int64 (Ext VName)) -> ExtLMAD -> ExtLMAD
forall {k} a (t :: k).
Ord a =>
Map a (TPrimExp t a) -> LMAD (TPrimExp t a) -> LMAD (TPrimExp t a)
LMAD.substitute Map (Ext VName) (TPrimExp Int64 (Ext VName))
ctx' (ExtLMAD -> ExtLMAD) -> (LMAD -> ExtLMAD) -> LMAD -> ExtLMAD
forall b c a. (b -> c) -> (a -> b) -> a -> c
forall {k} (cat :: k -> k -> *) (b :: k) (c :: k) (a :: k).
Category cat =>
cat b c -> cat a b -> cat a c
. (TPrimExp Int64 VName -> TPrimExp Int64 (Ext VName))
-> LMAD -> ExtLMAD
forall a b. (a -> b) -> LMAD a -> LMAD b
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap ((VName -> Ext VName)
-> TPrimExp Int64 VName -> TPrimExp Int64 (Ext VName)
forall a b. (a -> b) -> TPrimExp Int64 a -> TPrimExp Int64 b
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap VName -> Ext VName
forall a. a -> Ext a
Free)
  where
    ctx' :: Map (Ext VName) (TPrimExp Int64 (Ext VName))
ctx' = (Int -> TPrimExp Int64 (Ext VName))
-> Map (Ext VName) Int
-> Map (Ext VName) (TPrimExp Int64 (Ext VName))
forall a b k. (a -> b) -> Map k a -> Map k b
M.map Int -> TPrimExp Int64 (Ext VName)
forall a. Int -> TPrimExp Int64 (Ext a)
leafExp (Map (Ext VName) Int
 -> Map (Ext VName) (TPrimExp Int64 (Ext VName)))
-> Map (Ext VName) Int
-> Map (Ext VName) (TPrimExp Int64 (Ext VName))
forall a b. (a -> b) -> a -> b
$ [(Ext VName, Int)] -> Map (Ext VName) Int
forall k a. Ord k => [(k, a)] -> Map k a
M.fromList ([(Ext VName, Int)] -> Map (Ext VName) Int)
-> [(Ext VName, Int)] -> Map (Ext VName) Int
forall a b. (a -> b) -> a -> b
$ [Ext VName] -> [Int] -> [(Ext VName, Int)]
forall a b. [a] -> [b] -> [(a, b)]
zip ((VName -> Ext VName) -> [VName] -> [Ext VName]
forall a b. (a -> b) -> [a] -> [b]
map VName -> Ext VName
forall a. a -> Ext a
Free [VName]
ctx) [Int
0 ..]

instance PP.Pretty MemReturn where
  pretty :: forall ann. MemReturn -> Doc ann
pretty (ReturnsInBlock VName
v ExtLMAD
lmad) =
    VName -> Doc ann
forall a ann. Pretty a => a -> Doc ann
forall ann. VName -> Doc ann
pretty VName
v Doc ann -> Doc ann -> Doc ann
forall ann. Doc ann -> Doc ann -> Doc ann
<+> Doc ann
"->" Doc ann -> Doc ann -> Doc ann
forall ann. Doc ann -> Doc ann -> Doc ann
PP.</> ExtLMAD -> Doc ann
forall a ann. Pretty a => a -> Doc ann
forall ann. ExtLMAD -> Doc ann
PP.pretty ExtLMAD
lmad
  pretty (ReturnsNewBlock Space
space Int
i ExtLMAD
lmad) =
    Doc ann
"?" Doc ann -> Doc ann -> Doc ann
forall a. Semigroup a => a -> a -> a
<> Int -> Doc ann
forall ann. Int -> Doc ann
forall a ann. Pretty a => a -> Doc ann
pretty Int
i Doc ann -> Doc ann -> Doc ann
forall a. Semigroup a => a -> a -> a
<> Space -> Doc ann
forall a ann. Pretty a => a -> Doc ann
forall ann. Space -> Doc ann
PP.pretty Space
space Doc ann -> Doc ann -> Doc ann
forall ann. Doc ann -> Doc ann -> Doc ann
<+> Doc ann
"->" Doc ann -> Doc ann -> Doc ann
forall ann. Doc ann -> Doc ann -> Doc ann
PP.</> ExtLMAD -> Doc ann
forall a ann. Pretty a => a -> Doc ann
forall ann. ExtLMAD -> Doc ann
PP.pretty ExtLMAD
lmad

instance FreeIn MemReturn where
  freeIn' :: MemReturn -> FV
freeIn' (ReturnsInBlock VName
v ExtLMAD
lmad) = VName -> FV
forall a. FreeIn a => a -> FV
freeIn' VName
v FV -> FV -> FV
forall a. Semigroup a => a -> a -> a
<> ExtLMAD -> FV
forall a. FreeIn a => a -> FV
freeIn' ExtLMAD
lmad
  freeIn' (ReturnsNewBlock Space
space Int
_ ExtLMAD
lmad) = Space -> FV
forall a. FreeIn a => a -> FV
freeIn' Space
space FV -> FV -> FV
forall a. Semigroup a => a -> a -> a
<> ExtLMAD -> FV
forall a. FreeIn a => a -> FV
freeIn' ExtLMAD
lmad

instance Engine.Simplifiable MemReturn where
  simplify :: forall rep.
SimplifiableRep rep =>
MemReturn -> SimpleM rep MemReturn
simplify (ReturnsNewBlock Space
space Int
i ExtLMAD
lmad) =
    Space -> Int -> ExtLMAD -> MemReturn
ReturnsNewBlock Space
space Int
i (ExtLMAD -> MemReturn)
-> SimpleM rep ExtLMAD -> SimpleM rep MemReturn
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> ExtLMAD -> SimpleM rep ExtLMAD
forall rep. SimplifiableRep rep => ExtLMAD -> SimpleM rep ExtLMAD
simplifyExtLMAD ExtLMAD
lmad
  simplify (ReturnsInBlock VName
v ExtLMAD
lmad) =
    VName -> ExtLMAD -> MemReturn
ReturnsInBlock (VName -> ExtLMAD -> MemReturn)
-> SimpleM rep VName -> SimpleM rep (ExtLMAD -> MemReturn)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> VName -> SimpleM rep VName
forall rep. SimplifiableRep rep => VName -> SimpleM rep VName
forall e rep.
(Simplifiable e, SimplifiableRep rep) =>
e -> SimpleM rep e
Engine.simplify VName
v SimpleM rep (ExtLMAD -> MemReturn)
-> SimpleM rep ExtLMAD -> SimpleM rep MemReturn
forall a b. SimpleM rep (a -> b) -> SimpleM rep a -> SimpleM rep b
forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> ExtLMAD -> SimpleM rep ExtLMAD
forall rep. SimplifiableRep rep => ExtLMAD -> SimpleM rep ExtLMAD
simplifyExtLMAD ExtLMAD
lmad

instance Engine.Simplifiable MemBind where
  simplify :: forall rep. SimplifiableRep rep => MemBind -> SimpleM rep MemBind
simplify (ArrayIn VName
mem LMAD
lmad) =
    VName -> LMAD -> MemBind
ArrayIn (VName -> LMAD -> MemBind)
-> SimpleM rep VName -> SimpleM rep (LMAD -> MemBind)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> VName -> SimpleM rep VName
forall rep. SimplifiableRep rep => VName -> SimpleM rep VName
forall e rep.
(Simplifiable e, SimplifiableRep rep) =>
e -> SimpleM rep e
Engine.simplify VName
mem SimpleM rep (LMAD -> MemBind)
-> SimpleM rep LMAD -> SimpleM rep MemBind
forall a b. SimpleM rep (a -> b) -> SimpleM rep a -> SimpleM rep b
forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> LMAD -> SimpleM rep LMAD
forall rep. SimplifiableRep rep => LMAD -> SimpleM rep LMAD
simplifyLMAD LMAD
lmad

instance Engine.Simplifiable [FunReturns] where
  simplify :: forall rep.
SimplifiableRep rep =>
[FunReturns] -> SimpleM rep [FunReturns]
simplify = (FunReturns -> SimpleM rep FunReturns)
-> [FunReturns] -> SimpleM rep [FunReturns]
forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
(a -> m b) -> t a -> m (t b)
forall (m :: * -> *) a b. Monad m => (a -> m b) -> [a] -> m [b]
mapM FunReturns -> SimpleM rep FunReturns
forall rep.
SimplifiableRep rep =>
FunReturns -> SimpleM rep FunReturns
forall e rep.
(Simplifiable e, SimplifiableRep rep) =>
e -> SimpleM rep e
Engine.simplify

-- | The memory return of an expression.  An array is annotated with
-- @Maybe MemReturn@, which can be interpreted as the expression
-- either dictating exactly where the array is located when it is
-- returned (if 'Just'), or able to put it whereever the binding
-- prefers (if 'Nothing').
--
-- This is necessary to capture the difference between an expression
-- that is just an array-typed variable, in which the array being
-- "returned" is located where it already is, and a @copy@ expression,
-- whose entire purpose is to store an existing array in some
-- arbitrary location.  This is a consequence of the design decision
-- never to have implicit memory copies.
type ExpReturns = MemInfo ExtSize NoMode (Maybe MemReturn)

-- | The return of a body, which must always indicate where
-- returned arrays are located.
type BodyReturns = MemInfo ExtSize NoMode MemReturn

-- | The memory return of a function, which must always indicate where
-- returned arrays are located.
type FunReturns = MemInfo ExtSize NoMode MemReturn

maybeReturns :: MemInfo d o r -> MemInfo d o (Maybe r)
maybeReturns :: forall d o r. MemInfo d o r -> MemInfo d o (Maybe r)
maybeReturns (MemArray PrimType
bt ShapeBase d
shape o
o r
ret) =
  PrimType -> ShapeBase d -> o -> Maybe r -> MemInfo d o (Maybe r)
forall d o ret.
PrimType -> ShapeBase d -> o -> ret -> MemInfo d o ret
MemArray PrimType
bt ShapeBase d
shape o
o (Maybe r -> MemInfo d o (Maybe r))
-> Maybe r -> MemInfo d o (Maybe r)
forall a b. (a -> b) -> a -> b
$ r -> Maybe r
forall a. a -> Maybe a
Just r
ret
maybeReturns (MemPrim PrimType
bt) =
  PrimType -> MemInfo d o (Maybe r)
forall d o ret. PrimType -> MemInfo d o ret
MemPrim PrimType
bt
maybeReturns (MemMem Space
space) =
  Space -> MemInfo d o (Maybe r)
forall d o ret. Space -> MemInfo d o ret
MemMem Space
space
maybeReturns (MemAcc VName
acc ShapeBase SubExp
ispace [Type]
ts) =
  VName -> ShapeBase SubExp -> [Type] -> MemInfo d o (Maybe r)
forall d o ret.
VName -> ShapeBase SubExp -> [Type] -> MemInfo d o ret
MemAcc VName
acc ShapeBase SubExp
ispace [Type]
ts

noModeReturns :: MemInfo d o r -> MemInfo d NoMode r
noModeReturns :: forall d o r. MemInfo d o r -> MemInfo d NoMode r
noModeReturns (MemArray PrimType
bt ShapeBase d
shape o
_ r
r) =
  PrimType -> ShapeBase d -> NoMode -> r -> MemInfo d NoMode r
forall d o ret.
PrimType -> ShapeBase d -> o -> ret -> MemInfo d o ret
MemArray PrimType
bt ShapeBase d
shape NoMode
NoMode r
r
noModeReturns (MemPrim PrimType
bt) =
  PrimType -> MemInfo d NoMode r
forall d o ret. PrimType -> MemInfo d o ret
MemPrim PrimType
bt
noModeReturns (MemMem Space
space) =
  Space -> MemInfo d NoMode r
forall d o ret. Space -> MemInfo d o ret
MemMem Space
space
noModeReturns (MemAcc VName
acc ShapeBase SubExp
ispace [Type]
ts) =
  VName -> ShapeBase SubExp -> [Type] -> MemInfo d NoMode r
forall d o ret.
VName -> ShapeBase SubExp -> [Type] -> MemInfo d o ret
MemAcc VName
acc ShapeBase SubExp
ispace [Type]
ts

funReturnsToExpReturns :: FunReturns -> ExpReturns
funReturnsToExpReturns :: FunReturns -> ExpReturns
funReturnsToExpReturns = ExpReturns -> ExpReturns
forall d o r. MemInfo d o r -> MemInfo d NoMode r
noModeReturns (ExpReturns -> ExpReturns)
-> (FunReturns -> ExpReturns) -> FunReturns -> ExpReturns
forall b c a. (b -> c) -> (a -> b) -> a -> c
forall {k} (cat :: k -> k -> *) (b :: k) (c :: k) (a :: k).
Category cat =>
cat b c -> cat a b -> cat a c
. FunReturns -> ExpReturns
forall d o r. MemInfo d o r -> MemInfo d o (Maybe r)
maybeReturns

bodyReturnsToExpReturns :: BodyReturns -> ExpReturns
bodyReturnsToExpReturns :: FunReturns -> ExpReturns
bodyReturnsToExpReturns = ExpReturns -> ExpReturns
forall d o r. MemInfo d o r -> MemInfo d NoMode r
noModeReturns (ExpReturns -> ExpReturns)
-> (FunReturns -> ExpReturns) -> FunReturns -> ExpReturns
forall b c a. (b -> c) -> (a -> b) -> a -> c
forall {k} (cat :: k -> k -> *) (b :: k) (c :: k) (a :: k).
Category cat =>
cat b c -> cat a b -> cat a c
. FunReturns -> ExpReturns
forall d o r. MemInfo d o r -> MemInfo d o (Maybe r)
maybeReturns

varInfoToExpReturns :: MemInfo SubExp NoMode MemBind -> ExpReturns
varInfoToExpReturns :: LetDecMem -> ExpReturns
varInfoToExpReturns (MemArray PrimType
et ShapeBase SubExp
shape NoMode
o (ArrayIn VName
mem LMAD
lmad)) =
  PrimType
-> ShapeBase (Ext SubExp)
-> NoMode
-> Maybe MemReturn
-> ExpReturns
forall d o ret.
PrimType -> ShapeBase d -> o -> ret -> MemInfo d o ret
MemArray PrimType
et ((SubExp -> Ext SubExp)
-> ShapeBase SubExp -> ShapeBase (Ext SubExp)
forall a b. (a -> b) -> ShapeBase a -> ShapeBase b
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap SubExp -> Ext SubExp
forall a. a -> Ext a
Free ShapeBase SubExp
shape) NoMode
o (Maybe MemReturn -> ExpReturns) -> Maybe MemReturn -> ExpReturns
forall a b. (a -> b) -> a -> b
$
    MemReturn -> Maybe MemReturn
forall a. a -> Maybe a
Just (MemReturn -> Maybe MemReturn) -> MemReturn -> Maybe MemReturn
forall a b. (a -> b) -> a -> b
$
      VName -> ExtLMAD -> MemReturn
ReturnsInBlock VName
mem (ExtLMAD -> MemReturn) -> ExtLMAD -> MemReturn
forall a b. (a -> b) -> a -> b
$
        [VName] -> LMAD -> ExtLMAD
existentialiseLMAD [] LMAD
lmad
varInfoToExpReturns (MemPrim PrimType
pt) = PrimType -> ExpReturns
forall d o ret. PrimType -> MemInfo d o ret
MemPrim PrimType
pt
varInfoToExpReturns (MemAcc VName
acc ShapeBase SubExp
ispace [Type]
ts) = VName -> ShapeBase SubExp -> [Type] -> ExpReturns
forall d o ret.
VName -> ShapeBase SubExp -> [Type] -> MemInfo d o ret
MemAcc VName
acc ShapeBase SubExp
ispace [Type]
ts
varInfoToExpReturns (MemMem Space
space) = Space -> ExpReturns
forall d o ret. Space -> MemInfo d o ret
MemMem Space
space

matchRetTypeToResult ::
  (Mem rep inner, TC.Checkable rep) =>
  [FunReturns] ->
  Result ->
  TC.TypeM rep ()
matchRetTypeToResult :: forall rep (inner :: * -> *).
(Mem rep inner, Checkable rep) =>
[FunReturns] -> Result -> TypeM rep ()
matchRetTypeToResult [FunReturns]
rettype Result
result = do
  scope <- TypeM rep (Scope (Aliases rep))
forall rep (m :: * -> *). HasScope rep m => m (Scope rep)
askScope
  result_ts <- runReaderT (mapM (subExpMemInfo . resSubExp) result) $ removeScopeAliases scope
  matchReturnType rettype (map resSubExp result) result_ts

matchFunctionReturnType ::
  (Mem rep inner, TC.Checkable rep) =>
  [FunReturns] ->
  Result ->
  TC.TypeM rep ()
matchFunctionReturnType :: forall rep (inner :: * -> *).
(Mem rep inner, Checkable rep) =>
[FunReturns] -> Result -> TypeM rep ()
matchFunctionReturnType [FunReturns]
rettype Result
result = do
  [FunReturns] -> Result -> TypeM rep ()
forall rep (inner :: * -> *).
(Mem rep inner, Checkable rep) =>
[FunReturns] -> Result -> TypeM rep ()
matchRetTypeToResult [FunReturns]
rettype Result
result
  (SubExpRes -> TypeM rep ()) -> Result -> TypeM rep ()
forall (t :: * -> *) (m :: * -> *) a b.
(Foldable t, Monad m) =>
(a -> m b) -> t a -> m ()
mapM_ (SubExp -> TypeM rep ()
forall {rep} {inner :: * -> *}.
(BranchType rep ~ FunReturns, LParamInfo rep ~ LetDecMem,
 FParamInfo rep ~ FParamMem, RetType rep ~ FunReturns,
 OpC rep ~ MemOp inner, HasLetDecMem (LetDec rep), ASTRep rep,
 OpReturns inner, RephraseOp inner, Pretty (inner rep),
 Rename (inner rep), Show (inner rep), Ord (inner rep),
 Substitute (inner rep), FreeIn (inner rep)) =>
SubExp -> TypeM rep ()
checkResultSubExp (SubExp -> TypeM rep ())
-> (SubExpRes -> SubExp) -> SubExpRes -> TypeM rep ()
forall b c a. (b -> c) -> (a -> b) -> a -> c
forall {k} (cat :: k -> k -> *) (b :: k) (c :: k) (a :: k).
Category cat =>
cat b c -> cat a b -> cat a c
. SubExpRes -> SubExp
resSubExp) Result
result
  where
    checkResultSubExp :: SubExp -> TypeM rep ()
checkResultSubExp Constant {} =
      () -> TypeM rep ()
forall a. a -> TypeM rep a
forall (f :: * -> *) a. Applicative f => a -> f a
pure ()
    checkResultSubExp (Var VName
v) = do
      dec <- VName -> TypeM rep LetDecMem
forall rep (inner :: * -> *).
Mem rep inner =>
VName -> TypeM rep LetDecMem
varMemInfo VName
v
      case dec of
        MemPrim PrimType
_ -> () -> TypeM rep ()
forall a. a -> TypeM rep a
forall (f :: * -> *) a. Applicative f => a -> f a
pure ()
        MemMem {} -> () -> TypeM rep ()
forall a. a -> TypeM rep a
forall (f :: * -> *) a. Applicative f => a -> f a
pure ()
        MemAcc {} -> () -> TypeM rep ()
forall a. a -> TypeM rep a
forall (f :: * -> *) a. Applicative f => a -> f a
pure ()
        MemArray PrimType
_ ShapeBase SubExp
_ NoMode
_ (ArrayIn VName
_ LMAD
lmad)
          | LMAD -> Bool
forall num. (Eq num, IntegralExp num) => LMAD num -> Bool
LMAD.isDirect LMAD
lmad ->
              () -> TypeM rep ()
forall a. a -> TypeM rep a
forall (f :: * -> *) a. Applicative f => a -> f a
pure ()
          | Bool
otherwise ->
              ErrorCase rep -> TypeM rep ()
forall rep a. ErrorCase rep -> TypeM rep a
TC.bad (ErrorCase rep -> TypeM rep ())
-> (Text -> ErrorCase rep) -> Text -> TypeM rep ()
forall b c a. (b -> c) -> (a -> b) -> a -> c
forall {k} (cat :: k -> k -> *) (b :: k) (c :: k) (a :: k).
Category cat =>
cat b c -> cat a b -> cat a c
. Text -> ErrorCase rep
forall rep. Text -> ErrorCase rep
TC.TypeError (Text -> TypeM rep ()) -> Text -> TypeM rep ()
forall a b. (a -> b) -> a -> b
$
                Text
"Array "
                  Text -> Text -> Text
forall a. Semigroup a => a -> a -> a
<> VName -> Text
forall a. Pretty a => a -> Text
prettyText VName
v
                  Text -> Text -> Text
forall a. Semigroup a => a -> a -> a
<> Text
" returned by function, but has nontrivial index function:\n"
                  Text -> Text -> Text
forall a. Semigroup a => a -> a -> a
<> LMAD -> Text
forall a. Pretty a => a -> Text
prettyText LMAD
lmad

matchLoopResultMem ::
  (Mem rep inner, TC.Checkable rep) =>
  [FParam (Aliases rep)] ->
  Result ->
  TC.TypeM rep ()
matchLoopResultMem :: forall rep (inner :: * -> *).
(Mem rep inner, Checkable rep) =>
[FParam (Aliases rep)] -> Result -> TypeM rep ()
matchLoopResultMem [FParam (Aliases rep)]
params = [FunReturns] -> Result -> TypeM rep ()
forall rep (inner :: * -> *).
(Mem rep inner, Checkable rep) =>
[FunReturns] -> Result -> TypeM rep ()
matchRetTypeToResult [FunReturns]
rettype
  where
    param_names :: [VName]
param_names = (Param FParamMem -> VName) -> [Param FParamMem] -> [VName]
forall a b. (a -> b) -> [a] -> [b]
map Param FParamMem -> VName
forall dec. Param dec -> VName
paramName [FParam (Aliases rep)]
[Param FParamMem]
params

    -- Invent a ReturnType so we can pretend that the loop body is
    -- actually returning from a function.
    rettype :: [FunReturns]
rettype = (Param FParamMem -> FunReturns)
-> [Param FParamMem] -> [FunReturns]
forall a b. (a -> b) -> [a] -> [b]
map (FParamMem -> FunReturns
toRet (FParamMem -> FunReturns)
-> (Param FParamMem -> FParamMem) -> Param FParamMem -> FunReturns
forall b c a. (b -> c) -> (a -> b) -> a -> c
forall {k} (cat :: k -> k -> *) (b :: k) (c :: k) (a :: k).
Category cat =>
cat b c -> cat a b -> cat a c
. Param FParamMem -> FParamMem
forall dec. Param dec -> dec
paramDec) [FParam (Aliases rep)]
[Param FParamMem]
params

    toExtV :: VName -> Ext VName
toExtV VName
v
      | Just Int
i <- VName
v VName -> [VName] -> Maybe Int
forall a. Eq a => a -> [a] -> Maybe Int
`elemIndex` [VName]
param_names = Int -> Ext VName
forall a. Int -> Ext a
Ext Int
i
      | Bool
otherwise = VName -> Ext VName
forall a. a -> Ext a
Free VName
v

    toExtSE :: SubExp -> Ext SubExp
toExtSE (Var VName
v) = VName -> SubExp
Var (VName -> SubExp) -> Ext VName -> Ext SubExp
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> VName -> Ext VName
toExtV VName
v
    toExtSE (Constant PrimValue
v) = SubExp -> Ext SubExp
forall a. a -> Ext a
Free (SubExp -> Ext SubExp) -> SubExp -> Ext SubExp
forall a b. (a -> b) -> a -> b
$ PrimValue -> SubExp
Constant PrimValue
v

    toRet :: FParamMem -> FunReturns
toRet (MemPrim PrimType
t) =
      PrimType -> FunReturns
forall d o ret. PrimType -> MemInfo d o ret
MemPrim PrimType
t
    toRet (MemMem Space
space) =
      Space -> FunReturns
forall d o ret. Space -> MemInfo d o ret
MemMem Space
space
    toRet (MemAcc VName
acc ShapeBase SubExp
ispace [Type]
ts) =
      VName -> ShapeBase SubExp -> [Type] -> FunReturns
forall d o ret.
VName -> ShapeBase SubExp -> [Type] -> MemInfo d o ret
MemAcc VName
acc ShapeBase SubExp
ispace [Type]
ts
    toRet (MemArray PrimType
pt ShapeBase SubExp
shape Diet
_ (ArrayIn VName
mem LMAD
lmad))
      | Just Int
i <- VName
mem VName -> [VName] -> Maybe Int
forall a. Eq a => a -> [a] -> Maybe Int
`elemIndex` [VName]
param_names,
        Param Attrs
_ VName
_ (MemMem Space
space) : [Param FParamMem]
_ <- Int -> [Param FParamMem] -> [Param FParamMem]
forall a. Int -> [a] -> [a]
drop Int
i [FParam (Aliases rep)]
[Param FParamMem]
params =
          PrimType
-> ShapeBase (Ext SubExp) -> NoMode -> MemReturn -> FunReturns
forall d o ret.
PrimType -> ShapeBase d -> o -> ret -> MemInfo d o ret
MemArray PrimType
pt ShapeBase (Ext SubExp)
shape' NoMode
NoMode (MemReturn -> FunReturns) -> MemReturn -> FunReturns
forall a b. (a -> b) -> a -> b
$ Space -> Int -> ExtLMAD -> MemReturn
ReturnsNewBlock Space
space Int
i ExtLMAD
lmad'
      | Bool
otherwise =
          PrimType
-> ShapeBase (Ext SubExp) -> NoMode -> MemReturn -> FunReturns
forall d o ret.
PrimType -> ShapeBase d -> o -> ret -> MemInfo d o ret
MemArray PrimType
pt ShapeBase (Ext SubExp)
shape' NoMode
NoMode (MemReturn -> FunReturns) -> MemReturn -> FunReturns
forall a b. (a -> b) -> a -> b
$ VName -> ExtLMAD -> MemReturn
ReturnsInBlock VName
mem ExtLMAD
lmad'
      where
        shape' :: ShapeBase (Ext SubExp)
shape' = (SubExp -> Ext SubExp)
-> ShapeBase SubExp -> ShapeBase (Ext SubExp)
forall a b. (a -> b) -> ShapeBase a -> ShapeBase b
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap SubExp -> Ext SubExp
toExtSE ShapeBase SubExp
shape
        lmad' :: ExtLMAD
lmad' = [VName] -> LMAD -> ExtLMAD
existentialiseLMAD [VName]
param_names LMAD
lmad

matchBranchReturnType ::
  (Mem rep inner, TC.Checkable rep) =>
  [BodyReturns] ->
  Body (Aliases rep) ->
  TC.TypeM rep ()
matchBranchReturnType :: forall rep (inner :: * -> *).
(Mem rep inner, Checkable rep) =>
[FunReturns] -> Body (Aliases rep) -> TypeM rep ()
matchBranchReturnType [FunReturns]
rettype (Body BodyDec (Aliases rep)
_ Stms (Aliases rep)
stms Result
res) = do
  scope <- TypeM rep (Scope (Aliases rep))
forall rep (m :: * -> *). HasScope rep m => m (Scope rep)
askScope
  ts <- runReaderT (mapM (subExpMemInfo . resSubExp) res) $ removeScopeAliases (scope <> scopeOf stms)
  matchReturnType rettype (map resSubExp res) ts

-- | Helper function for index function unification.
--
-- The first return value maps a VName (wrapped in 'Free') to its Int
-- (wrapped in 'Ext').  In case of duplicates, it is mapped to the
-- *first* Int that occurs.
--
-- The second return value maps each Int (wrapped in an 'Ext') to a
-- 'LeafExp' 'Ext' with the Int at which its associated VName first
-- occurs.
getExtMaps ::
  [(VName, Int)] ->
  ( M.Map (Ext VName) (TPrimExp Int64 (Ext VName)),
    M.Map (Ext VName) (TPrimExp Int64 (Ext VName))
  )
getExtMaps :: [(VName, Int)]
-> (Map (Ext VName) (TPrimExp Int64 (Ext VName)),
    Map (Ext VName) (TPrimExp Int64 (Ext VName)))
getExtMaps [(VName, Int)]
ctx_lst_ids =
  ( (Int -> TPrimExp Int64 (Ext VName))
-> Map (Ext VName) Int
-> Map (Ext VName) (TPrimExp Int64 (Ext VName))
forall a b k. (a -> b) -> Map k a -> Map k b
M.map Int -> TPrimExp Int64 (Ext VName)
forall a. Int -> TPrimExp Int64 (Ext a)
leafExp (Map (Ext VName) Int
 -> Map (Ext VName) (TPrimExp Int64 (Ext VName)))
-> Map (Ext VName) Int
-> Map (Ext VName) (TPrimExp Int64 (Ext VName))
forall a b. (a -> b) -> a -> b
$ (VName -> Ext VName) -> Map VName Int -> Map (Ext VName) Int
forall k2 k1 a. Ord k2 => (k1 -> k2) -> Map k1 a -> Map k2 a
M.mapKeys VName -> Ext VName
forall a. a -> Ext a
Free (Map VName Int -> Map (Ext VName) Int)
-> Map VName Int -> Map (Ext VName) Int
forall a b. (a -> b) -> a -> b
$ (Int -> Int -> Int) -> [(VName, Int)] -> Map VName Int
forall k a. Ord k => (a -> a -> a) -> [(k, a)] -> Map k a
M.fromListWith ((Int -> Int) -> Int -> Int -> Int
forall a b. a -> b -> a
const Int -> Int
forall a. a -> a
forall {k} (cat :: k -> k -> *) (a :: k). Category cat => cat a a
id) [(VName, Int)]
ctx_lst_ids,
    [(Ext VName, TPrimExp Int64 (Ext VName))]
-> Map (Ext VName) (TPrimExp Int64 (Ext VName))
forall k a. Ord k => [(k, a)] -> Map k a
M.fromList ([(Ext VName, TPrimExp Int64 (Ext VName))]
 -> Map (Ext VName) (TPrimExp Int64 (Ext VName)))
-> [(Ext VName, TPrimExp Int64 (Ext VName))]
-> Map (Ext VName) (TPrimExp Int64 (Ext VName))
forall a b. (a -> b) -> a -> b
$
      ((VName, Int) -> Maybe (Ext VName, TPrimExp Int64 (Ext VName)))
-> [(VName, Int)] -> [(Ext VName, TPrimExp Int64 (Ext VName))]
forall a b. (a -> Maybe b) -> [a] -> [b]
mapMaybe
        ( (VName -> Maybe (TPrimExp Int64 (Ext VName)))
-> (Ext VName, VName)
-> Maybe (Ext VName, TPrimExp Int64 (Ext VName))
forall (t :: * -> *) (f :: * -> *) a b.
(Traversable t, Applicative f) =>
(a -> f b) -> t a -> f (t b)
forall (f :: * -> *) a b.
Applicative f =>
(a -> f b) -> (Ext VName, a) -> f (Ext VName, b)
traverse
            ( (Int -> TPrimExp Int64 (Ext VName))
-> Maybe Int -> Maybe (TPrimExp Int64 (Ext VName))
forall a b. (a -> b) -> Maybe a -> Maybe b
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap (\Int
i -> PrimExp (Ext VName) -> TPrimExp Int64 (Ext VName)
forall v. PrimExp v -> TPrimExp Int64 v
isInt64 (PrimExp (Ext VName) -> TPrimExp Int64 (Ext VName))
-> PrimExp (Ext VName) -> TPrimExp Int64 (Ext VName)
forall a b. (a -> b) -> a -> b
$ 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
int64)
                (Maybe Int -> Maybe (TPrimExp Int64 (Ext VName)))
-> (VName -> Maybe Int)
-> VName
-> Maybe (TPrimExp Int64 (Ext VName))
forall b c a. (b -> c) -> (a -> b) -> a -> c
forall {k} (cat :: k -> k -> *) (b :: k) (c :: k) (a :: k).
Category cat =>
cat b c -> cat a b -> cat a c
. (VName -> [(VName, Int)] -> Maybe Int
forall a b. Eq a => a -> [(a, b)] -> Maybe b
`lookup` [(VName, Int)]
ctx_lst_ids)
            )
            ((Ext VName, VName)
 -> Maybe (Ext VName, TPrimExp Int64 (Ext VName)))
-> ((VName, Int) -> (Ext VName, VName))
-> (VName, Int)
-> Maybe (Ext VName, TPrimExp Int64 (Ext VName))
forall b c a. (b -> c) -> (a -> b) -> a -> c
forall {k} (cat :: k -> k -> *) (b :: k) (c :: k) (a :: k).
Category cat =>
cat b c -> cat a b -> cat a c
. (VName -> Ext VName -> (Ext VName, VName))
-> (VName, Ext VName) -> (Ext VName, VName)
forall a b c. (a -> b -> c) -> (a, b) -> c
uncurry ((Ext VName -> VName -> (Ext VName, VName))
-> VName -> Ext VName -> (Ext VName, VName)
forall a b c. (a -> b -> c) -> b -> a -> c
flip (,))
            ((VName, Ext VName) -> (Ext VName, VName))
-> ((VName, Int) -> (VName, Ext VName))
-> (VName, Int)
-> (Ext VName, VName)
forall b c a. (b -> c) -> (a -> b) -> a -> c
forall {k} (cat :: k -> k -> *) (b :: k) (c :: k) (a :: k).
Category cat =>
cat b c -> cat a b -> cat a c
. (Int -> Ext VName) -> (VName, Int) -> (VName, Ext VName)
forall a b. (a -> b) -> (VName, a) -> (VName, b)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap Int -> Ext VName
forall a. Int -> Ext a
Ext
        )
        [(VName, Int)]
ctx_lst_ids
  )

matchReturnType ::
  (PP.Pretty o) =>
  [MemInfo ExtSize o MemReturn] ->
  [SubExp] ->
  [MemInfo SubExp NoMode MemBind] ->
  TC.TypeM rep ()
matchReturnType :: forall o rep.
Pretty o =>
[MemInfo (Ext SubExp) o MemReturn]
-> [SubExp] -> [LetDecMem] -> TypeM rep ()
matchReturnType [MemInfo (Ext SubExp) o MemReturn]
rettype [SubExp]
res [LetDecMem]
ts = do
  let existentialiseLMAD0 :: LMAD -> ExtLMAD
      existentialiseLMAD0 :: LMAD -> ExtLMAD
existentialiseLMAD0 = (TPrimExp Int64 VName -> TPrimExp Int64 (Ext VName))
-> LMAD -> ExtLMAD
forall a b. (a -> b) -> LMAD a -> LMAD b
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap ((TPrimExp Int64 VName -> TPrimExp Int64 (Ext VName))
 -> LMAD -> ExtLMAD)
-> (TPrimExp Int64 VName -> TPrimExp Int64 (Ext VName))
-> LMAD
-> ExtLMAD
forall a b. (a -> b) -> a -> b
$ (VName -> Ext VName)
-> TPrimExp Int64 VName -> TPrimExp Int64 (Ext VName)
forall a b. (a -> b) -> TPrimExp Int64 a -> TPrimExp Int64 b
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap VName -> Ext VName
forall a. a -> Ext a
Free

      fetchCtx :: Int -> ExceptT Text (TypeM rep) (SubExp, LetDecMem)
fetchCtx Int
i = case Int -> [(SubExp, LetDecMem)] -> Maybe (SubExp, LetDecMem)
forall int a. Integral int => int -> [a] -> Maybe a
maybeNth Int
i ([(SubExp, LetDecMem)] -> Maybe (SubExp, LetDecMem))
-> [(SubExp, LetDecMem)] -> Maybe (SubExp, LetDecMem)
forall a b. (a -> b) -> a -> b
$ [SubExp] -> [LetDecMem] -> [(SubExp, LetDecMem)]
forall a b. [a] -> [b] -> [(a, b)]
zip [SubExp]
res [LetDecMem]
ts of
        Maybe (SubExp, LetDecMem)
Nothing ->
          Text -> ExceptT Text (TypeM rep) (SubExp, LetDecMem)
forall a. Text -> ExceptT Text (TypeM rep) a
forall e (m :: * -> *) a. MonadError e m => e -> m a
throwError (Text -> ExceptT Text (TypeM rep) (SubExp, LetDecMem))
-> Text -> ExceptT Text (TypeM rep) (SubExp, LetDecMem)
forall a b. (a -> b) -> a -> b
$ Text
"Cannot find variable #" Text -> Text -> Text
forall a. Semigroup a => a -> a -> a
<> Int -> Text
forall a. Pretty a => a -> Text
prettyText Int
i Text -> Text -> Text
forall a. Semigroup a => a -> a -> a
<> Text
" in results: " Text -> Text -> Text
forall a. Semigroup a => a -> a -> a
<> [SubExp] -> Text
forall a. Pretty a => a -> Text
prettyText [SubExp]
res
        Just (SubExp
se, LetDecMem
t) -> (SubExp, LetDecMem) -> ExceptT Text (TypeM rep) (SubExp, LetDecMem)
forall a. a -> ExceptT Text (TypeM rep) a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (SubExp
se, LetDecMem
t)

      checkReturn :: MemInfo (Ext SubExp) o MemReturn
-> LetDecMem -> ExceptT Text (TypeM rep) ()
checkReturn (MemPrim PrimType
x) (MemPrim PrimType
y)
        | PrimType
x PrimType -> PrimType -> Bool
forall a. Eq a => a -> a -> Bool
== PrimType
y = () -> ExceptT Text (TypeM rep) ()
forall a. a -> ExceptT Text (TypeM rep) a
forall (f :: * -> *) a. Applicative f => a -> f a
pure ()
      checkReturn (MemMem Space
x) (MemMem Space
y)
        | Space
x Space -> Space -> Bool
forall a. Eq a => a -> a -> Bool
== Space
y = () -> ExceptT Text (TypeM rep) ()
forall a. a -> ExceptT Text (TypeM rep) a
forall (f :: * -> *) a. Applicative f => a -> f a
pure ()
      checkReturn (MemAcc VName
xacc ShapeBase SubExp
xispace [Type]
xts) (MemAcc VName
yacc ShapeBase SubExp
yispace [Type]
yts)
        | (VName
xacc, ShapeBase SubExp
xispace, [Type]
xts) (VName, ShapeBase SubExp, [Type])
-> (VName, ShapeBase SubExp, [Type]) -> Bool
forall a. Eq a => a -> a -> Bool
== (VName
yacc, ShapeBase SubExp
yispace, [Type]
yts) =
            () -> ExceptT Text (TypeM rep) ()
forall a. a -> ExceptT Text (TypeM rep) a
forall (f :: * -> *) a. Applicative f => a -> f a
pure ()
      checkReturn
        (MemArray PrimType
x_pt ShapeBase (Ext SubExp)
x_shape o
_ MemReturn
x_ret)
        (MemArray PrimType
y_pt ShapeBase SubExp
y_shape NoMode
_ MemBind
y_ret)
          | PrimType
x_pt PrimType -> PrimType -> Bool
forall a. Eq a => a -> a -> Bool
== PrimType
y_pt,
            ShapeBase (Ext SubExp) -> Int
forall a. ArrayShape a => a -> Int
shapeRank ShapeBase (Ext SubExp)
x_shape Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
== ShapeBase SubExp -> Int
forall a. ArrayShape a => a -> Int
shapeRank ShapeBase SubExp
y_shape = do
              (Ext SubExp -> SubExp -> ExceptT Text (TypeM rep) ())
-> [Ext SubExp] -> [SubExp] -> ExceptT Text (TypeM rep) ()
forall (m :: * -> *) a b c.
Applicative m =>
(a -> b -> m c) -> [a] -> [b] -> m ()
zipWithM_ Ext SubExp -> SubExp -> ExceptT Text (TypeM rep) ()
checkDim (ShapeBase (Ext SubExp) -> [Ext SubExp]
forall d. ShapeBase d -> [d]
shapeDims ShapeBase (Ext SubExp)
x_shape) (ShapeBase SubExp -> [SubExp]
forall d. ShapeBase d -> [d]
shapeDims ShapeBase SubExp
y_shape)
              MemReturn -> MemBind -> ExceptT Text (TypeM rep) ()
checkMemReturn MemReturn
x_ret MemBind
y_ret
      checkReturn MemInfo (Ext SubExp) o MemReturn
x LetDecMem
y =
        Text -> ExceptT Text (TypeM rep) ()
forall a. Text -> ExceptT Text (TypeM rep) a
forall e (m :: * -> *) a. MonadError e m => e -> m a
throwError (Text -> ExceptT Text (TypeM rep) ())
-> Text -> ExceptT Text (TypeM rep) ()
forall a b. (a -> b) -> a -> b
$ [Text] -> Text
T.unwords [Text
"Expected", MemInfo (Ext SubExp) o MemReturn -> Text
forall a. Pretty a => a -> Text
prettyText MemInfo (Ext SubExp) o MemReturn
x, Text
"but got", LetDecMem -> Text
forall a. Pretty a => a -> Text
prettyText LetDecMem
y]

      checkDim :: Ext SubExp -> SubExp -> ExceptT Text (TypeM rep) ()
checkDim (Free SubExp
x) SubExp
y
        | SubExp
x SubExp -> SubExp -> Bool
forall a. Eq a => a -> a -> Bool
== SubExp
y = () -> ExceptT Text (TypeM rep) ()
forall a. a -> ExceptT Text (TypeM rep) a
forall (f :: * -> *) a. Applicative f => a -> f a
pure ()
        | Bool
otherwise =
            Text -> ExceptT Text (TypeM rep) ()
forall a. Text -> ExceptT Text (TypeM rep) a
forall e (m :: * -> *) a. MonadError e m => e -> m a
throwError (Text -> ExceptT Text (TypeM rep) ())
-> Text -> ExceptT Text (TypeM rep) ()
forall a b. (a -> b) -> a -> b
$ [Text] -> Text
T.unwords [Text
"Expected dim", SubExp -> Text
forall a. Pretty a => a -> Text
prettyText SubExp
x, Text
"but got", SubExp -> Text
forall a. Pretty a => a -> Text
prettyText SubExp
y]
      checkDim (Ext Int
i) SubExp
y = do
        (x, _) <- Int -> ExceptT Text (TypeM rep) (SubExp, LetDecMem)
fetchCtx Int
i
        unless (x == y) . throwError . T.unwords $
          ["Expected ext dim", prettyText i, "=>", prettyText x, "but got", prettyText y]

      checkMemReturn :: MemReturn -> MemBind -> ExceptT Text (TypeM rep) ()
checkMemReturn (ReturnsInBlock VName
x_mem ExtLMAD
x_lmad) (ArrayIn VName
y_mem LMAD
y_lmad)
        | VName
x_mem VName -> VName -> Bool
forall a. Eq a => a -> a -> Bool
== VName
y_mem =
            Bool -> ExceptT Text (TypeM rep) () -> ExceptT Text (TypeM rep) ()
forall (f :: * -> *). Applicative f => Bool -> f () -> f ()
unless (ExtLMAD -> ExtLMAD -> Bool
forall num. LMAD num -> LMAD num -> Bool
LMAD.closeEnough ExtLMAD
x_lmad (ExtLMAD -> Bool) -> ExtLMAD -> Bool
forall a b. (a -> b) -> a -> b
$ LMAD -> ExtLMAD
existentialiseLMAD0 LMAD
y_lmad) (ExceptT Text (TypeM rep) () -> ExceptT Text (TypeM rep) ())
-> ExceptT Text (TypeM rep) () -> ExceptT Text (TypeM rep) ()
forall a b. (a -> b) -> a -> b
$
              Text -> ExceptT Text (TypeM rep) ()
forall a. Text -> ExceptT Text (TypeM rep) a
forall e (m :: * -> *) a. MonadError e m => e -> m a
throwError (Text -> ExceptT Text (TypeM rep) ())
-> ([Text] -> Text) -> [Text] -> ExceptT Text (TypeM rep) ()
forall b c a. (b -> c) -> (a -> b) -> a -> c
forall {k} (cat :: k -> k -> *) (b :: k) (c :: k) (a :: k).
Category cat =>
cat b c -> cat a b -> cat a c
. [Text] -> Text
T.unwords ([Text] -> ExceptT Text (TypeM rep) ())
-> [Text] -> ExceptT Text (TypeM rep) ()
forall a b. (a -> b) -> a -> b
$
                [ Text
"Index function unification failed (ReturnsInBlock)",
                  Text
"\nlmad of body result: ",
                  LMAD -> Text
forall a. Pretty a => a -> Text
prettyText LMAD
y_lmad,
                  Text
"\nlmad of return type: ",
                  ExtLMAD -> Text
forall a. Pretty a => a -> Text
prettyText ExtLMAD
x_lmad
                ]
      checkMemReturn
        (ReturnsNewBlock Space
x_space Int
x_ext ExtLMAD
x_lmad)
        (ArrayIn VName
y_mem LMAD
y_lmad) = do
          (x_mem, x_mem_type) <- Int -> ExceptT Text (TypeM rep) (SubExp, LetDecMem)
fetchCtx Int
x_ext
          unless (LMAD.closeEnough x_lmad $ existentialiseLMAD0 y_lmad) $
            throwError . docText $
              "Index function unification failed (ReturnsNewBlock)"
                </> "Lmad of body result:"
                </> indent 2 (pretty y_lmad)
                </> "Lmad of return type:"
                </> indent 2 (pretty x_lmad)
          case x_mem_type of
            MemMem Space
y_space ->
              Bool -> ExceptT Text (TypeM rep) () -> ExceptT Text (TypeM rep) ()
forall (f :: * -> *). Applicative f => Bool -> f () -> f ()
unless (Space
x_space Space -> Space -> Bool
forall a. Eq a => a -> a -> Bool
== Space
y_space) (ExceptT Text (TypeM rep) () -> ExceptT Text (TypeM rep) ())
-> ([Text] -> ExceptT Text (TypeM rep) ())
-> [Text]
-> ExceptT Text (TypeM rep) ()
forall b c a. (b -> c) -> (a -> b) -> a -> c
forall {k} (cat :: k -> k -> *) (b :: k) (c :: k) (a :: k).
Category cat =>
cat b c -> cat a b -> cat a c
. Text -> ExceptT Text (TypeM rep) ()
forall a. Text -> ExceptT Text (TypeM rep) a
forall e (m :: * -> *) a. MonadError e m => e -> m a
throwError (Text -> ExceptT Text (TypeM rep) ())
-> ([Text] -> Text) -> [Text] -> ExceptT Text (TypeM rep) ()
forall b c a. (b -> c) -> (a -> b) -> a -> c
forall {k} (cat :: k -> k -> *) (b :: k) (c :: k) (a :: k).
Category cat =>
cat b c -> cat a b -> cat a c
. [Text] -> Text
T.unwords ([Text] -> ExceptT Text (TypeM rep) ())
-> [Text] -> ExceptT Text (TypeM rep) ()
forall a b. (a -> b) -> a -> b
$
                [ Text
"Expected memory",
                  VName -> Text
forall a. Pretty a => a -> Text
prettyText VName
y_mem,
                  Text
"in space",
                  Space -> Text
forall a. Pretty a => a -> Text
prettyText Space
x_space,
                  Text
"but actually in space",
                  Space -> Text
forall a. Pretty a => a -> Text
prettyText Space
y_space
                ]
            LetDecMem
t ->
              Text -> ExceptT Text (TypeM rep) ()
forall a. Text -> ExceptT Text (TypeM rep) a
forall e (m :: * -> *) a. MonadError e m => e -> m a
throwError (Text -> ExceptT Text (TypeM rep) ())
-> ([Text] -> Text) -> [Text] -> ExceptT Text (TypeM rep) ()
forall b c a. (b -> c) -> (a -> b) -> a -> c
forall {k} (cat :: k -> k -> *) (b :: k) (c :: k) (a :: k).
Category cat =>
cat b c -> cat a b -> cat a c
. [Text] -> Text
T.unwords ([Text] -> ExceptT Text (TypeM rep) ())
-> [Text] -> ExceptT Text (TypeM rep) ()
forall a b. (a -> b) -> a -> b
$
                [Text
"Expected memory", Int -> Text
forall a. Pretty a => a -> Text
prettyText Int
x_ext, Text
"=>", SubExp -> Text
forall a. Pretty a => a -> Text
prettyText SubExp
x_mem, Text
"but but has type", LetDecMem -> Text
forall a. Pretty a => a -> Text
prettyText LetDecMem
t]
      checkMemReturn MemReturn
x MemBind
y =
        Text -> ExceptT Text (TypeM rep) ()
forall a. Text -> ExceptT Text (TypeM rep) a
forall e (m :: * -> *) a. MonadError e m => e -> m a
throwError (Text -> ExceptT Text (TypeM rep) ())
-> (Doc (ZonkAny 3) -> Text)
-> Doc (ZonkAny 3)
-> ExceptT Text (TypeM rep) ()
forall b c a. (b -> c) -> (a -> b) -> a -> c
forall {k} (cat :: k -> k -> *) (b :: k) (c :: k) (a :: k).
Category cat =>
cat b c -> cat a b -> cat a c
. Doc (ZonkAny 3) -> Text
forall a. Doc a -> Text
docText (Doc (ZonkAny 3) -> ExceptT Text (TypeM rep) ())
-> Doc (ZonkAny 3) -> ExceptT Text (TypeM rep) ()
forall a b. (a -> b) -> a -> b
$
          Doc (ZonkAny 3)
"Expected array in"
            Doc (ZonkAny 3) -> Doc (ZonkAny 3) -> Doc (ZonkAny 3)
forall ann. Doc ann -> Doc ann -> Doc ann
</> Int -> Doc (ZonkAny 3) -> Doc (ZonkAny 3)
forall ann. Int -> Doc ann -> Doc ann
indent Int
2 (MemReturn -> Doc (ZonkAny 3)
forall a ann. Pretty a => a -> Doc ann
forall ann. MemReturn -> Doc ann
pretty MemReturn
x)
            Doc (ZonkAny 3) -> Doc (ZonkAny 3) -> Doc (ZonkAny 3)
forall ann. Doc ann -> Doc ann -> Doc ann
</> Doc (ZonkAny 3)
"but array returned in"
            Doc (ZonkAny 3) -> Doc (ZonkAny 3) -> Doc (ZonkAny 3)
forall ann. Doc ann -> Doc ann -> Doc ann
</> Int -> Doc (ZonkAny 3) -> Doc (ZonkAny 3)
forall ann. Int -> Doc ann -> Doc ann
indent Int
2 (MemBind -> Doc (ZonkAny 3)
forall a ann. Pretty a => a -> Doc ann
forall ann. MemBind -> Doc ann
pretty MemBind
y)

      bad :: Text -> TypeM rep ()
bad Text
s =
        ErrorCase rep -> TypeM rep ()
forall rep a. ErrorCase rep -> TypeM rep a
TC.bad (ErrorCase rep -> TypeM rep ())
-> (Doc (ZonkAny 1) -> ErrorCase rep)
-> Doc (ZonkAny 1)
-> TypeM rep ()
forall b c a. (b -> c) -> (a -> b) -> a -> c
forall {k} (cat :: k -> k -> *) (b :: k) (c :: k) (a :: k).
Category cat =>
cat b c -> cat a b -> cat a c
. Text -> ErrorCase rep
forall rep. Text -> ErrorCase rep
TC.TypeError (Text -> ErrorCase rep)
-> (Doc (ZonkAny 1) -> Text) -> Doc (ZonkAny 1) -> ErrorCase rep
forall b c a. (b -> c) -> (a -> b) -> a -> c
forall {k} (cat :: k -> k -> *) (b :: k) (c :: k) (a :: k).
Category cat =>
cat b c -> cat a b -> cat a c
. Doc (ZonkAny 1) -> Text
forall a. Doc a -> Text
docText (Doc (ZonkAny 1) -> TypeM rep ())
-> Doc (ZonkAny 1) -> TypeM rep ()
forall a b. (a -> b) -> a -> b
$
          Doc (ZonkAny 1)
"Return type"
            Doc (ZonkAny 1) -> Doc (ZonkAny 1) -> Doc (ZonkAny 1)
forall ann. Doc ann -> Doc ann -> Doc ann
</> Int -> Doc (ZonkAny 1) -> Doc (ZonkAny 1)
forall ann. Int -> Doc ann -> Doc ann
indent Int
2 ([Doc (ZonkAny 1)] -> Doc (ZonkAny 1)
forall a. [Doc a] -> Doc a
ppTupleLines' ([Doc (ZonkAny 1)] -> Doc (ZonkAny 1))
-> [Doc (ZonkAny 1)] -> Doc (ZonkAny 1)
forall a b. (a -> b) -> a -> b
$ (MemInfo (Ext SubExp) o MemReturn -> Doc (ZonkAny 1))
-> [MemInfo (Ext SubExp) o MemReturn] -> [Doc (ZonkAny 1)]
forall a b. (a -> b) -> [a] -> [b]
map MemInfo (Ext SubExp) o MemReturn -> Doc (ZonkAny 1)
forall a ann. Pretty a => a -> Doc ann
forall ann. MemInfo (Ext SubExp) o MemReturn -> Doc ann
pretty [MemInfo (Ext SubExp) o MemReturn]
rettype)
            Doc (ZonkAny 1) -> Doc (ZonkAny 1) -> Doc (ZonkAny 1)
forall ann. Doc ann -> Doc ann -> Doc ann
</> Doc (ZonkAny 1)
"cannot match returns of results"
            Doc (ZonkAny 1) -> Doc (ZonkAny 1) -> Doc (ZonkAny 1)
forall ann. Doc ann -> Doc ann -> Doc ann
</> Int -> Doc (ZonkAny 1) -> Doc (ZonkAny 1)
forall ann. Int -> Doc ann -> Doc ann
indent Int
2 ([Doc (ZonkAny 1)] -> Doc (ZonkAny 1)
forall a. [Doc a] -> Doc a
ppTupleLines' ([Doc (ZonkAny 1)] -> Doc (ZonkAny 1))
-> [Doc (ZonkAny 1)] -> Doc (ZonkAny 1)
forall a b. (a -> b) -> a -> b
$ (LetDecMem -> Doc (ZonkAny 1)) -> [LetDecMem] -> [Doc (ZonkAny 1)]
forall a b. (a -> b) -> [a] -> [b]
map LetDecMem -> Doc (ZonkAny 1)
forall a ann. Pretty a => a -> Doc ann
forall ann. LetDecMem -> Doc ann
pretty [LetDecMem]
ts)
            Doc (ZonkAny 1) -> Doc (ZonkAny 1) -> Doc (ZonkAny 1)
forall ann. Doc ann -> Doc ann -> Doc ann
</> Text -> Doc (ZonkAny 1)
forall ann. Text -> Doc ann
forall a ann. Pretty a => a -> Doc ann
pretty Text
s

  Bool -> TypeM rep () -> TypeM rep ()
forall (f :: * -> *). Applicative f => Bool -> f () -> f ()
unless ([MemInfo (Ext SubExp) o MemReturn] -> Int
forall a. [a] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length [MemInfo (Ext SubExp) o MemReturn]
rettype Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
== [LetDecMem] -> Int
forall a. [a] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length [LetDecMem]
ts) (TypeM rep () -> TypeM rep ()) -> TypeM rep () -> TypeM rep ()
forall a b. (a -> b) -> a -> b
$
    ErrorCase rep -> TypeM rep ()
forall rep a. ErrorCase rep -> TypeM rep a
TC.bad (ErrorCase rep -> TypeM rep ())
-> (Doc (ZonkAny 4) -> ErrorCase rep)
-> Doc (ZonkAny 4)
-> TypeM rep ()
forall b c a. (b -> c) -> (a -> b) -> a -> c
forall {k} (cat :: k -> k -> *) (b :: k) (c :: k) (a :: k).
Category cat =>
cat b c -> cat a b -> cat a c
. Text -> ErrorCase rep
forall rep. Text -> ErrorCase rep
TC.TypeError (Text -> ErrorCase rep)
-> (Doc (ZonkAny 4) -> Text) -> Doc (ZonkAny 4) -> ErrorCase rep
forall b c a. (b -> c) -> (a -> b) -> a -> c
forall {k} (cat :: k -> k -> *) (b :: k) (c :: k) (a :: k).
Category cat =>
cat b c -> cat a b -> cat a c
. Doc (ZonkAny 4) -> Text
forall a. Doc a -> Text
docText (Doc (ZonkAny 4) -> TypeM rep ())
-> Doc (ZonkAny 4) -> TypeM rep ()
forall a b. (a -> b) -> a -> b
$
      Doc (ZonkAny 4)
"Return type"
        Doc (ZonkAny 4) -> Doc (ZonkAny 4) -> Doc (ZonkAny 4)
forall ann. Doc ann -> Doc ann -> Doc ann
</> Int -> Doc (ZonkAny 4) -> Doc (ZonkAny 4)
forall ann. Int -> Doc ann -> Doc ann
indent Int
2 ([Doc (ZonkAny 4)] -> Doc (ZonkAny 4)
forall a. [Doc a] -> Doc a
ppTupleLines' ([Doc (ZonkAny 4)] -> Doc (ZonkAny 4))
-> [Doc (ZonkAny 4)] -> Doc (ZonkAny 4)
forall a b. (a -> b) -> a -> b
$ (MemInfo (Ext SubExp) o MemReturn -> Doc (ZonkAny 4))
-> [MemInfo (Ext SubExp) o MemReturn] -> [Doc (ZonkAny 4)]
forall a b. (a -> b) -> [a] -> [b]
map MemInfo (Ext SubExp) o MemReturn -> Doc (ZonkAny 4)
forall a ann. Pretty a => a -> Doc ann
forall ann. MemInfo (Ext SubExp) o MemReturn -> Doc ann
pretty [MemInfo (Ext SubExp) o MemReturn]
rettype)
        Doc (ZonkAny 4) -> Doc (ZonkAny 4) -> Doc (ZonkAny 4)
forall ann. Doc ann -> Doc ann -> Doc ann
</> Doc (ZonkAny 4)
"does not have same number of elements as results"
        Doc (ZonkAny 4) -> Doc (ZonkAny 4) -> Doc (ZonkAny 4)
forall ann. Doc ann -> Doc ann -> Doc ann
</> Int -> Doc (ZonkAny 4) -> Doc (ZonkAny 4)
forall ann. Int -> Doc ann -> Doc ann
indent Int
2 ([Doc (ZonkAny 4)] -> Doc (ZonkAny 4)
forall a. [Doc a] -> Doc a
ppTupleLines' ([Doc (ZonkAny 4)] -> Doc (ZonkAny 4))
-> [Doc (ZonkAny 4)] -> Doc (ZonkAny 4)
forall a b. (a -> b) -> a -> b
$ (LetDecMem -> Doc (ZonkAny 4)) -> [LetDecMem] -> [Doc (ZonkAny 4)]
forall a b. (a -> b) -> [a] -> [b]
map LetDecMem -> Doc (ZonkAny 4)
forall a ann. Pretty a => a -> Doc ann
forall ann. LetDecMem -> Doc ann
pretty [LetDecMem]
ts)

  (Text -> TypeM rep ())
-> (() -> TypeM rep ()) -> Either Text () -> TypeM rep ()
forall a c b. (a -> c) -> (b -> c) -> Either a b -> c
either Text -> TypeM rep ()
bad () -> TypeM rep ()
forall a. a -> TypeM rep a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (Either Text () -> TypeM rep ())
-> TypeM rep (Either Text ()) -> TypeM rep ()
forall (m :: * -> *) a b. Monad m => (a -> m b) -> m a -> m b
=<< ExceptT Text (TypeM rep) () -> TypeM rep (Either Text ())
forall e (m :: * -> *) a. ExceptT e m a -> m (Either e a)
runExceptT ((MemInfo (Ext SubExp) o MemReturn
 -> LetDecMem -> ExceptT Text (TypeM rep) ())
-> [MemInfo (Ext SubExp) o MemReturn]
-> [LetDecMem]
-> ExceptT Text (TypeM rep) ()
forall (m :: * -> *) a b c.
Applicative m =>
(a -> b -> m c) -> [a] -> [b] -> m ()
zipWithM_ MemInfo (Ext SubExp) o MemReturn
-> LetDecMem -> ExceptT Text (TypeM rep) ()
checkReturn [MemInfo (Ext SubExp) o MemReturn]
rettype [LetDecMem]
ts)

matchPatToExp ::
  (Mem rep inner, LetDec rep ~ LetDecMem, TC.Checkable rep) =>
  Pat (LetDec (Aliases rep)) ->
  Exp (Aliases rep) ->
  TC.TypeM rep ()
matchPatToExp :: forall rep (inner :: * -> *).
(Mem rep inner, LetDec rep ~ LetDecMem, Checkable rep) =>
Pat (LetDec (Aliases rep)) -> Exp (Aliases rep) -> TypeM rep ()
matchPatToExp Pat (LetDec (Aliases rep))
pat Exp (Aliases rep)
e = do
  scope <- (Scope (Aliases rep) -> Scope rep) -> TypeM rep (Scope rep)
forall a. (Scope (Aliases rep) -> a) -> TypeM rep a
forall rep (m :: * -> *) a.
HasScope rep m =>
(Scope rep -> a) -> m a
asksScope Scope (Aliases rep) -> Scope rep
forall rep. Scope (Aliases rep) -> Scope rep
removeScopeAliases
  rt <- maybe illformed pure $ runReader (expReturns $ removeExpAliases e) scope

  let (ctx_ids, val_ts) = unzip $ bodyReturnsFromPat $ removePatAliases pat
      (ctx_map_ids, ctx_map_exts) = getExtMaps $ zip ctx_ids [0 .. 1]
      ok =
        [FunReturns] -> Int
forall a. [a] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length [FunReturns]
val_ts Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
== [ExpReturns] -> Int
forall a. [a] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length [ExpReturns]
rt
          Bool -> Bool -> Bool
&& [Bool] -> Bool
forall (t :: * -> *). Foldable t => t Bool -> Bool
and ((FunReturns -> ExpReturns -> Bool)
-> [FunReturns] -> [ExpReturns] -> [Bool]
forall a b c. (a -> b -> c) -> [a] -> [b] -> [c]
zipWith (Map (Ext VName) (TPrimExp Int64 (Ext VName))
-> Map (Ext VName) (TPrimExp Int64 (Ext VName))
-> FunReturns
-> ExpReturns
-> Bool
forall {d} {o} {o}.
Eq d =>
Map (Ext VName) (TPrimExp Int64 (Ext VName))
-> Map (Ext VName) (TPrimExp Int64 (Ext VName))
-> MemInfo d o MemReturn
-> MemInfo d o (Maybe MemReturn)
-> Bool
matches Map (Ext VName) (TPrimExp Int64 (Ext VName))
ctx_map_ids Map (Ext VName) (TPrimExp Int64 (Ext VName))
ctx_map_exts) [FunReturns]
val_ts [ExpReturns]
rt)

  unless ok . TC.bad . TC.TypeError . docText $
    "Expression type:"
      </> indent 2 (ppTupleLines' $ map pretty rt)
      </> "cannot match pattern type:"
      </> indent 2 (ppTupleLines' $ map pretty val_ts)
  where
    illformed :: TypeM rep [ExpReturns]
illformed =
      ErrorCase rep -> TypeM rep [ExpReturns]
forall rep a. ErrorCase rep -> TypeM rep a
TC.bad (ErrorCase rep -> TypeM rep [ExpReturns])
-> ErrorCase rep -> TypeM rep [ExpReturns]
forall a b. (a -> b) -> a -> b
$
        Text -> ErrorCase rep
forall rep. Text -> ErrorCase rep
TC.TypeError (Text -> ErrorCase rep)
-> (Doc (ZonkAny 5) -> Text) -> Doc (ZonkAny 5) -> ErrorCase rep
forall b c a. (b -> c) -> (a -> b) -> a -> c
forall {k} (cat :: k -> k -> *) (b :: k) (c :: k) (a :: k).
Category cat =>
cat b c -> cat a b -> cat a c
. Doc (ZonkAny 5) -> Text
forall a. Doc a -> Text
docText (Doc (ZonkAny 5) -> ErrorCase rep)
-> Doc (ZonkAny 5) -> ErrorCase rep
forall a b. (a -> b) -> a -> b
$
          Doc (ZonkAny 5)
"Expression"
            Doc (ZonkAny 5) -> Doc (ZonkAny 5) -> Doc (ZonkAny 5)
forall ann. Doc ann -> Doc ann -> Doc ann
</> Int -> Doc (ZonkAny 5) -> Doc (ZonkAny 5)
forall ann. Int -> Doc ann -> Doc ann
indent Int
2 (Exp (Aliases rep) -> Doc (ZonkAny 5)
forall a ann. Pretty a => a -> Doc ann
forall ann. Exp (Aliases rep) -> Doc ann
pretty Exp (Aliases rep)
e)
            Doc (ZonkAny 5) -> Doc (ZonkAny 5) -> Doc (ZonkAny 5)
forall ann. Doc ann -> Doc ann -> Doc ann
</> Doc (ZonkAny 5)
"cannot be assigned an index function."

    matches :: Map (Ext VName) (TPrimExp Int64 (Ext VName))
-> Map (Ext VName) (TPrimExp Int64 (Ext VName))
-> MemInfo d o MemReturn
-> MemInfo d o (Maybe MemReturn)
-> Bool
matches Map (Ext VName) (TPrimExp Int64 (Ext VName))
_ Map (Ext VName) (TPrimExp Int64 (Ext VName))
_ (MemPrim PrimType
x) (MemPrim PrimType
y) = PrimType
x PrimType -> PrimType -> Bool
forall a. Eq a => a -> a -> Bool
== PrimType
y
    matches Map (Ext VName) (TPrimExp Int64 (Ext VName))
_ Map (Ext VName) (TPrimExp Int64 (Ext VName))
_ (MemMem Space
x_space) (MemMem Space
y_space) =
      Space
x_space Space -> Space -> Bool
forall a. Eq a => a -> a -> Bool
== Space
y_space
    matches Map (Ext VName) (TPrimExp Int64 (Ext VName))
_ Map (Ext VName) (TPrimExp Int64 (Ext VName))
_ (MemAcc VName
x_accs ShapeBase SubExp
x_ispace [Type]
x_ts) (MemAcc VName
y_accs ShapeBase SubExp
y_ispace [Type]
y_ts) =
      (VName
x_accs, ShapeBase SubExp
x_ispace, [Type]
x_ts) (VName, ShapeBase SubExp, [Type])
-> (VName, ShapeBase SubExp, [Type]) -> Bool
forall a. Eq a => a -> a -> Bool
== (VName
y_accs, ShapeBase SubExp
y_ispace, [Type]
y_ts)
    matches Map (Ext VName) (TPrimExp Int64 (Ext VName))
ctxids Map (Ext VName) (TPrimExp Int64 (Ext VName))
ctxexts (MemArray PrimType
x_pt ShapeBase d
x_shape o
_ MemReturn
x_ret) (MemArray PrimType
y_pt ShapeBase d
y_shape o
_ Maybe MemReturn
y_ret) =
      PrimType
x_pt PrimType -> PrimType -> Bool
forall a. Eq a => a -> a -> Bool
== PrimType
y_pt
        Bool -> Bool -> Bool
&& ShapeBase d
x_shape ShapeBase d -> ShapeBase d -> Bool
forall a. Eq a => a -> a -> Bool
== ShapeBase d
y_shape
        Bool -> Bool -> Bool
&& case (MemReturn
x_ret, Maybe MemReturn
y_ret) of
          (ReturnsInBlock VName
_ ExtLMAD
x_lmad, Just (ReturnsInBlock VName
_ ExtLMAD
y_lmad)) ->
            let x_lmad' :: ExtLMAD
x_lmad' = Map (Ext VName) (TPrimExp Int64 (Ext VName)) -> ExtLMAD -> ExtLMAD
forall {k} a (t :: k).
Ord a =>
Map a (TPrimExp t a) -> LMAD (TPrimExp t a) -> LMAD (TPrimExp t a)
LMAD.substitute Map (Ext VName) (TPrimExp Int64 (Ext VName))
ctxids ExtLMAD
x_lmad
                y_lmad' :: ExtLMAD
y_lmad' = Map (Ext VName) (TPrimExp Int64 (Ext VName)) -> ExtLMAD -> ExtLMAD
forall {k} a (t :: k).
Ord a =>
Map a (TPrimExp t a) -> LMAD (TPrimExp t a) -> LMAD (TPrimExp t a)
LMAD.substitute Map (Ext VName) (TPrimExp Int64 (Ext VName))
ctxexts ExtLMAD
y_lmad
             in ExtLMAD -> ExtLMAD -> Bool
forall num. LMAD num -> LMAD num -> Bool
LMAD.closeEnough ExtLMAD
x_lmad' ExtLMAD
y_lmad'
          ( ReturnsInBlock VName
_ ExtLMAD
x_lmad,
            Just (ReturnsNewBlock Space
_ Int
_ ExtLMAD
y_lmad)
            ) ->
              let x_lmad' :: ExtLMAD
x_lmad' = Map (Ext VName) (TPrimExp Int64 (Ext VName)) -> ExtLMAD -> ExtLMAD
forall {k} a (t :: k).
Ord a =>
Map a (TPrimExp t a) -> LMAD (TPrimExp t a) -> LMAD (TPrimExp t a)
LMAD.substitute Map (Ext VName) (TPrimExp Int64 (Ext VName))
ctxids ExtLMAD
x_lmad
                  y_lmad' :: ExtLMAD
y_lmad' = Map (Ext VName) (TPrimExp Int64 (Ext VName)) -> ExtLMAD -> ExtLMAD
forall {k} a (t :: k).
Ord a =>
Map a (TPrimExp t a) -> LMAD (TPrimExp t a) -> LMAD (TPrimExp t a)
LMAD.substitute Map (Ext VName) (TPrimExp Int64 (Ext VName))
ctxexts ExtLMAD
y_lmad
               in ExtLMAD -> ExtLMAD -> Bool
forall num. LMAD num -> LMAD num -> Bool
LMAD.closeEnough ExtLMAD
x_lmad' ExtLMAD
y_lmad'
          ( ReturnsNewBlock Space
_ Int
x_i ExtLMAD
x_lmad,
            Just (ReturnsNewBlock Space
_ Int
y_i ExtLMAD
y_lmad)
            ) ->
              let x_lmad' :: ExtLMAD
x_lmad' = Map (Ext VName) (TPrimExp Int64 (Ext VName)) -> ExtLMAD -> ExtLMAD
forall {k} a (t :: k).
Ord a =>
Map a (TPrimExp t a) -> LMAD (TPrimExp t a) -> LMAD (TPrimExp t a)
LMAD.substitute Map (Ext VName) (TPrimExp Int64 (Ext VName))
ctxids ExtLMAD
x_lmad
                  y_lmad' :: ExtLMAD
y_lmad' = Map (Ext VName) (TPrimExp Int64 (Ext VName)) -> ExtLMAD -> ExtLMAD
forall {k} a (t :: k).
Ord a =>
Map a (TPrimExp t a) -> LMAD (TPrimExp t a) -> LMAD (TPrimExp t a)
LMAD.substitute Map (Ext VName) (TPrimExp Int64 (Ext VName))
ctxexts ExtLMAD
y_lmad
               in Int
x_i Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
== Int
y_i Bool -> Bool -> Bool
&& ExtLMAD -> ExtLMAD -> Bool
forall num. LMAD num -> LMAD num -> Bool
LMAD.closeEnough ExtLMAD
x_lmad' ExtLMAD
y_lmad'
          (MemReturn
_, Maybe MemReturn
Nothing) -> Bool
True
          (MemReturn, Maybe MemReturn)
_ -> Bool
False
    matches Map (Ext VName) (TPrimExp Int64 (Ext VName))
_ Map (Ext VName) (TPrimExp Int64 (Ext VName))
_ MemInfo d o MemReturn
_ MemInfo d o (Maybe MemReturn)
_ = Bool
False

varMemInfo ::
  (Mem rep inner) =>
  VName ->
  TC.TypeM rep (MemInfo SubExp NoMode MemBind)
varMemInfo :: forall rep (inner :: * -> *).
Mem rep inner =>
VName -> TypeM rep LetDecMem
varMemInfo VName
name = do
  dec <- VName -> TypeM rep (NameInfo (Aliases rep))
forall rep. VName -> TypeM rep (NameInfo (Aliases rep))
TC.lookupVar VName
name

  case dec of
    LetName (AliasDec
_, LetDec rep
summary) -> LetDecMem -> TypeM rep LetDecMem
forall a. a -> TypeM rep a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (LetDecMem -> TypeM rep LetDecMem)
-> LetDecMem -> TypeM rep LetDecMem
forall a b. (a -> b) -> a -> b
$ LetDec rep -> LetDecMem
forall t. HasLetDecMem t => t -> LetDecMem
letDecMem LetDec rep
summary
    FParamName FParamInfo (Aliases rep)
summary -> LetDecMem -> TypeM rep LetDecMem
forall a. a -> TypeM rep a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (LetDecMem -> TypeM rep LetDecMem)
-> LetDecMem -> TypeM rep LetDecMem
forall a b. (a -> b) -> a -> b
$ FParamMem -> LetDecMem
forall d o r. MemInfo d o r -> MemInfo d NoMode r
noModeReturns FParamInfo (Aliases rep)
FParamMem
summary
    LParamName LParamInfo (Aliases rep)
summary -> LetDecMem -> TypeM rep LetDecMem
forall a. a -> TypeM rep a
forall (f :: * -> *) a. Applicative f => a -> f a
pure LParamInfo (Aliases rep)
LetDecMem
summary
    IndexName IntType
it -> LetDecMem -> TypeM rep LetDecMem
forall a. a -> TypeM rep a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (LetDecMem -> TypeM rep LetDecMem)
-> LetDecMem -> TypeM rep LetDecMem
forall a b. (a -> b) -> a -> b
$ PrimType -> LetDecMem
forall d o ret. PrimType -> MemInfo d o ret
MemPrim (PrimType -> LetDecMem) -> PrimType -> LetDecMem
forall a b. (a -> b) -> a -> b
$ IntType -> PrimType
IntType IntType
it

-- | Turn info into memory information.
nameInfoToMemInfo :: (Mem rep inner) => NameInfo rep -> MemBound NoMode
nameInfoToMemInfo :: forall rep (inner :: * -> *).
Mem rep inner =>
NameInfo rep -> LetDecMem
nameInfoToMemInfo NameInfo rep
info =
  case NameInfo rep
info of
    FParamName FParamInfo rep
summary -> FParamMem -> LetDecMem
forall d o r. MemInfo d o r -> MemInfo d NoMode r
noModeReturns FParamInfo rep
FParamMem
summary
    LParamName LParamInfo rep
summary -> LParamInfo rep
LetDecMem
summary
    LetName LetDec rep
summary -> LetDec rep -> LetDecMem
forall t. HasLetDecMem t => t -> LetDecMem
letDecMem LetDec rep
summary
    IndexName IntType
it -> PrimType -> LetDecMem
forall d o ret. PrimType -> MemInfo d o ret
MemPrim (PrimType -> LetDecMem) -> PrimType -> LetDecMem
forall a b. (a -> b) -> a -> b
$ IntType -> PrimType
IntType IntType
it

-- | Look up information about the memory block with this name.
lookupMemInfo ::
  (HasScope rep m, Mem rep inner) =>
  VName ->
  m (MemInfo SubExp NoMode MemBind)
lookupMemInfo :: forall rep (m :: * -> *) (inner :: * -> *).
(HasScope rep m, Mem rep inner) =>
VName -> m LetDecMem
lookupMemInfo = (NameInfo rep -> LetDecMem) -> m (NameInfo rep) -> m LetDecMem
forall a b. (a -> b) -> m a -> m b
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap NameInfo rep -> LetDecMem
forall rep (inner :: * -> *).
Mem rep inner =>
NameInfo rep -> LetDecMem
nameInfoToMemInfo (m (NameInfo rep) -> m LetDecMem)
-> (VName -> m (NameInfo rep)) -> VName -> m LetDecMem
forall b c a. (b -> c) -> (a -> b) -> a -> c
forall {k} (cat :: k -> k -> *) (b :: k) (c :: k) (a :: k).
Category cat =>
cat b c -> cat a b -> cat a c
. VName -> m (NameInfo rep)
forall rep (m :: * -> *).
HasScope rep m =>
VName -> m (NameInfo rep)
lookupInfo

subExpMemInfo ::
  (HasScope rep m, Mem rep inner) =>
  SubExp ->
  m (MemInfo SubExp NoMode MemBind)
subExpMemInfo :: forall rep (m :: * -> *) (inner :: * -> *).
(HasScope rep m, Mem rep inner) =>
SubExp -> m LetDecMem
subExpMemInfo (Var VName
v) = VName -> m LetDecMem
forall rep (m :: * -> *) (inner :: * -> *).
(HasScope rep m, Mem rep inner) =>
VName -> m LetDecMem
lookupMemInfo VName
v
subExpMemInfo (Constant PrimValue
v) = LetDecMem -> m LetDecMem
forall a. a -> m a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (LetDecMem -> m LetDecMem) -> LetDecMem -> m LetDecMem
forall a b. (a -> b) -> a -> b
$ PrimType -> LetDecMem
forall d o ret. PrimType -> MemInfo d o ret
MemPrim (PrimType -> LetDecMem) -> PrimType -> LetDecMem
forall a b. (a -> b) -> a -> b
$ PrimValue -> PrimType
primValueType PrimValue
v

lookupArraySummary ::
  (Mem rep inner, HasScope rep m, Monad m) =>
  VName ->
  m (VName, LMAD.LMAD (TPrimExp Int64 VName))
lookupArraySummary :: forall rep (inner :: * -> *) (m :: * -> *).
(Mem rep inner, HasScope rep m, Monad m) =>
VName -> m (VName, LMAD)
lookupArraySummary VName
name = do
  summary <- VName -> m LetDecMem
forall rep (m :: * -> *) (inner :: * -> *).
(HasScope rep m, Mem rep inner) =>
VName -> m LetDecMem
lookupMemInfo VName
name
  case summary of
    MemArray PrimType
_ ShapeBase SubExp
_ NoMode
_ (ArrayIn VName
mem LMAD
lmad) ->
      (VName, LMAD) -> m (VName, LMAD)
forall a. a -> m a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (VName
mem, LMAD
lmad)
    LetDecMem
_ ->
      String -> m (VName, LMAD)
forall a. HasCallStack => String -> a
error (String -> m (VName, LMAD))
-> (Text -> String) -> Text -> m (VName, LMAD)
forall b c a. (b -> c) -> (a -> b) -> a -> c
forall {k} (cat :: k -> k -> *) (b :: k) (c :: k) (a :: k).
Category cat =>
cat b c -> cat a b -> cat a c
. Text -> String
T.unpack (Text -> m (VName, LMAD)) -> Text -> m (VName, LMAD)
forall a b. (a -> b) -> a -> b
$
        Text
"Expected "
          Text -> Text -> Text
forall a. Semigroup a => a -> a -> a
<> VName -> Text
forall a. Pretty a => a -> Text
prettyText VName
name
          Text -> Text -> Text
forall a. Semigroup a => a -> a -> a
<> Text
" to be array but bound to:\n"
          Text -> Text -> Text
forall a. Semigroup a => a -> a -> a
<> LetDecMem -> Text
forall a. Pretty a => a -> Text
prettyText LetDecMem
summary

lookupMemSpace ::
  (Mem rep inner, HasScope rep m, Monad m) =>
  VName ->
  m Space
lookupMemSpace :: forall rep (inner :: * -> *) (m :: * -> *).
(Mem rep inner, HasScope rep m, Monad m) =>
VName -> m Space
lookupMemSpace VName
name = do
  summary <- VName -> m LetDecMem
forall rep (m :: * -> *) (inner :: * -> *).
(HasScope rep m, Mem rep inner) =>
VName -> m LetDecMem
lookupMemInfo VName
name
  case summary of
    MemMem Space
space ->
      Space -> m Space
forall a. a -> m a
forall (f :: * -> *) a. Applicative f => a -> f a
pure Space
space
    LetDecMem
_ ->
      String -> m Space
forall a. HasCallStack => String -> a
error (String -> m Space) -> (Text -> String) -> Text -> m Space
forall b c a. (b -> c) -> (a -> b) -> a -> c
forall {k} (cat :: k -> k -> *) (b :: k) (c :: k) (a :: k).
Category cat =>
cat b c -> cat a b -> cat a c
. Text -> String
T.unpack (Text -> m Space) -> Text -> m Space
forall a b. (a -> b) -> a -> b
$
        Text
"Expected "
          Text -> Text -> Text
forall a. Semigroup a => a -> a -> a
<> VName -> Text
forall a. Pretty a => a -> Text
prettyText VName
name
          Text -> Text -> Text
forall a. Semigroup a => a -> a -> a
<> Text
" to be memory but bound to:\n"
          Text -> Text -> Text
forall a. Semigroup a => a -> a -> a
<> LetDecMem -> Text
forall a. Pretty a => a -> Text
prettyText LetDecMem
summary

checkMemInfo ::
  (TC.Checkable rep) =>
  VName ->
  MemInfo SubExp o MemBind ->
  TC.TypeM rep ()
checkMemInfo :: forall rep o.
Checkable rep =>
VName -> MemInfo SubExp o MemBind -> TypeM rep ()
checkMemInfo VName
_ (MemPrim PrimType
_) = () -> TypeM rep ()
forall a. a -> TypeM rep a
forall (f :: * -> *) a. Applicative f => a -> f a
pure ()
checkMemInfo VName
_ (MemMem (ScalarSpace [SubExp]
d PrimType
_)) = (SubExp -> TypeM rep ()) -> [SubExp] -> TypeM rep ()
forall (t :: * -> *) (m :: * -> *) a b.
(Foldable t, Monad m) =>
(a -> m b) -> t a -> m ()
mapM_ (Type -> SubExp -> TypeM rep ()
forall rep. Checkable rep => Type -> SubExp -> TypeM rep ()
TC.require (PrimType -> Type
forall shape o. PrimType -> TypeBase shape o
Prim PrimType
int64)) [SubExp]
d
checkMemInfo VName
_ (MemMem Space
_) = () -> TypeM rep ()
forall a. a -> TypeM rep a
forall (f :: * -> *) a. Applicative f => a -> f a
pure ()
checkMemInfo VName
_ (MemAcc VName
acc ShapeBase SubExp
ispace [Type]
ts) =
  TypeBase (ShapeBase SubExp) (ZonkAny 0) -> TypeM rep ()
forall rep o.
Checkable rep =>
TypeBase (ShapeBase SubExp) o -> TypeM rep ()
TC.checkType (TypeBase (ShapeBase SubExp) (ZonkAny 0) -> TypeM rep ())
-> TypeBase (ShapeBase SubExp) (ZonkAny 0) -> TypeM rep ()
forall a b. (a -> b) -> a -> b
$ VName
-> ShapeBase SubExp
-> [Type]
-> TypeBase (ShapeBase SubExp) (ZonkAny 0)
forall shape o.
VName -> ShapeBase SubExp -> [Type] -> TypeBase shape o
Acc VName
acc ShapeBase SubExp
ispace [Type]
ts
checkMemInfo VName
name (MemArray PrimType
_ ShapeBase SubExp
shape o
_ (ArrayIn VName
v LMAD
lmad)) = do
  t <- VName -> TypeM rep Type
forall rep (m :: * -> *). HasScope rep m => VName -> m Type
lookupType VName
v
  case t of
    Mem {} ->
      () -> TypeM rep ()
forall a. a -> TypeM rep a
forall (f :: * -> *) a. Applicative f => a -> f a
pure ()
    Type
_ ->
      ErrorCase rep -> TypeM rep ()
forall rep a. ErrorCase rep -> TypeM rep a
TC.bad (ErrorCase rep -> TypeM rep ()) -> ErrorCase rep -> TypeM rep ()
forall a b. (a -> b) -> a -> b
$
        Text -> ErrorCase rep
forall rep. Text -> ErrorCase rep
TC.TypeError (Text -> ErrorCase rep) -> Text -> ErrorCase rep
forall a b. (a -> b) -> a -> b
$
          Text
"Variable "
            Text -> Text -> Text
forall a. Semigroup a => a -> a -> a
<> VName -> Text
forall a. Pretty a => a -> Text
prettyText VName
v
            Text -> Text -> Text
forall a. Semigroup a => a -> a -> a
<> Text
" used as memory block, but is of type "
            Text -> Text -> Text
forall a. Semigroup a => a -> a -> a
<> Type -> Text
forall a. Pretty a => a -> Text
prettyText Type
t
            Text -> Text -> Text
forall a. Semigroup a => a -> a -> a
<> Text
"."

  TC.context ("in index function " <> prettyText lmad) $ do
    traverse_ (TC.requirePrimExp int64 . untyped) lmad
    unless (LMAD.shape lmad == map pe64 (shapeDims shape)) $
      TC.bad $
        TC.TypeError $
          "Shape of index function ("
            <> prettyText (LMAD.shape lmad)
            <> ") does not match shape of array "
            <> prettyText name
            <> " ("
            <> prettyText shape
            <> ")"

bodyReturnsFromPat ::
  Pat (MemBound NoMode) -> [(VName, BodyReturns)]
bodyReturnsFromPat :: Pat LetDecMem -> [(VName, FunReturns)]
bodyReturnsFromPat Pat LetDecMem
pat =
  (PatElem LetDecMem -> (VName, FunReturns))
-> [PatElem LetDecMem] -> [(VName, FunReturns)]
forall a b. (a -> b) -> [a] -> [b]
map PatElem LetDecMem -> (VName, FunReturns)
asReturns ([PatElem LetDecMem] -> [(VName, FunReturns)])
-> [PatElem LetDecMem] -> [(VName, FunReturns)]
forall a b. (a -> b) -> a -> b
$ Pat LetDecMem -> [PatElem LetDecMem]
forall dec. Pat dec -> [PatElem dec]
patElems Pat LetDecMem
pat
  where
    ctx :: [PatElem LetDecMem]
ctx = Pat LetDecMem -> [PatElem LetDecMem]
forall dec. Pat dec -> [PatElem dec]
patElems Pat LetDecMem
pat

    ext :: SubExp -> Ext SubExp
ext (Var VName
v)
      | Just (Int
i, PatElem LetDecMem
_) <- ((Int, PatElem LetDecMem) -> Bool)
-> [(Int, PatElem LetDecMem)] -> Maybe (Int, PatElem LetDecMem)
forall (t :: * -> *) a. Foldable t => (a -> Bool) -> t a -> Maybe a
find ((VName -> VName -> Bool
forall a. Eq a => a -> a -> Bool
== VName
v) (VName -> Bool)
-> ((Int, PatElem LetDecMem) -> VName)
-> (Int, PatElem LetDecMem)
-> Bool
forall b c a. (b -> c) -> (a -> b) -> a -> c
forall {k} (cat :: k -> k -> *) (b :: k) (c :: k) (a :: k).
Category cat =>
cat b c -> cat a b -> cat a c
. PatElem LetDecMem -> VName
forall dec. PatElem dec -> VName
patElemName (PatElem LetDecMem -> VName)
-> ((Int, PatElem LetDecMem) -> PatElem LetDecMem)
-> (Int, PatElem LetDecMem)
-> VName
forall b c a. (b -> c) -> (a -> b) -> a -> c
forall {k} (cat :: k -> k -> *) (b :: k) (c :: k) (a :: k).
Category cat =>
cat b c -> cat a b -> cat a c
. (Int, PatElem LetDecMem) -> PatElem LetDecMem
forall a b. (a, b) -> b
snd) ([(Int, PatElem LetDecMem)] -> Maybe (Int, PatElem LetDecMem))
-> [(Int, PatElem LetDecMem)] -> Maybe (Int, PatElem LetDecMem)
forall a b. (a -> b) -> a -> b
$ [Int] -> [PatElem LetDecMem] -> [(Int, PatElem LetDecMem)]
forall a b. [a] -> [b] -> [(a, b)]
zip [Int
0 ..] [PatElem LetDecMem]
ctx =
          Int -> Ext SubExp
forall a. Int -> Ext a
Ext Int
i
    ext SubExp
se = SubExp -> Ext SubExp
forall a. a -> Ext a
Free SubExp
se

    asReturns :: PatElem LetDecMem -> (VName, FunReturns)
asReturns PatElem LetDecMem
pe =
      ( PatElem LetDecMem -> VName
forall dec. PatElem dec -> VName
patElemName PatElem LetDecMem
pe,
        case PatElem LetDecMem -> LetDecMem
forall dec. PatElem dec -> dec
patElemDec PatElem LetDecMem
pe of
          MemPrim PrimType
pt -> PrimType -> FunReturns
forall d o ret. PrimType -> MemInfo d o ret
MemPrim PrimType
pt
          MemMem Space
space -> Space -> FunReturns
forall d o ret. Space -> MemInfo d o ret
MemMem Space
space
          MemArray PrimType
pt ShapeBase SubExp
shape NoMode
o (ArrayIn VName
mem LMAD
lmad) ->
            PrimType
-> ShapeBase (Ext SubExp) -> NoMode -> MemReturn -> FunReturns
forall d o ret.
PrimType -> ShapeBase d -> o -> ret -> MemInfo d o ret
MemArray PrimType
pt ([Ext SubExp] -> ShapeBase (Ext SubExp)
forall d. [d] -> ShapeBase d
Shape ([Ext SubExp] -> ShapeBase (Ext SubExp))
-> [Ext SubExp] -> ShapeBase (Ext SubExp)
forall a b. (a -> b) -> a -> b
$ (SubExp -> Ext SubExp) -> [SubExp] -> [Ext SubExp]
forall a b. (a -> b) -> [a] -> [b]
map SubExp -> Ext SubExp
ext ([SubExp] -> [Ext SubExp]) -> [SubExp] -> [Ext SubExp]
forall a b. (a -> b) -> a -> b
$ ShapeBase SubExp -> [SubExp]
forall d. ShapeBase d -> [d]
shapeDims ShapeBase SubExp
shape) NoMode
o (MemReturn -> FunReturns) -> MemReturn -> FunReturns
forall a b. (a -> b) -> a -> b
$
              case ((Int, PatElem LetDecMem) -> Bool)
-> [(Int, PatElem LetDecMem)] -> Maybe (Int, PatElem LetDecMem)
forall (t :: * -> *) a. Foldable t => (a -> Bool) -> t a -> Maybe a
find ((VName -> VName -> Bool
forall a. Eq a => a -> a -> Bool
== VName
mem) (VName -> Bool)
-> ((Int, PatElem LetDecMem) -> VName)
-> (Int, PatElem LetDecMem)
-> Bool
forall b c a. (b -> c) -> (a -> b) -> a -> c
forall {k} (cat :: k -> k -> *) (b :: k) (c :: k) (a :: k).
Category cat =>
cat b c -> cat a b -> cat a c
. PatElem LetDecMem -> VName
forall dec. PatElem dec -> VName
patElemName (PatElem LetDecMem -> VName)
-> ((Int, PatElem LetDecMem) -> PatElem LetDecMem)
-> (Int, PatElem LetDecMem)
-> VName
forall b c a. (b -> c) -> (a -> b) -> a -> c
forall {k} (cat :: k -> k -> *) (b :: k) (c :: k) (a :: k).
Category cat =>
cat b c -> cat a b -> cat a c
. (Int, PatElem LetDecMem) -> PatElem LetDecMem
forall a b. (a, b) -> b
snd) ([(Int, PatElem LetDecMem)] -> Maybe (Int, PatElem LetDecMem))
-> [(Int, PatElem LetDecMem)] -> Maybe (Int, PatElem LetDecMem)
forall a b. (a -> b) -> a -> b
$ [Int] -> [PatElem LetDecMem] -> [(Int, PatElem LetDecMem)]
forall a b. [a] -> [b] -> [(a, b)]
zip [Int
0 ..] [PatElem LetDecMem]
ctx of
                Just (Int
i, PatElem VName
_ (MemMem Space
space)) ->
                  Space -> Int -> ExtLMAD -> MemReturn
ReturnsNewBlock Space
space Int
i (ExtLMAD -> MemReturn) -> ExtLMAD -> MemReturn
forall a b. (a -> b) -> a -> b
$
                    [VName] -> LMAD -> ExtLMAD
existentialiseLMAD ((PatElem LetDecMem -> VName) -> [PatElem LetDecMem] -> [VName]
forall a b. (a -> b) -> [a] -> [b]
map PatElem LetDecMem -> VName
forall dec. PatElem dec -> VName
patElemName [PatElem LetDecMem]
ctx) LMAD
lmad
                Maybe (Int, PatElem LetDecMem)
_ -> VName -> ExtLMAD -> MemReturn
ReturnsInBlock VName
mem (ExtLMAD -> MemReturn) -> ExtLMAD -> MemReturn
forall a b. (a -> b) -> a -> b
$ [VName] -> LMAD -> ExtLMAD
existentialiseLMAD [] LMAD
lmad
          MemAcc VName
acc ShapeBase SubExp
ispace [Type]
ts -> VName -> ShapeBase SubExp -> [Type] -> FunReturns
forall d o ret.
VName -> ShapeBase SubExp -> [Type] -> MemInfo d o ret
MemAcc VName
acc ShapeBase SubExp
ispace [Type]
ts
      )

extReturns :: [ExtType] -> [ExpReturns]
extReturns :: [ExtType] -> [ExpReturns]
extReturns [ExtType]
ets =
  State Int [ExpReturns] -> Int -> [ExpReturns]
forall s a. State s a -> s -> a
evalState ((ExtType -> StateT Int Identity ExpReturns)
-> [ExtType] -> State Int [ExpReturns]
forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
(a -> m b) -> t a -> m (t b)
forall (m :: * -> *) a b. Monad m => (a -> m b) -> [a] -> m [b]
mapM ExtType -> StateT Int Identity ExpReturns
forall {f :: * -> *}. MonadState Int f => ExtType -> f ExpReturns
addDec [ExtType]
ets) Int
0
  where
    addDec :: ExtType -> f ExpReturns
addDec (Prim PrimType
bt) =
      ExpReturns -> f ExpReturns
forall a. a -> f a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (ExpReturns -> f ExpReturns) -> ExpReturns -> f ExpReturns
forall a b. (a -> b) -> a -> b
$ PrimType -> ExpReturns
forall d o ret. PrimType -> MemInfo d o ret
MemPrim PrimType
bt
    addDec (Mem Space
space) =
      ExpReturns -> f ExpReturns
forall a. a -> f a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (ExpReturns -> f ExpReturns) -> ExpReturns -> f ExpReturns
forall a b. (a -> b) -> a -> b
$ Space -> ExpReturns
forall d o ret. Space -> MemInfo d o ret
MemMem Space
space
    addDec t :: ExtType
t@(Array PrimType
bt ShapeBase (Ext SubExp)
shape NoMode
o)
      | ExtType -> Bool
existential ExtType
t = do
          i <- f Int
forall s (m :: * -> *). MonadState s m => m s
get f Int -> f () -> f Int
forall a b. f a -> f b -> f a
forall (f :: * -> *) a b. Applicative f => f a -> f b -> f a
<* (Int -> Int) -> f ()
forall s (m :: * -> *). MonadState s m => (s -> s) -> m ()
modify (Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
1)
          pure . MemArray bt shape o . Just $
            ReturnsNewBlock DefaultSpace i $
              LMAD.iota 0 (map convert $ shapeDims shape)
      | Bool
otherwise =
          ExpReturns -> f ExpReturns
forall a. a -> f a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (ExpReturns -> f ExpReturns) -> ExpReturns -> f ExpReturns
forall a b. (a -> b) -> a -> b
$ PrimType
-> ShapeBase (Ext SubExp)
-> NoMode
-> Maybe MemReturn
-> ExpReturns
forall d o ret.
PrimType -> ShapeBase d -> o -> ret -> MemInfo d o ret
MemArray PrimType
bt ShapeBase (Ext SubExp)
shape NoMode
o Maybe MemReturn
forall a. Maybe a
Nothing
    addDec (Acc VName
acc ShapeBase SubExp
ispace [Type]
ts) =
      ExpReturns -> f ExpReturns
forall a. a -> f a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (ExpReturns -> f ExpReturns) -> ExpReturns -> f ExpReturns
forall a b. (a -> b) -> a -> b
$ VName -> ShapeBase SubExp -> [Type] -> ExpReturns
forall d o ret.
VName -> ShapeBase SubExp -> [Type] -> MemInfo d o ret
MemAcc VName
acc ShapeBase SubExp
ispace [Type]
ts
    convert :: Ext SubExp -> TPrimExp Int64 (Ext VName)
convert (Ext Int
i) = Ext VName -> TPrimExp Int64 (Ext VName)
forall a. a -> TPrimExp Int64 a
le64 (Int -> Ext VName
forall a. Int -> Ext a
Ext Int
i)
    convert (Free SubExp
v) = VName -> Ext VName
forall a. a -> Ext a
Free (VName -> Ext VName)
-> TPrimExp Int64 VName -> TPrimExp Int64 (Ext VName)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> SubExp -> TPrimExp Int64 VName
pe64 SubExp
v

arrayVarReturns ::
  (HasScope rep m, Monad m, Mem rep inner) =>
  VName ->
  m (PrimType, Shape, VName, LMAD)
arrayVarReturns :: forall rep (m :: * -> *) (inner :: * -> *).
(HasScope rep m, Monad m, Mem rep inner) =>
VName -> m (PrimType, ShapeBase SubExp, VName, LMAD)
arrayVarReturns VName
v = do
  summary <- VName -> m LetDecMem
forall rep (m :: * -> *) (inner :: * -> *).
(HasScope rep m, Mem rep inner) =>
VName -> m LetDecMem
lookupMemInfo VName
v
  case summary of
    MemArray PrimType
et ShapeBase SubExp
shape NoMode
_ (ArrayIn VName
mem LMAD
lmad) ->
      (PrimType, ShapeBase SubExp, VName, LMAD)
-> m (PrimType, ShapeBase SubExp, VName, LMAD)
forall a. a -> m a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (PrimType
et, [SubExp] -> ShapeBase SubExp
forall d. [d] -> ShapeBase d
Shape ([SubExp] -> ShapeBase SubExp) -> [SubExp] -> ShapeBase SubExp
forall a b. (a -> b) -> a -> b
$ ShapeBase SubExp -> [SubExp]
forall d. ShapeBase d -> [d]
shapeDims ShapeBase SubExp
shape, VName
mem, LMAD
lmad)
    LetDecMem
_ ->
      String -> m (PrimType, ShapeBase SubExp, VName, LMAD)
forall a. HasCallStack => String -> a
error (String -> m (PrimType, ShapeBase SubExp, VName, LMAD))
-> (Text -> String)
-> Text
-> m (PrimType, ShapeBase SubExp, VName, LMAD)
forall b c a. (b -> c) -> (a -> b) -> a -> c
forall {k} (cat :: k -> k -> *) (b :: k) (c :: k) (a :: k).
Category cat =>
cat b c -> cat a b -> cat a c
. Text -> String
T.unpack (Text -> m (PrimType, ShapeBase SubExp, VName, LMAD))
-> Text -> m (PrimType, ShapeBase SubExp, VName, LMAD)
forall a b. (a -> b) -> a -> b
$ Text
"arrayVarReturns: " Text -> Text -> Text
forall a. Semigroup a => a -> a -> a
<> VName -> Text
forall a. Pretty a => a -> Text
prettyText VName
v Text -> Text -> Text
forall a. Semigroup a => a -> a -> a
<> Text
" is not an array."

varReturns ::
  (HasScope rep m, Monad m, Mem rep inner) =>
  VName ->
  m ExpReturns
varReturns :: forall rep (m :: * -> *) (inner :: * -> *).
(HasScope rep m, Monad m, Mem rep inner) =>
VName -> m ExpReturns
varReturns VName
v = do
  summary <- VName -> m LetDecMem
forall rep (m :: * -> *) (inner :: * -> *).
(HasScope rep m, Mem rep inner) =>
VName -> m LetDecMem
lookupMemInfo VName
v
  case summary of
    MemPrim PrimType
bt ->
      ExpReturns -> m ExpReturns
forall a. a -> m a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (ExpReturns -> m ExpReturns) -> ExpReturns -> m ExpReturns
forall a b. (a -> b) -> a -> b
$ PrimType -> ExpReturns
forall d o ret. PrimType -> MemInfo d o ret
MemPrim PrimType
bt
    MemArray PrimType
et ShapeBase SubExp
shape NoMode
_ (ArrayIn VName
mem LMAD
lmad) ->
      ExpReturns -> m ExpReturns
forall a. a -> m a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (ExpReturns -> m ExpReturns) -> ExpReturns -> m ExpReturns
forall a b. (a -> b) -> a -> b
$
        PrimType
-> ShapeBase (Ext SubExp)
-> NoMode
-> Maybe MemReturn
-> ExpReturns
forall d o ret.
PrimType -> ShapeBase d -> o -> ret -> MemInfo d o ret
MemArray PrimType
et ((SubExp -> Ext SubExp)
-> ShapeBase SubExp -> ShapeBase (Ext SubExp)
forall a b. (a -> b) -> ShapeBase a -> ShapeBase b
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap SubExp -> Ext SubExp
forall a. a -> Ext a
Free ShapeBase SubExp
shape) NoMode
NoMode (Maybe MemReturn -> ExpReturns) -> Maybe MemReturn -> ExpReturns
forall a b. (a -> b) -> a -> b
$
          MemReturn -> Maybe MemReturn
forall a. a -> Maybe a
Just (MemReturn -> Maybe MemReturn) -> MemReturn -> Maybe MemReturn
forall a b. (a -> b) -> a -> b
$
            VName -> ExtLMAD -> MemReturn
ReturnsInBlock VName
mem (ExtLMAD -> MemReturn) -> ExtLMAD -> MemReturn
forall a b. (a -> b) -> a -> b
$
              [VName] -> LMAD -> ExtLMAD
existentialiseLMAD [] LMAD
lmad
    MemMem Space
space ->
      ExpReturns -> m ExpReturns
forall a. a -> m a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (ExpReturns -> m ExpReturns) -> ExpReturns -> m ExpReturns
forall a b. (a -> b) -> a -> b
$ Space -> ExpReturns
forall d o ret. Space -> MemInfo d o ret
MemMem Space
space
    MemAcc VName
acc ShapeBase SubExp
ispace [Type]
ts ->
      ExpReturns -> m ExpReturns
forall a. a -> m a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (ExpReturns -> m ExpReturns) -> ExpReturns -> m ExpReturns
forall a b. (a -> b) -> a -> b
$ VName -> ShapeBase SubExp -> [Type] -> ExpReturns
forall d o ret.
VName -> ShapeBase SubExp -> [Type] -> MemInfo d o ret
MemAcc VName
acc ShapeBase SubExp
ispace [Type]
ts

subExpReturns :: (HasScope rep m, Monad m, Mem rep inner) => SubExp -> m ExpReturns
subExpReturns :: forall rep (m :: * -> *) (inner :: * -> *).
(HasScope rep m, Monad m, Mem rep inner) =>
SubExp -> m ExpReturns
subExpReturns (Var VName
v) =
  VName -> m ExpReturns
forall rep (m :: * -> *) (inner :: * -> *).
(HasScope rep m, Monad m, Mem rep inner) =>
VName -> m ExpReturns
varReturns VName
v
subExpReturns (Constant PrimValue
v) =
  ExpReturns -> m ExpReturns
forall a. a -> m a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (ExpReturns -> m ExpReturns) -> ExpReturns -> m ExpReturns
forall a b. (a -> b) -> a -> b
$ PrimType -> ExpReturns
forall d o ret. PrimType -> MemInfo d o ret
MemPrim (PrimType -> ExpReturns) -> PrimType -> ExpReturns
forall a b. (a -> b) -> a -> b
$ PrimValue -> PrimType
primValueType PrimValue
v

-- | The return information of an expression.  This can be seen as the
-- "return type with memory annotations" of the expression.
--
-- This can produce Nothing, which signifies that the result is an
-- array layout that is not expressible as an index function.
expReturns ::
  (LocalScope rep m, Mem rep inner) =>
  Exp rep ->
  m (Maybe [ExpReturns])
expReturns :: forall rep (m :: * -> *) (inner :: * -> *).
(LocalScope rep m, Mem rep inner) =>
Exp rep -> m (Maybe [ExpReturns])
expReturns (BasicOp (SubExp SubExp
se)) =
  [ExpReturns] -> Maybe [ExpReturns]
forall a. a -> Maybe a
Just ([ExpReturns] -> Maybe [ExpReturns])
-> (ExpReturns -> [ExpReturns]) -> ExpReturns -> Maybe [ExpReturns]
forall b c a. (b -> c) -> (a -> b) -> a -> c
forall {k} (cat :: k -> k -> *) (b :: k) (c :: k) (a :: k).
Category cat =>
cat b c -> cat a b -> cat a c
. ExpReturns -> [ExpReturns]
forall a. a -> [a]
forall (f :: * -> *) a. Applicative f => a -> f a
pure (ExpReturns -> Maybe [ExpReturns])
-> m ExpReturns -> m (Maybe [ExpReturns])
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> SubExp -> m ExpReturns
forall rep (m :: * -> *) (inner :: * -> *).
(HasScope rep m, Monad m, Mem rep inner) =>
SubExp -> m ExpReturns
subExpReturns SubExp
se
expReturns (BasicOp (Opaque OpaqueOp
_ (Var VName
v))) =
  [ExpReturns] -> Maybe [ExpReturns]
forall a. a -> Maybe a
Just ([ExpReturns] -> Maybe [ExpReturns])
-> (ExpReturns -> [ExpReturns]) -> ExpReturns -> Maybe [ExpReturns]
forall b c a. (b -> c) -> (a -> b) -> a -> c
forall {k} (cat :: k -> k -> *) (b :: k) (c :: k) (a :: k).
Category cat =>
cat b c -> cat a b -> cat a c
. ExpReturns -> [ExpReturns]
forall a. a -> [a]
forall (f :: * -> *) a. Applicative f => a -> f a
pure (ExpReturns -> Maybe [ExpReturns])
-> m ExpReturns -> m (Maybe [ExpReturns])
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> VName -> m ExpReturns
forall rep (m :: * -> *) (inner :: * -> *).
(HasScope rep m, Monad m, Mem rep inner) =>
VName -> m ExpReturns
varReturns VName
v
expReturns (BasicOp (Reshape VName
v NewShape SubExp
newshape)) = do
  (et, _, mem, lmad) <- VName -> m (PrimType, ShapeBase SubExp, VName, LMAD)
forall rep (m :: * -> *) (inner :: * -> *).
(HasScope rep m, Monad m, Mem rep inner) =>
VName -> m (PrimType, ShapeBase SubExp, VName, LMAD)
arrayVarReturns VName
v
  case reshaper (reshapeKind newshape) lmad $ map pe64 $ shapeDims $ newShape newshape of
    Just LMAD
lmad' ->
      Maybe [ExpReturns] -> m (Maybe [ExpReturns])
forall a. a -> m a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (Maybe [ExpReturns] -> m (Maybe [ExpReturns]))
-> ([ExpReturns] -> Maybe [ExpReturns])
-> [ExpReturns]
-> m (Maybe [ExpReturns])
forall b c a. (b -> c) -> (a -> b) -> a -> c
forall {k} (cat :: k -> k -> *) (b :: k) (c :: k) (a :: k).
Category cat =>
cat b c -> cat a b -> cat a c
. [ExpReturns] -> Maybe [ExpReturns]
forall a. a -> Maybe a
Just ([ExpReturns] -> m (Maybe [ExpReturns]))
-> [ExpReturns] -> m (Maybe [ExpReturns])
forall a b. (a -> b) -> a -> b
$
        [ PrimType
-> ShapeBase (Ext SubExp)
-> NoMode
-> Maybe MemReturn
-> ExpReturns
forall d o ret.
PrimType -> ShapeBase d -> o -> ret -> MemInfo d o ret
MemArray PrimType
et (SubExp -> Ext SubExp
forall a. a -> Ext a
Free (SubExp -> Ext SubExp)
-> ShapeBase SubExp -> ShapeBase (Ext SubExp)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> NewShape SubExp -> ShapeBase SubExp
forall d. NewShape d -> ShapeBase d
newShape NewShape SubExp
newshape) NoMode
NoMode (Maybe MemReturn -> ExpReturns)
-> (MemReturn -> Maybe MemReturn) -> MemReturn -> ExpReturns
forall b c a. (b -> c) -> (a -> b) -> a -> c
forall {k} (cat :: k -> k -> *) (b :: k) (c :: k) (a :: k).
Category cat =>
cat b c -> cat a b -> cat a c
. MemReturn -> Maybe MemReturn
forall a. a -> Maybe a
Just (MemReturn -> ExpReturns) -> MemReturn -> ExpReturns
forall a b. (a -> b) -> a -> b
$
            VName -> ExtLMAD -> MemReturn
ReturnsInBlock VName
mem ([VName] -> LMAD -> ExtLMAD
existentialiseLMAD [] LMAD
lmad')
        ]
    Maybe LMAD
Nothing -> Maybe [ExpReturns] -> m (Maybe [ExpReturns])
forall a. a -> m a
forall (f :: * -> *) a. Applicative f => a -> f a
pure Maybe [ExpReturns]
forall a. Maybe a
Nothing
  where
    reshaper :: ReshapeKind -> LMAD num -> Shape num -> Maybe (LMAD num)
reshaper ReshapeKind
ReshapeArbitrary LMAD num
lmad =
      LMAD num -> Shape num -> Maybe (LMAD num)
forall num.
(Eq num, IntegralExp num) =>
LMAD num -> Shape num -> Maybe (LMAD num)
LMAD.reshape LMAD num
lmad
    reshaper ReshapeKind
ReshapeCoerce LMAD num
lmad =
      LMAD num -> Maybe (LMAD num)
forall a. a -> Maybe a
Just (LMAD num -> Maybe (LMAD num))
-> (Shape num -> LMAD num) -> Shape num -> Maybe (LMAD num)
forall b c a. (b -> c) -> (a -> b) -> a -> c
forall {k} (cat :: k -> k -> *) (b :: k) (c :: k) (a :: k).
Category cat =>
cat b c -> cat a b -> cat a c
. LMAD num -> Shape num -> LMAD num
forall num. LMAD num -> Shape num -> LMAD num
LMAD.coerce LMAD num
lmad
expReturns (BasicOp (Rearrange VName
v [Int]
perm)) = do
  (et, Shape dims, mem, lmad) <- VName -> m (PrimType, ShapeBase SubExp, VName, LMAD)
forall rep (m :: * -> *) (inner :: * -> *).
(HasScope rep m, Monad m, Mem rep inner) =>
VName -> m (PrimType, ShapeBase SubExp, VName, LMAD)
arrayVarReturns VName
v
  let lmad' = LMAD -> [Int] -> LMAD
forall num. LMAD num -> [Int] -> LMAD num
LMAD.permute LMAD
lmad [Int]
perm
      dims' = [Int] -> [SubExp] -> [SubExp]
forall a. [Int] -> [a] -> [a]
rearrangeShape [Int]
perm [SubExp]
dims
  pure $
    Just
      [ MemArray et (Shape $ map Free dims') NoMode $
          Just $
            ReturnsInBlock mem $
              existentialiseLMAD [] lmad'
      ]
expReturns (BasicOp (Index VName
v Slice SubExp
slice)) = do
  [ExpReturns] -> Maybe [ExpReturns]
forall a. a -> Maybe a
Just ([ExpReturns] -> Maybe [ExpReturns])
-> (LetDecMem -> [ExpReturns]) -> LetDecMem -> Maybe [ExpReturns]
forall b c a. (b -> c) -> (a -> b) -> a -> c
forall {k} (cat :: k -> k -> *) (b :: k) (c :: k) (a :: k).
Category cat =>
cat b c -> cat a b -> cat a c
. ExpReturns -> [ExpReturns]
forall a. a -> [a]
forall (f :: * -> *) a. Applicative f => a -> f a
pure (ExpReturns -> [ExpReturns])
-> (LetDecMem -> ExpReturns) -> LetDecMem -> [ExpReturns]
forall b c a. (b -> c) -> (a -> b) -> a -> c
forall {k} (cat :: k -> k -> *) (b :: k) (c :: k) (a :: k).
Category cat =>
cat b c -> cat a b -> cat a c
. LetDecMem -> ExpReturns
varInfoToExpReturns (LetDecMem -> Maybe [ExpReturns])
-> m LetDecMem -> m (Maybe [ExpReturns])
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> VName -> Slice SubExp -> m LetDecMem
forall (m :: * -> *) rep (inner :: * -> *).
(Monad m, HasScope rep m, Mem rep inner) =>
VName -> Slice SubExp -> m LetDecMem
sliceInfo VName
v Slice SubExp
slice
expReturns (BasicOp (Update Safety
_ VName
v Slice SubExp
_ SubExp
_)) =
  [ExpReturns] -> Maybe [ExpReturns]
forall a. a -> Maybe a
Just ([ExpReturns] -> Maybe [ExpReturns])
-> (ExpReturns -> [ExpReturns]) -> ExpReturns -> Maybe [ExpReturns]
forall b c a. (b -> c) -> (a -> b) -> a -> c
forall {k} (cat :: k -> k -> *) (b :: k) (c :: k) (a :: k).
Category cat =>
cat b c -> cat a b -> cat a c
. ExpReturns -> [ExpReturns]
forall a. a -> [a]
forall (f :: * -> *) a. Applicative f => a -> f a
pure (ExpReturns -> Maybe [ExpReturns])
-> m ExpReturns -> m (Maybe [ExpReturns])
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> VName -> m ExpReturns
forall rep (m :: * -> *) (inner :: * -> *).
(HasScope rep m, Monad m, Mem rep inner) =>
VName -> m ExpReturns
varReturns VName
v
expReturns (BasicOp (FlatIndex VName
v FlatSlice SubExp
slice)) =
  [ExpReturns] -> Maybe [ExpReturns]
forall a. a -> Maybe a
Just ([ExpReturns] -> Maybe [ExpReturns])
-> (LetDecMem -> [ExpReturns]) -> LetDecMem -> Maybe [ExpReturns]
forall b c a. (b -> c) -> (a -> b) -> a -> c
forall {k} (cat :: k -> k -> *) (b :: k) (c :: k) (a :: k).
Category cat =>
cat b c -> cat a b -> cat a c
. ExpReturns -> [ExpReturns]
forall a. a -> [a]
forall (f :: * -> *) a. Applicative f => a -> f a
pure (ExpReturns -> [ExpReturns])
-> (LetDecMem -> ExpReturns) -> LetDecMem -> [ExpReturns]
forall b c a. (b -> c) -> (a -> b) -> a -> c
forall {k} (cat :: k -> k -> *) (b :: k) (c :: k) (a :: k).
Category cat =>
cat b c -> cat a b -> cat a c
. LetDecMem -> ExpReturns
varInfoToExpReturns (LetDecMem -> Maybe [ExpReturns])
-> m LetDecMem -> m (Maybe [ExpReturns])
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> VName -> FlatSlice SubExp -> m LetDecMem
forall (m :: * -> *) rep (inner :: * -> *).
(Monad m, HasScope rep m, Mem rep inner) =>
VName -> FlatSlice SubExp -> m LetDecMem
flatSliceInfo VName
v FlatSlice SubExp
slice
expReturns (BasicOp (FlatUpdate VName
v FlatSlice SubExp
_ VName
_)) =
  [ExpReturns] -> Maybe [ExpReturns]
forall a. a -> Maybe a
Just ([ExpReturns] -> Maybe [ExpReturns])
-> (ExpReturns -> [ExpReturns]) -> ExpReturns -> Maybe [ExpReturns]
forall b c a. (b -> c) -> (a -> b) -> a -> c
forall {k} (cat :: k -> k -> *) (b :: k) (c :: k) (a :: k).
Category cat =>
cat b c -> cat a b -> cat a c
. ExpReturns -> [ExpReturns]
forall a. a -> [a]
forall (f :: * -> *) a. Applicative f => a -> f a
pure (ExpReturns -> Maybe [ExpReturns])
-> m ExpReturns -> m (Maybe [ExpReturns])
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> VName -> m ExpReturns
forall rep (m :: * -> *) (inner :: * -> *).
(HasScope rep m, Monad m, Mem rep inner) =>
VName -> m ExpReturns
varReturns VName
v
expReturns (BasicOp BasicOp
op) =
  [ExpReturns] -> Maybe [ExpReturns]
forall a. a -> Maybe a
Just ([ExpReturns] -> Maybe [ExpReturns])
-> ([Type] -> [ExpReturns]) -> [Type] -> Maybe [ExpReturns]
forall b c a. (b -> c) -> (a -> b) -> a -> c
forall {k} (cat :: k -> k -> *) (b :: k) (c :: k) (a :: k).
Category cat =>
cat b c -> cat a b -> cat a c
. [ExtType] -> [ExpReturns]
extReturns ([ExtType] -> [ExpReturns])
-> ([Type] -> [ExtType]) -> [Type] -> [ExpReturns]
forall b c a. (b -> c) -> (a -> b) -> a -> c
forall {k} (cat :: k -> k -> *) (b :: k) (c :: k) (a :: k).
Category cat =>
cat b c -> cat a b -> cat a c
. [Type] -> [ExtType]
forall o.
[TypeBase (ShapeBase SubExp) o]
-> [TypeBase (ShapeBase (Ext SubExp)) o]
staticShapes ([Type] -> Maybe [ExpReturns])
-> m [Type] -> m (Maybe [ExpReturns])
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> BasicOp -> m [Type]
forall rep (m :: * -> *). HasScope rep m => BasicOp -> m [Type]
basicOpType BasicOp
op
expReturns e :: Exp rep
e@(Loop [(FParam rep, SubExp)]
merge LoopForm
_ Body rep
_) = do
  t <- Exp rep -> m [ExtType]
forall rep (m :: * -> *).
(HasScope rep m, TypedOp (OpC rep)) =>
Exp rep -> m [ExtType]
expExtType Exp rep
e
  Just <$> zipWithM typeWithDec t (map fst merge)
  where
    typeWithDec :: ExtType -> Param FParamMem -> m ExpReturns
typeWithDec ExtType
t Param FParamMem
p =
      case (ExtType
t, Param FParamMem -> FParamMem
forall dec. Param dec -> dec
paramDec Param FParamMem
p) of
        ( Array PrimType
pt ShapeBase (Ext SubExp)
shape NoMode
o,
          MemArray PrimType
_ ShapeBase SubExp
_ Diet
_ (ArrayIn VName
mem LMAD
lmad)
          )
            | Just (Int
i, Param FParamMem
mem_p) <- VName -> Maybe (Int, Param FParamMem)
isLoopVar VName
mem,
              Mem Space
space <- Param FParamMem -> Type
forall dec. Typed dec => Param dec -> Type
paramType Param FParamMem
mem_p ->
                ExpReturns -> m ExpReturns
forall a. a -> m a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (ExpReturns -> m ExpReturns) -> ExpReturns -> m ExpReturns
forall a b. (a -> b) -> a -> b
$ PrimType
-> ShapeBase (Ext SubExp)
-> NoMode
-> Maybe MemReturn
-> ExpReturns
forall d o ret.
PrimType -> ShapeBase d -> o -> ret -> MemInfo d o ret
MemArray PrimType
pt ShapeBase (Ext SubExp)
shape NoMode
o (Maybe MemReturn -> ExpReturns) -> Maybe MemReturn -> ExpReturns
forall a b. (a -> b) -> a -> b
$ MemReturn -> Maybe MemReturn
forall a. a -> Maybe a
Just (MemReturn -> Maybe MemReturn) -> MemReturn -> Maybe MemReturn
forall a b. (a -> b) -> a -> b
$ Space -> Int -> ExtLMAD -> MemReturn
ReturnsNewBlock Space
space Int
i ExtLMAD
lmad'
            | Bool
otherwise ->
                ExpReturns -> m ExpReturns
forall a. a -> m a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (ExpReturns -> m ExpReturns) -> ExpReturns -> m ExpReturns
forall a b. (a -> b) -> a -> b
$ PrimType
-> ShapeBase (Ext SubExp)
-> NoMode
-> Maybe MemReturn
-> ExpReturns
forall d o ret.
PrimType -> ShapeBase d -> o -> ret -> MemInfo d o ret
MemArray PrimType
pt ShapeBase (Ext SubExp)
shape NoMode
o (Maybe MemReturn -> ExpReturns) -> Maybe MemReturn -> ExpReturns
forall a b. (a -> b) -> a -> b
$ MemReturn -> Maybe MemReturn
forall a. a -> Maybe a
Just (MemReturn -> Maybe MemReturn) -> MemReturn -> Maybe MemReturn
forall a b. (a -> b) -> a -> b
$ VName -> ExtLMAD -> MemReturn
ReturnsInBlock VName
mem ExtLMAD
lmad'
            where
              lmad' :: ExtLMAD
lmad' = [VName] -> LMAD -> ExtLMAD
existentialiseLMAD ((Param FParamMem -> VName) -> [Param FParamMem] -> [VName]
forall a b. (a -> b) -> [a] -> [b]
map Param FParamMem -> VName
forall dec. Param dec -> VName
paramName [Param FParamMem]
mergevars) LMAD
lmad
        (Array {}, FParamMem
_) ->
          String -> m ExpReturns
forall a. HasCallStack => String -> a
error String
"expReturns: Array return type but not array merge variable."
        (Acc VName
acc ShapeBase SubExp
ispace [Type]
ts, FParamMem
_) ->
          ExpReturns -> m ExpReturns
forall a. a -> m a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (ExpReturns -> m ExpReturns) -> ExpReturns -> m ExpReturns
forall a b. (a -> b) -> a -> b
$ VName -> ShapeBase SubExp -> [Type] -> ExpReturns
forall d o ret.
VName -> ShapeBase SubExp -> [Type] -> MemInfo d o ret
MemAcc VName
acc ShapeBase SubExp
ispace [Type]
ts
        (Prim PrimType
pt, FParamMem
_) ->
          ExpReturns -> m ExpReturns
forall a. a -> m a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (ExpReturns -> m ExpReturns) -> ExpReturns -> m ExpReturns
forall a b. (a -> b) -> a -> b
$ PrimType -> ExpReturns
forall d o ret. PrimType -> MemInfo d o ret
MemPrim PrimType
pt
        (Mem Space
space, FParamMem
_) ->
          ExpReturns -> m ExpReturns
forall a. a -> m a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (ExpReturns -> m ExpReturns) -> ExpReturns -> m ExpReturns
forall a b. (a -> b) -> a -> b
$ Space -> ExpReturns
forall d o ret. Space -> MemInfo d o ret
MemMem Space
space
    isLoopVar :: VName -> Maybe (Int, Param FParamMem)
isLoopVar VName
v = ((Int, Param FParamMem) -> Bool)
-> [(Int, Param FParamMem)] -> Maybe (Int, Param FParamMem)
forall (t :: * -> *) a. Foldable t => (a -> Bool) -> t a -> Maybe a
find ((VName -> VName -> Bool
forall a. Eq a => a -> a -> Bool
== VName
v) (VName -> Bool)
-> ((Int, Param FParamMem) -> VName)
-> (Int, Param FParamMem)
-> Bool
forall b c a. (b -> c) -> (a -> b) -> a -> c
forall {k} (cat :: k -> k -> *) (b :: k) (c :: k) (a :: k).
Category cat =>
cat b c -> cat a b -> cat a c
. Param FParamMem -> VName
forall dec. Param dec -> VName
paramName (Param FParamMem -> VName)
-> ((Int, Param FParamMem) -> Param FParamMem)
-> (Int, Param FParamMem)
-> VName
forall b c a. (b -> c) -> (a -> b) -> a -> c
forall {k} (cat :: k -> k -> *) (b :: k) (c :: k) (a :: k).
Category cat =>
cat b c -> cat a b -> cat a c
. (Int, Param FParamMem) -> Param FParamMem
forall a b. (a, b) -> b
snd) ([(Int, Param FParamMem)] -> Maybe (Int, Param FParamMem))
-> [(Int, Param FParamMem)] -> Maybe (Int, Param FParamMem)
forall a b. (a -> b) -> a -> b
$ [Int] -> [Param FParamMem] -> [(Int, Param FParamMem)]
forall a b. [a] -> [b] -> [(a, b)]
zip [Int
0 ..] [Param FParamMem]
mergevars
    mergevars :: [Param FParamMem]
mergevars = ((Param FParamMem, SubExp) -> Param FParamMem)
-> [(Param FParamMem, SubExp)] -> [Param FParamMem]
forall a b. (a -> b) -> [a] -> [b]
map (Param FParamMem, SubExp) -> Param FParamMem
forall a b. (a, b) -> a
fst [(FParam rep, SubExp)]
[(Param FParamMem, SubExp)]
merge
expReturns (Apply Name
_ [(SubExp, Diet)]
_ [(RetType rep, RetAls)]
ret Safety
_) =
  Maybe [ExpReturns] -> m (Maybe [ExpReturns])
forall a. a -> m a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (Maybe [ExpReturns] -> m (Maybe [ExpReturns]))
-> Maybe [ExpReturns] -> m (Maybe [ExpReturns])
forall a b. (a -> b) -> a -> b
$ [ExpReturns] -> Maybe [ExpReturns]
forall a. a -> Maybe a
Just ([ExpReturns] -> Maybe [ExpReturns])
-> [ExpReturns] -> Maybe [ExpReturns]
forall a b. (a -> b) -> a -> b
$ ((FunReturns, RetAls) -> ExpReturns)
-> [(FunReturns, RetAls)] -> [ExpReturns]
forall a b. (a -> b) -> [a] -> [b]
map (FunReturns -> ExpReturns
funReturnsToExpReturns (FunReturns -> ExpReturns)
-> ((FunReturns, RetAls) -> FunReturns)
-> (FunReturns, RetAls)
-> ExpReturns
forall b c a. (b -> c) -> (a -> b) -> a -> c
forall {k} (cat :: k -> k -> *) (b :: k) (c :: k) (a :: k).
Category cat =>
cat b c -> cat a b -> cat a c
. (FunReturns, RetAls) -> FunReturns
forall a b. (a, b) -> a
fst) [(RetType rep, RetAls)]
[(FunReturns, RetAls)]
ret
expReturns (Match [SubExp]
_ [Case (Body rep)]
_ Body rep
_ (MatchDec [BranchType rep]
ret MatchSort
_)) =
  Maybe [ExpReturns] -> m (Maybe [ExpReturns])
forall a. a -> m a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (Maybe [ExpReturns] -> m (Maybe [ExpReturns]))
-> Maybe [ExpReturns] -> m (Maybe [ExpReturns])
forall a b. (a -> b) -> a -> b
$ [ExpReturns] -> Maybe [ExpReturns]
forall a. a -> Maybe a
Just ([ExpReturns] -> Maybe [ExpReturns])
-> [ExpReturns] -> Maybe [ExpReturns]
forall a b. (a -> b) -> a -> b
$ (FunReturns -> ExpReturns) -> [FunReturns] -> [ExpReturns]
forall a b. (a -> b) -> [a] -> [b]
map FunReturns -> ExpReturns
bodyReturnsToExpReturns [BranchType rep]
[FunReturns]
ret
expReturns (Op Op rep
op) =
  [ExpReturns] -> Maybe [ExpReturns]
forall a. a -> Maybe a
Just ([ExpReturns] -> Maybe [ExpReturns])
-> m [ExpReturns] -> m (Maybe [ExpReturns])
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> MemOp inner rep -> m [ExpReturns]
forall rep (inner :: * -> *) (m :: * -> *).
(Mem rep inner, Monad m, HasScope rep m) =>
MemOp inner rep -> m [ExpReturns]
forall (op :: * -> *) rep (inner :: * -> *) (m :: * -> *).
(OpReturns op, Mem rep inner, Monad m, HasScope rep m) =>
op rep -> m [ExpReturns]
opReturns Op rep
MemOp inner rep
op
expReturns (WithAcc [WithAccInput rep]
inputs Lambda rep
lam) =
  [ExpReturns] -> Maybe [ExpReturns]
forall a. a -> Maybe a
Just
    ([ExpReturns] -> Maybe [ExpReturns])
-> m [ExpReturns] -> m (Maybe [ExpReturns])
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> ( [ExpReturns] -> [ExpReturns] -> [ExpReturns]
forall a. Semigroup a => a -> a -> a
(<>)
            ([ExpReturns] -> [ExpReturns] -> [ExpReturns])
-> m [ExpReturns] -> m ([ExpReturns] -> [ExpReturns])
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> ([[ExpReturns]] -> [ExpReturns]
forall (t :: * -> *) a. Foldable t => t [a] -> [a]
concat ([[ExpReturns]] -> [ExpReturns])
-> m [[ExpReturns]] -> m [ExpReturns]
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> (WithAccInput rep -> m [ExpReturns])
-> [WithAccInput rep] -> m [[ExpReturns]]
forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
(a -> m b) -> t a -> m (t b)
forall (m :: * -> *) a b. Monad m => (a -> m b) -> [a] -> m [b]
mapM WithAccInput rep -> m [ExpReturns]
forall {rep} {inner :: * -> *} {m :: * -> *} {t :: * -> *} {a} {c}.
(BranchType rep ~ FunReturns, LParamInfo rep ~ LetDecMem,
 FParamInfo rep ~ FParamMem, RetType rep ~ FunReturns,
 OpC rep ~ MemOp inner, HasScope rep m, Monad m, Traversable t,
 HasLetDecMem (LetDec rep), ASTRep rep, OpReturns inner,
 RephraseOp inner, Pretty (inner rep), Rename (inner rep),
 Substitute (inner rep), FreeIn (inner rep), Show (inner rep),
 Ord (inner rep)) =>
(a, t VName, c) -> m (t ExpReturns)
inputReturns [WithAccInput rep]
inputs)
            m ([ExpReturns] -> [ExpReturns])
-> m [ExpReturns] -> m [ExpReturns]
forall a b. m (a -> b) -> m a -> m b
forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*>
            -- XXX: this is a bit dubious because it enforces extra copies.  I
            -- think WithAcc should perhaps have a return annotation like If.
            [ExpReturns] -> m [ExpReturns]
forall a. a -> m a
forall (f :: * -> *) a. Applicative f => a -> f a
pure ([ExtType] -> [ExpReturns]
extReturns ([ExtType] -> [ExpReturns]) -> [ExtType] -> [ExpReturns]
forall a b. (a -> b) -> a -> b
$ [Type] -> [ExtType]
forall o.
[TypeBase (ShapeBase SubExp) o]
-> [TypeBase (ShapeBase (Ext SubExp)) o]
staticShapes ([Type] -> [ExtType]) -> [Type] -> [ExtType]
forall a b. (a -> b) -> a -> b
$ Int -> [Type] -> [Type]
forall a. Int -> [a] -> [a]
drop Int
num_accs ([Type] -> [Type]) -> [Type] -> [Type]
forall a b. (a -> b) -> a -> b
$ Lambda rep -> [Type]
forall rep t. GLambda rep t -> [t]
lambdaReturnType Lambda rep
lam)
        )
  where
    inputReturns :: (a, t VName, c) -> m (t ExpReturns)
inputReturns (a
_, t VName
arrs, c
_) = (VName -> m ExpReturns) -> t VName -> m (t ExpReturns)
forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
(a -> m b) -> t a -> m (t b)
forall (m :: * -> *) a b. Monad m => (a -> m b) -> t a -> m (t b)
mapM VName -> m ExpReturns
forall rep (m :: * -> *) (inner :: * -> *).
(HasScope rep m, Monad m, Mem rep inner) =>
VName -> m ExpReturns
varReturns t VName
arrs
    num_accs :: Int
num_accs = [WithAccInput rep] -> Int
forall a. [a] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length [WithAccInput rep]
inputs

sliceInfo ::
  (Monad m, HasScope rep m, Mem rep inner) =>
  VName ->
  Slice SubExp ->
  m (MemInfo SubExp NoMode MemBind)
sliceInfo :: forall (m :: * -> *) rep (inner :: * -> *).
(Monad m, HasScope rep m, Mem rep inner) =>
VName -> Slice SubExp -> m LetDecMem
sliceInfo VName
v Slice SubExp
slice = do
  (et, _, mem, lmad) <- VName -> m (PrimType, ShapeBase SubExp, VName, LMAD)
forall rep (m :: * -> *) (inner :: * -> *).
(HasScope rep m, Monad m, Mem rep inner) =>
VName -> m (PrimType, ShapeBase SubExp, VName, LMAD)
arrayVarReturns VName
v
  case sliceDims slice of
    [] -> LetDecMem -> m LetDecMem
forall a. a -> m a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (LetDecMem -> m LetDecMem) -> LetDecMem -> m LetDecMem
forall a b. (a -> b) -> a -> b
$ PrimType -> LetDecMem
forall d o ret. PrimType -> MemInfo d o ret
MemPrim PrimType
et
    [SubExp]
dims ->
      LetDecMem -> m LetDecMem
forall a. a -> m a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (LetDecMem -> m LetDecMem) -> LetDecMem -> m LetDecMem
forall a b. (a -> b) -> a -> b
$
        PrimType -> ShapeBase SubExp -> NoMode -> MemBind -> LetDecMem
forall d o ret.
PrimType -> ShapeBase d -> o -> ret -> MemInfo d o ret
MemArray PrimType
et ([SubExp] -> ShapeBase SubExp
forall d. [d] -> ShapeBase d
Shape [SubExp]
dims) NoMode
NoMode (MemBind -> LetDecMem) -> (LMAD -> MemBind) -> LMAD -> LetDecMem
forall b c a. (b -> c) -> (a -> b) -> a -> c
forall {k} (cat :: k -> k -> *) (b :: k) (c :: k) (a :: k).
Category cat =>
cat b c -> cat a b -> cat a c
. VName -> LMAD -> MemBind
ArrayIn VName
mem (LMAD -> LetDecMem) -> LMAD -> LetDecMem
forall a b. (a -> b) -> a -> b
$
          LMAD -> Slice (TPrimExp Int64 VName) -> LMAD
forall num.
(Eq num, IntegralExp num) =>
LMAD num -> Slice num -> LMAD num
LMAD.slice LMAD
lmad ((SubExp -> TPrimExp Int64 VName)
-> Slice SubExp -> Slice (TPrimExp Int64 VName)
forall a b. (a -> b) -> Slice a -> Slice b
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap SubExp -> TPrimExp Int64 VName
pe64 Slice SubExp
slice)

flatSliceInfo ::
  (Monad m, HasScope rep m, Mem rep inner) =>
  VName ->
  FlatSlice SubExp ->
  m (MemInfo SubExp NoMode MemBind)
flatSliceInfo :: forall (m :: * -> *) rep (inner :: * -> *).
(Monad m, HasScope rep m, Mem rep inner) =>
VName -> FlatSlice SubExp -> m LetDecMem
flatSliceInfo VName
v slice :: FlatSlice SubExp
slice@(FlatSlice SubExp
offset [FlatDimIndex SubExp]
idxs) = do
  (et, _, mem, lmad) <- VName -> m (PrimType, ShapeBase SubExp, VName, LMAD)
forall rep (m :: * -> *) (inner :: * -> *).
(HasScope rep m, Monad m, Mem rep inner) =>
VName -> m (PrimType, ShapeBase SubExp, VName, LMAD)
arrayVarReturns VName
v
  map (fmap pe64) idxs
    & FlatSlice (pe64 offset)
    & LMAD.flatSlice lmad
    & MemArray et (Shape (flatSliceDims slice)) NoMode . ArrayIn mem
    & pure

class (IsOp op) => OpReturns op where
  opReturns :: (Mem rep inner, Monad m, HasScope rep m) => op rep -> m [ExpReturns]
  opReturns op rep
op = [ExtType] -> [ExpReturns]
extReturns ([ExtType] -> [ExpReturns]) -> m [ExtType] -> m [ExpReturns]
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> op rep -> m [ExtType]
forall rep (m :: * -> *). HasScope rep m => op rep -> m [ExtType]
forall (op :: * -> *) rep (m :: * -> *).
(TypedOp op, HasScope rep m) =>
op rep -> m [ExtType]
opType op rep
op

instance (OpReturns inner) => OpReturns (MemOp inner) where
  opReturns :: forall rep (inner :: * -> *) (m :: * -> *).
(Mem rep inner, Monad m, HasScope rep m) =>
MemOp inner rep -> m [ExpReturns]
opReturns (Alloc SubExp
_ Space
space) = [ExpReturns] -> m [ExpReturns]
forall a. a -> m a
forall (f :: * -> *) a. Applicative f => a -> f a
pure [Space -> ExpReturns
forall d o ret. Space -> MemInfo d o ret
MemMem Space
space]
  opReturns (EnsureDirect VName
v) = do
    space <- VName -> m Space
forall rep (inner :: * -> *) (m :: * -> *).
(Mem rep inner, HasScope rep m, Monad m) =>
VName -> m Space
lookupMemSpace (VName -> m Space)
-> ((VName, LMAD) -> VName) -> (VName, LMAD) -> m Space
forall b c a. (b -> c) -> (a -> b) -> a -> c
forall {k} (cat :: k -> k -> *) (b :: k) (c :: k) (a :: k).
Category cat =>
cat b c -> cat a b -> cat a c
. (VName, LMAD) -> VName
forall a b. (a, b) -> a
fst ((VName, LMAD) -> m Space) -> m (VName, LMAD) -> m Space
forall (m :: * -> *) a b. Monad m => (a -> m b) -> m a -> m b
=<< VName -> m (VName, LMAD)
forall rep (inner :: * -> *) (m :: * -> *).
(Mem rep inner, HasScope rep m, Monad m) =>
VName -> m (VName, LMAD)
lookupArraySummary VName
v
    summary <- lookupMemInfo v
    case summary of
      MemArray PrimType
et ShapeBase SubExp
shape NoMode
_ MemBind
_ ->
        [ExpReturns] -> m [ExpReturns]
forall a. a -> m a
forall (f :: * -> *) a. Applicative f => a -> f a
pure
          [ Space -> ExpReturns
forall d o ret. Space -> MemInfo d o ret
MemMem Space
space,
            PrimType
-> ShapeBase (Ext SubExp)
-> NoMode
-> Maybe MemReturn
-> ExpReturns
forall d o ret.
PrimType -> ShapeBase d -> o -> ret -> MemInfo d o ret
MemArray PrimType
et ((SubExp -> Ext SubExp)
-> ShapeBase SubExp -> ShapeBase (Ext SubExp)
forall a b. (a -> b) -> ShapeBase a -> ShapeBase b
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap SubExp -> Ext SubExp
forall a. a -> Ext a
Free ShapeBase SubExp
shape) NoMode
NoMode (Maybe MemReturn -> ExpReturns)
-> (MemReturn -> Maybe MemReturn) -> MemReturn -> ExpReturns
forall b c a. (b -> c) -> (a -> b) -> a -> c
forall {k} (cat :: k -> k -> *) (b :: k) (c :: k) (a :: k).
Category cat =>
cat b c -> cat a b -> cat a c
. MemReturn -> Maybe MemReturn
forall a. a -> Maybe a
Just (MemReturn -> ExpReturns) -> MemReturn -> ExpReturns
forall a b. (a -> b) -> a -> b
$
              Space -> Int -> ExtLMAD -> MemReturn
ReturnsNewBlock Space
space Int
0 (ExtLMAD -> MemReturn)
-> ([TPrimExp Int64 (Ext VName)] -> ExtLMAD)
-> [TPrimExp Int64 (Ext VName)]
-> MemReturn
forall b c a. (b -> c) -> (a -> b) -> a -> c
forall {k} (cat :: k -> k -> *) (b :: k) (c :: k) (a :: k).
Category cat =>
cat b c -> cat a b -> cat a c
. TPrimExp Int64 (Ext VName)
-> [TPrimExp Int64 (Ext VName)] -> ExtLMAD
forall num. IntegralExp num => num -> [num] -> LMAD num
LMAD.iota TPrimExp Int64 (Ext VName)
0 ([TPrimExp Int64 (Ext VName)] -> MemReturn)
-> [TPrimExp Int64 (Ext VName)] -> MemReturn
forall a b. (a -> b) -> a -> b
$
                (SubExp -> TPrimExp Int64 (Ext VName))
-> [SubExp] -> [TPrimExp Int64 (Ext VName)]
forall a b. (a -> b) -> [a] -> [b]
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap ((VName -> Ext VName)
-> TPrimExp Int64 VName -> TPrimExp Int64 (Ext VName)
forall a b. (a -> b) -> TPrimExp Int64 a -> TPrimExp Int64 b
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap VName -> Ext VName
forall a. a -> Ext a
Free (TPrimExp Int64 VName -> TPrimExp Int64 (Ext VName))
-> (SubExp -> TPrimExp Int64 VName)
-> SubExp
-> TPrimExp Int64 (Ext VName)
forall b c a. (b -> c) -> (a -> b) -> a -> c
forall {k} (cat :: k -> k -> *) (b :: k) (c :: k) (a :: k).
Category cat =>
cat b c -> cat a b -> cat a c
. SubExp -> TPrimExp Int64 VName
pe64) (ShapeBase SubExp -> [SubExp]
forall d. ShapeBase d -> [d]
shapeDims ShapeBase SubExp
shape)
          ]
      LetDecMem
_ -> String -> m [ExpReturns]
forall a. HasCallStack => String -> a
error String
"EnsureDirect applied to non-array"
  opReturns (Inner inner rep
op) = inner rep -> m [ExpReturns]
forall rep (inner :: * -> *) (m :: * -> *).
(Mem rep inner, Monad m, HasScope rep m) =>
inner rep -> m [ExpReturns]
forall (op :: * -> *) rep (inner :: * -> *) (m :: * -> *).
(OpReturns op, Mem rep inner, Monad m, HasScope rep m) =>
op rep -> m [ExpReturns]
opReturns inner rep
op

instance OpReturns NoOp where
  opReturns :: forall rep (inner :: * -> *) (m :: * -> *).
(Mem rep inner, Monad m, HasScope rep m) =>
NoOp rep -> m [ExpReturns]
opReturns NoOp rep
NoOp = [ExpReturns] -> m [ExpReturns]
forall a. a -> m a
forall (f :: * -> *) a. Applicative f => a -> f a
pure []

applyFunReturns ::
  (Typed dec) =>
  [FunReturns] ->
  [Param dec] ->
  [(SubExp, Type)] ->
  Maybe [FunReturns]
applyFunReturns :: forall dec.
Typed dec =>
[FunReturns]
-> [Param dec] -> [(SubExp, Type)] -> Maybe [FunReturns]
applyFunReturns [FunReturns]
rets [Param dec]
params [(SubExp, Type)]
args
  | Just [ExtType]
_ <- [ExtType] -> [Param dec] -> [(SubExp, Type)] -> Maybe [ExtType]
forall dec.
Typed dec =>
[ExtType] -> [Param dec] -> [(SubExp, Type)] -> Maybe [ExtType]
forall rt dec.
(IsRetType rt, Typed dec) =>
[rt] -> [Param dec] -> [(SubExp, Type)] -> Maybe [rt]
applyRetType [ExtType]
rettype [Param dec]
params [(SubExp, Type)]
args =
      [FunReturns] -> Maybe [FunReturns]
forall a. a -> Maybe a
Just ([FunReturns] -> Maybe [FunReturns])
-> [FunReturns] -> Maybe [FunReturns]
forall a b. (a -> b) -> a -> b
$ (FunReturns -> FunReturns) -> [FunReturns] -> [FunReturns]
forall a b. (a -> b) -> [a] -> [b]
map FunReturns -> FunReturns
forall {o}.
MemInfo (Ext SubExp) o MemReturn
-> MemInfo (Ext SubExp) o MemReturn
correctDims [FunReturns]
rets
  | Bool
otherwise =
      Maybe [FunReturns]
forall a. Maybe a
Nothing
  where
    rettype :: [ExtType]
rettype = (FunReturns -> ExtType) -> [FunReturns] -> [ExtType]
forall a b. (a -> b) -> [a] -> [b]
map FunReturns -> ExtType
forall t. ExtTyped t => t -> ExtType
extTypeOf [FunReturns]
rets
    parammap :: M.Map VName (SubExp, Type)
    parammap :: Map VName (SubExp, Type)
parammap =
      [(VName, (SubExp, Type))] -> Map VName (SubExp, Type)
forall k a. Ord k => [(k, a)] -> Map k a
M.fromList ([(VName, (SubExp, Type))] -> Map VName (SubExp, Type))
-> [(VName, (SubExp, Type))] -> Map VName (SubExp, Type)
forall a b. (a -> b) -> a -> b
$
        [VName] -> [(SubExp, Type)] -> [(VName, (SubExp, Type))]
forall a b. [a] -> [b] -> [(a, b)]
zip ((Param dec -> VName) -> [Param dec] -> [VName]
forall a b. (a -> b) -> [a] -> [b]
map Param dec -> VName
forall dec. Param dec -> VName
paramName [Param dec]
params) [(SubExp, Type)]
args

    substSubExp :: SubExp -> SubExp
substSubExp (Var VName
v)
      | Just (SubExp
se, Type
_) <- VName -> Map VName (SubExp, Type) -> Maybe (SubExp, Type)
forall k a. Ord k => k -> Map k a -> Maybe a
M.lookup VName
v Map VName (SubExp, Type)
parammap = SubExp
se
    substSubExp SubExp
se = SubExp
se

    correctDims :: MemInfo (Ext SubExp) o MemReturn
-> MemInfo (Ext SubExp) o MemReturn
correctDims (MemPrim PrimType
t) =
      PrimType -> MemInfo (Ext SubExp) o MemReturn
forall d o ret. PrimType -> MemInfo d o ret
MemPrim PrimType
t
    correctDims (MemMem Space
space) =
      Space -> MemInfo (Ext SubExp) o MemReturn
forall d o ret. Space -> MemInfo d o ret
MemMem Space
space
    correctDims (MemArray PrimType
et ShapeBase (Ext SubExp)
shape o
o MemReturn
memsummary) =
      PrimType
-> ShapeBase (Ext SubExp)
-> o
-> MemReturn
-> MemInfo (Ext SubExp) o MemReturn
forall d o ret.
PrimType -> ShapeBase d -> o -> ret -> MemInfo d o ret
MemArray PrimType
et (ShapeBase (Ext SubExp) -> ShapeBase (Ext SubExp)
correctShape ShapeBase (Ext SubExp)
shape) o
o (MemReturn -> MemInfo (Ext SubExp) o MemReturn)
-> MemReturn -> MemInfo (Ext SubExp) o MemReturn
forall a b. (a -> b) -> a -> b
$
        MemReturn -> MemReturn
correctSummary MemReturn
memsummary
    correctDims (MemAcc VName
acc ShapeBase SubExp
ispace [Type]
ts) =
      VName
-> ShapeBase SubExp -> [Type] -> MemInfo (Ext SubExp) o MemReturn
forall d o ret.
VName -> ShapeBase SubExp -> [Type] -> MemInfo d o ret
MemAcc VName
acc ShapeBase SubExp
ispace [Type]
ts

    correctShape :: ShapeBase (Ext SubExp) -> ShapeBase (Ext SubExp)
correctShape = [Ext SubExp] -> ShapeBase (Ext SubExp)
forall d. [d] -> ShapeBase d
Shape ([Ext SubExp] -> ShapeBase (Ext SubExp))
-> (ShapeBase (Ext SubExp) -> [Ext SubExp])
-> ShapeBase (Ext SubExp)
-> ShapeBase (Ext SubExp)
forall b c a. (b -> c) -> (a -> b) -> a -> c
forall {k} (cat :: k -> k -> *) (b :: k) (c :: k) (a :: k).
Category cat =>
cat b c -> cat a b -> cat a c
. (Ext SubExp -> Ext SubExp) -> [Ext SubExp] -> [Ext SubExp]
forall a b. (a -> b) -> [a] -> [b]
map Ext SubExp -> Ext SubExp
correctDim ([Ext SubExp] -> [Ext SubExp])
-> (ShapeBase (Ext SubExp) -> [Ext SubExp])
-> ShapeBase (Ext SubExp)
-> [Ext SubExp]
forall b c a. (b -> c) -> (a -> b) -> a -> c
forall {k} (cat :: k -> k -> *) (b :: k) (c :: k) (a :: k).
Category cat =>
cat b c -> cat a b -> cat a c
. ShapeBase (Ext SubExp) -> [Ext SubExp]
forall d. ShapeBase d -> [d]
shapeDims
    correctDim :: Ext SubExp -> Ext SubExp
correctDim (Ext Int
i) = Int -> Ext SubExp
forall a. Int -> Ext a
Ext Int
i
    correctDim (Free SubExp
se) = SubExp -> Ext SubExp
forall a. a -> Ext a
Free (SubExp -> Ext SubExp) -> SubExp -> Ext SubExp
forall a b. (a -> b) -> a -> b
$ SubExp -> SubExp
substSubExp SubExp
se

    correctSummary :: MemReturn -> MemReturn
correctSummary (ReturnsNewBlock Space
space Int
i ExtLMAD
lmad) =
      Space -> Int -> ExtLMAD -> MemReturn
ReturnsNewBlock Space
space Int
i ExtLMAD
lmad
    correctSummary (ReturnsInBlock VName
mem ExtLMAD
lmad) =
      -- FIXME: we should also do a replacement in lmad here.
      VName -> ExtLMAD -> MemReturn
ReturnsInBlock VName
mem' ExtLMAD
lmad
      where
        mem' :: VName
mem' = case VName -> Map VName (SubExp, Type) -> Maybe (SubExp, Type)
forall k a. Ord k => k -> Map k a -> Maybe a
M.lookup VName
mem Map VName (SubExp, Type)
parammap of
          Just (Var VName
v, Type
_) -> VName
v
          Maybe (SubExp, Type)
_ -> VName
mem