accelerate-cufft-0.0.1: example/Merged.hs

{- |
Simple manual implementation
of embedding cufft functionality in the @accelerate@ framework.
In this example, a plan is created for every transform
and is run within 'CUDA.run1'.
-}
module Main where

import qualified Data.Array.Accelerate.CUDA.Foreign as AF
import qualified Data.Array.Accelerate.CUDA as CUDA

import qualified Foreign.CUDA.FFT as CUFFT

import qualified Data.Array.Accelerate as A
import Data.Array.Accelerate (Acc, Vector, Z(Z), (:.)((:.)), )


transformForeign :: Vector Float -> AF.CIO (Vector Float)
transformForeign input =
   let (Z:.inlen) = A.arrayShape input
       outlen = (div inlen 2 + 1) * 2
   in  do
       output <- AF.allocateArray (Z:.outlen)
       ((), iptr) <- AF.devicePtrsOfArray input
       ((), optr) <- AF.devicePtrsOfArray output
       AF.liftIO $ do
          h <- CUFFT.plan1D inlen CUFFT.R2C 1
          CUFFT.execR2C h iptr optr
       return output

transform :: Acc (Vector Float) -> Acc (Vector Float)
transform =
   A.foreignAcc
      (AF.CUDAForeignAcc "transformForeign" $ const transformForeign)
      (error "no fft fallback implemented")

main :: IO ()
main =
   print $ CUDA.run1 transform $ A.fromList (Z:.5) [1,0,0,0,0]