{-# OPTIONS_GHC -Wno-incomplete-patterns #-}

-- | Multicore code generation for SegScan. Uses a fairly naive multipass
-- algorithm, with no particular locality optimisations.
module Futhark.CodeGen.ImpGen.Multicore.SegScan
  ( compileSegScan,
  )
where

import Control.Monad
import Data.List (zip4)
import Data.Maybe (isNothing)
import Futhark.CodeGen.ImpCode.Multicore qualified as Imp
import Futhark.CodeGen.ImpGen
import Futhark.CodeGen.ImpGen.Multicore.Base
import Futhark.IR.MCMem
import Futhark.Transform.Rename (renameBody)
import Futhark.Util.IntegralExp (ceilDiv)
import Prelude hiding (quot, rem)

-- This does not correspond with the actual cache size, but the actual cache
-- size may not really be the optimal value either. Ideally this should be
-- exposed as a tuning parameter.
cacheSize :: Imp.TExp Int64
cacheSize :: TExp Int64
cacheSize = TExp Int64
65536 -- 64 KB

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

lamBody :: SegBinOp MCMem -> Body MCMem
lamBody :: SegBinOp MCMem -> Body MCMem
lamBody = GLambda MCMem Type -> Body MCMem
forall rep t. GLambda rep t -> Body rep
lambdaBody (GLambda MCMem Type -> Body MCMem)
-> (SegBinOp MCMem -> GLambda MCMem Type)
-> SegBinOp MCMem
-> Body MCMem
forall b c a. (b -> c) -> (a -> b) -> a -> c
. SegBinOp MCMem -> GLambda MCMem Type
forall rep. SegBinOp rep -> Lambda rep
segBinOpLambda

genBinOpParams :: [SegBinOp MCMem] -> MulticoreGen ()
genBinOpParams :: [SegBinOp MCMem] -> MulticoreGen ()
genBinOpParams [SegBinOp MCMem]
scan_ops =
  Maybe (Exp MCMem) -> Scope MCMem -> MulticoreGen ()
forall rep (inner :: * -> *) r op.
Mem rep inner =>
Maybe (Exp rep) -> Scope rep -> ImpM rep r op ()
dScope Maybe (Exp MCMem)
forall a. Maybe a
Nothing (Scope MCMem -> MulticoreGen ()) -> Scope MCMem -> MulticoreGen ()
forall a b. (a -> b) -> a -> b
$
    [Param LParamMem] -> Scope MCMem
forall rep dec. (LParamInfo rep ~ dec) => [Param dec] -> Scope rep
scopeOfLParams ([Param LParamMem] -> Scope MCMem)
-> [Param LParamMem] -> Scope MCMem
forall a b. (a -> b) -> a -> b
$
      (SegBinOp MCMem -> [Param LParamMem])
-> [SegBinOp MCMem] -> [Param LParamMem]
forall (t :: * -> *) a b. Foldable t => (a -> [b]) -> t a -> [b]
concatMap (GLambda MCMem Type -> [LParam MCMem]
GLambda MCMem Type -> [Param LParamMem]
forall rep t. GLambda rep t -> [LParam rep]
lambdaParams (GLambda MCMem Type -> [Param LParamMem])
-> (SegBinOp MCMem -> GLambda MCMem Type)
-> SegBinOp MCMem
-> [Param LParamMem]
forall b c a. (b -> c) -> (a -> b) -> a -> c
. SegBinOp MCMem -> GLambda MCMem Type
forall rep. SegBinOp rep -> Lambda rep
segBinOpLambda) [SegBinOp MCMem]
scan_ops

initialiseLocalPrefixes :: [SegBinOp MCMem] -> [[VName]] -> MulticoreGen ()
initialiseLocalPrefixes :: [SegBinOp MCMem] -> [[VName]] -> MulticoreGen ()
initialiseLocalPrefixes [SegBinOp MCMem]
scan_ops [[VName]]
per_op_prefix_var = do
  scan_ops_renamed <- [SegBinOp MCMem] -> MulticoreGen [SegBinOp MCMem]
renameSegBinOp [SegBinOp MCMem]
scan_ops
  genBinOpParams scan_ops_renamed
  forM_ (zip scan_ops_renamed per_op_prefix_var) $ \(SegBinOp MCMem
scan_op, [VName]
prefix_vars) -> do
    ShapeBase SubExp
-> ([TExp Int64] -> MulticoreGen ()) -> MulticoreGen ()
forall rep r op.
ShapeBase SubExp
-> ([TExp Int64] -> ImpM rep r op ()) -> ImpM rep r op ()
sLoopNest (SegBinOp MCMem -> ShapeBase SubExp
forall rep. SegBinOp rep -> ShapeBase SubExp
segBinOpShape SegBinOp MCMem
scan_op) (([TExp Int64] -> MulticoreGen ()) -> MulticoreGen ())
-> ([TExp Int64] -> MulticoreGen ()) -> MulticoreGen ()
forall a b. (a -> b) -> a -> b
$ \[TExp Int64]
vec_is ->
      [(SubExp, VName)]
-> ((SubExp, VName) -> MulticoreGen ()) -> MulticoreGen ()
forall (t :: * -> *) (m :: * -> *) a b.
(Foldable t, Monad m) =>
t a -> (a -> m b) -> m ()
forM_ ([SubExp] -> [VName] -> [(SubExp, VName)]
forall a b. [a] -> [b] -> [(a, b)]
zip (SegBinOp MCMem -> [SubExp]
forall rep. SegBinOp rep -> [SubExp]
segBinOpNeutral SegBinOp MCMem
scan_op) [VName]
prefix_vars) (((SubExp, VName) -> MulticoreGen ()) -> MulticoreGen ())
-> ((SubExp, VName) -> MulticoreGen ()) -> MulticoreGen ()
forall a b. (a -> b) -> a -> b
$ \(SubExp
ne, VName
prefix_var) -> do
        VName -> [TExp Int64] -> SubExp -> [TExp Int64] -> MulticoreGen ()
forall rep r op.
VName -> [TExp Int64] -> SubExp -> [TExp Int64] -> ImpM rep r op ()
copyDWIMFix VName
prefix_var [TExp Int64]
vec_is SubExp
ne []

updateLocalPrefixes :: [SegBinOp MCMem] -> [[VName]] -> [[VName]] -> [Imp.TExp Int64] -> MulticoreGen ()
updateLocalPrefixes :: [SegBinOp MCMem]
-> [[VName]] -> [[VName]] -> [TExp Int64] -> MulticoreGen ()
updateLocalPrefixes [SegBinOp MCMem]
scan_ops [[VName]]
per_op_prefix_var [[VName]]
per_op_prefix_arr [TExp Int64]
index = do
  [(SegBinOp MCMem, [VName], [VName])]
-> ((SegBinOp MCMem, [VName], [VName]) -> MulticoreGen ())
-> MulticoreGen ()
forall (t :: * -> *) (m :: * -> *) a b.
(Foldable t, Monad m) =>
t a -> (a -> m b) -> m ()
forM_ ([SegBinOp MCMem]
-> [[VName]] -> [[VName]] -> [(SegBinOp MCMem, [VName], [VName])]
forall a b c. [a] -> [b] -> [c] -> [(a, b, c)]
zip3 [SegBinOp MCMem]
scan_ops [[VName]]
per_op_prefix_var [[VName]]
per_op_prefix_arr) (((SegBinOp MCMem, [VName], [VName]) -> MulticoreGen ())
 -> MulticoreGen ())
-> ((SegBinOp MCMem, [VName], [VName]) -> MulticoreGen ())
-> MulticoreGen ()
forall a b. (a -> b) -> a -> b
$ \(SegBinOp MCMem
scan_op, [VName]
prefix_vars, [VName]
prefix_arrs) -> do
    let shape :: ShapeBase SubExp
shape = SegBinOp MCMem -> ShapeBase SubExp
forall rep. SegBinOp rep -> ShapeBase SubExp
segBinOpShape SegBinOp MCMem
scan_op
    ShapeBase SubExp
-> ([TExp Int64] -> MulticoreGen ()) -> MulticoreGen ()
forall rep r op.
ShapeBase SubExp
-> ([TExp Int64] -> ImpM rep r op ()) -> ImpM rep r op ()
sLoopNest ShapeBase SubExp
shape (([TExp Int64] -> MulticoreGen ()) -> MulticoreGen ())
-> ([TExp Int64] -> MulticoreGen ()) -> MulticoreGen ()
forall a b. (a -> b) -> a -> b
$ \[TExp Int64]
vec_is ->
      [(VName, VName)]
-> ((VName, VName) -> MulticoreGen ()) -> MulticoreGen ()
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]
prefix_vars [VName]
prefix_arrs) (((VName, VName) -> MulticoreGen ()) -> MulticoreGen ())
-> ((VName, VName) -> MulticoreGen ()) -> MulticoreGen ()
forall a b. (a -> b) -> a -> b
$ \(VName
prefix_var, VName
prefix_arr) -> do
        VName -> [TExp Int64] -> SubExp -> [TExp Int64] -> MulticoreGen ()
forall rep r op.
VName -> [TExp Int64] -> SubExp -> [TExp Int64] -> ImpM rep r op ()
copyDWIMFix VName
prefix_var [TExp Int64]
vec_is (VName -> SubExp
Var VName
prefix_arr) ([TExp Int64]
index [TExp Int64] -> [TExp Int64] -> [TExp Int64]
forall a. [a] -> [a] -> [a]
++ [TExp Int64]
vec_is)

genArrays :: [SegBinOp MCMem] -> Name -> Shape -> MulticoreGen [[VName]]
genArrays :: [SegBinOp MCMem]
-> Name -> ShapeBase SubExp -> MulticoreGen [[VName]]
genArrays [SegBinOp MCMem]
scan_ops Name
arr_name ShapeBase SubExp
block_shape = do
  [SegBinOp MCMem]
-> (SegBinOp MCMem -> ImpM MCMem HostEnv Multicore [VName])
-> MulticoreGen [[VName]]
forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
t a -> (a -> m b) -> m (t b)
forM [SegBinOp MCMem]
scan_ops ((SegBinOp MCMem -> ImpM MCMem HostEnv Multicore [VName])
 -> MulticoreGen [[VName]])
-> (SegBinOp MCMem -> ImpM MCMem HostEnv Multicore [VName])
-> MulticoreGen [[VName]]
forall a b. (a -> b) -> a -> b
$ \SegBinOp MCMem
scan_op ->
    [Type]
-> (Type -> ImpM MCMem HostEnv Multicore VName)
-> ImpM MCMem HostEnv Multicore [VName]
forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
t a -> (a -> m b) -> m (t b)
forM (GLambda MCMem Type -> [Type]
forall rep t. GLambda rep t -> [t]
lambdaReturnType (GLambda MCMem Type -> [Type]) -> GLambda MCMem Type -> [Type]
forall a b. (a -> b) -> a -> b
$ SegBinOp MCMem -> GLambda MCMem Type
forall rep. SegBinOp rep -> Lambda rep
segBinOpLambda SegBinOp MCMem
scan_op) ((Type -> ImpM MCMem HostEnv Multicore VName)
 -> ImpM MCMem HostEnv Multicore [VName])
-> (Type -> ImpM MCMem HostEnv Multicore VName)
-> ImpM MCMem HostEnv Multicore [VName]
forall a b. (a -> b) -> a -> b
$ \Type
t -> do
      let shape :: ShapeBase SubExp
shape = ShapeBase SubExp
block_shape ShapeBase SubExp -> ShapeBase SubExp -> ShapeBase SubExp
forall a. Semigroup a => a -> a -> a
<> SegBinOp MCMem -> ShapeBase SubExp
forall rep. SegBinOp rep -> ShapeBase SubExp
segBinOpShape SegBinOp MCMem
scan_op 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
      Name
-> PrimType
-> ShapeBase SubExp
-> Space
-> ImpM MCMem HostEnv Multicore VName
forall rep r op.
Name
-> PrimType -> ShapeBase SubExp -> Space -> ImpM rep r op VName
sAllocArray Name
arr_name (Type -> PrimType
forall shape o. TypeBase shape o -> PrimType
elemType Type
t) ShapeBase SubExp
shape Space
DefaultSpace

genLocalArray :: [SegBinOp MCMem] -> MulticoreGen [[VName]]
genLocalArray :: [SegBinOp MCMem] -> MulticoreGen [[VName]]
genLocalArray [SegBinOp MCMem]
scan_ops = do
  [SegBinOp MCMem]
-> (SegBinOp MCMem -> ImpM MCMem HostEnv Multicore [VName])
-> MulticoreGen [[VName]]
forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
t a -> (a -> m b) -> m (t b)
forM [SegBinOp MCMem]
scan_ops ((SegBinOp MCMem -> ImpM MCMem HostEnv Multicore [VName])
 -> MulticoreGen [[VName]])
-> (SegBinOp MCMem -> ImpM MCMem HostEnv Multicore [VName])
-> MulticoreGen [[VName]]
forall a b. (a -> b) -> a -> b
$ \SegBinOp MCMem
scan_op -> do
    let shape :: ShapeBase SubExp
