{-# LANGUAGE ConstraintKinds #-}
{-# LANGUAGE FlexibleContexts #-}
{-# LANGUAGE FlexibleInstances #-}
{-# LANGUAGE MultiParamTypeClasses #-}
{-# LANGUAGE OverloadedStrings #-}
{-# LANGUAGE TypeFamilies #-}
{-# LANGUAGE UndecidableInstances #-}
module Futhark.IR.Mem
( LetDecMem,
FParamMem,
LParamMem,
RetTypeMem,
BranchTypeMem,
MemOp (..),
traverseMemOpStms,
MemInfo (..),
MemBound,
MemBind (..),
MemReturn (..),
IxFun,
ExtIxFun,
isStaticIxFun,
ExpReturns,
BodyReturns,
FunReturns,
noUniquenessReturns,
bodyReturnsToExpReturns,
Mem,
HasLetDecMem (..),
OpReturns (..),
varReturns,
expReturns,
extReturns,
lookupMemInfo,
subExpMemInfo,
lookupArraySummary,
existentialiseIxFun,
matchBranchReturnType,
matchPatToExp,
matchFunctionReturnType,
matchLoopResultMem,
bodyReturnsFromPat,
checkMemInfo,
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.Except
import Control.Monad.Reader
import Control.Monad.State
import Data.Foldable (traverse_)
import Data.Function ((&))
import Data.List (elemIndex, find)
import qualified Data.Map.Strict as M
import Data.Maybe
import Futhark.Analysis.Metrics
import Futhark.Analysis.PrimExp.Convert
import Futhark.Analysis.PrimExp.Simplify
import qualified Futhark.Analysis.SymbolTable as ST
import Futhark.IR.Aliases
( Aliases,
removeExpAliases,
removePatAliases,
removeScopeAliases,
)
import qualified Futhark.IR.Mem.IxFun as IxFun
import Futhark.IR.Pretty
import Futhark.IR.Prop
import Futhark.IR.Prop.Aliases
import Futhark.IR.Syntax
import Futhark.IR.Traversals
import qualified Futhark.IR.TypeCheck as TC
import qualified Futhark.Optimise.Simplify.Engine as Engine
import Futhark.Optimise.Simplify.Rep
import Futhark.Transform.Rename
import Futhark.Transform.Substitute
import Futhark.Util
import Futhark.Util.Pretty (indent, ppr, text, (<+>), (</>))
import qualified Futhark.Util.Pretty as PP
import Prelude hiding (id, (.))
type LetDecMem = MemInfo SubExp NoUniqueness MemBind
type FParamMem = MemInfo SubExp Uniqueness MemBind
type LParamMem = MemInfo SubExp NoUniqueness MemBind
type RetTypeMem = FunReturns
type BranchTypeMem = BodyReturns
class HasLetDecMem t where
letDecMem :: t -> LetDecMem
instance HasLetDecMem LetDecMem where
letDecMem :: LetDecMem -> LetDecMem
letDecMem = LetDecMem -> LetDecMem
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 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,
Op rep ~ MemOp inner
)
instance IsRetType FunReturns where
primRetType :: PrimType -> FunReturns
primRetType = PrimType -> FunReturns
forall d u ret. PrimType -> MemInfo d u ret
MemPrim
applyRetType :: [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 -> BodyReturns
primBodyType = PrimType -> BodyReturns
forall d u ret. PrimType -> MemInfo d u ret
MemPrim
data MemOp inner
=
Alloc SubExp Space
| Inner inner
deriving (MemOp inner -> MemOp inner -> Bool
(MemOp inner -> MemOp inner -> Bool)
-> (MemOp inner -> MemOp inner -> Bool) -> Eq (MemOp inner)
forall inner. Eq inner => MemOp inner -> MemOp inner -> Bool
forall a. (a -> a -> Bool) -> (a -> a -> Bool) -> Eq a
/= :: MemOp inner -> MemOp inner -> Bool
$c/= :: forall inner. Eq inner => MemOp inner -> MemOp inner -> Bool
== :: MemOp inner -> MemOp inner -> Bool
$c== :: forall inner. Eq inner => MemOp inner -> MemOp inner -> Bool
Eq, Eq (MemOp inner)
Eq (MemOp inner)
-> (MemOp inner -> MemOp inner -> Ordering)
-> (MemOp inner -> MemOp inner -> Bool)
-> (MemOp inner -> MemOp inner -> Bool)
-> (MemOp inner -> MemOp inner -> Bool)
-> (MemOp inner -> MemOp inner -> Bool)
-> (MemOp inner -> MemOp inner -> MemOp inner)
-> (MemOp inner -> MemOp inner -> MemOp inner)
-> Ord (MemOp inner)
MemOp inner -> MemOp inner -> Bool
MemOp inner -> MemOp inner -> Ordering
MemOp inner -> MemOp inner -> MemOp inner
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. Ord inner => Eq (MemOp inner)
forall inner. Ord inner => MemOp inner -> MemOp inner -> Bool
forall inner. Ord inner => MemOp inner -> MemOp inner -> Ordering
forall inner.
Ord inner =>
MemOp inner -> MemOp inner -> MemOp inner
min :: MemOp inner -> MemOp inner -> MemOp inner
$cmin :: forall inner.
Ord inner =>
MemOp inner -> MemOp inner -> MemOp inner
max :: MemOp inner -> MemOp inner -> MemOp inner
$cmax :: forall inner.
Ord inner =>
MemOp inner -> MemOp inner -> MemOp inner
>= :: MemOp inner -> MemOp inner -> Bool
$c>= :: forall inner. Ord inner => MemOp inner -> MemOp inner -> Bool
> :: MemOp inner -> MemOp inner -> Bool
$c> :: forall inner. Ord inner => MemOp inner -> MemOp inner -> Bool
<= :: MemOp inner -> MemOp inner -> Bool
$c<= :: forall inner. Ord inner => MemOp inner -> MemOp inner -> Bool
< :: MemOp inner -> MemOp inner -> Bool
$c< :: forall inner. Ord inner => MemOp inner -> MemOp inner -> Bool
compare :: MemOp inner -> MemOp inner -> Ordering
$ccompare :: forall inner. Ord inner => MemOp inner -> MemOp inner -> Ordering
$cp1Ord :: forall inner. Ord inner => Eq (MemOp inner)
Ord, Int -> MemOp inner -> ShowS
[MemOp inner] -> ShowS
MemOp inner -> String
(Int -> MemOp inner -> ShowS)
-> (MemOp inner -> String)
-> ([MemOp inner] -> ShowS)
-> Show (MemOp inner)
forall inner. Show inner => Int -> MemOp inner -> ShowS
forall inner. Show inner => [MemOp inner] -> ShowS
forall inner. Show inner => MemOp inner -> String
forall a.
(Int -> a -> ShowS) -> (a -> String) -> ([a] -> ShowS) -> Show a
showList :: [MemOp inner] -> ShowS
$cshowList :: forall inner. Show inner => [MemOp inner] -> ShowS
show :: MemOp inner -> String
$cshow :: forall inner. Show inner => MemOp inner -> String
showsPrec :: Int -> MemOp inner -> ShowS
$cshowsPrec :: forall inner. Show inner => Int -> MemOp inner -> ShowS
Show)
traverseMemOpStms ::
Monad m =>
OpStmsTraverser m inner rep ->
OpStmsTraverser m (MemOp inner) rep
traverseMemOpStms :: OpStmsTraverser m inner rep -> OpStmsTraverser m (MemOp inner) rep
traverseMemOpStms OpStmsTraverser m inner rep
_ Scope rep -> Stms rep -> m (Stms rep)
_ op :: MemOp inner
op@Alloc {} = MemOp inner -> m (MemOp inner)
forall (f :: * -> *) a. Applicative f => a -> f a
pure MemOp inner
op
traverseMemOpStms OpStmsTraverser m inner rep
onInner Scope rep -> Stms rep -> m (Stms rep)
f (Inner inner
inner) = inner -> MemOp inner
forall inner. inner -> MemOp inner
Inner (inner -> MemOp inner) -> m inner -> m (MemOp inner)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> OpStmsTraverser m inner rep
onInner Scope rep -> Stms rep -> m (Stms rep)
f inner
inner
instance FreeIn inner => FreeIn (MemOp inner) where
freeIn' :: MemOp inner -> FV
freeIn' (Alloc SubExp
size Space
_) = SubExp -> FV
forall a. FreeIn a => a -> FV
freeIn' SubExp
size
freeIn' (Inner inner
k) = inner -> FV
forall a. FreeIn a => a -> FV
freeIn' inner
k
instance TypedOp inner => TypedOp (MemOp inner) where
opType :: MemOp inner -> m [ExtType]
opType (Alloc SubExp
_ Space
space) = [ExtType] -> m [ExtType]
forall (f :: * -> *) a. Applicative f => a -> f a
pure [Space -> ExtType
forall shape u. Space -> TypeBase shape u
Mem Space
space]
opType (Inner inner
k) = inner -> m [ExtType]
forall op t (m :: * -> *).
(TypedOp op, HasScope t m) =>
op -> m [ExtType]
opType inner
k
instance AliasedOp inner => AliasedOp (MemOp inner) where
opAliases :: MemOp inner -> [Names]
opAliases Alloc {} = [Names
forall a. Monoid a => a
mempty]
opAliases (Inner inner
k) = inner -> [Names]
forall op. AliasedOp op => op -> [Names]
opAliases inner
k
consumedInOp :: MemOp inner -> Names
consumedInOp Alloc {} = Names
forall a. Monoid a => a
mempty
consumedInOp (Inner inner
k) = inner -> Names
forall op. AliasedOp op => op -> Names
consumedInOp inner
k
instance CanBeAliased inner => CanBeAliased (MemOp inner) where
type OpWithAliases (MemOp inner) = MemOp (OpWithAliases inner)
removeOpAliases :: OpWithAliases (MemOp inner) -> MemOp inner
removeOpAliases (Alloc se space) = SubExp -> Space -> MemOp inner
forall inner. SubExp -> Space -> MemOp inner
Alloc SubExp
se Space
space
removeOpAliases (Inner k) = inner -> MemOp inner
forall inner. inner -> MemOp inner
Inner (inner -> MemOp inner) -> inner -> MemOp inner
forall a b. (a -> b) -> a -> b
$ OpWithAliases inner -> inner
forall op. CanBeAliased op => OpWithAliases op -> op
removeOpAliases OpWithAliases inner
k
addOpAliases :: AliasTable -> MemOp inner -> OpWithAliases (MemOp inner)
addOpAliases AliasTable
_ (Alloc SubExp
se Space
space) = SubExp -> Space -> MemOp (OpWithAliases inner)
forall inner. SubExp -> Space -> MemOp inner
Alloc SubExp
se Space
space
addOpAliases AliasTable
aliases (Inner inner
k) = OpWithAliases inner -> MemOp (OpWithAliases inner)
forall inner. inner -> MemOp inner
Inner (OpWithAliases inner -> MemOp (OpWithAliases inner))
-> OpWithAliases inner -> MemOp (OpWithAliases inner)
forall a b. (a -> b) -> a -> b
$ AliasTable -> inner -> OpWithAliases inner
forall op. CanBeAliased op => AliasTable -> op -> OpWithAliases op
addOpAliases AliasTable
aliases inner
k
instance Rename inner => Rename (MemOp inner) where
rename :: MemOp inner -> RenameM (MemOp inner)
rename (Alloc SubExp
size Space
space) = SubExp -> Space -> MemOp inner
forall inner. SubExp -> Space -> MemOp inner
Alloc (SubExp -> Space -> MemOp inner)
-> RenameM SubExp -> RenameM (Space -> MemOp inner)
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)
-> RenameM Space -> RenameM (MemOp inner)
forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> Space -> RenameM Space
forall (f :: * -> *) a. Applicative f => a -> f a
pure Space
space
rename (Inner inner
k) = inner -> MemOp inner
forall inner. inner -> MemOp inner
Inner (inner -> MemOp inner) -> RenameM inner -> RenameM (MemOp inner)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> inner -> RenameM inner
forall a. Rename a => a -> RenameM a
rename inner
k
instance Substitute inner => Substitute (MemOp inner) where
substituteNames :: Map VName VName -> MemOp inner -> MemOp inner
substituteNames Map VName VName
subst (Alloc SubExp
size Space
space) = SubExp -> Space -> MemOp inner
forall inner. SubExp -> Space -> MemOp inner
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 (Inner inner
k) = inner -> MemOp inner
forall inner. inner -> MemOp inner
Inner (inner -> MemOp inner) -> inner -> MemOp inner
forall a b. (a -> b) -> a -> b
$ Map VName VName -> inner -> inner
forall a. Substitute a => Map VName VName -> a -> a
substituteNames Map VName VName
subst inner
k
instance PP.Pretty inner => PP.Pretty (MemOp inner) where
ppr :: MemOp inner -> Doc
ppr (Alloc SubExp
e Space
DefaultSpace) = String -> Doc
PP.text String
"alloc" Doc -> Doc -> Doc
forall a. Semigroup a => a -> a -> a
<> [Doc] -> Doc
PP.apply [SubExp -> Doc
forall a. Pretty a => a -> Doc
PP.ppr SubExp
e]
ppr (Alloc SubExp
e Space
s) = String -> Doc
PP.text String
"alloc" Doc -> Doc -> Doc
forall a. Semigroup a => a -> a -> a
<> [Doc] -> Doc
PP.apply [SubExp -> Doc
forall a. Pretty a => a -> Doc
PP.ppr SubExp
e, Space -> Doc
forall a. Pretty a => a -> Doc
PP.ppr Space
s]
ppr (Inner inner
k) = inner -> Doc
forall a. Pretty a => a -> Doc
PP.ppr inner
k
instance OpMetrics inner => OpMetrics (MemOp inner) where
opMetrics :: MemOp inner -> MetricsM ()
opMetrics Alloc {} = Text -> MetricsM ()
seen Text
"Alloc"
opMetrics (Inner inner
k) = inner -> MetricsM ()
forall op. OpMetrics op => op -> MetricsM ()
opMetrics inner
k
instance IsOp inner => IsOp (MemOp inner) where
safeOp :: MemOp inner -> 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 (Inner inner
k) = inner -> Bool
forall op. IsOp op => op -> Bool
safeOp inner
k
cheapOp :: MemOp inner -> Bool
cheapOp (Inner inner
k) = inner -> Bool
forall op. IsOp op => op -> Bool
cheapOp inner
k
cheapOp Alloc {} = Bool
True
instance CanBeWise inner => CanBeWise (MemOp inner) where
type OpWithWisdom (MemOp inner) = MemOp (OpWithWisdom inner)
removeOpWisdom :: OpWithWisdom (MemOp inner) -> MemOp inner
removeOpWisdom (Alloc size space) = SubExp -> Space -> MemOp inner
forall inner. SubExp -> Space -> MemOp inner
Alloc SubExp
size Space
space
removeOpWisdom (Inner k) = inner -> MemOp inner
forall inner. inner -> MemOp inner
Inner (inner -> MemOp inner) -> inner -> MemOp inner
forall a b. (a -> b) -> a -> b
$ OpWithWisdom inner -> inner
forall op. CanBeWise op => OpWithWisdom op -> op
removeOpWisdom OpWithWisdom inner
k
addOpWisdom :: MemOp inner -> OpWithWisdom (MemOp inner)
addOpWisdom (Alloc SubExp
size Space
space) = SubExp -> Space -> MemOp (OpWithWisdom inner)
forall inner. SubExp -> Space -> MemOp inner
Alloc SubExp
size Space
space
addOpWisdom (Inner inner
k) = OpWithWisdom inner -> MemOp (OpWithWisdom inner)
forall inner. inner -> MemOp inner
Inner (OpWithWisdom inner -> MemOp (OpWithWisdom inner))
-> OpWithWisdom inner -> MemOp (OpWithWisdom inner)
forall a b. (a -> b) -> a -> b
$ inner -> OpWithWisdom inner
forall op. CanBeWise op => op -> OpWithWisdom op
addOpWisdom inner
k
instance ST.IndexOp inner => ST.IndexOp (MemOp inner) where
indexOp :: SymbolTable rep
-> Int -> MemOp inner -> [TPrimExp Int64 VName] -> Maybe Indexed
indexOp SymbolTable rep
vtable Int
k (Inner inner
op) [TPrimExp Int64 VName]
is = SymbolTable rep
-> Int -> inner -> [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
op [TPrimExp Int64 VName]
is
indexOp SymbolTable rep
_ Int
_ MemOp inner
_ [TPrimExp Int64 VName]
_ = Maybe Indexed
forall a. Maybe a
Nothing
type IxFun = IxFun.IxFun (TPrimExp Int64 VName)
type ExtIxFun = IxFun.IxFun (TPrimExp Int64 (Ext VName))
data MemInfo d u ret
=
MemPrim PrimType
|
MemMem Space
|
MemArray PrimType (ShapeBase d) u ret
|
MemAcc VName Shape [Type] u
deriving (MemInfo d u ret -> MemInfo d u ret -> Bool
(MemInfo d u ret -> MemInfo d u ret -> Bool)
-> (MemInfo d u ret -> MemInfo d u ret -> Bool)
-> Eq (MemInfo d u ret)
forall a. (a -> a -> Bool) -> (a -> a -> Bool) -> Eq a
forall d u ret.
(Eq d, Eq u, Eq ret) =>
MemInfo d u ret -> MemInfo d u ret -> Bool
/= :: MemInfo d u ret -> MemInfo d u ret -> Bool
$c/= :: forall d u ret.
(Eq d, Eq u, Eq ret) =>
MemInfo d u ret -> MemInfo d u ret -> Bool
== :: MemInfo d u ret -> MemInfo d u ret -> Bool
$c== :: forall d u ret.
(Eq d, Eq u, Eq ret) =>
MemInfo d u ret -> MemInfo d u ret -> Bool
Eq, Int -> MemInfo d u ret -> ShowS
[MemInfo d u ret] -> ShowS
MemInfo d u ret -> String
(Int -> MemInfo d u ret -> ShowS)
-> (MemInfo d u ret -> String)
-> ([MemInfo d u ret] -> ShowS)
-> Show (MemInfo d u ret)
forall a.
(Int -> a -> ShowS) -> (a -> String) -> ([a] -> ShowS) -> Show a
forall d u ret.
(Show d, Show u, Show ret) =>
Int -> MemInfo d u ret -> ShowS
forall d u ret.
(Show d, Show u, Show ret) =>
[MemInfo d u ret] -> ShowS
forall d u ret.
(Show d, Show u, Show ret) =>
MemInfo d u ret -> String
showList :: [MemInfo d u ret] -> ShowS
$cshowList :: forall d u ret.
(Show d, Show u, Show ret) =>
[MemInfo d u ret] -> ShowS
show :: MemInfo d u ret -> String
$cshow :: forall d u ret.
(Show d, Show u, Show ret) =>
MemInfo d u ret -> String
showsPrec :: Int -> MemInfo d u ret -> ShowS
$cshowsPrec :: forall d u ret.
(Show d, Show u, Show ret) =>
Int -> MemInfo d u ret -> ShowS
Show, Eq (MemInfo d u ret)
Eq (MemInfo d u ret)
-> (MemInfo d u ret -> MemInfo d u ret -> Ordering)
-> (MemInfo d u ret -> MemInfo d u ret -> Bool)
-> (MemInfo d u ret -> MemInfo d u ret -> Bool)
-> (MemInfo d u ret -> MemInfo d u ret -> Bool)
-> (MemInfo d u ret -> MemInfo d u ret -> Bool)
-> (MemInfo d u ret -> MemInfo d u ret -> MemInfo d u ret)
-> (MemInfo d u ret -> MemInfo d u ret -> MemInfo d u ret)
-> Ord (MemInfo d u ret)
MemInfo d u ret -> MemInfo d u ret -> Bool
MemInfo d u ret -> MemInfo d u ret -> Ordering
MemInfo d u ret -> MemInfo d u ret -> MemInfo d u 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 u ret. (Ord d, Ord u, Ord ret) => Eq (MemInfo d u ret)
forall d u ret.
(Ord d, Ord u, Ord ret) =>
MemInfo d u ret -> MemInfo d u ret -> Bool
forall d u ret.
(Ord d, Ord u, Ord ret) =>
MemInfo d u ret -> MemInfo d u ret -> Ordering
forall d u ret.
(Ord d, Ord u, Ord ret) =>
MemInfo d u ret -> MemInfo d u ret -> MemInfo d u ret
min :: MemInfo d u ret -> MemInfo d u ret -> MemInfo d u ret
$cmin :: forall d u ret.
(Ord d, Ord u, Ord ret) =>
MemInfo d u ret -> MemInfo d u ret -> MemInfo d u ret
max :: MemInfo d u ret -> MemInfo d u ret -> MemInfo d u ret
$cmax :: forall d u ret.
(Ord d, Ord u, Ord ret) =>
MemInfo d u ret -> MemInfo d u ret -> MemInfo d u ret
>= :: MemInfo d u ret -> MemInfo d u ret -> Bool
$c>= :: forall d u ret.
(Ord d, Ord u, Ord ret) =>
MemInfo d u ret -> MemInfo d u ret -> Bool
> :: MemInfo d u ret -> MemInfo d u ret -> Bool
$c> :: forall d u ret.
(Ord d, Ord u, Ord ret) =>
MemInfo d u ret -> MemInfo d u ret -> Bool
<= :: MemInfo d u ret -> MemInfo d u ret -> Bool
$c<= :: forall d u ret.
(Ord d, Ord u, Ord ret) =>
MemInfo d u ret -> MemInfo d u ret -> Bool
< :: MemInfo d u ret -> MemInfo d u ret -> Bool
$c< :: forall d u ret.
(Ord d, Ord u, Ord ret) =>
MemInfo d u ret -> MemInfo d u ret -> Bool
compare :: MemInfo d u ret -> MemInfo d u ret -> Ordering
$ccompare :: forall d u ret.
(Ord d, Ord u, Ord ret) =>
MemInfo d u ret -> MemInfo d u ret -> Ordering
$cp1Ord :: forall d u ret. (Ord d, Ord u, Ord ret) => Eq (MemInfo d u ret)
Ord)
type MemBound u = MemInfo SubExp u MemBind
instance FixExt ret => DeclExtTyped (MemInfo ExtSize Uniqueness ret) where
declExtTypeOf :: MemInfo ExtSize Uniqueness ret -> DeclExtType
declExtTypeOf (MemPrim PrimType
pt) = PrimType -> DeclExtType
forall shape u. PrimType -> TypeBase shape u
Prim PrimType
pt
declExtTypeOf (MemMem Space
space) = Space -> DeclExtType
forall shape u. Space -> TypeBase shape u
Mem Space
space
declExtTypeOf (MemArray PrimType
pt ShapeBase ExtSize
shape Uniqueness
u ret
_) = PrimType -> ShapeBase ExtSize -> Uniqueness -> DeclExtType
forall shape u. PrimType -> shape -> u -> TypeBase shape u
Array PrimType
pt ShapeBase ExtSize
shape Uniqueness
u
declExtTypeOf (MemAcc VName
acc Shape
ispace [Type]
ts Uniqueness
u) = VName -> Shape -> [Type] -> Uniqueness -> DeclExtType
forall shape u. VName -> Shape -> [Type] -> u -> TypeBase shape u
Acc VName
acc Shape
ispace [Type]
ts Uniqueness
u
instance FixExt ret => ExtTyped (MemInfo ExtSize NoUniqueness ret) where
extTypeOf :: MemInfo ExtSize NoUniqueness ret -> ExtType
extTypeOf (MemPrim PrimType
pt) = PrimType -> ExtType
forall shape u. PrimType -> TypeBase shape u
Prim PrimType
pt
extTypeOf (MemMem Space
space) = Space -> ExtType
forall shape u. Space -> TypeBase shape u
Mem Space
space
extTypeOf (MemArray PrimType
pt ShapeBase ExtSize
shape NoUniqueness
u ret
_) = PrimType -> ShapeBase ExtSize -> NoUniqueness -> ExtType
forall shape u. PrimType -> shape -> u -> TypeBase shape u
Array PrimType
pt ShapeBase ExtSize
shape NoUniqueness
u
extTypeOf (MemAcc VName
acc Shape
ispace [Type]
ts NoUniqueness
u) = VName -> Shape -> [Type] -> NoUniqueness -> ExtType
forall shape u. VName -> Shape -> [Type] -> u -> TypeBase shape u
Acc VName
acc Shape
ispace [Type]
ts NoUniqueness
u
instance FixExt ret => FixExt (MemInfo ExtSize u ret) where
fixExt :: Int -> SubExp -> MemInfo ExtSize u ret -> MemInfo ExtSize u ret
fixExt Int
_ SubExp
_ (MemPrim PrimType
pt) = PrimType -> MemInfo ExtSize u ret
forall d u ret. PrimType -> MemInfo d u ret
MemPrim PrimType
pt
fixExt Int
_ SubExp
_ (MemMem Space
space) = Space -> MemInfo ExtSize u ret
forall d u ret. Space -> MemInfo d u ret
MemMem Space
space
fixExt Int
i SubExp
se (MemArray PrimType
pt ShapeBase ExtSize
shape u
u ret
ret) =
PrimType -> ShapeBase ExtSize -> u -> ret -> MemInfo ExtSize u ret
forall d u ret.
PrimType -> ShapeBase d -> u -> ret -> MemInfo d u ret
MemArray PrimType
pt (Int -> SubExp -> ShapeBase ExtSize -> ShapeBase ExtSize
forall t. FixExt t => Int -> SubExp -> t -> t
fixExt Int
i SubExp
se ShapeBase ExtSize
shape) u
u (Int -> SubExp -> ret -> ret
forall t. FixExt t => Int -> SubExp -> t -> t
fixExt Int
i SubExp
se ret
ret)
fixExt Int
_ SubExp
_ (MemAcc VName
acc Shape
ispace [Type]
ts u
u) = VName -> Shape -> [Type] -> u -> MemInfo ExtSize u ret
forall d u ret. VName -> Shape -> [Type] -> u -> MemInfo d u ret
MemAcc VName
acc Shape
ispace [Type]
ts u
u
instance Typed (MemInfo SubExp Uniqueness ret) where
typeOf :: MemInfo SubExp Uniqueness ret -> Type
typeOf = TypeBase Shape Uniqueness -> Type
forall shape.
TypeBase shape Uniqueness -> TypeBase shape NoUniqueness
fromDecl (TypeBase Shape Uniqueness -> Type)
-> (MemInfo SubExp Uniqueness ret -> TypeBase Shape Uniqueness)
-> MemInfo SubExp Uniqueness ret
-> Type
forall k (cat :: k -> k -> *) (b :: k) (c :: k) (a :: k).
Category cat =>
cat b c -> cat a b -> cat a c
. MemInfo SubExp Uniqueness ret -> TypeBase Shape Uniqueness
forall t. DeclTyped t => t -> TypeBase Shape Uniqueness
declTypeOf
instance Typed (MemInfo SubExp NoUniqueness ret) where
typeOf :: MemInfo SubExp NoUniqueness ret -> Type
typeOf (MemPrim PrimType
pt) = PrimType -> Type
forall shape u. PrimType -> TypeBase shape u
Prim PrimType
pt
typeOf (MemMem Space
space) = Space -> Type
forall shape u. Space -> TypeBase shape u
Mem Space
space
typeOf (MemArray PrimType
bt Shape
shape NoUniqueness
u ret
_) = PrimType -> Shape -> NoUniqueness -> Type
forall shape u. PrimType -> shape -> u -> TypeBase shape u
Array PrimType
bt Shape
shape NoUniqueness
u
typeOf (MemAcc VName
acc Shape
ispace [Type]
ts NoUniqueness
u) = VName -> Shape -> [Type] -> NoUniqueness -> Type
forall shape u. VName -> Shape -> [Type] -> u -> TypeBase shape u
Acc VName
acc Shape
ispace [Type]
ts NoUniqueness
u
instance DeclTyped (MemInfo SubExp Uniqueness ret) where
declTypeOf :: MemInfo SubExp Uniqueness ret -> TypeBase Shape Uniqueness
declTypeOf (MemPrim PrimType
bt) = PrimType -> TypeBase Shape Uniqueness
forall shape u. PrimType -> TypeBase shape u
Prim PrimType
bt
declTypeOf (MemMem Space
space) = Space -> TypeBase Shape Uniqueness
forall shape u. Space -> TypeBase shape u
Mem Space
space
declTypeOf (MemArray PrimType
bt Shape
shape Uniqueness
u ret
_) = PrimType -> Shape -> Uniqueness -> TypeBase Shape Uniqueness
forall shape u. PrimType -> shape -> u -> TypeBase shape u
Array PrimType
bt Shape
shape Uniqueness
u
declTypeOf (MemAcc VName
acc Shape
ispace [Type]
ts Uniqueness
u) = VName -> Shape -> [Type] -> Uniqueness -> TypeBase Shape Uniqueness
forall shape u. VName -> Shape -> [Type] -> u -> TypeBase shape u
Acc VName
acc Shape
ispace [Type]
ts Uniqueness
u
instance (FreeIn d, FreeIn ret) => FreeIn (MemInfo d u ret) where
freeIn' :: MemInfo d u ret -> FV
freeIn' (MemArray PrimType
_ ShapeBase d
shape u
_ 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 Shape
ispace [Type]
ts u
_) = (VName, Shape, [Type]) -> FV
forall a. FreeIn a => a -> FV
freeIn' (VName
acc, Shape
ispace, [Type]
ts)
instance (Substitute d, Substitute ret) => Substitute (MemInfo d u ret) where
substituteNames :: Map VName VName -> MemInfo d u ret -> MemInfo d u ret
substituteNames Map VName VName
subst (MemArray PrimType
bt ShapeBase d
shape u
u ret
ret) =
PrimType -> ShapeBase d -> u -> ret -> MemInfo d u ret
forall d u ret.
PrimType -> ShapeBase d -> u -> ret -> MemInfo d u 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)
u
u
(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 Shape
ispace [Type]
ts u
u) =
VName -> Shape -> [Type] -> u -> MemInfo d u ret
forall d u ret. VName -> Shape -> [Type] -> u -> MemInfo d u 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 -> Shape -> Shape
forall a. Substitute a => Map VName VName -> a -> a
substituteNames Map VName VName
substs Shape
ispace)
(Map VName VName -> [Type] -> [Type]
forall a. Substitute a => Map VName VName -> a -> a
substituteNames Map VName VName
substs [Type]
ts)
u
u
substituteNames Map VName VName
_ (MemMem Space
space) =
Space -> MemInfo d u ret
forall d u ret. Space -> MemInfo d u ret
MemMem Space
space
substituteNames Map VName VName
_ (MemPrim PrimType
bt) =
PrimType -> MemInfo d u ret
forall d u ret. PrimType -> MemInfo d u ret
MemPrim PrimType
bt
instance (Substitute d, Substitute ret) => Rename (MemInfo d u ret) where
rename :: MemInfo d u ret -> RenameM (MemInfo d u ret)
rename = MemInfo d u ret -> RenameM (MemInfo d u ret)
forall a. Substitute a => a -> RenameM a
substituteRename
simplifyIxFun ::
Engine.SimplifiableRep rep =>
IxFun ->
Engine.SimpleM rep IxFun
simplifyIxFun :: IxFun -> SimpleM rep IxFun
simplifyIxFun = (TPrimExp Int64 VName -> SimpleM rep (TPrimExp Int64 VName))
-> IxFun -> SimpleM rep IxFun
forall (t :: * -> *) (f :: * -> *) a b.
(Traversable t, Applicative f) =>
(a -> f b) -> t a -> f (t b)
traverse ((TPrimExp Int64 VName -> SimpleM rep (TPrimExp Int64 VName))
-> IxFun -> SimpleM rep IxFun)
-> (TPrimExp Int64 VName -> SimpleM rep (TPrimExp Int64 VName))
-> IxFun
-> SimpleM rep IxFun
forall a b. (a -> b) -> a -> b
$ (PrimExp VName -> TPrimExp Int64 VName)
-> SimpleM rep (PrimExp VName)
-> SimpleM rep (TPrimExp Int64 VName)
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 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 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 t v. TPrimExp t v -> PrimExp v
untyped
simplifyExtIxFun ::
Engine.SimplifiableRep rep =>
ExtIxFun ->
Engine.SimpleM rep ExtIxFun
simplifyExtIxFun :: ExtIxFun -> SimpleM rep ExtIxFun
simplifyExtIxFun = (TPrimExp Int64 (Ext VName)
-> SimpleM rep (TPrimExp Int64 (Ext VName)))
-> ExtIxFun -> SimpleM rep ExtIxFun
forall (t :: * -> *) (f :: * -> *) a b.
(Traversable t, Applicative f) =>
(a -> f b) -> t a -> f (t b)
traverse ((TPrimExp Int64 (Ext VName)
-> SimpleM rep (TPrimExp Int64 (Ext VName)))
-> ExtIxFun -> SimpleM rep ExtIxFun)
-> (TPrimExp Int64 (Ext VName)
-> SimpleM rep (TPrimExp Int64 (Ext VName)))
-> ExtIxFun
-> SimpleM rep ExtIxFun
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 (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 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 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 t v. TPrimExp t v -> PrimExp v
untyped
isStaticIxFun :: ExtIxFun -> Maybe IxFun
isStaticIxFun :: ExtIxFun -> Maybe IxFun
isStaticIxFun = (TPrimExp Int64 (Ext VName) -> Maybe (TPrimExp Int64 VName))
-> ExtIxFun -> Maybe IxFun
forall (t :: * -> *) (f :: * -> *) a b.
(Traversable t, Applicative f) =>
(a -> f b) -> t a -> f (t b)
traverse ((TPrimExp Int64 (Ext VName) -> Maybe (TPrimExp Int64 VName))
-> ExtIxFun -> Maybe IxFun)
-> (TPrimExp Int64 (Ext VName) -> Maybe (TPrimExp Int64 VName))
-> ExtIxFun
-> Maybe IxFun
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)
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 u ret)
where
simplify :: MemInfo d u ret -> SimpleM rep (MemInfo d u ret)
simplify (MemPrim PrimType
bt) =
MemInfo d u ret -> SimpleM rep (MemInfo d u ret)
forall (f :: * -> *) a. Applicative f => a -> f a
pure (MemInfo d u ret -> SimpleM rep (MemInfo d u ret))
-> MemInfo d u ret -> SimpleM rep (MemInfo d u ret)
forall a b. (a -> b) -> a -> b
$ PrimType -> MemInfo d u ret
forall d u ret. PrimType -> MemInfo d u ret
MemPrim PrimType
bt
simplify (MemMem Space
space) =
MemInfo d u ret -> SimpleM rep (MemInfo d u ret)
forall (f :: * -> *) a. Applicative f => a -> f a
pure (MemInfo d u ret -> SimpleM rep (MemInfo d u ret))
-> MemInfo d u ret -> SimpleM rep (MemInfo d u ret)
forall a b. (a -> b) -> a -> b
$ Space -> MemInfo d u ret
forall d u ret. Space -> MemInfo d u ret
MemMem Space
space
simplify (MemArray PrimType
bt ShapeBase d
shape u
u ret
ret) =
PrimType -> ShapeBase d -> u -> ret -> MemInfo d u ret
forall d u ret.
PrimType -> ShapeBase d -> u -> ret -> MemInfo d u ret
MemArray PrimType
bt (ShapeBase d -> u -> ret -> MemInfo d u ret)
-> SimpleM rep (ShapeBase d)
-> SimpleM rep (u -> ret -> MemInfo d u ret)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> ShapeBase d -> SimpleM rep (ShapeBase d)
forall e rep.
(Simplifiable e, SimplifiableRep rep) =>
e -> SimpleM rep e
Engine.simplify ShapeBase d
shape SimpleM rep (u -> ret -> MemInfo d u ret)
-> SimpleM rep u -> SimpleM rep (ret -> MemInfo d u ret)
forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> u -> SimpleM rep u
forall (f :: * -> *) a. Applicative f => a -> f a
pure u
u SimpleM rep (ret -> MemInfo d u ret)
-> SimpleM rep ret -> SimpleM rep (MemInfo d u ret)
forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> ret -> SimpleM rep ret
forall e rep.
(Simplifiable e, SimplifiableRep rep) =>
e -> SimpleM rep e
Engine.simplify ret
ret
simplify (MemAcc VName
acc Shape
ispace [Type]
ts u
u) =
VName -> Shape -> [Type] -> u -> MemInfo d u ret
forall d u ret. VName -> Shape -> [Type] -> u -> MemInfo d u ret
MemAcc (VName -> Shape -> [Type] -> u -> MemInfo d u ret)
-> SimpleM rep VName
-> SimpleM rep (Shape -> [Type] -> u -> MemInfo d u ret)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> VName -> SimpleM rep VName
forall e rep.
(Simplifiable e, SimplifiableRep rep) =>
e -> SimpleM rep e
Engine.simplify VName
acc SimpleM rep (Shape -> [Type] -> u -> MemInfo d u ret)
-> SimpleM rep Shape
-> SimpleM rep ([Type] -> u -> MemInfo d u ret)
forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> Shape -> SimpleM rep Shape
forall e rep.
(Simplifiable e, SimplifiableRep rep) =>
e -> SimpleM rep e
Engine.simplify Shape
ispace SimpleM rep ([Type] -> u -> MemInfo d u ret)
-> SimpleM rep [Type] -> SimpleM rep (u -> MemInfo d u ret)
forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> [Type] -> SimpleM rep [Type]
forall e rep.
(Simplifiable e, SimplifiableRep rep) =>
e -> SimpleM rep e
Engine.simplify [Type]
ts SimpleM rep (u -> MemInfo d u ret)
-> SimpleM rep u -> SimpleM rep (MemInfo d u ret)
forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> u -> SimpleM rep u
forall (f :: * -> *) a. Applicative f => a -> f a
pure u
u
instance
( PP.Pretty (ShapeBase d),
PP.Pretty (TypeBase (ShapeBase d) u),
PP.Pretty d,
PP.Pretty u,
PP.Pretty ret
) =>
PP.Pretty (MemInfo d u ret)
where
ppr :: MemInfo d u ret -> Doc
ppr (MemPrim PrimType
bt) = PrimType -> Doc
forall a. Pretty a => a -> Doc
PP.ppr PrimType
bt
ppr (MemMem Space
DefaultSpace) = String -> Doc
PP.text String
"mem"
ppr (MemMem Space
s) = String -> Doc
PP.text String
"mem" Doc -> Doc -> Doc
forall a. Semigroup a => a -> a -> a
<> Space -> Doc
forall a. Pretty a => a -> Doc
PP.ppr Space
s
ppr (MemArray PrimType
bt ShapeBase d
shape u
u ret
ret) =
TypeBase (ShapeBase d) u -> Doc
forall a. Pretty a => a -> Doc
PP.ppr (PrimType -> ShapeBase d -> u -> TypeBase (ShapeBase d) u
forall shape u. PrimType -> shape -> u -> TypeBase shape u
Array PrimType
bt ShapeBase d
shape u
u) Doc -> Doc -> Doc
<+> String -> Doc
PP.text String
"@" Doc -> Doc -> Doc
<+> ret -> Doc
forall a. Pretty a => a -> Doc
PP.ppr ret
ret
ppr (MemAcc VName
acc Shape
ispace [Type]
ts u
u) =
u -> Doc
forall a. Pretty a => a -> Doc
PP.ppr u
u Doc -> Doc -> Doc
forall a. Semigroup a => a -> a -> a
<> Type -> Doc
forall a. Pretty a => a -> Doc
PP.ppr (VName -> Shape -> [Type] -> NoUniqueness -> Type
forall shape u. VName -> Shape -> [Type] -> u -> TypeBase shape u
Acc VName
acc Shape
ispace [Type]
ts NoUniqueness
NoUniqueness :: Type)
data MemBind
=
ArrayIn VName IxFun
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
showList :: [MemBind] -> ShowS
$cshowList :: [MemBind] -> ShowS
show :: MemBind -> String
$cshow :: MemBind -> String
showsPrec :: Int -> MemBind -> ShowS
$cshowsPrec :: Int -> 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 IxFun
ixfun) =
VName -> IxFun -> 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 -> IxFun -> IxFun
forall a. Substitute a => Map VName VName -> a -> a
substituteNames Map VName VName
substs IxFun
ixfun)
instance PP.Pretty MemBind where
ppr :: MemBind -> Doc
ppr (ArrayIn VName
mem IxFun
ixfun) =
VName -> Doc
forall a. Pretty a => a -> Doc
PP.ppr VName
mem Doc -> Doc -> Doc
<+> Doc
"->" Doc -> Doc -> Doc
PP.</> IxFun -> Doc
forall a. Pretty a => a -> Doc
PP.ppr IxFun
ixfun
instance FreeIn MemBind where
freeIn' :: MemBind -> FV
freeIn' (ArrayIn VName
mem IxFun
ixfun) = VName -> FV
forall a. FreeIn a => a -> FV
freeIn' VName
mem FV -> FV -> FV
forall a. Semigroup a => a -> a -> a
<> IxFun -> FV
forall a. FreeIn a => a -> FV
freeIn' IxFun
ixfun
data MemReturn
=
ReturnsInBlock VName ExtIxFun
|
ReturnsNewBlock Space Int ExtIxFun
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
showList :: [MemReturn] -> ShowS
$cshowList :: [MemReturn] -> ShowS
show :: MemReturn -> String
$cshow :: MemReturn -> String
showsPrec :: Int -> MemReturn -> ShowS
$cshowsPrec :: Int -> 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 ExtIxFun
ixfun) =
VName -> ExtIxFun -> 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 -> ExtIxFun -> ExtIxFun
forall a. Substitute a => Map VName VName -> a -> a
substituteNames Map VName VName
substs ExtIxFun
ixfun)
substituteNames Map VName VName
substs (ReturnsNewBlock Space
space Int
i ExtIxFun
ixfun) =
Space -> Int -> ExtIxFun -> MemReturn
ReturnsNewBlock Space
space Int
i (Map VName VName -> ExtIxFun -> ExtIxFun
forall a. Substitute a => Map VName VName -> a -> a
substituteNames Map VName VName
substs ExtIxFun
ixfun)
instance FixExt MemReturn where
fixExt :: Int -> SubExp -> MemReturn -> MemReturn
fixExt Int
i (Var VName
v) (ReturnsNewBlock Space
_ Int
j ExtIxFun
ixfun)
| Int
j Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
== Int
i =
VName -> ExtIxFun -> MemReturn
ReturnsInBlock VName
v (ExtIxFun -> MemReturn) -> ExtIxFun -> MemReturn
forall a b. (a -> b) -> a -> b
$
Int -> PrimExp VName -> ExtIxFun -> ExtIxFun
fixExtIxFun
Int
i
(PrimType -> SubExp -> PrimExp VName
primExpFromSubExp PrimType
int64 (VName -> SubExp
Var VName
v))
ExtIxFun
ixfun
fixExt Int
i SubExp
se (ReturnsNewBlock Space
space Int
j ExtIxFun
ixfun) =
Space -> Int -> ExtIxFun -> MemReturn
ReturnsNewBlock
Space
space
Int
j'
(Int -> PrimExp VName -> ExtIxFun -> ExtIxFun
fixExtIxFun Int
i (PrimType -> SubExp -> PrimExp VName
primExpFromSubExp PrimType
int64 SubExp
se) ExtIxFun
ixfun)
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 ExtIxFun
ixfun) =
VName -> ExtIxFun -> MemReturn
ReturnsInBlock VName
mem (Int -> PrimExp VName -> ExtIxFun -> ExtIxFun
fixExtIxFun Int
i (PrimType -> SubExp -> PrimExp VName
primExpFromSubExp PrimType
int64 SubExp
se) ExtIxFun
ixfun)
fixExtIxFun :: Int -> PrimExp VName -> ExtIxFun -> ExtIxFun
fixExtIxFun :: Int -> PrimExp VName -> ExtIxFun -> ExtIxFun
fixExtIxFun Int
i PrimExp VName
e = (TPrimExp Int64 (Ext VName) -> TPrimExp Int64 (Ext VName))
-> ExtIxFun -> ExtIxFun
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap ((TPrimExp Int64 (Ext VName) -> TPrimExp Int64 (Ext VName))
-> ExtIxFun -> ExtIxFun)
-> (TPrimExp Int64 (Ext VName) -> TPrimExp Int64 (Ext VName))
-> ExtIxFun
-> ExtIxFun
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 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 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 t 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 (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 :: 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
existentialiseIxFun :: [VName] -> IxFun -> ExtIxFun
existentialiseIxFun :: [VName] -> IxFun -> ExtIxFun
existentialiseIxFun [VName]
ctx = Map (Ext VName) (TPrimExp Int64 (Ext VName))
-> ExtIxFun -> ExtIxFun
forall a t.
Ord a =>
Map a (TPrimExp t a)
-> IxFun (TPrimExp t a) -> IxFun (TPrimExp t a)
IxFun.substituteInIxFun Map (Ext VName) (TPrimExp Int64 (Ext VName))
ctx' (ExtIxFun -> ExtIxFun) -> (IxFun -> ExtIxFun) -> IxFun -> ExtIxFun
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))
-> IxFun -> ExtIxFun
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap ((VName -> Ext VName)
-> TPrimExp Int64 VName -> TPrimExp Int64 (Ext VName)
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
ppr :: MemReturn -> Doc
ppr (ReturnsInBlock VName
v ExtIxFun
ixfun) =
Doc -> Doc
PP.parens (Doc -> Doc) -> Doc -> Doc
forall a b. (a -> b) -> a -> b
$ VName -> Doc
forall a. Pretty a => a -> Doc
ppr VName
v Doc -> Doc -> Doc
<+> Doc
"->" Doc -> Doc -> Doc
PP.</> ExtIxFun -> Doc
forall a. Pretty a => a -> Doc
PP.ppr ExtIxFun
ixfun
ppr (ReturnsNewBlock Space
space Int
i ExtIxFun
ixfun) =
Doc
"?" Doc -> Doc -> Doc
forall a. Semigroup a => a -> a -> a
<> Int -> Doc
forall a. Pretty a => a -> Doc
ppr Int
i Doc -> Doc -> Doc
forall a. Semigroup a => a -> a -> a
<> Space -> Doc
forall a. Pretty a => a -> Doc
PP.ppr Space
space Doc -> Doc -> Doc
<+> Doc
"->" Doc -> Doc -> Doc
PP.</> ExtIxFun -> Doc
forall a. Pretty a => a -> Doc
PP.ppr ExtIxFun
ixfun
instance FreeIn MemReturn where
freeIn' :: MemReturn -> FV
freeIn' (ReturnsInBlock VName
v ExtIxFun
ixfun) = VName -> FV
forall a. FreeIn a => a -> FV
freeIn' VName
v FV -> FV -> FV
forall a. Semigroup a => a -> a -> a
<> ExtIxFun -> FV
forall a. FreeIn a => a -> FV
freeIn' ExtIxFun
ixfun
freeIn' (ReturnsNewBlock Space
space Int
_ ExtIxFun
ixfun) = Space -> FV
forall a. FreeIn a => a -> FV
freeIn' Space
space FV -> FV -> FV
forall a. Semigroup a => a -> a -> a
<> ExtIxFun -> FV
forall a. FreeIn a => a -> FV
freeIn' ExtIxFun
ixfun
instance Engine.Simplifiable MemReturn where
simplify :: MemReturn -> SimpleM rep MemReturn
simplify (ReturnsNewBlock Space
space Int
i ExtIxFun
ixfun) =
Space -> Int -> ExtIxFun -> MemReturn
ReturnsNewBlock Space
space Int
i (ExtIxFun -> MemReturn)
-> SimpleM rep ExtIxFun -> SimpleM rep MemReturn
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> ExtIxFun -> SimpleM rep ExtIxFun
forall rep. SimplifiableRep rep => ExtIxFun -> SimpleM rep ExtIxFun
simplifyExtIxFun ExtIxFun
ixfun
simplify (ReturnsInBlock VName
v ExtIxFun
ixfun) =
VName -> ExtIxFun -> MemReturn
ReturnsInBlock (VName -> ExtIxFun -> MemReturn)
-> SimpleM rep VName -> SimpleM rep (ExtIxFun -> MemReturn)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> VName -> SimpleM rep VName
forall e rep.
(Simplifiable e, SimplifiableRep rep) =>
e -> SimpleM rep e
Engine.simplify VName
v SimpleM rep (ExtIxFun -> MemReturn)
-> SimpleM rep ExtIxFun -> SimpleM rep MemReturn
forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> ExtIxFun -> SimpleM rep ExtIxFun
forall rep. SimplifiableRep rep => ExtIxFun -> SimpleM rep ExtIxFun
simplifyExtIxFun ExtIxFun
ixfun
instance Engine.Simplifiable MemBind where
simplify :: MemBind -> SimpleM rep MemBind
simplify (ArrayIn VName
mem IxFun
ixfun) =
VName -> IxFun -> MemBind
ArrayIn (VName -> IxFun -> MemBind)
-> SimpleM rep VName -> SimpleM rep (IxFun -> MemBind)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> VName -> SimpleM rep VName
forall e rep.
(Simplifiable e, SimplifiableRep rep) =>
e -> SimpleM rep e
Engine.simplify VName
mem SimpleM rep (IxFun -> MemBind)
-> SimpleM rep IxFun -> SimpleM rep MemBind
forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> IxFun -> SimpleM rep IxFun
forall rep. SimplifiableRep rep => IxFun -> SimpleM rep IxFun
simplifyIxFun IxFun
ixfun
instance Engine.Simplifiable [FunReturns] where
simplify :: [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)
mapM FunReturns -> SimpleM rep FunReturns
forall e rep.
(Simplifiable e, SimplifiableRep rep) =>
e -> SimpleM rep e
Engine.simplify
type ExpReturns = MemInfo ExtSize NoUniqueness (Maybe MemReturn)
type BodyReturns = MemInfo ExtSize NoUniqueness MemReturn
type FunReturns = MemInfo ExtSize Uniqueness MemReturn
maybeReturns :: MemInfo d u r -> MemInfo d u (Maybe r)
maybeReturns :: MemInfo d u r -> MemInfo d u (Maybe r)
maybeReturns (MemArray PrimType
bt ShapeBase d
shape u
u r
ret) =
PrimType -> ShapeBase d -> u -> Maybe r -> MemInfo d u (Maybe r)
forall d u ret.
PrimType -> ShapeBase d -> u -> ret -> MemInfo d u ret
MemArray PrimType
bt ShapeBase d
shape u
u (Maybe r -> MemInfo d u (Maybe r))
-> Maybe r -> MemInfo d u (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 u (Maybe r)
forall d u ret. PrimType -> MemInfo d u ret
MemPrim PrimType
bt
maybeReturns (MemMem Space
space) =
Space -> MemInfo d u (Maybe r)
forall d u ret. Space -> MemInfo d u ret
MemMem Space
space
maybeReturns (MemAcc VName
acc Shape
ispace [Type]
ts u
u) =
VName -> Shape -> [Type] -> u -> MemInfo d u (Maybe r)
forall d u ret. VName -> Shape -> [Type] -> u -> MemInfo d u ret
MemAcc VName
acc Shape
ispace [Type]
ts u
u
noUniquenessReturns :: MemInfo d u r -> MemInfo d NoUniqueness r
noUniquenessReturns :: MemInfo d u r -> MemInfo d NoUniqueness r
noUniquenessReturns (MemArray PrimType
bt ShapeBase d
shape u
_ r
r) =
PrimType
-> ShapeBase d -> NoUniqueness -> r -> MemInfo d NoUniqueness r
forall d u ret.
PrimType -> ShapeBase d -> u -> ret -> MemInfo d u ret
MemArray PrimType
bt ShapeBase d
shape NoUniqueness
NoUniqueness r
r
noUniquenessReturns (MemPrim PrimType
bt) =
PrimType -> MemInfo d NoUniqueness r
forall d u ret. PrimType -> MemInfo d u ret
MemPrim PrimType
bt
noUniquenessReturns (MemMem Space
space) =
Space -> MemInfo d NoUniqueness r
forall d u ret. Space -> MemInfo d u ret
MemMem Space
space
noUniquenessReturns (MemAcc VName
acc Shape
ispace [Type]
ts u
_) =
VName
-> Shape -> [Type] -> NoUniqueness -> MemInfo d NoUniqueness r
forall d u ret. VName -> Shape -> [Type] -> u -> MemInfo d u ret
MemAcc VName
acc Shape
ispace [Type]
ts NoUniqueness
NoUniqueness
funReturnsToExpReturns :: FunReturns -> ExpReturns
funReturnsToExpReturns :: FunReturns -> ExpReturns
funReturnsToExpReturns = MemInfo ExtSize Uniqueness (Maybe MemReturn) -> ExpReturns
forall d u r. MemInfo d u r -> MemInfo d NoUniqueness r
noUniquenessReturns (MemInfo ExtSize Uniqueness (Maybe MemReturn) -> ExpReturns)
-> (FunReturns -> MemInfo ExtSize Uniqueness (Maybe MemReturn))
-> FunReturns
-> ExpReturns
forall k (cat :: k -> k -> *) (b :: k) (c :: k) (a :: k).
Category cat =>
cat b c -> cat a b -> cat a c
. FunReturns -> MemInfo ExtSize Uniqueness (Maybe MemReturn)
forall d u r. MemInfo d u r -> MemInfo d u (Maybe r)
maybeReturns
bodyReturnsToExpReturns :: BodyReturns -> ExpReturns
bodyReturnsToExpReturns :: BodyReturns -> ExpReturns
bodyReturnsToExpReturns = ExpReturns -> ExpReturns
forall d u r. MemInfo d u r -> MemInfo d NoUniqueness r
noUniquenessReturns (ExpReturns -> ExpReturns)
-> (BodyReturns -> ExpReturns) -> BodyReturns -> ExpReturns
forall k (cat :: k -> k -> *) (b :: k) (c :: k) (a :: k).
Category cat =>
cat b c -> cat a b -> cat a c
. BodyReturns -> ExpReturns
forall d u r. MemInfo d u r -> MemInfo d u (Maybe r)
maybeReturns
varInfoToExpReturns :: MemInfo SubExp NoUniqueness MemBind -> ExpReturns
varInfoToExpReturns :: LetDecMem -> ExpReturns
varInfoToExpReturns (MemArray PrimType
et Shape
shape NoUniqueness
u (ArrayIn VName
mem IxFun
ixfun)) =
PrimType
-> ShapeBase ExtSize
-> NoUniqueness
-> Maybe MemReturn
-> ExpReturns
forall d u ret.
PrimType -> ShapeBase d -> u -> ret -> MemInfo d u ret
MemArray PrimType
et ((SubExp -> ExtSize) -> Shape -> ShapeBase ExtSize
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap SubExp -> ExtSize
forall a. a -> Ext a
Free Shape
shape) NoUniqueness
u (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 -> ExtIxFun -> MemReturn
ReturnsInBlock VName
mem (ExtIxFun -> MemReturn) -> ExtIxFun -> MemReturn
forall a b. (a -> b) -> a -> b
$
[VName] -> IxFun -> ExtIxFun
existentialiseIxFun [] IxFun
ixfun
varInfoToExpReturns (MemPrim PrimType
pt) = PrimType -> ExpReturns
forall d u ret. PrimType -> MemInfo d u ret
MemPrim PrimType
pt
varInfoToExpReturns (MemAcc VName
acc Shape
ispace [Type]
ts NoUniqueness
u) = VName -> Shape -> [Type] -> NoUniqueness -> ExpReturns
forall d u ret. VName -> Shape -> [Type] -> u -> MemInfo d u ret
MemAcc VName
acc Shape
ispace [Type]
ts NoUniqueness
u
varInfoToExpReturns (MemMem Space
space) = Space -> ExpReturns
forall d u ret. Space -> MemInfo d u ret
MemMem Space
space
matchRetTypeToResult ::
(Mem rep inner, TC.Checkable rep) =>
[FunReturns] ->
Result ->
TC.TypeM rep ()
matchRetTypeToResult :: [FunReturns] -> Result -> TypeM rep ()
matchRetTypeToResult [FunReturns]
rettype Result
result = do
Scope (Aliases rep)
scope <- TypeM rep (Scope (Aliases rep))
forall rep (m :: * -> *). HasScope rep m => m (Scope rep)
askScope
[LetDecMem]
result_ts <- ReaderT (Scope rep) (TypeM rep) [LetDecMem]
-> Scope rep -> TypeM rep [LetDecMem]
forall r (m :: * -> *) a. ReaderT r m a -> r -> m a
runReaderT ((SubExpRes -> ReaderT (Scope rep) (TypeM rep) LetDecMem)
-> Result -> ReaderT (Scope rep) (TypeM rep) [LetDecMem]
forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
(a -> m b) -> t a -> m (t b)
mapM (SubExp -> ReaderT (Scope rep) (TypeM rep) LetDecMem
forall rep (m :: * -> *) inner.
(HasScope rep m, Mem rep inner) =>
SubExp -> m LetDecMem
subExpMemInfo (SubExp -> ReaderT (Scope rep) (TypeM rep) LetDecMem)
-> (SubExpRes -> SubExp)
-> SubExpRes
-> ReaderT (Scope rep) (TypeM rep) LetDecMem
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) (Scope rep -> TypeM rep [LetDecMem])
-> Scope rep -> TypeM rep [LetDecMem]
forall a b. (a -> b) -> a -> b
$ Scope (Aliases rep) -> Scope rep
forall rep. Scope (Aliases rep) -> Scope rep
removeScopeAliases Scope (Aliases rep)
scope
[FunReturns] -> [SubExp] -> [LetDecMem] -> TypeM rep ()
forall u rep.
Pretty u =>
[MemInfo ExtSize u MemReturn]
-> [SubExp] -> [LetDecMem] -> TypeM rep ()
matchReturnType [FunReturns]
rettype ((SubExpRes -> SubExp) -> Result -> [SubExp]
forall a b. (a -> b) -> [a] -> [b]
map SubExpRes -> SubExp
resSubExp Result
result) [LetDecMem]
result_ts
matchFunctionReturnType ::
(Mem rep inner, TC.Checkable rep) =>
[FunReturns] ->
Result ->
TC.TypeM rep ()
matchFunctionReturnType :: [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.
(HasLetDecMem (LetDec rep), ASTRep rep, OpReturns inner,
FParamInfo rep ~ FParamMem, LParamInfo rep ~ LetDecMem,
RetType rep ~ FunReturns, BranchType rep ~ BodyReturns,
Op rep ~ MemOp inner) =>
SubExp -> TypeM rep ()
checkResultSubExp (SubExp -> TypeM rep ())
-> (SubExpRes -> SubExp) -> SubExpRes -> TypeM rep ()
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 (f :: * -> *) a. Applicative f => a -> f a
pure ()
checkResultSubExp (Var VName
v) = do
LetDecMem
dec <- VName -> TypeM rep LetDecMem
forall rep inner. Mem rep inner => VName -> TypeM rep LetDecMem
varMemInfo VName
v
case LetDecMem
dec of
MemPrim PrimType
_ -> () -> TypeM rep ()
forall (f :: * -> *) a. Applicative f => a -> f a
pure ()
MemMem {} -> () -> TypeM rep ()
forall (f :: * -> *) a. Applicative f => a -> f a
pure ()
MemAcc {} -> () -> TypeM rep ()
forall (f :: * -> *) a. Applicative f => a -> f a
pure ()
MemArray PrimType
_ Shape
_ NoUniqueness
_ (ArrayIn VName
_ IxFun
ixfun)
| IxFun -> Bool
forall num. (Eq num, IntegralExp num) => IxFun num -> Bool
IxFun.isLinear IxFun
ixfun ->
() -> TypeM rep ()
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 ()) -> ErrorCase rep -> TypeM rep ()
forall a b. (a -> b) -> a -> b
$
String -> ErrorCase rep
forall rep. String -> ErrorCase rep
TC.TypeError (String -> ErrorCase rep) -> String -> ErrorCase rep
forall a b. (a -> b) -> a -> b
$
String
"Array "
String -> ShowS
forall a. [a] -> [a] -> [a]
++ VName -> String
forall a. Pretty a => a -> String
pretty VName
v
String -> ShowS
forall a. [a] -> [a] -> [a]
++ String
" returned by function, but has nontrivial index function "
String -> ShowS
forall a. [a] -> [a] -> [a]
++ IxFun -> String
forall a. Pretty a => a -> String
pretty IxFun
ixfun
matchLoopResultMem ::
(Mem rep inner, TC.Checkable rep) =>
[FParam (Aliases rep)] ->
Result ->
TC.TypeM rep ()
matchLoopResultMem :: [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
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 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 -> ExtSize
toExtSE (Var VName
v) = VName -> SubExp
Var (VName -> SubExp) -> Ext VName -> ExtSize
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> VName -> Ext VName
toExtV VName
v
toExtSE (Constant PrimValue
v) = SubExp -> ExtSize
forall a. a -> Ext a
Free (SubExp -> ExtSize) -> SubExp -> ExtSize
forall a b. (a -> b) -> a -> b
$ PrimValue -> SubExp
Constant PrimValue
v
toRet :: FParamMem -> FunReturns
toRet (MemPrim PrimType
t) =
PrimType -> FunReturns
forall d u ret. PrimType -> MemInfo d u ret
MemPrim PrimType
t
toRet (MemMem Space
space) =
Space -> FunReturns
forall d u ret. Space -> MemInfo d u ret
MemMem Space
space
toRet (MemAcc VName
acc Shape
ispace [Type]
ts Uniqueness
u) =
VName -> Shape -> [Type] -> Uniqueness -> FunReturns
forall d u ret. VName -> Shape -> [Type] -> u -> MemInfo d u ret
MemAcc VName
acc Shape
ispace [Type]
ts Uniqueness
u
toRet (MemArray PrimType
pt Shape
shape Uniqueness
u (ArrayIn VName
mem IxFun
ixfun))
| 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 ExtSize -> Uniqueness -> MemReturn -> FunReturns
forall d u ret.
PrimType -> ShapeBase d -> u -> ret -> MemInfo d u ret
MemArray PrimType
pt ShapeBase ExtSize
shape' Uniqueness
u (MemReturn -> FunReturns) -> MemReturn -> FunReturns
forall a b. (a -> b) -> a -> b
$ Space -> Int -> ExtIxFun -> MemReturn
ReturnsNewBlock Space
space Int
i ExtIxFun
ixfun'
| Bool
otherwise =
PrimType
-> ShapeBase ExtSize -> Uniqueness -> MemReturn -> FunReturns
forall d u ret.
PrimType -> ShapeBase d -> u -> ret -> MemInfo d u ret
MemArray PrimType
pt ShapeBase ExtSize
shape' Uniqueness
u (MemReturn -> FunReturns) -> MemReturn -> FunReturns
forall a b. (a -> b) -> a -> b
$ VName -> ExtIxFun -> MemReturn
ReturnsInBlock VName
mem ExtIxFun
ixfun'
where
shape' :: ShapeBase ExtSize
shape' = (SubExp -> ExtSize) -> Shape -> ShapeBase ExtSize
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap SubExp -> ExtSize
toExtSE Shape
shape
ixfun' :: ExtIxFun
ixfun' = [VName] -> IxFun -> ExtIxFun
existentialiseIxFun [VName]
param_names IxFun
ixfun
matchBranchReturnType ::
(Mem rep inner, TC.Checkable rep) =>
[BodyReturns] ->
Body (Aliases rep) ->
TC.TypeM rep ()
matchBranchReturnType :: [BodyReturns] -> Body (Aliases rep) -> TypeM rep ()
matchBranchReturnType [BodyReturns]
rettype (Body BodyDec (Aliases rep)
_ Stms (Aliases rep)
stms Result
res) = do
Scope (Aliases rep)
scope <- TypeM rep (Scope (Aliases rep))
forall rep (m :: * -> *). HasScope rep m => m (Scope rep)
askScope
[LetDecMem]
ts <- ReaderT (Scope rep) (TypeM rep) [LetDecMem]
-> Scope rep -> TypeM rep [LetDecMem]
forall r (m :: * -> *) a. ReaderT r m a -> r -> m a
runReaderT ((SubExpRes -> ReaderT (Scope rep) (TypeM rep) LetDecMem)
-> Result -> ReaderT (Scope rep) (TypeM rep) [LetDecMem]
forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
(a -> m b) -> t a -> m (t b)
mapM (SubExp -> ReaderT (Scope rep) (TypeM rep) LetDecMem
forall rep (m :: * -> *) inner.
(HasScope rep m, Mem rep inner) =>
SubExp -> m LetDecMem
subExpMemInfo (SubExp -> ReaderT (Scope rep) (TypeM rep) LetDecMem)
-> (SubExpRes -> SubExp)
-> SubExpRes
-> ReaderT (Scope rep) (TypeM rep) LetDecMem
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
res) (Scope rep -> TypeM rep [LetDecMem])
-> Scope rep -> TypeM rep [LetDecMem]
forall a b. (a -> b) -> a -> b
$ Scope (Aliases rep) -> Scope rep
forall rep. Scope (Aliases rep) -> Scope rep
removeScopeAliases (Scope (Aliases rep)
scope Scope (Aliases rep) -> Scope (Aliases rep) -> Scope (Aliases rep)
forall a. Semigroup a => a -> a -> a
<> Stms (Aliases rep) -> Scope (Aliases rep)
forall rep a. Scoped rep a => a -> Scope rep
scopeOf Stms (Aliases rep)
stms)
[BodyReturns] -> [SubExp] -> [LetDecMem] -> TypeM rep ()
forall u rep.
Pretty u =>
[MemInfo ExtSize u MemReturn]
-> [SubExp] -> [LetDecMem] -> TypeM rep ()
matchReturnType [BodyReturns]
rettype ((SubExpRes -> SubExp) -> Result -> [SubExp]
forall a b. (a -> b) -> [a] -> [b]
map SubExpRes -> SubExp
resSubExp Result
res) [LetDecMem]
ts
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 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)
traverse
( (Int -> TPrimExp Int64 (Ext VName))
-> Maybe Int -> Maybe (TPrimExp Int64 (Ext VName))
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 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 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 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 (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 u =>
[MemInfo ExtSize u MemReturn] ->
[SubExp] ->
[MemInfo SubExp NoUniqueness MemBind] ->
TC.TypeM rep ()
matchReturnType :: [MemInfo ExtSize u MemReturn]
-> [SubExp] -> [LetDecMem] -> TypeM rep ()
matchReturnType [MemInfo ExtSize u MemReturn]
rettype [SubExp]
res [LetDecMem]
ts = do
let existentialiseIxFun0 :: IxFun -> ExtIxFun
existentialiseIxFun0 :: IxFun -> ExtIxFun
existentialiseIxFun0 = (TPrimExp Int64 VName -> TPrimExp Int64 (Ext VName))
-> IxFun -> ExtIxFun
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap ((TPrimExp Int64 VName -> TPrimExp Int64 (Ext VName))
-> IxFun -> ExtIxFun)
-> (TPrimExp Int64 VName -> TPrimExp Int64 (Ext VName))
-> IxFun
-> ExtIxFun
forall a b. (a -> b) -> a -> b
$ (VName -> Ext VName)
-> TPrimExp Int64 VName -> TPrimExp Int64 (Ext VName)
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 String (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 ->
String -> ExceptT String (TypeM rep) (SubExp, LetDecMem)
forall e (m :: * -> *) a. MonadError e m => e -> m a
throwError (String -> ExceptT String (TypeM rep) (SubExp, LetDecMem))
-> String -> ExceptT String (TypeM rep) (SubExp, LetDecMem)
forall a b. (a -> b) -> a -> b
$ String
"Cannot find variable #" String -> ShowS
forall a. [a] -> [a] -> [a]
++ Int -> String
forall a. Show a => a -> String
show Int
i String -> ShowS
forall a. [a] -> [a] -> [a]
++ String
" in results: " String -> ShowS
forall a. [a] -> [a] -> [a]
++ [SubExp] -> String
forall a. Pretty a => a -> String
pretty [SubExp]
res
Just (SubExp
se, LetDecMem
t) -> (SubExp, LetDecMem)
-> ExceptT String (TypeM rep) (SubExp, LetDecMem)
forall (f :: * -> *) a. Applicative f => a -> f a
pure (SubExp
se, LetDecMem
t)
checkReturn :: MemInfo ExtSize u MemReturn
-> LetDecMem -> ExceptT String (TypeM rep) ()
checkReturn (MemPrim PrimType
x) (MemPrim PrimType
y)
| PrimType
x PrimType -> PrimType -> Bool
forall a. Eq a => a -> a -> Bool
== PrimType
y = () -> ExceptT String (TypeM rep) ()
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 String (TypeM rep) ()
forall (f :: * -> *) a. Applicative f => a -> f a
pure ()
checkReturn (MemAcc VName
xacc Shape
xispace [Type]
xts u
_) (MemAcc VName
yacc Shape
yispace [Type]
yts NoUniqueness
_)
| (VName
xacc, Shape
xispace, [Type]
xts) (VName, Shape, [Type]) -> (VName, Shape, [Type]) -> Bool
forall a. Eq a => a -> a -> Bool
== (VName
yacc, Shape
yispace, [Type]
yts) =
() -> ExceptT String (TypeM rep) ()
forall (f :: * -> *) a. Applicative f => a -> f a
pure ()
checkReturn
(MemArray PrimType
x_pt ShapeBase ExtSize
x_shape u
_ MemReturn
x_ret)
(MemArray PrimType
y_pt Shape
y_shape NoUniqueness
_ MemBind
y_ret)
| PrimType
x_pt PrimType -> PrimType -> Bool
forall a. Eq a => a -> a -> Bool
== PrimType
y_pt,
ShapeBase ExtSize -> Int
forall a. ArrayShape a => a -> Int
shapeRank ShapeBase ExtSize
x_shape Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
== Shape -> Int
forall a. ArrayShape a => a -> Int
shapeRank Shape
y_shape = do
(ExtSize -> SubExp -> ExceptT String (TypeM rep) ())
-> [ExtSize] -> [SubExp] -> ExceptT String (TypeM rep) ()
forall (m :: * -> *) a b c.
Applicative m =>
(a -> b -> m c) -> [a] -> [b] -> m ()
zipWithM_ ExtSize -> SubExp -> ExceptT String (TypeM rep) ()
checkDim (ShapeBase ExtSize -> [ExtSize]
forall d. ShapeBase d -> [d]
shapeDims ShapeBase ExtSize
x_shape) (Shape -> [SubExp]
forall d. ShapeBase d -> [d]
shapeDims Shape
y_shape)
MemReturn -> MemBind -> ExceptT String (TypeM rep) ()
checkMemReturn MemReturn
x_ret MemBind
y_ret
checkReturn MemInfo ExtSize u MemReturn
x LetDecMem
y =
String -> ExceptT String (TypeM rep) ()
forall e (m :: * -> *) a. MonadError e m => e -> m a
throwError (String -> ExceptT String (TypeM rep) ())
-> String -> ExceptT String (TypeM rep) ()
forall a b. (a -> b) -> a -> b
$ [String] -> String
unwords [String
"Expected", MemInfo ExtSize u MemReturn -> String
forall a. Pretty a => a -> String
pretty MemInfo ExtSize u MemReturn
x, String
"but got", LetDecMem -> String
forall a. Pretty a => a -> String
pretty LetDecMem
y]
checkDim :: ExtSize -> SubExp -> ExceptT String (TypeM rep) ()
checkDim (Free SubExp
x) SubExp
y
| SubExp
x SubExp -> SubExp -> Bool
forall a. Eq a => a -> a -> Bool
== SubExp
y = () -> ExceptT String (TypeM rep) ()
forall (f :: * -> *) a. Applicative f => a -> f a
pure ()
| Bool
otherwise =
String -> ExceptT String (TypeM rep) ()
forall e (m :: * -> *) a. MonadError e m => e -> m a
throwError (String -> ExceptT String (TypeM rep) ())
-> String -> ExceptT String (TypeM rep) ()
forall a b. (a -> b) -> a -> b
$ [String] -> String
unwords [String
"Expected dim", SubExp -> String
forall a. Pretty a => a -> String
pretty SubExp
x, String
"but got", SubExp -> String
forall a. Pretty a => a -> String
pretty SubExp
y]
checkDim (Ext Int
i) SubExp
y = do
(SubExp
x, LetDecMem
_) <- Int -> ExceptT String (TypeM rep) (SubExp, LetDecMem)
fetchCtx Int
i
Bool
-> ExceptT String (TypeM rep) () -> ExceptT String (TypeM rep) ()
forall (f :: * -> *). Applicative f => Bool -> f () -> f ()
unless (SubExp
x SubExp -> SubExp -> Bool
forall a. Eq a => a -> a -> Bool
== SubExp
y) (ExceptT String (TypeM rep) () -> ExceptT String (TypeM rep) ())
-> ([String] -> ExceptT String (TypeM rep) ())
-> [String]
-> ExceptT String (TypeM rep) ()
forall k (cat :: k -> k -> *) (b :: k) (c :: k) (a :: k).
Category cat =>
cat b c -> cat a b -> cat a c
. String -> ExceptT String (TypeM rep) ()
forall e (m :: * -> *) a. MonadError e m => e -> m a
throwError (String -> ExceptT String (TypeM rep) ())
-> ([String] -> String)
-> [String]
-> ExceptT String (TypeM rep) ()
forall k (cat :: k -> k -> *) (b :: k) (c :: k) (a :: k).
Category cat =>
cat b c -> cat a b -> cat a c
. [String] -> String
unwords ([String] -> ExceptT String (TypeM rep) ())
-> [String] -> ExceptT String (TypeM rep) ()
forall a b. (a -> b) -> a -> b
$
[String
"Expected ext dim", Int -> String
forall a. Pretty a => a -> String
pretty Int
i, String
"=>", SubExp -> String
forall a. Pretty a => a -> String
pretty SubExp
x, String
"but got", SubExp -> String
forall a. Pretty a => a -> String
pretty SubExp
y]
checkMemReturn :: MemReturn -> MemBind -> ExceptT String (TypeM rep) ()
checkMemReturn (ReturnsInBlock VName
x_mem ExtIxFun
x_ixfun) (ArrayIn VName
y_mem IxFun
y_ixfun)
| VName
x_mem VName -> VName -> Bool
forall a. Eq a => a -> a -> Bool
== VName
y_mem =
Bool
-> ExceptT String (TypeM rep) () -> ExceptT String (TypeM rep) ()
forall (f :: * -> *). Applicative f => Bool -> f () -> f ()
unless (ExtIxFun -> ExtIxFun -> Bool
forall num. IxFun num -> IxFun num -> Bool
IxFun.closeEnough ExtIxFun
x_ixfun (ExtIxFun -> Bool) -> ExtIxFun -> Bool
forall a b. (a -> b) -> a -> b
$ IxFun -> ExtIxFun
existentialiseIxFun0 IxFun
y_ixfun) (ExceptT String (TypeM rep) () -> ExceptT String (TypeM rep) ())
-> ExceptT String (TypeM rep) () -> ExceptT String (TypeM rep) ()
forall a b. (a -> b) -> a -> b
$
String -> ExceptT String (TypeM rep) ()
forall e (m :: * -> *) a. MonadError e m => e -> m a
throwError (String -> ExceptT String (TypeM rep) ())
-> ([String] -> String)
-> [String]
-> ExceptT String (TypeM rep) ()
forall k (cat :: k -> k -> *) (b :: k) (c :: k) (a :: k).
Category cat =>
cat b c -> cat a b -> cat a c
. [String] -> String
unwords ([String] -> ExceptT String (TypeM rep) ())
-> [String] -> ExceptT String (TypeM rep) ()
forall a b. (a -> b) -> a -> b
$
[ String
"Index function unification failed (ReturnsInBlock)",
String
"\nixfun of body result: ",
IxFun -> String
forall a. Pretty a => a -> String
pretty IxFun
y_ixfun,
String
"\nixfun of return type: ",
ExtIxFun -> String
forall a. Pretty a => a -> String
pretty ExtIxFun
x_ixfun
]
checkMemReturn
(ReturnsNewBlock Space
x_space Int
x_ext ExtIxFun
x_ixfun)
(ArrayIn VName
y_mem IxFun
y_ixfun) = do
(SubExp
x_mem, LetDecMem
x_mem_type) <- Int -> ExceptT String (TypeM rep) (SubExp, LetDecMem)
fetchCtx Int
x_ext
Bool
-> ExceptT String (TypeM rep) () -> ExceptT String (TypeM rep) ()
forall (f :: * -> *). Applicative f => Bool -> f () -> f ()
unless (ExtIxFun -> ExtIxFun -> Bool
forall num. IxFun num -> IxFun num -> Bool
IxFun.closeEnough ExtIxFun
x_ixfun (ExtIxFun -> Bool) -> ExtIxFun -> Bool
forall a b. (a -> b) -> a -> b
$ IxFun -> ExtIxFun
existentialiseIxFun0 IxFun
y_ixfun) (ExceptT String (TypeM rep) () -> ExceptT String (TypeM rep) ())
-> ExceptT String (TypeM rep) () -> ExceptT String (TypeM rep) ()
forall a b. (a -> b) -> a -> b
$
String -> ExceptT String (TypeM rep) ()
forall e (m :: * -> *) a. MonadError e m => e -> m a
throwError (String -> ExceptT String (TypeM rep) ())
-> (Doc -> String) -> Doc -> ExceptT String (TypeM rep) ()
forall k (cat :: k -> k -> *) (b :: k) (c :: k) (a :: k).
Category cat =>
cat b c -> cat a b -> cat a c
. Doc -> String
forall a. Pretty a => a -> String
pretty (Doc -> ExceptT String (TypeM rep) ())
-> Doc -> ExceptT String (TypeM rep) ()
forall a b. (a -> b) -> a -> b
$
Doc
"Index function unification failed (ReturnsNewBlock)"
Doc -> Doc -> Doc
</> Doc
"Ixfun of body result:"
Doc -> Doc -> Doc
</> Int -> Doc -> Doc
indent Int
2 (IxFun -> Doc
forall a. Pretty a => a -> Doc
ppr IxFun
y_ixfun)
Doc -> Doc -> Doc
</> Doc
"Ixfun of return type:"
Doc -> Doc -> Doc
</> Int -> Doc -> Doc
indent Int
2 (ExtIxFun -> Doc
forall a. Pretty a => a -> Doc
ppr ExtIxFun
x_ixfun)
case LetDecMem
x_mem_type of
MemMem Space
y_space ->
Bool
-> ExceptT String (TypeM rep) () -> ExceptT String (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 String (TypeM rep) () -> ExceptT String (TypeM rep) ())
-> ([String] -> ExceptT String (TypeM rep) ())
-> [String]
-> ExceptT String (TypeM rep) ()
forall k (cat :: k -> k -> *) (b :: k) (c :: k) (a :: k).
Category cat =>
cat b c -> cat a b -> cat a c
. String -> ExceptT String (TypeM rep) ()
forall e (m :: * -> *) a. MonadError e m => e -> m a
throwError (String -> ExceptT String (TypeM rep) ())
-> ([String] -> String)
-> [String]
-> ExceptT String (TypeM rep) ()
forall k (cat :: k -> k -> *) (b :: k) (c :: k) (a :: k).
Category cat =>
cat b c -> cat a b -> cat a c
. [String] -> String
unwords ([String] -> ExceptT String (TypeM rep) ())
-> [String] -> ExceptT String (TypeM rep) ()
forall a b. (a -> b) -> a -> b
$
[ String
"Expected memory",
VName -> String
forall a. Pretty a => a -> String
pretty VName
y_mem,
String
"in space",
Space -> String
forall a. Pretty a => a -> String
pretty Space
x_space,
String
"but actually in space",
Space -> String
forall a. Pretty a => a -> String
pretty Space
y_space
]
LetDecMem
t ->
String -> ExceptT String (TypeM rep) ()
forall e (m :: * -> *) a. MonadError e m => e -> m a
throwError (String -> ExceptT String (TypeM rep) ())
-> ([String] -> String)
-> [String]
-> ExceptT String (TypeM rep) ()
forall k (cat :: k -> k -> *) (b :: k) (c :: k) (a :: k).
Category cat =>
cat b c -> cat a b -> cat a c
. [String] -> String
unwords ([String] -> ExceptT String (TypeM rep) ())
-> [String] -> ExceptT String (TypeM rep) ()
forall a b. (a -> b) -> a -> b
$
[String
"Expected memory", Int -> String
forall a. Pretty a => a -> String
pretty Int
x_ext, String
"=>", SubExp -> String
forall a. Pretty a => a -> String
pretty SubExp
x_mem, String
"but but has type", LetDecMem -> String
forall a. Pretty a => a -> String
pretty LetDecMem
t]
checkMemReturn MemReturn
x MemBind
y =
String -> ExceptT String (TypeM rep) ()
forall e (m :: * -> *) a. MonadError e m => e -> m a
throwError (String -> ExceptT String (TypeM rep) ())
-> (Doc -> String) -> Doc -> ExceptT String (TypeM rep) ()
forall k (cat :: k -> k -> *) (b :: k) (c :: k) (a :: k).
Category cat =>
cat b c -> cat a b -> cat a c
. Doc -> String
forall a. Pretty a => a -> String
pretty (Doc -> ExceptT String (TypeM rep) ())
-> Doc -> ExceptT String (TypeM rep) ()
forall a b. (a -> b) -> a -> b
$
Doc
"Expected array in"
Doc -> Doc -> Doc
</> Int -> Doc -> Doc
indent Int
2 (MemReturn -> Doc
forall a. Pretty a => a -> Doc
ppr MemReturn
x)
Doc -> Doc -> Doc
</> Doc
"but array returned in"
Doc -> Doc -> Doc
</> Int -> Doc -> Doc
indent Int
2 (MemBind -> Doc
forall a. Pretty a => a -> Doc
ppr MemBind
y)
bad :: String -> TC.TypeM rep a
bad :: String -> TypeM rep a
bad String
s =
ErrorCase rep -> TypeM rep a
forall rep a. ErrorCase rep -> TypeM rep a
TC.bad (ErrorCase rep -> TypeM rep a) -> ErrorCase rep -> TypeM rep a
forall a b. (a -> b) -> a -> b
$
String -> ErrorCase rep
forall rep. String -> ErrorCase rep
TC.TypeError (String -> ErrorCase rep) -> String -> ErrorCase rep
forall a b. (a -> b) -> a -> b
$
Doc -> String
forall a. Pretty a => a -> String
pretty (Doc -> String) -> Doc -> String
forall a b. (a -> b) -> a -> b
$
Doc
"Return type"
Doc -> Doc -> Doc
</> Int -> Doc -> Doc
indent Int
2 ([MemInfo ExtSize u MemReturn] -> Doc
forall a. Pretty a => [a] -> Doc
ppTuple' [MemInfo ExtSize u MemReturn]
rettype)
Doc -> Doc -> Doc
</> Doc
"cannot match returns of results"
Doc -> Doc -> Doc
</> Int -> Doc -> Doc
indent Int
2 ([LetDecMem] -> Doc
forall a. Pretty a => [a] -> Doc
ppTuple' [LetDecMem]
ts)
Doc -> Doc -> Doc
</> String -> Doc
text String
s
(String -> TypeM rep ())
-> (() -> TypeM rep ()) -> Either String () -> TypeM rep ()
forall a c b. (a -> c) -> (b -> c) -> Either a b -> c
either String -> TypeM rep ()
forall rep a. String -> TypeM rep a
bad () -> TypeM rep ()
forall (f :: * -> *) a. Applicative f => a -> f a
pure (Either String () -> TypeM rep ())
-> TypeM rep (Either String ()) -> TypeM rep ()
forall (m :: * -> *) a b. Monad m => (a -> m b) -> m a -> m b
=<< ExceptT String (TypeM rep) () -> TypeM rep (Either String ())
forall e (m :: * -> *) a. ExceptT e m a -> m (Either e a)
runExceptT ((MemInfo ExtSize u MemReturn
-> LetDecMem -> ExceptT String (TypeM rep) ())
-> [MemInfo ExtSize u MemReturn]
-> [LetDecMem]
-> ExceptT String (TypeM rep) ()
forall (m :: * -> *) a b c.
Applicative m =>
(a -> b -> m c) -> [a] -> [b] -> m ()
zipWithM_ MemInfo ExtSize u MemReturn
-> LetDecMem -> ExceptT String (TypeM rep) ()
checkReturn [MemInfo ExtSize u 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 :: Pat (LetDec (Aliases rep)) -> Exp (Aliases rep) -> TypeM rep ()
matchPatToExp Pat (LetDec (Aliases rep))
pat Exp (Aliases rep)
e = do
Scope rep
scope <- (Scope (Aliases rep) -> Scope rep) -> TypeM rep (Scope rep)
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
[ExpReturns]
rt <- ReaderT (Scope rep) (TypeM rep) [ExpReturns]
-> Scope rep -> TypeM rep [ExpReturns]
forall r (m :: * -> *) a. ReaderT r m a -> r -> m a
runReaderT (Exp rep -> ReaderT (Scope rep) (TypeM rep) [ExpReturns]
forall (m :: * -> *) rep inner.
(Monad m, LocalScope rep m, Mem rep inner) =>
Exp rep -> m [ExpReturns]
expReturns (Exp rep -> ReaderT (Scope rep) (TypeM rep) [ExpReturns])
-> Exp rep -> ReaderT (Scope rep) (TypeM rep) [ExpReturns]
forall a b. (a -> b) -> a -> b
$ Exp (Aliases rep) -> Exp rep
forall rep. CanBeAliased (Op rep) => Exp (Aliases rep) -> Exp rep
removeExpAliases Exp (Aliases rep)
e) Scope rep
scope
let ([VName]
ctx_ids, [BodyReturns]
val_ts) = [(VName, BodyReturns)] -> ([VName], [BodyReturns])
forall a b. [(a, b)] -> ([a], [b])
unzip ([(VName, BodyReturns)] -> ([VName], [BodyReturns]))
-> [(VName, BodyReturns)] -> ([VName], [BodyReturns])
forall a b. (a -> b) -> a -> b
$ Pat LetDecMem -> [(VName, BodyReturns)]
bodyReturnsFromPat (Pat LetDecMem -> [(VName, BodyReturns)])
-> Pat LetDecMem -> [(VName, BodyReturns)]
forall a b. (a -> b) -> a -> b
$ Pat (AliasDec, LetDecMem) -> Pat LetDecMem
forall a. Pat (AliasDec, a) -> Pat a
removePatAliases Pat (AliasDec, LetDecMem)
Pat (LetDec (Aliases rep))
pat
(Map (Ext VName) (TPrimExp Int64 (Ext VName))
ctx_map_ids, Map (Ext VName) (TPrimExp Int64 (Ext VName))
ctx_map_exts) = [(VName, Int)]
-> (Map (Ext VName) (TPrimExp Int64 (Ext VName)),
Map (Ext VName) (TPrimExp Int64 (Ext VName)))
getExtMaps ([(VName, Int)]
-> (Map (Ext VName) (TPrimExp Int64 (Ext VName)),
Map (Ext VName) (TPrimExp Int64 (Ext VName))))
-> [(VName, Int)]
-> (Map (Ext VName) (TPrimExp Int64 (Ext VName)),
Map (Ext VName) (TPrimExp Int64 (Ext VName)))
forall a b. (a -> b) -> a -> b
$ [VName] -> [Int] -> [(VName, Int)]
forall a b. [a] -> [b] -> [(a, b)]
zip [VName]
ctx_ids [Int
0 .. Int
1]
Bool -> TypeM rep () -> TypeM rep ()
forall (f :: * -> *). Applicative f => Bool -> f () -> f ()
unless
( [BodyReturns] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length [BodyReturns]
val_ts Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
== [ExpReturns] -> 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 ((BodyReturns -> ExpReturns -> Bool)
-> [BodyReturns] -> [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))
-> BodyReturns
-> ExpReturns
-> Bool
forall d u u.
Eq d =>
Map (Ext VName) (TPrimExp Int64 (Ext VName))
-> Map (Ext VName) (TPrimExp Int64 (Ext VName))
-> MemInfo d u MemReturn
-> MemInfo d u (Maybe MemReturn)
-> Bool
matches Map (Ext VName) (TPrimExp Int64 (Ext VName))
ctx_map_ids Map (Ext VName) (TPrimExp Int64 (Ext VName))
ctx_map_exts) [BodyReturns]
val_ts [ExpReturns]
rt)
)
(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 ()) -> ErrorCase rep -> TypeM rep ()
forall a b. (a -> b) -> a -> b
$ String -> ErrorCase rep
forall rep. String -> ErrorCase rep
TC.TypeError
(String -> ErrorCase rep) -> String -> ErrorCase rep
forall a b. (a -> b) -> a -> b
$ String
"Expression type:\n "
String -> ShowS
forall a. [a] -> [a] -> [a]
++ [ExpReturns] -> String
forall a. Pretty a => [a] -> String
prettyTuple [ExpReturns]
rt
String -> ShowS
forall a. [a] -> [a] -> [a]
++ String
"\ncannot match pattern type:\n "
String -> ShowS
forall a. [a] -> [a] -> [a]
++ [BodyReturns] -> String
forall a. Pretty a => [a] -> String
prettyTuple [BodyReturns]
val_ts
where
matches :: Map (Ext VName) (TPrimExp Int64 (Ext VName))
-> Map (Ext VName) (TPrimExp Int64 (Ext VName))
-> MemInfo d u MemReturn
-> MemInfo d u (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 Shape
x_ispace [Type]
x_ts u
_) (MemAcc VName
y_accs Shape
y_ispace [Type]
y_ts u
_) =
(VName
x_accs, Shape
x_ispace, [Type]
x_ts) (VName, Shape, [Type]) -> (VName, Shape, [Type]) -> Bool
forall a. Eq a => a -> a -> Bool
== (VName
y_accs, Shape
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 u
_ MemReturn
x_ret) (MemArray PrimType
y_pt ShapeBase d
y_shape u
_ 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
_ ExtIxFun
x_ixfun, Just (ReturnsInBlock VName
_ ExtIxFun
y_ixfun)) ->
let x_ixfun' :: ExtIxFun
x_ixfun' = Map (Ext VName) (TPrimExp Int64 (Ext VName))
-> ExtIxFun -> ExtIxFun
forall a t.
Ord a =>
Map a (TPrimExp t a)
-> IxFun (TPrimExp t a) -> IxFun (TPrimExp t a)
IxFun.substituteInIxFun Map (Ext VName) (TPrimExp Int64 (Ext VName))
ctxids ExtIxFun
x_ixfun
y_ixfun' :: ExtIxFun
y_ixfun' = Map (Ext VName) (TPrimExp Int64 (Ext VName))
-> ExtIxFun -> ExtIxFun
forall a t.
Ord a =>
Map a (TPrimExp t a)
-> IxFun (TPrimExp t a) -> IxFun (TPrimExp t a)
IxFun.substituteInIxFun Map (Ext VName) (TPrimExp Int64 (Ext VName))
ctxexts ExtIxFun
y_ixfun
in ExtIxFun -> ExtIxFun -> Bool
forall num. IxFun num -> IxFun num -> Bool
IxFun.closeEnough ExtIxFun
x_ixfun' ExtIxFun
y_ixfun'
( ReturnsInBlock VName
_ ExtIxFun
x_ixfun,
Just (ReturnsNewBlock Space
_ Int
_ ExtIxFun
y_ixfun)
) ->
let x_ixfun' :: ExtIxFun
x_ixfun' = Map (Ext VName) (TPrimExp Int64 (Ext VName))
-> ExtIxFun -> ExtIxFun
forall a t.
Ord a =>
Map a (TPrimExp t a)
-> IxFun (TPrimExp t a) -> IxFun (TPrimExp t a)
IxFun.substituteInIxFun Map (Ext VName) (TPrimExp Int64 (Ext VName))
ctxids ExtIxFun
x_ixfun
y_ixfun' :: ExtIxFun
y_ixfun' = Map (Ext VName) (TPrimExp Int64 (Ext VName))
-> ExtIxFun -> ExtIxFun
forall a t.
Ord a =>
Map a (TPrimExp t a)
-> IxFun (TPrimExp t a) -> IxFun (TPrimExp t a)
IxFun.substituteInIxFun Map (Ext VName) (TPrimExp Int64 (Ext VName))
ctxexts ExtIxFun
y_ixfun
in ExtIxFun -> ExtIxFun -> Bool
forall num. IxFun num -> IxFun num -> Bool
IxFun.closeEnough ExtIxFun
x_ixfun' ExtIxFun
y_ixfun'
( ReturnsNewBlock Space
_ Int
x_i ExtIxFun
x_ixfun,
Just (ReturnsNewBlock Space
_ Int
y_i ExtIxFun
y_ixfun)
) ->
let x_ixfun' :: ExtIxFun
x_ixfun' = Map (Ext VName) (TPrimExp Int64 (Ext VName))
-> ExtIxFun -> ExtIxFun
forall a t.
Ord a =>
Map a (TPrimExp t a)
-> IxFun (TPrimExp t a) -> IxFun (TPrimExp t a)
IxFun.substituteInIxFun Map (Ext VName) (TPrimExp Int64 (Ext VName))
ctxids ExtIxFun
x_ixfun
y_ixfun' :: ExtIxFun
y_ixfun' = Map (Ext VName) (TPrimExp Int64 (Ext VName))
-> ExtIxFun -> ExtIxFun
forall a t.
Ord a =>
Map a (TPrimExp t a)
-> IxFun (TPrimExp t a) -> IxFun (TPrimExp t a)
IxFun.substituteInIxFun Map (Ext VName) (TPrimExp Int64 (Ext VName))
ctxexts ExtIxFun
y_ixfun
in Int
x_i Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
== Int
y_i Bool -> Bool -> Bool
&& ExtIxFun -> ExtIxFun -> Bool
forall num. IxFun num -> IxFun num -> Bool
IxFun.closeEnough ExtIxFun
x_ixfun' ExtIxFun
y_ixfun'
(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 u MemReturn
_ MemInfo d u (Maybe MemReturn)
_ = Bool
False
varMemInfo ::
Mem rep inner =>
VName ->
TC.TypeM rep (MemInfo SubExp NoUniqueness MemBind)
varMemInfo :: VName -> TypeM rep LetDecMem
varMemInfo VName
name = do
NameInfo (Aliases rep)
dec <- VName -> TypeM rep (NameInfo (Aliases rep))
forall rep. VName -> TypeM rep (NameInfo (Aliases rep))
TC.lookupVar VName
name
case NameInfo (Aliases rep)
dec of
LetName (_, summary) -> LetDecMem -> TypeM rep LetDecMem
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 (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 u r. MemInfo d u r -> MemInfo d NoUniqueness r
noUniquenessReturns FParamInfo (Aliases rep)
FParamMem
summary
LParamName LParamInfo (Aliases rep)
summary -> LetDecMem -> TypeM rep LetDecMem
forall (f :: * -> *) a. Applicative f => a -> f a
pure LParamInfo (Aliases rep)
LetDecMem
summary
IndexName IntType
it -> LetDecMem -> TypeM rep LetDecMem
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 u ret. PrimType -> MemInfo d u ret
MemPrim (PrimType -> LetDecMem) -> PrimType -> LetDecMem
forall a b. (a -> b) -> a -> b
$ IntType -> PrimType
IntType IntType
it
nameInfoToMemInfo :: Mem rep inner => NameInfo rep -> MemBound NoUniqueness
nameInfoToMemInfo :: NameInfo rep -> LetDecMem
nameInfoToMemInfo NameInfo rep
info =
case NameInfo rep
info of
FParamName FParamInfo rep
summary -> FParamMem -> LetDecMem
forall d u r. MemInfo d u r -> MemInfo d NoUniqueness r
noUniquenessReturns 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 u ret. PrimType -> MemInfo d u ret
MemPrim (PrimType -> LetDecMem) -> PrimType -> LetDecMem
forall a b. (a -> b) -> a -> b
$ IntType -> PrimType
IntType IntType
it
lookupMemInfo ::
(HasScope rep m, Mem rep inner) =>
VName ->
m (MemInfo SubExp NoUniqueness MemBind)
lookupMemInfo :: VName -> m LetDecMem
lookupMemInfo = (NameInfo rep -> LetDecMem) -> m (NameInfo rep) -> m LetDecMem
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 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 NoUniqueness MemBind)
subExpMemInfo :: 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 (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 u ret. PrimType -> MemInfo d u 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, IxFun.IxFun (TPrimExp Int64 VName))
lookupArraySummary :: VName -> m (VName, IxFun)
lookupArraySummary VName
name = do
LetDecMem
summary <- VName -> m LetDecMem
forall rep (m :: * -> *) inner.
(HasScope rep m, Mem rep inner) =>
VName -> m LetDecMem
lookupMemInfo VName
name
case LetDecMem
summary of
MemArray PrimType
_ Shape
_ NoUniqueness
_ (ArrayIn VName
mem IxFun
ixfun) ->
(VName, IxFun) -> m (VName, IxFun)
forall (f :: * -> *) a. Applicative f => a -> f a
pure (VName
mem, IxFun
ixfun)
LetDecMem
_ ->
String -> m (VName, IxFun)
forall a. HasCallStack => String -> a
error (String -> m (VName, IxFun)) -> String -> m (VName, IxFun)
forall a b. (a -> b) -> a -> b
$
String
"Expected "
String -> ShowS
forall a. [a] -> [a] -> [a]
++ VName -> String
forall a. Pretty a => a -> String
pretty VName
name
String -> ShowS
forall a. [a] -> [a] -> [a]
++ String
" to be array but bound to:\n"
String -> ShowS
forall a. [a] -> [a] -> [a]
++ LetDecMem -> String
forall a. Pretty a => a -> String
pretty LetDecMem
summary
checkMemInfo ::
TC.Checkable rep =>
VName ->
MemInfo SubExp u MemBind ->
TC.TypeM rep ()
checkMemInfo :: VName -> MemInfo SubExp u MemBind -> TypeM rep ()
checkMemInfo VName
_ (MemPrim PrimType
_) = () -> TypeM rep ()
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 u. PrimType -> TypeBase shape u
Prim PrimType
int64]) [SubExp]
d
checkMemInfo VName
_ (MemMem Space
_) = () -> TypeM rep ()
forall (f :: * -> *) a. Applicative f => a -> f a
pure ()
checkMemInfo VName
_ (MemAcc VName
acc Shape
ispace [Type]
ts u
u) =
TypeBase Shape u -> TypeM rep ()
forall rep u. Checkable rep => TypeBase Shape u -> TypeM rep ()
TC.checkType (TypeBase Shape u -> TypeM rep ())
-> TypeBase Shape u -> TypeM rep ()
forall a b. (a -> b) -> a -> b
$ VName -> Shape -> [Type] -> u -> TypeBase Shape u
forall shape u. VName -> Shape -> [Type] -> u -> TypeBase shape u
Acc VName
acc Shape
ispace [Type]
ts u
u
checkMemInfo VName
name (MemArray PrimType
_ Shape
shape u
_ (ArrayIn VName
v IxFun
ixfun)) = do
Type
t <- VName -> TypeM rep Type
forall rep (m :: * -> *). HasScope rep m => VName -> m Type
lookupType VName
v
case Type
t of
Mem {} ->
() -> TypeM rep ()
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
$
String -> ErrorCase rep
forall rep. String -> ErrorCase rep
TC.TypeError (String -> ErrorCase rep) -> String -> ErrorCase rep
forall a b. (a -> b) -> a -> b
$
String
"Variable "
String -> ShowS
forall a. [a] -> [a] -> [a]
++ VName -> String
forall a. Pretty a => a -> String
pretty VName
v
String -> ShowS
forall a. [a] -> [a] -> [a]
++ String
" used as memory block, but is of type "
String -> ShowS
forall a. [a] -> [a] -> [a]
++ Type -> String
forall a. Pretty a => a -> String
pretty Type
t
String -> ShowS
forall a. [a] -> [a] -> [a]
++ String
"."
String -> TypeM rep () -> TypeM rep ()
forall rep a. String -> TypeM rep a -> TypeM rep a
TC.context (String
"in index function " String -> ShowS
forall a. [a] -> [a] -> [a]
++ IxFun -> String
forall a. Pretty a => a -> String
pretty IxFun
ixfun) (TypeM rep () -> TypeM rep ()) -> TypeM rep () -> TypeM rep ()
forall a b. (a -> b) -> a -> b
$ do
(TPrimExp Int64 VName -> TypeM rep ()) -> IxFun -> TypeM rep ()
forall (t :: * -> *) (f :: * -> *) a b.
(Foldable t, Applicative f) =>
(a -> f b) -> t a -> f ()
traverse_ (PrimType -> PrimExp VName -> TypeM rep ()
forall rep.
Checkable rep =>
PrimType -> PrimExp VName -> TypeM rep ()
TC.requirePrimExp PrimType
int64 (PrimExp VName -> TypeM rep ())
-> (TPrimExp Int64 VName -> PrimExp VName)
-> TPrimExp Int64 VName
-> TypeM rep ()
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 t v. TPrimExp t v -> PrimExp v
untyped) IxFun
ixfun
let ixfun_rank :: Int
ixfun_rank = IxFun -> Int
forall num. IntegralExp num => IxFun num -> Int
IxFun.rank IxFun
ixfun
ident_rank :: Int
ident_rank = Shape -> Int
forall a. ArrayShape a => a -> Int
shapeRank Shape
shape
Bool -> TypeM rep () -> TypeM rep ()
forall (f :: * -> *). Applicative f => Bool -> f () -> f ()
unless (Int
ixfun_rank Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
== Int
ident_rank) (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 ()) -> ErrorCase rep -> TypeM rep ()
forall a b. (a -> b) -> a -> b
$
String -> ErrorCase rep
forall rep. String -> ErrorCase rep
TC.TypeError (String -> ErrorCase rep) -> String -> ErrorCase rep
forall a b. (a -> b) -> a -> b
$
String
"Arity of index function ("
String -> ShowS
forall a. [a] -> [a] -> [a]
++ Int -> String
forall a. Pretty a => a -> String
pretty Int
ixfun_rank
String -> ShowS
forall a. [a] -> [a] -> [a]
++ String
") does not match rank of array "
String -> ShowS
forall a. [a] -> [a] -> [a]
++ VName -> String
forall a. Pretty a => a -> String
pretty VName
name
String -> ShowS
forall a. [a] -> [a] -> [a]
++ String
" ("
String -> ShowS
forall a. [a] -> [a] -> [a]
++ Int -> String
forall a. Show a => a -> String
show Int
ident_rank
String -> ShowS
forall a. [a] -> [a] -> [a]
++ String
")"
bodyReturnsFromPat ::
Pat (MemBound NoUniqueness) -> [(VName, BodyReturns)]
bodyReturnsFromPat :: Pat LetDecMem -> [(VName, BodyReturns)]
bodyReturnsFromPat Pat LetDecMem
pat =
(PatElem LetDecMem -> (VName, BodyReturns))
-> [PatElem LetDecMem] -> [(VName, BodyReturns)]
forall a b. (a -> b) -> [a] -> [b]
map PatElem LetDecMem -> (VName, BodyReturns)
asReturns ([PatElem LetDecMem] -> [(VName, BodyReturns)])
-> [PatElem LetDecMem] -> [(VName, BodyReturns)]
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 -> ExtSize
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 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 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 -> ExtSize
forall a. Int -> Ext a
Ext Int
i
ext SubExp
se = SubExp -> ExtSize
forall a. a -> Ext a
Free SubExp
se
asReturns :: PatElem LetDecMem -> (VName, BodyReturns)
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 -> BodyReturns
forall d u ret. PrimType -> MemInfo d u ret
MemPrim PrimType
pt
MemMem Space
space -> Space -> BodyReturns
forall d u ret. Space -> MemInfo d u ret
MemMem Space
space
MemArray PrimType
pt Shape
shape NoUniqueness
u (ArrayIn VName
mem IxFun
ixfun) ->
PrimType
-> ShapeBase ExtSize -> NoUniqueness -> MemReturn -> BodyReturns
forall d u ret.
PrimType -> ShapeBase d -> u -> ret -> MemInfo d u ret
MemArray PrimType
pt ([ExtSize] -> ShapeBase ExtSize
forall d. [d] -> ShapeBase d
Shape ([ExtSize] -> ShapeBase ExtSize) -> [ExtSize] -> ShapeBase ExtSize
forall a b. (a -> b) -> a -> b
$ (SubExp -> ExtSize) -> [SubExp] -> [ExtSize]
forall a b. (a -> b) -> [a] -> [b]
map SubExp -> ExtSize
ext ([SubExp] -> [ExtSize]) -> [SubExp] -> [ExtSize]
forall a b. (a -> b) -> a -> b
$ Shape -> [SubExp]
forall d. ShapeBase d -> [d]
shapeDims Shape
shape) NoUniqueness
u (MemReturn -> BodyReturns) -> MemReturn -> BodyReturns
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 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 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 -> ExtIxFun -> MemReturn
ReturnsNewBlock Space
space Int
i (ExtIxFun -> MemReturn) -> ExtIxFun -> MemReturn
forall a b. (a -> b) -> a -> b
$
[VName] -> IxFun -> ExtIxFun
existentialiseIxFun ((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) IxFun
ixfun
Maybe (Int, PatElem LetDecMem)
_ -> VName -> ExtIxFun -> MemReturn
ReturnsInBlock VName
mem (ExtIxFun -> MemReturn) -> ExtIxFun -> MemReturn
forall a b. (a -> b) -> a -> b
$ [VName] -> IxFun -> ExtIxFun
existentialiseIxFun [] IxFun
ixfun
MemAcc VName
acc Shape
ispace [Type]
ts NoUniqueness
u -> VName -> Shape -> [Type] -> NoUniqueness -> BodyReturns
forall d u ret. VName -> Shape -> [Type] -> u -> MemInfo d u ret
MemAcc VName
acc Shape
ispace [Type]
ts NoUniqueness
u
)
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)
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 (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 u ret. PrimType -> MemInfo d u ret
MemPrim PrimType
bt
addDec (Mem Space
space) =
ExpReturns -> f ExpReturns
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 u ret. Space -> MemInfo d u ret
MemMem Space
space
addDec t :: ExtType
t@(Array PrimType
bt ShapeBase ExtSize
shape NoUniqueness
u)
| ExtType -> Bool
existential ExtType
t = do
Int
i <- f Int
forall s (m :: * -> *). MonadState s m => m s
get f Int -> f () -> f Int
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)
ExpReturns -> f ExpReturns
forall (f :: * -> *) a. Applicative f => a -> f a
pure (ExpReturns -> f ExpReturns) -> ExpReturns -> f ExpReturns
forall a b. (a -> b) -> a -> b
$
PrimType
-> ShapeBase ExtSize
-> NoUniqueness
-> Maybe MemReturn
-> ExpReturns
forall d u ret.
PrimType -> ShapeBase d -> u -> ret -> MemInfo d u ret
MemArray PrimType
bt ShapeBase ExtSize
shape NoUniqueness
u (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 -> ExtIxFun -> MemReturn
ReturnsNewBlock Space
DefaultSpace Int
i (ExtIxFun -> MemReturn) -> ExtIxFun -> MemReturn
forall a b. (a -> b) -> a -> b
$
Shape (TPrimExp Int64 (Ext VName)) -> ExtIxFun
forall num. IntegralExp num => Shape num -> IxFun num
IxFun.iota (Shape (TPrimExp Int64 (Ext VName)) -> ExtIxFun)
-> Shape (TPrimExp Int64 (Ext VName)) -> ExtIxFun
forall a b. (a -> b) -> a -> b
$
(ExtSize -> TPrimExp Int64 (Ext VName))
-> [ExtSize] -> Shape (TPrimExp Int64 (Ext VName))
forall a b. (a -> b) -> [a] -> [b]
map ExtSize -> TPrimExp Int64 (Ext VName)
convert ([ExtSize] -> Shape (TPrimExp Int64 (Ext VName)))
-> [ExtSize] -> Shape (TPrimExp Int64 (Ext VName))
forall a b. (a -> b) -> a -> b
$
ShapeBase ExtSize -> [ExtSize]
forall d. ShapeBase d -> [d]
shapeDims ShapeBase ExtSize
shape
| Bool
otherwise =
ExpReturns -> f ExpReturns
forall (f :: * -> *) a. Applicative f => a -> f a
pure (ExpReturns -> f ExpReturns) -> ExpReturns -> f ExpReturns
forall a b. (a -> b) -> a -> b
$ PrimType
-> ShapeBase ExtSize
-> NoUniqueness
-> Maybe MemReturn
-> ExpReturns
forall d u ret.
PrimType -> ShapeBase d -> u -> ret -> MemInfo d u ret
MemArray PrimType
bt ShapeBase ExtSize
shape NoUniqueness
u Maybe MemReturn
forall a. Maybe a
Nothing
addDec (Acc VName
acc Shape
ispace [Type]
ts NoUniqueness
u) =
ExpReturns -> f ExpReturns
forall (f :: * -> *) a. Applicative f => a -> f a
pure (ExpReturns -> f ExpReturns) -> ExpReturns -> f ExpReturns
forall a b. (a -> b) -> a -> b
$ VName -> Shape -> [Type] -> NoUniqueness -> ExpReturns
forall d u ret. VName -> Shape -> [Type] -> u -> MemInfo d u ret
MemAcc VName
acc Shape
ispace [Type]
ts NoUniqueness
u
convert :: ExtSize -> 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, IxFun)
arrayVarReturns :: VName -> m (PrimType, Shape, VName, IxFun)
arrayVarReturns VName
v = do
LetDecMem
summary <- VName -> m LetDecMem
forall rep (m :: * -> *) inner.
(HasScope rep m, Mem rep inner) =>
VName -> m LetDecMem
lookupMemInfo VName
v
case LetDecMem
summary of
MemArray PrimType
et Shape
shape NoUniqueness
_ (ArrayIn VName
mem IxFun
ixfun) ->
(PrimType, Shape, VName, IxFun)
-> m (PrimType, Shape, VName, IxFun)
forall (f :: * -> *) a. Applicative f => a -> f a
pure (PrimType
et, [SubExp] -> Shape
forall d. [d] -> ShapeBase d
Shape ([SubExp] -> Shape) -> [SubExp] -> Shape
forall a b. (a -> b) -> a -> b
$ Shape -> [SubExp]
forall d. ShapeBase d -> [d]
shapeDims Shape
shape, VName
mem, IxFun
ixfun)
LetDecMem
_ ->
String -> m (PrimType, Shape, VName, IxFun)
forall a. HasCallStack => String -> a
error (String -> m (PrimType, Shape, VName, IxFun))
-> String -> m (PrimType, Shape, VName, IxFun)
forall a b. (a -> b) -> a -> b
$ String
"arrayVarReturns: " String -> ShowS
forall a. [a] -> [a] -> [a]
++ VName -> String
forall a. Pretty a => a -> String
pretty VName
v String -> ShowS
forall a. [a] -> [a] -> [a]
++ String
" is not an array."
varReturns ::
(HasScope rep m, Monad m, Mem rep inner) =>
VName ->
m ExpReturns
varReturns :: VName -> m ExpReturns
varReturns VName
v = do
LetDecMem
summary <- VName -> m LetDecMem
forall rep (m :: * -> *) inner.
(HasScope rep m, Mem rep inner) =>
VName -> m LetDecMem
lookupMemInfo VName
v
case LetDecMem
summary of
MemPrim PrimType
bt ->
ExpReturns -> m ExpReturns
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 u ret. PrimType -> MemInfo d u ret
MemPrim PrimType
bt
MemArray PrimType
et Shape
shape NoUniqueness
_ (ArrayIn VName
mem IxFun
ixfun) ->
ExpReturns -> m ExpReturns
forall (f :: * -> *) a. Applicative f => a -> f a
pure (ExpReturns -> m ExpReturns) -> ExpReturns -> m ExpReturns
forall a b. (a -> b) -> a -> b
$
PrimType
-> ShapeBase ExtSize
-> NoUniqueness
-> Maybe MemReturn
-> ExpReturns
forall d u ret.
PrimType -> ShapeBase d -> u -> ret -> MemInfo d u ret
MemArray PrimType
et ((SubExp -> ExtSize) -> Shape -> ShapeBase ExtSize
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap SubExp -> ExtSize
forall a. a -> Ext a
Free Shape
shape) NoUniqueness
NoUniqueness (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 -> ExtIxFun -> MemReturn
ReturnsInBlock VName
mem (ExtIxFun -> MemReturn) -> ExtIxFun -> MemReturn
forall a b. (a -> b) -> a -> b
$
[VName] -> IxFun -> ExtIxFun
existentialiseIxFun [] IxFun
ixfun
MemMem Space
space ->
ExpReturns -> m ExpReturns
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 u ret. Space -> MemInfo d u ret
MemMem Space
space
MemAcc VName
acc Shape
ispace [Type]
ts NoUniqueness
u ->
ExpReturns -> m ExpReturns
forall (f :: * -> *) a. Applicative f => a -> f a
pure (ExpReturns -> m ExpReturns) -> ExpReturns -> m ExpReturns
forall a b. (a -> b) -> a -> b
$ VName -> Shape -> [Type] -> NoUniqueness -> ExpReturns
forall d u ret. VName -> Shape -> [Type] -> u -> MemInfo d u ret
MemAcc VName
acc Shape
ispace [Type]
ts NoUniqueness
u
subExpReturns :: (HasScope rep m, Monad m, Mem rep inner) => SubExp -> m ExpReturns
subExpReturns :: 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 (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 u ret. PrimType -> MemInfo d u ret
MemPrim (PrimType -> ExpReturns) -> PrimType -> ExpReturns
forall a b. (a -> b) -> a -> b
$ PrimValue -> PrimType
primValueType PrimValue
v
expReturns ::
( Monad m,
LocalScope rep m,
Mem rep inner
) =>
Exp rep ->
m [ExpReturns]
expReturns :: Exp rep -> m [ExpReturns]
expReturns (BasicOp (SubExp SubExp
se)) =
ExpReturns -> [ExpReturns]
forall (f :: * -> *) a. Applicative f => a -> f a
pure (ExpReturns -> [ExpReturns]) -> m ExpReturns -> m [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 -> [ExpReturns]
forall (f :: * -> *) a. Applicative f => a -> f a
pure (ExpReturns -> [ExpReturns]) -> m ExpReturns -> m [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 ShapeChange SubExp
newshape VName
v)) = do
(PrimType
et, Shape
_, VName
mem, IxFun
ixfun) <- VName -> m (PrimType, Shape, VName, IxFun)
forall rep (m :: * -> *) inner.
(HasScope rep m, Monad m, Mem rep inner) =>
VName -> m (PrimType, Shape, VName, IxFun)
arrayVarReturns VName
v
[ExpReturns] -> m [ExpReturns]
forall (f :: * -> *) a. Applicative f => a -> f a
pure
[ PrimType
-> ShapeBase ExtSize
-> NoUniqueness
-> Maybe MemReturn
-> ExpReturns
forall d u ret.
PrimType -> ShapeBase d -> u -> ret -> MemInfo d u ret
MemArray PrimType
et ([ExtSize] -> ShapeBase ExtSize
forall d. [d] -> ShapeBase d
Shape ([ExtSize] -> ShapeBase ExtSize) -> [ExtSize] -> ShapeBase ExtSize
forall a b. (a -> b) -> a -> b
$ (DimChange SubExp -> ExtSize) -> ShapeChange SubExp -> [ExtSize]
forall a b. (a -> b) -> [a] -> [b]
map (SubExp -> ExtSize
forall a. a -> Ext a
Free (SubExp -> ExtSize)
-> (DimChange SubExp -> SubExp) -> DimChange SubExp -> ExtSize
forall k (cat :: k -> k -> *) (b :: k) (c :: k) (a :: k).
Category cat =>
cat b c -> cat a b -> cat a c
. DimChange SubExp -> SubExp
forall d. DimChange d -> d
newDim) ShapeChange SubExp
newshape) NoUniqueness
NoUniqueness (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 -> ExtIxFun -> MemReturn
ReturnsInBlock VName
mem (ExtIxFun -> MemReturn) -> ExtIxFun -> MemReturn
forall a b. (a -> b) -> a -> b
$
[VName] -> IxFun -> ExtIxFun
existentialiseIxFun [] (IxFun -> ExtIxFun) -> IxFun -> ExtIxFun
forall a b. (a -> b) -> a -> b
$
IxFun -> ShapeChange (TPrimExp Int64 VName) -> IxFun
forall num.
(Eq num, IntegralExp num) =>
IxFun num -> ShapeChange num -> IxFun num
IxFun.reshape IxFun
ixfun (ShapeChange (TPrimExp Int64 VName) -> IxFun)
-> ShapeChange (TPrimExp Int64 VName) -> IxFun
forall a b. (a -> b) -> a -> b
$
(DimChange SubExp -> DimChange (TPrimExp Int64 VName))
-> ShapeChange SubExp -> ShapeChange (TPrimExp Int64 VName)
forall a b. (a -> b) -> [a] -> [b]
map ((SubExp -> TPrimExp Int64 VName)
-> DimChange SubExp -> DimChange (TPrimExp Int64 VName)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap SubExp -> TPrimExp Int64 VName
pe64) ShapeChange SubExp
newshape
]
expReturns (BasicOp (Rearrange [Int]
perm VName
v)) = do
(PrimType
et, Shape [SubExp]
dims, VName
mem, IxFun
ixfun) <- VName -> m (PrimType, Shape, VName, IxFun)
forall rep (m :: * -> *) inner.
(HasScope rep m, Monad m, Mem rep inner) =>
VName -> m (PrimType, Shape, VName, IxFun)
arrayVarReturns VName
v
let ixfun' :: IxFun
ixfun' = IxFun -> [Int] -> IxFun
forall num. IntegralExp num => IxFun num -> [Int] -> IxFun num
IxFun.permute IxFun
ixfun [Int]
perm
dims' :: [SubExp]
dims' = [Int] -> [SubExp] -> [SubExp]
forall a. [Int] -> [a] -> [a]
rearrangeShape [Int]
perm [SubExp]
dims
[ExpReturns] -> m [ExpReturns]
forall (f :: * -> *) a. Applicative f => a -> f a
pure
[ PrimType
-> ShapeBase ExtSize
-> NoUniqueness
-> Maybe MemReturn
-> ExpReturns
forall d u ret.
PrimType -> ShapeBase d -> u -> ret -> MemInfo d u ret
MemArray PrimType
et ([ExtSize] -> ShapeBase ExtSize
forall d. [d] -> ShapeBase d
Shape ([ExtSize] -> ShapeBase ExtSize) -> [ExtSize] -> ShapeBase ExtSize
forall a b. (a -> b) -> a -> b
$ (SubExp -> ExtSize) -> [SubExp] -> [ExtSize]
forall a b. (a -> b) -> [a] -> [b]
map SubExp -> ExtSize
forall a. a -> Ext a
Free [SubExp]
dims') NoUniqueness
NoUniqueness (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 -> ExtIxFun -> MemReturn
ReturnsInBlock VName
mem (ExtIxFun -> MemReturn) -> ExtIxFun -> MemReturn
forall a b. (a -> b) -> a -> b
$
[VName] -> IxFun -> ExtIxFun
existentialiseIxFun [] IxFun
ixfun'
]
expReturns (BasicOp (Rotate [SubExp]
offsets VName
v)) = do
(PrimType
et, Shape [SubExp]
dims, VName
mem, IxFun
ixfun) <- VName -> m (PrimType, Shape, VName, IxFun)
forall rep (m :: * -> *) inner.
(HasScope rep m, Monad m, Mem rep inner) =>
VName -> m (PrimType, Shape, VName, IxFun)
arrayVarReturns VName
v
let offsets' :: [TPrimExp Int64 VName]
offsets' = (SubExp -> TPrimExp Int64 VName)
-> [SubExp] -> [TPrimExp Int64 VName]
forall a b. (a -> b) -> [a] -> [b]
map SubExp -> TPrimExp Int64 VName
pe64 [SubExp]
offsets
ixfun' :: IxFun
ixfun' = IxFun -> [TPrimExp Int64 VName] -> IxFun
forall num.
(Eq num, IntegralExp num) =>
IxFun num -> Indices num -> IxFun num
IxFun.rotate IxFun
ixfun [TPrimExp Int64 VName]
offsets'
[ExpReturns] -> m [ExpReturns]
forall (f :: * -> *) a. Applicative f => a -> f a
pure
[ PrimType
-> ShapeBase ExtSize
-> NoUniqueness
-> Maybe MemReturn
-> ExpReturns
forall d u ret.
PrimType -> ShapeBase d -> u -> ret -> MemInfo d u ret
MemArray PrimType
et ([ExtSize] -> ShapeBase ExtSize
forall d. [d] -> ShapeBase d
Shape ([ExtSize] -> ShapeBase ExtSize) -> [ExtSize] -> ShapeBase ExtSize
forall a b. (a -> b) -> a -> b
$ (SubExp -> ExtSize) -> [SubExp] -> [ExtSize]
forall a b. (a -> b) -> [a] -> [b]
map SubExp -> ExtSize
forall a. a -> Ext a
Free [SubExp]
dims) NoUniqueness
NoUniqueness (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 -> ExtIxFun -> MemReturn
ReturnsInBlock VName
mem (ExtIxFun -> MemReturn) -> ExtIxFun -> MemReturn
forall a b. (a -> b) -> a -> b
$
[VName] -> IxFun -> ExtIxFun
existentialiseIxFun [] IxFun
ixfun'
]
expReturns (BasicOp (Index VName
v Slice SubExp
slice)) = do
ExpReturns -> [ExpReturns]
forall (f :: * -> *) a. Applicative f => a -> f a
pure (ExpReturns -> [ExpReturns])
-> (LetDecMem -> ExpReturns) -> LetDecMem -> [ExpReturns]
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 -> [ExpReturns]) -> m LetDecMem -> m [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 -> [ExpReturns]
forall (f :: * -> *) a. Applicative f => a -> f a
pure (ExpReturns -> [ExpReturns]) -> m ExpReturns -> m [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)) = do
ExpReturns -> [ExpReturns]
forall (f :: * -> *) a. Applicative f => a -> f a
pure (ExpReturns -> [ExpReturns])
-> (LetDecMem -> ExpReturns) -> LetDecMem -> [ExpReturns]
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 -> [ExpReturns]) -> m LetDecMem -> m [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 -> [ExpReturns]
forall (f :: * -> *) a. Applicative f => a -> f a
pure (ExpReturns -> [ExpReturns]) -> m ExpReturns -> m [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) =
[ExtType] -> [ExpReturns]
extReturns ([ExtType] -> [ExpReturns])
-> ([Type] -> [ExtType]) -> [Type] -> [ExpReturns]
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 u. [TypeBase Shape u] -> [TypeBase (ShapeBase ExtSize) u]
staticShapes ([Type] -> [ExpReturns]) -> m [Type] -> m [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@(DoLoop [(FParam rep, SubExp)]
merge LoopForm rep
_ Body rep
_) = do
[ExtType]
t <- Exp rep -> m [ExtType]
forall rep (m :: * -> *).
(HasScope rep m, TypedOp (Op rep)) =>
Exp rep -> m [ExtType]
expExtType Exp rep
e
(ExtType -> Param FParamMem -> m ExpReturns)
-> [ExtType] -> [Param FParamMem] -> m [ExpReturns]
forall (m :: * -> *) a b c.
Applicative m =>
(a -> b -> m c) -> [a] -> [b] -> m [c]
zipWithM ExtType -> Param FParamMem -> m ExpReturns
typeWithDec [ExtType]
t ([Param FParamMem] -> m [ExpReturns])
-> [Param FParamMem] -> m [ExpReturns]
forall a b. (a -> b) -> a -> b
$ ((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
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 ExtSize
shape NoUniqueness
u,
MemArray PrimType
_ Shape
_ Uniqueness
_ (ArrayIn VName
mem IxFun
ixfun)
)
| Just (Int
i, Param FParamMem
mem_p) <- VName -> Maybe (Int, Param FParamMem)
isMergeVar VName
mem,
Mem Space
space <- Param FParamMem -> Type
forall dec. Typed dec => Param dec -> Type
paramType Param FParamMem
mem_p ->
ExpReturns -> m ExpReturns
forall (f :: * -> *) a. Applicative f => a -> f a
pure (ExpReturns -> m ExpReturns) -> ExpReturns -> m ExpReturns
forall a b. (a -> b) -> a -> b
$ PrimType
-> ShapeBase ExtSize
-> NoUniqueness
-> Maybe MemReturn
-> ExpReturns
forall d u ret.
PrimType -> ShapeBase d -> u -> ret -> MemInfo d u ret
MemArray PrimType
pt ShapeBase ExtSize
shape NoUniqueness
u (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 -> ExtIxFun -> MemReturn
ReturnsNewBlock Space
space Int
i ExtIxFun
ixfun'
| Bool
otherwise ->
ExpReturns -> m ExpReturns
forall (f :: * -> *) a. Applicative f => a -> f a
pure (ExpReturns -> m ExpReturns) -> ExpReturns -> m ExpReturns
forall a b. (a -> b) -> a -> b
$ PrimType
-> ShapeBase ExtSize
-> NoUniqueness
-> Maybe MemReturn
-> ExpReturns
forall d u ret.
PrimType -> ShapeBase d -> u -> ret -> MemInfo d u ret
MemArray PrimType
pt ShapeBase ExtSize
shape NoUniqueness
u (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 -> ExtIxFun -> MemReturn
ReturnsInBlock VName
mem ExtIxFun
ixfun'
where
ixfun' :: ExtIxFun
ixfun' = [VName] -> IxFun -> ExtIxFun
existentialiseIxFun ((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) IxFun
ixfun
(Array {}, FParamMem
_) ->
String -> m ExpReturns
forall a. HasCallStack => String -> a
error String
"expReturns: Array return type but not array merge variable."
(Acc VName
acc Shape
ispace [Type]
ts NoUniqueness
u, FParamMem
_) ->
ExpReturns -> m ExpReturns
forall (f :: * -> *) a. Applicative f => a -> f a
pure (ExpReturns -> m ExpReturns) -> ExpReturns -> m ExpReturns
forall a b. (a -> b) -> a -> b
$ VName -> Shape -> [Type] -> NoUniqueness -> ExpReturns
forall d u ret. VName -> Shape -> [Type] -> u -> MemInfo d u ret
MemAcc VName
acc Shape
ispace [Type]
ts NoUniqueness
u
(Prim PrimType
pt, FParamMem
_) ->
ExpReturns -> m ExpReturns
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 u ret. PrimType -> MemInfo d u ret
MemPrim PrimType
pt
(Mem Space
space, FParamMem
_) ->
ExpReturns -> m ExpReturns
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 u ret. Space -> MemInfo d u ret
MemMem Space
space
isMergeVar :: VName -> Maybe (Int, Param FParamMem)
isMergeVar 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 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 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]
ret (Safety, SrcLoc, [SrcLoc])
_) =
[ExpReturns] -> m [ExpReturns]
forall (f :: * -> *) a. Applicative f => a -> f a
pure ([ExpReturns] -> m [ExpReturns]) -> [ExpReturns] -> m [ExpReturns]
forall a b. (a -> b) -> a -> b
$ (FunReturns -> ExpReturns) -> [FunReturns] -> [ExpReturns]
forall a b. (a -> b) -> [a] -> [b]
map FunReturns -> ExpReturns
funReturnsToExpReturns [RetType rep]
[FunReturns]
ret
expReturns (If SubExp
_ Body rep
_ Body rep
_ (IfDec [BranchType rep]
ret IfSort
_)) =
[ExpReturns] -> m [ExpReturns]
forall (f :: * -> *) a. Applicative f => a -> f a
pure ([ExpReturns] -> m [ExpReturns]) -> [ExpReturns] -> m [ExpReturns]
forall a b. (a -> b) -> a -> b
$ (BodyReturns -> ExpReturns) -> [BodyReturns] -> [ExpReturns]
forall a b. (a -> b) -> [a] -> [b]
map BodyReturns -> ExpReturns
bodyReturnsToExpReturns [BranchType rep]
[BodyReturns]
ret
expReturns (Op Op rep
op) =
MemOp inner -> m [ExpReturns]
forall op rep inner (m :: * -> *).
(OpReturns op, Mem rep inner, Monad m, HasScope rep m) =>
op -> m [ExpReturns]
opReturns Op rep
MemOp inner
op
expReturns (WithAcc [WithAccInput rep]
inputs Lambda rep
lam) =
[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)
mapM WithAccInput rep -> m [ExpReturns]
forall rep (m :: * -> *) (t :: * -> *) inner a c.
(HasScope rep m, Monad m, Traversable t, HasLetDecMem (LetDec rep),
ASTRep rep, OpReturns inner, FParamInfo rep ~ FParamMem,
LParamInfo rep ~ LetDecMem, RetType rep ~ FunReturns,
BranchType rep ~ BodyReturns, Op rep ~ MemOp inner) =>
(a, t VName, c) -> m (t ExpReturns)
inputReturns [WithAccInput rep]
inputs)
m ([ExpReturns] -> [ExpReturns])
-> m [ExpReturns] -> m [ExpReturns]
forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*>
[ExpReturns] -> m [ExpReturns]
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 u. [TypeBase Shape u] -> [TypeBase (ShapeBase ExtSize) u]
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. Lambda rep -> [Type]
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)
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 (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 NoUniqueness MemBind)
sliceInfo :: VName -> Slice SubExp -> m LetDecMem
sliceInfo VName
v Slice SubExp
slice = do
(PrimType
et, Shape
_, VName
mem, IxFun
ixfun) <- VName -> m (PrimType, Shape, VName, IxFun)
forall rep (m :: * -> *) inner.
(HasScope rep m, Monad m, Mem rep inner) =>
VName -> m (PrimType, Shape, VName, IxFun)
arrayVarReturns VName
v
case Slice SubExp -> [SubExp]
forall d. Slice d -> [d]
sliceDims Slice SubExp
slice of
[] -> LetDecMem -> m LetDecMem
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 u ret. PrimType -> MemInfo d u ret
MemPrim PrimType
et
[SubExp]
dims ->
LetDecMem -> m LetDecMem
forall (f :: * -> *) a. Applicative f => a -> f a
pure (LetDecMem -> m LetDecMem) -> LetDecMem -> m LetDecMem
forall a b. (a -> b) -> a -> b
$
PrimType -> Shape -> NoUniqueness -> MemBind -> LetDecMem
forall d u ret.
PrimType -> ShapeBase d -> u -> ret -> MemInfo d u ret
MemArray PrimType
et ([SubExp] -> Shape
forall d. [d] -> ShapeBase d
Shape [SubExp]
dims) NoUniqueness
NoUniqueness (MemBind -> LetDecMem) -> (IxFun -> MemBind) -> IxFun -> LetDecMem
forall k (cat :: k -> k -> *) (b :: k) (c :: k) (a :: k).
Category cat =>
cat b c -> cat a b -> cat a c
. VName -> IxFun -> MemBind
ArrayIn VName
mem (IxFun -> LetDecMem) -> IxFun -> LetDecMem
forall a b. (a -> b) -> a -> b
$
IxFun -> Slice (TPrimExp Int64 VName) -> IxFun
forall num.
(Eq num, IntegralExp num) =>
IxFun num -> Slice num -> IxFun num
IxFun.slice IxFun
ixfun ((SubExp -> TPrimExp Int64 VName)
-> Slice SubExp -> Slice (TPrimExp Int64 VName)
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 NoUniqueness MemBind)
flatSliceInfo :: VName -> FlatSlice SubExp -> m LetDecMem
flatSliceInfo VName
v slice :: FlatSlice SubExp
slice@(FlatSlice SubExp
offset [FlatDimIndex SubExp]
idxs) = do
(PrimType
et, Shape
_, VName
mem, IxFun
ixfun) <- VName -> m (PrimType, Shape, VName, IxFun)
forall rep (m :: * -> *) inner.
(HasScope rep m, Monad m, Mem rep inner) =>
VName -> m (PrimType, Shape, VName, IxFun)
arrayVarReturns VName
v
(FlatDimIndex SubExp -> FlatDimIndex (TPrimExp Int64 VName))
-> [FlatDimIndex SubExp] -> [FlatDimIndex (TPrimExp Int64 VName)]
forall a b. (a -> b) -> [a] -> [b]
map ((SubExp -> TPrimExp Int64 VName)
-> FlatDimIndex SubExp -> FlatDimIndex (TPrimExp Int64 VName)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap SubExp -> TPrimExp Int64 VName
pe64) [FlatDimIndex SubExp]
idxs
[FlatDimIndex (TPrimExp Int64 VName)]
-> ([FlatDimIndex (TPrimExp Int64 VName)]
-> FlatSlice (TPrimExp Int64 VName))
-> FlatSlice (TPrimExp Int64 VName)
forall a b. a -> (a -> b) -> b
& TPrimExp Int64 VName
-> [FlatDimIndex (TPrimExp Int64 VName)]
-> FlatSlice (TPrimExp Int64 VName)
forall d. d -> [FlatDimIndex d] -> FlatSlice d
FlatSlice (SubExp -> TPrimExp Int64 VName
pe64 SubExp
offset)
FlatSlice (TPrimExp Int64 VName)
-> (FlatSlice (TPrimExp Int64 VName) -> IxFun) -> IxFun
forall a b. a -> (a -> b) -> b
& IxFun -> FlatSlice (TPrimExp Int64 VName) -> IxFun
forall num.
(Eq num, IntegralExp num) =>
IxFun num -> FlatSlice num -> IxFun num
IxFun.flatSlice IxFun
ixfun
IxFun -> (IxFun -> MemBind) -> MemBind
forall a b. a -> (a -> b) -> b
& VName -> IxFun -> MemBind
ArrayIn VName
mem
MemBind -> (MemBind -> LetDecMem) -> LetDecMem
forall a b. a -> (a -> b) -> b
& PrimType -> Shape -> NoUniqueness -> MemBind -> LetDecMem
forall d u ret.
PrimType -> ShapeBase d -> u -> ret -> MemInfo d u ret
MemArray PrimType
et ([SubExp] -> Shape
forall d. [d] -> ShapeBase d
Shape (FlatSlice SubExp -> [SubExp]
forall d. FlatSlice d -> [d]
flatSliceDims FlatSlice SubExp
slice)) NoUniqueness
NoUniqueness
LetDecMem -> (LetDecMem -> m LetDecMem) -> m LetDecMem
forall a b. a -> (a -> b) -> b
& LetDecMem -> m LetDecMem
forall (f :: * -> *) a. Applicative f => a -> f a
pure
class TypedOp op => OpReturns op where
opReturns :: (Mem rep inner, Monad m, HasScope rep m) => op -> m [ExpReturns]
opReturns op
op = [ExtType] -> [ExpReturns]
extReturns ([ExtType] -> [ExpReturns]) -> m [ExtType] -> m [ExpReturns]
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> op -> m [ExtType]
forall op t (m :: * -> *).
(TypedOp op, HasScope t m) =>
op -> m [ExtType]
opType op
op
instance OpReturns inner => OpReturns (MemOp inner) where
opReturns :: MemOp inner -> m [ExpReturns]
opReturns (Alloc SubExp
_ Space
space) = [ExpReturns] -> m [ExpReturns]
forall (f :: * -> *) a. Applicative f => a -> f a
pure [Space -> ExpReturns
forall d u ret. Space -> MemInfo d u ret
MemMem Space
space]
opReturns (Inner inner
op) = inner -> m [ExpReturns]
forall op rep inner (m :: * -> *).
(OpReturns op, Mem rep inner, Monad m, HasScope rep m) =>
op -> m [ExpReturns]
opReturns inner
op
instance OpReturns () where
opReturns :: () -> m [ExpReturns]
opReturns () = [ExpReturns] -> m [ExpReturns]
forall (f :: * -> *) a. Applicative f => a -> f a
pure []
applyFunReturns ::
Typed dec =>
[FunReturns] ->
[Param dec] ->
[(SubExp, Type)] ->
Maybe [FunReturns]
applyFunReturns :: [FunReturns]
-> [Param dec] -> [(SubExp, Type)] -> Maybe [FunReturns]
applyFunReturns [FunReturns]
rets [Param dec]
params [(SubExp, Type)]
args
| Just [DeclExtType]
_ <- [DeclExtType]
-> [Param dec] -> [(SubExp, Type)] -> Maybe [DeclExtType]
forall rt dec.
(IsRetType rt, Typed dec) =>
[rt] -> [Param dec] -> [(SubExp, Type)] -> Maybe [rt]
applyRetType [DeclExtType]
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 u.
MemInfo ExtSize u MemReturn -> MemInfo ExtSize u MemReturn
correctDims [FunReturns]
rets
| Bool
otherwise =
Maybe [FunReturns]
forall a. Maybe a
Nothing
where
rettype :: [DeclExtType]
rettype = (FunReturns -> DeclExtType) -> [FunReturns] -> [DeclExtType]
forall a b. (a -> b) -> [a] -> [b]
map FunReturns -> DeclExtType
forall t. DeclExtTyped t => t -> DeclExtType
declExtTypeOf [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 ExtSize u MemReturn -> MemInfo ExtSize u MemReturn
correctDims (MemPrim PrimType
t) =
PrimType -> MemInfo ExtSize u MemReturn
forall d u ret. PrimType -> MemInfo d u ret
MemPrim PrimType
t
correctDims (MemMem Space
space) =
Space -> MemInfo ExtSize u MemReturn
forall d u ret. Space -> MemInfo d u ret
MemMem Space
space
correctDims (MemArray PrimType
et ShapeBase ExtSize
shape u
u MemReturn
memsummary) =
PrimType
-> ShapeBase ExtSize
-> u
-> MemReturn
-> MemInfo ExtSize u MemReturn
forall d u ret.
PrimType -> ShapeBase d -> u -> ret -> MemInfo d u ret
MemArray PrimType
et (ShapeBase ExtSize -> ShapeBase ExtSize
correctShape ShapeBase ExtSize
shape) u
u (MemReturn -> MemInfo ExtSize u MemReturn)
-> MemReturn -> MemInfo ExtSize u MemReturn
forall a b. (a -> b) -> a -> b
$
MemReturn -> MemReturn
correctSummary MemReturn
memsummary
correctDims (MemAcc VName
acc Shape
ispace [Type]
ts u
u) =
VName -> Shape -> [Type] -> u -> MemInfo ExtSize u MemReturn
forall d u ret. VName -> Shape -> [Type] -> u -> MemInfo d u ret
MemAcc VName
acc Shape
ispace [Type]
ts u
u
correctShape :: ShapeBase ExtSize -> ShapeBase ExtSize
correctShape = [ExtSize] -> ShapeBase ExtSize
forall d. [d] -> ShapeBase d
Shape ([ExtSize] -> ShapeBase ExtSize)
-> (ShapeBase ExtSize -> [ExtSize])
-> ShapeBase ExtSize
-> ShapeBase ExtSize
forall k (cat :: k -> k -> *) (b :: k) (c :: k) (a :: k).
Category cat =>
cat b c -> cat a b -> cat a c
. (ExtSize -> ExtSize) -> [ExtSize] -> [ExtSize]
forall a b. (a -> b) -> [a] -> [b]
map ExtSize -> ExtSize
correctDim ([ExtSize] -> [ExtSize])
-> (ShapeBase ExtSize -> [ExtSize])
-> ShapeBase ExtSize
-> [ExtSize]
forall k (cat :: k -> k -> *) (b :: k) (c :: k) (a :: k).
Category cat =>
cat b c -> cat a b -> cat a c
. ShapeBase ExtSize -> [ExtSize]
forall d. ShapeBase d -> [d]
shapeDims
correctDim :: ExtSize -> ExtSize
correctDim (Ext Int
i) = Int -> ExtSize
forall a. Int -> Ext a
Ext Int
i
correctDim (Free SubExp
se) = SubExp -> ExtSize
forall a. a -> Ext a
Free (SubExp -> ExtSize) -> SubExp -> ExtSize
forall a b. (a -> b) -> a -> b
$ SubExp -> SubExp
substSubExp SubExp
se
correctSummary :: MemReturn -> MemReturn
correctSummary (ReturnsNewBlock Space
space Int
i ExtIxFun
ixfun) =
Space -> Int -> ExtIxFun -> MemReturn
ReturnsNewBlock Space
space Int
i ExtIxFun
ixfun
correctSummary (ReturnsInBlock VName
mem ExtIxFun
ixfun) =
VName -> ExtIxFun -> MemReturn
ReturnsInBlock VName
mem' ExtIxFun
ixfun
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