Atomics.hs

{-# LANGUAGE BlockArguments      #-}
{-# LANGUAGE GADTs               #-}
{-# LANGUAGE NamedFieldPuns      #-}
{-# LANGUAGE OverloadedStrings   #-}
{-# LANGUAGE TupleSections       #-}

{-|
Module: CodeGen.Atomics

Code generation for atomic operations.
-}

module CodeGen.Atomics
  ( atomicInstruction )
  where

-- base
import Control.Monad
  ( unless )

-- mtl
import Control.Monad.Except
  ( throwError )

-- lens
import Control.Lens
  ( view )

-- text-short
import Data.Text.Short
  ( ShortText )
import qualified Data.Text.Short as ShortText
  ( pack )

-- fir
import CodeGen.Binary
  ( instruction )
import CodeGen.IDs
  ( laidOutTypeID, constID )
import CodeGen.Instruction
  ( Args, toArgs
  , ID(ID), Instruction(..)
  , Pointer(..)
  )
import CodeGen.Monad
  ( CGMonad
  , MonadFresh(fresh)
  )
import CodeGen.Pointers
  ( relayout )
import CodeGen.State
  ( _backend
  , requireCapability, requireExtension
  )
import FIR.Prim.Atomic
  ( SAtomicOp(..), atomicOp, SAtomicShape(..), atomicShape )
import qualified SPIRV.Capability      as SPIRV
import qualified SPIRV.Extension       as SPIRV.Extension
import qualified SPIRV.Layout          as SPIRV
import qualified SPIRV.Operation       as SPIRV.Op
import qualified SPIRV.PrimTy          as SPIRV
import qualified SPIRV.ScalarTy        as SPIRV
import qualified SPIRV.Stage           as SPIRV
import qualified SPIRV.Storage         as Storage
import qualified SPIRV.Synchronisation as SPIRV

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

-- | Perform an atomic operation on the scalar a pointer points to,
-- with the given operands.
atomicInstruction
  :: SAtomicOp op -> Pointer -> [ ID ] -> CGMonad (ID, SPIRV.PrimTy)
atomicInstruction sOp (Pointer ptrID (SPIRV.PointerTy storage lay ty)) operandIDs
  = do
      bk <- view _backend
      unless ( storage `elem` atomicStorageClasses bk ) $
        throwError
          (  "codeGen: atomic operations are unsupported on pointers of storage "
          <> "class " <> ShortText.pack ( show storage ) <> " for backend "
          <> ShortText.pack ( show bk ) <> ".\n\
             \Supported storage classes: "
          <> ShortText.pack ( show ( atomicStorageClasses bk ) ) <> "."
          )
      op <- case ty of
        SPIRV.Scalar sTy -> atomicOperation sOp sTy
        _ -> throwError
               (  "codeGen: " <> opName <> " is unsupported on values of type "
               <> ShortText.pack ( show ty ) <> "."
               )
      -- Device scope and relaxed memory semantics, as for GLSL atomics.
      scopeID         <- constID ( SPIRV.synchronisationScope SPIRV.Device )
      semanticsID     <- constID ( SPIRV.memorySemanticsBitmask relaxed )
      laidOutOperands <-
        traverse ( \ v -> relayout SPIRV.logical lay (v, ty) ) operandIDs
      let
        args :: Args
        args = case ( atomicShape sOp, laidOutOperands ) of
          -- A compare-exchange has separate memory semantics for success and
          -- failure, and takes the new value before the expected value.
          ( SCompares, [ expectedID, newID ] )
            -> toArgs
                 [ ptrID, scopeID, semanticsID, semanticsID, newID, expectedID ]
          _ -> toArgs ( ptrID : scopeID : semanticsID : laidOutOperands )
      case atomicShape sOp of
        SWrites -> do
          instruction
            Instruction
              { operation = op
              , resTy = Nothing
              , resID = Nothing
              , args
              }
          pure (ID 0, SPIRV.Unit) -- ID should never be used
        _ -> do
          tyID <- laidOutTypeID lay ty
          v <- fresh
          instruction
            Instruction
              { operation = op
              , resTy = Just tyID
              , resID = Just v
              , args
              }
          (, ty) <$> relayout lay SPIRV.logical (v, ty)
  where
    opName :: ShortText
    opName = ShortText.pack ( show ( atomicOp sOp ) )
    relaxed :: SPIRV.MemorySemantics
    relaxed = SPIRV.MemorySemantics ( SPIRV.Lock [] ) []

-- | The SPIR-V instruction for an atomic operation on scalars of the given
-- type, requiring the capabilities and extensions it needs.
atomicOperation :: SAtomicOp op -> SPIRV.ScalarTy -> CGMonad SPIRV.Op.Operation
atomicOperation sOp sTy = case sTy of
  SPIRV.Integer sgn w -> do
    case w of
      SPIRV.W32 -> pure ()
      SPIRV.W64 -> requireCapability SPIRV.Int64Atomics
      _         -> unsupported
    pure case sOp of
      SAtomicLoad            -> SPIRV.Op.AtomicLoad
      SAtomicStore           -> SPIRV.Op.AtomicStore
      SAtomicExchange        -> SPIRV.Op.AtomicExchange
      SAtomicCompareExchange -> SPIRV.Op.AtomicCompareExchange
      SAtomicIncrement       -> SPIRV.Op.AtomicIIncrement
      SAtomicDecrement       -> SPIRV.Op.AtomicIDecrement
      SAtomicAdd             -> SPIRV.Op.AtomicIAdd
      SAtomicSub             -> SPIRV.Op.AtomicISub
      SAtomicMin             -> case sgn of
        SPIRV.Signed   -> SPIRV.Op.AtomicSMin
        SPIRV.Unsigned -> SPIRV.Op.AtomicUMin
      SAtomicMax             -> case sgn of
        SPIRV.Signed   -> SPIRV.Op.AtomicSMax
        SPIRV.Unsigned -> SPIRV.Op.AtomicUMax
      SAtomicAnd             -> SPIRV.Op.AtomicAnd
      SAtomicOr              -> SPIRV.Op.AtomicOr
      SAtomicXor             -> SPIRV.Op.AtomicXor
  SPIRV.Floating w -> case sOp of
    SAtomicLoad     -> pure SPIRV.Op.AtomicLoad
    SAtomicStore    -> pure SPIRV.Op.AtomicStore
    SAtomicExchange -> pure SPIRV.Op.AtomicExchange
    SAtomicAdd      -> do
      case w of
        SPIRV.W16 -> do
          requireCapability SPIRV.AtomicFloat16AddEXT
          requireExtension  SPIRV.Extension.SPV_EXT_shader_atomic_float16_add
        SPIRV.W32 -> do
          requireCapability SPIRV.AtomicFloat32AddEXT
          requireExtension  SPIRV.Extension.SPV_EXT_shader_atomic_float_add
        SPIRV.W64 -> do
          requireCapability SPIRV.AtomicFloat64AddEXT
          requireExtension  SPIRV.Extension.SPV_EXT_shader_atomic_float_add
        SPIRV.W8 -> unsupported
      pure SPIRV.Op.AtomicFAddEXT
    SAtomicMin -> SPIRV.Op.AtomicFMinEXT <$ floatMinMax
    SAtomicMax -> SPIRV.Op.AtomicFMaxEXT <$ floatMinMax
    _ -> unsupported
    where
      floatMinMax :: CGMonad ()
      floatMinMax = do
        case w of
          SPIRV.W16 -> requireCapability SPIRV.AtomicFloat16MinMaxEXT
          SPIRV.W32 -> requireCapability SPIRV.AtomicFloat32MinMaxEXT
          SPIRV.W64 -> requireCapability SPIRV.AtomicFloat64MinMaxEXT
          SPIRV.W8  -> unsupported
        requireExtension SPIRV.Extension.SPV_EXT_shader_atomic_float_min_max
  where
    unsupported :: CGMonad a
    unsupported = throwError
      (  "codeGen: " <> ShortText.pack ( show ( atomicOp sOp ) )
      <> " is unsupported on values of type "
      <> ShortText.pack ( show sTy ) <> "."
      )

-- | Storage classes of the pointers atomic operations accept, as specified by
-- the SPIR-V execution environment of each backend.
atomicStorageClasses :: SPIRV.Backend -> [ Storage.StorageClass ]
atomicStorageClasses SPIRV.Vulkan
  = [ Storage.Uniform, Storage.Workgroup, Storage.Image, Storage.StorageBuffer ]
atomicStorageClasses SPIRV.OpenCL
  = [ Storage.Function, Storage.Workgroup
    , Storage.CrossWorkgroup, Storage.Generic
    ]