{-# LANGUAGE BangPatterns             #-}
{-# LANGUAGE CPP                      #-}
{-# LANGUAGE FlexibleInstances        #-}
{-# LANGUAGE ForeignFunctionInterface #-}
{-# LANGUAGE GADTs                    #-}
{-# LANGUAGE ImpredicativeTypes       #-}
{-# LANGUAGE QuasiQuotes              #-}
{-# LANGUAGE RankNTypes               #-}
{-# LANGUAGE ScopedTypeVariables      #-}
{-# LANGUAGE TemplateHaskell          #-}
{-# LANGUAGE TypeFamilies             #-}
{-# LANGUAGE ViewPatterns             #-}
{-# OPTIONS_GHC -fno-warn-name-shadowing #-}
{-# OPTIONS_GHC -fno-warn-orphans        #-}
-- |
-- Module      : Data.Array.Accelerate.CUDA.Foreign.Export
-- Copyright   : [2013..2014] Manuel M T Chakravarty, Gabriele Keller, Trevor L. McDonell, Robert Clifton-Everest
-- License     : BSD3
--
-- Maintainer  : Robert Clifton-Everest <robertce@cse.unsw.edu.au>
-- Stability   : experimental
-- Portability : non-portable (GHC extensions)
--
-- This module is intended to allow the calling of Accelerate functions from
-- within CUDA C/C++code. See the nbody/visualizer example in the accelerate-examples
-- package to see how it is used.
--

module Data.Array.Accelerate.CUDA.Foreign.Export (

  -- ** Functions callable from foreign code
  -- In order to call these from from C, see the corresponding C function signature.
  accelerateCreate, accelerateDestroy, freeOutput, freeProgram,

  -- ** Exporting
  exportAfun, buildExported,

  -- ** Types
  InputArray, OutputArray, ShapeBuffer, DevicePtrBuffer,

) where

import Data.Functor
import Control.Applicative
import Control.Monad.State                              ( liftIO )
import Foreign.StablePtr
import Foreign.C.Types
import Foreign.Ptr
import Foreign.Storable                                 ( Storable(..) )
import Foreign.Marshal.Array                            ( peekArray, pokeArray, mallocArray )
import Foreign.Marshal.Alloc                            ( free )
import Language.Haskell.TH                              hiding ( ppr )
import qualified Foreign.CUDA.Driver                    as CUDA

import Prelude                                          as P

-- friends
import Data.Array.Accelerate.Smart                      ( Acc )
import Data.Array.Accelerate.Type
import Data.Array.Accelerate.Array.Data
import Data.Array.Accelerate.CUDA                       ( run1With )
import Data.Array.Accelerate.CUDA.Array.Sugar           hiding ( shape, size )
import Data.Array.Accelerate.CUDA.Array.Data            hiding ( pokeArray, peekArray, mallocArray )
import Data.Array.Accelerate.CUDA.State
import Data.Array.Accelerate.CUDA.Context

-- |A handle foreign code can use to call accelerate functions.
type AccHandle = StablePtr Context

-- |A foreign buffer that represents a shape as an array of ints.
type ShapeBuffer = Ptr CInt

-- |A buffer of device pointers
type DevicePtrBuffer = Ptr WordPtr

-- |The input required from foreign code.
type InputArray = (ShapeBuffer, DevicePtrBuffer)

-- |A result array from an accelerate program.
type OutputArray = (ShapeBuffer, DevicePtrBuffer, StablePtr EArray)

-- |Foreign exportable representation of a CUDA device
type Device = Int32

-- |Foreign representation of a CUDA context.
type ForeignContext = Ptr ()

-- We need to capture the Arrays constraint
data Afun where
  Afun :: (Arrays a, Arrays b)
       => (a -> b)
       -> a {- dummy -}
       -> b {- dummy -}
       -> Afun

data EArray where
  EArray :: (Shape sh, Elt e) => Array sh e -> EArray

-- We need to export these
foreign export ccall accelerateCreate  :: Device -> ForeignContext -> IO AccHandle
foreign export ccall accelerateDestroy :: AccHandle -> IO ()
foreign export ccall runProgram        :: AccHandle -> StablePtr Afun -> Ptr InputArray -> Ptr OutputArray -> IO ()
foreign export ccall freeOutput        :: Ptr OutputArray -> IO ()
foreign export ccall freeProgram       :: StablePtr a -> IO ()

instance Storable InputArray where
  sizeOf (sh, ptrs) = sizeOf sh + sizeOf ptrs

  alignment _ = 0

  peek ptr = do
    let p_sh   = castPtr ptr :: Ptr ShapeBuffer
    sh         <- peek p_sh
    let p_ptrs = (castPtr p_sh :: Ptr DevicePtrBuffer) `plusPtr` sizeOf sh
    ptrs       <- peek p_ptrs
    return (sh, ptrs)

  poke ptr (sh, ptrs) = do
    let p_sh   = castPtr ptr :: Ptr ShapeBuffer
        p_ptrs = (castPtr p_sh :: Ptr DevicePtrBuffer) `plusPtr` sizeOf sh
    poke p_sh sh
    poke p_ptrs ptrs

instance Storable OutputArray where
  sizeOf (sh, ptrs, sa) = sizeOf sh + sizeOf ptrs + sizeOf sa
  alignment _ = 0
  peek ptr = do
    let p_sh   = castPtr ptr :: Ptr ShapeBuffer
    sh         <- peek p_sh
    let p_ptrs = (castPtr p_sh :: Ptr DevicePtrBuffer) `plusPtr` sizeOf sh
    ptrs       <- peek p_ptrs
    let p_sa   = (castPtr p_ptrs :: Ptr (StablePtr a)) `plusPtr` sizeOf ptrs
    sa         <- peek p_sa
    return (sh, ptrs, sa)
  poke ptr (sh, ptrs, sa) = do
    let p_sh   = castPtr ptr :: Ptr ShapeBuffer
        p_ptrs = (castPtr p_sh :: Ptr DevicePtrBuffer) `plusPtr` sizeOf sh
        p_sa   = (castPtr p_ptrs :: Ptr (StablePtr a)) `plusPtr` sizeOf ptrs
    poke p_sh sh
    poke p_ptrs ptrs
    poke p_sa sa

-- |Create an Accelerate handle given a device and a cuda context.
--
-- @AccHandle accelerateCreate(int device, CUcontext ctx);@
accelerateCreate :: Device -> ForeignContext -> IO AccHandle
accelerateCreate dev ctx = fromDeviceContext (CUDA.Device $ CInt dev) (CUDA.Context ctx) >>= newStablePtr

-- |Releases all resources used by the accelerate library.
--
-- @void accelerateDestroy(AccHandle hndl);@
accelerateDestroy :: AccHandle -> IO ()
accelerateDestroy = freeStablePtr

-- |Function callable from foreign code to 'free' a OutputArray returned after executing
-- an Accelerate computation.
--
-- Once freed, the device pointers associated with an array are no longer valid.
--
-- @void freeOutput(OutputArray arr);@
freeOutput :: Ptr OutputArray -> IO ()
freeOutput o = do
  (sh, dptrs, sa) <- peek o
  free sh
  free dptrs
  freeStablePtr sa

-- |Free a compiled accelerate program.
--
-- @void freeProgram(Program prg);@
freeProgram :: StablePtr a -> IO ()
freeProgram = freeStablePtr

-- |Execute the given accelerate program with @is@ as the input and @os@ as the output.
--
-- @void runProgram(AccHandle hndl, AccProgram p, InputArray* is, OutputArray* os);@
runProgram :: AccHandle -> StablePtr Afun -> Ptr InputArray -> Ptr OutputArray -> IO ()
runProgram hndl fun input output = do
  ctx <- deRefStablePtr hndl
  af <- deRefStablePtr fun
  run ctx af
  where
    run :: Context -> Afun -> IO ()
    run ctx (Afun f (_ :: a) (_ :: b)) = do
      _ <- evalCUDA ctx $ do
        (a, _) <- marshalIn (arrays (undefined :: a)) input
        let !b = f (toArr a)
        marshalOut (arrays (undefined :: b)) (fromArr b) output
      return ()

    marshalIn :: ArraysR a -> Ptr InputArray -> CIO (a, Ptr InputArray)
    marshalIn ArraysRunit  ptr = return ((), ptr)
    marshalIn ArraysRarray ptr = do
      (sh, ptrs) <- liftIO (peek ptr)
      a <- arrayFromForeignData ptrs sh
      let ptr' = plusPtr ptr (sizeOf (sh, ptrs))
      return (a, ptr')
    marshalIn (ArraysRpair aR1 aR2) ptr = do
      (x, ptr')  <- marshalIn aR1 ptr
      (y, ptr'') <- marshalIn aR2 ptr'
      return ((x,y), ptr'')

    marshalOut :: ArraysR b -> b -> Ptr OutputArray -> CIO (Ptr OutputArray)
    marshalOut ArraysRunit  () ptr = return ptr
    marshalOut ArraysRarray a  ptr = do
      oarr <- mkOutput a
      liftIO $ poke ptr oarr
      return (plusPtr ptr (sizeOf oarr))
      where
        mkOutput :: forall sh e. Shape sh => Array sh e -> CIO OutputArray
        mkOutput (Array sh adata) = do
          let sh' = shapeToList (toElt sh :: sh)
          shbuf <- liftIO $ mallocArray (P.length sh')
          liftIO $ pokeArray shbuf (map fromIntegral sh')

          withDevicePtrs adata Nothing $ \dptrs -> do
            let wptrs = devicePtrsToWordPtrs adata dptrs
            pbuf  <- liftIO $ mallocArray (P.length wptrs)
            liftIO $ pokeArray pbuf wptrs

            sa <- liftIO $ newStablePtr (EArray a)
            return (shbuf, pbuf, sa)

    marshalOut (ArraysRpair aR1 aR2) (x,y) ptr = do
      ptr' <- marshalOut aR1 x ptr
      marshalOut aR2 y ptr'


-- |Given the 'Name' of an Accelerate function (a function of type ''Acc a -> Acc b'') generate a
-- a function callable from foreign code with the second argument specifying it's name.
exportAfun :: Name -> String -> Q [Dec]
exportAfun fname ename = do
#if __GLASGOW_HASKELL__ <= 710
  (VarI n ty _ _) <- reify fname
#else
  (VarI n ty _)   <- reify fname
#endif

  -- Generate initialisation function
  genCompileFun n ename ty

genCompileFun :: Name -> String -> Type -> Q [Dec]
genCompileFun fname ename (AppT (AppT ArrowT (AppT _ _)) (AppT _ _))
  = sequence [sig, dec, expt]
  where
    initName = mkName ename

    body = [| \hndl -> buildExported hndl $(varE fname) |]
    dec  = FunD initName . (:[]) <$> cls
    cls  = Clause [] <$> (NormalB <$> body) <*> return []
    sig  = SigD initName <$> ety
    expt = ForeignD <$> (ExportF cCall (nameBase initName) initName <$> ety)

    ety   = [t| AccHandle -> IO (StablePtr Afun) |]
genCompileFun _     _     _
  = error "Invalid accelerate function"

-- |Given a handle and an Accelerate function, generate an exportable version.
buildExported :: forall a b. (Arrays a, Arrays b) => AccHandle -> (Acc a -> Acc b) -> IO (StablePtr Afun)
buildExported hndl f = ef
  where
    ef :: IO (StablePtr Afun)
    ef = do
      ctx <- deRefStablePtr hndl
      newStablePtr (Afun (run1With ctx f) (undefined :: a) (undefined :: b))

-- Utility functions
-- ------------------

arrayFromForeignData :: forall sh e. (Shape sh, Elt e) => DevicePtrBuffer -> ShapeBuffer -> CIO (Array sh e)
arrayFromForeignData ptrs shape = do
   let d  = rank (ignore :: sh) -- Using ignore as using dim requires a non-dummy argument
   let sz = eltSize (eltType (undefined :: e))
   lst <- liftIO (peekArray d shape)
   let sh = listToShape (map fromIntegral lst) :: sh
   plst <- liftIO $ peekArray sz ptrs
   let ptrs' = devicePtrsFromList (arrayElt :: ArrayEltR (EltRepr e)) plst
   useDevicePtrs (fromElt sh) ptrs'

eltSize :: TupleType e -> Int
eltSize UnitTuple         = 0
eltSize (SingleTuple _  ) = 1
eltSize (PairTuple   a b) = eltSize a + eltSize b