{-# LANGUAGE TypeFamilies #-}

-- | Code generation for segmented and non-segmented scans.  Uses a
-- fairly inefficient two-pass algorithm, but can handle anything.
module Futhark.CodeGen.ImpGen.GPU.SegScan.TwoPass (compileSegScan) where

import Control.Monad
import Control.Monad.State
import Data.List qualified as L
import Data.Maybe
import Futhark.CodeGen.ImpCode.GPU qualified as Imp
import Futhark.CodeGen.ImpGen
import Futhark.CodeGen.ImpGen.GPU.Base
import Futhark.IR.GPUMem
import Futhark.IR.Mem.LMAD qualified as LMAD
import Futhark.Transform.Rename
import Futhark.Util.IntegralExp (ceilDiv, quot, rem)
import Prelude hiding (quot, rem)

-- Aggressively try to reuse memory for different SegBinOps, because
-- we will run them sequentially after another.
makeLocalArrays ::
  Imp.TExp Int64 ->
  Count BlockSize SubExp ->
  [SegBinOp GPUMem] ->
  InKernelGen [[VName]]
makeLocalArrays :: TPrimExp Int64 VName
-> Count BlockSize SubExp
-> [SegBinOp GPUMem]
-> InKernelGen [[VName]]
makeLocalArrays TPrimExp Int64 VName
tblock_id (Count SubExp
tblock_size) [SegBinOp GPUMem]
scans = do
  (arrs, mems_and_sizes) <- StateT
  [([Count Bytes (TPrimExp Int64 VName)], VName)]
  (ImpM GPUMem KernelEnv KernelOp)
  [[VName]]
-> [([Count Bytes (TPrimExp Int64 VName)], VName)]
-> ImpM
     GPUMem
     KernelEnv
     KernelOp
     ([[VName]], [([Count Bytes (TPrimExp Int64 VName)], VName)])
forall s (m :: * -> *) a. StateT s m a -> s -> m (a, s)
runStateT ((SegBinOp GPUMem
 -> StateT
      [([Count Bytes (TPrimExp Int64 VName)], VName)]
      (ImpM GPUMem KernelEnv KernelOp)
      [VName])
-> [SegBinOp GPUMem]
-> StateT
     [([Count Bytes (TPrimExp Int64 VName)], VName)]
     (ImpM GPUMem KernelEnv KernelOp)
     [[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 SegBinOp GPUMem
-> StateT
     [([Count Bytes (TPrimExp Int64 VName)], VName)]
     (ImpM GPUMem KernelEnv KernelOp)
     [VName]
onScan [SegBinOp GPUMem]
scans) [([Count Bytes (TPrimExp Int64 VName)], VName)]
forall a. Monoid a => a
mempty
  let maxSize [Count u (TPrimExp Int64 v)]
sizes = TPrimExp Int64 v -> Count Bytes (TPrimExp Int64 v)
forall a. a -> Count Bytes a
Imp.bytes (TPrimExp Int64 v -> Count Bytes (TPrimExp Int64 v))
-> TPrimExp Int64 v -> Count Bytes (TPrimExp Int64 v)
forall a b. (a -> b) -> a -> b
$ (TPrimExp Int64 v -> TPrimExp Int64 v -> TPrimExp Int64 v)
-> TPrimExp Int64 v -> [TPrimExp Int64 v] -> TPrimExp Int64 v
forall b a. (b -> a -> b) -> b -> [a] -> b
forall (t :: * -> *) b a.
Foldable t =>
(b -> a -> b) -> b -> t a -> b
L.foldl' TPrimExp Int64 v -> TPrimExp Int64 v -> TPrimExp Int64 v
forall v. TPrimExp Int64 v -> TPrimExp Int64 v -> TPrimExp Int64 v
sMax64 TPrimExp Int64 v
1 ([TPrimExp Int64 v] -> TPrimExp Int64 v)
-> [TPrimExp Int64 v] -> TPrimExp Int64 v
forall a b. (a -> b) -> a -> b
$ (Count u (TPrimExp Int64 v) -> TPrimExp Int64 v)
-> [Count u (TPrimExp Int64 v)] -> [TPrimExp Int64 v]
forall a b. (a -> b) -> [a] -> [b]
map Count u (TPrimExp Int64 v) -> TPrimExp Int64 v
forall {k} (u :: k) e. Count u e -> e
Imp.unCount [Count u (TPrimExp Int64 v)]
sizes
  forM_ mems_and_sizes $ \([Count Bytes (TPrimExp Int64 VName)]
sizes, VName
mem) ->
    VName
-> Count Bytes (TPrimExp Int64 VName)
-> Space
-> ImpM GPUMem KernelEnv KernelOp ()
forall rep r op.
VName
-> Count Bytes (TPrimExp Int64 VName) -> Space -> ImpM rep r op ()
sAlloc_ VName
mem ([Count Bytes (TPrimExp Int64 VName)]
-> Count Bytes (TPrimExp Int64 VName)
forall {k} {v} {u :: k}.
Pretty v =>
[Count u (TPrimExp Int64 v)] -> Count Bytes (TPrimExp Int64 v)
maxSize [Count Bytes (TPrimExp Int64 VName)]
sizes) (SpaceId -> Space
Space SpaceId
"shared")
  pure arrs
  where
    onScan :: SegBinOp GPUMem
-> StateT
     [([Count Bytes (TPrimExp Int64 VName)], VName)]
     (ImpM GPUMem KernelEnv KernelOp)
     [VName]
onScan (SegBinOp Commutativity
_ Lambda GPUMem
scan_op [SubExp]
nes Shape
_) = do
      let ([Param LParamMem]
scan_x_params, [Param LParamMem]
_scan_y_params) =
            Int -> [Param LParamMem] -> ([Param LParamMem], [Param LParamMem])
forall a. Int -> [a] -> ([a], [a])
splitAt ([SubExp] -> Int
forall a. [a] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length [SubExp]
nes) ([Param LParamMem] -> ([Param LParamMem], [Param LParamMem]))
-> [Param LParamMem] -> ([Param LParamMem], [Param LParamMem])
forall a b. (a -> b) -> a -> b
$ Lambda GPUMem -> [LParam GPUMem]
forall rep t. GLambda rep t -> [LParam rep]
lambdaParams Lambda GPUMem
scan_op
      (arrs, used_mems) <- ([(VName, [([Count Bytes (TPrimExp Int64 VName)], VName)])]
 -> ([VName], [[([Count Bytes (TPrimExp Int64 VName)], VName)]]))
-> StateT
     [([Count Bytes (TPrimExp Int64 VName)], VName)]
     (ImpM GPUMem KernelEnv KernelOp)
     [(VName, [([Count Bytes (TPrimExp Int64 VName)], VName)])]
-> StateT
     [([Count Bytes (TPrimExp Int64 VName)], VName)]
     (ImpM GPUMem KernelEnv KernelOp)
     ([VName], [[([Count Bytes (TPrimExp Int64 VName)], VName)]])
forall a b.
(a -> b)
-> StateT
     [([Count Bytes (TPrimExp Int64 VName)], VName)]
     (ImpM GPUMem KernelEnv KernelOp)
     a
-> StateT
     [([Count Bytes (TPrimExp Int64 VName)], VName)]
     (ImpM GPUMem KernelEnv KernelOp)
     b
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap [(VName, [([Count Bytes (TPrimExp Int64 VName)], VName)])]
-> ([VName], [[([Count Bytes (TPrimExp Int64 VName)], VName)]])
forall a b. [(a, b)] -> ([a], [b])
unzip (StateT
   [([Count Bytes (TPrimExp Int64 VName)], VName)]
   (ImpM GPUMem KernelEnv KernelOp)
   [(VName, [([Count Bytes (TPrimExp Int64 VName)], VName)])]
 -> StateT
      [([Count Bytes (TPrimExp Int64 VName)], VName)]
      (ImpM GPUMem KernelEnv KernelOp)
      ([VName], [[([Count Bytes (TPrimExp Int64 VName)], VName)]]))
-> ((Param LParamMem
     -> StateT
          [([Count Bytes (TPrimExp Int64 VName)], VName)]
          (ImpM GPUMem KernelEnv KernelOp)
          (VName, [([Count Bytes (TPrimExp Int64 VName)], VName)]))
    -> StateT
         [([Count Bytes (TPrimExp Int64 VName)], VName)]
         (ImpM GPUMem KernelEnv KernelOp)
         [(VName, [([Count Bytes (TPrimExp Int64 VName)], VName)])])
-> (Param LParamMem
    -> StateT
         [([Count Bytes (TPrimExp Int64 VName)], VName)]
         (ImpM GPUMem KernelEnv KernelOp)
         (VName, [([Count Bytes (TPrimExp Int64 VName)], VName)]))
-> StateT
     [([Count Bytes (TPrimExp Int64 VName)], VName)]
     (ImpM GPUMem KernelEnv KernelOp)
     ([VName], [[([Count Bytes (TPrimExp Int64 VName)], VName)]])
forall b c a. (b -> c) -> (a -> b) -> a -> c
. [Param LParamMem]
-> (Param LParamMem
    -> StateT
         [([Count Bytes (TPrimExp Int64 VName)], VName)]
         (ImpM GPUMem KernelEnv KernelOp)
         (VName, [([Count Bytes (TPrimExp Int64 VName)], VName)]))
-> StateT
     [([Count Bytes (TPrimExp Int64 VName)], VName)]
     (ImpM GPUMem KernelEnv KernelOp)
     [(VName, [([Count Bytes (TPrimExp Int64 VName)], VName)])]
forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
t a -> (a -> m b) -> m (t b)
forM [Param LParamMem]
scan_x_params ((Param LParamMem
  -> StateT
       [([Count Bytes (TPrimExp Int64 VName)], VName)]
       (ImpM GPUMem KernelEnv KernelOp)
       (VName, [([Count Bytes (TPrimExp Int64 VName)], VName)]))
 -> StateT
      [([Count Bytes (TPrimExp Int64 VName)], VName)]
      (ImpM GPUMem KernelEnv KernelOp)
      ([VName], [[([Count Bytes (TPrimExp Int64 VName)], VName)]]))
-> (Param LParamMem
    -> StateT
         [([Count Bytes (TPrimExp Int64 VName)], VName)]
         (ImpM GPUMem KernelEnv KernelOp)
         (VName, [([Count Bytes (TPrimExp Int64 VName)], VName)]))
-> StateT
     [([Count Bytes (TPrimExp Int64 VName)], VName)]
     (ImpM GPUMem KernelEnv KernelOp)
     ([VName], [[([Count Bytes (TPrimExp Int64 VName)], VName)]])
forall a b. (a -> b) -> a -> b
$ \Param LParamMem
p ->
        case Param LParamMem -> LParamMem
forall dec. Param dec -> dec
paramDec Param LParamMem
p of
          MemArray PrimType
pt Shape
shape NoMode
_ (ArrayIn VName
mem LMAD
_) -> do
            let shape' :: Shape
shape' = [SubExp] -> Shape
forall d. [d] -> ShapeBase d
Shape [SubExp
tblock_size] Shape -> Shape -> Shape
forall a. Semigroup a => a -> a -> a
<> Shape
shape
            let shape_E :: [TPrimExp Int64 VName]
shape_E = (SubExp -> TPrimExp Int64 VName)
-> [SubExp] -> [TPrimExp Int64 VName]
forall a b. (a -> b) -> [a] -> [b]
map SubExp -> TPrimExp Int64 VName
pe64 ([SubExp] -> [TPrimExp Int64 VName])
-> [SubExp] -> [TPrimExp Int64 VName]
forall a b. (a -> b) -> a -> b
$ Shape -> [SubExp]
forall d. ShapeBase d -> [d]
shapeDims Shape
shape'
            arr <-
              ImpM GPUMem KernelEnv KernelOp VName
-> StateT
     [([Count Bytes (TPrimExp Int64 VName)], VName)]
     (ImpM GPUMem KernelEnv KernelOp)
     VName
forall (m :: * -> *) a.
Monad m =>
m a -> StateT [([Count Bytes (TPrimExp Int64 VName)], VName)] m a
forall (t :: (* -> *) -> * -> *) (m :: * -> *) a.
(MonadTrans t, Monad m) =>
m a -> t m a
lift (ImpM GPUMem KernelEnv KernelOp VName
 -> StateT
      [([Count Bytes (TPrimExp Int64 VName)], VName)]
      (ImpM GPUMem KernelEnv KernelOp)
      VName)
-> (LMAD -> ImpM GPUMem KernelEnv KernelOp VName)
-> LMAD
-> StateT
     [([Count Bytes (TPrimExp Int64 VName)], VName)]
     (ImpM GPUMem KernelEnv KernelOp)
     VName
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Name
-> PrimType
-> Shape
-> VName
-> LMAD
-> ImpM GPUMem KernelEnv KernelOp VName
forall rep r op.
Name -> PrimType -> Shape -> VName -> LMAD -> ImpM rep r op VName
sArray Name
"scan_arr" PrimType
pt Shape
shape' VName
mem (LMAD
 -> StateT
      [([Count Bytes (TPrimExp Int64 VName)], VName)]
      (ImpM GPUMem KernelEnv KernelOp)
      VName)
-> LMAD
-> StateT
     [([Count Bytes (TPrimExp Int64 VName)], VName)]
     (ImpM GPUMem KernelEnv KernelOp)
     VName
forall a b. (a -> b) -> a -> b
$
                TPrimExp Int64 VName -> [TPrimExp Int64 VName] -> LMAD
forall num. IntegralExp num => num -> [num] -> LMAD num
LMAD.iota (TPrimExp Int64 VName
tblock_id TPrimExp Int64 VName
-> TPrimExp Int64 VName -> TPrimExp Int64 VName
forall a. Num a => a -> a -> a
* [TPrimExp Int64 VName] -> TPrimExp Int64 VName
forall a. Num a => [a] -> a
forall (t :: * -> *) a. (Foldable t, Num a) => t a -> a
product [TPrimExp Int64 VName]
shape_E) [TPrimExp Int64 VName]
shape_E
            pure (arr, [])
          LParamMem
_ -> do
            let pt :: PrimType
pt = Type -> PrimType
forall shape o. TypeBase shape o -> PrimType
elemType (Type -> PrimType) -> Type -> PrimType
forall a b. (a -> b) -> a -> b
$ Param LParamMem -> Type
forall dec. Typed dec => Param dec -> Type
paramType Param LParamMem
p
                shape :: Shape
shape = [SubExp] -> Shape
forall d. [d] -> ShapeBase d
Shape [SubExp
tblock_size]
            (sizes, mem') <- PrimType
-> Shape
-> StateT
     [([Count Bytes (TPrimExp Int64 VName)], VName)]
     (ImpM GPUMem KernelEnv KernelOp)
     ([Count Bytes (TPrimExp Int64 VName)], VName)
forall {t :: (* -> *) -> * -> *} {rep} {r} {op}.
(MonadState
   [([Count Bytes (TPrimExp Int64 VName)], VName)]
   (t (ImpM rep r op)),
 MonadTrans t) =>
PrimType
-> Shape
-> t (ImpM rep r op) ([Count Bytes (TPrimExp Int64 VName)], VName)
getMem PrimType
pt Shape
shape
            arr <- lift $ sArrayInMem "scan_arr" pt shape mem'
            pure (arr, [(sizes, mem')])
      modify (<> concat used_mems)
      pure arrs

    getMem :: PrimType
-> Shape
-> t (ImpM rep r op) ([Count Bytes (TPrimExp Int64 VName)], VName)
getMem PrimType
pt Shape
shape = do
      let size :: Count Bytes (TPrimExp Int64 VName)
size = Type -> Count Bytes (TPrimExp Int64 VName)
typeSize (Type -> Count Bytes (TPrimExp Int64 VName))
-> Type -> Count Bytes (TPrimExp Int64 VName)
forall a b. (a -> b) -> a -> b
$ PrimType -> Shape -> NoMode -> Type
forall shape o. PrimType -> shape -> o -> TypeBase shape o
Array PrimType
pt Shape
shape NoMode
NoMode
      mems <- t (ImpM rep r op) [([Count Bytes (TPrimExp Int64 VName)], VName)]
forall s (m :: * -> *). MonadState s m => m s
get
      case (L.find ((size `elem`) . fst) mems, mems) of
        (Just ([Count Bytes (TPrimExp Int64 VName)], VName)
mem, [([Count Bytes (TPrimExp Int64 VName)], VName)]
_) -> do
          ([([Count Bytes (TPrimExp Int64 VName)], VName)]
 -> [([Count Bytes (TPrimExp Int64 VName)], VName)])
-> t (ImpM rep r op) ()
forall s (m :: * -> *). MonadState s m => (s -> s) -> m ()
modify (([([Count Bytes (TPrimExp Int64 VName)], VName)]
  -> [([Count Bytes (TPrimExp Int64 VName)], VName)])
 -> t (ImpM rep r op) ())
-> ([([Count Bytes (TPrimExp Int64 VName)], VName)]
    -> [([Count Bytes (TPrimExp Int64 VName)], VName)])
-> t (ImpM rep r op) ()
forall a b. (a -> b) -> a -> b
$ ([Count Bytes (TPrimExp Int64 VName)], VName)
-> [([Count Bytes (TPrimExp Int64 VName)], VName)]
-> [([Count Bytes (TPrimExp Int64 VName)], VName)]
forall a. Eq a => a -> [a] -> [a]
L.delete ([Count Bytes (TPrimExp Int64 VName)], VName)
mem
          ([Count Bytes (TPrimExp Int64 VName)], VName)
-> t (ImpM rep r op) ([Count Bytes (TPrimExp Int64 VName)], VName)
forall a. a -> t (ImpM rep r op) a
forall (f :: * -> *) a. Applicative f => a -> f a
pure ([Count Bytes (TPrimExp Int64 VName)], VName)
mem
        (Maybe ([Count Bytes (TPrimExp Int64 VName)], VName)
Nothing, ([Count Bytes (TPrimExp Int64 VName)]
size', VName
mem) : [([Count Bytes (TPrimExp Int64 VName)], VName)]
mems') -> do
          [([Count Bytes (TPrimExp Int64 VName)], VName)]
-> t (ImpM rep r op) ()
forall s (m :: * -> *). MonadState s m => s -> m ()
put [([Count Bytes (TPrimExp Int64 VName)], VName)]
mems'
          ([Count Bytes (TPrimExp Int64 VName)], VName)
-> t (ImpM rep r op) ([Count Bytes (TPrimExp Int64 VName)], VName)
forall a. a -> t (ImpM rep r op) a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (Count Bytes (TPrimExp Int64 VName)
size Count Bytes (TPrimExp Int64 VName)
-> [Count Bytes (TPrimExp Int64 VName)]
-> [Count Bytes (TPrimExp Int64 VName)]
forall a. a -> [a] -> [a]
: [Count Bytes (TPrimExp Int64 VName)]
size', VName
mem)
        (Maybe ([Count Bytes (TPrimExp Int64 VName)], VName)
Nothing, []) -> do
          mem <- ImpM rep r op VName -> t (ImpM rep r op) VName
forall (m :: * -> *) a. Monad m => m a -> t m a
forall (t :: (* -> *) -> * -> *) (m :: * -> *) a.
(MonadTrans t, Monad m) =>
m a -> t m a
lift (ImpM rep r op VName -> t (ImpM rep r op) VName)
-> ImpM rep r op VName -> t (ImpM rep r op) VName
forall a b. (a -> b) -> a -> b
$ Name -> Space -> ImpM rep r op VName
forall rep r op. Name -> Space -> ImpM rep r op VName
sDeclareMem Name
"scan_arr_mem" (Space -> ImpM rep r op VName) -> Space -> ImpM rep r op VName
forall a b. (a -> b) -> a -> b
$ SpaceId -> Space
Space SpaceId
"shared"
          pure ([size], mem)

type CrossesSegment = Maybe (Imp.TExp Int64 -> Imp.TExp Int64 -> Imp.TExp Bool)

localArrayIndex :: KernelConstants -> Imp.TExp Int64
localArrayIndex :: KernelConstants -> TPrimExp Int64 VName
localArrayIndex KernelConstants
constants = TPrimExp Int32 VName -> TPrimExp Int64 VName
forall {k} (t :: k) v. IntExp t => TPrimExp t v -> TPrimExp Int64 v
sExt64 (KernelConstants -> TPrimExp Int32 VName
kernelLocalThreadId KernelConstants
constants)

barrierFor :: Lambda GPUMem -> (Bool, Imp.Fence, InKernelGen ())
barrierFor :: Lambda GPUMem -> (Bool, Fence, ImpM GPUMem KernelEnv KernelOp ())
barrierFor Lambda GPUMem
scan_op = (Bool
array_scan, Fence
fence, KernelOp -> ImpM GPUMem KernelEnv KernelOp ()
forall op rep r. op -> ImpM rep r op ()
sOp (KernelOp -> ImpM GPUMem KernelEnv KernelOp ())
-> KernelOp -> ImpM GPUMem KernelEnv KernelOp ()
forall a b. (a -> b) -> a -> b
$ Fence -> KernelOp
Imp.Barrier Fence
fence)
  where
    array_scan :: Bool
array_scan = Bool -> Bool
not (Bool -> Bool) -> Bool -> Bool
forall a b. (a -> b) -> a -> b
$ (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] -> Bool) -> [Type] -> Bool
forall a b. (a -> b) -> a -> b
$ Lambda GPUMem -> [Type]
forall rep t. GLambda rep t -> [t]
lambdaReturnType Lambda GPUMem
scan_op
    fence :: Fence
fence
      | Bool
array_scan = Fence
Imp.FenceGlobal
      | Bool
otherwise = Fence
Imp.FenceLocal

xParams, yParams :: SegBinOp GPUMem -> [LParam GPUMem]
xParams :: SegBinOp GPUMem -> [LParam GPUMem]
xParams SegBinOp GPUMem
scan =
  Int -> [Param LParamMem] -> [Param LParamMem]
forall a. Int -> [a] -> [a]
take ([SubExp] -> Int
forall a. [a] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length (SegBinOp GPUMem -> [SubExp]
forall rep. SegBinOp rep -> [SubExp]
segBinOpNeutral SegBinOp GPUMem
scan)) (Lambda GPUMem -> [LParam GPUMem]
forall rep t. GLambda rep t -> [LParam rep]
lambdaParams (SegBinOp GPUMem -> Lambda GPUMem
forall rep. SegBinOp rep -> Lambda rep
segBinOpLambda SegBinOp GPUMem
scan))
yParams :: SegBinOp GPUMem -> [LParam GPUMem]
yParams SegBinOp GPUMem
scan =
  Int -> [Param LParamMem] -> [Param LParamMem]
forall a. Int -> [a] -> [a]
drop ([SubExp] -> Int
forall a. [a] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length (SegBinOp GPUMem -> [SubExp]
forall rep. SegBinOp rep -> [SubExp]
segBinOpNeutral SegBinOp GPUMem
scan)) (Lambda GPUMem -> [LParam GPUMem]
forall rep t. GLambda rep t -> [LParam rep]
lambdaParams (SegBinOp GPUMem -> Lambda GPUMem
forall rep. SegBinOp rep -> Lambda rep
segBinOpLambda SegBinOp GPUMem
scan))

writeToScanValues ::
  [VName] ->
  ([VName], SegBinOp GPUMem, [KernelResult]) ->
  InKernelGen ()
writeToScanValues :: [VName]
-> ([VName], SegBinOp GPUMem, [KernelResult])
-> ImpM GPUMem KernelEnv KernelOp ()
writeToScanValues [VName]
gtids ([VName]
pes, SegBinOp GPUMem
scan, [KernelResult]
scan_res)
  | Shape -> Int
forall a. ArrayShape a => a -> Int
shapeRank (SegBinOp GPUMem -> Shape
forall rep. SegBinOp rep -> Shape
segBinOpShape SegBinOp GPUMem
scan) Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
> Int
0 =
      [(VName, KernelResult)]
-> ((VName, KernelResult) -> ImpM GPUMem KernelEnv KernelOp ())
-> ImpM GPUMem KernelEnv KernelOp ()
forall (t :: * -> *) (m :: * -> *) a b.
(Foldable t, Monad m) =>
t a -> (a -> m b) -> m ()
forM_ ([VName] -> [KernelResult] -> [(VName, KernelResult)]
forall a b. [a] -> [b] -> [(a, b)]
zip [VName]
pes [KernelResult]
scan_res) (((VName, KernelResult) -> ImpM GPUMem KernelEnv KernelOp ())
 -> ImpM GPUMem KernelEnv KernelOp ())
-> ((VName, KernelResult) -> ImpM GPUMem KernelEnv KernelOp ())
-> ImpM GPUMem KernelEnv KernelOp ()
forall a b. (a -> b) -> a -> b
$ \(VName
pe, KernelResult
res) ->
        VName
-> [TPrimExp Int64 VName]
-> SubExp
-> [TPrimExp Int64 VName]
-> ImpM GPUMem KernelEnv KernelOp ()
forall rep r op.
VName
-> [TPrimExp Int64 VName]
-> SubExp
-> [TPrimExp Int64 VName]
-> ImpM rep r op ()
copyDWIMFix
          VName
pe
          ((VName -> TPrimExp Int64 VName)
-> [VName] -> [TPrimExp Int64 VName]
forall a b. (a -> b) -> [a] -> [b]
map VName -> TPrimExp Int64 VName
forall a. a -> TPrimExp Int64 a
Imp.le64 [VName]
gtids)
          (KernelResult -> SubExp
kernelResultSubExp KernelResult
res)
          []
  | Bool
otherwise =
      [(Param LParamMem, KernelResult)]
-> ((Param LParamMem, KernelResult)
    -> ImpM GPUMem KernelEnv KernelOp ())
-> ImpM GPUMem KernelEnv KernelOp ()
forall (t :: * -> *) (m :: * -> *) a b.
(Foldable t, Monad m) =>
t a -> (a -> m b) -> m ()
forM_ ([Param LParamMem]
-> [KernelResult] -> [(Param LParamMem, KernelResult)]
forall a b. [a] -> [b] -> [(a, b)]
zip (SegBinOp GPUMem -> [LParam GPUMem]
yParams SegBinOp GPUMem
scan) [KernelResult]
scan_res) (((Param LParamMem, KernelResult)
  -> ImpM GPUMem KernelEnv KernelOp ())
 -> ImpM GPUMem KernelEnv KernelOp ())
-> ((Param LParamMem, KernelResult)
    -> ImpM GPUMem KernelEnv KernelOp ())
-> ImpM GPUMem KernelEnv KernelOp ()
forall a b. (a -> b) -> a -> b
$ \(Param LParamMem
p, KernelResult
res) ->
        VName
-> [TPrimExp Int64 VName]
-> SubExp
-> [TPrimExp Int64 VName]
-> ImpM GPUMem KernelEnv KernelOp ()
forall rep r op.
VName
-> [TPrimExp Int64 VName]
-> SubExp
-> [TPrimExp Int64 VName]
-> ImpM rep r op ()
copyDWIMFix (Param LParamMem -> VName
forall dec. Param dec -> VName
paramName Param LParamMem
p) [] (KernelResult -> SubExp
kernelResultSubExp KernelResult
res) []

readToScanValues ::
  [Imp.TExp Int64] ->
  [VName] ->
  SegBinOp GPUMem ->
  InKernelGen ()
readToScanValues :: [TPrimExp Int64 VName]
-> [VName] -> SegBinOp GPUMem -> ImpM GPUMem KernelEnv KernelOp ()
readToScanValues [TPrimExp Int64 VName]
is [VName]
pes SegBinOp GPUMem
scan
  | Shape -> Int
forall a. ArrayShape a => a -> Int
shapeRank (SegBinOp GPUMem -> Shape
forall rep. SegBinOp rep -> Shape
segBinOpShape SegBinOp GPUMem
scan) Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
> Int
0 =
      [(Param LParamMem, VName)]
-> ((Param LParamMem, VName) -> ImpM GPUMem KernelEnv KernelOp ())
-> ImpM GPUMem KernelEnv KernelOp ()
forall (t :: * -> *) (m :: * -> *) a b.
(Foldable t, Monad m) =>
t a -> (a -> m b) -> m ()
forM_ ([Param LParamMem] -> [VName] -> [(Param LParamMem, VName)]
forall a b. [a] -> [b] -> [(a, b)]
zip (SegBinOp GPUMem -> [LParam GPUMem]
yParams SegBinOp GPUMem
scan) [VName]
pes) (((Param LParamMem, VName) -> ImpM GPUMem KernelEnv KernelOp ())
 -> ImpM GPUMem KernelEnv KernelOp ())
-> ((Param LParamMem, VName) -> ImpM GPUMem KernelEnv KernelOp ())
-> ImpM GPUMem KernelEnv KernelOp ()
forall a b. (a -> b) -> a -> b
$ \(Param LParamMem
p, VName
pe) ->
        VName
-> [TPrimExp Int64 VName]
-> SubExp
-> [TPrimExp Int64 VName]
-> ImpM GPUMem KernelEnv KernelOp ()
forall rep r op.
VName
-> [TPrimExp Int64 VName]
-> SubExp
-> [TPrimExp Int64 VName]
-> ImpM rep r op ()
copyDWIMFix (Param LParamMem -> VName
forall dec. Param dec -> VName
paramName Param LParamMem
p) [] (VName -> SubExp
Var VName
pe) [TPrimExp Int64 VName]
is
  | Bool
otherwise =
      () -> ImpM GPUMem KernelEnv KernelOp ()
forall a. a -> ImpM GPUMem KernelEnv KernelOp a
forall (f :: * -> *) a. Applicative f => a -> f a
pure ()

readCarries ::
  Imp.TExp Int64 ->
  Imp.TExp Int64 ->
  [Imp.TExp Int64] ->
  [Imp.TExp Int64] ->
  [VName] ->
  SegBinOp GPUMem ->
  InKernelGen ()
readCarries :: TPrimExp Int64 VName
-> TPrimExp Int64 VName
-> [TPrimExp Int64 VName]
-> [TPrimExp Int64 VName]
-> [VName]
-> SegBinOp GPUMem
-> ImpM GPUMem KernelEnv KernelOp ()
readCarries TPrimExp Int64 VName
chunk_id TPrimExp Int64 VName
chunk_offset [TPrimExp Int64 VName]
dims' [TPrimExp Int64 VName]
vec_is [VName]
pes SegBinOp GPUMem
scan
  | Shape -> Int
forall a. ArrayShape a => a -> Int
shapeRank (SegBinOp GPUMem -> Shape
forall rep. SegBinOp rep -> Shape
segBinOpShape SegBinOp GPUMem
scan) Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
> Int
0 = do
      ltid <- KernelConstants -> TPrimExp Int32 VName
kernelLocalThreadId (KernelConstants -> TPrimExp Int32 VName)
-> (KernelEnv -> KernelConstants)
-> KernelEnv
-> TPrimExp Int32 VName
forall b c a. (b -> c) -> (a -> b) -> a -> c
. KernelEnv -> KernelConstants
kernelConstants (KernelEnv -> TPrimExp Int32 VName)
-> ImpM GPUMem KernelEnv KernelOp KernelEnv
-> ImpM GPUMem KernelEnv KernelOp (TPrimExp Int32 VName)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> ImpM GPUMem KernelEnv KernelOp KernelEnv
forall rep r op. ImpM rep r op r
askEnv
      -- We may have to reload the carries from the output of the
      -- previous chunk.
      sIf
        (chunk_id .>. 0 .&&. ltid .==. 0)
        ( do
            let is = [TPrimExp Int64 VName]
-> TPrimExp Int64 VName -> [TPrimExp Int64 VName]
forall num. IntegralExp num => [num] -> num -> [num]
unflattenIndex [TPrimExp Int64 VName]
dims' (TPrimExp Int64 VName -> [TPrimExp Int64 VName])
-> TPrimExp Int64 VName -> [TPrimExp Int64 VName]
forall a b. (a -> b) -> a -> b
$ TPrimExp Int64 VName
chunk_offset TPrimExp Int64 VName
-> TPrimExp Int64 VName -> TPrimExp Int64 VName
forall a. Num a => a -> a -> a
- TPrimExp Int64 VName
1
            forM_ (zip (xParams scan) pes) $ \(Param LParamMem
p, VName
pe) ->
              VName
-> [TPrimExp Int64 VName]
-> SubExp
-> [TPrimExp Int64 VName]
-> ImpM GPUMem KernelEnv KernelOp ()
forall rep r op.
VName
-> [TPrimExp Int64 VName]
-> SubExp
-> [TPrimExp Int64 VName]
-> ImpM rep r op ()
copyDWIMFix (Param LParamMem -> VName
forall dec. Param dec -> VName
paramName Param LParamMem
p) [] (VName -> SubExp
Var VName
pe) ([TPrimExp Int64 VName]
is [TPrimExp Int64 VName]
-> [TPrimExp Int64 VName] -> [TPrimExp Int64 VName]
forall a. [a] -> [a] -> [a]
++ [TPrimExp Int64 VName]
vec_is)
        )
        ( forM_ (zip (xParams scan) (segBinOpNeutral scan)) $ \(Param LParamMem
p, SubExp
ne) ->
            VName
-> [TPrimExp Int64 VName]
-> SubExp
-> [TPrimExp Int64 VName]
-> ImpM GPUMem KernelEnv KernelOp ()
forall rep r op.
VName
-> [TPrimExp Int64 VName]
-> SubExp
-> [TPrimExp Int64 VName]
-> ImpM rep r op ()
copyDWIMFix (Param LParamMem -> VName
forall dec. Param dec -> VName
paramName Param LParamMem
p) [] SubExp
ne []
        )
  | Bool
otherwise =
      () -> ImpM GPUMem KernelEnv KernelOp ()
forall a. a -> ImpM GPUMem KernelEnv KernelOp a
forall (f :: * -> *) a. Applicative f => a -> f a
pure ()

-- | Produce partially scanned intervals; one per threadblock.
--
-- The @Maybe VName@ for the result is to cater to accumulators, which become
-- @Nothing@.
scanStage1 ::
  [VName] ->
  [Maybe VName] ->
  Count NumBlocks SubExp ->
  Count BlockSize SubExp ->
  SegSpace ->
  [SegBinOp GPUMem] ->
  KernelBody GPUMem ->
  CallKernelGen (TV Int32, Imp.TExp Int64, CrossesSegment)
scanStage1 :: [VName]
-> [Maybe VName]
-> Count NumBlocks SubExp
-> Count BlockSize SubExp
-> SegSpace
-> [SegBinOp GPUMem]
-> KernelBody GPUMem
-> CallKernelGen (TV Int32, TPrimExp Int64 VName, CrossesSegment)
scanStage1 [VName]
scan_out [Maybe VName]
map_out Count NumBlocks SubExp
num_tblocks Count BlockSize SubExp
tblock_size SegSpace
space [SegBinOp GPUMem]
scans KernelBody GPUMem
kbody = do
  let num_tblocks' :: Count NumBlocks (TPrimExp Int64 VName)
num_tblocks' = (SubExp -> TPrimExp Int64 VName)
-> Count NumBlocks SubExp -> Count NumBlocks (TPrimExp Int64 VName)
forall a b. (a -> b) -> Count NumBlocks a -> Count NumBlocks b
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap SubExp -> TPrimExp Int64 VName
pe64 Count NumBlocks SubExp
num_tblocks
      tblock_size' :: Count BlockSize (TPrimExp Int64 VName)
tblock_size' = (SubExp -> TPrimExp Int64 VName)
-> Count BlockSize SubExp -> Count BlockSize (TPrimExp Int64 VName)
forall a b. (a -> b) -> Count BlockSize a -> Count BlockSize b
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap SubExp -> TPrimExp Int64 VName
pe64 Count BlockSize SubExp
tblock_size
  num_threads <- Name
-> TPrimExp Int32 VName -> ImpM GPUMem HostEnv HostOp (TV Int32)
forall {k} (t :: k) rep r op.
Name -> TExp t -> ImpM rep r op (TV t)
dPrimV Name
"num_threads" (TPrimExp Int32 VName -> ImpM GPUMem HostEnv HostOp (TV Int32))
-> TPrimExp Int32 VName -> ImpM GPUMem HostEnv HostOp (TV Int32)
forall a b. (a -> b) -> a -> b
$ TPrimExp Int64 VName -> TPrimExp Int32 VName
forall {k} (t :: k) v. IntExp t => TPrimExp t v -> TPrimExp Int32 v
sExt32 (TPrimExp Int64 VName -> TPrimExp Int32 VName)
-> TPrimExp Int64 VName -> TPrimExp Int32 VName
forall a b. (a -> b) -> a -> b
$ Count NumBlocks (TPrimExp Int64 VName) -> TPrimExp Int64 VName
forall {k} (u :: k) e. Count u e -> e
unCount Count NumBlocks (TPrimExp Int64 VName)
num_tblocks' TPrimExp Int64 VName
-> TPrimExp Int64 VName -> TPrimExp Int64 VName
forall a. Num a => a -> a -> a
* Count BlockSize (TPrimExp Int64 VName) -> TPrimExp Int64 VName
forall {k} (u :: k) e. Count u e -> e
unCount Count BlockSize (TPrimExp Int64 VName)
tblock_size'

  let (gtids, dims) = unzip $ unSegSpace space
      dims' = (SubExp -> TPrimExp Int64 VName)
-> [SubExp] -> [TPrimExp Int64 VName]
forall a b. (a -> b) -> [a] -> [b]
map SubExp -> TPrimExp Int64 VName
pe64 [SubExp]
dims
  let num_elements = [TPrimExp Int64 VName] -> TPrimExp Int64 VName
forall a. Num a => [a] -> a
forall (t :: * -> *) a. (Foldable t, Num a) => t a -> a
product [TPrimExp Int64 VName]
dims'
      elems_per_thread = TPrimExp Int64 VName
num_elements TPrimExp Int64 VName
-> TPrimExp Int64 VName -> TPrimExp Int64 VName
forall e. IntegralExp e => e -> e -> e
`ceilDiv` TPrimExp Int32 VName -> TPrimExp Int64 VName
forall {k} (t :: k) v. IntExp t => TPrimExp t v -> TPrimExp Int64 v
sExt64 (TV Int32 -> TPrimExp Int32 VName
forall {k} (t :: k). TV t -> TExp t
tvExp TV Int32
num_threads)
      elems_per_group = Count BlockSize (TPrimExp Int64 VName) -> TPrimExp Int64 VName
forall {k} (u :: k) e. Count u e -> e
unCount Count BlockSize (TPrimExp Int64 VName)
tblock_size' TPrimExp Int64 VName
-> TPrimExp Int64 VName -> TPrimExp Int64 VName
forall a. Num a => a -> a -> a
* TPrimExp Int64 VName
elems_per_thread

  let crossesSegment =
        case [TPrimExp Int64 VName] -> [TPrimExp Int64 VName]
forall a. [a] -> [a]
reverse [TPrimExp Int64 VName]
dims' of
          TPrimExp Int64 VName
segment_size : TPrimExp Int64 VName
_ : [TPrimExp Int64 VName]
_ -> (TPrimExp Int64 VName -> TPrimExp Int64 VName -> TExp Bool)
-> CrossesSegment
forall a. a -> Maybe a
Just ((TPrimExp Int64 VName -> TPrimExp Int64 VName -> TExp Bool)
 -> CrossesSegment)
-> (TPrimExp Int64 VName -> TPrimExp Int64 VName -> TExp Bool)
-> CrossesSegment
forall a b. (a -> b) -> a -> b
$ \TPrimExp Int64 VName
from TPrimExp Int64 VName
to ->
            (TPrimExp Int64 VName
to TPrimExp Int64 VName
-> TPrimExp Int64 VName -> TPrimExp Int64 VName
forall a. Num a => a -> a -> a
- TPrimExp Int64 VName
from) TPrimExp Int64 VName -> TPrimExp Int64 VName -> TExp Bool
forall {k} v (t :: k).
Eq v =>
TPrimExp t v -> TPrimExp t v -> TPrimExp Bool v
.>. (TPrimExp Int64 VName
to TPrimExp Int64 VName
-> TPrimExp Int64 VName -> TPrimExp Int64 VName
forall e. IntegralExp e => e -> e -> e
`rem` TPrimExp Int64 VName
segment_size)
          [TPrimExp Int64 VName]
_ -> CrossesSegment
forall a. Maybe a
Nothing

  sKernelThread "scan_stage1" (segFlat space) (defKernelAttrs num_tblocks tblock_size) $ do
    constants <- kernelConstants <$> askEnv
    let tblock_id = KernelConstants -> TPrimExp Int32 VName
kernelBlockId KernelConstants
constants
    all_local_arrs <- makeLocalArrays (sExt64 tblock_id) tblock_size scans

    -- The variables from scan_op will be used for the carry and such
    -- in the big chunking loop.
    forM_ scans $ \SegBinOp GPUMem
scan -> do
      Maybe (Exp GPUMem)
-> Scope GPUMem -> ImpM GPUMem KernelEnv KernelOp ()
forall rep (inner :: * -> *) r op.
Mem rep inner =>
Maybe (Exp rep) -> Scope rep -> ImpM rep r op ()
dScope Maybe (Exp GPUMem)
forall a. Maybe a
Nothing (Scope GPUMem -> ImpM GPUMem KernelEnv KernelOp ())
-> Scope GPUMem -> ImpM GPUMem KernelEnv KernelOp ()
forall a b. (a -> b) -> a -> b
$ [Param LParamMem] -> Scope GPUMem
forall rep dec. (LParamInfo rep ~ dec) => [Param dec] -> Scope rep
scopeOfLParams ([Param LParamMem] -> Scope GPUMem)
-> [Param LParamMem] -> Scope GPUMem
forall a b. (a -> b) -> a -> b
$ Lambda GPUMem -> [LParam GPUMem]
forall rep t. GLambda rep t -> [LParam rep]
lambdaParams (Lambda GPUMem -> [LParam GPUMem])
-> Lambda GPUMem -> [LParam GPUMem]
forall a b. (a -> b) -> a -> b
$ SegBinOp GPUMem -> Lambda GPUMem
forall rep. SegBinOp rep -> Lambda rep
segBinOpLambda SegBinOp GPUMem
scan
      [(Param LParamMem, SubExp)]
-> ((Param LParamMem, SubExp) -> ImpM GPUMem KernelEnv KernelOp ())
-> ImpM GPUMem KernelEnv KernelOp ()
forall (t :: * -> *) (m :: * -> *) a b.
(Foldable t, Monad m) =>
t a -> (a -> m b) -> m ()
forM_ ([Param LParamMem] -> [SubExp] -> [(Param LParamMem, SubExp)]
forall a b. [a] -> [b] -> [(a, b)]
zip (SegBinOp GPUMem -> [LParam GPUMem]
xParams SegBinOp GPUMem
scan) (SegBinOp GPUMem -> [SubExp]
forall rep. SegBinOp rep -> [SubExp]
segBinOpNeutral SegBinOp GPUMem
scan)) (((Param LParamMem, SubExp) -> ImpM GPUMem KernelEnv KernelOp ())
 -> ImpM GPUMem KernelEnv KernelOp ())
-> ((Param LParamMem, SubExp) -> ImpM GPUMem KernelEnv KernelOp ())
-> ImpM GPUMem KernelEnv KernelOp ()
forall a b. (a -> b) -> a -> b
$ \(Param LParamMem
p, SubExp
ne) ->
        VName
-> [TPrimExp Int64 VName]
-> SubExp
-> [TPrimExp Int64 VName]
-> ImpM GPUMem KernelEnv KernelOp ()
forall rep r op.
VName
-> [TPrimExp Int64 VName]
-> SubExp
-> [TPrimExp Int64 VName]
-> ImpM rep r op ()
copyDWIMFix (Param LParamMem -> VName
forall dec. Param dec -> VName
paramName Param LParamMem
p) [] SubExp
ne []

    sFor "j" elems_per_thread $ \TPrimExp Int64 VName
j -> do
      chunk_offset <-
        Name
-> TPrimExp Int64 VName
-> ImpM GPUMem KernelEnv KernelOp (TV Int64)
forall {k} (t :: k) rep r op.
Name -> TExp t -> ImpM rep r op (TV t)
dPrimV Name
"chunk_offset" (TPrimExp Int64 VName -> ImpM GPUMem KernelEnv KernelOp (TV Int64))
-> TPrimExp Int64 VName
-> ImpM GPUMem KernelEnv KernelOp (TV Int64)
forall a b. (a -> b) -> a -> b
$
          TPrimExp Int64 VName -> TPrimExp Int64 VName
forall {k} (t :: k) v. IntExp t => TPrimExp t v -> TPrimExp Int64 v
sExt64 (KernelConstants -> TPrimExp Int64 VName
kernelBlockSize KernelConstants
constants) TPrimExp Int64 VName
-> TPrimExp Int64 VName -> TPrimExp Int64 VName
forall a. Num a => a -> a -> a
* TPrimExp Int64 VName
j
            TPrimExp Int64 VName
-> TPrimExp Int64 VName -> TPrimExp Int64 VName
forall a. Num a => a -> a -> a
+ TPrimExp Int32 VName -> TPrimExp Int64 VName
forall {k} (t :: k) v. IntExp t => TPrimExp t v -> TPrimExp Int64 v
sExt64 (KernelConstants -> TPrimExp Int32 VName
kernelBlockId KernelConstants
constants) TPrimExp Int64 VName
-> TPrimExp Int64 VName -> TPrimExp Int64 VName
forall a. Num a => a -> a -> a
* TPrimExp Int64 VName
elems_per_group
      flat_idx <-
        dPrimV "flat_idx" $
          tvExp chunk_offset + sExt64 (kernelLocalThreadId constants)
      -- Construct segment indices.
      zipWithM_ dPrimV_ gtids $ unflattenIndex dims' $ tvExp flat_idx

      let per_scan_pes = [SegBinOp GPUMem] -> [VName] -> [[VName]]
forall rep a. [SegBinOp rep] -> [a] -> [[a]]
segBinOpChunks [SegBinOp GPUMem]
scans [VName]
scan_out

          in_bounds =
            (TExp Bool -> TExp Bool -> TExp Bool) -> [TExp Bool] -> TExp Bool
forall a. (a -> a -> a) -> [a] -> a
forall (t :: * -> *) a. Foldable t => (a -> a -> a) -> t a -> a
foldl1 TExp Bool -> TExp Bool -> TExp Bool
forall v.
Eq v =>
TPrimExp Bool v -> TPrimExp Bool v -> TPrimExp Bool v
(.&&.) ([TExp Bool] -> TExp Bool) -> [TExp Bool] -> TExp Bool
forall a b. (a -> b) -> a -> b
$ (TPrimExp Int64 VName -> TPrimExp Int64 VName -> TExp Bool)
-> [TPrimExp Int64 VName] -> [TPrimExp Int64 VName] -> [TExp Bool]
forall a b c. (a -> b -> c) -> [a] -> [b] -> [c]
zipWith TPrimExp Int64 VName -> TPrimExp Int64 VName -> TExp Bool
forall {k} v (t :: k).
Eq v =>
TPrimExp t v -> TPrimExp t v -> TPrimExp Bool v
(.<.) ((VName -> TPrimExp Int64 VName)
-> [VName] -> [TPrimExp Int64 VName]
forall a b. (a -> b) -> [a] -> [b]
map VName -> TPrimExp Int64 VName
forall a. a -> TPrimExp Int64 a
Imp.le64 [VName]
gtids) [TPrimExp Int64 VName]
dims'

          when_in_bounds =
            Names
-> Stms GPUMem
-> ImpM GPUMem KernelEnv KernelOp ()
-> ImpM GPUMem KernelEnv KernelOp ()
forall rep r op.
Names -> Stms rep -> ImpM rep r op () -> ImpM rep r op ()
compileStms Names
forall a. Monoid a => a
mempty (KernelBody GPUMem -> Stms GPUMem
forall rep res. GBody rep res -> Stms rep
bodyStms KernelBody GPUMem
kbody) (ImpM GPUMem KernelEnv KernelOp ()
 -> ImpM GPUMem KernelEnv KernelOp ())
-> ImpM GPUMem KernelEnv KernelOp ()
-> ImpM GPUMem KernelEnv KernelOp ()
forall a b. (a -> b) -> a -> b
$ do
              let ([KernelResult]
all_scan_res, [KernelResult]
map_res) =
                    Int -> [KernelResult] -> ([KernelResult], [KernelResult])
forall a. Int -> [a] -> ([a], [a])
splitAt ([SegBinOp GPUMem] -> Int
forall rep. [SegBinOp rep] -> Int
segBinOpResults [SegBinOp GPUMem]
scans) ([KernelResult] -> ([KernelResult], [KernelResult]))
-> [KernelResult] -> ([KernelResult], [KernelResult])
forall a b. (a -> b) -> a -> b
$ KernelBody GPUMem -> [KernelResult]
forall rep res. GBody rep res -> [res]
bodyResult KernelBody GPUMem
kbody
                  per_scan_res :: [[KernelResult]]
per_scan_res =
                    [SegBinOp GPUMem] -> [KernelResult] -> [[KernelResult]]
forall rep a. [SegBinOp rep] -> [a] -> [[a]]
segBinOpChunks [SegBinOp GPUMem]
scans [KernelResult]
all_scan_res

              Text
-> ImpM GPUMem KernelEnv KernelOp ()
-> ImpM GPUMem KernelEnv KernelOp ()
forall rep r op. Text -> ImpM rep r op () -> ImpM rep r op ()
sComment Text
"write to-scan values to parameters" (ImpM GPUMem KernelEnv KernelOp ()
 -> ImpM GPUMem KernelEnv KernelOp ())
-> ImpM GPUMem KernelEnv KernelOp ()
-> ImpM GPUMem KernelEnv KernelOp ()
forall a b. (a -> b) -> a -> b
$
                (([VName], SegBinOp GPUMem, [KernelResult])
 -> ImpM GPUMem KernelEnv KernelOp ())
-> [([VName], SegBinOp GPUMem, [KernelResult])]
-> ImpM GPUMem KernelEnv KernelOp ()
forall (t :: * -> *) (m :: * -> *) a b.
(Foldable t, Monad m) =>
(a -> m b) -> t a -> m ()
mapM_ ([VName]
-> ([VName], SegBinOp GPUMem, [KernelResult])
-> ImpM GPUMem KernelEnv KernelOp ()
writeToScanValues [VName]
gtids) ([([VName], SegBinOp GPUMem, [KernelResult])]
 -> ImpM GPUMem KernelEnv KernelOp ())
-> [([VName], SegBinOp GPUMem, [KernelResult])]
-> ImpM GPUMem KernelEnv KernelOp ()
forall a b. (a -> b) -> a -> b
$
                  [[VName]]
-> [SegBinOp GPUMem]
-> [[KernelResult]]
-> [([VName], SegBinOp GPUMem, [KernelResult])]
forall a b c. [a] -> [b] -> [c] -> [(a, b, c)]
zip3 [[VName]]
per_scan_pes [SegBinOp GPUMem]
scans [[KernelResult]]
per_scan_res

              Text
-> ImpM GPUMem KernelEnv KernelOp ()
-> ImpM GPUMem KernelEnv KernelOp ()
forall rep r op. Text -> ImpM rep r op () -> ImpM rep r op ()
sComment Text
"write mapped values results to global memory" (ImpM GPUMem KernelEnv KernelOp ()
 -> ImpM GPUMem KernelEnv KernelOp ())
-> ImpM GPUMem KernelEnv KernelOp ()
-> ImpM GPUMem KernelEnv KernelOp ()
forall a b. (a -> b) -> a -> b
$
                [(Maybe VName, KernelResult)]
-> ((Maybe VName, KernelResult)
    -> ImpM GPUMem KernelEnv KernelOp (Maybe ()))
-> ImpM GPUMem KernelEnv KernelOp ()
forall (t :: * -> *) (m :: * -> *) a b.
(Foldable t, Monad m) =>
t a -> (a -> m b) -> m ()
forM_ ([Maybe VName] -> [KernelResult] -> [(Maybe VName, KernelResult)]
forall a b. [a] -> [b] -> [(a, b)]
zip [Maybe VName]
map_out [KernelResult]
map_res) (((Maybe VName, KernelResult)
  -> ImpM GPUMem KernelEnv KernelOp (Maybe ()))
 -> ImpM GPUMem KernelEnv KernelOp ())
-> ((Maybe VName, KernelResult)
    -> ImpM GPUMem KernelEnv KernelOp (Maybe ()))
-> ImpM GPUMem KernelEnv KernelOp ()
forall a b. (a -> b) -> a -> b
$ \(Maybe VName
pe, KernelResult
se) ->
                  Maybe VName
-> (VName -> ImpM GPUMem KernelEnv KernelOp ())
-> ImpM GPUMem KernelEnv KernelOp (Maybe ())
forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
t a -> (a -> m b) -> m (t b)
forM Maybe VName
pe ((VName -> ImpM GPUMem KernelEnv KernelOp ())
 -> ImpM GPUMem KernelEnv KernelOp (Maybe ()))
-> (VName -> ImpM GPUMem KernelEnv KernelOp ())
-> ImpM GPUMem KernelEnv KernelOp (Maybe ())
forall a b. (a -> b) -> a -> b
$ \VName
p ->
                    VName
-> [TPrimExp Int64 VName]
-> SubExp
-> [TPrimExp Int64 VName]
-> ImpM GPUMem KernelEnv KernelOp ()
forall rep r op.
VName
-> [TPrimExp Int64 VName]
-> SubExp
-> [TPrimExp Int64 VName]
-> ImpM rep r op ()
copyDWIMFix
                      VName
p
                      ((VName -> TPrimExp Int64 VName)
-> [VName] -> [TPrimExp Int64 VName]
forall a b. (a -> b) -> [a] -> [b]
map VName -> TPrimExp Int64 VName
forall a. a -> TPrimExp Int64 a
Imp.le64 [VName]
gtids)
                      (KernelResult -> SubExp
kernelResultSubExp KernelResult
se)
                      []

      sComment "threads in bounds read input" $
        sWhen in_bounds when_in_bounds

      unless (all (null . segBinOpShape) scans) $
        sOp $
          Imp.Barrier Imp.FenceGlobal

      forM_ (zip3 per_scan_pes scans all_local_arrs) $
        \([VName]
pes, scan :: SegBinOp GPUMem
scan@(SegBinOp Commutativity
_ Lambda GPUMem
scan_op [SubExp]
nes Shape
vec_shape), [VName]
local_arrs) ->
          Text
-> ImpM GPUMem KernelEnv KernelOp ()
-> ImpM GPUMem KernelEnv KernelOp ()
forall rep r op. Text -> ImpM rep r op () -> ImpM rep r op ()
sComment Text
"do one intra-group scan operation" (ImpM GPUMem KernelEnv KernelOp ()
 -> ImpM GPUMem KernelEnv KernelOp ())
-> ImpM GPUMem KernelEnv KernelOp ()
-> ImpM GPUMem KernelEnv KernelOp ()
forall a b. (a -> b) -> a -> b
$ do
            let scan_x_params :: [LParam GPUMem]
scan_x_params = SegBinOp GPUMem -> [LParam GPUMem]
xParams SegBinOp GPUMem
scan
                (Bool
array_scan, Fence
fence, ImpM GPUMem KernelEnv KernelOp ()
barrier) = Lambda GPUMem -> (Bool, Fence, ImpM GPUMem KernelEnv KernelOp ())
barrierFor Lambda GPUMem
scan_op

            Bool
-> ImpM GPUMem KernelEnv KernelOp ()
-> ImpM GPUMem KernelEnv KernelOp ()
forall (f :: * -> *). Applicative f => Bool -> f () -> f ()
when Bool
array_scan ImpM GPUMem KernelEnv KernelOp ()
barrier

            Shape
-> ([TPrimExp Int64 VName] -> ImpM GPUMem KernelEnv KernelOp ())
-> ImpM GPUMem KernelEnv KernelOp ()
forall rep r op.
Shape
-> ([TPrimExp Int64 VName] -> ImpM rep r op ()) -> ImpM rep r op ()
sLoopNest Shape
vec_shape (([TPrimExp Int64 VName] -> ImpM GPUMem KernelEnv KernelOp ())
 -> ImpM GPUMem KernelEnv KernelOp ())
-> ([TPrimExp Int64 VName] -> ImpM GPUMem KernelEnv KernelOp ())
-> ImpM GPUMem KernelEnv KernelOp ()
forall a b. (a -> b) -> a -> b
$ \[TPrimExp Int64 VName]
vec_is -> do
              Text
-> ImpM GPUMem KernelEnv KernelOp ()
-> ImpM GPUMem KernelEnv KernelOp ()
forall rep r op. Text -> ImpM rep r op () -> ImpM rep r op ()
sComment Text
"maybe restore some to-scan values to parameters, or read neutral" (ImpM GPUMem KernelEnv KernelOp ()
 -> ImpM GPUMem KernelEnv KernelOp ())
-> ImpM GPUMem KernelEnv KernelOp ()
-> ImpM GPUMem KernelEnv KernelOp ()
forall a b. (a -> b) -> a -> b
$
                TExp Bool
-> ImpM GPUMem KernelEnv KernelOp ()
-> ImpM GPUMem KernelEnv KernelOp ()
-> ImpM GPUMem KernelEnv KernelOp ()
forall rep r op.
TExp Bool
-> ImpM rep r op () -> ImpM rep r op () -> ImpM rep r op ()
sIf
                  TExp Bool
in_bounds
                  ( do
                      [TPrimExp Int64 VName]
-> [VName] -> SegBinOp GPUMem -> ImpM GPUMem KernelEnv KernelOp ()
readToScanValues ((VName -> TPrimExp Int64 VName)
-> [VName] -> [TPrimExp Int64 VName]
forall a b. (a -> b) -> [a] -> [b]
map VName -> TPrimExp Int64 VName
forall a. a -> TPrimExp Int64 a
Imp.le64 [VName]
gtids [TPrimExp Int64 VName]
-> [TPrimExp Int64 VName] -> [TPrimExp Int64 VName]
forall a. [a] -> [a] -> [a]
++ [TPrimExp Int64 VName]
vec_is) [VName]
pes SegBinOp GPUMem
scan
                      TPrimExp Int64 VName
-> TPrimExp Int64 VName
-> [TPrimExp Int64 VName]
-> [TPrimExp Int64 VName]
-> [VName]
-> SegBinOp GPUMem
-> ImpM GPUMem KernelEnv KernelOp ()
readCarries TPrimExp Int64 VName
j (TV Int64 -> TPrimExp Int64 VName
forall {k} (t :: k). TV t -> TExp t
tvExp TV Int64
chunk_offset) [TPrimExp Int64 VName]
dims' [TPrimExp Int64 VName]
vec_is [VName]
pes SegBinOp GPUMem
scan
                  )
                  ( [(Param LParamMem, SubExp)]
-> ((Param LParamMem, SubExp) -> ImpM GPUMem KernelEnv KernelOp ())
-> ImpM GPUMem KernelEnv KernelOp ()
forall (t :: * -> *) (m :: * -> *) a b.
(Foldable t, Monad m) =>
t a -> (a -> m b) -> m ()
forM_ ([Param LParamMem] -> [SubExp] -> [(Param LParamMem, SubExp)]
forall a b. [a] -> [b] -> [(a, b)]
zip (SegBinOp GPUMem -> [LParam GPUMem]
yParams SegBinOp GPUMem
scan) (SegBinOp GPUMem -> [SubExp]
forall rep. SegBinOp rep -> [SubExp]
segBinOpNeutral SegBinOp GPUMem
scan)) (((Param LParamMem, SubExp) -> ImpM GPUMem KernelEnv KernelOp ())
 -> ImpM GPUMem KernelEnv KernelOp ())
-> ((Param LParamMem, SubExp) -> ImpM GPUMem KernelEnv KernelOp ())
-> ImpM GPUMem KernelEnv KernelOp ()
forall a b. (a -> b) -> a -> b
$ \(Param LParamMem
p, SubExp
ne) ->
                      VName
-> [TPrimExp Int64 VName]
-> SubExp
-> [TPrimExp Int64 VName]
-> ImpM GPUMem KernelEnv KernelOp ()
forall rep r op.
VName
-> [TPrimExp Int64 VName]
-> SubExp
-> [TPrimExp Int64 VName]
-> ImpM rep r op ()
copyDWIMFix (Param LParamMem -> VName
forall dec. Param dec -> VName
paramName Param LParamMem
p) [] SubExp
ne []
                  )

              Text
-> ImpM GPUMem KernelEnv KernelOp ()
-> ImpM GPUMem KernelEnv KernelOp ()
forall rep r op. Text -> ImpM rep r op () -> ImpM rep r op ()
sComment Text
"combine with carry and write to shared memory" (ImpM GPUMem KernelEnv KernelOp ()
 -> ImpM GPUMem KernelEnv KernelOp ())
-> ImpM GPUMem KernelEnv KernelOp ()
-> ImpM GPUMem KernelEnv KernelOp ()
forall a b. (a -> b) -> a -> b
$
                Names
-> Stms GPUMem
-> ImpM GPUMem KernelEnv KernelOp ()
-> ImpM GPUMem KernelEnv KernelOp ()
forall rep r op.
Names -> Stms rep -> ImpM rep r op () -> ImpM rep r op ()
compileStms Names
forall a. Monoid a => a
mempty (GBody GPUMem SubExpRes -> Stms GPUMem
forall rep res. GBody rep res -> Stms rep
bodyStms (GBody GPUMem SubExpRes -> Stms GPUMem)
-> GBody GPUMem SubExpRes -> Stms GPUMem
forall a b. (a -> b) -> a -> b
$ Lambda GPUMem -> GBody GPUMem SubExpRes
forall rep t. GLambda rep t -> Body rep
lambdaBody Lambda GPUMem
scan_op) (ImpM GPUMem KernelEnv KernelOp ()
 -> ImpM GPUMem KernelEnv KernelOp ())
-> ImpM GPUMem KernelEnv KernelOp ()
-> ImpM GPUMem KernelEnv KernelOp ()
forall a b. (a -> b) -> a -> b
$
                  [(VName, SubExp)]
-> ((VName, SubExp) -> ImpM GPUMem KernelEnv KernelOp ())
-> ImpM GPUMem KernelEnv KernelOp ()
forall (t :: * -> *) (m :: * -> *) a b.
(Foldable t, Monad m) =>
t a -> (a -> m b) -> m ()
forM_ ([VName] -> [SubExp] -> [(VName, SubExp)]
forall a b. [a] -> [b] -> [(a, b)]
zip [VName]
local_arrs ([SubExp] -> [(VName, SubExp)]) -> [SubExp] -> [(VName, SubExp)]
forall a b. (a -> b) -> a -> b
$ (SubExpRes -> SubExp) -> [SubExpRes] -> [SubExp]
forall a b. (a -> b) -> [a] -> [b]
map SubExpRes -> SubExp
resSubExp ([SubExpRes] -> [SubExp]) -> [SubExpRes] -> [SubExp]
forall a b. (a -> b) -> a -> b
$ GBody GPUMem SubExpRes -> [SubExpRes]
forall rep res. GBody rep res -> [res]
bodyResult (GBody GPUMem SubExpRes -> [SubExpRes])
-> GBody GPUMem SubExpRes -> [SubExpRes]
forall a b. (a -> b) -> a -> b
$ Lambda GPUMem -> GBody GPUMem SubExpRes
forall rep t. GLambda rep t -> Body rep
lambdaBody Lambda GPUMem
scan_op) (((VName, SubExp) -> ImpM GPUMem KernelEnv KernelOp ())
 -> ImpM GPUMem KernelEnv KernelOp ())
-> ((VName, SubExp) -> ImpM GPUMem KernelEnv KernelOp ())
-> ImpM GPUMem KernelEnv KernelOp ()
forall a b. (a -> b) -> a -> b
$
                    \(VName
arr, SubExp
se) ->
                      VName
-> [TPrimExp Int64 VName]
-> SubExp
-> [TPrimExp Int64 VName]
-> ImpM GPUMem KernelEnv KernelOp ()
forall rep r op.
VName
-> [TPrimExp Int64 VName]
-> SubExp
-> [TPrimExp Int64 VName]
-> ImpM rep r op ()
copyDWIMFix VName
arr [KernelConstants -> TPrimExp Int64 VName
localArrayIndex KernelConstants
constants] SubExp
se []

              let crossesSegment' :: Maybe (TPrimExp Int32 VName -> TPrimExp Int32 VName -> TExp Bool)
crossesSegment' = do
                    f <- CrossesSegment
crossesSegment
                    Just $ \TPrimExp Int32 VName
from TPrimExp Int32 VName
to ->
                      let from' :: TPrimExp Int64 VName
from' = TPrimExp Int32 VName -> TPrimExp Int64 VName
forall {k} (t :: k) v. IntExp t => TPrimExp t v -> TPrimExp Int64 v
sExt64 TPrimExp Int32 VName
from TPrimExp Int64 VName
-> TPrimExp Int64 VName -> TPrimExp Int64 VName
forall a. Num a => a -> a -> a
+ TV Int64 -> TPrimExp Int64 VName
forall {k} (t :: k). TV t -> TExp t
tvExp TV Int64
chunk_offset
                          to' :: TPrimExp Int64 VName
to' = TPrimExp Int32 VName -> TPrimExp Int64 VName
forall {k} (t :: k) v. IntExp t => TPrimExp t v -> TPrimExp Int64 v
sExt64 TPrimExp Int32 VName
to TPrimExp Int64 VName
-> TPrimExp Int64 VName -> TPrimExp Int64 VName
forall a. Num a => a -> a -> a
+ TV Int64 -> TPrimExp Int64 VName
forall {k} (t :: k). TV t -> TExp t
tvExp TV Int64
chunk_offset
                       in TPrimExp Int64 VName -> TPrimExp Int64 VName -> TExp Bool
f TPrimExp Int64 VName
from' TPrimExp Int64 VName
to'

              KernelOp -> ImpM GPUMem KernelEnv KernelOp ()
forall op rep r. op -> ImpM rep r op ()
sOp (KernelOp -> ImpM GPUMem KernelEnv KernelOp ())
-> KernelOp -> ImpM GPUMem KernelEnv KernelOp ()
forall a b. (a -> b) -> a -> b
$ Fence -> KernelOp
Imp.ErrorSync Fence
fence

              -- We need to avoid parameter name clashes.
              scan_op_renamed <- Lambda GPUMem -> ImpM GPUMem KernelEnv KernelOp (Lambda GPUMem)
forall rep (m :: * -> *) t.
(Renameable rep, MonadFreshNames m, Rename t) =>
GLambda rep t -> m (GLambda rep t)
renameLambda Lambda GPUMem
scan_op
              blockScan
                crossesSegment'
                (sExt64 $ tvExp num_threads)
                (sExt64 $ kernelBlockSize constants)
                scan_op_renamed
                local_arrs

              sComment "threads in bounds write partial scan result" $
                sWhen in_bounds $
                  forM_ (zip pes local_arrs) $ \(VName
pe, VName
arr) ->
                    VName
-> [TPrimExp Int64 VName]
-> SubExp
-> [TPrimExp Int64 VName]
-> ImpM GPUMem KernelEnv KernelOp ()
forall rep r op.
VName
-> [TPrimExp Int64 VName]
-> SubExp
-> [TPrimExp Int64 VName]
-> ImpM rep r op ()
copyDWIMFix
                      VName
pe
                      ((VName -> TPrimExp Int64 VName)
-> [VName] -> [TPrimExp Int64 VName]
forall a b. (a -> b) -> [a] -> [b]
map VName -> TPrimExp Int64 VName
forall a. a -> TPrimExp Int64 a
Imp.le64 [VName]
gtids [TPrimExp Int64 VName]
-> [TPrimExp Int64 VName] -> [TPrimExp Int64 VName]
forall a. [a] -> [a] -> [a]
++ [TPrimExp Int64 VName]
vec_is)
                      (VName -> SubExp
Var VName
arr)
                      [KernelConstants -> TPrimExp Int64 VName
localArrayIndex KernelConstants
constants]

              barrier

              let load_carry =
                    [(VName, Param LParamMem)]
-> ((VName, Param LParamMem) -> ImpM GPUMem KernelEnv KernelOp ())
-> ImpM GPUMem KernelEnv KernelOp ()
forall (t :: * -> *) (m :: * -> *) a b.
(Foldable t, Monad m) =>
t a -> (a -> m b) -> m ()
forM_ ([VName] -> [Param LParamMem] -> [(VName, Param LParamMem)]
forall a b. [a] -> [b] -> [(a, b)]
zip [VName]
local_arrs [LParam GPUMem]
[Param LParamMem]
scan_x_params) (((VName, Param LParamMem) -> ImpM GPUMem KernelEnv KernelOp ())
 -> ImpM GPUMem KernelEnv KernelOp ())
-> ((VName, Param LParamMem) -> ImpM GPUMem KernelEnv KernelOp ())
-> ImpM GPUMem KernelEnv KernelOp ()
forall a b. (a -> b) -> a -> b
$ \(VName
arr, Param LParamMem
p) ->
                      VName
-> [TPrimExp Int64 VName]
-> SubExp
-> [TPrimExp Int64 VName]
-> ImpM GPUMem KernelEnv KernelOp ()
forall rep r op.
VName
-> [TPrimExp Int64 VName]
-> SubExp
-> [TPrimExp Int64 VName]
-> ImpM rep r op ()
copyDWIMFix
                        (Param LParamMem -> VName
forall dec. Param dec -> VName
paramName Param LParamMem
p)
                        []
                        (VName -> SubExp
Var VName
arr)
                        [ TPrimExp Int64 VName -> TPrimExp Int64 VName
forall {k} (t :: k) v. IntExp t => TPrimExp t v -> TPrimExp Int64 v
sExt64 (KernelConstants -> TPrimExp Int64 VName
kernelBlockSize KernelConstants
constants) TPrimExp Int64 VName
-> TPrimExp Int64 VName -> TPrimExp Int64 VName
forall a. Num a => a -> a -> a
- TPrimExp Int64 VName
1
                        ]
                  load_neutral =
                    [(SubExp, Param LParamMem)]
-> ((SubExp, Param LParamMem) -> ImpM GPUMem KernelEnv KernelOp ())
-> ImpM GPUMem KernelEnv KernelOp ()
forall (t :: * -> *) (m :: * -> *) a b.
(Foldable t, Monad m) =>
t a -> (a -> m b) -> m ()
forM_ ([SubExp] -> [Param LParamMem] -> [(SubExp, Param LParamMem)]
forall a b. [a] -> [b] -> [(a, b)]
zip [SubExp]
nes [LParam GPUMem]
[Param LParamMem]
scan_x_params) (((SubExp, Param LParamMem) -> ImpM GPUMem KernelEnv KernelOp ())
 -> ImpM GPUMem KernelEnv KernelOp ())
-> ((SubExp, Param LParamMem) -> ImpM GPUMem KernelEnv KernelOp ())
-> ImpM GPUMem KernelEnv KernelOp ()
forall a b. (a -> b) -> a -> b
$ \(SubExp
ne, Param LParamMem
p) ->
                      VName
-> [TPrimExp Int64 VName]
-> SubExp
-> [TPrimExp Int64 VName]
-> ImpM GPUMem KernelEnv KernelOp ()
forall rep r op.
VName
-> [TPrimExp Int64 VName]
-> SubExp
-> [TPrimExp Int64 VName]
-> ImpM rep r op ()
copyDWIMFix (Param LParamMem -> VName
forall dec. Param dec -> VName
paramName Param LParamMem
p) [] SubExp
ne []

              sComment "first thread reads last element as carry-in for next iteration" $ do
                crosses_segment <- dPrimVE "crosses_segment" $
                  case crossesSegment of
                    CrossesSegment
Nothing -> TExp Bool
forall v. TPrimExp Bool v
false
                    Just TPrimExp Int64 VName -> TPrimExp Int64 VName -> TExp Bool
f ->
                      TPrimExp Int64 VName -> TPrimExp Int64 VName -> TExp Bool
f
                        ( TV Int64 -> TPrimExp Int64 VName
forall {k} (t :: k). TV t -> TExp t
tvExp TV Int64
chunk_offset
                            TPrimExp Int64 VName
-> TPrimExp Int64 VName -> TPrimExp Int64 VName
forall a. Num a => a -> a -> a
+ TPrimExp Int64 VName -> TPrimExp Int64 VName
forall {k} (t :: k) v. IntExp t => TPrimExp t v -> TPrimExp Int64 v
sExt64 (KernelConstants -> TPrimExp Int64 VName
kernelBlockSize KernelConstants
constants)
                            TPrimExp Int64 VName
-> TPrimExp Int64 VName -> TPrimExp Int64 VName
forall a. Num a => a -> a -> a
- TPrimExp Int64 VName
1
                        )
                        ( TV Int64 -> TPrimExp Int64 VName
forall {k} (t :: k). TV t -> TExp t
tvExp TV Int64
chunk_offset
                            TPrimExp Int64 VName
-> TPrimExp Int64 VName -> TPrimExp Int64 VName
forall a. Num a => a -> a -> a
+ TPrimExp Int64 VName -> TPrimExp Int64 VName
forall {k} (t :: k) v. IntExp t => TPrimExp t v -> TPrimExp Int64 v
sExt64 (KernelConstants -> TPrimExp Int64 VName
kernelBlockSize KernelConstants
constants)
                        )
                should_load_carry <-
                  dPrimVE "should_load_carry" $
                    kernelLocalThreadId constants .==. 0 .&&. bNot crosses_segment
                sWhen should_load_carry load_carry
                when array_scan barrier
                sUnless should_load_carry load_neutral

              barrier

  pure (num_threads, elems_per_group, crossesSegment)

scanStage2 ::
  [VName] ->
  TV Int32 ->
  Imp.TExp Int64 ->
  Count NumBlocks SubExp ->
  Count BlockSize SubExp ->
  CrossesSegment ->
  SegSpace ->
  [SegBinOp GPUMem] ->
  CallKernelGen ()
scanStage2 :: [VName]
-> TV Int32
-> TPrimExp Int64 VName
-> Count NumBlocks SubExp
-> Count BlockSize SubExp
-> CrossesSegment
-> SegSpace
-> [SegBinOp GPUMem]
-> CallKernelGen ()
scanStage2 [VName]
scan_out TV Int32
stage1_num_threads TPrimExp Int64 VName
elems_per_group Count NumBlocks SubExp
stage1_num_tblocks Count BlockSize SubExp
stage2_tblock_size CrossesSegment
crossesSegment SegSpace
space [SegBinOp GPUMem]
scans = do
  let ([VName]
gtids, [SubExp]
dims) = [(VName, SubExp)] -> ([VName], [SubExp])
forall a b. [(a, b)] -> ([a], [b])
unzip ([(VName, SubExp)] -> ([VName], [SubExp]))
-> [(VName, SubExp)] -> ([VName], [SubExp])
forall a b. (a -> b) -> a -> b
$ SegSpace -> [(VName, SubExp)]
unSegSpace SegSpace
space
      dims' :: [TPrimExp Int64 VName]
dims' = (SubExp -> TPrimExp Int64 VName)
-> [SubExp] -> [TPrimExp Int64 VName]
forall a b. (a -> b) -> [a] -> [b]
map SubExp -> TPrimExp Int64 VName
pe64 [SubExp]
dims
      stage1_num_tblocks_e :: TPrimExp Int64 VName
stage1_num_tblocks_e = SubExp -> TPrimExp Int64 VName
pe64 (SubExp -> TPrimExp Int64 VName) -> SubExp -> TPrimExp Int64 VName
forall a b. (a -> b) -> a -> b
$ Count NumBlocks SubExp -> SubExp
forall {k} (u :: k) e. Count u e -> e
unCount Count NumBlocks SubExp
stage1_num_tblocks
      stage2_tblock_size_e :: TPrimExp Int64 VName
stage2_tblock_size_e = SubExp -> TPrimExp Int64 VName
pe64 (SubExp -> TPrimExp Int64 VName) -> SubExp -> TPrimExp Int64 VName
forall a b. (a -> b) -> a -> b
$ Count BlockSize SubExp -> SubExp
forall {k} (u :: k) e. Count u e -> e
unCount Count BlockSize SubExp
stage2_tblock_size

  -- Number of chunks needed to cover all stage-1 blocks.
  num_chunks <-
    Name
-> TPrimExp Int64 VName
-> ImpM GPUMem HostEnv HostOp (TPrimExp Int64 VName)
forall {k} (t :: k) rep r op.
Name -> TExp t -> ImpM rep r op (TExp t)
dPrimVE Name
"stage2_num_chunks" (TPrimExp Int64 VName
 -> ImpM GPUMem HostEnv HostOp (TPrimExp Int64 VName))
-> TPrimExp Int64 VName
-> ImpM GPUMem HostEnv HostOp (TPrimExp Int64 VName)
forall a b. (a -> b) -> a -> b
$
      TPrimExp Int64 VName
stage1_num_tblocks_e TPrimExp Int64 VName
-> TPrimExp Int64 VName -> TPrimExp Int64 VName
forall e. IntegralExp e => e -> e -> e
`ceilDiv` TPrimExp Int64 VName
stage2_tblock_size_e

  sKernelThread "scan_stage2" (segFlat space) (defKernelAttrs (Count (intConst Int64 1)) stage2_tblock_size) $ do
    constants <- kernelConstants <$> askEnv
    let tblock_id = KernelConstants -> TPrimExp Int32 VName
kernelBlockId KernelConstants
constants
    per_scan_local_arrs <- makeLocalArrays (sExt64 tblock_id) stage2_tblock_size scans
    let per_scan_pes = [SegBinOp GPUMem] -> [VName] -> [[VName]]
forall rep a. [SegBinOp rep] -> [a] -> [[a]]
segBinOpChunks [SegBinOp GPUMem]
scans [VName]
scan_out

    -- Declare lambda params and initialise carries (xParams) to the
    -- neutral element.  For scalar scans these persist across chunk
    -- iterations as registers; for array scans they are reloaded from
    -- global memory at the start of each chunk.
    forM_ scans $ \SegBinOp GPUMem
scan -> do
      Maybe (Exp GPUMem)
-> Scope GPUMem -> ImpM GPUMem KernelEnv KernelOp ()
forall rep (inner :: * -> *) r op.
Mem rep inner =>
Maybe (Exp rep) -> Scope rep -> ImpM rep r op ()
dScope Maybe (Exp GPUMem)
forall a. Maybe a
Nothing (Scope GPUMem -> ImpM GPUMem KernelEnv KernelOp ())
-> Scope GPUMem -> ImpM GPUMem KernelEnv KernelOp ()
forall a b. (a -> b) -> a -> b
$ [Param LParamMem] -> Scope GPUMem
forall rep dec. (LParamInfo rep ~ dec) => [Param dec] -> Scope rep
scopeOfLParams ([Param LParamMem] -> Scope GPUMem)
-> [Param LParamMem] -> Scope GPUMem
forall a b. (a -> b) -> a -> b
$ Lambda GPUMem -> [LParam GPUMem]
forall rep t. GLambda rep t -> [LParam rep]
lambdaParams (Lambda GPUMem -> [LParam GPUMem])
-> Lambda GPUMem -> [LParam GPUMem]
forall a b. (a -> b) -> a -> b
$ SegBinOp GPUMem -> Lambda GPUMem
forall rep. SegBinOp rep -> Lambda rep
segBinOpLambda SegBinOp GPUMem
scan
      [(Param LParamMem, SubExp)]
-> ((Param LParamMem, SubExp) -> ImpM GPUMem KernelEnv KernelOp ())
-> ImpM GPUMem KernelEnv KernelOp ()
forall (t :: * -> *) (m :: * -> *) a b.
(Foldable t, Monad m) =>
t a -> (a -> m b) -> m ()
forM_ ([Param LParamMem] -> [SubExp] -> [(Param LParamMem, SubExp)]
forall a b. [a] -> [b] -> [(a, b)]
zip (SegBinOp GPUMem -> [LParam GPUMem]
xParams SegBinOp GPUMem
scan) (SegBinOp GPUMem -> [SubExp]
forall rep. SegBinOp rep -> [SubExp]
segBinOpNeutral SegBinOp GPUMem
scan)) (((Param LParamMem, SubExp) -> ImpM GPUMem KernelEnv KernelOp ())
 -> ImpM GPUMem KernelEnv KernelOp ())
-> ((Param LParamMem, SubExp) -> ImpM GPUMem KernelEnv KernelOp ())
-> ImpM GPUMem KernelEnv KernelOp ()
forall a b. (a -> b) -> a -> b
$ \(Param LParamMem
p, SubExp
ne) ->
        VName
-> [TPrimExp Int64 VName]
-> SubExp
-> [TPrimExp Int64 VName]
-> ImpM GPUMem KernelEnv KernelOp ()
forall rep r op.
VName
-> [TPrimExp Int64 VName]
-> SubExp
-> [TPrimExp Int64 VName]
-> ImpM rep r op ()
copyDWIMFix (Param LParamMem -> VName
forall dec. Param dec -> VName
paramName Param LParamMem
p) [] SubExp
ne []

    sFor "chunk_id" num_chunks $ \TPrimExp Int64 VName
chunk_id -> do
      let chunk_offset :: TPrimExp Int64 VName
chunk_offset = TPrimExp Int64 VName
chunk_id TPrimExp Int64 VName
-> TPrimExp Int64 VName -> TPrimExp Int64 VName
forall a. Num a => a -> a -> a
* TPrimExp Int64 VName
stage2_tblock_size_e

      flat_idx <-
        Name
-> TPrimExp Int64 VName
-> ImpM GPUMem KernelEnv KernelOp (TV Int64)
forall {k} (t :: k) rep r op.
Name -> TExp t -> ImpM rep r op (TV t)
dPrimV Name
"flat_idx" (TPrimExp Int64 VName -> ImpM GPUMem KernelEnv KernelOp (TV Int64))
-> TPrimExp Int64 VName
-> ImpM GPUMem KernelEnv KernelOp (TV Int64)
forall a b. (a -> b) -> a -> b
$
          (TPrimExp Int64 VName
chunk_offset TPrimExp Int64 VName
-> TPrimExp Int64 VName -> TPrimExp Int64 VName
forall a. Num a => a -> a -> a
+ TPrimExp Int32 VName -> TPrimExp Int64 VName
forall {k} (t :: k) v. IntExp t => TPrimExp t v -> TPrimExp Int64 v
sExt64 (KernelConstants -> TPrimExp Int32 VName
kernelLocalThreadId KernelConstants
constants) TPrimExp Int64 VName
-> TPrimExp Int64 VName -> TPrimExp Int64 VName
forall a. Num a => a -> a -> a
+ TPrimExp Int64 VName
1) TPrimExp Int64 VName
-> TPrimExp Int64 VName -> TPrimExp Int64 VName
forall a. Num a => a -> a -> a
* TPrimExp Int64 VName
elems_per_group TPrimExp Int64 VName
-> TPrimExp Int64 VName -> TPrimExp Int64 VName
forall a. Num a => a -> a -> a
- TPrimExp Int64 VName
1
      -- Construct segment indices.
      zipWithM_ dPrimV_ gtids $ unflattenIndex dims' $ tvExp flat_idx

      forM_ (L.zip3 scans per_scan_local_arrs per_scan_pes) $
        \(scan :: SegBinOp GPUMem
scan@(SegBinOp Commutativity
_ Lambda GPUMem
scan_op [SubExp]
nes Shape
vec_shape), [VName]
local_arrs, [VName]
pes) ->
          Text
-> ImpM GPUMem KernelEnv KernelOp ()
-> ImpM GPUMem KernelEnv KernelOp ()
forall rep r op. Text -> ImpM rep r op () -> ImpM rep r op ()
sComment Text
"do one stage-2 scan chunk" (ImpM GPUMem KernelEnv KernelOp ()
 -> ImpM GPUMem KernelEnv KernelOp ())
-> ImpM GPUMem KernelEnv KernelOp ()
-> ImpM GPUMem KernelEnv KernelOp ()
forall a b. (a -> b) -> a -> b
$ do
            let (Bool
array_scan, Fence
fence, ImpM GPUMem KernelEnv KernelOp ()
barrier) = Lambda GPUMem -> (Bool, Fence, ImpM GPUMem KernelEnv KernelOp ())
barrierFor Lambda GPUMem
scan_op
                scan_x_params :: [LParam GPUMem]
scan_x_params = SegBinOp GPUMem -> [LParam GPUMem]
xParams SegBinOp GPUMem
scan
                scan_y_params :: [LParam GPUMem]
scan_y_params = SegBinOp GPUMem -> [LParam GPUMem]
yParams SegBinOp GPUMem
scan
                -- Scalar scans with a non-trivial vec_shape need per-vec-element
                -- carries reloaded from global memory, just like array scans.
                -- The scalar carry register cannot hold independent carries for
                -- each vector element across chunk iterations.
                use_global_carry :: Bool
use_global_carry = Bool
array_scan Bool -> Bool -> Bool
|| Bool -> Bool
not ([SubExp] -> Bool
forall a. [a] -> Bool
forall (t :: * -> *) a. Foldable t => t a -> Bool
null (Shape -> [SubExp]
forall d. ShapeBase d -> [d]
shapeDims Shape
vec_shape))

            Bool
-> ImpM GPUMem KernelEnv KernelOp ()
-> ImpM GPUMem KernelEnv KernelOp ()
forall (f :: * -> *). Applicative f => Bool -> f () -> f ()
when Bool
use_global_carry (KernelOp -> ImpM GPUMem KernelEnv KernelOp ()
forall op rep r. op -> ImpM rep r op ()
sOp (KernelOp -> ImpM GPUMem KernelEnv KernelOp ())
-> KernelOp -> ImpM GPUMem KernelEnv KernelOp ()
forall a b. (a -> b) -> a -> b
$ Fence -> KernelOp
Imp.Barrier Fence
Imp.FenceGlobal)

            Shape
-> ([TPrimExp Int64 VName] -> ImpM GPUMem KernelEnv KernelOp ())
-> ImpM GPUMem KernelEnv KernelOp ()
forall rep r op.
Shape
-> ([TPrimExp Int64 VName] -> ImpM rep r op ()) -> ImpM rep r op ()
sLoopNest Shape
vec_shape (([TPrimExp Int64 VName] -> ImpM GPUMem KernelEnv KernelOp ())
 -> ImpM GPUMem KernelEnv KernelOp ())
-> ([TPrimExp Int64 VName] -> ImpM GPUMem KernelEnv KernelOp ())
-> ImpM GPUMem KernelEnv KernelOp ()
forall a b. (a -> b) -> a -> b
$ \[TPrimExp Int64 VName]
vec_is -> do
              let glob_is :: [TPrimExp Int64 VName]
glob_is = (VName -> TPrimExp Int64 VName)
-> [VName] -> [TPrimExp Int64 VName]
forall a b. (a -> b) -> [a] -> [b]
map VName -> TPrimExp Int64 VName
forall a. a -> TPrimExp Int64 a
Imp.le64 [VName]
gtids [TPrimExp Int64 VName]
-> [TPrimExp Int64 VName] -> [TPrimExp Int64 VName]
forall a. [a] -> [a] -> [a]
++ [TPrimExp Int64 VName]
vec_is
                  in_bounds :: TExp Bool
in_bounds =
                    (TExp Bool -> TExp Bool -> TExp Bool) -> [TExp Bool] -> TExp Bool
forall a. (a -> a -> a) -> [a] -> a
forall (t :: * -> *) a. Foldable t => (a -> a -> a) -> t a -> a
foldl1 TExp Bool -> TExp Bool -> TExp Bool
forall v.
Eq v =>
TPrimExp Bool v -> TPrimExp Bool v -> TPrimExp Bool v
(.&&.) ([TExp Bool] -> TExp Bool) -> [TExp Bool] -> TExp Bool
forall a b. (a -> b) -> a -> b
$ (TPrimExp Int64 VName -> TPrimExp Int64 VName -> TExp Bool)
-> [TPrimExp Int64 VName] -> [TPrimExp Int64 VName] -> [TExp Bool]
forall a b c. (a -> b -> c) -> [a] -> [b] -> [c]
zipWith TPrimExp Int64 VName -> TPrimExp Int64 VName -> TExp Bool
forall {k} v (t :: k).
Eq v =>
TPrimExp t v -> TPrimExp t v -> TPrimExp Bool v
(.<.) ((VName -> TPrimExp Int64 VName)
-> [VName] -> [TPrimExp Int64 VName]
forall a b. (a -> b) -> [a] -> [b]
map VName -> TPrimExp Int64 VName
forall a. a -> TPrimExp Int64 a
Imp.le64 [VName]
gtids) [TPrimExp Int64 VName]
dims'
                  -- Element index of the last element of the previous chunk,
                  -- used to load the inter-chunk carry from global memory.
                  prev_chunk_last :: TPrimExp Int64 VName
prev_chunk_last = TPrimExp Int64 VName
chunk_offset TPrimExp Int64 VName
-> TPrimExp Int64 VName -> TPrimExp Int64 VName
forall a. Num a => a -> a -> a
* TPrimExp Int64 VName
elems_per_group TPrimExp Int64 VName
-> TPrimExp Int64 VName -> TPrimExp Int64 VName
forall a. Num a => a -> a -> a
- TPrimExp Int64 VName
1

              -- For array scans and scalar scans with non-trivial vec_shape,
              -- reload carry (xParams) from the scan output written by the
              -- previous chunk.  Thread 0 reads the last element of the
              -- previous chunk, unless a segment boundary falls between the
              -- chunks.
              Bool
-> ImpM GPUMem KernelEnv KernelOp ()
-> ImpM GPUMem KernelEnv KernelOp ()
forall (f :: * -> *). Applicative f => Bool -> f () -> f ()
when Bool
use_global_carry (ImpM GPUMem KernelEnv KernelOp ()
 -> ImpM GPUMem KernelEnv KernelOp ())
-> ImpM GPUMem KernelEnv KernelOp ()
-> ImpM GPUMem KernelEnv KernelOp ()
forall a b. (a -> b) -> a -> b
$ do
                crosses_seg <-
                  Name -> TExp Bool -> ImpM GPUMem KernelEnv KernelOp (TExp Bool)
forall {k} (t :: k) rep r op.
Name -> TExp t -> ImpM rep r op (TExp t)
dPrimVE Name
"crosses_seg" (TExp Bool -> ImpM GPUMem KernelEnv KernelOp (TExp Bool))
-> TExp Bool -> ImpM GPUMem KernelEnv KernelOp (TExp Bool)
forall a b. (a -> b) -> a -> b
$
                    case CrossesSegment
crossesSegment of
                      CrossesSegment
Nothing -> TExp Bool
forall v. TPrimExp Bool v
false
                      Just TPrimExp Int64 VName -> TPrimExp Int64 VName -> TExp Bool
f -> TPrimExp Int64 VName -> TPrimExp Int64 VName -> TExp Bool
f TPrimExp Int64 VName
prev_chunk_last (TPrimExp Int64 VName
prev_chunk_last TPrimExp Int64 VName
-> TPrimExp Int64 VName -> TPrimExp Int64 VName
forall a. Num a => a -> a -> a
+ TPrimExp Int64 VName
1)
                sIf
                  (chunk_id .>. 0 .&&. kernelLocalThreadId constants .==. 0 .&&. bNot crosses_seg)
                  ( do
                      let carry_is = [TPrimExp Int64 VName]
-> TPrimExp Int64 VName -> [TPrimExp Int64 VName]
forall num. IntegralExp num => [num] -> num -> [num]
unflattenIndex [TPrimExp Int64 VName]
dims' TPrimExp Int64 VName
prev_chunk_last
                      forM_ (zip scan_x_params pes) $ \(Param LParamMem
p, VName
pe) ->
                        VName
-> [TPrimExp Int64 VName]
-> SubExp
-> [TPrimExp Int64 VName]
-> ImpM GPUMem KernelEnv KernelOp ()
forall rep r op.
VName
-> [TPrimExp Int64 VName]
-> SubExp
-> [TPrimExp Int64 VName]
-> ImpM rep r op ()
copyDWIMFix (Param LParamMem -> VName
forall dec. Param dec -> VName
paramName Param LParamMem
p) [] (VName -> SubExp
Var VName
pe) ([TPrimExp Int64 VName]
carry_is [TPrimExp Int64 VName]
-> [TPrimExp Int64 VName] -> [TPrimExp Int64 VName]
forall a. [a] -> [a] -> [a]
++ [TPrimExp Int64 VName]
vec_is)
                  )
                  ( forM_ (zip scan_x_params nes) $ \(Param LParamMem
p, SubExp
ne) ->
                      VName
-> [TPrimExp Int64 VName]
-> SubExp
-> [TPrimExp Int64 VName]
-> ImpM GPUMem KernelEnv KernelOp ()
forall rep r op.
VName
-> [TPrimExp Int64 VName]
-> SubExp
-> [TPrimExp Int64 VName]
-> ImpM rep r op ()
copyDWIMFix (Param LParamMem -> VName
forall dec. Param dec -> VName
paramName Param LParamMem
p) [] SubExp
ne []
                  )

              -- Load the stage-1 partial-scan result for this thread's
              -- block into yParams (or the neutral element when out of
              -- bounds).
              TExp Bool
-> ImpM GPUMem KernelEnv KernelOp ()
-> ImpM GPUMem KernelEnv KernelOp ()
-> ImpM GPUMem KernelEnv KernelOp ()
forall rep r op.
TExp Bool
-> ImpM rep r op () -> ImpM rep r op () -> ImpM rep r op ()
sIf
                TExp Bool
in_bounds
                ( [(Param LParamMem, VName)]
-> ((Param LParamMem, VName) -> ImpM GPUMem KernelEnv KernelOp ())
-> ImpM GPUMem KernelEnv KernelOp ()
forall (t :: * -> *) (m :: * -> *) a b.
(Foldable t, Monad m) =>
t a -> (a -> m b) -> m ()
forM_ ([Param LParamMem] -> [VName] -> [(Param LParamMem, VName)]
forall a b. [a] -> [b] -> [(a, b)]
zip [LParam GPUMem]
[Param LParamMem]
scan_y_params [VName]
pes) (((Param LParamMem, VName) -> ImpM GPUMem KernelEnv KernelOp ())
 -> ImpM GPUMem KernelEnv KernelOp ())
-> ((Param LParamMem, VName) -> ImpM GPUMem KernelEnv KernelOp ())
-> ImpM GPUMem KernelEnv KernelOp ()
forall a b. (a -> b) -> a -> b
$ \(Param LParamMem
p, VName
pe) ->
                    VName
-> [TPrimExp Int64 VName]
-> SubExp
-> [TPrimExp Int64 VName]
-> ImpM GPUMem KernelEnv KernelOp ()
forall rep r op.
VName
-> [TPrimExp Int64 VName]
-> SubExp
-> [TPrimExp Int64 VName]
-> ImpM rep r op ()
copyDWIMFix (Param LParamMem -> VName
forall dec. Param dec -> VName
paramName Param LParamMem
p) [] (VName -> SubExp
Var VName
pe) [TPrimExp Int64 VName]
glob_is
                )
                ( [(Param LParamMem, SubExp)]
-> ((Param LParamMem, SubExp) -> ImpM GPUMem KernelEnv KernelOp ())
-> ImpM GPUMem KernelEnv KernelOp ()
forall (t :: * -> *) (m :: * -> *) a b.
(Foldable t, Monad m) =>
t a -> (a -> m b) -> m ()
forM_ ([Param LParamMem] -> [SubExp] -> [(Param LParamMem, SubExp)]
forall a b. [a] -> [b] -> [(a, b)]
zip [LParam GPUMem]
[Param LParamMem]
scan_y_params [SubExp]
nes) (((Param LParamMem, SubExp) -> ImpM GPUMem KernelEnv KernelOp ())
 -> ImpM GPUMem KernelEnv KernelOp ())
-> ((Param LParamMem, SubExp) -> ImpM GPUMem KernelEnv KernelOp ())
-> ImpM GPUMem KernelEnv KernelOp ()
forall a b. (a -> b) -> a -> b
$ \(Param LParamMem
p, SubExp
ne) ->
                    VName
-> [TPrimExp Int64 VName]
-> SubExp
-> [TPrimExp Int64 VName]
-> ImpM GPUMem KernelEnv KernelOp ()
forall rep r op.
VName
-> [TPrimExp Int64 VName]
-> SubExp
-> [TPrimExp Int64 VName]
-> ImpM rep r op ()
copyDWIMFix (Param LParamMem -> VName
forall dec. Param dec -> VName
paramName Param LParamMem
p) [] SubExp
ne []
                )

              -- Combine carry (xParams) with new value (yParams) and
              -- write the result to shared/local memory.  Thread 0
              -- incorporates the inter-chunk carry; other threads have
              -- neutral in xParams, so the op is a no-op for them.
              Names
-> Stms GPUMem
-> ImpM GPUMem KernelEnv KernelOp ()
-> ImpM GPUMem KernelEnv KernelOp ()
forall rep r op.
Names -> Stms rep -> ImpM rep r op () -> ImpM rep r op ()
compileStms Names
forall a. Monoid a => a
mempty (GBody GPUMem SubExpRes -> Stms GPUMem
forall rep res. GBody rep res -> Stms rep
bodyStms (GBody GPUMem SubExpRes -> Stms GPUMem)
-> GBody GPUMem SubExpRes -> Stms GPUMem
forall a b. (a -> b) -> a -> b
$ Lambda GPUMem -> GBody GPUMem SubExpRes
forall rep t. GLambda rep t -> Body rep
lambdaBody Lambda GPUMem
scan_op) (ImpM GPUMem KernelEnv KernelOp ()
 -> ImpM GPUMem KernelEnv KernelOp ())
-> ImpM GPUMem KernelEnv KernelOp ()
-> ImpM GPUMem KernelEnv KernelOp ()
forall a b. (a -> b) -> a -> b
$
                [(VName, SubExp)]
-> ((VName, SubExp) -> ImpM GPUMem KernelEnv KernelOp ())
-> ImpM GPUMem KernelEnv KernelOp ()
forall (t :: * -> *) (m :: * -> *) a b.
(Foldable t, Monad m) =>
t a -> (a -> m b) -> m ()
forM_ ([VName] -> [SubExp] -> [(VName, SubExp)]
forall a b. [a] -> [b] -> [(a, b)]
zip [VName]
local_arrs ([SubExp] -> [(VName, SubExp)]) -> [SubExp] -> [(VName, SubExp)]
forall a b. (a -> b) -> a -> b
$ (SubExpRes -> SubExp) -> [SubExpRes] -> [SubExp]
forall a b. (a -> b) -> [a] -> [b]
map SubExpRes -> SubExp
resSubExp ([SubExpRes] -> [SubExp]) -> [SubExpRes] -> [SubExp]
forall a b. (a -> b) -> a -> b
$ GBody GPUMem SubExpRes -> [SubExpRes]
forall rep res. GBody rep res -> [res]
bodyResult (GBody GPUMem SubExpRes -> [SubExpRes])
-> GBody GPUMem SubExpRes -> [SubExpRes]
forall a b. (a -> b) -> a -> b
$ Lambda GPUMem -> GBody GPUMem SubExpRes
forall rep t. GLambda rep t -> Body rep
lambdaBody Lambda GPUMem
scan_op) (((VName, SubExp) -> ImpM GPUMem KernelEnv KernelOp ())
 -> ImpM GPUMem KernelEnv KernelOp ())
-> ((VName, SubExp) -> ImpM GPUMem KernelEnv KernelOp ())
-> ImpM GPUMem KernelEnv KernelOp ()
forall a b. (a -> b) -> a -> b
$
                  \(VName
arr, SubExp
se) ->
                    VName
-> [TPrimExp Int64 VName]
-> SubExp
-> [TPrimExp Int64 VName]
-> ImpM GPUMem KernelEnv KernelOp ()
forall rep r op.
VName
-> [TPrimExp Int64 VName]
-> SubExp
-> [TPrimExp Int64 VName]
-> ImpM rep r op ()
copyDWIMFix VName
arr [KernelConstants -> TPrimExp Int64 VName
localArrayIndex KernelConstants
constants] SubExp
se []

              KernelOp -> ImpM GPUMem KernelEnv KernelOp ()
forall op rep r. op -> ImpM rep r op ()
sOp (KernelOp -> ImpM GPUMem KernelEnv KernelOp ())
-> KernelOp -> ImpM GPUMem KernelEnv KernelOp ()
forall a b. (a -> b) -> a -> b
$ Fence -> KernelOp
Imp.ErrorSync Fence
fence

              -- crossesSegment' maps block-local thread IDs to element
              -- indices, adjusting for the current chunk offset.
              let crossesSegment' :: Maybe (TPrimExp Int32 VName -> TPrimExp Int32 VName -> TExp Bool)
crossesSegment' =
                    CrossesSegment
crossesSegment CrossesSegment
-> ((TPrimExp Int64 VName -> TPrimExp Int64 VName -> TExp Bool)
    -> Maybe
         (TPrimExp Int32 VName -> TPrimExp Int32 VName -> TExp Bool))
-> Maybe
     (TPrimExp Int32 VName -> TPrimExp Int32 VName -> TExp Bool)
forall a b. Maybe a -> (a -> Maybe b) -> Maybe b
forall (m :: * -> *) a b. Monad m => m a -> (a -> m b) -> m b
>>= \TPrimExp Int64 VName -> TPrimExp Int64 VName -> TExp Bool
f ->
                      (TPrimExp Int32 VName -> TPrimExp Int32 VName -> TExp Bool)
-> Maybe
     (TPrimExp Int32 VName -> TPrimExp Int32 VName -> TExp Bool)
forall a. a -> Maybe a
Just ((TPrimExp Int32 VName -> TPrimExp Int32 VName -> TExp Bool)
 -> Maybe
      (TPrimExp Int32 VName -> TPrimExp Int32 VName -> TExp Bool))
-> (TPrimExp Int32 VName -> TPrimExp Int32 VName -> TExp Bool)
-> Maybe
     (TPrimExp Int32 VName -> TPrimExp Int32 VName -> TExp Bool)
forall a b. (a -> b) -> a -> b
$ \TPrimExp Int32 VName
from TPrimExp Int32 VName
to ->
                        TPrimExp Int64 VName -> TPrimExp Int64 VName -> TExp Bool
f
                          ((TPrimExp Int64 VName
chunk_offset TPrimExp Int64 VName
-> TPrimExp Int64 VName -> TPrimExp Int64 VName
forall a. Num a => a -> a -> a
+ TPrimExp Int32 VName -> TPrimExp Int64 VName
forall {k} (t :: k) v. IntExp t => TPrimExp t v -> TPrimExp Int64 v
sExt64 TPrimExp Int32 VName
from TPrimExp Int64 VName
-> TPrimExp Int64 VName -> TPrimExp Int64 VName
forall a. Num a => a -> a -> a
+ TPrimExp Int64 VName
1) TPrimExp Int64 VName
-> TPrimExp Int64 VName -> TPrimExp Int64 VName
forall a. Num a => a -> a -> a
* TPrimExp Int64 VName
elems_per_group TPrimExp Int64 VName
-> TPrimExp Int64 VName -> TPrimExp Int64 VName
forall a. Num a => a -> a -> a
- TPrimExp Int64 VName
1)
                          ((TPrimExp Int64 VName
chunk_offset TPrimExp Int64 VName
-> TPrimExp Int64 VName -> TPrimExp Int64 VName
forall a. Num a => a -> a -> a
+ TPrimExp Int32 VName -> TPrimExp Int64 VName
forall {k} (t :: k) v. IntExp t => TPrimExp t v -> TPrimExp Int64 v
sExt64 TPrimExp Int32 VName
to TPrimExp Int64 VName
-> TPrimExp Int64 VName -> TPrimExp Int64 VName
forall a. Num a => a -> a -> a
+ TPrimExp Int64 VName
1) TPrimExp Int64 VName
-> TPrimExp Int64 VName -> TPrimExp Int64 VName
forall a. Num a => a -> a -> a
* TPrimExp Int64 VName
elems_per_group TPrimExp Int64 VName
-> TPrimExp Int64 VName -> TPrimExp Int64 VName
forall a. Num a => a -> a -> a
- TPrimExp Int64 VName
1)

              scan_op_renamed <- Lambda GPUMem -> ImpM GPUMem KernelEnv KernelOp (Lambda GPUMem)
forall rep (m :: * -> *) t.
(Renameable rep, MonadFreshNames m, Rename t) =>
GLambda rep t -> m (GLambda rep t)
renameLambda Lambda GPUMem
scan_op
              blockScan
                crossesSegment'
                (sExt64 $ tvExp stage1_num_threads)
                (sExt64 $ kernelBlockSize constants)
                scan_op_renamed
                local_arrs

              sComment "threads in bounds write scanned carries" $
                sWhen in_bounds $
                  forM_ (zip pes local_arrs) $ \(VName
pe, VName
arr) ->
                    VName
-> [TPrimExp Int64 VName]
-> SubExp
-> [TPrimExp Int64 VName]
-> ImpM GPUMem KernelEnv KernelOp ()
forall rep r op.
VName
-> [TPrimExp Int64 VName]
-> SubExp
-> [TPrimExp Int64 VName]
-> ImpM rep r op ()
copyDWIMFix
                      VName
pe
                      [TPrimExp Int64 VName]
glob_is
                      (VName -> SubExp
Var VName
arr)
                      [KernelConstants -> TPrimExp Int64 VName
localArrayIndex KernelConstants
constants]

              barrier

              -- For scalar scans with trivial vec_shape (no vec loop), update
              -- the carry register (xParams) so the next chunk can use it.
              -- For array scans and scalar scans with non-trivial vec_shape,
              -- the carry is reloaded from global memory at the start of each
              -- chunk (above), so no register update is needed here.
              unless use_global_carry $ do
                let next_chunk_start = TPrimExp Int64 VName
chunk_offset TPrimExp Int64 VName
-> TPrimExp Int64 VName -> TPrimExp Int64 VName
forall a. Num a => a -> a -> a
+ TPrimExp Int64 VName
stage2_tblock_size_e
                crosses_seg2 <-
                  dPrimVE "crosses_seg" $
                    case crossesSegment of
                      CrossesSegment
Nothing -> TExp Bool
forall v. TPrimExp Bool v
false
                      Just TPrimExp Int64 VName -> TPrimExp Int64 VName -> TExp Bool
f ->
                        TPrimExp Int64 VName -> TPrimExp Int64 VName -> TExp Bool
f
                          (TPrimExp Int64 VName
next_chunk_start TPrimExp Int64 VName
-> TPrimExp Int64 VName -> TPrimExp Int64 VName
forall a. Num a => a -> a -> a
* TPrimExp Int64 VName
elems_per_group TPrimExp Int64 VName
-> TPrimExp Int64 VName -> TPrimExp Int64 VName
forall a. Num a => a -> a -> a
- TPrimExp Int64 VName
1)
                          ((TPrimExp Int64 VName
next_chunk_start TPrimExp Int64 VName
-> TPrimExp Int64 VName -> TPrimExp Int64 VName
forall a. Num a => a -> a -> a
+ TPrimExp Int64 VName
1) TPrimExp Int64 VName
-> TPrimExp Int64 VName -> TPrimExp Int64 VName
forall a. Num a => a -> a -> a
* TPrimExp Int64 VName
elems_per_group TPrimExp Int64 VName
-> TPrimExp Int64 VName -> TPrimExp Int64 VName
forall a. Num a => a -> a -> a
- TPrimExp Int64 VName
1)
                let should_load_carry =
                      KernelConstants -> TPrimExp Int32 VName
kernelLocalThreadId KernelConstants
constants TPrimExp Int32 VName -> TPrimExp Int32 VName -> TExp Bool
forall {k} v (t :: k).
Eq v =>
TPrimExp t v -> TPrimExp t v -> TPrimExp Bool v
.==. TPrimExp Int32 VName
0 TExp Bool -> TExp Bool -> TExp Bool
forall v.
Eq v =>
TPrimExp Bool v -> TPrimExp Bool v -> TPrimExp Bool v
.&&. TExp Bool -> TExp Bool
forall v. TPrimExp Bool v -> TPrimExp Bool v
bNot TExp Bool
crosses_seg2
                    load_carry =
                      [(VName, Param LParamMem)]
-> ((VName, Param LParamMem) -> ImpM GPUMem KernelEnv KernelOp ())
-> ImpM GPUMem KernelEnv KernelOp ()
forall (t :: * -> *) (m :: * -> *) a b.
(Foldable t, Monad m) =>
t a -> (a -> m b) -> m ()
forM_ ([VName] -> [Param LParamMem] -> [(VName, Param LParamMem)]
forall a b. [a] -> [b] -> [(a, b)]
zip [VName]
local_arrs [LParam GPUMem]
[Param LParamMem]
scan_x_params) (((VName, Param LParamMem) -> ImpM GPUMem KernelEnv KernelOp ())
 -> ImpM GPUMem KernelEnv KernelOp ())
-> ((VName, Param LParamMem) -> ImpM GPUMem KernelEnv KernelOp ())
-> ImpM GPUMem KernelEnv KernelOp ()
forall a b. (a -> b) -> a -> b
$ \(VName
arr, Param LParamMem
p) ->
                        VName
-> [TPrimExp Int64 VName]
-> SubExp
-> [TPrimExp Int64 VName]
-> ImpM GPUMem KernelEnv KernelOp ()
forall rep r op.
VName
-> [TPrimExp Int64 VName]
-> SubExp
-> [TPrimExp Int64 VName]
-> ImpM rep r op ()
copyDWIMFix
                          (Param LParamMem -> VName
forall dec. Param dec -> VName
paramName Param LParamMem
p)
                          []
                          (VName -> SubExp
Var VName
arr)
                          [TPrimExp Int64 VName -> TPrimExp Int64 VName
forall {k} (t :: k) v. IntExp t => TPrimExp t v -> TPrimExp Int64 v
sExt64 (KernelConstants -> TPrimExp Int64 VName
kernelBlockSize KernelConstants
constants) TPrimExp Int64 VName
-> TPrimExp Int64 VName -> TPrimExp Int64 VName
forall a. Num a => a -> a -> a
- TPrimExp Int64 VName
1]
                    load_neutral =
                      [(SubExp, Param LParamMem)]
-> ((SubExp, Param LParamMem) -> ImpM GPUMem KernelEnv KernelOp ())
-> ImpM GPUMem KernelEnv KernelOp ()
forall (t :: * -> *) (m :: * -> *) a b.
(Foldable t, Monad m) =>
t a -> (a -> m b) -> m ()
forM_ ([SubExp] -> [Param LParamMem] -> [(SubExp, Param LParamMem)]
forall a b. [a] -> [b] -> [(a, b)]
zip [SubExp]
nes [LParam GPUMem]
[Param LParamMem]
scan_x_params) (((SubExp, Param LParamMem) -> ImpM GPUMem KernelEnv KernelOp ())
 -> ImpM GPUMem KernelEnv KernelOp ())
-> ((SubExp, Param LParamMem) -> ImpM GPUMem KernelEnv KernelOp ())
-> ImpM GPUMem KernelEnv KernelOp ()
forall a b. (a -> b) -> a -> b
$ \(SubExp
ne, Param LParamMem
p) ->
                        VName
-> [TPrimExp Int64 VName]
-> SubExp
-> [TPrimExp Int64 VName]
-> ImpM GPUMem KernelEnv KernelOp ()
forall rep r op.
VName
-> [TPrimExp Int64 VName]
-> SubExp
-> [TPrimExp Int64 VName]
-> ImpM rep r op ()
copyDWIMFix (Param LParamMem -> VName
forall dec. Param dec -> VName
paramName Param LParamMem
p) [] SubExp
ne []
                sWhen should_load_carry load_carry
                sUnless should_load_carry load_neutral

              barrier

scanStage3 ::
  Pat LetDecMem ->
  [VName] ->
  [Maybe VName] ->
  Count NumBlocks SubExp ->
  Count BlockSize SubExp ->
  Imp.TExp Int64 ->
  CrossesSegment ->
  SegSpace ->
  [SegBinOp GPUMem] ->
  SegPostOp GPUMem ->
  CallKernelGen ()
scanStage3 :: Pat LParamMem
-> [VName]
-> [Maybe VName]
-> Count NumBlocks SubExp
-> Count BlockSize SubExp
-> TPrimExp Int64 VName
-> CrossesSegment
-> SegSpace
-> [SegBinOp GPUMem]
-> SegPostOp GPUMem
-> CallKernelGen ()
scanStage3 Pat LParamMem
pat [VName]
scan_out [Maybe VName]
map_out Count NumBlocks SubExp
num_tblocks Count BlockSize SubExp
tblock_size TPrimExp Int64 VName
elems_per_group CrossesSegment
crossesSegment SegSpace
space [SegBinOp GPUMem]
scans SegPostOp GPUMem
post_op = do
  let tblock_size' :: Count BlockSize (TPrimExp Int64 VName)
tblock_size' = (SubExp -> TPrimExp Int64 VName)
-> Count BlockSize SubExp -> Count BlockSize (TPrimExp Int64 VName)
forall a b. (a -> b) -> Count BlockSize a -> Count BlockSize b
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap SubExp -> TPrimExp Int64 VName
pe64 Count BlockSize SubExp
tblock_size
      ([VName]
gtids, [SubExp]
dims) = [(VName, SubExp)] -> ([VName], [SubExp])
forall a b. [(a, b)] -> ([a], [b])
unzip ([(VName, SubExp)] -> ([VName], [SubExp]))
-> [(VName, SubExp)] -> ([VName], [SubExp])
forall a b. (a -> b) -> a -> b
$ SegSpace -> [(VName, SubExp)]
unSegSpace SegSpace
space
      dims' :: [TPrimExp Int64 VName]
dims' = (SubExp -> TPrimExp Int64 VName)
-> [SubExp] -> [TPrimExp Int64 VName]
forall a b. (a -> b) -> [a] -> [b]
map SubExp -> TPrimExp Int64 VName
pe64 [SubExp]
dims
  required_groups <-
    Name
-> TPrimExp Int32 VName
-> ImpM GPUMem HostEnv HostOp (TPrimExp Int32 VName)
forall {k} (t :: k) rep r op.
Name -> TExp t -> ImpM rep r op (TExp t)
dPrimVE Name
"required_groups" (TPrimExp Int32 VName
 -> ImpM GPUMem HostEnv HostOp (TPrimExp Int32 VName))
-> TPrimExp Int32 VName
-> ImpM GPUMem HostEnv HostOp (TPrimExp Int32 VName)
forall a b. (a -> b) -> a -> b
$
      TPrimExp Int64 VName -> TPrimExp Int32 VName
forall {k} (t :: k) v. IntExp t => TPrimExp t v -> TPrimExp Int32 v
sExt32 (TPrimExp Int64 VName -> TPrimExp Int32 VName)
-> TPrimExp Int64 VName -> TPrimExp Int32 VName
forall a b. (a -> b) -> a -> b
$
        [TPrimExp Int64 VName] -> TPrimExp Int64 VName
forall a. Num a => [a] -> a
forall (t :: * -> *) a. (Foldable t, Num a) => t a -> a
product [TPrimExp Int64 VName]
dims' TPrimExp Int64 VName
-> TPrimExp Int64 VName -> TPrimExp Int64 VName
forall e. IntegralExp e => e -> e -> e
`ceilDiv` TPrimExp Int64 VName -> TPrimExp Int64 VName
forall {k} (t :: k) v. IntExp t => TPrimExp t v -> TPrimExp Int64 v
sExt64 (Count BlockSize (TPrimExp Int64 VName) -> TPrimExp Int64 VName
forall {k} (u :: k) e. Count u e -> e
unCount Count BlockSize (TPrimExp Int64 VName)
tblock_size')

  sKernelThread "scan_stage3" (segFlat space) (defKernelAttrs num_tblocks tblock_size) $
    virtualiseBlocks SegVirt required_groups $ \TPrimExp Int32 VName
virt_tblock_id -> do
      constants <- KernelEnv -> KernelConstants
kernelConstants (KernelEnv -> KernelConstants)
-> ImpM GPUMem KernelEnv KernelOp KernelEnv
-> ImpM GPUMem KernelEnv KernelOp KernelConstants
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> ImpM GPUMem KernelEnv KernelOp KernelEnv
forall rep r op. ImpM rep r op r
askEnv

      -- Compute our logical index.
      flat_idx <-
        dPrimVE "flat_idx" $
          sExt64 virt_tblock_id * sExt64 (unCount tblock_size')
            + sExt64 (kernelLocalThreadId constants)
      zipWithM_ dPrimV_ gtids $ unflattenIndex dims' flat_idx

      -- Figure out which group this element was originally in.
      orig_group <- dPrimV "orig_group" $ flat_idx `quot` elems_per_group
      -- Then the index of the carry-in of the preceding group.
      carry_in_flat_idx <-
        dPrimV "carry_in_flat_idx" $
          tvExp orig_group * elems_per_group - 1
      -- Figure out the logical index of the carry-in.
      let carry_in_idx = [TPrimExp Int64 VName]
-> TPrimExp Int64 VName -> [TPrimExp Int64 VName]
forall num. IntegralExp num => [num] -> num -> [num]
unflattenIndex [TPrimExp Int64 VName]
dims' (TPrimExp Int64 VName -> [TPrimExp Int64 VName])
-> TPrimExp Int64 VName -> [TPrimExp Int64 VName]
forall a b. (a -> b) -> a -> b
$ TV Int64 -> TPrimExp Int64 VName
forall {k} (t :: k). TV t -> TExp t
tvExp TV Int64
carry_in_flat_idx

      -- Apply the carry if we are not in the scan results for the first
      -- group, and are not the last element in such a group (because
      -- then the carry was updated in stage 2), and we are not crossing
      -- a segment boundary.
      let in_bounds =
            (TExp Bool -> TExp Bool -> TExp Bool) -> [TExp Bool] -> TExp Bool
forall a. (a -> a -> a) -> [a] -> a
forall (t :: * -> *) a. Foldable t => (a -> a -> a) -> t a -> a
foldl1 TExp Bool -> TExp Bool -> TExp Bool
forall v.
Eq v =>
TPrimExp Bool v -> TPrimExp Bool v -> TPrimExp Bool v
(.&&.) ([TExp Bool] -> TExp Bool) -> [TExp Bool] -> TExp Bool
forall a b. (a -> b) -> a -> b
$ (TPrimExp Int64 VName -> TPrimExp Int64 VName -> TExp Bool)
-> [TPrimExp Int64 VName] -> [TPrimExp Int64 VName] -> [TExp Bool]
forall a b c. (a -> b -> c) -> [a] -> [b] -> [c]
zipWith TPrimExp Int64 VName -> TPrimExp Int64 VName -> TExp Bool
forall {k} v (t :: k).
Eq v =>
TPrimExp t v -> TPrimExp t v -> TPrimExp Bool v
(.<.) ((VName -> TPrimExp Int64 VName)
-> [VName] -> [TPrimExp Int64 VName]
forall a b. (a -> b) -> [a] -> [b]
map VName -> TPrimExp Int64 VName
forall a. a -> TPrimExp Int64 a
Imp.le64 [VName]
gtids) [TPrimExp Int64 VName]
dims'
          crosses_segment =
            TExp Bool -> Maybe (TExp Bool) -> TExp Bool
forall a. a -> Maybe a -> a
fromMaybe TExp Bool
forall v. TPrimExp Bool v
false (Maybe (TExp Bool) -> TExp Bool) -> Maybe (TExp Bool) -> TExp Bool
forall a b. (a -> b) -> a -> b
$
              CrossesSegment
crossesSegment
                CrossesSegment
-> Maybe (TPrimExp Int64 VName)
-> Maybe (TPrimExp Int64 VName -> TExp Bool)
forall a b. Maybe (a -> b) -> Maybe a -> Maybe b
forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> TPrimExp Int64 VName -> Maybe (TPrimExp Int64 VName)
forall a. a -> Maybe a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (TV Int64 -> TPrimExp Int64 VName
forall {k} (t :: k). TV t -> TExp t
tvExp TV Int64
carry_in_flat_idx)
                Maybe (TPrimExp Int64 VName -> TExp Bool)
-> Maybe (TPrimExp Int64 VName) -> Maybe (TExp Bool)
forall a b. Maybe (a -> b) -> Maybe a -> Maybe b
forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> TPrimExp Int64 VName -> Maybe (TPrimExp Int64 VName)
forall a. a -> Maybe a
forall (f :: * -> *) a. Applicative f => a -> f a
pure TPrimExp Int64 VName
flat_idx
          is_a_carry = TPrimExp Int64 VName
flat_idx TPrimExp Int64 VName -> TPrimExp Int64 VName -> TExp Bool
forall {k} v (t :: k).
Eq v =>
TPrimExp t v -> TPrimExp t v -> TPrimExp Bool v
.==. (TV Int64 -> TPrimExp Int64 VName
forall {k} (t :: k). TV t -> TExp t
tvExp TV Int64
orig_group TPrimExp Int64 VName
-> TPrimExp Int64 VName -> TPrimExp Int64 VName
forall a. Num a => a -> a -> a
+ TPrimExp Int64 VName
1) TPrimExp Int64 VName
-> TPrimExp Int64 VName -> TPrimExp Int64 VName
forall a. Num a => a -> a -> a
* TPrimExp Int64 VName
elems_per_group TPrimExp Int64 VName
-> TPrimExp Int64 VName -> TPrimExp Int64 VName
forall a. Num a => a -> a -> a
- TPrimExp Int64 VName
1
          no_carry_in = TV Int64 -> TPrimExp Int64 VName
forall {k} (t :: k). TV t -> TExp t
tvExp TV Int64
orig_group TPrimExp Int64 VName -> TPrimExp Int64 VName -> TExp Bool
forall {k} v (t :: k).
Eq v =>
TPrimExp t v -> TPrimExp t v -> TPrimExp Bool v
.==. TPrimExp Int64 VName
0 TExp Bool -> TExp Bool -> TExp Bool
forall v.
Eq v =>
TPrimExp Bool v -> TPrimExp Bool v -> TPrimExp Bool v
.||. TExp Bool
is_a_carry TExp Bool -> TExp Bool -> TExp Bool
forall v.
Eq v =>
TPrimExp Bool v -> TPrimExp Bool v -> TPrimExp Bool v
.||. TExp Bool
crosses_segment

      let per_scan_pes = [SegBinOp GPUMem] -> [VName] -> [[VName]]
forall rep a. [SegBinOp rep] -> [a] -> [[a]]
segBinOpChunks [SegBinOp GPUMem]
scans [VName]
scan_out
      sWhen in_bounds $ do
        sUnless no_carry_in $
          forM_ (zip per_scan_pes scans) $
            \([VName]
pes, SegBinOp Commutativity
_ Lambda GPUMem
scan_op [SubExp]
nes Shape
vec_shape) -> do
              Maybe (Exp GPUMem)
-> Scope GPUMem -> ImpM GPUMem KernelEnv KernelOp ()
forall rep (inner :: * -> *) r op.
Mem rep inner =>
Maybe (Exp rep) -> Scope rep -> ImpM rep r op ()
dScope Maybe (Exp GPUMem)
forall a. Maybe a
Nothing (Scope GPUMem -> ImpM GPUMem KernelEnv KernelOp ())
-> Scope GPUMem -> ImpM GPUMem KernelEnv KernelOp ()
forall a b. (a -> b) -> a -> b
$ [Param LParamMem] -> Scope GPUMem
forall rep dec. (LParamInfo rep ~ dec) => [Param dec] -> Scope rep
scopeOfLParams ([Param LParamMem] -> Scope GPUMem)
-> [Param LParamMem] -> Scope GPUMem
forall a b. (a -> b) -> a -> b
$ Lambda GPUMem -> [LParam GPUMem]
forall rep t. GLambda rep t -> [LParam rep]
lambdaParams Lambda GPUMem
scan_op
              let ([Param LParamMem]
scan_x_params, [Param LParamMem]
scan_y_params) =
                    Int -> [Param LParamMem] -> ([Param LParamMem], [Param LParamMem])
forall a. Int -> [a] -> ([a], [a])
splitAt ([SubExp] -> Int
forall a. [a] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length [SubExp]
nes) ([Param LParamMem] -> ([Param LParamMem], [Param LParamMem]))
-> [Param LParamMem] -> ([Param LParamMem], [Param LParamMem])
forall a b. (a -> b) -> a -> b
$ Lambda GPUMem -> [LParam GPUMem]
forall rep t. GLambda rep t -> [LParam rep]
lambdaParams Lambda GPUMem
scan_op

              Shape
-> ([TPrimExp Int64 VName] -> ImpM GPUMem KernelEnv KernelOp ())
-> ImpM GPUMem KernelEnv KernelOp ()
forall rep r op.
Shape
-> ([TPrimExp Int64 VName] -> ImpM rep r op ()) -> ImpM rep r op ()
sLoopNest Shape
vec_shape (([TPrimExp Int64 VName] -> ImpM GPUMem KernelEnv KernelOp ())
 -> ImpM GPUMem KernelEnv KernelOp ())
-> ([TPrimExp Int64 VName] -> ImpM GPUMem KernelEnv KernelOp ())
-> ImpM GPUMem KernelEnv KernelOp ()
forall a b. (a -> b) -> a -> b
$ \[TPrimExp Int64 VName]
vec_is -> do
                [(Param LParamMem, VName)]
-> ((Param LParamMem, VName) -> ImpM GPUMem KernelEnv KernelOp ())
-> ImpM GPUMem KernelEnv KernelOp ()
forall (t :: * -> *) (m :: * -> *) a b.
(Foldable t, Monad m) =>
t a -> (a -> m b) -> m ()
forM_ ([Param LParamMem] -> [VName] -> [(Param LParamMem, VName)]
forall a b. [a] -> [b] -> [(a, b)]
zip [Param LParamMem]
scan_x_params [VName]
pes) (((Param LParamMem, VName) -> ImpM GPUMem KernelEnv KernelOp ())
 -> ImpM GPUMem KernelEnv KernelOp ())
-> ((Param LParamMem, VName) -> ImpM GPUMem KernelEnv KernelOp ())
-> ImpM GPUMem KernelEnv KernelOp ()
forall a b. (a -> b) -> a -> b
$ \(Param LParamMem
p, VName
pe) ->
                  VName
-> [TPrimExp Int64 VName]
-> SubExp
-> [TPrimExp Int64 VName]
-> ImpM GPUMem KernelEnv KernelOp ()
forall rep r op.
VName
-> [TPrimExp Int64 VName]
-> SubExp
-> [TPrimExp Int64 VName]
-> ImpM rep r op ()
copyDWIMFix
                    (Param LParamMem -> VName
forall dec. Param dec -> VName
paramName Param LParamMem
p)
                    []
                    (VName -> SubExp
Var VName
pe)
                    ([TPrimExp Int64 VName]
carry_in_idx [TPrimExp Int64 VName]
-> [TPrimExp Int64 VName] -> [TPrimExp Int64 VName]
forall a. [a] -> [a] -> [a]
++ [TPrimExp Int64 VName]
vec_is)

                [(Param LParamMem, VName)]
-> ((Param LParamMem, VName) -> ImpM GPUMem KernelEnv KernelOp ())
-> ImpM GPUMem KernelEnv KernelOp ()
forall (t :: * -> *) (m :: * -> *) a b.
(Foldable t, Monad m) =>
t a -> (a -> m b) -> m ()
forM_ ([Param LParamMem] -> [VName] -> [(Param LParamMem, VName)]
forall a b. [a] -> [b] -> [(a, b)]
zip [Param LParamMem]
scan_y_params [VName]
pes) (((Param LParamMem, VName) -> ImpM GPUMem KernelEnv KernelOp ())
 -> ImpM GPUMem KernelEnv KernelOp ())
-> ((Param LParamMem, VName) -> ImpM GPUMem KernelEnv KernelOp ())
-> ImpM GPUMem KernelEnv KernelOp ()
forall a b. (a -> b) -> a -> b
$ \(Param LParamMem
p, VName
pe) ->
                  VName
-> [TPrimExp Int64 VName]
-> SubExp
-> [TPrimExp Int64 VName]
-> ImpM GPUMem KernelEnv KernelOp ()
forall rep r op.
VName
-> [TPrimExp Int64 VName]
-> SubExp
-> [TPrimExp Int64 VName]
-> ImpM rep r op ()
copyDWIMFix
                    (Param LParamMem -> VName
forall dec. Param dec -> VName
paramName Param LParamMem
p)
                    []
                    (VName -> SubExp
Var VName
pe)
                    ((VName -> TPrimExp Int64 VName)
-> [VName] -> [TPrimExp Int64 VName]
forall a b. (a -> b) -> [a] -> [b]
map VName -> TPrimExp Int64 VName
forall a. a -> TPrimExp Int64 a
Imp.le64 [VName]
gtids [TPrimExp Int64 VName]
-> [TPrimExp Int64 VName] -> [TPrimExp Int64 VName]
forall a. [a] -> [a] -> [a]
++ [TPrimExp Int64 VName]
vec_is)

                [Param LParamMem]
-> GBody GPUMem SubExpRes -> ImpM GPUMem KernelEnv KernelOp ()
forall dec rep r op. [Param dec] -> Body rep -> ImpM rep r op ()
compileBody' [Param LParamMem]
scan_x_params (GBody GPUMem SubExpRes -> ImpM GPUMem KernelEnv KernelOp ())
-> GBody GPUMem SubExpRes -> ImpM GPUMem KernelEnv KernelOp ()
forall a b. (a -> b) -> a -> b
$ Lambda GPUMem -> GBody GPUMem SubExpRes
forall rep t. GLambda rep t -> Body rep
lambdaBody Lambda GPUMem
scan_op

                [(Param LParamMem, VName)]
-> ((Param LParamMem, VName) -> ImpM GPUMem KernelEnv KernelOp ())
-> ImpM GPUMem KernelEnv KernelOp ()
forall (t :: * -> *) (m :: * -> *) a b.
(Foldable t, Monad m) =>
t a -> (a -> m b) -> m ()
forM_ ([Param LParamMem] -> [VName] -> [(Param LParamMem, VName)]
forall a b. [a] -> [b] -> [(a, b)]
zip [Param LParamMem]
scan_x_params [VName]
pes) (((Param LParamMem, VName) -> ImpM GPUMem KernelEnv KernelOp ())
 -> ImpM GPUMem KernelEnv KernelOp ())
-> ((Param LParamMem, VName) -> ImpM GPUMem KernelEnv KernelOp ())
-> ImpM GPUMem KernelEnv KernelOp ()
forall a b. (a -> b) -> a -> b
$ \(Param LParamMem
p, VName
pe) ->
                  VName
-> [TPrimExp Int64 VName]
-> SubExp
-> [TPrimExp Int64 VName]
-> ImpM GPUMem KernelEnv KernelOp ()
forall rep r op.
VName
-> [TPrimExp Int64 VName]
-> SubExp
-> [TPrimExp Int64 VName]
-> ImpM rep r op ()
copyDWIMFix
                    VName
pe
                    ((VName -> TPrimExp Int64 VName)
-> [VName] -> [TPrimExp Int64 VName]
forall a b. (a -> b) -> [a] -> [b]
map VName -> TPrimExp Int64 VName
forall a. a -> TPrimExp Int64 a
Imp.le64 [VName]
gtids [TPrimExp Int64 VName]
-> [TPrimExp Int64 VName] -> [TPrimExp Int64 VName]
forall a. [a] -> [a] -> [a]
++ [TPrimExp Int64 VName]
vec_is)
                    (VName -> SubExp
Var (VName -> SubExp) -> VName -> SubExp
forall a b. (a -> b) -> a -> b
$ Param LParamMem -> VName
forall dec. Param dec -> VName
paramName Param LParamMem
p)
                    []

      sOp $ Imp.Barrier Imp.FenceLocal

      if isIdentityPostOp post_op
        then pure ()
        else do
          sWhen in_bounds $ do
            let (scan_pars, map_pars) = splitAt (segBinOpResults scans) $ lambdaParams $ segPostOpLambda post_op
            dScope Nothing $
              scopeOfLParams $
                lambdaParams $
                  segPostOpLambda post_op

            sComment "bind scan results to post lambda params" $ do
              forM_ (zip scan_pars scan_out) $ \(Param LParamMem
par, VName
acc) ->
                VName
-> [TPrimExp Int64 VName]
-> SubExp
-> [TPrimExp Int64 VName]
-> ImpM GPUMem KernelEnv KernelOp ()
forall rep r op.
VName
-> [TPrimExp Int64 VName]
-> SubExp
-> [TPrimExp Int64 VName]
-> ImpM rep r op ()
copyDWIMFix (Param LParamMem -> VName
forall dec. Param dec -> VName
paramName Param LParamMem
par) [] (VName -> SubExp
Var VName
acc) ((VName -> TPrimExp Int64 VName)
-> [VName] -> [TPrimExp Int64 VName]
forall a b. (a -> b) -> [a] -> [b]
map VName -> TPrimExp Int64 VName
forall a. a -> TPrimExp Int64 a
Imp.le64 [VName]
gtids)

            sComment "bind map results to post lamda params" $
              forM_ (zip map_pars map_out) $ \(Param LParamMem
par, Maybe VName
out) -> do
                ImpM GPUMem KernelEnv KernelOp ()
-> (VName -> ImpM GPUMem KernelEnv KernelOp ())
-> Maybe VName
-> ImpM GPUMem KernelEnv KernelOp ()
forall b a. b -> (a -> b) -> Maybe a -> b
maybe
                  (() -> ImpM GPUMem KernelEnv KernelOp ()
forall a. a -> ImpM GPUMem KernelEnv KernelOp a
forall (f :: * -> *) a. Applicative f => a -> f a
pure ())
                  ( \VName
o ->
                      VName
-> [TPrimExp Int64 VName]
-> SubExp
-> [TPrimExp Int64 VName]
-> ImpM GPUMem KernelEnv KernelOp ()
forall rep r op.
VName
-> [TPrimExp Int64 VName]
-> SubExp
-> [TPrimExp Int64 VName]
-> ImpM rep r op ()
copyDWIMFix (Param LParamMem -> VName
forall dec. Param dec -> VName
paramName Param LParamMem
par) [] (VName -> SubExp
Var VName
o) ((VName -> TPrimExp Int64 VName)
-> [VName] -> [TPrimExp Int64 VName]
forall a b. (a -> b) -> [a] -> [b]
map VName -> TPrimExp Int64 VName
forall a. a -> TPrimExp Int64 a
Imp.le64 [VName]
gtids)
                  )
                  Maybe VName
out

            let res = (SubExpRes -> SubExp) -> [SubExpRes] -> [SubExp]
forall a b. (a -> b) -> [a] -> [b]
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap SubExpRes -> SubExp
resSubExp ([SubExpRes] -> [SubExp]) -> [SubExpRes] -> [SubExp]
forall a b. (a -> b) -> a -> b
$ GBody GPUMem SubExpRes -> [SubExpRes]
forall rep res. GBody rep res -> [res]
bodyResult (GBody GPUMem SubExpRes -> [SubExpRes])
-> GBody GPUMem SubExpRes -> [SubExpRes]
forall a b. (a -> b) -> a -> b
$ Lambda GPUMem -> GBody GPUMem SubExpRes
forall rep t. GLambda rep t -> Body rep
lambdaBody (Lambda GPUMem -> GBody GPUMem SubExpRes)
-> Lambda GPUMem -> GBody GPUMem SubExpRes
forall a b. (a -> b) -> a -> b
$ SegPostOp GPUMem -> Lambda GPUMem
forall rep. SegPostOp rep -> Lambda rep
segPostOpLambda SegPostOp GPUMem
post_op
            sComment "compute post op." $
              compileStms mempty (bodyStms $ lambdaBody $ segPostOpLambda post_op) $
                sComment "write values" $
                  forM_ (zip (patElems pat) res) $ \(PatElem LParamMem
pe, SubExp
subexp) ->
                    VName
-> [TPrimExp Int64 VName]
-> SubExp
-> [TPrimExp Int64 VName]
-> ImpM GPUMem KernelEnv KernelOp ()
forall rep r op.
VName
-> [TPrimExp Int64 VName]
-> SubExp
-> [TPrimExp Int64 VName]
-> ImpM rep r op ()
copyDWIMFix (PatElem LParamMem -> VName
forall dec. PatElem dec -> VName
patElemName PatElem LParamMem
pe) ((VName -> TPrimExp Int64 VName)
-> [VName] -> [TPrimExp Int64 VName]
forall a b. (a -> b) -> [a] -> [b]
map VName -> TPrimExp Int64 VName
forall a. a -> TPrimExp Int64 a
Imp.le64 [VName]
gtids) SubExp
subexp []

      sOp $ Imp.Barrier Imp.FenceGlobal

isIdentityPostOp :: SegPostOp rep -> Bool
isIdentityPostOp :: forall rep. SegPostOp rep -> Bool
isIdentityPostOp SegPostOp rep
post_op =
  (SubExpRes -> SubExp) -> [SubExpRes] -> [SubExp]
forall a b. (a -> b) -> [a] -> [b]
map SubExpRes -> SubExp
resSubExp (GBody rep SubExpRes -> [SubExpRes]
forall rep res. GBody rep res -> [res]
bodyResult (GLambda rep Type -> GBody rep SubExpRes
forall rep t. GLambda rep t -> Body rep
lambdaBody GLambda rep Type
lam))
    [SubExp] -> [SubExp] -> Bool
forall a. Eq a => a -> a -> Bool
== (Param (LParamInfo rep) -> SubExp)
-> [Param (LParamInfo rep)] -> [SubExp]
forall a b. (a -> b) -> [a] -> [b]
map (VName -> SubExp
Var (VName -> SubExp)
-> (Param (LParamInfo rep) -> VName)
-> Param (LParamInfo rep)
-> SubExp
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Param (LParamInfo rep) -> VName
forall dec. Param dec -> VName
paramName) (GLambda rep Type -> [Param (LParamInfo rep)]
forall rep t. GLambda rep t -> [LParam rep]
lambdaParams GLambda rep Type
lam)
  where
    lam :: GLambda rep Type
lam = SegPostOp rep -> GLambda rep Type
forall rep. SegPostOp rep -> Lambda rep
segPostOpLambda SegPostOp rep
post_op

-- | Compile 'SegScan' instance to host-level code with calls to
-- various kernels.
compileSegScan ::
  Pat LetDecMem ->
  SegLevel ->
  SegSpace ->
  [Type] ->
  [SegBinOp GPUMem] ->
  KernelBody GPUMem ->
  SegPostOp GPUMem ->
  CallKernelGen ()
compileSegScan :: Pat LParamMem
-> SegLevel
-> SegSpace
-> [Type]
-> [SegBinOp GPUMem]
-> KernelBody GPUMem
-> SegPostOp GPUMem
-> CallKernelGen ()
compileSegScan Pat LParamMem
pat SegLevel
lvl SegSpace
space [Type]
ts [SegBinOp GPUMem]
scans KernelBody GPUMem
kbody SegPostOp GPUMem
post_op = do
  attrs <- SegLevel -> CallKernelGen KernelAttrs
lvlKernelAttrs SegLevel
lvl

  -- Stage 2 uses loop virtualization, so stage1_num_tblocks is no
  -- longer capped by the maximum thread block size.
  let stage1_num_tblocks = KernelAttrs -> Count NumBlocks SubExp
kAttrNumBlocks KernelAttrs
attrs

  -- The stage-2 block size is a tunable/user-settable parameter.  It
  -- is independent of stage1_num_tblocks so the autotuner can treat
  -- it as a fixed knob.
  stage2_tblock_size_param <- getSize "segscan_stage2_tblock_size" SizeThreadBlock
  stage2_tblock_size <-
    fmap (Imp.Count . tvSize) $
      dPrimV "stage2_tblock_size" $
        tvExp stage2_tblock_size_param

  let shpT SegBinOp rep
op = (SegBinOp rep -> Shape
forall rep. SegBinOp rep -> Shape
segBinOpShape SegBinOp rep
op,) (Type -> (Shape, Type)) -> [Type] -> [(Shape, Type)]
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> GLambda rep Type -> [Type]
forall rep t. GLambda rep t -> [t]
lambdaReturnType (SegBinOp rep -> GLambda rep Type
forall rep. SegBinOp rep -> Lambda rep
segBinOpLambda SegBinOp rep
op)
      scan_ts = (SegBinOp GPUMem -> [(Shape, Type)])
-> [SegBinOp GPUMem] -> [(Shape, Type)]
forall (t :: * -> *) a b. Foldable t => (a -> [b]) -> t a -> [b]
concatMap SegBinOp GPUMem -> [(Shape, Type)]
forall {rep}. SegBinOp rep -> [(Shape, Type)]
shpT [SegBinOp GPUMem]
scans
      shpOfT Type
t Shape
s =
        Type -> Shape
forall shape o. ArrayShape shape => TypeBase shape o -> shape
arrayShape (Type -> Shape) -> Type -> Shape
forall a b. (a -> b) -> a -> b
$
          (SubExp -> Type -> Type) -> Type -> [SubExp] -> Type
forall a b. (a -> b -> b) -> b -> [a] -> b
forall (t :: * -> *) a b.
Foldable t =>
(a -> b -> b) -> b -> t a -> b
foldr ((Type -> SubExp -> Type) -> SubExp -> Type -> Type
forall a b c. (a -> b -> c) -> b -> a -> c
flip Type -> SubExp -> Type
forall d.
TypeBase (ShapeBase d) NoMode -> d -> TypeBase (ShapeBase d) NoMode
arrayOfRow) (Type -> Shape -> Type
arrayOfShape Type
t Shape
s) ([SubExp] -> Type) -> [SubExp] -> Type
forall a b. (a -> b) -> a -> b
$
            SegSpace -> [SubExp]
segSpaceDims SegSpace
space

  (scan_out, map_out) <-
    if isIdentityPostOp post_op
      then
        let (scan_out, map_out) = splitAt (segBinOpResults scans) $ patElemName <$> patElems pat
            map_out' =
              (Type -> VName -> Maybe VName)
-> [Type] -> [VName] -> [Maybe VName]
forall a b c. (a -> b -> c) -> [a] -> [b] -> [c]
zipWith
                ( \Type
t VName
n ->
                    if Type -> Bool
forall shape o. TypeBase shape o -> Bool
isAcc Type
t then Maybe VName
forall a. Maybe a
Nothing else VName -> Maybe VName
forall a. a -> Maybe a
Just VName
n
                )
                (Int -> [Type] -> [Type]
forall a. Int -> [a] -> [a]
drop ([SegBinOp GPUMem] -> Int
forall rep. [SegBinOp rep] -> Int
segBinOpResults [SegBinOp GPUMem]
scans) [Type]
ts)
                [VName]
map_out
         in pure (scan_out, map_out')
      else do
        scan_out <- forM scan_ts $ \(Shape
s, Type
t) ->
          Name
-> PrimType -> Shape -> Space -> ImpM GPUMem HostEnv HostOp VName
forall rep r op.
Name -> PrimType -> Shape -> Space -> ImpM rep r op VName
sAllocArray Name
"scan_out" (Type -> PrimType
forall shape o. TypeBase shape o -> PrimType
elemType Type
t) (Type -> Shape -> Shape
shpOfT Type
t Shape
s) (SpaceId -> Space
Space SpaceId
"device")

        map_out <- forM (drop (segBinOpResults scans) ts) $ \Type
t ->
          if Type -> Bool
forall shape o. TypeBase shape o -> Bool
isAcc Type
t
            then Maybe VName -> ImpM GPUMem HostEnv HostOp (Maybe VName)
forall a. a -> ImpM GPUMem HostEnv HostOp a
forall (f :: * -> *) a. Applicative f => a -> f a
pure Maybe VName
forall a. Maybe a
Nothing
            else VName -> Maybe VName
forall a. a -> Maybe a
Just (VName -> Maybe VName)
-> ImpM GPUMem HostEnv HostOp VName
-> ImpM GPUMem HostEnv HostOp (Maybe VName)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Name
-> PrimType -> Shape -> Space -> ImpM GPUMem HostEnv HostOp VName
forall rep r op.
Name -> PrimType -> Shape -> Space -> ImpM rep r op VName
sAllocArray Name
"map_out" (Type -> PrimType
forall shape o. TypeBase shape o -> PrimType
elemType Type
t) (Type -> Shape -> Shape
shpOfT Type
t Shape
forall a. Monoid a => a
mempty) (SpaceId -> Space
Space SpaceId
"device")

        pure (scan_out, map_out)

  (stage1_num_threads, elems_per_group, crossesSegment) <-
    scanStage1 scan_out map_out stage1_num_tblocks (kAttrBlockSize attrs) space scans kbody

  emit $ Imp.DebugPrint "elems_per_group" $ Just $ untyped elems_per_group

  scanStage2 scan_out stage1_num_threads elems_per_group stage1_num_tblocks stage2_tblock_size crossesSegment space scans
  scanStage3 pat scan_out map_out (kAttrNumBlocks attrs) (kAttrBlockSize attrs) elems_per_group crossesSegment space scans post_op