    {-# LANGUAGE TypeFamilies #-}

    -- | Perform a restricted form of block+register tiling corresponding to
    --   the following pattern:
    --     * a redomap is quasi-perfectly nested inside a kernel with at
    --       least two parallel dimension (the perfectly nested restriction
    --       is relaxed a bit to allow for SGEMM);
    --     * all streamed arrays of redomap are one dimensional;
    --     * all streamed arrays are variant to exacly one of the two
    --       innermost parallel dimensions, and conversely for each of
    --       the two innermost parallel dimensions, there is at least
    --       one streamed array variant to it;
    --     * the stream's result is a tuple of scalar values, which are
    --       also the "thread-in-space" return of the kernel.
    --     * We have further restrictions that in principle can be relaxed:
    --          the redomap has exactly two array input
    --          the redomap produces one scalar result
    --          the kernel produces one scalar result
    module Futhark.Optimise.BlkRegTiling (mmBlkRegTiling, doRegTiling3D) where

    import Control.Monad
    import Data.List qualified as L
    import Data.Map.Strict qualified as M
    import Data.Maybe
    import Data.Sequence qualified as Seq
    import Futhark.IR.GPU
    import Futhark.IR.Mem.LMAD qualified as LMAD
    import Futhark.MonadFreshNames
    import Futhark.Optimise.TileLoops.Shared
    import Futhark.Tools
    import Futhark.Transform.Rename
    import Futhark.Transform.Substitute

    se0 :: SubExp
    se0 :: SubExp
    se0 = IntType -> Integer -> SubExp
    intConst IntType
    Int64 Integer
    0

    se1 :: SubExp
    se1 :: SubExp
    se1 = IntType -> Integer -> SubExp
    intConst IntType
    Int64 Integer
    1

    se2 :: SubExp
    se2 :: SubExp
    se2 = IntType -> Integer -> SubExp
    intConst IntType
    Int64 Integer
    2

    se4 :: SubExp
    se4 :: SubExp
    se4 = IntType -> Integer -> SubExp
    intConst IntType
    Int64 Integer
    4

    se8 :: SubExp
    se8 :: SubExp
    se8 = IntType -> Integer -> SubExp
    intConst IntType
    Int64 Integer
    8

    isInnerCoal :: Env -> VName -> Stm GPU -> Bool
    isInnerCoal :: Env -> VName -> Stm GPU -> Bool
    isInnerCoal (WithEnv
    _, IxFnEnv
    ixfn_env) VName
    slc_X (Let (Pat [PatElem (LetDec GPU)
    pe]) StmAux (ExpDec GPU)
    _ (BasicOp (Index VName
    x Slice SubExp
    _)))
      | VName
    slc_X VName -> VName -> Bool
    forall a. Eq a => a -> a -> Bool
    == PatElem (TypeBase Shape NoUniqueness) -> VName
    forall dec. PatElem dec -> VName
    patElemName PatElem (TypeBase Shape NoUniqueness)
    PatElem (LetDec GPU)
    pe =
          -- if not in the table, we assume not-transposed!
          Bool
    -> (LMAD (TPrimExp Int64 VName) -> Bool)
    -> Maybe (LMAD (TPrimExp Int64 VName))
    -> Bool
    forall b a. b -> (a -> b) -> Maybe a -> b
    maybe Bool
    True LMAD (TPrimExp Int64 VName) -> Bool
    innerHasStride1 (Maybe (LMAD (TPrimExp Int64 VName)) -> Bool)
    -> Maybe (LMAD (TPrimExp Int64 VName)) -> Bool
    forall a b. (a -> b) -> a -> b
    $ VName -> IxFnEnv -> Maybe (LMAD (TPrimExp Int64 VName))
    forall k a. Ord k => k -> Map k a -> Maybe a
    M.lookup VName
    x IxFnEnv
    ixfn_env
      where
        innerHasStride1 :: LMAD (TPrimExp Int64 VName) -> Bool
    innerHasStride1 LMAD (TPrimExp Int64 VName)
    lmad =
          let lmad_dims :: [LMADDim (TPrimExp Int64 VName)]
    lmad_dims = LMAD (TPrimExp Int64 VName) -> [LMADDim (TPrimExp Int64 VName)]
    forall num. LMAD num -> [LMADDim num]
    LMAD.dims LMAD (TPrimExp Int64 VName)
    lmad
              stride :: TPrimExp Int64 VName
    stride = LMADDim (TPrimExp Int64 VName) -> TPrimExp Int64 VName
    forall num. LMADDim num -> num
    LMAD.ldStride (LMADDim (TPrimExp Int64 VName) -> TPrimExp Int64 VName)
    -> LMADDim (TPrimExp Int64 VName) -> TPrimExp Int64 VName
    forall a b. (a -> b) -> a -> b
    $ [LMADDim (TPrimExp Int64 VName)] -> LMADDim (TPrimExp Int64 VName)
    forall a. HasCallStack => [a] -> a
    last [LMADDim (TPrimExp Int64 VName)]
    lmad_dims
           in TPrimExp Int64 VName
    stride TPrimExp Int64 VName -> TPrimExp Int64 VName -> Bool
    forall a. Eq a => a -> a -> Bool
    == SubExp -> TPrimExp Int64 VName
    pe64 (IntType -> Integer -> SubExp
    intConst IntType
    Int64 Integer
    1)
    isInnerCoal Env
    _ VName
    _ Stm GPU
    _ =
      [Char] -> Bool
    forall a. HasCallStack => [Char] -> a
    error [Char]
    "kkLoopBody.isInnerCoal: not an error, but I would like to know why!"

    scratch :: (MonadBuilder m) => String -> PrimType -> [SubExp] -> m VName
    scratch :: forall (m :: * -> *).
    MonadBuilder m =>
    [Char] -> PrimType -> [SubExp] -> m VName
    scratch [Char]
    se_name PrimType
    t [SubExp]
    shape = [Char] -> Exp (Rep m) -> m VName
    forall (m :: * -> *).
    MonadBuilder m =>
    [Char] -> Exp (Rep m) -> m VName
    letExp [Char]
    se_name (Exp (Rep m) -> m VName) -> Exp (Rep m) -> m VName
    forall a b. (a -> b) -> a -> b
    $ BasicOp -> Exp (Rep m)
    forall rep. BasicOp -> Exp rep
    BasicOp (BasicOp -> Exp (Rep m)) -> BasicOp -> Exp (Rep m)
    forall a b. (a -> b) -> a -> b
    $ PrimType -> [SubExp] -> BasicOp
    Scratch PrimType
    t [SubExp]
    shape

    -- | Main helper function for Register-and-Block Tiling
    kkLoopBody ::
      Env ->
      ( (SubExp, SubExp, SubExp, SubExp, SubExp, SubExp, SubExp, SubExp),
        SegLevel,
        [Int],
        (VName, SubExp, VName, SubExp, SubExp),
        (VName, VName),
        (Stm GPU, VName, PrimType, Stm GPU, VName, PrimType),
        (Lambda GPU, Lambda GPU)
      ) ->
      VName ->
      (VName, VName, VName) ->
      Bool ->
      Builder GPU [VName]
    kkLoopBody :: Env
    -> ((SubExp, SubExp, SubExp, SubExp, SubExp, SubExp, SubExp,
         SubExp),
        SegLevel, [Int], (VName, SubExp, VName, SubExp, SubExp),
        (VName, VName),
        (Stm GPU, VName, PrimType, Stm GPU, VName, PrimType),
        (Lambda GPU, Lambda GPU))
    -> VName
    -> (VName, VName, VName)
    -> Bool
    -> Builder GPU [VName]
    kkLoopBody
      Env
    env
      ( (SubExp
    rx, SubExp
    ry, SubExp
    tx, SubExp
    ty, SubExp
    tk, SubExp
    tk_div_tx, SubExp
    _tk_div_ty, SubExp
    tx_rx),
        SegLevel
    segthd_lvl,
        [Int]
    var_dims,
        (VName
    gtid_x, SubExp
    width_B, VName
    gtid_y, SubExp
    height_A, SubExp
    common_dim),
        (VName
    iii, VName
    jjj),
        (Stm GPU
    load_A, VName
    inp_A, PrimType
    pt_A, Stm GPU
    load_B, VName
    inp_B, PrimType
    pt_B),
        (Lambda GPU
    map_lam, Lambda GPU
    red_lam)
        )
      VName
    kk0
      (VName
    thd_res_merge, VName
    a_loc_init', VName
    b_loc_init')
      Bool
    epilogue = do
        let (PrimType
    map_t1, PrimType
    map_t2) = (PrimType
    pt_A, PrimType
    pt_B)
        kk <- [Char]
    -> Exp (Rep (BuilderT GPU (State VNameSource)))
    -> BuilderT GPU (State VNameSource) VName
    forall (m :: * -> *).
    MonadBuilder m =>
    [Char] -> Exp (Rep m) -> m VName
    letExp [Char]
    "kk" (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
    =<< 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
    kk0 TPrimExp Int64 VName
    -> TPrimExp Int64 VName -> TPrimExp Int64 VName
    forall a. Num a => a -> a -> a
    * SubExp -> TPrimExp Int64 VName
    pe64 SubExp
    tk)
        -- copy A to shared memory
        (a_loc, aCopyLoc2Reg) <-
          copyGlb2ShMem False kk (gtid_y, iii, map_t1, height_A, inp_A, load_A, a_loc_init')

        -- copy B from global to shared memory
        (b_loc, bCopyLoc2Reg) <-
          copyGlb2ShMem True kk (gtid_x, jjj, map_t2, width_B, inp_B, load_B, b_loc_init')

        -- inner loop updating this thread's accumulator (loop k in mmm_kernels).
        thd_acc <- mkRedomapOneTileBody kk thd_res_merge aCopyLoc2Reg bCopyLoc2Reg True
        pure [thd_acc, a_loc, b_loc]
        where
          mk_ik :: Bool
    -> Bool
    -> (VName, VName)
    -> (VName, VName)
    -> BuilderT
         GPU (State VNameSource) (VName, VName, TPrimExp Int64 VName)
    mk_ik Bool
    is_B Bool
    is_coal (VName
    thd_y, VName
    thd_x) (VName
    i0, VName
    k0)
            | Bool
    is_coal = do
                -- not-transposed case (i.e., already coalesced)
                let (SubExp
    t_par, SubExp
    t_seq) = (SubExp
    tx, SubExp
    tk)
                k <- [Char]
    -> Exp (Rep (BuilderT GPU (State VNameSource)))
    -> BuilderT GPU (State VNameSource) VName
    forall (m :: * -> *).
    MonadBuilder m =>
    [Char] -> Exp (Rep m) -> m VName
    letExp [Char]
    "k" (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
    =<< 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
    thd_x 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
    k0 TPrimExp Int64 VName
    -> TPrimExp Int64 VName -> TPrimExp Int64 VName
    forall a. Num a => a -> a -> a
    * SubExp -> TPrimExp Int64 VName
    pe64 SubExp
    t_par)
                i <- letExp "i" =<< toExp (le64 thd_y + le64 i0 * pe64 t_par)
                -- to optimize bank conflicts, we use padding only for B
                -- iff B has the last dimension permuted.
                let pad_term = if Bool
    is_B then SubExp -> TPrimExp Int64 VName
    pe64 SubExp
    se1 else SubExp -> TPrimExp Int64 VName
    pe64 SubExp
    se0
                let e = VName -> TPrimExp Int64 VName
    forall a. a -> TPrimExp Int64 a
    le64 VName
    k 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
    i TPrimExp Int64 VName
    -> TPrimExp Int64 VName -> TPrimExp Int64 VName
    forall a. Num a => a -> a -> a
    * (SubExp -> TPrimExp Int64 VName
    pe64 SubExp
    t_seq TPrimExp Int64 VName
    -> TPrimExp Int64 VName -> TPrimExp Int64 VName
    forall a. Num a => a -> a -> a
    + TPrimExp Int64 VName
    pad_term)
                pure (i, k, e)
          mk_ik Bool
    _ Bool
    _ (VName
    thd_y, VName
    thd_x) (VName
    i0, VName
    k0) = do
            -- matrix is transposed case (i.e., uncoalesced):
            let (SubExp
    t_par, SubExp
    tr_par) = (SubExp
    tx, SubExp
    tx_rx)
            k <- [Char]
    -> Exp (Rep (BuilderT GPU (State VNameSource)))
    -> BuilderT GPU (State VNameSource) VName
    forall (m :: * -> *).
    MonadBuilder m =>
    [Char] -> Exp (Rep m) -> m VName
    letExp [Char]
    "k" (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
    =<< 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
    thd_y 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
    k0 TPrimExp Int64 VName
    -> TPrimExp Int64 VName -> TPrimExp Int64 VName
    forall a. Num a => a -> a -> a
    * SubExp -> TPrimExp Int64 VName
    pe64 SubExp
    t_par)
            i <- letExp "i" =<< toExp (le64 thd_x + le64 i0 * pe64 t_par)
            -- no padding
            let e = VName -> TPrimExp Int64 VName
    forall a. a -> TPrimExp Int64 a
    le64 VName
    i 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
    k TPrimExp Int64 VName
    -> TPrimExp Int64 VName -> TPrimExp Int64 VName
    forall a. Num a => a -> a -> a
    * SubExp -> TPrimExp Int64 VName
    pe64 SubExp
    tr_par
            pure (i, k, e)
          --
          mkCompLoopRxRy :: Bool
    -> VName
    -> (VName -> VName -> BuilderT GPU (State VNameSource) VName,
        VName -> VName -> BuilderT GPU (State VNameSource) VName)
    -> (VName, VName)
    -> BuilderT GPU (State VNameSource) (Body GPU)
    mkCompLoopRxRy Bool
    fits_ij VName
    css_init (VName -> VName -> BuilderT GPU (State VNameSource) VName
    a_idx_fn, VName -> VName -> BuilderT GPU (State VNameSource) VName
    b_idx_fn) (VName
    ltid_y, VName
    ltid_x) = do
            css <- SubExp
    -> [VName]
    -> (VName
        -> [VName] -> BuilderT GPU (State VNameSource) (Body GPU))
    -> BuilderT GPU (State VNameSource) VName
    forLoop SubExp
    ry [VName
    css_init] ((VName -> [VName] -> BuilderT GPU (State VNameSource) (Body GPU))
     -> BuilderT GPU (State VNameSource) VName)
    -> (VName
        -> [VName] -> BuilderT GPU (State VNameSource) (Body GPU))
    -> BuilderT GPU (State VNameSource) VName
    forall a b. (a -> b) -> a -> b
    $ \VName
    i [VName
    css_merge] -> do
              css <- SubExp
    -> [VName]
    -> (VName
        -> [VName] -> BuilderT GPU (State VNameSource) (Body GPU))
    -> BuilderT GPU (State VNameSource) VName
    forLoop SubExp
    rx [VName
    css_merge] ((VName -> [VName] -> BuilderT GPU (State VNameSource) (Body GPU))
     -> BuilderT GPU (State VNameSource) VName)
    -> (VName
        -> [VName] -> BuilderT GPU (State VNameSource) (Body GPU))
    -> BuilderT GPU (State VNameSource) VName
    forall a b. (a -> b) -> a -> b
    $ \VName
    j [VName
    css_merge'] ->
                ([SubExp]
    -> BuilderT
         GPU
         (State VNameSource)
         (Body (Rep (BuilderT GPU (State VNameSource))))
    [SubExp] -> BuilderT GPU (State VNameSource) (Body GPU)
    forall (m :: * -> *).
    MonadBuilder m =>
    [SubExp] -> m (Body (Rep m))
    resultBodyM ([SubExp] -> BuilderT GPU (State VNameSource) (Body GPU))
    -> (Exp GPU -> BuilderT GPU (State VNameSource) [SubExp])
    -> Exp GPU
    -> BuilderT GPU (State VNameSource) (Body GPU)
    forall (m :: * -> *) b c a.
    Monad m =>
    (b -> m c) -> (a -> m b) -> a -> m c
    <=< [Char]
    -> Exp (Rep (BuilderT GPU (State VNameSource)))
    -> BuilderT GPU (State VNameSource) [SubExp]
    forall (m :: * -> *).
    MonadBuilder m =>
    [Char] -> Exp (Rep m) -> m [SubExp]
    letTupExp' [Char]
    "foo")
                  (Exp GPU -> BuilderT GPU (State VNameSource) (Body GPU))
    -> BuilderT GPU (State VNameSource) (Exp GPU)
    -> BuilderT GPU (State VNameSource) (Body GPU)
    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
    $
                        if Bool
    fits_ij
                          then TPrimExp Bool VName
    forall v. TPrimExp Bool v
    true
                          else -- this condition is never needed because
                          -- if i and j are out of range than css[i,j]
                          -- is garbage anyways and should not be written.
                          -- so fits_ij should be always true!!!
                            (VName -> TPrimExp Int64 VName
    forall a. a -> TPrimExp Int64 a
    le64 VName
    iii 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
    i TPrimExp Int64 VName
    -> TPrimExp Int64 VName -> TPrimExp Int64 VName
    forall a. Num a => a -> a -> a
    + SubExp -> TPrimExp Int64 VName
    pe64 SubExp
    ry 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 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
    height_A)
                              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
    jjj 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
    j TPrimExp Int64 VName
    -> TPrimExp Int64 VName -> TPrimExp Int64 VName
    forall a. Num a => a -> a -> a
    + SubExp -> TPrimExp Int64 VName
    pe64 SubExp
    rx 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 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
    width_B)
                    )
                    ( do
                        a <- VName -> VName -> BuilderT GPU (State VNameSource) VName
    a_idx_fn VName
    ltid_y VName
    i
                        b <- b_idx_fn ltid_x j
                        c <- index "c" css_merge' [i, j]

                        map_lam' <- renameLambda map_lam
                        red_lam' <- renameLambda red_lam

                        -- the inputs to map are supposed to be permutted with the
                        -- inverted permutation, so as to reach the original position;
                        -- it just so happens that the inverse of [a,b] is [b,a]
                        let map_inp_reg = if [Int]
    var_dims [Int] -> [Int] -> Bool
    forall a. Eq a => a -> a -> Bool
    == [Int
    0, Int
    1] then [VName
    a, VName
    b] else [VName
    b, VName
    a]

                        map_res <- eLambda map_lam' (map (eSubExp . Var) map_inp_reg)
                        ~[red_res] <- eLambda red_lam' (map eSubExp $ Var c : map resSubExp map_res)
                        css <- update "css" css_merge' [i, j] (resSubExp red_res)

                        resultBodyM [Var css]
                    )
                    ([SubExp]
    -> BuilderT
         GPU
         (State VNameSource)
         (Body (Rep (BuilderT GPU (State VNameSource))))
    forall (m :: * -> *).
    MonadBuilder m =>
    [SubExp] -> m (Body (Rep m))
    resultBodyM [VName -> SubExp
    Var VName
    css_merge'])
              resultBodyM [Var css]
            resultBodyM [Var css]
          --
          mkRedomapOneTileBody :: VName
    -> VName
    -> (VName
        -> VName -> VName -> BuilderT GPU (State VNameSource) VName)
    -> (VName
        -> VName -> VName -> BuilderT GPU (State VNameSource) VName)
    -> Bool
    -> BuilderT GPU (State VNameSource) VName
    mkRedomapOneTileBody VName
    kk VName
    css_merge VName -> VName -> VName -> BuilderT GPU (State VNameSource) VName
    a_idx_fn VName -> VName -> VName -> BuilderT GPU (State VNameSource) VName
    b_idx_fn Bool
    fits_ij = do
            -- the actual redomap.
            redomap_res <- [Char]
    -> SegLevel
    -> ResultManifest
    -> (SubExp, SubExp)
    -> ((VName, VName) -> BuilderT GPU (State VNameSource) [SubExpRes])
    -> Builder GPU [VName]
    segMap2D [Char]
    "redomap_res" SegLevel
    segthd_lvl ResultManifest
    ResultPrivate (SubExp
    ty, SubExp
    tx) (((VName, VName) -> BuilderT GPU (State VNameSource) [SubExpRes])
     -> Builder GPU [VName])
    -> ((VName, VName) -> BuilderT GPU (State VNameSource) [SubExpRes])
    -> Builder GPU [VName]
    forall a b. (a -> b) -> a -> b
    $
              \(VName
    ltid_y, VName
    ltid_x) -> do
                css_init <- [Char]
    -> VName -> [VName] -> BuilderT GPU (State VNameSource) VName
    forall (m :: * -> *).
    MonadBuilder m =>
    [Char] -> VName -> [VName] -> m VName
    index [Char]
    "css_init" VName
    css_merge [VName
    ltid_y, VName
    ltid_x]

                css <- forLoop tk [css_init] $ \VName
    k [VName
    acc_merge] ->
                  ([SubExp]
    -> BuilderT
         GPU
         (State VNameSource)
         (Body (Rep (BuilderT GPU (State VNameSource))))
    [SubExp] -> BuilderT GPU (State VNameSource) (Body GPU)
    forall (m :: * -> *).
    MonadBuilder m =>
    [SubExp] -> m (Body (Rep m))
    resultBodyM ([SubExp] -> BuilderT GPU (State VNameSource) (Body GPU))
    -> (Exp GPU -> BuilderT GPU (State VNameSource) [SubExp])
    -> Exp GPU
    -> BuilderT GPU (State VNameSource) (Body GPU)
    forall (m :: * -> *) b c a.
    Monad m =>
    (b -> m c) -> (a -> m b) -> a -> m c
    <=< [Char]
    -> Exp (Rep (BuilderT GPU (State VNameSource)))
    -> BuilderT GPU (State VNameSource) [SubExp]
    forall (m :: * -> *).
    MonadBuilder m =>
    [Char] -> Exp (Rep m) -> m [SubExp]
    letTupExp' [Char]
    "foo")
                    (Exp GPU -> BuilderT GPU (State VNameSource) (Body GPU))
    -> BuilderT GPU (State VNameSource) (Exp GPU)
    -> BuilderT GPU (State VNameSource) (Body GPU)
    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
    $
                          if Bool
    epilogue
                            then VName -> TPrimExp Int64 VName
    forall a. a -> TPrimExp Int64 a
    le64 VName
    kk 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
    k 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
    common_dim
                            else TPrimExp Bool VName
    forall v. TPrimExp Bool v
    true -- if in prologue, always compute redomap.
                      )
                      (Bool
    -> VName
    -> (VName -> VName -> BuilderT GPU (State VNameSource) VName,
        VName -> VName -> BuilderT GPU (State VNameSource) VName)
    -> (VName, VName)
    -> BuilderT GPU (State VNameSource) (Body GPU)
    mkCompLoopRxRy Bool
    fits_ij VName
    acc_merge (VName -> VName -> VName -> BuilderT GPU (State VNameSource) VName
    a_idx_fn VName
    k, VName -> VName -> VName -> BuilderT GPU (State VNameSource) VName
    b_idx_fn VName
    k) (VName
    ltid_y, VName
    ltid_x))
                      ([SubExp]
    -> BuilderT
         GPU
         (State VNameSource)
         (Body (Rep (BuilderT GPU (State VNameSource))))
    forall (m :: * -> *).
    MonadBuilder m =>
    [SubExp] -> m (Body (Rep m))
    resultBodyM [VName -> SubExp
    Var VName
    acc_merge])

                pure [varRes css]
            pure $ head redomap_res
          --
          copyGlb2ShMem ::
            Bool ->
            VName ->
            (VName, VName, PrimType, SubExp, VName, Stm GPU, VName) ->
            Builder GPU (VName, VName -> VName -> VName -> Builder GPU VName)
          copyGlb2ShMem :: Bool
    -> VName
    -> (VName, VName, PrimType, SubExp, VName, Stm GPU, VName)
    -> Builder
         GPU
         (VName,
          VName -> VName -> VName -> BuilderT GPU (State VNameSource) VName)
    copyGlb2ShMem Bool
    is_B VName
    kk (VName
    gtid, VName
    ii, PrimType
    ptp_X_el, SubExp
    parlen_X, VName
    inp_X, Stm GPU
    load_X, VName
    x_loc_init') = do
            let (SubExp
    t_par, SubExp
    r_par, SubExp
    tseq_div_tpar) = (SubExp
    tx, SubExp
    rx, SubExp
    tk_div_tx)
                is_inner_coal :: Bool
    is_inner_coal = Env -> VName -> Stm GPU -> Bool
    isInnerCoal Env
    env VName
    inp_X Stm GPU
    load_X
                str_A :: [Char]
    str_A = VName -> [Char]
    baseString VName
    inp_X
            x_loc <-
              [Char]
    -> VName
    -> [SubExp]
    -> (SubExp, SubExp)
    -> ([VName] -> (VName, VName) -> Builder GPU (SubExp, SubExp))
    -> BuilderT GPU (State VNameSource) VName
    segScatter2D ([Char]
    str_A [Char] -> [Char] -> [Char]
    forall a. [a] -> [a] -> [a]
    ++ [Char]
    "_glb2loc") VName
    x_loc_init' [SubExp
    r_par, SubExp
    tseq_div_tpar] (SubExp
    t_par, SubExp
    t_par) (([VName] -> (VName, VName) -> Builder GPU (SubExp, SubExp))
     -> BuilderT GPU (State VNameSource) VName)
    -> ([VName] -> (VName, VName) -> Builder GPU (SubExp, SubExp))
    -> BuilderT GPU (State VNameSource) VName
    forall a b. (a -> b) -> a -> b
    $
                Bool -> [VName] -> (VName, VName) -> Builder GPU (SubExp, SubExp)
    scatterFun Bool
    is_inner_coal
            pure (x_loc, indexLocMem is_inner_coal str_A x_loc)
            where
              indexLocMem ::
                Bool ->
                String ->
                VName ->
                VName ->
                VName ->
                VName ->
                Builder GPU VName
              indexLocMem :: Bool
    -> [Char]
    -> VName
    -> VName
    -> VName
    -> VName
    -> BuilderT GPU (State VNameSource) VName
    indexLocMem Bool
    is_inner_coal [Char]
    str_A VName
    x_loc VName
    k VName
    ltid_yx VName
    ij = do
                let (SubExp
    r_par, SubExp
    t_seq, SubExp
    tr_par) = (SubExp
    rx, SubExp
    tk, SubExp
    tx_rx)
                let pad_term :: TPrimExp Int64 VName
    pad_term = if Bool
    is_B then SubExp -> TPrimExp Int64 VName
    pe64 SubExp
    se1 else SubExp -> TPrimExp Int64 VName
    pe64 SubExp
    se0
                x_loc_ind_32 <-
                  [Char]
    -> Exp (Rep (BuilderT GPU (State VNameSource)))
    -> BuilderT GPU (State VNameSource) VName
    forall (m :: * -> *).
    MonadBuilder m =>
    [Char] -> Exp (Rep m) -> m VName
    letExp ([Char]
    str_A [Char] -> [Char] -> [Char]
    forall a. [a] -> [a] -> [a]
    ++ [Char]
    "_loc_ind_64")
                    (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
    =<< 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
                      ( if Bool
    is_inner_coal -- ToDo: check this is correct + turn to i32
                          then VName -> TPrimExp Int64 VName
    forall a. a -> TPrimExp Int64 a
    le64 VName
    k 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_yx TPrimExp Int64 VName
    -> TPrimExp Int64 VName -> TPrimExp Int64 VName
    forall a. Num a => a -> a -> a
    * SubExp -> TPrimExp Int64 VName
    pe64 SubExp
    r_par 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
    ij) TPrimExp Int64 VName
    -> TPrimExp Int64 VName -> TPrimExp Int64 VName
    forall a. Num a => a -> a -> a
    * (SubExp -> TPrimExp Int64 VName
    pe64 SubExp
    t_seq TPrimExp Int64 VName
    -> TPrimExp Int64 VName -> TPrimExp Int64 VName
    forall a. Num a => a -> a -> a
    + TPrimExp Int64 VName
    pad_term)
                          else VName -> TPrimExp Int64 VName
    forall a. a -> TPrimExp Int64 a
    le64 VName
    ij 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_yx TPrimExp Int64 VName
    -> TPrimExp Int64 VName -> TPrimExp Int64 VName
    forall a. Num a => a -> a -> a
    * SubExp -> TPrimExp Int64 VName
    pe64 SubExp
    r_par 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
    k TPrimExp Int64 VName
    -> TPrimExp Int64 VName -> TPrimExp Int64 VName
    forall a. Num a => a -> a -> a
    * SubExp -> TPrimExp Int64 VName
    pe64 SubExp
    tr_par
                      )
                index (str_A ++ "_loc_elem") x_loc [x_loc_ind_32]
              --
              scatterFun ::
                Bool ->
                [VName] ->
                (VName, VName) ->
                Builder GPU (SubExp, SubExp)
              scatterFun :: Bool -> [VName] -> (VName, VName) -> Builder GPU (SubExp, SubExp)
    scatterFun Bool
    is_inner_coal [VName
    i0, VName
    k0] (VName
    thd_y, VName
    thd_x) = do
                let str_A :: [Char]
    str_A = VName -> [Char]
    baseString VName
    inp_X
                    t_seq :: SubExp
    t_seq = SubExp
    tk
                (i, k, epx_loc_fi) <- Bool
    -> Bool
    -> (VName, VName)
    -> (VName, VName)
    -> BuilderT
         GPU (State VNameSource) (VName, VName, TPrimExp Int64 VName)
    mk_ik Bool
    is_B Bool
    is_inner_coal (VName
    thd_y, VName
    thd_x) (VName
    i0, VName
    k0)
                letBindNames [gtid] =<< toExp (le64 ii + le64 i)
                a_seqdim_idx <- letExp (str_A ++ "_seqdim_idx") =<< toExp (le64 kk + le64 k)

                a_elem <-
                  letSubExp (str_A ++ "_elem")
                    =<< eIf
                      ( toExp $
                          le64 gtid
                            .<. pe64 parlen_X
                            .&&. if epilogue
                              then le64 a_seqdim_idx .<. pe64 common_dim
                              else true
                      )
                      ( do
                          addStm load_X
                          res <- index "A_elem" inp_X [a_seqdim_idx]
                          resultBodyM [Var res]
                      )
                      (eBody [eBlank $ Prim ptp_X_el])

                a_loc_ind <-
                  letSubExp (str_A ++ "_loc_ind")
                    =<< eIf
                      (toExp $ le64 k .<. pe64 t_seq)
                      (eBody [toExp epx_loc_fi])
                      (eBody [eSubExp $ intConst Int64 (-1)])
                pure (a_elem, a_loc_ind)
              scatterFun Bool
    _ [VName]
    _ (VName, VName)
    _ = do
                [Char] -> Builder GPU (SubExp, SubExp)
    forall a. HasCallStack => [Char] -> a
    error [Char]
    "Function scatterFun in Shared.hs: 2nd arg should be an array with 2 elements!"

    -- ToDo: we need tx == ty (named t_par), and rx == ry (named r_par)
    --       in order to handle all the cases without transpositions.
    --       additionally, of course, we need that tk is a multiple of t_par.
    mmBlkRegTiling :: Env -> Stm GPU -> TileM (Maybe (Stms GPU, Stm GPU))
    mmBlkRegTiling :: Env -> Stm GPU -> TileM (Maybe (Stms GPU, Stm GPU))
    mmBlkRegTiling Env
    env Stm GPU
    stm = do
      res <- Env -> Stm GPU -> TileM (Maybe (Stms GPU, Stm GPU))
    mmBlkRegTilingAcc Env
    env Stm GPU
    stm
      case res of
        Maybe (Stms GPU, Stm GPU)
    Nothing -> Env -> Stm GPU -> TileM (Maybe (Stms GPU, Stm GPU))
    mmBlkRegTilingNrm Env
    env Stm GPU
    stm
        Maybe (Stms GPU, Stm GPU)
    _ -> Maybe (Stms GPU, Stm GPU) -> TileM (Maybe (Stms GPU, Stm GPU))
    forall a. a -> ReaderT (Scope GPU) (State VNameSource) a
    forall (f :: * -> *) a. Applicative f => a -> f a
    pure Maybe (Stms GPU, Stm GPU)
    res

    mmBlkRegTilingAcc :: Env -> Stm GPU -> TileM (Maybe (Stms GPU, Stm GPU))
    mmBlkRegTilingAcc :: Env -> Stm GPU -> TileM (Maybe (Stms GPU, Stm GPU))
    mmBlkRegTilingAcc Env
    env (Let Pat (LetDec GPU)
    pat StmAux (ExpDec GPU)
    aux (Op (SegOp (SegMap SegThread {} SegSpace
    seg_space [TypeBase Shape NoUniqueness]
    ts KernelBody GPU
    old_kbody))))
      | KernelBody () Stms GPU
    kstms [Returns ResultManifest
    ResultMaySimplify Certs
    cs (Var VName
    res_nm)] <- KernelBody GPU
    old_kbody,
        Certs
    cs Certs -> Certs -> Bool
    forall a. Eq a => a -> a -> Bool
    == Certs
    forall a. Monoid a => a
    mempty,
        -- check kernel has one result of primitive type
        [TypeBase Shape NoUniqueness
    res_tp] <- [TypeBase Shape NoUniqueness]
    ts,
        TypeBase Shape NoUniqueness -> Bool
    forall shape u. TypeBase shape u -> Bool
    isAcc TypeBase Shape NoUniqueness
    res_tp,
        -- we get the global-thread id for the two inner dimensions,
        --   as we are probably going to use it in code generation
        (VName
    gtid_x, SubExp
    width_B) : (VName
    gtid_y, SubExp
    height_A) : [(VName, SubExp)]
    rem_outer_dims_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
    seg_space,
        [(VName, SubExp)]
    rem_outer_dims <- [(VName, SubExp)] -> [(VName, SubExp)]
    forall a. [a] -> [a]
    reverse [(VName, SubExp)]
    rem_outer_dims_rev,
        Just
          ( Stms GPU
    code2',
            (Stm GPU
    load_A, VName
    inp_A, PrimType
    map_t1, Stm GPU
    load_B, VName
    inp_B, PrimType
    map_t2),
            SubExp
    common_dim,
            [Int]
    var_dims,
            (Lambda GPU
    map_lam, Lambda GPU
    red_lam, SubExp
    red_ne, VName
    redomap_orig_res, PrimType
    red_t)
            ) <-
          SegSpace
    -> Stms GPU
    -> Maybe
         (Stms GPU, (Stm GPU, VName, PrimType, Stm GPU, VName, PrimType),
          SubExp, [Int], (Lambda GPU, Lambda GPU, SubExp, VName, PrimType))
    matchesBlkRegTile SegSpace
    seg_space Stms GPU
    kstms,
        VName -> Stms GPU -> VName -> Bool
    forall {rep}. VName -> Stms rep -> VName -> Bool
    checkAccumulatesRedomapRes VName
    res_nm Stms GPU
    code2' VName
    redomap_orig_res = do
          -- Here we start the implementation --
          let is_B_coal :: Bool
    is_B_coal = Env -> VName -> Stm GPU -> Bool
    isInnerCoal Env
    env VName
    inp_B Stm GPU
    load_B
          ---- in this binder: host code and outer seggroup (ie. the new kernel) ----
          (new_kernel, host_stms) <- Builder GPU (Stm GPU)
    -> ReaderT (Scope GPU) (State VNameSource) (Stm GPU, 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 (Stm GPU)
     -> ReaderT (Scope GPU) (State VNameSource) (Stm GPU, Stms GPU))
    -> Builder GPU (Stm GPU)
    -> ReaderT (Scope GPU) (State VNameSource) (Stm GPU, Stms GPU)
    forall a b. (a -> b) -> a -> b
    $ do
            -- host code
            (rx, ry, tx, ty, tk, tk_div_tx, tk_div_ty, tx_rx, ty_ry, a_loc_sz, b_loc_sz) <-
              SubExp
    -> SubExp
    -> SubExp
    -> Bool
    -> Builder
         GPU
         (SubExp, SubExp, SubExp, SubExp, SubExp, SubExp, SubExp, SubExp,
          SubExp, SubExp, SubExp)
    mkTileMemSizes SubExp
    height_A SubExp
    width_B SubExp
    common_dim Bool
    is_B_coal

            rk <- letSubExp "rk" $ BasicOp $ SubExp $ intConst Int64 8 -- 16 and 8 seem good values
            tk_rk <- letSubExp "tk_rk" =<< toExp (pe64 tk * pe64 rk)

            gridDim_t <- letSubExp "gridDim_t" =<< ceilDiv common_dim tk_rk
            gridDim_y <- letSubExp "gridDim_y" =<< ceilDiv height_A ty_ry
            gridDim_x <- letSubExp "gridDim_x" =<< ceilDiv width_B tx_rx

            let gridxyt_pexp = SubExp -> TPrimExp Int64 VName
    pe64 SubExp
    gridDim_y TPrimExp Int64 VName
    -> TPrimExp Int64 VName -> TPrimExp Int64 VName
    forall a. Num a => a -> a -> a
    * SubExp -> TPrimExp Int64 VName
    pe64 SubExp
    gridDim_x TPrimExp Int64 VName
    -> TPrimExp Int64 VName -> TPrimExp Int64 VName
    forall a. Num a => a -> a -> a
    * SubExp -> TPrimExp Int64 VName
    pe64 SubExp
    gridDim_t
                grid_pexp =
                  (TPrimExp Int64 VName -> SubExp -> TPrimExp Int64 VName)
    -> TPrimExp Int64 VName -> [SubExp] -> TPrimExp Int64 VName
    forall b a. (b -> a -> b) -> b -> [a] -> b
    forall (t :: * -> *) b a.
    Foldable t =>
    (b -> a -> b) -> b -> t a -> b
    foldl (\TPrimExp Int64 VName
    x SubExp
    d -> SubExp -> TPrimExp Int64 VName
    pe64 SubExp
    d TPrimExp Int64 VName
    -> TPrimExp Int64 VName -> TPrimExp Int64 VName
    forall a. Num a => a -> a -> a
    * TPrimExp Int64 VName
    x) TPrimExp Int64 VName
    gridxyt_pexp ([SubExp] -> TPrimExp Int64 VName)
    -> [SubExp] -> TPrimExp Int64 VName
    forall a b. (a -> b) -> a -> b
    $
                    ((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)]
    rem_outer_dims_rev

            (grid_size, tblock_size, segthd_lvl) <- mkNewSegthdLvl tx ty grid_pexp
            (gid_x, gid_y, gid_flat) <- mkGidsXYF
            gid_t <- newVName "gid_t"

            ---- in this binder: outer seggroup ----
            (ret_seggroup, stms_seggroup) <- runBuilder $ do
              iii <- letExp "iii" =<< toExp (le64 gid_y * pe64 ty_ry)
              jjj <- letExp "jjj" =<< toExp (le64 gid_x * pe64 tx_rx)
              ttt <- letExp "ttt" =<< toExp (le64 gid_t * pe64 tk_rk)

              -- initialize register mem with neutral elements and create shmem
              (cssss, a_loc_init, b_loc_init) <-
                initRegShmem
                  (rx, tx, ry, ty, a_loc_sz, b_loc_sz)
                  (map_t1, map_t2, red_t)
                  segthd_lvl
                  red_ne

              -- build prologue.
              elems_on_t <- letSubExp "elems_on_t" =<< toExp (pe64 common_dim - le64 ttt)
              tiles_on_t <- letSubExp "tiles_on_t" $ BasicOp $ BinOp (SQuot Int64 Unsafe) elems_on_t tk
              full_tiles <- letExp "full_tiles" $ BasicOp $ BinOp (SMin Int64) rk tiles_on_t

              let ct_arg =
                    ( (SubExp
    rx, SubExp
    ry, SubExp
    tx, SubExp
    ty, SubExp
    tk, SubExp
    tk_div_tx, SubExp
    tk_div_ty, SubExp
    tx_rx),
                      SegLevel
    segthd_lvl,
                      [Int]
    var_dims,
                      (VName
    gtid_x, SubExp
    width_B, VName
    gtid_y, SubExp
    height_A, SubExp
    common_dim),
                      (VName
    iii, VName
    jjj),
                      (Stm GPU
    load_A, VName
    inp_A, PrimType
    map_t1, Stm GPU
    load_B, VName
    inp_B, PrimType
    map_t2),
                      (Lambda GPU
    map_lam, Lambda GPU
    red_lam)
                    )

              prologue_res_list <-
                forLoop' (Var full_tiles) [cssss, a_loc_init, b_loc_init] $
                  \VName
    kk0 [VName
    thd_res_merge, VName
    a_loc_merge, VName
    b_loc_merge] -> do
                    off_t <- [Char]
    -> Exp (Rep (BuilderT GPU (State VNameSource)))
    -> BuilderT GPU (State VNameSource) VName
    forall (m :: * -> *).
    MonadBuilder m =>
    [Char] -> Exp (Rep m) -> m VName
    letExp [Char]
    "off_t" (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
    =<< 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
    rk 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
    gid_t 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
    kk0)
                    process_full_tiles <-
                      kkLoopBody env ct_arg off_t (thd_res_merge, a_loc_merge, b_loc_merge) False

                    resultBodyM $ map Var process_full_tiles

              let prologue_res : a_loc_reuse : b_loc_reuse : _ = prologue_res_list

              redomap_res_lst <-
                letTupExp "redomap_res_if"
                  =<< eIf
                    ( toExp $
                        le64 full_tiles
                          .==. pe64 rk
                          .||. pe64 common_dim
                          .==. (pe64 tk * le64 full_tiles + le64 ttt)
                    )
                    (resultBodyM $ map Var prologue_res_list)
                    ( do
                        off_t <- letExp "off_t" =<< toExp (pe64 rk * le64 gid_t + le64 full_tiles)
                        process_sprs_tile <-
                          kkLoopBody env ct_arg off_t (prologue_res, a_loc_reuse, b_loc_reuse) True

                        resultBodyM $ map Var process_sprs_tile
                    )
              let redomap_res : _ = redomap_res_lst

              -- support for non-empty code2'
              --  segmap (ltid_y < ty, ltid_x < tx) {
              --    for i < ry do
              --      for j < rx do
              --        res = if (iii+ltid_y*ry+i < height_A && jjj+ltid_x*rx+j < width_B)
              --              then code2' else dummy
              --        final_res[i,j] = res
              mkEpilogueAccRes
                segthd_lvl
                (redomap_orig_res, redomap_res)
                (res_nm, res_tp)
                (ty, tx, ry, rx)
                (iii, jjj)
                (gtid_y, gtid_x)
                (height_A, width_B, rem_outer_dims)
                code2'

            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
    grid_size) (SubExp -> Count BlockSize SubExp
    forall {k} (u :: k) e. e -> Count u e
    Count SubExp
    tblock_size)
                level' = SegVirt -> Maybe KernelGrid -> SegLevel
    SegBlock SegVirt
    SegNoVirt (KernelGrid -> Maybe KernelGrid
    forall a. a -> Maybe a
    Just KernelGrid
    grid)
                space' = VName -> [(VName, SubExp)] -> SegSpace
    SegSpace VName
    gid_flat ([(VName, SubExp)]
    rem_outer_dims [(VName, SubExp)] -> [(VName, SubExp)] -> [(VName, SubExp)]
    forall a. [a] -> [a] -> [a]
    ++ [(VName
    gid_t, SubExp
    gridDim_t), (VName
    gid_y, SubExp
    gridDim_y), (VName
    gid_x, SubExp
    gridDim_x)])
                kbody' = BodyDec GPU -> Stms GPU -> [KernelResult] -> KernelBody GPU
    forall rep.
    BodyDec rep -> Stms rep -> [KernelResult] -> KernelBody rep
    KernelBody () Stms GPU
    stms_seggroup [KernelResult]
    ret_seggroup
            pure $ Let pat aux $ Op $ SegOp $ SegMap level' space' ts kbody'
          pure $ Just (host_stms, new_kernel)
      where
        sameAccType :: VName -> TypeBase shape u -> Bool
    sameAccType VName
    acc_sglton (Acc VName
    sglton Shape
    _ [TypeBase Shape NoUniqueness]
    _ u
    _) =
          VName
    acc_sglton VName -> VName -> Bool
    forall a. Eq a => a -> a -> Bool
    == VName
    sglton
        sameAccType VName
    _ TypeBase shape u
    _ = Bool
    False
        getAccumFV :: TypeBase Shape NoUniqueness
    -> BuilderT GPU (State VNameSource) VName
    getAccumFV (Acc VName
    singleton Shape
    _shp [TypeBase Shape NoUniqueness
    _eltp] NoUniqueness
    _) = do
          let fvs :: [VName]
    fvs = Names -> [VName]
    namesToList (Names -> [VName]) -> Names -> [VName]
    forall a b. (a -> b) -> a -> b
    $ KernelBody GPU -> Names
    forall a. FreeIn a => a -> Names
    freeIn KernelBody GPU
    old_kbody -- code
          tps <- Scope GPU
    -> BuilderT GPU (State VNameSource) [TypeBase Shape NoUniqueness]
    -> BuilderT GPU (State VNameSource) [TypeBase Shape NoUniqueness]
    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 (SegSpace -> Scope GPU
    forall rep. SegSpace -> Scope rep
    scopeOfSegSpace SegSpace
    seg_space) (BuilderT GPU (State VNameSource) [TypeBase Shape NoUniqueness]
     -> BuilderT GPU (State VNameSource) [TypeBase Shape NoUniqueness])
    -> BuilderT GPU (State VNameSource) [TypeBase Shape NoUniqueness]
    -> BuilderT GPU (State VNameSource) [TypeBase Shape NoUniqueness]
    forall a b. (a -> b) -> a -> b
    $ do
            (VName
     -> BuilderT GPU (State VNameSource) (TypeBase Shape NoUniqueness))
    -> [VName]
    -> BuilderT GPU (State VNameSource) [TypeBase Shape NoUniqueness]
    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) (TypeBase Shape NoUniqueness)
    forall rep (m :: * -> *).
    HasScope rep m =>
    VName -> m (TypeBase Shape NoUniqueness)
    lookupType [VName]
    fvs
          let (acc_0s, _) = unzip $ filter (sameAccType singleton . snd) $ zip fvs tps
          case acc_0s of
            [VName
    acc_0] -> VName -> BuilderT GPU (State VNameSource) VName
    forall a. a -> BuilderT GPU (State VNameSource) a
    forall (f :: * -> *) a. Applicative f => a -> f a
    pure VName
    acc_0
            [VName]
    _ -> [Char] -> BuilderT GPU (State VNameSource) VName
    forall a. HasCallStack => [Char] -> a
    error [Char]
    "Impossible case reached when treating accumulators!"
        getAccumFV TypeBase Shape NoUniqueness
    tp = [Char] -> BuilderT GPU (State VNameSource) VName
    forall a. HasCallStack => [Char] -> a
    error ([Char]
    "Should be an accumulator type at this point, given: " [Char] -> [Char] -> [Char]
    forall a. [a] -> [a] -> [a]
    ++ TypeBase Shape NoUniqueness -> [Char]
    forall a. Pretty a => a -> [Char]
    prettyString TypeBase Shape NoUniqueness
    tp)
        --
        -- checks that the redomap result is used directly as the accumulated value,
        -- in which case it is safe to parallelize the innermost dimension (of tile tk)
        checkAccumulatesRedomapRes :: VName -> Stms rep -> VName -> Bool
    checkAccumulatesRedomapRes VName
    res_nm Stms rep
    acc_code VName
    redomap_orig_res = do
          (Bool -> Stm rep -> Bool) -> Bool -> [Stm rep] -> Bool
    forall b a. (b -> a -> b) -> b -> [a] -> b
    forall (t :: * -> *) b a.
    Foldable t =>
    (b -> a -> b) -> b -> t a -> b
    foldl Bool -> Stm rep -> Bool
    getAccumStm Bool
    False ([Stm rep] -> Bool) -> [Stm rep] -> Bool
    forall a b. (a -> b) -> a -> b
    $ [Stm rep] -> [Stm rep]
    forall a. [a] -> [a]
    reverse ([Stm rep] -> [Stm rep]) -> [Stm rep] -> [Stm rep]
    forall a b. (a -> b) -> a -> b
    $ Stms rep -> [Stm rep]
    forall rep. Stms rep -> [Stm rep]
    stmsToList Stms rep
    acc_code
          where
            getAccumStm :: Bool -> Stm rep -> Bool
    getAccumStm Bool
    True Stm rep
    _ = Bool
    True
            getAccumStm Bool
    False (Let (Pat [PatElem (LetDec rep)
    pat_el]) StmAux (ExpDec rep)
    _aux (BasicOp (UpdateAcc Safety
    _ VName
    _acc_nm [SubExp]
    _ind [SubExp]
    vals)))
              | [SubExp
    v] <- [SubExp]
    vals,
                PatElem (LetDec rep) -> VName
    forall dec. PatElem dec -> VName
    patElemName PatElem (LetDec rep)
    pat_el VName -> VName -> Bool
    forall a. Eq a => a -> a -> Bool
    == VName
    res_nm =
                  SubExp
    v SubExp -> SubExp -> Bool
    forall a. Eq a => a -> a -> Bool
    == VName -> SubExp
    Var VName
    redomap_orig_res
            getAccumStm Bool
    False Stm rep
    _ = Bool
    False
        --
        -- epilogue for accumulator result type
        mkEpilogueAccRes :: SegLevel
    -> (VName, VName)
    -> (VName, TypeBase Shape NoUniqueness)
    -> (SubExp, SubExp, SubExp, SubExp)
    -> (VName, VName)
    -> (VName, VName)
    -> (SubExp, SubExp, [(VName, SubExp)])
    -> Stms GPU
    -> Builder GPU [KernelResult]
    mkEpilogueAccRes
          SegLevel
    segthd_lvl
          (VName
    redomap_orig_res, VName
    redomap_res)
          (VName
    res_nm, TypeBase Shape NoUniqueness
    res_tp)
          (SubExp
    ty, SubExp
    tx, SubExp
    ry, SubExp
    rx)
          (VName
    iii, VName
    jjj)
          (VName
    gtid_y, VName
    gtid_x)
          (SubExp
    height_A, SubExp
    width_B, [(VName, SubExp)]
    _rem_outer_dims)
          Stms GPU
    code2' = do
            rss_init <- TypeBase Shape NoUniqueness
    -> BuilderT GPU (State VNameSource) VName
    getAccumFV TypeBase Shape NoUniqueness
    res_tp
            rssss_list <- segMap2D "rssss" segthd_lvl ResultMaySimplify (ty, tx) $ \(VName
    ltid_y, VName
    ltid_x) -> do
              (css, ii, jj) <- (SubExp, SubExp)
    -> (VName, VName)
    -> (VName, VName, VName)
    -> Builder GPU (VName, VName, VName)
    getThdRedomapRes (SubExp
    rx, SubExp
    ry) (VName
    ltid_x, VName
    ltid_y) (VName
    iii, VName
    jjj, VName
    redomap_res)
              rss <- forLoop ry [rss_init] $ \VName
    i [VName
    rss_merge] -> do
                rss' <- SubExp
    -> [VName]
    -> (VName
        -> [VName] -> BuilderT GPU (State VNameSource) (Body GPU))
    -> BuilderT GPU (State VNameSource) VName
    forLoop SubExp
    rx [VName
    rss_merge] ((VName -> [VName] -> BuilderT GPU (State VNameSource) (Body GPU))
     -> BuilderT GPU (State VNameSource) VName)
    -> (VName
        -> [VName] -> BuilderT GPU (State VNameSource) (Body GPU))
    -> BuilderT GPU (State VNameSource) VName
    forall a b. (a -> b) -> a -> b
    $ \VName
    j [VName
    rss_merge'] -> do
                  (VName, VName)
    -> (VName, VName, VName, VName)
    -> (VName, VName)
    -> BuilderT GPU (State VNameSource) ()
    prereqAddCode2 (VName
    gtid_x, VName
    gtid_y) (VName
    ii, VName
    i, VName
    jj, VName
    j) (VName
    css, VName
    redomap_orig_res)
                  let code2_subs :: Stms GPU
    code2_subs = Map VName VName -> Stms GPU -> Stms GPU
    forall a. Substitute a => Map VName VName -> a -> a
    substituteNames (VName -> VName -> Map VName VName
    forall k a. k -> a -> Map k a
    M.singleton VName
    rss_init VName
    rss_merge') Stms GPU
    code2'

                  res_el <-
                    [Char]
    -> Exp (Rep (BuilderT GPU (State VNameSource)))
    -> BuilderT GPU (State VNameSource) SubExp
    forall (m :: * -> *).
    MonadBuilder m =>
    [Char] -> Exp (Rep m) -> m SubExp
    letSubExp [Char]
    "res_elem"
                      (Exp GPU -> BuilderT GPU (State VNameSource) SubExp)
    -> BuilderT GPU (State VNameSource) (Exp GPU)
    -> BuilderT GPU (State VNameSource) SubExp
    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
    $
                            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
    height_A
                              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_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
    width_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
    code2_subs
                            [SubExp]
    -> BuilderT
         GPU
         (State VNameSource)
         (Body (Rep (BuilderT GPU (State VNameSource))))
    forall (m :: * -> *).
    MonadBuilder m =>
    [SubExp] -> m (Body (Rep m))
    resultBodyM [VName -> SubExp
    Var VName
    res_nm]
                        )
                        ([SubExp]
    -> BuilderT
         GPU
         (State VNameSource)
         (Body (Rep (BuilderT GPU (State VNameSource))))
    forall (m :: * -> *).
    MonadBuilder m =>
    [SubExp] -> m (Body (Rep m))
    resultBodyM [VName -> SubExp
    Var VName
    rss_merge'])
                  resultBodyM [res_el]
                resultBodyM [Var rss']
              pure [varRes rss]
            let epilogue_res_acc : _ = rssss_list
            pure [Returns ResultMaySimplify (Certs []) $ Var epilogue_res_acc]
    mmBlkRegTilingAcc Env
    _ Stm GPU
    _ = Maybe (Stms GPU, Stm GPU) -> TileM (Maybe (Stms GPU, Stm GPU))
    forall a. a -> ReaderT (Scope GPU) (State VNameSource) a
    forall (f :: * -> *) a. Applicative f => a -> f a
    pure Maybe (Stms GPU, Stm GPU)
    forall a. Maybe a
    Nothing

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

    mmBlkRegTilingNrm :: Env -> Stm GPU -> TileM (Maybe (Stms GPU, Stm GPU))
    mmBlkRegTilingNrm :: Env -> Stm GPU -> TileM (Maybe (Stms GPU, Stm GPU))
    mmBlkRegTilingNrm Env
    env (Let Pat (LetDec GPU)
    pat StmAux (ExpDec GPU)
    aux (Op (SegOp (SegMap SegThread {} SegSpace
    seg_space [TypeBase Shape NoUniqueness]
    ts KernelBody GPU
    old_kbody))))
      | KernelBody () Stms GPU
    kstms [Returns ResultManifest
    ResultMaySimplify Certs
    cs (Var VName
    res_nm)] <- KernelBody GPU
    old_kbody,
        Certs
    cs Certs -> Certs -> Bool
    forall a. Eq a => a -> a -> Bool
    == Certs
    forall a. Monoid a => a
    mempty,
        -- check kernel has one result of primitive type
        [TypeBase Shape NoUniqueness
    res_tp] <- [TypeBase Shape NoUniqueness]
    ts,
        TypeBase Shape NoUniqueness -> Bool
    forall shape u. TypeBase shape u -> Bool
    primType TypeBase Shape NoUniqueness
    res_tp,
        -- we get the global-thread id for the two inner dimensions,
        --   as we are probably going to use it in code generation
        (VName
    gtid_x, SubExp
    width_B) : (VName
    gtid_y, SubExp
    height_A) : [(VName, SubExp)]
    rem_outer_dims_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
    seg_space,
        [(VName, SubExp)]
    rem_outer_dims <- [(VName, SubExp)] -> [(VName, SubExp)]
    forall a. [a] -> [a]
    reverse [(VName, SubExp)]
    rem_outer_dims_rev,
        Just
          ( Stms GPU
    code2',
            (Stm GPU
    load_A, VName
    inp_A, PrimType
    map_t1, Stm GPU
    load_B, VName
    inp_B, PrimType
    map_t2),
            SubExp
    common_dim,
            [Int]
    var_dims,
            (Lambda GPU
    map_lam, Lambda GPU
    red_lam, SubExp
    red_ne, VName
    redomap_orig_res, PrimType
    red_t)
            ) <-
          SegSpace
    -> Stms GPU
    -> Maybe
         (Stms GPU, (Stm GPU, VName, PrimType, Stm GPU, VName, PrimType),
          SubExp, [Int], (Lambda GPU, Lambda GPU, SubExp, VName, PrimType))
    matchesBlkRegTile SegSpace
    seg_space Stms GPU
    kstms = do
          -- Here we start the implementation
          let is_B_coal :: Bool
    is_B_coal = Env -> VName -> Stm GPU -> Bool
    isInnerCoal Env
    env VName
    inp_B Stm GPU
    load_B
          ---- in this binder: host code and outer seggroup (ie. the new kernel) ----
          (new_kernel, host_stms) <- Builder GPU (Stm GPU)
    -> ReaderT (Scope GPU) (State VNameSource) (Stm GPU, 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 (Stm GPU)
     -> ReaderT (Scope GPU) (State VNameSource) (Stm GPU, Stms GPU))
    -> Builder GPU (Stm GPU)
    -> ReaderT (Scope GPU) (State VNameSource) (Stm GPU, Stms GPU)
    forall a b. (a -> b) -> a -> b
    $ do
            -- host code
            (rx, ry, tx, ty, tk, tk_div_tx, tk_div_ty, tx_rx, ty_ry, a_loc_sz, b_loc_sz) <-
              SubExp
    -> SubExp
    -> SubExp
    -> Bool
    -> Builder
         GPU
         (SubExp, SubExp, SubExp, SubExp, SubExp, SubExp, SubExp, SubExp,
          SubExp, SubExp, SubExp)
    mkTileMemSizes SubExp
    height_A SubExp
    width_B SubExp
    common_dim Bool
    is_B_coal

            gridDim_x <- letSubExp "gridDim_x" =<< ceilDiv width_B tx_rx
            gridDim_y <- letSubExp "gridDim_y" =<< ceilDiv height_A ty_ry
            let gridxy_pexp = SubExp -> TPrimExp Int64 VName
    pe64 SubExp
    gridDim_y TPrimExp Int64 VName
    -> TPrimExp Int64 VName -> TPrimExp Int64 VName
    forall a. Num a => a -> a -> a
    * SubExp -> TPrimExp Int64 VName
    pe64 SubExp
    gridDim_x
            let grid_pexp =
                  (TPrimExp Int64 VName -> SubExp -> TPrimExp Int64 VName)
    -> TPrimExp Int64 VName -> [SubExp] -> TPrimExp Int64 VName
    forall b a. (b -> a -> b) -> b -> [a] -> b
    forall (t :: * -> *) b a.
    Foldable t =>
    (b -> a -> b) -> b -> t a -> b
    foldl (\TPrimExp Int64 VName
    x SubExp
    d -> SubExp -> TPrimExp Int64 VName
    pe64 SubExp
    d TPrimExp Int64 VName
    -> TPrimExp Int64 VName -> TPrimExp Int64 VName
    forall a. Num a => a -> a -> a
    * TPrimExp Int64 VName
    x) TPrimExp Int64 VName
    gridxy_pexp ([SubExp] -> TPrimExp Int64 VName)
    -> [SubExp] -> TPrimExp Int64 VName
    forall a b. (a -> b) -> a -> b
    $
                    ((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)]
    rem_outer_dims_rev
            (grid_size, tblock_size, segthd_lvl) <- mkNewSegthdLvl tx ty grid_pexp

            (gid_x, gid_y, gid_flat) <- mkGidsXYF

            ---- in this binder: outer seggroup ----
            (ret_seggroup, stms_seggroup) <- runBuilder $ do
              iii <- letExp "iii" =<< toExp (le64 gid_y * pe64 ty_ry)
              jjj <- letExp "jjj" =<< toExp (le64 gid_x * pe64 tx_rx)

              -- initialize register mem with neutral elements and create shmem
              (cssss, a_loc_init, b_loc_init) <-
                initRegShmem
                  (rx, tx, ry, ty, a_loc_sz, b_loc_sz)
                  (map_t1, map_t2, red_t)
                  segthd_lvl
                  red_ne

              -- build prologue.
              full_tiles <-
                letExp "full_tiles" $
                  BasicOp $
                    BinOp (SQuot Int64 Unsafe) common_dim tk

              let ct_arg =
                    ( (SubExp
    rx, SubExp
    ry, SubExp
    tx, SubExp
    ty, SubExp
    tk, SubExp
    tk_div_tx, SubExp
    tk_div_ty, SubExp
    tx_rx),
                      SegLevel
    segthd_lvl,
                      [Int]
    var_dims,
                      (VName
    gtid_x, SubExp
    width_B, VName
    gtid_y, SubExp
    height_A, SubExp
    common_dim),
                      (VName
    iii, VName
    jjj),
                      (Stm GPU
    load_A, VName
    inp_A, PrimType
    map_t1, Stm GPU
    load_B, VName
    inp_B, PrimType
    map_t2),
                      (Lambda GPU
    map_lam, Lambda GPU
    red_lam)
                    )

              prologue_res_list <-
                forLoop' (Var full_tiles) [cssss, a_loc_init, b_loc_init] $
                  \VName
    kk0 [VName
    thd_res_merge, VName
    a_loc_merge, VName
    b_loc_merge] -> do
                    process_full_tiles <-
                      Env
    -> ((SubExp, SubExp, SubExp, SubExp, SubExp, SubExp, SubExp,
         SubExp),
        SegLevel, [Int], (VName, SubExp, VName, SubExp, SubExp),
        (VName, VName),
        (Stm GPU, VName, PrimType, Stm GPU, VName, PrimType),
        (Lambda GPU, Lambda GPU))
    -> VName
    -> (VName, VName, VName)
    -> Bool
    -> Builder GPU [VName]
    kkLoopBody Env
    env ((SubExp, SubExp, SubExp, SubExp, SubExp, SubExp, SubExp, SubExp),
     SegLevel, [Int], (VName, SubExp, VName, SubExp, SubExp),
     (VName, VName),
     (Stm GPU, VName, PrimType, Stm GPU, VName, PrimType),
     (Lambda GPU, Lambda GPU))
    ct_arg VName
    kk0 (VName
    thd_res_merge, VName
    a_loc_merge, VName
    b_loc_merge) Bool
    False

                    resultBodyM $ map Var process_full_tiles

              let prologue_res : a_loc_reuse : b_loc_reuse : _ = prologue_res_list

              -- build epilogue.
              epilogue_res_list <- kkLoopBody env ct_arg full_tiles (prologue_res, a_loc_reuse, b_loc_reuse) True

              let redomap_res : _ = epilogue_res_list

              -- support for non-empty code2'
              --  segmap (ltid_y < ty, ltid_x < tx) {
              --    for i < ry do
              --      for j < rx do
              --        res = if (iii+ltid_y*ry+i < height_A && jjj+ltid_x*rx+j < width_B)
              --              then code2' else dummy
              --        final_res[i,j] = res
              mkEpiloguePrimRes
                segthd_lvl
                (redomap_orig_res, redomap_res)
                (res_nm, res_tp)
                (ty, tx, ry, rx)
                (iii, jjj)
                (gtid_y, gtid_x)
                (height_A, width_B, rem_outer_dims)
                code2'

            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
    grid_size) (SubExp -> Count BlockSize SubExp
    forall {k} (u :: k) e. e -> Count u e
    Count SubExp
    tblock_size)
                level' = SegVirt -> Maybe KernelGrid -> SegLevel
    SegBlock SegVirt
    SegNoVirt (KernelGrid -> Maybe KernelGrid
    forall a. a -> Maybe a
    Just KernelGrid
    grid)
                space' = VName -> [(VName, SubExp)] -> SegSpace
    SegSpace VName
    gid_flat ([(VName, SubExp)]
    rem_outer_dims [(VName, SubExp)] -> [(VName, SubExp)] -> [(VName, SubExp)]
    forall a. [a] -> [a] -> [a]
    ++ [(VName
    gid_y, SubExp
    gridDim_y), (VName
    gid_x, SubExp
    gridDim_x)])
                kbody' = BodyDec GPU -> Stms GPU -> [KernelResult] -> KernelBody GPU
    forall rep.
    BodyDec rep -> Stms rep -> [KernelResult] -> KernelBody rep
    KernelBody () Stms GPU
    stms_seggroup [KernelResult]
    ret_seggroup
            pure $ Let pat aux $ Op $ SegOp $ SegMap level' space' ts kbody'
          pure $ Just (host_stms, new_kernel)
      where
        mkEpiloguePrimRes :: SegLevel
    -> (VName, VName)
    -> (VName, TypeBase Shape NoUniqueness)
    -> (SubExp, SubExp, SubExp, SubExp)
    -> (VName, VName)
    -> (VName, VName)
    -> (SubExp, SubExp, [(a, SubExp)])
    -> Stms GPU
    -> Builder GPU [KernelResult]
    mkEpiloguePrimRes
          SegLevel
    segthd_lvl
          (VName
    redomap_orig_res, VName
    redomap_res)
          (VName
    res_nm, TypeBase Shape NoUniqueness
    res_tp)
          (SubExp
    ty, SubExp
    tx, SubExp
    ry, SubExp
    rx)
          (VName
    iii, VName
    jjj)
          (VName
    gtid_y, VName
    gtid_x)
          (SubExp
    height_A, SubExp
    width_B, [(a, SubExp)]
    rem_outer_dims)
          Stms GPU
    code2' = do
            epilogue_res <-
              if VName
    redomap_orig_res VName -> VName -> Bool
    forall a. Eq a => a -> a -> Bool
    == VName
    res_nm
                then VName -> BuilderT GPU (State VNameSource) VName
    forall a. a -> BuilderT GPU (State VNameSource) a
    forall (f :: * -> *) a. Applicative f => a -> f a
    pure VName
    redomap_res -- epilogue_res_list
                else do
                  rssss_list <- [Char]
    -> SegLevel
    -> ResultManifest
    -> (SubExp, SubExp)
    -> ((VName, VName) -> BuilderT GPU (State VNameSource) [SubExpRes])
    -> Builder GPU [VName]
    segMap2D [Char]
    "rssss" SegLevel
    segthd_lvl ResultManifest
    ResultPrivate (SubExp
    ty, SubExp
    tx) (((VName, VName) -> BuilderT GPU (State VNameSource) [SubExpRes])
     -> Builder GPU [VName])
    -> ((VName, VName) -> BuilderT GPU (State VNameSource) [SubExpRes])
    -> Builder GPU [VName]
    forall a b. (a -> b) -> a -> b
    $ \(VName
    ltid_y, VName
    ltid_x) -> do
                    rss_init <- [Char]
    -> PrimType -> [SubExp] -> BuilderT GPU (State VNameSource) VName
    forall (m :: * -> *).
    MonadBuilder m =>
    [Char] -> PrimType -> [SubExp] -> m VName
    scratch [Char]
    "rss_init" (TypeBase Shape NoUniqueness -> PrimType
    forall shape u. TypeBase shape u -> PrimType
    elemType TypeBase Shape NoUniqueness
    res_tp) [SubExp
    ry, SubExp
    rx]
                    (css, ii, jj) <- getThdRedomapRes (rx, ry) (ltid_x, ltid_y) (iii, jjj, redomap_res)
                    rss <- forLoop ry [rss_init] $ \VName
    i [VName
    rss_merge] -> do
                      rss' <- SubExp
    -> [VName]
    -> (VName
        -> [VName] -> BuilderT GPU (State VNameSource) (Body GPU))
    -> BuilderT GPU (State VNameSource) VName
    forLoop SubExp
    rx [VName
    rss_merge] ((VName -> [VName] -> BuilderT GPU (State VNameSource) (Body GPU))
     -> BuilderT GPU (State VNameSource) VName)
    -> (VName
        -> [VName] -> BuilderT GPU (State VNameSource) (Body GPU))
    -> BuilderT GPU (State VNameSource) VName
    forall a b. (a -> b) -> a -> b
    $ \VName
    j [VName
    rss_merge'] -> do
                        (VName, VName)
    -> (VName, VName, VName, VName)
    -> (VName, VName)
    -> BuilderT GPU (State VNameSource) ()
    prereqAddCode2 (VName
    gtid_x, VName
    gtid_y) (VName
    ii, VName
    i, VName
    jj, VName
    j) (VName
    css, VName
    redomap_orig_res)

                        res_el <-
                          [Char]
    -> Exp (Rep (BuilderT GPU (State VNameSource)))
    -> BuilderT GPU (State VNameSource) SubExp
    forall (m :: * -> *).
    MonadBuilder m =>
    [Char] -> Exp (Rep m) -> m SubExp
    letSubExp [Char]
    "res_elem"
                            (Exp GPU -> BuilderT GPU (State VNameSource) SubExp)
    -> BuilderT GPU (State VNameSource) (Exp GPU)
    -> BuilderT GPU (State VNameSource) SubExp
    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
    $
                                  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
    height_A
                                    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_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
    width_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
    code2'
                                  [SubExp]
    -> BuilderT
         GPU
         (State VNameSource)
         (Body (Rep (BuilderT GPU (State VNameSource))))
    forall (m :: * -> *).
    MonadBuilder m =>
    [SubExp] -> m (Body (Rep m))
    resultBodyM [VName -> SubExp
    Var VName
    res_nm]
                              )
                              ([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 [TypeBase Shape NoUniqueness
    -> BuilderT
         GPU
         (State VNameSource)
         (Exp (Rep (BuilderT GPU (State VNameSource))))
    forall (m :: * -> *).
    MonadBuilder m =>
    TypeBase Shape NoUniqueness -> m (Exp (Rep m))
    eBlank TypeBase Shape NoUniqueness
    res_tp])
                        rss'' <- update "rss" rss_merge' [i, j] res_el
                        resultBodyM [Var rss'']
                      resultBodyM [Var rss']
                    pure [varRes rss]
                  let rssss : _ = rssss_list
                  pure rssss

            let regtile_ret_dims =
                  ((a, SubExp) -> (SubExp, SubExp, SubExp))
    -> [(a, SubExp)] -> [(SubExp, SubExp, SubExp)]
    forall a b. (a -> b) -> [a] -> [b]
    map (\(a
    _, SubExp
    sz) -> (SubExp
    sz, SubExp
    se1, SubExp
    se1)) [(a, SubExp)]
    rem_outer_dims
                    [(SubExp, SubExp, SubExp)]
    -> [(SubExp, SubExp, SubExp)] -> [(SubExp, SubExp, SubExp)]
    forall a. [a] -> [a] -> [a]
    ++ [(SubExp
    height_A, SubExp
    ty, SubExp
    ry), (SubExp
    width_B, SubExp
    tx, SubExp
    rx)]

            -- Add dummy dimensions to tile to reflect the outer dimensions.
            epilogue_res' <-
              if null rem_outer_dims
                then pure epilogue_res
                else do
                  epilogue_t <- lookupType epilogue_res
                  let (block_dims, rest_dims) = splitAt 2 $ arrayDims epilogue_t
                      ones = ((a, SubExp) -> SubExp) -> [(a, SubExp)] -> [SubExp]
    forall a b. (a -> b) -> [a] -> [b]
    map (SubExp -> (a, SubExp) -> SubExp
    forall a b. a -> b -> a
    const (SubExp -> (a, SubExp) -> SubExp)
    -> SubExp -> (a, SubExp) -> SubExp
    forall a b. (a -> b) -> a -> b
    $ IntType -> Integer -> SubExp
    intConst IntType
    Int64 Integer
    1) [(a, SubExp)]
    rem_outer_dims
                      new_shape = [SubExp] -> Shape
    forall d. [d] -> ShapeBase d
    Shape ([SubExp] -> Shape) -> [SubExp] -> Shape
    forall a b. (a -> b) -> a -> b
    $ [[SubExp]] -> [SubExp]
    forall (t :: * -> *) a. Foldable t => t [a] -> [a]
    concat [[SubExp]
    ones, [SubExp]
    block_dims, [SubExp]
    ones, [SubExp]
    rest_dims]
                  letExp "res_reshaped" . BasicOp $
                    Reshape epilogue_res (reshapeAll (arrayShape epilogue_t) new_shape)
            pure [RegTileReturns mempty regtile_ret_dims epilogue_res']
    mmBlkRegTilingNrm Env
    _ Stm GPU
    _ = Maybe (Stms GPU, Stm GPU) -> TileM (Maybe (Stms GPU, Stm GPU))
    forall a. a -> ReaderT (Scope GPU) (State VNameSource) a
    forall (f :: * -> *) a. Applicative f => a -> f a
    pure Maybe (Stms GPU, Stm GPU)
    forall a. Maybe a
    Nothing

    -- pattern match the properties of the code that we look to
    -- tile: a redomap whose two input arrays are each invariant
    -- to one of the last two (innermost) parallel dimensions.
    matchesBlkRegTile ::
      SegSpace ->
      Stms GPU ->
      Maybe
        ( Stms GPU,
          (Stm GPU, VName, PrimType, Stm GPU, VName, PrimType),
          SubExp,
          [Int],
          (Lambda GPU, Lambda GPU, SubExp, VName, PrimType)
        )
    matchesBlkRegTile :: SegSpace
    -> Stms GPU
    -> Maybe
         (Stms GPU, (Stm GPU, VName, PrimType, Stm GPU, VName, PrimType),
          SubExp, [Int], (Lambda GPU, Lambda GPU, SubExp, VName, PrimType))
    matchesBlkRegTile SegSpace
    seg_space Stms GPU
    kstms
      | -- build the variance table, that records, for
        -- each variable name, the variables it depends on
        Map VName Names
    initial_variance <- (NameInfo (ZonkAny 0) -> Names)
    -> Map VName (NameInfo (ZonkAny 0)) -> Map VName Names
    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)) -> Map VName Names)
    -> Map VName (NameInfo (ZonkAny 0)) -> Map VName Names
    forall a b. (a -> b) -> a -> b
    $ SegSpace -> Map VName (NameInfo (ZonkAny 0))
    forall rep. SegSpace -> Scope rep
    scopeOfSegSpace SegSpace
    seg_space,
        Map VName Names
    variance <- Map VName Names -> Stms GPU -> Map VName Names
    varianceInStms Map VName Names
    initial_variance Stms GPU
    kstms,
        -- check that the code fits the pattern having:
        -- some `code1`, followed by one Screma SOAC, followed by some `code2`
        (Stms GPU
    code1, Just Stm GPU
    screma_stmt, Stms GPU
    code2) <- Stms GPU -> (Stms GPU, Maybe (Stm GPU), Stms GPU)
    matchCodeStreamCode Stms GPU
    kstms,
        Let Pat (LetDec GPU)
    pat_redomap StmAux (ExpDec GPU)
    _ (Op Op GPU
    _) <- Stm GPU
    screma_stmt,
        -- checks that the Screma SOAC is actually a redomap and normalizes it
        Just (SubExp
    common_dim, [VName]
    arrs, (Commutativity
    _, Lambda GPU
    red_lam, [SubExp]
    red_nes, Lambda GPU
    map_lam)) <- Stm GPU
    -> Maybe
         (SubExp, [VName],
          (Commutativity, Lambda GPU, [SubExp], Lambda GPU))
    isTileableRedomap Stm GPU
    screma_stmt,
        -- check that exactly two 1D arrays are streamed thorugh redomap,
        -- and the result of redomap is one scalar
        -- !!!I need to rearrange this whole thing!!! including inp_A and inp_B
        [VName] -> Int
    forall a. [a] -> Int
    forall (t :: * -> *) a. Foldable t => t a -> Int
    length [VName]
    arrs Int -> Int -> Bool
    forall a. Eq a => a -> a -> Bool
    == Int
    2,
        [SubExp
    red_ne] <- [SubExp]
    red_nes,
        [TypeBase Shape NoUniqueness
    map_t1t, TypeBase Shape NoUniqueness
    map_t2t] <- (Param (TypeBase Shape NoUniqueness)
     -> TypeBase Shape NoUniqueness)
    -> [Param (TypeBase Shape NoUniqueness)]
    -> [TypeBase Shape NoUniqueness]
    forall a b. (a -> b) -> [a] -> [b]
    map Param (TypeBase Shape NoUniqueness) -> TypeBase Shape NoUniqueness
    forall dec. Param dec -> dec
    paramDec ([Param (TypeBase Shape NoUniqueness)]
     -> [TypeBase Shape NoUniqueness])
    -> [Param (TypeBase Shape NoUniqueness)]
    -> [TypeBase Shape NoUniqueness]
    forall a b. (a -> b) -> a -> b
    $ Lambda GPU -> [LParam GPU]
    forall rep. Lambda rep -> [LParam rep]
    lambdaParams Lambda GPU
    map_lam,
        [TypeBase Shape NoUniqueness
    red_t1, TypeBase Shape NoUniqueness
    _] <- (Param (TypeBase Shape NoUniqueness)
     -> TypeBase Shape NoUniqueness)
    -> [Param (TypeBase Shape NoUniqueness)]
    -> [TypeBase Shape NoUniqueness]
    forall a b. (a -> b) -> [a] -> [b]
    map Param (TypeBase Shape NoUniqueness) -> TypeBase Shape NoUniqueness
    forall dec. Param dec -> dec
    paramDec ([Param (TypeBase Shape NoUniqueness)]
     -> [TypeBase Shape NoUniqueness])
    -> [Param (TypeBase Shape NoUniqueness)]
    -> [TypeBase Shape NoUniqueness]
    forall a b. (a -> b) -> a -> b
    $ Lambda GPU -> [LParam GPU]
    forall rep. Lambda rep -> [LParam rep]
    lambdaParams Lambda GPU
    red_lam,
        TypeBase Shape NoUniqueness -> Bool
    forall shape u. TypeBase shape u -> Bool
    primType TypeBase Shape NoUniqueness
    map_t1t Bool -> Bool -> Bool
    && TypeBase Shape NoUniqueness -> Bool
    forall shape u. TypeBase shape u -> Bool
    primType TypeBase Shape NoUniqueness
    map_t2t Bool -> Bool -> Bool
    && TypeBase Shape NoUniqueness -> Bool
    forall shape u. TypeBase shape u -> Bool
    primType TypeBase Shape NoUniqueness
    red_t1,
        PrimType
    map_t1_0 <- TypeBase Shape NoUniqueness -> PrimType
    forall shape u. TypeBase shape u -> PrimType
    elemType TypeBase Shape NoUniqueness
    map_t1t,
        PrimType
    map_t2_0 <- TypeBase Shape NoUniqueness -> PrimType
    forall shape u. TypeBase shape u -> PrimType
    elemType TypeBase Shape NoUniqueness
    map_t2t,
        -- checks that the input arrays to redomap are variant to
        -- exactly one of the two innermost dimensions of the kernel
        Just [Int]
    var_dims <- Names -> SegSpace -> Map VName Names -> [VName] -> Maybe [Int]
    isInvarTo1of2InnerDims Names
    forall a. Monoid a => a
    mempty SegSpace
    seg_space Map VName Names
    variance [VName]
    arrs,
        -- get the variables on which the first result of redomap depends on
        [VName
    redomap_orig_res] <- Pat (TypeBase Shape NoUniqueness) -> [VName]
    forall dec. Pat dec -> [VName]
    patNames Pat (TypeBase Shape NoUniqueness)
    Pat (LetDec GPU)
    pat_redomap,
        Just Names
    res_red_var <- VName -> Map VName Names -> Maybe Names
    forall k a. Ord k => k -> Map k a -> Maybe a
    M.lookup VName
    redomap_orig_res Map VName Names
    variance, -- variance of the reduce result

        -- we furthermore check that code1 is only formed by
        -- 1. statements that slice some globally-declared arrays
        --    to produce the input for the redomap, and
        -- 2. potentially some statements on which the redomap
        --    is independent; these are recorded in `code2''`
        Just (Stms GPU
    code2'', Map VName (Stm GPU)
    tab_inv_stm) <-
          (Maybe (Stms GPU, Map VName (Stm GPU))
     -> Stm GPU -> Maybe (Stms GPU, Map VName (Stm GPU)))
    -> Maybe (Stms GPU, Map VName (Stm GPU))
    -> Stms GPU
    -> Maybe (Stms GPU, Map VName (Stm GPU))
    forall b a. (b -> a -> b) -> b -> Seq a -> b
    forall (t :: * -> *) b a.
    Foldable t =>
    (b -> a -> b) -> b -> t a -> b
    foldl
            (Names
    -> Names
    -> Maybe (Stms GPU, Map VName (Stm GPU))
    -> Stm GPU
    -> Maybe (Stms GPU, Map VName (Stm GPU))
    processIndirections ([VName] -> Names
    namesFromList [VName]
    arrs) Names
    res_red_var)
            ((Stms GPU, Map VName (Stm GPU))
    -> Maybe (Stms GPU, Map VName (Stm GPU))
    forall a. a -> Maybe a
    Just (Stms GPU
    forall a. Seq a
    Seq.empty, Map VName (Stm GPU)
    forall k a. Map k a
    M.empty))
            Stms GPU
    code1,
        -- identify load_A, load_B
        [Stm GPU]
    tmp_stms <- (VName -> Maybe (Stm GPU)) -> [VName] -> [Stm GPU]
    forall a b. (a -> Maybe b) -> [a] -> [b]
    mapMaybe (VName -> Map VName (Stm GPU) -> Maybe (Stm GPU)
    forall k a. Ord k => k -> Map k a -> Maybe a
    `M.lookup` Map VName (Stm GPU)
    tab_inv_stm) [VName]
    arrs,
        [Stm GPU] -> Int
    forall a. [a] -> Int
    forall (t :: * -> *) a. Foldable t => t a -> Int
    length [Stm GPU]
    tmp_stms Int -> Int -> Bool
    forall a. Eq a => a -> a -> Bool
    == [VName] -> Int
    forall a. [a] -> Int
    forall (t :: * -> *) a. Foldable t => t a -> Int
    length [VName]
    arrs =
          let zip_AB :: [(Stm GPU, VName, PrimType)]
    zip_AB = [Stm GPU] -> [VName] -> [PrimType] -> [(Stm GPU, VName, PrimType)]
    forall a b c. [a] -> [b] -> [c] -> [(a, b, c)]
    zip3 [Stm GPU]
    tmp_stms [VName]
    arrs [PrimType
    map_t1_0, PrimType
    map_t2_0]
              [(Stm GPU
    load_A, VName
    inp_A, PrimType
    map_t1), (Stm GPU
    load_B, VName
    inp_B, PrimType
    map_t2)] =
                if [Int]
    var_dims [Int] -> [Int] -> Bool
    forall a. Eq a => a -> a -> Bool
    == [Int
    0, Int
    1]
                  then [(Stm GPU, VName, PrimType)]
    zip_AB
                  else [(Stm GPU, VName, PrimType)] -> [(Stm GPU, VName, PrimType)]
    forall a. [a] -> [a]
    reverse [(Stm GPU, VName, PrimType)]
    zip_AB
              code2' :: Stms GPU
    code2' = Stms GPU
    code2'' Stms GPU -> Stms GPU -> Stms GPU
    forall a. Semigroup a => a -> a -> a
    <> Stms GPU
    code2
           in (Stms GPU, (Stm GPU, VName, PrimType, Stm GPU, VName, PrimType),
     SubExp, [Int], (Lambda GPU, Lambda GPU, SubExp, VName, PrimType))
    -> Maybe
         (Stms GPU, (Stm GPU, VName, PrimType, Stm GPU, VName, PrimType),
          SubExp, [Int], (Lambda GPU, Lambda GPU, SubExp, VName, PrimType))
    forall a. a -> Maybe a
    Just
                ( Stms GPU
    code2',
                  (Stm GPU
    load_A, VName
    inp_A, PrimType
    map_t1, Stm GPU
    load_B, VName
    inp_B, PrimType
    map_t2),
                  SubExp
    common_dim,
                  [Int]
    var_dims,
                  (Lambda GPU
    map_lam, Lambda GPU
    red_lam, SubExp
    red_ne, VName
    redomap_orig_res, TypeBase Shape NoUniqueness -> PrimType
    forall shape u. TypeBase shape u -> PrimType
    elemType TypeBase Shape NoUniqueness
    red_t1)
                )
    matchesBlkRegTile SegSpace
    _ Stms GPU
    _ = Maybe
      (Stms GPU, (Stm GPU, VName, PrimType, Stm GPU, VName, PrimType),
       SubExp, [Int], (Lambda GPU, Lambda GPU, SubExp, VName, PrimType))
    forall a. Maybe a
    Nothing

    -- ceiled division expression
    ceilDiv :: (MonadBuilder m) => SubExp -> SubExp -> m (Exp (Rep m))
    ceilDiv :: forall (m :: * -> *).
    MonadBuilder m =>
    SubExp -> SubExp -> m (Exp (Rep m))
    ceilDiv SubExp
    x SubExp
    y = Exp (Rep m) -> m (Exp (Rep m))
    forall a. a -> m a
    forall (f :: * -> *) a. Applicative f => a -> f a
    pure (Exp (Rep m) -> m (Exp (Rep m))) -> Exp (Rep m) -> m (Exp (Rep m))
    forall a b. (a -> b) -> a -> b
    $ BasicOp -> Exp (Rep m)
    forall rep. BasicOp -> Exp rep
    BasicOp (BasicOp -> Exp (Rep m)) -> BasicOp -> Exp (Rep m)
    forall a b. (a -> b) -> a -> b
    $ BinOp -> SubExp -> SubExp -> BasicOp
    BinOp (IntType -> Safety -> BinOp
    SDivUp IntType
    Int64 Safety
    Unsafe) SubExp
    x SubExp
    y

    mkTileMemSizes ::
      SubExp ->
      SubExp ->
      SubExp ->
      Bool ->
      Builder
        GPU
        ( SubExp,
          SubExp,
          SubExp,
          SubExp,
          SubExp,
          SubExp,
          SubExp,
          SubExp,
          SubExp,
          SubExp,
          SubExp
        )
    mkTileMemSizes :: SubExp
    -> SubExp
    -> SubExp
    -> Bool
    -> Builder
         GPU
         (SubExp, SubExp, SubExp, SubExp, SubExp, SubExp, SubExp, SubExp,
          SubExp, SubExp, SubExp)
    mkTileMemSizes SubExp
    height_A SubExp
    _width_B SubExp
    common_dim Bool
    is_B_not_transp = do
      tk_name <- [Char] -> Name
    nameFromString ([Char] -> Name) -> (VName -> [Char]) -> VName -> Name
    forall b c a. (b -> c) -> (a -> b) -> a -> c
    . VName -> [Char]
    forall a. Pretty a => a -> [Char]
    prettyString (VName -> Name)
    -> BuilderT GPU (State VNameSource) VName
    -> BuilderT GPU (State VNameSource) Name
    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 [Char]
    "Tk"
      ty_name <- nameFromString . prettyString <$> newVName "Ty"
      ry_name <- nameFromString . prettyString <$> newVName "Ry"

      -- until we change the copying to use lmads we need to
      --   guarantee that Tx=Ty AND Rx = Ry AND Tx | Tk
      -- for matrix multiplication it would be safe if they aren't
      --   but not for any of the other three cases!
      (ty, ry) <- getParTiles ("Ty", "Ry") (ty_name, ry_name) height_A
      let (tx, rx) = (ty, ry)
      tk <- getSeqTile "Tk" tk_name common_dim tx ty

      tk_div_tx <- letSubExp "tk_div_tx" =<< ceilDiv tk tx
      tk_div_ty <- letSubExp "tk_div_ty" =<< ceilDiv tk ty

      tx_rx <- letSubExp "TxRx" =<< toExp (pe64 tx * pe64 rx)
      ty_ry <- letSubExp "TyRy" =<< toExp (pe64 ty * pe64 ry)

      -- let pad_term = sMax64 (pe64 tk) (pe64 ty * pe64 ry)
      let pad_term =
            if Bool
    is_B_not_transp
              then SubExp -> TPrimExp Int64 VName
    pe64 SubExp
    ty TPrimExp Int64 VName
    -> TPrimExp Int64 VName -> TPrimExp Int64 VName
    forall a. Num a => a -> a -> a
    * SubExp -> TPrimExp Int64 VName
    pe64 SubExp
    ry
              else SubExp -> TPrimExp Int64 VName
    pe64 SubExp
    se0
      a_loc_sz <-
        letSubExp "a_loc_sz"
          =<< toExp (pe64 ty * pe64 ry * pe64 tk)
      -- if B is transposed, its shmem should be [tk][tx*rx]
      -- we pad as above, by assuming tx*rx == ty*ry >= tk
      b_loc_sz <-
        letSubExp "b_loc_sz"
          =<< toExp (pe64 tx * pe64 rx * pe64 tk + pad_term)
      pure (rx, ry, tx, ty, tk, tk_div_tx, tk_div_ty, tx_rx, ty_ry, a_loc_sz, b_loc_sz)

    mkNewSegthdLvl ::
      SubExp ->
      SubExp ->
      TPrimExp Int64 VName ->
      Builder GPU (SubExp, SubExp, SegLevel)
    mkNewSegthdLvl :: SubExp
    -> SubExp
    -> TPrimExp Int64 VName
    -> Builder GPU (SubExp, SubExp, SegLevel)
    mkNewSegthdLvl SubExp
    tx SubExp
    ty TPrimExp Int64 VName
    grid_pexp = do
      grid_size <- [Char]
    -> Exp (Rep (BuilderT GPU (State VNameSource)))
    -> BuilderT GPU (State VNameSource) SubExp
    forall (m :: * -> *).
    MonadBuilder m =>
    [Char] -> Exp (Rep m) -> m SubExp
    letSubExp [Char]
    "grid_size" (Exp GPU -> BuilderT GPU (State VNameSource) SubExp)
    -> BuilderT GPU (State VNameSource) (Exp GPU)
    -> BuilderT GPU (State VNameSource) 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 TPrimExp Int64 VName
    grid_pexp
      tblock_size <- letSubExp "tblock_size" =<< toExp (pe64 ty * pe64 tx)
      let segthd_lvl = SegVirt -> SegLevel
    SegThreadInBlock (SegSeqDims -> SegVirt
    SegNoVirtFull ([Int] -> SegSeqDims
    SegSeqDims []))
      pure (grid_size, tblock_size, segthd_lvl)

    mkGidsXYF :: Builder GPU (VName, VName, VName)
    mkGidsXYF :: Builder GPU (VName, VName, VName)
    mkGidsXYF = do
      gid_y <- [Char] -> BuilderT GPU (State VNameSource) VName
    forall (m :: * -> *). MonadFreshNames m => [Char] -> m VName
    newVName [Char]
    "gid_y"
      gid_x <- newVName "gid_x"
      gid_flat <- newVName "gid_flat"
      pure (gid_x, gid_y, gid_flat)

    initRegShmem ::
      (SubExp, SubExp, SubExp, SubExp, SubExp, SubExp) ->
      (PrimType, PrimType, PrimType) ->
      SegLevel ->
      SubExp ->
      Builder GPU (VName, VName, VName)
    initRegShmem :: (SubExp, SubExp, SubExp, SubExp, SubExp, SubExp)
    -> (PrimType, PrimType, PrimType)
    -> SegLevel
    -> SubExp
    -> Builder GPU (VName, VName, VName)
    initRegShmem
      (SubExp
    rx, SubExp
    tx, SubExp
    ry, SubExp
    ty, SubExp
    a_loc_sz, SubExp
    b_loc_sz)
      (PrimType
    map_t1, PrimType
    map_t2, PrimType
    red_t)
      SegLevel
    segthd_lvl
      SubExp
    red_ne = do
        -- initialize register mem with neutral elements.
        cssss_list <- [Char]
    -> SegLevel
    -> ResultManifest
    -> (SubExp, SubExp)
    -> ((VName, VName) -> BuilderT GPU (State VNameSource) [SubExpRes])
    -> Builder GPU [VName]
    segMap2D [Char]
    "cssss" SegLevel
    segthd_lvl ResultManifest
    ResultPrivate (SubExp
    ty, SubExp
    tx) (((VName, VName) -> BuilderT GPU (State VNameSource) [SubExpRes])
     -> Builder GPU [VName])
    -> ((VName, VName) -> BuilderT GPU (State VNameSource) [SubExpRes])
    -> Builder GPU [VName]
    forall a b. (a -> b) -> a -> b
    $ \(VName, VName)
    _ -> do
          css_init <- [Char]
    -> PrimType -> [SubExp] -> BuilderT GPU (State VNameSource) VName
    forall (m :: * -> *).
    MonadBuilder m =>
    [Char] -> PrimType -> [SubExp] -> m VName
    scratch [Char]
    "css_init" PrimType
    red_t [SubExp
    ry, SubExp
    rx]
          css <- forLoop ry [css_init] $ \VName
    i [VName
    css_merge] -> do
            css' <- SubExp
    -> [VName]
    -> (VName
        -> [VName] -> BuilderT GPU (State VNameSource) (Body GPU))
    -> BuilderT GPU (State VNameSource) VName
    forLoop SubExp
    rx [VName
    css_merge] ((VName -> [VName] -> BuilderT GPU (State VNameSource) (Body GPU))
     -> BuilderT GPU (State VNameSource) VName)
    -> (VName
        -> [VName] -> BuilderT GPU (State VNameSource) (Body GPU))
    -> BuilderT GPU (State VNameSource) VName
    forall a b. (a -> b) -> a -> b
    $ \VName
    j [VName
    css_merge'] -> do
              css'' <- [Char]
    -> VName
    -> [VName]
    -> SubExp
    -> BuilderT GPU (State VNameSource) VName
    forall (m :: * -> *).
    MonadBuilder m =>
    [Char] -> VName -> [VName] -> SubExp -> m VName
    update [Char]
    "css" VName
    css_merge' [VName
    i, VName
    j] SubExp
    red_ne
              resultBodyM [Var css'']
            resultBodyM [Var css']
          pure [varRes css]
        let [cssss] = cssss_list
        -- scratch shared memory
        a_loc_init <- scratch "A_loc" map_t1 [a_loc_sz]
        b_loc_init <- scratch "B_loc" map_t2 [b_loc_sz]
        pure (cssss, a_loc_init, b_loc_init)

    getThdRedomapRes ::
      (SubExp, SubExp) ->
      (VName, VName) ->
      (VName, VName, VName) ->
      Builder GPU (VName, VName, VName)
    getThdRedomapRes :: (SubExp, SubExp)
    -> (VName, VName)
    -> (VName, VName, VName)
    -> Builder GPU (VName, VName, VName)
    getThdRedomapRes (SubExp
    rx, SubExp
    ry) (VName
    ltid_x, VName
    ltid_y) (VName
    iii, VName
    jjj, VName
    redomap_res) = do
      css <- [Char]
    -> VName -> [VName] -> BuilderT GPU (State VNameSource) VName
    forall (m :: * -> *).
    MonadBuilder m =>
    [Char] -> VName -> [VName] -> m VName
    index [Char]
    "redomap_thd" VName
    redomap_res [VName
    ltid_y, VName
    ltid_x]
      ii <- letExp "ii" =<< toExp (le64 iii + le64 ltid_y * pe64 ry)
      jj <- letExp "jj" =<< toExp (le64 jjj + le64 ltid_x * pe64 rx)
      pure (css, ii, jj)

    prereqAddCode2 ::
      (VName, VName) ->
      (VName, VName, VName, VName) ->
      (VName, VName) ->
      Builder GPU ()
    prereqAddCode2 :: (VName, VName)
    -> (VName, VName, VName, VName)
    -> (VName, VName)
    -> BuilderT GPU (State VNameSource) ()
    prereqAddCode2 (VName
    gtid_x, VName
    gtid_y) (VName
    ii, VName
    i, VName
    jj, VName
    j) (VName
    css, VName
    redomap_orig_res) = do
      c <- [Char]
    -> VName -> [VName] -> BuilderT GPU (State VNameSource) VName
    forall (m :: * -> *).
    MonadBuilder m =>
    [Char] -> VName -> [VName] -> m VName
    index [Char]
    "redomap_elm" VName
    css [VName
    i, VName
    j]
      cpy_stm <- mkLetNamesM [redomap_orig_res] $ BasicOp $ SubExp $ Var c
      addStm cpy_stm
      letBindNames [gtid_y] =<< toExp (le64 ii + le64 i)
      letBindNames [gtid_x] =<< toExp (le64 jj + le64 j)

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

    -- | Checks that all streamed arrays are variant to exacly one of
    --   the two innermost parallel dimensions, and conversely, for
    --   each of the two innermost parallel dimensions, there is at
    --   least one streamed array variant to it. The result is the
    --   number of the only variant parallel dimension for each array.
    isInvarTo1of2InnerDims ::
      Names ->
      SegSpace ->
      VarianceTable ->
      [VName] ->
      Maybe [Int]
    isInvarTo1of2InnerDims :: Names -> SegSpace -> Map VName Names -> [VName] -> Maybe [Int]
    isInvarTo1of2InnerDims Names
    branch_variant SegSpace
    kspace Map VName Names
    variance [VName]
    arrs =
      let inner_perm0 :: [Maybe Int]
    inner_perm0 = (VName -> Maybe Int) -> [VName] -> [Maybe Int]
    forall a b. (a -> b) -> [a] -> [b]
    map VName -> Maybe Int
    varToOnly1of2InnerDims [VName]
    arrs
          inner_perm :: [Int]
    inner_perm = [Maybe Int] -> [Int]
    forall a. [Maybe a] -> [a]
    catMaybes [Maybe Int]
    inner_perm0
          ok1 :: Bool
    ok1 = Int -> [Int] -> Bool
    forall a. Eq a => a -> [a] -> Bool
    forall (t :: * -> *) a. (Foldable t, Eq a) => a -> t a -> Bool
    elem Int
    0 [Int]
    inner_perm Bool -> Bool -> Bool
    && Int -> [Int] -> Bool
    forall a. Eq a => a -> [a] -> Bool
    forall (t :: * -> *) a. (Foldable t, Eq a) => a -> t a -> Bool
    elem Int
    1 [Int]
    inner_perm
          ok2 :: Bool
    ok2 = [Maybe Int] -> Int
    forall a. [a] -> Int
    forall (t :: * -> *) a. Foldable t => t a -> Int
    length [Maybe Int]
    inner_perm0 Int -> Int -> Bool
    forall a. Eq a => a -> a -> Bool
    == [Int] -> Int
    forall a. [a] -> Int
    forall (t :: * -> *) a. Foldable t => t a -> Int
    length [Int]
    inner_perm
       in if Bool
    ok1 Bool -> Bool -> Bool
    && Bool
    ok2 then [Int] -> Maybe [Int]
    forall a. a -> Maybe a
    Just [Int]
    inner_perm else Maybe [Int]
    forall a. Maybe a
    Nothing
      where
        varToOnly1of2InnerDims :: VName -> Maybe Int
        varToOnly1of2InnerDims :: VName -> Maybe Int
    varToOnly1of2InnerDims VName
    arr = do
          (j, _) : (i, _) : _ <- [(VName, SubExp)] -> Maybe [(VName, SubExp)]
    forall a. a -> Maybe a
    Just ([(VName, SubExp)] -> Maybe [(VName, SubExp)])
    -> [(VName, SubExp)] -> Maybe [(VName, SubExp)]
    forall a b. (a -> b) -> a -> b
    $ [(VName, SubExp)] -> [(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
    kspace
          let variant_to = Names -> VName -> Map VName Names -> Names
    forall k a. Ord k => a -> k -> Map k a -> a
    M.findWithDefault Names
    forall a. Monoid a => a
    mempty VName
    arr Map VName Names
    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 not branch_invariant
            then Nothing -- if i or j in branch_variant; return nothing
            else
              if nameIn i variant_to && j `notNameIn` variant_to
                then Just 0
                else
                  if nameIn j variant_to && i `notNameIn` variant_to
                    then Just 1
                    else Nothing

    processIndirections ::
      Names -> -- input arrays to redomap
      Names -> -- variables on which the result of redomap depends on.
      Maybe (Stms GPU, M.Map VName (Stm GPU)) ->
      Stm GPU ->
      Maybe (Stms GPU, M.Map VName (Stm GPU))
    processIndirections :: Names
    -> Names
    -> Maybe (Stms GPU, Map VName (Stm GPU))
    -> Stm GPU
    -> Maybe (Stms GPU, Map VName (Stm GPU))
    processIndirections Names
    arrs Names
    _ Maybe (Stms GPU, Map VName (Stm GPU))
    acc stm :: Stm GPU
    stm@(Let Pat (LetDec GPU)
    patt StmAux (ExpDec GPU)
    _ (BasicOp (Index VName
    _ Slice SubExp
    _)))
      | Just (Stms GPU
    ss, Map VName (Stm GPU)
    tab) <- Maybe (Stms GPU, Map VName (Stm GPU))
    acc,
        [PatElem (TypeBase Shape NoUniqueness)
    p] <- Pat (TypeBase Shape NoUniqueness)
    -> [PatElem (TypeBase Shape NoUniqueness)]
    forall dec. Pat dec -> [PatElem dec]
    patElems Pat (TypeBase Shape NoUniqueness)
    Pat (LetDec GPU)
    patt,
        VName
    p_nm <- PatElem (TypeBase Shape NoUniqueness) -> VName
    forall dec. PatElem dec -> VName
    patElemName PatElem (TypeBase Shape NoUniqueness)
    p,
        VName
    p_nm VName -> Names -> Bool
    `nameIn` Names
    arrs =
          (Stms GPU, Map VName (Stm GPU))
    -> Maybe (Stms GPU, Map VName (Stm GPU))
    forall a. a -> Maybe a
    Just (Stms GPU
    ss, VName -> Stm GPU -> Map VName (Stm GPU) -> Map VName (Stm GPU)
    forall k a. Ord k => k -> a -> Map k a -> Map k a
    M.insert VName
    p_nm Stm GPU
    stm Map VName (Stm GPU)
    tab)
    processIndirections Names
    _ Names
    res_red_var Maybe (Stms GPU, Map VName (Stm GPU))
    acc stm' :: Stm GPU
    stm'@(Let Pat (LetDec GPU)
    patt StmAux (ExpDec GPU)
    _ Exp GPU
    _)
      | Just (Stms GPU
    ss, Map VName (Stm GPU)
    tab) <- Maybe (Stms GPU, Map VName (Stm GPU))
    acc,
        [PatElem (TypeBase Shape NoUniqueness)]
    ps <- Pat (TypeBase Shape NoUniqueness)
    -> [PatElem (TypeBase Shape NoUniqueness)]
    forall dec. Pat dec -> [PatElem dec]
    patElems Pat (TypeBase Shape NoUniqueness)
    Pat (LetDec GPU)
    patt,
        (PatElem (TypeBase Shape NoUniqueness) -> Bool)
    -> [PatElem (TypeBase Shape NoUniqueness)] -> Bool
    forall (t :: * -> *) a. Foldable t => (a -> Bool) -> t a -> Bool
    all (\PatElem (TypeBase Shape NoUniqueness)
    p -> PatElem (TypeBase Shape NoUniqueness) -> VName
    forall dec. PatElem dec -> VName
    patElemName PatElem (TypeBase Shape NoUniqueness)
    p VName -> Names -> Bool
    `notNameIn` Names
    res_red_var) [PatElem (TypeBase Shape NoUniqueness)]
    ps =
          (Stms GPU, Map VName (Stm GPU))
    -> Maybe (Stms GPU, Map VName (Stm GPU))
    forall a. a -> Maybe a
    Just (Stms GPU
    ss Stms GPU -> Stm GPU -> Stms GPU
    forall a. Seq a -> a -> Seq a
    Seq.|> Stm GPU
    stm', Map VName (Stm GPU)
    tab)
      | Bool
    otherwise = Maybe (Stms GPU, Map VName (Stm GPU))
    forall a. Maybe a
    Nothing

    getParTiles :: (String, String) -> (Name, Name) -> SubExp -> Builder GPU (SubExp, SubExp)
    getParTiles :: ([Char], [Char])
    -> (Name, Name) -> SubExp -> Builder GPU (SubExp, SubExp)
    getParTiles ([Char]
    t_str, [Char]
    r_str) (Name
    t_name, Name
    r_name) SubExp
    len_dim =
      case SubExp
    len_dim of
        Constant (IntValue (Int64Value Int64
    8)) ->
          (SubExp, SubExp) -> Builder GPU (SubExp, SubExp)
    forall a. a -> BuilderT GPU (State VNameSource) a
    forall (f :: * -> *) a. Applicative f => a -> f a
    pure (SubExp
    se8, SubExp
    se1)
        Constant (IntValue (Int64Value Int64
    16)) ->
          (SubExp, SubExp) -> Builder GPU (SubExp, SubExp)
    forall a. a -> BuilderT GPU (State VNameSource) a
    forall (f :: * -> *) a. Applicative f => a -> f a
    pure (SubExp
    se8, SubExp
    se2)
        Constant (IntValue (Int64Value Int64
    32)) ->
          (SubExp, SubExp) -> Builder GPU (SubExp, SubExp)
    forall a. a -> BuilderT GPU (State VNameSource) a
    forall (f :: * -> *) a. Applicative f => a -> f a
    pure (SubExp
    se8, SubExp
    se4)
        SubExp
    _ -> do
          t <- [Char]
    -> Exp (Rep (BuilderT GPU (State VNameSource)))
    -> BuilderT GPU (State VNameSource) SubExp
    forall (m :: * -> *).
    MonadBuilder m =>
    [Char] -> Exp (Rep m) -> m SubExp
    letSubExp [Char]
    t_str (Exp (Rep (BuilderT GPU (State VNameSource)))
     -> BuilderT GPU (State VNameSource) SubExp)
    -> Exp (Rep (BuilderT GPU (State VNameSource)))
    -> BuilderT GPU (State VNameSource) SubExp
    forall a b. (a -> b) -> a -> b
    $ Op (Rep (BuilderT GPU (State VNameSource)))
    -> Exp (Rep (BuilderT GPU (State VNameSource)))
    forall rep. Op rep -> Exp rep
    Op (Op (Rep (BuilderT GPU (State VNameSource)))
     -> Exp (Rep (BuilderT GPU (State VNameSource))))
    -> Op (Rep (BuilderT GPU (State VNameSource)))
    -> Exp (Rep (BuilderT GPU (State VNameSource)))
    forall a b. (a -> b) -> a -> b
    $ SizeOp -> HostOp SOAC GPU
    forall (op :: * -> *) rep. SizeOp -> HostOp op rep
    SizeOp (SizeOp -> HostOp SOAC GPU) -> SizeOp -> HostOp SOAC GPU
    forall a b. (a -> b) -> a -> b
    $ Name -> SizeClass -> SizeOp
    GetSize Name
    t_name SizeClass
    SizeTile
          r <- letSubExp r_str $ Op $ SizeOp $ GetSize r_name SizeRegTile
          pure (t, r)

    getSeqTile :: String -> Name -> SubExp -> SubExp -> SubExp -> Builder GPU SubExp
    getSeqTile :: [Char]
    -> Name
    -> SubExp
    -> SubExp
    -> SubExp
    -> BuilderT GPU (State VNameSource) SubExp
    getSeqTile [Char]
    tk_str Name
    tk_name SubExp
    len_dim SubExp
    tx SubExp
    ty =
      case (SubExp
    tx, SubExp
    ty) of
        (Constant (IntValue (Int64Value Int64
    v_x)), Constant (IntValue (Int64Value Int64
    v_y))) ->
          [Char]
    -> Exp (Rep (BuilderT GPU (State VNameSource)))
    -> BuilderT GPU (State VNameSource) SubExp
    forall (m :: * -> *).
    MonadBuilder m =>
    [Char] -> Exp (Rep m) -> m SubExp
    letSubExp [Char]
    tk_str (Exp GPU -> BuilderT GPU (State VNameSource) SubExp)
    -> (Int64 -> Exp GPU)
    -> Int64
    -> BuilderT GPU (State VNameSource) SubExp
    forall b c a. (b -> c) -> (a -> b) -> a -> c
    . BasicOp -> Exp GPU
    forall rep. BasicOp -> Exp rep
    BasicOp (BasicOp -> Exp GPU) -> (Int64 -> BasicOp) -> Int64 -> Exp GPU
    forall b c a. (b -> c) -> (a -> b) -> a -> c
    . SubExp -> BasicOp
    SubExp (SubExp -> BasicOp) -> (Int64 -> SubExp) -> Int64 -> BasicOp
    forall b c a. (b -> c) -> (a -> b) -> a -> c
    . Int64 -> SubExp
    forall v. IsValue v => v -> SubExp
    constant (Int64 -> BuilderT GPU (State VNameSource) SubExp)
    -> Int64 -> BuilderT GPU (State VNameSource) SubExp
    forall a b. (a -> b) -> a -> b
    $
            case SubExp
    len_dim of
              Constant (IntValue (Int64Value Int64
    v_d)) -> Int64 -> Int64 -> Int64
    forall a. Ord a => a -> a -> a
    min Int64
    v_d (Int64 -> Int64) -> Int64 -> Int64
    forall a b. (a -> b) -> a -> b
    $ Int64 -> Int64 -> Int64
    forall a. Ord a => a -> a -> a
    min Int64
    v_x Int64
    v_y
              SubExp
    _ -> Int64 -> Int64 -> Int64
    forall a. Ord a => a -> a -> a
    min Int64
    v_x Int64
    v_y
        (SubExp, SubExp)
    _ ->
          [Char]
    -> Exp (Rep (BuilderT GPU (State VNameSource)))
    -> BuilderT GPU (State VNameSource) SubExp
    forall (m :: * -> *).
    MonadBuilder m =>
    [Char] -> Exp (Rep m) -> m SubExp
    letSubExp [Char]
    tk_str (Exp (Rep (BuilderT GPU (State VNameSource)))
     -> BuilderT GPU (State VNameSource) SubExp)
    -> Exp (Rep (BuilderT GPU (State VNameSource)))
    -> BuilderT GPU (State VNameSource) SubExp
    forall a b. (a -> b) -> a -> b
    $ Op (Rep (BuilderT GPU (State VNameSource)))
    -> Exp (Rep (BuilderT GPU (State VNameSource)))
    forall rep. Op rep -> Exp rep
    Op (Op (Rep (BuilderT GPU (State VNameSource)))
     -> Exp (Rep (BuilderT GPU (State VNameSource))))
    -> Op (Rep (BuilderT GPU (State VNameSource)))
    -> Exp (Rep (BuilderT GPU (State VNameSource)))
    forall a b. (a -> b) -> a -> b
    $ SizeOp -> HostOp SOAC GPU
    forall (op :: * -> *) rep. SizeOp -> HostOp op rep
    SizeOp (SizeOp -> HostOp SOAC GPU) -> SizeOp -> HostOp SOAC GPU
    forall a b. (a -> b) -> a -> b
    $ Name -> SizeClass -> SizeOp
    GetSize Name
    tk_name SizeClass
    SizeTile

    ----------------------------------------------------------------------------------------------
    --- 3D Tiling (RegTiling for the outermost dimension & Block tiling for the innermost two) ---
    ----------------------------------------------------------------------------------------------

    maxRegTile :: Int64
    maxRegTile :: Int64
    maxRegTile = Int64
    30

    mkRegTileSe :: Int64 -> SubExp
    mkRegTileSe :: Int64 -> SubExp
    mkRegTileSe = Int64 -> SubExp
    forall v. IsValue v => v -> SubExp
    constant

    variantToDim :: VarianceTable -> VName -> VName -> Bool
    variantToDim :: Map VName Names -> VName -> VName -> Bool
    variantToDim Map VName Names
    variance VName
    gid_outer VName
    nm =
      VName
    gid_outer VName -> VName -> Bool
    forall a. Eq a => a -> a -> Bool
    == VName
    nm Bool -> Bool -> Bool
    || VName -> Names -> Bool
    nameIn VName
    gid_outer (Names -> VName -> Map VName Names -> Names
    forall k a. Ord k => a -> k -> Map k a -> a
    M.findWithDefault Names
    forall a. Monoid a => a
    mempty VName
    nm Map VName Names
    variance)

    -- | Checks that all streamed arrays are variant to exacly one of
    --   the two innermost parallel dimensions, and conversely, for
    --   each of the two innermost parallel dimensions, there is at
    --   least one streamed array variant to it. The result is the
    --   number of the only variant parallel dimension for each array.
    isInvarTo2of3InnerDims ::
      Names ->
      SegSpace ->
      VarianceTable ->
      [VName] ->
      Maybe [Int]
    isInvarTo2of3InnerDims :: Names -> SegSpace -> Map VName Names -> [VName] -> Maybe [Int]
    isInvarTo2of3InnerDims Names
    branch_variant SegSpace
    kspace Map VName Names
    variance [VName]
    arrs =
      let inner_perm0 :: [Maybe Int]
    inner_perm0 = (VName -> Maybe Int) -> [VName] -> [Maybe Int]
    forall a b. (a -> b) -> [a] -> [b]
    map VName -> Maybe Int
    varToOnly1of3InnerDims [VName]
    arrs
          inner_perm :: [Int]
    inner_perm = [Maybe Int] -> [Int]
    forall a. [Maybe a] -> [a]
    catMaybes [Maybe Int]
    inner_perm0
          ok1 :: Bool
    ok1 = Int -> [Int] -> Bool
    forall a. Eq a => a -> [a] -> Bool
    forall (t :: * -> *) a. (Foldable t, Eq a) => a -> t a -> Bool
    elem Int
    0 [Int]
    inner_perm Bool -> Bool -> Bool
    && Int -> [Int] -> Bool
    forall a. Eq a => a -> [a] -> Bool
    forall (t :: * -> *) a. (Foldable t, Eq a) => a -> t a -> Bool
    elem Int
    1 [Int]
    inner_perm Bool -> Bool -> Bool
    && Int -> [Int] -> Bool
    forall a. Eq a => a -> [a] -> Bool
    forall (t :: * -> *) a. (Foldable t, Eq a) => a -> t a -> Bool
    elem Int
    2 [Int]
    inner_perm
          ok2 :: Bool
    ok2 = [Maybe Int] -> Int
    forall a. [a] -> Int
    forall (t :: * -> *) a. Foldable t => t a -> Int
    length [Maybe Int]
    inner_perm0 Int -> Int -> Bool
    forall a. Eq a => a -> a -> Bool
    == [Int] -> Int
    forall a. [a] -> Int
    forall (t :: * -> *) a. Foldable t => t a -> Int
    length [Int]
    inner_perm
       in if Bool
    ok1 Bool -> Bool -> Bool
    && Bool
    ok2 then [Int] -> Maybe [Int]
    forall a. a -> Maybe a
    Just [Int]
    inner_perm else Maybe [Int]
    forall a. Maybe a
    Nothing
      where
        varToOnly1of3InnerDims :: VName -> Maybe Int
        varToOnly1of3InnerDims :: VName -> Maybe Int
    varToOnly1of3InnerDims VName
    arr = do
          (k, _) : (j, _) : (i, _) : _ <- [(VName, SubExp)] -> Maybe [(VName, SubExp)]
    forall a. a -> Maybe a
    Just ([(VName, SubExp)] -> Maybe [(VName, SubExp)])
    -> [(VName, SubExp)] -> Maybe [(VName, SubExp)]
    forall a b. (a -> b) -> a -> b
    $ [(VName, SubExp)] -> [(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
    kspace
          let variant_to = Names -> VName -> Map VName Names -> Names
    forall k a. Ord k => a -> k -> Map k a -> a
    M.findWithDefault Names
    forall a. Monoid a => a
    mempty VName
    arr Map VName Names
    variance
              branch_invariant =
                Bool -> Bool
    not (Bool -> Bool) -> Bool -> Bool
    forall a b. (a -> b) -> a -> b
    $
                  VName -> Names -> Bool
    nameIn VName
    k Names
    branch_variant
                    Bool -> Bool -> Bool
    || VName -> Names -> Bool
    nameIn VName
    j Names
    branch_variant
                    Bool -> Bool -> Bool
    || VName -> Names -> Bool
    nameIn VName
    i Names
    branch_variant
          if not branch_invariant
            then Nothing -- if i or j or k in branch_variant; return nothing
            else
              if nameIn i variant_to && not (nameIn j variant_to || nameIn k variant_to)
                then Just 0
                else
                  if nameIn j variant_to && not (nameIn i variant_to || nameIn k variant_to)
                    then Just 1
                    else
                      if nameIn k variant_to && not (nameIn i variant_to || nameIn j variant_to)
                        then Just 2
                        else Nothing

    -- | Expects a kernel statement as argument.
    --   CONDITIONS for 3D tiling optimization to fire are:
    --     1. a) The kernel body can be broken into
    --              scalar-code-1 ++ [Redomap stmt] ++ scalar-code-2.
    --        b) The kernels has a per-thread result, and obviously
    --              the result is variant to the 3rd dimension
    --              (counted from innermost to outermost)
    --     2. For the Redomap:
    --          a) the streamed arrays are one dimensional
    --          b) each of the array arguments of Redomap are variant
    --              to exactly one of the three innermost-parallel dimension
    --              of the kernel. This condition can be relaxed by interchanging
    --              kernel dimensions whenever possible.
    --     3. For scalar-code-1:
    --          a) each of the statements is a slice that produces one of the
    --             streamed arrays
    --
    -- mmBlkRegTiling :: Stm GPU -> TileM (Maybe (Stms GPU, Stm GPU))
    -- mmBlkRegTiling (Let pat aux (Op (SegOp (SegMap SegThread{} seg_space ts old_kbody))))
    doRegTiling3D :: Stm GPU -> TileM (Maybe (Stms GPU, Stm GPU))
    doRegTiling3D :: Stm GPU -> TileM (Maybe (Stms GPU, Stm GPU))
    doRegTiling3D (Let Pat (LetDec GPU)
    pat StmAux (ExpDec GPU)
    aux (Op (SegOp SegOp SegLevel GPU
    old_kernel)))
      | SegMap SegThread {} SegSpace
    space [TypeBase Shape NoUniqueness]
    kertp (KernelBody () Stms GPU
    kstms [KernelResult]
    kres) <- SegOp SegLevel GPU
    old_kernel,
        -- build the variance table, that records, for
        -- each variable name, the variables it depends on
        Map VName Names
    initial_variance <- (NameInfo (ZonkAny 1) -> Names)
    -> Map VName (NameInfo (ZonkAny 1)) -> Map VName Names
    forall a b k. (a -> b) -> Map k a -> Map k b
    M.map NameInfo (ZonkAny 1) -> Names
    forall a. Monoid a => a
    mempty (Map VName (NameInfo (ZonkAny 1)) -> Map VName Names)
    -> Map VName (NameInfo (ZonkAny 1)) -> Map VName Names
    forall a b. (a -> b) -> a -> b
    $ SegSpace -> Map VName (NameInfo (ZonkAny 1))
    forall rep. SegSpace -> Scope rep
    scopeOfSegSpace SegSpace
    space,
        Map VName Names
    variance <- Map VName Names -> Stms GPU -> Map VName Names
    varianceInStms Map VName Names
    initial_variance Stms GPU
    kstms,
        -- we get the global-thread id for the two inner dimensions,
        --   as we are probably going to use it in code generation
        (VName
    gtid_x, SubExp
    d_Kx) : (VName
    gtid_y, SubExp
    d_Ky) : (VName
    gtid_z, SubExp
    d_M) : [(VName, SubExp)]
    rem_outer_dims_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
    space,
        [(VName, SubExp)]
    rem_outer_dims <- [(VName, SubExp)] -> [(VName, SubExp)]
    forall a. [a] -> [a]
    reverse [(VName, SubExp)]
    rem_outer_dims_rev,
        -- check that the code fits the pattern having:
        -- some `code1`, followed by one Screma SOAC, followed by some `code2`
        (Stms GPU
    code1, Just Stm GPU
    screma_stmt, Stms GPU
    code2) <- Stms GPU -> (Stms GPU, Maybe (Stm GPU), Stms GPU)
    matchCodeStreamCode Stms GPU
    kstms,
        Let Pat (LetDec GPU)
    pat_redomap StmAux (ExpDec GPU)
    _ (Op Op GPU
    _) <- Stm GPU
    screma_stmt,
        -- checks that the Screma SOAC is actually a redomap and normalize it
        Just (SubExp
    common_dim, [VName]
    inp_soac_arrs, (Commutativity
    _, Lambda GPU
    red_lam, [SubExp]
    red_nes, Lambda GPU
    map_lam)) <- Stm GPU
    -> Maybe
         (SubExp, [VName],
          (Commutativity, Lambda GPU, [SubExp], Lambda GPU))
    isTileableRedomap Stm GPU
    screma_stmt,
        Bool -> Bool
    not ([SubExp] -> Bool
    forall a. [a] -> Bool
    forall (t :: * -> *) a. Foldable t => t a -> Bool
    null [SubExp]
    red_nes),
        -- assuming we have a budget of maxRegTile registers, we distribute
        -- that budget across the result of redomap and the kernel result
        Int
    num_res <- Int -> Int -> Int
    forall a. Ord a => a -> a -> a
    max ([SubExp] -> Int
    forall a. [a] -> Int
    forall (t :: * -> *) a. Foldable t => t a -> Int
    length [SubExp]
    red_nes) ([KernelResult] -> Int
    forall a. [a] -> Int
    forall (t :: * -> *) a. Foldable t => t a -> Int
    length [KernelResult]
    kres),
        Int64
    reg_tile <- Int64
    maxRegTile Int64 -> Int64 -> Int64
    forall a. Integral a => a -> a -> a
    `quot` Int -> Int64
    forall a b. (Integral a, Num b) => a -> b
    fromIntegral Int
    num_res,
        SubExp
    reg_tile_se <- Int64 -> SubExp
    mkRegTileSe Int64
    reg_tile,
        -- check that the element-type of the map and reduce are scalars:
        (Param (TypeBase Shape NoUniqueness) -> Bool)
    -> [Param (TypeBase Shape NoUniqueness)] -> Bool
    forall (t :: * -> *) a. Foldable t => (a -> Bool) -> t a -> Bool
    all (TypeBase Shape NoUniqueness -> Bool
    forall shape u. TypeBase shape u -> Bool
    primType (TypeBase Shape NoUniqueness -> Bool)
    -> (Param (TypeBase Shape NoUniqueness)
        -> TypeBase Shape NoUniqueness)
    -> Param (TypeBase Shape NoUniqueness)
    -> Bool
    forall b c a. (b -> c) -> (a -> b) -> a -> c
    . Param (TypeBase Shape NoUniqueness) -> TypeBase Shape NoUniqueness
    forall dec. Param dec -> dec
    paramDec) ([Param (TypeBase Shape NoUniqueness)] -> Bool)
    -> [Param (TypeBase Shape NoUniqueness)] -> Bool
    forall a b. (a -> b) -> a -> b
    $ Lambda GPU -> [LParam GPU]
    forall rep. Lambda rep -> [LParam rep]
    lambdaParams Lambda GPU
    map_lam,
        [TypeBase Shape NoUniqueness]
    red_res_tps <- (Param (TypeBase Shape NoUniqueness)
     -> TypeBase Shape NoUniqueness)
    -> [Param (TypeBase Shape NoUniqueness)]
    -> [TypeBase Shape NoUniqueness]
    forall a b. (a -> b) -> [a] -> [b]
    map Param (TypeBase Shape NoUniqueness) -> TypeBase Shape NoUniqueness
    forall dec. Param dec -> dec
    paramDec ([Param (TypeBase Shape NoUniqueness)]
     -> [TypeBase Shape NoUniqueness])
    -> [Param (TypeBase Shape NoUniqueness)]
    -> [TypeBase Shape NoUniqueness]
    forall a b. (a -> b) -> a -> b
    $ Int
    -> [Param (TypeBase Shape NoUniqueness)]
    -> [Param (TypeBase Shape NoUniqueness)]
    forall a. Int -> [a] -> [a]
    take ([SubExp] -> Int
    forall a. [a] -> Int
    forall (t :: * -> *) a. Foldable t => t a -> Int
    length [SubExp]
    red_nes) ([Param (TypeBase Shape NoUniqueness)]
     -> [Param (TypeBase Shape NoUniqueness)])
    -> [Param (TypeBase Shape NoUniqueness)]
    -> [Param (TypeBase Shape NoUniqueness)]
    forall a b. (a -> b) -> a -> b
    $ Lambda GPU -> [LParam GPU]
    forall rep. Lambda rep -> [LParam rep]
    lambdaParams Lambda GPU
    red_lam,
        (TypeBase Shape NoUniqueness -> Bool)
    -> [TypeBase Shape NoUniqueness] -> Bool
    forall (t :: * -> *) a. Foldable t => (a -> Bool) -> t a -> Bool
    all TypeBase Shape NoUniqueness -> Bool
    forall shape u. TypeBase shape u -> Bool
    primType [TypeBase Shape NoUniqueness]
    red_res_tps,
        -- checks that the input arrays to redomap are variant to
        -- exactly one of the two innermost dimensions of the kernel
        Just [Int]
    _ <- Names -> SegSpace -> Map VName Names -> [VName] -> Maybe [Int]
    isInvarTo2of3InnerDims Names
    forall a. Monoid a => a
    mempty SegSpace
    space Map VName Names
    variance [VName]
    inp_soac_arrs,
        -- get the free variables on which the result of redomap depends on
        [PatElem (TypeBase Shape NoUniqueness)]
    redomap_orig_res <- Pat (TypeBase Shape NoUniqueness)
    -> [PatElem (TypeBase Shape NoUniqueness)]
    forall dec. Pat dec -> [PatElem dec]
    patElems Pat (TypeBase Shape NoUniqueness)
    Pat (LetDec GPU)
    pat_redomap,
        Names
    res_red_var <- -- variance of the reduce result
          [Names] -> Names
    forall a. Monoid a => [a] -> a
    mconcat ([Names] -> Names) -> [Names] -> Names
    forall a b. (a -> b) -> a -> b
    $ (PatElem (TypeBase Shape NoUniqueness) -> Maybe Names)
    -> [PatElem (TypeBase Shape NoUniqueness)] -> [Names]
    forall a b. (a -> Maybe b) -> [a] -> [b]
    mapMaybe ((VName -> Map VName Names -> Maybe Names
    forall k a. Ord k => k -> Map k a -> Maybe a
    `M.lookup` Map VName Names
    variance) (VName -> Maybe Names)
    -> (PatElem (TypeBase Shape NoUniqueness) -> VName)
    -> PatElem (TypeBase Shape NoUniqueness)
    -> Maybe Names
    forall b c a. (b -> c) -> (a -> b) -> a -> c
    . PatElem (TypeBase Shape NoUniqueness) -> VName
    forall dec. PatElem dec -> VName
    patElemName) [PatElem (TypeBase Shape NoUniqueness)]
    redomap_orig_res,
        Names
    forall a. Monoid a => a
    mempty Names -> Names -> Bool
    forall a. Eq a => a -> a -> Bool
    /= Names
    res_red_var,
        -- we furthermore check that code1 is only formed by
        -- 1. statements that slice some globally-declared arrays
        --    to produce the input for the redomap, and
        -- 2. potentially some statements on which the redomap
        --    is independent; these are recorded in `code2''`
        Just (Stms GPU
    code2'', Map VName (Stm GPU)
    arr_tab0) <-
          (Maybe (Stms GPU, Map VName (Stm GPU))
     -> Stm GPU -> Maybe (Stms GPU, Map VName (Stm GPU)))
    -> Maybe (Stms GPU, Map VName (Stm GPU))
    -> Stms GPU
    -> Maybe (Stms GPU, Map VName (Stm GPU))
    forall b a. (b -> a -> b) -> b -> Seq a -> b
    forall (t :: * -> *) b a.
    Foldable t =>
    (b -> a -> b) -> b -> t a -> b
    foldl
            (Names
    -> Names
    -> Maybe (Stms GPU, Map VName (Stm GPU))
    -> Stm GPU
    -> Maybe (Stms GPU, Map VName (Stm GPU))
    processIndirections ([VName] -> Names
    namesFromList [VName]
    inp_soac_arrs) Names
    res_red_var)
            ((Stms GPU, Map VName (Stm GPU))
    -> Maybe (Stms GPU, Map VName (Stm GPU))
    forall a. a -> Maybe a
    Just (Stms GPU
    forall a. Seq a
    Seq.empty, Map VName (Stm GPU)
    forall k a. Map k a
    M.empty))
            Stms GPU
    code1,
        -- check that code1 contains exacly one slice for each of the input array to redomap
        [Stm GPU]
    tmp_stms <- (VName -> Maybe (Stm GPU)) -> [VName] -> [Stm GPU]
    forall a b. (a -> Maybe b) -> [a] -> [b]
    mapMaybe (VName -> Map VName (Stm GPU) -> Maybe (Stm GPU)
    forall k a. Ord k => k -> Map k a -> Maybe a
    `M.lookup` Map VName (Stm GPU)
    arr_tab0) [VName]
    inp_soac_arrs,
        [Stm GPU] -> Int
    forall a. [a] -> Int
    forall (t :: * -> *) a. Foldable t => t a -> Int
    length [Stm GPU]
    tmp_stms Int -> Int -> Bool
    forall a. Eq a => a -> a -> Bool
    == [VName] -> Int
    forall a. [a] -> Int
    forall (t :: * -> *) a. Foldable t => t a -> Int
    length [VName]
    inp_soac_arrs,
        -- code1' <- stmsFromList $ stmsToList code1 \\ stmsToList code2'',
        Stms GPU
    code2' <- Stms GPU
    code2'' Stms GPU -> Stms GPU -> Stms GPU
    forall a. Semigroup a => a -> a -> a
    <> Stms GPU
    code2,
        -- we assume the kernel results are variant to the thrid-outer parallel dimension
        -- (for sanity sake, they should be)
        [VName]
    ker_res_nms <- (KernelResult -> Maybe VName) -> [KernelResult] -> [VName]
    forall a b. (a -> Maybe b) -> [a] -> [b]
    mapMaybe KernelResult -> Maybe VName
    getResNm [KernelResult]
    kres,
        [VName] -> Int
    forall a. [a] -> Int
    forall (t :: * -> *) a. Foldable t => t a -> Int
    length [VName]
    ker_res_nms Int -> Int -> Bool
    forall a. Eq a => a -> a -> Bool
    == [KernelResult] -> Int
    forall a. [a] -> Int
    forall (t :: * -> *) a. Foldable t => t a -> Int
    length [KernelResult]
    kres,
        (TypeBase Shape NoUniqueness -> Bool)
    -> [TypeBase Shape NoUniqueness] -> Bool
    forall (t :: * -> *) a. Foldable t => (a -> Bool) -> t a -> Bool
    all TypeBase Shape NoUniqueness -> Bool
    forall shape u. TypeBase shape u -> Bool
    primType [TypeBase Shape NoUniqueness]
    kertp,
        (VName -> Bool) -> [VName] -> Bool
    forall (t :: * -> *) a. Foldable t => (a -> Bool) -> t a -> Bool
    all (Map VName Names -> VName -> VName -> Bool
    variantToDim Map VName Names
    variance VName
    gtid_z) [VName]
    ker_res_nms = do
          -- HERE STARTS THE IMPLEMENTATION:
          (new_kernel, host_stms) <- Builder GPU (Stm GPU)
    -> ReaderT (Scope GPU) (State VNameSource) (Stm GPU, 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 (Stm GPU)
     -> ReaderT (Scope GPU) (State VNameSource) (Stm GPU, Stms GPU))
    -> Builder GPU (Stm GPU)
    -> ReaderT (Scope GPU) (State VNameSource) (Stm GPU, Stms GPU)
    forall a b. (a -> b) -> a -> b
    $ do
            -- host code
            -- process the z-variant arrays that need transposition;
            -- these "manifest" statements should come before the kernel
            (tab_inn, tab_out) <-
              ((Map VName (Stm GPU), Map VName (PrimType, Stm GPU))
     -> (VName, Stm GPU)
     -> BuilderT
          GPU
          (State VNameSource)
          (Map VName (Stm GPU), Map VName (PrimType, Stm GPU)))
    -> (Map VName (Stm GPU), Map VName (PrimType, Stm GPU))
    -> [(VName, Stm GPU)]
    -> BuilderT
         GPU
         (State VNameSource)
         (Map VName (Stm GPU), Map VName (PrimType, Stm GPU))
    forall (t :: * -> *) (m :: * -> *) b a.
    (Foldable t, Monad m) =>
    (b -> a -> m b) -> b -> t a -> m b
    foldM
                (Map VName Names
    -> (VName, SubExp)
    -> (Map VName (Stm GPU), Map VName (PrimType, Stm GPU))
    -> (VName, Stm GPU)
    -> BuilderT
         GPU
         (State VNameSource)
         (Map VName (Stm GPU), Map VName (PrimType, Stm GPU))
    insertTranspose Map VName Names
    variance (VName
    gtid_z, SubExp
    d_M))
                (Map VName (Stm GPU)
    forall k a. Map k a
    M.empty, Map VName (PrimType, Stm GPU)
    forall k a. Map k a
    M.empty)
                ([(VName, Stm GPU)]
     -> BuilderT
          GPU
          (State VNameSource)
          (Map VName (Stm GPU), Map VName (PrimType, Stm GPU)))
    -> [(VName, Stm GPU)]
    -> BuilderT
         GPU
         (State VNameSource)
         (Map VName (Stm GPU), Map VName (PrimType, Stm GPU))
    forall a b. (a -> b) -> a -> b
    $ Map VName (Stm GPU) -> [(VName, Stm GPU)]
    forall k a. Map k a -> [(k, a)]
    M.toList Map VName (Stm GPU)
    arr_tab0

            tx_name <- nameFromString . prettyString <$> newVName "Tx"
            ty_name <- nameFromString . prettyString <$> newVName "Ty"

            tx0 <- letSubExp "Tx" $ Op $ SizeOp $ GetSize tx_name SizeTile
            ty0 <- letSubExp "Ty" $ Op $ SizeOp $ GetSize ty_name SizeTile
            ty <- limitTile "Ty" ty0 d_Ky
            tx <- limitTile "Tx" tx0 d_Kx
            let rz = SubExp
    reg_tile_se

            gridDim_x <- letSubExp "gridDim_x" =<< ceilDiv d_Kx tx
            gridDim_y <- letSubExp "gridDim_y" =<< ceilDiv d_Ky ty
            gridDim_z <- letSubExp "gridDim_z" =<< ceilDiv d_M rz
            let gridxyz_pexp = SubExp -> TPrimExp Int64 VName
    pe64 SubExp
    gridDim_z TPrimExp Int64 VName
    -> TPrimExp Int64 VName -> TPrimExp Int64 VName
    forall a. Num a => a -> a -> a
    * SubExp -> TPrimExp Int64 VName
    pe64 SubExp
    gridDim_y TPrimExp Int64 VName
    -> TPrimExp Int64 VName -> TPrimExp Int64 VName
    forall a. Num a => a -> a -> a
    * SubExp -> TPrimExp Int64 VName
    pe64 SubExp
    gridDim_x
            let grid_pexp = [TPrimExp Int64 VName] -> TPrimExp Int64 VName
    forall a. Num a => [a] -> a
    forall (t :: * -> *) a. (Foldable t, Num a) => t a -> a
    product ([TPrimExp Int64 VName] -> TPrimExp Int64 VName)
    -> [TPrimExp Int64 VName] -> TPrimExp Int64 VName
    forall a b. (a -> b) -> a -> b
    $ TPrimExp Int64 VName
    gridxyz_pexp TPrimExp Int64 VName
    -> [TPrimExp Int64 VName] -> [TPrimExp Int64 VName]
    forall a. a -> [a] -> [a]
    : ((VName, SubExp) -> TPrimExp Int64 VName)
    -> [(VName, SubExp)] -> [TPrimExp Int64 VName]
    forall a b. (a -> b) -> [a] -> [b]
    map (SubExp -> TPrimExp Int64 VName
    pe64 (SubExp -> TPrimExp Int64 VName)
    -> ((VName, SubExp) -> SubExp)
    -> (VName, SubExp)
    -> TPrimExp Int64 VName
    forall b c a. (b -> c) -> (a -> b) -> a -> c
    . (VName, SubExp) -> SubExp
    forall a b. (a, b) -> b
    snd) [(VName, SubExp)]
    rem_outer_dims_rev
            grid_size <- letSubExp "grid_size_tile3d" =<< toExp grid_pexp
            tblock_size <- letSubExp "tblock_size_tile3d" =<< toExp (pe64 ty * pe64 tx)
            let segthd_lvl = SegVirt -> SegLevel
    SegThreadInBlock (SegSeqDims -> SegVirt
    SegNoVirtFull ([Int] -> SegSeqDims
    SegSeqDims []))

            count_shmem <- letSubExp "count_shmem" =<< ceilDiv rz tblock_size

            gid_x <- newVName "gid_x"
            gid_y <- newVName "gid_y"
            gid_z <- newVName "gid_z"
            gid_flat <- newVName "gid_flat"

            ---- in this binder: outer seggroup ----
            (ret_seggroup, stms_seggroup) <- runBuilder $ do
              ii <- letExp "ii" =<< toExp (le64 gid_z * pe64 rz)
              jj1 <- letExp "jj1" =<< toExp (le64 gid_y * pe64 ty)
              jj2 <- letExp "jj2" =<< toExp (le64 gid_x * pe64 tx)

              -- initialize the register arrays corresponding to the result of redomap;
              reg_arr_nms <- segMap2D "res" segthd_lvl ResultPrivate (ty, tx) $ \(VName, VName)
    _ ->
                [(SubExp, TypeBase Shape NoUniqueness)]
    -> ((SubExp, TypeBase Shape NoUniqueness)
        -> BuilderT GPU (State VNameSource) SubExpRes)
    -> BuilderT GPU (State VNameSource) [SubExpRes]
    forall (t :: * -> *) (m :: * -> *) a b.
    (Traversable t, Monad m) =>
    t a -> (a -> m b) -> m (t b)
    forM ([SubExp]
    -> [TypeBase Shape NoUniqueness]
    -> [(SubExp, TypeBase Shape NoUniqueness)]
    forall a b. [a] -> [b] -> [(a, b)]
    zip [SubExp]
    red_nes [TypeBase Shape NoUniqueness]
    red_res_tps) (((SubExp, TypeBase Shape NoUniqueness)
      -> BuilderT GPU (State VNameSource) SubExpRes)
     -> BuilderT GPU (State VNameSource) [SubExpRes])
    -> ((SubExp, TypeBase Shape NoUniqueness)
        -> BuilderT GPU (State VNameSource) SubExpRes)
    -> BuilderT GPU (State VNameSource) [SubExpRes]
    forall a b. (a -> b) -> a -> b
    $ \(SubExp
    red_ne, TypeBase Shape NoUniqueness
    red_t) -> do
                  css_init <- [Char]
    -> PrimType -> [SubExp] -> BuilderT GPU (State VNameSource) VName
    forall (m :: * -> *).
    MonadBuilder m =>
    [Char] -> PrimType -> [SubExp] -> m VName
    scratch [Char]
    "res_init" (TypeBase Shape NoUniqueness -> PrimType
    forall shape u. TypeBase shape u -> PrimType
    elemType TypeBase Shape NoUniqueness
    red_t) [SubExp
    rz]
                  css <- forLoop rz [css_init] $ \VName
    i [VName
    css_merge] -> do
                    css' <- [Char]
    -> VName
    -> [VName]
    -> SubExp
    -> BuilderT GPU (State VNameSource) VName
    forall (m :: * -> *).
    MonadBuilder m =>
    [Char] -> VName -> [VName] -> SubExp -> m VName
    update [Char]
    "css" VName
    css_merge [VName
    i] SubExp
    red_ne
                    resultBodyM [Var css']
                  pure $ varRes css

              -- scratch the shared-memory arrays corresponding to the arrays that are
              --   input to the redomap and are invariant to the outermost parallel dimension.
              loc_arr_nms <- forM (M.toList tab_out) $ \(VName
    nm, (PrimType
    ptp, Stm GPU
    _)) ->
                [Char]
    -> PrimType -> [SubExp] -> BuilderT GPU (State VNameSource) VName
    forall (m :: * -> *).
    MonadBuilder m =>
    [Char] -> PrimType -> [SubExp] -> m VName
    scratch (VName -> [Char]
    baseString VName
    nm [Char] -> [Char] -> [Char]
    forall a. [a] -> [a] -> [a]
    ++ [Char]
    "_loc") PrimType
    ptp [SubExp
    rz]

              prologue_res_list <-
                forLoop' common_dim (reg_arr_nms ++ loc_arr_nms) $
                  \VName
    q [VName]
    var_nms -> do
                    let reg_arr_merge_nms :: [VName]
    reg_arr_merge_nms = Int -> [VName] -> [VName]
    forall a. Int -> [a] -> [a]
    take ([SubExp] -> Int
    forall a. [a] -> Int
    forall (t :: * -> *) a. Foldable t => t a -> Int
    length [SubExp]
    red_nes) [VName]
    var_nms
                    let loc_arr_merge_nms :: [VName]
    loc_arr_merge_nms = Int -> [VName] -> [VName]
    forall a. Int -> [a] -> [a]
    drop ([SubExp] -> Int
    forall a. [a] -> Int
    forall (t :: * -> *) a. Foldable t => t a -> Int
    length [SubExp]
    red_nes) [VName]
    var_nms

                    -- collective copy from global to shared memory
                    loc_arr_nms' <-
                      SubExp
    -> [VName]
    -> (VName
        -> [VName] -> BuilderT GPU (State VNameSource) (Body GPU))
    -> Builder GPU [VName]
    forLoop' SubExp
    count_shmem [VName]
    loc_arr_merge_nms ((VName -> [VName] -> BuilderT GPU (State VNameSource) (Body GPU))
     -> Builder GPU [VName])
    -> (VName
        -> [VName] -> BuilderT GPU (State VNameSource) (Body GPU))
    -> Builder GPU [VName]
    forall a b. (a -> b) -> a -> b
    $ \VName
    tt [VName]
    loc_arr_merge2_nms -> do
                        loc_arr_merge2_nms' <-
                          [(VName, (VName, (PrimType, Stm GPU)))]
    -> ((VName, (VName, (PrimType, Stm GPU)))
        -> BuilderT GPU (State VNameSource) VName)
    -> Builder GPU [VName]
    forall (t :: * -> *) (m :: * -> *) a b.
    (Traversable t, Monad m) =>
    t a -> (a -> m b) -> m (t b)
    forM ([VName]
    -> [(VName, (PrimType, Stm GPU))]
    -> [(VName, (VName, (PrimType, Stm GPU)))]
    forall a b. [a] -> [b] -> [(a, b)]
    zip [VName]
    loc_arr_merge2_nms (Map VName (PrimType, Stm GPU) -> [(VName, (PrimType, Stm GPU))]
    forall k a. Map k a -> [(k, a)]
    M.toList Map VName (PrimType, Stm GPU)
    tab_out)) (((VName, (VName, (PrimType, Stm GPU)))
      -> BuilderT GPU (State VNameSource) VName)
     -> Builder GPU [VName])
    -> ((VName, (VName, (PrimType, Stm GPU)))
        -> BuilderT GPU (State VNameSource) VName)
    -> Builder GPU [VName]
    forall a b. (a -> b) -> a -> b
    $ \(VName
    loc_Y_nm, (VName
    glb_Y_nm, (PrimType
    ptp_Y, Stm GPU
    load_Y))) -> do
                            ltid_flat <- [Char] -> BuilderT GPU (State VNameSource) VName
    forall (m :: * -> *). MonadFreshNames m => [Char] -> m VName
    newVName [Char]
    "ltid_flat"
                            ltid <- newVName "ltid"
                            let segspace = VName -> [(VName, SubExp)] -> SegSpace
    SegSpace VName
    ltid_flat [(VName
    ltid, SubExp
    tblock_size)]
                            ((res_v, res_i), stms) <- runBuilder $ do
                              offs <- letExp "offs" =<< toExp (pe64 tblock_size * le64 tt)
                              loc_ind <- letExp "loc_ind" =<< toExp (le64 ltid + le64 offs)
                              letBindNames [gtid_z] =<< toExp (le64 ii + le64 loc_ind)
                              let glb_ind = VName
    gtid_z
                              y_elm <-
                                letSubExp "y_elem"
                                  =<< eIf
                                    (toExp $ le64 glb_ind .<. pe64 d_M)
                                    ( do
                                        addStm load_Y
                                        res <- index "Y_elem" glb_Y_nm [q]
                                        resultBodyM [Var res]
                                    )
                                    (eBody [eBlank $ Prim ptp_Y])
                              y_ind <-
                                letSubExp "y_loc_ind"
                                  =<< eIf
                                    (toExp $ le64 loc_ind .<. pe64 rz)
                                    (toExp loc_ind >>= letTupExp' "loc_fi" >>= resultBodyM)
                                    (eBody [pure $ BasicOp $ SubExp $ intConst Int64 (-1)])
                              -- y_tp  <- subExpType y_elm
                              pure (y_elm, y_ind)

                            let ret = Certs -> VName -> [(Slice SubExp, SubExp)] -> KernelResult
    WriteReturns Certs
    forall a. Monoid a => a
    mempty VName
    loc_Y_nm [([DimIndex SubExp] -> Slice SubExp
    forall d. [DimIndex d] -> Slice d
    Slice [SubExp -> DimIndex SubExp
    forall d. d -> DimIndex d
    DimFix SubExp
    res_i], SubExp
    res_v)]
                            let body = BodyDec GPU -> Stms GPU -> [KernelResult] -> KernelBody GPU
    forall rep.
    BodyDec rep -> Stms rep -> [KernelResult] -> KernelBody rep
    KernelBody () Stms GPU
    stms [KernelResult
    ret]
                            loc_Y_nm_t <- lookupType loc_Y_nm

                            res_nms <-
                              letTupExp "Y_glb2loc" <=< renameExp $
                                Op . SegOp $
                                  SegMap segthd_lvl segspace [loc_Y_nm_t] body
                            let res_nm : _ = res_nms
                            pure res_nm
                        resultBodyM $ map Var loc_arr_merge2_nms'

                    redomap_res <-
                      segMap2D "redomap_res" segthd_lvl ResultPrivate (ty, tx) $
                        \(VName
    ltid_y, VName
    ltid_x) -> do
                          [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
    jj1 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)
                          [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
    jj2 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)
                          reg_arr_merge_nms_slc <- [VName]
    -> (VName -> BuilderT GPU (State VNameSource) VName)
    -> Builder GPU [VName]
    forall (t :: * -> *) (m :: * -> *) a b.
    (Traversable t, Monad m) =>
    t a -> (a -> m b) -> m (t b)
    forM [VName]
    reg_arr_merge_nms ((VName -> BuilderT GPU (State VNameSource) VName)
     -> Builder GPU [VName])
    -> (VName -> BuilderT GPU (State VNameSource) VName)
    -> Builder GPU [VName]
    forall a b. (a -> b) -> a -> b
    $ \VName
    reg_arr_nm ->
                            [Char]
    -> VName -> [VName] -> BuilderT GPU (State VNameSource) VName
    forall (m :: * -> *).
    MonadBuilder m =>
    [Char] -> VName -> [VName] -> m VName
    index [Char]
    "res_reg_slc" VName
    reg_arr_nm [VName
    ltid_y, VName
    ltid_x]
                          fmap subExpsRes . letTupExp' "redomap_guarded"
                            =<< eIf
                              (toExp $ le64 gtid_y .<. pe64 d_Ky .&&. le64 gtid_x .<. pe64 d_Kx)
                              ( do
                                  inp_scals_invar_outer <-
                                    forM (M.toList tab_inn) $ \(VName
    inp_arr_nm, Stm GPU
    load_stm) -> do
                                      Stm (Rep (BuilderT GPU (State VNameSource)))
    -> BuilderT GPU (State VNameSource) ()
    forall (m :: * -> *). MonadBuilder m => Stm (Rep m) -> m ()
    addStm Stm (Rep (BuilderT GPU (State VNameSource)))
    Stm GPU
    load_stm
                                      [Char]
    -> VName -> [VName] -> BuilderT GPU (State VNameSource) VName
    forall (m :: * -> *).
    MonadBuilder m =>
    [Char] -> VName -> [VName] -> m VName
    index (VName -> [Char]
    baseString VName
    inp_arr_nm) VName
    inp_arr_nm [VName
    q]
                                  -- build the loop of count R whose body is semantically the redomap code
                                  reg_arr_merge_nms' <-
                                    forLoop' rz reg_arr_merge_nms_slc $ \VName
    i [VName]
    reg_arr_mm_nms -> do
                                      [VName]
    -> Exp (Rep (BuilderT GPU (State VNameSource)))
    -> BuilderT GPU (State VNameSource) ()
    forall (m :: * -> *).
    MonadBuilder m =>
    [VName] -> Exp (Rep m) -> m ()
    letBindNames [VName
    gtid_z] (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
    ii 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
    i)
                                      [SubExp]
    -> BuilderT
         GPU
         (State VNameSource)
         (Body (Rep (BuilderT GPU (State VNameSource))))
    [SubExp] -> BuilderT GPU (State VNameSource) (Body GPU)
    forall (m :: * -> *).
    MonadBuilder m =>
    [SubExp] -> m (Body (Rep m))
    resultBodyM
                                        ([SubExp] -> BuilderT GPU (State VNameSource) (Body GPU))
    -> BuilderT GPU (State VNameSource) [SubExp]
    -> BuilderT GPU (State VNameSource) (Body GPU)
    forall (m :: * -> *) a b. Monad m => (a -> m b) -> m a -> m b
    =<< [Char]
    -> Exp (Rep (BuilderT GPU (State VNameSource)))
    -> BuilderT GPU (State VNameSource) [SubExp]
    forall (m :: * -> *).
    MonadBuilder m =>
    [Char] -> Exp (Rep m) -> m [SubExp]
    letTupExp' [Char]
    "redomap_lam"
                                        (Exp GPU -> BuilderT GPU (State VNameSource) [SubExp])
    -> BuilderT GPU (State VNameSource) (Exp GPU)
    -> BuilderT GPU (State VNameSource) [SubExp]
    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
    $ VName -> TPrimExp Int64 VName
    forall a. a -> TPrimExp Int64 a
    le64 VName
    gtid_z 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
    d_M)
                                          ( do
                                              -- read from shared memory
                                              ys <- [VName]
    -> (VName -> BuilderT GPU (State VNameSource) VName)
    -> Builder GPU [VName]
    forall (t :: * -> *) (m :: * -> *) a b.
    (Traversable t, Monad m) =>
    t a -> (a -> m b) -> m (t b)
    forM [VName]
    loc_arr_nms' ((VName -> BuilderT GPU (State VNameSource) VName)
     -> Builder GPU [VName])
    -> (VName -> BuilderT GPU (State VNameSource) VName)
    -> Builder GPU [VName]
    forall a b. (a -> b) -> a -> b
    $ \VName
    loc_arr_nm ->
                                                [Char]
    -> VName -> [VName] -> BuilderT GPU (State VNameSource) VName
    forall (m :: * -> *).
    MonadBuilder m =>
    [Char] -> VName -> [VName] -> m VName
    index [Char]
    "inp_reg_var2z" VName
    loc_arr_nm [VName
    i]
                                              cs <- forM reg_arr_mm_nms $ \VName
    reg_arr_nm ->
                                                [Char]
    -> VName -> [VName] -> BuilderT GPU (State VNameSource) VName
    forall (m :: * -> *).
    MonadBuilder m =>
    [Char] -> VName -> [VName] -> m VName
    index [Char]
    "res_reg_var2z" VName
    reg_arr_nm [VName
    i]
                                              -- here we need to put in order the scalar inputs to map:
                                              let tab_scals =
                                                    [(VName, VName)] -> Map VName VName
    forall k a. Ord k => [(k, a)] -> Map k a
    M.fromList ([(VName, VName)] -> Map VName VName)
    -> [(VName, VName)] -> Map VName VName
    forall a b. (a -> b) -> a -> b
    $
                                                      [VName] -> [VName] -> [(VName, VName)]
    forall a b. [a] -> [b] -> [(a, b)]
    zip (((VName, (PrimType, Stm GPU)) -> VName)
    -> [(VName, (PrimType, Stm GPU))] -> [VName]
    forall a b. (a -> b) -> [a] -> [b]
    map (VName, (PrimType, Stm GPU)) -> VName
    forall a b. (a, b) -> a
    fst ([(VName, (PrimType, Stm GPU))] -> [VName])
    -> [(VName, (PrimType, Stm GPU))] -> [VName]
    forall a b. (a -> b) -> a -> b
    $ Map VName (PrimType, Stm GPU) -> [(VName, (PrimType, Stm GPU))]
    forall k a. Map k a -> [(k, a)]
    M.toList Map VName (PrimType, Stm GPU)
    tab_out) [VName]
    ys
                                                        [(VName, VName)] -> [(VName, VName)] -> [(VName, VName)]
    forall a. [a] -> [a] -> [a]
    ++ [VName] -> [VName] -> [(VName, VName)]
    forall a b. [a] -> [b] -> [(a, b)]
    zip (((VName, Stm GPU) -> VName) -> [(VName, Stm GPU)] -> [VName]
    forall a b. (a -> b) -> [a] -> [b]
    map (VName, Stm GPU) -> VName
    forall a b. (a, b) -> a
    fst ([(VName, Stm GPU)] -> [VName]) -> [(VName, Stm GPU)] -> [VName]
    forall a b. (a -> b) -> a -> b
    $ Map VName (Stm GPU) -> [(VName, Stm GPU)]
    forall k a. Map k a -> [(k, a)]
    M.toList Map VName (Stm GPU)
    tab_inn) [VName]
    inp_scals_invar_outer
                                              map_inp_scals <- forM inp_soac_arrs $ \VName
    arr_nm ->
                                                case VName -> Map VName VName -> Maybe VName
    forall k a. Ord k => k -> Map k a -> Maybe a
    M.lookup VName
    arr_nm Map VName VName
    tab_scals of
                                                  Maybe VName
    Nothing -> [Char] -> BuilderT GPU (State VNameSource) VName
    forall a. HasCallStack => [Char] -> a
    error [Char]
    "Impossible case reached in tiling3D\n"
                                                  Just VName
    nm -> VName -> BuilderT GPU (State VNameSource) VName
    forall a. a -> BuilderT GPU (State VNameSource) a
    forall (f :: * -> *) a. Applicative f => a -> f a
    pure VName
    nm
                                              map_lam' <- renameLambda map_lam
                                              red_lam' <- renameLambda red_lam
                                              map_res_scals <- eLambda map_lam' (map (eSubExp . Var) map_inp_scals)
                                              red_res <- eLambda red_lam' (map eSubExp (map Var cs ++ map resSubExp map_res_scals))
                                              css <- forM (zip reg_arr_mm_nms red_res) $ \(VName
    reg_arr_nm, SubExpRes
    c) ->
                                                [Char]
    -> VName
    -> [VName]
    -> SubExp
    -> BuilderT GPU (State VNameSource) VName
    forall (m :: * -> *).
    MonadBuilder m =>
    [Char] -> VName -> [VName] -> SubExp -> m VName
    update (VName -> [Char]
    baseString VName
    reg_arr_nm) VName
    reg_arr_nm [VName
    i] (SubExpRes -> SubExp
    resSubExp SubExpRes
    c)
                                              resultBodyM $ map Var css
                                          )
                                          ([SubExp]
    -> BuilderT
         GPU
         (State VNameSource)
         (Body (Rep (BuilderT GPU (State VNameSource))))
    forall (m :: * -> *).
    MonadBuilder m =>
    [SubExp] -> m (Body (Rep m))
    resultBodyM ([SubExp]
     -> BuilderT
          GPU
          (State VNameSource)
          (Body (Rep (BuilderT GPU (State VNameSource)))))
    -> [SubExp]
    -> BuilderT
         GPU
         (State VNameSource)
         (Body (Rep (BuilderT GPU (State VNameSource))))
    forall a b. (a -> b) -> a -> b
    $ (VName -> SubExp) -> [VName] -> [SubExp]
    forall a b. (a -> b) -> [a] -> [b]
    map VName -> SubExp
    Var [VName]
    reg_arr_mm_nms)
                                  resultBodyM $ map Var reg_arr_merge_nms'
                              )
                              (resultBodyM $ map Var reg_arr_merge_nms_slc)
                    resultBodyM $ map Var $ redomap_res ++ loc_arr_nms'

              -- support for non-empty code2'
              --  segmap (ltid_y < ty, ltid_x < tx) {
              --    for i < rz do
              --        res = if (ii+i < d_M && jj1+ltid_y < d_Ky && jj2 + ltid_x < d_Kx)
              --              then code2' else dummy
              --        final_res[i] = res
              let redomap_res = Int -> [VName] -> [VName]
    forall a. Int -> [a] -> [a]
    take ([SubExp] -> Int
    forall a. [a] -> Int
    forall (t :: * -> *) a. Foldable t => t a -> Int
    length [SubExp]
    red_nes) [VName]
    prologue_res_list
              epilogue_res <-
                if length redomap_orig_res == length ker_res_nms
                  && ker_res_nms == map patElemName redomap_orig_res
                  then segMap3D "rssss" segthd_lvl ResultPrivate (se1, ty, tx) $ \(VName
    _ltid_z, VName
    ltid_y, VName
    ltid_x) ->
                    [(TypeBase Shape NoUniqueness, VName)]
    -> ((TypeBase Shape NoUniqueness, VName)
        -> BuilderT GPU (State VNameSource) SubExpRes)
    -> BuilderT GPU (State VNameSource) [SubExpRes]
    forall (t :: * -> *) (m :: * -> *) a b.
    (Traversable t, Monad m) =>
    t a -> (a -> m b) -> m (t b)
    forM ([TypeBase Shape NoUniqueness]
    -> [VName] -> [(TypeBase Shape NoUniqueness, VName)]
    forall a b. [a] -> [b] -> [(a, b)]
    zip [TypeBase Shape NoUniqueness]
    kertp [VName]
    redomap_res) (((TypeBase Shape NoUniqueness, VName)
      -> BuilderT GPU (State VNameSource) SubExpRes)
     -> BuilderT GPU (State VNameSource) [SubExpRes])
    -> ((TypeBase Shape NoUniqueness, VName)
        -> BuilderT GPU (State VNameSource) SubExpRes)
    -> BuilderT GPU (State VNameSource) [SubExpRes]
    forall a b. (a -> b) -> a -> b
    $ \(TypeBase Shape NoUniqueness
    res_tp, VName
    res) -> do
                      rss_init <- [Char]
    -> PrimType -> [SubExp] -> BuilderT GPU (State VNameSource) VName
    forall (m :: * -> *).
    MonadBuilder m =>
    [Char] -> PrimType -> [SubExp] -> m VName
    scratch [Char]
    "rss_init" (TypeBase Shape NoUniqueness -> PrimType
    forall shape u. TypeBase shape u -> PrimType
    elemType TypeBase Shape NoUniqueness
    res_tp) [SubExp
    rz, SubExp
    se1, SubExp
    se1]
                      fmap varRes $
                        forLoop rz [rss_init] $ \VName
    i [VName
    rss] -> do
                          let slice :: Slice SubExp
    slice = [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
    i, SubExp -> DimIndex SubExp
    forall d. d -> DimIndex d
    DimFix SubExp
    se0, SubExp -> DimIndex SubExp
    forall d. d -> DimIndex d
    DimFix SubExp
    se0]
                          thread_res <- [Char]
    -> VName -> [VName] -> BuilderT GPU (State VNameSource) VName
    forall (m :: * -> *).
    MonadBuilder m =>
    [Char] -> VName -> [VName] -> m VName
    index [Char]
    "thread_res" VName
    res [VName
    ltid_y, VName
    ltid_x, VName
    i]
                          rss' <- letSubExp "rss" $ BasicOp $ Update Unsafe rss slice $ Var thread_res
                          resultBodyM [rss']
                  else segMap3D "rssss" segthd_lvl ResultPrivate (se1, ty, tx) $ \(VName
    _ltid_z, VName
    ltid_y, VName
    ltid_x) -> do
                    [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
    jj1 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)
                    [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
    jj2 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)
                    rss_init <- [TypeBase Shape NoUniqueness]
    -> (TypeBase Shape NoUniqueness
        -> BuilderT GPU (State VNameSource) VName)
    -> Builder GPU [VName]
    forall (t :: * -> *) (m :: * -> *) a b.
    (Traversable t, Monad m) =>
    t a -> (a -> m b) -> m (t b)
    forM [TypeBase Shape NoUniqueness]
    kertp ((TypeBase Shape NoUniqueness
      -> BuilderT GPU (State VNameSource) VName)
     -> Builder GPU [VName])
    -> (TypeBase Shape NoUniqueness
        -> BuilderT GPU (State VNameSource) VName)
    -> Builder GPU [VName]
    forall a b. (a -> b) -> a -> b
    $ \TypeBase Shape NoUniqueness
    res_tp ->
                      [Char]
    -> PrimType -> [SubExp] -> BuilderT GPU (State VNameSource) VName
    forall (m :: * -> *).
    MonadBuilder m =>
    [Char] -> PrimType -> [SubExp] -> m VName
    scratch [Char]
    "rss_init" (TypeBase Shape NoUniqueness -> PrimType
    forall shape u. TypeBase shape u -> PrimType
    elemType TypeBase Shape NoUniqueness
    res_tp) [SubExp
    rz, SubExp
    se1, SubExp
    se1]
                    rss <- forLoop' rz rss_init $ \VName
    i [VName]
    rss_merge -> do
                      [VName]
    -> Exp (Rep (BuilderT GPU (State VNameSource)))
    -> BuilderT GPU (State VNameSource) ()
    forall (m :: * -> *).
    MonadBuilder m =>
    [VName] -> Exp (Rep m) -> m ()
    letBindNames [VName
    gtid_z] (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
    ii 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
    i)
                      [(PatElem (TypeBase Shape NoUniqueness), VName)]
    -> ((PatElem (TypeBase Shape NoUniqueness), VName)
        -> BuilderT GPU (State VNameSource) VName)
    -> BuilderT GPU (State VNameSource) ()
    forall (t :: * -> *) (m :: * -> *) a b.
    (Foldable t, Monad m) =>
    t a -> (a -> m b) -> m ()
    forM_ ([PatElem (TypeBase Shape NoUniqueness)]
    -> [VName] -> [(PatElem (TypeBase Shape NoUniqueness), VName)]
    forall a b. [a] -> [b] -> [(a, b)]
    zip [PatElem (TypeBase Shape NoUniqueness)]
    redomap_orig_res [VName]
    redomap_res) (((PatElem (TypeBase Shape NoUniqueness), VName)
      -> BuilderT GPU (State VNameSource) VName)
     -> BuilderT GPU (State VNameSource) ())
    -> ((PatElem (TypeBase Shape NoUniqueness), VName)
        -> BuilderT GPU (State VNameSource) VName)
    -> BuilderT GPU (State VNameSource) ()
    forall a b. (a -> b) -> a -> b
    $ \(PatElem (TypeBase Shape NoUniqueness)
    o_res, VName
    n_res) -> do
                        c <- [Char]
    -> VName -> [VName] -> BuilderT GPU (State VNameSource) VName
    forall (m :: * -> *).
    MonadBuilder m =>
    [Char] -> VName -> [VName] -> m VName
    index [Char]
    "redomap_thd" VName
    n_res [VName
    ltid_y, VName
    ltid_x, VName
    i]
                        letBindNames [patElemName o_res] =<< toExp (le64 c)
                        pure c
                      res_els <-
                        [Char]
    -> Exp (Rep (BuilderT GPU (State VNameSource)))
    -> BuilderT GPU (State VNameSource) [SubExp]
    forall (m :: * -> *).
    MonadBuilder m =>
    [Char] -> Exp (Rep m) -> m [SubExp]
    letTupExp' [Char]
    "res_elem"
                          (Exp GPU -> BuilderT GPU (State VNameSource) [SubExp])
    -> BuilderT GPU (State VNameSource) (Exp GPU)
    -> BuilderT GPU (State VNameSource) [SubExp]
    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
    $
                                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
    d_Ky
                                  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_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
    d_Kx
                                  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_z
                                  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
    d_M
                            )
                            ( 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
    code2'
                                [SubExp]
    -> BuilderT
         GPU
         (State VNameSource)
         (Body (Rep (BuilderT GPU (State VNameSource))))
    forall (m :: * -> *).
    MonadBuilder m =>
    [SubExp] -> m (Body (Rep m))
    resultBodyM ([SubExp]
     -> BuilderT
          GPU
          (State VNameSource)
          (Body (Rep (BuilderT GPU (State VNameSource)))))
    -> [SubExp]
    -> BuilderT
         GPU
         (State VNameSource)
         (Body (Rep (BuilderT GPU (State VNameSource))))
    forall a b. (a -> b) -> a -> b
    $ (VName -> SubExp) -> [VName] -> [SubExp]
    forall a b. (a -> b) -> [a] -> [b]
    map VName -> SubExp
    Var [VName]
    ker_res_nms
                            )
                            ([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
    $ (TypeBase Shape NoUniqueness
     -> BuilderT GPU (State VNameSource) (Exp GPU))
    -> [TypeBase Shape NoUniqueness]
    -> [BuilderT GPU (State VNameSource) (Exp GPU)]
    forall a b. (a -> b) -> [a] -> [b]
    map TypeBase Shape NoUniqueness
    -> BuilderT
         GPU
         (State VNameSource)
         (Exp (Rep (BuilderT GPU (State VNameSource))))
    TypeBase Shape NoUniqueness
    -> BuilderT GPU (State VNameSource) (Exp GPU)
    forall (m :: * -> *).
    MonadBuilder m =>
    TypeBase Shape NoUniqueness -> m (Exp (Rep m))
    eBlank [TypeBase Shape NoUniqueness]
    kertp)
                      rss' <- forM (zip res_els rss_merge) $ \(SubExp
    res_el, VName
    rs_merge) -> do
                        let slice :: Slice SubExp
    slice = [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
    i, SubExp -> DimIndex SubExp
    forall d. d -> DimIndex d
    DimFix SubExp
    se0, SubExp -> DimIndex SubExp
    forall d. d -> DimIndex d
    DimFix SubExp
    se0]
                        [Char]
    -> Exp (Rep (BuilderT GPU (State VNameSource)))
    -> BuilderT GPU (State VNameSource) SubExp
    forall (m :: * -> *).
    MonadBuilder m =>
    [Char] -> Exp (Rep m) -> m SubExp
    letSubExp [Char]
    "rss" (Exp (Rep (BuilderT GPU (State VNameSource)))
     -> BuilderT GPU (State VNameSource) SubExp)
    -> Exp (Rep (BuilderT GPU (State VNameSource)))
    -> BuilderT GPU (State VNameSource) 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
    $ Safety -> VName -> Slice SubExp -> SubExp -> BasicOp
    Update Safety
    Unsafe VName
    rs_merge Slice SubExp
    slice SubExp
    res_el
                      resultBodyM rss'
                    pure $ varsRes rss

              ----------------------------------------------------------------
              -- Finally, reshape the result arrays for the RegTileReturn  ---
              ----------------------------------------------------------------
              let regtile_ret_dims =
                    ((VName, SubExp) -> (SubExp, SubExp, SubExp))
    -> [(VName, SubExp)] -> [(SubExp, SubExp, SubExp)]
    forall a b. (a -> b) -> [a] -> [b]
    map (\(VName
    _, SubExp
    sz) -> (SubExp
    sz, SubExp
    se1, SubExp
    se1)) [(VName, SubExp)]
    rem_outer_dims
                      [(SubExp, SubExp, SubExp)]
    -> [(SubExp, SubExp, SubExp)] -> [(SubExp, SubExp, SubExp)]
    forall a. [a] -> [a] -> [a]
    ++ [(SubExp
    d_M, SubExp
    se1, SubExp
    rz), (SubExp
    d_Ky, SubExp
    ty, SubExp
    se1), (SubExp
    d_Kx, SubExp
    tx, SubExp
    se1)]

              epilogue_res' <- forM epilogue_res $ \VName
    res ->
                if [(VName, SubExp)] -> Bool
    forall a. [a] -> Bool
    forall (t :: * -> *) a. Foldable t => t a -> Bool
    null [(VName, SubExp)]
    rem_outer_dims
                  then VName -> BuilderT GPU (State VNameSource) VName
    forall a. a -> BuilderT GPU (State VNameSource) a
    forall (f :: * -> *) a. Applicative f => a -> f a
    pure VName
    res
                  else do
                    -- Add dummy dimensions to tile to reflect the outer dimensions
                    res_tp' <- VName
    -> BuilderT GPU (State VNameSource) (TypeBase Shape NoUniqueness)
    forall rep (m :: * -> *).
    HasScope rep m =>
    VName -> m (TypeBase Shape NoUniqueness)
    lookupType VName
    res
                    let (block_dims, rest_dims) = splitAt 2 $ arrayDims res_tp'
                        ones = ((VName, SubExp) -> SubExp) -> [(VName, SubExp)] -> [SubExp]
    forall a b. (a -> b) -> [a] -> [b]
    map (SubExp -> (VName, SubExp) -> SubExp
    forall a b. a -> b -> a
    const SubExp
    se1) [(VName, SubExp)]
    rem_outer_dims
                        new_shape = [SubExp] -> Shape
    forall d. [d] -> ShapeBase d
    Shape ([SubExp] -> Shape) -> [SubExp] -> Shape
    forall a b. (a -> b) -> a -> b
    $ [[SubExp]] -> [SubExp]
    forall (t :: * -> *) a. Foldable t => t [a] -> [a]
    concat [[SubExp]
    ones, [SubExp]
    block_dims, [SubExp]
    ones, [SubExp]
    rest_dims]
                    letExp "res_reshaped" . BasicOp $
                      Reshape res (reshapeAll (arrayShape res_tp') new_shape)

              pure $ map (RegTileReturns mempty regtile_ret_dims) epilogue_res'
            -- END (ret_seggroup, stms_seggroup) <- runBuilder $ do
            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
    grid_size) (SubExp -> Count BlockSize SubExp
    forall {k} (u :: k) e. e -> Count u e
    Count SubExp
    tblock_size)
                level' = SegVirt -> Maybe KernelGrid -> SegLevel
    SegBlock SegVirt
    SegNoVirt (KernelGrid -> Maybe KernelGrid
    forall a. a -> Maybe a
    Just KernelGrid
    grid)
                space' = VName -> [(VName, SubExp)] -> SegSpace
    SegSpace VName
    gid_flat ([(VName, SubExp)]
    rem_outer_dims [(VName, SubExp)] -> [(VName, SubExp)] -> [(VName, SubExp)]
    forall a. [a] -> [a] -> [a]
    ++ [(VName
    gid_z, SubExp
    gridDim_z), (VName
    gid_y, SubExp
    gridDim_y), (VName
    gid_x, SubExp
    gridDim_x)])
                kbody' = BodyDec GPU -> Stms GPU -> [KernelResult] -> KernelBody GPU
    forall rep.
    BodyDec rep -> Stms rep -> [KernelResult] -> KernelBody rep
    KernelBody () Stms GPU
    stms_seggroup [KernelResult]
    ret_seggroup

            pure $ Let pat aux $ Op $ SegOp $ SegMap level' space' kertp kbody'
          -- END (new_kernel, host_stms) <- runBuilder $ do
          pure $ Just (host_stms, new_kernel)
      where
        getResNm :: KernelResult -> Maybe VName
    getResNm (Returns ResultManifest
    ResultMaySimplify Certs
    _ (Var VName
    res_nm)) = VName -> Maybe VName
    forall a. a -> Maybe a
    Just VName
    res_nm
        getResNm KernelResult
    _ = Maybe VName
    forall a. Maybe a
    Nothing

        limitTile :: String -> SubExp -> SubExp -> Builder GPU SubExp
        limitTile :: [Char]
    -> SubExp -> SubExp -> BuilderT GPU (State VNameSource) SubExp
    limitTile [Char]
    t_str SubExp
    t SubExp
    d_K = [Char]
    -> Exp (Rep (BuilderT GPU (State VNameSource)))
    -> BuilderT GPU (State VNameSource) SubExp
    forall (m :: * -> *).
    MonadBuilder m =>
    [Char] -> Exp (Rep m) -> m SubExp
    letSubExp [Char]
    t_str (Exp (Rep (BuilderT GPU (State VNameSource)))
     -> BuilderT GPU (State VNameSource) SubExp)
    -> Exp (Rep (BuilderT GPU (State VNameSource)))
    -> BuilderT GPU (State VNameSource) 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 -> BinOp
    SMin IntType
    Int64) SubExp
    t SubExp
    d_K
        insertTranspose ::
          VarianceTable ->
          (VName, SubExp) ->
          (M.Map VName (Stm GPU), M.Map VName (PrimType, Stm GPU)) ->
          (VName, Stm GPU) ->
          Builder GPU (M.Map VName (Stm GPU), M.Map VName (PrimType, Stm GPU))
        insertTranspose :: Map VName Names
    -> (VName, SubExp)
    -> (Map VName (Stm GPU), Map VName (PrimType, Stm GPU))
    -> (VName, Stm GPU)
    -> BuilderT
         GPU
         (State VNameSource)
         (Map VName (Stm GPU), Map VName (PrimType, Stm GPU))
    insertTranspose Map VName Names
    variance (VName
    gidz, SubExp
    _) (Map VName (Stm GPU)
    tab_inn, Map VName (PrimType, Stm GPU)
    tab_out) (VName
    p_nm, stm :: Stm GPU
    stm@(Let Pat (LetDec GPU)
    patt StmAux (ExpDec GPU)
    yy (BasicOp (Index VName
    arr_nm Slice SubExp
    slc))))
          | [PatElem (TypeBase Shape NoUniqueness)
    p] <- Pat (TypeBase Shape NoUniqueness)
    -> [PatElem (TypeBase Shape NoUniqueness)]
    forall dec. Pat dec -> [PatElem dec]
    patElems Pat (TypeBase Shape NoUniqueness)
    Pat (LetDec GPU)
    patt,
            PrimType
    ptp <- TypeBase Shape NoUniqueness -> PrimType
    forall shape u. TypeBase shape u -> PrimType
    elemType (TypeBase Shape NoUniqueness -> PrimType)
    -> TypeBase Shape NoUniqueness -> PrimType
    forall a b. (a -> b) -> a -> b
    $ PatElem (TypeBase Shape NoUniqueness)
    -> TypeBase Shape NoUniqueness
    forall dec. Typed dec => PatElem dec -> TypeBase Shape NoUniqueness
    patElemType PatElem (TypeBase Shape NoUniqueness)
    p,
            VName
    p_nm VName -> VName -> Bool
    forall a. Eq a => a -> a -> Bool
    == PatElem (TypeBase Shape NoUniqueness) -> VName
    forall dec. PatElem dec -> VName
    patElemName PatElem (TypeBase Shape NoUniqueness)
    p =
              case (DimIndex SubExp -> Bool) -> [DimIndex SubExp] -> [Int]
    forall a. (a -> Bool) -> [a] -> [Int]
    L.findIndices (Map VName Names -> VName -> DimIndex SubExp -> Bool
    variantSliceDim Map VName Names
    variance VName
    gidz) (Slice SubExp -> [DimIndex SubExp]
    forall d. Slice d -> [DimIndex d]
    unSlice Slice SubExp
    slc) of
                [] -> (Map VName (Stm GPU), Map VName (PrimType, Stm GPU))
    -> BuilderT
         GPU
         (State VNameSource)
         (Map VName (Stm GPU), Map VName (PrimType, Stm GPU))
    forall a. a -> BuilderT GPU (State VNameSource) a
    forall (f :: * -> *) a. Applicative f => a -> f a
    pure (VName -> Stm GPU -> Map VName (Stm GPU) -> Map VName (Stm GPU)
    forall k a. Ord k => k -> a -> Map k a -> Map k a
    M.insert VName
    p_nm Stm GPU
    stm Map VName (Stm GPU)
    tab_inn, Map VName (PrimType, Stm GPU)
    tab_out)
                Int
    i : [Int]
    _ -> do
                  arr_tp <- VName
    -> BuilderT GPU (State VNameSource) (TypeBase Shape NoUniqueness)
    forall rep (m :: * -> *).
    HasScope rep m =>
    VName -> m (TypeBase Shape NoUniqueness)
    lookupType VName
    arr_nm
                  let perm = [Int
    i Int -> Int -> Int
    forall a. Num a => a -> a -> a
    + Int
    1 .. TypeBase Shape NoUniqueness -> Int
    forall shape u. ArrayShape shape => TypeBase shape u -> Int
    arrayRank TypeBase Shape NoUniqueness
    arr_tp Int -> Int -> Int
    forall a. Num a => a -> a -> a
    - Int
    1] [Int] -> [Int] -> [Int]
    forall a. [a] -> [a] -> [a]
    ++ [Int
    0 .. Int
    i]
                  let arr_tr_str = VName -> [Char]
    baseString VName
    arr_nm [Char] -> [Char] -> [Char]
    forall a. [a] -> [a] -> [a]
    ++ [Char]
    "_transp"
                  arr_tr_nm <- letExp arr_tr_str $ BasicOp $ Manifest arr_nm perm
                  let e_ind' = BasicOp -> Exp GPU
    forall rep. BasicOp -> Exp rep
    BasicOp (BasicOp -> Exp GPU) -> BasicOp -> Exp GPU
    forall a b. (a -> b) -> a -> b
    $ VName -> Slice SubExp -> BasicOp
    Index VName
    arr_tr_nm Slice SubExp
    slc
                  let stm' = Pat (LetDec GPU) -> StmAux (ExpDec GPU) -> Exp GPU -> Stm GPU
    forall rep.
    Pat (LetDec rep) -> StmAux (ExpDec rep) -> Exp rep -> Stm rep
    Let Pat (LetDec GPU)
    patt StmAux (ExpDec GPU)
    yy Exp GPU
    e_ind'
                  pure (tab_inn, M.insert p_nm (ptp, stm') tab_out)
        insertTranspose Map VName Names
    _ (VName, SubExp)
    _ (Map VName (Stm GPU), Map VName (PrimType, Stm GPU))
    _ (VName, Stm GPU)
    _ = [Char]
    -> BuilderT
         GPU
         (State VNameSource)
         (Map VName (Stm GPU), Map VName (PrimType, Stm GPU))
    forall a. HasCallStack => [Char] -> a
    error [Char]
    "\nUnreachable case reached in insertTranspose case, doRegTiling3D\n"

        variantSliceDim :: VarianceTable -> VName -> DimIndex SubExp -> Bool
        variantSliceDim :: Map VName Names -> VName -> DimIndex SubExp -> Bool
    variantSliceDim Map VName Names
    variance VName
    gidz (DimFix (Var VName
    vnm)) = Map VName Names -> VName -> VName -> Bool
    variantToDim Map VName Names
    variance VName
    gidz VName
    vnm
        variantSliceDim Map VName Names
    _ VName
    _ DimIndex SubExp
    _ = Bool
    False
    doRegTiling3D Stm GPU
    _ = Maybe (Stms GPU, Stm GPU) -> TileM (Maybe (Stms GPU, Stm GPU))
    forall a. a -> ReaderT (Scope GPU) (State VNameSource) a
    forall (f :: * -> *) a. Applicative f => a -> f a
    pure Maybe (Stms GPU, Stm GPU)
    forall a. Maybe a
    Nothing