GPU.hs

-- | Variation of "Futhark.CodeGen.ImpCode" that contains the notion
-- of a kernel invocation.
module Futhark.CodeGen.ImpCode.GPU
  ( Program,
    HostCode,
    KernelCode,
    KernelConst (..),
    KernelConstExp,
    HostOp (..),
    KernelOp (..),
    Fence (..),
    AtomicOp (..),
    BlockDim,
    Kernel (..),
    KernelUse (..),
    module Futhark.CodeGen.ImpCode,
    module Futhark.IR.GPU.Sizes,
  )
where

import Futhark.CodeGen.ImpCode
import Futhark.IR.GPU.Sizes
import Futhark.IR.Pretty ()
import Futhark.Util.Pretty

-- | A program that calls kernels.
type Program = Definitions HostOp

-- | Host-level code that can call kernels.
type HostCode = Code HostOp

-- | Code inside a kernel.
type KernelCode = Code KernelOp

-- | A run-time constant related to kernels.
data KernelConst
  = SizeConst Name SizeClass
  | SizeMaxConst SizeClass
  deriving (Eq, Ord, Show)

-- | An expression whose variables are kernel constants.
type KernelConstExp = PrimExp KernelConst

-- | An operation that runs on the host (CPU).
data HostOp
  = CallKernel Kernel
  | GetSize VName Name SizeClass
  | CmpSizeLe VName Name SizeClass Exp
  | GetSizeMax VName SizeClass
  deriving (Show)

-- | The size of one dimension of a block.
type BlockDim = Either Exp KernelConstExp

-- | A generic kernel containing arbitrary kernel code.
data Kernel = Kernel
  { kernelBody :: Code KernelOp,
    -- | The host variables referenced by the kernel.
    kernelUses :: [KernelUse],
    kernelNumBlocks :: [Exp],
    kernelBlockSize :: [BlockDim],
    -- | A short descriptive and _unique_ name - should be
    -- alphanumeric and without spaces.
    kernelName :: Name,
    -- | If true, this kernel does not need to check whether we are in
    -- a failing state, as it can cope. Intuitively, it means that the
    -- kernel does not depend on any non-scalar parameters to make
    -- control flow decisions. Replication, transpose, and copy
    -- kernels are examples of this.
    kernelFailureTolerant :: Bool,
    -- | If true, multi-versioning branches will consider this kernel
    -- when considering the shared memory requirements. Set this to
    -- false for kernels that do their own checking.
    kernelCheckSharedMemory :: Bool
  }
  deriving (Show)

-- | Information about a host-level variable that is used inside this
-- kernel.  When generating the actual kernel code, this is used to
-- deduce which parameters are needed.
data KernelUse
  = ScalarUse VName PrimType
  | MemoryUse VName
  | ConstUse VName KernelConstExp
  deriving (Eq, Ord, Show)

instance Pretty KernelConst where
  pretty (SizeConst key size_class) =
    "get_size" <> parens (commasep [pretty key, pretty size_class])
  pretty (SizeMaxConst size_class) =
    "get_max_size" <> parens (pretty size_class)

instance FreeIn KernelConst where
  freeIn' SizeConst {} = mempty
  freeIn' (SizeMaxConst _) = mempty

instance Pretty KernelUse where
  pretty (ScalarUse name t) =
    oneLine $ "scalar_copy" <> parens (commasep [pretty name, pretty t])
  pretty (MemoryUse name) =
    oneLine $ "mem_copy" <> parens (commasep [pretty name])
  pretty (ConstUse name e) =
    oneLine $ "const" <> parens (commasep [pretty name, pretty e])

instance Pretty HostOp where
  pretty (GetSize dest key size_class) =
    pretty dest
      <+> "<-"
      <+> "get_size"
      <> parens (commasep [pretty key, pretty size_class])
  pretty (GetSizeMax dest size_class) =
    pretty dest <+> "<-" <+> "get_size_max" <> parens (pretty size_class)
  pretty (CmpSizeLe dest name size_class x) =
    pretty dest
      <+> "<-"
      <+> "get_size"
      <> parens (commasep [pretty name, pretty size_class])
        <+> "<"
        <+> pretty x
  pretty (CallKernel c) =
    pretty c

instance FreeIn HostOp where
  freeIn' (CallKernel c) =
    freeIn' c
  freeIn' (CmpSizeLe dest _ _ x) =
    freeIn' dest <> freeIn' x
  freeIn' (GetSizeMax dest _) =
    freeIn' dest
  freeIn' (GetSize dest _ _) =
    freeIn' dest

