    {-# LANGUAGE TypeFamilies #-}

    -- | Turn certain uses of accumulators into SegHists.
    module Futhark.Optimise.HistAccs (histAccsGPU) where

    import Control.Monad
    import Control.Monad.Reader
    import Control.Monad.State
    import Data.Map.Strict qualified as M
    import Futhark.IR.GPU
    import Futhark.MonadFreshNames
    import Futhark.Pass
    import Futhark.Tools
    import Futhark.Transform.Rename
    import Prelude hiding (quot)

    -- | A mapping from accumulator variables to their source.
    type Accs rep = M.Map VName (WithAccInput rep)

    type OptM = ReaderT (Scope GPU) (State VNameSource)

    optimiseBody :: Accs GPU -> Body GPU -> OptM (Body GPU)
    optimiseBody :: Map VName (WithAccInput GPU) -> Body GPU -> OptM (Body GPU)
    optimiseBody Map VName (WithAccInput GPU)
    accs Body GPU
    body = Stms GPU -> Result -> Body GPU
    forall rep. Buildable rep => Stms rep -> Result -> Body rep
    mkBody (Stms GPU -> Result -> Body GPU)
    -> ReaderT (Scope GPU) (State VNameSource) (Stms GPU)
    -> ReaderT (Scope GPU) (State VNameSource) (Result -> Body GPU)
    forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
    <$> Map VName (WithAccInput GPU)
    -> Stms GPU -> ReaderT (Scope GPU) (State VNameSource) (Stms GPU)
    optimiseStms Map VName (WithAccInput GPU)
    accs (Body GPU -> Stms GPU
    forall rep. Body rep -> Stms rep
    bodyStms Body GPU
    body) ReaderT (Scope GPU) (State VNameSource) (Result -> Body GPU)
    -> ReaderT (Scope GPU) (State VNameSource) Result
    -> OptM (Body GPU)
    forall a b.
    ReaderT (Scope GPU) (State VNameSource) (a -> b)
    -> ReaderT (Scope GPU) (State VNameSource) a
    -> ReaderT (Scope GPU) (State VNameSource) b
    forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
    <*> Result -> ReaderT (Scope GPU) (State VNameSource) Result
    forall a. a -> ReaderT (Scope GPU) (State VNameSource) a
    forall (f :: * -> *) a. Applicative f => a -> f a
    pure (Body GPU -> Result
    forall rep. Body rep -> Result
    bodyResult Body GPU
    body)

    optimiseExp :: Accs GPU -> Exp GPU -> OptM (Exp GPU)
    optimiseExp :: Map VName (WithAccInput GPU) -> Exp GPU -> OptM (Exp GPU)
    optimiseExp Map VName (WithAccInput GPU)
    accs = Mapper GPU GPU (ReaderT (Scope GPU) (State VNameSource))
    -> Exp GPU -> OptM (Exp GPU)
    forall (m :: * -> *) frep trep.
    Monad m =>
    Mapper frep trep m -> Exp frep -> m (Exp trep)
    mapExpM Mapper GPU GPU (ReaderT (Scope GPU) (State VNameSource))
    mapper
      where
        mapper :: Mapper GPU GPU (ReaderT (Scope GPU) (State VNameSource))
    mapper =
          Mapper GPU GPU (ReaderT (Scope GPU) (State VNameSource))
    forall rep (m :: * -> *). Monad m => Mapper rep rep m
    identityMapper
            { mapOnBody = \Scope GPU
    scope Body GPU
    body -> Scope GPU -> OptM (Body GPU) -> OptM (Body GPU)
    forall a.
    Scope GPU
    -> ReaderT (Scope GPU) (State VNameSource) a
    -> ReaderT (Scope GPU) (State VNameSource) a
    forall rep (m :: * -> *) a.
    LocalScope rep m =>
    Scope rep -> m a -> m a
    localScope Scope GPU
    scope (OptM (Body GPU) -> OptM (Body GPU))
    -> OptM (Body GPU) -> OptM (Body GPU)
    forall a b. (a -> b) -> a -> b
    $ Map VName (WithAccInput GPU) -> Body GPU -> OptM (Body GPU)
    optimiseBody Map VName (WithAccInput GPU)
    accs Body GPU
    body
            }

    extractUpdate ::
      Accs rep ->
      VName ->
      Stms rep ->
      Maybe ((WithAccInput rep, VName, [SubExp], [SubExp]), Stms rep)
    extractUpdate :: forall rep.
    Accs rep
    -> VName
    -> Stms rep
    -> Maybe ((WithAccInput rep, VName, [SubExp], [SubExp]), Stms rep)
    extractUpdate Accs rep
    accs VName
    v Stms rep
    stms = do
      (Stm rep
    stm, Stms rep
    stms') <- Stms rep -> Maybe (Stm rep, Stms rep)
    forall rep. Stms rep -> Maybe (Stm rep, Stms rep)
    stmsHead Stms rep
    stms
      case Stm rep
    stm of
        Let (Pat [PatElem VName
    pe_v LetDec rep
    _]) StmAux (ExpDec rep)
    _ (BasicOp (UpdateAcc Safety
    _ VName
    acc [SubExp]
    is [SubExp]
    vs))
          | VName
    pe_v VName -> VName -> Bool
    forall a. Eq a => a -> a -> Bool
    == VName
    v -> do
              WithAccInput rep
    acc_input <- VName -> Accs rep -> Maybe (WithAccInput rep)
    forall k a. Ord k => k -> Map k a -> Maybe a
    M.lookup VName
    acc Accs rep
    accs
              ((WithAccInput rep, VName, [SubExp], [SubExp]), Stms rep)
    -> Maybe ((WithAccInput rep, VName, [SubExp], [SubExp]), Stms rep)
    forall a. a -> Maybe a
    Just ((WithAccInput rep
    acc_input, VName
    acc, [SubExp]
    is, [SubExp]
    vs), Stms rep
    stms')
        Stm rep
    _ -> do
          ((WithAccInput rep, VName, [SubExp], [SubExp])
    x, Stms rep
    stms'') <- Accs rep
    -> VName
    -> Stms rep
    -> Maybe ((WithAccInput rep, VName, [SubExp], [SubExp]), Stms rep)
    forall rep.
    Accs rep
    -> VName
    -> Stms rep
    -> Maybe ((WithAccInput rep, VName, [SubExp], [SubExp]), Stms rep)
    extractUpdate Accs rep
    accs VName
    v Stms rep
    stms'
          ((WithAccInput rep, VName, [SubExp], [SubExp]), Stms rep)
    -> Maybe ((WithAccInput rep, VName, [SubExp], [SubExp]), Stms rep)
    forall a. a -> Maybe a
    forall (f :: * -> *) a. Applicative f => a -> f a
    pure ((WithAccInput rep, VName, [SubExp], [SubExp])
    x, Stm rep -> Stms rep
    forall rep. Stm rep -> Stms rep
    oneStm Stm rep
    stm Stms rep -> Stms rep -> Stms rep
    forall a. Semigroup a => a -> a -> a
    <> Stms rep
    stms'')

    mkHistBody :: Accs GPU -> KernelBody GPU -> Maybe (KernelBody GPU, WithAccInput GPU, VName)
    mkHistBody :: Map VName (WithAccInput GPU)
    -> KernelBody GPU
    -> Maybe (KernelBody GPU, WithAccInput GPU, VName)
    mkHistBody Map VName (WithAccInput GPU)
    accs (KernelBody () Stms GPU
    stms [Returns ResultManifest
    rm Certs
    cs (Var VName
    v)]) = do
      ((WithAccInput GPU
    acc_input, VName
    acc, [SubExp]
    is, [SubExp]
    vs), Stms GPU
    stms') <- Map VName (WithAccInput GPU)
    -> VName
    -> Stms GPU
    -> Maybe ((WithAccInput GPU, VName, [SubExp], [SubExp]), Stms GPU)
    forall rep.
    Accs rep
    -> VName
    -> Stms rep
    -> Maybe ((WithAccInput rep, VName, [SubExp], [SubExp]), Stms rep)
    extractUpdate Map VName (WithAccInput GPU)
    accs VName
    v Stms GPU
    stms
      (KernelBody GPU, WithAccInput GPU, VName)
    -> Maybe (KernelBody GPU, WithAccInput GPU, VName)
    forall a. a -> Maybe a
    forall (f :: * -> *) a. Applicative f => a -> f a
    pure
        ( BodyDec GPU -> Stms GPU -> [KernelResult] -> KernelBody GPU
    forall rep.
    BodyDec rep -> Stms rep -> [KernelResult] -> KernelBody rep
    KernelBody () Stms GPU
    stms' ([KernelResult] -> KernelBody GPU)
    -> [KernelResult] -> KernelBody GPU
    forall a b. (a -> b) -> a -> b
    $ (SubExp -> KernelResult) -> [SubExp] -> [KernelResult]
    forall a b. (a -> b) -> [a] -> [b]
    map (ResultManifest -> Certs -> SubExp -> KernelResult
    Returns ResultManifest
    rm Certs
    cs) [SubExp]
    is [KernelResult] -> [KernelResult] -> [KernelResult]
    forall a. [a] -> [a] -> [a]
    ++ (SubExp -> KernelResult) -> [SubExp] -> [KernelResult]
    forall a b. (a -> b) -> [a] -> [b]
    map (ResultManifest -> Certs -> SubExp -> KernelResult
    Returns ResultManifest
    rm Certs
    cs) [SubExp]
    vs,
          WithAccInput GPU
    acc_input,
          VName
    acc
        )
    mkHistBody Map VName (WithAccInput GPU)
    _ KernelBody GPU
    _ = Maybe (KernelBody GPU, WithAccInput GPU, VName)
    forall a. Maybe a
    Nothing

    withAccLamToHistLam :: (MonadFreshNames m) => Shape -> Lambda GPU -> m (Lambda GPU)
    withAccLamToHistLam :: forall (m :: * -> *).
    MonadFreshNames m =>
    Shape -> Lambda GPU -> m (Lambda GPU)
    withAccLamToHistLam Shape
    shape Lambda GPU
    lam =
      Lambda GPU -> m (Lambda GPU)
    forall rep (m :: * -> *).
    (Renameable rep, MonadFreshNames m) =>
    Lambda rep -> m (Lambda rep)
    renameLambda (Lambda GPU -> m (Lambda GPU)) -> Lambda GPU -> m (Lambda GPU)
    forall a b. (a -> b) -> a -> b
    $ Lambda GPU
    lam {lambdaParams = drop (shapeRank shape) (lambdaParams lam)}

    addArrsToAcc ::
      (MonadBuilder m, Rep m ~ GPU) =>
      SegLevel ->
      Shape ->
      [VName] ->
      VName ->
      m (Exp GPU)
    addArrsToAcc :: forall (m :: * -> *).
    (MonadBuilder m, Rep m ~ GPU) =>
    SegLevel -> Shape -> [VName] -> VName -> m (Exp GPU)
    addArrsToAcc SegLevel
    lvl Shape
    shape [VName]
    arrs VName
    acc = do
      VName
    flat <- String -> m VName
    forall (m :: * -> *). MonadFreshNames m => String -> m VName
    newVName String
    "phys_tid"
      [VName]
    gtids <- Int -> m VName -> m [VName]
    forall (m :: * -> *) a. Applicative m => Int -> m a -> m [a]
    replicateM (Shape -> Int
    forall a. ArrayShape a => a -> Int
    shapeRank Shape
    shape) (String -> m VName
    forall (m :: * -> *). MonadFreshNames m => String -> m VName
    newVName String
    "gtid")
      let space :: SegSpace
    space = VName -> [(VName, SubExp)] -> SegSpace
    SegSpace VName
    flat ([(VName, SubExp)] -> SegSpace) -> [(VName, SubExp)] -> SegSpace
    forall a b. (a -> b) -> a -> b
    $ [VName] -> [SubExp] -> [(VName, SubExp)]
    forall a b. [a] -> [b] -> [(a, b)]
    zip [VName]
    gtids ([SubExp] -> [(VName, SubExp)]) -> [SubExp] -> [(VName, SubExp)]
    forall a b. (a -> b) -> a -> b
    $ Shape -> [SubExp]
    forall d. ShapeBase d -> [d]
    shapeDims Shape
    shape

      (VName
    acc', Stms GPU
    stms) <- Scope GPU -> m (VName, Stms GPU) -> m (VName, Stms GPU)
    forall a. Scope GPU -> m a -> m a
    forall rep (m :: * -> *) a.
    LocalScope rep m =>
    Scope rep -> m a -> m a
    localScope (SegSpace -> Scope GPU
    forall rep. SegSpace -> Scope rep
    scopeOfSegSpace SegSpace
    space) (m (VName, Stms GPU) -> m (VName, Stms GPU))
    -> (m VName -> m (VName, Stms GPU))
    -> m VName
    -> m (VName, Stms GPU)
    forall b c a. (b -> c) -> (a -> b) -> a -> c
    . m VName -> m (VName, Stms (Rep m))
    m VName -> m (VName, Stms GPU)
    forall a. m a -> m (a, Stms (Rep m))
    forall (m :: * -> *) a.
    MonadBuilder m =>
    m a -> m (a, Stms (Rep m))
    collectStms (m VName -> m (VName, Stms GPU)) -> m VName -> m (VName, Stms GPU)
    forall a b. (a -> b) -> a -> b
    $ do
        [SubExp]
    vs <- [VName] -> (VName -> m SubExp) -> m [SubExp]
    forall (t :: * -> *) (m :: * -> *) a b.
    (Traversable t, Monad m) =>
    t a -> (a -> m b) -> m (t b)
    forM [VName]
    arrs ((VName -> m SubExp) -> m [SubExp])
    -> (VName -> m SubExp) -> m [SubExp]
    forall a b. (a -> b) -> a -> b
    $ \VName
    arr -> do
          Type
    arr_t <- VName -> m Type
    forall rep (m :: * -> *). HasScope rep m => VName -> m Type
    lookupType VName
    arr
          String -> Exp (Rep m) -> m SubExp
    forall (m :: * -> *).
    MonadBuilder m =>
    String -> Exp (Rep m) -> m SubExp
    letSubExp (VName -> String
    baseString VName
    arr String -> String -> String
    forall a. Semigroup a => a -> a -> a
    <> String
    "_elem") (Exp (Rep m) -> m SubExp) -> Exp (Rep m) -> m SubExp
    forall a b. (a -> b) -> a -> b
    $
            BasicOp -> Exp (Rep m)
    forall rep. BasicOp -> Exp rep
    BasicOp (BasicOp -> Exp (Rep m)) -> BasicOp -> Exp (Rep m)
    forall a b. (a -> b) -> a -> b
    $
              VName -> Slice SubExp -> BasicOp
    Index VName
    arr (Slice SubExp -> BasicOp) -> Slice SubExp -> BasicOp
    forall a b. (a -> b) -> a -> b
    $
                Type -> [DimIndex SubExp] -> Slice SubExp
    fullSlice Type
    arr_t ([DimIndex SubExp] -> Slice SubExp)
    -> [DimIndex SubExp] -> Slice SubExp
    forall a b. (a -> b) -> a -> b
    $
                  (VName -> DimIndex SubExp) -> [VName] -> [DimIndex SubExp]
    forall a b. (a -> b) -> [a] -> [b]
    map (SubExp -> DimIndex SubExp
    forall d. d -> DimIndex d
    DimFix (SubExp -> DimIndex SubExp)
    -> (VName -> SubExp) -> VName -> DimIndex SubExp
    forall b c a. (b -> c) -> (a -> b) -> a -> c
    . VName -> SubExp
    Var) [VName]
    gtids
        String -> Exp (Rep m) -> m VName
    forall (m :: * -> *).
    MonadBuilder m =>
    String -> Exp (Rep m) -> m VName
    letExp (VName -> String
    baseString VName
    acc String -> String -> String
    forall a. Semigroup a => a -> a -> a
    <> String
    "_upd") (Exp (Rep m) -> m VName) -> Exp (Rep m) -> m VName
    forall a b. (a -> b) -> a -> b
    $
          BasicOp -> Exp (Rep m)
    forall rep. BasicOp -> Exp rep
    BasicOp (BasicOp -> Exp (Rep m)) -> BasicOp -> Exp (Rep m)
    forall a b. (a -> b) -> a -> b
    $
            Safety -> VName -> [SubExp] -> [SubExp] -> BasicOp
    UpdateAcc Safety
    Safe VName
    acc ((VName -> SubExp) -> [VName] -> [SubExp]
    forall a b. (a -> b) -> [a] -> [b]
    map VName -> SubExp
    Var [VName]
    gtids) [SubExp]
    vs

      Type
    acc_t <- VName -> m Type
    forall rep (m :: * -> *). HasScope rep m => VName -> m Type
    lookupType VName
    acc
      Exp GPU -> m (Exp GPU)
    forall a. a -> m a
    forall (f :: * -> *) a. Applicative f => a -> f a
    pure (Exp GPU -> m (Exp GPU))
    -> (KernelBody GPU -> Exp GPU) -> KernelBody GPU -> m (Exp GPU)
    forall b c a. (b -> c) -> (a -> b) -> a -> c
    . Op GPU -> Exp GPU
    HostOp SOAC GPU -> Exp GPU
    forall rep. Op rep -> Exp rep
    Op (HostOp SOAC GPU -> Exp GPU)
    -> (KernelBody GPU -> HostOp SOAC GPU) -> KernelBody GPU -> Exp GPU
    forall b c a. (b -> c) -> (a -> b) -> a -> c
    . SegOp SegLevel GPU -> HostOp SOAC GPU
    forall (op :: * -> *) rep. SegOp SegLevel rep -> HostOp op rep
    SegOp (SegOp SegLevel GPU -> HostOp SOAC GPU)
    -> (KernelBody GPU -> SegOp SegLevel GPU)
    -> KernelBody GPU
    -> HostOp SOAC GPU
    forall b c a. (b -> c) -> (a -> b) -> a -> c
    . SegLevel
    -> SegSpace -> [Type] -> KernelBody GPU -> SegOp SegLevel GPU
    forall lvl rep.
    lvl -> SegSpace -> [Type] -> KernelBody rep -> SegOp lvl rep
    SegMap SegLevel
    lvl SegSpace
    space [Type
    acc_t] (KernelBody GPU -> m (Exp GPU)) -> KernelBody GPU -> m (Exp GPU)
    forall a b. (a -> b) -> a -> b
    $
        BodyDec GPU -> Stms GPU -> [KernelResult] -> KernelBody GPU
    forall rep.
    BodyDec rep -> Stms rep -> [KernelResult] -> KernelBody rep
    KernelBody () Stms GPU
    stms [ResultManifest -> Certs -> SubExp -> KernelResult
    Returns ResultManifest
    ResultMaySimplify Certs
    forall a. Monoid a => a
    mempty (VName -> SubExp
    Var VName
    acc')]

    flatKernelBody ::
      (MonadBuilder m) =>
      SegSpace ->
      KernelBody (Rep m) ->
      m (SegSpace, KernelBody (Rep m))
    flatKernelBody :: forall (m :: * -> *).
    MonadBuilder m =>
    SegSpace -> KernelBody (Rep m) -> m (SegSpace, KernelBody (Rep m))
    flatKernelBody SegSpace
    space KernelBody (Rep m)
    kbody = do
      VName
    gtid <- String -> m VName
    forall (m :: * -> *). MonadFreshNames m => String -> m VName
    newVName String
    "gtid"
      SubExp
    dims_prod <-
        String -> Exp (Rep m) -> m SubExp
    forall (m :: * -> *).
    MonadBuilder m =>
    String -> Exp (Rep m) -> m SubExp
    letSubExp String
    "dims_prod"
          (Exp (Rep m) -> m SubExp) -> m (Exp (Rep m)) -> m SubExp
    forall (m :: * -> *) a b. Monad m => (a -> m b) -> m a -> m b
    =<< BinOp -> SubExp -> [SubExp] -> m (Exp (Rep m))
    forall (m :: * -> *).
    MonadBuilder m =>
    BinOp -> SubExp -> [SubExp] -> m (Exp (Rep m))
    foldBinOp (IntType -> Overflow -> BinOp
    Mul IntType
    Int64 Overflow
    OverflowUndef) (IntType -> Integer -> SubExp
    intConst IntType
    Int64 Integer
    1) (SegSpace -> [SubExp]
    segSpaceDims SegSpace
    space)

      let space' :: SegSpace
    space' = VName -> [(VName, SubExp)] -> SegSpace
    SegSpace (SegSpace -> VName
    segFlat SegSpace
    space) [(VName
    gtid, SubExp
    dims_prod)]

      Stms (Rep m)
    kbody_stms <- Scope (Rep m) -> m (Stms (Rep m)) -> m (Stms (Rep m))
    forall a. Scope (Rep m) -> m a -> m a
    forall rep (m :: * -> *) a.
    LocalScope rep m =>
    Scope rep -> m a -> m a
    localScope (SegSpace -> Scope (Rep m)
    forall rep. SegSpace -> Scope rep
    scopeOfSegSpace SegSpace
    space') (m (Stms (Rep m)) -> m (Stms (Rep m)))
    -> (m () -> m (Stms (Rep m))) -> m () -> m (Stms (Rep m))
    forall b c a. (b -> c) -> (a -> b) -> a -> c
    . m () -> m (Stms (Rep m))
    forall (m :: * -> *) a. MonadBuilder m => m a -> m (Stms (Rep m))
    collectStms_ (m () -> m (Stms (Rep m))) -> m () -> m (Stms (Rep m))
    forall a b. (a -> b) -> a -> b
    $ do
        let new_inds :: [TPrimExp Int64 VName]
    new_inds =
              [TPrimExp Int64 VName]
    -> TPrimExp Int64 VName -> [TPrimExp Int64 VName]
    forall num. IntegralExp num => [num] -> num -> [num]
    unflattenIndex ((SubExp -> TPrimExp Int64 VName)
    -> [SubExp] -> [TPrimExp Int64 VName]
    forall a b. (a -> b) -> [a] -> [b]
    map SubExp -> TPrimExp Int64 VName
    pe64 (SegSpace -> [SubExp]
    segSpaceDims SegSpace
    space)) (SubExp -> TPrimExp Int64 VName
    pe64 (SubExp -> TPrimExp Int64 VName) -> SubExp -> TPrimExp Int64 VName
    forall a b. (a -> b) -> a -> b
    $ VName -> SubExp
    Var VName
    gtid)
        ([VName] -> Exp (Rep m) -> m ())
    -> [[VName]] -> [Exp (Rep m)] -> m ()
    forall (m :: * -> *) a b c.
    Applicative m =>
    (a -> b -> m c) -> [a] -> [b] -> m ()
    zipWithM_ [VName] -> Exp (Rep m) -> m ()
    forall (m :: * -> *).
    MonadBuilder m =>
    [VName] -> Exp (Rep m) -> m ()
    letBindNames (((VName, SubExp) -> [VName]) -> [(VName, SubExp)] -> [[VName]]
    forall a b. (a -> b) -> [a] -> [b]
    map (VName -> [VName]
    forall a. a -> [a]
    forall (f :: * -> *) a. Applicative f => a -> f a
    pure (VName -> [VName])
    -> ((VName, SubExp) -> VName) -> (VName, SubExp) -> [VName]
    forall b c a. (b -> c) -> (a -> b) -> a -> c
    . (VName, SubExp) -> VName
    forall a b. (a, b) -> a
    fst) (SegSpace -> [(VName, SubExp)]
    unSegSpace SegSpace
    space))
          ([Exp (Rep m)] -> m ()) -> m [Exp (Rep m)] -> m ()
    forall (m :: * -> *) a b. Monad m => (a -> m b) -> m a -> m b
    =<< (TPrimExp Int64 VName -> m (Exp (Rep m)))
    -> [TPrimExp Int64 VName] -> m [Exp (Rep m)]
    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 TPrimExp Int64 VName -> m (Exp (Rep m))
    forall a (m :: * -> *).
    (ToExp a, MonadBuilder m) =>
    a -> m (Exp (Rep m))
    forall (m :: * -> *).
    MonadBuilder m =>
    TPrimExp Int64 VName -> m (Exp (Rep m))
    toExp [TPrimExp Int64 VName]
    new_inds
        Stms (Rep m) -> m ()
    forall (m :: * -> *). MonadBuilder m => Stms (Rep m) -> m ()
    addStms (Stms (Rep m) -> m ()) -> Stms (Rep m) -> m ()
    forall a b. (a -> b) -> a -> b
    $ KernelBody (Rep m) -> Stms (Rep m)
    forall rep. KernelBody rep -> Stms rep
    kernelBodyStms KernelBody (Rep m)
    kbody

      (SegSpace, KernelBody (Rep m)) -> m (SegSpace, KernelBody (Rep m))
    forall a. a -> m a
    forall (f :: * -> *) a. Applicative f => a -> f a
    pure (SegSpace
    space', KernelBody (Rep m)
    kbody {kernelBodyStms = kbody_stms})

    optimiseStm :: Accs GPU -> Stm GPU -> OptM (Stms GPU)
    -- TODO: this is very restricted currently, but shows the idea.
    optimiseStm :: Map VName (WithAccInput GPU)
    -> Stm GPU -> ReaderT (Scope GPU) (State VNameSource) (Stms GPU)
    optimiseStm Map VName (WithAccInput GPU)
    accs (Let Pat (LetDec GPU)
    pat StmAux (ExpDec GPU)
    aux (WithAcc [WithAccInput GPU]
    inputs Lambda GPU
    lam)) = do
      Scope GPU
    -> ReaderT (Scope GPU) (State VNameSource) (Stms GPU)
    -> ReaderT (Scope GPU) (State VNameSource) (Stms GPU)
    forall a.
    Scope GPU
    -> ReaderT (Scope GPU) (State VNameSource) a
    -> ReaderT (Scope GPU) (State VNameSource) a
    forall rep (m :: * -> *) a.
    LocalScope rep m =>
    Scope rep -> m a -> m a
    localScope ([Param (LParamInfo GPU)] -> Scope GPU
    forall rep dec. (LParamInfo rep ~ dec) => [Param dec] -> Scope rep
    scopeOfLParams (Lambda GPU -> [Param (LParamInfo GPU)]
    forall rep. Lambda rep -> [LParam rep]
    lambdaParams Lambda GPU
    lam)) (ReaderT (Scope GPU) (State VNameSource) (Stms GPU)
     -> ReaderT (Scope GPU) (State VNameSource) (Stms GPU))
    -> ReaderT (Scope GPU) (State VNameSource) (Stms GPU)
    -> ReaderT (Scope GPU) (State VNameSource) (Stms GPU)
    forall a b. (a -> b) -> a -> b
    $ do
        Body GPU
    body' <- Map VName (WithAccInput GPU) -> Body GPU -> OptM (Body GPU)
    optimiseBody Map VName (WithAccInput GPU)
    accs' (Body GPU -> OptM (Body GPU)) -> Body GPU -> OptM (Body GPU)
    forall a b. (a -> b) -> a -> b
    $ Lambda GPU -> Body GPU
    forall rep. Lambda rep -> Body rep
    lambdaBody Lambda GPU
    lam
        let lam' :: Lambda GPU
    lam' = Lambda GPU
    lam {lambdaBody = body'}
        Stms GPU -> ReaderT (Scope GPU) (State VNameSource) (Stms GPU)
    forall a. a -> ReaderT (Scope GPU) (State VNameSource) a
    forall (f :: * -> *) a. Applicative f => a -> f a
    pure (Stms GPU -> ReaderT (Scope GPU) (State VNameSource) (Stms GPU))
    -> Stms GPU -> ReaderT (Scope GPU) (State VNameSource) (Stms GPU)
    forall a b. (a -> b) -> a -> b
    $ Stm GPU -> Stms GPU
    forall rep. Stm rep -> Stms rep
    oneStm (Stm GPU -> Stms GPU) -> Stm GPU -> Stms GPU
    forall a b. (a -> b) -> a -> b
    $ Pat (LetDec GPU) -> StmAux (ExpDec GPU) -> Exp GPU -> Stm GPU
    forall rep.
    Pat (LetDec rep) -> StmAux (ExpDec rep) -> Exp rep -> Stm rep
    Let Pat (LetDec GPU)
    pat StmAux (ExpDec GPU)
    aux (Exp GPU -> Stm GPU) -> Exp GPU -> Stm GPU
    forall a b. (a -> b) -> a -> b
    $ [WithAccInput GPU] -> Lambda GPU -> Exp GPU
    forall rep. [WithAccInput rep] -> Lambda rep -> Exp rep
    WithAcc [WithAccInput GPU]
    inputs Lambda GPU
    lam'
      where
        acc_names :: [VName]
    acc_names = (Param Type -> VName) -> [Param Type] -> [VName]
    forall a b. (a -> b) -> [a] -> [b]
    map Param Type -> VName
    forall dec. Param dec -> VName
    paramName ([Param Type] -> [VName]) -> [Param Type] -> [VName]
    forall a b. (a -> b) -> a -> b
    $ Int -> [Param Type] -> [Param Type]
    forall a. Int -> [a] -> [a]
    drop ([WithAccInput GPU] -> Int
    forall a. [a] -> Int
    forall (t :: * -> *) a. Foldable t => t a -> Int
    length [WithAccInput GPU]
    inputs) ([Param Type] -> [Param Type]) -> [Param Type] -> [Param Type]
    forall a b. (a -> b) -> a -> b
    $ Lambda GPU -> [Param (LParamInfo GPU)]
    forall rep. Lambda rep -> [LParam rep]
    lambdaParams Lambda GPU
    lam
        accs' :: Map VName (WithAccInput GPU)
    accs' = [(VName, WithAccInput GPU)] -> Map VName (WithAccInput GPU)
    forall k a. Ord k => [(k, a)] -> Map k a
    M.fromList ([VName] -> [WithAccInput GPU] -> [(VName, WithAccInput GPU)]
    forall a b. [a] -> [b] -> [(a, b)]
    zip [VName]
    acc_names [WithAccInput GPU]
    inputs) Map VName (WithAccInput GPU)
    -> Map VName (WithAccInput GPU) -> Map VName (WithAccInput GPU)
    forall a. Semigroup a => a -> a -> a
    <> Map VName (WithAccInput GPU)
    accs
    optimiseStm Map VName (WithAccInput GPU)
    accs (Let Pat (LetDec GPU)
    pat StmAux (ExpDec GPU)
    aux (Op (SegOp (SegMap SegLevel
    lvl SegSpace
    space [Type]
    _ KernelBody GPU
    kbody))))
      | Map VName (WithAccInput GPU)
    accs Map VName (WithAccInput GPU)
    -> Map VName (WithAccInput GPU) -> Bool
    forall a. Eq a => a -> a -> Bool
    /= Map VName (WithAccInput GPU)
    forall a. Monoid a => a
    mempty,
        Just (KernelBody GPU
    kbody', (Shape
    acc_shape, [VName]
    _, Just (Lambda GPU
    acc_lam, [SubExp]
    acc_nes)), VName
    acc) <-
          Map VName (WithAccInput GPU)
    -> KernelBody GPU
    -> Maybe (KernelBody GPU, WithAccInput GPU, VName)
    mkHistBody Map VName (WithAccInput GPU)
    accs KernelBody GPU
    kbody,
        (Type -> Bool) -> [Type] -> Bool
    forall (t :: * -> *) a. Foldable t => (a -> Bool) -> t a -> Bool
    all Type -> Bool
    forall shape u. TypeBase shape u -> Bool
    primType ([Type] -> Bool) -> [Type] -> Bool
    forall a b. (a -> b) -> a -> b
    $ Lambda GPU -> [Type]
    forall rep. Lambda rep -> [Type]
    lambdaReturnType Lambda GPU
    acc_lam = Builder GPU ()
    -> ReaderT (Scope GPU) (State VNameSource) (Stms GPU)
    forall (m :: * -> *) somerep rep a.
    (MonadFreshNames m, HasScope somerep m, SameScope somerep rep) =>
    Builder rep a -> m (Stms rep)
    runBuilder_ (Builder GPU ()
     -> ReaderT (Scope GPU) (State VNameSource) (Stms GPU))
    -> Builder GPU ()
    -> ReaderT (Scope GPU) (State VNameSource) (Stms GPU)
    forall a b. (a -> b) -> a -> b
    $ do
          [VName]
    hist_dests <- [SubExp]
    -> (SubExp -> BuilderT GPU (State VNameSource) VName)
    -> BuilderT GPU (State VNameSource) [VName]
    forall (t :: * -> *) (m :: * -> *) a b.
    (Traversable t, Monad m) =>
    t a -> (a -> m b) -> m (t b)
    forM [SubExp]
    acc_nes ((SubExp -> BuilderT GPU (State VNameSource) VName)
     -> BuilderT GPU (State VNameSource) [VName])
    -> (SubExp -> BuilderT GPU (State VNameSource) VName)
    -> BuilderT GPU (State VNameSource) [VName]
    forall a b. (a -> b) -> a -> b
    $ \SubExp
    ne ->
            String
    -> Exp (Rep (BuilderT GPU (State VNameSource)))
    -> BuilderT GPU (State VNameSource) VName
    forall (m :: * -> *).
    MonadBuilder m =>
    String -> Exp (Rep m) -> m VName
    letExp String
    "hist_dest" (Exp (Rep (BuilderT GPU (State VNameSource)))
     -> BuilderT GPU (State VNameSource) VName)
    -> Exp (Rep (BuilderT GPU (State VNameSource)))
    -> BuilderT GPU (State VNameSource) VName
    forall a b. (a -> b) -> a -> b
    $ BasicOp -> Exp (Rep (BuilderT GPU (State VNameSource)))
    forall rep. BasicOp -> Exp rep
    BasicOp (BasicOp -> Exp (Rep (BuilderT GPU (State VNameSource))))
    -> BasicOp -> Exp (Rep (BuilderT GPU (State VNameSource)))
    forall a b. (a -> b) -> a -> b
    $ Shape -> SubExp -> BasicOp
    Replicate Shape
    acc_shape SubExp
    ne

          Lambda GPU
    acc_lam' <- Shape
    -> Lambda GPU -> BuilderT GPU (State VNameSource) (Lambda GPU)
    forall (m :: * -> *).
    MonadFreshNames m =>
    Shape -> Lambda GPU -> m (Lambda GPU)
    withAccLamToHistLam Shape
    acc_shape Lambda GPU
    acc_lam

          let ts' :: [Type]
    ts' =
                Int -> Type -> [Type]
    forall a. Int -> a -> [a]
    replicate (Shape -> Int
    forall a. ArrayShape a => a -> Int
    shapeRank Shape
    acc_shape) (PrimType -> Type
    forall shape u. PrimType -> TypeBase shape u
    Prim PrimType
    int64)
                  [Type] -> [Type] -> [Type]
    forall a. [a] -> [a] -> [a]
    ++ Lambda GPU -> [Type]
    forall rep. Lambda rep -> [Type]
    lambdaReturnType Lambda GPU
    acc_lam
              histop :: HistOp GPU
    histop =
                HistOp
                  { histShape :: Shape
    histShape = Shape
    acc_shape,
                    histRaceFactor :: SubExp
    histRaceFactor = IntType -> Integer -> SubExp
    intConst IntType
    Int64 Integer
    1,
                    histDest :: [VName]
    histDest = [VName]
    hist_dests,
                    histNeutral :: [SubExp]
    histNeutral = [SubExp]
    acc_nes,
                    histOpShape :: Shape
    histOpShape = Shape
    forall a. Monoid a => a
    mempty,
                    histOp :: Lambda GPU
    histOp = Lambda GPU
    acc_lam'
                  }

          (SegSpace
    space', KernelBody GPU
    kbody'') <- SegSpace
    -> KernelBody (Rep (BuilderT GPU (State VNameSource)))
    -> BuilderT
         GPU
         (State VNameSource)
         (SegSpace, KernelBody (Rep (BuilderT GPU (State VNameSource))))
    forall (m :: * -> *).
    MonadBuilder m =>
    SegSpace -> KernelBody (Rep m) -> m (SegSpace, KernelBody (Rep m))
    flatKernelBody SegSpace
    space KernelBody (Rep (BuilderT GPU (State VNameSource)))
    KernelBody GPU
    kbody'

          [VName]
    hist_dest_upd <-
            String
    -> Exp (Rep (BuilderT GPU (State VNameSource)))
    -> BuilderT GPU (State VNameSource) [VName]
    forall (m :: * -> *).
    MonadBuilder m =>
    String -> Exp (Rep m) -> m [VName]
    letTupExp String
    "hist_dest_upd" (Exp (Rep (BuilderT GPU (State VNameSource)))
     -> BuilderT GPU (State VNameSource) [VName])
    -> Exp (Rep (BuilderT GPU (State VNameSource)))
    -> BuilderT GPU (State VNameSource) [VName]
    forall a b. (a -> b) -> a -> b
    $ Op (Rep (BuilderT GPU (State VNameSource)))
    -> Exp (Rep (BuilderT GPU (State VNameSource)))
    forall rep. Op rep -> Exp rep
    Op (Op (Rep (BuilderT GPU (State VNameSource)))
     -> Exp (Rep (BuilderT GPU (State VNameSource))))
    -> Op (Rep (BuilderT GPU (State VNameSource)))
    -> Exp (Rep (BuilderT GPU (State VNameSource)))
    forall a b. (a -> b) -> a -> b
    $ SegOp SegLevel GPU -> HostOp SOAC GPU
    forall (op :: * -> *) rep. SegOp SegLevel rep -> HostOp op rep
    SegOp (SegOp SegLevel GPU -> HostOp SOAC GPU)
    -> SegOp SegLevel GPU -> HostOp SOAC GPU
    forall a b. (a -> b) -> a -> b
    $ SegLevel
    -> SegSpace
    -> [Type]
    -> KernelBody GPU
    -> [HistOp GPU]
    -> SegOp SegLevel GPU
    forall lvl rep.
    lvl
    -> SegSpace
    -> [Type]
    -> KernelBody rep
    -> [HistOp rep]
    -> SegOp lvl rep
    SegHist SegLevel
    lvl SegSpace
    space' [Type]
    ts' KernelBody GPU
    kbody'' [HistOp GPU
    histop]

          Stm (Rep (BuilderT GPU (State VNameSource))) -> Builder GPU ()
    Stm GPU -> Builder GPU ()
    forall (m :: * -> *). MonadBuilder m => Stm (Rep m) -> m ()
    addStm (Stm GPU -> Builder GPU ())
    -> (Exp GPU -> Stm GPU) -> Exp GPU -> Builder GPU ()
    forall b c a. (b -> c) -> (a -> b) -> a -> c
    . Pat (LetDec GPU) -> StmAux (ExpDec GPU) -> Exp GPU -> Stm GPU
    forall rep.
    Pat (LetDec rep) -> StmAux (ExpDec rep) -> Exp rep -> Stm rep
    Let Pat (LetDec GPU)
    pat StmAux (ExpDec GPU)
    aux (Exp GPU -> Builder GPU ())
    -> BuilderT GPU (State VNameSource) (Exp GPU) -> Builder GPU ()
    forall (m :: * -> *) a b. Monad m => (a -> m b) -> m a -> m b
    =<< SegLevel
    -> Shape
    -> [VName]
    -> VName
    -> BuilderT GPU (State VNameSource) (Exp GPU)
    forall (m :: * -> *).
    (MonadBuilder m, Rep m ~ GPU) =>
    SegLevel -> Shape -> [VName] -> VName -> m (Exp GPU)
    addArrsToAcc SegLevel
    lvl Shape
    acc_shape [VName]
    hist_dest_upd VName
    acc
    optimiseStm Map VName (WithAccInput GPU)
    accs (Let Pat (LetDec GPU)
    pat StmAux (ExpDec GPU)
    aux Exp GPU
    e) =
      Stm GPU -> Stms GPU
    forall rep. Stm rep -> Stms rep
    oneStm (Stm GPU -> Stms GPU)
    -> (Exp GPU -> Stm GPU) -> Exp GPU -> Stms GPU
    forall b c a. (b -> c) -> (a -> b) -> a -> c
    . Pat (LetDec GPU) -> StmAux (ExpDec GPU) -> Exp GPU -> Stm GPU
    forall rep.
    Pat (LetDec rep) -> StmAux (ExpDec rep) -> Exp rep -> Stm rep
    Let Pat (LetDec GPU)
    pat StmAux (ExpDec GPU)
    aux (Exp GPU -> Stms GPU)
    -> OptM (Exp GPU)
    -> ReaderT (Scope GPU) (State VNameSource) (Stms GPU)
    forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
    <$> Map VName (WithAccInput GPU) -> Exp GPU -> OptM (Exp GPU)
    optimiseExp Map VName (WithAccInput GPU)
    accs Exp GPU
    e

    optimiseStms :: Accs GPU -> Stms GPU -> OptM (Stms GPU)
    optimiseStms :: Map VName (WithAccInput GPU)
    -> Stms GPU -> ReaderT (Scope GPU) (State VNameSource) (Stms GPU)
    optimiseStms Map VName (WithAccInput GPU)
    accs Stms GPU
    stms =
      Scope GPU
    -> ReaderT (Scope GPU) (State VNameSource) (Stms GPU)
    -> ReaderT (Scope GPU) (State VNameSource) (Stms GPU)
    forall a.
    Scope GPU
    -> ReaderT (Scope GPU) (State VNameSource) a
    -> ReaderT (Scope GPU) (State VNameSource) a
    forall rep (m :: * -> *) a.
    LocalScope rep m =>
    Scope rep -> m a -> m a
    localScope (Stms GPU -> Scope GPU
    forall rep a. Scoped rep a => a -> Scope rep
    scopeOf Stms GPU
    stms) (ReaderT (Scope GPU) (State VNameSource) (Stms GPU)
     -> ReaderT (Scope GPU) (State VNameSource) (Stms GPU))
    -> ReaderT (Scope GPU) (State VNameSource) (Stms GPU)
    -> ReaderT (Scope GPU) (State VNameSource) (Stms GPU)
    forall a b. (a -> b) -> a -> b
    $
        [Stms GPU] -> Stms GPU
    forall a. Monoid a => [a] -> a
    mconcat ([Stms GPU] -> Stms GPU)
    -> ReaderT (Scope GPU) (State VNameSource) [Stms GPU]
    -> ReaderT (Scope GPU) (State VNameSource) (Stms GPU)
    forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
    <$> (Stm GPU -> ReaderT (Scope GPU) (State VNameSource) (Stms GPU))
    -> [Stm GPU] -> ReaderT (Scope GPU) (State VNameSource) [Stms GPU]
    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 (Map VName (WithAccInput GPU)
    -> Stm GPU -> ReaderT (Scope GPU) (State VNameSource) (Stms GPU)
    optimiseStm Map VName (WithAccInput GPU)
    accs) (Stms GPU -> [Stm GPU]
    forall rep. Stms rep -> [Stm rep]
    stmsToList Stms GPU
    stms)

    -- | The pass for GPU kernels.
    histAccsGPU :: Pass GPU GPU
    histAccsGPU :: Pass GPU GPU
    histAccsGPU =
      String -> String -> (Prog GPU -> PassM (Prog GPU)) -> Pass GPU GPU
    forall fromrep torep.
    String
    -> String
    -> (Prog fromrep -> PassM (Prog torep))
    -> Pass fromrep torep
    Pass String
    "hist accs" String
    "Turn certain accumulations into histograms" ((Prog GPU -> PassM (Prog GPU)) -> Pass GPU GPU)
    -> (Prog GPU -> PassM (Prog GPU)) -> Pass GPU GPU
    forall a b. (a -> b) -> a -> b
    $
        (Scope GPU -> Stms GPU -> PassM (Stms GPU))
    -> Prog GPU -> PassM (Prog GPU)
    forall rep.
    (Scope rep -> Stms rep -> PassM (Stms rep))
    -> Prog rep -> PassM (Prog rep)
    intraproceduralTransformation Scope GPU -> Stms GPU -> PassM (Stms GPU)
    forall {m :: * -> *}.
    MonadFreshNames m =>
    Scope GPU -> Stms GPU -> m (Stms GPU)
    onStms
      where
        onStms :: Scope GPU -> Stms GPU -> m (Stms GPU)
    onStms Scope GPU
    scope Stms GPU
    stms =
          (VNameSource -> (Stms GPU, VNameSource)) -> m (Stms GPU)
    forall (m :: * -> *) a.
    MonadFreshNames m =>
    (VNameSource -> (a, VNameSource)) -> m a
    modifyNameSource ((VNameSource -> (Stms GPU, VNameSource)) -> m (Stms GPU))
    -> (State VNameSource (Stms GPU)
        -> VNameSource -> (Stms GPU, VNameSource))
    -> State VNameSource (Stms GPU)
    -> m (Stms GPU)
    forall b c a. (b -> c) -> (a -> b) -> a -> c
    . State VNameSource (Stms GPU)
    -> VNameSource -> (Stms GPU, VNameSource)
    forall s a. State s a -> s -> (a, s)
    runState (State VNameSource (Stms GPU) -> m (Stms GPU))
    -> State VNameSource (Stms GPU) -> m (Stms GPU)
    forall a b. (a -> b) -> a -> b
    $
            ReaderT (Scope GPU) (State VNameSource) (Stms GPU)
    -> Scope GPU -> State VNameSource (Stms GPU)
    forall r (m :: * -> *) a. ReaderT r m a -> r -> m a
    runReaderT (Map VName (WithAccInput GPU)
    -> Stms GPU -> ReaderT (Scope GPU) (State VNameSource) (Stms GPU)
    optimiseStms Map VName (WithAccInput GPU)
    forall a. Monoid a => a
    mempty Stms GPU
    stms) Scope GPU
    scope