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

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

import Control.Monad.Identity
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@.
--
-- Only handles a pattern with an empty 'patternContextElements'.
redomapToMapAndReduce ::
  ( MonadFreshNames m,
    Bindable rep,
    ExpDec rep ~ (),
    Op rep ~ SOAC rep
  ) =>
  Pattern rep ->
  ( SubExp,
    Commutativity,
    LambdaT rep,
    LambdaT rep,
    [SubExp],
    [VName]
  ) ->
  m (Stm rep, Stm rep)
redomapToMapAndReduce :: forall (m :: * -> *) rep.
(MonadFreshNames m, Bindable rep, ExpDec rep ~ (),
 Op rep ~ SOAC rep) =>
Pattern rep
-> (SubExp, Commutativity, LambdaT rep, LambdaT rep, [SubExp],
    [VName])
-> m (Stm rep, Stm rep)
redomapToMapAndReduce
  (Pattern [] [PatElemT (LetDec rep)]
patelems)
  (SubExp
w, Commutativity
comm, LambdaT rep
redlam, LambdaT rep
map_lam, [SubExp]
accs, [VName]
arrs) = do
    ([Ident]
map_pat, PatternT (LetDec rep)
red_pat, [(SubExp, VName)]
red_args) <-
      [PatElemT (LetDec rep)]
-> SubExp
-> LambdaT rep
-> [SubExp]
-> m ([Ident], PatternT (LetDec rep), [(SubExp, VName)])
forall dec (m :: * -> *) rep.
(Typed dec, MonadFreshNames m) =>
[PatElemT dec]
-> SubExp
-> LambdaT rep
-> [SubExp]
-> m ([Ident], PatternT dec, [(SubExp, VName)])
splitScanOrRedomap [PatElemT (LetDec rep)]
patelems SubExp
w LambdaT rep
map_lam [SubExp]
accs
    let map_bnd :: Stm rep
map_bnd = [Ident] -> [Ident] -> Exp rep -> Stm rep
forall rep.
Bindable rep =>
[Ident] -> [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 -> ExpT 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 (LambdaT rep -> ScremaForm rep
forall rep. Lambda rep -> ScremaForm rep
mapSOAC LambdaT rep
map_lam)
        ([SubExp]
nes, [VName]
red_arrs) = [(SubExp, VName)] -> ([SubExp], [VName])
forall a b. [(a, b)] -> ([a], [b])
unzip [(SubExp, VName)]
red_args
    Stm rep
red_bnd <-
      PatternT (LetDec rep) -> StmAux (ExpDec rep) -> Exp rep -> Stm rep
forall rep.
Pattern rep -> StmAux (ExpDec rep) -> Exp rep -> Stm rep
Let PatternT (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
. SOAC rep -> Exp rep
forall rep. Op rep -> ExpT 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 :: * -> *).
(Bindable rep, MonadFreshNames m) =>
[Reduce rep] -> m (ScremaForm rep)
reduceSOAC [Commutativity -> LambdaT rep -> [SubExp] -> Reduce rep
forall rep. Commutativity -> Lambda rep -> [SubExp] -> Reduce rep
Reduce Commutativity
comm LambdaT rep
redlam [SubExp]
nes])
    (Stm rep, Stm rep) -> m (Stm rep, Stm rep)
forall (m :: * -> *) a. Monad m => a -> m a
return (Stm rep
map_bnd, Stm rep
red_bnd)
redomapToMapAndReduce PatternT (LetDec rep)
_ (SubExp, Commutativity, LambdaT rep, LambdaT rep, [SubExp],
 [VName])
_ =
  [Char] -> m (Stm rep, Stm rep)
forall a. HasCallStack => [Char] -> a
error [Char]
"redomapToMapAndReduce does not handle a non-empty 'patternContextElements'"

splitScanOrRedomap ::
  (Typed dec, MonadFreshNames m) =>
  [PatElemT dec] ->
  SubExp ->
  LambdaT rep ->
  [SubExp] ->
  m ([Ident], PatternT dec, [(SubExp, VName)])
splitScanOrRedomap :: forall dec (m :: * -> *) rep.
(Typed dec, MonadFreshNames m) =>
[PatElemT dec]
-> SubExp
-> LambdaT rep
-> [SubExp]
-> m ([Ident], PatternT dec, [(SubExp, VName)])
splitScanOrRedomap [PatElemT dec]
patelems SubExp
w LambdaT rep
map_lam [SubExp]
accs = do
  let ([PatElemT dec]
acc_patelems, [PatElemT dec]
arr_patelems) = Int -> [PatElemT dec] -> ([PatElemT dec], [PatElemT dec])
forall a. Int -> [a] -> ([a], [a])
splitAt ([SubExp] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length [SubExp]
accs) [PatElemT dec]
patelems
      ([Type]
acc_ts, [Type]
_arr_ts) = Int -> [Type] -> ([Type], [Type])
forall a. Int -> [a] -> ([a], [a])
splitAt ([SubExp] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length [SubExp]
accs) ([Type] -> ([Type], [Type])) -> [Type] -> ([Type], [Type])
forall a b. (a -> b) -> a -> b
$ LambdaT rep -> [Type]
forall rep. LambdaT rep -> [Type]
lambdaReturnType LambdaT rep
map_lam
  [Ident]
