module Futhark.CodeGen.ImpGen.GPU.Transpose
( TransposeType (..),
TransposeArgs,
mapTransposeKernel,
)
where
import Futhark.CodeGen.ImpCode.GPU
import Futhark.IR.Prop.Types
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 =
( VName,
TExp Int32,
VName,
TExp Int32,
TExp Int32,
TExp Int32,
TExp Int32,
TExp Int32,
TExp Int32,
VName
)
elemsPerThread :: TExp Int32
elemsPerThread :: TExp Int32
elemsPerThread = TExp Int32
4
mapTranspose :: TExp Int32 -> TransposeArgs -> PrimType -> TransposeType -> KernelCode
mapTranspose :: TExp Int32
-> TransposeArgs -> PrimType -> TransposeType -> KernelCode
mapTranspose TExp Int32
block_dim TransposeArgs
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
$ forall a. a -> TPrimExp Int32 a
le32 VName
get_global_id_0 forall e. IntegralExp e => e -> e -> e
`quot` (TExp Int32
height forall a. Num a => a -> a -> a
* TExp Int32
width) forall a. Num a => a -> a -> a
* (TExp Int32
height forall a. Num a => a -> a -> a
* TExp Int32
width),
forall {k} {t :: k} {a}. VName -> TPrimExp t VName -> Code a
dec VName
x_index forall a b. (a -> b) -> a -> b
$ (forall a. a -> TPrimExp Int32 a
le32 VName
get_global_id_0 forall e. IntegralExp e => e -> e -> e
`rem` (TExp Int32
height forall a. Num a => a -> a -> a
* TExp Int32
width)) forall e. IntegralExp e => e -> e -> e
`quot` TExp Int32
height,
forall {k} {t :: k} {a}. VName -> TPrimExp t VName -> Code a
dec VName
y_index forall a b. (a -> b) -> a -> b
$ forall a. a -> TPrimExp Int32 a
le32 VName
get_global_id_0 forall e. IntegralExp e => e -> e -> e
`rem` TExp Int32
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
$
(TExp Int32
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
+ forall a. a -> TPrimExp Int32 a
le32 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
$
(TExp Int32
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
+ forall a. a -> TPrimExp Int32 a
le32 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
$ forall a. a -> TPrimExp Int32 a
le32 VName
y_index forall a. Num a => a -> a -> a
* TExp Int32
width forall a. Num a => a -> a -> a
+ forall a. a -> TPrimExp Int32 a
le32 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
$ forall a. a -> TPrimExp Int32 a
le32 VName
x_index forall a. Num a => a -> a -> a
* TExp Int32
height forall a. Num a => a -> a -> a
+ forall a. a -> TPrimExp Int32 a
le32 VName
y_index,
forall {a}. TExp Bool -> Code a -> Code a
when
(forall a. a -> TPrimExp Int32 a
le32 VName
get_global_id_0 forall {k} (t :: k) v.
TPrimExp t v -> TPrimExp t v -> TPrimExp Bool v
.<. TExp Int32
width forall a. Num a => a -> a -> a
* TExp Int32
height forall a. Num a => a -> a -> a
* TExp Int32
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
$ forall {k} (t :: k) v. IntExp t => TPrimExp t v -> TPrimExp Int64 v
sExt64 forall a b. (a -> b) -> a -> b
$ forall a. a -> TPrimExp Int32 a
le32 VName
idata_offset forall a. Num a => a -> a -> a
+ forall a. a -> TPrimExp Int32 a
le32 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
$ forall {k} (t :: k) v. IntExp t => TPrimExp t v -> TPrimExp Int64 v
sExt64 forall a b. (a -> b) -> a -> b
$ forall a. a -> TPrimExp Int32 a
le32 VName
odata_offset forall a. Num a => a -> a -> a
+ forall a. a -> TPrimExp Int32 a
le32 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
(forall a. a -> TPrimExp Int32 a
le32 VName
get_group_id_0 forall a. Num a => a -> a -> a
* TExp Int32
block_dim forall a. Num a => a -> a -> a
+ (forall a. a -> TPrimExp Int32 a
le32 VName
get_local_id_0 forall e. IntegralExp e => e -> e -> e
`quot` TExp Int32
muly))
( forall a. a -> TPrimExp Int32 a
le32 VName
get_group_id_1 forall a. Num a => a -> a -> a
* TExp Int32
block_dim forall a. Num a => a -> a -> a
* TExp Int32
muly
forall a. Num a => a -> a -> a
+ forall a. a -> TPrimExp Int32 a
le32 VName
get_local_id_1
forall a. Num a => a -> a -> a
+ (forall a. a -> TPrimExp Int32 a
le32 VName
get_local_id_0 forall e. IntegralExp e => e -> e -> e
`rem` TExp Int32
muly) forall a. Num a => a -> a -> a
* TExp Int32
block_dim
)
( forall a. a -> TPrimExp Int32 a
le32 VName
get_group_id_1 forall a. Num a => a -> a -> a
* TExp Int32
block_dim forall a. Num a => a -> a -> a
* TExp Int32
muly
forall a. Num a => a -> a -> a
+ forall a. a -> TPrimExp Int32 a
le32 VName
get_local_id_0
forall a. Num a => a -> a -> a
+ (forall a. a -> TPrimExp Int32 a
le32 VName
get_local_id_1 forall e. IntegralExp e => e -> e -> e
`rem` TExp Int32
muly) forall a. Num a => a -> a -> a
* TExp Int32
block_dim
)
(forall a. a -> TPrimExp Int32 a
le32 VName
get_group_id_0 forall a. Num a => a -> a -> a
* TExp Int32
block_dim forall a. Num a => a -> a -> a
+ (forall a. a -> TPrimExp Int32 a
le32 VName
get_local_id_1 forall e. IntegralExp e => e -> e -> e
`quot` TExp Int32
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
( forall a. a -> TPrimExp Int32 a
le32 VName
get_group_id_0 forall a. Num a => a -> a -> a
* TExp Int32
block_dim forall a. Num a => a -> a -> a
* TExp Int32
mulx
forall a. Num a => a -> a -> a
+ forall a. a -> TPrimExp Int32 a
le32 VName
get_local_id_0
forall a. Num a => a -> a -> a
+ (forall a. a -> TPrimExp Int32 a
le32 VName
get_local_id_1 forall e. IntegralExp e => e -> e -> e
`rem` TExp Int32
mulx) forall a. Num a => a -> a -> a
* TExp Int32
block_dim
)
(forall a. a -> TPrimExp Int32 a
le32 VName
get_group_id_1 forall a. Num a => a -> a -> a
* TExp Int32
block_dim forall a. Num a => a -> a -> a
+ (forall a. a -> TPrimExp Int32 a
le32 VName
get_local_id_1 forall e. IntegralExp e => e -> e -> e
`quot` TExp Int32
mulx))
(forall a. a -> TPrimExp Int32 a
le32 VName
get_group_id_1 forall a. Num a => a -> a -> a
* TExp Int32
block_dim forall a. Num a => a -> a -> a
+ (forall a. a -> TPrimExp Int32 a
le32 VName
get_local_id_0 forall e. IntegralExp e => e -> e -> e
`quot` TExp Int32
mulx))
( forall a. a -> TPrimExp Int32 a
le32 VName
get_group_id_0 forall a. Num a => a -> a -> a
* TExp Int32
block_dim forall a. Num a => a -> a -> a
* TExp Int32
mulx
forall a. Num a => a -> a -> a
+ forall a. a -> TPrimExp Int32 a
le32 VName
get_local_id_1
forall a. Num a => a -> a -> a
+ (forall a. a -> TPrimExp Int32 a
le32 VName
get_local_id_0 forall e. IntegralExp e => e -> e -> e
`rem` TExp Int32
mulx) forall a. Num a => a -> a -> a
* TExp Int32
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
$ forall a. a -> TPrimExp Int32 a
le32 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
$ forall a. a -> TPrimExp Int32 a
le32 VName
get_group_id_1 forall a. Num a => a -> a -> a
* TExp Int32
tile_dim forall a. Num a => a -> a -> a
+ forall a. a -> TPrimExp Int32 a
le32 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 (forall a. a -> TPrimExp Int32 a
le32 VName
x_index forall {k} (t :: k) v.
TPrimExp t v -> TPrimExp t v -> TPrimExp Bool v
.<. TExp Int32
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 TExp Int32
elemsPerThread) forall a b. (a -> b) -> a -> b
$
let i :: TExp Int32
i = forall a. a -> TPrimExp Int32 a
le32 VName
j forall a. Num a => a -> a -> a
* (TExp Int32
tile_dim forall e. IntegralExp e => e -> e -> e
`quot` TExp Int32
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
$ (forall a. a -> TPrimExp Int32 a
le32 VName
y_index forall a. Num a => a -> a -> a
+ TExp Int32
i) forall a. Num a => a -> a -> a
* TExp Int32
width forall a. Num a => a -> a -> a
+ forall a. a -> TPrimExp Int32 a
le32 VName
x_index,
forall {a}. TExp Bool -> Code a -> Code a
when (forall a. a -> TPrimExp Int32 a
le32 VName
y_index forall a. Num a => a -> a -> a
+ TExp Int32
i forall {k} (t :: k) v.
TPrimExp t v -> TPrimExp t v -> TPrimExp Bool v
.<. TExp Int32
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
$ forall a. a -> TPrimExp Int32 a
le32 VName
idata_offset forall a. Num a => a -> a -> a
+ forall a. a -> TPrimExp Int32 a
le32 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
$
(forall a. a -> TPrimExp Int32 a
le32 VName
get_local_id_1 forall a. Num a => a -> a -> a
+ TExp Int32
i) forall a. Num a => a -> a -> a
* (TExp Int32
tile_dim forall a. Num a => a -> a -> a
+ TExp Int32
1)
forall a. Num a => a -> a -> a
+ forall a. a -> TPrimExp Int32 a
le32 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
$ forall a. a -> TPrimExp Int32 a
le32 VName
get_group_id_1 forall a. Num a => a -> a -> a
* TExp Int32
tile_dim forall a. Num a => a -> a -> a
+ forall a. a -> TPrimExp Int32 a
le32 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
$ forall a. a -> TPrimExp Int32 a
le32 VName
get_group_id_0 forall a. Num a => a -> a -> a
* TExp Int32
tile_dim forall a. Num a => a -> a -> a
+ forall a. a -> TPrimExp Int32 a
le32 VName
get_local_id_1,
forall {a}. TExp Bool -> Code a -> Code a
when (forall a. a -> TPrimExp Int32 a
le32 VName
x_index forall {k} (t :: k) v.
TPrimExp t v -> TPrimExp t v -> TPrimExp Bool v
.<. TExp Int32
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 TExp Int32
elemsPerThread) forall a b. (a -> b) -> a -> b
$
let i :: TExp Int32
i = forall a. a -> TPrimExp Int32 a
le32 VName
j forall a. Num a => a -> a -> a
* (TExp Int32
tile_dim forall e. IntegralExp e => e -> e -> e
`quot` TExp Int32
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
$ (forall a. a -> TPrimExp Int32 a
le32 VName
y_index forall a. Num a => a -> a -> a
+ TExp Int32
i) forall a. Num a => a -> a -> a
* TExp Int32
height forall a. Num a => a -> a -> a
+ forall a. a -> TPrimExp Int32 a
le32 VName
x_index,
forall {a}. TExp Bool -> Code a -> Code a
when (forall a. a -> TPrimExp Int32 a
le32 VName
y_index forall a. Num a => a -> a -> a
+ TExp Int32
i forall {k} (t :: k) v.
TPrimExp t v -> TPrimExp t v -> TPrimExp Bool v
.<. TExp Int32
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
. forall {k} (t :: k) v. IntExp t => TPrimExp t v -> TPrimExp Int64 v
sExt64 forall a b. (a -> b) -> a -> b
$
forall a. a -> TPrimExp Int32 a
le32 VName
get_local_id_0 forall a. Num a => a -> a -> a
* (TExp Int32
tile_dim forall a. Num a => a -> a -> a
+ TExp Int32
1) forall a. Num a => a -> a -> a
+ forall a. a -> TPrimExp Int32 a
le32 VName
get_local_id_1 forall a. Num a => a -> a -> a
+ TExp Int32
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
$ forall {k} (t :: k) v. IntExp t => TPrimExp t v -> TPrimExp Int64 v
sExt64 forall a b. (a -> b) -> a -> b
$ forall a. a -> TPrimExp Int32 a
le32 VName
odata_offset forall a. Num a => a -> a -> a
+ forall a. a -> TPrimExp Int32 a
le32 VName
index_out)
PrimType
t
(String -> Space
Space String
"global")
Volatility
Nonvolatile
(VName -> PrimType -> Exp
var VName
val PrimType
t)
]
]
]
where
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 :: TExp Int32
tile_dim = TExp Int32
2 forall a. Num a => a -> a -> a
* TExp Int32
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,
TExp Int32
basic_odata_offset,
VName
idata,
TExp Int32
basic_idata_offset,
TExp Int32
width,
TExp Int32
height,
TExp Int32
mulx,
TExp Int32
muly,
TExp Int32
num_arrays,
VName
block
) = TransposeArgs
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
int32,
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
int32,
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
int32,
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
int32,
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
int32,
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
int32,
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
int32,
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
$ forall a. a -> TPrimExp Int32 a
le32 VName
get_group_id_0 forall a. Num a => a -> a -> a
* forall a. a -> TPrimExp Int32 a
le32 VName
get_local_size_0 forall a. Num a => a -> a -> a
+ forall a. a -> TPrimExp Int32 a
le32 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
$ forall a. a -> TPrimExp Int32 a
le32 VName
get_group_id_2 forall a. Num a => a -> a -> a
* TExp Int32
width forall a. Num a => a -> a -> a
* TExp Int32
height,
forall {k} {t :: k} {a}. VName -> TPrimExp t VName -> Code a
dec VName
odata_offset forall a b. (a -> b) -> a -> b
$
(TExp Int32
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
+ forall a. a -> TPrimExp Int32 a
le32 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
$
(TExp Int32
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
+ forall a. a -> TPrimExp Int32 a
le32 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
$ forall a. a -> TPrimExp Int32 a
le32 VName
y_index forall a. Num a => a -> a -> a
* TExp Int32
width forall a. Num a => a -> a -> a
+ forall a. a -> TPrimExp Int32 a
le32 VName
x_index,
forall {a}. TExp Bool -> Code a -> Code a
when (forall a. a -> TPrimExp Int32 a
le32 VName
x_index forall {k} (t :: k) v.
TPrimExp t v -> TPrimExp t v -> TPrimExp Bool v
.<. TExp Int32
width forall v. TPrimExp Bool v -> TPrimExp Bool v -> TPrimExp Bool v
.&&. forall a. a -> TPrimExp Int32 a
le32 VName
y_index forall {k} (t :: k) v.
TPrimExp t v -> TPrimExp t v -> TPrimExp Bool v
.<. TExp Int32
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
$ forall a. a -> TPrimExp Int32 a
le32 VName
idata_offset forall a. Num a => a -> a -> a
+ forall a. a -> TPrimExp Int32 a
le32 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
$ forall a. a -> TPrimExp Int32 a
le32 VName
get_local_id_1 forall a. Num a => a -> a -> a
* (TExp Int32
block_dim forall a. Num a => a -> a -> a
+ TExp Int32
1) forall a. Num a => a -> a -> a
+ forall a. a -> TPrimExp Int32 a
le32 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
$ forall a. a -> TPrimExp Int32 a
le32 VName
y_index forall a. Num a => a -> a -> a
* TExp Int32
height forall a. Num a => a -> a -> a
+ forall a. a -> TPrimExp Int32 a
le32 VName
x_index,
forall {a}. TExp Bool -> Code a -> Code a
when (forall a. a -> TPrimExp Int32 a
le32 VName
x_index forall {k} (t :: k) v.
TPrimExp t v -> TPrimExp t v -> TPrimExp Bool v
.<. TExp Int32
height forall v. TPrimExp Bool v -> TPrimExp Bool v -> TPrimExp Bool v
.&&. forall a. a -> TPrimExp Int32 a
le32 VName
y_index forall {k} (t :: k) v.
TPrimExp t v -> TPrimExp t v -> TPrimExp Bool v
.<. TExp Int32
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
$ forall {k} (t :: k) v. IntExp t => TPrimExp t v -> TPrimExp Int64 v
sExt64 forall a b. (a -> b) -> a -> b
$ forall a. a -> TPrimExp Int32 a
le32 VName
get_local_id_0 forall a. Num a => a -> a -> a
* (TExp Int32
block_dim forall a. Num a => a -> a -> a
+ TExp Int32
1) forall a. Num a => a -> a -> a
+ forall a. a -> TPrimExp Int32 a
le32 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
$ forall {k} (t :: k) v. IntExp t => TPrimExp t v -> TPrimExp Int64 v
sExt64 (forall a. a -> TPrimExp Int32 a
le32 VName
odata_offset forall a. Num a => a -> a -> a
+ forall a. a -> TPrimExp Int32 a
le32 VName
index_out))
PrimType
t
(String -> Space
Space String
"global")
Volatility
Nonvolatile
(VName -> PrimType -> Exp
var VName
val PrimType
t)
]
]
mapTransposeKernel ::
String ->
Integer ->
TransposeArgs ->
PrimType ->
TransposeType ->
Kernel
mapTransposeKernel :: String
-> Integer -> TransposeArgs -> PrimType -> TransposeType -> Kernel
mapTransposeKernel String
desc Integer
block_dim_int TransposeArgs
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
<> TExp Int32
-> TransposeArgs -> PrimType -> TransposeType -> KernelCode
mapTranspose TExp Int32
block_dim TransposeArgs
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 Int32]
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 Int32]
group_size,
kernelName :: Name
kernelName = String -> Name
nameFromString String
name,
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 =
forall a. a -> Count Bytes a
bytes forall a b. (a -> b) -> a -> b
$
case TransposeType
kind of
TransposeType
TransposeSmall -> TExp Int64
1 :: TExp Int64
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 Int32
block_dim = forall a. Num a => Integer -> a
fromInteger Integer
block_dim_int :: TExp Int32
( VName
odata,
TExp Int32
basic_odata_offset,
VName
idata,
TExp Int32
basic_idata_offset,
TExp Int32
width,
TExp Int32
height,
TExp Int32
mulx,
TExp Int32
muly,
TExp Int32
num_arrays,
VName
block
) = TransposeArgs
args
([TExp Int32]
num_groups, [TExp Int32]
group_size) =
case TransposeType
kind of
TransposeType
TransposeSmall ->
( [(TExp Int32
num_arrays forall a. Num a => a -> a -> a
* TExp Int32
width forall a. Num a => a -> a -> a
* TExp Int32
height) forall e. IntegralExp e => e -> e -> e
`divUp` (TExp Int32
block_dim forall a. Num a => a -> a -> a
* TExp Int32
block_dim)],
[TExp Int32
block_dim forall a. Num a => a -> a -> a
* TExp Int32
block_dim]
)
TransposeType
TransposeLowWidth ->
TExp Int32
-> TExp Int32
-> TExp Int32
-> TExp Int32
-> ([TExp Int32], [TExp Int32])
lowDimKernelAndGroupSize TExp Int32
block_dim TExp Int32
num_arrays TExp Int32
width forall a b. (a -> b) -> a -> b
$ TExp Int32
height forall e. IntegralExp e => e -> e -> e
`divUp` TExp Int32
muly
TransposeType
TransposeLowHeight ->
TExp Int32
-> TExp Int32
-> TExp Int32
-> TExp Int32
-> ([TExp Int32], [TExp Int32])
lowDimKernelAndGroupSize TExp Int32
block_dim TExp Int32
num_arrays (TExp Int32
width forall e. IntegralExp e => e -> e -> e
`divUp` TExp Int32
mulx) TExp Int32
height
TransposeType
TransposeNormal ->
let actual_dim :: TExp Int32
actual_dim = TExp Int32
block_dim forall a. Num a => a -> a -> a
* TExp Int32
2
in ( [ TExp Int32
width forall e. IntegralExp e => e -> e -> e
`divUp` TExp Int32
actual_dim,
TExp Int32
height forall e. IntegralExp e => e -> e -> e
`divUp` TExp Int32
actual_dim,
TExp Int32
num_arrays
],
[TExp Int32
actual_dim, TExp Int32
actual_dim forall e. IntegralExp e => e -> e -> e
`quot` TExp Int32
elemsPerThread, TExp Int32
1]
)
uses :: [KernelUse]
uses =
forall a b. (a -> b) -> [a] -> [b]
map
(VName -> PrimType -> KernelUse
`ScalarUse` PrimType
int32)
( 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 Int32
basic_odata_offset,
TExp Int32
basic_idata_offset,
TExp Int32
num_arrays,
TExp Int32
width,
TExp Int32
height,
TExp Int32
mulx,
TExp Int32
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
lowDimKernelAndGroupSize ::
TExp Int32 ->
TExp Int32 ->
TExp Int32 ->
TExp Int32 ->
([TExp Int32], [TExp Int32])
lowDimKernelAndGroupSize :: TExp Int32
-> TExp Int32
-> TExp Int32
-> TExp Int32
-> ([TExp Int32], [TExp Int32])
lowDimKernelAndGroupSize TExp Int32
block_dim TExp Int32
num_arrays TExp Int32
x_elems TExp Int32
y_elems =
( [ TExp Int32
x_elems forall e. IntegralExp e => e -> e -> e
`divUp` TExp Int32
block_dim,
TExp Int32
y_elems forall e. IntegralExp e => e -> e -> e
`divUp` TExp Int32
block_dim,
TExp Int32
num_arrays
],
[TExp Int32
block_dim, TExp Int32
block_dim, TExp Int32
1]
)