diff --git a/LICENSE b/LICENSE new file mode 100644 --- /dev/null +++ b/LICENSE @@ -0,0 +1,30 @@ +Copyright Alexey Kuleshevich (c) 2017 + +All rights reserved. + +Redistribution and use in source and binary forms, with or without +modification, are permitted provided that the following conditions are met: + + * Redistributions of source code must retain the above copyright + notice, this list of conditions and the following disclaimer. + + * Redistributions in binary form must reproduce the above + copyright notice, this list of conditions and the following + disclaimer in the documentation and/or other materials provided + with the distribution. + + * Neither the name of Author name here nor the names of other + contributors may be used to endorse or promote products derived + from this software without specific prior written permission. + +THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS +"AS IS" AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT +LIMITED TO, THE IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR +A PARTICULAR PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT +OWNER OR CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, +SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT +LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, +DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY +THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT +(INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE +OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. diff --git a/README.md b/README.md new file mode 100644 --- /dev/null +++ b/README.md @@ -0,0 +1,2 @@ +# massiv +Efficient Haskell Arrays featuring Parallel computation diff --git a/Setup.hs b/Setup.hs new file mode 100644 --- /dev/null +++ b/Setup.hs @@ -0,0 +1,2 @@ +import Distribution.Simple +main = defaultMain diff --git a/massiv.cabal b/massiv.cabal new file mode 100644 --- /dev/null +++ b/massiv.cabal @@ -0,0 +1,90 @@ +name: massiv +version: 0.1.0.0 +synopsis: Massiv (Массив) is an Array Library. +description: Multi-dimensional Arrays with fusion, stencils and parallel computation. +homepage: https://github.com/lehins/massiv +license: BSD3 +license-file: LICENSE +author: Alexey Kuleshevich +maintainer: alexey@kuleshevi.ch +copyright: 2018 Alexey Kuleshevich +category: Data, Data Structures +build-type: Simple +extra-source-files: README.md +cabal-version: >=1.10 + +library + hs-source-dirs: src + exposed-modules: Data.Massiv.Array + , Data.Massiv.Array.Delayed + , Data.Massiv.Array.Manifest + , Data.Massiv.Array.Manifest.Vector + , Data.Massiv.Array.Mutable + , Data.Massiv.Array.Numeric + , Data.Massiv.Array.Stencil + , Data.Massiv.Array.Unsafe + , Data.Massiv.Core + , Data.Massiv.Core.Scheduler + , Data.Massiv.Core.Index + + other-modules: Data.Massiv.Array.Delayed.Internal + , Data.Massiv.Array.Delayed.Interleaved + , Data.Massiv.Array.Delayed.Windowed + , Data.Massiv.Array.Manifest.BoxedNF + , Data.Massiv.Array.Manifest.BoxedStrict + , Data.Massiv.Array.Manifest.Internal + , Data.Massiv.Array.Manifest.List + , Data.Massiv.Array.Manifest.Primitive + , Data.Massiv.Array.Manifest.Storable + , Data.Massiv.Array.Manifest.Unboxed + , Data.Massiv.Array.Ops.Construct + , Data.Massiv.Array.Ops.Fold + , Data.Massiv.Array.Ops.Map + , Data.Massiv.Array.Ops.Slice + , Data.Massiv.Array.Ops.Transform + , Data.Massiv.Array.Stencil.Convolution + , Data.Massiv.Array.Stencil.Internal + , Data.Massiv.Core.Common + , Data.Massiv.Core.Computation + , Data.Massiv.Core.Index.Class + , Data.Massiv.Core.Index.Ix + , Data.Massiv.Core.Iterator + , Data.Massiv.Core.List + build-depends: base >= 4.7 && < 5 + , data-default-class + , deepseq + , ghc-prim + , primitive + , vector + default-language: Haskell2010 + ghc-options: -Wall + + +Test-Suite tests + Type: exitcode-stdio-1.0 + HS-Source-Dirs: tests + Main-Is: Spec.hs + Other-Modules: Data.Massiv.Array.DelayedSpec + , Data.Massiv.Array.Manifest.VectorSpec + , Data.Massiv.Array.Ops.ConstructSpec + , Data.Massiv.Array.Ops.FoldSpec + , Data.Massiv.Array.Ops.SliceSpec + , Data.Massiv.Array.Ops.TransformSpec + , Data.Massiv.Array.StencilSpec + , Data.Massiv.CoreArbitrary + , Data.Massiv.Core.IndexSpec + , Data.Massiv.Core.SchedulerSpec + Build-Depends: base >= 4.5 && < 5 + , deepseq + , data-default + , safe-exceptions + , massiv + , hspec + , QuickCheck + , vector + Default-Language: Haskell2010 + GHC-Options: -Wall -O2 -fno-warn-orphans -threaded -with-rtsopts=-N2 + +source-repository head + type: git + location: https://github.com/lehins/massiv diff --git a/src/Data/Massiv/Array.hs b/src/Data/Massiv/Array.hs new file mode 100644 --- /dev/null +++ b/src/Data/Massiv/Array.hs @@ -0,0 +1,145 @@ +{-# LANGUAGE FlexibleContexts #-} +{-# LANGUAGE MultiParamTypeClasses #-} +-- | +-- Module : Data.Massiv.Array +-- Copyright : (c) Alexey Kuleshevich 2018 +-- License : BSD3 +-- Maintainer : Alexey Kuleshevich +-- Stability : experimental +-- Portability : non-portable +-- +-- +-- Massiv is a library, that allows creation and manipulation of arrays in parallel and +-- sequentially. Depending on the representation (@__r__@), an @__`Array` r ix e__@ will have +-- certain properties that are unique to that particular representation, but all of them will share +-- the same trait, that an array is simply a mapping from an index (@__ix__@) of an arbitrary +-- dimension to an element (@__e__@) of some value. Which means that some of the array types are +-- pretty classic and are represented by a contiguous chunk of memory reserved for the elements, +-- namely arrays with `Manifest` representations: +-- +-- * `B` - The most basic type of array that can hold any type of element in a boxed form, i.e. each +-- element is a pointer to the actual value, therefore it is also the slowest +-- representation. Elements are kept in a Weak Head Normal Form (WHNF). +-- +-- * `N` - Similar to `B`, is also a boxed type, except it's elements are always kept in a Normal +-- Form (NF). This property is very useful for parallel processing, i.e. when calling +-- `compute` you do want all of your elements to be fully evaluated. +-- +-- * `S` - Is a type of array that is backed by pinned memory, therefore pointers to those arrays +-- can be passed to FFI calls, because Garbage Collector (GC) is guaranteed not to move +-- it. Elements must be an instance of `Storable` class. It is just as efficient as `P` and +-- `U` arrays, except it is subject to fragmentation. +-- +-- * `U` - Unboxed representation. Elements must be an instance of `Unbox` class. +-- +-- * `P` - Array that can hold Haskell primitives, such as `Int`, `Word`, `Double`, etc. Any element +-- must be an instance of `Prim` class. +-- +-- * `M` - General manifest array type, that any of the above representations can be converted to in +-- constant time using `toManifest`. +-- +-- While at the same time, there are arrays that only describe how values for it's elements can be +-- computed, and have no memory overhead on their own. +-- +-- * `D` - delayed array that is a mere function from an index to an element. Crucial representation +-- for fusing computation. Use `computeAs` in order to load array into `Manifest` +-- representation. +-- +-- * `DI` - delayed interleaved array. Same as `D`, but performced better with unbalanced +-- computation, when evaluation one element takes much longer than it's neighbor. +-- +-- * `DW` - delayed windowed array. This peculiar representation allows for very fast `Stencil` +-- computation. +-- +-- Other Array types: +-- +-- * `L` and `LN` - those types aren't particularly useful on their own, but because of their unique +-- ability to be converted to and from nested lists in constant time, provide an amazing +-- intermediary for list/array conversion. +-- +-- Most of the `Manifest` arrays are capable of in-place mutation. Check out +-- "Data.Massiv.Array.Mutable" module for available functionality. +-- +-- Many of the function names exported by this package will clash with the ones +-- from "Prelude", hence it can be more convenient to import like this: +-- +-- @ +-- import Prelude as P +-- import Data.Massiv.Array as A +-- @ +-- +module Data.Massiv.Array + ( -- * Construct + module Data.Massiv.Array.Ops.Construct + -- * Compute + , getComp + , setComp + , compute + , computeAs + , computeSource + , clone + , convert + , convertAs + -- * Size + , size + , Core.elemsCount + , Core.isEmpty + -- * Indexing + , (!?) + , (!) + , (??) + , index + , index' + , defaultIndex + , borderIndex + , evaluateAt + -- * Mapping + , module Data.Massiv.Array.Ops.Map + -- * Folding + + -- $folding + + , module Data.Massiv.Array.Ops.Fold + -- * Transforming + , module Data.Massiv.Array.Ops.Transform + -- * Slicing + , module Data.Massiv.Array.Ops.Slice + -- * Conversion + , module Data.Massiv.Array.Manifest.List + -- * Core + , module Data.Massiv.Core + -- * Representations + , module Data.Massiv.Array.Delayed + , module Data.Massiv.Array.Manifest + -- * Stencil + , module Data.Massiv.Array.Stencil + , module Data.Massiv.Array.Numeric + ) where + +import Data.Massiv.Array.Delayed +import Data.Massiv.Array.Manifest +import Data.Massiv.Array.Numeric +import Data.Massiv.Array.Manifest.Internal +import Data.Massiv.Array.Manifest.List +import Data.Massiv.Array.Mutable as A +import Data.Massiv.Array.Ops.Construct +import Data.Massiv.Array.Ops.Fold +import Data.Massiv.Array.Ops.Map +import Data.Massiv.Array.Ops.Slice +import Data.Massiv.Array.Ops.Transform +import Data.Massiv.Array.Stencil +import Data.Massiv.Core hiding (elemsCount, + isEmpty) +import qualified Data.Massiv.Core as Core (elemsCount, + isEmpty) +import Data.Massiv.Core.Common +import Prelude as P hiding (all, and, any, + foldl, foldr, + maximum, minimum, or, + product, splitAt, + sum) +{- $folding + +All folding is done in a row-major order. + +-} diff --git a/src/Data/Massiv/Array/Delayed.hs b/src/Data/Massiv/Array/Delayed.hs new file mode 100644 --- /dev/null +++ b/src/Data/Massiv/Array/Delayed.hs @@ -0,0 +1,23 @@ +{-# LANGUAGE FlexibleContexts #-} +{-# LANGUAGE FlexibleInstances #-} +{-# LANGUAGE MultiParamTypeClasses #-} +-- | +-- Module : Data.Massiv.Array.Delayed +-- Copyright : (c) Alexey Kuleshevich 2018 +-- License : BSD3 +-- Maintainer : Alexey Kuleshevich +-- Stability : experimental +-- Portability : non-portable +-- +module Data.Massiv.Array.Delayed + ( D(..) + , delay + , DI + , toInterleaved + , DW + ) where + +import Data.Massiv.Array.Delayed.Interleaved +import Data.Massiv.Array.Delayed.Internal +import Data.Massiv.Array.Delayed.Windowed + diff --git a/src/Data/Massiv/Array/Delayed/Interleaved.hs b/src/Data/Massiv/Array/Delayed/Interleaved.hs new file mode 100644 --- /dev/null +++ b/src/Data/Massiv/Array/Delayed/Interleaved.hs @@ -0,0 +1,73 @@ +{-# LANGUAGE BangPatterns #-} +{-# LANGUAGE FlexibleContexts #-} +{-# LANGUAGE FlexibleInstances #-} +{-# LANGUAGE MultiParamTypeClasses #-} +{-# LANGUAGE TypeFamilies #-} +{-# LANGUAGE UndecidableInstances #-} +-- | +-- Module : Data.Massiv.Array.Delayed.Interleaved +-- Copyright : (c) Alexey Kuleshevich 2018 +-- License : BSD3 +-- Maintainer : Alexey Kuleshevich +-- Stability : experimental +-- Portability : non-portable +-- +module Data.Massiv.Array.Delayed.Interleaved + ( DI + , toInterleaved + ) where + +import Data.Massiv.Array.Delayed.Internal +import Data.Massiv.Core.Common +import Data.Massiv.Core.Scheduler + + +-- | Delayed array that will be loaded in an interleaved fasion during parallel +-- computation. +data DI + +type instance EltRepr DI ix = DI + +newtype instance Array DI ix e = DIArray { idArray :: (Array D ix e) } + +instance Index ix => Construct DI ix e where + getComp = dComp . idArray + {-# INLINE getComp #-} + + setComp c arr = arr { idArray = (idArray arr) { dComp = c } } + {-# INLINE setComp #-} + + unsafeMakeArray c sz = DIArray . unsafeMakeArray c sz + {-# INLINE unsafeMakeArray #-} + +instance Functor (Array DI ix) where + fmap f (DIArray arr) = DIArray (fmap f arr) + +instance Index ix => Size DI ix e where + size (DIArray arr) = size arr + {-# INLINE size #-} + + unsafeResize sz = DIArray . unsafeResize sz . idArray + {-# INLINE unsafeResize #-} + + unsafeExtract sIx newSz = DIArray . unsafeExtract sIx newSz . idArray + {-# INLINE unsafeExtract #-} + + +instance Index ix => Load DI ix e where + loadS (DIArray arr) unsafeRead unsafeWrite = loadS arr unsafeRead unsafeWrite + {-# INLINE loadS #-} + loadP wIds (DIArray (DArray _ sz f)) _ unsafeWrite = + withScheduler_ wIds $ \ scheduler -> do + let !totalLength = totalElem sz + loopM_ 0 (< numWorkers scheduler) (+ 1) $ \ !start -> + scheduleWork scheduler $ + iterLinearM_ sz start totalLength (numWorkers scheduler) (<) $ \ !k !ix -> + unsafeWrite k $ f ix + {-# INLINE loadP #-} + +-- | Convert a source array into an array that, when computed, will have its elemets evaluated out +-- of order (interleaved amoungs cores), hence making unbalanced computation better parallelizable. +toInterleaved :: Source r ix e => Array r ix e -> Array DI ix e +toInterleaved = DIArray . delay +{-# INLINE toInterleaved #-} diff --git a/src/Data/Massiv/Array/Delayed/Internal.hs b/src/Data/Massiv/Array/Delayed/Internal.hs new file mode 100644 --- /dev/null +++ b/src/Data/Massiv/Array/Delayed/Internal.hs @@ -0,0 +1,245 @@ +{-# LANGUAGE BangPatterns #-} +{-# LANGUAGE FlexibleContexts #-} +{-# LANGUAGE FlexibleInstances #-} +{-# LANGUAGE MultiParamTypeClasses #-} +{-# LANGUAGE ScopedTypeVariables #-} +{-# LANGUAGE TypeFamilies #-} +-- | +-- Module : Data.Massiv.Array.Delayed.Internal +-- Copyright : (c) Alexey Kuleshevich 2018 +-- License : BSD3 +-- Maintainer : Alexey Kuleshevich +-- Stability : experimental +-- Portability : non-portable +-- +module Data.Massiv.Array.Delayed.Internal + ( D(..) + , Array(..) + , delay + , eq + , liftArray + , liftArray2 + ) where + +import Data.Foldable (Foldable (..)) +import Data.Massiv.Array.Ops.Fold as A +import Data.Massiv.Core.Common +import Data.Massiv.Core.Scheduler +import Data.Monoid ((<>)) +import GHC.Base (build) +import Prelude hiding (zipWith) + +-- | Delayed representation. +data D = D deriving Show + + +data instance Array D ix e = DArray { dComp :: !Comp + , dSize :: !ix + , dUnsafeIndex :: ix -> e } +type instance EltRepr D ix = D + +instance Index ix => Construct D ix e where + getComp = dComp + {-# INLINE getComp #-} + + setComp c arr = arr { dComp = c } + {-# INLINE setComp #-} + + unsafeMakeArray = DArray + {-# INLINE unsafeMakeArray #-} + + +instance Index ix => Source D ix e where + unsafeIndex = dUnsafeIndex + {-# INLINE unsafeIndex #-} + +instance Index ix => Size D ix e where + size = dSize + {-# INLINE size #-} + + unsafeResize !sz !arr = + DArray (getComp arr) sz $ \ !ix -> + unsafeIndex arr (fromLinearIndex (size arr) (toLinearIndex sz ix)) + {-# INLINE unsafeResize #-} + + unsafeExtract !sIx !newSz !arr = + DArray (getComp arr) newSz $ \ !ix -> + unsafeIndex arr (liftIndex2 (+) ix sIx) + {-# INLINE unsafeExtract #-} + +instance ( Index ix + , Index (Lower ix) + , Elt D ix e ~ Array D (Lower ix) e + ) => + Slice D ix e where + unsafeSlice arr start cutSz dim = do + newSz <- dropDim cutSz dim + return $ unsafeResize newSz (unsafeExtract start cutSz arr) + {-# INLINE unsafeSlice #-} + + +instance (Elt D ix e ~ Array D (Lower ix) e, Index ix) => OuterSlice D ix e where + + unsafeOuterSlice !arr !i = + DArray (getComp arr) (tailDim (size arr)) (\ !ix -> unsafeIndex arr (consDim i ix)) + {-# INLINE unsafeOuterSlice #-} + +instance (Elt D ix e ~ Array D (Lower ix) e, Index ix) => InnerSlice D ix e where + + unsafeInnerSlice !arr !(szL, _) !i = + DArray (getComp arr) szL (\ !ix -> unsafeIndex arr (snocDim ix i)) + {-# INLINE unsafeInnerSlice #-} + + +instance (Eq e, Index ix) => Eq (Array D ix e) where + (==) = eq (==) + {-# INLINE (==) #-} + + +instance Functor (Array D ix) where + fmap f (DArray c sz g) = DArray c sz (f . g) + {-# INLINE fmap #-} + + +instance Index ix => Applicative (Array D ix) where + pure a = DArray Seq (liftIndex (+ 1) zeroIndex) (const a) + {-# INLINE pure #-} + (<*>) (DArray c1 sz1 uIndex1) (DArray c2 sz2 uIndex2) = + DArray (c1 <> c2) (liftIndex2 min sz1 sz2) $ \ !ix -> + (uIndex1 ix) (uIndex2 ix) + {-# INLINE (<*>) #-} + + +-- | Row-major sequential folding over a delayed array. +instance Index ix => Foldable (Array D ix) where + foldl = lazyFoldlS + {-# INLINE foldl #-} + foldl' = foldlS + {-# INLINE foldl' #-} + foldr = foldrFB + {-# INLINE foldr #-} + foldr' = foldrS + {-# INLINE foldr' #-} + null (DArray _ sz _) = totalElem sz == 0 + {-# INLINE null #-} + sum = foldl' (+) 0 + {-# INLINE sum #-} + product = foldl' (*) 1 + {-# INLINE product #-} + length = totalElem . size + {-# INLINE length #-} + toList arr = build (\ c n -> foldrFB c n arr) + {-# INLINE toList #-} + + +instance Index ix => Load D ix e where + loadS (DArray _ sz f) _ unsafeWrite = + iterM_ zeroIndex sz 1 (<) $ \ !ix -> + unsafeWrite (toLinearIndex sz ix) (f ix) + {-# INLINE loadS #-} + loadP wIds (DArray _ sz f) _ unsafeWrite = do + divideWork_ wIds sz $ \ !scheduler !chunkLength !totalLength !slackStart -> do + loopM_ 0 (< slackStart) (+ chunkLength) $ \ !start -> + scheduleWork scheduler $ + iterLinearM_ sz start (start + chunkLength) 1 (<) $ \ !k !ix -> do + unsafeWrite k $ f ix + scheduleWork scheduler $ + iterLinearM_ sz slackStart totalLength 1 (<) $ \ !k !ix -> do + unsafeWrite k (f ix) + {-# INLINE loadP #-} + + +instance (Index ix, Num e) => Num (Array D ix e) where + (+) = liftArray2 (+) + {-# INLINE (+) #-} + (-) = liftArray2 (-) + {-# INLINE (-) #-} + (*) = liftArray2 (*) + {-# INLINE (*) #-} + abs = liftArray abs + {-# INLINE abs #-} + signum = liftArray signum + {-# INLINE signum #-} + fromInteger = singleton Seq . fromInteger + {-# INLINE fromInteger #-} + +instance (Index ix, Fractional e) => Fractional (Array D ix e) where + (/) = liftArray2 (/) + {-# INLINE (/) #-} + fromRational = singleton Seq . fromRational + {-# INLINE fromRational #-} + + +instance (Index ix, Floating e) => Floating (Array D ix e) where + pi = singleton Seq pi + {-# INLINE pi #-} + exp = liftArray exp + {-# INLINE exp #-} + log = liftArray log + {-# INLINE log #-} + sin = liftArray sin + {-# INLINE sin #-} + cos = liftArray cos + {-# INLINE cos #-} + asin = liftArray asin + {-# INLINE asin #-} + atan = liftArray atan + {-# INLINE atan #-} + acos = liftArray acos + {-# INLINE acos #-} + sinh = liftArray sinh + {-# INLINE sinh #-} + cosh = liftArray cosh + {-# INLINE cosh #-} + asinh = liftArray asinh + {-# INLINE asinh #-} + atanh = liftArray atanh + {-# INLINE atanh #-} + acosh = liftArray acosh + {-# INLINE acosh #-} + + + +-- | /O(1)/ Conversion from a source array to `D` representation. +delay :: Source r ix e => Array r ix e -> Array D ix e +delay arr = DArray (getComp arr) (size arr) (unsafeIndex arr) +{-# INLINE delay #-} + + +-- | /O(n1 + n2)/ - Compute array equality by applying a comparing function to each element. +eq :: (Source r1 ix e1, Source r2 ix e2) => + (e1 -> e2 -> Bool) -> Array r1 ix e1 -> Array r2 ix e2 -> Bool +eq f arr1 arr2 = + (size arr1 == size arr2) && + A.fold + (&&) + True + (DArray (getComp arr1 <> getComp arr2) (size arr1) $ \ix -> + f (unsafeIndex arr1 ix) (unsafeIndex arr2 ix)) +{-# INLINE eq #-} + + +liftArray :: Source r ix b => (b -> e) -> Array r ix b -> Array D ix e +liftArray f !arr = DArray (getComp arr) (size arr) (f . unsafeIndex arr) +{-# INLINE liftArray #-} + +-- | Similar to @zipWith@, except dimensions of both arrays either have to be the +-- same, or at least one of two array must be a singleton array, in which +-- case it will behave as @fmap@. +liftArray2 + :: (Source r1 ix a, Source r2 ix b) + => (a -> b -> e) -> Array r1 ix a -> Array r2 ix b -> Array D ix e +liftArray2 f !arr1 !arr2 + | sz1 == oneIndex = liftArray (f (unsafeIndex arr1 zeroIndex)) arr2 + | sz2 == oneIndex = liftArray (`f` (unsafeIndex arr2 zeroIndex)) arr1 + | sz1 == sz2 = + DArray (getComp arr1) sz1 (\ !ix -> f (unsafeIndex arr1 ix) (unsafeIndex arr2 ix)) + | otherwise = + error $ + "Array dimensions must be the same, instead got: " ++ + show (size arr1) ++ " and " ++ show (size arr2) + where + oneIndex = pureIndex 1 + sz1 = size arr1 + sz2 = size arr2 +{-# INLINE liftArray2 #-} diff --git a/src/Data/Massiv/Array/Delayed/Windowed.hs b/src/Data/Massiv/Array/Delayed/Windowed.hs new file mode 100644 --- /dev/null +++ b/src/Data/Massiv/Array/Delayed/Windowed.hs @@ -0,0 +1,373 @@ +{-# LANGUAGE RecordWildCards #-} +{-# LANGUAGE BangPatterns #-} +{-# LANGUAGE FlexibleContexts #-} +{-# LANGUAGE FlexibleInstances #-} +{-# LANGUAGE MultiParamTypeClasses #-} +{-# LANGUAGE TypeFamilies #-} +{-# LANGUAGE UndecidableInstances #-} +-- | +-- Module : Data.Massiv.Array.Delayed.Windowed +-- Copyright : (c) Alexey Kuleshevich 2018 +-- License : BSD3 +-- Maintainer : Alexey Kuleshevich +-- Stability : experimental +-- Portability : non-portable +-- +module Data.Massiv.Array.Delayed.Windowed + ( DW + , Array(..) + , makeWindowedArray + ) where + +import Control.Monad (when) +import Data.Massiv.Array.Delayed.Internal +import Data.Massiv.Core +import Data.Massiv.Core.Common +import Data.Massiv.Core.Scheduler + +-- | Delayed Windowed Array representation. +data DW + +type instance EltRepr DW ix = D + +data instance Array DW ix e = DWArray { wdArray :: !(Array D ix e) + , wdStencilSize :: Maybe ix + -- ^ Setting this value during stencil + -- application improves cache utilization + -- while computing an array + , wdWindowStartIndex :: !ix + , wdWindowSize :: !ix + , wdWindowUnsafeIndex :: ix -> e } + +instance Index ix => Construct DW ix e where + getComp = dComp . wdArray + {-# INLINE getComp #-} + + setComp c arr = arr { wdArray = (wdArray arr) { dComp = c } } + {-# INLINE setComp #-} + + unsafeMakeArray c sz f = DWArray (unsafeMakeArray c sz f) Nothing zeroIndex zeroIndex f + {-# INLINE unsafeMakeArray #-} + + +-- | Any resize or extract on Windowed Array will hurt the performance. +instance Index ix => Size DW ix e where + size = size . wdArray + {-# INLINE size #-} + unsafeResize sz DWArray {..} = + let dArr = unsafeResize sz wdArray + in DWArray + { wdArray = dArr + , wdStencilSize = Nothing + , wdWindowStartIndex = zeroIndex + , wdWindowSize = zeroIndex + , wdWindowUnsafeIndex = evaluateAt dArr + } + unsafeExtract sIx newSz = unsafeExtract sIx newSz . wdArray + + +instance Functor (Array DW ix) where + fmap f !arr = + arr + { wdArray = fmap f (wdArray arr) + , wdWindowUnsafeIndex = f . wdWindowUnsafeIndex arr + } + {-# INLINE fmap #-} + + +-- | Supply a separate generating function for interior of an array. This is +-- very usful for stencil mapping, where interior function does not perform +-- boundary checks, thus significantly speeding up computation process. +makeWindowedArray + :: Source r ix e + => Array r ix e -- ^ Source array that will have a window inserted into it + -> ix -- ^ Start index for the window + -> ix -- ^ Size of the window + -> (ix -> e) -- ^ Inside window indexing function + -> Array DW ix e +makeWindowedArray !arr !wIx !wSz wUnsafeIndex + | not (isSafeIndex sz wIx) = + error $ + "Incorrect window starting index: " ++ show wIx ++ " for: " ++ show (size arr) + | liftIndex2 (+) wIx wSz > sz = + error $ + "Incorrect window size: " ++ + show wSz ++ " and/or placement: " ++ show wIx ++ " for: " ++ show (size arr) + | otherwise = + DWArray + { wdArray = delay arr + , wdStencilSize = Nothing + , wdWindowStartIndex = wIx + , wdWindowSize = wSz + , wdWindowUnsafeIndex = wUnsafeIndex + } + where sz = size arr +{-# INLINE makeWindowedArray #-} + + + + +instance {-# OVERLAPPING #-} Load DW Ix1 e where + loadS (DWArray (DArray _ sz indexB) _ it wk indexW) _ unsafeWrite = do + iterM_ 0 it 1 (<) $ \ !i -> unsafeWrite i (indexB i) + iterM_ it wk 1 (<) $ \ !i -> unsafeWrite i (indexW i) + iterM_ wk sz 1 (<) $ \ !i -> unsafeWrite i (indexB i) + {-# INLINE loadS #-} + loadP wIds (DWArray (DArray _ sz indexB) _ it wk indexW) _ unsafeWrite = do + divideWork_ wIds wk $ \ !scheduler !chunkLength !totalLength !slackStart -> do + scheduleWork scheduler $ + iterM_ 0 it 1 (<) $ \ !ix -> + unsafeWrite (toLinearIndex sz ix) (indexB ix) + scheduleWork scheduler $ + iterM_ wk sz 1 (<) $ \ !ix -> + unsafeWrite (toLinearIndex sz ix) (indexB ix) + loopM_ it (< (slackStart + it)) (+ chunkLength) $ \ !start -> + scheduleWork scheduler $ + iterM_ start (start + chunkLength) 1 (<) $ \ !k -> + unsafeWrite k $ indexW k + scheduleWork scheduler $ + iterM_ (slackStart + it) (totalLength + it) 1 (<) $ \ !k -> + unsafeWrite k (indexW k) + {-# INLINE loadP #-} + + + +instance {-# OVERLAPPING #-} Load DW Ix2 e where + loadS arr _ unsafeWrite = do + let (DWArray (DArray _ sz@(m :. n) indexB) mStencilSz (it :. jt) (wm :. wn) indexW) = + arr + let (ib :. jb) = (wm + it) :. (wn + jt) + blockHeight = case mStencilSz of + Just (i :. _) -> i + _ -> 1 + iterM_ (0 :. 0) (it :. n) 1 (<) $ \ !ix -> + unsafeWrite (toLinearIndex sz ix) (indexB ix) + iterM_ (ib :. 0) (m :. n) 1 (<) $ \ !ix -> + unsafeWrite (toLinearIndex sz ix) (indexB ix) + iterM_ (it :. 0) (ib :. jt) 1 (<) $ \ !ix -> + unsafeWrite (toLinearIndex sz ix) (indexB ix) + iterM_ (it :. jb) (ib :. n) 1 (<) $ \ !ix -> + unsafeWrite (toLinearIndex sz ix) (indexB ix) + unrollAndJam blockHeight (it :. ib) (jt :. jb) $ \ !ix -> + unsafeWrite (toLinearIndex sz ix) (indexW ix) + {-# INLINE loadS #-} + loadP wIds arr _ unsafeWrite = do + let (DWArray (DArray _ sz@(m :. n) indexB) mStencilSz (it :. jt) (wm :. wn) indexW) = arr + withScheduler_ wIds $ \scheduler -> do + let (ib :. jb) = (wm + it) :. (wn + jt) + !blockHeight = case mStencilSz of + Just (i :. _) -> i + _ -> 1 + !(chunkHeight, slackHeight) = wm `quotRem` numWorkers scheduler + let loadBlock !it' !ib' = + unrollAndJam blockHeight (it' :. ib') (jt :. jb) $ \ !ix -> + unsafeWrite (toLinearIndex sz ix) (indexW ix) + {-# INLINE loadBlock #-} + scheduleWork scheduler $ + iterM_ (0 :. 0) (it :. n) 1 (<) $ \ !ix -> + unsafeWrite (toLinearIndex sz ix) (indexB ix) + scheduleWork scheduler $ + iterM_ (ib :. 0) (m :. n) 1 (<) $ \ !ix -> + unsafeWrite (toLinearIndex sz ix) (indexB ix) + scheduleWork scheduler $ + iterM_ (it :. 0) (ib :. jt) 1 (<) $ \ !ix -> + unsafeWrite (toLinearIndex sz ix) (indexB ix) + scheduleWork scheduler $ + iterM_ (it :. jb) (ib :. n) 1 (<) $ \ !ix -> + unsafeWrite (toLinearIndex sz ix) (indexB ix) + loopM_ 0 (< numWorkers scheduler) (+ 1) $ \ !wid -> do + let !it' = wid * chunkHeight + it + scheduleWork scheduler $ loadBlock it' (it' + chunkHeight) + when (slackHeight > 0) $ do + let !itSlack = (numWorkers scheduler) * chunkHeight + it + scheduleWork scheduler $ + loadBlock itSlack (itSlack + slackHeight) + {-# INLINE loadP #-} + + +-- instance {-# OVERLAPPING #-} Load DW Ix3 e where +-- loadS = loadWindowedSRec +-- {-# INLINE loadS #-} +-- loadP = loadWindowedPRec +-- {-# INLINE loadP #-} + + +instance {-# OVERLAPPABLE #-} (Index ix, Load DW (Lower ix) e) => Load DW ix e where + loadS = loadWindowedSRec + {-# INLINE loadS #-} + loadP = loadWindowedPRec + {-# INLINE loadP #-} + + +loadWindowedSRec :: (Index ix, Load DW (Lower ix) e, Monad m) => + Array DW ix e -> (Int -> m e) -> (Int -> e -> m ()) -> m () +loadWindowedSRec (DWArray darr mStencilSz tix wSz indexW) _unsafeRead unsafeWrite = do + let DArray _ sz indexB = darr + !szL = tailDim sz + !bix = liftIndex2 (+) tix wSz + !(t, tixL) = unconsDim tix + !pageElements = totalElem szL + unsafeWriteLower i k val = unsafeWrite (k + pageElements * i) val + {-# INLINE unsafeWriteLower #-} + iterM_ zeroIndex tix 1 (<) $ \ !ix -> + unsafeWrite (toLinearIndex sz ix) (indexB ix) + iterM_ bix sz 1 (<) $ \ !ix -> + unsafeWrite (toLinearIndex sz ix) (indexB ix) + loopM_ t (< headDim bix) (+ 1) $ \ !i -> + let !lowerArr = + (DWArray + (DArray Seq szL (indexB . consDim i)) + (tailDim <$> mStencilSz) -- can safely drop the dim, only + -- last 2 matter anyways + tixL + (tailDim wSz) + (indexW . consDim i)) + in loadS lowerArr _unsafeRead (unsafeWriteLower i) +{-# INLINE loadWindowedSRec #-} + + +loadWindowedPRec :: (Index ix, Load DW (Lower ix) e) => + [Int] -> Array DW ix e -> (Int -> IO e) -> (Int -> e -> IO ()) -> IO () +loadWindowedPRec wIds (DWArray darr mStencilSz tix wSz indexW) _unsafeRead unsafeWrite = do + withScheduler_ wIds $ \ scheduler -> do + let DArray _ sz indexB = darr + !szL = tailDim sz + !bix = liftIndex2 (+) tix wSz + !(t, tixL) = unconsDim tix + !pageElements = totalElem szL + unsafeWriteLower i k = unsafeWrite (k + pageElements * i) + {-# INLINE unsafeWriteLower #-} + scheduleWork scheduler $ + iterM_ zeroIndex tix 1 (<) $ \ !ix -> + unsafeWrite (toLinearIndex sz ix) (indexB ix) + scheduleWork scheduler $ + iterM_ bix sz 1 (<) $ \ !ix -> + unsafeWrite (toLinearIndex sz ix) (indexB ix) + loopM_ t (< headDim bix) (+ 1) $ \ !i -> + let !lowerArr = + (DWArray + (DArray Seq szL (indexB . consDim i)) + (tailDim <$> mStencilSz) -- can safely drop the dim, only + -- last 2 matter anyways + tixL + (tailDim wSz) + (indexW . consDim i)) + in scheduleWork scheduler $ + loadS + lowerArr + (_unsafeRead) + (unsafeWriteLower i) +{-# INLINE loadWindowedPRec #-} + + + +unrollAndJam :: Monad m => + Int -> Ix2 -> Ix2 -> (Ix2 -> m a) -> m () +unrollAndJam !bH (it :. ib) (jt :. jb) f = do + let !bH' = min (max 1 bH) 7 + let f2 (i :. j) = f (i :. j) >> f ((i + 1) :. j) + let f3 (i :. j) = f (i :. j) >> f2 ((i + 1) :. j) + let f4 (i :. j) = f (i :. j) >> f3 ((i + 1) :. j) + let f5 (i :. j) = f (i :. j) >> f4 ((i + 1) :. j) + let f6 (i :. j) = f (i :. j) >> f5 ((i + 1) :. j) + let f7 (i :. j) = f (i :. j) >> f6 ((i + 1) :. j) + let f' = case bH' of + 1 -> f + 2 -> f2 + 3 -> f3 + 4 -> f4 + 5 -> f5 + 6 -> f6 + _ -> f7 + let !ibS = ib - ((ib - it) `mod` bH') + loopM_ it (< ibS) (+ bH') $ \ !i -> + loopM_ jt (< jb) (+ 1) $ \ !j -> + f' (i :. j) + loopM_ ibS (< ib) (+ 1) $ \ !i -> + loopM_ jt (< jb) (+ 1) $ \ !j -> + f (i :. j) +{-# INLINE unrollAndJam #-} + + +-- TODO: Implement Hilbert curve + + +instance {-# OVERLAPPING #-} Load DW Ix2T e where + loadS arr _ unsafeWrite = do + let (DWArray (DArray _ sz@(m, n) indexB) mStencilSz (it, jt) (wm, wn) indexW) = + arr + let (ib, jb) = (wm + it, wn + jt) + blockHeight = case mStencilSz of + Just (i, _) -> i + _ -> 1 + iterM_ (0, 0) (it, n) 1 (<) $ \ !ix -> + unsafeWrite (toLinearIndex sz ix) (indexB ix) + iterM_ (ib, 0) (m, n) 1 (<) $ \ !ix -> + unsafeWrite (toLinearIndex sz ix) (indexB ix) + iterM_ (it, 0) (ib, jt) 1 (<) $ \ !ix -> + unsafeWrite (toLinearIndex sz ix) (indexB ix) + iterM_ (it, jb) (ib, n) 1 (<) $ \ !ix -> + unsafeWrite (toLinearIndex sz ix) (indexB ix) + unrollAndJamT blockHeight (it, ib) (jt, jb) $ \ !ix -> + unsafeWrite (toLinearIndex sz ix) (indexW ix) + {-# INLINE loadS #-} + loadP wIds arr _ unsafeWrite = do + let (DWArray (DArray _ sz@(m, n) indexB) mStencilSz (it, jt) (wm, wn) indexW) = arr + withScheduler_ wIds $ \ scheduler -> do + let (ib, jb) = (wm + it, wn + jt) + blockHeight = case mStencilSz of + Just (i, _) -> i + _ -> 1 + !(chunkHeight, slackHeight) = wm `quotRem` numWorkers scheduler + let loadBlock !it' !ib' = + unrollAndJamT blockHeight (it', ib') (jt, jb) $ \ !ix -> + unsafeWrite (toLinearIndex sz ix) (indexW ix) + {-# INLINE loadBlock #-} + scheduleWork scheduler $ + iterM_ (0, 0) (it, n) 1 (<) $ \ !ix -> + unsafeWrite (toLinearIndex sz ix) (indexB ix) + scheduleWork scheduler $ + iterM_ (ib, 0) (m, n) 1 (<) $ \ !ix -> + unsafeWrite (toLinearIndex sz ix) (indexB ix) + scheduleWork scheduler $ + iterM_ (it, 0) (ib, jt) 1 (<) $ \ !ix -> + unsafeWrite (toLinearIndex sz ix) (indexB ix) + scheduleWork scheduler $ + iterM_ (it, jb) (ib, n) 1 (<) $ \ !ix -> + unsafeWrite (toLinearIndex sz ix) (indexB ix) + loopM_ 0 (< numWorkers scheduler) (+ 1) $ \ !wid -> do + let !it' = wid * chunkHeight + it + scheduleWork scheduler $ loadBlock it' (it' + chunkHeight) + when (slackHeight > 0) $ do + let !itSlack = (numWorkers scheduler) * chunkHeight + it + scheduleWork scheduler $ loadBlock itSlack (itSlack + slackHeight) + {-# INLINE loadP #-} + + + +unrollAndJamT :: Monad m => + Int -> Ix2T -> Ix2T -> (Ix2T -> m a) -> m () +unrollAndJamT !bH (it, ib) (jt, jb) f = do + let !bH' = min (max 1 bH) 7 + let f2 !(i, j) = f (i, j) >> f (i+1, j) + let f3 !(i, j) = f (i, j) >> f2 (i+1, j) + let f4 !(i, j) = f (i, j) >> f3 (i+1, j) + let f5 !(i, j) = f (i, j) >> f4 (i+1, j) + let f6 !(i, j) = f (i, j) >> f5 (i+1, j) + let f7 !(i, j) = f (i, j) >> f6 (i+1, j) + let f' = case bH' of + 1 -> f + 2 -> f2 + 3 -> f3 + 4 -> f4 + 5 -> f5 + 6 -> f6 + _ -> f7 + let !ibS = ib - ((ib - it) `mod` bH') + loopM_ it (< ibS) (+ bH') $ \ !i -> + loopM_ jt (< jb) (+ 1) $ \ !j -> + f' (i, j) + loopM_ ibS (< ib) (+ 1) $ \ !i -> + loopM_ jt (< jb) (+ 1) $ \ !j -> + f (i, j) +{-# INLINE unrollAndJamT #-} diff --git a/src/Data/Massiv/Array/Manifest.hs b/src/Data/Massiv/Array/Manifest.hs new file mode 100644 --- /dev/null +++ b/src/Data/Massiv/Array/Manifest.hs @@ -0,0 +1,40 @@ +{-# LANGUAGE FlexibleContexts #-} +{-# LANGUAGE FlexibleInstances #-} +{-# LANGUAGE MultiParamTypeClasses #-} +{-# LANGUAGE ScopedTypeVariables #-} +{-# LANGUAGE TypeFamilies #-} +{-# LANGUAGE TypeOperators #-} +-- | +-- Module : Data.Massiv.Array.Manifest +-- Copyright : (c) Alexey Kuleshevich 2018 +-- License : BSD3 +-- Maintainer : Alexey Kuleshevich +-- Stability : experimental +-- Portability : non-portable +-- +module Data.Massiv.Array.Manifest + ( + -- * Manifest + Manifest + , toManifest + , M + -- * Boxed + , B(..) + , N(..) + -- * Primitive + , P(..) + , Prim + -- * Storable + , S(..) + , Storable + -- * Unboxed + , U(..) + , Unbox + ) where + +import Data.Massiv.Array.Manifest.BoxedStrict +import Data.Massiv.Array.Manifest.BoxedNF +import Data.Massiv.Array.Manifest.Internal +import Data.Massiv.Array.Manifest.Primitive +import Data.Massiv.Array.Manifest.Storable +import Data.Massiv.Array.Manifest.Unboxed diff --git a/src/Data/Massiv/Array/Manifest/BoxedNF.hs b/src/Data/Massiv/Array/Manifest/BoxedNF.hs new file mode 100644 --- /dev/null +++ b/src/Data/Massiv/Array/Manifest/BoxedNF.hs @@ -0,0 +1,214 @@ +{-# LANGUAGE BangPatterns #-} +{-# LANGUAGE FlexibleContexts #-} +{-# LANGUAGE FlexibleInstances #-} +{-# LANGUAGE MultiParamTypeClasses #-} +{-# LANGUAGE ScopedTypeVariables #-} +{-# LANGUAGE TypeFamilies #-} +{-# LANGUAGE UndecidableInstances #-} +-- | +-- Module : Data.Massiv.Array.Manifest.Boxed +-- Copyright : (c) Alexey Kuleshevich 2018 +-- License : BSD3 +-- Maintainer : Alexey Kuleshevich +-- Stability : experimental +-- Portability : non-portable +-- +module Data.Massiv.Array.Manifest.BoxedNF + ( N (..) + , Array(..) + , deepseqArray + , deepseqArrayP + , vectorFromArray + , vectorToArray + , castVectorToArray + ) where + +import Control.DeepSeq (NFData (..), deepseq) +import Control.Monad.ST (runST) +import Data.Massiv.Array.Delayed.Internal (eq) +import Data.Massiv.Array.Manifest.Internal (M, toManifest) +import Data.Massiv.Array.Manifest.List as A +import Data.Massiv.Array.Mutable +import Data.Massiv.Array.Unsafe (unsafeGenerateArray, + unsafeGenerateArrayP) +import Data.Massiv.Core.Common +import Data.Massiv.Core.List +import Data.Massiv.Core.Scheduler +import qualified Data.Primitive.Array as A +import qualified Data.Vector as VB +import qualified Data.Vector.Mutable as VB +import GHC.Exts as GHC (IsList (..)) +import Prelude hiding (mapM) +import System.IO.Unsafe (unsafePerformIO) + +-- | Array representation for Boxed elements. This structure is element and +-- spine strict, and elements are always in Normal Form (NF), therefore `NFData` +-- instance is required. +data N = N deriving Show + +type instance EltRepr N ix = M + +data instance Array N ix e = NArray { nComp :: Comp + , nSize :: !ix + , nData :: {-# UNPACK #-} !(A.Array e) + } + +instance (Index ix, NFData e) => NFData (Array N ix e) where + rnf (NArray comp sz arr) = -- comp `deepseq` sz `deepseq` a `seq` () + case comp of + Seq -> deepseqArray sz arr () + ParOn wIds -> deepseqArrayP wIds sz arr () + + +instance (Index ix, NFData e, Eq e) => Eq (Array N ix e) where + (==) = eq (==) + {-# INLINE (==) #-} + + +instance (Index ix, NFData e) => Construct N ix e where + getComp = nComp + {-# INLINE getComp #-} + + setComp c arr = arr { nComp = c } + {-# INLINE setComp #-} + + unsafeMakeArray Seq !sz f = unsafeGenerateArray sz f + unsafeMakeArray (ParOn wIds) !sz f = unsafeGenerateArrayP wIds sz f + {-# INLINE unsafeMakeArray #-} + +instance (Index ix, NFData e) => Source N ix e where + unsafeLinearIndex (NArray _ _ a) = A.indexArray a + {-# INLINE unsafeLinearIndex #-} + + +instance (Index ix, NFData e) => Size N ix e where + size = nSize + {-# INLINE size #-} + + unsafeResize !sz !arr = arr { nSize = sz } + {-# INLINE unsafeResize #-} + + unsafeExtract !sIx !newSz !arr = unsafeExtract sIx newSz (toManifest arr) + {-# INLINE unsafeExtract #-} + + +instance ( NFData e + , Index ix + , Index (Lower ix) + , Elt M ix e ~ Array M (Lower ix) e + , Elt N ix e ~ Array M (Lower ix) e + ) => + OuterSlice N ix e where + unsafeOuterSlice arr = unsafeOuterSlice (toManifest arr) + {-# INLINE unsafeOuterSlice #-} + +instance ( NFData e + , Index ix + , Index (Lower ix) + , Elt M ix e ~ Array M (Lower ix) e + , Elt N ix e ~ Array M (Lower ix) e + ) => + InnerSlice N ix e where + unsafeInnerSlice arr = unsafeInnerSlice (toManifest arr) + {-# INLINE unsafeInnerSlice #-} + + +instance (Index ix, NFData e) => Manifest N ix e where + + unsafeLinearIndexM (NArray _ _ a) = A.indexArray a + {-# INLINE unsafeLinearIndexM #-} + + +uninitialized :: a +uninitialized = error "Data.Array.Massiv.Manifest.BoxedNF: uninitialized element" + + +instance (Index ix, NFData e) => Mutable N ix e where + data MArray s N ix e = MNArray !ix {-# UNPACK #-} !(A.MutableArray s e) + + msize (MNArray sz _) = sz + {-# INLINE msize #-} + + unsafeThaw (NArray _ sz a) = MNArray sz <$> A.unsafeThawArray a + {-# INLINE unsafeThaw #-} + + unsafeFreeze comp (MNArray sz ma) = NArray comp sz <$> A.unsafeFreezeArray ma + {-# INLINE unsafeFreeze #-} + + unsafeNew sz = MNArray sz <$> A.newArray (totalElem sz) uninitialized + {-# INLINE unsafeNew #-} + + unsafeNewZero = unsafeNew + {-# INLINE unsafeNewZero #-} + + unsafeLinearRead (MNArray _ ma) i = A.readArray ma i + {-# INLINE unsafeLinearRead #-} + + unsafeLinearWrite (MNArray _ ma) i e = e `deepseq` A.writeArray ma i e + {-# INLINE unsafeLinearWrite #-} + + +deepseqArray :: (Index ix, NFData a) => ix -> A.Array a -> b -> b +deepseqArray sz arr b = + iter 0 (totalElem sz) 1 (<) b $ \ !i acc -> A.indexArray arr i `deepseq` acc +{-# INLINE deepseqArray #-} + + +deepseqArrayP :: (Index ix, NFData a) => [Int] -> ix -> A.Array a -> b -> b +deepseqArrayP wIds sz arr b = + unsafePerformIO $ do + divideWork_ wIds sz $ \ !scheduler !chunkLength !totalLength !slackStart -> do + loopM_ 0 (< slackStart) (+ chunkLength) $ \ !start -> + scheduleWork scheduler $ + loopM_ start (< (start + chunkLength)) (+ 1) $ \ !k -> + A.indexArray arr k `deepseq` return () + scheduleWork scheduler $ + loopM_ slackStart (< totalLength) (+ 1) $ \ !k -> + A.indexArray arr k `deepseq` return () + return b +{-# INLINE deepseqArrayP #-} + + +vectorFromArray :: Index ix => ix -> A.Array a -> VB.Vector a +vectorFromArray sz arr = runST $ do + marr <- A.unsafeThawArray arr + VB.unsafeFreeze $ VB.MVector 0 (totalElem sz) marr +{-# INLINE vectorFromArray #-} + + +vectorToArray :: VB.Vector a -> A.Array a +vectorToArray v = + runST $ do + VB.MVector start len marr <- VB.unsafeThaw v + marr' <- + if start == 0 + then return marr + else A.cloneMutableArray marr start len + A.unsafeFreezeArray marr' +{-# INLINE vectorToArray #-} + + +-- | Cast a Boxed Vector into an Array, but only if it wasn't previously sliced. +castVectorToArray :: VB.Vector a -> Maybe (A.Array a) +castVectorToArray v = + runST $ do + VB.MVector start _ marr <- VB.unsafeThaw v + if start == 0 + then Just <$> A.unsafeFreezeArray marr + else return Nothing +{-# INLINE castVectorToArray #-} + + + +instance ( NFData e + , IsList (Array L ix e) + , Nested LN ix e + , Nested L ix e + , Ragged L ix e + ) => + IsList (Array N ix e) where + type Item (Array N ix e) = Item (Array L ix e) + fromList = A.fromLists' Seq + {-# INLINE fromList #-} + toList = GHC.toList . toListArray + {-# INLINE toList #-} diff --git a/src/Data/Massiv/Array/Manifest/BoxedStrict.hs b/src/Data/Massiv/Array/Manifest/BoxedStrict.hs new file mode 100644 --- /dev/null +++ b/src/Data/Massiv/Array/Manifest/BoxedStrict.hs @@ -0,0 +1,176 @@ +{-# LANGUAGE BangPatterns #-} +{-# LANGUAGE FlexibleContexts #-} +{-# LANGUAGE FlexibleInstances #-} +{-# LANGUAGE MultiParamTypeClasses #-} +{-# LANGUAGE ScopedTypeVariables #-} +{-# LANGUAGE TypeFamilies #-} +{-# LANGUAGE UndecidableInstances #-} +-- | +-- Module : Data.Massiv.Array.Manifest.Boxed +-- Copyright : (c) Alexey Kuleshevich 2018 +-- License : BSD3 +-- Maintainer : Alexey Kuleshevich +-- Stability : experimental +-- Portability : non-portable +-- +module Data.Massiv.Array.Manifest.BoxedStrict + ( B (..) + , Array(..) + ) where + +import Control.DeepSeq (NFData (..)) +import qualified Data.Foldable as F (Foldable (..)) +import Data.Massiv.Array.Delayed.Internal (eq) +import Data.Massiv.Array.Manifest.BoxedNF (deepseqArray, + deepseqArrayP) +import Data.Massiv.Array.Unsafe (unsafeGenerateArray, + unsafeGenerateArrayP) +import Data.Massiv.Array.Manifest.Internal +import Data.Massiv.Array.Manifest.List as A +import Data.Massiv.Array.Mutable +import Data.Massiv.Array.Ops.Fold +import Data.Massiv.Core.Common +import Data.Massiv.Core.List +import qualified Data.Primitive.Array as A +import GHC.Base (build) +import GHC.Exts as GHC (IsList (..)) +import Prelude hiding (mapM) + + +-- | Array representation for Boxed elements. This structure is element and +-- spine strict, but elements are strict to Weak Head Normal Form (WHNF) only. +data B = B deriving Show + +type instance EltRepr B ix = M + +data instance Array B ix e = BArray { bComp :: !Comp + , bSize :: !ix + , bData :: {-# UNPACK #-} !(A.Array e) + } + +instance (Index ix, NFData e) => NFData (Array B ix e) where + rnf (BArray comp sz arr) = + case comp of + Seq -> deepseqArray sz arr () + ParOn wIds -> deepseqArrayP wIds sz arr () + +instance (Index ix, Eq e) => Eq (Array B ix e) where + (==) = eq (==) + {-# INLINE (==) #-} + +instance Index ix => Construct B ix e where + getComp = bComp + {-# INLINE getComp #-} + + setComp c arr = arr { bComp = c } + {-# INLINE setComp #-} + + unsafeMakeArray Seq !sz f = unsafeGenerateArray sz f + unsafeMakeArray (ParOn wIds) !sz f = unsafeGenerateArrayP wIds sz f + {-# INLINE unsafeMakeArray #-} + +instance Index ix => Source B ix e where + unsafeLinearIndex (BArray _ _ a) = A.indexArray a + {-# INLINE unsafeLinearIndex #-} + + +instance Index ix => Size B ix e where + size = bSize + {-# INLINE size #-} + + unsafeResize !sz !arr = arr { bSize = sz } + {-# INLINE unsafeResize #-} + + unsafeExtract !sIx !newSz !arr = unsafeExtract sIx newSz (toManifest arr) + {-# INLINE unsafeExtract #-} + + +instance ( NFData e + , Index ix + , Index (Lower ix) + , Elt M ix e ~ Array M (Lower ix) e + , Elt B ix e ~ Array M (Lower ix) e + ) => + OuterSlice B ix e where + unsafeOuterSlice arr = unsafeOuterSlice (toManifest arr) + {-# INLINE unsafeOuterSlice #-} + +instance ( NFData e + , Index ix + , Index (Lower ix) + , Elt M ix e ~ Array M (Lower ix) e + , Elt B ix e ~ Array M (Lower ix) e + ) => + InnerSlice B ix e where + unsafeInnerSlice arr = unsafeInnerSlice (toManifest arr) + {-# INLINE unsafeInnerSlice #-} + + +instance Index ix => Manifest B ix e where + + unsafeLinearIndexM (BArray _ _ a) = A.indexArray a + {-# INLINE unsafeLinearIndexM #-} + + +uninitialized :: a +uninitialized = error "Data.Array.Massiv.Manifest.BoxedStrict: uninitialized element" + + +instance Index ix => Mutable B ix e where + data MArray s B ix e = MBArray !ix {-# UNPACK #-} !(A.MutableArray s e) + + msize (MBArray sz _) = sz + {-# INLINE msize #-} + + unsafeThaw (BArray _ sz a) = MBArray sz <$> A.unsafeThawArray a + {-# INLINE unsafeThaw #-} + + unsafeFreeze comp (MBArray sz ma) = BArray comp sz <$> A.unsafeFreezeArray ma + {-# INLINE unsafeFreeze #-} + + unsafeNew sz = MBArray sz <$> A.newArray (totalElem sz) uninitialized + {-# INLINE unsafeNew #-} + + unsafeNewZero = unsafeNew + {-# INLINE unsafeNewZero #-} + + unsafeLinearRead (MBArray _ ma) i = A.readArray ma i + {-# INLINE unsafeLinearRead #-} + + unsafeLinearWrite (MBArray _ ma) i e = e `seq` A.writeArray ma i e + {-# INLINE unsafeLinearWrite #-} + + +-- | Row-major sequential folding over a Boxed array. +instance Index ix => Foldable (Array B ix) where + foldl = lazyFoldlS + {-# INLINE foldl #-} + foldl' = foldlS + {-# INLINE foldl' #-} + foldr = foldrFB + {-# INLINE foldr #-} + foldr' = foldrS + {-# INLINE foldr' #-} + null (BArray _ sz _) = totalElem sz == 0 + {-# INLINE null #-} + sum = F.foldl' (+) 0 + {-# INLINE sum #-} + product = F.foldl' (*) 1 + {-# INLINE product #-} + length = totalElem . size + {-# INLINE length #-} + toList arr = build (\ c n -> foldrFB c n arr) + {-# INLINE toList #-} + + +instance ( IsList (Array L ix e) + , Nested LN ix e + , Nested L ix e + , Ragged L ix e + ) => + IsList (Array B ix e) where + type Item (Array B ix e) = Item (Array L ix e) + fromList = A.fromLists' Seq + {-# INLINE fromList #-} + toList = GHC.toList . toListArray + {-# INLINE toList #-} diff --git a/src/Data/Massiv/Array/Manifest/Internal.hs b/src/Data/Massiv/Array/Manifest/Internal.hs new file mode 100644 --- /dev/null +++ b/src/Data/Massiv/Array/Manifest/Internal.hs @@ -0,0 +1,320 @@ +{-# LANGUAGE BangPatterns #-} +{-# LANGUAGE FlexibleContexts #-} +{-# LANGUAGE FlexibleInstances #-} +{-# LANGUAGE MultiParamTypeClasses #-} +{-# LANGUAGE ScopedTypeVariables #-} +{-# LANGUAGE TypeFamilies #-} +{-# LANGUAGE TypeOperators #-} +{-# LANGUAGE UndecidableInstances #-} +-- | +-- Module : Data.Massiv.Array.Manifest.Internal +-- Copyright : (c) Alexey Kuleshevich 2018 +-- License : BSD3 +-- Maintainer : Alexey Kuleshevich +-- Stability : experimental +-- Portability : non-portable +-- +module Data.Massiv.Array.Manifest.Internal + ( M + , Manifest(..) + , Array(..) + , makeBoxedVector + , toManifest + , compute + , computeAs + , computeSource + , clone + , convert + , convertAs + , gcastArr + , loadMutableS + , loadMutableOnP + , sequenceP + , sequenceOnP + , fromRaggedArray + , fromRaggedArray' + ) where + +import Control.Exception (try) +import Control.Monad.ST (runST) +import Data.Foldable (Foldable (..)) +import Data.Massiv.Array.Delayed.Internal +import Data.Massiv.Array.Ops.Fold as M +import Data.Massiv.Array.Ops.Map (iforM_) +import Data.Massiv.Array.Unsafe +import Data.Massiv.Core.Common +import Data.Massiv.Core.List +import Data.Massiv.Core.Scheduler +import Data.Maybe (fromMaybe) +import Data.Typeable +import qualified Data.Vector as V +import GHC.Base (build) +import System.IO.Unsafe (unsafePerformIO) + + +-- | General Manifest representation +data M + +data instance Array M ix e = MArray { mComp :: !Comp + , mSize :: !ix + , mUnsafeLinearIndex :: Int -> e } +type instance EltRepr M ix = M + +instance Index ix => Construct M ix e where + getComp = mComp + {-# INLINE getComp #-} + + setComp c arr = arr { mComp = c } + {-# INLINE setComp #-} + + unsafeMakeArray !c !sz f = MArray c sz (V.unsafeIndex (makeBoxedVector sz f)) + {-# INLINE unsafeMakeArray #-} + +-- | Create a boxed from usual size and index to element function +makeBoxedVector :: Index ix => ix -> (ix -> a) -> V.Vector a +makeBoxedVector !sz f = V.generate (totalElem sz) (f . fromLinearIndex sz) +{-# INLINE makeBoxedVector #-} + + +-- | /O(1)/ - Conversion of `Manifest` arrays to `M` representation. +toManifest :: Manifest r ix e => Array r ix e -> Array M ix e +toManifest !arr = MArray (getComp arr) (size arr) (unsafeLinearIndexM arr) where +{-# INLINE toManifest #-} + + +-- | Row-major sequential folding over a Manifest array. +instance Index ix => Foldable (Array M ix) where + foldl = lazyFoldlS + {-# INLINE foldl #-} + foldl' = foldlS + {-# INLINE foldl' #-} + foldr = foldrFB + {-# INLINE foldr #-} + foldr' = foldrS + {-# INLINE foldr' #-} + null (MArray _ sz _) = totalElem sz == 0 + {-# INLINE null #-} + sum = foldl' (+) 0 + {-# INLINE sum #-} + product = foldl' (*) 1 + {-# INLINE product #-} + length = totalElem . size + {-# INLINE length #-} + toList arr = build (\ c n -> foldrFB c n arr) + {-# INLINE toList #-} + + + +instance Index ix => Source M ix e where + unsafeLinearIndex = mUnsafeLinearIndex + {-# INLINE unsafeLinearIndex #-} + + +instance Index ix => Manifest M ix e where + + unsafeLinearIndexM = mUnsafeLinearIndex + {-# INLINE unsafeLinearIndexM #-} + + +instance Index ix => Size M ix e where + size = mSize + {-# INLINE size #-} + + unsafeResize !sz !arr = arr { mSize = sz } + {-# INLINE unsafeResize #-} + + unsafeExtract !sIx !newSz !arr = + MArray (getComp arr) newSz $ \ i -> + unsafeIndex arr (liftIndex2 (+) (fromLinearIndex newSz i) sIx) + {-# INLINE unsafeExtract #-} + + + +instance {-# OVERLAPPING #-} Slice M Ix1 e where + unsafeSlice arr i _ _ = Just (unsafeLinearIndex arr i) + {-# INLINE unsafeSlice #-} + +instance ( Index ix + , Index (Lower ix) + , Elt M ix e ~ Array M (Lower ix) e + ) => + Slice M ix e where + unsafeSlice arr start cutSz dim = do + newSz <- dropDim cutSz dim + return $ unsafeResize newSz (unsafeExtract start cutSz arr) + {-# INLINE unsafeSlice #-} + +instance {-# OVERLAPPING #-} OuterSlice M Ix1 e where + unsafeOuterSlice !arr = unsafeIndex arr + {-# INLINE unsafeOuterSlice #-} + +instance (Elt M ix e ~ Array M (Lower ix) e, Index ix, Index (Lower ix)) => OuterSlice M ix e where + unsafeOuterSlice !arr !i = + MArray (getComp arr) (tailDim (size arr)) (unsafeLinearIndex arr . (+ kStart)) + where + !kStart = toLinearIndex (size arr) (consDim i (zeroIndex :: Lower ix)) + {-# INLINE unsafeOuterSlice #-} + +instance {-# OVERLAPPING #-} InnerSlice M Ix1 e where + unsafeInnerSlice !arr _ = unsafeIndex arr + {-# INLINE unsafeInnerSlice #-} + +instance (Elt M ix e ~ Array M (Lower ix) e, Index ix, Index (Lower ix)) => InnerSlice M ix e where + unsafeInnerSlice !arr !(szL, m) !i = + MArray (getComp arr) szL (\k -> unsafeLinearIndex arr (k * m + kStart)) + where + !kStart = toLinearIndex (size arr) (snocDim (zeroIndex :: Lower ix) i) + {-# INLINE unsafeInnerSlice #-} + + +instance Index ix => Load M ix e where + loadS (MArray _ sz f) _ uWrite = + iterM_ 0 (totalElem sz) 1 (<) $ \ !i -> + uWrite i (f i) + {-# INLINE loadS #-} + loadP wIds (MArray _ sz f) _ uWrite = do + divideWork_ wIds (totalElem sz) $ \ !scheduler !chunkLength !totalLength !slackStart -> do + loopM_ 0 (< slackStart) (+ chunkLength) $ \ !start -> + scheduleWork scheduler $ + iterM_ start (start + chunkLength) 1 (<) $ \ !i -> + uWrite i (f i) + scheduleWork scheduler $ + iterM_ slackStart totalLength 1 (<) $ \ !i -> + uWrite i (f i) + {-# INLINE loadP #-} + + + +loadMutableS :: (Load r' ix e, Mutable r ix e) => + Array r' ix e -> Array r ix e +loadMutableS !arr = + runST $ do + mArr <- unsafeNew (size arr) + loadS arr (unsafeLinearRead mArr) (unsafeLinearWrite mArr) + unsafeFreeze Seq mArr +{-# INLINE loadMutableS #-} + +loadMutableOnP :: (Load r' ix e, Mutable r ix e) => + [Int] -> Array r' ix e -> IO (Array r ix e) +loadMutableOnP wIds !arr = do + mArr <- unsafeNew (size arr) + loadP wIds arr (unsafeLinearRead mArr) (unsafeLinearWrite mArr) + unsafeFreeze (ParOn wIds) mArr +{-# INLINE loadMutableOnP #-} + +-- | Ensure that Array is computed, i.e. represented with concrete elements in memory, hence is the +-- `Mutable` type class restriction. Use `setComp` if you'd like to change computation strategy +-- before calling @compute@ +compute :: (Load r' ix e, Mutable r ix e) => Array r' ix e -> Array r ix e +compute !arr = + case getComp arr of + Seq -> loadMutableS arr + ParOn wIds -> unsafePerformIO $ loadMutableOnP wIds arr +{-# INLINE compute #-} + +-- | Just as `compute`, but let's you supply resulting representation type as an argument. +computeAs :: (Load r' ix e, Mutable r ix e) => r -> Array r' ix e -> Array r ix e +computeAs _ = compute +{-# INLINE computeAs #-} + + +-- | This is just like `compute`, but can be applied to `Source` arrays and will be a noop if +-- resulting type is the same as the input. +computeSource :: forall r' r ix e . (Source r' ix e, Mutable r ix e) + => Array r' ix e -> Array r ix e +computeSource arr = + fromMaybe (compute $ delay arr) $ fmap (\Refl -> arr) (eqT :: Maybe (r' :~: r)) +{-# INLINE computeSource #-} + + +-- | /O(n)/ - Make an exact immutable copy of an Array. +clone :: Mutable r ix e => Array r ix e -> Array r ix e +clone = compute . toManifest +{-# INLINE clone #-} + + +-- | /O(1)/ - Cast over Array representation +gcastArr :: forall r' r ix e. (Typeable r, Typeable r') + => Array r' ix e -> Maybe (Array r ix e) +gcastArr arr = fmap (\Refl -> arr) (eqT :: Maybe (r :~: r')) + + +-- | /O(n)/ - conversion between manifest types, except when source and result arrays +-- are of the same representation, in which case it is an /O(1)/ operation. +convert :: (Manifest r' ix e, Mutable r ix e) + => Array r' ix e -> Array r ix e +convert arr = + fromMaybe (compute $ toManifest arr) (gcastArr arr) +{-# INLINE convert #-} + +-- | Same as `convert`, but let's you supply resulting representation type as an argument. +convertAs :: (Mutable r' ix e, Mutable r ix e, Typeable ix, Typeable e) + => r -> Array r' ix e -> Array r ix e +convertAs _ = convert +{-# INLINE convertAs #-} + + +sequenceOnP :: (Source r1 ix (IO e), Mutable r ix e) => + [Int] -> Array r1 ix (IO e) -> IO (Array r ix e) +sequenceOnP wIds !arr = do + resArrM <- unsafeNew (size arr) + withScheduler_ wIds $ \scheduler -> + iforM_ arr $ \ !ix action -> + scheduleWork scheduler $ action >>= unsafeWrite resArrM ix + unsafeFreeze (getComp arr) resArrM +{-# INLINE sequenceOnP #-} + + +sequenceP :: (Source r1 ix (IO e), Mutable r ix e) => Array r1 ix (IO e) -> IO (Array r ix e) +sequenceP = sequenceOnP [] +{-# INLINE sequenceP #-} + + + + + +-- sequenceOnP' :: (NFData e, Source r1 ix (IO e), Mutable r ix e) => +-- [Int] -> Array r1 ix (IO e) -> IO (Array r ix e) +-- sequenceOnP' wIds !arr = do +-- resArrM <- unsafeNew (size arr) +-- scheduler <- makeScheduler wIds +-- iforM_ arr $ \ !ix action -> +-- submitRequest scheduler $ JobRequest $ do +-- res <- action +-- res `deepseq` unsafeWrite resArrM ix res +-- waitTillDone scheduler +-- unsafeFreeze resArrM +-- {-# INLINE sequenceOnP' #-} + + +-- sequenceP' :: (NFData e, Source r1 ix (IO e), Mutable r ix e) +-- => Array r1 ix (IO e) -> IO (Array r ix e) +-- sequenceP' = sequenceOnP' [] +-- {-# INLINE sequenceP' #-} + +-- | Convert a ragged array into a usual rectangular shaped one. +fromRaggedArray :: (Ragged r' ix e, Mutable r ix e) => + Array r' ix e -> Either ShapeError (Array r ix e) +fromRaggedArray arr = unsafePerformIO $ do + let sz = edgeSize arr + mArr <- unsafeNew sz + let loadWith using = + loadRagged using (unsafeLinearWrite mArr) 0 (totalElem sz) (tailDim sz) arr + try $ case getComp arr of + Seq -> loadWith id >> unsafeFreeze (getComp arr) mArr + ParOn ss -> do + withScheduler_ ss (loadWith . scheduleWork) + unsafeFreeze (getComp arr) mArr +{-# INLINE fromRaggedArray #-} + +-- | Same as `fromRaggedArray`, but will throw an error if its shape is not +-- rectangular. +fromRaggedArray' :: (Ragged r' ix e, Mutable r ix e) => + Array r' ix e -> Array r ix e +fromRaggedArray' arr = + case fromRaggedArray arr of + Left RowTooShortError -> error "Not enough elements in a row" + Left RowTooLongError -> error "Too many elements in a row" + Right resArr -> resArr +{-# INLINE fromRaggedArray' #-} + diff --git a/src/Data/Massiv/Array/Manifest/List.hs b/src/Data/Massiv/Array/Manifest/List.hs new file mode 100644 --- /dev/null +++ b/src/Data/Massiv/Array/Manifest/List.hs @@ -0,0 +1,209 @@ +{-# LANGUAGE BangPatterns #-} +{-# LANGUAGE FlexibleContexts #-} +{-# LANGUAGE FlexibleInstances #-} +{-# LANGUAGE MultiParamTypeClasses #-} +{-# LANGUAGE ScopedTypeVariables #-} +{-# LANGUAGE TypeFamilies #-} +{-# LANGUAGE TypeOperators #-} +-- | +-- Module : Data.Massiv.Array.Manifest.List +-- Copyright : (c) Alexey Kuleshevich 2018 +-- License : BSD3 +-- Maintainer : Alexey Kuleshevich +-- Stability : experimental +-- Portability : non-portable +-- +module Data.Massiv.Array.Manifest.List + ( + -- ** List + fromList + , fromLists + , fromLists' + , toList + , toLists + , toLists2 + , toLists3 + , toLists4 + ) where + +import Data.Massiv.Array.Delayed (D (..)) +import Data.Massiv.Array.Manifest.Internal +import Data.Massiv.Array.Ops.Construct (makeArrayR) +import Data.Massiv.Array.Ops.Fold (foldrFB, foldrS) +import Data.Massiv.Core.Common +import Data.Massiv.Core.List +import GHC.Base (build) + +-- | Convert a flat list into a vector +fromList :: (Nested LN Ix1 e, Nested L Ix1 e, Ragged L Ix1 e, Mutable r Ix1 e) + => Comp -- ^ Computation startegy to use + -> [e] -- ^ Nested list + -> Array r Ix1 e +fromList = fromLists' +{-# INLINE fromList #-} + + + + +-- | /O(n)/ - Convert a nested list into an array. Nested list must be of a rectangular shape, +-- otherwise a runtime error will occur. Also, nestedness must match the rank of resulting array, +-- which should be specified through an explicit type signature. +-- +-- __Note__: This function is almost the same (modulo customizable computation strategy) if you +-- would turn on @{-# LANGUAGE OverloadedLists #-}@. For that reason you can also use +-- `GHC.Exts.fromList`. +-- +-- ==== __Examples__ +-- +-- >>> fromLists Seq [[1,2],[3,4]] :: Maybe (Array U Ix2 Int) +-- Just (Array U Seq (2 :. 2) +-- [ [ 1,2 ] +-- , [ 3,4 ] +-- ]) +-- +-- >>> fromLists Par [[[1,2,3]],[[4,5,6]]] :: Maybe (Array U Ix3 Int) +-- Just (Array U Par (2 :> 1 :. 3) +-- [ [ [ 1,2,3 ] +-- ] +-- , [ [ 4,5,6 ] +-- ] +-- ]) +-- +-- Elements of a boxed array could be lists themselves if necessary, but cannot be ragged: +-- +-- >>> fromLists Seq [[[1,2,3]],[[4,5]]] :: Maybe (Array B Ix2 [Int]) +-- Just (Array B Seq (2 :. 1) +-- [ [ [1,2,3] ] +-- , [ [4,5] ] +-- ]) +-- >>> fromLists Seq [[[1,2,3]],[[4,5]]] :: Maybe (Array B Ix3 Int) +-- Nothing +-- +fromLists :: (Nested LN ix e, Nested L ix e, Ragged L ix e, Mutable r ix e) + => Comp -> [ListItem ix e] -> Maybe (Array r ix e) +fromLists comp = either (const Nothing) Just . fromRaggedArray . setComp comp . throughNested +{-# INLINE fromLists #-} + + +-- | Same as `fromLists`, but will throw an error on irregular shaped lists. +-- +-- ===__Examples__ +-- +-- Convert a list of lists into a 2D Array +-- +-- >>> fromLists' Seq [[1,2],[3,4]] :: Array U Ix2 Int +-- (Array U Seq (2 :. 2) +-- [ [ 1,2 ] +-- , [ 3,4 ] +-- ]) +-- +-- Above example implemented using GHC's `OverloadedLists` extension: +-- +-- >>> :set -XOverloadedLists +-- >>> [[1,2],[3,4]] :: Array U Ix2 Int +-- (Array U Seq (2 :. 2) +-- [ [ 1,2 ] +-- , [ 3,4 ] +-- ]) +-- +-- Example of failure on ceonversion of an irregular nested list. +-- +-- >>> fromLists' Seq [[1],[3,4]] :: Array U Ix2 Int +-- (Array U *** Exception: Too many elements in a row +-- +fromLists' :: (Nested LN ix e, Nested L ix e, Ragged L ix e, Mutable r ix e) + => Comp -- ^ Computation startegy to use + -> [ListItem ix e] -- ^ Nested list + -> Array r ix e +fromLists' comp = fromRaggedArray' . setComp comp . throughNested +{-# INLINE fromLists' #-} + + +throughNested :: forall ix e . (Nested LN ix e, Nested L ix e) => [ListItem ix e] -> Array L ix e +throughNested xs = fromNested (fromNested xs :: Array LN ix e) +{-# INLINE throughNested #-} + + + +-- | Convert any array to a flat list. +-- +-- ==== __Examples__ +-- +-- >>> toList $ makeArrayR U Seq (2 :. 3) fromIx2 +-- [(0,0),(0,1),(0,2),(1,0),(1,1),(1,2)] +-- +toList :: Source r ix e => Array r ix e -> [e] +toList !arr = build (\ c n -> foldrFB c n arr) +{-# INLINE toList #-} + + +-- | /O(n)/ - Convert an array into a nested list. Array rank and list nestedness will always match, +-- but you can use `toList`, `toLists2`, etc. if flattening of inner dimensions is desired. +-- +-- __Note__: This function is almost the same as `GHC.Exts.toList`. +-- +-- ====__Examples__ +-- +-- >>> let arr = makeArrayR U Seq (2 :> 1 :. 3) fromIx3 +-- >>> print arr +-- (Array U Seq (2 :> 1 :. 3) +-- [ [ [ (0,0,0),(0,0,1),(0,0,2) ] +-- ] +-- , [ [ (1,0,0),(1,0,1),(1,0,2) ] +-- ] +-- ]) +-- >>> toList arr +-- [[[(0,0,0),(0,0,1),(0,0,2)]],[[(1,0,0),(1,0,1),(1,0,2)]]] +-- +toLists :: (Nested LN ix e, Nested L ix e, Construct L ix e, Source r ix e) + => Array r ix e + -> [ListItem ix e] +toLists = toNested . toNested . toListArray +{-# INLINE toLists #-} + + + +-- | Convert an array with at least 2 dimensions into a list of lists. Inner dimensions will get +-- flattened. +-- +-- ==== __Examples__ +-- +-- >>> toList2 $ makeArrayR U Seq (2 :. 3) fromIx2 +-- [[(0,0),(0,1),(0,2)],[(1,0),(1,1),(1,2)]] +-- >>> toList2 $ makeArrayR U Seq (2 :> 1 :. 3) fromIx3 +-- [[(0,0,0),(0,0,1),(0,0,2)],[(1,0,0),(1,0,1),(1,0,2)]] +-- +toLists2 :: (Source r ix e, Index (Lower ix)) => Array r ix e -> [[e]] +toLists2 = toList . foldrInner (:) [] +{-# INLINE toLists2 #-} + + +-- | Convert an array with at least 3 dimensions into a 3 deep nested list. Inner dimensions will +-- get flattened. +toLists3 :: (Index (Lower (Lower ix)), Index (Lower ix), Source r ix e) => Array r ix e -> [[[e]]] +toLists3 = toList . foldrInner (:) [] . foldrInner (:) [] +{-# INLINE toLists3 #-} + +-- | Convert an array with at least 4 dimensions into a 4 deep nested list. Inner dimensions will +-- get flattened. +toLists4 :: + ( Index (Lower (Lower (Lower ix))) + , Index (Lower (Lower ix)) + , Index (Lower ix) + , Source r ix e + ) + => Array r ix e + -> [[[[e]]]] +toLists4 = toList . foldrInner (:) [] . foldrInner (:) [] . foldrInner (:) [] +{-# INLINE toLists4 #-} + + +-- | Right fold with an index aware function of inner most dimension. +foldrInner :: (Source r ix e, Index (Lower ix)) => + (e -> a -> a) -> a -> Array r ix e -> Array D (Lower ix) a +foldrInner f !acc !arr = + unsafeMakeArray (getComp arr) szL $ \ !ix -> + foldrS f acc $ makeArrayR D Seq m (unsafeIndex arr . snocDim ix) + where + !(szL, m) = unsnocDim (size arr) +{-# INLINE foldrInner #-} diff --git a/src/Data/Massiv/Array/Manifest/Primitive.hs b/src/Data/Massiv/Array/Manifest/Primitive.hs new file mode 100644 --- /dev/null +++ b/src/Data/Massiv/Array/Manifest/Primitive.hs @@ -0,0 +1,190 @@ +{-# LANGUAGE BangPatterns #-} +{-# LANGUAGE FlexibleContexts #-} +{-# LANGUAGE FlexibleInstances #-} +{-# LANGUAGE MultiParamTypeClasses #-} +{-# LANGUAGE ScopedTypeVariables #-} +{-# LANGUAGE TypeFamilies #-} +{-# LANGUAGE UndecidableInstances #-} +-- | +-- Module : Data.Massiv.Array.Manifest.Primitive +-- Copyright : (c) Alexey Kuleshevich 2018 +-- License : BSD3 +-- Maintainer : Alexey Kuleshevich +-- Stability : experimental +-- Portability : non-portable +-- +module Data.Massiv.Array.Manifest.Primitive + ( P(..) + , Array(..) + , Prim + , vectorToByteArray + ) where + +import Control.DeepSeq (NFData (..), deepseq) +import Control.Monad.ST (runST) +import Data.Massiv.Array.Delayed.Internal (eq) +import Data.Massiv.Array.Manifest.Internal +import Data.Massiv.Array.Manifest.List as A +import Data.Massiv.Array.Mutable +import Data.Massiv.Array.Unsafe (unsafeGenerateArray, + unsafeGenerateArrayP) +import Data.Massiv.Core.Common +import Data.Massiv.Core.List +import Data.Primitive (sizeOf) +import Data.Primitive.ByteArray +import Data.Primitive.Types (Prim) +import qualified Data.Vector.Primitive as VP +import GHC.Exts as GHC (IsList (..)) +import Prelude hiding (mapM) + +-- | Representation for `Prim`itive elements +data P = P deriving Show + +type instance EltRepr P ix = M + +data instance Array P ix e = PArray { pComp :: !Comp + , pSize :: !ix + , pData :: {-# UNPACK #-} !ByteArray + } + +instance (Index ix, NFData e) => NFData (Array P ix e) where + rnf (PArray c sz a) = c `deepseq` sz `deepseq` a `seq` () + {-# INLINE rnf #-} + +instance (Prim e, Eq e, Index ix) => Eq (Array P ix e) where + (==) = eq (==) + {-# INLINE (==) #-} + + +instance (Prim e, Index ix) => Construct P ix e where + getComp = pComp + {-# INLINE getComp #-} + + setComp c arr = arr { pComp = c } + {-# INLINE setComp #-} + + unsafeMakeArray Seq !sz f = unsafeGenerateArray sz f + unsafeMakeArray (ParOn wIds) !sz f = unsafeGenerateArrayP wIds sz f + {-# INLINE unsafeMakeArray #-} + +instance (Prim e, Index ix) => Source P ix e where + unsafeLinearIndex (PArray _ _ a) = indexByteArray a + {-# INLINE unsafeLinearIndex #-} + + +instance (Prim e, Index ix) => Size P ix e where + size = pSize + {-# INLINE size #-} + + unsafeResize !sz !arr = arr { pSize = sz } + {-# INLINE unsafeResize #-} + + unsafeExtract !sIx !newSz !arr = unsafeExtract sIx newSz (toManifest arr) + {-# INLINE unsafeExtract #-} + + +instance {-# OVERLAPPING #-} Prim e => Slice P Ix1 e where + unsafeSlice arr i _ _ = Just (unsafeLinearIndex arr i) + {-# INLINE unsafeSlice #-} + + +instance ( Prim e + , Index ix + , Index (Lower ix) + , Elt P ix e ~ Elt M ix e + , Elt M ix e ~ Array M (Lower ix) e + ) => + Slice P ix e where + unsafeSlice arr = unsafeSlice (toManifest arr) + {-# INLINE unsafeSlice #-} + +instance {-# OVERLAPPING #-} Prim e => OuterSlice P Ix1 e where + unsafeOuterSlice = unsafeLinearIndex + {-# INLINE unsafeOuterSlice #-} + +instance ( Prim e + , Index ix + , Index (Lower ix) + , Elt M ix e ~ Array M (Lower ix) e + , Elt P ix e ~ Array M (Lower ix) e + ) => + OuterSlice P ix e where + unsafeOuterSlice arr = unsafeOuterSlice (toManifest arr) + {-# INLINE unsafeOuterSlice #-} + + +instance {-# OVERLAPPING #-} Prim e => InnerSlice P Ix1 e where + unsafeInnerSlice arr _ = unsafeLinearIndex arr + {-# INLINE unsafeInnerSlice #-} + +instance ( Prim e + , Index ix + , Index (Lower ix) + , Elt M ix e ~ Array M (Lower ix) e + , Elt P ix e ~ Array M (Lower ix) e + ) => + InnerSlice P ix e where + unsafeInnerSlice arr = unsafeInnerSlice (toManifest arr) + {-# INLINE unsafeInnerSlice #-} + +instance (Index ix, Prim e) => Manifest P ix e where + + unsafeLinearIndexM (PArray _ _ a) = indexByteArray a + {-# INLINE unsafeLinearIndexM #-} + + +instance (Index ix, Prim e) => Mutable P ix e where + data MArray s P ix e = MPArray !ix !(MutableByteArray s) + + msize (MPArray sz _) = sz + {-# INLINE msize #-} + + unsafeThaw (PArray _ sz a) = MPArray sz <$> unsafeThawByteArray a + {-# INLINE unsafeThaw #-} + + unsafeFreeze comp (MPArray sz a) = PArray comp sz <$> unsafeFreezeByteArray a + {-# INLINE unsafeFreeze #-} + + unsafeNew sz = MPArray sz <$> newByteArray (totalElem sz * sizeOf (undefined :: e)) + {-# INLINE unsafeNew #-} + + unsafeNewZero sz = do + let szBytes = totalElem sz * sizeOf (undefined :: e) + barr <- newByteArray szBytes + fillByteArray barr 0 szBytes 0 + return $ MPArray sz barr + {-# INLINE unsafeNewZero #-} + + unsafeLinearRead (MPArray _ a) = readByteArray a + {-# INLINE unsafeLinearRead #-} + + unsafeLinearWrite (MPArray _ v) = writeByteArray v + {-# INLINE unsafeLinearWrite #-} + + +instance ( VP.Prim e + , IsList (Array L ix e) + , Nested LN ix e + , Nested L ix e + , Ragged L ix e + ) => + IsList (Array P ix e) where + type Item (Array P ix e) = Item (Array L ix e) + fromList = A.fromLists' Seq + {-# INLINE fromList #-} + toList = GHC.toList . toListArray + {-# INLINE toList #-} + + +vectorToByteArray :: forall e . VP.Prim e => VP.Vector e -> ByteArray +vectorToByteArray (VP.Vector start len arr) = + if start == 0 + then arr + else runST $ do + marr <- newByteArray len + let elSize = sizeOf (undefined :: e) + copyByteArray marr 0 arr (start * elSize) (len * elSize) + unsafeFreezeByteArray marr +{-# INLINE vectorToByteArray #-} + + diff --git a/src/Data/Massiv/Array/Manifest/Storable.hs b/src/Data/Massiv/Array/Manifest/Storable.hs new file mode 100644 --- /dev/null +++ b/src/Data/Massiv/Array/Manifest/Storable.hs @@ -0,0 +1,145 @@ +{-# LANGUAGE BangPatterns #-} +{-# LANGUAGE FlexibleContexts #-} +{-# LANGUAGE FlexibleInstances #-} +{-# LANGUAGE MultiParamTypeClasses #-} +{-# LANGUAGE ScopedTypeVariables #-} +{-# LANGUAGE TypeFamilies #-} +{-# LANGUAGE UndecidableInstances #-} +-- | +-- Module : Data.Massiv.Array.Manifest.Storable +-- Copyright : (c) Alexey Kuleshevich 2018 +-- License : BSD3 +-- Maintainer : Alexey Kuleshevich +-- Stability : experimental +-- Portability : non-portable +-- +module Data.Massiv.Array.Manifest.Storable + ( S (..) + , Array(..) + , VS.Storable + ) where + +import Control.DeepSeq (NFData (..), deepseq) +import Data.Massiv.Array.Delayed.Internal (eq) +import Data.Massiv.Array.Manifest.Internal +import Data.Massiv.Array.Manifest.List as A +import Data.Massiv.Array.Mutable +import Data.Massiv.Array.Unsafe (unsafeGenerateArray, + unsafeGenerateArrayP) +import Data.Massiv.Core.Common +import Data.Massiv.Core.List +import qualified Data.Vector.Storable as VS +import qualified Data.Vector.Storable.Mutable as MVS +import GHC.Exts as GHC (IsList (..)) +import Prelude hiding (mapM) + +-- | Representation for `Storable` elements +data S = S deriving Show + +type instance EltRepr S ix = M + +data instance Array S ix e = SArray { sComp :: !Comp + , sSize :: !ix + , sData :: !(VS.Vector e) + } + +instance (Index ix, NFData e) => NFData (Array S ix e) where + rnf (SArray c sz v) = c `deepseq` sz `deepseq` v `deepseq` () + +instance (VS.Storable e, Eq e, Index ix) => Eq (Array S ix e) where + (==) = eq (==) + {-# INLINE (==) #-} + +instance (VS.Storable e, Index ix) => Construct S ix e where + getComp = sComp + {-# INLINE getComp #-} + + setComp c arr = arr { sComp = c } + {-# INLINE setComp #-} + + unsafeMakeArray Seq !sz f = unsafeGenerateArray sz f + unsafeMakeArray (ParOn wIds) !sz f = unsafeGenerateArrayP wIds sz f + {-# INLINE unsafeMakeArray #-} + + +instance (VS.Storable e, Index ix) => Source S ix e where + unsafeLinearIndex (SArray _ _ v) = VS.unsafeIndex v + {-# INLINE unsafeLinearIndex #-} + + +instance (VS.Storable e, Index ix) => Size S ix e where + size = sSize + {-# INLINE size #-} + + unsafeResize !sz !arr = arr { sSize = sz } + {-# INLINE unsafeResize #-} + + unsafeExtract !sIx !newSz !arr = unsafeExtract sIx newSz (toManifest arr) + {-# INLINE unsafeExtract #-} + + + +instance ( VS.Storable e + , Index ix + , Index (Lower ix) + , Elt M ix e ~ Array M (Lower ix) e + , Elt S ix e ~ Array M (Lower ix) e + ) => + OuterSlice S ix e where + unsafeOuterSlice arr = unsafeOuterSlice (toManifest arr) + {-# INLINE unsafeOuterSlice #-} + +instance ( VS.Storable e + , Index ix + , Index (Lower ix) + , Elt M ix e ~ Array M (Lower ix) e + , Elt S ix e ~ Array M (Lower ix) e + ) => + InnerSlice S ix e where + unsafeInnerSlice arr = unsafeInnerSlice (toManifest arr) + {-# INLINE unsafeInnerSlice #-} + + +instance (Index ix, VS.Storable e) => Manifest S ix e where + + unsafeLinearIndexM (SArray _ _ v) = VS.unsafeIndex v + {-# INLINE unsafeLinearIndexM #-} + + +instance (Index ix, VS.Storable e) => Mutable S ix e where + data MArray s S ix e = MSArray !ix !(VS.MVector s e) + + msize (MSArray sz _) = sz + {-# INLINE msize #-} + + unsafeThaw (SArray _ sz v) = MSArray sz <$> VS.unsafeThaw v + {-# INLINE unsafeThaw #-} + + unsafeFreeze comp (MSArray sz v) = SArray comp sz <$> VS.unsafeFreeze v + {-# INLINE unsafeFreeze #-} + + unsafeNew sz = MSArray sz <$> MVS.unsafeNew (totalElem sz) + {-# INLINE unsafeNew #-} + + unsafeNewZero sz = MSArray sz <$> MVS.new (totalElem sz) + {-# INLINE unsafeNewZero #-} + + unsafeLinearRead (MSArray _ v) i = MVS.unsafeRead v i + {-# INLINE unsafeLinearRead #-} + + unsafeLinearWrite (MSArray _ v) i = MVS.unsafeWrite v i + {-# INLINE unsafeLinearWrite #-} + + +instance ( VS.Storable e + , IsList (Array L ix e) + , Nested LN ix e + , Nested L ix e + , Ragged L ix e + ) => + IsList (Array S ix e) where + type Item (Array S ix e) = Item (Array L ix e) + fromList = A.fromLists' Seq + {-# INLINE fromList #-} + toList = GHC.toList . toListArray + {-# INLINE toList #-} diff --git a/src/Data/Massiv/Array/Manifest/Unboxed.hs b/src/Data/Massiv/Array/Manifest/Unboxed.hs new file mode 100644 --- /dev/null +++ b/src/Data/Massiv/Array/Manifest/Unboxed.hs @@ -0,0 +1,169 @@ +{-# LANGUAGE BangPatterns #-} +{-# LANGUAGE FlexibleContexts #-} +{-# LANGUAGE FlexibleInstances #-} +{-# LANGUAGE MultiParamTypeClasses #-} +{-# LANGUAGE ScopedTypeVariables #-} +{-# LANGUAGE TypeFamilies #-} +{-# LANGUAGE UndecidableInstances #-} +-- | +-- Module : Data.Massiv.Array.Manifest.Unboxed +-- Copyright : (c) Alexey Kuleshevich 2018 +-- License : BSD3 +-- Maintainer : Alexey Kuleshevich +-- Stability : experimental +-- Portability : non-portable +-- +module Data.Massiv.Array.Manifest.Unboxed + ( U (..) + , VU.Unbox + , Array(..) + ) where + +import Control.DeepSeq (NFData (..), deepseq) +import Data.Massiv.Array.Delayed.Internal (eq) +import Data.Massiv.Array.Manifest.Internal (M, toManifest) +import Data.Massiv.Array.Manifest.List as A +import Data.Massiv.Array.Mutable +import Data.Massiv.Array.Unsafe (unsafeGenerateArray, + unsafeGenerateArrayP) +import Data.Massiv.Core.Common +import Data.Massiv.Core.List +import qualified Data.Vector.Unboxed as VU +import qualified Data.Vector.Unboxed.Mutable as MVU +import GHC.Exts as GHC (IsList (..)) +import Prelude hiding (mapM) + +-- | Representation for `Unbox`ed elements +data U = U deriving Show + +type instance EltRepr U ix = M + +data instance Array U ix e = UArray { uComp :: !Comp + , uSize :: !ix + , uData :: !(VU.Vector e) + } + + +instance (Index ix, NFData e) => NFData (Array U ix e) where + rnf (UArray c sz v) = c `deepseq` sz `deepseq` v `deepseq` () + + +instance (VU.Unbox e, Index ix) => Construct U ix e where + getComp = uComp + {-# INLINE getComp #-} + + setComp c arr = arr { uComp = c } + {-# INLINE setComp #-} + + unsafeMakeArray Seq !sz f = unsafeGenerateArray sz f + unsafeMakeArray (ParOn wIds) !sz f = unsafeGenerateArrayP wIds sz f + {-# INLINE unsafeMakeArray #-} + + +instance (VU.Unbox e, Eq e, Index ix) => Eq (Array U ix e) where + (==) = eq (==) + {-# INLINE (==) #-} + + +instance (VU.Unbox e, Index ix) => Source U ix e where + unsafeLinearIndex (UArray _ _ v) = VU.unsafeIndex v + {-# INLINE unsafeLinearIndex #-} + + +instance (VU.Unbox e, Index ix) => Size U ix e where + size = uSize + {-# INLINE size #-} + + unsafeResize !sz !arr = arr { uSize = sz } + {-# INLINE unsafeResize #-} + + unsafeExtract !sIx !newSz !arr = unsafeExtract sIx newSz (toManifest arr) + {-# INLINE unsafeExtract #-} + + +instance {-# OVERLAPPING #-} VU.Unbox e => Slice U Ix1 e where + unsafeSlice arr i _ _ = Just (unsafeLinearIndex arr i) + {-# INLINE unsafeSlice #-} + + +instance ( VU.Unbox e + , Index ix + , Index (Lower ix) + , Elt U ix e ~ Elt M ix e + , Elt M ix e ~ Array M (Lower ix) e + ) => + Slice U ix e where + unsafeSlice arr = unsafeSlice (toManifest arr) + {-# INLINE unsafeSlice #-} + + +instance {-# OVERLAPPING #-} VU.Unbox e => OuterSlice U Ix1 e where + unsafeOuterSlice = unsafeLinearIndex + {-# INLINE unsafeOuterSlice #-} + +instance ( VU.Unbox e + , Index ix + , Index (Lower ix) + , Elt U ix e ~ Elt M ix e + , Elt M ix e ~ Array M (Lower ix) e + ) => + OuterSlice U ix e where + unsafeOuterSlice arr = unsafeOuterSlice (toManifest arr) + {-# INLINE unsafeOuterSlice #-} + +instance {-# OVERLAPPING #-} VU.Unbox e => InnerSlice U Ix1 e where + unsafeInnerSlice arr _ = unsafeLinearIndex arr + {-# INLINE unsafeInnerSlice #-} + +instance ( VU.Unbox e + , Index ix + , Index (Lower ix) + , Elt U ix e ~ Elt M ix e + , Elt M ix e ~ Array M (Lower ix) e + ) => + InnerSlice U ix e where + unsafeInnerSlice arr = unsafeInnerSlice (toManifest arr) + {-# INLINE unsafeInnerSlice #-} + +instance (VU.Unbox e, Index ix) => Manifest U ix e where + + unsafeLinearIndexM (UArray _ _ v) = VU.unsafeIndex v + {-# INLINE unsafeLinearIndexM #-} + +instance (VU.Unbox e, Index ix) => Mutable U ix e where + data MArray s U ix e = MUArray ix (VU.MVector s e) + + msize (MUArray sz _) = sz + {-# INLINE msize #-} + + unsafeThaw (UArray _ sz v) = MUArray sz <$> VU.unsafeThaw v + {-# INLINE unsafeThaw #-} + + unsafeFreeze comp (MUArray sz v) = UArray comp sz <$> VU.unsafeFreeze v + {-# INLINE unsafeFreeze #-} + + unsafeNew sz = MUArray sz <$> MVU.unsafeNew (totalElem sz) + {-# INLINE unsafeNew #-} + + unsafeNewZero sz = MUArray sz <$> MVU.new (totalElem sz) + {-# INLINE unsafeNewZero #-} + + unsafeLinearRead (MUArray _ v) i = MVU.unsafeRead v i + {-# INLINE unsafeLinearRead #-} + + unsafeLinearWrite (MUArray _ v) i = MVU.unsafeWrite v i + {-# INLINE unsafeLinearWrite #-} + + +instance ( VU.Unbox e + , IsList (Array L ix e) + , Nested LN ix e + , Nested L ix e + , Ragged L ix e + ) => + IsList (Array U ix e) where + type Item (Array U ix e) = Item (Array L ix e) + fromList = A.fromLists' Seq + {-# INLINE fromList #-} + toList = GHC.toList . toListArray + {-# INLINE toList #-} diff --git a/src/Data/Massiv/Array/Manifest/Vector.hs b/src/Data/Massiv/Array/Manifest/Vector.hs new file mode 100644 --- /dev/null +++ b/src/Data/Massiv/Array/Manifest/Vector.hs @@ -0,0 +1,166 @@ +{-# LANGUAGE FlexibleContexts #-} +{-# LANGUAGE FlexibleInstances #-} +{-# LANGUAGE MultiParamTypeClasses #-} +{-# LANGUAGE ScopedTypeVariables #-} +{-# LANGUAGE TypeFamilies #-} +{-# LANGUAGE TypeOperators #-} +-- | +-- Module : Data.Massiv.Array.Manifest.Vector +-- Copyright : (c) Alexey Kuleshevich 2018 +-- License : BSD3 +-- Maintainer : Alexey Kuleshevich +-- Stability : experimental +-- Portability : non-portable +-- +module Data.Massiv.Array.Manifest.Vector + ( fromVector + , castFromVector + , toVector + , castToVector + , ARepr + , VRepr + ) where + +import Control.Monad (guard, join, msum) +import Data.Massiv.Array.Manifest.BoxedNF +import Data.Massiv.Array.Manifest.BoxedStrict +import Data.Massiv.Array.Manifest.Internal +import Data.Massiv.Array.Manifest.Primitive +import Data.Massiv.Array.Manifest.Storable +import Data.Massiv.Array.Manifest.Unboxed +import Data.Massiv.Array.Mutable +import Data.Massiv.Core.Common +import Data.Typeable +import qualified Data.Vector as VB +import qualified Data.Vector.Generic as VG +import qualified Data.Vector.Primitive as VP +import qualified Data.Vector.Storable as VS +import qualified Data.Vector.Unboxed as VU + +-- | Match vector type to array representation +type family ARepr (v :: * -> *) :: * +-- | Match array representation to a vector type +type family VRepr r :: * -> * + +type instance ARepr VU.Vector = U +type instance ARepr VS.Vector = S +type instance ARepr VP.Vector = P +type instance ARepr VB.Vector = B +type instance VRepr U = VU.Vector +type instance VRepr S = VS.Vector +type instance VRepr P = VP.Vector +type instance VRepr B = VB.Vector +type instance VRepr N = VB.Vector + + +-- | /O(1)/ - conversion from vector to an array with a corresponding +-- representation. Will return `Nothing` if there is a size mismatch, vector has +-- been sliced before or if some non-standard vector type is supplied. +castFromVector :: forall v r ix e. (VG.Vector v e, Typeable v, Mutable r ix e, ARepr v ~ r) + => Comp + -> ix -- ^ Size of the result Array + -> v e -- ^ Source Vector + -> Maybe (Array r ix e) +castFromVector comp sz vector = do + guard (totalElem sz == VG.length vector) + msum + [ do Refl <- eqT :: Maybe (v :~: VU.Vector) + uVector <- join $ gcast1 (Just vector) + return $ UArray {uComp = comp, uSize = sz, uData = uVector} + , do Refl <- eqT :: Maybe (v :~: VS.Vector) + sVector <- join $ gcast1 (Just vector) + return $ SArray {sComp = comp, sSize = sz, sData = sVector} + , do Refl <- eqT :: Maybe (v :~: VP.Vector) + VP.Vector 0 _ arr <- join $ gcast1 (Just vector) + return $ PArray {pComp = comp, pSize = sz, pData = arr} + , do Refl <- eqT :: Maybe (v :~: VB.Vector) + bVector <- join $ gcast1 (Just vector) + arr <- castVectorToArray bVector + return $ BArray {bComp = comp, bSize = sz, bData = arr} + ] +{-# NOINLINE castFromVector #-} + + +-- | In case when resulting array representation matches the one of vector's it +-- will do a /O(1)/ - conversion using `castFromVector`, otherwise Vector elements +-- will be copied into a new array. Will throw an error if length of resulting +-- array doesn't match the source vector length. +fromVector :: + (Typeable v, VG.Vector v a, Mutable (ARepr v) ix a, Mutable r ix a) + => Comp + -> ix -- ^ Resulting size of the array + -> v a -- ^ Source Vector + -> Array r ix a +fromVector comp sz v = + case castFromVector comp sz v of + Just arr -> convert arr + Nothing -> + if (totalElem sz /= VG.length v) + then error $ + "Data.Array.Massiv.Manifest.fromVector: Supplied size: " ++ + show sz ++ " doesn't match vector length: " ++ show (VG.length v) + else unsafeMakeArray comp sz ((v VG.!) . toLinearIndex sz) +{-# NOINLINE fromVector #-} + + +-- | /O(1)/ - conversion from `Mutable` array to a corresponding vector. Will +-- return `Nothing` only if source array representation was not one of `B`, `N`, +-- `P`, `S` or `U`. +castToVector :: forall v r ix e . (VG.Vector v e, Mutable r ix e, VRepr r ~ v) + => Array r ix e -> Maybe (v e) +castToVector arr = + msum + [ do Refl <- eqT :: Maybe (r :~: U) + uArr <- gcastArr arr + return $ uData uArr + , do Refl <- eqT :: Maybe (r :~: S) + sArr <- gcastArr arr + return $ sData sArr + , do Refl <- eqT :: Maybe (r :~: P) + pArr <- gcastArr arr + return $ VP.Vector 0 (totalElem (size arr)) $ pData pArr + , do Refl <- eqT :: Maybe (r :~: B) + bArr <- gcastArr arr + return $ vectorFromArray (size arr) $ bData bArr + , do Refl <- eqT :: Maybe (r :~: N) + bArr <- gcastArr arr + return $ vectorFromArray (size arr) $ nData bArr + ] +{-# NOINLINE castToVector #-} + + +-- | Convert an array into a vector. Will perform a cast if resulting vector is +-- of compatible representation, otherwise memory copy will occur. +-- +-- ==== __Examples__ +-- +-- In this example a `S`torable Array is created and then casted into a Storable +-- `VS.Vector` in costant time: +-- +-- >>> import qualified Data.Vector.Storable as VS +-- >>> toVector (makeArrayR S Par (5 :. 6) (\(i :. j) -> i + j)) :: VS.Vector Int +-- [0,1,2,3,4,5,1,2,3,4,5,6,2,3,4,5,6,7,3,4,5,6,7,8,4,5,6,7,8,9] +-- +-- While in this example `S`torable Array will first be converted into `U`nboxed +-- representation in `Par`allel and only after that will be coverted into Unboxed +-- `VU.Vector` in constant time. +-- +-- >>> import qualified Data.Vector.Unboxed as VU +-- >>> toVector (makeArrayR S Par (5 :. 6) (\(i :. j) -> i + j)) :: VU.Vector Int +-- [0,1,2,3,4,5,1,2,3,4,5,6,2,3,4,5,6,7,3,4,5,6,7,8,4,5,6,7,8,9] +-- +toVector :: + forall r ix e v. + ( Manifest r ix e + , Mutable (ARepr v) ix e + , VG.Vector v e + , VRepr (ARepr v) ~ v + ) + => Array r ix e + -> v e +toVector arr = + case castToVector (convert arr :: Array (ARepr v) ix e) of + Just v -> v + Nothing -> VG.generate (totalElem (size arr)) (unsafeLinearIndex arr) +{-# NOINLINE toVector #-} + diff --git a/src/Data/Massiv/Array/Mutable.hs b/src/Data/Massiv/Array/Mutable.hs new file mode 100644 --- /dev/null +++ b/src/Data/Massiv/Array/Mutable.hs @@ -0,0 +1,158 @@ +{-# LANGUAGE FlexibleContexts #-} +{-# LANGUAGE MultiParamTypeClasses #-} +{-# LANGUAGE ScopedTypeVariables #-} +{-# LANGUAGE TypeFamilies #-} +{-# LANGUAGE TypeOperators #-} +-- | +-- Module : Data.Massiv.Array.Mutable +-- Copyright : (c) Alexey Kuleshevich 2018 +-- License : BSD3 +-- Maintainer : Alexey Kuleshevich +-- Stability : experimental +-- Portability : non-portable +-- +module Data.Massiv.Array.Mutable + ( Mutable + , MArray + , msize + , new + , thaw + , freeze + , read + , read' + , write + , write' + , modify + , modify' + , swap + , swap' + ) where + +import Prelude hiding (read) + +import Control.Monad (unless) +import Control.Monad.Primitive (PrimMonad (..)) +import Data.Massiv.Array.Manifest.Internal +import Data.Massiv.Array.Unsafe +import Data.Massiv.Core.Common + +-- errorSizeMismatch fName sz1 sz2 = +-- error $ fName ++ ": Size mismatch: " ++ show sz1 ++ " /= " ++ show sz2 +-- -- TODO: make sure copy is done in parallel as well as sequentially +-- copy mTargetArr sourceArr = do +-- unless (msize mTargetArr == size sourceArr) $ +-- errorSizeMismatch "Data.Massiv.Array.Mutable.copy" (msize mTargetArr) (size sourceArr) +-- mSourdceArray <- unsafeThaw sourceArray +-- -- comp from marr +-- -- TODO: use load +-- imapM_ (unsafeWrite) + + + +-- | Initialize a new mutable array. Negative size will result in an empty array. +new :: (Mutable r ix e, PrimMonad m) => ix -> m (MArray (PrimState m) r ix e) +new sz = unsafeNewZero (liftIndex (max 0) sz) +{-# INLINE new #-} + +-- | /O(n)/ - Yield a mutable copy of the immutable array +thaw :: (Mutable r ix e, PrimMonad m) => Array r ix e -> m (MArray (PrimState m) r ix e) +thaw = unsafeThaw . clone +{-# INLINE thaw #-} + +-- | /O(n)/ - Yield an immutable copy of the mutable array +freeze :: (Mutable r ix e, PrimMonad m) => Comp -> MArray (PrimState m) r ix e -> m (Array r ix e) +freeze comp marr = unsafeFreeze comp marr >>= (return . clone) +{-# INLINE freeze #-} + + +-- | /O(1)/ - Lookup an element in the mutable array. Return `Nothing` when index is out of bounds. +read :: (Mutable r ix e, PrimMonad m) => + MArray (PrimState m) r ix e -> ix -> m (Maybe e) +read marr ix = + if isSafeIndex (msize marr) ix + then Just <$> unsafeRead marr ix + else return Nothing +{-# INLINE read #-} + + +-- | /O(1)/ - Same as `read`, but throws an error if index is out of bounds. +read' :: (Mutable r ix e, PrimMonad m) => + MArray (PrimState m) r ix e -> ix -> m e +read' marr ix = do + mval <- read marr ix + case mval of + Just e -> return e + Nothing -> errorIx "Data.Massiv.Array.Mutable.read'" (msize marr) ix +{-# INLINE read' #-} + + +-- | /O(1)/ - Write an element into the cell of a mutable array. Returns `False` when index is out +-- of bounds. +write :: (Mutable r ix e, PrimMonad m) => + MArray (PrimState m) r ix e -> ix -> e -> m Bool +write marr ix e = + if isSafeIndex (msize marr) ix + then unsafeWrite marr ix e >> return True + else return False +{-# INLINE write #-} + + +-- | /O(1)/ - Same as `write`, but throws an error if index is out of bounds. +write' :: (Mutable r ix e, PrimMonad m) => + MArray (PrimState m) r ix e -> ix -> e -> m () +write' marr ix e = + write marr ix e >>= (`unless` errorIx "Data.Massiv.Array.Mutable.write'" (msize marr) ix) +{-# INLINE write' #-} + + +-- | /O(1)/ - Modify an element in the cell of a mutable array with a supplied function. Returns +-- `False` when index is out of bounds. +modify :: (Mutable r ix e, PrimMonad m) => + MArray (PrimState m) r ix e -> (e -> e) -> ix -> m Bool +modify marr f ix = + if isSafeIndex (msize marr) ix + then do + val <- unsafeRead marr ix + unsafeWrite marr ix $ f val + return True + else return False +{-# INLINE modify #-} + + +-- | /O(1)/ - Same as `modify`, but throws an error if index is out of bounds. +modify' :: (Mutable r ix e, PrimMonad m) => + MArray (PrimState m) r ix e -> (e -> e) -> ix -> m () +modify' marr f ix = + modify marr f ix >>= (`unless` errorIx "Data.Massiv.Array.Mutable.modify'" (msize marr) ix) +{-# INLINE modify' #-} + + +-- | /O(1)/ - Swap two elements in a mutable array by supplying their indices. Returns `False` when +-- either one of the indices is out of bounds. +swap :: (Mutable r ix e, PrimMonad m) => + MArray (PrimState m) r ix e -> ix -> ix -> m Bool +swap marr ix1 ix2 = do + let sz = msize marr + if isSafeIndex sz ix1 && isSafeIndex sz ix2 + then do + val1 <- unsafeRead marr ix1 + val2 <- unsafeRead marr ix2 + unsafeWrite marr ix1 val2 + unsafeWrite marr ix2 val1 + return True + else return False +{-# INLINE swap #-} + + +-- | /O(1)/ - Same as `swap`, but throws an error if index is out of bounds. +swap' :: (Mutable r ix e, PrimMonad m) => + MArray (PrimState m) r ix e -> ix -> ix -> m () +swap' marr ix1 ix2 = do + success <- swap marr ix1 ix2 + unless success $ + errorIx "Data.Massiv.Array.Mutable.swap'" (msize marr) $ + if isSafeIndex (msize marr) ix1 + then ix2 + else ix1 +{-# INLINE swap' #-} + diff --git a/src/Data/Massiv/Array/Numeric.hs b/src/Data/Massiv/Array/Numeric.hs new file mode 100644 --- /dev/null +++ b/src/Data/Massiv/Array/Numeric.hs @@ -0,0 +1,382 @@ +{-# LANGUAGE FlexibleContexts #-} +{-# LANGUAGE MultiParamTypeClasses #-} +{-# LANGUAGE ScopedTypeVariables #-} +{-# LANGUAGE TypeFamilies #-} +-- | +-- Module : Data.Massiv.Array.Numeric +-- Copyright : (c) Alexey Kuleshevich 2018 +-- License : BSD3 +-- Maintainer : Alexey Kuleshevich +-- Stability : experimental +-- Portability : non-portable +-- +module Data.Massiv.Array.Numeric + ( -- * Num + (.+) + , (.-) + , (.*) + , (.^) + , (|*|) + , negateA + , absA + , signumA + , fromIntegerA + -- * Integral + , quotA + , remA + , divA + , modA + , quotRemA + , divModA + -- * Fractional + , (./) + , (.^^) + , recipA + , fromRationalA + -- * Floating + , piA + , expA + , logA + , sqrtA + , (.**) + , logBaseA + , sinA + , cosA + , tanA + , asinA + , acosA + , atanA + , sinhA + , coshA + , tanhA + , asinhA + , acoshA + , atanhA + -- * RealFrac + , truncateA + , roundA + , ceilingA + , floorA + -- * RealFloat + , atan2A + ) where + +import Data.Massiv.Array.Delayed.Internal +import Data.Massiv.Array.Manifest.Internal (compute) +import Data.Massiv.Array.Ops.Fold as A +import Data.Massiv.Array.Ops.Map as A +import Data.Massiv.Array.Ops.Slice as A +import Data.Massiv.Array.Ops.Transform as A +import Data.Massiv.Core +import Data.Massiv.Core.Common +import Data.Monoid ((<>)) +import Prelude as P + + +infixr 8 .^, .^^ +infixl 7 .*, ./, `quotA`, `remA`, `divA`, `modA` +infixl 6 .+, .- + +(.+) + :: (Source r1 ix e, Source r2 ix e, Num e) + => Array r1 ix e -> Array r2 ix e -> Array D ix e +(.+) = liftArray2 (+) +{-# INLINE (.+) #-} + +(.-) + :: (Source r1 ix e, Source r2 ix e, Num e) + => Array r1 ix e -> Array r2 ix e -> Array D ix e +(.-) = liftArray2 (-) +{-# INLINE (.-) #-} + +(.*) + :: (Source r1 ix e, Source r2 ix e, Num e) + => Array r1 ix e -> Array r2 ix e -> Array D ix e +(.*) = liftArray2 (*) +{-# INLINE (.*) #-} + +(.^) + :: (Source r ix e, Num e, Integral b) + => Array r ix e -> b -> Array D ix e +(.^) arr n = liftArray (^ n) arr +{-# INLINE (.^) #-} + +-- | Perform matrix multiplication. Inner dimensions must agree, otherwise error. +(|*|) :: + ( Mutable r1 Ix2 e + , Mutable r2 Ix2 e + , OuterSlice r1 Ix2 e + , OuterSlice r2 Ix2 e + , Source (EltRepr r1 Ix2) Ix1 e + , Source (EltRepr r2 Ix2) Ix1 e + , Num e + ) + => Array r1 Ix2 e + -> Array r2 Ix2 e + -> Array D Ix2 e +(|*|) = multArrs +{-# INLINE (|*|) #-} + + +multArrs :: forall r1 r2 e. + ( Mutable r1 Ix2 e + , Mutable r2 Ix2 e + , OuterSlice r1 Ix2 e + , OuterSlice r2 Ix2 e + , Source (EltRepr r1 Ix2) Ix1 e + , Source (EltRepr r2 Ix2) Ix1 e + , Num e + ) + => Array r1 Ix2 e -> Array r2 Ix2 e -> Array D Ix2 e +multArrs arr1 arr2 + | n1 /= m2 = + error $ + "(|*|): Inner array dimensions must agree, but received: " ++ + show (size arr1) ++ " and " ++ show (size arr2) + | otherwise = + DArray (getComp arr1 <> getComp arr2) (m1 :. n2) $ \(i :. j) -> + A.sum ((arr1' !> i) .* (arr2' !> j)) + where + (m1 :. n1) = size arr1 + (m2 :. n2) = size arr2 + arr1' = setComp Seq arr1 + arr2' :: Array r2 Ix2 e + arr2' = setComp Seq $ compute $ transpose arr2 +{-# INLINE multArrs #-} + + +negateA + :: (Source r ix e, Num e) + => Array r ix e -> Array D ix e +negateA = liftArray negate +{-# INLINE negateA #-} + +absA + :: (Source r ix e, Num e) + => Array r ix e -> Array D ix e +absA = liftArray abs +{-# INLINE absA #-} + +signumA + :: (Source r ix e, Num e) + => Array r ix e -> Array D ix e +signumA = liftArray signum +{-# INLINE signumA #-} + +fromIntegerA + :: (Index ix, Num e) + => Integer -> Array D ix e +fromIntegerA = singleton Seq . fromInteger +{-# INLINE fromIntegerA #-} + +(./) + :: (Source r1 ix e, Source r2 ix e, Fractional e) + => Array r1 ix e -> Array r2 ix e -> Array D ix e +(./) = liftArray2 (/) +{-# INLINE (./) #-} + +(.^^) + :: (Source r ix e, Fractional e, Integral b) + => Array r ix e -> b -> Array D ix e +(.^^) arr n = liftArray (^^ n) arr +{-# INLINE (.^^) #-} + +recipA + :: (Source r ix e, Fractional e) + => Array r ix e -> Array D ix e +recipA = liftArray recip +{-# INLINE recipA #-} + + +fromRationalA + :: (Index ix, Fractional e) + => Rational -> Array D ix e +fromRationalA = singleton Seq . fromRational +{-# INLINE fromRationalA #-} + +piA + :: (Index ix, Floating e) + => Array D ix e +piA = singleton Seq pi +{-# INLINE piA #-} + +expA + :: (Source r ix e, Floating e) + => Array r ix e -> Array D ix e +expA = liftArray exp +{-# INLINE expA #-} + +sqrtA + :: (Source r ix e, Floating e) + => Array r ix e -> Array D ix e +sqrtA = liftArray exp +{-# INLINE sqrtA #-} + +logA + :: (Source r ix e, Floating e) + => Array r ix e -> Array D ix e +logA = liftArray log +{-# INLINE logA #-} + +logBaseA + :: (Source r1 ix e, Source r2 ix e, Floating e) + => Array r1 ix e -> Array r2 ix e -> Array D ix e +logBaseA = liftArray2 logBase +{-# INLINE logBaseA #-} + +(.**) + :: (Source r1 ix e, Source r2 ix e, Floating e) + => Array r1 ix e -> Array r2 ix e -> Array D ix e +(.**) = liftArray2 (**) +{-# INLINE (.**) #-} + + + +sinA + :: (Source r ix e, Floating e) + => Array r ix e -> Array D ix e +sinA = liftArray sin +{-# INLINE sinA #-} + +cosA + :: (Source r ix e, Floating e) + => Array r ix e -> Array D ix e +cosA = liftArray cos +{-# INLINE cosA #-} + +tanA + :: (Source r ix e, Floating e) + => Array r ix e -> Array D ix e +tanA = liftArray cos +{-# INLINE tanA #-} + +asinA + :: (Source r ix e, Floating e) + => Array r ix e -> Array D ix e +asinA = liftArray asin +{-# INLINE asinA #-} + +atanA + :: (Source r ix e, Floating e) + => Array r ix e -> Array D ix e +atanA = liftArray atan +{-# INLINE atanA #-} + +acosA + :: (Source r ix e, Floating e) + => Array r ix e -> Array D ix e +acosA = liftArray acos +{-# INLINE acosA #-} + +sinhA + :: (Source r ix e, Floating e) + => Array r ix e -> Array D ix e +sinhA = liftArray sinh +{-# INLINE sinhA #-} + +tanhA + :: (Source r ix e, Floating e) + => Array r ix e -> Array D ix e +tanhA = liftArray cos +{-# INLINE tanhA #-} + +coshA + :: (Source r ix e, Floating e) + => Array r ix e -> Array D ix e +coshA = liftArray cosh +{-# INLINE coshA #-} + +asinhA + :: (Source r ix e, Floating e) + => Array r ix e -> Array D ix e +asinhA = liftArray asinh +{-# INLINE asinhA #-} + +acoshA + :: (Source r ix e, Floating e) + => Array r ix e -> Array D ix e +acoshA = liftArray acosh +{-# INLINE acoshA #-} + +atanhA + :: (Source r ix e, Floating e) + => Array r ix e -> Array D ix e +atanhA = liftArray atanh +{-# INLINE atanhA #-} + + +quotA + :: (Source r1 ix e, Source r2 ix e, Integral e) + => Array r1 ix e -> Array r2 ix e -> Array D ix e +quotA = liftArray2 (quot) +{-# INLINE quotA #-} + + +remA + :: (Source r1 ix e, Source r2 ix e, Integral e) + => Array r1 ix e -> Array r2 ix e -> Array D ix e +remA = liftArray2 (rem) +{-# INLINE remA #-} + +divA + :: (Source r1 ix e, Source r2 ix e, Integral e) + => Array r1 ix e -> Array r2 ix e -> Array D ix e +divA = liftArray2 (div) +{-# INLINE divA #-} + +modA + :: (Source r1 ix e, Source r2 ix e, Integral e) + => Array r1 ix e -> Array r2 ix e -> Array D ix e +modA = liftArray2 (mod) +{-# INLINE modA #-} + + + +quotRemA + :: (Source r1 ix e, Source r2 ix e, Integral e) + => Array r1 ix e -> Array r2 ix e -> (Array D ix e, Array D ix e) +quotRemA arr1 = A.unzip . liftArray2 (quotRem) arr1 +{-# INLINE quotRemA #-} + + +divModA + :: (Source r1 ix e, Source r2 ix e, Integral e) + => Array r1 ix e -> Array r2 ix e -> (Array D ix e, Array D ix e) +divModA arr1 = A.unzip . liftArray2 (divMod) arr1 +{-# INLINE divModA #-} + + + +truncateA + :: (Source r ix a, RealFrac a, Integral b) + => Array r ix a -> Array D ix b +truncateA = liftArray truncate +{-# INLINE truncateA #-} + + +roundA + :: (Source r ix a, RealFrac a, Integral b) + => Array r ix a -> Array D ix b +roundA = liftArray round +{-# INLINE roundA #-} + + +ceilingA + :: (Source r ix a, RealFrac a, Integral b) + => Array r ix a -> Array D ix b +ceilingA = liftArray ceiling +{-# INLINE ceilingA #-} + + +floorA + :: (Source r ix a, RealFrac a, Integral b) + => Array r ix a -> Array D ix b +floorA = liftArray floor +{-# INLINE floorA #-} + +atan2A + :: (Source r ix e, RealFloat e) + => Array r ix e -> Array r ix e -> Array D ix e +atan2A = liftArray2 atan2 +{-# INLINE atan2A #-} + diff --git a/src/Data/Massiv/Array/Ops/Construct.hs b/src/Data/Massiv/Array/Ops/Construct.hs new file mode 100644 --- /dev/null +++ b/src/Data/Massiv/Array/Ops/Construct.hs @@ -0,0 +1,122 @@ +{-# LANGUAGE BangPatterns #-} +{-# LANGUAGE FlexibleContexts #-} +{-# LANGUAGE MultiParamTypeClasses #-} +{-# LANGUAGE TypeFamilies #-} +-- | +-- Module : Data.Massiv.Array.Ops.Construct +-- Copyright : (c) Alexey Kuleshevich 2018 +-- License : BSD3 +-- Maintainer : Alexey Kuleshevich +-- Stability : experimental +-- Portability : non-portable +-- +module Data.Massiv.Array.Ops.Construct + ( makeArray + , makeArrayR + , makeVectorR + , singleton + , range + , rangeStep + , enumFromN + , enumFromStepN + ) where + +import Data.Massiv.Array.Delayed.Internal +import Data.Massiv.Core.Common +import Prelude as P + + +-- | Just like `makeArray` but with ability to specify the result representation as an +-- argument. Note the `Data.Massiv.Array.U`nboxed type constructor in the below example. +-- +-- >>> makeArrayR U Par (2 :> 3 :. 4) (\ (i :> j :. k) -> i * i + j * j == k * k) +-- (Array U Par (2 :> 3 :. 4) +-- [ [ [ True,False,False,False ] +-- , [ False,True,False,False ] +-- , [ False,False,True,False ] +-- ] +-- , [ [ False,True,False,False ] +-- , [ False,False,False,False ] +-- , [ False,False,False,False ] +-- ] +-- ]) +-- +makeArrayR :: Construct r ix e => r -> Comp -> ix -> (ix -> e) -> Array r ix e +makeArrayR _ = makeArray +{-# INLINE makeArrayR #-} + + +-- | Same as `makeArrayR`, but restricted to 1-dimensional arrays. +makeVectorR :: Construct r Ix1 e => r -> Comp -> Ix1-> (Ix1 -> e) -> Array r Ix1 e +makeVectorR _ = makeArray +{-# INLINE makeVectorR #-} + + +-- | Create a vector with a range of @Int@s incremented by 1. +-- @range k0 k1 == rangeStep k0 k1 1@ +-- +-- >>> range Seq 1 6 +-- (Array D Seq (5) +-- [ 1,2,3,4,5 ]) +-- >>> range Seq (-2) 3 +-- (Array D Seq (5) +-- [ -2,-1,0,1,2 ]) +-- +range :: Comp -> Int -> Int -> Array D Ix1 Int +range comp !from !to = makeArray comp (max 0 (to - from)) (+ from) +{-# INLINE range #-} + + +-- | Same as `range`, but with a custom step. +-- +-- >>> rangeStep Seq 1 2 6 +-- (Array D Seq (3) +-- [ 1,3,5 ]) +-- +rangeStep :: Comp -- ^ Computation strategy + -> Int -- ^ Start + -> Int -- ^ Step (Can't be zero) + -> Int -- ^ End + -> Array D Ix1 Int +rangeStep comp !from !step !to + | step == 0 = error "rangeStep: Step can't be zero" + | otherwise = + let (sz, r) = (to - from) `divMod` step + in makeArray comp (sz + signum r) (\i -> from + i * step) +{-# INLINE rangeStep #-} + + +-- | Same as `enumFromStepN` with step @delta = 1@. +-- +-- >>> enumFromN Seq (5 :: Double) 3 +-- (Array D Seq (3) +-- [ 5.0,6.0,7.0 ]) +-- +enumFromN :: Num e => + Comp + -> e -- ^ @x@ - start value + -> Int -- ^ @n@ - length of resulting vector. + -> Array D Ix1 e +enumFromN comp !from !sz = makeArray comp sz $ \ i -> fromIntegral i + from +{-# INLINE enumFromN #-} + + +-- | Create a vector with length @n@ that has it's 0th value set to @x@ and gradually increasing +-- with @step@ delta until the end. Similar to: @`Data.Massiv.Array.fromList'` `Seq` $ `take` n [x, +-- x + delta ..]@. Major difference is that `fromList` constructs an `Array` with manifest +-- representation, while `enumFromStepN` is delayed. +-- +-- >>> enumFromStepN Seq 1 (0.1 :: Double) 5 +-- (Array D Seq (5) +-- [ 1.0,1.1,1.2,1.3,1.4 ]) +-- +enumFromStepN :: Num e => + Comp + -> e -- ^ @x@ - start value + -> e -- ^ @delta@ - step value + -> Int -- ^ @n@ - length of resulting vector + -> Array D Ix1 e +enumFromStepN comp !from !step !sz = makeArray comp sz $ \ i -> from + fromIntegral i * step +{-# INLINE enumFromStepN #-} + + diff --git a/src/Data/Massiv/Array/Ops/Fold.hs b/src/Data/Massiv/Array/Ops/Fold.hs new file mode 100644 --- /dev/null +++ b/src/Data/Massiv/Array/Ops/Fold.hs @@ -0,0 +1,463 @@ +{-# LANGUAGE BangPatterns #-} +{-# LANGUAGE FlexibleContexts #-} +{-# LANGUAGE MultiParamTypeClasses #-} +{-# LANGUAGE ScopedTypeVariables #-} +{-# LANGUAGE UndecidableInstances #-} +-- | +-- Module : Data.Massiv.Array.Ops.Fold +-- Copyright : (c) Alexey Kuleshevich 2018 +-- License : BSD3 +-- Maintainer : Alexey Kuleshevich +-- Stability : experimental +-- Portability : non-portable +-- +module Data.Massiv.Array.Ops.Fold + ( + -- ** Unstructured folds + + -- $unstruct_folds + + fold + , minimum + , maximum + , sum + , product + , and + , or + , all + , any + -- ** Sequential folds + + -- $seq_folds + + , foldlS + , foldrS + , ifoldlS + , ifoldrS + -- *** Monadic + , foldlM + , foldrM + , foldlM_ + , foldrM_ + , ifoldlM + , ifoldrM + , ifoldlM_ + , ifoldrM_ + -- *** Special folds + , foldrFB + , lazyFoldlS + , lazyFoldrS + -- ** Parallel folds + + -- $par_folds + + , foldlP + , foldrP + , ifoldlP + , ifoldrP + , foldlOnP + , ifoldlIO + , foldrOnP + , ifoldlOnP + , ifoldrOnP + , ifoldrIO + ) where + +import Control.Monad (void, when) +import qualified Data.Foldable as F +import Data.Functor.Identity (runIdentity) +import Data.Massiv.Core +import Data.Massiv.Core.Common +import Data.Massiv.Core.Scheduler +import Prelude hiding (all, and, any, foldl, foldr, + maximum, minimum, or, product, sum) +import System.IO.Unsafe (unsafePerformIO) + + +-- | /O(n)/ - Monadic left fold. +foldlM :: (Source r ix e, Monad m) => (a -> e -> m a) -> a -> Array r ix e -> m a +foldlM f = ifoldlM (\ a _ b -> f a b) +{-# INLINE foldlM #-} + + +-- | /O(n)/ - Monadic left fold, that discards the result. +foldlM_ :: (Source r ix e, Monad m) => (a -> e -> m a) -> a -> Array r ix e -> m () +foldlM_ f = ifoldlM_ (\ a _ b -> f a b) +{-# INLINE foldlM_ #-} + + +-- | /O(n)/ - Monadic left fold with an index aware function. +ifoldlM :: (Source r ix e, Monad m) => (a -> ix -> e -> m a) -> a -> Array r ix e -> m a +ifoldlM f !acc !arr = + iterM zeroIndex (size arr) 1 (<) acc $ \ !ix !a -> f a ix (unsafeIndex arr ix) +{-# INLINE ifoldlM #-} + + +-- | /O(n)/ - Monadic left fold with an index aware function, that discards the result. +ifoldlM_ :: (Source r ix e, Monad m) => (a -> ix -> e -> m a) -> a -> Array r ix e -> m () +ifoldlM_ f acc = void . ifoldlM f acc +{-# INLINE ifoldlM_ #-} + + +-- | /O(n)/ - Monadic right fold. +foldrM :: (Source r ix e, Monad m) => (e -> a -> m a) -> a -> Array r ix e -> m a +foldrM f = ifoldrM (\_ e a -> f e a) +{-# INLINE foldrM #-} + + +-- | /O(n)/ - Monadic right fold, that discards the result. +foldrM_ :: (Source r ix e, Monad m) => (e -> a -> m a) -> a -> Array r ix e -> m () +foldrM_ f = ifoldrM_ (\_ e a -> f e a) +{-# INLINE foldrM_ #-} + + +-- | /O(n)/ - Monadic right fold with an index aware function. +ifoldrM :: (Source r ix e, Monad m) => (ix -> e -> a -> m a) -> a -> Array r ix e -> m a +ifoldrM f !acc !arr = + iterM (liftIndex (subtract 1) (size arr)) zeroIndex (-1) (>=) acc $ \ !ix !acc0 -> + f ix (unsafeIndex arr ix) acc0 +{-# INLINE ifoldrM #-} + + +-- | /O(n)/ - Monadic right fold with an index aware function, that discards the result. +ifoldrM_ :: (Source r ix e, Monad m) => (ix -> e -> a -> m a) -> a -> Array r ix e -> m () +ifoldrM_ f !acc !arr = void $ ifoldrM f acc arr +{-# INLINE ifoldrM_ #-} + + + +-- | /O(n)/ - Left fold, computed sequentially with lazy accumulator. +lazyFoldlS :: Source r ix e => (a -> e -> a) -> a -> Array r ix e -> a +lazyFoldlS f initAcc arr = go initAcc 0 where + len = totalElem (size arr) + go acc k | k < len = go (f acc (unsafeLinearIndex arr k)) (k + 1) + | otherwise = acc +{-# INLINE lazyFoldlS #-} + + +-- | /O(n)/ - Right fold, computed sequentially with lazy accumulator. +lazyFoldrS :: Source r ix e => (e -> a -> a) -> a -> Array r ix e -> a +lazyFoldrS = foldrFB +{-# INLINE lazyFoldrS #-} + + +-- | /O(n)/ - Left fold, computed sequentially. +foldlS :: Source r ix e => (a -> e -> a) -> a -> Array r ix e -> a +foldlS f = ifoldlS (\ a _ e -> f a e) +{-# INLINE foldlS #-} + + +-- | /O(n)/ - Left fold with an index aware function, computed sequentially. +ifoldlS :: Source r ix e + => (a -> ix -> e -> a) -> a -> Array r ix e -> a +ifoldlS f acc = runIdentity . ifoldlM (\ a ix e -> return $ f a ix e) acc +{-# INLINE ifoldlS #-} + + +-- | /O(n)/ - Right fold, computed sequentially. +foldrS :: Source r ix e => (e -> a -> a) -> a -> Array r ix e -> a +foldrS f = ifoldrS (\_ e a -> f e a) +{-# INLINE foldrS #-} + + +-- | Version of foldr that supports @foldr/build@ list fusion implemented by GHC. +foldrFB :: Source r ix e => (e -> b -> b) -> b -> Array r ix e -> b +foldrFB c n arr = go 0 + where + !k = totalElem (size arr) + go !i + | i == k = n + | otherwise = let !v = unsafeLinearIndex arr i in v `c` go (i + 1) +{-# INLINE [0] foldrFB #-} + + + +-- | /O(n)/ - Right fold with an index aware function, computed sequentially. +ifoldrS :: Source r ix e => (ix -> e -> a -> a) -> a -> Array r ix e -> a +ifoldrS f acc = runIdentity . ifoldrM (\ ix e a -> return $ f ix e a) acc +{-# INLINE ifoldrS #-} + + + +-- | /O(n)/ - Left fold, computed in parallel. Parallelization of folding is implemented in such a +-- way that an array is split into a number of chunks of equal length, plus an extra one for the +-- left over. Number of chunks is the same as number of available cores (capabilities) plus one, and +-- each chunk is individually folded by a separate core with a function @g@. Results from folding +-- each chunk are further folded with another function @f@, thus allowing us to use information +-- about the structure of an array during folding. +-- +-- ===__Examples__ +-- +-- >>> foldlP (flip (:)) [] (flip (:)) [] $ makeArrayR U Seq (Ix1 11) id +-- [[10,9,8,7,6,5,4,3,2,1,0]] +-- +-- And this is how the result would look like if the above computation would be performed in a +-- program executed with @+RTS -N3@, i.e. with 3 capabilities: +-- +-- >>> foldlOnP [1,2,3] (flip (:)) [] (flip (:)) [] $ makeArrayR U Seq (Ix1 11) id +-- [[10,9],[8,7,6],[5,4,3],[2,1,0]] +-- +foldlP :: Source r ix e => + (a -> e -> a) -- ^ Folding function @g@. + -> a -- ^ Accumulator. Will be applied to @g@ multiple times, thus must be neutral. + -> (b -> a -> b) -- ^ Chunk results folding function @f@. + -> b -- ^ Accumulator for results of chunks folding. + -> Array r ix e -> IO b +foldlP f = ifoldlP (\ x _ -> f x) +{-# INLINE foldlP #-} + + +-- | Just like `foldlP`, but allows you to specify which cores (capabilities) to run computation +-- on. The order in which chunked results will be supplied to function @f@ is guaranteed to be +-- consecutive and aligned with the folding direction. +foldlOnP + :: Source r ix e + => [Int] -> (a -> e -> a) -> a -> (b -> a -> b) -> b -> Array r ix e -> IO b +foldlOnP wIds f = ifoldlOnP wIds (\ x _ -> f x) +{-# INLINE foldlOnP #-} + + + +-- | Parallel left fold. +ifoldlIO :: Source r ix e => + [Int] -- ^ List of capabilities + -> (a -> ix -> e -> IO a) -- ^ Index aware folding IO action + -> a -- ^ Accumulator + -> (b -> a -> IO b) -- ^ Folding action that is applied to results of parallel fold + -> b -- ^ Accumulator for chunks folding + -> Array r ix e -> IO b +ifoldlIO wIds f !initAcc g !tAcc !arr = do + let !sz = size arr + results <- + divideWork wIds sz $ \ !scheduler !chunkLength !totalLength !slackStart -> do + loopM_ 0 (< slackStart) (+ chunkLength) $ \ !start -> do + scheduleWork scheduler $ + iterLinearM sz start (start + chunkLength) 1 (<) initAcc $ \ !i ix !acc -> + f acc ix (unsafeLinearIndex arr i) + when (slackStart < totalLength) $ + scheduleWork scheduler $ + iterLinearM sz slackStart totalLength 1 (<) initAcc $ \ !i ix !acc -> + f acc ix (unsafeLinearIndex arr i) + F.foldlM g tAcc results +{-# INLINE ifoldlIO #-} + + +-- | Just like `ifoldlP`, but allows you to specify which cores to run +-- computation on. +ifoldlOnP :: Source r ix e => + [Int] -> (a -> ix -> e -> a) -> a -> (b -> a -> b) -> b -> Array r ix e -> IO b +ifoldlOnP wIds f initAcc g = + ifoldlIO wIds (\acc ix -> return . f acc ix) initAcc (\acc -> return . g acc) +{-# INLINE ifoldlOnP #-} + + + +-- | /O(n)/ - Left fold with an index aware function, computed in parallel. Just +-- like `foldlP`, except that folding function will receive an index of an +-- element it is being applied to. +ifoldlP :: Source r ix e => + (a -> ix -> e -> a) -> a -> (b -> a -> b) -> b -> Array r ix e -> IO b +ifoldlP = ifoldlOnP [] +{-# INLINE ifoldlP #-} + + +-- | /O(n)/ - Right fold, computed in parallel. Same as `foldlP`, except directed +-- from the last element in the array towards beginning. +-- +-- ==== __Examples__ +-- +-- >>> foldrP (++) [] (:) [] $ makeArray2D (3,4) id +-- [(0,0),(0,1),(0,2),(0,3),(1,0),(1,1),(1,2),(1,3),(2,0),(2,1),(2,2),(2,3)] +-- +foldrP :: Source r ix e => + (e -> a -> a) -> a -> (a -> b -> b) -> b -> Array r ix e -> IO b +foldrP f = ifoldrP (const f) +{-# INLINE foldrP #-} + + +-- | Just like `foldrP`, but allows you to specify which cores to run +-- computation on. +-- +-- ==== __Examples__ +-- +-- Number of wokers dictate the result structure: +-- +-- >>> foldrOnP [1,2,3] (:) [] (:) [] $ makeArray1D 9 id +-- [[0,1,2],[3,4,5],[6,7,8]] +-- >>> foldrOnP [1,2,3] (:) [] (:) [] $ makeArray1D 10 id +-- [[0,1,2],[3,4,5],[6,7,8],[9]] +-- >>> foldrOnP [1,2,3] (:) [] (:) [] $ makeArray1D 12 id +-- [[0,1,2,3],[4,5,6,7],[8,9,10,11]] +-- +-- But most of the time that structure is of no importance: +-- +-- >>> foldrOnP [1,2,3] (++) [] (:) [] $ makeArray1D 10 id +-- [0,1,2,3,4,5,6,7,8,9] +-- +-- Same as `foldlOnP`, order is guaranteed to be consecutive and in proper direction: +-- +-- >>> fmap snd $ foldrOnP [1,2,3] (\x (i, acc) -> (i + 1, (i, x):acc)) (1, []) (:) [] $ makeArray1D 11 id +-- [(4,[0,1,2]),(3,[3,4,5]),(2,[6,7,8]),(1,[9,10])] +-- >>> fmap (P.zip [4,3..]) <$> foldrOnP [1,2,3] (:) [] (:) [] $ makeArray1D 11 id +-- [(4,[0,1,2]),(3,[3,4,5]),(2,[6,7,8]),(1,[9,10])] +-- +foldrOnP :: Source r ix e => + [Int] -> (e -> a -> a) -> a -> (a -> b -> b) -> b -> Array r ix e -> IO b +foldrOnP wIds f = ifoldrOnP wIds (const f) +{-# INLINE foldrOnP #-} + + +-- | Parallel right fold. Differs from `ifoldrP` in that it accepts `IO` actions instead of the +-- usual pure functions as arguments. +ifoldrIO :: Source r ix e => + [Int] -> (ix -> e -> a -> IO a) -> a -> (a -> b -> IO b) -> b -> Array r ix e -> IO b +ifoldrIO wIds f !initAcc g !tAcc !arr = do + let !sz = size arr + results <- + divideWork wIds sz $ \ !scheduler !chunkLength !totalLength !slackStart -> do + when (slackStart < totalLength) $ + scheduleWork scheduler $ + iterLinearM sz (totalLength - 1) slackStart (-1) (>=) initAcc $ \ !i ix !acc -> + f ix (unsafeLinearIndex arr i) acc + loopM_ slackStart (> 0) (subtract chunkLength) $ \ !start -> + scheduleWork scheduler $ + iterLinearM sz (start - 1) (start - chunkLength) (-1) (>=) initAcc $ \ !i ix !acc -> + f ix (unsafeLinearIndex arr i) acc + F.foldlM (flip g) tAcc results +{-# INLINE ifoldrIO #-} + + +-- | /O(n)/ - Right fold with an index aware function, computed in parallel. +-- Same as `ifoldlP`, except directed from the last element in the array towards +-- beginning. +ifoldrOnP :: Source r ix e => + [Int] -> (ix -> e -> a -> a) -> a -> (a -> b -> b) -> b -> Array r ix e -> IO b +ifoldrOnP wIds f !initAcc g = + ifoldrIO wIds (\ix e -> return . f ix e) initAcc (\e -> return . g e) +{-# INLINE ifoldrOnP #-} + + +-- | Just like `ifoldrOnP`, but allows you to specify which cores to run computation on. +ifoldrP :: Source r ix e => + (ix -> e -> a -> a) -> a -> (a -> b -> b) -> b -> Array r ix e -> IO b +ifoldrP = ifoldrOnP [] +{-# INLINE ifoldrP #-} + + + +-- | /O(n)/ - Unstructured fold of an array. +fold :: Source r ix e => + (e -> e -> e) -- ^ Folding function (like with left fold, first argument + -- is an accumulator) + -> e -- ^ Initial element. Has to be neutral with respect to the folding + -- function. + -> Array r ix e -- ^ Source array + -> e +fold f initAcc = foldl f initAcc f initAcc +{-# INLINE fold #-} + + +-- | /O(n)/ - Compute maximum of all elements. +maximum :: (Source r ix e, Ord e) => + Array r ix e -> e +maximum = \arr -> + if isEmpty arr + then error "Data.Massiv.Array.maximum - empty" + else fold max (evaluateAt arr zeroIndex) arr +{-# INLINE maximum #-} + + +-- | /O(n)/ - Compute minimum of all elements. +minimum :: (Source r ix e, Ord e) => + Array r ix e -> e +minimum = \arr -> + if isEmpty arr + then error "Data.Massiv.Array.minimum - empty" + else fold max (evaluateAt arr zeroIndex) arr +{-# INLINE minimum #-} + + +-- | /O(n)/ - Compute sum of all elements. +sum :: (Source r ix e, Num e) => + Array r ix e -> e +sum = fold (+) 0 +{-# INLINE sum #-} + + +-- | /O(n)/ - Compute product of all elements. +product :: (Source r ix e, Num e) => + Array r ix e -> e +product = fold (*) 1 +{-# INLINE product #-} + + +-- | /O(n)/ - Compute conjunction of all elements. +and :: (Source r ix Bool) => + Array r ix Bool -> Bool +and = fold (&&) True +{-# INLINE and #-} + + +-- | /O(n)/ - Compute disjunction of all elements. +or :: Source r ix Bool => + Array r ix Bool -> Bool +or = fold (||) False +{-# INLINE or #-} + + +-- | Determines whether all element of the array satisfy the predicate. +all :: Source r ix e => + (e -> Bool) -> Array r ix e -> Bool +all f = foldl (\acc el -> acc && f el) True (&&) True +{-# INLINE all #-} + +-- | Determines whether any element of the array satisfies the predicate. +any :: Source r ix e => + (e -> Bool) -> Array r ix e -> Bool +any f = foldl (\acc el -> acc || f el) False (||) False +{-# INLINE any #-} + + +-- | This folding function breaks referencial transparency on some functions +-- @f@, therefore it is kept here for internal use only. +foldl :: Source r ix e => + (a -> e -> a) -> a -> (b -> a -> b) -> b -> Array r ix e -> b +foldl g initAcc f resAcc = \ arr -> + case getComp arr of + Seq -> f resAcc (foldlS g initAcc arr) + ParOn wIds -> unsafePerformIO $ foldlOnP wIds g initAcc f resAcc arr +{-# INLINE foldl #-} + + +{- $unstruct_folds + +Functions in this section will fold any `Source` array with respect to the inner +`Comp`utation strategy setting. + +-} + + +{- $seq_folds + +Functions in this section will fold any `Source` array sequentially, regardless of the inner +`Comp`utation strategy setting. + +-} + + +{- $par_folds + +__Note__ It is important to compile with @-threaded -with-rtsopts=-N@ flags, otherwise there will be +no parallelization. + +Functions in this section will fold any `Source` array in parallel, regardless of the inner +`Comp`utation strategy setting. All of the parallel structured folds are performed inside `IO` +monad, because referential transparency can't generally be preserved and results will depend on the +number of cores/capabilities that computation is being performed on. + +In contrast to sequential folds, each parallel folding function accepts two functions and two +initial elements as arguments. This is necessary because an array is first split into chunks, which +folded individually on separate cores with the first function, and the results of those folds are +further folded with the second function. + +-} diff --git a/src/Data/Massiv/Array/Ops/Map.hs b/src/Data/Massiv/Array/Ops/Map.hs new file mode 100644 --- /dev/null +++ b/src/Data/Massiv/Array/Ops/Map.hs @@ -0,0 +1,220 @@ +{-# LANGUAGE BangPatterns #-} +{-# LANGUAGE FlexibleContexts #-} +{-# LANGUAGE MultiParamTypeClasses #-} +-- | +-- Module : Data.Massiv.Array.Ops.Map +-- Copyright : (c) Alexey Kuleshevich 2018 +-- License : BSD3 +-- Maintainer : Alexey Kuleshevich +-- Stability : experimental +-- Portability : non-portable +-- +module Data.Massiv.Array.Ops.Map + ( map + , imap + -- ** Monadic + , mapM_ + , imapM_ + , forM_ + , iforM_ + , mapP_ + , imapP_ + -- ** Zipping + , zip + , zip3 + , unzip + , unzip3 + , zipWith + , zipWith3 + , izipWith + , izipWith3 + ) where + +import Control.Monad (void, when) +import Data.Massiv.Array.Delayed.Internal +import Data.Massiv.Core.Common +import Data.Massiv.Core.Scheduler +import Prelude hiding (map, mapM_, unzip, unzip3, + zip, zip3, zipWith, zipWith3) + + +-- | Map a function over an array +map :: Source r ix e' => (e' -> e) -> Array r ix e' -> Array D ix e +map f = imap (const f) +{-# INLINE map #-} + +-- | Map an index aware function over an array +imap :: Source r ix e' => (ix -> e' -> e) -> Array r ix e' -> Array D ix e +imap f !arr = DArray (getComp arr) (size arr) (\ !ix -> f ix (unsafeIndex arr ix)) +{-# INLINE imap #-} + +-- | Zip two arrays +zip :: (Source r1 ix e1, Source r2 ix e2) + => Array r1 ix e1 -> Array r2 ix e2 -> Array D ix (e1, e2) +zip = zipWith (,) +{-# INLINE zip #-} + +-- | Zip three arrays +zip3 :: (Source r1 ix e1, Source r2 ix e2, Source r3 ix e3) + => Array r1 ix e1 -> Array r2 ix e2 -> Array r3 ix e3 -> Array D ix (e1, e2, e3) +zip3 = zipWith3 (,,) +{-# INLINE zip3 #-} + +-- | Unzip two arrays +unzip :: Source r ix (e1, e2) => Array r ix (e1, e2) -> (Array D ix e1, Array D ix e2) +unzip arr = (map fst arr, map snd arr) +{-# INLINE unzip #-} + +-- | Unzip three arrays +unzip3 :: Source r ix (e1, e2, e3) + => Array r ix (e1, e2, e3) -> (Array D ix e1, Array D ix e2, Array D ix e3) +unzip3 arr = (map (\ (e, _, _) -> e) arr, map (\ (_, e, _) -> e) arr, map (\ (_, _, e) -> e) arr) +{-# INLINE unzip3 #-} + + + +-- | Zip two arrays with a function. Resulting array will be an intersection of +-- source arrays in case their dimensions do not match. +zipWith :: (Source r1 ix e1, Source r2 ix e2) + => (e1 -> e2 -> e) -> Array r1 ix e1 -> Array r2 ix e2 -> Array D ix e +zipWith f = izipWith (\ _ e1 e2 -> f e1 e2) +{-# INLINE zipWith #-} + + +-- | Just like `zipWith`, except with an index aware function. +izipWith :: (Source r1 ix e1, Source r2 ix e2) + => (ix -> e1 -> e2 -> e) -> Array r1 ix e1 -> Array r2 ix e2 -> Array D ix e +izipWith f arr1 arr2 = + DArray (getComp arr1) (liftIndex2 min (size arr1) (size arr1)) $ \ !ix -> + f ix (unsafeIndex arr1 ix) (unsafeIndex arr2 ix) +{-# INLINE izipWith #-} + + +-- | Just like `zipWith`, except zip three arrays with a function. +zipWith3 :: (Source r1 ix e1, Source r2 ix e2, Source r3 ix e3) + => (e1 -> e2 -> e3 -> e) -> Array r1 ix e1 -> Array r2 ix e2 -> Array r3 ix e3 -> Array D ix e +zipWith3 f = izipWith3 (\ _ e1 e2 e3 -> f e1 e2 e3) +{-# INLINE zipWith3 #-} + + +-- | Just like `zipWith3`, except with an index aware function. +izipWith3 + :: (Source r1 ix e1, Source r2 ix e2, Source r3 ix e3) + => (ix -> e1 -> e2 -> e3 -> e) + -> Array r1 ix e1 + -> Array r2 ix e2 + -> Array r3 ix e3 + -> Array D ix e +izipWith3 f arr1 arr2 arr3 = + DArray + (getComp arr1) + (liftIndex2 min (liftIndex2 min (size arr1) (size arr1)) (size arr3)) $ \ !ix -> + f ix (unsafeIndex arr1 ix) (unsafeIndex arr2 ix) (unsafeIndex arr3 ix) +{-# INLINE izipWith3 #-} + + + +-- | Map a monadic function over an array sequentially, while discarding the result. +-- +-- ==== __Examples__ +-- +-- >>> mapM_ print $ rangeStep 10 12 60 +-- 10 +-- 22 +-- 34 +-- 46 +-- 58 +-- +mapM_ :: (Source r ix a, Monad m) => (a -> m b) -> Array r ix a -> m () +mapM_ f !arr = iterM_ zeroIndex (size arr) 1 (<) (f . unsafeIndex arr) +{-# INLINE mapM_ #-} + + +-- | Just like `mapM_`, except with flipped arguments. +-- +-- ==== __Examples__ +-- +-- Here is a common way of iterating N times using a for loop in an imperative +-- language with mutation being an obvious side effect: +-- +-- >>> :m + Data.IORef +-- >>> var <- newIORef 0 :: IO (IORef Int) +-- >>> forM_ (range 0 1000) $ \ i -> modifyIORef' var (+i) +-- >>> readIORef var +-- 499500 +-- +forM_ :: (Source r ix a, Monad m) => Array r ix a -> (a -> m b) -> m () +forM_ = flip mapM_ +{-# INLINE forM_ #-} + + +-- | Map a monadic index aware function over an array sequentially, while discarding the result. +-- +-- ==== __Examples__ +-- +-- >>> imapM_ (curry print) $ range 10 15 +-- (0,10) +-- (1,11) +-- (2,12) +-- (3,13) +-- (4,14) +-- +imapM_ :: (Source r ix a, Monad m) => (ix -> a -> m b) -> Array r ix a -> m () +imapM_ f !arr = + iterM_ zeroIndex (size arr) 1 (<) $ \ !ix -> f ix (unsafeIndex arr ix) +{-# INLINE imapM_ #-} + +-- | Just like `imapM_`, except with flipped arguments. +iforM_ :: (Source r ix a, Monad m) => Array r ix a -> (ix -> a -> m b) -> m () +iforM_ = flip imapM_ +{-# INLINE iforM_ #-} + + + +-- | Map an IO action, over an array in parallel, while discarding the result. +mapP_ :: Source r ix a => (a -> IO b) -> Array r ix a -> IO () +mapP_ f = imapP_ (const f) +{-# INLINE mapP_ #-} + + +-- | Map an index aware IO action, over an array in parallel, while +-- discarding the result. +imapP_ :: Source r ix a => (ix -> a -> IO b) -> Array r ix a -> IO () +imapP_ f arr = do + let sz = size arr + wIds = + case getComp arr of + ParOn ids -> ids + _ -> [] + divideWork_ wIds sz $ \ !scheduler !chunkLength !totalLength !slackStart -> do + loopM_ 0 (< slackStart) (+ chunkLength) $ \ !start -> + scheduleWork scheduler $ + iterLinearM_ sz start (start + chunkLength) 1 (<) $ \ !i ix -> do + void $ f ix (unsafeLinearIndex arr i) + when (slackStart < totalLength) $ + scheduleWork scheduler $ + iterLinearM_ sz slackStart totalLength 1 (<) $ \ !i ix -> do + void $ f ix (unsafeLinearIndex arr i) +{-# INLINE imapP_ #-} + + + +-- -- | Map an IO action, that is index aware, over an array in parallel, while +-- -- discarding the result. +-- imapP_ :: (NFData b, Source r ix a) => (ix -> a -> IO b) -> Array r ix a -> IO () +-- imapP_ f !arr = do +-- let !sz = size arr +-- splitWork_ sz $ \ !scheduler !chunkLength !totalLength !slackStart -> do +-- loopM_ 0 (< slackStart) (+ chunkLength) $ \ !start -> +-- submitRequest scheduler $ +-- JobRequest 0 $ +-- iterLinearM_ sz start (start + chunkLength) 1 (<) $ \ !i ix -> do +-- res <- f ix (unsafeLinearIndex arr i) +-- res `deepseq` return () +-- when (slackStart < totalLength) $ +-- submitRequest scheduler $ +-- JobRequest 0 $ +-- iterLinearM_ sz slackStart totalLength 1 (<) $ \ !i ix -> do +-- res <- f ix (unsafeLinearIndex arr i) +-- res `deepseq` return () +-- {-# INLINE imapP_ #-} diff --git a/src/Data/Massiv/Array/Ops/Slice.hs b/src/Data/Massiv/Array/Ops/Slice.hs new file mode 100644 --- /dev/null +++ b/src/Data/Massiv/Array/Ops/Slice.hs @@ -0,0 +1,172 @@ +{-# LANGUAGE BangPatterns #-} +{-# LANGUAGE FlexibleContexts #-} +{-# LANGUAGE MultiParamTypeClasses #-} +{-# LANGUAGE TypeFamilies #-} +-- | +-- Module : Data.Massiv.Array.Ops.Slice +-- Copyright : (c) Alexey Kuleshevich 2018 +-- License : BSD3 +-- Maintainer : Alexey Kuleshevich +-- Stability : experimental +-- Portability : non-portable +-- +module Data.Massiv.Array.Ops.Slice + ( + -- ** From the outside + (!>) + , (!?>) + , (??>) + -- ** From the inside + , () + , () + , () + ) where + +import Control.Monad (guard) +import Data.Massiv.Core.Common + + +infixl 4 !>, !?>, ??>, , , + + +-- | /O(1)/ - Slices the array from the outside. For 2-dimensional array this will +-- be equivalent of taking a row. Throws an error when index is out of bounds. +-- +-- ===__Examples__ +-- +-- You could say that slicing from outside is synonymous to slicing from the end or slicing at the +-- highermost dimension. For example with rank-3 arrays outer slice would be equivalent to getting a +-- page: +-- +-- >>> let arr = makeArrayR U Seq (3 :> 2 :. 4) fromIx3 +-- >>> arr +-- (Array U Seq (3 :> 2 :. 4) +-- [ [ [ (0,0,0),(0,0,1),(0,0,2),(0,0,3) ] +-- , [ (0,1,0),(0,1,1),(0,1,2),(0,1,3) ] +-- ] +-- , [ [ (1,0,0),(1,0,1),(1,0,2),(1,0,3) ] +-- , [ (1,1,0),(1,1,1),(1,1,2),(1,1,3) ] +-- ] +-- , [ [ (2,0,0),(2,0,1),(2,0,2),(2,0,3) ] +-- , [ (2,1,0),(2,1,1),(2,1,2),(2,1,3) ] +-- ] +-- ]) +-- >>> arr !> 2 +-- (Array M Seq (2 :. 4) +-- [ [ (2,0,0),(2,0,1),(2,0,2),(2,0,3) ] +-- , [ (2,1,0),(2,1,1),(2,1,2),(2,1,3) ] +-- ]) +-- +-- There is nothing wrong with chaining, mixing and matching slicing operators, or even using them +-- to index arrays: +-- +-- >>> arr !> 2 !> 0 !> 3 +-- (2,0,3) +-- >>> arr !> 2 >> arr !> 2 !> 0 !> 3 == arr ! 2 :> 0 :. 3 +-- True +-- +(!>) :: OuterSlice r ix e => Array r ix e -> Int -> Elt r ix e +(!>) !arr !ix = + case arr !?> ix of + Just res -> res + Nothing -> errorIx "(!>)" (outerLength arr) ix +{-# INLINE (!>) #-} + + +-- | /O(1)/ - Just like `!>` slices the array from the outside, but returns +-- `Nothing` when index is out of bounds. +(!?>) :: OuterSlice r ix e => Array r ix e -> Int -> Maybe (Elt r ix e) +(!?>) !arr !i + | isSafeIndex (outerLength arr) i = Just $ unsafeOuterSlice arr i + | otherwise = Nothing +{-# INLINE (!?>) #-} + + +-- | /O(1)/ - Safe slicing continuation from the outside. Similarly to (`!>`) slices the array from +-- the outside, but takes `Maybe` array as input and returns `Nothing` when index is out of bounds. +-- +-- ===__Examples__ +-- +-- >>> let arr = makeArrayR U Seq (3 :> 2 :. 4) fromIx3 +-- >>> arr !?> 2 ??> 0 ??> 3 +-- Just (2,0,3) +-- >>> arr !?> 2 ??> 0 ??> -1 +-- Nothing +-- >>> arr !?> -2 ??> 0 ?? 1 +-- Nothing +-- +(??>) :: OuterSlice r ix e => Maybe (Array r ix e) -> Int -> Maybe (Elt r ix e) +(??>) Nothing _ = Nothing +(??>) (Just arr) !ix = arr !?> ix +{-# INLINE (??>) #-} + + +-- | /O(1)/ - Safe slice from the inside +( Array r ix e -> Int -> Maybe (Elt r ix e) +(`) slice an array from an opposite direction. +( Array r ix e -> Int -> Elt r ix e +( res + Nothing -> errorIx "( Maybe (Array r ix e) -> Int -> Maybe (Elt r ix e) +(`), but fails gracefully with a `Nothing`, instead of an error +() :: (Slice r ix e) + => Array r ix e -> (Dim, Int) -> Maybe (Elt r ix e) +() !arr !(dim, i) = do + m <- getIndex (size arr) dim + guard $ isSafeIndex m i + start <- setIndex zeroIndex dim i + cutSz <- setIndex (size arr) dim 1 + unsafeSlice arr start cutSz dim +{-# INLINE () #-} + + +-- | /O(1)/ - Slices the array in any available dimension. Throws an error when +-- index is out of bounds or dimensions is invalid. +-- +-- prop> arr !> i == arr (rank (size arr), i) +-- prop> arr (1,i) +-- +() :: Slice r ix e => Array r ix e -> (Dim, Int) -> Elt r ix e +() !arr !(dim, i) = + case arr (dim, i) of + Just res -> res + Nothing -> + let arrRank = rank (size arr) + in if dim < 1 || dim > arrRank + then error $ + "(): Invalid dimension: " ++ + show dim ++ " for Array of rank: " ++ show arrRank + else errorIx "()" (size arr) (dim, i) +{-# INLINE () #-} + + +-- | /O(1)/ - Safe slicing continuation from within. +() :: Slice r ix e => Maybe (Array r ix e) -> (Dim, Int) -> Maybe (Elt r ix e) +() Nothing _ = Nothing +() (Just arr) !ix = arr ix +{-# INLINE () #-} diff --git a/src/Data/Massiv/Array/Ops/Transform.hs b/src/Data/Massiv/Array/Ops/Transform.hs new file mode 100644 --- /dev/null +++ b/src/Data/Massiv/Array/Ops/Transform.hs @@ -0,0 +1,382 @@ +{-# LANGUAGE BangPatterns #-} +{-# LANGUAGE FlexibleContexts #-} +{-# LANGUAGE MultiParamTypeClasses #-} +{-# LANGUAGE TypeFamilies #-} +-- | +-- Module : Data.Massiv.Array.Ops.Transform +-- Copyright : (c) Alexey Kuleshevich 2018 +-- License : BSD3 +-- Maintainer : Alexey Kuleshevich +-- Stability : experimental +-- Portability : non-portable +-- +module Data.Massiv.Array.Ops.Transform + ( -- ** Transpose + transpose + , transposeInner + , transposeOuter + -- ** Backpermute + , backpermute + -- ** Resize + , resize + , resize' + -- ** Extract + , extract + , extract' + , extractFromTo + -- ** Append/Split + , append + , append' + , splitAt + , splitAt' + -- * Traverse + , traverse + , traverse2 + ) where + +import Control.Monad (guard) +import Data.Massiv.Array.Delayed.Internal +import Data.Massiv.Array.Ops.Construct +import Data.Massiv.Core.Common +import Data.Maybe (fromMaybe) +import Prelude hiding (splitAt, traverse) + + +-- | Extract a sub-array from within a larger source array. Array that is being extracted must be +-- fully encapsulated in a source array, otherwise `Nothing` is returned, +extract :: Size r ix e + => ix -- ^ Starting index + -> ix -- ^ Size fo the resulting array + -> Array r ix e -- ^ Source array + -> Maybe (Array (EltRepr r ix) ix e) +extract !sIx !newSz !arr + | isSafeIndex sz1 sIx && isSafeIndex eIx1 sIx && isSafeIndex sz1 eIx = + Just $ unsafeExtract sIx newSz arr + | otherwise = Nothing + where + sz1 = liftIndex (+1) (size arr) + eIx1 = liftIndex (+1) eIx + eIx = liftIndex2 (+) sIx newSz +{-# INLINE extract #-} + +-- | Same as `extract`, but will throw an error if supplied dimensions are incorrect. +extract' :: Size r ix e + => ix -- ^ Starting index + -> ix -- ^ Size fo the resulting array + -> Array r ix e -- ^ Source array + -> Array (EltRepr r ix) ix e +extract' !sIx !newSz !arr = + case extract sIx newSz arr of + Just arr' -> arr' + Nothing -> + error $ + "Data.Massiv.Array.extract': Cannot extract an array of size " ++ + show newSz ++ + " starting at " ++ show sIx ++ " from within an array of size: " ++ show (size arr) +{-# INLINE extract' #-} + +-- | Similar to `extract`, except it takes starting and ending index. Result array will not include +-- the ending index. +extractFromTo :: Size r ix e => + ix -- ^ Starting index + -> ix -- ^ Index up to which elmenets should be extracted. + -> Array r ix e -- ^ Source array. + -> Maybe (Array (EltRepr r ix) ix e) +extractFromTo sIx eIx = extract sIx $ liftIndex2 (-) eIx sIx +{-# INLINE extractFromTo #-} + +-- | /O(1)/ - Changes the shape of an array. Returns `Nothing` if total +-- number of elements does not match the source array. +resize :: (Index ix', Size r ix e) => ix' -> Array r ix e -> Maybe (Array r ix' e) +resize !sz !arr + | totalElem sz == totalElem (size arr) = Just $ unsafeResize sz arr + | otherwise = Nothing +{-# INLINE resize #-} + +-- | Same as `resize`, but will throw an error if supplied dimensions are incorrect. +resize' :: (Index ix', Size r ix e) => ix' -> Array r ix e -> Array r ix' e +resize' !sz !arr = + maybe + (error $ + "Total number of elements do not match: " ++ + show sz ++ " vs " ++ show (size arr)) + id $ + resize sz arr +{-# INLINE resize' #-} + + +-- | Transpose a 2-dimensional array +-- +-- ===__Examples__ +-- +-- >>> let arr = makeArrayR U Seq (2 :. 3) (toLinearIndex (2 :. 3)) +-- >>> arr +-- (ArrayU Seq (2 :. 3) +-- [ [ 0,1,2 ] +-- , [ 3,4,5 ] +-- ]) +-- >>> transpose arr +-- (Array D Seq (3 :. 2) +-- [ [ 0,3 ] +-- , [ 1,4 ] +-- , [ 2,5 ] +-- ]) +-- +transpose :: Source r Ix2 e => Array r Ix2 e -> Array D Ix2 e +transpose = transposeInner +{-# INLINE transpose #-} + + +-- | Transpose inner two dimensions of at least rank-2 array. +-- +-- ===__Examples__ +-- +-- >>> let arr = makeArrayR U Seq (2 :> 3 :. 4) fromIx3 +-- >>> arr +-- (Array U Seq (2 :> 3 :. 4) +-- [ [ [ (0,0,0),(0,0,1),(0,0,2),(0,0,3) ] +-- , [ (0,1,0),(0,1,1),(0,1,2),(0,1,3) ] +-- , [ (0,2,0),(0,2,1),(0,2,2),(0,2,3) ] +-- ] +-- , [ [ (1,0,0),(1,0,1),(1,0,2),(1,0,3) ] +-- , [ (1,1,0),(1,1,1),(1,1,2),(1,1,3) ] +-- , [ (1,2,0),(1,2,1),(1,2,2),(1,2,3) ] +-- ] +-- ]) +-- >>> transposeInner arr +-- (Array D Seq (3 :> 2 :. 4) +-- [ [ [ (0,0,0),(0,0,1),(0,0,2),(0,0,3) ] +-- , [ (1,0,0),(1,0,1),(1,0,2),(1,0,3) ] +-- ] +-- , [ [ (0,1,0),(0,1,1),(0,1,2),(0,1,3) ] +-- , [ (1,1,0),(1,1,1),(1,1,2),(1,1,3) ] +-- ] +-- , [ [ (0,2,0),(0,2,1),(0,2,2),(0,2,3) ] +-- , [ (1,2,0),(1,2,1),(1,2,2),(1,2,3) ] +-- ] +-- ]) +-- +transposeInner :: (Index (Lower ix), Source r' ix e) + => Array r' ix e -> Array D ix e +transposeInner !arr = unsafeMakeArray (getComp arr) (transInner (size arr)) newVal + where + transInner !ix = + fromMaybe (errorImpossible "transposeInner" ix) $ do + n <- getIndex ix (rank ix) + m <- getIndex ix (rank ix - 1) + ix' <- setIndex ix (rank ix) m + setIndex ix' (rank ix - 1) n + {-# INLINE transInner #-} + newVal = unsafeIndex arr . transInner + {-# INLINE newVal #-} +{-# INLINE transposeInner #-} + +-- | Transpose outer two dimensions of at least rank-2 array. +-- +-- ===__Examples__ +-- +-- >>> let arr = makeArrayR U Seq (2 :> 3 :. 4) fromIx3 +-- >>> arr +-- (Array U Seq (2 :> 3 :. 4) +-- [ [ [ (0,0,0),(0,0,1),(0,0,2),(0,0,3) ] +-- , [ (0,1,0),(0,1,1),(0,1,2),(0,1,3) ] +-- , [ (0,2,0),(0,2,1),(0,2,2),(0,2,3) ] +-- ] +-- , [ [ (1,0,0),(1,0,1),(1,0,2),(1,0,3) ] +-- , [ (1,1,0),(1,1,1),(1,1,2),(1,1,3) ] +-- , [ (1,2,0),(1,2,1),(1,2,2),(1,2,3) ] +-- ] +-- ]) +-- >>> transposeOuter arr +-- (Array D Seq (2 :> 4 :. 3) +-- [ [ [ (0,0,0),(0,1,0),(0,2,0) ] +-- , [ (0,0,1),(0,1,1),(0,2,1) ] +-- , [ (0,0,2),(0,1,2),(0,2,2) ] +-- , [ (0,0,3),(0,1,3),(0,2,3) ] +-- ] +-- , [ [ (1,0,0),(1,1,0),(1,2,0) ] +-- , [ (1,0,1),(1,1,1),(1,2,1) ] +-- , [ (1,0,2),(1,1,2),(1,2,2) ] +-- , [ (1,0,3),(1,1,3),(1,2,3) ] +-- ] +-- ]) +-- +transposeOuter :: (Index (Lower ix), Source r' ix e) + => Array r' ix e -> Array D ix e +transposeOuter !arr = unsafeMakeArray (getComp arr) (transOuter (size arr)) newVal + where + transOuter !ix = + fromMaybe (errorImpossible "transposeOuter" ix) $ do + n <- getIndex ix 1 + m <- getIndex ix 2 + ix' <- setIndex ix 1 m + setIndex ix' 2 n + {-# INLINE transOuter #-} + newVal = unsafeIndex arr . transOuter + {-# INLINE newVal #-} +{-# INLINE transposeOuter #-} + + +-- | Rearrange elements of an array into a new one. +-- +-- ===__Examples__ +-- +-- >>> let arr = makeArrayR U Seq (2 :> 3 :. 4) fromIx3 +-- >>> arr +-- (Array U Seq (2 :> 3 :. 4) +-- [ [ [ (0,0,0),(0,0,1),(0,0,2),(0,0,3) ] +-- , [ (0,1,0),(0,1,1),(0,1,2),(0,1,3) ] +-- , [ (0,2,0),(0,2,1),(0,2,2),(0,2,3) ] +-- ] +-- , [ [ (1,0,0),(1,0,1),(1,0,2),(1,0,3) ] +-- , [ (1,1,0),(1,1,1),(1,1,2),(1,1,3) ] +-- , [ (1,2,0),(1,2,1),(1,2,2),(1,2,3) ] +-- ] +-- ]) +-- >>> backpermute (4 :. 3) (\(i :. j) -> 0 :> j :. i) arr +-- (Array D Seq (4 :. 3) +-- [ [ (0,0,0),(0,1,0),(0,2,0) ] +-- , [ (0,0,1),(0,1,1),(0,2,1) ] +-- , [ (0,0,2),(0,1,2),(0,2,2) ] +-- , [ (0,0,3),(0,1,3),(0,2,3) ] +-- ]) +-- +backpermute :: (Source r' ix' e, Index ix) => + ix -- ^ Size of the result array + -> (ix -> ix') -- ^ A function that maps indices of old array into the source one. + -> Array r' ix' e -- ^ Source array. + -> Array D ix e +backpermute sz ixF !arr = makeArray (getComp arr) sz (evaluateAt arr . ixF) +{-# INLINE backpermute #-} + + +-- | Append two arrays together along a particular dimension. Sizes of both arrays must match, with +-- an allowed exception of the dimension they are being appended along, otherwise `Nothing` is +-- returned. +-- +-- ===__Examples__ +-- +-- Append two 2D arrays along both dimensions. Note that they have the same shape. +-- +-- >>> let arrA = makeArrayR U Seq (2 :. 3) (\(i :. j) -> ('A', i, j)) +-- >>> let arrB = makeArrayR U Seq (2 :. 3) (\(i :. j) -> ('B', i, j)) +-- >>> append 1 arrA arrB +-- Just (Array D Seq (2 :. 6) +-- [ [ ('A',0,0),('A',0,1),('A',0,2),('B',0,0),('B',0,1),('B',0,2) ] +-- , [ ('A',1,0),('A',1,1),('A',1,2),('B',1,0),('B',1,1),('B',1,2) ] +-- ]) +-- >>> append 2 arrA arrB +-- Just (Array D Seq (4 :. 3) +-- [ [ ('A',0,0),('A',0,1),('A',0,2) ] +-- , [ ('A',1,0),('A',1,1),('A',1,2) ] +-- , [ ('B',0,0),('B',0,1),('B',0,2) ] +-- , [ ('B',1,0),('B',1,1),('B',1,2) ] +-- ]) +-- +-- Now appending arrays with different sizes: +-- +-- >>> let arrC = makeArrayR U Seq (2 :. 4) (\(i :. j) -> ('C', i, j)) +-- >>> append 1 arrA arrC +-- Just (Array D Seq (2 :. 7) +-- [ [ ('A',0,0),('A',0,1),('A',0,2),('C',0,0),('C',0,1),('C',0,2),('C',0,3) ] +-- , [ ('A',1,0),('A',1,1),('A',1,2),('C',1,0),('C',1,1),('C',1,2),('C',1,3) ] +-- ]) +-- >>> append 2 arrA arrC +-- Nothing +-- +append :: (Source r1 ix e, Source r2 ix e) => + Dim -> Array r1 ix e -> Array r2 ix e -> Maybe (Array D ix e) +append n !arr1 !arr2 = do + let sz1 = size arr1 + sz2 = size arr2 + k1 <- getIndex sz1 n + k2 <- getIndex sz2 n + sz1' <- setIndex sz2 n k1 + guard $ sz1 == sz1' + newSz <- setIndex sz1 n (k1 + k2) + return $ + unsafeMakeArray (getComp arr1) newSz $ \ !ix -> + fromMaybe (errorImpossible "append" ix) $ do + k' <- getIndex ix n + if k' < k1 + then Just (unsafeIndex arr1 ix) + else do + i <- getIndex ix n + ix' <- setIndex ix n (i - k1) + return $ unsafeIndex arr2 ix' +{-# INLINE append #-} + +-- | Same as `append`, but will throw an error instead of returning `Nothing` on mismatched sizes. +append' :: (Source r1 ix e, Source r2 ix e) => + Dim -> Array r1 ix e -> Array r2 ix e -> Array D ix e +append' dim arr1 arr2 = + case append dim arr1 arr2 of + Just arr -> arr + Nothing -> + error $ + if 0 < dim && dim <= rank (size arr1) + then "append': Dimension mismatch: " ++ show (size arr1) ++ " and " ++ show (size arr2) + else "append': Invalid dimension: " ++ show dim +{-# INLINE append' #-} + +-- | /O(1)/ - Split an array at an index along a specified dimension. +splitAt :: + (Size r ix e, r' ~ EltRepr r ix) + => Dim -- ^ Dimension along which to split + -> Int -- ^ Index along the dimension to split at + -> Array r ix e -- ^ Source array + -> Maybe (Array r' ix e, Array r' ix e) +splitAt dim i arr = do + let sz = size arr + eIx <- setIndex sz dim i + sIx <- setIndex zeroIndex dim i + arr1 <- extractFromTo zeroIndex eIx arr + arr2 <- extractFromTo sIx sz arr + return (arr1, arr2) +{-# INLINE splitAt #-} + +-- | Same as `splitAt`, but will throw an error instead of returning `Nothing` on wrong dimension +-- and index out of bounds. +splitAt' :: (Size r ix e, r' ~ EltRepr r ix) => + Dim -> Int -> Array r ix e -> (Array r' ix e, Array r' ix e) +splitAt' dim i arr = + case splitAt dim i arr of + Just res -> res + Nothing -> + error $ + "Data.Massiv.Array.splitAt': " ++ + if 0 < dim && dim <= rank (size arr) + then "Index out of bounds: " ++ + show i ++ " for dimension: " ++ show dim ++ " and array with size: " ++ show (size arr) + else "Invalid dimension: " ++ show dim ++ " for array with size: " ++ show (size arr) +{-# INLINE splitAt' #-} + +-- | Create an array by traversing a source array. +traverse + :: (Source r1 ix1 e1, Index ix) + => ix -- ^ Size of the result array + -> ((ix1 -> e1) -> ix -> e) -- ^ Function that will receive a source array safe index function and + -- an index for an element it should return a value of. + -> Array r1 ix1 e1 -- ^ Source array + -> Array D ix e +traverse sz f arr1 = makeArray (getComp arr1) sz (f (evaluateAt arr1)) +{-# INLINE traverse #-} + + +-- | Create an array by traversing two source arrays. +traverse2 + :: (Source r1 ix1 e1, Source r2 ix2 e2, Index ix) + => ix + -> ((ix1 -> e1) -> (ix2 -> e2) -> ix -> e) + -> Array r1 ix1 e1 + -> Array r2 ix2 e2 + -> Array D ix e +traverse2 sz f arr1 arr2 = makeArray (getComp arr1) sz (f (evaluateAt arr1) (evaluateAt arr2)) +{-# INLINE traverse2 #-} + + +-- | Throw an impossible error on a `Nothing` +errorImpossible :: Show c => String -> c -> a +errorImpossible fName cause = + error $ "Data.Massiv.Array." ++ fName ++ ": Impossible happened " ++ show cause +{-# NOINLINE errorImpossible #-} diff --git a/src/Data/Massiv/Array/Stencil.hs b/src/Data/Massiv/Array/Stencil.hs new file mode 100644 --- /dev/null +++ b/src/Data/Massiv/Array/Stencil.hs @@ -0,0 +1,88 @@ +{-# LANGUAGE BangPatterns #-} +{-# LANGUAGE FlexibleContexts #-} +{-# LANGUAGE MultiParamTypeClasses #-} +-- | +-- Module : Data.Massiv.Array.Stencil +-- Copyright : (c) Alexey Kuleshevich 2018 +-- License : BSD3 +-- Maintainer : Alexey Kuleshevich +-- Stability : experimental +-- Portability : non-portable +-- +module Data.Massiv.Array.Stencil + ( Stencil + , Value + , mapStencil + , makeStencil + , makeConvolutionStencil + , makeConvolutionStencilFromKernel + ) where + +import Data.Default.Class (Default (def)) +import Data.Massiv.Array.Delayed.Windowed +import Data.Massiv.Array.Manifest +import Data.Massiv.Array.Stencil.Convolution +import Data.Massiv.Array.Stencil.Internal +import Data.Massiv.Core.Common +import GHC.Exts (inline) + + +-- | Map a constructed stencil over an array. Resulting array must be `compute`d in order to be +-- useful. +mapStencil :: (Source r ix e, Manifest r ix e) => + Stencil ix e a -> Array r ix e -> Array DW ix a +mapStencil (Stencil b sSz sCenter stencilF) !arr = + DWArray + (DArray (getComp arr) sz (unValue . stencilF (Value . borderIndex b arr))) + (Just sSz) + sCenter + (liftIndex2 (-) sz (liftIndex2 (-) sSz (pureIndex 1))) + (unValue . stencilF (Value . unsafeIndex arr)) + where + !sz = size arr +{-# INLINE mapStencil #-} + + +-- | Construct a stencil from a function, which describes how to calculate the +-- value at a point while having access to neighboring elements with a function +-- that accepts idices relative to the center of stencil. Trying to index +-- outside the stencil box will result in a runtime error upon stencil +-- creation. +-- +-- ==== __Example__ +-- +-- Below is an example of creating a `Stencil`, which, when mapped over a +-- 2-dimensional array, will compute an average of all elements in a 3x3 square +-- for each element in that array. /Note:/ Make sure to add @INLINE@ pragma, +-- otherwise performance will be terrible. +-- +-- > average3x3Stencil :: (Default a, Fractional a) => Border a -> Stencil Ix2 a a +-- > average3x3Stencil b = makeStencil b (3 :. 3) (1 :. 1) $ \ get -> +-- > ( get (-1 :. -1) + get (-1 :. 0) + get (-1 :. 1) + +-- > get ( 0 :. -1) + get ( 0 :. 0) + get ( 0 :. 1) + +-- > get ( 1 :. -1) + get ( 1 :. 0) + get ( 1 :. 1) ) / 9 +-- > {-# INLINE average3x3Stencil #-} +-- +makeStencil + :: (Index ix, Default e) + => Border e -- ^ Border resolution technique + -> ix -- ^ Size of the stencil + -> ix -- ^ Center of the stencil + -> ((ix -> Value e) -> Value a) + -- ^ Stencil function that receives a "get" function as it's argument that can + -- retrieve values of cells in the source array with respect to the center of + -- the stencil. Stencil function must return a value that will be assigned to + -- the cell in the result array. Offset supplied to the "get" function + -- cannot go outside the boundaries of the stencil, otherwise an error will be + -- raised during stencil creation. + -> Stencil ix e a +makeStencil b !sSz !sCenter relStencil = + validateStencil def $ Stencil b sSz sCenter stencil + where + stencil getVal !ix = + (inline relStencil $ \ !ixD -> getVal (liftIndex2 (-) ix ixD)) + {-# INLINE stencil #-} +{-# INLINE makeStencil #-} + + + diff --git a/src/Data/Massiv/Array/Stencil/Convolution.hs b/src/Data/Massiv/Array/Stencil/Convolution.hs new file mode 100644 --- /dev/null +++ b/src/Data/Massiv/Array/Stencil/Convolution.hs @@ -0,0 +1,67 @@ +{-# LANGUAGE FlexibleContexts #-} +{-# LANGUAGE BangPatterns #-} +-- | +-- Module : Data.Massiv.Array.Stencil.Convolution +-- Copyright : (c) Alexey Kuleshevich 2018 +-- License : BSD3 +-- Maintainer : Alexey Kuleshevich +-- Stability : experimental +-- Portability : non-portable +-- +module Data.Massiv.Array.Stencil.Convolution where + +import Data.Massiv.Core.Common +import Data.Massiv.Array.Ops.Fold (ifoldlS) +import Data.Massiv.Array.Stencil.Internal +import GHC.Exts (inline) + +-- | Create a convolution stencil by specifying border resolution technique and +-- an accumulator function. +-- +-- ==== __Examples__ +-- +-- Here is how to create a 2D horizontal Sobel Stencil: +-- +-- > sobelX :: Num e => Border e -> Stencil Ix2 e e +-- > sobelX b = makeConvolutionStencil b (3 :. 3) (1 :. 1) $ +-- > \f -> f (-1 :. -1) 1 . f (-1 :. 1) (-1) . +-- > f ( 0 :. -1) 2 . f ( 0 :. 1) (-2) . +-- > f ( 1 :. -1) 1 . f ( 1 :. 1) (-1) +-- > {-# INLINE sobelX #-} +-- +makeConvolutionStencil + :: (Index ix, Num e) + => Border e + -> ix + -> ix + -> ((ix -> Value e -> Value e -> Value e) -> Value e -> Value e) + -> Stencil ix e e +makeConvolutionStencil b !sSz !sCenter relStencil = + validateStencil 0 $ Stencil b sSz sCenter stencil + where + stencil getVal !ix = + ((inline relStencil $ \ !ixD !kVal !acc -> + (getVal (liftIndex2 (-) ix ixD)) * kVal + acc) + 0) + {-# INLINE stencil #-} +{-# INLINE makeConvolutionStencil #-} + + +-- | Make a stencil out of a Kernel Array. This `Stencil` will be slower than if +-- `makeConvolutionStencil` is used, but sometimes we just really don't know the +-- kernel at compile time. +makeConvolutionStencilFromKernel + :: (Manifest r ix e, Num e) + => Border e + -> Array r ix e + -> Stencil ix e e +makeConvolutionStencilFromKernel b kArr = Stencil b sz sCenter stencil + where + !sz = size kArr + !sCenter = (liftIndex (`div` 2) sz) + stencil getVal !ix = Value (ifoldlS accum 0 kArr) where + accum !acc !kIx !kVal = + unValue (getVal (liftIndex2 (+) ix (liftIndex2 (-) sCenter kIx))) * kVal + acc + {-# INLINE accum #-} + {-# INLINE stencil #-} +{-# INLINE makeConvolutionStencilFromKernel #-} diff --git a/src/Data/Massiv/Array/Stencil/Internal.hs b/src/Data/Massiv/Array/Stencil/Internal.hs new file mode 100644 --- /dev/null +++ b/src/Data/Massiv/Array/Stencil/Internal.hs @@ -0,0 +1,224 @@ +{-# LANGUAGE BangPatterns #-} +{-# LANGUAGE FlexibleContexts #-} +{-# LANGUAGE FlexibleInstances #-} +{-# LANGUAGE MultiParamTypeClasses #-} +{-# LANGUAGE RecordWildCards #-} +{-# LANGUAGE ScopedTypeVariables #-} +-- | +-- Module : Data.Massiv.Array.Stencil.Internal +-- Copyright : (c) Alexey Kuleshevich 2018 +-- License : BSD3 +-- Maintainer : Alexey Kuleshevich +-- Stability : experimental +-- Portability : non-portable +-- +module Data.Massiv.Array.Stencil.Internal where + +import Control.Applicative +import Control.DeepSeq +import Data.Massiv.Core.Common +import Data.Massiv.Array.Delayed.Internal +import Data.Default.Class (Default (def)) + +-- | Stencil is abstract description of how to handle elements in the neighborhood of every array +-- cell in order to compute a value for the cells in the new array. Use `Data.Array.makeStencil` and +-- `Data.Array.makeConvolutionStencil` in order to create a stencil. +data Stencil ix e a = Stencil + { stencilBorder :: Border e + , stencilSize :: !ix + , stencilCenter :: !ix + , stencilFunc :: (ix -> Value e) -> ix -> Value a + } + +instance (NFData e, Index ix) => NFData (Stencil ix e a) where + rnf (Stencil b sz ix f) = b `deepseq` sz `deepseq` ix `deepseq` f `seq` () + +-- | This is a simple wrapper for value of an array cell. It is used in order to improve safety of +-- `Stencil` mapping. Using various class instances, such as `Num` and `Functor` for example, make +-- it possible to manipulate the value, without having direct access to it. +newtype Value e = Value { unValue :: e } deriving (Show, Eq, Ord, Bounded) + + +instance Functor Value where + fmap f (Value e) = Value (f e) + {-# INLINE fmap #-} + +instance Applicative Value where + pure = Value + {-# INLINE pure #-} + (<*>) (Value f) (Value e) = Value (f e) + {-# INLINE (<*>) #-} + +instance Num e => Num (Value e) where + (+) = liftA2 (+) + {-# INLINE (+) #-} + (*) = liftA2 (*) + {-# INLINE (*) #-} + negate = fmap negate + {-# INLINE negate #-} + abs = fmap abs + {-# INLINE abs #-} + signum = fmap signum + {-# INLINE signum #-} + fromInteger = Value . fromInteger + {-# INLINE fromInteger #-} + +instance Fractional e => Fractional (Value e) where + (/) = liftA2 (/) + {-# INLINE (/) #-} + recip = fmap recip + {-# INLINE recip #-} + fromRational = pure . fromRational + {-# INLINE fromRational #-} + +instance Floating e => Floating (Value e) where + pi = pure pi + {-# INLINE pi #-} + exp = fmap exp + {-# INLINE exp #-} + log = fmap log + {-# INLINE log #-} + sqrt = fmap sqrt + {-# INLINE sqrt #-} + (**) = liftA2 (**) + {-# INLINE (**) #-} + logBase = liftA2 logBase + {-# INLINE logBase #-} + sin = fmap sin + {-# INLINE sin #-} + cos = fmap cos + {-# INLINE cos #-} + tan = fmap tan + {-# INLINE tan #-} + asin = fmap asin + {-# INLINE asin #-} + acos = fmap acos + {-# INLINE acos #-} + atan = fmap atan + {-# INLINE atan #-} + sinh = fmap sinh + {-# INLINE sinh #-} + cosh = fmap cosh + {-# INLINE cosh #-} + tanh = fmap tanh + {-# INLINE tanh #-} + asinh = fmap asinh + {-# INLINE asinh #-} + acosh = fmap acosh + {-# INLINE acosh #-} + atanh = fmap atanh + {-# INLINE atanh #-} + + + + +instance Functor (Stencil ix e) where + fmap f stencil@(Stencil {stencilFunc = g}) = stencil {stencilFunc = stF} + where + stF s = Value . f . unValue . g s + {-# INLINE stF #-} + {-# INLINE fmap #-} + + +-- TODO: Figure out interchange law (u <*> pure y = pure ($ y) <*> u) and issue +-- with discarding size and center. Best idea so far is to increase stencil size to +-- the maximum one and shift the center of the other stencil so that they both match +-- up. This approach would also remove requirement to validate the result +-- Stencil - both stencils are trusted, increasing the size will not affect the +-- safety. +instance (Default e, Index ix) => Applicative (Stencil ix e) where + pure a = Stencil Edge (pureIndex 1) zeroIndex (const (const (Value a))) + {-# INLINE pure #-} + (<*>) (Stencil _ sSz1 sC1 f1) (Stencil sB sSz2 sC2 f2) = + validateStencil def (Stencil sB newSz maxCenter stF) + where + stF gV !ix = Value ((unValue (f1 gV ix)) (unValue (f2 gV ix))) + {-# INLINE stF #-} + !newSz = + liftIndex2 + (+) + maxCenter + (liftIndex2 max (liftIndex2 (-) sSz1 sC1) (liftIndex2 (-) sSz2 sC2)) + !maxCenter = liftIndex2 max sC1 sC2 + {-# INLINE (<*>) #-} + +instance (Index ix, Default e, Num a) => Num (Stencil ix e a) where + (+) = liftA2 (+) + {-# INLINE (+) #-} + (-) = liftA2 (-) + {-# INLINE (-) #-} + (*) = liftA2 (*) + {-# INLINE (*) #-} + negate = fmap negate + {-# INLINE negate #-} + abs = fmap abs + {-# INLINE abs #-} + signum = fmap signum + {-# INLINE signum #-} + fromInteger = pure . fromInteger + {-# INLINE fromInteger #-} + +instance (Index ix, Default e, Fractional a) => Fractional (Stencil ix e a) where + (/) = liftA2 (/) + {-# INLINE (/) #-} + recip = fmap recip + {-# INLINE recip #-} + fromRational = pure . fromRational + {-# INLINE fromRational #-} + +instance (Index ix, Default e, Floating a) => Floating (Stencil ix e a) where + pi = pure pi + {-# INLINE pi #-} + exp = fmap exp + {-# INLINE exp #-} + log = fmap log + {-# INLINE log #-} + sqrt = fmap sqrt + {-# INLINE sqrt #-} + (**) = liftA2 (**) + {-# INLINE (**) #-} + logBase = liftA2 logBase + {-# INLINE logBase #-} + sin = fmap sin + {-# INLINE sin #-} + cos = fmap cos + {-# INLINE cos #-} + tan = fmap tan + {-# INLINE tan #-} + asin = fmap asin + {-# INLINE asin #-} + acos = fmap acos + {-# INLINE acos #-} + atan = fmap atan + {-# INLINE atan #-} + sinh = fmap sinh + {-# INLINE sinh #-} + cosh = fmap cosh + {-# INLINE cosh #-} + tanh = fmap tanh + {-# INLINE tanh #-} + asinh = fmap asinh + {-# INLINE asinh #-} + acosh = fmap acosh + {-# INLINE acosh #-} + atanh = fmap atanh + {-# INLINE atanh #-} + + +safeStencilIndex :: Index ix => Array D ix e -> ix -> e +safeStencilIndex DArray {..} ix + | isSafeIndex dSize ix = dUnsafeIndex ix + | otherwise = + error $ + "Index is out of bounds: " ++ show ix ++ " for stencil size: " ++ show dSize + + +-- | Make sure constructed stencil doesn't index outside the allowed stencil size boundary. +validateStencil + :: Index ix + => e -> Stencil ix e a -> Stencil ix e a +validateStencil d s@(Stencil _ sSz sCenter stencil) = + let valArr = DArray Seq sSz (const d) + in stencil (Value . safeStencilIndex valArr) sCenter `seq` s +{-# INLINE validateStencil #-} + diff --git a/src/Data/Massiv/Array/Unsafe.hs b/src/Data/Massiv/Array/Unsafe.hs new file mode 100644 --- /dev/null +++ b/src/Data/Massiv/Array/Unsafe.hs @@ -0,0 +1,118 @@ +{-# LANGUAGE BangPatterns #-} +{-# LANGUAGE FlexibleContexts #-} +{-# LANGUAGE FlexibleInstances #-} +{-# LANGUAGE MultiParamTypeClasses #-} +-- | +-- Module : Data.Massiv.Array.Ops.Unsafe +-- Copyright : (c) Alexey Kuleshevich 2018 +-- License : BSD3 +-- Maintainer : Alexey Kuleshevich +-- Stability : experimental +-- Portability : non-portable +-- +module Data.Massiv.Array.Unsafe + ( -- * Creation + unsafeMakeArray + , unsafeGenerateArray + , unsafeGenerateArrayP + , unsafeGenerateM + -- * Indexing + , unsafeIndex + , unsafeLinearIndex + , unsafeLinearIndexM + -- * Manipulations + , unsafeBackpermute + , unsafeTraverse + , unsafeTraverse2 + , unsafeResize + , unsafeExtract + -- * Slicing + , unsafeSlice + , unsafeOuterSlice + , unsafeInnerSlice + -- * Mutable interface + , unsafeThaw + , unsafeFreeze + , unsafeNew + , unsafeNewZero + , unsafeRead + , unsafeLinearRead + , unsafeWrite + , unsafeLinearWrite + ) where + +import Control.Monad.Primitive (PrimMonad (..)) +import Control.Monad.ST (runST) +import Data.Massiv.Array.Delayed.Internal (D) +import Data.Massiv.Core.Common +import Data.Massiv.Core.Scheduler +import System.IO.Unsafe (unsafePerformIO) + + +unsafeBackpermute :: (Source r' ix' e, Index ix) => + ix -> (ix -> ix') -> Array r' ix' e -> Array D ix e +unsafeBackpermute !sz ixF !arr = + unsafeMakeArray (getComp arr) sz $ \ !ix -> unsafeIndex arr (ixF ix) +{-# INLINE unsafeBackpermute #-} + + +unsafeTraverse + :: (Source r1 ix1 e1, Index ix) + => ix + -> ((ix1 -> e1) -> ix -> e) + -> Array r1 ix1 e1 + -> Array D ix e +unsafeTraverse sz f arr1 = + unsafeMakeArray (getComp arr1) sz (f (unsafeIndex arr1)) +{-# INLINE unsafeTraverse #-} + + +unsafeTraverse2 + :: (Source r1 ix1 e1, Source r2 ix2 e2, Index ix) + => ix + -> ((ix1 -> e1) -> (ix2 -> e2) -> ix -> e) + -> Array r1 ix1 e1 + -> Array r2 ix2 e2 + -> Array D ix e +unsafeTraverse2 sz f arr1 arr2 = + unsafeMakeArray (getComp arr1) sz (f (unsafeIndex arr1) (unsafeIndex arr2)) +{-# INLINE unsafeTraverse2 #-} + + +-- | Read an array element +unsafeRead :: (Mutable r ix e, PrimMonad m) => + MArray (PrimState m) r ix e -> ix -> m e +unsafeRead !marr !ix = unsafeLinearRead marr (toLinearIndex (msize marr) ix) +{-# INLINE unsafeRead #-} + +-- | Write an element into array +unsafeWrite :: (Mutable r ix e, PrimMonad m) => + MArray (PrimState m) r ix e -> ix -> e -> m () +unsafeWrite !marr !ix = unsafeLinearWrite marr (toLinearIndex (msize marr) ix) +{-# INLINE unsafeWrite #-} + + +-- | Create an array sequentially using mutable interface +unsafeGenerateArray :: Mutable r ix e => ix -> (ix -> e) -> Array r ix e +unsafeGenerateArray !sz f = runST $ do + marr <- unsafeNew sz + iterLinearM_ sz 0 (totalElem sz) 1 (<) $ \ !k !ix -> + unsafeLinearWrite marr k (f ix) + unsafeFreeze Seq marr +{-# INLINE unsafeGenerateArray #-} + + +-- | Create an array in parallel using mutable interface +unsafeGenerateArrayP :: Mutable r ix e => [Int] -> ix -> (ix -> e) -> Array r ix e +unsafeGenerateArrayP wIds !sz f = unsafePerformIO $ do + marr <- unsafeNew sz + divideWork_ wIds sz $ \ !scheduler !chunkLength !totalLength !slackStart -> do + loopM_ 0 (< slackStart) (+ chunkLength) $ \ !start -> + scheduleWork scheduler $ + iterLinearM_ sz start (start + chunkLength) 1 (<) $ \ !k !ix -> + unsafeLinearWrite marr k (f ix) + scheduleWork scheduler $ + iterLinearM_ sz slackStart totalLength 1 (<) $ \ !k !ix -> + unsafeLinearWrite marr k (f ix) + unsafeFreeze (ParOn wIds) marr +{-# INLINE unsafeGenerateArrayP #-} diff --git a/src/Data/Massiv/Core.hs b/src/Data/Massiv/Core.hs new file mode 100644 --- /dev/null +++ b/src/Data/Massiv/Core.hs @@ -0,0 +1,64 @@ +{-# LANGUAGE BangPatterns #-} +{-# LANGUAGE CPP #-} +{-# LANGUAGE FlexibleContexts #-} +{-# LANGUAGE FlexibleInstances #-} +{-# LANGUAGE MultiParamTypeClasses #-} +{-# LANGUAGE PatternSynonyms #-} +{-# LANGUAGE ScopedTypeVariables #-} +{-# LANGUAGE TypeFamilies #-} +{-# LANGUAGE UndecidableInstances #-} +#if __GLASGOW_HASKELL__ >= 800 + {-# OPTIONS_GHC -fno-warn-duplicate-exports #-} +#endif +-- | +-- Module : Data.Massiv.Core +-- Copyright : (c) Alexey Kuleshevich 2018 +-- License : BSD3 +-- Maintainer : Alexey Kuleshevich +-- Stability : experimental +-- Portability : non-portable +-- +module Data.Massiv.Core + ( Array(List, unList) + , Elt + , EltRepr + , Construct + , Source + , Load(..) + , Size + , Slice + , OuterSlice(outerLength) + , InnerSlice + , Manifest + , Mutable + , Ragged(..) + , Nested(..) + , NestedStruct + , L(..) + , LN + , ListItem +#if __GLASGOW_HASKELL__ >= 800 + , Comp(Seq, Par, ParOn) + , pattern Par -- already exported above and only needed for Haddock +#else + , Comp(..) + , pattern Par +#endif + , module Data.Massiv.Core.Index + , elemsCount + , isEmpty + ) where + +import Data.Massiv.Core.Common hiding (unsafeGenerateM) +import Data.Massiv.Core.List +import Data.Massiv.Core.Index + +-- | /O(1)/ - Get the number of elements in the array +elemsCount :: Size r ix e => Array r ix e -> Int +elemsCount = totalElem . size +{-# INLINE elemsCount #-} + +-- | /O(1)/ - Check if array has no elements. +isEmpty :: Size r ix e => Array r ix e -> Bool +isEmpty !arr = 0 == elemsCount arr +{-# INLINE isEmpty #-} diff --git a/src/Data/Massiv/Core/Common.hs b/src/Data/Massiv/Core/Common.hs new file mode 100644 --- /dev/null +++ b/src/Data/Massiv/Core/Common.hs @@ -0,0 +1,312 @@ +{-# LANGUAGE BangPatterns #-} +{-# LANGUAGE DefaultSignatures #-} +{-# LANGUAGE FlexibleContexts #-} +{-# LANGUAGE MultiParamTypeClasses #-} +{-# LANGUAGE TypeFamilies #-} +{-# LANGUAGE UndecidableInstances #-} +-- | +-- Module : Data.Massiv.Core.Common +-- Copyright : (c) Alexey Kuleshevich 2018 +-- License : BSD3 +-- Maintainer : Alexey Kuleshevich +-- Stability : experimental +-- Portability : non-portable +module Data.Massiv.Core.Common + ( Array + , Elt + , EltRepr + , Construct(..) + , Source(..) + , Load(..) + , Size(..) + , Slice(..) + , OuterSlice(..) + , InnerSlice(..) + , Manifest(..) + , Mutable(..) + , Ragged(..) + , Nested(..) + , NestedStruct + , makeArray + , singleton + -- * Indexing + , (!?) + , index + , (!) + , index' + , (??) + , defaultIndex + , borderIndex + , evaluateAt + , module Data.Massiv.Core.Index + , module Data.Massiv.Core.Computation + ) where + +import Control.Monad.Primitive (PrimMonad (..)) +import Data.Massiv.Core.Computation +import Data.Massiv.Core.Index +import Data.Typeable + +-- | The array family. Representations @r@ describes how data is arranged or computed. All arrays +-- have a common property that each index @ix@ always maps to the same unique element, even if that +-- element does not exist in memory and has to be computed upon lookup. Data is always arranged in a +-- nested fasion, depth of which is controlled by @`Rank` ix@. +data family Array r ix e :: * + +type family EltRepr r ix :: * + +type family Elt r ix e :: * where + Elt r Ix1 e = e + Elt r ix e = Array (EltRepr r ix) (Lower ix) e + +type family NestedStruct r ix e :: * + +-- | Array types that can be constructed. +class (Typeable r, Index ix) => Construct r ix e where + + -- | Get computation strategy of this array + getComp :: Array r ix e -> Comp + + -- | Set computation strategy for this array + setComp :: Comp -> Array r ix e -> Array r ix e + + -- | Construct an array. No size validation is performed. + unsafeMakeArray :: Comp -> ix -> (ix -> e) -> Array r ix e + + +-- | An array that contains size information. They can be resized and new arrays extracted from it +-- in constant time. +class Construct r ix e => Size r ix e where + + -- | /O(1)/ - Get the size of an array + size :: Array r ix e -> ix + + -- | /O(1)/ - Change the size of an array. New size is not validated. + unsafeResize :: Index ix' => ix' -> Array r ix e -> Array r ix' e + + -- | /O(1)/ - Extract a portion of an array. Staring index and new size are + -- not validated. + unsafeExtract :: ix -> ix -> Array r ix e -> Array (EltRepr r ix) ix e + + +-- | Arrays that can be used as source to practically any manipulation function. +class Size r ix e => Source r ix e where + + -- | Lookup element in the array. No bounds check is performed and access of + -- arbitrary memory is possible when invalid index is supplied. + unsafeIndex :: Array r ix e -> ix -> e + unsafeIndex !arr = unsafeLinearIndex arr . toLinearIndex (size arr) + {-# INLINE unsafeIndex #-} + + -- | Lookup element in the array using flat index in a row-major fasion. No + -- bounds check is performed + unsafeLinearIndex :: Array r ix e -> Int -> e + unsafeLinearIndex !arr = unsafeIndex arr . fromLinearIndex (size arr) + {-# INLINE unsafeLinearIndex #-} + +-- | Any array that can be computed +class Size r ix e => Load r ix e where + -- | Load an array into memory sequentially + loadS + :: Monad m => + Array r ix e -- ^ Array that is being loaded + -> (Int -> m e) -- ^ Function that reads an element from target array + -> (Int -> e -> m ()) -- ^ Function that writes an element into target array + -> m () + + -- | Load an array into memory in parallel + loadP + :: [Int] -- ^ List of capabilities to run workers on, as described in + -- `Control.Concurrent.forkOn`. Empty list will imply all + -- capabilities, i.e. run on all cores available through @+RTS -N@. + -> Array r ix e -- ^ Array that is being loaded + -> (Int -> IO e) -- ^ Function that reads an element from target array + -> (Int -> e -> IO ()) -- ^ Function that writes an element into target array + -> IO () + +class OuterSlice r ix e where + -- | /O(1)/ - Take a slice out of an array from the outside + unsafeOuterSlice :: Array r ix e -> Int -> Elt r ix e + + outerLength :: Array r ix e -> Int + default outerLength :: Size r ix e => Array r ix e -> Int + outerLength = headDim . size + +class Size r ix e => InnerSlice r ix e where + unsafeInnerSlice :: Array r ix e -> (Lower ix, Int) -> Int -> Elt r ix e + +class Size r ix e => Slice r ix e where + unsafeSlice :: Array r ix e -> ix -> ix -> Dim -> Maybe (Elt r ix e) + + +-- | Manifest arrays are backed by actual memory and values are looked up versus +-- computed as it is with delayed arrays. Because of this fact indexing functions +-- @(`!`)@, @(`!?`)@, etc. are constrained to manifest arrays only. +class Source r ix e => Manifest r ix e where + + unsafeLinearIndexM :: Array r ix e -> Int -> e + + +class Manifest r ix e => Mutable r ix e where + data MArray s r ix e :: * + + -- | Get the size of a mutable array. + msize :: MArray s r ix e -> ix + + unsafeThaw :: PrimMonad m => + Array r ix e -> m (MArray (PrimState m) r ix e) + + unsafeFreeze :: PrimMonad m => + Comp -> MArray (PrimState m) r ix e -> m (Array r ix e) + + -- | Create new mutable array, leaving it's elements uninitialized. Size isn't validated + -- either. + unsafeNew :: PrimMonad m => + ix -> m (MArray (PrimState m) r ix e) + + -- | Create new mutable array, leaving it's elements uninitialized. Size isn't validated + -- either. + unsafeNewZero :: PrimMonad m => + ix -> m (MArray (PrimState m) r ix e) + + unsafeLinearRead :: PrimMonad m => + MArray (PrimState m) r ix e -> Int -> m e + + unsafeLinearWrite :: PrimMonad m => + MArray (PrimState m) r ix e -> Int -> e -> m () + + +class Nested r ix e where + fromNested :: NestedStruct r ix e -> Array r ix e + + toNested :: Array r ix e -> NestedStruct r ix e + + +class Construct r ix e => Ragged r ix e where + + empty :: Comp -> Array r ix e + + isNull :: Array r ix e -> Bool + + cons :: Elt r ix e -> Array r ix e -> Array r ix e + + uncons :: Array r ix e -> Maybe (Elt r ix e, Array r ix e) + + -- head :: Array r ix e -> Maybe (Elt r ix e, Array r ix e) + + -- tail :: Array r ix e -> Maybe (Elt r ix e, Array r ix e) + + unsafeGenerateM :: Monad m => Comp -> ix -> (ix -> m e) -> m (Array r ix e) + + edgeSize :: Array r ix e -> ix + + --outerLength :: Array r ix e -> Int + + flatten :: Array r ix e -> Array r Ix1 e + + loadRagged :: + (IO () -> IO ()) -> (Int -> e -> IO a) -> Int -> Int -> Lower ix -> Array r ix e -> IO () + + -- TODO: test property: + -- (read $ raggedFormat show "\n" (ls :: Array L (IxN n) Int)) == ls + raggedFormat :: (e -> String) -> String -> Array r ix e -> String + + + +-- | Create an Array. Resulting type either has to be unambiguously inferred or restricted manually, +-- like in the example below. +-- +-- >>> makeArray Seq (3 :. 4) (\ (i :. j) -> if i == j then i else 0) :: Array D Ix2 Int +-- (Array D Seq (3 :. 4) +-- [ [ 0,0,0,0 ] +-- , [ 0,1,0,0 ] +-- , [ 0,0,2,0 ] +-- ]) +-- +makeArray :: Construct r ix e => + Comp -- ^ Computation strategy. Useful constructors are `Seq` and `Par` + -> ix -- ^ Size of the result array. Negative values will result in an empty array. + -> (ix -> e) -- ^ Function to generate elements at a particular index + -> Array r ix e +makeArray !c = unsafeMakeArray c . liftIndex (max 0) +{-# INLINE makeArray #-} + + +-- | Create an Array with a single element. +singleton :: Construct r ix e => + Comp -- ^ Computation strategy + -> e -- ^ The element + -> Array r ix e +singleton !c = unsafeMakeArray c (pureIndex 1) . const +{-# INLINE singleton #-} + + +infixl 4 !, !?, ?? + +-- | Infix version of `index'`. +(!) :: Manifest r ix e => Array r ix e -> ix -> e +(!) = index' +{-# INLINE (!) #-} + + +-- | Infix version of `index`. +(!?) :: Manifest r ix e => Array r ix e -> ix -> Maybe e +(!?) = index +{-# INLINE (!?) #-} + + +-- | /O(1)/ - Lookup an element in the array, where array can itself be +-- `Nothing`. This operator is useful when used together with slicing or other +-- functions that return `Maybe` array: +-- +-- >>> (fromList Seq [[[1,2,3]],[[4,5,6]]] :: Maybe (Array U Ix3 Int)) ??> 1 ?? (0 :. 2) +-- Just 6 +-- +(??) :: Manifest r ix e => Maybe (Array r ix e) -> ix -> Maybe e +(??) Nothing = const Nothing +(??) (Just arr) = (arr !?) +{-# INLINE (??) #-} + +-- | /O(1)/ - Lookup an element in the array. Returns `Nothing`, when index is out +-- of bounds, `Just` element otherwise. +index :: Manifest r ix e => Array r ix e -> ix -> Maybe e +index arr = handleBorderIndex (Fill Nothing) (size arr) (Just . unsafeIndex arr) +{-# INLINE index #-} + +-- | /O(1)/ - Lookup an element in the array, while using default element when +-- index is out of bounds. +defaultIndex :: Manifest r ix e => e -> Array r ix e -> ix -> e +defaultIndex defVal = borderIndex (Fill defVal) +{-# INLINE defaultIndex #-} + +-- | /O(1)/ - Lookup an element in the array. Use a border resolution technique +-- when index is out of bounds. +borderIndex :: Manifest r ix e => Border e -> Array r ix e -> ix -> e +borderIndex border arr = handleBorderIndex border (size arr) (unsafeIndex arr) +{-# INLINE borderIndex #-} + +-- | /O(1)/ - Lookup an element in the array. Throw an error if index is out of bounds. +index' :: Manifest r ix e => Array r ix e -> ix -> e +index' arr ix = + borderIndex (Fill (errorIx "Data.Massiv.Array.index" (size arr) ix)) arr ix +{-# INLINE index' #-} + + +-- | This is just like `index'` function, but it allows getting values from +-- delayed arrays as well as manifest. As the name suggests, indexing into a +-- delayed array at the same index multiple times will cause evaluation of the +-- value each time and can destroy the performace if used without care. +evaluateAt :: Source r ix e => Array r ix e -> ix -> e +evaluateAt !arr !ix = + handleBorderIndex + (Fill (errorIx "Data.Massiv.Array.evaluateAt" (size arr) ix)) + (size arr) + (unsafeIndex arr) + ix +{-# INLINE evaluateAt #-} + + +-- errorImpossible :: String -> a +-- errorImpossible loc = +-- error $ "Please report this error. Impossible happend at: " ++ loc +-- {-# NOINLINE errorImpossible #-} diff --git a/src/Data/Massiv/Core/Computation.hs b/src/Data/Massiv/Core/Computation.hs new file mode 100644 --- /dev/null +++ b/src/Data/Massiv/Core/Computation.hs @@ -0,0 +1,60 @@ +{-# LANGUAGE PatternSynonyms #-} +{-# LANGUAGE ScopedTypeVariables #-} +-- | +-- Module : Data.Massiv.Core.Computation +-- Copyright : (c) Alexey Kuleshevich 2018 +-- License : BSD3 +-- Maintainer : Alexey Kuleshevich +-- Stability : experimental +-- Portability : non-portable +-- +module Data.Massiv.Core.Computation + ( Comp(..) + , pattern Par + ) where + +import Control.DeepSeq (NFData (..), deepseq) + +-- | Computation type to use. +data Comp + = Seq -- ^ Sequential computation + | ParOn [Int] + -- ^ Use `Par` instead to use your CPU to the fullest. Also don't forget to compile + -- the program with @-threaded@ flag. + -- + -- Parallel computation with a list of capabilities to run computation + -- on. Specifying an empty list (@ParOn []@) or using `Par` will result in + -- utilization of all available capabilities, which are set at runtime by + -- @+RTS -Nx@ or at compile time by GHC flag @-with-rtsopts=-Nx@, + -- where @x@ is the number of capabilities. Ommiting @x@ in above flags + -- defaults to number available cores. + deriving (Show, Eq) + +-- | Parallel computation using all available cores. +pattern Par :: Comp +pattern Par <- ParOn [] where + Par = ParOn [] + +instance NFData Comp where + rnf comp = + case comp of + Seq -> () + Par -> () + ParOn wIds -> wIds `deepseq` () + {-# INLINE rnf #-} + +instance Monoid Comp where + mempty = Seq + {-# INLINE mempty #-} + mappend = joinComp + {-# INLINE mappend #-} + + +joinComp :: Comp -> Comp -> Comp +joinComp Par _ = Par +joinComp _ Par = Par +joinComp (ParOn w1) (ParOn w2) = ParOn $ w1 ++ w2 +joinComp c@(ParOn _) _ = c +joinComp _ c@(ParOn _) = c +joinComp _ _ = Seq +{-# INLINE joinComp #-} diff --git a/src/Data/Massiv/Core/Index.hs b/src/Data/Massiv/Core/Index.hs new file mode 100644 --- /dev/null +++ b/src/Data/Massiv/Core/Index.hs @@ -0,0 +1,164 @@ +{-# LANGUAGE BangPatterns #-} +{-# LANGUAGE DefaultSignatures #-} +{-# LANGUAGE FlexibleContexts #-} +{-# LANGUAGE FlexibleInstances #-} +{-# LANGUAGE GeneralizedNewtypeDeriving #-} +{-# LANGUAGE MultiParamTypeClasses #-} +{-# LANGUAGE TypeFamilies #-} +-- | +-- Module : Data.Massiv.Core.Index +-- Copyright : (c) Alexey Kuleshevich 2018 +-- License : BSD3 +-- Maintainer : Alexey Kuleshevich +-- Stability : experimental +-- Portability : non-portable +-- +module Data.Massiv.Core.Index + ( module Data.Massiv.Core.Index.Ix + , Border(..) + , handleBorderIndex + , module Data.Massiv.Core.Index.Class + , isSafeSize + , isNonEmpty + , headDim + , tailDim + , lastDim + , initDim + , iterLinearM + , iterLinearM_ + , module Data.Massiv.Core.Iterator + ) where + +import Control.DeepSeq +import Data.Massiv.Core.Index.Class +import Data.Massiv.Core.Index.Ix +import Data.Massiv.Core.Iterator + + +-- | Approach to be used near the borders during various transformations. +-- Whenever a function needs information not only about an element of interest, but +-- also about it's neighbours, it will go out of bounds around the image edges, +-- hence is this set of approaches that can be used in such situtation. +data Border e = + Fill e -- ^ Fill in a constant element. + -- + -- @ + -- outside | Image | outside + -- ('Fill' 0) : 0 0 0 0 | 1 2 3 4 | 0 0 0 0 + -- @ + -- + | Wrap -- ^ Wrap around from the opposite border of the array. + -- + -- @ + -- outside | Image | outside + -- 'Wrap' : 1 2 3 4 | 1 2 3 4 | 1 2 3 4 + -- @ + -- + | Edge -- ^ Replicate the element at the edge. + -- + -- @ + -- outside | Image | outside + -- 'Edge' : 1 1 1 1 | 1 2 3 4 | 4 4 4 4 + -- @ + -- + | Reflect -- ^ Mirror like reflection. + -- + -- @ + -- outside | Image | outside + -- 'Reflect' : 4 3 2 1 | 1 2 3 4 | 4 3 2 1 + -- @ + -- + | Continue -- ^ Also mirror like reflection, but without repeating the edge element. + -- + -- @ + -- outside | Image | outside + -- 'Continue' : 1 4 3 2 | 1 2 3 4 | 3 2 1 4 + -- @ + -- + deriving (Eq, Show) + +instance NFData e => NFData (Border e) where + rnf b = case b of + Fill e -> rnf e + Wrap -> () + Edge -> () + Reflect -> () + Continue -> () + + +-- | Apply a border resolution technique to an index +handleBorderIndex :: + Index ix + => Border e -- ^ Broder resolution technique + -> ix -- ^ Size + -> (ix -> e) -- ^ Index function that produces an element + -> ix -- ^ Index + -> e +handleBorderIndex border !sz getVal !ix = + case border of + Fill val -> if isSafeIndex sz ix then getVal ix else val + Wrap -> getVal (repairIndex sz ix (flip mod) (flip mod)) + Edge -> getVal (repairIndex sz ix (const (const 0)) (\ !k _ -> k - 1)) + Reflect -> getVal (repairIndex sz ix (\ !k !i -> (abs i - 1) `mod` k) + (\ !k !i -> (-i - 1) `mod` k)) + Continue -> getVal (repairIndex sz ix (\ !k !i -> abs i `mod` k) + (\ !k !i -> (-i - 2) `mod` k)) +{-# INLINE [1] handleBorderIndex #-} + + + +-- | Checks whether the size is valid. +isSafeSize :: Index ix => ix -> Bool +isSafeSize = (zeroIndex >=) +{-# INLINE [1] isSafeSize #-} + + +-- | Checks whether array with this size can hold at least one element. +isNonEmpty :: Index ix => ix -> Bool +isNonEmpty !sz = isSafeIndex sz zeroIndex +{-# INLINE [1] isNonEmpty #-} + + +headDim :: Index ix => ix -> Int +headDim = fst . unconsDim +{-# INLINE [1] headDim #-} + +tailDim :: Index ix => ix -> Lower ix +tailDim = snd . unconsDim +{-# INLINE [1] tailDim #-} + +lastDim :: Index ix => ix -> Int +lastDim = snd . unsnocDim +{-# INLINE [1] lastDim #-} + +initDim :: Index ix => ix -> Lower ix +initDim = fst . unsnocDim +{-# INLINE [1] initDim #-} + + +-- | Iterate over N-dimensional space from start to end with accumulator +iterLinearM :: (Index ix, Monad m) + => ix -- ^ Size + -> Int -- ^ Linear start + -> Int -- ^ Linear end + -> Int -- ^ Increment + -> (Int -> Int -> Bool) -- ^ Continuation condition (continue if True) + -> a -- ^ Accumulator + -> (Int -> ix -> a -> m a) + -> m a +iterLinearM !sz !k0 !k1 !inc cond !acc f = + loopM k0 (`cond` k1) (+ inc) acc $ \ !i !acc0 -> f i (fromLinearIndex sz i) acc0 +{-# INLINE iterLinearM #-} + +iterLinearM_ :: (Index ix, Monad m) => + ix -- ^ Size + -> Int -- ^ Start + -> Int -- ^ End + -> Int -- ^ Increment + -> (Int -> Int -> Bool) -- ^ Continuation condition + -> (Int -> ix -> m ()) -- ^ Monadic action that takes index in both forms + -> m () +iterLinearM_ !sz !k0 !k1 !inc cond f = + loopM_ k0 (`cond` k1) (+ inc) $ \ !i -> f i (fromLinearIndex sz i) +{-# INLINE iterLinearM_ #-} + diff --git a/src/Data/Massiv/Core/Index/Class.hs b/src/Data/Massiv/Core/Index/Class.hs new file mode 100644 --- /dev/null +++ b/src/Data/Massiv/Core/Index/Class.hs @@ -0,0 +1,381 @@ +{-# LANGUAGE DataKinds #-} +{-# LANGUAGE BangPatterns #-} +{-# LANGUAGE DefaultSignatures #-} +{-# LANGUAGE FlexibleContexts #-} +{-# LANGUAGE FlexibleInstances #-} +{-# LANGUAGE GeneralizedNewtypeDeriving #-} +{-# LANGUAGE MultiParamTypeClasses #-} +{-# LANGUAGE TypeFamilies #-} +-- | +-- Module : Data.Massiv.Core.Index.Class +-- Copyright : (c) Alexey Kuleshevich 2018 +-- License : BSD3 +-- Maintainer : Alexey Kuleshevich +-- Stability : experimental +-- Portability : non-portable +-- +module Data.Massiv.Core.Index.Class where + +import Control.DeepSeq (NFData (..)) +import Data.Functor.Identity (runIdentity) +import Data.Massiv.Core.Iterator +import GHC.TypeLits + +newtype Dim = Dim Int deriving (Show, Eq, Ord, Num, Real, Integral, Enum) + +data Ix0 = Ix0 deriving (Eq, Ord, Show) + +type Ix1T = Int + +type Ix2T = (Int, Int) + +type Ix3T = (Int, Int, Int) + +type Ix4T = (Int, Int, Int, Int) + +type Ix5T = (Int, Int, Int, Int, Int) + +type family Lower ix :: * + +type instance Lower Ix1T = Ix0 +type instance Lower Ix2T = Ix1T +type instance Lower Ix3T = Ix2T +type instance Lower Ix4T = Ix3T +type instance Lower Ix5T = Ix4T + + +class (Eq ix, Ord ix, Show ix, NFData ix) => Index ix where + type Rank ix :: Nat + + rank :: ix -> Dim + + -- | Total number of elements in an array of this size. + totalElem :: ix -> Int + + consDim :: Int -> Lower ix -> ix + + unconsDim :: ix -> (Int, Lower ix) + + snocDim :: Lower ix -> Int -> ix + + unsnocDim :: ix -> (Lower ix, Int) + + dropDim :: ix -> Dim -> Maybe (Lower ix) + + getIndex :: ix -> Dim -> Maybe Int + + setIndex :: ix -> Dim -> Int -> Maybe ix + + pureIndex :: Int -> ix + + -- | Zip together two indices with a function + liftIndex2 :: (Int -> Int -> Int) -> ix -> ix -> ix + + zeroIndex :: ix + zeroIndex = pureIndex 0 + {-# INLINE [1] zeroIndex #-} + + -- | Map a function over an index + liftIndex :: (Int -> Int) -> ix -> ix + liftIndex f = liftIndex2 (\_ i -> f i) zeroIndex + {-# INLINE [1] liftIndex #-} + + -- | Check whether index is within the size. + isSafeIndex :: ix -- ^ Size + -> ix -- ^ Index + -> Bool + default isSafeIndex :: Index (Lower ix) => ix -> ix -> Bool + isSafeIndex !sz !ix = isSafeIndex n0 i0 && isSafeIndex szL ixL + where + !(n0, szL) = unconsDim sz + !(i0, ixL) = unconsDim ix + {-# INLINE [1] isSafeIndex #-} + + -- | Produce linear index from size and index + toLinearIndex :: ix -- ^ Size + -> ix -- ^ Index + -> Int + + default toLinearIndex :: Index (Lower ix) => ix -> ix -> Int + toLinearIndex !sz !ix = toLinearIndex szL ixL * n + i + where !(szL, n) = unsnocDim sz + !(ixL, i) = unsnocDim ix + {-# INLINE [1] toLinearIndex #-} + + toLinearIndexAcc :: Int -> ix -> ix -> Int + default toLinearIndexAcc :: Index (Lower ix) => Int -> ix -> ix -> Int + toLinearIndexAcc !acc !sz !ix = toLinearIndexAcc (acc * n + i) szL ixL + where !(n, szL) = unconsDim sz + !(i, ixL) = unconsDim ix + {-# INLINE [1] toLinearIndexAcc #-} + + -- | Produce N Dim index from size and linear index + fromLinearIndex :: ix -> Int -> ix + default fromLinearIndex :: Index (Lower ix) => ix -> Int -> ix + fromLinearIndex sz k = consDim q ixL + where !(q, ixL) = fromLinearIndexAcc (snd (unconsDim sz)) k + {-# INLINE [1] fromLinearIndex #-} + + fromLinearIndexAcc :: ix -> Int -> (Int, ix) + default fromLinearIndexAcc :: Index (Lower ix) => ix -> Int -> (Int, ix) + fromLinearIndexAcc ix' !k = (q, consDim r ixL) + where !(m, ix) = unconsDim ix' + !(kL, ixL) = fromLinearIndexAcc ix k + !(q, r) = quotRem kL m + {-# INLINE [1] fromLinearIndexAcc #-} + + repairIndex :: ix -> ix -> (Int -> Int -> Int) -> (Int -> Int -> Int) -> ix + default repairIndex :: Index (Lower ix) + => ix -> ix -> (Int -> Int -> Int) -> (Int -> Int -> Int) -> ix + repairIndex !sz !ix rBelow rOver = + consDim (repairIndex n i rBelow rOver) (repairIndex szL ixL rBelow rOver) + where !(n, szL) = unconsDim sz + !(i, ixL) = unconsDim ix + {-# INLINE [1] repairIndex #-} + + iter :: ix -> ix -> Int -> (Int -> Int -> Bool) -> a -> (ix -> a -> a) -> a + iter sIx eIx inc cond acc f = + runIdentity $ iterM sIx eIx inc cond acc (\ix -> return . f ix) + {-# INLINE iter #-} + + iterM :: Monad m => + ix -- ^ Start index + -> ix -- ^ End index + -> Int -- ^ Increment + -> (Int -> Int -> Bool) -- ^ Continue iteration while predicate is True (eg. until end of row) + -> a -- ^ Initial value for an accumulator + -> (ix -> a -> m a) -- ^ Accumulator function + -> m a + default iterM :: (Index (Lower ix), Monad m) + => ix -> ix -> Int -> (Int -> Int -> Bool) -> a -> (ix -> a -> m a) -> m a + iterM !sIx !eIx !inc cond !acc f = + loopM k0 (`cond` k1) (+ inc) acc $ \ !i !acc0 -> + iterM sIxL eIxL inc cond acc0 $ \ !ix -> + f (consDim i ix) + where + !(k0, sIxL) = unconsDim sIx + !(k1, eIxL) = unconsDim eIx + {-# INLINE iterM #-} + + iterM_ :: Monad m => ix -> ix -> Int -> (Int -> Int -> Bool) -> (ix -> m a) -> m () + default iterM_ :: (Index (Lower ix), Monad m) + => ix -> ix -> Int -> (Int -> Int -> Bool) -> (ix -> m a) -> m () + iterM_ !sIx !eIx !inc cond f = + loopM_ k0 (`cond` k1) (+ inc) $ \ !i -> + iterM_ sIxL eIxL inc cond $ \ !ix -> + f (consDim i ix) + where + !(k0, sIxL) = unconsDim sIx + !(k1, eIxL) = unconsDim eIx + {-# INLINE iterM_ #-} + + +instance Index Ix1T where + type Rank Ix1T = 1 + rank _ = 1 + {-# INLINE [1] rank #-} + totalElem = id + {-# INLINE [1] totalElem #-} + isSafeIndex !k !i = 0 <= i && i < k + {-# INLINE [1] isSafeIndex #-} + toLinearIndex _ = id + {-# INLINE [1] toLinearIndex #-} + toLinearIndexAcc !acc m i = acc * m + i + {-# INLINE [1] toLinearIndexAcc #-} + fromLinearIndex _ = id + {-# INLINE [1] fromLinearIndex #-} + fromLinearIndexAcc n k = k `quotRem` n + {-# INLINE [1] fromLinearIndexAcc #-} + repairIndex !k !i rBelow rOver + | i < 0 = rBelow k i + | i >= k = rOver k i + | otherwise = i + {-# INLINE [1] repairIndex #-} + consDim i _ = i + {-# INLINE [1] consDim #-} + unconsDim i = (i, Ix0) + {-# INLINE [1] unconsDim #-} + snocDim _ i = i + {-# INLINE [1] snocDim #-} + unsnocDim i = (Ix0, i) + {-# INLINE [1] unsnocDim #-} + getIndex i 1 = Just i + getIndex _ _ = Nothing + {-# INLINE [1] getIndex #-} + setIndex _ 1 i = Just i + setIndex _ _ _ = Nothing + {-# INLINE [1] setIndex #-} + dropDim _ 1 = Just Ix0 + dropDim _ _ = Nothing + {-# INLINE [1] dropDim #-} + pureIndex i = i + {-# INLINE [1] pureIndex #-} + liftIndex f = f + {-# INLINE [1] liftIndex #-} + liftIndex2 f = f + {-# INLINE [1] liftIndex2 #-} + iter k0 k1 inc cond = loop k0 (`cond` k1) (+inc) + {-# INLINE iter #-} + iterM k0 k1 inc cond = loopM k0 (`cond` k1) (+inc) + {-# INLINE iterM #-} + iterM_ k0 k1 inc cond = loopM_ k0 (`cond` k1) (+inc) + {-# INLINE iterM_ #-} + + +instance Index Ix2T where + type Rank Ix2T = 2 + rank _ = 2 + {-# INLINE [1] rank #-} + totalElem !(m, n) = m * n + {-# INLINE [1] totalElem #-} + toLinearIndex !(_, n) !(i, j) = n * i + j + {-# INLINE [1] toLinearIndex #-} + fromLinearIndex (_, n) !k = k `quotRem` n + {-# INLINE [1] fromLinearIndex #-} + consDim = (,) + {-# INLINE [1] consDim #-} + unconsDim = id + {-# INLINE [1] unconsDim #-} + snocDim = (,) + {-# INLINE [1] snocDim #-} + unsnocDim = id + {-# INLINE [1] unsnocDim #-} + getIndex (i, _) 2 = Just i + getIndex (_, j) 1 = Just j + getIndex _ _ = Nothing + {-# INLINE [1] getIndex #-} + setIndex (_, j) 2 i = Just (i, j) + setIndex (i, _) 1 j = Just (i, j) + setIndex _ _ _ = Nothing + {-# INLINE [1] setIndex #-} + dropDim (_, j) 2 = Just j + dropDim (i, _) 1 = Just i + dropDim _ _ = Nothing + {-# INLINE [1] dropDim #-} + pureIndex i = (i, i) + {-# INLINE [1] pureIndex #-} + liftIndex2 f (i0, j0) (i1, j1) = (f i0 i1, f j0 j1) + {-# INLINE [1] liftIndex2 #-} + + +instance Index Ix3T where + type Rank Ix3T = 3 + rank _ = 3 + {-# INLINE [1] rank #-} + totalElem !(m, n, o) = m * n * o + {-# INLINE [1] totalElem #-} + consDim i (j, k) = (i, j, k) + {-# INLINE [1] consDim #-} + unconsDim (i, j, k) = (i, (j, k)) + {-# INLINE [1] unconsDim #-} + snocDim (i, j) k = (i, j, k) + {-# INLINE [1] snocDim #-} + unsnocDim (i, j, k) = ((i, j), k) + {-# INLINE [1] unsnocDim #-} + getIndex (i, _, _) 3 = Just i + getIndex (_, j, _) 2 = Just j + getIndex (_, _, k) 1 = Just k + getIndex _ _ = Nothing + {-# INLINE [1] getIndex #-} + setIndex (_, j, k) 3 i = Just (i, j, k) + setIndex (i, _, k) 2 j = Just (i, j, k) + setIndex (i, j, _) 1 k = Just (i, j, k) + setIndex _ _ _ = Nothing + {-# INLINE [1] setIndex #-} + dropDim (_, j, k) 3 = Just (j, k) + dropDim (i, _, k) 2 = Just (i, k) + dropDim (i, j, _) 1 = Just (i, j) + dropDim _ _ = Nothing + {-# INLINE [1] dropDim #-} + pureIndex i = (i, i, i) + {-# INLINE [1] pureIndex #-} + liftIndex2 f (i0, j0, k0) (i1, j1, k1) = (f i0 i1, f j0 j1, f k0 k1) + {-# INLINE [1] liftIndex2 #-} + + +instance Index Ix4T where + type Rank Ix4T = 4 + rank _ = 4 + {-# INLINE [1] rank #-} + totalElem !(n1, n2, n3, n4) = n1 * n2 * n3 * n4 + {-# INLINE [1] totalElem #-} + consDim i1 (i2, i3, i4) = (i1, i2, i3, i4) + {-# INLINE [1] consDim #-} + unconsDim (i1, i2, i3, i4) = (i1, (i2, i3, i4)) + {-# INLINE [1] unconsDim #-} + snocDim (i1, i2, i3) i4 = (i1, i2, i3, i4) + {-# INLINE [1] snocDim #-} + unsnocDim (i1, i2, i3, i4) = ((i1, i2, i3), i4) + {-# INLINE [1] unsnocDim #-} + getIndex (i1, _, _, _) 4 = Just i1 + getIndex ( _, i2, _, _) 3 = Just i2 + getIndex ( _, _, i3, _) 2 = Just i3 + getIndex ( _, _, _, i4) 1 = Just i4 + getIndex _ _ = Nothing + {-# INLINE [1] getIndex #-} + setIndex ( _, i2, i3, i4) 4 i1 = Just (i1, i2, i3, i4) + setIndex (i1, _, i3, i4) 3 i2 = Just (i1, i2, i3, i4) + setIndex (i1, i2, _, i4) 2 i3 = Just (i1, i2, i3, i4) + setIndex (i1, i2, i3, _) 1 i4 = Just (i1, i2, i3, i4) + setIndex _ _ _ = Nothing + {-# INLINE [1] setIndex #-} + dropDim ( _, i2, i3, i4) 4 = Just (i2, i3, i4) + dropDim (i1, _, i3, i4) 3 = Just (i1, i3, i4) + dropDim (i1, i2, _, i4) 2 = Just (i1, i2, i4) + dropDim (i1, i2, i3, _) 1 = Just (i1, i2, i3) + dropDim _ _ = Nothing + {-# INLINE [1] dropDim #-} + pureIndex i = (i, i, i, i) + {-# INLINE [1] pureIndex #-} + liftIndex2 f (i0, i1, i2, i3) (j0, j1, j2, j3) = (f i0 j0, f i1 j1, f i2 j2, f i3 j3) + {-# INLINE [1] liftIndex2 #-} + + +instance Index Ix5T where + type Rank Ix5T = 5 + rank _ = 5 + {-# INLINE [1] rank #-} + totalElem !(n1, n2, n3, n4, n5) = n1 * n2 * n3 * n4 * n5 + {-# INLINE [1] totalElem #-} + consDim i1 (i2, i3, i4, i5) = (i1, i2, i3, i4, i5) + {-# INLINE [1] consDim #-} + unconsDim (i1, i2, i3, i4, i5) = (i1, (i2, i3, i4, i5)) + {-# INLINE [1] unconsDim #-} + snocDim (i1, i2, i3, i4) i5 = (i1, i2, i3, i4, i5) + {-# INLINE [1] snocDim #-} + unsnocDim (i1, i2, i3, i4, i5) = ((i1, i2, i3, i4), i5) + {-# INLINE [1] unsnocDim #-} + getIndex (i1, _, _, _, _) 5 = Just i1 + getIndex ( _, i2, _, _, _) 4 = Just i2 + getIndex ( _, _, i3, _, _) 3 = Just i3 + getIndex ( _, _, _, i4, _) 2 = Just i4 + getIndex ( _, _, _, _, i5) 1 = Just i5 + getIndex _ _ = Nothing + {-# INLINE [1] getIndex #-} + setIndex ( _, i2, i3, i4, i5) 5 i1 = Just (i1, i2, i3, i4, i5) + setIndex (i1, _, i3, i4, i5) 4 i2 = Just (i1, i2, i3, i4, i5) + setIndex (i1, i2, _, i4, i5) 3 i3 = Just (i1, i2, i3, i4, i5) + setIndex (i1, i2, i3, _, i5) 2 i4 = Just (i1, i2, i3, i4, i5) + setIndex (i1, i2, i3, i4, _) 1 i5 = Just (i1, i2, i3, i4, i5) + setIndex _ _ _ = Nothing + {-# INLINE [1] setIndex #-} + dropDim ( _, i2, i3, i4, i5) 5 = Just (i2, i3, i4, i5) + dropDim (i1, _, i3, i4, i5) 4 = Just (i1, i3, i4, i5) + dropDim (i1, i2, _, i4, i5) 3 = Just (i1, i2, i4, i5) + dropDim (i1, i2, i3, _, i5) 2 = Just (i1, i2, i3, i5) + dropDim (i1, i2, i3, i4, _) 1 = Just (i1, i2, i3, i4) + dropDim _ _ = Nothing + {-# INLINE [1] dropDim #-} + pureIndex i = (i, i, i, i, i) + {-# INLINE [1] pureIndex #-} + liftIndex2 f (i0, i1, i2, i3, i4) (j0, j1, j2, j3, j4) = + (f i0 j0, f i1 j1, f i2 j2, f i3 j3, f i4 j4) + {-# INLINE [1] liftIndex2 #-} + + +errorIx :: (Show ix, Show ix') => String -> ix -> ix' -> a +errorIx fName sz ix = + error $ + fName ++ + ": Index out of bounds: " ++ show ix ++ " for Array of size: " ++ show sz +{-# NOINLINE errorIx #-} diff --git a/src/Data/Massiv/Core/Index/Ix.hs b/src/Data/Massiv/Core/Index/Ix.hs new file mode 100644 --- /dev/null +++ b/src/Data/Massiv/Core/Index/Ix.hs @@ -0,0 +1,528 @@ +{-# LANGUAGE BangPatterns #-} +{-# LANGUAGE CPP #-} +{-# LANGUAGE DataKinds #-} +{-# LANGUAGE FlexibleContexts #-} +{-# LANGUAGE FlexibleInstances #-} +{-# LANGUAGE GADTs #-} +{-# LANGUAGE MultiParamTypeClasses #-} +{-# LANGUAGE PatternSynonyms #-} +{-# LANGUAGE ScopedTypeVariables #-} +{-# LANGUAGE TypeFamilies #-} +{-# LANGUAGE TypeOperators #-} +{-# LANGUAGE UndecidableInstances #-} + +#if __GLASGOW_HASKELL__ >= 800 + +{-# LANGUAGE TypeFamilyDependencies #-} + +#endif +-- | +-- Module : Data.Massiv.Core.Index.Ix +-- Copyright : (c) Alexey Kuleshevich 2018 +-- License : BSD3 +-- Maintainer : Alexey Kuleshevich +-- Stability : experimental +-- Portability : non-portable +-- +module Data.Massiv.Core.Index.Ix where + +import Control.DeepSeq +import Control.Monad (liftM) +import Data.Massiv.Core.Index.Class +import Data.Monoid ((<>)) +import Data.Proxy +import qualified Data.Vector.Generic as V +import qualified Data.Vector.Generic.Mutable as VM +import qualified Data.Vector.Unboxed as VU +import GHC.TypeLits + + +infixr 5 :>, :. + +type Ix1 = Int + +pattern Ix1 :: Int -> Ix1 +pattern Ix1 i = i + +data Ix2 = (:.) {-# UNPACK #-} !Int {-# UNPACK #-} !Int +pattern Ix2 :: Int -> Int -> Ix2 +pattern Ix2 i j = i :. j + +type Ix3 = IxN 3 +pattern Ix3 :: Int -> Int -> Int -> Ix3 +pattern Ix3 i j k = i :> j :. k + +type Ix4 = IxN 4 +pattern Ix4 :: Int -> Int -> Int -> Int -> Ix4 +pattern Ix4 i j k l = i :> j :> k :. l + +type Ix5 = IxN 5 +pattern Ix5 :: Int -> Int -> Int -> Int -> Int -> Ix5 +pattern Ix5 i j k l m = i :> j :> k :> l :. m + + +#if __GLASGOW_HASKELL__ >= 800 + +data IxN (n :: Nat) where + (:>) :: {-# UNPACK #-} !Int -> !(Ix (n - 1)) -> IxN n + +type family Ix (n :: Nat) = r | r -> n where + Ix 0 = Ix0 + Ix 1 = Ix1 + Ix 2 = Ix2 + Ix n = IxN n + +#else + +data IxN (n :: Nat) where + (:>) :: Rank (Ix (n - 1)) ~ (n - 1) => {-# UNPACK #-} !Int -> !(Ix (n - 1)) -> IxN n + +type family Ix (n :: Nat) where + Ix 0 = Ix0 + Ix 1 = Ix1 + Ix 2 = Ix2 + Ix n = IxN n + +#endif + + +type instance Lower Ix2 = Ix1 +type instance Lower (IxN n) = Ix (n - 1) + + +instance Show Ix2 where + show (i :. j) = show i ++ " :. " ++ show j + +instance Show (Ix (n - 1)) => Show (IxN n) where + show (i :> ix) = show i ++ " :> " ++ show ix + + +instance Num Ix2 where + (+) = liftIndex2 (+) + {-# INLINE [1] (+) #-} + (-) = liftIndex2 (-) + {-# INLINE [1] (-) #-} + (*) = liftIndex2 (*) + {-# INLINE [1] (*) #-} + negate = liftIndex negate + {-# INLINE [1] negate #-} + abs = liftIndex abs + {-# INLINE [1] abs #-} + signum = liftIndex signum + {-# INLINE [1] signum #-} + fromInteger = pureIndex . fromInteger + {-# INLINE [1] fromInteger #-} + +instance Num Ix3 where + (+) = liftIndex2 (+) + {-# INLINE [1] (+) #-} + (-) = liftIndex2 (-) + {-# INLINE [1] (-) #-} + (*) = liftIndex2 (*) + {-# INLINE [1] (*) #-} + negate = liftIndex negate + {-# INLINE [1] negate #-} + abs = liftIndex abs + {-# INLINE [1] abs #-} + signum = liftIndex signum + {-# INLINE [1] signum #-} + fromInteger = pureIndex . fromInteger + {-# INLINE [1] fromInteger #-} + + +instance {-# OVERLAPPABLE #-} (4 <= n, + KnownNat n, + Index (Ix (n - 1)), +#if __GLASGOW_HASKELL__ < 800 + Rank (Ix ((n - 1) - 1)) ~ ((n - 1) - 1), +#endif + IxN (n - 1) ~ Ix (n - 1) + ) => Num (IxN n) where + (+) = liftIndex2 (+) + {-# INLINE [1] (+) #-} + (-) = liftIndex2 (-) + {-# INLINE [1] (-) #-} + (*) = liftIndex2 (*) + {-# INLINE [1] (*) #-} + negate = liftIndex negate + {-# INLINE [1] negate #-} + abs = liftIndex abs + {-# INLINE [1] abs #-} + signum = liftIndex signum + {-# INLINE [1] signum #-} + fromInteger = pureIndex . fromInteger + {-# INLINE [1] fromInteger #-} + + + +instance Bounded Ix2 where + minBound = pureIndex minBound + {-# INLINE minBound #-} + maxBound = pureIndex maxBound + {-# INLINE maxBound #-} + +instance Bounded Ix3 where + minBound = pureIndex minBound + {-# INLINE minBound #-} + maxBound = pureIndex maxBound + {-# INLINE maxBound #-} + +instance {-# OVERLAPPABLE #-} (4 <= n, + KnownNat n, + Index (Ix (n - 1)), +#if __GLASGOW_HASKELL__ < 800 + Rank (Ix ((n - 1) - 1)) ~ ((n - 1) - 1), +#endif + IxN (n - 1) ~ Ix (n - 1) + ) => Bounded (IxN n) where + minBound = pureIndex minBound + {-# INLINE minBound #-} + maxBound = pureIndex maxBound + {-# INLINE maxBound #-} + +instance NFData Ix2 where + rnf ix = ix `seq` () + +instance NFData (IxN n) where + rnf ix = ix `seq` () + + +instance Eq Ix2 where + (i1 :. j1) == (i2 :. j2) = i1 == i2 && j1 == j2 + +instance Eq (Ix (n - 1)) => Eq (IxN n) where + (i1 :> ix1) == (i2 :> ix2) = i1 == i2 && ix1 == ix2 + + +instance Ord Ix2 where + compare (i1 :. j1) (i2 :. j2) = compare i1 i2 <> compare j1 j2 + +instance Ord (Ix (n - 1)) => Ord (IxN n) where + compare (i1 :> ix1) (i2 :> ix2) = compare i1 i2 <> compare ix1 ix2 + + +toIx2 :: Ix2T -> Ix2 +toIx2 (i, j) = i :. j +{-# INLINE toIx2 #-} + +fromIx2 :: Ix2 -> Ix2T +fromIx2 (i :. j) = (i, j) +{-# INLINE fromIx2 #-} + +toIx3 :: Ix3T -> Ix3 +toIx3 (i, j, k) = i :> j :. k +{-# INLINE toIx3 #-} + +fromIx3 :: Ix3 -> Ix3T +fromIx3 (i :> j :. k) = (i, j, k) +{-# INLINE fromIx3 #-} + +toIx4 :: Ix4T -> Ix4 +toIx4 (i, j, k, l) = i :> j :> k :. l +{-# INLINE toIx4 #-} + +fromIx4 :: Ix4 -> Ix4T +fromIx4 (i :> j :> k :. l) = (i, j, k, l) +{-# INLINE fromIx4 #-} + +toIx5 :: Ix5T -> Ix5 +toIx5 (i, j, k, l, m) = i :> j :> k :> l :. m +{-# INLINE toIx5 #-} + +fromIx5 :: Ix5 -> Ix5T +fromIx5 (i :> j :> k :> l :. m) = (i, j, k, l, m) +{-# INLINE fromIx5 #-} + + +instance {-# OVERLAPPING #-} Index Ix2 where + type Rank Ix2 = 2 + rank _ = 2 + {-# INLINE [1] rank #-} + totalElem (m :. n) = m * n + {-# INLINE [1] totalElem #-} + isSafeIndex (m :. n) (i :. j) = 0 <= i && 0 <= j && i < m && j < n + {-# INLINE [1] isSafeIndex #-} + toLinearIndex (_ :. n) (i :. j) = n * i + j + {-# INLINE [1] toLinearIndex #-} + fromLinearIndex (_ :. n) k = case k `quotRem` n of + (i, j) -> i :. j + {-# INLINE [1] fromLinearIndex #-} + consDim = (:.) + {-# INLINE [1] consDim #-} + unconsDim (i :. ix) = (i, ix) + {-# INLINE [1] unconsDim #-} + snocDim i j = i :. j + {-# INLINE [1] snocDim #-} + unsnocDim (i :. j) = (i, j) + {-# INLINE [1] unsnocDim #-} + getIndex (i :. _) 2 = Just i + getIndex (_ :. j) 1 = Just j + getIndex _ _ = Nothing + {-# INLINE [1] getIndex #-} + setIndex (_ :. j) 2 i = Just (i :. j) + setIndex (i :. _) 1 j = Just (i :. j) + setIndex _ _ _ = Nothing + {-# INLINE [1] setIndex #-} + dropDim (_ :. j) 2 = Just j + dropDim (i :. _) 1 = Just i + dropDim _ _ = Nothing + {-# INLINE [1] dropDim #-} + pureIndex i = i :. i + {-# INLINE [1] pureIndex #-} + liftIndex f (i :. j) = f i :. f j + {-# INLINE [1] liftIndex #-} + liftIndex2 f (i0 :. j0) (i1 :. j1) = f i0 i1 :. f j0 j1 + {-# INLINE [1] liftIndex2 #-} + repairIndex (n :. szL) (i :. ixL) rBelow rOver = + repairIndex n i rBelow rOver :. repairIndex szL ixL rBelow rOver + {-# INLINE [1] repairIndex #-} + + +instance {-# OVERLAPPING #-} Index (IxN 3) where + type Rank Ix3 = 3 + rank _ = 3 + {-# INLINE [1] rank #-} + totalElem (m :> n :. o) = m * n * o + {-# INLINE [1] totalElem #-} + isSafeIndex (m :> n :. o) (i :> j :. k) = + 0 <= i && 0 <= j && 0 <= k && i < m && j < n && k < o + {-# INLINE [1] isSafeIndex #-} + toLinearIndex (_ :> n :. o) (i :> j :. k) = (n * i + j) * o + k + {-# INLINE [1] toLinearIndex #-} + fromLinearIndex (_ :> ix) k = let !(q, ixL) = fromLinearIndexAcc ix k in q :> ixL + {-# INLINE [1] fromLinearIndex #-} + consDim = (:>) + {-# INLINE [1] consDim #-} + unconsDim (i :> ix) = (i, ix) + {-# INLINE [1] unconsDim #-} + snocDim (i :. j) k = i :> j :. k + {-# INLINE [1] snocDim #-} + unsnocDim (i :> j :. k) = (i :. j, k) + {-# INLINE [1] unsnocDim #-} + getIndex (i :> _ :. _) 3 = Just i + getIndex (_ :> j :. _) 2 = Just j + getIndex (_ :> _ :. k) 1 = Just k + getIndex _ _ = Nothing + {-# INLINE [1] getIndex #-} + setIndex (_ :> j :. k) 3 i = Just (i :> j :. k) + setIndex (i :> _ :. k) 2 j = Just (i :> j :. k) + setIndex (i :> j :. _) 1 k = Just (i :> j :. k) + setIndex _ _ _ = Nothing + {-# INLINE [1] setIndex #-} + dropDim (_ :> j :. k) 3 = Just (j :. k) + dropDim (i :> _ :. k) 2 = Just (i :. k) + dropDim (i :> j :. _) 1 = Just (i :. j) + dropDim _ _ = Nothing + {-# INLINE [1] dropDim #-} + pureIndex i = i :> i :. i + {-# INLINE [1] pureIndex #-} + liftIndex f (i :> j :. k) = f i :> f j :. f k + {-# INLINE [1] liftIndex #-} + liftIndex2 f (i0 :> j0 :. k0) (i1 :> j1 :. k1) = f i0 i1 :> f j0 j1 :. f k0 k1 + {-# INLINE [1] liftIndex2 #-} + repairIndex (n :> szL) (i :> ixL) rBelow rOver = + repairIndex n i rBelow rOver :> repairIndex szL ixL rBelow rOver + {-# INLINE [1] repairIndex #-} + +instance {-# OVERLAPPABLE #-} (4 <= n, + KnownNat n, + Index (Ix (n - 1)), +#if __GLASGOW_HASKELL__ < 800 + Rank (Ix ((n - 1) - 1)) ~ ((n - 1) - 1), +#endif + IxN (n - 1) ~ Ix (n - 1) + ) => Index (IxN n) where + type Rank (IxN n) = n + rank _ = fromInteger $ natVal (Proxy :: Proxy n) + {-# INLINE [1] rank #-} + totalElem (i :> ix) = i * totalElem ix + {-# INLINE [1] totalElem #-} + consDim = (:>) + {-# INLINE [1] consDim #-} + unconsDim (i :> ix) = (i, ix) + {-# INLINE [1] unconsDim #-} + snocDim (i :> ix) k = i :> snocDim ix k + {-# INLINE [1] snocDim #-} + unsnocDim (i :> ix) = case unsnocDim ix of + (jx, j) -> (i :> jx, j) + {-# INLINE [1] unsnocDim #-} + getIndex ix@(j :> jx) k | k == rank ix = Just j + | otherwise = getIndex jx k + {-# INLINE [1] getIndex #-} + setIndex ix@(j :> jx) k o | k == rank ix = Just (o :> jx) + | otherwise = (j :>) <$> setIndex jx k o + {-# INLINE [1] setIndex #-} + dropDim ix@(j :> jx) k | k == rank ix = Just jx + | otherwise = (j :>) <$> dropDim jx k + {-# INLINE [1] dropDim #-} + pureIndex i = i :> (pureIndex i :: Ix (n - 1)) + {-# INLINE [1] pureIndex #-} + liftIndex f (i :> ix) = f i :> liftIndex f ix + {-# INLINE [1] liftIndex #-} + liftIndex2 f (i1 :> ix1) (i2 :> ix2) = f i1 i2 :> liftIndex2 f ix1 ix2 + {-# INLINE [1] liftIndex2 #-} + repairIndex (n :> szL) (i :> ixL) rBelow rOver = + repairIndex n i rBelow rOver :> repairIndex szL ixL rBelow rOver + {-# INLINE [1] repairIndex #-} + + + +---- Unbox Ix + +-- | Unboxing of a `Ix2`. +instance VU.Unbox Ix2 + +newtype instance VU.MVector s Ix2 = MV_Ix2 (VU.MVector s Ix2T) + +instance VM.MVector VU.MVector Ix2 where + basicLength (MV_Ix2 mvec) = VM.basicLength mvec + {-# INLINE basicLength #-} + basicUnsafeSlice idx len (MV_Ix2 mvec) = MV_Ix2 (VM.basicUnsafeSlice idx len mvec) + {-# INLINE basicUnsafeSlice #-} + basicOverlaps (MV_Ix2 mvec) (MV_Ix2 mvec') = VM.basicOverlaps mvec mvec' + {-# INLINE basicOverlaps #-} + basicUnsafeNew len = MV_Ix2 `liftM` VM.basicUnsafeNew len + {-# INLINE basicUnsafeNew #-} + basicUnsafeReplicate len val = MV_Ix2 `liftM` VM.basicUnsafeReplicate len (fromIx2 val) + {-# INLINE basicUnsafeReplicate #-} + basicUnsafeRead (MV_Ix2 mvec) idx = toIx2 `liftM` VM.basicUnsafeRead mvec idx + {-# INLINE basicUnsafeRead #-} + basicUnsafeWrite (MV_Ix2 mvec) idx val = VM.basicUnsafeWrite mvec idx (fromIx2 val) + {-# INLINE basicUnsafeWrite #-} + basicClear (MV_Ix2 mvec) = VM.basicClear mvec + {-# INLINE basicClear #-} + basicSet (MV_Ix2 mvec) val = VM.basicSet mvec (fromIx2 val) + {-# INLINE basicSet #-} + basicUnsafeCopy (MV_Ix2 mvec) (MV_Ix2 mvec') = VM.basicUnsafeCopy mvec mvec' + {-# INLINE basicUnsafeCopy #-} + basicUnsafeMove (MV_Ix2 mvec) (MV_Ix2 mvec') = VM.basicUnsafeMove mvec mvec' + {-# INLINE basicUnsafeMove #-} + basicUnsafeGrow (MV_Ix2 mvec) len = MV_Ix2 `liftM` VM.basicUnsafeGrow mvec len + {-# INLINE basicUnsafeGrow #-} +#if MIN_VERSION_vector(0,11,0) + basicInitialize (MV_Ix2 mvec) = VM.basicInitialize mvec + {-# INLINE basicInitialize #-} +#endif + + +newtype instance VU.Vector Ix2 = V_Ix2 (VU.Vector Ix2T) + +instance V.Vector VU.Vector Ix2 where + basicUnsafeFreeze (MV_Ix2 mvec) = V_Ix2 `liftM` V.basicUnsafeFreeze mvec + {-# INLINE basicUnsafeFreeze #-} + basicUnsafeThaw (V_Ix2 vec) = MV_Ix2 `liftM` V.basicUnsafeThaw vec + {-# INLINE basicUnsafeThaw #-} + basicLength (V_Ix2 vec) = V.basicLength vec + {-# INLINE basicLength #-} + basicUnsafeSlice idx len (V_Ix2 vec) = V_Ix2 (V.basicUnsafeSlice idx len vec) + {-# INLINE basicUnsafeSlice #-} + basicUnsafeIndexM (V_Ix2 vec) idx = toIx2 `liftM` V.basicUnsafeIndexM vec idx + {-# INLINE basicUnsafeIndexM #-} + basicUnsafeCopy (MV_Ix2 mvec) (V_Ix2 vec) = V.basicUnsafeCopy mvec vec + {-# INLINE basicUnsafeCopy #-} + elemseq _ = seq + {-# INLINE elemseq #-} + + + +---- Unbox Ix + + + +-- | Unboxing of a `IxN`. +instance (3 <= n, +#if __GLASGOW_HASKELL__ < 800 + Rank (Ix (n - 1)) ~ (n - 1), +#endif + VU.Unbox (Ix (n-1))) => VU.Unbox (IxN n) + +newtype instance VU.MVector s (IxN n) = MV_IxN (VU.MVector s Int, VU.MVector s (Ix (n-1))) + +instance (3 <= n, +#if __GLASGOW_HASKELL__ < 800 + Rank (Ix (n - 1)) ~ (n - 1), +#endif + VU.Unbox (Ix (n - 1))) => + VM.MVector VU.MVector (IxN n) where + basicLength (MV_IxN (_, mvec)) = VM.basicLength mvec + {-# INLINE basicLength #-} + basicUnsafeSlice idx len (MV_IxN (mvec1, mvec)) = + MV_IxN (VM.basicUnsafeSlice idx len mvec1, VM.basicUnsafeSlice idx len mvec) + {-# INLINE basicUnsafeSlice #-} + basicOverlaps (MV_IxN (mvec1, mvec)) (MV_IxN (mvec1', mvec')) = + VM.basicOverlaps mvec1 mvec1' && VM.basicOverlaps mvec mvec' + {-# INLINE basicOverlaps #-} + basicUnsafeNew len = do + iv <- VM.basicUnsafeNew len + ivs <- VM.basicUnsafeNew len + return $ MV_IxN (iv, ivs) + {-# INLINE basicUnsafeNew #-} + basicUnsafeReplicate len (i :> ix) = do + iv <- VM.basicUnsafeReplicate len i + ivs <- VM.basicUnsafeReplicate len ix + return $ MV_IxN (iv, ivs) + {-# INLINE basicUnsafeReplicate #-} + basicUnsafeRead (MV_IxN (mvec1, mvec)) idx = do + i <- VM.basicUnsafeRead mvec1 idx + ix <- VM.basicUnsafeRead mvec idx + return (i :> ix) + {-# INLINE basicUnsafeRead #-} + basicUnsafeWrite (MV_IxN (mvec1, mvec)) idx (i :> ix) = do + VM.basicUnsafeWrite mvec1 idx i + VM.basicUnsafeWrite mvec idx ix + {-# INLINE basicUnsafeWrite #-} + basicClear (MV_IxN (mvec1, mvec)) = VM.basicClear mvec1 >> VM.basicClear mvec + {-# INLINE basicClear #-} + basicSet (MV_IxN (mvec1, mvec)) (i :> ix) = VM.basicSet mvec1 i >> VM.basicSet mvec ix + {-# INLINE basicSet #-} + basicUnsafeCopy (MV_IxN (mvec1, mvec)) (MV_IxN (mvec1', mvec')) = + VM.basicUnsafeCopy mvec1 mvec1' >> VM.basicUnsafeCopy mvec mvec' + {-# INLINE basicUnsafeCopy #-} + basicUnsafeMove (MV_IxN (mvec1, mvec)) (MV_IxN (mvec1', mvec')) = + VM.basicUnsafeMove mvec1 mvec1' >> VM.basicUnsafeMove mvec mvec' + {-# INLINE basicUnsafeMove #-} + basicUnsafeGrow (MV_IxN (mvec1, mvec)) len = do + iv <- VM.basicUnsafeGrow mvec1 len + ivs <- VM.basicUnsafeGrow mvec len + return $ MV_IxN (iv, ivs) + {-# INLINE basicUnsafeGrow #-} +#if MIN_VERSION_vector(0,11,0) + basicInitialize (MV_IxN (mvec1, mvec)) = + VM.basicInitialize mvec1 >> VM.basicInitialize mvec + {-# INLINE basicInitialize #-} +#endif + + +newtype instance VU.Vector (IxN n) = V_IxN (VU.Vector Int, VU.Vector (Ix (n-1))) + +instance (3 <= n, +#if __GLASGOW_HASKELL__ < 800 + Rank (Ix (n - 1)) ~ (n - 1), +#endif + VU.Unbox (Ix (n-1))) => V.Vector VU.Vector (IxN n) where + basicUnsafeFreeze (MV_IxN (mvec1, mvec)) = do + iv <- V.basicUnsafeFreeze mvec1 + ivs <- V.basicUnsafeFreeze mvec + return $ V_IxN (iv, ivs) + {-# INLINE basicUnsafeFreeze #-} + basicUnsafeThaw (V_IxN (vec1, vec)) = do + imv <- V.basicUnsafeThaw vec1 + imvs <- V.basicUnsafeThaw vec + return $ MV_IxN (imv, imvs) + {-# INLINE basicUnsafeThaw #-} + basicLength (V_IxN (_, vec)) = V.basicLength vec + {-# INLINE basicLength #-} + basicUnsafeSlice idx len (V_IxN (vec1, vec)) = do + V_IxN (V.basicUnsafeSlice idx len vec1, V.basicUnsafeSlice idx len vec) + {-# INLINE basicUnsafeSlice #-} + basicUnsafeIndexM (V_IxN (vec1, vec)) idx = do + i <- V.basicUnsafeIndexM vec1 idx + ix <- V.basicUnsafeIndexM vec idx + return (i :> ix) + {-# INLINE basicUnsafeIndexM #-} + basicUnsafeCopy (MV_IxN (mvec1, mvec)) (V_IxN (vec1, vec)) = + V.basicUnsafeCopy mvec1 vec1 >> V.basicUnsafeCopy mvec vec + {-# INLINE basicUnsafeCopy #-} + elemseq _ = seq + {-# INLINE elemseq #-} + diff --git a/src/Data/Massiv/Core/Iterator.hs b/src/Data/Massiv/Core/Iterator.hs new file mode 100644 --- /dev/null +++ b/src/Data/Massiv/Core/Iterator.hs @@ -0,0 +1,44 @@ +{-# LANGUAGE BangPatterns #-} +-- | +-- Module : Data.Massiv.Core.Iterator +-- Copyright : (c) Alexey Kuleshevich 2018 +-- License : BSD3 +-- Maintainer : Alexey Kuleshevich +-- Stability : experimental +-- Portability : non-portable +-- +module Data.Massiv.Core.Iterator + ( loop + , loopM + , loopM_ + ) where + + +-- | Efficient loop with an accumulator +loop :: Int -> (Int -> Bool) -> (Int -> Int) -> a -> (Int -> a -> a) -> a +loop !init' condition increment !initAcc f = go init' initAcc where + go !step !acc = + case condition step of + False -> acc + True -> go (increment step) (f step acc) +{-# INLINE loop #-} + + +-- | Very efficient monadic loop with an accumulator +loopM :: Monad m => Int -> (Int -> Bool) -> (Int -> Int) -> a -> (Int -> a -> m a) -> m a +loopM !init' condition increment !initAcc f = go init' initAcc where + go !step !acc = + case condition step of + False -> return acc + True -> f step acc >>= go (increment step) +{-# INLINE loopM #-} + + +-- | Efficient monadic loop. Result of each iteration is discarded. +loopM_ :: Monad m => Int -> (Int -> Bool) -> (Int -> Int) -> (Int -> m a) -> m () +loopM_ !init' condition increment f = go init' where + go !step = + case condition step of + False -> return () + True -> f step >> go (increment step) +{-# INLINE loopM_ #-} diff --git a/src/Data/Massiv/Core/List.hs b/src/Data/Massiv/Core/List.hs new file mode 100644 --- /dev/null +++ b/src/Data/Massiv/Core/List.hs @@ -0,0 +1,325 @@ +{-# LANGUAGE BangPatterns #-} +{-# OPTIONS_GHC -fno-warn-orphans #-} +{-# LANGUAGE FlexibleContexts #-} +{-# LANGUAGE FlexibleInstances #-} +{-# LANGUAGE MultiParamTypeClasses #-} +{-# LANGUAGE RankNTypes #-} +{-# LANGUAGE RecordWildCards #-} +{-# LANGUAGE ScopedTypeVariables #-} +{-# LANGUAGE TypeFamilies #-} +{-# LANGUAGE UndecidableInstances #-} +-- | +-- Module : Data.Massiv.Core.List +-- Copyright : (c) Alexey Kuleshevich 2018 +-- License : BSD3 +-- Maintainer : Alexey Kuleshevich +-- Stability : experimental +-- Portability : non-portable +-- +module Data.Massiv.Core.List + ( LN + , L(..) + , Array(..) + , toListArray + , ListItem + , ShapeError(..) + ) where + +import Control.Exception +import Control.Monad (unless) +import Data.Coerce +import Data.Foldable (foldr') +import Data.Functor.Identity +import qualified Data.List as L +import Data.Massiv.Core.Common +import Data.Massiv.Core.Scheduler +import Data.Proxy +import Data.Typeable +import GHC.Exts +import System.IO.Unsafe (unsafePerformIO) + +data LN + +type instance EltRepr LN ix = LN + +type family ListItem ix e :: * where + ListItem Ix1 e = e + ListItem ix e = [ListItem (Lower ix) e] + +type instance NestedStruct LN ix e = [ListItem ix e] + +newtype instance Array LN ix e = List { unList :: [Elt LN ix e] } + + +instance {-# OVERLAPPING #-} Nested LN Ix1 e where + fromNested = coerce + {-# INLINE fromNested #-} + toNested = coerce + {-# INLINE toNested #-} + +instance ( Elt LN ix e ~ Array LN (Lower ix) e + , ListItem ix e ~ [ListItem (Lower ix) e] + , Coercible (Elt LN ix e) (ListItem ix e) + ) => + Nested LN ix e where + fromNested = coerce + {-# INLINE fromNested #-} + toNested = coerce + {-# INLINE toNested #-} + + +instance Nested LN ix e => IsList (Array LN ix e) where + type Item (Array LN ix e) = ListItem ix e + fromList = fromNested + {-# INLINE fromList #-} + toList = toNested + {-# INLINE toList #-} + + +data L = L +type instance EltRepr L ix = L + +type instance NestedStruct L ix e = Array LN ix e + +data instance Array L ix e = LArray { lComp :: Comp + , lData :: !(Array LN ix e) } + + + +data ShapeError = RowTooShortError + | RowTooLongError + deriving Show + +instance Exception ShapeError + + +instance Nested L ix e where + fromNested = LArray Seq + {-# INLINE fromNested #-} + toNested = lData + {-# INLINE toNested #-} + + +instance Nested LN ix e => IsList (Array L ix e) where + type Item (Array L ix e) = ListItem ix e + fromList = LArray Seq . fromNested + {-# INLINE fromList #-} + toList = toNested . lData + {-# INLINE toList #-} + + + +instance {-# OVERLAPPING #-} Ragged L Ix1 e where + isNull = null . unList . lData + {-# INLINE isNull #-} + empty comp = LArray comp (List []) + {-# INLINE empty #-} + edgeSize = length . unList . lData + {-# INLINE edgeSize #-} + cons x arr = arr { lData = coerce (x : coerce (lData arr)) } + {-# INLINE cons #-} + uncons LArray {..} = + case L.uncons $ coerce lData of + Nothing -> Nothing + Just (x, xs) -> Just (x, LArray lComp (coerce xs)) + {-# INLINE uncons #-} + flatten = id + {-# INLINE flatten #-} + unsafeGenerateM !comp !k f = do + xs <- loopM (k - 1) (>= 0) (subtract 1) [] $ \i acc -> do + e <- f i + return (e:acc) + return $ LArray comp $ coerce xs + {-# INLINE unsafeGenerateM #-} + loadRagged using uWrite start end _ xs = + using $ do + leftOver <- + loopM start (< end) (+ 1) xs $ \i xs' -> + case uncons xs' of + Nothing -> throwIO RowTooShortError + Just (y, ys) -> uWrite i y >> return ys + unless (isNull leftOver) $ throwIO RowTooLongError + {-# INLINE loadRagged #-} + raggedFormat f _ arr = L.concat $ "[ " : (L.intersperse "," $ map f (coerce (lData arr))) ++ [" ]"] + + +instance ( Index ix + , Index (Lower ix) + , Ragged L (Lower ix) e + , Elt L ix e ~ Array L (Lower ix) e + , Elt LN ix e ~ Array LN (Lower ix) e + , Coercible (Elt LN ix e) [Elt LN (Lower ix) e] + ) => + Ragged L ix e where + isNull = null . unList . lData + {-# INLINE isNull #-} + empty comp = LArray comp (List []) + {-# INLINE empty #-} + edgeSize arr = + consDim (length (unList (lData arr))) $ + case uncons arr of + Nothing -> zeroIndex + Just (x, _) -> edgeSize x + {-# INLINE edgeSize #-} + cons (LArray _ x) arr = newArr + where + newArr = + arr {lData = coerce (x : coerce (lData arr))} + {-# INLINE cons #-} + uncons LArray {..} = + case L.uncons (coerce lData) of + Nothing -> Nothing + Just (x, xs) -> + let newArr = LArray lComp (coerce xs) + newX = LArray lComp x + in Just (newX, newArr) + {-# INLINE uncons #-} + unsafeGenerateM !comp !sz f = do + let !(k, szL) = unconsDim sz + loopM (k - 1) (>= 0) (subtract 1) (empty comp) $ \i acc -> do + e <- unsafeGenerateM comp szL (\ !ixL -> f (consDim i ixL)) + return (cons e acc) + {-# INLINE unsafeGenerateM #-} + flatten arr = LArray {lComp = lComp arr, lData = coerce xs} + where + xs = + concatMap + (unList . lData . flatten . LArray (lComp arr)) + (unList (lData arr)) + {-# INLINE flatten #-} + loadRagged using uWrite start end sz xs = do + let step = totalElem sz + szL = tailDim sz + leftOver <- + loopM start (< end) (+ step) xs $ \i zs -> + case uncons zs of + Nothing -> throwIO RowTooShortError + Just (y, ys) -> do + _ <- loadRagged using uWrite i (i + step) szL y + return ys + unless (isNull (flatten leftOver)) $ throwIO RowTooLongError + {-# INLINE loadRagged #-} + raggedFormat f sep (LArray comp xs) = + showN + (\s y -> raggedFormat f s (LArray comp y :: Array L (Lower ix) e)) + sep + (coerce xs) + + +instance {-# OVERLAPPING #-} Construct L Ix1 e where + getComp = lComp + {-# INLINE getComp #-} + setComp c arr = arr { lComp = c } + {-# INLINE setComp #-} + unsafeMakeArray Seq sz f = runIdentity $ unsafeGenerateM Seq sz (return . f) + unsafeMakeArray (ParOn wss) sz f = LArray (ParOn wss) $ List $ unsafePerformIO $ do + withScheduler' wss $ \scheduler -> + loopM_ 0 (< sz) (+ 1) (scheduleWork scheduler . return . f) + {-# INLINE unsafeMakeArray #-} + + +instance ( Index ix + , Ragged L ix e + , Ragged L (Lower ix) e + , Elt L ix e ~ Array L (Lower ix) e + ) => + Construct L ix e where + getComp = lComp + {-# INLINE getComp #-} + setComp c arr = arr {lComp = c} + {-# INLINE setComp #-} + unsafeMakeArray comp sz f = unsafeGenerateN comp sz f + {-# INLINE unsafeMakeArray #-} + + +unsafeGenerateN :: + ( Index ix + , Ragged r ix e + , Ragged r (Lower ix) e + , Elt r ix e ~ Array r (Lower ix) e ) + => Comp + -> ix + -> (ix -> e) + -> Array r ix e +unsafeGenerateN Seq sz f = runIdentity $ unsafeGenerateM Seq sz (return . f) +unsafeGenerateN c@(ParOn wss) sz f = unsafePerformIO $ do + let !(m, szL) = unconsDim sz + xs <- withScheduler' wss $ \scheduler -> do + loopM_ 0 (< m) (+ 1) $ \i -> scheduleWork scheduler $ do + unsafeGenerateM c szL $ \ix -> return $ f (consDim i ix) + return $! foldr' cons (empty c) xs +{-# INLINE unsafeGenerateN #-} + + +toListArray :: (Construct L ix e, Source r ix e) + => Array r ix e + -> Array L ix e +toListArray !arr = + unsafeMakeArray (getComp arr) (size arr) (unsafeIndex arr) +{-# INLINE toListArray #-} + + + + +-- -- | Version of foldr that supports foldr/build list fusion implemented by GHC. +-- foldrFB :: (e -> b -> b) -> b -> Int -> (Int -> e) -> b +-- --foldrFB c n k f = loop (k - 1) (>= 0) (subtract 1) n $ \i acc -> f i `c` acc +-- foldrFB c n k f = go 0 +-- where +-- go !i +-- | i == k = n +-- | otherwise = let !v = f i in v `c` go (i + 1) +-- {-# INLINE [0] foldrFB #-} + + +instance {-# OVERLAPPING #-} (Ragged L ix e, Show e) => Show (Array L ix e) where + show arr = " " ++ raggedFormat show "\n " arr + +instance {-# OVERLAPPING #-} (Ragged L ix e, Nested LN ix e, Show e) => + Show (Array LN ix e) where + show arr = show (fromNested arr :: Array L ix e) + + +showN :: (String -> a -> String) -> String -> [a] -> String +showN _ _ [] = "[ ]" +showN fShow lnPrefix ls = + L.concat + (["[ "] ++ + (L.intersperse (lnPrefix ++ ", ") $ map (fShow (lnPrefix ++ " ")) ls) ++ [lnPrefix, "]"]) + +instance ( Ragged L ix e + , Construct L ix e + , Source r ix e + , Show e + ) => + Show (Array r ix e) where + show arr = + "(Array " ++ showsTypeRep (typeRep (Proxy :: Proxy r)) " " ++ + showComp (getComp arr) ++ " (" ++ + (show (size arr)) ++ ")\n" ++ + show (makeArray (getComp arr) (size arr) (evaluateAt arr) :: Array L ix e) ++ ")" + where showComp Seq = "Seq" + showComp Par = "Par" + showComp c = "(" ++ show c ++ ")" + + + + +instance {-# OVERLAPPING #-} OuterSlice L Ix1 e where + unsafeOuterSlice (LArray _ xs) = (coerce xs !!) + {-# INLINE unsafeOuterSlice #-} + outerLength = length . (coerce :: Array LN Ix1 e -> [e]). lData + {-# INLINE outerLength #-} + + +instance Ragged L ix e => OuterSlice L ix e where + unsafeOuterSlice arr' i = go 0 arr' + where + go n arr = + case uncons arr of + Nothing -> errorIx "Data.Massiv.Core.List.unsafeOuterSlice" (outerLength arr') i + Just (x, _) | n == i -> x + Just (_, xs) -> go (n + 1) xs + {-# INLINE unsafeOuterSlice #-} + outerLength = length . (coerce :: Array LN ix e -> [Elt LN ix e]) . lData + {-# INLINE outerLength #-} diff --git a/src/Data/Massiv/Core/Scheduler.hs b/src/Data/Massiv/Core/Scheduler.hs new file mode 100644 --- /dev/null +++ b/src/Data/Massiv/Core/Scheduler.hs @@ -0,0 +1,268 @@ +{-# LANGUAGE BangPatterns #-} +{-# LANGUAGE ExistentialQuantification #-} +{-# LANGUAGE RecordWildCards #-} +{-# LANGUAGE ScopedTypeVariables #-} +-- | +-- Module : Data.Array.Massiv.Scheduler +-- Copyright : (c) Alexey Kuleshevich 2018 +-- License : BSD3 +-- Maintainer : Alexey Kuleshevich +-- Stability : experimental +-- Portability : non-portable +-- +module Data.Massiv.Core.Scheduler + ( Scheduler + , numWorkers + , scheduleWork + , withScheduler + , withScheduler' + , withScheduler_ + , divideWork + , divideWork_ + ) where + +import Control.Concurrent (ThreadId, forkOnWithUnmask, + getNumCapabilities, killThread) +import Control.Concurrent.MVar +import Control.DeepSeq +import Control.Exception (SomeException, catch, mask, + mask_, throwIO, try, + uninterruptibleMask_) +import Control.Monad (forM) +import Control.Monad.Primitive (RealWorld) +import Data.IORef (IORef, atomicModifyIORef', + newIORef, readIORef) +import Data.Massiv.Core.Index.Class (Index (totalElem)) +import Data.Massiv.Core.Iterator (loop) +import Data.Primitive.Array (Array, MutableArray, indexArray, + newArray, unsafeFreezeArray, + writeArray) +import System.IO.Unsafe (unsafePerformIO) +import System.Mem.Weak + +data Job = Job (IO ()) + | Retire + +data Scheduler a = Scheduler + { jobsCountIORef :: !(IORef Int) + , jobQueueMVar :: !(MVar [Job]) + , resultsMVar :: !(MVar (MutableArray RealWorld a)) + , workers :: !Workers + , numWorkers :: {-# UNPACK #-} !Int + } + + +data Workers = Workers { workerThreadIds :: ![ThreadId] + , workerJobDone :: !(MVar (Maybe SomeException)) + , workerJobQueue :: !(MVar [Job]) + } + + +-- | Helper function that allows scheduling work to be done in parallel. Use +-- `withScheduler` to be able to get to a `Scheduler`. +scheduleWork :: Scheduler a -- ^ Scheduler to use + -> IO a -- ^ Action to hand of to a worker + -> IO () +scheduleWork Scheduler {..} jobAction = do + modifyMVar_ jobQueueMVar $ \jobs -> do + jix <- atomicModifyIORef' jobsCountIORef $ \jc -> (jc + 1, jc) + let job = + Job $ do + jobResult <- jobAction + withMVar resultsMVar $ \resArray -> do + writeArray resArray jix jobResult + putMVar (workerJobDone workers) Nothing + return (job : jobs) + + +uninitialized :: a +uninitialized = error "Data.Array.Massiv.Scheduler: uncomputed job result" + + +-- | Execute some action that needs a resource. Perform different cleanup +-- actions depending if thataction resulted in an error or was successful. Sort +-- of like `bracket` and `bracketOnError` with info about exception combined. +bracketWithException :: forall a b c d . + IO a -- ^ Acquire resource + -> (a -> IO b) -- ^ Run after successfull execution + -> (SomeException -> a -> IO c) -- ^ Run if execution resulted in exception. + -> (a -> IO d) -- ^ Execute an action that actually needs that resource. + -> IO d +bracketWithException before afterSuccess afterError thing = mask $ \restore -> do + x <- before + eRes <- try $ restore (thing x) + case eRes of + Left (exc :: SomeException) -> do + _ :: Either SomeException c <- try $ uninterruptibleMask_ $ afterError exc x + throwIO exc + Right y -> do + _ <- uninterruptibleMask_ $ afterSuccess x + return y + +-- | Run arbitrary computations in parallel. A pool of workers is initialized, +-- unless Worker Stations list is empty and a global worker pool is currently +-- available. All of those workers will be stealing work that you can schedule +-- using `scheduleWork`. The order in which work is scheduled will be the same +-- as the order of the resuts of those computations, stored withing the +-- resulting array. Size of the array, which is also the first element in the +-- returned tuple, will match the number of times `scheduleWork` has been +-- invoked. This function blocks until all of the submitted jobs has finished or +-- one of them resulted in an exception, which will be re-thrown here. +-- +-- __Important__: In order to get work done truly in parallel, program needs to be +-- compiled with @-threaded@ GHC flag and executed with @+RTS -N@. +-- +withScheduler :: [Int] -- ^ Worker Stations, i.e. capabilities. Empty list will + -- result in utilization of all available capabilities. + -> (Scheduler a -> IO b) -- ^ Action that will be scheduling all + -- the work. + -> IO (Int, Array a) +withScheduler wss submitJobs = do + jobsCountIORef <- newIORef 0 + jobQueueMVar <- newMVar [] + resultsMVar <- newEmptyMVar + bracketWithException + (do mWeakWorkers <- + if null wss + then tryTakeMVar globalWorkersMVar + else return Nothing + mGlobalWorkers <- maybe (return Nothing) deRefWeak mWeakWorkers + let toWorkers w = return (mWeakWorkers, w) + maybe (hireWorkers wss >>= toWorkers) toWorkers mGlobalWorkers) + (\(mWeakWorkers, workers) -> do + case mWeakWorkers of + Nothing -> + putMVar (workerJobQueue workers) $ + replicate (length (workerThreadIds workers)) Retire + Just weak -> putMVar globalWorkersMVar weak) + (\_ (mWeakWorkers, workers) -> do + case mWeakWorkers of + Nothing -> mapM_ killThread (workerThreadIds workers) + Just weakWorkers -> do + finalize weakWorkers + newWeakWorkers <- hireWeakWorkers globalWorkersMVar + putMVar globalWorkersMVar newWeakWorkers) + (\(_, workers) -> do + let scheduler = + Scheduler {numWorkers = length $ workerThreadIds workers, ..} + _ <- submitJobs scheduler + jobCount <- readIORef jobsCountIORef + marr <- newArray jobCount uninitialized + putMVar resultsMVar marr + jobQueue <- takeMVar jobQueueMVar + putMVar (workerJobQueue workers) $ reverse jobQueue + waitTillDone scheduler + arr <- unsafeFreezeArray marr + return (jobCount, arr)) + + +-- | Just like `withScheduler`, but returns computed results in a list, instead +-- of an array. +withScheduler' :: [Int] -> (Scheduler a -> IO b) -> IO [a] +withScheduler' wss submitJobs = do + (jc, arr) <- withScheduler wss submitJobs + return $ + loop (jc - 1) (>= 0) (subtract 1) [] $ \i acc -> indexArray arr i : acc + + +-- | Just like `withScheduler`, but discards the results. +withScheduler_ :: [Int] -> (Scheduler a -> IO b) -> IO () +withScheduler_ wss submitJobs = withScheduler wss submitJobs >> return () + + +-- | Same as `divideWork`, but discard the result. +divideWork_ :: Index ix + => [Int] -> ix -> (Scheduler a -> Int -> Int -> Int -> IO b) -> IO () +divideWork_ wss sz submit = divideWork wss sz submit >> return () + + +-- | Linearly (row-major first) and equally divide work among available +-- workers. Submit function will receive a `Scheduler`, length of each chunk, +-- total number of elements, as well as where chunks end and slack begins. Slack +-- work will get picked up by the first worker, that has finished working on his +-- chunk. Returns list with results in the same order that work was submitted +divideWork :: Index ix + => [Int] -- ^ Worker Stations (capabilities) + -> ix -- ^ Size + -> (Scheduler a -> Int -> Int -> Int -> IO b) -- ^ Submit function + -> IO [a] +divideWork wss sz submit + | totalElem sz == 0 = return [] + | otherwise = do + withScheduler' wss $ \scheduler -> do + let !totalLength = totalElem sz + !chunkLength = totalLength `quot` numWorkers scheduler + !slackStart = chunkLength * numWorkers scheduler + submit scheduler chunkLength totalLength slackStart + +-- | Wait till workers finished with all submitted jobs, but raise an exception +-- if either of them has died. Raised exception is the same one that was the +-- cause of worker's death. +waitTillDone :: Scheduler a -> IO () +waitTillDone (Scheduler {..}) = readIORef jobsCountIORef >>= waitTill 0 + where + waitTill jobsDone jobsCount + | jobsDone == jobsCount = return () + | otherwise = do + mExc <- takeMVar (workerJobDone workers) + case mExc of + Just exc -> throwIO exc + Nothing -> waitTill (jobsDone + 1) jobsCount + + +-- | Worker can either be doing work, waiting for a job, or going into +-- retirement. Temp workers are rarely in waiting state, unless there is simply +-- not enough work for all workers in the pool. Unlike temp workers, global +-- workers do spend quite a bit of time waiting for work and they are never +-- retired, but ruthlessly killed. +runWorker :: MVar [Job] -> IO () +runWorker jobsMVar = do + jobs <- takeMVar jobsMVar + case jobs of + (Job job:rest) -> putMVar jobsMVar rest >> job >> runWorker jobsMVar + (Retire:rest) -> putMVar jobsMVar rest + [] -> runWorker jobsMVar + + +-- | Used whenever a pool of new workers is needed. If list is empty all +-- capabilities are utilized, otherwise each element in the list will be an +-- argument to `forkOn`. +hireWorkers :: [Int] -> IO Workers +hireWorkers wss = do + wss' <- + if null wss + then do + wNum <- getNumCapabilities + return [0 .. wNum - 1] + else return wss + workerJobQueue <- newEmptyMVar + workerJobDone <- newEmptyMVar + workerThreadIds <- + forM wss' $ \ws -> + mask_ $ + forkOnWithUnmask ws $ \unmask -> do + catch + (unmask $ runWorker workerJobQueue) + (unmask . putMVar workerJobDone . Just) + workerThreadIds `deepseq` return Workers {..} + +-- | Global workers are the most utilized ones, therefore they are rarily +-- restarted, in particular, only in case when one of them dies of an +-- exception. Weak reference is used so workers don't continue running after +-- MVar has been cleaned up by the GC. Each global worker has his own station, +-- i.e. global workers always span all available capabilities. +globalWorkersMVar :: MVar (Weak Workers) +globalWorkersMVar = unsafePerformIO $ do + workersMVar <- newEmptyMVar + weakWorkers <- hireWeakWorkers workersMVar + putMVar workersMVar weakWorkers + return workersMVar +{-# NOINLINE globalWorkersMVar #-} + + +-- | Hire workers under weak pointers. Finilizer will kill all the +-- workers. These will be used as global workers +hireWeakWorkers :: key -> IO (Weak Workers) +hireWeakWorkers k = do + workers <- hireWorkers [] + mkWeak k workers (Just (mapM_ killThread (workerThreadIds workers))) diff --git a/tests/Data/Massiv/Array/DelayedSpec.hs b/tests/Data/Massiv/Array/DelayedSpec.hs new file mode 100644 --- /dev/null +++ b/tests/Data/Massiv/Array/DelayedSpec.hs @@ -0,0 +1,22 @@ +{-# LANGUAGE FlexibleContexts #-} +{-# LANGUAGE FlexibleInstances #-} +{-# LANGUAGE MultiParamTypeClasses #-} +module Data.Massiv.Array.DelayedSpec (spec) where + +import Test.Hspec +--import Test.QuickCheck + + + + +spec :: Spec +spec = return () + -- describe "DIM1" $ do + -- specShapeN (Nothing :: Maybe (D, DIM1, Int)) + -- describe "DIM2" $ do + -- specShapeN (Nothing :: Maybe (D, DIM2, Int)) + -- specSliceN (Nothing :: Maybe (D, DIM2, Int)) + -- specSliceDim2 + -- describe "DIM3" $ do + -- specShapeN (Nothing :: Maybe (D, DIM3, Int)) + -- specSliceN (Nothing :: Maybe (D, DIM3, Int)) diff --git a/tests/Data/Massiv/Array/Manifest/VectorSpec.hs b/tests/Data/Massiv/Array/Manifest/VectorSpec.hs new file mode 100644 --- /dev/null +++ b/tests/Data/Massiv/Array/Manifest/VectorSpec.hs @@ -0,0 +1,87 @@ +{-# LANGUAGE FlexibleContexts #-} +{-# LANGUAGE FlexibleInstances #-} +{-# LANGUAGE MultiParamTypeClasses #-} +{-# LANGUAGE ScopedTypeVariables #-} +{-# LANGUAGE TypeFamilies #-} +module Data.Massiv.Array.Manifest.VectorSpec (spec) where + +import Data.Massiv.CoreArbitrary +import Data.Massiv.Array.Manifest.Vector +import Data.Proxy +import Data.Typeable +import qualified Data.Vector as VB +import qualified Data.Vector.Generic as VG +import qualified Data.Vector.Primitive as VP +import qualified Data.Vector.Storable as VS +import qualified Data.Vector.Unboxed as VU +import Test.Hspec +import Test.QuickCheck + +prop_castToFromVector + :: ( VG.Vector (VRepr r) Int + , Mutable r ix Int + , Typeable (VRepr r) + , ARepr (VRepr r) ~ r + , Eq (Array r ix Int) + , Show (Array r ix Int) + ) + => proxy ix -> r -> Arr r ix Int -> Property +prop_castToFromVector _ _ (Arr arr) = + Just arr === (castToVector arr >>= castFromVector (getComp arr) (size arr)) + + +prop_toFromVector :: + forall r ix v. + ( Mutable r ix Int + , Mutable (ARepr v) ix Int + , VRepr (ARepr v) ~ v + , Eq (Array r ix Int) + , VG.Vector v Int + , Show (Array r ix Int) + , Typeable v + ) + => Proxy v + -> Proxy ix + -> r + -> Arr r ix Int + -> Property +prop_toFromVector _ _ _ (Arr arr) = + arr === fromVector (getComp arr) (size arr) (toVector arr :: v Int) + +toFromVectorSpec :: Spec +toFromVectorSpec = do + let it_prop name r = describe name $ do + describe "Through Boxed Vector" $ do + it "Ix1" $ property $ prop_toFromVector (Proxy :: Proxy VB.Vector) (Proxy :: Proxy Ix1) r + it "Ix2" $ property $ prop_toFromVector (Proxy :: Proxy VB.Vector) (Proxy :: Proxy Ix2) r + describe "Through Unboxed Vector" $ do + it "Ix1" $ property $ prop_toFromVector (Proxy :: Proxy VU.Vector) (Proxy :: Proxy Ix1) r + it "Ix2" $ property $ prop_toFromVector (Proxy :: Proxy VU.Vector) (Proxy :: Proxy Ix2) r + describe "Through Primitive Vector" $ do + it "Ix1" $ property $ prop_toFromVector (Proxy :: Proxy VP.Vector) (Proxy :: Proxy Ix1) r + it "Ix2" $ property $ prop_toFromVector (Proxy :: Proxy VP.Vector) (Proxy :: Proxy Ix2) r + describe "Through Storable Vector" $ do + it "Ix1" $ property $ prop_toFromVector (Proxy :: Proxy VS.Vector) (Proxy :: Proxy Ix1) r + it "Ix2" $ property $ prop_toFromVector (Proxy :: Proxy VS.Vector) (Proxy :: Proxy Ix2) r + it_prop "Unboxed" U + it_prop "Primitive" P + it_prop "Storable" S + it_prop "BoxedStrict" B + + +castToFromVectorSpec :: Spec +castToFromVectorSpec = do + let it_prop name r = describe name $ do + it "Ix1" $ property $ prop_castToFromVector (Proxy :: Proxy Ix1) r + it "Ix2" $ property $ prop_castToFromVector (Proxy :: Proxy Ix2) r + it "Ix3" $ property $ prop_castToFromVector (Proxy :: Proxy Ix3) r + it_prop "Unboxed" U + it_prop "Primitive" P + it_prop "Storable" S + it_prop "BoxedStrict" B + + +spec :: Spec +spec = do + describe "toFromVector" toFromVectorSpec + describe "castToFromVector" castToFromVectorSpec diff --git a/tests/Data/Massiv/Array/Ops/ConstructSpec.hs b/tests/Data/Massiv/Array/Ops/ConstructSpec.hs new file mode 100644 --- /dev/null +++ b/tests/Data/Massiv/Array/Ops/ConstructSpec.hs @@ -0,0 +1,96 @@ +{-# LANGUAGE FlexibleContexts #-} +{-# LANGUAGE ScopedTypeVariables #-} +{-# LANGUAGE GADTs #-} +module Data.Massiv.Array.Ops.ConstructSpec (spec) where + +import Data.Massiv.CoreArbitrary as A +import Data.Proxy +import qualified GHC.Exts as GHC (IsList (..)) +import Prelude as P +import Prelude hiding (map) +import Test.Hspec +import Test.QuickCheck + + +prop_rangeEqRangeStep1 :: Int -> Int -> Property +prop_rangeEqRangeStep1 from to = range Seq from to === rangeStep Par from 1 to + +prop_rangeEqEnumFromN :: Int -> Int -> Property +prop_rangeEqEnumFromN from to = range Seq from to === enumFromN Par from (to - from) + +prop_rangeStepEqEnumFromStepN :: Int -> NonZero Int -> Int -> Property +prop_rangeStepEqEnumFromStepN from (NonZero step) sz = + rangeStep Seq from step (from + step * sz) === enumFromStepN Par from step sz + + +prop_rangeStepExc :: Int -> Int -> Property +prop_rangeStepExc from to = + assertSomeException (computeAs U (rangeStep Seq from 0 to)) + +prop_toFromListIsList :: + (Show (Array U ix Int), GHC.IsList (Array U ix Int), Index ix) + => Proxy ix + -> Arr U ix Int + -> Property +prop_toFromListIsList _ (Arr arr) = arr === GHC.fromList (GHC.toList arr) + + +prop_toFromList :: + forall ix . (Show (Array U ix Int), Nested LN ix Int, Nested L ix Int, Ragged L ix Int) + => Proxy ix + -> Arr U ix Int + -> Property +prop_toFromList _ (Arr arr) = arr === fromLists' (getComp arr) (toLists arr :: [ListItem ix Int]) + + +prop_excFromToListIx2 :: Comp -> [[Int]] -> Property +prop_excFromToListIx2 comp ls2 = + if P.null lsL || P.all (head lsL ==) lsL + then label "Expected Success" $ resultLs === ls2 + else label "Expected Failure" $ assertSomeException resultLs + where + lsL = P.map P.length ls2 + resultLs = toLists (fromLists' comp ls2 :: Array U Ix2 Int) + + +-- prop_excFromToListIx3 :: Comp -> [[[Int]]] -> Property +-- prop_excFromToListIx3 comp ls3 = +-- if P.null lsL || +-- (P.all (head lsL ==) lsL && +-- (P.null (head lsLL) || P.and (P.map (P.all (head (head lsLL) ==)) lsLL))) +-- then classify True "Expected Success" $ resultLs === ls3 +-- else classify True "Expected Failure" $ +-- assertSomeException resultLs +-- where +-- resultLs = toList (fromList' comp ls3 :: Array U Ix3 Int) +-- lsL = P.map P.length ls3 +-- lsLL = P.map (P.map P.length) ls3 + + +specIx1 :: Spec +specIx1 = do + it "toFromList" $ property (prop_toFromList (Proxy :: Proxy Ix1)) + it "toFromListIsList" $ property (prop_toFromListIsList (Proxy :: Proxy Ix1)) + it "rangeEqRangeStep1" $ property prop_rangeEqRangeStep1 + it "rangeEqEnumFromN" $ property prop_rangeEqEnumFromN + it "rangeStepEqEnumFromStepN" $ property prop_rangeStepEqEnumFromStepN + it "rangeStepExc" $ property prop_rangeStepExc + +specIx2 :: Spec +specIx2 = do + it "toFromList" $ property (prop_toFromList (Proxy :: Proxy Ix2)) + it "toFromListIsList" $ property (prop_toFromListIsList (Proxy :: Proxy Ix2)) + it "excFromToListIx2" $ property prop_excFromToListIx2 + +specIx3 :: Spec +specIx3 = do + it "toFromList" $ property (prop_toFromList (Proxy :: Proxy Ix3)) + it "toFromListIsList" $ property (prop_toFromListIsList (Proxy :: Proxy Ix3)) + --it "excFromToListIx3" $ property prop_excFromToListIx3 + + +spec :: Spec +spec = do + describe "Ix1" specIx1 + describe "Ix2" specIx2 + describe "Ix3" specIx3 diff --git a/tests/Data/Massiv/Array/Ops/FoldSpec.hs b/tests/Data/Massiv/Array/Ops/FoldSpec.hs new file mode 100644 --- /dev/null +++ b/tests/Data/Massiv/Array/Ops/FoldSpec.hs @@ -0,0 +1,64 @@ +{-# LANGUAGE FlexibleContexts #-} +{-# LANGUAGE FlexibleInstances #-} +{-# LANGUAGE MonoLocalBinds #-} +{-# LANGUAGE MultiParamTypeClasses #-} +module Data.Massiv.Array.Ops.FoldSpec (spec) where + +import Data.Massiv.CoreArbitrary +import Prelude hiding (map, product, sum) +import qualified Prelude as P (length, sum) +import Test.Hspec +import Test.QuickCheck +import Test.QuickCheck.Monadic + + +prop_SumSEqSumP :: Index ix => proxy ix -> Array D ix Int -> Bool +prop_SumSEqSumP _ arr = sum arr == sum (setComp Par arr) + + +prop_ProdSEqProdP :: Index ix => proxy ix -> Array D ix Int -> Bool +prop_ProdSEqProdP _ arr = product arr == product (setComp Par arr) + +prop_NestedFoldP :: Array D Ix1 (Array D Ix1 Int) -> Bool +prop_NestedFoldP arr = sum (setComp Par (map sum $ setComp Par arr)) == sum (map sum arr) + +prop_FoldrOnP :: Int -> [Int] -> ArrP D Ix1 Int -> Property +prop_FoldrOnP wId wIds (ArrP arr) = + P.length arr > P.length wIds ==> monadicIO $ do + res <- run $ ifoldrOnP wIdsNE (\_ -> (+)) 0 (:) [] arr + if P.length arr `mod` P.length wIdsNE == 0 + then assert (P.length res == P.length wIdsNE) + else assert (P.length res == P.length wIdsNE + 1) + assert (P.sum res == sum arr) + where + wIdsNE = wId : wIds + +prop_FoldlOnP :: Int -> [Int] -> ArrP D Ix1 Int -> Property +prop_FoldlOnP wId wIds (ArrP arr) = + P.length arr > P.length wIds ==> monadicIO $ do + res <- run $ ifoldlOnP wIdsNE (\a _ x -> a + x) 0 (flip (:)) [] arr + if P.length arr `mod` P.length wIdsNE == 0 + then assert (P.length res == P.length wIdsNE) + else assert (P.length res == P.length wIdsNE + 1) + assert (P.sum res == sum arr) + where + wIdsNE = wId : wIds + + +specFold :: + (Arbitrary ix, CoArbitrary ix, Index ix, Show (Array D ix Int)) + => proxy ix + -> String + -> Spec +specFold proxy dimStr = do + describe dimStr $ do + it "sumS Eq sumP" $ property $ prop_SumSEqSumP proxy + it "prodS Eq prodP" $ property $ prop_ProdSEqProdP proxy + +spec :: Spec +spec = do + specFold (Nothing :: Maybe Ix1) "Ix1" + specFold (Nothing :: Maybe Ix2) "Ix2" + it "Nested Parallel Fold" $ property prop_NestedFoldP + it "FoldrOnP" $ property $ prop_FoldrOnP + it "FoldlOnP" $ property $ prop_FoldlOnP diff --git a/tests/Data/Massiv/Array/Ops/SliceSpec.hs b/tests/Data/Massiv/Array/Ops/SliceSpec.hs new file mode 100644 --- /dev/null +++ b/tests/Data/Massiv/Array/Ops/SliceSpec.hs @@ -0,0 +1,220 @@ +{-# LANGUAGE FlexibleContexts #-} +{-# LANGUAGE FlexibleInstances #-} +{-# LANGUAGE MonoLocalBinds #-} +{-# LANGUAGE MultiParamTypeClasses #-} +{-# LANGUAGE TypeFamilies #-} +module Data.Massiv.Array.Ops.SliceSpec (spec) where + +import Data.Massiv.Array.Unsafe +import Data.Massiv.CoreArbitrary +import Test.Hspec +import Test.QuickCheck + +----------- +-- Size -- +----------- + +-- extract + +prop_ExtractEqualsExtractFromTo + :: (Eq (Array (EltRepr r ix) ix e), Arbitrary (Array r ix e), Size r ix e) + => proxy (r, ix, e) -> SzIx ix -> Array r ix e -> Bool +prop_ExtractEqualsExtractFromTo _ (SzIx (Sz eIx) sIx) arr = + extractFromTo sIx eIx arr == extract sIx (liftIndex2 (-) eIx sIx) arr + + + + +specSizeN + :: (Eq (Array (EltRepr r ix) ix e), Arbitrary (Array r ix e), Show (Array r ix e), Arbitrary ix, Size r ix e) + => proxy (r, ix, e) -> Spec +specSizeN proxy = do + describe "extract" $ do + it "ExtractEqualsExtractFromTo" $ property $ prop_ExtractEqualsExtractFromTo proxy + + +----------- +-- Slice -- +----------- + + +prop_SliceRight :: (Arbitrary (Array r ix e), Slice r ix e, OuterSlice r ix e, Eq (Elt r ix e)) + => proxy (r, ix, e) -> Int -> Array r ix e -> Bool +prop_SliceRight _ i arr = (arr !?> i) == (arr (rank (size arr), i)) + + +prop_SliceLeft :: (Arbitrary (Array r ix e), Slice r ix e, InnerSlice r ix e, Eq (Elt r ix e)) + => proxy (r, ix, e) -> Int -> Array r ix e -> Bool +prop_SliceLeft _ i arr = (arr (1, i)) + + +prop_SliceIndexDim2D :: ArrIx D Ix2 Int -> Bool +prop_SliceIndexDim2D (ArrIx arr ix@(i :. j)) = + val == evaluateAt (arr i) j + where + val = unsafeIndex arr ix + + +prop_SliceIndexDim2RankD :: ArrIx D Ix2 Int -> Bool +prop_SliceIndexDim2RankD (ArrIx arr ix@(i :. j)) = + val == evaluateAt (arr (2, i)) j && + val == evaluateAt (arr (1, j)) i + where + val = unsafeIndex arr ix + + +prop_SliceIndexDim3D :: ArrIx D Ix3 Int -> Bool +prop_SliceIndexDim3D (ArrIx arr ix@(i :> j :. k)) = + val == evaluateAt (arr i !> j) k && + val == evaluateAt (arr i) j && + val == evaluateAt (arr !> i Bool +prop_SliceIndexDim3RankD (ArrIx arr ix@(i :> j :. k)) = + val == evaluateAt (arr (3, i) (2, j)) k && + val == evaluateAt (arr (3, i) (1, k)) j && + val == evaluateAt (arr (2, j) (2, i)) k && + val == evaluateAt (arr (2, j) (1, k)) i && + val == evaluateAt (arr (1, k) (2, i)) j && + val == evaluateAt (arr (1, k) (1, j)) i + where + val = unsafeIndex arr ix + + +prop_SliceIndexDim2M :: ArrIx M Ix2 Int -> Bool +prop_SliceIndexDim2M (ArrIx arr ix@(i :. j)) = + val == (arr !> i ! j) && + val == (arr Bool +prop_SliceIndexDim2RankM (ArrIx arr ix@(i :. j)) = + val == (arr (2, i) ! j) && + val == (arr (1, j) ! i) + where + val = unsafeIndex arr ix + + +prop_SliceIndexDim3M :: ArrIx M Ix3 Int -> Bool +prop_SliceIndexDim3M (ArrIx arr ix@(i :> j :. k)) = + val == (arr i !> j ! k) && + val == (arr i ! j) && + val == (arr !> i Bool +prop_SliceIndexDim3RankM (ArrIx arr ix@(i :> j :. k)) = + val == (arr (3, i) (2, j) ! k) && + val == (arr (3, i) (1, k) ! j) && + val == (arr (2, j) (2, i) ! k) && + val == (arr (2, j) (1, k) ! i) && + val == (arr (1, k) (2, i) ! j) && + val == (arr (1, k) (1, j) ! i) + where + val = unsafeIndex arr ix + + +prop_SliceIndexDim4D :: ArrIx D Ix4 Int -> Bool +prop_SliceIndexDim4D (ArrIx arr ix@(i1 :> i2 :> i3 :. i4)) = + val == evaluateAt (arr !> i1 !> i2 !> i3) i4 && + val == evaluateAt (arr !> i1 !> i2 i1 i1 i2) i3 && + val == evaluateAt (arr i1 !> i2) i3 && + val == evaluateAt (arr i1 i1) i2 + where + val = unsafeIndex arr ix + +prop_SliceIndexDim4RankD :: ArrIx D Ix4 Int -> Bool +prop_SliceIndexDim4RankD (ArrIx arr ix@(i1 :> i2 :> i3 :. i4)) = + val == unsafeIndex (arr (4, i1) (3, i2) (2, i3)) i4 && + val == unsafeIndex (arr (4, i1) (2, i3) (3, i2) (3, i1)) (i3 :. i4) && + val == unsafeIndex (arr (2, i3) (2, i2)) (i1 :. i4) && + val == unsafeIndex (arr (2, i3) (1, i4) !> i1) i2 && + val == unsafeIndex (arr (1, i4) !> i1 !> i2) i3 + where + val = evaluateAt arr ix + + +prop_SliceIndexDim4RankM :: ArrIx M Ix4 Int -> Bool +prop_SliceIndexDim4RankM (ArrIx arr ix@(i1 :> i2 :> i3 :. i4)) = + val == (arr (4, i1) (3, i2) (2, i3) ! i4) && + val == (arr (4, i1) (2, i3) (3, i2) (3, i1) ! (i3 :. i4)) && + val == (arr (2, i3) (2, i2) ! (i1 :. i4)) && + val == (arr (2, i3) (1, i4) !> i1 ! i2) && + val == (arr (1, i4) !> i1 !> i2 ! i3) + where + val = unsafeIndex arr ix + + +prop_SliceIndexDim4M :: ArrIx M Ix4 Int -> Bool +prop_SliceIndexDim4M (ArrIx arr ix@(i1 :> i2 :> i3 :. i4)) = + val == (arr !> i1 !> i2 !> i3 ! i4) && + val == (arr !> i1 !> i2 i1 i1 i2 ! i3) && + val == (arr i1 !> i2 ! i3) && + val == (arr i1 i1 ! i2) + where + val = unsafeIndex arr ix + + + +specSliceN :: ( Arbitrary (Array r ix e) + , Show (Array r ix e) + , Arbitrary ix + , Slice r ix e + , OuterSlice r ix e + , InnerSlice r ix e + , Eq (Elt r ix e) + ) + => proxy (r, ix, e) -> Spec +specSliceN proxy = do + describe "Slice" $ do + it "SliceRight" $ property $ prop_SliceRight proxy + it "SliceLeft" $ property $ prop_SliceLeft proxy + + + +spec :: Spec +spec = do + describe "Ix1" $ do + specSizeN (Nothing :: Maybe (D, Ix1, Int)) + describe "Ix2" $ do + specSizeN (Nothing :: Maybe (D, Ix2, Int)) + specSliceN (Nothing :: Maybe (D, Ix2, Int)) + describe "SliceIndex" $ do + it "Delayed" $ property $ prop_SliceIndexDim2D + it "Rank - Delayed" $ property $ prop_SliceIndexDim2RankD + it "Manifest" $ property $ prop_SliceIndexDim2M + it "Rank - Manifest" $ property $ prop_SliceIndexDim2RankM + describe "Ix3" $ do + specSizeN (Nothing :: Maybe (D, Ix3, Int)) + specSliceN (Nothing :: Maybe (D, Ix3, Int)) + describe "SliceIndex" $ do + it "Delayed" $ property $ prop_SliceIndexDim3D + it "Rank - Delayed" $ property $ prop_SliceIndexDim3RankD + it "Manifest" $ property $ prop_SliceIndexDim3M + it "Rank - Manifest" $ property $ prop_SliceIndexDim3RankM + describe "Ix4" $ do + specSizeN (Nothing :: Maybe (D, Ix4, Int)) + specSliceN (Nothing :: Maybe (D, Ix4, Int)) + describe "SliceIndex" $ do + it "Delayed" $ property $ prop_SliceIndexDim4D + it "Rank - Delayed" $ property $ prop_SliceIndexDim4RankD + it "Manifest" $ property $ prop_SliceIndexDim4M + it "Rank - Manifest" $ property $ prop_SliceIndexDim4RankM diff --git a/tests/Data/Massiv/Array/Ops/TransformSpec.hs b/tests/Data/Massiv/Array/Ops/TransformSpec.hs new file mode 100644 --- /dev/null +++ b/tests/Data/Massiv/Array/Ops/TransformSpec.hs @@ -0,0 +1,52 @@ +{-# LANGUAGE FlexibleContexts #-} +{-# LANGUAGE MonoLocalBinds #-} +{-# LANGUAGE MultiParamTypeClasses #-} +module Data.Massiv.Array.Ops.TransformSpec (spec) where + +import Data.Massiv.CoreArbitrary as A +import Data.Maybe (fromJust) +import Data.Typeable (Typeable) +import Test.Hspec +import Test.QuickCheck + + +prop_ExtractAppend + :: (Eq e, Size r ix e, Source r ix e, Source (EltRepr r ix) ix e, Arbitrary (ArrIx r ix e)) + => proxy (r, ix, e) -> DimIx ix -> ArrIx r ix e -> Bool +prop_ExtractAppend _ (DimIx dim) (ArrIx arr ix) = + maybe False ((delay arr ==) . uncurry (append' dim)) $ + A.splitAt dim (fromJust (getIndex ix dim)) arr + + +prop_transposeOuterInner :: Arr D Ix2 Int -> Property +prop_transposeOuterInner (Arr arr) = transposeOuter arr === transpose arr + + +specN :: + ( Eq e + , Size r ix e + , Source r ix e + , Source (EltRepr r ix) ix e + , Typeable e + , Show (Array r ix e) + , Arbitrary (ArrIx r ix e) + ) + => proxy (r, ix, e) + -> Spec +specN r = do + it "ExtractAppend" $ property $ prop_ExtractAppend r + + +spec :: Spec +spec = do + it "transposeOuterInner" $ property prop_transposeOuterInner + describe "Delayed" $ do + describe "Ix1" $ specN (Nothing :: Maybe (D, Ix1, Int)) + describe "Ix2" $ specN (Nothing :: Maybe (D, Ix2, Int)) + describe "Ix3" $ specN (Nothing :: Maybe (D, Ix3, Int)) + describe "Ix4" $ specN (Nothing :: Maybe (D, Ix4, Int)) + describe "Unboxed" $ do + describe "Ix1" $ specN (Nothing :: Maybe (U, Ix1, Int)) + describe "Ix2" $ specN (Nothing :: Maybe (U, Ix2, Int)) + describe "Ix3" $ specN (Nothing :: Maybe (U, Ix3, Int)) + describe "Ix4" $ specN (Nothing :: Maybe (U, Ix4, Int)) diff --git a/tests/Data/Massiv/Array/StencilSpec.hs b/tests/Data/Massiv/Array/StencilSpec.hs new file mode 100644 --- /dev/null +++ b/tests/Data/Massiv/Array/StencilSpec.hs @@ -0,0 +1,75 @@ +{-# LANGUAGE FlexibleContexts #-} +{-# LANGUAGE FlexibleInstances #-} +{-# LANGUAGE MonoLocalBinds #-} +{-# LANGUAGE MultiParamTypeClasses #-} +module Data.Massiv.Array.StencilSpec (spec) where + +import Control.DeepSeq (deepseq) +import Data.Massiv.Array.Stencil +import Data.Massiv.CoreArbitrary as A +import Data.Maybe (fromJust) +import Data.Proxy +import Test.Hspec +import Test.QuickCheck +import Test.QuickCheck.Function +import Data.Default () +-- sum3x3Stencil :: (Default a, Fractional a) => Border a -> Stencil Ix2 a a +-- sum3x3Stencil b = mkConvolutionStencil b (3 :. 3) (1 :. 1) $ \ get -> +-- get (-1 :. -1) 1 . get (-1 :. 0) 1 . get (-1 :. 1) 1 . +-- get ( 0 :. -1) 1 . get ( 0 :. 0) 1 . get ( 0 :. 1) 1 . +-- get ( 1 :. -1) 1 . get ( 1 :. 0) 1 . get ( 1 :. 1) 1 +-- {-# INLINE sum3x3Stencil #-} + + +singletonStencil :: (Num ix, Index ix) => (Int -> Int) -> Border Int -> Stencil ix Int Int +singletonStencil f b = makeStencil b 1 0 $ \ get -> fmap f (get zeroIndex) +{-# INLINE singletonStencil #-} + + +prop_MapSingletonStencil :: (Load DW ix Int, Manifest U ix Int, Num ix) => + Proxy ix -> Fun Int Int -> Border Int -> ArrP U ix Int -> Bool +prop_MapSingletonStencil _ f b (ArrP arr) = + computeAs U (mapStencil (singletonStencil (apply f) b) arr) == computeAs U (A.map (apply f) arr) + +-- Tests out of bounds stencil indexing +prop_DangerousStencil :: + Index ix => Proxy ix -> NonZero Int -> DimIx ix -> Border Int -> SzIx ix -> Property +prop_DangerousStencil _ (NonZero s) (DimIx r) b (SzIx (Sz sz) ix) = + ix' `deepseq` assertSomeException $ makeStencil b sz ix $ \get -> get ix' + where + ix' = + liftIndex (* signum s) $ + fromJust $ do + i <- getIndex sz r + setIndex zeroIndex r i + + +stencilSpec :: Spec +stencilSpec = do + describe "MapSingletonStencil" $ do + it "Ix1" $ property $ prop_MapSingletonStencil (Proxy :: Proxy Ix1) + it "Ix2" $ property $ prop_MapSingletonStencil (Proxy :: Proxy Ix2) + it "Ix3" $ property $ prop_MapSingletonStencil (Proxy :: Proxy Ix3) + it "Ix4" $ property $ prop_MapSingletonStencil (Proxy :: Proxy Ix4) + describe "DangerousStencil" $ do + it "Ix1" $ property $ prop_DangerousStencil (Proxy :: Proxy Ix1) + it "Ix2" $ property $ prop_DangerousStencil (Proxy :: Proxy Ix2) + it "Ix3" $ property $ prop_DangerousStencil (Proxy :: Proxy Ix3) + it "Ix4" $ property $ prop_DangerousStencil (Proxy :: Proxy Ix4) +-- describe "Storable" $ do +-- it "Ix1" $ property $ prop_toFromVector (Nothing :: Maybe Ix1) S +-- it "Ix2" $ property $ prop_toFromVector (Nothing :: Maybe Ix2) S +-- it "Ix3" $ property $ prop_toFromVector (Nothing :: Maybe Ix3) S +-- describe "Primitive" $ do +-- it "Ix1" $ property $ prop_toFromVector (Nothing :: Maybe Ix1) P +-- it "Ix2" $ property $ prop_toFromVector (Nothing :: Maybe Ix2) P +-- it "Ix3" $ property $ prop_toFromVector (Nothing :: Maybe Ix3) P +-- describe "Boxed" $ do +-- it "Ix1" $ property $ prop_toFromVector (Nothing :: Maybe Ix1) B +-- it "Ix2" $ property $ prop_toFromVector (Nothing :: Maybe Ix2) B +-- it "Ix3" $ property $ prop_toFromVector (Nothing :: Maybe Ix3) B + + + +spec :: Spec +spec = describe "Stencil" stencilSpec diff --git a/tests/Data/Massiv/Core/IndexSpec.hs b/tests/Data/Massiv/Core/IndexSpec.hs new file mode 100644 --- /dev/null +++ b/tests/Data/Massiv/Core/IndexSpec.hs @@ -0,0 +1,319 @@ +{-# LANGUAGE FlexibleContexts #-} +{-# LANGUAGE FlexibleInstances #-} +{-# LANGUAGE GADTs #-} +{-# LANGUAGE ScopedTypeVariables #-} +{-# LANGUAGE TypeOperators #-} +module Data.Massiv.Core.IndexSpec (Sz(..), SzZ(..), SzIx(..), DimIx(..), spec) where + +import Control.Monad +import Data.Massiv.Core.Index +import Data.Functor.Identity +import Test.Hspec +import Test.QuickCheck + +-- | Size that will result in a non-empty array +newtype Sz ix = Sz ix deriving Show + +-- | Size that can have zero elements +newtype SzZ ix = SzZ ix deriving Show + +-- | Dimension that is always within bounds of an index +newtype DimIx ix = DimIx Dim deriving Show + +instance Functor Sz where + fmap f (Sz sz) = Sz (f sz) + +instance Functor SzZ where + fmap f (SzZ sz) = SzZ (f sz) + +data SzIx ix = SzIx (Sz ix) ix deriving Show + +instance (Index ix, Arbitrary ix) => Arbitrary (Sz ix) where + arbitrary = do + sz <- liftIndex ((+1) . abs) <$> arbitrary + if totalElem sz > 200000 + then arbitrary + else return $ Sz sz + +instance (Index ix, Arbitrary ix) => Arbitrary (SzZ ix) where + arbitrary = SzZ <$> liftIndex abs <$> arbitrary + + +instance (Index ix, Arbitrary ix) => Arbitrary (SzIx ix) where + arbitrary = do + Sz sz <- arbitrary + -- Make sure index is within bounds: + SzIx (Sz sz) <$> flip (liftIndex2 mod) sz <$> arbitrary + + +instance Arbitrary e => Arbitrary (Border e) where + arbitrary = + oneof + [ Fill <$> arbitrary + , return Wrap + , return Edge + , return Reflect + , return Continue + ] + + +instance Index ix => Arbitrary (DimIx ix) where + arbitrary = do + n <- arbitrary + return $ DimIx (1 + (Dim n `mod` (rank (undefined :: ix)))) + +instance Arbitrary Ix2 where + arbitrary = (:.) <$> arbitrary <*> arbitrary + +instance Arbitrary Ix3 where + arbitrary = (:>) <$> arbitrary <*> ((:.) <$> arbitrary <*> arbitrary) + +instance Arbitrary Ix4 where + arbitrary = (:>) <$> arbitrary <*> arbitrary + +instance Arbitrary Ix5 where + arbitrary = (:>) <$> arbitrary <*> arbitrary + +instance CoArbitrary Ix2 where + coarbitrary (i :. j) = coarbitrary i . coarbitrary j + +instance CoArbitrary Ix3 where + coarbitrary (i :> ix) = coarbitrary i . coarbitrary ix + +instance CoArbitrary Ix4 where + coarbitrary (i :> ix) = coarbitrary i . coarbitrary ix + +instance CoArbitrary Ix5 where + coarbitrary (i :> ix) = coarbitrary i . coarbitrary ix + + +prop_IsSafeIx :: Index ix => proxy ix -> SzIx ix -> Bool +prop_IsSafeIx _ (SzIx (Sz sz) ix) = isSafeIndex sz ix + +prop_RepairSafeIx :: Index ix => proxy ix -> SzIx ix -> Bool +prop_RepairSafeIx _ (SzIx (Sz sz) ix) = + ix == repairIndex sz ix (error "Impossible") (error "Impossible") + +prop_UnconsCons :: Index ix => proxy ix -> ix -> Bool +prop_UnconsCons _ ix = ix == uncurry consDim (unconsDim ix) + +prop_UnsnocSnoc :: Index ix => proxy ix -> ix -> Bool +prop_UnsnocSnoc _ ix = ix == uncurry snocDim (unsnocDim ix) + +prop_ToFromLinearIndex :: Index ix => proxy ix -> SzIx ix -> Property +prop_ToFromLinearIndex _ (SzIx (Sz sz) ix) = + isSafeIndex sz ix ==> ix == fromLinearIndex sz (toLinearIndex sz ix) + +prop_FromToLinearIndex :: Index ix => proxy ix -> Sz ix -> Int -> Property +prop_FromToLinearIndex _ (Sz sz) i = + totalElem sz >= i ==> i == toLinearIndex sz (fromLinearIndex sz i) + +prop_CountElements :: Index ix => proxy ix -> Int -> Sz ix -> Property +prop_CountElements _ thresh (Sz sz) = + totalElem sz < thresh ==> totalElem sz == iter zeroIndex sz 1 (<) 0 (\ _ acc -> (acc + 1)) + +prop_IterMonotonic :: Index ix => proxy ix -> Int -> Sz ix -> Property +prop_IterMonotonic _ thresh (Sz sz) = + totalElem sz < thresh ==> fst $ + iter (liftIndex succ zeroIndex) sz 1 (<) (True, zeroIndex) $ \ curIx (prevMono, prevIx) -> + let isMono = prevMono && prevIx < curIx in isMono `seq` (isMono, curIx) + + +prop_IterMonotonic' :: Index ix => proxy ix -> Int -> Sz ix -> Property +prop_IterMonotonic' _ thresh (Sz sz) = + totalElem sz < + thresh ==> + if isM + then isM + else error (show a) + where + (isM, a, _) = + iter (liftIndex succ zeroIndex) sz 1 (<) (True, [], zeroIndex) $ + \ curIx (prevMono, acc, prevIx) -> + let nAcc = (prevIx, curIx, prevIx < curIx) : acc + isMono = prevMono && prevIx < curIx + in isMono `seq` (isMono, nAcc, curIx) + + +prop_IterMonotonicBackwards' :: Index ix => proxy ix -> Int -> Sz ix -> Property +prop_IterMonotonicBackwards' _ thresh (Sz sz) = + totalElem sz < + thresh ==> + if isM + then isM + else error (show a) + where + (isM, a, _) = + iter (liftIndex pred sz) zeroIndex (-1) (>=) (True, [], sz) $ + \ curIx (prevMono, acc, prevIx) -> + let isMono = prevMono && prevIx > curIx + nAcc = (prevIx, curIx, prevIx > curIx) : acc + in isMono `seq` (isMono, nAcc, curIx) + +prop_IterMonotonicM :: Index ix => proxy ix -> Int -> Sz ix -> Property +prop_IterMonotonicM _ thresh (Sz sz) = + totalElem sz < thresh ==> fst $ + runIdentity $ + iterM (liftIndex succ zeroIndex) sz 1 (<) (True, zeroIndex) $ \curIx (prevMono, prevIx) -> + let isMono = prevMono && prevIx < curIx + in return $ isMono `seq` (isMono, curIx) + + +prop_IterMonotonicBackwards :: Index ix => proxy ix -> Int -> Sz ix -> Property +prop_IterMonotonicBackwards _ thresh (Sz sz) = + totalElem sz < thresh ==> fst $ + iter (liftIndex pred sz) zeroIndex (-1) (>=) (True, sz) $ \ curIx (prevMono, prevIx) -> + let isMono = prevMono && prevIx > curIx in isMono `seq` (isMono, curIx) + +prop_IterMonotonicBackwardsM :: Index ix => proxy ix -> Int -> Sz ix -> Property +prop_IterMonotonicBackwardsM _ thresh (Sz sz) = + totalElem sz < thresh ==> fst $ runIdentity $ + iterM (liftIndex pred sz) zeroIndex (-1) (>=) (True, sz) $ \ curIx (prevMono, prevIx) -> + let isMono = prevMono && prevIx > curIx in return $ isMono `seq` (isMono, curIx) + +prop_LiftLift2 :: Index ix => proxy ix -> ix -> Int -> Bool +prop_LiftLift2 _ ix delta = liftIndex2 (+) ix (liftIndex (+delta) zeroIndex) == + liftIndex (+delta) ix + + +instance Show (Ix1 -> Double) where + show _ = "Index Func: Ix1 -> Double" + + +prop_BorderRepairSafe :: Index ix => proxy ix -> Border ix -> Sz ix -> ix -> Property +prop_BorderRepairSafe _ border@(Fill defIx) (Sz sz) ix = + not (isSafeIndex sz ix) ==> handleBorderIndex border sz id ix == defIx +prop_BorderRepairSafe _ border (Sz sz) ix = + not (isSafeIndex sz ix) ==> isSafeIndex sz (handleBorderIndex border sz id ix) + + +prop_UnconsGetDrop :: (Index (Lower ix), Index ix) => proxy ix -> ix -> Bool +prop_UnconsGetDrop _ ix = + Just (unconsDim ix) == do + i <- getIndex ix (rank ix) + ixL <- dropDim ix (rank ix) + return (i, ixL) + +prop_UnsnocGetDrop :: (Index (Lower ix), Index ix) => proxy ix -> ix -> Bool +prop_UnsnocGetDrop _ ix = + Just (unsnocDim ix) == do + i <- getIndex ix 1 + ixL <- dropDim ix 1 + return (ixL, i) + +prop_SetAll :: Index ix => proxy ix -> ix -> Int -> Bool +prop_SetAll _ ix i = + foldM (\cix d -> setIndex cix d i) ix ([1 .. rank ix] :: [Dim]) == + Just (pureIndex i) + + +prop_SetGet :: Index ix => proxy ix -> ix -> DimIx ix -> Int -> Bool +prop_SetGet _ ix (DimIx dim) n = Just n == (setIndex ix dim n >>= (`getIndex` dim)) + + +prop_BorderIx1 :: Positive Int -> Border Double -> (Ix1 -> Double) -> Sz Ix1 -> Ix1 -> Bool +prop_BorderIx1 (Positive period) border getVal (Sz sz) ix = + if isSafeIndex sz ix + then getVal ix == val + else case border of + Fill defVal -> defVal == val + Wrap -> + val == + handleBorderIndex + border + sz + getVal + (liftIndex2 (+) (liftIndex (* period) sz) ix) + Edge -> + if ix < 0 + then val == getVal (liftIndex (max 0) ix) + else val == + getVal (liftIndex2 min (liftIndex (subtract 1) sz) ix) + Reflect -> + val == + handleBorderIndex + border + sz + getVal + (liftIndex2 (+) (liftIndex (* (2 * signum ix * period)) sz) ix) + Continue -> + val == + handleBorderIndex + Reflect + sz + getVal + (if ix < 0 + then ix - 1 + else ix + 1) + where + val = handleBorderIndex border sz getVal ix + +specDimN :: (Index ix, Ord ix, CoArbitrary ix, Arbitrary ix) => proxy ix -> Spec +specDimN proxy = do + describe "Safety" $ do + it "isSafeIndex" $ property $ prop_IsSafeIx proxy + it "RepairSafeIx" $ property $ prop_RepairSafeIx proxy + describe "Lifting" $ do + it "Lift/Lift2" $ property $ prop_LiftLift2 proxy + describe "Linear" $ do + it "ToFromLinearIndex" $ property $ prop_ToFromLinearIndex proxy + it "FromToLinearIndex" $ property $ prop_FromToLinearIndex proxy + describe "Iterator" $ do + it "CountElements" $ property $ prop_CountElements proxy (2000000) + it "Monotonic" $ property $ prop_IterMonotonic proxy (2000000) + it "MonotonicBackwards" $ property $ prop_IterMonotonicBackwards proxy (2000000) + it "MonotonicM" $ property $ prop_IterMonotonicM proxy (2000000) + it "MonotonicBackwardsM" $ property $ prop_IterMonotonicBackwardsM proxy (2000000) + describe "Border" $ do + it "BorderRepairSafe" $ property $ prop_BorderRepairSafe proxy + describe "SetGetDrop" $ do + it "SetAll" $ property $ prop_SetAll proxy + it "SetGet" $ property $ prop_SetGet proxy + +specDim2AndUp + :: (Index ix, Index (Lower ix), Ord ix, CoArbitrary ix, Arbitrary ix) + => proxy ix -> Spec +specDim2AndUp proxy = do + describe "Higher/Lower" $ do + it "UnconsCons" $ property $ prop_UnconsCons proxy + it "UnsnocSnoc" $ property $ prop_UnsnocSnoc proxy + it "UnconsGetDrop" $ property $ prop_UnconsGetDrop proxy + it "UnsnocGetDrop" $ property $ prop_UnsnocGetDrop proxy + + +spec :: Spec +spec = do + describe "Tuple based indices" $ do + describe "Ix1T" $ do + specDimN (Nothing :: Maybe Ix1T) + it "BorderIndex" $ property $ prop_BorderIx1 + describe "Ix2T" $ do + specDimN (Nothing :: Maybe Ix2T) + specDim2AndUp (Nothing :: Maybe Ix2T) + describe "Ix3T" $ do + specDimN (Nothing :: Maybe Ix3T) + specDim2AndUp (Nothing :: Maybe Ix3T) + describe "Ix4T" $ do + specDimN (Nothing :: Maybe Ix4T) + specDim2AndUp (Nothing :: Maybe Ix4T) + describe "Ix5T" $ do + specDimN (Nothing :: Maybe Ix5T) + specDim2AndUp (Nothing :: Maybe Ix5T) + describe "Specialized indices" $ do + describe "Ix2" $ do + -- These can be used to quickly debug monotonicity + it "Monotonic'" $ + property $ prop_IterMonotonic' (Nothing :: Maybe Ix2) (20000) + it "MonotonicBackwards'" $ + property $ prop_IterMonotonicBackwards' (Nothing :: Maybe Ix2) (20000) + specDimN (Nothing :: Maybe Ix2) + specDim2AndUp (Nothing :: Maybe Ix2) + describe "Ix3" $ do + specDimN (Nothing :: Maybe Ix3) + specDim2AndUp (Nothing :: Maybe Ix3) + describe "Ix4" $ do + specDimN (Nothing :: Maybe Ix4) + specDim2AndUp (Nothing :: Maybe Ix4) + describe "Ix5" $ do + specDimN (Nothing :: Maybe Ix5) + specDim2AndUp (Nothing :: Maybe Ix5) diff --git a/tests/Data/Massiv/Core/SchedulerSpec.hs b/tests/Data/Massiv/Core/SchedulerSpec.hs new file mode 100644 --- /dev/null +++ b/tests/Data/Massiv/Core/SchedulerSpec.hs @@ -0,0 +1,78 @@ +{-# LANGUAGE FlexibleContexts #-} +{-# LANGUAGE FlexibleInstances #-} +{-# LANGUAGE MultiParamTypeClasses #-} +module Data.Massiv.Core.SchedulerSpec (spec) where + +import Control.Concurrent +import Control.Exception.Base (ArithException (DivideByZero), + AsyncException (ThreadKilled)) +import Data.Massiv.Core.Scheduler +import Data.Massiv.CoreArbitrary as A +import Prelude as P +import Test.Hspec +import Test.QuickCheck +import Test.QuickCheck.Monadic + + +-- | Ensure proper exception handling. +prop_CatchDivideByZero :: ArrIx D Ix2 Int -> [Int] -> Property +prop_CatchDivideByZero (ArrIx arr ix) caps = + assertException + (== DivideByZero) + (A.sum $ + A.imap + (\ix' x -> + if ix == ix' + then x `div` 0 + else x) + (setComp (ParOn caps) arr)) + +-- | Ensure proper exception handling in nested parallel computation +prop_CatchNested :: ArrIx D Ix1 (ArrIxP D Ix1 Int) -> [Int] -> Property +prop_CatchNested (ArrIx arr ix) caps = + assertException + (== DivideByZero) + (computeAs U $ + A.map A.sum $ + A.imap + (\ix' (ArrIxP iarr ixi) -> + if ix == ix' + then A.imap + (\ixi' e -> + if ixi == ixi' + then e `div` 0 + else e) + iarr + else iarr) + (setComp (ParOn caps) arr)) + +-- | Make sure there is no deadlock if all workers get killed +prop_AllWorkersDied :: [Int] -> (Int, [Int]) -> Property +prop_AllWorkersDied wIds (hId, ids) = + assertExceptionIO + (== ThreadKilled) + (withScheduler_ [] $ \scheduler1 -> + scheduleWork + scheduler1 + (withScheduler_ wIds $ \scheduler -> + P.mapM_ + (\_ -> scheduleWork scheduler (myThreadId >>= killThread)) + (hId : ids))) + + +-- | Check weather all jobs have been completed and returned order is correct +prop_SchedulerAllJobsProcessed :: [Int] -> OrderedList Int -> Property +prop_SchedulerAllJobsProcessed wIds (Ordered jobs) = + monadicIO $ do + res <- (run $ withScheduler' wIds $ \scheduler -> + P.mapM_ (scheduleWork scheduler . return) jobs) + return (res === jobs) + + +spec :: Spec +spec = do + describe "Exceptions" $ do + it "CatchDivideByZero" $ property prop_CatchDivideByZero + it "CatchNested" $ property prop_CatchNested + it "AllWorkersDied" $ property prop_AllWorkersDied + it "SchedulerAllJobsProcessed" $ property prop_SchedulerAllJobsProcessed diff --git a/tests/Data/Massiv/CoreArbitrary.hs b/tests/Data/Massiv/CoreArbitrary.hs new file mode 100644 --- /dev/null +++ b/tests/Data/Massiv/CoreArbitrary.hs @@ -0,0 +1,140 @@ +{-# LANGUAGE FlexibleContexts #-} +{-# LANGUAGE FlexibleInstances #-} +{-# LANGUAGE MultiParamTypeClasses #-} +{-# LANGUAGE ScopedTypeVariables #-} +{-# LANGUAGE StandaloneDeriving #-} +{-# LANGUAGE UndecidableInstances #-} +module Data.Massiv.CoreArbitrary + ( Arr(..) + , ArrIx(..) + , ArrP(..) + , ArrIxP(..) + , Sz(..) + , SzIx(..) + , SzZ(..) + , DimIx(..) + , assertException + , assertSomeException + , assertExceptionIO + , assertSomeExceptionIO + , module Data.Massiv.Array + ) where + +import Control.DeepSeq (NFData, deepseq) +import Control.Exception (Exception, SomeException, catch) +--import Data.Massiv.Array.Ops.Construct +import Data.Massiv.Array +import Data.Massiv.Core.IndexSpec hiding (spec) +import Data.Typeable +import Test.QuickCheck +import Test.QuickCheck.Monadic + +data Arr r ix e = Arr (Array r ix e) + +data ArrS r ix e = ArrS (Array r ix e) + +data ArrP r ix e = ArrP (Array r ix e) + +data ArrIx r ix e = ArrIx (Array r ix e) ix + +data ArrIxS r ix e = ArrIxS (Array r ix e) ix + +data ArrIxP r ix e = ArrIxP (Array r ix e) ix + +deriving instance (Show (Array r ix e)) => Show (Arr r ix e) +deriving instance (Show (Array r ix e)) => Show (ArrS r ix e) +deriving instance (Show (Array r ix e)) => Show (ArrP r ix e) +deriving instance (Show (Array r ix e), Show ix) => Show (ArrIx r ix e) +deriving instance (Show (Array r ix e), Show ix) => Show (ArrIxS r ix e) +deriving instance (Show (Array r ix e), Show ix) => Show (ArrIxP r ix e) + +instance Arbitrary Comp where + arbitrary = oneof [pure Seq, fmap ParOn arbitrary] + + +-- | Arbitrary array +instance (CoArbitrary ix, Arbitrary ix, Typeable e, Construct r ix e, Arbitrary e) => + Arbitrary (Array r ix e) where + arbitrary = do + SzZ sz <- arbitrary + func <- arbitrary + comp <- oneof [pure Seq, pure Par] + return $ makeArray comp sz func + + +-- | Arbitrary non-empty array. Computation strategy can be either `Seq` or `Par`. +instance (CoArbitrary ix, Arbitrary ix, Typeable e, Construct r ix e, Arbitrary e) => + Arbitrary (Arr r ix e) where + arbitrary = do + Sz sz <- arbitrary + func <- arbitrary + comp <- oneof [pure Seq, pure Par] + return $ Arr $ makeArray comp sz func + +-- | Arbitrary non-empty array +instance (CoArbitrary ix, Arbitrary ix, Typeable e, Construct r ix e, Arbitrary e) => + Arbitrary (ArrS r ix e) where + arbitrary = do + Sz sz <- arbitrary + func <- arbitrary + return $ ArrS $ makeArray Seq sz func + +instance (CoArbitrary ix, Arbitrary ix, Typeable e, Construct r ix e, Arbitrary e) => + Arbitrary (ArrP r ix e) where + arbitrary = do + Arr arr <- arbitrary + return $ ArrP (setComp Par arr) + +-- | Arbitrary non-empty array with a valid index +instance (CoArbitrary ix, Arbitrary ix, Typeable e, Construct r ix e, Arbitrary e) => + Arbitrary (ArrIx r ix e) where + arbitrary = do + SzIx (Sz sz) ix <- arbitrary + func <- arbitrary + comp <- arbitrary + return $ ArrIx (makeArray comp sz func) ix + +-- | Arbitrary non-empty array with a valid index +instance (CoArbitrary ix, Arbitrary ix, Typeable e, Construct r ix e, Arbitrary e) => + Arbitrary (ArrIxS r ix e) where + arbitrary = do + SzIx (Sz sz) ix <- arbitrary + func <- arbitrary + return $ ArrIxS (makeArray Seq sz func) ix + + +-- | Arbitrary non-empty array with a valid index +instance (CoArbitrary ix, Arbitrary ix, Typeable e, Construct r ix e, Arbitrary e) => + Arbitrary (ArrIxP r ix e) where + arbitrary = do + ArrIx arrIx ix <- arbitrary + return $ ArrIxP (setComp Par arrIx) ix + + +assertException :: (NFData a, Exception exc) => + (exc -> Bool) -- ^ Return True if that is the exception that was expected + -> a -- ^ Value that should throw an exception, when fully evaluated + -> Property +assertException isExc action = assertExceptionIO isExc (return action) + + +assertSomeException :: NFData a => a -> Property +assertSomeException = assertSomeExceptionIO . return + + +assertExceptionIO :: (NFData a, Exception exc) => + (exc -> Bool) -- ^ Return True if that is the exception that was expected + -> IO a -- ^ IO Action that should throw an exception + -> Property +assertExceptionIO isExc action = + monadicIO $ do + hasFailed <- + run + (catch + (do res <- action + res `deepseq` return False) $ \exc -> + show exc `deepseq` return (isExc exc)) + assert hasFailed + +assertSomeExceptionIO :: NFData a => IO a -> Property +assertSomeExceptionIO = assertExceptionIO (\exc -> const True (exc :: SomeException)) diff --git a/tests/Spec.hs b/tests/Spec.hs new file mode 100644 --- /dev/null +++ b/tests/Spec.hs @@ -0,0 +1,32 @@ +module Main where + +import Data.Massiv.Array.DelayedSpec as Delayed +import Data.Massiv.Array.Manifest.VectorSpec as Vector +import Data.Massiv.Array.Ops.ConstructSpec as Construct +import Data.Massiv.Array.Ops.FoldSpec as Fold +import Data.Massiv.Array.Ops.SliceSpec as Slice +import Data.Massiv.Array.Ops.TransformSpec as Transform +import Data.Massiv.Array.StencilSpec as Stencil +import Data.Massiv.Core.IndexSpec as Index +import Data.Massiv.Core.SchedulerSpec as Scheduler +import System.IO (BufferMode (LineBuffering), + hSetBuffering, stdout) +import Test.Hspec + + +-- | Main entry point. Returns ExitFailure if a test fails. +main :: IO () +main = do + hSetBuffering stdout LineBuffering + hspec $ do + describe "Core" $ do + Scheduler.spec + Index.spec + describe "Ops" $ do + Construct.spec + Fold.spec + Slice.spec + Transform.spec + Delayed.spec + Stencil.spec + Vector.spec