{-# LANGUAGE TypeFamilies #-}

-- | An unstructured grab-bag of various tools and inspection
-- functions that didn't really fit anywhere else.
module Futhark.Tools
  ( module Futhark.Construct,
    redomapToMapAndReduce,
    scanomapToMapAndScan,
    maposcanomapToMapScanAndMap,
    maposcanomapToMaposcanAndMap,
    dissectScrema,
    extractPostLambda,
    sequentialStreamWholeArray,
    partitionChunkedFoldParameters,
    withAcc,
    doScatter,
    doHist,
    addBinOp,
    addLambda,

    -- * Primitive expressions
    module Futhark.Analysis.PrimExp.Convert,
  )
where

import Control.Monad
import Data.List qualified as L
import Data.Maybe
import Futhark.Analysis.PrimExp.Convert
import Futhark.Construct
import Futhark.IR
import Futhark.IR.SOACS.SOAC
import Futhark.Util (chunks, mapAccumLM)

splitScanOrRedomap ::
  (MonadFreshNames m) =>
  [PatElem Type] ->
  SubExp ->
  Lambda rep ->
  [[SubExp]] ->
  m (Pat Type, Pat Type, [VName], Lambda rep)
splitScanOrRedomap :: forall (m :: * -> *) rep.
MonadFreshNames m =>
[PatElem Type]
-> SubExp
-> Lambda rep
-> [[SubExp]]
-> m (Pat Type, Pat Type, [VName], Lambda rep)
splitScanOrRedomap [PatElem Type]
pes SubExp
w Lambda rep
map_lam [[SubExp]]
nes = do
  let ([PatElem Type]
nonmap_pes, [PatElem Type]
map_pes) =
        Int -> [PatElem Type] -> ([PatElem Type], [PatElem Type])
forall a. Int -> [a] -> ([a], [a])
splitAt ([SubExp] -> Int
forall a. [a] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length ([SubExp] -> Int) -> [SubExp] -> Int
forall a b. (a -> b) -> a -> b
$ [[SubExp]] -> [SubExp]
forall (t :: * -> *) a. Foldable t => t [a] -> [a]
concat [[SubExp]]
nes) [PatElem Type]
pes
      ([Type]
nonmap_ts, [Type]
map_ts) =
        Int -> [Type] -> ([Type], [Type])
forall a. Int -> [a] -> ([a], [a])
splitAt ([SubExp] -> Int
forall a. [a] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length ([[SubExp]] -> [SubExp]
forall (t :: * -> *) a. Foldable t => t [a] -> [a]
concat [[SubExp]]
nes)) ([Type] -> ([Type], [Type])) -> [Type] -> ([Type], [Type])
forall a b. (a -> b) -> a -> b
$ Lambda rep -> [Type]
forall rep t. GLambda rep t -> [t]
lambdaReturnType Lambda rep
map_lam
      ([SubExpRes]
nonmap_res, [SubExpRes]
map_res) =
        Int -> [SubExpRes] -> ([SubExpRes], [SubExpRes])
forall a. Int -> [a] -> ([a], [a])
splitAt ([SubExp] -> Int
forall a. [a] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length ([[SubExp]] -> [SubExp]
forall (t :: * -> *) a. Foldable t => t [a] -> [a]
concat [[SubExp]]
nes)) ([SubExpRes] -> ([SubExpRes], [SubExpRes]))
-> [SubExpRes] -> ([SubExpRes], [SubExpRes])
forall a b. (a -> b) -> a -> b
$ GBody rep SubExpRes -> [SubExpRes]
forall rep res. GBody rep res -> [res]
bodyResult (GBody rep SubExpRes -> [SubExpRes])
-> GBody rep SubExpRes -> [SubExpRes]
forall a b. (a -> b) -> a -> b
$ Lambda rep -> GBody rep SubExpRes
forall rep t. GLambda rep t -> Body rep
lambdaBody Lambda rep
map_lam

  -- Put some care into not having duplicate results from the map function.
  (red_arrs, acc_info) <-
    [(VName, Maybe (PatElem Type, Type, SubExpRes))]
-> ([VName], [Maybe (PatElem Type, Type, SubExpRes)])
forall a b. [(a, b)] -> ([a], [b])
unzip ([(VName, Maybe (PatElem Type, Type, SubExpRes))]
 -> ([VName], [Maybe (PatElem Type, Type, SubExpRes)]))
-> (([(SubExpRes, PatElem Type)],
     [(VName, Maybe (PatElem Type, Type, SubExpRes))])
    -> [(VName, Maybe (PatElem Type, Type, SubExpRes))])
-> ([(SubExpRes, PatElem Type)],
    [(VName, Maybe (PatElem Type, Type, SubExpRes))])
-> ([VName], [Maybe (PatElem Type, Type, SubExpRes)])
forall b c a. (b -> c) -> (a -> b) -> a -> c
. ([(SubExpRes, PatElem Type)],
 [(VName, Maybe (PatElem Type, Type, SubExpRes))])
-> [(VName, Maybe (PatElem Type, Type, SubExpRes))]
forall a b. (a, b) -> b
snd
      (([(SubExpRes, PatElem Type)],
  [(VName, Maybe (PatElem Type, Type, SubExpRes))])
 -> ([VName], [Maybe (PatElem Type, Type, SubExpRes)]))
-> m ([(SubExpRes, PatElem Type)],
      [(VName, Maybe (PatElem Type, Type, SubExpRes))])
-> m ([VName], [Maybe (PatElem Type, Type, SubExpRes)])
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> ([(SubExpRes, PatElem Type)]
 -> (PatElem Type, Type, SubExpRes)
 -> m ([(SubExpRes, PatElem Type)],
       (VName, Maybe (PatElem Type, Type, SubExpRes))))
-> [(SubExpRes, PatElem Type)]
-> [(PatElem Type, Type, SubExpRes)]
-> m ([(SubExpRes, PatElem Type)],
      [(VName, Maybe (PatElem Type, Type, SubExpRes))])
forall (m :: * -> *) (t :: * -> *) acc x y.
(Monad m, Traversable t) =>
(acc -> x -> m (acc, y)) -> acc -> t x -> m (acc, t y)
mapAccumLM
        [(SubExpRes, PatElem Type)]
-> (PatElem Type, Type, SubExpRes)
-> m ([(SubExpRes, PatElem Type)],
      (VName, Maybe (PatElem Type, Type, SubExpRes)))
accMapPatElem
        ([SubExpRes] -> [PatElem Type] -> [(SubExpRes, PatElem Type)]
forall a b. [a] -> [b] -> [(a, b)]
zip [SubExpRes]
map_res [PatElem Type]
map_pes)
        ([PatElem Type]
-> [Type] -> [SubExpRes] -> [(PatElem Type, Type, SubExpRes)]
forall a b c. [a] -> [b] -> [c] -> [(a, b, c)]
zip3 [PatElem Type]
nonmap_pes [Type]
nonmap_ts [SubExpRes]
nonmap_res)

  let (nonmap_tmppes, nonmap_ts', nonmap_res') =
        L.unzip3 $ catMaybes acc_info
      map_lam' =
        Lambda rep
map_lam
          { lambdaBody = (lambdaBody map_lam) {bodyResult = nonmap_res' <> map_res},
            lambdaReturnType = nonmap_ts' <> map_ts
          }
      map_pat = [PatElem Type]
nonmap_tmppes [PatElem Type] -> [PatElem Type] -> [PatElem Type]
forall a. Semigroup a => a -> a -> a
<> [PatElem Type]
map_pes

  pure (Pat map_pat, Pat nonmap_pes, red_arrs, map_lam')
  where
    accMapPatElem :: [(SubExpRes, PatElem Type)]
-> (PatElem Type, Type, SubExpRes)
-> m ([(SubExpRes, PatElem Type)],
      (VName, Maybe (PatElem Type, Type, SubExpRes)))
accMapPatElem [(SubExpRes, PatElem Type)]
res_to_pe (PatElem Type
pe, Type
nonmap_t, SubExpRes
res) =
      case SubExpRes
res SubExpRes -> [(SubExpRes, PatElem Type)] -> Maybe (PatElem Type)
forall a b. Eq a => a -> [(a, b)] -> Maybe b
`L.lookup` [(SubExpRes, PatElem Type)]
res_to_pe of
        Just PatElem Type
pe' -> ([(SubExpRes, PatElem Type)],
 (VName, Maybe (PatElem Type, Type, SubExpRes)))
-> m ([(SubExpRes, PatElem Type)],
      (VName, Maybe (PatElem Type, Type, SubExpRes)))
forall a. a -> m a
forall (f :: * -> *) a. Applicative f => a -> f a
pure ([(SubExpRes, PatElem Type)]
res_to_pe, (PatElem Type -> VName
forall dec. PatElem dec -> VName
patElemName PatElem Type
pe', Maybe (PatElem Type, Type, SubExpRes)
forall a. Maybe a
Nothing))
        Maybe (PatElem Type)
Nothing -> do
          pe' <-
            VName -> Type -> PatElem Type
forall dec. VName -> dec -> PatElem dec
PatElem
              (VName -> Type -> PatElem Type)
-> m VName -> m (Type -> PatElem Type)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Name -> m VName
forall (m :: * -> *). MonadFreshNames m => Name -> m VName
newVName (VName -> Name
baseName (PatElem Type -> VName
forall dec. PatElem dec -> VName
patElemName PatElem Type
pe) Name -> Name -> Name
forall a. Semigroup a => a -> a -> a
<> Name
"_map_acc")
              m (Type -> PatElem Type) -> m Type -> m (PatElem Type)
forall a b. m (a -> b) -> m a -> m b
forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> Type -> m Type
forall a. a -> m a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (Type
nonmap_t Type -> SubExp -> Type
forall d.
TypeBase (ShapeBase d) NoMode -> d -> TypeBase (ShapeBase d) NoMode
`arrayOfRow` SubExp
w)
          pure ((res, pe') : res_to_pe, (patElemName pe', Just (pe', nonmap_t, res)))

-- | Turns a binding of a @redomap@ into two seperate bindings, a
-- @map@ binding and a @reduce@ binding (returned in that order).
--
-- Reuses the original pattern for the @reduce@, and creates a new
-- pattern with new 'Ident's for the result of the @map@.
redomapToMapAndReduce ::
  ( MonadFreshNames m,
    LetDec rep ~ Type,
    Buildable rep,
    ExpDec rep ~ (),
    Op rep ~ SOAC rep
  ) =>
  Pat (LetDec rep) ->
  ( SubExp,
    [Reduce rep],
    Lambda rep,
    [VName]
  ) ->
  m (Stm rep, Stm rep)
redomapToMapAndReduce :: forall (m :: * -> *) rep.
(MonadFreshNames m, LetDec rep ~ Type, Buildable rep,
 ExpDec rep ~ (), Op rep ~ SOAC rep) =>
Pat (LetDec rep)
-> (SubExp, [Reduce rep], Lambda rep, [VName])
-> m (Stm rep, Stm rep)
redomapToMapAndReduce (Pat [PatElem (LetDec rep)]
pes) (SubExp
w, [Reduce rep]
reds, Lambda rep
map_lam, [VName]
arrs) = do
  (map_pat, red_pat, red_arrs, map_lam') <-
    [PatElem Type]
-> SubExp
-> Lambda rep
-> [[SubExp]]
-> m (Pat Type, Pat Type, [VName], Lambda rep)
forall (m :: * -> *) rep.
MonadFreshNames m =>
[PatElem Type]
-> SubExp
-> Lambda rep
-> [[SubExp]]
-> m (Pat Type, Pat Type, [VName], Lambda rep)
splitScanOrRedomap [PatElem Type]
[PatElem (LetDec rep)]
pes SubExp
w Lambda rep
map_lam ([[SubExp]] -> m (Pat Type, Pat Type, [VName], Lambda rep))
-> [[SubExp]] -> m (Pat Type, Pat Type, [VName], Lambda rep)
forall a b. (a -> b) -> a -> b
$ (Reduce rep -> [SubExp]) -> [Reduce rep] -> [[SubExp]]
forall a b. (a -> b) -> [a] -> [b]
map Reduce rep -> [SubExp]
forall rep. Reduce rep -> [SubExp]
redNeutral [Reduce rep]
reds

  map_stm <- Let map_pat (defAux ()) . Op . Screma w arrs <$> mapSOAC map_lam'
  red_stm <-
    Let red_pat (defAux ()) . Op
      <$> (Screma w red_arrs <$> reduceSOAC reds)
  pure (map_stm, red_stm)

scanomapToMapAndScan ::
  ( MonadFreshNames m,
    Buildable rep,
    LetDec rep ~ Type,
    ExpDec rep ~ (),
    Op rep ~ SOAC rep
  ) =>
  Pat (LetDec rep) ->
  ( SubExp,
    [Scan rep],
    Lambda rep,
    [VName]
  ) ->
  m (Stm rep, Stm rep)
scanomapToMapAndScan :: forall (m :: * -> *) rep.
(MonadFreshNames m, Buildable rep, LetDec rep ~ Type,
 ExpDec rep ~ (), Op rep ~ SOAC rep) =>
Pat (LetDec rep)
-> (SubExp, [Scan rep], Lambda rep, [VName])
-> m (Stm rep, Stm rep)
scanomapToMapAndScan (Pat [PatElem (LetDec rep)]
pes) (SubExp
w, [Scan rep]
scans, Lambda rep
map_lam, [VName]
arrs) = do
  (map_pat, scan_pat, scan_arrs, map_lam') <-
    [PatElem Type]
-> SubExp
-> Lambda rep
-> [[SubExp]]
-> m (Pat Type, Pat Type, [VName], Lambda rep)
forall (m :: * -> *) rep.
MonadFreshNames m =>
[PatElem Type]
-> SubExp
-> Lambda rep
-> [[SubExp]]
-> m (Pat Type, Pat Type, [VName], Lambda rep)
splitScanOrRedomap [PatElem Type]
[PatElem (LetDec rep)]
pes SubExp
w Lambda rep
map_lam ([[SubExp]] -> m (Pat Type, Pat Type, [VName], Lambda rep))
-> [[SubExp]] -> m (Pat Type, Pat Type, [VName], Lambda rep)
forall a b. (a -> b) -> a -> b
$ (Scan rep -> [SubExp]) -> [Scan rep] -> [[SubExp]]
forall a b. (a -> b) -> [a] -> [b]
map Scan rep -> [SubExp]
forall rep. Scan rep -> [SubExp]
scanNeutral [Scan rep]
scans
  map_stm <- Let map_pat (defAux ()) . Op . Screma w arrs <$> mapSOAC map_lam'
  scan_stm <-
    Let scan_pat (defAux ()) . Op
      <$> (Screma w scan_arrs <$> scanSOAC scans)
  pure (map_stm, scan_stm)

maposcanomapToMapScanAndMap ::
  ( MonadFreshNames m,
    Buildable rep,
    Op rep ~ SOAC rep
  ) =>
  Pat (LetDec rep) ->
  ( SubExp,
    Lambda rep,
    [Scan rep],
    Lambda rep,
    [VName]
  ) ->
  m (Stm rep, Stm rep, Stm rep)
maposcanomapToMapScanAndMap :: forall (m :: * -> *) rep.
(MonadFreshNames m, Buildable rep, Op rep ~ SOAC rep) =>
Pat (LetDec rep)
-> (SubExp, Lambda rep, [Scan rep], Lambda rep, [VName])
-> m (Stm rep, Stm rep, Stm rep)
maposcanomapToMapScanAndMap (Pat [PatElem (LetDec rep)]
pes) (SubExp
w, Lambda rep
post_lam, [Scan rep]
scans, Lambda rep
map_lam, [VName]
arrs) = do
  map_res <- (Type -> m Ident) -> [Type] -> m [Ident]
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 Ident
tempRes ([Type] -> m [Ident]) -> [Type] -> m [Ident]
forall a b. (a -> b) -> a -> b
$ Lambda rep -> [Type]
forall rep t. GLambda rep t -> [t]
lambdaReturnType Lambda rep
map_lam
  map_stm <- mkLet map_res . Op . Screma w arrs <$> mapSOAC map_lam
  scan_res <- mapM tempRes $ take (scanResults scans) $ lambdaReturnType map_lam
  let (scan_arrs, map_arrs) = splitAt (scanResults scans) $ map identName map_res
  scan_stm <- mkLet scan_res . Op . Screma w scan_arrs <$> scanSOAC scans
  let post_arrs = (Ident -> VName) -> [Ident] -> [VName]
forall a b. (a -> b) -> [a] -> [b]
map Ident -> VName
identName [Ident]
scan_res [VName] -> [VName] -> [VName]
forall a. Semigroup a => a -> a -> a
<> [VName]
map_arrs
      res = PatElem (LetDec rep) -> Ident
forall dec. Typed dec => PatElem dec -> Ident
patElemIdent (PatElem (LetDec rep) -> Ident)
-> [PatElem (LetDec rep)] -> [Ident]
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> [PatElem (LetDec rep)]
pes
  post_stm <- mkLet res . Op . Screma w post_arrs <$> mapSOAC post_lam
  pure (map_stm, scan_stm, post_stm)
  where
    tempRes :: Type -> m Ident
tempRes Type
res = Name -> Type -> m Ident
forall (m :: * -> *). MonadFreshNames m => Name -> Type -> m Ident
newIdent Name
"temp_res" (Type -> m Ident) -> Type -> m Ident
forall a b. (a -> b) -> a -> b
$ Type
res Type -> SubExp -> Type
forall d.
TypeBase (ShapeBase d) NoMode -> d -> TypeBase (ShapeBase d) NoMode
`arrayOfRow` SubExp
w

maposcanomapToMaposcanAndMap ::
  ( MonadFreshNames m,
    Buildable rep,
    Op rep ~ SOAC rep
  ) =>
  Pat (LetDec rep) ->
  ( SubExp,
    Lambda rep,
    [Scan rep],
    Lambda rep,
    [VName]
  ) ->
  m (Stm rep, Stm rep)
maposcanomapToMaposcanAndMap :: forall (m :: * -> *) rep.
(MonadFreshNames m, Buildable rep, Op rep ~ SOAC rep) =>
Pat (LetDec rep)
-> (SubExp, Lambda rep, [Scan rep], Lambda rep, [VName])
-> m (Stm rep, Stm rep)
maposcanomapToMaposcanAndMap (Pat [PatElem (LetDec rep)]
pes) (SubExp
w, Lambda rep
post_lam, [Scan rep]
scans, Lambda rep
map_lam, [VName]
arrs) = do
  map_res <- (Type -> m Ident) -> [Type] -> m [Ident]
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 Ident
tempRes ([Type] -> m [Ident]) -> [Type] -> m [Ident]
forall a b. (a -> b) -> a -> b
$ Lambda rep -> [Type]
forall rep t. GLambda rep t -> [t]
lambdaReturnType Lambda rep
map_lam
  map_stm <- mkLet map_res . Op . Screma w arrs <$> mapSOAC map_lam
  id_lam <- mkIdentityLambda $ lambdaReturnType map_lam
  let res = PatElem (LetDec rep) -> Ident
forall dec. Typed dec => PatElem dec -> Ident
patElemIdent (PatElem (LetDec rep) -> Ident)
-> [PatElem (LetDec rep)] -> [Ident]
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> [PatElem (LetDec rep)]
pes
      pre_arrs = (Ident -> VName) -> [Ident] -> [VName]
forall a b. (a -> b) -> [a] -> [b]
map Ident -> VName
identName [Ident]
map_res
  post_stm <- mkLet res . Op . Screma w pre_arrs <$> maposcanomapSOAC id_lam scans post_lam
  pure (map_stm, post_stm)
  where
    tempRes :: Type -> m Ident
tempRes Type
res = Name -> Type -> m Ident
forall (m :: * -> *). MonadFreshNames m => Name -> Type -> m Ident
newIdent Name
"temp_res" (Type -> m Ident) -> Type -> m Ident
forall a b. (a -> b) -> a -> b
$ Type
res Type -> SubExp -> Type
forall d.
TypeBase (ShapeBase d) NoMode -> d -> TypeBase (ShapeBase d) NoMode
`arrayOfRow` SubExp
w

-- | Turn a Screma into a maposcanomap (possibly with mapout parts) and a
-- Redomap.  This is used to handle Scremas that are so complicated
-- that we cannot directly generate efficient parallel code for them.
-- In essense, what happens is the opposite of horisontal fusion.
dissectScrema ::
  ( MonadBuilder m,
    Op (Rep m) ~ SOAC (Rep m),
    Buildable (Rep m)
  ) =>
  Pat (LetDec (Rep m)) ->
  SubExp ->
  ScremaForm (Rep m) ->
  [VName] ->
  m ()
dissectScrema :: forall (m :: * -> *).
(MonadBuilder m, Op (Rep m) ~ SOAC (Rep m), Buildable (Rep m)) =>
Pat (LetDec (Rep m))
-> SubExp -> ScremaForm (Rep m) -> [VName] -> m ()
dissectScrema Pat (LetDec (Rep m))
pat SubExp
w (ScremaForm Lambda (Rep m)
map_lam [Scan (Rep m)]
scans [Reduce (Rep m)]
reds Lambda (Rep m)
post_lam) [VName]
arrs = do
  let num_reds :: Int
num_reds = [Reduce (Rep m)] -> Int
forall rep. [Reduce rep] -> Int
redResults [Reduce (Rep m)]
reds
      num_scans :: Int
num_scans = [Scan (Rep m)] -> Int
forall rep. [Scan rep] -> Int
scanResults [Scan (Rep m)]
scans
      reds_ts :: [Type]
reds_ts = (Reduce (Rep m) -> [Type]) -> [Reduce (Rep m)] -> [Type]
forall (t :: * -> *) a b. Foldable t => (a -> [b]) -> t a -> [b]
concatMap (Lambda (Rep m) -> [Type]
forall rep t. GLambda rep t -> [t]
lambdaReturnType (Lambda (Rep m) -> [Type])
-> (Reduce (Rep m) -> Lambda (Rep m)) -> Reduce (Rep m) -> [Type]
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Reduce (Rep m) -> Lambda (Rep m)
forall rep. Reduce rep -> Lambda rep
redLambda) [Reduce (Rep m)]
reds
      ([VName]
red_res, [VName]
scan_map_res) = Int -> [VName] -> ([VName], [VName])
forall a. Int -> [a] -> ([a], [a])
splitAt Int
num_reds ([VName] -> ([VName], [VName])) -> [VName] -> ([VName], [VName])
forall a b. (a -> b) -> a -> b
$ Pat (LetDec (Rep m)) -> [VName]
forall dec. Pat dec -> [VName]
patNames Pat (LetDec (Rep m))
pat
      ([Param Type]
scan_pars, [Param Type]
map_pars) = Int -> [Param Type] -> ([Param Type], [Param Type])
forall a. Int -> [a] -> ([a], [a])
splitAt Int
num_scans ([Param Type] -> ([Param Type], [Param Type]))
-> [Param Type] -> ([Param Type], [Param Type])
forall a b. (a -> b) -> a -> b
$ Lambda (Rep m) -> [LParam (Rep m)]
forall rep t. GLambda rep t -> [LParam rep]
lambdaParams Lambda (Rep m)
post_lam
      post_res :: [SubExpRes]
post_res = GBody (Rep m) SubExpRes -> [SubExpRes]
forall rep res. GBody rep res -> [res]
bodyResult (GBody (Rep m) SubExpRes -> [SubExpRes])
-> GBody (Rep m) SubExpRes -> [SubExpRes]
forall a b. (a -> b) -> a -> b
$ Lambda (Rep m) -> GBody (Rep m) SubExpRes
forall rep t. GLambda rep t -> Body rep
lambdaBody Lambda (Rep m)
post_lam

  to_red <- Int -> m VName -> m [VName]
forall (m :: * -> *) a. Applicative m => Int -> m a -> m [a]
replicateM Int
num_reds (m VName -> m [VName]) -> m VName -> m [VName]
forall a b. (a -> b) -> a -> b
$ Name -> m VName
forall (m :: * -> *). MonadFreshNames m => Name -> m VName
newVName Name
"to_red"
  red_pars <- mapM (newParam "x") reds_ts
  let red_post_res = Param Type -> VName
forall dec. Param dec -> VName
paramName (Param Type -> VName) -> [Param Type] -> [VName]
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> [Param Type]
red_pars

  let post_lam' =
        Lambda (Rep m)
post_lam
          { lambdaParams = scan_pars <> red_pars <> map_pars,
            lambdaBody =
              (lambdaBody post_lam)
                { bodyResult = varsRes red_post_res <> post_res
                },
            lambdaReturnType = reds_ts <> lambdaReturnType post_lam
          }

  maposcanomap <- maposcanomapSOAC map_lam scans post_lam'
  letBindNames (to_red <> scan_map_res) $ Op (Screma w arrs maposcanomap)

  reduce <- reduceSOAC reds
  letBindNames red_res $ Op $ Screma w to_red reduce

-- | Remove the post lambda from a screma, producing the screma with an identity
-- post-lambda, and a new map screma that is just the post-lambda. You can apply
-- this indefinitely in case the post-lambda is already an identity lambda, so
-- be careful.
extractPostLambda ::
  ( MonadBuilder m,
    Op (Rep m) ~ SOAC (Rep m),
    Buildable (Rep m)
  ) =>
  Pat (LetDec (Rep m)) ->
  SubExp ->
  [VName] ->
  ScremaForm (Rep m) ->
  m ()
extractPostLambda :: forall (m :: * -> *).
(MonadBuilder m, Op (Rep m) ~ SOAC (Rep m), Buildable (Rep m)) =>
Pat (LetDec (Rep m))
-> SubExp -> [VName] -> ScremaForm (Rep m) -> m ()
extractPostLambda Pat (LetDec (Rep m))
pat SubExp
w [VName]
arrs (ScremaForm Lambda (Rep m)
pre_lam [Scan (Rep m)]
scans [Reduce (Rep m)]
reds Lambda (Rep m)
post_lam) = do
  tmp_names <-
    (Param (LParamInfo (Rep m)) -> m VName)
-> [Param (LParamInfo (Rep m))] -> m [VName]
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 (Name -> m VName
forall (m :: * -> *). MonadFreshNames m => Name -> m VName
newVName (Name -> m VName)
-> (Param (LParamInfo (Rep m)) -> Name)
-> Param (LParamInfo (Rep m))
-> m VName
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (Name -> Name -> Name
forall a. Semigroup a => a -> a -> a
<> Name
"_extract") (Name -> Name)
-> (Param (LParamInfo (Rep m)) -> Name)
-> Param (LParamInfo (Rep m))
-> Name
forall b c a. (b -> c) -> (a -> b) -> a -> c
. VName -> Name
baseName (VName -> Name)
-> (Param (LParamInfo (Rep m)) -> VName)
-> Param (LParamInfo (Rep m))
-> Name
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Param (LParamInfo (Rep m)) -> VName
forall dec. Param dec -> VName
paramName) (Lambda (Rep m) -> [Param (LParamInfo (Rep m))]
forall rep t. GLambda rep t -> [LParam rep]
lambdaParams Lambda (Rep m)
post_lam)
  id_lam <- mkIdentityLambda $ map paramType $ lambdaParams post_lam
  letBindNames (map patElemName red_res <> tmp_names) $
    Op (Screma w arrs $ ScremaForm pre_lam scans reds id_lam)
  letBind (Pat nonred_res) . Op . Screma w tmp_names =<< mapSOAC post_lam
  where
    ([PatElem (LetDec (Rep m))]
red_res, [PatElem (LetDec (Rep m))]
nonred_res) = Int
-> [PatElem (LetDec (Rep m))]
-> ([PatElem (LetDec (Rep m))], [PatElem (LetDec (Rep m))])
forall a. Int -> [a] -> ([a], [a])
splitAt ([Reduce (Rep m)] -> Int
forall rep. [Reduce rep] -> Int
redResults [Reduce (Rep m)]
reds) (Pat (LetDec (Rep m)) -> [PatElem (LetDec (Rep m))]
forall dec. Pat dec -> [PatElem dec]
patElems Pat (LetDec (Rep m))
pat)

-- | Turn a stream SOAC into statements that apply the stream lambda
-- to the entire input.
sequentialStreamWholeArray ::
  (MonadBuilder m, Buildable (Rep m)) =>
  Pat (LetDec (Rep m)) ->
  SubExp ->
  [SubExp] ->
  Lambda (Rep m) ->
  [VName] ->
  m ()
sequentialStreamWholeArray :: forall (m :: * -> *).
(MonadBuilder m, Buildable (Rep m)) =>
Pat (LetDec (Rep m))
-> SubExp -> [SubExp] -> Lambda (Rep m) -> [VName] -> m ()
sequentialStreamWholeArray Pat (LetDec (Rep m))
pat SubExp
w [SubExp]
nes Lambda (Rep m)
lam [VName]
arrs = do
  -- We just set the chunksize to w and inline the lambda body.  There
  -- is no difference between parallel and sequential streams here.
  let (Param Type
chunk_size_param, [Param Type]
fold_params, [Param Type]
arr_params) =
        Int -> [Param Type] -> (Param Type, [Param Type], [Param Type])
forall dec.
Int -> [Param dec] -> (Param dec, [Param dec], [Param dec])
partitionChunkedFoldParameters ([SubExp] -> Int
forall a. [a] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length [SubExp]
nes) ([Param Type] -> (Param Type, [Param Type], [Param Type]))
-> [Param Type] -> (Param Type, [Param Type], [Param Type])
forall a b. (a -> b) -> a -> b
$ Lambda (Rep m) -> [LParam (Rep m)]
forall rep t. GLambda rep t -> [LParam rep]
lambdaParams Lambda (Rep m)
lam

  -- The chunk size is the full size of the array.
  [VName] -> Exp (Rep m) -> m ()
forall (m :: * -> *).
MonadBuilder m =>
[VName] -> Exp (Rep m) -> m ()
letBindNames [Param Type -> VName
forall dec. Param dec -> VName
paramName Param Type
chunk_size_param] (Exp (Rep m) -> m ()) -> Exp (Rep m) -> m ()
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
w

  -- The accumulator parameters are initialised to the neutral element.
  [(Param Type, SubExp)] -> ((Param Type, SubExp) -> m ()) -> m ()
forall (t :: * -> *) (m :: * -> *) a b.
(Foldable t, Monad m) =>
t a -> (a -> m b) -> m ()
forM_ ([Param Type] -> [SubExp] -> [(Param Type, SubExp)]
forall a b. [a] -> [b] -> [(a, b)]
zip [Param Type]
fold_params [SubExp]
nes) (((Param Type, SubExp) -> m ()) -> m ())
-> ((Param Type, SubExp) -> m ()) -> m ()
forall a b. (a -> b) -> a -> b
$ \(Param Type
p, SubExp
ne) ->
    [VName] -> Exp (Rep m) -> m ()
forall (m :: * -> *).
MonadBuilder m =>
[VName] -> Exp (Rep m) -> m ()
letBindNames [Param Type -> VName
forall dec. Param dec -> VName
paramName Param Type
p] (Exp (Rep m) -> m ()) -> Exp (Rep m) -> m ()
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
ne

  -- Finally, the array parameters are set to the arrays (but reshaped
  -- to make the types work out; this will be simplified rapidly).
  [(Param Type, VName)] -> ((Param Type, VName) -> m ()) -> m ()
forall (t :: * -> *) (m :: * -> *) a b.
(Foldable t, Monad m) =>
t a -> (a -> m b) -> m ()
forM_ ([Param Type] -> [VName] -> [(Param Type, VName)]
forall a b. [a] -> [b] -> [(a, b)]
zip [Param Type]
arr_params [VName]
arrs) (((Param Type, VName) -> m ()) -> m ())
-> ((Param Type, VName) -> m ()) -> m ()
forall a b. (a -> b) -> a -> b
$ \(Param Type
p, VName
arr) ->
    [VName] -> Exp (Rep m) -> m ()
forall (m :: * -> *).
MonadBuilder m =>
[VName] -> Exp (Rep m) -> m ()
letBindNames [Param Type -> VName
forall dec. Param dec -> VName
paramName Param Type
p] (Exp (Rep m) -> m ()) -> Exp (Rep m) -> m ()
forall a b. (a -> b) -> a -> b
$
      if [SubExp] -> Bool
forall a. [a] -> Bool
forall (t :: * -> *) a. Foldable t => t a -> Bool
null (Type -> [SubExp]
forall d m. TypeBase (ShapeBase d) m -> [d]
arrayDims (Type -> [SubExp]) -> Type -> [SubExp]
forall a b. (a -> b) -> a -> b
$ Param Type -> Type
forall dec. Typed dec => Param dec -> Type
paramType Param Type
p)
        then 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
$ VName -> SubExp
Var VName
arr
        else [SubExp] -> VName -> Exp (Rep m)
forall rep. [SubExp] -> VName -> Exp rep
shapeCoerce (Type -> [SubExp]
forall d m. TypeBase (ShapeBase d) m -> [d]
arrayDims (Type -> [SubExp]) -> Type -> [SubExp]
forall a b. (a -> b) -> a -> b
$ Param Type -> Type
forall dec. Typed dec => Param dec -> Type
paramType Param Type
p) VName
arr

  -- Then we just inline the lambda body.
  (Stm (Rep m) -> m ()) -> Seq (Stm (Rep m)) -> m ()
forall (t :: * -> *) (m :: * -> *) a b.
(Foldable t, Monad m) =>
(a -> m b) -> t a -> m ()
mapM_ Stm (Rep m) -> m ()
forall (m :: * -> *). MonadBuilder m => Stm (Rep m) -> m ()
addStm (Seq (Stm (Rep m)) -> m ()) -> Seq (Stm (Rep m)) -> m ()
forall a b. (a -> b) -> a -> b
$ GBody (Rep m) SubExpRes -> Seq (Stm (Rep m))
forall rep res. GBody rep res -> Stms rep
bodyStms (GBody (Rep m) SubExpRes -> Seq (Stm (Rep m)))
-> GBody (Rep m) SubExpRes -> Seq (Stm (Rep m))
forall a b. (a -> b) -> a -> b
$ Lambda (Rep m) -> GBody (Rep m) SubExpRes
forall rep t. GLambda rep t -> Body rep
lambdaBody Lambda (Rep m)
lam

  -- The number of results in the body matches exactly the size (and
  -- order) of 'pat', so we bind them up here, again with a reshape to
  -- make the types work out.
  [(PatElem (LetDec (Rep m)), SubExpRes)]
-> ((PatElem (LetDec (Rep m)), SubExpRes) -> m ()) -> m ()
forall (t :: * -> *) (m :: * -> *) a b.
(Foldable t, Monad m) =>
t a -> (a -> m b) -> m ()
forM_ ([PatElem (LetDec (Rep m))]
-> [SubExpRes] -> [(PatElem (LetDec (Rep m)), SubExpRes)]
forall a b. [a] -> [b] -> [(a, b)]
zip (Pat (LetDec (Rep m)) -> [PatElem (LetDec (Rep m))]
forall dec. Pat dec -> [PatElem dec]
patElems Pat (LetDec (Rep m))
pat) ([SubExpRes] -> [(PatElem (LetDec (Rep m)), SubExpRes)])
-> [SubExpRes] -> [(PatElem (LetDec (Rep m)), SubExpRes)]
forall a b. (a -> b) -> a -> b
$ GBody (Rep m) SubExpRes -> [SubExpRes]
forall rep res. GBody rep res -> [res]
bodyResult (GBody (Rep m) SubExpRes -> [SubExpRes])
-> GBody (Rep m) SubExpRes -> [SubExpRes]
forall a b. (a -> b) -> a -> b
$ Lambda (Rep m) -> GBody (Rep m) SubExpRes
forall rep t. GLambda rep t -> Body rep
lambdaBody Lambda (Rep m)
lam) (((PatElem (LetDec (Rep m)), SubExpRes) -> m ()) -> m ())
-> ((PatElem (LetDec (Rep m)), SubExpRes) -> m ()) -> m ()
forall a b. (a -> b) -> a -> b
$ \(PatElem (LetDec (Rep m))
pe, SubExpRes Certs
cs SubExp
se) ->
    Certs -> m () -> m ()
forall a. Certs -> m a -> m a
forall (m :: * -> *) a. MonadBuilder m => Certs -> m a -> m a
certifying Certs
cs (m () -> m ()) -> m () -> m ()
forall a b. (a -> b) -> a -> b
$ case (Type -> [SubExp]
forall d m. TypeBase (ShapeBase d) m -> [d]
arrayDims (Type -> [SubExp]) -> Type -> [SubExp]
forall a b. (a -> b) -> a -> b
$ PatElem (LetDec (Rep m)) -> Type
forall dec. Typed dec => PatElem dec -> Type
patElemType PatElem (LetDec (Rep m))
pe, SubExp
se) of
      ([SubExp]
dims, Var VName
v)
        | Bool -> Bool
not (Bool -> Bool) -> Bool -> Bool
forall a b. (a -> b) -> a -> b
$ [SubExp] -> Bool
forall a. [a] -> Bool
forall (t :: * -> *) a. Foldable t => t a -> Bool
null [SubExp]
dims ->
            [VName] -> Exp (Rep m) -> m ()
forall (m :: * -> *).
MonadBuilder m =>
[VName] -> Exp (Rep m) -> m ()
letBindNames [PatElem (LetDec (Rep m)) -> VName
forall dec. PatElem dec -> VName
patElemName PatElem (LetDec (Rep m))
pe] (Exp (Rep m) -> m ()) -> Exp (Rep m) -> m ()
forall a b. (a -> b) -> a -> b
$ [SubExp] -> VName -> Exp (Rep m)
forall rep. [SubExp] -> VName -> Exp rep
shapeCoerce [SubExp]
dims VName
v
      ([SubExp], SubExp)
_ -> [VName] -> Exp (Rep m) -> m ()
forall (m :: * -> *).
MonadBuilder m =>
[VName] -> Exp (Rep m) -> m ()
letBindNames [PatElem (LetDec (Rep m)) -> VName
forall dec. PatElem dec -> VName
patElemName PatElem (LetDec (Rep m))
pe] (Exp (Rep m) -> m ()) -> Exp (Rep m) -> m ()
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
se

-- | Split the parameters of a stream reduction lambda into the chunk
-- size parameter, the accumulator parameters, and the input chunk
-- parameters.  The integer argument is how many accumulators are
-- used.
partitionChunkedFoldParameters ::
  Int ->
  [Param dec] ->
  (Param dec, [Param dec], [Param dec])
partitionChunkedFoldParameters :: forall dec.
Int -> [Param dec] -> (Param dec, [Param dec], [Param dec])
partitionChunkedFoldParameters Int
_ [] =
  String -> (Param dec, [Param dec], [Param dec])
forall a. HasCallStack => String -> a
error String
"partitionChunkedFoldParameters: lambda takes no parameters"
partitionChunkedFoldParameters Int
num_accs (Param dec
chunk_param : [Param dec]
params) =
  let ([Param dec]
acc_params, [Param dec]
arr_params) = Int -> [Param dec] -> ([Param dec], [Param dec])
forall a. Int -> [a] -> ([a], [a])
splitAt Int
num_accs [Param dec]
params
   in (Param dec
chunk_param, [Param dec]
acc_params, [Param dec]
arr_params)

-- | Construct a one-dimensional scatter-like 'WithAcc'. The closure is invoked
-- with the accumulators.
withAcc ::
  (MonadBuilder m, LParam (Rep m) ~ Param Type) =>
  [VName] ->
  Int ->
  ([VName] -> m [SubExp]) ->
  m (Exp (Rep m))
withAcc :: forall (m :: * -> *).
(MonadBuilder m, LParam (Rep m) ~ Param Type) =>
[VName] -> Int -> ([VName] -> m [SubExp]) -> m (Exp (Rep m))
withAcc [VName]
dest Int
rank [VName] -> m [SubExp]
mk = do
  cert_ps <- Int -> m (Param Type) -> m [Param Type]
forall (m :: * -> *) a. Applicative m => Int -> m a -> m [a]
replicateM ([VName] -> Int
forall a. [a] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length [VName]
dest) (m (Param Type) -> m [Param Type])
-> m (Param Type) -> m [Param Type]
forall a b. (a -> b) -> a -> b
$ Name -> Type -> m (Param Type)
forall (m :: * -> *) dec.
MonadFreshNames m =>
Name -> dec -> m (Param dec)
newParam Name
"acc_cert" (Type -> m (Param Type)) -> Type -> m (Param Type)
forall a b. (a -> b) -> a -> b
$ PrimType -> Type
forall shape o. PrimType -> TypeBase shape o
Prim PrimType
Unit
  dest_ts <- mapM lookupType dest
  let acc_shape = [SubExp] -> ShapeBase SubExp
forall d. [d] -> ShapeBase d
Shape ([SubExp] -> ShapeBase SubExp) -> [SubExp] -> ShapeBase SubExp
forall a b. (a -> b) -> a -> b
$ Int -> [SubExp] -> [SubExp]
forall a. Int -> [a] -> [a]
take Int
rank ([SubExp] -> [SubExp]) -> [SubExp] -> [SubExp]
forall a b. (a -> b) -> a -> b
$ Type -> [SubExp]
forall d m. TypeBase (ShapeBase d) m -> [d]
arrayDims (Type -> [SubExp]) -> Type -> [SubExp]
forall a b. (a -> b) -> a -> b
$ [Type] -> Type
forall a. HasCallStack => [a] -> a
head [Type]
dest_ts
      mkT VName
cert Type
elem_t = VName -> ShapeBase SubExp -> [Type] -> Type
forall shape o.
VName -> ShapeBase SubExp -> [Type] -> TypeBase shape o
Acc VName
cert ShapeBase SubExp
acc_shape [Type
elem_t]
      acc_ts =
        (VName -> Type -> Type) -> [VName] -> [Type] -> [Type]
forall a b c. (a -> b -> c) -> [a] -> [b] -> [c]
zipWith VName -> Type -> Type
mkT ((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]
cert_ps) ([Type] -> [Type]) -> [Type] -> [Type]
forall a b. (a -> b) -> a -> b
$
          (Type -> Type) -> [Type] -> [Type]
forall a b. (a -> b) -> [a] -> [b]
map (Int -> Type -> Type
forall d m.
Int -> TypeBase (ShapeBase d) m -> TypeBase (ShapeBase d) m
stripArray Int
rank) [Type]
dest_ts
  acc_ps <- mapM (newParam "acc_p") acc_ts

  withacc_lam <- mkLambda (cert_ps <> acc_ps) $ subExpsRes <$> mk (map paramName acc_ps)

  pure $ WithAcc [(acc_shape, [v], Nothing) | v <- dest] withacc_lam

-- | Perform a scatter-like operation using accumulators and map.
doScatter ::
  (MonadBuilder m, Buildable (Rep m), Op (Rep m) ~ SOAC (Rep m)) =>
  Name ->
  Int ->
  [VName] ->
  [VName] ->
  ([LParam (Rep m)] -> m [SubExp]) ->
  m [VName]
doScatter :: forall (m :: * -> *).
(MonadBuilder m, Buildable (Rep m), Op (Rep m) ~ SOAC (Rep m)) =>
Name
-> Int
-> [VName]
-> [VName]
-> ([LParam (Rep m)] -> m [SubExp])
-> m [VName]
doScatter Name
desc Int
rank [VName]
dest [VName]
arrs [LParam (Rep m)] -> m [SubExp]
mk = do
  cert_ps <- Int -> m (Param Type) -> m [Param Type]
forall (m :: * -> *) a. Applicative m => Int -> m a -> m [a]
replicateM ([VName] -> Int
forall a. [a] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length [VName]
dest) (m (Param Type) -> m [Param Type])
-> m (Param Type) -> m [Param Type]
forall a b. (a -> b) -> a -> b
$ Name -> Type -> m (Param Type)
forall (m :: * -> *) dec.
MonadFreshNames m =>
Name -> dec -> m (Param dec)
newParam Name
"acc_cert" (Type -> m (Param Type)) -> Type -> m (Param Type)
forall a b. (a -> b) -> a -> b
$ PrimType -> Type
forall shape o. PrimType -> TypeBase shape o
Prim PrimType
Unit
  dest_ts <- mapM lookupType dest
  let acc_shape = [SubExp] -> ShapeBase SubExp
forall d. [d] -> ShapeBase d
Shape ([SubExp] -> ShapeBase SubExp) -> [SubExp] -> ShapeBase SubExp
forall a b. (a -> b) -> a -> b
$ Int -> [SubExp] -> [SubExp]
forall a. Int -> [a] -> [a]
take Int
rank ([SubExp] -> [SubExp]) -> [SubExp] -> [SubExp]
forall a b. (a -> b) -> a -> b
$ Type -> [SubExp]
forall d m. TypeBase (ShapeBase d) m -> [d]
arrayDims (Type -> [SubExp]) -> Type -> [SubExp]
forall a b. (a -> b) -> a -> b
$ [Type] -> Type
forall a. HasCallStack => [a] -> a
head [Type]
dest_ts
      mkT VName
cert Type
elem_t = VName -> ShapeBase SubExp -> [Type] -> Type
forall shape o.
VName -> ShapeBase SubExp -> [Type] -> TypeBase shape o
Acc VName
cert ShapeBase SubExp
acc_shape [Type
elem_t]
      acc_ts =
        (VName -> Type -> Type) -> [VName] -> [Type] -> [Type]
forall a b c. (a -> b -> c) -> [a] -> [b] -> [c]
zipWith VName -> Type -> Type
mkT ((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]
cert_ps) ([Type] -> [Type]) -> [Type] -> [Type]
forall a b. (a -> b) -> a -> b
$
          (Type -> Type) -> [Type] -> [Type]
forall a b. (a -> b) -> [a] -> [b]
map (Int -> Type -> Type
forall d m.
Int -> TypeBase (ShapeBase d) m -> TypeBase (ShapeBase d) m
stripArray Int
rank) [Type]
dest_ts
  acc_ps <- mapM (newParam "acc_p") acc_ts
  arrs_ts <- mapM lookupType arrs

  withacc_lam <- mkLambda (cert_ps <> acc_ps) $ do
    acc_ps_inner <- mapM (newParam "acc_p") acc_ts
    params <- mapM (newParam "v" . stripArray 1) arrs_ts
    map_lam <-
      mkLambda (acc_ps_inner <> params) $ do
        (is, vs) <- splitAt rank <$> mk params
        fmap subExpsRes $ forM (zip acc_ps_inner vs) $ \(Param Type
acc_p_inner, SubExp
v) ->
          Name -> Exp (Rep m) -> m SubExp
forall (m :: * -> *).
MonadBuilder m =>
Name -> Exp (Rep m) -> m SubExp
letSubExp Name
"scatter_acc" (Exp (Rep m) -> m SubExp)
-> (BasicOp -> Exp (Rep m)) -> BasicOp -> m SubExp
forall b c a. (b -> c) -> (a -> b) -> a -> c
. BasicOp -> Exp (Rep m)
forall rep. BasicOp -> Exp rep
BasicOp (BasicOp -> m SubExp) -> BasicOp -> m SubExp
forall a b. (a -> b) -> a -> b
$
            Safety -> VName -> [SubExp] -> [SubExp] -> BasicOp
UpdateAcc Safety
Safe (Param Type -> VName
forall dec. Param dec -> VName
paramName Param Type
acc_p_inner) [SubExp]
is [SubExp
v]

    let w = Int -> [Type] -> SubExp
forall o. Int -> [TypeBase (ShapeBase SubExp) o] -> SubExp
arraysSize Int
0 [Type]
arrs_ts
    (fmap varsRes . letTupExp "acc_res")
      . Op
      . Screma w (map paramName acc_ps <> arrs)
      =<< mapSOAC map_lam

  letTupExp desc $ WithAcc [(acc_shape, [v], Nothing) | v <- dest] withacc_lam

-- | Perform a histogram-like operation using accumulators and map.
doHist ::
  (MonadBuilder m, Buildable (Rep m), Op (Rep m) ~ SOAC (Rep m)) =>
  Name ->
  [HistOp (Rep m)] ->
  [VName] ->
  ([LParam (Rep m)] -> m [SubExp]) ->
  m [[VName]]
doHist :: forall (m :: * -> *).
(MonadBuilder m, Buildable (Rep m), Op (Rep m) ~ SOAC (Rep m)) =>
Name
-> [HistOp (Rep m)]
-> [VName]
-> ([LParam (Rep m)] -> m [SubExp])
-> m [[VName]]
doHist Name
desc [HistOp (Rep m)]
ops [VName]
arrs [LParam (Rep m)] -> m [SubExp]
mk = do
  (inputs, cert_ps, acc_ts) <- [((ShapeBase SubExp, [VName],
   Maybe (GLambda (Rep m) Type, [SubExp])),
  Param Type, Type)]
-> ([(ShapeBase SubExp, [VName],
      Maybe (GLambda (Rep m) Type, [SubExp]))],
    [Param Type], [Type])
forall a b c. [(a, b, c)] -> ([a], [b], [c])
unzip3 ([((ShapeBase SubExp, [VName],
    Maybe (GLambda (Rep m) Type, [SubExp])),
   Param Type, Type)]
 -> ([(ShapeBase SubExp, [VName],
       Maybe (GLambda (Rep m) Type, [SubExp]))],
     [Param Type], [Type]))
-> m [((ShapeBase SubExp, [VName],
        Maybe (GLambda (Rep m) Type, [SubExp])),
       Param Type, Type)]
-> m ([(ShapeBase SubExp, [VName],
        Maybe (GLambda (Rep m) Type, [SubExp]))],
      [Param Type], [Type])
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> (HistOp (Rep m)
 -> m ((ShapeBase SubExp, [VName],
        Maybe (GLambda (Rep m) Type, [SubExp])),
       Param Type, Type))
-> [HistOp (Rep m)]
-> m [((ShapeBase SubExp, [VName],
        Maybe (GLambda (Rep m) Type, [SubExp])),
       Param Type, Type)]
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 HistOp (Rep m)
-> m ((ShapeBase SubExp, [VName],
       Maybe (GLambda (Rep m) Type, [SubExp])),
      Param Type, Type)
forall {rep} {shape} {o} {m :: * -> *} {rep} {shape} {o} {shape}
       {o}.
(LParamInfo rep ~ TypeBase shape o, MonadFreshNames m,
 HasScope rep m) =>
HistOp rep
-> m ((ShapeBase SubExp, [VName],
       Maybe (GLambda rep Type, [SubExp])),
      Param (TypeBase shape o), TypeBase shape o)
onOp [HistOp (Rep m)]
ops
  acc_ps <- mapM (newParam "acc_p") acc_ts
  arrs_ts <- mapM lookupType arrs

  withacc_lam <- mkLambda (cert_ps <> acc_ps) $ do
    acc_ps_inner <- mapM (newParam "acc_p") acc_ts
    params <- mapM (newParam "v" . stripArray 1) arrs_ts
    map_lam <-
      mkLambda (acc_ps_inner <> params) $ do
        let num_is_per_op = (HistOp (Rep m) -> Int) -> [HistOp (Rep m)] -> [Int]
forall a b. (a -> b) -> [a] -> [b]
map (ShapeBase SubExp -> Int
forall a. ArrayShape a => a -> Int
shapeRank (ShapeBase SubExp -> Int)
-> (HistOp (Rep m) -> ShapeBase SubExp) -> HistOp (Rep m) -> Int
forall b c a. (b -> c) -> (a -> b) -> a -> c
. HistOp (Rep m) -> ShapeBase SubExp
forall rep. HistOp rep -> ShapeBase SubExp
histShape) [HistOp (Rep m)]
ops
            num_vs_per_op = (HistOp (Rep m) -> Int) -> [HistOp (Rep m)] -> [Int]
forall a b. (a -> b) -> [a] -> [b]
map ([VName] -> Int
forall a. [a] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length ([VName] -> Int)
-> (HistOp (Rep m) -> [VName]) -> HistOp (Rep m) -> Int
forall b c a. (b -> c) -> (a -> b) -> a -> c
. HistOp (Rep m) -> [VName]
forall rep. HistOp rep -> [VName]
histDest) [HistOp (Rep m)]
ops
        (is, vs) <- splitAt (sum num_is_per_op) <$> mk params
        let is_per_op = [Int] -> [SubExp] -> [[SubExp]]
forall a. [Int] -> [a] -> [[a]]
chunks [Int]
num_is_per_op [SubExp]
is
            vs_per_op = [Int] -> [SubExp] -> [[SubExp]]
forall a. [Int] -> [a] -> [[a]]
chunks [Int]
num_vs_per_op [SubExp]
vs
        fmap subExpsRes $ forM (zip3 acc_ps_inner is_per_op vs_per_op) $ \(Param Type
acc_p_inner, [SubExp]
op_is, [SubExp]
op_vs) ->
          Name -> Exp (Rep m) -> m SubExp
forall (m :: * -> *).
MonadBuilder m =>
Name -> Exp (Rep m) -> m SubExp
letSubExp Name
"scatter_acc" (Exp (Rep m) -> m SubExp)
-> (BasicOp -> Exp (Rep m)) -> BasicOp -> m SubExp
forall b c a. (b -> c) -> (a -> b) -> a -> c
. BasicOp -> Exp (Rep m)
forall rep. BasicOp -> Exp rep
BasicOp (BasicOp -> m SubExp) -> BasicOp -> m SubExp
forall a b. (a -> b) -> a -> b
$
            Safety -> VName -> [SubExp] -> [SubExp] -> BasicOp
UpdateAcc Safety
Safe (Param Type -> VName
forall dec. Param dec -> VName
paramName Param Type
acc_p_inner) [SubExp]
op_is [SubExp]
op_vs

    let w = Int -> [Type] -> SubExp
forall o. Int -> [TypeBase (ShapeBase SubExp) o] -> SubExp
arraysSize Int
0 [Type]
arrs_ts
    (fmap varsRes . letTupExp "acc_res")
      . Op
      . Screma w (map paramName acc_ps <> arrs)
      =<< mapSOAC map_lam
  fmap (chunks (map (length . histDest) ops)) $
    letTupExp desc $
      WithAcc inputs withacc_lam
  where
    onOp :: HistOp rep
-> m ((ShapeBase SubExp, [VName],
       Maybe (GLambda rep Type, [SubExp])),
      Param (TypeBase shape o), TypeBase shape o)
onOp HistOp rep
op = do
      idx_params <- Int -> m (Param (TypeBase shape o)) -> m [Param (TypeBase shape o)]
forall (m :: * -> *) a. Applicative m => Int -> m a -> m [a]
replicateM (ShapeBase SubExp -> Int
forall a. ArrayShape a => a -> Int
shapeRank (HistOp rep -> ShapeBase SubExp
forall rep. HistOp rep -> ShapeBase SubExp
histShape HistOp rep
op)) (m (Param (TypeBase shape o)) -> m [Param (TypeBase shape o)])
-> m (Param (TypeBase shape o)) -> m [Param (TypeBase shape o)]
forall a b. (a -> b) -> a -> b
$ Name -> TypeBase shape o -> m (Param (TypeBase shape o))
forall (m :: * -> *) dec.
MonadFreshNames m =>
Name -> dec -> m (Param dec)
newParam Name
"idx" (TypeBase shape o -> m (Param (TypeBase shape o)))
-> TypeBase shape o -> m (Param (TypeBase shape o))
forall a b. (a -> b) -> a -> b
$ PrimType -> TypeBase shape o
forall shape o. PrimType -> TypeBase shape o
Prim PrimType
int64
      let addIdxParams GLambda rep Type
lam = GLambda rep Type
lam {lambdaParams = idx_params <> lambdaParams lam}
          input = (HistOp rep -> ShapeBase SubExp
forall rep. HistOp rep -> ShapeBase SubExp
histShape HistOp rep
op, HistOp rep -> [VName]
forall rep. HistOp rep -> [VName]
histDest HistOp rep
op, (GLambda rep Type, [SubExp]) -> Maybe (GLambda rep Type, [SubExp])
forall a. a -> Maybe a
Just (GLambda rep Type -> GLambda rep Type
addIdxParams (HistOp rep -> GLambda rep Type
forall rep. HistOp rep -> Lambda rep
histOp HistOp rep
op), HistOp rep -> [SubExp]
forall rep. HistOp rep -> [SubExp]
histNeutral HistOp rep
op))
          shape = HistOp rep -> ShapeBase SubExp
forall rep. HistOp rep -> ShapeBase SubExp
histShape HistOp rep
op
      elem_ts <- fmap (map (stripArray (shapeRank shape))) $ mapM lookupType $ histDest op
      cert_p <- newParam "acc_cert" $ Prim Unit
      let cert = Param (TypeBase shape o) -> VName
forall dec. Param dec -> VName
paramName Param (TypeBase shape o)
cert_p
      pure (input, cert_p, Acc cert shape elem_ts)

-- | The most addition-like binary operator for some primitive type.
addBinOp :: PrimType -> BinOp
addBinOp :: PrimType -> BinOp
addBinOp (IntType IntType
it) = IntType -> Overflow -> BinOp
Add IntType
it Overflow
OverflowWrap
addBinOp (FloatType FloatType
ft) = FloatType -> BinOp
FAdd FloatType
ft
addBinOp PrimType
Bool = BinOp
LogAnd
addBinOp PrimType
Unit = BinOp
LogAnd

-- | Construct a lambda for adding two values of the given type, Using SOACs to handle arrays.
addLambda ::
  ( OpC (Rep m) ~ SOAC,
    MonadBuilder m,
    Buildable (Rep m)
  ) =>
  TypeBase Shape NoMode ->
  m (Lambda (Rep m))
addLambda :: forall (m :: * -> *).
(OpC (Rep m) ~ SOAC, MonadBuilder m, Buildable (Rep m)) =>
Type -> m (Lambda (Rep m))
addLambda (Prim PrimType
pt) = BinOp -> PrimType -> m (GLambda (Rep m) Type)
forall (m :: * -> *).
(MonadBuilder m, Buildable (Rep m)) =>
BinOp -> PrimType -> m (Lambda (Rep m))
binOpLambda (PrimType -> BinOp
addBinOp PrimType
pt) PrimType
pt
addLambda t :: Type
t@Array {} = do
  xs_p <- Name -> Type -> m (Param Type)
forall (m :: * -> *) dec.
MonadFreshNames m =>
Name -> dec -> m (Param dec)
newParam Name
"xs" Type
t
  ys_p <- newParam "ys" t
  lam <- addLambda $ rowType t
  body <- insertStmsM $ do
    res <-
      letSubExp "lam_map"
        . Op
        . Screma (arraySize 0 t) [paramName xs_p, paramName ys_p]
        =<< mapSOAC lam
    pure $ resultBody [res]
  pure
    Lambda
      { lambdaParams = [xs_p, ys_p],
        lambdaReturnType = [t],
        lambdaBody = body
      }
addLambda Type
t =
  String -> m (GLambda (Rep m) Type)
forall a. HasCallStack => String -> a
error (String -> m (GLambda (Rep m) Type))
-> String -> m (GLambda (Rep m) Type)
forall a b. (a -> b) -> a -> b
$ String
"addLambda: " String -> String -> String
forall a. [a] -> [a] -> [a]
++ Type -> String
forall a. Show a => a -> String
show Type
t