{-# LANGUAGE TypeFamilies #-}
module Futhark.IR.Prop.TypeOf
( expExtType,
subExpType,
subExpResType,
basicOpType,
mapType,
module Futhark.IR.RetType,
module Futhark.IR.Prop.Scope,
TypedOp (..),
)
where
import Data.List.NonEmpty (NonEmpty (..))
import Futhark.IR.Prop.Constants
import Futhark.IR.Prop.Scope
import Futhark.IR.Prop.Types
import Futhark.IR.RetType
import Futhark.IR.Syntax
subExpType :: HasScope t m => SubExp -> m Type
subExpType :: forall t (m :: * -> *). HasScope t m => SubExp -> m Type
subExpType (Constant PrimValue
val) = forall (f :: * -> *) a. Applicative f => a -> f a
pure forall a b. (a -> b) -> a -> b
$ forall shape u. PrimType -> TypeBase shape u
Prim forall a b. (a -> b) -> a -> b
$ PrimValue -> PrimType
primValueType PrimValue
val
subExpType (Var VName
name) = forall rep (m :: * -> *). HasScope rep m => VName -> m Type
lookupType VName
name
subExpResType :: HasScope t m => SubExpRes -> m Type
subExpResType :: forall t (m :: * -> *). HasScope t m => SubExpRes -> m Type
subExpResType = forall t (m :: * -> *). HasScope t m => SubExp -> m Type
subExpType forall b c a. (b -> c) -> (a -> b) -> a -> c
. SubExpRes -> SubExp
resSubExp
mapType :: SubExp -> Lambda rep -> [Type]
mapType :: forall rep. SubExp -> Lambda rep -> [Type]
mapType SubExp
outersize Lambda rep
f =
[ forall shape u_unused u.
ArrayShape shape =>
TypeBase shape u_unused -> shape -> u -> TypeBase shape u
arrayOf Type
t (forall d. [d] -> ShapeBase d
Shape [SubExp
outersize]) NoUniqueness
NoUniqueness
| Type
t <- forall rep. Lambda rep -> [Type]
lambdaReturnType Lambda rep
f
]
basicOpType :: HasScope rep m => BasicOp -> m [Type]
basicOpType :: forall rep (m :: * -> *). HasScope rep m => BasicOp -> m [Type]
basicOpType (SubExp SubExp
se) =
forall (f :: * -> *) a. Applicative f => a -> f a
pure forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> forall t (m :: * -> *). HasScope t m => SubExp -> m Type
subExpType SubExp
se
basicOpType (Opaque OpaqueOp
_ SubExp
se) =
forall (f :: * -> *) a. Applicative f => a -> f a
pure forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> forall t (m :: * -> *). HasScope t m => SubExp -> m Type
subExpType SubExp
se
basicOpType (ArrayLit [SubExp]
es Type
rt) =
forall (f :: * -> *) a. Applicative f => a -> f a
pure [forall shape u_unused u.
ArrayShape shape =>
TypeBase shape u_unused -> shape -> u -> TypeBase shape u
arrayOf Type
rt (forall d. [d] -> ShapeBase d
Shape [SubExp
n]) NoUniqueness
NoUniqueness]
where
n :: SubExp
n = IntType -> Integer -> SubExp
intConst IntType
Int64 forall a b. (a -> b) -> a -> b
$ forall a. Integral a => a -> Integer
toInteger forall a b. (a -> b) -> a -> b
$ forall (t :: * -> *) a. Foldable t => t a -> Int
length [SubExp]
es
basicOpType (BinOp BinOp
bop SubExp
_ SubExp
_) =
forall (f :: * -> *) a. Applicative f => a -> f a
pure [forall shape u. PrimType -> TypeBase shape u
Prim forall a b. (a -> b) -> a -> b
$ BinOp -> PrimType
binOpType BinOp
bop]
basicOpType (UnOp UnOp
_ SubExp
x) =
forall (f :: * -> *) a. Applicative f => a -> f a
pure forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> forall t (m :: * -> *). HasScope t m => SubExp -> m Type
subExpType SubExp
x
basicOpType CmpOp {} =
forall (f :: * -> *) a. Applicative f => a -> f a
pure [forall shape u. PrimType -> TypeBase shape u
Prim PrimType
Bool]
basicOpType (ConvOp ConvOp
conv SubExp
_) =
forall (f :: * -> *) a. Applicative f => a -> f a
pure [forall shape u. PrimType -> TypeBase shape u
Prim forall a b. (a -> b) -> a -> b
$ forall a b. (a, b) -> b
snd forall a b. (a -> b) -> a -> b
$ ConvOp -> (PrimType, PrimType)
convOpType ConvOp
conv]
basicOpType (Index VName
ident Slice SubExp
slice) =
Type -> [Type]
result forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> forall rep (m :: * -> *). HasScope rep m => VName -> m Type
lookupType VName
ident
where
result :: Type -> [Type]
result Type
t = [forall shape u. PrimType -> TypeBase shape u
Prim (forall shape u. TypeBase shape u -> PrimType
elemType Type
t) Type -> Shape -> Type
`arrayOfShape` Shape
shape]
shape :: Shape
shape = forall d. [d] -> ShapeBase d
Shape forall a b. (a -> b) -> a -> b
$ forall d. Slice d -> [d]
sliceDims Slice SubExp
slice
basicOpType (Update Safety
_ VName
src Slice SubExp
_ SubExp
_) =
forall (f :: * -> *) a. Applicative f => a -> f a
pure forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> forall rep (m :: * -> *). HasScope rep m => VName -> m Type
lookupType VName
src
basicOpType (FlatIndex VName
ident FlatSlice SubExp
slice) =
Type -> [Type]
result forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> forall rep (m :: * -> *). HasScope rep m => VName -> m Type
lookupType VName
ident
where
result :: Type -> [Type]
result Type
t = [forall shape u. PrimType -> TypeBase shape u
Prim (forall shape u. TypeBase shape u -> PrimType
elemType Type
t) Type -> Shape -> Type
`arrayOfShape` Shape
shape]
shape :: Shape
shape = forall d. [d] -> ShapeBase d
Shape forall a b. (a -> b) -> a -> b
$ forall d. FlatSlice d -> [d]
flatSliceDims FlatSlice SubExp
slice
basicOpType (FlatUpdate VName
src FlatSlice SubExp
_ VName
_) =
forall (f :: * -> *) a. Applicative f => a -> f a
pure forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> forall rep (m :: * -> *). HasScope rep m => VName -> m Type
lookupType VName
src
basicOpType (Iota SubExp
n SubExp
_ SubExp
_ IntType
et) =
forall (f :: * -> *) a. Applicative f => a -> f a
pure [forall shape u_unused u.
ArrayShape shape =>
TypeBase shape u_unused -> shape -> u -> TypeBase shape u
arrayOf (forall shape u. PrimType -> TypeBase shape u
Prim (IntType -> PrimType
IntType IntType
et)) (forall d. [d] -> ShapeBase d
Shape [SubExp
n]) NoUniqueness
NoUniqueness]
basicOpType (Replicate (Shape []) SubExp
e) =
forall (f :: * -> *) a. Applicative f => a -> f a
pure forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> forall t (m :: * -> *). HasScope t m => SubExp -> m Type
subExpType SubExp
e
basicOpType (Replicate Shape
shape SubExp
e) =
forall (f :: * -> *) a. Applicative f => a -> f a
pure forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall a b c. (a -> b -> c) -> b -> a -> c
flip Type -> Shape -> Type
arrayOfShape Shape
shape forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> forall t (m :: * -> *). HasScope t m => SubExp -> m Type
subExpType SubExp
e
basicOpType (Scratch PrimType
t [SubExp]
shape) =
forall (f :: * -> *) a. Applicative f => a -> f a
pure [forall shape u_unused u.
ArrayShape shape =>
TypeBase shape u_unused -> shape -> u -> TypeBase shape u
arrayOf (forall shape u. PrimType -> TypeBase shape u
Prim PrimType
t) (forall d. [d] -> ShapeBase d
Shape [SubExp]
shape) NoUniqueness
NoUniqueness]
basicOpType (Reshape ReshapeKind
_ (Shape []) VName
e) =
forall {shape} {u} {shape} {u}.
TypeBase shape u -> [TypeBase shape u]
result forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> forall rep (m :: * -> *). HasScope rep m => VName -> m Type
lookupType VName
e
where
result :: TypeBase shape u -> [TypeBase shape u]
result TypeBase shape u
t = [forall shape u. PrimType -> TypeBase shape u
Prim forall a b. (a -> b) -> a -> b
$ forall shape u. TypeBase shape u -> PrimType
elemType TypeBase shape u
t]
basicOpType (Reshape ReshapeKind
_ Shape
shape VName
e) =
Type -> [Type]
result forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> forall rep (m :: * -> *). HasScope rep m => VName -> m Type
lookupType VName
e
where
result :: Type -> [Type]
result Type
t = [Type
t forall newshape oldshape u.
ArrayShape newshape =>
TypeBase oldshape u -> newshape -> TypeBase newshape u
`setArrayShape` Shape
shape]
basicOpType (Rearrange [Int]
perm VName
e) =
Type -> [Type]
result forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> forall rep (m :: * -> *). HasScope rep m => VName -> m Type
lookupType VName
e
where
result :: Type -> [Type]
result Type
t = [[Int] -> Type -> Type
rearrangeType [Int]
perm Type
t]
basicOpType (Rotate [SubExp]
_ VName
e) =
forall (f :: * -> *) a. Applicative f => a -> f a
pure forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> forall rep (m :: * -> *). HasScope rep m => VName -> m Type
lookupType VName
e
basicOpType (Concat Int
i (VName
x :| [VName]
_) SubExp
ressize) =
Type -> [Type]
result forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> forall rep (m :: * -> *). HasScope rep m => VName -> m Type
lookupType VName
x
where
result :: Type -> [Type]
result Type
xt = [forall d u.
ArrayShape (ShapeBase d) =>
Int -> TypeBase (ShapeBase d) u -> d -> TypeBase (ShapeBase d) u
setDimSize Int
i Type
xt SubExp
ressize]
basicOpType (Copy VName
v) =
forall (f :: * -> *) a. Applicative f => a -> f a
pure forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> forall rep (m :: * -> *). HasScope rep m => VName -> m Type
lookupType VName
v
basicOpType (Manifest [Int]
_ VName
v) =
forall (f :: * -> *) a. Applicative f => a -> f a
pure forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> forall rep (m :: * -> *). HasScope rep m => VName -> m Type
lookupType VName
v
basicOpType Assert {} =
forall (f :: * -> *) a. Applicative f => a -> f a
pure [forall shape u. PrimType -> TypeBase shape u
Prim PrimType
Unit]
basicOpType (UpdateAcc VName
v [SubExp]
_ [SubExp]
_) =
forall (f :: * -> *) a. Applicative f => a -> f a
pure forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> forall rep (m :: * -> *). HasScope rep m => VName -> m Type
lookupType VName
v
expExtType ::
(HasScope rep m, TypedOp (Op rep)) =>
Exp rep ->
m [ExtType]
expExtType :: forall rep (m :: * -> *).
(HasScope rep m, TypedOp (Op rep)) =>
Exp rep -> m [ExtType]
expExtType (Apply Name
_ [(SubExp, Diet)]
_ [RetType rep]
rt (Safety, SrcLoc, [SrcLoc])
_) = forall (f :: * -> *) a. Applicative f => a -> f a
pure forall a b. (a -> b) -> a -> b
$ forall a b. (a -> b) -> [a] -> [b]
map (forall shape.
TypeBase shape Uniqueness -> TypeBase shape NoUniqueness
fromDecl forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall t. DeclExtTyped t => t -> DeclExtType
declExtTypeOf) [RetType rep]
rt
expExtType (Match [SubExp]
_ [Case (Body rep)]
_ Body rep
_ MatchDec (BranchType rep)
rt) = forall (f :: * -> *) a. Applicative f => a -> f a
pure forall a b. (a -> b) -> a -> b
$ forall a b. (a -> b) -> [a] -> [b]
map forall t. ExtTyped t => t -> ExtType
extTypeOf forall a b. (a -> b) -> a -> b
$ forall rt. MatchDec rt -> [rt]
matchReturns MatchDec (BranchType rep)
rt
expExtType (DoLoop [(FParam rep, SubExp)]
merge LoopForm rep
_ Body rep
_) =
forall (f :: * -> *) a. Applicative f => a -> f a
pure forall a b. (a -> b) -> a -> b
$ forall dec. Typed dec => [Param dec] -> [ExtType]
loopExtType forall a b. (a -> b) -> a -> b
$ forall a b. (a -> b) -> [a] -> [b]
map forall a b. (a, b) -> a
fst [(FParam rep, SubExp)]
merge
expExtType (BasicOp BasicOp
op) = forall u. [TypeBase Shape u] -> [TypeBase ExtShape u]
staticShapes forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> forall rep (m :: * -> *). HasScope rep m => BasicOp -> m [Type]
basicOpType BasicOp
op
expExtType (WithAcc [WithAccInput rep]
inputs Lambda rep
lam) =
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap forall u. [TypeBase Shape u] -> [TypeBase ExtShape u]
staticShapes forall a b. (a -> b) -> a -> b
$
forall a. Semigroup a => a -> a -> a
(<>)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> (forall (t :: * -> *) a. Foldable t => t [a] -> [a]
concat forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> forall (t :: * -> *) (f :: * -> *) a b.
(Traversable t, Applicative f) =>
(a -> f b) -> t a -> f (t b)
traverse forall {t :: * -> *} {f :: * -> *} {rep} {a} {c}.
(Traversable t, HasScope rep f) =>
(a, t VName, c) -> f (t Type)
inputType [WithAccInput rep]
inputs)
forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> forall (f :: * -> *) a. Applicative f => a -> f a
pure (forall a. Int -> [a] -> [a]
drop Int
num_accs (forall rep. Lambda rep -> [Type]
lambdaReturnType Lambda rep
lam))
where
inputType :: (a, t VName, c) -> f (t Type)
inputType (a
_, t VName
arrs, c
_) = forall (t :: * -> *) (f :: * -> *) a b.
(Traversable t, Applicative f) =>
(a -> f b) -> t a -> f (t b)
traverse forall rep (m :: * -> *). HasScope rep m => VName -> m Type
lookupType t VName
arrs
num_accs :: Int
num_accs = forall (t :: * -> *) a. Foldable t => t a -> Int
length [WithAccInput rep]
inputs
expExtType (Op Op rep
op) = forall op t (m :: * -> *).
(TypedOp op, HasScope t m) =>
op -> m [ExtType]
opType Op rep
op
loopExtType :: Typed dec => [Param dec] -> [ExtType]
loopExtType :: forall dec. Typed dec => [Param dec] -> [ExtType]
loopExtType [Param dec]
params =
[VName] -> [ExtType] -> [ExtType]
existentialiseExtTypes [VName]
inaccessible forall a b. (a -> b) -> a -> b
$ forall u. [TypeBase Shape u] -> [TypeBase ExtShape u]
staticShapes forall a b. (a -> b) -> a -> b
$ forall a b. (a -> b) -> [a] -> [b]
map forall t. Typed t => t -> Type
typeOf [Param dec]
params
where
inaccessible :: [VName]
inaccessible = forall a b. (a -> b) -> [a] -> [b]
map forall dec. Param dec -> VName
paramName [Param dec]
params
class TypedOp op where
opType :: HasScope t m => op -> m [ExtType]
instance TypedOp (NoOp rep) where
opType :: forall t (m :: * -> *). HasScope t m => NoOp rep -> m [ExtType]
opType NoOp rep
NoOp = forall (f :: * -> *) a. Applicative f => a -> f a
pure []