{-# LANGUAGE TypeFamilies #-}

module Futhark.Pass.Flatten.Builtins
  ( BuiltinFn (..),
    builtinFunDef,
    determineReduceOp,
    genUniformSegHist,
    mkSegSpace,
    segMap,
    genFlags,
    genScan,
    genFilter,
    genSegScan,
    genSegScanomap,
    genSegScanomapWithPost,
    genNonSegRed,
    genUniformSegScanomapWithPost,
    genUniformSegRed,
    genSegRed,
    genSegRedomap,
    genScatter,
    genScatterND,
    genShapeIota,
    exScanAndSum,
    genExPrefixSum,
    genSegPrefixSum,
    doRepIota,
    doSegIota,
    doPrefixSum,
    doPartition,
  )
where

import Control.Monad (forM, forM_, (<=<))
import Control.Monad.State.Strict
import Data.Foldable (toList)
import Data.Maybe (fromMaybe)
import Futhark.IR.GPU
import Futhark.IR.SOACS as SOACS
import Futhark.MonadFreshNames
import Futhark.Pass.Flatten.Monad
import Futhark.Tools
import Futhark.Transform.Rename (renameBody, renameLambda)
import Futhark.Transform.ToGPU (getSize, soacsLambdaToGPU)
import Futhark.Util (unsnoc)

mkSegSpace :: (MonadFreshNames m) => [(VName, SubExp)] -> m SegSpace
mkSegSpace :: forall (m :: * -> *).
MonadFreshNames m =>
[(VName, SubExp)] -> m SegSpace
mkSegSpace [(VName, SubExp)]
dims = VName -> [(VName, SubExp)] -> SegSpace
SegSpace (VName -> [(VName, SubExp)] -> SegSpace)
-> m VName -> m ([(VName, SubExp)] -> SegSpace)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Name -> m VName
forall (m :: * -> *). MonadFreshNames m => Name -> m VName
newVName Name
"phys_tid" m ([(VName, SubExp)] -> SegSpace)
-> m [(VName, SubExp)] -> m SegSpace
forall a b. m (a -> b) -> m a -> m b
forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> [(VName, SubExp)] -> m [(VName, SubExp)]
forall a. a -> m a
forall (f :: * -> *) a. Applicative f => a -> f a
pure [(VName, SubExp)]
dims

segIotaName, repIotaName, prefixSumName, partitionName :: Name
segIotaName :: Name
segIotaName = Text -> Name
builtinName Text
"segiota"
repIotaName :: Name
repIotaName = Text -> Name
builtinName Text
"repiota"
prefixSumName :: Name
prefixSumName = Text -> Name
builtinName Text
"prefixsum"
partitionName :: Name
partitionName = Text -> Name
builtinName Text
"partition"

inlineBuiltinAtLevel :: SegLevel -> Bool
inlineBuiltinAtLevel :: SegLevel -> Bool
inlineBuiltinAtLevel SegThreadInBlock {} = Bool
True
inlineBuiltinAtLevel SegLevel
_ = Bool
False

topSegLevel :: SegLevel
topSegLevel :: SegLevel
topSegLevel = SegVirt -> Maybe KernelGrid -> SegLevel
SegThread SegVirt
SegVirt Maybe KernelGrid
forall a. Maybe a
Nothing

data ThreadRecommendation = ManyThreads | NoRecommendation SegVirt

numberOfBlocks ::
  (MonadBuilder m, Op (Rep m) ~ HostOp inner (Rep m)) =>
  Name ->
  SubExp ->
  SubExp ->
  m (SubExp, SubExp)
numberOfBlocks :: forall (m :: * -> *) (inner :: * -> *).
(MonadBuilder m, Op (Rep m) ~ HostOp inner (Rep m)) =>
Name -> SubExp -> SubExp -> m (SubExp, SubExp)
numberOfBlocks Name
desc SubExp
w SubExp
tblock_size = do
  max_num_tblocks_key <- Text -> Name
nameFromText (Text -> Name) -> (VName -> Text) -> VName -> Name
forall b c a. (b -> c) -> (a -> b) -> a -> c
. VName -> Text
forall a. Pretty a => a -> Text
prettyText (VName -> Name) -> m VName -> m Name
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Name -> m VName
forall (m :: * -> *). MonadFreshNames m => Name -> m VName
newVName (Name
desc Name -> Name -> Name
forall a. Semigroup a => a -> a -> a
<> Name
"_num_tblocks")
  num_tblocks <-
    letSubExp "num_tblocks" $
      Op $
        SizeOp $
          CalcNumBlocks w max_num_tblocks_key tblock_size
  num_threads <-
    letSubExp "num_threads" $
      BasicOp $
        BinOp (Mul Int64 OverflowUndef) num_tblocks tblock_size
  pure (num_tblocks, num_threads)

-- | Like 'segThread', but cap the thread count to the input size.
-- This is more efficient for small kernels, e.g. summing a small
-- array.
segThreadCapped ::
  (MonadBuilder m, Rep m ~ GPU) =>
  [SubExp] -> Name -> ThreadRecommendation -> m (SegOpLevel (Rep m))
segThreadCapped :: forall (m :: * -> *).
(MonadBuilder m, Rep m ~ GPU) =>
[SubExp] -> Name -> ThreadRecommendation -> m (SegOpLevel (Rep m))
segThreadCapped [SubExp]
ws Name
desc ThreadRecommendation
r = do
  w <-
    Name -> Exp (Rep m) -> m SubExp
forall (m :: * -> *).
MonadBuilder m =>
Name -> Exp (Rep m) -> m SubExp
letSubExp Name
"nest_size"
      (Exp GPU -> m SubExp) -> m (Exp GPU) -> m SubExp
forall (m :: * -> *) a b. Monad m => (a -> m b) -> m a -> m b
=<< BinOp -> SubExp -> [SubExp] -> m (Exp (Rep m))
forall (m :: * -> *).
MonadBuilder m =>
BinOp -> SubExp -> [SubExp] -> m (Exp (Rep m))
foldBinOp (IntType -> Overflow -> BinOp
Mul IntType
Int64 Overflow
OverflowUndef) (IntType -> Integer -> SubExp
intConst IntType
Int64 Integer
1) [SubExp]
ws
  tblock_size <- getSize (desc <> "_tblock_size") SizeThreadBlock

  case r of
    ThreadRecommendation
ManyThreads -> do
      usable_groups <-
        Name -> Exp (Rep m) -> m SubExp
forall (m :: * -> *).
MonadBuilder m =>
Name -> Exp (Rep m) -> m SubExp
letSubExp Name
"segmap_usable_groups"
          (Exp GPU -> m SubExp) -> m (Exp GPU) -> m SubExp
forall (m :: * -> *) a b. Monad m => (a -> m b) -> m a -> m b
=<< BinOp -> m (Exp (Rep m)) -> m (Exp (Rep m)) -> m (Exp (Rep m))
forall (m :: * -> *).
MonadBuilder m =>
BinOp -> m (Exp (Rep m)) -> m (Exp (Rep m)) -> m (Exp (Rep m))
eBinOp
            (IntType -> Safety -> BinOp
SCeilDiv IntType
Int64 Safety
Unsafe)
            (SubExp -> m (Exp (Rep m))
forall (m :: * -> *). MonadBuilder m => SubExp -> m (Exp (Rep m))
eSubExp SubExp
w)
            (SubExp -> m (Exp (Rep m))
SubExp -> m (Exp GPU)
forall (m :: * -> *). MonadBuilder m => SubExp -> m (Exp (Rep m))
eSubExp (SubExp -> m (Exp GPU)) -> m SubExp -> m (Exp GPU)
forall (m :: * -> *) a b. Monad m => (a -> m b) -> m a -> m b
=<< IntType -> SubExp -> m SubExp
forall (m :: * -> *).
MonadBuilder m =>
IntType -> SubExp -> m SubExp
asIntS IntType
Int64 SubExp
tblock_size)
      let grid = Count NumBlocks SubExp -> Count BlockSize SubExp -> KernelGrid
KernelGrid (SubExp -> Count NumBlocks SubExp
forall {k} (u :: k) e. e -> Count u e
Count SubExp
usable_groups) (SubExp -> Count BlockSize SubExp
forall {k} (u :: k) e. e -> Count u e
Count SubExp
tblock_size)
      pure $ SegThread SegNoVirt (Just grid)
    NoRecommendation SegVirt
v -> do
      (num_tblocks, _) <- Name -> SubExp -> SubExp -> m (SubExp, SubExp)
forall (m :: * -> *) (inner :: * -> *).
(MonadBuilder m, Op (Rep m) ~ HostOp inner (Rep m)) =>
Name -> SubExp -> SubExp -> m (SubExp, SubExp)
numberOfBlocks Name
desc SubExp
w SubExp
tblock_size
      let 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)
      pure $ SegThread v (Just grid)

-- FIXME: We use segThreadCapped here because otherwise we may get
-- out-of-bounds writes for SegOps with non-primitive return types.
capThreadSegLevel ::
  (MonadBuilder m, Rep m ~ GPU, Foldable t) =>
  t SubExp -> Name -> SegLevel -> ThreadRecommendation -> m SegLevel
capThreadSegLevel :: forall (m :: * -> *) (t :: * -> *).
(MonadBuilder m, Rep m ~ GPU, Foldable t) =>
t SubExp -> Name -> SegLevel -> ThreadRecommendation -> m SegLevel
capThreadSegLevel t SubExp
segments Name
desc SegLevel
lvl ThreadRecommendation
tr =
  case SegLevel
lvl of
    SegThread {} -> [SubExp] -> Name -> ThreadRecommendation -> m (SegOpLevel (Rep m))
forall (m :: * -> *).
(MonadBuilder m, Rep m ~ GPU) =>
[SubExp] -> Name -> ThreadRecommendation -> m (SegOpLevel (Rep m))
segThreadCapped (t SubExp -> [SubExp]
forall a. t a -> [a]
forall (t :: * -> *) a. Foldable t => t a -> [a]
toList t SubExp
segments) Name
desc ThreadRecommendation
tr
    SegLevel
_ -> SegLevel -> m SegLevel
forall a. a -> m a
forall (f :: * -> *) a. Applicative f => a -> f a
pure SegLevel
lvl

determineReduceOp ::
  (MonadBuilder m) =>
  Lambda SOACS ->
  [SubExp] ->
  m (Lambda SOACS, [SubExp], Shape)
determineReduceOp :: forall (m :: * -> *).
MonadBuilder m =>
Lambda SOACS -> [SubExp] -> m (Lambda SOACS, [SubExp], Shape)
determineReduceOp Lambda SOACS
lam [SubExp]
nes =
  -- We obtain the scalar neutral element by indexing the array-typed
  -- one at [0,...,0]. This is safe even if the array is not literally
  -- a replicate: every lane of a vectorised operator must have a
  -- neutral element, and neutral elements are unique, so all lanes
  -- are forced to hold the same value.
  case (SubExp -> Maybe VName) -> [SubExp] -> Maybe [VName]
forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
(a -> m b) -> t a -> m (t b)
forall (m :: * -> *) a b. Monad m => (a -> m b) -> [a] -> m [b]
mapM SubExp -> Maybe VName
subExpVar [SubExp]
nes of
    Just [VName]
