{-# LANGUAGE FlexibleContexts #-}
{-# LANGUAGE TupleSections #-}
{-# LANGUAGE TypeFamilies #-}

module Futhark.Analysis.HORep.MapNest
  ( Nesting (..),
    MapNest (..),
    typeOf,
    params,
    inputs,
    setInputs,
    fromSOAC,
    toSOAC,
  )
where

import Data.List (find)
import qualified Data.Map.Strict as M
import Data.Maybe
import Futhark.Analysis.HORep.SOAC (SOAC)
import qualified Futhark.Analysis.HORep.SOAC as SOAC
import Futhark.Construct
import Futhark.IR hiding (typeOf)
import qualified Futhark.IR.SOACS.SOAC as Futhark
import Futhark.Transform.Substitute

data Nesting lore = Nesting
  { forall lore. Nesting lore -> [VName]
nestingParamNames :: [VName],
    forall lore. Nesting lore -> [VName]
nestingResult :: [VName],
    forall lore. Nesting lore -> [Type]
nestingReturnType :: [Type],
    forall lore. Nesting lore -> SubExp
nestingWidth :: SubExp
  }
  deriving (Nesting lore -> Nesting lore -> Bool
(Nesting lore -> Nesting lore -> Bool)
-> (Nesting lore -> Nesting lore -> Bool) -> Eq (Nesting lore)
forall lore. Nesting lore -> Nesting lore -> Bool
forall a. (a -> a -> Bool) -> (a -> a -> Bool) -> Eq a
/= :: Nesting lore -> Nesting lore -> Bool
$c/= :: forall lore. Nesting lore -> Nesting lore -> Bool
== :: Nesting lore -> Nesting lore -> Bool
$c== :: forall lore. Nesting lore -> Nesting lore -> Bool
Eq, Eq (Nesting lore)
Eq (Nesting lore)
-> (Nesting lore -> Nesting lore -> Ordering)
-> (Nesting lore -> Nesting lore -> Bool)
-> (Nesting lore -> Nesting lore -> Bool)
-> (Nesting lore -> Nesting lore -> Bool)
-> (Nesting lore -> Nesting lore -> Bool)
-> (Nesting lore -> Nesting lore -> Nesting lore)
-> (Nesting lore -> Nesting lore -> Nesting lore)
-> Ord (Nesting lore)
Nesting lore -> Nesting lore -> Bool
Nesting lore -> Nesting lore -> Ordering
Nesting lore -> Nesting lore -> Nesting lore
forall lore. Eq (Nesting lore)
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 lore. Nesting lore -> Nesting lore -> Bool
forall lore. Nesting lore -> Nesting lore -> Ordering
forall lore. Nesting lore -> Nesting lore -> Nesting lore
min :: Nesting lore -> Nesting lore -> Nesting lore
$cmin :: forall lore. Nesting lore -> Nesting lore -> Nesting lore
max :: Nesting lore -> Nesting lore -> Nesting lore
$cmax :: forall lore. Nesting lore -> Nesting lore -> Nesting lore
>= :: Nesting lore -> Nesting lore -> Bool
$c>= :: forall lore. Nesting lore -> Nesting lore -> Bool
> :: Nesting lore -> Nesting lore -> Bool
$c> :: forall lore. Nesting lore -> Nesting lore -> Bool
<= :: Nesting lore -> Nesting lore -> Bool
$c<= :: forall lore. Nesting lore -> Nesting lore -> Bool
< :: Nesting lore -> Nesting lore -> Bool
$c< :: forall lore. Nesting lore -> Nesting lore -> Bool
compare :: Nesting lore -> Nesting lore -> Ordering
$ccompare :: forall lore. Nesting lore -> Nesting lore -> Ordering
Ord, Int -> Nesting lore -> ShowS
[Nesting lore] -> ShowS
Nesting lore -> String
(Int -> Nesting lore -> ShowS)
-> (Nesting lore -> String)
-> ([Nesting lore] -> ShowS)
-> Show (Nesting lore)
forall lore. Int -> Nesting lore -> ShowS
forall lore. [Nesting lore] -> ShowS
forall lore. Nesting lore -> String
forall a.
(Int -> a -> ShowS) -> (a -> String) -> ([a] -> ShowS) -> Show a
showList :: [Nesting lore] -> ShowS
$cshowList :: forall lore. [Nesting lore] -> ShowS
show :: Nesting lore -> String
$cshow :: forall lore. Nesting lore -> String
showsPrec :: Int -> Nesting lore -> ShowS
$cshowsPrec :: forall lore. Int -> Nesting lore -> ShowS
Show)

data MapNest lore = MapNest SubExp (Lambda lore) [Nesting lore] [SOAC.Input]
  deriving (Int -> MapNest lore -> ShowS
[MapNest lore] -> ShowS
MapNest lore -> String
(Int -> MapNest lore -> ShowS)
-> (MapNest lore -> String)
-> ([MapNest lore] -> ShowS)
-> Show (MapNest lore)
forall lore. Decorations lore => Int -> MapNest lore -> ShowS
forall lore. Decorations lore => [MapNest lore] -> ShowS
forall lore. Decorations lore => MapNest lore -> String
forall a.
(Int -> a -> ShowS) -> (a -> String) -> ([a] -> ShowS) -> Show a
showList :: [MapNest lore] -> ShowS
$cshowList :: forall lore. Decorations lore => [MapNest lore] -> ShowS
show :: MapNest lore -> String
$cshow :: forall lore. Decorations lore => MapNest lore -> String
showsPrec :: Int -> MapNest lore -> ShowS
$cshowsPrec :: forall lore. Decorations lore => Int -> MapNest lore -> ShowS
Show)

typeOf :: MapNest lore -> [Type]
typeOf :: forall lore. MapNest lore -> [Type]
typeOf (MapNest SubExp
w Lambda lore
lam [] [Input]
_) =
  (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
w) ([Type] -> [Type]) -> [Type] -> [Type]
forall a b. (a -> b) -> a -> b
$ Lambda lore -> [Type]
forall lore. LambdaT lore -> [Type]
lambdaReturnType Lambda lore
lam
typeOf (MapNest SubExp
w Lambda lore
_ (Nesting lore
nest : [Nesting lore]
_) [Input]
_) =
  (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
w) ([Type] -> [Type]) -> [Type] -> [Type]
forall a b. (a -> b) -> a -> b
$ Nesting lore -> [Type]
forall lore. Nesting lore -> [Type]
nestingReturnType Nesting lore
nest

params :: MapNest lore -> [VName]
params :: forall lore. MapNest lore -> [VName]
params (MapNest SubExp
_ Lambda lore
lam [] [Input]
_) =
  (Param (LParamInfo lore) -> VName)
-> [Param (LParamInfo lore)] -> [VName]
forall a b. (a -> b) -> [a] -> [b]
map Param (LParamInfo lore) -> VName
forall dec. Param dec -> VName
paramName ([Param (LParamInfo lore)] -> [VName])
-> [Param (LParamInfo lore)] -> [VName]
forall a b. (a -> b) -> a -> b
$ Lambda lore -> [Param (LParamInfo lore)]
forall lore. LambdaT lore -> [LParam lore]
lambdaParams Lambda lore
lam
params (MapNest SubExp
_ Lambda lore
_ (Nesting lore
nest : [Nesting lore]
_) [Input]
_) =
  Nesting lore -> [VName]
forall lore. Nesting lore -> [VName]
nestingParamNames Nesting lore
nest

inputs :: MapNest lore -> [SOAC.Input]
inputs :: forall lore. MapNest lore -> [Input]
inputs (MapNest SubExp
_ Lambda lore
_ [Nesting lore]
_ [Input]
inps) = [Input]
inps

setInputs :: [SOAC.Input] -> MapNest lore -> MapNest lore
setInputs :: forall lore. [Input] -> MapNest lore -> MapNest lore
setInputs [] (MapNest SubExp
w Lambda lore
body [Nesting lore]
ns [Input]
_) = SubExp -> Lambda lore -> [Nesting lore] -> [Input] -> MapNest lore
forall lore.
SubExp -> Lambda lore -> [Nesting lore] -> [Input] -> MapNest lore
MapNest SubExp
w Lambda lore
body [Nesting lore]
ns []
setInputs (Input
inp : [Input]
inps) (MapNest SubExp
_ Lambda lore
body [Nesting lore]
ns [Input]
_) = SubExp -> Lambda lore -> [Nesting lore] -> [Input] -> MapNest lore
forall lore.
SubExp -> Lambda lore -> [Nesting lore] -> [Input] -> MapNest lore
MapNest SubExp
w Lambda lore
body [Nesting lore]
ns' (Input
inp Input -> [Input] -> [Input]
forall a. a -> [a] -> [a]
: [Input]
inps)
  where
    w :: SubExp
w = Int -> Type -> SubExp
forall u. Int -> TypeBase (ShapeBase SubExp) u -> SubExp
arraySize Int
0 (Type -> SubExp) -> Type -> SubExp
forall a b. (a -> b) -> a -> b
$ Input -> Type
SOAC.inputType Input
inp
    ws :: [SubExp]
ws = Int -> [SubExp] -> [SubExp]
forall a. Int -> [a] -> [a]
drop Int
1 ([SubExp] -> [SubExp]) -> [SubExp] -> [SubExp]
forall a b. (a -> b) -> a -> b
$ Type -> [SubExp]
forall u. TypeBase (ShapeBase SubExp) u -> [SubExp]
arrayDims (Type -> [SubExp]) -> Type -> [SubExp]
forall a b. (a -> b) -> a -> b
$ Input -> Type
SOAC.inputType Input
inp
    ns' :: [Nesting lore]
ns' = (Nesting lore -> SubExp -> Nesting lore)
-> [Nesting lore] -> [SubExp] -> [Nesting lore]
forall a b c. (a -> b -> c) -> [a] -> [b] -> [c]
zipWith Nesting lore -> SubExp -> Nesting lore
forall {lore} {lore}. Nesting lore -> SubExp -> Nesting lore
setDepth [Nesting lore]
ns [SubExp]
ws
    setDepth :: Nesting lore -> SubExp -> Nesting lore
setDepth Nesting lore
n SubExp
nw = Nesting lore
n {nestingWidth :: SubExp
nestingWidth = SubExp
nw}

fromSOAC ::
  ( Bindable lore,
    MonadFreshNames m,
    LocalScope lore m,
    Op lore ~ Futhark.SOAC lore
  ) =>
  SOAC lore ->
  m (Maybe (MapNest lore))
fromSOAC :: forall lore (m :: * -> *).
(Bindable lore, MonadFreshNames m, LocalScope lore m,
 Op lore ~ SOAC lore) =>
SOAC lore -> m (Maybe (MapNest lore))
fromSOAC = [Ident] -> SOAC lore -> m (Maybe (MapNest lore))
forall lore (m :: * -> *).
(Bindable lore, MonadFreshNames m, LocalScope lore m,
 Op lore ~ SOAC lore) =>
[Ident] -> SOAC lore -> m (Maybe (MapNest lore))
fromSOAC' [Ident]
forall a. Monoid a => a
mempty

fromSOAC' ::
  ( Bindable lore,
    MonadFreshNames m,
    LocalScope lore m,
    Op lore ~ Futhark.SOAC lore
  ) =>
  [Ident] ->
  SOAC lore ->
  m (Maybe (MapNest lore))
fromSOAC' :: forall lore (m :: * -> *).
(Bindable lore, MonadFreshNames m, LocalScope lore m,
 Op lore ~ SOAC lore) =>
[Ident] -> SOAC lore -> m (Maybe (MapNest lore))
fromSOAC' [Ident]
bound (SOAC.Screma SubExp
w (SOAC.ScremaForm [] [] Lambda lore
lam) [Input]
inps) = do
  Either NotSOAC (Maybe (PatternT (LetDec lore), MapNest lore))
maybenest <- case ( Stms lore -> [Stm lore]
forall lore. Stms lore -> [Stm lore]
stmsToList (Stms lore -> [Stm lore]) -> Stms lore -> [Stm lore]
forall a b. (a -> b) -> a -> b
$ BodyT lore -> Stms lore
forall lore. BodyT lore -> Stms lore
bodyStms (BodyT lore -> Stms lore) -> BodyT lore -> Stms lore
forall a b. (a -> b) -> a -> b
$ Lambda lore -> BodyT lore
forall lore. LambdaT lore -> BodyT lore
lambdaBody Lambda lore
lam,
                      BodyT lore -> [SubExp]
forall lore. BodyT lore -> [SubExp]
bodyResult (BodyT lore -> [SubExp]) -> BodyT lore -> [SubExp]
forall a b. (a -> b) -> a -> b
$ Lambda lore -> BodyT lore
forall lore. LambdaT lore -> BodyT lore
lambdaBody Lambda lore
lam
                    ) of
    ([Let PatternT (LetDec lore)
pat StmAux (ExpDec lore)
_ Exp lore
e], [SubExp]
res)
      | [SubExp]
res [SubExp] -> [SubExp] -> Bool
forall a. Eq a => a -> a -> Bool
== (VName -> SubExp) -> [VName] -> [SubExp]
forall a b. (a -> b) -> [a] -> [b]
map VName -> SubExp
Var (PatternT (LetDec lore) -> [VName]
forall dec. PatternT dec -> [VName]
patternNames PatternT (LetDec lore)
pat) ->
        Scope lore
-> m (Either
        NotSOAC (Maybe (PatternT (LetDec lore), MapNest lore)))
-> m (Either
        NotSOAC (Maybe (PatternT (LetDec lore), MapNest lore)))
forall lore (m :: * -> *) a.
LocalScope lore m =>
Scope lore -> m a -> m a
localScope ([Param Type] -> Scope lore
forall lore dec.
(LParamInfo lore ~ dec) =>
[Param dec] -> Scope lore
scopeOfLParams ([Param Type] -> Scope lore) -> [Param Type] -> Scope lore
forall a b. (a -> b) -> a -> b
$ Lambda lore -> [LParam lore]
forall lore. LambdaT lore -> [LParam lore]
lambdaParams Lambda lore
lam) (m (Either NotSOAC (Maybe (PatternT (LetDec lore), MapNest lore)))
 -> m (Either
         NotSOAC (Maybe (PatternT (LetDec lore), MapNest lore))))
-> m (Either
        NotSOAC (Maybe (PatternT (LetDec lore), MapNest lore)))
-> m (Either
        NotSOAC (Maybe (PatternT (LetDec lore), MapNest lore)))
forall a b. (a -> b) -> a -> b
$
          Exp lore -> m (Either NotSOAC (SOAC lore))
forall lore (m :: * -> *).
(Op lore ~ SOAC lore, HasScope lore m) =>
Exp lore -> m (Either NotSOAC (SOAC lore))
SOAC.fromExp Exp lore
e
            m (Either NotSOAC (SOAC lore))
-> (Either NotSOAC (SOAC lore)
    -> m (Either
            NotSOAC (Maybe (PatternT (LetDec lore), MapNest lore))))
-> m (Either
        NotSOAC (Maybe (PatternT (LetDec lore), MapNest lore)))
forall (m :: * -> *) a b. Monad m => m a -> (a -> m b) -> m b
>>= (NotSOAC
 -> m (Either
         NotSOAC (Maybe (PatternT (LetDec lore), MapNest lore))))
-> (SOAC lore
    -> m (Either
            NotSOAC (Maybe (PatternT (LetDec lore), MapNest lore))))
-> Either NotSOAC (SOAC lore)
-> m (Either
        NotSOAC (Maybe (PatternT (LetDec lore), MapNest lore)))
forall a c b. (a -> c) -> (b -> c) -> Either a b -> c
either (Either NotSOAC (Maybe (PatternT (LetDec lore), MapNest lore))
-> m (Either
        NotSOAC (Maybe (PatternT (LetDec lore), MapNest lore)))
forall (m :: * -> *) a. Monad m => a -> m a
return (Either NotSOAC (Maybe (PatternT (LetDec lore), MapNest lore))
 -> m (Either
         NotSOAC (Maybe (PatternT (LetDec lore), MapNest lore))))
-> (NotSOAC
    -> Either NotSOAC (Maybe (PatternT (LetDec lore), MapNest lore)))
-> NotSOAC
-> m (Either
        NotSOAC (Maybe (PatternT (LetDec lore), MapNest lore)))
forall b c a. (b -> c) -> (a -> b) -> a -> c
. NotSOAC
-> Either NotSOAC (Maybe (PatternT (LetDec lore), MapNest lore))
forall a b. a -> Either a b
Left) ((Maybe (MapNest lore)
 -> Either NotSOAC (Maybe (PatternT (LetDec lore), MapNest lore)))
-> m (Maybe (MapNest lore))
-> m (Either
        NotSOAC (Maybe (PatternT (LetDec lore), MapNest lore)))
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap (Maybe (PatternT (LetDec lore), MapNest lore)
-> Either NotSOAC (Maybe (PatternT (LetDec lore), MapNest lore))
forall a b. b -> Either a b
Right (Maybe (PatternT (LetDec lore), MapNest lore)
 -> Either NotSOAC (Maybe (PatternT (LetDec lore), MapNest lore)))
-> (Maybe (MapNest lore)
    -> Maybe (PatternT (LetDec lore), MapNest lore))
-> Maybe (MapNest lore)
-> Either NotSOAC (Maybe (PatternT (LetDec lore), MapNest lore))
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (MapNest lore -> (PatternT (LetDec lore), MapNest lore))
-> Maybe (MapNest lore)
-> Maybe (PatternT (LetDec lore), MapNest lore)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap (PatternT (LetDec lore)
pat,)) (m (Maybe (MapNest lore))
 -> m (Either
         NotSOAC (Maybe (PatternT (LetDec lore), MapNest lore))))
-> (SOAC lore -> m (Maybe (MapNest lore)))
-> SOAC lore
-> m (Either
        NotSOAC (Maybe (PatternT (LetDec lore), MapNest lore)))
forall b c a. (b -> c) -> (a -> b) -> a -> c
. [Ident] -> SOAC lore -> m (Maybe (MapNest lore))
forall lore (m :: * -> *).
(Bindable lore, MonadFreshNames m, LocalScope lore m,
 Op lore ~ SOAC lore) =>
[Ident] -> SOAC lore -> m (Maybe (MapNest lore))
fromSOAC' [Ident]
bound')
    ([Stm lore], [SubExp])
_ ->
      Either NotSOAC (Maybe (PatternT (LetDec lore), MapNest lore))
-> m (Either
        NotSOAC (Maybe (PatternT (LetDec lore), MapNest lore)))
forall (m :: * -> *) a. Monad m => a -> m a
return (Either NotSOAC (Maybe (PatternT (LetDec lore), MapNest lore))
 -> m (Either
         NotSOAC (Maybe (PatternT (LetDec lore), MapNest lore))))
-> Either NotSOAC (Maybe (PatternT (LetDec lore), MapNest lore))
-> m (Either
        NotSOAC (Maybe (PatternT (LetDec lore), MapNest lore)))
forall a b. (a -> b) -> a -> b
$ Maybe (PatternT (LetDec lore), MapNest lore)
-> Either NotSOAC (Maybe (PatternT (LetDec lore), MapNest lore))
forall a b. b -> Either a b
Right Maybe (PatternT (LetDec lore), MapNest lore)
forall a. Maybe a
Nothing

  case Either NotSOAC (Maybe (PatternT (LetDec lore), MapNest lore))
maybenest of
    -- Do we have a nested MapNest?
    Right (Just (PatternT (LetDec lore)
pat, mn :: MapNest lore
mn@(MapNest SubExp
inner_w Lambda lore
body' [Nesting lore]
ns' [Input]
inps'))) -> do
      ([VName]
ps, [Input]
inps'') <-
        [(VName, Input)] -> ([VName], [Input])
forall a b. [(a, b)] -> ([a], [b])
unzip
          ([(VName, Input)] -> ([VName], [Input]))
-> m [(VName, Input)] -> m ([VName], [Input])
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> SubExp
-> [(VName, Input)] -> [(VName, Input)] -> m [(VName, Input)]
forall (m :: * -> *).
MonadFreshNames m =>
SubExp
-> [(VName, Input)] -> [(VName, Input)] -> m [(VName, Input)]
fixInputs
            SubExp
w
            ([VName] -> [Input] -> [(VName, Input)]
forall a b. [a] -> [b] -> [(a, b)]
zip ((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] -> [VName]) -> [Param Type] -> [VName]
forall a b. (a -> b) -> a -> b
$ Lambda lore -> [LParam lore]
forall lore. LambdaT lore -> [LParam lore]
lambdaParams Lambda lore
lam) [Input]
inps)
            ([VName] -> [Input] -> [(VName, Input)]
forall a b. [a] -> [b] -> [(a, b)]
zip (MapNest lore -> [VName]
forall lore. MapNest lore -> [VName]
params MapNest lore
mn) [Input]
inps')
      let n' :: Nesting lore
n' =
            Nesting :: forall lore. [VName] -> [VName] -> [Type] -> SubExp -> Nesting lore
Nesting
              { nestingParamNames :: [VName]
nestingParamNames = [VName]
ps,
                nestingResult :: [VName]
nestingResult = PatternT (LetDec lore) -> [VName]
forall dec. PatternT dec -> [VName]
patternNames PatternT (LetDec lore)
pat,
                nestingReturnType :: [Type]
nestingReturnType = MapNest lore -> [Type]
forall lore. MapNest lore -> [Type]
typeOf MapNest lore
mn,
                nestingWidth :: SubExp
nestingWidth = SubExp
inner_w
              }
      Maybe (MapNest lore) -> m (Maybe (MapNest lore))
forall (m :: * -> *) a. Monad m => a -> m a
return (Maybe (MapNest lore) -> m (Maybe (MapNest lore)))
-> Maybe (MapNest lore) -> m (Maybe (MapNest lore))
forall a b. (a -> b) -> a -> b
$ MapNest lore -> Maybe (MapNest lore)
forall a. a -> Maybe a
Just (MapNest lore -> Maybe (MapNest lore))
-> MapNest lore -> Maybe (MapNest lore)
forall a b. (a -> b) -> a -> b
$ SubExp -> Lambda lore -> [Nesting lore] -> [Input] -> MapNest lore
forall lore.
SubExp -> Lambda lore -> [Nesting lore] -> [Input] -> MapNest lore
MapNest SubExp
w Lambda lore
body' (Nesting lore
n' Nesting lore -> [Nesting lore] -> [Nesting lore]
forall a. a -> [a] -> [a]
: [Nesting lore]
ns') [Input]
inps''
    -- No nested MapNest it seems.
    Either NotSOAC (Maybe (PatternT (LetDec lore), MapNest lore))
_ -> do
      let isBound :: VName -> Maybe Ident
isBound VName
name
            | Just Ident
param <- (Ident -> Bool) -> [Ident] -> Maybe Ident
forall (t :: * -> *) a. Foldable t => (a -> Bool) -> t a -> Maybe a
find ((VName
name VName -> VName -> Bool
forall a. Eq a => a -> a -> Bool
==) (VName -> Bool) -> (Ident -> VName) -> Ident -> Bool
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Ident -> VName
identName) [Ident]
bound =
              Ident -> Maybe Ident
forall a. a -> Maybe a
Just Ident
param
            | Bool
otherwise =
              Maybe Ident
forall a. Maybe a
Nothing
          boundUsedInBody :: [Ident]
boundUsedInBody =
            (VName -> Maybe Ident) -> [VName] -> [Ident]
forall a b. (a -> Maybe b) -> [a] -> [b]
mapMaybe VName -> Maybe Ident
isBound ([VName] -> [Ident]) -> [VName] -> [Ident]
forall a b. (a -> b) -> a -> b
$ Names -> [VName]
namesToList (Names -> [VName]) -> Names -> [VName]
forall a b. (a -> b) -> a -> b
$ Lambda lore -> Names
forall a. FreeIn a => a -> Names
freeIn Lambda lore
lam
      [Ident]
newParams <- (Ident -> m Ident) -> [Ident] -> m [Ident]
forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
(a -> m b) -> t a -> m (t b)
mapM (ShowS -> Ident -> m Ident
forall (m :: * -> *).
MonadFreshNames m =>
ShowS -> Ident -> m Ident
newIdent' (String -> ShowS
forall a. [a] -> [a] -> [a]
++ String
"_wasfree")) [Ident]
boundUsedInBody
      let subst :: Map VName VName
subst =
            [(VName, VName)] -> Map VName VName
forall k a. Ord k => [(k, a)] -> Map k a
M.fromList ([(VName, VName)] -> Map VName VName)
-> [(VName, VName)] -> Map VName VName
forall a b. (a -> b) -> a -> b
$
              [VName] -> [VName] -> [(VName, VName)]
forall a b. [a] -> [b] -> [(a, b)]
zip ((Ident -> VName) -> [Ident] -> [VName]
forall a b. (a -> b) -> [a] -> [b]
map Ident -> VName
identName [Ident]
boundUsedInBody) ((Ident -> VName) -> [Ident] -> [VName]
forall a b. (a -> b) -> [a] -> [b]
map Ident -> VName
identName [Ident]
newParams)
          inps' :: [Input]
inps' =
            [Input]
inps
              [Input] -> [Input] -> [Input]
forall a. [a] -> [a] -> [a]
++ (Ident -> Input) -> [Ident] -> [Input]
forall a b. (a -> b) -> [a] -> [b]
map
                (ArrayTransform -> Input -> Input
SOAC.addTransform (Certificates -> ShapeBase SubExp -> ArrayTransform
SOAC.Replicate Certificates
forall a. Monoid a => a
mempty (ShapeBase SubExp -> ArrayTransform)
-> ShapeBase SubExp -> ArrayTransform
forall a b. (a -> b) -> a -> b
$ [SubExp] -> ShapeBase SubExp
forall d. [d] -> ShapeBase d
Shape [SubExp
w]) (Input -> Input) -> (Ident -> Input) -> Ident -> Input
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Ident -> Input
SOAC.identInput)
                [Ident]
boundUsedInBody
          lam' :: Lambda lore
lam' =
            Lambda lore
lam
              { lambdaBody :: BodyT lore
lambdaBody =
                  Map VName VName -> BodyT lore -> BodyT lore
forall a. Substitute a => Map VName VName -> a -> a
substituteNames Map VName VName
subst (BodyT lore -> BodyT lore) -> BodyT lore -> BodyT lore
forall a b. (a -> b) -> a -> b
$ Lambda lore -> BodyT lore
forall lore. LambdaT lore -> BodyT lore
lambdaBody Lambda lore
lam,
                lambdaParams :: [LParam lore]
lambdaParams =
                  Lambda lore -> [LParam lore]
forall lore. LambdaT lore -> [LParam lore]
lambdaParams Lambda lore
lam
                    [Param Type] -> [Param Type] -> [Param Type]
forall a. [a] -> [a] -> [a]
++ [ VName -> Type -> Param Type
forall dec. VName -> dec -> Param dec
Param VName
name Type
t
                         | Ident VName
name Type
t <- [Ident]
newParams
                       ]
              }
      Maybe (MapNest lore) -> m (Maybe (MapNest lore))
forall (m :: * -> *) a. Monad m => a -> m a
return (Maybe (MapNest lore) -> m (Maybe (MapNest lore)))
-> Maybe (MapNest lore) -> m (Maybe (MapNest lore))
forall a b. (a -> b) -> a -> b
$ MapNest lore -> Maybe (MapNest lore)
forall a. a -> Maybe a
Just (MapNest lore -> Maybe (MapNest lore))
-> MapNest lore -> Maybe (MapNest lore)
forall a b. (a -> b) -> a -> b
$ SubExp -> Lambda lore -> [Nesting lore] -> [Input] -> MapNest lore
forall lore.
SubExp -> Lambda lore -> [Nesting lore] -> [Input] -> MapNest lore
MapNest SubExp
w Lambda lore
lam' [] [Input]
inps'
  where
    bound' :: [Ident]
bound' = [Ident]
bound [Ident] -> [Ident] -> [Ident]
forall a. Semigroup a => a -> a -> a
<> (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 (Lambda lore -> [LParam lore]
forall lore. LambdaT lore -> [LParam lore]
lambdaParams Lambda lore
lam)
fromSOAC' [Ident]
_ SOAC lore
_ = Maybe (MapNest lore) -> m (Maybe (MapNest lore))
forall (m :: * -> *) a. Monad m => a -> m a
return Maybe (MapNest lore)
forall a. Maybe a
Nothing

toSOAC ::
  ( MonadFreshNames m,
    HasScope lore m,
    Bindable lore,
    BinderOps lore,
    Op lore ~ Futhark.SOAC lore
  ) =>
  MapNest lore ->
  m (SOAC lore)
toSOAC :: forall (m :: * -> *) lore.
(MonadFreshNames m, HasScope lore m, Bindable lore, BinderOps lore,
 Op lore ~ SOAC lore) =>
MapNest lore -> m (SOAC lore)
toSOAC (MapNest SubExp
w Lambda lore
lam [] [Input]
inps) =
  SOAC lore -> m (SOAC lore)
forall (m :: * -> *) a. Monad m => a -> m a
return (SOAC lore -> m (SOAC lore)) -> SOAC lore -> m (SOAC lore)
forall a b. (a -> b) -> a -> b
$ SubExp -> ScremaForm lore -> [Input] -> SOAC lore
forall lore. SubExp -> ScremaForm lore -> [Input] -> SOAC lore
SOAC.Screma SubExp
w (Lambda lore -> ScremaForm lore
forall lore. Lambda lore -> ScremaForm lore
Futhark.mapSOAC Lambda lore
lam) [Input]
inps
toSOAC (MapNest SubExp
w Lambda lore
lam (Nesting [VName]
npnames [VName]
nres [Type]
nrettype SubExp
nw : [Nesting lore]
ns) [Input]
inps) = do
  let nparams :: [Param Type]
nparams = (VName -> Type -> Param Type) -> [VName] -> [Type] -> [Param Type]
forall a b c. (a -> b -> c) -> [a] -> [b] -> [c]
zipWith VName -> Type -> Param Type
forall dec. VName -> dec -> Param dec
Param [VName]
npnames ([Type] -> [Param Type]) -> [Type] -> [Param Type]
forall a b. (a -> b) -> a -> b
$ (Input -> Type) -> [Input] -> [Type]
forall a b. (a -> b) -> [a] -> [b]
map Input -> Type
SOAC.inputRowType [Input]
inps
  Body lore
body <- Binder lore (Body lore) -> m (Body lore)
forall lore (m :: * -> *) somelore.
(Bindable lore, MonadFreshNames m, HasScope somelore m,
 SameScope somelore lore) =>
Binder lore (Body lore) -> m (Body lore)
runBodyBinder (Binder lore (Body lore) -> m (Body lore))
-> Binder lore (Body lore) -> m (Body lore)
forall a b. (a -> b) -> a -> b
$
    Scope lore -> Binder lore (Body lore) -> Binder lore (Body lore)
forall lore (m :: * -> *) a.
LocalScope lore m =>
Scope lore -> m a -> m a
localScope ([Param Type] -> Scope lore
forall lore dec.
(LParamInfo lore ~ dec) =>
[Param dec] -> Scope lore
scopeOfLParams [Param Type]
nparams) (Binder lore (Body lore) -> Binder lore (Body lore))
-> Binder lore (Body lore) -> Binder lore (Body lore)
forall a b. (a -> b) -> a -> b
$ do
      [VName]
-> Exp (Lore (BinderT lore (State VNameSource)))
-> BinderT lore (State VNameSource) ()
forall (m :: * -> *).
MonadBinder m =>
[VName] -> Exp (Lore m) -> m ()
letBindNames [VName]
nres (ExpT lore -> BinderT lore (State VNameSource) ())
-> BinderT lore (State VNameSource) (ExpT lore)
-> BinderT lore (State VNameSource) ()
forall (m :: * -> *) a b. Monad m => (a -> m b) -> m a -> m b
=<< SOAC lore -> BinderT lore (State VNameSource) (ExpT lore)
forall (m :: * -> *).
(MonadBinder m, Op (Lore m) ~ SOAC (Lore m)) =>
SOAC (Lore m) -> m (Exp (Lore m))
SOAC.toExp
        (SOAC lore -> BinderT lore (State VNameSource) (ExpT lore))
-> BinderT lore (State VNameSource) (SOAC lore)
-> BinderT lore (State VNameSource) (ExpT lore)
forall (m :: * -> *) a b. Monad m => (a -> m b) -> m a -> m b
=<< MapNest lore -> BinderT lore (State VNameSource) (SOAC lore)
forall (m :: * -> *) lore.
(MonadFreshNames m, HasScope lore m, Bindable lore, BinderOps lore,
 Op lore ~ SOAC lore) =>
MapNest lore -> m (SOAC lore)
toSOAC (SubExp -> Lambda lore -> [Nesting lore] -> [Input] -> MapNest lore
forall lore.
SubExp -> Lambda lore -> [Nesting lore] -> [Input] -> MapNest lore
MapNest SubExp
nw Lambda lore
lam [Nesting lore]
ns ([Input] -> MapNest lore) -> [Input] -> MapNest lore
forall a b. (a -> b) -> a -> b
$ (Param Type -> Input) -> [Param Type] -> [Input]
forall a b. (a -> b) -> [a] -> [b]
map (Ident -> Input
SOAC.identInput (Ident -> Input) -> (Param Type -> Ident) -> Param Type -> Input
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Param Type -> Ident
forall dec. Typed dec => Param dec -> Ident
paramIdent) [Param Type]
nparams)
      Body lore -> Binder lore (Body lore)
forall (m :: * -> *) a. Monad m => a -> m a
return (Body lore -> Binder lore (Body lore))
-> Body lore -> Binder lore (Body lore)
forall a b. (a -> b) -> a -> b
$ [SubExp] -> Body lore
forall lore. Bindable lore => [SubExp] -> Body lore
resultBody ([SubExp] -> Body lore) -> [SubExp] -> Body lore
forall a b. (a -> b) -> a -> b
$ (VName -> SubExp) -> [VName] -> [SubExp]
forall a b. (a -> b) -> [a] -> [b]
map VName -> SubExp
Var [VName]
nres
  let outerlam :: Lambda lore
outerlam =
        Lambda :: forall lore. [LParam lore] -> BodyT lore -> [Type] -> LambdaT lore
Lambda
          { lambdaParams :: [LParam lore]
lambdaParams = [Param Type]
[LParam lore]
nparams,
            lambdaBody :: Body lore
lambdaBody = Body lore
body,
            lambdaReturnType :: [Type]
lambdaReturnType = [Type]
nrettype
          }
  SOAC lore -> m (SOAC lore)
forall (m :: * -> *) a. Monad m => a -> m a
return (SOAC lore -> m (SOAC lore)) -> SOAC lore -> m (SOAC lore)
forall a b. (a -> b) -> a -> b
$ SubExp -> ScremaForm lore -> [Input] -> SOAC lore
forall lore. SubExp -> ScremaForm lore -> [Input] -> SOAC lore
SOAC.Screma SubExp
w (Lambda lore -> ScremaForm lore
forall lore. Lambda lore -> ScremaForm lore
Futhark.mapSOAC Lambda lore
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 = ((VName, Input) -> m (VName, Input))
-> [(VName, Input)] -> m [(VName, Input)]
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 a -> a -> Bool
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) <- ((VName, Input) -> Bool)
-> [(VName, Input)] -> Maybe (VName, Input)
forall (t :: * -> *) a. Foldable t => (a -> Bool) -> t a -> Maybe a
find (VName -> (VName, Input) -> Bool
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' <- String -> m VName
forall (m :: * -> *). MonadFreshNames m => String -> m VName
newNameFromString (String -> m VName) -> String -> m VName
forall a b. (a -> b) -> a -> b
$ VName -> String
baseString VName
p
        (VName, Input) -> m (VName, Input)
forall (m :: * -> *) a. Monad m => a -> m a
return (VName
p', Input
pInp')
    inspect (VName
param, SOAC.Input ArrayTransforms
ts VName
a Type
t) = do
      VName
param' <- String -> m VName
forall (m :: * -> *). MonadFreshNames m => String -> m VName
newNameFromString (VName -> String
baseString VName
param String -> ShowS
forall a. [a] -> [a] -> [a]
++ String
"_rep")
      (VName, Input) -> m (VName, Input)
forall (m :: * -> *) a. Monad m => a -> m a
return (VName
param', ArrayTransforms -> VName -> Type -> Input
SOAC.Input (ArrayTransforms
ts ArrayTransforms -> ArrayTransform -> ArrayTransforms
SOAC.|> Certificates -> ShapeBase SubExp -> ArrayTransform
SOAC.Replicate Certificates
forall a. Monoid a => a
mempty ([SubExp] -> ShapeBase SubExp
forall d. [d] -> ShapeBase d
Shape [SubExp
w])) VName
a Type
t)