    {-# LANGUAGE LambdaCase #-}
    {-# LANGUAGE TypeFamilies #-}

    -- | Perform a restricted form of loop tiling within SegMaps.  We only
    -- tile primitive types, to avoid excessive shared memory use.
    module Futhark.Optimise.TileLoops (tileLoops) where

    import Control.Monad
    import Control.Monad.Reader
    import Control.Monad.State
    import Data.Map.Strict qualified as M
    import Data.Maybe (mapMaybe)
    import Data.Sequence qualified as Seq
    import Futhark.Analysis.Alias qualified as Alias
    import Futhark.IR.GPU
    import Futhark.IR.Prop.Aliases (consumedInStm)
    import Futhark.MonadFreshNames
    import Futhark.Optimise.BlkRegTiling
    import Futhark.Optimise.TileLoops.Shared
    import Futhark.Pass
    import Futhark.Tools
    import Futhark.Transform.Rename
    import Prelude hiding (quot)

    -- | The pass definition.
    tileLoops :: Pass GPU GPU
    tileLoops :: Pass GPU GPU
    tileLoops =
      [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]
    "tile loops" [Char]
    "Tile stream loops inside kernels" ((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 (Map VName (Lambda GPU, [SubExp])
    forall k a. Map k a
    M.empty, Scope GPU -> IxFnEnv
    initialIxFnEnv Scope GPU
    scope) Stms GPU
    stms) Scope GPU
    scope

    optimiseBody :: Env -> Body GPU -> TileM (Body GPU)
    optimiseBody :: Env -> Body GPU -> TileM (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
    -> TileM (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 -> TileM (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 -> TileM (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 -> TileM (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)
    pat StmAux (ExpDec GPU)
    aux (Op (SegOp (SegMap lvl :: SegLevel
    lvl@SegThread {} SegSpace
    space [Type]
    ts KernelBody GPU
    kbody)))) = do
      res3dtiling <- Scope GPU
    -> ReaderT
         (Scope GPU) (State VNameSource) (Maybe (Stms GPU, Stm GPU))
    -> ReaderT
         (Scope GPU) (State VNameSource) (Maybe (Stms GPU, Stm 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 (SegSpace -> Scope GPU
    forall rep. SegSpace -> Scope rep
    scopeOfSegSpace SegSpace
    space) (ReaderT
       (Scope GPU) (State VNameSource) (Maybe (Stms GPU, Stm GPU))
     -> ReaderT
          (Scope GPU) (State VNameSource) (Maybe (Stms GPU, Stm GPU)))
    -> ReaderT
         (Scope GPU) (State VNameSource) (Maybe (Stms GPU, Stm GPU))
    -> ReaderT
         (Scope GPU) (State VNameSource) (Maybe (Stms GPU, Stm GPU))
    forall a b. (a -> b) -> a -> b
    $ Stm GPU
    -> ReaderT
         (Scope GPU) (State VNameSource) (Maybe (Stms GPU, Stm GPU))
    doRegTiling3D Stm GPU
    stm
      stms' <-
        case res3dtiling of
          Just (Stms GPU
    extra_stms, Stm GPU
    stmt') -> 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
    extra_stms 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
    stmt')
          Maybe (Stms GPU, Stm GPU)
    Nothing -> do
            blkRegTiling_res <- Env
    -> Stm GPU
    -> ReaderT
         (Scope GPU) (State VNameSource) (Maybe (Stms GPU, Stm GPU))
    mmBlkRegTiling Env
    env Stm GPU
    stm
            case blkRegTiling_res of
              Just (Stms GPU
    extra_stms, Stm GPU
    stmt') -> 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
    extra_stms 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
    stmt')
              Maybe (Stms GPU, Stm GPU)
    Nothing -> 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 (SegSpace -> Scope GPU
    forall rep. SegSpace -> Scope rep
    scopeOfSegSpace SegSpace
    space) (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
                (host_stms, (lvl', space', kbody')) <- Names
    -> AliasTable
    -> SegLevel
    -> SegSpace
    -> [Type]
    -> KernelBody GPU
    -> TileM (Stms GPU, (SegLevel, SegSpace, KernelBody GPU))
    tileInKernelBody Names
    forall a. Monoid a => a
    mempty AliasTable
    initial_variance SegLevel
    lvl SegSpace
    space [Type]
    ts KernelBody GPU
    kbody
                pure $ host_stms <> oneStm (Let pat aux $ Op $ SegOp $ SegMap lvl' space' ts kbody')
      pure (env, stms')
      where
        initial_variance :: AliasTable
    initial_variance = (NameInfo (ZonkAny 0) -> Names)
    -> Map VName (NameInfo (ZonkAny 0)) -> AliasTable
    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)) -> AliasTable)
    -> Map VName (NameInfo (ZonkAny 0)) -> AliasTable
    forall a b. (a -> b) -> a -> b
    $ SegSpace -> Map VName (NameInfo (ZonkAny 0))
    forall rep. SegSpace -> Scope rep
    scopeOfSegSpace SegSpace
    space
    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 -> TileM (Body GPU) -> TileM (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 (TileM (Body GPU) -> TileM (Body GPU))
    -> (Body GPU -> TileM (Body GPU)) -> Body GPU -> TileM (Body GPU)
    forall b c a. (b -> c) -> (a -> b) -> a -> c
    . Env -> Body GPU -> TileM (Body GPU)
    optimiseBody Env
    env'}

    tileInKernelBody ::
      Names ->
      VarianceTable ->
      SegLevel ->
      SegSpace ->
      [Type] ->
      KernelBody GPU ->
      TileM (Stms GPU, (SegLevel, SegSpace, KernelBody GPU))
    tileInKernelBody :: Names
    -> AliasTable
    -> SegLevel
    -> SegSpace
    -> [Type]
    -> KernelBody GPU
    -> TileM (Stms GPU, (SegLevel, SegSpace, KernelBody GPU))
    tileInKernelBody Names
    branch_variant AliasTable
    initial_variance SegLevel
    lvl SegSpace
    initial_kspace [Type]
    ts KernelBody GPU
    kbody
      | Just Result
    kbody_res <- (KernelResult -> Maybe SubExpRes) -> [KernelResult] -> Maybe Result
    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 KernelResult -> Maybe SubExpRes
    isSimpleResult ([KernelResult] -> Maybe Result) -> [KernelResult] -> Maybe Result
    forall a b. (a -> b) -> a -> b
    $ KernelBody GPU -> [KernelResult]
    forall rep. KernelBody rep -> [KernelResult]
    kernelBodyResult KernelBody GPU
    kbody = do
          maybe_tiled <-
            Names
    -> AliasTable
    -> SegLevel
    -> SegSpace
    -> [Type]
    -> Body GPU
    -> TileM (Maybe (Stms GPU, Tiling, TiledBody))
    tileInBody Names
    branch_variant AliasTable
    initial_variance SegLevel
    lvl SegSpace
    initial_kspace [Type]
    ts (Body GPU -> TileM (Maybe (Stms GPU, Tiling, TiledBody)))
    -> Body GPU -> TileM (Maybe (Stms GPU, Tiling, TiledBody))
    forall a b. (a -> b) -> a -> b
    $
              BodyDec GPU -> Stms GPU -> Result -> Body GPU
    forall rep. BodyDec rep -> Stms rep -> Result -> Body rep
    Body () (KernelBody GPU -> Stms GPU
    forall rep. KernelBody rep -> Stms rep
    kernelBodyStms KernelBody GPU
    kbody) Result
    kbody_res
          case maybe_tiled of
            Just (Stms GPU
    host_stms, Tiling
    tiling, TiledBody
    tiledBody) -> do
              (res', stms') <-
                Builder GPU [KernelResult]
    -> ReaderT
         (Scope GPU) (State VNameSource) ([KernelResult], Stms GPU)
    forall (m :: * -> *) somerep rep a.
    (MonadFreshNames m, HasScope somerep m, SameScope somerep rep) =>
    Builder rep a -> m (a, Stms rep)
    runBuilder (Builder GPU [KernelResult]
     -> ReaderT
          (Scope GPU) (State VNameSource) ([KernelResult], Stms GPU))
    -> Builder GPU [KernelResult]
    -> ReaderT
         (Scope GPU) (State VNameSource) ([KernelResult], Stms GPU)
    forall a b. (a -> b) -> a -> b
    $ (VName -> BuilderT GPU (State VNameSource) KernelResult)
    -> [VName] -> Builder GPU [KernelResult]
    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 (Tiling -> VName -> BuilderT GPU (State VNameSource) KernelResult
    tilingTileReturns Tiling
    tiling) ([VName] -> Builder GPU [KernelResult])
    -> BuilderT GPU (State VNameSource) [VName]
    -> Builder GPU [KernelResult]
    forall (m :: * -> *) a b. Monad m => (a -> m b) -> m a -> m b
    =<< TiledBody
    tiledBody Names
    forall a. Monoid a => a
    mempty PrivStms
    forall a. Monoid a => a
    mempty
              pure
                ( host_stms,
                  ( tilingLevel tiling,
                    tilingSpace tiling,
                    KernelBody () stms' res'
                  )
                )
            Maybe (Stms GPU, Tiling, TiledBody)
    Nothing ->
              (Stms GPU, (SegLevel, SegSpace, KernelBody GPU))
    -> TileM (Stms GPU, (SegLevel, SegSpace, KernelBody GPU))
    forall a. a -> ReaderT (Scope GPU) (State VNameSource) a
    forall (f :: * -> *) a. Applicative f => a -> f a
    pure (Stms GPU
    forall a. Monoid a => a
    mempty, (SegLevel
    lvl, SegSpace
    initial_kspace, KernelBody GPU
    kbody))
      | Bool
    otherwise =
          (Stms GPU, (SegLevel, SegSpace, KernelBody GPU))
    -> TileM (Stms GPU, (SegLevel, SegSpace, KernelBody GPU))
    forall a. a -> ReaderT (Scope GPU) (State VNameSource) a
    forall (f :: * -> *) a. Applicative f => a -> f a
    pure (Stms GPU
    forall a. Monoid a => a
    mempty, (SegLevel
    lvl, SegSpace
    initial_kspace, KernelBody GPU
    kbody))
      where
        isSimpleResult :: KernelResult -> Maybe SubExpRes
    isSimpleResult (Returns ResultManifest
    _ Certs
    cs SubExp
    se) = SubExpRes -> Maybe SubExpRes
    forall a. a -> Maybe a
    Just (SubExpRes -> Maybe SubExpRes) -> SubExpRes -> Maybe SubExpRes
    forall a b. (a -> b) -> a -> b
    $ Certs -> SubExp -> SubExpRes
    SubExpRes Certs
    cs SubExp
    se
        isSimpleResult KernelResult
    _ = Maybe SubExpRes
    forall a. Maybe a
    Nothing

    tileInBody ::
      Names ->
      VarianceTable ->
      SegLevel ->
      SegSpace ->
      [Type] ->
      Body GPU ->
      TileM (Maybe (Stms GPU, Tiling, TiledBody))
    tileInBody :: Names
    -> AliasTable
    -> SegLevel
    -> SegSpace
    -> [Type]
    -> Body GPU
    -> TileM (Maybe (Stms GPU, Tiling, TiledBody))
    tileInBody Names
    branch_variant AliasTable
    initial_variance SegLevel
    initial_lvl SegSpace
    initial_space [Type]
    res_ts (Body () Stms GPU
    initial_kstms Result
    stms_res) =
      Stms GPU
    -> [Stm GPU] -> TileM (Maybe (Stms GPU, Tiling, TiledBody))
    descend Stms GPU
    forall a. Monoid a => a
    mempty ([Stm GPU] -> TileM (Maybe (Stms GPU, Tiling, TiledBody)))
    -> [Stm GPU] -> TileM (Maybe (Stms GPU, Tiling, TiledBody))
    forall a b. (a -> b) -> a -> b
    $ Stms GPU -> [Stm GPU]
    forall rep. Stms rep -> [Stm rep]
    stmsToList Stms GPU
    initial_kstms
      where
        variance :: AliasTable
    variance = AliasTable -> Stms GPU -> AliasTable
    varianceInStms AliasTable
    initial_variance Stms GPU
    initial_kstms

        descend :: Stms GPU
    -> [Stm GPU] -> TileM (Maybe (Stms GPU, Tiling, TiledBody))
    descend Stms GPU
    _ [] =
          Maybe (Stms GPU, Tiling, TiledBody)
    -> TileM (Maybe (Stms GPU, Tiling, TiledBody))
    forall a. a -> ReaderT (Scope GPU) (State VNameSource) a
    forall (f :: * -> *) a. Applicative f => a -> f a
    pure Maybe (Stms GPU, Tiling, TiledBody)
    forall a. Maybe a
    Nothing
        descend Stms GPU
    prestms (Stm GPU
    stm_to_tile : [Stm GPU]
    poststms)
          -- 2D tiling of redomap.
          | ([VName]
    gtids, [SubExp]
    kdims) <- [(VName, SubExp)] -> ([VName], [SubExp])
    forall a b. [(a, b)] -> ([a], [b])
    unzip ([(VName, SubExp)] -> ([VName], [SubExp]))
    -> [(VName, SubExp)] -> ([VName], [SubExp])
    forall a b. (a -> b) -> a -> b
    $ SegSpace -> [(VName, SubExp)]
    unSegSpace SegSpace
    initial_space,
            Just (SubExp
    w, [VName]
    arrs, (Commutativity, Lambda GPU, [SubExp], Lambda GPU)
    form) <- Stm GPU
    -> Maybe
         (SubExp, [VName],
          (Commutativity, Lambda GPU, [SubExp], Lambda GPU))
    tileable Stm GPU
    stm_to_tile,
            Just [InputArray]
    inputs <-
              (VName -> Maybe InputArray) -> [VName] -> Maybe [InputArray]
    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 (Names -> AliasTable -> [VName] -> VName -> Maybe InputArray
    invariantToOneOfTwoInnerDims Names
    branch_variant AliasTable
    variance [VName]
    gtids) [VName]
    arrs,
            Bool -> Bool
    not (Bool -> Bool) -> Bool -> Bool
    forall a b. (a -> b) -> a -> b
    $ [(VName, [Int])] -> Bool
    forall a. [a] -> Bool
    forall (t :: * -> *) a. Foldable t => t a -> Bool
    null ([(VName, [Int])] -> Bool) -> [(VName, [Int])] -> Bool
    forall a b. (a -> b) -> a -> b
    $ [InputArray] -> [(VName, [Int])]
    tiledInputs [InputArray]
    inputs,
            VName
    gtid_y : VName
    gtid_x : [VName]
    top_gtids_rev <- [VName] -> [VName]
    forall a. [a] -> [a]
    reverse [VName]
    gtids,
            SubExp
    kdim_y : SubExp
    kdim_x : [SubExp]
    top_kdims_rev <- [SubExp] -> [SubExp]
    forall a. [a] -> [a]
    reverse [SubExp]
    kdims,
            Just (Stms GPU
    prestms', Stms GPU
    poststms') <-
              AliasTable
    -> Stms GPU -> Stm GPU -> Stms GPU -> Maybe (Stms GPU, Stms GPU)
    preludeToPostlude AliasTable
    variance Stms GPU
    prestms Stm GPU
    stm_to_tile ([Stm GPU] -> Stms GPU
    forall rep. [Stm rep] -> Stms rep
    stmsFromList [Stm GPU]
    poststms),
            Names
    used <- Stm GPU -> Names
    forall a. FreeIn a => a -> Names
    freeIn Stm GPU
    stm_to_tile Names -> Names -> Names
    forall a. Semigroup a => a -> a -> a
    <> Stms GPU -> Names
    forall a. FreeIn a => a -> Names
    freeIn Stms GPU
    poststms' Names -> Names -> Names
    forall a. Semigroup a => a -> a -> a
    <> Result -> Names
    forall a. FreeIn a => a -> Names
    freeIn Result
    stms_res =
              (Stms GPU, Tiling, TiledBody)
    -> Maybe (Stms GPU, Tiling, TiledBody)
    forall a. a -> Maybe a
    Just ((Stms GPU, Tiling, TiledBody)
     -> Maybe (Stms GPU, Tiling, TiledBody))
    -> ((Stms GPU, Tiling, TiledBody) -> (Stms GPU, Tiling, TiledBody))
    -> (Stms GPU, Tiling, TiledBody)
    -> Maybe (Stms GPU, Tiling, TiledBody)
    forall b c a. (b -> c) -> (a -> b) -> a -> c
    . SegSpace
    -> AliasTable
    -> Stms GPU
    -> Names
    -> (Stms GPU, Tiling, TiledBody)
    -> (Stms GPU, Tiling, TiledBody)
    injectPrelude SegSpace
    initial_space AliasTable
    variance Stms GPU
    prestms' Names
    used
                ((Stms GPU, Tiling, TiledBody)
     -> Maybe (Stms GPU, Tiling, TiledBody))
    -> ReaderT
         (Scope GPU) (State VNameSource) (Stms GPU, Tiling, TiledBody)
    -> TileM (Maybe (Stms GPU, Tiling, TiledBody))
    forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
    <$> DoTiling (VName, VName) (SubExp, SubExp)
    -> [Type]
    -> Pat Type
    -> (VName, VName)
    -> (SubExp, SubExp)
    -> SubExp
    -> (Commutativity, Lambda GPU, [SubExp], Lambda GPU)
    -> [InputArray]
    -> Stms GPU
    -> Result
    -> ReaderT
         (Scope GPU) (State VNameSource) (Stms GPU, Tiling, TiledBody)
    forall gtids kdims.
    DoTiling gtids kdims
    -> [Type]
    -> Pat Type
    -> gtids
    -> kdims
    -> SubExp
    -> (Commutativity, Lambda GPU, [SubExp], Lambda GPU)
    -> [InputArray]
    -> Stms GPU
    -> Result
    -> ReaderT
         (Scope GPU) (State VNameSource) (Stms GPU, Tiling, TiledBody)
    tileGeneric
                  ([(VName, SubExp)] -> DoTiling (VName, VName) (SubExp, SubExp)
    tiling2d ([(VName, SubExp)] -> DoTiling (VName, VName) (SubExp, SubExp))
    -> [(VName, SubExp)] -> DoTiling (VName, VName) (SubExp, SubExp)
    forall a b. (a -> b) -> a -> b
    $ [(VName, SubExp)] -> [(VName, SubExp)]
    forall a. [a] -> [a]
    reverse ([(VName, SubExp)] -> [(VName, SubExp)])
    -> [(VName, SubExp)] -> [(VName, SubExp)]
    forall a b. (a -> b) -> a -> b
    $ [VName] -> [SubExp] -> [(VName, SubExp)]
    forall a b. [a] -> [b] -> [(a, b)]
    zip [VName]
    top_gtids_rev [SubExp]
    top_kdims_rev)
                  [Type]
    res_ts
                  (Stm GPU -> Pat (LetDec GPU)
    forall rep. Stm rep -> Pat (LetDec rep)
    stmPat Stm GPU
    stm_to_tile)
                  (VName
    gtid_x, VName
    gtid_y)
                  (SubExp
    kdim_x, SubExp
    kdim_y)
                  SubExp
    w
                  (Commutativity, Lambda GPU, [SubExp], Lambda GPU)
    form
                  [InputArray]
    inputs
                  Stms GPU
    poststms'
                  Result
    stms_res
          -- 1D tiling of redomap.
          | (VName
    gtid, SubExp
    kdim) : [(VName, SubExp)]
    top_space_rev <- [(VName, SubExp)] -> [(VName, SubExp)]
    forall a. [a] -> [a]
    reverse ([(VName, SubExp)] -> [(VName, SubExp)])
    -> [(VName, SubExp)] -> [(VName, SubExp)]
    forall a b. (a -> b) -> a -> b
    $ SegSpace -> [(VName, SubExp)]
    unSegSpace SegSpace
    initial_space,
            Just (SubExp
    w, [VName]
    arrs, (Commutativity, Lambda GPU, [SubExp], Lambda GPU)
    form) <- Stm GPU
    -> Maybe
         (SubExp, [VName],
          (Commutativity, Lambda GPU, [SubExp], Lambda GPU))
    tileable Stm GPU
    stm_to_tile,
            [InputArray]
    inputs <- (VName -> InputArray) -> [VName] -> [InputArray]
    forall a b. (a -> b) -> [a] -> [b]
    map (VName -> AliasTable -> VName -> InputArray
    is1DTileable VName
    gtid AliasTable
    variance) [VName]
    arrs,
            Bool -> Bool
    not (Bool -> Bool) -> Bool -> Bool
    forall a b. (a -> b) -> a -> b
    $ [(VName, [Int])] -> Bool
    forall a. [a] -> Bool
    forall (t :: * -> *) a. Foldable t => t a -> Bool
    null ([(VName, [Int])] -> Bool) -> [(VName, [Int])] -> Bool
    forall a b. (a -> b) -> a -> b
    $ [InputArray] -> [(VName, [Int])]
    tiledInputs [InputArray]
    inputs,
            VName
    gtid VName -> Names -> Bool
    `notNameIn` Names
    branch_variant,
            Just (Stms GPU
    prestms', Stms GPU
    poststms') <-
              AliasTable
    -> Stms GPU -> Stm GPU -> Stms GPU -> Maybe (Stms GPU, Stms GPU)
    preludeToPostlude AliasTable
    variance Stms GPU
    prestms Stm GPU
    stm_to_tile ([Stm GPU] -> Stms GPU
    forall rep. [Stm rep] -> Stms rep
    stmsFromList [Stm GPU]
    poststms),
            Names
    used <- Stm GPU -> Names
    forall a. FreeIn a => a -> Names
    freeIn Stm GPU
    stm_to_tile Names -> Names -> Names
    forall a. Semigroup a => a -> a -> a
    <> Stms GPU -> Names
    forall a. FreeIn a => a -> Names
    freeIn Stms GPU
    poststms' Names -> Names -> Names
    forall a. Semigroup a => a -> a -> a
    <> Result -> Names
    forall a. FreeIn a => a -> Names
    freeIn Result
    stms_res =
              (Stms GPU, Tiling, TiledBody)
    -> Maybe (Stms GPU, Tiling, TiledBody)
    forall a. a -> Maybe a
    Just ((Stms GPU, Tiling, TiledBody)
     -> Maybe (Stms GPU, Tiling, TiledBody))
    -> ((Stms GPU, Tiling, TiledBody) -> (Stms GPU, Tiling, TiledBody))
    -> (Stms GPU, Tiling, TiledBody)
    -> Maybe (Stms GPU, Tiling, TiledBody)
    forall b c a. (b -> c) -> (a -> b) -> a -> c
    . SegSpace
    -> AliasTable
    -> Stms GPU
    -> Names
    -> (Stms GPU, Tiling, TiledBody)
    -> (Stms GPU, Tiling, TiledBody)
    injectPrelude SegSpace
    initial_space AliasTable
    variance Stms GPU
    prestms' Names
    used
                ((Stms GPU, Tiling, TiledBody)
     -> Maybe (Stms GPU, Tiling, TiledBody))
    -> ReaderT
         (Scope GPU) (State VNameSource) (Stms GPU, Tiling, TiledBody)
    -> TileM (Maybe (Stms GPU, Tiling, TiledBody))
    forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
    <$> DoTiling VName SubExp
    -> [Type]
    -> Pat Type
    -> VName
    -> SubExp
    -> SubExp
    -> (Commutativity, Lambda GPU, [SubExp], Lambda GPU)
    -> [InputArray]
    -> Stms GPU
    -> Result
    -> ReaderT
         (Scope GPU) (State VNameSource) (Stms GPU, Tiling, TiledBody)
    forall gtids kdims.
    DoTiling gtids kdims
    -> [Type]
    -> Pat Type
    -> gtids
    -> kdims
    -> SubExp
    -> (Commutativity, Lambda GPU, [SubExp], Lambda GPU)
    -> [InputArray]
    -> Stms GPU
    -> Result
    -> ReaderT
         (Scope GPU) (State VNameSource) (Stms GPU, Tiling, TiledBody)
    tileGeneric
                  ([(VName, SubExp)] -> DoTiling VName SubExp
    tiling1d ([(VName, SubExp)] -> DoTiling VName SubExp)
    -> [(VName, SubExp)] -> DoTiling VName SubExp
    forall a b. (a -> b) -> a -> b
    $ [(VName, SubExp)] -> [(VName, SubExp)]
    forall a. [a] -> [a]
    reverse [(VName, SubExp)]
    top_space_rev)
                  [Type]
    res_ts
                  (Stm GPU -> Pat (LetDec GPU)
    forall rep. Stm rep -> Pat (LetDec rep)
    stmPat Stm GPU
    stm_to_tile)
                  VName
    gtid
                  SubExp
    kdim
                  SubExp
    w
                  (Commutativity, Lambda GPU, [SubExp], Lambda GPU)
    form
                  [InputArray]
    inputs
                  Stms GPU
    poststms'
                  Result
    stms_res
          -- Tiling inside for-loop.
          | Loop [(FParam GPU, SubExp)]
    merge (ForLoop VName
    i IntType
    it SubExp
    bound) Body GPU
    loopbody <- Stm GPU -> Exp GPU
    forall rep. Stm rep -> Exp rep
    stmExp Stm GPU
    stm_to_tile,
            Bool -> Bool
    not (Bool -> Bool) -> Bool -> Bool
    forall a b. (a -> b) -> a -> b
    $ ((Param (TypeBase Shape Uniqueness), SubExp) -> Bool)
    -> [(Param (TypeBase Shape Uniqueness), SubExp)] -> Bool
    forall (t :: * -> *) a. Foldable t => (a -> Bool) -> t a -> Bool
    any ((VName -> Names -> Bool
    `nameIn` [(Param (TypeBase Shape Uniqueness), SubExp)] -> Names
    forall a. FreeIn a => a -> Names
    freeIn [(Param (TypeBase Shape Uniqueness), SubExp)]
    [(FParam GPU, SubExp)]
    merge) (VName -> Bool)
    -> ((Param (TypeBase Shape Uniqueness), SubExp) -> VName)
    -> (Param (TypeBase Shape Uniqueness), SubExp)
    -> Bool
    forall b c a. (b -> c) -> (a -> b) -> a -> c
    . Param (TypeBase Shape Uniqueness) -> VName
    forall dec. Param dec -> VName
    paramName (Param (TypeBase Shape Uniqueness) -> VName)
    -> ((Param (TypeBase Shape Uniqueness), SubExp)
        -> Param (TypeBase Shape Uniqueness))
    -> (Param (TypeBase Shape Uniqueness), SubExp)
    -> VName
    forall b c a. (b -> c) -> (a -> b) -> a -> c
    . (Param (TypeBase Shape Uniqueness), SubExp)
    -> Param (TypeBase Shape Uniqueness)
    forall a b. (a, b) -> a
    fst) [(Param (TypeBase Shape Uniqueness), SubExp)]
    [(FParam GPU, SubExp)]
    merge,
            Just (Stms GPU
    prestms', Stms GPU
    poststms') <-
              AliasTable
    -> Stms GPU -> Stm GPU -> Stms GPU -> Maybe (Stms GPU, Stms GPU)
    preludeToPostlude AliasTable
    variance Stms GPU
    prestms Stm GPU
    stm_to_tile ([Stm GPU] -> Stms GPU
    forall rep. [Stm rep] -> Stms rep
    stmsFromList [Stm GPU]
    poststms) = do
              let branch_variant' :: Names
    branch_variant' =
                    Names
    branch_variant
                      Names -> Names -> Names
    forall a. Semigroup a => a -> a -> a
    <> [Names] -> Names
    forall a. Monoid a => [a] -> a
    mconcat
                        ( (VName -> Names) -> [VName] -> [Names]
    forall a b. (a -> b) -> [a] -> [b]
    map
                            ((VName -> AliasTable -> Names) -> AliasTable -> VName -> Names
    forall a b c. (a -> b -> c) -> b -> a -> c
    flip (Names -> VName -> AliasTable -> Names
    forall k a. Ord k => a -> k -> Map k a -> a
    M.findWithDefault Names
    forall a. Monoid a => a
    mempty) AliasTable
    variance)
                            (Names -> [VName]
    namesToList (SubExp -> Names
    forall a. FreeIn a => a -> Names
    freeIn SubExp
    bound))
                        )
                  merge_params :: [Param (TypeBase Shape Uniqueness)]
    merge_params = ((Param (TypeBase Shape Uniqueness), SubExp)
     -> Param (TypeBase Shape Uniqueness))
    -> [(Param (TypeBase Shape Uniqueness), SubExp)]
    -> [Param (TypeBase Shape Uniqueness)]
    forall a b. (a -> b) -> [a] -> [b]
    map (Param (TypeBase Shape Uniqueness), SubExp)
    -> Param (TypeBase Shape Uniqueness)
    forall a b. (a, b) -> a
    fst [(Param (TypeBase Shape Uniqueness), SubExp)]
    [(FParam GPU, SubExp)]
    merge

              maybe_tiled <-
                Scope GPU
    -> TileM (Maybe (Stms GPU, Tiling, TiledBody))
    -> TileM (Maybe (Stms GPU, Tiling, TiledBody))
    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 (VName -> NameInfo GPU -> Scope GPU -> Scope GPU
    forall k a. Ord k => k -> a -> Map k a -> Map k a
    M.insert VName
    i (IntType -> NameInfo GPU
    forall rep. IntType -> NameInfo rep
    IndexName IntType
    it) (Scope GPU -> Scope GPU) -> Scope GPU -> Scope GPU
    forall a b. (a -> b) -> a -> b
    $ [Param (TypeBase Shape Uniqueness)] -> Scope GPU
    forall rep dec. (FParamInfo rep ~ dec) => [Param dec] -> Scope rep
    scopeOfFParams [Param (TypeBase Shape Uniqueness)]
    merge_params)
                  (TileM (Maybe (Stms GPU, Tiling, TiledBody))
     -> TileM (Maybe (Stms GPU, Tiling, TiledBody)))
    -> TileM (Maybe (Stms GPU, Tiling, TiledBody))
    -> TileM (Maybe (Stms GPU, Tiling, TiledBody))
    forall a b. (a -> b) -> a -> b
    $ Names
    -> AliasTable
    -> SegLevel
    -> SegSpace
    -> [Type]
    -> Body GPU
    -> TileM (Maybe (Stms GPU, Tiling, TiledBody))
    tileInBody
                    Names
    branch_variant'
                    AliasTable
    variance
                    SegLevel
    initial_lvl
                    SegSpace
    initial_space
                    ((Param (TypeBase Shape Uniqueness) -> Type)
    -> [Param (TypeBase Shape Uniqueness)] -> [Type]
    forall a b. (a -> b) -> [a] -> [b]
    map Param (TypeBase Shape Uniqueness) -> Type
    forall dec. Typed dec => Param dec -> Type
    paramType [Param (TypeBase Shape Uniqueness)]
    merge_params)
                  (Body GPU -> TileM (Maybe (Stms GPU, Tiling, TiledBody)))
    -> Body GPU -> TileM (Maybe (Stms GPU, Tiling, TiledBody))
    forall a b. (a -> b) -> a -> b
    $ Stms GPU -> Result -> Body GPU
    forall rep. Buildable rep => Stms rep -> Result -> Body rep
    mkBody (Body GPU -> Stms GPU
    forall rep. Body rep -> Stms rep
    bodyStms Body GPU
    loopbody) (Body GPU -> Result
    forall rep. Body rep -> Result
    bodyResult Body GPU
    loopbody)

              case maybe_tiled of
                Maybe (Stms GPU, Tiling, TiledBody)
    Nothing -> TileM (Maybe (Stms GPU, Tiling, TiledBody))
    next
                Just (Stms GPU, Tiling, TiledBody)
    tiled ->
                  (Stms GPU, Tiling, TiledBody)
    -> Maybe (Stms GPU, Tiling, TiledBody)
    forall a. a -> Maybe a
    Just
                    ((Stms GPU, Tiling, TiledBody)
     -> Maybe (Stms GPU, Tiling, TiledBody))
    -> ReaderT
         (Scope GPU) (State VNameSource) (Stms GPU, Tiling, TiledBody)
    -> TileM (Maybe (Stms GPU, Tiling, TiledBody))
    forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
    <$> SegSpace
    -> AliasTable
    -> Stms GPU
    -> Names
    -> (Stms GPU, Tiling, TiledBody)
    -> [Type]
    -> Pat Type
    -> StmAux (ExpDec GPU)
    -> [(FParam GPU, SubExp)]
    -> VName
    -> IntType
    -> SubExp
    -> Stms GPU
    -> Result
    -> ReaderT
         (Scope GPU) (State VNameSource) (Stms GPU, Tiling, TiledBody)
    tileLoop
                      SegSpace
    initial_space
                      AliasTable
    variance
                      Stms GPU
    prestms'
                      (Body GPU -> Names
    forall a. FreeIn a => a -> Names
    freeIn Body GPU
    loopbody Names -> Names -> Names
    forall a. Semigroup a => a -> a -> a
    <> [(Param (TypeBase Shape Uniqueness), SubExp)] -> Names
    forall a. FreeIn a => a -> Names
    freeIn [(Param (TypeBase Shape Uniqueness), SubExp)]
    [(FParam GPU, SubExp)]
    merge)
                      (Stms GPU, Tiling, TiledBody)
    tiled
                      [Type]
    res_ts
                      (Stm GPU -> Pat (LetDec GPU)
    forall rep. Stm rep -> Pat (LetDec rep)
    stmPat Stm GPU
    stm_to_tile)
                      (Stm GPU -> StmAux (ExpDec GPU)
    forall rep. Stm rep -> StmAux (ExpDec rep)
    stmAux Stm GPU
    stm_to_tile)
                      [(FParam GPU, SubExp)]
    merge
                      VName
    i
                      IntType
    it
                      SubExp
    bound
                      Stms GPU
    poststms'
                      Result
    stms_res
          | Bool
    otherwise = TileM (Maybe (Stms GPU, Tiling, TiledBody))
    next
          where
            next :: TileM (Maybe (Stms GPU, Tiling, TiledBody))
    next =
              Scope GPU
    -> TileM (Maybe (Stms GPU, Tiling, TiledBody))
    -> TileM (Maybe (Stms GPU, Tiling, TiledBody))
    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 (Stm GPU -> Scope GPU
    forall rep a. Scoped rep a => a -> Scope rep
    scopeOf Stm GPU
    stm_to_tile) (TileM (Maybe (Stms GPU, Tiling, TiledBody))
     -> TileM (Maybe (Stms GPU, Tiling, TiledBody)))
    -> TileM (Maybe (Stms GPU, Tiling, TiledBody))
    -> TileM (Maybe (Stms GPU, Tiling, TiledBody))
    forall a b. (a -> b) -> a -> b
    $
                Stms GPU
    -> [Stm GPU] -> TileM (Maybe (Stms GPU, Tiling, TiledBody))
    descend (Stms GPU
    prestms 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
    stm_to_tile) [Stm GPU]
    poststms

    -- | Move statements from prelude to postlude if they are not used in
    -- the tiled statement anyway.  Also, fail if the provided Stm uses
    -- anything from the resulting prelude whose size is not free in the
    -- prelude.
    preludeToPostlude ::
      VarianceTable ->
      Stms GPU ->
      Stm GPU ->
      Stms GPU ->
      Maybe (Stms GPU, Stms GPU)
    preludeToPostlude :: AliasTable
    -> Stms GPU -> Stm GPU -> Stms GPU -> Maybe (Stms GPU, Stms GPU)
    preludeToPostlude AliasTable
    variance Stms GPU
    prelude Stm GPU
    stm_to_tile Stms GPU
    postlude = do
      let prelude_sizes :: Names
    prelude_sizes =
            [Type] -> Names
    forall a. FreeIn a => a -> Names
    freeIn ([Type] -> Names) -> [Type] -> Names
    forall a b. (a -> b) -> a -> b
    $ (Stm GPU -> [Type]) -> Stms GPU -> [Type]
    forall m a. Monoid m => (a -> m) -> Seq a -> m
    forall (t :: * -> *) m a.
    (Foldable t, Monoid m) =>
    (a -> m) -> t a -> m
    foldMap (Pat Type -> [Type]
    forall dec. Typed dec => Pat dec -> [Type]
    patTypes (Pat Type -> [Type]) -> (Stm GPU -> Pat Type) -> Stm GPU -> [Type]
    forall b c a. (b -> c) -> (a -> b) -> a -> c
    . Stm GPU -> Pat Type
    Stm GPU -> Pat (LetDec GPU)
    forall rep. Stm rep -> Pat (LetDec rep)
    stmPat) Stms GPU
    prelude_used
          prelude_bound :: Names
    prelude_bound =
            [VName] -> Names
    namesFromList ([VName] -> Names) -> [VName] -> Names
    forall a b. (a -> b) -> a -> b
    $ (Stm GPU -> [VName]) -> Stms GPU -> [VName]
    forall m a. Monoid m => (a -> m) -> Seq a -> m
    forall (t :: * -> *) m a.
    (Foldable t, Monoid m) =>
    (a -> m) -> t a -> m
    foldMap (Pat Type -> [VName]
    forall dec. Pat dec -> [VName]
    patNames (Pat Type -> [VName])
    -> (Stm GPU -> Pat Type) -> Stm GPU -> [VName]
    forall b c a. (b -> c) -> (a -> b) -> a -> c
    . Stm GPU -> Pat Type
    Stm GPU -> Pat (LetDec GPU)
    forall rep. Stm rep -> Pat (LetDec rep)
    stmPat) Stms GPU
    prelude_used
      Bool -> Maybe ()
    forall (f :: * -> *). Alternative f => Bool -> f ()
    guard (Bool -> Maybe ()) -> Bool -> Maybe ()
    forall a b. (a -> b) -> a -> b
    $ Bool -> Bool
    not (Bool -> Bool) -> Bool -> Bool
    forall a b. (a -> b) -> a -> b
    $ Names
    prelude_sizes Names -> Names -> Bool
    `namesIntersect` Names
    prelude_bound
      (Stms GPU, Stms GPU) -> Maybe (Stms GPU, Stms GPU)
    forall a. a -> Maybe a
    Just (Stms GPU
    prelude_used, Stms GPU
    prelude_not_used Stms GPU -> Stms GPU -> Stms GPU
    forall a. Semigroup a => a -> a -> a
    <> Stms GPU
    postlude)
      where
        used_in_tiled :: Names
    used_in_tiled = Stm GPU -> Names
    forall a. FreeIn a => a -> Names
    freeIn Stm GPU
    stm_to_tile

        used_in_stm_variant :: Names
    used_in_stm_variant =
          (Names
    used_in_tiled Names -> Names -> Names
    forall a. Semigroup a => a -> a -> a
    <>) (Names -> Names) -> Names -> Names
    forall a b. (a -> b) -> a -> b
    $
            [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 ((VName -> AliasTable -> Names) -> AliasTable -> VName -> Names
    forall a b c. (a -> b -> c) -> b -> a -> c
    flip (Names -> VName -> AliasTable -> Names
    forall k a. Ord k => a -> k -> Map k a -> a
    M.findWithDefault Names
    forall a. Monoid a => a
    mempty) AliasTable
    variance) ([VName] -> [Names]) -> [VName] -> [Names]
    forall a b. (a -> b) -> a -> b
    $
                Names -> [VName]
    namesToList Names
    used_in_tiled

        used :: Stm GPU -> Bool
    used Stm GPU
    stm =
          (VName -> Bool) -> [VName] -> Bool
    forall (t :: * -> *) a. Foldable t => (a -> Bool) -> t a -> Bool
    any (VName -> Names -> Bool
    `nameIn` Names
    used_in_stm_variant) ([VName] -> Bool) -> [VName] -> Bool
    forall a b. (a -> b) -> a -> b
    $
            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

        (Stms GPU
    prelude_used, Stms GPU
    prelude_not_used) =
          (Stm GPU -> Bool) -> Stms GPU -> (Stms GPU, Stms GPU)
    forall a. (a -> Bool) -> Seq a -> (Seq a, Seq a)
    Seq.partition Stm GPU -> Bool
    used Stms GPU
    prelude

    -- | Partition prelude statements preceding a tiled loop (or something
    -- containing a tiled loop) into three categories:
    --
    -- 1) Group-level statements that are invariant to the threads in the group.
    --
    -- 2) Thread-variant statements that should be computed once with a segmap_thread_scalar.
    --
    -- 3) Thread-variant statements that should be recomputed whenever
    -- they are needed.
    --
    -- The third category duplicates computation, so we only want to do it
    -- when absolutely necessary.  Currently, this is necessary for
    -- results that are views of an array (slicing, rotate, etc) and which
    -- results are used after the prelude, because these cannot be
    -- efficiently represented by a scalar segmap (they'll be manifested
    -- in memory).  To avoid unnecessarily moving computation from
    -- category 2 to category 3 simply because they depend on a category 3
    -- result, everything in category 3 is also in category 2.  This is
    -- efficient only when category 3 contains exclusively "free" or at
    -- least very cheap expressions (e.g. index space transformations).
    partitionPrelude ::
      VarianceTable ->
      Stms GPU ->
      Names ->
      Names ->
      (Stms GPU, Stms GPU, Stms GPU)
    partitionPrelude :: AliasTable
    -> Stms GPU -> Names -> Names -> (Stms GPU, Stms GPU, Stms GPU)
    partitionPrelude AliasTable
    variance Stms GPU
    prestms Names
    private Names
    used_after =
      (Stms GPU
    invariant_prestms, Stms GPU
    variant_prestms, Stms GPU
    recomputed_variant_prestms)
      where
        invariantTo :: Names -> Stm GPU -> Bool
    invariantTo Names
    names Stm GPU
    stm =
          case Pat Type -> [VName]
    forall dec. Pat dec -> [VName]
    patNames (Stm GPU -> Pat (LetDec GPU)
    forall rep. Stm rep -> Pat (LetDec rep)
    stmPat Stm GPU
    stm) of
            [] -> Bool
    True -- Does not matter.
            VName
    v : [VName]
    _ -> (VName -> Bool) -> [VName] -> Bool
    forall (t :: * -> *) a. Foldable t => (a -> Bool) -> t a -> Bool
    all (VName -> Names -> Bool
    `notNameIn` Names
    names) (Names -> [VName]
    namesToList (Names -> [VName]) -> Names -> [VName]
    forall a b. (a -> b) -> a -> b
    $ Names -> VName -> AliasTable -> Names
    forall k a. Ord k => a -> k -> Map k a -> a
    M.findWithDefault Names
    forall a. Monoid a => a
    mempty VName
    v AliasTable
    variance)

        consumed_in_prestms :: Names
    consumed_in_prestms =
          (Stm (Aliases GPU) -> Names) -> Seq (Stm (Aliases GPU)) -> Names
    forall m a. Monoid m => (a -> m) -> Seq a -> m
    forall (t :: * -> *) m a.
    (Foldable t, Monoid m) =>
    (a -> m) -> t a -> m
    foldMap Stm (Aliases GPU) -> Names
    forall rep. Aliased rep => Stm rep -> Names
    consumedInStm (Seq (Stm (Aliases GPU)) -> Names)
    -> Seq (Stm (Aliases GPU)) -> Names
    forall a b. (a -> b) -> a -> b
    $ (Seq (Stm (Aliases GPU)), AliasesAndConsumed)
    -> Seq (Stm (Aliases GPU))
    forall a b. (a, b) -> a
    fst ((Seq (Stm (Aliases GPU)), AliasesAndConsumed)
     -> Seq (Stm (Aliases GPU)))
    -> (Seq (Stm (Aliases GPU)), AliasesAndConsumed)
    -> Seq (Stm (Aliases GPU))
    forall a b. (a -> b) -> a -> b
    $ AliasTable
    -> Stms GPU -> (Seq (Stm (Aliases GPU)), AliasesAndConsumed)
    forall rep.
    AliasableRep rep =>
    AliasTable -> Stms rep -> (Stms (Aliases rep), AliasesAndConsumed)
    Alias.analyseStms AliasTable
    forall a. Monoid a => a
    mempty Stms GPU
    prestms
        consumed :: VName -> Bool
    consumed VName
    v = VName
    v VName -> Names -> Bool
    `nameIn` Names
    consumed_in_prestms
        consumedStm :: Stm GPU -> Bool
    consumedStm Stm GPU
    stm = (VName -> Bool) -> [VName] -> Bool
    forall (t :: * -> *) a. Foldable t => (a -> Bool) -> t a -> Bool
    any VName -> Bool
    consumed (Pat Type -> [VName]
    forall dec. Pat dec -> [VName]
    patNames (Stm GPU -> Pat (LetDec GPU)
    forall rep. Stm rep -> Pat (LetDec rep)
    stmPat Stm GPU
    stm))
        later_consumed :: Names
    later_consumed =
          [VName] -> Names
    namesFromList ([VName] -> Names) -> [VName] -> Names
    forall a b. (a -> b) -> a -> b
    $ (Stm GPU -> [VName]) -> Stms GPU -> [VName]
    forall m a. Monoid m => (a -> m) -> Seq a -> m
    forall (t :: * -> *) m a.
    (Foldable t, Monoid m) =>
    (a -> m) -> t a -> m
    foldMap (Pat Type -> [VName]
    forall dec. Pat dec -> [VName]
    patNames (Pat Type -> [VName])
    -> (Stm GPU -> Pat Type) -> Stm GPU -> [VName]
    forall b c a. (b -> c) -> (a -> b) -> a -> c
    . Stm GPU -> Pat Type
    Stm GPU -> Pat (LetDec GPU)
    forall rep. Stm rep -> Pat (LetDec rep)
    stmPat) (Stms GPU -> [VName]) -> Stms GPU -> [VName]
    forall a b. (a -> b) -> a -> b
    $ (Stm GPU -> Bool) -> Stms GPU -> Stms GPU
    forall a. (a -> Bool) -> Seq a -> Seq a
    Seq.filter Stm GPU -> Bool
    consumedStm Stms GPU
    prestms

        groupInvariant :: Stm GPU -> Bool
    groupInvariant Stm GPU
    stm =
          Names -> Stm GPU -> Bool
    invariantTo Names
    private Stm GPU
    stm
            Bool -> Bool -> Bool
    && (VName -> Bool) -> [VName] -> Bool
    forall (t :: * -> *) a. Foldable t => (a -> Bool) -> t a -> Bool
    all (VName -> Names -> Bool
    `notNameIn` Names
    later_consumed) (Pat Type -> [VName]
    forall dec. Pat dec -> [VName]
    patNames (Stm GPU -> Pat (LetDec GPU)
    forall rep. Stm rep -> Pat (LetDec rep)
    stmPat Stm GPU
    stm))
            Bool -> Bool -> Bool
    && Names -> Stm GPU -> Bool
    invariantTo Names
    later_consumed Stm GPU
    stm
        (Stms GPU
    invariant_prestms, Stms GPU
    variant_prestms) =
          (Stm GPU -> Bool) -> Stms GPU -> (Stms GPU, Stms GPU)
    forall a. (a -> Bool) -> Seq a -> (Seq a, Seq a)
    Seq.partition Stm GPU -> Bool
    groupInvariant Stms GPU
    prestms

        mustBeInlinedExp :: Exp rep -> Bool
    mustBeInlinedExp (BasicOp (Index VName
    _ Slice SubExp
    slice)) = Bool -> Bool
    not (Bool -> Bool) -> Bool -> Bool
    forall a b. (a -> b) -> a -> b
    $ [SubExp] -> Bool
    forall a. [a] -> Bool
    forall (t :: * -> *) a. Foldable t => t a -> Bool
    null ([SubExp] -> Bool) -> [SubExp] -> Bool
    forall a b. (a -> b) -> a -> b
    $ Slice SubExp -> [SubExp]
    forall d. Slice d -> [d]
    sliceDims Slice SubExp
    slice
        mustBeInlinedExp (BasicOp Iota {}) = Bool
    True
        mustBeInlinedExp (BasicOp Rearrange {}) = Bool
    True
        mustBeInlinedExp (BasicOp Reshape {}) = Bool
    True
        mustBeInlinedExp Exp rep
    _ = Bool
    False
        mustBeInlined :: Stm GPU -> Bool
    mustBeInlined Stm GPU
    stm =
          Exp GPU -> Bool
    forall {rep}. Exp rep -> Bool
    mustBeInlinedExp (Stm GPU -> Exp GPU
    forall rep. Stm rep -> Exp rep
    stmExp Stm GPU
    stm)
            Bool -> Bool -> Bool
    && (VName -> Bool) -> [VName] -> Bool
    forall (t :: * -> *) a. Foldable t => (a -> Bool) -> t a -> Bool
    any (VName -> Names -> Bool
    `nameIn` Names
    used_after) (Pat Type -> [VName]
    forall dec. Pat dec -> [VName]
    patNames (Stm GPU -> Pat (LetDec GPU)
    forall rep. Stm rep -> Pat (LetDec rep)
    stmPat Stm GPU
    stm))
        must_be_inlined :: Names
    must_be_inlined =
          [VName] -> Names
    namesFromList ([VName] -> Names) -> [VName] -> Names
    forall a b. (a -> b) -> a -> b
    $
            (Stm GPU -> [VName]) -> Stms GPU -> [VName]
    forall m a. Monoid m => (a -> m) -> Seq a -> m
    forall (t :: * -> *) m a.
    (Foldable t, Monoid m) =>
    (a -> m) -> t a -> m
    foldMap (Pat Type -> [VName]
    forall dec. Pat dec -> [VName]
    patNames (Pat Type -> [VName])
    -> (Stm GPU -> Pat Type) -> Stm GPU -> [VName]
    forall b c a. (b -> c) -> (a -> b) -> a -> c
    . Stm GPU -> Pat Type
    Stm GPU -> Pat (LetDec GPU)
    forall rep. Stm rep -> Pat (LetDec rep)
    stmPat) (Stms GPU -> [VName]) -> Stms GPU -> [VName]
    forall a b. (a -> b) -> a -> b
    $
              (Stm GPU -> Bool) -> Stms GPU -> Stms GPU
    forall a. (a -> Bool) -> Seq a -> Seq a
    Seq.filter Stm GPU -> Bool
    mustBeInlined Stms GPU
    variant_prestms
        recompute :: Stm GPU -> Bool
    recompute Stm GPU
    stm =
          (VName -> Bool) -> [VName] -> Bool
    forall (t :: * -> *) a. Foldable t => (a -> Bool) -> t a -> Bool
    any (VName -> Names -> Bool
    `nameIn` Names
    must_be_inlined) (Pat Type -> [VName]
    forall dec. Pat dec -> [VName]
    patNames (Stm GPU -> Pat (LetDec GPU)
    forall rep. Stm rep -> Pat (LetDec rep)
    stmPat Stm GPU
    stm))
        recomputed_variant_prestms :: Stms GPU
    recomputed_variant_prestms =
          (Stm GPU -> Bool) -> Stms GPU -> Stms GPU
    forall a. (a -> Bool) -> Seq a -> Seq a
    Seq.filter Stm GPU -> Bool
    recompute Stms GPU
    variant_prestms

    -- Anything that is variant to the "private" names should be
    -- considered thread-local.
    injectPrelude ::
      SegSpace ->
      VarianceTable ->
      Stms GPU ->
      Names ->
      (Stms GPU, Tiling, TiledBody) ->
      (Stms GPU, Tiling, TiledBody)
    injectPrelude :: SegSpace
    -> AliasTable
    -> Stms GPU
    -> Names
    -> (Stms GPU, Tiling, TiledBody)
    -> (Stms GPU, Tiling, TiledBody)
    injectPrelude SegSpace
    initial_space AliasTable
    variance Stms GPU
    prestms Names
    used (Stms GPU
    host_stms, Tiling
    tiling, TiledBody
    tiledBody) =
      (Stms GPU
    host_stms, Tiling
    tiling, TiledBody
    tiledBody')
      where
        tiledBody' :: TiledBody
    tiledBody' Names
    private PrivStms
    privstms = do
          let nontiled :: (VName, SubExp) -> Bool
    nontiled = ((VName, SubExp) -> [(VName, SubExp)] -> Bool
    forall (t :: * -> *) a. (Foldable t, Eq a) => a -> t a -> Bool
    `notElem` SegSpace -> [(VName, SubExp)]
    unSegSpace (Tiling -> SegSpace
    tilingSpace Tiling
    tiling))
              private' :: Names
    private' =
                Names
    private
                  Names -> Names -> Names
    forall a. Semigroup a => a -> a -> a
    <> [VName] -> Names
    namesFromList (((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) -> Bool) -> [(VName, SubExp)] -> [(VName, SubExp)]
    forall a. (a -> Bool) -> [a] -> [a]
    filter (VName, SubExp) -> Bool
    nontiled ([(VName, SubExp)] -> [(VName, SubExp)])
    -> [(VName, SubExp)] -> [(VName, SubExp)]
    forall a b. (a -> b) -> a -> b
    $ SegSpace -> [(VName, SubExp)]
    unSegSpace SegSpace
    initial_space))
              ( Stms GPU
    invariant_prestms,
                Stms GPU
    precomputed_variant_prestms,
                Stms GPU
    recomputed_variant_prestms
                ) =
                  AliasTable
    -> Stms GPU -> Names -> Names -> (Stms GPU, Stms GPU, Stms GPU)
    partitionPrelude AliasTable
    variance Stms GPU
    prestms Names
    private' Names
    used

          Stms (Rep (BuilderT GPU (State VNameSource)))
    -> BuilderT GPU (State VNameSource) ()
    forall (m :: * -> *). MonadBuilder m => Stms (Rep m) -> m ()
    addStms Stms (Rep (BuilderT GPU (State VNameSource)))
    Stms GPU
    invariant_prestms

          let live_set :: [VName]
    live_set =
                Names -> [VName]
    namesToList (Names -> [VName]) -> Names -> [VName]
    forall a b. (a -> b) -> a -> b
    $
                  Stms GPU -> Names -> Names
    forall a. FreeIn a => Stms GPU -> a -> Names
    liveSet Stms GPU
    precomputed_variant_prestms (Names -> Names) -> Names -> Names
    forall a b. (a -> b) -> a -> b
    $
                    Names
    used Names -> Names -> Names
    forall a. Semigroup a => a -> a -> a
    <> Stms GPU -> Names
    forall a. FreeIn a => a -> Names
    freeIn Stms GPU
    recomputed_variant_prestms
          prelude_arrs <-
            Stms GPU
    -> BuilderT GPU (State VNameSource) [VName]
    -> BuilderT GPU (State VNameSource) [VName]
    forall rep a (m :: * -> *) b.
    (Scoped rep a, LocalScope rep m) =>
    a -> m b -> m b
    inScopeOf Stms GPU
    precomputed_variant_prestms (BuilderT GPU (State VNameSource) [VName]
     -> BuilderT GPU (State VNameSource) [VName])
    -> BuilderT GPU (State VNameSource) [VName]
    -> BuilderT GPU (State VNameSource) [VName]
    forall a b. (a -> b) -> a -> b
    $
              Tiling
    -> PrivStms
    -> Stms GPU
    -> [VName]
    -> BuilderT GPU (State VNameSource) [VName]
    doPrelude Tiling
    tiling PrivStms
    privstms Stms GPU
    precomputed_variant_prestms [VName]
    live_set

          let prelude_privstms =
                Stms GPU -> ReadPrelude -> PrivStms
    PrivStms Stms GPU
    recomputed_variant_prestms (ReadPrelude -> PrivStms) -> ReadPrelude -> PrivStms
    forall a b. (a -> b) -> a -> b
    $
                  [VName] -> [VName] -> ReadPrelude
    mkReadPreludeValues [VName]
    prelude_arrs [VName]
    live_set

          tiledBody private' (prelude_privstms <> privstms)

    tileLoop ::
      SegSpace ->
      VarianceTable ->
      Stms GPU ->
      Names ->
      (Stms GPU, Tiling, TiledBody) ->
      [Type] ->
      Pat Type ->
      StmAux (ExpDec GPU) ->
      [(FParam GPU, SubExp)] ->
      VName ->
      IntType ->
      SubExp ->
      Stms GPU ->
      Result ->
      TileM (Stms GPU, Tiling, TiledBody)
    tileLoop :: SegSpace
    -> AliasTable
    -> Stms GPU
    -> Names
    -> (Stms GPU, Tiling, TiledBody)
    -> [Type]
    -> Pat Type
    -> StmAux (ExpDec GPU)
    -> [(FParam GPU, SubExp)]
    -> VName
    -> IntType
    -> SubExp
    -> Stms GPU
    -> Result
    -> ReaderT
         (Scope GPU) (State VNameSource) (Stms GPU, Tiling, TiledBody)
    tileLoop SegSpace
    initial_space AliasTable
    variance Stms GPU
    prestms Names
    used_in_body (Stms GPU
    host_stms, Tiling
    tiling, TiledBody
    tiledBody) [Type]
    res_ts Pat Type
    pat StmAux (ExpDec GPU)
    aux [(FParam GPU, SubExp)]
    merge VName
    i IntType
    it SubExp
    bound Stms GPU
    poststms Result
    poststms_res = do
      let prestms_used :: Names
    prestms_used = Names
    used_in_body Names -> Names -> Names
    forall a. Semigroup a => a -> a -> a
    <> Stms GPU -> Names
    forall a. FreeIn a => a -> Names
    freeIn Stms GPU
    poststms Names -> Names -> Names
    forall a. Semigroup a => a -> a -> a
    <> Result -> Names
    forall a. FreeIn a => a -> Names
    freeIn Result
    poststms_res
          ( Stms GPU
    invariant_prestms,
            Stms GPU
    precomputed_variant_prestms,
            Stms GPU
    recomputed_variant_prestms
            ) =
              AliasTable
    -> Stms GPU -> Names -> Names -> (Stms GPU, Stms GPU, Stms GPU)
    partitionPrelude AliasTable
    variance Stms GPU
    prestms Names
    tiled_kdims Names
    prestms_used

      let ([Param (TypeBase Shape Uniqueness)]
    mergeparams, [SubExp]
    mergeinits) = [(Param (TypeBase Shape Uniqueness), SubExp)]
    -> ([Param (TypeBase Shape Uniqueness)], [SubExp])
    forall a b. [(a, b)] -> ([a], [b])
    unzip [(Param (TypeBase Shape Uniqueness), SubExp)]
    [(FParam GPU, SubExp)]
    merge

          -- Expand the loop merge parameters to be arrays.
          tileDim :: TypeBase Shape Uniqueness -> TypeBase Shape Uniqueness
    tileDim TypeBase Shape Uniqueness
    t = TypeBase Shape Uniqueness
    -> Shape -> Uniqueness -> TypeBase Shape Uniqueness
    forall shape u_unused u.
    ArrayShape shape =>
    TypeBase shape u_unused -> shape -> u -> TypeBase shape u
    arrayOf TypeBase Shape Uniqueness
    t (Tiling -> Shape
    tilingTileShape Tiling
    tiling) (Uniqueness -> TypeBase Shape Uniqueness)
    -> Uniqueness -> TypeBase Shape Uniqueness
    forall a b. (a -> b) -> a -> b
    $ TypeBase Shape Uniqueness -> Uniqueness
    forall shape. TypeBase shape Uniqueness -> Uniqueness
    uniqueness TypeBase Shape Uniqueness
    t

          merge_scope :: Scope GPU
    merge_scope = VName -> NameInfo GPU -> Scope GPU -> Scope GPU
    forall k a. Ord k => k -> a -> Map k a -> Map k a
    M.insert VName
    i (IntType -> NameInfo GPU
    forall rep. IntType -> NameInfo rep
    IndexName IntType
    it) (Scope GPU -> Scope GPU) -> Scope GPU -> Scope GPU
    forall a b. (a -> b) -> a -> b
    $ [Param (TypeBase Shape Uniqueness)] -> Scope GPU
    forall rep dec. (FParamInfo rep ~ dec) => [Param dec] -> Scope rep
    scopeOfFParams [Param (TypeBase Shape Uniqueness)]
    mergeparams

          tiledBody' :: TiledBody
    tiledBody' Names
    private PrivStms
    privstms = Scope GPU
    -> BuilderT GPU (State VNameSource) [VName]
    -> BuilderT GPU (State VNameSource) [VName]
    forall a.
    Scope GPU
    -> BuilderT GPU (State VNameSource) a
    -> BuilderT 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
    host_stms Scope GPU -> Scope GPU -> Scope GPU
    forall a. Semigroup a => a -> a -> a
    <> Scope GPU
    merge_scope) (BuilderT GPU (State VNameSource) [VName]
     -> BuilderT GPU (State VNameSource) [VName])
    -> BuilderT GPU (State VNameSource) [VName]
    -> BuilderT GPU (State VNameSource) [VName]
    forall a b. (a -> b) -> a -> b
    $ do
            Stms (Rep (BuilderT GPU (State VNameSource)))
    -> BuilderT GPU (State VNameSource) ()
    forall (m :: * -> *). MonadBuilder m => Stms (Rep m) -> m ()
    addStms Stms (Rep (BuilderT GPU (State VNameSource)))
    Stms GPU
    invariant_prestms

            let live_set :: [VName]
    live_set =
                  Names -> [VName]
    namesToList (Names -> [VName]) -> Names -> [VName]
    forall a b. (a -> b) -> a -> b
    $
                    Stms GPU -> Names -> Names
    forall a. FreeIn a => Stms GPU -> a -> Names
    liveSet Stms GPU
    precomputed_variant_prestms (Names -> Names) -> Names -> Names
    forall a b. (a -> b) -> a -> b
    $
                      Stms GPU -> Names
    forall a. FreeIn a => a -> Names
    freeIn Stms GPU
    recomputed_variant_prestms Names -> Names -> Names
    forall a. Semigroup a => a -> a -> a
    <> Names
    prestms_used

            prelude_arrs <-
              Stms GPU
    -> BuilderT GPU (State VNameSource) [VName]
    -> BuilderT GPU (State VNameSource) [VName]
    forall rep a (m :: * -> *) b.
    (Scoped rep a, LocalScope rep m) =>
    a -> m b -> m b
    inScopeOf Stms GPU
    precomputed_variant_prestms (BuilderT GPU (State VNameSource) [VName]
     -> BuilderT GPU (State VNameSource) [VName])
    -> BuilderT GPU (State VNameSource) [VName]
    -> BuilderT GPU (State VNameSource) [VName]
    forall a b. (a -> b) -> a -> b
    $
                Tiling
    -> PrivStms
    -> Stms GPU
    -> [VName]
    -> BuilderT GPU (State VNameSource) [VName]
    doPrelude Tiling
    tiling PrivStms
    privstms Stms GPU
    precomputed_variant_prestms [VName]
    live_set

            mergeparams' <- forM mergeparams $ \(Param Attrs
    attrs VName
    pname TypeBase Shape Uniqueness
    pt) ->
              Attrs
    -> VName
    -> TypeBase Shape Uniqueness
    -> Param (TypeBase Shape Uniqueness)
    forall dec. Attrs -> VName -> dec -> Param dec
    Param Attrs
    attrs (VName
     -> TypeBase Shape Uniqueness -> Param (TypeBase Shape Uniqueness))
    -> BuilderT GPU (State VNameSource) VName
    -> BuilderT
         GPU
         (State VNameSource)
         (TypeBase Shape Uniqueness -> Param (TypeBase Shape Uniqueness))
    forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
    <$> [Char] -> BuilderT GPU (State VNameSource) VName
    forall (m :: * -> *). MonadFreshNames m => [Char] -> m VName
    newVName (VName -> [Char]
    baseString VName
    pname [Char] -> [Char] -> [Char]
    forall a. [a] -> [a] -> [a]
    ++ [Char]
    "_group") BuilderT
      GPU
      (State VNameSource)
      (TypeBase Shape Uniqueness -> Param (TypeBase Shape Uniqueness))
    -> BuilderT GPU (State VNameSource) (TypeBase Shape Uniqueness)
    -> BuilderT
         GPU (State VNameSource) (Param (TypeBase Shape Uniqueness))
    forall a b.
    BuilderT GPU (State VNameSource) (a -> b)
    -> BuilderT GPU (State VNameSource) a
    -> BuilderT GPU (State VNameSource) b
    forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
    <*> TypeBase Shape Uniqueness
    -> BuilderT GPU (State VNameSource) (TypeBase Shape Uniqueness)
    forall a. a -> BuilderT GPU (State VNameSource) a
    forall (f :: * -> *) a. Applicative f => a -> f a
    pure (TypeBase Shape Uniqueness -> TypeBase Shape Uniqueness
    tileDim TypeBase Shape Uniqueness
    pt)

            let merge_ts = (Param (TypeBase Shape Uniqueness) -> Type)
    -> [Param (TypeBase Shape Uniqueness)] -> [Type]
    forall a b. (a -> b) -> [a] -> [b]
    map Param (TypeBase Shape Uniqueness) -> Type
    forall dec. Typed dec => Param dec -> Type
    paramType [Param (TypeBase Shape Uniqueness)]
    mergeparams

            let inloop_privstms =
                  Stms GPU -> ReadPrelude -> PrivStms
    PrivStms Stms GPU
    recomputed_variant_prestms (ReadPrelude -> PrivStms) -> ReadPrelude -> PrivStms
    forall a b. (a -> b) -> a -> b
    $
                    [VName] -> [VName] -> ReadPrelude
    mkReadPreludeValues [VName]
    prelude_arrs [VName]
    live_set

            mergeinit' <-
              fmap (map Var) $
                certifying (stmAuxCerts aux) $
                  tilingSegMap tiling "tiled_loopinit" ResultPrivate $
                    \PrimExp VName
    in_bounds [DimIndex SubExp]
    slice ->
                      ([VName] -> Result)
    -> BuilderT GPU (State VNameSource) [VName] -> Builder GPU Result
    forall a b.
    (a -> b)
    -> BuilderT GPU (State VNameSource) a
    -> BuilderT GPU (State VNameSource) b
    forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
    fmap [VName] -> Result
    varsRes (BuilderT GPU (State VNameSource) [VName] -> Builder GPU Result)
    -> BuilderT GPU (State VNameSource) [VName] -> Builder GPU Result
    forall a b. (a -> b) -> a -> b
    $
                        [Char]
    -> PrimExp VName
    -> [Type]
    -> Builder GPU Result
    -> BuilderT GPU (State VNameSource) [VName]
    protectOutOfBounds [Char]
    "loopinit" PrimExp VName
    in_bounds [Type]
    merge_ts (Builder GPU Result -> BuilderT GPU (State VNameSource) [VName])
    -> Builder GPU Result -> BuilderT GPU (State VNameSource) [VName]
    forall a b. (a -> b) -> a -> b
    $ do
                          [DimIndex SubExp]
    -> PrivStms -> BuilderT GPU (State VNameSource) ()
    addPrivStms [DimIndex SubExp]
    slice PrivStms
    inloop_privstms
                          [DimIndex SubExp]
    -> PrivStms -> BuilderT GPU (State VNameSource) ()
    addPrivStms [DimIndex SubExp]
    slice PrivStms
    privstms
                          Result -> Builder GPU Result
    forall a. a -> BuilderT GPU (State VNameSource) a
    forall (f :: * -> *) a. Applicative f => a -> f a
    pure (Result -> Builder GPU Result) -> Result -> Builder GPU Result
    forall a b. (a -> b) -> a -> b
    $ [SubExp] -> Result
    subExpsRes [SubExp]
    mergeinits

            let merge' = [Param (TypeBase Shape Uniqueness)]
    -> [SubExp] -> [(Param (TypeBase Shape Uniqueness), SubExp)]
    forall a b. [a] -> [b] -> [(a, b)]
    zip [Param (TypeBase Shape Uniqueness)]
    mergeparams' [SubExp]
    mergeinit'

            let indexLoopParams [DimIndex SubExp]
    slice =
                  Scope GPU
    -> BuilderT GPU (State VNameSource) ()
    -> BuilderT GPU (State VNameSource) ()
    forall a.
    Scope GPU
    -> BuilderT GPU (State VNameSource) a
    -> BuilderT GPU (State VNameSource) a
    forall rep (m :: * -> *) a.
    LocalScope rep m =>
    Scope rep -> m a -> m a
    localScope ([Param (TypeBase Shape Uniqueness)] -> Scope GPU
    forall rep dec. (FParamInfo rep ~ dec) => [Param dec] -> Scope rep
    scopeOfFParams [Param (TypeBase Shape Uniqueness)]
    mergeparams') (BuilderT GPU (State VNameSource) ()
     -> BuilderT GPU (State VNameSource) ())
    -> BuilderT GPU (State VNameSource) ()
    -> BuilderT GPU (State VNameSource) ()
    forall a b. (a -> b) -> a -> b
    $
                    [(Param (TypeBase Shape Uniqueness),
      Param (TypeBase Shape Uniqueness))]
    -> ((Param (TypeBase Shape Uniqueness),
         Param (TypeBase Shape Uniqueness))
        -> BuilderT GPU (State VNameSource) ())
    -> BuilderT GPU (State VNameSource) ()
    forall (t :: * -> *) (m :: * -> *) a b.
    (Foldable t, Monad m) =>
    t a -> (a -> m b) -> m ()
    forM_ ([Param (TypeBase Shape Uniqueness)]
    -> [Param (TypeBase Shape Uniqueness)]
    -> [(Param (TypeBase Shape Uniqueness),
         Param (TypeBase Shape Uniqueness))]
    forall a b. [a] -> [b] -> [(a, b)]
    zip [Param (TypeBase Shape Uniqueness)]
    mergeparams [Param (TypeBase Shape Uniqueness)]
    mergeparams') (((Param (TypeBase Shape Uniqueness),
       Param (TypeBase Shape Uniqueness))
      -> BuilderT GPU (State VNameSource) ())
     -> BuilderT GPU (State VNameSource) ())
    -> ((Param (TypeBase Shape Uniqueness),
         Param (TypeBase Shape Uniqueness))
        -> BuilderT GPU (State VNameSource) ())
    -> BuilderT GPU (State VNameSource) ()
    forall a b. (a -> b) -> a -> b
    $ \(Param (TypeBase Shape Uniqueness)
    to, Param (TypeBase Shape Uniqueness)
    from) ->
                      [VName]
    -> Exp (Rep (BuilderT GPU (State VNameSource)))
    -> BuilderT GPU (State VNameSource) ()
    forall (m :: * -> *).
    MonadBuilder m =>
    [VName] -> Exp (Rep m) -> m ()
    letBindNames [Param (TypeBase Shape Uniqueness) -> VName
    forall dec. Param dec -> VName
    paramName Param (TypeBase Shape Uniqueness)
    to] (Exp GPU -> BuilderT GPU (State VNameSource) ())
    -> (Slice SubExp -> Exp GPU)
    -> Slice SubExp
    -> BuilderT GPU (State VNameSource) ()
    forall b c a. (b -> c) -> (a -> b) -> a -> c
    . BasicOp -> Exp GPU
    forall rep. BasicOp -> Exp rep
    BasicOp (BasicOp -> Exp GPU)
    -> (Slice SubExp -> BasicOp) -> Slice SubExp -> Exp GPU
    forall b c a. (b -> c) -> (a -> b) -> a -> c
    . VName -> Slice SubExp -> BasicOp
    Index (Param (TypeBase Shape Uniqueness) -> VName
    forall dec. Param dec -> VName
    paramName Param (TypeBase Shape Uniqueness)
    from) (Slice SubExp -> BuilderT GPU (State VNameSource) ())
    -> Slice SubExp -> BuilderT GPU (State VNameSource) ()
    forall a b. (a -> b) -> a -> b
    $
                        Type -> [DimIndex SubExp] -> Slice SubExp
    fullSlice (Param (TypeBase Shape Uniqueness) -> Type
    forall dec. Typed dec => Param dec -> Type
    paramType Param (TypeBase Shape Uniqueness)
    from) [DimIndex SubExp]
    slice

                private' =
                  Names
    private Names -> Names -> Names
    forall a. Semigroup a => a -> a -> a
    <> [VName] -> Names
    namesFromList ((Param (TypeBase Shape Uniqueness) -> VName)
    -> [Param (TypeBase Shape Uniqueness)] -> [VName]
    forall a b. (a -> b) -> [a] -> [b]
    map Param (TypeBase Shape Uniqueness) -> VName
    forall dec. Param dec -> VName
    paramName [Param (TypeBase Shape Uniqueness)]
    mergeparams [VName] -> [VName] -> [VName]
    forall a. [a] -> [a] -> [a]
    ++ (Param (TypeBase Shape Uniqueness) -> VName)
    -> [Param (TypeBase Shape Uniqueness)] -> [VName]
    forall a b. (a -> b) -> [a] -> [b]
    map Param (TypeBase Shape Uniqueness) -> VName
    forall dec. Param dec -> VName
    paramName [Param (TypeBase Shape Uniqueness)]
    mergeparams')

                privstms' =
                  Stms GPU -> ReadPrelude -> PrivStms
    PrivStms Stms GPU
    forall a. Monoid a => a
    mempty ReadPrelude
    indexLoopParams PrivStms -> PrivStms -> PrivStms
    forall a. Semigroup a => a -> a -> a
    <> PrivStms
    privstms PrivStms -> PrivStms -> PrivStms
    forall a. Semigroup a => a -> a -> a
    <> PrivStms
    inloop_privstms

            loopbody' <-
              localScope (scopeOfFParams mergeparams') . runBodyBuilder $
                varsRes <$> tiledBody private' privstms'
            accs' <-
              letTupExp "tiled_inside_loop" $
                Loop merge' (ForLoop i it bound) loopbody'

            postludeGeneric tiling (privstms <> inloop_privstms) pat accs' poststms poststms_res res_ts

      (Stms GPU, Tiling, TiledBody)
    -> ReaderT
         (Scope GPU) (State VNameSource) (Stms GPU, Tiling, TiledBody)
    forall a. a -> ReaderT (Scope GPU) (State VNameSource) a
    forall (f :: * -> *) a. Applicative f => a -> f a
    pure (Stms GPU
    host_stms, Tiling
    tiling, TiledBody
    tiledBody')
      where
        tiled_kdims :: Names
    tiled_kdims =
          [VName] -> Names
    namesFromList ([VName] -> Names) -> [VName] -> Names
    forall a b. (a -> b) -> a -> b
    $
            ((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
    $
              ((VName, SubExp) -> Bool) -> [(VName, SubExp)] -> [(VName, SubExp)]
    forall a. (a -> Bool) -> [a] -> [a]
    filter ((VName, SubExp) -> [(VName, SubExp)] -> Bool
    forall (t :: * -> *) a. (Foldable t, Eq a) => a -> t a -> Bool
    `notElem` SegSpace -> [(VName, SubExp)]
    unSegSpace (Tiling -> SegSpace
    tilingSpace Tiling
    tiling)) ([(VName, SubExp)] -> [(VName, SubExp)])
    -> [(VName, SubExp)] -> [(VName, SubExp)]
    forall a b. (a -> b) -> a -> b
    $
                SegSpace -> [(VName, SubExp)]
    unSegSpace SegSpace
    initial_space

    doPrelude :: Tiling -> PrivStms -> Stms GPU -> [VName] -> Builder GPU [VName]
    doPrelude :: Tiling
    -> PrivStms
    -> Stms GPU
    -> [VName]
    -> BuilderT GPU (State VNameSource) [VName]
    doPrelude Tiling
    tiling PrivStms
    privstms Stms GPU
    prestms [VName]
    prestms_live =
      -- Create a SegMap that takes care of the prelude for every thread.
      Tiling
    -> [Char]
    -> ResultManifest
    -> (PrimExp VName -> [DimIndex SubExp] -> Builder GPU Result)
    -> BuilderT GPU (State VNameSource) [VName]
    tilingSegMap Tiling
    tiling [Char]
    "prelude" ResultManifest
    ResultPrivate ((PrimExp VName -> [DimIndex SubExp] -> Builder GPU Result)
     -> BuilderT GPU (State VNameSource) [VName])
    -> (PrimExp VName -> [DimIndex SubExp] -> Builder GPU Result)
    -> BuilderT GPU (State VNameSource) [VName]
    forall a b. (a -> b) -> a -> b
    $ \PrimExp VName
    in_bounds [DimIndex SubExp]
    slice -> do
        ts <- (VName -> BuilderT GPU (State VNameSource) Type)
    -> [VName] -> BuilderT GPU (State VNameSource) [Type]
    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 VName -> BuilderT GPU (State VNameSource) Type
    forall rep (m :: * -> *). HasScope rep m => VName -> m Type
    lookupType [VName]
    prestms_live
        fmap varsRes . protectOutOfBounds "pre" in_bounds ts $ do
          addPrivStms slice privstms
          addStms prestms
          pure $ varsRes prestms_live

    liveSet :: (FreeIn a) => Stms GPU -> a -> Names
    liveSet :: forall a. FreeIn a => Stms GPU -> a -> Names
    liveSet Stms GPU
    stms a
    after =
      [VName] -> Names
    namesFromList ((Stm GPU -> [VName]) -> Stms GPU -> [VName]
    forall (t :: * -> *) a b. Foldable t => (a -> [b]) -> t a -> [b]
    concatMap (Pat Type -> [VName]
    forall dec. Pat dec -> [VName]
    patNames (Pat Type -> [VName])
    -> (Stm GPU -> Pat Type) -> Stm GPU -> [VName]
    forall b c a. (b -> c) -> (a -> b) -> a -> c
    . Stm GPU -> Pat Type
    Stm GPU -> Pat (LetDec GPU)
    forall rep. Stm rep -> Pat (LetDec rep)
    stmPat) Stms GPU
    stms)
        Names -> Names -> Names
    `namesIntersection` a -> Names
    forall a. FreeIn a => a -> Names
    freeIn a
    after

    tileable ::
      Stm GPU ->
      Maybe
        ( SubExp,
          [VName],
          (Commutativity, Lambda GPU, [SubExp], Lambda GPU)
        )
    tileable :: Stm GPU
    -> Maybe
         (SubExp, [VName],
          (Commutativity, Lambda GPU, [SubExp], Lambda GPU))
    tileable Stm GPU
    stm
      | Op (OtherOp (Screma SubExp
    w [VName]
    arrs ScremaForm GPU
    form)) <- Stm GPU -> Exp GPU
    forall rep. Stm rep -> Exp rep
    stmExp Stm GPU
    stm,
        Just ([Reduce GPU]
    reds, Lambda GPU
    map_lam) <- ScremaForm GPU -> Maybe ([Reduce GPU], Lambda GPU)
    forall rep. ScremaForm rep -> Maybe ([Reduce rep], Lambda rep)
    isRedomapSOAC ScremaForm GPU
    form,
        Reduce Commutativity
    red_comm Lambda GPU
    red_lam [SubExp]
    red_nes <- [Reduce GPU] -> Reduce GPU
    forall rep. Buildable rep => [Reduce rep] -> Reduce rep
    singleReduce [Reduce GPU]
    reds,
        Lambda GPU -> [Type]
    forall rep. Lambda rep -> [Type]
    lambdaReturnType Lambda GPU
    map_lam [Type] -> [Type] -> Bool
    forall a. Eq a => a -> a -> Bool
    == Lambda GPU -> [Type]
    forall rep. Lambda rep -> [Type]
    lambdaReturnType Lambda GPU
    red_lam, -- No mapout arrays.
        Bool -> Bool
    not (Bool -> Bool) -> Bool -> Bool
    forall a b. (a -> b) -> a -> b
    $ [VName] -> Bool
    forall a. [a] -> Bool
    forall (t :: * -> *) a. Foldable t => t a -> Bool
    null [VName]
    arrs,
        (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
    map_lam,
        (Param Type -> Bool) -> [Param 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) -> (Param Type -> Type) -> Param Type -> Bool
    forall b c a. (b -> c) -> (a -> b) -> a -> c
    . Param Type -> Type
    forall dec. Typed dec => Param dec -> Type
    paramType) ([Param Type] -> Bool) -> [Param Type] -> Bool
    forall a b. (a -> b) -> a -> b
    $ Lambda GPU -> [LParam GPU]
    forall rep. Lambda rep -> [LParam rep]
    lambdaParams Lambda GPU
    map_lam,
        Bool -> Bool
    not (Bool -> Bool) -> Bool -> Bool
    forall a b. (a -> b) -> a -> b
    $ Attr
    "unroll" Attr -> Attrs -> Bool
    `inAttrs` StmAux () -> Attrs
    forall dec. StmAux dec -> Attrs
    stmAuxAttrs (Stm GPU -> StmAux (ExpDec GPU)
    forall rep. Stm rep -> StmAux (ExpDec rep)
    stmAux Stm GPU
    stm) =
          (SubExp, [VName],
     (Commutativity, Lambda GPU, [SubExp], Lambda GPU))
    -> Maybe
         (SubExp, [VName],
          (Commutativity, Lambda GPU, [SubExp], Lambda GPU))
    forall a. a -> Maybe a
    Just (SubExp
    w, [VName]
    arrs, (Commutativity
    red_comm, Lambda GPU
    red_lam, [SubExp]
    red_nes, Lambda GPU
    map_lam))
      | Bool
    otherwise =
          Maybe
      (SubExp, [VName],
       (Commutativity, Lambda GPU, [SubExp], Lambda GPU))
    forall a. Maybe a
    Nothing

    -- | We classify the inputs to the tiled loop as whether they are
    -- tileable (and with what permutation of the kernel indexes) or not.
    -- In practice, we should have at least one tileable array per loop,
    -- but this is not enforced in our representation.
    data InputArray
      = InputTile [Int] VName
      | InputDontTile VName

    tiledInputs :: [InputArray] -> [(VName, [Int])]
    tiledInputs :: [InputArray] -> [(VName, [Int])]
    tiledInputs = (InputArray -> Maybe (VName, [Int]))
    -> [InputArray] -> [(VName, [Int])]
    forall a b. (a -> Maybe b) -> [a] -> [b]
    mapMaybe InputArray -> Maybe (VName, [Int])
    f
      where
        f :: InputArray -> Maybe (VName, [Int])
    f (InputTile [Int]
    perm VName
    arr) = (VName, [Int]) -> Maybe (VName, [Int])
    forall a. a -> Maybe a
    Just (VName
    arr, [Int]
    perm)
        f InputDontTile {} = Maybe (VName, [Int])
    forall a. Maybe a
    Nothing

    -- | A tile (or an original untiled array).
    data InputTile
      = InputTiled [Int] VName
      | InputUntiled VName

    -- First VNames are the tiles, second are the untiled.
    inputsToTiles :: [InputArray] -> [VName] -> [InputTile]
    inputsToTiles :: [InputArray] -> [VName] -> [InputTile]
    inputsToTiles (InputTile [Int]
    perm VName
    _ : [InputArray]
    inputs) (VName
    tile : [VName]
    tiles) =
      [Int] -> VName -> InputTile
    InputTiled [Int]
    perm VName
    tile InputTile -> [InputTile] -> [InputTile]
    forall a. a -> [a] -> [a]
    : [InputArray] -> [VName] -> [InputTile]
    inputsToTiles [InputArray]
    inputs [VName]
    tiles
    inputsToTiles (InputDontTile VName
    arr : [InputArray]
    inputs) [VName]
    tiles =
      VName -> InputTile
    InputUntiled VName
    arr InputTile -> [InputTile] -> [InputTile]
    forall a. a -> [a] -> [a]
    : [InputArray] -> [VName] -> [InputTile]
    inputsToTiles [InputArray]
    inputs [VName]
    tiles
    inputsToTiles [InputArray]
    _ [VName]
    _ = []

    -- The atual tile size may be smaller for the last tile, so we have to
    -- be careful now.
    sliceUntiled ::
      (MonadBuilder m) =>
      VName ->
      SubExp ->
      SubExp ->
      SubExp ->
      m VName
    sliceUntiled :: forall (m :: * -> *).
    MonadBuilder m =>
    VName -> SubExp -> SubExp -> SubExp -> m VName
    sliceUntiled VName
    arr SubExp
    tile_id SubExp
    full_tile_size SubExp
    this_tile_size = do
      arr_t <- VName -> m Type
    forall rep (m :: * -> *). HasScope rep m => VName -> m Type
    lookupType VName
    arr
      slice_offset <-
        letSubExp "slice_offset" =<< toExp (pe64 tile_id * pe64 full_tile_size)
      let slice = SubExp -> SubExp -> SubExp -> DimIndex SubExp
    forall d. d -> d -> d -> DimIndex d
    DimSlice SubExp
    slice_offset SubExp
    this_tile_size (IntType -> Integer -> SubExp
    intConst IntType
    Int64 Integer
    1)
      letExp "untiled_slice" $
        BasicOp $
          Index arr $
            fullSlice arr_t [slice]

    -- | Statements that we insert directly into every thread-private
    -- SegMaps.  This is for things that cannot efficiently be computed
    -- once in advance in the prelude SegMap, primarily (exclusively?)
    -- array slicing operations.
    data PrivStms = PrivStms (Stms GPU) ReadPrelude

    privStms :: Stms GPU -> PrivStms
    privStms :: Stms GPU -> PrivStms
    privStms Stms GPU
    stms = Stms GPU -> ReadPrelude -> PrivStms
    PrivStms Stms GPU
    stms (ReadPrelude -> PrivStms) -> ReadPrelude -> PrivStms
    forall a b. (a -> b) -> a -> b
    $ BuilderT GPU (State VNameSource) () -> ReadPrelude
    forall a b. a -> b -> a
    const (BuilderT GPU (State VNameSource) () -> ReadPrelude)
    -> BuilderT GPU (State VNameSource) () -> ReadPrelude
    forall a b. (a -> b) -> a -> b
    $ () -> BuilderT GPU (State VNameSource) ()
    forall a. a -> BuilderT GPU (State VNameSource) a
    forall (f :: * -> *) a. Applicative f => a -> f a
    pure ()

    addPrivStms :: [DimIndex SubExp] -> PrivStms -> Builder GPU ()
    addPrivStms :: [DimIndex SubExp]
    -> PrivStms -> BuilderT GPU (State VNameSource) ()
    addPrivStms [DimIndex SubExp]
    local_slice (PrivStms Stms GPU
    stms ReadPrelude
    readPrelude) = do
      ReadPrelude
    readPrelude [DimIndex SubExp]
    local_slice
      Stms (Rep (BuilderT GPU (State VNameSource)))
    -> BuilderT GPU (State VNameSource) ()
    forall (m :: * -> *). MonadBuilder m => Stms (Rep m) -> m ()
    addStms Stms (Rep (BuilderT GPU (State VNameSource)))
    Stms GPU
    stms

    instance Semigroup PrivStms where
      PrivStms Stms GPU
    stms_x ReadPrelude
    readPrelude_x <> :: PrivStms -> PrivStms -> PrivStms
    <> PrivStms Stms GPU
    stms_y ReadPrelude
    readPrelude_y =
        Stms GPU -> ReadPrelude -> PrivStms
    PrivStms Stms GPU
    stms_z ReadPrelude
    readPrelude_z
        where
          stms_z :: Stms GPU
    stms_z = Stms GPU
    stms_x Stms GPU -> Stms GPU -> Stms GPU
    forall a. Semigroup a => a -> a -> a
    <> Stms GPU
    stms_y
          readPrelude_z :: ReadPrelude
    readPrelude_z [DimIndex SubExp]
    slice = ReadPrelude
    readPrelude_x [DimIndex SubExp]
    slice BuilderT GPU (State VNameSource) ()
    -> BuilderT GPU (State VNameSource) ()
    -> BuilderT GPU (State VNameSource) ()
    forall a b.
    BuilderT GPU (State VNameSource) a
    -> BuilderT GPU (State VNameSource) b
    -> BuilderT GPU (State VNameSource) b
    forall (m :: * -> *) a b. Monad m => m a -> m b -> m b
    >> ReadPrelude
    readPrelude_y [DimIndex SubExp]
    slice

    instance Monoid PrivStms where
      mempty :: PrivStms
    mempty = Stms GPU -> PrivStms
    privStms Stms GPU
    forall a. Monoid a => a
    mempty

    type ReadPrelude = [DimIndex SubExp] -> Builder GPU ()

    data ProcessTileArgs = ProcessTileArgs
      { ProcessTileArgs -> PrivStms
    processPrivStms :: PrivStms,
        ProcessTileArgs -> Commutativity
    processComm :: Commutativity,
        ProcessTileArgs -> Lambda GPU
    processRedLam :: Lambda GPU,
        ProcessTileArgs -> Lambda GPU
    processMapLam :: Lambda GPU,
        ProcessTileArgs -> [InputTile]
    processTiles :: [InputTile],
        ProcessTileArgs -> [VName]
    processAcc :: [VName],
        ProcessTileArgs -> SubExp
    processTileId :: SubExp
      }

    data ResidualTileArgs = ResidualTileArgs
      { ResidualTileArgs -> PrivStms
    residualPrivStms :: PrivStms,
        ResidualTileArgs -> Commutativity
    residualComm :: Commutativity,
        ResidualTileArgs -> Lambda GPU
    residualRedLam :: Lambda GPU,
        ResidualTileArgs -> Lambda GPU
    residualMapLam :: Lambda GPU,
        ResidualTileArgs -> [InputArray]
    residualInput :: [InputArray],
        ResidualTileArgs -> [VName]
    residualAcc :: [VName],
        ResidualTileArgs -> SubExp
    residualInputSize :: SubExp,
        ResidualTileArgs -> SubExp
    residualNumWholeTiles :: SubExp
      }

    -- | Information about a loop that has been tiled inside a kernel, as
    -- well as the kinds of changes that we would then like to perform on
    -- the kernel.
    data Tiling = Tiling
      { Tiling
    -> [Char]
    -> ResultManifest
    -> (PrimExp VName -> [DimIndex SubExp] -> Builder GPU Result)
    -> BuilderT GPU (State VNameSource) [VName]
    tilingSegMap ::
          String ->
          ResultManifest ->
          (PrimExp VName -> [DimIndex SubExp] -> Builder GPU Result) ->
          Builder GPU [VName],
        -- The boolean PrimExp indicates whether they are in-bounds.

        Tiling
    -> TileKind
    -> PrivStms
    -> SubExp
    -> [InputArray]
    -> Builder GPU [InputTile]
    tilingReadTile ::
          TileKind ->
          PrivStms ->
          SubExp ->
          [InputArray] ->
          Builder GPU [InputTile],
        Tiling
    -> ProcessTileArgs -> BuilderT GPU (State VNameSource) [VName]
    tilingProcessTile ::
          ProcessTileArgs ->
          Builder GPU [VName],
        Tiling
    -> ResidualTileArgs -> BuilderT GPU (State VNameSource) [VName]
    tilingProcessResidualTile ::
          ResidualTileArgs ->
          Builder GPU [VName],
        Tiling -> VName -> BuilderT GPU (State VNameSource) KernelResult
    tilingTileReturns :: VName -> Builder GPU KernelResult,
        Tiling -> SegSpace
    tilingSpace :: SegSpace,
        Tiling -> Shape
    tilingTileShape :: Shape,
        Tiling -> SegLevel
    tilingLevel :: SegLevel,
        Tiling -> Builder GPU SubExp
    tilingNumWholeTiles :: Builder GPU SubExp
      }

    type DoTiling gtids kdims =
      gtids -> kdims -> SubExp -> Builder GPU Tiling

    protectOutOfBounds ::
      String ->
      PrimExp VName ->
      [Type] ->
      Builder GPU Result ->
      Builder GPU [VName]
    protectOutOfBounds :: [Char]
    -> PrimExp VName
    -> [Type]
    -> Builder GPU Result
    -> BuilderT GPU (State VNameSource) [VName]
    protectOutOfBounds [Char]
    desc PrimExp VName
    in_bounds [Type]
    ts Builder GPU Result
    m = do
      -- This is more complicated than you might expect, because we need
      -- to be able to produce a blank accumulator, which eBlank cannot
      -- do.  By the linear type rules of accumulators, the body returns
      -- an accumulator of type 'acc_t', then a unique variable of type
      -- 'acc_t' must also be free in the body.  This means we can find it
      -- based just on the type.
      m_body <- BuilderT
      GPU
      (State VNameSource)
      (Body (Rep (BuilderT GPU (State VNameSource))))
    -> BuilderT
         GPU
         (State VNameSource)
         (Body (Rep (BuilderT GPU (State VNameSource))))
    forall (m :: * -> *).
    MonadBuilder m =>
    m (Body (Rep m)) -> m (Body (Rep m))
    insertStmsM (BuilderT
       GPU
       (State VNameSource)
       (Body (Rep (BuilderT GPU (State VNameSource))))
     -> BuilderT
          GPU
          (State VNameSource)
          (Body (Rep (BuilderT GPU (State VNameSource)))))
    -> BuilderT
         GPU
         (State VNameSource)
         (Body (Rep (BuilderT GPU (State VNameSource))))
    -> BuilderT
         GPU
         (State VNameSource)
         (Body (Rep (BuilderT GPU (State VNameSource))))
    forall a b. (a -> b) -> a -> b
    $ Stms GPU -> Result -> Body GPU
    forall rep. Buildable rep => Stms rep -> Result -> Body rep
    mkBody Stms GPU
    forall a. Monoid a => a
    mempty (Result -> Body GPU)
    -> Builder GPU Result
    -> BuilderT GPU (State VNameSource) (Body GPU)
    forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
    <$> Builder GPU Result
    m
      let m_body_free = Names -> [VName]
    namesToList (Names -> [VName]) -> Names -> [VName]
    forall a b. (a -> b) -> a -> b
    $ Body GPU -> Names
    forall a. FreeIn a => a -> Names
    freeIn Body GPU
    m_body
      t_to_v <-
        filter (isAcc . fst)
          <$> (zip <$> mapM lookupType m_body_free <*> pure m_body_free)
      let blank Type
    t = BuilderT
      GPU
      (State VNameSource)
      (Exp (Rep (BuilderT GPU (State VNameSource))))
    -> (VName
        -> BuilderT
             GPU
             (State VNameSource)
             (Exp (Rep (BuilderT GPU (State VNameSource)))))
    -> Maybe VName
    -> BuilderT
         GPU
         (State VNameSource)
         (Exp (Rep (BuilderT GPU (State VNameSource))))
    forall b a. b -> (a -> b) -> Maybe a -> b
    maybe (Type
    -> BuilderT
         GPU
         (State VNameSource)
         (Exp (Rep (BuilderT GPU (State VNameSource))))
    forall (m :: * -> *). MonadBuilder m => Type -> m (Exp (Rep m))
    eBlank Type
    t) (Exp (Rep (BuilderT GPU (State VNameSource)))
    -> BuilderT
         GPU
         (State VNameSource)
         (Exp (Rep (BuilderT GPU (State VNameSource))))
    forall a. a -> BuilderT GPU (State VNameSource) a
    forall (f :: * -> *) a. Applicative f => a -> f a
    pure (Exp (Rep (BuilderT GPU (State VNameSource)))
     -> BuilderT
          GPU
          (State VNameSource)
          (Exp (Rep (BuilderT GPU (State VNameSource)))))
    -> (VName -> Exp (Rep (BuilderT GPU (State VNameSource))))
    -> VName
    -> BuilderT
         GPU
         (State VNameSource)
         (Exp (Rep (BuilderT GPU (State VNameSource))))
    forall b c a. (b -> c) -> (a -> b) -> a -> c
    . BasicOp -> Exp (Rep (BuilderT GPU (State VNameSource)))
    forall rep. BasicOp -> Exp rep
    BasicOp (BasicOp -> Exp (Rep (BuilderT GPU (State VNameSource))))
    -> (VName -> BasicOp)
    -> VName
    -> Exp (Rep (BuilderT GPU (State VNameSource)))
    forall b c a. (b -> c) -> (a -> b) -> a -> c
    . SubExp -> BasicOp
    SubExp (SubExp -> BasicOp) -> (VName -> SubExp) -> VName -> BasicOp
    forall b c a. (b -> c) -> (a -> b) -> a -> c
    . VName -> SubExp
    Var) (Maybe VName
     -> BuilderT
          GPU
          (State VNameSource)
          (Exp (Rep (BuilderT GPU (State VNameSource)))))
    -> Maybe VName
    -> BuilderT
         GPU
         (State VNameSource)
         (Exp (Rep (BuilderT GPU (State VNameSource))))
    forall a b. (a -> b) -> a -> b
    $ Type -> [(Type, VName)] -> Maybe VName
    forall a b. Eq a => a -> [(a, b)] -> Maybe b
    lookup Type
    t [(Type, VName)]
    t_to_v
      letTupExp desc =<< eIf (toExp in_bounds) (pure m_body) (eBody $ map blank ts)

    postludeGeneric ::
      Tiling ->
      PrivStms ->
      Pat Type ->
      [VName] ->
      Stms GPU ->
      Result ->
      [Type] ->
      Builder GPU [VName]
    postludeGeneric :: Tiling
    -> PrivStms
    -> Pat Type
    -> [VName]
    -> Stms GPU
    -> Result
    -> [Type]
    -> BuilderT GPU (State VNameSource) [VName]
    postludeGeneric Tiling
    tiling PrivStms
    privstms Pat Type
    pat [VName]
    accs' Stms GPU
    poststms Result
    poststms_res [Type]
    res_ts =
      Tiling
    -> [Char]
    -> ResultManifest
    -> (PrimExp VName -> [DimIndex SubExp] -> Builder GPU Result)
    -> BuilderT GPU (State VNameSource) [VName]
    tilingSegMap Tiling
    tiling [Char]
    "thread_res" ResultManifest
    ResultPrivate ((PrimExp VName -> [DimIndex SubExp] -> Builder GPU Result)
     -> BuilderT GPU (State VNameSource) [VName])
    -> (PrimExp VName -> [DimIndex SubExp] -> Builder GPU Result)
    -> BuilderT GPU (State VNameSource) [VName]
    forall a b. (a -> b) -> a -> b
    $ \PrimExp VName
    in_bounds [DimIndex SubExp]
    slice -> do
        -- Read our per-thread result from the tiled loop.
        [(VName, VName)]
    -> ((VName, VName) -> BuilderT GPU (State VNameSource) ())
    -> BuilderT GPU (State VNameSource) ()
    forall (t :: * -> *) (m :: * -> *) a b.
    (Foldable t, Monad m) =>
    t a -> (a -> m b) -> m ()
    forM_ ([VName] -> [VName] -> [(VName, VName)]
    forall a b. [a] -> [b] -> [(a, b)]
    zip (Pat Type -> [VName]
    forall dec. Pat dec -> [VName]
    patNames Pat Type
    pat) [VName]
    accs') (((VName, VName) -> BuilderT GPU (State VNameSource) ())
     -> BuilderT GPU (State VNameSource) ())
    -> ((VName, VName) -> BuilderT GPU (State VNameSource) ())
    -> BuilderT GPU (State VNameSource) ()
    forall a b. (a -> b) -> a -> b
    $ \(VName
    us, VName
    everyone) -> do
          everyone_t <- VName -> BuilderT GPU (State VNameSource) Type
    forall rep (m :: * -> *). HasScope rep m => VName -> m Type
    lookupType VName
    everyone
          letBindNames [us] $ BasicOp $ Index everyone $ fullSlice everyone_t slice

        if Stms GPU
    poststms Stms GPU -> Stms GPU -> Bool
    forall a. Eq a => a -> a -> Bool
    == Stms GPU
    forall a. Monoid a => a
    mempty
          then do
            -- The privstms may still be necessary for the result.
            [DimIndex SubExp]
    -> PrivStms -> BuilderT GPU (State VNameSource) ()
    addPrivStms [DimIndex SubExp]
    slice PrivStms
    privstms
            Result -> Builder GPU Result
    forall a. a -> BuilderT GPU (State VNameSource) a
    forall (f :: * -> *) a. Applicative f => a -> f a
    pure Result
    poststms_res
          else ([VName] -> Result)
    -> BuilderT GPU (State VNameSource) [VName] -> Builder GPU Result
    forall a b.
    (a -> b)
    -> BuilderT GPU (State VNameSource) a
    -> BuilderT GPU (State VNameSource) b
    forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
    fmap [VName] -> Result
    varsRes (BuilderT GPU (State VNameSource) [VName] -> Builder GPU Result)
    -> BuilderT GPU (State VNameSource) [VName] -> Builder GPU Result
    forall a b. (a -> b) -> a -> b
    $
            [Char]
    -> PrimExp VName
    -> [Type]
    -> Builder GPU Result
    -> BuilderT GPU (State VNameSource) [VName]
    protectOutOfBounds [Char]
    "postlude" PrimExp VName
    in_bounds [Type]
    res_ts (Builder GPU Result -> BuilderT GPU (State VNameSource) [VName])
    -> Builder GPU Result -> BuilderT GPU (State VNameSource) [VName]
    forall a b. (a -> b) -> a -> b
    $ do
              [DimIndex SubExp]
    -> PrivStms -> BuilderT GPU (State VNameSource) ()
    addPrivStms [DimIndex SubExp]
    slice PrivStms
    privstms
              Stms (Rep (BuilderT GPU (State VNameSource)))
    -> BuilderT GPU (State VNameSource) ()
    forall (m :: * -> *). MonadBuilder m => Stms (Rep m) -> m ()
    addStms Stms (Rep (BuilderT GPU (State VNameSource)))
    Stms GPU
    poststms
              Result -> Builder GPU Result
    forall a. a -> BuilderT GPU (State VNameSource) a
    forall (f :: * -> *) a. Applicative f => a -> f a
    pure Result
    poststms_res

    type TiledBody = Names -> PrivStms -> Builder GPU [VName]

    tileGeneric ::
      DoTiling gtids kdims ->
      [Type] ->
      Pat Type ->
      gtids ->
      kdims ->
      SubExp ->
      (Commutativity, Lambda GPU, [SubExp], Lambda GPU) ->
      [InputArray] ->
      Stms GPU ->
      Result ->
      TileM (Stms GPU, Tiling, TiledBody)
    tileGeneric :: forall gtids kdims.
    DoTiling gtids kdims
    -> [Type]
    -> Pat Type
    -> gtids
    -> kdims
    -> SubExp
    -> (Commutativity, Lambda GPU, [SubExp], Lambda GPU)
    -> [InputArray]
    -> Stms GPU
    -> Result
    -> ReaderT
         (Scope GPU) (State VNameSource) (Stms GPU, Tiling, TiledBody)
    tileGeneric DoTiling gtids kdims
    doTiling [Type]
    res_ts Pat Type
    pat gtids
    gtids kdims
    kdims SubExp
    w (Commutativity, Lambda GPU, [SubExp], Lambda GPU)
    form [InputArray]
    inputs Stms GPU
    poststms Result
    poststms_res = do
      (tiling, tiling_stms) <- Builder GPU Tiling
    -> ReaderT (Scope GPU) (State VNameSource) (Tiling, Stms GPU)
    forall (m :: * -> *) somerep rep a.
    (MonadFreshNames m, HasScope somerep m, SameScope somerep rep) =>
    Builder rep a -> m (a, Stms rep)
    runBuilder (Builder GPU Tiling
     -> ReaderT (Scope GPU) (State VNameSource) (Tiling, Stms GPU))
    -> Builder GPU Tiling
    -> ReaderT (Scope GPU) (State VNameSource) (Tiling, Stms GPU)
    forall a b. (a -> b) -> a -> b
    $ DoTiling gtids kdims
    doTiling gtids
    gtids kdims
    kdims SubExp
    w

      pure (tiling_stms, tiling, tiledBody tiling)
      where
        (Commutativity
    red_comm, Lambda GPU
    red_lam, [SubExp]
    red_nes, Lambda GPU
    map_lam) = (Commutativity, Lambda GPU, [SubExp], Lambda GPU)
    form

        tiledBody :: Tiling -> Names -> PrivStms -> Builder GPU [VName]
        tiledBody :: Tiling -> TiledBody
    tiledBody Tiling
    tiling Names
    _private PrivStms
    privstms = do
          let tile_shape :: Shape
    tile_shape = Tiling -> Shape
    tilingTileShape Tiling
    tiling

          num_whole_tiles <- Tiling -> Builder GPU SubExp
    tilingNumWholeTiles Tiling
    tiling

          -- We don't use a Replicate here, because we want to enforce a
          -- scalar memory space.
          mergeinits <- tilingSegMap tiling "mergeinit" ResultPrivate $ \PrimExp VName
    in_bounds [DimIndex SubExp]
    slice ->
            -- Constant neutral elements (a common case) do not need protection from OOB.
            if [SubExp] -> Names
    forall a. FreeIn a => a -> Names
    freeIn [SubExp]
    red_nes Names -> Names -> Bool
    forall a. Eq a => a -> a -> Bool
    == Names
    forall a. Monoid a => a
    mempty
              then Result -> Builder GPU Result
    forall a. a -> BuilderT GPU (State VNameSource) a
    forall (f :: * -> *) a. Applicative f => a -> f a
    pure (Result -> Builder GPU Result) -> Result -> Builder GPU Result
    forall a b. (a -> b) -> a -> b
    $ [SubExp] -> Result
    subExpsRes [SubExp]
    red_nes
              else ([VName] -> Result)
    -> BuilderT GPU (State VNameSource) [VName] -> Builder GPU Result
    forall a b.
    (a -> b)
    -> BuilderT GPU (State VNameSource) a
    -> BuilderT GPU (State VNameSource) b
    forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
    fmap [VName] -> Result
    varsRes (BuilderT GPU (State VNameSource) [VName] -> Builder GPU Result)
    -> BuilderT GPU (State VNameSource) [VName] -> Builder GPU Result
    forall a b. (a -> b) -> a -> b
    $
                [Char]
    -> PrimExp VName
    -> [Type]
    -> Builder GPU Result
    -> BuilderT GPU (State VNameSource) [VName]
    protectOutOfBounds [Char]
    "neutral" PrimExp VName
    in_bounds (Lambda GPU -> [Type]
    forall rep. Lambda rep -> [Type]
    lambdaReturnType Lambda GPU
    red_lam) (Builder GPU Result -> BuilderT GPU (State VNameSource) [VName])
    -> Builder GPU Result -> BuilderT GPU (State VNameSource) [VName]
    forall a b. (a -> b) -> a -> b
    $ do
                  [DimIndex SubExp]
    -> PrivStms -> BuilderT GPU (State VNameSource) ()
    addPrivStms [DimIndex SubExp]
    slice PrivStms
    privstms
                  Result -> Builder GPU Result
    forall a. a -> BuilderT GPU (State VNameSource) a
    forall (f :: * -> *) a. Applicative f => a -> f a
    pure (Result -> Builder GPU Result) -> Result -> Builder GPU Result
    forall a b. (a -> b) -> a -> b
    $ [SubExp] -> Result
    subExpsRes [SubExp]
    red_nes

          merge <- forM (zip (lambdaParams red_lam) mergeinits) $ \(Param Type
    p, VName
    mergeinit) ->
            (,)
              (Param (TypeBase Shape Uniqueness)
     -> SubExp -> (Param (TypeBase Shape Uniqueness), SubExp))
    -> BuilderT
         GPU (State VNameSource) (Param (TypeBase Shape Uniqueness))
    -> BuilderT
         GPU
         (State VNameSource)
         (SubExp -> (Param (TypeBase Shape Uniqueness), SubExp))
    forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
    <$> [Char]
    -> TypeBase Shape Uniqueness
    -> BuilderT
         GPU (State VNameSource) (Param (TypeBase Shape Uniqueness))
    forall (m :: * -> *) dec.
    MonadFreshNames m =>
    [Char] -> dec -> m (Param dec)
    newParam
                (VName -> [Char]
    baseString (Param Type -> VName
    forall dec. Param dec -> VName
    paramName Param Type
    p) [Char] -> [Char] -> [Char]
    forall a. [a] -> [a] -> [a]
    ++ [Char]
    "_merge")
                (Param Type -> Type
    forall dec. Typed dec => Param dec -> Type
    paramType Param Type
    p Type -> Shape -> Type
    `arrayOfShape` Shape
    tile_shape Type -> Uniqueness -> TypeBase Shape Uniqueness
    forall shape.
    TypeBase shape NoUniqueness
    -> Uniqueness -> TypeBase shape Uniqueness
    `toDecl` Uniqueness
    Unique)
              BuilderT
      GPU
      (State VNameSource)
      (SubExp -> (Param (TypeBase Shape Uniqueness), SubExp))
    -> Builder GPU SubExp
    -> BuilderT
         GPU (State VNameSource) (Param (TypeBase Shape Uniqueness), SubExp)
    forall a b.
    BuilderT GPU (State VNameSource) (a -> b)
    -> BuilderT GPU (State VNameSource) a
    -> BuilderT GPU (State VNameSource) b
    forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
    <*> SubExp -> Builder GPU SubExp
    forall a. a -> BuilderT GPU (State VNameSource) a
    forall (f :: * -> *) a. Applicative f => a -> f a
    pure (VName -> SubExp
    Var VName
    mergeinit)

          tile_id <- newVName "tile_id"
          let loopform = VName -> IntType -> SubExp -> LoopForm
    ForLoop VName
    tile_id IntType
    Int64 SubExp
    num_whole_tiles
          loopbody <- renameBody <=< runBodyBuilder $
            localScope (scopeOfLoopForm loopform <> scopeOfFParams (map fst merge)) $ do
              -- Collectively read a tile.
              tile <- tilingReadTile tiling TilePartial privstms (Var tile_id) inputs

              -- Now each thread performs a traversal of the tile and
              -- updates its accumulator.
              let accs =
                    ((Param (TypeBase Shape Uniqueness), SubExp) -> VName)
    -> [(Param (TypeBase Shape Uniqueness), SubExp)] -> [VName]
    forall a b. (a -> b) -> [a] -> [b]
    map (Param (TypeBase Shape Uniqueness) -> VName
    forall dec. Param dec -> VName
    paramName (Param (TypeBase Shape Uniqueness) -> VName)
    -> ((Param (TypeBase Shape Uniqueness), SubExp)
        -> Param (TypeBase Shape Uniqueness))
    -> (Param (TypeBase Shape Uniqueness), SubExp)
    -> VName
    forall b c a. (b -> c) -> (a -> b) -> a -> c
    . (Param (TypeBase Shape Uniqueness), SubExp)
    -> Param (TypeBase Shape Uniqueness)
    forall a b. (a, b) -> a
    fst) [(Param (TypeBase Shape Uniqueness), SubExp)]
    merge
                  tile_args =
                    PrivStms
    -> Commutativity
    -> Lambda GPU
    -> Lambda GPU
    -> [InputTile]
    -> [VName]
    -> SubExp
    -> ProcessTileArgs
    ProcessTileArgs PrivStms
    privstms Commutativity
    red_comm Lambda GPU
    red_lam Lambda GPU
    map_lam [InputTile]
    tile [VName]
    accs (VName -> SubExp
    Var VName
    tile_id)
              varsRes <$> tilingProcessTile tiling tile_args

          accs <- letTupExp "accs" $ Loop merge loopform loopbody

          -- We possibly have to traverse a residual tile.
          red_lam' <- renameLambda red_lam
          map_lam' <- renameLambda map_lam
          let residual_args =
                PrivStms
    -> Commutativity
    -> Lambda GPU
    -> Lambda GPU
    -> [InputArray]
    -> [VName]
    -> SubExp
    -> SubExp
    -> ResidualTileArgs
    ResidualTileArgs PrivStms
    privstms Commutativity
    red_comm Lambda GPU
    red_lam' Lambda GPU
    map_lam' [InputArray]
    inputs [VName]
    accs SubExp
    w SubExp
    num_whole_tiles
          accs' <- tilingProcessResidualTile tiling residual_args

          -- Create a SegMap that takes care of the postlude for every thread.
          postludeGeneric tiling privstms pat accs' poststms poststms_res res_ts

    mkReadPreludeValues :: [VName] -> [VName] -> ReadPrelude
    mkReadPreludeValues :: [VName] -> [VName] -> ReadPrelude
    mkReadPreludeValues [VName]
    prestms_live_arrs [VName]
    prestms_live [DimIndex SubExp]
    slice =
      ([()] -> ())
    -> BuilderT GPU (State VNameSource) [()]
    -> BuilderT GPU (State VNameSource) ()
    forall a b.
    (a -> b)
    -> BuilderT GPU (State VNameSource) a
    -> BuilderT GPU (State VNameSource) b
    forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
    fmap [()] -> ()
    forall a. Monoid a => [a] -> a
    mconcat (BuilderT GPU (State VNameSource) [()]
     -> BuilderT GPU (State VNameSource) ())
    -> (((VName, VName) -> BuilderT GPU (State VNameSource) ())
        -> BuilderT GPU (State VNameSource) [()])
    -> ((VName, VName) -> BuilderT GPU (State VNameSource) ())
    -> BuilderT GPU (State VNameSource) ()
    forall b c a. (b -> c) -> (a -> b) -> a -> c
    . [(VName, VName)]
    -> ((VName, VName) -> BuilderT GPU (State VNameSource) ())
    -> BuilderT GPU (State VNameSource) [()]
    forall (t :: * -> *) (m :: * -> *) a b.
    (Traversable t, Monad m) =>
    t a -> (a -> m b) -> m (t b)
    forM ([VName] -> [VName] -> [(VName, VName)]
    forall a b. [a] -> [b] -> [(a, b)]
    zip [VName]
    prestms_live_arrs [VName]
    prestms_live) (((VName, VName) -> BuilderT GPU (State VNameSource) ())
     -> BuilderT GPU (State VNameSource) ())
    -> ((VName, VName) -> BuilderT GPU (State VNameSource) ())
    -> BuilderT GPU (State VNameSource) ()
    forall a b. (a -> b) -> a -> b
    $ \(VName
    arr, VName
    v) -> do
        arr_t <- VName -> BuilderT GPU (State VNameSource) Type
    forall rep (m :: * -> *). HasScope rep m => VName -> m Type
    lookupType VName
    arr
        letBindNames [v] $ BasicOp $ Index arr $ fullSlice arr_t slice

    tileReturns :: [(VName, SubExp)] -> [(SubExp, SubExp)] -> VName -> Builder GPU KernelResult
    tileReturns :: [(VName, SubExp)]
    -> [(SubExp, SubExp)]
    -> VName
    -> BuilderT GPU (State VNameSource) KernelResult
    tileReturns [(VName, SubExp)]
    dims_on_top [(SubExp, SubExp)]
    dims VName
    arr = do
      let unit_dims :: [SubExp]
    unit_dims = Int -> SubExp -> [SubExp]
    forall a. Int -> a -> [a]
    replicate ([(VName, SubExp)] -> Int
    forall a. [a] -> Int
    forall (t :: * -> *) a. Foldable t => t a -> Int
    length [(VName, SubExp)]
    dims_on_top) (IntType -> Integer -> SubExp
    intConst IntType
    Int64 Integer
    1)
      arr_t <- VName -> BuilderT GPU (State VNameSource) Type
    forall rep (m :: * -> *). HasScope rep m => VName -> m Type
    lookupType VName
    arr
      arr' <-
        if null dims_on_top || null (arrayDims arr_t) -- Second check is for accumulators.
          then pure arr
          else do
            let new_shape = [SubExp] -> Shape
    forall d. [d] -> ShapeBase d
    Shape ([SubExp] -> Shape) -> [SubExp] -> Shape
    forall a b. (a -> b) -> a -> b
    $ [SubExp]
    unit_dims [SubExp] -> [SubExp] -> [SubExp]
    forall a. [a] -> [a] -> [a]
    ++ Type -> [SubExp]
    forall u. TypeBase Shape u -> [SubExp]
    arrayDims Type
    arr_t
            letExp (baseString arr) . BasicOp $
              Reshape arr (reshapeAll (arrayShape arr_t) new_shape)
      let tile_dims = [SubExp] -> [SubExp] -> [(SubExp, SubExp)]
    forall a b. [a] -> [b] -> [(a, b)]
    zip (((VName, SubExp) -> SubExp) -> [(VName, SubExp)] -> [SubExp]
    forall a b. (a -> b) -> [a] -> [b]
    map (VName, SubExp) -> SubExp
    forall a b. (a, b) -> b
    snd [(VName, SubExp)]
    dims_on_top) [SubExp]
    unit_dims [(SubExp, SubExp)] -> [(SubExp, SubExp)] -> [(SubExp, SubExp)]
    forall a. [a] -> [a] -> [a]
    ++ [(SubExp, SubExp)]
    dims
      pure $ TileReturns mempty tile_dims arr'

    is1DTileable :: VName -> M.Map VName Names -> VName -> InputArray
    is1DTileable :: VName -> AliasTable -> VName -> InputArray
    is1DTileable VName
    gtid AliasTable
    variance VName
    arr
      | Bool -> Bool
    not (Bool -> Bool) -> Bool -> Bool
    forall a b. (a -> b) -> a -> b
    $ VName -> Names -> Bool
    nameIn VName
    gtid (Names -> Bool) -> Names -> Bool
    forall a b. (a -> b) -> a -> b
    $ Names -> VName -> AliasTable -> Names
    forall k a. Ord k => a -> k -> Map k a -> a
    M.findWithDefault Names
    forall a. Monoid a => a
    mempty VName
    arr AliasTable
    variance =
          [Int] -> VName -> InputArray
    InputTile [Int
    0] VName
    arr
      | Bool
    otherwise =
          VName -> InputArray
    InputDontTile VName
    arr

    reconstructGtids1D ::
      Count BlockSize SubExp ->
      VName ->
      VName ->
      VName ->
      Builder GPU ()
    reconstructGtids1D :: Count BlockSize SubExp
    -> VName -> VName -> VName -> BuilderT GPU (State VNameSource) ()
    reconstructGtids1D Count BlockSize SubExp
    tblock_size VName
    gtid VName
    gid VName
    ltid =
      [VName]
    -> Exp (Rep (BuilderT GPU (State VNameSource)))
    -> BuilderT GPU (State VNameSource) ()
    forall (m :: * -> *).
    MonadBuilder m =>
    [VName] -> Exp (Rep m) -> m ()
    letBindNames [VName
    gtid]
        (Exp GPU -> BuilderT GPU (State VNameSource) ())
    -> BuilderT GPU (State VNameSource) (Exp GPU)
    -> BuilderT GPU (State VNameSource) ()
    forall (m :: * -> *) a b. Monad m => (a -> m b) -> m a -> m b
    =<< TPrimExp Int64 VName
    -> BuilderT
         GPU
         (State VNameSource)
         (Exp (Rep (BuilderT GPU (State VNameSource))))
    forall a (m :: * -> *).
    (ToExp a, MonadBuilder m) =>
    a -> m (Exp (Rep m))
    forall (m :: * -> *).
    MonadBuilder m =>
    TPrimExp Int64 VName -> m (Exp (Rep m))
    toExp (VName -> TPrimExp Int64 VName
    forall a. a -> TPrimExp Int64 a
    le64 VName
    gid TPrimExp Int64 VName
    -> TPrimExp Int64 VName -> TPrimExp Int64 VName
    forall a. Num a => a -> a -> a
    * SubExp -> TPrimExp Int64 VName
    pe64 (Count BlockSize SubExp -> SubExp
    forall {k} (u :: k) e. Count u e -> e
    unCount Count BlockSize SubExp
    tblock_size) TPrimExp Int64 VName
    -> TPrimExp Int64 VName -> TPrimExp Int64 VName
    forall a. Num a => a -> a -> a
    + VName -> TPrimExp Int64 VName
    forall a. a -> TPrimExp Int64 a
    le64 VName
    ltid)

    readTile1D ::
      SubExp ->
      VName ->
      VName ->
      KernelGrid ->
      TileKind ->
      PrivStms ->
      SubExp ->
      [InputArray] ->
      Builder GPU [InputTile]
    readTile1D :: SubExp
    -> VName
    -> VName
    -> KernelGrid
    -> TileKind
    -> PrivStms
    -> SubExp
    -> [InputArray]
    -> Builder GPU [InputTile]
    readTile1D SubExp
    tile_size VName
    gid VName
    gtid (KernelGrid Count NumBlocks SubExp
    _num_tblocks Count BlockSize SubExp
    tblock_size) TileKind
    kind PrivStms
    privstms SubExp
    tile_id [InputArray]
    inputs =
      ([VName] -> [InputTile])
    -> BuilderT GPU (State VNameSource) [VName]
    -> Builder GPU [InputTile]
    forall a b.
    (a -> b)
    -> BuilderT GPU (State VNameSource) a
    -> BuilderT GPU (State VNameSource) b
    forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
    fmap ([InputArray] -> [VName] -> [InputTile]
    inputsToTiles [InputArray]
    inputs)
        (BuilderT GPU (State VNameSource) [VName]
     -> Builder GPU [InputTile])
    -> ((VName -> Builder GPU Result)
        -> BuilderT GPU (State VNameSource) [VName])
    -> (VName -> Builder GPU Result)
    -> Builder GPU [InputTile]
    forall b c a. (b -> c) -> (a -> b) -> a -> c
    . [Char]
    -> SegLevel
    -> ResultManifest
    -> SubExp
    -> (VName -> Builder GPU Result)
    -> BuilderT GPU (State VNameSource) [VName]
    segMap1D [Char]
    "full_tile" SegLevel
    lvl ResultManifest
    ResultNoSimplify SubExp
    tile_size
        ((VName -> Builder GPU Result) -> Builder GPU [InputTile])
    -> (VName -> Builder GPU Result) -> Builder GPU [InputTile]
    forall a b. (a -> b) -> a -> b
    $ \VName
    ltid -> do
          j <-
            [Char]
    -> Exp (Rep (BuilderT GPU (State VNameSource)))
    -> Builder GPU SubExp
    forall (m :: * -> *).
    MonadBuilder m =>
    [Char] -> Exp (Rep m) -> m SubExp
    letSubExp [Char]
    "j"
              (Exp GPU -> Builder GPU SubExp)
    -> BuilderT GPU (State VNameSource) (Exp GPU) -> Builder GPU SubExp
    forall (m :: * -> *) a b. Monad m => (a -> m b) -> m a -> m b
    =<< TPrimExp Int64 VName
    -> BuilderT
         GPU
         (State VNameSource)
         (Exp (Rep (BuilderT GPU (State VNameSource))))
    forall a (m :: * -> *).
    (ToExp a, MonadBuilder m) =>
    a -> m (Exp (Rep m))
    forall (m :: * -> *).
    MonadBuilder m =>
    TPrimExp Int64 VName -> m (Exp (Rep m))
    toExp (SubExp -> TPrimExp Int64 VName
    pe64 SubExp
    tile_id TPrimExp Int64 VName
    -> TPrimExp Int64 VName -> TPrimExp Int64 VName
    forall a. Num a => a -> a -> a
    * SubExp -> TPrimExp Int64 VName
    pe64 SubExp
    tile_size TPrimExp Int64 VName
    -> TPrimExp Int64 VName -> TPrimExp Int64 VName
    forall a. Num a => a -> a -> a
    + VName -> TPrimExp Int64 VName
    forall a. a -> TPrimExp Int64 a
    le64 VName
    ltid)

          reconstructGtids1D tblock_size gtid gid ltid
          addPrivStms [DimFix $ Var ltid] privstms

          let arrs = ((VName, [Int]) -> VName) -> [(VName, [Int])] -> [VName]
    forall a b. (a -> b) -> [a] -> [b]
    map (VName, [Int]) -> VName
    forall a b. (a, b) -> a
    fst ([(VName, [Int])] -> [VName]) -> [(VName, [Int])] -> [VName]
    forall a b. (a -> b) -> a -> b
    $ [InputArray] -> [(VName, [Int])]
    tiledInputs [InputArray]
    inputs
          arr_ts <- mapM lookupType arrs
          let tile_ts = (Type -> Type) -> [Type] -> [Type]
    forall a b. (a -> b) -> [a] -> [b]
    map Type -> Type
    forall u. TypeBase Shape u -> TypeBase Shape u
    rowType [Type]
    arr_ts
              w = Int -> [Type] -> SubExp
    forall u. Int -> [TypeBase Shape u] -> SubExp
    arraysSize Int
    0 [Type]
    arr_ts

          let readTileElem VName
    arr =
                -- No need for fullSlice because we are tiling only prims.
                [Char]
    -> Exp (Rep (BuilderT GPU (State VNameSource)))
    -> BuilderT GPU (State VNameSource) VName
    forall (m :: * -> *).
    MonadBuilder m =>
    [Char] -> Exp (Rep m) -> m VName
    letExp [Char]
    "tile_elem" (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
    $ VName -> Slice SubExp -> BasicOp
    Index VName
    arr (Slice SubExp -> BasicOp) -> Slice SubExp -> BasicOp
    forall a b. (a -> b) -> a -> b
    $ [DimIndex SubExp] -> Slice SubExp
    forall d. [DimIndex d] -> Slice d
    Slice [SubExp -> DimIndex SubExp
    forall d. d -> DimIndex d
    DimFix SubExp
    j])
          fmap varsRes $
            case kind of
              TileKind
    TilePartial ->
                [Char]
    -> Exp (Rep (BuilderT GPU (State VNameSource)))
    -> BuilderT GPU (State VNameSource) [VName]
    forall (m :: * -> *).
    MonadBuilder m =>
    [Char] -> Exp (Rep m) -> m [VName]
    letTupExp [Char]
    "pre1d"
                  (Exp GPU -> BuilderT GPU (State VNameSource) [VName])
    -> BuilderT GPU (State VNameSource) (Exp GPU)
    -> BuilderT GPU (State VNameSource) [VName]
    forall (m :: * -> *) a b. Monad m => (a -> m b) -> m a -> m b
    =<< BuilderT
      GPU
      (State VNameSource)
      (Exp (Rep (BuilderT GPU (State VNameSource))))
    -> BuilderT
         GPU
         (State VNameSource)
         (Body (Rep (BuilderT GPU (State VNameSource))))
    -> BuilderT
         GPU
         (State VNameSource)
         (Body (Rep (BuilderT GPU (State VNameSource))))
    -> BuilderT
         GPU
         (State VNameSource)
         (Exp (Rep (BuilderT GPU (State VNameSource))))
    forall (m :: * -> *).
    (MonadBuilder m, BranchType (Rep m) ~ ExtType) =>
    m (Exp (Rep m))
    -> m (Body (Rep m)) -> m (Body (Rep m)) -> m (Exp (Rep m))
    eIf
                    (TPrimExp Bool VName
    -> BuilderT
         GPU
         (State VNameSource)
         (Exp (Rep (BuilderT GPU (State VNameSource))))
    forall a (m :: * -> *).
    (ToExp a, MonadBuilder m) =>
    a -> m (Exp (Rep m))
    forall (m :: * -> *).
    MonadBuilder m =>
    TPrimExp Bool VName -> m (Exp (Rep m))
    toExp (TPrimExp Bool VName
     -> BuilderT
          GPU
          (State VNameSource)
          (Exp (Rep (BuilderT GPU (State VNameSource)))))
    -> TPrimExp Bool VName
    -> BuilderT
         GPU
         (State VNameSource)
         (Exp (Rep (BuilderT GPU (State VNameSource))))
    forall a b. (a -> b) -> a -> b
    $ SubExp -> TPrimExp Int64 VName
    pe64 SubExp
    j TPrimExp Int64 VName -> TPrimExp Int64 VName -> TPrimExp Bool VName
    forall {k} v (t :: k).
    Eq v =>
    TPrimExp t v -> TPrimExp t v -> TPrimExp Bool v
    .<. SubExp -> TPrimExp Int64 VName
    pe64 SubExp
    w)
                    ([SubExp] -> Body GPU
    forall rep. Buildable rep => [SubExp] -> Body rep
    resultBody ([SubExp] -> Body GPU)
    -> BuilderT GPU (State VNameSource) [SubExp]
    -> BuilderT GPU (State VNameSource) (Body GPU)
    forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
    <$> (VName -> Builder GPU SubExp)
    -> [VName] -> BuilderT GPU (State VNameSource) [SubExp]
    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 ((VName -> SubExp)
    -> BuilderT GPU (State VNameSource) VName -> Builder GPU SubExp
    forall a b.
    (a -> b)
    -> BuilderT GPU (State VNameSource) a
    -> BuilderT GPU (State VNameSource) b
    forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
    fmap VName -> SubExp
    Var (BuilderT GPU (State VNameSource) VName -> Builder GPU SubExp)
    -> (VName -> BuilderT GPU (State VNameSource) VName)
    -> VName
    -> Builder GPU SubExp
    forall b c a. (b -> c) -> (a -> b) -> a -> c
    . VName -> BuilderT GPU (State VNameSource) VName
    readTileElem) [VName]
    arrs)
                    ([BuilderT
       GPU
       (State VNameSource)
       (Exp (Rep (BuilderT GPU (State VNameSource))))]
    -> BuilderT
         GPU
         (State VNameSource)
         (Body (Rep (BuilderT GPU (State VNameSource))))
    forall (m :: * -> *).
    MonadBuilder m =>
    [m (Exp (Rep m))] -> m (Body (Rep m))
    eBody ([BuilderT
        GPU
        (State VNameSource)
        (Exp (Rep (BuilderT GPU (State VNameSource))))]
     -> BuilderT
          GPU
          (State VNameSource)
          (Body (Rep (BuilderT GPU (State VNameSource)))))
    -> [BuilderT
          GPU
          (State VNameSource)
          (Exp (Rep (BuilderT GPU (State VNameSource))))]
    -> BuilderT
         GPU
         (State VNameSource)
         (Body (Rep (BuilderT GPU (State VNameSource))))
    forall a b. (a -> b) -> a -> b
    $ (Type -> BuilderT GPU (State VNameSource) (Exp GPU))
    -> [Type] -> [BuilderT GPU (State VNameSource) (Exp GPU)]
    forall a b. (a -> b) -> [a] -> [b]
    map Type
    -> BuilderT
         GPU
         (State VNameSource)
         (Exp (Rep (BuilderT GPU (State VNameSource))))
    Type -> BuilderT GPU (State VNameSource) (Exp GPU)
    forall (m :: * -> *). MonadBuilder m => Type -> m (Exp (Rep m))
    eBlank [Type]
    tile_ts)
              TileKind
    TileFull ->
                (VName -> BuilderT GPU (State VNameSource) VName)
    -> [VName] -> BuilderT GPU (State VNameSource) [VName]
    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 VName -> BuilderT GPU (State VNameSource) VName
    readTileElem [VName]
    arrs
      where
        lvl :: SegLevel
    lvl = SegVirt -> SegLevel
    SegThreadInBlock SegVirt
    SegNoVirt

    processTile1D ::
      VName ->
      VName ->
      SubExp ->
      SubExp ->
      KernelGrid ->
      ProcessTileArgs ->
      Builder GPU [VName]
    processTile1D :: VName
    -> VName
    -> SubExp
    -> SubExp
    -> KernelGrid
    -> ProcessTileArgs
    -> BuilderT GPU (State VNameSource) [VName]
    processTile1D VName
    gid VName
    gtid SubExp
    kdim SubExp
    tile_size (KernelGrid Count NumBlocks SubExp
    _num_tblocks Count BlockSize SubExp
    tblock_size) ProcessTileArgs
    tile_args = do
      let red_comm :: Commutativity
    red_comm = ProcessTileArgs -> Commutativity
    processComm ProcessTileArgs
    tile_args
          privstms :: PrivStms
    privstms = ProcessTileArgs -> PrivStms
    processPrivStms ProcessTileArgs
    tile_args
          map_lam :: Lambda GPU
    map_lam = ProcessTileArgs -> Lambda GPU
    processMapLam ProcessTileArgs
    tile_args
          red_lam :: Lambda GPU
    red_lam = ProcessTileArgs -> Lambda GPU
    processRedLam ProcessTileArgs
    tile_args
          tiles :: [InputTile]
    tiles = ProcessTileArgs -> [InputTile]
    processTiles ProcessTileArgs
    tile_args
          tile_id :: SubExp
    tile_id = ProcessTileArgs -> SubExp
    processTileId ProcessTileArgs
    tile_args
          accs :: [VName]
    accs = ProcessTileArgs -> [VName]
    processAcc ProcessTileArgs
    tile_args

      [Char]
    -> SegLevel
    -> ResultManifest
    -> SubExp
    -> (VName -> Builder GPU Result)
    -> BuilderT GPU (State VNameSource) [VName]
    segMap1D [Char]
    "acc" SegLevel
    lvl ResultManifest
    ResultPrivate (Count BlockSize SubExp -> SubExp
    forall {k} (u :: k) e. Count u e -> e
    unCount Count BlockSize SubExp
    tblock_size) ((VName -> Builder GPU Result)
     -> BuilderT GPU (State VNameSource) [VName])
    -> (VName -> Builder GPU Result)
    -> BuilderT GPU (State VNameSource) [VName]
    forall a b. (a -> b) -> a -> b
    $ \VName
    ltid -> do
        Count BlockSize SubExp
    -> VName -> VName -> VName -> BuilderT GPU (State VNameSource) ()
    reconstructGtids1D Count BlockSize SubExp
    tblock_size VName
    gtid VName
    gid VName
    ltid
        [DimIndex SubExp]
    -> PrivStms -> BuilderT GPU (State VNameSource) ()
    addPrivStms [SubExp -> DimIndex SubExp
    forall d. d -> DimIndex d
    DimFix (SubExp -> DimIndex SubExp) -> SubExp -> DimIndex SubExp
    forall a b. (a -> b) -> a -> b
    $ VName -> SubExp
    Var VName
    ltid] PrivStms
    privstms

        -- We replace the neutral elements with the accumulators (this is
        -- OK because the parallel semantics are not used after this
        -- point).
        thread_accs <- [VName]
    -> (VName -> Builder GPU SubExp)
    -> BuilderT GPU (State VNameSource) [SubExp]
    forall (t :: * -> *) (m :: * -> *) a b.
    (Traversable t, Monad m) =>
    t a -> (a -> m b) -> m (t b)
    forM [VName]
    accs ((VName -> Builder GPU SubExp)
     -> BuilderT GPU (State VNameSource) [SubExp])
    -> (VName -> Builder GPU SubExp)
    -> BuilderT GPU (State VNameSource) [SubExp]
    forall a b. (a -> b) -> a -> b
    $ \VName
    acc ->
          [Char]
    -> Exp (Rep (BuilderT GPU (State VNameSource)))
    -> Builder GPU SubExp
    forall (m :: * -> *).
    MonadBuilder m =>
    [Char] -> Exp (Rep m) -> m SubExp
    letSubExp [Char]
    "acc" (Exp (Rep (BuilderT GPU (State VNameSource)))
     -> Builder GPU SubExp)
    -> Exp (Rep (BuilderT GPU (State VNameSource)))
    -> Builder GPU SubExp
    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
    $ VName -> Slice SubExp -> BasicOp
    Index VName
    acc (Slice SubExp -> BasicOp) -> Slice SubExp -> BasicOp
    forall a b. (a -> b) -> a -> b
    $ [DimIndex SubExp] -> Slice SubExp
    forall d. [DimIndex d] -> Slice d
    Slice [SubExp -> DimIndex SubExp
    forall d. d -> DimIndex d
    DimFix (SubExp -> DimIndex SubExp) -> SubExp -> DimIndex SubExp
    forall a b. (a -> b) -> a -> b
    $ VName -> SubExp
    Var VName
    ltid]
        let sliceTile (InputTiled [Int]
    _ VName
    arr) =
              VName -> BuilderT GPU (State VNameSource) VName
    forall a. a -> BuilderT GPU (State VNameSource) a
    forall (f :: * -> *) a. Applicative f => a -> f a
    pure VName
    arr
            sliceTile (InputUntiled VName
    arr) =
              VName
    -> SubExp
    -> SubExp
    -> SubExp
    -> BuilderT GPU (State VNameSource) VName
    forall (m :: * -> *).
    MonadBuilder m =>
    VName -> SubExp -> SubExp -> SubExp -> m VName
    sliceUntiled VName
    arr SubExp
    tile_id SubExp
    tile_size SubExp
    tile_size

        tiles' <- mapM sliceTile tiles

        let form' = [Reduce GPU] -> Lambda GPU -> ScremaForm GPU
    forall rep. [Reduce rep] -> Lambda rep -> ScremaForm rep
    redomapSOAC [Commutativity -> Lambda GPU -> [SubExp] -> Reduce GPU
    forall rep. Commutativity -> Lambda rep -> [SubExp] -> Reduce rep
    Reduce Commutativity
    red_comm Lambda GPU
    red_lam [SubExp]
    thread_accs] Lambda GPU
    map_lam
        fmap varsRes $
          letTupExp "acc"
            =<< eIf
              (toExp $ le64 gtid .<. pe64 kdim)
              (eBody [pure $ Op $ OtherOp $ Screma tile_size tiles' form'])
              (resultBodyM thread_accs)
      where
        lvl :: SegLevel
    lvl = SegVirt -> SegLevel
    SegThreadInBlock SegVirt
    SegNoVirt

    processResidualTile1D ::
      VName ->
      VName ->
      SubExp ->
      SubExp ->
      KernelGrid ->
      ResidualTileArgs ->
      Builder GPU [VName]
    processResidualTile1D :: VName
    -> VName
    -> SubExp
    -> SubExp
    -> KernelGrid
    -> ResidualTileArgs
    -> BuilderT GPU (State VNameSource) [VName]
    processResidualTile1D VName
    gid VName
    gtid SubExp
    kdim SubExp
    tile_size KernelGrid
    grid ResidualTileArgs
    args = do
      -- The number of residual elements that are not covered by
      -- the whole tiles.
      residual_input <-
        [Char]
    -> Exp (Rep (BuilderT GPU (State VNameSource)))
    -> Builder GPU SubExp
    forall (m :: * -> *).
    MonadBuilder m =>
    [Char] -> Exp (Rep m) -> m SubExp
    letSubExp [Char]
    "residual_input" (Exp (Rep (BuilderT GPU (State VNameSource)))
     -> Builder GPU SubExp)
    -> Exp (Rep (BuilderT GPU (State VNameSource)))
    -> Builder GPU SubExp
    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
    $
            BinOp -> SubExp -> SubExp -> BasicOp
    BinOp (IntType -> Safety -> BinOp
    SRem IntType
    Int64 Safety
    Unsafe) SubExp
    w SubExp
    tile_size

      letTupExp "acc_after_residual"
        =<< eIf
          (toExp $ pe64 residual_input .==. 0)
          (resultBodyM $ map Var accs)
          (nonemptyTile residual_input)
      where
        red_comm :: Commutativity
    red_comm = ResidualTileArgs -> Commutativity
    residualComm ResidualTileArgs
    args
        map_lam :: Lambda GPU
    map_lam = ResidualTileArgs -> Lambda GPU
    residualMapLam ResidualTileArgs
    args
        red_lam :: Lambda GPU
    red_lam = ResidualTileArgs -> Lambda GPU
    residualRedLam ResidualTileArgs
    args
        privstms :: PrivStms
    privstms = ResidualTileArgs -> PrivStms
    residualPrivStms ResidualTileArgs
    args
        inputs :: [InputArray]
    inputs = ResidualTileArgs -> [InputArray]
    residualInput ResidualTileArgs
    args
        accs :: [VName]
    accs = ResidualTileArgs -> [VName]
    residualAcc ResidualTileArgs
    args
        num_whole_tiles :: SubExp
    num_whole_tiles = ResidualTileArgs -> SubExp
    residualNumWholeTiles ResidualTileArgs
    args
        w :: SubExp
    w = ResidualTileArgs -> SubExp
    residualInputSize ResidualTileArgs
    args

        nonemptyTile :: SubExp -> BuilderT GPU (State VNameSource) (Body GPU)
    nonemptyTile SubExp
    residual_input = Builder GPU Result -> BuilderT GPU (State VNameSource) (Body GPU)
    forall rep (m :: * -> *) somerep.
    (Buildable rep, MonadFreshNames m, HasScope somerep m,
     SameScope somerep rep) =>
    Builder rep Result -> m (Body rep)
    runBodyBuilder (Builder GPU Result -> BuilderT GPU (State VNameSource) (Body GPU))
    -> Builder GPU Result
    -> BuilderT GPU (State VNameSource) (Body GPU)
    forall a b. (a -> b) -> a -> b
    $ do
          -- Collectively construct a tile.  Threads that are out-of-bounds
          -- provide a blank dummy value.
          full_tiles <-
            SubExp
    -> VName
    -> VName
    -> KernelGrid
    -> TileKind
    -> PrivStms
    -> SubExp
    -> [InputArray]
    -> Builder GPU [InputTile]
    readTile1D
              SubExp
    tile_size
              VName
    gid
              VName
    gtid
              KernelGrid
    grid
              TileKind
    TilePartial
              PrivStms
    privstms
              SubExp
    num_whole_tiles
              [InputArray]
    inputs

          let sliceTile (InputUntiled VName
    arr) =
                InputTile -> BuilderT GPU (State VNameSource) InputTile
    forall a. a -> BuilderT GPU (State VNameSource) a
    forall (f :: * -> *) a. Applicative f => a -> f a
    pure (InputTile -> BuilderT GPU (State VNameSource) InputTile)
    -> InputTile -> BuilderT GPU (State VNameSource) InputTile
    forall a b. (a -> b) -> a -> b
    $ VName -> InputTile
    InputUntiled VName
    arr
              sliceTile (InputTiled [Int]
    perm VName
    tile) = do
                let slice :: DimIndex SubExp
    slice =
                      SubExp -> SubExp -> SubExp -> DimIndex SubExp
    forall d. d -> d -> d -> DimIndex d
    DimSlice (IntType -> Integer -> SubExp
    intConst IntType
    Int64 Integer
    0) SubExp
    residual_input (IntType -> Integer -> SubExp
    intConst IntType
    Int64 Integer
    1)
                [Int] -> VName -> InputTile
    InputTiled [Int]
    perm
                  (VName -> InputTile)
    -> BuilderT GPU (State VNameSource) VName
    -> BuilderT GPU (State VNameSource) InputTile
    forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
    <$> [Char]
    -> Exp (Rep (BuilderT GPU (State VNameSource)))
    -> BuilderT GPU (State VNameSource) VName
    forall (m :: * -> *).
    MonadBuilder m =>
    [Char] -> Exp (Rep m) -> m VName
    letExp [Char]
    "partial_tile" (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
    $ VName -> Slice SubExp -> BasicOp
    Index VName
    tile (Slice SubExp -> BasicOp) -> Slice SubExp -> BasicOp
    forall a b. (a -> b) -> a -> b
    $ [DimIndex SubExp] -> Slice SubExp
    forall d. [DimIndex d] -> Slice d
    Slice [DimIndex SubExp
    slice])

          tiles <- mapM sliceTile full_tiles

          -- Now each thread performs a traversal of the tile and
          -- updates its accumulator.
          let tile_args =
                PrivStms
    -> Commutativity
    -> Lambda GPU
    -> Lambda GPU
    -> [InputTile]
    -> [VName]
    -> SubExp
    -> ProcessTileArgs
    ProcessTileArgs PrivStms
    privstms Commutativity
    red_comm Lambda GPU
    red_lam Lambda GPU
    map_lam [InputTile]
    tiles [VName]
    accs SubExp
    num_whole_tiles
          varsRes <$> processTile1D gid gtid kdim residual_input grid tile_args

    tiling1d :: [(VName, SubExp)] -> DoTiling VName SubExp
    tiling1d :: [(VName, SubExp)] -> DoTiling VName SubExp
    tiling1d [(VName, SubExp)]
    dims_on_top VName
    gtid SubExp
    kdim SubExp
    w = do
      gid <- [Char] -> BuilderT GPU (State VNameSource) VName
    forall (m :: * -> *). MonadFreshNames m => [Char] -> m VName
    newVName [Char]
    "gid"
      gid_flat <- newVName "gid_flat"

      tile_size_key <- nameFromString . prettyString <$> newVName "tile_size"
      tile_size <- letSubExp "tile_size" $ Op $ SizeOp $ GetSize tile_size_key SizeThreadBlock
      let tblock_size = SubExp
    tile_size

      (grid, space) <- do
        -- How many blocks we need to exhaust the innermost dimension.
        ldim <-
          letSubExp "ldim" . BasicOp $
            BinOp (SDivUp Int64 Unsafe) kdim tblock_size

        num_tblocks <-
          letSubExp "computed_num_tblocks"
            =<< foldBinOp (Mul Int64 OverflowUndef) ldim (map snd dims_on_top)

        pure
          ( KernelGrid (Count num_tblocks) (Count tblock_size),
            SegSpace gid_flat $ dims_on_top ++ [(gid, ldim)]
          )
      let tiling_lvl = SegVirt -> SegLevel
    SegThreadInBlock SegVirt
    SegNoVirt

      pure
        Tiling
          { tilingSegMap = \[Char]
    desc ResultManifest
    manifest PrimExp VName -> [DimIndex SubExp] -> Builder GPU Result
    f -> [Char]
    -> SegLevel
    -> ResultManifest
    -> SubExp
    -> (VName -> Builder GPU Result)
    -> BuilderT GPU (State VNameSource) [VName]
    segMap1D [Char]
    desc SegLevel
    tiling_lvl ResultManifest
    manifest SubExp
    tile_size ((VName -> Builder GPU Result)
     -> BuilderT GPU (State VNameSource) [VName])
    -> (VName -> Builder GPU Result)
    -> BuilderT GPU (State VNameSource) [VName]
    forall a b. (a -> b) -> a -> b
    $ \VName
    ltid -> do
              [VName]
    -> Exp (Rep (BuilderT GPU (State VNameSource)))
    -> BuilderT GPU (State VNameSource) ()
    forall (m :: * -> *).
    MonadBuilder m =>
    [VName] -> Exp (Rep m) -> m ()
    letBindNames [VName
    gtid]
                (Exp GPU -> BuilderT GPU (State VNameSource) ())
    -> BuilderT GPU (State VNameSource) (Exp GPU)
    -> BuilderT GPU (State VNameSource) ()
    forall (m :: * -> *) a b. Monad m => (a -> m b) -> m a -> m b
    =<< TPrimExp Int64 VName
    -> BuilderT
         GPU
         (State VNameSource)
         (Exp (Rep (BuilderT GPU (State VNameSource))))
    forall a (m :: * -> *).
    (ToExp a, MonadBuilder m) =>
    a -> m (Exp (Rep m))
    forall (m :: * -> *).
    MonadBuilder m =>
    TPrimExp Int64 VName -> m (Exp (Rep m))
    toExp (VName -> TPrimExp Int64 VName
    forall a. a -> TPrimExp Int64 a
    le64 VName
    gid TPrimExp Int64 VName
    -> TPrimExp Int64 VName -> TPrimExp Int64 VName
    forall a. Num a => a -> a -> a
    * SubExp -> TPrimExp Int64 VName
    pe64 SubExp
    tile_size TPrimExp Int64 VName
    -> TPrimExp Int64 VName -> TPrimExp Int64 VName
    forall a. Num a => a -> a -> a
    + VName -> TPrimExp Int64 VName
    forall a. a -> TPrimExp Int64 a
    le64 VName
    ltid)
              PrimExp VName -> [DimIndex SubExp] -> Builder GPU Result
    f (TPrimExp Bool VName -> PrimExp VName
    forall {k} (t :: k) v. TPrimExp t v -> PrimExp v
    untyped (TPrimExp Bool VName -> PrimExp VName)
    -> TPrimExp Bool VName -> PrimExp VName
    forall a b. (a -> b) -> a -> b
    $ VName -> TPrimExp Int64 VName
    forall a. a -> TPrimExp Int64 a
    le64 VName
    gtid TPrimExp Int64 VName -> TPrimExp Int64 VName -> TPrimExp Bool VName
    forall {k} v (t :: k).
    Eq v =>
    TPrimExp t v -> TPrimExp t v -> TPrimExp Bool v
    .<. SubExp -> TPrimExp Int64 VName
    pe64 SubExp
    kdim) [SubExp -> DimIndex SubExp
    forall d. d -> DimIndex d
    DimFix (SubExp -> DimIndex SubExp) -> SubExp -> DimIndex SubExp
    forall a b. (a -> b) -> a -> b
    $ VName -> SubExp
    Var VName
    ltid],
            tilingReadTile =
              readTile1D tile_size gid gtid grid,
            tilingProcessTile =
              processTile1D gid gtid kdim tile_size grid,
            tilingProcessResidualTile =
              processResidualTile1D gid gtid kdim tile_size grid,
            tilingTileReturns = tileReturns dims_on_top [(kdim, tile_size)],
            tilingTileShape = Shape [tile_size],
            tilingNumWholeTiles =
              letSubExp "num_whole_tiles" $
                BasicOp $
                  BinOp (SQuot Int64 Unsafe) w tile_size,
            tilingLevel = SegBlock SegNoVirt (Just grid),
            tilingSpace = space
          }

    invariantToOneOfTwoInnerDims ::
      Names ->
      M.Map VName Names ->
      [VName] ->
      VName ->
      Maybe InputArray
    invariantToOneOfTwoInnerDims :: Names -> AliasTable -> [VName] -> VName -> Maybe InputArray
    invariantToOneOfTwoInnerDims Names
    branch_variant AliasTable
    variance [VName]
    dims VName
    arr = do
      j : i : _ <- [VName] -> Maybe [VName]
    forall a. a -> Maybe a
    Just ([VName] -> Maybe [VName]) -> [VName] -> Maybe [VName]
    forall a b. (a -> b) -> a -> b
    $ [VName] -> [VName]
    forall a. [a] -> [a]
    reverse [VName]
    dims
      let variant_to = Names -> VName -> AliasTable -> Names
    forall k a. Ord k => a -> k -> Map k a -> a
    M.findWithDefault Names
    forall a. Monoid a => a
    mempty VName
    arr AliasTable
    variance
          branch_invariant = Bool -> Bool
    not (Bool -> Bool) -> Bool -> Bool
    forall a b. (a -> b) -> a -> b
    $ VName -> Names -> Bool
    nameIn VName
    j Names
    branch_variant Bool -> Bool -> Bool
    || VName -> Names -> Bool
    nameIn VName
    i Names
    branch_variant
      if branch_invariant && i `nameIn` variant_to && j `notNameIn` variant_to
        then Just $ InputTile [0, 1] arr
        else
          if branch_invariant && j `nameIn` variant_to && i `notNameIn` variant_to
            then Just $ InputTile [1, 0] arr
            else Just $ InputDontTile arr

    -- Reconstruct the original gtids from group and local IDs.
    reconstructGtids2D ::
      SubExp ->
      (VName, VName) ->
      (VName, VName) ->
      (VName, VName) ->
      Builder GPU ()
    reconstructGtids2D :: SubExp
    -> (VName, VName)
    -> (VName, VName)
    -> (VName, VName)
    -> BuilderT GPU (State VNameSource) ()
    reconstructGtids2D SubExp
    tile_size (VName
    gtid_x, VName
    gtid_y) (VName
    gid_x, VName
    gid_y) (VName
    ltid_x, VName
    ltid_y) = do
      -- Reconstruct the original gtids from gid_x/gid_y and ltid_x/ltid_y.
      [VName]
    -> Exp (Rep (BuilderT GPU (State VNameSource)))
    -> BuilderT GPU (State VNameSource) ()
    forall (m :: * -> *).
    MonadBuilder m =>
    [VName] -> Exp (Rep m) -> m ()
    letBindNames [VName
    gtid_x]
        (Exp GPU -> BuilderT GPU (State VNameSource) ())
    -> BuilderT GPU (State VNameSource) (Exp GPU)
    -> BuilderT GPU (State VNameSource) ()
    forall (m :: * -> *) a b. Monad m => (a -> m b) -> m a -> m b
    =<< TPrimExp Int64 VName
    -> BuilderT
         GPU
         (State VNameSource)
         (Exp (Rep (BuilderT GPU (State VNameSource))))
    forall a (m :: * -> *).
    (ToExp a, MonadBuilder m) =>
    a -> m (Exp (Rep m))
    forall (m :: * -> *).
    MonadBuilder m =>
    TPrimExp Int64 VName -> m (Exp (Rep m))
    toExp (VName -> TPrimExp Int64 VName
    forall a. a -> TPrimExp Int64 a
    le64 VName
    gid_x TPrimExp Int64 VName
    -> TPrimExp Int64 VName -> TPrimExp Int64 VName
    forall a. Num a => a -> a -> a
    * SubExp -> TPrimExp Int64 VName
    pe64 SubExp
    tile_size TPrimExp Int64 VName
    -> TPrimExp Int64 VName -> TPrimExp Int64 VName
    forall a. Num a => a -> a -> a
    + VName -> TPrimExp Int64 VName
    forall a. a -> TPrimExp Int64 a
    le64 VName
    ltid_x)
      [VName]
    -> Exp (Rep (BuilderT GPU (State VNameSource)))
    -> BuilderT GPU (State VNameSource) ()
    forall (m :: * -> *).
    MonadBuilder m =>
    [VName] -> Exp (Rep m) -> m ()
    letBindNames [VName
    gtid_y]
        (Exp GPU -> BuilderT GPU (State VNameSource) ())
    -> BuilderT GPU (State VNameSource) (Exp GPU)
    -> BuilderT GPU (State VNameSource) ()
    forall (m :: * -> *) a b. Monad m => (a -> m b) -> m a -> m b
    =<< TPrimExp Int64 VName
    -> BuilderT
         GPU
         (State VNameSource)
         (Exp (Rep (BuilderT GPU (State VNameSource))))
    forall a (m :: * -> *).
    (ToExp a, MonadBuilder m) =>
    a -> m (Exp (Rep m))
    forall (m :: * -> *).
    MonadBuilder m =>
    TPrimExp Int64 VName -> m (Exp (Rep m))
    toExp (VName -> TPrimExp Int64 VName
    forall a. a -> TPrimExp Int64 a
    le64 VName
    gid_y TPrimExp Int64 VName
    -> TPrimExp Int64 VName -> TPrimExp Int64 VName
    forall a. Num a => a -> a -> a
    * SubExp -> TPrimExp Int64 VName
    pe64 SubExp
    tile_size TPrimExp Int64 VName
    -> TPrimExp Int64 VName -> TPrimExp Int64 VName
    forall a. Num a => a -> a -> a
    + VName -> TPrimExp Int64 VName
    forall a. a -> TPrimExp Int64 a
    le64 VName
    ltid_y)

    readTile2D ::
      (SubExp, SubExp) ->
      (VName, VName) ->
      (VName, VName) ->
      SubExp ->
      TileKind ->
      PrivStms ->
      SubExp ->
      [InputArray] ->
      Builder GPU [InputTile]
    readTile2D :: (SubExp, SubExp)
    -> (VName, VName)
    -> (VName, VName)
    -> SubExp
    -> TileKind
    -> PrivStms
    -> SubExp
    -> [InputArray]
    -> Builder GPU [InputTile]
    readTile2D (SubExp
    kdim_x, SubExp
    kdim_y) (VName
    gtid_x, VName
    gtid_y) (VName
    gid_x, VName
    gid_y) SubExp
    tile_size TileKind
    kind PrivStms
    privstms SubExp
    tile_id [InputArray]
    inputs =
      ([VName] -> [InputTile])
    -> BuilderT GPU (State VNameSource) [VName]
    -> Builder GPU [InputTile]
    forall a b.
    (a -> b)
    -> BuilderT GPU (State VNameSource) a
    -> BuilderT GPU (State VNameSource) b
    forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
    fmap ([InputArray] -> [VName] -> [InputTile]
    inputsToTiles [InputArray]
    inputs)
        (BuilderT GPU (State VNameSource) [VName]
     -> Builder GPU [InputTile])
    -> (((VName, VName) -> Builder GPU Result)
        -> BuilderT GPU (State VNameSource) [VName])
    -> ((VName, VName) -> Builder GPU Result)
    -> Builder GPU [InputTile]
    forall b c a. (b -> c) -> (a -> b) -> a -> c
    . [Char]
    -> SegLevel
    -> ResultManifest
    -> (SubExp, SubExp)
    -> ((VName, VName) -> Builder GPU Result)
    -> BuilderT GPU (State VNameSource) [VName]
    segMap2D
          [Char]
    "full_tile"
          (SegVirt -> Maybe KernelGrid -> SegLevel
    SegThread (SegSeqDims -> SegVirt
    SegNoVirtFull ([Int] -> SegSeqDims
    SegSeqDims [])) Maybe KernelGrid
    forall a. Maybe a
    Nothing)
          ResultManifest
    ResultNoSimplify
          (SubExp
    tile_size, SubExp
    tile_size)
        (((VName, VName) -> Builder GPU Result) -> Builder GPU [InputTile])
    -> ((VName, VName) -> Builder GPU Result)
    -> Builder GPU [InputTile]
    forall a b. (a -> b) -> a -> b
    $ \(VName
    ltid_x, VName
    ltid_y) -> do
          i <-
            [Char]
    -> Exp (Rep (BuilderT GPU (State VNameSource)))
    -> Builder GPU SubExp
    forall (m :: * -> *).
    MonadBuilder m =>
    [Char] -> Exp (Rep m) -> m SubExp
    letSubExp [Char]
    "i"
              (Exp GPU -> Builder GPU SubExp)
    -> BuilderT GPU (State VNameSource) (Exp GPU) -> Builder GPU SubExp
    forall (m :: * -> *) a b. Monad m => (a -> m b) -> m a -> m b
    =<< TPrimExp Int64 VName
    -> BuilderT
         GPU
         (State VNameSource)
         (Exp (Rep (BuilderT GPU (State VNameSource))))
    forall a (m :: * -> *).
    (ToExp a, MonadBuilder m) =>
    a -> m (Exp (Rep m))
    forall (m :: * -> *).
    MonadBuilder m =>
    TPrimExp Int64 VName -> m (Exp (Rep m))
    toExp (SubExp -> TPrimExp Int64 VName
    pe64 SubExp
    tile_id TPrimExp Int64 VName
    -> TPrimExp Int64 VName -> TPrimExp Int64 VName
    forall a. Num a => a -> a -> a
    * SubExp -> TPrimExp Int64 VName
    pe64 SubExp
    tile_size TPrimExp Int64 VName
    -> TPrimExp Int64 VName -> TPrimExp Int64 VName
    forall a. Num a => a -> a -> a
    + VName -> TPrimExp Int64 VName
    forall a. a -> TPrimExp Int64 a
    le64 VName
    ltid_x)
          j <-
            letSubExp "j"
              =<< toExp (pe64 tile_id * pe64 tile_size + le64 ltid_y)

          reconstructGtids2D tile_size (gtid_x, gtid_y) (gid_x, gid_y) (ltid_x, ltid_y)
          addPrivStms [DimFix $ Var ltid_x, DimFix $ Var ltid_y] privstms

          let arrs_and_perms = [InputArray] -> [(VName, [Int])]
    tiledInputs [InputArray]
    inputs

              readTileElem (VName
    arr, [Int]
    perm) =
                -- No need for fullSlice because we are tiling only prims.
                [Char]
    -> Exp (Rep (BuilderT GPU (State VNameSource)))
    -> BuilderT GPU (State VNameSource) VName
    forall (m :: * -> *).
    MonadBuilder m =>
    [Char] -> Exp (Rep m) -> m VName
    letExp
                  [Char]
    "tile_elem"
                  ( BasicOp -> Exp (Rep (BuilderT GPU (State VNameSource)))
    BasicOp -> Exp GPU
    forall rep. BasicOp -> Exp rep
    BasicOp (BasicOp -> Exp (Rep (BuilderT GPU (State VNameSource))))
    -> (Slice SubExp -> BasicOp)
    -> Slice SubExp
    -> Exp (Rep (BuilderT GPU (State VNameSource)))
    forall b c a. (b -> c) -> (a -> b) -> a -> c
    . VName -> Slice SubExp -> BasicOp
    Index VName
    arr (Slice SubExp -> Exp (Rep (BuilderT GPU (State VNameSource))))
    -> Slice SubExp -> Exp (Rep (BuilderT GPU (State VNameSource)))
    forall a b. (a -> b) -> a -> b
    $
                      [DimIndex SubExp] -> Slice SubExp
    forall d. [DimIndex d] -> Slice d
    Slice [SubExp -> DimIndex SubExp
    forall d. d -> DimIndex d
    DimFix (SubExp -> DimIndex SubExp) -> SubExp -> DimIndex SubExp
    forall a b. (a -> b) -> a -> b
    $ [SubExp] -> SubExp
    forall a. HasCallStack => [a] -> a
    last ([SubExp] -> SubExp) -> [SubExp] -> SubExp
    forall a b. (a -> b) -> a -> b
    $ [Int] -> [SubExp] -> [SubExp]
    forall a. [Int] -> [a] -> [a]
    rearrangeShape [Int]
    perm [SubExp
    i, SubExp
    j]]
                  )

              readTileElemIfInBounds (VName
    arr, [Int]
    perm) = do
                arr_t <- VName -> BuilderT GPU (State VNameSource) Type
    forall rep (m :: * -> *). HasScope rep m => VName -> m Type
    lookupType VName
    arr
                let tile_t = Type -> Type
    forall u. TypeBase Shape u -> TypeBase Shape u
    rowType Type
    arr_t
                    w = Int -> Type -> SubExp
    forall u. Int -> TypeBase Shape u -> SubExp
    arraySize Int
    0 Type
    arr_t
                    idx = [SubExp] -> SubExp
    forall a. HasCallStack => [a] -> a
    last ([SubExp] -> SubExp) -> [SubExp] -> SubExp
    forall a b. (a -> b) -> a -> b
    $ [Int] -> [SubExp] -> [SubExp]
    forall a. [Int] -> [a] -> [a]
    rearrangeShape [Int]
    perm [SubExp
    i, SubExp
    j]
                    othercheck =
                      [TPrimExp Bool VName] -> TPrimExp Bool VName
    forall a. HasCallStack => [a] -> a
    last ([TPrimExp Bool VName] -> TPrimExp Bool VName)
    -> [TPrimExp Bool VName] -> TPrimExp Bool VName
    forall a b. (a -> b) -> a -> b
    $
                        [Int] -> [TPrimExp Bool VName] -> [TPrimExp Bool VName]
    forall a. [Int] -> [a] -> [a]
    rearrangeShape
                          [Int]
    perm
                          [ VName -> TPrimExp Int64 VName
    forall a. a -> TPrimExp Int64 a
    le64 VName
    gtid_y TPrimExp Int64 VName -> TPrimExp Int64 VName -> TPrimExp Bool VName
    forall {k} v (t :: k).
    Eq v =>
    TPrimExp t v -> TPrimExp t v -> TPrimExp Bool v
    .<. SubExp -> TPrimExp Int64 VName
    pe64 SubExp
    kdim_y,
                            VName -> TPrimExp Int64 VName
    forall a. a -> TPrimExp Int64 a
    le64 VName
    gtid_x TPrimExp Int64 VName -> TPrimExp Int64 VName -> TPrimExp Bool VName
    forall {k} v (t :: k).
    Eq v =>
    TPrimExp t v -> TPrimExp t v -> TPrimExp Bool v
    .<. SubExp -> TPrimExp Int64 VName
    pe64 SubExp
    kdim_x
                          ]
                eIf
                  (toExp $ pe64 idx .<. pe64 w .&&. othercheck)
                  (eBody [pure $ BasicOp $ Index arr $ Slice [DimFix idx]])
                  (eBody [eBlank tile_t])

          fmap varsRes $
            case kind of
              TileKind
    TilePartial ->
                ((VName, [Int]) -> BuilderT GPU (State VNameSource) VName)
    -> [(VName, [Int])] -> BuilderT GPU (State VNameSource) [VName]
    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 ([Char]
    -> Exp (Rep (BuilderT GPU (State VNameSource)))
    -> BuilderT GPU (State VNameSource) VName
    forall (m :: * -> *).
    MonadBuilder m =>
    [Char] -> Exp (Rep m) -> m VName
    letExp [Char]
    "pre2d" (Exp GPU -> BuilderT GPU (State VNameSource) VName)
    -> ((VName, [Int]) -> BuilderT GPU (State VNameSource) (Exp GPU))
    -> (VName, [Int])
    -> BuilderT GPU (State VNameSource) VName
    forall (m :: * -> *) b c a.
    Monad m =>
    (b -> m c) -> (a -> m b) -> a -> m c
    <=< (VName, [Int]) -> BuilderT GPU (State VNameSource) (Exp GPU)
    readTileElemIfInBounds) [(VName, [Int])]
    arrs_and_perms
              TileKind
    TileFull ->
                ((VName, [Int]) -> BuilderT GPU (State VNameSource) VName)
    -> [(VName, [Int])] -> BuilderT GPU (State VNameSource) [VName]
    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 (VName, [Int]) -> BuilderT GPU (State VNameSource) VName
    readTileElem [(VName, [Int])]
    arrs_and_perms

    findTileSize :: (HasScope rep m) => [InputTile] -> m SubExp
    findTileSize :: forall rep (m :: * -> *). HasScope rep m => [InputTile] -> m SubExp
    findTileSize [InputTile]
    tiles =
      case (InputTile -> Maybe VName) -> [InputTile] -> [VName]
    forall a b. (a -> Maybe b) -> [a] -> [b]
    mapMaybe InputTile -> Maybe VName
    isTiled [InputTile]
    tiles of
        VName
    v : [VName]
    _ -> Int -> Type -> SubExp
    forall u. Int -> TypeBase Shape u -> SubExp
    arraySize Int
    0 (Type -> SubExp) -> m Type -> m SubExp
    forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
    <$> VName -> m Type
    forall rep (m :: * -> *). HasScope rep m => VName -> m Type
    lookupType VName
    v
        [] -> SubExp -> m SubExp
    forall a. a -> m a
    forall (f :: * -> *) a. Applicative f => a -> f a
    pure (SubExp -> m SubExp) -> SubExp -> m SubExp
    forall a b. (a -> b) -> a -> b
    $ IntType -> Integer -> SubExp
    intConst IntType
    Int64 Integer
    0
      where
        isTiled :: InputTile -> Maybe VName
    isTiled InputUntiled {} = Maybe VName
    forall a. Maybe a
    Nothing
        isTiled (InputTiled [Int]
    _ VName
    tile) = VName -> Maybe VName
    forall a. a -> Maybe a
    Just VName
    tile

    processTile2D ::
      (VName, VName) ->
      (VName, VName) ->
      (SubExp, SubExp) ->
      SubExp ->
      ProcessTileArgs ->
      Builder GPU [VName]
    processTile2D :: (VName, VName)
    -> (VName, VName)
    -> (SubExp, SubExp)
    -> SubExp
    -> ProcessTileArgs
    -> BuilderT GPU (State VNameSource) [VName]
    processTile2D (VName
    gid_x, VName
    gid_y) (VName
    gtid_x, VName
    gtid_y) (SubExp
    kdim_x, SubExp
    kdim_y) SubExp
    tile_size ProcessTileArgs
    tile_args = do
      let privstms :: PrivStms
    privstms = ProcessTileArgs -> PrivStms
    processPrivStms ProcessTileArgs
    tile_args
          red_comm :: Commutativity
    red_comm = ProcessTileArgs -> Commutativity
    processComm ProcessTileArgs
    tile_args
          red_lam :: Lambda GPU
    red_lam = ProcessTileArgs -> Lambda GPU
    processRedLam ProcessTileArgs
    tile_args
          map_lam :: Lambda GPU
    map_lam = ProcessTileArgs -> Lambda GPU
    processMapLam ProcessTileArgs
    tile_args
          tiles :: [InputTile]
    tiles = ProcessTileArgs -> [InputTile]
    processTiles ProcessTileArgs
    tile_args
          accs :: [VName]
    accs = ProcessTileArgs -> [VName]
    processAcc ProcessTileArgs
    tile_args
          tile_id :: SubExp
    tile_id = ProcessTileArgs -> SubExp
    processTileId ProcessTileArgs
    tile_args

      -- Might be truncated in case of a partial tile.
      actual_tile_size <- [InputTile] -> Builder GPU SubExp
    forall rep (m :: * -> *). HasScope rep m => [InputTile] -> m SubExp
    findTileSize [InputTile]
    tiles

      segMap2D
        "acc"
        (SegThreadInBlock (SegNoVirtFull (SegSeqDims [])))
        ResultPrivate
        (tile_size, tile_size)
        $ \(VName
    ltid_x, VName
    ltid_y) -> do
          SubExp
    -> (VName, VName)
    -> (VName, VName)
    -> (VName, VName)
    -> BuilderT GPU (State VNameSource) ()
    reconstructGtids2D SubExp
    tile_size (VName
    gtid_x, VName
    gtid_y) (VName
    gid_x, VName
    gid_y) (VName
    ltid_x, VName
    ltid_y)

          [DimIndex SubExp]
    -> PrivStms -> BuilderT GPU (State VNameSource) ()
    addPrivStms [SubExp -> DimIndex SubExp
    forall d. d -> DimIndex d
    DimFix (SubExp -> DimIndex SubExp) -> SubExp -> DimIndex SubExp
    forall a b. (a -> b) -> a -> b
    $ VName -> SubExp
    Var VName
    ltid_x, SubExp -> DimIndex SubExp
    forall d. d -> DimIndex d
    DimFix (SubExp -> DimIndex SubExp) -> SubExp -> DimIndex SubExp
    forall a b. (a -> b) -> a -> b
    $ VName -> SubExp
    Var VName
    ltid_y] PrivStms
    privstms

          -- We replace the neutral elements with the accumulators (this is
          -- OK because the parallel semantics are not used after this
          -- point).
          thread_accs <- [VName]
    -> (VName -> Builder GPU SubExp)
    -> BuilderT GPU (State VNameSource) [SubExp]
    forall (t :: * -> *) (m :: * -> *) a b.
    (Traversable t, Monad m) =>
    t a -> (a -> m b) -> m (t b)
    forM [VName]
    accs ((VName -> Builder GPU SubExp)
     -> BuilderT GPU (State VNameSource) [SubExp])
    -> (VName -> Builder GPU SubExp)
    -> BuilderT GPU (State VNameSource) [SubExp]
    forall a b. (a -> b) -> a -> b
    $ \VName
    acc ->
            [Char]
    -> Exp (Rep (BuilderT GPU (State VNameSource)))
    -> Builder GPU SubExp
    forall (m :: * -> *).
    MonadBuilder m =>
    [Char] -> Exp (Rep m) -> m SubExp
    letSubExp [Char]
    "acc" (Exp (Rep (BuilderT GPU (State VNameSource)))
     -> Builder GPU SubExp)
    -> Exp (Rep (BuilderT GPU (State VNameSource)))
    -> Builder GPU SubExp
    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
    $ VName -> Slice SubExp -> BasicOp
    Index VName
    acc (Slice SubExp -> BasicOp) -> Slice SubExp -> BasicOp
    forall a b. (a -> b) -> a -> b
    $ [DimIndex SubExp] -> Slice SubExp
    forall d. [DimIndex d] -> Slice d
    Slice [SubExp -> DimIndex SubExp
    forall d. d -> DimIndex d
    DimFix (SubExp -> DimIndex SubExp) -> SubExp -> DimIndex SubExp
    forall a b. (a -> b) -> a -> b
    $ VName -> SubExp
    Var VName
    ltid_x, SubExp -> DimIndex SubExp
    forall d. d -> DimIndex d
    DimFix (SubExp -> DimIndex SubExp) -> SubExp -> DimIndex SubExp
    forall a b. (a -> b) -> a -> b
    $ VName -> SubExp
    Var VName
    ltid_y]
          let form' = [Reduce GPU] -> Lambda GPU -> ScremaForm GPU
    forall rep. [Reduce rep] -> Lambda rep -> ScremaForm rep
    redomapSOAC [Commutativity -> Lambda GPU -> [SubExp] -> Reduce GPU
    forall rep. Commutativity -> Lambda rep -> [SubExp] -> Reduce rep
    Reduce Commutativity
    red_comm Lambda GPU
    red_lam [SubExp]
    thread_accs] Lambda GPU
    map_lam

              sliceTile (InputUntiled VName
    arr) =
                VName
    -> SubExp
    -> SubExp
    -> SubExp
    -> BuilderT GPU (State VNameSource) VName
    forall (m :: * -> *).
    MonadBuilder m =>
    VName -> SubExp -> SubExp -> SubExp -> m VName
    sliceUntiled VName
    arr SubExp
    tile_id SubExp
    tile_size SubExp
    actual_tile_size
              sliceTile (InputTiled [Int]
    perm VName
    tile) = do
                tile_t <- VName -> BuilderT GPU (State VNameSource) Type
    forall rep (m :: * -> *). HasScope rep m => VName -> m Type
    lookupType VName
    tile
                let idx = SubExp -> DimIndex SubExp
    forall d. d -> DimIndex d
    DimFix (SubExp -> DimIndex SubExp) -> SubExp -> DimIndex SubExp
    forall a b. (a -> b) -> a -> b
    $ VName -> SubExp
    Var (VName -> SubExp) -> VName -> SubExp
    forall a b. (a -> b) -> a -> b
    $ [VName] -> VName
    forall a. HasCallStack => [a] -> a
    head ([VName] -> VName) -> [VName] -> VName
    forall a b. (a -> b) -> a -> b
    $ [Int] -> [VName] -> [VName]
    forall a. [Int] -> [a] -> [a]
    rearrangeShape [Int]
    perm [VName
    ltid_x, VName
    ltid_y]
                letExp "tile" $
                  BasicOp $
                    Index tile $
                      sliceAt tile_t (head perm) [idx]

          tiles' <- mapM sliceTile tiles

          fmap varsRes $
            letTupExp "acc"
              =<< eIf
                ( toExp $ le64 gtid_x .<. pe64 kdim_x .&&. le64 gtid_y .<. pe64 kdim_y
                )
                (eBody [pure $ Op $ OtherOp $ Screma actual_tile_size tiles' form'])
                (resultBodyM thread_accs)

    processResidualTile2D ::
      (VName, VName) ->
      (VName, VName) ->
      (SubExp, SubExp) ->
      SubExp ->
      ResidualTileArgs ->
      Builder GPU [VName]
    processResidualTile2D :: (VName, VName)
    -> (VName, VName)
    -> (SubExp, SubExp)
    -> SubExp
    -> ResidualTileArgs
    -> BuilderT GPU (State VNameSource) [VName]
    processResidualTile2D (VName, VName)
    gids (VName, VName)
    gtids (SubExp, SubExp)
    kdims SubExp
    tile_size ResidualTileArgs
    args = do
      -- The number of residual elements that are not covered by
      -- the whole tiles.
      residual_input <-
        [Char]
    -> Exp (Rep (BuilderT GPU (State VNameSource)))
    -> Builder GPU SubExp
    forall (m :: * -> *).
    MonadBuilder m =>
    [Char] -> Exp (Rep m) -> m SubExp
    letSubExp [Char]
    "residual_input" (Exp (Rep (BuilderT GPU (State VNameSource)))
     -> Builder GPU SubExp)
    -> Exp (Rep (BuilderT GPU (State VNameSource)))
    -> Builder GPU SubExp
    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
    $
            BinOp -> SubExp -> SubExp -> BasicOp
    BinOp (IntType -> Safety -> BinOp
    SRem IntType
    Int64 Safety
    Unsafe) SubExp
    w SubExp
    tile_size

      letTupExp "acc_after_residual"
        =<< eIf
          (toExp $ pe64 residual_input .==. 0)
          (resultBodyM $ map Var accs)
          (nonemptyTile residual_input)
      where
        privstms :: PrivStms
    privstms = ResidualTileArgs -> PrivStms
    residualPrivStms ResidualTileArgs
    args
        red_comm :: Commutativity
    red_comm = ResidualTileArgs -> Commutativity
    residualComm ResidualTileArgs
    args
        red_lam :: Lambda GPU
    red_lam = ResidualTileArgs -> Lambda GPU
    residualRedLam ResidualTileArgs
    args
        map_lam :: Lambda GPU
    map_lam = ResidualTileArgs -> Lambda GPU
    residualMapLam ResidualTileArgs
    args
        accs :: [VName]
    accs = ResidualTileArgs -> [VName]
    residualAcc ResidualTileArgs
    args
        inputs :: [InputArray]
    inputs = ResidualTileArgs -> [InputArray]
    residualInput ResidualTileArgs
    args
        num_whole_tiles :: SubExp
    num_whole_tiles = ResidualTileArgs -> SubExp
    residualNumWholeTiles ResidualTileArgs
    args
        w :: SubExp
    w = ResidualTileArgs -> SubExp
    residualInputSize ResidualTileArgs
    args

        nonemptyTile :: SubExp -> BuilderT GPU (State VNameSource) (Body GPU)
    nonemptyTile SubExp
    residual_input = Body GPU -> BuilderT GPU (State VNameSource) (Body GPU)
    forall rep (m :: * -> *).
    (Renameable rep, MonadFreshNames m) =>
    Body rep -> m (Body rep)
    renameBody (Body GPU -> BuilderT GPU (State VNameSource) (Body GPU))
    -> (Builder GPU Result
        -> BuilderT GPU (State VNameSource) (Body GPU))
    -> Builder GPU Result
    -> BuilderT GPU (State VNameSource) (Body GPU)
    forall (m :: * -> *) b c a.
    Monad m =>
    (b -> m c) -> (a -> m b) -> a -> m c
    <=< Builder GPU Result -> BuilderT GPU (State VNameSource) (Body GPU)
    forall rep (m :: * -> *) somerep.
    (Buildable rep, MonadFreshNames m, HasScope somerep m,
     SameScope somerep rep) =>
    Builder rep Result -> m (Body rep)
    runBodyBuilder (Builder GPU Result -> BuilderT GPU (State VNameSource) (Body GPU))
    -> Builder GPU Result
    -> BuilderT GPU (State VNameSource) (Body GPU)
    forall a b. (a -> b) -> a -> b
    $ do
          -- Collectively construct a tile.  Threads that are out-of-bounds
          -- provide a blank dummy value.
          full_tile <-
            (SubExp, SubExp)
    -> (VName, VName)
    -> (VName, VName)
    -> SubExp
    -> TileKind
    -> PrivStms
    -> SubExp
    -> [InputArray]
    -> Builder GPU [InputTile]
    readTile2D
              (SubExp, SubExp)
    kdims
              (VName, VName)
    gtids
              (VName, VName)
    gids
              SubExp
    tile_size
              TileKind
    TilePartial
              PrivStms
    privstms
              SubExp
    num_whole_tiles
              [InputArray]
    inputs

          let slice =
                SubExp -> SubExp -> SubExp -> DimIndex SubExp
    forall d. d -> d -> d -> DimIndex d
    DimSlice (IntType -> Integer -> SubExp
    intConst IntType
    Int64 Integer
    0) SubExp
    residual_input (IntType -> Integer -> SubExp
    intConst IntType
    Int64 Integer
    1)
          tiles <- forM full_tile $ \case
            InputTiled [Int]
    perm VName
    tile' ->
              [Int] -> VName -> InputTile
    InputTiled [Int]
    perm
                (VName -> InputTile)
    -> BuilderT GPU (State VNameSource) VName
    -> BuilderT GPU (State VNameSource) InputTile
    forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
    <$> [Char]
    -> Exp (Rep (BuilderT GPU (State VNameSource)))
    -> BuilderT GPU (State VNameSource) VName
    forall (m :: * -> *).
    MonadBuilder m =>
    [Char] -> Exp (Rep m) -> m VName
    letExp [Char]
    "partial_tile" (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
    $ VName -> Slice SubExp -> BasicOp
    Index VName
    tile' ([DimIndex SubExp] -> Slice SubExp
    forall d. [DimIndex d] -> Slice d
    Slice [DimIndex SubExp
    slice, DimIndex SubExp
    slice]))
            InputUntiled VName
    arr ->
              InputTile -> BuilderT GPU (State VNameSource) InputTile
    forall a. a -> BuilderT GPU (State VNameSource) a
    forall (f :: * -> *) a. Applicative f => a -> f a
    pure (InputTile -> BuilderT GPU (State VNameSource) InputTile)
    -> InputTile -> BuilderT GPU (State VNameSource) InputTile
    forall a b. (a -> b) -> a -> b
    $ VName -> InputTile
    InputUntiled VName
    arr

          let tile_args =
                PrivStms
    -> Commutativity
    -> Lambda GPU
    -> Lambda GPU
    -> [InputTile]
    -> [VName]
    -> SubExp
    -> ProcessTileArgs
    ProcessTileArgs PrivStms
    privstms Commutativity
    red_comm Lambda GPU
    red_lam Lambda GPU
    map_lam [InputTile]
    tiles [VName]
    accs SubExp
    num_whole_tiles

          -- Now each thread performs a traversal of the tile and
          -- updates its accumulator.
          varsRes <$> processTile2D gids gtids kdims tile_size tile_args

    tiling2d :: [(VName, SubExp)] -> DoTiling (VName, VName) (SubExp, SubExp)
    tiling2d :: [(VName, SubExp)] -> DoTiling (VName, VName) (SubExp, SubExp)
    tiling2d [(VName, SubExp)]
    dims_on_top (VName
    gtid_x, VName
    gtid_y) (SubExp
    kdim_x, SubExp
    kdim_y) SubExp
    w = do
      gid_x <- [Char] -> BuilderT GPU (State VNameSource) VName
    forall (m :: * -> *). MonadFreshNames m => [Char] -> m VName
    newVName [Char]
    "gid_x"
      gid_y <- newVName "gid_y"

      tile_size_key <- nameFromString . prettyString <$> newVName "tile_size"
      tile_size <- letSubExp "tile_size" $ Op $ SizeOp $ GetSize tile_size_key SizeTile
      tblock_size <- letSubExp "tblock_size" $ BasicOp $ BinOp (Mul Int64 OverflowUndef) tile_size tile_size

      num_tblocks_x <-
        letSubExp "num_tblocks_x" $
          BasicOp $
            BinOp (SDivUp Int64 Unsafe) kdim_x tile_size
      num_tblocks_y <-
        letSubExp "num_tblocks_y" $
          BasicOp $
            BinOp (SDivUp Int64 Unsafe) kdim_y tile_size

      num_tblocks <-
        letSubExp "num_tblocks_top"
          =<< foldBinOp
            (Mul Int64 OverflowUndef)
            num_tblocks_x
            (num_tblocks_y : map snd dims_on_top)

      gid_flat <- newVName "gid_flat"
      let grid = Count NumBlocks SubExp -> Count BlockSize SubExp -> KernelGrid
    KernelGrid (SubExp -> Count NumBlocks SubExp
    forall {k} (u :: k) e. e -> Count u e
    Count SubExp
    num_tblocks) (SubExp -> Count BlockSize SubExp
    forall {k} (u :: k) e. e -> Count u e
    Count SubExp
    tblock_size)
          lvl = SegVirt -> Maybe KernelGrid -> SegLevel
    SegBlock (SegSeqDims -> SegVirt
    SegNoVirtFull ([Int] -> SegSeqDims
    SegSeqDims [])) (KernelGrid -> Maybe KernelGrid
    forall a. a -> Maybe a
    Just KernelGrid
    grid)
          space =
            VName -> [(VName, SubExp)] -> SegSpace
    SegSpace VName
    gid_flat ([(VName, SubExp)] -> SegSpace) -> [(VName, SubExp)] -> SegSpace
    forall a b. (a -> b) -> a -> b
    $
              [(VName, SubExp)]
    dims_on_top [(VName, SubExp)] -> [(VName, SubExp)] -> [(VName, SubExp)]
    forall a. [a] -> [a] -> [a]
    ++ [(VName
    gid_x, SubExp
    num_tblocks_x), (VName
    gid_y, SubExp
    num_tblocks_y)]
          tiling_lvl = SegVirt -> SegLevel
    SegThreadInBlock SegVirt
    SegNoVirt

      pure
        Tiling
          { tilingSegMap = \[Char]
    desc ResultManifest
    manifest PrimExp VName -> [DimIndex SubExp] -> Builder GPU Result
    f ->
              [Char]
    -> SegLevel
    -> ResultManifest
    -> (SubExp, SubExp)
    -> ((VName, VName) -> Builder GPU Result)
    -> BuilderT GPU (State VNameSource) [VName]
    segMap2D [Char]
    desc SegLevel
    tiling_lvl ResultManifest
    manifest (SubExp
    tile_size, SubExp
    tile_size) (((VName, VName) -> Builder GPU Result)
     -> BuilderT GPU (State VNameSource) [VName])
    -> ((VName, VName) -> Builder GPU Result)
    -> BuilderT GPU (State VNameSource) [VName]
    forall a b. (a -> b) -> a -> b
    $ \(VName
    ltid_x, VName
    ltid_y) -> do
                SubExp
    -> (VName, VName)
    -> (VName, VName)
    -> (VName, VName)
    -> BuilderT GPU (State VNameSource) ()
    reconstructGtids2D SubExp
    tile_size (VName
    gtid_x, VName
    gtid_y) (VName
    gid_x, VName
    gid_y) (VName
    ltid_x, VName
    ltid_y)
                PrimExp VName -> [DimIndex SubExp] -> Builder GPU Result
    f
                  ( TPrimExp Bool VName -> PrimExp VName
    forall {k} (t :: k) v. TPrimExp t v -> PrimExp v
    untyped (TPrimExp Bool VName -> PrimExp VName)
    -> TPrimExp Bool VName -> PrimExp VName
    forall a b. (a -> b) -> a -> b
    $
                      VName -> TPrimExp Int64 VName
    forall a. a -> TPrimExp Int64 a
    le64 VName
    gtid_x TPrimExp Int64 VName -> TPrimExp Int64 VName -> TPrimExp Bool VName
    forall {k} v (t :: k).
    Eq v =>
    TPrimExp t v -> TPrimExp t v -> TPrimExp Bool v
    .<. SubExp -> TPrimExp Int64 VName
    pe64 SubExp
    kdim_x TPrimExp Bool VName -> TPrimExp Bool VName -> TPrimExp Bool VName
    forall v.
    Eq v =>
    TPrimExp Bool v -> TPrimExp Bool v -> TPrimExp Bool v
    .&&. VName -> TPrimExp Int64 VName
    forall a. a -> TPrimExp Int64 a
    le64 VName
    gtid_y TPrimExp Int64 VName -> TPrimExp Int64 VName -> TPrimExp Bool VName
    forall {k} v (t :: k).
    Eq v =>
    TPrimExp t v -> TPrimExp t v -> TPrimExp Bool v
    .<. SubExp -> TPrimExp Int64 VName
    pe64 SubExp
    kdim_y
                  )
                  [SubExp -> DimIndex SubExp
    forall d. d -> DimIndex d
    DimFix (SubExp -> DimIndex SubExp) -> SubExp -> DimIndex SubExp
    forall a b. (a -> b) -> a -> b
    $ VName -> SubExp
    Var VName
    ltid_x, SubExp -> DimIndex SubExp
    forall d. d -> DimIndex d
    DimFix (SubExp -> DimIndex SubExp) -> SubExp -> DimIndex SubExp
    forall a b. (a -> b) -> a -> b
    $ VName -> SubExp
    Var VName
    ltid_y],
            tilingReadTile = readTile2D (kdim_x, kdim_y) (gtid_x, gtid_y) (gid_x, gid_y) tile_size,
            tilingProcessTile = processTile2D (gid_x, gid_y) (gtid_x, gtid_y) (kdim_x, kdim_y) tile_size,
            tilingProcessResidualTile = processResidualTile2D (gid_x, gid_y) (gtid_x, gtid_y) (kdim_x, kdim_y) tile_size,
            tilingTileReturns = tileReturns dims_on_top [(kdim_x, tile_size), (kdim_y, tile_size)],
            tilingTileShape = Shape [tile_size, tile_size],
            tilingNumWholeTiles =
              letSubExp "num_whole_tiles" $
                BasicOp $
                  BinOp (SQuot Int64 Unsafe) w tile_size,
            tilingLevel = lvl,
            tilingSpace = space
          }