{-# LANGUAGE TypeFamilies #-}

-- | The code generator cannot handle the array combinators (@map@ and
-- friends), so this module was written to transform them into the
-- equivalent do-loops.  The transformation is currently rather naive,
-- and - it's certainly worth considering when we can express such
-- transformations in-place.
module Futhark.Transform.FirstOrderTransform
  ( transformFunDef,
    transformConsts,
    FirstOrderRep,
    Transformer,
    transformStmRecursively,
    transformLambda,
    transformSOAC,
    transformScrema,
    transformFlatMap,
  )
where

import Control.Monad
import Control.Monad.State
import Data.List (find, uncons, zip4)
import Data.Map.Strict qualified as M
import Data.Maybe
import Futhark.Analysis.Alias qualified as Alias
import Futhark.IR qualified as AST
import Futhark.IR.Prop.Aliases
import Futhark.IR.SOACS
import Futhark.MonadFreshNames
import Futhark.Tools
import Futhark.Util (chunks, splitAt3)

-- | The constraints that must hold for a rep in order to be the
-- target of first-order transformation.
type FirstOrderRep rep =
  ( Buildable rep,
    BuilderOps rep,
    LetDec SOACS ~ LetDec rep,
    LParamInfo SOACS ~ LParamInfo rep,
    Alias.AliasableRep rep
  )

-- | First-order-transform a single function, with the given scope
-- provided by top-level constants.
transformFunDef ::
  (MonadFreshNames m, FirstOrderRep torep) =>
  Scope torep ->
  FunDef SOACS ->
  m (AST.FunDef torep)
transformFunDef :: forall (m :: * -> *) torep.
(MonadFreshNames m, FirstOrderRep torep) =>
Scope torep -> FunDef SOACS -> m (FunDef torep)
transformFunDef Scope torep
consts_scope (FunDef Maybe EntryPoint
entry Attrs
attrs Name
fname [(RetType SOACS, RetAls)]
rettype [FParam SOACS]
params Body SOACS
body) = do
  (body', _) <- (VNameSource -> ((Body torep, Stms torep), VNameSource))
-> m (Body torep, Stms torep)
forall (m :: * -> *) a.
MonadFreshNames m =>
(VNameSource -> (a, VNameSource)) -> m a
modifyNameSource ((VNameSource -> ((Body torep, Stms torep), VNameSource))
 -> m (Body torep, Stms torep))
-> (VNameSource -> ((Body torep, Stms torep), VNameSource))
-> m (Body torep, Stms torep)
forall a b. (a -> b) -> a -> b
$ State VNameSource (Body torep, Stms torep)
-> VNameSource -> ((Body torep, Stms torep), VNameSource)
forall s a. State s a -> s -> (a, s)
runState (State VNameSource (Body torep, Stms torep)
 -> VNameSource -> ((Body torep, Stms torep), VNameSource))
-> State VNameSource (Body torep, Stms torep)
-> VNameSource
-> ((Body torep, Stms torep), VNameSource)
forall a b. (a -> b) -> a -> b
$ BuilderT torep (StateT VNameSource Identity) (Body torep)
-> Scope torep -> State VNameSource (Body torep, Stms torep)
forall (m :: * -> *) rep a.
MonadFreshNames m =>
BuilderT rep m a -> Scope rep -> m (a, Stms rep)
runBuilderT BuilderT torep (StateT VNameSource Identity) (Body torep)
BuilderT
  torep
  (StateT VNameSource Identity)
  (Body (Rep (BuilderT torep (StateT VNameSource Identity))))
m Scope torep
consts_scope
  pure $ FunDef entry attrs fname rettype params body'
  where
    m :: BuilderT
  torep
  (StateT VNameSource Identity)
  (Body (Rep (BuilderT torep (StateT VNameSource Identity))))
m = Scope torep
-> BuilderT
     torep
     (StateT VNameSource Identity)
     (Body (Rep (BuilderT torep (StateT VNameSource Identity))))
-> BuilderT
     torep
     (StateT VNameSource Identity)
     (Body (Rep (BuilderT torep (StateT VNameSource Identity))))
forall a.
Scope torep
-> BuilderT torep (StateT VNameSource Identity) a
-> BuilderT torep (StateT VNameSource Identity) a
forall rep (m :: * -> *) a.
LocalScope rep m =>
Scope rep -> m a -> m a
localScope ([FParam torep] -> Scope torep
forall rep dec. (FParamInfo rep ~ dec) => [Param dec] -> Scope rep
scopeOfFParams [FParam torep]
[FParam SOACS]
params) (BuilderT
   torep
   (StateT VNameSource Identity)
   (Body (Rep (BuilderT torep (StateT VNameSource Identity))))
 -> BuilderT
      torep
      (StateT VNameSource Identity)
      (Body (Rep (BuilderT torep (StateT VNameSource Identity)))))
-> BuilderT
     torep
     (StateT VNameSource Identity)
     (Body (Rep (BuilderT torep (StateT VNameSource Identity))))
-> BuilderT
     torep
     (StateT VNameSource Identity)
     (Body (Rep (BuilderT torep (StateT VNameSource Identity))))
forall a b. (a -> b) -> a -> b
$ Body SOACS
-> BuilderT
     torep
     (StateT VNameSource Identity)
     (Body (Rep (BuilderT torep (StateT VNameSource Identity))))
forall (m :: * -> *).
(Transformer m, LetDec (Rep m) ~ LetDec SOACS) =>
Body SOACS -> m (Body (Rep m))
transformBody Body SOACS
body

-- | First-order-transform these top-level constants.
transformConsts ::
  (MonadFreshNames m, FirstOrderRep torep) =>
  Stms SOACS ->
  m (AST.Stms torep)
transformConsts :: forall (m :: * -> *) torep.
(MonadFreshNames m, FirstOrderRep torep) =>
Stms SOACS -> m (Stms torep)
transformConsts Stms SOACS
stms =
  (((), Stms torep) -> Stms torep)
-> m ((), Stms torep) -> m (Stms torep)
forall a b. (a -> b) -> m a -> m b
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap ((), Stms torep) -> Stms torep
forall a b. (a, b) -> b
snd (m ((), Stms torep) -> m (Stms torep))
-> m ((), Stms torep) -> m (Stms torep)
forall a b. (a -> b) -> a -> b
$ (VNameSource -> (((), Stms torep), VNameSource))
-> m ((), Stms torep)
forall (m :: * -> *) a.
MonadFreshNames m =>
(VNameSource -> (a, VNameSource)) -> m a
modifyNameSource ((VNameSource -> (((), Stms torep), VNameSource))
 -> m ((), Stms torep))
-> (VNameSource -> (((), Stms torep), VNameSource))
-> m ((), Stms torep)
forall a b. (a -> b) -> a -> b
$ State VNameSource ((), Stms torep)
-> VNameSource -> (((), Stms torep), VNameSource)
forall s a. State s a -> s -> (a, s)
runState (State VNameSource ((), Stms torep)
 -> VNameSource -> (((), Stms torep), VNameSource))
-> State VNameSource ((), Stms torep)
-> VNameSource
-> (((), Stms torep), VNameSource)
forall a b. (a -> b) -> a -> b
$ BuilderT torep (StateT VNameSource Identity) ()
-> Scope torep -> State VNameSource ((), Stms torep)
forall (m :: * -> *) rep a.
MonadFreshNames m =>
BuilderT rep m a -> Scope rep -> m (a, Stms rep)
runBuilderT BuilderT torep (StateT VNameSource Identity) ()
m Scope torep
forall a. Monoid a => a
mempty
  where
    m :: BuilderT torep (StateT VNameSource Identity) ()
m = (Stm SOACS -> BuilderT torep (StateT VNameSource Identity) ())
-> Stms SOACS -> BuilderT torep (StateT VNameSource Identity) ()
forall (t :: * -> *) (m :: * -> *) a b.
(Foldable t, Monad m) =>
(a -> m b) -> t a -> m ()
mapM_ Stm SOACS -> BuilderT torep (StateT VNameSource Identity) ()
forall (m :: * -> *).
(Transformer m, LetDec (Rep m) ~ LetDec SOACS) =>
Stm SOACS -> m ()
transformStmRecursively Stms SOACS
stms

-- | The constraints that a monad must uphold in order to be used for
-- first-order transformation.
type Transformer m =
  ( MonadBuilder m,
    LocalScope (Rep m) m,
    Buildable (Rep m),
    BuilderOps (Rep m),
    LParamInfo SOACS ~ LParamInfo (Rep m),
    Alias.AliasableRep (Rep m)
  )

transformBody ::
  (Transformer m, LetDec (Rep m) ~ LetDec SOACS) =>
  Body SOACS ->
  m (AST.Body (Rep m))
transformBody :: forall (m :: * -> *).
(Transformer m, LetDec (Rep m) ~ LetDec SOACS) =>
Body SOACS -> m (Body (Rep m))
transformBody (Body () Stms SOACS
stms [SubExpRes]
res) = m [SubExpRes] -> m (GBody (Rep m) SubExpRes)
forall (m :: * -> *) res.
(MonadBuilder m, IsResult res) =>
m [res] -> m (GBody (Rep m) res)
buildBody_ (m [SubExpRes] -> m (GBody (Rep m) SubExpRes))
-> m [SubExpRes] -> m (GBody (Rep m) SubExpRes)
forall a b. (a -> b) -> a -> b
$ do
  (Stm SOACS -> m ()) -> Stms SOACS -> m ()
forall (t :: * -> *) (m :: * -> *) a b.
(Foldable t, Monad m) =>
(a -> m b) -> t a -> m ()
mapM_ Stm SOACS -> m ()
forall (m :: * -> *).
(Transformer m, LetDec (Rep m) ~ LetDec SOACS) =>
Stm SOACS -> m ()
transformStmRecursively Stms SOACS
stms
  [SubExpRes] -> m [SubExpRes]
forall a. a -> m a
forall (f :: * -> *) a. Applicative f => a -> f a
pure [SubExpRes]
res

-- | First transform any nested t'Body' or t'Lambda' elements, then
-- apply 'transformSOAC' if the expression is a SOAC.
transformStmRecursively ::
  (Transformer m, LetDec (Rep m) ~ LetDec SOACS) => Stm SOACS -> m ()
transformStmRecursively :: forall (m :: * -> *).
(Transformer m, LetDec (Rep m) ~ LetDec SOACS) =>
Stm SOACS -> m ()
transformStmRecursively (Let Pat (LetDec SOACS)
pat StmAux (ExpDec SOACS)
aux (Op Op SOACS
soac)) =
  StmAux () -> m () -> m ()
forall (m :: * -> *) anyrep a.
MonadBuilder m =>
StmAux anyrep -> m a -> m a
auxing StmAux ()
StmAux (ExpDec SOACS)
aux (m () -> m ()) -> m () -> m ()
forall a b. (a -> b) -> a -> b
$ Pat (LetDec (Rep m)) -> SOAC (Rep m) -> m ()
forall (m :: * -> *).
Transformer m =>
Pat (LetDec (Rep m)) -> SOAC (Rep m) -> m ()
transformSOAC Pat (LetDec (Rep m))
Pat (LetDec SOACS)
pat (SOAC (Rep m) -> m ()) -> m (SOAC (Rep m)) -> m ()
forall (m :: * -> *) a b. Monad m => (a -> m b) -> m a -> m b
=<< SOACMapper SOACS (Rep m) m -> SOAC SOACS -> m (SOAC (Rep m))
forall (m :: * -> *) frep trep.
Monad m =>
SOACMapper frep trep m -> SOAC frep -> m (SOAC trep)
mapSOACM SOACMapper SOACS (Rep m) m
soacTransform Op SOACS
SOAC SOACS
soac
  where
    soacTransform :: SOACMapper SOACS (Rep m) m
soacTransform =
      SOACMapper (ZonkAny 0) (ZonkAny 0) m
forall rep (m :: * -> *). Monad m => SOACMapper rep rep m
identitySOACMapper
        { mapOnSOACLambda = transformLambda,
          mapOnSOACExtLambda = transformLambda
        }
transformStmRecursively (Let Pat (LetDec SOACS)
pat StmAux (ExpDec SOACS)
aux Exp SOACS
e) =
  StmAux () -> m () -> m ()
forall (m :: * -> *) anyrep a.
MonadBuilder m =>
StmAux anyrep -> m a -> m a
auxing StmAux ()
StmAux (ExpDec SOACS)
aux (m () -> m ()) -> m () -> m ()
forall a b. (a -> b) -> a -> b
$ Pat (LetDec (Rep m)) -> Exp (Rep m) -> m ()
forall (m :: * -> *).
MonadBuilder m =>
Pat (LetDec (Rep m)) -> Exp (Rep m) -> m ()
letBind Pat (LetDec (Rep m))
Pat (LetDec SOACS)
pat (Exp (Rep m) -> m ()) -> m (Exp (Rep m)) -> m ()
forall (m :: * -> *) a b. Monad m => (a -> m b) -> m a -> m b
=<< Mapper SOACS (Rep m) m -> Exp SOACS -> m (Exp (Rep m))
forall (m :: * -> *) frep trep.
Monad m =>
Mapper frep trep m -> Exp frep -> m (Exp trep)
mapExpM Mapper SOACS (Rep m) m
transform Exp SOACS
e
  where
    transform :: Mapper SOACS (Rep m) m
transform =
      Mapper (ZonkAny 1) (ZonkAny 1) m
forall rep (m :: * -> *). Monad m => Mapper rep rep m
identityMapper
        { mapOnBody = \Scope (Rep m)
scope -> Scope (Rep m) -> m (Body (Rep m)) -> m (Body (Rep m))
forall a. Scope (Rep m) -> m a -> m a
forall rep (m :: * -> *) a.
LocalScope rep m =>
Scope rep -> m a -> m a
localScope Scope (Rep m)
scope (m (Body (Rep m)) -> m (Body (Rep m)))
-> (Body SOACS -> m (Body (Rep m)))
-> Body SOACS
-> m (Body (Rep m))
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Body SOACS -> m (Body (Rep m))
forall (m :: * -> *).
(Transformer m, LetDec (Rep m) ~ LetDec SOACS) =>
Body SOACS -> m (Body (Rep m))
transformBody,
          mapOnRetType = pure,
          mapOnBranchType = pure,
          mapOnFParam = pure,
          mapOnLParam = pure,
          mapOnOp = error "Unhandled Op in first order transform"
        }

-- Produce scratch "arrays" for the Map and Scan outputs of Screma.
-- "Arrays" is in quotes because some of those may be accumulators.
resultArray :: (Transformer m) => [VName] -> [Type] -> m [VName]
resultArray :: forall (m :: * -> *).
Transformer m =>
[VName] -> [Type] -> m [VName]
resultArray [VName]
arrs [Type]
ts = do
  arrs_ts <- (VName -> m Type) -> [VName] -> m [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 VName -> m Type
forall rep (m :: * -> *). HasScope rep m => VName -> m Type
lookupType [VName]
arrs
  let oneArray t :: Type
t@Acc {}
        | Just (VName
v, Type
_) <- ((VName, Type) -> Bool) -> [(VName, Type)] -> Maybe (VName, Type)
forall (t :: * -> *) a. Foldable t => (a -> Bool) -> t a -> Maybe a
find ((Type -> Type -> Bool
forall a. Eq a => a -> a -> Bool
== Type
t) (Type -> Bool) -> ((VName, Type) -> Type) -> (VName, Type) -> Bool
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (VName, Type) -> Type
forall a b. (a, b) -> b
snd) ([VName] -> [Type] -> [(VName, Type)]
forall a b. [a] -> [b] -> [(a, b)]
zip [VName]
arrs [Type]
arrs_ts) =
            VName -> m VName
forall a. a -> m a
forall (f :: * -> *) a. Applicative f => a -> f a
pure VName
v
      oneArray Type
t =
        Name -> Exp (Rep m) -> m VName
forall (m :: * -> *).
MonadBuilder m =>
Name -> Exp (Rep m) -> m VName
letExp Name
"result" (Exp (Rep m) -> m VName) -> m (Exp (Rep m)) -> m VName
forall (m :: * -> *) a b. Monad m => (a -> m b) -> m a -> m b
=<< Type -> m (Exp (Rep m))
forall (m :: * -> *). MonadBuilder m => Type -> m (Exp (Rep m))
eBlank Type
t
  mapM oneArray ts

-- | Sequentialise a single FlatMap. The size of the nonuniform results produced
-- by the lambda is not known until it has been run, so each of them is
-- accumulated in a scratch buffer that is doubled whenever it runs out of
-- space, and finally truncated to the actual size. The value results need no
-- such treatment, as there is exactly one per iteration. The shape and offset
-- arrays are filled in as we go, and the flag array is then a scatter of the
-- segment starts.
transformFlatMap ::
  (Transformer m) =>
  Pat (LetDec (Rep m)) ->
  SubExp ->
  [VName] ->
  ExtLambda (Rep m) ->
  m ()
transformFlatMap :: forall (m :: * -> *).
Transformer m =>
Pat (LetDec (Rep m))
-> SubExp -> [VName] -> ExtLambda (Rep m) -> m ()
transformFlatMap Pat (LetDec (Rep m))
pat SubExp
w [VName]
arrs ExtLambda (Rep m)
lam = do
  let irreg_ts :: [Type]
irreg_ts = ExtLambda (Rep m) -> [Type]
forall rep. ExtLambda rep -> [Type]
flatMapRowTypes ExtLambda (Rep m)
lam
      reg_ts :: [Type]
reg_ts = ExtLambda (Rep m) -> [Type]
forall rep. ExtLambda rep -> [Type]
flatMapUniformTypes ExtLambda (Rep m)
lam
  arrs_ts <- (VName -> m Type) -> [VName] -> m [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 VName -> m Type
forall rep (m :: * -> *). HasScope rep m => VName -> m Type
lookupType [VName]
arrs

  -- Loop parameters: the current filled size, the current capacity, the
  -- per-element shape and offset arrays, one scratch buffer per nonuniform
  -- result, and one array per uniform result.
  size_p <- newParam "flatmap_size" $ toDecl (Prim int64) Observe
  cap_p <- newParam "flatmap_cap" $ toDecl (Prim int64) Observe
  shape_p <- newParam "flatmap_shape" $ toDecl (arrayOfRow (Prim int64) w) Consume
  offset_p <- newParam "flatmap_offset" $ toDecl (arrayOfRow (Prim int64) w) Consume
  scratch_ps <-
    forM irreg_ts $ \Type
et ->
      Name -> DeclType -> m (Param DeclType)
forall (m :: * -> *) dec.
MonadFreshNames m =>
Name -> dec -> m (Param dec)
newParam Name
"flatmap_res" (DeclType -> m (Param DeclType)) -> DeclType -> m (Param DeclType)
forall a b. (a -> b) -> a -> b
$ Type -> Diet -> DeclType
forall shape. TypeBase shape NoMode -> Diet -> TypeBase shape Diet
toDecl (Type -> SubExp -> Type
forall d.
TypeBase (ShapeBase d) NoMode -> d -> TypeBase (ShapeBase d) NoMode
arrayOfRow Type
et (VName -> SubExp
Var (Param DeclType -> VName
forall dec. Param dec -> VName
paramName Param DeclType
cap_p))) Diet
Consume
  reg_ps <-
    forM reg_ts $ \Type
rt ->
      Name -> DeclType -> m (Param DeclType)
forall (m :: * -> *) dec.
MonadFreshNames m =>
Name -> dec -> m (Param dec)
newParam Name
"flatmap_reg" (DeclType -> m (Param DeclType)) -> DeclType -> m (Param DeclType)
forall a b. (a -> b) -> a -> b
$ Type -> Diet -> DeclType
forall shape. TypeBase shape NoMode -> Diet -> TypeBase shape Diet
toDecl (Type -> SubExp -> Type
forall d.
TypeBase (ShapeBase d) NoMode -> d -> TypeBase (ShapeBase d) NoMode
arrayOfRow Type
rt SubExp
w) Diet
Consume

  -- The capacity initially matches the input size.
  shape_init <- letExp "flatmap_shape" $ BasicOp $ Scratch int64 [w]
  offset_init <- letExp "flatmap_offset" $ BasicOp $ Scratch int64 [w]
  scratch_init <- forM irreg_ts $ \Type
et -> Name -> Exp (Rep m) -> m VName
forall (m :: * -> *).
MonadBuilder m =>
Name -> Exp (Rep m) -> m VName
letExp Name
"flatmap_res" (Exp (Rep m) -> m VName) -> m (Exp (Rep m)) -> m VName
forall (m :: * -> *) a b. Monad m => (a -> m b) -> m a -> m b
=<< Type -> m (Exp (Rep m))
forall (m :: * -> *). MonadBuilder m => Type -> m (Exp (Rep m))
eBlank (Type -> SubExp -> Type
forall d.
TypeBase (ShapeBase d) NoMode -> d -> TypeBase (ShapeBase d) NoMode
arrayOfRow Type
et SubExp
w)
  reg_init <- forM reg_ts $ \Type
rt -> Name -> Exp (Rep m) -> m VName
forall (m :: * -> *).
MonadBuilder m =>
Name -> Exp (Rep m) -> m VName
letExp Name
"flatmap_reg" (Exp (Rep m) -> m VName) -> m (Exp (Rep m)) -> m VName
forall (m :: * -> *) a b. Monad m => (a -> m b) -> m a -> m b
=<< Type -> m (Exp (Rep m))
forall (m :: * -> *). MonadBuilder m => Type -> m (Exp (Rep m))
eBlank (Type -> SubExp -> Type
forall d.
TypeBase (ShapeBase d) NoMode -> d -> TypeBase (ShapeBase d) NoMode
arrayOfRow Type
rt SubExp
w)

  let merge =
        (Param DeclType
size_p, IntType -> Integer -> SubExp
intConst IntType
Int64 Integer
0)
          (Param DeclType, SubExp)
-> [(Param DeclType, SubExp)] -> [(Param DeclType, SubExp)]
forall a. a -> [a] -> [a]
: (Param DeclType
cap_p, SubExp
w)
          (Param DeclType, SubExp)
-> [(Param DeclType, SubExp)] -> [(Param DeclType, SubExp)]
forall a. a -> [a] -> [a]
: (Param DeclType
shape_p, VName -> SubExp
Var VName
shape_init)
          (Param DeclType, SubExp)
-> [(Param DeclType, SubExp)] -> [(Param DeclType, SubExp)]
forall a. a -> [a] -> [a]
: (Param DeclType
offset_p, VName -> SubExp
Var VName
offset_init)
          (Param DeclType, SubExp)
-> [(Param DeclType, SubExp)] -> [(Param DeclType, SubExp)]
forall a. a -> [a] -> [a]
: [Param DeclType] -> [SubExp] -> [(Param DeclType, SubExp)]
forall a b. [a] -> [b] -> [(a, b)]
zip [Param DeclType]
scratch_ps ((VName -> SubExp) -> [VName] -> [SubExp]
forall a b. (a -> b) -> [a] -> [b]
map VName -> SubExp
Var [VName]
scratch_init)
            [(Param DeclType, SubExp)]
-> [(Param DeclType, SubExp)] -> [(Param DeclType, SubExp)]
forall a. Semigroup a => a -> a -> a
<> [Param DeclType] -> [SubExp] -> [(Param DeclType, SubExp)]
forall a b. [a] -> [b] -> [(a, b)]
zip [Param DeclType]
reg_ps ((VName -> SubExp) -> [VName] -> [SubExp]
forall a b. (a -> b) -> [a] -> [b]
map VName -> SubExp
Var [VName]
reg_init)
      merge_params = ((Param DeclType, SubExp) -> Param DeclType)
-> [(Param DeclType, SubExp)] -> [Param DeclType]
forall a b. (a -> b) -> [a] -> [b]
map (Param DeclType, SubExp) -> Param DeclType
forall a b. (a, b) -> a
fst [(Param DeclType, SubExp)]
merge

  i <- newVName "i"
  let loop_form = VName -> IntType -> SubExp -> LoopForm
ForLoop VName
i IntType
Int64 SubExp
w
      loop_scope = LoopForm -> Scope (Rep m)
forall rep. LoopForm -> Scope rep
scopeOfLoopForm LoopForm
loop_form Scope (Rep m) -> Scope (Rep m) -> Scope (Rep m)
forall a. Semigroup a => a -> a -> a
<> [Param DeclType] -> Scope (Rep m)
forall rep dec. (FParamInfo rep ~ dec) => [Param dec] -> Scope rep
scopeOfFParams [Param DeclType]
merge_params
  loop_body <- runBodyBuilder . localScope loop_scope $ do
    -- Apply the lambda to the current elements.
    let arg VName
arr Type
arr_t = 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
$ VName -> Slice SubExp -> BasicOp
Index VName
arr (Slice SubExp -> BasicOp) -> Slice SubExp -> BasicOp
forall a b. (a -> b) -> a -> b
$ Type -> [DimIndex SubExp] -> Slice SubExp
fullSlice Type
arr_t [SubExp -> DimIndex SubExp
forall d. d -> DimIndex d
DimFix (SubExp -> DimIndex SubExp) -> SubExp -> DimIndex SubExp
forall a b. (a -> b) -> a -> b
$ VName -> SubExp
Var VName
i]
        size = VName -> SubExp
Var (VName -> SubExp) -> VName -> SubExp
forall a b. (a -> b) -> a -> b
$ Param DeclType -> VName
forall dec. Param dec -> VName
paramName Param DeclType
size_p
    lam_res <- map resSubExp <$> bindLambda lam (zipWith arg arrs arrs_ts)

    -- The lambda produces the common length of its segment results first.
    let (k, ys) =
          fromMaybe (error "transformFlatMap: malformed FlatMap.") $
            uncons lam_res
        (irreg_ys, reg_ys) = flatMapSplitValues lam ys
    new_size <-
      letSubExp "flatmap_new_size" . BasicOp $
        BinOp (Add Int64 OverflowUndef) size k
    grow <-
      letSubExp "flatmap_grow" . BasicOp $
        CmpOp (CmpSlt Int64) (Var (paramName cap_p)) new_size

    -- New capacity: double it (but at least fit) when it overflows.
    new_cap <-
      letSubExp "flatmap_new_cap"
        =<< eIf
          (eSubExp grow)
          ( buildBody_ $ do
              doubled <-
                letSubExp "doubled" . BasicOp $
                  BinOp (Mul Int64 OverflowUndef) (Var (paramName cap_p)) (intConst Int64 2)
              fmap (pure . subExpRes) . letSubExp "atleast" . BasicOp $
                BinOp (SMax Int64) doubled new_size
          )
          (buildBody_ $ pure [subExpRes $ Var (paramName cap_p)])

    let lowSlice Type
t =
          Type -> [DimIndex SubExp] -> Slice SubExp
fullSlice Type
t [SubExp -> SubExp -> SubExp -> DimIndex SubExp
forall d. d -> d -> d -> DimIndex d
DimSlice (IntType -> Integer -> SubExp
intConst IntType
Int64 Integer
0) SubExp
size (IntType -> Integer -> SubExp
intConst IntType
Int64 Integer
1)]

    -- Grow (and copy) each scratch buffer when necessary, then write the
    -- new elements at the end.
    scratch_res <- forM (zip3 scratch_ps irreg_ts irreg_ys) $ \(Param DeclType
sp, Type
et, SubExp
ys_j) -> do
      let full_t :: Type
full_t = Type -> SubExp -> Type
forall d.
TypeBase (ShapeBase d) NoMode -> d -> TypeBase (ShapeBase d) NoMode
arrayOfRow Type
et SubExp
new_cap
      base <-
        Name
-> Exp (Rep (BuilderT (Rep m) (State VNameSource)))
-> BuilderT (Rep m) (State VNameSource) VName
forall (m :: * -> *).
MonadBuilder m =>
Name -> Exp (Rep m) -> m VName
letExp Name
"flatmap_grown"
          (Exp (Rep m) -> BuilderT (Rep m) (State VNameSource) VName)
-> BuilderT (Rep m) (State VNameSource) (Exp (Rep m))
-> BuilderT (Rep m) (State VNameSource) VName
forall (m :: * -> *) a b. Monad m => (a -> m b) -> m a -> m b
=<< BuilderT
  (Rep m)
  (State VNameSource)
  (Exp (Rep (BuilderT (Rep m) (State VNameSource))))
-> BuilderT
     (Rep m)
     (State VNameSource)
     (Body (Rep (BuilderT (Rep m) (State VNameSource))))
-> BuilderT
     (Rep m)
     (State VNameSource)
     (Body (Rep (BuilderT (Rep m) (State VNameSource))))
-> BuilderT
     (Rep m)
     (State VNameSource)
     (Exp (Rep (BuilderT (Rep m) (State VNameSource))))
forall (m :: * -> *).
(MonadBuilder m, BranchType (Rep m) ~ ExtType) =>
m (Exp (Rep m))
-> m (Body (Rep m)) -> m (Body (Rep m)) -> m (Exp (Rep m))
eIf
            (SubExp
-> BuilderT
     (Rep m)
     (State VNameSource)
     (Exp (Rep (BuilderT (Rep m) (State VNameSource))))
forall (m :: * -> *). MonadBuilder m => SubExp -> m (Exp (Rep m))
eSubExp SubExp
grow)
            ( Builder (Rep m) [SubExpRes]
-> BuilderT
     (Rep m)
     (State VNameSource)
     (Body (Rep (BuilderT (Rep m) (State VNameSource))))
forall (m :: * -> *) res.
(MonadBuilder m, IsResult res) =>
m [res] -> m (GBody (Rep m) res)
buildBody_ (Builder (Rep m) [SubExpRes]
 -> BuilderT
      (Rep m)
      (State VNameSource)
      (Body (Rep (BuilderT (Rep m) (State VNameSource)))))
-> Builder (Rep m) [SubExpRes]
-> BuilderT
     (Rep m)
     (State VNameSource)
     (Body (Rep (BuilderT (Rep m) (State VNameSource))))
forall a b. (a -> b) -> a -> b
$ do
                fresh <-
                  Name
-> Exp (Rep (BuilderT (Rep m) (State VNameSource)))
-> BuilderT (Rep m) (State VNameSource) VName
forall (m :: * -> *).
MonadBuilder m =>
Name -> Exp (Rep m) -> m VName
letExp Name
"flatmap_fresh" (Exp (Rep m) -> BuilderT (Rep m) (State VNameSource) VName)
-> (BasicOp -> Exp (Rep m))
-> BasicOp
-> BuilderT (Rep m) (State VNameSource) VName
forall b c a. (b -> c) -> (a -> b) -> a -> c
. BasicOp -> Exp (Rep m)
forall rep. BasicOp -> Exp rep
BasicOp (BasicOp -> BuilderT (Rep m) (State VNameSource) VName)
-> BasicOp -> BuilderT (Rep m) (State VNameSource) VName
forall a b. (a -> b) -> a -> b
$
                    PrimType -> [SubExp] -> BasicOp
Scratch (Type -> PrimType
forall shape o. TypeBase shape o -> PrimType
elemType Type
full_t) (Type -> [SubExp]
forall d m. TypeBase (ShapeBase d) m -> [d]
arrayDims Type
full_t)
                old_t <- lookupType $ paramName sp
                copied <-
                  letInPlace "flatmap_fresh" fresh (lowSlice full_t) $
                    BasicOp (Index (paramName sp) (lowSlice old_t))
                pure [varRes copied]
            )
            ( Builder (Rep m) [SubExpRes]
-> BuilderT
     (Rep m)
     (State VNameSource)
     (Body (Rep (BuilderT (Rep m) (State VNameSource))))
forall (m :: * -> *) res.
(MonadBuilder m, IsResult res) =>
m [res] -> m (GBody (Rep m) res)
buildBody_ (Builder (Rep m) [SubExpRes]
 -> BuilderT
      (Rep m)
      (State VNameSource)
      (Body (Rep (BuilderT (Rep m) (State VNameSource)))))
-> Builder (Rep m) [SubExpRes]
-> BuilderT
     (Rep m)
     (State VNameSource)
     (Body (Rep (BuilderT (Rep m) (State VNameSource))))
forall a b. (a -> b) -> a -> b
$
                (VName -> [SubExpRes])
-> BuilderT (Rep m) (State VNameSource) VName
-> Builder (Rep m) [SubExpRes]
forall a b.
(a -> b)
-> BuilderT (Rep m) (State VNameSource) a
-> BuilderT (Rep m) (State VNameSource) b
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap (SubExpRes -> [SubExpRes]
forall a. a -> [a]
forall (f :: * -> *) a. Applicative f => a -> f a
pure (SubExpRes -> [SubExpRes])
-> (VName -> SubExpRes) -> VName -> [SubExpRes]
forall b c a. (b -> c) -> (a -> b) -> a -> c
. VName -> SubExpRes
varRes) (BuilderT (Rep m) (State VNameSource) VName
 -> Builder (Rep m) [SubExpRes])
-> (Exp (Rep m) -> BuilderT (Rep m) (State VNameSource) VName)
-> Exp (Rep m)
-> Builder (Rep m) [SubExpRes]
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Name
-> Exp (Rep (BuilderT (Rep m) (State VNameSource)))
-> BuilderT (Rep m) (State VNameSource) VName
forall (m :: * -> *).
MonadBuilder m =>
Name -> Exp (Rep m) -> m VName
letExp Name
"flatmap_kept" (Exp (Rep m) -> Builder (Rep m) [SubExpRes])
-> Exp (Rep m) -> Builder (Rep m) [SubExpRes]
forall a b. (a -> b) -> a -> b
$
                  [SubExp] -> VName -> Exp (Rep m)
forall rep. [SubExp] -> VName -> Exp rep
shapeCoerce (Type -> [SubExp]
forall d m. TypeBase (ShapeBase d) m -> [d]
arrayDims Type
full_t) (Param DeclType -> VName
forall dec. Param dec -> VName
paramName Param DeclType
sp)
            )
      letInPlace "flatmap_res" base (fullSlice full_t [DimSlice size k (intConst Int64 1)]) $
        BasicOp (SubExp ys_j)

    -- The segment's size and its offset (the running total before it).
    shape' <-
      letInPlace "flatmap_shape" (paramName shape_p) (fullSlice (paramType shape_p) [DimFix $ Var i]) $
        BasicOp (SubExp k)
    offset' <-
      letInPlace "flatmap_offset" (paramName offset_p) (fullSlice (paramType offset_p) [DimFix $ Var i]) $
        BasicOp (SubExp size)

    -- The uniform results are simply written at this iteration's index.
    reg_res <- forM (zip reg_ps reg_ys) $ \(Param DeclType
rp, SubExp
reg_y) ->
      Name
-> VName
-> Slice SubExp
-> Exp (Rep (BuilderT (Rep m) (State VNameSource)))
-> BuilderT (Rep m) (State VNameSource) VName
forall (m :: * -> *).
MonadBuilder m =>
Name -> VName -> Slice SubExp -> Exp (Rep m) -> m VName
letInPlace Name
"flatmap_reg" (Param DeclType -> VName
forall dec. Param dec -> VName
paramName Param DeclType
rp) (Type -> [DimIndex SubExp] -> Slice SubExp
fullSlice (Param DeclType -> Type
forall dec. Typed dec => Param dec -> Type
paramType Param DeclType
rp) [SubExp -> DimIndex SubExp
forall d. d -> DimIndex d
DimFix (SubExp -> DimIndex SubExp) -> SubExp -> DimIndex SubExp
forall a b. (a -> b) -> a -> b
$ VName -> SubExp
Var VName
i]) (Exp (Rep (BuilderT (Rep m) (State VNameSource)))
 -> BuilderT (Rep m) (State VNameSource) VName)
-> Exp (Rep (BuilderT (Rep m) (State VNameSource)))
-> BuilderT (Rep m) (State VNameSource) VName
forall a b. (a -> b) -> a -> b
$
        BasicOp -> Exp (Rep (BuilderT (Rep m) (State VNameSource)))
forall rep. BasicOp -> Exp rep
BasicOp (SubExp -> BasicOp
SubExp SubExp
reg_y)

    pure $
      subExpsRes [new_size, new_cap]
        <> varsRes (shape' : offset' : scratch_res <> reg_res)

  loop_res <- letTupExp "flatmap" $ Loop merge loop_form loop_body
  case (loop_res, patNames pat) of
    (VName
size_res : VName
_cap_res : VName
shape_res : VName
offset_res : [VName]
value_res, VName
m_pat : VName
shape_pat : VName
flag_pat : VName
offset_pat : [VName]
out_pats) -> do
      -- Bind the total size and the shape/offset arrays, then truncate each
      -- buffer. The uniform results are already of the right size.
      [VName] -> Exp (Rep m) -> m ()
forall (m :: * -> *).
MonadBuilder m =>
[VName] -> Exp (Rep m) -> m ()
letBindNames [VName
m_pat] (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 -> BasicOp) -> SubExp -> BasicOp
forall a b. (a -> b) -> a -> b
$ VName -> SubExp
Var VName
size_res
      [VName] -> Exp (Rep m) -> m ()
forall (m :: * -> *).
MonadBuilder m =>
[VName] -> Exp (Rep m) -> m ()
letBindNames [VName
shape_pat] (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 -> BasicOp) -> SubExp -> BasicOp
forall a b. (a -> b) -> a -> b
$ VName -> SubExp
Var VName
shape_res
      [VName] -> Exp (Rep m) -> m ()
forall (m :: * -> *).
MonadBuilder m =>
[VName] -> Exp (Rep m) -> m ()
letBindNames [VName
offset_pat] (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 -> BasicOp) -> SubExp -> BasicOp
forall a b. (a -> b) -> a -> b
$ VName -> SubExp
Var VName
offset_res
      let ([VName]
scratch_res, [VName]
reg_res) = Int -> [VName] -> ([VName], [VName])
forall a. Int -> [a] -> ([a], [a])
splitAt ([Type] -> Int
forall a. [a] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length [Type]
irreg_ts) [VName]
value_res
          ([VName]
data_pats, [VName]
reg_pats) = ExtLambda (Rep m) -> [VName] -> ([VName], [VName])
forall rep a. ExtLambda rep -> [a] -> ([a], [a])
flatMapSplitValues ExtLambda (Rep m)
lam [VName]
out_pats
      [(VName, VName)] -> ((VName, VName) -> m ()) -> m ()
forall (t :: * -> *) (m :: * -> *) a b.
(Foldable t, Monad m) =>
t a -> (a -> m b) -> m ()
forM_ ([VName] -> [VName] -> [(VName, VName)]
forall a b. [a] -> [b] -> [(a, b)]
zip [VName]
data_pats [VName]
scratch_res) (((VName, VName) -> m ()) -> m ())
-> ((VName, VName) -> m ()) -> m ()
forall a b. (a -> b) -> a -> b
$ \(VName
out, VName
scratch) -> do
        scratch_t <- VName -> m Type
forall rep (m :: * -> *). HasScope rep m => VName -> m Type
lookupType VName
scratch
        letBindNames [out] . BasicOp . Index scratch $
          fullSlice scratch_t [DimSlice (intConst Int64 0) (Var m_pat) (intConst Int64 1)]
      [(VName, VName)] -> ((VName, VName) -> m ()) -> m ()
forall (t :: * -> *) (m :: * -> *) a b.
(Foldable t, Monad m) =>
t a -> (a -> m b) -> m ()
forM_ ([VName] -> [VName] -> [(VName, VName)]
forall a b. [a] -> [b] -> [(a, b)]
zip [VName]
reg_pats [VName]
reg_res) (((VName, VName) -> m ()) -> m ())
-> ((VName, VName) -> m ()) -> m ()
forall a b. (a -> b) -> a -> b
$ \(VName
out, VName
reg) ->
        [VName] -> Exp (Rep m) -> m ()
forall (m :: * -> *).
MonadBuilder m =>
[VName] -> Exp (Rep m) -> m ()
letBindNames [VName
out] (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 -> BasicOp) -> SubExp -> BasicOp
forall a b. (a -> b) -> a -> b
$ VName -> SubExp
Var VName
reg
      -- The flag array: scatter a 'true' at the offset of each non-empty
      -- segment, over an otherwise 'false' array.
      VName -> SubExp -> SubExp -> VName -> VName -> m ()
forall (m :: * -> *).
Transformer m =>
VName -> SubExp -> SubExp -> VName -> VName -> m ()
transformFlatMapFlags VName
flag_pat SubExp
w (VName -> SubExp
Var VName
m_pat) VName
shape_res VName
offset_res
    ([VName], [VName])
_ ->
      [Char] -> m ()
forall a. HasCallStack => [Char] -> a
error [Char]
"transformFlatMap: malformed FlatMap."

-- | Compute a 'FlatMap' flag array of length @m@: 'true' at the start of each
-- non-empty segment, 'false' elsewhere. Emitted as a sequential scatter loop.
transformFlatMapFlags ::
  (Transformer m) =>
  VName ->
  SubExp ->
  SubExp ->
  VName ->
  VName ->
  m ()
transformFlatMapFlags :: forall (m :: * -> *).
Transformer m =>
VName -> SubExp -> SubExp -> VName -> VName -> m ()
transformFlatMapFlags VName
flag_pat SubExp
w SubExp
m VName
shape VName
offset = do
  let flag_t :: Type
flag_t = Type -> SubExp -> Type
forall d.
TypeBase (ShapeBase d) NoMode -> d -> TypeBase (ShapeBase d) NoMode
arrayOfRow (PrimType -> Type
forall shape o. PrimType -> TypeBase shape o
Prim PrimType
Bool) SubExp
m
  flags_init <- Name -> Exp (Rep m) -> m VName
forall (m :: * -> *).
MonadBuilder m =>
Name -> Exp (Rep m) -> m VName
letExp Name
"flatmap_flags" (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
$ ShapeBase SubExp -> SubExp -> BasicOp
Replicate ([SubExp] -> ShapeBase SubExp
forall d. [d] -> ShapeBase d
Shape [SubExp
m]) (Bool -> SubExp
forall v. IsValue v => v -> SubExp
constant Bool
False)
  flags_p <- newParam "flatmap_flags" $ toDecl flag_t Consume
  j <- newVName "j"
  let flag_form = VName -> IntType -> SubExp -> LoopForm
ForLoop VName
j IntType
Int64 SubExp
w
  shape_t <- lookupType shape
  offset_t <- lookupType offset
  flag_body <- runBodyBuilder $
    localScope (scopeOfLoopForm flag_form <> scopeOfFParams [flags_p]) $ do
      sz <- letSubExp "flatmap_sz" $ BasicOp $ Index shape $ fullSlice shape_t [DimFix $ Var j]
      off <- letSubExp "flatmap_off" $ BasicOp $ Index offset $ fullSlice offset_t [DimFix $ Var j]
      nonempty <-
        letSubExp "flatmap_nonempty" . BasicOp $
          CmpOp (CmpSlt Int64) (intConst Int64 0) sz
      flags' <-
        letSubExp "flatmap_flags"
          =<< eIf
            (eSubExp nonempty)
            ( buildBody_
                $ fmap (pure . varRes)
                  . letInPlace "flatmap_flags" (paramName flags_p) (fullSlice flag_t [DimFix off])
                $ BasicOp (SubExp (constant True))
            )
            (buildBody_ $ pure [varRes $ paramName flags_p])
      pure [subExpRes flags']
  letBindNames [flag_pat] $ Loop [(flags_p, Var flags_init)] flag_form flag_body

-- | Sequentialise a single Screma.
transformScrema ::
  (Transformer m) =>
  Pat dec ->
  SubExp ->
  [VName] ->
  ScremaForm (Rep m) ->
  m ()
transformScrema :: forall (m :: * -> *) dec.
Transformer m =>
Pat dec -> SubExp -> [VName] -> ScremaForm (Rep m) -> m ()
transformScrema Pat dec
pat SubExp
w [VName]
arrs form :: ScremaForm (Rep m)
form@(ScremaForm Lambda (Rep m)
map_lam [Scan (Rep m)]
scans [Reduce (Rep m)]
reds Lambda (Rep m)
post_lam) = do
  -- See Note [Translation of Screma].
  --
  -- Start by combining all the reduction and scan parts into a single
  -- operator
  let Reduce Commutativity
_ Lambda (Rep m)
red_lam [SubExp]
red_nes = [Reduce (Rep m)] -> Reduce (Rep m)
forall rep. Buildable rep => [Reduce rep] -> Reduce rep
singleReduce [Reduce (Rep m)]
reds
      Scan Lambda (Rep m)
scan_lam [SubExp]
scan_nes = [Scan (Rep m)] -> Scan (Rep m)
forall rep. Buildable rep => [Scan rep] -> Scan rep
singleScan [Scan (Rep m)]
scans
      ([Type]
red_ts, [Type]
post_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]
red_nes) ([Type] -> ([Type], [Type])) -> [Type] -> ([Type], [Type])
forall a b. (a -> b) -> a -> b
$ SubExp -> ScremaForm (Rep m) -> [Type]
forall rep. SubExp -> ScremaForm rep -> [Type]
scremaType SubExp
w ScremaForm (Rep m)
form

  post_arrs <- [VName] -> [Type] -> m [VName]
forall (m :: * -> *).
Transformer m =>
[VName] -> [Type] -> m [VName]
resultArray [VName]
arrs [Type]
post_ts

  scanacc_params <- mapM (newParam "scanacc" . flip toDecl Observe) $ lambdaReturnType scan_lam
  redout_params <- mapM (newParam "redout" . flip toDecl Observe) $ lambdaReturnType red_lam
  out_params <- mapM (newParam "out" . flip toDecl Consume) post_ts

  arr_ts <- mapM lookupType arrs
  let paramForAcc (Acc VName
c ShapeBase SubExp
_ [Type]
_) = (Param DeclType -> Bool)
-> [Param DeclType] -> Maybe (Param DeclType)
forall (t :: * -> *) a. Foldable t => (a -> Bool) -> t a -> Maybe a
find (Type -> Bool
f (Type -> Bool)
-> (Param DeclType -> Type) -> Param DeclType -> Bool
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Param DeclType -> Type
forall dec. Typed dec => Param dec -> Type
paramType) [Param DeclType]
out_params
        where
          f :: Type -> Bool
f (Acc VName
c2 ShapeBase SubExp
_ [Type]
_) = VName
c VName -> VName -> Bool
forall a. Eq a => a -> a -> Bool
== VName
c2
          f Type
_ = Bool
False
      paramForAcc Type
_ = Maybe (Param DeclType)
forall a. Maybe a
Nothing

  let merge =
        [[(Param DeclType, SubExp)]] -> [(Param DeclType, SubExp)]
forall (t :: * -> *) a. Foldable t => t [a] -> [a]
concat
          [ [Param DeclType] -> [SubExp] -> [(Param DeclType, SubExp)]
forall a b. [a] -> [b] -> [(a, b)]
zip [Param DeclType]
scanacc_params [SubExp]
scan_nes,
            [Param DeclType] -> [SubExp] -> [(Param DeclType, SubExp)]
forall a b. [a] -> [b] -> [(a, b)]
zip [Param DeclType]
redout_params [SubExp]
red_nes,
            [Param DeclType] -> [SubExp] -> [(Param DeclType, SubExp)]
forall a b. [a] -> [b] -> [(a, b)]
zip [Param DeclType]
out_params ([SubExp] -> [(Param DeclType, SubExp)])
-> [SubExp] -> [(Param DeclType, SubExp)]
forall a b. (a -> b) -> a -> b
$ (VName -> SubExp) -> [VName] -> [SubExp]
forall a b. (a -> b) -> [a] -> [b]
map VName -> SubExp
Var [VName]
post_arrs
          ]
  i <- newVName "i"
  let loopform = VName -> IntType -> SubExp -> LoopForm
ForLoop VName
i IntType
Int64 SubExp
w
      lam_cons = Lambda (Aliases (Rep m)) -> Names
forall rep. Aliased rep => Lambda rep -> Names
consumedByLambda (Lambda (Aliases (Rep m)) -> Names)
-> Lambda (Aliases (Rep m)) -> Names
forall a b. (a -> b) -> a -> b
$ AliasTable -> Lambda (Rep m) -> Lambda (Aliases (Rep m))
forall rep t.
AliasableRep rep =>
AliasTable -> GLambda rep t -> GLambda (Aliases rep) t
Alias.analyseLambda AliasTable
forall a. Monoid a => a
mempty Lambda (Rep m)
map_lam

  loop_body <- runBodyBuilder
    . localScope (scopeOfFParams (map fst merge) <> scopeOfLoopForm loopform)
    $ do
      -- Bind the parameters to the lambda.
      forM_ (zip3 (lambdaParams map_lam) arrs arr_ts) $ \(Param (LParamInfo (Rep m))
p, VName
arr, Type
arr_t) ->
        case Type -> Maybe (Param DeclType)
paramForAcc Type
arr_t of
          Just Param DeclType
acc_out_p ->
            [VName]
-> Exp (Rep (BuilderT (Rep m) (State VNameSource)))
-> BuilderT (Rep m) (State VNameSource) ()
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) -> BuilderT (Rep m) (State VNameSource) ())
-> (BasicOp -> Exp (Rep m))
-> BasicOp
-> BuilderT (Rep m) (State VNameSource) ()
forall b c a. (b -> c) -> (a -> b) -> a -> c
. BasicOp -> Exp (Rep m)
forall rep. BasicOp -> Exp rep
BasicOp (BasicOp -> BuilderT (Rep m) (State VNameSource) ())
-> BasicOp -> BuilderT (Rep m) (State VNameSource) ()
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 -> SubExp) -> VName -> SubExp
forall a b. (a -> b) -> a -> b
$
                  Param DeclType -> VName
forall dec. Param dec -> VName
paramName Param DeclType
acc_out_p
          Maybe (Param DeclType)
Nothing
            | Param (LParamInfo (Rep m)) -> VName
forall dec. Param dec -> VName
paramName Param (LParamInfo (Rep m))
p VName -> Names -> Bool
`nameIn` Names
lam_cons -> do
                p' <-
                  Name
-> Exp (Rep (BuilderT (Rep m) (State VNameSource)))
-> BuilderT (Rep m) (State VNameSource) VName
forall (m :: * -> *).
MonadBuilder m =>
Name -> Exp (Rep m) -> m VName
letExp (VName -> Name
baseName (Param (LParamInfo (Rep m)) -> VName
forall dec. Param dec -> VName
paramName Param (LParamInfo (Rep m))
p)) (Exp (Rep m) -> BuilderT (Rep m) (State VNameSource) VName)
-> (BasicOp -> Exp (Rep m))
-> BasicOp
-> BuilderT (Rep m) (State VNameSource) VName
forall b c a. (b -> c) -> (a -> b) -> a -> c
. BasicOp -> Exp (Rep m)
forall rep. BasicOp -> Exp rep
BasicOp (BasicOp -> BuilderT (Rep m) (State VNameSource) VName)
-> BasicOp -> BuilderT (Rep m) (State VNameSource) VName
forall a b. (a -> b) -> a -> b
$
                    VName -> Slice SubExp -> BasicOp
Index VName
arr (Slice SubExp -> BasicOp) -> Slice SubExp -> BasicOp
forall a b. (a -> b) -> a -> b
$
                      Type -> [DimIndex SubExp] -> Slice SubExp
fullSlice Type
arr_t [SubExp -> DimIndex SubExp
forall d. d -> DimIndex d
DimFix (SubExp -> DimIndex SubExp) -> SubExp -> DimIndex SubExp
forall a b. (a -> b) -> a -> b
$ VName -> SubExp
Var VName
i]
                letBindNames [paramName p] $ BasicOp $ Replicate mempty $ Var p'
            | Bool
otherwise ->
                [VName]
-> Exp (Rep (BuilderT (Rep m) (State VNameSource)))
-> BuilderT (Rep m) (State VNameSource) ()
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) -> BuilderT (Rep m) (State VNameSource) ())
-> (Slice SubExp -> Exp (Rep m))
-> Slice SubExp
-> BuilderT (Rep m) (State VNameSource) ()
forall b c a. (b -> c) -> (a -> b) -> a -> c
. BasicOp -> Exp (Rep m)
forall rep. BasicOp -> Exp rep
BasicOp (BasicOp -> Exp (Rep m))
-> (Slice SubExp -> BasicOp) -> Slice SubExp -> Exp (Rep m)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. VName -> Slice SubExp -> BasicOp
Index VName
arr (Slice SubExp -> BuilderT (Rep m) (State VNameSource) ())
-> Slice SubExp -> BuilderT (Rep m) (State VNameSource) ()
forall a b. (a -> b) -> a -> b
$
                  Type -> [DimIndex SubExp] -> Slice SubExp