map_accpat <- (PatElemT dec -> Type -> m Ident)
-> [PatElemT dec] -> [Type] -> m [Ident]
forall (m :: * -> *) a b c.
Applicative m =>
(a -> b -> m c) -> [a] -> [b] -> m [c]
zipWithM PatElemT dec -> Type -> m Ident
accMapPatElem [PatElemT dec]
acc_patelems [Type]
acc_ts
  [Ident]
map_arrpat <- (PatElemT dec -> m Ident) -> [PatElemT dec] -> m [Ident]
forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
(a -> m b) -> t a -> m (t b)
mapM PatElemT dec -> m Ident
arrMapPatElem [PatElemT dec]
arr_patelems
  let map_pat :: [Ident]
map_pat = [Ident]
map_accpat [Ident] -> [Ident] -> [Ident]
forall a. [a] -> [a] -> [a]
++ [Ident]
map_arrpat
      red_args :: [(SubExp, VName)]
red_args = [SubExp] -> [VName] -> [(SubExp, VName)]
forall a b. [a] -> [b] -> [(a, b)]
zip [SubExp]
accs ([VName] -> [(SubExp, VName)]) -> [VName] -> [(SubExp, VName)]
forall a b. (a -> b) -> a -> b
$ (Ident -> VName) -> [Ident] -> [VName]
forall a b. (a -> b) -> [a] -> [b]
map Ident -> VName
identName [Ident]
map_accpat
  ([Ident], PatternT dec, [(SubExp, VName)])
-> m ([Ident], PatternT dec, [(SubExp, VName)])
forall (m :: * -> *) a. Monad m => a -> m a
return ([Ident]
map_pat, [PatElemT dec] -> [PatElemT dec] -> PatternT dec
forall dec. [PatElemT dec] -> [PatElemT dec] -> PatternT dec
Pattern [] [PatElemT dec]
acc_patelems, [(SubExp, VName)]
red_args)
  where
    accMapPatElem :: PatElemT dec -> Type -> m Ident
accMapPatElem PatElemT dec
pe Type
acc_t =
      [Char] -> Type -> m Ident
