{-# 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
import Futhark.Util.IntegralExp qualified as IE

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 -> Stms GPU -> Bool
isInnerCoal :: Env -> VName -> Stms GPU -> Bool
isInnerCoal (WithEnv
_, IxFnEnv
ixfn_env) VName
slc_X Stms GPU
load_X
  | Just (Stms GPU
_, Let (Pat [PatElem (LetDec GPU)
pe]) StmAux (ExpDec GPU)
_ (BasicOp (Index VName
x Slice SubExp
slice))) <- Stms GPU -> Maybe (Stms GPU, Stm GPU)
forall lore. Stms lore -> Maybe (Stms lore, Stm lore)
stmsLast Stms GPU
load_X,
    VName
slc_X VName -> VName -> Bool
forall a. Eq a => a -> a -> Bool
== PatElem Type -> VName
forall dec. PatElem dec -> VName
patElemName PatElem Type
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 (LMAD (TPrimExp Int64 VName) -> Bool)
-> (LMAD (TPrimExp Int64 VName) -> LMAD (TPrimExp Int64 VName))
-> LMAD (TPrimExp Int64 VName)
-> Bool
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (LMAD (TPrimExp Int64 VName)
-> Slice (TPrimExp Int64 VName) -> LMAD (TPrimExp Int64 VName)
forall num.
(Eq num, IntegralExp num) =>
LMAD num -> Slice num -> LMAD num
`LMAD.slice` (SubExp -> TPrimExp Int64 VName)
-> Slice SubExp -> Slice (TPrimExp Int64 VName)
forall a b. (a -> b) -> Slice a -> Slice b
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap SubExp -> TPrimExp Int64 VName
pe64 Slice SubExp
slice)) (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
_ Stms 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) => Name -> PrimType -> [SubExp] -> m VName
scratch :: forall (m :: * -> *).
MonadBuilder m =>
Name -> PrimType -> [SubExp] -> m VName
scratch Name
se_name PrimType
t [SubExp]
shape = Name -> Exp (Rep m) -> m VName
forall (m :: * -> *).
MonadBuilder m =>
Name -> Exp (Rep m) -> m VName
letExp Name
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),
    (Stms GPU, VName, PrimType, Stms 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),
    (Stms GPU, VName, PrimType, Stms 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
height_A, VName
gtid_y, SubExp
width_B, SubExp
common_dim),
    (VName
iii, VName
jjj),
    (Stms GPU
load_A, VName
inp_A, PrimType
pt_A, Stms 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 <- Name
-> Exp (Rep (BuilderT GPU (State VNameSource)))
-> BuilderT GPU (State VNameSource) VName
forall (m :: * -> *).
MonadBuilder m =>
Name -> Exp (Rep m) -> m VName
letExp Name
"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_x, 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_y, 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 :: (TPrimExp Int64 VName -> TPrimExp Int64 VName)
-> Bool
-> (VName, VName)
-> (VName, VName)
-> BuilderT
     GPU (State VNameSource) (VName, VName, TPrimExp Int64 VName)
mk_ik TPrimExp Int64 VName -> TPrimExp Int64 VName
pad Bool
is_coal (VName
thd_x, VName
thd_y) (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 <- Name
-> Exp (Rep (BuilderT GPU (State VNameSource)))
-> BuilderT GPU (State VNameSource) VName
forall (m :: * -> *).
MonadBuilder m =>
Name -> Exp (Rep m) -> m VName
letExp Name
"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)
            -- the rx rows of each thread form a block with a padded stride.
            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
+ VName -> TPrimExp Int64 VName
forall a. a -> TPrimExp Int64 a
le64 VName
i TPrimExp Int64 VName
-> TPrimExp Int64 VName -> TPrimExp Int64 VName
forall e. IntegralExp e => e -> e -> e
`IE.quot` SubExp -> TPrimExp Int64 VName
pe64 SubExp
rx TPrimExp Int64 VName
-> TPrimExp Int64 VName -> TPrimExp Int64 VName
forall a. Num a => a -> a -> a
* (TPrimExp Int64 VName -> TPrimExp Int64 VName
pad (SubExp -> TPrimExp Int64 VName
pe64 SubExp
rx 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
- SubExp -> TPrimExp Int64 VName
pe64 SubExp
rx TPrimExp Int64 VName
-> TPrimExp Int64 VName -> TPrimExp Int64 VName
forall a. Num a => a -> a -> a
* SubExp -> TPrimExp Int64 VName
pe64 SubExp
t_seq)
            pure (i, k, e)
      mk_ik TPrimExp Int64 VName -> TPrimExp Int64 VName
pad Bool
_ (VName
thd_x, VName
thd_y) (VName
i0, VName
k0) = do
        -- matrix is transposed case (i.e., uncoalesced):
        let t_par :: SubExp
t_par = SubExp
tx
        k <- Name
-> Exp (Rep (BuilderT GPU (State VNameSource)))
-> BuilderT GPU (State VNameSource) VName
forall (m :: * -> *).
MonadBuilder m =>
Name -> Exp (Rep m) -> m VName
letExp Name
"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)
        -- the rx elements of each thread form a block with a padded stride.
        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
i TPrimExp Int64 VName
-> TPrimExp Int64 VName -> TPrimExp Int64 VName
forall e. IntegralExp e => e -> e -> e
`IE.quot` SubExp -> TPrimExp Int64 VName
pe64 SubExp
rx TPrimExp Int64 VName
-> TPrimExp Int64 VName -> TPrimExp Int64 VName
forall a. Num a => a -> a -> a
* (TPrimExp Int64 VName -> TPrimExp Int64 VName
pad (SubExp -> TPrimExp Int64 VName
pe64 SubExp
rx) 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
k TPrimExp Int64 VName
-> TPrimExp Int64 VName -> TPrimExp Int64 VName
forall a. Num a => a -> a -> a
* SubExp -> TPrimExp Int64 VName
pe64 SubExp
tx TPrimExp Int64 VName
-> TPrimExp Int64 VName -> TPrimExp Int64 VName
forall a. Num a => a -> a -> a
* TPrimExp Int64 VName -> TPrimExp Int64 VName
pad (SubExp -> TPrimExp Int64 VName
pe64 SubExp
rx)
        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_x, VName
ltid_y) = do
        css <- SubExp
-> [VName]
-> (VName
    -> [VName] -> BuilderT GPU (State VNameSource) (Body GPU))
-> BuilderT GPU (State VNameSource) VName
forLoop SubExp
rx [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
ry [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
<=< Name
-> Exp (Rep (BuilderT GPU (State VNameSource)))
-> BuilderT GPU (State VNameSource) [SubExp]
forall (m :: * -> *).
MonadBuilder m =>
Name -> Exp (Rep m) -> m [SubExp]
letTupExp' Name
"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
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
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
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
width_B)
                )
                ( do
                    a <- VName -> VName -> BuilderT GPU (State VNameSource) VName
a_idx_fn VName
ltid_x VName
i
                    b <- b_idx_fn ltid_y 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 <- Name
-> SegLevel
-> ResultManifest
-> (SubExp, SubExp)
-> ((VName, VName) -> BuilderT GPU (State VNameSource) [SubExpRes])
-> Builder GPU [VName]
segMap2D Name
"redomap_res" SegLevel
segthd_lvl ResultManifest
ResultPrivate (SubExp
tx, SubExp
ty) (((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_x, VName
ltid_y) -> do
            css_init <- Name -> VName -> [VName] -> BuilderT GPU (State VNameSource) VName
forall (m :: * -> *).
MonadBuilder m =>
Name -> VName -> [VName] -> m VName
index Name
"css_init" VName
css_merge [VName
ltid_x, VName
ltid_y]

            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
<=< Name
-> Exp (Rep (BuilderT GPU (State VNameSource)))
-> BuilderT GPU (State VNameSource) [SubExp]
forall (m :: * -> *).
MonadBuilder m =>
Name -> Exp (Rep m) -> m [SubExp]
letTupExp' Name
"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_x, VName
ltid_y))
                  ([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, Stms GPU, VName) ->
        Builder GPU (VName, VName -> VName -> VName -> Builder GPU VName)
      copyGlb2ShMem :: Bool
-> VName
-> (VName, VName, PrimType, SubExp, VName, Stms 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, Stms 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 -> Stms GPU -> Bool
isInnerCoal Env
env VName
inp_X Stms GPU
load_X
            str_A :: Name
str_A = VName -> Name
baseName VName
inp_X
        x_loc <-
          Name
-> VName
-> [SubExp]
-> (SubExp, SubExp)
-> ([VName] -> (VName, VName) -> Builder GPU (SubExp, SubExp))
-> BuilderT GPU (State VNameSource) VName
segScatter2D (Name
str_A Name -> Name -> Name
forall a. Semigroup a => a -> a -> a
<> Name
"_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
          -- The stride between the blocks of consecutive threads is padded
          -- to be odd, avoiding bank conflicts.  A warp reads A at only a
          -- few distinct addresses, which only conflict in the coalesced
          -- layout.
          pad :: TPrimExp Int64 VName -> TPrimExp Int64 VName
pad
            | Bool
is_B Bool -> Bool -> Bool
|| Env -> VName -> Stms GPU -> Bool
isInnerCoal Env
env VName
inp_X Stms GPU
load_X = TPrimExp Int64 VName -> TPrimExp Int64 VName
oddUp
            | Bool
otherwise = TPrimExp Int64 VName -> TPrimExp Int64 VName
forall a. a -> a
id
          --
          indexLocMem ::
            Bool ->
            Name ->
            VName ->
            VName ->
            VName ->
            VName ->
            Builder GPU VName
          indexLocMem :: Bool
-> Name
-> VName
-> VName
-> VName
-> VName
-> BuilderT GPU (State VNameSource) VName
indexLocMem Bool
is_inner_coal Name
str_A VName
x_loc VName
k VName
ltid VName
ij = do
            let (SubExp
r_par, SubExp
t_seq) = (SubExp
rx, SubExp
tk)
            x_loc_ind_32 <-
              Name
-> Exp (Rep (BuilderT GPU (State VNameSource)))
-> BuilderT GPU (State VNameSource) VName
forall (m :: * -> *).
MonadBuilder m =>
Name -> Exp (Rep m) -> m VName
letExp (Name
str_A Name -> Name -> Name
forall a. Semigroup a => a -> a -> a
<> Name
"_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
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
+ VName -> TPrimExp Int64 VName
forall a. a -> TPrimExp Int64 a
le64 VName
ltid TPrimExp Int64 VName
-> TPrimExp Int64 VName -> TPrimExp Int64 VName
forall a. Num a => a -> a -> a
* TPrimExp Int64 VName -> TPrimExp Int64 VName
pad (SubExp -> TPrimExp Int64 VName
pe64 SubExp
r_par TPrimExp Int64 VName
-> TPrimExp Int64 VName -> TPrimExp Int64 VName
forall a. Num a => a -> a -> a
* SubExp -> TPrimExp Int64 VName
pe64 SubExp
t_seq)
                      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 TPrimExp Int64 VName
-> TPrimExp Int64 VName -> TPrimExp Int64 VName
forall a. Num a => a -> a -> a
* TPrimExp Int64 VName -> TPrimExp Int64 VName
pad (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
tx TPrimExp Int64 VName
-> TPrimExp Int64 VName -> TPrimExp Int64 VName
forall a. Num a => a -> a -> a
* TPrimExp Int64 VName -> TPrimExp Int64 VName
pad (SubExp -> TPrimExp Int64 VName
pe64 SubExp
r_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_x, VName
thd_y) = do
            let str_A :: Name
str_A = VName -> Name
baseName VName
inp_X
                t_seq :: SubExp
t_seq = SubExp
tk
            (i, k, epx_loc_fi) <- (TPrimExp Int64 VName -> TPrimExp Int64 VName)
-> Bool
-> (VName, VName)
-> (VName, VName)
-> BuilderT
     GPU (State VNameSource) (VName, VName, TPrimExp Int64 VName)
mk_ik TPrimExp Int64 VName -> TPrimExp Int64 VName
pad Bool
is_inner_coal (VName
thd_x, VName
thd_y) (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
                      addStms 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 [Type]
ts KernelBody GPU
old_kbody))))
  | Body () 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
    [Type
res_tp] <- [Type]
ts,
    Type -> Bool
forall shape o. TypeBase shape o -> Bool
isAcc Type
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_y, SubExp
width_B) : (VName
gtid_x, 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',
        (Stms GPU
load_A, VName
inp_A, PrimType
map_t1, Stms 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, (Stms GPU, VName, PrimType, Stms 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 --
      ---- 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
-> Builder
     GPU
     (SubExp, SubExp, SubExp, SubExp, SubExp, SubExp, SubExp, SubExp,
      SubExp, SubExp, SubExp)
mkTileMemSizes SubExp
height_A SubExp
width_B SubExp
common_dim

        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_x <- letSubExp "gridDim_x" =<< ceilDiv height_A tx_rx
        gridDim_y <- letSubExp "gridDim_y" =<< ceilDiv width_B ty_ry

        let gridxyt_pexp = 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_y 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_x * pe64 tx_rx)
          jjj <- letExp "jjj" =<< toExp (le64 gid_y * pe64 ty_ry)
          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
height_A, VName
gtid_y, SubExp
width_B, SubExp
common_dim),
                  (VName
iii, VName
jjj),
                  (Stms GPU
load_A, VName
inp_A, PrimType
map_t1, Stms 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 <- Name
-> Exp (Rep (BuilderT GPU (State VNameSource)))
-> BuilderT GPU (State VNameSource) VName
forall (m :: * -> *).
MonadBuilder m =>
Name -> Exp (Rep m) -> m VName
letExp Name
"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_x < tx, ltid_y < ty) {
          --    for i < rx do
          --      for j < ry do
          --        res = if (iii+ltid_x*rx+i < height_A && jjj+ltid_y*ry+j < width_B)
          --              then code2' else dummy
          --        final_res[i,j] = res
          mkEpilogueAccRes
            segthd_lvl
            (redomap_orig_res, redomap_res)
            (res_nm, res_tp)
            (tx, ty, rx, ry)
            (iii, jjj)
            (gtid_x, gtid_y)
            (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_x, SubExp
gridDim_x), (VName
gid_y, SubExp
gridDim_y)])
            kbody' = BodyDec GPU -> Stms GPU -> [KernelResult] -> KernelBody GPU
forall rep res. BodyDec rep -> Stms rep -> [res] -> GBody rep res
Body () 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 o -> Bool
sameAccType VName
acc_sglton (Acc VName
sglton Shape
_ [Type]
_) =
      VName
acc_sglton VName -> VName -> Bool
forall a. Eq a => a -> a -> Bool
== VName
sglton
    sameAccType VName
_ TypeBase shape o
_ = Bool
False
    getAccumFV :: Type -> BuilderT GPU (State VNameSource) VName
getAccumFV (Acc VName
singleton Shape
_shp [Type
_eltp]) = 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) [Type]
-> BuilderT GPU (State VNameSource) [Type]
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) [Type]
 -> BuilderT GPU (State VNameSource) [Type])
-> BuilderT GPU (State VNameSource) [Type]
-> BuilderT GPU (State VNameSource) [Type]
forall a b. (a -> b) -> a -> b
$ do
        (VName -> BuilderT GPU (State VNameSource) Type)
-> [VName] -> BuilderT GPU (State VNameSource) [Type]
forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
(a -> m b) -> t a -> m (t b)
forall (m :: * -> *) a b. Monad m => (a -> m b) -> [a] -> m [b]
mapM VName -> BuilderT GPU (State VNameSource) Type
forall rep (m :: * -> *). HasScope rep m => VName -> m Type
lookupType [VName]
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 Type
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]
++ Type -> [Char]
forall a. Pretty a => a -> [Char]
prettyString Type
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, Type)
-> (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, Type
res_tp)
      (SubExp
tx, SubExp
ty, SubExp
rx, SubExp
ry)
      (VName
iii, VName
jjj)
      (VName
gtid_x, VName
gtid_y)
      (SubExp
height_A, SubExp
width_B, [(VName, SubExp)]
_rem_outer_dims)
      Stms GPU