fullSlice Type
arr_t [SubExp -> DimIndex SubExp
forall d. d -> DimIndex d
DimFix (SubExp -> DimIndex SubExp) -> SubExp -> DimIndex SubExp
forall a b. (a -> b) -> a -> b
$ VName -> SubExp
Var VName
i]

      -- Insert the statements of the lambda.  We have taken care to
      -- ensure that the parameters are bound at this point.
      mapM_ addStm $ bodyStms $ lambdaBody map_lam
      -- Split into scan results, reduce results, and map results.
      let (scan_res, red_res, map_res) =
            splitAt3 (length scan_nes) (length red_nes) $
              bodyResult $
                lambdaBody map_lam

      scan_res' <-
        eLambda scan_lam $
          map (pure . BasicOp . SubExp) $
            map (Var . paramName) scanacc_params ++ map resSubExp scan_res
      red_res' <-
        eLambda red_lam $
          map (pure . BasicOp . SubExp) $
            map (Var . paramName) redout_params ++ map resSubExp red_res

      let res = [SubExpRes]
scan_res' [SubExpRes] -> [SubExpRes] -> [SubExpRes]
forall a. Semigroup a => a -> a -> a
<> [SubExpRes]
map_res
          param_bind = SubExpRes -> SubExp
resSubExp (SubExpRes -> SubExp) -> [SubExpRes] -> [SubExp]
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> [SubExpRes]
res
          certs = SubExpRes -> Certs
