{-# LANGUAGE FlexibleInstances #-}
{-# LANGUAGE TypeFamilies #-}
module Futhark.Analysis.HORep.SOAC
(
SOAC (..),
Futhark.ScremaForm (..),
inputs,
setInputs,
lambda,
setLambda,
typeOf,
width,
NotSOAC (..),
fromExp,
toExp,
toSOAC,
Input (..),
varInput,
identInput,
isVarInput,
isVarishInput,
addTransform,
addInitialTransforms,
inputArray,
inputRank,
inputType,
inputRowType,
transformRows,
transposeInput,
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,
Var,
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 (ShapeChange SubExp)
|
ReshapeOuter Certs (ShapeChange SubExp)
|
ReshapeInner Certs (ShapeChange SubExp)
|
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
$cp1Ord :: Eq ArrayTransform
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 ShapeChange SubExp
ses) =
Certs -> ShapeChange SubExp -> ArrayTransform
Reshape (Map VName VName -> Certs -> Certs
forall a. Substitute a => Map VName VName -> a -> a
substituteNames Map VName VName
substs Certs
cs) (Map VName VName -> ShapeChange SubExp -> ShapeChange SubExp
forall a. Substitute a => Map VName VName -> a -> a
substituteNames Map VName VName
substs ShapeChange SubExp
ses)
substituteNames Map VName VName
substs (ReshapeOuter Certs
cs ShapeChange SubExp
ses) =
Certs -> ShapeChange SubExp -> ArrayTransform
ReshapeOuter (Map VName VName -> Certs -> Certs
forall a. Substitute a => Map VName VName -> a -> a
substituteNames Map VName VName
substs Certs
cs) (Map VName VName -> ShapeChange SubExp -> ShapeChange SubExp
forall a. Substitute a => Map VName VName -> a -> a
substituteNames Map VName VName
substs ShapeChange SubExp
ses)
substituteNames Map VName VName
substs (ReshapeInner Certs
cs ShapeChange SubExp
ses) =
Certs -> ShapeChange SubExp -> ArrayTransform
ReshapeInner (Map VName VName -> Certs -> Certs
forall a. Substitute a => Map VName VName -> a -> a
substituteNames Map VName VName
substs Certs
cs) (Map VName VName -> ShapeChange SubExp -> ShapeChange SubExp
forall a. Substitute a => Map VName VName -> a -> a
substituteNames Map VName VName
substs ShapeChange SubExp
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
$cp1Ord :: Eq ArrayTransforms
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 :: 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 ShapeChange SubExp
shape VName
v)) =
(VName, ArrayTransform) -> Maybe (VName, ArrayTransform)
forall a. a -> Maybe a
Just (VName
v, Certs -> ShapeChange SubExp -> ArrayTransform
Reshape Certs
cs ShapeChange SubExp
shape)
transformFromExp Certs
cs (BasicOp (Futhark.Replicate Shape
shape (Futhark.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
$cp1Ord :: Eq Input
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 :: 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 [DimCoercion 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
inputsToSubExps ::
(MonadBuilder m) =>
[Input] ->
m [VName]
inputsToSubExps :: [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
inputToExp'
where
inputToExp' :: Input -> m VName
inputToExp' (Input (ArrayTransforms Seq ArrayTransform
ts) VName
a Type
_) =
(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 VName -> ArrayTransform -> m VName
forall (m :: * -> *).
MonadBuilder m =>
VName -> ArrayTransform -> m VName
transform VName
a Seq ArrayTransform
ts
transform :: VName -> ArrayTransform -> m VName
transform VName
ia (Replicate Certs
cs Shape
n) =
Certs -> m VName -> m VName
forall (m :: * -> *) a. MonadBuilder m => Certs -> m a -> m a
certifying Certs
cs (m VName -> m VName) -> m VName -> m VName
forall a b. (a -> b) -> a -> b
$
String -> Exp (Rep m) -> m VName
forall (m :: * -> *).
MonadBuilder m =>
String -> Exp (Rep m) -> m VName
letExp String
"repeat" (Exp (Rep m) -> m VName) -> Exp (Rep m) -> m VName
forall a b. (a -> b) -> a -> b
$ BasicOp -> Exp (Rep m)
forall rep. BasicOp -> ExpT rep
BasicOp (BasicOp -> Exp (Rep m)) -> BasicOp -> Exp (Rep m)
forall a b. (a -> b) -> a -> b
$ Shape -> SubExp -> BasicOp
Futhark.Replicate Shape
n (VName -> SubExp
Futhark.Var VName
ia)
transform VName
ia (Rearrange Certs
cs [Int]
perm) =
Certs -> m VName -> m VName
forall (m :: * -> *) a. MonadBuilder m => Certs -> m a -> m a
certifying Certs
cs (m VName -> m VName) -> m VName -> m VName
forall a b. (a -> b) -> a -> b
$
String -> Exp (Rep m) -> m VName
forall (m :: * -> *).
MonadBuilder m =>
String -> Exp (Rep m) -> m VName
letExp String
"rearrange" (Exp (Rep m) -> m VName) -> Exp (Rep m) -> m VName
forall a b. (a -> b) -> a -> b
$ BasicOp -> Exp (Rep m)
forall rep. BasicOp -> ExpT rep
BasicOp (BasicOp -> Exp (Rep m)) -> BasicOp -> Exp (Rep m)
forall a b. (a -> b) -> a -> b
$ [Int] -> VName -> BasicOp
Futhark.Rearrange [Int]
perm VName
ia
transform VName
ia (Reshape Certs
cs ShapeChange SubExp
shape) =
Certs -> m VName -> m VName
forall (m :: * -> *) a. MonadBuilder m => Certs -> m a -> m a
certifying Certs
cs (m VName -> m VName) -> m VName -> m VName
forall a b. (a -> b) -> a -> b
$
String -> Exp (Rep m) -> m VName
forall (m :: * -> *).
MonadBuilder m =>
String -> Exp (Rep m) -> m VName
letExp String
"reshape" (Exp (Rep m) -> m VName) -> Exp (Rep m) -> m VName
forall a b. (a -> b) -> a -> b
$ BasicOp -> Exp (Rep m)
forall rep. BasicOp -> ExpT rep
BasicOp (BasicOp -> Exp (Rep m)) -> BasicOp -> Exp (Rep m)
forall a b. (a -> b) -> a -> b
$ ShapeChange SubExp -> VName -> BasicOp
Futhark.Reshape ShapeChange SubExp
shape VName
ia
transform VName
ia (ReshapeOuter Certs
cs ShapeChange SubExp
shape) = do
ShapeChange SubExp
shape' <- ShapeChange SubExp -> Int -> Shape -> ShapeChange SubExp
reshapeOuter ShapeChange SubExp
shape Int
1 (Shape -> ShapeChange SubExp)
-> (Type -> Shape) -> Type -> ShapeChange SubExp
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Type -> Shape
forall shape u. ArrayShape shape => TypeBase shape u -> shape
arrayShape (Type -> ShapeChange SubExp) -> m Type -> m (ShapeChange SubExp)
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) -> m VName -> m VName
forall a b. (a -> b) -> a -> b
$
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) -> Exp (Rep m) -> m VName
forall a b. (a -> b) -> a -> b
$ BasicOp -> Exp (Rep m)
forall rep. BasicOp -> ExpT rep
BasicOp (BasicOp -> Exp (Rep m)) -> BasicOp -> Exp (Rep m)
forall a b. (a -> b) -> a -> b
$ ShapeChange SubExp -> VName -> BasicOp
Futhark.Reshape ShapeChange SubExp
shape' VName
ia
transform VName
ia (ReshapeInner Certs
cs ShapeChange SubExp
shape) = do
ShapeChange SubExp
shape' <- ShapeChange SubExp -> Int -> Shape -> ShapeChange SubExp
reshapeInner ShapeChange SubExp
shape Int
1 (Shape -> ShapeChange SubExp)
-> (Type -> Shape) -> Type -> ShapeChange SubExp
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Type -> Shape
forall shape u. ArrayShape shape => TypeBase shape u -> shape
arrayShape (Type -> ShapeChange SubExp) -> m Type -> m (ShapeChange SubExp)
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) -> m VName -> m VName
forall a b. (a -> b) -> a -> b
$
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) -> Exp (Rep m) -> m VName
forall a b. (a -> b) -> a -> b
$ BasicOp -> Exp (Rep m)
forall rep. BasicOp -> ExpT rep
BasicOp (BasicOp -> Exp (Rep m)) -> BasicOp -> Exp (Rep m)
forall a b. (a -> b) -> a -> b
$ ShapeChange SubExp -> VName -> BasicOp
Futhark.Reshape ShapeChange SubExp
shape' VName
ia
inputArray :: Input -> VName
inputArray :: Input -> VName
inputArray (Input ArrayTransforms
_ VName
v Type
_) = VName
v
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
_ ShapeChange SubExp
shape) =
Type
t Type -> Shape -> Type
forall newshape oldshape u.
ArrayShape newshape =>
TypeBase oldshape u -> newshape -> TypeBase newshape u
`setArrayShape` ShapeChange SubExp -> Shape
newShape ShapeChange SubExp
shape
transformType Type
t (ReshapeOuter Certs
_ ShapeChange SubExp
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 (ShapeChange SubExp -> [SubExp]
forall d. ShapeChange d -> [d]
newDims ShapeChange SubExp
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
_ ShapeChange SubExp
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]
++ ShapeChange SubExp -> [SubExp]
forall d. ShapeChange d -> [d]
newDims ShapeChange SubExp
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 ShapeChange SubExp
shape) =
ArrayTransform -> Input -> Input
addTransform (Certs -> ShapeChange SubExp -> ArrayTransform
ReshapeInner Certs
cs ShapeChange SubExp
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 ShapeChange SubExp
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 [[Doc] -> Doc
PP.apply ((DimChange SubExp -> Doc) -> ShapeChange SubExp -> [Doc]
forall a b. (a -> b) -> [a] -> [b]
map DimChange SubExp -> Doc
forall a. Pretty a => a -> Doc
ppr ShapeChange SubExp
shape), Doc
e]
f Doc
e (ReshapeOuter Certs
cs ShapeChange SubExp
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 [[Doc] -> Doc
PP.apply ((DimChange SubExp -> Doc) -> ShapeChange SubExp -> [Doc]
forall a b. (a -> b) -> [a] -> [b]
map DimChange SubExp -> Doc
forall a. Pretty a => a -> Doc
ppr ShapeChange SubExp
shape), Doc
e]
f Doc
e (ReshapeInner Certs
cs ShapeChange SubExp
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 [[Doc] -> Doc
PP.apply ((DimChange SubExp -> Doc) -> ShapeChange SubExp -> [Doc]
forall a b. (a -> b) -> [a] -> [b]
map DimChange SubExp -> Doc
forall a. Pretty a => a -> Doc
ppr ShapeChange SubExp
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 SOAC rep
soac = String -> Doc
text (String -> Doc) -> String -> Doc
forall a b. (a -> b) -> a -> b
$ SOAC rep -> String
forall a. Show a => a -> String
show SOAC rep
soac
inputs :: SOAC rep -> [Input]
inputs :: 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 :: [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 :: 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 :: 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 :: 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. LambdaT 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. LambdaT 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. LambdaT 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. LambdaT 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 :: 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 :: SOAC (Rep m) -> m (Exp (Rep m))
toExp SOAC (Rep m)
soac = SOAC (Rep m) -> Exp (Rep m)
forall rep. Op rep -> ExpT 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 :: 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 :: Exp rep -> m (Either NotSOAC (SOAC rep))
fromExp (Op (Futhark.Stream w as form nes 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 w ivs lam 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 w arrs 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 w arrs ops 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 ::
(MonadFreshNames m, Buildable rep, Op rep ~ Futhark.SOAC rep) =>
SOAC rep ->
m (SOAC rep, [Ident])
soacToStream :: 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
Futhark.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 -> ExpT 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
Futhark.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] -> BodyT rep -> [Type] -> LambdaT 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] -> BodyT rep -> [Type] -> LambdaT 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 (m :: * -> *) a. Monad m => a -> m a
return (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. LambdaT rep -> [Type]
lambdaReturnType Lambda rep
scan_lam
map_arr_ts :: [Type]
map_arr_ts = Int -> [Type] -> [Type]
forall a. Int -> [a] -> [a]
drop ([SubExp] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length [SubExp]
nes) [Type]
loutps
accrtps :: [Type]
accrtps = Lambda rep -> [Type]
forall rep. LambdaT rep -> [Type]
lambdaReturnType Lambda rep
scan_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]
scan_arr_ts
[Ident]
scan0_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
"resarr0") [Type]
scan_arr_ts
[Ident]
map_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
"map_res") [Type]
map_arr_ts
[Ident]
lastel_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
"lstel") [Type]
accrtps
[Ident]
lastel_tmp_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
"lstel_tmp") [Type]
accrtps
Ident
empty_arr <- String -> Type -> m Ident
forall (m :: * -> *).
MonadFreshNames m =>
String -> Type -> m Ident
newIdent String
"empty_arr" (Type -> m Ident) -> Type -> m Ident
forall a b. (a -> b) -> a -> b
$ PrimType -> Type
forall shape u. PrimType -> TypeBase shape u
Prim PrimType
Bool
[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
outszm1id <- String -> Type -> m Ident
forall (m :: * -> *).
MonadFreshNames m =>
String -> Type -> m Ident
newIdent String
"szm1" (Type -> m Ident) -> Type -> m Ident
forall a b. (a -> b) -> a -> b
$ PrimType -> Type
forall shape u. PrimType -> TypeBase shape u
Prim PrimType
int64
let insstm :: Stm rep
insstm =
[Ident] -> Exp rep -> Stm rep
forall rep. Buildable rep => [Ident] -> Exp rep -> Stm rep
mkLet ([Ident]
scan0_ids [Ident] -> [Ident] -> [Ident]
forall a. [a] -> [a] -> [a]
++ [Ident]
map_resids) (Exp rep -> Stm rep)
-> (SOAC rep -> Exp rep) -> SOAC rep -> Stm rep
forall b c a. (b -> c) -> (a -> b) -> a -> c
. SOAC rep -> Exp rep
forall rep. Op rep -> ExpT rep
Op (SOAC rep -> Stm rep) -> SOAC rep -> Stm rep
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'
outszm1stm :: Stm rep
outszm1stm =
[Ident] -> Exp rep -> Stm rep
forall rep. Buildable rep => [Ident] -> Exp rep -> Stm rep
mkLet [Ident
outszm1id] (Exp rep -> Stm rep) -> (BasicOp -> Exp rep) -> BasicOp -> Stm rep
forall b c a. (b -> c) -> (a -> b) -> a -> c
. BasicOp -> Exp rep
forall rep. BasicOp -> ExpT rep
BasicOp (BasicOp -> Stm rep) -> BasicOp -> Stm rep
forall a b. (a -> b) -> a -> b
$
BinOp -> SubExp -> SubExp -> BasicOp
BinOp
(IntType -> Overflow -> BinOp
Sub IntType
Int64 Overflow
OverflowUndef)
(VName -> SubExp
Futhark.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))
empty_arr_stm :: Stm rep
empty_arr_stm =
[Ident] -> Exp rep -> Stm rep
forall rep. Buildable rep => [Ident] -> Exp rep -> Stm rep
mkLet [Ident
empty_arr] (Exp rep -> Stm rep) -> (BasicOp -> Exp rep) -> BasicOp -> Stm rep
forall b c a. (b -> c) -> (a -> b) -> a -> c
. BasicOp -> Exp rep
forall rep. BasicOp -> ExpT rep
BasicOp (BasicOp -> Stm rep) -> BasicOp -> Stm rep
forall a b. (a -> b) -> a -> b
$
CmpOp -> SubExp -> SubExp -> BasicOp
CmpOp
(IntType -> CmpOp
CmpSlt IntType
Int64)
(VName -> SubExp
Futhark.Var (VName -> SubExp) -> VName -> SubExp
forall a b. (a -> b) -> a -> b
$ Ident -> VName
identName Ident
outszm1id)
(Int64 -> SubExp
forall v. IsValue v => v -> SubExp
constant (Int64
0 :: Int64))
leltmpstms :: [Stm rep]
leltmpstms =
(Ident -> Ident -> Stm rep) -> [Ident] -> [Ident] -> [Stm rep]
forall a b c. (a -> b -> c) -> [a] -> [b] -> [c]
zipWith
( \Ident
lid Ident
arrid ->
[Ident] -> Exp rep -> Stm rep
forall rep. Buildable rep => [Ident] -> Exp rep -> Stm rep
mkLet [Ident
lid] (Exp rep -> Stm rep) -> (BasicOp -> Exp rep) -> BasicOp -> Stm rep
forall b c a. (b -> c) -> (a -> b) -> a -> c
. BasicOp -> Exp rep
forall rep. BasicOp -> ExpT rep
BasicOp (BasicOp -> Stm rep) -> BasicOp -> Stm rep
forall a b. (a -> b) -> a -> b
$
VName -> Slice SubExp -> BasicOp
Index (Ident -> VName
identName Ident
arrid) (Slice SubExp -> BasicOp) -> Slice SubExp -> BasicOp
forall a b. (a -> b) -> a -> b
$
Type -> [DimIndex SubExp] -> Slice SubExp
fullSlice
(Ident -> Type
identType Ident
arrid)
[SubExp -> DimIndex SubExp
forall d. d -> DimIndex d
DimFix (SubExp -> DimIndex SubExp) -> SubExp -> DimIndex SubExp
forall a b. (a -> b) -> a -> b
$ VName -> SubExp
Futhark.Var (VName -> SubExp) -> VName -> SubExp
forall a b. (a -> b) -> a -> b
$ Ident -> VName
identName Ident
outszm1id]
)
[Ident]
lastel_tmp_ids
[Ident]
scan0_ids
lelstm :: Stm rep
lelstm =
[Ident] -> Exp rep -> Stm rep
forall rep. Buildable rep => [Ident] -> Exp rep -> Stm rep
mkLet [Ident]
lastel_ids (Exp rep -> Stm rep) -> Exp rep -> Stm rep
forall a b. (a -> b) -> a -> b
$
SubExp -> Body rep -> Body rep -> IfDec (BranchType rep) -> Exp rep
forall rep.
SubExp
-> BodyT rep -> BodyT rep -> IfDec (BranchType rep) -> ExpT rep
If
(VName -> SubExp
Futhark.Var (VName -> SubExp) -> VName -> SubExp
forall a b. (a -> b) -> a -> b
$ Ident -> VName
identName Ident
empty_arr)
(Stms rep -> Result -> Body rep
forall rep. Buildable rep => Stms rep -> Result -> Body rep
mkBody Stms rep
forall a. Monoid a => a
mempty (Result -> Body rep) -> Result -> Body rep
forall a b. (a -> b) -> a -> b
$ [SubExp] -> Result
subExpsRes [SubExp]
nes)
( 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
stmsFromList [Stm rep]
leltmpstms) (Result -> Body rep) -> Result -> Body rep
forall a b. (a -> b) -> a -> b
$
[VName] -> Result
varsRes ([VName] -> Result) -> [VName] -> Result
forall a b. (a -> b) -> a -> b
$ (Ident -> VName) -> [Ident] -> [VName]
forall a b. (a -> b) -> [a] -> [b]
map Ident -> VName
identName [Ident]
lastel_tmp_ids
)
(IfDec (BranchType rep) -> Exp rep)
-> IfDec (BranchType rep) -> Exp rep
forall a b. (a -> b) -> a -> b
$ [Type] -> IfDec ExtType
ifCommon ([Type] -> IfDec ExtType) -> [Type] -> IfDec ExtType
forall a b. (a -> b) -> a -> b
$ (Ident -> Type) -> [Ident] -> [Type]
forall a b. (a -> b) -> [a] -> [b]
map Ident -> Type
identType [Ident]
lastel_tmp_ids
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
Futhark.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 mapstm :: Stm rep
mapstm =
[Ident] -> Exp rep -> Stm rep
forall rep. Buildable rep => [Ident] -> Exp rep -> Stm rep
mkLet [Ident]
strm_resids (Exp rep -> Stm rep)
-> (SOAC rep -> Exp rep) -> SOAC rep -> Stm rep
forall b c a. (b -> c) -> (a -> b) -> a -> c
. SOAC rep -> Exp rep
forall rep. Op rep -> ExpT rep
Op (SOAC rep -> Stm rep) -> SOAC rep -> Stm rep
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 ((Ident -> VName) -> [Ident] -> [VName]
forall a b. (a -> b) -> [a] -> [b]
map Ident -> VName
identName [Ident]
scan0_ids) (Lambda rep -> ScremaForm rep
forall rep. Lambda rep -> ScremaForm rep
Futhark.mapSOAC Lambda rep
maplam)
Body rep
addlelbdy <-
Lambda rep -> [SubExp] -> m (Body rep)
forall (m :: * -> *) rep.
(MonadFreshNames m, Buildable rep) =>
Lambda rep -> [SubExp] -> m (Body rep)
mkPlusBnds Lambda rep
scan_lam ([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
Futhark.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]
lastel_ids
let (Stms rep
addlelstm, Result
addlelres) = (Body rep -> Stms rep
forall rep. BodyT rep -> Stms rep
bodyStms Body rep
addlelbdy, Body rep -> Result
forall rep. BodyT rep -> Result
bodyResult Body rep
addlelbdy)
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
stmsFromList [Stm rep
insstm, Stm rep
outszm1stm, Stm rep
empty_arr_stm, Stm rep
lelstm, Stm rep
mapstm] Stms rep -> Stms rep -> Stms rep
forall a. Semigroup a => a -> a -> a
<> Stms rep
addlelstm) (Result -> Body rep) -> Result -> Body rep
forall a b. (a -> b) -> a -> b
$
Result
addlelres 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
Futhark.Var (VName -> SubExp) -> (Ident -> VName) -> Ident -> SubExp
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Ident -> VName
identName) ([Ident]
strm_resids [Ident] -> [Ident] -> [Ident]
forall a. [a] -> [a] -> [a]
++ [Ident]
map_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] -> BodyT rep -> [Type] -> LambdaT rep
Lambda [Param Type]
[LParam rep]
strmpar Body rep
strmbdy ([Type]
accrtps [Type] -> [Type] -> [Type]
forall a. [a] -> [a] -> [a]
++ [Type]
loutps)
(SOAC rep, [Ident]) -> m (SOAC rep, [Ident])
forall (m :: * -> *) a. Monad m => a -> m a
return
( 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. LambdaT 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 -> ExpT 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
Futhark.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. BodyT rep -> Stms rep
bodyStms Body rep
addaccbdy, Body rep -> Result
forall rep. BodyT 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
Futhark.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] -> BodyT rep -> [Type] -> LambdaT 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 (m :: * -> *) a. Monad m => a -> m a
return (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 (m :: * -> *) a. Monad m => a -> m a
return (SOAC rep
soac, [])
where
mkMapPlusAccLam ::
(MonadFreshNames m, Buildable rep) =>
[SubExp] ->
Lambda rep ->
m (Lambda rep)
mkMapPlusAccLam :: [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. LambdaT 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 -> ExpT 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 :: BodyT rep
plus_bdy = Lambda rep -> BodyT rep
forall rep. LambdaT rep -> BodyT rep
lambdaBody Lambda rep
plus
newlambdy :: BodyT rep
newlambdy =
BodyDec rep -> Stms rep -> Result -> BodyT rep
forall rep. BodyDec rep -> Stms rep -> Result -> BodyT rep
Body
(BodyT rep -> BodyDec rep
forall rep. BodyT rep -> BodyDec rep
bodyDec BodyT 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
<> BodyT rep -> Stms rep
forall rep. BodyT rep -> Stms rep
bodyStms BodyT rep
plus_bdy)
(BodyT rep -> Result
forall rep. BodyT rep -> Result
bodyResult BodyT 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] -> BodyT rep -> [Type] -> Lambda rep
forall rep. [LParam rep] -> BodyT rep -> [Type] -> LambdaT rep
Lambda [Param Type]
[LParam rep]
rempars BodyT rep
newlambdy ([Type] -> Lambda rep) -> [Type] -> Lambda rep
forall a b. (a -> b) -> a -> b
$ Lambda rep -> [Type]
forall rep. LambdaT rep -> [Type]
lambdaReturnType Lambda rep
plus
mkPlusBnds ::
(MonadFreshNames m, Buildable rep) =>
Lambda rep ->
[SubExp] ->
m (Body rep)
mkPlusBnds :: 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 -> ExpT 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. LambdaT rep -> [LParam rep]
lambdaParams Lambda rep
plus')
[SubExp]
accels
body :: Body rep
body = Lambda rep -> Body rep
forall rep. LambdaT rep -> BodyT rep
lambdaBody Lambda rep
plus'
Body rep -> m (Body rep)
forall (m :: * -> *) a. Monad m => a -> m a
return (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. BodyT rep -> Stms rep
bodyStms Body rep
body}