shape = SegBinOp MCMem -> ShapeBase SubExp
forall rep. SegBinOp rep -> ShapeBase SubExp
segBinOpShape SegBinOp MCMem
scan_op
        ts :: [Type]
ts = GLambda MCMem Type -> [Type]
forall rep t. GLambda rep t -> [t]
lambdaReturnType (GLambda MCMem Type -> [Type]) -> GLambda MCMem Type -> [Type]
forall a b. (a -> b) -> a -> b
$ SegBinOp MCMem -> GLambda MCMem Type
forall rep. SegBinOp rep -> Lambda rep
segBinOpLambda SegBinOp MCMem
scan_op
    [(Param LParamMem, Type)]
-> ((Param LParamMem, Type) -> ImpM MCMem HostEnv Multicore VName)
-> ImpM MCMem HostEnv Multicore [VName]
forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
t a -> (a -> m b) -> m (t b)
forM ([Param LParamMem] -> [Type] -> [(Param LParamMem, Type)]
forall a b. [a] -> [b] -> [(a, b)]
zip (SegBinOp MCMem -> [LParam MCMem]
xParams SegBinOp MCMem
scan_op) [Type]
ts) (((Param LParamMem, Type) -> ImpM MCMem HostEnv Multicore VName)
 -> ImpM MCMem HostEnv Multicore [VName])
-> ((Param LParamMem, Type) -> ImpM MCMem HostEnv Multicore VName)
-> ImpM MCMem HostEnv Multicore [VName]
forall a b. (a -> b) -> a -> b
$ \(Param LParamMem
p, Type
t) -> do
      case ShapeBase SubExp -> [SubExp]
forall d. ShapeBase d -> [d]
shapeDims ShapeBase SubExp
shape of
        [] -> VName -> ImpM MCMem HostEnv Multicore VName
forall a. a -> ImpM MCMem HostEnv Multicore a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (VName -> ImpM MCMem HostEnv Multicore VName)
-> VName -> ImpM MCMem HostEnv Multicore VName
forall a b. (a -> b) -> a -> b
$ Param LParamMem -> VName
forall dec. Param dec -> VName
paramName Param LParamMem
p
        [SubExp]
_ -> do
          let pt :: PrimType
pt = Type -> PrimType
forall shape o. TypeBase shape o -> PrimType
elemType Type
t
          Name
-> PrimType
-> ShapeBase SubExp
-> Space
-> ImpM MCMem HostEnv Multicore VName
forall rep r op.
Name
-> PrimType -> ShapeBase SubExp -> Space -> ImpM rep r op VName
sAllocArray Name
"local_acc" PrimType
pt (ShapeBase SubExp
shape 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) Space
DefaultSpace

totalBytes :: [SegBinOp MCMem] -> Imp.TExp Int64
totalBytes :: [SegBinOp MCMem] -> TExp Int64
totalBytes [SegBinOp MCMem]
scan_ops =
  [TExp Int64] -> TExp Int64
forall a. Num a => [a] -> a
forall (t :: * -> *) a. (Foldable t, Num a) => t a -> a
sum
    [ PrimType -> ShapeBase SubExp -> TExp Int64
bytesOfShape (Type -> PrimType
forall shape o. TypeBase shape o -> PrimType
elemType Type
t) (Type -> ShapeBase SubExp
forall shape o. ArrayShape shape => TypeBase shape o -> shape
arrayShape Type
t ShapeBase SubExp -> ShapeBase SubExp -> ShapeBase SubExp
forall a. Semigroup a => a -> a -> a
<> ShapeBase SubExp
scan_shape)
    | SegBinOp MCMem
op <- [SegBinOp MCMem]
scan_ops,
      let lam :: GLambda MCMem Type
lam = SegBinOp MCMem -> GLambda MCMem Type
forall rep. SegBinOp rep -> Lambda rep
segBinOpLambda SegBinOp MCMem
op,
      let scan_shape :: ShapeBase SubExp
scan_shape = SegBinOp MCMem -> ShapeBase SubExp
forall rep. SegBinOp rep -> ShapeBase SubExp
segBinOpShape SegBinOp MCMem
op,
      Type
t <- GLambda MCMem Type -> [Type]
forall rep t. GLambda rep t -> [t]
lambdaReturnType GLambda MCMem Type
lam
    ]
  where
    bytesOfShape :: PrimType -> ShapeBase SubExp -> TExp Int64
bytesOfShape PrimType
pt ShapeBase SubExp
sh = PrimType -> TExp Int64
forall a. Num a => PrimType -> a
primByteSize PrimType
pt TExp Int64 -> TExp Int64 -> TExp Int64
forall a. Num a => a -> a -> a
* [TExp Int64] -> TExp Int64
forall a. Num a => [a] -> a
forall (t :: * -> *) a. (Foldable t, Num a) => t a -> a
product ((SubExp -> TExp Int64) -> [SubExp] -> [TExp Int64]
forall a b. (a -> b) -> [a] -> [b]
map SubExp -> TExp Int64
pe64 (ShapeBase SubExp -> [SubExp]
forall d. ShapeBase d -> [d]
shapeDims ShapeBase SubExp
sh))

bodyHas :: (Exp MCMem -> Bool) -> GBody MCMem res -> Bool
bodyHas :: forall res. (Exp MCMem -> Bool) -> GBody MCMem res -> Bool
bodyHas Exp MCMem -> Bool
f = (Stm MCMem -> Bool) -> Seq (Stm MCMem) -> Bool
forall (t :: * -> *) a. Foldable t => (a -> Bool) -> t a -> Bool
any (Exp MCMem -> Bool
f' (Exp MCMem -> Bool)
-> (Stm MCMem -> Exp MCMem) -> Stm MCMem -> Bool
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Stm MCMem -> Exp MCMem
forall rep. Stm rep -> Exp rep
stmExp) (Seq (Stm MCMem) -> Bool)
-> (GBody MCMem res -> Seq (Stm MCMem)) -> GBody MCMem res -> Bool
forall b c a. (b -> c) -> (a -> b) -> a -> c
. GBody MCMem res -> Seq (Stm MCMem)
forall rep res. GBody rep res -> Stms rep
bodyStms
  where
    f' :: Exp MCMem -> Bool
f' Exp MCMem
e
      | Exp MCMem -> Bool
f Exp MCMem
e = Bool
True
      | Bool
otherwise = Maybe () -> Bool
forall a. Maybe a -> Bool
isNothing (Maybe () -> Bool) -> Maybe () -> Bool
forall a b. (a -> b) -> a -> b
$ Walker MCMem Maybe -> Exp MCMem -> Maybe ()
forall (m :: * -> *) rep.
Monad m =>
Walker rep m -> Exp rep -> m ()
walkExpM Walker MCMem Maybe
walker Exp MCMem
e
    walker :: Walker MCMem Maybe
walker =
      Walker MCMem Maybe
forall rep (m :: * -> *). Monad m => Walker rep m
identityWalker
        { walkOnBody = const $ guard . not . bodyHas f
        }

-- | Determine whether this kernel body should be recomputed. Involves both
-- correctness checks and (crude) efficiency checks. Basically, recomputation is
-- only safe if nothing is consumed in the body, and considered efficient only
-- if the body has no loops.
shouldRecompute :: KernelBody MCMem -> Bool
shouldRecompute :: KernelBody MCMem -> Bool
shouldRecompute = Bool -> Bool
not (Bool -> Bool)
-> (KernelBody MCMem -> Bool) -> KernelBody MCMem -> Bool
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (Exp MCMem -> Bool) -> KernelBody MCMem -> Bool
forall res. (Exp MCMem -> Bool) -> GBody MCMem res -> Bool
bodyHas Exp MCMem -> Bool
forall {rep}. Exp rep -> Bool
bad
  where
    bad :: Exp rep -> Bool
bad (BasicOp Update {}) = Bool
True
    bad (BasicOp UpdateAcc {}) = Bool
True
    bad (WithAcc {}) = Bool
True
    bad Loop {} = Bool
True
    bad Exp rep
_ = Bool
False

hasAssert :: GBody MCMem res -> Bool
hasAssert :: forall res. GBody MCMem res -> Bool
hasAssert = (Exp MCMem -> Bool) -> GBody MCMem res -> Bool
forall res. (Exp MCMem -> Bool) -> GBody MCMem res -> Bool
bodyHas Exp MCMem -> Bool
forall {rep}. Exp rep -> Bool
isAssert
  where
    isAssert :: Exp rep -> Bool
isAssert (BasicOp Assert {}) = Bool
True
    isAssert Exp rep
_ = Bool
False

copyFromDescToLocal :: [SegBinOp MCMem] -> [[VName]] -> [[VName]] -> TV Int64 -> MulticoreGen ()
copyFromDescToLocal :: [SegBinOp MCMem]
-> [[VName]] -> [[VName]] -> TV Int64 -> MulticoreGen ()
copyFromDescToLocal [SegBinOp MCMem]
scan_ops [[VName]]
per_op_local_vars [[VName]]
per_op_description_arrays TV Int64
index = do
  [(SegBinOp MCMem, [VName], [VName])]
-> ((SegBinOp MCMem, [VName], [VName]) -> MulticoreGen ())
-> MulticoreGen ()
forall (t :: * -> *) (m :: * -> *) a b.
(Foldable t, Monad m) =>
t a -> (a -> m b) -> m ()
forM_ ([SegBinOp MCMem]
-> [[VName]] -> [[VName]] -> [(SegBinOp MCMem, [VName], [VName])]
forall a b c. [a] -> [b] -> [c] -> [(a, b, c)]
zip3 [SegBinOp MCMem]
scan_ops [[VName]]
per_op_local_vars [[VName]]
per_op_description_arrays) (((SegBinOp MCMem, [VName], [VName]) -> MulticoreGen ())
 -> MulticoreGen ())
-> ((SegBinOp MCMem, [VName], [VName]) -> MulticoreGen ())
-> MulticoreGen ()
forall a b. (a -> b) -> a -> b
$ \(SegBinOp MCMem
scan_op, [VName]
local_vars, [VName]
desc_arrs) -> do
    ShapeBase SubExp
-> ([TExp Int64] -> MulticoreGen ()) -> MulticoreGen ()
forall rep r op.
ShapeBase SubExp
-> ([TExp Int64] -> ImpM rep r op ()) -> ImpM rep r op ()
sLoopNest (SegBinOp MCMem -> ShapeBase SubExp
forall rep. SegBinOp rep -> ShapeBase SubExp
segBinOpShape SegBinOp MCMem
scan_op) (([TExp Int64] -> MulticoreGen ()) -> MulticoreGen ())
-> ([TExp Int64] -> MulticoreGen ()) -> MulticoreGen ()
forall a b. (a -> b) -> a -> b
$ \[TExp Int64]
vec_is -> do
      [(VName, VName)]
-> ((VName, VName) -> MulticoreGen ()) -> MulticoreGen ()
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]
local_vars [VName]
desc_arrs) (((VName, VName) -> MulticoreGen ()) -> MulticoreGen ())
-> ((VName, VName) -> MulticoreGen ()) -> MulticoreGen ()
forall a b. (a -> b) -> a -> b
$ \(VName
local_var, VName
desc_arr) -> do
        VName -> [TExp Int64] -> SubExp -> [TExp Int64] -> MulticoreGen ()
forall rep r op.
VName -> [TExp Int64] -> SubExp -> [TExp Int64] -> ImpM rep r op ()
copyDWIMFix VName
local_var [TExp Int64]
vec_is (VName -> SubExp
Var VName
desc_arr) (TV Int64 -> TExp Int64
forall {k} (t :: k). TV t -> TExp t
tvExp TV Int64
index TExp Int64 -> [TExp Int64] -> [TExp Int64]
forall a. a -> [a] -> [a]
: [TExp Int64]
vec_is)

applyScanFromDescToLocal ::
  [SegBinOp MCMem] ->
  [[VName]] ->
  [[VName]] ->
  TV Int64 ->
  MulticoreGen ()
applyScanFromDescToLocal :: [SegBinOp MCMem]
-> [[VName]] -> [[VName]] -> TV Int64 -> MulticoreGen ()
applyScanFromDescToLocal [SegBinOp MCMem]
scan_ops [[VName]]
per_op_local_vars [[VName]]
per_op_description_arrays TV Int64
index = do
  scan_ops_renamed <- [SegBinOp MCMem] -> MulticoreGen [SegBinOp MCMem]
renameSegBinOp [SegBinOp MCMem]
scan_ops
  genBinOpParams scan_ops_renamed
  forM_ (zip3 scan_ops_renamed per_op_local_vars per_op_description_arrays) $ \(SegBinOp MCMem
scan_op, [VName]
local_vars, [VName]
description_array) -> do
    ShapeBase SubExp