resCerts (SubExpRes -> Certs) -> [SubExpRes] -> [Certs]
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> [SubExpRes]
res
      forM_ (zip3 (paramName <$> lambdaParams post_lam) param_bind certs) $
        \(VName
par, SubExp
v, Certs
cs) -> do
          Certs
-> BuilderT (Rep m) (State VNameSource) ()
-> BuilderT (Rep m) (State VNameSource) ()
forall a.
Certs
-> BuilderT (Rep m) (State VNameSource) a
-> BuilderT (Rep m) (State VNameSource) a
forall (m :: * -> *) a. MonadBuilder m => Certs -> m a -> m a
certifying Certs
cs (BuilderT (Rep m) (State VNameSource) ()
 -> BuilderT (Rep m) (State VNameSource) ())
-> BuilderT (Rep m) (State VNameSource) ()
-> BuilderT (Rep m) (State VNameSource) ()
forall a b. (a -> b) -> a -> b
$ [VName]
-> Exp (Rep (BuilderT (Rep m) (State VNameSource)))
-> BuilderT (Rep m) (State VNameSource) ()
forall (m :: * -> *).
MonadBuilder m =>
[VName] -> Exp (Rep m) -> m ()
letBindNames [VName
par] (Exp (Rep (BuilderT (Rep m) (State VNameSource)))
 -> BuilderT (Rep m) (State VNameSource) ())
