    module Futhark.CodeGen.ImpGen.Multicore.SegRed
      ( compileSegRed,
        compileSegRed',
        DoSegBody,
      )
    where

    import Control.Monad
    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 (renameLambda)
    import Prelude hiding (quot, rem)

    type DoSegBody = (([[(SubExp, [Imp.TExp Int64])]] -> MulticoreGen ()) -> MulticoreGen ())

    -- | Generate code for a SegRed construct
    compileSegRed ::
      Pat LetDecMem ->
      SegSpace ->
      [SegBinOp MCMem] ->
      KernelBody MCMem ->
      TV Int32 ->
      MulticoreGen Imp.MCCode
    compileSegRed :: Pat LParamMem
    -> SegSpace
    -> [SegBinOp MCMem]
    -> KernelBody MCMem
    -> TV Int32
    -> MulticoreGen MCCode
    compileSegRed Pat LParamMem
    pat SegSpace
    space [SegBinOp MCMem]
    reds KernelBody MCMem
    kbody TV Int32
    nsubtasks =
      Pat LParamMem
    -> SegSpace
    -> [SegBinOp MCMem]
    -> TV Int32
    -> DoSegBody
    -> MulticoreGen MCCode
    compileSegRed' Pat LParamMem
    pat SegSpace
    space [SegBinOp MCMem]
    reds TV Int32
    nsubtasks (DoSegBody -> MulticoreGen MCCode)
    -> DoSegBody -> MulticoreGen MCCode
    forall a b. (a -> b) -> a -> b
    $ \[[(SubExp, [TExp Int64])]] -> MulticoreGen ()
    red_cont ->
        Names -> Stms 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 (KernelBody MCMem -> Stms MCMem
    forall rep. KernelBody rep -> Stms rep
    kernelBodyStms KernelBody MCMem
    kbody) (MulticoreGen () -> MulticoreGen ())
    -> MulticoreGen () -> MulticoreGen ()
    forall a b. (a -> b) -> a -> b
    $ do
          let ([KernelResult]
    red_res, [KernelResult]
    map_res) = Int -> [KernelResult] -> ([KernelResult], [KernelResult])
    forall a. Int -> [a] -> ([a], [a])
    splitAt ([SegBinOp MCMem] -> Int
    forall rep. [SegBinOp rep] -> Int
    segBinOpResults [SegBinOp MCMem]
    reds) ([KernelResult] -> ([KernelResult], [KernelResult]))
    -> [KernelResult] -> ([KernelResult], [KernelResult])
    forall a b. (a -> b) -> a -> b
    $ KernelBody MCMem -> [KernelResult]
    forall rep. KernelBody rep -> [KernelResult]
    kernelBodyResult KernelBody MCMem
    kbody

          Text -> MulticoreGen () -> MulticoreGen ()
    forall rep r op. Text -> ImpM rep r op () -> ImpM rep r op ()
    sComment Text
    "save map-out results" (MulticoreGen () -> MulticoreGen ())
    -> MulticoreGen () -> MulticoreGen ()
    forall a b. (a -> b) -> a -> b
    $ do
            let map_arrs :: [PatElem LParamMem]
    map_arrs = Int -> [PatElem LParamMem] -> [PatElem LParamMem]
    forall a. Int -> [a] -> [a]
    drop ([SegBinOp MCMem] -> Int
    forall rep. [SegBinOp rep] -> Int
    segBinOpResults [SegBinOp MCMem]
    reds) ([PatElem LParamMem] -> [PatElem LParamMem])
    -> [PatElem LParamMem] -> [PatElem LParamMem]
    forall a b. (a -> b) -> a -> b
    $ Pat LParamMem -> [PatElem LParamMem]
    forall dec. Pat dec -> [PatElem dec]
    patElems Pat LParamMem
    pat
            (PatElem LParamMem -> KernelResult -> MulticoreGen ())
    -> [PatElem LParamMem] -> [KernelResult] -> MulticoreGen ()
    forall (m :: * -> *) a b c.
    Applicative m =>
    (a -> b -> m c) -> [a] -> [b] -> m ()
    zipWithM_ (SegSpace -> PatElem LParamMem -> KernelResult -> MulticoreGen ()
    compileThreadResult SegSpace
    space) [PatElem LParamMem]
    map_arrs [KernelResult]
    map_res

          [[(SubExp, [TExp Int64])]] -> MulticoreGen ()
    red_cont ([[(SubExp, [TExp Int64])]] -> MulticoreGen ())
    -> [[(SubExp, [TExp Int64])]] -> MulticoreGen ()
    forall a b. (a -> b) -> a -> b
    $ [SegBinOp MCMem]
    -> [(SubExp, [TExp Int64])] -> [[(SubExp, [TExp Int64])]]
    forall rep a. [SegBinOp rep] -> [a] -> [[a]]
    segBinOpChunks [SegBinOp MCMem]
    reds ([(SubExp, [TExp Int64])] -> [[(SubExp, [TExp Int64])]])
    -> [(SubExp, [TExp Int64])] -> [[(SubExp, [TExp Int64])]]
    forall a b. (a -> b) -> a -> b
    $ (KernelResult -> (SubExp, [TExp Int64]))
    -> [KernelResult] -> [(SubExp, [TExp Int64])]
    forall a b. (a -> b) -> [a] -> [b]
    map ((,[]) (SubExp -> (SubExp, [TExp Int64]))
    -> (KernelResult -> SubExp)
    -> KernelResult
    -> (SubExp, [TExp Int64])
    forall b c a. (b -> c) -> (a -> b) -> a -> c
    . KernelResult -> SubExp
    kernelResultSubExp) [KernelResult]
    red_res

    -- | Like 'compileSegRed', but where the body is a monadic action.
    compileSegRed' ::
      Pat LetDecMem ->
      SegSpace ->
      [SegBinOp MCMem] ->
      TV Int32 ->
      DoSegBody ->
      MulticoreGen Imp.MCCode
    compileSegRed' :: Pat LParamMem
    -> SegSpace
    -> [SegBinOp MCMem]
    -> TV Int32
    -> DoSegBody
    -> MulticoreGen MCCode
    compileSegRed' Pat LParamMem
    pat SegSpace
    space [SegBinOp MCMem]
    reds TV Int32
    nsubtasks DoSegBody
    kbody
      | [(VName, SubExp)
    _] <- SegSpace -> [(VName, SubExp)]
    unSegSpace SegSpace
    space =
          Pat LParamMem
    -> SegSpace
    -> [SegBinOp MCMem]
    -> TV Int32
    -> DoSegBody
    -> MulticoreGen MCCode
    nonsegmentedReduction Pat LParamMem
    pat SegSpace
    space [SegBinOp MCMem]
    reds TV Int32
    nsubtasks DoSegBody
    kbody
      | Bool
    otherwise =
          Pat LParamMem
    -> SegSpace -> [SegBinOp MCMem] -> DoSegBody -> MulticoreGen MCCode
    segmentedReduction Pat LParamMem
    pat SegSpace
    space [SegBinOp MCMem]
    reds DoSegBody
    kbody

    -- | A SegBinOp with auxiliary information.
    data SegBinOpSlug = SegBinOpSlug
      { SegBinOpSlug -> SegBinOp MCMem
    slugOp :: SegBinOp MCMem,
        -- | The array in which we write the intermediate results, indexed
        -- by the flat/physical thread ID.
        SegBinOpSlug -> [VName]
    slugResArrs :: [VName]
      }

    slugBody :: SegBinOpSlug -> Body MCMem
    slugBody :: SegBinOpSlug -> Body MCMem
    slugBody = Lambda MCMem -> Body MCMem
    forall rep. Lambda rep -> Body rep
    lambdaBody (Lambda MCMem -> Body MCMem)
    -> (SegBinOpSlug -> Lambda MCMem) -> SegBinOpSlug -> Body MCMem
    forall b c a. (b -> c) -> (a -> b) -> a -> c
    . SegBinOp MCMem -> Lambda MCMem
    forall rep. SegBinOp rep -> Lambda rep
    segBinOpLambda (SegBinOp MCMem -> Lambda MCMem)
    -> (SegBinOpSlug -> SegBinOp MCMem) -> SegBinOpSlug -> Lambda MCMem
    forall b c a. (b -> c) -> (a -> b) -> a -> c
    . SegBinOpSlug -> SegBinOp MCMem
    slugOp

    slugParams :: SegBinOpSlug -> [LParam MCMem]
    slugParams :: SegBinOpSlug -> [LParam MCMem]
    slugParams = Lambda MCMem -> [LParam MCMem]
    Lambda MCMem -> [Param LParamMem]
    forall rep. Lambda rep -> [LParam rep]
    lambdaParams (Lambda MCMem -> [Param LParamMem])
    -> (SegBinOpSlug -> Lambda MCMem)
    -> SegBinOpSlug
    -> [Param LParamMem]
    forall b c a. (b -> c) -> (a -> b) -> a -> c
    . SegBinOp MCMem -> Lambda MCMem
    forall rep. SegBinOp rep -> Lambda rep
    segBinOpLambda (SegBinOp MCMem -> Lambda MCMem)
    -> (SegBinOpSlug -> SegBinOp MCMem) -> SegBinOpSlug -> Lambda MCMem
    forall b c a. (b -> c) -> (a -> b) -> a -> c
    . SegBinOpSlug -> SegBinOp MCMem
    slugOp

    slugNeutral :: SegBinOpSlug -> [SubExp]
    slugNeutral :: SegBinOpSlug -> [SubExp]
    slugNeutral = SegBinOp MCMem -> [SubExp]
    forall rep. SegBinOp rep -> [SubExp]
    segBinOpNeutral (SegBinOp MCMem -> [SubExp])
    -> (SegBinOpSlug -> SegBinOp MCMem) -> SegBinOpSlug -> [SubExp]
    forall b c a. (b -> c) -> (a -> b) -> a -> c
    . SegBinOpSlug -> SegBinOp MCMem
    slugOp

    slugShape :: SegBinOpSlug -> Shape
    slugShape :: SegBinOpSlug -> Shape
    slugShape = SegBinOp MCMem -> Shape
    forall rep. SegBinOp rep -> Shape
    segBinOpShape (SegBinOp MCMem -> Shape)
    -> (SegBinOpSlug -> SegBinOp MCMem) -> SegBinOpSlug -> Shape
    forall b c a. (b -> c) -> (a -> b) -> a -> c
    . SegBinOpSlug -> SegBinOp MCMem
    slugOp

    accParams, nextParams :: SegBinOpSlug -> [LParam MCMem]
    accParams :: SegBinOpSlug -> [LParam MCMem]
    accParams SegBinOpSlug
    slug = Int -> [LParam MCMem] -> [LParam MCMem]
    forall a. Int -> [a] -> [a]
    take ([SubExp] -> Int
    forall a. [a] -> Int
    forall (t :: * -> *) a. Foldable t => t a -> Int
    length (SegBinOpSlug -> [SubExp]
    slugNeutral SegBinOpSlug
    slug)) ([LParam MCMem] -> [LParam MCMem])
    -> [LParam MCMem] -> [LParam MCMem]
    forall a b. (a -> b) -> a -> b
    $ SegBinOpSlug -> [LParam MCMem]
    slugParams SegBinOpSlug
    slug
    nextParams :: SegBinOpSlug -> [LParam MCMem]
    nextParams SegBinOpSlug
    slug = Int -> [LParam MCMem] -> [LParam MCMem]
    forall a. Int -> [a] -> [a]
    drop ([SubExp] -> Int
    forall a. [a] -> Int
    forall (t :: * -> *) a. Foldable t => t a -> Int
    length (SegBinOpSlug -> [SubExp]
    slugNeutral SegBinOpSlug
    slug)) ([LParam MCMem] -> [LParam MCMem])
    -> [LParam MCMem] -> [LParam MCMem]
    forall a b. (a -> b) -> a -> b
    $ SegBinOpSlug -> [LParam MCMem]
    slugParams SegBinOpSlug
    slug

    renameSlug :: SegBinOpSlug -> MulticoreGen SegBinOpSlug
    renameSlug :: SegBinOpSlug -> MulticoreGen SegBinOpSlug
    renameSlug SegBinOpSlug
    slug = do
      let op :: SegBinOp MCMem
    op = SegBinOpSlug -> SegBinOp MCMem
    slugOp SegBinOpSlug
    slug
      let lambda :: Lambda MCMem
    lambda = SegBinOp MCMem -> Lambda MCMem
    forall rep. SegBinOp rep -> Lambda rep
    segBinOpLambda SegBinOp MCMem
    op
      lambda' <- Lambda MCMem -> ImpM MCMem HostEnv Multicore (Lambda MCMem)
    forall rep (m :: * -> *).
    (Renameable rep, MonadFreshNames m) =>
    Lambda rep -> m (Lambda rep)
    renameLambda Lambda MCMem
    lambda
      let op' = SegBinOp MCMem
    op {segBinOpLambda = lambda'}
      pure slug {slugOp = op'}

    -- | Arrays for storing group results shared between threads
    groupResultArrays ::
      String ->
      SubExp ->
      [SegBinOp MCMem] ->
      MulticoreGen [[VName]]
    groupResultArrays :: String -> SubExp -> [SegBinOp MCMem] -> MulticoreGen [[VName]]
    groupResultArrays String
    s SubExp
    num_threads [SegBinOp MCMem]
    reds =
      [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]
    reds ((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 Commutativity
    _ Lambda MCMem
    lam [SubExp]
    _ Shape
    shape) ->
        [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 (Lambda MCMem -> [Type]
    forall rep. Lambda rep -> [Type]
    lambdaReturnType Lambda MCMem
    lam) ((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 full_shape :: Shape
    full_shape = [SubExp] -> Shape
    forall d. [d] -> ShapeBase d
    Shape [SubExp
    num_threads] Shape -> Shape -> Shape
    forall a. Semigroup a => a -> a -> a
    <> Shape
    shape Shape -> Shape -> Shape
    forall a. Semigroup a => a -> a -> a
    <> Type -> Shape
    forall shape u. ArrayShape shape => TypeBase shape u -> shape
    arrayShape Type
    t
          String
    -> PrimType -> Shape -> Space -> ImpM MCMem HostEnv Multicore VName
    forall rep r op.
    String -> PrimType -> Shape -> Space -> ImpM rep r op VName
    sAllocArray String
    s (Type -> PrimType
    forall shape u. TypeBase shape u -> PrimType
    elemType Type
    t) Shape
    full_shape Space
    DefaultSpace

    nonsegmentedReduction ::
      Pat LetDecMem ->
      SegSpace ->
      [SegBinOp MCMem] ->
      TV Int32 ->
      DoSegBody ->
      MulticoreGen Imp.MCCode
    nonsegmentedReduction :: Pat LParamMem
    -> SegSpace
    -> [SegBinOp MCMem]
    -> TV Int32
    -> DoSegBody
    -> MulticoreGen MCCode
    nonsegmentedReduction Pat LParamMem
    pat SegSpace
    space [SegBinOp MCMem]
    reds TV Int32
    nsubtasks DoSegBody
    kbody = MulticoreGen () -> MulticoreGen MCCode
    forall rep r op. ImpM rep r op () -> ImpM rep r op (Code op)
    collect (MulticoreGen () -> MulticoreGen MCCode)
    -> MulticoreGen () -> MulticoreGen MCCode
    forall a b. (a -> b) -> a -> b
    $ do
      thread_res_arrs <- String -> SubExp -> [SegBinOp MCMem] -> MulticoreGen [[VName]]
    groupResultArrays String
    "reduce_stage_1_tid_res_arr" (TV Int32 -> SubExp
    forall {k} (t :: k). TV t -> SubExp
    tvSize TV Int32
    nsubtasks) [SegBinOp MCMem]
    reds
      let slugs1 = (SegBinOp MCMem -> [VName] -> SegBinOpSlug)
    -> [SegBinOp MCMem] -> [[VName]] -> [SegBinOpSlug]
    forall a b c. (a -> b -> c) -> [a] -> [b] -> [c]
    zipWith SegBinOp MCMem -> [VName] -> SegBinOpSlug
    SegBinOpSlug [SegBinOp MCMem]
    reds [[VName]]
    thread_res_arrs
          nsubtasks' = TV Int32 -> TExp Int32
    forall {k} (t :: k). TV t -> TExp t
    tvExp TV Int32
    nsubtasks

      -- Are all the operators commutative?
      let comm = (SegBinOp MCMem -> Bool) -> [SegBinOp MCMem] -> Bool
    forall (t :: * -> *) a. Foldable t => (a -> Bool) -> t a -> Bool
    all ((Commutativity -> Commutativity -> Bool
    forall a. Eq a => a -> a -> Bool
    == Commutativity
    Commutative) (Commutativity -> Bool)
    -> (SegBinOp MCMem -> Commutativity) -> SegBinOp MCMem -> Bool
    forall b c a. (b -> c) -> (a -> b) -> a -> c
    . SegBinOp MCMem -> Commutativity
    forall rep. SegBinOp rep -> Commutativity
    segBinOpComm) [SegBinOp MCMem]
    reds
      let dims = (SegBinOpSlug -> [SubExp]) -> [SegBinOpSlug] -> [[SubExp]]
    forall a b. (a -> b) -> [a] -> [b]
    map (Shape -> [SubExp]
    forall d. ShapeBase d -> [d]
    shapeDims (Shape -> [SubExp])
    -> (SegBinOpSlug -> Shape) -> SegBinOpSlug -> [SubExp]
    forall b c a. (b -> c) -> (a -> b) -> a -> c
    . SegBinOpSlug -> Shape
    slugShape) [SegBinOpSlug]
    slugs1
      let isScalar MemInfo d u ret
    x = case MemInfo d u ret
    x of MemPrim PrimType
    _ -> Bool
    True; MemInfo d u ret
    _ -> Bool
    False
      -- Are we only working on scalar arrays?
      let scalars = (SegBinOpSlug -> Bool) -> [SegBinOpSlug] -> Bool
    forall (t :: * -> *) a. Foldable t => (a -> Bool) -> t a -> Bool
    all ((Param LParamMem -> Bool) -> [Param LParamMem] -> Bool
    forall (t :: * -> *) a. Foldable t => (a -> Bool) -> t a -> Bool
    all (LParamMem -> Bool
    forall {d} {u} {ret}. MemInfo d u ret -> Bool
    isScalar (LParamMem -> Bool)
    -> (Param LParamMem -> LParamMem) -> Param LParamMem -> Bool
    forall b c a. (b -> c) -> (a -> b) -> a -> c
    . Param LParamMem -> LParamMem
    forall dec. Param dec -> dec
    paramDec) ([Param LParamMem] -> Bool)
    -> (SegBinOpSlug -> [Param LParamMem]) -> SegBinOpSlug -> Bool
    forall b c a. (b -> c) -> (a -> b) -> a -> c
    . SegBinOpSlug -> [LParam MCMem]
    SegBinOpSlug -> [Param LParamMem]
    slugParams) [SegBinOpSlug]
    slugs1 Bool -> Bool -> Bool
    && ([SubExp] -> Bool) -> [[SubExp]] -> Bool
    forall (t :: * -> *) a. Foldable t => (a -> Bool) -> t a -> Bool
    all ([SubExp] -> [SubExp] -> Bool
    forall a. Eq a => a -> a -> Bool
    == []) [[SubExp]]
    dims
      -- Are we working with vectorized inner maps?
      let inner_map = [] [SubExp] -> [[SubExp]] -> Bool
    forall (t :: * -> *) a. (Foldable t, Eq a) => a -> t a -> Bool
    `notElem` [[SubExp]]
    dims

      let path
            | Bool
    comm Bool -> Bool -> Bool
    && Bool
    scalars = ReductionStage1
    reductionStage1CommScalar
            | Bool
    inner_map = ReductionStage1
    reductionStage1Array
            | Bool
    scalars = ReductionStage1
    reductionStage1NonCommScalar
            | Bool
    otherwise = ReductionStage1
    reductionStage1Fallback
      path space slugs1 kbody

      reds2 <- renameSegBinOp reds
      let slugs2 = (SegBinOp MCMem -> [VName] -> SegBinOpSlug)
    -> [SegBinOp MCMem] -> [[VName]] -> [SegBinOpSlug]
    forall a b c. (a -> b -> c) -> [a] -> [b] -> [c]
    zipWith SegBinOp MCMem -> [VName] -> SegBinOpSlug
    SegBinOpSlug [SegBinOp MCMem]
    reds2 [[VName]]
    thread_res_arrs
      reductionStage2 pat space nsubtasks' slugs2

    -- Generate code that declares the params for the binop
    genBinOpParams :: [SegBinOpSlug] -> MulticoreGen ()
    genBinOpParams :: [SegBinOpSlug] -> MulticoreGen ()
    genBinOpParams [SegBinOpSlug]
    slugs =
      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
    $ (SegBinOpSlug -> [Param LParamMem])
    -> [SegBinOpSlug] -> [Param LParamMem]
    forall (t :: * -> *) a b. Foldable t => (a -> [b]) -> t a -> [b]
    concatMap SegBinOpSlug -> [LParam MCMem]
    SegBinOpSlug -> [Param LParamMem]
    slugParams [SegBinOpSlug]
    slugs

    -- Generate code that declares accumulators, return a list of these
    genAccumulators :: [SegBinOpSlug] -> MulticoreGen [[VName]]
    genAccumulators :: [SegBinOpSlug] -> MulticoreGen [[VName]]
    genAccumulators [SegBinOpSlug]
    slugs =
      [SegBinOpSlug]
    -> (SegBinOpSlug -> 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 [SegBinOpSlug]
    slugs ((SegBinOpSlug -> ImpM MCMem HostEnv Multicore [VName])
     -> MulticoreGen [[VName]])
    -> (SegBinOpSlug -> ImpM MCMem HostEnv Multicore [VName])
    -> MulticoreGen [[VName]]
    forall a b. (a -> b) -> a -> b
    $ \SegBinOpSlug
    slug -> do
        let shape :: Shape
    shape = SegBinOp MCMem -> Shape
    forall rep. SegBinOp rep -> Shape
    segBinOpShape (SegBinOp MCMem -> Shape) -> SegBinOp MCMem -> Shape
    forall a b. (a -> b) -> a -> b
    $ SegBinOpSlug -> SegBinOp MCMem
    slugOp SegBinOpSlug
    slug
        [(Param LParamMem, SubExp)]
    -> ((Param LParamMem, SubExp)
        -> 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] -> [SubExp] -> [(Param LParamMem, SubExp)]
    forall a b. [a] -> [b] -> [(a, b)]
    zip (SegBinOpSlug -> [LParam MCMem]
    accParams SegBinOpSlug
    slug) (SegBinOpSlug -> [SubExp]
    slugNeutral SegBinOpSlug
    slug)) (((Param LParamMem, SubExp) -> ImpM MCMem HostEnv Multicore VName)
     -> ImpM MCMem HostEnv Multicore [VName])
    -> ((Param LParamMem, SubExp)
        -> ImpM MCMem HostEnv Multicore VName)
    -> ImpM MCMem HostEnv Multicore [VName]
    forall a b. (a -> b) -> a -> b
    $ \(Param LParamMem
    p, SubExp
    ne) -> do
          -- Declare accumulator variable.
          acc <-
            case Param LParamMem -> Type
    forall dec. Typed dec => Param dec -> Type
    paramType Param LParamMem
    p of
              Prim PrimType
    pt
                | Shape
    shape Shape -> Shape -> Bool
    forall a. Eq a => a -> a -> Bool
    == Shape
    forall a. Monoid a => a
    mempty ->
                    String -> PrimType -> ImpM MCMem HostEnv Multicore VName
    forall rep r op. String -> PrimType -> ImpM rep r op VName
    dPrimS String
    "local_acc" PrimType
    pt
                | Bool
    otherwise ->
                    String
    -> PrimType -> Shape -> Space -> ImpM MCMem HostEnv Multicore VName
    forall rep r op.
    String -> PrimType -> Shape -> Space -> ImpM rep r op VName
    sAllocArray String
    "local_acc" PrimType
    pt Shape
    shape Space
    DefaultSpace
              Type
    _ ->
                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

          -- Now neutral-initialise the accumulator.
          sLoopNest (slugShape slug) $ \[TExp Int64]
    vec_is ->
            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
    ne []

          pure acc

    -- Datatype to represent all the different ways we can generate
    -- code for a reduction.
    data RedLoopType
      = RedSeq -- Fully sequential
      | RedComm -- Commutative scalar
      | RedNonComm -- Noncommutative scalar
      | RedNested -- Nested vectorized operator
      | RedUniformize -- Uniformize over scalar acc

    -- Given a type of reduction and the loop index, should we wrap
    -- the loop body in some extra code?
    getRedLoop ::
      RedLoopType ->
      Imp.TExp Int64 ->
      (Imp.TExp Int64 -> MulticoreGen ()) ->
      MulticoreGen ()
    getRedLoop :: RedLoopType
    -> TExp Int64 -> (TExp Int64 -> MulticoreGen ()) -> MulticoreGen ()
    getRedLoop RedLoopType
    RedNonComm TExp Int64
    _ = (TExp Int64 -> MulticoreGen ()) -> MulticoreGen ()
    generateUniformizeLoop
    getRedLoop RedLoopType
    RedUniformize TExp Int64
    uni = \TExp Int64 -> MulticoreGen ()
    body -> TExp Int64 -> MulticoreGen ()
    body TExp Int64
    uni
    getRedLoop RedLoopType
    _ TExp Int64
    _ = \TExp Int64 -> MulticoreGen ()
    body -> TExp Int64 -> MulticoreGen ()
    body TExp Int64
    0

    -- Given a type of reduction, should we perform extracts on
    -- the accumulator?
    getExtract ::
      RedLoopType ->
      Imp.TExp Int64 ->
      MulticoreGen Imp.MCCode ->
      MulticoreGen ()
    getExtract :: RedLoopType -> TExp Int64 -> MulticoreGen MCCode -> MulticoreGen ()
    getExtract RedLoopType
    RedNonComm = TExp Int64 -> MulticoreGen MCCode -> MulticoreGen ()
    extractVectorLane
    getExtract RedLoopType
    RedUniformize = TExp Int64 -> MulticoreGen MCCode -> MulticoreGen ()
    extractVectorLane
    getExtract RedLoopType
    _ = \TExp Int64
    _ MulticoreGen MCCode
    body -> MulticoreGen MCCode
    body MulticoreGen MCCode
    -> (MCCode -> MulticoreGen ()) -> MulticoreGen ()
    forall a b.
    ImpM MCMem HostEnv Multicore a
    -> (a -> ImpM MCMem HostEnv Multicore b)
    -> ImpM MCMem HostEnv Multicore b
    forall (m :: * -> *) a b. Monad m => m a -> (a -> m b) -> m b
    >>= MCCode -> MulticoreGen ()
    forall op rep r. Code op -> ImpM rep r op ()
    emit

    -- Given a type of reduction, should we vectorize the inner
    -- map, if it exists?
    getNestLoop ::
      RedLoopType ->
      Shape ->
      ([Imp.TExp Int64] -> MulticoreGen ()) ->
      MulticoreGen ()
    getNestLoop :: RedLoopType
    -> Shape -> ([TExp Int64] -> MulticoreGen ()) -> MulticoreGen ()
    getNestLoop RedLoopType
    RedNested = Shape -> ([TExp Int64] -> MulticoreGen ()) -> MulticoreGen ()
    sLoopNestVectorized
    getNestLoop RedLoopType
    _ = Shape -> ([TExp Int64] -> MulticoreGen ()) -> MulticoreGen ()
    forall rep r op.
    Shape -> ([TExp Int64] -> ImpM rep r op ()) -> ImpM rep r op ()
    sLoopNest

    -- Given a list of accumulators, use them as the source
    -- data for reduction.
    redSourceAccs :: [[VName]] -> DoSegBody
    redSourceAccs :: [[VName]] -> DoSegBody
    redSourceAccs [[VName]]
    slug_local_accs [[(SubExp, [TExp Int64])]] -> MulticoreGen ()
    m =
      [[(SubExp, [TExp Int64])]] -> MulticoreGen ()
    m ([[(SubExp, [TExp Int64])]] -> MulticoreGen ())
    -> [[(SubExp, [TExp Int64])]] -> MulticoreGen ()
    forall a b. (a -> b) -> a -> b
    $ ([VName] -> [(SubExp, [TExp Int64])])
    -> [[VName]] -> [[(SubExp, [TExp Int64])]]
    forall a b. (a -> b) -> [a] -> [b]
    map ((VName -> (SubExp, [TExp Int64]))
    -> [VName] -> [(SubExp, [TExp Int64])]
    forall a b. (a -> b) -> [a] -> [b]
    map (\VName
    x -> (VName -> SubExp
    Var VName
    x, []))) [[VName]]
    slug_local_accs

    -- Generate a reduction loop for uniformizing vectors
    genPostbodyReductionLoop ::
      [[VName]] ->
      [SegBinOpSlug] ->
      [[VName]] ->
      SegSpace ->
      Imp.TExp Int64 ->
      MulticoreGen ()
    genPostbodyReductionLoop :: [[VName]]
    -> [SegBinOpSlug]
    -> [[VName]]
    -> SegSpace
    -> TExp Int64
    -> MulticoreGen ()
    genPostbodyReductionLoop [[VName]]
    accs =
      RedLoopType
    -> DoSegBody
    -> [SegBinOpSlug]
    -> [[VName]]
    -> SegSpace
    -> TExp Int64
    -> MulticoreGen ()
    genReductionLoop RedLoopType
    RedUniformize ([[VName]] -> DoSegBody
    redSourceAccs [[VName]]
    accs)

    -- Generate a potentially vectorized body of code that performs reduction
    -- when put inside a chunked loop.
    genReductionLoop ::
      RedLoopType ->
      DoSegBody ->
      [SegBinOpSlug] ->
      [[VName]] ->
      SegSpace ->
      Imp.TExp Int64 ->
      MulticoreGen ()
    genReductionLoop :: RedLoopType
    -> DoSegBody
    -> [SegBinOpSlug]
    -> [[VName]]
    -> SegSpace
    -> TExp Int64
    -> MulticoreGen ()
    genReductionLoop RedLoopType
    typ DoSegBody
    kbodymap [SegBinOpSlug]
    slugs [[VName]]
    slug_local_accs SegSpace
    space TExp Int64
    i = do
      let ([VName]
    is, [SubExp]
    ns) = [(VName, SubExp)] -> ([VName], [SubExp])
    forall a b. [(a, b)] -> ([a], [b])
    unzip ([(VName, SubExp)] -> ([VName], [SubExp]))
    -> [(VName, SubExp)] -> ([VName], [SubExp])
    forall a b. (a -> b) -> a -> b
    $ SegSpace -> [(VName, SubExp)]
    unSegSpace SegSpace
    space
          ns' :: [TExp Int64]
    ns' = (SubExp -> TExp Int64) -> [SubExp] -> [TExp Int64]
    forall a b. (a -> b) -> [a] -> [b]
    map SubExp -> TExp Int64
    pe64 [SubExp]
    ns
      (VName -> TExp Int64 -> MulticoreGen ())
    -> [VName] -> [TExp Int64] -> MulticoreGen ()
    forall (m :: * -> *) a b c.
    Applicative m =>
    (a -> b -> m c) -> [a] -> [b] -> m ()
    zipWithM_ VName -> TExp Int64 -> MulticoreGen ()
    forall {k} (t :: k) rep r op. VName -> TExp t -> ImpM rep r op ()
    dPrimV_ [VName]
    is ([TExp Int64] -> MulticoreGen ())
    -> [TExp Int64] -> MulticoreGen ()
    forall a b. (a -> b) -> a -> b
    $ [TExp Int64] -> TExp Int64 -> [TExp Int64]
    forall num. IntegralExp num => [num] -> num -> [num]
    unflattenIndex [TExp Int64]
    ns' TExp Int64
    i
      DoSegBody
    kbodymap DoSegBody -> DoSegBody
    forall a b. (a -> b) -> a -> b
    $ \[[(SubExp, [TExp Int64])]]
    all_red_res' -> do
        [([(SubExp, [TExp Int64])], SegBinOpSlug, [VName])]
    -> (([(SubExp, [TExp Int64])], SegBinOpSlug, [VName])
        -> MulticoreGen ())
    -> MulticoreGen ()
    forall (t :: * -> *) (m :: * -> *) a b.
    (Foldable t, Monad m) =>
    t a -> (a -> m b) -> m ()
    forM_ ([[(SubExp, [TExp Int64])]]
    -> [SegBinOpSlug]
    -> [[VName]]
    -> [([(SubExp, [TExp Int64])], SegBinOpSlug, [VName])]
    forall a b c. [a] -> [b] -> [c] -> [(a, b, c)]
    zip3 [[(SubExp, [TExp Int64])]]
    all_red_res' [SegBinOpSlug]
    slugs [[VName]]
    slug_local_accs) ((([(SubExp, [TExp Int64])], SegBinOpSlug, [VName])
      -> MulticoreGen ())
     -> MulticoreGen ())
    -> (([(SubExp, [TExp Int64])], SegBinOpSlug, [VName])
        -> MulticoreGen ())
    -> MulticoreGen ()
    forall a b. (a -> b) -> a -> b
    $ \([(SubExp, [TExp Int64])]
    red_res, SegBinOpSlug
    slug, [VName]
    local_accs) ->
          RedLoopType
    -> Shape -> ([TExp Int64] -> MulticoreGen ()) -> MulticoreGen ()
    getNestLoop RedLoopType
    typ (SegBinOpSlug -> Shape
    slugShape SegBinOpSlug
    slug) (([TExp Int64] -> MulticoreGen ()) -> MulticoreGen ())
    -> ([TExp Int64] -> MulticoreGen ()) -> MulticoreGen ()
    forall a b. (a -> b) -> a -> b
    $ \[TExp Int64]
    vec_is -> do
            let lamtypes :: [Type]
    lamtypes = Lambda MCMem -> [Type]
    forall rep. Lambda rep -> [Type]
    lambdaReturnType (Lambda MCMem -> [Type]) -> Lambda MCMem -> [Type]
    forall a b. (a -> b) -> a -> b
    $ SegBinOp MCMem -> Lambda MCMem
    forall rep. SegBinOp rep -> Lambda rep
    segBinOpLambda (SegBinOp MCMem -> Lambda MCMem) -> SegBinOp MCMem -> Lambda MCMem
    forall a b. (a -> b) -> a -> b
    $ SegBinOpSlug -> SegBinOp MCMem
    slugOp SegBinOpSlug
    slug
            RedLoopType
    -> TExp Int64 -> (TExp Int64 -> MulticoreGen ()) -> MulticoreGen ()
    getRedLoop RedLoopType
    typ TExp Int64
    i ((TExp Int64 -> MulticoreGen ()) -> MulticoreGen ())
    -> (TExp Int64 -> MulticoreGen ()) -> MulticoreGen ()
    forall a b. (a -> b) -> a -> b
    $ \TExp Int64
    uni -> do
              Text -> MulticoreGen () -> MulticoreGen ()
    forall rep r op. Text -> ImpM rep r op () -> ImpM rep r op ()
    sComment Text
    "Load accum params" (MulticoreGen () -> MulticoreGen ())
    -> MulticoreGen () -> MulticoreGen ()
    forall a b. (a -> b) -> a -> b
    $
                [(Param LParamMem, VName, Type)]
    -> ((Param LParamMem, VName, Type) -> MulticoreGen ())
    -> MulticoreGen ()
    forall (t :: * -> *) (m :: * -> *) a b.
    (Foldable t, Monad m) =>
    t a -> (a -> m b) -> m ()
    forM_ ([Param LParamMem]
    -> [VName] -> [Type] -> [(Param LParamMem, VName, Type)]
    forall a b c. [a] -> [b] -> [c] -> [(a, b, c)]
    zip3 (SegBinOpSlug -> [LParam MCMem]
    accParams SegBinOpSlug
    slug) [VName]
    local_accs [Type]
    lamtypes) (((Param LParamMem, VName, Type) -> MulticoreGen ())
     -> MulticoreGen ())
    -> ((Param LParamMem, VName, Type) -> MulticoreGen ())
    -> MulticoreGen ()
    forall a b. (a -> b) -> a -> b
    $
                  \(Param LParamMem
    p, VName
    local_acc, Type
    t) ->
                    Bool -> MulticoreGen () -> MulticoreGen ()
    forall (f :: * -> *). Applicative f => Bool -> f () -> f ()
    when (Type -> Bool
    forall shape u. TypeBase shape u -> Bool
    primType Type
    t) (MulticoreGen () -> MulticoreGen ())
    -> MulticoreGen () -> MulticoreGen ()
    forall a b. (a -> b) -> a -> b
    $ 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
    p) [] (VName -> SubExp
    Var VName
    local_acc) [TExp Int64]
    vec_is

              Text -> MulticoreGen () -> MulticoreGen ()
    forall rep r op. Text -> ImpM rep r op () -> ImpM rep r op ()
    sComment Text
    "Load next params" (MulticoreGen () -> MulticoreGen ())
    -> MulticoreGen () -> MulticoreGen ()
    forall a b. (a -> b) -> a -> b
    $
                [(Param LParamMem, (SubExp, [TExp Int64]))]
    -> ((Param LParamMem, (SubExp, [TExp Int64])) -> MulticoreGen ())
    -> MulticoreGen ()
    forall (t :: * -> *) (m :: * -> *) a b.
    (Foldable t, Monad m) =>
    t a -> (a -> m b) -> m ()
    forM_ ([Param LParamMem]
    -> [(SubExp, [TExp Int64])]
    -> [(Param LParamMem, (SubExp, [TExp Int64]))]
    forall a b. [a] -> [b] -> [(a, b)]
    zip (SegBinOpSlug -> [LParam MCMem]
    nextParams SegBinOpSlug
    slug) [(SubExp, [TExp Int64])]
    red_res) (((Param LParamMem, (SubExp, [TExp Int64])) -> MulticoreGen ())
     -> MulticoreGen ())
    -> ((Param LParamMem, (SubExp, [TExp Int64])) -> MulticoreGen ())
    -> MulticoreGen ()
    forall a b. (a -> b) -> a -> b
    $ \(Param LParamMem
    p, (SubExp
    res, [TExp Int64]
    res_is)) -> do
                  RedLoopType -> TExp Int64 -> MulticoreGen MCCode -> MulticoreGen ()
    getExtract RedLoopType
    typ TExp Int64
    uni (MulticoreGen MCCode -> MulticoreGen ())
    -> MulticoreGen MCCode -> MulticoreGen ()
    forall a b. (a -> b) -> a -> b
    $
                    MulticoreGen () -> MulticoreGen MCCode
    forall rep r op. ImpM rep r op () -> ImpM rep r op (Code op)
    collect (MulticoreGen () -> MulticoreGen MCCode)
    -> MulticoreGen () -> MulticoreGen MCCode
    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 (Param LParamMem -> VName
    forall dec. Param dec -> VName
    paramName Param LParamMem
    p) [] SubExp
    res ([TExp Int64]
    res_is [TExp Int64] -> [TExp Int64] -> [TExp Int64]
    forall a. [a] -> [a] -> [a]
    ++ [TExp Int64]
    vec_is)

              Text -> MulticoreGen () -> MulticoreGen ()
    forall rep r op. Text -> ImpM rep r op () -> ImpM rep r op ()
    sComment Text
    "SegRed body" (MulticoreGen () -> MulticoreGen ())
    -> MulticoreGen () -> MulticoreGen ()
    forall a b. (a -> b) -> a -> b
    $
                Names -> Stms 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 -> Stms MCMem
    forall rep. Body rep -> Stms rep
    bodyStms (Body MCMem -> Stms MCMem) -> Body MCMem -> Stms MCMem
    forall a b. (a -> b) -> a -> b
    $ SegBinOpSlug -> Body MCMem
    slugBody SegBinOpSlug
    slug) (MulticoreGen () -> MulticoreGen ())
    -> MulticoreGen () -> MulticoreGen ()
    forall a b. (a -> b) -> a -> b
    $
                  [(VName, SubExp)]
    -> ((VName, SubExp) -> MulticoreGen ()) -> MulticoreGen ()
    forall (t :: * -> *) (m :: * -> *) a b.
    (Foldable t, Monad m) =>
    t a -> (a -> m b) -> m ()
    forM_ ([VName] -> [SubExp] -> [(VName, SubExp)]
    forall a b. [a] -> [b] -> [(a, b)]
    zip [VName]
    local_accs ([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 MCMem -> [SubExpRes]
    forall rep. Body rep -> [SubExpRes]
    bodyResult (Body MCMem -> [SubExpRes]) -> Body MCMem -> [SubExpRes]
    forall a b. (a -> b) -> a -> b
    $ SegBinOpSlug -> Body MCMem
    slugBody SegBinOpSlug
    slug) (((VName, SubExp) -> MulticoreGen ()) -> MulticoreGen ())
    -> ((VName, SubExp) -> MulticoreGen ()) -> MulticoreGen ()
    forall a b. (a -> b) -> a -> b
    $
                    \(VName
    local_acc, 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
    local_acc [TExp Int64]
    vec_is SubExp
    se []

    -- Generate code to write back results from the accumulators
    genWriteBack :: [SegBinOpSlug] -> [[VName]] -> SegSpace -> MulticoreGen ()
    genWriteBack :: [SegBinOpSlug] -> [[VName]] -> SegSpace -> MulticoreGen ()
    genWriteBack [SegBinOpSlug]
    slugs [[VName]]
    slug_local_accs SegSpace
    space =
      [(SegBinOpSlug, [VName])]
    -> ((SegBinOpSlug, [VName]) -> ImpM MCMem HostEnv Multicore [()])
    -> MulticoreGen ()
    forall (t :: * -> *) (m :: * -> *) a b.
    (Foldable t, Monad m) =>
    t a -> (a -> m b) -> m ()
    forM_ ([SegBinOpSlug] -> [[VName]] -> [(SegBinOpSlug, [VName])]
    forall a b. [a] -> [b] -> [(a, b)]
    zip [SegBinOpSlug]
    slugs [[VName]]
    slug_local_accs) (((SegBinOpSlug, [VName]) -> ImpM MCMem HostEnv Multicore [()])
     -> MulticoreGen ())
    -> ((SegBinOpSlug, [VName]) -> ImpM MCMem HostEnv Multicore [()])
    -> MulticoreGen ()
    forall a b. (a -> b) -> a -> b
    $ \(SegBinOpSlug
    slug, [VName]
    local_accs) ->
        [(VName, VName)]
    -> ((VName, VName) -> MulticoreGen ())
    -> ImpM MCMem HostEnv Multicore [()]
    forall (t :: * -> *) (m :: * -> *) a b.
    (Traversable t, Monad m) =>
    t a -> (a -> m b) -> m (t b)
    forM ([VName] -> [VName] -> [(VName, VName)]
    forall a b. [a] -> [b] -> [(a, b)]
    zip (SegBinOpSlug -> [VName]
    slugResArrs SegBinOpSlug
    slug) [VName]
    local_accs) (((VName, VName) -> MulticoreGen ())
     -> ImpM MCMem HostEnv Multicore [()])
    -> ((VName, VName) -> MulticoreGen ())
    -> ImpM MCMem HostEnv Multicore [()]
    forall a b. (a -> b) -> a -> b
    $ \(VName
    acc, VName
    local_acc) ->
          VName -> [TExp Int64] -> SubExp -> [TExp Int64] -> MulticoreGen ()
    forall rep r op.
    VName -> [TExp Int64] -> SubExp -> [TExp Int64] -> ImpM rep r op ()
    copyDWIMFix VName
    acc [VName -> TExp Int64
    forall a. a -> TPrimExp Int64 a
    Imp.le64 (VName -> TExp Int64) -> VName -> TExp Int64
    forall a b. (a -> b) -> a -> b
    $ SegSpace -> VName
    segFlat SegSpace
    space] (VName -> SubExp
    Var VName
    local_acc) []

    type ReductionStage1 = SegSpace -> [SegBinOpSlug] -> DoSegBody -> MulticoreGen ()

    -- Pure sequential codegen with no fancy vectorization
    reductionStage1Fallback :: ReductionStage1
    reductionStage1Fallback :: ReductionStage1
    reductionStage1Fallback SegSpace
    space [SegBinOpSlug]
    slugs DoSegBody
    kbody = do
      fbody <- MulticoreGen () -> MulticoreGen MCCode
    forall rep r op. ImpM rep r op () -> ImpM rep r op (Code op)
    collect (MulticoreGen () -> MulticoreGen MCCode)
    -> MulticoreGen () -> MulticoreGen MCCode
    forall a b. (a -> b) -> a -> b
    $ do
        VName -> PrimType -> MulticoreGen ()
    forall rep r op. VName -> PrimType -> ImpM rep r op ()
    dPrim_ (SegSpace -> VName
    segFlat SegSpace
    space) PrimType
    int64
        Multicore -> MulticoreGen ()
    forall op rep r. op -> ImpM rep r op ()
    sOp (Multicore -> MulticoreGen ()) -> Multicore -> MulticoreGen ()
    forall a b. (a -> b) -> a -> b
    $ VName -> Multicore
    Imp.GetTaskId (SegSpace -> VName
    segFlat SegSpace
    space)
        -- Declare params
        [SegBinOpSlug] -> MulticoreGen ()
    genBinOpParams [SegBinOpSlug]
    slugs
        slug_local_accs <- [SegBinOpSlug] -> MulticoreGen [[VName]]
    genAccumulators [SegBinOpSlug]
    slugs
        -- Generate main reduction loop
        generateChunkLoop "SegRed" Scalar $
          genReductionLoop RedSeq kbody slugs slug_local_accs space
        -- Write back results
        genWriteBack slugs slug_local_accs space
      free_params <- freeParams fbody
      emit $ Imp.Op $ Imp.ParLoop "segred_stage_1" fbody free_params

    -- Codegen for noncommutative scalar reduction. We vectorize the
    -- kernel body, and do the reduction sequentially.
    reductionStage1NonCommScalar :: ReductionStage1
    reductionStage1NonCommScalar :: ReductionStage1
    reductionStage1NonCommScalar SegSpace
    space [SegBinOpSlug]
    slugs DoSegBody
    kbody = do
      fbody <- MulticoreGen () -> MulticoreGen MCCode
    forall rep r op. ImpM rep r op () -> ImpM rep r op (Code op)
    collect (MulticoreGen () -> MulticoreGen MCCode)
    -> MulticoreGen () -> MulticoreGen MCCode
    forall a b. (a -> b) -> a -> b
    $ do
        VName -> PrimType -> MulticoreGen ()
    forall rep r op. VName -> PrimType -> ImpM rep r op ()
    dPrim_ (SegSpace -> VName
    segFlat SegSpace
    space) PrimType
    int64
        Multicore -> MulticoreGen ()
    forall op rep r. op -> ImpM rep r op ()
    sOp (Multicore -> MulticoreGen ()) -> Multicore -> MulticoreGen ()
    forall a b. (a -> b) -> a -> b
    $ VName -> Multicore
    Imp.GetTaskId (SegSpace -> VName
    segFlat SegSpace
    space)
        MulticoreGen () -> MulticoreGen ()
    inISPC (MulticoreGen () -> MulticoreGen ())
    -> MulticoreGen () -> MulticoreGen ()
    forall a b. (a -> b) -> a -> b
    $ do
          -- Declare params
          [SegBinOpSlug] -> MulticoreGen ()
    genBinOpParams [SegBinOpSlug]
    slugs
          slug_local_accs <- [SegBinOpSlug] -> MulticoreGen [[VName]]
    genAccumulators [SegBinOpSlug]
    slugs
          -- Generate main reduction loop
          generateChunkLoop "SegRed" Vectorized $
            genReductionLoop RedNonComm kbody slugs slug_local_accs space
          -- Write back results
          genWriteBack slugs slug_local_accs space
      free_params <- freeParams fbody
      emit $ Imp.Op $ Imp.ParLoop "segred_stage_1" fbody free_params

    -- Codegen for a commutative reduction on scalar arrays
    -- In this case, we can generate an efficient interleaved reduction
    reductionStage1CommScalar :: ReductionStage1
    reductionStage1CommScalar :: ReductionStage1
    reductionStage1CommScalar SegSpace
    space [SegBinOpSlug]
    slugs DoSegBody
    kbody = do
      fbody <- MulticoreGen () -> MulticoreGen MCCode
    forall rep r op. ImpM rep r op () -> ImpM rep r op (Code op)
    collect (MulticoreGen () -> MulticoreGen MCCode)
    -> MulticoreGen () -> MulticoreGen MCCode
    forall a b. (a -> b) -> a -> b
    $ do
        VName -> PrimType -> MulticoreGen ()
    forall rep r op. VName -> PrimType -> ImpM rep r op ()
    dPrim_ (SegSpace -> VName
    segFlat SegSpace
    space) PrimType
    int64
        Multicore -> MulticoreGen ()
    forall op rep r. op -> ImpM rep r op ()
    sOp (Multicore -> MulticoreGen ()) -> Multicore -> MulticoreGen ()
    forall a b. (a -> b) -> a -> b
    $ VName -> Multicore
    Imp.GetTaskId (SegSpace -> VName
    segFlat SegSpace
    space)
        -- Rename lambda params in slugs to get a new set of them
        slugs' <- (SegBinOpSlug -> MulticoreGen SegBinOpSlug)
    -> [SegBinOpSlug] -> ImpM MCMem HostEnv Multicore [SegBinOpSlug]
    forall (t :: * -> *) (m :: * -> *) a b.
    (Traversable t, Monad m) =>
    (a -> m b) -> t a -> m (t b)
    forall (m :: * -> *) a b. Monad m => (a -> m b) -> [a] -> m [b]
    mapM SegBinOpSlug -> MulticoreGen SegBinOpSlug
    renameSlug [SegBinOpSlug]
    slugs
        inISPC $ do
          -- Declare one set of params uniform
          genBinOpParams slugs'
          slug_local_accs_uni <- genAccumulators slugs'
          -- Declare the other varying
          genBinOpParams slugs
          slug_local_accs <- genAccumulators slugs
          -- Generate the main reduction loop over vectors
          generateChunkLoop "SegRed" Vectorized $
            genReductionLoop RedComm kbody slugs slug_local_accs space
          -- Now reduce over those vector accumulators to get scalar results
          generateUniformizeLoop $
            genPostbodyReductionLoop slug_local_accs slugs' slug_local_accs_uni space
          -- And write back the results
          genWriteBack slugs slug_local_accs_uni space
      free_params <- freeParams fbody
      emit $ Imp.Op $ Imp.ParLoop "segred_stage_1" fbody free_params

    -- Codegen for a reduction on arrays, where the body is a perfect nested map.
    -- We vectorize just the inner map.
    reductionStage1Array :: ReductionStage1
    reductionStage1Array :: ReductionStage1
    reductionStage1Array SegSpace
    space [SegBinOpSlug]
    slugs DoSegBody
    kbody = do
      fbody <- MulticoreGen () -> MulticoreGen MCCode
    forall rep r op. ImpM rep r op () -> ImpM rep r op (Code op)
    collect (MulticoreGen () -> MulticoreGen MCCode)
    -> MulticoreGen () -> MulticoreGen MCCode
    forall a b. (a -> b) -> a -> b
    $ do
        VName -> PrimType -> MulticoreGen ()
    forall rep r op. VName -> PrimType -> ImpM rep r op ()
    dPrim_ (SegSpace -> VName
    segFlat SegSpace
    space) PrimType
    int64
        Multicore -> MulticoreGen ()
    forall op rep r. op -> ImpM rep r op ()
    sOp (Multicore -> MulticoreGen ()) -> Multicore -> MulticoreGen ()
    forall a b. (a -> b) -> a -> b
    $ VName -> Multicore
    Imp.GetTaskId (SegSpace -> VName
    segFlat SegSpace
    space)
        -- Declare params
        lparams <- MulticoreGen () -> MulticoreGen MCCode
    forall rep r op. ImpM rep r op () -> ImpM rep r op (Code op)
    collect (MulticoreGen () -> MulticoreGen MCCode)
    -> MulticoreGen () -> MulticoreGen MCCode
    forall a b. (a -> b) -> a -> b
    $ [SegBinOpSlug] -> MulticoreGen ()
    genBinOpParams [SegBinOpSlug]
    slugs
        (slug_local_accs, uniform_prebody) <- collect' $ genAccumulators slugs
        -- Put the accumulators outside of the kernel, so they are forced uniform
        emit uniform_prebody
        inISPC $ do
          -- Put the lambda params inside the kernel so they are varying
          emit lparams
          -- Generate the main reduction loop
          generateChunkLoop "SegRed" Scalar $
            genReductionLoop RedNested kbody slugs slug_local_accs space
          -- Write back results
          genWriteBack slugs slug_local_accs space
      free_params <- freeParams fbody
      emit $ Imp.Op $ Imp.ParLoop "segred_stage_1" fbody free_params

    reductionStage2 ::
      Pat LetDecMem ->
      SegSpace ->
      Imp.TExp Int32 ->
      [SegBinOpSlug] ->
      MulticoreGen ()
    reductionStage2 :: Pat LParamMem
    -> SegSpace -> TExp Int32 -> [SegBinOpSlug] -> MulticoreGen ()
    reductionStage2 Pat LParamMem
    pat SegSpace
    space TExp Int32
    nsubtasks [SegBinOpSlug]
    slugs = do
      let per_red_pes :: [[PatElem LParamMem]]
    per_red_pes = [SegBinOp MCMem] -> [PatElem LParamMem] -> [[PatElem LParamMem]]
    forall rep a. [SegBinOp rep] -> [a] -> [[a]]
    segBinOpChunks ((SegBinOpSlug -> SegBinOp MCMem)
    -> [SegBinOpSlug] -> [SegBinOp MCMem]
    forall a b. (a -> b) -> [a] -> [b]
    map SegBinOpSlug -> SegBinOp MCMem
    slugOp [SegBinOpSlug]
    slugs) ([PatElem LParamMem] -> [[PatElem LParamMem]])
    -> [PatElem LParamMem] -> [[PatElem LParamMem]]
    forall a b. (a -> b) -> a -> b
    $ Pat LParamMem -> [PatElem LParamMem]
    forall dec. Pat dec -> [PatElem dec]
    patElems Pat LParamMem
    pat
          phys_id :: TExp Int64
    phys_id = VName -> TExp Int64
    forall a. a -> TPrimExp Int64 a
    Imp.le64 (SegSpace -> VName
    segFlat SegSpace
    space)
      Text -> MulticoreGen () -> MulticoreGen ()
    forall rep r op. Text -> ImpM rep r op () -> ImpM rep r op ()
    sComment Text
    "neutral-initialise the output" (MulticoreGen () -> MulticoreGen ())
    -> MulticoreGen () -> MulticoreGen ()
    forall a b. (a -> b) -> a -> b
    $
        [(SegBinOp MCMem, [PatElem LParamMem])]
    -> ((SegBinOp MCMem, [PatElem LParamMem]) -> MulticoreGen ())
    -> MulticoreGen ()
    forall (t :: * -> *) (m :: * -> *) a b.
    (Foldable t, Monad m) =>
    t a -> (a -> m b) -> m ()
    forM_ ([SegBinOp MCMem]
    -> [[PatElem LParamMem]] -> [(SegBinOp MCMem, [PatElem LParamMem])]
    forall a b. [a] -> [b] -> [(a, b)]
    zip ((SegBinOpSlug -> SegBinOp MCMem)
    -> [SegBinOpSlug] -> [SegBinOp MCMem]
    forall a b. (a -> b) -> [a] -> [b]
    map SegBinOpSlug -> SegBinOp MCMem
    slugOp [SegBinOpSlug]
    slugs) [[PatElem LParamMem]]
    per_red_pes) (((SegBinOp MCMem, [PatElem LParamMem]) -> MulticoreGen ())
     -> MulticoreGen ())
    -> ((SegBinOp MCMem, [PatElem LParamMem]) -> MulticoreGen ())
    -> MulticoreGen ()
    forall a b. (a -> b) -> a -> b
    $ \(SegBinOp MCMem
    red, [PatElem LParamMem]
    red_res) ->
          [(PatElem LParamMem, SubExp)]
    -> ((PatElem LParamMem, SubExp) -> MulticoreGen ())
    -> MulticoreGen ()
    forall (t :: * -> *) (m :: * -> *) a b.
    (Foldable t, Monad m) =>
    t a -> (a -> m b) -> m ()
    forM_ ([PatElem LParamMem] -> [SubExp] -> [(PatElem LParamMem, SubExp)]
    forall a b. [a] -> [b] -> [(a, b)]
    zip [PatElem LParamMem]
    red_res ([SubExp] -> [(PatElem LParamMem, SubExp)])
    -> [SubExp] -> [(PatElem LParamMem, SubExp)]
    forall a b. (a -> b) -> a -> b
    $ SegBinOp MCMem -> [SubExp]
    forall rep. SegBinOp rep -> [SubExp]
    segBinOpNeutral SegBinOp MCMem
    red) (((PatElem LParamMem, SubExp) -> MulticoreGen ())
     -> MulticoreGen ())
    -> ((PatElem LParamMem, SubExp) -> MulticoreGen ())
    -> MulticoreGen ()
    forall a b. (a -> b) -> a -> b
    $ \(PatElem LParamMem
    pe, SubExp
    ne) ->
            Shape -> ([TExp Int64] -> MulticoreGen ()) -> MulticoreGen ()
    forall rep r op.
    Shape -> ([TExp Int64] -> ImpM rep r op ()) -> ImpM rep r op ()
    sLoopNest (SegBinOp MCMem -> Shape
    forall rep. SegBinOp rep -> Shape
    segBinOpShape SegBinOp MCMem
    red) (([TExp Int64] -> MulticoreGen ()) -> MulticoreGen ())
    -> ([TExp Int64] -> MulticoreGen ()) -> MulticoreGen ()
    forall a b. (a -> b) -> a -> b
    $ \[TExp Int64]
    vec_is ->
              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) [TExp Int64]
    vec_is SubExp
    ne []

      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
    $ (SegBinOpSlug -> [Param LParamMem])
    -> [SegBinOpSlug] -> [Param LParamMem]
    forall (t :: * -> *) a b. Foldable t => (a -> [b]) -> t a -> [b]
    concatMap SegBinOpSlug -> [LParam MCMem]
    SegBinOpSlug -> [Param LParamMem]
    slugParams [SegBinOpSlug]
    slugs

      String
    -> TExp Int32 -> (TExp Int32 -> MulticoreGen ()) -> MulticoreGen ()
    forall {k} (t :: k) rep r op.
    String
    -> TExp t -> (TExp t -> ImpM rep r op ()) -> ImpM rep r op ()
    sFor String
    "i" TExp Int32
    nsubtasks ((TExp Int32 -> MulticoreGen ()) -> MulticoreGen ())
    -> (TExp Int32 -> MulticoreGen ()) -> MulticoreGen ()
    forall a b. (a -> b) -> a -> b
    $ \TExp Int32
    i' -> do
        VName -> TV Int32
    forall {k} (t :: k). MkTV t => VName -> TV t
    mkTV (SegSpace -> VName
    segFlat SegSpace
    space) TV Int32 -> TExp Int32 -> MulticoreGen ()
    forall {k} (t :: k) rep r op. TV t -> TExp t -> ImpM rep r op ()
    <-- TExp Int32
    i'
        Text -> MulticoreGen () -> MulticoreGen ()
    forall rep r op. Text -> ImpM rep r op () -> ImpM rep r op ()
    sComment Text
    "Apply main thread reduction" (MulticoreGen () -> MulticoreGen ())
    -> MulticoreGen () -> MulticoreGen ()
    forall a b. (a -> b) -> a -> b
    $
          [(SegBinOpSlug, [PatElem LParamMem])]
    -> ((SegBinOpSlug, [PatElem LParamMem]) -> MulticoreGen ())
    -> MulticoreGen ()
    forall (t :: * -> *) (m :: * -> *) a b.
    (Foldable t, Monad m) =>
    t a -> (a -> m b) -> m ()
    forM_ ([SegBinOpSlug]
    -> [[PatElem LParamMem]] -> [(SegBinOpSlug, [PatElem LParamMem])]
    forall a b. [a] -> [b] -> [(a, b)]
    zip [SegBinOpSlug]
    slugs [[PatElem LParamMem]]
    per_red_pes) (((SegBinOpSlug, [PatElem LParamMem]) -> MulticoreGen ())
     -> MulticoreGen ())
    -> ((SegBinOpSlug, [PatElem LParamMem]) -> MulticoreGen ())
    -> MulticoreGen ()
    forall a b. (a -> b) -> a -> b
    $ \(SegBinOpSlug
    slug, [PatElem LParamMem]
    red_res) ->
            Shape -> ([TExp Int64] -> MulticoreGen ()) -> MulticoreGen ()
    forall rep r op.
    Shape -> ([TExp Int64] -> ImpM rep r op ()) -> ImpM rep r op ()
    sLoopNest (SegBinOpSlug -> Shape
    slugShape SegBinOpSlug
    slug) (([TExp Int64] -> MulticoreGen ()) -> MulticoreGen ())
    -> ([TExp Int64] -> MulticoreGen ()) -> MulticoreGen ()
    forall a b. (a -> b) -> a -> b
    $ \[TExp Int64]
    vec_is -> do
              Text -> MulticoreGen () -> MulticoreGen ()
    forall rep r op. Text -> ImpM rep r op () -> ImpM rep r op ()
    sComment Text
    "load acc params" (MulticoreGen () -> MulticoreGen ())
    -> MulticoreGen () -> MulticoreGen ()
    forall a b. (a -> b) -> a -> b
    $
                [(Param LParamMem, PatElem LParamMem)]
    -> ((Param LParamMem, PatElem LParamMem) -> MulticoreGen ())
    -> MulticoreGen ()
    forall (t :: * -> *) (m :: * -> *) a b.
    (Foldable t, Monad m) =>
    t a -> (a -> m b) -> m ()
    forM_ ([Param LParamMem]
    -> [PatElem LParamMem] -> [(Param LParamMem, PatElem LParamMem)]
    forall a b. [a] -> [b] -> [(a, b)]
    zip (SegBinOpSlug -> [LParam MCMem]
    accParams SegBinOpSlug
    slug) [PatElem LParamMem]
    red_res) (((Param LParamMem, PatElem LParamMem) -> MulticoreGen ())
     -> MulticoreGen ())
    -> ((Param LParamMem, PatElem LParamMem) -> MulticoreGen ())
    -> MulticoreGen ()
    forall a b. (a -> b) -> a -> b
    $ \(Param LParamMem
    p, PatElem LParamMem
    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
    p) [] (VName -> SubExp
    Var (VName -> SubExp) -> VName -> SubExp
    forall a b. (a -> b) -> a -> b
    $ PatElem LParamMem -> VName
    forall dec. PatElem dec -> VName
    patElemName PatElem LParamMem
    pe) [TExp Int64]
    vec_is
              Text -> MulticoreGen () -> MulticoreGen ()
    forall rep r op. Text -> ImpM rep r op () -> ImpM rep r op ()
    sComment Text
    "load next params" (MulticoreGen () -> MulticoreGen ())
    -> MulticoreGen () -> MulticoreGen ()
    forall a b. (a -> b) -> a -> b
    $
                [(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 (SegBinOpSlug -> [LParam MCMem]
    nextParams SegBinOpSlug
    slug) (SegBinOpSlug -> [VName]
    slugResArrs SegBinOpSlug
    slug)) (((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
    phys_id TExp Int64 -> [TExp Int64] -> [TExp Int64]
    forall a. a -> [a] -> [a]
    : [TExp Int64]
    vec_is)
              Text -> MulticoreGen () -> MulticoreGen ()
    forall rep r op. Text -> ImpM rep r op () -> ImpM rep r op ()
    sComment Text
    "red body" (MulticoreGen () -> MulticoreGen ())
    -> MulticoreGen () -> MulticoreGen ()
    forall a b. (a -> b) -> a -> b
    $
                Names -> Stms 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 -> Stms MCMem
    forall rep. Body rep -> Stms rep
    bodyStms (Body MCMem -> Stms MCMem) -> Body MCMem -> Stms MCMem
    forall a b. (a -> b) -> a -> b
    $ SegBinOpSlug -> Body MCMem
    slugBody SegBinOpSlug
    slug) (MulticoreGen () -> MulticoreGen ())
    -> MulticoreGen () -> MulticoreGen ()
    forall a b. (a -> b) -> a -> b
    $
                  [(PatElem LParamMem, SubExp)]
    -> ((PatElem LParamMem, SubExp) -> MulticoreGen ())
    -> MulticoreGen ()
    forall (t :: * -> *) (m :: * -> *) a b.
    (Foldable t, Monad m) =>
    t a -> (a -> m b) -> m ()
    forM_ ([PatElem LParamMem] -> [SubExp] -> [(PatElem LParamMem, SubExp)]
    forall a b. [a] -> [b] -> [(a, b)]
    zip [PatElem LParamMem]
    red_res ([SubExp] -> [(PatElem LParamMem, SubExp)])
    -> [SubExp] -> [(PatElem LParamMem, 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 MCMem -> [SubExpRes]
    forall rep. Body rep -> [SubExpRes]
    bodyResult (Body MCMem -> [SubExpRes]) -> Body MCMem -> [SubExpRes]
    forall a b. (a -> b) -> a -> b
    $ SegBinOpSlug -> Body MCMem
    slugBody SegBinOpSlug
    slug) (((PatElem LParamMem, SubExp) -> MulticoreGen ())
     -> MulticoreGen ())
    -> ((PatElem LParamMem, SubExp) -> MulticoreGen ())
    -> MulticoreGen ()
    forall a b. (a -> b) -> a -> b
    $
                    \(PatElem LParamMem
    pe, SubExp
    se') -> 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) [TExp Int64]
    vec_is SubExp
    se' []

    -- Each thread reduces over the number of segments
    -- each of which is done sequentially
    -- Maybe we should select the work of the inner loop
    -- based on n_segments and dimensions etc.
    segmentedReduction ::
      Pat LetDecMem ->
      SegSpace ->
      [SegBinOp MCMem] ->
      DoSegBody ->
      MulticoreGen Imp.MCCode
    segmentedReduction :: Pat LParamMem
    -> SegSpace -> [SegBinOp MCMem] -> DoSegBody -> MulticoreGen MCCode
    segmentedReduction Pat LParamMem
    pat SegSpace
    space [SegBinOp MCMem]
    reds DoSegBody
    kbody =
      MulticoreGen () -> MulticoreGen MCCode
    forall rep r op. ImpM rep r op () -> ImpM rep r op (Code op)
    collect (MulticoreGen () -> MulticoreGen MCCode)
    -> MulticoreGen () -> MulticoreGen MCCode
    forall a b. (a -> b) -> a -> b
    $ do
        body <- Pat LParamMem
    -> SegSpace -> [SegBinOp MCMem] -> DoSegBody -> MulticoreGen MCCode
    compileSegRedBody Pat LParamMem
    pat SegSpace
    space [SegBinOp MCMem]
    reds DoSegBody
    kbody
        free_params <- freeParams body
        emit $ Imp.Op $ Imp.ParLoop "segmented_segred" body free_params

    -- Currently, this is only used as part of SegHist calculations, never alone.
    compileSegRedBody ::
      Pat LetDecMem ->
      SegSpace ->
      [SegBinOp MCMem] ->
      DoSegBody ->
      MulticoreGen Imp.MCCode
    compileSegRedBody :: Pat LParamMem
    -> SegSpace -> [SegBinOp MCMem] -> DoSegBody -> MulticoreGen MCCode
    compileSegRedBody Pat LParamMem
    pat SegSpace
    space [SegBinOp MCMem]
    reds DoSegBody
    kbody = do
      let ([VName]
    is, [SubExp]
    ns) = [(VName, SubExp)] -> ([VName], [SubExp])
    forall a b. [(a, b)] -> ([a], [b])
    unzip ([(VName, SubExp)] -> ([VName], [SubExp]))
    -> [(VName, SubExp)] -> ([VName], [SubExp])
    forall a b. (a -> b) -> a -> b
    $ SegSpace -> [(VName, SubExp)]
    unSegSpace SegSpace
    space
          ns_64 :: [TExp Int64]
    ns_64 = (SubExp -> TExp Int64) -> [SubExp] -> [TExp Int64]
    forall a b. (a -> b) -> [a] -> [b]
    map SubExp -> TExp Int64
    pe64 [SubExp]
    ns
          inner_bound :: TExp Int64
    inner_bound = [TExp Int64] -> TExp Int64
    forall a. HasCallStack => [a] -> a
    last [TExp Int64]
    ns_64
      VName -> PrimType -> MulticoreGen ()
    forall rep r op. VName -> PrimType -> ImpM rep r op ()
    dPrim_ (SegSpace -> VName
    segFlat SegSpace
    space) PrimType
    int64
      Multicore -> MulticoreGen ()
    forall op rep r. op -> ImpM rep r op ()
    sOp (Multicore -> MulticoreGen ()) -> Multicore -> MulticoreGen ()
    forall a b. (a -> b) -> a -> b
    $ VName -> Multicore
    Imp.GetTaskId (SegSpace -> VName
    segFlat SegSpace
    space)

      let per_red_pes :: [[PatElem LParamMem]]
    per_red_pes = [SegBinOp MCMem] -> [PatElem LParamMem] -> [[PatElem LParamMem]]
    forall rep a. [SegBinOp rep] -> [a] -> [[a]]
    segBinOpChunks [SegBinOp MCMem]
    reds ([PatElem LParamMem] -> [[PatElem LParamMem]])
    -> [PatElem LParamMem] -> [[PatElem LParamMem]]
    forall a b. (a -> b) -> a -> b
    $ Pat LParamMem -> [PatElem LParamMem]
    forall dec. Pat dec -> [PatElem dec]
    patElems Pat LParamMem
    pat
      -- Perform sequential reduce on inner most dimension
      MulticoreGen () -> MulticoreGen MCCode
    forall rep r op. ImpM rep r op () -> ImpM rep r op (Code op)
    collect (MulticoreGen () -> MulticoreGen MCCode)
    -> (MulticoreGen () -> MulticoreGen ())
    -> MulticoreGen ()
    -> MulticoreGen MCCode
    forall b c a. (b -> c) -> (a -> b) -> a -> c
    . MulticoreGen () -> MulticoreGen ()
    inISPC (MulticoreGen () -> MulticoreGen MCCode)
    -> MulticoreGen () -> MulticoreGen MCCode
    forall a b. (a -> b) -> a -> b
    $
        String
    -> ChunkLoopVectorization
    -> (TExp Int64 -> MulticoreGen ())
    -> MulticoreGen ()
    generateChunkLoop String
    "SegRed" ChunkLoopVectorization
    Vectorized ((TExp Int64 -> MulticoreGen ()) -> MulticoreGen ())
    -> (TExp Int64 -> MulticoreGen ()) -> MulticoreGen ()
    forall a b. (a -> b) -> a -> b
    $ \TExp Int64
    n_segments -> do
          flat_idx <- String -> TExp Int64 -> ImpM MCMem HostEnv Multicore (TExp Int64)
    forall {k} (t :: k) rep r op.
    String -> TExp t -> ImpM rep r op (TExp t)
    dPrimVE String
    "flat_idx" (TExp Int64 -> ImpM MCMem HostEnv Multicore (TExp Int64))
    -> TExp Int64 -> ImpM MCMem HostEnv Multicore (TExp Int64)
    forall a b. (a -> b) -> a -> b
    $ TExp Int64
    n_segments TExp Int64 -> TExp Int64 -> TExp Int64
    forall a. Num a => a -> a -> a
    * TExp Int64
    inner_bound
          zipWithM_ dPrimV_ is $ unflattenIndex ns_64 flat_idx
          sComment "neutral-initialise the accumulators" $
            forM_ (zip per_red_pes reds) $ \([PatElem LParamMem]
    pes, SegBinOp MCMem
    red) ->
              [(PatElem LParamMem, SubExp)]
    -> ((PatElem LParamMem, SubExp) -> MulticoreGen ())
    -> MulticoreGen ()
    forall (t :: * -> *) (m :: * -> *) a b.
    (Foldable t, Monad m) =>
    t a -> (a -> m b) -> m ()
    forM_ ([PatElem LParamMem] -> [SubExp] -> [(PatElem LParamMem, SubExp)]
    forall a b. [a] -> [b] -> [(a, b)]
    zip [PatElem LParamMem]
    pes (SegBinOp MCMem -> [SubExp]
    forall rep. SegBinOp rep -> [SubExp]
    segBinOpNeutral SegBinOp MCMem
    red)) (((PatElem LParamMem, SubExp) -> MulticoreGen ())
     -> MulticoreGen ())
    -> ((PatElem LParamMem, SubExp) -> MulticoreGen ())
    -> MulticoreGen ()
    forall a b. (a -> b) -> a -> b
    $ \(PatElem LParamMem
    pe, SubExp
    ne) ->
                Shape -> ([TExp Int64] -> MulticoreGen ()) -> MulticoreGen ()
    forall rep r op.
    Shape -> ([TExp Int64] -> ImpM rep r op ()) -> ImpM rep r op ()
    sLoopNest (SegBinOp MCMem -> Shape
    forall rep. SegBinOp rep -> Shape
    segBinOpShape SegBinOp MCMem
    red) (([TExp Int64] -> MulticoreGen ()) -> MulticoreGen ())
    -> ([TExp Int64] -> MulticoreGen ()) -> MulticoreGen ()
    forall a b. (a -> b) -> a -> b
    $ \[TExp Int64]
    vec_is ->
                  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) ((VName -> TExp Int64) -> [VName] -> [TExp Int64]
    forall a b. (a -> b) -> [a] -> [b]
    map VName -> TExp Int64
    forall a. a -> TPrimExp Int64 a
    Imp.le64 ([VName] -> [VName]
    forall a. HasCallStack => [a] -> [a]
    init [VName]
    is) [TExp Int64] -> [TExp Int64] -> [TExp Int64]
    forall a. [a] -> [a] -> [a]
    ++ [TExp Int64]
    vec_is) SubExp
    ne []

          sComment "main body" $ do
            dScope Nothing $ scopeOfLParams $ concatMap (lambdaParams . segBinOpLambda) reds
            sFor "i" inner_bound $ \TExp Int64
    i -> do
              (TV Int64 -> TExp Int64 -> MulticoreGen ())
    -> [TV Int64] -> [TExp Int64] -> MulticoreGen ()
    forall (m :: * -> *) a b c.
    Applicative m =>
    (a -> b -> m c) -> [a] -> [b] -> m ()
    zipWithM_
                TV Int64 -> TExp Int64 -> MulticoreGen ()
    forall {k} (t :: k) rep r op. TV t -> TExp t -> ImpM rep r op ()
    (<--)
                ((VName -> TV Int64) -> [VName] -> [TV Int64]
    forall a b. (a -> b) -> [a] -> [b]
    map VName -> TV Int64
    forall {k} (t :: k). MkTV t => VName -> TV t
    mkTV ([VName] -> [TV Int64]) -> [VName] -> [TV Int64]
    forall a b. (a -> b) -> a -> b
    $ [VName] -> [VName]
    forall a. HasCallStack => [a] -> [a]
    init [VName]
    is)
                ([TExp Int64] -> TExp Int64 -> [TExp Int64]
    forall num. IntegralExp num => [num] -> num -> [num]
    unflattenIndex ([TExp Int64] -> [TExp Int64]
    forall a. HasCallStack => [a] -> [a]
    init [TExp Int64]
    ns_64) (TExp Int64 -> TExp Int64
    forall {k} (t :: k) v. IntExp t => TPrimExp t v -> TPrimExp Int64 v
    sExt64 TExp Int64
    n_segments))
              VName -> TExp Int64 -> MulticoreGen ()
    forall {k} (t :: k) rep r op. VName -> TExp t -> ImpM rep r op ()
    dPrimV_ ([VName] -> VName
    forall a. HasCallStack => [a] -> a
    last [VName]
    is) TExp Int64
    i
              DoSegBody
    kbody DoSegBody -> DoSegBody
    forall a b. (a -> b) -> a -> b
    $ \[[(SubExp, [TExp Int64])]]
    red_res' -> do
                [([PatElem LParamMem], SegBinOp MCMem, [(SubExp, [TExp Int64])])]
    -> (([PatElem LParamMem], SegBinOp MCMem, [(SubExp, [TExp Int64])])
        -> MulticoreGen ())
    -> MulticoreGen ()
    forall (t :: * -> *) (m :: * -> *) a b.
    (Foldable t, Monad m) =>
    t a -> (a -> m b) -> m ()
    forM_ ([[PatElem LParamMem]]
    -> [SegBinOp MCMem]
    -> [[(SubExp, [TExp Int64])]]
    -> [([PatElem LParamMem], SegBinOp MCMem,
         [(SubExp, [TExp Int64])])]
    forall a b c. [a] -> [b] -> [c] -> [(a, b, c)]
    zip3 [[PatElem LParamMem]]
    per_red_pes [SegBinOp MCMem]
    reds [[(SubExp, [TExp Int64])]]
    red_res') ((([PatElem LParamMem], SegBinOp MCMem, [(SubExp, [TExp Int64])])
      -> MulticoreGen ())
     -> MulticoreGen ())
    -> (([PatElem LParamMem], SegBinOp MCMem, [(SubExp, [TExp Int64])])
        -> MulticoreGen ())
    -> MulticoreGen ()
    forall a b. (a -> b) -> a -> b
    $ \([PatElem LParamMem]
    pes, SegBinOp MCMem
    red, [(SubExp, [TExp Int64])]
    res') ->
                  Shape -> ([TExp Int64] -> MulticoreGen ()) -> MulticoreGen ()
    forall rep r op.
    Shape -> ([TExp Int64] -> ImpM rep r op ()) -> ImpM rep r op ()
    sLoopNest (SegBinOp MCMem -> Shape
    forall rep. SegBinOp rep -> Shape
    segBinOpShape SegBinOp MCMem
    red) (([TExp Int64] -> MulticoreGen ()) -> MulticoreGen ())
    -> ([TExp Int64] -> MulticoreGen ()) -> MulticoreGen ()
    forall a b. (a -> b) -> a -> b
    $ \[TExp Int64]
    vec_is -> do
                    Text -> MulticoreGen () -> MulticoreGen ()
    forall rep r op. Text -> ImpM rep r op () -> ImpM rep r op ()
    sComment Text
    "load accum" (MulticoreGen () -> MulticoreGen ())
    -> MulticoreGen () -> MulticoreGen ()
    forall a b. (a -> b) -> a -> b
    $ do
                      let acc_params :: [Param LParamMem]
    acc_params = 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
    red)) ([Param LParamMem] -> [Param LParamMem])
    -> [Param LParamMem] -> [Param LParamMem]
    forall a b. (a -> b) -> a -> b
    $ (Lambda MCMem -> [LParam MCMem]
    Lambda MCMem -> [Param LParamMem]
    forall rep. Lambda rep -> [LParam rep]
    lambdaParams (Lambda MCMem -> [Param LParamMem])
    -> (SegBinOp MCMem -> Lambda MCMem)
    -> SegBinOp MCMem
    -> [Param LParamMem]
    forall b c a. (b -> c) -> (a -> b) -> a -> c
    . SegBinOp MCMem -> Lambda MCMem
    forall rep. SegBinOp rep -> Lambda rep
    segBinOpLambda) SegBinOp MCMem
    red
                      [(Param LParamMem, PatElem LParamMem)]
    -> ((Param LParamMem, PatElem LParamMem) -> MulticoreGen ())
    -> MulticoreGen ()
    forall (t :: * -> *) (m :: * -> *) a b.
    (Foldable t, Monad m) =>
    t a -> (a -> m b) -> m ()
    forM_ ([Param LParamMem]
    -> [PatElem LParamMem] -> [(Param LParamMem, PatElem LParamMem)]
    forall a b. [a] -> [b] -> [(a, b)]
    zip [Param LParamMem]
    acc_params [PatElem LParamMem]
    pes) (((Param LParamMem, PatElem LParamMem) -> MulticoreGen ())
     -> MulticoreGen ())
    -> ((Param LParamMem, PatElem LParamMem) -> MulticoreGen ())
    -> MulticoreGen ()
    forall a b. (a -> b) -> a -> b
    $ \(Param LParamMem
    p, PatElem LParamMem
    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
    p) [] (VName -> SubExp
    Var (VName -> SubExp) -> VName -> SubExp
    forall a b. (a -> b) -> a -> b
    $ PatElem LParamMem -> VName
    forall dec. PatElem dec -> VName
    patElemName PatElem LParamMem
    pe) ((VName -> TExp Int64) -> [VName] -> [TExp Int64]
    forall a b. (a -> b) -> [a] -> [b]
    map VName -> TExp Int64
    forall a. a -> TPrimExp Int64 a
    Imp.le64 ([VName] -> [VName]
    forall a. HasCallStack => [a] -> [a]
    init [VName]
    is) [TExp Int64] -> [TExp Int64] -> [TExp Int64]
    forall a. [a] -> [a] -> [a]
    ++ [TExp Int64]
    vec_is)

                    Text -> MulticoreGen () -> MulticoreGen ()
    forall rep r op. Text -> ImpM rep r op () -> ImpM rep r op ()
    sComment Text
    "load new val" (MulticoreGen () -> MulticoreGen ())
    -> MulticoreGen () -> MulticoreGen ()
    forall a b. (a -> b) -> a -> b
    $ do
                      let next_params :: [Param LParamMem]
    next_params = 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
    red)) ([Param LParamMem] -> [Param LParamMem])
    -> [Param LParamMem] -> [Param LParamMem]
    forall a b. (a -> b) -> a -> b
    $ (Lambda MCMem -> [LParam MCMem]
    Lambda MCMem -> [Param LParamMem]
    forall rep. Lambda rep -> [LParam rep]
    lambdaParams (Lambda MCMem -> [Param LParamMem])
    -> (SegBinOp MCMem -> Lambda MCMem)
    -> SegBinOp MCMem
    -> [Param LParamMem]
    forall b c a. (b -> c) -> (a -> b) -> a -> c
    . SegBinOp MCMem -> Lambda MCMem
    forall rep. SegBinOp rep -> Lambda rep
    segBinOpLambda) SegBinOp MCMem
    red
                      [(Param LParamMem, (SubExp, [TExp Int64]))]
    -> ((Param LParamMem, (SubExp, [TExp Int64])) -> MulticoreGen ())
    -> MulticoreGen ()
    forall (t :: * -> *) (m :: * -> *) a b.
    (Foldable t, Monad m) =>
    t a -> (a -> m b) -> m ()
    forM_ ([Param LParamMem]
    -> [(SubExp, [TExp Int64])]
    -> [(Param LParamMem, (SubExp, [TExp Int64]))]
    forall a b. [a] -> [b] -> [(a, b)]
    zip [Param LParamMem]
    next_params [(SubExp, [TExp Int64])]
    res') (((Param LParamMem, (SubExp, [TExp Int64])) -> MulticoreGen ())
     -> MulticoreGen ())
    -> ((Param LParamMem, (SubExp, [TExp Int64])) -> MulticoreGen ())
    -> MulticoreGen ()
    forall a b. (a -> b) -> a -> b
    $ \(Param LParamMem
    p, (SubExp
    res, [TExp Int64]
    res_is)) ->
                        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) [] SubExp
    res ([TExp Int64]
    res_is [TExp Int64] -> [TExp Int64] -> [TExp Int64]
    forall a. [a] -> [a] -> [a]
    ++ [TExp Int64]
    vec_is)

                    Text -> MulticoreGen () -> MulticoreGen ()
    forall rep r op. Text -> ImpM rep r op () -> ImpM rep r op ()
    sComment Text
    "apply reduction" (MulticoreGen () -> MulticoreGen ())
    -> MulticoreGen () -> MulticoreGen ()
    forall a b. (a -> b) -> a -> b
    $ do
                      let lbody :: Body MCMem
    lbody = (Lambda MCMem -> Body MCMem
    forall rep. Lambda rep -> Body rep
    lambdaBody (Lambda MCMem -> Body MCMem)
    -> (SegBinOp MCMem -> Lambda MCMem) -> SegBinOp MCMem -> Body MCMem
    forall b c a. (b -> c) -> (a -> b) -> a -> c
    . SegBinOp MCMem -> Lambda MCMem
    forall rep. SegBinOp rep -> Lambda rep
    segBinOpLambda) SegBinOp MCMem
    red
                      Names -> Stms 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 -> Stms MCMem
    forall rep. Body rep -> Stms rep
    bodyStms Body MCMem
    lbody) (MulticoreGen () -> MulticoreGen ())
    -> MulticoreGen () -> MulticoreGen ()
    forall a b. (a -> b) -> a -> b
    $
                        Text -> MulticoreGen () -> MulticoreGen ()
    forall rep r op. Text -> ImpM rep r op () -> ImpM rep r op ()
    sComment Text
    "write back to res" (MulticoreGen () -> MulticoreGen ())
    -> MulticoreGen () -> MulticoreGen ()
    forall a b. (a -> b) -> a -> b
    $
                          [(PatElem LParamMem, SubExp)]
    -> ((PatElem LParamMem, SubExp) -> MulticoreGen ())
    -> MulticoreGen ()
    forall (t :: * -> *) (m :: * -> *) a b.
    (Foldable t, Monad m) =>
    t a -> (a -> m b) -> m ()
    forM_ ([PatElem LParamMem] -> [SubExp] -> [(PatElem LParamMem, SubExp)]
    forall a b. [a] -> [b] -> [(a, b)]
    zip [PatElem LParamMem]
    pes ([SubExp] -> [(PatElem LParamMem, SubExp)])
    -> [SubExp] -> [(PatElem LParamMem, 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 MCMem -> [SubExpRes]
    forall rep. Body rep -> [SubExpRes]
    bodyResult Body MCMem
    lbody) (((PatElem LParamMem, SubExp) -> MulticoreGen ())
     -> MulticoreGen ())
    -> ((PatElem LParamMem, SubExp) -> MulticoreGen ())
    -> MulticoreGen ()
    forall a b. (a -> b) -> a -> b
    $
                            \(PatElem LParamMem
    pe, SubExp
    se') -> 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) ((VName -> TExp Int64) -> [VName] -> [TExp Int64]
    forall a b. (a -> b) -> [a] -> [b]
    map VName -> TExp Int64
    forall a. a -> TPrimExp Int64 a
    Imp.le64 ([VName] -> [VName]
    forall a. HasCallStack => [a] -> [a]
    init [VName]
    is) [TExp Int64] -> [TExp Int64] -> [TExp Int64]
    forall a. [a] -> [a] -> [a]
    ++ [TExp Int64]
    vec_is) SubExp
    se' []