{-# 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,
    dissectScrema,
    sequentialStreamWholeArray,
    partitionChunkedFoldParameters,

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

import Control.Monad
import Futhark.Analysis.PrimExp.Convert
import Futhark.Construct
import Futhark.IR
import Futhark.IR.SOACS.SOAC
import Futhark.Util

-- | 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,
    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, 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
  ([Ident]
map_pat, Pat (LetDec rep)
red_pat, [VName]
red_arrs) <-
    [PatElem (LetDec rep)]
-> SubExp
-> Lambda rep
-> [[SubExp]]
-> m ([Ident], Pat (LetDec rep), [VName])
forall dec (m :: * -> *) rep.
(Typed dec, MonadFreshNames m) =>
[PatElem dec]
-> SubExp
-> Lambda rep
-> [[SubExp]]
-> m ([Ident], Pat dec, [VName])
splitScanOrRedomap [PatElem (LetDec rep)]
pes SubExp
w Lambda rep
map_lam ([[SubExp]] -> m ([Ident], Pat (LetDec rep), [VName]))
-> [[SubExp]] -> m ([Ident], Pat (LetDec rep), [VName])
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
  let map_stm :: Stm rep
map_stm = [Ident] -> Exp rep -> Stm rep
forall rep. Buildable rep => [Ident] -> Exp rep -> Stm rep
mkLet [Ident]
map_pat (Exp rep -> Stm rep) -> Exp rep -> Stm rep
forall a b. (a -> b) -> a -> b
$ Op rep -> Exp rep
forall rep. Op rep -> Exp rep
Op (Op rep -> Exp rep) -> Op rep -> Exp rep
forall a b. (a -> b) -> a -> b
$ SubExp -> [VName] -> ScremaForm rep -> SOAC rep
forall rep. SubExp -> [VName] -> ScremaForm rep -> SOAC rep
Screma SubExp
w [VName]
arrs (Lambda rep -> ScremaForm rep
forall rep. Lambda rep -> ScremaForm rep
mapSOAC Lambda rep
map_lam)
  Stm rep
red_stm <-
    Pat (LetDec rep) -> StmAux (ExpDec rep) -> Exp rep -> Stm rep
forall rep.
Pat (LetDec rep) -> StmAux (ExpDec rep) -> Exp rep -> Stm rep
Let Pat (LetDec rep)
red_pat (() -> StmAux ()
forall dec. dec -> StmAux dec
defAux ()) (Exp rep -> Stm rep)
-> (SOAC rep -> Exp rep) -> SOAC rep -> Stm rep
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Op rep -> Exp rep
SOAC rep -> Exp rep
forall rep. Op rep -> Exp rep
Op
      (SOAC rep -> Stm rep) -> m (SOAC rep) -> m (Stm rep)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> (SubExp -> [VName] -> ScremaForm rep -> SOAC rep
forall rep. SubExp -> [VName] -> ScremaForm rep -> SOAC rep
Screma SubExp
w [VName]
red_arrs (ScremaForm rep -> SOAC rep) -> m (ScremaForm rep) -> m (SOAC rep)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> [Reduce rep] -> m (ScremaForm rep)
forall rep (m :: * -> *).
(Buildable rep, MonadFreshNames m) =>
[Reduce rep] -> m (ScremaForm rep)
reduceSOAC [Reduce rep]
reds)
  (Stm rep, Stm rep) -> m (Stm rep, Stm rep)
forall a. a -> m a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (Stm rep
map_stm, Stm rep
red_stm)

scanomapToMapAndScan ::
  ( MonadFreshNames m,
    Buildable rep,
    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, 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
  ([Ident]
map_pat, Pat (LetDec rep)
scan_pat, [VName]
scan_arrs) <-
    [PatElem (LetDec rep)]
-> SubExp
-> Lambda rep
-> [[SubExp]]
-> m ([Ident], Pat (LetDec rep), [VName])
forall dec (m :: * -> *) rep.
(Typed dec, MonadFreshNames m) =>
[PatElem dec]
-> SubExp
-> Lambda rep
-> [[SubExp]]
-> m ([Ident], Pat dec, [VName])
splitScanOrRedomap [PatElem (LetDec rep)]
pes SubExp
w Lambda rep
map_lam ([[SubExp]] -> m ([Ident], Pat (LetDec rep), [VName]))
-> [[SubExp]] -> m ([Ident], Pat (LetDec rep), [VName])
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
  let map_stm :: Stm rep
map_stm = [Ident] -> Exp rep -> Stm rep
forall rep. Buildable rep => [Ident] -> Exp rep -> Stm rep
mkLet [Ident]
map_pat (Exp rep -> Stm rep) -> Exp rep -> Stm rep
forall a b. (a -> b) -> a -> b
$ Op rep -> Exp rep
forall rep. Op rep -> Exp rep
Op (Op rep -> Exp rep) -> Op rep -> Exp rep
forall a b. (a -> b) -> a -> b
$ SubExp -> [VName] -> ScremaForm rep -> SOAC rep
forall rep. SubExp -> [VName] -> ScremaForm rep -> SOAC rep
Screma SubExp
w [VName]
arrs (Lambda rep -> ScremaForm rep
forall rep. Lambda rep -> ScremaForm rep
mapSOAC Lambda rep
map_lam)
  Stm rep
scan_stm <-
    Pat (LetDec rep) -> StmAux (ExpDec rep) -> Exp rep -> Stm rep
forall rep.
Pat (LetDec rep) -> StmAux (ExpDec rep) -> Exp rep -> Stm rep
Let Pat (LetDec rep)
scan_pat (() -> StmAux ()
forall dec. dec -> StmAux dec
defAux ()) (Exp rep -> Stm rep)
-> (SOAC rep -> Exp rep) -> SOAC rep -> Stm rep
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Op rep -> Exp rep
SOAC rep -> Exp rep
forall rep. Op rep -> Exp rep
Op
      (SOAC rep -> Stm rep) -> m (SOAC rep) -> m (Stm rep)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> (SubExp -> [VName] -> ScremaForm rep -> SOAC rep
forall rep. SubExp -> [VName] -> ScremaForm rep -> SOAC rep
Screma SubExp
w [VName]
scan_arrs (ScremaForm rep -> SOAC rep) -> m (ScremaForm rep) -> m (SOAC rep)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> [Scan rep] -> m (ScremaForm rep)
forall rep (m :: * -> *).
(Buildable rep, MonadFreshNames m) =>
[Scan rep] -> m (ScremaForm rep)
scanSOAC [Scan rep]
scans)
  (Stm rep, Stm rep) -> m (Stm rep, Stm rep)
forall a. a -> m a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (Stm rep
map_stm, Stm rep
scan_stm)

splitScanOrRedomap ::
  (Typed dec, MonadFreshNames m) =>
  [PatElem dec] ->
  SubExp ->
  Lambda rep ->
  [[SubExp]] ->
  m ([Ident], Pat dec, [VName])
splitScanOrRedomap :: forall dec (m :: * -> *) rep.
(Typed dec, MonadFreshNames m) =>
[PatElem dec]
-> SubExp
-> Lambda rep
-> [[SubExp]]
-> m ([Ident], Pat dec, [VName])
splitScanOrRedomap [PatElem dec]
pes SubExp
w Lambda rep
map_lam [[SubExp]]
nes = do
  let ([PatElem dec]
acc_pes, [PatElem dec]
arr_pes) =
        Int -> [PatElem dec] -> ([PatElem dec], [PatElem dec])
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 dec]
pes
      ([Type]
acc_ts, [Type]
_arr_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. Lambda rep -> [Type]
lambdaReturnType Lambda rep
map_lam
  [Ident]
map_accpat <- (PatElem dec -> Type -> m Ident)
-> [PatElem dec] -> [Type] -> m [Ident]
forall (m :: * -> *) a b c.
Applicative m =>
(a -> b -> m c) -> [a] -> [b] -> m [c]
zipWithM PatElem dec -> Type -> m Ident
accMapPatElem [PatElem dec]
acc_pes [Type]
acc_ts
  [Ident]
map_arrpat <- (PatElem dec -> m Ident) -> [PatElem dec] -> 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 PatElem dec -> m Ident
arrMapPatElem [PatElem dec]
arr_pes
  let map_pat :: [Ident]
map_pat = [Ident]
map_accpat [Ident] -> [Ident] -> [Ident]
forall a. [a] -> [a] -> [a]
++ [Ident]
map_arrpat
  ([Ident], Pat dec, [VName]) -> m ([Ident], Pat dec, [VName])
forall a. a -> m a
forall (f :: * -> *) a. Applicative f => a -> f a
pure ([Ident]
map_pat, [PatElem dec] -> Pat dec
forall dec. [PatElem dec] -> Pat dec
Pat [PatElem dec]
acc_pes, (Ident -> VName) -> [Ident] -> [VName]
forall a b. (a -> b) -> [a] -> [b]
map Ident -> VName
identName [Ident]
map_accpat)
  where
    accMapPatElem :: PatElem dec -> Type -> m Ident
accMapPatElem PatElem dec
pe Type
acc_t =
      [Char] -> Type -> m Ident
forall (m :: * -> *).
MonadFreshNames m =>
[Char] -> Type -> m Ident
newIdent (VName -> [Char]
baseString (PatElem dec -> VName
forall dec. PatElem dec -> VName
patElemName PatElem dec
pe) [Char] -> [Char] -> [Char]
forall a. [a] -> [a] -> [a]
++ [Char]
"_map_acc") (Type -> m Ident) -> Type -> m Ident
forall a b. (a -> b) -> a -> b
$ Type
acc_t Type -> SubExp -> Type
forall d.
ArrayShape (ShapeBase d) =>
TypeBase (ShapeBase d) NoUniqueness
-> d -> TypeBase (ShapeBase d) NoUniqueness
`arrayOfRow` SubExp
w
    arrMapPatElem :: PatElem dec -> m Ident
arrMapPatElem = Ident -> m Ident
forall a. a -> m a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (Ident -> m Ident)
-> (PatElem dec -> Ident) -> PatElem dec -> m Ident
forall b c a. (b -> c) -> (a -> b) -> a -> c
. PatElem dec -> Ident
forall dec. Typed dec => PatElem dec -> Ident
patElemIdent

-- | Turn a Screma into a Scanomap (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) [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
      ([VName]
scan_res, [VName]
red_res, [VName]
map_res) = Int -> Int -> [VName] -> ([VName], [VName], [VName])
forall a. Int -> Int -> [a] -> ([a], [a], [a])
splitAt3 Int
num_scans Int
num_reds ([VName] -> ([VName], [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

  [VName]
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
$ [Char] -> m VName
forall (m :: * -> *). MonadFreshNames m => [Char] -> m VName
newVName [Char]
"to_red"

  let scanomap :: ScremaForm (Rep m)
scanomap = [Scan (Rep m)] -> Lambda (Rep m) -> ScremaForm (Rep m)
forall rep. [Scan rep] -> Lambda rep -> ScremaForm rep
scanomapSOAC [Scan (Rep m)]
scans Lambda (Rep m)
map_lam
  [VName] -> Exp (Rep m) -> m ()
forall (m :: * -> *).
MonadBuilder m =>
[VName] -> Exp (Rep m) -> m ()
letBindNames ([VName]
scan_res [VName] -> [VName] -> [VName]
forall a. Semigroup a => a -> a -> a
<> [VName]
to_red [VName] -> [VName] -> [VName]
forall a. Semigroup a => a -> a -> a
<> [VName]
map_res) (Exp (Rep m) -> m ()) -> Exp (Rep m) -> m ()
forall a b. (a -> b) -> a -> b
$
    Op (Rep m) -> Exp (Rep m)
forall rep. Op rep -> Exp rep
Op (SubExp -> [VName] -> ScremaForm (Rep m) -> SOAC (Rep m)
forall rep. SubExp -> [VName] -> ScremaForm rep -> SOAC rep
Screma SubExp
w [VName]
arrs ScremaForm (Rep m)
scanomap)

  ScremaForm (Rep m)
reduce <- [Reduce (Rep m)] -> m (ScremaForm (Rep m))
forall rep (m :: * -> *).
(Buildable rep, MonadFreshNames m) =>
[Reduce rep] -> m (ScremaForm rep)
reduceSOAC [Reduce (Rep m)]
reds
  [VName] -> Exp (Rep m) -> m ()
forall (m :: * -> *).
MonadBuilder m =>
[VName] -> Exp (Rep m) -> m ()
letBindNames [VName]
red_res (Exp (Rep m) -> m ()) -> Exp (Rep m) -> m ()
forall a b. (a -> b) -> a -> b
$ Op (Rep m) -> Exp (Rep m)
forall rep. Op rep -> Exp rep
Op (Op (Rep m) -> Exp (Rep m)) -> Op (Rep m) -> Exp (Rep m)
forall a b. (a -> b) -> a -> b
$ SubExp -> [VName] -> ScremaForm (Rep m) -> SOAC (Rep m)
forall rep. SubExp -> [VName] -> ScremaForm rep -> SOAC rep
Screma SubExp
w [VName]
to_red ScremaForm (Rep m)
reduce

-- | 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 (LParamInfo (Rep m))
chunk_size_param, [Param (LParamInfo (Rep m))]
fold_params, [Param (LParamInfo (Rep m))]
arr_params) =
        Int
-> [Param (LParamInfo (Rep m))]
-> (Param (LParamInfo (Rep m)), [Param (LParamInfo (Rep m))],
    [Param (LParamInfo (Rep m))])
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 (LParamInfo (Rep m))]
 -> (Param (LParamInfo (Rep m)), [Param (LParamInfo (Rep m))],
     [Param (LParamInfo (Rep m))]))
-> [Param (LParamInfo (Rep m))]
-> (Param (LParamInfo (Rep m)), [Param (LParamInfo (Rep m))],
    [Param (LParamInfo (Rep m))])
forall a b. (a -> b) -> a -> b
$ Lambda (Rep m) -> [Param (LParamInfo (Rep m))]
forall rep. Lambda rep -> [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 (LParamInfo (Rep m)) -> VName
forall dec. Param dec -> VName
paramName Param (LParamInfo (Rep m))
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 (LParamInfo (Rep m)), SubExp)]
-> ((Param (LParamInfo (Rep m)), SubExp) -> m ()) -> m ()
forall (t :: * -> *) (m :: * -> *) a b.
(Foldable t, Monad m) =>
t a -> (a -> m b) -> m ()
forM_ ([Param (LParamInfo (Rep m))]
-> [SubExp] -> [(Param (LParamInfo (Rep m)), SubExp)]
forall a b. [a] -> [b] -> [(a, b)]
zip [Param (LParamInfo (Rep m))]
fold_params [SubExp]
nes) (((Param (LParamInfo (Rep m)), SubExp) -> m ()) -> m ())
-> ((Param (LParamInfo (Rep m)), SubExp) -> m ()) -> m ()
forall a b. (a -> b) -> a -> b
$ \(Param (LParamInfo (Rep m))
p, SubExp
ne) ->
    [VName] -> Exp (Rep m) -> m ()
forall (m :: * -> *).
MonadBuilder m =>
[VName] -> Exp (Rep m) -> m ()
letBindNames [Param (LParamInfo (Rep m)) -> VName
forall dec. Param dec -> VName
paramName Param (LParamInfo (Rep m))
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 (LParamInfo (Rep m)), VName)]
-> ((Param (LParamInfo (Rep m)), VName) -> m ()) -> m ()
forall (t :: * -> *) (m :: * -> *) a b.
(Foldable t, Monad m) =>
t a -> (a -> m b) -> m ()
forM_ ([Param (LParamInfo (Rep m))]
-> [VName] -> [(Param (LParamInfo (Rep m)), VName)]
forall a b. [a] -> [b] -> [(a, b)]
zip [Param (LParamInfo (Rep m))]
arr_params [VName]
arrs) (((Param (LParamInfo (Rep m)), VName) -> m ()) -> m ())
-> ((Param (LParamInfo (Rep m)), VName) -> m ()) -> m ()
forall a b. (a -> b) -> a -> b
$ \(Param (LParamInfo (Rep m))
p, VName
arr) ->
    [VName] -> Exp (Rep m) -> m ()
forall (m :: * -> *).
MonadBuilder m =>
[VName] -> Exp (Rep m) -> m ()
letBindNames [Param (LParamInfo (Rep m)) -> VName
forall dec. Param dec -> VName
paramName Param (LParamInfo (Rep m))
p] (Exp (Rep m) -> m ())
-> (BasicOp -> Exp (Rep m)) -> BasicOp -> m ()
forall b c a. (b -> c) -> (a -> b) -> a -> c
. BasicOp -> Exp (Rep m)
forall rep. BasicOp -> Exp rep
BasicOp (BasicOp -> m ()) -> BasicOp -> m ()
forall a b. (a -> b) -> a -> b
$
      ReshapeKind -> ShapeBase SubExp -> VName -> BasicOp
Reshape ReshapeKind
ReshapeCoerce (Type -> ShapeBase SubExp
forall shape u. ArrayShape shape => TypeBase shape u -> shape
arrayShape (Type -> ShapeBase SubExp) -> Type -> ShapeBase SubExp
forall a b. (a -> b) -> a -> b
$ Param (LParamInfo (Rep m)) -> Type
forall dec. Typed dec => Param dec -> Type
paramType Param (LParamInfo (Rep m))
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
$ Body (Rep m) -> Seq (Stm (Rep m))
forall rep. Body rep -> Stms rep
bodyStms (Body (Rep m) -> Seq (Stm (Rep m)))
-> Body (Rep m) -> Seq (Stm (Rep m))
forall a b. (a -> b) -> a -> b
$ Lambda (Rep m) -> Body (Rep m)
forall rep. Lambda rep -> 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
$ Body (Rep m) -> [SubExpRes]
forall rep. Body rep -> [SubExpRes]
bodyResult (Body (Rep m) -> [SubExpRes]) -> Body (Rep m) -> [SubExpRes]
forall a b. (a -> b) -> a -> b
$ Lambda (Rep m) -> Body (Rep m)
forall rep. Lambda rep -> 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 u. TypeBase (ShapeBase SubExp) u -> [SubExp]
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
$ 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
$ ReshapeKind -> ShapeBase SubExp -> VName -> BasicOp
Reshape ReshapeKind
ReshapeCoerce ([SubExp] -> ShapeBase SubExp
forall d. [d] -> ShapeBase d
Shape [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
_ [] =
  [Char] -> (Param dec, [Param dec], [Param dec])
forall a. HasCallStack => [Char] -> a
error [Char]
"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)