-> ([TExp Int64] -> MulticoreGen ()) -> MulticoreGen ()
forall rep r op.
ShapeBase SubExp
-> ([TExp Int64] -> ImpM rep r op ()) -> ImpM rep r op ()
sLoopNest (SegBinOp MCMem -> ShapeBase SubExp
forall rep. SegBinOp rep -> ShapeBase SubExp
segBinOpShape SegBinOp MCMem
scan_op) (([TExp Int64] -> MulticoreGen ()) -> MulticoreGen ())
-> ([TExp Int64] -> MulticoreGen ()) -> MulticoreGen ()
forall a b. (a -> b) -> a -> b
$ \[TExp Int64]
vec_is -> do
      [(Param LParamMem, VName)]
-> ((Param LParamMem, VName) -> MulticoreGen ()) -> MulticoreGen ()
forall (t :: * -> *) (m :: * -> *) a b.
(Foldable t, Monad m) =>
t a -> (a -> m b) -> m ()
forM_ ([Param LParamMem] -> [VName] -> [(Param LParamMem, VName)]
forall a b. [a] -> [b] -> [(a, b)]
zip (SegBinOp MCMem -> [LParam MCMem]
xParams SegBinOp MCMem
scan_op) [VName]
description_array) (((Param LParamMem, VName) -> MulticoreGen ()) -> MulticoreGen ())
-> ((Param LParamMem, VName) -> MulticoreGen ()) -> MulticoreGen ()
forall a b. (a -> b) -> a -> b
$ \(Param LParamMem
xs, VName
desc_var) -> do
        VName -> [TExp Int64] -> SubExp -> [TExp Int64] -> MulticoreGen ()
forall rep r op.
VName -> [TExp Int64] -> SubExp -> [TExp Int64] -> ImpM rep r op ()
copyDWIMFix (Param LParamMem -> VName
forall dec. Param dec -> VName
paramName Param LParamMem
xs) [] (VName -> SubExp
Var VName
desc_var) (TV Int64 -> TExp Int64
forall {k} (t :: k). TV t -> TExp t
tvExp TV Int64
index TExp Int64 -> [TExp Int64] -> [TExp Int64]
forall a. a -> [a] -> [a]
: [TExp Int64]
vec_is)
      [(Param LParamMem, VName)]
-> ((Param LParamMem, VName) -> MulticoreGen ()) -> MulticoreGen ()
forall (t :: * -> *) (m :: * -> *) a b.
(Foldable t, Monad m) =>
t a -> (a -> m b) -> m ()
forM_ ([Param LParamMem] -> [VName] -> [(Param LParamMem, VName)]
forall a b. [a] -> [b] -> [(a, b)]
zip (SegBinOp MCMem -> [LParam MCMem]
yParams SegBinOp MCMem
scan_op) [VName]
local_vars) (((Param LParamMem, VName) -> MulticoreGen ()) -> MulticoreGen ())
-> ((Param LParamMem, VName) -> MulticoreGen ()) -> MulticoreGen ()
forall a b. (a -> b) -> a -> b
$ \(Param LParamMem
ys, VName
local_var) -> do
        VName -> [TExp Int64] -> SubExp -> [TExp Int64] -> MulticoreGen ()
forall rep r op.
VName -> [TExp Int64] -> SubExp -> [TExp Int64] -> ImpM rep r op ()
copyDWIMFix (Param LParamMem -> VName
forall dec. Param dec -> VName
paramName Param LParamMem
ys) [] (VName -> SubExp
Var VName
local_var) [TExp Int64]
vec_is
      Names -> Seq (Stm MCMem) -> MulticoreGen () -> MulticoreGen ()
forall rep r op.
Names -> Stms rep -> ImpM rep r op () -> ImpM rep r op ()
compileStms Names
forall a. Monoid a => a
mempty (Body MCMem -> Seq (Stm MCMem)
forall rep res. GBody rep res -> Stms rep
bodyStms (Body MCMem -> Seq (Stm MCMem)) -> Body MCMem -> Seq (Stm MCMem)
forall a b. (a -> b) -> a -> b
$ SegBinOp MCMem -> Body MCMem
lamBody SegBinOp MCMem
scan_op) (MulticoreGen () -> MulticoreGen ())
-> MulticoreGen () -> MulticoreGen ()
forall a b. (a -> b) -> a -> b
$ do
        [(SubExp, VName)]
-> ((SubExp, VName) -> MulticoreGen ()) -> MulticoreGen ()
forall (t :: * -> *) (m :: * -> *) a b.
(Foldable t, Monad m) =>
t a -> (a -> m b) -> m ()
forM_ ([SubExp] -> [VName] -> [(SubExp, VName)]
forall a b. [a] -> [b] -> [(a, b)]
zip ((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 MCMem -> [SubExpRes]
forall rep res. GBody rep res -> [res]
bodyResult (Body MCMem -> [SubExpRes]) -> Body MCMem -> [SubExpRes]
forall a b. (a -> b) -> a -> b
$ SegBinOp MCMem -> Body MCMem
lamBody SegBinOp MCMem
scan_op) [VName]
local_vars) (((SubExp, VName) -> MulticoreGen ()) -> MulticoreGen ())
-> ((SubExp, VName) -> MulticoreGen ()) -> MulticoreGen ()
forall a b. (a -> b) -> a -> b
$ \(SubExp
se, VName
local_var) -> do
          VName -> [TExp Int64] -> SubExp -> [TExp Int64] -> MulticoreGen ()
forall rep r op.
VName -> [TExp Int64] -> SubExp -> [TExp Int64] -> ImpM rep r op ()
copyDWIMFix VName
local_var [TExp Int64]
vec_is SubExp
se []

applyScanAggregateToPrefix ::
  [SegBinOp MCMem] ->
  [[VName]] ->
  [[VName]] ->
  [[VName]] ->
  TV Int64 ->
  MulticoreGen ()
applyScanAggregateToPrefix :: [SegBinOp MCMem]
-> [[VName]]
-> [[VName]]
-> [[VName]]
-> TV Int64
-> MulticoreGen ()
applyScanAggregateToPrefix [SegBinOp MCMem]
scan_ops [[VName]]
per_op_local_vars [[VName]]
per_op_aggr_arrs [[VName]]
prefArrs TV Int64
block_idx = do
  scan_ops_renamed <- [SegBinOp MCMem] -> MulticoreGen [SegBinOp MCMem]
renameSegBinOp [SegBinOp MCMem]
scan_ops
  genBinOpParams scan_ops_renamed
  forM_ (zip4 scan_ops_renamed per_op_local_vars per_op_aggr_arrs prefArrs) $ \(SegBinOp MCMem
scan_op, [VName]
local_vars, [VName]
aggr_arrs, [VName]
prifix_arrs) -> do
    ShapeBase SubExp
-> ([TExp Int64] -> MulticoreGen ()) -> MulticoreGen ()
forall rep r op.
ShapeBase SubExp
-> ([TExp Int64] -> ImpM rep r op ()) -> ImpM rep r op ()
sLoopNest (SegBinOp MCMem -> ShapeBase SubExp
forall rep. SegBinOp rep -> ShapeBase SubExp
segBinOpShape SegBinOp MCMem
scan_op) (([TExp Int64] -> MulticoreGen ()) -> MulticoreGen ())
-> ([TExp Int64] -> MulticoreGen ()) -> MulticoreGen ()
forall a b. (a -> b) -> a -> b
$ \[TExp Int64]
vec_is -> do
      [(Param LParamMem, VName)]
-> ((Param LParamMem, VName) -> MulticoreGen ()) -> MulticoreGen ()
forall (t :: * -> *) (m :: * -> *) a b.
(Foldable t, Monad m) =>
t a -> (a -> m b) -> m ()
forM_ ([Param LParamMem] -> [VName] -> [(Param LParamMem, VName)]
forall a b. [a] -> [b] -> [(a, b)]
zip (SegBinOp MCMem -> [LParam MCMem]
xParams SegBinOp MCMem
scan_op) [VName]
local_vars) (((Param LParamMem, VName) -> MulticoreGen ()) -> MulticoreGen ())
-> ((Param LParamMem, VName) -> MulticoreGen ()) -> MulticoreGen ()
forall a b. (a -> b) -> a -> b
$ \(Param LParamMem
xs, VName
local_var) -> do
        VName -> [TExp Int64] -> SubExp -> [TExp Int64] -> MulticoreGen ()
forall rep r op.
VName -> [TExp Int64] -> SubExp -> [TExp Int64] -> ImpM rep r op ()
copyDWIMFix (Param LParamMem -> VName
forall dec. Param dec -> VName
paramName Param LParamMem
xs) [] (VName -> SubExp
Var VName
local_var) [TExp Int64]
vec_is
      [(Param LParamMem, VName)]
-> ((Param LParamMem, VName) -> MulticoreGen ()) -> MulticoreGen ()
forall (t :: * -> *) (m :: * -> *) a b.
(Foldable t, Monad m) =>
t a -> (a -> m b) -> m ()
forM_ ([Param LParamMem] -> [VName] -> [(Param LParamMem, VName)]
forall a b. [a] -> [b] -> [(a, b)]
zip (SegBinOp MCMem -> [LParam MCMem]
yParams SegBinOp MCMem
scan_op) [VName]
aggr_arrs) (((Param LParamMem, VName) -> MulticoreGen ()) -> MulticoreGen ())
-> ((Param LParamMem, VName) -> MulticoreGen ()) -> MulticoreGen ()
forall a b. (a -> b) -> a -> b
$ \(Param LParamMem
ys, VName
aggrArr) -> do
        VName -> [TExp Int64] -> SubExp -> [TExp Int64] -> MulticoreGen ()
forall rep r op.
VName -> [TExp Int64] -> SubExp -> [TExp Int64] -> ImpM rep r op ()
copyDWIMFix (Param LParamMem -> VName
forall dec. Param dec -> VName
paramName Param LParamMem
ys) [] (VName -> SubExp
Var VName
aggrArr) (TV Int64 -> TExp Int64
forall {k} (t :: k). TV t -> TExp t
tvExp TV Int64
block_idx TExp Int64 -> [TExp Int64] -> [TExp Int64]
forall a. a -> [a] -> [a]
: [TExp Int64]
vec_is)

      Names -> Seq (Stm MCMem) -> MulticoreGen () -> MulticoreGen ()
forall rep r op.
Names -> Stms rep -> ImpM rep r op () -> ImpM rep r op ()
compileStms Names
forall a. Monoid a => a
mempty (Body MCMem -> Seq (Stm MCMem)
forall rep res. GBody rep res -> Stms rep
bodyStms (Body MCMem -> Seq (Stm MCMem)) -> Body MCMem -> Seq (Stm MCMem)
forall a b. (a -> b) -> a -> b
$ SegBinOp MCMem -> Body MCMem
lamBody SegBinOp MCMem
scan_op) (MulticoreGen () -> MulticoreGen ())
-> MulticoreGen () -> MulticoreGen ()
forall a b. (a -> b) -> a -> b
$ do
        [(SubExp, VName)]
-> ((SubExp, VName) -> MulticoreGen ()) -> MulticoreGen ()
forall (t :: * -> *) (m :: * -> *) a b.
(Foldable t, Monad m) =>
t a -> (a -> m b) -> m ()
forM_ ([SubExp] -> [VName] -> [(SubExp, VName)]
forall a b. [a] -> [b] -> [(a, b)]
zip ((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 MCMem -> [SubExpRes]
forall rep res. GBody rep res -> [res]
bodyResult (Body MCMem -> [SubExpRes]) -> Body MCMem -> [SubExpRes]
forall a b. (a -> b) -> a -> b
$ SegBinOp MCMem -> Body MCMem
lamBody SegBinOp MCMem
scan_op) [VName]
prifix_arrs) (((SubExp, VName) -> MulticoreGen ()) -> MulticoreGen ())
-> ((SubExp, VName) -> MulticoreGen ()) -> MulticoreGen ()
forall a b. (a -> b) -> a -> b
$ \(SubExp
se, VName
prefix_arr) -> do
          VName -> [TExp Int64] -> SubExp -> [TExp Int64] -> MulticoreGen ()
forall rep r op.
VName -> [TExp Int64] -> SubExp -> [TExp Int64] -> ImpM rep r op ()
copyDWIMFix VName
prefix_arr (TV Int64 -> TExp Int64
forall {k} (t :: k). TV t -> TExp t
tvExp TV Int64
block_idx TExp Int64 -> [TExp Int64] -> [TExp Int64]
forall a. a -> [a] -> [a]
: [TExp Int64]
vec_is) SubExp
se []

seqScanFastPath ::
  [[VName]] ->
  [Maybe VName] ->
  VName ->
  [SegBinOp MCMem] ->
  KernelBody MCMem ->
  [[VName]] ->
  TV Int64 ->
  TV Int64 ->
  [[VName]] ->
  TV Int64 ->
  TV Int64 ->
  MulticoreGen ()
seqScanFastPath :: [[VName]]
-> [Maybe VName]
-> VName
-> [SegBinOp MCMem]
-> KernelBody MCMem
-> [[VName]]
-> TV Int64
-> TV Int64
-> [[VName]]
-> TV Int64
-> TV Int64
-> MulticoreGen ()
seqScanFastPath [[VName]]
scan_out [Maybe VName]
map_out VName
i [SegBinOp MCMem]
scan_ops KernelBody MCMem
kbody [[VName]]
per_op_prefixes_var TV Int64
start TV Int64
chunk_length [[VName]]
per_op_prefix_arr TV Int64
block_idx TV Int64
task_id = do
  kbody_renamed <- KernelBody MCMem -> ImpM MCMem HostEnv Multicore (KernelBody MCMem)
forall rep res (m :: * -> *).
(Renameable rep, Rename res, MonadFreshNames m) =>
GBody rep res -> m (GBody rep res)
renameBody KernelBody MCMem
kbody
  scan_ops_renamed <- renameSegBinOp scan_ops
  genBinOpParams scan_ops_renamed
  per_op_local_accum <- genLocalArray scan_ops_renamed

  let results = KernelBody MCMem -> [KernelResult]
forall rep res. GBody rep res -> [res]
bodyResult KernelBody MCMem
kbody_renamed
  let n_scan = [SegBinOp MCMem] -> Int
forall rep. [SegBinOp rep] -> Int
segBinOpResults [SegBinOp MCMem]
scan_ops_renamed
  let (all_scan_res, map_res) = splitAt n_scan results

  let per_scan_res = [SegBinOp MCMem] -> [KernelResult] -> [[KernelResult]]
forall rep a. [SegBinOp rep] -> [a] -> [[a]]
segBinOpChunks [SegBinOp MCMem]
scan_ops_renamed [KernelResult]
all_scan_res

  forM_ (zip3 scan_ops_renamed per_op_prefixes_var per_op_local_accum) $ \(SegBinOp MCMem
scan_op, [VName]
prefix_vars, [VName]
local_accums) ->
    ShapeBase SubExp
-> ([TExp Int64] -> MulticoreGen ()) -> MulticoreGen ()
forall rep r op.
ShapeBase SubExp
-> ([TExp Int64] -> ImpM rep r op ()) -> ImpM rep r op ()
sLoopNest (SegBinOp MCMem -> ShapeBase SubExp
forall rep. SegBinOp rep -> ShapeBase SubExp
segBinOpShape SegBinOp MCMem
scan_op) (([TExp Int64] -> MulticoreGen ()) -> MulticoreGen ())
-> ([TExp Int64] -> MulticoreGen ()) -> MulticoreGen ()
forall a b. (a -> b) -> a -> b
$ \[TExp Int64]
vec_is -> do
      [(VName, VName)]
-> ((VName, VName) -> MulticoreGen ()) -> MulticoreGen ()
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]
local_accums [VName]
prefix_vars) (((VName, VName) -> MulticoreGen ()) -> MulticoreGen ())
-> ((VName, VName) -> MulticoreGen ()) -> MulticoreGen ()
forall a b. (a -> b) -> a -> b
$ \(VName
acc, VName
prefix) -> do
        VName -> [TExp Int64] -> SubExp -> [TExp Int64] -> MulticoreGen ()
