{-# LANGUAGE TypeFamilies #-}

-- | Code generation for segmented and non-segmented scans.  Uses a
-- fast single-pass algorithm, but which only works on NVIDIA GPUs and
-- with some constraints on the operator.  We use this when we can.
module Futhark.CodeGen.ImpGen.GPU.SegScan.SinglePass (compileSegScan) where

import Control.Monad
import Data.List (zip4, zip5, zip7)
import Data.Map qualified as M
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 (mapAccumLM)
import Futhark.Util.IntegralExp (IntegralExp (mod, rem), ceilDiv, nextMul, quot)
import Prelude hiding (mod, quot, rem)

xParams, yParams :: SegBinOp GPUMem -> [LParam GPUMem]
xParams :: SegBinOp GPUMem -> [LParam GPUMem]
xParams SegBinOp GPUMem
scan =
  Int -> [Param LetDecMem] -> [Param LetDecMem]
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)) (GLambda GPUMem Type -> [LParam GPUMem]
forall rep t. GLambda rep t -> [LParam rep]
lambdaParams (SegBinOp GPUMem -> GLambda GPUMem Type
forall rep. SegBinOp rep -> Lambda rep
segBinOpLambda SegBinOp GPUMem
scan))
yParams :: SegBinOp GPUMem -> [LParam GPUMem]
yParams SegBinOp GPUMem
scan =
  Int -> [Param LetDecMem] -> [Param LetDecMem]
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)) (GLambda GPUMem Type -> [LParam GPUMem]
forall rep t. GLambda rep t -> [LParam rep]
lambdaParams (SegBinOp GPUMem -> GLambda GPUMem Type
forall rep. SegBinOp rep -> Lambda rep
segBinOpLambda SegBinOp GPUMem
scan))

-- | Given available register, thread block size, scan parameter
-- types, and map parameter types, compute the largest available chunk
-- size given the parameters for which we want chunking and the
-- available resources.
getScanChunkSize :: [Type] -> Imp.KernelConstExp
getScanChunkSize :: [Type] -> KernelConstExp
getScanChunkSize [Type]
types =
  let max_tblock_size :: KernelConst
max_tblock_size = SizeClass -> KernelConst
Imp.SizeMaxConst SizeClass
SizeThreadBlock
      max_block_mem :: KernelConst
max_block_mem = SizeClass -> KernelConst
Imp.SizeMaxConst SizeClass
SizeSharedMemory
      max_block_reg :: KernelConst
max_block_reg = SizeClass -> KernelConst
Imp.SizeMaxConst SizeClass
SizeRegisters
      k_mem :: TPrimExp Int64 KernelConst
k_mem = KernelConst -> TPrimExp Int64 KernelConst
forall a. a -> TPrimExp Int64 a
le64 KernelConst
max_block_mem TPrimExp Int64 KernelConst
-> TPrimExp Int64 KernelConst -> TPrimExp Int64 KernelConst
forall e. IntegralExp e => e -> e -> e
`quot` KernelConst -> TPrimExp Int64 KernelConst
forall a. a -> TPrimExp Int64 a
le64 KernelConst
max_tblock_size
      k_reg :: TPrimExp Int64 KernelConst
k_reg = KernelConst -> TPrimExp Int64 KernelConst
forall a. a -> TPrimExp Int64 a
le64 KernelConst
max_block_reg TPrimExp Int64 KernelConst
-> TPrimExp Int64 KernelConst -> TPrimExp Int64 KernelConst
forall e. IntegralExp e => e -> e -> e
`quot` KernelConst -> TPrimExp Int64 KernelConst
forall a. a -> TPrimExp Int64 a
le64 KernelConst
max_tblock_size
      types' :: [PrimType]
types' = (Type -> PrimType) -> [Type] -> [PrimType]
forall a b. (a -> b) -> [a] -> [b]
map Type -> PrimType
forall shape o. TypeBase shape o -> PrimType
elemType ([Type] -> [PrimType]) -> [Type] -> [PrimType]
forall a b. (a -> b) -> a -> b
$ (Type -> Bool) -> [Type] -> [Type]
forall a. (a -> Bool) -> [a] -> [a]
filter Type -> Bool
forall shape o. TypeBase shape o -> Bool
primType [Type]
types
      sizes :: [TPrimExp Int64 KernelConst]
sizes = (PrimType -> TPrimExp Int64 KernelConst)
-> [PrimType] -> [TPrimExp Int64 KernelConst]
forall a b. (a -> b) -> [a] -> [b]
map PrimType -> TPrimExp Int64 KernelConst
forall a. Num a => PrimType -> a
primByteSize [PrimType]
types'

      sum_sizes :: TPrimExp Int64 KernelConst
sum_sizes = [TPrimExp Int64 KernelConst] -> TPrimExp Int64 KernelConst
forall a. Num a => [a] -> a
forall (t :: * -> *) a. (Foldable t, Num a) => t a -> a
sum [TPrimExp Int64 KernelConst]
sizes
      sum_sizes' :: TPrimExp Int64 KernelConst