-> Exp (Rep (BuilderT (Rep m) (State VNameSource)))
-> BuilderT (Rep m) (State VNameSource) ()
forall a b. (a -> b) -> a -> b
$ BasicOp -> Exp (Rep (BuilderT (Rep m) (State VNameSource)))
forall rep. BasicOp -> Exp rep
BasicOp (BasicOp -> Exp (Rep (BuilderT (Rep m) (State VNameSource))))
-> BasicOp -> Exp (Rep (BuilderT (Rep m) (State VNameSource)))
forall a b. (a -> b) -> a -> b
$ SubExp -> BasicOp
SubExp SubExp
v

      mapM_ addStm $ bodyStms $ lambdaBody post_lam

      let 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
      outarrs <-
        certifying (foldMap resCerts post_res) $
          letwith (map paramName out_params) (Var i) $
            map resSubExp post_res

      pure . concat $
        [ scan_res',
          red_res',
          varsRes outarrs
        ]

  discards <- replicateM (length scanacc_params) (newVName "discard")

  -- Screma requires alias-free results, so reduction results are bound to fresh
  -- names, so that we can copy the array-typed ones.
  let (red_pat_names, post_pat_names) = splitAt (length red_nes) $ patNames pat
  red_names <- mapM newName red_pat_names

  letBindNames (discards ++ red_names ++ post_pat_names) $
    Loop merge loopform loop_body

  forM_ (zip3 red_pat_names red_names red_ts) $ \(VName
to, VName
from, Type
t) ->
    [VName] -> Exp (Rep m) -> m ()
