{-# LANGUAGE TypeFamilies #-}
module Futhark.Analysis.HORep.SOAC
(
SOAC (..),
Futhark.ScremaForm (..),
inputs,
setInputs,
lambda,
setLambda,
typeOf,
width,
NotSOAC (..),
fromExp,
toExp,
toSOAC,
Input (..),
varInput,
inputTransforms,
identInput,
isVarInput,
isVarishInput,
addTransform,
addInitialTransforms,
inputArray,
inputRank,
inputType,
inputRowType,
transformRows,
transposeInput,
applyTransforms,
ArrayTransforms,
noTransforms,
nullTransforms,
(|>),
(<|),
viewf,
ViewF (..),
viewl,
ViewL (..),
ArrayTransform (..),
transformFromExp,
soacToStream,
)
where
import Data.Foldable as Foldable
import Data.Maybe
import qualified Data.Sequence as Seq
import Futhark.Construct hiding (toExp)
import Futhark.IR hiding
( Iota,
Rearrange,
Replicate,
Reshape,
typeOf,
)
import qualified Futhark.IR as Futhark
import Futhark.IR.SOACS.SOAC
( HistOp (..),
ScremaForm (..),
StreamForm (..),
StreamOrd (..),
scremaType,
)
import qualified Futhark.IR.SOACS.SOAC as Futhark
import Futhark.Transform.Rename (renameLambda)
import Futhark.Transform.Substitute
import Futhark.Util.Pretty (ppr, text)
import qualified Futhark.Util.Pretty as PP
data ArrayTransform
=
Rearrange Certs [Int]
|
Reshape Certs ReshapeKind Shape
|
ReshapeOuter Certs ReshapeKind Shape
|
ReshapeInner Certs ReshapeKind Shape
|
Replicate Certs Shape
deriving (Int -> ArrayTransform -> ShowS
[ArrayTransform] -> ShowS
ArrayTransform -> String
(Int -> ArrayTransform -> ShowS)
-> (ArrayTransform -> String)
-> ([ArrayTransform] -> ShowS)
-> Show ArrayTransform
forall a.
(Int -> a -> ShowS) -> (a -> String) -> ([a] -> ShowS) -> Show a
showList :: [ArrayTransform] -> ShowS
$cshowList :: [ArrayTransform] -> ShowS
show :: ArrayTransform -> String
$cshow :: ArrayTransform -> String
showsPrec :: Int -> ArrayTransform -> ShowS
$cshowsPrec :: Int -> ArrayTransform -> ShowS
Show, ArrayTransform -> ArrayTransform -> Bool
(ArrayTransform -> ArrayTransform -> Bool)
-> (ArrayTransform -> ArrayTransform -> Bool) -> Eq ArrayTransform
forall a. (a -> a -> Bool) -> (a -> a -> Bool) -> Eq a
/= :: ArrayTransform -> ArrayTransform -> Bool
$c/= :: ArrayTransform -> ArrayTransform -> Bool
== :: ArrayTransform -> ArrayTransform -> Bool
$c== :: ArrayTransform -> ArrayTransform -> Bool
Eq, Eq ArrayTransform
Eq ArrayTransform
-> (ArrayTransform -> ArrayTransform -> Ordering)
-> (ArrayTransform -> ArrayTransform -> Bool)
-> (ArrayTransform -> ArrayTransform -> Bool)
-> (ArrayTransform -> ArrayTransform -> Bool)
-> (ArrayTransform -> ArrayTransform -> Bool)
-> (ArrayTransform -> ArrayTransform -> ArrayTransform)
-> (ArrayTransform -> ArrayTransform -> ArrayTransform)
-> Ord ArrayTransform
ArrayTransform -> ArrayTransform -> Bool
ArrayTransform -> ArrayTransform -> Ordering
ArrayTransform -> ArrayTransform -> ArrayTransform
forall a.
Eq a
-> (a -> a -> Ordering)
-> (a -> a -> Bool)
-> (a -> a -> Bool)
-> (a -> a -> Bool)
-> (a -> a -> Bool)
-> (a -> a -> a)
-> (a -> a -> a)
-> Ord a
min :: ArrayTransform -> ArrayTransform -> ArrayTransform
$cmin :: ArrayTransform -> ArrayTransform -> ArrayTransform
max :: ArrayTransform -> ArrayTransform -> ArrayTransform
$cmax :: ArrayTransform -> ArrayTransform -> ArrayTransform
>= :: ArrayTransform -> ArrayTransform -> Bool
$c>= :: ArrayTransform -> ArrayTransform -> Bool
> :: ArrayTransform -> ArrayTransform -> Bool
$c> :: ArrayTransform -> ArrayTransform -> Bool
<= :: ArrayTransform -> ArrayTransform -> Bool
$c<= :: ArrayTransform -> ArrayTransform -> Bool
< :: ArrayTransform -> ArrayTransform -> Bool
$c< :: ArrayTransform -> ArrayTransform -> Bool
compare :: ArrayTransform -> ArrayTransform -> Ordering
$ccompare :: ArrayTransform -> ArrayTransform -> Ordering
Ord)
instance Substitute ArrayTransform where
substituteNames :: Map VName VName -> ArrayTransform -> ArrayTransform
substituteNames Map VName VName
substs (Rearrange Certs
cs [Int]
xs) =
Certs -> [Int] -> ArrayTransform
Rearrange (Map VName VName -> Certs -> Certs
forall a. Substitute a => Map VName VName -> a -> a
substituteNames Map VName VName
substs Certs
cs) [Int]
xs
substituteNames Map VName VName
substs (Reshape Certs
cs ReshapeKind
k Shape
ses) =
Certs -> ReshapeKind -> Shape -> ArrayTransform
Reshape (Map VName VName -> Certs -> Certs
forall a. Substitute a => Map VName VName -> a -> a
substituteNames Map VName VName
substs Certs
cs) ReshapeKind
k (Map VName VName -> Shape -> Shape
forall a. Substitute a => Map VName VName -> a -> a
substituteNames Map VName VName
substs Shape
ses)
substituteNames Map VName VName
substs (ReshapeOuter Certs
cs ReshapeKind
k Shape
ses) =
Certs -> ReshapeKind -> Shape -> ArrayTransform
ReshapeOuter (Map VName VName -> Certs -> Certs
forall a. Substitute a => Map VName VName -> a -> a
substituteNames Map VName VName
substs Certs
cs) ReshapeKind
k (Map VName VName -> Shape -> Shape
forall a. Substitute a => Map VName VName -> a -> a
substituteNames Map VName VName
substs Shape
ses)
substituteNames Map VName VName
substs (ReshapeInner Certs
cs ReshapeKind
k Shape
ses) =
Certs -> ReshapeKind -> Shape -> ArrayTransform
ReshapeInner (Map VName VName -> Certs -> Certs
forall a. Substitute a => Map VName VName -> a -> a
substituteNames Map VName VName
substs Certs
cs) ReshapeKind
k (Map VName VName -> Shape -> Shape
forall a. Substitute a => Map VName VName -> a -> a
substituteNames Map VName VName
substs Shape
ses)
substituteNames Map VName VName
substs (Replicate Certs
cs Shape
se) =
Certs -> Shape -> ArrayTransform
Replicate (Map VName VName -> Certs -> Certs
forall a. Substitute a => Map VName VName -> a -> a
substituteNames Map VName VName
substs Certs
cs) (Map VName VName -> Shape -> Shape
forall a. Substitute a => Map VName VName -> a -> a
substituteNames Map VName VName
substs Shape
se)
newtype ArrayTransforms = ArrayTransforms (Seq.Seq ArrayTransform)
deriving (ArrayTransforms -> ArrayTransforms -> Bool
(ArrayTransforms -> ArrayTransforms -> Bool)
-> (ArrayTransforms -> ArrayTransforms -> Bool)
-> Eq ArrayTransforms
forall a. (a -> a -> Bool) -> (a -> a -> Bool) -> Eq a
/= :: ArrayTransforms -> ArrayTransforms -> Bool
$c/= :: ArrayTransforms -> ArrayTransforms -> Bool
== :: ArrayTransforms -> ArrayTransforms -> Bool
$c== :: ArrayTransforms -> ArrayTransforms -> Bool
Eq, Eq ArrayTransforms
Eq ArrayTransforms
-> (ArrayTransforms -> ArrayTransforms -> Ordering)
-> (ArrayTransforms -> ArrayTransforms -> Bool)
-> (ArrayTransforms -> ArrayTransforms -> Bool)
-> (ArrayTransforms -> ArrayTransforms -> Bool)
-> (ArrayTransforms -> ArrayTransforms -> Bool)
-> (ArrayTransforms -> ArrayTransforms -> ArrayTransforms)
-> (ArrayTransforms -> ArrayTransforms -> ArrayTransforms)
-> Ord ArrayTransforms
ArrayTransforms -> ArrayTransforms -> Bool
ArrayTransforms -> ArrayTransforms -> Ordering
ArrayTransforms -> ArrayTransforms -> ArrayTransforms
forall a.
Eq a
-> (a -> a -> Ordering)
-> (a -> a -> Bool)
-> (a -> a -> Bool)
-> (a -> a -> Bool)
-> (a -> a -> Bool)
-> (a -> a -> a)
-> (a -> a -> a)
-> Ord a
min :: ArrayTransforms -> ArrayTransforms -> ArrayTransforms
$cmin :: ArrayTransforms -> ArrayTransforms -> ArrayTransforms
max :: ArrayTransforms -> ArrayTransforms -> ArrayTransforms
$cmax :: ArrayTransforms -> ArrayTransforms -> ArrayTransforms
>= :: ArrayTransforms -> ArrayTransforms -> Bool
$c>= :: ArrayTransforms -> ArrayTransforms -> Bool
> :: ArrayTransforms -> ArrayTransforms -> Bool
$c> :: ArrayTransforms -> ArrayTransforms -> Bool
<= :: ArrayTransforms -> ArrayTransforms -> Bool
$c<= :: ArrayTransforms -> ArrayTransforms -> Bool
< :: ArrayTransforms -> ArrayTransforms -> Bool
$c< :: ArrayTransforms -> ArrayTransforms -> Bool
compare :: ArrayTransforms -> ArrayTransforms -> Ordering
$ccompare :: ArrayTransforms -> ArrayTransforms -> Ordering
Ord, Int -> ArrayTransforms -> ShowS
[ArrayTransforms] -> ShowS
ArrayTransforms -> String
(Int -> ArrayTransforms -> ShowS)
-> (ArrayTransforms -> String)
-> ([ArrayTransforms] -> ShowS)
-> Show ArrayTransforms
forall a.
(Int -> a -> ShowS) -> (a -> String) -> ([a] -> ShowS) -> Show a
showList :: [ArrayTransforms] -> ShowS
$cshowList :: [ArrayTransforms] -> ShowS
show :: ArrayTransforms -> String
$cshow :: ArrayTransforms -> String
showsPrec :: Int -> ArrayTransforms -> ShowS
$cshowsPrec :: Int -> ArrayTransforms -> ShowS
Show)
instance Semigroup ArrayTransforms where
ArrayTransforms
ts1 <> :: ArrayTransforms -> ArrayTransforms -> ArrayTransforms
<> ArrayTransforms
ts2 = case ArrayTransforms -> ViewF
viewf ArrayTransforms
ts2 of
ArrayTransform
t :< ArrayTransforms
ts2' -> (ArrayTransforms
ts1 ArrayTransforms -> ArrayTransform -> ArrayTransforms
|> ArrayTransform
t) ArrayTransforms -> ArrayTransforms -> ArrayTransforms
forall a. Semigroup a => a -> a -> a
<> ArrayTransforms
ts2'
ViewF
EmptyF -> ArrayTransforms
ts1
instance Monoid ArrayTransforms where
mempty :: ArrayTransforms
mempty = ArrayTransforms
noTransforms
instance Substitute ArrayTransforms where
substituteNames :: Map VName VName -> ArrayTransforms -> ArrayTransforms
substituteNames Map VName VName
substs (ArrayTransforms Seq ArrayTransform
ts) =
Seq ArrayTransform -> ArrayTransforms
ArrayTransforms (Seq ArrayTransform -> ArrayTransforms)
-> Seq ArrayTransform -> ArrayTransforms
forall a b. (a -> b) -> a -> b
$ Map VName VName -> ArrayTransform -> ArrayTransform
forall a. Substitute a => Map VName VName -> a -> a
substituteNames Map VName VName
substs (ArrayTransform -> ArrayTransform)
-> Seq ArrayTransform -> Seq ArrayTransform
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Seq ArrayTransform
ts
noTransforms :: ArrayTransforms
noTransforms :: ArrayTransforms
noTransforms = Seq ArrayTransform -> ArrayTransforms
ArrayTransforms Seq ArrayTransform
forall a. Seq a
Seq.empty
nullTransforms :: ArrayTransforms -> Bool
nullTransforms :: ArrayTransforms -> Bool
nullTransforms (ArrayTransforms Seq ArrayTransform
s) = Seq ArrayTransform -> Bool
forall a. Seq a -> Bool
Seq.null Seq ArrayTransform
s
viewf :: ArrayTransforms -> ViewF
viewf :: ArrayTransforms -> ViewF
viewf (ArrayTransforms Seq ArrayTransform
s) = case Seq ArrayTransform -> ViewL ArrayTransform
forall a. Seq a -> ViewL a
Seq.viewl Seq ArrayTransform
s of
ArrayTransform
t Seq.:< Seq ArrayTransform
s' -> ArrayTransform
t ArrayTransform -> ArrayTransforms -> ViewF
:< Seq ArrayTransform -> ArrayTransforms
ArrayTransforms Seq ArrayTransform
s'
ViewL ArrayTransform
Seq.EmptyL -> ViewF
EmptyF
data ViewF
= EmptyF
| ArrayTransform :< ArrayTransforms
viewl :: ArrayTransforms -> ViewL
viewl :: ArrayTransforms -> ViewL
viewl (ArrayTransforms Seq ArrayTransform
s) = case Seq ArrayTransform -> ViewR ArrayTransform
forall a. Seq a -> ViewR a
Seq.viewr Seq ArrayTransform
s of
Seq ArrayTransform
s' Seq.:> ArrayTransform
t -> Seq ArrayTransform -> ArrayTransforms
ArrayTransforms Seq ArrayTransform
s' ArrayTransforms -> ArrayTransform -> ViewL
:> ArrayTransform
t
ViewR ArrayTransform
Seq.EmptyR -> ViewL
EmptyL
data ViewL
= EmptyL
| ArrayTransforms :> ArrayTransform
(|>) :: ArrayTransforms -> ArrayTransform -> ArrayTransforms
|> :: ArrayTransforms -> ArrayTransform -> ArrayTransforms
(|>) = (ArrayTransform -> ArrayTransforms -> ArrayTransforms)
-> ArrayTransforms -> ArrayTransform -> ArrayTransforms
forall a b c. (a -> b -> c) -> b -> a -> c
flip ((ArrayTransform -> ArrayTransforms -> ArrayTransforms)
-> ArrayTransforms -> ArrayTransform -> ArrayTransforms)
-> (ArrayTransform -> ArrayTransforms -> ArrayTransforms)
-> ArrayTransforms
-> ArrayTransform
-> ArrayTransforms
forall a b. (a -> b) -> a -> b
$ (ArrayTransforms -> Maybe (ArrayTransform, ArrayTransforms))
-> (ArrayTransform -> ArrayTransforms -> ArrayTransforms)
-> ((ArrayTransform, ArrayTransform)
-> (ArrayTransform, ArrayTransform))
-> ArrayTransform
-> ArrayTransforms
-> ArrayTransforms
addTransform' ArrayTransforms -> Maybe (ArrayTransform, ArrayTransforms)
extract ArrayTransform -> ArrayTransforms -> ArrayTransforms
add (((ArrayTransform, ArrayTransform)
-> (ArrayTransform, ArrayTransform))
-> ArrayTransform -> ArrayTransforms -> ArrayTransforms)
-> ((ArrayTransform, ArrayTransform)
-> (ArrayTransform, ArrayTransform))
-> ArrayTransform
-> ArrayTransforms
-> ArrayTransforms
forall a b. (a -> b) -> a -> b
$ (ArrayTransform
-> ArrayTransform -> (ArrayTransform, ArrayTransform))
-> (ArrayTransform, ArrayTransform)
-> (ArrayTransform, ArrayTransform)
forall a b c. (a -> b -> c) -> (a, b) -> c
uncurry ((ArrayTransform
-> ArrayTransform -> (ArrayTransform, ArrayTransform))
-> ArrayTransform
-> ArrayTransform
-> (ArrayTransform, ArrayTransform)
forall a b c. (a -> b -> c) -> b -> a -> c
flip (,))
where
extract :: ArrayTransforms -> Maybe (ArrayTransform, ArrayTransforms)
extract ArrayTransforms
ts' = case ArrayTransforms -> ViewL
viewl ArrayTransforms
ts' of
ViewL
EmptyL -> Maybe (ArrayTransform, ArrayTransforms)
forall a. Maybe a
Nothing
ArrayTransforms
ts'' :> ArrayTransform
t' -> (ArrayTransform, ArrayTransforms)
-> Maybe (ArrayTransform, ArrayTransforms)
forall a. a -> Maybe a
Just (ArrayTransform
t', ArrayTransforms
ts'')
add :: ArrayTransform -> ArrayTransforms -> ArrayTransforms
add ArrayTransform
t' (ArrayTransforms Seq ArrayTransform
ts') = Seq ArrayTransform -> ArrayTransforms
ArrayTransforms (Seq ArrayTransform -> ArrayTransforms)
-> Seq ArrayTransform -> ArrayTransforms
forall a b. (a -> b) -> a -> b
$ Seq ArrayTransform
ts' Seq ArrayTransform -> ArrayTransform -> Seq ArrayTransform
forall a. Seq a -> a -> Seq a
Seq.|> ArrayTransform
t'
(<|) :: ArrayTransform -> ArrayTransforms -> ArrayTransforms
<| :: ArrayTransform -> ArrayTransforms -> ArrayTransforms
(<|) = (ArrayTransforms -> Maybe (ArrayTransform, ArrayTransforms))
-> (ArrayTransform -> ArrayTransforms -> ArrayTransforms)
-> ((ArrayTransform, ArrayTransform)
-> (ArrayTransform, ArrayTransform))
-> ArrayTransform
-> ArrayTransforms
-> ArrayTransforms
addTransform' ArrayTransforms -> Maybe (ArrayTransform, ArrayTransforms)
extract ArrayTransform -> ArrayTransforms -> ArrayTransforms
add (ArrayTransform, ArrayTransform)
-> (ArrayTransform, ArrayTransform)
forall a. a -> a
id
where
extract :: ArrayTransforms -> Maybe (ArrayTransform, ArrayTransforms)
extract ArrayTransforms
ts' = case ArrayTransforms -> ViewF
viewf ArrayTransforms
ts' of
ViewF
EmptyF -> Maybe (ArrayTransform, ArrayTransforms)
forall a. Maybe a
Nothing
ArrayTransform
t' :< ArrayTransforms
ts'' -> (ArrayTransform, ArrayTransforms)
-> Maybe (ArrayTransform, ArrayTransforms)
forall a. a -> Maybe a
Just (ArrayTransform
t', ArrayTransforms
ts'')
add :: ArrayTransform -> ArrayTransforms -> ArrayTransforms
add ArrayTransform
t' (ArrayTransforms Seq ArrayTransform
ts') = Seq ArrayTransform -> ArrayTransforms
ArrayTransforms (Seq ArrayTransform -> ArrayTransforms)
-> Seq ArrayTransform -> ArrayTransforms
forall a b. (a -> b) -> a -> b
$ ArrayTransform
t' ArrayTransform -> Seq ArrayTransform -> Seq ArrayTransform
forall a. a -> Seq a -> Seq a
Seq.<| Seq ArrayTransform
ts'
addTransform' ::
(ArrayTransforms -> Maybe (ArrayTransform, ArrayTransforms)) ->
(ArrayTransform -> ArrayTransforms -> ArrayTransforms) ->
((ArrayTransform, ArrayTransform) -> (ArrayTransform, ArrayTransform)) ->
ArrayTransform ->
ArrayTransforms ->
ArrayTransforms
addTransform' :: (ArrayTransforms -> Maybe (ArrayTransform, ArrayTransforms))
-> (ArrayTransform -> ArrayTransforms -> ArrayTransforms)
-> ((ArrayTransform, ArrayTransform)
-> (ArrayTransform, ArrayTransform))
-> ArrayTransform
-> ArrayTransforms
-> ArrayTransforms
addTransform' ArrayTransforms -> Maybe (ArrayTransform, ArrayTransforms)
extract ArrayTransform -> ArrayTransforms -> ArrayTransforms
add (ArrayTransform, ArrayTransform)
-> (ArrayTransform, ArrayTransform)
swap ArrayTransform
t ArrayTransforms
ts =
ArrayTransforms -> Maybe ArrayTransforms -> ArrayTransforms
forall a. a -> Maybe a -> a
fromMaybe (ArrayTransform
t ArrayTransform -> ArrayTransforms -> ArrayTransforms
`add` ArrayTransforms
ts) (Maybe ArrayTransforms -> ArrayTransforms)
-> Maybe ArrayTransforms -> ArrayTransforms
forall a b. (a -> b) -> a -> b
$ do
(ArrayTransform
t', ArrayTransforms
ts') <- ArrayTransforms -> Maybe (ArrayTransform, ArrayTransforms)
extract ArrayTransforms
ts
ArrayTransform
combined <- (ArrayTransform -> ArrayTransform -> Maybe ArrayTransform)
-> (ArrayTransform, ArrayTransform) -> Maybe ArrayTransform
forall a b c. (a -> b -> c) -> (a, b) -> c
uncurry ArrayTransform -> ArrayTransform -> Maybe ArrayTransform
combineTransforms ((ArrayTransform, ArrayTransform) -> Maybe ArrayTransform)
-> (ArrayTransform, ArrayTransform) -> Maybe ArrayTransform
forall a b. (a -> b) -> a -> b
$ (ArrayTransform, ArrayTransform)
-> (ArrayTransform, ArrayTransform)
swap (ArrayTransform
t', ArrayTransform
t)
ArrayTransforms -> Maybe ArrayTransforms
forall a. a -> Maybe a
Just (ArrayTransforms -> Maybe ArrayTransforms)
-> ArrayTransforms -> Maybe ArrayTransforms
forall a b. (a -> b) -> a -> b
$
if ArrayTransform -> Bool
identityTransform ArrayTransform
combined
then ArrayTransforms
ts'
else (ArrayTransforms -> Maybe (ArrayTransform, ArrayTransforms))
-> (ArrayTransform -> ArrayTransforms -> ArrayTransforms)
-> ((ArrayTransform, ArrayTransform)
-> (ArrayTransform, ArrayTransform))
-> ArrayTransform
-> ArrayTransforms
-> ArrayTransforms
addTransform' ArrayTransforms -> Maybe (ArrayTransform, ArrayTransforms)
extract ArrayTransform -> ArrayTransforms -> ArrayTransforms
add (ArrayTransform, ArrayTransform)
-> (ArrayTransform, ArrayTransform)
swap ArrayTransform
combined ArrayTransforms
ts'
identityTransform :: ArrayTransform -> Bool
identityTransform :: ArrayTransform -> Bool
identityTransform (Rearrange Certs
_ [Int]
perm) =
[Bool] -> Bool
forall (t :: * -> *). Foldable t => t Bool -> Bool
Foldable.and ([Bool] -> Bool) -> [Bool] -> Bool
forall a b. (a -> b) -> a -> b
$ (Int -> Int -> Bool) -> [Int] -> [Int] -> [Bool]
forall a b c. (a -> b -> c) -> [a] -> [b] -> [c]
zipWith Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
(==) [Int]
perm [Int
0 ..]
identityTransform ArrayTransform
_ = Bool
False
combineTransforms :: ArrayTransform -> ArrayTransform -> Maybe ArrayTransform
combineTransforms :: ArrayTransform -> ArrayTransform -> Maybe ArrayTransform
combineTransforms (Rearrange Certs
cs2 [Int]
perm2) (Rearrange Certs
cs1 [Int]
perm1) =
ArrayTransform -> Maybe ArrayTransform
forall a. a -> Maybe a
Just (ArrayTransform -> Maybe ArrayTransform)
-> ArrayTransform -> Maybe ArrayTransform
forall a b. (a -> b) -> a -> b
$ Certs -> [Int] -> ArrayTransform
Rearrange (Certs
cs1 Certs -> Certs -> Certs
forall a. Semigroup a => a -> a -> a
<> Certs
cs2) ([Int] -> ArrayTransform) -> [Int] -> ArrayTransform
forall a b. (a -> b) -> a -> b
$ [Int]
perm2 [Int] -> [Int] -> [Int]
`rearrangeCompose` [Int]
perm1
combineTransforms ArrayTransform
_ ArrayTransform
_ = Maybe ArrayTransform
forall a. Maybe a
Nothing
transformFromExp :: Certs -> Exp rep -> Maybe (VName, ArrayTransform)
transformFromExp :: forall rep. Certs -> Exp rep -> Maybe (VName, ArrayTransform)
transformFromExp Certs
cs (BasicOp (Futhark.Rearrange [Int]
perm VName
v)) =
(VName, ArrayTransform) -> Maybe (VName, ArrayTransform)
forall a. a -> Maybe a
Just (VName
v, Certs -> [Int] -> ArrayTransform
Rearrange Certs
cs [Int]
perm)
transformFromExp Certs
cs (BasicOp (Futhark.Reshape ReshapeKind
k Shape
shape VName
v)) =
(VName, ArrayTransform) -> Maybe (VName, ArrayTransform)
forall a. a -> Maybe a
Just (VName
v, Certs -> ReshapeKind -> Shape -> ArrayTransform
Reshape Certs
cs ReshapeKind
k Shape
shape)
transformFromExp Certs
cs (BasicOp (Futhark.Replicate Shape
shape (Var VName
v))) =
(VName, ArrayTransform) -> Maybe (VName, ArrayTransform)
forall a. a -> Maybe a
Just (VName
v, Certs -> Shape -> ArrayTransform
Replicate Certs
cs Shape
shape)
transformFromExp Certs
_ Exp rep
_ = Maybe (VName, ArrayTransform)
forall a. Maybe a
Nothing
data Input = Input ArrayTransforms VName Type
deriving (Int -> Input -> ShowS
[Input] -> ShowS
Input -> String
(Int -> Input -> ShowS)
-> (Input -> String) -> ([Input] -> ShowS) -> Show Input
forall a.
(Int -> a -> ShowS) -> (a -> String) -> ([a] -> ShowS) -> Show a
showList :: [Input] -> ShowS
$cshowList :: [Input] -> ShowS
show :: Input -> String
$cshow :: Input -> String
showsPrec :: Int -> Input -> ShowS
$cshowsPrec :: Int -> Input -> ShowS
Show, Input -> Input -> Bool
(Input -> Input -> Bool) -> (Input -> Input -> Bool) -> Eq Input
forall a. (a -> a -> Bool) -> (a -> a -> Bool) -> Eq a
/= :: Input -> Input -> Bool
$c/= :: Input -> Input -> Bool
== :: Input -> Input -> Bool
$c== :: Input -> Input -> Bool
Eq, Eq Input
Eq Input
-> (Input -> Input -> Ordering)
-> (Input -> Input -> Bool)
-> (Input -> Input -> Bool)
-> (Input -> Input -> Bool)
-> (Input -> Input -> Bool)
-> (Input -> Input -> Input)
-> (Input -> Input -> Input)
-> Ord Input
Input -> Input -> Bool
Input -> Input -> Ordering
Input -> Input -> Input
forall a.
Eq a
-> (a -> a -> Ordering)
-> (a -> a -> Bool)
-> (a -> a -> Bool)
-> (a -> a -> Bool)
-> (a -> a -> Bool)
-> (a -> a -> a)
-> (a -> a -> a)
-> Ord a
min :: Input -> Input -> Input
$cmin :: Input -> Input -> Input
max :: Input -> Input -> Input
$cmax :: Input -> Input -> Input
>= :: Input -> Input -> Bool
$c>= :: Input -> Input -> Bool
> :: Input -> Input -> Bool
$c> :: Input -> Input -> Bool
<= :: Input -> Input -> Bool
$c<= :: Input -> Input -> Bool
< :: Input -> Input -> Bool
$c< :: Input -> Input -> Bool
compare :: Input -> Input -> Ordering
$ccompare :: Input -> Input -> Ordering
Ord)
instance Substitute Input where
substituteNames :: Map VName VName -> Input -> Input
substituteNames Map VName VName
substs (Input ArrayTransforms
ts VName
v Type
t) =
ArrayTransforms -> VName -> Type -> Input
Input
(Map VName VName -> ArrayTransforms -> ArrayTransforms
forall a. Substitute a => Map VName VName -> a -> a
substituteNames Map VName VName
substs ArrayTransforms
ts)
(Map VName VName -> VName -> VName
forall a. Substitute a => Map VName VName -> a -> a
substituteNames Map VName VName
substs VName
v)
(Map VName VName -> Type -> Type
forall a. Substitute a => Map VName VName -> a -> a
substituteNames Map VName VName
substs Type
t)
varInput :: HasScope t f => VName -> f Input
varInput :: forall t (f :: * -> *). HasScope t f => VName -> f Input
varInput VName
v = Type -> Input
withType (Type -> Input) -> f Type -> f Input
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> VName -> f Type
forall rep (m :: * -> *). HasScope rep m => VName -> m Type
lookupType VName
v
where
withType :: Type -> Input
withType = ArrayTransforms -> VName -> Type -> Input
Input (Seq ArrayTransform -> ArrayTransforms
ArrayTransforms Seq ArrayTransform
forall a. Seq a
Seq.empty) VName
v
identInput :: Ident -> Input
identInput :: Ident -> Input
identInput Ident
v = ArrayTransforms -> VName -> Type -> Input
Input (Seq ArrayTransform -> ArrayTransforms
ArrayTransforms Seq ArrayTransform
forall a. Seq a
Seq.empty) (Ident -> VName
identName Ident
v) (Ident -> Type
identType Ident
v)
isVarInput :: Input -> Maybe VName
isVarInput :: Input -> Maybe VName
isVarInput (Input ArrayTransforms
ts VName
v Type
_) | ArrayTransforms -> Bool
nullTransforms ArrayTransforms
ts = VName -> Maybe VName
forall a. a -> Maybe a
Just VName
v
isVarInput Input
_ = Maybe VName
forall a. Maybe a
Nothing
isVarishInput :: Input -> Maybe VName
isVarishInput :: Input -> Maybe VName
isVarishInput (Input ArrayTransforms
ts VName
v Type
t)
| ArrayTransforms -> Bool
nullTransforms ArrayTransforms
ts = VName -> Maybe VName
forall a. a -> Maybe a
Just VName
v
| Reshape Certs
cs ReshapeKind
ReshapeCoerce (Shape [SubExp
_]) :< ArrayTransforms
ts' <- ArrayTransforms -> ViewF
viewf ArrayTransforms
ts,
Certs
cs Certs -> Certs -> Bool
forall a. Eq a => a -> a -> Bool
== Certs
forall a. Monoid a => a
mempty =
Input -> Maybe VName
isVarishInput (Input -> Maybe VName) -> Input -> Maybe VName
forall a b. (a -> b) -> a -> b
$ ArrayTransforms -> VName -> Type -> Input
Input ArrayTransforms
ts' VName
v Type
t
isVarishInput Input
_ = Maybe VName
forall a. Maybe a
Nothing
addTransform :: ArrayTransform -> Input -> Input
addTransform :: ArrayTransform -> Input -> Input
addTransform ArrayTransform
tr (Input ArrayTransforms
trs VName
a Type
t) =
ArrayTransforms -> VName -> Type -> Input
Input (ArrayTransforms
trs ArrayTransforms -> ArrayTransform -> ArrayTransforms
|> ArrayTransform
tr) VName
a Type
t
addInitialTransforms :: ArrayTransforms -> Input -> Input
addInitialTransforms :: ArrayTransforms -> Input -> Input
addInitialTransforms ArrayTransforms
ts (Input ArrayTransforms
ots VName
a Type
t) = ArrayTransforms -> VName -> Type -> Input
Input (ArrayTransforms
ts ArrayTransforms -> ArrayTransforms -> ArrayTransforms
forall a. Semigroup a => a -> a -> a
<> ArrayTransforms
ots) VName
a Type
t
applyTransform :: MonadBuilder m => ArrayTransform -> VName -> m VName
applyTransform :: forall (m :: * -> *).
MonadBuilder m =>
ArrayTransform -> VName -> m VName
applyTransform (Replicate Certs
cs Shape
n) VName
ia =
Certs -> m VName -> m VName
forall (m :: * -> *) a. MonadBuilder m => Certs -> m a -> m a
certifying Certs
cs (m VName -> m VName) -> (BasicOp -> m VName) -> BasicOp -> m VName
forall b c a. (b -> c) -> (a -> b) -> a -> c
. String -> Exp (Rep m) -> m VName
forall (m :: * -> *).
MonadBuilder m =>
String -> Exp (Rep m) -> m VName
letExp String
"repeat" (Exp (Rep m) -> m VName)
-> (BasicOp -> Exp (Rep m)) -> BasicOp -> m VName
forall b c a. (b -> c) -> (a -> b) -> a -> c
. BasicOp -> Exp (Rep m)
forall rep. BasicOp -> Exp rep
BasicOp (BasicOp -> m VName) -> BasicOp -> m VName
forall a b. (a -> b) -> a -> b
$
Shape -> SubExp -> BasicOp
Futhark.Replicate Shape
n (VName -> SubExp
Var VName
ia)
applyTransform (Rearrange Certs
cs [Int]
perm) VName
ia = do
Int
r <- Type -> Int
forall shape u. ArrayShape shape => TypeBase shape u -> Int
arrayRank (Type -> Int) -> m Type -> m Int
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> VName -> m Type
forall rep (m :: * -> *). HasScope rep m => VName -> m Type
lookupType VName
ia
Certs -> m VName -> m VName
forall (m :: * -> *) a. MonadBuilder m => Certs -> m a -> m a
certifying Certs
cs (m VName -> m VName) -> (BasicOp -> m VName) -> BasicOp -> m VName
forall b c a. (b -> c) -> (a -> b) -> a -> c
. String -> Exp (Rep m) -> m VName
forall (m :: * -> *).
MonadBuilder m =>
String -> Exp (Rep m) -> m VName
letExp String
"rearrange" (Exp (Rep m) -> m VName)
-> (BasicOp -> Exp (Rep m)) -> BasicOp -> m VName
forall b c a. (b -> c) -> (a -> b) -> a -> c
. BasicOp -> Exp (Rep m)
forall rep. BasicOp -> Exp rep
BasicOp (BasicOp -> m VName) -> BasicOp -> m VName
forall a b. (a -> b) -> a -> b
$
[Int] -> VName -> BasicOp
Futhark.Rearrange ([Int]
perm [Int] -> [Int] -> [Int]
forall a. [a] -> [a] -> [a]
++ [[Int] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length [Int]
perm .. Int
r Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
1]) VName
ia
applyTransform (Reshape Certs
cs ReshapeKind
k Shape
shape) VName
ia =
Certs -> m VName -> m VName
forall (m :: * -> *) a. MonadBuilder m => Certs -> m a -> m a
certifying Certs
cs (m VName -> m VName) -> (BasicOp -> m VName) -> BasicOp -> m VName
forall b c a. (b -> c) -> (a -> b) -> a -> c
. String -> Exp (Rep m) -> m VName
forall (m :: * -> *).
MonadBuilder m =>
String -> Exp (Rep m) -> m VName
letExp String
"reshape" (Exp (Rep m) -> m VName)
-> (BasicOp -> Exp (Rep m)) -> BasicOp -> m VName
forall b c a. (b -> c) -> (a -> b) -> a -> c
. BasicOp -> Exp (Rep m)
forall rep. BasicOp -> Exp rep
BasicOp (BasicOp -> m VName) -> BasicOp -> m VName
forall a b. (a -> b) -> a -> b
$
ReshapeKind -> Shape -> VName -> BasicOp
Futhark.Reshape ReshapeKind
k Shape
shape VName
ia
applyTransform (ReshapeOuter Certs
cs ReshapeKind
k Shape
shape) VName
ia = do
Shape
shape' <- Shape -> Int -> Shape -> Shape
reshapeOuter Shape
shape Int
1 (Shape -> Shape) -> (Type -> Shape) -> Type -> Shape
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Type -> Shape
forall shape u. ArrayShape shape => TypeBase shape u -> shape
arrayShape (Type -> Shape) -> m Type -> m Shape
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> VName -> m Type
forall rep (m :: * -> *). HasScope rep m => VName -> m Type
lookupType VName
ia
Certs -> m VName -> m VName
forall (m :: * -> *) a. MonadBuilder m => Certs -> m a -> m a
certifying Certs
cs (m VName -> m VName) -> (BasicOp -> m VName) -> BasicOp -> m VName
forall b c a. (b -> c) -> (a -> b) -> a -> c
. String -> Exp (Rep m) -> m VName
forall (m :: * -> *).
MonadBuilder m =>
String -> Exp (Rep m) -> m VName
letExp String
"reshape_outer" (Exp (Rep m) -> m VName)
-> (BasicOp -> Exp (Rep m)) -> BasicOp -> m VName
forall b c a. (b -> c) -> (a -> b) -> a -> c
. BasicOp -> Exp (Rep m)
forall rep. BasicOp -> Exp rep
BasicOp (BasicOp -> m VName) -> BasicOp -> m VName
forall a b. (a -> b) -> a -> b
$
ReshapeKind -> Shape -> VName -> BasicOp
Futhark.Reshape ReshapeKind
k Shape
shape' VName
ia
applyTransform (ReshapeInner Certs
cs ReshapeKind
k Shape
shape) VName
ia = do
Shape
shape' <- Shape -> Int -> Shape -> Shape
reshapeInner Shape
shape Int
1 (Shape -> Shape) -> (Type -> Shape) -> Type -> Shape
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Type -> Shape
forall shape u. ArrayShape shape => TypeBase shape u -> shape
arrayShape (Type -> Shape) -> m Type -> m Shape
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> VName -> m Type
forall rep (m :: * -> *). HasScope rep m => VName -> m Type
lookupType VName
ia
Certs -> m VName -> m VName
forall (m :: * -> *) a. MonadBuilder m => Certs -> m a -> m a
certifying Certs
cs (m VName -> m VName) -> (BasicOp -> m VName) -> BasicOp -> m VName
forall b c a. (b -> c) -> (a -> b) -> a -> c
. String -> Exp (Rep m) -> m VName
forall (m :: * -> *).
MonadBuilder m =>
String -> Exp (Rep m) -> m VName
letExp String
"reshape_inner" (Exp (Rep m) -> m VName)
-> (BasicOp -> Exp (Rep m)) -> BasicOp -> m VName
forall b c a. (b -> c) -> (a -> b) -> a -> c
. BasicOp -> Exp (Rep m)
forall rep. BasicOp -> Exp rep
BasicOp (BasicOp -> m VName) -> BasicOp -> m VName
forall a b. (a -> b) -> a -> b
$
ReshapeKind -> Shape -> VName -> BasicOp
Futhark.Reshape ReshapeKind
k Shape
shape' VName
ia
applyTransforms :: MonadBuilder m => ArrayTransforms -> VName -> m VName
applyTransforms :: forall (m :: * -> *).
MonadBuilder m =>
ArrayTransforms -> VName -> m VName
applyTransforms (ArrayTransforms Seq ArrayTransform
ts) VName
a = (VName -> ArrayTransform -> m VName)
-> VName -> Seq ArrayTransform -> m VName
forall (t :: * -> *) (m :: * -> *) b a.
(Foldable t, Monad m) =>
(b -> a -> m b) -> b -> t a -> m b
foldlM ((ArrayTransform -> VName -> m VName)
-> VName -> ArrayTransform -> m VName
forall a b c. (a -> b -> c) -> b -> a -> c
flip ArrayTransform -> VName -> m VName
forall (m :: * -> *).
MonadBuilder m =>
ArrayTransform -> VName -> m VName
applyTransform) VName
a Seq ArrayTransform
ts
inputsToSubExps ::
(MonadBuilder m) =>
[Input] ->
m [VName]
inputsToSubExps :: forall (m :: * -> *). MonadBuilder m => [Input] -> m [VName]
inputsToSubExps = (Input -> m VName) -> [Input] -> m [VName]
forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
(a -> m b) -> t a -> m (t b)
mapM Input -> m VName
forall {m :: * -> *}. MonadBuilder m => Input -> m VName
f
where
f :: Input -> m VName
f (Input ArrayTransforms
ts VName
a Type
_) = ArrayTransforms -> VName -> m VName
forall (m :: * -> *).
MonadBuilder m =>
ArrayTransforms -> VName -> m VName
applyTransforms ArrayTransforms
ts VName
a
inputArray :: Input -> VName
inputArray :: Input -> VName
inputArray (Input ArrayTransforms
_ VName
v Type
_) = VName
v
inputTransforms :: Input -> ArrayTransforms
inputTransforms :: Input -> ArrayTransforms
inputTransforms (Input ArrayTransforms
ts VName
_ Type
_) = ArrayTransforms
ts
inputType :: Input -> Type
inputType :: Input -> Type
inputType (Input (ArrayTransforms Seq ArrayTransform
ts) VName
_ Type
at) =
(Type -> ArrayTransform -> Type)
-> Type -> Seq ArrayTransform -> Type
forall (t :: * -> *) b a.
Foldable t =>
(b -> a -> b) -> b -> t a -> b
Foldable.foldl Type -> ArrayTransform -> Type
transformType Type
at Seq ArrayTransform
ts
where
transformType :: Type -> ArrayTransform -> Type
transformType Type
t (Replicate Certs
_ Shape
shape) =
Type -> Shape -> Type
arrayOfShape Type
t Shape
shape
transformType Type
t (Rearrange Certs
_ [Int]
perm) =
[Int] -> Type -> Type
rearrangeType [Int]
perm Type
t
transformType Type
t (Reshape Certs
_ ReshapeKind
_ Shape
shape) =
Type
t Type -> Shape -> Type
forall newshape oldshape u.
ArrayShape newshape =>
TypeBase oldshape u -> newshape -> TypeBase newshape u
`setArrayShape` Shape
shape
transformType Type
t (ReshapeOuter Certs
_ ReshapeKind
_ Shape
shape) =
let Shape [SubExp]
oldshape = Type -> Shape
forall shape u. ArrayShape shape => TypeBase shape u -> shape
arrayShape Type
t
in Type
t Type -> Shape -> Type
forall newshape oldshape u.
ArrayShape newshape =>
TypeBase oldshape u -> newshape -> TypeBase newshape u
`setArrayShape` [SubExp] -> Shape
forall d. [d] -> ShapeBase d
Shape (Shape -> [SubExp]
forall d. ShapeBase d -> [d]
shapeDims Shape
shape [SubExp] -> [SubExp] -> [SubExp]
forall a. [a] -> [a] -> [a]
++ Int -> [SubExp] -> [SubExp]
forall a. Int -> [a] -> [a]
drop Int
1 [SubExp]
oldshape)
transformType Type
t (ReshapeInner Certs
_ ReshapeKind
_ Shape
shape) =
let Shape [SubExp]
oldshape = Type -> Shape
forall shape u. ArrayShape shape => TypeBase shape u -> shape
arrayShape Type
t
in Type
t Type -> Shape -> Type
forall newshape oldshape u.
ArrayShape newshape =>
TypeBase oldshape u -> newshape -> TypeBase newshape u
`setArrayShape` [SubExp] -> Shape
forall d. [d] -> ShapeBase d
Shape (Int -> [SubExp] -> [SubExp]
forall a. Int -> [a] -> [a]
take Int
1 [SubExp]
oldshape [SubExp] -> [SubExp] -> [SubExp]
forall a. [a] -> [a] -> [a]
++ Shape -> [SubExp]
forall d. ShapeBase d -> [d]
shapeDims Shape
shape)
inputRowType :: Input -> Type
inputRowType :: Input -> Type
inputRowType = Type -> Type
forall u. TypeBase Shape u -> TypeBase Shape u
rowType (Type -> Type) -> (Input -> Type) -> Input -> Type
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Input -> Type
inputType
inputRank :: Input -> Int
inputRank :: Input -> Int
inputRank = Type -> Int
forall shape u. ArrayShape shape => TypeBase shape u -> Int
arrayRank (Type -> Int) -> (Input -> Type) -> Input -> Int
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Input -> Type
inputType
transformRows :: ArrayTransforms -> Input -> Input
transformRows :: ArrayTransforms -> Input -> Input
transformRows (ArrayTransforms Seq ArrayTransform
ts) =
(Input -> Seq ArrayTransform -> Input)
-> Seq ArrayTransform -> Input -> Input
forall a b c. (a -> b -> c) -> b -> a -> c
flip ((Input -> ArrayTransform -> Input)
-> Input -> Seq ArrayTransform -> Input
forall (t :: * -> *) b a.
Foldable t =>
(b -> a -> b) -> b -> t a -> b
Foldable.foldl Input -> ArrayTransform -> Input
transformRows') Seq ArrayTransform
ts
where
transformRows' :: Input -> ArrayTransform -> Input
transformRows' Input
inp (Rearrange Certs
cs [Int]
perm) =
ArrayTransform -> Input -> Input
addTransform (Certs -> [Int] -> ArrayTransform
Rearrange Certs
cs (Int
0 Int -> [Int] -> [Int]
forall a. a -> [a] -> [a]
: (Int -> Int) -> [Int] -> [Int]
forall a b. (a -> b) -> [a] -> [b]
map (Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
1) [Int]
perm)) Input
inp
transformRows' Input
inp (Reshape Certs
cs ReshapeKind
k Shape
shape) =
ArrayTransform -> Input -> Input
addTransform (Certs -> ReshapeKind -> Shape -> ArrayTransform
ReshapeInner Certs
cs ReshapeKind
k Shape
shape) Input
inp
transformRows' Input
inp (Replicate Certs
cs Shape
n)
| Input -> Int
inputRank Input
inp Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
== Int
1 =
Certs -> [Int] -> ArrayTransform
Rearrange Certs
forall a. Monoid a => a
mempty [Int
1, Int
0]
ArrayTransform -> Input -> Input
`addTransform` (Certs -> Shape -> ArrayTransform
Replicate Certs
cs Shape
n ArrayTransform -> Input -> Input
`addTransform` Input
inp)
| Bool
otherwise =
Certs -> [Int] -> ArrayTransform
Rearrange Certs
forall a. Monoid a => a
mempty (Int
2 Int -> [Int] -> [Int]
forall a. a -> [a] -> [a]
: Int
0 Int -> [Int] -> [Int]
forall a. a -> [a] -> [a]
: Int
1 Int -> [Int] -> [Int]
forall a. a -> [a] -> [a]
: [Int
3 .. Input -> Int
inputRank Input
inp])
ArrayTransform -> Input -> Input
`addTransform` ( Certs -> Shape -> ArrayTransform
Replicate Certs
cs Shape
n
ArrayTransform -> Input -> Input
`addTransform` (Certs -> [Int] -> ArrayTransform
Rearrange Certs
forall a. Monoid a => a
mempty (Int
1 Int -> [Int] -> [Int]
forall a. a -> [a] -> [a]
: Int
0 Int -> [Int] -> [Int]
forall a. a -> [a] -> [a]
: [Int
2 .. Input -> Int
inputRank Input
inp Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
1]) ArrayTransform -> Input -> Input
`addTransform` Input
inp)
)
transformRows' Input
inp ArrayTransform
nts =
String -> Input
forall a. HasCallStack => String -> a
error (String -> Input) -> String -> Input
forall a b. (a -> b) -> a -> b
$ String
"transformRows: Cannot transform this yet:\n" String -> ShowS
forall a. [a] -> [a] -> [a]
++ ArrayTransform -> String
forall a. Show a => a -> String
show ArrayTransform
nts String -> ShowS
forall a. [a] -> [a] -> [a]
++ String
"\n" String -> ShowS
forall a. [a] -> [a] -> [a]
++ Input -> String
forall a. Show a => a -> String
show Input
inp
transposeInput :: Int -> Int -> Input -> Input
transposeInput :: Int -> Int -> Input -> Input
transposeInput Int
k Int
n Input
inp =
ArrayTransform -> Input -> Input
addTransform (Certs -> [Int] -> ArrayTransform
Rearrange Certs
forall a. Monoid a => a
mempty ([Int] -> ArrayTransform) -> [Int] -> ArrayTransform
forall a b. (a -> b) -> a -> b
$ Int -> Int -> [Int] -> [Int]
forall a. Int -> Int -> [a] -> [a]
transposeIndex Int
k Int
n [Int
0 .. Input -> Int
inputRank Input
inp Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
1]) Input
inp
data SOAC rep
= Stream SubExp (StreamForm rep) (Lambda rep) [SubExp] [Input]
| Scatter SubExp (Lambda rep) [Input] [(Shape, Int, VName)]
| Screma SubExp (ScremaForm rep) [Input]
| Hist SubExp [HistOp rep] (Lambda rep) [Input]
deriving (SOAC rep -> SOAC rep -> Bool
(SOAC rep -> SOAC rep -> Bool)
-> (SOAC rep -> SOAC rep -> Bool) -> Eq (SOAC rep)
forall rep. RepTypes rep => SOAC rep -> SOAC rep -> Bool
forall a. (a -> a -> Bool) -> (a -> a -> Bool) -> Eq a
/= :: SOAC rep -> SOAC rep -> Bool
$c/= :: forall rep. RepTypes rep => SOAC rep -> SOAC rep -> Bool
== :: SOAC rep -> SOAC rep -> Bool
$c== :: forall rep. RepTypes rep => SOAC rep -> SOAC rep -> Bool
Eq, Int -> SOAC rep -> ShowS
[SOAC rep] -> ShowS
SOAC rep -> String
(Int -> SOAC rep -> ShowS)
-> (SOAC rep -> String) -> ([SOAC rep] -> ShowS) -> Show (SOAC rep)
forall rep. RepTypes rep => Int -> SOAC rep -> ShowS
forall rep. RepTypes rep => [SOAC rep] -> ShowS
forall rep. RepTypes rep => SOAC rep -> String
forall a.
(Int -> a -> ShowS) -> (a -> String) -> ([a] -> ShowS) -> Show a
showList :: [SOAC rep] -> ShowS
$cshowList :: forall rep. RepTypes rep => [SOAC rep] -> ShowS
show :: SOAC rep -> String
$cshow :: forall rep. RepTypes rep => SOAC rep -> String
showsPrec :: Int -> SOAC rep -> ShowS
$cshowsPrec :: forall rep. RepTypes rep => Int -> SOAC rep -> ShowS
Show)
instance PP.Pretty Input where
ppr :: Input -> Doc
ppr (Input (ArrayTransforms Seq ArrayTransform
ts) VName
arr Type
_) = (Doc -> ArrayTransform -> Doc) -> Doc -> Seq ArrayTransform -> Doc
forall (t :: * -> *) b a.
Foldable t =>
(b -> a -> b) -> b -> t a -> b
foldl Doc -> ArrayTransform -> Doc
f (VName -> Doc
forall a. Pretty a => a -> Doc
ppr VName
arr) Seq ArrayTransform
ts
where
f :: Doc -> ArrayTransform -> Doc
f Doc
e (Rearrange Certs
cs [Int]
perm) =
String -> Doc
text String
"rearrange" Doc -> Doc -> Doc
forall a. Semigroup a => a -> a -> a
<> Certs -> Doc
forall a. Pretty a => a -> Doc
ppr Certs
cs Doc -> Doc -> Doc
forall a. Semigroup a => a -> a -> a
<> [Doc] -> Doc
PP.apply [[Doc] -> Doc
PP.apply ((Int -> Doc) -> [Int] -> [Doc]
forall a b. (a -> b) -> [a] -> [b]
map Int -> Doc
forall a. Pretty a => a -> Doc
ppr [Int]
perm), Doc
e]
f Doc
e (Reshape Certs
cs ReshapeKind
ReshapeArbitrary Shape
shape) =
String -> Doc
text String
"reshape" Doc -> Doc -> Doc
forall a. Semigroup a => a -> a -> a
<> Certs -> Doc
forall a. Pretty a => a -> Doc
ppr Certs
cs Doc -> Doc -> Doc
forall a. Semigroup a => a -> a -> a
<> [Doc] -> Doc
PP.apply [Shape -> Doc
forall a. Pretty a => a -> Doc
ppr Shape
shape, Doc
e]
f Doc
e (ReshapeOuter Certs
cs ReshapeKind
ReshapeArbitrary Shape
shape) =
String -> Doc
text String
"reshape_outer" Doc -> Doc -> Doc
forall a. Semigroup a => a -> a -> a
<> Certs -> Doc
forall a. Pretty a => a -> Doc
ppr Certs
cs Doc -> Doc -> Doc
forall a. Semigroup a => a -> a -> a
<> [Doc] -> Doc
PP.apply [Shape -> Doc
forall a. Pretty a => a -> Doc
ppr Shape
shape, Doc
e]
f Doc
e (ReshapeInner Certs
cs ReshapeKind
ReshapeArbitrary Shape
shape) =
String -> Doc
text String
"reshape_inner" Doc -> Doc -> Doc
forall a. Semigroup a => a -> a -> a
<> Certs -> Doc
forall a. Pretty a => a -> Doc
ppr Certs
cs Doc -> Doc -> Doc
forall a. Semigroup a => a -> a -> a
<> [Doc] -> Doc
PP.apply [Shape -> Doc
forall a. Pretty a => a -> Doc
ppr Shape
shape, Doc
e]
f Doc
e (Reshape Certs
cs ReshapeKind
ReshapeCoerce Shape
shape) =
String -> Doc
text String
"coerce" Doc -> Doc -> Doc
forall a. Semigroup a => a -> a -> a
<> Certs -> Doc
forall a. Pretty a => a -> Doc
ppr Certs
cs Doc -> Doc -> Doc
forall a. Semigroup a => a -> a -> a
<> [Doc] -> Doc
PP.apply [Shape -> Doc
forall a. Pretty a => a -> Doc
ppr Shape
shape, Doc
e]
f Doc
e (ReshapeOuter Certs
cs ReshapeKind
ReshapeCoerce Shape
shape) =
String -> Doc
text String
"coerce_outer" Doc -> Doc -> Doc
forall a. Semigroup a => a -> a -> a
<> Certs -> Doc
forall a. Pretty a => a -> Doc
ppr Certs
cs Doc -> Doc -> Doc
forall a. Semigroup a => a -> a -> a
<> [Doc] -> Doc
PP.apply [Shape -> Doc
forall a. Pretty a => a -> Doc
ppr Shape
shape, Doc
e]
f Doc
e (ReshapeInner Certs
cs ReshapeKind
ReshapeCoerce Shape
shape) =
String -> Doc
text String
"coerce_inner" Doc -> Doc -> Doc
forall a. Semigroup a => a -> a -> a
<> Certs -> Doc
forall a. Pretty a => a -> Doc
ppr Certs
cs Doc -> Doc -> Doc
forall a. Semigroup a => a -> a -> a
<> [Doc] -> Doc
PP.apply [Shape -> Doc
forall a. Pretty a => a -> Doc
ppr Shape
shape, Doc
e]
f Doc
e (Replicate Certs
cs Shape
ne) =
String -> Doc
text String
"replicate" Doc -> Doc -> Doc
forall a. Semigroup a => a -> a -> a
<> Certs -> Doc
forall a. Pretty a => a -> Doc
ppr Certs
cs Doc -> Doc -> Doc
forall a. Semigroup a => a -> a -> a
<> [Doc] -> Doc
PP.apply [Shape -> Doc
forall a. Pretty a => a -> Doc
ppr Shape
ne, Doc
e]
instance PrettyRep rep => PP.Pretty (SOAC rep) where
ppr :: SOAC rep -> Doc
ppr (Screma SubExp
w ScremaForm rep
form [Input]
arrs) = SubExp -> [Input] -> ScremaForm rep -> Doc
forall rep inp.
(PrettyRep rep, Pretty inp) =>
SubExp -> [inp] -> ScremaForm rep -> Doc
Futhark.ppScrema SubExp
w [Input]
arrs ScremaForm rep
form
ppr (Hist SubExp
len [HistOp rep]
ops Lambda rep
bucket_fun [Input]
imgs) = SubExp -> [Input] -> [HistOp rep] -> Lambda rep -> Doc
forall rep inp.
(PrettyRep rep, Pretty inp) =>
SubExp -> [inp] -> [HistOp rep] -> Lambda rep -> Doc
Futhark.ppHist SubExp
len [Input]
imgs [HistOp rep]
ops Lambda rep
bucket_fun
ppr (Stream SubExp
w StreamForm rep
form Lambda rep
lam [SubExp]
nes [Input]
arrs) = SubExp
-> [Input] -> StreamForm rep -> [SubExp] -> Lambda rep -> Doc
forall rep inp.
(PrettyRep rep, Pretty inp) =>
SubExp -> [inp] -> StreamForm rep -> [SubExp] -> Lambda rep -> Doc
Futhark.ppStream SubExp
w [Input]
arrs StreamForm rep
form [SubExp]
nes Lambda rep
lam
ppr (Scatter SubExp
w Lambda rep
lam [Input]
arrs [(Shape, Int, VName)]
dests) = SubExp -> [Input] -> Lambda rep -> [(Shape, Int, VName)] -> Doc
forall rep inp.
(PrettyRep rep, Pretty inp) =>
SubExp -> [inp] -> Lambda rep -> [(Shape, Int, VName)] -> Doc
Futhark.ppScatter SubExp
w [Input]
arrs Lambda rep
lam [(Shape, Int, VName)]
dests
inputs :: SOAC rep -> [Input]
inputs :: forall rep. SOAC rep -> [Input]
inputs (Stream SubExp
_ StreamForm rep
_ Lambda rep
_ [SubExp]
_ [Input]
arrs) = [Input]
arrs
inputs (Scatter SubExp
_len Lambda rep
_lam [Input]
ivs [(Shape, Int, VName)]
_as) = [Input]
ivs
inputs (Screma SubExp
_ ScremaForm rep
_ [Input]
arrs) = [Input]
arrs
inputs (Hist SubExp
_ [HistOp rep]
_ Lambda rep
_ [Input]
inps) = [Input]
inps
setInputs :: [Input] -> SOAC rep -> SOAC rep
setInputs :: forall rep. [Input] -> SOAC rep -> SOAC rep
setInputs [Input]
arrs (Stream SubExp
w StreamForm rep
form Lambda rep
lam [SubExp]
nes [Input]
_) =
SubExp
-> StreamForm rep -> Lambda rep -> [SubExp] -> [Input] -> SOAC rep
forall rep.
SubExp
-> StreamForm rep -> Lambda rep -> [SubExp] -> [Input] -> SOAC rep
Stream ([Input] -> SubExp -> SubExp
newWidth [Input]
arrs SubExp
w) StreamForm rep
form Lambda rep
lam [SubExp]
nes [Input]
arrs
setInputs [Input]
arrs (Scatter SubExp
w Lambda rep
lam [Input]
_ivs [(Shape, Int, VName)]
as) =
SubExp
-> Lambda rep -> [Input] -> [(Shape, Int, VName)] -> SOAC rep
forall rep.
SubExp
-> Lambda rep -> [Input] -> [(Shape, Int, VName)] -> SOAC rep
Scatter ([Input] -> SubExp -> SubExp
newWidth [Input]
arrs SubExp
w) Lambda rep
lam [Input]
arrs [(Shape, Int, VName)]
as
setInputs [Input]
arrs (Screma SubExp
w ScremaForm rep
form [Input]
_) =
SubExp -> ScremaForm rep -> [Input] -> SOAC rep
forall rep. SubExp -> ScremaForm rep -> [Input] -> SOAC rep
Screma SubExp
w ScremaForm rep
form [Input]
arrs
setInputs [Input]
inps (Hist SubExp
w [HistOp rep]
ops Lambda rep
lam [Input]
_) =
SubExp -> [HistOp rep] -> Lambda rep -> [Input] -> SOAC rep
forall rep.
SubExp -> [HistOp rep] -> Lambda rep -> [Input] -> SOAC rep
Hist SubExp
w [HistOp rep]
ops Lambda rep
lam [Input]
inps
newWidth :: [Input] -> SubExp -> SubExp
newWidth :: [Input] -> SubExp -> SubExp
newWidth [] SubExp
w = SubExp
w
newWidth (Input
inp : [Input]
_) SubExp
_ = Int -> Type -> SubExp
forall u. Int -> TypeBase Shape u -> SubExp
arraySize Int
0 (Type -> SubExp) -> Type -> SubExp
forall a b. (a -> b) -> a -> b
$ Input -> Type
inputType Input
inp
lambda :: SOAC rep -> Lambda rep
lambda :: forall rep. SOAC rep -> Lambda rep
lambda (Stream SubExp
_ StreamForm rep
_ Lambda rep
lam [SubExp]
_ [Input]
_) = Lambda rep
lam
lambda (Scatter SubExp
_len Lambda rep
lam [Input]
_ivs [(Shape, Int, VName)]
_as) = Lambda rep
lam
lambda (Screma SubExp
_ (ScremaForm [Scan rep]
_ [Reduce rep]
_ Lambda rep
lam) [Input]
_) = Lambda rep
lam
lambda (Hist SubExp
_ [HistOp rep]
_ Lambda rep
lam [Input]
_) = Lambda rep
lam
setLambda :: Lambda rep -> SOAC rep -> SOAC rep
setLambda :: forall rep. Lambda rep -> SOAC rep -> SOAC rep
setLambda Lambda rep
lam (Stream SubExp
w StreamForm rep
form Lambda rep
_ [SubExp]
nes [Input]
arrs) =
SubExp
-> StreamForm rep -> Lambda rep -> [SubExp] -> [Input] -> SOAC rep
forall rep.
SubExp
-> StreamForm rep -> Lambda rep -> [SubExp] -> [Input] -> SOAC rep
Stream SubExp
w StreamForm rep
form Lambda rep
lam [SubExp]
nes [Input]
arrs
setLambda Lambda rep
lam (Scatter SubExp
len Lambda rep
_lam [Input]
ivs [(Shape, Int, VName)]
as) =
SubExp
-> Lambda rep -> [Input] -> [(Shape, Int, VName)] -> SOAC rep
forall rep.
SubExp
-> Lambda rep -> [Input] -> [(Shape, Int, VName)] -> SOAC rep
Scatter SubExp
len Lambda rep
lam [Input]
ivs [(Shape, Int, VName)]
as
setLambda Lambda rep
lam (Screma SubExp
w (ScremaForm [Scan rep]
scan [Reduce rep]
red Lambda rep
_) [Input]
arrs) =
SubExp -> ScremaForm rep -> [Input] -> SOAC rep
forall rep. SubExp -> ScremaForm rep -> [Input] -> SOAC rep
Screma SubExp
w ([Scan rep] -> [Reduce rep] -> Lambda rep -> ScremaForm rep
forall rep.
[Scan rep] -> [Reduce rep] -> Lambda rep -> ScremaForm rep
ScremaForm [Scan rep]
scan [Reduce rep]
red Lambda rep
lam) [Input]
arrs
setLambda Lambda rep
lam (Hist SubExp
w [HistOp rep]
ops Lambda rep
_ [Input]
inps) =
SubExp -> [HistOp rep] -> Lambda rep -> [Input] -> SOAC rep
forall rep.
SubExp -> [HistOp rep] -> Lambda rep -> [Input] -> SOAC rep
Hist SubExp
w [HistOp rep]
ops Lambda rep
lam [Input]
inps
typeOf :: SOAC rep -> [Type]
typeOf :: forall rep. SOAC rep -> [Type]
typeOf (Stream SubExp
w StreamForm rep
_ Lambda rep
lam [SubExp]
nes [Input]
_) =
let accrtps :: [Type]
accrtps = Int -> [Type] -> [Type]
forall a. Int -> [a] -> [a]
take ([SubExp] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length [SubExp]
nes) ([Type] -> [Type]) -> [Type] -> [Type]
forall a b. (a -> b) -> a -> b
$ Lambda rep -> [Type]
forall rep. Lambda rep -> [Type]
lambdaReturnType Lambda rep
lam
arrtps :: [Type]
arrtps =
[ Type -> Shape -> NoUniqueness -> Type
forall shape u_unused u.
ArrayShape shape =>
TypeBase shape u_unused -> shape -> u -> TypeBase shape u
arrayOf (Int -> Type -> Type
forall u. Int -> TypeBase Shape u -> TypeBase Shape u
stripArray Int
1 Type
t) ([SubExp] -> Shape
forall d. [d] -> ShapeBase d
Shape [SubExp
w]) NoUniqueness
NoUniqueness
| Type
t <- Int -> [Type] -> [Type]
forall a. Int -> [a] -> [a]
drop ([SubExp] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length [SubExp]
nes) (Lambda rep -> [Type]
forall rep. Lambda rep -> [Type]
lambdaReturnType Lambda rep
lam)
]
in [Type]
accrtps [Type] -> [Type] -> [Type]
forall a. [a] -> [a] -> [a]
++ [Type]
arrtps
typeOf (Scatter SubExp
_w Lambda rep
lam [Input]
_ivs [(Shape, Int, VName)]
dests) =
(Type -> Shape -> Type) -> [Type] -> [Shape] -> [Type]
forall a b c. (a -> b -> c) -> [a] -> [b] -> [c]
zipWith Type -> Shape -> Type
arrayOfShape [Type]
val_ts [Shape]
ws
where
indexes :: Int
indexes = [Int] -> Int
forall (t :: * -> *) a. (Foldable t, Num a) => t a -> a
sum ([Int] -> Int) -> [Int] -> Int
forall a b. (a -> b) -> a -> b
$ (Int -> Int -> Int) -> [Int] -> [Int] -> [Int]
forall a b c. (a -> b -> c) -> [a] -> [b] -> [c]
zipWith Int -> Int -> Int
forall a. Num a => a -> a -> a
(*) [Int]
ns ([Int] -> [Int]) -> [Int] -> [Int]
forall a b. (a -> b) -> a -> b
$ (Shape -> Int) -> [Shape] -> [Int]
forall a b. (a -> b) -> [a] -> [b]
map Shape -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length [Shape]
ws
val_ts :: [Type]
val_ts = Int -> [Type] -> [Type]
forall a. Int -> [a] -> [a]
drop Int
indexes ([Type] -> [Type]) -> [Type] -> [Type]
forall a b. (a -> b) -> a -> b
$ Lambda rep -> [Type]
forall rep. Lambda rep -> [Type]
lambdaReturnType Lambda rep
lam
([Shape]
ws, [Int]
ns, [VName]
_) = [(Shape, Int, VName)] -> ([Shape], [Int], [VName])
forall a b c. [(a, b, c)] -> ([a], [b], [c])
unzip3 [(Shape, Int, VName)]
dests
typeOf (Screma SubExp
w ScremaForm rep
form [Input]
_) =
SubExp -> ScremaForm rep -> [Type]
forall rep. SubExp -> ScremaForm rep -> [Type]
scremaType SubExp
w ScremaForm rep
form
typeOf (Hist SubExp
_ [HistOp rep]
ops Lambda rep
_ [Input]
_) = do
HistOp rep
op <- [HistOp rep]
ops
(Type -> Type) -> [Type] -> [Type]
forall a b. (a -> b) -> [a] -> [b]
map (Type -> Shape -> Type
`arrayOfShape` HistOp rep -> Shape
forall rep. HistOp rep -> Shape
histShape HistOp rep
op) (Lambda rep -> [Type]
forall rep. Lambda rep -> [Type]
lambdaReturnType (Lambda rep -> [Type]) -> Lambda rep -> [Type]
forall a b. (a -> b) -> a -> b
$ HistOp rep -> Lambda rep
forall rep. HistOp rep -> Lambda rep
histOp HistOp rep
op)
width :: SOAC rep -> SubExp
width :: forall rep. SOAC rep -> SubExp
width (Stream SubExp
w StreamForm rep
_ Lambda rep
_ [SubExp]
_ [Input]
_) = SubExp
w
width (Scatter SubExp
len Lambda rep
_lam [Input]
_ivs [(Shape, Int, VName)]
_as) = SubExp
len
width (Screma SubExp
w ScremaForm rep
_ [Input]
_) = SubExp
w
width (Hist SubExp
w [HistOp rep]
_ Lambda rep
_ [Input]
_) = SubExp
w
toExp ::
(MonadBuilder m, Op (Rep m) ~ Futhark.SOAC (Rep m)) =>
SOAC (Rep m) ->
m (Exp (Rep m))
toExp :: forall (m :: * -> *).
(MonadBuilder m, Op (Rep m) ~ SOAC (Rep m)) =>
SOAC (Rep m) -> m (Exp (Rep m))
toExp SOAC (Rep m)
soac = SOAC (Rep m) -> Exp (Rep m)
forall rep. Op rep -> Exp rep
Op (SOAC (Rep m) -> Exp (Rep m))
-> m (SOAC (Rep m)) -> m (Exp (Rep m))
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> SOAC (Rep m) -> m (SOAC (Rep m))
forall (m :: * -> *).
MonadBuilder m =>
SOAC (Rep m) -> m (SOAC (Rep m))
toSOAC SOAC (Rep m)
soac
toSOAC :: MonadBuilder m => SOAC (Rep m) -> m (Futhark.SOAC (Rep m))
toSOAC :: forall (m :: * -> *).
MonadBuilder m =>
SOAC (Rep m) -> m (SOAC (Rep m))
toSOAC (Stream SubExp
w StreamForm (Rep m)
form Lambda (Rep m)
lam [SubExp]
nes [Input]
inps) =
SubExp
-> [VName]
-> StreamForm (Rep m)
-> [SubExp]
-> Lambda (Rep m)
-> SOAC (Rep m)
forall rep.
SubExp
-> [VName] -> StreamForm rep -> [SubExp] -> Lambda rep -> SOAC rep
Futhark.Stream SubExp
w ([VName]
-> StreamForm (Rep m)
-> [SubExp]
-> Lambda (Rep m)
-> SOAC (Rep m))
-> m [VName]
-> m (StreamForm (Rep m)
-> [SubExp] -> Lambda (Rep m) -> SOAC (Rep m))
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> [Input] -> m [VName]
forall (m :: * -> *). MonadBuilder m => [Input] -> m [VName]
inputsToSubExps [Input]
inps m (StreamForm (Rep m)
-> [SubExp] -> Lambda (Rep m) -> SOAC (Rep m))
-> m (StreamForm (Rep m))
-> m ([SubExp] -> Lambda (Rep m) -> SOAC (Rep m))
forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> StreamForm (Rep m) -> m (StreamForm (Rep m))
forall (f :: * -> *) a. Applicative f => a -> f a
pure StreamForm (Rep m)
form m ([SubExp] -> Lambda (Rep m) -> SOAC (Rep m))
-> m [SubExp] -> m (Lambda (Rep m) -> SOAC (Rep m))
forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> [SubExp] -> m [SubExp]
forall (f :: * -> *) a. Applicative f => a -> f a
pure [SubExp]
nes m (Lambda (Rep m) -> SOAC (Rep m))
-> m (Lambda (Rep m)) -> m (SOAC (Rep m))
forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> Lambda (Rep m) -> m (Lambda (Rep m))
forall (f :: * -> *) a. Applicative f => a -> f a
pure Lambda (Rep m)
lam
toSOAC (Scatter SubExp
w Lambda (Rep m)
lam [Input]
ivs [(Shape, Int, VName)]
dests) =
SubExp
-> [VName]
-> Lambda (Rep m)
-> [(Shape, Int, VName)]
-> SOAC (Rep m)
forall rep.
SubExp
-> [VName] -> Lambda rep -> [(Shape, Int, VName)] -> SOAC rep
Futhark.Scatter SubExp
w ([VName]
-> Lambda (Rep m) -> [(Shape, Int, VName)] -> SOAC (Rep m))
-> m [VName]
-> m (Lambda (Rep m) -> [(Shape, Int, VName)] -> SOAC (Rep m))
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> [Input] -> m [VName]
forall (m :: * -> *). MonadBuilder m => [Input] -> m [VName]
inputsToSubExps [Input]
ivs m (Lambda (Rep m) -> [(Shape, Int, VName)] -> SOAC (Rep m))
-> m (Lambda (Rep m)) -> m ([(Shape, Int, VName)] -> SOAC (Rep m))
forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> Lambda (Rep m) -> m (Lambda (Rep m))
forall (f :: * -> *) a. Applicative f => a -> f a
pure Lambda (Rep m)
lam m ([(Shape, Int, VName)] -> SOAC (Rep m))
-> m [(Shape, Int, VName)] -> m (SOAC (Rep m))
forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> [(Shape, Int, VName)] -> m [(Shape, Int, VName)]
forall (f :: * -> *) a. Applicative f => a -> f a
pure [(Shape, Int, VName)]
dests
toSOAC (Screma SubExp
w ScremaForm (Rep m)
form [Input]
arrs) =
SubExp -> [VName] -> ScremaForm (Rep m) -> SOAC (Rep m)
forall rep. SubExp -> [VName] -> ScremaForm rep -> SOAC rep
Futhark.Screma SubExp
w ([VName] -> ScremaForm (Rep m) -> SOAC (Rep m))
-> m [VName] -> m (ScremaForm (Rep m) -> SOAC (Rep m))
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> [Input] -> m [VName]
forall (m :: * -> *). MonadBuilder m => [Input] -> m [VName]
inputsToSubExps [Input]
arrs m (ScremaForm (Rep m) -> SOAC (Rep m))
-> m (ScremaForm (Rep m)) -> m (SOAC (Rep m))
forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> ScremaForm (Rep m) -> m (ScremaForm (Rep m))
forall (f :: * -> *) a. Applicative f => a -> f a
pure ScremaForm (Rep m)
form
toSOAC (Hist SubExp
w [HistOp (Rep m)]
ops Lambda (Rep m)
lam [Input]
arrs) =
SubExp
-> [VName] -> [HistOp (Rep m)] -> Lambda (Rep m) -> SOAC (Rep m)
forall rep.
SubExp -> [VName] -> [HistOp rep] -> Lambda rep -> SOAC rep
Futhark.Hist SubExp
w ([VName] -> [HistOp (Rep m)] -> Lambda (Rep m) -> SOAC (Rep m))
-> m [VName]
-> m ([HistOp (Rep m)] -> Lambda (Rep m) -> SOAC (Rep m))
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> [Input] -> m [VName]
forall (m :: * -> *). MonadBuilder m => [Input] -> m [VName]
inputsToSubExps [Input]
arrs m ([HistOp (Rep m)] -> Lambda (Rep m) -> SOAC (Rep m))
-> m [HistOp (Rep m)] -> m (Lambda (Rep m) -> SOAC (Rep m))
forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> [HistOp (Rep m)] -> m [HistOp (Rep m)]
forall (f :: * -> *) a. Applicative f => a -> f a
pure [HistOp (Rep m)]
ops m (Lambda (Rep m) -> SOAC (Rep m))
-> m (Lambda (Rep m)) -> m (SOAC (Rep m))
forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> Lambda (Rep m) -> m (Lambda (Rep m))
forall (f :: * -> *) a. Applicative f => a -> f a
pure Lambda (Rep m)
lam
data NotSOAC
=
NotSOAC
deriving (Int -> NotSOAC -> ShowS
[NotSOAC] -> ShowS
NotSOAC -> String
(Int -> NotSOAC -> ShowS)
-> (NotSOAC -> String) -> ([NotSOAC] -> ShowS) -> Show NotSOAC
forall a.
(Int -> a -> ShowS) -> (a -> String) -> ([a] -> ShowS) -> Show a
showList :: [NotSOAC] -> ShowS
$cshowList :: [NotSOAC] -> ShowS
show :: NotSOAC -> String
$cshow :: NotSOAC -> String
showsPrec :: Int -> NotSOAC -> ShowS
$cshowsPrec :: Int -> NotSOAC -> ShowS
Show)
fromExp ::
(Op rep ~ Futhark.SOAC rep, HasScope rep m) =>
Exp rep ->
m (Either NotSOAC (SOAC rep))
fromExp :: forall rep (m :: * -> *).
(Op rep ~ SOAC rep, HasScope rep m) =>
Exp rep -> m (Either NotSOAC (SOAC rep))
fromExp (Op (Futhark.Stream SubExp
w [VName]
as StreamForm rep
form [SubExp]
nes Lambda rep
lam)) =
SOAC rep -> Either NotSOAC (SOAC rep)
forall a b. b -> Either a b
Right (SOAC rep -> Either NotSOAC (SOAC rep))
-> ([Input] -> SOAC rep) -> [Input] -> Either NotSOAC (SOAC rep)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. SubExp
-> StreamForm rep -> Lambda rep -> [SubExp] -> [Input] -> SOAC rep
forall rep.
SubExp
-> StreamForm rep -> Lambda rep -> [SubExp] -> [Input] -> SOAC rep
Stream SubExp
w StreamForm rep
form Lambda rep
lam [SubExp]
nes ([Input] -> Either NotSOAC (SOAC rep))
-> m [Input] -> m (Either NotSOAC (SOAC rep))
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> (VName -> m Input) -> [VName] -> m [Input]
forall (t :: * -> *) (f :: * -> *) a b.
(Traversable t, Applicative f) =>
(a -> f b) -> t a -> f (t b)
traverse VName -> m Input
forall t (f :: * -> *). HasScope t f => VName -> f Input
varInput [VName]
as
fromExp (Op (Futhark.Scatter SubExp
w [VName]
ivs Lambda rep
lam [(Shape, Int, VName)]
as)) =
SOAC rep -> Either NotSOAC (SOAC rep)
forall a b. b -> Either a b
Right (SOAC rep -> Either NotSOAC (SOAC rep))
-> m (SOAC rep) -> m (Either NotSOAC (SOAC rep))
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> (SubExp
-> Lambda rep -> [Input] -> [(Shape, Int, VName)] -> SOAC rep
forall rep.
SubExp
-> Lambda rep -> [Input] -> [(Shape, Int, VName)] -> SOAC rep
Scatter SubExp
w Lambda rep
lam ([Input] -> [(Shape, Int, VName)] -> SOAC rep)
-> m [Input] -> m ([(Shape, Int, VName)] -> SOAC rep)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> (VName -> m Input) -> [VName] -> m [Input]
forall (t :: * -> *) (f :: * -> *) a b.
(Traversable t, Applicative f) =>
(a -> f b) -> t a -> f (t b)
traverse VName -> m Input
forall t (f :: * -> *). HasScope t f => VName -> f Input
varInput [VName]
ivs m ([(Shape, Int, VName)] -> SOAC rep)
-> m [(Shape, Int, VName)] -> m (SOAC rep)
forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> [(Shape, Int, VName)] -> m [(Shape, Int, VName)]
forall (f :: * -> *) a. Applicative f => a -> f a
pure [(Shape, Int, VName)]
as)
fromExp (Op (Futhark.Screma SubExp
w [VName]
arrs ScremaForm rep
form)) =
SOAC rep -> Either NotSOAC (SOAC rep)
forall a b. b -> Either a b
Right (SOAC rep -> Either NotSOAC (SOAC rep))
-> ([Input] -> SOAC rep) -> [Input] -> Either NotSOAC (SOAC rep)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. SubExp -> ScremaForm rep -> [Input] -> SOAC rep
forall rep. SubExp -> ScremaForm rep -> [Input] -> SOAC rep
Screma SubExp
w ScremaForm rep
form ([Input] -> Either NotSOAC (SOAC rep))
-> m [Input] -> m (Either NotSOAC (SOAC rep))
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> (VName -> m Input) -> [VName] -> m [Input]
forall (t :: * -> *) (f :: * -> *) a b.
(Traversable t, Applicative f) =>
(a -> f b) -> t a -> f (t b)
traverse VName -> m Input
forall t (f :: * -> *). HasScope t f => VName -> f Input
varInput [VName]
arrs
fromExp (Op (Futhark.Hist SubExp
w [VName]
arrs [HistOp rep]
ops Lambda rep
lam)) =
SOAC rep -> Either NotSOAC (SOAC rep)
forall a b. b -> Either a b
Right (SOAC rep -> Either NotSOAC (SOAC rep))
-> ([Input] -> SOAC rep) -> [Input] -> Either NotSOAC (SOAC rep)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. SubExp -> [HistOp rep] -> Lambda rep -> [Input] -> SOAC rep
forall rep.
SubExp -> [HistOp rep] -> Lambda rep -> [Input] -> SOAC rep
Hist SubExp
w [HistOp rep]
ops Lambda rep
lam ([Input] -> Either NotSOAC (SOAC rep))
-> m [Input] -> m (Either NotSOAC (SOAC rep))
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> (VName -> m Input) -> [VName] -> m [Input]
forall (t :: * -> *) (f :: * -> *) a b.
(Traversable t, Applicative f) =>
(a -> f b) -> t a -> f (t b)
traverse VName -> m Input
forall t (f :: * -> *). HasScope t f => VName -> f Input
varInput [VName]
arrs
fromExp Exp rep
_ = Either NotSOAC (SOAC rep) -> m (Either NotSOAC (SOAC rep))
forall (f :: * -> *) a. Applicative f => a -> f a
pure (Either NotSOAC (SOAC rep) -> m (Either NotSOAC (SOAC rep)))
-> Either NotSOAC (SOAC rep) -> m (Either NotSOAC (SOAC rep))
forall a b. (a -> b) -> a -> b
$ NotSOAC -> Either NotSOAC (SOAC rep)
forall a b. a -> Either a b
Left NotSOAC
NotSOAC
soacToStream ::
( HasScope rep m,
MonadFreshNames m,
Buildable rep,
BuilderOps rep,
Op rep ~ Futhark.SOAC rep
) =>
SOAC rep ->
m (SOAC rep, [Ident])
soacToStream :: forall rep (m :: * -> *).
(HasScope rep m, MonadFreshNames m, Buildable rep, BuilderOps rep,
Op rep ~ SOAC rep) =>
SOAC rep -> m (SOAC rep, [Ident])
soacToStream SOAC rep
soac = do
Param Type
chunk_param <- String -> Type -> m (Param Type)
forall (m :: * -> *) dec.
MonadFreshNames m =>
String -> dec -> m (Param dec)
newParam String
"chunk" (Type -> m (Param Type)) -> Type -> m (Param Type)
forall a b. (a -> b) -> a -> b
$ PrimType -> Type
forall shape u. PrimType -> TypeBase shape u
Prim PrimType
int64
let chvar :: SubExp
chvar = VName -> SubExp
Var (VName -> SubExp) -> VName -> SubExp
forall a b. (a -> b) -> a -> b
$ Param Type -> VName
forall dec. Param dec -> VName
paramName Param Type
chunk_param
(Lambda rep
lam, [Input]
inps) = (SOAC rep -> Lambda rep
forall rep. SOAC rep -> Lambda rep
lambda SOAC rep
soac, SOAC rep -> [Input]
forall rep. SOAC rep -> [Input]
inputs SOAC rep
soac)
w :: SubExp
w = SOAC rep -> SubExp
forall rep. SOAC rep -> SubExp
width SOAC rep
soac
Lambda rep
lam' <- Lambda rep -> m (Lambda rep)
forall rep (m :: * -> *).
(Renameable rep, MonadFreshNames m) =>
Lambda rep -> m (Lambda rep)
renameLambda Lambda rep
lam
let arrrtps :: [Type]
arrrtps = SubExp -> Lambda rep -> [Type]
forall rep. SubExp -> Lambda rep -> [Type]
mapType SubExp
w Lambda rep
lam
loutps :: [Type]
loutps = [Type -> SubExp -> Type
forall d.
ArrayShape (ShapeBase d) =>
TypeBase (ShapeBase d) NoUniqueness
-> d -> TypeBase (ShapeBase d) NoUniqueness
arrayOfRow Type
t SubExp
chvar | Type
t <- (Type -> Type) -> [Type] -> [Type]
forall a b. (a -> b) -> [a] -> [b]
map Type -> Type
forall u. TypeBase Shape u -> TypeBase Shape u
rowType [Type]
arrrtps]
lintps :: [Type]
lintps = [Type -> SubExp -> Type
forall d.
ArrayShape (ShapeBase d) =>
TypeBase (ShapeBase d) NoUniqueness
-> d -> TypeBase (ShapeBase d) NoUniqueness
arrayOfRow Type
t SubExp
chvar | Type
t <- (Input -> Type) -> [Input] -> [Type]
forall a b. (a -> b) -> [a] -> [b]
map Input -> Type
inputRowType [Input]
inps]
[Param Type]
strm_inpids <- (Type -> m (Param Type)) -> [Type] -> m [Param Type]
forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
(a -> m b) -> t a -> m (t b)
mapM (String -> Type -> m (Param Type)
forall (m :: * -> *) dec.
MonadFreshNames m =>
String -> dec -> m (Param dec)
newParam String
"inp") [Type]
lintps
case SOAC rep
soac of
Screma SubExp
_ ScremaForm rep
form [Input]
_
| Just Lambda rep
_ <- ScremaForm rep -> Maybe (Lambda rep)
forall rep. ScremaForm rep -> Maybe (Lambda rep)
Futhark.isMapSOAC ScremaForm rep
form -> do
[Ident]
strm_resids <- (Type -> m Ident) -> [Type] -> m [Ident]
forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
(a -> m b) -> t a -> m (t b)
mapM (String -> Type -> m Ident
forall (m :: * -> *).
MonadFreshNames m =>
String -> Type -> m Ident
newIdent String
"res") [Type]
loutps
let insoac :: SOAC rep
insoac =
SubExp -> [VName] -> ScremaForm rep -> SOAC rep
forall rep. SubExp -> [VName] -> ScremaForm rep -> SOAC rep
Futhark.Screma SubExp
chvar ((Param Type -> VName) -> [Param Type] -> [VName]
forall a b. (a -> b) -> [a] -> [b]
map Param Type -> VName
forall dec. Param dec -> VName
paramName [Param Type]
strm_inpids) (ScremaForm rep -> SOAC rep) -> ScremaForm rep -> SOAC rep
forall a b. (a -> b) -> a -> b
$
Lambda rep -> ScremaForm rep
forall rep. Lambda rep -> ScremaForm rep
Futhark.mapSOAC Lambda rep
lam'
insstm :: Stm rep
insstm = [Ident] -> Exp rep -> Stm rep
forall rep. Buildable rep => [Ident] -> Exp rep -> Stm rep
mkLet [Ident]
strm_resids (Exp rep -> Stm rep) -> Exp rep -> Stm rep
forall a b. (a -> b) -> a -> b
$ Op rep -> Exp rep
forall rep. Op rep -> Exp rep
Op Op rep
SOAC rep
insoac
strmbdy :: Body rep
strmbdy = Stms rep -> Result -> Body rep
forall rep. Buildable rep => Stms rep -> Result -> Body rep
mkBody (Stm rep -> Stms rep
forall rep. Stm rep -> Stms rep
oneStm Stm rep
insstm) (Result -> Body rep) -> Result -> Body rep
forall a b. (a -> b) -> a -> b
$ (Ident -> SubExpRes) -> [Ident] -> Result
forall a b. (a -> b) -> [a] -> [b]
map (SubExp -> SubExpRes
subExpRes (SubExp -> SubExpRes) -> (Ident -> SubExp) -> Ident -> SubExpRes
forall b c a. (b -> c) -> (a -> b) -> a -> c
. VName -> SubExp
Var (VName -> SubExp) -> (Ident -> VName) -> Ident -> SubExp
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Ident -> VName
identName) [Ident]
strm_resids
strmpar :: [Param Type]
strmpar = Param Type
chunk_param Param Type -> [Param Type] -> [Param Type]
forall a. a -> [a] -> [a]
: [Param Type]
strm_inpids
strmlam :: Lambda rep
strmlam = [LParam rep] -> Body rep -> [Type] -> Lambda rep
forall rep. [LParam rep] -> Body rep -> [Type] -> Lambda rep
Lambda [Param Type]
[LParam rep]
strmpar Body rep
strmbdy [Type]
loutps
empty_lam :: Lambda rep
empty_lam = [LParam rep] -> Body rep -> [Type] -> Lambda rep
forall rep. [LParam rep] -> Body rep -> [Type] -> Lambda rep
Lambda [] (Stms rep -> Result -> Body rep
forall rep. Buildable rep => Stms rep -> Result -> Body rep
mkBody Stms rep
forall a. Monoid a => a
mempty []) []
(SOAC rep, [Ident]) -> m (SOAC rep, [Ident])
forall (f :: * -> *) a. Applicative f => a -> f a
pure (SubExp
-> StreamForm rep -> Lambda rep -> [SubExp] -> [Input] -> SOAC rep
forall rep.
SubExp
-> StreamForm rep -> Lambda rep -> [SubExp] -> [Input] -> SOAC rep
Stream SubExp
w (StreamOrd -> Commutativity -> Lambda rep -> StreamForm rep
forall rep.
StreamOrd -> Commutativity -> Lambda rep -> StreamForm rep
Parallel StreamOrd
Disorder Commutativity
Commutative Lambda rep
empty_lam) Lambda rep
strmlam [] [Input]
inps, [])
| Just ([Scan rep]
scans, Lambda rep
_) <- ScremaForm rep -> Maybe ([Scan rep], Lambda rep)
forall rep. ScremaForm rep -> Maybe ([Scan rep], Lambda rep)
Futhark.isScanomapSOAC ScremaForm rep
form,
Futhark.Scan Lambda rep
scan_lam [SubExp]
nes <- [Scan rep] -> Scan rep
forall rep. Buildable rep => [Scan rep] -> Scan rep
Futhark.singleScan [Scan rep]
scans -> do
let scan_arr_ts :: [Type]
scan_arr_ts = (Type -> Type) -> [Type] -> [Type]
forall a b. (a -> b) -> [a] -> [b]
map (Type -> SubExp -> Type
forall d.
ArrayShape (ShapeBase d) =>
TypeBase (ShapeBase d) NoUniqueness
-> d -> TypeBase (ShapeBase d) NoUniqueness
`arrayOfRow` SubExp
chvar) ([Type] -> [Type]) -> [Type] -> [Type]
forall a b. (a -> b) -> a -> b
$ Lambda rep -> [Type]
forall rep. Lambda rep -> [Type]
lambdaReturnType Lambda rep
scan_lam
accrtps :: [Type]
accrtps = Lambda rep -> [Type]
forall rep. Lambda rep -> [Type]
lambdaReturnType Lambda rep
scan_lam
[Param Type]
inpacc_ids <- (Type -> m (Param Type)) -> [Type] -> m [Param Type]
forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
(a -> m b) -> t a -> m (t b)
mapM (String -> Type -> m (Param Type)
forall (m :: * -> *) dec.
MonadFreshNames m =>
String -> dec -> m (Param dec)
newParam String
"inpacc") [Type]
accrtps
Lambda rep
maplam <- [SubExp] -> Lambda rep -> m (Lambda rep)
forall (m :: * -> *) rep.
(MonadFreshNames m, Buildable rep) =>
[SubExp] -> Lambda rep -> m (Lambda rep)
mkMapPlusAccLam ((Param Type -> SubExp) -> [Param Type] -> [SubExp]
forall a b. (a -> b) -> [a] -> [b]
map (VName -> SubExp
Var (VName -> SubExp) -> (Param Type -> VName) -> Param Type -> SubExp
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Param Type -> VName
forall dec. Param dec -> VName
paramName) [Param Type]
inpacc_ids) Lambda rep
scan_lam
let strmpar :: [Param Type]
strmpar = Param Type
chunk_param Param Type -> [Param Type] -> [Param Type]
forall a. a -> [a] -> [a]
: [Param Type]
inpacc_ids [Param Type] -> [Param Type] -> [Param Type]
forall a. [a] -> [a] -> [a]
++ [Param Type]
strm_inpids
Lambda rep
strmlam <- ((Lambda rep, Stms rep) -> Lambda rep)
-> m (Lambda rep, Stms rep) -> m (Lambda rep)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap (Lambda rep, Stms rep) -> Lambda rep
forall a b. (a, b) -> a
fst (m (Lambda rep, Stms rep) -> m (Lambda rep))
-> (BuilderT rep (State VNameSource) Result
-> m (Lambda rep, Stms rep))
-> BuilderT rep (State VNameSource) Result
-> m (Lambda rep)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Builder rep (Lambda rep) -> m (Lambda rep, Stms rep)
forall (m :: * -> *) somerep rep a.
(MonadFreshNames m, HasScope somerep m, SameScope somerep rep) =>
Builder rep a -> m (a, Stms rep)
runBuilder (Builder rep (Lambda rep) -> m (Lambda rep, Stms rep))
-> (BuilderT rep (State VNameSource) Result
-> Builder rep (Lambda rep))
-> BuilderT rep (State VNameSource) Result
-> m (Lambda rep, Stms rep)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. [LParam (Rep (BuilderT rep (State VNameSource)))]
-> BuilderT rep (State VNameSource) Result
-> BuilderT
rep
(State VNameSource)
(Lambda (Rep (BuilderT rep (State VNameSource))))
forall (m :: * -> *).
MonadBuilder m =>
[LParam (Rep m)] -> m Result -> m (Lambda (Rep m))
mkLambda [Param Type]
[LParam (Rep (BuilderT rep (State VNameSource)))]
strmpar (BuilderT rep (State VNameSource) Result -> m (Lambda rep))
-> BuilderT rep (State VNameSource) Result -> m (Lambda rep)
forall a b. (a -> b) -> a -> b
$ do
([VName]
scan0_ids, [VName]
map_resids) <-
([VName] -> ([VName], [VName]))
-> BuilderT rep (State VNameSource) [VName]
-> BuilderT rep (State VNameSource) ([VName], [VName])
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap (Int -> [VName] -> ([VName], [VName])
forall a. Int -> [a] -> ([a], [a])
splitAt ([Type] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length [Type]
scan_arr_ts)) (BuilderT rep (State VNameSource) [VName]
-> BuilderT rep (State VNameSource) ([VName], [VName]))
-> (SOAC rep -> BuilderT rep (State VNameSource) [VName])
-> SOAC rep
-> BuilderT rep (State VNameSource) ([VName], [VName])
forall b c a. (b -> c) -> (a -> b) -> a -> c
. String
-> Exp (Rep (BuilderT rep (State VNameSource)))
-> BuilderT rep (State VNameSource) [VName]
forall (m :: * -> *).
MonadBuilder m =>
String -> Exp (Rep m) -> m [VName]
letTupExp String
"scan" (Exp rep -> BuilderT rep (State VNameSource) [VName])
-> (SOAC rep -> Exp rep)
-> SOAC rep
-> BuilderT rep (State VNameSource) [VName]
forall b c a. (b -> c) -> (a -> b) -> a -> c
. SOAC rep -> Exp rep
forall rep. Op rep -> Exp rep
Op (SOAC rep -> BuilderT rep (State VNameSource) ([VName], [VName]))
-> SOAC rep -> BuilderT rep (State VNameSource) ([VName], [VName])
forall a b. (a -> b) -> a -> b
$
SubExp -> [VName] -> ScremaForm rep -> SOAC rep
forall rep. SubExp -> [VName] -> ScremaForm rep -> SOAC rep
Futhark.Screma SubExp
chvar ((Param Type -> VName) -> [Param Type] -> [VName]
forall a b. (a -> b) -> [a] -> [b]
map Param Type -> VName
forall dec. Param dec -> VName
paramName [Param Type]
strm_inpids) (ScremaForm rep -> SOAC rep) -> ScremaForm rep -> SOAC rep
forall a b. (a -> b) -> a -> b
$
[Scan rep] -> Lambda rep -> ScremaForm rep
forall rep. [Scan rep] -> Lambda rep -> ScremaForm rep
Futhark.scanomapSOAC [Lambda rep -> [SubExp] -> Scan rep
forall rep. Lambda rep -> [SubExp] -> Scan rep
Futhark.Scan Lambda rep
scan_lam [SubExp]
nes] Lambda rep
lam'
SubExp
outszm1id <-
String
-> Exp (Rep (BuilderT rep (State VNameSource)))
-> BuilderT rep (State VNameSource) SubExp
forall (m :: * -> *).
MonadBuilder m =>
String -> Exp (Rep m) -> m SubExp
letSubExp String
"outszm1" (Exp rep -> BuilderT rep (State VNameSource) SubExp)
-> (BasicOp -> Exp rep)
-> BasicOp
-> BuilderT rep (State VNameSource) SubExp
forall b c a. (b -> c) -> (a -> b) -> a -> c
. BasicOp -> Exp rep
forall rep. BasicOp -> Exp rep
BasicOp (BasicOp -> BuilderT rep (State VNameSource) SubExp)
-> BasicOp -> BuilderT rep (State VNameSource) SubExp
forall a b. (a -> b) -> a -> b
$
BinOp -> SubExp -> SubExp -> BasicOp
BinOp
(IntType -> Overflow -> BinOp
Sub IntType
Int64 Overflow
OverflowUndef)
(VName -> SubExp
Var (VName -> SubExp) -> VName -> SubExp
forall a b. (a -> b) -> a -> b
$ Param Type -> VName
forall dec. Param dec -> VName
paramName Param Type
chunk_param)
(Int64 -> SubExp
forall v. IsValue v => v -> SubExp
constant (Int64
1 :: Int64))
VName
empty_arr <-
String
-> Exp (Rep (BuilderT rep (State VNameSource)))
-> BuilderT rep (State VNameSource) VName
forall (m :: * -> *).
MonadBuilder m =>
String -> Exp (Rep m) -> m VName
letExp String
"empty_arr" (Exp rep -> BuilderT rep (State VNameSource) VName)
-> (BasicOp -> Exp rep)
-> BasicOp
-> BuilderT rep (State VNameSource) VName
forall b c a. (b -> c) -> (a -> b) -> a -> c
. BasicOp -> Exp rep
forall rep. BasicOp -> Exp rep
BasicOp (BasicOp -> BuilderT rep (State VNameSource) VName)
-> BasicOp -> BuilderT rep (State VNameSource) VName
forall a b. (a -> b) -> a -> b
$
CmpOp -> SubExp -> SubExp -> BasicOp
CmpOp
(IntType -> CmpOp
CmpSlt IntType
Int64)
SubExp
outszm1id
(Int64 -> SubExp
forall v. IsValue v => v -> SubExp
constant (Int64
0 :: Int64))
let indexLast :: VName -> BuilderT rep (State VNameSource) (Exp rep)
indexLast VName
arr = do
Type
arr_t <- VName -> BuilderT rep (State VNameSource) Type
forall rep (m :: * -> *). HasScope rep m => VName -> m Type
lookupType VName
arr
Exp rep -> BuilderT rep (State VNameSource) (Exp rep)
forall (f :: * -> *) a. Applicative f => a -> f a
pure (Exp rep -> BuilderT rep (State VNameSource) (Exp rep))
-> Exp rep -> BuilderT rep (State VNameSource) (Exp rep)
forall a b. (a -> b) -> a -> b
$ BasicOp -> Exp rep
forall rep. BasicOp -> Exp rep
BasicOp (BasicOp -> Exp rep)
-> (Slice SubExp -> BasicOp) -> Slice SubExp -> Exp rep
forall b c a. (b -> c) -> (a -> b) -> a -> c
. VName -> Slice SubExp -> BasicOp
Index VName
arr (Slice SubExp -> Exp rep) -> Slice SubExp -> Exp rep
forall a b. (a -> b) -> a -> b
$ Type -> [DimIndex SubExp] -> Slice SubExp
fullSlice Type
arr_t [SubExp -> DimIndex SubExp
forall d. d -> DimIndex d
DimFix SubExp
outszm1id]
[VName]
lastel_ids <-
String
-> Exp (Rep (BuilderT rep (State VNameSource)))
-> BuilderT rep (State VNameSource) [VName]
forall (m :: * -> *).
MonadBuilder m =>
String -> Exp (Rep m) -> m [VName]
letTupExp String
"lastel"
(Exp rep -> BuilderT rep (State VNameSource) [VName])
-> BuilderT rep (State VNameSource) (Exp rep)
-> BuilderT rep (State VNameSource) [VName]
forall (m :: * -> *) a b. Monad m => (a -> m b) -> m a -> m b
=<< BuilderT
rep
(State VNameSource)
(Exp (Rep (BuilderT rep (State VNameSource))))
-> BuilderT
rep
(State VNameSource)
(Body (Rep (BuilderT rep (State VNameSource))))
-> BuilderT
rep
(State VNameSource)
(Body (Rep (BuilderT rep (State VNameSource))))
-> BuilderT
rep
(State VNameSource)
(Exp (Rep (BuilderT rep (State VNameSource))))
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
-> BuilderT
rep
(State VNameSource)
(Exp (Rep (BuilderT rep (State VNameSource))))
forall (m :: * -> *). MonadBuilder m => SubExp -> m (Exp (Rep m))
eSubExp (SubExp
-> BuilderT
rep
(State VNameSource)
(Exp (Rep (BuilderT rep (State VNameSource)))))
-> SubExp
-> BuilderT
rep
(State VNameSource)
(Exp (Rep (BuilderT rep (State VNameSource))))
forall a b. (a -> b) -> a -> b
$ VName -> SubExp
Var VName
empty_arr)
([SubExp]
-> BuilderT
rep
(State VNameSource)
(Body (Rep (BuilderT rep (State VNameSource))))
forall (m :: * -> *).
MonadBuilder m =>
[SubExp] -> m (Body (Rep m))
resultBodyM [SubExp]
nes)
([BuilderT
rep
(State VNameSource)
(Exp (Rep (BuilderT rep (State VNameSource))))]
-> BuilderT
rep
(State VNameSource)
(Body (Rep (BuilderT rep (State VNameSource))))
forall (m :: * -> *).
MonadBuilder m =>
[m (Exp (Rep m))] -> m (Body (Rep m))
eBody ([BuilderT
rep
(State VNameSource)
(Exp (Rep (BuilderT rep (State VNameSource))))]
-> BuilderT
rep
(State VNameSource)
(Body (Rep (BuilderT rep (State VNameSource)))))
-> [BuilderT
rep
(State VNameSource)
(Exp (Rep (BuilderT rep (State VNameSource))))]
-> BuilderT
rep
(State VNameSource)
(Body (Rep (BuilderT rep (State VNameSource))))
forall a b. (a -> b) -> a -> b
$ (VName -> BuilderT rep (State VNameSource) (Exp rep))
-> [VName] -> [BuilderT rep (State VNameSource) (Exp rep)]
forall a b. (a -> b) -> [a] -> [b]
map VName -> BuilderT rep (State VNameSource) (Exp rep)
indexLast [VName]
scan0_ids)
Body rep
addlelbdy <-
Lambda rep
-> [SubExp] -> BuilderT rep (State VNameSource) (Body rep)
forall (m :: * -> *) rep.
(MonadFreshNames m, Buildable rep) =>
Lambda rep -> [SubExp] -> m (Body rep)
mkPlusBnds Lambda rep
scan_lam ([SubExp] -> BuilderT rep (State VNameSource) (Body rep))
-> [SubExp] -> BuilderT rep (State VNameSource) (Body rep)
forall a b. (a -> b) -> a -> b
$ (VName -> SubExp) -> [VName] -> [SubExp]
forall a b. (a -> b) -> [a] -> [b]
map VName -> SubExp
Var ([VName] -> [SubExp]) -> [VName] -> [SubExp]
forall a b. (a -> b) -> a -> b
$ (Param Type -> VName) -> [Param Type] -> [VName]
forall a b. (a -> b) -> [a] -> [b]
map Param Type -> VName
forall dec. Param dec -> VName
paramName [Param Type]
inpacc_ids [VName] -> [VName] -> [VName]
forall a. [a] -> [a] -> [a]
++ [VName]
lastel_ids
let (Stms rep
addlelstm, Result
addlelres) = (Body rep -> Stms rep
forall rep. Body rep -> Stms rep
bodyStms Body rep
addlelbdy, Body rep -> Result
forall rep. Body rep -> Result
bodyResult Body rep
addlelbdy)
[VName]
strm_resids <-
String
-> Exp (Rep (BuilderT rep (State VNameSource)))
-> BuilderT rep (State VNameSource) [VName]
forall (m :: * -> *).
MonadBuilder m =>
String -> Exp (Rep m) -> m [VName]
letTupExp String
"strm_res" (Exp rep -> BuilderT rep (State VNameSource) [VName])
-> (SOAC rep -> Exp rep)
-> SOAC rep
-> BuilderT rep (State VNameSource) [VName]
forall b c a. (b -> c) -> (a -> b) -> a -> c
. SOAC rep -> Exp rep
forall rep. Op rep -> Exp rep
Op (SOAC rep -> BuilderT rep (State VNameSource) [VName])
-> SOAC rep -> BuilderT rep (State VNameSource) [VName]
forall a b. (a -> b) -> a -> b
$
SubExp -> [VName] -> ScremaForm rep -> SOAC rep
forall rep. SubExp -> [VName] -> ScremaForm rep -> SOAC rep
Futhark.Screma SubExp
chvar [VName]
scan0_ids (Lambda rep -> ScremaForm rep
forall rep. Lambda rep -> ScremaForm rep
Futhark.mapSOAC Lambda rep
maplam)
Stms (Rep (BuilderT rep (State VNameSource)))
-> BuilderT rep (State VNameSource) ()
forall (m :: * -> *). MonadBuilder m => Stms (Rep m) -> m ()
addStms Stms rep
Stms (Rep (BuilderT rep (State VNameSource)))
addlelstm
Result -> BuilderT rep (State VNameSource) Result
forall (f :: * -> *) a. Applicative f => a -> f a
pure (Result -> BuilderT rep (State VNameSource) Result)
-> Result -> BuilderT rep (State VNameSource) Result
forall a b. (a -> b) -> a -> b
$ Result
addlelres Result -> Result -> Result
forall a. [a] -> [a] -> [a]
++ (VName -> SubExpRes) -> [VName] -> Result
forall a b. (a -> b) -> [a] -> [b]
map (SubExp -> SubExpRes
subExpRes (SubExp -> SubExpRes) -> (VName -> SubExp) -> VName -> SubExpRes
forall b c a. (b -> c) -> (a -> b) -> a -> c
. VName -> SubExp
Var) ([VName]
strm_resids [VName] -> [VName] -> [VName]
forall a. [a] -> [a] -> [a]
++ [VName]
map_resids)
(SOAC rep, [Ident]) -> m (SOAC rep, [Ident])
forall (f :: * -> *) a. Applicative f => a -> f a
pure
( SubExp
-> StreamForm rep -> Lambda rep -> [SubExp] -> [Input] -> SOAC rep
forall rep.
SubExp
-> StreamForm rep -> Lambda rep -> [SubExp] -> [Input] -> SOAC rep
Stream SubExp
w StreamForm rep
forall rep. StreamForm rep
Sequential Lambda rep
strmlam [SubExp]
nes [Input]
inps,
(Param Type -> Ident) -> [Param Type] -> [Ident]
forall a b. (a -> b) -> [a] -> [b]
map Param Type -> Ident
forall dec. Typed dec => Param dec -> Ident
paramIdent [Param Type]
inpacc_ids
)
| Just ([Reduce rep]
reds, Lambda rep
_) <- ScremaForm rep -> Maybe ([Reduce rep], Lambda rep)
forall rep. ScremaForm rep -> Maybe ([Reduce rep], Lambda rep)
Futhark.isRedomapSOAC ScremaForm rep
form,
Futhark.Reduce Commutativity
comm Lambda rep
lamin [SubExp]
nes <- [Reduce rep] -> Reduce rep
forall rep. Buildable rep => [Reduce rep] -> Reduce rep
Futhark.singleReduce [Reduce rep]
reds -> do
let accrtps :: [Type]
accrtps = Int -> [Type] -> [Type]
forall a. Int -> [a] -> [a]
take ([SubExp] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length [SubExp]
nes) ([Type] -> [Type]) -> [Type] -> [Type]
forall a b. (a -> b) -> a -> b
$ Lambda rep -> [Type]
forall rep. Lambda rep -> [Type]
lambdaReturnType Lambda rep
lam
loutps' :: [Type]
loutps' = Int -> [Type] -> [Type]
forall a. Int -> [a] -> [a]
drop ([SubExp] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length [SubExp]
nes) [Type]
loutps
foldlam :: Lambda rep
foldlam = Lambda rep
lam'
[Ident]
strm_resids <- (Type -> m Ident) -> [Type] -> m [Ident]
forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
(a -> m b) -> t a -> m (t b)
mapM (String -> Type -> m Ident
forall (m :: * -> *).
MonadFreshNames m =>
String -> Type -> m Ident
newIdent String
"res") [Type]
loutps'
[Param Type]
inpacc_ids <- (Type -> m (Param Type)) -> [Type] -> m [Param Type]
forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
(a -> m b) -> t a -> m (t b)
mapM (String -> Type -> m (Param Type)
forall (m :: * -> *) dec.
MonadFreshNames m =>
String -> dec -> m (Param dec)
newParam String
"inpacc") [Type]
accrtps
[Ident]
acc0_ids <- (Type -> m Ident) -> [Type] -> m [Ident]
forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
(a -> m b) -> t a -> m (t b)
mapM (String -> Type -> m Ident
forall (m :: * -> *).
MonadFreshNames m =>
String -> Type -> m Ident
newIdent String
"acc0") [Type]
accrtps
let insoac :: SOAC rep
insoac =
SubExp -> [VName] -> ScremaForm rep -> SOAC rep
forall rep. SubExp -> [VName] -> ScremaForm rep -> SOAC rep
Futhark.Screma
SubExp
chvar
((Param Type -> VName) -> [Param Type] -> [VName]
forall a b. (a -> b) -> [a] -> [b]
map Param Type -> VName
forall dec. Param dec -> VName
paramName [Param Type]
strm_inpids)
(ScremaForm rep -> SOAC rep) -> ScremaForm rep -> SOAC rep
forall a b. (a -> b) -> a -> b
$ [Reduce rep] -> Lambda rep -> ScremaForm rep
forall rep. [Reduce rep] -> Lambda rep -> ScremaForm rep
Futhark.redomapSOAC [Commutativity -> Lambda rep -> [SubExp] -> Reduce rep
forall rep. Commutativity -> Lambda rep -> [SubExp] -> Reduce rep
Futhark.Reduce Commutativity
comm Lambda rep
lamin [SubExp]
nes] Lambda rep
foldlam
insstm :: Stm rep
insstm = [Ident] -> Exp rep -> Stm rep
forall rep. Buildable rep => [Ident] -> Exp rep -> Stm rep
mkLet ([Ident]
acc0_ids [Ident] -> [Ident] -> [Ident]
forall a. [a] -> [a] -> [a]
++ [Ident]
strm_resids) (Exp rep -> Stm rep) -> Exp rep -> Stm rep
forall a b. (a -> b) -> a -> b
$ Op rep -> Exp rep
forall rep. Op rep -> Exp rep
Op Op rep
SOAC rep
insoac
Body rep
addaccbdy <-
Lambda rep -> [SubExp] -> m (Body rep)
forall (m :: * -> *) rep.
(MonadFreshNames m, Buildable rep) =>
Lambda rep -> [SubExp] -> m (Body rep)
mkPlusBnds Lambda rep
lamin ([SubExp] -> m (Body rep)) -> [SubExp] -> m (Body rep)
forall a b. (a -> b) -> a -> b
$
(VName -> SubExp) -> [VName] -> [SubExp]
forall a b. (a -> b) -> [a] -> [b]
map VName -> SubExp
Var ([VName] -> [SubExp]) -> [VName] -> [SubExp]
forall a b. (a -> b) -> a -> b
$
(Param Type -> VName) -> [Param Type] -> [VName]
forall a b. (a -> b) -> [a] -> [b]
map Param Type -> VName
forall dec. Param dec -> VName
paramName [Param Type]
inpacc_ids [VName] -> [VName] -> [VName]
forall a. [a] -> [a] -> [a]
++ (Ident -> VName) -> [Ident] -> [VName]
forall a b. (a -> b) -> [a] -> [b]
map Ident -> VName
identName [Ident]
acc0_ids
let (Stms rep
addaccstm, Result
addaccres) = (Body rep -> Stms rep
forall rep. Body rep -> Stms rep
bodyStms Body rep
addaccbdy, Body rep -> Result
forall rep. Body rep -> Result
bodyResult Body rep
addaccbdy)
strmbdy :: Body rep
strmbdy =
Stms rep -> Result -> Body rep
forall rep. Buildable rep => Stms rep -> Result -> Body rep
mkBody (Stm rep -> Stms rep
forall rep. Stm rep -> Stms rep
oneStm Stm rep
insstm Stms rep -> Stms rep -> Stms rep
forall a. Semigroup a => a -> a -> a
<> Stms rep
addaccstm) (Result -> Body rep) -> Result -> Body rep
forall a b. (a -> b) -> a -> b
$
Result
addaccres Result -> Result -> Result
forall a. [a] -> [a] -> [a]
++ (Ident -> SubExpRes) -> [Ident] -> Result
forall a b. (a -> b) -> [a] -> [b]
map (SubExp -> SubExpRes
subExpRes (SubExp -> SubExpRes) -> (Ident -> SubExp) -> Ident -> SubExpRes
forall b c a. (b -> c) -> (a -> b) -> a -> c
. VName -> SubExp
Var (VName -> SubExp) -> (Ident -> VName) -> Ident -> SubExp
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Ident -> VName
identName) [Ident]
strm_resids
strmpar :: [Param Type]
strmpar = Param Type
chunk_param Param Type -> [Param Type] -> [Param Type]
forall a. a -> [a] -> [a]
: [Param Type]
inpacc_ids [Param Type] -> [Param Type] -> [Param Type]
forall a. [a] -> [a] -> [a]
++ [Param Type]
strm_inpids
strmlam :: Lambda rep
strmlam = [LParam rep] -> Body rep -> [Type] -> Lambda rep
forall rep. [LParam rep] -> Body rep -> [Type] -> Lambda rep
Lambda [Param Type]
[LParam rep]
strmpar Body rep
strmbdy ([Type]
accrtps [Type] -> [Type] -> [Type]
forall a. [a] -> [a] -> [a]
++ [Type]
loutps')
Lambda rep
lam0 <- Lambda rep -> m (Lambda rep)
forall rep (m :: * -> *).
(Renameable rep, MonadFreshNames m) =>
Lambda rep -> m (Lambda rep)
renameLambda Lambda rep
lamin
(SOAC rep, [Ident]) -> m (SOAC rep, [Ident])
forall (f :: * -> *) a. Applicative f => a -> f a
pure (SubExp
-> StreamForm rep -> Lambda rep -> [SubExp] -> [Input] -> SOAC rep
forall rep.
SubExp
-> StreamForm rep -> Lambda rep -> [SubExp] -> [Input] -> SOAC rep
Stream SubExp
w (StreamOrd -> Commutativity -> Lambda rep -> StreamForm rep
forall rep.
StreamOrd -> Commutativity -> Lambda rep -> StreamForm rep
Parallel StreamOrd
InOrder Commutativity
comm Lambda rep
lam0) Lambda rep
strmlam [SubExp]
nes [Input]
inps, [])
SOAC rep
_ -> (SOAC rep, [Ident]) -> m (SOAC rep, [Ident])
forall (f :: * -> *) a. Applicative f => a -> f a
pure (SOAC rep
soac, [])
where
mkMapPlusAccLam ::
(MonadFreshNames m, Buildable rep) =>
[SubExp] ->
Lambda rep ->
m (Lambda rep)
mkMapPlusAccLam :: forall (m :: * -> *) rep.
(MonadFreshNames m, Buildable rep) =>
[SubExp] -> Lambda rep -> m (Lambda rep)
mkMapPlusAccLam [SubExp]
accs Lambda rep
plus = do
let ([Param Type]
accpars, [Param Type]
rempars) = Int -> [Param Type] -> ([Param Type], [Param Type])
forall a. Int -> [a] -> ([a], [a])
splitAt ([SubExp] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length [SubExp]
accs) ([Param Type] -> ([Param Type], [Param Type]))
-> [Param Type] -> ([Param Type], [Param Type])
forall a b. (a -> b) -> a -> b
$ Lambda rep -> [LParam rep]
forall rep. Lambda rep -> [LParam rep]
lambdaParams Lambda rep
plus
parstms :: [Stm rep]
parstms =
(Param Type -> SubExp -> Stm rep)
-> [Param Type] -> [SubExp] -> [Stm rep]
forall a b c. (a -> b -> c) -> [a] -> [b] -> [c]
zipWith
(\Param Type
par SubExp
se -> [Ident] -> Exp rep -> Stm rep
forall rep. Buildable rep => [Ident] -> Exp rep -> Stm rep
mkLet [Param Type -> Ident
forall dec. Typed dec => Param dec -> Ident
paramIdent Param Type
par] (BasicOp -> Exp rep
forall rep. BasicOp -> Exp rep
BasicOp (BasicOp -> Exp rep) -> BasicOp -> Exp rep
forall a b. (a -> b) -> a -> b
$ SubExp -> BasicOp
SubExp SubExp
se))
[Param Type]
accpars
[SubExp]
accs
plus_bdy :: Body rep
plus_bdy = Lambda rep -> Body rep
forall rep. Lambda rep -> Body rep
lambdaBody Lambda rep
plus
newlambdy :: Body rep
newlambdy =
BodyDec rep -> Stms rep -> Result -> Body rep
forall rep. BodyDec rep -> Stms rep -> Result -> Body rep
Body
(Body rep -> BodyDec rep
forall rep. Body rep -> BodyDec rep
bodyDec Body rep
plus_bdy)
([Stm rep] -> Stms rep
forall rep. [Stm rep] -> Stms rep
stmsFromList [Stm rep]
parstms Stms rep -> Stms rep -> Stms rep
forall a. Semigroup a => a -> a -> a
<> Body rep -> Stms rep
forall rep. Body rep -> Stms rep
bodyStms Body rep
plus_bdy)
(Body rep -> Result
forall rep. Body rep -> Result
bodyResult Body rep
plus_bdy)
Lambda rep -> m (Lambda rep)
forall rep (m :: * -> *).
(Renameable rep, MonadFreshNames m) =>
Lambda rep -> m (Lambda rep)
renameLambda (Lambda rep -> m (Lambda rep)) -> Lambda rep -> m (Lambda rep)
forall a b. (a -> b) -> a -> b
$ [LParam rep] -> Body rep -> [Type] -> Lambda rep
forall rep. [LParam rep] -> Body rep -> [Type] -> Lambda rep
Lambda [Param Type]
[LParam rep]
rempars Body rep
newlambdy ([Type] -> Lambda rep) -> [Type] -> Lambda rep
forall a b. (a -> b) -> a -> b
$ Lambda rep -> [Type]
forall rep. Lambda rep -> [Type]
lambdaReturnType Lambda rep
plus
mkPlusBnds ::
(MonadFreshNames m, Buildable rep) =>
Lambda rep ->
[SubExp] ->
m (Body rep)
mkPlusBnds :: forall (m :: * -> *) rep.
(MonadFreshNames m, Buildable rep) =>
Lambda rep -> [SubExp] -> m (Body rep)
mkPlusBnds Lambda rep
plus [SubExp]
accels = do
Lambda rep
plus' <- Lambda rep -> m (Lambda rep)
forall rep (m :: * -> *).
(Renameable rep, MonadFreshNames m) =>
Lambda rep -> m (Lambda rep)
renameLambda Lambda rep
plus
let parstms :: [Stm rep]
parstms =
(Param Type -> SubExp -> Stm rep)
-> [Param Type] -> [SubExp] -> [Stm rep]
forall a b c. (a -> b -> c) -> [a] -> [b] -> [c]
zipWith
(\Param Type
par SubExp
se -> [Ident] -> Exp rep -> Stm rep
forall rep. Buildable rep => [Ident] -> Exp rep -> Stm rep
mkLet [Param Type -> Ident
forall dec. Typed dec => Param dec -> Ident
paramIdent Param Type
par] (BasicOp -> Exp rep
forall rep. BasicOp -> Exp rep
BasicOp (BasicOp -> Exp rep) -> BasicOp -> Exp rep
forall a b. (a -> b) -> a -> b
$ SubExp -> BasicOp
SubExp SubExp
se))
(Lambda rep -> [LParam rep]
forall rep. Lambda rep -> [LParam rep]
lambdaParams Lambda rep
plus')
[SubExp]
accels
body :: Body rep
body = Lambda rep -> Body rep
forall rep. Lambda rep -> Body rep
lambdaBody Lambda rep
plus'
Body rep -> m (Body rep)
forall (f :: * -> *) a. Applicative f => a -> f a
pure (Body rep -> m (Body rep)) -> Body rep -> m (Body rep)
forall a b. (a -> b) -> a -> b
$ Body rep
body {bodyStms :: Stms rep
bodyStms = [Stm rep] -> Stms rep
forall rep. [Stm rep] -> Stms rep
stmsFromList [Stm rep]
parstms Stms rep -> Stms rep -> Stms rep
forall a. Semigroup a => a -> a -> a
<> Body rep -> Stms rep
forall rep. Body rep -> Stms rep
bodyStms Body rep
body}