-- | In the context of this module, a "size" is any kind of tunable
-- (run-time) constant.
module Futhark.IR.GPU.Sizes
  ( SizeClass (..),
    sizeDefault,
    KernelPath,
    Count (..),
    NumBlocks,
    BlockSize,
    NumThreads,
  )
where

import Data.Int (Int64)
import Data.Traversable
import Futhark.IR.Prop.Names (FreeIn)
import Futhark.Transform.Substitute
import Futhark.Util.IntegralExp (IntegralExp)
import Futhark.Util.Pretty
import Language.Futhark.Core (Name)
import Prelude hiding (id, (.))

-- | An indication of which comparisons have been performed to get to
-- this point, as well as the result of each comparison.
type KernelPath = [(Name, Bool)]

-- | The class of some kind of configurable size.  Each class may
-- impose constraints on the valid values.
data SizeClass
  = -- | A threshold with an optional default.
    SizeThreshold KernelPath (Maybe Int64)
  | SizeThreadBlock
  | -- | The number of thread blocks.
    SizeGrid
  | SizeTile
  | SizeRegTile
  | -- | Likely not useful on its own, but querying the
    -- maximum can be handy.
    SizeSharedMemory
  | -- | A bespoke size with a default.
    SizeBespoke Name Int64
  | -- | Amount of registers available per threadblock. Mostly
    -- meaningful for querying the maximum.
    SizeRegisters
  | -- | Amount of L2 cache memory, in bytes. Mostly meaningful for
    -- querying the maximum.
    SizeCache
  deriving (SizeClass -> SizeClass -> Bool
(SizeClass -> SizeClass -> Bool)
-> (SizeClass -> SizeClass -> Bool) -> Eq SizeClass
forall a. (a -> a -> Bool) -> (a -> a -> Bool) -> Eq a
$c== :: SizeClass -> SizeClass -> Bool
== :: SizeClass -> SizeClass -> Bool
$c/= :: SizeClass -> SizeClass -> Bool
/= :: SizeClass -> SizeClass -> Bool
Eq, Eq SizeClass
Eq SizeClass =>
(SizeClass -> SizeClass -> Ordering)
-> (SizeClass -> SizeClass -> Bool)
-> (SizeClass -> SizeClass -> Bool)
-> (SizeClass -> SizeClass -> Bool)
-> (SizeClass -> SizeClass -> Bool)
-> (SizeClass -> SizeClass -> SizeClass)
-> (SizeClass -> SizeClass -> SizeClass)
-> Ord SizeClass
SizeClass -> SizeClass -> Bool
SizeClass -> SizeClass -> Ordering
SizeClass -> SizeClass -> SizeClass
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
$ccompare :: SizeClass -> SizeClass -> Ordering
compare :: SizeClass -> SizeClass -> Ordering
$c< :: SizeClass -> SizeClass -> Bool
< :: SizeClass -> SizeClass -> Bool
$c<= :: SizeClass -> SizeClass -> Bool
<= :: SizeClass -> SizeClass -> Bool
$c> :: SizeClass -> SizeClass -> Bool
> :: SizeClass -> SizeClass -> Bool
$c>= :: SizeClass -> SizeClass -> Bool
>= :: SizeClass -> SizeClass -> Bool
$cmax :: SizeClass -> SizeClass -> SizeClass
max :: SizeClass -> SizeClass -> SizeClass
$cmin :: SizeClass -> SizeClass -> SizeClass
min :: SizeClass -> SizeClass -> SizeClass
Ord, Int -> SizeClass -> ShowS
[SizeClass] -> ShowS
SizeClass -> String
(Int -> SizeClass -> ShowS)
-> (SizeClass -> String)
-> ([SizeClass] -> ShowS)
-> Show SizeClass
forall a.
(Int -> a -> ShowS) -> (a -> String) -> ([a] -> ShowS) -> Show a
$cshowsPrec :: Int -> SizeClass -> ShowS
showsPrec :: Int -> SizeClass -> ShowS
$cshow :: SizeClass -> String
show :: SizeClass -> String
$cshowList :: [SizeClass] -> ShowS
showList :: [SizeClass] -> ShowS
Show)

instance Pretty SizeClass where
  pretty :: forall ann. SizeClass -> Doc ann
pretty (SizeThreshold KernelPath
path Maybe Int64
def) =
    Doc ann