code2' = do
        rss_init <- Type -> BuilderT GPU (State VNameSource) VName
getAccumFV Type
res_tp
        rssss_list <- segMap2D "rssss" segthd_lvl ResultMaySimplify (tx, ty) $ \(VName
ltid_x, VName
ltid_y) -> 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 rx [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
ry [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 <-
                Name
-> Exp (Rep (BuilderT GPU (State VNameSource)))
-> BuilderT GPU (State VNameSource) SubExp
forall (m :: * -> *).
MonadBuilder m =>
Name -> Exp (Rep m) -> m SubExp
letSubExp Name
"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_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
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_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
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 [Type]
ts KernelBody GPU
old_kbody))))
  | Body () 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
    [Type
res_tp] <- [Type]
ts,
    Type -> Bool
forall shape o. TypeBase shape o -> Bool
primType Type
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_y, SubExp
width_B) : (VName
gtid_x, 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',
        (Stms GPU
load_A, VName
inp_A, PrimType
map_t1, Stms 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, (Stms GPU, VName, PrimType, Stms 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
      ---- 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
-> Builder
     GPU
     (SubExp, SubExp, SubExp, SubExp, SubExp, SubExp, SubExp, SubExp,
      SubExp, SubExp, SubExp)
mkTileMemSizes SubExp
height_A SubExp
width_B SubExp
common_dim

        gridDim_x <- letSubExp "gridDim_x" =<< ceilDiv height_A tx_rx
        gridDim_y <- letSubExp "gridDim_y" =<< ceilDiv width_B ty_ry
        let gridxy_pexp = 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_y
        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_x * pe64 tx_rx)
          jjj <- letExp "jjj" =<< toExp (le64 gid_y * pe64 ty_ry)

          -- 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
height_A, VName
gtid_y, SubExp
width_B, SubExp
common_dim),
                  (VName
iii, VName
jjj),
                  (Stms GPU
load_A, VName
inp_A, PrimType
map_t1, Stms 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),
    (Stms GPU, VName, PrimType, Stms 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),
 (Stms GPU, VName, PrimType, Stms 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_x < tx, ltid_y < ty) {
          --    for i < rx do
          --      for j < ry do
          --        res = if (iii+ltid_x*rx+i < height_A && jjj+ltid_y*ry+j < width_B)
          --              then code2' else dummy
          --        final_res[i,j] = res
          mkEpiloguePrimRes
            segthd_lvl
            (redomap_orig_res, redomap_res)
            (res_nm, res_tp)
            (tx, ty, rx, ry)
            (iii, jjj)
            (gtid_x, gtid_y)
            (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_x, SubExp
gridDim_x), (VName
gid_y, SubExp
gridDim_y)])
            kbody' = BodyDec GPU -> Stms GPU -> [KernelResult] -> KernelBody GPU
