Operations.hs

{-# LANGUAGE BlockArguments   #-}
{-# LANGUAGE DataKinds        #-}
{-# LANGUAGE PolyKinds        #-}
{-# LANGUAGE RebindableSyntax #-}
{-# LANGUAGE TypeApplications #-}
{-# LANGUAGE TypeOperators    #-}

module Tests.Atomics.Operations where

-- fir
import FIR
import Math.Linear

------------------------------------------------
-- program

type Defs =
  '[ "floats" ':-> StorageBuffer '[ DescriptorSet 0, Binding 0 ]
                    ( Struct '[ "arr" ':-> RuntimeArray Float ] )
   , "words"  ':-> StorageBuffer '[ DescriptorSet 0, Binding 1 ]
                    ( Struct '[ "count" ':-> Word32, "arr" ':-> RuntimeArray Word32 ] )
   , "ints"   ':-> StorageBuffer '[ DescriptorSet 0, Binding 2 ]
                    ( Struct '[ "arr" ':-> RuntimeArray Int64 ] )
   , "shared" ':-> Workgroup '[] ( Array 4 Int32 )
   , "main"   ':-> EntryPoint '[ LocalSize 8 8 1 ] Compute
   ]

program :: Module Defs
program = Module $ entryPoint @"main" @Compute do

  ~( Vec3 x y _ ) <- get @"gl_GlobalInvocationID"
  i <- let' ( x + 64 * y )

  -- Floating-point operations on elements of a storage buffer.
  previous <- atomicAdd      @( Name "floats" :.: Name "arr" :.: AnIndex Word32 ) i 0.5
  _        <- atomicMin      @( Name "floats" :.: Name "arr" :.: AnIndex Word32 ) i previous
  _        <- atomicMax      @( Name "floats" :.: Name "arr" :.: AnIndex Word32 ) i 1
  _        <- atomicExchange @( Name "floats" :.: Name "arr" :.: AnIndex Word32 ) i 2
  f        <- atomicLoad     @( Name "floats" :.: Name "arr" :.: AnIndex Word32 ) i
  atomicStore @( Name "floats" :.: Name "arr" :.: AnIndex Word32 ) i ( f + 1 )

  -- Unsigned integer operations on a storage buffer member and on elements
  -- of its array.
  slot <- atomicAdd       @( Name "words" :.: Name "count" ) 1
  _    <- atomicIncrement @( Name "words" :.: Name "count" )
  _    <- atomicDecrement @( Name "words" :.: Name "count" )
  _    <- atomicSub       @( Name "words" :.: Name "count" ) 1
  _    <- atomicMin @( Name "words" :.: Name "arr" :.: AnIndex Word32 ) slot x
  _    <- atomicMax @( Name "words" :.: Name "arr" :.: AnIndex Word32 ) slot y
  _    <- atomicAnd @( Name "words" :.: Name "arr" :.: AnIndex Word32 ) slot 0xff
  _    <- atomicOr  @( Name "words" :.: Name "arr" :.: AnIndex Word32 ) slot 0x100
  _    <- atomicXor @( Name "words" :.: Name "arr" :.: AnIndex Word32 ) slot i
  old  <- atomicCompareExchange @( Name "words" :.: Name "arr" :.: AnIndex Word32 ) slot 0 x

  -- 64-bit signed integer operations.
  _ <- atomicAdd @( Name "ints" :.: Name "arr" :.: AnIndex Word32 ) i 1
  _ <- atomicMin @( Name "ints" :.: Name "arr" :.: AnIndex Word32 ) i (-1)

  -- Signed integer operations on a workgroup variable.
  _ <- atomicAdd @( Name "shared" :.: AnIndex Word32 ) 0 ( if previous > 1 then 1 else -1 )
  _ <- atomicMax @( Name "shared" :.: AnIndex Word32 ) 1 ( bitcast old )

  pure ( Lit () )