-- | Carefully optimised implementations of GPU transpositions.
-- Written in ImpCode so we can compile it to both CUDA and OpenCL.
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)

-- | Which form of transposition to generate code for.
data TransposeType
  = TransposeNormal
  | TransposeLowWidth
  | TransposeLowHeight
  | -- | For small arrays that do not
    -- benefit from coalescing.
    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)

-- | The types of the arguments accepted by a transposition function.
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

    -- Be extremely careful when editing this list to ensure that
    -- the names match up.  Also, be careful that the tags on
    -- these names do not conflict with the tags of the
    -- surrounding code.  We accomplish the latter by using very
    -- low tags (normal variables start at least in the low
    -- hundreds).
    [ 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)
              ]
        ]

-- | Generate a transpose kernel.  There is special support to handle
-- input arrays with low width, low height, or both.
--
-- Normally when transposing a @[2][n]@ array we would use a @FUT_BLOCK_DIM x
-- FUT_BLOCK_DIM@ group to process a @[2][FUT_BLOCK_DIM]@ slice of the input
-- array. This would mean that many of the threads in a group would be inactive.
-- We try to remedy this by using a special kernel that will process a larger
-- part of the input, by using more complex indexing. In our example, we could
-- use all threads in a group if we are processing @(2/FUT_BLOCK_DIM)@ as large
-- a slice of each rows per group. The variable @mulx@ contains this factor for
-- the kernel to handle input arrays with low height.
--
-- See issue #308 on GitHub for more details.
--
-- These kernels are optimized to ensure all global reads and writes
-- are coalesced, and to avoid bank conflicts in shared memory.  Each
-- thread group transposes a 2D tile of block_dim*2 by block_dim*2
-- elements. The size of a thread group is block_dim/2 by
-- block_dim*2, meaning that each thread will process 4 elements in a
-- 2D tile.  The shared memory array containing the 2D tile consists
-- of block_dim*2 by block_dim*2+1 elements. Padding each row with
-- an additional element prevents bank conflicts from occuring when
-- the tile is accessed column-wise.
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 :: [Exp]
kernelGroupSize = forall a b. (a -> b) -> [a] -> [b]
map 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
          -- Not used, but AMD's OpenCL
          -- does not like zero-size
          -- local memory.
          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]
  )