forall rep res. BodyDec rep -> Stms rep -> [res] -> GBody rep res
Body () 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, Type)
-> (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, Type
res_tp)
      (SubExp
tx, SubExp
ty, SubExp
rx, SubExp
ry)
      (VName
iii, VName
jjj)
      (VName
gtid_x, VName
gtid_y)
      (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 <- Name
-> SegLevel
-> ResultManifest
-> (SubExp, SubExp)
-> ((VName, VName) -> BuilderT GPU (State VNameSource) [SubExpRes])
-> Builder GPU [VName]
segMap2D Name
"rssss" SegLevel
segthd_lvl ResultManifest
ResultPrivate (SubExp
tx, SubExp
ty) (((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_x, VName
ltid_y) -> do
                rss_init <- Name
-> PrimType -> [SubExp] -> BuilderT GPU (State VNameSource) VName
forall (m :: * -> *).
MonadBuilder m =>
Name -> PrimType -> [SubExp] -> m VName
scratch Name
"rss_init" (Type -> PrimType
forall shape o. TypeBase shape o -> PrimType
elemType Type
res_tp) [SubExp
rx, SubExp
ry]
                (css, ii, jj) <- getThdRedomapRes (rx, ry) (ltid_x, ltid_y) (iii, jjj, redomap_res)
                rss <- forLoop rx [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
ry [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 <-
                      Name
-> Exp (Rep (BuilderT GPU (State VNameSource)))
-> BuilderT GPU (State VNameSource) SubExp
forall (m :: * -> *).
MonadBuilder m =>
Name -> Exp (Rep m) -> m SubExp
letSubExp Name
"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_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
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_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
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 [Type
-> BuilderT
     GPU
     (State VNameSource)
     (Exp (Rep (BuilderT GPU (State VNameSource))))
forall (m :: * -> *). MonadBuilder m => Type -> m (Exp (Rep m))
eBlank Type
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
tx, SubExp
rx), (SubExp
width_B, SubExp
ty, SubExp
ry)]

        -- 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,
      (Stms GPU, VName, PrimType, Stms GPU, VName, PrimType),
      SubExp,
      [Int],
      (Lambda GPU, Lambda GPU, SubExp, VName, PrimType)
    )
matchesBlkRegTile :: SegSpace
-> Stms GPU
-> Maybe
     (Stms GPU, (Stms GPU, VName, PrimType, Stms 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 OpC GPU 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,
    [Type
map_t1t, Type
map_t2t] <- (Param Type -> Type) -> [Param Type] -> [Type]
forall a b. (a -> b) -> [a] -> [b]
map Param Type -> Type
forall dec. Param dec -> dec
paramDec ([Param Type] -> [Type]) -> [Param Type] -> [Type]
forall a b. (a -> b) -> a -> b
$ Lambda GPU -> [LParam GPU]
forall rep t. GLambda rep t -> [LParam rep]
lambdaParams Lambda GPU
map_lam,
    [Type
red_t1, Type
_] <- (Param Type -> Type) -> [Param Type] -> [Type]
forall a b. (a -> b) -> [a] -> [b]
map Param Type -> Type
forall dec. Param dec -> dec
paramDec ([Param Type] -> [Type]) -> [Param Type] -> [Type]
forall a b. (a -> b) -> a -> b
$ Lambda GPU -> [LParam GPU]
forall rep t. GLambda rep t -> [LParam rep]
lambdaParams Lambda GPU
red_lam,
    Type -> Bool
forall shape o. TypeBase shape o -> Bool
primType Type
map_t1t Bool -> Bool -> Bool
&& Type -> Bool
forall shape o. TypeBase shape o -> Bool
primType Type
map_t2t Bool -> Bool -> Bool
&& Type -> Bool
forall shape o. TypeBase shape o -> Bool
primType Type
red_t1,
    PrimType
map_t1_0 <- Type -> PrimType
forall shape o. TypeBase shape o -> PrimType
elemType Type
map_t1t,
    PrimType
map_t2_0 <- Type -> PrimType
forall shape o. TypeBase shape o -> PrimType
elemType Type
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 Type -> [VName]
forall dec. Pat dec -> [VName]
patNames Pat Type
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 can be split into the
    -- statements producing the input for the redomap, and some
    -- statements that are moved after it; see `processIndirections`.
    Just (Stms GPU
code2'', Map VName (Stms GPU)
tab_inv_stm) <-
      Names
-> Names
-> Names
-> Stms GPU
-> Maybe (Stms GPU, Map VName (Stms GPU))
processIndirections ([VName] -> Names
namesFromList [VName]
arrs) Names
res_red_var (Stm GPU -> Names
forall a. FreeIn a => a -> Names
freeIn Stm GPU
screma_stmt) Stms GPU
code1,
    -- identify load_A, load_B
    [Stms GPU]
tmp_stms <- (VName -> Maybe (Stms GPU)) -> [VName] -> [Stms GPU]
forall a b. (a -> Maybe b) -> [a] -> [b]
mapMaybe (VName -> Map VName (Stms GPU) -> Maybe (Stms GPU)
forall k a. Ord k => k -> Map k a -> Maybe a
`M.lookup` Map VName (Stms GPU)
tab_inv_stm) [VName]
arrs,
    [Stms GPU] -> Int
forall a. [a] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length [Stms 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,
    -- If any tiled-loop input array is also used in the postlude code,
    -- we cannot safely apply this optimization (the array won't be in
    -- scope in the epilogue).  See issue #2467.
    Bool -> Bool
not (Bool -> Bool) -> Bool -> Bool
forall a b. (a -> b) -> a -> b
$
      [VName] -> Names
namesFromList (Map VName (Stms GPU) -> [VName]
forall k a. Map k a -> [k]
M.keys Map VName (Stms GPU)
tab_inv_stm)
        Names -> Names -> Bool
`namesIntersect` Stms GPU -> Names
forall a. FreeIn a => a -> Names
freeIn (Stms GPU
code2'' Stms GPU -> Stms GPU -> Stms GPU
forall a. Semigroup a => a -> a -> a
<> Stms GPU
code2) =
      let zip_AB :: [(Stms GPU, VName, PrimType)]
zip_AB = [Stms GPU]
-> [VName] -> [PrimType] -> [(Stms GPU, VName, PrimType)]
forall a b c. [a] -> [b] -> [c] -> [(a, b, c)]
zip3 [Stms GPU]
tmp_stms [VName]
arrs [PrimType
map_t1_0, PrimType
map_t2_0]
          [(Stms GPU
load_A, VName
inp_A, PrimType
map_t1), (Stms 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 [(Stms GPU, VName, PrimType)]
zip_AB
              else [(Stms GPU, VName, PrimType)] -> [(Stms GPU, VName, PrimType)]
forall a. [a] -> [a]
reverse [(Stms 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, (Stms GPU, VName, PrimType, Stms GPU, VName, PrimType),
 SubExp, [Int], (Lambda GPU, Lambda GPU, SubExp, VName, PrimType))
-> Maybe
     (Stms GPU, (Stms GPU, VName, PrimType, Stms GPU, VName, PrimType),
      SubExp, [Int], (Lambda GPU, Lambda GPU, SubExp, VName, PrimType))
forall a. a -> Maybe a
Just
            ( Stms GPU
code2',
              (Stms GPU
load_A, VName
inp_A, PrimType
map_t1, Stms 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, Type -> PrimType
forall shape o. TypeBase shape o -> PrimType
elemType Type
red_t1)
            )
matchesBlkRegTile SegSpace
_ Stms GPU
_ = Maybe
  (Stms GPU, (Stms GPU, VName, PrimType, Stms GPU, VName, PrimType),
   SubExp, [Int], (Lambda GPU, Lambda GPU, SubExp, VName, PrimType))
forall a. Maybe a
Nothing

-- | Round up to an odd number.  Shared memory strides are padded like
-- this to avoid bank conflicts.
oddUp :: TPrimExp Int64 VName -> TPrimExp Int64 VName
oddUp :: TPrimExp Int64 VName -> TPrimExp Int64 VName
oddUp TPrimExp Int64 VName
x = TPrimExp Int64 VName
2 TPrimExp Int64 VName
-> TPrimExp Int64 VName -> TPrimExp Int64 VName
forall a. Num a => a -> a -> a
* (TPrimExp Int64 VName
x TPrimExp Int64 VName
-> TPrimExp Int64 VName -> TPrimExp Int64 VName
forall e. IntegralExp e => e -> e -> e
`IE.quot` TPrimExp Int64 VName
2) TPrimExp Int64 VName
-> TPrimExp Int64 VName -> TPrimExp Int64 VName
forall a. Num a => a -> a -> a
+ TPrimExp Int64 VName
1

-- | 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
SCeilDiv IntType
Int64 Safety
Unsafe) SubExp
x SubExp
y

mkTileMemSizes ::
  SubExp ->
  SubExp ->
  SubExp ->
  Builder
    GPU
    ( SubExp,
      SubExp,
      SubExp,
      SubExp,
      SubExp,
      SubExp,
      SubExp,
      SubExp,
      SubExp,
      SubExp,
      SubExp
    )
mkTileMemSizes :: SubExp
-> SubExp
-> SubExp
-> Builder
     GPU
     (SubExp, SubExp, SubExp, SubExp, SubExp, SubExp, SubExp, SubExp,
      SubExp, SubExp, SubExp)
mkTileMemSizes SubExp
height_A SubExp
_width_B SubExp
common_dim = do
  tk_name <- Text -> Name
nameFromText (Text -> Name) -> (VName -> Text) -> VName -> Name
forall b c a. (b -> c) -> (a -> b) -> a -> c
. VName -> Text
forall a. Pretty a => a -> Text
prettyText (VName -> Name)
-> BuilderT GPU (State VNameSource) VName
-> BuilderT GPU (State VNameSource) Name
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Name -> BuilderT GPU (State VNameSource) VName
forall (m :: * -> *). MonadFreshNames m => Name -> m VName
newVName Name
"Tk"
  tx_name <- nameFromText . prettyText <$> newVName "Tx"
  rx_name <- nameFromText . prettyText <$> newVName "Rx"

  -- 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!
  (tx, rx) <- getParTiles ("Tx", "Rx") (tx_name, rx_name) height_A
  let (ty, ry) = (tx, rx)
  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)

  -- Large enough for either layout, including padding (see mk_ik).
  a_loc_sz <-
    letSubExp "a_loc_sz"
      =<< toExp (pe64 tx * oddUp (pe64 rx * pe64 tk))
  b_loc_sz <-
    letSubExp "b_loc_sz"
      =<< toExp
        ( pe64 ty
            * sMax64 (oddUp (pe64 ry * pe64 tk)) (oddUp (pe64 ry) * pe64 tk)
        )
  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 <- Name
-> Exp (Rep (BuilderT GPU (State VNameSource)))
-> BuilderT GPU (State VNameSource) SubExp
forall (m :: * -> *).
MonadBuilder m =>
Name -> Exp (Rep m) -> m SubExp
letSubExp Name
"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 tx * pe64 ty)
  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_x <- Name -> BuilderT GPU (State VNameSource) VName
forall (m :: * -> *). MonadFreshNames m => Name -> m VName
newVName Name
"gid_x"
  gid_y <- newVName "gid_y"
  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 <- Name
-> SegLevel
-> ResultManifest
-> (SubExp, SubExp)
-> ((VName, VName) -> BuilderT GPU (State VNameSource) [SubExpRes])
-> Builder GPU [VName]
segMap2D Name
"cssss" SegLevel
segthd_lvl ResultManifest
ResultPrivate (SubExp
tx, SubExp
ty) (((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 <- Name
-> PrimType -> [SubExp] -> BuilderT GPU (State VNameSource) VName
forall (m :: * -> *).
MonadBuilder m =>
Name -> PrimType -> [SubExp] -> m VName
scratch Name
"css_init" PrimType
red_t [SubExp
rx, SubExp
ry]
      css <- forLoop rx [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
ry [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'' <- Name
-> VName
-> [VName]
-> SubExp
-> BuilderT GPU (State VNameSource) VName
forall (m :: * -> *).
MonadBuilder m =>
Name -> VName -> [VName] -> SubExp -> m VName
update Name
"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 <- Name -> VName -> [VName] -> BuilderT GPU (State VNameSource) VName
forall (m :: * -> *).
MonadBuilder m =>
Name -> VName -> [VName] -> m VName
index Name
"redomap_thd" VName
redomap_res [VName
ltid_x, VName
ltid_y]
  ii <- letExp "ii" =<< toExp (le64 iii + le64 ltid_x * pe64 rx)
  jj <- letExp "jj" =<< toExp (le64 jjj + le64 ltid_y * pe64 ry)
  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 <- Name -> VName -> [VName] -> BuilderT GPU (State VNameSource) VName
forall (m :: * -> *).
MonadBuilder m =>
Name -> VName -> [VName] -> m VName
index Name
"redomap_elm" VName
css [VName
i, VName
j]
  cpy_stm <- mkLetNamesM [redomap_orig_res] $ BasicOp $ SubExp $ Var c
  addStm cpy_stm
  letBindNames [gtid_x] =<< toExp (le64 ii + le64 i)
  letBindNames [gtid_y] =<< 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

-- | Split the statements preceding the redomap into
--
-- 1. the slices that produce the input arrays of the redomap, each
--    preceded by the scalar statements computing its indices (these
--    are re-executed when reading tiles), and
--
-- 2. the remaining statements, which are moved after the redomap.
--    Except for those computing slice indices, the redomap must not
--    depend on these.
processIndirections ::
  Names -> -- input arrays to redomap
  Names -> -- variables on which the result of redomap depends on.
  Names -> -- free variables of the redomap
  Stms GPU ->
  Maybe (Stms GPU, M.Map VName (Stms GPU))
processIndirections :: Names
-> Names
-> Names
-> Stms GPU
-> Maybe (Stms GPU, Map VName (Stms GPU))
processIndirections Names
arrs Names
res_red_var Names
red_free Stms GPU
code1 = do
  let (Stms GPU
loads, Stms GPU
rest) = (Stm GPU -> Bool) -> Stms GPU -> (Stms GPU, Stms GPU)
forall a. (a -> Bool) -> Seq a -> (Seq a, Seq a)
Seq.partition Stm GPU -> Bool
isLoad Stms GPU
code1
      withDeps :: Stm GPU -> Stms GPU
withDeps Stm GPU
stm = Stms GPU -> Stm GPU -> Stms GPU
dependencies Stms GPU
rest Stm GPU
stm Stms GPU -> Stm GPU -> Stms GPU
forall a. Seq a -> a -> Seq a
Seq.|> Stm GPU
stm
      deps :: Stms GPU
deps = (Stm GPU -> Stms GPU) -> Stms GPU -> Stms GPU
forall m a. Monoid m => (a -> m) -> Seq a -> m
forall (t :: * -> *) m a.
(Foldable t, Monoid m) =>
(a -> m) -> t a -> m
foldMap (Stms GPU -> Stm GPU -> Stms GPU
dependencies Stms GPU
rest) Stms GPU
loads
  Bool -> Maybe ()
forall (f :: * -> *). Alternative f => Bool -> f ()
guard (Bool -> Maybe ()) -> Bool -> Maybe ()
forall a b. (a -> b) -> a -> b
$ (Stm GPU -> Bool) -> Stms GPU -> Bool
forall (t :: * -> *) a. Foldable t => (a -> Bool) -> t a -> Bool
all Stm GPU -> Bool
forall {rep}. Typed (LetDec rep) => Stm rep -> Bool
isScalarBasicOp Stms GPU
deps
  Bool -> Maybe ()
forall (f :: * -> *). Alternative f => Bool -> f ()
guard (Bool -> Maybe ()) -> Bool -> Maybe ()
forall a b. (a -> b) -> a -> b
$ Bool -> Bool
not (Bool -> Bool) -> Bool -> Bool
forall a b. (a -> b) -> a -> b
$ Stms GPU -> Names
forall rep. Stms rep -> Names
boundByStms Stms GPU
deps Names -> Names -> Bool
`namesIntersect` Names
red_free
  Bool -> Maybe ()
forall (f :: * -> *). Alternative f => Bool -> f ()
guard (Bool -> Maybe ()) -> Bool -> Maybe ()
forall a b. (a -> b) -> a -> b
$ Bool -> Bool
not (Bool -> Bool) -> Bool -> Bool
forall a b. (a -> b) -> a -> b
$ (Stms GPU -> Names
forall rep. Stms rep -> Names
boundByStms Stms GPU
rest Names -> Names -> Names
`namesSubtract` Stms GPU -> Names
forall rep. Stms rep -> Names
boundByStms Stms GPU
deps) Names -> Names -> Bool
`namesIntersect` Names
res_red_var
  (Stms GPU, Map VName (Stms GPU))
-> Maybe (Stms GPU, Map VName (Stms GPU))
forall a. a -> Maybe a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (Stms GPU
rest, [(VName, Stms GPU)] -> Map VName (Stms GPU)
forall k a. Ord k => [(k, a)] -> Map k a
M.fromList [(VName
v, Stm GPU -> Stms GPU
withDeps Stm GPU
stm) | Stm GPU
stm <- Stms GPU -> [Stm GPU]
forall rep. Stms rep -> [Stm rep]
stmsToList Stms GPU
loads, VName
v <- Pat Type -> [VName]
forall dec. Pat dec -> [VName]
patNames (Pat Type -> [VName]) -> Pat Type -> [VName]
forall a b. (a -> b) -> a -> b
$ Stm GPU -> Pat (LetDec GPU)
forall rep. Stm rep -> Pat (LetDec rep)
stmPat Stm GPU
stm])
  where
    isLoad :: Stm GPU -> Bool
isLoad (Let (Pat [PatElem (LetDec GPU)
pe]) StmAux (ExpDec GPU)
_ (BasicOp Index {})) = PatElem Type -> VName
forall dec. PatElem dec -> VName
patElemName PatElem Type
PatElem (LetDec GPU)
pe VName -> Names -> Bool
`nameIn` Names
arrs
    isLoad Stm GPU
_ = Bool
False
    isScalarBasicOp :: Stm rep -> Bool
isScalarBasicOp (Let Pat (LetDec rep)
pat StmAux (ExpDec rep)
_ BasicOp {}) = (Type -> Bool) -> [Type] -> Bool
forall (t :: * -> *) a. Foldable t => (a -> Bool) -> t a -> Bool
all Type -> Bool
forall shape o. TypeBase shape o -> Bool
primType ([Type] -> Bool) -> [Type] -> Bool
forall a b. (a -> b) -> a -> b
$ Pat (LetDec rep) -> [Type]
forall dec. Typed dec => Pat dec -> [Type]
patTypes Pat (LetDec rep)
pat
    isScalarBasicOp Stm rep
_ = Bool
False

-- | The statements among the given ones that the statement
-- transitively depends on.
dependencies :: Stms GPU -> Stm GPU -> Stms GPU
dependencies :: Stms GPU -> Stm GPU -> Stms GPU
dependencies Stms GPU
stms Stm GPU
stm = (Names, Stms GPU) -> Stms GPU
forall a b. (a, b) -> b
snd ((Names, Stms GPU) -> Stms GPU) -> (Names, Stms GPU) -> Stms GPU
forall a b. (a -> b) -> a -> b
$ (Stm GPU -> (Names, Stms GPU) -> (Names, Stms GPU))
-> (Names, Stms GPU) -> Stms GPU -> (Names, Stms GPU)
forall a b. (a -> b -> b) -> b -> Seq a -> b
forall (t :: * -> *) a b.
Foldable t =>
(a -> b -> b) -> b -> t a -> b
foldr Stm GPU -> (Names, Stms GPU) -> (Names, Stms GPU)
forall {rep}.
(FreeDec (ExpDec rep), FreeDec (BodyDec rep), FreeIn (OpC rep rep),
 FreeIn (FParamInfo rep), FreeIn (LParamInfo rep),
 FreeIn (LetDec rep), FreeIn (RetType rep),
 FreeIn (BranchType rep)) =>
Stm rep -> (Names, Seq (Stm rep)) -> (Names, Seq (Stm rep))
f (Stm GPU -> Names
forall a. FreeIn a => a -> Names
freeIn Stm GPU
stm, Stms GPU
forall a. Monoid a => a
mempty) Stms GPU
stms
  where
    f :: Stm rep -> (Names, Seq (Stm rep)) -> (Names, Seq (Stm rep))
f Stm rep
s (Names
needed, Seq (Stm rep)
deps)
      | Stm rep -> Names
forall rep. Stm rep -> Names
boundByStm Stm rep
s Names -> Names -> Bool
`namesIntersect` Names
needed = (Names
needed Names -> Names -> Names
forall a. Semigroup a => a -> a -> a
<> Stm rep -> Names
forall a. FreeIn a => a -> Names
freeIn Stm rep
s, Stm rep
s Stm rep -> Seq (Stm rep) -> Seq (Stm rep)
forall a. a -> Seq a -> Seq a
Seq.<| Seq (Stm rep)
deps)
      | Bool
otherwise = (Names
needed, Seq (Stm rep)
deps)

getParTiles :: (Name, Name) -> (Name, Name) -> SubExp -> Builder GPU (SubExp, SubExp)
getParTiles :: (Name, Name)
-> (Name, Name) -> SubExp -> Builder GPU (SubExp, SubExp)
getParTiles (Name
t_str, Name
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 <- Name
-> Exp (Rep (BuilderT GPU (State VNameSource)))
-> BuilderT GPU (State VNameSource) SubExp
forall (m :: * -> *).
MonadBuilder m =>
Name -> Exp (Rep m) -> m SubExp
letSubExp Name
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)

-- | The sequential tile size. It is rounded up to a multiple of @tx@, as the
-- tile loops assume that @tx@ divides it.
getSeqTile :: Name -> Name -> SubExp -> SubExp -> SubExp -> Builder GPU SubExp
getSeqTile :: Name
-> Name
-> SubExp
-> SubExp
-> SubExp
-> BuilderT GPU (State VNameSource) SubExp
getSeqTile Name
tk_str Name
tk_name SubExp
len_dim SubExp
tx SubExp
ty = do
  tk <- case (SubExp
tx, SubExp
ty) of
    (Constant (IntValue (Int64Value Int64
v_x)), Constant (IntValue (Int64Value Int64
v_y))) ->
      Name
-> Exp (Rep (BuilderT GPU (State VNameSource)))
-> BuilderT GPU (State VNameSource) SubExp
forall (m :: * -> *).
MonadBuilder m =>
Name -> Exp (Rep m) -> m SubExp
letSubExp Name
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)
_ ->
      Name
-> Exp (Rep (BuilderT GPU (State VNameSource)))
-> BuilderT GPU (State VNameSource) SubExp
forall (m :: * -> *).
MonadBuilder m =>
Name -> Exp (Rep m) -> m SubExp
letSubExp Name
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
  tk_div_tx <- letSubExp "tk_div_tx" =<< ceilDiv tk tx
  letSubExp tk_str =<< toExp (pe64 tk_div_tx * pe64 tx)

----------------------------------------------------------------------------------------------
--- 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: see `processIndirections`.
--
-- 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 [Type]
kertp (Body () 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_z, SubExp
d_Kz) : (VName
gtid_y, SubExp
d_Ky) : (VName
gtid_x, SubExp
d_Kx) : [(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 OpC GPU 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 Type -> Bool) -> [Param Type] -> Bool
forall (t :: * -> *) a. Foldable t => (a -> Bool) -> t a -> Bool
all (Type -> Bool
forall shape o. TypeBase shape o -> Bool
primType (Type -> Bool) -> (Param Type -> Type) -> Param Type -> Bool
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Param Type -> Type
forall dec. Param dec -> dec
paramDec) ([Param Type] -> Bool) -> [Param Type] -> Bool
forall a b. (a -> b) -> a -> b
$ Lambda GPU -> [LParam GPU]
forall rep t. GLambda rep t -> [LParam rep]
lambdaParams Lambda GPU
map_lam,
    [Type]
red_res_tps <- (Param Type -> Type) -> [Param Type] -> [Type]
forall a b. (a -> b) -> [a] -> [b]
map Param Type -> Type
forall dec. Param dec -> dec
paramDec ([Param Type] -> [Type]) -> [Param Type] -> [Type]
forall a b. (a -> b) -> a -> b
$ Int -> [Param Type] -> [Param Type]
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 Type] -> [Param Type]) -> [Param Type] -> [Param Type]
forall a b. (a -> b) -> a -> b
$ Lambda GPU -> [LParam GPU]
forall rep t. GLambda rep t -> [LParam rep]
lambdaParams Lambda GPU
red_lam,
    (Type -> Bool) -> [Type] -> Bool
forall (t :: * -> *) a. Foldable t => (a -> Bool) -> t a -> Bool
all Type -> Bool
forall shape o. TypeBase shape o -> Bool
primType [Type]
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 Type]
redomap_orig_res <- Pat Type -> [PatElem Type]
forall dec. Pat dec -> [PatElem dec]
patElems Pat Type
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 Type -> Maybe Names) -> [PatElem Type] -> [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 Type -> VName) -> PatElem Type -> Maybe Names
forall b c a. (b -> c) -> (a -> b) -> a -> c
. PatElem Type -> VName
forall dec. PatElem dec -> VName
patElemName) [PatElem Type]
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 can be split into the
    -- statements producing the input for the redomap, and some
    -- statements that are moved after it; see `processIndirections`.
    Just (Stms GPU
code2'', Map VName (Stms GPU)
arr_tab0) <-
      Names
-> Names
-> Names
-> Stms GPU
-> Maybe (Stms GPU, Map VName (Stms GPU))
processIndirections ([VName] -> Names
namesFromList [VName]
inp_soac_arrs) Names
res_red_var (Stm GPU -> Names
forall a. FreeIn a => a -> Names
freeIn Stm GPU
screma_stmt) Stms GPU
code1,
    -- check that code1 contains exacly one slice for each of the input array to redomap
    [Stms GPU]
tmp_stms <- (VName -> Maybe (Stms GPU)) -> [VName] -> [Stms GPU]
forall a b. (a -> Maybe b) -> [a] -> [b]
mapMaybe (VName -> Map VName (Stms GPU) -> Maybe (Stms GPU)
forall k a. Ord k => k -> Map k a -> Maybe a
`M.lookup` Map VName (Stms GPU)
arr_tab0) [VName]
inp_soac_arrs,
    [Stms GPU] -> Int
forall a. [a] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length [Stms 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 outermost of the three
    -- innermost parallel dimensions (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,
    (Type -> Bool) -> [Type] -> Bool
forall (t :: * -> *) a. Foldable t => (a -> Bool) -> t a -> Bool
all Type -> Bool
forall shape o. TypeBase shape o -> Bool
primType [Type]
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_x) [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 x-variant arrays that need transposition;
        -- these "manifest" statements should come before the kernel
        (tab_inn, tab_out) <-
          ((Map VName (Stms GPU), Map VName (PrimType, Stms GPU))
 -> (VName, Stms GPU)
 -> BuilderT
      GPU
      (State VNameSource)
      (Map VName (Stms GPU), Map VName (PrimType, Stms GPU)))
-> (Map VName (Stms GPU), Map VName (PrimType, Stms GPU))
-> [(VName, Stms GPU)]
-> BuilderT
     GPU
     (State VNameSource)
     (Map VName (Stms GPU), Map VName (PrimType, Stms 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 (Stms GPU), Map VName (PrimType, Stms GPU))
-> (VName, Stms GPU)
-> BuilderT
     GPU
     (State VNameSource)
     (Map VName (Stms GPU), Map VName (PrimType, Stms GPU))
insertTranspose Map VName Names
variance (VName
gtid_x, SubExp
d_Kx))
            (Map VName (Stms GPU)
forall k a. Map k a
M.empty, Map VName (PrimType, Stms GPU)
forall k a. Map k a
M.empty)
            ([(VName, Stms GPU)]
 -> BuilderT
      GPU
      (State VNameSource)
      (Map VName (Stms GPU), Map VName (PrimType, Stms GPU)))
-> [(VName, Stms GPU)]
-> BuilderT
     GPU
     (State VNameSource)
     (Map VName (Stms GPU), Map VName (PrimType, Stms GPU))
forall a b. (a -> b) -> a -> b
$ Map VName (Stms GPU) -> [(VName, Stms GPU)]
forall k a. Map k a -> [(k, a)]
M.toList Map VName (Stms GPU)
arr_tab0

        tz_name <- nameFromText . prettyText <$> newVName "Tz"
        ty_name <- nameFromText . prettyText <$> newVName "Ty"

        tz0 <- letSubExp "Tz" $ Op $ SizeOp $ GetSize tz_name SizeTile
        ty0 <- letSubExp "Ty" $ Op $ SizeOp $ GetSize ty_name SizeTile
        tz <- limitTile "Tz" tz0 d_Kz
        ty <- limitTile "Ty" ty0 d_Ky
        let rx = SubExp
reg_tile_se

        gridDim_z <- letSubExp "gridDim_z" =<< ceilDiv d_Kz tz
        gridDim_y <- letSubExp "gridDim_y" =<< ceilDiv d_Ky ty
        gridDim_x <- letSubExp "gridDim_x" =<< ceilDiv d_Kx rx
        let gridxyz_pexp = 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_y TPrimExp Int64 VName
-> TPrimExp Int64 VName -> TPrimExp Int64 VName
forall a. Num a => a -> a -> a
* SubExp -> TPrimExp Int64 VName
pe64 SubExp
gridDim_z
        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 tz)
        let segthd_lvl = SegVirt -> SegLevel
SegThreadInBlock (SegSeqDims -> SegVirt
SegNoVirtFull ([Int] -> SegSeqDims
SegSeqDims []))

        count_shmem <- letSubExp "count_shmem" =<< ceilDiv rx 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_x * pe64 rx)
          jj1 <- letExp "jj1" =<< toExp (le64 gid_y * pe64 ty)
          jj2 <- letExp "jj2" =<< toExp (le64 gid_z * pe64 tz)

          -- initialize the register arrays corresponding to the result of redomap;
          reg_arr_nms <- segMap2D "res" segthd_lvl ResultPrivate (ty, tz) $ \(VName, VName)
_ ->
            [(SubExp, Type)]
-> ((SubExp, Type) -> 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] -> [Type] -> [(SubExp, Type)]
forall a b. [a] -> [b] -> [(a, b)]
zip [SubExp]
red_nes [Type]
red_res_tps) (((SubExp, Type) -> BuilderT GPU (State VNameSource) SubExpRes)
 -> BuilderT GPU (State VNameSource) [SubExpRes])
-> ((SubExp, Type) -> BuilderT GPU (State VNameSource) SubExpRes)
-> BuilderT GPU (State VNameSource) [SubExpRes]
forall a b. (a -> b) -> a -> b
$ \(SubExp
red_ne, Type
red_t) -> do
              css_init <- Name
-> PrimType -> [SubExp] -> BuilderT GPU (State VNameSource) VName
forall (m :: * -> *).
MonadBuilder m =>
Name -> PrimType -> [SubExp] -> m VName
scratch Name
"res_init" (Type -> PrimType
forall shape o. TypeBase shape o -> PrimType
elemType Type
red_t) [SubExp
rx]
              css <- forLoop rx [css_init] $ \VName
i [VName
css_merge] -> do
                css' <- Name
-> VName
-> [VName]
-> SubExp
-> BuilderT GPU (State VNameSource) VName
forall (m :: * -> *).
MonadBuilder m =>
Name -> VName -> [VName] -> SubExp -> m VName
update Name
"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, Stms GPU
_)) ->
            Name
-> PrimType -> [SubExp] -> BuilderT GPU (State VNameSource) VName
forall (m :: * -> *).
MonadBuilder m =>
Name -> PrimType -> [SubExp] -> m VName
scratch (VName -> Name
baseName VName
nm Name -> Name -> Name
forall a. Semigroup a => a -> a -> a
<> Name
"_loc") PrimType
ptp [SubExp
rx]

          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, Stms GPU)))]
