-- | Various definitions used for both forward and reverse mode.
module Futhark.AD.Shared
  ( vecPerm,
    asVName,
    mapNest,
    mkMap,
    accAddLambda,
  )
where

import Control.Monad
import Data.Bifunctor (bimap)
import Data.Foldable
import Futhark.Construct
import Futhark.IR.SOACS
import Futhark.Tools

-- | A permutation for transposing the vector shape past the next dimension.
--
-- That is, converts @[vec...][d][elem...]@ to @[d][vec...][elem...]@.
vecPerm :: (ArrayShape s) => Shape -> TypeBase s o -> [Int]
vecPerm :: forall s o. ArrayShape s => Shape -> TypeBase s o -> [Int]
vecPerm Shape
vec_shape TypeBase s o
t =
  [Shape -> Int
forall a. ArrayShape a => a -> Int
shapeRank Shape
vec_shape]
    [Int] -> [Int] -> [Int]
forall a. [a] -> [a] -> [a]
++ [Int
0 .. Shape -> Int
forall a. ArrayShape a => a -> Int
shapeRank Shape
vec_shape Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
1]
    [Int] -> [Int] -> [Int]
forall a. [a] -> [a] -> [a]
++ [Shape -> Int
forall a. ArrayShape a => a -> Int
shapeRank Shape
vec_shape Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
1 .. TypeBase s o -> Int
forall shape o. ArrayShape shape => TypeBase shape o -> Int
arrayRank TypeBase s o
t Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
1]

asVName :: (MonadBuilder m) => SubExp -> m VName
asVName :: forall (m :: * -> *). MonadBuilder m => SubExp -> m VName
asVName (Var VName
v) = VName -> m VName
forall a. a -> m a
forall (f :: * -> *) a. Applicative f => a -> f a
pure VName
v
asVName (Constant PrimValue
x) = Name -> Exp (Rep m) -> m VName
forall (m :: * -> *).
MonadBuilder m =>
Name -> Exp (Rep m) -> m VName
letExp Name
"asv" (Exp (Rep m) -> m VName) -> Exp (Rep m) -> m VName
forall a b. (a -> b) -> a -> b
$ BasicOp -> Exp (Rep m)
forall rep. BasicOp -> Exp rep
BasicOp (BasicOp -> Exp (Rep m)) -> BasicOp -> Exp (Rep m)
forall a b. (a -> b) -> a -> b
$ SubExp -> BasicOp
SubExp (SubExp -> BasicOp) -> SubExp -> BasicOp
forall a b. (a -> b) -> a -> b
$ PrimValue -> SubExp
Constant PrimValue
x

-- | An addition operator for an accumulator whose index space has the given
-- rank and whose elements have the given types.  The operators for the
-- individual element types are fused horizontally, as an accumulator operator
-- must handle all of them at once.
accAddLambda ::
  (MonadBuilder m, Rep m ~ SOACS) =>
  Int ->
  [Type] ->
  m (Lambda SOACS)
accAddLambda :: forall (m :: * -> *).
(MonadBuilder m, Rep m ~ SOACS) =>
Int -> [Type] -> m (Lambda SOACS)
accAddLambda Int
n [Type]
ts = do
  lams <- (Type -> m (Lambda SOACS)) -> [Type] -> m [Lambda SOACS]
forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
(a -> m b) -> t a -> m (t b)
forall (m :: * -> *) a b. Monad m => (a -> m b) -> [a] -> m [b]
mapM Type -> m (Lambda (Rep m))
Type -> m (Lambda SOACS)
forall (m :: * -> *).
(OpC (Rep m) ~ SOAC, MonadBuilder m, Buildable (Rep m)) =>
Type -> m (Lambda (Rep m))
addLambda [Type]
ts
  idx_params <- replicateM n $ newParam "idx" $ Prim int64
  let (xs, ys) = bimap concat concat $ unzip $ map (splitAt 1 . lambdaParams) lams
  mkLambda (idx_params <> xs <> ys) $
    mconcat <$> mapM (bodyBind . lambdaBody) lams