forall rep r op.
VName -> [TExp Int64] -> SubExp -> [TExp Int64] -> ImpM rep r op ()
copyDWIMFix VName
acc [TExp Int64]
vec_is (VName -> SubExp
Var VName
prefix) [TExp Int64]
vec_is

  z <- dPrimV "z" (0 :: Imp.TExp Int64)
  sWhile (tvExp z .<. tvExp chunk_length) $ do
    dPrimV_ i (tvExp start + tvExp z)
    compileStms mempty (bodyStms kbody_renamed) $ do
      sComment "write mapped values results to memory" $
        forM_ (zip map_out (map kernelResultSubExp map_res)) $ \(Maybe VName
marr, SubExp
res) ->
          Maybe VName -> (VName -> MulticoreGen ()) -> MulticoreGen ()
forall (t :: * -> *) (m :: * -> *) a b.
(Foldable t, Monad m) =>
t a -> (a -> m b) -> m ()
forM_ Maybe VName
marr ((VName -> MulticoreGen ()) -> MulticoreGen ())
-> (VName -> MulticoreGen ()) -> MulticoreGen ()
forall a b. (a -> b) -> a -> b
$ \VName
arr -> VName -> [TExp Int64] -> SubExp -> [TExp Int64] -> MulticoreGen ()
forall rep r op.
VName -> [TExp Int64] -> SubExp -> [TExp Int64] -> ImpM rep r op ()
copyDWIMFix VName
arr [TV Int64 -> TExp Int64
forall {k} (t :: k). TV t -> TExp t
tvExp TV Int64
task_id, TV Int64 -> TExp Int64
forall {k} (t :: k). TV t -> TExp t
tvExp TV Int64
z] SubExp
res []

      forM_ (zip4 scan_out scan_ops_renamed per_scan_res per_op_local_accum) $ \([VName]
pes, SegBinOp MCMem
scan_op, [KernelResult]
scan_res, [VName]
local_accums) ->
        ShapeBase SubExp
-> ([TExp Int64] -> MulticoreGen ()) -> MulticoreGen ()
forall rep r op.
ShapeBase SubExp
-> ([TExp Int64] -> ImpM rep r op ()) -> ImpM rep r op ()
sLoopNest (SegBinOp MCMem -> ShapeBase SubExp
forall rep. SegBinOp rep -> ShapeBase SubExp
segBinOpShape SegBinOp MCMem
scan_op) (([TExp Int64] -> MulticoreGen ()) -> MulticoreGen ())
-> ([TExp Int64] -> MulticoreGen ()) -> MulticoreGen ()
forall a b. (a -> b) -> a -> b
$ \[TExp Int64]
vec_is -> do
          [(Param LParamMem, VName)]
-> ((Param LParamMem, VName) -> MulticoreGen ()) -> MulticoreGen ()
forall (t :: * -> *) (m :: * -> *) a b.
(Foldable t, Monad m) =>
t a -> (a -> m b) -> m ()
forM_ ([Param LParamMem] -> [VName] -> [(Param LParamMem, VName)]
forall a b. [a] -> [b] -> [(a, b)]
zip (SegBinOp MCMem -> [LParam MCMem]
xParams SegBinOp MCMem
scan_op) [VName]
local_accums) (((Param LParamMem, VName) -> MulticoreGen ()) -> MulticoreGen ())
-> ((Param LParamMem, VName) -> MulticoreGen ()) -> MulticoreGen ()
forall a b. (a -> b) -> a -> b
$ \(Param LParamMem
p, VName
acc) ->
            VName -> [TExp Int64] -> SubExp -> [TExp Int64] -> MulticoreGen ()
forall rep r op.
VName -> [TExp Int64] -> SubExp -> [TExp Int64] -> ImpM rep r op ()
copyDWIMFix (Param LParamMem -> VName
forall dec. Param dec -> VName
paramName Param LParamMem
p) [] (VName -> SubExp
Var VName
acc) [TExp Int64]
vec_is
          [(Param LParamMem, KernelResult)]
-> ((Param LParamMem, KernelResult) -> MulticoreGen ())
-> MulticoreGen ()
forall (t :: * -> *) (m :: * -> *) a b.
(Foldable t, Monad m) =>
t a -> (a -> m b) -> m ()
forM_ ([Param LParamMem]
-> [KernelResult] -> [(Param LParamMem, KernelResult)]
forall a b. [a] -> [b] -> [(a, b)]
zip (SegBinOp MCMem -> [LParam MCMem]
yParams SegBinOp MCMem
scan_op) [KernelResult]
scan_res) (((Param LParamMem, KernelResult) -> MulticoreGen ())
 -> MulticoreGen ())
-> ((Param LParamMem, KernelResult) -> MulticoreGen ())
-> MulticoreGen ()
forall a b. (a -> b) -> a -> b
$ \(Param LParamMem
py, KernelResult
kr) ->
            VName -> [TExp Int64] -> SubExp -> [TExp Int64] -> MulticoreGen ()
forall rep r op.
VName -> [TExp Int64] -> SubExp -> [TExp Int64] -> ImpM rep r op ()
copyDWIMFix (Param LParamMem -> VName
forall dec. Param dec -> VName
paramName Param LParamMem
py) [] (KernelResult -> SubExp
kernelResultSubExp KernelResult
kr) [TExp Int64]
vec_is
          Names -> Seq (Stm MCMem) -> MulticoreGen () -> MulticoreGen ()
forall rep r op.
Names -> Stms rep -> ImpM rep r op () -> ImpM rep r op ()
compileStms Names
forall a. Monoid a => a
mempty (Body MCMem -> Seq (Stm MCMem)
forall rep res. GBody rep res -> Stms rep
bodyStms (Body MCMem -> Seq (Stm MCMem)) -> Body MCMem -> Seq (Stm MCMem)
forall a b. (a -> b) -> a -> b
$ SegBinOp MCMem -> Body MCMem
lamBody SegBinOp MCMem
scan_op) (MulticoreGen () -> MulticoreGen ())
-> MulticoreGen () -> MulticoreGen ()
forall a b. (a -> b) -> a -> b
$
            [(VName, SubExp, VName)]
-> ((VName, SubExp, VName) -> MulticoreGen ()) -> MulticoreGen ()
forall (t :: * -> *) (m :: * -> *) a b.
(Foldable t, Monad m) =>
t a -> (a -> m b) -> m ()
forM_ ([VName] -> [SubExp] -> [VName] -> [(VName, SubExp, VName)]
forall a b c. [a] -> [b] -> [c] -> [(a, b, c)]
zip3 [VName]
local_accums ((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 MCMem -> [SubExpRes]
forall rep res. GBody rep res -> [res]
bodyResult (Body MCMem -> [SubExpRes]) -> Body MCMem -> [SubExpRes]
forall a b. (a -> b) -> a -> b
$ SegBinOp MCMem -> Body MCMem
lamBody SegBinOp MCMem
scan_op) [VName]
pes) (((VName, SubExp, VName) -> MulticoreGen ()) -> MulticoreGen ())
-> ((VName, SubExp, VName) -> MulticoreGen ()) -> MulticoreGen ()
forall a b. (a -> b) -> a -> b
$ \(VName
acc, SubExp
se, VName
pe) -> do
              VName -> [TExp Int64] -> SubExp -> [TExp Int64] -> MulticoreGen ()
forall rep r op.
VName -> [TExp Int64] -> SubExp -> [TExp Int64] -> ImpM rep r op ()
copyDWIMFix VName
pe (TV Int64 -> TExp Int64
forall {k} (t :: k). TV t -> TExp t
tvExp TV Int64
task_id TExp Int64 -> [TExp Int64] -> [TExp Int64]
forall a. a -> [a] -> [a]
: TV Int64 -> TExp Int64
forall {k} (t :: k). TV t -> TExp t
tvExp TV Int64
z TExp Int64 -> [TExp Int64] -> [TExp Int64]
forall a. a -> [a] -> [a]
: [TExp Int64]
vec_is) SubExp
se []
              VName -> [TExp Int64] -> SubExp -> [TExp Int64] -> MulticoreGen ()
forall rep r op.
VName -> [TExp Int64] -> SubExp -> [TExp Int64] -> ImpM rep r op ()
copyDWIMFix VName
acc [TExp Int64]
vec_is SubExp
se []
    z <-- tvExp z + 1

    -- write back local accumulators to prefix arrays
    forM_ (zip3 scan_ops_renamed per_op_prefix_arr per_op_local_accum) $ \(SegBinOp MCMem
scan_op, [VName]
prefix_arrs, [VName]
local_accums) ->
      ShapeBase SubExp
-> ([TExp Int64] -> MulticoreGen ()) -> MulticoreGen ()
forall rep r op.
ShapeBase SubExp
-> ([TExp Int64] -> ImpM rep r op ()) -> ImpM rep r op ()
sLoopNest (SegBinOp MCMem -> ShapeBase SubExp
forall rep. SegBinOp rep -> ShapeBase SubExp
segBinOpShape SegBinOp MCMem
scan_op) (([TExp Int64] -> MulticoreGen ()) -> MulticoreGen ())
-> ([TExp Int64] -> MulticoreGen ()) -> MulticoreGen ()
forall a b. (a -> b) -> a -> b
$ \[TExp Int64]
vec_is ->
        [(VName, VName)]
-> ((VName, VName) -> MulticoreGen ()) -> MulticoreGen ()
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]
local_accums [VName]
prefix_arrs) (((VName, VName) -> MulticoreGen ()) -> MulticoreGen ())
-> ((VName, VName) -> MulticoreGen ()) -> MulticoreGen ()
forall a b. (a -> b) -> a -> b
$ \(VName
acc, VName
prefix_arr) ->
          VName -> [TExp Int64] -> SubExp -> [TExp Int64] -> MulticoreGen ()
forall rep r op.
VName -> [TExp Int64] -> SubExp -> [TExp Int64] -> ImpM rep r op ()
copyDWIMFix VName
prefix_arr (TV Int64 -> TExp Int64
forall {k} (t :: k). TV t -> TExp t
tvExp TV Int64
block_idx TExp Int64 -> [TExp Int64] -> [TExp Int64]
forall a. a -> [a] -> [a]
: [TExp Int64]
vec_is) (VName -> SubExp
Var VName
acc) [TExp Int64]
vec_is