forall (m :: * -> *).
MonadBuilder m =>
[VName] -> Exp (Rep m) -> m ()
letBindNames [VName
to] (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
$
      case Type
t of
        Array {} -> ShapeBase SubExp -> SubExp -> BasicOp
Replicate ShapeBase SubExp
forall a. Monoid a => a
mempty (SubExp -> BasicOp) -> SubExp -> BasicOp
forall a b. (a -> b) -> a -> b
$ VName -> SubExp
Var VName
from
        Type
_ -> SubExp -> BasicOp
SubExp (SubExp -> BasicOp) -> SubExp -> BasicOp
forall a b. (a -> b) -> a -> b
$ VName -> SubExp
Var VName
from

-- | Transform a single 'SOAC' into a do-loop.  The body of the lambda
-- is untouched, and may or may not contain further 'SOAC's depending
-- on the given rep.
transformSOAC ::
  (Transformer m) =>
  Pat (LetDec (Rep m)) ->
  SOAC (Rep m) ->
  m ()
transformSOAC :: forall (m :: * -> *).
Transformer m =>
Pat (LetDec (Rep m)) -> SOAC (Rep m) -> m ()
transformSOAC Pat (LetDec (Rep m))
_ JVP {} =
  [Char] -> m ()
forall a. HasCallStack => [Char] -> a
error [Char]
"transformSOAC: unhandled JVP"
transformSOAC Pat (LetDec (Rep m))
_ VJP {} =
  [Char] -> m ()
forall a. HasCallStack => [Char] -> a
error [Char]
"transformSOAC: unhandled VJP"
transformSOAC Pat (LetDec (Rep m))
_ WithVJP {} =
  [Char] -> m ()
forall a. HasCallStack => [Char] -> a
error [Char]
"transformSOAC: unhandled WithVJP"
transformSOAC Pat (LetDec (Rep m))
pat (FlatMap SubExp
w [VName]
arrs ExtLambda (Rep m)
lam) =
  Pat (LetDec (Rep m))
-> SubExp -> [VName] -> ExtLambda (Rep m) -> m ()
forall (m :: * -> *).
Transformer m =>
Pat (LetDec (Rep m))
-> SubExp -> [VName] -> ExtLambda (Rep m) -> m ()
transformFlatMap Pat (LetDec (Rep m))
pat SubExp
w [VName]
arrs ExtLambda (Rep m)
lam
transformSOAC Pat (LetDec (Rep m))
pat (Screma SubExp
w [VName]
arrs ScremaForm (Rep m)
form) =
  Pat (LetDec (Rep m))
-> SubExp -> [VName] -> ScremaForm (Rep m) -> m ()
forall (m :: * -> *) dec.
Transformer m =>
Pat dec -> SubExp -> [VName] -> ScremaForm (Rep m) -> m ()
transformScrema Pat (LetDec (Rep m))
pat SubExp
w [VName]
arrs ScremaForm (Rep m)
form
transformSOAC Pat (LetDec (Rep m))
pat (Stream SubExp
w [VName]
arrs [SubExp]
nes Lambda (Rep m)
lam) = do
  -- Create a loop that repeatedly applies the lambda body to a
  -- chunksize of 1.  Hopefully this will lead to this outer loop
  -- being the only one, as all the innermost one can be simplified
  -- away (as they will have one iteration each).
  let (Param Type
chunk_size_param, [Param Type]
fold_params, [Param Type]
chunk_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
      mapout_ts :: [Type]
mapout_ts = (Type -> Type) -> [Type] -> [Type]
forall a b. (a -> b) -> [a] -> [b]
map (Type -> SubExp -> Type
forall d m.
TypeBase (ShapeBase d) m -> d -> TypeBase (ShapeBase d) m
`setOuterSize` SubExp
w) ([Type] -> [Type]) -> [Type] -> [Type]
forall a b. (a -> b) -> a -> b
$ Int -> [Type] -> [Type]
forall a. Int -> [a] -> [a]
drop ([SubExp] -> Int
forall a. [a] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length [SubExp]
nes) ([Type] -> [Type]) -> [Type] -> [Type]
forall a b. (a -> b) -> a -> b
$ Lambda (Rep m) -> [Type]
forall rep t. GLambda rep t -> [t]
lambdaReturnType Lambda (Rep m)
lam

  mapout_initial <- [VName] -> [Type] -> m [VName]
forall (m :: * -> *).
Transformer m =>
[VName] -> [Type] -> m [VName]
resultArray [VName]
arrs [Type]
mapout_ts
  mapout_params <- forM mapout_ts $ \Type
t ->
    Name -> DeclType -> m (Param DeclType)
forall (m :: * -> *) dec.
MonadFreshNames m =>
Name -> dec -> m (Param dec)
newParam Name
"stream_mapout" (DeclType -> m (Param DeclType)) -> DeclType -> m (Param DeclType)
forall a b. (a -> b) -> a -> b
$ Type -> Diet -> DeclType
forall shape. TypeBase shape NoMode -> Diet -> TypeBase shape Diet
toDecl Type
t Diet
Consume
  let mapout_merge = [Param DeclType] -> [SubExp] -> [(Param DeclType, SubExp)]
forall a b. [a] -> [b] -> [(a, b)]
zip [Param DeclType]
mapout_params ([SubExp] -> [(Param DeclType, SubExp)])
-> [SubExp] -> [(Param DeclType, SubExp)]
forall a b. (a -> b) -> a -> b
$ (VName -> SubExp) -> [VName] -> [SubExp]
forall a b. (a -> b) -> [a] -> [b]
map VName -> SubExp
Var [VName]
mapout_initial

  let paramForAcc (Acc VName
c ShapeBase SubExp
_ [Type]
_) = (Param DeclType -> Bool)
-> [Param DeclType] -> Maybe (Param DeclType)
forall (t :: * -> *) a. Foldable t => (a -> Bool) -> t a -> Maybe a
find (Type -> Bool
f (Type -> Bool)
-> (Param DeclType -> Type) -> Param DeclType -> Bool
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Param DeclType -> Type
forall dec. Typed dec => Param dec -> Type
paramType) [Param DeclType]
mapout_params
        where
          f :: Type -> Bool
f (Acc VName
c2 ShapeBase SubExp
_ [Type]
_) = VName
c VName -> VName -> Bool
forall a. Eq a => a -> a -> Bool
== VName
c2
          f Type
_ = Bool
False
      paramForAcc Type
_ = Maybe (Param DeclType)
forall a. Maybe a
Nothing

  -- We need to copy the neutral elements because they may be consumed
  -- in the body of the Stream.
  let copyIfArray SubExp
se = do
        se_t <- SubExp -> m Type
forall t (m :: * -> *). HasScope t m => SubExp -> m Type
subExpType SubExp
se
        case (se_t, se) of
          (Array {}, Var VName
v) ->
            Name -> Exp (Rep m) -> m SubExp
forall (m :: * -> *).
MonadBuilder m =>
Name -> Exp (Rep m) -> m SubExp
letSubExp (VName -> Name
baseName VName
v) (Exp (Rep m) -> m SubExp) -> Exp (Rep m) -> m SubExp
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
$ ShapeBase SubExp -> SubExp -> BasicOp
Replicate ShapeBase SubExp
forall a. Monoid a => a
mempty SubExp
se
          (Type, SubExp)
_ -> SubExp -> m SubExp
forall a. a -> m a
forall (f :: * -> *) a. Applicative f => a -> f a
pure SubExp
se
  nes' <- mapM copyIfArray nes

  let onType TypeBase shape NoMode
t = TypeBase shape NoMode
t TypeBase shape NoMode -> Diet -> TypeBase shape Diet
forall shape. TypeBase shape NoMode -> Diet -> TypeBase shape Diet
`toDecl` Diet
Consume
      merge = [Param DeclType] -> [SubExp] -> [(Param DeclType, SubExp)]
forall a b. [a] -> [b] -> [(a, b)]
zip ((Param Type -> Param DeclType) -> [Param Type] -> [Param DeclType]
forall a b. (a -> b) -> [a] -> [b]
map ((Type -> DeclType) -> Param Type -> Param DeclType
forall a b. (a -> b) -> Param a -> Param b
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap Type -> DeclType
forall {shape}. TypeBase shape NoMode -> TypeBase shape Diet
onType) [Param Type]
fold_params) [SubExp]
nes' [(Param DeclType, SubExp)]
-> [(Param DeclType, SubExp)] -> [(Param DeclType, SubExp)]
forall a. [a] -> [a] -> [a]
++ [(Param DeclType, SubExp)]
mapout_merge
      merge_params = ((Param DeclType, SubExp) -> Param DeclType)
-> [(Param DeclType, SubExp)] -> [Param DeclType]
forall a b. (a -> b) -> [a] -> [b]
map (Param DeclType, SubExp) -> Param DeclType
forall a b. (a, b) -> a
fst [(Param DeclType, SubExp)]
merge

  i <- newVName "i"

  let loop_form = VName -> IntType -> SubExp -> LoopForm
ForLoop VName
i IntType
Int64 SubExp
w

  letBindNames [paramName chunk_size_param] . BasicOp . SubExp $
    intConst Int64 1

  arrs_ts <- mapM lookupType arrs
  loop_body <- runBodyBuilder $
    localScope (scopeOfLoopForm loop_form <> scopeOfFParams merge_params) $ do
      let slice = [SubExp -> SubExp -> SubExp -> DimIndex SubExp
forall d. d -> d -> d -> DimIndex d
DimSlice (VName -> SubExp
Var VName
i) (VName -> SubExp
Var (Param Type -> VName
forall dec. Param dec -> VName
paramName Param Type
chunk_size_param)) (IntType -> Integer -> SubExp
intConst IntType
Int64 Integer
1)]
      forM_ (zip3 chunk_params arrs arrs_ts) $ \(Param Type
p, VName
arr, Type
arr_t) ->
        case Type -> Maybe (Param DeclType)
paramForAcc Type
arr_t of
          Just Param DeclType
acc_out_p ->
            [VName]
-> Exp (Rep (BuilderT (Rep m) (State VNameSource)))
-> BuilderT (Rep m) (State VNameSource) ()
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) -> BuilderT (Rep m) (State VNameSource) ())
-> (SubExp -> Exp (Rep m))
-> SubExp
-> BuilderT (Rep m) (State VNameSource) ()
forall b c a. (b -> c) -> (a -> b) -> a -> c
. BasicOp -> Exp (Rep m)
forall rep. BasicOp -> Exp rep
BasicOp (BasicOp -> Exp (Rep m))
-> (SubExp -> BasicOp) -> SubExp -> Exp (Rep m)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. SubExp -> BasicOp
SubExp (SubExp -> BuilderT (Rep m) (State VNameSource) ())
-> SubExp -> BuilderT (Rep m) (State VNameSource) ()
forall a b. (a -> b) -> a -> b
$
              VName -> SubExp