instance FreeIn Kernel where
  freeIn' kernel =
    freeIn'
      ( kernelBody kernel,
        kernelNumBlocks kernel,
        kernelBlockSize kernel
      )

instance Pretty Kernel where
  pretty kernel =
    "kernel"
      <+> brace
        ( "blocks"
            <+> brace (pretty $ kernelNumBlocks kernel)
            </> "tblock_size"
            <+> brace (list $ map pSize $ kernelBlockSize kernel)
            </> "uses"
            <+> brace (stack $ map pretty $ kernelUses kernel)
            </> "failure_tolerant"
            <+> brace (pretty $ kernelFailureTolerant kernel)
            </> "check_shared_memory"
            <+> brace (pretty $ kernelCheckSharedMemory kernel)
            </> "body"
            <+> brace (pretty $ kernelBody kernel)
        )
    where
      pSize (Left x) = "dyn" <+> pretty x
      pSize (Right x) = "const" <+> pretty x

-- | When we do a barrier or fence, is it at the local or global
-- level?  By the 'Ord' instance, global is greater than local.
data Fence = FenceLocal | FenceGlobal
  deriving (Show, Eq, Ord)

-- | An operation that occurs within a kernel body.
data KernelOp
  = GetBlockId VName Int
  | GetLocalId VName Int
  | GetLocalSize VName Int
  | GetLockstepWidth VName
  | Atomic Space AtomicOp
  | Barrier Fence
  | MemFence Fence
  | SharedAlloc VName (Count Bytes (TExp Int64))
  | -- | Perform a barrier and also check whether any
    -- threads have failed an assertion.  Make sure all
    -- threads would reach all 'ErrorSync's if any of them
    -- do.  A failing assertion will jump to the next
    -- following 'ErrorSync', so make sure it's not inside
    -- control flow or similar.
    ErrorSync Fence
  deriving (Show)

-- | Atomic operations return the value stored before the update. This
-- old value is stored in the first 'VName' (except for
-- 'AtomicWrite'). The second 'VName' is the memory block to update.
-- The 'Exp' is the new value.
data AtomicOp
  = AtomicAdd IntType VName VName (Count Elements (TExp Int64)) Exp
  | AtomicFAdd FloatType VName VName (Count Elements (TExp Int64)) Exp
  | AtomicSMax IntType VName VName (Count Elements (TExp Int64)) Exp
  | AtomicSMin IntType VName VName (Count Elements (TExp Int64)) Exp
  | AtomicUMax IntType VName VName (Count Elements (TExp Int64)) Exp
  | AtomicUMin IntType VName VName (Count Elements (TExp Int64)) Exp
  | AtomicAnd IntType VName VName (Count Elements (TExp Int64)) Exp
  | AtomicOr IntType VName VName (Count Elements (TExp Int64)) Exp
  | AtomicXor IntType VName VName (Count Elements (TExp Int64)) Exp
  | AtomicCmpXchg PrimType VName VName (Count Elements (TExp Int64)) Exp Exp
  | AtomicXchg PrimType VName VName (Count Elements (TExp Int64)) Exp
  | -- | Corresponds to a write followed by a memory fence. The old
    -- value is not read.
    AtomicWrite PrimType VName (Count Elements (TExp Int64)) Exp
  deriving (Show)

instance FreeIn AtomicOp where
  freeIn' (AtomicAdd _ _ arr i x) = freeIn' arr <> freeIn' i <> freeIn' x
  freeIn' (AtomicFAdd _ _ arr i x) = freeIn' arr <> freeIn' i <> freeIn' x
  freeIn' (AtomicSMax _ _ arr i x) = freeIn' arr <> freeIn' i <> freeIn' x
  freeIn' (AtomicSMin _ _ arr i x) = freeIn' arr <> freeIn' i <> freeIn' x
  freeIn' (AtomicUMax _ _ arr i x) = freeIn' arr <> freeIn' i <> freeIn' x
  freeIn' (AtomicUMin _ _ arr i x) = freeIn' arr <> freeIn' i <> freeIn' x
  freeIn' (AtomicAnd _ _ arr i x) = freeIn' arr <> freeIn' i <> freeIn' x
  freeIn' (AtomicOr _ _ arr i x) = freeIn' arr <> freeIn' i <> freeIn' x
  freeIn' (AtomicXor _ _ arr i x) = freeIn' arr <> freeIn' i <> freeIn' x
  freeIn' (AtomicCmpXchg _ _ arr i x y) = freeIn' arr <> freeIn' i <> freeIn' x <> freeIn' y
  freeIn' (AtomicXchg _ _ arr i x) = freeIn' arr <> freeIn' i <> freeIn' x
  freeIn' (AtomicWrite _ arr i x) = freeIn' arr <> freeIn' i <> freeIn' x