-> ((VName, (VName, (PrimType, Stms 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, Stms GPU))]
-> [(VName, (VName, (PrimType, Stms GPU)))]
forall a b. [a] -> [b] -> [(a, b)]
zip [VName]
loc_arr_merge2_nms (Map VName (PrimType, Stms GPU) -> [(VName, (PrimType, Stms GPU))]
forall k a. Map k a -> [(k, a)]
M.toList Map VName (PrimType, Stms GPU)
tab_out)) (((VName, (VName, (PrimType, Stms GPU)))
  -> BuilderT GPU (State VNameSource) VName)
 -> Builder GPU [VName])
-> ((VName, (VName, (PrimType, Stms 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, Stms GPU
load_Y))) ->
                        Name
-> Exp (Rep (BuilderT GPU (State VNameSource)))
-> BuilderT GPU (State VNameSource) VName
forall (m :: * -> *).
MonadBuilder m =>
Name -> Exp (Rep m) -> m VName
letExp Name
"Y_glb2loc" (Exp GPU -> BuilderT GPU (State VNameSource) VName)
-> (([VName] -> BuilderT GPU (State VNameSource) [SubExp])
    -> BuilderT GPU (State VNameSource) (Exp GPU))
-> ([VName] -> BuilderT GPU (State VNameSource) [SubExp])
-> BuilderT GPU (State VNameSource) VName
forall (m :: * -> *) b c a.
Monad m =>
(b -> m c) -> (a -> m b) -> a -> m c
<=< [VName]
-> Int
-> ([VName] -> BuilderT GPU (State VNameSource) [SubExp])
-> BuilderT
     GPU
     (State VNameSource)
     (Exp (Rep (BuilderT GPU (State VNameSource))))
forall (m :: * -> *).
(MonadBuilder m, LParam (Rep m) ~ Param Type) =>
[VName] -> Int -> ([VName] -> m [SubExp]) -> m (Exp (Rep m))
withAcc [VName
loc_Y_nm] Int
1 (([VName] -> BuilderT GPU (State VNameSource) [SubExp])
 -> BuilderT GPU (State VNameSource) VName)
-> ([VName] -> BuilderT GPU (State VNameSource) [SubExp])
-> BuilderT GPU (State VNameSource) VName
forall a b. (a -> b) -> a -> b
$ \ ~[VName
acc] -> do
                          ltid_flat <- Name -> BuilderT GPU (State VNameSource) VName
forall (m :: * -> *). MonadFreshNames m => Name -> m VName
newVName Name
"ltid_flat"
                          ltid <- newVName "ltid"
                          let segspace = VName -> [(VName, SubExp)] -> SegSpace
SegSpace VName
ltid_flat [(VName
ltid, SubExp
tblock_size)]
                          body <- runBodyBuilder $ do
                            offs <- letExp "offs" =<< toExp (pe64 tblock_size * le64 tt)
                            loc_ind <- letExp "loc_ind" =<< toExp (le64 ltid + le64 offs)
                            letBindNames [gtid_x] =<< toExp (le64 ii + le64 loc_ind)
                            let glb_ind = VName
gtid_x
                            y_elm <-
                              letSubExp "y_elem"
                                =<< eIf
                                  (toExp $ le64 glb_ind .<. pe64 d_Kx)
                                  ( do
                                      addStms 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 rx)
                                  (toExp loc_ind >>= letTupExp' "loc_fi" >>= resultBodyM)
                                  (eBody [pure $ BasicOp $ SubExp $ intConst Int64 (-1)])
                            acc' <- letExp (baseName acc) $ BasicOp $ UpdateAcc Safe acc [y_ind] [y_elm]
                            pure [Returns ResultMaySimplify mempty $ Var acc']

                          acc_t <- lookupType acc

                          letTupExp' "Y_glb2loc" <=< renameExp $
                            Op . SegOp $
                              SegMap segthd_lvl segspace [acc_t] body
                    resultBodyM $ map Var loc_arr_merge2_nms'

                redomap_res <-
                  segMap2D "redomap_res" segthd_lvl ResultPrivate (ty, tz) $
                    \(VName
ltid_y, VName
ltid_z) -> 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_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
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_z)
                      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 ->
                        Name -> VName -> [VName] -> BuilderT GPU (State VNameSource) VName
forall (m :: * -> *).
MonadBuilder m =>
Name -> VName -> [VName] -> m VName
index Name
"res_reg_slc" VName
reg_arr_nm [VName
ltid_y, VName
ltid_z]
                      fmap subExpsRes . letTupExp' "redomap_guarded"
                        =<< eIf
                          (toExp $ le64 gtid_y .<. pe64 d_Ky .&&. le64 gtid_z .<. pe64 d_Kz)
                          ( do
                              inp_scals_invar_outer <-
                                forM (M.toList tab_inn) $ \(VName
inp_arr_nm, Stms GPU
load_stm) -> 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
load_stm
                                  Name -> VName -> [VName] -> BuilderT GPU (State VNameSource) VName
forall (m :: * -> *).
MonadBuilder m =>
Name -> VName -> [VName] -> m VName
index (VName -> Name
baseName 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' rx 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_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
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
=<< Name
-> Exp (Rep (BuilderT GPU (State VNameSource)))
-> BuilderT GPU (State VNameSource) [SubExp]
forall (m :: * -> *).
MonadBuilder m =>
Name -> Exp (Rep m) -> m [SubExp]
letTupExp' Name
"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_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)
                                      ( 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 ->
                                            Name -> VName -> [VName] -> BuilderT GPU (State VNameSource) VName
forall (m :: * -> *).
MonadBuilder m =>
Name -> VName -> [VName] -> m VName
index Name
"inp_reg_var2x" VName
loc_arr_nm [VName
i]
                                          cs <- forM reg_arr_mm_nms $ \VName
reg_arr_nm ->
                                            Name -> VName -> [VName] -> BuilderT GPU (State VNameSource) VName
forall (m :: * -> *).
MonadBuilder m =>
Name -> VName -> [VName] -> m VName
index Name
"res_reg_var2x" 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, Stms GPU)) -> VName)
-> [(VName, (PrimType, Stms GPU))] -> [VName]
forall a b. (a -> b) -> [a] -> [b]
map (VName, (PrimType, Stms GPU)) -> VName
forall a b. (a, b) -> a
fst ([(VName, (PrimType, Stms GPU))] -> [VName])
-> [(VName, (PrimType, Stms GPU))] -> [VName]
forall a b. (a -> b) -> a -> b
$ Map VName (PrimType, Stms GPU) -> [(VName, (PrimType, Stms GPU))]
forall k a. Map k a -> [(k, a)]
M.toList Map VName (PrimType, Stms 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, Stms GPU) -> VName) -> [(VName, Stms GPU)] -> [VName]
forall a b. (a -> b) -> [a] -> [b]
map (VName, Stms GPU) -> VName
forall a b. (a, b) -> a
fst ([(VName, Stms GPU)] -> [VName]) -> [(VName, Stms GPU)] -> [VName]
forall a b. (a -> b) -> a -> b
$ Map VName (Stms GPU) -> [(VName, Stms GPU)]
forall k a. Map k a -> [(k, a)]
M.toList Map VName (Stms 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) ->
                                            Name
-> VName
-> [VName]
-> SubExp
-> BuilderT GPU (State VNameSource) VName
forall (m :: * -> *).
MonadBuilder m =>
Name -> VName -> [VName] -> SubExp -> m VName
update (VName -> Name
baseName 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_z < tz) {
          --    for i < rx do
          --        res = if (ii+i < d_Kx && jj1+ltid_y < d_Ky && jj2 + ltid_z < d_Kz)
          --              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, tz) $ \(VName
_ltid_x, VName
ltid_y, VName
ltid_z) ->
                [(Type, VName)]
-> ((Type, 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 ([Type] -> [VName] -> [(Type, VName)]
forall a b. [a] -> [b] -> [(a, b)]
zip [Type]
kertp [VName]
redomap_res) (((Type, VName) -> BuilderT GPU (State VNameSource) SubExpRes)
 -> BuilderT GPU (State VNameSource) [SubExpRes])
-> ((Type, VName) -> BuilderT GPU (State VNameSource) SubExpRes)
-> BuilderT GPU (State VNameSource) [SubExpRes]
forall a b. (a -> b) -> a -> b
$ \(Type
res_tp, VName
res) -> do
                  rss_init <- Name
-> PrimType -> [SubExp] -> BuilderT GPU (State VNameSource) VName
forall (m :: * -> *).
MonadBuilder m =>
Name -> PrimType -> [SubExp] -> m VName
scratch Name
"rss_init" (Type -> PrimType
forall shape o. TypeBase shape o -> PrimType
elemType Type
res_tp) [SubExp
rx, SubExp
se1, SubExp
se1]
                  fmap varRes $
                    forLoop rx [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 <- Name -> VName -> [VName] -> BuilderT GPU (State VNameSource) VName
forall (m :: * -> *).
MonadBuilder m =>
Name -> VName -> [VName] -> m VName
index Name
"thread_res" VName
res [VName
ltid_y, VName
ltid_z, VName
i]
                      rss' <- letSubExp "rss" $ BasicOp $ Update Unsafe rss slice $ Var thread_res
                      resultBodyM [rss']
              else segMap3D "rssss" segthd_lvl ResultPrivate (se1, ty, tz) $ \(VName
_ltid_x, VName
ltid_y, VName
ltid_z) -> 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_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
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_z)
                rss_init <- [Type]
-> (Type -> 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 [Type]
kertp ((Type -> BuilderT GPU (State VNameSource) VName)
 -> Builder GPU [VName])
-> (Type -> BuilderT GPU (State VNameSource) VName)
-> Builder GPU [VName]
forall a b. (a -> b) -> a -> b
$ \Type
res_tp ->
                  Name
-> PrimType -> [SubExp] -> BuilderT GPU (State VNameSource) VName
forall (m :: * -> *).
MonadBuilder m =>
Name -> PrimType -> [SubExp] -> m VName
scratch Name
"rss_init" (Type -> PrimType
forall shape o. TypeBase shape o -> PrimType
elemType Type
res_tp) [SubExp
rx, SubExp
se1, SubExp
se1]
                rss <- forLoop' rx 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_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
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 Type, VName)]
-> ((PatElem Type, 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 Type] -> [VName] -> [(PatElem Type, VName)]
forall a b. [a] -> [b] -> [(a, b)]
zip [PatElem Type]
redomap_orig_res [VName]
redomap_res) (((PatElem Type, VName) -> BuilderT GPU (State VNameSource) VName)
 -> BuilderT GPU (State VNameSource) ())
-> ((PatElem Type, VName)
    -> BuilderT GPU (State VNameSource) VName)
-> BuilderT GPU (State VNameSource) ()
forall a b. (a -> b) -> a -> b
$ \(PatElem Type
o_res, VName
n_res) -> do
                    c <- Name -> VName -> [VName] -> BuilderT GPU (State VNameSource) VName
forall (m :: * -> *).
MonadBuilder m =>
Name -> VName -> [VName] -> m VName
index Name
"redomap_thd" VName
n_res [VName
ltid_y, VName
ltid_z, VName
i]
                    letBindNames [patElemName o_res] =<< toExp (le64 c)
                    pure c
                  res_els <-
                    Name
-> Exp (Rep (BuilderT GPU (State VNameSource)))
-> BuilderT GPU (State VNameSource) [SubExp]
forall (m :: * -> *).
MonadBuilder m =>
Name -> Exp (Rep m) -> m [SubExp]
letTupExp' Name
"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_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_Kz
                              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
                        )
                        ( 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
$ (Type
 -> BuilderT
      GPU
      (State VNameSource)
      (Exp (Rep (BuilderT GPU (State VNameSource)))))
-> [Type]
-> [BuilderT
      GPU
      (State VNameSource)
      (Exp (Rep (BuilderT GPU (State VNameSource))))]
forall a b. (a -> b) -> [a] -> [b]
map Type
-> BuilderT
     GPU
     (State VNameSource)
     (Exp (Rep (BuilderT GPU (State VNameSource))))
forall (m :: * -> *). MonadBuilder m => Type -> m (Exp (Rep m))
eBlank [Type]
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]
                    Name
-> Exp (Rep (BuilderT GPU (State VNameSource)))
-> BuilderT GPU (State VNameSource) SubExp
forall (m :: * -> *).
MonadBuilder m =>
Name -> Exp (Rep m) -> m SubExp
letSubExp Name
"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_Kx, SubExp
se1, SubExp
rx), (SubExp
d_Ky, SubExp
ty, SubExp
se1), (SubExp
d_Kz, SubExp
tz, 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) Type
forall rep (m :: * -> *). HasScope rep m => VName -> m Type
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_x, SubExp
gridDim_x), (VName
gid_y, SubExp
gridDim_y), (VName
gid_z, SubExp
gridDim_z)])
            kbody' = BodyDec GPU -> Stms GPU -> [KernelResult] -> KernelBody GPU