seqScanLB ::
  [[VName]] ->
  VName ->
  [SegBinOp MCMem] ->
  KernelBody MCMem ->
  [[VName]] ->
  TV Int64 ->
  TV Int64 ->
  TV Int64 ->
  MulticoreGen ()
seqScanLB :: [[VName]]
-> VName
-> [SegBinOp MCMem]
-> KernelBody MCMem
-> [[VName]]
-> TV Int64
-> TV Int64
-> TV Int64
-> MulticoreGen ()
seqScanLB [[VName]]
scan_out VName
i [SegBinOp MCMem]
scan_ops KernelBody MCMem
kbody [[VName]]
per_op_prefixes_var TV Int64
start TV Int64
chunk_length TV Int64
task_id = do
  kbody_renamed <- KernelBody MCMem -> ImpM MCMem HostEnv Multicore (KernelBody MCMem)
forall rep res (m :: * -> *).
(Renameable rep, Rename res, MonadFreshNames m) =>
GBody rep res -> m (GBody rep res)
renameBody KernelBody MCMem
kbody
  scan_ops_renamed <- renameSegBinOp scan_ops
  genBinOpParams scan_ops_renamed
  per_op_local_accum <- genLocalArray scan_ops_renamed

  let results = KernelBody MCMem -> [KernelResult]
forall rep res. GBody rep res -> [res]
bodyResult KernelBody MCMem
kbody_renamed
  let n_scan = [SegBinOp MCMem] -> Int
forall rep. [SegBinOp rep] -> Int
segBinOpResults [SegBinOp MCMem]
scan_ops_renamed
  let (all_scan_res, _) = splitAt n_scan results

  let per_scan_res = [SegBinOp MCMem] -> [KernelResult] -> [[KernelResult]]
forall rep a. [SegBinOp rep] -> [a] -> [[a]]
segBinOpChunks [SegBinOp MCMem]
scan_ops_renamed [KernelResult]
all_scan_res

  forM_ (zip3 scan_ops_renamed per_op_prefixes_var per_op_local_accum) $ \(SegBinOp MCMem
scan_op, [VName]
prefix_vars, [VName]
local_accums) ->
    ShapeBase SubExp
-> ([TExp Int64] -> MulticoreGen ()) -> MulticoreGen ()
forall rep r op.
ShapeBase SubExp
-> ([TExp Int64] -> ImpM rep r op ()) -> ImpM rep r op ()
sLoopNest (SegBinOp MCMem -> ShapeBase SubExp
forall rep. SegBinOp rep -> ShapeBase SubExp
segBinOpShape SegBinOp MCMem
scan_op) (([TExp Int64] -> MulticoreGen ()) -> MulticoreGen ())
-> ([TExp Int64] -> MulticoreGen ()) -> MulticoreGen ()
forall a b. (a -> b) -> a -> b
$ \[TExp Int64]
vec_is ->
      [(VName, VName)]
-> ((VName, VName) -> MulticoreGen ()) -> MulticoreGen ()
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]
local_accums [VName]
prefix_vars) (((VName, VName) -> MulticoreGen ()) -> MulticoreGen ())
-> ((VName, VName) -> MulticoreGen ()) -> MulticoreGen ()
forall a b. (a -> b) -> a -> b
$ \(VName
acc, VName
prefix) ->
        VName -> [TExp Int64] -> SubExp -> [TExp Int64] -> MulticoreGen ()
forall rep r op.
VName -> [TExp Int64] -> SubExp -> [TExp Int64] -> ImpM rep r op ()
copyDWIMFix VName
acc [TExp Int64]
vec_is (VName -> SubExp
Var VName
prefix) [TExp Int64]
vec_is

  z <- dPrimV "z" (0 :: Imp.TExp Int64)
  sWhile (tvExp z .<. tvExp chunk_length) $ do
    dPrimV_ i (tvExp start + tvExp z)
    if shouldRecompute kbody_renamed
      then do
        compileStms mempty (bodyStms kbody_renamed) $ do
          forM_ (zip4 scan_out scan_ops_renamed per_scan_res per_op_local_accum) $ \([VName]
pes, SegBinOp MCMem
scan_op, [KernelResult]
scan_res, [VName]
local_accums) ->
            ShapeBase SubExp
-> ([TExp Int64] -> MulticoreGen ()) -> MulticoreGen ()
forall rep r op.
ShapeBase SubExp
-> ([TExp Int64] -> ImpM rep r op ()) -> ImpM rep r op ()
sLoopNest (SegBinOp MCMem -> ShapeBase SubExp
forall rep. SegBinOp rep -> ShapeBase SubExp
segBinOpShape SegBinOp MCMem
scan_op) (([TExp Int64] -> MulticoreGen ()) -> MulticoreGen ())
-> ([TExp Int64] -> MulticoreGen ()) -> MulticoreGen ()
forall a b. (a -> b) -> a -> b
$ \[TExp Int64]
vec_is -> do
              [(Param LParamMem, VName)]
-> ((Param LParamMem, VName) -> MulticoreGen ()) -> MulticoreGen ()
forall (t :: * -> *) (m :: * -> *) a b.
(Foldable t, Monad m) =>
t a -> (a -> m b) -> m ()
forM_ ([Param LParamMem] -> [VName] -> [(Param LParamMem, VName)]
forall a b. [a] -> [b] -> [(a, b)]
zip (SegBinOp MCMem -> [LParam MCMem]
xParams SegBinOp MCMem
scan_op) [VName]
local_accums) (((Param LParamMem, VName) -> MulticoreGen ()) -> MulticoreGen ())
-> ((Param LParamMem, VName) -> MulticoreGen ()) -> MulticoreGen ()
forall a b. (a -> b) -> a -> b
$ \(Param LParamMem
px, VName
acc) ->
                VName -> [TExp Int64] -> SubExp -> [TExp Int64] -> MulticoreGen ()
forall rep r op.
VName -> [TExp Int64] -> SubExp -> [TExp Int64] -> ImpM rep r op ()
copyDWIMFix (Param LParamMem -> VName
forall dec. Param dec -> VName
paramName Param LParamMem
px) [] (VName -> SubExp
Var VName
acc) [TExp Int64]
vec_is
              [(Param LParamMem, KernelResult)]
-> ((Param LParamMem, KernelResult) -> MulticoreGen ())
-> MulticoreGen ()
forall (t :: * -> *) (m :: * -> *) a b.
(Foldable t, Monad m) =>
t a -> (a -> m b) -> m ()
forM_ ([Param LParamMem]
-> [KernelResult] -> [(Param LParamMem, KernelResult)]
forall a b. [a] -> [b] -> [(a, b)]
zip (SegBinOp MCMem -> [LParam MCMem]
yParams SegBinOp MCMem
scan_op) [KernelResult]
scan_res) (((Param LParamMem, KernelResult) -> MulticoreGen ())
 -> MulticoreGen ())
-> ((Param LParamMem, KernelResult) -> MulticoreGen ())
-> MulticoreGen ()
forall a b. (a -> b) -> a -> b
$ \(Param LParamMem
py, KernelResult
kr) ->
                VName -> [TExp Int64] -> SubExp -> [TExp Int64] -> MulticoreGen ()
forall rep r op.
VName -> [TExp Int64] -> SubExp -> [TExp Int64] -> ImpM rep r op ()
copyDWIMFix (Param LParamMem -> VName
forall dec. Param dec -> VName
paramName Param LParamMem
py) [] (KernelResult -> SubExp
kernelResultSubExp KernelResult
kr) [TExp Int64]
vec_is
              Names -> Seq (Stm MCMem) -> MulticoreGen () -> MulticoreGen ()
forall rep r op.
Names -> Stms rep -> ImpM rep r op () -> ImpM rep r op ()
compileStms Names
forall a. Monoid a => a
mempty (Body MCMem -> Seq (Stm MCMem)
forall rep res. GBody rep res -> Stms rep
bodyStms (Body MCMem -> Seq (Stm MCMem)) -> Body MCMem -> Seq (Stm MCMem)
forall a b. (a -> b) -> a -> b
$ SegBinOp MCMem -> Body MCMem
lamBody SegBinOp MCMem
scan_op) (MulticoreGen () -> MulticoreGen ())
-> MulticoreGen () -> MulticoreGen ()
forall a b. (a -> b) -> a -> b
$
                [(SubExp, VName, VName)]
