{-# LANGUAGE TypeFamilies #-}

module Futhark.Pass.Flatten.Intrablock
  ( IntrablockResult (..),
    intrablockParallelise,
    intraMinInnerPar,
    InBlockMapTransformer,
  )
where

import Control.Monad
import Control.Monad.RWS
import Control.Monad.Writer
import Data.Map qualified as M
import Data.Set qualified as S
import Futhark.Analysis.PrimExp.Convert
import Futhark.IR.GPU hiding (HistOp)
import Futhark.IR.SOACS
import Futhark.MonadFreshNames
import Futhark.Pass.Flatten.Distribute
import Futhark.Pass.Flatten.General
import Futhark.Pass.Flatten.PreProcess (preprocessLambda)
import Futhark.Tools
import Futhark.Transform.FirstOrderTransform qualified as FOT
import Futhark.Transform.Rename
import Futhark.Transform.ToGPU
import Prelude hiding (log)

-- | The minimum amount of inner parallelism we require (by default)
-- in intra-group versions.
intraMinInnerPar :: Int64
intraMinInnerPar :: Int64
intraMinInnerPar = Int64
32

data IntrablockResult = IntrablockResult
  { IntrablockResult -> SubExp
intraMinPar :: SubExp,
    IntrablockResult -> SubExp
intraAvailPar :: SubExp,
    IntrablockResult -> SubExp
intraThreadBlockSize :: SubExp,
    IntrablockResult -> Stms GPU
intraPreludeStms :: Stms GPU,
    IntrablockResult -> Stms GPU
intraKernelStms :: Stms GPU,
    IntrablockResult -> [VName]
intraResultNames :: [VName]
  }

type InBlockMapTransformer =
  Pat Type ->
  SubExp ->
  [VName] ->
  Lambda SOACS ->
  FlattenM ()

foldBinOp' :: (MonadBuilder m) => BinOp -> [SubExp] -> m (Exp (Rep m))
foldBinOp' :: forall (m :: * -> *).
MonadBuilder m =>
BinOp -> Segments -> m (Exp (Rep m))
foldBinOp' BinOp
_ [] = SubExp -> m (Exp (Rep m))
forall (m :: * -> *). MonadBuilder m => SubExp -> m (Exp (Rep m))
eSubExp (SubExp -> m (Exp (Rep m))) -> SubExp -> m (Exp (Rep m))
forall a b. (a -> b) -> a -> b
$ IntType -> Integer -> SubExp
intConst IntType
Int64 Integer
1
foldBinOp' BinOp
bop (SubExp
x : Segments
xs) = BinOp -> SubExp -> Segments -> m (Exp (Rep m))
forall (m :: * -> *).
MonadBuilder m =>
BinOp -> SubExp -> Segments -> m (Exp (Rep m))
foldBinOp BinOp
bop SubExp
x Segments
xs

-- | Extract the intra-block parallelism used by the body.
findParallelism :: KernelBody GPU -> [[SubExp]]
findParallelism :: KernelBody GPU -> [Segments]
findParallelism = Set Segments -> [Segments]
forall a. Set a -> [a]
S.toList (Set Segments -> [Segments])
-> (KernelBody GPU -> Set Segments) -> KernelBody GPU -> [Segments]
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Writer (Set Segments) () -> Set Segments
forall w a. Writer w a -> w
execWriter (Writer (Set Segments) () -> Set Segments)
-> (KernelBody GPU -> Writer (Set Segments) ())
-> KernelBody GPU
-> Set Segments
forall b c a. (b -> c) -> (a -> b) -> a -> c
. KernelBody GPU -> Writer (Set Segments) ()
onKernelBody
  where
    onKernelBody :: KernelBody GPU -> Writer (Set Segments) ()
onKernelBody = (Stm GPU -> Writer (Set Segments) ())
-> Stms GPU -> Writer (Set Segments) ()
forall (t :: * -> *) (m :: * -> *) a b.
(Foldable t, Monad m) =>
(a -> m b) -> t a -> m ()
mapM_ Stm GPU -> Writer (Set Segments) ()
onStm (Stms GPU -> Writer (Set Segments) ())
-> (KernelBody GPU -> Stms GPU)
-> KernelBody GPU
-> Writer (Set Segments) ()
forall b c a. (b -> c) -> (a -> b) -> a -> c
. KernelBody GPU -> Stms GPU
forall rep res. GBody rep res -> Stms rep
bodyStms
    onBody :: Body GPU -> Writer (Set Segments) ()
onBody = (Stm GPU -> Writer (Set Segments) ())
-> Stms GPU -> Writer (Set Segments) ()
forall (t :: * -> *) (m :: * -> *) a b.
(Foldable t, Monad m) =>
(a -> m b) -> t a -> m ()
mapM_ Stm GPU -> Writer (Set Segments) ()
onStm (Stms GPU -> Writer (Set Segments) ())
-> (Body GPU -> Stms GPU) -> Body GPU -> Writer (Set Segments) ()
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Body GPU -> Stms GPU
forall rep res. GBody rep res -> Stms rep
bodyStms
    onStm :: Stm GPU -> Writer (Set Segments) ()
onStm Stm GPU
stm = Walker GPU (WriterT (Set Segments) Identity)
-> Exp GPU -> Writer (Set Segments) ()
forall (m :: * -> *) rep.
Monad m =>
Walker rep m -> Exp rep -> m ()
walkExpM Walker GPU (WriterT (Set Segments) Identity)
walker (Exp GPU -> Writer (Set Segments) ())
-> Exp GPU -> Writer (Set Segments) ()
forall a b. (a -> b) -> a -> b
$ Stm GPU -> Exp GPU
forall rep. Stm rep -> Exp rep
stmExp Stm GPU
stm
    walker :: Walker GPU (WriterT (Set Segments) Identity)
walker =
      (forall rep (m :: * -> *). Monad m => Walker rep m
identityWalker @GPU)
        { walkOnBody = const onBody,
          walkOnOp = onOp
        }
    onOp :: HostOp SOAC GPU -> Writer (Set Segments) ()
onOp (SegOp SegOp SegLevel GPU
op) = do
      Set Segments -> Writer (Set Segments) ()
forall w (m :: * -> *). MonadWriter w m => w -> m ()
tell (Set Segments -> Writer (Set Segments) ())
-> Set Segments -> Writer (Set Segments) ()
forall a b. (a -> b) -> a -> b
$ Segments -> Set Segments
forall a. a -> Set a
S.singleton (Segments -> Set Segments) -> Segments -> Set Segments
forall a b. (a -> b) -> a -> b
$ SegSpace -> Segments
segSpaceDims (SegSpace -> Segments) -> SegSpace -> Segments
forall a b. (a -> b) -> a -> b
$ SegOp SegLevel GPU -> SegSpace
forall lvl rep. SegOp lvl rep -> SegSpace
segSpace SegOp SegLevel GPU
op
      KernelBody GPU -> Writer (Set Segments) ()
onKernelBody (KernelBody GPU -> Writer (Set Segments) ())
-> KernelBody GPU -> Writer (Set Segments) ()
forall a b. (a -> b) -> a -> b
$ SegOp SegLevel GPU -> KernelBody GPU
forall lvl rep. SegOp lvl rep -> KernelBody rep
segBody SegOp SegLevel GPU
op
    onOp HostOp SOAC GPU
_ = () -> Writer (Set Segments) ()
forall a. a -> WriterT (Set Segments) Identity a
forall (f :: * -> *) a. Applicative f => a -> f a
pure ()

computeThreadBlockSize :: [[SubExp]] -> [[SubExp]] -> FlattenM (SubExp, SubExp)
computeThreadBlockSize :: [Segments] -> [Segments] -> FlattenM (SubExp, SubExp)
computeThreadBlockSize [Segments]
wss_min [Segments]
wss_avail = do
  ws_min <-
    (Segments -> FlattenM SubExp) -> [Segments] -> FlattenM Segments
forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
(a -> m b) -> t a -> m (t b)
forall (m :: * -> *) a b. Monad m => (a -> m b) -> [a] -> m [b]
mapM (Name -> Exp (Rep FlattenM) -> FlattenM SubExp
forall (m :: * -> *).
MonadBuilder m =>
Name -> Exp (Rep m) -> m SubExp
letSubExp Name
"one_intra_par_min" (Exp GPU -> FlattenM SubExp)
-> (Segments -> FlattenM (Exp GPU)) -> Segments -> FlattenM SubExp
forall (m :: * -> *) b c a.
Monad m =>
(b -> m c) -> (a -> m b) -> a -> m c
<=< BinOp -> Segments -> FlattenM (Exp (Rep FlattenM))
forall (m :: * -> *).
MonadBuilder m =>
BinOp -> Segments -> m (Exp (Rep m))
foldBinOp' (IntType -> Overflow -> BinOp
Mul IntType
Int64 Overflow
OverflowUndef)) ([Segments] -> FlattenM Segments)
-> [Segments] -> FlattenM Segments
forall a b. (a -> b) -> a -> b
$
      (Segments -> Bool) -> [Segments] -> [Segments]
