{-# LANGUAGE FlexibleContexts #-}
module Futhark.Internalise.Lambdas
( InternaliseLambda,
internaliseStreamMapLambda,
internaliseFoldLambda,
internaliseStreamLambda,
internalisePartitionLambda,
)
where
import Futhark.IR.SOACS as I
import Futhark.Internalise.AccurateSizes
import Futhark.Internalise.Monad
import Language.Futhark as E
type InternaliseLambda =
E.Exp -> [I.Type] -> InternaliseM ([I.LParam SOACS], I.Body SOACS, [I.Type])
internaliseStreamMapLambda ::
InternaliseLambda ->
E.Exp ->
[I.SubExp] ->
InternaliseM (I.Lambda SOACS)
internaliseStreamMapLambda :: InternaliseLambda -> Exp -> [SubExp] -> InternaliseM (Lambda SOACS)
internaliseStreamMapLambda InternaliseLambda
internaliseLambda Exp
lam [SubExp]
args = do
VName
chunk_size <- String -> InternaliseM VName
forall (m :: * -> *). MonadFreshNames m => String -> m VName
newVName String
"chunk_size"
let chunk_param :: Param (TypeBase shape u)
chunk_param = Attrs -> VName -> TypeBase shape u -> Param (TypeBase shape u)
forall dec. Attrs -> VName -> dec -> Param dec
I.Param Attrs
forall a. Monoid a => a
mempty VName
chunk_size (PrimType -> TypeBase shape u
forall shape u. PrimType -> TypeBase shape u
I.Prim PrimType
int64)
outer :: TypeBase (ShapeBase SubExp) u -> TypeBase (ShapeBase SubExp) u
outer = (TypeBase (ShapeBase SubExp) u
-> SubExp -> TypeBase (ShapeBase SubExp) u
forall d u.
ArrayShape (ShapeBase d) =>
TypeBase (ShapeBase d) u -> d -> TypeBase (ShapeBase d) u
`setOuterSize` VName -> SubExp
I.Var VName
chunk_size)
Scope SOACS
-> InternaliseM (Lambda SOACS) -> InternaliseM (Lambda SOACS)
forall rep (m :: * -> *) a.
LocalScope rep m =>
Scope rep -> m a -> m a
localScope ([Param (TypeBase (ShapeBase SubExp) NoUniqueness)] -> Scope SOACS
forall rep dec. (LParamInfo rep ~ dec) => [Param dec] -> Scope rep
scopeOfLParams [Param (TypeBase (ShapeBase SubExp) NoUniqueness)
forall shape u. Param (TypeBase shape u)
chunk_param]) (InternaliseM (Lambda SOACS) -> InternaliseM (Lambda SOACS))
-> InternaliseM (Lambda SOACS) -> InternaliseM (Lambda SOACS)
forall a b. (a -> b) -> a -> b
$ do
[TypeBase (ShapeBase SubExp) NoUniqueness]
argtypes <- (SubExp -> InternaliseM (TypeBase (ShapeBase SubExp) NoUniqueness))
-> [SubExp]
-> InternaliseM [TypeBase (ShapeBase SubExp) NoUniqueness]
forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
(a -> m b) -> t a -> m (t b)
mapM SubExp -> InternaliseM (TypeBase (ShapeBase SubExp) NoUniqueness)
forall t (m :: * -> *).
HasScope t m =>
SubExp -> m (TypeBase (ShapeBase SubExp) NoUniqueness)
I.subExpType [SubExp]
args
([Param (TypeBase (ShapeBase SubExp) NoUniqueness)]
lam_params, Body SOACS
orig_body, [TypeBase (ShapeBase SubExp) NoUniqueness]
rettype) <-
InternaliseLambda
internaliseLambda Exp
lam ([TypeBase (ShapeBase SubExp) NoUniqueness]
-> InternaliseM
([LParam SOACS], Body SOACS,
[TypeBase (ShapeBase SubExp) NoUniqueness]))
-> [TypeBase (ShapeBase SubExp) NoUniqueness]
-> InternaliseM
([LParam SOACS], Body SOACS,
[TypeBase (ShapeBase SubExp) NoUniqueness])
forall a b. (a -> b) -> a -> b
$ PrimType -> TypeBase (ShapeBase SubExp) NoUniqueness
forall shape u. PrimType -> TypeBase shape u
I.Prim PrimType
int64 TypeBase (ShapeBase SubExp) NoUniqueness
-> [TypeBase (ShapeBase SubExp) NoUniqueness]
-> [TypeBase (ShapeBase SubExp) NoUniqueness]
forall a. a -> [a] -> [a]
: (TypeBase (ShapeBase SubExp) NoUniqueness
-> TypeBase (ShapeBase SubExp) NoUniqueness)
-> [TypeBase (ShapeBase SubExp) NoUniqueness]
-> [TypeBase (ShapeBase SubExp) NoUniqueness]
forall a b. (a -> b) -> [a] -> [b]
map TypeBase (ShapeBase SubExp) NoUniqueness
-> TypeBase (ShapeBase SubExp) NoUniqueness
forall u.
TypeBase (ShapeBase SubExp) u -> TypeBase (ShapeBase SubExp) u
outer [TypeBase (ShapeBase SubExp) NoUniqueness]
argtypes
let Param (TypeBase (ShapeBase SubExp) NoUniqueness)
orig_chunk_param : [Param (TypeBase (ShapeBase SubExp) NoUniqueness)]
params = [Param (TypeBase (ShapeBase SubExp) NoUniqueness)]
lam_params
Body SOACS
body <- Builder SOACS (Body SOACS) -> InternaliseM (Body SOACS)
forall rep (m :: * -> *) somerep.
(Buildable rep, MonadFreshNames m, HasScope somerep m,
SameScope somerep rep) =>
Builder rep (Body rep) -> m (Body rep)
runBodyBuilder (Builder SOACS (Body SOACS) -> InternaliseM (Body SOACS))
-> Builder SOACS (Body SOACS) -> InternaliseM (Body SOACS)
forall a b. (a -> b) -> a -> b
$ do
[VName]
-> Exp (Rep (BuilderT SOACS (State VNameSource)))
-> BuilderT SOACS (State VNameSource) ()
forall (m :: * -> *).
MonadBuilder m =>
[VName] -> Exp (Rep m) -> m ()
letBindNames [Param (TypeBase (ShapeBase SubExp) NoUniqueness) -> VName
forall dec. Param dec -> VName
paramName Param (TypeBase (ShapeBase SubExp) NoUniqueness)
orig_chunk_param] (Exp (Rep (BuilderT SOACS (State VNameSource)))
-> BuilderT SOACS (State VNameSource) ())
-> Exp (Rep (BuilderT SOACS (State VNameSource)))
-> BuilderT SOACS (State VNameSource) ()
forall a b. (a -> b) -> a -> b
$ BasicOp -> Exp SOACS
forall rep. BasicOp -> Exp rep
I.BasicOp (BasicOp -> Exp SOACS) -> BasicOp -> Exp SOACS
forall a b. (a -> b) -> a -> b
$ SubExp -> BasicOp
I.SubExp (SubExp -> BasicOp) -> SubExp -> BasicOp
forall a b. (a -> b) -> a -> b
$ VName -> SubExp
I.Var VName
chunk_size
Body SOACS -> Builder SOACS (Body SOACS)
forall (f :: * -> *) a. Applicative f => a -> f a
pure Body SOACS
orig_body
[LParam (Rep InternaliseM)]
-> InternaliseM Result -> InternaliseM (Lambda (Rep InternaliseM))
forall (m :: * -> *).
MonadBuilder m =>
[LParam (Rep m)] -> m Result -> m (Lambda (Rep m))
mkLambda (Param (TypeBase (ShapeBase SubExp) NoUniqueness)
forall shape u. Param (TypeBase shape u)
chunk_param Param (TypeBase (ShapeBase SubExp) NoUniqueness)
-> [Param (TypeBase (ShapeBase SubExp) NoUniqueness)]
-> [Param (TypeBase (ShapeBase SubExp) NoUniqueness)]
forall a. a -> [a] -> [a]
: [Param (TypeBase (ShapeBase SubExp) NoUniqueness)]
params) (InternaliseM Result -> InternaliseM (Lambda (Rep InternaliseM)))
-> InternaliseM Result -> InternaliseM (Lambda (Rep InternaliseM))
forall a b. (a -> b) -> a -> b
$ do
[VName] -> Exp (Rep InternaliseM) -> InternaliseM ()
forall (m :: * -> *).
MonadBuilder m =>
[VName] -> Exp (Rep m) -> m ()
letBindNames [Param (TypeBase (ShapeBase SubExp) NoUniqueness) -> VName
forall dec. Param dec -> VName
paramName Param (TypeBase (ShapeBase SubExp) NoUniqueness)
orig_chunk_param] (Exp (Rep InternaliseM) -> InternaliseM ())
-> Exp (Rep InternaliseM) -> InternaliseM ()
forall a b. (a -> b) -> a -> b
$ BasicOp -> Exp SOACS
forall rep. BasicOp -> Exp rep
I.BasicOp (BasicOp -> Exp SOACS) -> BasicOp -> Exp SOACS
forall a b. (a -> b) -> a -> b
$ SubExp -> BasicOp
I.SubExp (SubExp -> BasicOp) -> SubExp -> BasicOp
forall a b. (a -> b) -> a -> b
$ VName -> SubExp
I.Var VName
chunk_size
ErrorMsg SubExp
-> SrcLoc
-> [TypeBase (ShapeBase SubExp) NoUniqueness]
-> Result
-> InternaliseM Result
ensureResultShape
([ErrorMsgPart SubExp] -> ErrorMsg SubExp
forall a. [ErrorMsgPart a] -> ErrorMsg a
ErrorMsg [String -> ErrorMsgPart SubExp
forall a. String -> ErrorMsgPart a
ErrorString String
"not all iterations produce same shape"])
(Exp -> SrcLoc
forall a. Located a => a -> SrcLoc
srclocOf Exp
lam)
((TypeBase (ShapeBase SubExp) NoUniqueness
-> TypeBase (ShapeBase SubExp) NoUniqueness)
-> [TypeBase (ShapeBase SubExp) NoUniqueness]
-> [TypeBase (ShapeBase SubExp) NoUniqueness]
forall a b. (a -> b) -> [a] -> [b]
map TypeBase (ShapeBase SubExp) NoUniqueness
-> TypeBase (ShapeBase SubExp) NoUniqueness
forall u.
TypeBase (ShapeBase SubExp) u -> TypeBase (ShapeBase SubExp) u
outer [TypeBase (ShapeBase SubExp) NoUniqueness]
rettype)
(Result -> InternaliseM Result)
-> InternaliseM Result -> InternaliseM Result
forall (m :: * -> *) a b. Monad m => (a -> m b) -> m a -> m b
=<< Body (Rep InternaliseM) -> InternaliseM Result
forall (m :: * -> *). MonadBuilder m => Body (Rep m) -> m Result
bodyBind Body (Rep InternaliseM)
Body SOACS
body
internaliseFoldLambda ::
InternaliseLambda ->
E.Exp ->
[I.Type] ->
[I.Type] ->
InternaliseM (I.Lambda SOACS)
internaliseFoldLambda :: InternaliseLambda
-> Exp
-> [TypeBase (ShapeBase SubExp) NoUniqueness]
-> [TypeBase (ShapeBase SubExp) NoUniqueness]
-> InternaliseM (Lambda SOACS)
internaliseFoldLambda InternaliseLambda
internaliseLambda Exp
lam [TypeBase (ShapeBase SubExp) NoUniqueness]
acctypes [TypeBase (ShapeBase SubExp) NoUniqueness]
arrtypes = do
let rowtypes :: [TypeBase (ShapeBase SubExp) NoUniqueness]
rowtypes = (TypeBase (ShapeBase SubExp) NoUniqueness
-> TypeBase (ShapeBase SubExp) NoUniqueness)
-> [TypeBase (ShapeBase SubExp) NoUniqueness]
-> [TypeBase (ShapeBase SubExp) NoUniqueness]
forall a b. (a -> b) -> [a] -> [b]
map TypeBase (ShapeBase SubExp) NoUniqueness
-> TypeBase (ShapeBase SubExp) NoUniqueness
forall u.
TypeBase (ShapeBase SubExp) u -> TypeBase (ShapeBase SubExp) u
I.rowType [TypeBase (ShapeBase SubExp) NoUniqueness]
arrtypes
([Param (TypeBase (ShapeBase SubExp) NoUniqueness)]
params, Body SOACS
body, [TypeBase (ShapeBase SubExp) NoUniqueness]
rettype) <- InternaliseLambda
internaliseLambda Exp
lam ([TypeBase (ShapeBase SubExp) NoUniqueness]
-> InternaliseM
([LParam SOACS], Body SOACS,
[TypeBase (ShapeBase SubExp) NoUniqueness]))
-> [TypeBase (ShapeBase SubExp) NoUniqueness]
-> InternaliseM
([LParam SOACS], Body SOACS,
[TypeBase (ShapeBase SubExp) NoUniqueness])
forall a b. (a -> b) -> a -> b
$ [TypeBase (ShapeBase SubExp) NoUniqueness]
acctypes [TypeBase (ShapeBase SubExp) NoUniqueness]
-> [TypeBase (ShapeBase SubExp) NoUniqueness]
-> [TypeBase (ShapeBase SubExp) NoUniqueness]
forall a. [a] -> [a] -> [a]
++ [TypeBase (ShapeBase SubExp) NoUniqueness]
rowtypes
let rettype' :: [TypeBase (ShapeBase SubExp) NoUniqueness]
rettype' =
[ TypeBase (ShapeBase SubExp) NoUniqueness
t TypeBase (ShapeBase SubExp) NoUniqueness
-> ShapeBase SubExp -> TypeBase (ShapeBase SubExp) NoUniqueness
forall newshape oldshape u.
ArrayShape newshape =>
TypeBase oldshape u -> newshape -> TypeBase newshape u
`I.setArrayShape` TypeBase (ShapeBase SubExp) NoUniqueness -> ShapeBase SubExp
forall shape u. ArrayShape shape => TypeBase shape u -> shape
I.arrayShape TypeBase (ShapeBase SubExp) NoUniqueness
shape
| (TypeBase (ShapeBase SubExp) NoUniqueness
t, TypeBase (ShapeBase SubExp) NoUniqueness
shape) <- [TypeBase (ShapeBase SubExp) NoUniqueness]
-> [TypeBase (ShapeBase SubExp) NoUniqueness]
-> [(TypeBase (ShapeBase SubExp) NoUniqueness,
TypeBase (ShapeBase SubExp) NoUniqueness)]
forall a b. [a] -> [b] -> [(a, b)]
zip [TypeBase (ShapeBase SubExp) NoUniqueness]
rettype [TypeBase (ShapeBase SubExp) NoUniqueness]
acctypes
]
[LParam (Rep InternaliseM)]
-> InternaliseM Result -> InternaliseM (Lambda (Rep InternaliseM))
forall (m :: * -> *).
MonadBuilder m =>
[LParam (Rep m)] -> m Result -> m (Lambda (Rep m))
mkLambda [Param (TypeBase (ShapeBase SubExp) NoUniqueness)]
[LParam (Rep InternaliseM)]
params (InternaliseM Result -> InternaliseM (Lambda (Rep InternaliseM)))
-> InternaliseM Result -> InternaliseM (Lambda (Rep InternaliseM))
forall a b. (a -> b) -> a -> b
$
ErrorMsg SubExp
-> SrcLoc
-> [TypeBase (ShapeBase SubExp) NoUniqueness]
-> Result
-> InternaliseM Result
ensureResultShape
([ErrorMsgPart SubExp] -> ErrorMsg SubExp
forall a. [ErrorMsgPart a] -> ErrorMsg a
ErrorMsg [String -> ErrorMsgPart SubExp
forall a. String -> ErrorMsgPart a
ErrorString String
"shape of result does not match shape of initial value"])
(Exp -> SrcLoc
forall a. Located a => a -> SrcLoc
srclocOf Exp
lam)
[TypeBase (ShapeBase SubExp) NoUniqueness]
rettype'
(Result -> InternaliseM Result)
-> InternaliseM Result -> InternaliseM Result
forall (m :: * -> *) a b. Monad m => (a -> m b) -> m a -> m b
=<< Body (Rep InternaliseM) -> InternaliseM Result
forall (m :: * -> *). MonadBuilder m => Body (Rep m) -> m Result
bodyBind Body (Rep InternaliseM)
Body SOACS
body
internaliseStreamLambda ::
InternaliseLambda ->
E.Exp ->
[I.Type] ->
InternaliseM ([LParam SOACS], Body SOACS)
internaliseStreamLambda :: InternaliseLambda
-> Exp
-> [TypeBase (ShapeBase SubExp) NoUniqueness]
-> InternaliseM ([LParam SOACS], Body SOACS)
internaliseStreamLambda InternaliseLambda
internaliseLambda Exp
lam [TypeBase (ShapeBase SubExp) NoUniqueness]
rowts = do
VName
chunk_size <- String -> InternaliseM VName
forall (m :: * -> *). MonadFreshNames m => String -> m VName
newVName String
"chunk_size"
let chunk_param :: Param (TypeBase shape u)
chunk_param = Attrs -> VName -> TypeBase shape u -> Param (TypeBase shape u)
forall dec. Attrs -> VName -> dec -> Param dec
I.Param Attrs
forall a. Monoid a => a
mempty VName
chunk_size (TypeBase shape u -> Param (TypeBase shape u))
-> TypeBase shape u -> Param (TypeBase shape u)
forall a b. (a -> b) -> a -> b
$ PrimType -> TypeBase shape u
forall shape u. PrimType -> TypeBase shape u
I.Prim PrimType
int64
chunktypes :: [TypeBase (ShapeBase SubExp) NoUniqueness]
chunktypes = (TypeBase (ShapeBase SubExp) NoUniqueness
-> TypeBase (ShapeBase SubExp) NoUniqueness)
-> [TypeBase (ShapeBase SubExp) NoUniqueness]
-> [TypeBase (ShapeBase SubExp) NoUniqueness]
forall a b. (a -> b) -> [a] -> [b]
map (TypeBase (ShapeBase SubExp) NoUniqueness
-> SubExp -> TypeBase (ShapeBase SubExp) NoUniqueness
forall d.
ArrayShape (ShapeBase d) =>
TypeBase (ShapeBase d) NoUniqueness
-> d -> TypeBase (ShapeBase d) NoUniqueness
`arrayOfRow` VName -> SubExp
I.Var VName
chunk_size) [TypeBase (ShapeBase SubExp) NoUniqueness]
rowts
Scope SOACS
-> InternaliseM
([Param (TypeBase (ShapeBase SubExp) NoUniqueness)], Body SOACS)
-> InternaliseM
([Param (TypeBase (ShapeBase SubExp) NoUniqueness)], Body SOACS)
forall rep (m :: * -> *) a.
LocalScope rep m =>
Scope rep -> m a -> m a
localScope ([Param (TypeBase (ShapeBase SubExp) NoUniqueness)] -> Scope SOACS
forall rep dec. (LParamInfo rep ~ dec) => [Param dec] -> Scope rep
scopeOfLParams [Param (TypeBase (ShapeBase SubExp) NoUniqueness)
forall shape u. Param (TypeBase shape u)
chunk_param]) (InternaliseM
([Param (TypeBase (ShapeBase SubExp) NoUniqueness)], Body SOACS)
-> InternaliseM
([Param (TypeBase (ShapeBase SubExp) NoUniqueness)], Body SOACS))
-> InternaliseM
([Param (TypeBase (ShapeBase SubExp) NoUniqueness)], Body SOACS)
-> InternaliseM
([Param (TypeBase (ShapeBase SubExp) NoUniqueness)], Body SOACS)
forall a b. (a -> b) -> a -> b
$ do
([Param (TypeBase (ShapeBase SubExp) NoUniqueness)]
lam_params, Body SOACS
orig_body, [TypeBase (ShapeBase SubExp) NoUniqueness]
_) <-
InternaliseLambda
internaliseLambda Exp
lam ([TypeBase (ShapeBase SubExp) NoUniqueness]
-> InternaliseM
([LParam SOACS], Body SOACS,
[TypeBase (ShapeBase SubExp) NoUniqueness]))
-> [TypeBase (ShapeBase SubExp) NoUniqueness]
-> InternaliseM
([LParam SOACS], Body SOACS,
[TypeBase (ShapeBase SubExp) NoUniqueness])
forall a b. (a -> b) -> a -> b
$ PrimType -> TypeBase (ShapeBase SubExp) NoUniqueness
forall shape u. PrimType -> TypeBase shape u
I.Prim PrimType
int64 TypeBase (ShapeBase SubExp) NoUniqueness
-> [TypeBase (ShapeBase SubExp) NoUniqueness]
-> [TypeBase (ShapeBase SubExp) NoUniqueness]
forall a. a -> [a] -> [a]
: [TypeBase (ShapeBase SubExp) NoUniqueness]
chunktypes
let Param (TypeBase (ShapeBase SubExp) NoUniqueness)
orig_chunk_param : [Param (TypeBase (ShapeBase SubExp) NoUniqueness)]
params = [Param (TypeBase (ShapeBase SubExp) NoUniqueness)]
lam_params
Body SOACS
body <- Builder SOACS (Body SOACS) -> InternaliseM (Body SOACS)
forall rep (m :: * -> *) somerep.
(Buildable rep, MonadFreshNames m, HasScope somerep m,
SameScope somerep rep) =>
Builder rep (Body rep) -> m (Body rep)
runBodyBuilder (Builder SOACS (Body SOACS) -> InternaliseM (Body SOACS))
-> Builder SOACS (Body SOACS) -> InternaliseM (Body SOACS)
forall a b. (a -> b) -> a -> b
$ do
[VName]
-> Exp (Rep (BuilderT SOACS (State VNameSource)))
-> BuilderT SOACS (State VNameSource) ()
forall (m :: * -> *).
MonadBuilder m =>
[VName] -> Exp (Rep m) -> m ()
letBindNames [Param (TypeBase (ShapeBase SubExp) NoUniqueness) -> VName
forall dec. Param dec -> VName
paramName Param (TypeBase (ShapeBase SubExp) NoUniqueness)
orig_chunk_param] (Exp (Rep (BuilderT SOACS (State VNameSource)))
-> BuilderT SOACS (State VNameSource) ())
-> Exp (Rep (BuilderT SOACS (State VNameSource)))
-> BuilderT SOACS (State VNameSource) ()
forall a b. (a -> b) -> a -> b
$ BasicOp -> Exp SOACS
forall rep. BasicOp -> Exp rep
I.BasicOp (BasicOp -> Exp SOACS) -> BasicOp -> Exp SOACS
forall a b. (a -> b) -> a -> b
$ SubExp -> BasicOp
I.SubExp (SubExp -> BasicOp) -> SubExp -> BasicOp
forall a b. (a -> b) -> a -> b
$ VName -> SubExp
I.Var VName
chunk_size
Body SOACS -> Builder SOACS (Body SOACS)
forall (f :: * -> *) a. Applicative f => a -> f a
pure Body SOACS
orig_body
([Param (TypeBase (ShapeBase SubExp) NoUniqueness)], Body SOACS)
-> InternaliseM
([Param (TypeBase (ShapeBase SubExp) NoUniqueness)], Body SOACS)
forall (f :: * -> *) a. Applicative f => a -> f a
pure (Param (TypeBase (ShapeBase SubExp) NoUniqueness)
forall shape u. Param (TypeBase shape u)
chunk_param Param (TypeBase (ShapeBase SubExp) NoUniqueness)
-> [Param (TypeBase (ShapeBase SubExp) NoUniqueness)]
-> [Param (TypeBase (ShapeBase SubExp) NoUniqueness)]
forall a. a -> [a] -> [a]
: [Param (TypeBase (ShapeBase SubExp) NoUniqueness)]
params, Body SOACS
body)
internalisePartitionLambda ::
InternaliseLambda ->
Int ->
E.Exp ->
[I.SubExp] ->
InternaliseM (I.Lambda SOACS)
internalisePartitionLambda :: InternaliseLambda
-> Int -> Exp -> [SubExp] -> InternaliseM (Lambda SOACS)
internalisePartitionLambda InternaliseLambda
internaliseLambda Int
k Exp
lam [SubExp]
args = do
[TypeBase (ShapeBase SubExp) NoUniqueness]
argtypes <- (SubExp -> InternaliseM (TypeBase (ShapeBase SubExp) NoUniqueness))
-> [SubExp]
-> InternaliseM [TypeBase (ShapeBase SubExp) NoUniqueness]
forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
(a -> m b) -> t a -> m (t b)
mapM SubExp -> InternaliseM (TypeBase (ShapeBase SubExp) NoUniqueness)
forall t (m :: * -> *).
HasScope t m =>
SubExp -> m (TypeBase (ShapeBase SubExp) NoUniqueness)
I.subExpType [SubExp]
args
let rowtypes :: [TypeBase (ShapeBase SubExp) NoUniqueness]
rowtypes = (TypeBase (ShapeBase SubExp) NoUniqueness
-> TypeBase (ShapeBase SubExp) NoUniqueness)
-> [TypeBase (ShapeBase SubExp) NoUniqueness]
-> [TypeBase (ShapeBase SubExp) NoUniqueness]
forall a b. (a -> b) -> [a] -> [b]
map TypeBase (ShapeBase SubExp) NoUniqueness
-> TypeBase (ShapeBase SubExp) NoUniqueness
forall u.
TypeBase (ShapeBase SubExp) u -> TypeBase (ShapeBase SubExp) u
I.rowType [TypeBase (ShapeBase SubExp) NoUniqueness]
argtypes
([Param (TypeBase (ShapeBase SubExp) NoUniqueness)]
params, Body SOACS
body, [TypeBase (ShapeBase SubExp) NoUniqueness]
_) <- InternaliseLambda
internaliseLambda Exp
lam [TypeBase (ShapeBase SubExp) NoUniqueness]
rowtypes
Body SOACS
body' <-
Scope SOACS
-> InternaliseM (Body SOACS) -> InternaliseM (Body SOACS)
forall rep (m :: * -> *) a.
LocalScope rep m =>
Scope rep -> m a -> m a
localScope ([Param (TypeBase (ShapeBase SubExp) NoUniqueness)] -> Scope SOACS
forall rep dec. (LParamInfo rep ~ dec) => [Param dec] -> Scope rep
scopeOfLParams [Param (TypeBase (ShapeBase SubExp) NoUniqueness)]
params) (InternaliseM (Body SOACS) -> InternaliseM (Body SOACS))
-> InternaliseM (Body SOACS) -> InternaliseM (Body SOACS)
forall a b. (a -> b) -> a -> b
$
Body SOACS -> InternaliseM (Body SOACS)
lambdaWithIncrement Body SOACS
body
Lambda SOACS -> InternaliseM (Lambda SOACS)
forall (f :: * -> *) a. Applicative f => a -> f a
pure (Lambda SOACS -> InternaliseM (Lambda SOACS))
-> Lambda SOACS -> InternaliseM (Lambda SOACS)
forall a b. (a -> b) -> a -> b
$ [LParam SOACS]
-> Body SOACS
-> [TypeBase (ShapeBase SubExp) NoUniqueness]
-> Lambda SOACS
forall rep.
[LParam rep]
-> Body rep
-> [TypeBase (ShapeBase SubExp) NoUniqueness]
-> Lambda rep
I.Lambda [Param (TypeBase (ShapeBase SubExp) NoUniqueness)]
[LParam SOACS]
params Body SOACS
body' [TypeBase (ShapeBase SubExp) NoUniqueness]
forall shape u. [TypeBase shape u]
rettype
where
rettype :: [TypeBase shape u]
rettype = Int -> TypeBase shape u -> [TypeBase shape u]
forall a. Int -> a -> [a]
replicate (Int
k Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
2) (TypeBase shape u -> [TypeBase shape u])
-> TypeBase shape u -> [TypeBase shape u]
forall a b. (a -> b) -> a -> b
$ PrimType -> TypeBase shape u
forall shape u. PrimType -> TypeBase shape u
I.Prim PrimType
int64
result :: Int -> [SubExp]
result Int
i =
(Int64 -> SubExp) -> [Int64] -> [SubExp]
forall a b. (a -> b) -> [a] -> [b]
map Int64 -> SubExp
forall v. IsValue v => v -> SubExp
constant ([Int64] -> [SubExp]) -> [Int64] -> [SubExp]
forall a b. (a -> b) -> a -> b
$
Int -> Int64
forall a b. (Integral a, Num b) => a -> b
fromIntegral Int
i
Int64 -> [Int64] -> [Int64]
forall a. a -> [a] -> [a]
: (Int -> Int64 -> [Int64]
forall a. Int -> a -> [a]
replicate Int
i Int64
0 [Int64] -> [Int64] -> [Int64]
forall a. [a] -> [a] -> [a]
++ [Int64
1 :: Int64] [Int64] -> [Int64] -> [Int64]
forall a. [a] -> [a] -> [a]
++ Int -> Int64 -> [Int64]
forall a. Int -> a -> [a]
replicate (Int
k Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
i) Int64
0)
mkResult :: SubExp -> Int -> f [SubExp]
mkResult SubExp
_ Int
i | Int
i Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
>= Int
k = [SubExp] -> f [SubExp]
forall (f :: * -> *) a. Applicative f => a -> f a
pure ([SubExp] -> f [SubExp]) -> [SubExp] -> f [SubExp]
forall a b. (a -> b) -> a -> b
$ Int -> [SubExp]
result Int
i
mkResult SubExp
eq_class Int
i = do
SubExp
is_i <-
String -> Exp (Rep f) -> f SubExp
forall (m :: * -> *).
MonadBuilder m =>
String -> Exp (Rep m) -> m SubExp
letSubExp String
"is_i" (Exp (Rep f) -> f SubExp) -> Exp (Rep f) -> f SubExp
forall a b. (a -> b) -> a -> b
$
BasicOp -> Exp (Rep f)
forall rep. BasicOp -> Exp rep
BasicOp (BasicOp -> Exp (Rep f)) -> BasicOp -> Exp (Rep f)
forall a b. (a -> b) -> a -> b
$
CmpOp -> SubExp -> SubExp -> BasicOp
CmpOp (PrimType -> CmpOp
CmpEq PrimType
int64) SubExp
eq_class (SubExp -> BasicOp) -> SubExp -> BasicOp
forall a b. (a -> b) -> a -> b
$
IntType -> Integer -> SubExp
intConst IntType
Int64 (Integer -> SubExp) -> Integer -> SubExp
forall a b. (a -> b) -> a -> b
$
Int -> Integer
forall a. Integral a => a -> Integer
toInteger Int
i
String -> Exp (Rep f) -> f [SubExp]
forall (m :: * -> *).
MonadBuilder m =>
String -> Exp (Rep m) -> m [SubExp]
letTupExp' String
"part_res"
(Exp (Rep f) -> f [SubExp]) -> f (Exp (Rep f)) -> f [SubExp]
forall (m :: * -> *) a b. Monad m => (a -> m b) -> m a -> m b
=<< f (Exp (Rep f))
-> f (Body (Rep f)) -> f (Body (Rep f)) -> f (Exp (Rep f))
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
(SubExp -> f (Exp (Rep f))
forall (m :: * -> *). MonadBuilder m => SubExp -> m (Exp (Rep m))
eSubExp SubExp
is_i)
(Body (Rep f) -> f (Body (Rep f))
forall (f :: * -> *) a. Applicative f => a -> f a
pure (Body (Rep f) -> f (Body (Rep f)))
-> Body (Rep f) -> f (Body (Rep f))
forall a b. (a -> b) -> a -> b
$ [SubExp] -> Body (Rep f)
forall rep. Buildable rep => [SubExp] -> Body rep
resultBody ([SubExp] -> Body (Rep f)) -> [SubExp] -> Body (Rep f)
forall a b. (a -> b) -> a -> b
$ Int -> [SubExp]
result Int
i)
([SubExp] -> Body (Rep f)
forall rep. Buildable rep => [SubExp] -> Body rep
resultBody ([SubExp] -> Body (Rep f)) -> f [SubExp] -> f (Body (Rep f))
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> SubExp -> Int -> f [SubExp]
mkResult SubExp
eq_class (Int
i Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
1))
lambdaWithIncrement :: I.Body SOACS -> InternaliseM (I.Body SOACS)
lambdaWithIncrement :: Body SOACS -> InternaliseM (Body SOACS)
lambdaWithIncrement Body SOACS
lam_body = Builder SOACS (Body SOACS) -> InternaliseM (Body SOACS)
forall rep (m :: * -> *) somerep.
(Buildable rep, MonadFreshNames m, HasScope somerep m,
SameScope somerep rep) =>
Builder rep (Body rep) -> m (Body rep)
runBodyBuilder (Builder SOACS (Body SOACS) -> InternaliseM (Body SOACS))
-> Builder SOACS (Body SOACS) -> InternaliseM (Body SOACS)
forall a b. (a -> b) -> a -> b
$ do
SubExp
eq_class <- SubExpRes -> SubExp
resSubExp (SubExpRes -> SubExp) -> (Result -> SubExpRes) -> Result -> SubExp
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Result -> SubExpRes
forall a. [a] -> a
head (Result -> SubExp)
-> BuilderT SOACS (State VNameSource) Result
-> BuilderT SOACS (State VNameSource) SubExp
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Body (Rep (BuilderT SOACS (State VNameSource)))
-> BuilderT SOACS (State VNameSource) Result
forall (m :: * -> *). MonadBuilder m => Body (Rep m) -> m Result
bodyBind Body (Rep (BuilderT SOACS (State VNameSource)))
Body SOACS
lam_body
[SubExp] -> Body SOACS
forall rep. Buildable rep => [SubExp] -> Body rep
resultBody ([SubExp] -> Body SOACS)
-> BuilderT SOACS (State VNameSource) [SubExp]
-> Builder SOACS (Body SOACS)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> SubExp -> Int -> BuilderT SOACS (State VNameSource) [SubExp]
forall (f :: * -> *).
(MonadBuilder f, Buildable (Rep f)) =>
SubExp -> Int -> f [SubExp]
mkResult SubExp
eq_class Int
0