"threshold" Doc ann -> Doc ann -> Doc ann
forall a. Semigroup a => a -> a -> a
<> Doc ann -> Doc ann
forall ann. Doc ann -> Doc ann
parens (Doc ann
forall {ann}. Doc ann
def' Doc ann -> Doc ann -> Doc ann
forall a. Semigroup a => a -> a -> a
<> Doc ann
forall {ann}. Doc ann
comma Doc ann -> Doc ann -> Doc ann
forall ann. Doc ann -> Doc ann -> Doc ann
<+> [Doc ann] -> Doc ann
forall ann. [Doc ann] -> Doc ann
hsep (((Name, Bool) -> Doc ann) -> KernelPath -> [Doc ann]
forall a b. (a -> b) -> [a] -> [b]
map (Name, Bool) -> Doc ann
forall {a} {ann}. Pretty a => (a, Bool) -> Doc ann
pStep KernelPath
path))
    where
      pStep :: (a, Bool) -> Doc ann
pStep (a
v, Bool
True) = a -> Doc ann
forall ann. a -> Doc ann
forall a ann. Pretty a => a -> Doc ann
pretty a
v
      pStep (a
v, Bool
False) = Doc ann
"!" Doc ann -> Doc ann -> Doc ann
forall a. Semigroup a => a -> a -> a
<> a -> Doc ann
forall ann. a -> Doc ann
forall a ann. Pretty a => a -> Doc ann
pretty a
v
      def' :: Doc ann
def' = Doc ann -> (Int64 -> Doc ann) -> Maybe Int64 -> Doc ann
forall b a. b -> (a -> b) -> Maybe a -> b
maybe Doc ann
"def" Int64 -> Doc ann
forall ann. Int64 -> Doc ann
forall a ann. Pretty a => a -> Doc ann
pretty Maybe Int64
def
  pretty SizeClass
SizeThreadBlock = Doc ann
"thread_block_size"
  pretty SizeClass
SizeGrid = Doc ann
"grid_size"
  pretty SizeClass
SizeTile = Doc ann
"tile_size"
  pretty SizeClass
SizeRegTile = Doc ann
"reg_tile_size"
  pretty SizeClass
SizeSharedMemory = Doc ann
"shared_memory"
  pretty (SizeBespoke Name
k Int64
def) =
    Doc ann
"bespoke" Doc ann -> Doc ann -> Doc ann
forall a. Semigroup a => a -> a -> a
<> Doc ann -> Doc ann
forall ann. Doc ann -> Doc ann
parens (Name -> Doc ann
forall a ann. Pretty a => a -> Doc ann
forall ann. Name -> Doc ann
pretty Name
k Doc ann -> Doc ann -> Doc ann
forall a. Semigroup a => a -> a -> a
<> Doc ann
forall {ann}. Doc ann
comma Doc ann -> Doc ann -> Doc ann
forall ann. Doc ann -> Doc ann -> Doc ann
<+> Int64 -> Doc ann
forall ann. Int64 -> Doc ann
forall a ann. Pretty a => a -> Doc ann
pretty Int64
def)
  pretty SizeClass
SizeRegisters = Doc ann
"registers"
  pretty SizeClass
SizeCache = Doc ann
"cache"

-- | The default value for the size.  If 'Nothing', that means the backend gets to decide.
sizeDefault :: SizeClass -> Maybe Int64
sizeDefault :: SizeClass -> Maybe Int64
sizeDefault (SizeThreshold KernelPath
_ Maybe Int64
x) = Maybe Int64
x
sizeDefault (SizeBespoke Name
_ Int64
x) = Int64 -> Maybe Int64
forall a. a -> Maybe a
Just Int64
x
sizeDefault SizeClass
_ = Maybe Int64
forall a. Maybe a
Nothing