forall a. (a -> Bool) -> [a] -> [a]
filter (Bool -> Bool
not (Bool -> Bool) -> (Segments -> Bool) -> Segments -> Bool
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Segments -> Bool
forall a. [a] -> Bool
forall (t :: * -> *) a. Foldable t => t a -> Bool
null) [Segments]
wss_min
  ws_avail <-
    mapM (letSubExp "one_intra_par_avail" <=< foldBinOp' (Mul Int64 OverflowUndef)) $
      filter (not . null) wss_avail

  -- The amount of parallelism available *in the worst case* is
  -- equal to the smallest parallel loop, or *at least* 1.
  intra_avail_par <-
    letSubExp "intra_avail_par" =<< foldBinOp' (SMin Int64) ws_avail

  tblock_size <- newVName "computed_tblock_size"
  -- The group size is either the maximum of the minimum parallelism
  -- exploited, or the desired parallelism (bounded by the max group
  -- size) in case there is no minimum.
  letBindNames [tblock_size]
    =<< if null ws_min
      then
        eBinOp
          (SMin Int64)
          (eSubExp =<< letSubExp "max_tblock_size" (Op $ SizeOp $ GetSizeMax SizeThreadBlock))
          (eSubExp intra_avail_par)
      else foldBinOp' (SMax Int64) ws_min
  pure (intra_avail_par, Var tblock_size)

-- | Check whether this result is actually something that is acceptable to use:
-- the parallel dimensions must all be bound outside the kernel, as nonuniform
-- parallelism cannot be exploited inside a thread block.
noNonuniformPar :: Names -> FlattenM IntrablockResult -> FlattenM (Maybe IntrablockResult)
noNonuniformPar :: Names
-> FlattenM IntrablockResult -> FlattenM (Maybe IntrablockResult)
noNonuniformPar Names
pars FlattenM IntrablockResult
m = do
  outside_scope <- FlattenM (Scope GPU)
forall rep (m :: * -> *). HasScope rep m => m (Scope rep)
askScope
  if allNames (`M.member` outside_scope) pars
    then Just <$> m
    else pure Nothing

intrablockParallelise ::
  InBlockMapTransformer ->
  Segments ->
  DistEnv ->
  DistInputs ->
  [DistResult] ->
  Pat Type ->
  StmAux () ->
  SubExp ->
  [VName] ->
  Lambda SOACS ->
  FlattenM (Maybe IntrablockResult)
intrablockParallelise :: InBlockMapTransformer
-> Segments
-> DistEnv
-> DistInputs
-> [DistResult]
-> Pat Type
-> StmAux ()
-> SubExp
-> [VName]
-> Lambda SOACS
-> FlattenM (Maybe IntrablockResult)
intrablockParallelise InBlockMapTransformer
map_in_block Segments
segments DistEnv
env DistInputs
inps [DistResult]
dist_res Pat Type
_pat StmAux ()
aux SubExp
w [VName]
arrs Lambda SOACS
lam0 = do
  gpu_scope <- FlattenM (Scope GPU)
forall rep (m :: * -> *). HasScope rep m => m (Scope rep)
askScope
  let pp_scope = Scope GPU -> Scope SOACS
forall fromrep torep.
SameScope fromrep torep =>
Scope fromrep -> Scope torep
castScope (Scope GPU -> Scope SOACS) -> Scope GPU -> Scope SOACS
forall a b. (a -> b) -> a -> b
$ DistInputs -> Scope GPU
scopeOfDistInputs DistInputs
inps Scope GPU -> Scope GPU -> Scope GPU
forall a. Semigroup a => a -> a -> a
<> Scope GPU
gpu_scope
  lam <- renameLambda =<< preprocessLambda pp_scope lam0

  let result_ts =
        [ Type
t Type -> Shape -> Type
`arrayOfShape` Segments -> Shape
segmentsShape Segments
segments
        | DistResult ResTag
_ (DistType Segments
_ Rank
_ Type
t) VName
_ <- [DistResult]
dist_res
        ]

  -- Reconstruct the per-enclosing-segment inputs, so they can be read at the
  -- appropriate indices inside the kernel.
  (param_inputs, input_prelude_stms) <-
    collectStms $
      zipWithM (prepareRegularMapInput segments env inps) (lambdaParams lam) arrs

  -- A free variable of the map body is a per-enclosing-segment value: it is
  -- brought into the body's scope as a distributed input and reconstructed at
  -- the outer indices inside the kernel (via 'readInBlockInputs' below). With no
  -- enclosing segments, however, it is an ordinary top-level value already in
  -- the kernel's outer scope, so there is nothing to reconstruct - hence the
  -- empty 'free_inputs', which makes both the scope extension and the read below
  -- vanish.
  free_inputs <- if null segments then pure [] else freeInputsFor inps lam
  kbody <-
    localScope (scopeOfDistInputs free_inputs <> scopeOfLParams (lambdaParams lam)) $
      intrablockParalleliseBody map_in_block (lambdaBody lam)

  nested_pat <-
    fmap Pat $
      zipWith PatElem
        <$> mapM (newName . distResName) dist_res
        <*> pure result_ts

  let new_segments = Segments
segments Segments -> Segments -> Segments
forall a. Semigroup a => a -> a -> a
<> SubExp -> Segments
forall a. a -> [a]
forall (f :: * -> *) a. Applicative f => a -> f a
pure SubExp
w
      wss = KernelBody GPU -> [Segments]
findParallelism KernelBody GPU
kbody

  noNonuniformPar (freeIn wss) $ do
    ((intra_avail_par, tblock_size, kspace, num_tblocks), prelude_stms) <-
      collectStms $ do
        num_tblocks <-
          letSubExp "intra_num_tblocks"
            =<< foldBinOp' (Mul Int64 OverflowUndef) new_segments
        (intra_avail_par, tblock_size) <- computeThreadBlockSize wss wss
        gtids <- mapM (const $ newVName "gtid") new_segments
        kspace <- mkSegSpace $ zip gtids new_segments
        pure (intra_avail_par, tblock_size, kspace, num_tblocks)

    read_input_stms <-
      collectStms_ . localScope (scopeOfSegSpace kspace <> scopeOf input_prelude_stms <> scopeOf prelude_stms) $ do
        let SegSpace _ gtids_and_dims = kspace
            full_is = ((VName, SubExp) -> SubExp) -> [(VName, SubExp)] -> Segments
forall a b. (a -> b) -> [a] -> [b]
map (VName -> SubExp
Var (VName -> SubExp)
-> ((VName, SubExp) -> VName) -> (VName, SubExp) -> SubExp
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (VName, SubExp) -> VName
forall a b. (a, b) -> a
fst) [(VName, SubExp)]
gtids_and_dims
            outer_is = Int -> Segments -> Segments
forall a. Int -> [a] -> [a]
take (Segments -> Int
segmentsRank Segments
segments) Segments
full_is
        readInBlockInputs segments env outer_is free_inputs
        readInBlockInputs new_segments mempty full_is param_inputs

    let kbody' = KernelBody GPU
kbody {bodyStms = read_input_stms <> bodyStms kbody}
        rts = (Type -> Type) -> [Type] -> [Type]
forall a b. (a -> b) -> [a] -> [b]
map (Segments -> Int
forall a. [a] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length Segments
new_segments Int -> Type -> Type
forall d m.
Int -> TypeBase (ShapeBase d) m -> TypeBase (ShapeBase d) m
`stripArray`) [Type]
result_ts
        grid = Count NumBlocks SubExp -> Count BlockSize SubExp -> KernelGrid
KernelGrid (SubExp -> Count NumBlocks SubExp
forall {k} (u :: k) e. e -> Count u e
Count SubExp
num_tblocks) (SubExp -> Count BlockSize SubExp
forall {k} (u :: k) e. e -> Count u e
Count SubExp
tblock_size)
        lvl = SegVirt -> Maybe KernelGrid -> SegLevel
SegBlock SegVirt
SegNoVirt (KernelGrid -> Maybe KernelGrid
forall a. a -> Maybe a
Just KernelGrid
grid)
        kstm = Pat (LetDec GPU) -> StmAux (ExpDec GPU) -> Exp GPU -> Stm GPU
forall rep.
Pat (LetDec rep) -> StmAux (ExpDec rep) -> Exp rep -> Stm rep
Let Pat Type
Pat (LetDec GPU)
nested_pat StmAux ()
StmAux (ExpDec GPU)
aux (Exp GPU -> Stm GPU) -> Exp GPU -> Stm GPU
forall a b. (a -> b) -> a -> b
$ Op GPU -> Exp GPU
forall rep. Op rep -> Exp rep
Op (Op GPU -> Exp GPU) -> Op GPU -> Exp GPU
forall a b. (a -> b) -> a -> b
$ SegOp SegLevel GPU -> HostOp SOAC GPU
forall (op :: * -> *) rep. SegOp SegLevel rep -> HostOp op rep
SegOp (SegOp SegLevel GPU -> HostOp SOAC GPU)
-> SegOp SegLevel GPU -> HostOp SOAC GPU
forall a b. (a -> b) -> a -> b
$ SegLevel
-> SegSpace -> [Type] -> KernelBody GPU -> SegOp SegLevel GPU
forall lvl rep.
lvl -> SegSpace -> [Type] -> KernelBody rep -> SegOp lvl rep
SegMap SegLevel
lvl SegSpace
kspace [Type]
rts KernelBody GPU
kbody'

    pure $
      IntrablockResult
        { intraMinPar = intra_avail_par,
          intraAvailPar = intra_avail_par,
          intraThreadBlockSize = tblock_size,
          intraPreludeStms = input_prelude_stms <> prelude_stms,
          intraKernelStms = oneStm kstm,
          intraResultNames = patNames nested_pat
        }

readInBlockInputs :: Segments -> DistEnv -> [SubExp] -> DistInputs -> FlattenM ()
readInBlockInputs :: Segments -> DistEnv -> Segments -> DistInputs -> FlattenM ()
readInBlockInputs Segments
segments DistEnv
env Segments
is DistInputs
inputs =
  ((VName, DistInput) -> FlattenM ()) -> DistInputs -> FlattenM ()
forall (t :: * -> *) (m :: * -> *) a b.
(Foldable t, Monad m) =>
(a -> m b) -> t a -> m ()
mapM_ (VName, DistInput) -> FlattenM ()
onInput DistInputs
inputs
  where
    onInput :: (VName, DistInput) -> FlattenM ()
onInput (VName
v, DistInput
inp) = do
      v' <- Segments
-> DistEnv -> Segments -> DistInputs -> VName -> FlattenM VName
readInputVar Segments
segments DistEnv
env Segments
is DistInputs
inputs VName
v
      let t = DistInput -> Type
distInputType DistInput
inp
      if isAcc t
        then
          letBindNames [v] $ BasicOp $ SubExp $ Var v'
        else
          if arrayRank t > 0
            then
              letBindNames [v] $ BasicOp $ Replicate mempty $ Var v'
            else
              letBindNames [v] $ BasicOp $ SubExp $ Var v'

prepareRegularMapInput ::
  Segments ->
  DistEnv ->
  DistInputs ->
  Param Type ->
  VName ->
  FlattenM (VName, DistInput)
prepareRegularMapInput :: Segments
-> DistEnv
-> DistInputs
-> Param Type
-> VName
-> FlattenM (VName, DistInput)
prepareRegularMapInput Segments
segments DistEnv
env DistInputs
inps Param Type
p VName
arr = do
  t <- DistInputs -> VName -> FlattenM Type
lookupInputType DistInputs
inps VName
arr
  let expectedShape = Segments -> Shape
segmentsShape Segments
segments Shape -> Shape -> Shape
forall a. Semigroup a => a -> a -> a
<> Type -> Shape
forall shape o. ArrayShape shape => TypeBase shape o -> shape
arrayShape Type
t
      lvl = SegVirt -> Maybe KernelGrid -> SegLevel
SegThread SegVirt
SegVirt Maybe KernelGrid
forall a. Maybe a
Nothing
  arr_rep <- liftVarRegular lvl segments inps env expectedShape arr
  pure (paramName p, DistInputFree arr_rep (paramType p))

freeInputsFor :: DistInputs -> Lambda SOACS -> FlattenM DistInputs
freeInputsFor :: DistInputs -> Lambda SOACS -> FlattenM DistInputs
freeInputsFor DistInputs
inps Lambda SOACS
lam =
  do
    let free :: Names
free = Lambda SOACS -> Names
forall a. FreeIn a => a -> Names
freeIn Lambda SOACS
lam
    free_sizes <-
      (Type -> Names) -> [Type] -> Names
forall m a. Monoid m => (a -> m) -> [a] -> m
forall (t :: * -> *) m a.
(Foldable t, Monoid m) =>
(a -> m) -> t a -> m
foldMap Type -> Names
forall a. FreeIn a => a -> Names
freeIn ([Type] -> Names) -> FlattenM [Type] -> FlattenM Names
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> (VName -> FlattenM Type) -> [VName] -> FlattenM [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 (DistInputs -> VName -> FlattenM Type
lookupInputType DistInputs
inps) (Names -> [VName]
namesToList Names
free)
    pure
      [ (v, inp)
      | v <- namesToList $ free <> free_sizes,
        Just inp <- [lookup v inps]
      ]

intrablockBody :: InBlockMapTransformer -> Body SOACS -> FlattenM (Body GPU)
intrablockBody :: InBlockMapTransformer -> Body SOACS -> FlattenM (Body GPU)
intrablockBody InBlockMapTransformer
map_in_block Body SOACS
body = do
  stms <- FlattenM () -> FlattenM (Stms (Rep FlattenM))
forall (m :: * -> *) a. MonadBuilder m => m a -> m (Stms (Rep m))
collectStms_ (FlattenM () -> FlattenM (Stms (Rep FlattenM)))
-> FlattenM () -> FlattenM (Stms (Rep FlattenM))
forall a b. (a -> b) -> a -> b
$ InBlockMapTransformer -> Stms SOACS -> FlattenM ()
intrablockStms InBlockMapTransformer
map_in_block (Stms SOACS -> FlattenM ()) -> Stms SOACS -> FlattenM ()
forall a b. (a -> b) -> a -> b
$ Body SOACS -> Stms SOACS
forall rep res. GBody rep res -> Stms rep
bodyStms Body SOACS
body
  pure $ mkBody stms $ bodyResult body

intrablockLambda :: InBlockMapTransformer -> Lambda SOACS -> FlattenM (Lambda GPU)
intrablockLambda :: InBlockMapTransformer -> Lambda SOACS -> FlattenM (Lambda GPU)
intrablockLambda InBlockMapTransformer
map_in_block Lambda SOACS
lam =
  [LParam (Rep FlattenM)]
-> FlattenM [SubExpRes] -> FlattenM (Lambda (Rep FlattenM))
forall (m :: * -> *).
MonadBuilder m =>
[LParam (Rep m)] -> m [SubExpRes] -> m (Lambda (Rep m))
mkLambda (Lambda SOACS -> [LParam SOACS]
forall rep t. GLambda rep t -> [LParam rep]
lambdaParams Lambda SOACS
lam) (FlattenM [SubExpRes] -> FlattenM (Lambda (Rep FlattenM)))
-> FlattenM [SubExpRes] -> FlattenM (Lambda (Rep FlattenM))
forall a b. (a -> b) -> a -> b
$
    Body (Rep FlattenM) -> FlattenM [SubExpRes]
Body GPU -> FlattenM [SubExpRes]
forall (m :: * -> *).
MonadBuilder m =>
Body (Rep m) -> m [SubExpRes]
bodyBind (Body GPU -> FlattenM [SubExpRes])
-> FlattenM (Body GPU) -> FlattenM [SubExpRes]
forall (m :: * -> *) a b. Monad m => (a -> m b) -> m a -> m b
=<< InBlockMapTransformer -> Body SOACS -> FlattenM (Body GPU)
intrablockBody InBlockMapTransformer
map_in_block (Lambda SOACS -> Body SOACS
forall rep t. GLambda rep t -> Body rep
lambdaBody Lambda SOACS
lam)

intrablockWithAccInput :: InBlockMapTransformer -> WithAccInput SOACS -> FlattenM (WithAccInput GPU)
intrablockWithAccInput :: InBlockMapTransformer
-> WithAccInput SOACS -> FlattenM (WithAccInput GPU)
intrablockWithAccInput InBlockMapTransformer
_ (Shape
shape, [VName]
arrs, Maybe (Lambda SOACS, Segments)
Nothing) =
  WithAccInput GPU -> FlattenM (WithAccInput GPU)
forall a. a -> FlattenM a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (Shape
shape, [VName]
arrs, Maybe (Lambda GPU, Segments)
forall a. Maybe a
Nothing)
intrablockWithAccInput InBlockMapTransformer
map_in_block (Shape
shape, [VName]
arrs, Just (Lambda SOACS
lam, Segments
nes)) = do
  lam' <- InBlockMapTransformer -> Lambda SOACS -> FlattenM (Lambda GPU)
intrablockLambda InBlockMapTransformer
map_in_block Lambda SOACS
lam
  pure (shape, arrs, Just (lam', nes))

intrablockStm :: InBlockMapTransformer -> Stm SOACS -> FlattenM ()
intrablockStm :: InBlockMapTransformer -> Stm SOACS -> FlattenM ()
intrablockStm InBlockMapTransformer
map_in_block stm :: Stm SOACS
stm@(Let Pat (LetDec SOACS)
pat StmAux (ExpDec SOACS)
aux Exp SOACS
e) = do
  scope <- FlattenM (Scope GPU)
forall rep (m :: * -> *). HasScope rep m => m (Scope rep)
askScope
  let lvl = SegVirt -> SegLevel
SegThreadInBlock SegVirt
SegNoVirt

  case e of
    Loop [(FParam SOACS, SubExp)]
merge LoopForm
form Body SOACS
loopbody ->
      Scope GPU -> FlattenM () -> FlattenM ()
forall a. Scope GPU -> FlattenM a -> FlattenM a
forall rep (m :: * -> *) a.
LocalScope rep m =>
Scope rep -> m a -> m a
localScope (LoopForm -> Scope GPU
forall rep. LoopForm -> Scope rep
scopeOfLoopForm LoopForm
form Scope GPU -> Scope GPU -> Scope GPU
forall a. Semigroup a => a -> a -> a
<> [Param (FParamInfo GPU)] -> Scope GPU
forall rep dec. (FParamInfo rep ~ dec) => [Param dec] -> Scope rep
scopeOfFParams (((Param (FParamInfo GPU), SubExp) -> Param (FParamInfo GPU))
-> [(Param (FParamInfo GPU), SubExp)] -> [Param (FParamInfo GPU)]
forall a b. (a -> b) -> [a] -> [b]
map (Param (FParamInfo GPU), SubExp) -> Param (FParamInfo GPU)
forall a b. (a, b) -> a
fst [(FParam SOACS, SubExp)]
[(Param (FParamInfo GPU), SubExp)]
merge)) (FlattenM () -> FlattenM ()) -> FlattenM () -> FlattenM ()
forall a b. (a -> b) -> a -> b
$ do
        loopbody' <- InBlockMapTransformer -> Body SOACS -> FlattenM (Body GPU)
intrablockBody InBlockMapTransformer
map_in_block Body SOACS
loopbody
        certifying (stmAuxCerts aux) . letBind pat $
          Loop merge form loopbody'
    Match Segments
cond [Case (Body SOACS)]
cases Body SOACS
defbody MatchDec (BranchType SOACS)
ifdec -> do
      cases' <- (Case (Body SOACS) -> FlattenM (Case (Body GPU)))
-> [Case (Body SOACS)] -> FlattenM [Case (Body GPU)]
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 ((Body SOACS -> FlattenM (Body GPU))
-> Case (Body SOACS) -> FlattenM (Case (Body GPU))
forall (t :: * -> *) (f :: * -> *) a b.
(Traversable t, Applicative f) =>
(a -> f b) -> t a -> f (t b)
forall (f :: * -> *) a b.
Applicative f =>
(a -> f b) -> Case a -> f (Case b)
traverse ((Body SOACS -> FlattenM (Body GPU))
 -> Case (Body SOACS) -> FlattenM (Case (Body GPU)))
-> (Body SOACS -> FlattenM (Body GPU))
-> Case (Body SOACS)
-> FlattenM (Case (Body GPU))
forall a b. (a -> b) -> a -> b
$ InBlockMapTransformer -> Body SOACS -> FlattenM (Body GPU)
intrablockBody InBlockMapTransformer
map_in_block) [Case (Body SOACS)]
cases
      defbody' <- intrablockBody map_in_block defbody
      certifying (stmAuxCerts aux) . letBind pat $
        Match cond cases' defbody' ifdec
    WithAcc [WithAccInput SOACS]
inputs Lambda SOACS
lam -> do
      inputs' <- (WithAccInput SOACS -> FlattenM (WithAccInput GPU))
-> [WithAccInput SOACS] -> FlattenM [WithAccInput GPU]
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 (InBlockMapTransformer
-> WithAccInput SOACS -> FlattenM (WithAccInput GPU)
intrablockWithAccInput InBlockMapTransformer
map_in_block) [WithAccInput SOACS]
inputs
      lam' <- intrablockLambda map_in_block lam
      certifying (stmAuxCerts aux) . letBind pat $ WithAcc inputs' lam'
    Op Op SOACS
soac
      | Attr
"sequential_outer" Attr -> Attrs -> Bool
`inAttrs` StmAux () -> Attrs
forall dec. StmAux dec -> Attrs
stmAuxAttrs StmAux ()
StmAux (ExpDec SOACS)
aux ->
          InBlockMapTransformer -> Stms SOACS -> FlattenM ()
intrablockStms InBlockMapTransformer
map_in_block (Stms SOACS -> FlattenM ())
-> (Stms SOACS -> Stms SOACS) -> Stms SOACS -> FlattenM ()
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (Stm SOACS -> Stm SOACS) -> Stms SOACS -> Stms SOACS
forall a b. (a -> b) -> Seq a -> Seq b
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap (Certs -> Stm SOACS -> Stm SOACS
forall rep. Certs -> Stm rep -> Stm rep
certify (StmAux () -> Certs
forall dec. StmAux dec -> Certs
stmAuxCerts StmAux ()
StmAux (ExpDec SOACS)
aux))
            (Stms SOACS -> FlattenM ()) -> FlattenM (Stms SOACS) -> FlattenM ()
forall (m :: * -> *) a b. Monad m => (a -> m b) -> m a -> m b
=<< Builder SOACS () -> FlattenM (Stms SOACS)
forall (m :: * -> *) somerep rep a.
(MonadFreshNames m, HasScope somerep m, SameScope somerep rep) =>
Builder rep a -> m (Stms rep)
runBuilder_ (Pat (LetDec (Rep (BuilderT SOACS (State VNameSource))))
-> SOAC (Rep (BuilderT SOACS (State VNameSource)))
-> Builder SOACS ()
forall (m :: * -> *).
Transformer m =>
Pat (LetDec (Rep m)) -> SOAC (Rep m) -> m ()
FOT.transformSOAC Pat (LetDec (Rep (BuilderT SOACS (State VNameSource))))
Pat (LetDec SOACS)
pat Op SOACS
SOAC (Rep (BuilderT SOACS (State VNameSource)))
soac)
    Op (Screma SubExp
w [VName]
arrs ScremaForm SOACS
form)
      | Just Lambda SOACS
lam <- ScremaForm SOACS -> Maybe (Lambda SOACS)
forall rep. ScremaForm rep -> Maybe (Lambda rep)
isMapSOAC ScremaForm SOACS
form ->
          InBlockMapTransformer
map_in_block Pat Type
Pat (LetDec SOACS)
pat SubExp
w [VName]
arrs Lambda SOACS
lam
    Op (Screma SubExp
w [VName]
arrs ScremaForm SOACS
form)
      | Just (Lambda SOACS
post_lam, [Scan SOACS]
scans, Lambda SOACS
mapfun) <- ScremaForm SOACS
-> Maybe (Lambda SOACS, [Scan SOACS], Lambda SOACS)
forall rep.
ScremaForm rep -> Maybe (Lambda rep, [Scan rep], Lambda rep)
isMaposcanomapSOAC ScremaForm SOACS
form,
        -- FIXME: Futhark.CodeGen.ImpGen.GPU.Block.compileGroupOp
        -- cannot handle multiple scan operators yet.
        Scan Lambda SOACS
scanfun Segments
nes <- [Scan SOACS] -> Scan SOACS
forall rep. Buildable rep => [Scan rep] -> Scan rep
singleScan [Scan SOACS]
scans -> do
          let scanfun' :: Lambda GPU
scanfun' = Lambda SOACS -> Lambda GPU
soacsLambdaToGPU Lambda SOACS
scanfun
              mapfun' :: Lambda GPU
mapfun' = Lambda SOACS -> Lambda GPU
soacsLambdaToGPU Lambda SOACS
mapfun
              post_op :: Lambda GPU
post_op = Lambda SOACS -> Lambda GPU
soacsLambdaToGPU Lambda SOACS
post_lam
          scan_res <- Certs -> FlattenM [VName] -> FlattenM [VName]
forall a. Certs -> FlattenM a -> FlattenM a
forall (m :: * -> *) a. MonadBuilder m => Certs -> m a -> m a
certifying (StmAux () -> Certs
forall dec. StmAux dec -> Certs
stmAuxCerts StmAux ()
StmAux (ExpDec SOACS)
aux) (FlattenM [VName] -> FlattenM [VName])
-> FlattenM [VName] -> FlattenM [VName]
forall a b. (a -> b) -> a -> b
$ SegLevel
-> Segments
-> Name
-> Lambda GPU
-> Shape
-> Segments
-> Lambda GPU
-> Lambda GPU
-> [VName]
-> (Segments -> FlattenM ())
-> FlattenM [VName]
forall (m :: * -> *).
(MonadBuilder m, Rep m ~ GPU) =>
SegLevel
-> Segments
-> Name
-> Lambda GPU
-> Shape
-> Segments
-> Lambda GPU
-> Lambda GPU
-> [VName]
-> (Segments -> m ())
-> m [VName]
genUniformSegScanomapWithPost SegLevel
lvl (SubExp -> Segments
forall a. a -> [a]
forall (f :: * -> *) a. Applicative f => a -> f a
pure SubExp
w) Name
"intra_maposcanomap" Lambda GPU
scanfun' Shape
forall a. Monoid a => a
mempty Segments
nes Lambda GPU
post_op Lambda GPU
mapfun' [VName]
arrs (FlattenM () -> Segments -> FlattenM ()
forall a b. a -> b -> a
const (FlattenM () -> Segments -> FlattenM ())
-> FlattenM () -> Segments -> FlattenM ()
forall a b. (a -> b) -> a -> b
$ () -> FlattenM ()
forall a. a -> FlattenM a
forall (f :: * -> *) a. Applicative f => a -> f a
pure ())
          zipWithM_
            ( \PatElem Type
pe VName
v ->
                [VName] -> Exp (Rep FlattenM) -> FlattenM ()
forall (m :: * -> *).
MonadBuilder m =>
[VName] -> Exp (Rep m) -> m ()
letBindNames [PatElem Type -> VName
forall dec. PatElem dec -> VName
patElemName PatElem Type
pe] (Exp (Rep FlattenM) -> FlattenM ())
-> Exp (Rep FlattenM) -> FlattenM ()
forall a b. (a -> b) -> a -> b
$ BasicOp -> Exp (Rep FlattenM)
forall rep. BasicOp -> Exp rep
BasicOp (BasicOp -> Exp (Rep FlattenM)) -> BasicOp -> Exp (Rep FlattenM)
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
v
            )
            (patElems pat)
            scan_res
    Op (Screma SubExp
w [VName]
arrs ScremaForm SOACS
form)
      | Just ([Reduce SOACS]
reds, Lambda SOACS
map_lam) <- ScremaForm SOACS -> Maybe ([Reduce SOACS], Lambda SOACS)
forall rep. ScremaForm rep -> Maybe ([Reduce rep], Lambda rep)
isRedomapSOAC ScremaForm SOACS
form -> do
          let onRed :: Reduce SOACS -> Reduce GPU
onRed Reduce SOACS
red =
                let red_lam :: Lambda SOACS
red_lam = Reduce SOACS -> Lambda SOACS
forall rep. Reduce rep -> Lambda rep
redLambda Reduce SOACS
red
                    comm :: Commutativity
comm
                      | Lambda SOACS -> Bool
forall rep. Lambda rep -> Bool
commutativeLambda Lambda SOACS
red_lam = Commutativity
Commutative
                      | Bool
otherwise = Reduce SOACS -> Commutativity
forall rep. Reduce rep -> Commutativity
redComm Reduce SOACS
red
                 in Commutativity -> Lambda GPU -> Segments -> Reduce GPU
forall rep. Commutativity -> Lambda rep -> Segments -> Reduce rep
Reduce Commutativity
comm (Lambda SOACS -> Lambda GPU
soacsLambdaToGPU Lambda SOACS
red_lam) (Reduce SOACS -> Segments
forall rep. Reduce rep -> Segments
redNeutral Reduce SOACS
red)
              reds_gpu :: [Reduce GPU]
reds_gpu = (Reduce SOACS -> Reduce GPU) -> [Reduce SOACS] -> [Reduce GPU]
forall a b. (a -> b) -> [a] -> [b]
map Reduce SOACS -> Reduce GPU
onRed [Reduce SOACS]
reds
              map_lam' :: Lambda GPU
map_lam' = Lambda SOACS -> Lambda GPU
soacsLambdaToGPU Lambda SOACS
map_lam
          (red_res, stms) <- Builder GPU [VName] -> FlattenM ([VName], Stms GPU)
forall (m :: * -> *) somerep rep a.
(MonadFreshNames m, HasScope somerep m, SameScope somerep rep) =>
Builder rep a -> m (a, Stms rep)
runBuilder (SegLevel
-> Name
-> Segments
-> [Reduce GPU]
-> Shape
-> Lambda GPU
-> [VName]
-> (Segments -> BuilderT GPU (State VNameSource) ())
-> Builder GPU [VName]
forall (m :: * -> *).
(MonadBuilder m, Rep m ~ GPU) =>
SegLevel
-> Name
-> Segments
-> [Reduce GPU]
-> Shape
-> Lambda GPU
-> [VName]
-> (Segments -> m ())
-> m [VName]
genUniformSegRed SegLevel
lvl Name
"intra_redomap" (SubExp -> Segments
forall a. a -> [a]
forall (f :: * -> *) a. Applicative f => a -> f a
pure SubExp
w) [Reduce GPU]
reds_gpu Shape
forall a. Monoid a => a
mempty Lambda GPU
map_lam' [VName]
arrs (BuilderT GPU (State VNameSource) ()
-> Segments -> BuilderT GPU (State VNameSource) ()
forall a b. a -> b -> a
const (BuilderT GPU (State VNameSource) ()
 -> Segments -> BuilderT GPU (State VNameSource) ())
-> BuilderT GPU (State VNameSource) ()
-> Segments
-> BuilderT GPU (State VNameSource) ()
forall a b. (a -> b) -> a -> b
$ () -> BuilderT GPU (State VNameSource) ()
forall a. a -> BuilderT GPU (State VNameSource) a
forall (f :: * -> *) a. Applicative f => a -> f a
pure ()))
          certifying (stmAuxCerts aux) $ do
            addStms stms
            zipWithM_
              ( \PatElem Type
pe VName
v ->
                  [VName] -> Exp (Rep FlattenM) -> FlattenM ()
forall (m :: * -> *).
MonadBuilder m =>
[VName] -> Exp (Rep m) -> m ()
letBindNames [PatElem Type -> VName
forall dec. PatElem dec -> VName
patElemName PatElem Type
pe] (Exp (Rep FlattenM) -> FlattenM ())
-> Exp (Rep FlattenM) -> FlattenM ()
forall a b. (a -> b) -> a -> b
$ BasicOp -> Exp (Rep FlattenM)
forall rep. BasicOp -> Exp rep
BasicOp (BasicOp -> Exp (Rep FlattenM)) -> BasicOp -> Exp (Rep FlattenM)
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
v
              )
              (patElems pat)
              red_res
    Op (Screma SubExp
w [VName]
arrs ScremaForm SOACS
form) ->
      -- This screma is too complicated for us to immediately do
      -- anything, so split it up and try again.
      (Stm SOACS -> FlattenM ()) -> Stms SOACS -> FlattenM ()
forall (t :: * -> *) (m :: * -> *) a b.
(Foldable t, Monad m) =>
(a -> m b) -> t a -> m ()
mapM_ (InBlockMapTransformer -> Stm SOACS -> FlattenM ()
intrablockStm InBlockMapTransformer
map_in_block) (Stms SOACS -> FlattenM ())
-> (((), Stms SOACS) -> Stms SOACS)
-> ((), Stms SOACS)
-> FlattenM ()
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (Stm SOACS -> Stm SOACS) -> Stms SOACS -> Stms SOACS
forall a b. (a -> b) -> Seq a -> Seq b
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap (Certs -> Stm SOACS -> Stm SOACS
forall rep. Certs -> Stm rep -> Stm rep
certify (StmAux () -> Certs
forall dec. StmAux dec -> Certs
stmAuxCerts StmAux ()
StmAux (ExpDec SOACS)
aux)) (Stms SOACS -> Stms SOACS)
-> (((), Stms SOACS) -> Stms SOACS)
-> ((), Stms SOACS)
-> Stms SOACS
forall b c a. (b -> c) -> (a -> b) -> a -> c
. ((), Stms SOACS) -> Stms SOACS
forall a b. (a, b) -> b
snd
        (((), Stms SOACS) -> FlattenM ())
-> FlattenM ((), Stms SOACS) -> FlattenM ()
forall (m :: * -> *) a b. Monad m => (a -> m b) -> m a -> m b
=<< BuilderT SOACS FlattenM ()
-> Scope SOACS -> FlattenM ((), Stms SOACS)
forall (m :: * -> *) rep a.
MonadFreshNames m =>
BuilderT rep m a -> Scope rep -> m (a, Stms rep)
runBuilderT (Pat (LetDec (Rep (BuilderT SOACS FlattenM)))
-> SubExp
-> ScremaForm (Rep (BuilderT SOACS FlattenM))
-> [VName]
-> BuilderT SOACS FlattenM ()
forall (m :: * -> *).
(MonadBuilder m, Op (Rep m) ~ SOAC (Rep m), Buildable (Rep m)) =>
Pat (LetDec (Rep m))
-> SubExp -> ScremaForm (Rep m) -> [VName] -> m ()
dissectScrema Pat (LetDec (Rep (BuilderT SOACS FlattenM)))
Pat (LetDec SOACS)
pat SubExp
w ScremaForm (Rep (BuilderT SOACS FlattenM))
ScremaForm SOACS
form [VName]
arrs) (Scope GPU -> Scope SOACS
scopeForSOACs Scope GPU
scope)
    Op (Hist SubExp
w [VName]
arrs [HistOp SOACS]
ops Lambda SOACS
bucket_fun) -> do
      let bucket_fun' :: Lambda GPU
bucket_fun' = Lambda SOACS -> Lambda GPU
soacsLambdaToGPU Lambda SOACS
bucket_fun

      (hist_res, stms) <- Builder GPU [VName] -> FlattenM ([VName], Stms GPU)
forall (m :: * -> *) somerep rep a.
(MonadFreshNames m, HasScope somerep m, SameScope somerep rep) =>
Builder rep a -> m (a, Stms rep)
runBuilder (SegLevel
-> Name
-> Segments
-> [HistOp SOACS]
-> Lambda GPU
-> [VName]
-> (Segments -> BuilderT GPU (State VNameSource) ())
-> Builder GPU [VName]
forall (m :: * -> *).
(MonadBuilder m, Rep m ~ GPU) =>
SegLevel
-> Name
-> Segments
-> [HistOp SOACS]
-> Lambda GPU
-> [VName]
-> (Segments -> m ())
-> m [VName]
genUniformSegHist SegLevel
lvl Name
"Uniform_segHist" (SubExp -> Segments
forall a. a -> [a]
forall (f :: * -> *) a. Applicative f => a -> f a
pure SubExp
w) [HistOp SOACS]
ops Lambda GPU
bucket_fun' [VName]
arrs (BuilderT GPU (State VNameSource) ()
-> Segments -> BuilderT GPU (State VNameSource) ()
forall a b. a -> b -> a
const (BuilderT GPU (State VNameSource) ()
 -> Segments -> BuilderT GPU (State VNameSource) ())
-> BuilderT GPU (State VNameSource) ()
-> Segments
-> BuilderT GPU (State VNameSource) ()
forall a b. (a -> b) -> a -> b
$ () -> BuilderT GPU (State VNameSource) ()
forall a. a -> BuilderT GPU (State VNameSource) a
forall (f :: * -> *) a. Applicative f => a -> f a
pure ()))
      certifying (stmAuxCerts aux) $ do
        addStms stms
        zipWithM_
          ( \PatElem Type
pe VName
v ->
              [VName] -> Exp (Rep FlattenM) -> FlattenM ()
forall (m :: * -> *).
MonadBuilder m =>
[VName] -> Exp (Rep m) -> m ()
letBindNames [PatElem Type -> VName
forall dec. PatElem dec -> VName
patElemName PatElem Type
pe] (Exp (Rep FlattenM) -> FlattenM ())
-> Exp (Rep FlattenM) -> FlattenM ()
forall a b. (a -> b) -> a -> b
$
                BasicOp -> Exp (Rep FlattenM)
forall rep. BasicOp -> Exp rep
BasicOp (BasicOp -> Exp (Rep FlattenM)) -> BasicOp -> Exp (Rep FlattenM)
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
v
          )
          (patElems pat)
          hist_res
    Op (Stream SubExp
w [VName]
arrs Segments
accs Lambda SOACS
lam) -> do
      types <- (Scope GPU -> Scope SOACS) -> FlattenM (Scope SOACS)
forall a. (Scope GPU -> a) -> FlattenM a
forall rep (m :: * -> *) a.
HasScope rep m =>
(Scope rep -> a) -> m a
asksScope Scope GPU -> Scope SOACS
forall fromrep torep.
SameScope fromrep torep =>
Scope fromrep -> Scope torep
castScope
      ((), stream_stms) <-
        runBuilderT (sequentialStreamWholeArray pat w accs lam arrs) types
      intrablockStms map_in_block stream_stms
    Exp SOACS
_ ->
      Stm (Rep FlattenM) -> FlattenM ()
forall (m :: * -> *). MonadBuilder m => Stm (Rep m) -> m ()
addStm (Stm (Rep FlattenM) -> FlattenM ())
-> Stm (Rep FlattenM) -> FlattenM ()
forall a b. (a -> b) -> a -> b
$ Stm SOACS -> Stm GPU
soacsStmToGPU Stm SOACS
stm

intrablockStms :: InBlockMapTransformer -> Stms SOACS -> FlattenM ()
intrablockStms :: InBlockMapTransformer -> Stms SOACS -> FlattenM ()
intrablockStms InBlockMapTransformer
map_in_block = (Stm SOACS -> FlattenM ()) -> Stms SOACS -> FlattenM ()
forall (t :: * -> *) (m :: * -> *) a b.
(Foldable t, Monad m) =>
(a -> m b) -> t a -> m ()
mapM_ ((Stm SOACS -> FlattenM ()) -> Stms SOACS -> FlattenM ())
-> (Stm SOACS -> FlattenM ()) -> Stms SOACS -> FlattenM ()
forall a b. (a -> b) -> a -> b
$ InBlockMapTransformer -> Stm SOACS -> FlattenM ()
intrablockStm InBlockMapTransformer
map_in_block

intrablockParalleliseBody ::
  InBlockMapTransformer ->
  Body SOACS ->
  FlattenM (KernelBody GPU)
intrablockParalleliseBody :: InBlockMapTransformer -> Body SOACS -> FlattenM (KernelBody GPU)
intrablockParalleliseBody InBlockMapTransformer
map_in_block Body SOACS
body = do
  kstms <- FlattenM () -> FlattenM (Stms (Rep FlattenM))
forall (m :: * -> *) a. MonadBuilder m => m a -> m (Stms (Rep m))
collectStms_ (FlattenM () -> FlattenM (Stms (Rep FlattenM)))
-> FlattenM () -> FlattenM (Stms (Rep FlattenM))
forall a b. (a -> b) -> a -> b
$ InBlockMapTransformer -> Stms SOACS -> FlattenM ()
intrablockStms InBlockMapTransformer
map_in_block (Stms SOACS -> FlattenM ()) -> Stms SOACS -> FlattenM ()
forall a b. (a -> b) -> a -> b
$ Body SOACS -> Stms SOACS
forall rep res. GBody rep res -> Stms rep
bodyStms Body SOACS
body
  pure $ Body () kstms $ map ret $ bodyResult body
  where
    ret :: SubExpRes -> KernelResult
ret (SubExpRes Certs
cs SubExp
se) = ResultManifest -> Certs -> SubExp -> KernelResult
Returns ResultManifest
ResultMaySimplify Certs
cs SubExp
se