    {-# LANGUAGE TypeFamilies #-}

    -- | Tries to turn a generalized reduction kernel into
    --     a more specialized construct, for example:
    --       (a) a map nest with a sequential redomap ripe for tiling
    --       (b) a SegRed kernel followed by a smallish accumulation kernel.
    --       (c) a histogram (for this we need to track the withAccs)
    --   The idea is to identify the first accumulation and
    --     to separate the initial kernels into two:
    --     1. the code up to and including the accumulation,
    --        which is optimized to turn the accumulation either
    --        into a map-reduce composition or a histogram, and
    --     2. the remaining code, which is recursively optimized.
    --   Since this is mostly prototyping, when the accumulation
    --     can be rewritten as a map-reduce, we sequentialize the
    --     map-reduce, as to potentially enable tiling oportunities.
    module Futhark.Optimise.GenRedOpt (optimiseGenRed) where

    import Control.Monad
    import Control.Monad.Reader
    import Control.Monad.State
    import Data.List qualified as L
    import Data.Map.Strict qualified as M
    import Data.Maybe
    import Futhark.Builder
    import Futhark.IR.GPU
    import Futhark.Optimise.TileLoops.Shared
    import Futhark.Pass
    import Futhark.Tools
    import Futhark.Transform.Rename

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

    -- | The pass definition.
    optimiseGenRed :: Pass GPU GPU
    optimiseGenRed :: Pass GPU GPU
    optimiseGenRed =
      [Char] -> [Char] -> (Prog GPU -> PassM (Prog GPU)) -> Pass GPU GPU
    forall fromrep torep.
    [Char]
    -> [Char]
    -> (Prog fromrep -> PassM (Prog torep))
    -> Pass fromrep torep
    Pass [Char]
    "optimise generalized reductions" [Char]
    "Specializes generalized reductions into map-reductions or 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))
    -> (VNameSource -> (Stms GPU, VNameSource)) -> m (Stms GPU)
    forall a b. (a -> b) -> a -> b
    $
            State VNameSource (Stms GPU)
    -> VNameSource -> (Stms GPU, VNameSource)
    forall s a. State s a -> s -> (a, s)
    runState (State VNameSource (Stms GPU)
     -> VNameSource -> (Stms GPU, VNameSource))
    -> State VNameSource (Stms GPU)
    -> VNameSource
    -> (Stms GPU, VNameSource)
    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 (Env
    -> Stms GPU -> ReaderT (Scope GPU) (State VNameSource) (Stms GPU)
    optimiseStms (WithEnv
    forall k a. Map k a
    M.empty, Map VName LMAD
    forall k a. Map k a
    M.empty) Stms GPU
    stms) Scope GPU
    scope

    optimiseBody :: Env -> Body GPU -> GenRedM (Body GPU)
    optimiseBody :: Env -> Body GPU -> GenRedM (Body GPU)
    optimiseBody Env
    env (Body () Stms GPU
    stms Result
    res) =
      BodyDec GPU -> Stms GPU -> Result -> Body GPU
    forall rep. BodyDec rep -> Stms rep -> Result -> Body rep
    Body () (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
    <$> Env
    -> Stms GPU -> ReaderT (Scope GPU) (State VNameSource) (Stms GPU)
    optimiseStms Env
    env Stms GPU
    stms ReaderT (Scope GPU) (State VNameSource) (Result -> Body GPU)
    -> ReaderT (Scope GPU) (State VNameSource) Result
    -> GenRedM (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 Result
    res

    optimiseStms :: Env -> Stms GPU -> GenRedM (Stms GPU)
    optimiseStms :: Env
    -> Stms GPU -> ReaderT (Scope GPU) (State VNameSource) (Stms GPU)
    optimiseStms Env
    env 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
    $ do
        (_, stms') <- ((Env, Stms GPU)
     -> Stm GPU
     -> ReaderT (Scope GPU) (State VNameSource) (Env, Stms GPU))
    -> (Env, Stms GPU)
    -> [Stm GPU]
    -> ReaderT (Scope GPU) (State VNameSource) (Env, Stms GPU)
    forall (t :: * -> *) (m :: * -> *) b a.
    (Foldable t, Monad m) =>
    (b -> a -> m b) -> b -> t a -> m b
    foldM (Env, Stms GPU)
    -> Stm GPU
    -> ReaderT (Scope GPU) (State VNameSource) (Env, Stms GPU)
    foldfun (Env
    env, Stms GPU
    forall a. Monoid a => a
    mempty) ([Stm GPU]
     -> ReaderT (Scope GPU) (State VNameSource) (Env, Stms GPU))
    -> [Stm GPU]
    -> ReaderT (Scope GPU) (State VNameSource) (Env, Stms GPU)
    forall a b. (a -> b) -> a -> b
    $ Stms GPU -> [Stm GPU]
    forall rep. Stms rep -> [Stm rep]
    stmsToList Stms GPU
    stms
        pure stms'
      where
        foldfun :: (Env, Stms GPU) -> Stm GPU -> GenRedM (Env, Stms GPU)
        foldfun :: (Env, Stms GPU)
    -> Stm GPU
    -> ReaderT (Scope GPU) (State VNameSource) (Env, Stms GPU)
    foldfun (Env
    e, Stms GPU
    ss) Stm GPU
    s = do
          (e', s') <- Env
    -> Stm GPU
    -> ReaderT (Scope GPU) (State VNameSource) (Env, Stms GPU)
    optimiseStm Env
    e Stm GPU
    s
          pure (e', ss <> s')

    optimiseStm :: Env -> Stm GPU -> GenRedM (Env, Stms GPU)
    optimiseStm :: Env
    -> Stm GPU
    -> ReaderT (Scope GPU) (State VNameSource) (Env, Stms GPU)
    optimiseStm Env
    env stm :: Stm GPU
    stm@(Let Pat (LetDec GPU)
    _ StmAux (ExpDec GPU)
    _ (Op (SegOp (SegMap SegThread {} SegSpace
    _ [Type]
    _ KernelBody GPU
    _)))) = do
      res_genred_opt <- Env -> Stm GPU -> GenRedM (Maybe (Stms GPU))
    genRedOpts Env
    env Stm GPU
    stm
      let stms' =
            case Maybe (Stms GPU)
    res_genred_opt of
              Just Stms GPU
    stms -> Stms GPU
    stms
              Maybe (Stms GPU)
    Nothing -> Stm GPU -> Stms GPU
    forall rep. Stm rep -> Stms rep
    oneStm Stm GPU
    stm
      pure (env, stms')
    optimiseStm Env
    env (Let Pat (LetDec GPU)
    pat StmAux (ExpDec GPU)
    aux Exp GPU
    e) = do
      env' <- Env -> VName -> Exp GPU -> TileM Env
    changeEnv Env
    env ([VName] -> VName
    forall a. HasCallStack => [a] -> a
    head ([VName] -> VName) -> [VName] -> VName
    forall a b. (a -> b) -> a -> b
    $ Pat Type -> [VName]
    forall dec. Pat dec -> [VName]
    patNames Pat Type
    Pat (LetDec GPU)
    pat) Exp GPU
    e
      e' <- mapExpM (optimise env') e
      pure (env', oneStm $ Let pat aux e')
      where
        optimise :: Env -> Mapper GPU GPU (ReaderT (Scope GPU) (State VNameSource))
    optimise Env
    env' = Mapper GPU GPU (ReaderT (Scope GPU) (State VNameSource))
    forall rep (m :: * -> *). Monad m => Mapper rep rep m
    identityMapper {mapOnBody = \Scope GPU
    scope -> Scope GPU -> GenRedM (Body GPU) -> GenRedM (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 (GenRedM (Body GPU) -> GenRedM (Body GPU))
    -> (Body GPU -> GenRedM (Body GPU))
    -> Body GPU
    -> GenRedM (Body GPU)
    forall b c a. (b -> c) -> (a -> b) -> a -> c
    . Env -> Body GPU -> GenRedM (Body GPU)
    optimiseBody Env
    env'}

    ------------------------

    genRedOpts :: Env -> Stm GPU -> GenRedM (Maybe (Stms GPU))
    genRedOpts :: Env -> Stm GPU -> GenRedM (Maybe (Stms GPU))
    genRedOpts Env
    env Stm GPU
    ker = do
      res_tile <- Env -> Stm GPU -> GenRedM (Maybe (Stms GPU, Stm GPU))
    genRed2Tile2d Env
    env Stm GPU
    ker
      case res_tile of
        Maybe (Stms GPU, Stm GPU)
    Nothing -> do
          res_sgrd <- Env -> Stm GPU -> GenRedM (Maybe (Stms GPU, Stm GPU))
    genRed2SegRed Env
    env Stm GPU
    ker
          helperGenRed res_sgrd
        Maybe (Stms GPU, Stm GPU)
    _ -> Maybe (Stms GPU, Stm GPU) -> GenRedM (Maybe (Stms GPU))
    helperGenRed Maybe (Stms GPU, Stm GPU)
    res_tile
      where
        helperGenRed :: Maybe (Stms GPU, Stm GPU) -> GenRedM (Maybe (Stms GPU))
    helperGenRed Maybe (Stms GPU, Stm GPU)
    Nothing = Maybe (Stms GPU) -> GenRedM (Maybe (Stms GPU))
    forall a. a -> ReaderT (Scope GPU) (State VNameSource) a
    forall (f :: * -> *) a. Applicative f => a -> f a
    pure Maybe (Stms GPU)
    forall a. Maybe a
    Nothing
        helperGenRed (Just (Stms GPU
    stms_before, Stm GPU
    ker_snd)) = do
          mb_stms_after <- Env -> Stm GPU -> GenRedM (Maybe (Stms GPU))
    genRedOpts Env
    env Stm GPU
    ker_snd
          case mb_stms_after of
            Just Stms GPU
    stms_after -> Maybe (Stms GPU) -> GenRedM (Maybe (Stms GPU))
    forall a. a -> ReaderT (Scope GPU) (State VNameSource) a
    forall (f :: * -> *) a. Applicative f => a -> f a
    pure (Maybe (Stms GPU) -> GenRedM (Maybe (Stms GPU)))
    -> Maybe (Stms GPU) -> GenRedM (Maybe (Stms GPU))
    forall a b. (a -> b) -> a -> b
    $ Stms GPU -> Maybe (Stms GPU)
    forall a. a -> Maybe a
    Just (Stms GPU -> Maybe (Stms GPU)) -> Stms GPU -> Maybe (Stms GPU)
    forall a b. (a -> b) -> a -> b
    $ Stms GPU
    stms_before Stms GPU -> Stms GPU -> Stms GPU
    forall a. Semigroup a => a -> a -> a
    <> Stms GPU
    stms_after
            Maybe (Stms GPU)
    Nothing -> Maybe (Stms GPU) -> GenRedM (Maybe (Stms GPU))
    forall a. a -> ReaderT (Scope GPU) (State VNameSource) a
    forall (f :: * -> *) a. Applicative f => a -> f a
    pure (Maybe (Stms GPU) -> GenRedM (Maybe (Stms GPU)))
    -> Maybe (Stms GPU) -> GenRedM (Maybe (Stms GPU))
    forall a b. (a -> b) -> a -> b
    $ Stms GPU -> Maybe (Stms GPU)
    forall a. a -> Maybe a
    Just (Stms GPU -> Maybe (Stms GPU)) -> Stms GPU -> Maybe (Stms GPU)
    forall a b. (a -> b) -> a -> b
    $ Stms GPU
    stms_before Stms GPU -> Stms GPU -> Stms GPU
    forall a. Semigroup a => a -> a -> a
    <> Stm GPU -> Stms GPU
    forall rep. Stm rep -> Stms rep
    oneStm Stm GPU
    ker_snd

    genRed2Tile2d :: Env -> Stm GPU -> GenRedM (Maybe (Stms GPU, Stm GPU))
    genRed2Tile2d :: Env -> Stm GPU -> GenRedM (Maybe (Stms GPU, Stm GPU))
    genRed2Tile2d Env
    env kerstm :: Stm GPU
    kerstm@(Let Pat (LetDec GPU)
    pat_ker StmAux (ExpDec GPU)
    aux (Op (SegOp (SegMap SegLevel
    seg_thd SegSpace
    seg_space [Type]
    kres_tps KernelBody GPU
    old_kbody))))
      | SegThread SegVirt
    _novirt Maybe KernelGrid
    _ <- SegLevel
    seg_thd,
        -- novirt == SegNoVirtFull || novirt == SegNoVirt,
        KernelBody () Stms GPU
    kstms [KernelResult]
    kres <- KernelBody GPU
    old_kbody,
        Just ([VName]
    css, [SubExp]
    r_ses) <- [KernelResult] -> Maybe ([VName], [SubExp])
    allGoodReturns [KernelResult]
    kres,
        [VName] -> Bool
    forall a. [a] -> Bool
    forall (t :: * -> *) a. Foldable t => t a -> Bool
    null [VName]
    css,
        -- build the variance table, that records, for
        -- each variable name, the variables it depends on
        VarianceTable
    initial_variance <- (NameInfo (ZonkAny 0) -> Names)
    -> Map VName (NameInfo (ZonkAny 0)) -> VarianceTable
    forall a b k. (a -> b) -> Map k a -> Map k b
    M.map NameInfo (ZonkAny 0) -> Names
    forall a. Monoid a => a
    mempty (Map VName (NameInfo (ZonkAny 0)) -> VarianceTable)
    -> Map VName (NameInfo (ZonkAny 0)) -> VarianceTable
    forall a b. (a -> b) -> a -> b
    $ SegSpace -> Map VName (NameInfo (ZonkAny 0))
    forall rep. SegSpace -> Scope rep
    scopeOfSegSpace SegSpace
    seg_space,
        VarianceTable
    variance <- VarianceTable -> Stms GPU -> VarianceTable
    varianceInStms VarianceTable
    initial_variance Stms GPU
    kstms,
        -- check that the code fits the pattern having:
        -- some `code1`, followed by one accumulation, followed by some `code2`
        -- UpdateAcc VName [SubExp] [SubExp]
        (Stms GPU
    code1, Just Stm GPU
    accum_stmt, Stms GPU
    code2) <- Stms GPU -> (Stms GPU, Maybe (Stm GPU), Stms GPU)
    matchCodeAccumCode Stms GPU
    kstms,
        Let Pat (LetDec GPU)
    pat_accum StmAux (ExpDec GPU)
    _aux_acc (BasicOp (UpdateAcc Safety
    safety VName
    acc_nm [SubExp]
    acc_inds [SubExp]
    acc_vals)) <- Stm GPU
    accum_stmt,
        [VName
    pat_acc_nm] <- Pat Type -> [VName]
    forall dec. Pat dec -> [VName]
    patNames Pat Type
    Pat (LetDec GPU)
    pat_accum,
        -- check that the `acc_inds` are invariant to at least one
        -- parallel kernel dimensions, and return the innermost such one:
        Just (VName
    invar_gid, Int
    gid_ind) <- Names
    -> SegSpace -> VarianceTable -> [SubExp] -> Maybe (VName, Int)
    isInvarToParDim Names
    forall a. Monoid a => a
    mempty SegSpace
    seg_space VarianceTable
    variance [SubExp]
    acc_inds,
        [(VName, SubExp)]
    gid_dims_new_0 <- ((VName, SubExp) -> Bool) -> [(VName, SubExp)] -> [(VName, SubExp)]
    forall a. (a -> Bool) -> [a] -> [a]
    filter (\(VName, SubExp)
    x -> VName
    invar_gid VName -> VName -> Bool
    forall a. Eq a => a -> a -> Bool
    /= (VName, SubExp) -> VName
    forall a b. (a, b) -> a
    fst (VName, SubExp)
    x) (SegSpace -> [(VName, SubExp)]
    unSegSpace SegSpace
    seg_space),
        -- reorder the variant dimensions such that inner(most) accum-indices
        -- correspond to inner(most) parallel dimensions, so that the babysitter
        -- does not introduce transpositions
        -- gid_dims_new <- gid_dims_new_0,
        [(VName, SubExp)]
    gid_dims_new <- VarianceTable -> [SubExp] -> [(VName, SubExp)] -> [(VName, SubExp)]
    forall {b}.
    VarianceTable -> [SubExp] -> [(VName, b)] -> [(VName, b)]
    reorderParDims VarianceTable
    variance [SubExp]
    acc_inds [(VName, SubExp)]
    gid_dims_new_0,
        -- check that all global-memory accesses in `code1` on which
        --   `accum_stmt` depends on are invariant to at least one of
        --   the remaining parallel dimensions (i.e., excluding `invar_gid`)
        (Stm GPU -> Bool) -> [Stm GPU] -> Bool
    forall (t :: * -> *) a. Foldable t => (a -> Bool) -> t a -> Bool
    all (VName
    -> [(VName, SubExp)] -> VarianceTable -> VName -> Stm GPU -> Bool
    isTileable VName
    invar_gid [(VName, SubExp)]
    gid_dims_new VarianceTable
    variance VName
    pat_acc_nm) (Stms GPU -> [Stm GPU]
    forall rep. Stms rep -> [Stm rep]
    stmsToList Stms GPU
    code1),
        -- need to establish a cost model for the stms that would now
        --   be redundantly executed by the two kernels. If any recurence
        --   is redundant than it is a no go. Otherwise we need to look at
        --   memory accesses: if more than two are re-executed, then we
        --   should abort.
        Cost
    cost <- VarianceTable -> VName -> [SubExp] -> Stms GPU -> Cost
    costRedundantExecution VarianceTable
    variance VName
    pat_acc_nm [SubExp]
    r_ses Stms GPU
    kstms,
        Cost -> Cost -> Cost
    maxCost Cost
    cost (Int -> Cost
    Small Int
    2) Cost -> Cost -> Bool
    forall a. Eq a => a -> a -> Bool
    == Int -> Cost
    Small Int
    2 = do
          -- 1. create the first kernel
          acc_tp <- VName -> ReaderT (Scope GPU) (State VNameSource) Type
    forall rep (m :: * -> *). HasScope rep m => VName -> m Type
    lookupType VName
    acc_nm
          let inv_dim_len = SegSpace -> [SubExp]
    segSpaceDims SegSpace
    seg_space [SubExp] -> Int -> SubExp
    forall a. HasCallStack => [a] -> Int -> a
    !! Int
    gid_ind
              -- 1.1. get the accumulation operator
              ((redop0, neutral), el_tps) = getAccLambda acc_tp
          redop <- renameLambda redop0
          let red =
                Reduce
                  { redComm :: Commutativity
    redComm = Commutativity
    Commutative,
                    redLambda :: Lambda GPU
    redLambda = Lambda GPU
    redop,
                    redNeutral :: [SubExp]
    redNeutral = [SubExp]
    neutral
                  }
              -- 1.2. build the sequential map-reduce screma
              code1' =
                [Stm GPU] -> Stms GPU
    forall rep. [Stm rep] -> Stms rep
    stmsFromList ([Stm GPU] -> Stms GPU) -> [Stm GPU] -> Stms GPU
    forall a b. (a -> b) -> a -> b
    $
                  (Stm GPU -> Bool) -> [Stm GPU] -> [Stm GPU]
    forall a. (a -> Bool) -> [a] -> [a]
    filter (VName -> VarianceTable -> Stm GPU -> Bool
    forall {k} {rep}. Ord k => k -> Map k Names -> Stm rep -> Bool
    dependsOnAcc VName
    pat_acc_nm VarianceTable
    variance) ([Stm GPU] -> [Stm GPU]) -> [Stm GPU] -> [Stm GPU]
    forall a b. (a -> b) -> a -> b
    $
                    Stms GPU -> [Stm GPU]
    forall rep. Stms rep -> [Stm rep]
    stmsToList Stms GPU
    code1
          (code1'', code1_tr_host) <- transposeFVs (freeIn kerstm) variance invar_gid code1'
          let map_lam_body = Stms GPU -> Result -> Body GPU
    forall rep. Buildable rep => Stms rep -> Result -> Body rep
    mkBody Stms GPU
    code1'' (Result -> Body GPU) -> Result -> Body GPU
    forall a b. (a -> b) -> a -> b
    $ (SubExp -> SubExpRes) -> [SubExp] -> Result
    forall a b. (a -> b) -> [a] -> [b]
    map (Certs -> SubExp -> SubExpRes
    SubExpRes ([VName] -> Certs
    Certs [])) [SubExp]
    acc_vals
              map_lam0 = [LParam GPU] -> [Type] -> Body GPU -> Lambda GPU
    forall rep. [LParam rep] -> [Type] -> Body rep -> Lambda rep
    Lambda [Attrs -> VName -> Type -> Param Type
    forall dec. Attrs -> VName -> dec -> Param dec
    Param Attrs
    forall a. Monoid a => a
    mempty VName
    invar_gid (PrimType -> Type
    forall shape u. PrimType -> TypeBase shape u
    Prim PrimType
    int64)] [Type]
    el_tps Body GPU
    map_lam_body
          map_lam <- renameLambda map_lam0
          (k1_res, ker1_stms) <- runBuilderT' $ do
            iota <- letExp "iota" $ BasicOp $ Iota inv_dim_len (intConst Int64 0) (intConst Int64 1) Int64
            let op_exp = OpC GPU GPU -> Exp GPU
    forall rep. Op rep -> Exp rep
    Op (SOAC GPU -> HostOp SOAC GPU
    forall (op :: * -> *) rep. op rep -> HostOp op rep
    OtherOp (SubExp -> [VName] -> ScremaForm GPU -> SOAC GPU
    forall rep. SubExp -> [VName] -> ScremaForm rep -> SOAC rep
    Screma SubExp
    inv_dim_len [VName
    iota] (Lambda GPU -> [Scan GPU] -> [Reduce GPU] -> ScremaForm GPU
    forall rep.
    Lambda rep -> [Scan rep] -> [Reduce rep] -> ScremaForm rep
    ScremaForm Lambda GPU
    map_lam [] [Reduce GPU
    red])))
            res_redmap <- letTupExp "res_mapred" op_exp
            letSubExp (baseString pat_acc_nm ++ "_big_update") $
              BasicOp (UpdateAcc safety acc_nm acc_inds $ map Var res_redmap)

          -- 1.3. build the kernel expression and rename it!
          gid_flat_1 <- newVName "gid_flat"
          let space1 = VName -> [(VName, SubExp)] -> SegSpace
    SegSpace VName
    gid_flat_1 [(VName, SubExp)]
    gid_dims_new

          let level1 = SegVirt -> Maybe KernelGrid -> SegLevel
    SegThread (SegSeqDims -> SegVirt
    SegNoVirtFull ([Int] -> SegSeqDims
    SegSeqDims [])) Maybe KernelGrid
    forall a. Maybe a
    Nothing -- novirt ?
              kbody1 = BodyDec GPU -> Stms GPU -> [KernelResult] -> KernelBody GPU
    forall rep.
    BodyDec rep -> Stms rep -> [KernelResult] -> KernelBody rep
    KernelBody () Stms GPU
    ker1_stms [ResultManifest -> Certs -> SubExp -> KernelResult
    Returns ResultManifest
    ResultMaySimplify ([VName] -> Certs
    Certs []) SubExp
    k1_res]

          -- is it OK here to use the "aux" from the parrent kernel?
          ker_exp <- renameExp $ Op (SegOp (SegMap level1 space1 [acc_tp] kbody1))
          let ker1 = 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_accum StmAux (ExpDec GPU)
    aux Exp GPU
    ker_exp

          -- 2 build the second kernel
          let ker2_body = KernelBody GPU
    old_kbody {kernelBodyStms = code1 <> code2}
          ker2_exp <- renameExp $ Op (SegOp (SegMap seg_thd seg_space kres_tps ker2_body))
          let ker2 = 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_ker StmAux (ExpDec GPU)
    aux Exp GPU
    ker2_exp
          pure $
            Just (code1_tr_host <> oneStm ker1, ker2)
      where
        isIndVarToParDim :: VarianceTable -> SubExp -> (VName, b) -> Bool
    isIndVarToParDim VarianceTable
    _ (Constant PrimValue
    _) (VName, b)
    _ = Bool
    False
        isIndVarToParDim VarianceTable
    variance (Var VName
    acc_ind) (VName, b)
    par_dim =
          VName
    acc_ind VName -> VName -> Bool
    forall a. Eq a => a -> a -> Bool
    == (VName, b) -> VName
    forall a b. (a, b) -> a
    fst (VName, b)
    par_dim
            Bool -> Bool -> Bool
    || VName -> Names -> Bool
    nameIn ((VName, b) -> VName
    forall a b. (a, b) -> a
    fst (VName, b)
    par_dim) (Names -> VName -> VarianceTable -> Names
    forall k a. Ord k => a -> k -> Map k a -> a
    M.findWithDefault Names
    forall a. Monoid a => a
    mempty VName
    acc_ind VarianceTable
    variance)
        foldfunReorder :: VarianceTable
    -> ([(VName, b)], [(VName, b)])
    -> SubExp
    -> ([(VName, b)], [(VName, b)])
    foldfunReorder VarianceTable
    variance ([(VName, b)]
    unused_dims, [(VName, b)]
    inner_dims) SubExp
    acc_ind =
          case ((VName, b) -> Bool) -> [(VName, b)] -> Maybe Int
    forall a. (a -> Bool) -> [a] -> Maybe Int
    L.findIndex (VarianceTable -> SubExp -> (VName, b) -> Bool
    forall {b}. VarianceTable -> SubExp -> (VName, b) -> Bool
    isIndVarToParDim VarianceTable
    variance SubExp
    acc_ind) [(VName, b)]
    unused_dims of
            Maybe Int
    Nothing -> ([(VName, b)]
    unused_dims, [(VName, b)]
    inner_dims)
            Just Int
    i ->
              ( Int -> [(VName, b)] -> [(VName, b)]
    forall a. Int -> [a] -> [a]
    take Int
    i [(VName, b)]
    unused_dims [(VName, b)] -> [(VName, b)] -> [(VName, b)]
    forall a. [a] -> [a] -> [a]
    ++ Int -> [(VName, b)] -> [(VName, b)]
    forall a. Int -> [a] -> [a]
    drop (Int
    i Int -> Int -> Int
    forall a. Num a => a -> a -> a
    + Int
    1) [(VName, b)]
    unused_dims,
                ([(VName, b)]
    unused_dims [(VName, b)] -> Int -> (VName, b)
    forall a. HasCallStack => [a] -> Int -> a
    !! Int
    i) (VName, b) -> [(VName, b)] -> [(VName, b)]
    forall a. a -> [a] -> [a]
    : [(VName, b)]
    inner_dims
              )
        reorderParDims :: VarianceTable -> [SubExp] -> [(VName, b)] -> [(VName, b)]
    reorderParDims VarianceTable
    variance [SubExp]
    acc_inds [(VName, b)]
    gid_dims_new_0 =
          let ([(VName, b)]
    invar_dims, [(VName, b)]
    inner_dims) =
                (([(VName, b)], [(VName, b)])
     -> SubExp -> ([(VName, b)], [(VName, b)]))
    -> ([(VName, b)], [(VName, b)])
    -> [SubExp]
    -> ([(VName, b)], [(VName, b)])
    forall b a. (b -> a -> b) -> b -> [a] -> b
    forall (t :: * -> *) b a.
    Foldable t =>
    (b -> a -> b) -> b -> t a -> b
    foldl
                  (VarianceTable
    -> ([(VName, b)], [(VName, b)])
    -> SubExp
    -> ([(VName, b)], [(VName, b)])
    forall {b}.
    VarianceTable
    -> ([(VName, b)], [(VName, b)])
    -> SubExp
    -> ([(VName, b)], [(VName, b)])
    foldfunReorder VarianceTable
    variance)
                  ([(VName, b)]
    gid_dims_new_0, [])
                  ([SubExp] -> [SubExp]
    forall a. [a] -> [a]
    reverse [SubExp]
    acc_inds)
           in [(VName, b)]
    invar_dims [(VName, b)] -> [(VName, b)] -> [(VName, b)]
    forall a. [a] -> [a] -> [a]
    ++ [(VName, b)]
    inner_dims
        --
        getAccLambda :: Type -> ((Lambda GPU, [SubExp]), [Type])
    getAccLambda Type
    acc_tp =
          case Type
    acc_tp of
            (Acc VName
    tp_id ShapeBase SubExp
    _shp [Type]
    el_tps NoUniqueness
    _) ->
              case VName -> WithEnv -> Maybe (Lambda GPU, [SubExp])
    forall k a. Ord k => k -> Map k a -> Maybe a
    M.lookup VName
    tp_id (Env -> WithEnv
    forall a b. (a, b) -> a
    fst Env
    env) of
                Just (Lambda GPU, [SubExp])
    lam -> ((Lambda GPU, [SubExp])
    lam, [Type]
    el_tps)
                Maybe (Lambda GPU, [SubExp])
    _ -> [Char] -> ((Lambda GPU, [SubExp]), [Type])
    forall a. HasCallStack => [Char] -> a
    error ([Char] -> ((Lambda GPU, [SubExp]), [Type]))
    -> [Char] -> ((Lambda GPU, [SubExp]), [Type])
    forall a b. (a -> b) -> a -> b
    $ [Char]
    "Lookup in environment failed! " [Char] -> [Char] -> [Char]
    forall a. [a] -> [a] -> [a]
    ++ VName -> [Char]
    forall a. Pretty a => a -> [Char]
    prettyString VName
    tp_id [Char] -> [Char] -> [Char]
    forall a. [a] -> [a] -> [a]
    ++ [Char]
    " env: " [Char] -> [Char] -> [Char]
    forall a. [a] -> [a] -> [a]
    ++ WithEnv -> [Char]
    forall a. Show a => a -> [Char]
    show (Env -> WithEnv
    forall a b. (a, b) -> a
    fst Env
    env)
            Type
    _ -> [Char] -> ((Lambda GPU, [SubExp]), [Type])
    forall a. HasCallStack => [Char] -> a
    error [Char]
    "Illegal accumulator type!"
        -- is a subexp invariant to a gid of a parallel dimension?
        isSeInvar2 :: VarianceTable -> VName -> SubExp -> Bool
    isSeInvar2 VarianceTable
    variance VName
    gid (Var VName
    x) =
          let x_deps :: Names
    x_deps = Names -> VName -> VarianceTable -> Names
    forall k a. Ord k => a -> k -> Map k a -> a
    M.findWithDefault Names
    forall a. Monoid a => a
    mempty VName
    x VarianceTable
    variance
           in VName
    gid VName -> VName -> Bool
    forall a. Eq a => a -> a -> Bool
    /= VName
    x Bool -> Bool -> Bool
    && VName
    gid VName -> Names -> Bool
    `notNameIn` Names
    x_deps
        isSeInvar2 VarianceTable
    _ VName
    _ SubExp
    _ = Bool
    True
        -- is a DimIndex invar to a gid of a parallel dimension?
        isDimIdxInvar2 :: VarianceTable -> VName -> DimIndex SubExp -> Bool
    isDimIdxInvar2 VarianceTable
    variance VName
    gid (DimFix SubExp
    d) =
          VarianceTable -> VName -> SubExp -> Bool
    isSeInvar2 VarianceTable
    variance VName
    gid SubExp
    d
        isDimIdxInvar2 VarianceTable
    variance VName
    gid (DimSlice SubExp
    d1 SubExp
    d2 SubExp
    d3) =
          (SubExp -> Bool) -> [SubExp] -> Bool
    forall (t :: * -> *) a. Foldable t => (a -> Bool) -> t a -> Bool
    all (VarianceTable -> VName -> SubExp -> Bool
    isSeInvar2 VarianceTable
    variance VName
    gid) [SubExp
    d1, SubExp
    d2, SubExp
    d3]
        -- is an entire slice invariant to at least one gid of a parallel dimension
        isSliceInvar2 :: VarianceTable -> Slice SubExp -> t VName -> Bool
    isSliceInvar2 VarianceTable
    variance Slice SubExp
    slc =
          (VName -> Bool) -> t VName -> Bool
    forall (t :: * -> *) a. Foldable t => (a -> Bool) -> t a -> Bool
    any (\VName
    gid -> (DimIndex SubExp -> Bool) -> [DimIndex SubExp] -> Bool
    forall (t :: * -> *) a. Foldable t => (a -> Bool) -> t a -> Bool
    all (VarianceTable -> VName -> DimIndex SubExp -> Bool
    isDimIdxInvar2 VarianceTable
    variance VName
    gid) (Slice SubExp -> [DimIndex SubExp]
    forall d. Slice d -> [DimIndex d]
    unSlice Slice SubExp
    slc))
        -- are all statements that touch memory invariant to at least one parallel dimension?
        isTileable :: VName -> [(VName, SubExp)] -> VarianceTable -> VName -> Stm GPU -> Bool
        isTileable :: VName
    -> [(VName, SubExp)] -> VarianceTable -> VName -> Stm GPU -> Bool
    isTileable VName
    seq_gid [(VName, SubExp)]
    gid_dims VarianceTable
    variance VName
    acc_nm (Let (Pat [PatElem (LetDec GPU)
    pel]) StmAux (ExpDec GPU)
    _ (BasicOp (Index VName
    _ Slice SubExp
    slc)))
          | Names
    acc_deps <- Names -> VName -> VarianceTable -> Names
    forall k a. Ord k => a -> k -> Map k a -> a
    M.findWithDefault Names
    forall a. Monoid a => a
    mempty VName
    acc_nm VarianceTable
    variance,
            PatElem Type -> VName
    forall dec. PatElem dec -> VName
    patElemName PatElem Type
    PatElem (LetDec GPU)
    pel VName -> Names -> Bool
    `nameIn` Names
    acc_deps =
              let invar_par :: Bool
    invar_par = VarianceTable -> Slice SubExp -> [VName] -> Bool
    forall {t :: * -> *}.
    Foldable t =>
    VarianceTable -> Slice SubExp -> t VName -> Bool
    isSliceInvar2 VarianceTable
    variance Slice SubExp
    slc (((VName, SubExp) -> VName) -> [(VName, SubExp)] -> [VName]
    forall a b. (a -> b) -> [a] -> [b]
    map (VName, SubExp) -> VName
    forall a b. (a, b) -> a
    fst [(VName, SubExp)]
    gid_dims)
                  invar_seq :: Bool
    invar_seq = VarianceTable -> Slice SubExp -> [VName] -> Bool
    forall {t :: * -> *}.
    Foldable t =>
    VarianceTable -> Slice SubExp -> t VName -> Bool
    isSliceInvar2 VarianceTable
    variance Slice SubExp
    slc [VName
    seq_gid]
               in Bool
    invar_par Bool -> Bool -> Bool
    || Bool
    invar_seq
        -- this relies on the cost model, that currently accepts only
        -- global-memory reads, and for example rejects in-place updates
        -- or loops inside the code that is transformed in a redomap.
        isTileable VName
    _ [(VName, SubExp)]
    _ VarianceTable
    _ VName
    _ Stm GPU
    _ = Bool
    True
        -- does the to-be-reduced accumulator depends on this statement?
        dependsOnAcc :: k -> Map k Names -> Stm rep -> Bool
    dependsOnAcc k
    pat_acc_nm Map k Names
    variance (Let Pat (LetDec rep)
    pat StmAux (ExpDec rep)
    _ Exp rep
    _) =
          let acc_deps :: Names
    acc_deps = Names -> k -> Map k Names -> Names
    forall k a. Ord k => a -> k -> Map k a -> a
    M.findWithDefault Names
    forall a. Monoid a => a
    mempty k
    pat_acc_nm Map k Names
    variance
           in (VName -> Bool) -> [VName] -> Bool
    forall (t :: * -> *) a. Foldable t => (a -> Bool) -> t a -> Bool
    any (VName -> Names -> Bool
    `nameIn` Names
    acc_deps) ([VName] -> Bool) -> [VName] -> Bool
    forall a b. (a -> b) -> a -> b
    $ Pat (LetDec rep) -> [VName]
    forall dec. Pat dec -> [VName]
    patNames Pat (LetDec rep)
    pat
    genRed2Tile2d Env
    _ Stm GPU
    _ =
      Maybe (Stms GPU, Stm GPU) -> GenRedM (Maybe (Stms GPU, Stm GPU))
    forall a. a -> ReaderT (Scope GPU) (State VNameSource) a
    forall (f :: * -> *) a. Applicative f => a -> f a
    pure Maybe (Stms GPU, Stm GPU)
    forall a. Maybe a
    Nothing

    genRed2SegRed :: Env -> Stm GPU -> GenRedM (Maybe (Stms GPU, Stm GPU))
    genRed2SegRed :: Env -> Stm GPU -> GenRedM (Maybe (Stms GPU, Stm GPU))
    genRed2SegRed Env
    _ Stm GPU
    _ =
      Maybe (Stms GPU, Stm GPU) -> GenRedM (Maybe (Stms GPU, Stm GPU))
    forall a. a -> ReaderT (Scope GPU) (State VNameSource) a
    forall (f :: * -> *) a. Applicative f => a -> f a
    pure Maybe (Stms GPU, Stm GPU)
    forall a. Maybe a
    Nothing

    transposeFVs ::
      Names ->
      VarianceTable ->
      VName ->
      Stms GPU ->
      GenRedM (Stms GPU, Stms GPU)
    transposeFVs :: Names
    -> VarianceTable
    -> VName
    -> Stms GPU
    -> GenRedM (Stms GPU, Stms GPU)
    transposeFVs Names
    fvs VarianceTable
    variance VName
    gid Stms GPU
    stms = do
      (tab, stms') <- ((Map VName ([Int], VName, Stms GPU), Stms GPU)
     -> Stm GPU
     -> ReaderT
          (Scope GPU)
          (State VNameSource)
          (Map VName ([Int], VName, Stms GPU), Stms GPU))
    -> (Map VName ([Int], VName, Stms GPU), Stms GPU)
    -> [Stm GPU]
    -> ReaderT
         (Scope GPU)
         (State VNameSource)
         (Map VName ([Int], VName, Stms GPU), Stms GPU)
    forall (t :: * -> *) (m :: * -> *) b a.
    (Foldable t, Monad m) =>
    (b -> a -> m b) -> b -> t a -> m b
    foldM (Map VName ([Int], VName, Stms GPU), Stms GPU)
    -> Stm GPU
    -> ReaderT
         (Scope GPU)
         (State VNameSource)
         (Map VName ([Int], VName, Stms GPU), Stms GPU)
    foldfun (Map VName ([Int], VName, Stms GPU)
    forall k a. Map k a
    M.empty, Stms GPU
    forall a. Monoid a => a
    mempty) ([Stm GPU]
     -> ReaderT
          (Scope GPU)
          (State VNameSource)
          (Map VName ([Int], VName, Stms GPU), Stms GPU))
    -> [Stm GPU]
    -> ReaderT
         (Scope GPU)
         (State VNameSource)
         (Map VName ([Int], VName, Stms GPU), Stms GPU)
    forall a b. (a -> b) -> a -> b
    $ Stms GPU -> [Stm GPU]
    forall rep. Stms rep -> [Stm rep]
    stmsToList Stms GPU
    stms
      let stms_host = (([Int], VName, Stms GPU) -> Stms GPU -> Stms GPU)
    -> Stms GPU -> Map VName ([Int], VName, Stms GPU) -> Stms GPU
    forall a b k. (a -> b -> b) -> b -> Map k a -> b
    M.foldr (\([Int]
    _, VName
    _, Stms GPU
    s) Stms GPU
    ss -> Stms GPU
    ss Stms GPU -> Stms GPU -> Stms GPU
    forall a. Semigroup a => a -> a -> a
    <> Stms GPU
    s) Stms GPU
    forall a. Monoid a => a
    mempty Map VName ([Int], VName, Stms GPU)
    tab
      pure (stms', stms_host)
      where
        foldfun :: (Map VName ([Int], VName, Stms GPU), Stms GPU)
    -> Stm GPU
    -> ReaderT
         (Scope GPU)
         (State VNameSource)
         (Map VName ([Int], VName, Stms GPU), Stms GPU)
    foldfun (Map VName ([Int], VName, Stms GPU)
    tab, Stms GPU
    all_stms) Stm GPU
    stm = do
          (tab', stm') <- (Map VName ([Int], VName, Stms GPU), Stm GPU)
    -> ReaderT
         (Scope GPU)
         (State VNameSource)
         (Map VName ([Int], VName, Stms GPU), Stm GPU)
    transposeFV (Map VName ([Int], VName, Stms GPU)
    tab, Stm GPU
    stm)
          pure (tab', all_stms <> oneStm stm')
        -- ToDo: currently handles only 2-dim arrays, please generalize
        transposeFV :: (Map VName ([Int], VName, Stms GPU), Stm GPU)
    -> ReaderT
         (Scope GPU)
         (State VNameSource)
         (Map VName ([Int], VName, Stms GPU), Stm GPU)
    transposeFV (Map VName ([Int], VName, Stms GPU)
    tab, Let Pat (LetDec GPU)
    pat StmAux (ExpDec GPU)
    aux (BasicOp (Index VName
    arr Slice SubExp
    slc)))
          | [DimIndex SubExp]
    dims <- Slice SubExp -> [DimIndex SubExp]
    forall d. Slice d -> [DimIndex d]
    unSlice Slice SubExp
    slc,
            (DimIndex SubExp -> Bool) -> [DimIndex SubExp] -> Bool
    forall (t :: * -> *) a. Foldable t => (a -> Bool) -> t a -> Bool
    all DimIndex SubExp -> Bool
    forall {d}. DimIndex d -> Bool
    isFixDim [DimIndex SubExp]
    dims,
            VName
    arr VName -> Names -> Bool
    `nameIn` Names
    fvs,
            [Int]
    iis <- (DimIndex SubExp -> Bool) -> [DimIndex SubExp] -> [Int]
    forall a. (a -> Bool) -> [a] -> [Int]
    L.findIndices DimIndex SubExp -> Bool
    depOnGid [DimIndex SubExp]
    dims,
            [Int
    ii] <- [Int]
    iis,
            -- generalize below: treat any rearange and add to tab if not there.
            Maybe ([Int], VName, Stms GPU)
    Nothing <- VName
    -> Map VName ([Int], VName, Stms GPU)
    -> Maybe ([Int], VName, Stms GPU)
    forall k a. Ord k => k -> Map k a -> Maybe a
    M.lookup VName
    arr Map VName ([Int], VName, Stms GPU)
    tab,
            Int
    ii Int -> Int -> Bool
    forall a. Eq a => a -> a -> Bool
    /= [DimIndex SubExp] -> Int
    forall a. [a] -> Int
    forall (t :: * -> *) a. Foldable t => t a -> Int
    length [DimIndex SubExp]
    dims Int -> Int -> Int
    forall a. Num a => a -> a -> a
    - Int
    1,
            [Int]
    perm <- [Int
    0 .. Int
    ii Int -> Int -> Int
    forall a. Num a => a -> a -> a
    - Int
    1] [Int] -> [Int] -> [Int]
    forall a. [a] -> [a] -> [a]
    ++ [Int
    ii Int -> Int -> Int
    forall a. Num a => a -> a -> a
    + Int
    1 .. [DimIndex SubExp] -> Int
    forall a. [a] -> Int
    forall (t :: * -> *) a. Foldable t => t a -> Int
    length [DimIndex SubExp]
    dims Int -> Int -> Int
    forall a. Num a => a -> a -> a
    - Int
    1] [Int] -> [Int] -> [Int]
    forall a. [a] -> [a] -> [a]
    ++ [Int
    ii] = do
              (arr_tr, stms_tr) <- BuilderT GPU (ReaderT (Scope GPU) (State VNameSource)) VName
    -> ReaderT (Scope GPU) (State VNameSource) (VName, Stms GPU)
    forall (m :: * -> *) somerep rep a.
    (MonadFreshNames m, HasScope somerep m, SameScope somerep rep) =>
    BuilderT rep m a -> m (a, Stms rep)
    runBuilderT' (BuilderT GPU (ReaderT (Scope GPU) (State VNameSource)) VName
     -> ReaderT (Scope GPU) (State VNameSource) (VName, Stms GPU))
    -> BuilderT GPU (ReaderT (Scope GPU) (State VNameSource)) VName
    -> ReaderT (Scope GPU) (State VNameSource) (VName, Stms GPU)
    forall a b. (a -> b) -> a -> b
    $ do
                arr' <- [Char]
    -> Exp
         (Rep (BuilderT GPU (ReaderT (Scope GPU) (State VNameSource))))
    -> BuilderT GPU (ReaderT (Scope GPU) (State VNameSource)) VName
    forall (m :: * -> *).
    MonadBuilder m =>
    [Char] -> Exp (Rep m) -> m VName
    letExp (VName -> [Char]
    baseString VName
    arr [Char] -> [Char] -> [Char]
    forall a. [a] -> [a] -> [a]
    ++ [Char]
    "_trsp") (Exp (Rep (BuilderT GPU (ReaderT (Scope GPU) (State VNameSource))))
     -> BuilderT GPU (ReaderT (Scope GPU) (State VNameSource)) VName)
    -> Exp
         (Rep (BuilderT GPU (ReaderT (Scope GPU) (State VNameSource))))
    -> BuilderT GPU (ReaderT (Scope GPU) (State VNameSource)) VName
    forall a b. (a -> b) -> a -> b
    $ BasicOp
    -> Exp
         (Rep (BuilderT GPU (ReaderT (Scope GPU) (State VNameSource))))
    forall rep. BasicOp -> Exp rep
    BasicOp (BasicOp
     -> Exp
          (Rep (BuilderT GPU (ReaderT (Scope GPU) (State VNameSource)))))
    -> BasicOp
    -> Exp
         (Rep (BuilderT GPU (ReaderT (Scope GPU) (State VNameSource))))
    forall a b. (a -> b) -> a -> b
    $ VName -> [Int] -> BasicOp
    Rearrange VName
    arr [Int]
    perm
                letExp (baseString arr' ++ "_opaque") $ BasicOp $ Opaque OpaqueNil $ Var arr'
              let tab' = VName
    -> ([Int], VName, Stms GPU)
    -> Map VName ([Int], VName, Stms GPU)
    -> Map VName ([Int], VName, Stms GPU)
    forall k a. Ord k => k -> a -> Map k a -> Map k a
    M.insert VName
    arr ([Int]
    perm, VName
    arr_tr, Stms GPU
    stms_tr) Map VName ([Int], VName, Stms GPU)
    tab
                  slc' = [DimIndex SubExp] -> Slice SubExp
    forall d. [DimIndex d] -> Slice d
    Slice ([DimIndex SubExp] -> Slice SubExp)
    -> [DimIndex SubExp] -> Slice SubExp
    forall a b. (a -> b) -> a -> b
    $ (Int -> DimIndex SubExp) -> [Int] -> [DimIndex SubExp]
    forall a b. (a -> b) -> [a] -> [b]
    map ([DimIndex SubExp]
    dims [DimIndex SubExp] -> Int -> DimIndex SubExp
    forall a. HasCallStack => [a] -> Int -> a
    !!) [Int]
    perm
                  stm' = 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
    $ BasicOp -> Exp GPU
    forall rep. BasicOp -> Exp rep
    BasicOp (BasicOp -> Exp GPU) -> BasicOp -> Exp GPU
    forall a b. (a -> b) -> a -> b
    $ VName -> Slice SubExp -> BasicOp
    Index VName
    arr_tr Slice SubExp
    slc'
              pure (tab', stm')
          where
            isFixDim :: DimIndex d -> Bool
    isFixDim DimFix {} = Bool
    True
            isFixDim DimIndex d
    _ = Bool
    False
            depOnGid :: DimIndex SubExp -> Bool
    depOnGid (DimFix (Var VName
    nm)) =
              VName
    gid VName -> VName -> Bool
    forall a. Eq a => a -> a -> Bool
    == VName
    nm Bool -> Bool -> Bool
    || VName -> Names -> Bool
    nameIn VName
    gid (Names -> VName -> VarianceTable -> Names
    forall k a. Ord k => a -> k -> Map k a -> a
    M.findWithDefault Names
    forall a. Monoid a => a
    mempty VName
    nm VarianceTable
    variance)
            depOnGid DimIndex SubExp
    _ = Bool
    False
        transposeFV (Map VName ([Int], VName, Stms GPU), Stm GPU)
    r = (Map VName ([Int], VName, Stms GPU), Stm GPU)
    -> ReaderT
         (Scope GPU)
         (State VNameSource)
         (Map VName ([Int], VName, Stms GPU), Stm GPU)
    forall a. a -> ReaderT (Scope GPU) (State VNameSource) a
    forall (f :: * -> *) a. Applicative f => a -> f a
    pure (Map VName ([Int], VName, Stms GPU), Stm GPU)
    r

    -- | Tries to identify the following pattern:
    --   code followed by some UpdateAcc-statement
    --   followed by more code.
    matchCodeAccumCode ::
      Stms GPU ->
      (Stms GPU, Maybe (Stm GPU), Stms GPU)
    matchCodeAccumCode :: Stms GPU -> (Stms GPU, Maybe (Stm GPU), Stms GPU)
    matchCodeAccumCode Stms GPU
    kstms =
      let ([Stm GPU]
    code1, Maybe (Stm GPU)
    screma, [Stm GPU]
    code2) =
            (([Stm GPU], Maybe (Stm GPU), [Stm GPU])
     -> Stm GPU -> ([Stm GPU], Maybe (Stm GPU), [Stm GPU]))
    -> ([Stm GPU], Maybe (Stm GPU), [Stm GPU])
    -> [Stm GPU]
    -> ([Stm GPU], Maybe (Stm GPU), [Stm GPU])
    forall b a. (b -> a -> b) -> b -> [a] -> b
    forall (t :: * -> *) b a.
    Foldable t =>
    (b -> a -> b) -> b -> t a -> b
    foldl
              ( \([Stm GPU], Maybe (Stm GPU), [Stm GPU])
    acc Stm GPU
    stmt ->
                  case (([Stm GPU], Maybe (Stm GPU), [Stm GPU])
    acc, Stm GPU
    stmt) of
                    (([Stm GPU]
    cd1, Maybe (Stm GPU)
    Nothing, [Stm GPU]
    cd2), Let Pat (LetDec GPU)
    _ StmAux (ExpDec GPU)
    _ (BasicOp UpdateAcc {})) ->
                      ([Stm GPU]
    cd1, Stm GPU -> Maybe (Stm GPU)
    forall a. a -> Maybe a
    Just Stm GPU
    stmt, [Stm GPU]
    cd2)
                    (([Stm GPU]
    cd1, Maybe (Stm GPU)
    Nothing, [Stm GPU]
    cd2), Stm GPU
    _) ->
                      ([Stm GPU]
    cd1 [Stm GPU] -> [Stm GPU] -> [Stm GPU]
    forall a. [a] -> [a] -> [a]
    ++ [Stm GPU
    stmt], Maybe (Stm GPU)
    forall a. Maybe a
    Nothing, [Stm GPU]
    cd2)
                    (([Stm GPU]
    cd1, Just Stm GPU
    strm, [Stm GPU]
    cd2), Stm GPU
    _) ->
                      ([Stm GPU]
    cd1, Stm GPU -> Maybe (Stm GPU)
    forall a. a -> Maybe a
    Just Stm GPU
    strm, [Stm GPU]
    cd2 [Stm GPU] -> [Stm GPU] -> [Stm GPU]
    forall a. [a] -> [a] -> [a]
    ++ [Stm GPU
    stmt])
              )
              ([], Maybe (Stm GPU)
    forall a. Maybe a
    Nothing, [])
              (Stms GPU -> [Stm GPU]
    forall rep. Stms rep -> [Stm rep]
    stmsToList Stms GPU
    kstms)
       in ([Stm GPU] -> Stms GPU
    forall rep. [Stm rep] -> Stms rep
    stmsFromList [Stm GPU]
    code1, Maybe (Stm GPU)
    screma, [Stm GPU] -> Stms GPU
    forall rep. [Stm rep] -> Stms rep
    stmsFromList [Stm GPU]
    code2)

    -- | Checks that there exist a parallel dimension (among @kids@),
    --     to which all the indices (@acc_inds@) are invariant to.
    --   It returns the innermost such parallel dimension, as a tuple
    --     of the pardim gid ('VName') and its index ('Int') in the
    --     parallel space.
    isInvarToParDim ::
      Names ->
      SegSpace ->
      VarianceTable ->
      [SubExp] ->
      Maybe (VName, Int)
    isInvarToParDim :: Names
    -> SegSpace -> VarianceTable -> [SubExp] -> Maybe (VName, Int)
    isInvarToParDim Names
    branch_variant SegSpace
    kspace VarianceTable
    variance [SubExp]
    acc_inds =
      let ker_gids :: [VName]
    ker_gids = ((VName, SubExp) -> VName) -> [(VName, SubExp)] -> [VName]
    forall a b. (a -> b) -> [a] -> [b]
    map (VName, SubExp) -> VName
    forall a b. (a, b) -> a
    fst ([(VName, SubExp)] -> [VName]) -> [(VName, SubExp)] -> [VName]
    forall a b. (a -> b) -> a -> b
    $ SegSpace -> [(VName, SubExp)]
    unSegSpace SegSpace
    kspace
          branch_invariant :: Bool
    branch_invariant = (VName -> Bool) -> [VName] -> Bool
    forall (t :: * -> *) a. Foldable t => (a -> Bool) -> t a -> Bool
    all (VName -> Names -> Bool
    `notNameIn` Names
    branch_variant) [VName]
    ker_gids
          allvar2 :: Names
    allvar2 = [SubExp] -> [VName] -> Names
    allvariant2 [SubExp]
    acc_inds [VName]
    ker_gids
          last_invar_dim :: Maybe (VName, Int)
    last_invar_dim =
            (Maybe (VName, Int) -> (VName, Int) -> Maybe (VName, Int))
    -> Maybe (VName, Int) -> [(VName, Int)] -> Maybe (VName, Int)
    forall b a. (b -> a -> b) -> b -> [a] -> b
    forall (t :: * -> *) b a.
    Foldable t =>
    (b -> a -> b) -> b -> t a -> b
    foldl (Names -> Maybe (VName, Int) -> (VName, Int) -> Maybe (VName, Int)
    forall {b}.
    Names -> Maybe (VName, b) -> (VName, b) -> Maybe (VName, b)
    lastNotIn Names
    allvar2) Maybe (VName, Int)
    forall a. Maybe a
    Nothing ([(VName, Int)] -> Maybe (VName, Int))
    -> [(VName, Int)] -> Maybe (VName, Int)
    forall a b. (a -> b) -> a -> b
    $
              [VName] -> [Int] -> [(VName, Int)]
    forall a b. [a] -> [b] -> [(a, b)]
    zip [VName]
    ker_gids [Int
    0 .. [VName] -> Int
    forall a. [a] -> Int
    forall (t :: * -> *) a. Foldable t => t a -> Int
    length [VName]
    ker_gids Int -> Int -> Int
    forall a. Num a => a -> a -> a
    - Int
    1]
       in if Bool
    branch_invariant
            then Maybe (VName, Int)
    last_invar_dim
            else Maybe (VName, Int)
    forall a. Maybe a
    Nothing
      where
        variant2 :: SubExp -> [VName] -> [VName]
    variant2 (Var VName
    ind) [VName]
    kids =
          let variant_to :: Names
    variant_to =
                Names -> VName -> VarianceTable -> Names
    forall k a. Ord k => a -> k -> Map k a -> a
    M.findWithDefault Names
    forall a. Monoid a => a
    mempty VName
    ind VarianceTable
    variance
                  Names -> Names -> Names
    forall a. Semigroup a => a -> a -> a
    <> (if VName
    ind VName -> [VName] -> Bool
    forall a. Eq a => a -> [a] -> Bool
    forall (t :: * -> *) a. (Foldable t, Eq a) => a -> t a -> Bool
    `elem` [VName]
    kids then VName -> Names
    oneName VName
    ind else Names
    forall a. Monoid a => a
    mempty)
           in (VName -> Bool) -> [VName] -> [VName]
    forall a. (a -> Bool) -> [a] -> [a]
    filter (VName -> Names -> Bool
    `nameIn` Names
    variant_to) [VName]
    kids
        variant2 SubExp
    _ [VName]
    _ = []
        allvariant2 :: [SubExp] -> [VName] -> Names
    allvariant2 [SubExp]
    ind_ses [VName]
    kids =
          [VName] -> Names
    namesFromList ([VName] -> Names) -> [VName] -> Names
    forall a b. (a -> b) -> a -> b
    $ (SubExp -> [VName]) -> [SubExp] -> [VName]
    forall (t :: * -> *) a b. Foldable t => (a -> [b]) -> t a -> [b]
    concatMap (SubExp -> [VName] -> [VName]
    `variant2` [VName]
    kids) [SubExp]
    ind_ses
        lastNotIn :: Names -> Maybe (VName, b) -> (VName, b) -> Maybe (VName, b)
    lastNotIn Names
    allvar2 Maybe (VName, b)
    acc (VName
    kid, b
    k) =
          if VName
    kid VName -> Names -> Bool
    `nameIn` Names
    allvar2 then Maybe (VName, b)
    acc else (VName, b) -> Maybe (VName, b)
    forall a. a -> Maybe a
    Just (VName
    kid, b
    k)

    allGoodReturns :: [KernelResult] -> Maybe ([VName], [SubExp])
    allGoodReturns :: [KernelResult] -> Maybe ([VName], [SubExp])
    allGoodReturns [KernelResult]
    kres
      | (KernelResult -> Bool) -> [KernelResult] -> Bool
    forall (t :: * -> *) a. Foldable t => (a -> Bool) -> t a -> Bool
    all KernelResult -> Bool
    goodReturn [KernelResult]
    kres = do
          ([VName], [SubExp]) -> Maybe ([VName], [SubExp])
    forall a. a -> Maybe a
    Just (([VName], [SubExp]) -> Maybe ([VName], [SubExp]))
    -> ([VName], [SubExp]) -> Maybe ([VName], [SubExp])
    forall a b. (a -> b) -> a -> b
    $ (([VName], [SubExp]) -> KernelResult -> ([VName], [SubExp]))
    -> ([VName], [SubExp]) -> [KernelResult] -> ([VName], [SubExp])
    forall b a. (b -> a -> b) -> b -> [a] -> b
    forall (t :: * -> *) b a.
    Foldable t =>
    (b -> a -> b) -> b -> t a -> b
    foldl ([VName], [SubExp]) -> KernelResult -> ([VName], [SubExp])
    addCertAndRes ([], []) [KernelResult]
    kres
      where
        goodReturn :: KernelResult -> Bool
    goodReturn (Returns ResultManifest
    ResultMaySimplify Certs
    _ SubExp
    _) = Bool
    True
        goodReturn KernelResult
    _ = Bool
    False
        addCertAndRes :: ([VName], [SubExp]) -> KernelResult -> ([VName], [SubExp])
    addCertAndRes ([VName]
    cs, [SubExp]
    rs) (Returns ResultManifest
    ResultMaySimplify Certs
    c SubExp
    r_se) =
          ([VName]
    cs [VName] -> [VName] -> [VName]
    forall a. [a] -> [a] -> [a]
    ++ Certs -> [VName]
    unCerts Certs
    c, [SubExp]
    rs [SubExp] -> [SubExp] -> [SubExp]
    forall a. [a] -> [a] -> [a]
    ++ [SubExp
    r_se])
        addCertAndRes ([VName], [SubExp])
    _ KernelResult
    _ =
          [Char] -> ([VName], [SubExp])
    forall a. HasCallStack => [Char] -> a
    error [Char]
    "Impossible case reached in GenRedOpt.hs, function allGoodReturns!"
    allGoodReturns [KernelResult]
    _ = Maybe ([VName], [SubExp])
    forall a. Maybe a
    Nothing

    --------------------------
    --- Cost Model Helpers ---
    --------------------------

    costRedundantExecution ::
      VarianceTable ->
      VName ->
      [SubExp] ->
      Stms GPU ->
      Cost
    costRedundantExecution :: VarianceTable -> VName -> [SubExp] -> Stms GPU -> Cost
    costRedundantExecution VarianceTable
    variance VName
    pat_acc_nm [SubExp]
    r_ses Stms GPU
    kstms =
      let acc_deps :: Names
    acc_deps = Names -> VName -> VarianceTable -> Names
    forall k a. Ord k => a -> k -> Map k a -> a
    M.findWithDefault Names
    forall a. Monoid a => a
    mempty VName
    pat_acc_nm VarianceTable
    variance
          vartab_cut_acc :: VarianceTable
    vartab_cut_acc = Names -> VarianceTable -> Stms GPU -> VarianceTable
    varianceInStmsWithout (VName -> Names
    oneName VName
    pat_acc_nm) VarianceTable
    forall a. Monoid a => a
    mempty Stms GPU
    kstms
          res_deps :: Names
    res_deps = [Names] -> Names
    forall a. Monoid a => [a] -> a
    mconcat ([Names] -> Names) -> [Names] -> Names
    forall a b. (a -> b) -> a -> b
    $ (VName -> Names) -> [VName] -> [Names]
    forall a b. (a -> b) -> [a] -> [b]
    map (VarianceTable -> VName -> Names
    forall {k} {a}. (Ord k, Monoid a) => Map k a -> k -> a
    findDeps VarianceTable
    vartab_cut_acc) ([VName] -> [Names]) -> [VName] -> [Names]
    forall a b. (a -> b) -> a -> b
    $ (SubExp -> Maybe VName) -> [SubExp] -> [VName]
    forall a b. (a -> Maybe b) -> [a] -> [b]
    mapMaybe SubExp -> Maybe VName
    se2nm [SubExp]
    r_ses
          common_deps :: Names
    common_deps = Names -> Names -> Names
    namesIntersection Names
    res_deps Names
    acc_deps
       in (Cost -> Stm GPU -> Cost) -> Cost -> Stms GPU -> Cost
    forall b a. (b -> a -> b) -> b -> Seq a -> b
    forall (t :: * -> *) b a.
    Foldable t =>
    (b -> a -> b) -> b -> t a -> b
    foldl (Names -> Cost -> Stm GPU -> Cost
    addCostOfStmt Names
    common_deps) (Int -> Cost
    Small Int
    0) Stms GPU
    kstms
      where
        se2nm :: SubExp -> Maybe VName
    se2nm (Var VName
    nm) = VName -> Maybe VName
    forall a. a -> Maybe a
    Just VName
    nm
        se2nm SubExp
    _ = Maybe VName
    forall a. Maybe a
    Nothing
        findDeps :: Map k a -> k -> a
    findDeps Map k a
    vartab k
    nm = a -> k -> Map k a -> a
    forall k a. Ord k => a -> k -> Map k a -> a
    M.findWithDefault a
    forall a. Monoid a => a
    mempty k
    nm Map k a
    vartab
        addCostOfStmt :: Names -> Cost -> Stm GPU -> Cost
    addCostOfStmt Names
    common_deps Cost
    cur_cost Stm GPU
    stm =
          let pat_nms :: [VName]
    pat_nms = Pat Type -> [VName]
    forall dec. Pat dec -> [VName]
    patNames (Pat Type -> [VName]) -> Pat Type -> [VName]
    forall a b. (a -> b) -> a -> b
    $ Stm GPU -> Pat (LetDec GPU)
    forall rep. Stm rep -> Pat (LetDec rep)
    stmPat Stm GPU
    stm
           in if Names -> Names -> Bool
    namesIntersect ([VName] -> Names
    namesFromList [VName]
    pat_nms) Names
    common_deps
                then Cost -> Cost -> Cost
    addCosts Cost
    cur_cost (Cost -> Cost) -> Cost -> Cost
    forall a b. (a -> b) -> a -> b
    $ Stm GPU -> Cost
    costRedundantStmt Stm GPU
    stm
                else Cost
    cur_cost
        varianceInStmsWithout :: Names -> VarianceTable -> Stms GPU -> VarianceTable
        varianceInStmsWithout :: Names -> VarianceTable -> Stms GPU -> VarianceTable
    varianceInStmsWithout Names
    nms = (VarianceTable -> Stm GPU -> VarianceTable)
    -> VarianceTable -> Stms GPU -> VarianceTable
    forall b a. (b -> a -> b) -> b -> Seq a -> b
    forall (t :: * -> *) b a.
    Foldable t =>
    (b -> a -> b) -> b -> t a -> b
    L.foldl' (Names -> VarianceTable -> Stm GPU -> VarianceTable
    forall {rep}.
    (FreeDec (ExpDec rep), FreeDec (BodyDec rep), FreeIn (OpC rep rep),
     FreeIn (FParamInfo rep), FreeIn (LParamInfo rep),
     FreeIn (LetDec rep), FreeIn (RetType rep),
     FreeIn (BranchType rep)) =>
    Names -> VarianceTable -> Stm rep -> VarianceTable
    varianceInStmWithout Names
    nms)
        varianceInStmWithout :: Names -> VarianceTable -> Stm rep -> VarianceTable
    varianceInStmWithout Names
    cuts VarianceTable
    vartab Stm rep
    stm =
          let pat_nms :: [VName]
    pat_nms = Pat (LetDec rep) -> [VName]
    forall dec. Pat dec -> [VName]
    patNames (Pat (LetDec rep) -> [VName]) -> Pat (LetDec rep) -> [VName]
    forall a b. (a -> b) -> a -> b
    $ Stm rep -> Pat (LetDec rep)
    forall rep. Stm rep -> Pat (LetDec rep)
    stmPat Stm rep
    stm
           in if Names -> Names -> Bool
    namesIntersect ([VName] -> Names
    namesFromList [VName]
    pat_nms) Names
    cuts
                then VarianceTable
    vartab
                else (VarianceTable -> VName -> VarianceTable)
    -> VarianceTable -> [VName] -> VarianceTable
    forall b a. (b -> a -> b) -> b -> [a] -> b
    forall (t :: * -> *) b a.
    Foldable t =>
    (b -> a -> b) -> b -> t a -> b
    L.foldl' VarianceTable -> VName -> VarianceTable
    add VarianceTable
    vartab [VName]
    pat_nms
          where
            add :: VarianceTable -> VName -> VarianceTable
    add VarianceTable
    variance' VName
    v = VName -> Names -> VarianceTable -> VarianceTable
    forall k a. Ord k => k -> a -> Map k a -> Map k a
    M.insert VName
    v Names
    binding_variance VarianceTable
    variance'
            look :: VarianceTable -> VName -> Names
    look VarianceTable
    variance' VName
    v = VName -> Names
    oneName VName
    v Names -> Names -> Names
    forall a. Semigroup a => a -> a -> a
    <> Names -> VName -> VarianceTable -> Names
    forall k a. Ord k => a -> k -> Map k a -> a
    M.findWithDefault Names
    forall a. Monoid a => a
    mempty VName
    v VarianceTable
    variance'
            binding_variance :: Names
    binding_variance = [Names] -> Names
    forall a. Monoid a => [a] -> a
    mconcat ([Names] -> Names) -> [Names] -> Names
    forall a b. (a -> b) -> a -> b
    $ (VName -> Names) -> [VName] -> [Names]
    forall a b. (a -> b) -> [a] -> [b]
    map (VarianceTable -> VName -> Names
    look VarianceTable
    vartab) ([VName] -> [Names]) -> [VName] -> [Names]
    forall a b. (a -> b) -> a -> b
    $ Names -> [VName]
    namesToList (Stm rep -> Names
    forall a. FreeIn a => a -> Names
    freeIn Stm rep
    stm)

    data Cost = Small Int | Big | Break
      deriving (Cost -> Cost -> Bool
    (Cost -> Cost -> Bool) -> (Cost -> Cost -> Bool) -> Eq Cost
    forall a. (a -> a -> Bool) -> (a -> a -> Bool) -> Eq a
    $c== :: Cost -> Cost -> Bool
    == :: Cost -> Cost -> Bool
    $c/= :: Cost -> Cost -> Bool
    /= :: Cost -> Cost -> Bool
    Eq)

    addCosts :: Cost -> Cost -> Cost
    addCosts :: Cost -> Cost -> Cost
    addCosts Cost
    Break Cost
    _ = Cost
    Break
    addCosts Cost
    _ Cost
    Break = Cost
    Break
    addCosts Cost
    Big Cost
    _ = Cost
    Big
    addCosts Cost
    _ Cost
    Big = Cost
    Big
    addCosts (Small Int
    c1) (Small Int
    c2) = Int -> Cost
    Small (Int
    c1 Int -> Int -> Int
    forall a. Num a => a -> a -> a
    + Int
    c2)

    maxCost :: Cost -> Cost -> Cost
    maxCost :: Cost -> Cost -> Cost
    maxCost (Small Int
    c1) (Small Int
    c2) = Int -> Cost
    Small (Int -> Int -> Int
    forall a. Ord a => a -> a -> a
    max Int
    c1 Int
    c2)
    maxCost Cost
    c1 Cost
    c2 = Cost -> Cost -> Cost
    addCosts Cost
    c1 Cost
    c2

    costBody :: Body GPU -> Cost
    costBody :: Body GPU -> Cost
    costBody Body GPU
    bdy =
      (Cost -> Cost -> Cost) -> Cost -> [Cost] -> Cost
    forall b a. (b -> a -> b) -> b -> [a] -> b
    forall (t :: * -> *) b a.
    Foldable t =>
    (b -> a -> b) -> b -> t a -> b
    foldl Cost -> Cost -> Cost
    addCosts (Int -> Cost
    Small Int
    0) ([Cost] -> Cost) -> [Cost] -> Cost
    forall a b. (a -> b) -> a -> b
    $
        (Stm GPU -> Cost) -> [Stm GPU] -> [Cost]
    forall a b. (a -> b) -> [a] -> [b]
    map Stm GPU -> Cost
    costRedundantStmt ([Stm GPU] -> [Cost]) -> [Stm GPU] -> [Cost]
    forall a b. (a -> b) -> a -> b
    $
          Stms GPU -> [Stm GPU]
    forall rep. Stms rep -> [Stm rep]
    stmsToList (Stms GPU -> [Stm GPU]) -> Stms GPU -> [Stm GPU]
    forall a b. (a -> b) -> a -> b
    $
            Body GPU -> Stms GPU
    forall rep. Body rep -> Stms rep
    bodyStms Body GPU
    bdy

    costRedundantStmt :: Stm GPU -> Cost
    costRedundantStmt :: Stm GPU -> Cost
    costRedundantStmt (Let Pat (LetDec GPU)
    _ StmAux (ExpDec GPU)
    _ (Op OpC GPU GPU
    _)) = Cost
    Big
    costRedundantStmt (Let Pat (LetDec GPU)
    _ StmAux (ExpDec GPU)
    _ Loop {}) = Cost
    Big
    costRedundantStmt (Let Pat (LetDec GPU)
    _ StmAux (ExpDec GPU)
    _ Apply {}) = Cost
    Big
    costRedundantStmt (Let Pat (LetDec GPU)
    _ StmAux (ExpDec GPU)
    _ WithAcc {}) = Cost
    Big
    costRedundantStmt (Let Pat (LetDec GPU)
    _ StmAux (ExpDec GPU)
    _ (Match [SubExp]
    _ [Case (Body GPU)]
    cases Body GPU
    defbody MatchDec (BranchType GPU)
    _)) =
      (Cost -> Cost -> Cost) -> Cost -> [Cost] -> Cost
    forall b a. (b -> a -> b) -> b -> [a] -> b
    forall (t :: * -> *) b a.
    Foldable t =>
    (b -> a -> b) -> b -> t a -> b
    L.foldl' Cost -> Cost -> Cost
    maxCost (Body GPU -> Cost
    costBody Body GPU
    defbody) ([Cost] -> Cost) -> [Cost] -> Cost
    forall a b. (a -> b) -> a -> b
    $ (Case (Body GPU) -> Cost) -> [Case (Body GPU)] -> [Cost]
    forall a b. (a -> b) -> [a] -> [b]
    map (Body GPU -> Cost
    costBody (Body GPU -> Cost)
    -> (Case (Body GPU) -> Body GPU) -> Case (Body GPU) -> Cost
    forall b c a. (b -> c) -> (a -> b) -> a -> c
    . Case (Body GPU) -> Body GPU
    forall body. Case body -> body
    caseBody) [Case (Body GPU)]
    cases
    costRedundantStmt (Let Pat (LetDec GPU)
    _ StmAux (ExpDec GPU)
    _ (BasicOp (ArrayLit [SubExp]
    _ Array {}))) = Cost
    Big
    costRedundantStmt (Let Pat (LetDec GPU)
    _ StmAux (ExpDec GPU)
    _ (BasicOp (ArrayLit [SubExp]
    _ Type
    _))) = Int -> Cost
    Small Int
    1
    costRedundantStmt (Let Pat (LetDec GPU)
    _ StmAux (ExpDec GPU)
    _ (BasicOp (Index VName
    _ Slice SubExp
    slc))) =
      if (DimIndex SubExp -> Bool) -> [DimIndex SubExp] -> Bool
    forall (t :: * -> *) a. Foldable t => (a -> Bool) -> t a -> Bool
    all DimIndex SubExp -> Bool
    forall {d}. DimIndex d -> Bool
    isFixDim (Slice SubExp -> [DimIndex SubExp]
    forall d. Slice d -> [DimIndex d]
    unSlice Slice SubExp
    slc) then Int -> Cost
    Small Int
    1 else Int -> Cost
    Small Int
    0
      where
        isFixDim :: DimIndex d -> Bool
    isFixDim DimFix {} = Bool
    True
        isFixDim DimIndex d
    _ = Bool
    False
    costRedundantStmt (Let Pat (LetDec GPU)
    _ StmAux (ExpDec GPU)
    _ (BasicOp FlatIndex {})) = Int -> Cost
    Small Int
    0
    costRedundantStmt (Let Pat (LetDec GPU)
    _ StmAux (ExpDec GPU)
    _ (BasicOp Update {})) = Cost
    Break
    costRedundantStmt (Let Pat (LetDec GPU)
    _ StmAux (ExpDec GPU)
    _ (BasicOp FlatUpdate {})) = Cost
    Break
    costRedundantStmt (Let Pat (LetDec GPU)
    _ StmAux (ExpDec GPU)
    _ (BasicOp Concat {})) = Cost
    Big
    costRedundantStmt (Let Pat (LetDec GPU)
    _ StmAux (ExpDec GPU)
    _ (BasicOp Manifest {})) = Cost
    Big
    costRedundantStmt (Let Pat (LetDec GPU)
    _ StmAux (ExpDec GPU)
    _ (BasicOp Replicate {})) = Cost
    Big
    costRedundantStmt (Let Pat (LetDec GPU)
    _ StmAux (ExpDec GPU)
    _ (BasicOp UpdateAcc {})) = Cost
    Break
    costRedundantStmt (Let Pat (LetDec GPU)
    _ StmAux (ExpDec GPU)
    _ (BasicOp BasicOp
    _)) = Int -> Cost
    Small Int
    0