-- | A wrapper supporting a phantom type for indicating what we are counting.
newtype Count u e = Count {forall {k} (u :: k) e. Count u e -> e
unCount :: e}
  deriving (Count u e -> Count u e -> Bool
(Count u e -> Count u e -> Bool)
-> (Count u e -> Count u e -> Bool) -> Eq (Count u e)
forall a. (a -> a -> Bool) -> (a -> a -> Bool) -> Eq a
forall k (u :: k) e. Eq e => Count u e -> Count u e -> Bool
$c== :: forall k (u :: k) e. Eq e => Count u e -> Count u e -> Bool
== :: Count u e -> Count u e -> Bool
$c/= :: forall k (u :: k) e. Eq e => Count u e -> Count u e -> Bool
/= :: Count u e -> Count u e -> Bool
Eq, Eq (Count u e)
Eq (Count u e) =>
(Count u e -> Count u e -> Ordering)
-> (Count u e -> Count u e -> Bool)
-> (Count u e -> Count u e -> Bool)
-> (Count u e -> Count u e -> Bool)
-> (Count u e -> Count u e -> Bool)
-> (Count u e -> Count u e -> Count u e)
-> (Count u e -> Count u e -> Count u e)
-> Ord (Count u e)
Count u e -> Count u e -> Bool
Count u e -> Count u e -> Ordering
Count u e -> Count u e -> Count u e
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
forall k (u :: k) e. Ord e => Eq (Count u e)
forall k (u :: k) e. Ord e => Count u e -> Count u e -> Bool
forall k (u :: k) e. Ord e => Count u e -> Count u e -> Ordering
forall k (u :: k) e. Ord e => Count u e -> Count u e -> Count u e
$ccompare :: forall k (u :: k) e. Ord e => Count u e -> Count u e -> Ordering
compare :: Count u e -> Count u e -> Ordering
$c< :: forall k (u :: k) e. Ord e => Count u e -> Count u e -> Bool
< :: Count u e -> Count u e -> Bool
$c<= :: forall k (u :: k) e. Ord e => Count u e -> Count u e -> Bool
<= :: Count u e -> Count u e -> Bool
$c> :: forall k (u :: k) e. Ord e => Count u e -> Count u e -> Bool
> :: Count u e -> Count u e -> Bool
$c>= :: forall k (u :: k) e. Ord e => Count u e -> Count u e -> Bool
>= :: Count u e -> Count u e -> Bool
$cmax :: forall k (u :: k) e. Ord e => Count u e -> Count u e -> Count u e
max :: Count u e -> Count u e -> Count u e
$cmin :: forall k (u :: k) e. Ord e => Count u e -> Count u e -> Count u e
min :: Count u e -> Count u e -> Count u e
Ord, Int -> Count u e -> ShowS
[Count u e] -> ShowS
Count u e -> String
(Int -> Count u e -> ShowS)
-> (Count u e -> String)
-> ([Count u e] -> ShowS)
-> Show (Count u e)
forall a.
(Int -> a -> ShowS) -> (a -> String) -> ([a] -> ShowS) -> Show a
forall k (u :: k) e. Show e => Int -> Count u e -> ShowS
forall k (u :: k) e. Show e => [Count u e] -> ShowS
forall k (u :: k) e. Show e => Count u e -> String
$cshowsPrec :: forall k (u :: k) e. Show e => Int -> Count u e -> ShowS
showsPrec :: Int -> Count u e -> ShowS
$cshow :: forall k (u :: k) e. Show e => Count u e -> String
show :: Count u e -> String
$cshowList :: forall k (u :: k) e. Show e => [Count u e] -> ShowS
showList :: [Count u e] -> ShowS
Show, Integer -> Count u e
Count u e -> Count u e
Count u e -> Count u e -> Count u e
(Count u e -> Count u e -> Count u e)
-> (Count u e -> Count u e -> Count u e)
-> (Count u e -> Count u e -> Count u e)
-> (Count u e -> Count u e)
-> (Count u e -> Count u e)
-> (Count u e -> Count u e)
-> (Integer -> Count u e)
-> Num (Count u e)
forall a.
(a -> a -> a)
-> (a -> a -> a)
-> (a -> a -> a)
-> (a -> a)
-> (a -> a)
-> (a -> a)
-> (Integer -> a)
-> Num a
forall k (u :: k) e. Num e => Integer -> Count u e
forall k (u :: k) e. Num e => Count u e -> Count u e
forall k (u :: k) e. Num e => Count u e -> Count u e -> Count u e
$c+ :: forall k (u :: k) e. Num e => Count u e -> Count u e -> Count u e
+ :: Count u e -> Count u e -> Count u e
$c- :: forall k (u :: k) e. Num e => Count u e -> Count u e -> Count u e
- :: Count u e -> Count u e -> Count u e
$c* :: forall k (u :: k) e. Num e => Count u e -> Count u e -> Count u e
* :: Count u e -> Count u e -> Count u e
$cnegate :: forall k (u :: k) e. Num e => Count u e -> Count u e
negate :: Count u e -> Count u e
$cabs :: forall k (u :: k) e. Num e => Count u e -> Count u e
abs :: Count u e -> Count u e
$csignum :: forall k (u :: k) e. Num e => Count u e -> Count u e
signum :: Count u e -> Count u e
$cfromInteger :: forall k (u :: k) e. Num e => Integer -> Count u e
fromInteger :: Integer -> Count u e
Num, Num (Count u e)
Num (Count u e) =>
(Count u e -> Count u e -> Count u e)
-> (Count u e -> Count u e -> Count u e)
-> (Count u e -> Count u e -> Count u e)
-> (Count u e -> Count u e -> Count u e)
-> (Count u e -> Maybe Int)
-> (Count u e -> Count u e -> Count u e)
-> (Count u e -> Count u e -> Count u e)
-> (Count u e -> Count u e -> Count u e)
-> IntegralExp (Count u e)
Count u e -> Maybe Int
Count u e -> Count u e -> Count u e
forall e.
Num e =>
(e -> e -> e)
-> (e -> e -> e)
-> (e -> e -> e)
-> (e -> e -> e)
-> (e -> Maybe Int)
-> (e -> e -> e)
-> (e -> e -> e)
-> (e -> e -> e)
-> IntegralExp e
forall k (u :: k) e. IntegralExp e => Num (Count u e)
forall k (u :: k) e. IntegralExp e => Count u e -> Maybe Int
forall k (u :: k) e.
IntegralExp e =>
Count u e -> Count u e -> Count u e
$cquot :: forall k (u :: k) e.
IntegralExp e =>
Count u e -> Count u e -> Count u e
quot :: Count u e -> Count u e -> Count u e
$crem :: forall k (u :: k) e.
IntegralExp e =>
Count u e -> Count u e -> Count u e
rem :: Count u e -> Count u e -> Count u e
$cdiv :: forall k (u :: k) e.
IntegralExp e =>
Count u e -> Count u e -> Count u e
div :: Count u e -> Count u e -> Count u e
$cmod :: forall k (u :: k) e.
IntegralExp e =>
Count u e -> Count u e -> Count u e
mod :: Count u e -> Count u e -> Count u e
$csgn :: forall k (u :: k) e. IntegralExp e => Count u e -> Maybe Int
sgn :: Count u e -> Maybe Int
$cpow :: forall k (u :: k) e.
IntegralExp e =>
Count u e -> Count u e -> Count u e
pow :: Count u e -> Count u e -> Count u e
$cdivUp :: forall k (u :: k) e.
IntegralExp e =>
Count u e -> Count u e -> Count u e
divUp :: Count u e -> Count u e -> Count u e
$cnextMul :: forall k (u :: k) e.
IntegralExp e =>
Count u e -> Count u e -> Count u e
nextMul :: Count u e -> Count u e -> Count u e
IntegralExp, Count u e -> FV
(Count u e -> FV) -> FreeIn (Count u e)
forall a. (a -> FV) -> FreeIn a
forall k (u :: k) e. FreeIn e => Count u e -> FV
$cfreeIn' :: forall k (u :: k) e. FreeIn e => Count u e -> FV
freeIn' :: Count u e -> FV
FreeIn, (forall ann. Count u e -> Doc ann)
-> (forall ann. [Count u e] -> Doc ann) -> Pretty (Count u e)
forall ann. [Count u e] -> Doc ann
forall ann. Count u e -> Doc ann
forall a.
(forall ann. a -> Doc ann)
-> (forall ann. [a] -> Doc ann) -> Pretty a
forall k (u :: k) e ann. Pretty e => [Count u e] -> Doc ann
forall k (u :: k) e ann. Pretty e => Count u e -> Doc ann
$cpretty :: forall k (u :: k) e ann. Pretty e => Count u e -> Doc ann
pretty :: forall ann. Count u e -> Doc ann
$cprettyList :: forall k (u :: k) e ann. Pretty e => [Count u e] -> Doc ann
prettyList :: forall ann. [Count u e] -> Doc ann
Pretty, Map VName VName -> Count u e -> Count u e
(Map VName VName -> Count u e -> Count u e)
-> Substitute (Count u e)
forall a. (Map VName VName -> a -> a) -> Substitute a
forall k (u :: k) e.
Substitute e =>
Map VName VName -> Count u e -> Count u e
$csubstituteNames :: forall k (u :: k) e.
Substitute e =>
Map VName VName -> Count u e -> Count u e
substituteNames :: Map VName VName -> Count u e -> Count u e
Substitute)

instance Functor (Count u) where
  fmap :: forall a b. (a -> b) -> Count u a -> Count u b
fmap = (a -> b) -> Count u a -> Count u b
forall (t :: * -> *) a b. Traversable t => (a -> b) -> t a -> t b
fmapDefault

instance Foldable (Count u) where
  foldMap :: forall m a. Monoid m => (a -> m) -> Count u a -> m
foldMap = (a -> m) -> Count u a -> m
forall (t :: * -> *) m a.
(Traversable t, Monoid m) =>
(a -> m) -> t a -> m
foldMapDefault

instance Traversable (Count u) where
  traverse :: forall (f :: * -> *) a b.
Applicative f =>
(a -> f b) -> Count u a -> f (Count u b)
traverse a -> f b
f (Count a
x) = b -> Count u b
forall {k} (u :: k) e. e -> Count u e
Count (b -> Count u b) -> f b -> f (Count u b)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> a -> f b
f a
x

-- | Phantom type for the number of blocks of some kernel.
data NumBlocks

-- | Phantom type for the block size of some kernel.
data BlockSize

-- | Phantom type for number of threads.
data NumThreads