{-# LANGUAGE FlexibleContexts #-}
{-# LANGUAGE LambdaCase #-}
{-# LANGUAGE TypeFamilies #-}
module Futhark.Optimise.TileLoops (tileLoops) where
import Control.Monad.Reader
import Control.Monad.State
import qualified Data.Map.Strict as M
import Data.Maybe (mapMaybe)
import qualified Data.Sequence as Seq
import qualified Futhark.Analysis.Alias as Alias
import Futhark.IR.GPU
import Futhark.IR.Prop.Aliases (consumedInStm)
import Futhark.MonadFreshNames
import Futhark.Optimise.BlkRegTiling
import Futhark.Optimise.TileLoops.Shared
import Futhark.Pass
import Futhark.Tools
import Futhark.Transform.Rename
import Prelude hiding (quot)
tileLoops :: Pass GPU GPU
tileLoops :: Pass GPU GPU
tileLoops =
String -> String -> (Prog GPU -> PassM (Prog GPU)) -> Pass GPU GPU
forall fromrep torep.
String
-> String
-> (Prog fromrep -> PassM (Prog torep))
-> Pass fromrep torep
Pass String
"tile loops" String
"Tile stream loops inside kernels" ((Prog GPU -> PassM (Prog GPU)) -> Pass GPU GPU)
-> (Prog GPU -> PassM (Prog GPU)) -> Pass GPU GPU
forall a b. (a -> b) -> a -> b
$
(Scope GPU -> Stms GPU -> PassM (Stms GPU))
-> Prog GPU -> PassM (Prog GPU)
forall rep.
(Scope rep -> Stms rep -> PassM (Stms rep))
-> Prog rep -> PassM (Prog rep)
intraproceduralTransformation Scope GPU -> Stms GPU -> PassM (Stms GPU)
forall {m :: * -> *}.
MonadFreshNames m =>
Scope GPU -> Stms GPU -> m (Stms GPU)
onStms
where
onStms :: Scope GPU -> Stms GPU -> m (Stms GPU)
onStms Scope GPU
scope Stms GPU
stms =
(VNameSource -> (Stms GPU, VNameSource)) -> m (Stms GPU)
forall (m :: * -> *) a.
MonadFreshNames m =>
(VNameSource -> (a, VNameSource)) -> m a
modifyNameSource ((VNameSource -> (Stms GPU, VNameSource)) -> m (Stms GPU))
-> (VNameSource -> (Stms GPU, VNameSource)) -> m (Stms GPU)
forall a b. (a -> b) -> a -> b
$
State VNameSource (Stms GPU)
-> VNameSource -> (Stms GPU, VNameSource)
forall s a. State s a -> s -> (a, s)
runState (State VNameSource (Stms GPU)
-> VNameSource -> (Stms GPU, VNameSource))
-> State VNameSource (Stms GPU)
-> VNameSource
-> (Stms GPU, VNameSource)
forall a b. (a -> b) -> a -> b
$
ReaderT (Scope GPU) (State VNameSource) (Stms GPU)
-> Scope GPU -> State VNameSource (Stms GPU)
forall r (m :: * -> *) a. ReaderT r m a -> r -> m a
runReaderT (Stms GPU -> ReaderT (Scope GPU) (State VNameSource) (Stms GPU)
optimiseStms Stms GPU
stms) Scope GPU
scope
optimiseBody :: Body GPU -> TileM (Body GPU)
optimiseBody :: Body GPU -> TileM (Body GPU)
optimiseBody (Body () Stms GPU
stms Result
res) =
BodyDec GPU -> Stms GPU -> Result -> Body GPU
forall rep. BodyDec rep -> Stms rep -> Result -> BodyT rep
Body () (Stms GPU -> Result -> Body GPU)
-> ReaderT (Scope GPU) (State VNameSource) (Stms GPU)
-> ReaderT (Scope GPU) (State VNameSource) (Result -> Body GPU)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Stms GPU -> ReaderT (Scope GPU) (State VNameSource) (Stms GPU)
optimiseStms Stms GPU
stms ReaderT (Scope GPU) (State VNameSource) (Result -> Body GPU)
-> ReaderT (Scope GPU) (State VNameSource) Result
-> TileM (Body GPU)
forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> Result -> ReaderT (Scope GPU) (State VNameSource) Result
forall (f :: * -> *) a. Applicative f => a -> f a
pure Result
res
optimiseStms :: Stms GPU -> TileM (Stms GPU)
optimiseStms :: Stms GPU -> ReaderT (Scope GPU) (State VNameSource) (Stms GPU)
optimiseStms Stms GPU
stms =
Scope GPU
-> ReaderT (Scope GPU) (State VNameSource) (Stms GPU)
-> ReaderT (Scope GPU) (State VNameSource) (Stms GPU)
forall rep (m :: * -> *) a.
LocalScope rep m =>
Scope rep -> m a -> m a
localScope (Stms GPU -> Scope GPU
forall rep a. Scoped rep a => a -> Scope rep
scopeOf Stms GPU
stms) (ReaderT (Scope GPU) (State VNameSource) (Stms GPU)
-> ReaderT (Scope GPU) (State VNameSource) (Stms GPU))
-> ReaderT (Scope GPU) (State VNameSource) (Stms GPU)
-> ReaderT (Scope GPU) (State VNameSource) (Stms GPU)
forall a b. (a -> b) -> a -> b
$
[Stms GPU] -> Stms GPU
forall a. Monoid a => [a] -> a
mconcat ([Stms GPU] -> Stms GPU)
-> ReaderT (Scope GPU) (State VNameSource) [Stms GPU]
-> ReaderT (Scope GPU) (State VNameSource) (Stms GPU)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> (Stm GPU -> ReaderT (Scope GPU) (State VNameSource) (Stms GPU))
-> [Stm GPU] -> ReaderT (Scope GPU) (State VNameSource) [Stms GPU]
forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
(a -> m b) -> t a -> m (t b)
mapM Stm GPU -> ReaderT (Scope GPU) (State VNameSource) (Stms GPU)
optimiseStm (Stms GPU -> [Stm GPU]
forall rep. Stms rep -> [Stm rep]
stmsToList Stms GPU
stms)
optimiseStm :: Stm GPU -> TileM (Stms GPU)
optimiseStm :: Stm GPU -> ReaderT (Scope GPU) (State VNameSource) (Stms GPU)
optimiseStm stm :: Stm GPU
stm@(Let Pattern GPU
pat StmAux (ExpDec GPU)
aux (Op (SegOp (SegMap lvl :: SegLevel
lvl@SegThread {} SegSpace
space [Type]
ts KernelBody GPU
kbody)))) = do
Maybe (Stms GPU, Stm GPU)
res3dtiling <- Stm GPU -> TileM (Maybe (Stms GPU, Stm GPU))
doRegTiling3D Stm GPU
stm
case Maybe (Stms GPU, Stm GPU)
res3dtiling of
Just (Stms GPU
extra_bnds, Stm GPU
stmt') -> Stms GPU -> ReaderT (Scope GPU) (State VNameSource) (Stms GPU)
forall (m :: * -> *) a. Monad m => a -> m a
return (Stms GPU
extra_bnds Stms GPU -> Stms GPU -> Stms GPU
forall a. Semigroup a => a -> a -> a
<> Stm GPU -> Stms GPU
forall rep. Stm rep -> Stms rep
oneStm Stm GPU
stmt')
Maybe (Stms GPU, Stm GPU)
Nothing -> do
Maybe (Stms GPU, Stm GPU)
blkRegTiling_res <- Stm GPU -> TileM (Maybe (Stms GPU, Stm GPU))
mmBlkRegTiling Stm GPU
stm
case Maybe (Stms GPU, Stm GPU)
blkRegTiling_res of
Just (Stms GPU
extra_bnds, Stm GPU
stmt') -> Stms GPU -> ReaderT (Scope GPU) (State VNameSource) (Stms GPU)
forall (m :: * -> *) a. Monad m => a -> m a
return (Stms GPU
extra_bnds Stms GPU -> Stms GPU -> Stms GPU
forall a. Semigroup a => a -> a -> a
<> Stm GPU -> Stms GPU
forall rep. Stm rep -> Stms rep
oneStm Stm GPU
stmt')
Maybe (Stms GPU, Stm GPU)
Nothing -> Scope GPU
-> ReaderT (Scope GPU) (State VNameSource) (Stms GPU)
-> ReaderT (Scope GPU) (State VNameSource) (Stms GPU)
forall rep (m :: * -> *) a.
LocalScope rep m =>
Scope rep -> m a -> m a
localScope (SegSpace -> Scope GPU
forall rep. SegSpace -> Scope rep
scopeOfSegSpace SegSpace
space) (ReaderT (Scope GPU) (State VNameSource) (Stms GPU)
-> ReaderT (Scope GPU) (State VNameSource) (Stms GPU))
-> ReaderT (Scope GPU) (State VNameSource) (Stms GPU)
-> ReaderT (Scope GPU) (State VNameSource) (Stms GPU)
forall a b. (a -> b) -> a -> b
$ do
(Stms GPU
host_stms, (SegLevel
lvl', SegSpace
space', KernelBody GPU
kbody')) <- Names
-> AliasTable
-> SegLevel
-> SegSpace
-> [Type]
-> KernelBody GPU
-> TileM (Stms GPU, (SegLevel, SegSpace, KernelBody GPU))
tileInKernelBody Names
forall a. Monoid a => a
mempty AliasTable
initial_variance SegLevel
lvl SegSpace
space [Type]
ts KernelBody GPU
kbody
Stms GPU -> ReaderT (Scope GPU) (State VNameSource) (Stms GPU)
forall (m :: * -> *) a. Monad m => a -> m a
return (Stms GPU -> ReaderT (Scope GPU) (State VNameSource) (Stms GPU))
-> Stms GPU -> ReaderT (Scope GPU) (State VNameSource) (Stms GPU)
forall a b. (a -> b) -> a -> b
$ Stms GPU
host_stms Stms GPU -> Stms GPU -> Stms GPU
forall a. Semigroup a => a -> a -> a
<> Stm GPU -> Stms GPU
forall rep. Stm rep -> Stms rep
oneStm (Pattern GPU -> StmAux (ExpDec GPU) -> ExpT GPU -> Stm GPU
forall rep.
Pattern rep -> StmAux (ExpDec rep) -> Exp rep -> Stm rep
Let Pattern GPU
pat StmAux (ExpDec GPU)
aux (ExpT GPU -> Stm GPU) -> ExpT GPU -> Stm GPU
forall a b. (a -> b) -> a -> b
$ Op GPU -> ExpT GPU
forall rep. Op rep -> ExpT rep
Op (Op GPU -> ExpT GPU) -> Op GPU -> ExpT GPU
forall a b. (a -> b) -> a -> b
$ SegOp SegLevel GPU -> HostOp GPU (SOAC GPU)
forall rep op. SegOp SegLevel rep -> HostOp rep op
SegOp (SegOp SegLevel GPU -> HostOp GPU (SOAC GPU))
-> SegOp SegLevel GPU -> HostOp GPU (SOAC GPU)
forall a b. (a -> b) -> a -> b
$ SegLevel
-> SegSpace -> [Type] -> KernelBody GPU -> SegOp SegLevel GPU
forall lvl rep.
lvl -> SegSpace -> [Type] -> KernelBody rep -> SegOp lvl rep
SegMap SegLevel
lvl' SegSpace
space' [Type]
ts KernelBody GPU
kbody')
where
initial_variance :: AliasTable
initial_variance = (NameInfo Any -> Names) -> Map VName (NameInfo Any) -> AliasTable
forall a b k. (a -> b) -> Map k a -> Map k b
M.map NameInfo Any -> Names
forall a. Monoid a => a
mempty (Map VName (NameInfo Any) -> AliasTable)
-> Map VName (NameInfo Any) -> AliasTable
forall a b. (a -> b) -> a -> b
$ SegSpace -> Map VName (NameInfo Any)
forall rep. SegSpace -> Scope rep
scopeOfSegSpace SegSpace
space
optimiseStm (Let Pattern GPU
pat StmAux (ExpDec GPU)
aux ExpT GPU
e) =
Stm GPU -> Stms GPU
forall (f :: * -> *) a. Applicative f => a -> f a
pure (Stm GPU -> Stms GPU)
-> ReaderT (Scope GPU) (State VNameSource) (Stm GPU)
-> ReaderT (Scope GPU) (State VNameSource) (Stms GPU)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> (Pattern GPU -> StmAux (ExpDec GPU) -> ExpT GPU -> Stm GPU
forall rep.
Pattern rep -> StmAux (ExpDec rep) -> Exp rep -> Stm rep
Let Pattern GPU
pat StmAux (ExpDec GPU)
aux (ExpT GPU -> Stm GPU)
-> ReaderT (Scope GPU) (State VNameSource) (ExpT GPU)
-> ReaderT (Scope GPU) (State VNameSource) (Stm GPU)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Mapper GPU GPU (ReaderT (Scope GPU) (State VNameSource))
-> ExpT GPU -> ReaderT (Scope GPU) (State VNameSource) (ExpT GPU)
forall (m :: * -> *) frep trep.
(Applicative m, Monad m) =>
Mapper frep trep m -> Exp frep -> m (Exp trep)
mapExpM Mapper GPU GPU (ReaderT (Scope GPU) (State VNameSource))
optimise ExpT GPU
e)
where
optimise :: Mapper GPU GPU (ReaderT (Scope GPU) (State VNameSource))
optimise = Mapper GPU GPU (ReaderT (Scope GPU) (State VNameSource))
forall (m :: * -> *) rep. Monad m => Mapper rep rep m
identityMapper {mapOnBody :: Scope GPU -> Body GPU -> TileM (Body GPU)
mapOnBody = \Scope GPU
scope -> Scope GPU -> TileM (Body GPU) -> TileM (Body GPU)
forall rep (m :: * -> *) a.
LocalScope rep m =>
Scope rep -> m a -> m a
localScope Scope GPU
scope (TileM (Body GPU) -> TileM (Body GPU))
-> (Body GPU -> TileM (Body GPU)) -> Body GPU -> TileM (Body GPU)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Body GPU -> TileM (Body GPU)
optimiseBody}
tileInKernelBody ::
Names ->
VarianceTable ->
SegLevel ->
SegSpace ->
[Type] ->
KernelBody GPU ->
TileM (Stms GPU, (SegLevel, SegSpace, KernelBody GPU))
tileInKernelBody :: Names
-> AliasTable
-> SegLevel
-> SegSpace
-> [Type]
-> KernelBody GPU
-> TileM (Stms GPU, (SegLevel, SegSpace, KernelBody GPU))
tileInKernelBody Names
branch_variant AliasTable
initial_variance SegLevel
lvl SegSpace
initial_kspace [Type]
ts KernelBody GPU
kbody
| Just Result
kbody_res <- (KernelResult -> Maybe SubExp) -> [KernelResult] -> Maybe Result
forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
(a -> m b) -> t a -> m (t b)
mapM KernelResult -> Maybe SubExp
isSimpleResult ([KernelResult] -> Maybe Result) -> [KernelResult] -> Maybe Result
forall a b. (a -> b) -> a -> b
$ KernelBody GPU -> [KernelResult]
forall rep. KernelBody rep -> [KernelResult]
kernelBodyResult KernelBody GPU
kbody = do
Maybe (Stms GPU, Tiling, TiledBody)
maybe_tiled <-
Names
-> AliasTable
-> SegLevel
-> SegSpace
-> [Type]
-> Body GPU
-> TileM (Maybe (Stms GPU, Tiling, TiledBody))
tileInBody Names
branch_variant AliasTable
initial_variance SegLevel
lvl SegSpace
initial_kspace [Type]
ts (Body GPU -> TileM (Maybe (Stms GPU, Tiling, TiledBody)))
-> Body GPU -> TileM (Maybe (Stms GPU, Tiling, TiledBody))
forall a b. (a -> b) -> a -> b
$
BodyDec GPU -> Stms GPU -> Result -> Body GPU
forall rep. BodyDec rep -> Stms rep -> Result -> BodyT rep
Body () (KernelBody GPU -> Stms GPU
forall rep. KernelBody rep -> Stms rep
kernelBodyStms KernelBody GPU
kbody) Result
kbody_res
case Maybe (Stms GPU, Tiling, TiledBody)
maybe_tiled of
Just (Stms GPU
host_stms, Tiling
tiling, TiledBody
tiledBody) -> do
([KernelResult]
res', Stms GPU
stms') <-
Binder GPU [KernelResult]
-> ReaderT
(Scope GPU) (State VNameSource) ([KernelResult], Stms GPU)
forall (m :: * -> *) somerep rep a.
(MonadFreshNames m, HasScope somerep m, SameScope somerep rep) =>
Binder rep a -> m (a, Stms rep)
runBinder (Binder GPU [KernelResult]
-> ReaderT
(Scope GPU) (State VNameSource) ([KernelResult], Stms GPU))
-> Binder GPU [KernelResult]
-> ReaderT
(Scope GPU) (State VNameSource) ([KernelResult], Stms GPU)
forall a b. (a -> b) -> a -> b
$ (VName -> BinderT GPU (State VNameSource) KernelResult)
-> [VName] -> Binder GPU [KernelResult]
forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
(a -> m b) -> t a -> m (t b)
mapM (Tiling -> VName -> BinderT GPU (State VNameSource) KernelResult
tilingTileReturns Tiling
tiling) ([VName] -> Binder GPU [KernelResult])
-> BinderT GPU (State VNameSource) [VName]
-> Binder GPU [KernelResult]
forall (m :: * -> *) a b. Monad m => (a -> m b) -> m a -> m b
=<< TiledBody
tiledBody Names
forall a. Monoid a => a
mempty PrivStms
forall a. Monoid a => a
mempty
(Stms GPU, (SegLevel, SegSpace, KernelBody GPU))
-> TileM (Stms GPU, (SegLevel, SegSpace, KernelBody GPU))
forall (m :: * -> *) a. Monad m => a -> m a
return
( Stms GPU
host_stms,
( Tiling -> SegLevel
tilingLevel Tiling
tiling,
Tiling -> SegSpace
tilingSpace Tiling
tiling,
BodyDec GPU -> Stms GPU -> [KernelResult] -> KernelBody GPU
forall rep.
BodyDec rep -> Stms rep -> [KernelResult] -> KernelBody rep
KernelBody () Stms GPU
stms' [KernelResult]
res'
)
)
Maybe (Stms GPU, Tiling, TiledBody)
Nothing ->
(Stms GPU, (SegLevel, SegSpace, KernelBody GPU))
-> TileM (Stms GPU, (SegLevel, SegSpace, KernelBody GPU))
forall (m :: * -> *) a. Monad m => a -> m a
return (Stms GPU
forall a. Monoid a => a
mempty, (SegLevel
lvl, SegSpace
initial_kspace, KernelBody GPU
kbody))
| Bool
otherwise =
(Stms GPU, (SegLevel, SegSpace, KernelBody GPU))
-> TileM (Stms GPU, (SegLevel, SegSpace, KernelBody GPU))
forall (m :: * -> *) a. Monad m => a -> m a
return (Stms GPU
forall a. Monoid a => a
mempty, (SegLevel
lvl, SegSpace
initial_kspace, KernelBody GPU
kbody))
where
isSimpleResult :: KernelResult -> Maybe SubExp
isSimpleResult (Returns ResultManifest
_ SubExp
se) = SubExp -> Maybe SubExp
forall a. a -> Maybe a
Just SubExp
se
isSimpleResult KernelResult
_ = Maybe SubExp
forall a. Maybe a
Nothing
tileInBody ::
Names ->
VarianceTable ->
SegLevel ->
SegSpace ->
[Type] ->
Body GPU ->
TileM (Maybe (Stms GPU, Tiling, TiledBody))
tileInBody :: Names
-> AliasTable
-> SegLevel
-> SegSpace
-> [Type]
-> Body GPU
-> TileM (Maybe (Stms GPU, Tiling, TiledBody))
tileInBody Names
branch_variant AliasTable
initial_variance SegLevel
initial_lvl SegSpace
initial_space [Type]
res_ts (Body () Stms GPU
initial_kstms Result
stms_res) =
Stms GPU
-> [Stm GPU] -> TileM (Maybe (Stms GPU, Tiling, TiledBody))
descend Stms GPU
forall a. Monoid a => a
mempty ([Stm GPU] -> TileM (Maybe (Stms GPU, Tiling, TiledBody)))
-> [Stm GPU] -> TileM (Maybe (Stms GPU, Tiling, TiledBody))
forall a b. (a -> b) -> a -> b
$ Stms GPU -> [Stm GPU]
forall rep. Stms rep -> [Stm rep]
stmsToList Stms GPU
initial_kstms
where
variance :: AliasTable
variance = AliasTable -> Stms GPU -> AliasTable
varianceInStms AliasTable
initial_variance Stms GPU
initial_kstms
descend :: Stms GPU
-> [Stm GPU] -> TileM (Maybe (Stms GPU, Tiling, TiledBody))
descend Stms GPU
_ [] =
Maybe (Stms GPU, Tiling, TiledBody)
-> TileM (Maybe (Stms GPU, Tiling, TiledBody))
forall (m :: * -> *) a. Monad m => a -> m a
return Maybe (Stms GPU, Tiling, TiledBody)
forall a. Maybe a
Nothing
descend Stms GPU
prestms (Stm GPU
stm_to_tile : [Stm GPU]
poststms)
| ([VName]
gtids, Result
kdims) <- [(VName, SubExp)] -> ([VName], Result)
forall a b. [(a, b)] -> ([a], [b])
unzip ([(VName, SubExp)] -> ([VName], Result))
-> [(VName, SubExp)] -> ([VName], Result)
forall a b. (a -> b) -> a -> b
$ SegSpace -> [(VName, SubExp)]
unSegSpace SegSpace
initial_space,
Just (SubExp
w, [VName]
arrs, (Commutativity, Lambda GPU, Result, Lambda GPU)
form) <- Stm GPU
-> Maybe
(SubExp, [VName], (Commutativity, Lambda GPU, Result, Lambda GPU))
tileable Stm GPU
stm_to_tile,
Just [InputArray]
inputs <-
(VName -> Maybe InputArray) -> [VName] -> Maybe [InputArray]
forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
(a -> m b) -> t a -> m (t b)
mapM (Names -> AliasTable -> [VName] -> VName -> Maybe InputArray
invariantToOneOfTwoInnerDims Names
branch_variant AliasTable
variance [VName]
gtids) [VName]
arrs,
Bool -> Bool
not (Bool -> Bool) -> Bool -> Bool
forall a b. (a -> b) -> a -> b
$ [(VName, [Int])] -> Bool
forall (t :: * -> *) a. Foldable t => t a -> Bool
null ([(VName, [Int])] -> Bool) -> [(VName, [Int])] -> Bool
forall a b. (a -> b) -> a -> b
$ [InputArray] -> [(VName, [Int])]
tiledInputs [InputArray]
inputs,
VName
gtid_y : VName
gtid_x : [VName]
top_gtids_rev <- [VName] -> [VName]
forall a. [a] -> [a]
reverse [VName]
gtids,
SubExp
kdim_y : SubExp
kdim_x : Result
top_kdims_rev <- Result -> Result
forall a. [a] -> [a]
reverse Result
kdims,
(Stms GPU
prestms', Stms GPU
poststms') <-
AliasTable
-> Stms GPU -> Stm GPU -> Stms GPU -> (Stms GPU, Stms GPU)
preludeToPostlude AliasTable
variance Stms GPU
prestms Stm GPU
stm_to_tile ([Stm GPU] -> Stms GPU
forall rep. [Stm rep] -> Stms rep
stmsFromList [Stm GPU]
poststms),
Names
used <- Stm GPU -> Names
forall a. FreeIn a => a -> Names
freeIn Stm GPU
stm_to_tile Names -> Names -> Names
forall a. Semigroup a => a -> a -> a
<> Stms GPU -> Names
forall a. FreeIn a => a -> Names
freeIn Stms GPU
poststms' Names -> Names -> Names
forall a. Semigroup a => a -> a -> a
<> Result -> Names
forall a. FreeIn a => a -> Names
freeIn Result
stms_res =
(Stms GPU, Tiling, TiledBody)
-> Maybe (Stms GPU, Tiling, TiledBody)
forall a. a -> Maybe a
Just ((Stms GPU, Tiling, TiledBody)
-> Maybe (Stms GPU, Tiling, TiledBody))
-> ((Stms GPU, Tiling, TiledBody) -> (Stms GPU, Tiling, TiledBody))
-> (Stms GPU, Tiling, TiledBody)
-> Maybe (Stms GPU, Tiling, TiledBody)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. SegSpace
-> AliasTable
-> Stms GPU
-> Names
-> (Stms GPU, Tiling, TiledBody)
-> (Stms GPU, Tiling, TiledBody)
injectPrelude SegSpace
initial_space AliasTable
variance Stms GPU
prestms' Names
used
((Stms GPU, Tiling, TiledBody)
-> Maybe (Stms GPU, Tiling, TiledBody))
-> ReaderT
(Scope GPU) (State VNameSource) (Stms GPU, Tiling, TiledBody)
-> TileM (Maybe (Stms GPU, Tiling, TiledBody))
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> DoTiling (VName, VName) (SubExp, SubExp)
-> SegLevel
-> [Type]
-> Pattern GPU
-> (VName, VName)
-> (SubExp, SubExp)
-> SubExp
-> (Commutativity, Lambda GPU, Result, Lambda GPU)
-> [InputArray]
-> Stms GPU
-> Result
-> ReaderT
(Scope GPU) (State VNameSource) (Stms GPU, Tiling, TiledBody)
forall gtids kdims.
DoTiling gtids kdims
-> SegLevel
-> [Type]
-> Pattern GPU
-> gtids
-> kdims
-> SubExp
-> (Commutativity, Lambda GPU, Result, Lambda GPU)
-> [InputArray]
-> Stms GPU
-> Result
-> ReaderT
(Scope GPU) (State VNameSource) (Stms GPU, Tiling, TiledBody)
tileGeneric
([(VName, SubExp)] -> DoTiling (VName, VName) (SubExp, SubExp)
tiling2d ([(VName, SubExp)] -> DoTiling (VName, VName) (SubExp, SubExp))
-> [(VName, SubExp)] -> DoTiling (VName, VName) (SubExp, SubExp)
forall a b. (a -> b) -> a -> b
$ [(VName, SubExp)] -> [(VName, SubExp)]
forall a. [a] -> [a]
reverse ([(VName, SubExp)] -> [(VName, SubExp)])
-> [(VName, SubExp)] -> [(VName, SubExp)]
forall a b. (a -> b) -> a -> b
$ [VName] -> Result -> [(VName, SubExp)]
forall a b. [a] -> [b] -> [(a, b)]
zip [VName]
top_gtids_rev Result
top_kdims_rev)
SegLevel
initial_lvl
[Type]
res_ts
(Stm GPU -> Pattern GPU
forall rep. Stm rep -> Pattern rep
stmPattern Stm GPU
stm_to_tile)
(VName
gtid_x, VName
gtid_y)
(SubExp
kdim_x, SubExp
kdim_y)
SubExp
w
(Commutativity, Lambda GPU, Result, Lambda GPU)
form
[InputArray]
inputs
Stms GPU
poststms'
Result
stms_res
| (VName
gtid, SubExp
kdim) : [(VName, SubExp)]
top_space_rev <- [(VName, SubExp)] -> [(VName, SubExp)]
forall a. [a] -> [a]
reverse ([(VName, SubExp)] -> [(VName, SubExp)])
-> [(VName, SubExp)] -> [(VName, SubExp)]
forall a b. (a -> b) -> a -> b
$ SegSpace -> [(VName, SubExp)]
unSegSpace SegSpace
initial_space,
Just (SubExp
w, [VName]
arrs, (Commutativity, Lambda GPU, Result, Lambda GPU)
form) <- Stm GPU
-> Maybe
(SubExp, [VName], (Commutativity, Lambda GPU, Result, Lambda GPU))
tileable Stm GPU
stm_to_tile,
[InputArray]
inputs <- (VName -> InputArray) -> [VName] -> [InputArray]
forall a b. (a -> b) -> [a] -> [b]
map (VName -> AliasTable -> VName -> InputArray
is1DTileable VName
gtid AliasTable
variance) [VName]
arrs,
Bool -> Bool
not (Bool -> Bool) -> Bool -> Bool
forall a b. (a -> b) -> a -> b
$ [(VName, [Int])] -> Bool
forall (t :: * -> *) a. Foldable t => t a -> Bool
null ([(VName, [Int])] -> Bool) -> [(VName, [Int])] -> Bool
forall a b. (a -> b) -> a -> b
$ [InputArray] -> [(VName, [Int])]
tiledInputs [InputArray]
inputs,
Bool -> Bool
not (Bool -> Bool) -> Bool -> Bool
forall a b. (a -> b) -> a -> b
$ VName
gtid VName -> Names -> Bool
`nameIn` Names
branch_variant,
(Stms GPU
prestms', Stms GPU
poststms') <-
AliasTable
-> Stms GPU -> Stm GPU -> Stms GPU -> (Stms GPU, Stms GPU)
preludeToPostlude AliasTable
variance Stms GPU
prestms Stm GPU
stm_to_tile ([Stm GPU] -> Stms GPU
forall rep. [Stm rep] -> Stms rep
stmsFromList [Stm GPU]
poststms),
Names
used <- Stm GPU -> Names
forall a. FreeIn a => a -> Names
freeIn Stm GPU
stm_to_tile Names -> Names -> Names
forall a. Semigroup a => a -> a -> a
<> Stms GPU -> Names
forall a. FreeIn a => a -> Names
freeIn Stms GPU
poststms' Names -> Names -> Names
forall a. Semigroup a => a -> a -> a
<> Result -> Names
forall a. FreeIn a => a -> Names
freeIn Result
stms_res =
(Stms GPU, Tiling, TiledBody)
-> Maybe (Stms GPU, Tiling, TiledBody)
forall a. a -> Maybe a
Just ((Stms GPU, Tiling, TiledBody)
-> Maybe (Stms GPU, Tiling, TiledBody))
-> ((Stms GPU, Tiling, TiledBody) -> (Stms GPU, Tiling, TiledBody))
-> (Stms GPU, Tiling, TiledBody)
-> Maybe (Stms GPU, Tiling, TiledBody)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. SegSpace
-> AliasTable
-> Stms GPU
-> Names
-> (Stms GPU, Tiling, TiledBody)
-> (Stms GPU, Tiling, TiledBody)
injectPrelude SegSpace
initial_space AliasTable
variance Stms GPU
prestms' Names
used
((Stms GPU, Tiling, TiledBody)
-> Maybe (Stms GPU, Tiling, TiledBody))
-> ReaderT
(Scope GPU) (State VNameSource) (Stms GPU, Tiling, TiledBody)
-> TileM (Maybe (Stms GPU, Tiling, TiledBody))
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> DoTiling VName SubExp
-> SegLevel
-> [Type]
-> Pattern GPU
-> VName
-> SubExp
-> SubExp
-> (Commutativity, Lambda GPU, Result, Lambda GPU)
-> [InputArray]
-> Stms GPU
-> Result
-> ReaderT
(Scope GPU) (State VNameSource) (Stms GPU, Tiling, TiledBody)
forall gtids kdims.
DoTiling gtids kdims
-> SegLevel
-> [Type]
-> Pattern GPU
-> gtids
-> kdims
-> SubExp
-> (Commutativity, Lambda GPU, Result, Lambda GPU)
-> [InputArray]
-> Stms GPU
-> Result
-> ReaderT
(Scope GPU) (State VNameSource) (Stms GPU, Tiling, TiledBody)
tileGeneric
([(VName, SubExp)] -> DoTiling VName SubExp
tiling1d ([(VName, SubExp)] -> DoTiling VName SubExp)
-> [(VName, SubExp)] -> DoTiling VName SubExp
forall a b. (a -> b) -> a -> b
$ [(VName, SubExp)] -> [(VName, SubExp)]
forall a. [a] -> [a]
reverse [(VName, SubExp)]
top_space_rev)
SegLevel
initial_lvl
[Type]
res_ts
(Stm GPU -> Pattern GPU
forall rep. Stm rep -> Pattern rep
stmPattern Stm GPU
stm_to_tile)
VName
gtid
SubExp
kdim
SubExp
w
(Commutativity, Lambda GPU, Result, Lambda GPU)
form
[InputArray]
inputs
Stms GPU
poststms'
Result
stms_res
| DoLoop [] [(FParam GPU, SubExp)]
merge (ForLoop VName
i IntType
it SubExp
bound []) Body GPU
loopbody <- Stm GPU -> ExpT GPU
forall rep. Stm rep -> Exp rep
stmExp Stm GPU
stm_to_tile,
(Stms GPU
prestms', Stms GPU
poststms') <-
AliasTable
-> Stms GPU -> Stm GPU -> Stms GPU -> (Stms GPU, Stms GPU)
preludeToPostlude AliasTable
variance Stms GPU
prestms Stm GPU
stm_to_tile ([Stm GPU] -> Stms GPU
forall rep. [Stm rep] -> Stms rep
stmsFromList [Stm GPU]
poststms) = do
let branch_variant' :: Names
branch_variant' =
Names
branch_variant
Names -> Names -> Names
forall a. Semigroup a => a -> a -> a
<> [Names] -> Names
forall a. Monoid a => [a] -> a
mconcat
( (VName -> Names) -> [VName] -> [Names]
forall a b. (a -> b) -> [a] -> [b]
map
((VName -> AliasTable -> Names) -> AliasTable -> VName -> Names
forall a b c. (a -> b -> c) -> b -> a -> c
flip (Names -> VName -> AliasTable -> Names
forall k a. Ord k => a -> k -> Map k a -> a
M.findWithDefault Names
forall a. Monoid a => a
mempty) AliasTable
variance)
(Names -> [VName]
namesToList (SubExp -> Names
forall a. FreeIn a => a -> Names
freeIn SubExp
bound))
)
merge_params :: [Param (TypeBase Shape Uniqueness)]
merge_params = ((Param (TypeBase Shape Uniqueness), SubExp)
-> Param (TypeBase Shape Uniqueness))
-> [(Param (TypeBase Shape Uniqueness), SubExp)]
-> [Param (TypeBase Shape Uniqueness)]
forall a b. (a -> b) -> [a] -> [b]
map (Param (TypeBase Shape Uniqueness), SubExp)
-> Param (TypeBase Shape Uniqueness)
forall a b. (a, b) -> a
fst [(Param (TypeBase Shape Uniqueness), SubExp)]
[(FParam GPU, SubExp)]
merge
Maybe (Stms GPU, Tiling, TiledBody)
maybe_tiled <-
Scope GPU
-> TileM (Maybe (Stms GPU, Tiling, TiledBody))
-> TileM (Maybe (Stms GPU, Tiling, TiledBody))
forall rep (m :: * -> *) a.
LocalScope rep m =>
Scope rep -> m a -> m a
localScope (VName -> NameInfo GPU -> Scope GPU -> Scope GPU
forall k a. Ord k => k -> a -> Map k a -> Map k a
M.insert VName
i (IntType -> NameInfo GPU
forall rep. IntType -> NameInfo rep
IndexName IntType
it) (Scope GPU -> Scope GPU) -> Scope GPU -> Scope GPU
forall a b. (a -> b) -> a -> b
$ [Param (TypeBase Shape Uniqueness)] -> Scope GPU
forall rep dec. (FParamInfo rep ~ dec) => [Param dec] -> Scope rep
scopeOfFParams [Param (TypeBase Shape Uniqueness)]
merge_params) (TileM (Maybe (Stms GPU, Tiling, TiledBody))
-> TileM (Maybe (Stms GPU, Tiling, TiledBody)))
-> TileM (Maybe (Stms GPU, Tiling, TiledBody))
-> TileM (Maybe (Stms GPU, Tiling, TiledBody))
forall a b. (a -> b) -> a -> b
$
Names
-> AliasTable
-> SegLevel
-> SegSpace
-> [Type]
-> Body GPU
-> TileM (Maybe (Stms GPU, Tiling, TiledBody))
tileInBody
Names
branch_variant'
AliasTable
variance
SegLevel
initial_lvl
SegSpace
initial_space
((Param (TypeBase Shape Uniqueness) -> Type)
-> [Param (TypeBase Shape Uniqueness)] -> [Type]
forall a b. (a -> b) -> [a] -> [b]
map Param (TypeBase Shape Uniqueness) -> Type
forall dec. Typed dec => Param dec -> Type
paramType [Param (TypeBase Shape Uniqueness)]
merge_params)
(Body GPU -> TileM (Maybe (Stms GPU, Tiling, TiledBody)))
-> Body GPU -> TileM (Maybe (Stms GPU, Tiling, TiledBody))
forall a b. (a -> b) -> a -> b
$ Stms GPU -> Result -> Body GPU
forall rep. Bindable rep => Stms rep -> Result -> Body rep
mkBody (Body GPU -> Stms GPU
forall rep. BodyT rep -> Stms rep
bodyStms Body GPU
loopbody) (Body GPU -> Result
forall rep. BodyT rep -> Result
bodyResult Body GPU
loopbody)
case Maybe (Stms GPU, Tiling, TiledBody)
maybe_tiled of
Maybe (Stms GPU, Tiling, TiledBody)
Nothing -> TileM (Maybe (Stms GPU, Tiling, TiledBody))
next
Just (Stms GPU, Tiling, TiledBody)
tiled ->
(Stms GPU, Tiling, TiledBody)
-> Maybe (Stms GPU, Tiling, TiledBody)
forall a. a -> Maybe a
Just
((Stms GPU, Tiling, TiledBody)
-> Maybe (Stms GPU, Tiling, TiledBody))
-> ReaderT
(Scope GPU) (State VNameSource) (Stms GPU, Tiling, TiledBody)
-> TileM (Maybe (Stms GPU, Tiling, TiledBody))
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> SegSpace
-> AliasTable
-> Stms GPU
-> Names
-> (Stms GPU, Tiling, TiledBody)
-> [Type]
-> Pattern GPU
-> StmAux (ExpDec GPU)
-> [(FParam GPU, SubExp)]
-> VName
-> IntType
-> SubExp
-> Stms GPU
-> Result
-> ReaderT
(Scope GPU) (State VNameSource) (Stms GPU, Tiling, TiledBody)
tileDoLoop
SegSpace
initial_space
AliasTable
variance
Stms GPU
prestms'
(Body GPU -> Names
forall a. FreeIn a => a -> Names
freeIn Body GPU
loopbody Names -> Names -> Names
forall a. Semigroup a => a -> a -> a
<> [(Param (TypeBase Shape Uniqueness), SubExp)] -> Names
forall a. FreeIn a => a -> Names
freeIn [(Param (TypeBase Shape Uniqueness), SubExp)]
[(FParam GPU, SubExp)]
merge)
(Stms GPU, Tiling, TiledBody)
tiled
[Type]
res_ts
(Stm GPU -> Pattern GPU
forall rep. Stm rep -> Pattern rep
stmPattern Stm GPU
stm_to_tile)
(Stm GPU -> StmAux (ExpDec GPU)
forall rep. Stm rep -> StmAux (ExpDec rep)
stmAux Stm GPU
stm_to_tile)
[(FParam GPU, SubExp)]
merge
VName
i
IntType
it
SubExp
bound
Stms GPU
poststms'
Result
stms_res
| Bool
otherwise = TileM (Maybe (Stms GPU, Tiling, TiledBody))
next
where
next :: TileM (Maybe (Stms GPU, Tiling, TiledBody))
next =
Scope GPU
-> TileM (Maybe (Stms GPU, Tiling, TiledBody))
-> TileM (Maybe (Stms GPU, Tiling, TiledBody))
forall rep (m :: * -> *) a.
LocalScope rep m =>
Scope rep -> m a -> m a
localScope (Stm GPU -> Scope GPU
forall rep a. Scoped rep a => a -> Scope rep
scopeOf Stm GPU
stm_to_tile) (TileM (Maybe (Stms GPU, Tiling, TiledBody))
-> TileM (Maybe (Stms GPU, Tiling, TiledBody)))
-> TileM (Maybe (Stms GPU, Tiling, TiledBody))
-> TileM (Maybe (Stms GPU, Tiling, TiledBody))
forall a b. (a -> b) -> a -> b
$
Stms GPU
-> [Stm GPU] -> TileM (Maybe (Stms GPU, Tiling, TiledBody))
descend (Stms GPU
prestms Stms GPU -> Stms GPU -> Stms GPU
forall a. Semigroup a => a -> a -> a
<> Stm GPU -> Stms GPU
forall rep. Stm rep -> Stms rep
oneStm Stm GPU
stm_to_tile) [Stm GPU]
poststms
preludeToPostlude ::
VarianceTable ->
Stms GPU ->
Stm GPU ->
Stms GPU ->
(Stms GPU, Stms GPU)
preludeToPostlude :: AliasTable
-> Stms GPU -> Stm GPU -> Stms GPU -> (Stms GPU, Stms GPU)
preludeToPostlude AliasTable
variance Stms GPU
prelude Stm GPU
stm_to_tile Stms GPU
postlude =
(Stms GPU
prelude_used, Stms GPU
prelude_not_used Stms GPU -> Stms GPU -> Stms GPU
forall a. Semigroup a => a -> a -> a
<> Stms GPU
postlude)
where
used_in_tiled :: Names
used_in_tiled = Stm GPU -> Names
forall a. FreeIn a => a -> Names
freeIn Stm GPU
stm_to_tile
used_in_stm_variant :: Names
used_in_stm_variant =
(Names
used_in_tiled Names -> Names -> Names
forall a. Semigroup a => a -> a -> a
<>) (Names -> Names) -> Names -> Names
forall a b. (a -> b) -> a -> b
$
[Names] -> Names
forall a. Monoid a => [a] -> a
mconcat ([Names] -> Names) -> [Names] -> Names
forall a b. (a -> b) -> a -> b
$
(VName -> Names) -> [VName] -> [Names]
forall a b. (a -> b) -> [a] -> [b]
map ((VName -> AliasTable -> Names) -> AliasTable -> VName -> Names
forall a b c. (a -> b -> c) -> b -> a -> c
flip (Names -> VName -> AliasTable -> Names
forall k a. Ord k => a -> k -> Map k a -> a
M.findWithDefault Names
forall a. Monoid a => a
mempty) AliasTable
variance) ([VName] -> [Names]) -> [VName] -> [Names]
forall a b. (a -> b) -> a -> b
$
Names -> [VName]
namesToList Names
used_in_tiled
used :: Stm GPU -> Bool
used Stm GPU
stm =
(VName -> Bool) -> [VName] -> Bool
forall (t :: * -> *) a. Foldable t => (a -> Bool) -> t a -> Bool
any (VName -> Names -> Bool
`nameIn` Names
used_in_stm_variant) ([VName] -> Bool) -> [VName] -> Bool
forall a b. (a -> b) -> a -> b
$
PatternT Type -> [VName]
forall dec. PatternT dec -> [VName]
patternNames (PatternT Type -> [VName]) -> PatternT Type -> [VName]
forall a b. (a -> b) -> a -> b
$ Stm GPU -> Pattern GPU
forall rep. Stm rep -> Pattern rep
stmPattern Stm GPU
stm
(Stms GPU
prelude_used, Stms GPU
prelude_not_used) =
(Stm GPU -> Bool) -> Stms GPU -> (Stms GPU, Stms GPU)
forall a. (a -> Bool) -> Seq a -> (Seq a, Seq a)
Seq.partition Stm GPU -> Bool
used Stms GPU
prelude
partitionPrelude ::
VarianceTable ->
Stms GPU ->
Names ->
Names ->
(Stms GPU, Stms GPU, Stms GPU)
partitionPrelude :: AliasTable
-> Stms GPU -> Names -> Names -> (Stms GPU, Stms GPU, Stms GPU)
partitionPrelude AliasTable
variance Stms GPU
prestms Names
private Names
used_after =
(Stms GPU
invariant_prestms, Stms GPU
precomputed_variant_prestms, Stms GPU
recomputed_variant_prestms)
where
invariantTo :: Names -> Stm GPU -> Bool
invariantTo Names
names Stm GPU
stm =
case PatternT Type -> [VName]
forall dec. PatternT dec -> [VName]
patternNames (Stm GPU -> Pattern GPU
forall rep. Stm rep -> Pattern rep
stmPattern Stm GPU
stm) of
[] -> Bool
True
VName
v : [VName]
_ -> Bool -> Bool
not (Bool -> Bool) -> Bool -> Bool
forall a b. (a -> b) -> a -> b
$ (VName -> Bool) -> [VName] -> Bool
forall (t :: * -> *) a. Foldable t => (a -> Bool) -> t a -> Bool
any (VName -> Names -> Bool
`nameIn` Names
names) ([VName] -> Bool) -> [VName] -> Bool
forall a b. (a -> b) -> a -> b
$ Names -> [VName]
namesToList (Names -> [VName]) -> Names -> [VName]
forall a b. (a -> b) -> a -> b
$ Names -> VName -> AliasTable -> Names
forall k a. Ord k => a -> k -> Map k a -> a
M.findWithDefault Names
forall a. Monoid a => a
mempty VName
v AliasTable
variance
consumed :: VName -> Bool
consumed VName
v = VName
v VName -> Names -> Bool
`nameIn` Names
consumed_in_prestms
consumedStm :: Stm GPU -> Bool
consumedStm Stm GPU
stm = (VName -> Bool) -> [VName] -> Bool
forall (t :: * -> *) a. Foldable t => (a -> Bool) -> t a -> Bool
any VName -> Bool
consumed (PatternT Type -> [VName]
forall dec. PatternT dec -> [VName]
patternNames (Stm GPU -> Pattern GPU
forall rep. Stm rep -> Pattern rep
stmPattern Stm GPU
stm))
later_consumed :: Names
later_consumed =
[VName] -> Names
namesFromList ([VName] -> Names) -> [VName] -> Names
forall a b. (a -> b) -> a -> b
$
(Stm GPU -> [VName]) -> [Stm GPU] -> [VName]
forall (t :: * -> *) a b. Foldable t => (a -> [b]) -> t a -> [b]
concatMap (PatternT Type -> [VName]
forall dec. PatternT dec -> [VName]
patternNames (PatternT Type -> [VName])
-> (Stm GPU -> PatternT Type) -> Stm GPU -> [VName]
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Stm GPU -> PatternT Type
forall rep. Stm rep -> Pattern rep
stmPattern) ([Stm GPU] -> [VName]) -> [Stm GPU] -> [VName]
forall a b. (a -> b) -> a -> b
$
Stms GPU -> [Stm GPU]
forall rep. Stms rep -> [Stm rep]
stmsToList (Stms GPU -> [Stm GPU]) -> Stms GPU -> [Stm GPU]
forall a b. (a -> b) -> a -> b
$ (Stm GPU -> Bool) -> Stms GPU -> Stms GPU
forall a. (a -> Bool) -> Seq a -> Seq a
Seq.filter Stm GPU -> Bool
consumedStm Stms GPU
prestms
groupInvariant :: Stm GPU -> Bool
groupInvariant Stm GPU
stm =
Names -> Stm GPU -> Bool
invariantTo Names
private Stm GPU
stm
Bool -> Bool -> Bool
&& Bool -> Bool
not ((VName -> Bool) -> [VName] -> Bool
forall (t :: * -> *) a. Foldable t => (a -> Bool) -> t a -> Bool
any (VName -> Names -> Bool
`nameIn` Names
later_consumed) (PatternT Type -> [VName]
forall dec. PatternT dec -> [VName]
patternNames (Stm GPU -> Pattern GPU
forall rep. Stm rep -> Pattern rep
stmPattern Stm GPU
stm)))
Bool -> Bool -> Bool
&& Names -> Stm GPU -> Bool
invariantTo Names
later_consumed Stm GPU
stm
(Stms GPU
invariant_prestms, Stms GPU
variant_prestms) =
(Stm GPU -> Bool) -> Stms GPU -> (Stms GPU, Stms GPU)
forall a. (a -> Bool) -> Seq a -> (Seq a, Seq a)
Seq.partition Stm GPU -> Bool
groupInvariant Stms GPU
prestms
consumed_in_prestms :: Names
consumed_in_prestms =
(Stm (Aliases GPU) -> Names) -> Seq (Stm (Aliases GPU)) -> Names
forall (t :: * -> *) m a.
(Foldable t, Monoid m) =>
(a -> m) -> t a -> m
foldMap Stm (Aliases GPU) -> Names
forall rep. Aliased rep => Stm rep -> Names
consumedInStm (Seq (Stm (Aliases GPU)) -> Names)
-> Seq (Stm (Aliases GPU)) -> Names
forall a b. (a -> b) -> a -> b
$ (Seq (Stm (Aliases GPU)), AliasesAndConsumed)
-> Seq (Stm (Aliases GPU))
forall a b. (a, b) -> a
fst ((Seq (Stm (Aliases GPU)), AliasesAndConsumed)
-> Seq (Stm (Aliases GPU)))
-> (Seq (Stm (Aliases GPU)), AliasesAndConsumed)
-> Seq (Stm (Aliases GPU))
forall a b. (a -> b) -> a -> b
$ AliasTable
-> Stms GPU -> (Seq (Stm (Aliases GPU)), AliasesAndConsumed)
forall rep.
(ASTRep rep, CanBeAliased (Op rep)) =>
AliasTable -> Stms rep -> (Stms (Aliases rep), AliasesAndConsumed)
Alias.analyseStms AliasTable
forall a. Monoid a => a
mempty Stms GPU
prestms
mustBeInlinedExp :: ExpT rep -> Bool
mustBeInlinedExp (BasicOp (Index VName
_ Slice SubExp
slice)) = Bool -> Bool
not (Bool -> Bool) -> Bool -> Bool
forall a b. (a -> b) -> a -> b
$ Result -> Bool
forall (t :: * -> *) a. Foldable t => t a -> Bool
null (Result -> Bool) -> Result -> Bool
forall a b. (a -> b) -> a -> b
$ Slice SubExp -> Result
forall d. Slice d -> [d]
sliceDims Slice SubExp
slice
mustBeInlinedExp (BasicOp Rotate {}) = Bool
True
mustBeInlinedExp (BasicOp Rearrange {}) = Bool
True
mustBeInlinedExp (BasicOp Reshape {}) = Bool
True
mustBeInlinedExp ExpT rep
_ = Bool
False
mustBeInlined :: Stm GPU -> Bool
mustBeInlined Stm GPU
stm =
ExpT GPU -> Bool
forall {rep}. ExpT rep -> Bool
mustBeInlinedExp (Stm GPU -> ExpT GPU
forall rep. Stm rep -> Exp rep
stmExp Stm GPU
stm)
Bool -> Bool -> Bool
&& (VName -> Bool) -> [VName] -> Bool
forall (t :: * -> *) a. Foldable t => (a -> Bool) -> t a -> Bool
any (VName -> Names -> Bool
`nameIn` Names
used_after) (PatternT Type -> [VName]
forall dec. PatternT dec -> [VName]
patternNames (Stm GPU -> Pattern GPU
forall rep. Stm rep -> Pattern rep
stmPattern Stm GPU
stm))
must_be_inlined :: Names
must_be_inlined =
[VName] -> Names
namesFromList ([VName] -> Names) -> [VName] -> Names
forall a b. (a -> b) -> a -> b
$
(Stm GPU -> [VName]) -> [Stm GPU] -> [VName]
forall (t :: * -> *) a b. Foldable t => (a -> [b]) -> t a -> [b]
concatMap (PatternT Type -> [VName]
forall dec. PatternT dec -> [VName]
patternNames (PatternT Type -> [VName])
-> (Stm GPU -> PatternT Type) -> Stm GPU -> [VName]
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Stm GPU -> PatternT Type
forall rep. Stm rep -> Pattern rep
stmPattern) ([Stm GPU] -> [VName]) -> [Stm GPU] -> [VName]
forall a b. (a -> b) -> a -> b
$
Stms GPU -> [Stm GPU]
forall rep. Stms rep -> [Stm rep]
stmsToList (Stms GPU -> [Stm GPU]) -> Stms GPU -> [Stm GPU]
forall a b. (a -> b) -> a -> b
$ (Stm GPU -> Bool) -> Stms GPU -> Stms GPU
forall a. (a -> Bool) -> Seq a -> Seq a
Seq.filter Stm GPU -> Bool
mustBeInlined Stms GPU
variant_prestms
recompute :: Stm GPU -> Bool
recompute Stm GPU
stm =
(VName -> Bool) -> [VName] -> Bool
forall (t :: * -> *) a. Foldable t => (a -> Bool) -> t a -> Bool
any (VName -> Names -> Bool
`nameIn` Names
must_be_inlined) (PatternT Type -> [VName]
forall dec. PatternT dec -> [VName]
patternNames (Stm GPU -> Pattern GPU
forall rep. Stm rep -> Pattern rep
stmPattern Stm GPU
stm))
Bool -> Bool -> Bool
|| Bool -> Bool
not (Names -> Stm GPU -> Bool
invariantTo Names
must_be_inlined Stm GPU
stm)
(Stms GPU
recomputed_variant_prestms, Stms GPU
precomputed_variant_prestms) =
(Stm GPU -> Bool) -> Stms GPU -> (Stms GPU, Stms GPU)
forall a. (a -> Bool) -> Seq a -> (Seq a, Seq a)
Seq.partition Stm GPU -> Bool
recompute Stms GPU
variant_prestms
injectPrelude ::
SegSpace ->
VarianceTable ->
Stms GPU ->
Names ->
(Stms GPU, Tiling, TiledBody) ->
(Stms GPU, Tiling, TiledBody)
injectPrelude :: SegSpace
-> AliasTable
-> Stms GPU
-> Names
-> (Stms GPU, Tiling, TiledBody)
-> (Stms GPU, Tiling, TiledBody)
injectPrelude SegSpace
initial_space AliasTable
variance Stms GPU
prestms Names
used (Stms GPU
host_stms, Tiling
tiling, TiledBody
tiledBody) =
(Stms GPU
host_stms, Tiling
tiling, TiledBody
tiledBody')
where
tiledBody' :: TiledBody
tiledBody' Names
private PrivStms
privstms = do
let nontiled :: (VName, SubExp) -> Bool
nontiled = ((VName, SubExp) -> [(VName, SubExp)] -> Bool
forall (t :: * -> *) a. (Foldable t, Eq a) => a -> t a -> Bool
`notElem` SegSpace -> [(VName, SubExp)]
unSegSpace (Tiling -> SegSpace
tilingSpace Tiling
tiling))
private' :: Names
private' =
Names
private
Names -> Names -> Names
forall a. Semigroup a => a -> a -> a
<> [VName] -> Names
namesFromList (((VName, SubExp) -> VName) -> [(VName, SubExp)] -> [VName]
forall a b. (a -> b) -> [a] -> [b]
map (VName, SubExp) -> VName
forall a b. (a, b) -> a
fst (((VName, SubExp) -> Bool) -> [(VName, SubExp)] -> [(VName, SubExp)]
forall a. (a -> Bool) -> [a] -> [a]
filter (VName, SubExp) -> Bool
nontiled ([(VName, SubExp)] -> [(VName, SubExp)])
-> [(VName, SubExp)] -> [(VName, SubExp)]
forall a b. (a -> b) -> a -> b
$ SegSpace -> [(VName, SubExp)]
unSegSpace SegSpace
initial_space))
( Stms GPU
invariant_prestms,
Stms GPU
precomputed_variant_prestms,
Stms GPU
recomputed_variant_prestms
) =
AliasTable
-> Stms GPU -> Names -> Names -> (Stms GPU, Stms GPU, Stms GPU)
partitionPrelude AliasTable
variance Stms GPU
prestms Names
private' Names
used
Stms (Rep (BinderT GPU (State VNameSource)))
-> BinderT GPU (State VNameSource) ()
forall (m :: * -> *). MonadBinder m => Stms (Rep m) -> m ()
addStms Stms (Rep (BinderT GPU (State VNameSource)))
Stms GPU
invariant_prestms
let live_set :: [VName]
live_set =
Names -> [VName]
namesToList (Names -> [VName]) -> Names -> [VName]
forall a b. (a -> b) -> a -> b
$
Stms GPU -> Names -> Names
forall a. FreeIn a => Stms GPU -> a -> Names
liveSet Stms GPU
precomputed_variant_prestms (Names -> Names) -> Names -> Names
forall a b. (a -> b) -> a -> b
$
Names
used Names -> Names -> Names
forall a. Semigroup a => a -> a -> a
<> Stms GPU -> Names
forall a. FreeIn a => a -> Names
freeIn Stms GPU
recomputed_variant_prestms
[VName]
prelude_arrs <-
Stms GPU
-> BinderT GPU (State VNameSource) [VName]
-> BinderT GPU (State VNameSource) [VName]
forall rep a (m :: * -> *) b.
(Scoped rep a, LocalScope rep m) =>
a -> m b -> m b
inScopeOf Stms GPU
precomputed_variant_prestms (BinderT GPU (State VNameSource) [VName]
-> BinderT GPU (State VNameSource) [VName])
-> BinderT GPU (State VNameSource) [VName]
-> BinderT GPU (State VNameSource) [VName]
forall a b. (a -> b) -> a -> b
$
Tiling
-> PrivStms
-> Stms GPU
-> [VName]
-> BinderT GPU (State VNameSource) [VName]
doPrelude Tiling
tiling PrivStms
privstms Stms GPU
precomputed_variant_prestms [VName]
live_set
let prelude_privstms :: PrivStms
prelude_privstms =
Stms GPU -> ReadPrelude -> PrivStms
PrivStms Stms GPU
recomputed_variant_prestms (ReadPrelude -> PrivStms) -> ReadPrelude -> PrivStms
forall a b. (a -> b) -> a -> b
$
[VName] -> [VName] -> ReadPrelude
mkReadPreludeValues [VName]
prelude_arrs [VName]
live_set
TiledBody
tiledBody Names
private' (PrivStms
prelude_privstms PrivStms -> PrivStms -> PrivStms
forall a. Semigroup a => a -> a -> a
<> PrivStms
privstms)
tileDoLoop ::
SegSpace ->
VarianceTable ->
Stms GPU ->
Names ->
(Stms GPU, Tiling, TiledBody) ->
[Type] ->
Pattern GPU ->
StmAux (ExpDec GPU) ->
[(FParam GPU, SubExp)] ->
VName ->
IntType ->
SubExp ->
Stms GPU ->
Result ->
TileM (Stms GPU, Tiling, TiledBody)
tileDoLoop :: SegSpace
-> AliasTable
-> Stms GPU
-> Names
-> (Stms GPU, Tiling, TiledBody)
-> [Type]
-> Pattern GPU
-> StmAux (ExpDec GPU)
-> [(FParam GPU, SubExp)]
-> VName
-> IntType
-> SubExp
-> Stms GPU
-> Result
-> ReaderT
(Scope GPU) (State VNameSource) (Stms GPU, Tiling, TiledBody)
tileDoLoop SegSpace
initial_space AliasTable
variance Stms GPU
prestms Names
used_in_body (Stms GPU
host_stms, Tiling
tiling, TiledBody
tiledBody) [Type]
res_ts Pattern GPU
pat StmAux (ExpDec GPU)
aux [(FParam GPU, SubExp)]
merge VName
i IntType
it SubExp
bound Stms GPU
poststms Result
poststms_res = do
let prestms_used :: Names
prestms_used = Names
used_in_body Names -> Names -> Names
forall a. Semigroup a => a -> a -> a
<> Stms GPU -> Names
forall a. FreeIn a => a -> Names
freeIn Stms GPU
poststms Names -> Names -> Names
forall a. Semigroup a => a -> a -> a
<> Result -> Names
forall a. FreeIn a => a -> Names
freeIn Result
poststms_res
( Stms GPU
invariant_prestms,
Stms GPU
precomputed_variant_prestms,
Stms GPU
recomputed_variant_prestms
) =
AliasTable
-> Stms GPU -> Names -> Names -> (Stms GPU, Stms GPU, Stms GPU)
partitionPrelude AliasTable
variance Stms GPU
prestms Names
tiled_kdims Names
prestms_used
let ([Param (TypeBase Shape Uniqueness)]
mergeparams, Result
mergeinits) = [(Param (TypeBase Shape Uniqueness), SubExp)]
-> ([Param (TypeBase Shape Uniqueness)], Result)
forall a b. [(a, b)] -> ([a], [b])
unzip [(Param (TypeBase Shape Uniqueness), SubExp)]
[(FParam GPU, SubExp)]
merge
tileDim :: TypeBase Shape Uniqueness -> TypeBase Shape Uniqueness
tileDim TypeBase Shape Uniqueness
t = TypeBase Shape Uniqueness
-> Shape -> Uniqueness -> TypeBase Shape Uniqueness
forall shape u_unused u.
ArrayShape shape =>
TypeBase shape u_unused -> shape -> u -> TypeBase shape u
arrayOf TypeBase Shape Uniqueness
t (Tiling -> Shape
tilingTileShape Tiling
tiling) (Uniqueness -> TypeBase Shape Uniqueness)
-> Uniqueness -> TypeBase Shape Uniqueness
forall a b. (a -> b) -> a -> b
$ TypeBase Shape Uniqueness -> Uniqueness
forall shape. TypeBase shape Uniqueness -> Uniqueness
uniqueness TypeBase Shape Uniqueness
t
merge_scope :: Scope GPU
merge_scope = VName -> NameInfo GPU -> Scope GPU -> Scope GPU
forall k a. Ord k => k -> a -> Map k a -> Map k a
M.insert VName
i (IntType -> NameInfo GPU
forall rep. IntType -> NameInfo rep
IndexName IntType
it) (Scope GPU -> Scope GPU) -> Scope GPU -> Scope GPU
forall a b. (a -> b) -> a -> b
$ [Param (TypeBase Shape Uniqueness)] -> Scope GPU
forall rep dec. (FParamInfo rep ~ dec) => [Param dec] -> Scope rep
scopeOfFParams [Param (TypeBase Shape Uniqueness)]
mergeparams
tiledBody' :: TiledBody
tiledBody' Names
private PrivStms
privstms = Scope GPU
-> BinderT GPU (State VNameSource) [VName]
-> BinderT GPU (State VNameSource) [VName]
forall rep (m :: * -> *) a.
LocalScope rep m =>
Scope rep -> m a -> m a
localScope (Stms GPU -> Scope GPU
forall rep a. Scoped rep a => a -> Scope rep
scopeOf Stms GPU
host_stms Scope GPU -> Scope GPU -> Scope GPU
forall a. Semigroup a => a -> a -> a
<> Scope GPU
merge_scope) (BinderT GPU (State VNameSource) [VName]
-> BinderT GPU (State VNameSource) [VName])
-> BinderT GPU (State VNameSource) [VName]
-> BinderT GPU (State VNameSource) [VName]
forall a b. (a -> b) -> a -> b
$ do
Stms (Rep (BinderT GPU (State VNameSource)))
-> BinderT GPU (State VNameSource) ()
forall (m :: * -> *). MonadBinder m => Stms (Rep m) -> m ()
addStms Stms (Rep (BinderT GPU (State VNameSource)))
Stms GPU
invariant_prestms
let live_set :: [VName]
live_set =
Names -> [VName]
namesToList (Names -> [VName]) -> Names -> [VName]
forall a b. (a -> b) -> a -> b
$
Stms GPU -> Names -> Names
forall a. FreeIn a => Stms GPU -> a -> Names
liveSet Stms GPU
precomputed_variant_prestms (Names -> Names) -> Names -> Names
forall a b. (a -> b) -> a -> b
$
Stms GPU -> Names
forall a. FreeIn a => a -> Names
freeIn Stms GPU
recomputed_variant_prestms Names -> Names -> Names
forall a. Semigroup a => a -> a -> a
<> Names
prestms_used
[VName]
prelude_arrs <-
Stms GPU
-> BinderT GPU (State VNameSource) [VName]
-> BinderT GPU (State VNameSource) [VName]
forall rep a (m :: * -> *) b.
(Scoped rep a, LocalScope rep m) =>
a -> m b -> m b
inScopeOf Stms GPU
precomputed_variant_prestms (BinderT GPU (State VNameSource) [VName]
-> BinderT GPU (State VNameSource) [VName])
-> BinderT GPU (State VNameSource) [VName]
-> BinderT GPU (State VNameSource) [VName]
forall a b. (a -> b) -> a -> b
$
Tiling
-> PrivStms
-> Stms GPU
-> [VName]
-> BinderT GPU (State VNameSource) [VName]
doPrelude Tiling
tiling PrivStms
privstms Stms GPU
precomputed_variant_prestms [VName]
live_set
[Param (TypeBase Shape Uniqueness)]
mergeparams' <- [Param (TypeBase Shape Uniqueness)]
-> (Param (TypeBase Shape Uniqueness)
-> BinderT
GPU (State VNameSource) (Param (TypeBase Shape Uniqueness)))
-> BinderT
GPU (State VNameSource) [Param (TypeBase Shape Uniqueness)]
forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
t a -> (a -> m b) -> m (t b)
forM [Param (TypeBase Shape Uniqueness)]
mergeparams ((Param (TypeBase Shape Uniqueness)
-> BinderT
GPU (State VNameSource) (Param (TypeBase Shape Uniqueness)))
-> BinderT
GPU (State VNameSource) [Param (TypeBase Shape Uniqueness)])
-> (Param (TypeBase Shape Uniqueness)
-> BinderT
GPU (State VNameSource) (Param (TypeBase Shape Uniqueness)))
-> BinderT
GPU (State VNameSource) [Param (TypeBase Shape Uniqueness)]
forall a b. (a -> b) -> a -> b
$ \(Param VName
pname TypeBase Shape Uniqueness
pt) ->
VName
-> TypeBase Shape Uniqueness -> Param (TypeBase Shape Uniqueness)
forall dec. VName -> dec -> Param dec
Param (VName
-> TypeBase Shape Uniqueness -> Param (TypeBase Shape Uniqueness))
-> BinderT GPU (State VNameSource) VName
-> BinderT
GPU
(State VNameSource)
(TypeBase Shape Uniqueness -> Param (TypeBase Shape Uniqueness))
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> String -> BinderT GPU (State VNameSource) VName
forall (m :: * -> *). MonadFreshNames m => String -> m VName
newVName (VName -> String
baseString VName
pname String -> String -> String
forall a. [a] -> [a] -> [a]
++ String
"_group") BinderT
GPU
(State VNameSource)
(TypeBase Shape Uniqueness -> Param (TypeBase Shape Uniqueness))
-> BinderT GPU (State VNameSource) (TypeBase Shape Uniqueness)
-> BinderT
GPU (State VNameSource) (Param (TypeBase Shape Uniqueness))
forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> TypeBase Shape Uniqueness
-> BinderT GPU (State VNameSource) (TypeBase Shape Uniqueness)
forall (f :: * -> *) a. Applicative f => a -> f a
pure (TypeBase Shape Uniqueness -> TypeBase Shape Uniqueness
tileDim TypeBase Shape Uniqueness
pt)
let merge_ts :: [Type]
merge_ts = (Param (TypeBase Shape Uniqueness) -> Type)
-> [Param (TypeBase Shape Uniqueness)] -> [Type]
forall a b. (a -> b) -> [a] -> [b]
map Param (TypeBase Shape Uniqueness) -> Type
forall dec. Typed dec => Param dec -> Type
paramType [Param (TypeBase Shape Uniqueness)]
mergeparams
let inloop_privstms :: PrivStms
inloop_privstms =
Stms GPU -> ReadPrelude -> PrivStms
PrivStms Stms GPU
recomputed_variant_prestms (ReadPrelude -> PrivStms) -> ReadPrelude -> PrivStms
forall a b. (a -> b) -> a -> b
$
[VName] -> [VName] -> ReadPrelude
mkReadPreludeValues [VName]
prelude_arrs [VName]
live_set
Result
mergeinit' <-
([VName] -> Result)
-> BinderT GPU (State VNameSource) [VName]
-> BinderT GPU (State VNameSource) Result
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap ((VName -> SubExp) -> [VName] -> Result
forall a b. (a -> b) -> [a] -> [b]
map VName -> SubExp
Var) (BinderT GPU (State VNameSource) [VName]
-> BinderT GPU (State VNameSource) Result)
-> BinderT GPU (State VNameSource) [VName]
-> BinderT GPU (State VNameSource) Result
forall a b. (a -> b) -> a -> b
$
Certificates
-> BinderT GPU (State VNameSource) [VName]
-> BinderT GPU (State VNameSource) [VName]
forall (m :: * -> *) a. MonadBinder m => Certificates -> m a -> m a
certifying (StmAux () -> Certificates
forall dec. StmAux dec -> Certificates
stmAuxCerts StmAux ()
StmAux (ExpDec GPU)
aux) (BinderT GPU (State VNameSource) [VName]
-> BinderT GPU (State VNameSource) [VName])
-> BinderT GPU (State VNameSource) [VName]
-> BinderT GPU (State VNameSource) [VName]
forall a b. (a -> b) -> a -> b
$
Tiling
-> String
-> SegLevel
-> ResultManifest
-> (PrimExp VName
-> Slice SubExp -> BinderT GPU (State VNameSource) Result)
-> BinderT GPU (State VNameSource) [VName]
tilingSegMap Tiling
tiling String
"tiled_loopinit" (Tiling -> SegLevel
scalarLevel Tiling
tiling) ResultManifest
ResultPrivate ((PrimExp VName
-> Slice SubExp -> BinderT GPU (State VNameSource) Result)
-> BinderT GPU (State VNameSource) [VName])
-> (PrimExp VName
-> Slice SubExp -> BinderT GPU (State VNameSource) Result)
-> BinderT GPU (State VNameSource) [VName]
forall a b. (a -> b) -> a -> b
$
\PrimExp VName
in_bounds Slice SubExp
slice ->
([VName] -> Result)
-> BinderT GPU (State VNameSource) [VName]
-> BinderT GPU (State VNameSource) Result
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap ((VName -> SubExp) -> [VName] -> Result
forall a b. (a -> b) -> [a] -> [b]
map VName -> SubExp
Var) (BinderT GPU (State VNameSource) [VName]
-> BinderT GPU (State VNameSource) Result)
-> BinderT GPU (State VNameSource) [VName]
-> BinderT GPU (State VNameSource) Result
forall a b. (a -> b) -> a -> b
$
String
-> PrimExp VName
-> [Type]
-> BinderT GPU (State VNameSource) Result
-> BinderT GPU (State VNameSource) [VName]
protectOutOfBounds String
"loopinit" PrimExp VName
in_bounds [Type]
merge_ts (BinderT GPU (State VNameSource) Result
-> BinderT GPU (State VNameSource) [VName])
-> BinderT GPU (State VNameSource) Result
-> BinderT GPU (State VNameSource) [VName]
forall a b. (a -> b) -> a -> b
$ do
Slice SubExp -> PrivStms -> BinderT GPU (State VNameSource) ()
addPrivStms Slice SubExp
slice PrivStms
inloop_privstms
Slice SubExp -> PrivStms -> BinderT GPU (State VNameSource) ()
addPrivStms Slice SubExp
slice PrivStms
privstms
Result -> BinderT GPU (State VNameSource) Result
forall (m :: * -> *) a. Monad m => a -> m a
return Result
mergeinits
let merge' :: [(Param (TypeBase Shape Uniqueness), SubExp)]
merge' = [Param (TypeBase Shape Uniqueness)]
-> Result -> [(Param (TypeBase Shape Uniqueness), SubExp)]
forall a b. [a] -> [b] -> [(a, b)]
zip [Param (TypeBase Shape Uniqueness)]
mergeparams' Result
mergeinit'
let indexMergeParams :: ReadPrelude
indexMergeParams Slice SubExp
slice =
Scope GPU
-> BinderT GPU (State VNameSource) ()
-> BinderT GPU (State VNameSource) ()
forall rep (m :: * -> *) a.
LocalScope rep m =>
Scope rep -> m a -> m a
localScope ([Param (TypeBase Shape Uniqueness)] -> Scope GPU
forall rep dec. (FParamInfo rep ~ dec) => [Param dec] -> Scope rep
scopeOfFParams [Param (TypeBase Shape Uniqueness)]
mergeparams') (BinderT GPU (State VNameSource) ()
-> BinderT GPU (State VNameSource) ())
-> BinderT GPU (State VNameSource) ()
-> BinderT GPU (State VNameSource) ()
forall a b. (a -> b) -> a -> b
$
[(Param (TypeBase Shape Uniqueness),
Param (TypeBase Shape Uniqueness))]
-> ((Param (TypeBase Shape Uniqueness),
Param (TypeBase Shape Uniqueness))
-> BinderT GPU (State VNameSource) ())
-> BinderT GPU (State VNameSource) ()
forall (t :: * -> *) (m :: * -> *) a b.
(Foldable t, Monad m) =>
t a -> (a -> m b) -> m ()
forM_ ([Param (TypeBase Shape Uniqueness)]
-> [Param (TypeBase Shape Uniqueness)]
-> [(Param (TypeBase Shape Uniqueness),
Param (TypeBase Shape Uniqueness))]
forall a b. [a] -> [b] -> [(a, b)]
zip [Param (TypeBase Shape Uniqueness)]
mergeparams [Param (TypeBase Shape Uniqueness)]
mergeparams') (((Param (TypeBase Shape Uniqueness),
Param (TypeBase Shape Uniqueness))
-> BinderT GPU (State VNameSource) ())
-> BinderT GPU (State VNameSource) ())
-> ((Param (TypeBase Shape Uniqueness),
Param (TypeBase Shape Uniqueness))
-> BinderT GPU (State VNameSource) ())
-> BinderT GPU (State VNameSource) ()
forall a b. (a -> b) -> a -> b
$ \(Param (TypeBase Shape Uniqueness)
to, Param (TypeBase Shape Uniqueness)
from) ->
[VName]
-> Exp (Rep (BinderT GPU (State VNameSource)))
-> BinderT GPU (State VNameSource) ()
forall (m :: * -> *).
MonadBinder m =>
[VName] -> Exp (Rep m) -> m ()
letBindNames [Param (TypeBase Shape Uniqueness) -> VName
forall dec. Param dec -> VName
paramName Param (TypeBase Shape Uniqueness)
to] (Exp (Rep (BinderT GPU (State VNameSource)))
-> BinderT GPU (State VNameSource) ())
-> Exp (Rep (BinderT GPU (State VNameSource)))
-> BinderT GPU (State VNameSource) ()
forall a b. (a -> b) -> a -> b
$
BasicOp -> ExpT GPU
forall rep. BasicOp -> ExpT rep
BasicOp (BasicOp -> ExpT GPU) -> BasicOp -> ExpT GPU
forall a b. (a -> b) -> a -> b
$
VName -> Slice SubExp -> BasicOp
Index (Param (TypeBase Shape Uniqueness) -> VName
forall dec. Param dec -> VName
paramName Param (TypeBase Shape Uniqueness)
from) (Slice SubExp -> BasicOp) -> Slice SubExp -> BasicOp
forall a b. (a -> b) -> a -> b
$
Type -> Slice SubExp -> Slice SubExp
fullSlice (Param (TypeBase Shape Uniqueness) -> Type
forall dec. Typed dec => Param dec -> Type
paramType Param (TypeBase Shape Uniqueness)
from) Slice SubExp
slice
private' :: Names
private' =
Names
private Names -> Names -> Names
forall a. Semigroup a => a -> a -> a
<> [VName] -> Names
namesFromList ((Param (TypeBase Shape Uniqueness) -> VName)
-> [Param (TypeBase Shape Uniqueness)] -> [VName]
forall a b. (a -> b) -> [a] -> [b]
map Param (TypeBase Shape Uniqueness) -> VName
forall dec. Param dec -> VName
paramName [Param (TypeBase Shape Uniqueness)]
mergeparams [VName] -> [VName] -> [VName]
forall a. [a] -> [a] -> [a]
++ (Param (TypeBase Shape Uniqueness) -> VName)
-> [Param (TypeBase Shape Uniqueness)] -> [VName]
forall a b. (a -> b) -> [a] -> [b]
map Param (TypeBase Shape Uniqueness) -> VName
forall dec. Param dec -> VName
paramName [Param (TypeBase Shape Uniqueness)]
mergeparams')
privstms' :: PrivStms
privstms' =
Stms GPU -> ReadPrelude -> PrivStms
PrivStms Stms GPU
forall a. Monoid a => a
mempty ReadPrelude
indexMergeParams PrivStms -> PrivStms -> PrivStms
forall a. Semigroup a => a -> a -> a
<> PrivStms
privstms PrivStms -> PrivStms -> PrivStms
forall a. Semigroup a => a -> a -> a
<> PrivStms
inloop_privstms
Body GPU
loopbody' <-
Scope GPU
-> BinderT GPU (State VNameSource) (Body GPU)
-> BinderT GPU (State VNameSource) (Body GPU)
forall rep (m :: * -> *) a.
LocalScope rep m =>
Scope rep -> m a -> m a
localScope ([Param (TypeBase Shape Uniqueness)] -> Scope GPU
forall rep dec. (FParamInfo rep ~ dec) => [Param dec] -> Scope rep
scopeOfFParams [Param (TypeBase Shape Uniqueness)]
mergeparams') (BinderT GPU (State VNameSource) (Body GPU)
-> BinderT GPU (State VNameSource) (Body GPU))
-> (BinderT GPU (State VNameSource) (Body GPU)
-> BinderT GPU (State VNameSource) (Body GPU))
-> BinderT GPU (State VNameSource) (Body GPU)
-> BinderT GPU (State VNameSource) (Body GPU)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. BinderT GPU (State VNameSource) (Body GPU)
-> BinderT GPU (State VNameSource) (Body GPU)
forall rep (m :: * -> *) somerep.
(Bindable rep, MonadFreshNames m, HasScope somerep m,
SameScope somerep rep) =>
Binder rep (Body rep) -> m (Body rep)
runBodyBinder (BinderT GPU (State VNameSource) (Body GPU)
-> BinderT GPU (State VNameSource) (Body GPU))
-> BinderT GPU (State VNameSource) (Body GPU)
-> BinderT GPU (State VNameSource) (Body GPU)
forall a b. (a -> b) -> a -> b
$
Result -> Body GPU
forall rep. Bindable rep => Result -> Body rep
resultBody (Result -> Body GPU) -> ([VName] -> Result) -> [VName] -> Body GPU
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (VName -> SubExp) -> [VName] -> Result
forall a b. (a -> b) -> [a] -> [b]
map VName -> SubExp
Var
([VName] -> Body GPU)
-> BinderT GPU (State VNameSource) [VName]
-> BinderT GPU (State VNameSource) (Body GPU)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> TiledBody
tiledBody Names
private' PrivStms
privstms'
[VName]
accs' <-
String
-> Exp (Rep (BinderT GPU (State VNameSource)))
-> BinderT GPU (State VNameSource) [VName]
forall (m :: * -> *).
MonadBinder m =>
String -> Exp (Rep m) -> m [VName]
letTupExp String
"tiled_inside_loop" (Exp (Rep (BinderT GPU (State VNameSource)))
-> BinderT GPU (State VNameSource) [VName])
-> Exp (Rep (BinderT GPU (State VNameSource)))
-> BinderT GPU (State VNameSource) [VName]
forall a b. (a -> b) -> a -> b
$
[(FParam GPU, SubExp)]
-> [(FParam GPU, SubExp)] -> LoopForm GPU -> Body GPU -> ExpT GPU
forall rep.
[(FParam rep, SubExp)]
-> [(FParam rep, SubExp)] -> LoopForm rep -> BodyT rep -> ExpT rep
DoLoop [] [(Param (TypeBase Shape Uniqueness), SubExp)]
[(FParam GPU, SubExp)]
merge' (VName -> IntType -> SubExp -> [(LParam GPU, VName)] -> LoopForm GPU
forall rep.
VName -> IntType -> SubExp -> [(LParam rep, VName)] -> LoopForm rep
ForLoop VName
i IntType
it SubExp
bound []) Body GPU
loopbody'
Tiling
-> PrivStms
-> Pattern GPU
-> [VName]
-> Stms GPU
-> Result
-> [Type]
-> BinderT GPU (State VNameSource) [VName]
postludeGeneric Tiling
tiling (PrivStms
privstms PrivStms -> PrivStms -> PrivStms
forall a. Semigroup a => a -> a -> a
<> PrivStms
inloop_privstms) Pattern GPU
pat [VName]
accs' Stms GPU
poststms Result
poststms_res [Type]
res_ts
(Stms GPU, Tiling, TiledBody)
-> ReaderT
(Scope GPU) (State VNameSource) (Stms GPU, Tiling, TiledBody)
forall (m :: * -> *) a. Monad m => a -> m a
return (Stms GPU
host_stms, Tiling
tiling, TiledBody
tiledBody')
where
tiled_kdims :: Names
tiled_kdims =
[VName] -> Names
namesFromList ([VName] -> Names) -> [VName] -> Names
forall a b. (a -> b) -> a -> b
$
((VName, SubExp) -> VName) -> [(VName, SubExp)] -> [VName]
forall a b. (a -> b) -> [a] -> [b]
map (VName, SubExp) -> VName
forall a b. (a, b) -> a
fst ([(VName, SubExp)] -> [VName]) -> [(VName, SubExp)] -> [VName]
forall a b. (a -> b) -> a -> b
$
((VName, SubExp) -> Bool) -> [(VName, SubExp)] -> [(VName, SubExp)]
forall a. (a -> Bool) -> [a] -> [a]
filter ((VName, SubExp) -> [(VName, SubExp)] -> Bool
forall (t :: * -> *) a. (Foldable t, Eq a) => a -> t a -> Bool
`notElem` SegSpace -> [(VName, SubExp)]
unSegSpace (Tiling -> SegSpace
tilingSpace Tiling
tiling)) ([(VName, SubExp)] -> [(VName, SubExp)])
-> [(VName, SubExp)] -> [(VName, SubExp)]
forall a b. (a -> b) -> a -> b
$
SegSpace -> [(VName, SubExp)]
unSegSpace SegSpace
initial_space
doPrelude :: Tiling -> PrivStms -> Stms GPU -> [VName] -> Binder GPU [VName]
doPrelude :: Tiling
-> PrivStms
-> Stms GPU
-> [VName]
-> BinderT GPU (State VNameSource) [VName]
doPrelude Tiling
tiling PrivStms
privstms Stms GPU
prestms [VName]
prestms_live =
Tiling
-> String
-> SegLevel
-> ResultManifest
-> (PrimExp VName
-> Slice SubExp -> BinderT GPU (State VNameSource) Result)
-> BinderT GPU (State VNameSource) [VName]
tilingSegMap Tiling
tiling String
"prelude" (Tiling -> SegLevel
scalarLevel Tiling
tiling) ResultManifest
ResultPrivate ((PrimExp VName
-> Slice SubExp -> BinderT GPU (State VNameSource) Result)
-> BinderT GPU (State VNameSource) [VName])
-> (PrimExp VName
-> Slice SubExp -> BinderT GPU (State VNameSource) Result)
-> BinderT GPU (State VNameSource) [VName]
forall a b. (a -> b) -> a -> b
$
\PrimExp VName
in_bounds Slice SubExp
slice -> do
[Type]
ts <- (VName -> BinderT GPU (State VNameSource) Type)
-> [VName] -> BinderT GPU (State VNameSource) [Type]
forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
(a -> m b) -> t a -> m (t b)
mapM VName -> BinderT GPU (State VNameSource) Type
forall rep (m :: * -> *). HasScope rep m => VName -> m Type
lookupType [VName]
prestms_live
([VName] -> Result)
-> BinderT GPU (State VNameSource) [VName]
-> BinderT GPU (State VNameSource) Result
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap ((VName -> SubExp) -> [VName] -> Result
forall a b. (a -> b) -> [a] -> [b]
map VName -> SubExp
Var) (BinderT GPU (State VNameSource) [VName]
-> BinderT GPU (State VNameSource) Result)
-> BinderT GPU (State VNameSource) [VName]
-> BinderT GPU (State VNameSource) Result
forall a b. (a -> b) -> a -> b
$
String
-> PrimExp VName
-> [Type]
-> BinderT GPU (State VNameSource) Result
-> BinderT GPU (State VNameSource) [VName]
protectOutOfBounds String
"pre" PrimExp VName
in_bounds [Type]
ts (BinderT GPU (State VNameSource) Result
-> BinderT GPU (State VNameSource) [VName])
-> BinderT GPU (State VNameSource) Result
-> BinderT GPU (State VNameSource) [VName]
forall a b. (a -> b) -> a -> b
$ do
Slice SubExp -> PrivStms -> BinderT GPU (State VNameSource) ()
addPrivStms Slice SubExp
slice PrivStms
privstms
Stms (Rep (BinderT GPU (State VNameSource)))
-> BinderT GPU (State VNameSource) ()
forall (m :: * -> *). MonadBinder m => Stms (Rep m) -> m ()
addStms Stms (Rep (BinderT GPU (State VNameSource)))
Stms GPU
prestms
Result -> BinderT GPU (State VNameSource) Result
forall (f :: * -> *) a. Applicative f => a -> f a
pure (Result -> BinderT GPU (State VNameSource) Result)
-> Result -> BinderT GPU (State VNameSource) Result
forall a b. (a -> b) -> a -> b
$ (VName -> SubExp) -> [VName] -> Result
forall a b. (a -> b) -> [a] -> [b]
map VName -> SubExp
Var [VName]
prestms_live
liveSet :: FreeIn a => Stms GPU -> a -> Names
liveSet :: forall a. FreeIn a => Stms GPU -> a -> Names
liveSet Stms GPU
stms a
after =
[VName] -> Names
namesFromList ((Stm GPU -> [VName]) -> Stms GPU -> [VName]
forall (t :: * -> *) a b. Foldable t => (a -> [b]) -> t a -> [b]
concatMap (PatternT Type -> [VName]
forall dec. PatternT dec -> [VName]
patternNames (PatternT Type -> [VName])
-> (Stm GPU -> PatternT Type) -> Stm GPU -> [VName]
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Stm GPU -> PatternT Type
forall rep. Stm rep -> Pattern rep
stmPattern) Stms GPU
stms)
Names -> Names -> Names
`namesIntersection` a -> Names
forall a. FreeIn a => a -> Names
freeIn a
after
tileable ::
Stm GPU ->
Maybe
( SubExp,
[VName],
(Commutativity, Lambda GPU, [SubExp], Lambda GPU)
)
tileable :: Stm GPU
-> Maybe
(SubExp, [VName], (Commutativity, Lambda GPU, Result, Lambda GPU))
tileable Stm GPU
stm
| Op (OtherOp (Screma SubExp
w [VName]
arrs ScremaForm GPU
form)) <- Stm GPU -> ExpT GPU
forall rep. Stm rep -> Exp rep
stmExp Stm GPU
stm,
Just ([Reduce GPU]
reds, Lambda GPU
map_lam) <- ScremaForm GPU -> Maybe ([Reduce GPU], Lambda GPU)
forall rep. ScremaForm rep -> Maybe ([Reduce rep], Lambda rep)
isRedomapSOAC ScremaForm GPU
form,
Reduce Commutativity
red_comm Lambda GPU
red_lam Result
red_nes <- [Reduce GPU] -> Reduce GPU
forall rep. Bindable rep => [Reduce rep] -> Reduce rep
singleReduce [Reduce GPU]
reds,
Lambda GPU -> [Type]
forall rep. LambdaT rep -> [Type]
lambdaReturnType Lambda GPU
map_lam [Type] -> [Type] -> Bool
forall a. Eq a => a -> a -> Bool
== Lambda GPU -> [Type]
forall rep. LambdaT rep -> [Type]
lambdaReturnType Lambda GPU
red_lam,
Bool -> Bool
not (Bool -> Bool) -> Bool -> Bool
forall a b. (a -> b) -> a -> b
$ [VName] -> Bool
forall (t :: * -> *) a. Foldable t => t a -> Bool
null [VName]
arrs,
(Type -> Bool) -> [Type] -> Bool
forall (t :: * -> *) a. Foldable t => (a -> Bool) -> t a -> Bool
all Type -> Bool
forall shape u. TypeBase shape u -> Bool
primType ([Type] -> Bool) -> [Type] -> Bool
forall a b. (a -> b) -> a -> b
$ Lambda GPU -> [Type]
forall rep. LambdaT rep -> [Type]
lambdaReturnType Lambda GPU
map_lam,
(Param Type -> Bool) -> [Param Type] -> Bool
forall (t :: * -> *) a. Foldable t => (a -> Bool) -> t a -> Bool
all (Type -> Bool
forall shape u. TypeBase shape u -> Bool
primType (Type -> Bool) -> (Param Type -> Type) -> Param Type -> Bool
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Param Type -> Type
forall dec. Typed dec => Param dec -> Type
paramType) ([Param Type] -> Bool) -> [Param Type] -> Bool
forall a b. (a -> b) -> a -> b
$ Lambda GPU -> [LParam GPU]
forall rep. LambdaT rep -> [LParam rep]
lambdaParams Lambda GPU
map_lam =
(SubExp, [VName], (Commutativity, Lambda GPU, Result, Lambda GPU))
-> Maybe
(SubExp, [VName], (Commutativity, Lambda GPU, Result, Lambda GPU))
forall a. a -> Maybe a
Just (SubExp
w, [VName]
arrs, (Commutativity
red_comm, Lambda GPU
red_lam, Result
red_nes, Lambda GPU
map_lam))
| Bool
otherwise =
Maybe
(SubExp, [VName], (Commutativity, Lambda GPU, Result, Lambda GPU))
forall a. Maybe a
Nothing
data InputArray
= InputTile [Int] VName
| InputDontTile VName
tiledInputs :: [InputArray] -> [(VName, [Int])]
tiledInputs :: [InputArray] -> [(VName, [Int])]
tiledInputs = (InputArray -> Maybe (VName, [Int]))
-> [InputArray] -> [(VName, [Int])]
forall a b. (a -> Maybe b) -> [a] -> [b]
mapMaybe InputArray -> Maybe (VName, [Int])
f
where
f :: InputArray -> Maybe (VName, [Int])
f (InputTile [Int]
perm VName
arr) = (VName, [Int]) -> Maybe (VName, [Int])
forall a. a -> Maybe a
Just (VName
arr, [Int]
perm)
f InputDontTile {} = Maybe (VName, [Int])
forall a. Maybe a
Nothing
data InputTile
= InputTiled [Int] VName
| InputUntiled VName
inputsToTiles :: [InputArray] -> [VName] -> [InputTile]
inputsToTiles :: [InputArray] -> [VName] -> [InputTile]
inputsToTiles (InputTile [Int]
perm VName
_ : [InputArray]
inputs) (VName
tile : [VName]
tiles) =
[Int] -> VName -> InputTile
InputTiled [Int]
perm VName
tile InputTile -> [InputTile] -> [InputTile]
forall a. a -> [a] -> [a]
: [InputArray] -> [VName] -> [InputTile]
inputsToTiles [InputArray]
inputs [VName]
tiles
inputsToTiles (InputDontTile VName
arr : [InputArray]
inputs) [VName]
tiles =
VName -> InputTile
InputUntiled VName
arr InputTile -> [InputTile] -> [InputTile]
forall a. a -> [a] -> [a]
: [InputArray] -> [VName] -> [InputTile]
inputsToTiles [InputArray]
inputs [VName]
tiles
inputsToTiles [InputArray]
_ [VName]
_ = []
sliceUntiled ::
MonadBinder m =>
VName ->
SubExp ->
SubExp ->
SubExp ->
m VName
sliceUntiled :: forall (m :: * -> *).
MonadBinder m =>
VName -> SubExp -> SubExp -> SubExp -> m VName
sliceUntiled VName
arr SubExp
tile_id SubExp
full_tile_size SubExp
this_tile_size = do
Type
arr_t <- VName -> m Type
forall rep (m :: * -> *). HasScope rep m => VName -> m Type
lookupType VName
arr
SubExp
slice_offset <-
String -> ExpT (Rep m) -> m SubExp
forall (m :: * -> *).
MonadBinder m =>
String -> Exp (Rep m) -> m SubExp
letSubExp String
"slice_offset" (ExpT (Rep m) -> m SubExp) -> m (ExpT (Rep m)) -> m SubExp
forall (m :: * -> *) a b. Monad m => (a -> m b) -> m a -> m b
=<< TPrimExp Int64 VName -> m (ExpT (Rep m))
forall a (m :: * -> *).
(ToExp a, MonadBinder m) =>
a -> m (Exp (Rep m))
toExp (SubExp -> TPrimExp Int64 VName
pe64 SubExp
tile_id TPrimExp Int64 VName
-> TPrimExp Int64 VName -> TPrimExp Int64 VName
forall a. Num a => a -> a -> a
* SubExp -> TPrimExp Int64 VName
pe64 SubExp
full_tile_size)
let slice :: DimIndex SubExp
slice = SubExp -> SubExp -> SubExp -> DimIndex SubExp
forall d. d -> d -> d -> DimIndex d
DimSlice SubExp
slice_offset SubExp
this_tile_size (IntType -> Integer -> SubExp
intConst IntType
Int64 Integer
1)
String -> ExpT (Rep m) -> m VName
forall (m :: * -> *).
MonadBinder m =>
String -> Exp (Rep m) -> m VName
letExp String
"untiled_slice" (ExpT (Rep m) -> m VName) -> ExpT (Rep m) -> m VName
forall a b. (a -> b) -> a -> b
$
BasicOp -> ExpT (Rep m)
forall rep. BasicOp -> ExpT rep
BasicOp (BasicOp -> ExpT (Rep m)) -> BasicOp -> ExpT (Rep m)
forall a b. (a -> b) -> a -> b
$ VName -> Slice SubExp -> BasicOp
Index VName
arr (Slice SubExp -> BasicOp) -> Slice SubExp -> BasicOp
forall a b. (a -> b) -> a -> b
$ Type -> Slice SubExp -> Slice SubExp
fullSlice Type
arr_t [DimIndex SubExp
slice]
data PrivStms = PrivStms (Stms GPU) ReadPrelude
privStms :: Stms GPU -> PrivStms
privStms :: Stms GPU -> PrivStms
privStms Stms GPU
stms = Stms GPU -> ReadPrelude -> PrivStms
PrivStms Stms GPU
stms (ReadPrelude -> PrivStms) -> ReadPrelude -> PrivStms
forall a b. (a -> b) -> a -> b
$ BinderT GPU (State VNameSource) () -> ReadPrelude
forall a b. a -> b -> a
const (BinderT GPU (State VNameSource) () -> ReadPrelude)
-> BinderT GPU (State VNameSource) () -> ReadPrelude
forall a b. (a -> b) -> a -> b
$ () -> BinderT GPU (State VNameSource) ()
forall (m :: * -> *) a. Monad m => a -> m a
return ()
addPrivStms :: Slice SubExp -> PrivStms -> Binder GPU ()
addPrivStms :: Slice SubExp -> PrivStms -> BinderT GPU (State VNameSource) ()
addPrivStms Slice SubExp
local_slice (PrivStms Stms GPU
stms ReadPrelude
readPrelude) = do
ReadPrelude
readPrelude Slice SubExp
local_slice
Stms (Rep (BinderT GPU (State VNameSource)))
-> BinderT GPU (State VNameSource) ()
forall (m :: * -> *). MonadBinder m => Stms (Rep m) -> m ()
addStms Stms (Rep (BinderT GPU (State VNameSource)))
Stms GPU
stms
instance Semigroup PrivStms where
PrivStms Stms GPU
stms_x ReadPrelude
readPrelude_x <> :: PrivStms -> PrivStms -> PrivStms
<> PrivStms Stms GPU
stms_y ReadPrelude
readPrelude_y =
Stms GPU -> ReadPrelude -> PrivStms
PrivStms Stms GPU
stms_z ReadPrelude
readPrelude_z
where
stms_z :: Stms GPU
stms_z = Stms GPU
stms_x Stms GPU -> Stms GPU -> Stms GPU
forall a. Semigroup a => a -> a -> a
<> Stms GPU
stms_y
readPrelude_z :: ReadPrelude
readPrelude_z Slice SubExp
slice = ReadPrelude
readPrelude_x Slice SubExp
slice BinderT GPU (State VNameSource) ()
-> BinderT GPU (State VNameSource) ()
-> BinderT GPU (State VNameSource) ()
forall (m :: * -> *) a b. Monad m => m a -> m b -> m b
>> ReadPrelude
readPrelude_y Slice SubExp
slice
instance Monoid PrivStms where
mempty :: PrivStms
mempty = Stms GPU -> PrivStms
privStms Stms GPU
forall a. Monoid a => a
mempty
type ReadPrelude = Slice SubExp -> Binder GPU ()
data ProcessTileArgs = ProcessTileArgs
{ ProcessTileArgs -> PrivStms
processPrivStms :: PrivStms,
ProcessTileArgs -> Commutativity
processComm :: Commutativity,
ProcessTileArgs -> Lambda GPU
processRedLam :: Lambda GPU,
ProcessTileArgs -> Lambda GPU
processMapLam :: Lambda GPU,
ProcessTileArgs -> [InputTile]
processTiles :: [InputTile],
ProcessTileArgs -> [VName]
processAcc :: [VName],
ProcessTileArgs -> SubExp
processTileId :: SubExp
}
data ResidualTileArgs = ResidualTileArgs
{ ResidualTileArgs -> PrivStms
residualPrivStms :: PrivStms,
ResidualTileArgs -> Commutativity
residualComm :: Commutativity,
ResidualTileArgs -> Lambda GPU
residualRedLam :: Lambda GPU,
ResidualTileArgs -> Lambda GPU
residualMapLam :: Lambda GPU,
ResidualTileArgs -> [InputArray]
residualInput :: [InputArray],
ResidualTileArgs -> [VName]
residualAcc :: [VName],
ResidualTileArgs -> SubExp
residualInputSize :: SubExp,
ResidualTileArgs -> SubExp
residualNumWholeTiles :: SubExp
}
data Tiling = Tiling
{ Tiling
-> String
-> SegLevel
-> ResultManifest
-> (PrimExp VName
-> Slice SubExp -> BinderT GPU (State VNameSource) Result)
-> BinderT GPU (State VNameSource) [VName]
tilingSegMap ::
String ->
SegLevel ->
ResultManifest ->
(PrimExp VName -> Slice SubExp -> Binder GPU [SubExp]) ->
Binder GPU [VName],
Tiling
-> TileKind
-> PrivStms
-> SubExp
-> [InputArray]
-> Binder GPU [InputTile]
tilingReadTile ::
TileKind ->
PrivStms ->
SubExp ->
[InputArray] ->
Binder GPU [InputTile],
Tiling
-> ProcessTileArgs -> BinderT GPU (State VNameSource) [VName]
tilingProcessTile ::
ProcessTileArgs ->
Binder GPU [VName],
Tiling
-> ResidualTileArgs -> BinderT GPU (State VNameSource) [VName]
tilingProcessResidualTile ::
ResidualTileArgs ->
Binder GPU [VName],
Tiling -> VName -> BinderT GPU (State VNameSource) KernelResult
tilingTileReturns :: VName -> Binder GPU KernelResult,
Tiling -> SegSpace
tilingSpace :: SegSpace,
Tiling -> Shape
tilingTileShape :: Shape,
Tiling -> SegLevel
tilingLevel :: SegLevel,
Tiling -> Binder GPU SubExp
tilingNumWholeTiles :: Binder GPU SubExp
}
type DoTiling gtids kdims =
SegLevel -> gtids -> kdims -> SubExp -> Binder GPU Tiling
scalarLevel :: Tiling -> SegLevel
scalarLevel :: Tiling -> SegLevel
scalarLevel Tiling
tiling =
Count NumGroups SubExp
-> Count GroupSize SubExp -> SegVirt -> SegLevel
SegThread (SegLevel -> Count NumGroups SubExp
segNumGroups SegLevel
lvl) (SegLevel -> Count GroupSize SubExp
segGroupSize SegLevel
lvl) SegVirt
SegNoVirt
where
lvl :: SegLevel
lvl = Tiling -> SegLevel
tilingLevel Tiling
tiling
protectOutOfBounds ::
String ->
PrimExp VName ->
[Type] ->
Binder GPU [SubExp] ->
Binder GPU [VName]
protectOutOfBounds :: String
-> PrimExp VName
-> [Type]
-> BinderT GPU (State VNameSource) Result
-> BinderT GPU (State VNameSource) [VName]
protectOutOfBounds String
desc PrimExp VName
in_bounds [Type]
ts BinderT GPU (State VNameSource) Result
m = do
Body GPU
m_body <- BinderT
GPU
(State VNameSource)
(Body (Rep (BinderT GPU (State VNameSource))))
-> BinderT
GPU
(State VNameSource)
(Body (Rep (BinderT GPU (State VNameSource))))
forall (m :: * -> *).
MonadBinder m =>
m (Body (Rep m)) -> m (Body (Rep m))
insertStmsM (BinderT
GPU
(State VNameSource)
(Body (Rep (BinderT GPU (State VNameSource))))
-> BinderT
GPU
(State VNameSource)
(Body (Rep (BinderT GPU (State VNameSource)))))
-> BinderT
GPU
(State VNameSource)
(Body (Rep (BinderT GPU (State VNameSource))))
-> BinderT
GPU
(State VNameSource)
(Body (Rep (BinderT GPU (State VNameSource))))
forall a b. (a -> b) -> a -> b
$ Result -> Body GPU
forall rep. Bindable rep => Result -> Body rep
resultBody (Result -> Body GPU)
-> BinderT GPU (State VNameSource) Result
-> BinderT GPU (State VNameSource) (Body GPU)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> BinderT GPU (State VNameSource) Result
m
let m_body_free :: [VName]
m_body_free = Names -> [VName]
namesToList (Names -> [VName]) -> Names -> [VName]
forall a b. (a -> b) -> a -> b
$ Body GPU -> Names
forall a. FreeIn a => a -> Names
freeIn Body GPU
m_body
[(Type, VName)]
t_to_v <-
((Type, VName) -> Bool) -> [(Type, VName)] -> [(Type, VName)]
forall a. (a -> Bool) -> [a] -> [a]
filter (Type -> Bool
forall shape u. TypeBase shape u -> Bool
isAcc (Type -> Bool) -> ((Type, VName) -> Type) -> (Type, VName) -> Bool
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (Type, VName) -> Type
forall a b. (a, b) -> a
fst)
([(Type, VName)] -> [(Type, VName)])
-> BinderT GPU (State VNameSource) [(Type, VName)]
-> BinderT GPU (State VNameSource) [(Type, VName)]
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> ([Type] -> [VName] -> [(Type, VName)]
forall a b. [a] -> [b] -> [(a, b)]
zip ([Type] -> [VName] -> [(Type, VName)])
-> BinderT GPU (State VNameSource) [Type]
-> BinderT GPU (State VNameSource) ([VName] -> [(Type, VName)])
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> (VName -> BinderT GPU (State VNameSource) Type)
-> [VName] -> BinderT GPU (State VNameSource) [Type]
forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
(a -> m b) -> t a -> m (t b)
mapM VName -> BinderT GPU (State VNameSource) Type
forall rep (m :: * -> *). HasScope rep m => VName -> m Type
lookupType [VName]
m_body_free BinderT GPU (State VNameSource) ([VName] -> [(Type, VName)])
-> BinderT GPU (State VNameSource) [VName]
-> BinderT GPU (State VNameSource) [(Type, VName)]
forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> [VName] -> BinderT GPU (State VNameSource) [VName]
forall (f :: * -> *) a. Applicative f => a -> f a
pure [VName]
m_body_free)
let blank :: Type -> BinderT GPU (State VNameSource) (ExpT GPU)
blank Type
t = BinderT GPU (State VNameSource) (ExpT GPU)
-> (VName -> BinderT GPU (State VNameSource) (ExpT GPU))
-> Maybe VName
-> BinderT GPU (State VNameSource) (ExpT GPU)
forall b a. b -> (a -> b) -> Maybe a -> b
maybe (Type
-> BinderT
GPU
(State VNameSource)
(Exp (Rep (BinderT GPU (State VNameSource))))
forall (m :: * -> *). MonadBinder m => Type -> m (Exp (Rep m))
eBlank Type
t) (ExpT GPU -> BinderT GPU (State VNameSource) (ExpT GPU)
forall (f :: * -> *) a. Applicative f => a -> f a
pure (ExpT GPU -> BinderT GPU (State VNameSource) (ExpT GPU))
-> (VName -> ExpT GPU)
-> VName
-> BinderT GPU (State VNameSource) (ExpT GPU)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. BasicOp -> ExpT GPU
forall rep. BasicOp -> ExpT rep
BasicOp (BasicOp -> ExpT GPU) -> (VName -> BasicOp) -> VName -> ExpT GPU
forall b c a. (b -> c) -> (a -> b) -> a -> c
. SubExp -> BasicOp
SubExp (SubExp -> BasicOp) -> (VName -> SubExp) -> VName -> BasicOp
forall b c a. (b -> c) -> (a -> b) -> a -> c
. VName -> SubExp
Var) (Maybe VName -> BinderT GPU (State VNameSource) (ExpT GPU))
-> Maybe VName -> BinderT GPU (State VNameSource) (ExpT GPU)
forall a b. (a -> b) -> a -> b
$ Type -> [(Type, VName)] -> Maybe VName
forall a b. Eq a => a -> [(a, b)] -> Maybe b
lookup Type
t [(Type, VName)]
t_to_v
String
-> Exp (Rep (BinderT GPU (State VNameSource)))
-> BinderT GPU (State VNameSource) [VName]
forall (m :: * -> *).
MonadBinder m =>
String -> Exp (Rep m) -> m [VName]
letTupExp String
desc (ExpT GPU -> BinderT GPU (State VNameSource) [VName])
-> BinderT GPU (State VNameSource) (ExpT GPU)
-> BinderT GPU (State VNameSource) [VName]
forall (m :: * -> *) a b. Monad m => (a -> m b) -> m a -> m b
=<< BinderT
GPU
(State VNameSource)
(Exp (Rep (BinderT GPU (State VNameSource))))
-> BinderT
GPU
(State VNameSource)
(Body (Rep (BinderT GPU (State VNameSource))))
-> BinderT
GPU
(State VNameSource)
(Body (Rep (BinderT GPU (State VNameSource))))
-> BinderT
GPU
(State VNameSource)
(Exp (Rep (BinderT GPU (State VNameSource))))
forall (m :: * -> *).
(MonadBinder m, BranchType (Rep m) ~ ExtType) =>
m (Exp (Rep m))
-> m (Body (Rep m)) -> m (Body (Rep m)) -> m (Exp (Rep m))
eIf (PrimExp VName
-> BinderT
GPU
(State VNameSource)
(Exp (Rep (BinderT GPU (State VNameSource))))
forall a (m :: * -> *).
(ToExp a, MonadBinder m) =>
a -> m (Exp (Rep m))
toExp PrimExp VName
in_bounds) (Body GPU -> BinderT GPU (State VNameSource) (Body GPU)
forall (f :: * -> *) a. Applicative f => a -> f a
pure Body GPU
m_body) ([BinderT
GPU
(State VNameSource)
(Exp (Rep (BinderT GPU (State VNameSource))))]
-> BinderT
GPU
(State VNameSource)
(Body (Rep (BinderT GPU (State VNameSource))))
forall (m :: * -> *).
MonadBinder m =>
[m (Exp (Rep m))] -> m (Body (Rep m))
eBody ([BinderT
GPU
(State VNameSource)
(Exp (Rep (BinderT GPU (State VNameSource))))]
-> BinderT
GPU
(State VNameSource)
(Body (Rep (BinderT GPU (State VNameSource)))))
-> [BinderT
GPU
(State VNameSource)
(Exp (Rep (BinderT GPU (State VNameSource))))]
-> BinderT
GPU
(State VNameSource)
(Body (Rep (BinderT GPU (State VNameSource))))
forall a b. (a -> b) -> a -> b
$ (Type -> BinderT GPU (State VNameSource) (ExpT GPU))
-> [Type] -> [BinderT GPU (State VNameSource) (ExpT GPU)]
forall a b. (a -> b) -> [a] -> [b]
map Type -> BinderT GPU (State VNameSource) (ExpT GPU)
blank [Type]
ts)
where
isAcc :: TypeBase shape u -> Bool
isAcc Acc {} = Bool
True
isAcc TypeBase shape u
_ = Bool
False
postludeGeneric ::
Tiling ->
PrivStms ->
Pattern GPU ->
[VName] ->
Stms GPU ->
Result ->
[Type] ->
Binder GPU [VName]
postludeGeneric :: Tiling
-> PrivStms
-> Pattern GPU
-> [VName]
-> Stms GPU
-> Result
-> [Type]
-> BinderT GPU (State VNameSource) [VName]
postludeGeneric Tiling
tiling PrivStms
privstms Pattern GPU
pat [VName]
accs' Stms GPU
poststms Result
poststms_res [Type]
res_ts =
Tiling
-> String
-> SegLevel
-> ResultManifest
-> (PrimExp VName
-> Slice SubExp -> BinderT GPU (State VNameSource) Result)
-> BinderT GPU (State VNameSource) [VName]
tilingSegMap Tiling
tiling String
"thread_res" (Tiling -> SegLevel
scalarLevel Tiling
tiling) ResultManifest
ResultPrivate ((PrimExp VName
-> Slice SubExp -> BinderT GPU (State VNameSource) Result)
-> BinderT GPU (State VNameSource) [VName])
-> (PrimExp VName
-> Slice SubExp -> BinderT GPU (State VNameSource) Result)
-> BinderT GPU (State VNameSource) [VName]
forall a b. (a -> b) -> a -> b
$ \PrimExp VName
in_bounds Slice SubExp
slice -> do
[(VName, VName)]
-> ((VName, VName) -> BinderT GPU (State VNameSource) ())
-> BinderT GPU (State VNameSource) ()
forall (t :: * -> *) (m :: * -> *) a b.
(Foldable t, Monad m) =>
t a -> (a -> m b) -> m ()
forM_ ([VName] -> [VName] -> [(VName, VName)]
forall a b. [a] -> [b] -> [(a, b)]
zip (PatternT Type -> [VName]
forall dec. PatternT dec -> [VName]
patternNames PatternT Type
Pattern GPU
pat) [VName]
accs') (((VName, VName) -> BinderT GPU (State VNameSource) ())
-> BinderT GPU (State VNameSource) ())
-> ((VName, VName) -> BinderT GPU (State VNameSource) ())
-> BinderT GPU (State VNameSource) ()
forall a b. (a -> b) -> a -> b
$ \(VName
us, VName
everyone) -> do
Type
everyone_t <- VName -> BinderT GPU (State VNameSource) Type
forall rep (m :: * -> *). HasScope rep m => VName -> m Type
lookupType VName
everyone
[VName]
-> Exp (Rep (BinderT GPU (State VNameSource)))
-> BinderT GPU (State VNameSource) ()
forall (m :: * -> *).
MonadBinder m =>
[VName] -> Exp (Rep m) -> m ()
letBindNames [VName
us] (Exp (Rep (BinderT GPU (State VNameSource)))
-> BinderT GPU (State VNameSource) ())
-> Exp (Rep (BinderT GPU (State VNameSource)))
-> BinderT GPU (State VNameSource) ()
forall a b. (a -> b) -> a -> b
$ BasicOp -> ExpT GPU
forall rep. BasicOp -> ExpT rep
BasicOp (BasicOp -> ExpT GPU) -> BasicOp -> ExpT GPU
forall a b. (a -> b) -> a -> b
$ VName -> Slice SubExp -> BasicOp
Index VName
everyone (Slice SubExp -> BasicOp) -> Slice SubExp -> BasicOp
forall a b. (a -> b) -> a -> b
$ Type -> Slice SubExp -> Slice SubExp
fullSlice Type
everyone_t Slice SubExp
slice
if Stms GPU
poststms Stms GPU -> Stms GPU -> Bool
forall a. Eq a => a -> a -> Bool
== Stms GPU
forall a. Monoid a => a
mempty
then do
Slice SubExp -> PrivStms -> BinderT GPU (State VNameSource) ()
addPrivStms Slice SubExp
slice PrivStms
privstms
Result -> BinderT GPU (State VNameSource) Result
forall (m :: * -> *) a. Monad m => a -> m a
return Result
poststms_res
else ([VName] -> Result)
-> BinderT GPU (State VNameSource) [VName]
-> BinderT GPU (State VNameSource) Result
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap ((VName -> SubExp) -> [VName] -> Result
forall a b. (a -> b) -> [a] -> [b]
map VName -> SubExp
Var) (BinderT GPU (State VNameSource) [VName]
-> BinderT GPU (State VNameSource) Result)
-> BinderT GPU (State VNameSource) [VName]
-> BinderT GPU (State VNameSource) Result
forall a b. (a -> b) -> a -> b
$
String
-> PrimExp VName
-> [Type]
-> BinderT GPU (State VNameSource) Result
-> BinderT GPU (State VNameSource) [VName]
protectOutOfBounds String
"postlude" PrimExp VName
in_bounds [Type]
res_ts (BinderT GPU (State VNameSource) Result
-> BinderT GPU (State VNameSource) [VName])
-> BinderT GPU (State VNameSource) Result
-> BinderT GPU (State VNameSource) [VName]
forall a b. (a -> b) -> a -> b
$ do
Slice SubExp -> PrivStms -> BinderT GPU (State VNameSource) ()
addPrivStms Slice SubExp
slice PrivStms
privstms
Stms (Rep (BinderT GPU (State VNameSource)))
-> BinderT GPU (State VNameSource) ()
forall (m :: * -> *). MonadBinder m => Stms (Rep m) -> m ()
addStms Stms (Rep (BinderT GPU (State VNameSource)))
Stms GPU
poststms
Result -> BinderT GPU (State VNameSource) Result
forall (m :: * -> *) a. Monad m => a -> m a
return Result
poststms_res
type TiledBody = Names -> PrivStms -> Binder GPU [VName]
tileGeneric ::
DoTiling gtids kdims ->
SegLevel ->
[Type] ->
Pattern GPU ->
gtids ->
kdims ->
SubExp ->
(Commutativity, Lambda GPU, [SubExp], Lambda GPU) ->
[InputArray] ->
Stms GPU ->
Result ->
TileM (Stms GPU, Tiling, TiledBody)
tileGeneric :: forall gtids kdims.
DoTiling gtids kdims
-> SegLevel
-> [Type]
-> Pattern GPU
-> gtids
-> kdims
-> SubExp
-> (Commutativity, Lambda GPU, Result, Lambda GPU)
-> [InputArray]
-> Stms GPU
-> Result
-> ReaderT
(Scope GPU) (State VNameSource) (Stms GPU, Tiling, TiledBody)
tileGeneric DoTiling gtids kdims
doTiling SegLevel
initial_lvl [Type]
res_ts Pattern GPU
pat gtids
gtids kdims
kdims SubExp
w (Commutativity, Lambda GPU, Result, Lambda GPU)
form [InputArray]
inputs Stms GPU
poststms Result
poststms_res = do
(Tiling
tiling, Stms GPU
tiling_stms) <- Binder GPU Tiling
-> ReaderT (Scope GPU) (State VNameSource) (Tiling, Stms GPU)
forall (m :: * -> *) somerep rep a.
(MonadFreshNames m, HasScope somerep m, SameScope somerep rep) =>
Binder rep a -> m (a, Stms rep)
runBinder (Binder GPU Tiling
-> ReaderT (Scope GPU) (State VNameSource) (Tiling, Stms GPU))
-> Binder GPU Tiling
-> ReaderT (Scope GPU) (State VNameSource) (Tiling, Stms GPU)
forall a b. (a -> b) -> a -> b
$ DoTiling gtids kdims
doTiling SegLevel
initial_lvl gtids
gtids kdims
kdims SubExp
w
(Stms GPU, Tiling, TiledBody)
-> ReaderT
(Scope GPU) (State VNameSource) (Stms GPU, Tiling, TiledBody)
forall (m :: * -> *) a. Monad m => a -> m a
return (Stms GPU
tiling_stms, Tiling
tiling, Tiling -> TiledBody
tiledBody Tiling
tiling)
where
(Commutativity
red_comm, Lambda GPU
red_lam, Result
red_nes, Lambda GPU
map_lam) = (Commutativity, Lambda GPU, Result, Lambda GPU)
form
tiledBody :: Tiling -> Names -> PrivStms -> Binder GPU [VName]
tiledBody :: Tiling -> TiledBody
tiledBody Tiling
tiling Names
_private PrivStms
privstms = do
let tile_shape :: Shape
tile_shape = Tiling -> Shape
tilingTileShape Tiling
tiling
SubExp
num_whole_tiles <- Tiling -> Binder GPU SubExp
tilingNumWholeTiles Tiling
tiling
[VName]
mergeinits <- Tiling
-> String
-> SegLevel
-> ResultManifest
-> (PrimExp VName
-> Slice SubExp -> BinderT GPU (State VNameSource) Result)
-> BinderT GPU (State VNameSource) [VName]
tilingSegMap Tiling
tiling String
"mergeinit" (Tiling -> SegLevel
scalarLevel Tiling
tiling) ResultManifest
ResultPrivate ((PrimExp VName
-> Slice SubExp -> BinderT GPU (State VNameSource) Result)
-> BinderT GPU (State VNameSource) [VName])
-> (PrimExp VName
-> Slice SubExp -> BinderT GPU (State VNameSource) Result)
-> BinderT GPU (State VNameSource) [VName]
forall a b. (a -> b) -> a -> b
$ \PrimExp VName
in_bounds Slice SubExp
slice ->
if Result -> Names
forall a. FreeIn a => a -> Names
freeIn Result
red_nes Names -> Names -> Bool
forall a. Eq a => a -> a -> Bool
== Names
forall a. Monoid a => a
mempty
then Result -> BinderT GPU (State VNameSource) Result
forall (m :: * -> *) a. Monad m => a -> m a
return Result
red_nes
else ([VName] -> Result)
-> BinderT GPU (State VNameSource) [VName]
-> BinderT GPU (State VNameSource) Result
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap ((VName -> SubExp) -> [VName] -> Result
forall a b. (a -> b) -> [a] -> [b]
map VName -> SubExp
Var) (BinderT GPU (State VNameSource) [VName]
-> BinderT GPU (State VNameSource) Result)
-> BinderT GPU (State VNameSource) [VName]
-> BinderT GPU (State VNameSource) Result
forall a b. (a -> b) -> a -> b
$
String
-> PrimExp VName
-> [Type]
-> BinderT GPU (State VNameSource) Result
-> BinderT GPU (State VNameSource) [VName]
protectOutOfBounds String
"neutral" PrimExp VName
in_bounds (Lambda GPU -> [Type]
forall rep. LambdaT rep -> [Type]
lambdaReturnType Lambda GPU
red_lam) (BinderT GPU (State VNameSource) Result
-> BinderT GPU (State VNameSource) [VName])
-> BinderT GPU (State VNameSource) Result
-> BinderT GPU (State VNameSource) [VName]
forall a b. (a -> b) -> a -> b
$ do
Slice SubExp -> PrivStms -> BinderT GPU (State VNameSource) ()
addPrivStms Slice SubExp
slice PrivStms
privstms
Result -> BinderT GPU (State VNameSource) Result
forall (m :: * -> *) a. Monad m => a -> m a
return Result
red_nes
[(Param (TypeBase Shape Uniqueness), SubExp)]
merge <- [(Param Type, VName)]
-> ((Param Type, VName)
-> BinderT
GPU
(State VNameSource)
(Param (TypeBase Shape Uniqueness), SubExp))
-> BinderT
GPU
(State VNameSource)
[(Param (TypeBase Shape Uniqueness), SubExp)]
forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
t a -> (a -> m b) -> m (t b)
forM ([Param Type] -> [VName] -> [(Param Type, VName)]
forall a b. [a] -> [b] -> [(a, b)]
zip (Lambda GPU -> [LParam GPU]
forall rep. LambdaT rep -> [LParam rep]
lambdaParams Lambda GPU
red_lam) [VName]
mergeinits) (((Param Type, VName)
-> BinderT
GPU
(State VNameSource)
(Param (TypeBase Shape Uniqueness), SubExp))
-> BinderT
GPU
(State VNameSource)
[(Param (TypeBase Shape Uniqueness), SubExp)])
-> ((Param Type, VName)
-> BinderT
GPU
(State VNameSource)
(Param (TypeBase Shape Uniqueness), SubExp))
-> BinderT
GPU
(State VNameSource)
[(Param (TypeBase Shape Uniqueness), SubExp)]
forall a b. (a -> b) -> a -> b
$ \(Param Type
p, VName
mergeinit) ->
(,)
(Param (TypeBase Shape Uniqueness)
-> SubExp -> (Param (TypeBase Shape Uniqueness), SubExp))
-> BinderT
GPU (State VNameSource) (Param (TypeBase Shape Uniqueness))
-> BinderT
GPU
(State VNameSource)
(SubExp -> (Param (TypeBase Shape Uniqueness), SubExp))
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> String
-> TypeBase Shape Uniqueness
-> BinderT
GPU (State VNameSource) (Param (TypeBase Shape Uniqueness))
forall (m :: * -> *) dec.
MonadFreshNames m =>
String -> dec -> m (Param dec)
newParam
(VName -> String
baseString (Param Type -> VName
forall dec. Param dec -> VName
paramName Param Type
p) String -> String -> String
forall a. [a] -> [a] -> [a]
++ String
"_merge")
(Param Type -> Type
forall dec. Typed dec => Param dec -> Type
paramType Param Type
p Type -> Shape -> Type
`arrayOfShape` Shape
tile_shape Type -> Uniqueness -> TypeBase Shape Uniqueness
forall shape.
TypeBase shape NoUniqueness
-> Uniqueness -> TypeBase shape Uniqueness
`toDecl` Uniqueness
Unique)
BinderT
GPU
(State VNameSource)
(SubExp -> (Param (TypeBase Shape Uniqueness), SubExp))
-> Binder GPU SubExp
-> BinderT
GPU (State VNameSource) (Param (TypeBase Shape Uniqueness), SubExp)
forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> SubExp -> Binder GPU SubExp
forall (f :: * -> *) a. Applicative f => a -> f a
pure (VName -> SubExp
Var VName
mergeinit)
VName
tile_id <- String -> BinderT GPU (State VNameSource) VName
forall (m :: * -> *). MonadFreshNames m => String -> m VName
newVName String
"tile_id"
let loopform :: LoopForm GPU
loopform = VName -> IntType -> SubExp -> [(LParam GPU, VName)] -> LoopForm GPU
forall rep.
VName -> IntType -> SubExp -> [(LParam rep, VName)] -> LoopForm rep
ForLoop VName
tile_id IntType
Int64 SubExp
num_whole_tiles []
Body GPU
loopbody <- Body GPU -> BinderT GPU (State VNameSource) (Body GPU)
forall rep (m :: * -> *).
(Renameable rep, MonadFreshNames m) =>
Body rep -> m (Body rep)
renameBody (Body GPU -> BinderT GPU (State VNameSource) (Body GPU))
-> (BinderT GPU (State VNameSource) (Body GPU)
-> BinderT GPU (State VNameSource) (Body GPU))
-> BinderT GPU (State VNameSource) (Body GPU)
-> BinderT GPU (State VNameSource) (Body GPU)
forall (m :: * -> *) b c a.
Monad m =>
(b -> m c) -> (a -> m b) -> a -> m c
<=< BinderT GPU (State VNameSource) (Body GPU)
-> BinderT GPU (State VNameSource) (Body GPU)
forall rep (m :: * -> *) somerep.
(Bindable rep, MonadFreshNames m, HasScope somerep m,
SameScope somerep rep) =>
Binder rep (Body rep) -> m (Body rep)
runBodyBinder (BinderT GPU (State VNameSource) (Body GPU)
-> BinderT GPU (State VNameSource) (Body GPU))
-> BinderT GPU (State VNameSource) (Body GPU)
-> BinderT GPU (State VNameSource) (Body GPU)
forall a b. (a -> b) -> a -> b
$
LoopForm GPU
-> BinderT GPU (State VNameSource) (Body GPU)
-> BinderT GPU (State VNameSource) (Body GPU)
forall rep a (m :: * -> *) b.
(Scoped rep a, LocalScope rep m) =>
a -> m b -> m b
inScopeOf LoopForm GPU
loopform (BinderT GPU (State VNameSource) (Body GPU)
-> BinderT GPU (State VNameSource) (Body GPU))
-> BinderT GPU (State VNameSource) (Body GPU)
-> BinderT GPU (State VNameSource) (Body GPU)
forall a b. (a -> b) -> a -> b
$
Scope GPU
-> BinderT GPU (State VNameSource) (Body GPU)
-> BinderT GPU (State VNameSource) (Body GPU)
forall rep (m :: * -> *) a.
LocalScope rep m =>
Scope rep -> m a -> m a
localScope ([Param (TypeBase Shape Uniqueness)] -> Scope GPU
forall rep dec. (FParamInfo rep ~ dec) => [Param dec] -> Scope rep
scopeOfFParams ([Param (TypeBase Shape Uniqueness)] -> Scope GPU)
-> [Param (TypeBase Shape Uniqueness)] -> Scope GPU
forall a b. (a -> b) -> a -> b
$ ((Param (TypeBase Shape Uniqueness), SubExp)
-> Param (TypeBase Shape Uniqueness))
-> [(Param (TypeBase Shape Uniqueness), SubExp)]
-> [Param (TypeBase Shape Uniqueness)]
forall a b. (a -> b) -> [a] -> [b]
map (Param (TypeBase Shape Uniqueness), SubExp)
-> Param (TypeBase Shape Uniqueness)
forall a b. (a, b) -> a
fst [(Param (TypeBase Shape Uniqueness), SubExp)]
merge) (BinderT GPU (State VNameSource) (Body GPU)
-> BinderT GPU (State VNameSource) (Body GPU))
-> BinderT GPU (State VNameSource) (Body GPU)
-> BinderT GPU (State VNameSource) (Body GPU)
forall a b. (a -> b) -> a -> b
$ do
[InputTile]
tile <- Tiling
-> TileKind
-> PrivStms
-> SubExp
-> [InputArray]
-> Binder GPU [InputTile]
tilingReadTile Tiling
tiling TileKind
TilePartial PrivStms
privstms (VName -> SubExp
Var VName
tile_id) [InputArray]
inputs
let accs :: [VName]
accs =
((Param (TypeBase Shape Uniqueness), SubExp) -> VName)
-> [(Param (TypeBase Shape Uniqueness), SubExp)] -> [VName]
forall a b. (a -> b) -> [a] -> [b]
map (Param (TypeBase Shape Uniqueness) -> VName
forall dec. Param dec -> VName
paramName (Param (TypeBase Shape Uniqueness) -> VName)
-> ((Param (TypeBase Shape Uniqueness), SubExp)
-> Param (TypeBase Shape Uniqueness))
-> (Param (TypeBase Shape Uniqueness), SubExp)
-> VName
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (Param (TypeBase Shape Uniqueness), SubExp)
-> Param (TypeBase Shape Uniqueness)
forall a b. (a, b) -> a
fst) [(Param (TypeBase Shape Uniqueness), SubExp)]
merge
tile_args :: ProcessTileArgs
tile_args =
PrivStms
-> Commutativity
-> Lambda GPU
-> Lambda GPU
-> [InputTile]
-> [VName]
-> SubExp
-> ProcessTileArgs
ProcessTileArgs PrivStms
privstms Commutativity
red_comm Lambda GPU
red_lam Lambda GPU
map_lam [InputTile]
tile [VName]
accs (VName -> SubExp
Var VName
tile_id)
Result -> Body GPU
forall rep. Bindable rep => Result -> Body rep
resultBody (Result -> Body GPU) -> ([VName] -> Result) -> [VName] -> Body GPU
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (VName -> SubExp) -> [VName] -> Result
forall a b. (a -> b) -> [a] -> [b]
map VName -> SubExp
Var ([VName] -> Body GPU)
-> BinderT GPU (State VNameSource) [VName]
-> BinderT GPU (State VNameSource) (Body GPU)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Tiling
-> ProcessTileArgs -> BinderT GPU (State VNameSource) [VName]
tilingProcessTile Tiling
tiling ProcessTileArgs
tile_args
[VName]
accs <- String
-> Exp (Rep (BinderT GPU (State VNameSource)))
-> BinderT GPU (State VNameSource) [VName]
forall (m :: * -> *).
MonadBinder m =>
String -> Exp (Rep m) -> m [VName]
letTupExp String
"accs" (Exp (Rep (BinderT GPU (State VNameSource)))
-> BinderT GPU (State VNameSource) [VName])
-> Exp (Rep (BinderT GPU (State VNameSource)))
-> BinderT GPU (State VNameSource) [VName]
forall a b. (a -> b) -> a -> b
$ [(FParam GPU, SubExp)]
-> [(FParam GPU, SubExp)] -> LoopForm GPU -> Body GPU -> ExpT GPU
forall rep.
[(FParam rep, SubExp)]
-> [(FParam rep, SubExp)] -> LoopForm rep -> BodyT rep -> ExpT rep
DoLoop [] [(Param (TypeBase Shape Uniqueness), SubExp)]
[(FParam GPU, SubExp)]
merge LoopForm GPU
loopform Body GPU
loopbody
Lambda GPU
red_lam' <- Lambda GPU -> BinderT GPU (State VNameSource) (Lambda GPU)
forall rep (m :: * -> *).
(Renameable rep, MonadFreshNames m) =>
Lambda rep -> m (Lambda rep)
renameLambda Lambda GPU
red_lam
Lambda GPU
map_lam' <- Lambda GPU -> BinderT GPU (State VNameSource) (Lambda GPU)
forall rep (m :: * -> *).
(Renameable rep, MonadFreshNames m) =>
Lambda rep -> m (Lambda rep)
renameLambda Lambda GPU
map_lam
let residual_args :: ResidualTileArgs
residual_args =
PrivStms
-> Commutativity
-> Lambda GPU
-> Lambda GPU
-> [InputArray]
-> [VName]
-> SubExp
-> SubExp
-> ResidualTileArgs
ResidualTileArgs PrivStms
privstms Commutativity
red_comm Lambda GPU
red_lam' Lambda GPU
map_lam' [InputArray]
inputs [VName]
accs SubExp
w SubExp
num_whole_tiles
[VName]
accs' <- Tiling
-> ResidualTileArgs -> BinderT GPU (State VNameSource) [VName]
tilingProcessResidualTile Tiling
tiling ResidualTileArgs
residual_args
Tiling
-> PrivStms
-> Pattern GPU
-> [VName]
-> Stms GPU
-> Result
-> [Type]
-> BinderT GPU (State VNameSource) [VName]
postludeGeneric Tiling
tiling PrivStms
privstms Pattern GPU
pat [VName]
accs' Stms GPU
poststms Result
poststms_res [Type]
res_ts
data TileKind = TilePartial | TileFull
mkReadPreludeValues :: [VName] -> [VName] -> ReadPrelude
mkReadPreludeValues :: [VName] -> [VName] -> ReadPrelude
mkReadPreludeValues [VName]
prestms_live_arrs [VName]
prestms_live Slice SubExp
slice =
([()] -> ())
-> BinderT GPU (State VNameSource) [()]
-> BinderT GPU (State VNameSource) ()
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap [()] -> ()
forall a. Monoid a => [a] -> a
mconcat (BinderT GPU (State VNameSource) [()]
-> BinderT GPU (State VNameSource) ())
-> BinderT GPU (State VNameSource) [()]
-> BinderT GPU (State VNameSource) ()
forall a b. (a -> b) -> a -> b
$
[(VName, VName)]
-> ((VName, VName) -> BinderT GPU (State VNameSource) ())
-> BinderT GPU (State VNameSource) [()]
forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
t a -> (a -> m b) -> m (t b)
forM ([VName] -> [VName] -> [(VName, VName)]
forall a b. [a] -> [b] -> [(a, b)]
zip [VName]
prestms_live_arrs [VName]
prestms_live) (((VName, VName) -> BinderT GPU (State VNameSource) ())
-> BinderT GPU (State VNameSource) [()])
-> ((VName, VName) -> BinderT GPU (State VNameSource) ())
-> BinderT GPU (State VNameSource) [()]
forall a b. (a -> b) -> a -> b
$ \(VName
arr, VName
v) -> do
Type
arr_t <- VName -> BinderT GPU (State VNameSource) Type
forall rep (m :: * -> *). HasScope rep m => VName -> m Type
lookupType VName
arr
[VName]
-> Exp (Rep (BinderT GPU (State VNameSource)))
-> BinderT GPU (State VNameSource) ()
forall (m :: * -> *).
MonadBinder m =>
[VName] -> Exp (Rep m) -> m ()
letBindNames [VName
v] (Exp (Rep (BinderT GPU (State VNameSource)))
-> BinderT GPU (State VNameSource) ())
-> Exp (Rep (BinderT GPU (State VNameSource)))
-> BinderT GPU (State VNameSource) ()
forall a b. (a -> b) -> a -> b
$ BasicOp -> ExpT GPU
forall rep. BasicOp -> ExpT rep
BasicOp (BasicOp -> ExpT GPU) -> BasicOp -> ExpT GPU
forall a b. (a -> b) -> a -> b
$ VName -> Slice SubExp -> BasicOp
Index VName
arr (Slice SubExp -> BasicOp) -> Slice SubExp -> BasicOp
forall a b. (a -> b) -> a -> b
$ Type -> Slice SubExp -> Slice SubExp
fullSlice Type
arr_t Slice SubExp
slice
tileReturns :: [(VName, SubExp)] -> [(SubExp, SubExp)] -> VName -> Binder GPU KernelResult
tileReturns :: [(VName, SubExp)]
-> [(SubExp, SubExp)]
-> VName
-> BinderT GPU (State VNameSource) KernelResult
tileReturns [(VName, SubExp)]
dims_on_top [(SubExp, SubExp)]
dims VName
arr = do
let unit_dims :: Result
unit_dims = Int -> SubExp -> Result
forall a. Int -> a -> [a]
replicate ([(VName, SubExp)] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length [(VName, SubExp)]
dims_on_top) (IntType -> Integer -> SubExp
intConst IntType
Int64 Integer
1)
Type
arr_t <- VName -> BinderT GPU (State VNameSource) Type
forall rep (m :: * -> *). HasScope rep m => VName -> m Type
lookupType VName
arr
VName
arr' <-
if [(VName, SubExp)] -> Bool
forall (t :: * -> *) a. Foldable t => t a -> Bool
null [(VName, SubExp)]
dims_on_top Bool -> Bool -> Bool
|| Result -> Bool
forall (t :: * -> *) a. Foldable t => t a -> Bool
null (Type -> Result
forall u. TypeBase Shape u -> Result
arrayDims Type
arr_t)
then VName -> BinderT GPU (State VNameSource) VName
forall (m :: * -> *) a. Monad m => a -> m a
return VName
arr
else do
let new_shape :: Result
new_shape = Result
unit_dims Result -> Result -> Result
forall a. [a] -> [a] -> [a]
++ Type -> Result
forall u. TypeBase Shape u -> Result
arrayDims Type
arr_t
String
-> Exp (Rep (BinderT GPU (State VNameSource)))
-> BinderT GPU (State VNameSource) VName
forall (m :: * -> *).
MonadBinder m =>
String -> Exp (Rep m) -> m VName
letExp (VName -> String
baseString VName
arr) (Exp (Rep (BinderT GPU (State VNameSource)))
-> BinderT GPU (State VNameSource) VName)
-> Exp (Rep (BinderT GPU (State VNameSource)))
-> BinderT GPU (State VNameSource) VName
forall a b. (a -> b) -> a -> b
$ BasicOp -> ExpT GPU
forall rep. BasicOp -> ExpT rep
BasicOp (BasicOp -> ExpT GPU) -> BasicOp -> ExpT GPU
forall a b. (a -> b) -> a -> b
$ ShapeChange SubExp -> VName -> BasicOp
Reshape ((SubExp -> DimChange SubExp) -> Result -> ShapeChange SubExp
forall a b. (a -> b) -> [a] -> [b]
map SubExp -> DimChange SubExp
forall d. d -> DimChange d
DimNew Result
new_shape) VName
arr
let tile_dims :: [(SubExp, SubExp)]
tile_dims = Result -> Result -> [(SubExp, SubExp)]
forall a b. [a] -> [b] -> [(a, b)]
zip (((VName, SubExp) -> SubExp) -> [(VName, SubExp)] -> Result
forall a b. (a -> b) -> [a] -> [b]
map (VName, SubExp) -> SubExp
forall a b. (a, b) -> b
snd [(VName, SubExp)]
dims_on_top) Result
unit_dims [(SubExp, SubExp)] -> [(SubExp, SubExp)] -> [(SubExp, SubExp)]
forall a. [a] -> [a] -> [a]
++ [(SubExp, SubExp)]
dims
KernelResult -> BinderT GPU (State VNameSource) KernelResult
forall (m :: * -> *) a. Monad m => a -> m a
return (KernelResult -> BinderT GPU (State VNameSource) KernelResult)
-> KernelResult -> BinderT GPU (State VNameSource) KernelResult
forall a b. (a -> b) -> a -> b
$ [(SubExp, SubExp)] -> VName -> KernelResult
TileReturns [(SubExp, SubExp)]
tile_dims VName
arr'
is1DTileable :: VName -> M.Map VName Names -> VName -> InputArray
is1DTileable :: VName -> AliasTable -> VName -> InputArray
is1DTileable VName
gtid AliasTable
variance VName
arr
| Bool -> Bool
not (Bool -> Bool) -> Bool -> Bool
forall a b. (a -> b) -> a -> b
$ VName -> Names -> Bool
nameIn VName
gtid (Names -> Bool) -> Names -> Bool
forall a b. (a -> b) -> a -> b
$ Names -> VName -> AliasTable -> Names
forall k a. Ord k => a -> k -> Map k a -> a
M.findWithDefault Names
forall a. Monoid a => a
mempty VName
arr AliasTable
variance =
[Int] -> VName -> InputArray
InputTile [Int
0] VName
arr
| Bool
otherwise =
VName -> InputArray
InputDontTile VName
arr
reconstructGtids1D ::
Count GroupSize SubExp ->
VName ->
VName ->
VName ->
Binder GPU ()
reconstructGtids1D :: Count GroupSize SubExp
-> VName -> VName -> VName -> BinderT GPU (State VNameSource) ()
reconstructGtids1D Count GroupSize SubExp
group_size VName
gtid VName
gid VName
ltid =
[VName]
-> Exp (Rep (BinderT GPU (State VNameSource)))
-> BinderT GPU (State VNameSource) ()
forall (m :: * -> *).
MonadBinder m =>
[VName] -> Exp (Rep m) -> m ()
letBindNames [VName
gtid]
(ExpT GPU -> BinderT GPU (State VNameSource) ())
-> BinderT GPU (State VNameSource) (ExpT GPU)
-> BinderT GPU (State VNameSource) ()
forall (m :: * -> *) a b. Monad m => (a -> m b) -> m a -> m b
=<< TPrimExp Int64 VName
-> BinderT
GPU
(State VNameSource)
(Exp (Rep (BinderT GPU (State VNameSource))))
forall a (m :: * -> *).
(ToExp a, MonadBinder m) =>
a -> m (Exp (Rep m))
toExp (VName -> TPrimExp Int64 VName
forall a. a -> TPrimExp Int64 a
le64 VName
gid TPrimExp Int64 VName
-> TPrimExp Int64 VName -> TPrimExp Int64 VName
forall a. Num a => a -> a -> a
* SubExp -> TPrimExp Int64 VName
pe64 (Count GroupSize SubExp -> SubExp
forall u e. Count u e -> e
unCount Count GroupSize SubExp
group_size) TPrimExp Int64 VName
-> TPrimExp Int64 VName -> TPrimExp Int64 VName
forall a. Num a => a -> a -> a
+ VName -> TPrimExp Int64 VName
forall a. a -> TPrimExp Int64 a
le64 VName
ltid)
readTile1D ::
SubExp ->
VName ->
VName ->
Count NumGroups SubExp ->
Count GroupSize SubExp ->
TileKind ->
PrivStms ->
SubExp ->
[InputArray] ->
Binder GPU [InputTile]
readTile1D :: SubExp
-> VName
-> VName
-> Count NumGroups SubExp
-> Count GroupSize SubExp
-> TileKind
-> PrivStms
-> SubExp
-> [InputArray]
-> Binder GPU [InputTile]
readTile1D SubExp
tile_size VName
gid VName
gtid Count NumGroups SubExp
num_groups Count GroupSize SubExp
group_size TileKind
kind PrivStms
privstms SubExp
tile_id [InputArray]
inputs =
([VName] -> [InputTile])
-> BinderT GPU (State VNameSource) [VName]
-> Binder GPU [InputTile]
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap ([InputArray] -> [VName] -> [InputTile]
inputsToTiles [InputArray]
inputs)
(BinderT GPU (State VNameSource) [VName] -> Binder GPU [InputTile])
-> ((VName -> BinderT GPU (State VNameSource) Result)
-> BinderT GPU (State VNameSource) [VName])
-> (VName -> BinderT GPU (State VNameSource) Result)
-> Binder GPU [InputTile]
forall b c a. (b -> c) -> (a -> b) -> a -> c
. String
-> SegLevel
-> ResultManifest
-> (VName -> BinderT GPU (State VNameSource) Result)
-> BinderT GPU (State VNameSource) [VName]
segMap1D String
"full_tile" SegLevel
lvl ResultManifest
ResultNoSimplify
((VName -> BinderT GPU (State VNameSource) Result)
-> Binder GPU [InputTile])
-> (VName -> BinderT GPU (State VNameSource) Result)
-> Binder GPU [InputTile]
forall a b. (a -> b) -> a -> b
$ \VName
ltid -> do
SubExp
j <-
String
-> Exp (Rep (BinderT GPU (State VNameSource))) -> Binder GPU SubExp
forall (m :: * -> *).
MonadBinder m =>
String -> Exp (Rep m) -> m SubExp
letSubExp String
"j"
(ExpT GPU -> Binder GPU SubExp)
-> BinderT GPU (State VNameSource) (ExpT GPU) -> Binder GPU SubExp
forall (m :: * -> *) a b. Monad m => (a -> m b) -> m a -> m b
=<< TPrimExp Int64 VName
-> BinderT
GPU
(State VNameSource)
(Exp (Rep (BinderT GPU (State VNameSource))))
forall a (m :: * -> *).
(ToExp a, MonadBinder m) =>
a -> m (Exp (Rep m))
toExp (SubExp -> TPrimExp Int64 VName
pe64 SubExp
tile_id TPrimExp Int64 VName
-> TPrimExp Int64 VName -> TPrimExp Int64 VName
forall a. Num a => a -> a -> a
* SubExp -> TPrimExp Int64 VName
pe64 SubExp
tile_size TPrimExp Int64 VName
-> TPrimExp Int64 VName -> TPrimExp Int64 VName
forall a. Num a => a -> a -> a
+ VName -> TPrimExp Int64 VName
forall a. a -> TPrimExp Int64 a
le64 VName
ltid)
Count GroupSize SubExp
-> VName -> VName -> VName -> BinderT GPU (State VNameSource) ()
reconstructGtids1D Count GroupSize SubExp
group_size VName
gtid VName
gid VName
ltid
Slice SubExp -> PrivStms -> BinderT GPU (State VNameSource) ()
addPrivStms [SubExp -> DimIndex SubExp
forall d. d -> DimIndex d
DimFix (SubExp -> DimIndex SubExp) -> SubExp -> DimIndex SubExp
forall a b. (a -> b) -> a -> b
$ VName -> SubExp
Var VName
ltid] PrivStms
privstms
let arrs :: [VName]
arrs = ((VName, [Int]) -> VName) -> [(VName, [Int])] -> [VName]
forall a b. (a -> b) -> [a] -> [b]
map (VName, [Int]) -> VName
forall a b. (a, b) -> a
fst ([(VName, [Int])] -> [VName]) -> [(VName, [Int])] -> [VName]
forall a b. (a -> b) -> a -> b
$ [InputArray] -> [(VName, [Int])]
tiledInputs [InputArray]
inputs
[Type]
arr_ts <- (VName -> BinderT GPU (State VNameSource) Type)
-> [VName] -> BinderT GPU (State VNameSource) [Type]
forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
(a -> m b) -> t a -> m (t b)
mapM VName -> BinderT GPU (State VNameSource) Type
forall rep (m :: * -> *). HasScope rep m => VName -> m Type
lookupType [VName]
arrs
let tile_ts :: [Type]
tile_ts = (Type -> Type) -> [Type] -> [Type]
forall a b. (a -> b) -> [a] -> [b]
map Type -> Type
forall shape u.
ArrayShape shape =>
TypeBase shape u -> TypeBase shape u
rowType [Type]
arr_ts
w :: SubExp
w = Int -> [Type] -> SubExp
forall u. Int -> [TypeBase Shape u] -> SubExp
arraysSize Int
0 [Type]
arr_ts
let readTileElem :: VName -> BinderT GPU (State VNameSource) VName
readTileElem VName
arr =
String
-> Exp (Rep (BinderT GPU (State VNameSource)))
-> BinderT GPU (State VNameSource) VName
forall (m :: * -> *).
MonadBinder m =>
String -> Exp (Rep m) -> m VName
letExp String
"tile_elem" (BasicOp -> ExpT GPU
forall rep. BasicOp -> ExpT rep
BasicOp (BasicOp -> ExpT GPU) -> BasicOp -> ExpT GPU
forall a b. (a -> b) -> a -> b
$ VName -> Slice SubExp -> BasicOp
Index VName
arr [SubExp -> DimIndex SubExp
forall d. d -> DimIndex d
DimFix SubExp
j])
([VName] -> Result)
-> BinderT GPU (State VNameSource) [VName]
-> BinderT GPU (State VNameSource) Result
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap ((VName -> SubExp) -> [VName] -> Result
forall a b. (a -> b) -> [a] -> [b]
map VName -> SubExp
Var) (BinderT GPU (State VNameSource) [VName]
-> BinderT GPU (State VNameSource) Result)
-> BinderT GPU (State VNameSource) [VName]
-> BinderT GPU (State VNameSource) Result
forall a b. (a -> b) -> a -> b
$
case TileKind
kind of
TileKind
TilePartial ->
String
-> Exp (Rep (BinderT GPU (State VNameSource)))
-> BinderT GPU (State VNameSource) [VName]
forall (m :: * -> *).
MonadBinder m =>
String -> Exp (Rep m) -> m [VName]
letTupExp String
"pre"
(ExpT GPU -> BinderT GPU (State VNameSource) [VName])
-> BinderT GPU (State VNameSource) (ExpT GPU)
-> BinderT GPU (State VNameSource) [VName]
forall (m :: * -> *) a b. Monad m => (a -> m b) -> m a -> m b
=<< BinderT
GPU
(State VNameSource)
(Exp (Rep (BinderT GPU (State VNameSource))))
-> BinderT
GPU
(State VNameSource)
(Body (Rep (BinderT GPU (State VNameSource))))
-> BinderT
GPU
(State VNameSource)
(Body (Rep (BinderT GPU (State VNameSource))))
-> BinderT
GPU
(State VNameSource)
(Exp (Rep (BinderT GPU (State VNameSource))))
forall (m :: * -> *).
(MonadBinder 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
-> BinderT
GPU
(State VNameSource)
(Exp (Rep (BinderT GPU (State VNameSource))))
forall a (m :: * -> *).
(ToExp a, MonadBinder m) =>
a -> m (Exp (Rep m))
toExp (TPrimExp Bool VName
-> BinderT
GPU
(State VNameSource)
(Exp (Rep (BinderT GPU (State VNameSource)))))
-> TPrimExp Bool VName
-> BinderT
GPU
(State VNameSource)
(Exp (Rep (BinderT GPU (State VNameSource))))
forall a b. (a -> b) -> a -> b
$ SubExp -> TPrimExp Int64 VName
pe64 SubExp
j TPrimExp Int64 VName -> TPrimExp Int64 VName -> TPrimExp Bool VName
forall t v. TPrimExp t v -> TPrimExp t v -> TPrimExp Bool v
.<. SubExp -> TPrimExp Int64 VName
pe64 SubExp
w)
(Result -> Body GPU
forall rep. Bindable rep => Result -> Body rep
resultBody (Result -> Body GPU)
-> BinderT GPU (State VNameSource) Result
-> BinderT GPU (State VNameSource) (Body GPU)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> (VName -> Binder GPU SubExp)
-> [VName] -> BinderT GPU (State VNameSource) Result
forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
(a -> m b) -> t a -> m (t b)
mapM ((VName -> SubExp)
-> BinderT GPU (State VNameSource) VName -> Binder GPU SubExp
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap VName -> SubExp
Var (BinderT GPU (State VNameSource) VName -> Binder GPU SubExp)
-> (VName -> BinderT GPU (State VNameSource) VName)
-> VName
-> Binder GPU SubExp
forall b c a. (b -> c) -> (a -> b) -> a -> c
. VName -> BinderT GPU (State VNameSource) VName
readTileElem) [VName]
arrs)
([BinderT
GPU
(State VNameSource)
(Exp (Rep (BinderT GPU (State VNameSource))))]
-> BinderT
GPU
(State VNameSource)
(Body (Rep (BinderT GPU (State VNameSource))))
forall (m :: * -> *).
MonadBinder m =>
[m (Exp (Rep m))] -> m (Body (Rep m))
eBody ([BinderT
GPU
(State VNameSource)
(Exp (Rep (BinderT GPU (State VNameSource))))]
-> BinderT
GPU
(State VNameSource)
(Body (Rep (BinderT GPU (State VNameSource)))))
-> [BinderT
GPU
(State VNameSource)
(Exp (Rep (BinderT GPU (State VNameSource))))]
-> BinderT
GPU
(State VNameSource)
(Body (Rep (BinderT GPU (State VNameSource))))
forall a b. (a -> b) -> a -> b
$ (Type -> BinderT GPU (State VNameSource) (ExpT GPU))
-> [Type] -> [BinderT GPU (State VNameSource) (ExpT GPU)]
forall a b. (a -> b) -> [a] -> [b]
map Type -> BinderT GPU (State VNameSource) (ExpT GPU)
forall (m :: * -> *). MonadBinder m => Type -> m (Exp (Rep m))
eBlank [Type]
tile_ts)
TileKind
TileFull ->
(VName -> BinderT GPU (State VNameSource) VName)
-> [VName] -> BinderT GPU (State VNameSource) [VName]
forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
(a -> m b) -> t a -> m (t b)
mapM VName -> BinderT GPU (State VNameSource) VName
readTileElem [VName]
arrs
where
lvl :: SegLevel
lvl = Count NumGroups SubExp
-> Count GroupSize SubExp -> SegVirt -> SegLevel
SegThread Count NumGroups SubExp
num_groups Count GroupSize SubExp
group_size SegVirt
SegNoVirt
processTile1D ::
VName ->
VName ->
SubExp ->
SubExp ->
Count NumGroups SubExp ->
Count GroupSize SubExp ->
ProcessTileArgs ->
Binder GPU [VName]
processTile1D :: VName
-> VName
-> SubExp
-> SubExp
-> Count NumGroups SubExp
-> Count GroupSize SubExp
-> ProcessTileArgs
-> BinderT GPU (State VNameSource) [VName]
processTile1D VName
gid VName
gtid SubExp
kdim SubExp
tile_size Count NumGroups SubExp
num_groups Count GroupSize SubExp
group_size ProcessTileArgs
tile_args = do
let red_comm :: Commutativity
red_comm = ProcessTileArgs -> Commutativity
processComm ProcessTileArgs
tile_args
privstms :: PrivStms
privstms = ProcessTileArgs -> PrivStms
processPrivStms ProcessTileArgs
tile_args
map_lam :: Lambda GPU
map_lam = ProcessTileArgs -> Lambda GPU
processMapLam ProcessTileArgs
tile_args
red_lam :: Lambda GPU
red_lam = ProcessTileArgs -> Lambda GPU
processRedLam ProcessTileArgs
tile_args
tiles :: [InputTile]
tiles = ProcessTileArgs -> [InputTile]
processTiles ProcessTileArgs
tile_args
tile_id :: SubExp
tile_id = ProcessTileArgs -> SubExp
processTileId ProcessTileArgs
tile_args
accs :: [VName]
accs = ProcessTileArgs -> [VName]
processAcc ProcessTileArgs
tile_args
String
-> SegLevel
-> ResultManifest
-> (VName -> BinderT GPU (State VNameSource) Result)
-> BinderT GPU (State VNameSource) [VName]
segMap1D String
"acc" SegLevel
lvl ResultManifest
ResultPrivate ((VName -> BinderT GPU (State VNameSource) Result)
-> BinderT GPU (State VNameSource) [VName])
-> (VName -> BinderT GPU (State VNameSource) Result)
-> BinderT GPU (State VNameSource) [VName]
forall a b. (a -> b) -> a -> b
$ \VName
ltid -> do
Count GroupSize SubExp
-> VName -> VName -> VName -> BinderT GPU (State VNameSource) ()
reconstructGtids1D Count GroupSize SubExp
group_size VName
gtid VName
gid VName
ltid
Slice SubExp -> PrivStms -> BinderT GPU (State VNameSource) ()
addPrivStms [SubExp -> DimIndex SubExp
forall d. d -> DimIndex d
DimFix (SubExp -> DimIndex SubExp) -> SubExp -> DimIndex SubExp
forall a b. (a -> b) -> a -> b
$ VName -> SubExp
Var VName
ltid] PrivStms
privstms
Result
thread_accs <- [VName]
-> (VName -> Binder GPU SubExp)
-> BinderT GPU (State VNameSource) Result
forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
t a -> (a -> m b) -> m (t b)
forM [VName]
accs ((VName -> Binder GPU SubExp)
-> BinderT GPU (State VNameSource) Result)
-> (VName -> Binder GPU SubExp)
-> BinderT GPU (State VNameSource) Result
forall a b. (a -> b) -> a -> b
$ \VName
acc ->
String
-> Exp (Rep (BinderT GPU (State VNameSource))) -> Binder GPU SubExp
forall (m :: * -> *).
MonadBinder m =>
String -> Exp (Rep m) -> m SubExp
letSubExp String
"acc" (Exp (Rep (BinderT GPU (State VNameSource))) -> Binder GPU SubExp)
-> Exp (Rep (BinderT GPU (State VNameSource))) -> Binder GPU SubExp
forall a b. (a -> b) -> a -> b
$ BasicOp -> ExpT GPU
forall rep. BasicOp -> ExpT rep
BasicOp (BasicOp -> ExpT GPU) -> BasicOp -> ExpT GPU
forall a b. (a -> b) -> a -> b
$ VName -> Slice SubExp -> BasicOp
Index VName
acc [SubExp -> DimIndex SubExp
forall d. d -> DimIndex d
DimFix (SubExp -> DimIndex SubExp) -> SubExp -> DimIndex SubExp
forall a b. (a -> b) -> a -> b
$ VName -> SubExp
Var VName
ltid]
let sliceTile :: InputTile -> BinderT GPU (State VNameSource) VName
sliceTile (InputTiled [Int]
_ VName
arr) =
VName -> BinderT GPU (State VNameSource) VName
forall (f :: * -> *) a. Applicative f => a -> f a
pure VName
arr
sliceTile (InputUntiled VName
arr) =
VName
-> SubExp
-> SubExp
-> SubExp
-> BinderT GPU (State VNameSource) VName
forall (m :: * -> *).
MonadBinder m =>
VName -> SubExp -> SubExp -> SubExp -> m VName
sliceUntiled VName
arr SubExp
tile_id SubExp
tile_size SubExp
tile_size
[VName]
tiles' <- (InputTile -> BinderT GPU (State VNameSource) VName)
-> [InputTile] -> BinderT GPU (State VNameSource) [VName]
forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
(a -> m b) -> t a -> m (t b)
mapM InputTile -> BinderT GPU (State VNameSource) VName
sliceTile [InputTile]
tiles
let form' :: ScremaForm GPU
form' = [Reduce GPU] -> Lambda GPU -> ScremaForm GPU
forall rep. [Reduce rep] -> Lambda rep -> ScremaForm rep
redomapSOAC [Commutativity -> Lambda GPU -> Result -> Reduce GPU
forall rep. Commutativity -> Lambda rep -> Result -> Reduce rep
Reduce Commutativity
red_comm Lambda GPU
red_lam Result
thread_accs] Lambda GPU
map_lam
([VName] -> Result)
-> BinderT GPU (State VNameSource) [VName]
-> BinderT GPU (State VNameSource) Result
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap ((VName -> SubExp) -> [VName] -> Result
forall a b. (a -> b) -> [a] -> [b]
map VName -> SubExp
Var) (BinderT GPU (State VNameSource) [VName]
-> BinderT GPU (State VNameSource) Result)
-> BinderT GPU (State VNameSource) [VName]
-> BinderT GPU (State VNameSource) Result
forall a b. (a -> b) -> a -> b
$
String
-> Exp (Rep (BinderT GPU (State VNameSource)))
-> BinderT GPU (State VNameSource) [VName]
forall (m :: * -> *).
MonadBinder m =>
String -> Exp (Rep m) -> m [VName]
letTupExp String
"acc"
(ExpT GPU -> BinderT GPU (State VNameSource) [VName])
-> BinderT GPU (State VNameSource) (ExpT GPU)
-> BinderT GPU (State VNameSource) [VName]
forall (m :: * -> *) a b. Monad m => (a -> m b) -> m a -> m b
=<< BinderT
GPU
(State VNameSource)
(Exp (Rep (BinderT GPU (State VNameSource))))
-> BinderT
GPU
(State VNameSource)
(Body (Rep (BinderT GPU (State VNameSource))))
-> BinderT
GPU
(State VNameSource)
(Body (Rep (BinderT GPU (State VNameSource))))
-> BinderT
GPU
(State VNameSource)
(Exp (Rep (BinderT GPU (State VNameSource))))
forall (m :: * -> *).
(MonadBinder 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
-> BinderT
GPU
(State VNameSource)
(Exp (Rep (BinderT GPU (State VNameSource))))
forall a (m :: * -> *).
(ToExp a, MonadBinder m) =>
a -> m (Exp (Rep m))
toExp (TPrimExp Bool VName
-> BinderT
GPU
(State VNameSource)
(Exp (Rep (BinderT GPU (State VNameSource)))))
-> TPrimExp Bool VName
-> BinderT
GPU
(State VNameSource)
(Exp (Rep (BinderT GPU (State VNameSource))))
forall a b. (a -> b) -> a -> b
$ VName -> TPrimExp Int64 VName
forall a. a -> TPrimExp Int64 a
le64 VName
gtid TPrimExp Int64 VName -> TPrimExp Int64 VName -> TPrimExp Bool VName
forall t v. TPrimExp t v -> TPrimExp t v -> TPrimExp Bool v
.<. SubExp -> TPrimExp Int64 VName
pe64 SubExp
kdim)
([BinderT
GPU
(State VNameSource)
(Exp (Rep (BinderT GPU (State VNameSource))))]
-> BinderT
GPU
(State VNameSource)
(Body (Rep (BinderT GPU (State VNameSource))))
forall (m :: * -> *).
MonadBinder m =>
[m (Exp (Rep m))] -> m (Body (Rep m))
eBody [ExpT GPU -> BinderT GPU (State VNameSource) (ExpT GPU)
forall (f :: * -> *) a. Applicative f => a -> f a
pure (ExpT GPU -> BinderT GPU (State VNameSource) (ExpT GPU))
-> ExpT GPU -> BinderT GPU (State VNameSource) (ExpT GPU)
forall a b. (a -> b) -> a -> b
$ Op GPU -> ExpT GPU
forall rep. Op rep -> ExpT rep
Op (Op GPU -> ExpT GPU) -> Op GPU -> ExpT GPU
forall a b. (a -> b) -> a -> b
$ SOAC GPU -> HostOp GPU (SOAC GPU)
forall rep op. op -> HostOp rep op
OtherOp (SOAC GPU -> HostOp GPU (SOAC GPU))
-> SOAC GPU -> HostOp GPU (SOAC GPU)
forall a b. (a -> b) -> a -> b
$ SubExp -> [VName] -> ScremaForm GPU -> SOAC GPU
forall rep. SubExp -> [VName] -> ScremaForm rep -> SOAC rep
Screma SubExp
tile_size [VName]
tiles' ScremaForm GPU
form'])
(Result
-> BinderT
GPU
(State VNameSource)
(Body (Rep (BinderT GPU (State VNameSource))))
forall (m :: * -> *). MonadBinder m => Result -> m (Body (Rep m))
resultBodyM Result
thread_accs)
where
lvl :: SegLevel
lvl = Count NumGroups SubExp
-> Count GroupSize SubExp -> SegVirt -> SegLevel
SegThread Count NumGroups SubExp
num_groups Count GroupSize SubExp
group_size SegVirt
SegNoVirt
processResidualTile1D ::
VName ->
VName ->
SubExp ->
SubExp ->
Count NumGroups SubExp ->
Count GroupSize SubExp ->
ResidualTileArgs ->
Binder GPU [VName]
processResidualTile1D :: VName
-> VName
-> SubExp
-> SubExp
-> Count NumGroups SubExp
-> Count GroupSize SubExp
-> ResidualTileArgs
-> BinderT GPU (State VNameSource) [VName]
processResidualTile1D VName
gid VName
gtid SubExp
kdim SubExp
tile_size Count NumGroups SubExp
num_groups Count GroupSize SubExp
group_size ResidualTileArgs
args = do
SubExp
residual_input <-
String
-> Exp (Rep (BinderT GPU (State VNameSource))) -> Binder GPU SubExp
forall (m :: * -> *).
MonadBinder m =>
String -> Exp (Rep m) -> m SubExp
letSubExp String
"residual_input" (Exp (Rep (BinderT GPU (State VNameSource))) -> Binder GPU SubExp)
-> Exp (Rep (BinderT GPU (State VNameSource))) -> Binder GPU SubExp
forall a b. (a -> b) -> a -> b
$
BasicOp -> ExpT GPU
forall rep. BasicOp -> ExpT rep
BasicOp (BasicOp -> ExpT GPU) -> BasicOp -> ExpT GPU
forall a b. (a -> b) -> a -> b
$ BinOp -> SubExp -> SubExp -> BasicOp
BinOp (IntType -> Safety -> BinOp
SRem IntType
Int64 Safety
Unsafe) SubExp
w SubExp
tile_size
String
-> Exp (Rep (BinderT GPU (State VNameSource)))
-> BinderT GPU (State VNameSource) [VName]
forall (m :: * -> *).
MonadBinder m =>
String -> Exp (Rep m) -> m [VName]
letTupExp String
"acc_after_residual"
(ExpT GPU -> BinderT GPU (State VNameSource) [VName])
-> BinderT GPU (State VNameSource) (ExpT GPU)
-> BinderT GPU (State VNameSource) [VName]
forall (m :: * -> *) a b. Monad m => (a -> m b) -> m a -> m b
=<< BinderT
GPU
(State VNameSource)
(Exp (Rep (BinderT GPU (State VNameSource))))
-> BinderT
GPU
(State VNameSource)
(Body (Rep (BinderT GPU (State VNameSource))))
-> BinderT
GPU
(State VNameSource)
(Body (Rep (BinderT GPU (State VNameSource))))
-> BinderT
GPU
(State VNameSource)
(Exp (Rep (BinderT GPU (State VNameSource))))
forall (m :: * -> *).
(MonadBinder 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
-> BinderT
GPU
(State VNameSource)
(Exp (Rep (BinderT GPU (State VNameSource))))
forall a (m :: * -> *).
(ToExp a, MonadBinder m) =>
a -> m (Exp (Rep m))
toExp (TPrimExp Bool VName
-> BinderT
GPU
(State VNameSource)
(Exp (Rep (BinderT GPU (State VNameSource)))))
-> TPrimExp Bool VName
-> BinderT
GPU
(State VNameSource)
(Exp (Rep (BinderT GPU (State VNameSource))))
forall a b. (a -> b) -> a -> b
$ SubExp -> TPrimExp Int64 VName
pe64 SubExp
residual_input TPrimExp Int64 VName -> TPrimExp Int64 VName -> TPrimExp Bool VName
forall t v. TPrimExp t v -> TPrimExp t v -> TPrimExp Bool v
.==. TPrimExp Int64 VName
0)
(Result
-> BinderT
GPU
(State VNameSource)
(Body (Rep (BinderT GPU (State VNameSource))))
forall (m :: * -> *). MonadBinder m => Result -> m (Body (Rep m))
resultBodyM (Result
-> BinderT
GPU
(State VNameSource)
(Body (Rep (BinderT GPU (State VNameSource)))))
-> Result
-> BinderT
GPU
(State VNameSource)
(Body (Rep (BinderT GPU (State VNameSource))))
forall a b. (a -> b) -> a -> b
$ (VName -> SubExp) -> [VName] -> Result
forall a b. (a -> b) -> [a] -> [b]
map VName -> SubExp
Var [VName]
accs)
(SubExp -> BinderT GPU (State VNameSource) (Body GPU)
nonemptyTile SubExp
residual_input)
where
red_comm :: Commutativity
red_comm = ResidualTileArgs -> Commutativity
residualComm ResidualTileArgs
args
map_lam :: Lambda GPU
map_lam = ResidualTileArgs -> Lambda GPU
residualMapLam ResidualTileArgs
args
red_lam :: Lambda GPU
red_lam = ResidualTileArgs -> Lambda GPU
residualRedLam ResidualTileArgs
args
privstms :: PrivStms
privstms = ResidualTileArgs -> PrivStms
residualPrivStms ResidualTileArgs
args
inputs :: [InputArray]
inputs = ResidualTileArgs -> [InputArray]
residualInput ResidualTileArgs
args
accs :: [VName]
accs = ResidualTileArgs -> [VName]
residualAcc ResidualTileArgs
args
num_whole_tiles :: SubExp
num_whole_tiles = ResidualTileArgs -> SubExp
residualNumWholeTiles ResidualTileArgs
args
w :: SubExp
w = ResidualTileArgs -> SubExp
residualInputSize ResidualTileArgs
args
nonemptyTile :: SubExp -> BinderT GPU (State VNameSource) (Body GPU)
nonemptyTile SubExp
residual_input = BinderT GPU (State VNameSource) (Body GPU)
-> BinderT GPU (State VNameSource) (Body GPU)
forall rep (m :: * -> *) somerep.
(Bindable rep, MonadFreshNames m, HasScope somerep m,
SameScope somerep rep) =>
Binder rep (Body rep) -> m (Body rep)
runBodyBinder (BinderT GPU (State VNameSource) (Body GPU)
-> BinderT GPU (State VNameSource) (Body GPU))
-> BinderT GPU (State VNameSource) (Body GPU)
-> BinderT GPU (State VNameSource) (Body GPU)
forall a b. (a -> b) -> a -> b
$ do
[InputTile]
full_tiles <-
SubExp
-> VName
-> VName
-> Count NumGroups SubExp
-> Count GroupSize SubExp
-> TileKind
-> PrivStms
-> SubExp
-> [InputArray]
-> Binder GPU [InputTile]
readTile1D
SubExp
tile_size
VName
gid
VName
gtid
Count NumGroups SubExp
num_groups
Count GroupSize SubExp
group_size
TileKind
TilePartial
PrivStms
privstms
SubExp
num_whole_tiles
[InputArray]
inputs
let sliceTile :: InputTile -> BinderT GPU (State VNameSource) InputTile
sliceTile (InputUntiled VName
arr) =
InputTile -> BinderT GPU (State VNameSource) InputTile
forall (f :: * -> *) a. Applicative f => a -> f a
pure (InputTile -> BinderT GPU (State VNameSource) InputTile)
-> InputTile -> BinderT GPU (State VNameSource) InputTile
forall a b. (a -> b) -> a -> b
$ VName -> InputTile
InputUntiled VName
arr
sliceTile (InputTiled [Int]
perm VName
tile) = do
let slice :: DimIndex SubExp
slice =
SubExp -> SubExp -> SubExp -> DimIndex SubExp
forall d. d -> d -> d -> DimIndex d
DimSlice (IntType -> Integer -> SubExp
intConst IntType
Int64 Integer
0) SubExp
residual_input (IntType -> Integer -> SubExp
intConst IntType
Int64 Integer
1)
[Int] -> VName -> InputTile
InputTiled [Int]
perm
(VName -> InputTile)
-> BinderT GPU (State VNameSource) VName
-> BinderT GPU (State VNameSource) InputTile
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> String
-> Exp (Rep (BinderT GPU (State VNameSource)))
-> BinderT GPU (State VNameSource) VName
forall (m :: * -> *).
MonadBinder m =>
String -> Exp (Rep m) -> m VName
letExp String
"partial_tile" (BasicOp -> ExpT GPU
forall rep. BasicOp -> ExpT rep
BasicOp (BasicOp -> ExpT GPU) -> BasicOp -> ExpT GPU
forall a b. (a -> b) -> a -> b
$ VName -> Slice SubExp -> BasicOp
Index VName
tile [DimIndex SubExp
slice])
[InputTile]
tiles <- (InputTile -> BinderT GPU (State VNameSource) InputTile)
-> [InputTile] -> Binder GPU [InputTile]
forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
(a -> m b) -> t a -> m (t b)
mapM InputTile -> BinderT GPU (State VNameSource) InputTile
sliceTile [InputTile]
full_tiles
let tile_args :: ProcessTileArgs
tile_args =
PrivStms
-> Commutativity
-> Lambda GPU
-> Lambda GPU
-> [InputTile]
-> [VName]
-> SubExp
-> ProcessTileArgs
ProcessTileArgs PrivStms
privstms Commutativity
red_comm Lambda GPU
red_lam Lambda GPU
map_lam [InputTile]
tiles [VName]
accs SubExp
num_whole_tiles
Result -> Body GPU
forall rep. Bindable rep => Result -> Body rep
resultBody (Result -> Body GPU) -> ([VName] -> Result) -> [VName] -> Body GPU
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (VName -> SubExp) -> [VName] -> Result
forall a b. (a -> b) -> [a] -> [b]
map VName -> SubExp
Var
([VName] -> Body GPU)
-> BinderT GPU (State VNameSource) [VName]
-> BinderT GPU (State VNameSource) (Body GPU)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> VName
-> VName
-> SubExp
-> SubExp
-> Count NumGroups SubExp
-> Count GroupSize SubExp
-> ProcessTileArgs
-> BinderT GPU (State VNameSource) [VName]
processTile1D VName
gid VName
gtid SubExp
kdim SubExp
residual_input Count NumGroups SubExp
num_groups Count GroupSize SubExp
group_size ProcessTileArgs
tile_args
tiling1d :: [(VName, SubExp)] -> DoTiling VName SubExp
tiling1d :: [(VName, SubExp)] -> DoTiling VName SubExp
tiling1d [(VName, SubExp)]
dims_on_top SegLevel
initial_lvl VName
gtid SubExp
kdim SubExp
w = do
VName
gid <- String -> BinderT GPU (State VNameSource) VName
forall (m :: * -> *). MonadFreshNames m => String -> m VName
newVName String
"gid"
VName
gid_flat <- String -> BinderT GPU (State VNameSource) VName
forall (m :: * -> *). MonadFreshNames m => String -> m VName
newVName String
"gid_flat"
(SegLevel
lvl, SegSpace
space) <-
if [(VName, SubExp)] -> Bool
forall (t :: * -> *) a. Foldable t => t a -> Bool
null [(VName, SubExp)]
dims_on_top
then
(SegLevel, SegSpace)
-> BinderT GPU (State VNameSource) (SegLevel, SegSpace)
forall (m :: * -> *) a. Monad m => a -> m a
return
( Count NumGroups SubExp
-> Count GroupSize SubExp -> SegVirt -> SegLevel
SegGroup (SegLevel -> Count NumGroups SubExp
segNumGroups SegLevel
initial_lvl) (SegLevel -> Count GroupSize SubExp
segGroupSize SegLevel
initial_lvl) (SegVirt -> SegLevel) -> SegVirt -> SegLevel
forall a b. (a -> b) -> a -> b
$ SegLevel -> SegVirt
segVirt SegLevel
initial_lvl,
VName -> [(VName, SubExp)] -> SegSpace
SegSpace VName
gid_flat [(VName
gid, Count NumGroups SubExp -> SubExp
forall u e. Count u e -> e
unCount (Count NumGroups SubExp -> SubExp)
-> Count NumGroups SubExp -> SubExp
forall a b. (a -> b) -> a -> b
$ SegLevel -> Count NumGroups SubExp
segNumGroups SegLevel
initial_lvl)]
)
else do
SubExp
group_size <-
String
-> Exp (Rep (BinderT GPU (State VNameSource))) -> Binder GPU SubExp
forall (m :: * -> *).
MonadBinder m =>
String -> Exp (Rep m) -> m SubExp
letSubExp String
"computed_group_size" (Exp (Rep (BinderT GPU (State VNameSource))) -> Binder GPU SubExp)
-> Exp (Rep (BinderT GPU (State VNameSource))) -> Binder GPU SubExp
forall a b. (a -> b) -> a -> b
$
BasicOp -> ExpT GPU
forall rep. BasicOp -> ExpT rep
BasicOp (BasicOp -> ExpT GPU) -> BasicOp -> ExpT GPU
forall a b. (a -> b) -> a -> b
$ BinOp -> SubExp -> SubExp -> BasicOp
BinOp (IntType -> BinOp
SMin IntType
Int64) (Count GroupSize SubExp -> SubExp
forall u e. Count u e -> e
unCount (SegLevel -> Count GroupSize SubExp
segGroupSize SegLevel
initial_lvl)) SubExp
kdim
SubExp
ldim <-
String
-> Exp (Rep (BinderT GPU (State VNameSource))) -> Binder GPU SubExp
forall (m :: * -> *).
MonadBinder m =>
String -> Exp (Rep m) -> m SubExp
letSubExp String
"ldim" (Exp (Rep (BinderT GPU (State VNameSource))) -> Binder GPU SubExp)
-> Exp (Rep (BinderT GPU (State VNameSource))) -> Binder GPU SubExp
forall a b. (a -> b) -> a -> b
$
BasicOp -> ExpT GPU
forall rep. BasicOp -> ExpT rep
BasicOp (BasicOp -> ExpT GPU) -> BasicOp -> ExpT GPU
forall a b. (a -> b) -> a -> b
$ BinOp -> SubExp -> SubExp -> BasicOp
BinOp (IntType -> Safety -> BinOp
SDivUp IntType
Int64 Safety
Unsafe) SubExp
kdim SubExp
group_size
SubExp
num_groups <-
String
-> Exp (Rep (BinderT GPU (State VNameSource))) -> Binder GPU SubExp
forall (m :: * -> *).
MonadBinder m =>
String -> Exp (Rep m) -> m SubExp
letSubExp String
"computed_num_groups"
(ExpT GPU -> Binder GPU SubExp)
-> BinderT GPU (State VNameSource) (ExpT GPU) -> Binder GPU SubExp
forall (m :: * -> *) a b. Monad m => (a -> m b) -> m a -> m b
=<< BinOp
-> SubExp
-> Result
-> BinderT
GPU
(State VNameSource)
(Exp (Rep (BinderT GPU (State VNameSource))))
forall (m :: * -> *).
MonadBinder m =>
BinOp -> SubExp -> Result -> m (Exp (Rep m))
foldBinOp (IntType -> Overflow -> BinOp
Mul IntType
Int64 Overflow
OverflowUndef) SubExp
ldim (((VName, SubExp) -> SubExp) -> [(VName, SubExp)] -> Result
forall a b. (a -> b) -> [a] -> [b]
map (VName, SubExp) -> SubExp
forall a b. (a, b) -> b
snd [(VName, SubExp)]
dims_on_top)
(SegLevel, SegSpace)
-> BinderT GPU (State VNameSource) (SegLevel, SegSpace)
forall (m :: * -> *) a. Monad m => a -> m a
return
( Count NumGroups SubExp
-> Count GroupSize SubExp -> SegVirt -> SegLevel
SegGroup (SubExp -> Count NumGroups SubExp
forall u e. e -> Count u e
Count SubExp
num_groups) (SubExp -> Count GroupSize SubExp
forall u e. e -> Count u e
Count SubExp
group_size) SegVirt
SegNoVirt,
VName -> [(VName, SubExp)] -> SegSpace
SegSpace VName
gid_flat ([(VName, SubExp)] -> SegSpace) -> [(VName, SubExp)] -> SegSpace
forall a b. (a -> b) -> a -> b
$ [(VName, SubExp)]
dims_on_top [(VName, SubExp)] -> [(VName, SubExp)] -> [(VName, SubExp)]
forall a. [a] -> [a] -> [a]
++ [(VName
gid, SubExp
ldim)]
)
let tile_size :: SubExp
tile_size = Count GroupSize SubExp -> SubExp
forall u e. Count u e -> e
unCount (Count GroupSize SubExp -> SubExp)
-> Count GroupSize SubExp -> SubExp
forall a b. (a -> b) -> a -> b
$ SegLevel -> Count GroupSize SubExp
segGroupSize SegLevel
lvl
Tiling -> Binder GPU Tiling
forall (m :: * -> *) a. Monad m => a -> m a
return
Tiling :: (String
-> SegLevel
-> ResultManifest
-> (PrimExp VName
-> Slice SubExp -> BinderT GPU (State VNameSource) Result)
-> BinderT GPU (State VNameSource) [VName])
-> (TileKind
-> PrivStms -> SubExp -> [InputArray] -> Binder GPU [InputTile])
-> (ProcessTileArgs -> BinderT GPU (State VNameSource) [VName])
-> (ResidualTileArgs -> BinderT GPU (State VNameSource) [VName])
-> (VName -> BinderT GPU (State VNameSource) KernelResult)
-> SegSpace
-> Shape
-> SegLevel
-> Binder GPU SubExp
-> Tiling
Tiling
{ tilingSegMap :: String
-> SegLevel
-> ResultManifest
-> (PrimExp VName
-> Slice SubExp -> BinderT GPU (State VNameSource) Result)
-> BinderT GPU (State VNameSource) [VName]
tilingSegMap = \String
desc SegLevel
lvl' ResultManifest
manifest PrimExp VName
-> Slice SubExp -> BinderT GPU (State VNameSource) Result
f -> String
-> SegLevel
-> ResultManifest
-> (VName -> BinderT GPU (State VNameSource) Result)
-> BinderT GPU (State VNameSource) [VName]
segMap1D String
desc SegLevel
lvl' ResultManifest
manifest ((VName -> BinderT GPU (State VNameSource) Result)
-> BinderT GPU (State VNameSource) [VName])
-> (VName -> BinderT GPU (State VNameSource) Result)
-> BinderT GPU (State VNameSource) [VName]
forall a b. (a -> b) -> a -> b
$ \VName
ltid -> do
[VName]
-> Exp (Rep (BinderT GPU (State VNameSource)))
-> BinderT GPU (State VNameSource) ()
forall (m :: * -> *).
MonadBinder m =>
[VName] -> Exp (Rep m) -> m ()
letBindNames [VName
gtid]
(ExpT GPU -> BinderT GPU (State VNameSource) ())
-> BinderT GPU (State VNameSource) (ExpT GPU)
-> BinderT GPU (State VNameSource) ()
forall (m :: * -> *) a b. Monad m => (a -> m b) -> m a -> m b
=<< TPrimExp Int64 VName
-> BinderT
GPU
(State VNameSource)
(Exp (Rep (BinderT GPU (State VNameSource))))
forall a (m :: * -> *).
(ToExp a, MonadBinder m) =>
a -> m (Exp (Rep m))
toExp (VName -> TPrimExp Int64 VName
forall a. a -> TPrimExp Int64 a
le64 VName
gid TPrimExp Int64 VName
-> TPrimExp Int64 VName -> TPrimExp Int64 VName
forall a. Num a => a -> a -> a
* SubExp -> TPrimExp Int64 VName
pe64 SubExp
tile_size TPrimExp Int64 VName
-> TPrimExp Int64 VName -> TPrimExp Int64 VName
forall a. Num a => a -> a -> a
+ VName -> TPrimExp Int64 VName
forall a. a -> TPrimExp Int64 a
le64 VName
ltid)
PrimExp VName
-> Slice SubExp -> BinderT GPU (State VNameSource) Result
f (TPrimExp Bool VName -> PrimExp VName
forall t v. TPrimExp t v -> PrimExp v
untyped (TPrimExp Bool VName -> PrimExp VName)
-> TPrimExp Bool VName -> PrimExp VName
forall a b. (a -> b) -> a -> b
$ VName -> TPrimExp Int64 VName
forall a. a -> TPrimExp Int64 a
le64 VName
gtid TPrimExp Int64 VName -> TPrimExp Int64 VName -> TPrimExp Bool VName
forall t v. TPrimExp t v -> TPrimExp t v -> TPrimExp Bool v
.<. SubExp -> TPrimExp Int64 VName
pe64 SubExp
kdim) [SubExp -> DimIndex SubExp
forall d. d -> DimIndex d
DimFix (SubExp -> DimIndex SubExp) -> SubExp -> DimIndex SubExp
forall a b. (a -> b) -> a -> b
$ VName -> SubExp
Var VName
ltid],
tilingReadTile :: TileKind
-> PrivStms -> SubExp -> [InputArray] -> Binder GPU [InputTile]
tilingReadTile =
SubExp
-> VName
-> VName
-> Count NumGroups SubExp
-> Count GroupSize SubExp
-> TileKind
-> PrivStms
-> SubExp
-> [InputArray]
-> Binder GPU [InputTile]
readTile1D SubExp
tile_size VName
gid VName
gtid (SegLevel -> Count NumGroups SubExp
segNumGroups SegLevel
lvl) (SegLevel -> Count GroupSize SubExp
segGroupSize SegLevel
lvl),
tilingProcessTile :: ProcessTileArgs -> BinderT GPU (State VNameSource) [VName]
tilingProcessTile =
VName
-> VName
-> SubExp
-> SubExp
-> Count NumGroups SubExp
-> Count GroupSize SubExp
-> ProcessTileArgs
-> BinderT GPU (State VNameSource) [VName]
processTile1D VName
gid VName
gtid SubExp
kdim SubExp
tile_size (SegLevel -> Count NumGroups SubExp
segNumGroups SegLevel
lvl) (SegLevel -> Count GroupSize SubExp
segGroupSize SegLevel
lvl),
tilingProcessResidualTile :: ResidualTileArgs -> BinderT GPU (State VNameSource) [VName]
tilingProcessResidualTile =
VName
-> VName
-> SubExp
-> SubExp
-> Count NumGroups SubExp
-> Count GroupSize SubExp
-> ResidualTileArgs
-> BinderT GPU (State VNameSource) [VName]
processResidualTile1D VName
gid VName
gtid SubExp
kdim SubExp
tile_size (SegLevel -> Count NumGroups SubExp
segNumGroups SegLevel
lvl) (SegLevel -> Count GroupSize SubExp
segGroupSize SegLevel
lvl),
tilingTileReturns :: VName -> BinderT GPU (State VNameSource) KernelResult
tilingTileReturns = [(VName, SubExp)]
-> [(SubExp, SubExp)]
-> VName
-> BinderT GPU (State VNameSource) KernelResult
tileReturns [(VName, SubExp)]
dims_on_top [(SubExp
kdim, SubExp
tile_size)],
tilingTileShape :: Shape
tilingTileShape = Result -> Shape
forall d. [d] -> ShapeBase d
Shape [SubExp
tile_size],
tilingNumWholeTiles :: Binder GPU SubExp
tilingNumWholeTiles =
String
-> Exp (Rep (BinderT GPU (State VNameSource))) -> Binder GPU SubExp
forall (m :: * -> *).
MonadBinder m =>
String -> Exp (Rep m) -> m SubExp
letSubExp String
"num_whole_tiles" (Exp (Rep (BinderT GPU (State VNameSource))) -> Binder GPU SubExp)
-> Exp (Rep (BinderT GPU (State VNameSource))) -> Binder GPU SubExp
forall a b. (a -> b) -> a -> b
$
BasicOp -> ExpT GPU
forall rep. BasicOp -> ExpT rep
BasicOp (BasicOp -> ExpT GPU) -> BasicOp -> ExpT GPU
forall a b. (a -> b) -> a -> b
$ BinOp -> SubExp -> SubExp -> BasicOp
BinOp (IntType -> Safety -> BinOp
SQuot IntType
Int64 Safety
Unsafe) SubExp
w SubExp
tile_size,
tilingLevel :: SegLevel
tilingLevel = SegLevel
lvl,
tilingSpace :: SegSpace
tilingSpace = SegSpace
space
}
invariantToOneOfTwoInnerDims ::
Names ->
M.Map VName Names ->
[VName] ->
VName ->
Maybe InputArray
invariantToOneOfTwoInnerDims :: Names -> AliasTable -> [VName] -> VName -> Maybe InputArray
invariantToOneOfTwoInnerDims Names
branch_variant AliasTable
variance [VName]
dims VName
arr = do
VName
j : VName
i : [VName]
_ <- [VName] -> Maybe [VName]
forall a. a -> Maybe a
Just ([VName] -> Maybe [VName]) -> [VName] -> Maybe [VName]
forall a b. (a -> b) -> a -> b
$ [VName] -> [VName]
forall a. [a] -> [a]
reverse [VName]
dims
let variant_to :: Names
variant_to = Names -> VName -> AliasTable -> Names
forall k a. Ord k => a -> k -> Map k a -> a
M.findWithDefault Names
forall a. Monoid a => a
mempty VName
arr AliasTable
variance
branch_invariant :: Bool
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 Bool
branch_invariant Bool -> Bool -> Bool
&& VName
i VName -> Names -> Bool
`nameIn` Names
variant_to Bool -> Bool -> Bool
&& Bool -> Bool
not (VName
j VName -> Names -> Bool
`nameIn` Names
variant_to)
then InputArray -> Maybe InputArray
forall a. a -> Maybe a
Just (InputArray -> Maybe InputArray) -> InputArray -> Maybe InputArray
forall a b. (a -> b) -> a -> b
$ [Int] -> VName -> InputArray
InputTile [Int
0, Int
1] VName
arr
else
if Bool
branch_invariant Bool -> Bool -> Bool
&& VName
j VName -> Names -> Bool
`nameIn` Names
variant_to Bool -> Bool -> Bool
&& Bool -> Bool
not (VName
i VName -> Names -> Bool
`nameIn` Names
variant_to)
then InputArray -> Maybe InputArray
forall a. a -> Maybe a
Just (InputArray -> Maybe InputArray) -> InputArray -> Maybe InputArray
forall a b. (a -> b) -> a -> b
$ [Int] -> VName -> InputArray
InputTile [Int
1, Int
0] VName
arr
else InputArray -> Maybe InputArray
forall a. a -> Maybe a
Just (InputArray -> Maybe InputArray) -> InputArray -> Maybe InputArray
forall a b. (a -> b) -> a -> b
$ VName -> InputArray
InputDontTile VName
arr
reconstructGtids2D ::
SubExp ->
(VName, VName) ->
(VName, VName) ->
(VName, VName) ->
Binder GPU ()
reconstructGtids2D :: SubExp
-> (VName, VName)
-> (VName, VName)
-> (VName, VName)
-> BinderT GPU (State VNameSource) ()
reconstructGtids2D SubExp
tile_size (VName
gtid_x, VName
gtid_y) (VName
gid_x, VName
gid_y) (VName
ltid_x, VName
ltid_y) = do
[VName]
-> Exp (Rep (BinderT GPU (State VNameSource)))
-> BinderT GPU (State VNameSource) ()
forall (m :: * -> *).
MonadBinder m =>
[VName] -> Exp (Rep m) -> m ()
letBindNames [VName
gtid_x]
(ExpT GPU -> BinderT GPU (State VNameSource) ())
-> BinderT GPU (State VNameSource) (ExpT GPU)
-> BinderT GPU (State VNameSource) ()
forall (m :: * -> *) a b. Monad m => (a -> m b) -> m a -> m b
=<< TPrimExp Int64 VName
-> BinderT
GPU
(State VNameSource)
(Exp (Rep (BinderT GPU (State VNameSource))))
forall a (m :: * -> *).
(ToExp a, MonadBinder m) =>
a -> m (Exp (Rep m))
toExp (VName -> TPrimExp Int64 VName
forall a. a -> TPrimExp Int64 a
le64 VName
gid_x TPrimExp Int64 VName
-> TPrimExp Int64 VName -> TPrimExp Int64 VName
forall a. Num a => a -> a -> a
* SubExp -> TPrimExp Int64 VName
pe64 SubExp
tile_size TPrimExp Int64 VName
-> TPrimExp Int64 VName -> TPrimExp Int64 VName
forall a. Num a => a -> a -> a
+ VName -> TPrimExp Int64 VName
forall a. a -> TPrimExp Int64 a
le64 VName
ltid_x)
[VName]
-> Exp (Rep (BinderT GPU (State VNameSource)))
-> BinderT GPU (State VNameSource) ()
forall (m :: * -> *).
MonadBinder m =>
[VName] -> Exp (Rep m) -> m ()
letBindNames [VName
gtid_y]
(ExpT GPU -> BinderT GPU (State VNameSource) ())
-> BinderT GPU (State VNameSource) (ExpT GPU)
-> BinderT GPU (State VNameSource) ()
forall (m :: * -> *) a b. Monad m => (a -> m b) -> m a -> m b
=<< TPrimExp Int64 VName
-> BinderT
GPU
(State VNameSource)
(Exp (Rep (BinderT GPU (State VNameSource))))
forall a (m :: * -> *).
(ToExp a, MonadBinder m) =>
a -> m (Exp (Rep m))
toExp (VName -> TPrimExp Int64 VName
forall a. a -> TPrimExp Int64 a
le64 VName
gid_y TPrimExp Int64 VName
-> TPrimExp Int64 VName -> TPrimExp Int64 VName
forall a. Num a => a -> a -> a
* SubExp -> TPrimExp Int64 VName
pe64 SubExp
tile_size TPrimExp Int64 VName
-> TPrimExp Int64 VName -> TPrimExp Int64 VName
forall a. Num a => a -> a -> a
+ VName -> TPrimExp Int64 VName
forall a. a -> TPrimExp Int64 a
le64 VName
ltid_y)
readTile2D ::
(SubExp, SubExp) ->
(VName, VName) ->
(VName, VName) ->
SubExp ->
Count NumGroups SubExp ->
Count GroupSize SubExp ->
TileKind ->
PrivStms ->
SubExp ->
[InputArray] ->
Binder GPU [InputTile]
readTile2D :: (SubExp, SubExp)
-> (VName, VName)
-> (VName, VName)
-> SubExp
-> Count NumGroups SubExp
-> Count GroupSize SubExp
-> TileKind
-> PrivStms
-> SubExp
-> [InputArray]
-> Binder GPU [InputTile]
readTile2D (SubExp
kdim_x, SubExp
kdim_y) (VName
gtid_x, VName
gtid_y) (VName
gid_x, VName
gid_y) SubExp
tile_size Count NumGroups SubExp
num_groups Count GroupSize SubExp
group_size TileKind
kind PrivStms
privstms SubExp
tile_id [InputArray]
inputs =
([VName] -> [InputTile])
-> BinderT GPU (State VNameSource) [VName]
-> Binder GPU [InputTile]
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap ([InputArray] -> [VName] -> [InputTile]
inputsToTiles [InputArray]
inputs)
(BinderT GPU (State VNameSource) [VName] -> Binder GPU [InputTile])
-> (((VName, VName) -> BinderT GPU (State VNameSource) Result)
-> BinderT GPU (State VNameSource) [VName])
-> ((VName, VName) -> BinderT GPU (State VNameSource) Result)
-> Binder GPU [InputTile]
forall b c a. (b -> c) -> (a -> b) -> a -> c
. String
-> SegLevel
-> ResultManifest
-> (SubExp, SubExp)
-> ((VName, VName) -> BinderT GPU (State VNameSource) Result)
-> BinderT GPU (State VNameSource) [VName]
segMap2D
String
"full_tile"
(Count NumGroups SubExp
-> Count GroupSize SubExp -> SegVirt -> SegLevel
SegThread Count NumGroups SubExp
num_groups Count GroupSize SubExp
group_size SegVirt
SegNoVirtFull)
ResultManifest
ResultNoSimplify
(SubExp
tile_size, SubExp
tile_size)
(((VName, VName) -> BinderT GPU (State VNameSource) Result)
-> Binder GPU [InputTile])
-> ((VName, VName) -> BinderT GPU (State VNameSource) Result)
-> Binder GPU [InputTile]
forall a b. (a -> b) -> a -> b
$ \(VName
ltid_x, VName
ltid_y) -> do
SubExp
i <-
String
-> Exp (Rep (BinderT GPU (State VNameSource))) -> Binder GPU SubExp
forall (m :: * -> *).
MonadBinder m =>
String -> Exp (Rep m) -> m SubExp
letSubExp String
"i"
(ExpT GPU -> Binder GPU SubExp)
-> BinderT GPU (State VNameSource) (ExpT GPU) -> Binder GPU SubExp
forall (m :: * -> *) a b. Monad m => (a -> m b) -> m a -> m b
=<< TPrimExp Int64 VName
-> BinderT
GPU
(State VNameSource)
(Exp (Rep (BinderT GPU (State VNameSource))))
forall a (m :: * -> *).
(ToExp a, MonadBinder m) =>
a -> m (Exp (Rep m))
toExp (SubExp -> TPrimExp Int64 VName
pe64 SubExp
tile_id TPrimExp Int64 VName
-> TPrimExp Int64 VName -> TPrimExp Int64 VName
forall a. Num a => a -> a -> a
* SubExp -> TPrimExp Int64 VName
pe64 SubExp
tile_size TPrimExp Int64 VName
-> TPrimExp Int64 VName -> TPrimExp Int64 VName
forall a. Num a => a -> a -> a
+ VName -> TPrimExp Int64 VName
forall a. a -> TPrimExp Int64 a
le64 VName
ltid_x)
SubExp
j <-
String
-> Exp (Rep (BinderT GPU (State VNameSource))) -> Binder GPU SubExp
forall (m :: * -> *).
MonadBinder m =>
String -> Exp (Rep m) -> m SubExp
letSubExp String
"j"
(ExpT GPU -> Binder GPU SubExp)
-> BinderT GPU (State VNameSource) (ExpT GPU) -> Binder GPU SubExp
forall (m :: * -> *) a b. Monad m => (a -> m b) -> m a -> m b
=<< TPrimExp Int64 VName
-> BinderT
GPU
(State VNameSource)
(Exp (Rep (BinderT GPU (State VNameSource))))
forall a (m :: * -> *).
(ToExp a, MonadBinder m) =>
a -> m (Exp (Rep m))
toExp (SubExp -> TPrimExp Int64 VName
pe64 SubExp
tile_id TPrimExp Int64 VName
-> TPrimExp Int64 VName -> TPrimExp Int64 VName
forall a. Num a => a -> a -> a
* SubExp -> TPrimExp Int64 VName
pe64 SubExp
tile_size TPrimExp Int64 VName
-> TPrimExp Int64 VName -> TPrimExp Int64 VName
forall a. Num a => a -> a -> a
+ VName -> TPrimExp Int64 VName
forall a. a -> TPrimExp Int64 a
le64 VName
ltid_y)
SubExp
-> (VName, VName)
-> (VName, VName)
-> (VName, VName)
-> BinderT GPU (State VNameSource) ()
reconstructGtids2D SubExp
tile_size (VName
gtid_x, VName
gtid_y) (VName
gid_x, VName
gid_y) (VName
ltid_x, VName
ltid_y)
Slice SubExp -> PrivStms -> BinderT GPU (State VNameSource) ()
addPrivStms [SubExp -> DimIndex SubExp
forall d. d -> DimIndex d
DimFix (SubExp -> DimIndex SubExp) -> SubExp -> DimIndex SubExp
forall a b. (a -> b) -> a -> b
$ VName -> SubExp
Var VName
ltid_x, SubExp -> DimIndex SubExp
forall d. d -> DimIndex d
DimFix (SubExp -> DimIndex SubExp) -> SubExp -> DimIndex SubExp
forall a b. (a -> b) -> a -> b
$ VName -> SubExp
Var VName
ltid_y] PrivStms
privstms
let arrs_and_perms :: [(VName, [Int])]
arrs_and_perms = [InputArray] -> [(VName, [Int])]
tiledInputs [InputArray]
inputs
readTileElem :: (VName, [Int]) -> BinderT GPU (State VNameSource) VName
readTileElem (VName
arr, [Int]
perm) =
String
-> Exp (Rep (BinderT GPU (State VNameSource)))
-> BinderT GPU (State VNameSource) VName
forall (m :: * -> *).
MonadBinder m =>
String -> Exp (Rep m) -> m VName
letExp
String
"tile_elem"
( BasicOp -> ExpT GPU
forall rep. BasicOp -> ExpT rep
BasicOp (BasicOp -> ExpT GPU) -> BasicOp -> ExpT GPU
forall a b. (a -> b) -> a -> b
$
VName -> Slice SubExp -> BasicOp
Index
VName
arr
[SubExp -> DimIndex SubExp
forall d. d -> DimIndex d
DimFix (SubExp -> DimIndex SubExp) -> SubExp -> DimIndex SubExp
forall a b. (a -> b) -> a -> b
$ Result -> SubExp
forall a. [a] -> a
last (Result -> SubExp) -> Result -> SubExp
forall a b. (a -> b) -> a -> b
$ [Int] -> Result -> Result
forall a. [Int] -> [a] -> [a]
rearrangeShape [Int]
perm [SubExp
i, SubExp
j]]
)
readTileElemIfInBounds :: (VName, [Int]) -> BinderT GPU (State VNameSource) (ExpT GPU)
readTileElemIfInBounds (VName
arr, [Int]
perm) = do
Type
arr_t <- VName -> BinderT GPU (State VNameSource) Type
forall rep (m :: * -> *). HasScope rep m => VName -> m Type
lookupType VName
arr
let tile_t :: Type
tile_t = Type -> Type
forall shape u.
ArrayShape shape =>
TypeBase shape u -> TypeBase shape u
rowType Type
arr_t
w :: SubExp
w = Int -> Type -> SubExp
forall u. Int -> TypeBase Shape u -> SubExp
arraySize Int
0 Type
arr_t
idx :: SubExp
idx = Result -> SubExp
forall a. [a] -> a
last (Result -> SubExp) -> Result -> SubExp
forall a b. (a -> b) -> a -> b
$ [Int] -> Result -> Result
forall a. [Int] -> [a] -> [a]
rearrangeShape [Int]
perm [SubExp
i, SubExp
j]
othercheck :: TPrimExp Bool VName
othercheck =
[TPrimExp Bool VName] -> TPrimExp Bool VName
forall a. [a] -> a
last ([TPrimExp Bool VName] -> TPrimExp Bool VName)
-> [TPrimExp Bool VName] -> TPrimExp Bool VName
forall a b. (a -> b) -> a -> b
$
[Int] -> [TPrimExp Bool VName] -> [TPrimExp Bool VName]
forall a. [Int] -> [a] -> [a]
rearrangeShape
[Int]
perm
[ VName -> TPrimExp Int64 VName
forall a. a -> TPrimExp Int64 a
le64 VName
gtid_y TPrimExp Int64 VName -> TPrimExp Int64 VName -> TPrimExp Bool VName
forall t v. TPrimExp t v -> TPrimExp t v -> TPrimExp Bool v
.<. SubExp -> TPrimExp Int64 VName
pe64 SubExp
kdim_y,
VName -> TPrimExp Int64 VName
forall a. a -> TPrimExp Int64 a
le64 VName
gtid_x TPrimExp Int64 VName -> TPrimExp Int64 VName -> TPrimExp Bool VName
forall t v. TPrimExp t v -> TPrimExp t v -> TPrimExp Bool v
.<. SubExp -> TPrimExp Int64 VName
pe64 SubExp
kdim_x
]
BinderT
GPU
(State VNameSource)
(Exp (Rep (BinderT GPU (State VNameSource))))
-> BinderT
GPU
(State VNameSource)
(Body (Rep (BinderT GPU (State VNameSource))))
-> BinderT
GPU
(State VNameSource)
(Body (Rep (BinderT GPU (State VNameSource))))
-> BinderT
GPU
(State VNameSource)
(Exp (Rep (BinderT GPU (State VNameSource))))
forall (m :: * -> *).
(MonadBinder 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
-> BinderT
GPU
(State VNameSource)
(Exp (Rep (BinderT GPU (State VNameSource))))
forall a (m :: * -> *).
(ToExp a, MonadBinder m) =>
a -> m (Exp (Rep m))
toExp (TPrimExp Bool VName
-> BinderT
GPU
(State VNameSource)
(Exp (Rep (BinderT GPU (State VNameSource)))))
-> TPrimExp Bool VName
-> BinderT
GPU
(State VNameSource)
(Exp (Rep (BinderT GPU (State VNameSource))))
forall a b. (a -> b) -> a -> b
$ SubExp -> TPrimExp Int64 VName
pe64 SubExp
idx TPrimExp Int64 VName -> TPrimExp Int64 VName -> TPrimExp Bool VName
forall t v. TPrimExp t v -> TPrimExp t v -> TPrimExp Bool v
.<. SubExp -> TPrimExp Int64 VName
pe64 SubExp
w TPrimExp Bool VName -> TPrimExp Bool VName -> TPrimExp Bool VName
forall v. TPrimExp Bool v -> TPrimExp Bool v -> TPrimExp Bool v
.&&. TPrimExp Bool VName
othercheck)
([BinderT
GPU
(State VNameSource)
(Exp (Rep (BinderT GPU (State VNameSource))))]
-> BinderT
GPU
(State VNameSource)
(Body (Rep (BinderT GPU (State VNameSource))))
forall (m :: * -> *).
MonadBinder m =>
[m (Exp (Rep m))] -> m (Body (Rep m))
eBody [ExpT GPU -> BinderT GPU (State VNameSource) (ExpT GPU)
forall (m :: * -> *) a. Monad m => a -> m a
return (ExpT GPU -> BinderT GPU (State VNameSource) (ExpT GPU))
-> ExpT GPU -> BinderT GPU (State VNameSource) (ExpT GPU)
forall a b. (a -> b) -> a -> b
$ BasicOp -> ExpT GPU
forall rep. BasicOp -> ExpT rep
BasicOp (BasicOp -> ExpT GPU) -> BasicOp -> ExpT GPU
forall a b. (a -> b) -> a -> b
$ VName -> Slice SubExp -> BasicOp
Index VName
arr [SubExp -> DimIndex SubExp
forall d. d -> DimIndex d
DimFix SubExp
idx]])
([BinderT
GPU
(State VNameSource)
(Exp (Rep (BinderT GPU (State VNameSource))))]
-> BinderT
GPU
(State VNameSource)
(Body (Rep (BinderT GPU (State VNameSource))))
forall (m :: * -> *).
MonadBinder m =>
[m (Exp (Rep m))] -> m (Body (Rep m))
eBody [Type
-> BinderT
GPU
(State VNameSource)
(Exp (Rep (BinderT GPU (State VNameSource))))
forall (m :: * -> *). MonadBinder m => Type -> m (Exp (Rep m))
eBlank Type
tile_t])
([VName] -> Result)
-> BinderT GPU (State VNameSource) [VName]
-> BinderT GPU (State VNameSource) Result
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap ((VName -> SubExp) -> [VName] -> Result
forall a b. (a -> b) -> [a] -> [b]
map VName -> SubExp
Var) (BinderT GPU (State VNameSource) [VName]
-> BinderT GPU (State VNameSource) Result)
-> BinderT GPU (State VNameSource) [VName]
-> BinderT GPU (State VNameSource) Result
forall a b. (a -> b) -> a -> b
$
case TileKind
kind of
TileKind
TilePartial ->
((VName, [Int]) -> BinderT GPU (State VNameSource) VName)
-> [(VName, [Int])] -> BinderT GPU (State VNameSource) [VName]
forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
(a -> m b) -> t a -> m (t b)
mapM (String
-> Exp (Rep (BinderT GPU (State VNameSource)))
-> BinderT GPU (State VNameSource) VName
forall (m :: * -> *).
MonadBinder m =>
String -> Exp (Rep m) -> m VName
letExp String
"pre" (ExpT GPU -> BinderT GPU (State VNameSource) VName)
-> ((VName, [Int]) -> BinderT GPU (State VNameSource) (ExpT GPU))
-> (VName, [Int])
-> BinderT GPU (State VNameSource) VName
forall (m :: * -> *) b c a.
Monad m =>
(b -> m c) -> (a -> m b) -> a -> m c
<=< (VName, [Int]) -> BinderT GPU (State VNameSource) (ExpT GPU)
readTileElemIfInBounds) [(VName, [Int])]
arrs_and_perms
TileKind
TileFull ->
((VName, [Int]) -> BinderT GPU (State VNameSource) VName)
-> [(VName, [Int])] -> BinderT GPU (State VNameSource) [VName]
forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
(a -> m b) -> t a -> m (t b)
mapM (VName, [Int]) -> BinderT GPU (State VNameSource) VName
readTileElem [(VName, [Int])]
arrs_and_perms
findTileSize :: HasScope rep m => [InputTile] -> m SubExp
findTileSize :: forall rep (m :: * -> *). HasScope rep m => [InputTile] -> m SubExp
findTileSize [InputTile]
tiles =
case (InputTile -> Maybe VName) -> [InputTile] -> [VName]
forall a b. (a -> Maybe b) -> [a] -> [b]
mapMaybe InputTile -> Maybe VName
isTiled [InputTile]
tiles of
VName
v : [VName]
_ -> Int -> Type -> SubExp
forall u. Int -> TypeBase Shape u -> SubExp
arraySize Int
0 (Type -> SubExp) -> m Type -> m SubExp
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> VName -> m Type
forall rep (m :: * -> *). HasScope rep m => VName -> m Type
lookupType VName
v
[] -> SubExp -> m SubExp
forall (f :: * -> *) a. Applicative f => a -> f a
pure (SubExp -> m SubExp) -> SubExp -> m SubExp
forall a b. (a -> b) -> a -> b
$ IntType -> Integer -> SubExp
intConst IntType
Int64 Integer
0
where
isTiled :: InputTile -> Maybe VName
isTiled InputUntiled {} = Maybe VName
forall a. Maybe a
Nothing
isTiled (InputTiled [Int]
_ VName
tile) = VName -> Maybe VName
forall a. a -> Maybe a
Just VName
tile
processTile2D ::
(VName, VName) ->
(VName, VName) ->
(SubExp, SubExp) ->
SubExp ->
Count NumGroups SubExp ->
Count GroupSize SubExp ->
ProcessTileArgs ->
Binder GPU [VName]
processTile2D :: (VName, VName)
-> (VName, VName)
-> (SubExp, SubExp)
-> SubExp
-> Count NumGroups SubExp
-> Count GroupSize SubExp
-> ProcessTileArgs
-> BinderT GPU (State VNameSource) [VName]
processTile2D (VName
gid_x, VName
gid_y) (VName
gtid_x, VName
gtid_y) (SubExp
kdim_x, SubExp
kdim_y) SubExp
tile_size Count NumGroups SubExp
num_groups Count GroupSize SubExp
group_size ProcessTileArgs
tile_args = do
let privstms :: PrivStms
privstms = ProcessTileArgs -> PrivStms
processPrivStms ProcessTileArgs
tile_args
red_comm :: Commutativity
red_comm = ProcessTileArgs -> Commutativity
processComm ProcessTileArgs
tile_args
red_lam :: Lambda GPU
red_lam = ProcessTileArgs -> Lambda GPU
processRedLam ProcessTileArgs
tile_args
map_lam :: Lambda GPU
map_lam = ProcessTileArgs -> Lambda GPU
processMapLam ProcessTileArgs
tile_args
tiles :: [InputTile]
tiles = ProcessTileArgs -> [InputTile]
processTiles ProcessTileArgs
tile_args
accs :: [VName]
accs = ProcessTileArgs -> [VName]
processAcc ProcessTileArgs
tile_args
tile_id :: SubExp
tile_id = ProcessTileArgs -> SubExp
processTileId ProcessTileArgs
tile_args
SubExp
actual_tile_size <- [InputTile] -> Binder GPU SubExp
forall rep (m :: * -> *). HasScope rep m => [InputTile] -> m SubExp
findTileSize [InputTile]
tiles
String
-> SegLevel
-> ResultManifest
-> (SubExp, SubExp)
-> ((VName, VName) -> BinderT GPU (State VNameSource) Result)
-> BinderT GPU (State VNameSource) [VName]
segMap2D
String
"acc"
(Count NumGroups SubExp
-> Count GroupSize SubExp -> SegVirt -> SegLevel
SegThread Count NumGroups SubExp
num_groups Count GroupSize SubExp
group_size SegVirt
SegNoVirtFull)
ResultManifest
ResultPrivate
(SubExp
tile_size, SubExp
tile_size)
(((VName, VName) -> BinderT GPU (State VNameSource) Result)
-> BinderT GPU (State VNameSource) [VName])
-> ((VName, VName) -> BinderT GPU (State VNameSource) Result)
-> BinderT GPU (State VNameSource) [VName]
forall a b. (a -> b) -> a -> b
$ \(VName
ltid_x, VName
ltid_y) -> do
SubExp
-> (VName, VName)
-> (VName, VName)
-> (VName, VName)
-> BinderT GPU (State VNameSource) ()
reconstructGtids2D SubExp
tile_size (VName
gtid_x, VName
gtid_y) (VName
gid_x, VName
gid_y) (VName
ltid_x, VName
ltid_y)
Slice SubExp -> PrivStms -> BinderT GPU (State VNameSource) ()
addPrivStms [SubExp -> DimIndex SubExp
forall d. d -> DimIndex d
DimFix (SubExp -> DimIndex SubExp) -> SubExp -> DimIndex SubExp
forall a b. (a -> b) -> a -> b
$ VName -> SubExp
Var VName
ltid_x, SubExp -> DimIndex SubExp
forall d. d -> DimIndex d
DimFix (SubExp -> DimIndex SubExp) -> SubExp -> DimIndex SubExp
forall a b. (a -> b) -> a -> b
$ VName -> SubExp
Var VName
ltid_y] PrivStms
privstms
Result
thread_accs <- [VName]
-> (VName -> Binder GPU SubExp)
-> BinderT GPU (State VNameSource) Result
forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
t a -> (a -> m b) -> m (t b)
forM [VName]
accs ((VName -> Binder GPU SubExp)
-> BinderT GPU (State VNameSource) Result)
-> (VName -> Binder GPU SubExp)
-> BinderT GPU (State VNameSource) Result
forall a b. (a -> b) -> a -> b
$ \VName
acc ->
String
-> Exp (Rep (BinderT GPU (State VNameSource))) -> Binder GPU SubExp
forall (m :: * -> *).
MonadBinder m =>
String -> Exp (Rep m) -> m SubExp
letSubExp String
"acc" (Exp (Rep (BinderT GPU (State VNameSource))) -> Binder GPU SubExp)
-> Exp (Rep (BinderT GPU (State VNameSource))) -> Binder GPU SubExp
forall a b. (a -> b) -> a -> b
$ BasicOp -> ExpT GPU
forall rep. BasicOp -> ExpT rep
BasicOp (BasicOp -> ExpT GPU) -> BasicOp -> ExpT GPU
forall a b. (a -> b) -> a -> b
$ VName -> Slice SubExp -> BasicOp
Index VName
acc [SubExp -> DimIndex SubExp
forall d. d -> DimIndex d
DimFix (SubExp -> DimIndex SubExp) -> SubExp -> DimIndex SubExp
forall a b. (a -> b) -> a -> b
$ VName -> SubExp
Var VName
ltid_x, SubExp -> DimIndex SubExp
forall d. d -> DimIndex d
DimFix (SubExp -> DimIndex SubExp) -> SubExp -> DimIndex SubExp
forall a b. (a -> b) -> a -> b
$ VName -> SubExp
Var VName
ltid_y]
let form' :: ScremaForm GPU
form' = [Reduce GPU] -> Lambda GPU -> ScremaForm GPU
forall rep. [Reduce rep] -> Lambda rep -> ScremaForm rep
redomapSOAC [Commutativity -> Lambda GPU -> Result -> Reduce GPU
forall rep. Commutativity -> Lambda rep -> Result -> Reduce rep
Reduce Commutativity
red_comm Lambda GPU
red_lam Result
thread_accs] Lambda GPU
map_lam
sliceTile :: InputTile -> BinderT GPU (State VNameSource) VName
sliceTile (InputUntiled VName
arr) =
VName
-> SubExp
-> SubExp
-> SubExp
-> BinderT GPU (State VNameSource) VName
forall (m :: * -> *).
MonadBinder m =>
VName -> SubExp -> SubExp -> SubExp -> m VName
sliceUntiled VName
arr SubExp
tile_id SubExp
tile_size SubExp
actual_tile_size
sliceTile (InputTiled [Int]
perm VName
tile) = do
Type
tile_t <- VName -> BinderT GPU (State VNameSource) Type
forall rep (m :: * -> *). HasScope rep m => VName -> m Type
lookupType VName
tile
let idx :: DimIndex SubExp
idx = SubExp -> DimIndex SubExp
forall d. d -> DimIndex d
DimFix (SubExp -> DimIndex SubExp) -> SubExp -> DimIndex SubExp
forall a b. (a -> b) -> a -> b
$ VName -> SubExp
Var (VName -> SubExp) -> VName -> SubExp
forall a b. (a -> b) -> a -> b
$ [VName] -> VName
forall a. [a] -> a
head ([VName] -> VName) -> [VName] -> VName
forall a b. (a -> b) -> a -> b
$ [Int] -> [VName] -> [VName]
forall a. [Int] -> [a] -> [a]
rearrangeShape [Int]
perm [VName
ltid_x, VName
ltid_y]
String
-> Exp (Rep (BinderT GPU (State VNameSource)))
-> BinderT GPU (State VNameSource) VName
forall (m :: * -> *).
MonadBinder m =>
String -> Exp (Rep m) -> m VName
letExp String
"tile" (Exp (Rep (BinderT GPU (State VNameSource)))
-> BinderT GPU (State VNameSource) VName)
-> Exp (Rep (BinderT GPU (State VNameSource)))
-> BinderT GPU (State VNameSource) VName
forall a b. (a -> b) -> a -> b
$
BasicOp -> ExpT GPU
forall rep. BasicOp -> ExpT rep
BasicOp (BasicOp -> ExpT GPU) -> BasicOp -> ExpT GPU
forall a b. (a -> b) -> a -> b
$ VName -> Slice SubExp -> BasicOp
Index VName
tile (Slice SubExp -> BasicOp) -> Slice SubExp -> BasicOp
forall a b. (a -> b) -> a -> b
$ Type -> Int -> Slice SubExp -> Slice SubExp
sliceAt Type
tile_t ([Int] -> Int
forall a. [a] -> a
head [Int]
perm) [DimIndex SubExp
idx]
[VName]
tiles' <- (InputTile -> BinderT GPU (State VNameSource) VName)
-> [InputTile] -> BinderT GPU (State VNameSource) [VName]
forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
(a -> m b) -> t a -> m (t b)
mapM InputTile -> BinderT GPU (State VNameSource) VName
sliceTile [InputTile]
tiles
([VName] -> Result)
-> BinderT GPU (State VNameSource) [VName]
-> BinderT GPU (State VNameSource) Result
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap ((VName -> SubExp) -> [VName] -> Result
forall a b. (a -> b) -> [a] -> [b]
map VName -> SubExp
Var) (BinderT GPU (State VNameSource) [VName]
-> BinderT GPU (State VNameSource) Result)
-> BinderT GPU (State VNameSource) [VName]
-> BinderT GPU (State VNameSource) Result
forall a b. (a -> b) -> a -> b
$
String
-> Exp (Rep (BinderT GPU (State VNameSource)))
-> BinderT GPU (State VNameSource) [VName]
forall (m :: * -> *).
MonadBinder m =>
String -> Exp (Rep m) -> m [VName]
letTupExp String
"acc"
(ExpT GPU -> BinderT GPU (State VNameSource) [VName])
-> BinderT GPU (State VNameSource) (ExpT GPU)
-> BinderT GPU (State VNameSource) [VName]
forall (m :: * -> *) a b. Monad m => (a -> m b) -> m a -> m b
=<< BinderT
GPU
(State VNameSource)
(Exp (Rep (BinderT GPU (State VNameSource))))
-> BinderT
GPU
(State VNameSource)
(Body (Rep (BinderT GPU (State VNameSource))))
-> BinderT
GPU
(State VNameSource)
(Body (Rep (BinderT GPU (State VNameSource))))
-> BinderT
GPU
(State VNameSource)
(Exp (Rep (BinderT GPU (State VNameSource))))
forall (m :: * -> *).
(MonadBinder 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
-> BinderT
GPU
(State VNameSource)
(Exp (Rep (BinderT GPU (State VNameSource))))
forall a (m :: * -> *).
(ToExp a, MonadBinder m) =>
a -> m (Exp (Rep m))
toExp (TPrimExp Bool VName
-> BinderT
GPU
(State VNameSource)
(Exp (Rep (BinderT GPU (State VNameSource)))))
-> TPrimExp Bool VName
-> BinderT
GPU
(State VNameSource)
(Exp (Rep (BinderT 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 t v. TPrimExp t v -> TPrimExp t v -> TPrimExp Bool v
.<. SubExp -> TPrimExp Int64 VName
pe64 SubExp
kdim_x TPrimExp Bool VName -> TPrimExp Bool VName -> TPrimExp Bool VName
forall v. 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 t v. TPrimExp t v -> TPrimExp t v -> TPrimExp Bool v
.<. SubExp -> TPrimExp Int64 VName
pe64 SubExp
kdim_y
)
([BinderT
GPU
(State VNameSource)
(Exp (Rep (BinderT GPU (State VNameSource))))]
-> BinderT
GPU
(State VNameSource)
(Body (Rep (BinderT GPU (State VNameSource))))
forall (m :: * -> *).
MonadBinder m =>
[m (Exp (Rep m))] -> m (Body (Rep m))
eBody [ExpT GPU -> BinderT GPU (State VNameSource) (ExpT GPU)
forall (f :: * -> *) a. Applicative f => a -> f a
pure (ExpT GPU -> BinderT GPU (State VNameSource) (ExpT GPU))
-> ExpT GPU -> BinderT GPU (State VNameSource) (ExpT GPU)
forall a b. (a -> b) -> a -> b
$ Op GPU -> ExpT GPU
forall rep. Op rep -> ExpT rep
Op (Op GPU -> ExpT GPU) -> Op GPU -> ExpT GPU
forall a b. (a -> b) -> a -> b
$ SOAC GPU -> HostOp GPU (SOAC GPU)
forall rep op. op -> HostOp rep op
OtherOp (SOAC GPU -> HostOp GPU (SOAC GPU))
-> SOAC GPU -> HostOp GPU (SOAC GPU)
forall a b. (a -> b) -> a -> b
$ SubExp -> [VName] -> ScremaForm GPU -> SOAC GPU
forall rep. SubExp -> [VName] -> ScremaForm rep -> SOAC rep
Screma SubExp
actual_tile_size [VName]
tiles' ScremaForm GPU
form'])
(Result
-> BinderT
GPU
(State VNameSource)
(Body (Rep (BinderT GPU (State VNameSource))))
forall (m :: * -> *). MonadBinder m => Result -> m (Body (Rep m))
resultBodyM Result
thread_accs)
processResidualTile2D ::
(VName, VName) ->
(VName, VName) ->
(SubExp, SubExp) ->
SubExp ->
Count NumGroups SubExp ->
Count GroupSize SubExp ->
ResidualTileArgs ->
Binder GPU [VName]
processResidualTile2D :: (VName, VName)
-> (VName, VName)
-> (SubExp, SubExp)
-> SubExp
-> Count NumGroups SubExp
-> Count GroupSize SubExp
-> ResidualTileArgs
-> BinderT GPU (State VNameSource) [VName]
processResidualTile2D
(VName, VName)
gids
(VName, VName)
gtids
(SubExp, SubExp)
kdims
SubExp
tile_size
Count NumGroups SubExp
num_groups
Count GroupSize SubExp
group_size
ResidualTileArgs
args = do
SubExp
residual_input <-
String
-> Exp (Rep (BinderT GPU (State VNameSource))) -> Binder GPU SubExp
forall (m :: * -> *).
MonadBinder m =>
String -> Exp (Rep m) -> m SubExp
letSubExp String
"residual_input" (Exp (Rep (BinderT GPU (State VNameSource))) -> Binder GPU SubExp)
-> Exp (Rep (BinderT GPU (State VNameSource))) -> Binder GPU SubExp
forall a b. (a -> b) -> a -> b
$
BasicOp -> ExpT GPU
forall rep. BasicOp -> ExpT rep
BasicOp (BasicOp -> ExpT GPU) -> BasicOp -> ExpT GPU
forall a b. (a -> b) -> a -> b
$ BinOp -> SubExp -> SubExp -> BasicOp
BinOp (IntType -> Safety -> BinOp
SRem IntType
Int64 Safety
Unsafe) SubExp
w SubExp
tile_size
String
-> Exp (Rep (BinderT GPU (State VNameSource)))
-> BinderT GPU (State VNameSource) [VName]
forall (m :: * -> *).
MonadBinder m =>
String -> Exp (Rep m) -> m [VName]
letTupExp String
"acc_after_residual"
(ExpT GPU -> BinderT GPU (State VNameSource) [VName])
-> BinderT GPU (State VNameSource) (ExpT GPU)
-> BinderT GPU (State VNameSource) [VName]
forall (m :: * -> *) a b. Monad m => (a -> m b) -> m a -> m b
=<< BinderT
GPU
(State VNameSource)
(Exp (Rep (BinderT GPU (State VNameSource))))
-> BinderT
GPU
(State VNameSource)
(Body (Rep (BinderT GPU (State VNameSource))))
-> BinderT
GPU
(State VNameSource)
(Body (Rep (BinderT GPU (State VNameSource))))
-> BinderT
GPU
(State VNameSource)
(Exp (Rep (BinderT GPU (State VNameSource))))
forall (m :: * -> *).
(MonadBinder 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
-> BinderT
GPU
(State VNameSource)
(Exp (Rep (BinderT GPU (State VNameSource))))
forall a (m :: * -> *).
(ToExp a, MonadBinder m) =>
a -> m (Exp (Rep m))
toExp (TPrimExp Bool VName
-> BinderT
GPU
(State VNameSource)
(Exp (Rep (BinderT GPU (State VNameSource)))))
-> TPrimExp Bool VName
-> BinderT
GPU
(State VNameSource)
(Exp (Rep (BinderT GPU (State VNameSource))))
forall a b. (a -> b) -> a -> b
$ SubExp -> TPrimExp Int64 VName
pe64 SubExp
residual_input TPrimExp Int64 VName -> TPrimExp Int64 VName -> TPrimExp Bool VName
forall t v. TPrimExp t v -> TPrimExp t v -> TPrimExp Bool v
.==. TPrimExp Int64 VName
0)
(Result
-> BinderT
GPU
(State VNameSource)
(Body (Rep (BinderT GPU (State VNameSource))))
forall (m :: * -> *). MonadBinder m => Result -> m (Body (Rep m))
resultBodyM (Result
-> BinderT
GPU
(State VNameSource)
(Body (Rep (BinderT GPU (State VNameSource)))))
-> Result
-> BinderT
GPU
(State VNameSource)
(Body (Rep (BinderT GPU (State VNameSource))))
forall a b. (a -> b) -> a -> b
$ (VName -> SubExp) -> [VName] -> Result
forall a b. (a -> b) -> [a] -> [b]
map VName -> SubExp
Var [VName]
accs)
(SubExp -> BinderT GPU (State VNameSource) (Body GPU)
nonemptyTile SubExp
residual_input)
where
privstms :: PrivStms
privstms = ResidualTileArgs -> PrivStms
residualPrivStms ResidualTileArgs
args
red_comm :: Commutativity
red_comm = ResidualTileArgs -> Commutativity
residualComm ResidualTileArgs
args
red_lam :: Lambda GPU
red_lam = ResidualTileArgs -> Lambda GPU
residualRedLam ResidualTileArgs
args
map_lam :: Lambda GPU
map_lam = ResidualTileArgs -> Lambda GPU
residualMapLam ResidualTileArgs
args
accs :: [VName]
accs = ResidualTileArgs -> [VName]
residualAcc ResidualTileArgs
args
inputs :: [InputArray]
inputs = ResidualTileArgs -> [InputArray]
residualInput ResidualTileArgs
args
num_whole_tiles :: SubExp
num_whole_tiles = ResidualTileArgs -> SubExp
residualNumWholeTiles ResidualTileArgs
args
w :: SubExp
w = ResidualTileArgs -> SubExp
residualInputSize ResidualTileArgs
args
nonemptyTile :: SubExp -> BinderT GPU (State VNameSource) (Body GPU)
nonemptyTile SubExp
residual_input = Body GPU -> BinderT GPU (State VNameSource) (Body GPU)
forall rep (m :: * -> *).
(Renameable rep, MonadFreshNames m) =>
Body rep -> m (Body rep)
renameBody (Body GPU -> BinderT GPU (State VNameSource) (Body GPU))
-> (BinderT GPU (State VNameSource) (Body GPU)
-> BinderT GPU (State VNameSource) (Body GPU))
-> BinderT GPU (State VNameSource) (Body GPU)
-> BinderT GPU (State VNameSource) (Body GPU)
forall (m :: * -> *) b c a.
Monad m =>
(b -> m c) -> (a -> m b) -> a -> m c
<=< BinderT GPU (State VNameSource) (Body GPU)
-> BinderT GPU (State VNameSource) (Body GPU)
forall rep (m :: * -> *) somerep.
(Bindable rep, MonadFreshNames m, HasScope somerep m,
SameScope somerep rep) =>
Binder rep (Body rep) -> m (Body rep)
runBodyBinder (BinderT GPU (State VNameSource) (Body GPU)
-> BinderT GPU (State VNameSource) (Body GPU))
-> BinderT GPU (State VNameSource) (Body GPU)
-> BinderT GPU (State VNameSource) (Body GPU)
forall a b. (a -> b) -> a -> b
$ do
[InputTile]
full_tile <-
(SubExp, SubExp)
-> (VName, VName)
-> (VName, VName)
-> SubExp
-> Count NumGroups SubExp
-> Count GroupSize SubExp
-> TileKind
-> PrivStms
-> SubExp
-> [InputArray]
-> Binder GPU [InputTile]
readTile2D
(SubExp, SubExp)
kdims
(VName, VName)
gtids
(VName, VName)
gids
SubExp
tile_size
Count NumGroups SubExp
num_groups
Count GroupSize SubExp
group_size
TileKind
TilePartial
PrivStms
privstms
SubExp
num_whole_tiles
[InputArray]
inputs
let slice :: DimIndex SubExp
slice =
SubExp -> SubExp -> SubExp -> DimIndex SubExp
forall d. d -> d -> d -> DimIndex d
DimSlice (IntType -> Integer -> SubExp
intConst IntType
Int64 Integer
0) SubExp
residual_input (IntType -> Integer -> SubExp
intConst IntType
Int64 Integer
1)
[InputTile]
tiles <- [InputTile]
-> (InputTile -> BinderT GPU (State VNameSource) InputTile)
-> Binder GPU [InputTile]
forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
t a -> (a -> m b) -> m (t b)
forM [InputTile]
full_tile ((InputTile -> BinderT GPU (State VNameSource) InputTile)
-> Binder GPU [InputTile])
-> (InputTile -> BinderT GPU (State VNameSource) InputTile)
-> Binder GPU [InputTile]
forall a b. (a -> b) -> a -> b
$ \case
InputTiled [Int]
perm VName
tile' ->
[Int] -> VName -> InputTile
InputTiled [Int]
perm
(VName -> InputTile)
-> BinderT GPU (State VNameSource) VName
-> BinderT GPU (State VNameSource) InputTile
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> String
-> Exp (Rep (BinderT GPU (State VNameSource)))
-> BinderT GPU (State VNameSource) VName
forall (m :: * -> *).
MonadBinder m =>
String -> Exp (Rep m) -> m VName
letExp String
"partial_tile" (BasicOp -> ExpT GPU
forall rep. BasicOp -> ExpT rep
BasicOp (BasicOp -> ExpT GPU) -> BasicOp -> ExpT GPU
forall a b. (a -> b) -> a -> b
$ VName -> Slice SubExp -> BasicOp
Index VName
tile' [DimIndex SubExp
slice, DimIndex SubExp
slice])
InputUntiled VName
arr ->
InputTile -> BinderT GPU (State VNameSource) InputTile
forall (f :: * -> *) a. Applicative f => a -> f a
pure (InputTile -> BinderT GPU (State VNameSource) InputTile)
-> InputTile -> BinderT GPU (State VNameSource) InputTile
forall a b. (a -> b) -> a -> b
$ VName -> InputTile
InputUntiled VName
arr
let tile_args :: ProcessTileArgs
tile_args =
PrivStms
-> Commutativity
-> Lambda GPU
-> Lambda GPU
-> [InputTile]
-> [VName]
-> SubExp
-> ProcessTileArgs
ProcessTileArgs PrivStms
privstms Commutativity
red_comm Lambda GPU
red_lam Lambda GPU
map_lam [InputTile]
tiles [VName]
accs SubExp
num_whole_tiles
Result -> Body GPU
forall rep. Bindable rep => Result -> Body rep
resultBody (Result -> Body GPU) -> ([VName] -> Result) -> [VName] -> Body GPU
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (VName -> SubExp) -> [VName] -> Result
forall a b. (a -> b) -> [a] -> [b]
map VName -> SubExp
Var
([VName] -> Body GPU)
-> BinderT GPU (State VNameSource) [VName]
-> BinderT GPU (State VNameSource) (Body GPU)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> (VName, VName)
-> (VName, VName)
-> (SubExp, SubExp)
-> SubExp
-> Count NumGroups SubExp
-> Count GroupSize SubExp
-> ProcessTileArgs
-> BinderT GPU (State VNameSource) [VName]
processTile2D
(VName, VName)
gids
(VName, VName)
gtids
(SubExp, SubExp)
kdims
SubExp
tile_size
Count NumGroups SubExp
num_groups
Count GroupSize SubExp
group_size
ProcessTileArgs
tile_args
tiling2d :: [(VName, SubExp)] -> DoTiling (VName, VName) (SubExp, SubExp)
tiling2d :: [(VName, SubExp)] -> DoTiling (VName, VName) (SubExp, SubExp)
tiling2d [(VName, SubExp)]
dims_on_top SegLevel
_initial_lvl (VName
gtid_x, VName
gtid_y) (SubExp
kdim_x, SubExp
kdim_y) SubExp
w = do
VName
gid_x <- String -> BinderT GPU (State VNameSource) VName
forall (m :: * -> *). MonadFreshNames m => String -> m VName
newVName String
"gid_x"
VName
gid_y <- String -> BinderT GPU (State VNameSource) VName
forall (m :: * -> *). MonadFreshNames m => String -> m VName
newVName String
"gid_y"
Name
tile_size_key <- String -> Name
nameFromString (String -> Name) -> (VName -> String) -> VName -> Name
forall b c a. (b -> c) -> (a -> b) -> a -> c
. VName -> String
forall a. Pretty a => a -> String
pretty (VName -> Name)
-> BinderT GPU (State VNameSource) VName
-> BinderT GPU (State VNameSource) Name
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> String -> BinderT GPU (State VNameSource) VName
forall (m :: * -> *). MonadFreshNames m => String -> m VName
newVName String
"tile_size"
SubExp
tile_size <- String
-> Exp (Rep (BinderT GPU (State VNameSource))) -> Binder GPU SubExp
forall (m :: * -> *).
MonadBinder m =>
String -> Exp (Rep m) -> m SubExp
letSubExp String
"tile_size" (Exp (Rep (BinderT GPU (State VNameSource))) -> Binder GPU SubExp)
-> Exp (Rep (BinderT GPU (State VNameSource))) -> Binder GPU SubExp
forall a b. (a -> b) -> a -> b
$ Op GPU -> ExpT GPU
forall rep. Op rep -> ExpT rep
Op (Op GPU -> ExpT GPU) -> Op GPU -> ExpT GPU
forall a b. (a -> b) -> a -> b
$ SizeOp -> HostOp GPU (SOAC GPU)
forall rep op. SizeOp -> HostOp rep op
SizeOp (SizeOp -> HostOp GPU (SOAC GPU))
-> SizeOp -> HostOp GPU (SOAC GPU)
forall a b. (a -> b) -> a -> b
$ Name -> SizeClass -> SizeOp
GetSize Name
tile_size_key SizeClass
SizeTile
SubExp
group_size <- String
-> Exp (Rep (BinderT GPU (State VNameSource))) -> Binder GPU SubExp
forall (m :: * -> *).
MonadBinder m =>
String -> Exp (Rep m) -> m SubExp
letSubExp String
"group_size" (Exp (Rep (BinderT GPU (State VNameSource))) -> Binder GPU SubExp)
-> Exp (Rep (BinderT GPU (State VNameSource))) -> Binder GPU SubExp
forall a b. (a -> b) -> a -> b
$ BasicOp -> ExpT GPU
forall rep. BasicOp -> ExpT rep
BasicOp (BasicOp -> ExpT GPU) -> BasicOp -> ExpT GPU
forall a b. (a -> b) -> a -> b
$ BinOp -> SubExp -> SubExp -> BasicOp
BinOp (IntType -> Overflow -> BinOp
Mul IntType
Int64 Overflow
OverflowUndef) SubExp
tile_size SubExp
tile_size
SubExp
num_groups_x <-
String
-> Exp (Rep (BinderT GPU (State VNameSource))) -> Binder GPU SubExp
forall (m :: * -> *).
MonadBinder m =>
String -> Exp (Rep m) -> m SubExp
letSubExp String
"num_groups_x" (Exp (Rep (BinderT GPU (State VNameSource))) -> Binder GPU SubExp)
-> Exp (Rep (BinderT GPU (State VNameSource))) -> Binder GPU SubExp
forall a b. (a -> b) -> a -> b
$
BasicOp -> ExpT GPU
forall rep. BasicOp -> ExpT rep
BasicOp (BasicOp -> ExpT GPU) -> BasicOp -> ExpT GPU
forall a b. (a -> b) -> a -> b
$ BinOp -> SubExp -> SubExp -> BasicOp
BinOp (IntType -> Safety -> BinOp
SDivUp IntType
Int64 Safety
Unsafe) SubExp
kdim_x SubExp
tile_size
SubExp
num_groups_y <-
String
-> Exp (Rep (BinderT GPU (State VNameSource))) -> Binder GPU SubExp
forall (m :: * -> *).
MonadBinder m =>
String -> Exp (Rep m) -> m SubExp
letSubExp String
"num_groups_y" (Exp (Rep (BinderT GPU (State VNameSource))) -> Binder GPU SubExp)
-> Exp (Rep (BinderT GPU (State VNameSource))) -> Binder GPU SubExp
forall a b. (a -> b) -> a -> b
$
BasicOp -> ExpT GPU
forall rep. BasicOp -> ExpT rep
BasicOp (BasicOp -> ExpT GPU) -> BasicOp -> ExpT GPU
forall a b. (a -> b) -> a -> b
$ BinOp -> SubExp -> SubExp -> BasicOp
BinOp (IntType -> Safety -> BinOp
SDivUp IntType
Int64 Safety
Unsafe) SubExp
kdim_y SubExp
tile_size
SubExp
num_groups <-
String
-> Exp (Rep (BinderT GPU (State VNameSource))) -> Binder GPU SubExp
forall (m :: * -> *).
MonadBinder m =>
String -> Exp (Rep m) -> m SubExp
letSubExp String
"num_groups_top"
(ExpT GPU -> Binder GPU SubExp)
-> BinderT GPU (State VNameSource) (ExpT GPU) -> Binder GPU SubExp
forall (m :: * -> *) a b. Monad m => (a -> m b) -> m a -> m b
=<< BinOp
-> SubExp
-> Result
-> BinderT
GPU
(State VNameSource)
(Exp (Rep (BinderT GPU (State VNameSource))))
forall (m :: * -> *).
MonadBinder m =>
BinOp -> SubExp -> Result -> m (Exp (Rep m))
foldBinOp
(IntType -> Overflow -> BinOp
Mul IntType
Int64 Overflow
OverflowUndef)
SubExp
num_groups_x
(SubExp
num_groups_y SubExp -> Result -> Result
forall a. a -> [a] -> [a]
: ((VName, SubExp) -> SubExp) -> [(VName, SubExp)] -> Result
forall a b. (a -> b) -> [a] -> [b]
map (VName, SubExp) -> SubExp
forall a b. (a, b) -> b
snd [(VName, SubExp)]
dims_on_top)
VName
gid_flat <- String -> BinderT GPU (State VNameSource) VName
forall (m :: * -> *). MonadFreshNames m => String -> m VName
newVName String
"gid_flat"
let lvl :: SegLevel
lvl = Count NumGroups SubExp
-> Count GroupSize SubExp -> SegVirt -> SegLevel
SegGroup (SubExp -> Count NumGroups SubExp
forall u e. e -> Count u e
Count SubExp
num_groups) (SubExp -> Count GroupSize SubExp
forall u e. e -> Count u e
Count SubExp
group_size) SegVirt
SegNoVirtFull
space :: SegSpace
space =
VName -> [(VName, SubExp)] -> SegSpace
SegSpace VName
gid_flat ([(VName, SubExp)] -> SegSpace) -> [(VName, SubExp)] -> SegSpace
forall a b. (a -> b) -> a -> b
$
[(VName, SubExp)]
dims_on_top [(VName, SubExp)] -> [(VName, SubExp)] -> [(VName, SubExp)]
forall a. [a] -> [a] -> [a]
++ [(VName
gid_x, SubExp
num_groups_x), (VName
gid_y, SubExp
num_groups_y)]
Tiling -> Binder GPU Tiling
forall (m :: * -> *) a. Monad m => a -> m a
return
Tiling :: (String
-> SegLevel
-> ResultManifest
-> (PrimExp VName
-> Slice SubExp -> BinderT GPU (State VNameSource) Result)
-> BinderT GPU (State VNameSource) [VName])
-> (TileKind
-> PrivStms -> SubExp -> [InputArray] -> Binder GPU [InputTile])
-> (ProcessTileArgs -> BinderT GPU (State VNameSource) [VName])
-> (ResidualTileArgs -> BinderT GPU (State VNameSource) [VName])
-> (VName -> BinderT GPU (State VNameSource) KernelResult)
-> SegSpace
-> Shape
-> SegLevel
-> Binder GPU SubExp
-> Tiling
Tiling
{ tilingSegMap :: String
-> SegLevel
-> ResultManifest
-> (PrimExp VName
-> Slice SubExp -> BinderT GPU (State VNameSource) Result)
-> BinderT GPU (State VNameSource) [VName]
tilingSegMap = \String
desc SegLevel
lvl' ResultManifest
manifest PrimExp VName
-> Slice SubExp -> BinderT GPU (State VNameSource) Result
f ->
String
-> SegLevel
-> ResultManifest
-> (SubExp, SubExp)
-> ((VName, VName) -> BinderT GPU (State VNameSource) Result)
-> BinderT GPU (State VNameSource) [VName]
segMap2D String
desc SegLevel
lvl' ResultManifest
manifest (SubExp
tile_size, SubExp
tile_size) (((VName, VName) -> BinderT GPU (State VNameSource) Result)
-> BinderT GPU (State VNameSource) [VName])
-> ((VName, VName) -> BinderT GPU (State VNameSource) Result)
-> BinderT GPU (State VNameSource) [VName]
forall a b. (a -> b) -> a -> b
$ \(VName
ltid_x, VName
ltid_y) -> do
SubExp
-> (VName, VName)
-> (VName, VName)
-> (VName, VName)
-> BinderT GPU (State VNameSource) ()
reconstructGtids2D SubExp
tile_size (VName
gtid_x, VName
gtid_y) (VName
gid_x, VName
gid_y) (VName
ltid_x, VName
ltid_y)
PrimExp VName
-> Slice SubExp -> BinderT GPU (State VNameSource) Result
f
( TPrimExp Bool VName -> PrimExp VName
forall t v. TPrimExp t v -> PrimExp v
untyped (TPrimExp Bool VName -> PrimExp VName)
-> TPrimExp Bool VName -> PrimExp VName
forall a b. (a -> b) -> a -> b
$
VName -> TPrimExp Int64 VName
forall a. a -> TPrimExp Int64 a
le64 VName
gtid_x TPrimExp Int64 VName -> TPrimExp Int64 VName -> TPrimExp Bool VName
forall t v. TPrimExp t v -> TPrimExp t v -> TPrimExp Bool v
.<. SubExp -> TPrimExp Int64 VName
pe64 SubExp
kdim_x
TPrimExp Bool VName -> TPrimExp Bool VName -> TPrimExp Bool VName
forall v. 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 t v. TPrimExp t v -> TPrimExp t v -> TPrimExp Bool v
.<. SubExp -> TPrimExp Int64 VName
pe64 SubExp
kdim_y
)
[SubExp -> DimIndex SubExp
forall d. d -> DimIndex d
DimFix (SubExp -> DimIndex SubExp) -> SubExp -> DimIndex SubExp
forall a b. (a -> b) -> a -> b
$ VName -> SubExp
Var VName
ltid_x, SubExp -> DimIndex SubExp
forall d. d -> DimIndex d
DimFix (SubExp -> DimIndex SubExp) -> SubExp -> DimIndex SubExp
forall a b. (a -> b) -> a -> b
$ VName -> SubExp
Var VName
ltid_y],
tilingReadTile :: TileKind
-> PrivStms -> SubExp -> [InputArray] -> Binder GPU [InputTile]
tilingReadTile = (SubExp, SubExp)
-> (VName, VName)
-> (VName, VName)
-> SubExp
-> Count NumGroups SubExp
-> Count GroupSize SubExp
-> TileKind
-> PrivStms
-> SubExp
-> [InputArray]
-> Binder GPU [InputTile]
readTile2D (SubExp
kdim_x, SubExp
kdim_y) (VName
gtid_x, VName
gtid_y) (VName
gid_x, VName
gid_y) SubExp
tile_size (SegLevel -> Count NumGroups SubExp
segNumGroups SegLevel
lvl) (SegLevel -> Count GroupSize SubExp
segGroupSize SegLevel
lvl),
tilingProcessTile :: ProcessTileArgs -> BinderT GPU (State VNameSource) [VName]
tilingProcessTile = (VName, VName)
-> (VName, VName)
-> (SubExp, SubExp)
-> SubExp
-> Count NumGroups SubExp
-> Count GroupSize SubExp
-> ProcessTileArgs
-> BinderT GPU (State VNameSource) [VName]
processTile2D (VName
gid_x, VName
gid_y) (VName
gtid_x, VName
gtid_y) (SubExp
kdim_x, SubExp
kdim_y) SubExp
tile_size (SegLevel -> Count NumGroups SubExp
segNumGroups SegLevel
lvl) (SegLevel -> Count GroupSize SubExp
segGroupSize SegLevel
lvl),
tilingProcessResidualTile :: ResidualTileArgs -> BinderT GPU (State VNameSource) [VName]
tilingProcessResidualTile = (VName, VName)
-> (VName, VName)
-> (SubExp, SubExp)
-> SubExp
-> Count NumGroups SubExp
-> Count GroupSize SubExp
-> ResidualTileArgs
-> BinderT GPU (State VNameSource) [VName]
processResidualTile2D (VName
gid_x, VName
gid_y) (VName
gtid_x, VName
gtid_y) (SubExp
kdim_x, SubExp
kdim_y) SubExp
tile_size (SegLevel -> Count NumGroups SubExp
segNumGroups SegLevel
lvl) (SegLevel -> Count GroupSize SubExp
segGroupSize SegLevel
lvl),
tilingTileReturns :: VName -> BinderT GPU (State VNameSource) KernelResult
tilingTileReturns = [(VName, SubExp)]
-> [(SubExp, SubExp)]
-> VName
-> BinderT GPU (State VNameSource) KernelResult
tileReturns [(VName, SubExp)]
dims_on_top [(SubExp
kdim_x, SubExp
tile_size), (SubExp
kdim_y, SubExp
tile_size)],
tilingTileShape :: Shape
tilingTileShape = Result -> Shape
forall d. [d] -> ShapeBase d
Shape [SubExp
tile_size, SubExp
tile_size],
tilingNumWholeTiles :: Binder GPU SubExp
tilingNumWholeTiles =
String
-> Exp (Rep (BinderT GPU (State VNameSource))) -> Binder GPU SubExp
forall (m :: * -> *).
MonadBinder m =>
String -> Exp (Rep m) -> m SubExp
letSubExp String
"num_whole_tiles" (Exp (Rep (BinderT GPU (State VNameSource))) -> Binder GPU SubExp)
-> Exp (Rep (BinderT GPU (State VNameSource))) -> Binder GPU SubExp
forall a b. (a -> b) -> a -> b
$
BasicOp -> ExpT GPU
forall rep. BasicOp -> ExpT rep
BasicOp (BasicOp -> ExpT GPU) -> BasicOp -> ExpT GPU
forall a b. (a -> b) -> a -> b
$ BinOp -> SubExp -> SubExp -> BasicOp
BinOp (IntType -> Safety -> BinOp
SQuot IntType
Int64 Safety
Unsafe) SubExp
w SubExp
tile_size,
tilingLevel :: SegLevel
tilingLevel = SegLevel
lvl,
tilingSpace :: SegSpace
tilingSpace = SegSpace
space
}