mapNest ::
  (MonadBuilder m, Rep m ~ SOACS, Traversable f) =>
  Shape ->
  f SubExp ->
  (f SubExp -> m (Exp SOACS)) ->
  m (Exp SOACS)
mapNest :: forall (m :: * -> *) (f :: * -> *).
(MonadBuilder m, Rep m ~ SOACS, Traversable f) =>
Shape -> f SubExp -> (f SubExp -> m (Exp SOACS)) -> m (Exp SOACS)
mapNest Shape
shape f SubExp
x f SubExp -> m (Exp SOACS)
f = do
  if Shape
shape Shape -> Shape -> Bool
forall a. Eq a => a -> a -> Bool
== Shape
forall a. Monoid a => a
mempty
    then f SubExp -> m (Exp SOACS)
f f SubExp
x
    else do
      let w :: SubExp
w = Int -> Shape -> SubExp
shapeSize Int
0 Shape
shape
      x_v <- (SubExp -> m VName) -> f SubExp -> m (f VName)
forall (t :: * -> *) (f :: * -> *) a b.
(Traversable t, Applicative f) =>
(a -> f b) -> t a -> f (t b)
forall (f :: * -> *) a b.
Applicative f =>
(a -> f b) -> f a -> f (f b)
traverse SubExp -> m VName
forall (m :: * -> *). MonadBuilder m => SubExp -> m VName
asVName f SubExp
x
      x_p <- traverse (newParam "xp" . rowType <=< lookupType) x_v
      lam <- mkLambda (toList x_p) $ do
        fmap (subExpsRes . pure) . letSubExp "mapnest_res"
          =<< f (fmap (Var . paramName) x_p)
      Op . Screma w (toList x_v) <$> mapSOAC lam

-- | Construct a map over the given arrays, which must have the provided outer
-- shape. The purpose of the 'Shape' argument is to handle the case where no
-- arrays are provided.
mkMap ::
  (MonadBuilder m, Rep m ~ SOACS, Traversable f) =>
  Name ->
  Shape ->
  f VName ->
  -- | Action for building the body, passed names
  -- corresponding to elements of the arrays.
  (f VName -> m [VName]) ->
  m [VName]
mkMap :: forall (m :: * -> *) (f :: * -> *).
(MonadBuilder m, Rep m ~ SOACS, Traversable f) =>
Name -> Shape -> f VName -> (f VName -> m [VName]) -> m [VName]
mkMap Name
desc Shape
shape f VName
arrs f VName -> m [VName]
f = do
  let w :: SubExp
w = Int -> Shape -> SubExp
shapeSize Int
0 Shape
shape
  x_p <- (VName -> m (Param Type)) -> f VName -> m (f (Param Type))
forall (t :: * -> *) (f :: * -> *) a b.
(Traversable t, Applicative f) =>
(a -> f b) -> t a -> f (t b)
forall (f :: * -> *) a b.
Applicative f =>
(a -> f b) -> f a -> f (f b)
traverse (Name -> Type -> m (Param Type)
forall (m :: * -> *) dec.
MonadFreshNames m =>
Name -> dec -> m (Param dec)
newParam Name
"xp" (Type -> m (Param Type))
-> (Type -> Type) -> Type -> m (Param Type)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Type -> Type
forall d m. TypeBase (ShapeBase d) m -> TypeBase (ShapeBase d) m
rowType (Type -> m (Param Type))
-> (VName -> m Type) -> VName -> m (Param Type)
forall (m :: * -> *) b c a.
Monad m =>
(b -> m c) -> (a -> m b) -> a -> m c
<=< VName -> m Type
forall rep (m :: * -> *). HasScope rep m => VName -> m Type
lookupType) f VName
arrs
  lam <- mkLambda (toList x_p) $ varsRes <$> f (fmap paramName x_p)
  letTupExp desc . Op . Screma w (toList arrs) =<< mapSOAC lam