repa-3.0.0.1: Data/Array/Repa/Eval/Cursored.hs
{-# LANGUAGE MagicHash #-}
-- | Evaluate an array by dividing it into rectangular blocks and filling
-- each block in parallel.
module Data.Array.Repa.Eval.Cursored
( fillBlock2P
, fillBlock2S
, fillCursoredBlock2P
, fillCursoredBlock2S )
where
import Data.Array.Repa.Index
import Data.Array.Repa.Shape
import Data.Array.Repa.Eval.Elt
import Data.Array.Repa.Eval.Gang
import GHC.Base (remInt, quotInt)
import Prelude as P
import GHC.Exts
-- Non-cursored interface -----------------------------------------------------
-- | Fill a block in a rank-2 array in parallel.
--
-- * Blockwise filling can be more cache-efficient than linear filling for
-- rank-2 arrays.
--
-- * Coordinates given are of the filled edges of the block.
--
-- * We divide the block into columns, and give one column to each thread.
--
-- * Each column is filled in row major order from top to bottom.
--
fillBlock2P
:: Elt a
=> (Int -> a -> IO ()) -- ^ Update function to write into result buffer.
-> (DIM2 -> a) -- ^ Function to evaluate the element at an index.
-> Int -- ^ Width of the whole array.
-> Int -- ^ x0 lower left corner of block to fill
-> Int -- ^ y0
-> Int -- ^ x1 upper right corner of block to fill
-> Int -- ^ y1
-> IO ()
{-# INLINE [0] fillBlock2P #-}
fillBlock2P !write !getElem !imageWidth !x0 !y0 !x1 !y1
= fillCursoredBlock2P
write id addDim getElem
imageWidth x0 y0 x1 y1
-- | Fill a block in a rank-2 array sequentially.
--
-- * Blockwise filling can be more cache-efficient than linear filling for
-- rank-2 arrays.
--
-- * Coordinates given are of the filled edges of the block.
--
-- * The block is filled in row major order from top to bottom.
--
fillBlock2S
:: Elt a
=> (Int -> a -> IO ()) -- ^ Update function to write into result buffer.
-> (DIM2 -> a) -- ^ Function to evaluate the element at an index.
-> Int# -- ^ Width of the whole array.
-> Int# -- ^ x0 lower left corner of block to fill
-> Int# -- ^ y0
-> Int# -- ^ x1 upper right corner of block to fill
-> Int# -- ^ y1
-> IO ()
{-# INLINE [0] fillBlock2S #-}
fillBlock2S !write !getElem imageWidth x0 y0 x1 y1
= fillCursoredBlock2S
write id addDim getElem
imageWidth x0 y0 x1 y1
-- Block filling ----------------------------------------------------------------------------------
-- | Fill a block in a rank-2 array in parallel.
--
-- * Blockwise filling can be more cache-efficient than linear filling for rank-2 arrays.
--
-- * Using cursor functions can help to expose inter-element indexing computations to
-- the GHC and LLVM optimisers.
--
-- * Coordinates given are of the filled edges of the block.
--
-- * We divide the block into columns, and give one column to each thread.
--
-- * Each column is filled in row major order from top to bottom.
--
fillCursoredBlock2P
:: Elt a
=> (Int -> a -> IO ()) -- ^ Update function to write into result buffer.
-> (DIM2 -> cursor) -- ^ Make a cursor to a particular element.
-> (DIM2 -> cursor -> cursor) -- ^ Shift the cursor by an offset.
-> (cursor -> a) -- ^ Function to evaluate the element at an index.
-> Int -- ^ Width of the whole array.
-> Int -- ^ x0 lower left corner of block to fill
-> Int -- ^ y0
-> Int -- ^ x1 upper right corner of block to fill
-> Int -- ^ y1
-> IO ()
{-# INLINE [0] fillCursoredBlock2P #-}
fillCursoredBlock2P
!write
!makeCursorFCB !shiftCursorFCB !getElemFCB
!(I# imageWidth) !x0 !y0 !x1 !y1
= gangIO theGang fillBlock
where !threads = gangSize theGang
!blockWidth = x1 - x0 + 1
-- All columns have at least this many pixels.
!colChunkLen = blockWidth `quotInt` threads
-- Extra pixels that we have to divide between some of the threads.
!colChunkSlack = blockWidth `remInt` threads
-- Get the starting pixel of a column in the image.
{-# INLINE colIx #-}
colIx !ix
| ix < colChunkSlack = x0 + ix * (colChunkLen + 1)
| otherwise = x0 + ix * colChunkLen + colChunkSlack
-- Give one column to each thread
{-# INLINE fillBlock #-}
fillBlock :: Int -> IO ()
fillBlock !ix
= let !(I# x0') = colIx ix
!(I# x1') = colIx (ix + 1) - 1
!(I# y0') = y0
!(I# y1') = y1
in fillCursoredBlock2S
write
makeCursorFCB shiftCursorFCB getElemFCB
imageWidth x0' y0' x1' y1'
-- | Fill a block in a rank-2 array, sequentially.
--
-- * Blockwise filling can be more cache-efficient than linear filling for rank-2 arrays.
--
-- * Using cursor functions can help to expose inter-element indexing computations to
-- the GHC and LLVM optimisers.
--
-- * Coordinates given are of the filled edges of the block.
--
-- * The block is filled in row major order from top to bottom.
--
fillCursoredBlock2S
:: Elt a
=> (Int -> a -> IO ()) -- ^ Update function to write into result buffer.
-> (DIM2 -> cursor) -- ^ Make a cursor to a particular element.
-> (DIM2 -> cursor -> cursor) -- ^ Shift the cursor by an offset.
-> (cursor -> a) -- ^ Function to evaluate an element at the given index.
-> Int# -- ^ Width of the whole array.
-> Int# -- ^ x0 lower left corner of block to fill.
-> Int# -- ^ y0
-> Int# -- ^ x1 upper right corner of block to fill.
-> Int# -- ^ y1
-> IO ()
{-# INLINE [0] fillCursoredBlock2S #-}
fillCursoredBlock2S
!write
!makeCursor !shiftCursor !getElem
!imageWidth !x0 !y0 !x1 !y1
= fillBlock y0
where {-# INLINE fillBlock #-}
fillBlock !y
| y ># y1 = return ()
| otherwise
= do fillLine4 x0
fillBlock (y +# 1#)
where {-# INLINE fillLine4 #-}
fillLine4 !x
| x +# 4# ># x1 = fillLine1 x
| otherwise
= do -- Compute each source cursor based on the previous one so that
-- the variable live ranges in the generated code are shorter.
let srcCur0 = makeCursor (Z :. (I# y) :. (I# x))
let srcCur1 = shiftCursor (Z :. 0 :. 1) srcCur0
let srcCur2 = shiftCursor (Z :. 0 :. 1) srcCur1
let srcCur3 = shiftCursor (Z :. 0 :. 1) srcCur2
-- Get the result value for each cursor.
let val0 = getElem srcCur0
let val1 = getElem srcCur1
let val2 = getElem srcCur2
let val3 = getElem srcCur3
-- Ensure that we've computed each of the result values before we
-- write into the array. If the backend code generator can't tell
-- our destination array doesn't alias with the source then writing
-- to it can prevent the sharing of intermediate computations.
touch val0
touch val1
touch val2
touch val3
-- Compute cursor into destination array.
let !dstCur0 = x +# (y *# imageWidth)
write (I# dstCur0) val0
write (I# (dstCur0 +# 1#)) val1
write (I# (dstCur0 +# 2#)) val2
write (I# (dstCur0 +# 3#)) val3
fillLine4 (x +# 4#)
{-# INLINE fillLine1 #-}
fillLine1 !x
| x ># x1 = return ()
| otherwise
= do write (I# (x +# (y *# imageWidth)))
(getElem $ makeCursor (Z :. (I# y) :. (I# x)))
fillLine1 (x +# 1#)