Var (Param DeclType -> VName
forall dec. Param dec -> VName
paramName Param DeclType
acc_out_p)
          Maybe (Param DeclType)
Nothing ->
            [VName]
-> Exp (Rep (BuilderT (Rep m) (State VNameSource)))
-> BuilderT (Rep m) (State VNameSource) ()
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) -> BuilderT (Rep m) (State VNameSource) ())
-> (BasicOp -> Exp (Rep m))
-> BasicOp
-> BuilderT (Rep m) (State VNameSource) ()
forall b c a. (b -> c) -> (a -> b) -> a -> c
. BasicOp -> Exp (Rep m)
forall rep. BasicOp -> Exp rep
BasicOp (BasicOp -> BuilderT (Rep m) (State VNameSource) ())
-> BasicOp -> BuilderT (Rep m) (State VNameSource) ()
forall a b. (a -> b) -> a -> b
$
              VName -> Slice SubExp -> BasicOp
Index VName
arr (Type -> [DimIndex SubExp] -> Slice SubExp
fullSlice (Param Type -> Type
forall dec. Typed dec => Param dec -> Type
paramType Param Type
p) [DimIndex SubExp]
slice)

      (res, mapout_res) <- splitAt (length nes) <$> bodyBind (lambdaBody lam)

      res' <- mapM (copyIfArray . resSubExp) res

      mapout_res' <- forM (zip mapout_params mapout_res) $ \(Param DeclType
p, SubExpRes Certs
cs SubExp
se) ->
        Certs
-> BuilderT (Rep m) (State VNameSource) SubExp
-> BuilderT (Rep m) (State VNameSource) SubExp
forall a.
Certs
-> BuilderT (Rep m) (State VNameSource) a
-> BuilderT (Rep m) (State VNameSource) a
forall (m :: * -> *) a. MonadBuilder m => Certs -> m a -> m a
certifying Certs
cs (BuilderT (Rep m) (State VNameSource) SubExp
 -> BuilderT (Rep m) (State VNameSource) SubExp)
-> (BasicOp -> BuilderT (Rep m) (State VNameSource) SubExp)
-> BasicOp
-> BuilderT (Rep m) (State VNameSource) SubExp
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Name
-> Exp (Rep (BuilderT (Rep m) (State VNameSource)))
-> BuilderT (Rep m) (State VNameSource) SubExp
forall (m :: * -> *).
MonadBuilder m =>
Name -> Exp (Rep m) -> m SubExp
letSubExp Name
"mapout_res" (Exp (Rep m) -> BuilderT (Rep m) (State VNameSource) SubExp)
-> (BasicOp -> Exp (Rep m))
-> BasicOp
-> BuilderT (Rep m) (State VNameSource) SubExp
forall b c a. (b -> c) -> (a -> b) -> a -> c
. BasicOp -> Exp (Rep m)
forall rep. BasicOp -> Exp rep
BasicOp (BasicOp -> BuilderT (Rep m) (State VNameSource) SubExp)
-> BasicOp -> BuilderT (Rep m) (State VNameSource) SubExp
forall a b. (a -> b) -> a -> b
$
          if Type -> Bool
forall shape o. TypeBase shape o -> Bool
isAcc (Param DeclType -> Type
forall dec. Typed dec => Param dec -> Type
paramType Param DeclType
p)
            then SubExp -> BasicOp
SubExp SubExp
se
            else Safety -> VName -> Slice SubExp -> SubExp -> BasicOp
Update Safety
Unsafe (Param DeclType -> VName
forall dec. Param dec -> VName
paramName Param DeclType
p) (Type -> [DimIndex SubExp] -> Slice SubExp
fullSlice (Param DeclType -> Type
forall dec. Typed dec => Param dec -> Type
paramType Param DeclType
p) [DimIndex SubExp]
slice) SubExp
se

      pure $ subExpsRes $ res' ++ mapout_res'

  letBind pat $ Loop merge loop_form loop_body
transformSOAC Pat (LetDec (Rep m))
pat (Hist SubExp
len [VName]
imgs [HistOp (Rep m)]
ops Lambda (Rep m)
bucket_fun) = do
  iter <- Name -> m VName
forall (m :: * -> *). MonadFreshNames m => Name -> m VName
newVName Name
"iter"

  -- Bind arguments to parameters for the merge-variables.
  hists_ts <- mapM lookupType $ concatMap histDest ops
  hists_out <- mapM (newIdent "dests") hists_ts
  let merge = [Ident] -> [SubExp] -> [(Param DeclType, SubExp)]
loopMerge [Ident]
hists_out ([SubExp] -> [(Param DeclType, SubExp)])
-> [SubExp] -> [(Param DeclType, SubExp)]
forall a b. (a -> b) -> a -> b
$ (HistOp (Rep m) -> [SubExp]) -> [HistOp (Rep m)] -> [SubExp]
forall (t :: * -> *) a b. Foldable t => (a -> [b]) -> t a -> [b]
concatMap ((VName -> SubExp) -> [VName] -> [SubExp]
forall a b. (a -> b) -> [a] -> [b]
map VName -> SubExp
Var ([VName] -> [SubExp])
-> (HistOp (Rep m) -> [VName]) -> HistOp (Rep m) -> [SubExp]
forall b c a. (b -> c) -> (a -> b) -> a -> c
. HistOp (Rep m) -> [VName]
forall rep. HistOp rep -> [VName]
histDest) [HistOp (Rep m)]
ops

  -- Bind lambda-bodies for operators.
  let iter_scope = VName -> NameInfo (Rep m) -> Scope (Rep m) -> Scope (Rep m)
