module Futhark.CodeGen.ImpGen.GPU.Transpose
( TransposeType (..),
TransposeArgs,
mapTransposeKernel,
)
where
import Futhark.CodeGen.ImpCode.GPU
import Futhark.Util.IntegralExp (divUp, quot, rem)
import Prelude hiding (quot, rem)
data TransposeType
= TransposeNormal
| TransposeLowWidth
| TransposeLowHeight
|
TransposeSmall
deriving (TransposeType -> TransposeType -> Bool
forall a. (a -> a -> Bool) -> (a -> a -> Bool) -> Eq a
/= :: TransposeType -> TransposeType -> Bool
$c/= :: TransposeType -> TransposeType -> Bool
== :: TransposeType -> TransposeType -> Bool
$c== :: TransposeType -> TransposeType -> Bool
Eq, Eq TransposeType
TransposeType -> TransposeType -> Bool
TransposeType -> TransposeType -> Ordering
TransposeType -> TransposeType -> TransposeType
forall a.
Eq a
-> (a -> a -> Ordering)
-> (a -> a -> Bool)
-> (a -> a -> Bool)
-> (a -> a -> Bool)
-> (a -> a -> Bool)
-> (a -> a -> a)
-> (a -> a -> a)
-> Ord a
min :: TransposeType -> TransposeType -> TransposeType
$cmin :: TransposeType -> TransposeType -> TransposeType
max :: TransposeType -> TransposeType -> TransposeType
$cmax :: TransposeType -> TransposeType -> TransposeType
>= :: TransposeType -> TransposeType -> Bool
$c>= :: TransposeType -> TransposeType -> Bool
> :: TransposeType -> TransposeType -> Bool
$c> :: TransposeType -> TransposeType -> Bool
<= :: TransposeType -> TransposeType -> Bool
$c<= :: TransposeType -> TransposeType -> Bool
< :: TransposeType -> TransposeType -> Bool
$c< :: TransposeType -> TransposeType -> Bool
compare :: TransposeType -> TransposeType -> Ordering
$ccompare :: TransposeType -> TransposeType -> Ordering
Ord, Int -> TransposeType -> ShowS
[TransposeType] -> ShowS
TransposeType -> String
forall a.
(Int -> a -> ShowS) -> (a -> String) -> ([a] -> ShowS) -> Show a
showList :: [TransposeType] -> ShowS
$cshowList :: [TransposeType] -> ShowS
show :: TransposeType -> String
$cshow :: TransposeType -> String
showsPrec :: Int -> TransposeType -> ShowS
$cshowsPrec :: Int -> TransposeType -> ShowS
Show)
type TransposeArgs int =
( VName,
TExp int,
VName,
TExp int,
TExp int,
TExp int,
TExp int,
TExp int,
TExp int,
VName
)
elemsPerThread :: Num a => a
elemsPerThread :: forall a. Num a => a
elemsPerThread = a
8
mapTranspose :: forall int. IntExp int => (PrimType, VName -> TExp int) -> TExp int -> TransposeArgs int -> PrimType -> TransposeType -> KernelCode
mapTranspose :: forall {k} (int :: k).
IntExp int =>
(PrimType, VName -> TExp int)
-> TExp int
-> TransposeArgs int
-> PrimType
-> TransposeType
-> KernelCode
mapTranspose (PrimType
int, VName -> TPrimExp int VName
le) TPrimExp int VName
block_dim TransposeArgs int
args PrimType
t TransposeType
kind =
case TransposeType
kind of
TransposeType
TransposeSmall ->
forall a. Monoid a => [a] -> a
mconcat
[ KernelCode
get_ids,
forall {k} {t :: k} {a}. VName -> TPrimExp t VName -> Code a
dec VName
our_array_offset forall a b. (a -> b) -> a -> b
$ VName -> TPrimExp int VName
le VName
get_global_id_0 forall e. IntegralExp e => e -> e -> e
`quot` (TPrimExp int VName
height forall a. Num a => a -> a -> a
* TPrimExp int VName
width) forall a. Num a => a -> a -> a
* (TPrimExp int VName
height forall a. Num a => a -> a -> a
* TPrimExp int VName
width),
forall {k} {t :: k} {a}. VName -> TPrimExp t VName -> Code a
dec VName
x_index forall a b. (a -> b) -> a -> b
$ (VName -> TPrimExp int VName
le VName
get_global_id_0 forall e. IntegralExp e => e -> e -> e
`rem` (TPrimExp int VName
height forall a. Num a => a -> a -> a
* TPrimExp int VName
width)) forall e. IntegralExp e => e -> e -> e
`quot` TPrimExp int VName
height,
forall {k} {t :: k} {a}. VName -> TPrimExp t VName -> Code a
dec VName
y_index forall a b. (a -> b) -> a -> b
$ VName -> TPrimExp int VName
le VName
get_global_id_0 forall e. IntegralExp e => e -> e -> e
`rem` TPrimExp int VName
height,
forall a. VName -> Volatility -> PrimType -> Code a
DeclareScalar VName
val Volatility
Nonvolatile PrimType
t,
forall {k} {t :: k} {a}. VName -> TPrimExp t VName -> Code a
dec VName
odata_offset forall a b. (a -> b) -> a -> b
$
(TPrimExp int VName
basic_odata_offset forall e. IntegralExp e => e -> e -> e
`quot` forall a. Num a => PrimType -> a
primByteSize PrimType
t) forall a. Num a => a -> a -> a
+ VName -> TPrimExp int VName
le VName
our_array_offset,
forall {k} {t :: k} {a}. VName -> TPrimExp t VName -> Code a
dec VName
idata_offset forall a b. (a -> b) -> a -> b
$
(TPrimExp int VName
basic_idata_offset forall e. IntegralExp e => e -> e -> e
`quot` forall a. Num a => PrimType -> a
primByteSize PrimType
t) forall a. Num a => a -> a -> a
+ VName -> TPrimExp int VName
le VName
our_array_offset,
forall {k} {t :: k} {a}. VName -> TPrimExp t VName -> Code a
dec VName
index_in forall a b. (a -> b) -> a -> b
$ VName -> TPrimExp int VName
le VName
y_index forall a. Num a => a -> a -> a
* TPrimExp int VName
width forall a. Num a => a -> a -> a
+ VName -> TPrimExp int VName
le VName
x_index,
forall {k} {t :: k} {a}. VName -> TPrimExp t VName -> Code a
dec VName
index_out forall a b. (a -> b) -> a -> b
$ VName -> TPrimExp int VName
le VName
x_index forall a. Num a => a -> a -> a
* TPrimExp int VName
height forall a. Num a => a -> a -> a
+ VName -> TPrimExp int VName
le VName
y_index,
forall {a}. TExp Bool -> Code a -> Code a
when
(VName -> TPrimExp int VName
le VName
get_global_id_0 forall {k} v (t :: k).
Eq v =>
TPrimExp t v -> TPrimExp t v -> TPrimExp Bool v
.<. TPrimExp int VName
width forall a. Num a => a -> a -> a
* TPrimExp int VName
height forall a. Num a => a -> a -> a
* TPrimExp int VName
num_arrays)
( forall a. Monoid a => [a] -> a
mconcat
[ forall a.
VName
-> VName
-> Count Elements (TExp Int64)
-> PrimType
-> Space
-> Volatility
-> Code a
Read VName
val VName
idata (forall a. a -> Count Elements a
elements forall a b. (a -> b) -> a -> b
$ TPrimExp int VName -> TExp Int64
toOffset forall a b. (a -> b) -> a -> b
$ VName -> TPrimExp int VName
le VName
idata_offset forall a. Num a => a -> a -> a
+ VName -> TPrimExp int VName
le VName
index_in) PrimType
t (String -> Space
Space String
"global") Volatility
Nonvolatile,
forall a.
VName
-> Count Elements (TExp Int64)
-> PrimType
-> Space
-> Volatility
-> Exp
-> Code a
Write VName
odata (forall a. a -> Count Elements a
elements forall a b. (a -> b) -> a -> b
$ TPrimExp int VName -> TExp Int64
toOffset forall a b. (a -> b) -> a -> b
$ VName -> TPrimExp int VName
le VName
odata_offset forall a. Num a => a -> a -> a
+ VName -> TPrimExp int VName
le VName
index_out) PrimType
t (String -> Space
Space String
"global") Volatility
Nonvolatile (VName -> PrimType -> Exp
var VName
val PrimType
t)
]
)
]
TransposeType
TransposeLowWidth ->
KernelCode -> KernelCode
mkTranspose forall a b. (a -> b) -> a -> b
$
forall {k} {k} {k} {k} {t :: k} {t :: k} {t :: k} {t :: k}.
TPrimExp t VName
-> TPrimExp t VName
-> TPrimExp t VName
-> TPrimExp t VName
-> KernelCode
lowDimBody
(VName -> TPrimExp int VName
le VName
get_group_id_0 forall a. Num a => a -> a -> a
* TPrimExp int VName
block_dim forall a. Num a => a -> a -> a
+ (VName -> TPrimExp int VName
le VName
get_local_id_0 forall e. IntegralExp e => e -> e -> e
`quot` TPrimExp int VName
muly))
( VName -> TPrimExp int VName
le VName
get_group_id_1 forall a. Num a => a -> a -> a
* TPrimExp int VName
block_dim forall a. Num a => a -> a -> a
* TPrimExp int VName
muly
forall a. Num a => a -> a -> a
+ VName -> TPrimExp int VName
le VName
get_local_id_1
forall a. Num a => a -> a -> a
+ (VName -> TPrimExp int VName
le VName
get_local_id_0 forall e. IntegralExp e => e -> e -> e
`rem` TPrimExp int VName
muly) forall a. Num a => a -> a -> a
* TPrimExp int VName
block_dim
)
( VName -> TPrimExp int VName
le VName
get_group_id_1 forall a. Num a => a -> a -> a
* TPrimExp int VName
block_dim forall a. Num a => a -> a -> a
* TPrimExp int VName
muly
forall a. Num a => a -> a -> a
+ VName -> TPrimExp int VName
le VName
get_local_id_0
forall a. Num a => a -> a -> a
+ (VName -> TPrimExp int VName
le VName
get_local_id_1 forall e. IntegralExp e => e -> e -> e
`rem` TPrimExp int VName
muly) forall a. Num a => a -> a -> a
* TPrimExp int VName
block_dim
)
(VName -> TPrimExp int VName
le VName
get_group_id_0 forall a. Num a => a -> a -> a
* TPrimExp int VName
block_dim forall a. Num a => a -> a -> a
+ (VName -> TPrimExp int VName
le VName
get_local_id_1 forall e. IntegralExp e => e -> e -> e
`quot` TPrimExp int VName
muly))
TransposeType
TransposeLowHeight ->
KernelCode -> KernelCode
mkTranspose forall a b. (a -> b) -> a -> b
$
forall {k} {k} {k} {k} {t :: k} {t :: k} {t :: k} {t :: k}.
TPrimExp t VName
-> TPrimExp t VName
-> TPrimExp t VName
-> TPrimExp t VName
-> KernelCode
lowDimBody
( VName -> TPrimExp int VName
le VName
get_group_id_0 forall a. Num a => a -> a -> a
* TPrimExp int VName
block_dim forall a. Num a => a -> a -> a
* TPrimExp int VName
mulx
forall a. Num a => a -> a -> a
+ VName -> TPrimExp int VName
le VName
get_local_id_0
forall a. Num a => a -> a -> a
+ (VName -> TPrimExp int VName
le VName
get_local_id_1 forall e. IntegralExp e => e -> e -> e
`rem` TPrimExp int VName
mulx) forall a. Num a => a -> a -> a
* TPrimExp int VName
block_dim
)
(VName -> TPrimExp int VName
le VName
get_group_id_1 forall a. Num a => a -> a -> a
* TPrimExp int VName
block_dim forall a. Num a => a -> a -> a
+ (VName -> TPrimExp int VName
le VName
get_local_id_1 forall e. IntegralExp e => e -> e -> e
`quot` TPrimExp int VName
mulx))
(VName -> TPrimExp int VName
le VName
get_group_id_1 forall a. Num a => a -> a -> a
* TPrimExp int VName
block_dim forall a. Num a => a -> a -> a
+ (VName -> TPrimExp int VName
le VName
get_local_id_0 forall e. IntegralExp e => e -> e -> e
`quot` TPrimExp int VName
mulx))
( VName -> TPrimExp int VName
le VName
get_group_id_0 forall a. Num a => a -> a -> a
* TPrimExp int VName
block_dim forall a. Num a => a -> a -> a
* TPrimExp int VName
mulx
forall a. Num a => a -> a -> a
+ VName -> TPrimExp int VName
le VName
get_local_id_1
forall a. Num a => a -> a -> a
+ (VName -> TPrimExp int VName
le VName
get_local_id_0 forall e. IntegralExp e => e -> e -> e
`rem` TPrimExp int VName
mulx) forall a. Num a => a -> a -> a
* TPrimExp int VName
block_dim
)
TransposeType
TransposeNormal ->
KernelCode -> KernelCode
mkTranspose forall a b. (a -> b) -> a -> b
$
forall a. Monoid a => [a] -> a
mconcat
[ forall {k} {t :: k} {a}. VName -> TPrimExp t VName -> Code a
dec VName
x_index forall a b. (a -> b) -> a -> b
$ VName -> TPrimExp int VName
le VName
get_global_id_0,
forall {k} {t :: k} {a}. VName -> TPrimExp t VName -> Code a
dec VName
y_index forall a b. (a -> b) -> a -> b
$ VName -> TPrimExp int VName
le VName
get_group_id_1 forall a. Num a => a -> a -> a
* TPrimExp int VName
tile_dim forall a. Num a => a -> a -> a
+ VName -> TPrimExp int VName
le VName
get_local_id_1,
forall a. VName -> Volatility -> PrimType -> Code a
DeclareScalar VName
val Volatility
Nonvolatile PrimType
t,
forall {a}. TExp Bool -> Code a -> Code a
when (VName -> TPrimExp int VName
le VName
x_index forall {k} v (t :: k).
Eq v =>
TPrimExp t v -> TPrimExp t v -> TPrimExp Bool v
.<. TPrimExp int VName
width) forall a b. (a -> b) -> a -> b
$
forall a. VName -> Exp -> Code a -> Code a
For VName
j (forall {k} (t :: k) v. TPrimExp t v -> PrimExp v
untyped (forall a. Num a => a
elemsPerThread :: TExp int)) forall a b. (a -> b) -> a -> b
$
let i :: TPrimExp int VName
i = VName -> TPrimExp int VName
le VName
j forall a. Num a => a -> a -> a
* (TPrimExp int VName
tile_dim forall e. IntegralExp e => e -> e -> e
`quot` forall a. Num a => a
elemsPerThread)
in forall a. Monoid a => [a] -> a
mconcat
[ forall {k} {t :: k} {a}. VName -> TPrimExp t VName -> Code a
dec VName
index_in forall a b. (a -> b) -> a -> b
$ (VName -> TPrimExp int VName
le VName
y_index forall a. Num a => a -> a -> a
+ TPrimExp int VName
i) forall a. Num a => a -> a -> a
* TPrimExp int VName
width forall a. Num a => a -> a -> a
+ VName -> TPrimExp int VName
le VName
x_index,
forall {a}. TExp Bool -> Code a -> Code a
when (VName -> TPrimExp int VName
le VName
y_index forall a. Num a => a -> a -> a
+ TPrimExp int VName
i forall {k} v (t :: k).
Eq v =>
TPrimExp t v -> TPrimExp t v -> TPrimExp Bool v
.<. TPrimExp int VName
height) forall a b. (a -> b) -> a -> b
$
forall a. Monoid a => [a] -> a
mconcat
[ forall a.
VName
-> VName
-> Count Elements (TExp Int64)
-> PrimType
-> Space
-> Volatility
-> Code a
Read
VName
val
VName
idata
(forall a. a -> Count Elements a
elements forall a b. (a -> b) -> a -> b
$ TPrimExp int VName -> TExp Int64
toOffset forall a b. (a -> b) -> a -> b
$ VName -> TPrimExp int VName
le VName
idata_offset forall a. Num a => a -> a -> a
+ VName -> TPrimExp int VName
le VName
index_in)
PrimType
t
(String -> Space
Space String
"global")
Volatility
Nonvolatile,
forall a.
VName
-> Count Elements (TExp Int64)
-> PrimType
-> Space
-> Volatility
-> Exp
-> Code a
Write
VName
block
( forall a. a -> Count Elements a
elements forall a b. (a -> b) -> a -> b
$
TPrimExp int VName -> TExp Int64
toOffset forall a b. (a -> b) -> a -> b
$
(VName -> TPrimExp int VName
le VName
get_local_id_1 forall a. Num a => a -> a -> a
+ TPrimExp int VName
i) forall a. Num a => a -> a -> a
* (TPrimExp int VName
tile_dim forall a. Num a => a -> a -> a
+ TPrimExp int VName
1)
forall a. Num a => a -> a -> a
+ VName -> TPrimExp int VName
le VName
get_local_id_0
)
PrimType
t
(String -> Space
Space String
"local")
Volatility
Nonvolatile
(VName -> PrimType -> Exp
var VName
val PrimType
t)
]
],
forall a. a -> Code a
Op forall a b. (a -> b) -> a -> b
$ Fence -> KernelOp
Barrier Fence
FenceLocal,
forall a. VName -> Exp -> Code a
SetScalar VName
x_index forall a b. (a -> b) -> a -> b
$ forall {k} (t :: k) v. TPrimExp t v -> PrimExp v
untyped forall a b. (a -> b) -> a -> b
$ VName -> TPrimExp int VName
le VName
get_group_id_1 forall a. Num a => a -> a -> a
* TPrimExp int VName
tile_dim forall a. Num a => a -> a -> a
+ VName -> TPrimExp int VName
le VName
get_local_id_0,
forall a. VName -> Exp -> Code a
SetScalar VName
y_index forall a b. (a -> b) -> a -> b
$ forall {k} (t :: k) v. TPrimExp t v -> PrimExp v
untyped forall a b. (a -> b) -> a -> b
$ VName -> TPrimExp int VName
le VName
get_group_id_0 forall a. Num a => a -> a -> a
* TPrimExp int VName
tile_dim forall a. Num a => a -> a -> a
+ VName -> TPrimExp int VName
le VName
get_local_id_1,
forall {a}. TExp Bool -> Code a -> Code a
when (VName -> TPrimExp int VName
le VName
x_index forall {k} v (t :: k).
Eq v =>
TPrimExp t v -> TPrimExp t v -> TPrimExp Bool v
.<. TPrimExp int VName
height) forall a b. (a -> b) -> a -> b
$
forall a. VName -> Exp -> Code a -> Code a
For VName
j (forall {k} (t :: k) v. TPrimExp t v -> PrimExp v
untyped (forall a. Num a => a
elemsPerThread :: TExp int)) forall a b. (a -> b) -> a -> b
$
let i :: TPrimExp int VName
i = VName -> TPrimExp int VName
le VName
j forall a. Num a => a -> a -> a
* (TPrimExp int VName
tile_dim forall e. IntegralExp e => e -> e -> e
`quot` forall a. Num a => a
elemsPerThread)
in forall a. Monoid a => [a] -> a
mconcat
[ forall {k} {t :: k} {a}. VName -> TPrimExp t VName -> Code a
dec VName
index_out forall a b. (a -> b) -> a -> b
$ (VName -> TPrimExp int VName
le VName
y_index forall a. Num a => a -> a -> a
+ TPrimExp int VName
i) forall a. Num a => a -> a -> a
* TPrimExp int VName
height forall a. Num a => a -> a -> a
+ VName -> TPrimExp int VName
le VName
x_index,
forall {a}. TExp Bool -> Code a -> Code a
when (VName -> TPrimExp int VName
le VName
y_index forall a. Num a => a -> a -> a
+ TPrimExp int VName
i forall {k} v (t :: k).
Eq v =>
TPrimExp t v -> TPrimExp t v -> TPrimExp Bool v
.<. TPrimExp int VName
width) forall a b. (a -> b) -> a -> b
$
forall a. Monoid a => [a] -> a
mconcat
[ forall a.
VName
-> VName
-> Count Elements (TExp Int64)
-> PrimType
-> Space
-> Volatility
-> Code a
Read
VName
val
VName
block
( forall a. a -> Count Elements a
elements forall b c a. (b -> c) -> (a -> b) -> a -> c
. TPrimExp int VName -> TExp Int64
toOffset forall a b. (a -> b) -> a -> b
$
VName -> TPrimExp int VName
le VName
get_local_id_0 forall a. Num a => a -> a -> a
* (TPrimExp int VName
tile_dim forall a. Num a => a -> a -> a
+ TPrimExp int VName
1) forall a. Num a => a -> a -> a
+ VName -> TPrimExp int VName
le VName
get_local_id_1 forall a. Num a => a -> a -> a
+ TPrimExp int VName
i
)
PrimType
t
(String -> Space
Space String
"local")
Volatility
Nonvolatile,
forall a.
VName
-> Count Elements (TExp Int64)
-> PrimType
-> Space
-> Volatility
-> Exp
-> Code a
Write
VName
odata
(forall a. a -> Count Elements a
elements forall a b. (a -> b) -> a -> b
$ TPrimExp int VName -> TExp Int64
toOffset forall a b. (a -> b) -> a -> b
$ VName -> TPrimExp int VName
le VName
odata_offset forall a. Num a => a -> a -> a
+ VName -> TPrimExp int VName
le VName
index_out)
PrimType
t
(String -> Space
Space String
"global")
Volatility
Nonvolatile
(VName -> PrimType -> Exp
var VName
val PrimType
t)
]
]
]
where
toOffset :: TExp int -> TExp Int64
toOffset :: TPrimExp int VName -> TExp Int64
toOffset = forall {k} (t :: k) v. IntExp t => TPrimExp t v -> TPrimExp Int64 v
sExt64
dec :: VName -> TPrimExp t VName -> Code a
dec VName
v (TPrimExp Exp
e) =
forall a. VName -> Volatility -> PrimType -> Code a
DeclareScalar VName
v Volatility
Nonvolatile (forall v. PrimExp v -> PrimType
primExpType Exp
e) forall a. Semigroup a => a -> a -> a
<> forall a. VName -> Exp -> Code a
SetScalar VName
v Exp
e
tile_dim :: TPrimExp int VName
tile_dim = TPrimExp int VName
2 forall a. Num a => a -> a -> a
* TPrimExp int VName
block_dim
when :: TExp Bool -> Code a -> Code a
when TExp Bool
a Code a
b = forall a. TExp Bool -> Code a -> Code a -> Code a
If TExp Bool
a Code a
b forall a. Monoid a => a
mempty
( VName
odata,
TPrimExp int VName
basic_odata_offset,
VName
idata,
TPrimExp int VName
basic_idata_offset,
TPrimExp int VName
width,
TPrimExp int VName
height,
TPrimExp int VName
mulx,
TPrimExp int VName
muly,
TPrimExp int VName
num_arrays,
VName
block
) = TransposeArgs int
args
[ VName
our_array_offset,
VName
x_index,
VName
y_index,
VName
odata_offset,
VName
idata_offset,
VName
index_in,
VName
index_out,
VName
get_global_id_0,
VName
get_local_id_0,
VName
get_local_id_1,
VName
get_local_size_0,
VName
get_group_id_0,
VName
get_group_id_1,
VName
get_group_id_2,
VName
j,
VName
val
] =
forall a b c. (a -> b -> c) -> [a] -> [b] -> [c]
zipWith (forall a b c. (a -> b -> c) -> b -> a -> c
flip Name -> Int -> VName
VName) [Int
30 ..] forall a b. (a -> b) -> a -> b
$
forall a b. (a -> b) -> [a] -> [b]
map
String -> Name
nameFromString
[ String
"our_array_offset",
String
"x_index",
String
"y_index",
String
"odata_offset",
String
"idata_offset",
String
"index_in",
String
"index_out",
String
"get_global_id_0",
String
"get_local_id_0",
String
"get_local_id_1",
String
"get_local_size_0",
String
"get_group_id_0",
String
"get_group_id_1",
String
"get_group_id_2",
String
"j",
String
"val"
]
get_ids :: KernelCode
get_ids =
forall a. Monoid a => [a] -> a
mconcat
[ forall a. VName -> Volatility -> PrimType -> Code a
DeclareScalar VName
get_local_id_0 Volatility
Nonvolatile PrimType
int,
forall a. a -> Code a
Op forall a b. (a -> b) -> a -> b
$ VName -> Int -> KernelOp
GetLocalId VName
get_local_id_0 Int
0,
forall a. VName -> Volatility -> PrimType -> Code a
DeclareScalar VName
get_local_id_1 Volatility
Nonvolatile PrimType
int,
forall a. a -> Code a
Op forall a b. (a -> b) -> a -> b
$ VName -> Int -> KernelOp
GetLocalId VName
get_local_id_1 Int
1,
forall a. VName -> Volatility -> PrimType -> Code a
DeclareScalar VName
get_group_id_0 Volatility
Nonvolatile PrimType
int,
forall a. a -> Code a
Op forall a b. (a -> b) -> a -> b
$ VName -> Int -> KernelOp
GetGroupId VName
get_group_id_0 Int
0,
forall a. VName -> Volatility -> PrimType -> Code a
DeclareScalar VName
get_group_id_1 Volatility
Nonvolatile PrimType
int,
forall a. a -> Code a
Op forall a b. (a -> b) -> a -> b
$ VName -> Int -> KernelOp
GetGroupId VName
get_group_id_1 Int
1,
forall a. VName -> Volatility -> PrimType -> Code a
DeclareScalar VName
get_group_id_2 Volatility
Nonvolatile PrimType
int,
forall a. a -> Code a
Op forall a b. (a -> b) -> a -> b
$ VName -> Int -> KernelOp
GetGroupId VName
get_group_id_2 Int
2,
forall a. VName -> Volatility -> PrimType -> Code a
DeclareScalar VName
get_local_size_0 Volatility
Nonvolatile PrimType
int,
forall a. a -> Code a
Op forall a b. (a -> b) -> a -> b
$ VName -> Int -> KernelOp
GetLocalSize VName
get_local_size_0 Int
0,
forall a. VName -> Volatility -> PrimType -> Code a
DeclareScalar VName
get_global_id_0 Volatility
Nonvolatile PrimType
int,
forall a. VName -> Exp -> Code a
SetScalar VName
get_global_id_0 forall a b. (a -> b) -> a -> b
$ forall {k} (t :: k) v. TPrimExp t v -> PrimExp v
untyped forall a b. (a -> b) -> a -> b
$ VName -> TPrimExp int VName
le VName
get_group_id_0 forall a. Num a => a -> a -> a
* VName -> TPrimExp int VName
le VName
get_local_size_0 forall a. Num a => a -> a -> a
+ VName -> TPrimExp int VName
le VName
get_local_id_0
]
mkTranspose :: KernelCode -> KernelCode
mkTranspose KernelCode
body =
forall a. Monoid a => [a] -> a
mconcat
[ KernelCode
get_ids,
forall {k} {t :: k} {a}. VName -> TPrimExp t VName -> Code a
dec VName
our_array_offset forall a b. (a -> b) -> a -> b
$ VName -> TPrimExp int VName
le VName
get_group_id_2 forall a. Num a => a -> a -> a
* TPrimExp int VName
width forall a. Num a => a -> a -> a
* TPrimExp int VName
height,
forall {k} {t :: k} {a}. VName -> TPrimExp t VName -> Code a
dec VName
odata_offset forall a b. (a -> b) -> a -> b
$
(TPrimExp int VName
basic_odata_offset forall e. IntegralExp e => e -> e -> e
`quot` forall a. Num a => PrimType -> a
primByteSize PrimType
t) forall a. Num a => a -> a -> a
+ VName -> TPrimExp int VName
le VName
our_array_offset,
forall {k} {t :: k} {a}. VName -> TPrimExp t VName -> Code a
dec VName
idata_offset forall a b. (a -> b) -> a -> b
$
(TPrimExp int VName
basic_idata_offset forall e. IntegralExp e => e -> e -> e
`quot` forall a. Num a => PrimType -> a
primByteSize PrimType
t) forall a. Num a => a -> a -> a
+ VName -> TPrimExp int VName
le VName
our_array_offset,
KernelCode
body
]
lowDimBody :: TPrimExp t VName
-> TPrimExp t VName
-> TPrimExp t VName
-> TPrimExp t VName
-> KernelCode
lowDimBody TPrimExp t VName
x_in_index TPrimExp t VName
y_in_index TPrimExp t VName
x_out_index TPrimExp t VName
y_out_index =
forall a. Monoid a => [a] -> a
mconcat
[ forall {k} {t :: k} {a}. VName -> TPrimExp t VName -> Code a
dec VName
x_index TPrimExp t VName
x_in_index,
forall {k} {t :: k} {a}. VName -> TPrimExp t VName -> Code a
dec VName
y_index TPrimExp t VName
y_in_index,
forall a. VName -> Volatility -> PrimType -> Code a
DeclareScalar VName
val Volatility
Nonvolatile PrimType
t,
forall {k} {t :: k} {a}. VName -> TPrimExp t VName -> Code a
dec VName
index_in forall a b. (a -> b) -> a -> b
$ VName -> TPrimExp int VName
le VName
y_index forall a. Num a => a -> a -> a
* TPrimExp int VName
width forall a. Num a => a -> a -> a
+ VName -> TPrimExp int VName
le VName
x_index,
forall {a}. TExp Bool -> Code a -> Code a
when (VName -> TPrimExp int VName
le VName
x_index forall {k} v (t :: k).
Eq v =>
TPrimExp t v -> TPrimExp t v -> TPrimExp Bool v
.<. TPrimExp int VName
width forall v.
Eq v =>
TPrimExp Bool v -> TPrimExp Bool v -> TPrimExp Bool v
.&&. VName -> TPrimExp int VName
le VName
y_index forall {k} v (t :: k).
Eq v =>
TPrimExp t v -> TPrimExp t v -> TPrimExp Bool v
.<. TPrimExp int VName
height) forall a b. (a -> b) -> a -> b
$
forall a. Monoid a => [a] -> a
mconcat
[ forall a.
VName
-> VName
-> Count Elements (TExp Int64)
-> PrimType
-> Space
-> Volatility
-> Code a
Read
VName
val
VName
idata
(forall a. a -> Count Elements a
elements forall a b. (a -> b) -> a -> b
$ forall {k} (t :: k) v. IntExp t => TPrimExp t v -> TPrimExp Int64 v
sExt64 forall a b. (a -> b) -> a -> b
$ VName -> TPrimExp int VName
le VName
idata_offset forall a. Num a => a -> a -> a
+ VName -> TPrimExp int VName
le VName
index_in)
PrimType
t
(String -> Space
Space String
"global")
Volatility
Nonvolatile,
forall a.
VName
-> Count Elements (TExp Int64)
-> PrimType
-> Space
-> Volatility
-> Exp
-> Code a
Write
VName
block
(forall a. a -> Count Elements a
elements forall a b. (a -> b) -> a -> b
$ forall {k} (t :: k) v. IntExp t => TPrimExp t v -> TPrimExp Int64 v
sExt64 forall a b. (a -> b) -> a -> b
$ VName -> TPrimExp int VName
le VName
get_local_id_1 forall a. Num a => a -> a -> a
* (TPrimExp int VName
block_dim forall a. Num a => a -> a -> a
+ TPrimExp int VName
1) forall a. Num a => a -> a -> a
+ VName -> TPrimExp int VName
le VName
get_local_id_0)
PrimType
t
(String -> Space
Space String
"local")
Volatility
Nonvolatile
(VName -> PrimType -> Exp
var VName
val PrimType
t)
],
forall a. a -> Code a
Op forall a b. (a -> b) -> a -> b
$ Fence -> KernelOp
Barrier Fence
FenceLocal,
forall a. VName -> Exp -> Code a
SetScalar VName
x_index forall a b. (a -> b) -> a -> b
$ forall {k} (t :: k) v. TPrimExp t v -> PrimExp v
untyped TPrimExp t VName
x_out_index,
forall a. VName -> Exp -> Code a
SetScalar VName
y_index forall a b. (a -> b) -> a -> b
$ forall {k} (t :: k) v. TPrimExp t v -> PrimExp v
untyped TPrimExp t VName
y_out_index,
forall {k} {t :: k} {a}. VName -> TPrimExp t VName -> Code a
dec VName
index_out forall a b. (a -> b) -> a -> b
$ VName -> TPrimExp int VName
le VName
y_index forall a. Num a => a -> a -> a
* TPrimExp int VName
height forall a. Num a => a -> a -> a
+ VName -> TPrimExp int VName
le VName
x_index,
forall {a}. TExp Bool -> Code a -> Code a
when (VName -> TPrimExp int VName
le VName
x_index forall {k} v (t :: k).
Eq v =>
TPrimExp t v -> TPrimExp t v -> TPrimExp Bool v
.<. TPrimExp int VName
height forall v.
Eq v =>
TPrimExp Bool v -> TPrimExp Bool v -> TPrimExp Bool v
.&&. VName -> TPrimExp int VName
le VName
y_index forall {k} v (t :: k).
Eq v =>
TPrimExp t v -> TPrimExp t v -> TPrimExp Bool v
.<. TPrimExp int VName
width) forall a b. (a -> b) -> a -> b
$
forall a. Monoid a => [a] -> a
mconcat
[ forall a.
VName
-> VName
-> Count Elements (TExp Int64)
-> PrimType
-> Space
-> Volatility
-> Code a
Read
VName
val
VName
block
(forall a. a -> Count Elements a
elements forall a b. (a -> b) -> a -> b
$ TPrimExp int VName -> TExp Int64
toOffset forall a b. (a -> b) -> a -> b
$ VName -> TPrimExp int VName
le VName
get_local_id_0 forall a. Num a => a -> a -> a
* (TPrimExp int VName
block_dim forall a. Num a => a -> a -> a
+ TPrimExp int VName
1) forall a. Num a => a -> a -> a
+ VName -> TPrimExp int VName
le VName
get_local_id_1)
PrimType
t
(String -> Space
Space String
"local")
Volatility
Nonvolatile,
forall a.
VName
-> Count Elements (TExp Int64)
-> PrimType
-> Space
-> Volatility
-> Exp
-> Code a
Write
VName
odata
(forall a. a -> Count Elements a
elements forall a b. (a -> b) -> a -> b
$ TPrimExp int VName -> TExp Int64
toOffset (VName -> TPrimExp int VName
le VName
odata_offset forall a. Num a => a -> a -> a
+ VName -> TPrimExp int VName
le VName
index_out))
PrimType
t
(String -> Space
Space String
"global")
Volatility
Nonvolatile
(VName -> PrimType -> Exp
var VName
val PrimType
t)
]
]
lowDimKernelAndGroupSize ::
IntExp int =>
TExp int ->
TExp int ->
TExp int ->
TExp int ->
([TExp int], [TExp int])
lowDimKernelAndGroupSize :: forall {k} (int :: k).
IntExp int =>
TExp int
-> TExp int -> TExp int -> TExp int -> ([TExp int], [TExp int])
lowDimKernelAndGroupSize TExp int
block_dim TExp int
num_arrays TExp int
x_elems TExp int
y_elems =
( [ TExp int
x_elems forall e. IntegralExp e => e -> e -> e
`divUp` TExp int
block_dim,
TExp int
y_elems forall e. IntegralExp e => e -> e -> e
`divUp` TExp int
block_dim,
TExp int
num_arrays
],
[TExp int
block_dim, TExp int
block_dim, TExp int
1]
)
mapTransposeKernel ::
forall int.
IntExp int =>
(PrimType, VName -> TExp int) ->
String ->
Integer ->
TransposeArgs int ->
PrimType ->
TransposeType ->
Kernel
mapTransposeKernel :: forall {k} (int :: k).
IntExp int =>
(PrimType, VName -> TExp int)
-> String
-> Integer
-> TransposeArgs int
-> PrimType
-> TransposeType
-> Kernel
mapTransposeKernel (PrimType
int, VName -> TExp int
le) String
desc Integer
block_dim_int TransposeArgs int
args PrimType
t TransposeType
kind =
Kernel
{ kernelBody :: KernelCode
kernelBody =
forall a. VName -> Space -> Code a
DeclareMem VName
block (String -> Space
Space String
"local")
forall a. Semigroup a => a -> a -> a
<> forall a. a -> Code a
Op (VName -> Count Bytes (TExp Int64) -> KernelOp
LocalAlloc VName
block Count Bytes (TExp Int64)
block_size)
forall a. Semigroup a => a -> a -> a
<> forall {k} (int :: k).
IntExp int =>
(PrimType, VName -> TExp int)
-> TExp int
-> TransposeArgs int
-> PrimType
-> TransposeType
-> KernelCode
mapTranspose (PrimType
int, VName -> TExp int
le) TExp int
block_dim TransposeArgs int
args PrimType
t TransposeType
kind,
kernelUses :: [KernelUse]
kernelUses = [KernelUse]
uses,
kernelNumGroups :: [Exp]
kernelNumGroups = forall a b. (a -> b) -> [a] -> [b]
map forall {k} (t :: k) v. TPrimExp t v -> PrimExp v
untyped [TExp int]
num_groups,
kernelGroupSize :: [GroupDim]
kernelGroupSize = forall a b. (a -> b) -> [a] -> [b]
map (forall a b. a -> Either a b
Left forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall {k} (t :: k) v. TPrimExp t v -> PrimExp v
untyped) [TExp int]
group_size,
kernelName :: Name
kernelName = String -> Name
nameFromString (String
name forall a. Semigroup a => a -> a -> a
<> String
"_" forall a. Semigroup a => a -> a -> a
<> forall a. Pretty a => a -> String
prettyString PrimType
int),
kernelFailureTolerant :: Bool
kernelFailureTolerant = Bool
True,
kernelCheckLocalMemory :: Bool
kernelCheckLocalMemory = Bool
False
}
where
pad2DBytes :: a -> a
pad2DBytes a
k = a
k forall a. Num a => a -> a -> a
* (a
k forall a. Num a => a -> a -> a
+ a
1) forall a. Num a => a -> a -> a
* forall a. Num a => PrimType -> a
primByteSize PrimType
t
block_size :: Count Bytes (TExp Int64)
block_size :: Count Bytes (TExp Int64)
block_size =
forall a. a -> Count Bytes a
bytes forall a b. (a -> b) -> a -> b
$
case TransposeType
kind of
TransposeType
TransposeSmall -> TExp Int64
1
TransposeType
TransposeNormal -> forall a. Num a => Integer -> a
fromInteger forall a b. (a -> b) -> a -> b
$ forall {a}. Num a => a -> a
pad2DBytes forall a b. (a -> b) -> a -> b
$ Integer
2 forall a. Num a => a -> a -> a
* Integer
block_dim_int
TransposeType
TransposeLowWidth -> forall a. Num a => Integer -> a
fromInteger forall a b. (a -> b) -> a -> b
$ forall {a}. Num a => a -> a
pad2DBytes Integer
block_dim_int
TransposeType
TransposeLowHeight -> forall a. Num a => Integer -> a
fromInteger forall a b. (a -> b) -> a -> b
$ forall {a}. Num a => a -> a
pad2DBytes Integer
block_dim_int
block_dim :: TExp int
block_dim = forall a. Num a => Integer -> a
fromInteger Integer
block_dim_int :: TExp int
( VName
odata,
TExp int
basic_odata_offset,
VName
idata,
TExp int
basic_idata_offset,
TExp int
width,
TExp int
height,
TExp int
mulx,
TExp int
muly,
TExp int
num_arrays,
VName
block
) = TransposeArgs int
args
([TExp int]
num_groups, [TExp int]
group_size) =
case TransposeType
kind of
TransposeType
TransposeSmall ->
( [(TExp int
num_arrays forall a. Num a => a -> a -> a
* TExp int
width forall a. Num a => a -> a -> a
* TExp int
height) forall e. IntegralExp e => e -> e -> e
`divUp` (TExp int
block_dim forall a. Num a => a -> a -> a
* TExp int
block_dim)],
[TExp int
block_dim forall a. Num a => a -> a -> a
* TExp int
block_dim]
)
TransposeType
TransposeLowWidth ->
forall {k} (int :: k).
IntExp int =>
TExp int
-> TExp int -> TExp int -> TExp int -> ([TExp int], [TExp int])
lowDimKernelAndGroupSize TExp int
block_dim TExp int
num_arrays TExp int
width forall a b. (a -> b) -> a -> b
$ TExp int
height forall e. IntegralExp e => e -> e -> e
`divUp` TExp int
muly
TransposeType
TransposeLowHeight ->
forall {k} (int :: k).
IntExp int =>
TExp int
-> TExp int -> TExp int -> TExp int -> ([TExp int], [TExp int])
lowDimKernelAndGroupSize TExp int
block_dim TExp int
num_arrays (TExp int
width forall e. IntegralExp e => e -> e -> e
`divUp` TExp int
mulx) TExp int
height
TransposeType
TransposeNormal ->
let actual_dim :: TExp int
actual_dim = TExp int
block_dim forall a. Num a => a -> a -> a
* TExp int
2
in ( [ TExp int
width forall e. IntegralExp e => e -> e -> e
`divUp` TExp int
actual_dim,
TExp int
height forall e. IntegralExp e => e -> e -> e
`divUp` TExp int
actual_dim,
TExp int
num_arrays
],
[TExp int
actual_dim, TExp int
actual_dim forall e. IntegralExp e => e -> e -> e
`quot` forall a. Num a => a
elemsPerThread, TExp int
1]
)
uses :: [KernelUse]
uses =
forall a b. (a -> b) -> [a] -> [b]
map
(VName -> PrimType -> KernelUse
`ScalarUse` IntType -> PrimType
IntType IntType
Int64)
( Names -> [VName]
namesToList forall a b. (a -> b) -> a -> b
$
forall a. Monoid a => [a] -> a
mconcat forall a b. (a -> b) -> a -> b
$
forall a b. (a -> b) -> [a] -> [b]
map
forall a. FreeIn a => a -> Names
freeIn
[ TExp int
basic_odata_offset,
TExp int
basic_idata_offset,
TExp int
num_arrays,
TExp int
width,
TExp int
height,
TExp int
mulx,
TExp int
muly
]
)
forall a. [a] -> [a] -> [a]
++ forall a b. (a -> b) -> [a] -> [b]
map VName -> KernelUse
MemoryUse [VName
odata, VName
idata]
name :: String
name =
case TransposeType
kind of
TransposeType
TransposeSmall -> String
desc forall a. [a] -> [a] -> [a]
++ String
"_small"
TransposeType
TransposeLowHeight -> String
desc forall a. [a] -> [a] -> [a]
++ String
"_low_height"
TransposeType
TransposeLowWidth -> String
desc forall a. [a] -> [a] -> [a]
++ String
"_low_width"
TransposeType
TransposeNormal -> String
desc