accelerate-llvm-native-1.4.0.0: src/Data/Array/Accelerate/LLVM/Native/CodeGen/Permute.hs
{-# LANGUAGE GADTs #-}
{-# LANGUAGE OverloadedStrings #-}
{-# LANGUAGE RecordWildCards #-}
{-# LANGUAGE ScopedTypeVariables #-}
{-# LANGUAGE TypeApplications #-}
-- |
-- Module : Data.Array.Accelerate.LLVM.Native.CodeGen.Permute
-- Copyright : [2016..2020] The Accelerate Team
-- License : BSD3
--
-- Maintainer : Trevor L. McDonell <trevor.mcdonell@gmail.com>
-- Stability : experimental
-- Portability : non-portable (GHC extensions)
--
module Data.Array.Accelerate.LLVM.Native.CodeGen.Permute
where
import Data.Array.Accelerate.AST ( PrimMaybe )
import Data.Array.Accelerate.Error
import Data.Array.Accelerate.Representation.Array
import Data.Array.Accelerate.Representation.Shape
import Data.Array.Accelerate.Representation.Type
import Data.Array.Accelerate.LLVM.CodeGen.Arithmetic as A
import Data.Array.Accelerate.LLVM.CodeGen.Array
import Data.Array.Accelerate.LLVM.CodeGen.Base
import Data.Array.Accelerate.LLVM.CodeGen.Constant
import Data.Array.Accelerate.LLVM.CodeGen.Environment
import Data.Array.Accelerate.LLVM.CodeGen.Exp
import Data.Array.Accelerate.LLVM.CodeGen.IR
import Data.Array.Accelerate.LLVM.CodeGen.Monad
import Data.Array.Accelerate.LLVM.CodeGen.Permute
import Data.Array.Accelerate.LLVM.CodeGen.Ptr
import Data.Array.Accelerate.LLVM.CodeGen.Sugar
import Data.Array.Accelerate.LLVM.Compile.Cache
import Data.Array.Accelerate.LLVM.Native.Target ( Native )
import Data.Array.Accelerate.LLVM.Native.CodeGen.Base
import Data.Array.Accelerate.LLVM.Native.CodeGen.Loop
import LLVM.AST.Type.GetElementPtr
import LLVM.AST.Type.Instruction
import LLVM.AST.Type.Instruction.Atomic
import LLVM.AST.Type.Instruction.RMW as RMW
import LLVM.AST.Type.Instruction.Volatile
import LLVM.AST.Type.Representation
import Control.Applicative
import Control.Monad ( void )
import Prelude
-- Forward permutation specified by an indexing mapping. The resulting array is
-- initialised with the given defaults, and any further values that are permuted
-- into the result array are added to the current value using the combination
-- function.
--
-- The combination function must be /associative/ and /commutative/. Elements
-- that are mapped to the magic index 'ignore' are dropped.
--
mkPermute
:: HasCallStack
=> UID
-> Gamma aenv
-> ArrayR (Array sh e)
-> ShapeR sh'
-> IRPermuteFun Native aenv (e -> e -> e)
-> IRFun1 Native aenv (sh -> PrimMaybe sh')
-> MIRDelayed Native aenv (Array sh e)
-> CodeGen Native (IROpenAcc Native aenv (Array sh' e))
mkPermute uid aenv repr shr combine project arr =
(+++) <$> mkPermuteS uid aenv repr shr combine project arr
<*> mkPermuteP uid aenv repr shr combine project arr
-- Forward permutation which does not require locking the output array. This
-- could be because we are executing sequentially with a single thread, or
-- because the default values are unused (e.g. for a filter).
--
-- We could also use this method if we can prove that the mapping function is
-- injective (distinct elements in the domain map to distinct elements in the
-- co-domain).
--
mkPermuteS
:: UID
-> Gamma aenv
-> ArrayR (Array sh e)
-> ShapeR sh'
-> IRPermuteFun Native aenv (e -> e -> e)
-> IRFun1 Native aenv (sh -> PrimMaybe sh')
-> MIRDelayed Native aenv (Array sh e)
-> CodeGen Native (IROpenAcc Native aenv (Array sh' e))
mkPermuteS uid aenv repr shr IRPermuteFun{..} project marr =
let
(start, end, paramGang) = gangParam (arrayRshape repr)
(arrOut, paramOut) = mutableArray (reprOut repr shr) "out"
(arrIn, paramIn) = delayedArray "in" marr
paramEnv = envParam aenv
in
makeOpenAcc uid "permuteS" (paramGang ++ paramOut ++ paramIn ++ paramEnv) $ do
sh <- delayedExtent arrIn
imapNestFromTo (arrayRshape repr) start end sh $ \ix _ -> do
ix' <- app1 project ix
when (isJust ix') $ do
i <- fromJust ix'
j <- intOfIndex shr (irArrayShape arrOut) i
-- project element onto the destination array and update
x <- app1 (delayedIndex arrIn) ix
y <- readArray TypeInt arrOut j
r <- app2 combine x y
writeArray TypeInt arrOut j r
-- Parallel forward permutation has to take special care because different
-- threads could concurrently try to update the same memory location. Where
-- available we make use of special atomic instructions and other optimisations,
-- but in the general case each element of the output array has a lock which
-- must be obtained by the thread before it can update that memory location.
--
-- TODO: After too many failures to acquire the lock on an element, the thread
-- should back off and try a different element, adding this failed element to
-- a queue or some such.
--
mkPermuteP
:: HasCallStack
=> UID
-> Gamma aenv
-> ArrayR (Array sh e)
-> ShapeR sh'
-> IRPermuteFun Native aenv (e -> e -> e)
-> IRFun1 Native aenv (sh -> PrimMaybe sh')
-> MIRDelayed Native aenv (Array sh e)
-> CodeGen Native (IROpenAcc Native aenv (Array sh' e))
mkPermuteP uid aenv repr shr IRPermuteFun{..} project arr =
case atomicRMW of
Nothing -> mkPermuteP_mutex uid aenv repr shr combine project arr
Just (rmw, f) -> mkPermuteP_rmw uid aenv repr shr rmw f project arr
-- Parallel forward permutation function which uses atomic instructions to
-- implement lock-free array updates.
--
mkPermuteP_rmw
:: HasCallStack
=> UID
-> Gamma aenv
-> ArrayR (Array sh e)
-> ShapeR sh'
-> RMWOperation
-> IRFun1 Native aenv (e -> e)
-> IRFun1 Native aenv (sh -> PrimMaybe sh')
-> MIRDelayed Native aenv (Array sh e)
-> CodeGen Native (IROpenAcc Native aenv (Array sh' e))
mkPermuteP_rmw uid aenv repr shr rmw update project marr =
let
(start, end, paramGang) = gangParam (arrayRshape repr)
(arrOut, paramOut) = mutableArray (reprOut repr shr) "out"
(arrIn, paramIn) = delayedArray "in" marr
paramEnv = envParam aenv
in
makeOpenAcc uid "permuteP_rmw" (paramGang ++ paramOut ++ paramIn ++ paramEnv) $ do
sh <- delayedExtent arrIn
imapNestFromTo (arrayRshape repr) start end sh $ \ix _ -> do
ix' <- app1 project ix
when (isJust ix') $ do
i <- fromJust ix'
j <- intOfIndex shr (irArrayShape arrOut) i
x <- app1 (delayedIndex arrIn) ix
r <- app1 update x
case rmw of
Exchange
-> writeArray TypeInt arrOut j r
--
_ | TupRsingle (SingleScalarType s) <- arrayRtype repr
, adata <- irArrayData arrOut
-> do
addr <- instr' $ GetElementPtr $ GEP1 (SingleScalarType s) (asPtr defaultAddrSpace (op s adata)) (op integralType j)
--
case s of
NumSingleType t -> void . instr' $ AtomicRMW t NonVolatile rmw addr (op t r) (CrossThread, AcquireRelease)
--
_ -> internalError "unexpected transition"
-- Parallel forward permutation function which uses a spinlock to acquire
-- a mutex before updating the value at that location.
--
mkPermuteP_mutex
:: UID
-> Gamma aenv
-> ArrayR (Array sh e)
-> ShapeR sh'
-> IRFun2 Native aenv (e -> e -> e)
-> IRFun1 Native aenv (sh -> PrimMaybe sh')
-> MIRDelayed Native aenv (Array sh e)
-> CodeGen Native (IROpenAcc Native aenv (Array sh' e))
mkPermuteP_mutex uid aenv repr shr combine project marr =
let
(start, end, paramGang) = gangParam (arrayRshape repr)
(arrOut, paramOut) = mutableArray (reprOut repr shr) "out"
(arrLock, paramLock) = mutableArray reprLock "lock"
(arrIn, paramIn) = delayedArray "in" marr
paramEnv = envParam aenv
in
makeOpenAcc uid "permuteP_mutex" (paramGang ++ paramOut ++ paramLock ++ paramIn ++ paramEnv) $ do
sh <- delayedExtent arrIn
imapNestFromTo (arrayRshape repr) start end sh $ \ix _ -> do
ix' <- app1 project ix
-- project element onto the destination array and (atomically) update
when (isJust ix') $ do
i <- fromJust ix'
j <- intOfIndex shr (irArrayShape arrOut) i
x <- app1 (delayedIndex arrIn) ix
atomically arrLock j $ do
y <- readArray TypeInt arrOut j
r <- app2 combine x y
writeArray TypeInt arrOut j r
-- Atomically execute the critical section only when the lock at the given array
-- index is obtained. The thread spins waiting for the lock to be released and
-- there is no backoff strategy in case the lock is contended.
--
-- It is important that the thread loops trying to acquire the lock without
-- writing data anything until the lock value changes. Then, because of MESI
-- caching protocols there will be no bus traffic while the CPU waits for the
-- value to change.
--
-- <https://en.wikipedia.org/wiki/Spinlock#Significant_optimizations>
--
atomically
:: IRArray (Vector Word8)
-> Operands Int
-> CodeGen Native a
-> CodeGen Native a
atomically barriers i action = do
let
lock = integral integralType 1
unlock = integral integralType 0
unlocked = ir TypeWord8 unlock
--
spin <- newBlock "spinlock.entry"
crit <- newBlock "spinlock.critical-section"
exit <- newBlock "spinlock.exit"
addr <- instr' $ GetElementPtr $ GEP1 scalarTypeWord8 (asPtr defaultAddrSpace (op integralType (irArrayData barriers))) (op integralType i)
_ <- br spin
-- Atomically (attempt to) set the lock slot to the locked state. If the slot
-- was unlocked we just acquired it, otherwise the state remains unchanged and
-- we spin until it becomes available.
setBlock spin
old <- instr $ AtomicRMW numType NonVolatile Exchange addr lock (CrossThread, Acquire)
ok <- A.eq singleType old unlocked
_ <- cbr ok crit spin
-- We just acquired the lock; perform the critical section then release the
-- lock and exit. For ("some") x86 processors, an unlocked MOV instruction
-- could be used rather than the slower XCHG, due to subtle memory ordering
-- rules.
setBlock crit
r <- action
_ <- instr $ AtomicRMW numType NonVolatile Exchange addr unlock (CrossThread, Release)
_ <- br exit
setBlock exit
return r
-- Helper functions
-- ----------------
reprOut :: ArrayR (Array sh e) -> ShapeR sh' -> ArrayR (Array sh' e)
reprOut (ArrayR _ tp) shr = ArrayR shr tp
reprLock :: ArrayR (Array ((), Int) Word8)
reprLock = ArrayR (ShapeRsnoc ShapeRz) $ TupRsingle scalarTypeWord8