dph-common/Data/Array/Parallel/Lifted/PArray.hs

   1 {-# LANGUAGE CPP #-}
   2
   3 #include "fusion-phases.h"
   4
   5 module Data.Array.Parallel.Lifted.PArray (
   6   PArray(..), PData,
   7
   8   PA(..),
   9   lengthPA#, dataPA#, replicatePA#, replicatelPA#, repeatPA#, repeatcPA#,
  10   emptyPA, indexPA#, extractPA#, bpermutePA#, appPA#, applPA#,
  11   packPA#, combine2PA#, fromListPA#, fromListPA, nfPA,
  12
  13   replicatePD, replicatelPD, repeatPD, repeatcPD, emptyPD,
  14   indexPD, extractPD, bpermutePD, appPD, applPD,
  15   packPD, combine2PD, fromListPD, fromListPD, nfPD,
  16
  17   PRepr, PR(..), mkPR, mkReprPA,
  18
  19   T_replicatePR, T_replicatelPR, T_repeatPR, T_repeatcPR, T_emptyPR,
  20   T_indexPR, T_extractPR, T_bpermutePR, T_appPR, T_applPR,
  21   T_packPR, T_combine2PR, T_fromListPR, T_fromListPR, T_nfPR
  22 ) where
  23
  24 import qualified Data.Array.Parallel.Unlifted as U
  25 import Data.Array.Parallel.Lifted.Selector
  26 import Data.Array.Parallel.Lifted.Unboxed ( elementsSegd# )
  27 import GHC.Exts (Int#, Int(..), (+#), (*#))
  28
  29 -- |Lifted parallel arrays
  30 --
  31 data PArray a = PArray Int# (PData a)
  32 data family PData a
  33
  34 -- |Representation types
  35 --
  36 type family PRepr a
  37
  38 -- |Dictionaries
  39 --
  40
  41 data PA a = PA {
  42               toPRepr      :: a               -> PRepr a
  43             , fromPRepr    :: PRepr a         -> a
  44             , toArrPRepr   :: PData a         -> PData (PRepr a)
  45             , fromArrPRepr :: PData (PRepr a) -> PData a
  46             , dictPRepr    :: PR (PRepr a)
  47             }
  48
  49
  50 type T_emptyPR      a =  PData a
  51
  52 type T_replicatePR  a =  Int# -> a -> PData a
  53
  54 type T_replicatelPR a =  U.Segd            -- segd of result array
  55                       -> PData a -> PData a
  56
  57 type T_repeatPR     a =  Int#              -- number of times to repeat
  58                       -> Int#              -- length of src array
  59                       -> PData a -> PData a
  60
  61 type T_repeatcPR    a =  Int#              -- length of result array
  62                       -> U.Array Int       -- number of times each segment
  63                                            -- is repeated
  64                       -> U.Segd            -- src segd
  65                       -> PData a -> PData a
  66
  67 type T_indexPR      a =  PData a -> Int# -> a
  68
  69 type T_extractPR    a =  PData a
  70                       -> Int#              -- starting index
  71                       -> Int#              -- length of result array
  72                       -> PData a
  73
  74 type T_bpermutePR   a =  PData a
  75                       -> Int#              -- result length
  76                       -> U.Array Int       -- indices
  77                       -> PData a
  78
  79 type T_appPR        a = PData a -> PData a -> PData a
  80
  81 type T_applPR       a =  U.Segd -> PData a   -- src segd/data 1
  82                       -> U.Segd -> PData a   -- src segd/data 2
  83                       -> PData a
  84
  85 type T_packPR       a =  PData a
  86                       -> Int#              -- result length
  87                       -> U.Array Bool      -- flags
  88                       -> PData a
  89
  90 type T_combine2PR   a =  Int#              -- result length
  91                       -> Sel2              -- selector
  92                       -> PData a -> PData a -> PData a
  93
  94 type T_fromListPR a = Int# -> [a] -> PData a
  95
  96 type T_nfPR a = PData a -> ()
  97
  98 data PR a = PR {
  99               emptyPR      :: T_emptyPR a
 100             , replicatePR  :: T_replicatePR a
 101             , replicatelPR :: T_replicatelPR a
 102             , repeatPR     :: T_repeatPR a
 103             , repeatcPR    :: T_repeatcPR a
 104             , indexPR      :: T_indexPR a
 105             , extractPR    :: T_extractPR a
 106             , bpermutePR   :: T_bpermutePR a
 107             , appPR        :: T_appPR a
 108             , applPR       :: T_applPR a
 109             , packPR       :: T_packPR a
 110             , combine2PR   :: T_combine2PR a
 111             , fromListPR   :: T_fromListPR a
 112             , nfPR         :: T_nfPR a
 113             }
 114
 115 emptyPD :: PA a -> T_emptyPR a
 116 {-# INLINE_PA emptyPD #-}
 117 emptyPD pa = fromArrPRepr pa
 118           $ emptyPR (dictPRepr pa)
 119
 120 replicatePD :: PA a -> T_replicatePR a
 121 {-# INLINE_PA replicatePD #-}
 122 replicatePD pa n# x = fromArrPRepr pa
 123                     . replicatePR (dictPRepr pa) n#
 124                     $ toPRepr pa x
 125
 126 replicatelPD :: PA a -> T_replicatelPR a
 127 {-# INLINE_PA replicatelPD #-}
 128 replicatelPD pa segd xs = fromArrPRepr pa
 129                         . replicatelPR (dictPRepr pa) segd
 130                         $ toArrPRepr pa xs
 131
 132 repeatPD :: PA a -> T_repeatPR a
 133 {-# INLINE_PA repeatPD #-}
 134 repeatPD pa n# len# xs = fromArrPRepr pa
 135                        . repeatPR (dictPRepr pa) n# len#
 136                        $ toArrPRepr pa xs
 137
 138 repeatcPD :: PA a -> T_repeatcPR a
 139 {-# INLINE_PA repeatcPD #-}
 140 repeatcPD pa n# ns segd xs = fromArrPRepr pa
 141                            . repeatcPR (dictPRepr pa) n# ns segd
 142                            $ toArrPRepr pa xs
 143
 144 indexPD :: PA a -> T_indexPR a
 145 {-# INLINE_PA indexPD #-}
 146 indexPD pa xs i# = fromPRepr pa
 147                  $ indexPR (dictPRepr pa) (toArrPRepr pa xs) i#
 148
 149 extractPD :: PA a -> T_extractPR a
 150 {-# INLINE_PA extractPD #-}
 151 extractPD pa xs i# m# = fromArrPRepr pa
 152                       $ extractPR (dictPRepr pa) (toArrPRepr pa xs) i# m#
 153
 154 bpermutePD :: PA a -> T_bpermutePR a
 155 {-# INLINE bpermutePD #-}
 156 bpermutePD pa xs n# is = fromArrPRepr pa
 157                        $ bpermutePR (dictPRepr pa) (toArrPRepr pa xs) n# is
 158
 159 appPD :: PA a -> T_appPR a
 160 {-# INLINE_PA appPD #-}
 161 appPD pa xs ys = fromArrPRepr pa
 162                $ appPR (dictPRepr pa) (toArrPRepr pa xs) (toArrPRepr pa ys)
 163
 164 applPD :: PA a -> T_applPR a
 165 {-# INLINE_PA applPD #-}
 166 applPD pa is xs js ys = fromArrPRepr pa
 167                       $ applPR (dictPRepr pa) is (toArrPRepr pa xs)
 168                                               js (toArrPRepr pa ys)
 169
 170 packPD :: PA a -> T_packPR a
 171 {-# INLINE_PA packPD #-}
 172 packPD pa xs n# bs = fromArrPRepr pa
 173                    $ packPR (dictPRepr pa) (toArrPRepr pa xs) n# bs
 174
 175 combine2PD :: PA a -> T_combine2PR a
 176 {-# INLINE_PA combine2PD #-}
 177 combine2PD pa n# sel as bs
 178   = fromArrPRepr pa
 179   $ combine2PR (dictPRepr pa) n# sel (toArrPRepr pa as)
 180                                      (toArrPRepr pa bs)
 181
 182 fromListPD :: PA a -> T_fromListPR a
 183 {-# INLINE_PA fromListPD #-}
 184 fromListPD pa n# xs = fromArrPRepr pa
 185                     $ fromListPR (dictPRepr pa) n# (map (toPRepr pa) xs)
 186
 187 nfPD :: PA a -> T_nfPR a
 188 {-# INLINE nfPD #-}
 189 nfPD pa xs = nfPR (dictPRepr pa) (toArrPRepr pa xs)
 190
 191
 192
 193
 194 lengthPA# :: PArray a -> Int#
 195 {-# INLINE_PA lengthPA# #-}
 196 lengthPA# (PArray n# _) = n#
 197
 198 dataPA# :: PArray a -> PData a
 199 {-# INLINE_PA dataPA# #-}
 200 dataPA# (PArray _ d) = d
 201
 202 emptyPA :: PA a -> PArray a
 203 {-# INLINE_PA emptyPA #-}
 204 emptyPA pa = PArray 0# (emptyPD pa)
 205
 206 replicatePA# :: PA a -> Int# -> a -> PArray a
 207 {-# INLINE_PA replicatePA# #-}
 208 replicatePA# pa n# x = PArray n# (replicatePD pa n# x)
 209
 210 replicatelPA# :: PA a -> U.Segd -> PArray a -> PArray a
 211 {-# INLINE_PA replicatelPA# #-}
 212 replicatelPA# pa segd (PArray n# xs)
 213   = PArray (elementsSegd# segd) (replicatelPD pa segd xs)
 214
 215 repeatPA# :: PA a -> Int# -> PArray a -> PArray a
 216 {-# INLINE_PA repeatPA# #-}
 217 repeatPA# pa m# (PArray n# xs) = PArray (m# *# n#) (repeatPD pa m# n# xs)
 218
 219 repeatcPA# :: PA a -> U.Array Int -> U.Segd -> PArray a -> PArray a
 220 {-# INLINE_PA repeatcPA# #-}
 221 repeatcPA# pa ns segd (PArray n# xs)
 222   = case U.sum (U.zipWith (*) ns (U.lengthsSegd segd)) of
 223       I# m# -> PArray m# (repeatcPD pa m# ns segd xs)
 224
 225 indexPA# :: PA a -> PArray a -> Int# -> a
 226 {-# INLINE_PA indexPA# #-}
 227 indexPA# pa (PArray _ xs) i# = indexPD pa xs i#
 228
 229 extractPA# :: PA a -> PArray a -> Int# -> Int# -> PArray a
 230 {-# INLINE_PA extractPA# #-}
 231 extractPA# pa (PArray _ xs) i# n# = PArray n# (extractPD pa xs i# n#)
 232
 233 bpermutePA# :: PA a -> PArray a -> Int# -> U.Array Int -> PArray a
 234 {-# INLINE bpermutePA# #-}
 235 bpermutePA# pa (PArray _ xs) n# is = PArray n# (bpermutePD pa xs n# is)
 236
 237 appPA# :: PA a -> PArray a -> PArray a -> PArray a
 238 {-# INLINE_PA appPA# #-}
 239 appPA# pa (PArray m# xs) (PArray n# ys) = PArray (m# +# n#) (appPD pa xs ys)
 240
 241 applPA# :: PA a -> U.Segd -> PArray a -> U.Segd -> PArray a -> PArray a
 242 {-# INLINE_PA applPA# #-}
 243 applPA# pa is (PArray m# xs) js (PArray n# ys)
 244   = PArray (m# +# n#) (applPD pa is xs js ys)
 245
 246 packPA# :: PA a -> PArray a -> Int# -> U.Array Bool -> PArray a
 247 {-# INLINE_PA packPA# #-}
 248 packPA# pa (PArray _ xs) n# bs = PArray n# (packPD pa xs n# bs)
 249
 250 combine2PA# :: PA a -> Int# -> Sel2 -> PArray a -> PArray a -> PArray a
 251 {-# INLINE_PA combine2PA# #-}
 252 combine2PA# pa n# sel (PArray _ as) (PArray _ bs)
 253   = PArray n# (combine2PD pa n# sel as bs)
 254
 255 fromListPA# :: PA a -> Int# -> [a] -> PArray a
 256 {-# INLINE_PA fromListPA# #-}
 257 fromListPA# pa n# xs = PArray n# (fromListPD pa n# xs)
 258
 259 fromListPA :: PA a -> [a] -> PArray a
 260 {-# INLINE fromListPA #-}
 261 fromListPA pa xs = case length xs of
 262                      I# n# -> fromListPA# pa n# xs
 263
 264 nfPA :: PA a -> PArray a -> ()
 265 {-# INLINE nfPA #-}
 266 nfPA pa (PArray _ xs) = nfPD pa xs
 267
 268 mkPR :: PA a -> PR a
 269 {-# INLINE mkPR #-}
 270 mkPR pa = PR {
 271             emptyPR      = emptyPD pa
 272           , replicatePR  = replicatePD pa
 273           , replicatelPR = replicatelPD pa
 274           , repeatPR     = repeatPD pa
 275           , repeatcPR    = repeatcPD pa
 276           , indexPR      = indexPD pa
 277           , bpermutePR   = bpermutePD pa
 278           , appPR        = appPD pa
 279           , applPR       = applPD pa
 280           , packPR       = packPD pa
 281           , combine2PR   = combine2PD pa
 282           , fromListPR   = fromListPD pa
 283           , nfPR         = nfPD pa
 284           }
 285
 286 mkReprPA :: (a ~ PRepr a) => PR a -> PA a
 287 {-# INLINE mkReprPA #-}
 288 mkReprPA pr = PA {
 289                 toPRepr      = id
 290               , fromPRepr    = id
 291               , toArrPRepr   = id
 292               , fromArrPRepr = id
 293               , dictPRepr    = pr
 294               }
 295