{-# LANGUAGE Strict #-}
{-# LANGUAGE TypeFamilies #-}
module Futhark.Internalise.Exps (transformProg) where
import Control.Monad.Reader
import Data.List (elemIndex, find, intercalate, intersperse, transpose)
import Data.List.NonEmpty (NonEmpty (..))
import Data.List.NonEmpty qualified as NE
import Data.Map.Strict qualified as M
import Data.Set qualified as S
import Data.Text qualified as T
import Futhark.IR.SOACS as I hiding (stmPat)
import Futhark.Internalise.AccurateSizes
import Futhark.Internalise.Bindings
import Futhark.Internalise.Entry
import Futhark.Internalise.Lambdas
import Futhark.Internalise.Monad as I
import Futhark.Internalise.TypesValues
import Futhark.Transform.Rename as I
import Futhark.Util (splitAt3)
import Futhark.Util.Pretty (align, docText, pretty)
import Language.Futhark as E hiding (TypeArg)
transformProg :: MonadFreshNames m => Bool -> VisibleTypes -> [E.ValBind] -> m (I.Prog SOACS)
transformProg :: forall (m :: * -> *).
MonadFreshNames m =>
Bool -> VisibleTypes -> [ValBind] -> m (Prog SOACS)
transformProg Bool
always_safe VisibleTypes
types [ValBind]
vbinds = do
(OpaqueTypes
opaques, Stms SOACS
consts, [FunDef SOACS]
funs) <-
forall (m :: * -> *).
MonadFreshNames m =>
Bool
-> InternaliseM () -> m (OpaqueTypes, Stms SOACS, [FunDef SOACS])
runInternaliseM Bool
always_safe (VisibleTypes -> [ValBind] -> InternaliseM ()
internaliseValBinds VisibleTypes
types [ValBind]
vbinds)
forall rep (m :: * -> *).
(Renameable rep, MonadFreshNames m) =>
Prog rep -> m (Prog rep)
I.renameProg forall a b. (a -> b) -> a -> b
$ forall rep. OpaqueTypes -> Stms rep -> [FunDef rep] -> Prog rep
I.Prog OpaqueTypes
opaques Stms SOACS
consts [FunDef SOACS]
funs
internaliseValBinds :: VisibleTypes -> [E.ValBind] -> InternaliseM ()
internaliseValBinds :: VisibleTypes -> [ValBind] -> InternaliseM ()
internaliseValBinds VisibleTypes
types = forall (t :: * -> *) (m :: * -> *) a b.
(Foldable t, Monad m) =>
(a -> m b) -> t a -> m ()
mapM_ forall a b. (a -> b) -> a -> b
$ VisibleTypes -> ValBind -> InternaliseM ()
internaliseValBind VisibleTypes
types
internaliseFunName :: VName -> Name
internaliseFunName :: VName -> Name
internaliseFunName = [Char] -> Name
nameFromString forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall a. Pretty a => a -> [Char]
prettyString
internaliseValBind :: VisibleTypes -> E.ValBind -> InternaliseM ()
internaliseValBind :: VisibleTypes -> ValBind -> InternaliseM ()
internaliseValBind VisibleTypes
types fb :: ValBind
fb@(E.ValBind Maybe (Info EntryPoint)
entry VName
fname Maybe (TypeExp VName)
retdecl (Info StructRetType
rettype) [TypeParamBase VName]
tparams [PatBase Info VName]
params Exp
body Maybe DocComment
_ [AttrInfo VName]
attrs SrcLoc
loc) = do
forall a.
[TypeParamBase VName]
-> [PatBase Info VName]
-> ([FParam SOACS] -> [[FParam SOACS]] -> InternaliseM a)
-> InternaliseM a
bindingFParams [TypeParamBase VName]
tparams [PatBase Info VName]
params forall a b. (a -> b) -> a -> b
$ \[FParam SOACS]
shapeparams [[FParam SOACS]]
params' -> do
let shapenames :: [VName]
shapenames = forall a b. (a -> b) -> [a] -> [b]
map forall dec. Param dec -> VName
I.paramName [FParam SOACS]
shapeparams
ErrorMsg SubExp
msg <- case Maybe (TypeExp VName)
retdecl of
Just TypeExp VName
dt ->
forall a. [ErrorMsgPart a] -> ErrorMsg a
errorMsg
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (ErrorMsgPart SubExp
"Function return value does not match shape of type " :)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> TypeExp VName -> InternaliseM [ErrorMsgPart SubExp]
typeExpForError TypeExp VName
dt
Maybe (TypeExp VName)
Nothing -> forall (f :: * -> *) a. Applicative f => a -> f a
pure forall a b. (a -> b) -> a -> b
$ forall a. [ErrorMsgPart a] -> ErrorMsg a
errorMsg [ErrorMsgPart SubExp
"Function return value does not match shape of declared return type."]
(Body SOACS
body', [TypeBase ExtShape Uniqueness]
rettype') <- forall (m :: * -> *) a.
MonadBuilder m =>
m (Result, a) -> m (Body (Rep m), a)
buildBody forall a b. (a -> b) -> a -> b
$ do
[SubExp]
body_res <- [Char] -> Exp -> InternaliseM [SubExp]
internaliseExp (VName -> [Char]
baseString VName
fname forall a. Semigroup a => a -> a -> a
<> [Char]
"_res") Exp
body
[TypeBase ExtShape Uniqueness]
rettype' <-
forall {u}. [TypeBase ExtShape u] -> [TypeBase ExtShape u]
zeroExts forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall shape u.
StructRetType
-> [TypeBase shape u] -> [TypeBase ExtShape Uniqueness]
internaliseReturnType StructRetType
rettype forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
(a -> m b) -> t a -> m (t b)
mapM forall t (m :: * -> *).
HasScope t m =>
SubExp -> m (TypeBase Shape NoUniqueness)
subExpType [SubExp]
body_res
Result
body_res' <-
ErrorMsg SubExp
-> SrcLoc -> [ExtType] -> Result -> InternaliseM Result
ensureResultExtShape ErrorMsg SubExp
msg SrcLoc
loc (forall a b. (a -> b) -> [a] -> [b]
map forall shape.
TypeBase shape Uniqueness -> TypeBase shape NoUniqueness
I.fromDecl [TypeBase ExtShape Uniqueness]
rettype') forall a b. (a -> b) -> a -> b
$ [SubExp] -> Result
subExpsRes [SubExp]
body_res
forall (f :: * -> *) a. Applicative f => a -> f a
pure
( Result
body_res',
forall a. Int -> a -> [a]
replicate (forall (t :: * -> *) a. Foldable t => t a -> Int
length (forall u. [TypeBase ExtShape u] -> Set Int
shapeContext [TypeBase ExtShape Uniqueness]
rettype')) (forall shape u. PrimType -> TypeBase shape u
I.Prim PrimType
int64) forall a. [a] -> [a] -> [a]
++ [TypeBase ExtShape Uniqueness]
rettype'
)
let all_params :: [Param DeclType]
all_params = [FParam SOACS]
shapeparams forall a. [a] -> [a] -> [a]
++ forall (t :: * -> *) a. Foldable t => t [a] -> [a]
concat [[FParam SOACS]]
params'
Attrs
attrs' <- [AttrInfo VName] -> InternaliseM Attrs
internaliseAttrs [AttrInfo VName]
attrs
let fd :: FunDef SOACS
fd =
forall rep.
Maybe EntryPoint
-> Attrs
-> Name
-> [RetType rep]
-> [FParam rep]
-> Body rep
-> FunDef rep
I.FunDef
forall a. Maybe a
Nothing
Attrs
attrs'
(VName -> Name
internaliseFunName VName
fname)
[TypeBase ExtShape Uniqueness]
rettype'
[Param DeclType]
all_params
Body SOACS
body'
if forall (t :: * -> *) a. Foldable t => t a -> Bool
null [[FParam SOACS]]
params'
then VName -> FunDef SOACS -> InternaliseM ()
bindConstant VName
fname FunDef SOACS
fd
else
VName -> FunDef SOACS -> FunInfo -> InternaliseM ()
bindFunction
VName
fname
FunDef SOACS
fd
( [VName]
shapenames,
forall a b. (a -> b) -> [a] -> [b]
map forall t. DeclTyped t => t -> DeclType
declTypeOf forall a b. (a -> b) -> a -> b
$ forall (t :: * -> *) a. Foldable t => t [a] -> [a]
concat [[FParam SOACS]]
params',
[Param DeclType]
all_params,
forall rt dec.
(IsRetType rt, Typed dec) =>
[rt]
-> [Param dec]
-> [(SubExp, TypeBase Shape NoUniqueness)]
-> Maybe [rt]
applyRetType [TypeBase ExtShape Uniqueness]
rettype' [Param DeclType]
all_params
)
case Maybe (Info EntryPoint)
entry of
Just (Info EntryPoint
entry') -> VisibleTypes -> EntryPoint -> ValBind -> InternaliseM ()
generateEntryPoint VisibleTypes
types EntryPoint
entry' ValBind
fb
Maybe (Info EntryPoint)
Nothing -> forall (f :: * -> *) a. Applicative f => a -> f a
pure ()
where
zeroExts :: [TypeBase ExtShape u] -> [TypeBase ExtShape u]
zeroExts [TypeBase ExtShape u]
ts = forall u.
[TypeBase ExtShape u]
-> [TypeBase ExtShape u] -> [TypeBase ExtShape u]
generaliseExtTypes [TypeBase ExtShape u]
ts [TypeBase ExtShape u]
ts
generateEntryPoint :: VisibleTypes -> E.EntryPoint -> E.ValBind -> InternaliseM ()
generateEntryPoint :: VisibleTypes -> EntryPoint -> ValBind -> InternaliseM ()
generateEntryPoint VisibleTypes
types (E.EntryPoint [EntryParam]
e_params EntryType
e_rettype) ValBind
vb = do
let (E.ValBind Maybe (Info EntryPoint)
_ VName
ofname Maybe (TypeExp VName)
_ (Info StructRetType
rettype) [TypeParamBase VName]
tparams [PatBase Info VName]
params Exp
_ Maybe DocComment
_ [AttrInfo VName]
attrs SrcLoc
loc) = ValBind
vb
forall a.
[TypeParamBase VName]
-> [PatBase Info VName]
-> ([FParam SOACS] -> [[FParam SOACS]] -> InternaliseM a)
-> InternaliseM a
bindingFParams [TypeParamBase VName]
tparams [PatBase Info VName]
params forall a b. (a -> b) -> a -> b
$ \[FParam SOACS]
shapeparams [[FParam SOACS]]
params' -> do
let entry_rettype :: [[TypeBase ExtShape Uniqueness]]
entry_rettype = StructRetType -> [[TypeBase ExtShape Uniqueness]]
internaliseEntryReturnType StructRetType
rettype
(EntryPoint
entry', OpaqueTypes
opaques) =
VisibleTypes
-> Name
-> [(EntryParam, [Param DeclType])]
-> (EntryType, [[TypeBase Rank Uniqueness]])
-> (EntryPoint, OpaqueTypes)
entryPoint
VisibleTypes
types
(VName -> Name
baseName VName
ofname)
(forall a b. [a] -> [b] -> [(a, b)]
zip [EntryParam]
e_params [[FParam SOACS]]
params')
(EntryType
e_rettype, forall a b. (a -> b) -> [a] -> [b]
map (forall a b. (a -> b) -> [a] -> [b]
map forall shape u.
ArrayShape shape =>
TypeBase shape u -> TypeBase Rank u
I.rankShaped) [[TypeBase ExtShape Uniqueness]]
entry_rettype)
args :: [SubExp]
args = forall a b. (a -> b) -> [a] -> [b]
map (VName -> SubExp
I.Var forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall dec. Param dec -> VName
I.paramName) forall a b. (a -> b) -> a -> b
$ forall (t :: * -> *) a. Foldable t => t [a] -> [a]
concat [[FParam SOACS]]
params'
OpaqueTypes -> InternaliseM ()
addOpaques OpaqueTypes
opaques
(Body SOACS
entry_body, [TypeBase ExtShape Uniqueness]
ctx_ts) <- forall (m :: * -> *) a.
MonadBuilder m =>
m (Result, a) -> m (Body (Rep m), a)
buildBody forall a b. (a -> b) -> a -> b
$ do
Maybe [SubExp]
maybe_const <- VName -> InternaliseM (Maybe [SubExp])
lookupConst VName
ofname
[SubExp]
vals <- case Maybe [SubExp]
maybe_const of
Just [SubExp]
ses ->
forall (f :: * -> *) a. Applicative f => a -> f a
pure [SubExp]
ses
Maybe [SubExp]
Nothing ->
forall a b. (a, b) -> a
fst forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> [Char]
-> QualName VName
-> [SubExp]
-> SrcLoc
-> InternaliseM ([SubExp], [ExtType])
funcall [Char]
"entry_result" (forall v. v -> QualName v
E.qualName VName
ofname) [SubExp]
args SrcLoc
loc
[SubExp]
ctx <-
forall u a. [TypeBase ExtShape u] -> [[a]] -> [a]
extractShapeContext (forall {u}. [TypeBase ExtShape u] -> [TypeBase ExtShape u]
zeroExts forall a b. (a -> b) -> a -> b
$ forall (t :: * -> *) a. Foldable t => t [a] -> [a]
concat [[TypeBase ExtShape Uniqueness]]
entry_rettype)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
(a -> m b) -> t a -> m (t b)
mapM (forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap forall u. TypeBase Shape u -> [SubExp]
I.arrayDims forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall t (m :: * -> *).
HasScope t m =>
SubExp -> m (TypeBase Shape NoUniqueness)
subExpType) [SubExp]
vals
forall (f :: * -> *) a. Applicative f => a -> f a
pure ([SubExp] -> Result
subExpsRes forall a b. (a -> b) -> a -> b
$ [SubExp]
ctx forall a. [a] -> [a] -> [a]
++ [SubExp]
vals, forall a b. (a -> b) -> [a] -> [b]
map (forall a b. a -> b -> a
const (forall shape u. PrimType -> TypeBase shape u
I.Prim PrimType
int64)) [SubExp]
ctx)
Attrs
attrs' <- [AttrInfo VName] -> InternaliseM Attrs
internaliseAttrs [AttrInfo VName]
attrs
FunDef SOACS -> InternaliseM ()
addFunDef forall a b. (a -> b) -> a -> b
$
forall rep.
Maybe EntryPoint
-> Attrs
-> Name
-> [RetType rep]
-> [FParam rep]
-> Body rep
-> FunDef rep
I.FunDef
(forall a. a -> Maybe a
Just EntryPoint
entry')
Attrs
attrs'
(Name
"entry_" forall a. Semigroup a => a -> a -> a
<> VName -> Name
baseName VName
ofname)
([TypeBase ExtShape Uniqueness]
ctx_ts forall a. [a] -> [a] -> [a]
++ forall {u}. [TypeBase ExtShape u] -> [TypeBase ExtShape u]
zeroExts (forall (t :: * -> *) a. Foldable t => t [a] -> [a]
concat [[TypeBase ExtShape Uniqueness]]
entry_rettype))
([FParam SOACS]
shapeparams forall a. [a] -> [a] -> [a]
++ forall (t :: * -> *) a. Foldable t => t [a] -> [a]
concat [[FParam SOACS]]
params')
Body SOACS
entry_body
where
zeroExts :: [TypeBase ExtShape u] -> [TypeBase ExtShape u]
zeroExts [TypeBase ExtShape u]
ts = forall u.
[TypeBase ExtShape u]
-> [TypeBase ExtShape u] -> [TypeBase ExtShape u]
generaliseExtTypes [TypeBase ExtShape u]
ts [TypeBase ExtShape u]
ts
internaliseBody :: String -> E.Exp -> InternaliseM (Body SOACS)
internaliseBody :: [Char] -> Exp -> InternaliseM (Body SOACS)
internaliseBody [Char]
desc Exp
e =
forall (m :: * -> *).
MonadBuilder m =>
m Result -> m (Body (Rep m))
buildBody_ forall a b. (a -> b) -> a -> b
$ [SubExp] -> Result
subExpsRes forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> [Char] -> Exp -> InternaliseM [SubExp]
internaliseExp ([Char]
desc forall a. Semigroup a => a -> a -> a
<> [Char]
"_res") Exp
e
bodyFromStms ::
InternaliseM (Result, a) ->
InternaliseM (Body SOACS, a)
bodyFromStms :: forall a. InternaliseM (Result, a) -> InternaliseM (Body SOACS, a)
bodyFromStms InternaliseM (Result, a)
m = do
((Result
res, a
a), Stms SOACS
stms) <- forall (m :: * -> *) a.
MonadBuilder m =>
m a -> m (a, Stms (Rep m))
collectStms InternaliseM (Result, a)
m
(,a
a) forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> forall (m :: * -> *).
MonadBuilder m =>
Stms (Rep m) -> Result -> m (Body (Rep m))
mkBodyM Stms SOACS
stms Result
res
letValExp :: String -> I.Exp SOACS -> InternaliseM [VName]
letValExp :: [Char] -> Exp SOACS -> InternaliseM [VName]
letValExp [Char]
name Exp SOACS
e = do
[ExtType]
e_t <- forall rep (m :: * -> *).
(HasScope rep m, TypedOp (Op rep)) =>
Exp rep -> m [ExtType]
expExtType Exp SOACS
e
[VName]
names <- forall (m :: * -> *) a. Applicative m => Int -> m a -> m [a]
replicateM (forall (t :: * -> *) a. Foldable t => t a -> Int
length [ExtType]
e_t) forall a b. (a -> b) -> a -> b
$ forall (m :: * -> *). MonadFreshNames m => [Char] -> m VName
newVName [Char]
name
forall (m :: * -> *).
MonadBuilder m =>
[VName] -> Exp (Rep m) -> m ()
letBindNames [VName]
names Exp SOACS
e
let ctx :: Set Int
ctx = forall u. [TypeBase ExtShape u] -> Set Int
shapeContext [ExtType]
e_t
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 a b. (a, b) -> a
fst forall a b. (a -> b) -> a -> b
$ forall a. (a -> Bool) -> [a] -> [a]
filter ((forall a. Ord a => a -> Set a -> Bool
`S.notMember` Set Int
ctx) forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall a b. (a, b) -> b
snd) forall a b. (a -> b) -> a -> b
$ forall a b. [a] -> [b] -> [(a, b)]
zip [VName]
names [Int
0 ..]
letValExp' :: String -> I.Exp SOACS -> InternaliseM [SubExp]
letValExp' :: [Char] -> Exp SOACS -> InternaliseM [SubExp]
letValExp' [Char]
_ (BasicOp (SubExp SubExp
se)) = forall (f :: * -> *) a. Applicative f => a -> f a
pure [SubExp
se]
letValExp' [Char]
name Exp SOACS
ses = forall a b. (a -> b) -> [a] -> [b]
map VName -> SubExp
I.Var forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> [Char] -> Exp SOACS -> InternaliseM [VName]
letValExp [Char]
name Exp SOACS
ses
internaliseAppExp :: String -> E.AppRes -> E.AppExp -> InternaliseM [I.SubExp]
internaliseAppExp :: [Char] -> AppRes -> AppExp -> InternaliseM [SubExp]
internaliseAppExp [Char]
desc AppRes
_ (E.Index Exp
e SliceBase Info VName
idxs SrcLoc
loc) = do
[VName]
vs <- [Char] -> Exp -> InternaliseM [VName]
internaliseExpToVars [Char]
"indexed" Exp
e
[SubExp]
dims <- case [VName]
vs of
[] -> forall (f :: * -> *) a. Applicative f => a -> f a
pure []
VName
v : [VName]
_ -> forall u. TypeBase Shape u -> [SubExp]
I.arrayDims forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> forall rep (m :: * -> *).
HasScope rep m =>
VName -> m (TypeBase Shape NoUniqueness)
lookupType VName
v
([DimIndex SubExp]
idxs', Certs
cs) <- SrcLoc
-> [SubExp]
-> SliceBase Info VName
-> InternaliseM ([DimIndex SubExp], Certs)
internaliseSlice SrcLoc
loc [SubExp]
dims SliceBase Info VName
idxs
let index :: VName -> InternaliseM (Exp SOACS)
index VName
v = do
TypeBase Shape NoUniqueness
v_t <- forall rep (m :: * -> *).
HasScope rep m =>
VName -> m (TypeBase Shape NoUniqueness)
lookupType VName
v
forall (f :: * -> *) a. Applicative f => a -> f a
pure forall a b. (a -> b) -> a -> b
$ forall rep. BasicOp -> Exp rep
I.BasicOp forall a b. (a -> b) -> a -> b
$ VName -> Slice SubExp -> BasicOp
I.Index VName
v forall a b. (a -> b) -> a -> b
$ TypeBase Shape NoUniqueness -> [DimIndex SubExp] -> Slice SubExp
fullSlice TypeBase Shape NoUniqueness
v_t [DimIndex SubExp]
idxs'
forall (m :: * -> *) a. MonadBuilder m => Certs -> m a -> m a
certifying Certs
cs forall a b. (a -> b) -> a -> b
$ forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
(a -> m b) -> t a -> m (t b)
mapM (forall (m :: * -> *).
MonadBuilder m =>
[Char] -> Exp (Rep m) -> m SubExp
letSubExp [Char]
desc forall (m :: * -> *) b c a.
Monad m =>
(b -> m c) -> (a -> m b) -> a -> m c
<=< VName -> InternaliseM (Exp SOACS)
index) [VName]
vs
internaliseAppExp [Char]
desc AppRes
_ (E.Range Exp
start Maybe Exp
maybe_second Inclusiveness Exp
end SrcLoc
loc) = do
SubExp
start' <- [Char] -> Exp -> InternaliseM SubExp
internaliseExp1 [Char]
"range_start" Exp
start
SubExp
end' <- [Char] -> Exp -> InternaliseM SubExp
internaliseExp1 [Char]
"range_end" forall a b. (a -> b) -> a -> b
$ case Inclusiveness Exp
end of
DownToExclusive Exp
e -> Exp
e
ToInclusive Exp
e -> Exp
e
UpToExclusive Exp
e -> Exp
e
Maybe SubExp
maybe_second' <-
forall (t :: * -> *) (f :: * -> *) a b.
(Traversable t, Applicative f) =>
(a -> f b) -> t a -> f (t b)
traverse ([Char] -> Exp -> InternaliseM SubExp
internaliseExp1 [Char]
"range_second") Maybe Exp
maybe_second
let conv :: SubExp -> InternaliseM SubExp
conv = case Exp -> PatType
E.typeOf Exp
start of
E.Scalar (E.Prim (E.Unsigned IntType
_)) -> forall (m :: * -> *).
MonadBuilder m =>
IntType -> SubExp -> m SubExp
asIntZ IntType
Int64
PatType
_ -> forall (m :: * -> *).
MonadBuilder m =>
IntType -> SubExp -> m SubExp
asIntS IntType
Int64
SubExp
start'_i64 <- SubExp -> InternaliseM SubExp
conv SubExp
start'
SubExp
end'_i64 <- SubExp -> InternaliseM SubExp
conv SubExp
end'
Maybe SubExp
maybe_second'_i64 <- forall (t :: * -> *) (f :: * -> *) a b.
(Traversable t, Applicative f) =>
(a -> f b) -> t a -> f (t b)
traverse SubExp -> InternaliseM SubExp
conv Maybe SubExp
maybe_second'
let errmsg :: ErrorMsg SubExp
errmsg =
forall a. [ErrorMsgPart a] -> ErrorMsg a
errorMsg forall a b. (a -> b) -> a -> b
$
[ErrorMsgPart SubExp
"Range "]
forall a. [a] -> [a] -> [a]
++ [forall a. PrimType -> a -> ErrorMsgPart a
ErrorVal PrimType
int64 SubExp
start'_i64]
forall a. [a] -> [a] -> [a]
++ ( case Maybe SubExp
maybe_second'_i64 of
Maybe SubExp
Nothing -> []
Just SubExp
second_i64 -> [ErrorMsgPart SubExp
"..", forall a. PrimType -> a -> ErrorMsgPart a
ErrorVal PrimType
int64 SubExp
second_i64]
)
forall a. [a] -> [a] -> [a]
++ ( case Inclusiveness Exp
end of
DownToExclusive {} -> [ErrorMsgPart SubExp
"..>"]
ToInclusive {} -> [ErrorMsgPart SubExp
"..."]
UpToExclusive {} -> [ErrorMsgPart SubExp
"..<"]
)
forall a. [a] -> [a] -> [a]
++ [forall a. PrimType -> a -> ErrorMsgPart a
ErrorVal PrimType
int64 SubExp
end'_i64, ErrorMsgPart SubExp
" is invalid."]
(IntType
it, CmpOp
le_op, CmpOp
lt_op) <-
case Exp -> PatType
E.typeOf Exp
start of
E.Scalar (E.Prim (E.Signed IntType
it)) -> forall (f :: * -> *) a. Applicative f => a -> f a
pure (IntType
it, IntType -> CmpOp
CmpSle IntType
it, IntType -> CmpOp
CmpSlt IntType
it)
E.Scalar (E.Prim (E.Unsigned IntType
it)) -> forall (f :: * -> *) a. Applicative f => a -> f a
pure (IntType
it, IntType -> CmpOp
CmpUle IntType
it, IntType -> CmpOp
CmpUlt IntType
it)
PatType
start_t -> forall a. HasCallStack => [Char] -> a
error forall a b. (a -> b) -> a -> b
$ [Char]
"Start value in range has type " forall a. [a] -> [a] -> [a]
++ forall a. Pretty a => a -> [Char]
prettyString PatType
start_t
let one :: SubExp
one = IntType -> Integer -> SubExp
intConst IntType
it Integer
1
negone :: SubExp
negone = IntType -> Integer -> SubExp
intConst IntType
it (-Integer
1)
default_step :: SubExp
default_step = case Inclusiveness Exp
end of
DownToExclusive {} -> SubExp
negone
ToInclusive {} -> SubExp
one
UpToExclusive {} -> SubExp
one
(SubExp
step, SubExp
step_zero) <- case Maybe SubExp
maybe_second' of
Just SubExp
second' -> do
SubExp
subtracted_step <-
forall (m :: * -> *).
MonadBuilder m =>
[Char] -> Exp (Rep m) -> m SubExp
letSubExp [Char]
"subtracted_step" forall a b. (a -> b) -> a -> b
$
forall rep. BasicOp -> Exp rep
I.BasicOp forall a b. (a -> b) -> a -> b
$
BinOp -> SubExp -> SubExp -> BasicOp
I.BinOp (IntType -> Overflow -> BinOp
I.Sub IntType
it Overflow
I.OverflowWrap) SubExp
second' SubExp
start'
SubExp
step_zero <- forall (m :: * -> *).
MonadBuilder m =>
[Char] -> Exp (Rep m) -> m SubExp
letSubExp [Char]
"step_zero" forall a b. (a -> b) -> a -> b
$ forall rep. BasicOp -> Exp rep
I.BasicOp forall a b. (a -> b) -> a -> b
$ CmpOp -> SubExp -> SubExp -> BasicOp
I.CmpOp (PrimType -> CmpOp
I.CmpEq forall a b. (a -> b) -> a -> b
$ IntType -> PrimType
IntType IntType
it) SubExp
start' SubExp
second'
forall (f :: * -> *) a. Applicative f => a -> f a
pure (SubExp
subtracted_step, SubExp
step_zero)
Maybe SubExp
Nothing ->
forall (f :: * -> *) a. Applicative f => a -> f a
pure (SubExp
default_step, forall v. IsValue v => v -> SubExp
constant Bool
False)
SubExp
step_sign <- forall (m :: * -> *).
MonadBuilder m =>
[Char] -> Exp (Rep m) -> m SubExp
letSubExp [Char]
"s_sign" forall a b. (a -> b) -> a -> b
$ forall rep. BasicOp -> Exp rep
BasicOp forall a b. (a -> b) -> a -> b
$ UnOp -> SubExp -> BasicOp
I.UnOp (IntType -> UnOp
I.SSignum IntType
it) SubExp
step
SubExp
step_sign_i64 <- forall (m :: * -> *).
MonadBuilder m =>
IntType -> SubExp -> m SubExp
asIntS IntType
Int64 SubExp
step_sign
SubExp
bounds_invalid_downwards <-
forall (m :: * -> *).
MonadBuilder m =>
[Char] -> Exp (Rep m) -> m SubExp
letSubExp [Char]
"bounds_invalid_downwards" forall a b. (a -> b) -> a -> b
$
forall rep. BasicOp -> Exp rep
I.BasicOp forall a b. (a -> b) -> a -> b
$
CmpOp -> SubExp -> SubExp -> BasicOp
I.CmpOp CmpOp
le_op SubExp
start' SubExp
end'
SubExp
bounds_invalid_upwards <-
forall (m :: * -> *).
MonadBuilder m =>
[Char] -> Exp (Rep m) -> m SubExp
letSubExp [Char]
"bounds_invalid_upwards" forall a b. (a -> b) -> a -> b
$
forall rep. BasicOp -> Exp rep
I.BasicOp forall a b. (a -> b) -> a -> b
$
CmpOp -> SubExp -> SubExp -> BasicOp
I.CmpOp CmpOp
lt_op SubExp
end' SubExp
start'
(SubExp
distance, SubExp
step_wrong_dir, SubExp
bounds_invalid) <- case Inclusiveness Exp
end of
DownToExclusive {} -> do
SubExp
step_wrong_dir <-
forall (m :: * -> *).
MonadBuilder m =>
[Char] -> Exp (Rep m) -> m SubExp
letSubExp [Char]
"step_wrong_dir" forall a b. (a -> b) -> a -> b
$
forall rep. BasicOp -> Exp rep
I.BasicOp forall a b. (a -> b) -> a -> b
$
CmpOp -> SubExp -> SubExp -> BasicOp
I.CmpOp (PrimType -> CmpOp
I.CmpEq forall a b. (a -> b) -> a -> b
$ IntType -> PrimType
IntType IntType
it) SubExp
step_sign SubExp
one
SubExp
distance <-
forall (m :: * -> *).
MonadBuilder m =>
[Char] -> Exp (Rep m) -> m SubExp
letSubExp [Char]
"distance" forall a b. (a -> b) -> a -> b
$
forall rep. BasicOp -> Exp rep
I.BasicOp forall a b. (a -> b) -> a -> b
$
BinOp -> SubExp -> SubExp -> BasicOp
I.BinOp (IntType -> Overflow -> BinOp
Sub IntType
it Overflow
I.OverflowWrap) SubExp
start' SubExp
end'
SubExp
distance_i64 <- forall (m :: * -> *).
MonadBuilder m =>
IntType -> SubExp -> m SubExp
asIntS IntType
Int64 SubExp
distance
forall (f :: * -> *) a. Applicative f => a -> f a
pure (SubExp
distance_i64, SubExp
step_wrong_dir, SubExp
bounds_invalid_downwards)
UpToExclusive {} -> do
SubExp
step_wrong_dir <-
forall (m :: * -> *).
MonadBuilder m =>
[Char] -> Exp (Rep m) -> m SubExp
letSubExp [Char]
"step_wrong_dir" forall a b. (a -> b) -> a -> b
$
forall rep. BasicOp -> Exp rep
I.BasicOp forall a b. (a -> b) -> a -> b
$
CmpOp -> SubExp -> SubExp -> BasicOp
I.CmpOp (PrimType -> CmpOp
I.CmpEq forall a b. (a -> b) -> a -> b
$ IntType -> PrimType
IntType IntType
it) SubExp
step_sign SubExp
negone
SubExp
distance <- forall (m :: * -> *).
MonadBuilder m =>
[Char] -> Exp (Rep m) -> m SubExp
letSubExp [Char]
"distance" forall a b. (a -> b) -> a -> b
$ forall rep. BasicOp -> Exp rep
I.BasicOp forall a b. (a -> b) -> a -> b
$ BinOp -> SubExp -> SubExp -> BasicOp
I.BinOp (IntType -> Overflow -> BinOp
Sub IntType
it Overflow
I.OverflowWrap) SubExp
end' SubExp
start'
SubExp
distance_i64 <- forall (m :: * -> *).
MonadBuilder m =>
IntType -> SubExp -> m SubExp
asIntS IntType
Int64 SubExp
distance
forall (f :: * -> *) a. Applicative f => a -> f a
pure (SubExp
distance_i64, SubExp
step_wrong_dir, SubExp
bounds_invalid_upwards)
ToInclusive {} -> do
SubExp
downwards <-
forall (m :: * -> *).
MonadBuilder m =>
[Char] -> Exp (Rep m) -> m SubExp
letSubExp [Char]
"downwards" forall a b. (a -> b) -> a -> b
$
forall rep. BasicOp -> Exp rep
I.BasicOp forall a b. (a -> b) -> a -> b
$
CmpOp -> SubExp -> SubExp -> BasicOp
I.CmpOp (PrimType -> CmpOp
I.CmpEq forall a b. (a -> b) -> a -> b
$ IntType -> PrimType
IntType IntType
it) SubExp
step_sign SubExp
negone
SubExp
distance_downwards_exclusive <-
forall (m :: * -> *).
MonadBuilder m =>
[Char] -> Exp (Rep m) -> m SubExp
letSubExp [Char]
"distance_downwards_exclusive" forall a b. (a -> b) -> a -> b
$
forall rep. BasicOp -> Exp rep
I.BasicOp forall a b. (a -> b) -> a -> b
$
BinOp -> SubExp -> SubExp -> BasicOp
I.BinOp (IntType -> Overflow -> BinOp
Sub IntType
it Overflow
I.OverflowWrap) SubExp
start' SubExp
end'
SubExp
distance_upwards_exclusive <-
forall (m :: * -> *).
MonadBuilder m =>
[Char] -> Exp (Rep m) -> m SubExp
letSubExp [Char]
"distance_upwards_exclusive" forall a b. (a -> b) -> a -> b
$
forall rep. BasicOp -> Exp rep
I.BasicOp forall a b. (a -> b) -> a -> b
$
BinOp -> SubExp -> SubExp -> BasicOp
I.BinOp (IntType -> Overflow -> BinOp
Sub IntType
it Overflow
I.OverflowWrap) SubExp
end' SubExp
start'
SubExp
bounds_invalid <-
forall (m :: * -> *).
MonadBuilder m =>
[Char] -> Exp (Rep m) -> m SubExp
letSubExp [Char]
"bounds_invalid"
forall (m :: * -> *) a b. Monad m => (a -> m b) -> m a -> m b
=<< forall (m :: * -> *).
(MonadBuilder m, BranchType (Rep m) ~ ExtType) =>
m (Exp (Rep m))
-> m (Body (Rep m)) -> m (Body (Rep m)) -> m (Exp (Rep m))
eIf
(forall (m :: * -> *). MonadBuilder m => SubExp -> m (Exp (Rep m))
eSubExp SubExp
downwards)
(forall (m :: * -> *).
MonadBuilder m =>
[SubExp] -> m (Body (Rep m))
resultBodyM [SubExp
bounds_invalid_downwards])
(forall (m :: * -> *).
MonadBuilder m =>
[SubExp] -> m (Body (Rep m))
resultBodyM [SubExp
bounds_invalid_upwards])
SubExp
distance_exclusive <-
forall (m :: * -> *).
MonadBuilder m =>
[Char] -> Exp (Rep m) -> m SubExp
letSubExp [Char]
"distance_exclusive"
forall (m :: * -> *) a b. Monad m => (a -> m b) -> m a -> m b
=<< forall (m :: * -> *).
(MonadBuilder m, BranchType (Rep m) ~ ExtType) =>
m (Exp (Rep m))
-> m (Body (Rep m)) -> m (Body (Rep m)) -> m (Exp (Rep m))
eIf
(forall (m :: * -> *). MonadBuilder m => SubExp -> m (Exp (Rep m))
eSubExp SubExp
downwards)
(forall (m :: * -> *).
MonadBuilder m =>
[SubExp] -> m (Body (Rep m))
resultBodyM [SubExp
distance_downwards_exclusive])
(forall (m :: * -> *).
MonadBuilder m =>
[SubExp] -> m (Body (Rep m))
resultBodyM [SubExp
distance_upwards_exclusive])
SubExp
distance_exclusive_i64 <- forall (m :: * -> *).
MonadBuilder m =>
IntType -> SubExp -> m SubExp
asIntS IntType
Int64 SubExp
distance_exclusive
SubExp
distance <-
forall (m :: * -> *).
MonadBuilder m =>
[Char] -> Exp (Rep m) -> m SubExp
letSubExp [Char]
"distance" forall a b. (a -> b) -> a -> b
$
forall rep. BasicOp -> Exp rep
I.BasicOp forall a b. (a -> b) -> a -> b
$
BinOp -> SubExp -> SubExp -> BasicOp
I.BinOp
(IntType -> Overflow -> BinOp
Add IntType
Int64 Overflow
I.OverflowWrap)
SubExp
distance_exclusive_i64
(IntType -> Integer -> SubExp
intConst IntType
Int64 Integer
1)
forall (f :: * -> *) a. Applicative f => a -> f a
pure (SubExp
distance, forall v. IsValue v => v -> SubExp
constant Bool
False, SubExp
bounds_invalid)
SubExp
step_invalid <-
forall (m :: * -> *).
MonadBuilder m =>
[Char] -> Exp (Rep m) -> m SubExp
letSubExp [Char]
"step_invalid" forall a b. (a -> b) -> a -> b
$
forall rep. BasicOp -> Exp rep
I.BasicOp forall a b. (a -> b) -> a -> b
$
BinOp -> SubExp -> SubExp -> BasicOp
I.BinOp BinOp
I.LogOr SubExp
step_wrong_dir SubExp
step_zero
SubExp
invalid <-
forall (m :: * -> *).
MonadBuilder m =>
[Char] -> Exp (Rep m) -> m SubExp
letSubExp [Char]
"range_invalid" forall a b. (a -> b) -> a -> b
$
forall rep. BasicOp -> Exp rep
I.BasicOp forall a b. (a -> b) -> a -> b
$
BinOp -> SubExp -> SubExp -> BasicOp
I.BinOp BinOp
I.LogOr SubExp
step_invalid SubExp
bounds_invalid
SubExp
valid <- forall (m :: * -> *).
MonadBuilder m =>
[Char] -> Exp (Rep m) -> m SubExp
letSubExp [Char]
"valid" forall a b. (a -> b) -> a -> b
$ forall rep. BasicOp -> Exp rep
I.BasicOp forall a b. (a -> b) -> a -> b
$ UnOp -> SubExp -> BasicOp
I.UnOp UnOp
I.Not SubExp
invalid
Certs
cs <- [Char] -> SubExp -> ErrorMsg SubExp -> SrcLoc -> InternaliseM Certs
assert [Char]
"range_valid_c" SubExp
valid ErrorMsg SubExp
errmsg SrcLoc
loc
SubExp
step_i64 <- forall (m :: * -> *).
MonadBuilder m =>
IntType -> SubExp -> m SubExp
asIntS IntType
Int64 SubExp
step
SubExp
pos_step <-
forall (m :: * -> *).
MonadBuilder m =>
[Char] -> Exp (Rep m) -> m SubExp
letSubExp [Char]
"pos_step" forall a b. (a -> b) -> a -> b
$
forall rep. BasicOp -> Exp rep
I.BasicOp forall a b. (a -> b) -> a -> b
$
BinOp -> SubExp -> SubExp -> BasicOp
I.BinOp (IntType -> Overflow -> BinOp
Mul IntType
Int64 Overflow
I.OverflowWrap) SubExp
step_i64 SubExp
step_sign_i64
SubExp
num_elems <-
forall (m :: * -> *) a. MonadBuilder m => Certs -> m a -> m a
certifying Certs
cs forall a b. (a -> b) -> a -> b
$
forall (m :: * -> *).
MonadBuilder m =>
[Char] -> Exp (Rep m) -> m SubExp
letSubExp [Char]
"num_elems" forall a b. (a -> b) -> a -> b
$
forall rep. BasicOp -> Exp rep
I.BasicOp forall a b. (a -> b) -> a -> b
$
BinOp -> SubExp -> SubExp -> BasicOp
I.BinOp (IntType -> Safety -> BinOp
SDivUp IntType
Int64 Safety
I.Unsafe) SubExp
distance SubExp
pos_step
SubExp
se <- forall (m :: * -> *).
MonadBuilder m =>
[Char] -> Exp (Rep m) -> m SubExp
letSubExp [Char]
desc (forall rep. BasicOp -> Exp rep
I.BasicOp forall a b. (a -> b) -> a -> b
$ SubExp -> SubExp -> SubExp -> IntType -> BasicOp
I.Iota SubExp
num_elems SubExp
start' SubExp
step IntType
it)
forall (f :: * -> *) a. Applicative f => a -> f a
pure [SubExp
se]
internaliseAppExp [Char]
desc (E.AppRes PatType
et [VName]
ext) (E.Coerce Exp
e TypeExp VName
dt SrcLoc
loc) = do
[SubExp]
ses <- [Char] -> Exp -> InternaliseM [SubExp]
internaliseExp [Char]
desc Exp
e
[TypeBase ExtShape Uniqueness]
ts <- forall shape u.
StructRetType
-> [TypeBase shape u] -> [TypeBase ExtShape Uniqueness]
internaliseReturnType (forall dim as. [VName] -> TypeBase dim as -> RetTypeBase dim as
E.RetType [VName]
ext (forall dim as. TypeBase dim as -> TypeBase dim ()
E.toStruct PatType
et)) forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
(a -> m b) -> t a -> m (t b)
mapM forall t (m :: * -> *).
HasScope t m =>
SubExp -> m (TypeBase Shape NoUniqueness)
subExpType [SubExp]
ses
[ErrorMsgPart SubExp]
dt' <- TypeExp VName -> InternaliseM [ErrorMsgPart SubExp]
typeExpForError TypeExp VName
dt
forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
t a -> (a -> m b) -> m (t b)
forM (forall a b. [a] -> [b] -> [(a, b)]
zip [SubExp]
ses [TypeBase ExtShape Uniqueness]
ts) forall a b. (a -> b) -> a -> b
$ \(SubExp
e', TypeBase ExtShape Uniqueness
t') -> do
[SubExp]
dims <- forall u. TypeBase Shape u -> [SubExp]
arrayDims forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> forall t (m :: * -> *).
HasScope t m =>
SubExp -> m (TypeBase Shape NoUniqueness)
subExpType SubExp
e'
let parts :: [ErrorMsgPart SubExp]
parts =
[ErrorMsgPart SubExp
"Value of (core language) shape ("]
forall a. [a] -> [a] -> [a]
++ forall a. a -> [a] -> [a]
intersperse ErrorMsgPart SubExp
", " (forall a b. (a -> b) -> [a] -> [b]
map (forall a. PrimType -> a -> ErrorMsgPart a
ErrorVal PrimType
int64) [SubExp]
dims)
forall a. [a] -> [a] -> [a]
++ [ErrorMsgPart SubExp
") cannot match shape of type `"]
forall a. [a] -> [a] -> [a]
++ [ErrorMsgPart SubExp]
dt'
forall a. [a] -> [a] -> [a]
++ [ErrorMsgPart SubExp
"`."]
ErrorMsg SubExp
-> SrcLoc -> ExtType -> [Char] -> SubExp -> InternaliseM SubExp
ensureExtShape (forall a. [ErrorMsgPart a] -> ErrorMsg a
errorMsg [ErrorMsgPart SubExp]
parts) SrcLoc
loc (forall shape.
TypeBase shape Uniqueness -> TypeBase shape NoUniqueness
I.fromDecl TypeBase ExtShape Uniqueness
t') [Char]
desc SubExp
e'
internaliseAppExp [Char]
desc AppRes
_ e :: AppExp
e@E.Apply {} =
case AppExp -> (Function, [(Exp, Maybe VName)])
findFuncall AppExp
e of
(FunctionHole PatType
t SrcLoc
loc, [(Exp, Maybe VName)]
_args) -> do
[Char] -> Exp -> InternaliseM [SubExp]
internaliseExp [Char]
desc (forall (f :: * -> *) vn. f PatType -> SrcLoc -> ExpBase f vn
E.Hole (forall a. a -> Info a
Info (forall dim as. TypeBase dim as -> TypeBase dim (Set Alias)
fromStruct forall a b. (a -> b) -> a -> b
$ forall a b. (a, b) -> b
snd forall a b. (a -> b) -> a -> b
$ forall dim as.
TypeBase dim as -> ([TypeBase dim ()], TypeBase dim ())
E.unfoldFunType PatType
t)) SrcLoc
loc)
(FunctionName QualName VName
qfname, [(Exp, Maybe VName)]
args) -> do
let fname :: Name
fname = [Char] -> Name
nameFromString forall a b. (a -> b) -> a -> b
$ forall a. Pretty a => a -> [Char]
prettyString forall a b. (a -> b) -> a -> b
$ VName -> Name
baseName forall a b. (a -> b) -> a -> b
$ forall vn. QualName vn -> vn
qualLeaf QualName VName
qfname
loc :: SrcLoc
loc = forall a. Located a => a -> SrcLoc
srclocOf AppExp
e
arg_desc :: [Char]
arg_desc = Name -> [Char]
nameToString Name
fname forall a. [a] -> [a] -> [a]
++ [Char]
"_arg"
case () of
()
| VName -> Int
baseTag (forall vn. QualName vn -> vn
qualLeaf QualName VName
qfname) forall a. Ord a => a -> a -> Bool
<= Int
maxIntrinsicTag,
VName -> [Char]
baseString (forall vn. QualName vn -> vn
qualLeaf QualName VName
qfname) forall a. Eq a => a -> a -> Bool
== [Char]
"&&",
[(Exp
x, Maybe VName
_), (Exp
y, Maybe VName
_)] <- [(Exp, Maybe VName)]
args ->
[Char] -> Exp -> InternaliseM [SubExp]
internaliseExp [Char]
desc forall a b. (a -> b) -> a -> b
$
forall (f :: * -> *) vn.
AppExpBase f vn -> f AppRes -> ExpBase f vn
E.AppExp
(forall (f :: * -> *) vn.
ExpBase f vn
-> ExpBase f vn -> ExpBase f vn -> SrcLoc -> AppExpBase f vn
E.If Exp
x Exp
y (forall (f :: * -> *) vn. PrimValue -> SrcLoc -> ExpBase f vn
E.Literal (Bool -> PrimValue
E.BoolValue Bool
False) forall a. Monoid a => a
mempty) forall a. Monoid a => a
mempty)
(forall a. a -> Info a
Info forall a b. (a -> b) -> a -> b
$ PatType -> [VName] -> AppRes
AppRes (forall dim as. ScalarTypeBase dim as -> TypeBase dim as
E.Scalar forall a b. (a -> b) -> a -> b
$ forall dim as. PrimType -> ScalarTypeBase dim as
E.Prim PrimType
E.Bool) [])
| VName -> Int
baseTag (forall vn. QualName vn -> vn
qualLeaf QualName VName
qfname) forall a. Ord a => a -> a -> Bool
<= Int
maxIntrinsicTag,
VName -> [Char]
baseString (forall vn. QualName vn -> vn
qualLeaf QualName VName
qfname) forall a. Eq a => a -> a -> Bool
== [Char]
"||",
[(Exp
x, Maybe VName
_), (Exp
y, Maybe VName
_)] <- [(Exp, Maybe VName)]
args ->
[Char] -> Exp -> InternaliseM [SubExp]
internaliseExp [Char]
desc forall a b. (a -> b) -> a -> b
$
forall (f :: * -> *) vn.
AppExpBase f vn -> f AppRes -> ExpBase f vn
E.AppExp
(forall (f :: * -> *) vn.
ExpBase f vn
-> ExpBase f vn -> ExpBase f vn -> SrcLoc -> AppExpBase f vn
E.If Exp
x (forall (f :: * -> *) vn. PrimValue -> SrcLoc -> ExpBase f vn
E.Literal (Bool -> PrimValue
E.BoolValue Bool
True) forall a. Monoid a => a
mempty) Exp
y forall a. Monoid a => a
mempty)
(forall a. a -> Info a
Info forall a b. (a -> b) -> a -> b
$ PatType -> [VName] -> AppRes
AppRes (forall dim as. ScalarTypeBase dim as -> TypeBase dim as
E.Scalar forall a b. (a -> b) -> a -> b
$ forall dim as. PrimType -> ScalarTypeBase dim as
E.Prim PrimType
E.Bool) [])
| Just [(TypeBase Size (), [SubExp])] -> InternaliseM [SubExp]
internalise <- QualName VName
-> [Char]
-> SrcLoc
-> Maybe ([(TypeBase Size (), [SubExp])] -> InternaliseM [SubExp])
isOverloadedFunction QualName VName
qfname [Char]
desc SrcLoc
loc -> do
let prepareArg :: (Exp, b) -> InternaliseM (TypeBase Size (), [SubExp])
prepareArg (Exp
arg, b
_) =
(forall dim as. TypeBase dim as -> TypeBase dim ()
E.toStruct (Exp -> PatType
E.typeOf Exp
arg),) forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> [Char] -> Exp -> InternaliseM [SubExp]
internaliseExp [Char]
"arg" Exp
arg
[(TypeBase Size (), [SubExp])] -> InternaliseM [SubExp]
internalise forall (m :: * -> *) a b. Monad m => (a -> m b) -> m a -> m b
=<< forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
(a -> m b) -> t a -> m (t b)
mapM forall {b}. (Exp, b) -> InternaliseM (TypeBase Size (), [SubExp])
prepareArg [(Exp, Maybe VName)]
args
| Just [Char] -> InternaliseM [SubExp]
internalise <- QualName VName
-> [Exp] -> SrcLoc -> Maybe ([Char] -> InternaliseM [SubExp])
isIntrinsicFunction QualName VName
qfname (forall a b. (a -> b) -> [a] -> [b]
map forall a b. (a, b) -> a
fst [(Exp, Maybe VName)]
args) SrcLoc
loc ->
[Char] -> InternaliseM [SubExp]
internalise [Char]
desc
| VName -> Int
baseTag (forall vn. QualName vn -> vn
qualLeaf QualName VName
qfname) forall a. Ord a => a -> a -> Bool
<= Int
maxIntrinsicTag,
Just (PrimType
rettype, [PrimType]
_) <- forall k a. Ord k => k -> Map k a -> Maybe a
M.lookup Name
fname Map Name (PrimType, [PrimType])
I.builtInFunctions -> do
let tag :: [a] -> [(a, Diet)]
tag [a]
ses = [(a
se, Diet
I.Observe) | a
se <- [a]
ses]
[[SubExp]]
args' <- forall a. [a] -> [a]
reverse forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
(a -> m b) -> t a -> m (t b)
mapM ([Char] -> (Exp, Maybe VName) -> InternaliseM [SubExp]
internaliseArg [Char]
arg_desc) (forall a. [a] -> [a]
reverse [(Exp, Maybe VName)]
args)
let args'' :: [(SubExp, Diet)]
args'' = forall (t :: * -> *) a b. Foldable t => (a -> [b]) -> t a -> [b]
concatMap forall {a}. [a] -> [(a, Diet)]
tag [[SubExp]]
args'
[Char] -> Exp SOACS -> InternaliseM [SubExp]
letValExp' [Char]
desc forall a b. (a -> b) -> a -> b
$ forall rep.
Name
-> [(SubExp, Diet)]
-> [RetType rep]
-> (Safety, SrcLoc, [SrcLoc])
-> Exp rep
I.Apply Name
fname [(SubExp, Diet)]
args'' [forall shape u. PrimType -> TypeBase shape u
I.Prim PrimType
rettype] (Safety
Safe, SrcLoc
loc, [])
| Bool
otherwise -> do
[SubExp]
args' <- forall (t :: * -> *) a. Foldable t => t [a] -> [a]
concat forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall a. [a] -> [a]
reverse forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
(a -> m b) -> t a -> m (t b)
mapM ([Char] -> (Exp, Maybe VName) -> InternaliseM [SubExp]
internaliseArg [Char]
arg_desc) (forall a. [a] -> [a]
reverse [(Exp, Maybe VName)]
args)
forall a b. (a, b) -> a
fst forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> [Char]
-> QualName VName
-> [SubExp]
-> SrcLoc
-> InternaliseM ([SubExp], [ExtType])
funcall [Char]
desc QualName VName
qfname [SubExp]
args' SrcLoc
loc
internaliseAppExp [Char]
desc AppRes
_ (E.LetPat [SizeBinder VName]
sizes PatBase Info VName
pat Exp
e Exp
body SrcLoc
_) =
forall a.
[Char]
-> [SizeBinder VName]
-> PatBase Info VName
-> Exp
-> InternaliseM a
-> InternaliseM a
internalisePat [Char]
desc [SizeBinder VName]
sizes PatBase Info VName
pat Exp
e forall a b. (a -> b) -> a -> b
$ [Char] -> Exp -> InternaliseM [SubExp]
internaliseExp [Char]
desc Exp
body
internaliseAppExp [Char]
_ AppRes
_ (E.LetFun VName
ofname ([TypeParamBase VName], [PatBase Info VName],
Maybe (TypeExp VName), Info StructRetType, Exp)
_ Exp
_ SrcLoc
_) =
forall a. HasCallStack => [Char] -> a
error forall a b. (a -> b) -> a -> b
$ [Char]
"Unexpected LetFun " forall a. [a] -> [a] -> [a]
++ forall a. Pretty a => a -> [Char]
prettyString VName
ofname
internaliseAppExp [Char]
desc AppRes
_ (E.DoLoop [VName]
sparams PatBase Info VName
mergepat Exp
mergeexp LoopFormBase Info VName
form Exp
loopbody SrcLoc
loc) = do
[SubExp]
ses <- [Char] -> Exp -> InternaliseM [SubExp]
internaliseExp [Char]
"loop_init" Exp
mergeexp
((Body SOACS
loopbody', (LoopForm SOACS
form', [Param DeclType]
shapepat, [Param DeclType]
mergepat', [SubExp]
mergeinit')), Stms SOACS
initstms) <-
forall (m :: * -> *) a.
MonadBuilder m =>
m a -> m (a, Stms (Rep m))
collectStms forall a b. (a -> b) -> a -> b
$ [SubExp]
-> LoopFormBase Info VName
-> InternaliseM
(Body SOACS,
(LoopForm SOACS, [Param DeclType], [Param DeclType], [SubExp]))
handleForm [SubExp]
ses LoopFormBase Info VName
form
forall (m :: * -> *). MonadBuilder m => Stms (Rep m) -> m ()
addStms Stms SOACS
initstms
[TypeBase Shape NoUniqueness]
mergeinit_ts' <- forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
(a -> m b) -> t a -> m (t b)
mapM forall t (m :: * -> *).
HasScope t m =>
SubExp -> m (TypeBase Shape NoUniqueness)
subExpType [SubExp]
mergeinit'
[SubExp]
ctxinit <- [VName]
-> [FParam SOACS]
-> [TypeBase Shape NoUniqueness]
-> InternaliseM [SubExp]
argShapes (forall a b. (a -> b) -> [a] -> [b]
map forall dec. Param dec -> VName
I.paramName [Param DeclType]
shapepat) [Param DeclType]
mergepat' [TypeBase Shape NoUniqueness]
mergeinit_ts'
let args :: [SubExp]
args = [SubExp]
ctxinit forall a. [a] -> [a] -> [a]
++ [SubExp]
mergeinit'
[SubExp]
args' <-
forall u.
Typed (TypeBase Shape u) =>
ErrorMsg SubExp
-> SrcLoc
-> [VName]
-> [TypeBase Shape u]
-> [SubExp]
-> InternaliseM [SubExp]
ensureArgShapes
ErrorMsg SubExp
"initial loop values have right shape"
SrcLoc
loc
(forall a b. (a -> b) -> [a] -> [b]
map forall dec. Param dec -> VName
I.paramName [Param DeclType]
shapepat)
(forall a b. (a -> b) -> [a] -> [b]
map forall dec. Typed dec => Param dec -> TypeBase Shape NoUniqueness
paramType forall a b. (a -> b) -> a -> b
$ [Param DeclType]
shapepat forall a. [a] -> [a] -> [a]
++ [Param DeclType]
mergepat')
[SubExp]
args
let dropCond :: [VName] -> [VName]
dropCond = case LoopFormBase Info VName
form of
E.While {} -> forall a. Int -> [a] -> [a]
drop Int
1
LoopFormBase Info VName
_ -> forall a. a -> a
id
let merge :: [(Param DeclType, SubExp)]
merge = forall a b. [a] -> [b] -> [(a, b)]
zip ([Param DeclType]
shapepat forall a. [a] -> [a] -> [a]
++ [Param DeclType]
mergepat') [SubExp]
args'
merge_ts :: [TypeBase Shape NoUniqueness]
merge_ts = forall a b. (a -> b) -> [a] -> [b]
map (forall dec. Typed dec => Param dec -> TypeBase Shape NoUniqueness
I.paramType forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall a b. (a, b) -> a
fst) [(Param DeclType, SubExp)]
merge
Body SOACS
loopbody'' <-
forall rep (m :: * -> *) a.
LocalScope rep m =>
Scope rep -> m a -> m a
localScope (forall rep dec. (FParamInfo rep ~ dec) => [Param dec] -> Scope rep
scopeOfFParams forall a b. (a -> b) -> a -> b
$ forall a b. (a -> b) -> [a] -> [b]
map forall a b. (a, b) -> a
fst [(Param DeclType, SubExp)]
merge) forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall rep a (m :: * -> *) b.
(Scoped rep a, LocalScope rep m) =>
a -> m b -> m b
inScopeOf LoopForm SOACS
form' forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall (m :: * -> *).
MonadBuilder m =>
m Result -> m (Body (Rep m))
buildBody_ forall a b. (a -> b) -> a -> b
$
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap [SubExp] -> Result
subExpsRes
forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall u.
Typed (TypeBase Shape u) =>
ErrorMsg SubExp
-> SrcLoc
-> [VName]
-> [TypeBase Shape u]
-> [SubExp]
-> InternaliseM [SubExp]
ensureArgShapes
ErrorMsg SubExp
"shape of loop result does not match shapes in loop parameter"
SrcLoc
loc
(forall a b. (a -> b) -> [a] -> [b]
map (forall dec. Param dec -> VName
I.paramName forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall a b. (a, b) -> a
fst) [(Param DeclType, SubExp)]
merge)
[TypeBase Shape NoUniqueness]
merge_ts
forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall a b. (a -> b) -> [a] -> [b]
map SubExpRes -> SubExp
resSubExp
forall (m :: * -> *) a b. Monad m => (a -> m b) -> m a -> m b
=<< forall (m :: * -> *). MonadBuilder m => Body (Rep m) -> m Result
bodyBind Body SOACS
loopbody'
Attrs
attrs <- forall r (m :: * -> *) a. MonadReader r m => (r -> a) -> m a
asks InternaliseEnv -> Attrs
envAttrs
forall a b. (a -> b) -> [a] -> [b]
map VName -> SubExp
I.Var forall b c a. (b -> c) -> (a -> b) -> a -> c
. [VName] -> [VName]
dropCond
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> forall (m :: * -> *) a. MonadBuilder m => Attrs -> m a -> m a
attributing
Attrs
attrs
([Char] -> Exp SOACS -> InternaliseM [VName]
letValExp [Char]
desc (forall rep.
[(FParam rep, SubExp)] -> LoopForm rep -> Body rep -> Exp rep
I.DoLoop [(Param DeclType, SubExp)]
merge LoopForm SOACS
form' Body SOACS
loopbody''))
where
sparams' :: [TypeParamBase VName]
sparams' = forall a b. (a -> b) -> [a] -> [b]
map (forall vn. vn -> SrcLoc -> TypeParamBase vn
`TypeParamDim` forall a. Monoid a => a
mempty) [VName]
sparams
forLoop :: [Param DeclType]
-> [Param DeclType]
-> [SubExp]
-> LoopForm SOACS
-> InternaliseM
(Body SOACS,
(LoopForm SOACS, [Param DeclType], [Param DeclType], [SubExp]))
forLoop [Param DeclType]
mergepat' [Param DeclType]
shapepat [SubExp]
mergeinit LoopForm SOACS
form' =
forall a. InternaliseM (Result, a) -> InternaliseM (Body SOACS, a)
bodyFromStms forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall rep a (m :: * -> *) b.
(Scoped rep a, LocalScope rep m) =>
a -> m b -> m b
inScopeOf LoopForm SOACS
form' forall a b. (a -> b) -> a -> b
$ do
[SubExp]
ses <- [Char] -> Exp -> InternaliseM [SubExp]
internaliseExp [Char]
"loopres" Exp
loopbody
[TypeBase Shape NoUniqueness]
sets <- forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
(a -> m b) -> t a -> m (t b)
mapM forall t (m :: * -> *).
HasScope t m =>
SubExp -> m (TypeBase Shape NoUniqueness)
subExpType [SubExp]
ses
[SubExp]
shapeargs <- [VName]
-> [FParam SOACS]
-> [TypeBase Shape NoUniqueness]
-> InternaliseM [SubExp]
argShapes (forall a b. (a -> b) -> [a] -> [b]
map forall dec. Param dec -> VName
I.paramName [Param DeclType]
shapepat) [Param DeclType]
mergepat' [TypeBase Shape NoUniqueness]
sets
forall (f :: * -> *) a. Applicative f => a -> f a
pure
( [SubExp] -> Result
subExpsRes forall a b. (a -> b) -> a -> b
$ [SubExp]
shapeargs forall a. [a] -> [a] -> [a]
++ [SubExp]
ses,
( LoopForm SOACS
form',
[Param DeclType]
shapepat,
[Param DeclType]
mergepat',
[SubExp]
mergeinit
)
)
handleForm :: [SubExp]
-> LoopFormBase Info VName
-> InternaliseM
(Body SOACS,
(LoopForm SOACS, [Param DeclType], [Param DeclType], [SubExp]))
handleForm [SubExp]
mergeinit (E.ForIn PatBase Info VName
x Exp
arr) = do
[VName]
arr' <- [Char] -> Exp -> InternaliseM [VName]
internaliseExpToVars [Char]
"for_in_arr" Exp
arr
[TypeBase Shape NoUniqueness]
arr_ts <- forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
(a -> m b) -> t a -> m (t b)
mapM forall rep (m :: * -> *).
HasScope rep m =>
VName -> m (TypeBase Shape NoUniqueness)
lookupType [VName]
arr'
let w :: SubExp
w = forall u. Int -> [TypeBase Shape u] -> SubExp
arraysSize Int
0 [TypeBase Shape NoUniqueness]
arr_ts
VName
i <- forall (m :: * -> *). MonadFreshNames m => [Char] -> m VName
newVName [Char]
"i"
[TypeBase Shape NoUniqueness]
ts <- forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
(a -> m b) -> t a -> m (t b)
mapM forall t (m :: * -> *).
HasScope t m =>
SubExp -> m (TypeBase Shape NoUniqueness)
subExpType [SubExp]
mergeinit
forall a.
[TypeParamBase VName]
-> PatBase Info VName
-> [TypeBase Shape NoUniqueness]
-> ([FParam SOACS] -> [FParam SOACS] -> InternaliseM a)
-> InternaliseM a
bindingLoopParams [TypeParamBase VName]
sparams' PatBase Info VName
mergepat [TypeBase Shape NoUniqueness]
ts forall a b. (a -> b) -> a -> b
$ \[FParam SOACS]
shapepat [FParam SOACS]
mergepat' ->
forall a.
[PatBase Info VName]
-> [TypeBase Shape NoUniqueness]
-> ([LParam SOACS] -> InternaliseM a)
-> InternaliseM a
bindingLambdaParams [PatBase Info VName
x] (forall a b. (a -> b) -> [a] -> [b]
map forall u. TypeBase Shape u -> TypeBase Shape u
rowType [TypeBase Shape NoUniqueness]
arr_ts) forall a b. (a -> b) -> a -> b
$ \[LParam SOACS]
x_params -> do
let loopvars :: [(Param (TypeBase Shape NoUniqueness), VName)]
loopvars = forall a b. [a] -> [b] -> [(a, b)]
zip [LParam SOACS]
x_params [VName]
arr'
[Param DeclType]
-> [Param DeclType]
-> [SubExp]
-> LoopForm SOACS
-> InternaliseM
(Body SOACS,
(LoopForm SOACS, [Param DeclType], [Param DeclType], [SubExp]))
forLoop [FParam SOACS]
mergepat' [FParam SOACS]
shapepat [SubExp]
mergeinit forall a b. (a -> b) -> a -> b
$
forall rep.
VName -> IntType -> SubExp -> [(LParam rep, VName)] -> LoopForm rep
I.ForLoop VName
i IntType
Int64 SubExp
w [(Param (TypeBase Shape NoUniqueness), VName)]
loopvars
handleForm [SubExp]
mergeinit (E.For IdentBase Info VName
i Exp
num_iterations) = do
SubExp
num_iterations' <- [Char] -> Exp -> InternaliseM SubExp
internaliseExp1 [Char]
"upper_bound" Exp
num_iterations
TypeBase Shape NoUniqueness
num_iterations_t <- forall t (m :: * -> *).
HasScope t m =>
SubExp -> m (TypeBase Shape NoUniqueness)
I.subExpType SubExp
num_iterations'
IntType
it <- case TypeBase Shape NoUniqueness
num_iterations_t of
I.Prim (IntType IntType
it) -> forall (f :: * -> *) a. Applicative f => a -> f a
pure IntType
it
TypeBase Shape NoUniqueness
_ -> forall a. HasCallStack => [Char] -> a
error [Char]
"internaliseExp DoLoop: invalid type"
[TypeBase Shape NoUniqueness]
ts <- forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
(a -> m b) -> t a -> m (t b)
mapM forall t (m :: * -> *).
HasScope t m =>
SubExp -> m (TypeBase Shape NoUniqueness)
subExpType [SubExp]
mergeinit
forall a.
[TypeParamBase VName]
-> PatBase Info VName
-> [TypeBase Shape NoUniqueness]
-> ([FParam SOACS] -> [FParam SOACS] -> InternaliseM a)
-> InternaliseM a
bindingLoopParams [TypeParamBase VName]
sparams' PatBase Info VName
mergepat [TypeBase Shape NoUniqueness]
ts forall a b. (a -> b) -> a -> b
$
\[FParam SOACS]
shapepat [FParam SOACS]
mergepat' ->
[Param DeclType]
-> [Param DeclType]
-> [SubExp]
-> LoopForm SOACS
-> InternaliseM
(Body SOACS,
(LoopForm SOACS, [Param DeclType], [Param DeclType], [SubExp]))
forLoop [FParam SOACS]
mergepat' [FParam SOACS]
shapepat [SubExp]
mergeinit forall a b. (a -> b) -> a -> b
$
forall rep.
VName -> IntType -> SubExp -> [(LParam rep, VName)] -> LoopForm rep
I.ForLoop (forall (f :: * -> *) vn. IdentBase f vn -> vn
E.identName IdentBase Info VName
i) IntType
it SubExp
num_iterations' []
handleForm [SubExp]
mergeinit (E.While Exp
cond) = do
[TypeBase Shape NoUniqueness]
ts <- forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
(a -> m b) -> t a -> m (t b)
mapM forall t (m :: * -> *).
HasScope t m =>
SubExp -> m (TypeBase Shape NoUniqueness)
subExpType [SubExp]
mergeinit
forall a.
[TypeParamBase VName]
-> PatBase Info VName
-> [TypeBase Shape NoUniqueness]
-> ([FParam SOACS] -> [FParam SOACS] -> InternaliseM a)
-> InternaliseM a
bindingLoopParams [TypeParamBase VName]
sparams' PatBase Info VName
mergepat [TypeBase Shape NoUniqueness]
ts forall a b. (a -> b) -> a -> b
$ \[FParam SOACS]
shapepat [FParam SOACS]
mergepat' -> do
[TypeBase Shape NoUniqueness]
mergeinit_ts <- forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
(a -> m b) -> t a -> m (t b)
mapM forall t (m :: * -> *).
HasScope t m =>
SubExp -> m (TypeBase Shape NoUniqueness)
subExpType [SubExp]
mergeinit
[SubExp]
shapeinit <- [VName]
-> [FParam SOACS]
-> [TypeBase Shape NoUniqueness]
-> InternaliseM [SubExp]
argShapes (forall a b. (a -> b) -> [a] -> [b]
map forall dec. Param dec -> VName
I.paramName [FParam SOACS]
shapepat) [FParam SOACS]
mergepat' [TypeBase Shape NoUniqueness]
mergeinit_ts
(SubExp
loop_initial_cond, Stms SOACS
init_loop_cond_stms) <- forall (m :: * -> *) a.
MonadBuilder m =>
m a -> m (a, Stms (Rep m))
collectStms forall a b. (a -> b) -> a -> b
$ do
forall (t :: * -> *) (m :: * -> *) a b.
(Foldable t, Monad m) =>
t a -> (a -> m b) -> m ()
forM_ (forall a b. [a] -> [b] -> [(a, b)]
zip [FParam SOACS]
shapepat [SubExp]
shapeinit) forall a b. (a -> b) -> a -> b
$ \(Param DeclType
p, SubExp
se) ->
forall (m :: * -> *).
MonadBuilder m =>
[VName] -> Exp (Rep m) -> m ()
letBindNames [forall dec. Param dec -> VName
I.paramName Param DeclType
p] forall a b. (a -> b) -> a -> b
$ forall rep. BasicOp -> Exp rep
BasicOp forall a b. (a -> b) -> a -> b
$ SubExp -> BasicOp
SubExp SubExp
se
forall (t :: * -> *) (m :: * -> *) a b.
(Foldable t, Monad m) =>
t a -> (a -> m b) -> m ()
forM_ (forall a b. [a] -> [b] -> [(a, b)]
zip [FParam SOACS]
mergepat' [SubExp]
mergeinit) forall a b. (a -> b) -> a -> b
$ \(Param DeclType
p, SubExp
se) ->
forall (f :: * -> *). Applicative f => Bool -> f () -> f ()
unless (SubExp
se forall a. Eq a => a -> a -> Bool
== VName -> SubExp
I.Var (forall dec. Param dec -> VName
I.paramName Param DeclType
p)) forall a b. (a -> b) -> a -> b
$
forall (m :: * -> *).
MonadBuilder m =>
[VName] -> Exp (Rep m) -> m ()
letBindNames [forall dec. Param dec -> VName
I.paramName Param DeclType
p] forall a b. (a -> b) -> a -> b
$
forall rep. BasicOp -> Exp rep
BasicOp forall a b. (a -> b) -> a -> b
$
case SubExp
se of
I.Var VName
v
| Bool -> Bool
not forall a b. (a -> b) -> a -> b
$ forall shape u. TypeBase shape u -> Bool
primType forall a b. (a -> b) -> a -> b
$ forall dec. Typed dec => Param dec -> TypeBase Shape NoUniqueness
paramType Param DeclType
p ->
ReshapeKind -> Shape -> VName -> BasicOp
Reshape ReshapeKind
I.ReshapeCoerce (forall shape u. ArrayShape shape => TypeBase shape u -> shape
I.arrayShape forall a b. (a -> b) -> a -> b
$ forall dec. Typed dec => Param dec -> TypeBase Shape NoUniqueness
paramType Param DeclType
p) VName
v
SubExp
_ -> SubExp -> BasicOp
SubExp SubExp
se
[Char] -> Exp -> InternaliseM SubExp
internaliseExp1 [Char]
"loop_cond" Exp
cond
forall (m :: * -> *). MonadBuilder m => Stms (Rep m) -> m ()
addStms Stms SOACS
init_loop_cond_stms
forall a. InternaliseM (Result, a) -> InternaliseM (Body SOACS, a)
bodyFromStms forall a b. (a -> b) -> a -> b
$ do
[SubExp]
ses <- [Char] -> Exp -> InternaliseM [SubExp]
internaliseExp [Char]
"loopres" Exp
loopbody
[TypeBase Shape NoUniqueness]
sets <- forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
(a -> m b) -> t a -> m (t b)
mapM forall t (m :: * -> *).
HasScope t m =>
SubExp -> m (TypeBase Shape NoUniqueness)
subExpType [SubExp]
ses
Param DeclType
loop_while <- forall (m :: * -> *) dec.
MonadFreshNames m =>
[Char] -> dec -> m (Param dec)
newParam [Char]
"loop_while" forall a b. (a -> b) -> a -> b
$ forall shape u. PrimType -> TypeBase shape u
I.Prim PrimType
I.Bool
[SubExp]
shapeargs <- [VName]
-> [FParam SOACS]
-> [TypeBase Shape NoUniqueness]
-> InternaliseM [SubExp]
argShapes (forall a b. (a -> b) -> [a] -> [b]
map forall dec. Param dec -> VName
I.paramName [FParam SOACS]
shapepat) [FParam SOACS]
mergepat' [TypeBase Shape NoUniqueness]
sets
Body SOACS
loop_end_cond_body <- forall rep (m :: * -> *).
(Renameable rep, MonadFreshNames m) =>
Body rep -> m (Body rep)
renameBody forall (m :: * -> *) b c a.
Monad m =>
(b -> m c) -> (a -> m b) -> a -> m c
<=< forall (m :: * -> *).
MonadBuilder m =>
m Result -> m (Body (Rep m))
buildBody_ forall a b. (a -> b) -> a -> b
$ do
forall (t :: * -> *) (m :: * -> *) a b.
(Foldable t, Monad m) =>
t a -> (a -> m b) -> m ()
forM_ (forall a b. [a] -> [b] -> [(a, b)]
zip [FParam SOACS]
shapepat [SubExp]
shapeargs) forall a b. (a -> b) -> a -> b
$ \(Param DeclType
p, SubExp
se) ->
forall (f :: * -> *). Applicative f => Bool -> f () -> f ()
unless (SubExp
se forall a. Eq a => a -> a -> Bool
== VName -> SubExp
I.Var (forall dec. Param dec -> VName
I.paramName Param DeclType
p)) forall a b. (a -> b) -> a -> b
$
forall (m :: * -> *).
MonadBuilder m =>
[VName] -> Exp (Rep m) -> m ()
letBindNames [forall dec. Param dec -> VName
I.paramName Param DeclType
p] forall a b. (a -> b) -> a -> b
$
forall rep. BasicOp -> Exp rep
BasicOp forall a b. (a -> b) -> a -> b
$
SubExp -> BasicOp
SubExp SubExp
se
forall (t :: * -> *) (m :: * -> *) a b.
(Foldable t, Monad m) =>
t a -> (a -> m b) -> m ()
forM_ (forall a b. [a] -> [b] -> [(a, b)]
zip [FParam SOACS]
mergepat' [SubExp]
ses) forall a b. (a -> b) -> a -> b
$ \(Param DeclType
p, SubExp
se) ->
forall (f :: * -> *). Applicative f => Bool -> f () -> f ()
unless (SubExp
se forall a. Eq a => a -> a -> Bool
== VName -> SubExp
I.Var (forall dec. Param dec -> VName
I.paramName Param DeclType
p)) forall a b. (a -> b) -> a -> b
$
forall (m :: * -> *).
MonadBuilder m =>
[VName] -> Exp (Rep m) -> m ()
letBindNames [forall dec. Param dec -> VName
I.paramName Param DeclType
p] forall a b. (a -> b) -> a -> b
$
forall rep. BasicOp -> Exp rep
BasicOp forall a b. (a -> b) -> a -> b
$
case SubExp
se of
I.Var VName
v
| Bool -> Bool
not forall a b. (a -> b) -> a -> b
$ forall shape u. TypeBase shape u -> Bool
primType forall a b. (a -> b) -> a -> b
$ forall dec. Typed dec => Param dec -> TypeBase Shape NoUniqueness
paramType Param DeclType
p ->
ReshapeKind -> Shape -> VName -> BasicOp
Reshape ReshapeKind
I.ReshapeCoerce (forall shape u. ArrayShape shape => TypeBase shape u -> shape
I.arrayShape forall a b. (a -> b) -> a -> b
$ forall dec. Typed dec => Param dec -> TypeBase Shape NoUniqueness
paramType Param DeclType
p) VName
v
SubExp
_ -> SubExp -> BasicOp
SubExp SubExp
se
[SubExp] -> Result
subExpsRes forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> [Char] -> Exp -> InternaliseM [SubExp]
internaliseExp [Char]
"loop_cond" Exp
cond
Result
loop_end_cond <- forall (m :: * -> *). MonadBuilder m => Body (Rep m) -> m Result
bodyBind Body SOACS
loop_end_cond_body
forall (f :: * -> *) a. Applicative f => a -> f a
pure
( [SubExp] -> Result
subExpsRes [SubExp]
shapeargs forall a. [a] -> [a] -> [a]
++ Result
loop_end_cond forall a. [a] -> [a] -> [a]
++ [SubExp] -> Result
subExpsRes [SubExp]
ses,
( forall rep. VName -> LoopForm rep
I.WhileLoop forall a b. (a -> b) -> a -> b
$ forall dec. Param dec -> VName
I.paramName Param DeclType
loop_while,
[FParam SOACS]
shapepat,
Param DeclType
loop_while forall a. a -> [a] -> [a]
: [FParam SOACS]
mergepat',
SubExp
loop_initial_cond forall a. a -> [a] -> [a]
: [SubExp]
mergeinit
)
)
internaliseAppExp [Char]
desc AppRes
_ (E.LetWith IdentBase Info VName
name IdentBase Info VName
src SliceBase Info VName
idxs Exp
ve Exp
body SrcLoc
loc) = do
let pat :: PatBase Info VName
pat = forall (f :: * -> *) vn. vn -> f PatType -> SrcLoc -> PatBase f vn
E.Id (forall (f :: * -> *) vn. IdentBase f vn -> vn
E.identName IdentBase Info VName
name) (forall (f :: * -> *) vn. IdentBase f vn -> f PatType
E.identType IdentBase Info VName
name) SrcLoc
loc
src_t :: Info PatType
src_t = forall dim as. TypeBase dim as -> TypeBase dim (Set Alias)
E.fromStruct forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> forall (f :: * -> *) vn. IdentBase f vn -> f PatType
E.identType IdentBase Info VName
src
e :: Exp
e = forall (f :: * -> *) vn.
ExpBase f vn
-> SliceBase f vn -> ExpBase f vn -> SrcLoc -> ExpBase f vn
E.Update (forall (f :: * -> *) vn.
QualName vn -> f PatType -> SrcLoc -> ExpBase f vn
E.Var (forall v. v -> QualName v
E.qualName forall a b. (a -> b) -> a -> b
$ forall (f :: * -> *) vn. IdentBase f vn -> vn
E.identName IdentBase Info VName
src) Info PatType
src_t SrcLoc
loc) SliceBase Info VName
idxs Exp
ve SrcLoc
loc
[Char] -> Exp -> InternaliseM [SubExp]
internaliseExp [Char]
desc forall a b. (a -> b) -> a -> b
$
forall (f :: * -> *) vn.
AppExpBase f vn -> f AppRes -> ExpBase f vn
E.AppExp
(forall (f :: * -> *) vn.
[SizeBinder vn]
-> PatBase f vn
-> ExpBase f vn
-> ExpBase f vn
-> SrcLoc
-> AppExpBase f vn
E.LetPat [] PatBase Info VName
pat Exp
e Exp
body SrcLoc
loc)
(forall a. a -> Info a
Info (PatType -> [VName] -> AppRes
AppRes (Exp -> PatType
E.typeOf Exp
body) forall a. Monoid a => a
mempty))
internaliseAppExp [Char]
desc AppRes
_ (E.Match Exp
e NonEmpty (CaseBase Info VName)
orig_cs SrcLoc
_) = do
[SubExp]
ses <- [Char] -> Exp -> InternaliseM [SubExp]
internaliseExp ([Char]
desc forall a. [a] -> [a] -> [a]
++ [Char]
"_scrutinee") Exp
e
NonEmpty (Case (InternaliseM (Body SOACS)))
cs <- forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
(a -> m b) -> t a -> m (t b)
mapM ([SubExp]
-> CaseBase Info VName
-> InternaliseM (Case (InternaliseM (Body SOACS)))
onCase [SubExp]
ses) NonEmpty (CaseBase Info VName)
orig_cs
case forall a. NonEmpty a -> (a, Maybe (NonEmpty a))
NE.uncons NonEmpty (Case (InternaliseM (Body SOACS)))
cs of
(I.Case [Maybe PrimValue]
_ InternaliseM (Body SOACS)
body, Maybe (NonEmpty (Case (InternaliseM (Body SOACS))))
Nothing) ->
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap (forall a b. (a -> b) -> [a] -> [b]
map SubExpRes -> SubExp
resSubExp) forall a b. (a -> b) -> a -> b
$ forall (m :: * -> *). MonadBuilder m => Body (Rep m) -> m Result
bodyBind forall (m :: * -> *) a b. Monad m => (a -> m b) -> m a -> m b
=<< InternaliseM (Body SOACS)
body
(Case (InternaliseM (Body SOACS)),
Maybe (NonEmpty (Case (InternaliseM (Body SOACS)))))
_ -> do
[Char] -> Exp SOACS -> InternaliseM [SubExp]
letValExp' [Char]
desc forall (m :: * -> *) a b. Monad m => (a -> m b) -> m a -> m b
=<< forall (m :: * -> *).
(MonadBuilder m, BranchType (Rep m) ~ ExtType) =>
[SubExp]
-> [Case (m (Body (Rep m)))] -> m (Body (Rep m)) -> m (Exp (Rep m))
eMatch [SubExp]
ses (forall a. NonEmpty a -> [a]
NE.init NonEmpty (Case (InternaliseM (Body SOACS)))
cs) (forall body. Case body -> body
I.caseBody forall a b. (a -> b) -> a -> b
$ forall a. NonEmpty a -> a
NE.last NonEmpty (Case (InternaliseM (Body SOACS)))
cs)
where
onCase :: [SubExp]
-> CaseBase Info VName
-> InternaliseM (Case (InternaliseM (Body SOACS)))
onCase [SubExp]
ses (E.CasePat PatBase Info VName
p Exp
case_e SrcLoc
_) = do
([Maybe PrimValue]
cmps, [SubExp]
pertinent) <- PatBase Info VName
-> [SubExp] -> InternaliseM ([Maybe PrimValue], [SubExp])
generateCond PatBase Info VName
p [SubExp]
ses
forall (f :: * -> *) a. Applicative f => a -> f a
pure forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall body. [Maybe PrimValue] -> body -> Case body
I.Case [Maybe PrimValue]
cmps forall a b. (a -> b) -> a -> b
$
forall a.
[SizeBinder VName]
-> PatBase Info VName
-> [SubExp]
-> InternaliseM a
-> InternaliseM a
internalisePat' [] PatBase Info VName
p [SubExp]
pertinent forall a b. (a -> b) -> a -> b
$
[Char] -> Exp -> InternaliseM (Body SOACS)
internaliseBody [Char]
"case" Exp
case_e
internaliseAppExp [Char]
desc AppRes
_ (E.If Exp
ce Exp
te Exp
fe SrcLoc
_) =
[Char] -> Exp SOACS -> InternaliseM [SubExp]
letValExp' [Char]
desc
forall (m :: * -> *) a b. Monad m => (a -> m b) -> m a -> m b
=<< forall (m :: * -> *).
(MonadBuilder m, BranchType (Rep m) ~ ExtType) =>
m (Exp (Rep m))
-> m (Body (Rep m)) -> m (Body (Rep m)) -> m (Exp (Rep m))
eIf
(forall rep. BasicOp -> Exp rep
BasicOp forall b c a. (b -> c) -> (a -> b) -> a -> c
. SubExp -> BasicOp
SubExp forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> [Char] -> Exp -> InternaliseM SubExp
internaliseExp1 [Char]
"cond" Exp
ce)
([Char] -> Exp -> InternaliseM (Body SOACS)
internaliseBody ([Char]
desc forall a. Semigroup a => a -> a -> a
<> [Char]
"_t") Exp
te)
([Char] -> Exp -> InternaliseM (Body SOACS)
internaliseBody ([Char]
desc forall a. Semigroup a => a -> a -> a
<> [Char]
"_f") Exp
fe)
internaliseAppExp [Char]
_ AppRes
_ e :: AppExp
e@E.BinOp {} =
forall a. HasCallStack => [Char] -> a
error forall a b. (a -> b) -> a -> b
$ [Char]
"internaliseAppExp: Unexpected BinOp " forall a. [a] -> [a] -> [a]
++ forall a. Pretty a => a -> [Char]
prettyString AppExp
e
internaliseExp :: String -> E.Exp -> InternaliseM [I.SubExp]
internaliseExp :: [Char] -> Exp -> InternaliseM [SubExp]
internaliseExp [Char]
desc (E.Parens Exp
e SrcLoc
_) =
[Char] -> Exp -> InternaliseM [SubExp]
internaliseExp [Char]
desc Exp
e
internaliseExp [Char]
desc (E.Hole (Info PatType
t) SrcLoc
loc) = do
let msg :: Text
msg = forall a. Doc a -> Text
docText forall a b. (a -> b) -> a -> b
$ Doc Any
"Reached hole of type: " forall a. Semigroup a => a -> a -> a
<> forall ann. Doc ann -> Doc ann
align (forall a ann. Pretty a => a -> Doc ann
pretty PatType
t)
ts :: [TypeBase ExtShape Uniqueness]
ts = TypeBase Size () -> [TypeBase ExtShape Uniqueness]
internaliseType (forall dim as. TypeBase dim as -> TypeBase dim ()
E.toStruct PatType
t)
Certs
c <- [Char] -> SubExp -> ErrorMsg SubExp -> SrcLoc -> InternaliseM Certs
assert [Char]
"hole_c" (forall v. IsValue v => v -> SubExp
constant Bool
False) (forall a. [ErrorMsgPart a] -> ErrorMsg a
errorMsg [forall a. Text -> ErrorMsgPart a
ErrorString Text
msg]) SrcLoc
loc
case forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
(a -> m b) -> t a -> m (t b)
mapM forall u. TypeBase ExtShape u -> Maybe (TypeBase Shape u)
hasStaticShape [TypeBase ExtShape Uniqueness]
ts of
Maybe [DeclType]
Nothing ->
forall a. HasCallStack => [Char] -> a
error forall a b. (a -> b) -> a -> b
$ [Char]
"Hole at " forall a. Semigroup a => a -> a -> a
<> forall a. Located a => a -> [Char]
locStr SrcLoc
loc forall a. Semigroup a => a -> a -> a
<> [Char]
" has existential type:\n" forall a. Semigroup a => a -> a -> a
<> forall a. Show a => a -> [Char]
show [TypeBase ExtShape Uniqueness]
ts
Just [DeclType]
ts' ->
forall (m :: * -> *) a. MonadBuilder m => Certs -> m a -> m a
certifying Certs
c forall a b. (a -> b) -> a -> b
$ forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
(a -> m b) -> t a -> m (t b)
mapM (forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap VName -> SubExp
I.Var forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall (m :: * -> *).
MonadBuilder m =>
[Char] -> Exp (Rep m) -> m VName
letExp [Char]
desc forall (m :: * -> *) b c a.
Monad m =>
(b -> m c) -> (a -> m b) -> a -> m c
<=< forall (m :: * -> *).
MonadBuilder m =>
TypeBase Shape NoUniqueness -> m (Exp (Rep m))
eBlank forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall shape.
TypeBase shape Uniqueness -> TypeBase shape NoUniqueness
I.fromDecl) [DeclType]
ts'
internaliseExp [Char]
desc (E.QualParens (QualName VName, SrcLoc)
_ Exp
e SrcLoc
_) =
[Char] -> Exp -> InternaliseM [SubExp]
internaliseExp [Char]
desc Exp
e
internaliseExp [Char]
desc (E.StringLit [Word8]
vs SrcLoc
_) =
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap forall (f :: * -> *) a. Applicative f => a -> f a
pure forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall (m :: * -> *).
MonadBuilder m =>
[Char] -> Exp (Rep m) -> m SubExp
letSubExp [Char]
desc forall a b. (a -> b) -> a -> b
$
forall rep. BasicOp -> Exp rep
I.BasicOp forall a b. (a -> b) -> a -> b
$
[SubExp] -> TypeBase Shape NoUniqueness -> BasicOp
I.ArrayLit (forall a b. (a -> b) -> [a] -> [b]
map forall v. IsValue v => v -> SubExp
constant [Word8]
vs) forall a b. (a -> b) -> a -> b
$
forall shape u. PrimType -> TypeBase shape u
I.Prim PrimType
int8
internaliseExp [Char]
_ (E.Var (E.QualName [VName]
_ VName
name) Info PatType
_ SrcLoc
_) = do
Maybe [SubExp]
subst <- VName -> InternaliseM (Maybe [SubExp])
lookupSubst VName
name
case Maybe [SubExp]
subst of
Just [SubExp]
substs -> forall (f :: * -> *) a. Applicative f => a -> f a
pure [SubExp]
substs
Maybe [SubExp]
Nothing -> forall (f :: * -> *) a. Applicative f => a -> f a
pure [VName -> SubExp
I.Var VName
name]
internaliseExp [Char]
desc (E.AppExp AppExp
e (Info AppRes
appres)) = do
[SubExp]
ses <- [Char] -> AppRes -> AppExp -> InternaliseM [SubExp]
internaliseAppExp [Char]
desc AppRes
appres AppExp
e
AppRes -> [SubExp] -> InternaliseM ()
bindExtSizes AppRes
appres [SubExp]
ses
forall (f :: * -> *) a. Applicative f => a -> f a
pure [SubExp]
ses
internaliseExp [Char]
_ (E.TupLit [] SrcLoc
_) =
forall (f :: * -> *) a. Applicative f => a -> f a
pure [forall v. IsValue v => v -> SubExp
constant PrimValue
UnitValue]
internaliseExp [Char]
_ (E.RecordLit [] SrcLoc
_) =
forall (f :: * -> *) a. Applicative f => a -> f a
pure [forall v. IsValue v => v -> SubExp
constant PrimValue
UnitValue]
internaliseExp [Char]
desc (E.TupLit [Exp]
es SrcLoc
_) = forall (t :: * -> *) a. Foldable t => t [a] -> [a]
concat forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
(a -> m b) -> t a -> m (t b)
mapM ([Char] -> Exp -> InternaliseM [SubExp]
internaliseExp [Char]
desc) [Exp]
es
internaliseExp [Char]
desc (E.RecordLit [FieldBase Info VName]
orig_fields SrcLoc
_) =
forall (t :: * -> *) a b. Foldable t => (a -> [b]) -> t a -> [b]
concatMap forall a b. (a, b) -> b
snd forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall a. Map Name a -> [(Name, a)]
sortFields forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall (f :: * -> *) k a.
(Foldable f, Ord k) =>
f (Map k a) -> Map k a
M.unions forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
(a -> m b) -> t a -> m (t b)
mapM FieldBase Info VName -> InternaliseM (Map Name [SubExp])
internaliseField [FieldBase Info VName]
orig_fields
where
internaliseField :: FieldBase Info VName -> InternaliseM (Map Name [SubExp])
internaliseField (E.RecordFieldExplicit Name
name Exp
e SrcLoc
_) =
forall k a. k -> a -> Map k a
M.singleton Name
name forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> [Char] -> Exp -> InternaliseM [SubExp]
internaliseExp [Char]
desc Exp
e
internaliseField (E.RecordFieldImplicit VName
name Info PatType
t SrcLoc
loc) =
FieldBase Info VName -> InternaliseM (Map Name [SubExp])
internaliseField forall a b. (a -> b) -> a -> b
$
forall (f :: * -> *) vn.
Name -> ExpBase f vn -> SrcLoc -> FieldBase f vn
E.RecordFieldExplicit
(VName -> Name
baseName VName
name)
(forall (f :: * -> *) vn.
QualName vn -> f PatType -> SrcLoc -> ExpBase f vn
E.Var (forall v. v -> QualName v
E.qualName VName
name) Info PatType
t SrcLoc
loc)
SrcLoc
loc
internaliseExp [Char]
desc (E.ArrayLit [Exp]
es (Info PatType
arr_t) SrcLoc
loc)
| Just (([Int]
eshape, [Exp]
e') : [([Int], [Exp])]
es') <- forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
(a -> m b) -> t a -> m (t b)
mapM Exp -> Maybe ([Int], [Exp])
isArrayLiteral [Exp]
es,
Bool -> Bool
not forall a b. (a -> b) -> a -> b
$ forall (t :: * -> *) a. Foldable t => t a -> Bool
null [Int]
eshape,
forall (t :: * -> *) a. Foldable t => (a -> Bool) -> t a -> Bool
all (([Int]
eshape ==) forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall a b. (a, b) -> a
fst) [([Int], [Exp])]
es',
Just PatType
basetype <- forall dim as. Int -> TypeBase dim as -> Maybe (TypeBase dim as)
E.peelArray (forall (t :: * -> *) a. Foldable t => t a -> Int
length [Int]
eshape) PatType
arr_t = do
let flat_lit :: Exp
flat_lit = forall (f :: * -> *) vn.
[ExpBase f vn] -> f PatType -> SrcLoc -> ExpBase f vn
E.ArrayLit ([Exp]
e' forall a. [a] -> [a] -> [a]
++ forall (t :: * -> *) a b. Foldable t => (a -> [b]) -> t a -> [b]
concatMap forall a b. (a, b) -> b
snd [([Int], [Exp])]
es') (forall a. a -> Info a
Info PatType
basetype) SrcLoc
loc
new_shape :: [Int]
new_shape = forall (t :: * -> *) a. Foldable t => t a -> Int
length [Exp]
es forall a. a -> [a] -> [a]
: [Int]
eshape
[VName]
flat_arrs <- [Char] -> Exp -> InternaliseM [VName]
internaliseExpToVars [Char]
"flat_literal" Exp
flat_lit
forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
t a -> (a -> m b) -> m (t b)
forM [VName]
flat_arrs forall a b. (a -> b) -> a -> b
$ \VName
flat_arr -> do
TypeBase Shape NoUniqueness
flat_arr_t <- forall rep (m :: * -> *).
HasScope rep m =>
VName -> m (TypeBase Shape NoUniqueness)
lookupType VName
flat_arr
let new_shape' :: Shape
new_shape' =
Shape -> Int -> Shape -> Shape
reshapeOuter
(forall d. [d] -> ShapeBase d
I.Shape forall a b. (a -> b) -> a -> b
$ forall a b. (a -> b) -> [a] -> [b]
map (IntType -> Integer -> SubExp
intConst IntType
Int64 forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall a. Integral a => a -> Integer
toInteger) [Int]
new_shape)
Int
1
forall a b. (a -> b) -> a -> b
$ forall shape u. ArrayShape shape => TypeBase shape u -> shape
I.arrayShape TypeBase Shape NoUniqueness
flat_arr_t
forall (m :: * -> *).
MonadBuilder m =>
[Char] -> Exp (Rep m) -> m SubExp
letSubExp [Char]
desc forall a b. (a -> b) -> a -> b
$ forall rep. BasicOp -> Exp rep
I.BasicOp forall a b. (a -> b) -> a -> b
$ ReshapeKind -> Shape -> VName -> BasicOp
I.Reshape ReshapeKind
I.ReshapeArbitrary Shape
new_shape' VName
flat_arr
| Bool
otherwise = do
[[SubExp]]
es' <- forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
(a -> m b) -> t a -> m (t b)
mapM ([Char] -> Exp -> InternaliseM [SubExp]
internaliseExp [Char]
"arr_elem") [Exp]
es
let arr_t_ext :: [TypeBase ExtShape Uniqueness]
arr_t_ext = TypeBase Size () -> [TypeBase ExtShape Uniqueness]
internaliseType forall a b. (a -> b) -> a -> b
$ forall dim as. TypeBase dim as -> TypeBase dim ()
E.toStruct PatType
arr_t
[TypeBase Shape NoUniqueness]
rowtypes <-
case forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
(a -> m b) -> t a -> m (t b)
mapM (forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap forall u. TypeBase Shape u -> TypeBase Shape u
rowType forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall u. TypeBase ExtShape u -> Maybe (TypeBase Shape u)
hasStaticShape forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall shape.
TypeBase shape Uniqueness -> TypeBase shape NoUniqueness
I.fromDecl) [TypeBase ExtShape Uniqueness]
arr_t_ext of
Just [TypeBase Shape NoUniqueness]
ts -> forall (f :: * -> *) a. Applicative f => a -> f a
pure [TypeBase Shape NoUniqueness]
ts
Maybe [TypeBase Shape NoUniqueness]
Nothing ->
case [[SubExp]]
es' of
[] -> forall a. HasCallStack => [Char] -> a
error forall a b. (a -> b) -> a -> b
$ [Char]
"internaliseExp ArrayLit: existential type: " forall a. [a] -> [a] -> [a]
++ forall a. Pretty a => a -> [Char]
prettyString PatType
arr_t
[SubExp]
e' : [[SubExp]]
_ -> forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
(a -> m b) -> t a -> m (t b)
mapM forall t (m :: * -> *).
HasScope t m =>
SubExp -> m (TypeBase Shape NoUniqueness)
subExpType [SubExp]
e'
let arraylit :: [SubExp] -> TypeBase Shape NoUniqueness -> InternaliseM (Exp SOACS)
arraylit [SubExp]
ks TypeBase Shape NoUniqueness
rt = do
[SubExp]
ks' <-
forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
(a -> m b) -> t a -> m (t b)
mapM
( ErrorMsg SubExp
-> SrcLoc
-> TypeBase Shape NoUniqueness
-> [Char]
-> SubExp
-> InternaliseM SubExp
ensureShape
ErrorMsg SubExp
"shape of element differs from shape of first element"
SrcLoc
loc
TypeBase Shape NoUniqueness
rt
[Char]
"elem_reshaped"
)
[SubExp]
ks
forall (f :: * -> *) a. Applicative f => a -> f a
pure forall a b. (a -> b) -> a -> b
$ forall rep. BasicOp -> Exp rep
I.BasicOp forall a b. (a -> b) -> a -> b
$ [SubExp] -> TypeBase Shape NoUniqueness -> BasicOp
I.ArrayLit [SubExp]
ks' TypeBase Shape NoUniqueness
rt
forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
(a -> m b) -> t a -> m (t b)
mapM (forall (m :: * -> *).
MonadBuilder m =>
[Char] -> Exp (Rep m) -> m SubExp
letSubExp [Char]
desc)
forall (m :: * -> *) a b. Monad m => (a -> m b) -> m a -> m b
=<< if forall (t :: * -> *) a. Foldable t => t a -> Bool
null [[SubExp]]
es'
then forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
(a -> m b) -> t a -> m (t b)
mapM ([SubExp] -> TypeBase Shape NoUniqueness -> InternaliseM (Exp SOACS)
arraylit []) [TypeBase Shape NoUniqueness]
rowtypes
else forall (m :: * -> *) a b c.
Applicative m =>
(a -> b -> m c) -> [a] -> [b] -> m [c]
zipWithM [SubExp] -> TypeBase Shape NoUniqueness -> InternaliseM (Exp SOACS)
arraylit (forall a. [[a]] -> [[a]]
transpose [[SubExp]]
es') [TypeBase Shape NoUniqueness]
rowtypes
where
isArrayLiteral :: E.Exp -> Maybe ([Int], [E.Exp])
isArrayLiteral :: Exp -> Maybe ([Int], [Exp])
isArrayLiteral (E.ArrayLit [Exp]
inner_es Info PatType
_ SrcLoc
_) = do
([Int]
eshape, [Exp]
e) : [([Int], [Exp])]
inner_es' <- forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
(a -> m b) -> t a -> m (t b)
mapM Exp -> Maybe ([Int], [Exp])
isArrayLiteral [Exp]
inner_es
forall (f :: * -> *). Alternative f => Bool -> f ()
guard forall a b. (a -> b) -> a -> b
$ forall (t :: * -> *) a. Foldable t => (a -> Bool) -> t a -> Bool
all (([Int]
eshape ==) forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall a b. (a, b) -> a
fst) [([Int], [Exp])]
inner_es'
forall (f :: * -> *) a. Applicative f => a -> f a
pure (forall (t :: * -> *) a. Foldable t => t a -> Int
length [Exp]
inner_es forall a. a -> [a] -> [a]
: [Int]
eshape, [Exp]
e forall a. [a] -> [a] -> [a]
++ forall (t :: * -> *) a b. Foldable t => (a -> [b]) -> t a -> [b]
concatMap forall a b. (a, b) -> b
snd [([Int], [Exp])]
inner_es')
isArrayLiteral Exp
e =
forall a. a -> Maybe a
Just ([], [Exp
e])
internaliseExp [Char]
desc (E.Ascript Exp
e TypeExp VName
_ SrcLoc
_) =
[Char] -> Exp -> InternaliseM [SubExp]
internaliseExp [Char]
desc Exp
e
internaliseExp [Char]
desc (E.Negate Exp
e SrcLoc
_) = do
SubExp
e' <- [Char] -> Exp -> InternaliseM SubExp
internaliseExp1 [Char]
"negate_arg" Exp
e
TypeBase Shape NoUniqueness
et <- forall t (m :: * -> *).
HasScope t m =>
SubExp -> m (TypeBase Shape NoUniqueness)
subExpType SubExp
e'
case TypeBase Shape NoUniqueness
et of
I.Prim (I.IntType IntType
t) ->
forall (m :: * -> *).
MonadBuilder m =>
[Char] -> Exp (Rep m) -> m [SubExp]
letTupExp' [Char]
desc forall a b. (a -> b) -> a -> b
$ forall rep. BasicOp -> Exp rep
I.BasicOp forall a b. (a -> b) -> a -> b
$ BinOp -> SubExp -> SubExp -> BasicOp
I.BinOp (IntType -> Overflow -> BinOp
I.Sub IntType
t Overflow
I.OverflowWrap) (IntType -> Integer -> SubExp
I.intConst IntType
t Integer
0) SubExp
e'
I.Prim (I.FloatType FloatType
t) ->
forall (m :: * -> *).
MonadBuilder m =>
[Char] -> Exp (Rep m) -> m [SubExp]
letTupExp' [Char]
desc forall a b. (a -> b) -> a -> b
$ forall rep. BasicOp -> Exp rep
I.BasicOp forall a b. (a -> b) -> a -> b
$ BinOp -> SubExp -> SubExp -> BasicOp
I.BinOp (FloatType -> BinOp
I.FSub FloatType
t) (FloatType -> Double -> SubExp
I.floatConst FloatType
t Double
0) SubExp
e'
TypeBase Shape NoUniqueness
_ -> forall a. HasCallStack => [Char] -> a
error [Char]
"Futhark.Internalise.internaliseExp: non-numeric type in Negate"
internaliseExp [Char]
desc (E.Not Exp
e SrcLoc
_) = do
SubExp
e' <- [Char] -> Exp -> InternaliseM SubExp
internaliseExp1 [Char]
"not_arg" Exp
e
TypeBase Shape NoUniqueness
et <- forall t (m :: * -> *).
HasScope t m =>
SubExp -> m (TypeBase Shape NoUniqueness)
subExpType SubExp
e'
case TypeBase Shape NoUniqueness
et of
I.Prim (I.IntType IntType
t) ->
forall (m :: * -> *).
MonadBuilder m =>
[Char] -> Exp (Rep m) -> m [SubExp]
letTupExp' [Char]
desc forall a b. (a -> b) -> a -> b
$ forall rep. BasicOp -> Exp rep
I.BasicOp forall a b. (a -> b) -> a -> b
$ UnOp -> SubExp -> BasicOp
I.UnOp (IntType -> UnOp
I.Complement IntType
t) SubExp
e'
I.Prim PrimType
I.Bool ->
forall (m :: * -> *).
MonadBuilder m =>
[Char] -> Exp (Rep m) -> m [SubExp]
letTupExp' [Char]
desc forall a b. (a -> b) -> a -> b
$ forall rep. BasicOp -> Exp rep
I.BasicOp forall a b. (a -> b) -> a -> b
$ UnOp -> SubExp -> BasicOp
I.UnOp UnOp
I.Not SubExp
e'
TypeBase Shape NoUniqueness
_ ->
forall a. HasCallStack => [Char] -> a
error [Char]
"Futhark.Internalise.internaliseExp: non-int/bool type in Not"
internaliseExp [Char]
desc (E.Update Exp
src SliceBase Info VName
slice Exp
ve SrcLoc
loc) = do
[SubExp]
ves <- [Char] -> Exp -> InternaliseM [SubExp]
internaliseExp [Char]
"lw_val" Exp
ve
[VName]
srcs <- [Char] -> Exp -> InternaliseM [VName]
internaliseExpToVars [Char]
"src" Exp
src
[SubExp]
dims <- case [VName]
srcs of
[] -> forall (f :: * -> *) a. Applicative f => a -> f a
pure []
VName
v : [VName]
_ -> forall u. TypeBase Shape u -> [SubExp]
I.arrayDims forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> forall rep (m :: * -> *).
HasScope rep m =>
VName -> m (TypeBase Shape NoUniqueness)
lookupType VName
v
([DimIndex SubExp]
idxs', Certs
cs) <- SrcLoc
-> [SubExp]
-> SliceBase Info VName
-> InternaliseM ([DimIndex SubExp], Certs)
internaliseSlice SrcLoc
loc [SubExp]
dims SliceBase Info VName
slice
let comb :: VName -> SubExp -> InternaliseM VName
comb VName
sname SubExp
ve' = do
TypeBase Shape NoUniqueness
sname_t <- forall rep (m :: * -> *).
HasScope rep m =>
VName -> m (TypeBase Shape NoUniqueness)
lookupType VName
sname
let full_slice :: Slice SubExp
full_slice = TypeBase Shape NoUniqueness -> [DimIndex SubExp] -> Slice SubExp
fullSlice TypeBase Shape NoUniqueness
sname_t [DimIndex SubExp]
idxs'
rowtype :: TypeBase Shape NoUniqueness
rowtype = TypeBase Shape NoUniqueness
sname_t forall oldshape u.
TypeBase oldshape u -> [SubExp] -> TypeBase Shape u
`setArrayDims` forall d. Slice d -> [d]
sliceDims Slice SubExp
full_slice
SubExp
ve'' <-
ErrorMsg SubExp
-> SrcLoc
-> TypeBase Shape NoUniqueness
-> [Char]
-> SubExp
-> InternaliseM SubExp
ensureShape
ErrorMsg SubExp
"shape of value does not match shape of source array"
SrcLoc
loc
TypeBase Shape NoUniqueness
rowtype
[Char]
"lw_val_correct_shape"
SubExp
ve'
forall (m :: * -> *).
MonadBuilder m =>
[Char] -> VName -> Slice SubExp -> Exp (Rep m) -> m VName
letInPlace [Char]
desc VName
sname Slice SubExp
full_slice forall a b. (a -> b) -> a -> b
$ forall rep. BasicOp -> Exp rep
BasicOp forall a b. (a -> b) -> a -> b
$ SubExp -> BasicOp
SubExp SubExp
ve''
forall (m :: * -> *) a. MonadBuilder m => Certs -> m a -> m a
certifying Certs
cs forall a b. (a -> b) -> a -> b
$ forall a b. (a -> b) -> [a] -> [b]
map VName -> SubExp
I.Var forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> forall (m :: * -> *) a b c.
Applicative m =>
(a -> b -> m c) -> [a] -> [b] -> m [c]
zipWithM VName -> SubExp -> InternaliseM VName
comb [VName]
srcs [SubExp]
ves
internaliseExp [Char]
desc (E.RecordUpdate Exp
src [Name]
fields Exp
ve Info PatType
_ SrcLoc
_) = do
[SubExp]
src' <- [Char] -> Exp -> InternaliseM [SubExp]
internaliseExp [Char]
desc Exp
src
[SubExp]
ve' <- [Char] -> Exp -> InternaliseM [SubExp]
internaliseExp [Char]
desc Exp
ve
forall {m :: * -> *} {als} {a}.
Monad m =>
TypeBase Size als -> [Name] -> [a] -> [a] -> m [a]
replace (Exp -> PatType
E.typeOf Exp
src forall dim asf ast. TypeBase dim asf -> ast -> TypeBase dim ast
`setAliases` ()) [Name]
fields [SubExp]
ve' [SubExp]
src'
where
replace :: TypeBase Size als -> [Name] -> [a] -> [a] -> m [a]
replace (E.Scalar (E.Record Map Name (TypeBase Size als)
m)) (Name
f : [Name]
fs) [a]
ve' [a]
src'
| Just TypeBase Size als
t <- forall k a. Ord k => k -> Map k a -> Maybe a
M.lookup Name
f Map Name (TypeBase Size als)
m = do
let i :: Int
i =
forall (t :: * -> *) a. (Foldable t, Num a) => t a -> a
sum forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall a b. (a -> b) -> [a] -> [b]
map (forall als. TypeBase Size als -> Int
internalisedTypeSize forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall a b. (a, b) -> b
snd) forall a b. (a -> b) -> a -> b
$
forall a. (a -> Bool) -> [a] -> [a]
takeWhile ((forall a. Eq a => a -> a -> Bool
/= Name
f) forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall a b. (a, b) -> a
fst) forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall a. Map Name a -> [(Name, a)]
sortFields forall a b. (a -> b) -> a -> b
$
Map Name (TypeBase Size als)
m
k :: Int
k = forall als. TypeBase Size als -> Int
internalisedTypeSize TypeBase Size als
t
([a]
bef, [a]
to_update, [a]
aft) = forall a. Int -> Int -> [a] -> ([a], [a], [a])
splitAt3 Int
i Int
k [a]
src'
[a]
src'' <- TypeBase Size als -> [Name] -> [a] -> [a] -> m [a]
replace TypeBase Size als
t [Name]
fs [a]
ve' [a]
to_update
forall (f :: * -> *) a. Applicative f => a -> f a
pure forall a b. (a -> b) -> a -> b
$ [a]
bef forall a. [a] -> [a] -> [a]
++ [a]
src'' forall a. [a] -> [a] -> [a]
++ [a]
aft
replace TypeBase Size als
_ [Name]
_ [a]
ve' [a]
_ = forall (f :: * -> *) a. Applicative f => a -> f a
pure [a]
ve'
internaliseExp [Char]
desc (E.Attr AttrInfo VName
attr Exp
e SrcLoc
loc) = do
Attr
attr' <- AttrInfo VName -> InternaliseM Attr
internaliseAttr AttrInfo VName
attr
[SubExp]
e' <- forall r (m :: * -> *) a. MonadReader r m => (r -> r) -> m a -> m a
local (Attr -> InternaliseEnv -> InternaliseEnv
f Attr
attr') forall a b. (a -> b) -> a -> b
$ [Char] -> Exp -> InternaliseM [SubExp]
internaliseExp [Char]
desc Exp
e
case Attr
attr' of
Attr
"trace" ->
Text -> [SubExp] -> InternaliseM [SubExp]
traceRes ([Char] -> Text
T.pack forall a b. (a -> b) -> a -> b
$ forall a. Located a => a -> [Char]
locStr SrcLoc
loc) [SubExp]
e'
I.AttrComp Name
"trace" [I.AttrName Name
tag] ->
Text -> [SubExp] -> InternaliseM [SubExp]
traceRes (Name -> Text
nameToText Name
tag) [SubExp]
e'
Attr
"opaque" ->
forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
(a -> m b) -> t a -> m (t b)
mapM (forall (m :: * -> *).
MonadBuilder m =>
[Char] -> Exp (Rep m) -> m SubExp
letSubExp [Char]
desc forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall rep. BasicOp -> Exp rep
BasicOp forall b c a. (b -> c) -> (a -> b) -> a -> c
. OpaqueOp -> SubExp -> BasicOp
Opaque OpaqueOp
OpaqueNil) [SubExp]
e'
Attr
_ ->
forall (f :: * -> *) a. Applicative f => a -> f a
pure [SubExp]
e'
where
traceRes :: Text -> [SubExp] -> InternaliseM [SubExp]
traceRes Text
tag' =
forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
(a -> m b) -> t a -> m (t b)
mapM (forall (m :: * -> *).
MonadBuilder m =>
[Char] -> Exp (Rep m) -> m SubExp
letSubExp [Char]
desc forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall rep. BasicOp -> Exp rep
BasicOp forall b c a. (b -> c) -> (a -> b) -> a -> c
. OpaqueOp -> SubExp -> BasicOp
Opaque (Text -> OpaqueOp
OpaqueTrace Text
tag'))
f :: Attr -> InternaliseEnv -> InternaliseEnv
f Attr
attr' InternaliseEnv
env
| Attr
attr' forall a. Eq a => a -> a -> Bool
== Attr
"unsafe",
Bool -> Bool
not forall a b. (a -> b) -> a -> b
$ InternaliseEnv -> Bool
envSafe InternaliseEnv
env =
InternaliseEnv
env {envDoBoundsChecks :: Bool
envDoBoundsChecks = Bool
False}
| Bool
otherwise =
InternaliseEnv
env {envAttrs :: Attrs
envAttrs = InternaliseEnv -> Attrs
envAttrs InternaliseEnv
env forall a. Semigroup a => a -> a -> a
<> Attr -> Attrs
oneAttr Attr
attr'}
internaliseExp [Char]
desc (E.Assert Exp
e1 Exp
e2 (Info Text
check) SrcLoc
loc) = do
SubExp
e1' <- [Char] -> Exp -> InternaliseM SubExp
internaliseExp1 [Char]
"assert_cond" Exp
e1
Certs
c <- [Char] -> SubExp -> ErrorMsg SubExp -> SrcLoc -> InternaliseM Certs
assert [Char]
"assert_c" SubExp
e1' (forall a. [ErrorMsgPart a] -> ErrorMsg a
errorMsg [forall a. Text -> ErrorMsgPart a
ErrorString forall a b. (a -> b) -> a -> b
$ Text
"Assertion is false: " forall a. Semigroup a => a -> a -> a
<> Text
check]) SrcLoc
loc
forall (m :: * -> *) a. MonadBuilder m => Certs -> m a -> m a
certifying Certs
c forall a b. (a -> b) -> a -> b
$ forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
(a -> m b) -> t a -> m (t b)
mapM forall {m :: * -> *}. MonadBuilder m => SubExp -> m SubExp
rebind forall (m :: * -> *) a b. Monad m => (a -> m b) -> m a -> m b
=<< [Char] -> Exp -> InternaliseM [SubExp]
internaliseExp [Char]
desc Exp
e2
where
rebind :: SubExp -> m SubExp
rebind SubExp
v = do
VName
v' <- forall (m :: * -> *). MonadFreshNames m => [Char] -> m VName
newVName [Char]
"assert_res"
forall (m :: * -> *).
MonadBuilder m =>
[VName] -> Exp (Rep m) -> m ()
letBindNames [VName
v'] forall a b. (a -> b) -> a -> b
$ forall rep. BasicOp -> Exp rep
I.BasicOp forall a b. (a -> b) -> a -> b
$ SubExp -> BasicOp
I.SubExp SubExp
v
forall (f :: * -> *) a. Applicative f => a -> f a
pure forall a b. (a -> b) -> a -> b
$ VName -> SubExp
I.Var VName
v'
internaliseExp [Char]
_ (E.Constr Name
c [Exp]
es (Info (E.Scalar (E.Sum Map Name [PatType]
fs))) SrcLoc
_) = do
([TypeBase ExtShape Uniqueness]
ts, Map Name (Int, [Int])
constr_map) <- Map Name [TypeBase Size ()]
-> InternaliseM
([TypeBase ExtShape Uniqueness], Map Name (Int, [Int]))
internaliseSumType forall a b. (a -> b) -> a -> b
$ forall a b k. (a -> b) -> Map k a -> Map k b
M.map (forall a b. (a -> b) -> [a] -> [b]
map forall dim as. TypeBase dim as -> TypeBase dim ()
E.toStruct) Map Name [PatType]
fs
[SubExp]
es' <- forall (t :: * -> *) a. Foldable t => t [a] -> [a]
concat forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
(a -> m b) -> t a -> m (t b)
mapM ([Char] -> Exp -> InternaliseM [SubExp]
internaliseExp [Char]
"payload") [Exp]
es
let noExt :: p -> f SubExp
noExt p
_ = forall (f :: * -> *) a. Applicative f => a -> f a
pure forall a b. (a -> b) -> a -> b
$ IntType -> Integer -> SubExp
intConst IntType
Int64 Integer
0
[TypeBase Shape NoUniqueness]
ts' <- forall (m :: * -> *) u.
Monad m =>
(Int -> m SubExp) -> [TypeBase ExtShape u] -> m [TypeBase Shape u]
instantiateShapes forall {f :: * -> *} {p}. Applicative f => p -> f SubExp
noExt forall a b. (a -> b) -> a -> b
$ forall a b. (a -> b) -> [a] -> [b]
map forall shape.
TypeBase shape Uniqueness -> TypeBase shape NoUniqueness
fromDecl [TypeBase ExtShape Uniqueness]
ts
case forall k a. Ord k => k -> Map k a -> Maybe a
M.lookup Name
c Map Name (Int, [Int])
constr_map of
Just (Int
i, [Int]
js) ->
(IntType -> Integer -> SubExp
intConst IntType
Int8 (forall a. Integral a => a -> Integer
toInteger Int
i) :) forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> forall {f :: * -> *} {t}.
(Num t, MonadBuilder f, Eq t) =>
t -> [TypeBase Shape NoUniqueness] -> [(t, SubExp)] -> f [SubExp]
clauses Int
0 [TypeBase Shape NoUniqueness]
ts' (forall a b. [a] -> [b] -> [(a, b)]
zip [Int]
js [SubExp]
es')
Maybe (Int, [Int])
Nothing ->
forall a. HasCallStack => [Char] -> a
error [Char]
"internaliseExp Constr: missing constructor"
where
clauses :: t -> [TypeBase Shape NoUniqueness] -> [(t, SubExp)] -> f [SubExp]
clauses t
j (TypeBase Shape NoUniqueness
t : [TypeBase Shape NoUniqueness]
ts) [(t, SubExp)]
js_to_es
| Just SubExp
e <- t
j forall a b. Eq a => a -> [(a, b)] -> Maybe b
`lookup` [(t, SubExp)]
js_to_es =
(SubExp
e :) forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> t -> [TypeBase Shape NoUniqueness] -> [(t, SubExp)] -> f [SubExp]
clauses (t
j forall a. Num a => a -> a -> a
+ t
1) [TypeBase Shape NoUniqueness]
ts [(t, SubExp)]
js_to_es
| Bool
otherwise = do
SubExp
blank <- forall (m :: * -> *).
MonadBuilder m =>
[Char] -> Exp (Rep m) -> m SubExp
letSubExp [Char]
"zero" forall (m :: * -> *) a b. Monad m => (a -> m b) -> m a -> m b
=<< forall (m :: * -> *).
MonadBuilder m =>
TypeBase Shape NoUniqueness -> m (Exp (Rep m))
eBlank TypeBase Shape NoUniqueness
t
(SubExp
blank :) forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> t -> [TypeBase Shape NoUniqueness] -> [(t, SubExp)] -> f [SubExp]
clauses (t
j forall a. Num a => a -> a -> a
+ t
1) [TypeBase Shape NoUniqueness]
ts [(t, SubExp)]
js_to_es
clauses t
_ [] [(t, SubExp)]
_ =
forall (f :: * -> *) a. Applicative f => a -> f a
pure []
internaliseExp [Char]
_ (E.Constr Name
_ [Exp]
_ (Info PatType
t) SrcLoc
loc) =
forall a. HasCallStack => [Char] -> a
error forall a b. (a -> b) -> a -> b
$ [Char]
"internaliseExp: constructor with type " forall a. [a] -> [a] -> [a]
++ forall a. Pretty a => a -> [Char]
prettyString PatType
t forall a. [a] -> [a] -> [a]
++ [Char]
" at " forall a. [a] -> [a] -> [a]
++ forall a. Located a => a -> [Char]
locStr SrcLoc
loc
internaliseExp [Char]
_ (E.Literal PrimValue
v SrcLoc
_) =
forall (f :: * -> *) a. Applicative f => a -> f a
pure [PrimValue -> SubExp
I.Constant forall a b. (a -> b) -> a -> b
$ PrimValue -> PrimValue
internalisePrimValue PrimValue
v]
internaliseExp [Char]
_ (E.IntLit Integer
v (Info PatType
t) SrcLoc
_) =
case PatType
t of
E.Scalar (E.Prim (E.Signed IntType
it)) ->
forall (f :: * -> *) a. Applicative f => a -> f a
pure [PrimValue -> SubExp
I.Constant forall a b. (a -> b) -> a -> b
$ IntValue -> PrimValue
I.IntValue forall a b. (a -> b) -> a -> b
$ forall int. Integral int => IntType -> int -> IntValue
intValue IntType
it Integer
v]
E.Scalar (E.Prim (E.Unsigned IntType
it)) ->
forall (f :: * -> *) a. Applicative f => a -> f a
pure [PrimValue -> SubExp
I.Constant forall a b. (a -> b) -> a -> b
$ IntValue -> PrimValue
I.IntValue forall a b. (a -> b) -> a -> b
$ forall int. Integral int => IntType -> int -> IntValue
intValue IntType
it Integer
v]
E.Scalar (E.Prim (E.FloatType FloatType
ft)) ->
forall (f :: * -> *) a. Applicative f => a -> f a
pure [PrimValue -> SubExp
I.Constant forall a b. (a -> b) -> a -> b
$ FloatValue -> PrimValue
I.FloatValue forall a b. (a -> b) -> a -> b
$ forall num. Real num => FloatType -> num -> FloatValue
floatValue FloatType
ft Integer
v]
PatType
_ -> forall a. HasCallStack => [Char] -> a
error forall a b. (a -> b) -> a -> b
$ [Char]
"internaliseExp: nonsensical type for integer literal: " forall a. [a] -> [a] -> [a]
++ forall a. Pretty a => a -> [Char]
prettyString PatType
t
internaliseExp [Char]
_ (E.FloatLit Double
v (Info PatType
t) SrcLoc
_) =
case PatType
t of
E.Scalar (E.Prim (E.FloatType FloatType
ft)) ->
forall (f :: * -> *) a. Applicative f => a -> f a
pure [PrimValue -> SubExp
I.Constant forall a b. (a -> b) -> a -> b
$ FloatValue -> PrimValue
I.FloatValue forall a b. (a -> b) -> a -> b
$ forall num. Real num => FloatType -> num -> FloatValue
floatValue FloatType
ft Double
v]
PatType
_ -> forall a. HasCallStack => [Char] -> a
error forall a b. (a -> b) -> a -> b
$ [Char]
"internaliseExp: nonsensical type for float literal: " forall a. [a] -> [a] -> [a]
++ forall a. Pretty a => a -> [Char]
prettyString PatType
t
internaliseExp [Char]
desc (E.Project Name
k Exp
e (Info PatType
rt) SrcLoc
_) = do
let i' :: Int
i' = forall (t :: * -> *) a. (Foldable t, Num a) => t a -> a
sum forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall a b. (a -> b) -> [a] -> [b]
map forall als. TypeBase Size als -> Int
internalisedTypeSize forall a b. (a -> b) -> a -> b
$
case Exp -> PatType
E.typeOf Exp
e forall dim asf ast. TypeBase dim asf -> ast -> TypeBase dim ast
`setAliases` () of
E.Scalar (Record Map Name (TypeBase Size ())
fs) ->
forall a b. (a -> b) -> [a] -> [b]
map forall a b. (a, b) -> b
snd forall a b. (a -> b) -> a -> b
$ forall a. (a -> Bool) -> [a] -> [a]
takeWhile ((forall a. Eq a => a -> a -> Bool
/= Name
k) forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall a b. (a, b) -> a
fst) forall a b. (a -> b) -> a -> b
$ forall a. Map Name a -> [(Name, a)]
sortFields Map Name (TypeBase Size ())
fs
TypeBase Size ()
t -> [TypeBase Size ()
t]
forall a. Int -> [a] -> [a]
take (forall als. TypeBase Size als -> Int
internalisedTypeSize forall a b. (a -> b) -> a -> b
$ PatType
rt forall dim asf ast. TypeBase dim asf -> ast -> TypeBase dim ast
`setAliases` ()) forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall a. Int -> [a] -> [a]
drop Int
i'
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> [Char] -> Exp -> InternaliseM [SubExp]
internaliseExp [Char]
desc Exp
e
internaliseExp [Char]
_ e :: Exp
e@E.Lambda {} =
forall a. HasCallStack => [Char] -> a
error forall a b. (a -> b) -> a -> b
$ [Char]
"internaliseExp: Unexpected lambda at " forall a. [a] -> [a] -> [a]
++ forall a. Located a => a -> [Char]
locStr (forall a. Located a => a -> SrcLoc
srclocOf Exp
e)
internaliseExp [Char]
_ e :: Exp
e@E.OpSection {} =
forall a. HasCallStack => [Char] -> a
error forall a b. (a -> b) -> a -> b
$ [Char]
"internaliseExp: Unexpected operator section at " forall a. [a] -> [a] -> [a]
++ forall a. Located a => a -> [Char]
locStr (forall a. Located a => a -> SrcLoc
srclocOf Exp
e)
internaliseExp [Char]
_ e :: Exp
e@E.OpSectionLeft {} =
forall a. HasCallStack => [Char] -> a
error forall a b. (a -> b) -> a -> b
$ [Char]
"internaliseExp: Unexpected left operator section at " forall a. [a] -> [a] -> [a]
++ forall a. Located a => a -> [Char]
locStr (forall a. Located a => a -> SrcLoc
srclocOf Exp
e)
internaliseExp [Char]
_ e :: Exp
e@E.OpSectionRight {} =
forall a. HasCallStack => [Char] -> a
error forall a b. (a -> b) -> a -> b
$ [Char]
"internaliseExp: Unexpected right operator section at " forall a. [a] -> [a] -> [a]
++ forall a. Located a => a -> [Char]
locStr (forall a. Located a => a -> SrcLoc
srclocOf Exp
e)
internaliseExp [Char]
_ e :: Exp
e@E.ProjectSection {} =
forall a. HasCallStack => [Char] -> a
error forall a b. (a -> b) -> a -> b
$ [Char]
"internaliseExp: Unexpected projection section at " forall a. [a] -> [a] -> [a]
++ forall a. Located a => a -> [Char]
locStr (forall a. Located a => a -> SrcLoc
srclocOf Exp
e)
internaliseExp [Char]
_ e :: Exp
e@E.IndexSection {} =
forall a. HasCallStack => [Char] -> a
error forall a b. (a -> b) -> a -> b
$ [Char]
"internaliseExp: Unexpected index section at " forall a. [a] -> [a] -> [a]
++ forall a. Located a => a -> [Char]
locStr (forall a. Located a => a -> SrcLoc
srclocOf Exp
e)
internaliseArg :: String -> (E.Exp, Maybe VName) -> InternaliseM [SubExp]
internaliseArg :: [Char] -> (Exp, Maybe VName) -> InternaliseM [SubExp]
internaliseArg [Char]
desc (Exp
arg, Maybe VName
argdim) = do
Scope SOACS
exists <- forall rep (m :: * -> *). HasScope rep m => m (Scope rep)
askScope
case Maybe VName
argdim of
Just VName
d | VName
d forall k a. Ord k => k -> Map k a -> Bool
`M.member` Scope SOACS
exists -> forall (f :: * -> *) a. Applicative f => a -> f a
pure [VName -> SubExp
I.Var VName
d]
Maybe VName
_ -> do
[SubExp]
arg' <- [Char] -> Exp -> InternaliseM [SubExp]
internaliseExp [Char]
desc Exp
arg
case ([SubExp]
arg', Maybe VName
argdim) of
([SubExp
se], Just VName
d) -> do
forall (m :: * -> *).
MonadBuilder m =>
[VName] -> Exp (Rep m) -> m ()
letBindNames [VName
d] forall a b. (a -> b) -> a -> b
$ forall rep. BasicOp -> Exp rep
BasicOp forall a b. (a -> b) -> a -> b
$ SubExp -> BasicOp
SubExp SubExp
se
([SubExp], Maybe VName)
_ -> forall (f :: * -> *) a. Applicative f => a -> f a
pure ()
forall (f :: * -> *) a. Applicative f => a -> f a
pure [SubExp]
arg'
internalisePatLit :: E.PatLit -> E.PatType -> I.PrimValue
internalisePatLit :: PatLit -> PatType -> PrimValue
internalisePatLit (E.PatLitPrim PrimValue
v) PatType
_ =
PrimValue -> PrimValue
internalisePrimValue PrimValue
v
internalisePatLit (E.PatLitInt Integer
x) (E.Scalar (E.Prim (E.Signed IntType
it))) =
IntValue -> PrimValue
I.IntValue forall a b. (a -> b) -> a -> b
$ forall int. Integral int => IntType -> int -> IntValue
intValue IntType
it Integer
x
internalisePatLit (E.PatLitInt Integer
x) (E.Scalar (E.Prim (E.Unsigned IntType
it))) =
IntValue -> PrimValue
I.IntValue forall a b. (a -> b) -> a -> b
$ forall int. Integral int => IntType -> int -> IntValue
intValue IntType
it Integer
x
internalisePatLit (E.PatLitFloat Double
x) (E.Scalar (E.Prim (E.FloatType FloatType
ft))) =
FloatValue -> PrimValue
I.FloatValue forall a b. (a -> b) -> a -> b
$ forall num. Real num => FloatType -> num -> FloatValue
floatValue FloatType
ft Double
x
internalisePatLit PatLit
l PatType
t =
forall a. HasCallStack => [Char] -> a
error forall a b. (a -> b) -> a -> b
$ [Char]
"Nonsensical pattern and type: " forall a. [a] -> [a] -> [a]
++ forall a. Show a => a -> [Char]
show (PatLit
l, PatType
t)
generateCond ::
E.Pat ->
[I.SubExp] ->
InternaliseM ([Maybe I.PrimValue], [I.SubExp])
generateCond :: PatBase Info VName
-> [SubExp] -> InternaliseM ([Maybe PrimValue], [SubExp])
generateCond PatBase Info VName
orig_p [SubExp]
orig_ses = do
([Maybe PrimValue]
cmps, [SubExp]
pertinent, [SubExp]
_) <- forall {vn} {a}.
(Eq vn, IsName vn) =>
PatBase Info vn
-> [a] -> InternaliseM ([Maybe PrimValue], [a], [a])
compares PatBase Info VName
orig_p [SubExp]
orig_ses
forall (f :: * -> *) a. Applicative f => a -> f a
pure ([Maybe PrimValue]
cmps, [SubExp]
pertinent)
where
compares :: PatBase Info vn
-> [a] -> InternaliseM ([Maybe PrimValue], [a], [a])
compares (E.PatLit PatLit
l (Info PatType
t) SrcLoc
_) (a
se : [a]
ses) =
forall (f :: * -> *) a. Applicative f => a -> f a
pure ([forall a. a -> Maybe a
Just forall a b. (a -> b) -> a -> b
$ PatLit -> PatType -> PrimValue
internalisePatLit PatLit
l PatType
t], [a
se], [a]
ses)
compares (E.PatConstr Name
c (Info (E.Scalar (E.Sum Map Name [PatType]
fs))) [PatBase Info vn]
pats SrcLoc
_) (a
_ : [a]
ses) = do
([TypeBase ExtShape Uniqueness]
payload_ts, Map Name (Int, [Int])
m) <- Map Name [TypeBase Size ()]
-> InternaliseM
([TypeBase ExtShape Uniqueness], Map Name (Int, [Int]))
internaliseSumType forall a b. (a -> b) -> a -> b
$ forall a b k. (a -> b) -> Map k a -> Map k b
M.map (forall a b. (a -> b) -> [a] -> [b]
map forall dim as. TypeBase dim as -> TypeBase dim ()
toStruct) Map Name [PatType]
fs
case forall k a. Ord k => k -> Map k a -> Maybe a
M.lookup Name
c Map Name (Int, [Int])
m of
Just (Int
tag, [Int]
payload_is) -> do
let ([a]
payload_ses, [a]
ses') = forall a. Int -> [a] -> ([a], [a])
splitAt (forall (t :: * -> *) a. Foldable t => t a -> Int
length [TypeBase ExtShape Uniqueness]
payload_ts) [a]
ses
([Maybe PrimValue]
cmps, [a]
pertinent, [a]
_) <-
[PatBase Info vn]
-> [a] -> InternaliseM ([Maybe PrimValue], [a], [a])
comparesMany [PatBase Info vn]
pats forall a b. (a -> b) -> a -> b
$ forall a b. (a -> b) -> [a] -> [b]
map ([a]
payload_ses !!) [Int]
payload_is
let missingCmps :: Int -> a -> Maybe PrimValue
missingCmps Int
i a
_ =
case Int
i forall a. Eq a => a -> [a] -> Maybe Int
`elemIndex` [Int]
payload_is of
Just Int
j -> [Maybe PrimValue]
cmps forall a. [a] -> Int -> a
!! Int
j
Maybe Int
Nothing -> forall a. Maybe a
Nothing
forall (f :: * -> *) a. Applicative f => a -> f a
pure
( forall a. a -> Maybe a
Just (IntValue -> PrimValue
I.IntValue forall a b. (a -> b) -> a -> b
$ forall int. Integral int => IntType -> int -> IntValue
intValue IntType
Int8 forall a b. (a -> b) -> a -> b
$ forall a. Integral a => a -> Integer
toInteger Int
tag)
forall a. a -> [a] -> [a]
: forall a b c. (a -> b -> c) -> [a] -> [b] -> [c]
zipWith Int -> a -> Maybe PrimValue
missingCmps [Int
0 ..] [a]
payload_ses,
[a]
pertinent,
[a]
ses'
)
Maybe (Int, [Int])
Nothing ->
forall a. HasCallStack => [Char] -> a
error [Char]
"generateCond: missing constructor"
compares (E.PatConstr Name
_ (Info PatType
t) [PatBase Info vn]
_ SrcLoc
_) [a]
_ =
forall a. HasCallStack => [Char] -> a
error forall a b. (a -> b) -> a -> b
$ [Char]
"generateCond: PatConstr has nonsensical type: " forall a. [a] -> [a] -> [a]
++ forall a. Pretty a => a -> [Char]
prettyString PatType
t
compares (E.Id vn
_ Info PatType
t SrcLoc
loc) [a]
ses =
PatBase Info vn
-> [a] -> InternaliseM ([Maybe PrimValue], [a], [a])
compares (forall (f :: * -> *) vn. f PatType -> SrcLoc -> PatBase f vn
E.Wildcard Info PatType
t SrcLoc
loc) [a]
ses
compares (E.Wildcard (Info PatType
t) SrcLoc
_) [a]
ses = do
let ([a]
id_ses, [a]
rest_ses) = forall a. Int -> [a] -> ([a], [a])
splitAt (forall als. TypeBase Size als -> Int
internalisedTypeSize forall a b. (a -> b) -> a -> b
$ forall dim as. TypeBase dim as -> TypeBase dim ()
E.toStruct PatType
t) [a]
ses
forall (f :: * -> *) a. Applicative f => a -> f a
pure (forall a b. (a -> b) -> [a] -> [b]
map (forall a b. a -> b -> a
const forall a. Maybe a
Nothing) [a]
id_ses, [a]
id_ses, [a]
rest_ses)
compares (E.PatParens PatBase Info vn
pat SrcLoc
_) [a]
ses =
PatBase Info vn
-> [a] -> InternaliseM ([Maybe PrimValue], [a], [a])
compares PatBase Info vn
pat [a]
ses
compares (E.PatAttr AttrInfo vn
_ PatBase Info vn
pat SrcLoc
_) [a]
ses =
PatBase Info vn
-> [a] -> InternaliseM ([Maybe PrimValue], [a], [a])
compares PatBase Info vn
pat [a]
ses
compares (E.TuplePat [] SrcLoc
loc) [a]
ses =
PatBase Info vn
-> [a] -> InternaliseM ([Maybe PrimValue], [a], [a])
compares (forall (f :: * -> *) vn. f PatType -> SrcLoc -> PatBase f vn
E.Wildcard (forall a. a -> Info a
Info forall a b. (a -> b) -> a -> b
$ forall dim as. ScalarTypeBase dim as -> TypeBase dim as
E.Scalar forall a b. (a -> b) -> a -> b
$ forall dim as. Map Name (TypeBase dim as) -> ScalarTypeBase dim as
E.Record forall a. Monoid a => a
mempty) SrcLoc
loc) [a]
ses
compares (E.RecordPat [] SrcLoc
loc) [a]
ses =
PatBase Info vn
-> [a] -> InternaliseM ([Maybe PrimValue], [a], [a])
compares (forall (f :: * -> *) vn. f PatType -> SrcLoc -> PatBase f vn
E.Wildcard (forall a. a -> Info a
Info forall a b. (a -> b) -> a -> b
$ forall dim as. ScalarTypeBase dim as -> TypeBase dim as
E.Scalar forall a b. (a -> b) -> a -> b
$ forall dim as. Map Name (TypeBase dim as) -> ScalarTypeBase dim as
E.Record forall a. Monoid a => a
mempty) SrcLoc
loc) [a]
ses
compares (E.TuplePat [PatBase Info vn]
pats SrcLoc
_) [a]
ses =
[PatBase Info vn]
-> [a] -> InternaliseM ([Maybe PrimValue], [a], [a])
comparesMany [PatBase Info vn]
pats [a]
ses
compares (E.RecordPat [(Name, PatBase Info vn)]
fs SrcLoc
_) [a]
ses =
[PatBase Info vn]
-> [a] -> InternaliseM ([Maybe PrimValue], [a], [a])
comparesMany (forall a b. (a -> b) -> [a] -> [b]
map forall a b. (a, b) -> b
snd forall a b. (a -> b) -> a -> b
$ forall a. Map Name a -> [(Name, a)]
E.sortFields forall a b. (a -> b) -> a -> b
$ forall k a. Ord k => [(k, a)] -> Map k a
M.fromList [(Name, PatBase Info vn)]
fs) [a]
ses
compares (E.PatAscription PatBase Info vn
pat TypeExp vn
_ SrcLoc
_) [a]
ses =
PatBase Info vn
-> [a] -> InternaliseM ([Maybe PrimValue], [a], [a])
compares PatBase Info vn
pat [a]
ses
compares PatBase Info vn
pat [] =
forall a. HasCallStack => [Char] -> a
error forall a b. (a -> b) -> a -> b
$ [Char]
"generateCond: No values left for pattern " forall a. [a] -> [a] -> [a]
++ forall a. Pretty a => a -> [Char]
prettyString PatBase Info vn
pat
comparesMany :: [PatBase Info vn]
-> [a] -> InternaliseM ([Maybe PrimValue], [a], [a])
comparesMany [] [a]
ses = forall (f :: * -> *) a. Applicative f => a -> f a
pure ([], [], [a]
ses)
comparesMany (PatBase Info vn
pat : [PatBase Info vn]
pats) [a]
ses = do
([Maybe PrimValue]
cmps1, [a]
pertinent1, [a]
ses') <- PatBase Info vn
-> [a] -> InternaliseM ([Maybe PrimValue], [a], [a])
compares PatBase Info vn
pat [a]
ses
([Maybe PrimValue]
cmps2, [a]
pertinent2, [a]
ses'') <- [PatBase Info vn]
-> [a] -> InternaliseM ([Maybe PrimValue], [a], [a])
comparesMany [PatBase Info vn]
pats [a]
ses'
forall (f :: * -> *) a. Applicative f => a -> f a
pure
( [Maybe PrimValue]
cmps1 forall a. Semigroup a => a -> a -> a
<> [Maybe PrimValue]
cmps2,
[a]
pertinent1 forall a. Semigroup a => a -> a -> a
<> [a]
pertinent2,
[a]
ses''
)
internalisePat ::
String ->
[E.SizeBinder VName] ->
E.Pat ->
E.Exp ->
InternaliseM a ->
InternaliseM a
internalisePat :: forall a.
[Char]
-> [SizeBinder VName]
-> PatBase Info VName
-> Exp
-> InternaliseM a
-> InternaliseM a
internalisePat [Char]
desc [SizeBinder VName]
sizes PatBase Info VName
p Exp
e InternaliseM a
m = do
[SubExp]
ses <- [Char] -> Exp -> InternaliseM [SubExp]
internaliseExp [Char]
desc' Exp
e
forall a.
[SizeBinder VName]
-> PatBase Info VName
-> [SubExp]
-> InternaliseM a
-> InternaliseM a
internalisePat' [SizeBinder VName]
sizes PatBase Info VName
p [SubExp]
ses InternaliseM a
m
where
desc' :: [Char]
desc' = case forall a. Set a -> [a]
S.toList forall a b. (a -> b) -> a -> b
$ forall (f :: * -> *) vn.
(Functor f, Ord vn) =>
PatBase f vn -> Set (IdentBase f vn)
E.patIdents PatBase Info VName
p of
[IdentBase Info VName
v] -> VName -> [Char]
baseString forall a b. (a -> b) -> a -> b
$ forall (f :: * -> *) vn. IdentBase f vn -> vn
E.identName IdentBase Info VName
v
[IdentBase Info VName]
_ -> [Char]
desc
internalisePat' ::
[E.SizeBinder VName] ->
E.Pat ->
[I.SubExp] ->
InternaliseM a ->
InternaliseM a
internalisePat' :: forall a.
[SizeBinder VName]
-> PatBase Info VName
-> [SubExp]
-> InternaliseM a
-> InternaliseM a
internalisePat' [SizeBinder VName]
sizes PatBase Info VName
p [SubExp]
ses InternaliseM a
m = do
[TypeBase Shape NoUniqueness]
ses_ts <- forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
(a -> m b) -> t a -> m (t b)
mapM forall t (m :: * -> *).
HasScope t m =>
SubExp -> m (TypeBase Shape NoUniqueness)
subExpType [SubExp]
ses
forall a.
PatBase Info VName
-> [TypeBase Shape NoUniqueness]
-> ([VName] -> InternaliseM a)
-> InternaliseM a
stmPat PatBase Info VName
p [TypeBase Shape NoUniqueness]
ses_ts forall a b. (a -> b) -> a -> b
$ \[VName]
pat_names -> do
AppRes -> [SubExp] -> InternaliseM ()
bindExtSizes (PatType -> [VName] -> AppRes
AppRes (PatBase Info VName -> PatType
E.patternType PatBase Info VName
p) (forall a b. (a -> b) -> [a] -> [b]
map forall vn. SizeBinder vn -> vn
E.sizeName [SizeBinder VName]
sizes)) [SubExp]
ses
forall (t :: * -> *) (m :: * -> *) a b.
(Foldable t, Monad m) =>
t a -> (a -> m b) -> m ()
forM_ (forall a b. [a] -> [b] -> [(a, b)]
zip [VName]
pat_names [SubExp]
ses) forall a b. (a -> b) -> a -> b
$ \(VName
v, SubExp
se) ->
forall (m :: * -> *).
MonadBuilder m =>
[VName] -> Exp (Rep m) -> m ()
letBindNames [VName
v] forall a b. (a -> b) -> a -> b
$ forall rep. BasicOp -> Exp rep
I.BasicOp forall a b. (a -> b) -> a -> b
$ SubExp -> BasicOp
I.SubExp SubExp
se
InternaliseM a
m
internaliseSlice ::
SrcLoc ->
[SubExp] ->
[E.DimIndex] ->
InternaliseM ([I.DimIndex SubExp], Certs)
internaliseSlice :: SrcLoc
-> [SubExp]
-> SliceBase Info VName
-> InternaliseM ([DimIndex SubExp], Certs)
internaliseSlice SrcLoc
loc [SubExp]
dims SliceBase Info VName
idxs = do
([DimIndex SubExp]
idxs', [SubExp]
oks, [[ErrorMsgPart SubExp]]
parts) <- forall a b c. [(a, b, c)] -> ([a], [b], [c])
unzip3 forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> forall (m :: * -> *) a b c.
Applicative m =>
(a -> b -> m c) -> [a] -> [b] -> m [c]
zipWithM SubExp
-> DimIndex
-> InternaliseM (DimIndex SubExp, SubExp, [ErrorMsgPart SubExp])
internaliseDimIndex [SubExp]
dims SliceBase Info VName
idxs
SubExp
ok <- forall (m :: * -> *).
MonadBuilder m =>
[Char] -> Exp (Rep m) -> m SubExp
letSubExp [Char]
"index_ok" forall (m :: * -> *) a b. Monad m => (a -> m b) -> m a -> m b
=<< forall (m :: * -> *). MonadBuilder m => [SubExp] -> m (Exp (Rep m))
eAll [SubExp]
oks
let msg :: ErrorMsg SubExp
msg =
forall a. [ErrorMsgPart a] -> ErrorMsg a
errorMsg forall a b. (a -> b) -> a -> b
$
[ErrorMsgPart SubExp
"Index ["]
forall a. [a] -> [a] -> [a]
++ forall a. [a] -> [[a]] -> [a]
intercalate [ErrorMsgPart SubExp
", "] [[ErrorMsgPart SubExp]]
parts
forall a. [a] -> [a] -> [a]
++ [ErrorMsgPart SubExp
"] out of bounds for array of shape ["]
forall a. [a] -> [a] -> [a]
++ forall a. a -> [a] -> [a]
intersperse ErrorMsgPart SubExp
"][" (forall a b. (a -> b) -> [a] -> [b]
map (forall a. PrimType -> a -> ErrorMsgPart a
ErrorVal PrimType
int64) forall a b. (a -> b) -> a -> b
$ forall a. Int -> [a] -> [a]
take (forall (t :: * -> *) a. Foldable t => t a -> Int
length SliceBase Info VName
idxs) [SubExp]
dims)
forall a. [a] -> [a] -> [a]
++ [ErrorMsgPart SubExp
"]."]
Certs
c <- [Char] -> SubExp -> ErrorMsg SubExp -> SrcLoc -> InternaliseM Certs
assert [Char]
"index_certs" SubExp
ok ErrorMsg SubExp
msg SrcLoc
loc
forall (f :: * -> *) a. Applicative f => a -> f a
pure ([DimIndex SubExp]
idxs', Certs
c)
internaliseDimIndex ::
SubExp ->
E.DimIndex ->
InternaliseM (I.DimIndex SubExp, SubExp, [ErrorMsgPart SubExp])
internaliseDimIndex :: SubExp
-> DimIndex
-> InternaliseM (DimIndex SubExp, SubExp, [ErrorMsgPart SubExp])
internaliseDimIndex SubExp
w (E.DimFix Exp
i) = do
(SubExp
i', IntType
_) <- [Char] -> Exp -> InternaliseM (SubExp, IntType)
internaliseSizeExp [Char]
"i" Exp
i
let lowerBound :: Exp SOACS
lowerBound =
forall rep. BasicOp -> Exp rep
I.BasicOp forall a b. (a -> b) -> a -> b
$ CmpOp -> SubExp -> SubExp -> BasicOp
I.CmpOp (IntType -> CmpOp
I.CmpSle IntType
I.Int64) (forall v. IsValue v => v -> SubExp
I.constant (Int64
0 :: I.Int64)) SubExp
i'
upperBound :: Exp SOACS
upperBound =
forall rep. BasicOp -> Exp rep
I.BasicOp forall a b. (a -> b) -> a -> b
$ CmpOp -> SubExp -> SubExp -> BasicOp
I.CmpOp (IntType -> CmpOp
I.CmpSlt IntType
I.Int64) SubExp
i' SubExp
w
SubExp
ok <- forall (m :: * -> *).
MonadBuilder m =>
[Char] -> Exp (Rep m) -> m SubExp
letSubExp [Char]
"bounds_check" forall (m :: * -> *) a b. Monad m => (a -> m b) -> m a -> m b
=<< forall (m :: * -> *).
MonadBuilder m =>
BinOp -> m (Exp (Rep m)) -> m (Exp (Rep m)) -> m (Exp (Rep m))
eBinOp BinOp
I.LogAnd (forall (f :: * -> *) a. Applicative f => a -> f a
pure Exp SOACS
lowerBound) (forall (f :: * -> *) a. Applicative f => a -> f a
pure Exp SOACS
upperBound)
forall (f :: * -> *) a. Applicative f => a -> f a
pure (forall d. d -> DimIndex d
I.DimFix SubExp
i', SubExp
ok, [forall a. PrimType -> a -> ErrorMsgPart a
ErrorVal PrimType
int64 SubExp
i'])
internaliseDimIndex
SubExp
w
( E.DimSlice
Maybe Exp
Nothing
Maybe Exp
Nothing
(Just (E.Negate (E.IntLit Integer
1 Info PatType
_ SrcLoc
_) SrcLoc
_))
) = do
SubExp
w_minus_1 <-
forall (m :: * -> *).
MonadBuilder m =>
[Char] -> Exp (Rep m) -> m SubExp
letSubExp [Char]
"w_minus_1" forall a b. (a -> b) -> a -> b
$
forall rep. BasicOp -> Exp rep
BasicOp forall a b. (a -> b) -> a -> b
$
BinOp -> SubExp -> SubExp -> BasicOp
I.BinOp (IntType -> Overflow -> BinOp
Sub IntType
Int64 Overflow
I.OverflowWrap) SubExp
w SubExp
one
forall (f :: * -> *) a. Applicative f => a -> f a
pure
( forall d. d -> d -> d -> DimIndex d
I.DimSlice SubExp
w_minus_1 SubExp
w forall a b. (a -> b) -> a -> b
$ IntType -> Integer -> SubExp
intConst IntType
Int64 (-Integer
1),
forall v. IsValue v => v -> SubExp
constant Bool
True,
forall a. Monoid a => a
mempty
)
where
one :: SubExp
one = forall v. IsValue v => v -> SubExp
constant (Int64
1 :: Int64)
internaliseDimIndex SubExp
w (E.DimSlice Maybe Exp
i Maybe Exp
j Maybe Exp
s) = do
SubExp
s' <- forall b a. b -> (a -> b) -> Maybe a -> b
maybe (forall (f :: * -> *) a. Applicative f => a -> f a
pure SubExp
one) (forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap forall a b. (a, b) -> a
fst forall b c a. (b -> c) -> (a -> b) -> a -> c
. [Char] -> Exp -> InternaliseM (SubExp, IntType)
internaliseSizeExp [Char]
"s") Maybe Exp
s
SubExp
s_sign <- forall (m :: * -> *).
MonadBuilder m =>
[Char] -> Exp (Rep m) -> m SubExp
letSubExp [Char]
"s_sign" forall a b. (a -> b) -> a -> b
$ forall rep. BasicOp -> Exp rep
BasicOp forall a b. (a -> b) -> a -> b
$ UnOp -> SubExp -> BasicOp
I.UnOp (IntType -> UnOp
I.SSignum IntType
Int64) SubExp
s'
SubExp
backwards <- forall (m :: * -> *).
MonadBuilder m =>
[Char] -> Exp (Rep m) -> m SubExp
letSubExp [Char]
"backwards" forall a b. (a -> b) -> a -> b
$ forall rep. BasicOp -> Exp rep
I.BasicOp forall a b. (a -> b) -> a -> b
$ CmpOp -> SubExp -> SubExp -> BasicOp
I.CmpOp (PrimType -> CmpOp
I.CmpEq PrimType
int64) SubExp
s_sign SubExp
negone
SubExp
w_minus_1 <- forall (m :: * -> *).
MonadBuilder m =>
[Char] -> Exp (Rep m) -> m SubExp
letSubExp [Char]
"w_minus_1" forall a b. (a -> b) -> a -> b
$ forall rep. BasicOp -> Exp rep
BasicOp forall a b. (a -> b) -> a -> b
$ BinOp -> SubExp -> SubExp -> BasicOp
I.BinOp (IntType -> Overflow -> BinOp
Sub IntType
Int64 Overflow
I.OverflowWrap) SubExp
w SubExp
one
let i_def :: InternaliseM SubExp
i_def =
forall (m :: * -> *).
MonadBuilder m =>
[Char] -> Exp (Rep m) -> m SubExp
letSubExp [Char]
"i_def"
forall (m :: * -> *) a b. Monad m => (a -> m b) -> m a -> m b
=<< forall (m :: * -> *).
(MonadBuilder m, BranchType (Rep m) ~ ExtType) =>
m (Exp (Rep m))
-> m (Body (Rep m)) -> m (Body (Rep m)) -> m (Exp (Rep m))
eIf
(forall (m :: * -> *). MonadBuilder m => SubExp -> m (Exp (Rep m))
eSubExp SubExp
backwards)
(forall (m :: * -> *).
MonadBuilder m =>
[SubExp] -> m (Body (Rep m))
resultBodyM [SubExp
w_minus_1])
(forall (m :: * -> *).
MonadBuilder m =>
[SubExp] -> m (Body (Rep m))
resultBodyM [SubExp
zero])
j_def :: InternaliseM SubExp
j_def =
forall (m :: * -> *).
MonadBuilder m =>
[Char] -> Exp (Rep m) -> m SubExp
letSubExp [Char]
"j_def"
forall (m :: * -> *) a b. Monad m => (a -> m b) -> m a -> m b
=<< forall (m :: * -> *).
(MonadBuilder m, BranchType (Rep m) ~ ExtType) =>
m (Exp (Rep m))
-> m (Body (Rep m)) -> m (Body (Rep m)) -> m (Exp (Rep m))
eIf
(forall (m :: * -> *). MonadBuilder m => SubExp -> m (Exp (Rep m))
eSubExp SubExp
backwards)
(forall (m :: * -> *).
MonadBuilder m =>
[SubExp] -> m (Body (Rep m))
resultBodyM [SubExp
negone])
(forall (m :: * -> *).
MonadBuilder m =>
[SubExp] -> m (Body (Rep m))
resultBodyM [SubExp
w])
SubExp
i' <- forall b a. b -> (a -> b) -> Maybe a -> b
maybe InternaliseM SubExp
i_def (forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap forall a b. (a, b) -> a
fst forall b c a. (b -> c) -> (a -> b) -> a -> c
. [Char] -> Exp -> InternaliseM (SubExp, IntType)
internaliseSizeExp [Char]
"i") Maybe Exp
i
SubExp
j' <- forall b a. b -> (a -> b) -> Maybe a -> b
maybe InternaliseM SubExp
j_def (forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap forall a b. (a, b) -> a
fst forall b c a. (b -> c) -> (a -> b) -> a -> c
. [Char] -> Exp -> InternaliseM (SubExp, IntType)
internaliseSizeExp [Char]
"j") Maybe Exp
j
SubExp
j_m_i <- forall (m :: * -> *).
MonadBuilder m =>
[Char] -> Exp (Rep m) -> m SubExp
letSubExp [Char]
"j_m_i" forall a b. (a -> b) -> a -> b
$ forall rep. BasicOp -> Exp rep
BasicOp forall a b. (a -> b) -> a -> b
$ BinOp -> SubExp -> SubExp -> BasicOp
I.BinOp (IntType -> Overflow -> BinOp
Sub IntType
Int64 Overflow
I.OverflowWrap) SubExp
j' SubExp
i'
let divRounding :: InternaliseM (Exp SOACS)
-> InternaliseM (Exp SOACS)
-> InternaliseM (Exp (Rep InternaliseM))
divRounding InternaliseM (Exp SOACS)
x InternaliseM (Exp SOACS)
y =
forall (m :: * -> *).
MonadBuilder m =>
BinOp -> m (Exp (Rep m)) -> m (Exp (Rep m)) -> m (Exp (Rep m))
eBinOp
(IntType -> Safety -> BinOp
SQuot IntType
Int64 Safety
Safe)
( forall (m :: * -> *).
MonadBuilder m =>
BinOp -> m (Exp (Rep m)) -> m (Exp (Rep m)) -> m (Exp (Rep m))
eBinOp
(IntType -> Overflow -> BinOp
Add IntType
Int64 Overflow
I.OverflowWrap)
InternaliseM (Exp SOACS)
x
(forall (m :: * -> *).
MonadBuilder m =>
BinOp -> m (Exp (Rep m)) -> m (Exp (Rep m)) -> m (Exp (Rep m))
eBinOp (IntType -> Overflow -> BinOp
Sub IntType
Int64 Overflow
I.OverflowWrap) InternaliseM (Exp SOACS)
y (forall (m :: * -> *).
MonadBuilder m =>
m (Exp (Rep m)) -> m (Exp (Rep m))
eSignum forall a b. (a -> b) -> a -> b
$ forall a (m :: * -> *).
(ToExp a, MonadBuilder m) =>
a -> m (Exp (Rep m))
toExp SubExp
s'))
)
InternaliseM (Exp SOACS)
y
SubExp
n <- forall (m :: * -> *).
MonadBuilder m =>
[Char] -> Exp (Rep m) -> m SubExp
letSubExp [Char]
"n" forall (m :: * -> *) a b. Monad m => (a -> m b) -> m a -> m b
=<< InternaliseM (Exp SOACS)
-> InternaliseM (Exp SOACS)
-> InternaliseM (Exp (Rep InternaliseM))
divRounding (forall a (m :: * -> *).
(ToExp a, MonadBuilder m) =>
a -> m (Exp (Rep m))
toExp SubExp
j_m_i) (forall a (m :: * -> *).
(ToExp a, MonadBuilder m) =>
a -> m (Exp (Rep m))
toExp SubExp
s')
SubExp
zero_stride <- forall (m :: * -> *).
MonadBuilder m =>
[Char] -> Exp (Rep m) -> m SubExp
letSubExp [Char]
"zero_stride" forall a b. (a -> b) -> a -> b
$ forall rep. BasicOp -> Exp rep
I.BasicOp forall a b. (a -> b) -> a -> b
$ CmpOp -> SubExp -> SubExp -> BasicOp
I.CmpOp (PrimType -> CmpOp
CmpEq PrimType
int64) SubExp
s_sign SubExp
zero
SubExp
nonzero_stride <- forall (m :: * -> *).
MonadBuilder m =>
[Char] -> Exp (Rep m) -> m SubExp
letSubExp [Char]
"nonzero_stride" forall a b. (a -> b) -> a -> b
$ forall rep. BasicOp -> Exp rep
I.BasicOp forall a b. (a -> b) -> a -> b
$ UnOp -> SubExp -> BasicOp
I.UnOp UnOp
I.Not SubExp
zero_stride
SubExp
empty_slice <- forall (m :: * -> *).
MonadBuilder m =>
[Char] -> Exp (Rep m) -> m SubExp
letSubExp [Char]
"empty_slice" forall a b. (a -> b) -> a -> b
$ forall rep. BasicOp -> Exp rep
I.BasicOp forall a b. (a -> b) -> a -> b
$ CmpOp -> SubExp -> SubExp -> BasicOp
I.CmpOp (PrimType -> CmpOp
CmpEq PrimType
int64) SubExp
n SubExp
zero
SubExp
m <- forall (m :: * -> *).
MonadBuilder m =>
[Char] -> Exp (Rep m) -> m SubExp
letSubExp [Char]
"m" forall a b. (a -> b) -> a -> b
$ forall rep. BasicOp -> Exp rep
I.BasicOp forall a b. (a -> b) -> a -> b
$ BinOp -> SubExp -> SubExp -> BasicOp
I.BinOp (IntType -> Overflow -> BinOp
Sub IntType
Int64 Overflow
I.OverflowWrap) SubExp
n SubExp
one
SubExp
m_t_s <- forall (m :: * -> *).
MonadBuilder m =>
[Char] -> Exp (Rep m) -> m SubExp
letSubExp [Char]
"m_t_s" forall a b. (a -> b) -> a -> b
$ forall rep. BasicOp -> Exp rep
I.BasicOp forall a b. (a -> b) -> a -> b
$ BinOp -> SubExp -> SubExp -> BasicOp
I.BinOp (IntType -> Overflow -> BinOp
Mul IntType
Int64 Overflow
I.OverflowWrap) SubExp
m SubExp
s'
SubExp
i_p_m_t_s <- forall (m :: * -> *).
MonadBuilder m =>
[Char] -> Exp (Rep m) -> m SubExp
letSubExp [Char]
"i_p_m_t_s" forall a b. (a -> b) -> a -> b
$ forall rep. BasicOp -> Exp rep
I.BasicOp forall a b. (a -> b) -> a -> b
$ BinOp -> SubExp -> SubExp -> BasicOp
I.BinOp (IntType -> Overflow -> BinOp
Add IntType
Int64 Overflow
I.OverflowWrap) SubExp
i' SubExp
m_t_s
SubExp
zero_leq_i_p_m_t_s <-
forall (m :: * -> *).
MonadBuilder m =>
[Char] -> Exp (Rep m) -> m SubExp
letSubExp [Char]
"zero_leq_i_p_m_t_s" forall a b. (a -> b) -> a -> b
$
forall rep. BasicOp -> Exp rep
I.BasicOp forall a b. (a -> b) -> a -> b
$
CmpOp -> SubExp -> SubExp -> BasicOp
I.CmpOp (IntType -> CmpOp
I.CmpSle IntType
Int64) SubExp
zero SubExp
i_p_m_t_s
SubExp
i_p_m_t_s_leq_w <-
forall (m :: * -> *).
MonadBuilder m =>
[Char] -> Exp (Rep m) -> m SubExp
letSubExp [Char]
"i_p_m_t_s_leq_w" forall a b. (a -> b) -> a -> b
$
forall rep. BasicOp -> Exp rep
I.BasicOp forall a b. (a -> b) -> a -> b
$
CmpOp -> SubExp -> SubExp -> BasicOp
I.CmpOp (IntType -> CmpOp
I.CmpSle IntType
Int64) SubExp
i_p_m_t_s SubExp
w
SubExp
i_p_m_t_s_lth_w <-
forall (m :: * -> *).
MonadBuilder m =>
[Char] -> Exp (Rep m) -> m SubExp
letSubExp [Char]
"i_p_m_t_s_leq_w" forall a b. (a -> b) -> a -> b
$
forall rep. BasicOp -> Exp rep
I.BasicOp forall a b. (a -> b) -> a -> b
$
CmpOp -> SubExp -> SubExp -> BasicOp
I.CmpOp (IntType -> CmpOp
I.CmpSlt IntType
Int64) SubExp
i_p_m_t_s SubExp
w
SubExp
zero_lte_i <- forall (m :: * -> *).
MonadBuilder m =>
[Char] -> Exp (Rep m) -> m SubExp
letSubExp [Char]
"zero_lte_i" forall a b. (a -> b) -> a -> b
$ forall rep. BasicOp -> Exp rep
I.BasicOp forall a b. (a -> b) -> a -> b
$ CmpOp -> SubExp -> SubExp -> BasicOp
I.CmpOp (IntType -> CmpOp
I.CmpSle IntType
Int64) SubExp
zero SubExp
i'
SubExp
i_lte_j <- forall (m :: * -> *).
MonadBuilder m =>
[Char] -> Exp (Rep m) -> m SubExp
letSubExp [Char]
"i_lte_j" forall a b. (a -> b) -> a -> b
$ forall rep. BasicOp -> Exp rep
I.BasicOp forall a b. (a -> b) -> a -> b
$ CmpOp -> SubExp -> SubExp -> BasicOp
I.CmpOp (IntType -> CmpOp
I.CmpSle IntType
Int64) SubExp
i' SubExp
j'
SubExp
forwards_ok <-
forall (m :: * -> *).
MonadBuilder m =>
[Char] -> Exp (Rep m) -> m SubExp
letSubExp [Char]
"forwards_ok"
forall (m :: * -> *) a b. Monad m => (a -> m b) -> m a -> m b
=<< forall (m :: * -> *). MonadBuilder m => [SubExp] -> m (Exp (Rep m))
eAll [SubExp
zero_lte_i, SubExp
zero_lte_i, SubExp
i_lte_j, SubExp
zero_leq_i_p_m_t_s, SubExp
i_p_m_t_s_lth_w]
SubExp
negone_lte_j <- forall (m :: * -> *).
MonadBuilder m =>
[Char] -> Exp (Rep m) -> m SubExp
letSubExp [Char]
"negone_lte_j" forall a b. (a -> b) -> a -> b
$ forall rep. BasicOp -> Exp rep
I.BasicOp forall a b. (a -> b) -> a -> b
$ CmpOp -> SubExp -> SubExp -> BasicOp
I.CmpOp (IntType -> CmpOp
I.CmpSle IntType
Int64) SubExp
negone SubExp
j'
SubExp
j_lte_i <- forall (m :: * -> *).
MonadBuilder m =>
[Char] -> Exp (Rep m) -> m SubExp
letSubExp [Char]
"j_lte_i" forall a b. (a -> b) -> a -> b
$ forall rep. BasicOp -> Exp rep
I.BasicOp forall a b. (a -> b) -> a -> b
$ CmpOp -> SubExp -> SubExp -> BasicOp
I.CmpOp (IntType -> CmpOp
I.CmpSle IntType
Int64) SubExp
j' SubExp
i'
SubExp
backwards_ok <-
forall (m :: * -> *).
MonadBuilder m =>
[Char] -> Exp (Rep m) -> m SubExp
letSubExp [Char]
"backwards_ok"
forall (m :: * -> *) a b. Monad m => (a -> m b) -> m a -> m b
=<< forall (m :: * -> *). MonadBuilder m => [SubExp] -> m (Exp (Rep m))
eAll
[SubExp
negone_lte_j, SubExp
negone_lte_j, SubExp
j_lte_i, SubExp
zero_leq_i_p_m_t_s, SubExp
i_p_m_t_s_leq_w]
SubExp
slice_ok <-
forall (m :: * -> *).
MonadBuilder m =>
[Char] -> Exp (Rep m) -> m SubExp
letSubExp [Char]
"slice_ok"
forall (m :: * -> *) a b. Monad m => (a -> m b) -> m a -> m b
=<< forall (m :: * -> *).
(MonadBuilder m, BranchType (Rep m) ~ ExtType) =>
m (Exp (Rep m))
-> m (Body (Rep m)) -> m (Body (Rep m)) -> m (Exp (Rep m))
eIf
(forall (m :: * -> *). MonadBuilder m => SubExp -> m (Exp (Rep m))
eSubExp SubExp
backwards)
(forall (m :: * -> *).
MonadBuilder m =>
[SubExp] -> m (Body (Rep m))
resultBodyM [SubExp
backwards_ok])
(forall (m :: * -> *).
MonadBuilder m =>
[SubExp] -> m (Body (Rep m))
resultBodyM [SubExp
forwards_ok])
SubExp
ok_or_empty <-
forall (m :: * -> *).
MonadBuilder m =>
[Char] -> Exp (Rep m) -> m SubExp
letSubExp [Char]
"ok_or_empty" forall a b. (a -> b) -> a -> b
$
forall rep. BasicOp -> Exp rep
I.BasicOp forall a b. (a -> b) -> a -> b
$
BinOp -> SubExp -> SubExp -> BasicOp
I.BinOp BinOp
I.LogOr SubExp
empty_slice SubExp
slice_ok
SubExp
acceptable <-
forall (m :: * -> *).
MonadBuilder m =>
[Char] -> Exp (Rep m) -> m SubExp
letSubExp [Char]
"slice_acceptable" forall a b. (a -> b) -> a -> b
$
forall rep. BasicOp -> Exp rep
I.BasicOp forall a b. (a -> b) -> a -> b
$
BinOp -> SubExp -> SubExp -> BasicOp
I.BinOp BinOp
I.LogAnd SubExp
nonzero_stride SubExp
ok_or_empty
let parts :: [ErrorMsgPart SubExp]
parts = case (Maybe Exp
i, Maybe Exp
j, Maybe Exp
s) of
(Maybe Exp
_, Maybe Exp
_, Just {}) ->
[ forall b a. b -> (a -> b) -> Maybe a -> b
maybe ErrorMsgPart SubExp
"" (forall a b. a -> b -> a
const forall a b. (a -> b) -> a -> b
$ forall a. PrimType -> a -> ErrorMsgPart a
ErrorVal PrimType
int64 SubExp
i') Maybe Exp
i,
ErrorMsgPart SubExp
":",
forall b a. b -> (a -> b) -> Maybe a -> b
maybe ErrorMsgPart SubExp
"" (forall a b. a -> b -> a
const forall a b. (a -> b) -> a -> b
$ forall a. PrimType -> a -> ErrorMsgPart a
ErrorVal PrimType
int64 SubExp
j') Maybe Exp
j,
ErrorMsgPart SubExp
":",
forall a. PrimType -> a -> ErrorMsgPart a
ErrorVal PrimType
int64 SubExp
s'
]
(Maybe Exp
_, Just {}, Maybe Exp
_) ->
[ forall b a. b -> (a -> b) -> Maybe a -> b
maybe ErrorMsgPart SubExp
"" (forall a b. a -> b -> a
const forall a b. (a -> b) -> a -> b
$ forall a. PrimType -> a -> ErrorMsgPart a
ErrorVal PrimType
int64 SubExp
i') Maybe Exp
i,
ErrorMsgPart SubExp
":",
forall a. PrimType -> a -> ErrorMsgPart a
ErrorVal PrimType
int64 SubExp
j'
]
forall a. [a] -> [a] -> [a]
++ forall b a. b -> (a -> b) -> Maybe a -> b
maybe forall a. Monoid a => a
mempty (forall a b. a -> b -> a
const [ErrorMsgPart SubExp
":", forall a. PrimType -> a -> ErrorMsgPart a
ErrorVal PrimType
int64 SubExp
s']) Maybe Exp
s
(Maybe Exp
_, Maybe Exp
Nothing, Maybe Exp
Nothing) ->
[forall a. PrimType -> a -> ErrorMsgPart a
ErrorVal PrimType
int64 SubExp
i', ErrorMsgPart SubExp
":"]
forall (f :: * -> *) a. Applicative f => a -> f a
pure (forall d. d -> d -> d -> DimIndex d
I.DimSlice SubExp
i' SubExp
n SubExp
s', SubExp
acceptable, [ErrorMsgPart SubExp]
parts)
where
zero :: SubExp
zero = forall v. IsValue v => v -> SubExp
constant (Int64
0 :: Int64)
negone :: SubExp
negone = forall v. IsValue v => v -> SubExp
constant (-Int64
1 :: Int64)
one :: SubExp
one = forall v. IsValue v => v -> SubExp
constant (Int64
1 :: Int64)
internaliseScanOrReduce ::
String ->
String ->
(SubExp -> I.Lambda SOACS -> [SubExp] -> [VName] -> InternaliseM (SOAC SOACS)) ->
(E.Exp, E.Exp, E.Exp, SrcLoc) ->
InternaliseM [SubExp]
internaliseScanOrReduce :: [Char]
-> [Char]
-> (SubExp
-> Lambda SOACS
-> [SubExp]
-> [VName]
-> InternaliseM (SOAC SOACS))
-> (Exp, Exp, Exp, SrcLoc)
-> InternaliseM [SubExp]
internaliseScanOrReduce [Char]
desc [Char]
what SubExp
-> Lambda SOACS -> [SubExp] -> [VName] -> InternaliseM (SOAC SOACS)
f (Exp
lam, Exp
ne, Exp
arr, SrcLoc
loc) = do
[VName]
arrs <- [Char] -> Exp -> InternaliseM [VName]
internaliseExpToVars ([Char]
what forall a. [a] -> [a] -> [a]
++ [Char]
"_arr") Exp
arr
[SubExp]
nes <- [Char] -> Exp -> InternaliseM [SubExp]
internaliseExp ([Char]
what forall a. [a] -> [a] -> [a]
++ [Char]
"_ne") Exp
ne
[SubExp]
nes' <- forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
t a -> (a -> m b) -> m (t b)
forM (forall a b. [a] -> [b] -> [(a, b)]
zip [SubExp]
nes [VName]
arrs) forall a b. (a -> b) -> a -> b
$ \(SubExp
ne', VName
arr') -> do
TypeBase Shape NoUniqueness
rowtype <- forall u. Int -> TypeBase Shape u -> TypeBase Shape u
I.stripArray Int
1 forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> forall rep (m :: * -> *).
HasScope rep m =>
VName -> m (TypeBase Shape NoUniqueness)
lookupType VName
arr'
ErrorMsg SubExp
-> SrcLoc
-> TypeBase Shape NoUniqueness
-> [Char]
-> SubExp
-> InternaliseM SubExp
ensureShape
ErrorMsg SubExp
"Row shape of input array does not match shape of neutral element"
SrcLoc
loc
TypeBase Shape NoUniqueness
rowtype
([Char]
what forall a. [a] -> [a] -> [a]
++ [Char]
"_ne_right_shape")
SubExp
ne'
[TypeBase Shape NoUniqueness]
nests <- forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
(a -> m b) -> t a -> m (t b)
mapM forall t (m :: * -> *).
HasScope t m =>
SubExp -> m (TypeBase Shape NoUniqueness)
I.subExpType [SubExp]
nes'
[TypeBase Shape NoUniqueness]
arrts <- forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
(a -> m b) -> t a -> m (t b)
mapM forall rep (m :: * -> *).
HasScope rep m =>
VName -> m (TypeBase Shape NoUniqueness)
lookupType [VName]
arrs
Lambda SOACS
lam' <- InternaliseLambda
-> Exp
-> [TypeBase Shape NoUniqueness]
-> [TypeBase Shape NoUniqueness]
-> InternaliseM (Lambda SOACS)
internaliseFoldLambda InternaliseLambda
internaliseLambda Exp
lam [TypeBase Shape NoUniqueness]
nests [TypeBase Shape NoUniqueness]
arrts
SubExp
w <- forall u. Int -> [TypeBase Shape u] -> SubExp
arraysSize Int
0 forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
(a -> m b) -> t a -> m (t b)
mapM forall rep (m :: * -> *).
HasScope rep m =>
VName -> m (TypeBase Shape NoUniqueness)
lookupType [VName]
arrs
[Char] -> Exp SOACS -> InternaliseM [SubExp]
letValExp' [Char]
desc forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall rep. Op rep -> Exp rep
I.Op forall (m :: * -> *) a b. Monad m => (a -> m b) -> m a -> m b
=<< SubExp
-> Lambda SOACS -> [SubExp] -> [VName] -> InternaliseM (SOAC SOACS)
f SubExp
w Lambda SOACS
lam' [SubExp]
nes' [VName]
arrs
internaliseHist ::
Int ->
String ->
E.Exp ->
E.Exp ->
E.Exp ->
E.Exp ->
E.Exp ->
E.Exp ->
SrcLoc ->
InternaliseM [SubExp]
internaliseHist :: Int
-> [Char]
-> Exp
-> Exp
-> Exp
-> Exp
-> Exp
-> Exp
-> SrcLoc
-> InternaliseM [SubExp]
internaliseHist Int
dim [Char]
desc Exp
rf Exp
hist Exp
op Exp
ne Exp
buckets Exp
img SrcLoc
loc = do
SubExp
rf' <- [Char] -> Exp -> InternaliseM SubExp
internaliseExp1 [Char]
"hist_rf" Exp
rf
[SubExp]
ne' <- [Char] -> Exp -> InternaliseM [SubExp]
internaliseExp [Char]
"hist_ne" Exp
ne
[VName]
hist' <- [Char] -> Exp -> InternaliseM [VName]
internaliseExpToVars [Char]
"hist_hist" Exp
hist
[VName]
buckets' <- [Char] -> Exp -> InternaliseM [VName]
internaliseExpToVars [Char]
"hist_buckets" Exp
buckets
[VName]
img' <- [Char] -> Exp -> InternaliseM [VName]
internaliseExpToVars [Char]
"hist_img" Exp
img
[SubExp]
ne_shp <- forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
t a -> (a -> m b) -> m (t b)
forM (forall a b. [a] -> [b] -> [(a, b)]
zip [SubExp]
ne' [VName]
hist') forall a b. (a -> b) -> a -> b
$ \(SubExp
n, VName
h) -> do
TypeBase Shape NoUniqueness
rowtype <- forall u. Int -> TypeBase Shape u -> TypeBase Shape u
I.stripArray Int
1 forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> forall rep (m :: * -> *).
HasScope rep m =>
VName -> m (TypeBase Shape NoUniqueness)
lookupType VName
h
ErrorMsg SubExp
-> SrcLoc
-> TypeBase Shape NoUniqueness
-> [Char]
-> SubExp
-> InternaliseM SubExp
ensureShape
ErrorMsg SubExp
"Row shape of destination array does not match shape of neutral element"
SrcLoc
loc
TypeBase Shape NoUniqueness
rowtype
[Char]
"hist_ne_right_shape"
SubExp
n
[TypeBase Shape NoUniqueness]
ne_ts <- forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
(a -> m b) -> t a -> m (t b)
mapM forall t (m :: * -> *).
HasScope t m =>
SubExp -> m (TypeBase Shape NoUniqueness)
I.subExpType [SubExp]
ne_shp
[TypeBase Shape NoUniqueness]
his_ts <- forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
(a -> m b) -> t a -> m (t b)
mapM (forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap (forall u. Int -> TypeBase Shape u -> TypeBase Shape u
I.stripArray (Int
dim forall a. Num a => a -> a -> a
- Int
1)) forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall rep (m :: * -> *).
HasScope rep m =>
VName -> m (TypeBase Shape NoUniqueness)
lookupType) [VName]
hist'
Lambda SOACS
op' <- InternaliseLambda
-> Exp
-> [TypeBase Shape NoUniqueness]
-> [TypeBase Shape NoUniqueness]
-> InternaliseM (Lambda SOACS)
internaliseFoldLambda InternaliseLambda
internaliseLambda Exp
op [TypeBase Shape NoUniqueness]
ne_ts [TypeBase Shape NoUniqueness]
his_ts
[Param (TypeBase Shape NoUniqueness)]
bucket_params <- forall (m :: * -> *) a. Applicative m => Int -> m a -> m [a]
replicateM Int
dim (forall (m :: * -> *) dec.
MonadFreshNames m =>
[Char] -> dec -> m (Param dec)
newParam [Char]
"bucket_p" forall a b. (a -> b) -> a -> b
$ forall shape u. PrimType -> TypeBase shape u
I.Prim PrimType
int64)
[Param (TypeBase Shape NoUniqueness)]
img_params <- forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
(a -> m b) -> t a -> m (t b)
mapM (forall (m :: * -> *) dec.
MonadFreshNames m =>
[Char] -> dec -> m (Param dec)
newParam [Char]
"img_p" forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall u. TypeBase Shape u -> TypeBase Shape u
rowType) forall (m :: * -> *) a b. Monad m => (a -> m b) -> m a -> m b
=<< forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
(a -> m b) -> t a -> m (t b)
mapM forall rep (m :: * -> *).
HasScope rep m =>
VName -> m (TypeBase Shape NoUniqueness)
lookupType [VName]
img'
let params :: [Param (TypeBase Shape NoUniqueness)]
params = [Param (TypeBase Shape NoUniqueness)]
bucket_params forall a. [a] -> [a] -> [a]
++ [Param (TypeBase Shape NoUniqueness)]
img_params
rettype :: [TypeBase Shape NoUniqueness]
rettype = forall a. Int -> a -> [a]
replicate Int
dim (forall shape u. PrimType -> TypeBase shape u
I.Prim PrimType
int64) forall a. [a] -> [a] -> [a]
++ [TypeBase Shape NoUniqueness]
ne_ts
body :: Body SOACS
body = forall rep. Buildable rep => Stms rep -> Result -> Body rep
mkBody forall a. Monoid a => a
mempty forall a b. (a -> b) -> a -> b
$ [VName] -> Result
varsRes forall a b. (a -> b) -> a -> b
$ forall a b. (a -> b) -> [a] -> [b]
map forall dec. Param dec -> VName
I.paramName [Param (TypeBase Shape NoUniqueness)]
params
Lambda SOACS
lam' <-
forall (m :: * -> *).
MonadBuilder m =>
[LParam (Rep m)] -> m Result -> m (Lambda (Rep m))
mkLambda [Param (TypeBase Shape NoUniqueness)]
params forall a b. (a -> b) -> a -> b
$
ErrorMsg SubExp
-> SrcLoc
-> [TypeBase Shape NoUniqueness]
-> Result
-> InternaliseM Result
ensureResultShape
ErrorMsg SubExp
"Row shape of value array does not match row shape of hist target"
(forall a. Located a => a -> SrcLoc
srclocOf Exp
img)
[TypeBase Shape NoUniqueness]
rettype
forall (m :: * -> *) a b. Monad m => (a -> m b) -> m a -> m b
=<< forall (m :: * -> *). MonadBuilder m => Body (Rep m) -> m Result
bodyBind Body SOACS
body
Shape
shape_hist <- forall d. [d] -> ShapeBase d
I.Shape forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall a. Int -> [a] -> [a]
take Int
dim forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall u. TypeBase Shape u -> [SubExp]
I.arrayDims forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> forall rep (m :: * -> *).
HasScope rep m =>
VName -> m (TypeBase Shape NoUniqueness)
lookupType (forall a. [a] -> a
head [VName]
hist')
SubExp
w_img <- forall u. Int -> TypeBase Shape u -> SubExp
I.arraySize Int
0 forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> forall rep (m :: * -> *).
HasScope rep m =>
VName -> m (TypeBase Shape NoUniqueness)
lookupType (forall a. [a] -> a
head [VName]
img')
[Char] -> Exp SOACS -> InternaliseM [SubExp]
letValExp' [Char]
desc forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall rep. Op rep -> Exp rep
I.Op forall a b. (a -> b) -> a -> b
$
forall rep.
SubExp -> [VName] -> [HistOp rep] -> Lambda rep -> SOAC rep
I.Hist SubExp
w_img ([VName]
buckets' forall a. [a] -> [a] -> [a]
++ [VName]
img') [forall rep.
Shape -> SubExp -> [VName] -> [SubExp] -> Lambda rep -> HistOp rep
HistOp Shape
shape_hist SubExp
rf' [VName]
hist' [SubExp]
ne_shp Lambda SOACS
op'] Lambda SOACS
lam'
internaliseStreamAcc ::
String ->
E.Exp ->
Maybe (E.Exp, E.Exp) ->
E.Exp ->
E.Exp ->
InternaliseM [SubExp]
internaliseStreamAcc :: [Char]
-> Exp -> Maybe (Exp, Exp) -> Exp -> Exp -> InternaliseM [SubExp]
internaliseStreamAcc [Char]
desc Exp
dest Maybe (Exp, Exp)
op Exp
lam Exp
bs = do
[VName]
dest' <- [Char] -> Exp -> InternaliseM [VName]
internaliseExpToVars [Char]
"scatter_dest" Exp
dest
[VName]
bs' <- [Char] -> Exp -> InternaliseM [VName]
internaliseExpToVars [Char]
"scatter_input" Exp
bs
VName
acc_cert_v <- forall (m :: * -> *). MonadFreshNames m => [Char] -> m VName
newVName [Char]
"acc_cert"
[TypeBase Shape NoUniqueness]
dest_ts <- forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
(a -> m b) -> t a -> m (t b)
mapM forall rep (m :: * -> *).
HasScope rep m =>
VName -> m (TypeBase Shape NoUniqueness)
lookupType [VName]
dest'
let dest_w :: SubExp
dest_w = forall u. Int -> [TypeBase Shape u] -> SubExp
arraysSize Int
0 [TypeBase Shape NoUniqueness]
dest_ts
acc_t :: TypeBase Shape NoUniqueness
acc_t = forall shape u.
VName
-> Shape -> [TypeBase Shape NoUniqueness] -> u -> TypeBase shape u
Acc VName
acc_cert_v (forall d. [d] -> ShapeBase d
I.Shape [SubExp
dest_w]) (forall a b. (a -> b) -> [a] -> [b]
map forall u. TypeBase Shape u -> TypeBase Shape u
rowType [TypeBase Shape NoUniqueness]
dest_ts) NoUniqueness
NoUniqueness
Param (TypeBase Shape NoUniqueness)
acc_p <- forall (m :: * -> *) dec.
MonadFreshNames m =>
[Char] -> dec -> m (Param dec)
newParam [Char]
"acc_p" TypeBase Shape NoUniqueness
acc_t
Lambda SOACS
withacc_lam <- forall (m :: * -> *).
MonadBuilder m =>
[LParam (Rep m)] -> m Result -> m (Lambda (Rep m))
mkLambda [forall dec. Attrs -> VName -> dec -> Param dec
Param forall a. Monoid a => a
mempty VName
acc_cert_v (forall shape u. PrimType -> TypeBase shape u
I.Prim PrimType
I.Unit), Param (TypeBase Shape NoUniqueness)
acc_p] forall a b. (a -> b) -> a -> b
$ do
[TypeBase Shape NoUniqueness]
bs_ts <- forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
(a -> m b) -> t a -> m (t b)
mapM forall rep (m :: * -> *).
HasScope rep m =>
VName -> m (TypeBase Shape NoUniqueness)
lookupType [VName]
bs'
Lambda SOACS
lam' <- Exp -> [TypeBase Shape NoUniqueness] -> InternaliseM (Lambda SOACS)
internaliseLambdaCoerce Exp
lam forall a b. (a -> b) -> a -> b
$ forall a b. (a -> b) -> [a] -> [b]
map forall u. TypeBase Shape u -> TypeBase Shape u
rowType forall a b. (a -> b) -> a -> b
$ forall dec. Typed dec => Param dec -> TypeBase Shape NoUniqueness
paramType Param (TypeBase Shape NoUniqueness)
acc_p forall a. a -> [a] -> [a]
: [TypeBase Shape NoUniqueness]
bs_ts
let w :: SubExp
w = forall u. Int -> [TypeBase Shape u] -> SubExp
arraysSize Int
0 [TypeBase Shape NoUniqueness]
bs_ts
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap [SubExp] -> Result
subExpsRes forall b c a. (b -> c) -> (a -> b) -> a -> c
. [Char] -> Exp SOACS -> InternaliseM [SubExp]
letValExp' [Char]
"acc_res" forall a b. (a -> b) -> a -> b
$
forall rep. Op rep -> Exp rep
I.Op forall a b. (a -> b) -> a -> b
$
forall rep. SubExp -> [VName] -> ScremaForm rep -> SOAC rep
I.Screma SubExp
w (forall dec. Param dec -> VName
I.paramName Param (TypeBase Shape NoUniqueness)
acc_p forall a. a -> [a] -> [a]
: [VName]
bs') (forall rep. Lambda rep -> ScremaForm rep
I.mapSOAC Lambda SOACS
lam')
Maybe (Lambda SOACS, [SubExp])
op' <-
case Maybe (Exp, Exp)
op of
Just (Exp
op_lam, Exp
ne) -> do
[SubExp]
ne' <- [Char] -> Exp -> InternaliseM [SubExp]
internaliseExp [Char]
"hist_ne" Exp
ne
[TypeBase Shape NoUniqueness]
ne_ts <- forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
(a -> m b) -> t a -> m (t b)
mapM forall t (m :: * -> *).
HasScope t m =>
SubExp -> m (TypeBase Shape NoUniqueness)
I.subExpType [SubExp]
ne'
([Param (TypeBase Shape NoUniqueness)]
lam_params, Body SOACS
lam_body, [TypeBase Shape NoUniqueness]
lam_rettype) <-
InternaliseLambda
internaliseLambda Exp
op_lam forall a b. (a -> b) -> a -> b
$ [TypeBase Shape NoUniqueness]
ne_ts forall a. [a] -> [a] -> [a]
++ [TypeBase Shape NoUniqueness]
ne_ts
Param (TypeBase Shape NoUniqueness)
idxp <- forall (m :: * -> *) dec.
MonadFreshNames m =>
[Char] -> dec -> m (Param dec)
newParam [Char]
"idx" forall a b. (a -> b) -> a -> b
$ forall shape u. PrimType -> TypeBase shape u
I.Prim PrimType
int64
let op_lam' :: Lambda SOACS
op_lam' = forall rep.
[LParam rep]
-> Body rep -> [TypeBase Shape NoUniqueness] -> Lambda rep
I.Lambda (Param (TypeBase Shape NoUniqueness)
idxp forall a. a -> [a] -> [a]
: [Param (TypeBase Shape NoUniqueness)]
lam_params) Body SOACS
lam_body [TypeBase Shape NoUniqueness]
lam_rettype
forall (f :: * -> *) a. Applicative f => a -> f a
pure forall a b. (a -> b) -> a -> b
$ forall a. a -> Maybe a
Just (Lambda SOACS
op_lam', [SubExp]
ne')
Maybe (Exp, Exp)
Nothing ->
forall (f :: * -> *) a. Applicative f => a -> f a
pure forall a. Maybe a
Nothing
SubExp
destw <- forall u. Int -> [TypeBase Shape u] -> SubExp
arraysSize Int
0 forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
(a -> m b) -> t a -> m (t b)
mapM forall rep (m :: * -> *).
HasScope rep m =>
VName -> m (TypeBase Shape NoUniqueness)
lookupType [VName]
dest'
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap (forall a b. (a -> b) -> [a] -> [b]
map VName -> SubExp
I.Var) forall a b. (a -> b) -> a -> b
$
forall (m :: * -> *).
MonadBuilder m =>
[Char] -> Exp (Rep m) -> m [VName]
letTupExp [Char]
desc forall a b. (a -> b) -> a -> b
$
forall rep. [WithAccInput rep] -> Lambda rep -> Exp rep
WithAcc [(forall d. [d] -> ShapeBase d
I.Shape [SubExp
destw], [VName]
dest', Maybe (Lambda SOACS, [SubExp])
op')] Lambda SOACS
withacc_lam
internaliseExp1 :: String -> E.Exp -> InternaliseM I.SubExp
internaliseExp1 :: [Char] -> Exp -> InternaliseM SubExp
internaliseExp1 [Char]
desc Exp
e = do
[SubExp]
vs <- [Char] -> Exp -> InternaliseM [SubExp]
internaliseExp [Char]
desc Exp
e
case [SubExp]
vs of
[SubExp
se] -> forall (f :: * -> *) a. Applicative f => a -> f a
pure SubExp
se
[SubExp]
_ -> forall a. HasCallStack => [Char] -> a
error [Char]
"Internalise.internaliseExp1: was passed not just a single subexpression"
internaliseSizeExp :: String -> E.Exp -> InternaliseM (I.SubExp, IntType)
internaliseSizeExp :: [Char] -> Exp -> InternaliseM (SubExp, IntType)
internaliseSizeExp [Char]
s Exp
e = do
SubExp
e' <- [Char] -> Exp -> InternaliseM SubExp
internaliseExp1 [Char]
s Exp
e
case Exp -> PatType
E.typeOf Exp
e of
E.Scalar (E.Prim (E.Signed IntType
it)) -> (,IntType
it) forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> forall (m :: * -> *).
MonadBuilder m =>
IntType -> SubExp -> m SubExp
asIntS IntType
Int64 SubExp
e'
PatType
_ -> forall a. HasCallStack => [Char] -> a
error [Char]
"internaliseSizeExp: bad type"
internaliseExpToVars :: String -> E.Exp -> InternaliseM [I.VName]
internaliseExpToVars :: [Char] -> Exp -> InternaliseM [VName]
internaliseExpToVars [Char]
desc Exp
e =
forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
(a -> m b) -> t a -> m (t b)
mapM SubExp -> InternaliseM VName
asIdent forall (m :: * -> *) a b. Monad m => (a -> m b) -> m a -> m b
=<< [Char] -> Exp -> InternaliseM [SubExp]
internaliseExp [Char]
desc Exp
e
where
asIdent :: SubExp -> InternaliseM VName
asIdent (I.Var VName
v) = forall (f :: * -> *) a. Applicative f => a -> f a
pure VName
v
asIdent SubExp
se = forall (m :: * -> *).
MonadBuilder m =>
[Char] -> Exp (Rep m) -> m VName
letExp [Char]
desc forall a b. (a -> b) -> a -> b
$ forall rep. BasicOp -> Exp rep
I.BasicOp forall a b. (a -> b) -> a -> b
$ SubExp -> BasicOp
I.SubExp SubExp
se
internaliseOperation ::
String ->
E.Exp ->
(I.VName -> InternaliseM I.BasicOp) ->
InternaliseM [I.SubExp]
internaliseOperation :: [Char]
-> Exp -> (VName -> InternaliseM BasicOp) -> InternaliseM [SubExp]
internaliseOperation [Char]
s Exp
e VName -> InternaliseM BasicOp
op = do
[VName]
vs <- [Char] -> Exp -> InternaliseM [VName]
internaliseExpToVars [Char]
s Exp
e
forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
(a -> m b) -> t a -> m (t b)
mapM (forall (m :: * -> *).
MonadBuilder m =>
[Char] -> Exp (Rep m) -> m SubExp
letSubExp [Char]
s forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall rep. BasicOp -> Exp rep
I.BasicOp forall (m :: * -> *) b c a.
Monad m =>
(b -> m c) -> (a -> m b) -> a -> m c
<=< VName -> InternaliseM BasicOp
op) [VName]
vs
certifyingNonzero ::
SrcLoc ->
IntType ->
SubExp ->
InternaliseM a ->
InternaliseM a
certifyingNonzero :: forall a.
SrcLoc -> IntType -> SubExp -> InternaliseM a -> InternaliseM a
certifyingNonzero SrcLoc
loc IntType
t SubExp
x InternaliseM a
m = do
SubExp
zero <-
forall (m :: * -> *).
MonadBuilder m =>
[Char] -> Exp (Rep m) -> m SubExp
letSubExp [Char]
"zero" forall a b. (a -> b) -> a -> b
$
forall rep. BasicOp -> Exp rep
I.BasicOp forall a b. (a -> b) -> a -> b
$
CmpOp -> SubExp -> SubExp -> BasicOp
CmpOp (PrimType -> CmpOp
CmpEq (IntType -> PrimType
IntType IntType
t)) SubExp
x (IntType -> Integer -> SubExp
intConst IntType
t Integer
0)
SubExp
nonzero <- forall (m :: * -> *).
MonadBuilder m =>
[Char] -> Exp (Rep m) -> m SubExp
letSubExp [Char]
"nonzero" forall a b. (a -> b) -> a -> b
$ forall rep. BasicOp -> Exp rep
I.BasicOp forall a b. (a -> b) -> a -> b
$ UnOp -> SubExp -> BasicOp
UnOp UnOp
I.Not SubExp
zero
Certs
c <- [Char] -> SubExp -> ErrorMsg SubExp -> SrcLoc -> InternaliseM Certs
assert [Char]
"nonzero_cert" SubExp
nonzero ErrorMsg SubExp
"division by zero" SrcLoc
loc
forall (m :: * -> *) a. MonadBuilder m => Certs -> m a -> m a
certifying Certs
c InternaliseM a
m
certifyingNonnegative ::
SrcLoc ->
IntType ->
SubExp ->
InternaliseM a ->
InternaliseM a
certifyingNonnegative :: forall a.
SrcLoc -> IntType -> SubExp -> InternaliseM a -> InternaliseM a
certifyingNonnegative SrcLoc
loc IntType
t SubExp
x InternaliseM a
m = do
SubExp
nonnegative <-
forall (m :: * -> *).
MonadBuilder m =>
[Char] -> Exp (Rep m) -> m SubExp
letSubExp [Char]
"nonnegative" forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall rep. BasicOp -> Exp rep
I.BasicOp forall a b. (a -> b) -> a -> b
$
CmpOp -> SubExp -> SubExp -> BasicOp
CmpOp (IntType -> CmpOp
CmpSle IntType
t) (IntType -> Integer -> SubExp
intConst IntType
t Integer
0) SubExp
x
Certs
c <- [Char] -> SubExp -> ErrorMsg SubExp -> SrcLoc -> InternaliseM Certs
assert [Char]
"nonzero_cert" SubExp
nonnegative ErrorMsg SubExp
"negative exponent" SrcLoc
loc
forall (m :: * -> *) a. MonadBuilder m => Certs -> m a -> m a
certifying Certs
c InternaliseM a
m
internaliseBinOp ::
SrcLoc ->
String ->
E.BinOp ->
I.SubExp ->
I.SubExp ->
E.PrimType ->
E.PrimType ->
InternaliseM [I.SubExp]
internaliseBinOp :: SrcLoc
-> [Char]
-> BinOp
-> SubExp
-> SubExp
-> PrimType
-> PrimType
-> InternaliseM [SubExp]
internaliseBinOp SrcLoc
_ [Char]
desc BinOp
E.LogAnd SubExp
x SubExp
y PrimType
E.Bool PrimType
_ =
[Char] -> BinOp -> SubExp -> SubExp -> InternaliseM [SubExp]
simpleBinOp [Char]
desc BinOp
I.LogAnd SubExp
x SubExp
y
internaliseBinOp SrcLoc
_ [Char]
desc BinOp
E.LogOr SubExp
x SubExp
y PrimType
E.Bool PrimType
_ =
[Char] -> BinOp -> SubExp -> SubExp -> InternaliseM [SubExp]
simpleBinOp [Char]
desc BinOp
I.LogOr SubExp
x SubExp
y
internaliseBinOp SrcLoc
_ [Char]
desc BinOp
E.Plus SubExp
x SubExp
y (E.Signed IntType
t) PrimType
_ =
[Char] -> BinOp -> SubExp -> SubExp -> InternaliseM [SubExp]
simpleBinOp [Char]
desc (IntType -> Overflow -> BinOp
I.Add IntType
t Overflow
I.OverflowWrap) SubExp
x SubExp
y
internaliseBinOp SrcLoc
_ [Char]
desc BinOp
E.Plus SubExp
x SubExp
y (E.Unsigned IntType
t) PrimType
_ =
[Char] -> BinOp -> SubExp -> SubExp -> InternaliseM [SubExp]
simpleBinOp [Char]
desc (IntType -> Overflow -> BinOp
I.Add IntType
t Overflow
I.OverflowWrap) SubExp
x SubExp
y
internaliseBinOp SrcLoc
_ [Char]
desc BinOp
E.Plus SubExp
x SubExp
y (E.FloatType FloatType
t) PrimType
_ =
[Char] -> BinOp -> SubExp -> SubExp -> InternaliseM [SubExp]
simpleBinOp [Char]
desc (FloatType -> BinOp
I.FAdd FloatType
t) SubExp
x SubExp
y
internaliseBinOp SrcLoc
_ [Char]
desc BinOp
E.Minus SubExp
x SubExp
y (E.Signed IntType
t) PrimType
_ =
[Char] -> BinOp -> SubExp -> SubExp -> InternaliseM [SubExp]
simpleBinOp [Char]
desc (IntType -> Overflow -> BinOp
I.Sub IntType
t Overflow
I.OverflowWrap) SubExp
x SubExp
y
internaliseBinOp SrcLoc
_ [Char]
desc BinOp
E.Minus SubExp
x SubExp
y (E.Unsigned IntType
t) PrimType
_ =
[Char] -> BinOp -> SubExp -> SubExp -> InternaliseM [SubExp]
simpleBinOp [Char]
desc (IntType -> Overflow -> BinOp
I.Sub IntType
t Overflow
I.OverflowWrap) SubExp
x SubExp
y
internaliseBinOp SrcLoc
_ [Char]
desc BinOp
E.Minus SubExp
x SubExp
y (E.FloatType FloatType
t) PrimType
_ =
[Char] -> BinOp -> SubExp -> SubExp -> InternaliseM [SubExp]
simpleBinOp [Char]
desc (FloatType -> BinOp
I.FSub FloatType
t) SubExp
x SubExp
y
internaliseBinOp SrcLoc
_ [Char]
desc BinOp
E.Times SubExp
x SubExp
y (E.Signed IntType
t) PrimType
_ =
[Char] -> BinOp -> SubExp -> SubExp -> InternaliseM [SubExp]
simpleBinOp [Char]
desc (IntType -> Overflow -> BinOp
I.Mul IntType
t Overflow
I.OverflowWrap) SubExp
x SubExp
y
internaliseBinOp SrcLoc
_ [Char]
desc BinOp
E.Times SubExp
x SubExp
y (E.Unsigned IntType
t) PrimType
_ =
[Char] -> BinOp -> SubExp -> SubExp -> InternaliseM [SubExp]
simpleBinOp [Char]
desc (IntType -> Overflow -> BinOp
I.Mul IntType
t Overflow
I.OverflowWrap) SubExp
x SubExp
y
internaliseBinOp SrcLoc
_ [Char]
desc BinOp
E.Times SubExp
x SubExp
y (E.FloatType FloatType
t) PrimType
_ =
[Char] -> BinOp -> SubExp -> SubExp -> InternaliseM [SubExp]
simpleBinOp [Char]
desc (FloatType -> BinOp
I.FMul FloatType
t) SubExp
x SubExp
y
internaliseBinOp SrcLoc
loc [Char]
desc BinOp
E.Divide SubExp
x SubExp
y (E.Signed IntType
t) PrimType
_ =
forall a.
SrcLoc -> IntType -> SubExp -> InternaliseM a -> InternaliseM a
certifyingNonzero SrcLoc
loc IntType
t SubExp
y forall a b. (a -> b) -> a -> b
$
[Char] -> BinOp -> SubExp -> SubExp -> InternaliseM [SubExp]
simpleBinOp [Char]
desc (IntType -> Safety -> BinOp
I.SDiv IntType
t Safety
I.Unsafe) SubExp
x SubExp
y
internaliseBinOp SrcLoc
loc [Char]
desc BinOp
E.Divide SubExp
x SubExp
y (E.Unsigned IntType
t) PrimType
_ =
forall a.
SrcLoc -> IntType -> SubExp -> InternaliseM a -> InternaliseM a
certifyingNonzero SrcLoc
loc IntType
t SubExp
y forall a b. (a -> b) -> a -> b
$
[Char] -> BinOp -> SubExp -> SubExp -> InternaliseM [SubExp]
simpleBinOp [Char]
desc (IntType -> Safety -> BinOp
I.UDiv IntType
t Safety
I.Unsafe) SubExp
x SubExp
y
internaliseBinOp SrcLoc
_ [Char]
desc BinOp
E.Divide SubExp
x SubExp
y (E.FloatType FloatType
t) PrimType
_ =
[Char] -> BinOp -> SubExp -> SubExp -> InternaliseM [SubExp]
simpleBinOp [Char]
desc (FloatType -> BinOp
I.FDiv FloatType
t) SubExp
x SubExp
y
internaliseBinOp SrcLoc
_ [Char]
desc BinOp
E.Pow SubExp
x SubExp
y (E.FloatType FloatType
t) PrimType
_ =
[Char] -> BinOp -> SubExp -> SubExp -> InternaliseM [SubExp]
simpleBinOp [Char]
desc (FloatType -> BinOp
I.FPow FloatType
t) SubExp
x SubExp
y
internaliseBinOp SrcLoc
loc [Char]
desc BinOp
E.Pow SubExp
x SubExp
y (E.Signed IntType
t) PrimType
_ =
forall a.
SrcLoc -> IntType -> SubExp -> InternaliseM a -> InternaliseM a
certifyingNonnegative SrcLoc
loc IntType
t SubExp
y forall a b. (a -> b) -> a -> b
$
[Char] -> BinOp -> SubExp -> SubExp -> InternaliseM [SubExp]
simpleBinOp [Char]
desc (IntType -> BinOp
I.Pow IntType
t) SubExp
x SubExp
y
internaliseBinOp SrcLoc
_ [Char]
desc BinOp
E.Pow SubExp
x SubExp
y (E.Unsigned IntType
t) PrimType
_ =
[Char] -> BinOp -> SubExp -> SubExp -> InternaliseM [SubExp]
simpleBinOp [Char]
desc (IntType -> BinOp
I.Pow IntType
t) SubExp
x SubExp
y
internaliseBinOp SrcLoc
loc [Char]
desc BinOp
E.Mod SubExp
x SubExp
y (E.Signed IntType
t) PrimType
_ =
forall a.
SrcLoc -> IntType -> SubExp -> InternaliseM a -> InternaliseM a
certifyingNonzero SrcLoc
loc IntType
t SubExp
y forall a b. (a -> b) -> a -> b
$
[Char] -> BinOp -> SubExp -> SubExp -> InternaliseM [SubExp]
simpleBinOp [Char]
desc (IntType -> Safety -> BinOp
I.SMod IntType
t Safety
I.Unsafe) SubExp
x SubExp
y
internaliseBinOp SrcLoc
loc [Char]
desc BinOp
E.Mod SubExp
x SubExp
y (E.Unsigned IntType
t) PrimType
_ =
forall a.
SrcLoc -> IntType -> SubExp -> InternaliseM a -> InternaliseM a
certifyingNonzero SrcLoc
loc IntType
t SubExp
y forall a b. (a -> b) -> a -> b
$
[Char] -> BinOp -> SubExp -> SubExp -> InternaliseM [SubExp]
simpleBinOp [Char]
desc (IntType -> Safety -> BinOp
I.UMod IntType
t Safety
I.Unsafe) SubExp
x SubExp
y
internaliseBinOp SrcLoc
_ [Char]
desc BinOp
E.Mod SubExp
x SubExp
y (E.FloatType FloatType
t) PrimType
_ =
[Char] -> BinOp -> SubExp -> SubExp -> InternaliseM [SubExp]
simpleBinOp [Char]
desc (FloatType -> BinOp
I.FMod FloatType
t) SubExp
x SubExp
y
internaliseBinOp SrcLoc
loc [Char]
desc BinOp
E.Quot SubExp
x SubExp
y (E.Signed IntType
t) PrimType
_ =
forall a.
SrcLoc -> IntType -> SubExp -> InternaliseM a -> InternaliseM a
certifyingNonzero SrcLoc
loc IntType
t SubExp
y forall a b. (a -> b) -> a -> b
$
[Char] -> BinOp -> SubExp -> SubExp -> InternaliseM [SubExp]
simpleBinOp [Char]
desc (IntType -> Safety -> BinOp
I.SQuot IntType
t Safety
I.Unsafe) SubExp
x SubExp
y
internaliseBinOp SrcLoc
loc [Char]
desc BinOp
E.Quot SubExp
x SubExp
y (E.Unsigned IntType
t) PrimType
_ =
forall a.
SrcLoc -> IntType -> SubExp -> InternaliseM a -> InternaliseM a
certifyingNonzero SrcLoc
loc IntType
t SubExp
y forall a b. (a -> b) -> a -> b
$
[Char] -> BinOp -> SubExp -> SubExp -> InternaliseM [SubExp]
simpleBinOp [Char]
desc (IntType -> Safety -> BinOp
I.UDiv IntType
t Safety
I.Unsafe) SubExp
x SubExp
y
internaliseBinOp SrcLoc
loc [Char]
desc BinOp
E.Rem SubExp
x SubExp
y (E.Signed IntType
t) PrimType
_ =
forall a.
SrcLoc -> IntType -> SubExp -> InternaliseM a -> InternaliseM a
certifyingNonzero SrcLoc
loc IntType
t SubExp
y forall a b. (a -> b) -> a -> b
$
[Char] -> BinOp -> SubExp -> SubExp -> InternaliseM [SubExp]
simpleBinOp [Char]
desc (IntType -> Safety -> BinOp
I.SRem IntType
t Safety
I.Unsafe) SubExp
x SubExp
y
internaliseBinOp SrcLoc
loc [Char]
desc BinOp
E.Rem SubExp
x SubExp
y (E.Unsigned IntType
t) PrimType
_ =
forall a.
SrcLoc -> IntType -> SubExp -> InternaliseM a -> InternaliseM a
certifyingNonzero SrcLoc
loc IntType
t SubExp
y forall a b. (a -> b) -> a -> b
$
[Char] -> BinOp -> SubExp -> SubExp -> InternaliseM [SubExp]
simpleBinOp [Char]
desc (IntType -> Safety -> BinOp
I.UMod IntType
t Safety
I.Unsafe) SubExp
x SubExp
y
internaliseBinOp SrcLoc
_ [Char]
desc BinOp
E.ShiftR SubExp
x SubExp
y (E.Signed IntType
t) PrimType
_ =
[Char] -> BinOp -> SubExp -> SubExp -> InternaliseM [SubExp]
simpleBinOp [Char]
desc (IntType -> BinOp
I.AShr IntType
t) SubExp
x SubExp
y
internaliseBinOp SrcLoc
_ [Char]
desc BinOp
E.ShiftR SubExp
x SubExp
y (E.Unsigned IntType
t) PrimType
_ =
[Char] -> BinOp -> SubExp -> SubExp -> InternaliseM [SubExp]
simpleBinOp [Char]
desc (IntType -> BinOp
I.LShr IntType
t) SubExp
x SubExp
y
internaliseBinOp SrcLoc
_ [Char]
desc BinOp
E.ShiftL SubExp
x SubExp
y (E.Signed IntType
t) PrimType
_ =
[Char] -> BinOp -> SubExp -> SubExp -> InternaliseM [SubExp]
simpleBinOp [Char]
desc (IntType -> BinOp
I.Shl IntType
t) SubExp
x SubExp
y
internaliseBinOp SrcLoc
_ [Char]
desc BinOp
E.ShiftL SubExp
x SubExp
y (E.Unsigned IntType
t) PrimType
_ =
[Char] -> BinOp -> SubExp -> SubExp -> InternaliseM [SubExp]
simpleBinOp [Char]
desc (IntType -> BinOp
I.Shl IntType
t) SubExp
x SubExp
y
internaliseBinOp SrcLoc
_ [Char]
desc BinOp
E.Band SubExp
x SubExp
y (E.Signed IntType
t) PrimType
_ =
[Char] -> BinOp -> SubExp -> SubExp -> InternaliseM [SubExp]
simpleBinOp [Char]
desc (IntType -> BinOp
I.And IntType
t) SubExp
x SubExp
y
internaliseBinOp SrcLoc
_ [Char]
desc BinOp
E.Band SubExp
x SubExp
y (E.Unsigned IntType
t) PrimType
_ =
[Char] -> BinOp -> SubExp -> SubExp -> InternaliseM [SubExp]
simpleBinOp [Char]
desc (IntType -> BinOp
I.And IntType
t) SubExp
x SubExp
y
internaliseBinOp SrcLoc
_ [Char]
desc BinOp
E.Xor SubExp
x SubExp
y (E.Signed IntType
t) PrimType
_ =
[Char] -> BinOp -> SubExp -> SubExp -> InternaliseM [SubExp]
simpleBinOp [Char]
desc (IntType -> BinOp
I.Xor IntType
t) SubExp
x SubExp
y
internaliseBinOp SrcLoc
_ [Char]
desc BinOp
E.Xor SubExp
x SubExp
y (E.Unsigned IntType
t) PrimType
_ =
[Char] -> BinOp -> SubExp -> SubExp -> InternaliseM [SubExp]
simpleBinOp [Char]
desc (IntType -> BinOp
I.Xor IntType
t) SubExp
x SubExp
y
internaliseBinOp SrcLoc
_ [Char]
desc BinOp
E.Bor SubExp
x SubExp
y (E.Signed IntType
t) PrimType
_ =
[Char] -> BinOp -> SubExp -> SubExp -> InternaliseM [SubExp]
simpleBinOp [Char]
desc (IntType -> BinOp
I.Or IntType
t) SubExp
x SubExp
y
internaliseBinOp SrcLoc
_ [Char]
desc BinOp
E.Bor SubExp
x SubExp
y (E.Unsigned IntType
t) PrimType
_ =
[Char] -> BinOp -> SubExp -> SubExp -> InternaliseM [SubExp]
simpleBinOp [Char]
desc (IntType -> BinOp
I.Or IntType
t) SubExp
x SubExp
y
internaliseBinOp SrcLoc
_ [Char]
desc BinOp
E.Equal SubExp
x SubExp
y PrimType
t PrimType
_ =
[Char] -> CmpOp -> SubExp -> SubExp -> InternaliseM [SubExp]
simpleCmpOp [Char]
desc (PrimType -> CmpOp
I.CmpEq forall a b. (a -> b) -> a -> b
$ PrimType -> PrimType
internalisePrimType PrimType
t) SubExp
x SubExp
y
internaliseBinOp SrcLoc
_ [Char]
desc BinOp
E.NotEqual SubExp
x SubExp
y PrimType
t PrimType
_ = do
SubExp
eq <- forall (m :: * -> *).
MonadBuilder m =>
[Char] -> Exp (Rep m) -> m SubExp
letSubExp ([Char]
desc forall a. [a] -> [a] -> [a]
++ [Char]
"true") forall a b. (a -> b) -> a -> b
$ forall rep. BasicOp -> Exp rep
I.BasicOp forall a b. (a -> b) -> a -> b
$ CmpOp -> SubExp -> SubExp -> BasicOp
I.CmpOp (PrimType -> CmpOp
I.CmpEq forall a b. (a -> b) -> a -> b
$ PrimType -> PrimType
internalisePrimType PrimType
t) SubExp
x SubExp
y
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap forall (f :: * -> *) a. Applicative f => a -> f a
pure forall a b. (a -> b) -> a -> b
$ forall (m :: * -> *).
MonadBuilder m =>
[Char] -> Exp (Rep m) -> m SubExp
letSubExp [Char]
desc forall a b. (a -> b) -> a -> b
$ forall rep. BasicOp -> Exp rep
I.BasicOp forall a b. (a -> b) -> a -> b
$ UnOp -> SubExp -> BasicOp
I.UnOp UnOp
I.Not SubExp
eq
internaliseBinOp SrcLoc
_ [Char]
desc BinOp
E.Less SubExp
x SubExp
y (E.Signed IntType
t) PrimType
_ =
[Char] -> CmpOp -> SubExp -> SubExp -> InternaliseM [SubExp]
simpleCmpOp [Char]
desc (IntType -> CmpOp
I.CmpSlt IntType
t) SubExp
x SubExp
y
internaliseBinOp SrcLoc
_ [Char]
desc BinOp
E.Less SubExp
x SubExp
y (E.Unsigned IntType
t) PrimType
_ =
[Char] -> CmpOp -> SubExp -> SubExp -> InternaliseM [SubExp]
simpleCmpOp [Char]
desc (IntType -> CmpOp
I.CmpUlt IntType
t) SubExp
x SubExp
y
internaliseBinOp SrcLoc
_ [Char]
desc BinOp
E.Leq SubExp
x SubExp
y (E.Signed IntType
t) PrimType
_ =
[Char] -> CmpOp -> SubExp -> SubExp -> InternaliseM [SubExp]
simpleCmpOp [Char]
desc (IntType -> CmpOp
I.CmpSle IntType
t) SubExp
x SubExp
y
internaliseBinOp SrcLoc
_ [Char]
desc BinOp
E.Leq SubExp
x SubExp
y (E.Unsigned IntType
t) PrimType
_ =
[Char] -> CmpOp -> SubExp -> SubExp -> InternaliseM [SubExp]
simpleCmpOp [Char]
desc (IntType -> CmpOp
I.CmpUle IntType
t) SubExp
x SubExp
y
internaliseBinOp SrcLoc
_ [Char]
desc BinOp
E.Greater SubExp
x SubExp
y (E.Signed IntType
t) PrimType
_ =
[Char] -> CmpOp -> SubExp -> SubExp -> InternaliseM [SubExp]
simpleCmpOp [Char]
desc (IntType -> CmpOp
I.CmpSlt IntType
t) SubExp
y SubExp
x
internaliseBinOp SrcLoc
_ [Char]
desc BinOp
E.Greater SubExp
x SubExp
y (E.Unsigned IntType
t) PrimType
_ =
[Char] -> CmpOp -> SubExp -> SubExp -> InternaliseM [SubExp]
simpleCmpOp [Char]
desc (IntType -> CmpOp
I.CmpUlt IntType
t) SubExp
y SubExp
x
internaliseBinOp SrcLoc
_ [Char]
desc BinOp
E.Geq SubExp
x SubExp
y (E.Signed IntType
t) PrimType
_ =
[Char] -> CmpOp -> SubExp -> SubExp -> InternaliseM [SubExp]
simpleCmpOp [Char]
desc (IntType -> CmpOp
I.CmpSle IntType
t) SubExp
y SubExp
x
internaliseBinOp SrcLoc
_ [Char]
desc BinOp
E.Geq SubExp
x SubExp
y (E.Unsigned IntType
t) PrimType
_ =
[Char] -> CmpOp -> SubExp -> SubExp -> InternaliseM [SubExp]
simpleCmpOp [Char]
desc (IntType -> CmpOp
I.CmpUle IntType
t) SubExp
y SubExp
x
internaliseBinOp SrcLoc
_ [Char]
desc BinOp
E.Less SubExp
x SubExp
y (E.FloatType FloatType
t) PrimType
_ =
[Char] -> CmpOp -> SubExp -> SubExp -> InternaliseM [SubExp]
simpleCmpOp [Char]
desc (FloatType -> CmpOp
I.FCmpLt FloatType
t) SubExp
x SubExp
y
internaliseBinOp SrcLoc
_ [Char]
desc BinOp
E.Leq SubExp
x SubExp
y (E.FloatType FloatType
t) PrimType
_ =
[Char] -> CmpOp -> SubExp -> SubExp -> InternaliseM [SubExp]
simpleCmpOp [Char]
desc (FloatType -> CmpOp
I.FCmpLe FloatType
t) SubExp
x SubExp
y
internaliseBinOp SrcLoc
_ [Char]
desc BinOp
E.Greater SubExp
x SubExp
y (E.FloatType FloatType
t) PrimType
_ =
[Char] -> CmpOp -> SubExp -> SubExp -> InternaliseM [SubExp]
simpleCmpOp [Char]
desc (FloatType -> CmpOp
I.FCmpLt FloatType
t) SubExp
y SubExp
x
internaliseBinOp SrcLoc
_ [Char]
desc BinOp
E.Geq SubExp
x SubExp
y (E.FloatType FloatType
t) PrimType
_ =
[Char] -> CmpOp -> SubExp -> SubExp -> InternaliseM [SubExp]
simpleCmpOp [Char]
desc (FloatType -> CmpOp
I.FCmpLe FloatType
t) SubExp
y SubExp
x
internaliseBinOp SrcLoc
_ [Char]
desc BinOp
E.Less SubExp
x SubExp
y PrimType
E.Bool PrimType
_ =
[Char] -> CmpOp -> SubExp -> SubExp -> InternaliseM [SubExp]
simpleCmpOp [Char]
desc CmpOp
I.CmpLlt SubExp
x SubExp
y
internaliseBinOp SrcLoc
_ [Char]
desc BinOp
E.Leq SubExp
x SubExp
y PrimType
E.Bool PrimType
_ =
[Char] -> CmpOp -> SubExp -> SubExp -> InternaliseM [SubExp]
simpleCmpOp [Char]
desc CmpOp
I.CmpLle SubExp
x SubExp
y
internaliseBinOp SrcLoc
_ [Char]
desc BinOp
E.Greater SubExp
x SubExp
y PrimType
E.Bool PrimType
_ =
[Char] -> CmpOp -> SubExp -> SubExp -> InternaliseM [SubExp]
simpleCmpOp [Char]
desc CmpOp
I.CmpLlt SubExp
y SubExp
x
internaliseBinOp SrcLoc
_ [Char]
desc BinOp
E.Geq SubExp
x SubExp
y PrimType
E.Bool PrimType
_ =
[Char] -> CmpOp -> SubExp -> SubExp -> InternaliseM [SubExp]
simpleCmpOp [Char]
desc CmpOp
I.CmpLle SubExp
y SubExp
x
internaliseBinOp SrcLoc
_ [Char]
_ BinOp
op SubExp
_ SubExp
_ PrimType
t1 PrimType
t2 =
forall a. HasCallStack => [Char] -> a
error forall a b. (a -> b) -> a -> b
$
[Char]
"Invalid binary operator "
forall a. [a] -> [a] -> [a]
++ forall a. Pretty a => a -> [Char]
prettyString BinOp
op
forall a. [a] -> [a] -> [a]
++ [Char]
" with operand types "
forall a. [a] -> [a] -> [a]
++ forall a. Pretty a => a -> [Char]
prettyString PrimType
t1
forall a. [a] -> [a] -> [a]
++ [Char]
", "
forall a. [a] -> [a] -> [a]
++ forall a. Pretty a => a -> [Char]
prettyString PrimType
t2
simpleBinOp ::
String ->
I.BinOp ->
I.SubExp ->
I.SubExp ->
InternaliseM [I.SubExp]
simpleBinOp :: [Char] -> BinOp -> SubExp -> SubExp -> InternaliseM [SubExp]
simpleBinOp [Char]
desc BinOp
bop SubExp
x SubExp
y =
forall (m :: * -> *).
MonadBuilder m =>
[Char] -> Exp (Rep m) -> m [SubExp]
letTupExp' [Char]
desc forall a b. (a -> b) -> a -> b
$ forall rep. BasicOp -> Exp rep
I.BasicOp forall a b. (a -> b) -> a -> b
$ BinOp -> SubExp -> SubExp -> BasicOp
I.BinOp BinOp
bop SubExp
x SubExp
y
simpleCmpOp ::
String ->
I.CmpOp ->
I.SubExp ->
I.SubExp ->
InternaliseM [I.SubExp]
simpleCmpOp :: [Char] -> CmpOp -> SubExp -> SubExp -> InternaliseM [SubExp]
simpleCmpOp [Char]
desc CmpOp
op SubExp
x SubExp
y =
forall (m :: * -> *).
MonadBuilder m =>
[Char] -> Exp (Rep m) -> m [SubExp]
letTupExp' [Char]
desc forall a b. (a -> b) -> a -> b
$ forall rep. BasicOp -> Exp rep
I.BasicOp forall a b. (a -> b) -> a -> b
$ CmpOp -> SubExp -> SubExp -> BasicOp
I.CmpOp CmpOp
op SubExp
x SubExp
y
data Function
= FunctionName (E.QualName VName)
| FunctionHole E.PatType SrcLoc
deriving (Int -> Function -> ShowS
[Function] -> ShowS
Function -> [Char]
forall a.
(Int -> a -> ShowS) -> (a -> [Char]) -> ([a] -> ShowS) -> Show a
showList :: [Function] -> ShowS
$cshowList :: [Function] -> ShowS
show :: Function -> [Char]
$cshow :: Function -> [Char]
showsPrec :: Int -> Function -> ShowS
$cshowsPrec :: Int -> Function -> ShowS
Show)
findFuncall :: E.AppExp -> (Function, [(E.Exp, Maybe VName)])
findFuncall :: AppExp -> (Function, [(Exp, Maybe VName)])
findFuncall (E.Apply Exp
f Exp
arg (Info (Diet
_, Maybe VName
argext)) SrcLoc
_)
| E.AppExp AppExp
f_e Info AppRes
_ <- Exp
f =
let (Function
f_e', [(Exp, Maybe VName)]
args) = AppExp -> (Function, [(Exp, Maybe VName)])
findFuncall AppExp
f_e
in (Function
f_e', [(Exp, Maybe VName)]
args forall a. [a] -> [a] -> [a]
++ [(Exp
arg, Maybe VName
argext)])
| E.Var QualName VName
fname Info PatType
_ SrcLoc
_ <- Exp
f =
(QualName VName -> Function
FunctionName QualName VName
fname, [(Exp
arg, Maybe VName
argext)])
| E.Hole (Info PatType
t) SrcLoc
loc <- Exp
f =
(PatType -> SrcLoc -> Function
FunctionHole PatType
t SrcLoc
loc, [(Exp
arg, Maybe VName
argext)])
findFuncall AppExp
e =
forall a. HasCallStack => [Char] -> a
error forall a b. (a -> b) -> a -> b
$ [Char]
"Invalid function expression in application:\n" forall a. [a] -> [a] -> [a]
++ forall a. Pretty a => a -> [Char]
prettyString AppExp
e
bodyExtType :: Body SOACS -> InternaliseM [ExtType]
bodyExtType :: Body SOACS -> InternaliseM [ExtType]
bodyExtType (Body BodyDec SOACS
_ Stms SOACS
stms Result
res) =
[VName] -> [ExtType] -> [ExtType]
existentialiseExtTypes (forall k a. Map k a -> [k]
M.keys Scope SOACS
stmsscope) forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall u. [TypeBase Shape u] -> [TypeBase ExtShape u]
staticShapes
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> forall rep (m :: * -> *) a.
ExtendedScope rep m a -> Scope rep -> m a
extendedScope (forall (t :: * -> *) (f :: * -> *) a b.
(Traversable t, Applicative f) =>
(a -> f b) -> t a -> f (t b)
traverse forall t (m :: * -> *).
HasScope t m =>
SubExpRes -> m (TypeBase Shape NoUniqueness)
subExpResType Result
res) Scope SOACS
stmsscope
where
stmsscope :: Scope SOACS
stmsscope = forall rep a. Scoped rep a => a -> Scope rep
scopeOf Stms SOACS
stms
internaliseLambda :: InternaliseLambda
internaliseLambda :: InternaliseLambda
internaliseLambda (E.Parens Exp
e SrcLoc
_) [TypeBase Shape NoUniqueness]
rowtypes =
InternaliseLambda
internaliseLambda Exp
e [TypeBase Shape NoUniqueness]
rowtypes
internaliseLambda (E.Lambda [PatBase Info VName]
params Exp
body Maybe (TypeExp VName)
_ (Info (Set Alias
_, RetType [VName]
_ TypeBase Size ()
rettype)) SrcLoc
_) [TypeBase Shape NoUniqueness]
rowtypes =
forall a.
[PatBase Info VName]
-> [TypeBase Shape NoUniqueness]
-> ([LParam SOACS] -> InternaliseM a)
-> InternaliseM a
bindingLambdaParams [PatBase Info VName]
params [TypeBase Shape NoUniqueness]
rowtypes forall a b. (a -> b) -> a -> b
$ \[LParam SOACS]
params' -> do
Body SOACS
body' <- [Char] -> Exp -> InternaliseM (Body SOACS)
internaliseBody [Char]
"lam" Exp
body
[TypeBase Shape NoUniqueness]
rettype' <- forall shape u.
TypeBase Size ()
-> [TypeBase shape u] -> InternaliseM [TypeBase Shape NoUniqueness]
internaliseLambdaReturnType TypeBase Size ()
rettype forall (m :: * -> *) a b. Monad m => (a -> m b) -> m a -> m b
=<< Body SOACS -> InternaliseM [ExtType]
bodyExtType Body SOACS
body'
forall (f :: * -> *) a. Applicative f => a -> f a
pure ([LParam SOACS]
params', Body SOACS
body', [TypeBase Shape NoUniqueness]
rettype')
internaliseLambda Exp
e [TypeBase Shape NoUniqueness]
_ = forall a. HasCallStack => [Char] -> a
error forall a b. (a -> b) -> a -> b
$ [Char]
"internaliseLambda: unexpected expression:\n" forall a. [a] -> [a] -> [a]
++ forall a. Pretty a => a -> [Char]
prettyString Exp
e
internaliseLambdaCoerce :: E.Exp -> [Type] -> InternaliseM (I.Lambda SOACS)
internaliseLambdaCoerce :: Exp -> [TypeBase Shape NoUniqueness] -> InternaliseM (Lambda SOACS)
internaliseLambdaCoerce Exp
lam [TypeBase Shape NoUniqueness]
argtypes = do
([Param (TypeBase Shape NoUniqueness)]
params, Body SOACS
body, [TypeBase Shape NoUniqueness]
rettype) <- InternaliseLambda
internaliseLambda Exp
lam [TypeBase Shape NoUniqueness]
argtypes
forall (m :: * -> *).
MonadBuilder m =>
[LParam (Rep m)] -> m Result -> m (Lambda (Rep m))
mkLambda [Param (TypeBase Shape NoUniqueness)]
params forall a b. (a -> b) -> a -> b
$
ErrorMsg SubExp
-> SrcLoc
-> [TypeBase Shape NoUniqueness]
-> Result
-> InternaliseM Result
ensureResultShape
(forall a. [ErrorMsgPart a] -> ErrorMsg a
ErrorMsg [forall a. Text -> ErrorMsgPart a
ErrorString Text
"unexpected lambda result size"])
(forall a. Located a => a -> SrcLoc
srclocOf Exp
lam)
[TypeBase Shape NoUniqueness]
rettype
forall (m :: * -> *) a b. Monad m => (a -> m b) -> m a -> m b
=<< forall (m :: * -> *). MonadBuilder m => Body (Rep m) -> m Result
bodyBind Body SOACS
body
isOverloadedFunction ::
E.QualName VName ->
String ->
SrcLoc ->
Maybe ([(E.StructType, [SubExp])] -> InternaliseM [SubExp])
isOverloadedFunction :: QualName VName
-> [Char]
-> SrcLoc
-> Maybe ([(TypeBase Size (), [SubExp])] -> InternaliseM [SubExp])
isOverloadedFunction QualName VName
qname [Char]
desc SrcLoc
loc = do
forall (f :: * -> *). Alternative f => Bool -> f ()
guard forall a b. (a -> b) -> a -> b
$ VName -> Int
baseTag (forall vn. QualName vn -> vn
qualLeaf QualName VName
qname) forall a. Ord a => a -> a -> Bool
<= Int
maxIntrinsicTag
[Char]
-> Maybe ([(TypeBase Size (), [SubExp])] -> InternaliseM [SubExp])
handle forall a b. (a -> b) -> a -> b
$ VName -> [Char]
baseString forall a b. (a -> b) -> a -> b
$ forall vn. QualName vn -> vn
qualLeaf QualName VName
qname
where
handle :: [Char]
-> Maybe ([(TypeBase Size (), [SubExp])] -> InternaliseM [SubExp])
handle [Char]
op
| Just SubExp -> InternaliseM [SubExp]
cmp_f <- [Char] -> Maybe (SubExp -> InternaliseM [SubExp])
isEqlOp [Char]
op = forall a. a -> Maybe a
Just forall a b. (a -> b) -> a -> b
$ \[(TypeBase Size ()
_, [SubExp]
xe'), (TypeBase Size ()
_, [SubExp]
ye')] -> do
[SubExp]
rs <- forall (m :: * -> *) a b c.
Applicative m =>
(a -> b -> m c) -> [a] -> [b] -> m [c]
zipWithM SubExp -> SubExp -> InternaliseM SubExp
doComparison [SubExp]
xe' [SubExp]
ye'
SubExp -> InternaliseM [SubExp]
cmp_f forall (m :: * -> *) a b. Monad m => (a -> m b) -> m a -> m b
=<< forall (m :: * -> *).
MonadBuilder m =>
[Char] -> Exp (Rep m) -> m SubExp
letSubExp [Char]
"eq" forall (m :: * -> *) a b. Monad m => (a -> m b) -> m a -> m b
=<< forall (m :: * -> *). MonadBuilder m => [SubExp] -> m (Exp (Rep m))
eAll [SubExp]
rs
where
isEqlOp :: [Char] -> Maybe (SubExp -> InternaliseM [SubExp])
isEqlOp [Char]
"!=" = forall a. a -> Maybe a
Just forall a b. (a -> b) -> a -> b
$ \SubExp
eq ->
forall (m :: * -> *).
MonadBuilder m =>
[Char] -> Exp (Rep m) -> m [SubExp]
letTupExp' [Char]
desc forall a b. (a -> b) -> a -> b
$ forall rep. BasicOp -> Exp rep
I.BasicOp forall a b. (a -> b) -> a -> b
$ UnOp -> SubExp -> BasicOp
I.UnOp UnOp
I.Not SubExp
eq
isEqlOp [Char]
"==" = forall a. a -> Maybe a
Just forall a b. (a -> b) -> a -> b
$ \SubExp
eq ->
forall (f :: * -> *) a. Applicative f => a -> f a
pure [SubExp
eq]
isEqlOp [Char]
_ = forall a. Maybe a
Nothing
doComparison :: SubExp -> SubExp -> InternaliseM SubExp
doComparison SubExp
x SubExp
y = do
TypeBase Shape NoUniqueness
x_t <- forall t (m :: * -> *).
HasScope t m =>
SubExp -> m (TypeBase Shape NoUniqueness)
I.subExpType SubExp
x
TypeBase Shape NoUniqueness
y_t <- forall t (m :: * -> *).
HasScope t m =>
SubExp -> m (TypeBase Shape NoUniqueness)
I.subExpType SubExp
y
case TypeBase Shape NoUniqueness
x_t of
I.Prim PrimType
t -> forall (m :: * -> *).
MonadBuilder m =>
[Char] -> Exp (Rep m) -> m SubExp
letSubExp [Char]
desc forall a b. (a -> b) -> a -> b
$ forall rep. BasicOp -> Exp rep
I.BasicOp forall a b. (a -> b) -> a -> b
$ CmpOp -> SubExp -> SubExp -> BasicOp
I.CmpOp (PrimType -> CmpOp
I.CmpEq PrimType
t) SubExp
x SubExp
y
TypeBase Shape NoUniqueness
_ -> do
let x_dims :: [SubExp]
x_dims = forall u. TypeBase Shape u -> [SubExp]
I.arrayDims TypeBase Shape NoUniqueness
x_t
y_dims :: [SubExp]
y_dims = forall u. TypeBase Shape u -> [SubExp]
I.arrayDims TypeBase Shape NoUniqueness
y_t
[SubExp]
dims_match <- forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
t a -> (a -> m b) -> m (t b)
forM (forall a b. [a] -> [b] -> [(a, b)]
zip [SubExp]
x_dims [SubExp]
y_dims) forall a b. (a -> b) -> a -> b
$ \(SubExp
x_dim, SubExp
y_dim) ->
forall (m :: * -> *).
MonadBuilder m =>
[Char] -> Exp (Rep m) -> m SubExp
letSubExp [Char]
"dim_eq" forall a b. (a -> b) -> a -> b
$ forall rep. BasicOp -> Exp rep
I.BasicOp forall a b. (a -> b) -> a -> b
$ CmpOp -> SubExp -> SubExp -> BasicOp
I.CmpOp (PrimType -> CmpOp
I.CmpEq PrimType
int64) SubExp
x_dim SubExp
y_dim
SubExp
shapes_match <- forall (m :: * -> *).
MonadBuilder m =>
[Char] -> Exp (Rep m) -> m SubExp
letSubExp [Char]
"shapes_match" forall (m :: * -> *) a b. Monad m => (a -> m b) -> m a -> m b
=<< forall (m :: * -> *). MonadBuilder m => [SubExp] -> m (Exp (Rep m))
eAll [SubExp]
dims_match
let compare_elems_body :: InternaliseM (Body SOACS)
compare_elems_body = forall rep (m :: * -> *) somerep.
(Buildable rep, MonadFreshNames m, HasScope somerep m,
SameScope somerep rep) =>
Builder rep (Body rep) -> m (Body rep)
runBodyBuilder forall a b. (a -> b) -> a -> b
$ do
SubExp
x_num_elems <-
forall (m :: * -> *).
MonadBuilder m =>
[Char] -> Exp (Rep m) -> m SubExp
letSubExp [Char]
"x_num_elems"
forall (m :: * -> *) a b. Monad m => (a -> m b) -> m a -> m b
=<< forall (m :: * -> *).
MonadBuilder m =>
BinOp -> SubExp -> [SubExp] -> m (Exp (Rep m))
foldBinOp (IntType -> Overflow -> BinOp
I.Mul IntType
Int64 Overflow
I.OverflowUndef) (forall v. IsValue v => v -> SubExp
constant (Int64
1 :: Int64)) [SubExp]
x_dims
VName
x' <- forall (m :: * -> *).
MonadBuilder m =>
[Char] -> Exp (Rep m) -> m VName
letExp [Char]
"x" forall a b. (a -> b) -> a -> b
$ forall rep. BasicOp -> Exp rep
I.BasicOp forall a b. (a -> b) -> a -> b
$ SubExp -> BasicOp
I.SubExp SubExp
x
VName
y' <- forall (m :: * -> *).
MonadBuilder m =>
[Char] -> Exp (Rep m) -> m VName
letExp [Char]
"x" forall a b. (a -> b) -> a -> b
$ forall rep. BasicOp -> Exp rep
I.BasicOp forall a b. (a -> b) -> a -> b
$ SubExp -> BasicOp
I.SubExp SubExp
y
VName
x_flat <-
forall (m :: * -> *).
MonadBuilder m =>
[Char] -> Exp (Rep m) -> m VName
letExp [Char]
"x_flat" forall a b. (a -> b) -> a -> b
$ forall rep. BasicOp -> Exp rep
I.BasicOp forall a b. (a -> b) -> a -> b
$ ReshapeKind -> Shape -> VName -> BasicOp
I.Reshape ReshapeKind
I.ReshapeArbitrary (forall d. [d] -> ShapeBase d
I.Shape [SubExp
x_num_elems]) VName
x'
VName
y_flat <-
forall (m :: * -> *).
MonadBuilder m =>
[Char] -> Exp (Rep m) -> m VName
letExp [Char]
"y_flat" forall a b. (a -> b) -> a -> b
$ forall rep. BasicOp -> Exp rep
I.BasicOp forall a b. (a -> b) -> a -> b
$ ReshapeKind -> Shape -> VName -> BasicOp
I.Reshape ReshapeKind
I.ReshapeArbitrary (forall d. [d] -> ShapeBase d
I.Shape [SubExp
x_num_elems]) VName
y'
Lambda SOACS
cmp_lam <- forall (m :: * -> *).
(MonadBuilder m, Buildable (Rep m)) =>
CmpOp -> m (Lambda (Rep m))
cmpOpLambda forall a b. (a -> b) -> a -> b
$ PrimType -> CmpOp
I.CmpEq (forall shape u. TypeBase shape u -> PrimType
elemType TypeBase Shape NoUniqueness
x_t)
VName
cmps <-
forall (m :: * -> *).
MonadBuilder m =>
[Char] -> Exp (Rep m) -> m VName
letExp [Char]
"cmps" forall a b. (a -> b) -> a -> b
$
forall rep. Op rep -> Exp rep
I.Op forall a b. (a -> b) -> a -> b
$
forall rep. SubExp -> [VName] -> ScremaForm rep -> SOAC rep
I.Screma SubExp
x_num_elems [VName
x_flat, VName
y_flat] (forall rep. Lambda rep -> ScremaForm rep
I.mapSOAC Lambda SOACS
cmp_lam)
Lambda SOACS
and_lam <- forall (m :: * -> *).
(MonadBuilder m, Buildable (Rep m)) =>
BinOp -> PrimType -> m (Lambda (Rep m))
binOpLambda BinOp
I.LogAnd PrimType
I.Bool
ScremaForm SOACS
reduce <- forall rep (m :: * -> *).
(Buildable rep, MonadFreshNames m) =>
[Reduce rep] -> m (ScremaForm rep)
I.reduceSOAC [forall rep. Commutativity -> Lambda rep -> [SubExp] -> Reduce rep
Reduce Commutativity
Commutative Lambda SOACS
and_lam [forall v. IsValue v => v -> SubExp
constant Bool
True]]
SubExp
all_equal <- forall (m :: * -> *).
MonadBuilder m =>
[Char] -> Exp (Rep m) -> m SubExp
letSubExp [Char]
"all_equal" forall a b. (a -> b) -> a -> b
$ forall rep. Op rep -> Exp rep
I.Op forall a b. (a -> b) -> a -> b
$ forall rep. SubExp -> [VName] -> ScremaForm rep -> SOAC rep
I.Screma SubExp
x_num_elems [VName
cmps] ScremaForm SOACS
reduce
forall (f :: * -> *) a. Applicative f => a -> f a
pure forall a b. (a -> b) -> a -> b
$ forall rep. Buildable rep => [SubExp] -> Body rep
resultBody [SubExp
all_equal]
forall (m :: * -> *).
MonadBuilder m =>
[Char] -> Exp (Rep m) -> m SubExp
letSubExp [Char]
"arrays_equal"
forall (m :: * -> *) a b. Monad m => (a -> m b) -> m a -> m b
=<< forall (m :: * -> *).
(MonadBuilder m, BranchType (Rep m) ~ ExtType) =>
m (Exp (Rep m))
-> m (Body (Rep m)) -> m (Body (Rep m)) -> m (Exp (Rep m))
eIf (forall (m :: * -> *). MonadBuilder m => SubExp -> m (Exp (Rep m))
eSubExp SubExp
shapes_match) InternaliseM (Body SOACS)
compare_elems_body (forall (m :: * -> *).
MonadBuilder m =>
[SubExp] -> m (Body (Rep m))
resultBodyM [forall v. IsValue v => v -> SubExp
constant Bool
False])
handle [Char]
name
| Just BinOp
bop <- forall (t :: * -> *) a. Foldable t => (a -> Bool) -> t a -> Maybe a
find (([Char]
name ==) forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall a. Pretty a => a -> [Char]
prettyString) [forall a. Bounded a => a
minBound .. forall a. Bounded a => a
maxBound :: E.BinOp] =
forall a. a -> Maybe a
Just forall a b. (a -> b) -> a -> b
$ \[(TypeBase Size ()
x_t, [SubExp
x']), (TypeBase Size ()
y_t, [SubExp
y'])] ->
case (TypeBase Size ()
x_t, TypeBase Size ()
y_t) of
(E.Scalar (E.Prim PrimType
t1), E.Scalar (E.Prim PrimType
t2)) ->
SrcLoc
-> [Char]
-> BinOp
-> SubExp
-> SubExp
-> PrimType
-> PrimType
-> InternaliseM [SubExp]
internaliseBinOp SrcLoc
loc [Char]
desc BinOp
bop SubExp
x' SubExp
y' PrimType
t1 PrimType
t2
(TypeBase Size (), TypeBase Size ())
_ -> forall a. HasCallStack => [Char] -> a
error [Char]
"Futhark.Internalise.internaliseExp: non-primitive type in BinOp."
handle [Char]
_ = forall a. Maybe a
Nothing
isIntrinsicFunction ::
E.QualName VName ->
[E.Exp] ->
SrcLoc ->
Maybe (String -> InternaliseM [SubExp])
isIntrinsicFunction :: QualName VName
-> [Exp] -> SrcLoc -> Maybe ([Char] -> InternaliseM [SubExp])
isIntrinsicFunction QualName VName
qname [Exp]
args SrcLoc
loc = do
forall (f :: * -> *). Alternative f => Bool -> f ()
guard forall a b. (a -> b) -> a -> b
$ VName -> Int
baseTag (forall vn. QualName vn -> vn
qualLeaf QualName VName
qname) forall a. Ord a => a -> a -> Bool
<= Int
maxIntrinsicTag
let handlers :: [[Exp] -> [Char] -> Maybe ([Char] -> InternaliseM [SubExp])]
handlers =
[ forall {a}.
(Eq a, IsString a) =>
[Exp] -> a -> Maybe ([Char] -> InternaliseM [SubExp])
handleSign,
forall {f :: * -> *}.
Applicative f =>
[Exp] -> [Char] -> Maybe ([Char] -> InternaliseM (f SubExp))
handleOps,
[Exp] -> [Char] -> Maybe ([Char] -> InternaliseM [SubExp])
handleSOACs,
forall {a}.
(Eq a, IsString a) =>
[Exp] -> a -> Maybe ([Char] -> InternaliseM [SubExp])
handleAccs,
forall {a}.
(IsString a, Eq a) =>
[Exp] -> a -> Maybe ([Char] -> InternaliseM [SubExp])
handleAD,
[Exp] -> [Char] -> Maybe ([Char] -> InternaliseM [SubExp])
handleRest
]
forall (t :: * -> *) (m :: * -> *) a.
(Foldable t, MonadPlus m) =>
t (m a) -> m a
msum [[Exp] -> [Char] -> Maybe ([Char] -> InternaliseM [SubExp])
h [Exp]
args forall a b. (a -> b) -> a -> b
$ VName -> [Char]
baseString forall a b. (a -> b) -> a -> b
$ forall vn. QualName vn -> vn
qualLeaf QualName VName
qname | [Exp] -> [Char] -> Maybe ([Char] -> InternaliseM [SubExp])
h <- [[Exp] -> [Char] -> Maybe ([Char] -> InternaliseM [SubExp])]
handlers]
where
handleSign :: [Exp] -> a -> Maybe ([Char] -> InternaliseM [SubExp])
handleSign [Exp
x] a
"sign_i8" = forall a. a -> Maybe a
Just forall a b. (a -> b) -> a -> b
$ IntType -> Exp -> [Char] -> InternaliseM [SubExp]
toSigned IntType
I.Int8 Exp
x
handleSign [Exp
x] a
"sign_i16" = forall a. a -> Maybe a
Just forall a b. (a -> b) -> a -> b
$ IntType -> Exp -> [Char] -> InternaliseM [SubExp]
toSigned IntType
I.Int16 Exp
x
handleSign [Exp
x] a
"sign_i32" = forall a. a -> Maybe a
Just forall a b. (a -> b) -> a -> b
$ IntType -> Exp -> [Char] -> InternaliseM [SubExp]
toSigned IntType
I.Int32 Exp
x
handleSign [Exp
x] a
"sign_i64" = forall a. a -> Maybe a
Just forall a b. (a -> b) -> a -> b
$ IntType -> Exp -> [Char] -> InternaliseM [SubExp]
toSigned IntType
I.Int64 Exp
x
handleSign [Exp
x] a
"unsign_i8" = forall a. a -> Maybe a
Just forall a b. (a -> b) -> a -> b
$ IntType -> Exp -> [Char] -> InternaliseM [SubExp]
toUnsigned IntType
I.Int8 Exp
x
handleSign [Exp
x] a
"unsign_i16" = forall a. a -> Maybe a
Just forall a b. (a -> b) -> a -> b
$ IntType -> Exp -> [Char] -> InternaliseM [SubExp]
toUnsigned IntType
I.Int16 Exp
x
handleSign [Exp
x] a
"unsign_i32" = forall a. a -> Maybe a
Just forall a b. (a -> b) -> a -> b
$ IntType -> Exp -> [Char] -> InternaliseM [SubExp]
toUnsigned IntType
I.Int32 Exp
x
handleSign [Exp
x] a
"unsign_i64" = forall a. a -> Maybe a
Just forall a b. (a -> b) -> a -> b
$ IntType -> Exp -> [Char] -> InternaliseM [SubExp]
toUnsigned IntType
I.Int64 Exp
x
handleSign [Exp]
_ a
_ = forall a. Maybe a
Nothing
handleOps :: [Exp] -> [Char] -> Maybe ([Char] -> InternaliseM (f SubExp))
handleOps [Exp
x] [Char]
s
| Just UnOp
unop <- forall (t :: * -> *) a. Foldable t => (a -> Bool) -> t a -> Maybe a
find ((forall a. Eq a => a -> a -> Bool
== [Char]
s) forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall a. Pretty a => a -> [Char]
prettyString) [UnOp]
allUnOps = forall a. a -> Maybe a
Just forall a b. (a -> b) -> a -> b
$ \[Char]
desc -> do
SubExp
x' <- [Char] -> Exp -> InternaliseM SubExp
internaliseExp1 [Char]
"x" Exp
x
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap forall (f :: * -> *) a. Applicative f => a -> f a
pure forall a b. (a -> b) -> a -> b
$ forall (m :: * -> *).
MonadBuilder m =>
[Char] -> Exp (Rep m) -> m SubExp
letSubExp [Char]
desc forall a b. (a -> b) -> a -> b
$ forall rep. BasicOp -> Exp rep
I.BasicOp forall a b. (a -> b) -> a -> b
$ UnOp -> SubExp -> BasicOp
I.UnOp UnOp
unop SubExp
x'
handleOps [TupLit [Exp
x, Exp
y] SrcLoc
_] [Char]
s
| Just BinOp
bop <- forall (t :: * -> *) a. Foldable t => (a -> Bool) -> t a -> Maybe a
find ((forall a. Eq a => a -> a -> Bool
== [Char]
s) forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall a. Pretty a => a -> [Char]
prettyString) [BinOp]
allBinOps = forall a. a -> Maybe a
Just forall a b. (a -> b) -> a -> b
$ \[Char]
desc -> do
SubExp
x' <- [Char] -> Exp -> InternaliseM SubExp
internaliseExp1 [Char]
"x" Exp
x
SubExp
y' <- [Char] -> Exp -> InternaliseM SubExp
internaliseExp1 [Char]
"y" Exp
y
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap forall (f :: * -> *) a. Applicative f => a -> f a
pure forall a b. (a -> b) -> a -> b
$ forall (m :: * -> *).
MonadBuilder m =>
[Char] -> Exp (Rep m) -> m SubExp
letSubExp [Char]
desc forall a b. (a -> b) -> a -> b
$ forall rep. BasicOp -> Exp rep
I.BasicOp forall a b. (a -> b) -> a -> b
$ BinOp -> SubExp -> SubExp -> BasicOp
I.BinOp BinOp
bop SubExp
x' SubExp
y'
| Just CmpOp
cmp <- forall (t :: * -> *) a. Foldable t => (a -> Bool) -> t a -> Maybe a
find ((forall a. Eq a => a -> a -> Bool
== [Char]
s) forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall a. Pretty a => a -> [Char]
prettyString) [CmpOp]
allCmpOps = forall a. a -> Maybe a
Just forall a b. (a -> b) -> a -> b
$ \[Char]
desc -> do
SubExp
x' <- [Char] -> Exp -> InternaliseM SubExp
internaliseExp1 [Char]
"x" Exp
x
SubExp
y' <- [Char] -> Exp -> InternaliseM SubExp
internaliseExp1 [Char]
"y" Exp
y
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap forall (f :: * -> *) a. Applicative f => a -> f a
pure forall a b. (a -> b) -> a -> b
$ forall (m :: * -> *).
MonadBuilder m =>
[Char] -> Exp (Rep m) -> m SubExp
letSubExp [Char]
desc forall a b. (a -> b) -> a -> b
$ forall rep. BasicOp -> Exp rep
I.BasicOp forall a b. (a -> b) -> a -> b
$ CmpOp -> SubExp -> SubExp -> BasicOp
I.CmpOp CmpOp
cmp SubExp
x' SubExp
y'
handleOps [Exp
x] [Char]
s
| Just ConvOp
conv <- forall (t :: * -> *) a. Foldable t => (a -> Bool) -> t a -> Maybe a
find ((forall a. Eq a => a -> a -> Bool
== [Char]
s) forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall a. Pretty a => a -> [Char]
prettyString) [ConvOp]
allConvOps = forall a. a -> Maybe a
Just forall a b. (a -> b) -> a -> b
$ \[Char]
desc -> do
SubExp
x' <- [Char] -> Exp -> InternaliseM SubExp
internaliseExp1 [Char]
"x" Exp
x
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap forall (f :: * -> *) a. Applicative f => a -> f a
pure forall a b. (a -> b) -> a -> b
$ forall (m :: * -> *).
MonadBuilder m =>
[Char] -> Exp (Rep m) -> m SubExp
letSubExp [Char]
desc forall a b. (a -> b) -> a -> b
$ forall rep. BasicOp -> Exp rep
I.BasicOp forall a b. (a -> b) -> a -> b
$ ConvOp -> SubExp -> BasicOp
I.ConvOp ConvOp
conv SubExp
x'
handleOps [Exp]
_ [Char]
_ = forall a. Maybe a
Nothing
handleSOACs :: [Exp] -> [Char] -> Maybe ([Char] -> InternaliseM [SubExp])
handleSOACs [TupLit [Exp
lam, Exp
arr] SrcLoc
_] [Char]
"map" = forall a. a -> Maybe a
Just forall a b. (a -> b) -> a -> b
$ \[Char]
desc -> do
[VName]
arr' <- [Char] -> Exp -> InternaliseM [VName]
internaliseExpToVars [Char]
"map_arr" Exp
arr
[TypeBase Shape NoUniqueness]
arr_ts <- forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
(a -> m b) -> t a -> m (t b)
mapM forall rep (m :: * -> *).
HasScope rep m =>
VName -> m (TypeBase Shape NoUniqueness)
lookupType [VName]
arr'
Lambda SOACS
lam' <- Exp -> [TypeBase Shape NoUniqueness] -> InternaliseM (Lambda SOACS)
internaliseLambdaCoerce Exp
lam forall a b. (a -> b) -> a -> b
$ forall a b. (a -> b) -> [a] -> [b]
map forall u. TypeBase Shape u -> TypeBase Shape u
rowType [TypeBase Shape NoUniqueness]
arr_ts
let w :: SubExp
w = forall u. Int -> [TypeBase Shape u] -> SubExp
arraysSize Int
0 [TypeBase Shape NoUniqueness]
arr_ts
forall (m :: * -> *).
MonadBuilder m =>
[Char] -> Exp (Rep m) -> m [SubExp]
letTupExp' [Char]
desc forall a b. (a -> b) -> a -> b
$ forall rep. Op rep -> Exp rep
I.Op forall a b. (a -> b) -> a -> b
$ forall rep. SubExp -> [VName] -> ScremaForm rep -> SOAC rep
I.Screma SubExp
w [VName]
arr' (forall rep. Lambda rep -> ScremaForm rep
I.mapSOAC Lambda SOACS
lam')
handleSOACs [TupLit [Exp
k, Exp
lam, Exp
arr] SrcLoc
_] [Char]
"partition" = do
Int
k' <- forall a b. (Integral a, Num b) => a -> b
fromIntegral forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> forall {vn}. ExpBase Info vn -> Maybe Int32
fromInt32 Exp
k
forall a. a -> Maybe a
Just forall a b. (a -> b) -> a -> b
$ \[Char]
_desc -> do
[VName]
arrs <- [Char] -> Exp -> InternaliseM [VName]
internaliseExpToVars [Char]
"partition_input" Exp
arr
Lambda SOACS
lam' <- InternaliseLambda
-> Int -> Exp -> [SubExp] -> InternaliseM (Lambda SOACS)
internalisePartitionLambda InternaliseLambda
internaliseLambda Int
k' Exp
lam forall a b. (a -> b) -> a -> b
$ forall a b. (a -> b) -> [a] -> [b]
map VName -> SubExp
I.Var [VName]
arrs
forall a b c. (a -> b -> c) -> (a, b) -> c
uncurry forall a. [a] -> [a] -> [a]
(++) forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Int -> Lambda SOACS -> [VName] -> InternaliseM ([SubExp], [SubExp])
partitionWithSOACS (forall a b. (Integral a, Num b) => a -> b
fromIntegral Int
k') Lambda SOACS
lam' [VName]
arrs
where
fromInt32 :: ExpBase Info vn -> Maybe Int32
fromInt32 (Literal (SignedValue (Int32Value Int32
k')) SrcLoc
_) = forall a. a -> Maybe a
Just Int32
k'
fromInt32 (IntLit Integer
k' (Info (E.Scalar (E.Prim (E.Signed IntType
Int32)))) SrcLoc
_) = forall a. a -> Maybe a
Just forall a b. (a -> b) -> a -> b
$ forall a. Num a => Integer -> a
fromInteger Integer
k'
fromInt32 ExpBase Info vn
_ = forall a. Maybe a
Nothing
handleSOACs [TupLit [Exp
lam, Exp
ne, Exp
arr] SrcLoc
_] [Char]
"reduce" = forall a. a -> Maybe a
Just forall a b. (a -> b) -> a -> b
$ \[Char]
desc ->
[Char]
-> [Char]
-> (SubExp
-> Lambda SOACS
-> [SubExp]
-> [VName]
-> InternaliseM (SOAC SOACS))
-> (Exp, Exp, Exp, SrcLoc)
-> InternaliseM [SubExp]
internaliseScanOrReduce [Char]
desc [Char]
"reduce" forall {f :: * -> *} {rep}.
(Buildable rep, MonadFreshNames f) =>
SubExp -> Lambda rep -> [SubExp] -> [VName] -> f (SOAC rep)
reduce (Exp
lam, Exp
ne, Exp
arr, SrcLoc
loc)
where
reduce :: SubExp -> Lambda rep -> [SubExp] -> [VName] -> f (SOAC rep)
reduce SubExp
w Lambda rep
red_lam [SubExp]
nes [VName]
arrs =
forall rep. SubExp -> [VName] -> ScremaForm rep -> SOAC rep
I.Screma SubExp
w [VName]
arrs
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> forall rep (m :: * -> *).
(Buildable rep, MonadFreshNames m) =>
[Reduce rep] -> m (ScremaForm rep)
I.reduceSOAC [forall rep. Commutativity -> Lambda rep -> [SubExp] -> Reduce rep
Reduce Commutativity
Noncommutative Lambda rep
red_lam [SubExp]
nes]
handleSOACs [TupLit [Exp
lam, Exp
ne, Exp
arr] SrcLoc
_] [Char]
"reduce_comm" = forall a. a -> Maybe a
Just forall a b. (a -> b) -> a -> b
$ \[Char]
desc ->
[Char]
-> [Char]
-> (SubExp
-> Lambda SOACS
-> [SubExp]
-> [VName]
-> InternaliseM (SOAC SOACS))
-> (Exp, Exp, Exp, SrcLoc)
-> InternaliseM [SubExp]
internaliseScanOrReduce [Char]
desc [Char]
"reduce" forall {f :: * -> *} {rep}.
(Buildable rep, MonadFreshNames f) =>
SubExp -> Lambda rep -> [SubExp] -> [VName] -> f (SOAC rep)
reduce (Exp
lam, Exp
ne, Exp
arr, SrcLoc
loc)
where
reduce :: SubExp -> Lambda rep -> [SubExp] -> [VName] -> f (SOAC rep)
reduce SubExp
w Lambda rep
red_lam [SubExp]
nes [VName]
arrs =
forall rep. SubExp -> [VName] -> ScremaForm rep -> SOAC rep
I.Screma SubExp
w [VName]
arrs
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> forall rep (m :: * -> *).
(Buildable rep, MonadFreshNames m) =>
[Reduce rep] -> m (ScremaForm rep)
I.reduceSOAC [forall rep. Commutativity -> Lambda rep -> [SubExp] -> Reduce rep
Reduce Commutativity
Commutative Lambda rep
red_lam [SubExp]
nes]
handleSOACs [TupLit [Exp
lam, Exp
ne, Exp
arr] SrcLoc
_] [Char]
"scan" = forall a. a -> Maybe a
Just forall a b. (a -> b) -> a -> b
$ \[Char]
desc ->
[Char]
-> [Char]
-> (SubExp
-> Lambda SOACS
-> [SubExp]
-> [VName]
-> InternaliseM (SOAC SOACS))
-> (Exp, Exp, Exp, SrcLoc)
-> InternaliseM [SubExp]
internaliseScanOrReduce [Char]
desc [Char]
"scan" forall {f :: * -> *} {rep}.
(Buildable rep, MonadFreshNames f) =>
SubExp -> Lambda rep -> [SubExp] -> [VName] -> f (SOAC rep)
reduce (Exp
lam, Exp
ne, Exp
arr, SrcLoc
loc)
where
reduce :: SubExp -> Lambda rep -> [SubExp] -> [VName] -> f (SOAC rep)
reduce SubExp
w Lambda rep
scan_lam [SubExp]
nes [VName]
arrs =
forall rep. SubExp -> [VName] -> ScremaForm rep -> SOAC rep
I.Screma SubExp
w [VName]
arrs forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> forall rep (m :: * -> *).
(Buildable rep, MonadFreshNames m) =>
[Scan rep] -> m (ScremaForm rep)
I.scanSOAC [forall rep. Lambda rep -> [SubExp] -> Scan rep
Scan Lambda rep
scan_lam [SubExp]
nes]
handleSOACs [TupLit [Exp
rf, Exp
dest, Exp
op, Exp
ne, Exp
buckets, Exp
img] SrcLoc
_] [Char]
"hist_1d" = forall a. a -> Maybe a
Just forall a b. (a -> b) -> a -> b
$ \[Char]
desc ->
Int
-> [Char]
-> Exp
-> Exp
-> Exp
-> Exp
-> Exp
-> Exp
-> SrcLoc
-> InternaliseM [SubExp]
internaliseHist Int
1 [Char]
desc Exp
rf Exp
dest Exp
op Exp
ne Exp
buckets Exp
img SrcLoc
loc
handleSOACs [TupLit [Exp
rf, Exp
dest, Exp
op, Exp
ne, Exp
buckets, Exp
img] SrcLoc
_] [Char]
"hist_2d" = forall a. a -> Maybe a
Just forall a b. (a -> b) -> a -> b
$ \[Char]
desc ->
Int
-> [Char]
-> Exp
-> Exp
-> Exp
-> Exp
-> Exp
-> Exp
-> SrcLoc
-> InternaliseM [SubExp]
internaliseHist Int
2 [Char]
desc Exp
rf Exp
dest Exp
op Exp
ne Exp
buckets Exp
img SrcLoc
loc
handleSOACs [TupLit [Exp
rf, Exp
dest, Exp
op, Exp
ne, Exp
buckets, Exp
img] SrcLoc
_] [Char]
"hist_3d" = forall a. a -> Maybe a
Just forall a b. (a -> b) -> a -> b
$ \[Char]
desc ->
Int
-> [Char]
-> Exp
-> Exp
-> Exp
-> Exp
-> Exp
-> Exp
-> SrcLoc
-> InternaliseM [SubExp]
internaliseHist Int
3 [Char]
desc Exp
rf Exp
dest Exp
op Exp
ne Exp
buckets Exp
img SrcLoc
loc
handleSOACs [Exp]
_ [Char]
_ = forall a. Maybe a
Nothing
handleAccs :: [Exp] -> a -> Maybe ([Char] -> InternaliseM [SubExp])
handleAccs [TupLit [Exp
dest, Exp
f, Exp
bs] SrcLoc
_] a
"scatter_stream" = forall a. a -> Maybe a
Just forall a b. (a -> b) -> a -> b
$ \[Char]
desc ->
[Char]
-> Exp -> Maybe (Exp, Exp) -> Exp -> Exp -> InternaliseM [SubExp]
internaliseStreamAcc [Char]
desc Exp
dest forall a. Maybe a
Nothing Exp
f Exp
bs
handleAccs [TupLit [Exp
dest, Exp
op, Exp
ne, Exp
f, Exp
bs] SrcLoc
_] a
"hist_stream" = forall a. a -> Maybe a
Just forall a b. (a -> b) -> a -> b
$ \[Char]
desc ->
[Char]
-> Exp -> Maybe (Exp, Exp) -> Exp -> Exp -> InternaliseM [SubExp]
internaliseStreamAcc [Char]
desc Exp
dest (forall a. a -> Maybe a
Just (Exp
op, Exp
ne)) Exp
f Exp
bs
handleAccs [TupLit [Exp
acc, Exp
i, Exp
v] SrcLoc
_] a
"acc_write" = forall a. a -> Maybe a
Just forall a b. (a -> b) -> a -> b
$ \[Char]
desc -> do
VName
acc' <- forall a. [a] -> a
head forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> [Char] -> Exp -> InternaliseM [VName]
internaliseExpToVars [Char]
"acc" Exp
acc
SubExp
i' <- [Char] -> Exp -> InternaliseM SubExp
internaliseExp1 [Char]
"acc_i" Exp
i
[SubExp]
vs <- [Char] -> Exp -> InternaliseM [SubExp]
internaliseExp [Char]
"acc_v" Exp
v
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap forall (f :: * -> *) a. Applicative f => a -> f a
pure forall a b. (a -> b) -> a -> b
$ forall (m :: * -> *).
MonadBuilder m =>
[Char] -> Exp (Rep m) -> m SubExp
letSubExp [Char]
desc forall a b. (a -> b) -> a -> b
$ forall rep. BasicOp -> Exp rep
BasicOp forall a b. (a -> b) -> a -> b
$ VName -> [SubExp] -> [SubExp] -> BasicOp
UpdateAcc VName
acc' [SubExp
i'] [SubExp]
vs
handleAccs [Exp]
_ a
_ = forall a. Maybe a
Nothing
handleAD :: [Exp] -> a -> Maybe ([Char] -> InternaliseM [SubExp])
handleAD [TupLit [Exp
f, Exp
x, Exp
v] SrcLoc
_] a
fname
| a
fname forall (t :: * -> *) a. (Foldable t, Eq a) => a -> t a -> Bool
`elem` [a
"jvp2", a
"vjp2"] = forall a. a -> Maybe a
Just forall a b. (a -> b) -> a -> b
$ \[Char]
desc -> do
[SubExp]
x' <- [Char] -> Exp -> InternaliseM [SubExp]
internaliseExp [Char]
"ad_x" Exp
x
[SubExp]
v' <- [Char] -> Exp -> InternaliseM [SubExp]
internaliseExp [Char]
"ad_v" Exp
v
Lambda SOACS
lam <- Exp -> [TypeBase Shape NoUniqueness] -> InternaliseM (Lambda SOACS)
internaliseLambdaCoerce Exp
f forall (m :: * -> *) a b. Monad m => (a -> m b) -> m a -> m b
=<< forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
(a -> m b) -> t a -> m (t b)
mapM forall t (m :: * -> *).
HasScope t m =>
SubExp -> m (TypeBase Shape NoUniqueness)
subExpType [SubExp]
x'
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap (forall a b. (a -> b) -> [a] -> [b]
map VName -> SubExp
I.Var) forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall (m :: * -> *).
MonadBuilder m =>
[Char] -> Exp (Rep m) -> m [VName]
letTupExp [Char]
desc forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall rep. Op rep -> Exp rep
Op forall a b. (a -> b) -> a -> b
$
case a
fname of
a
"jvp2" -> forall rep. Lambda rep -> [SubExp] -> [SubExp] -> SOAC rep
JVP Lambda SOACS
lam [SubExp]
x' [SubExp]
v'
a
_ -> forall rep. Lambda rep -> [SubExp] -> [SubExp] -> SOAC rep
VJP Lambda SOACS
lam [SubExp]
x' [SubExp]
v'
handleAD [Exp]
_ a
_ = forall a. Maybe a
Nothing
handleRest :: [Exp] -> [Char] -> Maybe ([Char] -> InternaliseM [SubExp])
handleRest [E.TupLit [Exp
a, Exp
si, Exp
v] SrcLoc
_] [Char]
"scatter" = forall a. a -> Maybe a
Just forall a b. (a -> b) -> a -> b
$ Int -> Exp -> Exp -> Exp -> [Char] -> InternaliseM [SubExp]
scatterF Int
1 Exp
a Exp
si Exp
v
handleRest [E.TupLit [Exp
a, Exp
si, Exp
v] SrcLoc
_] [Char]
"scatter_2d" = forall a. a -> Maybe a
Just forall a b. (a -> b) -> a -> b
$ Int -> Exp -> Exp -> Exp -> [Char] -> InternaliseM [SubExp]
scatterF Int
2 Exp
a Exp
si Exp
v
handleRest [E.TupLit [Exp
a, Exp
si, Exp
v] SrcLoc
_] [Char]
"scatter_3d" = forall a. a -> Maybe a
Just forall a b. (a -> b) -> a -> b
$ Int -> Exp -> Exp -> Exp -> [Char] -> InternaliseM [SubExp]
scatterF Int
3 Exp
a Exp
si Exp
v
handleRest [E.TupLit [Exp
n, Exp
m, Exp
arr] SrcLoc
_] [Char]
"unflatten" = forall a. a -> Maybe a
Just forall a b. (a -> b) -> a -> b
$ \[Char]
desc -> do
[VName]
arrs <- [Char] -> Exp -> InternaliseM [VName]
internaliseExpToVars [Char]
"unflatten_arr" Exp
arr
SubExp
n' <- [Char] -> Exp -> InternaliseM SubExp
internaliseExp1 [Char]
"n" Exp
n
SubExp
m' <- [Char] -> Exp -> InternaliseM SubExp
internaliseExp1 [Char]
"m" Exp
m
SubExp
old_dim <- forall u. Int -> [TypeBase Shape u] -> SubExp
I.arraysSize Int
0 forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
(a -> m b) -> t a -> m (t b)
mapM forall rep (m :: * -> *).
HasScope rep m =>
VName -> m (TypeBase Shape NoUniqueness)
lookupType [VName]
arrs
SubExp
dim_ok <-
forall (m :: * -> *).
MonadBuilder m =>
[Char] -> Exp (Rep m) -> m SubExp
letSubExp [Char]
"dim_ok"
forall (m :: * -> *) a b. Monad m => (a -> m b) -> m a -> m b
=<< forall (m :: * -> *).
MonadBuilder m =>
CmpOp -> m (Exp (Rep m)) -> m (Exp (Rep m)) -> m (Exp (Rep m))
eCmpOp
(PrimType -> CmpOp
I.CmpEq PrimType
I.int64)
(forall (m :: * -> *).
MonadBuilder m =>
BinOp -> m (Exp (Rep m)) -> m (Exp (Rep m)) -> m (Exp (Rep m))
eBinOp (IntType -> Overflow -> BinOp
I.Mul IntType
Int64 Overflow
I.OverflowUndef) (forall (m :: * -> *). MonadBuilder m => SubExp -> m (Exp (Rep m))
eSubExp SubExp
n') (forall (m :: * -> *). MonadBuilder m => SubExp -> m (Exp (Rep m))
eSubExp SubExp
m'))
(forall (m :: * -> *). MonadBuilder m => SubExp -> m (Exp (Rep m))
eSubExp SubExp
old_dim)
Certs
dim_ok_cert <-
[Char] -> SubExp -> ErrorMsg SubExp -> SrcLoc -> InternaliseM Certs
assert
[Char]
"dim_ok_cert"
SubExp
dim_ok
( forall a. [ErrorMsgPart a] -> ErrorMsg a
ErrorMsg
[ ErrorMsgPart SubExp
"Cannot unflatten array of shape [",
forall a. PrimType -> a -> ErrorMsgPart a
ErrorVal PrimType
int64 SubExp
old_dim,
ErrorMsgPart SubExp
"] to array of shape [",
forall a. PrimType -> a -> ErrorMsgPart a
ErrorVal PrimType
int64 SubExp
n',
ErrorMsgPart SubExp
"][",
forall a. PrimType -> a -> ErrorMsgPart a
ErrorVal PrimType
int64 SubExp
m',
ErrorMsgPart SubExp
"]"
]
)
SrcLoc
loc
forall (m :: * -> *) a. MonadBuilder m => Certs -> m a -> m a
certifying Certs
dim_ok_cert forall a b. (a -> b) -> a -> b
$
forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
t a -> (a -> m b) -> m (t b)
forM [VName]
arrs forall a b. (a -> b) -> a -> b
$ \VName
arr' -> do
TypeBase Shape NoUniqueness
arr_t <- forall rep (m :: * -> *).
HasScope rep m =>
VName -> m (TypeBase Shape NoUniqueness)
lookupType VName
arr'
forall (m :: * -> *).
MonadBuilder m =>
[Char] -> Exp (Rep m) -> m SubExp
letSubExp [Char]
desc forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall rep. BasicOp -> Exp rep
I.BasicOp forall a b. (a -> b) -> a -> b
$
ReshapeKind -> Shape -> VName -> BasicOp
I.Reshape
ReshapeKind
I.ReshapeArbitrary
(Shape -> Int -> Shape -> Shape
reshapeOuter (forall d. [d] -> ShapeBase d
I.Shape [SubExp
n', SubExp
m']) Int
1 forall a b. (a -> b) -> a -> b
$ forall shape u. ArrayShape shape => TypeBase shape u -> shape
I.arrayShape TypeBase Shape NoUniqueness
arr_t)
VName
arr'
handleRest [Exp
arr] [Char]
"flatten" = forall a. a -> Maybe a
Just forall a b. (a -> b) -> a -> b
$ \[Char]
desc -> do
[VName]
arrs <- [Char] -> Exp -> InternaliseM [VName]
internaliseExpToVars [Char]
"flatten_arr" Exp
arr
forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
t a -> (a -> m b) -> m (t b)
forM [VName]
arrs forall a b. (a -> b) -> a -> b
$ \VName
arr' -> do
TypeBase Shape NoUniqueness
arr_t <- forall rep (m :: * -> *).
HasScope rep m =>
VName -> m (TypeBase Shape NoUniqueness)
lookupType VName
arr'
let n :: SubExp
n = forall u. Int -> TypeBase Shape u -> SubExp
arraySize Int
0 TypeBase Shape NoUniqueness
arr_t
m :: SubExp
m = forall u. Int -> TypeBase Shape u -> SubExp
arraySize Int
1 TypeBase Shape NoUniqueness
arr_t
SubExp
k <- forall (m :: * -> *).
MonadBuilder m =>
[Char] -> Exp (Rep m) -> m SubExp
letSubExp [Char]
"flat_dim" forall a b. (a -> b) -> a -> b
$ forall rep. BasicOp -> Exp rep
I.BasicOp forall a b. (a -> b) -> a -> b
$ BinOp -> SubExp -> SubExp -> BasicOp
I.BinOp (IntType -> Overflow -> BinOp
Mul IntType
Int64 Overflow
I.OverflowUndef) SubExp
n SubExp
m
forall (m :: * -> *).
MonadBuilder m =>
[Char] -> Exp (Rep m) -> m SubExp
letSubExp [Char]
desc forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall rep. BasicOp -> Exp rep
I.BasicOp forall a b. (a -> b) -> a -> b
$
ReshapeKind -> Shape -> VName -> BasicOp
I.Reshape
ReshapeKind
I.ReshapeArbitrary
(Shape -> Int -> Shape -> Shape
reshapeOuter (forall d. [d] -> ShapeBase d
I.Shape [SubExp
k]) Int
2 forall a b. (a -> b) -> a -> b
$ forall shape u. ArrayShape shape => TypeBase shape u -> shape
I.arrayShape TypeBase Shape NoUniqueness
arr_t)
VName
arr'
handleRest [TupLit [Exp
x, Exp
y] SrcLoc
_] [Char]
"concat" = forall a. a -> Maybe a
Just forall a b. (a -> b) -> a -> b
$ \[Char]
desc -> do
[VName]
xs <- [Char] -> Exp -> InternaliseM [VName]
internaliseExpToVars [Char]
"concat_x" Exp
x
[VName]
ys <- [Char] -> Exp -> InternaliseM [VName]
internaliseExpToVars [Char]
"concat_y" Exp
y
SubExp
outer_size <- forall u. Int -> [TypeBase Shape u] -> SubExp
arraysSize Int
0 forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
(a -> m b) -> t a -> m (t b)
mapM forall rep (m :: * -> *).
HasScope rep m =>
VName -> m (TypeBase Shape NoUniqueness)
lookupType [VName]
xs
let sumdims :: SubExp -> SubExp -> m SubExp
sumdims SubExp
xsize SubExp
ysize =
forall (m :: * -> *).
MonadBuilder m =>
[Char] -> Exp (Rep m) -> m SubExp
letSubExp [Char]
"conc_tmp" forall a b. (a -> b) -> a -> b
$
forall rep. BasicOp -> Exp rep
I.BasicOp forall a b. (a -> b) -> a -> b
$
BinOp -> SubExp -> SubExp -> BasicOp
I.BinOp (IntType -> Overflow -> BinOp
I.Add IntType
I.Int64 Overflow
I.OverflowUndef) SubExp
xsize SubExp
ysize
SubExp
ressize <-
forall (t :: * -> *) (m :: * -> *) b a.
(Foldable t, Monad m) =>
(b -> a -> m b) -> b -> t a -> m b
foldM forall {m :: * -> *}.
MonadBuilder m =>
SubExp -> SubExp -> m SubExp
sumdims SubExp
outer_size
forall (m :: * -> *) a b. Monad m => (a -> m b) -> m a -> m b
=<< forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
(a -> m b) -> t a -> m (t b)
mapM (forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap (forall u. Int -> [TypeBase Shape u] -> SubExp
arraysSize Int
0) forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
(a -> m b) -> t a -> m (t b)
mapM forall rep (m :: * -> *).
HasScope rep m =>
VName -> m (TypeBase Shape NoUniqueness)
lookupType) [[VName]
ys]
let conc :: VName -> VName -> Exp SOACS
conc VName
xarr VName
yarr =
forall rep. BasicOp -> Exp rep
I.BasicOp forall a b. (a -> b) -> a -> b
$ Int -> NonEmpty VName -> SubExp -> BasicOp
I.Concat Int
0 (VName
xarr forall a. a -> [a] -> NonEmpty a
:| [VName
yarr]) SubExp
ressize
forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
(a -> m b) -> t a -> m (t b)
mapM (forall (m :: * -> *).
MonadBuilder m =>
[Char] -> Exp (Rep m) -> m SubExp
letSubExp [Char]
desc) forall a b. (a -> b) -> a -> b
$ forall a b c. (a -> b -> c) -> [a] -> [b] -> [c]
zipWith VName -> VName -> Exp SOACS
conc [VName]
xs [VName]
ys
handleRest [TupLit [Exp
offset, Exp
e] SrcLoc
_] [Char]
"rotate" = forall a. a -> Maybe a
Just forall a b. (a -> b) -> a -> b
$ \[Char]
desc -> do
SubExp
offset' <- [Char] -> Exp -> InternaliseM SubExp
internaliseExp1 [Char]
"rotation_offset" Exp
offset
[Char]
-> Exp -> (VName -> InternaliseM BasicOp) -> InternaliseM [SubExp]
internaliseOperation [Char]
desc Exp
e forall a b. (a -> b) -> a -> b
$ \VName
v -> do
Int
r <- forall shape u. ArrayShape shape => TypeBase shape u -> Int
I.arrayRank forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> forall rep (m :: * -> *).
HasScope rep m =>
VName -> m (TypeBase Shape NoUniqueness)
lookupType VName
v
let zero :: SubExp
zero = IntType -> Integer -> SubExp
intConst IntType
Int64 Integer
0
offsets :: [SubExp]
offsets = SubExp
offset' forall a. a -> [a] -> [a]
: forall a. Int -> a -> [a]
replicate (Int
r forall a. Num a => a -> a -> a
- Int
1) SubExp
zero
forall (f :: * -> *) a. Applicative f => a -> f a
pure forall a b. (a -> b) -> a -> b
$ [SubExp] -> VName -> BasicOp
I.Rotate [SubExp]
offsets VName
v
handleRest [Exp
e] [Char]
"transpose" = forall a. a -> Maybe a
Just forall a b. (a -> b) -> a -> b
$ \[Char]
desc ->
[Char]
-> Exp -> (VName -> InternaliseM BasicOp) -> InternaliseM [SubExp]
internaliseOperation [Char]
desc Exp
e forall a b. (a -> b) -> a -> b
$ \VName
v -> do
Int
r <- forall shape u. ArrayShape shape => TypeBase shape u -> Int
I.arrayRank forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> forall rep (m :: * -> *).
HasScope rep m =>
VName -> m (TypeBase Shape NoUniqueness)
lookupType VName
v
forall (f :: * -> *) a. Applicative f => a -> f a
pure forall a b. (a -> b) -> a -> b
$ [Int] -> VName -> BasicOp
I.Rearrange ([Int
1, Int
0] forall a. [a] -> [a] -> [a]
++ [Int
2 .. Int
r forall a. Num a => a -> a -> a
- Int
1]) VName
v
handleRest [TupLit [Exp
x, Exp
y] SrcLoc
_] [Char]
"zip" = forall a. a -> Maybe a
Just forall a b. (a -> b) -> a -> b
$ \[Char]
desc ->
forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
(a -> m b) -> t a -> m (t b)
mapM (forall (m :: * -> *).
MonadBuilder m =>
[Char] -> Exp (Rep m) -> m SubExp
letSubExp [Char]
"zip_copy" forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall rep. BasicOp -> Exp rep
BasicOp forall b c a. (b -> c) -> (a -> b) -> a -> c
. VName -> BasicOp
Copy)
forall (m :: * -> *) a b. Monad m => (a -> m b) -> m a -> m b
=<< ( forall a. [a] -> [a] -> [a]
(++)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> [Char] -> Exp -> InternaliseM [VName]
internaliseExpToVars ([Char]
desc forall a. [a] -> [a] -> [a]
++ [Char]
"_zip_x") Exp
x
forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> [Char] -> Exp -> InternaliseM [VName]
internaliseExpToVars ([Char]
desc forall a. [a] -> [a] -> [a]
++ [Char]
"_zip_y") Exp
y
)
handleRest [Exp
x] [Char]
"unzip" = forall a. a -> Maybe a
Just forall a b. (a -> b) -> a -> b
$ forall a b c. (a -> b -> c) -> b -> a -> c
flip [Char] -> Exp -> InternaliseM [SubExp]
internaliseExp Exp
x
handleRest [TupLit [Exp
arr, Exp
offset, Exp
n1, Exp
s1, Exp
n2, Exp
s2] SrcLoc
_] [Char]
"flat_index_2d" = forall a. a -> Maybe a
Just forall a b. (a -> b) -> a -> b
$ \[Char]
desc -> do
[Char]
-> SrcLoc -> Exp -> Exp -> [(Exp, Exp)] -> InternaliseM [SubExp]
flatIndexHelper [Char]
desc SrcLoc
loc Exp
arr Exp
offset [(Exp
n1, Exp
s1), (Exp
n2, Exp
s2)]
handleRest [TupLit [Exp
arr1, Exp
offset, Exp
s1, Exp
s2, Exp
arr2] SrcLoc
_] [Char]
"flat_update_2d" = forall a. a -> Maybe a
Just forall a b. (a -> b) -> a -> b
$ \[Char]
desc -> do
[Char]
-> SrcLoc -> Exp -> Exp -> [Exp] -> Exp -> InternaliseM [SubExp]
flatUpdateHelper [Char]
desc SrcLoc
loc Exp
arr1 Exp
offset [Exp
s1, Exp
s2] Exp
arr2
handleRest [TupLit [Exp
arr, Exp
offset, Exp
n1, Exp
s1, Exp
n2, Exp
s2, Exp
n3, Exp
s3] SrcLoc
_] [Char]
"flat_index_3d" = forall a. a -> Maybe a
Just forall a b. (a -> b) -> a -> b
$ \[Char]
desc -> do
[Char]
-> SrcLoc -> Exp -> Exp -> [(Exp, Exp)] -> InternaliseM [SubExp]
flatIndexHelper [Char]
desc SrcLoc
loc Exp
arr Exp
offset [(Exp
n1, Exp
s1), (Exp
n2, Exp
s2), (Exp
n3, Exp
s3)]
handleRest [TupLit [Exp
arr1, Exp
offset, Exp
s1, Exp
s2, Exp
s3, Exp
arr2] SrcLoc
_] [Char]
"flat_update_3d" = forall a. a -> Maybe a
Just forall a b. (a -> b) -> a -> b
$ \[Char]
desc -> do
[Char]
-> SrcLoc -> Exp -> Exp -> [Exp] -> Exp -> InternaliseM [SubExp]
flatUpdateHelper [Char]
desc SrcLoc
loc Exp
arr1 Exp
offset [Exp
s1, Exp
s2, Exp
s3] Exp
arr2
handleRest [TupLit [Exp
arr, Exp
offset, Exp
n1, Exp
s1, Exp
n2, Exp
s2, Exp
n3, Exp
s3, Exp
n4, Exp
s4] SrcLoc
_] [Char]
"flat_index_4d" = forall a. a -> Maybe a
Just forall a b. (a -> b) -> a -> b
$ \[Char]
desc -> do
[Char]
-> SrcLoc -> Exp -> Exp -> [(Exp, Exp)] -> InternaliseM [SubExp]
flatIndexHelper [Char]
desc SrcLoc
loc Exp
arr Exp
offset [(Exp
n1, Exp
s1), (Exp
n2, Exp
s2), (Exp
n3, Exp
s3), (Exp
n4, Exp
s4)]
handleRest [TupLit [Exp
arr1, Exp
offset, Exp
s1, Exp
s2, Exp
s3, Exp
s4, Exp
arr2] SrcLoc
_] [Char]
"flat_update_4d" = forall a. a -> Maybe a
Just forall a b. (a -> b) -> a -> b
$ \[Char]
desc -> do
[Char]
-> SrcLoc -> Exp -> Exp -> [Exp] -> Exp -> InternaliseM [SubExp]
flatUpdateHelper [Char]
desc SrcLoc
loc Exp
arr1 Exp
offset [Exp
s1, Exp
s2, Exp
s3, Exp
s4] Exp
arr2
handleRest [Exp]
_ [Char]
_ = forall a. Maybe a
Nothing
toSigned :: IntType -> Exp -> [Char] -> InternaliseM [SubExp]
toSigned IntType
int_to Exp
e [Char]
desc = do
SubExp
e' <- [Char] -> Exp -> InternaliseM SubExp
internaliseExp1 [Char]
"trunc_arg" Exp
e
case Exp -> PatType
E.typeOf Exp
e of
E.Scalar (E.Prim PrimType
E.Bool) ->
forall (m :: * -> *).
MonadBuilder m =>
[Char] -> Exp (Rep m) -> m [SubExp]
letTupExp' [Char]
desc
forall (m :: * -> *) a b. Monad m => (a -> m b) -> m a -> m b
=<< forall (m :: * -> *).
(MonadBuilder m, BranchType (Rep m) ~ ExtType) =>
m (Exp (Rep m))
-> m (Body (Rep m)) -> m (Body (Rep m)) -> m (Exp (Rep m))
eIf
(forall (m :: * -> *). MonadBuilder m => SubExp -> m (Exp (Rep m))
eSubExp SubExp
e')
(forall (m :: * -> *).
MonadBuilder m =>
[SubExp] -> m (Body (Rep m))
resultBodyM [IntType -> Integer -> SubExp
intConst IntType
int_to Integer
1])
(forall (m :: * -> *).
MonadBuilder m =>
[SubExp] -> m (Body (Rep m))
resultBodyM [IntType -> Integer -> SubExp
intConst IntType
int_to Integer
0])
E.Scalar (E.Prim (E.Signed IntType
int_from)) ->
forall (m :: * -> *).
MonadBuilder m =>
[Char] -> Exp (Rep m) -> m [SubExp]
letTupExp' [Char]
desc forall a b. (a -> b) -> a -> b
$ forall rep. BasicOp -> Exp rep
I.BasicOp forall a b. (a -> b) -> a -> b
$ ConvOp -> SubExp -> BasicOp
I.ConvOp (IntType -> IntType -> ConvOp
I.SExt IntType
int_from IntType
int_to) SubExp
e'
E.Scalar (E.Prim (E.Unsigned IntType
int_from)) ->
forall (m :: * -> *).
MonadBuilder m =>
[Char] -> Exp (Rep m) -> m [SubExp]
letTupExp' [Char]
desc forall a b. (a -> b) -> a -> b
$ forall rep. BasicOp -> Exp rep
I.BasicOp forall a b. (a -> b) -> a -> b
$ ConvOp -> SubExp -> BasicOp
I.ConvOp (IntType -> IntType -> ConvOp
I.ZExt IntType
int_from IntType
int_to) SubExp
e'
E.Scalar (E.Prim (E.FloatType FloatType
float_from)) ->
forall (m :: * -> *).
MonadBuilder m =>
[Char] -> Exp (Rep m) -> m [SubExp]
letTupExp' [Char]
desc forall a b. (a -> b) -> a -> b
$ forall rep. BasicOp -> Exp rep
I.BasicOp forall a b. (a -> b) -> a -> b
$ ConvOp -> SubExp -> BasicOp
I.ConvOp (FloatType -> IntType -> ConvOp
I.FPToSI FloatType
float_from IntType
int_to) SubExp
e'
PatType
_ -> forall a. HasCallStack => [Char] -> a
error [Char]
"Futhark.Internalise: non-numeric type in ToSigned"
toUnsigned :: IntType -> Exp -> [Char] -> InternaliseM [SubExp]
toUnsigned IntType
int_to Exp
e [Char]
desc = do
SubExp
e' <- [Char] -> Exp -> InternaliseM SubExp
internaliseExp1 [Char]
"trunc_arg" Exp
e
case Exp -> PatType
E.typeOf Exp
e of
E.Scalar (E.Prim PrimType
E.Bool) ->
forall (m :: * -> *).
MonadBuilder m =>
[Char] -> Exp (Rep m) -> m [SubExp]
letTupExp' [Char]
desc
forall (m :: * -> *) a b. Monad m => (a -> m b) -> m a -> m b
=<< forall (m :: * -> *).
(MonadBuilder m, BranchType (Rep m) ~ ExtType) =>
m (Exp (Rep m))
-> m (Body (Rep m)) -> m (Body (Rep m)) -> m (Exp (Rep m))
eIf
(forall (m :: * -> *). MonadBuilder m => SubExp -> m (Exp (Rep m))
eSubExp SubExp
e')
(forall (m :: * -> *).
MonadBuilder m =>
[SubExp] -> m (Body (Rep m))
resultBodyM [IntType -> Integer -> SubExp
intConst IntType
int_to Integer
1])
(forall (m :: * -> *).
MonadBuilder m =>
[SubExp] -> m (Body (Rep m))
resultBodyM [IntType -> Integer -> SubExp
intConst IntType
int_to Integer
0])
E.Scalar (E.Prim (E.Signed IntType
int_from)) ->
forall (m :: * -> *).
MonadBuilder m =>
[Char] -> Exp (Rep m) -> m [SubExp]
letTupExp' [Char]
desc forall a b. (a -> b) -> a -> b
$ forall rep. BasicOp -> Exp rep
I.BasicOp forall a b. (a -> b) -> a -> b
$ ConvOp -> SubExp -> BasicOp
I.ConvOp (IntType -> IntType -> ConvOp
I.ZExt IntType
int_from IntType
int_to) SubExp
e'
E.Scalar (E.Prim (E.Unsigned IntType
int_from)) ->
forall (m :: * -> *).
MonadBuilder m =>
[Char] -> Exp (Rep m) -> m [SubExp]
letTupExp' [Char]
desc forall a b. (a -> b) -> a -> b
$ forall rep. BasicOp -> Exp rep
I.BasicOp forall a b. (a -> b) -> a -> b
$ ConvOp -> SubExp -> BasicOp
I.ConvOp (IntType -> IntType -> ConvOp
I.ZExt IntType
int_from IntType
int_to) SubExp
e'
E.Scalar (E.Prim (E.FloatType FloatType
float_from)) ->
forall (m :: * -> *).
MonadBuilder m =>
[Char] -> Exp (Rep m) -> m [SubExp]
letTupExp' [Char]
desc forall a b. (a -> b) -> a -> b
$ forall rep. BasicOp -> Exp rep
I.BasicOp forall a b. (a -> b) -> a -> b
$ ConvOp -> SubExp -> BasicOp
I.ConvOp (FloatType -> IntType -> ConvOp
I.FPToUI FloatType
float_from IntType
int_to) SubExp
e'
PatType
_ -> forall a. HasCallStack => [Char] -> a
error [Char]
"Futhark.Internalise.internaliseExp: non-numeric type in ToUnsigned"
scatterF :: Int -> Exp -> Exp -> Exp -> [Char] -> InternaliseM [SubExp]
scatterF Int
dim Exp
a Exp
si Exp
v [Char]
desc = do
[VName]
si' <- [Char] -> Exp -> InternaliseM [VName]
internaliseExpToVars [Char]
"write_arg_i" Exp
si
[VName]
svs <- [Char] -> Exp -> InternaliseM [VName]
internaliseExpToVars [Char]
"write_arg_v" Exp
v
[VName]
sas <- [Char] -> Exp -> InternaliseM [VName]
internaliseExpToVars [Char]
"write_arg_a" Exp
a
SubExp
si_w <- forall u. Int -> [TypeBase Shape u] -> SubExp
I.arraysSize Int
0 forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
(a -> m b) -> t a -> m (t b)
mapM forall rep (m :: * -> *).
HasScope rep m =>
VName -> m (TypeBase Shape NoUniqueness)
lookupType [VName]
si'
[TypeBase Shape NoUniqueness]
sv_ts <- forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
(a -> m b) -> t a -> m (t b)
mapM forall rep (m :: * -> *).
HasScope rep m =>
VName -> m (TypeBase Shape NoUniqueness)
lookupType [VName]
svs
[VName]
svs' <- forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
t a -> (a -> m b) -> m (t b)
forM (forall a b. [a] -> [b] -> [(a, b)]
zip [VName]
svs [TypeBase Shape NoUniqueness]
sv_ts) forall a b. (a -> b) -> a -> b
$ \(VName
sv, TypeBase Shape NoUniqueness
sv_t) -> do
let sv_shape :: Shape
sv_shape = forall shape u. ArrayShape shape => TypeBase shape u -> shape
I.arrayShape TypeBase Shape NoUniqueness
sv_t
sv_w :: SubExp
sv_w = forall u. Int -> TypeBase Shape u -> SubExp
arraySize Int
0 TypeBase Shape NoUniqueness
sv_t
SubExp
cmp <-
forall (m :: * -> *).
MonadBuilder m =>
[Char] -> Exp (Rep m) -> m SubExp
letSubExp [Char]
"write_cmp" forall a b. (a -> b) -> a -> b
$
forall rep. BasicOp -> Exp rep
I.BasicOp forall a b. (a -> b) -> a -> b
$
CmpOp -> SubExp -> SubExp -> BasicOp
I.CmpOp (PrimType -> CmpOp
I.CmpEq PrimType
I.int64) SubExp
si_w SubExp
sv_w
Certs
c <-
[Char] -> SubExp -> ErrorMsg SubExp -> SrcLoc -> InternaliseM Certs
assert
[Char]
"write_cert"
SubExp
cmp
ErrorMsg SubExp
"length of index and value array does not match"
SrcLoc
loc
forall (m :: * -> *) a. MonadBuilder m => Certs -> m a -> m a
certifying Certs
c forall a b. (a -> b) -> a -> b
$
forall (m :: * -> *).
MonadBuilder m =>
[Char] -> Exp (Rep m) -> m VName
letExp (VName -> [Char]
baseString VName
sv forall a. [a] -> [a] -> [a]
++ [Char]
"_write_sv") forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall rep. BasicOp -> Exp rep
I.BasicOp forall a b. (a -> b) -> a -> b
$
ReshapeKind -> Shape -> VName -> BasicOp
I.Reshape ReshapeKind
I.ReshapeCoerce (Shape -> Int -> Shape -> Shape
reshapeOuter (forall d. [d] -> ShapeBase d
I.Shape [SubExp
si_w]) Int
1 Shape
sv_shape) VName
sv
[TypeBase Shape NoUniqueness]
indexType <- forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap forall u. TypeBase Shape u -> TypeBase Shape u
rowType forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
(a -> m b) -> t a -> m (t b)
mapM forall rep (m :: * -> *).
HasScope rep m =>
VName -> m (TypeBase Shape NoUniqueness)
lookupType [VName]
si'
[VName]
indexName <- forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
(a -> m b) -> t a -> m (t b)
mapM (\TypeBase Shape NoUniqueness
_ -> forall (m :: * -> *). MonadFreshNames m => [Char] -> m VName
newVName [Char]
"write_index") [TypeBase Shape NoUniqueness]
indexType
[VName]
valueNames <- forall (m :: * -> *) a. Applicative m => Int -> m a -> m [a]
replicateM (forall (t :: * -> *) a. Foldable t => t a -> Int
length [TypeBase Shape NoUniqueness]
sv_ts) forall a b. (a -> b) -> a -> b
$ forall (m :: * -> *). MonadFreshNames m => [Char] -> m VName
newVName [Char]
"write_value"
[TypeBase Shape NoUniqueness]
sa_ts <- forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
(a -> m b) -> t a -> m (t b)
mapM forall rep (m :: * -> *).
HasScope rep m =>
VName -> m (TypeBase Shape NoUniqueness)
lookupType [VName]
sas
let bodyTypes :: [TypeBase Shape NoUniqueness]
bodyTypes = forall (t :: * -> *) a. Foldable t => t [a] -> [a]
concat (forall a. Int -> a -> [a]
replicate (forall (t :: * -> *) a. Foldable t => t a -> Int
length [TypeBase Shape NoUniqueness]
sv_ts) [TypeBase Shape NoUniqueness]
indexType) forall a. [a] -> [a] -> [a]
++ forall a b. (a -> b) -> [a] -> [b]
map (forall u. Int -> TypeBase Shape u -> TypeBase Shape u
I.stripArray Int
dim) [TypeBase Shape NoUniqueness]
sa_ts
paramTypes :: [TypeBase Shape NoUniqueness]
paramTypes = [TypeBase Shape NoUniqueness]
indexType forall a. Semigroup a => a -> a -> a
<> forall a b. (a -> b) -> [a] -> [b]
map forall u. TypeBase Shape u -> TypeBase Shape u
rowType [TypeBase Shape NoUniqueness]
sv_ts
bodyNames :: [VName]
bodyNames = [VName]
indexName forall a. Semigroup a => a -> a -> a
<> [VName]
valueNames
bodyParams :: [Param (TypeBase Shape NoUniqueness)]
bodyParams = forall a b c. (a -> b -> c) -> [a] -> [b] -> [c]
zipWith (forall dec. Attrs -> VName -> dec -> Param dec
I.Param forall a. Monoid a => a
mempty) [VName]
bodyNames [TypeBase Shape NoUniqueness]
paramTypes
Body SOACS
body <- forall rep (m :: * -> *) a.
LocalScope rep m =>
Scope rep -> m a -> m a
localScope (forall rep dec. (LParamInfo rep ~ dec) => [Param dec] -> Scope rep
scopeOfLParams [Param (TypeBase Shape NoUniqueness)]
bodyParams) forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall (m :: * -> *).
MonadBuilder m =>
m Result -> m (Body (Rep m))
buildBody_ forall a b. (a -> b) -> a -> b
$ do
let outs :: [VName]
outs = forall (t :: * -> *) a. Foldable t => t [a] -> [a]
concat (forall a. Int -> a -> [a]
replicate (forall (t :: * -> *) a. Foldable t => t a -> Int
length [VName]
valueNames) [VName]
indexName) forall a. [a] -> [a] -> [a]
++ [VName]
valueNames
[SubExp]
results <- forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
t a -> (a -> m b) -> m (t b)
forM [VName]
outs forall a b. (a -> b) -> a -> b
$ \VName
name ->
forall (m :: * -> *).
MonadBuilder m =>
[Char] -> Exp (Rep m) -> m SubExp
letSubExp [Char]
"write_res" forall a b. (a -> b) -> a -> b
$ forall rep. BasicOp -> Exp rep
I.BasicOp forall a b. (a -> b) -> a -> b
$ SubExp -> BasicOp
I.SubExp forall a b. (a -> b) -> a -> b
$ VName -> SubExp
I.Var VName
name
ErrorMsg SubExp
-> SrcLoc
-> [TypeBase Shape NoUniqueness]
-> Result
-> InternaliseM Result
ensureResultShape
ErrorMsg SubExp
"scatter value has wrong size"
SrcLoc
loc
[TypeBase Shape NoUniqueness]
bodyTypes
([SubExp] -> Result
subExpsRes [SubExp]
results)
let lam :: Lambda SOACS
lam =
I.Lambda
{ lambdaParams :: [LParam SOACS]
I.lambdaParams = [Param (TypeBase Shape NoUniqueness)]
bodyParams,
lambdaReturnType :: [TypeBase Shape NoUniqueness]
I.lambdaReturnType = [TypeBase Shape NoUniqueness]
bodyTypes,
lambdaBody :: Body SOACS
I.lambdaBody = Body SOACS
body
}
sivs :: [VName]
sivs = [VName]
si' forall a. Semigroup a => a -> a -> a
<> [VName]
svs'
let sa_ws :: [Shape]
sa_ws = forall a b. (a -> b) -> [a] -> [b]
map (forall d. [d] -> ShapeBase d
I.Shape forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall a. Int -> [a] -> [a]
take Int
dim forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall u. TypeBase Shape u -> [SubExp]
arrayDims) [TypeBase Shape NoUniqueness]
sa_ts
forall (m :: * -> *).
MonadBuilder m =>
[Char] -> Exp (Rep m) -> m [SubExp]
letTupExp' [Char]
desc forall a b. (a -> b) -> a -> b
$ forall rep. Op rep -> Exp rep
I.Op forall a b. (a -> b) -> a -> b
$ forall rep.
SubExp
-> [VName] -> Lambda rep -> [(Shape, Int, VName)] -> SOAC rep
I.Scatter SubExp
si_w [VName]
sivs Lambda SOACS
lam forall a b. (a -> b) -> a -> b
$ forall a b c. [a] -> [b] -> [c] -> [(a, b, c)]
zip3 [Shape]
sa_ws (forall a. a -> [a]
repeat Int
1) [VName]
sas
flatIndexHelper :: String -> SrcLoc -> E.Exp -> E.Exp -> [(E.Exp, E.Exp)] -> InternaliseM [SubExp]
flatIndexHelper :: [Char]
-> SrcLoc -> Exp -> Exp -> [(Exp, Exp)] -> InternaliseM [SubExp]
flatIndexHelper [Char]
desc SrcLoc
loc Exp
arr Exp
offset [(Exp, Exp)]
slices = do
[VName]
arrs <- [Char] -> Exp -> InternaliseM [VName]
internaliseExpToVars [Char]
"arr" Exp
arr
SubExp
offset' <- [Char] -> Exp -> InternaliseM SubExp
internaliseExp1 [Char]
"offset" Exp
offset
SubExp
old_dim <- forall u. Int -> [TypeBase Shape u] -> SubExp
I.arraysSize Int
0 forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
(a -> m b) -> t a -> m (t b)
mapM forall rep (m :: * -> *).
HasScope rep m =>
VName -> m (TypeBase Shape NoUniqueness)
lookupType [VName]
arrs
SubExp
offset_inbounds_down <- forall (m :: * -> *).
MonadBuilder m =>
[Char] -> Exp (Rep m) -> m SubExp
letSubExp [Char]
"offset_inbounds_down" forall a b. (a -> b) -> a -> b
$ forall rep. BasicOp -> Exp rep
I.BasicOp forall a b. (a -> b) -> a -> b
$ CmpOp -> SubExp -> SubExp -> BasicOp
I.CmpOp (IntType -> CmpOp
I.CmpUle IntType
Int64) (IntType -> Integer -> SubExp
intConst IntType
Int64 Integer
0) SubExp
offset'
SubExp
offset_inbounds_up <- forall (m :: * -> *).
MonadBuilder m =>
[Char] -> Exp (Rep m) -> m SubExp
letSubExp [Char]
"offset_inbounds_up" forall a b. (a -> b) -> a -> b
$ forall rep. BasicOp -> Exp rep
I.BasicOp forall a b. (a -> b) -> a -> b
$ CmpOp -> SubExp -> SubExp -> BasicOp
I.CmpOp (IntType -> CmpOp
I.CmpUlt IntType
Int64) SubExp
offset' SubExp
old_dim
[(SubExp, SubExp)]
slices' <-
forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
(a -> m b) -> t a -> m (t b)
mapM
( \(Exp
n, Exp
s) -> do
SubExp
n' <- [Char] -> Exp -> InternaliseM SubExp
internaliseExp1 [Char]
"n" Exp
n
SubExp
s' <- [Char] -> Exp -> InternaliseM SubExp
internaliseExp1 [Char]
"s" Exp
s
forall (f :: * -> *) a. Applicative f => a -> f a
pure (SubExp
n', SubExp
s')
)
[(Exp, Exp)]
slices
(SubExp
min_bound, SubExp
max_bound) <-
forall (t :: * -> *) (m :: * -> *) b a.
(Foldable t, Monad m) =>
(b -> a -> m b) -> b -> t a -> m b
foldM
( \(SubExp
lower, SubExp
upper) (SubExp
n, SubExp
s) -> do
SubExp
n_m1 <- forall (m :: * -> *).
MonadBuilder m =>
[Char] -> Exp (Rep m) -> m SubExp
letSubExp [Char]
"span" forall a b. (a -> b) -> a -> b
$ forall rep. BasicOp -> Exp rep
I.BasicOp forall a b. (a -> b) -> a -> b
$ BinOp -> SubExp -> SubExp -> BasicOp
I.BinOp (IntType -> Overflow -> BinOp
I.Sub IntType
Int64 Overflow
I.OverflowUndef) SubExp
n (IntType -> Integer -> SubExp
intConst IntType
Int64 Integer
1)
SubExp
spn <- forall (m :: * -> *).
MonadBuilder m =>
[Char] -> Exp (Rep m) -> m SubExp
letSubExp [Char]
"span" forall a b. (a -> b) -> a -> b
$ forall rep. BasicOp -> Exp rep
I.BasicOp forall a b. (a -> b) -> a -> b
$ BinOp -> SubExp -> SubExp -> BasicOp
I.BinOp (IntType -> Overflow -> BinOp
I.Mul IntType
Int64 Overflow
I.OverflowUndef) SubExp
n_m1 SubExp
s
SubExp
span_and_lower <- forall (m :: * -> *).
MonadBuilder m =>
[Char] -> Exp (Rep m) -> m SubExp
letSubExp [Char]
"span_and_lower" forall a b. (a -> b) -> a -> b
$ forall rep. BasicOp -> Exp rep
I.BasicOp forall a b. (a -> b) -> a -> b
$ BinOp -> SubExp -> SubExp -> BasicOp
I.BinOp (IntType -> Overflow -> BinOp
I.Add IntType
Int64 Overflow
I.OverflowUndef) SubExp
spn SubExp
lower
SubExp
span_and_upper <- forall (m :: * -> *).
MonadBuilder m =>
[Char] -> Exp (Rep m) -> m SubExp
letSubExp [Char]
"span_and_upper" forall a b. (a -> b) -> a -> b
$ forall rep. BasicOp -> Exp rep
I.BasicOp forall a b. (a -> b) -> a -> b
$ BinOp -> SubExp -> SubExp -> BasicOp
I.BinOp (IntType -> Overflow -> BinOp
I.Add IntType
Int64 Overflow
I.OverflowUndef) SubExp
spn SubExp
upper
SubExp
lower' <- forall (m :: * -> *).
MonadBuilder m =>
[Char] -> Exp (Rep m) -> m SubExp
letSubExp [Char]
"minimum" forall a b. (a -> b) -> a -> b
$ forall rep. BasicOp -> Exp rep
I.BasicOp forall a b. (a -> b) -> a -> b
$ BinOp -> SubExp -> SubExp -> BasicOp
I.BinOp (IntType -> BinOp
I.UMin IntType
Int64) SubExp
span_and_lower SubExp
lower
SubExp
upper' <- forall (m :: * -> *).
MonadBuilder m =>
[Char] -> Exp (Rep m) -> m SubExp
letSubExp [Char]
"maximum" forall a b. (a -> b) -> a -> b
$ forall rep. BasicOp -> Exp rep
I.BasicOp forall a b. (a -> b) -> a -> b
$ BinOp -> SubExp -> SubExp -> BasicOp
I.BinOp (IntType -> BinOp
I.UMax IntType
Int64) SubExp
span_and_upper SubExp
upper
forall (f :: * -> *) a. Applicative f => a -> f a
pure (SubExp
lower', SubExp
upper')
)
(SubExp
offset', SubExp
offset')
[(SubExp, SubExp)]
slices'
SubExp
min_in_bounds <- forall (m :: * -> *).
MonadBuilder m =>
[Char] -> Exp (Rep m) -> m SubExp
letSubExp [Char]
"min_in_bounds" forall a b. (a -> b) -> a -> b
$ forall rep. BasicOp -> Exp rep
I.BasicOp forall a b. (a -> b) -> a -> b
$ CmpOp -> SubExp -> SubExp -> BasicOp
I.CmpOp (IntType -> CmpOp
I.CmpUle IntType
Int64) (IntType -> Integer -> SubExp
intConst IntType
Int64 Integer
0) SubExp
min_bound
SubExp
max_in_bounds <- forall (m :: * -> *).
MonadBuilder m =>
[Char] -> Exp (Rep m) -> m SubExp
letSubExp [Char]
"max_in_bounds" forall a b. (a -> b) -> a -> b
$ forall rep. BasicOp -> Exp rep
I.BasicOp forall a b. (a -> b) -> a -> b
$ CmpOp -> SubExp -> SubExp -> BasicOp
I.CmpOp (IntType -> CmpOp
I.CmpUlt IntType
Int64) SubExp
max_bound SubExp
old_dim
SubExp
all_bounds <-
forall (t :: * -> *) (m :: * -> *) b a.
(Foldable t, Monad m) =>
(b -> a -> m b) -> b -> t a -> m b
foldM
(\SubExp
x SubExp
y -> forall (m :: * -> *).
MonadBuilder m =>
[Char] -> Exp (Rep m) -> m SubExp
letSubExp [Char]
"inBounds" forall a b. (a -> b) -> a -> b
$ forall rep. BasicOp -> Exp rep
I.BasicOp forall a b. (a -> b) -> a -> b
$ BinOp -> SubExp -> SubExp -> BasicOp
I.BinOp BinOp
I.LogAnd SubExp
x SubExp
y)
SubExp
offset_inbounds_down
[SubExp
offset_inbounds_up, SubExp
min_in_bounds, SubExp
max_in_bounds]
Certs
c <- [Char] -> SubExp -> ErrorMsg SubExp -> SrcLoc -> InternaliseM Certs
assert [Char]
"bounds_cert" SubExp
all_bounds (forall a. [ErrorMsgPart a] -> ErrorMsg a
ErrorMsg [forall a. Text -> ErrorMsgPart a
ErrorString forall a b. (a -> b) -> a -> b
$ Text
"Flat slice out of bounds: " forall a. Semigroup a => a -> a -> a
<> forall a. Pretty a => a -> Text
prettyText SubExp
old_dim forall a. Semigroup a => a -> a -> a
<> Text
" and " forall a. Semigroup a => a -> a -> a
<> forall a. Pretty a => a -> Text
prettyText [(SubExp, SubExp)]
slices']) SrcLoc
loc
let slice :: FlatSlice SubExp
slice = forall d. d -> [FlatDimIndex d] -> FlatSlice d
I.FlatSlice SubExp
offset' forall a b. (a -> b) -> a -> b
$ forall a b. (a -> b) -> [a] -> [b]
map (forall a b c. (a -> b -> c) -> (a, b) -> c
uncurry forall d. d -> d -> FlatDimIndex d
FlatDimIndex) [(SubExp, SubExp)]
slices'
forall (m :: * -> *) a. MonadBuilder m => Certs -> m a -> m a
certifying Certs
c forall a b. (a -> b) -> a -> b
$
forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
t a -> (a -> m b) -> m (t b)
forM [VName]
arrs forall a b. (a -> b) -> a -> b
$ \VName
arr' ->
forall (m :: * -> *).
MonadBuilder m =>
[Char] -> Exp (Rep m) -> m SubExp
letSubExp [Char]
desc forall a b. (a -> b) -> a -> b
$ forall rep. BasicOp -> Exp rep
I.BasicOp forall a b. (a -> b) -> a -> b
$ VName -> FlatSlice SubExp -> BasicOp
I.FlatIndex VName
arr' FlatSlice SubExp
slice
flatUpdateHelper :: String -> SrcLoc -> E.Exp -> E.Exp -> [E.Exp] -> E.Exp -> InternaliseM [SubExp]
flatUpdateHelper :: [Char]
-> SrcLoc -> Exp -> Exp -> [Exp] -> Exp -> InternaliseM [SubExp]
flatUpdateHelper [Char]
desc SrcLoc
loc Exp
arr1 Exp
offset [Exp]
slices Exp
arr2 = do
[VName]
arrs1 <- [Char] -> Exp -> InternaliseM [VName]
internaliseExpToVars [Char]
"arr" Exp
arr1
SubExp
offset' <- [Char] -> Exp -> InternaliseM SubExp
internaliseExp1 [Char]
"offset" Exp
offset
SubExp
old_dim <- forall u. Int -> [TypeBase Shape u] -> SubExp
I.arraysSize Int
0 forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
(a -> m b) -> t a -> m (t b)
mapM forall rep (m :: * -> *).
HasScope rep m =>
VName -> m (TypeBase Shape NoUniqueness)
lookupType [VName]
arrs1
SubExp
offset_inbounds_down <- forall (m :: * -> *).
MonadBuilder m =>
[Char] -> Exp (Rep m) -> m SubExp
letSubExp [Char]
"offset_inbounds_down" forall a b. (a -> b) -> a -> b
$ forall rep. BasicOp -> Exp rep
I.BasicOp forall a b. (a -> b) -> a -> b
$ CmpOp -> SubExp -> SubExp -> BasicOp
I.CmpOp (IntType -> CmpOp
I.CmpUle IntType
Int64) (IntType -> Integer -> SubExp
intConst IntType
Int64 Integer
0) SubExp
offset'
SubExp
offset_inbounds_up <- forall (m :: * -> *).
MonadBuilder m =>
[Char] -> Exp (Rep m) -> m SubExp
letSubExp [Char]
"offset_inbounds_up" forall a b. (a -> b) -> a -> b
$ forall rep. BasicOp -> Exp rep
I.BasicOp forall a b. (a -> b) -> a -> b
$ CmpOp -> SubExp -> SubExp -> BasicOp
I.CmpOp (IntType -> CmpOp
I.CmpUlt IntType
Int64) SubExp
offset' SubExp
old_dim
[VName]
arrs2 <- [Char] -> Exp -> InternaliseM [VName]
internaliseExpToVars [Char]
"arr" Exp
arr2
[TypeBase Shape NoUniqueness]
ts <- forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
(a -> m b) -> t a -> m (t b)
mapM forall rep (m :: * -> *).
HasScope rep m =>
VName -> m (TypeBase Shape NoUniqueness)
lookupType [VName]
arrs2
[(SubExp, SubExp)]
slices' <-
forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
(a -> m b) -> t a -> m (t b)
mapM
( \(Exp
s, Int
i) -> do
SubExp
s' <- [Char] -> Exp -> InternaliseM SubExp
internaliseExp1 [Char]
"s" Exp
s
let n :: SubExp
n = forall u. Int -> [TypeBase Shape u] -> SubExp
arraysSize Int
i [TypeBase Shape NoUniqueness]
ts
forall (f :: * -> *) a. Applicative f => a -> f a
pure (SubExp
n, SubExp
s')
)
forall a b. (a -> b) -> a -> b
$ forall a b. [a] -> [b] -> [(a, b)]
zip [Exp]
slices [Int
0 ..]
(SubExp
min_bound, SubExp
max_bound) <-
forall (t :: * -> *) (m :: * -> *) b a.
(Foldable t, Monad m) =>
(b -> a -> m b) -> b -> t a -> m b
foldM
( \(SubExp
lower, SubExp
upper) (SubExp
n, SubExp
s) -> do
SubExp
n_m1 <- forall (m :: * -> *).
MonadBuilder m =>
[Char] -> Exp (Rep m) -> m SubExp
letSubExp [Char]
"span" forall a b. (a -> b) -> a -> b
$ forall rep. BasicOp -> Exp rep
I.BasicOp forall a b. (a -> b) -> a -> b
$ BinOp -> SubExp -> SubExp -> BasicOp
I.BinOp (IntType -> Overflow -> BinOp
I.Sub IntType
Int64 Overflow
I.OverflowUndef) SubExp
n (IntType -> Integer -> SubExp
intConst IntType
Int64 Integer
1)
SubExp
spn <- forall (m :: * -> *).
MonadBuilder m =>
[Char] -> Exp (Rep m) -> m SubExp
letSubExp [Char]
"span" forall a b. (a -> b) -> a -> b
$ forall rep. BasicOp -> Exp rep
I.BasicOp forall a b. (a -> b) -> a -> b
$ BinOp -> SubExp -> SubExp -> BasicOp
I.BinOp (IntType -> Overflow -> BinOp
I.Mul IntType
Int64 Overflow
I.OverflowUndef) SubExp
n_m1 SubExp
s
SubExp
span_and_lower <- forall (m :: * -> *).
MonadBuilder m =>
[Char] -> Exp (Rep m) -> m SubExp
letSubExp [Char]
"span_and_lower" forall a b. (a -> b) -> a -> b
$ forall rep. BasicOp -> Exp rep
I.BasicOp forall a b. (a -> b) -> a -> b
$ BinOp -> SubExp -> SubExp -> BasicOp
I.BinOp (IntType -> Overflow -> BinOp
I.Add IntType
Int64 Overflow
I.OverflowUndef) SubExp
spn SubExp
lower
SubExp
span_and_upper <- forall (m :: * -> *).
MonadBuilder m =>
[Char] -> Exp (Rep m) -> m SubExp
letSubExp [Char]
"span_and_upper" forall a b. (a -> b) -> a -> b
$ forall rep. BasicOp -> Exp rep
I.BasicOp forall a b. (a -> b) -> a -> b
$ BinOp -> SubExp -> SubExp -> BasicOp
I.BinOp (IntType -> Overflow -> BinOp
I.Add IntType
Int64 Overflow
I.OverflowUndef) SubExp
spn SubExp
upper
SubExp
lower' <- forall (m :: * -> *).
MonadBuilder m =>
[Char] -> Exp (Rep m) -> m SubExp
letSubExp [Char]
"minimum" forall a b. (a -> b) -> a -> b
$ forall rep. BasicOp -> Exp rep
I.BasicOp forall a b. (a -> b) -> a -> b
$ BinOp -> SubExp -> SubExp -> BasicOp
I.BinOp (IntType -> BinOp
I.UMin IntType
Int64) SubExp
span_and_lower SubExp
lower
SubExp
upper' <- forall (m :: * -> *).
MonadBuilder m =>
[Char] -> Exp (Rep m) -> m SubExp
letSubExp [Char]
"maximum" forall a b. (a -> b) -> a -> b
$ forall rep. BasicOp -> Exp rep
I.BasicOp forall a b. (a -> b) -> a -> b
$ BinOp -> SubExp -> SubExp -> BasicOp
I.BinOp (IntType -> BinOp
I.UMax IntType
Int64) SubExp
span_and_upper SubExp
upper
forall (f :: * -> *) a. Applicative f => a -> f a
pure (SubExp
lower', SubExp
upper')
)
(SubExp
offset', SubExp
offset')
[(SubExp, SubExp)]
slices'
SubExp
min_in_bounds <- forall (m :: * -> *).
MonadBuilder m =>
[Char] -> Exp (Rep m) -> m SubExp
letSubExp [Char]
"min_in_bounds" forall a b. (a -> b) -> a -> b
$ forall rep. BasicOp -> Exp rep
I.BasicOp forall a b. (a -> b) -> a -> b
$ CmpOp -> SubExp -> SubExp -> BasicOp
I.CmpOp (IntType -> CmpOp
I.CmpUle IntType
Int64) (IntType -> Integer -> SubExp
intConst IntType
Int64 Integer
0) SubExp
min_bound
SubExp
max_in_bounds <- forall (m :: * -> *).
MonadBuilder m =>
[Char] -> Exp (Rep m) -> m SubExp
letSubExp [Char]
"max_in_bounds" forall a b. (a -> b) -> a -> b
$ forall rep. BasicOp -> Exp rep
I.BasicOp forall a b. (a -> b) -> a -> b
$ CmpOp -> SubExp -> SubExp -> BasicOp
I.CmpOp (IntType -> CmpOp
I.CmpUlt IntType
Int64) SubExp
max_bound SubExp
old_dim
SubExp
all_bounds <-
forall (t :: * -> *) (m :: * -> *) b a.
(Foldable t, Monad m) =>
(b -> a -> m b) -> b -> t a -> m b
foldM
(\SubExp
x SubExp
y -> forall (m :: * -> *).
MonadBuilder m =>
[Char] -> Exp (Rep m) -> m SubExp
letSubExp [Char]
"inBounds" forall a b. (a -> b) -> a -> b
$ forall rep. BasicOp -> Exp rep
I.BasicOp forall a b. (a -> b) -> a -> b
$ BinOp -> SubExp -> SubExp -> BasicOp
I.BinOp BinOp
I.LogAnd SubExp
x SubExp
y)
SubExp
offset_inbounds_down
[SubExp
offset_inbounds_up, SubExp
min_in_bounds, SubExp
max_in_bounds]
Certs
c <- [Char] -> SubExp -> ErrorMsg SubExp -> SrcLoc -> InternaliseM Certs
assert [Char]
"bounds_cert" SubExp
all_bounds (forall a. [ErrorMsgPart a] -> ErrorMsg a
ErrorMsg [forall a. Text -> ErrorMsgPart a
ErrorString forall a b. (a -> b) -> a -> b
$ Text
"Flat slice out of bounds: " forall a. Semigroup a => a -> a -> a
<> forall a. Pretty a => a -> Text
prettyText SubExp
old_dim forall a. Semigroup a => a -> a -> a
<> Text
" and " forall a. Semigroup a => a -> a -> a
<> forall a. Pretty a => a -> Text
prettyText [(SubExp, SubExp)]
slices']) SrcLoc
loc
let slice :: FlatSlice SubExp
slice = forall d. d -> [FlatDimIndex d] -> FlatSlice d
I.FlatSlice SubExp
offset' forall a b. (a -> b) -> a -> b
$ forall a b. (a -> b) -> [a] -> [b]
map (forall a b c. (a -> b -> c) -> (a, b) -> c
uncurry forall d. d -> d -> FlatDimIndex d
FlatDimIndex) [(SubExp, SubExp)]
slices'
forall (m :: * -> *) a. MonadBuilder m => Certs -> m a -> m a
certifying Certs
c forall a b. (a -> b) -> a -> b
$
forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
t a -> (a -> m b) -> m (t b)
forM (forall a b. [a] -> [b] -> [(a, b)]
zip [VName]
arrs1 [VName]
arrs2) forall a b. (a -> b) -> a -> b
$ \(VName
arr1', VName
arr2') ->
forall (m :: * -> *).
MonadBuilder m =>
[Char] -> Exp (Rep m) -> m SubExp
letSubExp [Char]
desc forall a b. (a -> b) -> a -> b
$ forall rep. BasicOp -> Exp rep
I.BasicOp forall a b. (a -> b) -> a -> b
$ VName -> FlatSlice SubExp -> VName -> BasicOp
I.FlatUpdate VName
arr1' FlatSlice SubExp
slice VName
arr2'
funcall ::
String ->
QualName VName ->
[SubExp] ->
SrcLoc ->
InternaliseM ([SubExp], [I.ExtType])
funcall :: [Char]
-> QualName VName
-> [SubExp]
-> SrcLoc
-> InternaliseM ([SubExp], [ExtType])
funcall [Char]
desc (QualName [VName]
_ VName
fname) [SubExp]
args SrcLoc
loc = do
([VName]
shapes, [DeclType]
value_paramts, [Param DeclType]
fun_params, [(SubExp, TypeBase Shape NoUniqueness)]
-> Maybe [TypeBase ExtShape Uniqueness]
rettype_fun) <-
VName -> InternaliseM FunInfo
lookupFunction VName
fname
[TypeBase Shape NoUniqueness]
argts <- forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
(a -> m b) -> t a -> m (t b)
mapM forall t (m :: * -> *).
HasScope t m =>
SubExp -> m (TypeBase Shape NoUniqueness)
subExpType [SubExp]
args
[SubExp]
shapeargs <- [VName]
-> [FParam SOACS]
-> [TypeBase Shape NoUniqueness]
-> InternaliseM [SubExp]
argShapes [VName]
shapes [Param DeclType]
fun_params [TypeBase Shape NoUniqueness]
argts
let diets :: [Diet]
diets =
forall a. Int -> a -> [a]
replicate (forall (t :: * -> *) a. Foldable t => t a -> Int
length [SubExp]
shapeargs) Diet
I.ObservePrim
forall a. [a] -> [a] -> [a]
++ forall a b. (a -> b) -> [a] -> [b]
map forall shape. TypeBase shape Uniqueness -> Diet
I.diet [DeclType]
value_paramts
[SubExp]
args' <-
forall u.
Typed (TypeBase Shape u) =>
ErrorMsg SubExp
-> SrcLoc
-> [VName]
-> [TypeBase Shape u]
-> [SubExp]
-> InternaliseM [SubExp]
ensureArgShapes
ErrorMsg SubExp
"function arguments of wrong shape"
SrcLoc
loc
(forall a b. (a -> b) -> [a] -> [b]
map forall dec. Param dec -> VName
I.paramName [Param DeclType]
fun_params)
(forall a b. (a -> b) -> [a] -> [b]
map forall dec. Typed dec => Param dec -> TypeBase Shape NoUniqueness
I.paramType [Param DeclType]
fun_params)
([SubExp]
shapeargs forall a. [a] -> [a] -> [a]
++ [SubExp]
args)
[TypeBase Shape NoUniqueness]
argts' <- forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
(a -> m b) -> t a -> m (t b)
mapM forall t (m :: * -> *).
HasScope t m =>
SubExp -> m (TypeBase Shape NoUniqueness)
subExpType [SubExp]
args'
case [(SubExp, TypeBase Shape NoUniqueness)]
-> Maybe [TypeBase ExtShape Uniqueness]
rettype_fun forall a b. (a -> b) -> a -> b
$ forall a b. [a] -> [b] -> [(a, b)]
zip [SubExp]
args' [TypeBase Shape NoUniqueness]
argts' of
Maybe [TypeBase ExtShape Uniqueness]
Nothing ->
forall a. HasCallStack => [Char] -> a
error forall a b. (a -> b) -> a -> b
$
forall (t :: * -> *) a. Foldable t => t [a] -> [a]
concat
[ [Char]
"Cannot apply ",
forall a. Pretty a => a -> [Char]
prettyString VName
fname,
[Char]
" to ",
forall a. Show a => a -> [Char]
show (forall (t :: * -> *) a. Foldable t => t a -> Int
length [SubExp]
args'),
[Char]
" arguments\n ",
forall a. Pretty a => a -> [Char]
prettyString [SubExp]
args',
[Char]
"\nof types\n ",
forall a. Pretty a => a -> [Char]
prettyString [TypeBase Shape NoUniqueness]
argts',
[Char]
"\nFunction has ",
forall a. Show a => a -> [Char]
show (forall (t :: * -> *) a. Foldable t => t a -> Int
length [Param DeclType]
fun_params),
[Char]
" parameters\n ",
forall a. Pretty a => a -> [Char]
prettyString [Param DeclType]
fun_params
]
Just [TypeBase ExtShape Uniqueness]
ts -> do
Safety
safety <- InternaliseM Safety
askSafety
Attrs
attrs <- forall r (m :: * -> *) a. MonadReader r m => (r -> a) -> m a
asks InternaliseEnv -> Attrs
envAttrs
[SubExp]
ses <-
forall (m :: * -> *) a. MonadBuilder m => Attrs -> m a -> m a
attributing Attrs
attrs forall b c a. (b -> c) -> (a -> b) -> a -> c
. [Char] -> Exp SOACS -> InternaliseM [SubExp]
letValExp' [Char]
desc forall a b. (a -> b) -> a -> b
$
forall rep.
Name
-> [(SubExp, Diet)]
-> [RetType rep]
-> (Safety, SrcLoc, [SrcLoc])
-> Exp rep
I.Apply (VName -> Name
internaliseFunName VName
fname) (forall a b. [a] -> [b] -> [(a, b)]
zip [SubExp]
args' [Diet]
diets) [TypeBase ExtShape Uniqueness]
ts (Safety
safety, SrcLoc
loc, forall a. Monoid a => a
mempty)
forall (f :: * -> *) a. Applicative f => a -> f a
pure ([SubExp]
ses, forall a b. (a -> b) -> [a] -> [b]
map forall shape.
TypeBase shape Uniqueness -> TypeBase shape NoUniqueness
I.fromDecl [TypeBase ExtShape Uniqueness]
ts)
bindExtSizes :: AppRes -> [SubExp] -> InternaliseM ()
bindExtSizes :: AppRes -> [SubExp] -> InternaliseM ()
bindExtSizes (AppRes PatType
ret [VName]
retext) [SubExp]
ses = do
let ts :: [TypeBase ExtShape Uniqueness]
ts = TypeBase Size () -> [TypeBase ExtShape Uniqueness]
internaliseType forall a b. (a -> b) -> a -> b
$ forall dim as. TypeBase dim as -> TypeBase dim ()
E.toStruct PatType
ret
[TypeBase Shape NoUniqueness]
ses_ts <- forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
(a -> m b) -> t a -> m (t b)
mapM forall t (m :: * -> *).
HasScope t m =>
SubExp -> m (TypeBase Shape NoUniqueness)
subExpType [SubExp]
ses
let combine :: TypeBase ExtShape Uniqueness
-> TypeBase Shape NoUniqueness -> Map VName SubExp
combine TypeBase ExtShape Uniqueness
t1 TypeBase Shape NoUniqueness
t2 =
forall a. Monoid a => [a] -> a
mconcat forall a b. (a -> b) -> a -> b
$ forall a b c. (a -> b -> c) -> [a] -> [b] -> [c]
zipWith Ext SubExp -> SubExp -> Map VName SubExp
combine' (forall u. TypeBase ExtShape u -> [Ext SubExp]
arrayExtDims TypeBase ExtShape Uniqueness
t1) (forall u. TypeBase Shape u -> [SubExp]
arrayDims TypeBase Shape NoUniqueness
t2)
combine' :: Ext SubExp -> SubExp -> Map VName SubExp
combine' (I.Free (I.Var VName
v)) SubExp
se
| VName
v forall (t :: * -> *) a. (Foldable t, Eq a) => a -> t a -> Bool
`elem` [VName]
retext = forall k a. k -> a -> Map k a
M.singleton VName
v SubExp
se
combine' Ext SubExp
_ SubExp
_ = forall a. Monoid a => a
mempty
forall (t :: * -> *) (m :: * -> *) a b.
(Foldable t, Monad m) =>
t a -> (a -> m b) -> m ()
forM_ (forall k a. Map k a -> [(k, a)]
M.toList forall a b. (a -> b) -> a -> b
$ forall a. Monoid a => [a] -> a
mconcat forall a b. (a -> b) -> a -> b
$ forall a b c. (a -> b -> c) -> [a] -> [b] -> [c]
zipWith TypeBase ExtShape Uniqueness
-> TypeBase Shape NoUniqueness -> Map VName SubExp
combine [TypeBase ExtShape Uniqueness]
ts [TypeBase Shape NoUniqueness]
ses_ts) forall a b. (a -> b) -> a -> b
$ \(VName
v, SubExp
se) ->
forall (m :: * -> *).
MonadBuilder m =>
[VName] -> Exp (Rep m) -> m ()
letBindNames [VName
v] forall a b. (a -> b) -> a -> b
$ forall rep. BasicOp -> Exp rep
BasicOp forall a b. (a -> b) -> a -> b
$ SubExp -> BasicOp
SubExp SubExp
se
askSafety :: InternaliseM Safety
askSafety :: InternaliseM Safety
askSafety = do
Bool
check <- forall r (m :: * -> *) a. MonadReader r m => (r -> a) -> m a
asks InternaliseEnv -> Bool
envDoBoundsChecks
forall (f :: * -> *) a. Applicative f => a -> f a
pure forall a b. (a -> b) -> a -> b
$ if Bool
check then Safety
I.Safe else Safety
I.Unsafe
partitionWithSOACS :: Int -> I.Lambda SOACS -> [I.VName] -> InternaliseM ([I.SubExp], [I.SubExp])
partitionWithSOACS :: Int -> Lambda SOACS -> [VName] -> InternaliseM ([SubExp], [SubExp])
partitionWithSOACS Int
k Lambda SOACS
lam [VName]
arrs = do
[TypeBase Shape NoUniqueness]
arr_ts <- forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
(a -> m b) -> t a -> m (t b)
mapM forall rep (m :: * -> *).
HasScope rep m =>
VName -> m (TypeBase Shape NoUniqueness)
lookupType [VName]
arrs
let w :: SubExp
w = forall u. Int -> [TypeBase Shape u] -> SubExp
arraysSize Int
0 [TypeBase Shape NoUniqueness]
arr_ts
[VName]
classes_and_increments <- forall (m :: * -> *).
MonadBuilder m =>
[Char] -> Exp (Rep m) -> m [VName]
letTupExp [Char]
"increments" forall a b. (a -> b) -> a -> b
$ forall rep. Op rep -> Exp rep
I.Op forall a b. (a -> b) -> a -> b
$ forall rep. SubExp -> [VName] -> ScremaForm rep -> SOAC rep
I.Screma SubExp
w [VName]
arrs (forall rep. Lambda rep -> ScremaForm rep
mapSOAC Lambda SOACS
lam)
(VName
classes, [VName]
increments) <- case [VName]
classes_and_increments of
VName
classes : [VName]
increments -> forall (f :: * -> *) a. Applicative f => a -> f a
pure (VName
classes, forall a. Int -> [a] -> [a]
take Int
k [VName]
increments)
[VName]
_ -> forall a. HasCallStack => [Char] -> a
error [Char]
"partitionWithSOACS"
[Param (TypeBase Shape NoUniqueness)]
add_lam_x_params <-
forall (m :: * -> *) a. Applicative m => Int -> m a -> m [a]
replicateM Int
k forall a b. (a -> b) -> a -> b
$ forall (m :: * -> *) dec.
MonadFreshNames m =>
[Char] -> dec -> m (Param dec)
newParam [Char]
"x" (forall shape u. PrimType -> TypeBase shape u
I.Prim PrimType
int64)
[Param (TypeBase Shape NoUniqueness)]
add_lam_y_params <-
forall (m :: * -> *) a. Applicative m => Int -> m a -> m [a]
replicateM Int
k forall a b. (a -> b) -> a -> b
$ forall (m :: * -> *) dec.
MonadFreshNames m =>
[Char] -> dec -> m (Param dec)
newParam [Char]
"y" (forall shape u. PrimType -> TypeBase shape u
I.Prim PrimType
int64)
Body SOACS
add_lam_body <- forall rep (m :: * -> *) somerep.
(Buildable rep, MonadFreshNames m, HasScope somerep m,
SameScope somerep rep) =>
Builder rep (Body rep) -> m (Body rep)
runBodyBuilder forall a b. (a -> b) -> a -> b
$
forall rep (m :: * -> *) a.
LocalScope rep m =>
Scope rep -> m a -> m a
localScope (forall rep dec. (LParamInfo rep ~ dec) => [Param dec] -> Scope rep
scopeOfLParams forall a b. (a -> b) -> a -> b
$ [Param (TypeBase Shape NoUniqueness)]
add_lam_x_params forall a. [a] -> [a] -> [a]
++ [Param (TypeBase Shape NoUniqueness)]
add_lam_y_params) forall a b. (a -> b) -> a -> b
$
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap forall rep. Buildable rep => [SubExp] -> Body rep
resultBody forall a b. (a -> b) -> a -> b
$
forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
t a -> (a -> m b) -> m (t b)
forM (forall a b. [a] -> [b] -> [(a, b)]
zip [Param (TypeBase Shape NoUniqueness)]
add_lam_x_params [Param (TypeBase Shape NoUniqueness)]
add_lam_y_params) forall a b. (a -> b) -> a -> b
$ \(Param (TypeBase Shape NoUniqueness)
x, Param (TypeBase Shape NoUniqueness)
y) ->
forall (m :: * -> *).
MonadBuilder m =>
[Char] -> Exp (Rep m) -> m SubExp
letSubExp [Char]
"z" forall a b. (a -> b) -> a -> b
$
forall rep. BasicOp -> Exp rep
I.BasicOp forall a b. (a -> b) -> a -> b
$
BinOp -> SubExp -> SubExp -> BasicOp
I.BinOp
(IntType -> Overflow -> BinOp
I.Add IntType
Int64 Overflow
I.OverflowUndef)
(VName -> SubExp
I.Var forall a b. (a -> b) -> a -> b
$ forall dec. Param dec -> VName
I.paramName Param (TypeBase Shape NoUniqueness)
x)
(VName -> SubExp
I.Var forall a b. (a -> b) -> a -> b
$ forall dec. Param dec -> VName
I.paramName Param (TypeBase Shape NoUniqueness)
y)
let add_lam :: Lambda SOACS
add_lam =
I.Lambda
{ lambdaBody :: Body SOACS
I.lambdaBody = Body SOACS
add_lam_body,
lambdaParams :: [LParam SOACS]
I.lambdaParams = [Param (TypeBase Shape NoUniqueness)]
add_lam_x_params forall a. [a] -> [a] -> [a]
++ [Param (TypeBase Shape NoUniqueness)]
add_lam_y_params,
lambdaReturnType :: [TypeBase Shape NoUniqueness]
I.lambdaReturnType = forall a. Int -> a -> [a]
replicate Int
k forall a b. (a -> b) -> a -> b
$ forall shape u. PrimType -> TypeBase shape u
I.Prim PrimType
int64
}
nes :: [SubExp]
nes = forall a. Int -> a -> [a]
replicate (forall (t :: * -> *) a. Foldable t => t a -> Int
length [VName]
increments) forall a b. (a -> b) -> a -> b
$ IntType -> Integer -> SubExp
intConst IntType
Int64 Integer
0
ScremaForm SOACS
scan <- forall rep (m :: * -> *).
(Buildable rep, MonadFreshNames m) =>
[Scan rep] -> m (ScremaForm rep)
I.scanSOAC [forall rep. Lambda rep -> [SubExp] -> Scan rep
I.Scan Lambda SOACS
add_lam [SubExp]
nes]
[VName]
all_offsets <- forall (m :: * -> *).
MonadBuilder m =>
[Char] -> Exp (Rep m) -> m [VName]
letTupExp [Char]
"offsets" forall a b. (a -> b) -> a -> b
$ forall rep. Op rep -> Exp rep
I.Op forall a b. (a -> b) -> a -> b
$ forall rep. SubExp -> [VName] -> ScremaForm rep -> SOAC rep
I.Screma SubExp
w [VName]
increments ScremaForm SOACS
scan
SubExp
last_index <- forall (m :: * -> *).
MonadBuilder m =>
[Char] -> Exp (Rep m) -> m SubExp
letSubExp [Char]
"last_index" forall a b. (a -> b) -> a -> b
$ forall rep. BasicOp -> Exp rep
I.BasicOp forall a b. (a -> b) -> a -> b
$ BinOp -> SubExp -> SubExp -> BasicOp
I.BinOp (IntType -> Overflow -> BinOp
I.Sub IntType
Int64 Overflow
OverflowUndef) SubExp
w forall a b. (a -> b) -> a -> b
$ forall v. IsValue v => v -> SubExp
constant (Int64
1 :: Int64)
let nonempty_body :: InternaliseM (Body SOACS)
nonempty_body = forall rep (m :: * -> *) somerep.
(Buildable rep, MonadFreshNames m, HasScope somerep m,
SameScope somerep rep) =>
Builder rep (Body rep) -> m (Body rep)
runBodyBuilder forall a b. (a -> b) -> a -> b
$
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap forall rep. Buildable rep => [SubExp] -> Body rep
resultBody forall a b. (a -> b) -> a -> b
$
forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
t a -> (a -> m b) -> m (t b)
forM [VName]
all_offsets forall a b. (a -> b) -> a -> b
$ \VName
offset_array ->
forall (m :: * -> *).
MonadBuilder m =>
[Char] -> Exp (Rep m) -> m SubExp
letSubExp [Char]
"last_offset" forall a b. (a -> b) -> a -> b
$ forall rep. BasicOp -> Exp rep
I.BasicOp forall a b. (a -> b) -> a -> b
$ VName -> Slice SubExp -> BasicOp
I.Index VName
offset_array forall a b. (a -> b) -> a -> b
$ forall d. [DimIndex d] -> Slice d
Slice [forall d. d -> DimIndex d
I.DimFix SubExp
last_index]
empty_body :: InternaliseM (Body (Rep InternaliseM))
empty_body = forall (m :: * -> *).
MonadBuilder m =>
[SubExp] -> m (Body (Rep m))
resultBodyM forall a b. (a -> b) -> a -> b
$ forall a. Int -> a -> [a]
replicate Int
k forall a b. (a -> b) -> a -> b
$ forall v. IsValue v => v -> SubExp
constant (Int64
0 :: Int64)
SubExp
is_empty <- forall (m :: * -> *).
MonadBuilder m =>
[Char] -> Exp (Rep m) -> m SubExp
letSubExp [Char]
"is_empty" forall a b. (a -> b) -> a -> b
$ forall rep. BasicOp -> Exp rep
I.BasicOp forall a b. (a -> b) -> a -> b
$ CmpOp -> SubExp -> SubExp -> BasicOp
I.CmpOp (PrimType -> CmpOp
CmpEq PrimType
int64) SubExp
w forall a b. (a -> b) -> a -> b
$ forall v. IsValue v => v -> SubExp
constant (Int64
0 :: Int64)
[VName]
sizes <-
forall (m :: * -> *).
MonadBuilder m =>
[Char] -> Exp (Rep m) -> m [VName]
letTupExp [Char]
"partition_size" forall (m :: * -> *) a b. Monad m => (a -> m b) -> m a -> m b
=<< forall (m :: * -> *).
(MonadBuilder m, BranchType (Rep m) ~ ExtType) =>
m (Exp (Rep m))
-> m (Body (Rep m)) -> m (Body (Rep m)) -> m (Exp (Rep m))
eIf (forall (m :: * -> *). MonadBuilder m => SubExp -> m (Exp (Rep m))
eSubExp SubExp
is_empty) InternaliseM (Body (Rep InternaliseM))
empty_body InternaliseM (Body SOACS)
nonempty_body
[VName]
blanks <- forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
t a -> (a -> m b) -> m (t b)
forM [TypeBase Shape NoUniqueness]
arr_ts forall a b. (a -> b) -> a -> b
$ \TypeBase Shape NoUniqueness
arr_t ->
forall (m :: * -> *).
MonadBuilder m =>
[Char] -> Exp (Rep m) -> m VName
letExp [Char]
"partition_dest" forall a b. (a -> b) -> a -> b
$
forall rep. BasicOp -> Exp rep
I.BasicOp forall a b. (a -> b) -> a -> b
$
PrimType -> [SubExp] -> BasicOp
Scratch (forall shape u. TypeBase shape u -> PrimType
I.elemType TypeBase Shape NoUniqueness
arr_t) (SubExp
w forall a. a -> [a] -> [a]
: forall a. Int -> [a] -> [a]
drop Int
1 (forall u. TypeBase Shape u -> [SubExp]
I.arrayDims TypeBase Shape NoUniqueness
arr_t))
Lambda SOACS
write_lam <- do
Param (TypeBase Shape NoUniqueness)
c_param <- forall (m :: * -> *) dec.
MonadFreshNames m =>
[Char] -> dec -> m (Param dec)
newParam [Char]
"c" (forall shape u. PrimType -> TypeBase shape u
I.Prim PrimType
int64)
[Param (TypeBase Shape NoUniqueness)]
offset_params <- forall (m :: * -> *) a. Applicative m => Int -> m a -> m [a]
replicateM Int
k forall a b. (a -> b) -> a -> b
$ forall (m :: * -> *) dec.
MonadFreshNames m =>
[Char] -> dec -> m (Param dec)
newParam [Char]
"offset" (forall shape u. PrimType -> TypeBase shape u
I.Prim PrimType
int64)
[Param (TypeBase Shape NoUniqueness)]
value_params <- forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
(a -> m b) -> t a -> m (t b)
mapM (forall (m :: * -> *) dec.
MonadFreshNames m =>
[Char] -> dec -> m (Param dec)
newParam [Char]
"v" forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall u. TypeBase Shape u -> TypeBase Shape u
I.rowType) [TypeBase Shape NoUniqueness]
arr_ts
(SubExp
offset, Stms SOACS
offset_stms) <-
forall (m :: * -> *) a.
MonadBuilder m =>
m a -> m (a, Stms (Rep m))
collectStms forall a b. (a -> b) -> a -> b
$
[SubExp] -> SubExp -> Int -> [LParam SOACS] -> InternaliseM SubExp
mkOffsetLambdaBody
(forall a b. (a -> b) -> [a] -> [b]
map VName -> SubExp
I.Var [VName]
sizes)
(VName -> SubExp
I.Var forall a b. (a -> b) -> a -> b
$ forall dec. Param dec -> VName
I.paramName Param (TypeBase Shape NoUniqueness)
c_param)
Int
0
[Param (TypeBase Shape NoUniqueness)]
offset_params
forall (f :: * -> *) a. Applicative f => a -> f a
pure
I.Lambda
{ lambdaParams :: [LParam SOACS]
I.lambdaParams = Param (TypeBase Shape NoUniqueness)
c_param forall a. a -> [a] -> [a]
: [Param (TypeBase Shape NoUniqueness)]
offset_params forall a. [a] -> [a] -> [a]
++ [Param (TypeBase Shape NoUniqueness)]
value_params,
lambdaReturnType :: [TypeBase Shape NoUniqueness]
I.lambdaReturnType =
forall a. Int -> a -> [a]
replicate (forall (t :: * -> *) a. Foldable t => t a -> Int
length [TypeBase Shape NoUniqueness]
arr_ts) (forall shape u. PrimType -> TypeBase shape u
I.Prim PrimType
int64)
forall a. [a] -> [a] -> [a]
++ forall a b. (a -> b) -> [a] -> [b]
map forall u. TypeBase Shape u -> TypeBase Shape u
I.rowType [TypeBase Shape NoUniqueness]
arr_ts,
lambdaBody :: Body SOACS
I.lambdaBody =
forall rep. Buildable rep => Stms rep -> Result -> Body rep
mkBody Stms SOACS
offset_stms forall a b. (a -> b) -> a -> b
$
forall a. Int -> a -> [a]
replicate (forall (t :: * -> *) a. Foldable t => t a -> Int
length [TypeBase Shape NoUniqueness]
arr_ts) (SubExp -> SubExpRes
subExpRes SubExp
offset)
forall a. [a] -> [a] -> [a]
++ [VName] -> Result
I.varsRes (forall a b. (a -> b) -> [a] -> [b]
map forall dec. Param dec -> VName
I.paramName [Param (TypeBase Shape NoUniqueness)]
value_params)
}
[VName]
results <-
forall (m :: * -> *).
MonadBuilder m =>
[Char] -> Exp (Rep m) -> m [VName]
letTupExp [Char]
"partition_res" forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall rep. Op rep -> Exp rep
I.Op forall a b. (a -> b) -> a -> b
$
forall rep.
SubExp
-> [VName] -> Lambda rep -> [(Shape, Int, VName)] -> SOAC rep
I.Scatter SubExp
w (VName
classes forall a. a -> [a] -> [a]
: [VName]
all_offsets forall a. [a] -> [a] -> [a]
++ [VName]
arrs) Lambda SOACS
write_lam forall a b. (a -> b) -> a -> b
$
forall a b c. [a] -> [b] -> [c] -> [(a, b, c)]
zip3 (forall a. a -> [a]
repeat forall a b. (a -> b) -> a -> b
$ forall d. [d] -> ShapeBase d
I.Shape [SubExp
w]) (forall a. a -> [a]
repeat Int
1) [VName]
blanks
SubExp
sizes' <-
forall (m :: * -> *).
MonadBuilder m =>
[Char] -> Exp (Rep m) -> m SubExp
letSubExp [Char]
"partition_sizes" forall a b. (a -> b) -> a -> b
$
forall rep. BasicOp -> Exp rep
I.BasicOp forall a b. (a -> b) -> a -> b
$
[SubExp] -> TypeBase Shape NoUniqueness -> BasicOp
I.ArrayLit (forall a b. (a -> b) -> [a] -> [b]
map VName -> SubExp
I.Var [VName]
sizes) forall a b. (a -> b) -> a -> b
$
forall shape u. PrimType -> TypeBase shape u
I.Prim PrimType
int64
forall (f :: * -> *) a. Applicative f => a -> f a
pure (forall a b. (a -> b) -> [a] -> [b]
map VName -> SubExp
I.Var [VName]
results, [SubExp
sizes'])
where
mkOffsetLambdaBody ::
[SubExp] ->
SubExp ->
Int ->
[I.LParam SOACS] ->
InternaliseM SubExp
mkOffsetLambdaBody :: [SubExp] -> SubExp -> Int -> [LParam SOACS] -> InternaliseM SubExp
mkOffsetLambdaBody [SubExp]
_ SubExp
_ Int
_ [] =
forall (f :: * -> *) a. Applicative f => a -> f a
pure forall a b. (a -> b) -> a -> b
$ forall v. IsValue v => v -> SubExp
constant (-Int64
1 :: Int64)
mkOffsetLambdaBody [SubExp]
sizes SubExp
c Int
i (LParam SOACS
p : [LParam SOACS]
ps) = do
SubExp
is_this_one <-
forall (m :: * -> *).
MonadBuilder m =>
[Char] -> Exp (Rep m) -> m SubExp
letSubExp [Char]
"is_this_one" forall a b. (a -> b) -> a -> b
$
forall rep. BasicOp -> Exp rep
I.BasicOp forall a b. (a -> b) -> a -> b
$
CmpOp -> SubExp -> SubExp -> BasicOp
I.CmpOp (PrimType -> CmpOp
CmpEq PrimType
int64) SubExp
c forall a b. (a -> b) -> a -> b
$
IntType -> Integer -> SubExp
intConst IntType
Int64 forall a b. (a -> b) -> a -> b
$
forall a. Integral a => a -> Integer
toInteger Int
i
SubExp
next_one <- [SubExp] -> SubExp -> Int -> [LParam SOACS] -> InternaliseM SubExp
mkOffsetLambdaBody [SubExp]
sizes SubExp
c (Int
i forall a. Num a => a -> a -> a
+ Int
1) [LParam SOACS]
ps
SubExp
this_one <-
forall (m :: * -> *).
MonadBuilder m =>
[Char] -> Exp (Rep m) -> m SubExp
letSubExp [Char]
"this_offset"
forall (m :: * -> *) a b. Monad m => (a -> m b) -> m a -> m b
=<< forall (m :: * -> *).
MonadBuilder m =>
BinOp -> SubExp -> [SubExp] -> m (Exp (Rep m))
foldBinOp
(IntType -> Overflow -> BinOp
Add IntType
Int64 Overflow
OverflowUndef)
(forall v. IsValue v => v -> SubExp
constant (-Int64
1 :: Int64))
(VName -> SubExp
I.Var (forall dec. Param dec -> VName
I.paramName LParam SOACS
p) forall a. a -> [a] -> [a]
: forall a. Int -> [a] -> [a]
take Int
i [SubExp]
sizes)
forall (m :: * -> *).
MonadBuilder m =>
[Char] -> Exp (Rep m) -> m SubExp
letSubExp [Char]
"total_res"
forall (m :: * -> *) a b. Monad m => (a -> m b) -> m a -> m b
=<< forall (m :: * -> *).
(MonadBuilder m, BranchType (Rep m) ~ ExtType) =>
m (Exp (Rep m))
-> m (Body (Rep m)) -> m (Body (Rep m)) -> m (Exp (Rep m))
eIf
(forall (m :: * -> *). MonadBuilder m => SubExp -> m (Exp (Rep m))
eSubExp SubExp
is_this_one)
(forall (m :: * -> *).
MonadBuilder m =>
[SubExp] -> m (Body (Rep m))
resultBodyM [SubExp
this_one])
(forall (m :: * -> *).
MonadBuilder m =>
[SubExp] -> m (Body (Rep m))
resultBodyM [SubExp
next_one])
typeExpForError :: E.TypeExp VName -> InternaliseM [ErrorMsgPart SubExp]
typeExpForError :: TypeExp VName -> InternaliseM [ErrorMsgPart SubExp]
typeExpForError (E.TEVar QualName VName
qn SrcLoc
_) =
forall (f :: * -> *) a. Applicative f => a -> f a
pure [forall a. Text -> ErrorMsgPart a
ErrorString forall a b. (a -> b) -> a -> b
$ forall a. Pretty a => a -> Text
prettyText QualName VName
qn]
typeExpForError (E.TEUnique TypeExp VName
te SrcLoc
_) =
(ErrorMsgPart SubExp
"*" :) forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> TypeExp VName -> InternaliseM [ErrorMsgPart SubExp]
typeExpForError TypeExp VName
te
typeExpForError (E.TEDim [VName]
dims TypeExp VName
te SrcLoc
_) =
(forall a. Text -> ErrorMsgPart a
ErrorString (Text
"?" forall a. Semigroup a => a -> a -> a
<> Text
dims' forall a. Semigroup a => a -> a -> a
<> Text
".") :) forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> TypeExp VName -> InternaliseM [ErrorMsgPart SubExp]
typeExpForError TypeExp VName
te
where
dims' :: Text
dims' = forall a. Monoid a => [a] -> a
mconcat (forall a b. (a -> b) -> [a] -> [b]
map forall a. Pretty a => a -> Text
onDim [VName]
dims)
onDim :: a -> Text
onDim a
d = Text
"[" forall a. Semigroup a => a -> a -> a
<> forall a. Pretty a => a -> Text
prettyText a
d forall a. Semigroup a => a -> a -> a
<> Text
"]"
typeExpForError (E.TEArray SizeExp VName
d TypeExp VName
te SrcLoc
_) = do
ErrorMsgPart SubExp
d' <- SizeExp VName -> InternaliseM (ErrorMsgPart SubExp)
dimExpForError SizeExp VName
d
[ErrorMsgPart SubExp]
te' <- TypeExp VName -> InternaliseM [ErrorMsgPart SubExp]
typeExpForError TypeExp VName
te
forall (f :: * -> *) a. Applicative f => a -> f a
pure forall a b. (a -> b) -> a -> b
$ [ErrorMsgPart SubExp
"[", ErrorMsgPart SubExp
d', ErrorMsgPart SubExp
"]"] forall a. [a] -> [a] -> [a]
++ [ErrorMsgPart SubExp]
te'
typeExpForError (E.TETuple [TypeExp VName]
tes SrcLoc
_) = do
[[ErrorMsgPart SubExp]]
tes' <- forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
(a -> m b) -> t a -> m (t b)
mapM TypeExp VName -> InternaliseM [ErrorMsgPart SubExp]
typeExpForError [TypeExp VName]
tes
forall (f :: * -> *) a. Applicative f => a -> f a
pure forall a b. (a -> b) -> a -> b
$ [ErrorMsgPart SubExp
"("] forall a. [a] -> [a] -> [a]
++ forall a. [a] -> [[a]] -> [a]
intercalate [ErrorMsgPart SubExp
", "] [[ErrorMsgPart SubExp]]
tes' forall a. [a] -> [a] -> [a]
++ [ErrorMsgPart SubExp
")"]
typeExpForError (E.TERecord [(Name, TypeExp VName)]
fields SrcLoc
_) = do
[[ErrorMsgPart SubExp]]
fields' <- forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
(a -> m b) -> t a -> m (t b)
mapM forall {a}.
Pretty a =>
(a, TypeExp VName) -> InternaliseM [ErrorMsgPart SubExp]
onField [(Name, TypeExp VName)]
fields
forall (f :: * -> *) a. Applicative f => a -> f a
pure forall a b. (a -> b) -> a -> b
$ [ErrorMsgPart SubExp
"{"] forall a. [a] -> [a] -> [a]
++ forall a. [a] -> [[a]] -> [a]
intercalate [ErrorMsgPart SubExp
", "] [[ErrorMsgPart SubExp]]
fields' forall a. [a] -> [a] -> [a]
++ [ErrorMsgPart SubExp
"}"]
where
onField :: (a, TypeExp VName) -> InternaliseM [ErrorMsgPart SubExp]
onField (a
k, TypeExp VName
te) =
(forall a. Text -> ErrorMsgPart a
ErrorString (forall a. Pretty a => a -> Text
prettyText a
k forall a. Semigroup a => a -> a -> a
<> Text
": ") :) forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> TypeExp VName -> InternaliseM [ErrorMsgPart SubExp]
typeExpForError TypeExp VName
te
typeExpForError (E.TEArrow Maybe VName
_ TypeExp VName
t1 TypeExp VName
t2 SrcLoc
_) = do
[ErrorMsgPart SubExp]
t1' <- TypeExp VName -> InternaliseM [ErrorMsgPart SubExp]
typeExpForError TypeExp VName
t1
[ErrorMsgPart SubExp]
t2' <- TypeExp VName -> InternaliseM [ErrorMsgPart SubExp]
typeExpForError TypeExp VName
t2
forall (f :: * -> *) a. Applicative f => a -> f a
pure forall a b. (a -> b) -> a -> b
$ [ErrorMsgPart SubExp]
t1' forall a. [a] -> [a] -> [a]
++ [ErrorMsgPart SubExp
" -> "] forall a. [a] -> [a] -> [a]
++ [ErrorMsgPart SubExp]
t2'
typeExpForError (E.TEApply TypeExp VName
t TypeArgExp VName
arg SrcLoc
_) = do
[ErrorMsgPart SubExp]
t' <- TypeExp VName -> InternaliseM [ErrorMsgPart SubExp]
typeExpForError TypeExp VName
t
[ErrorMsgPart SubExp]
arg' <- case TypeArgExp VName
arg of
TypeArgExpType TypeExp VName
argt -> TypeExp VName -> InternaliseM [ErrorMsgPart SubExp]
typeExpForError TypeExp VName
argt
TypeArgExpDim SizeExp VName
d SrcLoc
_ -> forall (f :: * -> *) a. Applicative f => a -> f a
pure forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> SizeExp VName -> InternaliseM (ErrorMsgPart SubExp)
dimExpForError SizeExp VName
d
forall (f :: * -> *) a. Applicative f => a -> f a
pure forall a b. (a -> b) -> a -> b
$ [ErrorMsgPart SubExp]
t' forall a. [a] -> [a] -> [a]
++ [ErrorMsgPart SubExp
" "] forall a. [a] -> [a] -> [a]
++ [ErrorMsgPart SubExp]
arg'
typeExpForError (E.TESum [(Name, [TypeExp VName])]
cs SrcLoc
_) = do
[[ErrorMsgPart SubExp]]
cs' <- forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
(a -> m b) -> t a -> m (t b)
mapM ([TypeExp VName] -> InternaliseM [ErrorMsgPart SubExp]
onClause forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall a b. (a, b) -> b
snd) [(Name, [TypeExp VName])]
cs
forall (f :: * -> *) a. Applicative f => a -> f a
pure forall a b. (a -> b) -> a -> b
$ forall a. [a] -> [[a]] -> [a]
intercalate [ErrorMsgPart SubExp
" | "] [[ErrorMsgPart SubExp]]
cs'
where
onClause :: [TypeExp VName] -> InternaliseM [ErrorMsgPart SubExp]
onClause [TypeExp VName]
c = do
[[ErrorMsgPart SubExp]]
c' <- forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
(a -> m b) -> t a -> m (t b)
mapM TypeExp VName -> InternaliseM [ErrorMsgPart SubExp]
typeExpForError [TypeExp VName]
c
forall (f :: * -> *) a. Applicative f => a -> f a
pure forall a b. (a -> b) -> a -> b
$ forall a. [a] -> [[a]] -> [a]
intercalate [ErrorMsgPart SubExp
" "] [[ErrorMsgPart SubExp]]
c'
dimExpForError :: E.SizeExp VName -> InternaliseM (ErrorMsgPart SubExp)
dimExpForError :: SizeExp VName -> InternaliseM (ErrorMsgPart SubExp)
dimExpForError (SizeExpNamed QualName VName
d SrcLoc
_) = do
Maybe [SubExp]
substs <- VName -> InternaliseM (Maybe [SubExp])
lookupSubst forall a b. (a -> b) -> a -> b
$ forall vn. QualName vn -> vn
E.qualLeaf QualName VName
d
SubExp
d' <- case Maybe [SubExp]
substs of
Just [SubExp
v] -> forall (f :: * -> *) a. Applicative f => a -> f a
pure SubExp
v
Maybe [SubExp]
_ -> forall (f :: * -> *) a. Applicative f => a -> f a
pure forall a b. (a -> b) -> a -> b
$ VName -> SubExp
I.Var forall a b. (a -> b) -> a -> b
$ forall vn. QualName vn -> vn
E.qualLeaf QualName VName
d
forall (f :: * -> *) a. Applicative f => a -> f a
pure forall a b. (a -> b) -> a -> b
$ forall a. PrimType -> a -> ErrorMsgPart a
ErrorVal PrimType
int64 SubExp
d'
dimExpForError (SizeExpConst Int
d SrcLoc
_) =
forall (f :: * -> *) a. Applicative f => a -> f a
pure forall a b. (a -> b) -> a -> b
$ forall a. Text -> ErrorMsgPart a
ErrorString forall a b. (a -> b) -> a -> b
$ forall a. Pretty a => a -> Text
prettyText Int
d
dimExpForError SizeExp VName
SizeExpAny = forall (f :: * -> *) a. Applicative f => a -> f a
pure ErrorMsgPart SubExp
""
errorMsg :: [ErrorMsgPart a] -> ErrorMsg a
errorMsg :: forall a. [ErrorMsgPart a] -> ErrorMsg a
errorMsg = forall a. [ErrorMsgPart a] -> ErrorMsg a
ErrorMsg forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall {a}. [ErrorMsgPart a] -> [ErrorMsgPart a]
compact
where
compact :: [ErrorMsgPart a] -> [ErrorMsgPart a]
compact [] = []
compact (ErrorString Text
x : ErrorString Text
y : [ErrorMsgPart a]
parts) =
[ErrorMsgPart a] -> [ErrorMsgPart a]
compact (forall a. Text -> ErrorMsgPart a
ErrorString (Text
x forall a. Semigroup a => a -> a -> a
<> Text
y) forall a. a -> [a] -> [a]
: [ErrorMsgPart a]
parts)
compact (ErrorMsgPart a
x : [ErrorMsgPart a]
y) = ErrorMsgPart a
x forall a. a -> [a] -> [a]
: [ErrorMsgPart a] -> [ErrorMsgPart a]
compact [ErrorMsgPart a]
y