forall rep res. BodyDec rep -> Stms rep -> [res] -> GBody rep res
Body () 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 :: Name -> SubExp -> SubExp -> Builder GPU SubExp
    limitTile :: Name -> SubExp -> SubExp -> BuilderT GPU (State VNameSource) SubExp
limitTile Name
t_str SubExp
t SubExp
d_K = Name
-> Exp (Rep (BuilderT GPU (State VNameSource)))
-> BuilderT GPU (State VNameSource) SubExp
forall (m :: * -> *).
MonadBuilder m =>
Name -> Exp (Rep m) -> m SubExp
letSubExp Name
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 (Stms GPU), M.Map VName (PrimType, Stms GPU)) ->
      (VName, Stms GPU) ->
      Builder GPU (M.Map VName (Stms GPU), M.Map VName (PrimType, Stms GPU))
    insertTranspose :: Map VName Names
-> (VName, SubExp)
-> (Map VName (Stms GPU), Map VName (PrimType, Stms GPU))
-> (VName, Stms GPU)
-> BuilderT
     GPU
     (State VNameSource)
     (Map VName (Stms GPU), Map VName (PrimType, Stms GPU))
insertTranspose Map VName Names
variance (VName
gidz, SubExp
_) (Map VName (Stms GPU)
tab_inn, Map VName (PrimType, Stms GPU)
tab_out) (VName
p_nm, Stms GPU
stms)
      | Just (Stms GPU
deps, Let Pat (LetDec GPU)
patt StmAux (ExpDec GPU)
yy (BasicOp (Index VName
arr_nm Slice SubExp
slc))) <- Stms GPU -> Maybe (Stms GPU, Stm GPU)
forall lore. Stms lore -> Maybe (Stms lore, Stm lore)
stmsLast Stms GPU
stms,
        [PatElem Type
p] <- Pat Type -> [PatElem Type]
forall dec. Pat dec -> [PatElem dec]
patElems Pat Type
Pat (LetDec GPU)
patt,
        PrimType
