Atomic.hs

{-# LANGUAGE ConstraintKinds          #-}
{-# LANGUAGE DataKinds                #-}
{-# LANGUAGE DerivingStrategies       #-}
{-# LANGUAGE GADTs                    #-}
{-# LANGUAGE StandaloneKindSignatures #-}
{-# LANGUAGE TypeFamilies             #-}
{-# LANGUAGE UndecidableInstances     #-}

{-|
Module: FIR.Prim.Atomic

Atomic operations (loads, stores, exchanges, and arithmetic and bitwise
read-modify-write operations), and the scalar types that support them.
-}

module FIR.Prim.Atomic
  ( AtomicOp(..), SAtomicOp(..), atomicOp
  , AtomicShape(..), SAtomicShape(..), AtomicOpShape, atomicShape
  , AtomicTy, AtomicIntegralTy, AtomicOpTy
  )
  where

-- base
import Data.Int
  ( Int32, Int64 )
import Data.Kind
  ( Type, Constraint )
import Data.Word
  ( Word32, Word64 )

-- half
import Numeric.Half
  ( Half )

-- fir
import FIR.Prim.Types
  ( ScalarTy, IntegralTy )

------------------------------------------------------------

-- | Atomic operations on the scalar a pointer points to.
data AtomicOp
  = AtomicLoad            -- ^ Read the value.
  | AtomicStore           -- ^ Write a value.
  | AtomicExchange        -- ^ Write a value, returning the previous value.
  | AtomicCompareExchange -- ^ Write a value if the current value equals
                          --   an expected value.
  | AtomicIncrement       -- ^ Add one.
  | AtomicDecrement       -- ^ Subtract one.
  | AtomicAdd             -- ^ Add a value.
  | AtomicSub             -- ^ Subtract a value.
  | AtomicMin             -- ^ Take the minimum with a value.
  | AtomicMax             -- ^ Take the maximum with a value.
  | AtomicAnd             -- ^ Bitwise and with a value.
  | AtomicOr              -- ^ Bitwise or with a value.
  | AtomicXor             -- ^ Bitwise exclusive or with a value.
  deriving stock ( Eq, Show )

-- | Singleton for 'AtomicOp'.
data SAtomicOp ( op :: AtomicOp ) where
  SAtomicLoad            :: SAtomicOp 'AtomicLoad
  SAtomicStore           :: SAtomicOp 'AtomicStore
  SAtomicExchange        :: SAtomicOp 'AtomicExchange
  SAtomicCompareExchange :: SAtomicOp 'AtomicCompareExchange
  SAtomicIncrement       :: SAtomicOp 'AtomicIncrement
  SAtomicDecrement       :: SAtomicOp 'AtomicDecrement
  SAtomicAdd             :: SAtomicOp 'AtomicAdd
  SAtomicSub             :: SAtomicOp 'AtomicSub
  SAtomicMin             :: SAtomicOp 'AtomicMin
  SAtomicMax             :: SAtomicOp 'AtomicMax
  SAtomicAnd             :: SAtomicOp 'AtomicAnd
  SAtomicOr              :: SAtomicOp 'AtomicOr
  SAtomicXor             :: SAtomicOp 'AtomicXor

atomicOp :: SAtomicOp op -> AtomicOp
atomicOp SAtomicLoad            = AtomicLoad
atomicOp SAtomicStore           = AtomicStore
atomicOp SAtomicExchange        = AtomicExchange
atomicOp SAtomicCompareExchange = AtomicCompareExchange
atomicOp SAtomicIncrement       = AtomicIncrement
atomicOp SAtomicDecrement       = AtomicDecrement
atomicOp SAtomicAdd             = AtomicAdd
atomicOp SAtomicSub             = AtomicSub
atomicOp SAtomicMin             = AtomicMin
atomicOp SAtomicMax             = AtomicMax
atomicOp SAtomicAnd             = AtomicAnd
atomicOp SAtomicOr              = AtomicOr
atomicOp SAtomicXor             = AtomicXor

-- | The shape of an atomic operation: its operands besides the pointer,
-- and its result.
data AtomicShape
  = Reads    -- ^ No operand; returns the value.
  | Writes   -- ^ One operand; returns nothing.
  | Modifies -- ^ One operand; returns the previous value.
  | Compares -- ^ Two operands (the expected value, then the new value);
             --   returns the previous value.

-- | Singleton for 'AtomicShape'.
data SAtomicShape ( sh :: AtomicShape ) where
  SReads    :: SAtomicShape 'Reads
  SWrites   :: SAtomicShape 'Writes
  SModifies :: SAtomicShape 'Modifies
  SCompares :: SAtomicShape 'Compares

type AtomicOpShape :: AtomicOp -> AtomicShape
type family AtomicOpShape op where
  AtomicOpShape 'AtomicLoad            = 'Reads
  AtomicOpShape 'AtomicIncrement       = 'Reads
  AtomicOpShape 'AtomicDecrement       = 'Reads
  AtomicOpShape 'AtomicStore           = 'Writes
  AtomicOpShape 'AtomicCompareExchange = 'Compares
  AtomicOpShape _                      = 'Modifies

atomicShape :: SAtomicOp op -> SAtomicShape ( AtomicOpShape op )
atomicShape SAtomicLoad            = SReads
atomicShape SAtomicIncrement       = SReads
atomicShape SAtomicDecrement       = SReads
atomicShape SAtomicStore           = SWrites
atomicShape SAtomicCompareExchange = SCompares
atomicShape SAtomicExchange        = SModifies
atomicShape SAtomicAdd             = SModifies
atomicShape SAtomicSub             = SModifies
atomicShape SAtomicMin             = SModifies
atomicShape SAtomicMax             = SModifies
atomicShape SAtomicAnd             = SModifies
atomicShape SAtomicOr              = SModifies
atomicShape SAtomicXor             = SModifies

-- | Scalar types on which atomic operations are supported:
-- 32-bit and 64-bit integers, and floating-point numbers.
--
-- Under Vulkan, 64-bit integer atomics require the @shaderBufferInt64Atomics@
-- or @shaderSharedInt64Atomics@ features, and floating-point atomics require
-- features of @VK_EXT_shader_atomic_float@ (and of @VK_EXT_shader_atomic_float2@
-- for 16-bit floating-point numbers, and for minimum and maximum).
class ScalarTy ty => AtomicTy ty
instance AtomicTy Word32
instance AtomicTy Word64
instance AtomicTy Int32
instance AtomicTy Int64
instance AtomicTy Half
instance AtomicTy Float
instance AtomicTy Double

-- | Integral types on which atomic operations are supported:
-- 32-bit and 64-bit integers.
class ( AtomicTy ty, IntegralTy ty ) => AtomicIntegralTy ty
instance AtomicIntegralTy Word32
instance AtomicIntegralTy Word64
instance AtomicIntegralTy Int32
instance AtomicIntegralTy Int64

-- | Scalar types on which a given atomic operation is supported: bitwise
-- operations, subtraction, increment, decrement and compare-exchange are only
-- supported on integers.
type AtomicOpTy :: AtomicOp -> Type -> Constraint
type family AtomicOpTy op ty where
  AtomicOpTy 'AtomicCompareExchange ty = AtomicIntegralTy ty
  AtomicOpTy 'AtomicIncrement       ty = AtomicIntegralTy ty
  AtomicOpTy 'AtomicDecrement       ty = AtomicIntegralTy ty
  AtomicOpTy 'AtomicSub             ty = AtomicIntegralTy ty
  AtomicOpTy 'AtomicAnd             ty = AtomicIntegralTy ty
  AtomicOpTy 'AtomicOr              ty = AtomicIntegralTy ty
  AtomicOpTy 'AtomicXor             ty = AtomicIntegralTy ty
  AtomicOpTy _                      ty = AtomicTy ty