-> ((SubExp, VName, VName) -> MulticoreGen ()) -> MulticoreGen ()
forall (t :: * -> *) (m :: * -> *) a b.
(Foldable t, Monad m) =>
t a -> (a -> m b) -> m ()
forM_ ([SubExp] -> [VName] -> [VName] -> [(SubExp, VName, VName)]
forall a b c. [a] -> [b] -> [c] -> [(a, b, c)]
zip3 ((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 MCMem -> [SubExpRes]
forall rep res. GBody rep res -> [res]
bodyResult (Body MCMem -> [SubExpRes]) -> Body MCMem -> [SubExpRes]
forall a b. (a -> b) -> a -> b
$ SegBinOp MCMem -> Body MCMem
lamBody SegBinOp MCMem
scan_op) [VName]
pes [VName]
local_accums) (((SubExp, VName, VName) -> MulticoreGen ()) -> MulticoreGen ())
-> ((SubExp, VName, VName) -> MulticoreGen ()) -> MulticoreGen ()
forall a b. (a -> b) -> a -> b
$ \(SubExp
se, VName
pe, VName
acc) -> do
                  VName -> [TExp Int64] -> SubExp -> [TExp Int64] -> MulticoreGen ()
forall rep r op.
VName -> [TExp Int64] -> SubExp -> [TExp Int64] -> ImpM rep r op ()
copyDWIMFix VName
acc [TExp Int64]
vec_is SubExp
se []
                  VName -> [TExp Int64] -> SubExp -> [TExp Int64] -> MulticoreGen ()
forall rep r op.
VName -> [TExp Int64] -> SubExp -> [TExp Int64] -> ImpM rep r op ()
copyDWIMFix VName
pe (TV Int64 -> TExp Int64
forall {k} (t :: k). TV t -> TExp t
tvExp TV Int64
task_id TExp Int64 -> [TExp Int64] -> [TExp Int64]
forall a. a -> [a] -> [a]
: TV Int64 -> TExp Int64
forall {k} (t :: k). TV t -> TExp t
tvExp TV Int64
z TExp Int64 -> [TExp Int64] -> [TExp Int64]
forall a. a -> [a] -> [a]
: [TExp Int64]
vec_is) SubExp
se []
        z <-- tvExp z + 1
      else do
        forM_ (zip4 scan_out scan_ops_renamed per_scan_res per_op_local_accum) $ \([VName]
pes, SegBinOp MCMem
scan_op, [KernelResult]
_, [VName]
local_accums) ->
          ShapeBase SubExp
-> ([TExp Int64] -> MulticoreGen ()) -> MulticoreGen ()
forall rep r op.
ShapeBase SubExp
-> ([TExp Int64] -> ImpM rep r op ()) -> ImpM rep r op ()
sLoopNest (SegBinOp MCMem -> ShapeBase SubExp
forall rep. SegBinOp rep -> ShapeBase SubExp
segBinOpShape SegBinOp MCMem
scan_op) (([TExp Int64] -> MulticoreGen ()) -> MulticoreGen ())
-> ([TExp Int64] -> MulticoreGen ()) -> MulticoreGen ()
forall a b. (a -> b) -> a -> b
$ \[TExp Int64]
vec_is -> do
            [(Param LParamMem, VName)]
-> ((Param LParamMem, VName) -> MulticoreGen ()) -> MulticoreGen ()
forall (t :: * -> *) (m :: * -> *) a b.
(Foldable t, Monad m) =>
t a -> (a -> m b) -> m ()
forM_ ([Param LParamMem] -> [VName] -> [(Param LParamMem, VName)]
forall a b. [a] -> [b] -> [(a, b)]
zip (SegBinOp MCMem -> [LParam MCMem]
xParams SegBinOp MCMem
scan_op) [VName]
local_accums) (((Param LParamMem, VName) -> MulticoreGen ()) -> MulticoreGen ())
-> ((Param LParamMem, VName) -> MulticoreGen ()) -> MulticoreGen ()
forall a b. (a -> b) -> a -> b
$ \(Param LParamMem
px, VName
acc) ->
              VName -> [TExp Int64] -> SubExp -> [TExp Int64] -> MulticoreGen ()
forall rep r op.
VName -> [TExp Int64] -> SubExp -> [TExp Int64] -> ImpM rep r op ()
copyDWIMFix (Param LParamMem -> VName
forall dec. Param dec -> VName
paramName Param LParamMem
px) [] (VName -> SubExp
Var VName
acc) [TExp Int64]
vec_is
            [(Param LParamMem, VName)]
-> ((Param LParamMem, VName) -> MulticoreGen ()) -> MulticoreGen ()
forall (t :: * -> *) (m :: * -> *) a b.
(Foldable t, Monad m) =>
t a -> (a -> m b) -> m ()
forM_ ([Param LParamMem] -> [VName] -> [(Param LParamMem, VName)]
forall a b. [a] -> [b] -> [(a, b)]
zip (SegBinOp MCMem -> [LParam MCMem]
yParams SegBinOp MCMem
scan_op) [VName]
pes) (((Param LParamMem, VName) -> MulticoreGen ()) -> MulticoreGen ())
-> ((Param LParamMem, VName) -> MulticoreGen ()) -> MulticoreGen ()
forall a b. (a -> b) -> a -> b
$ \(Param LParamMem
py, VName
pe) ->
              VName -> [TExp Int64] -> SubExp -> [TExp Int64] -> MulticoreGen ()
forall rep r op.
VName -> [TExp Int64] -> SubExp -> [TExp Int64] -> ImpM rep r op ()
copyDWIMFix (Param LParamMem -> VName
forall dec. Param dec -> VName
paramName Param LParamMem
py) [] (VName -> SubExp
Var VName
pe) (TV Int64 -> TExp Int64
forall {k} (t :: k). TV t -> TExp t
tvExp TV Int64
task_id TExp Int64 -> [TExp Int64] -> [TExp Int64]
forall a. a -> [a] -> [a]
: TV Int64 -> TExp Int64
forall {k} (t :: k). TV t -> TExp t
tvExp TV Int64
z TExp Int64 -> [TExp Int64] -> [TExp Int64]
forall a. a -> [a] -> [a]
: [TExp Int64]
vec_is)
            Names -> Seq (Stm MCMem) -> MulticoreGen () -> MulticoreGen ()
forall rep r op.
Names -> Stms rep -> ImpM rep r op () -> ImpM rep r op ()
compileStms Names
forall a. Monoid a => a
mempty (Body MCMem -> Seq (Stm MCMem)
forall rep res. GBody rep res -> Stms rep
bodyStms (Body MCMem -> Seq (Stm MCMem)) -> Body MCMem -> Seq (Stm MCMem)
forall a b. (a -> b) -> a -> b
$ SegBinOp MCMem -> Body MCMem
lamBody SegBinOp MCMem
scan_op) (MulticoreGen () -> MulticoreGen ())
-> MulticoreGen () -> MulticoreGen ()
forall a b. (a -> b) -> a -> b
$ do
              [(SubExp, VName)]
-> ((SubExp, VName) -> MulticoreGen ()) -> MulticoreGen ()
forall (t :: * -> *) (m :: * -> *) a b.
(Foldable t, Monad m) =>
t a -> (a -> m b) -> m ()
forM_ ([SubExp] -> [VName] -> [(SubExp, VName)]
forall a b. [a] -> [b] -> [(a, b)]
zip ((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 MCMem -> [SubExpRes]
forall rep res. GBody rep res -> [res]
bodyResult (Body MCMem -> [SubExpRes]) -> Body MCMem -> [SubExpRes]
forall a b. (a -> b) -> a -> b
$ SegBinOp MCMem -> Body MCMem
lamBody SegBinOp MCMem
scan_op) [VName]
pes) (((SubExp, VName) -> MulticoreGen ()) -> MulticoreGen ())
-> ((SubExp, VName) -> MulticoreGen ()) -> MulticoreGen ()
forall a b. (a -> b) -> a -> b
$ \(SubExp
se, VName
pe) ->
                VName -> [TExp Int64] -> SubExp -> [TExp Int64] -> MulticoreGen ()
forall rep r op.
VName -> [TExp Int64] -> SubExp -> [TExp Int64] -> ImpM rep r op ()
copyDWIMFix VName
pe (TV Int64 -> TExp Int64
forall {k} (t :: k). TV t -> TExp t
tvExp TV Int64
task_id TExp Int64 -> [TExp Int64] -> [TExp Int64]
forall a. a -> [a] -> [a]
: TV Int64 -> TExp Int64
forall {k} (t :: k). TV t -> TExp t
tvExp TV Int64
z TExp Int64 -> [TExp Int64] -> [TExp Int64]
forall a. a -> [a] -> [a]
: [TExp Int64]
vec_is) SubExp
se []
        z <-- tvExp z + 1

seqAggregate ::
  [[VName]] ->
  [Maybe VName] ->
  VName ->
  [SegBinOp MCMem] ->
  KernelBody MCMem ->
  TV Int64 ->
  TV Int64 ->
  [[VName]] ->
  TV Int64 ->
  TV Int64 ->
  MulticoreGen ()
seqAggregate :: [[VName]]
-> [Maybe VName]
-> VName
-> [SegBinOp MCMem]
-> KernelBody MCMem
-> TV Int64
-> TV Int64
-> [[VName]]
-> TV Int64
-> TV Int64
-> MulticoreGen ()
seqAggregate [[VName]]
scan_out [Maybe VName]
map_out VName
i [SegBinOp MCMem]
scan_ops KernelBody MCMem
kbody TV Int64
start TV Int64
chunk_length [[VName]]
per_op_aggr_arrs TV Int64
block_idx TV Int64
task_id = do
  scan_ops_renamed <- [SegBinOp MCMem] -> MulticoreGen [SegBinOp MCMem]
renameSegBinOp [SegBinOp MCMem]
scan_ops
  kbody_renamed <- renameBody kbody
  genBinOpParams scan_ops_renamed
  let results = KernelBody MCMem -> [KernelResult]
forall rep res. GBody rep res -> [res]
bodyResult KernelBody MCMem
kbody_renamed
  let n_scan = [SegBinOp MCMem] -> Int
forall rep. [SegBinOp rep] -> Int
segBinOpResults [SegBinOp MCMem]
scan_ops_renamed
  let (all_scan_res, map_res) = splitAt n_scan results
  let per_scan_res = [SegBinOp MCMem] -> [KernelResult] -> [[KernelResult]]
forall rep a. [SegBinOp rep] -> [a] -> [[a]]
segBinOpChunks [SegBinOp MCMem]
scan_ops_renamed [KernelResult]
all_scan_res

  per_op_local_accum <- genLocalArray scan_ops_renamed

  j <- dPrimV "j" (0 :: Imp.TExp Int64)
  sWhile (tvExp j .<. tvExp chunk_length) $ do
    dPrimV_ i (tvExp start + tvExp j)
    compileStms mempty (bodyStms kbody_renamed) $ do
      sComment "write mapped values results to memory" $
        forM_ (zip map_out (map kernelResultSubExp map_res)) $ \(Maybe VName
marr, SubExp
res) ->
          Maybe VName -> (VName -> MulticoreGen ()) -> MulticoreGen ()
forall (t :: * -> *) (m :: * -> *) a b.
(Foldable t, Monad m) =>
t a -> (a -> m b) -> m ()
forM_ Maybe VName
marr ((VName -> MulticoreGen ()) -> MulticoreGen ())
-> (VName -> MulticoreGen ()) -> MulticoreGen ()
forall a b. (a -> b) -> a -> b
$ \VName
arr -> VName -> [TExp Int64] -> SubExp -> [TExp Int64] -> MulticoreGen ()
forall rep r op.
VName -> [TExp Int64] -> SubExp -> [TExp Int64] -> ImpM rep r op ()
copyDWIMFix VName
arr [TV Int64 -> TExp Int64
forall {k} (t :: k). TV t -> TExp t
tvExp TV Int64
task_id, TV Int64 -> TExp Int64
forall {k} (t :: k). TV t -> TExp t
tvExp TV Int64
j] SubExp
res []

      sIf
        (tvExp j .==. 0)
        ( forM_ (zip4 scan_ops_renamed per_scan_res per_op_local_accum scan_out) $
            \(SegBinOp MCMem
scan_op, [KernelResult]
scan_res_op, [VName]
local_accums, [VName]
pes) -> do
              let shape :: ShapeBase SubExp
shape = SegBinOp MCMem -> ShapeBase SubExp
forall rep. SegBinOp rep -> ShapeBase SubExp
segBinOpShape SegBinOp MCMem
scan_op
              ShapeBase SubExp
-> ([TExp Int64] -> MulticoreGen ()) -> MulticoreGen ()
forall rep r op.
ShapeBase SubExp
-> ([TExp Int64] -> ImpM rep r op ()) -> ImpM rep r op ()
sLoopNest ShapeBase SubExp
shape (([TExp Int64] -> MulticoreGen ()) -> MulticoreGen ())
-> ([TExp Int64] -> MulticoreGen ()) -> MulticoreGen ()
forall a b. (a -> b) -> a -> b
$ \[TExp Int64]
vec_is -> do
                [(KernelResult, VName, VName)]
-> ((KernelResult, VName, VName) -> MulticoreGen ())
-> MulticoreGen ()
forall (t :: * -> *) (m :: * -> *) a b.
(Foldable t, Monad m) =>
t a -> (a -> m b) -> m ()
forM_ ([KernelResult]
-> [VName] -> [VName] -> [(KernelResult, VName, VName)]
forall a b c. [a] -> [b] -> [c] -> [(a, b, c)]
zip3 [KernelResult]
scan_res_op [VName]
local_accums [VName]
pes) (((KernelResult, VName, VName) -> MulticoreGen ())
 -> MulticoreGen ())
-> ((KernelResult, VName, VName) -> MulticoreGen ())
-> MulticoreGen ()
forall a b. (a -> b) -> a -> b
$ \(KernelResult
kr, VName
acc, VName
pe) -> do
                  VName -> [TExp Int64] -> SubExp -> [TExp Int64] -> MulticoreGen ()
forall rep r op.
VName -> [TExp Int64] -> SubExp -> [TExp Int64] -> ImpM rep r op ()
copyDWIMFix VName
acc [TExp Int64]
vec_is (KernelResult -> SubExp
kernelResultSubExp KernelResult
kr) [TExp Int64]
vec_is
                  Bool -> MulticoreGen () -> MulticoreGen ()
forall (f :: * -> *). Applicative f => Bool -> f () -> f ()
unless (KernelBody MCMem -> Bool
shouldRecompute KernelBody MCMem
kbody_renamed) (MulticoreGen () -> MulticoreGen ())
-> MulticoreGen () -> MulticoreGen ()
forall a b. (a -> b) -> a -> b
$
                    VName -> [TExp Int64] -> SubExp -> [TExp Int64] -> MulticoreGen ()
forall rep r op.
VName -> [TExp Int64] -> SubExp -> [TExp Int64] -> ImpM rep r op ()
copyDWIMFix VName
pe (TV Int64 -> TExp Int64
forall {k} (t :: k). TV t -> TExp t
tvExp TV Int64
task_id TExp Int64 -> [TExp Int64] -> [TExp Int64]
forall a. a -> [a] -> [a]
: TV Int64 -> TExp Int64
forall {k} (t :: k). TV t -> TExp t
tvExp TV Int64
j TExp Int64 -> [TExp Int64] -> [TExp Int64]
forall a. a -> [a] -> [a]
: [TExp Int64]
vec_is) (KernelResult -> SubExp
kernelResultSubExp KernelResult
kr) [TExp Int64]
vec_is
        )
        ( forM_ (zip4 scan_ops_renamed per_scan_res per_op_local_accum scan_out) $
            \(SegBinOp MCMem
scan_op, [KernelResult]
scan_res, [VName]
local_accums, [VName]
pes) ->
              ShapeBase SubExp
-> ([TExp Int64] -> MulticoreGen ()) -> MulticoreGen ()
forall rep r op.
ShapeBase SubExp
-> ([TExp Int64] -> ImpM rep r op ()) -> ImpM rep r op ()
sLoopNest (SegBinOp MCMem -> ShapeBase SubExp
forall rep. SegBinOp rep -> ShapeBase SubExp
segBinOpShape SegBinOp MCMem
scan_op) (([TExp Int64] -> MulticoreGen ()) -> MulticoreGen ())
-> ([TExp Int64] -> MulticoreGen ()) -> MulticoreGen ()
forall a b. (a -> b) -> a -> b
$ \[TExp Int64]
vec_is -> do
                [(Param LParamMem, VName)]
-> ((Param LParamMem, VName) -> MulticoreGen ()) -> MulticoreGen ()
forall (t :: * -> *) (m :: * -> *) a b.
(Foldable t, Monad m) =>
t a -> (a -> m b) -> m ()
forM_ ([Param LParamMem] -> [VName] -> [(Param LParamMem, VName)]
forall a b. [a] -> [b] -> [(a, b)]
zip (SegBinOp MCMem -> [LParam MCMem]
xParams SegBinOp MCMem
scan_op) [VName]
local_accums) (((Param LParamMem, VName) -> MulticoreGen ()) -> MulticoreGen ())
-> ((Param LParamMem, VName) -> MulticoreGen ()) -> MulticoreGen ()
forall a b. (a -> b) -> a -> b
$ \(Param LParamMem
px, VName
acc) ->
                  VName -> [TExp Int64] -> SubExp -> [TExp Int64] -> MulticoreGen ()
forall rep r op.
VName -> [TExp Int64] -> SubExp -> [TExp Int64] -> ImpM rep r op ()
copyDWIMFix (Param LParamMem -> VName
forall dec. Param dec -> VName
paramName Param LParamMem
px) [] (VName -> SubExp
Var VName
acc) [TExp Int64]
vec_is
                [(Param LParamMem, KernelResult)]
-> ((Param LParamMem, KernelResult) -> MulticoreGen ())
-> MulticoreGen ()
forall (t :: * -> *) (m :: * -> *) a b.
(Foldable t, Monad m) =>
t a -> (a -> m b) -> m ()
forM_ ([Param LParamMem]
-> [KernelResult] -> [(Param LParamMem, KernelResult)]
forall a b. [a] -> [b] -> [(a, b)]
zip (SegBinOp MCMem -> [LParam MCMem]
yParams SegBinOp MCMem
scan_op) [KernelResult]
scan_res) (((Param LParamMem, KernelResult) -> MulticoreGen ())
 -> MulticoreGen ())
-> ((Param LParamMem, KernelResult) -> MulticoreGen ())
-> MulticoreGen ()
forall a b. (a -> b) -> a -> b
$ \(Param LParamMem
py, KernelResult
kr) ->
                  VName -> [TExp Int64] -> SubExp -> [TExp Int64] -> MulticoreGen ()
forall rep r op.
VName -> [TExp Int64] -> SubExp -> [TExp Int64] -> ImpM rep r op ()
copyDWIMFix (Param LParamMem -> VName
forall dec. Param dec -> VName
paramName Param LParamMem
py) [] (KernelResult -> SubExp
kernelResultSubExp KernelResult
kr) [TExp Int64]
vec_is

                Names -> Seq (Stm MCMem) -> MulticoreGen () -> MulticoreGen ()
forall rep r op.
Names -> Stms rep -> ImpM rep r op () -> ImpM rep r op ()
compileStms Names
forall a. Monoid a => a
mempty (Body MCMem -> Seq (Stm MCMem)
forall rep res. GBody rep res -> Stms rep
bodyStms (Body MCMem -> Seq (Stm MCMem)) -> Body MCMem -> Seq (Stm MCMem)
forall a b. (a -> b) -> a -> b
$ SegBinOp MCMem -> Body MCMem
lamBody SegBinOp MCMem
scan_op) (MulticoreGen () -> MulticoreGen ())
-> MulticoreGen () -> MulticoreGen ()
forall a b. (a -> b) -> a -> b
$ do
                  [(SubExp, VName, VName)]
-> ((SubExp, VName, VName) -> MulticoreGen ()) -> MulticoreGen ()
forall (t :: * -> *) (m :: * -> *) a b.
(Foldable t, Monad m) =>
t a -> (a -> m b) -> m ()
forM_ ([SubExp] -> [VName] -> [VName] -> [(SubExp, VName, VName)]
forall a b c. [a] -> [b] -> [c] -> [(a, b, c)]
zip3 ((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 MCMem -> [SubExpRes]
forall rep res. GBody rep res -> [res]
bodyResult (Body MCMem -> [SubExpRes]) -> Body MCMem -> [SubExpRes]
forall a b. (a -> b) -> a -> b
$ SegBinOp MCMem -> Body MCMem
lamBody SegBinOp MCMem
scan_op) [VName]
local_accums [VName]
pes) (((SubExp, VName, VName) -> MulticoreGen ()) -> MulticoreGen ())
-> ((SubExp, VName, VName) -> MulticoreGen ()) -> MulticoreGen ()
forall a b. (a -> b) -> a -> b
$ \(SubExp
se, VName
acc, VName
pe) -> do
                    VName -> [TExp Int64] -> SubExp -> [TExp Int64] -> MulticoreGen ()
forall rep r op.
VName -> [TExp Int64] -> SubExp -> [TExp Int64] -> ImpM rep r op ()
copyDWIMFix VName
acc [TExp Int64]
vec_is SubExp
se []
                    Bool -> MulticoreGen () -> MulticoreGen ()
forall (f :: * -> *). Applicative f => Bool -> f () -> f ()
unless (KernelBody MCMem -> Bool
shouldRecompute KernelBody MCMem
kbody_renamed) (MulticoreGen () -> MulticoreGen ())
-> MulticoreGen () -> MulticoreGen ()
forall a b. (a -> b) -> a -> b
$ VName -> [TExp Int64] -> SubExp -> [TExp Int64] -> MulticoreGen ()
forall rep r op.
VName -> [TExp Int64] -> SubExp -> [TExp Int64] -> ImpM rep r op ()
copyDWIMFix VName
pe (TV Int64 -> TExp Int64
forall {k} (t :: k). TV t -> TExp t
tvExp TV Int64
task_id TExp Int64 -> [TExp Int64] -> [TExp Int64]
forall a. a -> [a] -> [a]
: TV Int64 -> TExp Int64
forall {k} (t :: k). TV t -> TExp t
tvExp TV Int64
j TExp Int64 -> [TExp Int64] -> [TExp Int64]
forall a. a -> [a] -> [a]
: [TExp Int64]
vec_is) SubExp
se []
        )
    j <-- tvExp j + 1

  --  write local accumulators to aggregate arrays
  forM_ (zip3 scan_ops_renamed per_op_local_accum per_op_aggr_arrs) $ \(SegBinOp MCMem
scan_op, [VName]
local_accums, [VName]
aggr_arrs) ->
    ShapeBase SubExp
-> ([TExp Int64] -> MulticoreGen ()) -> MulticoreGen ()
forall rep r op.
ShapeBase SubExp
-> ([TExp Int64] -> ImpM rep r op ()) -> ImpM rep r op ()
sLoopNest (SegBinOp MCMem -> ShapeBase SubExp
forall rep. SegBinOp rep -> ShapeBase SubExp
segBinOpShape SegBinOp MCMem
scan_op) (([TExp Int64] -> MulticoreGen ()) -> MulticoreGen ())
-> ([TExp Int64] -> MulticoreGen ()) -> MulticoreGen ()
forall a b. (a -> b) -> a -> b
$ \[TExp Int64]
vec_is -> do
      [(VName, VName)]
-> ((VName, VName) -> MulticoreGen ()) -> MulticoreGen ()
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]
local_accums [VName]
aggr_arrs) (((VName, VName) -> MulticoreGen ()) -> MulticoreGen ())
-> ((VName, VName) -> MulticoreGen ()) -> MulticoreGen ()
forall a b. (a -> b) -> a -> b
$ \(VName
acc, VName
agg) -> do
        VName -> [TExp Int64] -> SubExp -> [TExp Int64] -> MulticoreGen ()
forall rep r op.
VName -> [TExp Int64] -> SubExp -> [TExp Int64] -> ImpM rep r op ()
copyDWIMFix VName
agg (TV Int64 -> TExp Int64
forall {k} (t :: k). TV t -> TExp t
tvExp TV Int64
block_idx TExp Int64 -> [TExp Int64] -> [TExp Int64]
forall a. a -> [a] -> [a]
: [TExp Int64]
vec_is) (VName -> SubExp
Var VName
acc) [TExp Int64]
vec_is

load64 ::
  VName ->
  VName ->
  Imp.Count Imp.Elements (Imp.TExp Int32) ->
  MulticoreGen ()
load64 :: VName -> VName -> Count Elements (TExp Int32) -> MulticoreGen ()
load64 VName
v VName
arr Count Elements (TExp Int32)
i = do
  space <- MemEntry -> Space
entryMemSpace (MemEntry -> Space)
-> ImpM MCMem HostEnv Multicore MemEntry
-> ImpM MCMem HostEnv Multicore Space
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> VName -> ImpM MCMem HostEnv Multicore MemEntry
forall rep r op. VName -> ImpM rep r op MemEntry
lookupMemory VName
arr
  sOp $ Imp.Atomic space $ Imp.AtomicLoad (IntType Int64) v arr i

store64 ::
  VName ->
  Imp.Count Imp.Elements (Imp.TExp Int32) ->
  Imp.TExp Int64 ->
  MulticoreGen ()
store64 :: VName
-> Count Elements (TExp Int32) -> TExp Int64 -> MulticoreGen ()
store64 VName
arr Count Elements (TExp Int32)
i TExp Int64
x = do
  space <- MemEntry -> Space
entryMemSpace (MemEntry -> Space)
-> ImpM MCMem HostEnv Multicore MemEntry
-> ImpM MCMem HostEnv Multicore Space
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> VName -> ImpM MCMem HostEnv Multicore MemEntry
forall rep r op. VName -> ImpM rep r op MemEntry
lookupMemory VName
arr
  sOp $ Imp.Atomic space $ Imp.AtomicStore (IntType Int64) arr i (untyped x)

add64 ::
  TV Int64 ->
  VName ->
  Imp.Count Imp.Elements (Imp.TExp Int32) ->
  Imp.TExp Int64 ->
  MulticoreGen ()
add64 :: TV Int64
-> VName
-> Count Elements (TExp Int32)
-> TExp Int64
-> MulticoreGen ()
add64 TV Int64
v VName
arr Count Elements (TExp Int32)
i TExp Int64
x = do
  space <- MemEntry -> Space
entryMemSpace (MemEntry -> Space)
-> ImpM MCMem HostEnv Multicore MemEntry
-> ImpM MCMem HostEnv Multicore Space
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> VName -> ImpM MCMem HostEnv Multicore MemEntry
forall rep r op. VName -> ImpM rep r op MemEntry
lookupMemory VName
arr
  sOp $ Imp.Atomic space $ Imp.AtomicAdd Int64 (tvVar v) arr i (untyped x)

applyPostOp ::
  Pat LetDecMem ->
  [[VName]] ->
  [Maybe VName] ->
  [SegBinOp MCMem] ->
  SegPostOp MCMem ->
  TV Int64 ->
  TV Int64 ->
  TV Int64 ->
  MulticoreGen ()
applyPostOp :: Pat LParamMem
-> [[VName]]
-> [Maybe VName]
-> [SegBinOp MCMem]
-> SegPostOp MCMem
-> TV Int64
-> TV Int64
-> TV Int64
-> MulticoreGen ()
applyPostOp Pat LParamMem
pat [[VName]]
scan_out [Maybe VName]
map_out [SegBinOp MCMem]
scan_ops SegPostOp MCMem
post_op TV Int64
start TV Int64
chunk_length TV Int64
task_id = do
  z <- Name -> TExp Int64 -> ImpM MCMem HostEnv Multicore (TV Int64)
forall {k} (t :: k) rep r op.
Name -> TExp t -> ImpM rep r op (TV t)
dPrimV Name
"z" (TExp Int64
0 :: Imp.TExp Int64)
  let (scan_pars, map_pars) = splitAt (segBinOpResults scan_ops) $ lambdaParams $ segPostOpLambda post_op
  dScope Nothing $
    scopeOfLParams $
      lambdaParams $
        segPostOpLambda post_op

  sWhile (tvExp z .<. tvExp chunk_length) $ do
    sComment "bind scan results to post lambda params" $ do
      forM_ (zip scan_pars $ mconcat scan_out) $ \(Param LParamMem
par, VName
acc) ->
        VName -> [TExp Int64] -> SubExp -> [TExp Int64] -> MulticoreGen ()
forall rep r op.
VName -> [TExp Int64] -> SubExp -> [TExp Int64] -> ImpM rep r op ()
copyDWIMFix (Param LParamMem -> VName
forall dec. Param dec -> VName
paramName Param LParamMem
par) [] (VName -> SubExp
Var VName
acc) [TV Int64 -> TExp Int64
forall {k} (t :: k). TV t -> TExp t
tvExp TV Int64
task_id, TV Int64 -> TExp Int64
forall {k} (t :: k). TV t -> TExp t
tvExp TV Int64
z]

    sComment "bind map results to post lamda params" $
      forM_ (zip map_pars map_out) $ \(Param LParamMem
par, Maybe VName
out) -> do
        Maybe VName -> (VName -> MulticoreGen ()) -> MulticoreGen ()
forall (t :: * -> *) (m :: * -> *) a b.
(Foldable t, Monad m) =>
t a -> (a -> m b) -> m ()
forM_ Maybe VName
out ((VName -> MulticoreGen ()) -> MulticoreGen ())
-> (VName -> MulticoreGen ()) -> MulticoreGen ()
forall a b. (a -> b) -> a -> b
$ \VName
o ->
          VName -> [TExp Int64] -> SubExp -> [TExp Int64] -> MulticoreGen ()
forall rep r op.
VName -> [TExp Int64] -> SubExp -> [TExp Int64] -> ImpM rep r op ()
copyDWIMFix (Param LParamMem -> VName
forall dec. Param dec -> VName
paramName Param LParamMem
par) [] (VName -> SubExp
Var VName
o) [TV Int64 -> TExp Int64
forall {k} (t :: k). TV t -> TExp t
tvExp TV Int64
task_id, TV Int64 -> TExp Int64
forall {k} (t :: k). TV t -> TExp t
tvExp TV Int64
z]

    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 MCMem -> [SubExpRes]
forall rep res. GBody rep res -> [res]
bodyResult (Body MCMem -> [SubExpRes]) -> Body MCMem -> [SubExpRes]
forall a b. (a -> b) -> a -> b
$ GLambda MCMem Type -> Body MCMem
forall rep t. GLambda rep t -> Body rep
lambdaBody (GLambda MCMem Type -> Body MCMem)
-> GLambda MCMem Type -> Body MCMem
forall a b. (a -> b) -> a -> b
$ SegPostOp MCMem -> GLambda MCMem Type
forall rep. SegPostOp rep -> Lambda rep
segPostOpLambda SegPostOp MCMem
post_op
    sComment "compute post op." $
      compileStms mempty (bodyStms $ lambdaBody $ segPostOpLambda post_op) $
        sComment "write values" $
          forM_ (zip (patElems pat) res) $ \(PatElem LParamMem
pe, SubExp
subexp) ->
            VName -> [TExp Int64] -> SubExp -> [TExp Int64] -> MulticoreGen ()
forall rep r op.
VName -> [TExp Int64] -> SubExp -> [TExp Int64] -> ImpM rep r op ()
copyDWIMFix (PatElem LParamMem -> VName
forall dec. PatElem dec -> VName
patElemName PatElem LParamMem
pe) [TV Int64 -> TExp Int64
forall {k} (t :: k). TV t -> TExp t
tvExp TV Int64
start TExp Int64 -> TExp Int64 -> TExp Int64
forall a. Num a => a -> a -> a
+ TV Int64 -> TExp Int64
forall {k} (t :: k). TV t -> TExp t
tvExp TV Int64
z] SubExp
subexp []

    z <-- tvExp z + 1

getTaskId :: MulticoreGen (TV Int64)
getTaskId :: ImpM MCMem HostEnv Multicore (TV Int64)
getTaskId = do
  task_id <- Name -> ImpM MCMem HostEnv Multicore (TV Int64)
forall {k} (t :: k) rep r op.
MkTV t =>
Name -> ImpM rep r op (TV t)
dPrim Name
"task_id"
  sOp $ Imp.GetTaskId (tvVar task_id)
  pure task_id

nonsegmentedScan ::
  Pat LetDecMem ->
  SegSpace ->
  [Type] ->
  [SegBinOp MCMem] ->
  KernelBody MCMem ->
  SegPostOp MCMem ->
  TV Int32 ->
  MulticoreGen ()
nonsegmentedScan :: Pat LParamMem
-> SegSpace
-> [Type]
-> [SegBinOp MCMem]
-> KernelBody MCMem
-> SegPostOp MCMem
-> TV Int32
-> MulticoreGen ()
nonsegmentedScan
  Pat LParamMem
pat
  (SegSpace VName
fid [(VName
i, SubExp
n)])
  [Type]
ts
  [SegBinOp MCMem]
scan_ops
  KernelBody MCMem
kbody
  SegPostOp MCMem
post_op
  TV Int32
nsubtasks = do
    let multiplier :: TExp Int64
multiplier = TExp Int64
1 -- For playing with.
        blockSize :: TExp Int64
blockSize = TExp Int64
cacheSize TExp Int64 -> TExp Int64 -> TExp Int64
forall e. IntegralExp e => e -> e -> e
`ceilDiv` ([SegBinOp MCMem] -> TExp Int64
totalBytes [SegBinOp MCMem]
scan_ops TExp Int64 -> TExp Int64 -> TExp Int64
forall a. Num a => a -> a -> a
* TExp Int64
multiplier)

    block_no <- Name -> TExp Int64 -> ImpM MCMem HostEnv Multicore (TV Int64)
forall {k} (t :: k) rep r op.
Name -> TExp t -> ImpM rep r op (TV t)
dPrimV Name
"nblocks" (SubExp -> TExp Int64
pe64 SubExp
n TExp Int64 -> TExp Int64 -> TExp Int64
forall e. IntegralExp e => e -> e -> e
`ceilDiv` TExp Int64
blockSize)

    -- allocate flags/aggr/prefix arrays of length nblocks
    flagsArr <- sAllocArray "scan_flags" int64 (Shape [Var (tvVar block_no)]) DefaultSpace

    nsubtasks_i64 <- dPrimV "nsubtasks_i64" $ sExt64 (tvExp nsubtasks)

    blockSize_var <- dPrimV "block_size" blockSize

    let block_shape = [SubExp] -> ShapeBase SubExp
forall d. [d] -> ShapeBase d
Shape [VName -> SubExp
Var (TV Int64 -> VName
forall {k} (t :: k). TV t -> VName
tvVar TV Int64
block_no)]
        nsubtasks_shape = [SubExp] -> ShapeBase SubExp
forall d. [d] -> ShapeBase d
Shape [VName -> SubExp
Var (TV Int64 -> VName
forall {k} (t :: k). TV t -> VName
tvVar TV Int64
nsubtasks_i64), VName -> SubExp
Var (TV Int64 -> VName
forall {k} (t :: k). TV t -> VName
tvVar TV Int64
blockSize_var)]

    aggrArrs <- genArrays scan_ops "scan_aggr" block_shape

    prefArrs <- genArrays scan_ops "scan_pref" block_shape

    scan_out <- genArrays scan_ops "scan_out" nsubtasks_shape

    map_out <- forM (drop (segBinOpResults scan_ops) ts) $ \Type
t ->
      if Type -> Bool
forall shape o. TypeBase shape o -> Bool
isAcc Type
t
        then Maybe VName -> ImpM MCMem HostEnv Multicore (Maybe VName)
forall a. a -> ImpM MCMem HostEnv Multicore 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 MCMem HostEnv Multicore VName
-> ImpM MCMem HostEnv Multicore (Maybe VName)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Name
-> PrimType
-> ShapeBase SubExp
-> Space
-> ImpM MCMem HostEnv Multicore VName
forall rep r op.
Name
-> PrimType -> ShapeBase SubExp -> Space -> ImpM rep r op VName
sAllocArray Name
"map_out" (Type -> PrimType
forall shape o. TypeBase shape o -> PrimType
elemType Type
t) (ShapeBase SubExp
nsubtasks_shape 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) Space
DefaultSpace

    work_index <- sAllocArray "work_index" int64 (Shape [intConst Int64 1]) DefaultSpace

    sFor "init" (tvExp block_no) $ \TExp Int64
j -> do
      VName -> [TExp Int64] -> SubExp -> [TExp Int64] -> MulticoreGen ()
forall rep r op.
VName -> [TExp Int64] -> SubExp -> [TExp Int64] -> ImpM rep r op ()
copyDWIMFix VName
flagsArr [TExp Int64
j] (IntType -> Integer -> SubExp
intConst IntType
Int64 Integer
0) []

    copyDWIMFix work_index [0] (intConst Int64 0) []

    fbody <- collect $ do
      seq_flag <- dPrimV "seq_flag" true

      sOp $ Imp.GetTaskId fid

      block_idx <- dPrim "block_idx"
      work_index_loc <- entryArrayLoc <$> lookupArray work_index
      let work_index_loc_name = MemLoc -> VName
memLocName MemLoc
work_index_loc

      add64 block_idx work_index_loc_name (Imp.elements 0) 1

      sWhile (tvExp block_idx .<. tvExp block_no) $ do
        start <- dPrimV "start" (tvExp block_idx * blockSize)

        diff_start <- dPrimV "diff_start" (pe64 n - tvExp start)

        chunk_length <- dPrim "chunk_length"
        sIf
          (tvExp diff_start .<. blockSize)
          (chunk_length <-- tvExp diff_start)
          (chunk_length <-- blockSize)

        prefix_seqs <- genArrays scan_ops "seq_prefix" (Shape [])

        flags_loc <- entryArrayLoc <$> lookupArray flagsArr
        let flag_loc_name = MemLoc -> VName
memLocName MemLoc
flags_loc
        let block_idx_32 = TExp Int64 -> TExp Int32
forall {k} (t :: k) v. IntExp t => TPrimExp t v -> TPrimExp Int32 v
sExt32 (TV Int64 -> TExp Int64
forall {k} (t :: k). TV t -> TExp t
tvExp TV Int64
block_idx)

        task_id <- getTaskId

        sWhen
          (tvExp seq_flag .==. true)
          ( sIf
              (tvExp block_idx .==. 0)
              (initialiseLocalPrefixes scan_ops prefix_seqs)
              ( do
                  prev_flag <- dPrim "prev_flag" :: MulticoreGen (TV Int64)
                  load64 (tvVar prev_flag) flag_loc_name (Imp.elements $ block_idx_32 - 1)
                  sIf
                    (tvExp prev_flag .==. 2)
                    (updateLocalPrefixes scan_ops prefix_seqs prefArrs [tvExp block_idx - 1])
                    (seq_flag <-- false)
              )
          )

        sIf
          (tvExp seq_flag .==. true)
          ( do
              seqScanFastPath scan_out map_out i scan_ops kbody prefix_seqs start chunk_length prefArrs block_idx task_id

              store64 flag_loc_name (Imp.elements block_idx_32) 2
          )
          ( do
              seqAggregate scan_out map_out i scan_ops kbody start chunk_length aggrArrs block_idx task_id

              -- write flag as 1
              store64 flag_loc_name (Imp.elements block_idx_32) 1

              old_flag <- dPrim "old_flag" :: MulticoreGen (TV Int64)
              old_flag <-- 0
              prefix_vars <- genArrays scan_ops "par_prefix" (Shape [])

              lb <- dPrimV "lb" (tvExp block_idx - 1)
              has_acc <- dPrimV "has_acc" (false :: Imp.TExp Bool)

              error_flag <- dPrim "error_flag" :: MulticoreGen (TV Bool)
              sOp $ Imp.GetError (tvVar error_flag)

              sWhile (bNot (tvExp old_flag .==. 2 .||. tvExp error_flag)) $ do
                load64 (tvVar old_flag) flag_loc_name (Imp.elements $ sExt32 $ tvExp lb)

                -- One of the other operators may fail, in which case we have to
                -- bail out to avoid an infinite wait. This is a very rare case,
                -- so only do it when necessary.
                when (any (hasAssert . lambdaBody . segBinOpLambda) scan_ops || hasAssert kbody) $
                  sOp $
                    Imp.GetError (tvVar error_flag)

                sWhen
                  (tvExp old_flag .==. 2)
                  ( do
                      sIf
                        (tvExp has_acc .==. false)
                        (copyFromDescToLocal scan_ops prefix_vars prefArrs lb)
                        (applyScanFromDescToLocal scan_ops prefix_vars prefArrs lb)
                  )
                sWhen
                  (tvExp old_flag .==. 1)
                  ( do
                      sIf
                        (tvExp has_acc .==. false)
                        (copyFromDescToLocal scan_ops prefix_vars aggrArrs lb)
                        (applyScanFromDescToLocal scan_ops prefix_vars aggrArrs lb)
                      has_acc <-- true
                      lb <-- tvExp lb - 1
                  )

              applyScanAggregateToPrefix scan_ops prefix_vars aggrArrs prefArrs block_idx

              store64 flag_loc_name (Imp.elements block_idx_32) 2

              seqScanLB scan_out i scan_ops kbody prefix_vars start chunk_length task_id
          )

        applyPostOp pat scan_out map_out scan_ops post_op start chunk_length task_id

        add64 block_idx work_index_loc_name (Imp.elements 0) 1

    free_params <- freeParams fbody
    emit $ Imp.Op $ Imp.ParLoop "segmap" fbody free_params

-- | Compile a SegScan construct.
compileSegScan ::
  Pat LetDecMem ->
  SegSpace ->
  [Type] ->
  KernelBody MCMem ->
  [SegBinOp MCMem] ->
  SegPostOp MCMem ->
  TV Int32 ->
  MulticoreGen ()
compileSegScan :: Pat LParamMem
-> SegSpace
-> [Type]
-> KernelBody MCMem
-> [SegBinOp MCMem]
-> SegPostOp MCMem
-> TV Int32
-> MulticoreGen ()
compileSegScan Pat LParamMem
pat SegSpace
space [Type]
ts KernelBody MCMem
kbody [SegBinOp MCMem]
reds SegPostOp MCMem
post_op TV Int32
nsubtasks
  | [(VName, SubExp)
_] <- SegSpace -> [(VName, SubExp)]
unSegSpace SegSpace
space =
      Pat LParamMem
-> SegSpace
-> [Type]
-> [SegBinOp MCMem]
-> KernelBody MCMem
-> SegPostOp MCMem
-> TV Int32
-> MulticoreGen ()
nonsegmentedScan Pat LParamMem
pat SegSpace
space [Type]
ts [SegBinOp MCMem]
reds KernelBody MCMem
kbody SegPostOp MCMem
post_op TV Int32
nsubtasks
  | Bool
otherwise =
      String -> MulticoreGen ()
forall a. HasCallStack => String -> a
error String
"only nonsegmented scans for now"