instance Pretty KernelOp where
  pretty (GetBlockId dest i) =
    pretty dest
      <+> "<-"
      <+> "get_tblock_id"
      <> parens (pretty i)
  pretty (GetLocalId dest i) =
    pretty dest
      <+> "<-"
      <+> "get_local_id"
      <> parens (pretty i)
  pretty (GetLocalSize dest i) =
    pretty dest
      <+> "<-"
      <+> "get_local_size"
      <> parens (pretty i)
  pretty (GetLockstepWidth dest) =
    pretty dest
      <+> "<-"
      <+> "get_lockstep_width()"
  pretty (Barrier FenceLocal) =
    "local_barrier()"
  pretty (Barrier FenceGlobal) =
    "global_barrier()"
  pretty (MemFence FenceLocal) =
    "mem_fence_local()"
  pretty (MemFence FenceGlobal) =
    "mem_fence_global()"
  pretty (SharedAlloc name size) =
    pretty name <+> equals <+> "shared_alloc" <> parens (pretty size)
  pretty (ErrorSync FenceLocal) =
    "error_sync_local()"
  pretty (ErrorSync FenceGlobal) =
    "error_sync_global()"
  pretty (Atomic _ (AtomicAdd t old arr ind x)) =
    pretty old
      <+> "<-"
      <+> "atomic_add_"
      <> pretty t
      <> parens (commasep [pretty arr <> brackets (pretty ind), pretty x])
  pretty (Atomic _ (AtomicFAdd t old arr ind x)) =
    pretty old
      <+> "<-"
      <+> "atomic_fadd_"
      <> pretty t
      <> parens (commasep [pretty arr <> brackets (pretty ind), pretty x])
  pretty (Atomic _ (AtomicSMax t old arr ind x)) =
    pretty old
      <+> "<-"
      <+> "atomic_smax"
      <> pretty t
      <> parens (commasep [pretty arr <> brackets (pretty ind), pretty x])
  pretty (Atomic _ (AtomicSMin t old arr ind x)) =
    pretty old
      <+> "<-"
      <+> "atomic_smin"
      <> pretty t
      <> parens (commasep [pretty arr <> brackets (pretty ind), pretty x])
  pretty (Atomic _ (AtomicUMax t old arr ind x)) =
    pretty old
      <+> "<-"
      <+> "atomic_umax"
      <> pretty t
      <> parens (commasep [pretty arr <> brackets (pretty ind), pretty x])
  pretty (Atomic _ (AtomicUMin t old arr ind x)) =
    pretty old
      <+> "<-"
      <+> "atomic_umin"
      <> pretty t
      <> parens (commasep [pretty arr <> brackets (pretty ind), pretty x])
  pretty (Atomic _ (AtomicAnd t old arr ind x)) =
    pretty old
      <+> "<-"
      <+> "atomic_and"
      <> pretty t
      <> parens (commasep [pretty arr <> brackets (pretty ind), pretty x])
  pretty (Atomic _ (AtomicOr t old arr ind x)) =
    pretty old
      <+> "<-"
      <+> "atomic_or"
      <> pretty t
      <> parens (commasep [pretty arr <> brackets (pretty ind), pretty x])
  pretty (Atomic _ (AtomicXor t old arr ind x)) =
    pretty old
      <+> "<-"
      <+> "atomic_xor"
      <> pretty t
      <> parens (commasep [pretty arr <> brackets (pretty ind), pretty x])
  pretty (Atomic _ (AtomicCmpXchg t old arr ind x y)) =
    pretty old
      <+> "<-"
      <+> "atomic_cmp_xchg"
      <> pretty t
      <> parens (commasep [pretty arr <> brackets (pretty ind), pretty x, pretty y])
  pretty (Atomic _ (AtomicXchg t old arr ind x)) =
    pretty old
      <+> "<-"
      <+> "atomic_xchg"
      <> pretty t
      <> parens (commasep [pretty arr <> brackets (pretty ind), pretty x])
  pretty (Atomic _ (AtomicWrite t arr ind x)) =
    "atomic_write"
      <> pretty t
      <> parens (commasep [pretty arr <> brackets (pretty ind), pretty x])

instance FreeIn KernelOp where
  freeIn' (Atomic _ op) = freeIn' op
  freeIn' (SharedAlloc _ size) = freeIn' size
  freeIn' _ = mempty

brace :: Doc a -> Doc a
brace body = " {" </> indent 2 body </> "}"