accelerate-llvm-native-1.3.0.0: src/Data/Array/Accelerate/LLVM/Native/Execute/Divide.hs
{-# LANGUAGE BangPatterns #-}
{-# LANGUAGE GADTs #-}
{-# LANGUAGE PatternGuards #-}
{-# LANGUAGE ScopedTypeVariables #-}
{-# LANGUAGE TypeApplications #-}
-- |
-- Module : Data.Array.Accelerate.LLVM.Native.Execute.Divide
-- Copyright : [2018..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.Execute.Divide (
divideWork, divideWork1
) where
import Data.Array.Accelerate.Representation.Shape
import Data.Bits
import Data.Sequence ( Seq )
import qualified Data.Sequence as Seq
import qualified Data.Vector.Unboxed as U
import qualified Data.Vector.Unboxed.Mutable as M
-- Divide the given multidimensional index range into a sequence of work pieces.
-- Splits will be made on the outermost (left-most) index preferentially, so
-- that spans are longest on the innermost dimension (because caches).
--
-- No dimension will be made smaller than the given minimum.
--
-- The number of subdivisions a hint (at most, it should generate a number of
-- pieces rounded up to the next power-of-two).
--
-- Full pieces will occur first in the resulting sequence, with smaller pieces
-- at the end (suitable for work-stealing). Note that the pieces are not sorted
-- according by size, and are ordered in the resulting sequence depending only
-- on whether all dimensions are above the minimum threshold or not. The integer
-- parameter to the apply action can be used to access the chunks linearly (for
-- example, this is useful when evaluating non-commutative operations).
--
-- {-# INLINABLE divideWork #-}
divideWork
:: ShapeR sh
-> Int -- #subdivisions (hint)
-> Int -- minimum size of a dimension (must be a power of two)
-> sh -- start index (e.g. top-left)
-> sh -- end index (e.g. bottom-right)
-> (Int -> sh -> sh -> a) -- action given start/end index range, and split number in the range [0..]
-> Seq a
divideWork ShapeRz = divideWork0
divideWork (ShapeRsnoc ShapeRz) = divideWork1
divideWork shr = divideWorkN shr
--
-- It is slightly faster to use lists instead of a Sequence here (though the
-- difference is <1us on 'divideWork empty (Z:.2000) nop 8 32'). However,
-- later operations will benefit from more efficient append, etc.
divideWork0 :: Int -> Int -> DIM0 -> DIM0 -> (Int -> DIM0 -> DIM0 -> a) -> Seq a
divideWork0 _ _ () () action = Seq.singleton (action 0 () ())
divideWork1 :: Int -> Int -> DIM1 -> DIM1 -> (Int -> DIM1 -> DIM1 -> a) -> Seq a
divideWork1 !pieces !minsize ((), (!from)) ((), (!to)) action =
let
split 0 !u !v !i !f !s
| v - u < minsize = (i+1, f, s Seq.|> apply i u v)
| otherwise = (i+1, f Seq.|> apply i u v, s)
--
split !s !u !v !i0 !f0 !s0 =
case findSplitPoint1 u v minsize of
Nothing -> (i0+1, f0, s0 Seq.|> apply i0 u v)
Just (u', v') ->
let s' = unsafeShiftR s 1
(i1,f1,s1) = split s' u v' i0 f0 s0
(i2,f2,s2) = split s' u' v i1 f1 s1
in
(i2, f2, s2)
apply i u v = action i ((), u) ((), v)
(_, fs, ss) = split pieces from to 0 Seq.empty Seq.empty
in
fs Seq.>< ss
{-# INLINE findSplitPoint1 #-}
findSplitPoint1
:: Int
-> Int
-> Int
-> Maybe (Int, Int)
findSplitPoint1 !u !v !minsize =
let a = v - u in
if a <= minsize
then Nothing
else
let b = unsafeShiftR (a+1) 1
c = minsize - 1
d = (b+c) .&. complement c
in
Just (d+u, v-a+d)
divideWorkN :: ShapeR sh -> Int -> Int -> sh -> sh -> (Int -> sh -> sh -> a) -> Seq a
divideWorkN !shr !pieces !minsize !from !to action =
let
-- Is it worth checking whether the piece is full? Doing so ensures that
-- full pieces are assigned to threads first, with the non-full blocks
-- being the ones at the end of the work queue to be stolen.
--
split 0 !u !v !i !f !s
| U.any (< minsize) (U.zipWith (-) v u) = (i+1, f, s Seq.|> apply i u v)
| otherwise = (i+1, f Seq.|> apply i u v, s)
--
split !s !u !v !i0 !f0 !s0 =
case findSplitPointN u v minsize of
Nothing -> (i0+1, f0, s0 Seq.|> apply i0 u v)
Just (u', v') ->
let s' = unsafeShiftR s 1
(i1,f1,s1) = split s' u v' i0 f0 s0
(i2,f2,s2) = split s' u' v i1 f1 s1
in
(i2, f2, s2)
apply i u v = action i (vecToShape shr u) (vecToShape shr v)
(_, fs, ss) = split pieces (shapeToVec shr from) (shapeToVec shr to) 0 Seq.empty Seq.empty
in
fs Seq.>< ss
-- Determine if and where to split the given index range. Returns new start and
-- end indices if found.
--
{-# INLINE findSplitPointN #-}
findSplitPointN
:: U.Vector Int -- start
-> U.Vector Int -- end
-> Int -- minimum size of a dimension (must be power of 2)
-> Maybe (U.Vector Int, U.Vector Int)
findSplitPointN !from !to !minsize =
let
mix = U.ifoldr' combine Nothing
$ U.zipWith (-) to from
combine i v old =
if v <= minsize
then old
else case old of
Nothing -> Just (i,v)
Just (_,u) -> if v < u
then Just (i,v)
else old
in
case mix of
Nothing -> Nothing
Just (i,a) ->
let b = unsafeShiftR (a+1) 1 -- divide by 2 (rounded up)
c = minsize - 1
d = (b+c) .&. complement c -- round up to next multiple of chunk size
e = U.unsafeIndex from i
f = U.unsafeIndex to i
--
from' = U.modify (\mv -> M.unsafeWrite mv i (d+e)) from
to' = U.modify (\mv -> M.unsafeWrite mv i (f-a+d)) to
in
Just (from', to')
{-# INLINE vecToShape #-}
vecToShape :: ShapeR sh -> U.Vector Int -> sh
vecToShape shr = listToShape shr . U.toList
{-# INLINE shapeToVec #-}
shapeToVec :: ShapeR sh -> sh -> U.Vector Int
shapeToVec shr sh = U.fromListN (rank shr) (shapeToList shr sh)