{-# LANGUAGE TypeFamilies #-}
module Futhark.Analysis.HORep.MapNest
( Nesting (..),
MapNest (..),
typeOf,
params,
inputs,
setInputs,
fromSOAC,
toSOAC,
)
where
import Data.List (find)
import Data.Map.Strict qualified as M
import Data.Maybe
import Futhark.Analysis.HORep.SOAC (SOAC)
import Futhark.Analysis.HORep.SOAC qualified as SOAC
import Futhark.Construct
import Futhark.IR hiding (typeOf)
import Futhark.IR.SOACS.SOAC qualified as Futhark
import Futhark.Transform.Substitute
data Nesting rep = Nesting
{ forall {k} (rep :: k). Nesting rep -> [VName]
nestingParamNames :: [VName],
forall {k} (rep :: k). Nesting rep -> [VName]
nestingResult :: [VName],
forall {k} (rep :: k). Nesting rep -> [Type]
nestingReturnType :: [Type],
forall {k} (rep :: k). Nesting rep -> SubExp
nestingWidth :: SubExp
}
deriving (Nesting rep -> Nesting rep -> Bool
forall a. (a -> a -> Bool) -> (a -> a -> Bool) -> Eq a
forall k (rep :: k). Nesting rep -> Nesting rep -> Bool
/= :: Nesting rep -> Nesting rep -> Bool
$c/= :: forall k (rep :: k). Nesting rep -> Nesting rep -> Bool
== :: Nesting rep -> Nesting rep -> Bool
$c== :: forall k (rep :: k). Nesting rep -> Nesting rep -> Bool
Eq, Nesting rep -> Nesting rep -> Bool
Nesting rep -> Nesting rep -> Ordering
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
forall k (rep :: k). Eq (Nesting rep)
forall k (rep :: k). Nesting rep -> Nesting rep -> Bool
forall k (rep :: k). Nesting rep -> Nesting rep -> Ordering
forall k (rep :: k). Nesting rep -> Nesting rep -> Nesting rep
min :: Nesting rep -> Nesting rep -> Nesting rep
$cmin :: forall k (rep :: k). Nesting rep -> Nesting rep -> Nesting rep
max :: Nesting rep -> Nesting rep -> Nesting rep
$cmax :: forall k (rep :: k). Nesting rep -> Nesting rep -> Nesting rep
>= :: Nesting rep -> Nesting rep -> Bool
$c>= :: forall k (rep :: k). Nesting rep -> Nesting rep -> Bool
> :: Nesting rep -> Nesting rep -> Bool
$c> :: forall k (rep :: k). Nesting rep -> Nesting rep -> Bool
<= :: Nesting rep -> Nesting rep -> Bool
$c<= :: forall k (rep :: k). Nesting rep -> Nesting rep -> Bool
< :: Nesting rep -> Nesting rep -> Bool
$c< :: forall k (rep :: k). Nesting rep -> Nesting rep -> Bool
compare :: Nesting rep -> Nesting rep -> Ordering
$ccompare :: forall k (rep :: k). Nesting rep -> Nesting rep -> Ordering
Ord, Int -> Nesting rep -> ShowS
forall a.
(Int -> a -> ShowS) -> (a -> String) -> ([a] -> ShowS) -> Show a
forall k (rep :: k). Int -> Nesting rep -> ShowS
forall k (rep :: k). [Nesting rep] -> ShowS
forall k (rep :: k). Nesting rep -> String
showList :: [Nesting rep] -> ShowS
$cshowList :: forall k (rep :: k). [Nesting rep] -> ShowS
show :: Nesting rep -> String
$cshow :: forall k (rep :: k). Nesting rep -> String
showsPrec :: Int -> Nesting rep -> ShowS
$cshowsPrec :: forall k (rep :: k). Int -> Nesting rep -> ShowS
Show)
data MapNest rep = MapNest SubExp (Lambda rep) [Nesting rep] [SOAC.Input]
deriving (Int -> MapNest rep -> ShowS
forall a.
(Int -> a -> ShowS) -> (a -> String) -> ([a] -> ShowS) -> Show a
forall k (rep :: k). RepTypes rep => Int -> MapNest rep -> ShowS
forall k (rep :: k). RepTypes rep => [MapNest rep] -> ShowS
forall k (rep :: k). RepTypes rep => MapNest rep -> String
showList :: [MapNest rep] -> ShowS
$cshowList :: forall k (rep :: k). RepTypes rep => [MapNest rep] -> ShowS
show :: MapNest rep -> String
$cshow :: forall k (rep :: k). RepTypes rep => MapNest rep -> String
showsPrec :: Int -> MapNest rep -> ShowS
$cshowsPrec :: forall k (rep :: k). RepTypes rep => Int -> MapNest rep -> ShowS
Show)
typeOf :: MapNest rep -> [Type]
typeOf :: forall {k} (rep :: k). MapNest rep -> [Type]
typeOf (MapNest SubExp
w Lambda rep
lam [] [Input]
_) =
forall a b. (a -> b) -> [a] -> [b]
map (forall d.
ArrayShape (ShapeBase d) =>
TypeBase (ShapeBase d) NoUniqueness
-> d -> TypeBase (ShapeBase d) NoUniqueness
`arrayOfRow` SubExp
w) forall a b. (a -> b) -> a -> b
$ forall {k} (rep :: k). Lambda rep -> [Type]
lambdaReturnType Lambda rep
lam
typeOf (MapNest SubExp
w Lambda rep
_ (Nesting rep
nest : [Nesting rep]
_) [Input]
_) =
forall a b. (a -> b) -> [a] -> [b]
map (forall d.
ArrayShape (ShapeBase d) =>
TypeBase (ShapeBase d) NoUniqueness
-> d -> TypeBase (ShapeBase d) NoUniqueness
`arrayOfRow` SubExp
w) forall a b. (a -> b) -> a -> b
$ forall {k} (rep :: k). Nesting rep -> [Type]
nestingReturnType Nesting rep
nest
params :: MapNest rep -> [VName]
params :: forall {k} (rep :: k). MapNest rep -> [VName]
params (MapNest SubExp
_ Lambda rep
lam [] [Input]
_) =
forall a b. (a -> b) -> [a] -> [b]
map forall dec. Param dec -> VName
paramName forall a b. (a -> b) -> a -> b
$ forall {k} (rep :: k). Lambda rep -> [LParam rep]
lambdaParams Lambda rep
lam
params (MapNest SubExp
_ Lambda rep
_ (Nesting rep
nest : [Nesting rep]
_) [Input]
_) =
forall {k} (rep :: k). Nesting rep -> [VName]
nestingParamNames Nesting rep
nest
inputs :: MapNest rep -> [SOAC.Input]
inputs :: forall {k} (rep :: k). MapNest rep -> [Input]
inputs (MapNest SubExp
_ Lambda rep
_ [Nesting rep]
_ [Input]
inps) = [Input]
inps
setInputs :: [SOAC.Input] -> MapNest rep -> MapNest rep
setInputs :: forall {k} (rep :: k). [Input] -> MapNest rep -> MapNest rep
setInputs [] (MapNest SubExp
w Lambda rep
body [Nesting rep]
ns [Input]
_) = forall {k} (rep :: k).
SubExp -> Lambda rep -> [Nesting rep] -> [Input] -> MapNest rep
MapNest SubExp
w Lambda rep
body [Nesting rep]
ns []
setInputs (Input
inp : [Input]
inps) (MapNest SubExp
_ Lambda rep
body [Nesting rep]
ns [Input]
_) = forall {k} (rep :: k).
SubExp -> Lambda rep -> [Nesting rep] -> [Input] -> MapNest rep
MapNest SubExp
w Lambda rep
body [Nesting rep]
ns' (Input
inp forall a. a -> [a] -> [a]
: [Input]
inps)
where
w :: SubExp
w = forall u. Int -> TypeBase (ShapeBase SubExp) u -> SubExp
arraySize Int
0 forall a b. (a -> b) -> a -> b
$ Input -> Type
SOAC.inputType Input
inp
ws :: [SubExp]
ws = forall a. Int -> [a] -> [a]
drop Int
1 forall a b. (a -> b) -> a -> b
$ forall u. TypeBase (ShapeBase SubExp) u -> [SubExp]
arrayDims forall a b. (a -> b) -> a -> b
$ Input -> Type
SOAC.inputType Input
inp
ns' :: [Nesting rep]
ns' = forall a b c. (a -> b -> c) -> [a] -> [b] -> [c]
zipWith forall {k} {k} {rep :: k} {rep :: k}.
Nesting rep -> SubExp -> Nesting rep
setDepth [Nesting rep]
ns [SubExp]
ws
setDepth :: Nesting rep -> SubExp -> Nesting rep
setDepth Nesting rep
n SubExp
nw = Nesting rep
n {nestingWidth :: SubExp
nestingWidth = SubExp
nw}
fromSOAC ::
( Buildable rep,
MonadFreshNames m,
LocalScope rep m,
Op rep ~ Futhark.SOAC rep
) =>
SOAC rep ->
m (Maybe (MapNest rep))
fromSOAC :: forall {k} (rep :: k) (m :: * -> *).
(Buildable rep, MonadFreshNames m, LocalScope rep m,
Op rep ~ SOAC rep) =>
SOAC rep -> m (Maybe (MapNest rep))
fromSOAC = forall {k} (rep :: k) (m :: * -> *).
(Buildable rep, MonadFreshNames m, LocalScope rep m,
Op rep ~ SOAC rep) =>
[Ident] -> SOAC rep -> m (Maybe (MapNest rep))
fromSOAC' forall a. Monoid a => a
mempty
fromSOAC' ::
( Buildable rep,
MonadFreshNames m,
LocalScope rep m,
Op rep ~ Futhark.SOAC rep
) =>
[Ident] ->
SOAC rep ->
m (Maybe (MapNest rep))
fromSOAC' :: forall {k} (rep :: k) (m :: * -> *).
(Buildable rep, MonadFreshNames m, LocalScope rep m,
Op rep ~ SOAC rep) =>
[Ident] -> SOAC rep -> m (Maybe (MapNest rep))
fromSOAC' [Ident]
bound (SOAC.Screma SubExp
w (SOAC.ScremaForm [] [] Lambda rep
lam) [Input]
inps) = do
Either NotSOAC (Maybe (Pat (LetDec rep), MapNest rep))
maybenest <- case ( forall {k} (rep :: k). Stms rep -> [Stm rep]
stmsToList forall a b. (a -> b) -> a -> b
$ forall {k} (rep :: k). Body rep -> Stms rep
bodyStms forall a b. (a -> b) -> a -> b
$ forall {k} (rep :: k). Lambda rep -> Body rep
lambdaBody Lambda rep
lam,
forall {k} (rep :: k). Body rep -> Result
bodyResult forall a b. (a -> b) -> a -> b
$ forall {k} (rep :: k). Lambda rep -> Body rep
lambdaBody Lambda rep
lam
) of
([Let Pat (LetDec rep)
pat StmAux (ExpDec rep)
_ Exp rep
e], Result
res)
| forall a b. (a -> b) -> [a] -> [b]
map SubExpRes -> SubExp
resSubExp Result
res forall a. Eq a => a -> a -> Bool
== forall a b. (a -> b) -> [a] -> [b]
map VName -> SubExp
Var (forall dec. Pat dec -> [VName]
patNames Pat (LetDec rep)
pat) ->
forall {k} (rep :: k) (m :: * -> *) a.
LocalScope rep m =>
Scope rep -> m a -> m a
localScope (forall {k} (rep :: k) dec.
(LParamInfo rep ~ dec) =>
[Param dec] -> Scope rep
scopeOfLParams forall a b. (a -> b) -> a -> b
$ forall {k} (rep :: k). Lambda rep -> [LParam rep]
lambdaParams Lambda rep
lam) forall a b. (a -> b) -> a -> b
$
forall {k} (rep :: k) (m :: * -> *).
(Op rep ~ SOAC rep, HasScope rep m) =>
Exp rep -> m (Either NotSOAC (SOAC rep))
SOAC.fromExp Exp rep
e
forall (m :: * -> *) a b. Monad m => m a -> (a -> m b) -> m b
>>= forall a c b. (a -> c) -> (b -> c) -> Either a b -> c
either (forall (f :: * -> *) a. Applicative f => a -> f a
pure forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall a b. a -> Either a b
Left) (forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap (forall a b. b -> Either a b
Right forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap (Pat (LetDec rep)
pat,)) forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall {k} (rep :: k) (m :: * -> *).
(Buildable rep, MonadFreshNames m, LocalScope rep m,
Op rep ~ SOAC rep) =>
[Ident] -> SOAC rep -> m (Maybe (MapNest rep))
fromSOAC' [Ident]
bound')
([Stm rep], Result)
_ ->
forall (f :: * -> *) a. Applicative f => a -> f a
pure forall a b. (a -> b) -> a -> b
$ forall a b. b -> Either a b
Right forall a. Maybe a
Nothing
case Either NotSOAC (Maybe (Pat (LetDec rep), MapNest rep))
maybenest of
Right (Just (Pat (LetDec rep)
pat, mn :: MapNest rep
mn@(MapNest SubExp
inner_w Lambda rep
body' [Nesting rep]
ns' [Input]
inps'))) -> do
([VName]
ps, [Input]
inps'') <-
forall a b. [(a, b)] -> ([a], [b])
unzip
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> forall (m :: * -> *).
MonadFreshNames m =>
SubExp
-> [(VName, Input)] -> [(VName, Input)] -> m [(VName, Input)]
fixInputs
SubExp
w
(forall a b. [a] -> [b] -> [(a, b)]
zip (forall a b. (a -> b) -> [a] -> [b]
map forall dec. Param dec -> VName
paramName forall a b. (a -> b) -> a -> b
$ forall {k} (rep :: k). Lambda rep -> [LParam rep]
lambdaParams Lambda rep
lam) [Input]
inps)
(forall a b. [a] -> [b] -> [(a, b)]
zip (forall {k} (rep :: k). MapNest rep -> [VName]
params MapNest rep
mn) [Input]
inps')
let n' :: Nesting rep
n' =
Nesting
{ nestingParamNames :: [VName]
nestingParamNames = [VName]
ps,
nestingResult :: [VName]
nestingResult = forall dec. Pat dec -> [VName]
patNames Pat (LetDec rep)
pat,
nestingReturnType :: [Type]
nestingReturnType = forall {k} (rep :: k). MapNest rep -> [Type]
typeOf MapNest rep
mn,
nestingWidth :: SubExp
nestingWidth = SubExp
inner_w
}
forall (f :: * -> *) a. Applicative f => a -> f a
pure forall a b. (a -> b) -> a -> b
$ forall a. a -> Maybe a
Just forall a b. (a -> b) -> a -> b
$ forall {k} (rep :: k).
SubExp -> Lambda rep -> [Nesting rep] -> [Input] -> MapNest rep
MapNest SubExp
w Lambda rep
body' (Nesting rep
n' forall a. a -> [a] -> [a]
: [Nesting rep]
ns') [Input]
inps''
Either NotSOAC (Maybe (Pat (LetDec rep), MapNest rep))
_ -> do
let isBound :: VName -> Maybe Ident
isBound VName
name
| Just Ident
param <- forall (t :: * -> *) a. Foldable t => (a -> Bool) -> t a -> Maybe a
find ((VName
name ==) forall b c a. (b -> c) -> (a -> b) -> a -> c
. Ident -> VName
identName) [Ident]
bound =
forall a. a -> Maybe a
Just Ident
param
| Bool
otherwise =
forall a. Maybe a
Nothing
boundUsedInBody :: [Ident]
boundUsedInBody =
forall a b. (a -> Maybe b) -> [a] -> [b]
mapMaybe VName -> Maybe Ident
isBound forall a b. (a -> b) -> a -> b
$ Names -> [VName]
namesToList forall a b. (a -> b) -> a -> b
$ forall a. FreeIn a => a -> Names
freeIn Lambda rep
lam
[Ident]
newParams <- forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
(a -> m b) -> t a -> m (t b)
mapM (forall (m :: * -> *).
MonadFreshNames m =>
ShowS -> Ident -> m Ident
newIdent' (forall a. [a] -> [a] -> [a]
++ String
"_wasfree")) [Ident]
boundUsedInBody
let subst :: Map VName VName
subst =
forall k a. Ord k => [(k, a)] -> Map k a
M.fromList forall a b. (a -> b) -> a -> b
$
forall a b. [a] -> [b] -> [(a, b)]
zip (forall a b. (a -> b) -> [a] -> [b]
map Ident -> VName
identName [Ident]
boundUsedInBody) (forall a b. (a -> b) -> [a] -> [b]
map Ident -> VName
identName [Ident]
newParams)
inps' :: [Input]
inps' =
[Input]
inps
forall a. [a] -> [a] -> [a]
++ forall a b. (a -> b) -> [a] -> [b]
map
(ArrayTransform -> Input -> Input
SOAC.addTransform (Certs -> ShapeBase SubExp -> ArrayTransform
SOAC.Replicate forall a. Monoid a => a
mempty forall a b. (a -> b) -> a -> b
$ forall d. [d] -> ShapeBase d
Shape [SubExp
w]) forall b c a. (b -> c) -> (a -> b) -> a -> c
. Ident -> Input
SOAC.identInput)
[Ident]
boundUsedInBody
lam' :: Lambda rep
lam' =
Lambda rep
lam
{ lambdaBody :: Body rep
lambdaBody =
forall a. Substitute a => Map VName VName -> a -> a
substituteNames Map VName VName
subst forall a b. (a -> b) -> a -> b
$ forall {k} (rep :: k). Lambda rep -> Body rep
lambdaBody Lambda rep
lam,
lambdaParams :: [LParam rep]
lambdaParams =
forall {k} (rep :: k). Lambda rep -> [LParam rep]
lambdaParams Lambda rep
lam
forall a. [a] -> [a] -> [a]
++ [forall dec. Attrs -> VName -> dec -> Param dec
Param forall a. Monoid a => a
mempty VName
name Type
t | Ident VName
name Type
t <- [Ident]
newParams]
}
forall (f :: * -> *) a. Applicative f => a -> f a
pure forall a b. (a -> b) -> a -> b
$ forall a. a -> Maybe a
Just forall a b. (a -> b) -> a -> b
$ forall {k} (rep :: k).
SubExp -> Lambda rep -> [Nesting rep] -> [Input] -> MapNest rep
MapNest SubExp
w Lambda rep
lam' [] [Input]
inps'
where
bound' :: [Ident]
bound' = [Ident]
bound forall a. Semigroup a => a -> a -> a
<> forall a b. (a -> b) -> [a] -> [b]
map forall dec. Typed dec => Param dec -> Ident
paramIdent (forall {k} (rep :: k). Lambda rep -> [LParam rep]
lambdaParams Lambda rep
lam)
fromSOAC' [Ident]
_ SOAC rep
_ = forall (f :: * -> *) a. Applicative f => a -> f a
pure forall a. Maybe a
Nothing
toSOAC ::
( MonadFreshNames m,
HasScope rep m,
Buildable rep,
BuilderOps rep,
Op rep ~ Futhark.SOAC rep
) =>
MapNest rep ->
m (SOAC rep)
toSOAC :: forall (m :: * -> *) rep.
(MonadFreshNames m, HasScope rep m, Buildable rep, BuilderOps rep,
Op rep ~ SOAC rep) =>
MapNest rep -> m (SOAC rep)
toSOAC (MapNest SubExp
w Lambda rep
lam [] [Input]
inps) =
forall (f :: * -> *) a. Applicative f => a -> f a
pure forall a b. (a -> b) -> a -> b
$ forall {k} (rep :: k).
SubExp -> ScremaForm rep -> [Input] -> SOAC rep
SOAC.Screma SubExp
w (forall {k} (rep :: k). Lambda rep -> ScremaForm rep
Futhark.mapSOAC Lambda rep
lam) [Input]
inps
toSOAC (MapNest SubExp
w Lambda rep
lam (Nesting [VName]
npnames [VName]
nres [Type]
nrettype SubExp
nw : [Nesting rep]
ns) [Input]
inps) = do
let nparams :: [Param Type]
nparams = forall a b c. (a -> b -> c) -> [a] -> [b] -> [c]
zipWith (forall dec. Attrs -> VName -> dec -> Param dec
Param forall a. Monoid a => a
mempty) [VName]
npnames forall a b. (a -> b) -> a -> b
$ forall a b. (a -> b) -> [a] -> [b]
map Input -> Type
SOAC.inputRowType [Input]
inps
Body rep
body <- forall {k1} {k2} (rep :: k1) (m :: * -> *) (somerep :: k2).
(Buildable rep, MonadFreshNames m, HasScope somerep m,
SameScope somerep rep) =>
Builder rep (Body rep) -> m (Body rep)
runBodyBuilder forall a b. (a -> b) -> a -> b
$
forall {k} (rep :: k) (m :: * -> *) a.
LocalScope rep m =>
Scope rep -> m a -> m a
localScope (forall {k} (rep :: k) dec.
(LParamInfo rep ~ dec) =>
[Param dec] -> Scope rep
scopeOfLParams [Param Type]
nparams) forall a b. (a -> b) -> a -> b
$ do
forall (m :: * -> *).
MonadBuilder m =>
[VName] -> Exp (Rep m) -> m ()
letBindNames [VName]
nres
forall (m :: * -> *) a b. Monad m => (a -> m b) -> m a -> m b
=<< forall (m :: * -> *).
(MonadBuilder m, Op (Rep m) ~ SOAC (Rep m)) =>
SOAC (Rep m) -> m (Exp (Rep m))
SOAC.toExp
forall (m :: * -> *) a b. Monad m => (a -> m b) -> m a -> m b
=<< forall (m :: * -> *) rep.
(MonadFreshNames m, HasScope rep m, Buildable rep, BuilderOps rep,
Op rep ~ SOAC rep) =>
MapNest rep -> m (SOAC rep)
toSOAC (forall {k} (rep :: k).
SubExp -> Lambda rep -> [Nesting rep] -> [Input] -> MapNest rep
MapNest SubExp
nw Lambda rep
lam [Nesting rep]
ns forall a b. (a -> b) -> a -> b
$ forall a b. (a -> b) -> [a] -> [b]
map (Ident -> Input
SOAC.identInput forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall dec. Typed dec => Param dec -> Ident
paramIdent) [Param Type]
nparams)
forall (f :: * -> *) a. Applicative f => a -> f a
pure forall a b. (a -> b) -> a -> b
$ forall {k} (rep :: k). Buildable rep => [SubExp] -> Body rep
resultBody forall a b. (a -> b) -> a -> b
$ forall a b. (a -> b) -> [a] -> [b]
map VName -> SubExp
Var [VName]
nres
let outerlam :: Lambda rep
outerlam =
Lambda
{ lambdaParams :: [LParam rep]
lambdaParams = [Param Type]
nparams,
lambdaBody :: Body rep
lambdaBody = Body rep
body,
lambdaReturnType :: [Type]
lambdaReturnType = [Type]
nrettype
}
forall (f :: * -> *) a. Applicative f => a -> f a
pure forall a b. (a -> b) -> a -> b
$ forall {k} (rep :: k).
SubExp -> ScremaForm rep -> [Input] -> SOAC rep
SOAC.Screma SubExp
w (forall {k} (rep :: k). Lambda rep -> ScremaForm rep
Futhark.mapSOAC Lambda rep
outerlam) [Input]
inps
fixInputs ::
MonadFreshNames m =>
SubExp ->
[(VName, SOAC.Input)] ->
[(VName, SOAC.Input)] ->
m [(VName, SOAC.Input)]
fixInputs :: forall (m :: * -> *).
MonadFreshNames m =>
SubExp
-> [(VName, Input)] -> [(VName, Input)] -> m [(VName, Input)]
fixInputs SubExp
w [(VName, Input)]
ourInps = forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
(a -> m b) -> t a -> m (t b)
mapM (VName, Input) -> m (VName, Input)
inspect
where
isParam :: a -> (a, b) -> Bool
isParam a
x (a
y, b
_) = a
x forall a. Eq a => a -> a -> Bool
== a
y
inspect :: (VName, Input) -> m (VName, Input)
inspect (VName
_, SOAC.Input ArrayTransforms
ts VName
v Type
_)
| Just (VName
p, Input
pInp) <- forall (t :: * -> *) a. Foldable t => (a -> Bool) -> t a -> Maybe a
find (forall {a} {b}. Eq a => a -> (a, b) -> Bool
isParam VName
v) [(VName, Input)]
ourInps = do
let pInp' :: Input
pInp' = ArrayTransforms -> Input -> Input
SOAC.transformRows ArrayTransforms
ts Input
pInp
VName
p' <- forall (m :: * -> *). MonadFreshNames m => String -> m VName
newNameFromString forall a b. (a -> b) -> a -> b
$ VName -> String
baseString VName
p
forall (f :: * -> *) a. Applicative f => a -> f a
pure (VName
p', Input
pInp')
inspect (VName
param, SOAC.Input ArrayTransforms
ts VName
a Type
t) = do
VName
param' <- forall (m :: * -> *). MonadFreshNames m => String -> m VName
newNameFromString (VName -> String
baseString VName
param forall a. [a] -> [a] -> [a]
++ String
"_rep")
forall (f :: * -> *) a. Applicative f => a -> f a
pure (VName
param', ArrayTransforms -> VName -> Type -> Input
SOAC.Input (ArrayTransforms
ts ArrayTransforms -> ArrayTransform -> ArrayTransforms
SOAC.|> Certs -> ShapeBase SubExp -> ArrayTransform
SOAC.Replicate forall a. Monoid a => a
mempty (forall d. [d] -> ShapeBase d
Shape [SubExp
w])) VName
a Type
t)