futhark-0.26.3: src/Futhark/IR/Mem.hs
{-# LANGUAGE TypeFamilies #-}
{-# LANGUAGE UndecidableInstances #-}
-- | Building blocks for defining representations where every array
-- is given information about which memory block is it based in, and
-- how array elements map to memory block offsets.
--
-- There are two primary concepts you will need to understand:
--
-- 1. Memory blocks, which are Futhark values of type v'Mem'
-- (parametrized with their size). These correspond to arbitrary
-- blocks of memory, and are created using the 'Alloc' operation.
--
-- 2. Index functions, which describe a mapping from the index space
-- of an array (eg. a two-dimensional space for an array of type
-- @[[int]]@) to a one-dimensional offset into a memory block.
-- Thus, index functions describe how arbitrary-dimensional arrays
-- are mapped to the single-dimensional world of memory.
--
-- At a conceptual level, imagine that we have a two-dimensional array
-- @a@ of 32-bit integers, consisting of @n@ rows of @m@ elements
-- each. This array could be represented in classic row-major format
-- with an index function like the following:
--
-- @
-- f(i,j) = i * m + j
-- @
--
-- When we want to know the location of element @a[2,3]@, we simply
-- call the index function as @f(2,3)@ and obtain @2*m+3@. We could
-- also have chosen another index function, one that represents the
-- array in column-major (or "transposed") format:
--
-- @
-- f(i,j) = j * n + i
-- @
--
-- Index functions are not Futhark-level functions, but a special
-- construct that the final code generator will eventually use to
-- generate concrete access code. By modifying the index functions we
-- can change how an array is represented in memory, which can permit
-- memory access pattern optimisations.
--
-- Every time we bind an array, whether in a @let@-binding, @loop@
-- merge parameter, or @lambda@ parameter, we have an annotation
-- specifying a memory block and an index function. In some cases,
-- such as @let@-bindings for many expressions, we are free to specify
-- an arbitrary index function and memory block - for example, we get
-- to decide where 'Copy' stores its result - but in other cases the
-- type rules of the expression chooses for us. For example, 'Index'
-- always produces an array in the same memory block as its input, and
-- with the same index function, except with some indices fixed.
module Futhark.IR.Mem
( LetDecMem,
FParamMem,
LParamMem,
RetTypeMem,
BranchTypeMem,
MemOp (..),
traverseMemOpStms,
MemInfo (..),
MemBound,
MemBind (..),
MemReturn (..),
LMAD,
ExtLMAD,
isStaticLMAD,
ExpReturns,
BodyReturns,
FunReturns,
noUniquenessReturns,
bodyReturnsToExpReturns,
Mem,
HasLetDecMem (..),
OpReturns (..),
varReturns,
expReturns,
extReturns,
nameInfoToMemInfo,
lookupMemInfo,
subExpMemInfo,
lookupArraySummary,
lookupMemSpace,
existentialiseLMAD,
-- * Type checking parts
matchBranchReturnType,
matchPatToExp,
matchFunctionReturnType,
matchLoopResultMem,
bodyReturnsFromPat,
checkMemInfo,
-- * Module re-exports
module Futhark.IR.Prop,
module Futhark.IR.Traversals,
module Futhark.IR.Pretty,
module Futhark.IR.Syntax,
module Futhark.Analysis.PrimExp.Convert,
)
where
import Control.Category
import Control.Monad
import Control.Monad.Except
import Control.Monad.Reader
import Control.Monad.State
import Data.Foldable (traverse_)
import Data.Function ((&))
import Data.Kind qualified
import Data.List (elemIndex, find)
import Data.Map.Strict qualified as M
import Data.Maybe
import Data.Text qualified as T
import Futhark.Analysis.Metrics
import Futhark.Analysis.PrimExp.Convert
import Futhark.Analysis.PrimExp.Simplify
import Futhark.Analysis.SymbolTable qualified as ST
import Futhark.IR.Aliases
( Aliases,
CanBeAliased (..),
removeExpAliases,
removePatAliases,
removeScopeAliases,
)
import Futhark.IR.Mem.LMAD qualified as LMAD
import Futhark.IR.Pretty
import Futhark.IR.Prop
import Futhark.IR.Prop.Aliases
import Futhark.IR.Syntax
import Futhark.IR.Traversals
import Futhark.IR.TypeCheck qualified as TC
import Futhark.Optimise.Simplify.Engine qualified as Engine
import Futhark.Optimise.Simplify.Rep
import Futhark.Transform.Rename
import Futhark.Transform.Substitute
import Futhark.Util
import Futhark.Util.Pretty (docText, indent, ppTupleLines', pretty, (<+>), (</>))
import Futhark.Util.Pretty qualified as PP
import Prelude hiding (id, (.))
type LetDecMem = MemInfo SubExp NoUniqueness MemBind
type FParamMem = MemInfo SubExp Uniqueness MemBind
type LParamMem = MemInfo SubExp NoUniqueness MemBind
type RetTypeMem = FunReturns
type BranchTypeMem = BodyReturns
-- | The class of pattern element decorators that contain memory
-- information.
class HasLetDecMem t where
letDecMem :: t -> LetDecMem
instance HasLetDecMem LetDecMem where
letDecMem = id
instance (HasLetDecMem b) => HasLetDecMem (a, b) where
letDecMem = letDecMem . snd
type Mem rep inner =
( FParamInfo rep ~ FParamMem,
LParamInfo rep ~ LParamMem,
HasLetDecMem (LetDec rep),
RetType rep ~ RetTypeMem,
BranchType rep ~ BranchTypeMem,
ASTRep rep,
OpReturns inner,
RephraseOp inner,
ASTConstraints (inner rep),
FreeIn (inner rep),
OpC rep ~ MemOp inner
)
instance IsRetType FunReturns where
primRetType = MemPrim
applyRetType = applyFunReturns
instance IsBodyType BodyReturns where
primBodyType = MemPrim
data MemOp (inner :: Data.Kind.Type -> Data.Kind.Type) (rep :: Data.Kind.Type)
= -- | Allocate a memory block.
Alloc SubExp Space
| -- | If the given array has a direct index function (as 'LMAD.isDirect'),
-- returns the array unchanged. Otherwise, make a copy of the array. The
-- choice is made dynamically, and the goal is to handle cases where we must
-- return something that is direct (mostly entry point return types), but we
-- cannot statically be sure of the layout of some array.
EnsureDirect VName
| Inner (inner rep)
deriving (Eq, Ord, Show)
-- | A helper for defining 'TraverseOpStms'.
traverseMemOpStms ::
(Monad m) =>
OpStmsTraverser m (inner rep) rep ->
OpStmsTraverser m (MemOp inner rep) rep
traverseMemOpStms _ _ op@Alloc {} = pure op
traverseMemOpStms _ _ op@EnsureDirect {} = pure op
traverseMemOpStms onInner f (Inner inner) = Inner <$> onInner f inner
instance (RephraseOp inner) => RephraseOp (MemOp inner) where
rephraseInOp _ (Alloc e space) = pure (Alloc e space)
rephraseInOp _ (EnsureDirect v) = pure (EnsureDirect v)
rephraseInOp r (Inner x) = Inner <$> rephraseInOp r x
instance (FreeIn (inner rep)) => FreeIn (MemOp inner rep) where
freeIn' (Alloc size _) = freeIn' size
freeIn' (EnsureDirect v) = freeIn' v
freeIn' (Inner k) = freeIn' k
instance (TypedOp inner) => TypedOp (MemOp inner) where
opType (Alloc _ space) = pure [Mem space]
opType (EnsureDirect v) = f <$> lookupType v
where
f (Array pt shape _) =
[Mem DefaultSpace, Array pt (fmap Free shape) NoUniqueness]
f _ = error $ "EnsureDirect applied to non-array: " ++ show v
opType (Inner k) = opType k
instance (AliasedOp inner) => AliasedOp (MemOp inner) where
opAliases Alloc {} = [mempty]
-- XXX - why not mem aliases for EnsureDirect?
opAliases (EnsureDirect v) = [mempty, oneName v]
opAliases (Inner k) = opAliases k
consumedInOp Alloc {} = mempty
consumedInOp EnsureDirect {} = mempty
consumedInOp (Inner k) = consumedInOp k
instance (CanBeAliased inner) => CanBeAliased (MemOp inner) where
addOpAliases _ (Alloc se space) = Alloc se space
addOpAliases _ (EnsureDirect v) = EnsureDirect v
addOpAliases aliases (Inner k) = Inner $ addOpAliases aliases k
instance (Rename (inner rep)) => Rename (MemOp inner rep) where
rename (Alloc size space) = Alloc <$> rename size <*> pure space
rename (EnsureDirect v) = EnsureDirect <$> rename v
rename (Inner k) = Inner <$> rename k
instance (Substitute (inner rep)) => Substitute (MemOp inner rep) where
substituteNames subst (Alloc size space) = Alloc (substituteNames subst size) space
substituteNames subst (EnsureDirect v) = EnsureDirect (substituteNames subst v)
substituteNames subst (Inner k) = Inner $ substituteNames subst k
instance (PP.Pretty (inner rep)) => PP.Pretty (MemOp inner rep) where
pretty (Alloc e DefaultSpace) = "alloc" <> PP.apply [PP.pretty e]
pretty (Alloc e s) = "alloc" <> PP.apply [PP.pretty e, PP.pretty s]
pretty (EnsureDirect v) = "ensure_direct" <> PP.apply [PP.pretty v]
pretty (Inner k) = PP.pretty k
instance (OpMetrics (inner rep)) => OpMetrics (MemOp inner rep) where
opMetrics Alloc {} = seen "Alloc"
opMetrics EnsureDirect {} = seen "EnsureDirect"
opMetrics (Inner k) = opMetrics k
instance (IsOp inner) => IsOp (MemOp inner) where
safeOp (Alloc (Constant (IntValue (Int64Value k))) _) = k >= 0
safeOp Alloc {} = False
safeOp EnsureDirect {} = True
safeOp (Inner k) = safeOp k
cheapOp (Inner k) = cheapOp k
cheapOp Alloc {} = True
cheapOp EnsureDirect {} = False
opDependencies (Alloc _ e) = [freeIn e]
opDependencies (EnsureDirect v) = [mempty, freeIn v]
opDependencies (Inner op) = opDependencies op
instance (CanBeWise inner) => CanBeWise (MemOp inner) where
addOpWisdom (Alloc size space) = Alloc size space
addOpWisdom (EnsureDirect v) = EnsureDirect v
addOpWisdom (Inner k) = Inner $ addOpWisdom k
instance (ST.IndexOp (inner rep)) => ST.IndexOp (MemOp inner rep) where
indexOp vtable k (Inner op) is = ST.indexOp vtable k op is
indexOp _ _ _ _ = Nothing
-- | The LMAD representation used for memory annotations.
type LMAD = LMAD.LMAD (TPrimExp Int64 VName)
-- | An index function that may contain existential variables.
type ExtLMAD = LMAD.LMAD (TPrimExp Int64 (Ext VName))
-- | A summary of the memory information for every let-bound
-- identifier, function parameter, and return value. Parameterisered
-- over uniqueness, dimension, and auxiliary array information.
data MemInfo d u ret
= -- | A primitive value.
MemPrim PrimType
| -- | A memory block.
MemMem Space
| -- | The array is stored in the named memory block, and with the
-- given index function. The index function maps indices in the
-- array to /element/ offset, /not/ byte offsets! To translate to
-- byte offsets, multiply the offset with the size of the array
-- element type.
MemArray PrimType (ShapeBase d) u ret
| -- | An accumulator, which is not stored anywhere.
MemAcc VName Shape [Type] u
deriving (Eq, Show, Ord) --- XXX Ord?
type MemBound u = MemInfo SubExp u MemBind
instance (FixExt ret) => DeclExtTyped (MemInfo ExtSize Uniqueness ret) where
declExtTypeOf (MemPrim pt) = Prim pt
declExtTypeOf (MemMem space) = Mem space
declExtTypeOf (MemArray pt shape u _) = Array pt shape u
declExtTypeOf (MemAcc acc ispace ts u) = Acc acc ispace ts u
instance (FixExt ret) => ExtTyped (MemInfo ExtSize Uniqueness ret) where
extTypeOf = fromDecl . declExtTypeOf
instance (FixExt ret) => ExtTyped (MemInfo ExtSize NoUniqueness ret) where
extTypeOf (MemPrim pt) = Prim pt
extTypeOf (MemMem space) = Mem space
extTypeOf (MemArray pt shape u _) = Array pt shape u
extTypeOf (MemAcc acc ispace ts u) = Acc acc ispace ts u
instance (FixExt ret) => FixExt (MemInfo ExtSize u ret) where
fixExt _ _ (MemPrim pt) = MemPrim pt
fixExt _ _ (MemMem space) = MemMem space
fixExt _ _ (MemAcc acc ispace ts u) = MemAcc acc ispace ts u
fixExt i se (MemArray pt shape u ret) =
MemArray pt (fixExt i se shape) u (fixExt i se ret)
mapExt _ (MemPrim pt) = MemPrim pt
mapExt _ (MemMem space) = MemMem space
mapExt _ (MemAcc acc ispace ts u) = MemAcc acc ispace ts u
mapExt f (MemArray pt shape u ret) =
MemArray pt (mapExt f shape) u (mapExt f ret)
instance Typed (MemInfo SubExp Uniqueness ret) where
typeOf = fromDecl . declTypeOf
instance Typed (MemInfo SubExp NoUniqueness ret) where
typeOf (MemPrim pt) = Prim pt
typeOf (MemMem space) = Mem space
typeOf (MemArray bt shape u _) = Array bt shape u
typeOf (MemAcc acc ispace ts u) = Acc acc ispace ts u
instance DeclTyped (MemInfo SubExp Uniqueness ret) where
declTypeOf (MemPrim bt) = Prim bt
declTypeOf (MemMem space) = Mem space
declTypeOf (MemArray bt shape u _) = Array bt shape u
declTypeOf (MemAcc acc ispace ts u) = Acc acc ispace ts u
instance (FreeIn d, FreeIn ret) => FreeIn (MemInfo d u ret) where
freeIn' (MemArray _ shape _ ret) = freeIn' shape <> freeIn' ret
freeIn' (MemMem s) = freeIn' s
freeIn' MemPrim {} = mempty
freeIn' (MemAcc acc ispace ts _) = freeIn' (acc, ispace, ts)
instance (Substitute d, Substitute ret) => Substitute (MemInfo d u ret) where
substituteNames subst (MemArray bt shape u ret) =
MemArray
bt
(substituteNames subst shape)
u
(substituteNames subst ret)
substituteNames substs (MemAcc acc ispace ts u) =
MemAcc
(substituteNames substs acc)
(substituteNames substs ispace)
(substituteNames substs ts)
u
substituteNames _ (MemMem space) =
MemMem space
substituteNames _ (MemPrim bt) =
MemPrim bt
instance (Substitute d, Substitute ret) => Rename (MemInfo d u ret) where
rename = substituteRename
simplifyLMAD ::
(Engine.SimplifiableRep rep) =>
LMAD ->
Engine.SimpleM rep LMAD
simplifyLMAD = traverse $ fmap isInt64 . simplifyPrimExp . untyped
simplifyExtLMAD ::
(Engine.SimplifiableRep rep) =>
ExtLMAD ->
Engine.SimpleM rep ExtLMAD
simplifyExtLMAD = traverse $ fmap isInt64 . simplifyExtPrimExp . untyped
isStaticLMAD :: ExtLMAD -> Maybe LMAD
isStaticLMAD = traverse $ traverse inst
where
inst Ext {} = Nothing
inst (Free x) = Just x
instance
(Engine.Simplifiable d, Engine.Simplifiable ret) =>
Engine.Simplifiable (MemInfo d u ret)
where
simplify (MemPrim bt) =
pure $ MemPrim bt
simplify (MemMem space) =
pure $ MemMem space
simplify (MemArray bt shape u ret) =
MemArray bt <$> Engine.simplify shape <*> pure u <*> Engine.simplify ret
simplify (MemAcc acc ispace ts u) =
MemAcc <$> Engine.simplify acc <*> Engine.simplify ispace <*> Engine.simplify ts <*> pure 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
pretty (MemPrim bt) = PP.pretty bt
pretty (MemMem DefaultSpace) = "mem"
pretty (MemMem s) = "mem" <> PP.pretty s
pretty (MemArray bt shape u ret) =
PP.pretty (Array bt shape u) <+> "@" <+> PP.pretty ret
pretty (MemAcc acc ispace ts u) =
PP.pretty u <> PP.pretty (Acc acc ispace ts NoUniqueness :: Type)
-- | Memory information for an array bound somewhere in the program.
data MemBind
= -- | Located in this memory block with this index
-- function.
ArrayIn VName LMAD
deriving (Show)
instance Eq MemBind where
_ == _ = True
instance Ord MemBind where
_ `compare` _ = EQ
instance Rename MemBind where
rename = substituteRename
instance Substitute MemBind where
substituteNames substs (ArrayIn ident lmad) =
ArrayIn (substituteNames substs ident) (substituteNames substs lmad)
instance PP.Pretty MemBind where
pretty (ArrayIn mem lmad) =
PP.pretty mem <+> "->" PP.</> PP.pretty lmad
instance FreeIn MemBind where
freeIn' (ArrayIn mem lmad) = freeIn' mem <> freeIn' lmad
-- | A description of the memory properties of an array being returned
-- by an operation.
data MemReturn
= -- | The array is located in a memory block that is
-- already in scope.
ReturnsInBlock VName ExtLMAD
| -- | The operation returns a new (existential) memory
-- block.
ReturnsNewBlock Space Int ExtLMAD
deriving (Show)
instance Eq MemReturn where
_ == _ = True
instance Ord MemReturn where
_ `compare` _ = EQ
instance Rename MemReturn where
rename = substituteRename
instance Substitute MemReturn where
substituteNames substs (ReturnsInBlock ident lmad) =
ReturnsInBlock (substituteNames substs ident) (substituteNames substs lmad)
substituteNames substs (ReturnsNewBlock space i lmad) =
ReturnsNewBlock space i (substituteNames substs lmad)
instance FixExt MemReturn where
fixExt i (Var v) (ReturnsNewBlock _ j lmad)
| j == i =
ReturnsInBlock v $
fixExtLMAD
i
(primExpFromSubExp int64 (Var v))
lmad
fixExt i se (ReturnsNewBlock space j lmad) =
ReturnsNewBlock
space
j'
(fixExtLMAD i (primExpFromSubExp int64 se) lmad)
where
j'
| i < j = j - 1
| otherwise = j
fixExt i se (ReturnsInBlock mem lmad) =
ReturnsInBlock mem (fixExtLMAD i (primExpFromSubExp int64 se) lmad)
mapExt f (ReturnsNewBlock space i lmad) =
ReturnsNewBlock space (f i) lmad
mapExt f (ReturnsInBlock mem lmad) =
ReturnsInBlock mem (fmap (fmap f') lmad)
where
f' (Ext i) = Ext $ f i
f' v = v
fixExtLMAD :: Int -> PrimExp VName -> ExtLMAD -> ExtLMAD
fixExtLMAD i e = fmap $ isInt64 . replaceInPrimExp update . untyped
where
update (Ext j) t
| j > i = LeafExp (Ext $ j - 1) t
| j == i = fmap Free e
| otherwise = LeafExp (Ext j) t
update (Free x) t = LeafExp (Free x) t
leafExp :: Int -> TPrimExp Int64 (Ext a)
leafExp i = isInt64 $ LeafExp (Ext i) int64
existentialiseLMAD :: [VName] -> LMAD -> ExtLMAD
existentialiseLMAD ctx = LMAD.substitute ctx' . fmap (fmap Free)
where
ctx' = M.map leafExp $ M.fromList $ zip (map Free ctx) [0 ..]
instance PP.Pretty MemReturn where
pretty (ReturnsInBlock v lmad) =
pretty v <+> "->" PP.</> PP.pretty lmad
pretty (ReturnsNewBlock space i lmad) =
"?" <> pretty i <> PP.pretty space <+> "->" PP.</> PP.pretty lmad
instance FreeIn MemReturn where
freeIn' (ReturnsInBlock v lmad) = freeIn' v <> freeIn' lmad
freeIn' (ReturnsNewBlock space _ lmad) = freeIn' space <> freeIn' lmad
instance Engine.Simplifiable MemReturn where
simplify (ReturnsNewBlock space i lmad) =
ReturnsNewBlock space i <$> simplifyExtLMAD lmad
simplify (ReturnsInBlock v lmad) =
ReturnsInBlock <$> Engine.simplify v <*> simplifyExtLMAD lmad
instance Engine.Simplifiable MemBind where
simplify (ArrayIn mem lmad) =
ArrayIn <$> Engine.simplify mem <*> simplifyLMAD lmad
instance Engine.Simplifiable [FunReturns] where
simplify = mapM Engine.simplify
-- | The memory return of an expression. An array is annotated with
-- @Maybe MemReturn@, which can be interpreted as the expression
-- either dictating exactly where the array is located when it is
-- returned (if 'Just'), or able to put it whereever the binding
-- prefers (if 'Nothing').
--
-- This is necessary to capture the difference between an expression
-- that is just an array-typed variable, in which the array being
-- "returned" is located where it already is, and a @copy@ expression,
-- whose entire purpose is to store an existing array in some
-- arbitrary location. This is a consequence of the design decision
-- never to have implicit memory copies.
type ExpReturns = MemInfo ExtSize NoUniqueness (Maybe MemReturn)
-- | The return of a body, which must always indicate where
-- returned arrays are located.
type BodyReturns = MemInfo ExtSize NoUniqueness MemReturn
-- | The memory return of a function, which must always indicate where
-- returned arrays are located.
type FunReturns = MemInfo ExtSize Uniqueness MemReturn
maybeReturns :: MemInfo d u r -> MemInfo d u (Maybe r)
maybeReturns (MemArray bt shape u ret) =
MemArray bt shape u $ Just ret
maybeReturns (MemPrim bt) =
MemPrim bt
maybeReturns (MemMem space) =
MemMem space
maybeReturns (MemAcc acc ispace ts u) =
MemAcc acc ispace ts u
noUniquenessReturns :: MemInfo d u r -> MemInfo d NoUniqueness r
noUniquenessReturns (MemArray bt shape _ r) =
MemArray bt shape NoUniqueness r
noUniquenessReturns (MemPrim bt) =
MemPrim bt
noUniquenessReturns (MemMem space) =
MemMem space
noUniquenessReturns (MemAcc acc ispace ts _) =
MemAcc acc ispace ts NoUniqueness
funReturnsToExpReturns :: FunReturns -> ExpReturns
funReturnsToExpReturns = noUniquenessReturns . maybeReturns
bodyReturnsToExpReturns :: BodyReturns -> ExpReturns
bodyReturnsToExpReturns = noUniquenessReturns . maybeReturns
varInfoToExpReturns :: MemInfo SubExp NoUniqueness MemBind -> ExpReturns
varInfoToExpReturns (MemArray et shape u (ArrayIn mem lmad)) =
MemArray et (fmap Free shape) u $
Just $
ReturnsInBlock mem $
existentialiseLMAD [] lmad
varInfoToExpReturns (MemPrim pt) = MemPrim pt
varInfoToExpReturns (MemAcc acc ispace ts u) = MemAcc acc ispace ts u
varInfoToExpReturns (MemMem space) = MemMem space
matchRetTypeToResult ::
(Mem rep inner, TC.Checkable rep) =>
[FunReturns] ->
Result ->
TC.TypeM rep ()
matchRetTypeToResult rettype result = do
scope <- askScope
result_ts <- runReaderT (mapM (subExpMemInfo . resSubExp) result) $ removeScopeAliases scope
matchReturnType rettype (map resSubExp result) result_ts
matchFunctionReturnType ::
(Mem rep inner, TC.Checkable rep) =>
[FunReturns] ->
Result ->
TC.TypeM rep ()
matchFunctionReturnType rettype result = do
matchRetTypeToResult rettype result
mapM_ (checkResultSubExp . resSubExp) result
where
checkResultSubExp Constant {} =
pure ()
checkResultSubExp (Var v) = do
dec <- varMemInfo v
case dec of
MemPrim _ -> pure ()
MemMem {} -> pure ()
MemAcc {} -> pure ()
MemArray _ _ _ (ArrayIn _ lmad)
| LMAD.isDirect lmad ->
pure ()
| otherwise ->
TC.bad . TC.TypeError $
"Array "
<> prettyText v
<> " returned by function, but has nontrivial index function:\n"
<> prettyText lmad
matchLoopResultMem ::
(Mem rep inner, TC.Checkable rep) =>
[FParam (Aliases rep)] ->
Result ->
TC.TypeM rep ()
matchLoopResultMem params = matchRetTypeToResult rettype
where
param_names = map paramName params
-- Invent a ReturnType so we can pretend that the loop body is
-- actually returning from a function.
rettype = map (toRet . paramDec) params
toExtV v
| Just i <- v `elemIndex` param_names = Ext i
| otherwise = Free v
toExtSE (Var v) = Var <$> toExtV v
toExtSE (Constant v) = Free $ Constant v
toRet (MemPrim t) =
MemPrim t
toRet (MemMem space) =
MemMem space
toRet (MemAcc acc ispace ts u) =
MemAcc acc ispace ts u
toRet (MemArray pt shape u (ArrayIn mem lmad))
| Just i <- mem `elemIndex` param_names,
Param _ _ (MemMem space) : _ <- drop i params =
MemArray pt shape' u $ ReturnsNewBlock space i lmad'
| otherwise =
MemArray pt shape' u $ ReturnsInBlock mem lmad'
where
shape' = fmap toExtSE shape
lmad' = existentialiseLMAD param_names lmad
matchBranchReturnType ::
(Mem rep inner, TC.Checkable rep) =>
[BodyReturns] ->
Body (Aliases rep) ->
TC.TypeM rep ()
matchBranchReturnType rettype (Body _ stms res) = do
scope <- askScope
ts <- runReaderT (mapM (subExpMemInfo . resSubExp) res) $ removeScopeAliases (scope <> scopeOf stms)
matchReturnType rettype (map resSubExp res) ts
-- | Helper function for index function unification.
--
-- The first return value maps a VName (wrapped in 'Free') to its Int
-- (wrapped in 'Ext'). In case of duplicates, it is mapped to the
-- *first* Int that occurs.
--
-- The second return value maps each Int (wrapped in an 'Ext') to a
-- 'LeafExp' 'Ext' with the Int at which its associated VName first
-- occurs.
getExtMaps ::
[(VName, Int)] ->
( M.Map (Ext VName) (TPrimExp Int64 (Ext VName)),
M.Map (Ext VName) (TPrimExp Int64 (Ext VName))
)
getExtMaps ctx_lst_ids =
( M.map leafExp $ M.mapKeys Free $ M.fromListWith (const id) ctx_lst_ids,
M.fromList $
mapMaybe
( traverse
( fmap (\i -> isInt64 $ LeafExp (Ext i) int64)
. (`lookup` ctx_lst_ids)
)
. uncurry (flip (,))
. fmap Ext
)
ctx_lst_ids
)
matchReturnType ::
(PP.Pretty u) =>
[MemInfo ExtSize u MemReturn] ->
[SubExp] ->
[MemInfo SubExp NoUniqueness MemBind] ->
TC.TypeM rep ()
matchReturnType rettype res ts = do
let existentialiseLMAD0 :: LMAD -> ExtLMAD
existentialiseLMAD0 = fmap $ fmap Free
fetchCtx i = case maybeNth i $ zip res ts of
Nothing ->
throwError $ "Cannot find variable #" <> prettyText i <> " in results: " <> prettyText res
Just (se, t) -> pure (se, t)
checkReturn (MemPrim x) (MemPrim y)
| x == y = pure ()
checkReturn (MemMem x) (MemMem y)
| x == y = pure ()
checkReturn (MemAcc xacc xispace xts _) (MemAcc yacc yispace yts _)
| (xacc, xispace, xts) == (yacc, yispace, yts) =
pure ()
checkReturn
(MemArray x_pt x_shape _ x_ret)
(MemArray y_pt y_shape _ y_ret)
| x_pt == y_pt,
shapeRank x_shape == shapeRank y_shape = do
zipWithM_ checkDim (shapeDims x_shape) (shapeDims y_shape)
checkMemReturn x_ret y_ret
checkReturn x y =
throwError $ T.unwords ["Expected", prettyText x, "but got", prettyText y]
checkDim (Free x) y
| x == y = pure ()
| otherwise =
throwError $ T.unwords ["Expected dim", prettyText x, "but got", prettyText y]
checkDim (Ext i) y = do
(x, _) <- fetchCtx i
unless (x == y) . throwError . T.unwords $
["Expected ext dim", prettyText i, "=>", prettyText x, "but got", prettyText y]
checkMemReturn (ReturnsInBlock x_mem x_lmad) (ArrayIn y_mem y_lmad)
| x_mem == y_mem =
unless (LMAD.closeEnough x_lmad $ existentialiseLMAD0 y_lmad) $
throwError . T.unwords $
[ "Index function unification failed (ReturnsInBlock)",
"\nlmad of body result: ",
prettyText y_lmad,
"\nlmad of return type: ",
prettyText x_lmad
]
checkMemReturn
(ReturnsNewBlock x_space x_ext x_lmad)
(ArrayIn y_mem y_lmad) = do
(x_mem, x_mem_type) <- fetchCtx x_ext
unless (LMAD.closeEnough x_lmad $ existentialiseLMAD0 y_lmad) $
throwError . docText $
"Index function unification failed (ReturnsNewBlock)"
</> "Lmad of body result:"
</> indent 2 (pretty y_lmad)
</> "Lmad of return type:"
</> indent 2 (pretty x_lmad)
case x_mem_type of
MemMem y_space ->
unless (x_space == y_space) . throwError . T.unwords $
[ "Expected memory",
prettyText y_mem,
"in space",
prettyText x_space,
"but actually in space",
prettyText y_space
]
t ->
throwError . T.unwords $
["Expected memory", prettyText x_ext, "=>", prettyText x_mem, "but but has type", prettyText t]
checkMemReturn x y =
throwError . docText $
"Expected array in"
</> indent 2 (pretty x)
</> "but array returned in"
</> indent 2 (pretty y)
bad s =
TC.bad . TC.TypeError . docText $
"Return type"
</> indent 2 (ppTupleLines' $ map pretty rettype)
</> "cannot match returns of results"
</> indent 2 (ppTupleLines' $ map pretty ts)
</> pretty s
unless (length rettype == length ts) $
TC.bad . TC.TypeError . docText $
"Return type"
</> indent 2 (ppTupleLines' $ map pretty rettype)
</> "does not have same number of elements as results"
</> indent 2 (ppTupleLines' $ map pretty ts)
either bad pure =<< runExceptT (zipWithM_ checkReturn rettype ts)
matchPatToExp ::
(Mem rep inner, LetDec rep ~ LetDecMem, TC.Checkable rep) =>
Pat (LetDec (Aliases rep)) ->
Exp (Aliases rep) ->
TC.TypeM rep ()
matchPatToExp pat e = do
scope <- asksScope removeScopeAliases
rt <- maybe illformed pure $ runReader (expReturns $ removeExpAliases e) scope
let (ctx_ids, val_ts) = unzip $ bodyReturnsFromPat $ removePatAliases pat
(ctx_map_ids, ctx_map_exts) = getExtMaps $ zip ctx_ids [0 .. 1]
ok =
length val_ts == length rt
&& and (zipWith (matches ctx_map_ids ctx_map_exts) val_ts rt)
unless ok . TC.bad . TC.TypeError . docText $
"Expression type:"
</> indent 2 (ppTupleLines' $ map pretty rt)
</> "cannot match pattern type:"
</> indent 2 (ppTupleLines' $ map pretty val_ts)
where
illformed =
TC.bad $
TC.TypeError . docText $
"Expression"
</> indent 2 (pretty e)
</> "cannot be assigned an index function."
matches _ _ (MemPrim x) (MemPrim y) = x == y
matches _ _ (MemMem x_space) (MemMem y_space) =
x_space == y_space
matches _ _ (MemAcc x_accs x_ispace x_ts _) (MemAcc y_accs y_ispace y_ts _) =
(x_accs, x_ispace, x_ts) == (y_accs, y_ispace, y_ts)
matches ctxids ctxexts (MemArray x_pt x_shape _ x_ret) (MemArray y_pt y_shape _ y_ret) =
x_pt == y_pt
&& x_shape == y_shape
&& case (x_ret, y_ret) of
(ReturnsInBlock _ x_lmad, Just (ReturnsInBlock _ y_lmad)) ->
let x_lmad' = LMAD.substitute ctxids x_lmad
y_lmad' = LMAD.substitute ctxexts y_lmad
in LMAD.closeEnough x_lmad' y_lmad'
( ReturnsInBlock _ x_lmad,
Just (ReturnsNewBlock _ _ y_lmad)
) ->
let x_lmad' = LMAD.substitute ctxids x_lmad
y_lmad' = LMAD.substitute ctxexts y_lmad
in LMAD.closeEnough x_lmad' y_lmad'
( ReturnsNewBlock _ x_i x_lmad,
Just (ReturnsNewBlock _ y_i y_lmad)
) ->
let x_lmad' = LMAD.substitute ctxids x_lmad
y_lmad' = LMAD.substitute ctxexts y_lmad
in x_i == y_i && LMAD.closeEnough x_lmad' y_lmad'
(_, Nothing) -> True
_ -> False
matches _ _ _ _ = False
varMemInfo ::
(Mem rep inner) =>
VName ->
TC.TypeM rep (MemInfo SubExp NoUniqueness MemBind)
varMemInfo name = do
dec <- TC.lookupVar name
case dec of
LetName (_, summary) -> pure $ letDecMem summary
FParamName summary -> pure $ noUniquenessReturns summary
LParamName summary -> pure summary
IndexName it -> pure $ MemPrim $ IntType it
-- | Turn info into memory information.
nameInfoToMemInfo :: (Mem rep inner) => NameInfo rep -> MemBound NoUniqueness
nameInfoToMemInfo info =
case info of
FParamName summary -> noUniquenessReturns summary
LParamName summary -> summary
LetName summary -> letDecMem summary
IndexName it -> MemPrim $ IntType it
-- | Look up information about the memory block with this name.
lookupMemInfo ::
(HasScope rep m, Mem rep inner) =>
VName ->
m (MemInfo SubExp NoUniqueness MemBind)
lookupMemInfo = fmap nameInfoToMemInfo . lookupInfo
subExpMemInfo ::
(HasScope rep m, Mem rep inner) =>
SubExp ->
m (MemInfo SubExp NoUniqueness MemBind)
subExpMemInfo (Var v) = lookupMemInfo v
subExpMemInfo (Constant v) = pure $ MemPrim $ primValueType v
lookupArraySummary ::
(Mem rep inner, HasScope rep m, Monad m) =>
VName ->
m (VName, LMAD.LMAD (TPrimExp Int64 VName))
lookupArraySummary name = do
summary <- lookupMemInfo name
case summary of
MemArray _ _ _ (ArrayIn mem lmad) ->
pure (mem, lmad)
_ ->
error . T.unpack $
"Expected "
<> prettyText name
<> " to be array but bound to:\n"
<> prettyText summary
lookupMemSpace ::
(Mem rep inner, HasScope rep m, Monad m) =>
VName ->
m Space
lookupMemSpace name = do
summary <- lookupMemInfo name
case summary of
MemMem space ->
pure space
_ ->
error . T.unpack $
"Expected "
<> prettyText name
<> " to be memory but bound to:\n"
<> prettyText summary
checkMemInfo ::
(TC.Checkable rep) =>
VName ->
MemInfo SubExp u MemBind ->
TC.TypeM rep ()
checkMemInfo _ (MemPrim _) = pure ()
checkMemInfo _ (MemMem (ScalarSpace d _)) = mapM_ (TC.require (Prim int64)) d
checkMemInfo _ (MemMem _) = pure ()
checkMemInfo _ (MemAcc acc ispace ts u) =
TC.checkType $ Acc acc ispace ts u
checkMemInfo name (MemArray _ shape _ (ArrayIn v lmad)) = do
t <- lookupType v
case t of
Mem {} ->
pure ()
_ ->
TC.bad $
TC.TypeError $
"Variable "
<> prettyText v
<> " used as memory block, but is of type "
<> prettyText t
<> "."
TC.context ("in index function " <> prettyText lmad) $ do
traverse_ (TC.requirePrimExp int64 . untyped) lmad
unless (LMAD.shape lmad == map pe64 (shapeDims shape)) $
TC.bad $
TC.TypeError $
"Shape of index function ("
<> prettyText (LMAD.shape lmad)
<> ") does not match shape of array "
<> prettyText name
<> " ("
<> prettyText shape
<> ")"
bodyReturnsFromPat ::
Pat (MemBound NoUniqueness) -> [(VName, BodyReturns)]
bodyReturnsFromPat pat =
map asReturns $ patElems pat
where
ctx = patElems pat
ext (Var v)
| Just (i, _) <- find ((== v) . patElemName . snd) $ zip [0 ..] ctx =
Ext i
ext se = Free se
asReturns pe =
( patElemName pe,
case patElemDec pe of
MemPrim pt -> MemPrim pt
MemMem space -> MemMem space
MemArray pt shape u (ArrayIn mem lmad) ->
MemArray pt (Shape $ map ext $ shapeDims shape) u $
case find ((== mem) . patElemName . snd) $ zip [0 ..] ctx of
Just (i, PatElem _ (MemMem space)) ->
ReturnsNewBlock space i $
existentialiseLMAD (map patElemName ctx) lmad
_ -> ReturnsInBlock mem $ existentialiseLMAD [] lmad
MemAcc acc ispace ts u -> MemAcc acc ispace ts u
)
extReturns :: [ExtType] -> [ExpReturns]
extReturns ets =
evalState (mapM addDec ets) 0
where
addDec (Prim bt) =
pure $ MemPrim bt
addDec (Mem space) =
pure $ MemMem space
addDec t@(Array bt shape u)
| existential t = do
i <- get <* modify (+ 1)
pure . MemArray bt shape u . Just $
ReturnsNewBlock DefaultSpace i $
LMAD.iota 0 (map convert $ shapeDims shape)
| otherwise =
pure $ MemArray bt shape u Nothing
addDec (Acc acc ispace ts u) =
pure $ MemAcc acc ispace ts u
convert (Ext i) = le64 (Ext i)
convert (Free v) = Free <$> pe64 v
arrayVarReturns ::
(HasScope rep m, Monad m, Mem rep inner) =>
VName ->
m (PrimType, Shape, VName, LMAD)
arrayVarReturns v = do
summary <- lookupMemInfo v
case summary of
MemArray et shape _ (ArrayIn mem lmad) ->
pure (et, Shape $ shapeDims shape, mem, lmad)
_ ->
error . T.unpack $ "arrayVarReturns: " <> prettyText v <> " is not an array."
varReturns ::
(HasScope rep m, Monad m, Mem rep inner) =>
VName ->
m ExpReturns
varReturns v = do
summary <- lookupMemInfo v
case summary of
MemPrim bt ->
pure $ MemPrim bt
MemArray et shape _ (ArrayIn mem lmad) ->
pure $
MemArray et (fmap Free shape) NoUniqueness $
Just $
ReturnsInBlock mem $
existentialiseLMAD [] lmad
MemMem space ->
pure $ MemMem space
MemAcc acc ispace ts u ->
pure $ MemAcc acc ispace ts u
subExpReturns :: (HasScope rep m, Monad m, Mem rep inner) => SubExp -> m ExpReturns
subExpReturns (Var v) =
varReturns v
subExpReturns (Constant v) =
pure $ MemPrim $ primValueType v
-- | The return information of an expression. This can be seen as the
-- "return type with memory annotations" of the expression.
--
-- This can produce Nothing, which signifies that the result is an
-- array layout that is not expressible as an index function.
expReturns ::
(LocalScope rep m, Mem rep inner) =>
Exp rep ->
m (Maybe [ExpReturns])
expReturns (BasicOp (SubExp se)) =
Just . pure <$> subExpReturns se
expReturns (BasicOp (Opaque _ (Var v))) =
Just . pure <$> varReturns v
expReturns (BasicOp (Reshape v newshape)) = do
(et, _, mem, lmad) <- arrayVarReturns v
case reshaper (reshapeKind newshape) lmad $ map pe64 $ shapeDims $ newShape newshape of
Just lmad' ->
pure . Just $
[ MemArray et (Free <$> newShape newshape) NoUniqueness . Just $
ReturnsInBlock mem (existentialiseLMAD [] lmad')
]
Nothing -> pure Nothing
where
reshaper ReshapeArbitrary lmad =
LMAD.reshape lmad
reshaper ReshapeCoerce lmad =
Just . LMAD.coerce lmad
expReturns (BasicOp (Rearrange v perm)) = do
(et, Shape dims, mem, lmad) <- arrayVarReturns v
let lmad' = LMAD.permute lmad perm
dims' = rearrangeShape perm dims
pure $
Just
[ MemArray et (Shape $ map Free dims') NoUniqueness $
Just $
ReturnsInBlock mem $
existentialiseLMAD [] lmad'
]
expReturns (BasicOp (Index v slice)) = do
Just . pure . varInfoToExpReturns <$> sliceInfo v slice
expReturns (BasicOp (Update _ v _ _)) =
Just . pure <$> varReturns v
expReturns (BasicOp (FlatIndex v slice)) =
Just . pure . varInfoToExpReturns <$> flatSliceInfo v slice
expReturns (BasicOp (FlatUpdate v _ _)) =
Just . pure <$> varReturns v
expReturns (BasicOp op) =
Just . extReturns . staticShapes <$> basicOpType op
expReturns e@(Loop merge _ _) = do
t <- expExtType e
Just <$> zipWithM typeWithDec t (map fst merge)
where
typeWithDec t p =
case (t, paramDec p) of
( Array pt shape u,
MemArray _ _ _ (ArrayIn mem lmad)
)
| Just (i, mem_p) <- isLoopVar mem,
Mem space <- paramType mem_p ->
pure $ MemArray pt shape u $ Just $ ReturnsNewBlock space i lmad'
| otherwise ->
pure $ MemArray pt shape u $ Just $ ReturnsInBlock mem lmad'
where
lmad' = existentialiseLMAD (map paramName mergevars) lmad
(Array {}, _) ->
error "expReturns: Array return type but not array merge variable."
(Acc acc ispace ts u, _) ->
pure $ MemAcc acc ispace ts u
(Prim pt, _) ->
pure $ MemPrim pt
(Mem space, _) ->
pure $ MemMem space
isLoopVar v = find ((== v) . paramName . snd) $ zip [0 ..] mergevars
mergevars = map fst merge
expReturns (Apply _ _ ret _) =
pure $ Just $ map (funReturnsToExpReturns . fst) ret
expReturns (Match _ _ _ (MatchDec ret _)) =
pure $ Just $ map bodyReturnsToExpReturns ret
expReturns (Op op) =
Just <$> opReturns op
expReturns (WithAcc inputs lam) =
Just
<$> ( (<>)
<$> (concat <$> mapM inputReturns inputs)
<*>
-- XXX: this is a bit dubious because it enforces extra copies. I
-- think WithAcc should perhaps have a return annotation like If.
pure (extReturns $ staticShapes $ drop num_accs $ lambdaReturnType lam)
)
where
inputReturns (_, arrs, _) = mapM varReturns arrs
num_accs = length inputs
sliceInfo ::
(Monad m, HasScope rep m, Mem rep inner) =>
VName ->
Slice SubExp ->
m (MemInfo SubExp NoUniqueness MemBind)
sliceInfo v slice = do
(et, _, mem, lmad) <- arrayVarReturns v
case sliceDims slice of
[] -> pure $ MemPrim et
dims ->
pure $
MemArray et (Shape dims) NoUniqueness . ArrayIn mem $
LMAD.slice lmad (fmap pe64 slice)
flatSliceInfo ::
(Monad m, HasScope rep m, Mem rep inner) =>
VName ->
FlatSlice SubExp ->
m (MemInfo SubExp NoUniqueness MemBind)
flatSliceInfo v slice@(FlatSlice offset idxs) = do
(et, _, mem, lmad) <- arrayVarReturns v
map (fmap pe64) idxs
& FlatSlice (pe64 offset)
& LMAD.flatSlice lmad
& MemArray et (Shape (flatSliceDims slice)) NoUniqueness . ArrayIn mem
& pure
class (IsOp op) => OpReturns op where
opReturns :: (Mem rep inner, Monad m, HasScope rep m) => op rep -> m [ExpReturns]
opReturns op = extReturns <$> opType op
instance (OpReturns inner) => OpReturns (MemOp inner) where
opReturns (Alloc _ space) = pure [MemMem space]
opReturns (EnsureDirect v) = do
space <- lookupMemSpace . fst =<< lookupArraySummary v
summary <- lookupMemInfo v
case summary of
MemArray et shape _ _ ->
pure
[ MemMem space,
MemArray et (fmap Free shape) NoUniqueness . Just $
ReturnsNewBlock space 0 . LMAD.iota 0 $
fmap (fmap Free . pe64) (shapeDims shape)
]
_ -> error "EnsureDirect applied to non-array"
opReturns (Inner op) = opReturns op
instance OpReturns NoOp where
opReturns NoOp = pure []
applyFunReturns ::
(Typed dec) =>
[FunReturns] ->
[Param dec] ->
[(SubExp, Type)] ->
Maybe [FunReturns]
applyFunReturns rets params args
| Just _ <- applyRetType rettype params args =
Just $ map correctDims rets
| otherwise =
Nothing
where
rettype = map declExtTypeOf rets
parammap :: M.Map VName (SubExp, Type)
parammap =
M.fromList $
zip (map paramName params) args
substSubExp (Var v)
| Just (se, _) <- M.lookup v parammap = se
substSubExp se = se
correctDims (MemPrim t) =
MemPrim t
correctDims (MemMem space) =
MemMem space
correctDims (MemArray et shape u memsummary) =
MemArray et (correctShape shape) u $
correctSummary memsummary
correctDims (MemAcc acc ispace ts u) =
MemAcc acc ispace ts u
correctShape = Shape . map correctDim . shapeDims
correctDim (Ext i) = Ext i
correctDim (Free se) = Free $ substSubExp se
correctSummary (ReturnsNewBlock space i lmad) =
ReturnsNewBlock space i lmad
correctSummary (ReturnsInBlock mem lmad) =
-- FIXME: we should also do a replacement in lmad here.
ReturnsInBlock mem' lmad
where
mem' = case M.lookup mem parammap of
Just (Var v, _) -> v
_ -> mem