diff --git a/Data/Array/Repa/Algorithms/Convolve.hs b/Data/Array/Repa/Algorithms/Convolve.hs --- a/Data/Array/Repa/Algorithms/Convolve.hs +++ b/Data/Array/Repa/Algorithms/Convolve.hs @@ -13,16 +13,17 @@ -- module Data.Array.Repa.Algorithms.Convolve ( -- * Arbitrary boundary handling - convolve + convolveP -- * Specialised boundary handling , GetOut , outAs , outClamp - , convolveOut ) + , convolveOutP ) where -import Data.Array.Repa as A -import Data.Array.Repa.Repr.Unboxed as A +import Data.Array.Repa as R +import Data.Array.Repa.Unsafe as R +import Data.Array.Repa.Repr.Unboxed as R import qualified Data.Vector.Unboxed as V import qualified Data.Array.Repa.Shape as S import Prelude as P @@ -32,15 +33,15 @@ -- | Image-kernel convolution, -- which takes a function specifying what value to return when the -- kernel doesn't apply. -convolve - :: (Num a, Unbox a) +convolveP + :: (Num a, Unbox a, Monad m) => (DIM2 -> a) -- ^ Function to get border elements when -- the stencil does not apply. -> Array U DIM2 a -- ^ Stencil to use in the convolution. -> Array U DIM2 a -- ^ Input image. - -> Array U DIM2 a + -> m (Array U DIM2 a) -convolve makeOut kernel image +convolveP makeOut kernel image = kernel `deepSeqArray` image `deepSeqArray` computeP $ unsafeTraverse image id update where @@ -87,7 +88,7 @@ krnZ = krnVec `V.unsafeIndex` krnCur here = imgZ * krnZ in integrate (acc + here) (x + 1) y (imgCur + 1) (krnCur + 1) -{-# INLINE convolve #-} +{-# INLINE convolveP #-} -- Convolve Out ----------------------------------------------------------------------------------- @@ -126,15 +127,14 @@ -- | Image-kernel convolution, -- which takes a function specifying what value to use for out-of-range elements. -convolveOut - :: (Num a, Unbox a) +convolveOutP + :: (Num a, Unbox a, Monad m) => GetOut a -- ^ How to handle out-of-range elements. -> Array U DIM2 a -- ^ Stencil to use in the convolution. -> Array U DIM2 a -- ^ Input image. - -> Array U DIM2 a + -> m (Array U DIM2 a) -{-# INLINE convolveOut #-} -convolveOut getOut kernel image +convolveOutP getOut kernel image = kernel `deepSeqArray` image `deepSeqArray` computeP $ unsafeTraverse image id stencil where @@ -177,4 +177,5 @@ in integrate (count + 1) (acc + here) in integrate 0 0 +{-# INLINE convolveOutP #-} diff --git a/Data/Array/Repa/Algorithms/DFT.hs b/Data/Array/Repa/Algorithms/DFT.hs --- a/Data/Array/Repa/Algorithms/DFT.hs +++ b/Data/Array/Repa/Algorithms/DFT.hs @@ -14,53 +14,53 @@ -- -- You can also compute single values of the transform using `dftWithRootsSingle`. module Data.Array.Repa.Algorithms.DFT - ( dft - , idft - , dftWithRoots - , dftWithRootsSingle) + ( dftP + , idftP + , dftWithRootsP + , dftWithRootsSingleS) where import Data.Array.Repa.Algorithms.DFT.Roots import Data.Array.Repa.Algorithms.Complex -import Data.Array.Repa as A +import Data.Array.Repa as R import Prelude as P -- | Compute the DFT along the low order dimension of an array. -dft :: forall sh - . Shape sh +dftP :: (Shape sh, Monad m) => Array U (sh :. Int) Complex - -> Array U (sh :. Int) Complex + -> m (Array U (sh :. Int) Complex) -dft v - = let rofu = calcRootsOfUnity (extent v) - in dftWithRoots rofu v +dftP v + = do rofu <- calcRootsOfUnityP (extent v) + dftWithRootsP rofu v +{-# INLINE dftP #-} -- | Compute the inverse DFT along the low order dimension of an array. -idft :: forall sh - . Shape sh +idftP :: (Shape sh, Monad m) => Array U (sh :. Int) Complex - -> Array U (sh :. Int) Complex + -> m (Array U (sh :. Int) Complex) -idft v - = let _ :. len = extent v - scale = (fromIntegral len, 0) - rofu = calcInverseRootsOfUnity (extent v) - in computeP $ A.map (/ scale) $ dftWithRoots rofu v +idftP v + = do let _ :. len = extent v + let scale = (fromIntegral len, 0) + rofu <- calcInverseRootsOfUnityP (extent v) + roots <- dftWithRootsP rofu v + computeP $ R.map (/ scale) roots +{-# INLINE idftP #-} -- | Generic function for computation of forward or inverse DFT. -- This function is also useful if you transform many arrays with the same extent, -- and don't want to recompute the roots for each one. -- The extent of the given roots must match that of the input array, else `error`. -dftWithRoots - :: forall sh - . Shape sh +dftWithRootsP + :: (Shape sh, Monad m) => Array U (sh :. Int) Complex -- ^ Roots of unity. -> Array U (sh :. Int) Complex -- ^ Input array. - -> Array U (sh :. Int) Complex + -> m (Array U (sh :. Int) Complex) -dftWithRoots rofu arr +dftWithRootsP rofu arr | _ :. rLen <- extent rofu , _ :. vLen <- extent arr , rLen /= vLen @@ -68,21 +68,20 @@ P.++ " does not match the length of the roots (" P.++ show rLen P.++ ")" | otherwise - = computeP $ traverse arr id (\_ k -> dftWithRootsSingle rofu arr k) - + = computeP $ traverse arr id (\_ k -> dftWithRootsSingleS rofu arr k) +{-# INLINE dftWithRootsP #-} + -- | Compute a single value of the DFT. -- The extent of the given roots must match that of the input array, else `error`. -dftWithRootsSingle - :: forall sh - . Shape sh +dftWithRootsSingleS + :: Shape sh => Array U (sh :. Int) Complex -- ^ Roots of unity. -> Array U (sh :. Int) Complex -- ^ Input array. -> (sh :. Int) -- ^ Index of the value we want. -> Complex -{-# INLINE dftWithRootsSingle #-} -dftWithRootsSingle rofu arrX (_ :. k) +dftWithRootsSingleS rofu arrX (_ :. k) | _ :. rLen <- extent rofu , _ :. vLen <- extent arrX , rLen /= vLen @@ -97,6 +96,7 @@ elemFn (sh' :. n) = rofu ! (sh' :. (k * n) `mod` len) - in A.sumAllP $ A.zipWith (*) arrX wroots + in R.sumAllS $ R.zipWith (*) arrX wroots +{-# INLINE dftWithRootsSingleS #-} diff --git a/Data/Array/Repa/Algorithms/DFT/Roots.hs b/Data/Array/Repa/Algorithms/DFT/Roots.hs --- a/Data/Array/Repa/Algorithms/DFT/Roots.hs +++ b/Data/Array/Repa/Algorithms/DFT/Roots.hs @@ -2,21 +2,20 @@ -- | Calculation of roots of unity for the forward and inverse DFT\/FFT. module Data.Array.Repa.Algorithms.DFT.Roots - ( calcRootsOfUnity - , calcInverseRootsOfUnity) + ( calcRootsOfUnityP + , calcInverseRootsOfUnityP) where import Data.Array.Repa import Data.Array.Repa.Algorithms.Complex -- | Calculate roots of unity for the forward transform. -calcRootsOfUnity - :: forall sh - . Shape sh +calcRootsOfUnityP + :: (Shape sh, Monad m) => (sh :. Int) -- ^ Length of lowest dimension of result. - -> Array U (sh :. Int) Complex + -> m (Array U (sh :. Int) Complex) -calcRootsOfUnity sh@(_ :. n) +calcRootsOfUnityP sh@(_ :. n) = computeP $ fromFunction sh f where f :: Shape sh => (sh :. Int) -> Complex @@ -28,13 +27,12 @@ -- | Calculate roots of unity for the inverse transform. -calcInverseRootsOfUnity - :: forall sh - . Shape sh +calcInverseRootsOfUnityP + :: (Shape sh, Monad m) => (sh :. Int) -- ^ Length of lowest dimension of result. - -> Array U (sh :. Int) Complex + -> m (Array U (sh :. Int) Complex) -calcInverseRootsOfUnity sh@(_ :. n) +calcInverseRootsOfUnityP sh@(_ :. n) = computeP $ fromFunction sh f where f :: Shape sh => (sh :. Int) -> Complex diff --git a/Data/Array/Repa/Algorithms/FFT.hs b/Data/Array/Repa/Algorithms/FFT.hs --- a/Data/Array/Repa/Algorithms/FFT.hs +++ b/Data/Array/Repa/Algorithms/FFT.hs @@ -10,12 +10,14 @@ module Data.Array.Repa.Algorithms.FFT ( Mode(..) , isPowerOfTwo - , fft3d - , fft2d - , fft1d) + , fft3dP + , fft2dP + , fft1dP) where import Data.Array.Repa.Algorithms.Complex -import Data.Array.Repa as A +import Data.Array.Repa as R +import Data.Array.Repa.Eval as R +import Data.Array.Repa.Unsafe as R import Prelude as P @@ -26,16 +28,15 @@ deriving (Show, Eq) -{-# INLINE signOfMode #-} signOfMode :: Mode -> Double signOfMode mode = case mode of Forward -> (-1) Reverse -> 1 Inverse -> 1 +{-# INLINE signOfMode #-} -{-# INLINE isPowerOfTwo #-} -- | Check if an `Int` is a power of two. isPowerOfTwo :: Int -> Bool isPowerOfTwo n @@ -43,16 +44,16 @@ | 2 <- n = True | n `mod` 2 == 0 = isPowerOfTwo (n `div` 2) | otherwise = False +{-# INLINE isPowerOfTwo #-} -- 3D Transform ----------------------------------------------------------------------------------- -- | Compute the DFT of a 3d array. Array dimensions must be powers of two else `error`. -fft3d :: Repr r Complex +fft3dP :: (Repr r Complex, Monad m) => Mode -> Array r DIM3 Complex - -> Array U DIM3 Complex -{-# INLINE fft3d #-} -fft3d mode arr + -> m (Array U DIM3 Complex) +fft3dP mode arr = let _ :. depth :. height :. width = extent arr !sign = signOfMode mode !scale = fromIntegral (depth * width * height) @@ -66,43 +67,44 @@ else arr `deepSeqArray` case mode of - Forward -> fftTrans3d sign $ fftTrans3d sign $ fftTrans3d sign arr - Reverse -> fftTrans3d sign $ fftTrans3d sign $ fftTrans3d sign arr + Forward -> now $ fftTrans3d sign $ fftTrans3d sign $ fftTrans3d sign arr + Reverse -> now $ fftTrans3d sign $ fftTrans3d sign $ fftTrans3d sign arr Inverse -> computeP - $ A.map (/ scale) + $ R.map (/ scale) $ fftTrans3d sign $ fftTrans3d sign $ fftTrans3d sign arr +{-# INLINE fft3dP #-} + fftTrans3d :: Repr r Complex => Double -> Array r DIM3 Complex -> Array U DIM3 Complex -{-# INLINE fftTrans3d #-} fftTrans3d sign arr = let (sh :. len) = extent arr - in computeP $ rotate3d $ fft sign sh len arr + in suspendedComputeP $ rotate3d $ fft sign sh len arr +{-# INLINE fftTrans3d #-} rotate3d :: Repr r Complex => Array r DIM3 Complex -> Array D DIM3 Complex -{-# INLINE rotate3d #-} rotate3d arr = backpermute (sh :. m :. k :. l) f arr where (sh :. k :. l :. m) = extent arr f (sh' :. m' :. k' :. l') = sh' :. k' :. l' :. m' +{-# INLINE rotate3d #-} -- Matrix Transform ------------------------------------------------------------------------------- -- | Compute the DFT of a matrix. Array dimensions must be powers of two else `error`. -fft2d :: Repr r Complex +fft2dP :: (Repr r Complex, Monad m) => Mode -> Array r DIM2 Complex - -> Array U DIM2 Complex -{-# INLINE fft2d #-} -fft2d mode arr + -> m (Array U DIM2 Complex) +fft2dP mode arr = let _ :. height :. width = extent arr sign = signOfMode mode scale = fromIntegral (width * height) @@ -115,30 +117,31 @@ else arr `deepSeqArray` case mode of - Forward -> fftTrans2d sign $ fftTrans2d sign arr - Reverse -> fftTrans2d sign $ fftTrans2d sign arr - Inverse -> computeP $ A.map (/ scale) $ fftTrans2d sign $ fftTrans2d sign arr + Forward -> now $ fftTrans2d sign $ fftTrans2d sign arr + Reverse -> now $ fftTrans2d sign $ fftTrans2d sign arr + Inverse -> computeP $ R.map (/ scale) $ fftTrans2d sign $ fftTrans2d sign arr +{-# INLINE fft2dP #-} -fftTrans2d + +fftTrans2d :: Repr r Complex => Double -> Array r DIM2 Complex -> Array U DIM2 Complex -{-# INLINE fftTrans2d #-} fftTrans2d sign arr = let (sh :. len) = extent arr - in computeP $ transpose $ fft sign sh len arr + in suspendedComputeP $ transpose $ fft sign sh len arr +{-# INLINE fftTrans2d #-} -- Vector Transform ------------------------------------------------------------------------------- -- | Compute the DFT of a vector. Array dimensions must be powers of two else `error`. -fft1d :: Repr r Complex +fft1dP :: (Repr r Complex, Monad m) => Mode -> Array r DIM1 Complex - -> Array U DIM1 Complex -{-# INLINE fft1d #-} -fft1d mode arr + -> m (Array U DIM1 Complex) +fft1dP mode arr = let _ :. len = extent arr sign = signOfMode mode scale = fromIntegral len @@ -151,20 +154,22 @@ else arr `deepSeqArray` case mode of - Forward -> fftTrans1d sign arr - Reverse -> fftTrans1d sign arr - Inverse -> computeP $ A.map (/ scale) $ fftTrans1d sign arr + Forward -> now $ fftTrans1d sign arr + Reverse -> now $ fftTrans1d sign arr + Inverse -> computeP $ R.map (/ scale) $ fftTrans1d sign arr +{-# INLINE fft1dP #-} + fftTrans1d :: Repr r Complex => Double -> Array r DIM1 Complex -> Array U DIM1 Complex -{-# INLINE fftTrans1d #-} fftTrans1d sign arr = let (sh :. len) = extent arr in fft sign sh len arr +{-# INLINE fftTrans1d #-} -- Rank Generalised Worker ------------------------------------------------------------------------ @@ -173,12 +178,11 @@ -> Array r (sh :. Int) Complex -> Array U (sh :. Int) Complex -{-# INLINE fft #-} fft !sign !sh !lenVec !vec = go lenVec 0 1 where go !len !offset !stride | len == 2 - = computeP $ fromFunction (sh :. 2) swivel + = suspendedComputeP $ fromFunction (sh :. 2) swivel | otherwise = combine len @@ -194,7 +198,8 @@ combine !len' evens odds = evens `deepSeqArray` odds `deepSeqArray` let odds' = unsafeTraverse odds id (\get ix@(_ :. k) -> twiddle sign k len' * get ix) - in computeP $ (evens +^ odds') A.++ (evens -^ odds') + in suspendedComputeP $ (evens +^ odds') R.++ (evens -^ odds') +{-# INLINE fft #-} -- Compute a twiddle factor. @@ -203,8 +208,9 @@ -> Int -- length -> Complex -{-# INLINE twiddle #-} twiddle sign k' n' = (cos (2 * pi * k / n), sign * sin (2 * pi * k / n)) where k = fromIntegral k' n = fromIntegral n' +{-# INLINE twiddle #-} + diff --git a/Data/Array/Repa/Algorithms/Matrix.hs b/Data/Array/Repa/Algorithms/Matrix.hs --- a/Data/Array/Repa/Algorithms/Matrix.hs +++ b/Data/Array/Repa/Algorithms/Matrix.hs @@ -21,6 +21,9 @@ , transpose2P, transpose2S) where import Data.Array.Repa as R +import Data.Array.Repa.Eval as R +import Data.Array.Repa.Unsafe as R +import Control.Monad.ST.Strict -- Projections ---------------------------------------------------------------- @@ -38,19 +41,22 @@ -- MMult ---------------------------------------------------------------------- -- | Matrix matrix multiply, in parallel. -mmultP :: Array U DIM2 Double +mmultP :: Monad m + => Array U DIM2 Double -> Array U DIM2 Double - -> Array U DIM2 Double + -> m (Array U DIM2 Double) -mmultP arr' brr - = mmult' arr' (transpose2P brr) - where mmult' arr trr - = trr `deepSeqArray` computeP - $ fromFunction (extent arr) - $ \ix -> R.sumAllS - $ R.zipWith (*) - (slice arr (Any :. (row ix) :. All)) - (slice trr (Any :. (col ix) :. All)) +mmultP arr brr + = [arr, brr] `deepSeqArrays` + do trr <- transpose2P brr + let (Z :. h1 :. _) = extent arr + let (Z :. _ :. w2) = extent brr + computeP + $ fromFunction (Z :. h1 :. w2) + $ \ix -> R.sumAllS + $ R.zipWith (*) + (unsafeSlice arr (Any :. (row ix) :. All)) + (unsafeSlice trr (Any :. (col ix) :. All)) {-# NOINLINE mmultP #-} @@ -59,28 +65,31 @@ -> Array U DIM2 Double -> Array U DIM2 Double -mmultS arr' brr - = mmult' arr' (transpose2S brr) - where mmult' arr trr - = trr `deepSeqArray` computeS - $ fromFunction (extent arr) - $ \ix -> R.sumAllS - $ R.zipWith (*) - (slice arr (Any :. (row ix) :. All)) - (slice trr (Any :. (col ix) :. All)) +mmultS arr brr + = [arr, brr] `deepSeqArrays` runST $ + do trr <- R.now $ transpose2S brr + let (Z :. h1 :. _) = extent arr + let (Z :. _ :. w2) = extent brr + return $ computeS + $ fromFunction (Z :. h1 :. w2) + $ \ix -> R.sumAllS + $ R.zipWith (*) + (unsafeSlice arr (Any :. (row ix) :. All)) + (unsafeSlice trr (Any :. (col ix) :. All)) {-# NOINLINE mmultS #-} -- Transpose ------------------------------------------------------------------ -- | Transpose a 2D matrix, in parallel. transpose2P - :: Array U DIM2 Double - -> Array U DIM2 Double + :: Monad m + => Array U DIM2 Double + -> m (Array U DIM2 Double) transpose2P arr = arr `deepSeqArray` - computeUnboxedP - $ unsafeBackpermute new_extent swap arr + do computeUnboxedP + $ unsafeBackpermute new_extent swap arr where swap (Z :. i :. j) = Z :. j :. i new_extent = swap (extent arr) {-# NOINLINE transpose2P #-} @@ -93,8 +102,8 @@ transpose2S arr = arr `deepSeqArray` - computeUnboxedS - $ unsafeBackpermute new_extent swap arr + do computeUnboxedS + $ unsafeBackpermute new_extent swap arr where swap (Z :. i :. j) = Z :. j :. i new_extent = swap (extent arr) {-# NOINLINE transpose2S #-} diff --git a/repa-algorithms.cabal b/repa-algorithms.cabal --- a/repa-algorithms.cabal +++ b/repa-algorithms.cabal @@ -1,5 +1,5 @@ Name: repa-algorithms -Version: 3.0.0.1 +Version: 3.1.0.1 License: BSD3 License-file: LICENSE Author: The DPH Team @@ -20,7 +20,7 @@ Build-Depends: base == 4.*, vector == 0.9.*, - repa == 3.0.* + repa == 3.1.* ghc-options: -Wall -fno-warn-missing-signatures