sum_sizes' = [TPrimExp Int64 KernelConst] -> TPrimExp Int64 KernelConst
forall a. Num a => [a] -> a
forall (t :: * -> *) a. (Foldable t, Num a) => t a -> a
sum ((PrimType -> TPrimExp Int64 KernelConst)
-> [PrimType] -> [TPrimExp Int64 KernelConst]
forall a b. (a -> b) -> [a] -> [b]
map (TPrimExp Int64 KernelConst
-> TPrimExp Int64 KernelConst -> TPrimExp Int64 KernelConst
forall v. TPrimExp Int64 v -> TPrimExp Int64 v -> TPrimExp Int64 v
sMax64 TPrimExp Int64 KernelConst
4 (TPrimExp Int64 KernelConst -> TPrimExp Int64 KernelConst)
-> (PrimType -> TPrimExp Int64 KernelConst)
-> PrimType
-> TPrimExp Int64 KernelConst
forall b c a. (b -> c) -> (a -> b) -> a -> c
. PrimType -> TPrimExp Int64 KernelConst
forall a. Num a => PrimType -> a
primByteSize) [PrimType]
types') TPrimExp Int64 KernelConst
-> TPrimExp Int64 KernelConst -> TPrimExp Int64 KernelConst
forall e. IntegralExp e => e -> e -> e
`quot` TPrimExp Int64 KernelConst
4
      max_size :: TPrimExp Int64 KernelConst
max_size = [TPrimExp Int64 KernelConst] -> TPrimExp Int64 KernelConst
forall a. Ord a => [a] -> a
forall (t :: * -> *) a. (Foldable t, Ord a) => t a -> a
maximum [TPrimExp Int64 KernelConst]
sizes

      mem_constraint :: TPrimExp Int64 KernelConst
mem_constraint = TPrimExp Int64 KernelConst
-> TPrimExp Int64 KernelConst -> TPrimExp Int64 KernelConst
forall a. Ord a => a -> a -> a
max TPrimExp Int64 KernelConst
k_mem TPrimExp Int64 KernelConst
sum_sizes TPrimExp Int64 KernelConst
-> TPrimExp Int64 KernelConst -> TPrimExp Int64 KernelConst
forall e. IntegralExp e => e -> e -> e
`quot` TPrimExp Int64 KernelConst
max_size
      reg_constraint :: TPrimExp Int64 KernelConst
reg_constraint = (TPrimExp Int64 KernelConst
k_reg TPrimExp Int64 KernelConst
-> TPrimExp Int64 KernelConst -> TPrimExp Int64 KernelConst
forall a. Num a => a -> a -> a
- TPrimExp Int64 KernelConst
1 TPrimExp Int64 KernelConst
-> TPrimExp Int64 KernelConst -> TPrimExp Int64 KernelConst
forall a. Num a => a -> a -> a
- TPrimExp Int64 KernelConst
sum_sizes') TPrimExp Int64 KernelConst
-> TPrimExp Int64 KernelConst -> TPrimExp Int64 KernelConst
forall e. IntegralExp e => e -> e -> e
`quot` (TPrimExp Int64 KernelConst
2 TPrimExp Int64 KernelConst
-> TPrimExp Int64 KernelConst -> TPrimExp Int64 KernelConst
forall a. Num a => a -> a -> a
* TPrimExp Int64 KernelConst
sum_sizes')
   in TPrimExp Int64 KernelConst -> KernelConstExp
forall {k} (t :: k) v. TPrimExp t v -> PrimExp v
untyped (TPrimExp Int64 KernelConst -> KernelConstExp)
-> TPrimExp Int64 KernelConst -> KernelConstExp
forall a b. (a -> b) -> a -> b
$ TPrimExp Int64 KernelConst
-> TPrimExp Int64 KernelConst -> TPrimExp Int64 KernelConst
forall v. TPrimExp Int64 v -> TPrimExp Int64 v -> TPrimExp Int64 v
sMax64 TPrimExp Int64 KernelConst
1 (TPrimExp Int64 KernelConst -> TPrimExp Int64 KernelConst)
-> TPrimExp Int64 KernelConst -> TPrimExp Int64 KernelConst
forall a b. (a -> b) -> a -> b
$ TPrimExp Int64 KernelConst
-> TPrimExp Int64 KernelConst -> TPrimExp Int64 KernelConst
forall v. TPrimExp Int64 v -> TPrimExp Int64 v -> TPrimExp Int64 v
sMin64 TPrimExp Int64 KernelConst
mem_constraint TPrimExp Int64 KernelConst
reg_constraint

createLocalArrays ::
  Count BlockSize SubExp ->
  SubExp ->
  [PrimType] ->
  InKernelGen (VName, [VName], [VName], VName, [VName])
createLocalArrays :: Count BlockSize SubExp
-> SubExp
-> [PrimType]
-> InKernelGen (VName, [VName], [VName], VName, [VName])
createLocalArrays (Count SubExp
block_size) SubExp
chunk [PrimType]
types = do
  let block_sizeE :: TExp Int64
block_sizeE = SubExp -> TExp Int64
pe64 SubExp
block_size
      workSize :: TExp Int64
workSize = SubExp -> TExp Int64
pe64 SubExp
chunk TExp Int64 -> TExp Int64 -> TExp Int64
forall a. Num a => a -> a -> a
* TExp Int64
block_sizeE
      prefixArraysSize :: TExp Int64
prefixArraysSize =
        (TExp Int64 -> TExp Int64 -> TExp Int64)
-> TExp Int64 -> [TExp Int64] -> TExp Int64
forall b a. (b -> a -> b) -> b -> [a] -> b
forall (t :: * -> *) b a.
Foldable t =>
(b -> a -> b) -> b -> t a -> b
foldl (\TExp Int64
acc TExp Int64
tySize -> TExp Int64 -> TExp Int64 -> TExp Int64
forall e. IntegralExp e => e -> e -> e
nextMul TExp Int64
acc TExp Int64
tySize TExp Int64 -> TExp Int64 -> TExp Int64
forall a. Num a => a -> a -> a
+ TExp Int64
tySize TExp Int64 -> TExp Int64 -> TExp Int64
forall a. Num a => a -> a -> a
* TExp Int64
block_sizeE) TExp Int64
0 ([TExp Int64] -> TExp Int64) -> [TExp Int64] -> TExp Int64
forall a b. (a -> b) -> a -> b
$
          (PrimType -> TExp Int64) -> [PrimType] -> [TExp Int64]
forall a b. (a -> b) -> [a] -> [b]
map PrimType -> TExp Int64
forall a. Num a => PrimType -> a
primByteSize [PrimType]
types
      maxTransposedArraySize :: TExp Int64
maxTransposedArraySize =
        (TExp Int64 -> TExp Int64 -> TExp Int64)
-> [TExp Int64] -> TExp Int64
forall a. (a -> a -> a) -> [a] -> a
forall (t :: * -> *) a. Foldable t => (a -> a -> a) -> t a -> a
foldl1 TExp Int64 -> TExp Int64 -> TExp Int64
forall v. TPrimExp Int64 v -> TPrimExp Int64 v -> TPrimExp Int64 v
sMax64 ([TExp Int64] -> TExp Int64) -> [TExp Int64] -> TExp Int64
forall a b. (a -> b) -> a -> b
$ (PrimType -> TExp Int64) -> [PrimType] -> [TExp Int64]
forall a b. (a -> b) -> [a] -> [b]
map (\PrimType
ty -> TExp Int64
workSize TExp Int64 -> TExp Int64 -> TExp Int64
forall a. Num a => a -> a -> a
* PrimType -> TExp Int64
forall a. Num a => PrimType -> a
primByteSize PrimType
ty) [PrimType]
types

      warp_size :: (Num a) => a
      warp_size :: forall a. Num a => a
warp_size = a
32
      maxWarpExchangeSize :: TExp Int64
maxWarpExchangeSize =
        (TExp Int64 -> TExp Int64 -> TExp Int64)
-> TExp Int64 -> [TExp Int64] -> TExp Int64
forall b a. (b -> a -> b) -> b -> [a] -> b
forall (t :: * -> *) b a.
Foldable t =>
(b -> a -> b) -> b -> t a -> b
foldl (\TExp Int64
acc TExp Int64
tySize -> TExp Int64 -> TExp Int64 -> TExp Int64
forall e. IntegralExp e => e -> e -> e
nextMul TExp Int64
acc TExp Int64
tySize TExp Int64 -> TExp Int64 -> TExp Int64
forall a. Num a => a -> a -> a
+ TExp Int64
tySize TExp Int64 -> TExp Int64 -> TExp Int64
forall a. Num a => a -> a -> a
* Integer -> TExp Int64
forall a. Num a => Integer -> a
fromInteger Integer
forall a. Num a => a
warp_size) TExp Int64
0 ([TExp Int64] -> TExp Int64) -> [TExp Int64] -> TExp Int64
forall a b. (a -> b) -> a -> b
$
          (PrimType -> TExp Int64) -> [PrimType] -> [TExp Int64]
forall a b. (a -> b) -> [a] -> [b]
map PrimType -> TExp Int64
forall a. Num a => PrimType -> a
primByteSize [PrimType]
types
      maxLookbackSize :: TExp Int64
maxLookbackSize = TExp Int64
maxWarpExchangeSize TExp Int64 -> TExp Int64 -> TExp Int64
forall a. Num a => a -> a -> a
+ TExp Int64
forall a. Num a => a
warp_size
      size :: Count Bytes (TExp Int64)
size = TExp Int64 -> Count Bytes (TExp Int64)
forall a. a -> Count Bytes a
Imp.bytes (TExp Int64 -> Count Bytes (TExp Int64))
-> TExp Int64 -> Count Bytes (TExp Int64)
forall a b. (a -> b) -> a -> b
$ TExp Int64
maxLookbackSize TExp Int64 -> TExp Int64 -> TExp Int64
forall v. TPrimExp Int64 v -> TPrimExp Int64 v -> TPrimExp Int64 v
`sMax64` TExp Int64
prefixArraysSize TExp Int64 -> TExp Int64 -> TExp Int64
forall v. TPrimExp Int64 v -> TPrimExp Int64 v -> TPrimExp Int64 v
`sMax64` TExp Int64
maxTransposedArraySize

  (_, byteOffsets) <-
    (TExp Int64
 -> TExp Int64
 -> ImpM GPUMem KernelEnv KernelOp (TExp Int64, TExp Int64))
-> TExp Int64
-> [TExp Int64]
-> ImpM GPUMem KernelEnv KernelOp (TExp Int64, [TExp Int64])
forall (m :: * -> *) (t :: * -> *) acc x y.
(Monad m, Traversable t) =>
(acc -> x -> m (acc, y)) -> acc -> t x -> m (acc, t y)
mapAccumLM
      ( \TExp Int64
off TExp Int64
tySize -> do
          off' <- Name -> TExp Int64 -> ImpM GPUMem KernelEnv KernelOp (TExp Int64)
forall {k} (t :: k) rep r op.
Name -> TExp t -> ImpM rep r op (TExp t)
dPrimVE Name
"byte_offsets" (TExp Int64 -> ImpM GPUMem KernelEnv KernelOp (TExp Int64))
-> TExp Int64 -> ImpM GPUMem KernelEnv KernelOp (TExp Int64)
forall a b. (a -> b) -> a -> b
$ TExp Int64 -> TExp Int64 -> TExp Int64
forall e. IntegralExp e => e -> e -> e
nextMul TExp Int64
off TExp Int64
tySize TExp Int64 -> TExp Int64 -> TExp Int64
forall a. Num a => a -> a -> a
+ SubExp -> TExp Int64
pe64 SubExp
block_size TExp Int64 -> TExp Int64 -> TExp Int64
forall a. Num a => a -> a -> a
* TExp Int64
tySize
          pure (off', off)
      )
      TExp Int64
0
      ([TExp Int64]
 -> ImpM GPUMem KernelEnv KernelOp (TExp Int64, [TExp Int64]))
-> [TExp Int64]
-> ImpM GPUMem KernelEnv KernelOp (TExp Int64, [TExp Int64])
forall a b. (a -> b) -> a -> b
$ (PrimType -> TExp Int64) -> [PrimType] -> [TExp Int64]
forall a b. (a -> b) -> [a] -> [b]
map PrimType -> TExp Int64
forall a. Num a => PrimType -> a
primByteSize [PrimType]
types

  (_, warpByteOffsets) <-
    mapAccumLM
      ( \TExp Int64
off TExp Int64
tySize -> do
          off' <- Name -> TExp Int64 -> ImpM GPUMem KernelEnv KernelOp (TExp Int64)
forall {k} (t :: k) rep r op.
Name -> TExp t -> ImpM rep r op (TExp t)
dPrimVE Name
"warp_byte_offset" (TExp Int64 -> ImpM GPUMem KernelEnv KernelOp (TExp Int64))
-> TExp Int64 -> ImpM GPUMem KernelEnv KernelOp (TExp Int64)
forall a b. (a -> b) -> a -> b
$ TExp Int64 -> TExp Int64 -> TExp Int64
forall e. IntegralExp e => e -> e -> e
nextMul TExp Int64
off TExp Int64
tySize TExp Int64 -> TExp Int64 -> TExp Int64
forall a. Num a => a -> a -> a
+ TExp Int64
forall a. Num a => a
warp_size TExp Int64 -> TExp Int64 -> TExp Int64
forall a. Num a => a -> a -> a
* TExp Int64
tySize
          pure (off', off)
      )
      warp_size
      $ map primByteSize types

  sComment "Allocate reusable shared memory" $ pure ()

  localMem <- sAlloc "local_mem" size (Space "shared")
  transposeArrayLength <- dPrimV "trans_arr_len" workSize

  sharedId <- sArrayInMem "shared_id" int32 (Shape [constant (1 :: Int32)]) localMem

  transposedArrays <-
    forM types $ \PrimType
ty ->
      Name
-> PrimType
-> ShapeBase SubExp
-> VName
-> ImpM GPUMem KernelEnv KernelOp VName
forall rep r op.
Name
-> PrimType -> ShapeBase SubExp -> VName -> ImpM rep r op VName
sArrayInMem
        Name
"local_transpose_arr"
        PrimType
ty
        ([SubExp] -> ShapeBase SubExp
forall d. [d] -> ShapeBase d
Shape [TV Int64 -> SubExp
forall {k} (t :: k). TV t -> SubExp
tvSize TV Int64
transposeArrayLength])
        VName
localMem

  prefixArrays <-
    forM (zip byteOffsets types) $ \(TExp Int64
off, PrimType
ty) -> do
      let off' :: TExp Int64
off' = TExp Int64
off TExp Int64 -> TExp Int64 -> TExp Int64
forall e. IntegralExp e => e -> e -> e
`quot` PrimType -> TExp Int64
forall a. Num a => PrimType -> a
primByteSize PrimType
ty
      Name
-> PrimType
-> ShapeBase SubExp
-> VName
-> LMAD
-> ImpM GPUMem KernelEnv KernelOp VName
forall rep r op.
Name
-> PrimType
-> ShapeBase SubExp
-> VName
-> LMAD
-> ImpM rep r op VName
sArray
        Name
"local_prefix_arr"
        PrimType
ty
        ([SubExp] -> ShapeBase SubExp
forall d. [d] -> ShapeBase d
Shape [SubExp
block_size])
        VName
localMem
        (LMAD -> ImpM GPUMem KernelEnv KernelOp VName)
-> LMAD -> ImpM GPUMem KernelEnv KernelOp VName
forall a b. (a -> b) -> a -> b
$ TExp Int64 -> [TExp Int64] -> LMAD
forall num. IntegralExp num => num -> [num] -> LMAD num
LMAD.iota TExp Int64
off' [SubExp -> TExp Int64
pe64 SubExp
block_size]

  warpscan <- sArrayInMem "warpscan" int8 (Shape [constant (warp_size :: Int64)]) localMem
  warpExchanges <-
    forM (zip warpByteOffsets types) $ \(TExp Int64
off, PrimType
ty) -> do
      let off' :: TExp Int64
off' = TExp Int64
off TExp Int64 -> TExp Int64 -> TExp Int64
forall e. IntegralExp e => e -> e -> e
`quot` PrimType -> TExp Int64
forall a. Num a => PrimType -> a
primByteSize PrimType
ty
      Name
-> PrimType
-> ShapeBase SubExp
-> VName
-> LMAD
-> ImpM GPUMem KernelEnv KernelOp VName
forall rep r op.
Name
-> PrimType
-> ShapeBase SubExp
-> VName
-> LMAD
-> ImpM rep r op VName
sArray
        Name
"warp_exchange"
        PrimType
ty
        ([SubExp] -> ShapeBase SubExp
forall d. [d] -> ShapeBase d
Shape [Int64 -> SubExp
forall v. IsValue v => v -> SubExp
constant (Int64
forall a. Num a => a
warp_size :: Int64)])
        VName
localMem
        (LMAD -> ImpM GPUMem KernelEnv KernelOp VName)
-> LMAD -> ImpM GPUMem KernelEnv KernelOp VName
forall a b. (a -> b) -> a -> b
$ TExp Int64 -> [TExp Int64] -> LMAD
forall num. IntegralExp num => num -> [num] -> LMAD num
LMAD.iota TExp Int64
off' [TExp Int64
forall a. Num a => a
warp_size]

  pure (sharedId, transposedArrays, prefixArrays, warpscan, warpExchanges)

statusX, statusA, statusP :: (Num a) => a
statusX :: forall a. Num a => a
statusX = a
0
statusA :: forall a. Num a => a
statusA = a
1
statusP :: forall a. Num a => a
statusP = a
2

inBlockScanLookback ::
  KernelConstants ->
  Imp.TExp Int64 ->
  VName ->
  [VName] ->
  Lambda GPUMem ->
  InKernelGen ()
inBlockScanLookback :: KernelConstants
-> TExp Int64
-> VName
-> [VName]
-> GLambda GPUMem Type
-> ImpM GPUMem KernelEnv KernelOp ()
inBlockScanLookback KernelConstants
constants TExp Int64
arrs_full_size VName
flag_arr [VName]
arrs GLambda GPUMem Type
scan_lam = ImpM GPUMem KernelEnv KernelOp ()
-> ImpM GPUMem KernelEnv KernelOp ()
forall rep r op a. ImpM rep r op a -> ImpM rep r op a
everythingVolatile (ImpM GPUMem KernelEnv KernelOp ()
 -> ImpM GPUMem KernelEnv KernelOp ())
-> ImpM GPUMem KernelEnv KernelOp ()
-> ImpM GPUMem KernelEnv KernelOp ()
forall a b. (a -> b) -> a -> b
$ do
  flg_x :: TV Int8 <- Name -> ImpM GPUMem KernelEnv KernelOp (TV Int8)
forall {k} (t :: k) rep r op.
MkTV t =>
Name -> ImpM rep r op (TV t)
dPrim Name
"flg_x"
  flg_y :: TV Int8 <- dPrim "flg_y"
  let flg_param_x = Attrs -> VName -> LetDecMem -> Param LetDecMem
forall dec. Attrs -> VName -> dec -> Param dec
Param Attrs
forall a. Monoid a => a
mempty (TV Int8 -> VName
forall {k} (t :: k). TV t -> VName
tvVar TV Int8
flg_x) (PrimType -> LetDecMem
forall d o ret. PrimType -> MemInfo d o ret
MemPrim PrimType
p_int8)
      flg_param_y = Attrs -> VName -> LetDecMem -> Param LetDecMem
forall dec. Attrs -> VName -> dec -> Param dec
Param Attrs
forall a. Monoid a => a
mempty (TV Int8 -> VName
forall {k} (t :: k). TV t -> VName
tvVar TV Int8
flg_y) (PrimType -> LetDecMem
forall d o ret. PrimType -> MemInfo d o ret
MemPrim PrimType
p_int8)
      flg_y_exp = TV Int8 -> TPrimExp Int8 VName
forall {k} (t :: k). TV t -> TExp t
tvExp TV Int8
flg_y
      statusP_e = TPrimExp Int8 VName
forall a. Num a => a
statusP :: Imp.TExp Int8
      statusX_e = TPrimExp Int8 VName
forall a. Num a => a
statusX :: Imp.TExp Int8

  dLParams (lambdaParams scan_lam)

  skip_threads <- dPrim "skip_threads"
  let in_block_thread_active =
        TV Int32 -> TExp Int32
forall {k} (t :: k). TV t -> TExp t
tvExp TV Int32
skip_threads TExp Int32 -> TExp Int32 -> TPrimExp Bool VName
forall {k} v (t :: k).
Eq v =>
TPrimExp t v -> TPrimExp t v -> TPrimExp Bool v
.<=. TExp Int32
in_block_id
      actual_params = GLambda GPUMem Type -> [LParam GPUMem]
forall rep t. GLambda rep t -> [LParam rep]
lambdaParams GLambda GPUMem Type
scan_lam
      (x_params, y_params) =
        splitAt (length actual_params `div` 2) actual_params
      y_to_x =
        [(Param LetDecMem, Param LetDecMem)]
-> ((Param LetDecMem, Param LetDecMem)
    -> 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 LetDecMem]
-> [Param LetDecMem] -> [(Param LetDecMem, Param LetDecMem)]
forall a b. [a] -> [b] -> [(a, b)]
zip [Param LetDecMem]
x_params [Param LetDecMem]
y_params) (((Param LetDecMem, Param LetDecMem)
  -> ImpM GPUMem KernelEnv KernelOp ())
 -> ImpM GPUMem KernelEnv KernelOp ())
-> ((Param LetDecMem, Param LetDecMem)
    -> ImpM GPUMem KernelEnv KernelOp ())
-> ImpM GPUMem KernelEnv KernelOp ()
forall a b. (a -> b) -> a -> b
$ \(Param LetDecMem
x, Param LetDecMem
y) ->
          Bool
-> ImpM GPUMem KernelEnv KernelOp ()
-> ImpM GPUMem KernelEnv KernelOp ()
forall (f :: * -> *). Applicative f => Bool -> f () -> f ()
when (Type -> Bool
forall shape o. TypeBase shape o -> Bool
primType (Param LetDecMem -> Type
forall dec. Typed dec => Param dec -> Type
paramType Param LetDecMem
x)) (ImpM GPUMem KernelEnv KernelOp ()
 -> ImpM GPUMem KernelEnv KernelOp ())
-> ImpM GPUMem KernelEnv KernelOp ()
-> ImpM GPUMem KernelEnv KernelOp ()
forall a b. (a -> b) -> a -> b
$
            VName
-> [DimIndex (TExp Int64)]
-> SubExp
-> [DimIndex (TExp Int64)]
-> ImpM GPUMem KernelEnv KernelOp ()
forall rep r op.
VName
-> [DimIndex (TExp Int64)]
-> SubExp
-> [DimIndex (TExp Int64)]
-> ImpM rep r op ()
copyDWIM (Param LetDecMem -> VName
forall dec. Param dec -> VName
paramName Param LetDecMem
x) [] (VName -> SubExp
Var (Param LetDecMem -> VName
forall dec. Param dec -> VName
paramName Param LetDecMem
y)) []
      y_to_x_flg =
        VName
-> [DimIndex (TExp Int64)]
-> SubExp
-> [DimIndex (TExp Int64)]
-> ImpM GPUMem KernelEnv KernelOp ()
forall rep r op.
VName
-> [DimIndex (TExp Int64)]
-> SubExp
-> [DimIndex (TExp Int64)]
-> ImpM rep r op ()
copyDWIM (TV Int8 -> VName
forall {k} (t :: k). TV t -> VName
tvVar TV Int8
flg_x) [] (VName -> SubExp
Var (TV Int8 -> VName
forall {k} (t :: k). TV t -> VName
tvVar TV Int8
flg_y)) []

  -- Set initial y values
  sComment "read input for in-block scan" $ do
    zipWithM_ readInitial (flg_param_y : y_params) (flag_arr : arrs)
    -- Since the final result is expected to be in x_params, we may
    -- need to copy it there for the first thread in the block.
    sWhen (in_block_id .==. 0) $ do
      y_to_x
      y_to_x_flg

  when array_scan barrier

  let op_to_x = do
        TPrimExp Bool VName
-> ImpM GPUMem KernelEnv KernelOp ()
-> ImpM GPUMem KernelEnv KernelOp ()
-> ImpM GPUMem KernelEnv KernelOp ()
forall rep r op.
TPrimExp Bool VName
-> ImpM rep r op () -> ImpM rep r op () -> ImpM rep r op ()
sIf
          (TPrimExp Int8 VName
flg_y_exp TPrimExp Int8 VName -> TPrimExp Int8 VName -> TPrimExp Bool VName
forall {k} v (t :: k).
Eq v =>
TPrimExp t v -> TPrimExp t v -> TPrimExp Bool v
.==. TPrimExp Int8 VName
statusP_e TPrimExp Bool VName -> TPrimExp Bool VName -> TPrimExp Bool VName
forall v.
Eq v =>
TPrimExp Bool v -> TPrimExp Bool v -> TPrimExp Bool v
.||. TPrimExp Int8 VName
flg_y_exp TPrimExp Int8 VName -> TPrimExp Int8 VName -> TPrimExp Bool VName
forall {k} v (t :: k).
Eq v =>
TPrimExp t v -> TPrimExp t v -> TPrimExp Bool v
.==. TPrimExp Int8 VName
statusX_e)
          ( do
              ImpM GPUMem KernelEnv KernelOp ()
y_to_x_flg
              ImpM GPUMem KernelEnv KernelOp ()
y_to_x
          )
          ([Param LetDecMem]
-> Body GPUMem -> ImpM GPUMem KernelEnv KernelOp ()
forall dec rep r op. [Param dec] -> Body rep -> ImpM rep r op ()
compileBody' [Param LetDecMem]
x_params (Body GPUMem -> ImpM GPUMem KernelEnv KernelOp ())
-> Body GPUMem -> ImpM GPUMem KernelEnv KernelOp ()
forall a b. (a -> b) -> a -> b
$ GLambda GPUMem Type -> Body GPUMem
forall rep t. GLambda rep t -> Body rep
lambdaBody GLambda GPUMem Type
scan_lam)

  sComment "in-block scan (hopefully no barriers needed)" $ do
    skip_threads <-- 1

    sWhile (tvExp skip_threads .<. block_size) $ do
      sWhen in_block_thread_active $ do
        sComment "read operands" $
          zipWithM_
            (readParam (sExt64 $ tvExp skip_threads))
            (flg_param_x : x_params)
            (flag_arr : arrs)
        sComment "perform operation" op_to_x

        sComment "write result" $
          sequence_ $
            zipWith3
              writeResult
              (flg_param_x : x_params)
              (flg_param_y : y_params)
              (flag_arr : arrs)

      skip_threads <-- tvExp skip_threads * 2
  where
    p_int8 :: PrimType
p_int8 = IntType -> PrimType
IntType IntType
Int8
    block_size :: TExp Int32
block_size = TExp Int32
32
    block_id :: TExp Int32
block_id = TExp Int32
ltid32 TExp Int32 -> TExp Int32 -> TExp Int32
forall e. IntegralExp e => e -> e -> e
`quot` TExp Int32
block_size
    in_block_id :: TExp Int32
in_block_id = TExp Int32
ltid32 TExp Int32 -> TExp Int32 -> TExp Int32
forall a. Num a => a -> a -> a
- TExp Int32
block_id TExp Int32 -> TExp Int32 -> TExp Int32
forall a. Num a => a -> a -> a
* TExp Int32
block_size
    ltid32 :: TExp Int32
ltid32 = KernelConstants -> TExp Int32
kernelLocalThreadId KernelConstants
constants
    ltid :: TExp Int64
ltid = TExp Int32 -> TExp Int64
forall {k} (t :: k) v. IntExp t => TPrimExp t v -> TPrimExp Int64 v
sExt64 TExp Int32
ltid32
    gtid :: TExp Int64
gtid = TExp Int32 -> TExp Int64
forall {k} (t :: k) v. IntExp t => TPrimExp t v -> TPrimExp Int64 v
sExt64 (TExp Int32 -> TExp Int64) -> TExp Int32 -> TExp Int64
forall a b. (a -> b) -> a -> b
$ KernelConstants -> TExp Int32
kernelGlobalThreadId KernelConstants
constants
    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
$ GLambda GPUMem Type -> [Type]
forall rep t. GLambda rep t -> [t]
lambdaReturnType GLambda GPUMem Type
scan_lam
    barrier :: ImpM GPUMem KernelEnv KernelOp ()
barrier
      | Bool
array_scan =
          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
      | Bool
otherwise =
          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.FenceLocal

    readInitial :: Param LetDecMem -> VName -> ImpM GPUMem KernelEnv KernelOp ()
readInitial Param LetDecMem
p VName
arr
      | Type -> Bool
forall shape o. TypeBase shape o -> Bool
primType (Type -> Bool) -> Type -> Bool
forall a b. (a -> b) -> a -> b
$ Param LetDecMem -> Type
forall dec. Typed dec => Param dec -> Type
paramType Param LetDecMem
p =
          VName
-> [TExp Int64]
-> SubExp
-> [TExp Int64]
-> ImpM GPUMem KernelEnv KernelOp ()
forall rep r op.
VName -> [TExp Int64] -> SubExp -> [TExp Int64] -> ImpM rep r op ()
copyDWIMFix (Param LetDecMem -> VName
forall dec. Param dec -> VName
paramName Param LetDecMem
p) [] (VName -> SubExp
Var VName
arr) [TExp Int64
ltid]
      | Bool
otherwise =
          VName
-> [TExp Int64]
-> SubExp
-> [TExp Int64]
-> ImpM GPUMem KernelEnv KernelOp ()
forall rep r op.
VName -> [TExp Int64] -> SubExp -> [TExp Int64] -> ImpM rep r op ()
copyDWIMFix (Param LetDecMem -> VName
forall dec. Param dec -> VName
paramName Param LetDecMem
p) [] (VName -> SubExp
Var VName
arr) [TExp Int64
gtid]
    readParam :: TExp Int64
-> Param LetDecMem -> VName -> ImpM GPUMem KernelEnv KernelOp ()
readParam TExp Int64
behind Param LetDecMem
p VName
arr
      | Type -> Bool
forall shape o. TypeBase shape o -> Bool
primType (Type -> Bool) -> Type -> Bool
forall a b. (a -> b) -> a -> b
$ Param LetDecMem -> Type
forall dec. Typed dec => Param dec -> Type
paramType Param LetDecMem
p =
          VName
-> [TExp Int64]
-> SubExp
-> [TExp Int64]
-> ImpM GPUMem KernelEnv KernelOp ()
forall rep r op.
VName -> [TExp Int64] -> SubExp -> [TExp Int64] -> ImpM rep r op ()
copyDWIMFix (Param LetDecMem -> VName
forall dec. Param dec -> VName
paramName Param LetDecMem
p) [] (VName -> SubExp
Var VName
arr) [TExp Int64
ltid TExp Int64 -> TExp Int64 -> TExp Int64
forall a. Num a => a -> a -> a
- TExp Int64
behind]
      | Bool
otherwise =
          VName
-> [TExp Int64]
-> SubExp
-> [TExp Int64]
-> ImpM GPUMem KernelEnv KernelOp ()
forall rep r op.
VName -> [TExp Int64] -> SubExp -> [TExp Int64] -> ImpM rep r op ()
copyDWIMFix (Param LetDecMem -> VName
forall dec. Param dec -> VName
paramName Param LetDecMem
p) [] (VName -> SubExp
Var VName
arr) [TExp Int64
gtid TExp Int64 -> TExp Int64 -> TExp Int64
forall a. Num a => a -> a -> a
- TExp Int64
behind TExp Int64 -> TExp Int64 -> TExp Int64
forall a. Num a => a -> a -> a
+ TExp Int64
arrs_full_size]

    writeResult :: Param LetDecMem
-> Param LetDecMem -> VName -> ImpM GPUMem KernelEnv KernelOp ()
writeResult Param LetDecMem
x Param LetDecMem
y VName
arr = do
      Bool
-> ImpM GPUMem KernelEnv KernelOp ()
-> ImpM GPUMem KernelEnv KernelOp ()
forall (f :: * -> *). Applicative f => Bool -> f () -> f ()
when (Param LetDecMem -> Bool
forall p. Typed p => Param p -> Bool
isPrimParam Param LetDecMem
x) (ImpM GPUMem KernelEnv KernelOp ()
 -> ImpM GPUMem KernelEnv KernelOp ())
-> ImpM GPUMem KernelEnv KernelOp ()
-> ImpM GPUMem KernelEnv KernelOp ()
forall a b. (a -> b) -> a -> b
$
        VName
-> [TExp Int64]
-> SubExp
-> [TExp Int64]
-> ImpM GPUMem KernelEnv KernelOp ()
forall rep r op.
VName -> [TExp Int64] -> SubExp -> [TExp Int64] -> ImpM rep r op ()
copyDWIMFix VName
arr [TExp Int64
ltid] (VName -> SubExp
Var (VName -> SubExp) -> VName -> SubExp
forall a b. (a -> b) -> a -> b
$ Param LetDecMem -> VName
forall dec. Param dec -> VName
paramName Param LetDecMem
x) []
      VName
-> [DimIndex (TExp Int64)]
-> SubExp
-> [DimIndex (TExp Int64)]
-> ImpM GPUMem KernelEnv KernelOp ()
forall rep r op.
VName
-> [DimIndex (TExp Int64)]
-> SubExp
-> [DimIndex (TExp Int64)]
-> ImpM rep r op ()
copyDWIM (Param LetDecMem -> VName
forall dec. Param dec -> VName
paramName Param LetDecMem
y) [] (VName -> SubExp
Var (VName -> SubExp) -> VName -> SubExp
forall a b. (a -> b) -> a -> b
$ Param LetDecMem -> VName
forall dec. Param dec -> VName
paramName Param LetDecMem
x) []

-- | Calculate the number of u64 words needed to store n bits
bitArrayWords :: Imp.KernelConstExp -> Imp.KernelConstExp
bitArrayWords :: KernelConstExp -> KernelConstExp
bitArrayWords KernelConstExp
n = TPrimExp Int64 KernelConst -> KernelConstExp
forall {k} (t :: k) v. TPrimExp t v -> PrimExp v
untyped (TPrimExp Int64 KernelConst -> KernelConstExp)
-> TPrimExp Int64 KernelConst -> KernelConstExp
forall a b. (a -> b) -> a -> b
$ KernelConstExp -> TPrimExp Int64 KernelConst
forall v. PrimExp v -> TPrimExp Int64 v
isInt64 KernelConstExp
n TPrimExp Int64 KernelConst
-> TPrimExp Int64 KernelConst -> TPrimExp Int64 KernelConst
forall e. IntegralExp e => e -> e -> e
`ceilDiv` TPrimExp Int64 KernelConst
64

-- | Set a bit in a bit array stored as u64 words
setBitInBitArray :: Imp.TExp Int64 -> VName -> Imp.TExp Int64 -> Imp.TExp Bool -> InKernelGen ()
setBitInBitArray :: TExp Int64
-> VName
-> TExp Int64
-> TPrimExp Bool VName
-> ImpM GPUMem KernelEnv KernelOp ()
setBitInBitArray TExp Int64
chunk VName
bit_array TExp Int64
idx TPrimExp Bool VName
bool_val = do
  word_idx <- Name -> TExp Int64 -> ImpM GPUMem KernelEnv KernelOp (TV Int64)
forall {k} (t :: k) rep r op.
Name -> TExp t -> ImpM rep r op (TV t)
dPrimV Name
"word_idx" (TExp Int64
0 :: Imp.TExp Int64)
  bit_idx <- dPrimV "bit_idx" idx

  sIf
    (chunk .<=. 64)
    ( do
        word_idx <-- 0
        bit_idx <-- idx
    )
    ( do
        word_idx <-- idx `quot` 64
        bit_idx <-- idx `rem` 64
    )

  -- Read current word
  current_word <- dPrimV "current_word" (0 :: Imp.TExp Int64)
  copyDWIMFix (tvVar current_word) [] (Var bit_array) [tvExp word_idx]

  bool_as_int <- dPrimVE "bool_as_int" $ fromBoolExp bool_val

  -- Create mask and update word
  let bit_mask = TExp Int64
1 TExp Int64 -> TExp Int64 -> TExp Int64
forall {k} v (t :: k).
Eq v =>
TPrimExp t v -> TPrimExp t v -> TPrimExp t v
.<<. TV Int64 -> TExp Int64
forall {k} (t :: k). TV t -> TExp t
tvExp TV Int64
bit_idx
      cleared_word = TV Int64 -> TExp Int64
forall {k} (t :: k). TV t -> TExp t
tvExp TV Int64
current_word TExp Int64 -> TExp Int64 -> TExp Int64
forall {k} v (t :: k).
Eq v =>
TPrimExp t v -> TPrimExp t v -> TPrimExp t v
.&. (TExp Int64
bit_mask TExp Int64 -> TExp Int64 -> TExp Int64
forall {k} v (t :: k).
Eq v =>
TPrimExp t v -> TPrimExp t v -> TPrimExp t v
.^. (-TExp Int64
1))
      set_bit = (TExp Int64
bool_as_int TExp Int64 -> TExp Int64 -> TExp Int64
forall {k} v (t :: k).
Eq v =>
TPrimExp t v -> TPrimExp t v -> TPrimExp t v
.&. TExp Int64
1) TExp Int64 -> TExp Int64 -> TExp Int64
forall {k} v (t :: k).
Eq v =>
TPrimExp t v -> TPrimExp t v -> TPrimExp t v
.<<. TV Int64 -> TExp Int64
forall {k} (t :: k). TV t -> TExp t
tvExp TV Int64
bit_idx
      new_word = TExp Int64
cleared_word TExp Int64 -> TExp Int64 -> TExp Int64
forall {k} v (t :: k).
Eq v =>
TPrimExp t v -> TPrimExp t v -> TPrimExp t v
.|. TExp Int64
set_bit

  new_word_var <- dPrimV "new_word" new_word
  copyDWIMFix bit_array [tvExp word_idx] (Var $ tvVar new_word_var) []

-- | Get a bit from a bit array stored as u64 words, storing result in destination
getBitFromBitArray :: Imp.TExp Int64 -> VName -> VName -> Imp.TExp Int64 -> InKernelGen ()
getBitFromBitArray :: TExp Int64
-> VName
-> VName
-> TExp Int64
-> ImpM GPUMem KernelEnv KernelOp ()
getBitFromBitArray TExp Int64
chunk VName
dest VName
bit_array TExp Int64
idx = do
  word_idx <- Name -> TExp Int64 -> ImpM GPUMem KernelEnv KernelOp (TV Int64)
forall {k} (t :: k) rep r op.
Name -> TExp t -> ImpM rep r op (TV t)
dPrimV Name
"word_idx" (TExp Int64
0 :: Imp.TExp Int64)
  bit_idx <- dPrimV "bit_idx" idx

  sIf
    (chunk .<=. 64)
    ( do
        word_idx <-- 0
        bit_idx <-- idx
    )
    ( do
        word_idx <-- idx `quot` 64
        bit_idx <-- idx `rem` 64
    )

  -- Read word containing the bit
  word <- dPrimV "bit_word_read" (0 :: Imp.TExp Int64)
  copyDWIMFix (tvVar word) [] (Var bit_array) [tvExp word_idx]

  -- Extract bit: (word >> bitIdx) & 1
  let extracted_bit = (TV Int64 -> TExp Int64
forall {k} (t :: k). TV t -> TExp t
tvExp TV Int64
word TExp Int64 -> TExp Int64 -> TExp Int64
forall {k} v (t :: k).
Eq v =>
TPrimExp t v -> TPrimExp t v -> TPrimExp t v
.>>. TV Int64 -> TExp Int64
forall {k} (t :: k). TV t -> TExp t
tvExp TV Int64
bit_idx) TExp Int64 -> TExp Int64 -> TExp Int64
forall {k} v (t :: k).
Eq v =>
TPrimExp t v -> TPrimExp t v -> TPrimExp t v
.&. TExp Int64
1
      bool_val = TExp Int64
extracted_bit TExp Int64 -> TExp Int64 -> TPrimExp Bool VName
forall {k} v (t :: k).
Eq v =>
TPrimExp t v -> TPrimExp t v -> TPrimExp Bool v
.==. TExp Int64
1

  bool_var <- dPrimV "bool_val" bool_val
  copyDWIMFix dest [] (Var $ tvVar bool_var) []

-- | Helper to determine if a type should use bit array representation
shouldUseBitArray :: Type -> Bool
shouldUseBitArray :: Type -> Bool
shouldUseBitArray Type
t = Type -> Bool
forall shape o. TypeBase shape o -> Bool
primType Type
t Bool -> Bool -> Bool
&& Type -> PrimType
forall shape o. TypeBase shape o -> PrimType
elemType Type
t PrimType -> PrimType -> Bool
forall a. Eq a => a -> a -> Bool
== PrimType
Bool

earlyReturn ::
  Pat LetDecMem ->
  SegBinOp GPUMem ->
  SegPostOp GPUMem ->
  ([[PatElem LetDecMem]], [Bool])
earlyReturn :: Pat LetDecMem
-> SegBinOp GPUMem
-> SegPostOp GPUMem
-> ([[PatElem LetDecMem]], [Bool])
earlyReturn Pat LetDecMem
pat SegBinOp GPUMem
scan_op SegPostOp GPUMem
post_op = ([[PatElem LetDecMem]]
early_write_pats, [Bool]
res_flags)
  where
    post_lam :: GLambda GPUMem Type
post_lam = SegPostOp GPUMem -> GLambda GPUMem Type
forall rep. SegPostOp rep -> Lambda rep
segPostOpLambda SegPostOp GPUMem
post_op
    num_scan_res :: Int
num_scan_res = [SubExp] -> Int
forall a. [a] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length ([SubExp] -> Int) -> [SubExp] -> Int
forall a b. (a -> b) -> a -> b
$ SegBinOp GPUMem -> [SubExp]
forall rep. SegBinOp rep -> [SubExp]
segBinOpNeutral SegBinOp GPUMem
scan_op
    earlyWrite :: Int -> Param LetDecMem -> [(SubExpRes, PatElem LetDecMem)]
earlyWrite Int
i Param LetDecMem
par =
      if Int
num_scan_res Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
< Int
i Bool -> Bool -> Bool
&& Param LetDecMem -> VName
forall dec. Param dec -> VName
paramName Param LetDecMem
par VName -> Names -> Bool
`notNameIn` Names
free_in_post
        then
          ((SubExpRes, PatElem LetDecMem) -> Bool)
-> [(SubExpRes, PatElem LetDecMem)]
-> [(SubExpRes, PatElem LetDecMem)]
forall a. (a -> Bool) -> [a] -> [a]
filter ((Maybe VName -> Maybe VName -> Bool
forall a. Eq a => a -> a -> Bool
== VName -> Maybe VName
forall a. a -> Maybe a
Just (Param LetDecMem -> VName
forall dec. Param dec -> VName
paramName Param LetDecMem
par)) (Maybe VName -> Bool)
-> ((SubExpRes, PatElem LetDecMem) -> Maybe VName)
-> (SubExpRes, PatElem LetDecMem)
-> Bool
forall b c a. (b -> c) -> (a -> b) -> a -> c
. SubExpRes -> Maybe VName
subExpResVName (SubExpRes -> Maybe VName)
-> ((SubExpRes, PatElem LetDecMem) -> SubExpRes)
-> (SubExpRes, PatElem LetDecMem)
-> Maybe VName
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (SubExpRes, PatElem LetDecMem) -> SubExpRes
forall a b. (a, b) -> a
fst)
            ([(SubExpRes, PatElem LetDecMem)]
 -> [(SubExpRes, PatElem LetDecMem)])
-> ([PatElem LetDecMem] -> [(SubExpRes, PatElem LetDecMem)])
-> [PatElem LetDecMem]
-> [(SubExpRes, PatElem LetDecMem)]
forall b c a. (b -> c) -> (a -> b) -> a -> c
. [SubExpRes]
-> [PatElem LetDecMem] -> [(SubExpRes, PatElem LetDecMem)]
forall a b. [a] -> [b] -> [(a, b)]
zip (Body GPUMem -> [SubExpRes]
forall rep res. GBody rep res -> [res]
bodyResult (Body GPUMem -> [SubExpRes]) -> Body GPUMem -> [SubExpRes]
forall a b. (a -> b) -> a -> b
$ GLambda GPUMem Type -> Body GPUMem
forall rep t. GLambda rep t -> Body rep
lambdaBody GLambda GPUMem Type
post_lam)
            ([PatElem LetDecMem] -> [(SubExpRes, PatElem LetDecMem)])
-> [PatElem LetDecMem] -> [(SubExpRes, PatElem LetDecMem)]
forall a b. (a -> b) -> a -> b
$ Pat LetDecMem -> [PatElem LetDecMem]
forall dec. Pat dec -> [PatElem dec]
patElems Pat LetDecMem
pat
        else []
    early_writes :: [[(SubExpRes, PatElem LetDecMem)]]
early_writes = (Int -> Param LetDecMem -> [(SubExpRes, PatElem LetDecMem)])
-> [Int] -> [Param LetDecMem] -> [[(SubExpRes, PatElem LetDecMem)]]
forall a b c. (a -> b -> c) -> [a] -> [b] -> [c]
zipWith Int -> Param LetDecMem -> [(SubExpRes, PatElem LetDecMem)]
earlyWrite [Int
1 ..] ([Param LetDecMem] -> [[(SubExpRes, PatElem LetDecMem)]])
-> [Param LetDecMem] -> [[(SubExpRes, PatElem LetDecMem)]]
forall a b. (a -> b) -> a -> b
$ GLambda GPUMem Type -> [LParam GPUMem]
forall rep t. GLambda rep t -> [LParam rep]
lambdaParams GLambda GPUMem Type
post_lam
    res_flags :: [Bool]
res_flags = (PatElem LetDecMem -> [PatElem LetDecMem] -> Bool
forall a. Eq a => a -> [a] -> Bool
forall (t :: * -> *) a. (Foldable t, Eq a) => a -> t a -> Bool
`elem` [[PatElem LetDecMem]] -> [PatElem LetDecMem]
forall a. Monoid a => [a] -> a
mconcat [[PatElem LetDecMem]]
early_write_pats) (PatElem LetDecMem -> Bool) -> [PatElem LetDecMem] -> [Bool]
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Pat LetDecMem -> [PatElem LetDecMem]
forall dec. Pat dec -> [PatElem dec]
patElems Pat LetDecMem
pat
    early_write_pats :: [[PatElem LetDecMem]]
early_write_pats = Int -> [[PatElem LetDecMem]] -> [[PatElem LetDecMem]]
forall a. Int -> [a] -> [a]
drop Int
num_scan_res ([[PatElem LetDecMem]] -> [[PatElem LetDecMem]])
-> [[PatElem LetDecMem]] -> [[PatElem LetDecMem]]
forall a b. (a -> b) -> a -> b
$ ((SubExpRes, PatElem LetDecMem) -> PatElem LetDecMem)
-> [(SubExpRes, PatElem LetDecMem)] -> [PatElem LetDecMem]
forall a b. (a -> b) -> [a] -> [b]
map (SubExpRes, PatElem LetDecMem) -> PatElem LetDecMem
forall a b. (a, b) -> b
snd ([(SubExpRes, PatElem LetDecMem)] -> [PatElem LetDecMem])
-> [[(SubExpRes, PatElem LetDecMem)]] -> [[PatElem LetDecMem]]
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> [[(SubExpRes, PatElem LetDecMem)]]
early_writes
    free_in_post :: Names
free_in_post = Stms GPUMem -> Names
forall a. FreeIn a => a -> Names
freeIn (Stms GPUMem -> Names) -> Stms GPUMem -> Names
forall a b. (a -> b) -> a -> b
$ Body GPUMem -> Stms GPUMem
forall rep res. GBody rep res -> Stms rep
bodyStms (Body GPUMem -> Stms GPUMem) -> Body GPUMem -> Stms GPUMem
forall a b. (a -> b) -> a -> b
$ GLambda GPUMem Type -> Body GPUMem
forall rep t. GLambda rep t -> Body rep
lambdaBody GLambda GPUMem Type
post_lam

-- | Compile 'SegScan' instance to host-level code with calls to a
-- single-pass kernel.
compileSegScan ::
  Pat LetDecMem ->
  SegLevel ->
  SegSpace ->
  [Type] ->
  SegBinOp GPUMem ->
  KernelBody GPUMem ->
  SegPostOp GPUMem ->
  CallKernelGen ()
compileSegScan :: Pat LetDecMem
-> SegLevel
-> SegSpace
-> [Type]
-> SegBinOp GPUMem
-> KernelBody GPUMem
-> SegPostOp GPUMem
-> CallKernelGen ()
compileSegScan Pat LetDecMem
pat SegLevel
lvl SegSpace
space [Type]
ts SegBinOp GPUMem
scan_op KernelBody GPUMem
map_kbody SegPostOp GPUMem
post_op = do
  attrs <- SegLevel -> CallKernelGen KernelAttrs
lvlKernelAttrs SegLevel
lvl
  let scanop_nes = SegBinOp GPUMem -> [SubExp]
forall rep. SegBinOp rep -> [SubExp]
segBinOpNeutral SegBinOp GPUMem
scan_op

      n = [TExp Int64] -> TExp Int64
forall a. Num a => [a] -> a
forall (t :: * -> *) a. (Foldable t, Num a) => t a -> a
product ([TExp Int64] -> TExp Int64) -> [TExp Int64] -> TExp Int64
forall a b. (a -> b) -> a -> b
$ (SubExp -> TExp Int64) -> [SubExp] -> [TExp Int64]
forall a b. (a -> b) -> [a] -> [b]
map SubExp -> TExp Int64
pe64 ([SubExp] -> [TExp Int64]) -> [SubExp] -> [TExp Int64]
forall a b. (a -> b) -> a -> b
$ SegSpace -> [SubExp]
segSpaceDims SegSpace
space

      scan_tys' = GLambda GPUMem Type -> [Type]
forall rep t. GLambda rep t -> [t]
lambdaReturnType (GLambda GPUMem Type -> [Type]) -> GLambda GPUMem Type -> [Type]
forall a b. (a -> b) -> a -> b
$ SegBinOp GPUMem -> GLambda GPUMem Type
forall rep. SegBinOp rep -> Lambda rep
segBinOpLambda SegBinOp GPUMem
scan_op
      map_tys' = Int -> [Type] -> [Type]
forall a. Int -> [a] -> [a]
drop ([SubExp] -> Int
forall a. [a] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length ([SubExp] -> Int) -> [SubExp] -> Int
forall a b. (a -> b) -> a -> b
$ SegBinOp GPUMem -> [SubExp]
forall rep. SegBinOp rep -> [SubExp]
segBinOpNeutral SegBinOp GPUMem
scan_op) [Type]
ts

      scan_tys = (Type -> PrimType) -> [Type] -> [PrimType]
forall a b. (a -> b) -> [a] -> [b]
map Type -> PrimType
forall shape o. TypeBase shape o -> PrimType
elemType [Type]
scan_tys'

      tblock_size_e = SubExp -> TExp Int64
pe64 (SubExp -> TExp Int64) -> SubExp -> TExp Int64
forall a b. (a -> b) -> a -> b
$ Count BlockSize SubExp -> SubExp
forall {k} (u :: k) e. Count u e -> e
unCount (Count BlockSize SubExp -> SubExp)
-> Count BlockSize SubExp -> SubExp
forall a b. (a -> b) -> a -> b
$ KernelAttrs -> Count BlockSize SubExp
kAttrBlockSize KernelAttrs
attrs
      num_phys_blocks_e = SubExp -> TExp Int64
pe64 (SubExp -> TExp Int64) -> SubExp -> TExp Int64
forall a b. (a -> b) -> a -> b
$ Count NumBlocks SubExp -> SubExp
forall {k} (u :: k) e. Count u e -> e
unCount (Count NumBlocks SubExp -> SubExp)
-> Count NumBlocks SubExp -> SubExp
forall a b. (a -> b) -> a -> b
$ KernelAttrs -> Count NumBlocks SubExp
kAttrNumBlocks KernelAttrs
attrs
      (early_write_pats, res_flags) = earlyReturn pat scan_op post_op

  let chunk_const = [Type] -> KernelConstExp
getScanChunkSize [Type]
scan_tys'
  chunk_v <- dPrim "chunk_size"
  let chunk_name = Text -> Name
nameFromText (Text -> Name) -> Text -> Name
forall a b. (a -> b) -> a -> b
$ VName -> Text
forall a. Pretty a => a -> Text
prettyText (VName -> Text) -> VName -> Text
forall a b. (a -> b) -> a -> b
$ TV Int64 -> VName
forall {k} (t :: k). TV t -> VName
tvVar TV Int64
chunk_v
  addTuningParam chunk_name Nothing
  emit . Imp.GetUserParam (tvVar chunk_v) chunk_name . isInt64
    =<< kernelConstToExp chunk_const
  let chunk_constexp = KernelConst -> PrimType -> KernelConstExp
forall v. v -> PrimType -> PrimExp v
LeafExp (Name -> KernelConstExp -> KernelConst
Imp.SizeUserParam Name
chunk_name KernelConstExp
chunk_const) PrimType
int64
      chunk = TV Int64 -> TExp Int64
forall {k} (t :: k). TV t -> TExp t
tvExp TV Int64
chunk_v

  let num_words_const = KernelConstExp -> KernelConstExp
bitArrayWords KernelConstExp
chunk_const
  num_words <-
    dPrimV "num_bit_words" . isInt64 =<< kernelConstToExp num_words_const

  num_virt_blocks <-
    tvSize <$> dPrimV "num_virt_blocks" (n `ceilDiv` (tblock_size_e * chunk))
  let num_virt_blocks_e = SubExp -> TExp Int64
pe64 SubExp
num_virt_blocks

  num_virt_threads <-
    dPrimVE "num_virt_threads" $ num_virt_blocks_e * tblock_size_e

  let (gtids, dims) = unzip $ unSegSpace space
      dims' = (SubExp -> TExp Int64) -> [SubExp] -> [TExp Int64]
forall a b. (a -> b) -> [a] -> [b]
map SubExp -> TExp Int64
pe64 [SubExp]
dims
      segmented = [TExp Int64] -> Int
forall a. [a] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length [TExp Int64]
dims' Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
> Int
1
      not_segmented_e = Bool -> TPrimExp Bool VName
forall v. Bool -> TPrimExp Bool v
fromBool (Bool -> TPrimExp Bool VName) -> Bool -> TPrimExp Bool VName
forall a b. (a -> b) -> a -> b
$ Bool -> Bool
not Bool
segmented
      segment_size = [TExp Int64] -> TExp Int64
forall a. HasCallStack => [a] -> a
last [TExp Int64]
dims'

  emit $ Imp.DebugPrint "Sequential elements per thread (chunk)" $ Just $ untyped chunk

  statusFlags <- sAllocArray "status_flags" int8 (Shape [num_virt_blocks]) (Space "device")
  sReplicate statusFlags $ intConst Int8 statusX

  (aggregateArrays, incprefixArrays) <-
    fmap unzip $
      forM scan_tys $ \PrimType
ty ->
        (,)
          (VName -> VName -> (VName, VName))
-> ImpM GPUMem HostEnv HostOp VName
-> ImpM GPUMem HostEnv HostOp (VName -> (VName, VName))
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Name
-> PrimType
-> ShapeBase SubExp
-> Space
-> ImpM GPUMem HostEnv HostOp VName
forall rep r op.
Name
-> PrimType -> ShapeBase SubExp -> Space -> ImpM rep r op VName
sAllocArray Name
"aggregates" PrimType
ty ([SubExp] -> ShapeBase SubExp
forall d. [d] -> ShapeBase d
Shape [SubExp
num_virt_blocks]) (String -> Space
Space String
"device")
          ImpM GPUMem HostEnv HostOp (VName -> (VName, VName))
-> ImpM GPUMem HostEnv HostOp VName
-> ImpM GPUMem HostEnv HostOp (VName, VName)
forall a b.
ImpM GPUMem HostEnv HostOp (a -> b)
-> ImpM GPUMem HostEnv HostOp a -> ImpM GPUMem HostEnv HostOp b
forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> Name
-> PrimType
-> ShapeBase SubExp
-> Space
-> ImpM GPUMem HostEnv HostOp VName
forall rep r op.
Name
-> PrimType -> ShapeBase SubExp -> Space -> ImpM rep r op VName
sAllocArray Name
"incprefixes" PrimType
ty ([SubExp] -> ShapeBase SubExp
forall d. [d] -> ShapeBase d
Shape [SubExp
num_virt_blocks]) (String -> Space
Space String
"device")

  global_id <- genZeroes "global_dynid" 1

  let attrs' =
        KernelAttrs
attrs
          { kAttrConstExps =
              M.fromList
                [ (tvVar chunk_v, chunk_constexp),
                  (tvVar num_words, num_words_const)
                ]
          }

  map_global_chunks <-
    forM (zip map_tys' early_write_pats) $ \(Type
t, [PatElem LetDecMem]
pats) ->
      if Type -> Bool
forall shape o. TypeBase shape o -> Bool
isAcc Type
t Bool -> Bool -> Bool
|| Type -> Bool
forall shape o. TypeBase shape o -> Bool
primType Type
t Bool -> Bool -> Bool
|| Bool -> Bool
not ([PatElem LetDecMem] -> Bool
forall a. [a] -> Bool
forall (t :: * -> *) a. Foldable t => t a -> Bool
null [PatElem LetDecMem]
pats)
        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
-> ShapeBase SubExp
-> Space
-> ImpM GPUMem HostEnv HostOp VName
forall rep r op.
Name
-> PrimType -> ShapeBase SubExp -> Space -> ImpM rep r op VName
sAllocArray
              Name
"global"
              (Type -> PrimType
forall shape o. TypeBase shape o -> PrimType
elemType Type
t)
              ( [SubExp] -> ShapeBase SubExp
forall d. [d] -> ShapeBase d
Shape
                  [ Count NumBlocks SubExp -> SubExp
forall {k} (u :: k) e. Count u e -> e
unCount (Count NumBlocks SubExp -> SubExp)
-> Count NumBlocks SubExp -> SubExp
forall a b. (a -> b) -> a -> b
$ KernelAttrs -> Count NumBlocks SubExp
kAttrNumBlocks KernelAttrs
attrs,
                    Count BlockSize SubExp -> SubExp
forall {k} (u :: k) e. Count u e -> e
unCount (Count BlockSize SubExp -> SubExp)
-> Count BlockSize SubExp -> SubExp
forall a b. (a -> b) -> a -> b
$ KernelAttrs -> Count BlockSize SubExp
kAttrBlockSize KernelAttrs
attrs,
                    TV Int64 -> SubExp
forall {k} (t :: k). TV t -> SubExp
tvSize TV Int64
chunk_v
                  ]
                  ShapeBase SubExp -> ShapeBase SubExp -> ShapeBase SubExp
forall a. Semigroup a => a -> a -> a
<> Type -> ShapeBase SubExp
forall shape o. ArrayShape shape => TypeBase shape o -> shape
arrayShape Type
t
              )
              (String -> Space
Space String
"device")
  sKernelThread "segscan" (segFlat space) attrs' $ do
    chunk32 <- dPrimVE "chunk_size_32b" $ sExt32 $ tvExp chunk_v

    constants <- kernelConstants <$> askEnv

    let ltid32 = KernelConstants -> TExp Int32
kernelLocalThreadId KernelConstants
constants
        ltid = TExp Int32 -> TExp Int64
forall {k} (t :: k) v. IntExp t => TPrimExp t v -> TPrimExp Int64 v
sExt64 TExp Int32
ltid32

    (sharedId, transposedArrays, prefixArrays, warpscan, exchanges) <-
      createLocalArrays (kAttrBlockSize attrs) (tvSize chunk_v) scan_tys

    -- We wrap the entire kernel body in a virtualisation loop to
    -- handle the case where we do not have enough thread blocks to
    -- cover the iteration space. Dynamic block indexing has no
    -- implication on this, since each block simply fetches a new
    -- dynamic ID upon entry into the virtualisation loop.
    --
    -- We could use virtualiseBlocks, but this introduces a barrier which is
    -- redundant in this case, and also we don't need to base virtual block IDs
    -- on the loop variable, but rather on the dynamic IDs.
    phys_block_id <- dPrim "phys_block_id"
    sOp $ Imp.GetBlockId (tvVar phys_block_id) 0
    iters <-
      dPrimVE "virtloop_bound" $
        (num_virt_blocks_e - tvExp phys_block_id) `ceilDiv` num_phys_blocks_e

    sFor "virtloop_i" iters $ const $ do
      dyn_id <- dPrim "dynamic_id"
      sComment "First thread in block fetches this block's dynamic_id" $ do
        sWhen (ltid32 .==. 0) $ do
          (globalIdMem, _, globalIdOff) <- fullyIndexArray global_id [0]
          sOp $
            Imp.Atomic DefaultSpace $
              Imp.AtomicAdd
                Int32
                (tvVar dyn_id)
                globalIdMem
                (Count $ unCount globalIdOff)
                (untyped (1 :: Imp.TExp Int32))
          sComment "Set dynamic id for this block" $ do
            copyDWIMFix sharedId [0] (tvSize dyn_id) []

          sComment "First thread in last (virtual) block resets global dynamic_id" $ do
            sWhen (tvExp dyn_id .==. num_virt_blocks_e - 1) $
              copyDWIMFix global_id [0] (intConst Int32 0) []

      let local_barrier = Fence -> KernelOp
Imp.Barrier Fence
Imp.FenceLocal
          local_fence = Fence -> KernelOp
Imp.MemFence Fence
Imp.FenceLocal
          global_fence = Fence -> KernelOp
Imp.MemFence Fence
Imp.FenceGlobal

      sOp local_barrier
      copyDWIMFix (tvVar dyn_id) [] (Var sharedId) [0]
      sOp local_barrier

      block_offset <-
        dPrimVE "block_offset" $
          sExt64 (tvExp dyn_id) * chunk * tblock_size_e
      sgm_idx <- dPrimVE "sgm_idx" $ block_offset `mod` segment_size
      boundary <-
        dPrimVE "boundary" $
          sExt32 $
            sMin64 (chunk * tblock_size_e) (segment_size - sgm_idx)
      segsize_compact <-
        dPrimVE "segsize_compact" $
          sExt32 $
            sMin64 (chunk * tblock_size_e) segment_size
      scan_private_chunks <-
        forM scan_tys $ \PrimType
ty ->
          Name
-> PrimType
-> ShapeBase SubExp
-> Space
-> ImpM GPUMem KernelEnv KernelOp VName
forall rep r op.
Name
-> PrimType -> ShapeBase SubExp -> Space -> ImpM rep r op VName
sAllocArray
            Name
"private"
            PrimType
ty
            ([SubExp] -> ShapeBase SubExp
forall d. [d] -> ShapeBase d
Shape [TV Int64 -> SubExp
forall {k} (t :: k). TV t -> SubExp
tvSize TV Int64
chunk_v])
            ([SubExp] -> PrimType -> Space
ScalarSpace [TV Int64 -> SubExp
forall {k} (t :: k). TV t -> SubExp
tvSize TV Int64
chunk_v] PrimType
ty)

      map_private_chunks <-
        forM (zip map_tys' early_write_pats) $ \(Type
t, [PatElem LetDecMem]
pats) ->
          if Type -> Bool
forall shape o. TypeBase shape o -> Bool
isAcc Type
t Bool -> Bool -> Bool
|| Bool -> Bool
not (Type -> Bool
forall shape o. TypeBase shape o -> Bool
primType Type
t) Bool -> Bool -> Bool
|| Bool -> Bool
not ([PatElem LetDecMem] -> Bool
forall a. [a] -> Bool
forall (t :: * -> *) a. Foldable t => t a -> Bool
null [PatElem LetDecMem]
pats)
            then Maybe VName -> ImpM GPUMem KernelEnv KernelOp (Maybe VName)
forall a. a -> ImpM GPUMem KernelEnv KernelOp a
forall (f :: * -> *) a. Applicative f => a -> f a
pure Maybe VName
forall a. Maybe a
Nothing
            else
              if Type -> Bool
shouldUseBitArray Type
t
                then do
                  VName -> Maybe VName
forall a. a -> Maybe a
Just
                    (VName -> Maybe VName)
-> ImpM GPUMem KernelEnv KernelOp VName
-> ImpM GPUMem KernelEnv KernelOp (Maybe VName)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Name
-> PrimType
-> ShapeBase SubExp
-> Space
-> ImpM GPUMem KernelEnv KernelOp VName
forall rep r op.
Name
-> PrimType -> ShapeBase SubExp -> Space -> ImpM rep r op VName
sAllocArray
                      Name
"private_bits"
                      PrimType
int64
                      ([SubExp] -> ShapeBase SubExp
forall d. [d] -> ShapeBase d
Shape [TV Int64 -> SubExp
forall {k} (t :: k). TV t -> SubExp
tvSize TV Int64
num_words])
                      ([SubExp] -> PrimType -> Space
ScalarSpace [TV Int64 -> SubExp
forall {k} (t :: k). TV t -> SubExp
tvSize TV Int64
num_words] PrimType
int64)
                else
                  VName -> Maybe VName
forall a. a -> Maybe a
Just
                    (VName -> Maybe VName)
-> ImpM GPUMem KernelEnv KernelOp VName
-> ImpM GPUMem KernelEnv KernelOp (Maybe VName)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Name
-> PrimType
-> ShapeBase SubExp
-> Space
-> ImpM GPUMem KernelEnv KernelOp VName
forall rep r op.
Name
-> PrimType -> ShapeBase SubExp -> Space -> ImpM rep r op VName
sAllocArray
                      Name
"private"
                      (Type -> PrimType
forall shape o. TypeBase shape o -> PrimType
elemType Type
t)
                      ([SubExp] -> ShapeBase SubExp
forall d. [d] -> ShapeBase d
Shape [TV Int64 -> SubExp
forall {k} (t :: k). TV t -> SubExp
tvSize TV Int64
chunk_v])
                      ([SubExp] -> PrimType -> Space
ScalarSpace [TV Int64 -> SubExp
forall {k} (t :: k). TV t -> SubExp
tvSize TV Int64
chunk_v] (Type -> PrimType
forall shape o. TypeBase shape o -> PrimType
elemType Type
t))

      thd_offset <- dPrimVE "thd_offset" $ block_offset + ltid

      sComment "Load and map" $
        sFor "i" chunk $ \TExp Int64
i -> do
          -- The map's input index
          virt_tid <- Name -> TExp Int64 -> ImpM GPUMem KernelEnv KernelOp (TExp Int64)
forall {k} (t :: k) rep r op.
Name -> TExp t -> ImpM rep r op (TExp t)
dPrimVE Name
"virt_tid" (TExp Int64 -> ImpM GPUMem KernelEnv KernelOp (TExp Int64))
-> TExp Int64 -> ImpM GPUMem KernelEnv KernelOp (TExp Int64)
forall a b. (a -> b) -> a -> b
$ TExp Int64
thd_offset TExp Int64 -> TExp Int64 -> TExp Int64
forall a. Num a => a -> a -> a
+ TExp Int64
i TExp Int64 -> TExp Int64 -> TExp Int64
forall a. Num a => a -> a -> a
* TExp Int64
tblock_size_e
          dIndexSpace (zip gtids dims') virt_tid
          -- Perform the map
          let 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
map_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
scan_op]) ([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
map_kbody

                  -- Write map results to memory.
                  [(Maybe VName, Maybe VName, SubExp, Type, [PatElem LetDecMem])]
-> ((Maybe VName, Maybe VName, SubExp, Type, [PatElem LetDecMem])
    -> ImpM GPUMem KernelEnv KernelOp ())
-> ImpM GPUMem KernelEnv KernelOp ()
forall (t :: * -> *) (m :: * -> *) a b.
(Foldable t, Monad m) =>
t a -> (a -> m b) -> m ()
forM_ ([Maybe VName]
-> [Maybe VName]
-> [SubExp]
-> [Type]
-> [[PatElem LetDecMem]]
-> [(Maybe VName, Maybe VName, SubExp, Type, [PatElem LetDecMem])]
forall a b c d e.
[a] -> [b] -> [c] -> [d] -> [e] -> [(a, b, c, d, e)]
zip5 [Maybe VName]
map_private_chunks [Maybe VName]
map_global_chunks ((KernelResult -> SubExp) -> [KernelResult] -> [SubExp]
forall a b. (a -> b) -> [a] -> [b]
map KernelResult -> SubExp
kernelResultSubExp [KernelResult]
map_res) [Type]
map_tys' [[PatElem LetDecMem]]
early_write_pats) (((Maybe VName, Maybe VName, SubExp, Type, [PatElem LetDecMem])
  -> ImpM GPUMem KernelEnv KernelOp ())
 -> ImpM GPUMem KernelEnv KernelOp ())
-> ((Maybe VName, Maybe VName, SubExp, Type, [PatElem LetDecMem])
    -> ImpM GPUMem KernelEnv KernelOp ())
-> ImpM GPUMem KernelEnv KernelOp ()
forall a b. (a -> b) -> a -> b
$
                    \(Maybe VName
priv_dest, Maybe VName
glob_dest, SubExp
src, Type
ty, [PatElem LetDecMem]
pats) -> do
                      case Maybe VName
priv_dest of
                        Just VName
d
                          | Type -> Bool
shouldUseBitArray Type
ty ->
                              TExp Int64
-> VName
-> TExp Int64
-> TPrimExp Bool VName
-> ImpM GPUMem KernelEnv KernelOp ()
setBitInBitArray TExp Int64
chunk VName
d TExp Int64
i (TPrimExp Bool VName -> ImpM GPUMem KernelEnv KernelOp ())
-> TPrimExp Bool VName -> ImpM GPUMem KernelEnv KernelOp ()
forall a b. (a -> b) -> a -> b
$ PrimExp VName -> TPrimExp Bool VName
forall v. PrimExp v -> TPrimExp Bool v
isBool (PrimExp VName -> TPrimExp Bool VName)
-> PrimExp VName -> TPrimExp Bool VName
forall a b. (a -> b) -> a -> b
$ PrimType -> SubExp -> PrimExp VName
forall a. ToExp a => PrimType -> a -> PrimExp VName
toExp' PrimType
Bool SubExp
src
                        Just VName
d ->
                          VName
-> [TExp Int64]
-> SubExp
-> [TExp Int64]
-> ImpM GPUMem KernelEnv KernelOp ()
forall rep r op.
VName -> [TExp Int64] -> SubExp -> [TExp Int64] -> ImpM rep r op ()
copyDWIMFix VName
d [TExp Int64
i] SubExp
src []
                        Maybe VName
Nothing -> () -> ImpM GPUMem KernelEnv KernelOp ()
forall a. a -> ImpM GPUMem KernelEnv KernelOp a
forall (f :: * -> *) a. Applicative f => a -> f a
pure ()

                      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
d -> VName
-> [TExp Int64]
-> SubExp
-> [TExp Int64]
-> ImpM GPUMem KernelEnv KernelOp ()
forall rep r op.
VName -> [TExp Int64] -> SubExp -> [TExp Int64] -> ImpM rep r op ()
copyDWIMFix VName
d [TV Int64 -> TExp Int64
forall {k} (t :: k). TV t -> TExp t
tvExp TV Int64
phys_block_id, TExp Int64
ltid, TExp Int64
i] SubExp
src []) Maybe VName
glob_dest

                      if [PatElem LetDecMem] -> Bool
forall a. [a] -> Bool
forall (t :: * -> *) a. Foldable t => t a -> Bool
null [PatElem LetDecMem]
pats
                        then () -> ImpM GPUMem KernelEnv KernelOp ()
forall a. a -> ImpM GPUMem KernelEnv KernelOp a
forall (f :: * -> *) a. Applicative f => a -> f a
pure ()
                        else [PatElem LetDecMem]
-> (PatElem LetDecMem -> ImpM GPUMem KernelEnv KernelOp ())
-> ImpM GPUMem KernelEnv KernelOp ()
forall (t :: * -> *) (m :: * -> *) a b.
(Foldable t, Monad m) =>
t a -> (a -> m b) -> m ()
forM_ [PatElem LetDecMem]
pats ((PatElem LetDecMem -> ImpM GPUMem KernelEnv KernelOp ())
 -> ImpM GPUMem KernelEnv KernelOp ())
-> (PatElem LetDecMem -> ImpM GPUMem KernelEnv KernelOp ())
-> ImpM GPUMem KernelEnv KernelOp ()
forall a b. (a -> b) -> a -> b
$ \PatElem LetDecMem
pe -> VName
-> [TExp Int64]
-> SubExp
-> [TExp Int64]
-> ImpM GPUMem KernelEnv KernelOp ()
forall rep r op.
VName -> [TExp Int64] -> SubExp -> [TExp Int64] -> ImpM rep r op ()
copyDWIMFix (PatElem LetDecMem -> VName
forall dec. PatElem dec -> VName
patElemName PatElem LetDecMem
pe) ((VName -> TExp Int64) -> [VName] -> [TExp Int64]
forall a b. (a -> b) -> [a] -> [b]
map VName -> TExp Int64
forall a. a -> TPrimExp Int64 a
le64 [VName]
gtids) SubExp
src []

                  -- Write to-scan results to private memory.
                  [(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]
scan_private_chunks ([SubExp] -> [(VName, SubExp)]) -> [SubExp] -> [(VName, SubExp)]
forall a b. (a -> b) -> a -> b
$ (KernelResult -> SubExp) -> [KernelResult] -> [SubExp]
forall a b. (a -> b) -> [a] -> [b]
map KernelResult -> SubExp
kernelResultSubExp [KernelResult]
all_scan_res) (((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
dest, SubExp
src) ->
                    VName
-> [TExp Int64]
-> SubExp
-> [TExp Int64]
-> ImpM GPUMem KernelEnv KernelOp ()
forall rep r op.
VName -> [TExp Int64] -> SubExp -> [TExp Int64] -> ImpM rep r op ()
copyDWIMFix VName
dest [TExp Int64
i] SubExp
src []

              out_of_bounds =
                [(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]
scan_private_chunks [SubExp]
scanop_nes) (((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
dest, SubExp
ne) ->
                  VName
-> [TExp Int64]
-> SubExp
-> [TExp Int64]
-> ImpM GPUMem KernelEnv KernelOp ()
forall rep r op.
VName -> [TExp Int64] -> SubExp -> [TExp Int64] -> ImpM rep r op ()
copyDWIMFix VName
dest [TExp Int64
i] SubExp
ne []

          sIf (virt_tid .<. n) in_bounds out_of_bounds

      sOp $ Imp.ErrorSync Imp.FenceLocal
      sComment "Transpose scan inputs" $ do
        forM_ (zip transposedArrays scan_private_chunks) $ \(VName
trans, VName
priv) -> do
          Name
-> TExp Int64
-> (TExp Int64 -> ImpM GPUMem KernelEnv KernelOp ())
-> ImpM GPUMem KernelEnv KernelOp ()
forall {k} (t :: k) rep r op.
Name -> TExp t -> (TExp t -> ImpM rep r op ()) -> ImpM rep r op ()
sFor Name
"i" TExp Int64
chunk ((TExp Int64 -> ImpM GPUMem KernelEnv KernelOp ())
 -> ImpM GPUMem KernelEnv KernelOp ())
-> (TExp Int64 -> ImpM GPUMem KernelEnv KernelOp ())
-> ImpM GPUMem KernelEnv KernelOp ()
forall a b. (a -> b) -> a -> b
$ \TExp Int64
i -> do
            sharedIdx <- Name -> TExp Int64 -> ImpM GPUMem KernelEnv KernelOp (TExp Int64)
forall {k} (t :: k) rep r op.
Name -> TExp t -> ImpM rep r op (TExp t)
dPrimVE Name
"sharedIdx" (TExp Int64 -> ImpM GPUMem KernelEnv KernelOp (TExp Int64))
-> TExp Int64 -> ImpM GPUMem KernelEnv KernelOp (TExp Int64)
forall a b. (a -> b) -> a -> b
$ TExp Int64
ltid TExp Int64 -> TExp Int64 -> TExp Int64
forall a. Num a => a -> a -> a
+ TExp Int64
i TExp Int64 -> TExp Int64 -> TExp Int64
forall a. Num a => a -> a -> a
* TExp Int64
tblock_size_e
            copyDWIMFix trans [sharedIdx] (Var priv) [i]
          KernelOp -> ImpM GPUMem KernelEnv KernelOp ()
forall op rep r. op -> ImpM rep r op ()
sOp KernelOp
local_barrier
          Name
-> TExp Int64
-> (TExp Int64 -> ImpM GPUMem KernelEnv KernelOp ())
-> ImpM GPUMem KernelEnv KernelOp ()
forall {k} (t :: k) rep r op.
Name -> TExp t -> (TExp t -> ImpM rep r op ()) -> ImpM rep r op ()
sFor Name
"i" TExp Int64
chunk ((TExp Int64 -> ImpM GPUMem KernelEnv KernelOp ())
 -> ImpM GPUMem KernelEnv KernelOp ())
-> (TExp Int64 -> ImpM GPUMem KernelEnv KernelOp ())
-> ImpM GPUMem KernelEnv KernelOp ()
forall a b. (a -> b) -> a -> b
$ \TExp Int64
i -> do
            sharedIdx <- Name -> TExp Int64 -> ImpM GPUMem KernelEnv KernelOp (TV Int64)
forall {k} (t :: k) rep r op.
Name -> TExp t -> ImpM rep r op (TV t)
dPrimV Name
"sharedIdx" (TExp Int64 -> ImpM GPUMem KernelEnv KernelOp (TV Int64))
-> TExp Int64 -> ImpM GPUMem KernelEnv KernelOp (TV Int64)
forall a b. (a -> b) -> a -> b
$ TExp Int64
ltid TExp Int64 -> TExp Int64 -> TExp Int64
forall a. Num a => a -> a -> a
* TExp Int64
chunk TExp Int64 -> TExp Int64 -> TExp Int64
forall a. Num a => a -> a -> a
+ TExp Int64
i
            copyDWIMFix priv [sExt64 i] (Var trans) [sExt64 $ tvExp sharedIdx]
          KernelOp -> ImpM GPUMem KernelEnv KernelOp ()
forall op rep r. op -> ImpM rep r op ()
sOp KernelOp
local_barrier

      sComment "Per thread scan" $ do
        -- We don't need to touch the first element, so only m-1
        -- iterations here.
        sFor "i" (chunk - 1) $ \TExp Int64
i -> do
          let xs :: [VName]
xs = (LParam GPUMem -> VName) -> [LParam GPUMem] -> [VName]
forall a b. (a -> b) -> [a] -> [b]
map LParam GPUMem -> VName
Param LetDecMem -> VName
forall dec. Param dec -> VName
paramName ([LParam GPUMem] -> [VName]) -> [LParam GPUMem] -> [VName]
forall a b. (a -> b) -> a -> b
$ SegBinOp GPUMem -> [LParam GPUMem]
xParams SegBinOp GPUMem
scan_op
              ys :: [VName]
ys = (LParam GPUMem -> VName) -> [LParam GPUMem] -> [VName]
forall a b. (a -> b) -> [a] -> [b]
map LParam GPUMem -> VName
Param LetDecMem -> VName
forall dec. Param dec -> VName
paramName ([LParam GPUMem] -> [VName]) -> [LParam GPUMem] -> [VName]
forall a b. (a -> b) -> a -> b
$ SegBinOp GPUMem -> [LParam GPUMem]
yParams SegBinOp GPUMem
scan_op
          -- determine if start of segment
          new_sgm <-
            if Bool
segmented
              then do
                gidx <- Name -> TExp Int32 -> ImpM GPUMem KernelEnv KernelOp (TExp Int32)
forall {k} (t :: k) rep r op.
Name -> TExp t -> ImpM rep r op (TExp t)
dPrimVE Name
"gidx" (TExp Int32 -> ImpM GPUMem KernelEnv KernelOp (TExp Int32))
-> TExp Int32 -> ImpM GPUMem KernelEnv KernelOp (TExp Int32)
forall a b. (a -> b) -> a -> b
$ (TExp Int32
ltid32 TExp Int32 -> TExp Int32 -> TExp Int32
forall a. Num a => a -> a -> a
* TExp Int32
chunk32) TExp Int32 -> TExp Int32 -> TExp Int32
forall a. Num a => a -> a -> a
+ TExp Int32
1
                dPrimVE "new_sgm" $ (gidx + sExt32 i - boundary) `mod` segsize_compact .==. 0
              else TPrimExp Bool VName
-> ImpM GPUMem KernelEnv KernelOp (TPrimExp Bool VName)
forall a. a -> ImpM GPUMem KernelEnv KernelOp a
forall (f :: * -> *) a. Applicative f => a -> f a
pure TPrimExp Bool VName
forall v. TPrimExp Bool v
false
          -- skip scan of first element in segment
          sUnless new_sgm $ do
            forM_ (zip4 scan_private_chunks xs ys scan_tys) $ \(VName
src, VName
x, VName
y, PrimType
ty) -> do
              VName -> PrimType -> ImpM GPUMem KernelEnv KernelOp ()
forall rep r op. VName -> PrimType -> ImpM rep r op ()
dPrim_ VName
x PrimType
ty
              VName -> PrimType -> ImpM GPUMem KernelEnv KernelOp ()
forall rep r op. VName -> PrimType -> ImpM rep r op ()
dPrim_ VName
y PrimType
ty
              VName
-> [TExp Int64]
-> SubExp
-> [TExp Int64]
-> ImpM GPUMem KernelEnv KernelOp ()
forall rep r op.
VName -> [TExp Int64] -> SubExp -> [TExp Int64] -> ImpM rep r op ()
copyDWIMFix VName
x [] (VName -> SubExp
Var VName
src) [TExp Int64
i]
              VName
-> [TExp Int64]
-> SubExp
-> [TExp Int64]
-> ImpM GPUMem KernelEnv KernelOp ()
forall rep r op.
VName -> [TExp Int64] -> SubExp -> [TExp Int64] -> ImpM rep r op ()
copyDWIMFix VName
y [] (VName -> SubExp
Var VName
src) [TExp Int64
i TExp Int64 -> TExp Int64 -> TExp Int64
forall a. Num a => a -> a -> a
+ TExp Int64
1]

            compileStms mempty (bodyStms $ lambdaBody $ segBinOpLambda scan_op) $
              forM_ (zip scan_private_chunks $ map resSubExp $ bodyResult $ lambdaBody $ segBinOpLambda scan_op) $ \(VName
dest, SubExp
res) ->
                VName
-> [TExp Int64]
-> SubExp
-> [TExp Int64]
-> ImpM GPUMem KernelEnv KernelOp ()
forall rep r op.
VName -> [TExp Int64] -> SubExp -> [TExp Int64] -> ImpM rep r op ()
copyDWIMFix VName
dest [TExp Int64
i TExp Int64 -> TExp Int64 -> TExp Int64
forall a. Num a => a -> a -> a
+ TExp Int64
1] SubExp
res []

      sComment "Publish results in shared memory" $ do
        forM_ (zip prefixArrays scan_private_chunks) $ \(VName
dest, VName
src) ->
          VName
-> [TExp Int64]
-> SubExp
-> [TExp Int64]
-> ImpM GPUMem KernelEnv KernelOp ()
forall rep r op.
VName -> [TExp Int64] -> SubExp -> [TExp Int64] -> ImpM rep r op ()
copyDWIMFix VName
dest [TExp Int64
ltid] (VName -> SubExp
Var VName
src) [TExp Int64
chunk TExp Int64 -> TExp Int64 -> TExp Int64
forall a. Num a => a -> a -> a
- TExp Int64
1]
        sOp local_barrier

      let crossesSegment = do
            Bool -> Maybe ()
forall (f :: * -> *). Alternative f => Bool -> f ()
guard Bool
segmented
            (TExp Int32 -> TExp Int32 -> TPrimExp Bool VName)
-> Maybe (TExp Int32 -> TExp Int32 -> TPrimExp Bool VName)
forall a. a -> Maybe a
Just ((TExp Int32 -> TExp Int32 -> TPrimExp Bool VName)
 -> Maybe (TExp Int32 -> TExp Int32 -> TPrimExp Bool VName))
-> (TExp Int32 -> TExp Int32 -> TPrimExp Bool VName)
-> Maybe (TExp Int32 -> TExp Int32 -> TPrimExp Bool VName)
forall a b. (a -> b) -> a -> b
$ \TExp Int32
from TExp Int32
to ->
              let from' :: TExp Int32
from' = (TExp Int32
from TExp Int32 -> TExp Int32 -> TExp Int32
forall a. Num a => a -> a -> a
+ TExp Int32
1) TExp Int32 -> TExp Int32 -> TExp Int32
forall a. Num a => a -> a -> a
* TExp Int32
chunk32 TExp Int32 -> TExp Int32 -> TExp Int32
forall a. Num a => a -> a -> a
- TExp Int32
1
                  to' :: TExp Int32
to' = (TExp Int32
to TExp Int32 -> TExp Int32 -> TExp Int32
forall a. Num a => a -> a -> a
+ TExp Int32
1) TExp Int32 -> TExp Int32 -> TExp Int32
forall a. Num a => a -> a -> a
* TExp Int32
chunk32 TExp Int32 -> TExp Int32 -> TExp Int32
forall a. Num a => a -> a -> a
- TExp Int32
1
               in (TExp Int32
to' TExp Int32 -> TExp Int32 -> TExp Int32
forall a. Num a => a -> a -> a
- TExp Int32
from') TExp Int32 -> TExp Int32 -> TPrimExp Bool VName
forall {k} v (t :: k).
Eq v =>
TPrimExp t v -> TPrimExp t v -> TPrimExp Bool v
.>. (TExp Int32
to' TExp Int32 -> TExp Int32 -> TExp Int32
forall a. Num a => a -> a -> a
+ TExp Int32
segsize_compact TExp Int32 -> TExp Int32 -> TExp Int32
forall a. Num a => a -> a -> a
- TExp Int32
boundary) TExp Int32 -> TExp Int32 -> TExp Int32
forall e. IntegralExp e => e -> e -> e
`mod` TExp Int32
segsize_compact

      scan_op1 <- renameLambda $ segBinOpLambda scan_op

      accs <- mapM (dPrimSV "acc") scan_tys
      sComment "Scan results (with warp scan)" $ do
        blockScan
          crossesSegment
          tblock_size_e
          num_virt_threads
          scan_op1
          prefixArrays

        sOp $ Imp.ErrorSync Imp.FenceLocal

        let firstThread TV (ZonkAny 1)
acc VName
prefixes =
              VName
-> [TExp Int64]
-> SubExp
-> [TExp Int64]
-> ImpM GPUMem KernelEnv KernelOp ()
forall rep r op.
VName -> [TExp Int64] -> SubExp -> [TExp Int64] -> ImpM rep r op ()
copyDWIMFix (TV (ZonkAny 1) -> VName
forall {k} (t :: k). TV t -> VName
tvVar TV (ZonkAny 1)
acc) [] (VName -> SubExp
Var VName
prefixes) [TExp Int64 -> TExp Int64
forall {k} (t :: k) v. IntExp t => TPrimExp t v -> TPrimExp Int64 v
sExt64 TExp Int64
tblock_size_e TExp Int64 -> TExp Int64 -> TExp Int64
forall a. Num a => a -> a -> a
- TExp Int64
1]
            notFirstThread TV (ZonkAny 1)
acc VName
prefixes =
              VName
-> [TExp Int64]
-> SubExp
-> [TExp Int64]
-> ImpM GPUMem KernelEnv KernelOp ()
forall rep r op.
VName -> [TExp Int64] -> SubExp -> [TExp Int64] -> ImpM rep r op ()
copyDWIMFix (TV (ZonkAny 1) -> VName
forall {k} (t :: k). TV t -> VName
tvVar TV (ZonkAny 1)
acc) [] (VName -> SubExp
Var VName
prefixes) [TExp Int64
ltid TExp Int64 -> TExp Int64 -> TExp Int64
forall a. Num a => a -> a -> a
- TExp Int64
1]
        sIf
          (ltid32 .==. 0)
          (zipWithM_ firstThread accs prefixArrays)
          (zipWithM_ notFirstThread accs prefixArrays)

        sOp local_barrier

      prefixes <- forM (zip scanop_nes scan_tys) $ \(SubExp
ne, PrimType
ty) ->
        Name
-> TExp (ZonkAny 3)
-> ImpM GPUMem KernelEnv KernelOp (TV (ZonkAny 3))
forall {k} (t :: k) rep r op.
Name -> TExp t -> ImpM rep r op (TV t)
dPrimV Name
"prefix" (TExp (ZonkAny 3)
 -> ImpM GPUMem KernelEnv KernelOp (TV (ZonkAny 3)))
-> TExp (ZonkAny 3)
-> ImpM GPUMem KernelEnv KernelOp (TV (ZonkAny 3))
forall a b. (a -> b) -> a -> b
$ PrimExp VName -> TExp (ZonkAny 3)
forall {k} (t :: k) v. PrimExp v -> TPrimExp t v
TPrimExp (PrimExp VName -> TExp (ZonkAny 3))
-> PrimExp VName -> TExp (ZonkAny 3)
forall a b. (a -> b) -> a -> b
$ PrimType -> SubExp -> PrimExp VName
forall a. ToExp a => PrimType -> a -> PrimExp VName
toExp' PrimType
ty SubExp
ne
      blockNewSgm <- dPrimVE "block_new_sgm" $ sgm_idx .==. 0
      sComment "Perform lookback" $ do
        sWhen (blockNewSgm .&&. ltid32 .==. 0) $ do
          everythingVolatile $
            forM_ (zip accs incprefixArrays) $ \(TV (ZonkAny 1)
acc, VName
incprefixArray) ->
              VName
-> [TExp Int64]
-> SubExp
-> [TExp Int64]
-> ImpM GPUMem KernelEnv KernelOp ()
forall rep r op.
VName -> [TExp Int64] -> SubExp -> [TExp Int64] -> ImpM rep r op ()
copyDWIMFix VName
incprefixArray [TV Int64 -> TExp Int64
forall {k} (t :: k). TV t -> TExp t
tvExp TV Int64
dyn_id] (TV (ZonkAny 1) -> SubExp
forall {k} (t :: k). TV t -> SubExp
tvSize TV (ZonkAny 1)
acc) []
          sOp global_fence
          everythingVolatile $
            copyDWIMFix statusFlags [tvExp dyn_id] (intConst Int8 statusP) []
          forM_ (zip scanop_nes accs) $ \(SubExp
ne, TV (ZonkAny 1)
acc) ->
            VName
-> [TExp Int64]
-> SubExp
-> [TExp Int64]
-> ImpM GPUMem KernelEnv KernelOp ()
forall rep r op.
VName -> [TExp Int64] -> SubExp -> [TExp Int64] -> ImpM rep r op ()
copyDWIMFix (TV (ZonkAny 1) -> VName
forall {k} (t :: k). TV t -> VName
tvVar TV (ZonkAny 1)
acc) [] SubExp
ne []
        -- end sWhen

        let warp_size = KernelConstants -> TExp Int32
kernelWaveSize KernelConstants
constants
        sWhen (bNot blockNewSgm .&&. ltid32 .<. warp_size) $ do
          sWhen (ltid32 .==. 0) $ do
            sIf
              (not_segmented_e .||. boundary .==. sExt32 (tblock_size_e * chunk))
              ( do
                  everythingVolatile $
                    forM_ (zip aggregateArrays accs) $ \(VName
aggregateArray, TV (ZonkAny 1)
acc) ->
                      VName
-> [TExp Int64]
-> SubExp
-> [TExp Int64]
-> ImpM GPUMem KernelEnv KernelOp ()
forall rep r op.
VName -> [TExp Int64] -> SubExp -> [TExp Int64] -> ImpM rep r op ()
copyDWIMFix VName
aggregateArray [TV Int64 -> TExp Int64
forall {k} (t :: k). TV t -> TExp t
tvExp TV Int64
dyn_id] (TV (ZonkAny 1) -> SubExp
forall {k} (t :: k). TV t -> SubExp
tvSize TV (ZonkAny 1)
acc) []
                  sOp global_fence
                  everythingVolatile $
                    copyDWIMFix statusFlags [tvExp dyn_id] (intConst Int8 statusA) []
              )
              ( do
                  everythingVolatile $
                    forM_ (zip incprefixArrays accs) $ \(VName
incprefixArray, TV (ZonkAny 1)
acc) ->
                      VName
-> [TExp Int64]
-> SubExp
-> [TExp Int64]
-> ImpM GPUMem KernelEnv KernelOp ()
forall rep r op.
VName -> [TExp Int64] -> SubExp -> [TExp Int64] -> ImpM rep r op ()
copyDWIMFix VName
incprefixArray [TV Int64 -> TExp Int64
forall {k} (t :: k). TV t -> TExp t
tvExp TV Int64
dyn_id] (TV (ZonkAny 1) -> SubExp
forall {k} (t :: k). TV t -> SubExp
tvSize TV (ZonkAny 1)
acc) []
                  sOp global_fence
                  everythingVolatile $
                    copyDWIMFix statusFlags [tvExp dyn_id] (intConst Int8 statusP) []
              )
            everythingVolatile $
              copyDWIMFix warpscan [0] (Var statusFlags) [tvExp dyn_id - 1]
          -- sWhen
          sOp local_fence

          status :: TV Int8 <- dPrim "status"
          copyDWIMFix (tvVar status) [] (Var warpscan) [0]

          sIf
            (tvExp status .==. statusP)
            ( sWhen (ltid32 .==. 0) $
                everythingVolatile $
                  forM_ (zip prefixes incprefixArrays) $ \(TV (ZonkAny 3)
prefix, VName
incprefixArray) ->
                    VName
-> [TExp Int64]
-> SubExp
-> [TExp Int64]
-> ImpM GPUMem KernelEnv KernelOp ()
forall rep r op.
VName -> [TExp Int64] -> SubExp -> [TExp Int64] -> ImpM rep r op ()
copyDWIMFix (TV (ZonkAny 3) -> VName
forall {k} (t :: k). TV t -> VName
tvVar TV (ZonkAny 3)
prefix) [] (VName -> SubExp
Var VName
incprefixArray) [TV Int64 -> TExp Int64
forall {k} (t :: k). TV t -> TExp t
tvExp TV Int64
dyn_id TExp Int64 -> TExp Int64 -> TExp Int64
forall a. Num a => a -> a -> a
- TExp Int64
1]
            )
            ( do
                readOffset <-
                  dPrimV "readOffset" $
                    sExt32 $
                      tvExp dyn_id - sExt64 (kernelWaveSize constants)
                let loopStop = TExp Int32
warp_size TExp Int32 -> TExp Int32 -> TExp Int32
forall a. Num a => a -> a -> a
* (-TExp Int32
1)
                    sameSegment TV Int32
readIdx
                      | Bool
segmented =
                          let startIdx :: TExp Int64
startIdx = TExp Int32 -> TExp Int64
forall {k} (t :: k) v. IntExp t => TPrimExp t v -> TPrimExp Int64 v
sExt64 (TV Int32 -> TExp Int32
forall {k} (t :: k). TV t -> TExp t
tvExp TV Int32
readIdx TExp Int32 -> TExp Int32 -> TExp Int32
forall a. Num a => a -> a -> a
+ TExp Int32
1) TExp Int64 -> TExp Int64 -> TExp Int64
forall a. Num a => a -> a -> a
* TExp Int64
tblock_size_e TExp Int64 -> TExp Int64 -> TExp Int64
forall a. Num a => a -> a -> a
* TExp Int64
chunk TExp Int64 -> TExp Int64 -> TExp Int64
forall a. Num a => a -> a -> a
- TExp Int64
1
                           in TExp Int64
block_offset TExp Int64 -> TExp Int64 -> TExp Int64
forall a. Num a => a -> a -> a
- TExp Int64
startIdx TExp Int64 -> TExp Int64 -> TPrimExp Bool VName
forall {k} v (t :: k).
Eq v =>
TPrimExp t v -> TPrimExp t v -> TPrimExp Bool v
.<=. TExp Int64
sgm_idx
                      | Bool
otherwise = TPrimExp Bool VName
forall v. TPrimExp Bool v
true
                sWhile (tvExp readOffset .>. loopStop) $ do
                  readI <- dPrimV "read_i" $ tvExp readOffset + ltid32
                  aggrs <- forM (zip scanop_nes scan_tys) $ \(SubExp
ne, PrimType
ty) ->
                    Name
-> TExp (ZonkAny 5)
-> ImpM GPUMem KernelEnv KernelOp (TV (ZonkAny 5))
forall {k} (t :: k) rep r op.
Name -> TExp t -> ImpM rep r op (TV t)
dPrimV Name
"aggr" (TExp (ZonkAny 5)
 -> ImpM GPUMem KernelEnv KernelOp (TV (ZonkAny 5)))
-> TExp (ZonkAny 5)
-> ImpM GPUMem KernelEnv KernelOp (TV (ZonkAny 5))
forall a b. (a -> b) -> a -> b
$ PrimExp VName -> TExp (ZonkAny 5)
forall {k} (t :: k) v. PrimExp v -> TPrimExp t v
TPrimExp (PrimExp VName -> TExp (ZonkAny 5))
-> PrimExp VName -> TExp (ZonkAny 5)
forall a b. (a -> b) -> a -> b
$ PrimType -> SubExp -> PrimExp VName
forall a. ToExp a => PrimType -> a -> PrimExp VName
toExp' PrimType
ty SubExp
ne
                  flag <- dPrimV "flag" (statusX :: Imp.TExp Int8)
                  everythingVolatile . sWhen (tvExp readI .>=. 0) $ do
                    sIf
                      (sameSegment readI)
                      ( do
                          copyDWIMFix (tvVar flag) [] (Var statusFlags) [sExt64 $ tvExp readI]
                          sIf
                            (tvExp flag .==. statusP)
                            ( forM_ (zip incprefixArrays aggrs) $ \(VName
incprefix, TV (ZonkAny 5)
aggr) ->
                                VName
-> [TExp Int64]
-> SubExp
-> [TExp Int64]
-> ImpM GPUMem KernelEnv KernelOp ()
forall rep r op.
VName -> [TExp Int64] -> SubExp -> [TExp Int64] -> ImpM rep r op ()
copyDWIMFix (TV (ZonkAny 5) -> VName
forall {k} (t :: k). TV t -> VName
tvVar TV (ZonkAny 5)
aggr) [] (VName -> SubExp
Var VName
incprefix) [TExp Int32 -> TExp Int64
forall {k} (t :: k) v. IntExp t => TPrimExp t v -> TPrimExp Int64 v
sExt64 (TExp Int32 -> TExp Int64) -> TExp Int32 -> TExp Int64
forall a b. (a -> b) -> a -> b
$ TV Int32 -> TExp Int32
forall {k} (t :: k). TV t -> TExp t
tvExp TV Int32
readI]
                            )
                            ( sWhen (tvExp flag .==. statusA) $ do
                                forM_ (zip aggrs aggregateArrays) $ \(TV (ZonkAny 5)
aggr, VName
aggregate) ->
                                  VName
-> [TExp Int64]
-> SubExp
-> [TExp Int64]
-> ImpM GPUMem KernelEnv KernelOp ()
forall rep r op.
VName -> [TExp Int64] -> SubExp -> [TExp Int64] -> ImpM rep r op ()
copyDWIMFix (TV (ZonkAny 5) -> VName
forall {k} (t :: k). TV t -> VName
tvVar TV (ZonkAny 5)
aggr) [] (VName -> SubExp
Var VName
aggregate) [TExp Int32 -> TExp Int64
forall {k} (t :: k) v. IntExp t => TPrimExp t v -> TPrimExp Int64 v
sExt64 (TExp Int32 -> TExp Int64) -> TExp Int32 -> TExp Int64
forall a b. (a -> b) -> a -> b
$ TV Int32 -> TExp Int32
forall {k} (t :: k). TV t -> TExp t
tvExp TV Int32
readI]
                            )
                      )
                      (copyDWIMFix (tvVar flag) [] (intConst Int8 statusP) [])
                  -- end sIf
                  -- end sWhen

                  forM_ (zip exchanges aggrs) $ \(VName
exchange, TV (ZonkAny 5)
aggr) ->
                    VName
-> [TExp Int64]
-> SubExp
-> [TExp Int64]
-> ImpM GPUMem KernelEnv KernelOp ()
forall rep r op.
VName -> [TExp Int64] -> SubExp -> [TExp Int64] -> ImpM rep r op ()
copyDWIMFix VName
exchange [TExp Int64
ltid] (TV (ZonkAny 5) -> SubExp
forall {k} (t :: k). TV t -> SubExp
tvSize TV (ZonkAny 5)
aggr) []
                  copyDWIMFix warpscan [ltid] (tvSize flag) []

                  -- execute warp-parallel reduction but only if the last read flag in not STATUS_P
                  copyDWIMFix (tvVar flag) [] (Var warpscan) [sExt64 warp_size - 1]
                  sWhen (tvExp flag .<. statusP) $ do
                    lam' <- renameLambda scan_op1
                    inBlockScanLookback
                      constants
                      num_virt_threads
                      warpscan
                      exchanges
                      lam'

                  -- all threads of the warp read the result of reduction
                  copyDWIMFix (tvVar flag) [] (Var warpscan) [sExt64 warp_size - 1]
                  forM_ (zip aggrs exchanges) $ \(TV (ZonkAny 5)
aggr, VName
exchange) ->
                    VName
-> [TExp Int64]
-> SubExp
-> [TExp Int64]
-> ImpM GPUMem KernelEnv KernelOp ()
forall rep r op.
VName -> [TExp Int64] -> SubExp -> [TExp Int64] -> ImpM rep r op ()
copyDWIMFix (TV (ZonkAny 5) -> VName
forall {k} (t :: k). TV t -> VName
tvVar TV (ZonkAny 5)
aggr) [] (VName -> SubExp
Var VName
exchange) [TExp Int32 -> TExp Int64
forall {k} (t :: k) v. IntExp t => TPrimExp t v -> TPrimExp Int64 v
sExt64 TExp Int32
warp_size TExp Int64 -> TExp Int64 -> TExp Int64
forall a. Num a => a -> a -> a
- TExp Int64
1]
                  -- update read offset
                  sIf
                    (tvExp flag .==. statusP)
                    (readOffset <-- loopStop)
                    ( sWhen (tvExp flag .==. statusA) $ do
                        readOffset <-- tvExp readOffset - zExt32 warp_size
                    )

                  -- update prefix if flag different than STATUS_X:
                  sWhen (tvExp flag .>. statusX) $ do
                    lam <- renameLambda scan_op1
                    let (xs, ys) = splitAt (length scan_tys) $ map paramName $ lambdaParams lam
                    forM_ (zip xs aggrs) $ \(VName
x, TV (ZonkAny 5)
aggr) -> VName -> TExp (ZonkAny 5) -> ImpM GPUMem KernelEnv KernelOp ()
forall {k} (t :: k) rep r op. VName -> TExp t -> ImpM rep r op ()
dPrimV_ VName
x (TV (ZonkAny 5) -> TExp (ZonkAny 5)
forall {k} (t :: k). TV t -> TExp t
tvExp TV (ZonkAny 5)
aggr)
                    forM_ (zip ys prefixes) $ \(VName
y, TV (ZonkAny 3)
prefix) -> VName -> TExp (ZonkAny 3) -> ImpM GPUMem KernelEnv KernelOp ()
forall {k} (t :: k) rep r op. VName -> TExp t -> ImpM rep r op ()
dPrimV_ VName
y (TV (ZonkAny 3) -> TExp (ZonkAny 3)
forall {k} (t :: k). TV t -> TExp t
tvExp TV (ZonkAny 3)
prefix)
                    compileStms mempty (bodyStms $ lambdaBody lam) $
                      forM_ (zip3 prefixes scan_tys $ map resSubExp $ bodyResult $ lambdaBody lam) $
                        \(TV (ZonkAny 3)
prefix, PrimType
ty, SubExp
res) -> TV (ZonkAny 3)
prefix TV (ZonkAny 3)
-> TExp (ZonkAny 3) -> ImpM GPUMem KernelEnv KernelOp ()
forall {k} (t :: k) rep r op. TV t -> TExp t -> ImpM rep r op ()
<-- PrimExp VName -> TExp (ZonkAny 3)
forall {k} (t :: k) v. PrimExp v -> TPrimExp t v
TPrimExp (PrimType -> SubExp -> PrimExp VName
forall a. ToExp a => PrimType -> a -> PrimExp VName
toExp' PrimType
ty SubExp
res)
                  sOp local_fence
            )

          -- end sWhile
          -- end sIf
          sWhen (ltid32 .==. 0) $ do
            scan_op2 <- renameLambda scan_op1
            let xs = (Param LetDecMem -> VName) -> [Param LetDecMem] -> [VName]
forall a b. (a -> b) -> [a] -> [b]
map Param LetDecMem -> VName
forall dec. Param dec -> VName
paramName ([Param LetDecMem] -> [VName]) -> [Param LetDecMem] -> [VName]
forall a b. (a -> b) -> a -> b
$ Int -> [Param LetDecMem] -> [Param LetDecMem]
forall a. Int -> [a] -> [a]
take ([PrimType] -> Int
forall a. [a] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length [PrimType]
scan_tys) ([Param LetDecMem] -> [Param LetDecMem])
-> [Param LetDecMem] -> [Param LetDecMem]
forall a b. (a -> b) -> a -> b
$ GLambda GPUMem Type -> [LParam GPUMem]
forall rep t. GLambda rep t -> [LParam rep]
lambdaParams GLambda GPUMem Type
scan_op2
                ys = (Param LetDecMem -> VName) -> [Param LetDecMem] -> [VName]
forall a b. (a -> b) -> [a] -> [b]
map Param LetDecMem -> VName
forall dec. Param dec -> VName
paramName ([Param LetDecMem] -> [VName]) -> [Param LetDecMem] -> [VName]
forall a b. (a -> b) -> a -> b
$ Int -> [Param LetDecMem] -> [Param LetDecMem]
forall a. Int -> [a] -> [a]
drop ([PrimType] -> Int
forall a. [a] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length [PrimType]
scan_tys) ([Param LetDecMem] -> [Param LetDecMem])
-> [Param LetDecMem] -> [Param LetDecMem]
forall a b. (a -> b) -> a -> b
$ GLambda GPUMem Type -> [LParam GPUMem]
forall rep t. GLambda rep t -> [LParam rep]
lambdaParams GLambda GPUMem Type
scan_op2
            sWhen (boundary .==. sExt32 (tblock_size_e * chunk)) $ do
              forM_ (zip xs prefixes) $ \(VName
x, TV (ZonkAny 3)
prefix) -> VName -> TExp (ZonkAny 3) -> ImpM GPUMem KernelEnv KernelOp ()
forall {k} (t :: k) rep r op. VName -> TExp t -> ImpM rep r op ()
dPrimV_ VName
x (TExp (ZonkAny 3) -> ImpM GPUMem KernelEnv KernelOp ())
-> TExp (ZonkAny 3) -> ImpM GPUMem KernelEnv KernelOp ()
forall a b. (a -> b) -> a -> b
$ TV (ZonkAny 3) -> TExp (ZonkAny 3)
forall {k} (t :: k). TV t -> TExp t
tvExp TV (ZonkAny 3)
prefix
              forM_ (zip ys accs) $ \(VName
y, TV (ZonkAny 1)
acc) -> VName -> TExp (ZonkAny 1) -> ImpM GPUMem KernelEnv KernelOp ()
forall {k} (t :: k) rep r op. VName -> TExp t -> ImpM rep r op ()
dPrimV_ VName
y (TExp (ZonkAny 1) -> ImpM GPUMem KernelEnv KernelOp ())
-> TExp (ZonkAny 1) -> ImpM GPUMem KernelEnv KernelOp ()
forall a b. (a -> b) -> a -> b
$ TV (ZonkAny 1) -> TExp (ZonkAny 1)
forall {k} (t :: k). TV t -> TExp t
tvExp TV (ZonkAny 1)
acc
              compileStms mempty (bodyStms $ lambdaBody scan_op2) $
                everythingVolatile $
                  forM_ (zip incprefixArrays $ map resSubExp $ bodyResult $ lambdaBody scan_op2) $
                    \(VName
incprefixArray, SubExp
res) -> VName
-> [TExp Int64]
-> SubExp
-> [TExp Int64]
-> ImpM GPUMem KernelEnv KernelOp ()
forall rep r op.
VName -> [TExp Int64] -> SubExp -> [TExp Int64] -> ImpM rep r op ()
copyDWIMFix VName
incprefixArray [TV Int64 -> TExp Int64
forall {k} (t :: k). TV t -> TExp t
tvExp TV Int64
dyn_id] SubExp
res []
              sOp global_fence
              everythingVolatile $ copyDWIMFix statusFlags [tvExp dyn_id] (intConst Int8 statusP) []
            forM_ (zip exchanges prefixes) $ \(VName
exchange, TV (ZonkAny 3)
prefix) ->
              VName
-> [TExp Int64]
-> SubExp
-> [TExp Int64]
-> ImpM GPUMem KernelEnv KernelOp ()
forall rep r op.
VName -> [TExp Int64] -> SubExp -> [TExp Int64] -> ImpM rep r op ()
copyDWIMFix VName
exchange [TExp Int64
0] (TV (ZonkAny 3) -> SubExp
forall {k} (t :: k). TV t -> SubExp
tvSize TV (ZonkAny 3)
prefix) []
            forM_ (zip3 accs scan_tys scanop_nes) $ \(TV (ZonkAny 1)
acc, PrimType
ty, SubExp
ne) ->
              TV (ZonkAny 1) -> VName
forall {k} (t :: k). TV t -> VName
tvVar TV (ZonkAny 1)
acc VName -> PrimExp VName -> ImpM GPUMem KernelEnv KernelOp ()
forall rep r op. VName -> PrimExp VName -> ImpM rep r op ()
<~~ PrimType -> SubExp -> PrimExp VName
forall a. ToExp a => PrimType -> a -> PrimExp VName
toExp' PrimType
ty SubExp
ne
        -- end sWhen
        -- end sWhen

        sWhen (bNot $ tvExp dyn_id .==. 0) $ do
          sOp local_barrier
          forM_ (zip exchanges prefixes) $ \(VName
exchange, TV (ZonkAny 3)
prefix) ->
            VName
-> [TExp Int64]
-> SubExp
-> [TExp Int64]
-> ImpM GPUMem KernelEnv KernelOp ()
forall rep r op.
VName -> [TExp Int64] -> SubExp -> [TExp Int64] -> ImpM rep r op ()
copyDWIMFix (TV (ZonkAny 3) -> VName
forall {k} (t :: k). TV t -> VName
tvVar TV (ZonkAny 3)
prefix) [] (VName -> SubExp
Var VName
exchange) [TExp Int64
0]
          sOp local_barrier
      -- end sWhen
      -- end sComment

      scan_op3 <- renameLambda scan_op1
      scan_op4 <- renameLambda scan_op1

      sComment "Distribute results" $ do
        let (xs, ys) = splitAt (length scan_tys) $ map paramName $ lambdaParams scan_op3
            (xs', ys') = splitAt (length scan_tys) $ map paramName $ lambdaParams scan_op4

        forM_ (zip7 prefixes accs xs xs' ys ys' scan_tys) $
          \(TV (ZonkAny 3)
prefix, TV (ZonkAny 1)
acc, VName
x, VName
x', VName
y, VName
y', PrimType
ty) -> do
            VName -> PrimType -> ImpM GPUMem KernelEnv KernelOp ()
forall rep r op. VName -> PrimType -> ImpM rep r op ()
dPrim_ VName
x PrimType
ty
            VName -> PrimType -> ImpM GPUMem KernelEnv KernelOp ()
forall rep r op. VName -> PrimType -> ImpM rep r op ()
dPrim_ VName
y PrimType
ty
            VName -> TExp (ZonkAny 3) -> ImpM GPUMem KernelEnv KernelOp ()
forall {k} (t :: k) rep r op. VName -> TExp t -> ImpM rep r op ()
dPrimV_ VName
x' (TExp (ZonkAny 3) -> ImpM GPUMem KernelEnv KernelOp ())
-> TExp (ZonkAny 3) -> ImpM GPUMem KernelEnv KernelOp ()
forall a b. (a -> b) -> a -> b
$ TV (ZonkAny 3) -> TExp (ZonkAny 3)
forall {k} (t :: k). TV t -> TExp t
tvExp TV (ZonkAny 3)
prefix
            VName -> TExp (ZonkAny 1) -> ImpM GPUMem KernelEnv KernelOp ()
forall {k} (t :: k) rep r op. VName -> TExp t -> ImpM rep r op ()
dPrimV_ VName
y' (TExp (ZonkAny 1) -> ImpM GPUMem KernelEnv KernelOp ())
-> TExp (ZonkAny 1) -> ImpM GPUMem KernelEnv KernelOp ()
forall a b. (a -> b) -> a -> b
$ TV (ZonkAny 1) -> TExp (ZonkAny 1)
forall {k} (t :: k). TV t -> TExp t
tvExp TV (ZonkAny 1)
acc

        sIf
          (ltid32 * chunk32 .<. boundary .&&. bNot blockNewSgm)
          ( compileStms mempty (bodyStms $ lambdaBody scan_op4) $
              forM_ (zip3 xs scan_tys $ map resSubExp $ bodyResult $ lambdaBody scan_op4) $
                \(VName
x, PrimType
ty, SubExp
res) -> VName
x VName -> PrimExp VName -> ImpM GPUMem KernelEnv KernelOp ()
forall rep r op. VName -> PrimExp VName -> ImpM rep r op ()
<~~ PrimType -> SubExp -> PrimExp VName
forall a. ToExp a => PrimType -> a -> PrimExp VName
toExp' PrimType
ty SubExp
res
          )
          (forM_ (zip xs accs) $ \(VName
x, TV (ZonkAny 1)
acc) -> VName
-> [TExp Int64]
-> SubExp
-> [TExp Int64]
-> ImpM GPUMem KernelEnv KernelOp ()
forall rep r op.
VName -> [TExp Int64] -> SubExp -> [TExp Int64] -> ImpM rep r op ()
copyDWIMFix VName
x [] (VName -> SubExp
Var (VName -> SubExp) -> VName -> SubExp
forall a b. (a -> b) -> a -> b
$ TV (ZonkAny 1) -> VName
forall {k} (t :: k). TV t -> VName
tvVar TV (ZonkAny 1)
acc) [])
        -- calculate where previous thread stopped, to determine number of
        -- elements left before new segment.
        stop <-
          dPrimVE "stopping_point" $
            segsize_compact - (ltid32 * chunk32 - 1 + segsize_compact - boundary) `rem` segsize_compact
        sFor "i" chunk $ \TExp Int64
i -> do
          TPrimExp Bool VName
-> ImpM GPUMem KernelEnv KernelOp ()
-> ImpM GPUMem KernelEnv KernelOp ()
forall rep r op.
TPrimExp Bool VName -> ImpM rep r op () -> ImpM rep r op ()
sWhen (TExp Int64 -> TExp Int32
forall {k} (t :: k) v. IntExp t => TPrimExp t v -> TPrimExp Int32 v
sExt32 TExp Int64
i TExp Int32 -> TExp Int32 -> TPrimExp Bool VName
forall {k} v (t :: k).
Eq v =>
TPrimExp t v -> TPrimExp t v -> TPrimExp Bool v
.<. TExp Int32
stop TExp Int32 -> TExp Int32 -> TExp Int32
forall a. Num a => a -> a -> a
- TExp Int32
1) (ImpM GPUMem KernelEnv KernelOp ()
 -> ImpM GPUMem KernelEnv KernelOp ())
-> ImpM GPUMem KernelEnv KernelOp ()
-> ImpM GPUMem KernelEnv KernelOp ()
forall a b. (a -> b) -> a -> b
$ do
            [(VName, VName)]
-> ((VName, 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_ ([VName] -> [VName] -> [(VName, VName)]
forall a b. [a] -> [b] -> [(a, b)]
zip [VName]
scan_private_chunks [VName]
ys) (((VName, VName) -> ImpM GPUMem KernelEnv KernelOp ())
 -> ImpM GPUMem KernelEnv KernelOp ())
-> ((VName, VName) -> ImpM GPUMem KernelEnv KernelOp ())
-> ImpM GPUMem KernelEnv KernelOp ()
forall a b. (a -> b) -> a -> b
$ \(VName
src, VName
y) ->
              -- only include prefix for the first segment part per thread
              VName
-> [TExp Int64]
-> SubExp
-> [TExp Int64]
-> ImpM GPUMem KernelEnv KernelOp ()
forall rep r op.
VName -> [TExp Int64] -> SubExp -> [TExp Int64] -> ImpM rep r op ()
copyDWIMFix VName
y [] (VName -> SubExp
Var VName
src) [TExp Int64
i]
            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 (Body GPUMem -> Stms GPUMem
forall rep res. GBody rep res -> Stms rep
bodyStms (Body GPUMem -> Stms GPUMem) -> Body GPUMem -> Stms GPUMem
forall a b. (a -> b) -> a -> b
$ GLambda GPUMem Type -> Body GPUMem
forall rep t. GLambda rep t -> Body rep
lambdaBody GLambda GPUMem Type
scan_op3) (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]
scan_private_chunks ([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
$ Body GPUMem -> [SubExpRes]
forall rep res. GBody rep res -> [res]
bodyResult (Body GPUMem -> [SubExpRes]) -> Body GPUMem -> [SubExpRes]
forall a b. (a -> b) -> a -> b
$ GLambda GPUMem Type -> Body GPUMem
forall rep t. GLambda rep t -> Body rep
lambdaBody GLambda GPUMem Type
scan_op3) (((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
dest, SubExp
res) ->
                  VName
-> [TExp Int64]
-> SubExp
-> [TExp Int64]
-> ImpM GPUMem KernelEnv KernelOp ()
forall rep r op.
VName -> [TExp Int64] -> SubExp -> [TExp Int64] -> ImpM rep r op ()
copyDWIMFix VName
dest [TExp Int64
i] SubExp
res []

      sComment "Transpose scan output and to write it later in coalesced fashion to global memory" $ do
        forM_ (zip transposedArrays scan_private_chunks) $ \(VName
locmem, VName
priv) -> do
          Name
-> TExp Int64
-> (TExp Int64 -> ImpM GPUMem KernelEnv KernelOp ())
-> ImpM GPUMem KernelEnv KernelOp ()
forall {k} (t :: k) rep r op.
Name -> TExp t -> (TExp t -> ImpM rep r op ()) -> ImpM rep r op ()
sFor Name
"i" TExp Int64
chunk ((TExp Int64 -> ImpM GPUMem KernelEnv KernelOp ())
 -> ImpM GPUMem KernelEnv KernelOp ())
-> (TExp Int64 -> ImpM GPUMem KernelEnv KernelOp ())
-> ImpM GPUMem KernelEnv KernelOp ()
forall a b. (a -> b) -> a -> b
$ \TExp Int64
i -> do
            sharedIdx <-
              Name -> TExp Int64 -> ImpM GPUMem KernelEnv KernelOp (TV Int64)
forall {k} (t :: k) rep r op.
Name -> TExp t -> ImpM rep r op (TV t)
dPrimV Name
"sharedIdx" (TExp Int64 -> ImpM GPUMem KernelEnv KernelOp (TV Int64))
-> TExp Int64 -> ImpM GPUMem KernelEnv KernelOp (TV Int64)
forall a b. (a -> b) -> a -> b
$
                TExp Int64 -> TExp Int64
forall {k} (t :: k) v. IntExp t => TPrimExp t v -> TPrimExp Int64 v
sExt64 (TExp Int64
ltid TExp Int64 -> TExp Int64 -> TExp Int64
forall a. Num a => a -> a -> a
* TExp Int64
chunk) TExp Int64 -> TExp Int64 -> TExp Int64
forall a. Num a => a -> a -> a
+ TExp Int64
i
            copyDWIMFix locmem [tvExp sharedIdx] (Var priv) [i]
          KernelOp -> ImpM GPUMem KernelEnv KernelOp ()
forall op rep r. op -> ImpM rep r op ()
sOp KernelOp
local_barrier
          Name
-> TExp Int64
-> (TExp Int64 -> ImpM GPUMem KernelEnv KernelOp ())
-> ImpM GPUMem KernelEnv KernelOp ()
forall {k} (t :: k) rep r op.
Name -> TExp t -> (TExp t -> ImpM rep r op ()) -> ImpM rep r op ()
sFor Name
"i" TExp Int64
chunk ((TExp Int64 -> ImpM GPUMem KernelEnv KernelOp ())
 -> ImpM GPUMem KernelEnv KernelOp ())
-> (TExp Int64 -> ImpM GPUMem KernelEnv KernelOp ())
-> ImpM GPUMem KernelEnv KernelOp ()
forall a b. (a -> b) -> a -> b
$ \TExp Int64
i -> do
            flat_idx <- Name -> TExp Int64 -> ImpM GPUMem KernelEnv KernelOp (TExp Int64)
forall {k} (t :: k) rep r op.
Name -> TExp t -> ImpM rep r op (TExp t)
dPrimVE Name
"flat_idx" (TExp Int64 -> ImpM GPUMem KernelEnv KernelOp (TExp Int64))
-> TExp Int64 -> ImpM GPUMem KernelEnv KernelOp (TExp Int64)
forall a b. (a -> b) -> a -> b
$ TExp Int64
thd_offset TExp Int64 -> TExp Int64 -> TExp Int64
forall a. Num a => a -> a -> a
+ TExp Int64
i TExp Int64 -> TExp Int64 -> TExp Int64
forall a. Num a => a -> a -> a
* TExp Int64
tblock_size_e
            dIndexSpace (zip gtids dims') flat_idx
            sWhen (flat_idx .<. n) $ do
              copyDWIMFix
                priv
                [i]
                (Var locmem)
                [sExt64 $ flat_idx - block_offset]
          KernelOp -> ImpM GPUMem KernelEnv KernelOp ()
forall op rep r. op -> ImpM rep r op ()
sOp KernelOp
local_barrier

      let (scan_pars, map_pars) =
            splitAt (length $ segBinOpNeutral scan_op) $
              map paramName $
                lambdaParams $
                  segPostOpLambda post_op

      sComment "Compute post op and write to global memory." $ do
        sFor "i" chunk $ \TExp Int64
i -> do
          flat_idx <- Name -> TExp Int64 -> ImpM GPUMem KernelEnv KernelOp (TExp Int64)
forall {k} (t :: k) rep r op.
Name -> TExp t -> ImpM rep r op (TExp t)
dPrimVE Name
"flat_idx" (TExp Int64 -> ImpM GPUMem KernelEnv KernelOp (TExp Int64))
-> TExp Int64 -> ImpM GPUMem KernelEnv KernelOp (TExp Int64)
forall a b. (a -> b) -> a -> b
$ TExp Int64
thd_offset TExp Int64 -> TExp Int64 -> TExp Int64
forall a. Num a => a -> a -> a
+ TExp Int64
i TExp Int64 -> TExp Int64 -> TExp Int64
forall a. Num a => a -> a -> a
* TExp Int64
tblock_size_e
          sWhen (flat_idx .<. n) $ do
            dIndexSpace (zip gtids dims') flat_idx

            dScope Nothing $
              scopeOfLParams $
                lambdaParams $
                  segPostOpLambda post_op

            sComment "bind scan results to post lambda params" $ do
              forM_ (zip scan_pars scan_private_chunks) $ \(VName
par, VName
priv) ->
                VName
-> [TExp Int64]
-> SubExp
-> [TExp Int64]
-> ImpM GPUMem KernelEnv KernelOp ()
forall rep r op.
VName -> [TExp Int64] -> SubExp -> [TExp Int64] -> ImpM rep r op ()
copyDWIMFix VName
par [] (VName -> SubExp
Var VName
priv) [TExp Int64
i]

            sComment "bind map results to post lamda params" $
              forM_ (zip4 map_pars map_private_chunks map_global_chunks map_tys') $
                \(VName
par, Maybe VName
priv, Maybe VName
glob, Type
ty) -> do
                  case Maybe VName
priv of
                    Just VName
p
                      | Type -> Bool
shouldUseBitArray Type
ty ->
                          TExp Int64
-> VName
-> VName
-> TExp Int64
-> ImpM GPUMem KernelEnv KernelOp ()
getBitFromBitArray TExp Int64
chunk VName
par VName
p TExp Int64
i
                    Just VName
p ->
                      VName
-> [TExp Int64]
-> SubExp
-> [TExp Int64]
-> ImpM GPUMem KernelEnv KernelOp ()
forall rep r op.
VName -> [TExp Int64] -> SubExp -> [TExp Int64] -> ImpM rep r op ()
copyDWIMFix VName
par [] (VName -> SubExp
Var VName
p) [TExp Int64
i]
                    Maybe VName
Nothing -> () -> ImpM GPUMem KernelEnv KernelOp ()
forall a. a -> ImpM GPUMem KernelEnv KernelOp a
forall (f :: * -> *) a. Applicative f => a -> f a
pure ()

                  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
g -> VName
-> [TExp Int64]
-> SubExp
-> [TExp Int64]
-> ImpM GPUMem KernelEnv KernelOp ()
forall rep r op.
VName -> [TExp Int64] -> SubExp -> [TExp Int64] -> ImpM rep r op ()
copyDWIMFix VName
par [] (VName -> SubExp
Var VName
g) [TV Int64 -> TExp Int64
forall {k} (t :: k). TV t -> TExp t
tvExp TV Int64
phys_block_id, TExp Int64
ltid, TExp Int64
i]) Maybe VName
glob

            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
$ Body GPUMem -> [SubExpRes]
forall rep res. GBody rep res -> [res]
bodyResult (Body GPUMem -> [SubExpRes]) -> Body GPUMem -> [SubExpRes]
forall a b. (a -> b) -> a -> b
$ GLambda GPUMem Type -> Body GPUMem
forall rep t. GLambda rep t -> Body rep
lambdaBody (GLambda GPUMem Type -> Body GPUMem)
-> GLambda GPUMem Type -> Body GPUMem
forall a b. (a -> b) -> a -> b
$ SegPostOp GPUMem -> GLambda GPUMem Type
forall rep. SegPostOp rep -> Lambda rep
segPostOpLambda SegPostOp GPUMem
post_op
            sComment "compute post op." $
              compileStms mempty (bodyStms $ lambdaBody $ segPostOpLambda post_op) $
                sComment "write mapped values" $
                  forM_ (zip3 (patElems pat) res res_flags) $ \(PatElem LetDecMem
pe, SubExp
subexp, Bool
flag) ->
                    if Bool
flag
                      then () -> ImpM GPUMem KernelEnv KernelOp ()
forall a. a -> ImpM GPUMem KernelEnv KernelOp a
forall (f :: * -> *) a. Applicative f => a -> f a
pure ()
                      else VName
-> [TExp Int64]
-> SubExp
-> [TExp Int64]
-> ImpM GPUMem KernelEnv KernelOp ()
forall rep r op.
VName -> [TExp Int64] -> SubExp -> [TExp Int64] -> ImpM rep r op ()
copyDWIMFix (PatElem LetDecMem -> VName
forall dec. PatElem dec -> VName
patElemName PatElem LetDecMem
pe) ((VName -> TExp Int64) -> [VName] -> [TExp Int64]
forall a b. (a -> b) -> [a] -> [b]
map VName -> TExp Int64
forall a. a -> TPrimExp Int64 a
le64 [VName]
gtids) SubExp
subexp []

      sOp local_barrier
{-# NOINLINE compileSegScan #-}