{-# LANGUAGE BangPatterns         #-}
{-# LANGUAGE GADTs                #-}
{-# LANGUAGE MagicHash            #-}
{-# LANGUAGE ScopedTypeVariables  #-}
{-# LANGUAGE TemplateHaskell      #-}
{-# LANGUAGE TypeFamilies         #-}
{-# LANGUAGE UnboxedTuples        #-}
{-# LANGUAGE UndecidableInstances #-}
{-# OPTIONS_HADDOCK hide #-}
-- |
-- Module      : Data.Array.Accelerate.Array.Data
-- Copyright   : [2008..2020] The Accelerate Team
-- License     : BSD3
--
-- Maintainer  : Trevor L. McDonell <trevor.mcdonell@gmail.com>
-- Stability   : experimental
-- Portability : non-portable (GHC extensions)
--
-- This module fixes the concrete representation of Accelerate arrays.  We
-- allocate all arrays using pinned memory to enable safe direct-access by
-- non-Haskell code in multi-threaded code.  In particular, we can safely pass
-- pointers to an array's payload to foreign code.
--

module Data.Array.Accelerate.Array.Data (

  -- * Array operations and representations
  ArrayData, MutableArrayData, ScalarArrayData, GArrayDataR, ScalarArrayDataR,
  runArrayData,
  newArrayData,
  indexArrayData, readArrayData, writeArrayData,
  unsafeArrayDataPtr,
  touchArrayData,
  rnfArrayData,

  -- * Type macros
  HTYPE_INT, HTYPE_WORD, HTYPE_CLONG, HTYPE_CULONG, HTYPE_CCHAR,

  -- * Allocator internals
  registerForeignPtrAllocator,

  -- * Utilities for type classes
  ScalarArrayDict(..), scalarArrayDict,
  SingleArrayDict(..), singleArrayDict,

  -- * TemplateHaskell
  liftArrayData,

) where

-- friends
import Data.Array.Accelerate.Array.Unique
import Data.Array.Accelerate.Representation.Type
import Data.Array.Accelerate.Error
import Data.Array.Accelerate.Type
import Data.Primitive.Vec

import Data.Array.Accelerate.Debug.Flags
import Data.Array.Accelerate.Debug.Monitoring
import Data.Array.Accelerate.Debug.Trace


-- standard libraries
import Control.Applicative
import Control.DeepSeq
import Control.Monad                                                ( (<=<) )
import Data.Bits
import Data.IORef
import Data.Primitive                                               ( sizeOf# )
import Foreign.ForeignPtr
import Foreign.Storable
import Language.Haskell.TH                                          hiding ( Type )
import System.IO.Unsafe
import Text.Printf
import Prelude                                                      hiding ( mapM )

import GHC.Base
import GHC.ForeignPtr
import GHC.Ptr


-- | Immutable array representation
--
type ArrayData e = MutableArrayData e

-- | Mutable array representation
--
type MutableArrayData e = GArrayDataR UniqueArray e

-- | Underlying array representation.
--
-- NOTE: We use a standard (non-strict) pair to enable lazy device-host data transfers
--
type family GArrayDataR ba a where
  GArrayDataR ba ()     = ()
  GArrayDataR ba (a, b) = (GArrayDataR ba a, GArrayDataR ba b)
  GArrayDataR ba a      = ba (ScalarArrayDataR a)

type ScalarArrayData a = UniqueArray (ScalarArrayDataR a)

-- | Mapping from scalar type to the type as represented in memory in an
-- array.
--
type family ScalarArrayDataR t where
  ScalarArrayDataR Int       = Int
  ScalarArrayDataR Int8      = Int8
  ScalarArrayDataR Int16     = Int16
  ScalarArrayDataR Int32     = Int32
  ScalarArrayDataR Int64     = Int64
  ScalarArrayDataR Word      = Word
  ScalarArrayDataR Word8     = Word8
  ScalarArrayDataR Word16    = Word16
  ScalarArrayDataR Word32    = Word32
  ScalarArrayDataR Word64    = Word64
  ScalarArrayDataR Half      = Half
  ScalarArrayDataR Float     = Float
  ScalarArrayDataR Double    = Double
  ScalarArrayDataR (Vec n t) = ScalarArrayDataR t


data ScalarArrayDict a where
  ScalarArrayDict :: ( ArrayData a ~ ScalarArrayData a, ScalarArrayDataR a ~ ScalarArrayDataR b )
                  => {-# UNPACK #-} !Int    -- vector width
                  -> SingleType b           -- base type
                  -> ScalarArrayDict a

data SingleArrayDict a where
  SingleArrayDict :: ( ArrayData a ~ ScalarArrayData a, ScalarArrayDataR a ~ a )
                  => SingleArrayDict a

scalarArrayDict :: ScalarType a -> ScalarArrayDict a
scalarArrayDict :: ScalarType a -> ScalarArrayDict a
scalarArrayDict = ScalarType a -> ScalarArrayDict a
forall a. ScalarType a -> ScalarArrayDict a
scalar
  where
    scalar :: ScalarType a -> ScalarArrayDict a
    scalar :: ScalarType a -> ScalarArrayDict a
scalar (VectorScalarType VectorType (Vec n a)
t) = VectorType (Vec n a) -> ScalarArrayDict (Vec n a)
forall a. VectorType a -> ScalarArrayDict a
vector VectorType (Vec n a)
t
    scalar (SingleScalarType SingleType a
t)
      | SingleArrayDict a
SingleArrayDict <- SingleType a -> SingleArrayDict a
forall a. SingleType a -> SingleArrayDict a
singleArrayDict SingleType a
t
      = Int -> SingleType a -> ScalarArrayDict a
forall a b.
(ArrayData a ~ ScalarArrayData a,
 ScalarArrayDataR a ~ ScalarArrayDataR b) =>
Int -> SingleType b -> ScalarArrayDict a
ScalarArrayDict Int
1 SingleType a
t

    vector :: VectorType a -> ScalarArrayDict a
    vector :: VectorType a -> ScalarArrayDict a
vector (VectorType Int
w SingleType a
s)
      | SingleArrayDict a
SingleArrayDict <- SingleType a -> SingleArrayDict a
forall a. SingleType a -> SingleArrayDict a
singleArrayDict SingleType a
s
      = Int -> SingleType a -> ScalarArrayDict a
forall a b.
(ArrayData a ~ ScalarArrayData a,
 ScalarArrayDataR a ~ ScalarArrayDataR b) =>
Int -> SingleType b -> ScalarArrayDict a
ScalarArrayDict Int
w SingleType a
s

singleArrayDict :: SingleType a -> SingleArrayDict a
singleArrayDict :: SingleType a -> SingleArrayDict a
singleArrayDict = SingleType a -> SingleArrayDict a
forall a. SingleType a -> SingleArrayDict a
single
  where
    single :: SingleType a -> SingleArrayDict a
    single :: SingleType a -> SingleArrayDict a
single (NumSingleType NumType a
t) = NumType a -> SingleArrayDict a
forall a. NumType a -> SingleArrayDict a
num NumType a
t

    num :: NumType a -> SingleArrayDict a
    num :: NumType a -> SingleArrayDict a
num (IntegralNumType IntegralType a
t) = IntegralType a -> SingleArrayDict a
forall a. IntegralType a -> SingleArrayDict a
integral IntegralType a
t
    num (FloatingNumType FloatingType a
t) = FloatingType a -> SingleArrayDict a
forall a. FloatingType a -> SingleArrayDict a
floating FloatingType a
t

    integral :: IntegralType a -> SingleArrayDict a
    integral :: IntegralType a -> SingleArrayDict a
integral IntegralType a
TypeInt    = SingleArrayDict a
forall a.
(ArrayData a ~ ScalarArrayData a, ScalarArrayDataR a ~ a) =>
SingleArrayDict a
SingleArrayDict
    integral IntegralType a
TypeInt8   = SingleArrayDict a
forall a.
(ArrayData a ~ ScalarArrayData a, ScalarArrayDataR a ~ a) =>
SingleArrayDict a
SingleArrayDict
    integral IntegralType a
TypeInt16  = SingleArrayDict a
forall a.
(ArrayData a ~ ScalarArrayData a, ScalarArrayDataR a ~ a) =>
SingleArrayDict a
SingleArrayDict
    integral IntegralType a
TypeInt32  = SingleArrayDict a
forall a.
(ArrayData a ~ ScalarArrayData a, ScalarArrayDataR a ~ a) =>
SingleArrayDict a
SingleArrayDict
    integral IntegralType a
TypeInt64  = SingleArrayDict a
forall a.
(ArrayData a ~ ScalarArrayData a, ScalarArrayDataR a ~ a) =>
SingleArrayDict a
SingleArrayDict
    integral IntegralType a
TypeWord   = SingleArrayDict a
forall a.
(ArrayData a ~ ScalarArrayData a, ScalarArrayDataR a ~ a) =>
SingleArrayDict a
SingleArrayDict
    integral IntegralType a
TypeWord8  = SingleArrayDict a
forall a.
(ArrayData a ~ ScalarArrayData a, ScalarArrayDataR a ~ a) =>
SingleArrayDict a
SingleArrayDict
    integral IntegralType a
TypeWord16 = SingleArrayDict a
forall a.
(ArrayData a ~ ScalarArrayData a, ScalarArrayDataR a ~ a) =>
SingleArrayDict a
SingleArrayDict
    integral IntegralType a
TypeWord32 = SingleArrayDict a
forall a.
(ArrayData a ~ ScalarArrayData a, ScalarArrayDataR a ~ a) =>
SingleArrayDict a
SingleArrayDict
    integral IntegralType a
TypeWord64 = SingleArrayDict a
forall a.
(ArrayData a ~ ScalarArrayData a, ScalarArrayDataR a ~ a) =>
SingleArrayDict a
SingleArrayDict

    floating :: FloatingType a -> SingleArrayDict a
    floating :: FloatingType a -> SingleArrayDict a
floating FloatingType a
TypeHalf   = SingleArrayDict a
forall a.
(ArrayData a ~ ScalarArrayData a, ScalarArrayDataR a ~ a) =>
SingleArrayDict a
SingleArrayDict
    floating FloatingType a
TypeFloat  = SingleArrayDict a
forall a.
(ArrayData a ~ ScalarArrayData a, ScalarArrayDataR a ~ a) =>
SingleArrayDict a
SingleArrayDict
    floating FloatingType a
TypeDouble = SingleArrayDict a
forall a.
(ArrayData a ~ ScalarArrayData a, ScalarArrayDataR a ~ a) =>
SingleArrayDict a
SingleArrayDict


-- Array operations
-- ----------------

newArrayData :: HasCallStack => TupR ScalarType e -> Int -> IO (MutableArrayData e)
newArrayData :: TupR ScalarType e -> Int -> IO (MutableArrayData e)
newArrayData TupR ScalarType e
TupRunit         !Int
_    = () -> IO ()
forall (m :: * -> *) a. Monad m => a -> m a
return ()
newArrayData (TupRpair TupR ScalarType a
t1 TupR ScalarType b
t2) !Int
size = (,) (GArrayDataR UniqueArray a
 -> GArrayDataR UniqueArray b
 -> (GArrayDataR UniqueArray a, GArrayDataR UniqueArray b))
-> IO (GArrayDataR UniqueArray a)
-> IO
     (GArrayDataR UniqueArray b
      -> (GArrayDataR UniqueArray a, GArrayDataR UniqueArray b))
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> TupR ScalarType a -> Int -> IO (GArrayDataR UniqueArray a)
forall e.
HasCallStack =>
TupR ScalarType e -> Int -> IO (MutableArrayData e)
newArrayData TupR ScalarType a
t1 Int
size IO
  (GArrayDataR UniqueArray b
   -> (GArrayDataR UniqueArray a, GArrayDataR UniqueArray b))
-> IO (GArrayDataR UniqueArray b)
-> IO (GArrayDataR UniqueArray a, GArrayDataR UniqueArray b)
forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> TupR ScalarType b -> Int -> IO (GArrayDataR UniqueArray b)
forall e.
HasCallStack =>
TupR ScalarType e -> Int -> IO (MutableArrayData e)
newArrayData TupR ScalarType b
t2 Int
size
newArrayData (TupRsingle ScalarType e
t)   !Int
size
  | SingleScalarType SingleType e
s <- ScalarType e
t
  , SingleDict e
SingleDict         <- SingleType e -> SingleDict e
forall a. SingleType a -> SingleDict a
singleDict SingleType e
s
  , SingleArrayDict e
SingleArrayDict    <- SingleType e -> SingleArrayDict e
forall a. SingleType a -> SingleArrayDict a
singleArrayDict SingleType e
s
  = Int -> IO (UniqueArray e)
forall e. (HasCallStack, Storable e) => Int -> IO (UniqueArray e)
allocateArray Int
size
  --
  | VectorScalarType VectorType (Vec n a)
v <- ScalarType e
t
  , VectorType Int
w SingleType a
s     <- VectorType (Vec n a)
v
  , SingleDict a
SingleDict         <- SingleType a -> SingleDict a
forall a. SingleType a -> SingleDict a
singleDict SingleType a
s
  , SingleArrayDict a
SingleArrayDict    <- SingleType a -> SingleArrayDict a
forall a. SingleType a -> SingleArrayDict a
singleArrayDict SingleType a
s
  = Int -> IO (UniqueArray a)
forall e. (HasCallStack, Storable e) => Int -> IO (UniqueArray e)
allocateArray (Int
w Int -> Int -> Int
forall a. Num a => a -> a -> a
* Int
size)

indexArrayData :: TupR ScalarType e -> ArrayData e -> Int -> e
indexArrayData :: TupR ScalarType e -> ArrayData e -> Int -> e
indexArrayData TupR ScalarType e
tR ArrayData e
arr Int
ix = IO e -> e
forall a. IO a -> a
unsafePerformIO (IO e -> e) -> IO e -> e
forall a b. (a -> b) -> a -> b
$ TupR ScalarType e -> ArrayData e -> Int -> IO e
forall e. TupR ScalarType e -> MutableArrayData e -> Int -> IO e
readArrayData TupR ScalarType e
tR ArrayData e
arr Int
ix

readArrayData :: forall e. TupR ScalarType e -> MutableArrayData e -> Int -> IO e
readArrayData :: TupR ScalarType e -> MutableArrayData e -> Int -> IO e
readArrayData TupR ScalarType e
TupRunit         ()       !Int
_  = () -> IO ()
forall (m :: * -> *) a. Monad m => a -> m a
return ()
readArrayData (TupRpair TupR ScalarType a
t1 TupR ScalarType b
t2) (a1, a2) !Int
ix = (,) (a -> b -> (a, b)) -> IO a -> IO (b -> (a, b))
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> TupR ScalarType a -> MutableArrayData a -> Int -> IO a
forall e. TupR ScalarType e -> MutableArrayData e -> Int -> IO e
readArrayData TupR ScalarType a
t1 MutableArrayData a
a1 Int
ix IO (b -> (a, b)) -> IO b -> IO (a, b)
forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> TupR ScalarType b -> MutableArrayData b -> Int -> IO b
forall e. TupR ScalarType e -> MutableArrayData e -> Int -> IO e
readArrayData TupR ScalarType b
t2 MutableArrayData b
a2 Int
ix
readArrayData (TupRsingle ScalarType e
t)   MutableArrayData e
arr      !Int
ix
  | SingleScalarType SingleType e
s <- ScalarType e
t
  , SingleDict e
SingleDict         <- SingleType e -> SingleDict e
forall a. SingleType a -> SingleDict a
singleDict SingleType e
s
  , SingleArrayDict e
SingleArrayDict    <- SingleType e -> SingleArrayDict e
forall a. SingleType a -> SingleArrayDict a
singleArrayDict SingleType e
s
  = UniqueArray e -> Int -> IO e
forall e. Storable e => UniqueArray e -> Int -> IO e
unsafeReadArray UniqueArray e
MutableArrayData e
arr Int
ix
  --
  | VectorScalarType VectorType (Vec n a)
v <- ScalarType e
t
  , VectorType Int
w SingleType a
s     <- VectorType (Vec n a)
v
  , I# Int#
w#              <- Int
w
  , I# Int#
ix#             <- Int
ix
  , SingleDict a
SingleDict         <- SingleType a -> SingleDict a
forall a. SingleType a -> SingleDict a
singleDict SingleType a
s
  , SingleArrayDict a
SingleArrayDict    <- SingleType a -> SingleArrayDict a
forall a. SingleType a -> SingleArrayDict a
singleArrayDict SingleType a
s
  = let
        !bytes# :: Int#
bytes# = Int#
w# Int# -> Int# -> Int#
*# a -> Int#
forall a. Prim a => a -> Int#
sizeOf# (ScalarArrayDataR e
forall a. HasCallStack => a
undefined :: ScalarArrayDataR e)
        !addr# :: Addr#
addr#  = Ptr a -> Addr#
forall a. Ptr a -> Addr#
unPtr# (UniqueArray a -> Ptr a
forall a. UniqueArray a -> Ptr a
unsafeUniqueArrayPtr UniqueArray a
MutableArrayData e
arr) Addr# -> Int# -> Addr#
`plusAddr#` (Int#
ix# Int# -> Int# -> Int#
*# Int#
bytes#)
     in
     (State# RealWorld -> (# State# RealWorld, Vec n a #))
-> IO (Vec n a)
forall a. (State# RealWorld -> (# State# RealWorld, a #)) -> IO a
IO ((State# RealWorld -> (# State# RealWorld, Vec n a #))
 -> IO (Vec n a))
-> (State# RealWorld -> (# State# RealWorld, Vec n a #))
-> IO (Vec n a)
forall a b. (a -> b) -> a -> b
$ \State# RealWorld
s0 ->
       case Int#
-> Int#
-> State# RealWorld
-> (# State# RealWorld, MutableByteArray# RealWorld #)
forall d.
Int# -> Int# -> State# d -> (# State# d, MutableByteArray# d #)
newAlignedPinnedByteArray# Int#
bytes# Int#
16# State# RealWorld
s0     of { (# State# RealWorld
s1, MutableByteArray# RealWorld
mba# #) ->
       case Addr#
-> MutableByteArray# RealWorld
-> Int#
-> Int#
-> State# RealWorld
-> State# RealWorld
forall d.
Addr#
-> MutableByteArray# d -> Int# -> Int# -> State# d -> State# d
copyAddrToByteArray# Addr#
addr# MutableByteArray# RealWorld
mba# Int#
0# Int#
bytes# State# RealWorld
s1 of { State# RealWorld
s2             ->
       case MutableByteArray# RealWorld
-> State# RealWorld -> (# State# RealWorld, ByteArray# #)
forall d.
MutableByteArray# d -> State# d -> (# State# d, ByteArray# #)
unsafeFreezeByteArray# MutableByteArray# RealWorld
mba# State# RealWorld
s2               of { (# State# RealWorld
s3, ByteArray#
ba# #)  ->
         (# State# RealWorld
s3, ByteArray# -> Vec n a
forall (n :: Nat) a. ByteArray# -> Vec n a
Vec ByteArray#
ba# #)
       }}}

writeArrayData :: forall e. TupR ScalarType e -> MutableArrayData e -> Int -> e -> IO ()
writeArrayData :: TupR ScalarType e -> MutableArrayData e -> Int -> e -> IO ()
writeArrayData TupR ScalarType e
TupRunit         ()       !Int
_  ()       = () -> IO ()
forall (m :: * -> *) a. Monad m => a -> m a
return ()
writeArrayData (TupRpair TupR ScalarType a
t1 TupR ScalarType b
t2) (a1, a2) !Int
ix (v1, v2) = TupR ScalarType a -> MutableArrayData a -> Int -> a -> IO ()
forall e.
TupR ScalarType e -> MutableArrayData e -> Int -> e -> IO ()
writeArrayData TupR ScalarType a
t1 MutableArrayData a
a1 Int
ix a
v1 IO () -> IO () -> IO ()
forall (m :: * -> *) a b. Monad m => m a -> m b -> m b
>> TupR ScalarType b -> MutableArrayData b -> Int -> b -> IO ()
forall e.
TupR ScalarType e -> MutableArrayData e -> Int -> e -> IO ()
writeArrayData TupR ScalarType b
t2 MutableArrayData b
a2 Int
ix b
v2
writeArrayData (TupRsingle ScalarType e
t)   MutableArrayData e
arr      !Int
ix !e
val
  | SingleScalarType SingleType e
s <- ScalarType e
t
  , SingleDict e
SingleDict         <- SingleType e -> SingleDict e
forall a. SingleType a -> SingleDict a
singleDict SingleType e
s
  , SingleArrayDict e
SingleArrayDict    <- SingleType e -> SingleArrayDict e
forall a. SingleType a -> SingleArrayDict a
singleArrayDict SingleType e
s
  = UniqueArray e -> Int -> e -> IO ()
forall e. Storable e => UniqueArray e -> Int -> e -> IO ()
unsafeWriteArray UniqueArray e
MutableArrayData e
arr Int
ix e
val
  --
  | VectorScalarType VectorType (Vec n a)
v <- ScalarType e
t
  , VectorType Int
w SingleType a
s     <- VectorType (Vec n a)
v
  , Vec ba#            <- e
val
  , I# Int#
w#              <- Int
w
  , I# Int#
ix#             <- Int
ix
  , SingleDict a
SingleDict         <- SingleType a -> SingleDict a
forall a. SingleType a -> SingleDict a
singleDict SingleType a
s
  , SingleArrayDict a
SingleArrayDict    <- SingleType a -> SingleArrayDict a
forall a. SingleType a -> SingleArrayDict a
singleArrayDict SingleType a
s
  = let
       !bytes# :: Int#
bytes# = Int#
w# Int# -> Int# -> Int#
*# a -> Int#
forall a. Prim a => a -> Int#
sizeOf# (ScalarArrayDataR e
forall a. HasCallStack => a
undefined :: ScalarArrayDataR e)
       !addr# :: Addr#
addr#  = Ptr a -> Addr#
forall a. Ptr a -> Addr#
unPtr# (UniqueArray a -> Ptr a
forall a. UniqueArray a -> Ptr a
unsafeUniqueArrayPtr UniqueArray a
MutableArrayData e
arr) Addr# -> Int# -> Addr#
`plusAddr#` (Int#
ix# Int# -> Int# -> Int#
*# Int#
bytes#)
     in
     (State# RealWorld -> (# State# RealWorld, () #)) -> IO ()
forall a. (State# RealWorld -> (# State# RealWorld, a #)) -> IO a
IO ((State# RealWorld -> (# State# RealWorld, () #)) -> IO ())
-> (State# RealWorld -> (# State# RealWorld, () #)) -> IO ()
forall a b. (a -> b) -> a -> b
$ \State# RealWorld
s0 -> case ByteArray#
-> Int# -> Addr# -> Int# -> State# RealWorld -> State# RealWorld
forall d.
ByteArray# -> Int# -> Addr# -> Int# -> State# d -> State# d
copyByteArrayToAddr# ByteArray#
ba# Int#
0# Addr#
addr# Int#
bytes# State# RealWorld
s0 of
                   State# RealWorld
s1 -> (# State# RealWorld
s1, () #)


unsafeArrayDataPtr :: ScalarType e -> ArrayData e -> Ptr (ScalarArrayDataR e)
unsafeArrayDataPtr :: ScalarType e -> ArrayData e -> Ptr (ScalarArrayDataR e)
unsafeArrayDataPtr ScalarType e
t ArrayData e
arr
  | ScalarArrayDict{} <- ScalarType e -> ScalarArrayDict e
forall a. ScalarType a -> ScalarArrayDict a
scalarArrayDict ScalarType e
t
  = UniqueArray (ScalarArrayDataR b) -> Ptr (ScalarArrayDataR b)
forall a. UniqueArray a -> Ptr a
unsafeUniqueArrayPtr UniqueArray (ScalarArrayDataR b)
ArrayData e
arr

touchArrayData :: TupR ScalarType e -> ArrayData e -> IO ()
touchArrayData :: TupR ScalarType e -> ArrayData e -> IO ()
touchArrayData TupR ScalarType e
TupRunit         ()       = () -> IO ()
forall (m :: * -> *) a. Monad m => a -> m a
return ()
touchArrayData (TupRpair TupR ScalarType a
t1 TupR ScalarType b
t2) (a1, a2) = TupR ScalarType a -> ArrayData a -> IO ()
forall e. TupR ScalarType e -> ArrayData e -> IO ()
touchArrayData TupR ScalarType a
t1 ArrayData a
a1 IO () -> IO () -> IO ()
forall (m :: * -> *) a b. Monad m => m a -> m b -> m b
>> TupR ScalarType b -> ArrayData b -> IO ()
forall e. TupR ScalarType e -> ArrayData e -> IO ()
touchArrayData TupR ScalarType b
t2 ArrayData b
a2
touchArrayData (TupRsingle ScalarType e
t)   ArrayData e
arr
  | ScalarArrayDict{} <- ScalarType e -> ScalarArrayDict e
forall a. ScalarType a -> ScalarArrayDict a
scalarArrayDict ScalarType e
t
  = UniqueArray (ScalarArrayDataR b) -> IO ()
forall a. UniqueArray a -> IO ()
touchUniqueArray UniqueArray (ScalarArrayDataR b)
ArrayData e
arr

rnfArrayData :: TupR ScalarType e -> ArrayData e -> ()
rnfArrayData :: TupR ScalarType e -> ArrayData e -> ()
rnfArrayData TupR ScalarType e
TupRunit         ()       = ()
rnfArrayData (TupRpair TupR ScalarType a
t1 TupR ScalarType b
t2) (a1, a2) = TupR ScalarType a -> ArrayData a -> ()
forall e. TupR ScalarType e -> ArrayData e -> ()
rnfArrayData TupR ScalarType a
t1 ArrayData a
a1 () -> () -> ()
`seq` TupR ScalarType b -> ArrayData b -> ()
forall e. TupR ScalarType e -> ArrayData e -> ()
rnfArrayData TupR ScalarType b
t2 ArrayData b
a2 () -> () -> ()
`seq` ()
rnfArrayData (TupRsingle ScalarType e
t)   ArrayData e
arr      = Ptr (ScalarArrayDataR e) -> ()
forall a. NFData a => a -> ()
rnf (ScalarType e -> ArrayData e -> Ptr (ScalarArrayDataR e)
forall e. ScalarType e -> ArrayData e -> Ptr (ScalarArrayDataR e)
unsafeArrayDataPtr ScalarType e
t ArrayData e
arr)

unPtr# :: Ptr a -> Addr#
unPtr# :: Ptr a -> Addr#
unPtr# (Ptr Addr#
addr#) = Addr#
addr#

-- | Safe combination of creating and fast freezing of array data.
--
runArrayData
    :: IO (MutableArrayData e, e)
    -> (ArrayData e, e)
runArrayData :: IO (MutableArrayData e, e) -> (MutableArrayData e, e)
runArrayData IO (MutableArrayData e, e)
st = IO (MutableArrayData e, e) -> (MutableArrayData e, e)
forall a. IO a -> a
unsafePerformIO (IO (MutableArrayData e, e) -> (MutableArrayData e, e))
-> IO (MutableArrayData e, e) -> (MutableArrayData e, e)
forall a b. (a -> b) -> a -> b
$ do
  (MutableArrayData e
mad, e
r) <- IO (MutableArrayData e, e)
st
  (MutableArrayData e, e) -> IO (MutableArrayData e, e)
forall (m :: * -> *) a. Monad m => a -> m a
return (MutableArrayData e
mad, e
r)

-- Allocate a new array with enough storage to hold the given number of
-- elements.
--
-- The array is uninitialised and, in particular, allocated lazily. The latter
-- is important because it means that for backends that have discrete memory
-- spaces (e.g. GPUs), we will not increase host memory pressure simply to track
-- intermediate arrays that contain meaningful data only on the device.
--
allocateArray :: forall e. (HasCallStack, Storable e) => Int -> IO (UniqueArray e)
allocateArray :: Int -> IO (UniqueArray e)
allocateArray !Int
size
  = String -> Bool -> IO (UniqueArray e) -> IO (UniqueArray e)
forall a. HasCallStack => String -> Bool -> a -> a
internalCheck String
"size must be >= 0" (Int
size Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
>= Int
0)
  (IO (UniqueArray e) -> IO (UniqueArray e))
-> IO (UniqueArray e) -> IO (UniqueArray e)
forall a b. (a -> b) -> a -> b
$ ForeignPtr e -> IO (UniqueArray e)
forall e. ForeignPtr e -> IO (UniqueArray e)
newUniqueArray (ForeignPtr e -> IO (UniqueArray e))
-> (IO (ForeignPtr e) -> IO (ForeignPtr e))
-> IO (ForeignPtr e)
-> IO (UniqueArray e)
forall (m :: * -> *) b c a.
Monad m =>
(b -> m c) -> (a -> m b) -> a -> m c
<=< IO (ForeignPtr e) -> IO (ForeignPtr e)
forall a. IO a -> IO a
unsafeInterleaveIO (IO (ForeignPtr e) -> IO (UniqueArray e))
-> IO (ForeignPtr e) -> IO (UniqueArray e)
forall a b. (a -> b) -> a -> b
$ do
      let bytes :: Int
bytes = Int
size Int -> Int -> Int
forall a. Num a => a -> a -> a
* e -> Int
forall a. Storable a => a -> Int
sizeOf (e
forall a. HasCallStack => a
undefined :: e)
      Int -> IO (ForeignPtr Word8)
new <- IORef (Int -> IO (ForeignPtr Word8))
-> IO (Int -> IO (ForeignPtr Word8))
forall a. IORef a -> IO a
readIORef IORef (Int -> IO (ForeignPtr Word8))
__mallocForeignPtrBytes
      ForeignPtr Word8
ptr <- Int -> IO (ForeignPtr Word8)
new Int
bytes
      Flag -> String -> IO ()
traceIO Flag
dump_gc (String -> IO ()) -> String -> IO ()
forall a b. (a -> b) -> a -> b
$ String -> Int -> String -> String
forall r. PrintfType r => String -> r
printf String
"gc: allocated new host array (size=%d, ptr=%s)" Int
bytes (ForeignPtr Word8 -> String
forall a. Show a => a -> String
show ForeignPtr Word8
ptr)
      Int64 -> IO ()
didAllocateBytesLocal (Int -> Int64
forall a b. (Integral a, Num b) => a -> b
fromIntegral Int
bytes)
      ForeignPtr e -> IO (ForeignPtr e)
forall (m :: * -> *) a. Monad m => a -> m a
return (ForeignPtr Word8 -> ForeignPtr e
forall a b. ForeignPtr a -> ForeignPtr b
castForeignPtr ForeignPtr Word8
ptr)

-- | Register the given function as the callback to use to allocate new array
-- data on the host containing the specified number of bytes. The returned array
-- must be pinned (with respect to Haskell's GC), so that it can be passed to
-- foreign code.
--
registerForeignPtrAllocator
    :: (Int -> IO (ForeignPtr Word8))
    -> IO ()
registerForeignPtrAllocator :: (Int -> IO (ForeignPtr Word8)) -> IO ()
registerForeignPtrAllocator Int -> IO (ForeignPtr Word8)
new = do
  Flag -> String -> IO ()
traceIO Flag
dump_gc String
"registering new array allocator"
  IORef (Int -> IO (ForeignPtr Word8))
-> (Int -> IO (ForeignPtr Word8)) -> IO ()
forall a. IORef a -> a -> IO ()
atomicWriteIORef IORef (Int -> IO (ForeignPtr Word8))
__mallocForeignPtrBytes Int -> IO (ForeignPtr Word8)
new

{-# NOINLINE __mallocForeignPtrBytes #-}
__mallocForeignPtrBytes :: IORef (Int -> IO (ForeignPtr Word8))
__mallocForeignPtrBytes :: IORef (Int -> IO (ForeignPtr Word8))
__mallocForeignPtrBytes = IO (IORef (Int -> IO (ForeignPtr Word8)))
-> IORef (Int -> IO (ForeignPtr Word8))
forall a. IO a -> a
unsafePerformIO (IO (IORef (Int -> IO (ForeignPtr Word8)))
 -> IORef (Int -> IO (ForeignPtr Word8)))
-> IO (IORef (Int -> IO (ForeignPtr Word8)))
-> IORef (Int -> IO (ForeignPtr Word8))
forall a b. (a -> b) -> a -> b
$! (Int -> IO (ForeignPtr Word8))
-> IO (IORef (Int -> IO (ForeignPtr Word8)))
forall a. a -> IO (IORef a)
newIORef Int -> IO (ForeignPtr Word8)
forall a. Int -> IO (ForeignPtr a)
mallocPlainForeignPtrBytesAligned

-- | Allocate the given number of bytes with 16-byte alignment. This is
-- essential for SIMD instructions.
--
-- Additionally, we return a plain ForeignPtr, which unlike a regular ForeignPtr
-- created with 'mallocForeignPtr' carries no finalisers. It is an error to try
-- to add a finaliser to the plain ForeignPtr. For our purposes this is fine,
-- since in Accelerate finalisers are handled using Lifetime
--
mallocPlainForeignPtrBytesAligned :: Int -> IO (ForeignPtr a)
mallocPlainForeignPtrBytesAligned :: Int -> IO (ForeignPtr a)
mallocPlainForeignPtrBytesAligned (I# Int#
size) = (State# RealWorld -> (# State# RealWorld, ForeignPtr a #))
-> IO (ForeignPtr a)
forall a. (State# RealWorld -> (# State# RealWorld, a #)) -> IO a
IO ((State# RealWorld -> (# State# RealWorld, ForeignPtr a #))
 -> IO (ForeignPtr a))
-> (State# RealWorld -> (# State# RealWorld, ForeignPtr a #))
-> IO (ForeignPtr a)
forall a b. (a -> b) -> a -> b
$ \State# RealWorld
s ->
  case Int#
-> Int#
-> State# RealWorld
-> (# State# RealWorld, MutableByteArray# RealWorld #)
forall d.
Int# -> Int# -> State# d -> (# State# d, MutableByteArray# d #)
newAlignedPinnedByteArray# Int#
size Int#
64# State# RealWorld
s of
    (# State# RealWorld
s', MutableByteArray# RealWorld
mbarr# #) -> (# State# RealWorld
s', Addr# -> ForeignPtrContents -> ForeignPtr a
forall a. Addr# -> ForeignPtrContents -> ForeignPtr a
ForeignPtr (ByteArray# -> Addr#
byteArrayContents# (MutableByteArray# RealWorld -> ByteArray#
unsafeCoerce# MutableByteArray# RealWorld
mbarr#)) (MutableByteArray# RealWorld -> ForeignPtrContents
PlainPtr MutableByteArray# RealWorld
mbarr#) #)


liftArrayData :: Int -> TypeR e -> ArrayData e -> Q (TExp (ArrayData e))
liftArrayData :: Int -> TypeR e -> ArrayData e -> Q (TExp (ArrayData e))
liftArrayData Int
n = TypeR e -> ArrayData e -> Q (TExp (ArrayData e))
forall e. TypeR e -> ArrayData e -> Q (TExp (ArrayData e))
tuple
  where
    tuple :: TypeR e -> ArrayData e -> Q (TExp (ArrayData e))
    tuple :: TypeR e -> ArrayData e -> Q (TExp (ArrayData e))
tuple TypeR e
TupRunit         ()       = [|| () ||]
    tuple (TupRpair TupR ScalarType a
t1 TupR ScalarType b
t2) (a1, a2) = [|| ($$(tuple t1 a1), $$(tuple t2 a2)) ||]
    tuple (TupRsingle ScalarType e
s) ArrayData e
adata      = ScalarType e -> ArrayData e -> Q (TExp (ArrayData e))
forall e. ScalarType e -> ArrayData e -> Q (TExp (ArrayData e))
scalar ScalarType e
s ArrayData e
adata

    scalar :: ScalarType e -> ArrayData e -> Q (TExp (ArrayData e))
    scalar :: ScalarType e -> ArrayData e -> Q (TExp (ArrayData e))
scalar (SingleScalarType SingleType e
t) = SingleType e -> ArrayData e -> Q (TExp (ArrayData e))
forall e. SingleType e -> ArrayData e -> Q (TExp (ArrayData e))
single SingleType e
t
    scalar (VectorScalarType VectorType (Vec n a)
t) = VectorType (Vec n a)
-> ArrayData (Vec n a) -> Q (TExp (ArrayData (Vec n a)))
forall (n :: Nat) e.
VectorType (Vec n e)
-> ArrayData (Vec n e) -> Q (TExp (ArrayData (Vec n e)))
vector VectorType (Vec n a)
t

    vector :: forall n e. VectorType (Vec n e) -> ArrayData (Vec n e) -> Q (TExp (ArrayData (Vec n e)))
    vector :: VectorType (Vec n e)
-> ArrayData (Vec n e) -> Q (TExp (ArrayData (Vec n e)))
vector (VectorType Int
w SingleType a
t)
      | SingleArrayDict a
SingleArrayDict <- SingleType a -> SingleArrayDict a
forall a. SingleType a -> SingleArrayDict a
singleArrayDict SingleType a
t
      = Int -> TypeR a -> ArrayData a -> Q (TExp (ArrayData a))
forall e. Int -> TypeR e -> ArrayData e -> Q (TExp (ArrayData e))
liftArrayData (Int
w Int -> Int -> Int
forall a. Num a => a -> a -> a
* Int
n) (ScalarType a -> TypeR a
forall (s :: * -> *) a. s a -> TupR s a
TupRsingle (SingleType a -> ScalarType a
forall a. SingleType a -> ScalarType a
SingleScalarType SingleType a
t))

    single :: SingleType e -> ArrayData e -> Q (TExp (ArrayData e))
    single :: SingleType e -> ArrayData e -> Q (TExp (ArrayData e))
single (NumSingleType NumType e
t) = NumType e -> ArrayData e -> Q (TExp (ArrayData e))
forall e. NumType e -> ArrayData e -> Q (TExp (ArrayData e))
num NumType e
t

    num :: NumType e -> ArrayData e -> Q (TExp (ArrayData e))
    num :: NumType e -> ArrayData e -> Q (TExp (ArrayData e))
num (IntegralNumType IntegralType e
t) = IntegralType e -> ArrayData e -> Q (TExp (ArrayData e))
forall e. IntegralType e -> ArrayData e -> Q (TExp (ArrayData e))
integral IntegralType e
t
    num (FloatingNumType FloatingType e
t) = FloatingType e -> ArrayData e -> Q (TExp (ArrayData e))
forall e. FloatingType e -> ArrayData e -> Q (TExp (ArrayData e))
floating FloatingType e
t

    integral :: IntegralType e -> ArrayData e -> Q (TExp (ArrayData e))
    integral :: IntegralType e -> ArrayData e -> Q (TExp (ArrayData e))
integral IntegralType e
TypeInt    = Int -> UniqueArray Int -> Q (TExp (UniqueArray Int))
forall a.
Storable a =>
Int -> UniqueArray a -> Q (TExp (UniqueArray a))
liftUniqueArray Int
n
    integral IntegralType e
TypeInt8   = Int -> UniqueArray Int8 -> Q (TExp (UniqueArray Int8))
forall a.
Storable a =>
Int -> UniqueArray a -> Q (TExp (UniqueArray a))
liftUniqueArray Int
n
    integral IntegralType e
TypeInt16  = Int -> UniqueArray Int16 -> Q (TExp (UniqueArray Int16))
forall a.
Storable a =>
Int -> UniqueArray a -> Q (TExp (UniqueArray a))
liftUniqueArray Int
n
    integral IntegralType e
TypeInt32  = Int -> UniqueArray Int32 -> Q (TExp (UniqueArray Int32))
forall a.
Storable a =>
Int -> UniqueArray a -> Q (TExp (UniqueArray a))
liftUniqueArray Int
n
    integral IntegralType e
TypeInt64  = Int -> UniqueArray Int64 -> Q (TExp (UniqueArray Int64))
forall a.
Storable a =>
Int -> UniqueArray a -> Q (TExp (UniqueArray a))
liftUniqueArray Int
n
    integral IntegralType e
TypeWord   = Int -> UniqueArray Word -> Q (TExp (UniqueArray Word))
forall a.
Storable a =>
Int -> UniqueArray a -> Q (TExp (UniqueArray a))
liftUniqueArray Int
n
    integral IntegralType e
TypeWord8  = Int -> UniqueArray Word8 -> Q (TExp (UniqueArray Word8))
forall a.
Storable a =>
Int -> UniqueArray a -> Q (TExp (UniqueArray a))
liftUniqueArray Int
n
    integral IntegralType e
TypeWord16 = Int -> UniqueArray Word16 -> Q (TExp (UniqueArray Word16))
forall a.
Storable a =>
Int -> UniqueArray a -> Q (TExp (UniqueArray a))
liftUniqueArray Int
n
    integral IntegralType e
TypeWord32 = Int -> UniqueArray Word32 -> Q (TExp (UniqueArray Word32))
forall a.
Storable a =>
Int -> UniqueArray a -> Q (TExp (UniqueArray a))
liftUniqueArray Int
n
    integral IntegralType e
TypeWord64 = Int -> UniqueArray Word64 -> Q (TExp (UniqueArray Word64))
forall a.
Storable a =>
Int -> UniqueArray a -> Q (TExp (UniqueArray a))
liftUniqueArray Int
n

    floating :: FloatingType e -> ArrayData e -> Q (TExp (ArrayData e))
    floating :: FloatingType e -> ArrayData e -> Q (TExp (ArrayData e))
floating FloatingType e
TypeHalf   = Int -> UniqueArray Half -> Q (TExp (UniqueArray Half))
forall a.
Storable a =>
Int -> UniqueArray a -> Q (TExp (UniqueArray a))
liftUniqueArray Int
n
    floating FloatingType e
TypeFloat  = Int -> UniqueArray Float -> Q (TExp (UniqueArray Float))
forall a.
Storable a =>
Int -> UniqueArray a -> Q (TExp (UniqueArray a))
liftUniqueArray Int
n
    floating FloatingType e
TypeDouble = Int -> UniqueArray Double -> Q (TExp (UniqueArray Double))
forall a.
Storable a =>
Int -> UniqueArray a -> Q (TExp (UniqueArray a))
liftUniqueArray Int
n

-- Determine the underlying type of a Haskell CLong or CULong.
--
runQ [d| type HTYPE_INT = $(
              case finiteBitSize (undefined::Int) of
                32 -> [t| Int32 |]
                64 -> [t| Int64 |]
                _  -> error "I don't know what architecture I am" ) |]

runQ [d| type HTYPE_WORD = $(
              case finiteBitSize (undefined::Word) of
                32 -> [t| Word32 |]
                64 -> [t| Word64 |]
                _  -> error "I don't know what architecture I am" ) |]

runQ [d| type HTYPE_CLONG = $(
              case finiteBitSize (undefined::CLong) of
                32 -> [t| Int32 |]
                64 -> [t| Int64 |]
                _  -> error "I don't know what architecture I am" ) |]

runQ [d| type HTYPE_CULONG = $(
              case finiteBitSize (undefined::CULong) of
                32 -> [t| Word32 |]
                64 -> [t| Word64 |]
                _  -> error "I don't know what architecture I am" ) |]

runQ [d| type HTYPE_CCHAR = $(
              if isSigned (undefined::CChar)
                then [t| Int8  |]
                else [t| Word8 |] ) |]