ne_vs' -> do
      let (Shape
shape, Lambda SOACS
lam') = Lambda SOACS -> (Shape, Lambda SOACS)
isVectorMap Lambda SOACS
lam
      nes' <- [VName] -> (VName -> m SubExp) -> m [SubExp]
forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
t a -> (a -> m b) -> m (t b)
forM [VName]
ne_vs' ((VName -> m SubExp) -> m [SubExp])
-> (VName -> m SubExp) -> m [SubExp]
forall a b. (a -> b) -> a -> b
$ \VName
ne_v -> do
        ne_v_t <- VName -> m Type
forall rep (m :: * -> *). HasScope rep m => VName -> m Type
lookupType VName
ne_v
        letSubExp "hist_ne" $
          BasicOp $
            Index ne_v $
              fullSlice ne_v_t $
                replicate (shapeRank shape) $
                  DimFix $
                    intConst Int64 0
      pure (lam', nes', shape)
    Maybe [VName]
Nothing ->
      (Lambda SOACS, [SubExp], Shape)
-> m (Lambda SOACS, [SubExp], Shape)
forall a. a -> m a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (Lambda SOACS
lam, [SubExp]
nes, Shape
forall a. Monoid a => a
mempty)

isVectorMap :: Lambda SOACS -> (Shape, Lambda SOACS)
isVectorMap :: Lambda SOACS -> (Shape, Lambda SOACS)
isVectorMap Lambda SOACS
lam
  | [Let (Pat [PatElem (LetDec SOACS)]
pes) StmAux (ExpDec SOACS)
_ (Op (Screma SubExp
w [VName]
arrs ScremaForm SOACS
form))] <-
      Stms SOACS -> [Stm SOACS]
forall rep. Stms rep -> [Stm rep]
stmsToList (Stms SOACS -> [Stm SOACS]) -> Stms SOACS -> [Stm SOACS]
forall a b. (a -> b) -> a -> b
$ GBody SOACS SubExpRes -> Stms SOACS
forall rep res. GBody rep res -> Stms rep
bodyStms (GBody SOACS SubExpRes -> Stms SOACS)
-> GBody SOACS SubExpRes -> Stms SOACS
forall a b. (a -> b) -> a -> b
$ Lambda SOACS -> GBody SOACS SubExpRes
forall rep t. GLambda rep t -> Body rep
lambdaBody Lambda SOACS
lam,
    (SubExpRes -> SubExp) -> [SubExpRes] -> [SubExp]
forall a b. (a -> b) -> [a] -> [b]
map SubExpRes -> SubExp
resSubExp (GBody SOACS SubExpRes -> [SubExpRes]
forall rep res. GBody rep res -> [res]
bodyResult (Lambda SOACS -> GBody SOACS SubExpRes
forall rep t. GLambda rep t -> Body rep
lambdaBody Lambda SOACS
lam)) [SubExp] -> [SubExp] -> Bool
forall a. Eq a => a -> a -> Bool
== (PatElem Type -> SubExp) -> [PatElem Type] -> [SubExp]
forall a b. (a -> b) -> [a] -> [b]
map (VName -> SubExp
Var (VName -> SubExp)
-> (PatElem Type -> VName) -> PatElem Type -> SubExp
forall b c a. (b -> c) -> (a -> b) -> a -> c
. PatElem Type -> VName
forall dec. PatElem dec -> VName
patElemName) [PatElem Type]
[PatElem (LetDec SOACS)]
pes,
    Just Lambda SOACS
map_lam <- ScremaForm SOACS -> Maybe (Lambda SOACS)
forall rep. ScremaForm rep -> Maybe (Lambda rep)
isMapSOAC ScremaForm SOACS
form,
    [VName]
arrs [VName] -> [VName] -> Bool
forall a. Eq a => a -> a -> Bool
== (Param Type -> VName) -> [Param Type] -> [VName]
forall a b. (a -> b) -> [a] -> [b]
map Param Type -> VName
forall dec. Param dec -> VName
paramName (Lambda SOACS -> [LParam SOACS]
forall rep t. GLambda rep t -> [LParam rep]
lambdaParams Lambda SOACS
lam) =
      let (Shape
shape, Lambda SOACS
lam') = Lambda SOACS -> (Shape, Lambda SOACS)
isVectorMap Lambda SOACS
map_lam
       in ([SubExp] -> Shape
forall d. [d] -> ShapeBase d
Shape [SubExp
w] Shape -> Shape -> Shape
forall a. Semigroup a => a -> a -> a
<> Shape
shape, Lambda SOACS
lam')
  | Bool
otherwise = (Shape
forall a. Monoid a => a
mempty, Lambda SOACS
lam)

segMap ::
  (MonadBuilder m, Rep m ~ GPU, Traversable f) =>
  SegLevel -> f SubExp -> (f SubExp -> m Result) -> m (Exp GPU)
segMap :: forall (m :: * -> *) (f :: * -> *).
(MonadBuilder m, Rep m ~ GPU, Traversable f) =>
SegLevel -> f SubExp -> (f SubExp -> m [SubExpRes]) -> m (Exp GPU)
segMap SegLevel
lvl f SubExp
segments f SubExp -> m [SubExpRes]
f = do
  gtids <- (SubExp -> m VName) -> f SubExp -> m (f VName)
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) -> f a -> f (f b)
traverse (m VName -> SubExp -> m VName
forall a b. a -> b -> a
const (m VName -> SubExp -> m VName) -> m VName -> SubExp -> m VName
forall a b. (a -> b) -> a -> b
$ Name -> m VName
forall (m :: * -> *). MonadFreshNames m => Name -> m VName
newVName Name
"gtid") f SubExp
segments
  space <- mkSegSpace $ zip (toList gtids) (toList segments)
  ((res, ts), stms) <- collectStms $ localScope (scopeOfSegSpace space) $ do
    res <- f $ fmap Var gtids
    ts <- mapM (subExpType . resSubExp) res
    pure (map mkResult res, ts)
  let kbody = BodyDec GPU -> Stms GPU -> [KernelResult] -> GBody GPU KernelResult
forall rep res. BodyDec rep -> Stms rep -> [res] -> GBody rep res
Body () Stms GPU
stms [KernelResult]
res
  let tr = if (Type -> Bool) -> [Type] -> Bool
forall (t :: * -> *) a. Foldable t => (a -> Bool) -> t a -> Bool
all Type -> Bool
forall shape o. TypeBase shape o -> Bool
primType [Type]
ts then ThreadRecommendation
ManyThreads else SegVirt -> ThreadRecommendation
NoRecommendation SegVirt
SegVirt
  lvl' <- capThreadSegLevel segments "segmap" lvl tr
  pure $ Op $ SegOp $ SegMap lvl' space ts kbody
  where
    mkResult :: SubExpRes -> KernelResult
mkResult (SubExpRes Certs
cs SubExp
se) = ResultManifest -> Certs -> SubExp -> KernelResult
Returns ResultManifest
ResultMaySimplify Certs
cs SubExp
se

genScanWithKernelBody ::
  (MonadBuilder m, Rep m ~ GPU, Traversable f) =>
  SegLevel ->
  Name ->
  f SubExp ->
  Lambda GPU ->
  [SubExp] ->
  (f SubExp -> m Result) ->
  m [VName]
genScanWithKernelBody :: forall (m :: * -> *) (f :: * -> *).
(MonadBuilder m, Rep m ~ GPU, Traversable f) =>
SegLevel
-> Name
-> f SubExp
-> Lambda GPU
-> [SubExp]
-> (f SubExp -> m [SubExpRes])
-> m [VName]
genScanWithKernelBody SegLevel
lvl Name
desc f SubExp
segments Lambda GPU
lam [SubExp]
nes =
  SegLevel
-> Name
-> f SubExp
-> (f SubExp -> m (Lambda GPU))
-> Shape
-> [SubExp]
-> (f SubExp -> [Type] -> m (Lambda GPU))
-> (f SubExp -> m [SubExpRes])
-> m [VName]
forall (m :: * -> *) (f :: * -> *).
(MonadBuilder m, Rep m ~ GPU, Traversable f) =>
SegLevel
-> Name
-> f SubExp
-> (f SubExp -> m (Lambda GPU))
-> Shape
-> [SubExp]
-> (f SubExp -> [Type] -> m (Lambda GPU))
-> (f SubExp -> m [SubExpRes])
-> m [VName]
genScanWithKernelBodyAndPost
    SegLevel
lvl
    Name
desc
    f SubExp
segments
    (\f SubExp
_ -> Lambda GPU -> m (Lambda GPU)
forall a. a -> m a
forall (f :: * -> *) a. Applicative f => a -> f a
pure Lambda GPU
lam)
    Shape
forall a. Monoid a => a
mempty
    [SubExp]
nes
    (\f SubExp
_ [Type]
res_t -> [Type] -> m (Lambda GPU)
forall rep (m :: * -> *).
(Buildable rep, MonadFreshNames m) =>
[Type] -> m (Lambda rep)
mkIdentityLambda [Type]
res_t)

-- The SegSpace of a SegRed must always have at least two dimensions, so that
-- its result is an array. We therefore prepend a dummy unit dimension, and
-- afterwards index out the single element of the unit-sized result.
genNonSegRed ::
  (MonadBuilder m, Rep m ~ GPU) =>
  SegLevel ->
  Name ->
  [SubExp] ->
  Reduce GPU ->
  Shape ->
  Lambda GPU ->
  [VName] ->
  m [VName]
genNonSegRed :: forall (m :: * -> *).
(MonadBuilder m, Rep m ~ GPU) =>
SegLevel
-> Name
-> [SubExp]
-> Reduce GPU
-> Shape
-> Lambda GPU
-> [VName]
-> m [VName]
genNonSegRed SegLevel
lvl Name
desc [SubExp]
segments Reduce GPU
red_op Shape
shape Lambda GPU
map_lam [VName]
arrs = do
  let red_lam :: Lambda GPU
red_lam = Reduce GPU -> Lambda GPU
forall rep. Reduce rep -> Lambda rep
redLambda Reduce GPU
red_op
      nes :: [SubExp]
nes = Reduce GPU -> [SubExp]
forall rep. Reduce rep -> [SubExp]
redNeutral Reduce GPU
red_op
      comm :: Commutativity
comm = Reduce GPU -> Commutativity
forall rep. Reduce rep -> Commutativity
redComm Reduce GPU
red_op
  let dummy :: SubExp
dummy = IntType -> Integer -> SubExp
intConst IntType
Int64 Integer
1
  gtids_dummy <- Name -> m VName
forall (m :: * -> *). MonadFreshNames m => Name -> m VName
newVName Name
"dummy"
  gtids_original <- traverse (const $ newVName "gtid") segments
  let gtids = VName
gtids_dummy VName -> [VName] -> [VName]
forall a. a -> [a] -> [a]
: [VName]
gtids_original
  let new_segment = SubExp
dummy SubExp -> [SubExp] -> [SubExp]
forall a. a -> [a] -> [a]
: [SubExp]
segments
  space <- mkSegSpace $ zip (toList gtids) (toList new_segment)
  let gtids' = (VName -> SubExp) -> [VName] -> [SubExp]
forall a b. (a -> b) -> [a] -> [b]
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap VName -> SubExp
Var [VName]
gtids
  (kbody, res_t) <- buildBody . localScope (scopeOfSegSpace space) $ do
    bindLambdaInputArrays (drop 1 gtids') map_lam arrs
    res <- bodyBind (lambdaBody map_lam)
    res_t <- mapM (subExpType . resSubExp) res
    pure (map mkResult res, res_t)
  kbody' <- renameBody kbody
  red_lam' <- renameLambda red_lam
  let op = Commutativity -> Lambda GPU -> [SubExp] -> Shape -> SegBinOp GPU
forall rep.
Commutativity -> Lambda rep -> [SubExp] -> Shape -> SegBinOp rep
SegBinOp Commutativity
comm Lambda GPU
red_lam' [SubExp]
nes Shape
shape
  lvl' <- capThreadSegLevel new_segment "uniform_nonsegred" lvl $ NoRecommendation SegNoVirt
  ress <- letTupExp desc $ Op $ SegOp $ SegRed lvl' space res_t kbody' [op]
  forM ress $ \VName
res_d -> do
    res_dt <- VName -> m Type
forall rep (m :: * -> *). HasScope rep m => VName -> m Type
lookupType VName
res_d
    letExp desc . BasicOp $
      case res_dt of
        Acc {} -> SubExp -> BasicOp
SubExp (SubExp -> BasicOp) -> SubExp -> BasicOp
forall a b. (a -> b) -> a -> b
$ VName -> SubExp
Var VName
res_d
        Type
_ -> VName -> Slice SubExp -> BasicOp
Index VName
res_d (Slice SubExp -> BasicOp) -> Slice SubExp -> BasicOp
forall a b. (a -> b) -> a -> b
$ Type -> [DimIndex SubExp] -> Slice SubExp
fullSlice Type
res_dt [SubExp -> DimIndex SubExp
forall d. d -> DimIndex d
DimFix (SubExp -> DimIndex SubExp) -> SubExp -> DimIndex SubExp
forall a b. (a -> b) -> a -> b
$ IntType -> Integer -> SubExp
intConst IntType
Int64 Integer
0]
  where
    mkResult :: SubExpRes -> KernelResult
mkResult (SubExpRes Certs
cs SubExp
se) = ResultManifest -> Certs -> SubExp -> KernelResult
Returns ResultManifest
ResultMaySimplify Certs
cs SubExp
se

genUniformSegHist ::
  (MonadBuilder m, Rep m ~ GPU) =>
  SegLevel ->
  Name ->
  [SubExp] ->
  [SOACS.HistOp SOACS] ->
  Lambda GPU ->
  [VName] ->
  ([SubExp] -> m ()) ->
  m [VName]
genUniformSegHist :: forall (m :: * -> *).
(MonadBuilder m, Rep m ~ GPU) =>
SegLevel
-> Name
-> [SubExp]
-> [HistOp SOACS]
-> Lambda GPU
-> [VName]
-> ([SubExp] -> m ())
-> m [VName]
genUniformSegHist SegLevel
lvl Name
desc [SubExp]
segments [HistOp SOACS]
ops Lambda GPU
bucket_fun [VName]
arrs [SubExp] -> m ()
readFree = do
  ops' <- [HistOp SOACS]
-> (HistOp SOACS -> m (HistOp GPU)) -> m [HistOp GPU]
forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
t a -> (a -> m b) -> m (t b)
forM [HistOp SOACS]
ops ((HistOp SOACS -> m (HistOp GPU)) -> m [HistOp GPU])
-> (HistOp SOACS -> m (HistOp GPU)) -> m [HistOp GPU]
forall a b. (a -> b) -> a -> b
$ \(SOACS.HistOp Shape
dest_shape SubExp
rf [VName]
dests [SubExp]
nes Lambda SOACS
op) -> do
    (op', nes', shape) <- Lambda SOACS -> [SubExp] -> m (Lambda SOACS, [SubExp], Shape)
forall (m :: * -> *).
MonadBuilder m =>
Lambda SOACS -> [SubExp] -> m (Lambda SOACS, [SubExp], Shape)
determineReduceOp Lambda SOACS
op [SubExp]
nes
    let op'' = Lambda SOACS -> Lambda GPU
soacsLambdaToGPU Lambda SOACS
op'
    pure $ Futhark.IR.GPU.HistOp dest_shape rf dests nes' shape op''
  gtids <- traverse (const $ newVName "gtid") segments
  space <- mkSegSpace $ zip (toList gtids) (toList segments)
  let gtids' = (VName -> SubExp) -> [VName] -> [SubExp]
forall a b. (a -> b) -> [a] -> [b]
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap VName -> SubExp
Var [VName]
gtids
  (kbody, res_t) <- buildBody . localScope (scopeOfSegSpace space) $ do
    readFree gtids'
    bindLambdaInputArrays gtids' bucket_fun arrs
    res <- bodyBind (lambdaBody bucket_fun)
    res_t <- mapM (subExpType . resSubExp) res
    pure (map mkResult res, res_t)
  kbody' <- renameBody kbody
  lvl' <- capThreadSegLevel segments "uniform_seghist" lvl $ NoRecommendation SegNoVirt
  letTupExp desc $ Op $ SegOp $ SegHist lvl' space res_t kbody' ops'
  where
    mkResult :: SubExpRes -> KernelResult
mkResult (SubExpRes Certs
cs SubExp
se) = ResultManifest -> Certs -> SubExp -> KernelResult
Returns ResultManifest
ResultMaySimplify Certs
cs SubExp
se

genUniformSegRed ::
  (MonadBuilder m, Rep m ~ GPU) =>
  SegLevel ->
  Name ->
  [SubExp] ->
  [Reduce GPU] ->
  Shape ->
  Lambda GPU ->
  [VName] ->
  ([SubExp] -> m ()) ->
  m [VName]
genUniformSegRed :: forall (m :: * -> *).
(MonadBuilder m, Rep m ~ GPU) =>
SegLevel
-> Name
-> [SubExp]
-> [Reduce GPU]
-> Shape
-> Lambda GPU
-> [VName]
-> ([SubExp] -> m ())
-> m [VName]
genUniformSegRed SegLevel
lvl Name
desc [SubExp]
segments [Reduce GPU]
red_ops Shape
shape Lambda GPU
map_lam [VName]
arrs [SubExp] -> m ()
readFree = do
  gtids <- (SubExp -> m VName) -> [SubExp] -> m [VName]
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) -> [a] -> f [b]
traverse (m VName -> SubExp -> m VName
forall a b. a -> b -> a
const (m VName -> SubExp -> m VName) -> m VName -> SubExp -> m VName
forall a b. (a -> b) -> a -> b
$ Name -> m VName
forall (m :: * -> *). MonadFreshNames m => Name -> m VName
newVName Name
"gtid") [SubExp]
segments
  space <- mkSegSpace $ zip (toList gtids) (toList segments)
  let gtids' = (VName -> SubExp) -> [VName] -> [SubExp]
forall a b. (a -> b) -> [a] -> [b]
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap VName -> SubExp
Var [VName]
gtids
  (kbody, res_t) <- buildBody . localScope (scopeOfSegSpace space) $ do
    readFree gtids'
    bindLambdaInputArrays gtids' map_lam arrs
    res <- bodyBind (lambdaBody map_lam)
    res_t <- mapM (subExpType . resSubExp) res
    pure (map mkResult res, res_t)

  ops <- forM red_ops $ \Reduce GPU
red_op -> do
    red_lam' <- Lambda GPU -> m (Lambda GPU)
forall rep (m :: * -> *) t.
(Renameable rep, MonadFreshNames m, Rename t) =>
GLambda rep t -> m (GLambda rep t)
renameLambda (Lambda GPU -> m (Lambda GPU)) -> Lambda GPU -> m (Lambda GPU)
forall a b. (a -> b) -> a -> b
$ Reduce GPU -> Lambda GPU
forall rep. Reduce rep -> Lambda rep
redLambda Reduce GPU
red_op
    pure $ SegBinOp (redComm red_op) red_lam' (redNeutral red_op) shape

  kbody' <- renameBody kbody
  lvl' <- capThreadSegLevel segments "uniform_segred" lvl $ NoRecommendation SegNoVirt
  letTupExp desc $ Op $ SegOp $ SegRed lvl' space res_t kbody' ops
  where
    mkResult :: SubExpRes -> KernelResult
mkResult (SubExpRes Certs
cs SubExp
se) = ResultManifest -> Certs -> SubExp -> KernelResult
Returns ResultManifest
ResultMaySimplify Certs
cs SubExp
se

genScanWithKernelBodyAndPost ::
  (MonadBuilder m, Rep m ~ GPU, Traversable f) =>
  SegLevel ->
  Name ->
  f SubExp ->
  (f SubExp -> m (Lambda GPU)) ->
  Shape ->
  [SubExp] ->
  (f SubExp -> [Type] -> m (Lambda GPU)) ->
  (f SubExp -> m Result) ->
  m [VName]
genScanWithKernelBodyAndPost :: forall (m :: * -> *) (f :: * -> *).
(MonadBuilder m, Rep m ~ GPU, Traversable f) =>
SegLevel
-> Name
-> f SubExp
-> (f SubExp -> m (Lambda GPU))
-> Shape
-> [SubExp]
-> (f SubExp -> [Type] -> m (Lambda GPU))
-> (f SubExp -> m [SubExpRes])
-> m [VName]
genScanWithKernelBodyAndPost SegLevel
lvl Name
desc f SubExp
segments f SubExp -> m (Lambda GPU)
mkScanLam Shape
shape [SubExp]
nes f SubExp -> [Type] -> m (Lambda GPU)
mkPostLam f SubExp -> m [SubExpRes]
m = do
  gtids <- (SubExp -> m VName) -> f SubExp -> m (f VName)
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) -> f a -> f (f b)
traverse (m VName -> SubExp -> m VName
forall a b. a -> b -> a
const (m VName -> SubExp -> m VName) -> m VName -> SubExp -> m VName
forall a b. (a -> b) -> a -> b
$ Name -> m VName
forall (m :: * -> *). MonadFreshNames m => Name -> m VName
newVName Name
"gtid") f SubExp
segments
  space <- mkSegSpace $ zip (toList gtids) (toList segments)
  let gtids' = (VName -> SubExp) -> f VName -> f SubExp
forall a b. (a -> b) -> f a -> f b
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap VName -> SubExp
Var f VName
gtids
  (kbody, res_t) <- buildBody . localScope (scopeOfSegSpace space) $ do
    res <- m gtids'
    res_t <- mapM (subExpType . resSubExp) res
    pure (map mkResult res, res_t)

  scan_lam <- mkScanLam gtids'
  post_lam <-
    localScope (scopeOfSegSpace space) $
      mkPostLam gtids' res_t
  -- We have to rename since we are using a global readFree
  scan_lam' <- renameLambda scan_lam
  post_lam' <- renameLambda post_lam

  kbody' <- renameBody kbody
  let op = Commutativity -> Lambda GPU -> [SubExp] -> Shape -> SegBinOp GPU
forall rep.
Commutativity -> Lambda rep -> [SubExp] -> Shape -> SegBinOp rep
SegBinOp Commutativity
Noncommutative Lambda GPU
scan_lam' [SubExp]
nes Shape
shape
  lvl' <- capThreadSegLevel segments "uniform_segscan" lvl $ NoRecommendation SegNoVirt
  letTupExp desc $ Op $ SegOp $ SegScan lvl' space res_t kbody' [op] (SegPostOp post_lam')
  where
    mkResult :: SubExpRes -> KernelResult
mkResult (SubExpRes Certs
cs SubExp
se) = ResultManifest -> Certs -> SubExp -> KernelResult
Returns ResultManifest
ResultMaySimplify Certs
cs SubExp
se

bindLambdaInputArrays ::
  (MonadBuilder m, Traversable f) =>
  f SubExp ->
  Lambda GPU ->
  [VName] ->
  m ()
bindLambdaInputArrays :: forall (m :: * -> *) (f :: * -> *).
(MonadBuilder m, Traversable f) =>
f SubExp -> Lambda GPU -> [VName] -> m ()
bindLambdaInputArrays f SubExp
gtids Lambda GPU
lam [VName]
arrs = do
  let idxs :: [SubExp]
idxs = f SubExp -> [SubExp]
forall a. f a -> [a]
forall (t :: * -> *) a. Foldable t => t a -> [a]
toList f SubExp
gtids
  [(Param Type, VName)] -> ((Param Type, VName) -> m ()) -> m ()
forall (t :: * -> *) (m :: * -> *) a b.
(Foldable t, Monad m) =>
t a -> (a -> m b) -> m ()
forM_ ([Param Type] -> [VName] -> [(Param Type, VName)]
forall a b. [a] -> [b] -> [(a, b)]
zip (Lambda GPU -> [LParam GPU]
forall rep t. GLambda rep t -> [LParam rep]
lambdaParams Lambda GPU
lam) [VName]
arrs) (((Param Type, VName) -> m ()) -> m ())
-> ((Param Type, VName) -> m ()) -> m ()
forall a b. (a -> b) -> a -> b
$ \(Param Type
p, VName
arr) ->
    [VName] -> Exp (Rep m) -> m ()
forall (m :: * -> *).
MonadBuilder m =>
[VName] -> Exp (Rep m) -> m ()
letBindNames [Param Type -> VName
forall dec. Param dec -> VName
paramName Param Type
p]
      (Exp (Rep m) -> m ()) -> m (Exp (Rep m)) -> m ()
forall (m :: * -> *) a b. Monad m => (a -> m b) -> m a -> m b
=<< case Param Type -> Type
forall dec. Typed dec => Param dec -> Type
paramType Param Type
p of
        Acc {} ->
          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
$ VName -> SubExp
Var VName
arr
        Type
_ ->
          VName -> [m (Exp (Rep m))] -> m (Exp (Rep m))
forall (m :: * -> *).
MonadBuilder m =>
VName -> [m (Exp (Rep m))] -> m (Exp (Rep m))
eIndex VName
arr ([m (Exp (Rep m))] -> m (Exp (Rep m)))
-> [m (Exp (Rep m))] -> m (Exp (Rep m))
forall a b. (a -> b) -> a -> b
$ (SubExp -> m (Exp (Rep m))) -> [SubExp] -> [m (Exp (Rep m))]
forall a b. (a -> b) -> [a] -> [b]
map SubExp -> m (Exp (Rep m))
forall (m :: * -> *). MonadBuilder m => SubExp -> m (Exp (Rep m))
eSubExp [SubExp]
idxs

genScan ::
  (Traversable f, MonadBuilder m, Rep m ~ GPU) =>
  SegLevel -> Name -> f SubExp -> Lambda GPU -> [SubExp] -> [VName] -> m [VName]
genScan :: forall (f :: * -> *) (m :: * -> *).
(Traversable f, MonadBuilder m, Rep m ~ GPU) =>
SegLevel
-> Name
-> f SubExp
-> Lambda GPU
-> [SubExp]
-> [VName]
-> m [VName]
genScan SegLevel
lvl Name
desc f SubExp
segments Lambda GPU
lam [SubExp]
nes [VName]
arrs =
  SegLevel
-> Name
-> f SubExp
-> Lambda GPU
-> [SubExp]
-> (f SubExp -> m [SubExpRes])
-> m [VName]
forall (m :: * -> *) (f :: * -> *).
(MonadBuilder m, Rep m ~ GPU, Traversable f) =>
SegLevel
-> Name
-> f SubExp
-> Lambda GPU
-> [SubExp]
-> (f SubExp -> m [SubExpRes])
-> m [VName]
genScanWithKernelBody SegLevel
lvl Name
desc f SubExp
segments Lambda GPU
lam [SubExp]
nes ((f SubExp -> m [SubExpRes]) -> m [VName])
-> (f SubExp -> m [SubExpRes]) -> m [VName]
forall a b. (a -> b) -> a -> b
$ \f SubExp
gtids ->
    ([SubExp] -> [SubExpRes]) -> m [SubExp] -> m [SubExpRes]
forall a b. (a -> b) -> m a -> m b
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap [SubExp] -> [SubExpRes]
subExpsRes (m [SubExp] -> m [SubExpRes])
-> ((VName -> m SubExp) -> m [SubExp])
-> (VName -> m SubExp)
-> m [SubExpRes]
forall b c a. (b -> c) -> (a -> b) -> a -> c
. [VName] -> (VName -> m SubExp) -> m [SubExp]
forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
t a -> (a -> m b) -> m (t b)
forM [VName]
arrs ((VName -> m SubExp) -> m [SubExpRes])
-> (VName -> m SubExp) -> m [SubExpRes]
forall a b. (a -> b) -> a -> b
$ \VName
arr ->
      Name -> Exp (Rep m) -> m SubExp
forall (m :: * -> *).
MonadBuilder m =>
Name -> Exp (Rep m) -> m SubExp
letSubExp (VName -> Name
baseName VName
arr Name -> Name -> Name
forall a. Semigroup a => a -> a -> a
<> Name
"_elem") (Exp GPU -> m SubExp) -> m (Exp GPU) -> m SubExp
forall (m :: * -> *) a b. Monad m => (a -> m b) -> m a -> m b
=<< VName -> [m (Exp (Rep m))] -> m (Exp (Rep m))
forall (m :: * -> *).
MonadBuilder m =>
VName -> [m (Exp (Rep m))] -> m (Exp (Rep m))
eIndex VName
arr (f (m (Exp (Rep m))) -> [m (Exp (Rep m))]
forall a. f a -> [a]
forall (t :: * -> *) a. Foldable t => t a -> [a]
toList (f (m (Exp (Rep m))) -> [m (Exp (Rep m))])
-> f (m (Exp (Rep m))) -> [m (Exp (Rep m))]
forall a b. (a -> b) -> a -> b
$ (SubExp -> m (Exp (Rep m))) -> f SubExp -> f (m (Exp (Rep m)))
forall a b. (a -> b) -> f a -> f b
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap SubExp -> m (Exp (Rep m))
forall (m :: * -> *). MonadBuilder m => SubExp -> m (Exp (Rep m))
eSubExp f SubExp
gtids)

genExScan ::
  (MonadBuilder m, Rep m ~ GPU, Traversable f) =>
  SegLevel -> Name -> f SubExp -> Lambda GPU -> [SubExp] -> [VName] -> m [VName]
genExScan :: forall (m :: * -> *) (f :: * -> *).
(MonadBuilder m, Rep m ~ GPU, Traversable f) =>
SegLevel
-> Name
-> f SubExp
-> Lambda GPU
-> [SubExp]
-> [VName]
-> m [VName]
genExScan SegLevel
lvl Name
desc f SubExp
segments Lambda GPU
lam [SubExp]
nes [VName]
arrs =
  SegLevel
-> Name
-> f SubExp
-> Lambda GPU
-> [SubExp]
-> (f SubExp -> m [SubExpRes])
-> m [VName]
forall (m :: * -> *) (f :: * -> *).
(MonadBuilder m, Rep m ~ GPU, Traversable f) =>
SegLevel
-> Name
-> f SubExp
-> Lambda GPU
-> [SubExp]
-> (f SubExp -> m [SubExpRes])
-> m [VName]
genScanWithKernelBody SegLevel
lvl Name
desc f SubExp
segments Lambda GPU
lam [SubExp]
nes ((f SubExp -> m [SubExpRes]) -> m [VName])
-> (f SubExp -> m [SubExpRes]) -> m [VName]
forall a b. (a -> b) -> a -> b
$ \f SubExp
gtids ->
    let Just ([SubExp]
outerDims, SubExp
innerDim) = [SubExp] -> Maybe ([SubExp], SubExp)
forall a. [a] -> Maybe ([a], a)
unsnoc ([SubExp] -> Maybe ([SubExp], SubExp))
-> [SubExp] -> Maybe ([SubExp], SubExp)
forall a b. (a -> b) -> a -> b
$ f SubExp -> [SubExp]
forall a. f a -> [a]
forall (t :: * -> *) a. Foldable t => t a -> [a]
toList f SubExp
gtids
     in do
          prescan <-
            Name -> Exp (Rep m) -> m [SubExp]
forall (m :: * -> *).
MonadBuilder m =>
Name -> Exp (Rep m) -> m [SubExp]
letTupExp' Name
"to_prescan"
              (Exp GPU -> m [SubExp]) -> m (Exp GPU) -> m [SubExp]
forall (m :: * -> *) a b. Monad m => (a -> m b) -> m a -> m b
=<< m (Exp (Rep m))
-> m (Body (Rep m)) -> m (Body (Rep m)) -> m (Exp (Rep m))
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
                (TPrimExp Bool VName -> m (Exp (Rep m))
forall a (m :: * -> *).
(ToExp a, MonadBuilder m) =>
a -> m (Exp (Rep m))
forall (m :: * -> *).
MonadBuilder m =>
TPrimExp Bool VName -> m (Exp (Rep m))
toExp (TPrimExp Bool VName -> m (Exp (Rep m)))
-> TPrimExp Bool VName -> m (Exp (Rep m))
forall a b. (a -> b) -> a -> b
$ SubExp -> TPrimExp Int64 VName
pe64 SubExp
innerDim TPrimExp Int64 VName -> TPrimExp Int64 VName -> TPrimExp Bool VName
forall {k} v (t :: k).
Eq v =>
TPrimExp t v -> TPrimExp t v -> TPrimExp Bool v
.==. TPrimExp Int64 VName
0)
                ([m (Exp (Rep m))] -> m (Body (Rep m))
forall (m :: * -> *).
MonadBuilder m =>
[m (Exp (Rep m))] -> m (Body (Rep m))
eBody ((SubExp -> m (Exp GPU)) -> [SubExp] -> [m (Exp GPU)]
forall a b. (a -> b) -> [a] -> [b]
map SubExp -> m (Exp (Rep m))
SubExp -> m (Exp GPU)
forall (m :: * -> *). MonadBuilder m => SubExp -> m (Exp (Rep m))
eSubExp [SubExp]
nes))
                ([m (Exp (Rep m))] -> m (Body (Rep m))
forall (m :: * -> *).
MonadBuilder m =>
[m (Exp (Rep m))] -> m (Body (Rep m))
eBody ((VName -> m (Exp GPU)) -> [VName] -> [m (Exp GPU)]
forall a b. (a -> b) -> [a] -> [b]
map (VName -> [m (Exp (Rep m))] -> m (Exp (Rep m))
forall (m :: * -> *).
MonadBuilder m =>
VName -> [m (Exp (Rep m))] -> m (Exp (Rep m))
`eIndex` ((SubExp -> m (Exp (Rep m))) -> [SubExp] -> [m (Exp (Rep m))]
forall a b. (a -> b) -> [a] -> [b]
map SubExp -> m (Exp (Rep m))
forall a (m :: * -> *).
(ToExp a, MonadBuilder m) =>
a -> m (Exp (Rep m))
forall (m :: * -> *). MonadBuilder m => SubExp -> m (Exp (Rep m))
toExp [SubExp]
outerDims [m (Exp (Rep m))] -> [m (Exp (Rep m))] -> [m (Exp (Rep m))]
forall a. [a] -> [a] -> [a]
++ [TPrimExp Int64 VName -> m (Exp (Rep m))
forall a (m :: * -> *).
(ToExp a, MonadBuilder m) =>
a -> m (Exp (Rep m))
forall (m :: * -> *).
MonadBuilder m =>
TPrimExp Int64 VName -> m (Exp (Rep m))
toExp (TPrimExp Int64 VName -> m (Exp (Rep m)))
-> TPrimExp Int64 VName -> m (Exp (Rep m))
forall a b. (a -> b) -> a -> b
$ SubExp -> TPrimExp Int64 VName
pe64 SubExp
innerDim TPrimExp Int64 VName
-> TPrimExp Int64 VName -> TPrimExp Int64 VName
forall a. Num a => a -> a -> a
- TPrimExp Int64 VName
1])) [VName]
arrs))
          pure $ subExpsRes prescan

segScanLambda ::
  (MonadBuilder m, Rep m ~ GPU) =>
  Lambda GPU ->
  ([SubExp] -> m ()) ->
  [SubExp] ->
  m (Lambda GPU)
segScanLambda :: forall (m :: * -> *).
(MonadBuilder m, Rep m ~ GPU) =>
Lambda GPU -> ([SubExp] -> m ()) -> [SubExp] -> m (Lambda GPU)
segScanLambda Lambda GPU
lam [SubExp] -> m ()
_readFree [SubExp]
_gtids = do
  x_flag_p <- Name -> Type -> m (Param Type)
forall (m :: * -> *) dec.
MonadFreshNames m =>
Name -> dec -> m (Param dec)
newParam Name
"x_flag" (Type -> m (Param Type)) -> Type -> m (Param Type)
forall a b. (a -> b) -> a -> b
$ PrimType -> Type
forall shape o. PrimType -> TypeBase shape o
Prim PrimType
Bool
  y_flag_p <- newParam "y_flag" $ Prim Bool
  let ts = Lambda GPU -> [Type]
forall rep t. GLambda rep t -> [t]
lambdaReturnType Lambda GPU
lam
      (xps, yps) = splitAt (length ts) $ lambdaParams lam
  mkLambda ([x_flag_p] ++ xps ++ [y_flag_p] ++ yps) $
    bodyBind
      =<< eBody
        [ eBinOp LogOr (eParam x_flag_p) (eParam y_flag_p),
          eIf
            (eParam y_flag_p)
            (eBody (map eParam yps))
            (pure $ lambdaBody lam)
        ]

genSegScan ::
  (MonadBuilder m, Rep m ~ GPU) =>
  SegLevel -> Name -> Lambda GPU -> [SubExp] -> VName -> [VName] -> m [VName]
genSegScan :: forall (m :: * -> *).
(MonadBuilder m, Rep m ~ GPU) =>
SegLevel
-> Name -> Lambda GPU -> [SubExp] -> VName -> [VName] -> m [VName]
genSegScan SegLevel
lvl Name
desc Lambda GPU
lam [SubExp]
nes VName
flags [VName]
arrs = do
  w <- Int -> Type -> SubExp
forall o. Int -> TypeBase Shape o -> SubExp
arraySize Int
0 (Type -> SubExp) -> m Type -> m SubExp
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> VName -> m Type
forall rep (m :: * -> *). HasScope rep m => VName -> m Type
lookupType VName
flags
  lam' <- segScanLambda lam (const $ pure ()) []
  drop 1 <$> genScan lvl desc [w] lam' (constant False : nes) (flags : arrs)

segScanomapPostLambda ::
  (MonadBuilder m, Rep m ~ GPU) =>
  Lambda GPU ->
  ([SubExp] -> m ()) ->
  [SubExp] ->
  m (Lambda GPU)
segScanomapPostLambda :: forall (m :: * -> *).
(MonadBuilder m, Rep m ~ GPU) =>
Lambda GPU -> ([SubExp] -> m ()) -> [SubExp] -> m (Lambda GPU)
segScanomapPostLambda Lambda GPU
lam [SubExp] -> m ()
readFree [SubExp]
gtids = do
  flag_p <- Name -> Type -> m (Param Type)
forall (m :: * -> *) dec.
MonadFreshNames m =>
Name -> dec -> m (Param dec)
newParam Name
"seg_flag" (Type -> m (Param Type)) -> Type -> m (Param Type)
forall a b. (a -> b) -> a -> b
$ PrimType -> Type
forall shape o. PrimType -> TypeBase shape o
Prim PrimType
Bool
  mkLambda (flag_p : lambdaParams lam) $ do
    readFree gtids
    bodyBind $ lambdaBody lam

genSegScanomap ::
  (MonadBuilder m, Rep m ~ GPU) =>
  SegLevel ->
  Name ->
  Lambda GPU ->
  [SubExp] ->
  VName ->
  Lambda GPU ->
  [VName] ->
  ([SubExp] -> m ()) ->
  m [VName]
genSegScanomap :: forall (m :: * -> *).
(MonadBuilder m, Rep m ~ GPU) =>
SegLevel
-> Name
-> Lambda GPU
-> [SubExp]
-> VName
-> Lambda GPU
-> [VName]
-> ([SubExp] -> m ())
-> m [VName]
genSegScanomap SegLevel
lvl Name
desc Lambda GPU
scan_lam [SubExp]
nes VName
flags Lambda GPU
map_lam [VName]
arrs [SubExp] -> m ()
readFree = do
  post_lam <- [Type] -> m (Lambda GPU)
forall rep (m :: * -> *).
(Buildable rep, MonadFreshNames m) =>
[Type] -> m (Lambda rep)
mkIdentityLambda ([Type] -> m (Lambda GPU)) -> [Type] -> m (Lambda GPU)
forall a b. (a -> b) -> a -> b
$ Lambda GPU -> [Type]
forall rep t. GLambda rep t -> [t]
lambdaReturnType Lambda GPU
map_lam
  genSegScanomapWithPost lvl desc scan_lam nes flags post_lam map_lam arrs readFree

genSegScanomapWithPost ::
  (MonadBuilder m, Rep m ~ GPU) =>
  SegLevel ->
  Name ->
  Lambda GPU ->
  [SubExp] ->
  VName ->
  Lambda GPU ->
  Lambda GPU ->
  [VName] ->
  ([SubExp] -> m ()) ->
  m [VName]
genSegScanomapWithPost :: forall (m :: * -> *).
(MonadBuilder m, Rep m ~ GPU) =>
SegLevel
-> Name
-> Lambda GPU
-> [SubExp]
-> VName
-> Lambda GPU
-> Lambda GPU
-> [VName]
-> ([SubExp] -> m ())
-> m [VName]
genSegScanomapWithPost SegLevel
lvl Name
desc Lambda GPU
scan_lam [SubExp]
nes VName
flags Lambda GPU
post_lam Lambda GPU
map_lam [VName]
arrs [SubExp] -> m ()
readFree = do
  w <- Int -> Type -> SubExp
forall o. Int -> TypeBase Shape o -> SubExp
arraySize Int
0 (Type -> SubExp) -> m Type -> m SubExp
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> VName -> m Type
forall rep (m :: * -> *). HasScope rep m => VName -> m Type
lookupType VName
flags

  genScanWithKernelBodyAndPost
    lvl
    desc
    [w]
    (segScanLambda scan_lam readFree)
    mempty
    (constant False : nes)
    ( \[SubExp]
gtids [Type]
_res_t ->
        Lambda GPU -> ([SubExp] -> m ()) -> [SubExp] -> m (Lambda GPU)
forall (m :: * -> *).
(MonadBuilder m, Rep m ~ GPU) =>
Lambda GPU -> ([SubExp] -> m ()) -> [SubExp] -> m (Lambda GPU)
segScanomapPostLambda Lambda GPU
post_lam [SubExp] -> m ()
readFree [SubExp]
gtids
    )
    ( \[SubExp]
gtids -> do
        let [SubExp
gtid] = [SubExp] -> [SubExp]
forall a. [a] -> [a]
forall (t :: * -> *) a. Foldable t => t a -> [a]
toList [SubExp]
gtids
        flag <- Name -> Exp (Rep m) -> m SubExp
forall (m :: * -> *).
MonadBuilder m =>
Name -> Exp (Rep m) -> m SubExp
letSubExp Name
"flag" (Exp GPU -> m SubExp) -> m (Exp GPU) -> m SubExp
forall (m :: * -> *) a b. Monad m => (a -> m b) -> m a -> m b
=<< VName -> [m (Exp (Rep m))] -> m (Exp (Rep m))
forall (m :: * -> *).
MonadBuilder m =>
VName -> [m (Exp (Rep m))] -> m (Exp (Rep m))
eIndex VName
flags [SubExp -> m (Exp (Rep m))
forall (m :: * -> *). MonadBuilder m => SubExp -> m (Exp (Rep m))
eSubExp SubExp
gtid]
        readFree gtids
        bindLambdaInputArrays gtids map_lam arrs
        map_res <- bodyBind (lambdaBody map_lam)
        pure (subExpRes flag : map_res)
    )

withReadFree ::
  (MonadBuilder m, Rep m ~ GPU) =>
  Lambda GPU ->
  ([SubExp] -> m ()) ->
  [SubExp] ->
  m (Lambda GPU)
withReadFree :: forall (m :: * -> *).
(MonadBuilder m, Rep m ~ GPU) =>
Lambda GPU -> ([SubExp] -> m ()) -> [SubExp] -> m (Lambda GPU)
withReadFree Lambda GPU
lam [SubExp] -> m ()
readFree [SubExp]
gtids =
  [LParam (Rep m)] -> m [SubExpRes] -> m (Lambda (Rep m))
forall (m :: * -> *).
MonadBuilder m =>
[LParam (Rep m)] -> m [SubExpRes] -> m (Lambda (Rep m))
mkLambda (Lambda GPU -> [LParam GPU]
forall rep t. GLambda rep t -> [LParam rep]
lambdaParams Lambda GPU
lam) (m [SubExpRes] -> m (Lambda (Rep m)))
-> m [SubExpRes] -> m (Lambda (Rep m))
forall a b. (a -> b) -> a -> b
$ do
    [SubExp] -> m ()
readFree [SubExp]
gtids
    Body (Rep m) -> m [SubExpRes]
forall (m :: * -> *).
MonadBuilder m =>
Body (Rep m) -> m [SubExpRes]
bodyBind (Body (Rep m) -> m [SubExpRes]) -> Body (Rep m) -> m [SubExpRes]
forall a b. (a -> b) -> a -> b
$ Lambda (Rep m) -> Body (Rep m)
forall rep t. GLambda rep t -> Body rep
lambdaBody Lambda (Rep m)
Lambda GPU
lam

genUniformSegScanomapWithPost ::
  (MonadBuilder m, Rep m ~ GPU) =>
  SegLevel ->
  [SubExp] ->
  Name ->
  Lambda GPU ->
  Shape ->
  [SubExp] ->
  Lambda GPU ->
  Lambda GPU ->
  [VName] ->
  ([SubExp] -> m ()) ->
  m [VName]
genUniformSegScanomapWithPost :: forall (m :: * -> *).
(MonadBuilder m, Rep m ~ GPU) =>
SegLevel
-> [SubExp]
-> Name
-> Lambda GPU
-> Shape
-> [SubExp]
-> Lambda GPU
-> Lambda GPU
-> [VName]
-> ([SubExp] -> m ())
-> m [VName]
genUniformSegScanomapWithPost SegLevel
lvl [SubExp]
segments Name
desc Lambda GPU
scan_lam Shape
shape [SubExp]
nes Lambda GPU
post_lam Lambda GPU
map_lam [VName]
arrs [SubExp] -> m ()
readFree = do
  SegLevel
-> Name
-> [SubExp]
-> ([SubExp] -> m (Lambda GPU))
-> Shape
-> [SubExp]
-> ([SubExp] -> [Type] -> m (Lambda GPU))
-> ([SubExp] -> m [SubExpRes])
-> m [VName]
forall (m :: * -> *) (f :: * -> *).
(MonadBuilder m, Rep m ~ GPU, Traversable f) =>
SegLevel
-> Name
-> f SubExp
-> (f SubExp -> m (Lambda GPU))
-> Shape
-> [SubExp]
-> (f SubExp -> [Type] -> m (Lambda GPU))
-> (f SubExp -> m [SubExpRes])
-> m [VName]
genScanWithKernelBodyAndPost
    SegLevel
lvl
    Name
desc
    [SubExp]
segments
    (m (Lambda GPU) -> [SubExp] -> m (Lambda GPU)
forall a b. a -> b -> a
const (m (Lambda GPU) -> [SubExp] -> m (Lambda GPU))
-> m (Lambda GPU) -> [SubExp] -> m (Lambda GPU)
forall a b. (a -> b) -> a -> b
$ Lambda GPU -> m (Lambda GPU)
forall a. a -> m a
forall (f :: * -> *) a. Applicative f => a -> f a
pure Lambda GPU
scan_lam)
    Shape
shape
    [SubExp]
nes
    (\[SubExp]
gtids [Type]
_res_t -> Lambda GPU -> ([SubExp] -> m ()) -> [SubExp] -> m (Lambda GPU)
forall (m :: * -> *).
(MonadBuilder m, Rep m ~ GPU) =>
Lambda GPU -> ([SubExp] -> m ()) -> [SubExp] -> m (Lambda GPU)
withReadFree Lambda GPU
post_lam [SubExp] -> m ()
readFree [SubExp]
gtids)
    ( \[SubExp]
gtids -> do
        [SubExp] -> m ()
readFree [SubExp]
gtids
        [SubExp] -> Lambda GPU -> [VName] -> m ()
forall (m :: * -> *) (f :: * -> *).
(MonadBuilder m, Traversable f) =>
f SubExp -> Lambda GPU -> [VName] -> m ()
bindLambdaInputArrays [SubExp]
gtids Lambda GPU
map_lam [VName]
arrs
        Body (Rep m) -> m [SubExpRes]
forall (m :: * -> *).
MonadBuilder m =>
Body (Rep m) -> m [SubExpRes]
bodyBind (Lambda GPU -> Body GPU
forall rep t. GLambda rep t -> Body rep
lambdaBody Lambda GPU
map_lam)
    )

genPrefixSum ::
  (MonadBuilder m, Rep m ~ GPU) =>
  SegLevel -> Name -> VName -> m VName
genPrefixSum :: forall (m :: * -> *).
(MonadBuilder m, Rep m ~ GPU) =>
SegLevel -> Name -> VName -> m VName
genPrefixSum SegLevel
lvl Name
desc VName
ns = do
  ws <- Type -> [SubExp]
forall d m. TypeBase (ShapeBase d) m -> [d]
arrayDims (Type -> [SubExp]) -> m Type -> m [SubExp]
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> VName -> m Type
forall rep (m :: * -> *). HasScope rep m => VName -> m Type
lookupType VName
ns
  add_lam <- binOpLambda (Add Int64 OverflowUndef) int64
  head <$> genScan lvl desc ws add_lam [intConst Int64 0] [ns]

genExPrefixSum ::
  (MonadBuilder m, Rep m ~ GPU) =>
  SegLevel -> Name -> VName -> m VName
genExPrefixSum :: forall (m :: * -> *).
(MonadBuilder m, Rep m ~ GPU) =>
SegLevel -> Name -> VName -> m VName
genExPrefixSum SegLevel
lvl Name
desc VName
ns = do
  ws <- Type -> [SubExp]
forall d m. TypeBase (ShapeBase d) m -> [d]
arrayDims (Type -> [SubExp]) -> m Type -> m [SubExp]
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> VName -> m Type
forall rep (m :: * -> *). HasScope rep m => VName -> m Type
lookupType VName
ns
  add_lam <- binOpLambda (Add Int64 OverflowUndef) int64
  head <$> genExScan lvl desc ws add_lam [intConst Int64 0] [ns]

genSegPrefixSum ::
  (MonadBuilder m, Rep m ~ GPU) =>
  SegLevel -> Name -> VName -> VName -> m VName
genSegPrefixSum :: forall (m :: * -> *).
(MonadBuilder m, Rep m ~ GPU) =>
SegLevel -> Name -> VName -> VName -> m VName
genSegPrefixSum SegLevel
lvl Name
desc VName
flags VName
ns = do
  add_lam <- BinOp -> PrimType -> m (Lambda (Rep m))
forall (m :: * -> *).
(MonadBuilder m, Buildable (Rep m)) =>
BinOp -> PrimType -> m (Lambda (Rep m))
binOpLambda (IntType -> Overflow -> BinOp
Add IntType
Int64 Overflow
OverflowUndef) PrimType
int64
  head <$> genSegScan lvl desc add_lam [intConst Int64 0] flags [ns]

-- | Convenience wrapper around 'genScatterND' for one-dimensional
-- destinations.
genScatter ::
  (MonadBuilder m, Rep m ~ GPU) =>
  SegLevel -> VName -> SubExp -> (SubExp -> m (VName, SubExp)) -> m (Exp GPU)
genScatter :: forall (m :: * -> *).
(MonadBuilder m, Rep m ~ GPU) =>
SegLevel
-> VName -> SubExp -> (SubExp -> m (VName, SubExp)) -> m (Exp GPU)
genScatter SegLevel
lvl VName
dest SubExp
n SubExp -> m (VName, SubExp)
f =
  SegLevel
-> VName
-> [SubExp]
-> ([SubExp] -> m ([SubExp], SubExp))
-> m (Exp GPU)
forall (m :: * -> *).
(MonadBuilder m, Rep m ~ GPU) =>
SegLevel
-> VName
-> [SubExp]
-> ([SubExp] -> m ([SubExp], SubExp))
-> m (Exp GPU)
genScatterND SegLevel
lvl VName
dest [SubExp
n] (([SubExp] -> m ([SubExp], SubExp)) -> m (Exp GPU))
-> ([SubExp] -> m ([SubExp], SubExp)) -> m (Exp GPU)
forall a b. (a -> b) -> a -> b
$ \ ~[SubExp
gtid] -> do
    (i, v) <- SubExp -> m (VName, SubExp)
f SubExp
gtid
    pure ([Var i], v)

genScatterND ::
  (MonadBuilder m, Rep m ~ GPU) =>
  SegLevel -> VName -> [SubExp] -> ([SubExp] -> m ([SubExp], SubExp)) -> m (Exp GPU)
genScatterND :: forall (m :: * -> *).
(MonadBuilder m, Rep m ~ GPU) =>
SegLevel
-> VName
-> [SubExp]
-> ([SubExp] -> m ([SubExp], SubExp))
-> m (Exp GPU)
genScatterND SegLevel
lvl VName
dest [SubExp]
grid [SubExp] -> m ([SubExp], SubExp)
f = do
  gtids <- (SubExp -> m VName) -> [SubExp] -> m [VName]
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) -> [a] -> f [b]
traverse (m VName -> SubExp -> m VName
forall a b. a -> b -> a
const (m VName -> SubExp -> m VName) -> m VName -> SubExp -> m VName
forall a b. (a -> b) -> a -> b
$ Name -> m VName
forall (m :: * -> *). MonadFreshNames m => Name -> m VName
newVName Name
"gtid") [SubExp]
grid
  space <- mkSegSpace $ zip gtids grid
  dest_t <- lookupType dest
  let accRank = Type -> Int
forall shape o. ArrayShape shape => TypeBase shape o -> Int
arrayRank Type
dest_t
  withAcc [dest] accRank $ \ ~[VName
acc] -> do
    kbody <- m [KernelResult] -> m (GBody (Rep m) KernelResult)
forall (m :: * -> *) res.
(MonadBuilder m, IsResult res) =>
m [res] -> m (GBody (Rep m) res)
buildBody_ (m [KernelResult] -> m (GBody (Rep m) KernelResult))
-> m [KernelResult] -> m (GBody (Rep m) KernelResult)
forall a b. (a -> b) -> a -> b
$ Scope GPU -> m [KernelResult] -> m [KernelResult]
forall a. Scope GPU -> m a -> m a
forall rep (m :: * -> *) a.
LocalScope rep m =>
Scope rep -> m a -> m a
localScope (SegSpace -> Scope GPU
forall rep. SegSpace -> Scope rep
scopeOfSegSpace SegSpace
space) (m [KernelResult] -> m [KernelResult])
-> m [KernelResult] -> m [KernelResult]
forall a b. (a -> b) -> a -> b
$ do
      (idxs, v) <- [SubExp] -> m ([SubExp], SubExp)
f ([SubExp] -> m ([SubExp], SubExp))
-> [SubExp] -> m ([SubExp], SubExp)
forall a b. (a -> b) -> a -> b
$ (VName -> SubExp) -> [VName] -> [SubExp]
forall a b. (a -> b) -> [a] -> [b]
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap VName -> SubExp
Var [VName]
gtids
      acc' <-
        letExp (baseName acc) $
          BasicOp $
            UpdateAcc Safe acc idxs [v]
      pure [Returns ResultMaySimplify mempty $ Var acc']
    acc_t <- lookupType acc
    lvl' <- capThreadSegLevel grid "genScatterND" lvl $ NoRecommendation SegVirt
    letTupExp' "scatter" $ Op $ SegOp $ SegMap lvl' space [acc_t] kbody

genTabulate ::
  (MonadBuilder m, Rep m ~ GPU) =>
  SegLevel -> SubExp -> (SubExp -> m [SubExp]) -> m (Exp GPU)
genTabulate :: forall (m :: * -> *).
(MonadBuilder m, Rep m ~ GPU) =>
SegLevel -> SubExp -> (SubExp -> m [SubExp]) -> m (Exp GPU)
genTabulate SegLevel
lvl SubExp
w SubExp -> m [SubExp]
m = do
  gtid <- Name -> m VName
forall (m :: * -> *). MonadFreshNames m => Name -> m VName
newVName Name
"gtid"
  space <- mkSegSpace [(gtid, w)]
  ((res, ts), stms) <- collectStms $ localScope (scopeOfSegSpace space) $ do
    ses <- m $ Var gtid
    ts <- mapM subExpType ses
    pure (map (Returns ResultMaySimplify mempty) ses, ts)
  let kbody = BodyDec GPU -> Stms GPU -> [KernelResult] -> GBody GPU KernelResult
forall rep res. BodyDec rep -> Stms rep -> [res] -> GBody rep res
Body () Stms GPU
stms [KernelResult]
res
  lvl' <- capThreadSegLevel [w] "genTabulate" lvl $ NoRecommendation SegVirt
  pure $ Op $ SegOp $ SegMap lvl' space ts kbody

genFlags ::
  (MonadBuilder m, Rep m ~ GPU) =>
  SegLevel -> SubExp -> VName -> m VName
genFlags :: forall (m :: * -> *).
(MonadBuilder m, Rep m ~ GPU) =>
SegLevel -> SubExp -> VName -> m VName
genFlags SegLevel
lvl SubExp
m VName
offsets = do
  flags_allfalse <-
    Name -> Exp (Rep m) -> m VName
forall (m :: * -> *).
MonadBuilder m =>
Name -> Exp (Rep m) -> m VName
letExp Name
"flags_allfalse" (Exp GPU -> m VName) -> (BasicOp -> Exp GPU) -> BasicOp -> m VName
forall b c a. (b -> c) -> (a -> b) -> a -> c
. BasicOp -> Exp GPU
forall rep. BasicOp -> Exp rep
BasicOp (BasicOp -> m VName) -> BasicOp -> m VName
forall a b. (a -> b) -> a -> b
$
      Shape -> SubExp -> BasicOp
Replicate ([SubExp] -> Shape
forall d. [d] -> ShapeBase d
Shape [SubExp
m]) (Bool -> SubExp
forall v. IsValue v => v -> SubExp
constant Bool
False)
  n <- arraySize 0 <$> lookupType offsets
  letExp "flags" <=< genScatter lvl flags_allfalse n $ \SubExp
gtid -> do
    i <- Name -> Exp (Rep m) -> m VName
forall (m :: * -> *).
MonadBuilder m =>
Name -> Exp (Rep m) -> m VName
letExp Name
"i" (Exp GPU -> m VName) -> m (Exp GPU) -> m VName
forall (m :: * -> *) a b. Monad m => (a -> m b) -> m a -> m b
=<< VName -> [m (Exp (Rep m))] -> m (Exp (Rep m))
forall (m :: * -> *).
MonadBuilder m =>
VName -> [m (Exp (Rep m))] -> m (Exp (Rep m))
eIndex VName
offsets [SubExp -> m (Exp (Rep m))
forall (m :: * -> *). MonadBuilder m => SubExp -> m (Exp (Rep m))
eSubExp SubExp
gtid]
    pure (i, constant True)

genSegRed ::
  (MonadBuilder m, Rep m ~ GPU) =>
  SegLevel -> VName -> VName -> VName -> [VName] -> Reduce SOACS -> m [VName]
genSegRed :: forall (m :: * -> *).
(MonadBuilder m, Rep m ~ GPU) =>
SegLevel
-> VName -> VName -> VName -> [VName] -> Reduce SOACS -> m [VName]
genSegRed SegLevel
lvl VName
segments VName
flags VName
offsets [VName]
elems Reduce SOACS
red = do
  scanned <-
    SegLevel
-> Name -> Lambda GPU -> [SubExp] -> VName -> [VName] -> m [VName]
forall (m :: * -> *).
(MonadBuilder m, Rep m ~ GPU) =>
SegLevel
-> Name -> Lambda GPU -> [SubExp] -> VName -> [VName] -> m [VName]
genSegScan
      SegLevel
lvl
      Name
"red"
      (Lambda SOACS -> Lambda GPU
soacsLambdaToGPU (Lambda SOACS -> Lambda GPU) -> Lambda SOACS -> Lambda GPU
forall a b. (a -> b) -> a -> b
$ Reduce SOACS -> Lambda SOACS
forall rep. Reduce rep -> Lambda rep
redLambda Reduce SOACS
red)
      (Reduce SOACS -> [SubExp]
forall rep. Reduce rep -> [SubExp]
redNeutral Reduce SOACS
red)
      VName
flags
      [VName]
elems
  num_segments <- arraySize 0 <$> lookupType offsets
  letTupExp "segred" <=< genTabulate lvl num_segments $ \SubExp
i -> do
    n <- Name -> Exp (Rep m) -> m SubExp
forall (m :: * -> *).
MonadBuilder m =>
Name -> Exp (Rep m) -> m SubExp
letSubExp Name
"n" (Exp GPU -> m SubExp) -> m (Exp GPU) -> m SubExp
forall (m :: * -> *) a b. Monad m => (a -> m b) -> m a -> m b
=<< VName -> [m (Exp (Rep m))] -> m (Exp (Rep m))
forall (m :: * -> *).
MonadBuilder m =>
VName -> [m (Exp (Rep m))] -> m (Exp (Rep m))
eIndex VName
segments [SubExp -> m (Exp (Rep m))
forall (m :: * -> *). MonadBuilder m => SubExp -> m (Exp (Rep m))
eSubExp SubExp
i]
    offset <- letSubExp "offset" =<< eIndex offsets [toExp (pe64 i)]
    letTupExp' "segment_res" <=< eIf (toExp $ pe64 n .==. 0) (eBody $ map eSubExp nes) $
      eBody $
        map (`eIndex` [toExp $ pe64 offset + pe64 n - 1]) scanned
  where
    nes :: [SubExp]
nes = Reduce SOACS -> [SubExp]
forall rep. Reduce rep -> [SubExp]
redNeutral Reduce SOACS
red

genSegRedomap ::
  (MonadBuilder m, Rep m ~ GPU) =>
  SegLevel ->
  VName ->
  VName ->
  VName ->
  [VName] ->
  Reduce SOACS ->
  Lambda GPU ->
  ([SubExp] -> m ()) ->
  m ([VName], [VName])
genSegRedomap :: forall (m :: * -> *).
(MonadBuilder m, Rep m ~ GPU) =>
SegLevel
-> VName
-> VName
-> VName
-> [VName]
-> Reduce SOACS
-> Lambda GPU
-> ([SubExp] -> m ())
-> m ([VName], [VName])
genSegRedomap SegLevel
lvl VName
segments VName
flags VName
offsets [VName]
elems Reduce SOACS
red Lambda GPU
map_lam [SubExp] -> m ()
readFree = do
  scanned_and_map <-
    SegLevel
-> Name
-> Lambda GPU
-> [SubExp]
-> VName
-> Lambda GPU
-> [VName]
-> ([SubExp] -> m ())
-> m [VName]
forall (m :: * -> *).
(MonadBuilder m, Rep m ~ GPU) =>
SegLevel
-> Name
-> Lambda GPU
-> [SubExp]
-> VName
-> Lambda GPU
-> [VName]
-> ([SubExp] -> m ())
-> m [VName]
genSegScanomap
      SegLevel
lvl
      Name
"redomap"
      (Lambda SOACS -> Lambda GPU
soacsLambdaToGPU (Lambda SOACS -> Lambda GPU) -> Lambda SOACS -> Lambda GPU
forall a b. (a -> b) -> a -> b
$ Reduce SOACS -> Lambda SOACS
forall rep. Reduce rep -> Lambda rep
redLambda Reduce SOACS
red)
      (Reduce SOACS -> [SubExp]
forall rep. Reduce rep -> [SubExp]
redNeutral Reduce SOACS
red)
      VName
flags
      Lambda GPU
map_lam
      [VName]
elems
      [SubExp] -> m ()
readFree
  let (scanned, mapout) = splitAt (length nes) scanned_and_map
  num_segments <- arraySize 0 <$> lookupType offsets
  reds <- letTupExp "segred" <=< genTabulate lvl num_segments $ \SubExp
i -> do
    n <- Name -> Exp (Rep m) -> m SubExp
forall (m :: * -> *).
MonadBuilder m =>
Name -> Exp (Rep m) -> m SubExp
letSubExp Name
"n" (Exp GPU -> m SubExp) -> m (Exp GPU) -> m SubExp
forall (m :: * -> *) a b. Monad m => (a -> m b) -> m a -> m b
=<< VName -> [m (Exp (Rep m))] -> m (Exp (Rep m))
forall (m :: * -> *).
MonadBuilder m =>
VName -> [m (Exp (Rep m))] -> m (Exp (Rep m))
eIndex VName
segments [SubExp -> m (Exp (Rep m))
forall (m :: * -> *). MonadBuilder m => SubExp -> m (Exp (Rep m))
eSubExp SubExp
i]
    offset <- letSubExp "offset" =<< eIndex offsets [toExp (pe64 i)]
    letTupExp' "segment_res" <=< eIf (toExp $ pe64 n .==. 0) (eBody $ map eSubExp nes) $
      eBody $
        map (`eIndex` [toExp $ pe64 offset + pe64 n - 1]) scanned
  pure (reds, mapout)
  where
    nes :: [SubExp]
nes = Reduce SOACS -> [SubExp]
forall rep. Reduce rep -> [SubExp]
redNeutral Reduce SOACS
red

-- | Produces a multidimensional iota for the given shape.
genShapeIota ::
  (MonadBuilder m, Rep m ~ GPU) =>
  SegLevel -> Shape -> m VName
genShapeIota :: forall (m :: * -> *).
(MonadBuilder m, Rep m ~ GPU) =>
SegLevel -> Shape -> m VName
genShapeIota SegLevel
lvl Shape
shape = do
  let dims :: [SubExp]
dims = Shape -> [SubExp]
forall d. ShapeBase d -> [d]
shapeDims Shape
shape
  Name -> Exp (Rep m) -> m VName
forall (m :: * -> *).
MonadBuilder m =>
Name -> Exp (Rep m) -> m VName
letExp Name
"shape_iota" (Exp GPU -> m VName)
-> (([SubExp] -> m [SubExpRes]) -> m (Exp GPU))
-> ([SubExp] -> m [SubExpRes])
-> m VName
forall (m :: * -> *) b c a.
Monad m =>
(b -> m c) -> (a -> m b) -> a -> m c
<=< SegLevel -> [SubExp] -> ([SubExp] -> m [SubExpRes]) -> m (Exp GPU)
forall (m :: * -> *) (f :: * -> *).
(MonadBuilder m, Rep m ~ GPU, Traversable f) =>
SegLevel -> f SubExp -> (f SubExp -> m [SubExpRes]) -> m (Exp GPU)
segMap SegLevel
lvl [SubExp]
dims (([SubExp] -> m [SubExpRes]) -> m VName)
-> ([SubExp] -> m [SubExpRes]) -> m VName
forall a b. (a -> b) -> a -> b
$ \[SubExp]
gtids -> do
    i <-
      Name -> TPrimExp Int64 VName -> m SubExp
forall (m :: * -> *) a.
(MonadBuilder m, ToExp a) =>
Name -> a -> m SubExp
toSubExp Name
"shape_iota_elem" (TPrimExp Int64 VName -> m SubExp)
-> TPrimExp Int64 VName -> m SubExp
forall a b. (a -> b) -> a -> b
$
        [TPrimExp Int64 VName]
-> [TPrimExp Int64 VName] -> TPrimExp Int64 VName
forall num. IntegralExp num => [num] -> [num] -> num
flattenIndex ((SubExp -> TPrimExp Int64 VName)
-> [SubExp] -> [TPrimExp Int64 VName]
forall a b. (a -> b) -> [a] -> [b]
map SubExp -> TPrimExp Int64 VName
pe64 [SubExp]
dims) ((SubExp -> TPrimExp Int64 VName)
-> [SubExp] -> [TPrimExp Int64 VName]
forall a b. (a -> b) -> [a] -> [b]
map SubExp -> TPrimExp Int64 VName
pe64 [SubExp]
gtids)
    pure [subExpRes i]

-- Returns (#segments, segment start offsets, sum of segment sizes)
-- Note: If given a multi-dimensional array,
-- `#segments` and `sum of segment sizes` will be arrays, not scalars.
-- `segment start offsets` will always have the same shape as `ks`.
exScanAndSum ::
  (MonadBuilder m, Rep m ~ GPU) =>
  SegLevel -> VName -> m (SubExp, VName, SubExp)
exScanAndSum :: forall (m :: * -> *).
(MonadBuilder m, Rep m ~ GPU) =>
SegLevel -> VName -> m (SubExp, VName, SubExp)
exScanAndSum SegLevel
lvl VName
ks = do
  ns <- Type -> [SubExp]
forall d m. TypeBase (ShapeBase d) m -> [d]
arrayDims (Type -> [SubExp]) -> m Type -> m [SubExp]
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> VName -> m Type
forall rep (m :: * -> *). HasScope rep m => VName -> m Type
lookupType VName
ks
  -- If `ks` only has a single dimension
  -- the size will be a scalar, otherwise it's an array.
  ns' <- letExp "ns" $ BasicOp $ case ns of
    [] -> String -> BasicOp
forall a. HasCallStack => String -> a
error (String -> BasicOp) -> String -> BasicOp
forall a b. (a -> b) -> a -> b
$ String
"exScanAndSum: Given non-array argument: " String -> String -> String
forall a. [a] -> [a] -> [a]
++ VName -> String
forall a. Pretty a => a -> String
prettyString VName
ks
    [SubExp
n] -> SubExp -> BasicOp
SubExp SubExp
n
    [SubExp]
_ -> [SubExp] -> Type -> BasicOp
ArrayLit [SubExp]
ns (PrimType -> Type
forall shape o. PrimType -> TypeBase shape o
Prim PrimType
int64)
  -- Check if the innermost dimension is empty.
  is_empty <-
    letExp "is_empty"
      =<< ( case ns of
              [SubExp
n] -> TPrimExp Bool VName -> m (Exp (Rep m))
forall a (m :: * -> *).
(ToExp a, MonadBuilder m) =>
a -> m (Exp (Rep m))
forall (m :: * -> *).
MonadBuilder m =>
TPrimExp Bool VName -> m (Exp (Rep m))
toExp (SubExp -> TPrimExp Int64 VName
pe64 SubExp
n TPrimExp Int64 VName -> TPrimExp Int64 VName -> TPrimExp Bool VName
forall {k} v (t :: k).
Eq v =>
TPrimExp t v -> TPrimExp t v -> TPrimExp Bool v
.==. TPrimExp Int64 VName
0)
              [SubExp]
_ -> VName -> m (Exp (Rep m))
forall (m :: * -> *). MonadBuilder m => VName -> m (Exp (Rep m))
eLast VName
ns' m (Exp GPU) -> (Exp GPU -> m SubExp) -> m SubExp
forall a b. m a -> (a -> m b) -> m b
forall (m :: * -> *) a b. Monad m => m a -> (a -> m b) -> m b
>>= Name -> Exp (Rep m) -> m SubExp
forall (m :: * -> *).
MonadBuilder m =>
Name -> Exp (Rep m) -> m SubExp
letSubExp Name
"n" m SubExp -> (SubExp -> m (Exp (Rep m))) -> m (Exp (Rep m))
forall a b. m a -> (a -> m b) -> m b
forall (m :: * -> *) a b. Monad m => m a -> (a -> m b) -> m b
>>= (\SubExp
n -> TPrimExp Bool VName -> m (Exp (Rep m))
forall a (m :: * -> *).
(ToExp a, MonadBuilder m) =>
a -> m (Exp (Rep m))
forall (m :: * -> *).
MonadBuilder m =>
TPrimExp Bool VName -> m (Exp (Rep m))
toExp (TPrimExp Bool VName -> m (Exp (Rep m)))
-> TPrimExp Bool VName -> m (Exp (Rep m))
forall a b. (a -> b) -> a -> b
$ SubExp -> TPrimExp Int64 VName
pe64 SubExp
n TPrimExp Int64 VName -> TPrimExp Int64 VName -> TPrimExp Bool VName
forall {k} v (t :: k).
Eq v =>
TPrimExp t v -> TPrimExp t v -> TPrimExp Bool v
.==. TPrimExp Int64 VName
0)
          )
  offsets <- letExp "offsets" =<< toExp =<< genExPrefixSum lvl "offsets" ks
  ms <- letExp "ms" <=< segMap lvl (init ns) $ \[SubExp]
gtids -> do
    let idxs :: [m (Exp (Rep m))]
idxs = (SubExp -> m (Exp (Rep m))) -> [SubExp] -> [m (Exp (Rep m))]
forall a b. (a -> b) -> [a] -> [b]
map SubExp -> m (Exp (Rep m))
forall a (m :: * -> *).
(ToExp a, MonadBuilder m) =>
a -> m (Exp (Rep m))
forall (m :: * -> *). MonadBuilder m => SubExp -> m (Exp (Rep m))
toExp [SubExp]
gtids
    offset <- Name -> Exp (Rep m) -> m VName
forall (m :: * -> *).
MonadBuilder m =>
Name -> Exp (Rep m) -> m VName
letExp Name
"offset" (Exp GPU -> m VName) -> m (Exp GPU) -> m VName
forall (m :: * -> *) a b. Monad m => (a -> m b) -> m a -> m b
=<< VName -> [m (Exp (Rep m))] -> m (Exp (Rep m))
forall (m :: * -> *).
MonadBuilder m =>
VName -> [m (Exp (Rep m))] -> m (Exp (Rep m))
eIndex VName
offsets [m (Exp (Rep m))]
idxs
    k <- letExp "k" =<< eIndex ks idxs
    m <-
      letSubExp "m"
        =<< eIf
          (toExp is_empty)
          (eBody [eSubExp $ intConst Int64 0])
          -- Add last size because 'offsets' is an *exclusive* prefix
          -- sum.
          (eBody [eBinOp (Add Int64 OverflowUndef) (eLast offset) (eLast k)])
    pure [subExpRes m]
  pure (Var ns', offsets, Var ms)

genSegIota ::
  (MonadBuilder m, Rep m ~ GPU) =>
  SegLevel -> VName -> m (VName, VName, VName)
genSegIota :: forall (m :: * -> *).
(MonadBuilder m, Rep m ~ GPU) =>
SegLevel -> VName -> m (VName, VName, VName)
genSegIota SegLevel
lvl VName
ks = do
  (_n, offsets, m) <- SegLevel -> VName -> m (SubExp, VName, SubExp)
forall (m :: * -> *).
(MonadBuilder m, Rep m ~ GPU) =>
SegLevel -> VName -> m (SubExp, VName, SubExp)
exScanAndSum SegLevel
lvl VName
ks
  flags <- genFlags lvl m offsets
  ones <- letExp "ones" $ BasicOp $ Replicate (Shape [m]) one
  iotas <- genSegPrefixSum lvl "iotas" flags ones
  res <- letExp "res" <=< genTabulate lvl m $ \SubExp
i -> do
    x <- Name -> Exp (Rep m) -> m SubExp
forall (m :: * -> *).
MonadBuilder m =>
Name -> Exp (Rep m) -> m SubExp
letSubExp Name
"x" (Exp GPU -> m SubExp) -> m (Exp GPU) -> m SubExp
forall (m :: * -> *) a b. Monad m => (a -> m b) -> m a -> m b
=<< VName -> [m (Exp (Rep m))] -> m (Exp (Rep m))
forall (m :: * -> *).
MonadBuilder m =>
VName -> [m (Exp (Rep m))] -> m (Exp (Rep m))
eIndex VName
iotas [SubExp -> m (Exp (Rep m))
forall (m :: * -> *). MonadBuilder m => SubExp -> m (Exp (Rep m))
eSubExp SubExp
i]
    letTupExp' "xm1" $ BasicOp $ BinOp (Sub Int64 OverflowUndef) x one
  pure (flags, offsets, res)
  where
    one :: SubExp
one = IntType -> Integer -> SubExp
intConst IntType
Int64 Integer
1

genRepIota ::
  (MonadBuilder m, Rep m ~ GPU) =>
  SegLevel ->
  VName ->
  m (VName, VName, VName)
genRepIota :: forall (m :: * -> *).
(MonadBuilder m, Rep m ~ GPU) =>
SegLevel -> VName -> m (VName, VName, VName)
genRepIota SegLevel
lvl VName
ks = do
  (n, offsets, m) <- SegLevel -> VName -> m (SubExp, VName, SubExp)
forall (m :: * -> *).
(MonadBuilder m, Rep m ~ GPU) =>
SegLevel -> VName -> m (SubExp, VName, SubExp)
exScanAndSum SegLevel
lvl VName
ks
  is <- letExp "is" <=< genTabulate lvl n $ \SubExp
i -> do
    o <- Name -> Exp (Rep m) -> m SubExp
forall (m :: * -> *).
MonadBuilder m =>
Name -> Exp (Rep m) -> m SubExp
letSubExp Name
"o" (Exp GPU -> m SubExp) -> m (Exp GPU) -> m SubExp
forall (m :: * -> *) a b. Monad m => (a -> m b) -> m a -> m b
=<< VName -> [m (Exp (Rep m))] -> m (Exp (Rep m))
forall (m :: * -> *).
MonadBuilder m =>
VName -> [m (Exp (Rep m))] -> m (Exp (Rep m))
eIndex VName
offsets [SubExp -> m (Exp (Rep m))
forall (m :: * -> *). MonadBuilder m => SubExp -> m (Exp (Rep m))
eSubExp SubExp
i]
    k <- letSubExp "n" =<< eIndex ks [eSubExp i]
    letTupExp' "i"
      =<< eIf
        (toExp (pe64 k .==. 0))
        (eBody [eSubExp negone])
        (eBody [toExp $ pe64 o])
  zeroes <- letExp "zeroes" $ BasicOp $ Replicate (Shape [m]) zero
  starts <-
    letExp "starts" <=< genScatter lvl zeroes n $ \SubExp
gtid -> do
      i <- Name -> Exp (Rep m) -> m VName
forall (m :: * -> *).
MonadBuilder m =>
Name -> Exp (Rep m) -> m VName
letExp Name
"i" (Exp GPU -> m VName) -> m (Exp GPU) -> m VName
forall (m :: * -> *) a b. Monad m => (a -> m b) -> m a -> m b
=<< VName -> [m (Exp (Rep m))] -> m (Exp (Rep m))
forall (m :: * -> *).
MonadBuilder m =>
VName -> [m (Exp (Rep m))] -> m (Exp (Rep m))
eIndex VName
is [SubExp -> m (Exp (Rep m))
forall (m :: * -> *). MonadBuilder m => SubExp -> m (Exp (Rep m))
eSubExp SubExp
gtid]
      pure (i, gtid)
  flags <- genFlags lvl m offsets
  res <- genSegPrefixSum lvl "res" flags starts
  pure (flags, offsets, res)
  where
    zero :: SubExp
zero = IntType -> Integer -> SubExp
intConst IntType
Int64 Integer
0
    negone :: SubExp
negone = IntType -> Integer -> SubExp
intConst IntType
Int64 (-Integer
1)

genPartition ::
  (MonadBuilder m, Rep m ~ GPU) =>
  SegLevel -> VName -> VName -> VName -> m (VName, VName, VName)
genPartition :: forall (m :: * -> *).
(MonadBuilder m, Rep m ~ GPU) =>
SegLevel -> VName -> VName -> VName -> m (VName, VName, VName)
genPartition SegLevel
lvl VName
n VName
k VName
cls = do
  let n' :: SubExp
n' = VName -> SubExp
Var VName
n
  let k' :: SubExp
k' = VName -> SubExp
Var VName
k
  let dims :: [SubExp]
dims = [SubExp
k', SubExp
n']
  -- Create a `[k][n]` array of flags such that `cls_flags[i][j]`
  -- is equal 1 if the j'th element is a member of equivalence class `i` i.e.
  -- the `i`th row is a flag array for equivalence class `i`.
  cls_flags <-
    Name -> Exp (Rep m) -> m VName
forall (m :: * -> *).
MonadBuilder m =>
Name -> Exp (Rep m) -> m VName
letExp Name
"flags"
      (Exp GPU -> m VName)
-> (([SubExp] -> m [SubExpRes]) -> m (Exp GPU))
-> ([SubExp] -> m [SubExpRes])
-> m VName
forall (m :: * -> *) b c a.
Monad m =>
(b -> m c) -> (a -> m b) -> a -> m c
<=< SegLevel -> [SubExp] -> ([SubExp] -> m [SubExpRes]) -> m (Exp GPU)
forall (m :: * -> *) (f :: * -> *).
(MonadBuilder m, Rep m ~ GPU, Traversable f) =>
SegLevel -> f SubExp -> (f SubExp -> m [SubExpRes]) -> m (Exp GPU)
segMap SegLevel
lvl [SubExp]
dims
      (([SubExp] -> m [SubExpRes]) -> m VName)
-> ([SubExp] -> m [SubExpRes]) -> m VName
forall a b. (a -> b) -> a -> b
$ \[SubExp
i, SubExp
j] -> do
        c <- Name -> Exp (Rep m) -> m SubExp
forall (m :: * -> *).
MonadBuilder m =>
Name -> Exp (Rep m) -> m SubExp
letSubExp Name
"c" (Exp GPU -> m SubExp) -> m (Exp GPU) -> m SubExp
forall (m :: * -> *) a b. Monad m => (a -> m b) -> m a -> m b
=<< VName -> [m (Exp (Rep m))] -> m (Exp (Rep m))
forall (m :: * -> *).
MonadBuilder m =>
VName -> [m (Exp (Rep m))] -> m (Exp (Rep m))
eIndex VName
cls [SubExp -> m (Exp (Rep m))
forall a (m :: * -> *).
(ToExp a, MonadBuilder m) =>
a -> m (Exp (Rep m))
forall (m :: * -> *). MonadBuilder m => SubExp -> m (Exp (Rep m))
toExp SubExp
j]
        cls_flag <-
          letSubExp "cls_flag"
            =<< eIf
              (toExp $ pe64 i .==. pe64 c)
              (eBody [toExp $ intConst Int64 1])
              (eBody [toExp $ intConst Int64 0])
        pure [subExpRes cls_flag]

  -- Offsets of each of the individual equivalence classes.
  (_, local_offs, _counts) <- exScanAndSum lvl cls_flags
  -- The number of elems in each class
  counts <- letExp "counts" =<< toExp _counts
  -- Offsets of the whole equivalence classes
  global_offs <- genExPrefixSum lvl "global_offs" counts
  -- Offsets over all of the equivalence classes.
  cls_offs <-
    letExp "cls_offs" =<< do
      segMap lvl dims $ \[SubExp
i, SubExp
j] -> do
        global_offset <- Name -> Exp (Rep m) -> m VName
forall (m :: * -> *).
MonadBuilder m =>
Name -> Exp (Rep m) -> m VName
letExp Name
"global_offset" (Exp GPU -> m VName) -> m (Exp GPU) -> m VName
forall (m :: * -> *) a b. Monad m => (a -> m b) -> m a -> m b
=<< VName -> [m (Exp (Rep m))] -> m (Exp (Rep m))
forall (m :: * -> *).
MonadBuilder m =>
VName -> [m (Exp (Rep m))] -> m (Exp (Rep m))
eIndex VName
global_offs [SubExp -> m (Exp (Rep m))
forall a (m :: * -> *).
(ToExp a, MonadBuilder m) =>
a -> m (Exp (Rep m))
forall (m :: * -> *). MonadBuilder m => SubExp -> m (Exp (Rep m))
toExp SubExp
i]
        offset <-
          letSubExp "offset"
            =<< eBinOp
              (Add Int64 OverflowUndef)
              (eIndex local_offs [toExp i, toExp j])
              (toExp global_offset)
        pure [subExpRes offset]

  scratch <- letExp "scratch" $ BasicOp $ Scratch int64 [n']
  res <- letExp "scatter_res" <=< genScatter lvl scratch n' $ \SubExp
gtid -> do
    c <- Name -> Exp (Rep m) -> m VName
forall (m :: * -> *).
MonadBuilder m =>
Name -> Exp (Rep m) -> m VName
letExp Name
"c" (Exp GPU -> m VName) -> m (Exp GPU) -> m VName
forall (m :: * -> *) a b. Monad m => (a -> m b) -> m a -> m b
=<< VName -> [m (Exp (Rep m))] -> m (Exp (Rep m))
forall (m :: * -> *).
MonadBuilder m =>
VName -> [m (Exp (Rep m))] -> m (Exp (Rep m))
eIndex VName
cls [SubExp -> m (Exp (Rep m))
forall a (m :: * -> *).
(ToExp a, MonadBuilder m) =>
a -> m (Exp (Rep m))
forall (m :: * -> *). MonadBuilder m => SubExp -> m (Exp (Rep m))
toExp SubExp
gtid]
    ind <- letExp "ind" =<< eIndex cls_offs [toExp c, toExp gtid]
    i <- letSubExp "i" =<< toExp gtid
    pure (ind, i)
  pure (counts, global_offs, res)

genFilter ::
  (MonadBuilder m, Rep m ~ GPU) =>
  SegLevel -> VName -> m (SubExp, VName)
genFilter :: forall (m :: * -> *).
(MonadBuilder m, Rep m ~ GPU) =>
SegLevel -> VName -> m (SubExp, VName)
genFilter SegLevel
lvl VName
flags = do
  w <- Int -> Type -> SubExp
forall o. Int -> TypeBase Shape o -> SubExp
arraySize Int
0 (Type -> SubExp) -> m Type -> m SubExp
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> VName -> m Type
forall rep (m :: * -> *). HasScope rep m => VName -> m Type
lookupType VName
flags
  flags_int <- letExp "flags_int" <=< segMap lvl [w] $ \[SubExp
i] -> do
    b <- Name -> Exp (Rep m) -> m SubExp
forall (m :: * -> *).
MonadBuilder m =>
Name -> Exp (Rep m) -> m SubExp
letSubExp Name
"b" (Exp GPU -> m SubExp) -> m (Exp GPU) -> m SubExp
forall (m :: * -> *) a b. Monad m => (a -> m b) -> m a -> m b
=<< VName -> [m (Exp (Rep m))] -> m (Exp (Rep m))
forall (m :: * -> *).
MonadBuilder m =>
VName -> [m (Exp (Rep m))] -> m (Exp (Rep m))
eIndex VName
flags [SubExp -> m (Exp (Rep m))
forall (m :: * -> *). MonadBuilder m => SubExp -> m (Exp (Rep m))
eSubExp SubExp
i]
    v <-
      letSubExp "v"
        =<< eIf
          (eSubExp b)
          (eBody [toExp $ intConst Int64 1])
          (eBody [toExp $ intConst Int64 0])
    pure [subExpRes v]
  -- offsets <- genExPrefixSum "filter_offs" flags_int
  (_n, offsets, num_true) <- exScanAndSum lvl flags_int
  -- num_true <- letSubExp "num_true"  =<< eIndex flags_int [toExp $ pe64 w - 1]
  scratch <- letExp "scratch" $ BasicOp $ Scratch int64 [num_true]
  -- is this efficient or do i need to do something smarter? like scatter with guard?
  -- offsets' <- letExp "offset" <=< segMap [w] $ \[i] -> do
  --   b' <- letSubExp "b" =<< eIndex flags [eSubExp i]
  --   v' <-
  --     letSubExp "v'"
  --       =<< eIf
  --         (eSubExp b')
  --         (eBody [eIndex offsets [eSubExp i]] )
  --         (eBody [toExp $ intConst Int64 (-1)])
  --   pure [subExpRes v']

  filtered <- letExp "filtered" <=< genScatter lvl scratch w $ \SubExp
gtid -> do
    b <- Name -> Exp (Rep m) -> m SubExp
forall (m :: * -> *).
MonadBuilder m =>
Name -> Exp (Rep m) -> m SubExp
letSubExp Name
"b" (Exp GPU -> m SubExp) -> m (Exp GPU) -> m SubExp
forall (m :: * -> *) a b. Monad m => (a -> m b) -> m a -> m b
=<< VName -> [m (Exp (Rep m))] -> m (Exp (Rep m))
forall (m :: * -> *).
MonadBuilder m =>
VName -> [m (Exp (Rep m))] -> m (Exp (Rep m))
eIndex VName
flags [SubExp -> m (Exp (Rep m))
forall (m :: * -> *). MonadBuilder m => SubExp -> m (Exp (Rep m))
eSubExp SubExp
gtid]
    -- idx <- letExp "idx" =<< eIndex offsets' [eSubExp gtid]
    idx_se <-
      letSubExp "idx"
        =<< eIf
          (eSubExp b)
          (eBody [eIndex offsets [eSubExp gtid]])
          (eBody [toExp $ intConst Int64 (-1)])
    -- maybe cleaner?
    idx <- letExp "idx" =<< toExp idx_se
    pure (idx, gtid)
  pure (num_true, filtered)

buildingBuiltin :: Builder GPU (FunDef GPU) -> FunDef GPU
buildingBuiltin :: Builder GPU (FunDef GPU) -> FunDef GPU
buildingBuiltin Builder GPU (FunDef GPU)
m = (FunDef GPU, Stms GPU) -> FunDef GPU
forall a b. (a, b) -> a
fst ((FunDef GPU, Stms GPU) -> FunDef GPU)
-> (FunDef GPU, Stms GPU) -> FunDef GPU
forall a b. (a -> b) -> a -> b
$ State VNameSource (FunDef GPU, Stms GPU)
-> VNameSource -> (FunDef GPU, Stms GPU)
forall s a. State s a -> s -> a
evalState (Builder GPU (FunDef GPU)
-> Scope GPU -> State VNameSource (FunDef GPU, Stms GPU)
forall (m :: * -> *) rep a.
MonadFreshNames m =>
BuilderT rep m a -> Scope rep -> m (a, Stms rep)
runBuilderT Builder GPU (FunDef GPU)
m Scope GPU
forall a. Monoid a => a
mempty) VNameSource
blankNameSource

segIotaBuiltin :: FunDef GPU
segIotaBuiltin :: FunDef GPU
segIotaBuiltin = Builder GPU (FunDef GPU) -> FunDef GPU
buildingBuiltin (Builder GPU (FunDef GPU) -> FunDef GPU)
-> Builder GPU (FunDef GPU) -> FunDef GPU
forall a b. (a -> b) -> a -> b
$ do
  np <- Name
-> TypeBase Shape Diet
-> BuilderT
     GPU (StateT VNameSource Identity) (Param (TypeBase Shape Diet))
forall (m :: * -> *) dec.
MonadFreshNames m =>
Name -> dec -> m (Param dec)
newParam Name
"n" (TypeBase Shape Diet
 -> BuilderT
      GPU (StateT VNameSource Identity) (Param (TypeBase Shape Diet)))
-> TypeBase Shape Diet
-> BuilderT
     GPU (StateT VNameSource Identity) (Param (TypeBase Shape Diet))
forall a b. (a -> b) -> a -> b
$ PrimType -> TypeBase Shape Diet
forall shape o. PrimType -> TypeBase shape o
Prim PrimType
int64
  nsp <- newParam "ns" $ Array int64 (Shape [Var (paramName np)]) Observe
  body <-
    localScope (scopeOfFParams [np, nsp]) . buildBody_ $ do
      (flags, offsets, res) <- genSegIota topSegLevel (paramName nsp)
      m <- arraySize 0 <$> lookupType res
      pure $ subExpsRes [m, Var flags, Var offsets, Var res]
  pure
    FunDef
      { funDefEntryPoint = Nothing,
        funDefAttrs = mempty,
        funDefName = segIotaName,
        funDefRetType =
          map
            (,mempty)
            [ Prim int64,
              Array Bool (Shape [Ext 0]) NoMode,
              Array int64 (Shape [Free $ Var $ paramName np]) NoMode,
              Array int64 (Shape [Ext 0]) NoMode
            ],
        funDefParams = [np, nsp],
        funDefBody = body
      }

repIotaBuiltin :: FunDef GPU
repIotaBuiltin :: FunDef GPU
repIotaBuiltin = Builder GPU (FunDef GPU) -> FunDef GPU
buildingBuiltin (Builder GPU (FunDef GPU) -> FunDef GPU)
-> Builder GPU (FunDef GPU) -> FunDef GPU
forall a b. (a -> b) -> a -> b
$ do
  np <- Name
-> TypeBase Shape Diet
-> BuilderT
     GPU (StateT VNameSource Identity) (Param (TypeBase Shape Diet))
forall (m :: * -> *) dec.
MonadFreshNames m =>
Name -> dec -> m (Param dec)
newParam Name
"n" (TypeBase Shape Diet
 -> BuilderT
      GPU (StateT VNameSource Identity) (Param (TypeBase Shape Diet)))
-> TypeBase Shape Diet
-> BuilderT
     GPU (StateT VNameSource Identity) (Param (TypeBase Shape Diet))
forall a b. (a -> b) -> a -> b
$ PrimType -> TypeBase Shape Diet
forall shape o. PrimType -> TypeBase shape o
Prim PrimType
int64
  nsp <- newParam "ns" $ Array int64 (Shape [Var (paramName np)]) Observe
  body <-
    localScope (scopeOfFParams [np, nsp]) . buildBody_ $ do
      (flags, offsets, res) <- genRepIota topSegLevel (paramName nsp)
      m <- arraySize 0 <$> lookupType res
      pure $ subExpsRes [m, Var flags, Var offsets, Var res]
  pure
    FunDef
      { funDefEntryPoint = Nothing,
        funDefAttrs = mempty,
        funDefName = repIotaName,
        funDefRetType =
          map
            (,mempty)
            [ Prim int64,
              Array Bool (Shape [Ext 0]) NoMode,
              Array int64 (Shape [Free $ Var $ paramName np]) NoMode,
              Array int64 (Shape [Ext 0]) NoMode
            ],
        funDefParams = [np, nsp],
        funDefBody = body
      }

prefixSumBuiltin :: FunDef GPU
prefixSumBuiltin :: FunDef GPU
prefixSumBuiltin = Builder GPU (FunDef GPU) -> FunDef GPU
buildingBuiltin (Builder GPU (FunDef GPU) -> FunDef GPU)
-> Builder GPU (FunDef GPU) -> FunDef GPU
forall a b. (a -> b) -> a -> b
$ do
  np <- Name
-> TypeBase Shape Diet
-> BuilderT
     GPU (StateT VNameSource Identity) (Param (TypeBase Shape Diet))
forall (m :: * -> *) dec.
MonadFreshNames m =>
Name -> dec -> m (Param dec)
newParam Name
"n" (TypeBase Shape Diet
 -> BuilderT
      GPU (StateT VNameSource Identity) (Param (TypeBase Shape Diet)))
-> TypeBase Shape Diet
-> BuilderT
     GPU (StateT VNameSource Identity) (Param (TypeBase Shape Diet))
forall a b. (a -> b) -> a -> b
$ PrimType -> TypeBase Shape Diet
forall shape o. PrimType -> TypeBase shape o
Prim PrimType
int64
  nsp <- newParam "ns" $ Array int64 (Shape [Var (paramName np)]) Observe
  body <-
    localScope (scopeOfFParams [np, nsp]) . buildBody_ $
      varsRes . pure <$> genPrefixSum topSegLevel "res" (paramName nsp)
  pure
    FunDef
      { funDefEntryPoint = Nothing,
        funDefAttrs = mempty,
        funDefName = prefixSumName,
        funDefRetType =
          [(Array int64 (Shape [Free $ Var $ paramName np]) NoMode, mempty)],
        funDefParams = [np, nsp],
        funDefBody = body
      }

partitionBuiltin :: FunDef GPU
partitionBuiltin :: FunDef GPU
partitionBuiltin = Builder GPU (FunDef GPU) -> FunDef GPU
buildingBuiltin (Builder GPU (FunDef GPU) -> FunDef GPU)
-> Builder GPU (FunDef GPU) -> FunDef GPU
forall a b. (a -> b) -> a -> b
$ do
  np <- Name
-> TypeBase Shape Diet
-> BuilderT
     GPU (StateT VNameSource Identity) (Param (TypeBase Shape Diet))
forall (m :: * -> *) dec.
MonadFreshNames m =>
Name -> dec -> m (Param dec)
newParam Name
"n" (TypeBase Shape Diet
 -> BuilderT
      GPU (StateT VNameSource Identity) (Param (TypeBase Shape Diet)))
-> TypeBase Shape Diet
-> BuilderT
     GPU (StateT VNameSource Identity) (Param (TypeBase Shape Diet))
forall a b. (a -> b) -> a -> b
$ PrimType -> TypeBase Shape Diet
forall shape o. PrimType -> TypeBase shape o
Prim PrimType
int64
  kp <- newParam "k" $ Prim int64
  csp <- newParam "cs" $ Array int64 (Shape [Var (paramName np)]) Observe
  body <-
    localScope (scopeOfFParams [np, kp, csp]) . buildBody_ $ do
      (counts, offsets, res) <- genPartition topSegLevel (paramName np) (paramName kp) (paramName csp)
      pure $ varsRes [counts, offsets, res]
  pure
    FunDef
      { funDefEntryPoint = Nothing,
        funDefAttrs = mempty,
        funDefName = partitionName,
        funDefRetType =
          map
            (,mempty)
            [ Array int64 (Shape [Free $ Var $ paramName kp]) NoMode,
              Array int64 (Shape [Free $ Var $ paramName kp]) NoMode,
              Array int64 (Shape [Free $ Var $ paramName np]) NoMode
            ],
        funDefParams = [np, kp, csp],
        funDefBody = body
      }

-- | Retrieve the function definition corresponding to a builtin.
builtinFunDef :: BuiltinFn -> FunDef GPU
builtinFunDef :: BuiltinFn -> FunDef GPU
builtinFunDef BuiltinFn
BuiltinSegIota = FunDef GPU
segIotaBuiltin
builtinFunDef BuiltinFn
BuiltinRepIota = FunDef GPU
repIotaBuiltin
builtinFunDef BuiltinFn
BuiltinPrefixSum = FunDef GPU
prefixSumBuiltin
builtinFunDef BuiltinFn
BuiltinPartition = FunDef GPU
partitionBuiltin

-- | @[0,1,2,0,1,0,1,2,3,4,...]@.  Returns @(flags,offsets,elems)@.
doSegIota ::
  SegLevel -> VName -> FlattenM (VName, VName, VName)
doSegIota :: SegLevel -> VName -> FlattenM (VName, VName, VName)
doSegIota SegLevel
lvl VName
ns
  | SegLevel -> Bool
inlineBuiltinAtLevel SegLevel
lvl =
      SegLevel -> VName -> FlattenM (VName, VName, VName)
forall (m :: * -> *).
(MonadBuilder m, Rep m ~ GPU) =>
SegLevel -> VName -> m (VName, VName, VName)
genSegIota SegLevel
lvl VName
ns
  | Bool
otherwise = do
      BuiltinFn -> FlattenM ()
demandBuiltin BuiltinFn
BuiltinSegIota
      ns_t <- VName -> FlattenM Type
forall rep (m :: * -> *). HasScope rep m => VName -> m Type
lookupType VName
ns
      let n = Int -> Type -> SubExp
forall o. Int -> TypeBase Shape o -> SubExp
arraySize Int
0 Type
ns_t
      m <- newVName "m"
      flags <- newVName "segiota_flags"
      offsets <- newVName "segiota_offsets"
      elems <- newVName "segiota_elems"
      let args = [(SubExp
n, PrimType -> Type
forall shape o. PrimType -> TypeBase shape o
Prim PrimType
int64), (VName -> SubExp
Var VName
ns, Type
ns_t)]
          restype =
            [ExtType] -> Maybe [ExtType] -> [ExtType]
forall a. a -> Maybe a -> a
fromMaybe (String -> [ExtType]
forall a. HasCallStack => String -> a
error String
"doSegIota: bad application") (Maybe [ExtType] -> [ExtType]) -> Maybe [ExtType] -> [ExtType]
forall a b. (a -> b) -> a -> b
$
              [ExtType]
-> [Param (TypeBase Shape Diet)]
-> [(SubExp, Type)]
-> Maybe [ExtType]
forall dec.
Typed dec =>
[ExtType] -> [Param dec] -> [(SubExp, Type)] -> Maybe [ExtType]
forall rt dec.
(IsRetType rt, Typed dec) =>
[rt] -> [Param dec] -> [(SubExp, Type)] -> Maybe [rt]
applyRetType
                (((ExtType, RetAls) -> ExtType) -> [(ExtType, RetAls)] -> [ExtType]
forall a b. (a -> b) -> [a] -> [b]
map (ExtType, RetAls) -> ExtType
forall a b. (a, b) -> a
fst ([(ExtType, RetAls)] -> [ExtType])
-> [(ExtType, RetAls)] -> [ExtType]
forall a b. (a -> b) -> a -> b
$ FunDef GPU -> [(RetType GPU, RetAls)]
forall rep. FunDef rep -> [(RetType rep, RetAls)]
funDefRetType FunDef GPU
segIotaBuiltin)
                (FunDef GPU -> [FParam GPU]
forall rep. FunDef rep -> [FParam rep]
funDefParams FunDef GPU
segIotaBuiltin)
                [(SubExp, Type)]
args
      letBindNames [m, flags, offsets, elems] $
        Apply
          (funDefName segIotaBuiltin)
          [(n, Observe), (Var ns, Observe)]
          (map (,mempty) restype)
          Safe
      pure (flags, offsets, elems)

-- | Produces @[0,0,0,1,1,2,2,2,...]@.  Returns @(flags, offsets,
-- elems)@.
doRepIota ::
  SegLevel -> VName -> FlattenM (VName, VName, VName)
doRepIota :: SegLevel -> VName -> FlattenM (VName, VName, VName)
doRepIota SegLevel
lvl VName
ns
  | SegLevel -> Bool
inlineBuiltinAtLevel SegLevel
lvl =
      SegLevel -> VName -> FlattenM (VName, VName, VName)
forall (m :: * -> *).
(MonadBuilder m, Rep m ~ GPU) =>
SegLevel -> VName -> m (VName, VName, VName)
genRepIota SegLevel
lvl VName
ns
  | Bool
otherwise = do
      BuiltinFn -> FlattenM ()
demandBuiltin BuiltinFn
BuiltinRepIota
      ns_t <- VName -> FlattenM Type
forall rep (m :: * -> *). HasScope rep m => VName -> m Type
lookupType VName
ns
      let n = Int -> Type -> SubExp
forall o. Int -> TypeBase Shape o -> SubExp
arraySize Int
0 Type
ns_t
      m <- newVName "m"
      flags <- newVName "repiota_flags"
      offsets <- newVName "repiota_offsets"
      elems <- newVName "repiota_elems"
      let args = [(SubExp
n, PrimType -> Type
forall shape o. PrimType -> TypeBase shape o
Prim PrimType
int64), (VName -> SubExp
Var VName
ns, Type
ns_t)]
          restype =
            [ExtType] -> Maybe [ExtType] -> [ExtType]
forall a. a -> Maybe a -> a
fromMaybe (String -> [ExtType]
forall a. HasCallStack => String -> a
error String
"doRepIota: bad application") (Maybe [ExtType] -> [ExtType]) -> Maybe [ExtType] -> [ExtType]
forall a b. (a -> b) -> a -> b
$
              [ExtType]
-> [Param (TypeBase Shape Diet)]
-> [(SubExp, Type)]
-> Maybe [ExtType]
forall dec.
Typed dec =>
[ExtType] -> [Param dec] -> [(SubExp, Type)] -> Maybe [ExtType]
forall rt dec.
(IsRetType rt, Typed dec) =>
[rt] -> [Param dec] -> [(SubExp, Type)] -> Maybe [rt]
applyRetType
                (((ExtType, RetAls) -> ExtType) -> [(ExtType, RetAls)] -> [ExtType]
forall a b. (a -> b) -> [a] -> [b]
map (ExtType, RetAls) -> ExtType
forall a b. (a, b) -> a
fst ([(ExtType, RetAls)] -> [ExtType])
-> [(ExtType, RetAls)] -> [ExtType]
forall a b. (a -> b) -> a -> b
$ FunDef GPU -> [(RetType GPU, RetAls)]
forall rep. FunDef rep -> [(RetType rep, RetAls)]
funDefRetType FunDef GPU
repIotaBuiltin)
                (FunDef GPU -> [FParam GPU]
forall rep. FunDef rep -> [FParam rep]
funDefParams FunDef GPU
repIotaBuiltin)
                [(SubExp, Type)]
args
      letBindNames [m, flags, offsets, elems] $
        Apply
          (funDefName repIotaBuiltin)
          [(n, Observe), (Var ns, Observe)]
          (map (,mempty) restype)
          Safe
      pure (flags, offsets, elems)

doPrefixSum ::
  SegLevel -> VName -> FlattenM VName
doPrefixSum :: SegLevel -> VName -> FlattenM VName
doPrefixSum SegLevel
lvl VName
ns
  | SegLevel -> Bool
inlineBuiltinAtLevel SegLevel
lvl =
      SegLevel -> Name -> VName -> FlattenM VName
forall (m :: * -> *).
(MonadBuilder m, Rep m ~ GPU) =>
SegLevel -> Name -> VName -> m VName
genPrefixSum SegLevel
lvl Name
"prefix_sum" VName
ns
  | Bool
otherwise = do
      BuiltinFn -> FlattenM ()
demandBuiltin BuiltinFn
BuiltinPrefixSum
      ns_t <- VName -> FlattenM Type
forall rep (m :: * -> *). HasScope rep m => VName -> m Type
lookupType VName
ns
      let n = Int -> Type -> SubExp
forall o. Int -> TypeBase Shape o -> SubExp
arraySize Int
0 Type
ns_t
      letExp "prefix_sum" $
        Apply
          (funDefName prefixSumBuiltin)
          [(n, Observe), (Var ns, Observe)]
          [(staticShapes1 ns_t, mempty)]
          Safe

doPartition ::
  SegLevel -> VName -> VName -> FlattenM (VName, VName, VName)
doPartition :: SegLevel -> VName -> VName -> FlattenM (VName, VName, VName)
doPartition SegLevel
lvl VName
k VName
cs
  | SegLevel -> Bool
inlineBuiltinAtLevel SegLevel
lvl = do
      cs_t <- VName -> FlattenM Type
forall rep (m :: * -> *). HasScope rep m => VName -> m Type
lookupType VName
cs
      n <- letExp "n" $ BasicOp $ SubExp $ arraySize 0 cs_t
      genPartition lvl n k cs
  | Bool
otherwise = do
      BuiltinFn -> FlattenM ()
demandBuiltin BuiltinFn
BuiltinPartition
      cs_t <- VName -> FlattenM Type
forall rep (m :: * -> *). HasScope rep m => VName -> m Type
lookupType VName
cs
      let n = Int -> Type -> SubExp
forall o. Int -> TypeBase Shape o -> SubExp
arraySize Int
0 Type
cs_t
      counts <- newVName "partition_counts"
      offsets <- newVName "partition_offsets"
      res <- newVName "partition_res"
      let args = [(SubExp
n, PrimType -> Type
forall shape o. PrimType -> TypeBase shape o
Prim PrimType
int64), (VName -> SubExp
Var VName
k, PrimType -> Type
forall shape o. PrimType -> TypeBase shape o
Prim PrimType
int64), (VName -> SubExp
Var VName
cs, Type
cs_t)]
          restype =
            [ExtType] -> Maybe [ExtType] -> [ExtType]
forall a. a -> Maybe a -> a
fromMaybe (String -> [ExtType]
forall a. HasCallStack => String -> a
error String
"doPartition: bad application") (Maybe [ExtType] -> [ExtType]) -> Maybe [ExtType] -> [ExtType]
forall a b. (a -> b) -> a -> b
$
              [ExtType]
-> [Param (TypeBase Shape Diet)]
-> [(SubExp, Type)]
-> Maybe [ExtType]
forall dec.
Typed dec =>
[ExtType] -> [Param dec] -> [(SubExp, Type)] -> Maybe [ExtType]
forall rt dec.
(IsRetType rt, Typed dec) =>
[rt] -> [Param dec] -> [(SubExp, Type)] -> Maybe [rt]
applyRetType
                (((ExtType, RetAls) -> ExtType) -> [(ExtType, RetAls)] -> [ExtType]
forall a b. (a -> b) -> [a] -> [b]
map (ExtType, RetAls) -> ExtType
forall a b. (a, b) -> a
fst ([(ExtType, RetAls)] -> [ExtType])
-> [(ExtType, RetAls)] -> [ExtType]
forall a b. (a -> b) -> a -> b
$ FunDef GPU -> [(RetType GPU, RetAls)]
forall rep. FunDef rep -> [(RetType rep, RetAls)]
funDefRetType FunDef GPU
partitionBuiltin)
                (FunDef GPU -> [FParam GPU]
forall rep. FunDef rep -> [FParam rep]
funDefParams FunDef GPU
partitionBuiltin)
                [(SubExp, Type)]
args
      letBindNames [counts, offsets, res] $
        Apply
          (funDefName partitionBuiltin)
          [(n, Observe), (Var k, Observe), (Var cs, Observe)]
          (map (,mempty) restype)
          Safe
      pure (counts, offsets, res)