ptp <- Type -> PrimType
forall shape o. TypeBase shape o -> PrimType
elemType (Type -> PrimType) -> Type -> PrimType
forall a b. (a -> b) -> a -> b
$ PatElem Type -> Type
forall dec. Typed dec => PatElem dec -> Type
patElemType PatElem Type
p,
        VName
p_nm VName -> VName -> Bool
forall a. Eq a => a -> a -> Bool
== PatElem Type -> VName
forall dec. PatElem dec -> VName
patElemName PatElem Type
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 (Stms GPU), Map VName (PrimType, Stms GPU))
-> BuilderT
     GPU
     (State VNameSource)
     (Map VName (Stms GPU), Map VName (PrimType, Stms GPU))
forall a. a -> BuilderT GPU (State VNameSource) a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (VName -> Stms GPU -> Map VName (Stms GPU) -> Map VName (Stms GPU)
forall k a. Ord k => k -> a -> Map k a -> Map k a
M.insert VName
p_nm Stms GPU
stms Map VName (Stms GPU)
tab_inn, Map VName (PrimType, Stms GPU)
tab_out)
            Int
i : [Int]
_ -> do
              arr_tp <- VName -> BuilderT GPU (State VNameSource) Type
forall rep (m :: * -> *). HasScope rep m => VName -> m Type
lookupType VName
arr_nm
              let perm = [Int
i Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
1 .. Type -> Int
forall shape o. ArrayShape shape => TypeBase shape o -> Int
arrayRank Type
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 -> Name
baseName VName
arr_nm Name -> Name -> Name
forall a. Semigroup a => a -> a -> a
<> Name
"_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, deps Seq.|> stm') tab_out)
    insertTranspose Map VName Names
_ (VName, SubExp)
_ (Map VName (Stms GPU), Map VName (PrimType, Stms GPU))
_ (VName, Stms GPU)
_ = [Char]
-> BuilderT
     GPU
     (State VNameSource)
     (Map VName (Stms GPU), Map VName (PrimType, Stms 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