accelerate-cufft-0.0: example/Separate.hs

{- |
Simple manual implementation
of embedding cufft functionality in the @accelerate@ framework.
In this example, a plan is created once globally
and must be run 'inDefaultContext'.
-}
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 :: CUFFT.Handle -> Vector Float -> AF.CIO (Vector Float)
transformForeign h 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 $ CUFFT.execR2C h iptr optr
       return output

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

main :: IO ()
main = do
   let inlen = 5
   h <- AF.inDefaultContext $ CUFFT.plan1D inlen CUFFT.R2C 1
   print $ CUDA.run1 (transform h) $ A.fromList (Z:.inlen) [1,0,0,0,0]