{-# LANGUAGE GADTs               #-}
{-# LANGUAGE MagicHash           #-}
{-# LANGUAGE PatternGuards       #-}
{-# LANGUAGE ScopedTypeVariables #-}
{-# LANGUAGE TypeApplications    #-}
{-# LANGUAGE TypeOperators       #-}
-- |
-- Module      : Data.Array.Accelerate.Math.FFT.LLVM.Native.Base
-- Copyright   : [2017..2020] The Accelerate Team
-- License     : BSD3
--
-- Maintainer  : Trevor L. McDonell <trevor.mcdonell@gmail.com>
-- Stability   : experimental
-- Portability : non-portable (GHC extensions)
--

module Data.Array.Accelerate.Math.FFT.LLVM.Native.Base
  where

import Data.Array.Accelerate.Array.Data
import Data.Array.Accelerate.Array.Unique
import Data.Array.Accelerate.Data.Complex
import Data.Array.Accelerate.Lifetime
import Data.Array.Accelerate.Representation.Array
import Data.Array.Accelerate.Representation.Shape
import Data.Array.Accelerate.Sugar.Elt
import Data.Primitive.Vec

import Data.Array.Accelerate.Math.FFT.Mode
import Data.Array.Accelerate.Math.FFT.Type
import Data.Array.Accelerate.Math.FFT.LLVM.Native.Ix

import Data.Array.CArray.Base                                       ( CArray(..) )
import Math.FFT.Base                                                ( Sign(..), Flag, measure, preserveInput )

import Data.Bits
import Foreign.ForeignPtr
import Text.Printf
import Prelude                                                      as P


signOf :: Mode -> Sign
signOf :: Mode -> Sign
signOf Mode
Forward = Sign
DFTForward
signOf Mode
_       = Sign
DFTBackward

flags :: Flag
flags :: Flag
flags = Flag
measure Flag -> Flag -> Flag
forall a. Bits a => a -> a -> a
.|. Flag
preserveInput

nameOf :: forall sh. Mode -> ShapeR sh -> String
nameOf :: Mode -> ShapeR sh -> String
nameOf Mode
Forward ShapeR sh
shR = String -> Int -> String
forall r. PrintfType r => String -> r
printf String
"FFTW.dft%dD"  (ShapeR sh -> Int
forall sh. ShapeR sh -> Int
rank ShapeR sh
shR)
nameOf Mode
_       ShapeR sh
shR = String -> Int -> String
forall r. PrintfType r => String -> r
printf String
"FFTW.idft%dD" (ShapeR sh -> Int
forall sh. ShapeR sh -> Int
rank ShapeR sh
shR)


-- /O(1)/ Convert a CArray to an Accelerate array
--
{-# INLINE fromCArray #-}
fromCArray
    :: forall ix sh e. (IxShapeR (EltR ix) ~ sh, Elt ix)
    => ShapeR sh
    -> NumericR e
    -> CArray ix (Complex e)
    -> IO (Array sh (Vec2 e))
fromCArray :: ShapeR sh
-> NumericR e -> CArray ix (Complex e) -> IO (Array sh (Vec2 e))
fromCArray ShapeR sh
shR NumericR e
eR (CArray ix
lo ix
hi Int
_ ForeignPtr (Complex e)
fp) = do
  --
  sh
sh <- sh -> IO sh
forall (m :: * -> *) a. Monad m => a -> m a
return (sh -> IO sh) -> sh -> IO sh
forall a b. (a -> b) -> a -> b
$ ShapeR sh -> (sh, sh) -> sh
forall sh. ShapeR sh -> (sh, sh) -> sh
rangeToShape ShapeR sh
shR (ix -> sh
forall ix sh.
(HasCallStack, IxShapeR (EltR ix) ~ sh, Elt ix) =>
ix -> sh
toIxShapeR ix
lo, ix -> sh
forall ix sh.
(HasCallStack, IxShapeR (EltR ix) ~ sh, Elt ix) =>
ix -> sh
toIxShapeR ix
hi) :: IO sh
  UniqueArray e
ua <- ForeignPtr e -> IO (UniqueArray e)
forall e. ForeignPtr e -> IO (UniqueArray e)
newUniqueArray (ForeignPtr (Complex e) -> ForeignPtr e
forall a b. ForeignPtr a -> ForeignPtr b
castForeignPtr ForeignPtr (Complex e)
fp :: ForeignPtr e)
  --
  case NumericR e
eR of
    NumericR e
NumericRfloat32 -> Array sh (Vec2 e) -> IO (Array sh (Vec2 e))
forall (m :: * -> *) a. Monad m => a -> m a
return (Array sh (Vec2 e) -> IO (Array sh (Vec2 e)))
-> Array sh (Vec2 e) -> IO (Array sh (Vec2 e))
forall a b. (a -> b) -> a -> b
$ sh -> ArrayData (Vec2 e) -> Array sh (Vec2 e)
forall sh e. sh -> ArrayData e -> Array sh e
Array sh
sh ArrayData (Vec2 e)
UniqueArray e
ua
    NumericR e
NumericRfloat64 -> Array sh (Vec2 e) -> IO (Array sh (Vec2 e))
forall (m :: * -> *) a. Monad m => a -> m a
return (Array sh (Vec2 e) -> IO (Array sh (Vec2 e)))
-> Array sh (Vec2 e) -> IO (Array sh (Vec2 e))
forall a b. (a -> b) -> a -> b
$ sh -> ArrayData (Vec2 e) -> Array sh (Vec2 e)
forall sh e. sh -> ArrayData e -> Array sh e
Array sh
sh ArrayData (Vec2 e)
UniqueArray e
ua

-- /O(1)/ Use an Accelerate array as a CArray
--
{-# INLINE withCArray #-}
withCArray
    :: forall ix sh e a. (IxShapeR (EltR ix) ~ sh, Elt ix)
    => ShapeR sh
    -> NumericR e
    -> Array sh (Vec2 e)
    -> (CArray ix (Complex e) -> IO a)
    -> IO a
withCArray :: ShapeR sh
-> NumericR e
-> Array sh (Vec2 e)
-> (CArray ix (Complex e) -> IO a)
-> IO a
withCArray ShapeR sh
shR NumericR e
eR (Array sh
sh ArrayData (Vec2 e)
adata) CArray ix (Complex e) -> IO a
f =
  let
      (sh
lo, sh
hi)  = ShapeR sh -> sh -> (sh, sh)
forall sh. ShapeR sh -> sh -> (sh, sh)
shapeToRange ShapeR sh
shR sh
sh
      wrap :: ForeignPtr e -> CArray ix (Complex e)
wrap ForeignPtr e
fp   = ix -> ix -> Int -> ForeignPtr (Complex e) -> CArray ix (Complex e)
forall i e. i -> i -> Int -> ForeignPtr e -> CArray i e
CArray (sh -> ix
forall ix sh.
(HasCallStack, IxShapeR (EltR ix) ~ sh, Elt ix) =>
sh -> ix
fromIxShapeR sh
lo) (sh -> ix
forall ix sh.
(HasCallStack, IxShapeR (EltR ix) ~ sh, Elt ix) =>
sh -> ix
fromIxShapeR sh
hi) (ShapeR sh -> sh -> Int
forall sh. ShapeR sh -> sh -> Int
size ShapeR sh
shR sh
sh) (ForeignPtr e -> ForeignPtr (Complex e)
forall a b. ForeignPtr a -> ForeignPtr b
castForeignPtr ForeignPtr e
fp)
  in
  NumericR e -> ArrayData (Vec2 e) -> (ForeignPtr e -> IO a) -> IO a
forall e a.
NumericR e -> ArrayData (Vec2 e) -> (ForeignPtr e -> IO a) -> IO a
withArrayData NumericR e
eR ArrayData (Vec2 e)
adata (CArray ix (Complex e) -> IO a
f (CArray ix (Complex e) -> IO a)
-> (ForeignPtr e -> CArray ix (Complex e)) -> ForeignPtr e -> IO a
forall b c a. (b -> c) -> (a -> b) -> a -> c
. ForeignPtr e -> CArray ix (Complex e)
wrap)


-- Use underlying array pointers
--
{-# INLINE withArrayData #-}
withArrayData
    :: NumericR e
    -> ArrayData (Vec2 e)
    -> (ForeignPtr e -> IO a)
    -> IO a
withArrayData :: NumericR e -> ArrayData (Vec2 e) -> (ForeignPtr e -> IO a) -> IO a
withArrayData NumericR e
NumericRfloat32 ArrayData (Vec2 e)
ua = Lifetime (ForeignPtr e) -> (ForeignPtr e -> IO a) -> IO a
forall a b. Lifetime a -> (a -> IO b) -> IO b
withLifetime (UniqueArray e -> Lifetime (ForeignPtr e)
forall e. UniqueArray e -> Lifetime (ForeignPtr e)
uniqueArrayData ArrayData (Vec2 e)
UniqueArray e
ua)
withArrayData NumericR e
NumericRfloat64 ArrayData (Vec2 e)
ua = Lifetime (ForeignPtr e) -> (ForeignPtr e -> IO a) -> IO a
forall a b. Lifetime a -> (a -> IO b) -> IO b
withLifetime (UniqueArray e -> Lifetime (ForeignPtr e)
forall e. UniqueArray e -> Lifetime (ForeignPtr e)
uniqueArrayData ArrayData (Vec2 e)
UniqueArray e
ua)