forall (m :: * -> *).
MonadFreshNames m =>
[Char] -> Type -> m Ident
newIdent (VName -> [Char]
baseString (PatElemT dec -> VName
forall dec. PatElemT dec -> VName
patElemName PatElemT 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 :: PatElemT dec -> m Ident
arrMapPatElem = Ident -> m Ident
forall (m :: * -> *) a. Monad m => a -> m a
return (Ident -> m Ident)
-> (PatElemT dec -> Ident) -> PatElemT dec -> m Ident
forall b c a. (b -> c) -> (a -> b) -> a -> c
. PatElemT dec -> Ident
forall dec. Typed dec => PatElemT 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 ::
  ( MonadBinder m,
    Op (Rep m) ~ SOAC (Rep m),
    Bindable (Rep m)
  ) =>
  Pattern (Rep m) ->
  SubExp ->
  ScremaForm (Rep m) ->
  [VName] ->
  m ()
dissectScrema :: forall (m :: * -> *).
(MonadBinder m, Op (Rep m) ~ SOAC (Rep m), Bindable (Rep m)) =>
Pattern (Rep m) -> SubExp -> ScremaForm (Rep m) -> [VName] -> m ()
dissectScrema Pattern (Rep m)
pat SubExp
w (ScremaForm [Scan (Rep m)]
scans [Reduce (Rep m)]
reds Lambda (Rep m)
map_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
      ([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
$ Pattern (Rep m) -> [VName]
forall dec. PatternT dec -> [VName]
patternNames Pattern (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 :: * -> *).
MonadBinder 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 -> ExpT 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]
arrs ScremaForm (Rep m)
scanomap

  ScremaForm (Rep m)
reduce <- [Reduce (Rep m)] -> m (ScremaForm (Rep m))
forall rep (m :: * -> *).
(Bindable rep, MonadFreshNames m) =>
[Reduce rep] -> m (ScremaForm rep)
reduceSOAC [Reduce (Rep m)]
reds
  [VName] -> Exp (Rep m) -> m ()
forall (m :: * -> *).
MonadBinder 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 -> ExpT 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 ::
  (MonadBinder m, Bindable (Rep m)) =>
  Pattern (Rep m) ->
  SubExp ->
  [SubExp] ->
  LambdaT (Rep m) ->
  [VName] ->
  m ()
sequentialStreamWholeArray :: forall (m :: * -> *).
(MonadBinder m, Bindable (Rep m)) =>
Pattern (Rep m)
-> SubExp -> [SubExp] -> LambdaT (Rep m) -> [VName] -> m ()
sequentialStreamWholeArray Pattern (Rep m)
pat SubExp
w [SubExp]
nes LambdaT (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 (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
$ LambdaT (Rep m) -> [LParam (Rep m)]
forall rep. LambdaT rep -> [LParam rep]
lambdaParams LambdaT (Rep m)
lam

  -- The chunk size is the full size of the array.
  [VName] -> Exp (Rep m) -> m ()
forall (m :: * -> *).
MonadBinder 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 -> ExpT 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 :: * -> *).
MonadBinder 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 -> ExpT 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 :: * -> *).
MonadBinder 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 -> ExpT rep
BasicOp (BasicOp -> Exp (Rep m)) -> BasicOp -> Exp (Rep m)
forall a b. (a -> b) -> a -> b
$
        ShapeChange SubExp -> VName -> BasicOp
Reshape ((SubExp -> DimChange SubExp) -> [SubExp] -> ShapeChange SubExp
forall a b. (a -> b) -> [a] -> [b]
map SubExp -> DimChange SubExp
forall d. d -> DimChange d
DimCoercion ([SubExp] -> ShapeChange SubExp) -> [SubExp] -> ShapeChange 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
$ 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 :: * -> *). MonadBinder m => Stm (Rep m) -> m ()
addStm (Seq (Stm (Rep m)) -> m ()) -> Seq (Stm (Rep m)) -> m ()
forall a b. (a -> b) -> a -> b
$ BodyT (Rep m) -> Seq (Stm (Rep m))
forall rep. BodyT rep -> Stms rep
bodyStms (BodyT (Rep m) -> Seq (Stm (Rep m)))
-> BodyT (Rep m) -> Seq (Stm (Rep m))
forall a b. (a -> b) -> a -> b
$ LambdaT (Rep m) -> BodyT (Rep m)
forall rep. LambdaT rep -> BodyT rep
lambdaBody LambdaT (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.
  [(PatElemT (LetDec (Rep m)), SubExp)]
-> ((PatElemT (LetDec (Rep m)), SubExp) -> m ()) -> m ()
forall (t :: * -> *) (m :: * -> *) a b.
(Foldable t, Monad m) =>
t a -> (a -> m b) -> m ()
forM_ ([PatElemT (LetDec (Rep m))]
-> [SubExp] -> [(PatElemT (LetDec (Rep m)), SubExp)]
forall a b. [a] -> [b] -> [(a, b)]
zip (Pattern (Rep m) -> [PatElemT (LetDec (Rep m))]
forall dec. PatternT dec -> [PatElemT dec]
patternElements Pattern (Rep m)
pat) ([SubExp] -> [(PatElemT (LetDec (Rep m)), SubExp)])
-> [SubExp] -> [(PatElemT (LetDec (Rep m)), SubExp)]
forall a b. (a -> b) -> a -> b
$ BodyT (Rep m) -> [SubExp]
forall rep. BodyT rep -> [SubExp]
bodyResult (BodyT (Rep m) -> [SubExp]) -> BodyT (Rep m) -> [SubExp]
forall a b. (a -> b) -> a -> b
$ LambdaT (Rep m) -> BodyT (Rep m)
forall rep. LambdaT rep -> BodyT rep
lambdaBody LambdaT (Rep m)
lam) (((PatElemT (LetDec (Rep m)), SubExp) -> m ()) -> m ())
-> ((PatElemT (LetDec (Rep m)), SubExp) -> m ()) -> m ()
forall a b. (a -> b) -> a -> b
$ \(PatElemT (LetDec (Rep m))
pe, SubExp
se) ->
    case (Type -> [SubExp]
forall u. TypeBase (ShapeBase SubExp) u -> [SubExp]
arrayDims (Type -> [SubExp]) -> Type -> [SubExp]
forall a b. (a -> b) -> a -> b
$ PatElemT (LetDec (Rep m)) -> Type
forall dec. Typed dec => PatElemT dec -> Type
patElemType PatElemT (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 (t :: * -> *) a. Foldable t => t a -> Bool
null [SubExp]
dims ->
          [VName] -> Exp (Rep m) -> m ()
forall (m :: * -> *).
MonadBinder m =>
[VName] -> Exp (Rep m) -> m ()
letBindNames [PatElemT (LetDec (Rep m)) -> VName
forall dec. PatElemT dec -> VName
patElemName PatElemT (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 -> ExpT rep
BasicOp (BasicOp -> Exp (Rep m)) -> BasicOp -> Exp (Rep m)
forall a b. (a -> b) -> a -> b
$ ShapeChange SubExp -> VName -> BasicOp
Reshape ((SubExp -> DimChange SubExp) -> [SubExp] -> ShapeChange SubExp
forall a b. (a -> b) -> [a] -> [b]
map SubExp -> DimChange SubExp
forall d. d -> DimChange d
DimCoercion [SubExp]
dims) VName
v
      ([SubExp], SubExp)
_ -> [VName] -> Exp (Rep m) -> m ()
forall (m :: * -> *).
MonadBinder m =>
[VName] -> Exp (Rep m) -> m ()
letBindNames [PatElemT (LetDec (Rep m)) -> VName
forall dec. PatElemT dec -> VName
patElemName PatElemT (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 -> ExpT 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)