forall k a. Ord k => k -> a -> Map k a -> Map k a
M.insert VName
iter (IntType -> NameInfo (Rep m)
forall rep. IntType -> NameInfo rep
IndexName IntType
Int64) (Scope (Rep m) -> Scope (Rep m)) -> Scope (Rep m) -> Scope (Rep m)
forall a b. (a -> b) -> a -> b
$ [Param DeclType] -> Scope (Rep m)
forall rep dec. (FParamInfo rep ~ dec) => [Param dec] -> Scope rep
scopeOfFParams ([Param DeclType] -> Scope (Rep m))
-> [Param DeclType] -> Scope (Rep m)
forall a b. (a -> b) -> a -> b
$ ((Param DeclType, SubExp) -> Param DeclType)
-> [(Param DeclType, SubExp)] -> [Param DeclType]
forall a b. (a -> b) -> [a] -> [b]
map (Param DeclType, SubExp) -> Param DeclType
forall a b. (a, b) -> a
fst [(Param DeclType, SubExp)]
merge
  loopBody <- runBodyBuilder . localScope iter_scope $ do
    -- Bind images to parameters of bucket function.
    imgs' <- forM imgs $ \VName
img -> do
      img_t <- VName -> BuilderT (Rep m) (State VNameSource) Type
forall rep (m :: * -> *). HasScope rep m => VName -> m Type
lookupType VName
img
      letSubExp "pixel" $ BasicOp $ Index img $ fullSlice img_t [DimFix $ Var iter]
    imgs'' <- map resSubExp <$> bindLambda bucket_fun (map (BasicOp . SubExp) imgs')

    -- Split out values from bucket function.
    let lens = [Int] -> Int
forall a. Num a => [a] -> a
forall (t :: * -> *) a. (Foldable t, Num a) => t a -> a
sum ([Int] -> Int) -> [Int] -> Int
forall a b. (a -> b) -> a -> b
$ (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
        ops_inds = [Int] -> [SubExp] -> [[SubExp]]
forall a. [Int] -> [a] -> [[a]]
chunks ((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) (Int -> [SubExp] -> [SubExp]
forall a. Int -> [a] -> [a]
take Int
lens [SubExp]
imgs'')
        vals = [Int] -> [SubExp] -> [[SubExp]]
forall a. [Int] -> [a] -> [[a]]
chunks ((HistOp (Rep m) -> Int) -> [HistOp (Rep m)] -> [Int]
forall a b. (a -> b) -> [a] -> [b]
map ([Type] -> Int
forall a. [a] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length ([Type] -> Int)
-> (HistOp (Rep m) -> [Type]) -> HistOp (Rep m) -> Int
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Lambda (Rep m) -> [Type]
forall rep t. GLambda rep t -> [t]
lambdaReturnType (Lambda (Rep m) -> [Type])
-> (HistOp (Rep m) -> Lambda (Rep m)) -> HistOp (Rep m) -> [Type]
forall b c a. (b -> c) -> (a -> b) -> a -> c
. HistOp (Rep m) -> Lambda (Rep m)
forall rep. HistOp rep -> Lambda rep
histOp) [HistOp (Rep m)]
ops) ([SubExp] -> [[SubExp]]) -> [SubExp] -> [[SubExp]]
forall a b. (a -> b) -> a -> b
$ Int -> [SubExp] -> [SubExp]
forall a. Int -> [a] -> [a]
drop Int
lens [SubExp]
imgs''
        hists_out' =
          [Int] -> [VName] -> [[VName]]
forall a. [Int] -> [a] -> [[a]]
chunks ((HistOp (Rep m) -> Int) -> [HistOp (Rep m)] -> [Int]
forall a b. (a -> b) -> [a] -> [b]
map ([Type] -> Int
forall a. [a] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length ([Type] -> Int)
-> (HistOp (Rep m) -> [Type]) -> HistOp (Rep m) -> Int
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Lambda (Rep m) -> [Type]
forall rep t. GLambda rep t -> [t]
lambdaReturnType (Lambda (Rep m) -> [Type])
-> (HistOp (Rep m) -> Lambda (Rep m)) -> HistOp (Rep m) -> [Type]
forall b c a. (b -> c) -> (a -> b) -> a -> c
. HistOp (Rep m) -> Lambda (Rep m)
forall rep. HistOp rep -> Lambda rep
histOp) [HistOp (Rep m)]
ops) ([VName] -> [[VName]]) -> [VName] -> [[VName]]
forall a b. (a -> b) -> a -> b
$
            (Ident -> VName) -> [Ident] -> [VName]
forall a b. (a -> b) -> [a] -> [b]
map Ident -> VName
identName [Ident]
hists_out

    hists_out'' <- forM (zip4 hists_out' ops ops_inds vals) $ \([VName]
hist, HistOp (Rep m)
op, [SubExp]
idxs, [SubExp]
val) -> do
      -- Check whether the indexes are in-bound.  If they are not, we
      -- return the histograms unchanged.
      let outside_bounds_branch :: BuilderT
  (Rep m)
  (State VNameSource)
  (Body (Rep (BuilderT (Rep m) (State VNameSource))))
outside_bounds_branch = Builder (Rep m) [SubExpRes]
-> BuilderT
     (Rep m)
     (State VNameSource)
     (Body (Rep (BuilderT (Rep m) (State VNameSource))))
forall (m :: * -> *) res.
(MonadBuilder m, IsResult res) =>
m [res] -> m (GBody (Rep m) res)
buildBody_ (Builder (Rep m) [SubExpRes]
 -> BuilderT
      (Rep m)
      (State VNameSource)
      (Body (Rep (BuilderT (Rep m) (State VNameSource)))))
-> Builder (Rep m) [SubExpRes]
-> BuilderT
     (Rep m)
     (State VNameSource)
     (Body (Rep (BuilderT (Rep m) (State VNameSource))))
forall a b. (a -> b) -> a -> b
$ [SubExpRes] -> Builder (Rep m) [SubExpRes]
forall a. a -> BuilderT (Rep m) (State VNameSource) a
forall (f :: * -> *) a. Applicative f => a -> f a
pure ([SubExpRes] -> Builder (Rep m) [SubExpRes])
-> [SubExpRes] -> Builder (Rep m) [SubExpRes]
forall a b. (a -> b) -> a -> b
$ [VName] -> [SubExpRes]
varsRes [VName]
hist
          oob :: BuilderT
  (Rep m)
  (State VNameSource)
  (Exp (Rep (BuilderT (Rep m) (State VNameSource))))
oob = case [VName]
hist of
            [] -> SubExp
-> BuilderT
     (Rep m)
     (State VNameSource)
     (Exp (Rep (BuilderT (Rep m) (State VNameSource))))
forall (m :: * -> *). MonadBuilder m => SubExp -> m (Exp (Rep m))
eSubExp (SubExp
 -> BuilderT
      (Rep m)
      (State VNameSource)
      (Exp (Rep (BuilderT (Rep m) (State VNameSource)))))
-> SubExp
-> BuilderT
     (Rep m)
     (State VNameSource)
     (Exp (Rep (BuilderT (Rep m) (State VNameSource))))
forall a b. (a -> b) -> a -> b
$ Bool -> SubExp
forall v. IsValue v => v -> SubExp
constant Bool
True
            VName
arr : [VName]
_ -> VName
-> [BuilderT
      (Rep m)
      (State VNameSource)
      (Exp (Rep (BuilderT (Rep m) (State VNameSource))))]
-> BuilderT
     (Rep m)
     (State VNameSource)
     (Exp (Rep (BuilderT (Rep m) (State VNameSource))))
forall (m :: * -> *).
MonadBuilder m =>
VName -> [m (Exp (Rep m))] -> m (Exp (Rep m))
eOutOfBounds VName
arr ([BuilderT
    (Rep m)
    (State VNameSource)
    (Exp (Rep (BuilderT (Rep m) (State VNameSource))))]
 -> BuilderT
      (Rep m)
      (State VNameSource)
      (Exp (Rep (BuilderT (Rep m) (State VNameSource)))))
-> [BuilderT
      (Rep m)
      (State VNameSource)
      (Exp (Rep (BuilderT (Rep m) (State VNameSource))))]
-> BuilderT
     (Rep m)
     (State VNameSource)
     (Exp (Rep (BuilderT (Rep m) (State VNameSource))))
forall a b. (a -> b) -> a -> b
$ (SubExp
 -> BuilderT
      (Rep m)
      (State VNameSource)
      (Exp (Rep (BuilderT (Rep m) (State VNameSource)))))
-> [SubExp]
-> [BuilderT
      (Rep m)
      (State VNameSource)
      (Exp (Rep (BuilderT (Rep m) (State VNameSource))))]
forall a b. (a -> b) -> [a] -> [b]
map SubExp
-> BuilderT
     (Rep m)
     (State VNameSource)
     (Exp (Rep (BuilderT (Rep m) (State VNameSource))))
forall (m :: * -> *). MonadBuilder m => SubExp -> m (Exp (Rep m))
eSubExp [SubExp]
idxs

      Name
-> Exp (Rep (BuilderT (Rep m) (State VNameSource)))
-> BuilderT (Rep m) (State VNameSource) [VName]
forall (m :: * -> *).
MonadBuilder m =>
Name -> Exp (Rep m) -> m [VName]
letTupExp Name
"new_histo" (Exp (Rep m) -> BuilderT (Rep m) (State VNameSource) [VName])
-> (BuilderT (Rep m) (State VNameSource) (GBody (Rep m) SubExpRes)
    -> BuilderT (Rep m) (State VNameSource) (Exp (Rep m)))
-> BuilderT (Rep m) (State VNameSource) (GBody (Rep m) SubExpRes)
-> BuilderT (Rep m) (State VNameSource) [VName]
forall (m :: * -> *) b c a.
Monad m =>
(b -> m c) -> (a -> m b) -> a -> m c
<=< BuilderT
  (Rep m)
  (State VNameSource)
  (Exp (Rep (BuilderT (Rep m) (State VNameSource))))
-> BuilderT
     (Rep m)
     (State VNameSource)
     (Body (Rep (BuilderT (Rep m) (State VNameSource))))
-> BuilderT
     (Rep m)
     (State VNameSource)
     (Body (Rep (BuilderT (Rep m) (State VNameSource))))
-> BuilderT
     (Rep m)
     (State VNameSource)
     (Exp (Rep (BuilderT (Rep m) (State VNameSource))))
forall (m :: * -> *).
(MonadBuilder m, BranchType (Rep m) ~ ExtType) =>
m (Exp (Rep m))
-> m (Body (Rep m)) -> m (Body (Rep m)) -> m (Exp (Rep m))
eIf BuilderT
  (Rep m)
  (State VNameSource)
  (Exp (Rep (BuilderT (Rep m) (State VNameSource))))
oob BuilderT
  (Rep m)
  (State VNameSource)
  (Body (Rep (BuilderT (Rep m) (State VNameSource))))
outside_bounds_branch (BuilderT (Rep m) (State VNameSource) (GBody (Rep m) SubExpRes)
 -> BuilderT (Rep m) (State VNameSource) [VName])
-> BuilderT (Rep m) (State VNameSource) (GBody (Rep m) SubExpRes)
-> BuilderT (Rep m) (State VNameSource) [VName]
forall a b. (a -> b) -> a -> b
$
        Builder (Rep m) [SubExpRes]
-> BuilderT
     (Rep m)
     (State VNameSource)
     (Body (Rep (BuilderT (Rep m) (State VNameSource))))
forall (m :: * -> *) res.
(MonadBuilder m, IsResult res) =>
m [res] -> m (GBody (Rep m) res)
buildBody_ (Builder (Rep m) [SubExpRes]
 -> BuilderT
      (Rep m)
      (State VNameSource)
      (Body (Rep (BuilderT (Rep m) (State VNameSource)))))
-> Builder (Rep m) [SubExpRes]
-> BuilderT
     (Rep m)
     (State VNameSource)
     (Body (Rep (BuilderT (Rep m) (State VNameSource))))
forall a b. (a -> b) -> a -> b
$ do
          -- Read values from histogram.
          h_val <- [VName]
-> (VName -> BuilderT (Rep m) (State VNameSource) SubExp)
-> BuilderT (Rep m) (State VNameSource) [SubExp]
forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
t a -> (a -> m b) -> m (t b)
forM [VName]
hist ((VName -> BuilderT (Rep m) (State VNameSource) SubExp)
 -> BuilderT (Rep m) (State VNameSource) [SubExp])
-> (VName -> BuilderT (Rep m) (State VNameSource) SubExp)
-> BuilderT (Rep m) (State VNameSource) [SubExp]
forall a b. (a -> b) -> a -> b
$ \VName
arr -> do
            arr_t <- VName -> BuilderT (Rep m) (State VNameSource) Type
forall rep (m :: * -> *). HasScope rep m => VName -> m Type
lookupType VName
arr
            letSubExp "read_hist" $ BasicOp $ Index arr $ fullSlice arr_t $ map DimFix idxs

          -- Apply operator.
          h_val' <- bindLambda (histOp op) $ map (BasicOp . SubExp) $ h_val ++ val

          -- Write values back to histograms.
          hist' <- forM (zip hist h_val') $ \(VName
arr, SubExpRes Certs
cs SubExp
v) -> do
            arr_t <- VName -> BuilderT (Rep m) (State VNameSource) Type
forall rep (m :: * -> *). HasScope rep m => VName -> m Type
lookupType VName
arr
            certifying cs . letInPlace "hist_out" arr (fullSlice arr_t $ map DimFix idxs) $
              BasicOp $
                SubExp v

          pure $ varsRes hist'

    pure $ varsRes $ concat hists_out''

  -- Wrap up the above into a for-loop.
  letBind pat $ Loop merge (ForLoop iter Int64 len) loopBody

-- | Recursively first-order-transform a lambda.
transformLambda ::
  ( MonadFreshNames m,
    Buildable rep,
    BuilderOps rep,
    LocalScope somerep m,
    SameScope somerep rep,
    LetDec rep ~ LetDec SOACS,
    Alias.AliasableRep rep
  ) =>
  GLambda SOACS t ->
  m (AST.GLambda rep t)
transformLambda :: forall (m :: * -> *) rep somerep t.
(MonadFreshNames m, Buildable rep, BuilderOps rep,
 LocalScope somerep m, SameScope somerep rep,
 LetDec rep ~ LetDec SOACS, AliasableRep rep) =>
GLambda SOACS t -> m (GLambda rep t)
transformLambda (Lambda [Param (LParamInfo SOACS)]
params [t]
rettype Body SOACS
body) = do
  body' <-
    ((Body (Rep (BuilderT rep (State VNameSource))), Stms rep)
 -> Body (Rep (BuilderT rep (State VNameSource))))
-> m (Body (Rep (BuilderT rep (State VNameSource))), Stms rep)
-> m (Body (Rep (BuilderT rep (State VNameSource))))
forall a b. (a -> b) -> m a -> m b
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap (Body (Rep (BuilderT rep (State VNameSource))), Stms rep)
-> Body (Rep (BuilderT rep (State VNameSource)))
forall a b. (a, b) -> a
fst (m (Body (Rep (BuilderT rep (State VNameSource))), Stms rep)
 -> m (Body (Rep (BuilderT rep (State VNameSource)))))
-> (Builder rep (Body (Rep (BuilderT rep (State VNameSource))))
    -> m (Body (Rep (BuilderT rep (State VNameSource))), Stms rep))
-> Builder rep (Body (Rep (BuilderT rep (State VNameSource))))
-> m (Body (Rep (BuilderT rep (State VNameSource))))
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Builder rep (Body (Rep (BuilderT rep (State VNameSource))))
-> m (Body (Rep (BuilderT rep (State VNameSource))), Stms rep)
forall (m :: * -> *) somerep rep a.
(MonadFreshNames m, HasScope somerep m, SameScope somerep rep) =>
Builder rep a -> m (a, Stms rep)
runBuilder (Builder rep (Body (Rep (BuilderT rep (State VNameSource))))
 -> m (Body (Rep (BuilderT rep (State VNameSource)))))
-> Builder rep (Body (Rep (BuilderT rep (State VNameSource))))
-> m (Body (Rep (BuilderT rep (State VNameSource))))
forall a b. (a -> b) -> a -> b
$
      Scope rep
-> Builder rep (Body (Rep (BuilderT rep (State VNameSource))))
-> Builder rep (Body (Rep (BuilderT rep (State VNameSource))))
forall a.
Scope rep
-> BuilderT rep (State VNameSource) a
-> BuilderT rep (State VNameSource) a
forall rep (m :: * -> *) a.
LocalScope rep m =>
Scope rep -> m a -> m a
localScope ([Param (LParamInfo rep)] -> Scope rep
forall rep dec. (LParamInfo rep ~ dec) => [Param dec] -> Scope rep
scopeOfLParams [Param (LParamInfo rep)]
[Param (LParamInfo SOACS)]
params) (Builder rep (Body (Rep (BuilderT rep (State VNameSource))))
 -> Builder rep (Body (Rep (BuilderT rep (State VNameSource)))))
-> Builder rep (Body (Rep (BuilderT rep (State VNameSource))))
-> Builder rep (Body (Rep (BuilderT rep (State VNameSource))))
forall a b. (a -> b) -> a -> b
$
        Body SOACS
-> Builder rep (Body (Rep (BuilderT rep (State VNameSource))))
forall (m :: * -> *).
(Transformer m, LetDec (Rep m) ~ LetDec SOACS) =>
Body SOACS -> m (Body (Rep m))
transformBody Body SOACS
body
  pure $ Lambda params rettype body'

letwith :: (Transformer m) => [VName] -> SubExp -> [SubExp] -> m [VName]
letwith :: forall (m :: * -> *).
Transformer m =>
[VName] -> SubExp -> [SubExp] -> m [VName]
letwith [VName]
ks SubExp
i [SubExp]
vs = do
  let update :: VName -> SubExp -> m VName
update VName
k SubExp
v = do
        k_t <- VName -> m Type
forall rep (m :: * -> *). HasScope rep m => VName -> m Type
lookupType VName
k
        case k_t of
          Acc {} ->
            Name -> Exp (Rep m) -> m VName
forall (m :: * -> *).
MonadBuilder m =>
Name -> Exp (Rep m) -> m VName
letExp Name
"lw_acc" (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
v
          Type
_ ->
            Name -> VName -> Slice SubExp -> Exp (Rep m) -> m VName
forall (m :: * -> *).
MonadBuilder m =>
Name -> VName -> Slice SubExp -> Exp (Rep m) -> m VName
letInPlace Name
"lw_dest" VName
k (Type -> [DimIndex SubExp] -> Slice SubExp
fullSlice Type
k_t [SubExp -> DimIndex SubExp
forall d. d -> DimIndex d
DimFix SubExp
i]) (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
v
  (VName -> SubExp -> m VName) -> [VName] -> [SubExp] -> m [VName]
forall (m :: * -> *) a b c.
Applicative m =>
(a -> b -> m c) -> [a] -> [b] -> m [c]
zipWithM VName -> SubExp -> m VName
update [VName]
ks [SubExp]
vs

bindLambda ::
  (Transformer m) =>
  AST.GLambda (Rep m) t ->
  [AST.Exp (Rep m)] ->
  m Result
bindLambda :: forall (m :: * -> *) t.
Transformer m =>
GLambda (Rep m) t -> [Exp (Rep m)] -> m [SubExpRes]
bindLambda (Lambda [LParam (Rep m)]
params [t]
_ Body (Rep m)
body) [Exp (Rep m)]
args = do
  [(Param Type, Exp (Rep m))]
-> ((Param Type, Exp (Rep m)) -> m ()) -> m ()
forall (t :: * -> *) (m :: * -> *) a b.
(Foldable t, Monad m) =>
t a -> (a -> m b) -> m ()
forM_ ([Param Type] -> [Exp (Rep m)] -> [(Param Type, Exp (Rep m))]
forall a b. [a] -> [b] -> [(a, b)]
zip [Param Type]
[LParam (Rep m)]
params [Exp (Rep m)]
args) (((Param Type, Exp (Rep m)) -> m ()) -> m ())
-> ((Param Type, Exp (Rep m)) -> m ()) -> m ()
forall a b. (a -> b) -> a -> b
$ \(Param Type
param, Exp (Rep m)
arg) ->
    if Type -> Bool
forall shape o. TypeBase shape o -> Bool
primType (Type -> Bool) -> Type -> Bool
forall a b. (a -> b) -> a -> b
$ Param Type -> Type
forall dec. Typed dec => Param dec -> Type
paramType Param Type
param
      then [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
param] Exp (Rep m)
arg
      else [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
param] (Exp (Rep m) -> m ()) -> m (Exp (Rep m)) -> m ()
forall (m :: * -> *) a b. Monad m => (a -> m b) -> m a -> m b
=<< m (Exp (Rep m)) -> m (Exp (Rep m))
forall (m :: * -> *).
MonadBuilder m =>
m (Exp (Rep m)) -> m (Exp (Rep m))
eCopy (Exp (Rep m) -> m (Exp (Rep m))
forall a. a -> m a
forall (f :: * -> *) a. Applicative f => a -> f a
pure Exp (Rep m)
arg)
  Body (Rep m) -> m [SubExpRes]
forall (m :: * -> *).
MonadBuilder m =>
Body (Rep m) -> m [SubExpRes]
bodyBind Body (Rep m)
body

loopMerge :: [Ident] -> [SubExp] -> [(Param DeclType, SubExp)]
loopMerge :: [Ident] -> [SubExp] -> [(Param DeclType, SubExp)]
loopMerge [Ident]
vars = [(Ident, Diet)] -> [SubExp] -> [(Param DeclType, SubExp)]
loopMerge' ([(Ident, Diet)] -> [SubExp] -> [(Param DeclType, SubExp)])
-> [(Ident, Diet)] -> [SubExp] -> [(Param DeclType, SubExp)]
forall a b. (a -> b) -> a -> b
$ (Ident -> (Ident, Diet)) -> [Ident] -> [(Ident, Diet)]
forall a b. (a -> b) -> [a] -> [b]
map (,Diet
Consume) [Ident]
vars

loopMerge' :: [(Ident, Diet)] -> [SubExp] -> [(Param DeclType, SubExp)]
loopMerge' :: [(Ident, Diet)] -> [SubExp] -> [(Param DeclType, SubExp)]
loopMerge' [(Ident, Diet)]
vars [SubExp]
vals =
  [ (Attrs -> VName -> DeclType -> Param DeclType
forall dec. Attrs -> VName -> dec -> Param dec
Param Attrs
forall a. Monoid a => a
mempty VName
pname (DeclType -> Param DeclType) -> DeclType -> Param DeclType
forall a b. (a -> b) -> a -> b
$ Type -> Diet -> DeclType
forall shape. TypeBase shape NoMode -> Diet -> TypeBase shape Diet
toDecl Type
ptype Diet
o, SubExp
val)
  | ((Ident VName
pname Type
ptype, Diet
o), SubExp
val) <- [(Ident, Diet)] -> [SubExp] -> [((Ident, Diet), SubExp)]
forall a b. [a] -> [b] -> [(a, b)]
zip [(Ident, Diet)]
vars [SubExp]
vals
  ]

-- Note [Translation of Screma]
-- ~~~~~~~~~~~~~~~~~~~~~~~~~~~~
--
-- Screma is the most general SOAC.  It is translated by constructing
-- a loop that contains several groups of parameters, in this order:
--
-- (0) Scan accumulator, initialised with neutral element.
-- (1) Scan results, initialised with Scratch.
-- (2) Reduce results (also functioning as accumulators),
--     initialised with neutral element.
-- (3) Map results, mostly initialised with Scratch.
--
-- However, category (3) is a little more tricky in the case where one
-- of the results is an Acc.  In that case, the result is not an
-- array, but another Acc.  Any Acc result of a Map must correspond to
-- an Acc that is an to the map, and the result is initialised
-- to be that input.  This requires a 1:1 relationship between Acc
-- inputs and Acc outputs, which the type checker should enforce.
-- There is no guarantee that the map results appear in any particular
-- order (e.g. accumulator results before non-accumulator results), so
-- we need to do a little sleuthing to establish the relationship.
--
-- Inside the loop, the non-Acc parameters to map_lam for-in
-- parameters.  Acc parameters refer to loop parameters for the
-- corresponding Map result instead.
--
-- Intuitively, a Screma(w,
--                       (scan_op, scan_ne),
--                       (red_op, red_ne),
--                       map_fn,
--                       {acc_input, arr_input})
--
-- then becomes
--
-- loop (scan_acc, scan_arr, red_acc, map_acc, map_arr) =
--   for i < w, x in arr_input do
--     let (a,b,map_acc',d) = map_fn(map_acc, x)
--     let scan_acc' = scan_op(scan_acc, a)
--     let scan_arr[i] = scan_acc'
--     let red_acc' = red_op(red_acc, b)
--     let map_arr[i] = d
--     in (scan_acc', scan_arr', red_acc', map_acc', map_arr)
--
-- A similar operation is done for Stream.