{-# LANGUAGE CPP, GADTs, TupleSections #-}
-- |
-- Module      : Data.Array.Accelerate.CUDA.Execute
-- Copyright   : [2008..2010] Manuel M T Chakravarty, Gabriele Keller, Sean Lee, Trevor L. McDonell
-- License     : BSD3
--
-- Maintainer  : Manuel M T Chakravarty <chak@cse.unsw.edu.au>
-- Stability   : experimental
-- Portability : non-partable (GHC extensions)
--

module Data.Array.Accelerate.CUDA.Execute (executeAcc)
  where

import Prelude hiding (id, (.), sum)
import Control.Category

import Data.Maybe
import Control.Monad
import Control.Monad.Trans
import Control.Applicative                              hiding (Const)
import qualified Data.HashTable                         as HT

import System.FilePath
import System.Posix.Process
import System.Exit                                      (ExitCode(..))
import System.Posix.Types                               (ProcessID)
import Text.PrettyPrint

import Data.Array.Accelerate.AST
import Data.Array.Accelerate.Tuple
import Data.Array.Accelerate.Array.Data
import Data.Array.Accelerate.Array.Representation       hiding (sliceIndex)
import Data.Array.Accelerate.Array.Sugar                (Array(..))
import qualified Data.Array.Accelerate.Array.Sugar      as Sugar
import qualified Data.Array.Accelerate.Interpreter      as I

import Data.Array.Accelerate.CUDA.State
import Data.Array.Accelerate.CUDA.Array.Data
import Data.Array.Accelerate.CUDA.Analysis.Hash
import Data.Array.Accelerate.CUDA.Analysis.Launch

import qualified Foreign.CUDA.Driver                    as CUDA

#include "accelerate.h"


-- Expression evaluation
-- ~~~~~~~~~~~~~~~~~~~~~

-- Evaluate a closed scalar expression. Expressions are evaluated on the host,
-- but may require some interaction with the device (such as array indexing).
--
-- TLM 2010-07-13:
--   We sneakily use the Interpreter backend for primitive operations, but maybe
--   we don't want to...?
--
executeExp :: Exp aenv t -> Val aenv -> CIO t
executeExp e = executeOpenExp e Empty

executeOpenExp :: OpenExp env aenv t -> Val env -> Val aenv -> CIO t
executeOpenExp (Var idx) env _            = return . Sugar.toElem $ prj idx env
executeOpenExp (Const c) _ _              = return $ Sugar.toElem c
executeOpenExp (PrimConst c) _ _          = return $ I.evalPrimConst c
executeOpenExp (PrimApp fun arg) env aenv = I.evalPrim fun <$> executeOpenExp arg env aenv
executeOpenExp (Prj idx e) env aenv       = I.evalPrj idx . fromTuple <$> executeOpenExp e env aenv
executeOpenExp (Tuple tup) env aenv       = toTuple                   <$> executeTuple tup env aenv
executeOpenExp (IndexScalar a e) env aenv = do
  (Array sh ad) <- executeOpenAcc a aenv
  ix            <- executeOpenExp e env aenv
  Sugar.toElem <$> ad `indexArray` index sh (Sugar.fromElem ix)

executeOpenExp (Shape a) _ aenv = do
  (Array sh _)  <- executeOpenAcc a aenv
  return (Sugar.toElem sh)

executeOpenExp (Cond c t e) env aenv = do
  p <- executeOpenExp c env aenv
  if p then executeOpenExp t env aenv
       else executeOpenExp e env aenv


executeTuple :: Tuple (OpenExp env aenv) t -> Val env -> Val aenv -> CIO t
executeTuple NilTup            _    _     = return ()
executeTuple (tup `SnocTup` e) env  aenv  = (,) <$> executeTuple tup env aenv <*> executeOpenExp e env aenv


-- Array evaluation
-- ~~~~~~~~~~~~~~~

-- | Execute a closed array expression
--
executeAcc :: Acc a -> CIO a
executeAcc acc = executeOpenAcc acc Empty

-- |
-- Execute an embedded array program using the CUDA backend. The function will
-- block if the compilation has not yet completed, but subsequent invocations of
-- the same kernel will be able to extract the loaded kernel directly.
--
executeOpenAcc :: OpenAcc aenv a -> Val aenv -> CIO a
executeOpenAcc (Use  xs) _env = return xs
executeOpenAcc (Avar ix)  env = return $ prj ix env
executeOpenAcc (Let  x y) env = do
  ax <- executeOpenAcc x env
  executeOpenAcc y (env `Push` ax)

executeOpenAcc (Let2 x y) env = do
  (ax1,ax2) <- executeOpenAcc x env
  executeOpenAcc y (env `Push` ax1 `Push` ax2)

executeOpenAcc (Unit e)   env = do
  v  <- executeExp e env
  let ad = fst . runArrayData $ (,undefined) <$> do
        arr <- newArrayData 1
        writeArrayData  arr 0 (Sugar.fromElem v)
        return arr

  mallocArray    ad 1
  pokeArrayAsync ad 1 Nothing
  return (Array (Sugar.fromElem ()) ad)


executeOpenAcc acc env = do
  tab <- getM kernelTable
  krn <- fromMaybe (error "code generation failed") <$> liftIO (HT.lookup tab key)
  mdl <- either' (getL kernelStatus krn) return $ \pid -> do
    liftIO (waitFor pid)
    mdl <- liftIO $ CUDA.loadFile (getL kernelName krn `replaceExtension` ".cubin")
    liftIO        $ HT.insert tab key (setL kernelStatus (Right mdl) krn)
    return mdl

  -- determine dispatch pattern, extract parameters, allocate storage, run
  --
  dispatch acc env mdl

  where
    key           = accToKey acc
    either' e r l = either l r e


-- Dispatch
-- ~~~~~~~~

data FVar where
  FArr :: Array dim e -> FVar

-- Lift free array variables out of scalar computations. Returns a list of the
-- arrays in the order that they are encountered, which also corresponds to the
-- texture reference index they should be bound to.
--
liftFun :: OpenFun env aenv a -> Val aenv -> CIO [FVar]
liftFun (Body e) = liftExp e
liftFun (Lam f)  = liftFun f

liftExp :: OpenExp env aenv a -> Val aenv -> CIO [FVar]
liftExp (Tuple t)         aenv = liftTup t aenv
liftExp (Prj _ e)         aenv = liftExp e aenv
liftExp (PrimApp _ e)     aenv = liftExp e aenv
liftExp (Cond e1 e2 e3)   aenv = concat <$> sequence [liftExp e1 aenv, liftExp e2 aenv, liftExp e3 aenv]
liftExp (IndexScalar a e) aenv = (:) . FArr <$> executeOpenAcc a aenv <*> liftExp e aenv
liftExp _ _ =
  return []

liftTup :: Tuple (OpenExp env aenv) a -> Val aenv -> CIO [FVar]
liftTup NilTup _             = return []
liftTup (t `SnocTup` e) aenv = (++) <$> liftTup t aenv <*> liftExp e aenv


-- Extract texture references from the compiled module and bind an array to it
--
bind :: CUDA.Module -> [FVar] -> CIO [CUDA.FunParam]
bind mdl var =
  let tex n (FArr (Array sh ad)) = textureRefs ad mdl (size sh) n
  in  foldM (\texs farr -> (texs ++) <$> tex (length texs) farr) [] var

release :: [FVar] -> CIO ()
release = mapM_ (\(FArr (Array _ ad)) -> freeArray ad)


-- Setup and initiate the computation. This may require several kernel launches.
--
-- A NOTE ON TUPLES TYPES
--
--   The astute reader may be wondering, if arrays of tuples are stored as a
--   tuple of arrays, how exactly are we telling the kernel function about this?
--
--   From Haskell land, the tuple components are passed to the kernel function
--   individually. The C function, however, interprets these into a structure of
--   pointers. While this is a nasty sleight-of-hand, it should indeed be safe.
--   Pointer types will all have the same size and alignment, and C structures
--   are defined to keep their fields adjacent in memory (modulo alignment
--   restrictions, which don't concern us).
--
dispatch :: OpenAcc aenv a -> Val aenv -> CUDA.Module -> CIO a
dispatch acc@(Map f ad) env mdl = do
  fn             <- liftIO $ CUDA.getFun mdl "map"
  (Array sh in0) <- executeOpenAcc ad env
  let res@(Array sh' out) = newArray (Sugar.toElem sh)
      n                   = size sh'

  mallocArray out n
  d_out <- devicePtrs out
  d_in0 <- devicePtrs in0
  f_var <- liftFun f env
  t_var <- bind mdl f_var

  launch acc n fn (d_out ++ d_in0 ++ t_var ++ [CUDA.IArg n])
  freeArray in0
  release f_var
  return res

dispatch acc@(ZipWith f ad1 ad0) env mdl = do
  fn              <- liftIO $ CUDA.getFun mdl "zipWith"
  (Array sh1 in1) <- executeOpenAcc ad1 env
  (Array sh0 in0) <- executeOpenAcc ad0 env
  let res@(Array sh' out) = newArray (Sugar.toElem (sh0 `intersect` sh1))
      n                   = size sh'

  mallocArray out n
  d_out <- devicePtrs out
  d_in1 <- devicePtrs in1
  d_in0 <- devicePtrs in0
  f_var <- liftFun f env
  t_var <- bind mdl f_var

  launch acc n fn (d_out ++ d_in1 ++ d_in0 ++ t_var ++ convertIx sh' ++ convertIx sh1 ++ convertIx sh0)
  freeArray in0
  freeArray in1
  release f_var
  return res

dispatch acc@(Fold f x ad) env mdl = do
  fn              <- liftIO $ CUDA.getFun mdl "fold"
  (Array sh in0)  <- executeOpenAcc ad env
  (cta,grid,smem) <- launchConfig acc (size sh) fn
  let res@(Array _ out) = newArray grid

  mallocArray out grid
  d_out <- devicePtrs out
  d_in0 <- devicePtrs in0
  f_arr <- (++) <$> liftExp x env <*> liftFun f env
  t_var <- bind mdl f_arr

  launch' (cta,grid,smem) fn (d_out ++ d_in0 ++ t_var ++ [CUDA.IArg (size sh)])
  freeArray in0
  release f_arr
  if grid > 1 then dispatch (Fold f x (Use res)) env mdl
              else return (Array (Sugar.fromElem ()) out)

dispatch acc@(FoldSeg f x ad sd) env mdl = do
  fn              <- liftIO $ CUDA.getFun mdl "fold_segmented"
  (Array sh  in0) <- executeOpenAcc ad env
  (Array sh' seg) <- executeOpenAcc sd env
  (cta,grid,smem) <- launchConfig acc (size sh') fn
  let res@(Array _ out) = newArray (size sh')

  mallocArray out (size sh')
  d_out <- devicePtrs out
  d_in0 <- devicePtrs in0
  d_seg <- devicePtrs seg
  f_arr <- (++) <$> liftExp x env <*> liftFun f env
  t_var <- bind mdl f_arr

  launch' (cta,grid,smem) fn (d_out ++ d_in0 ++ d_seg ++ t_var ++ map (CUDA.IArg . size) [sh', sh])
  freeArray in0
  freeArray seg
  release f_arr
  return res


dispatch acc@(Scanl _ _ _) env mdl = dispatchScan acc env mdl
dispatch acc@(Scanr _ _ _) env mdl = dispatchScan acc env mdl

dispatch acc@(Permute f1 df f2 ad) env mdl = do
  fn              <- liftIO $ CUDA.getFun mdl "permute"
  (Array sh  def) <- executeOpenAcc df env
  (Array sh0 in0) <- executeOpenAcc ad env
  let res@(Array _ out) = newArray (Sugar.toElem sh)
      n                 = size sh0

  mallocArray out n
  copyArray def out n
  d_out <- devicePtrs out
  d_in0 <- devicePtrs in0
  f_arr <- (++) <$> liftFun f1 env <*> liftFun f2 env
  t_var <- bind mdl f_arr

  launch acc n fn (d_out ++ d_in0 ++ t_var ++ convertIx sh ++ convertIx sh0)
  freeArray def
  freeArray in0
  release f_arr
  return res

dispatch acc@(Backpermute e f ad) env mdl = do
  fn              <- liftIO $ CUDA.getFun mdl "backpermute"
  sh              <- executeExp e env
  (Array sh0 in0) <- executeOpenAcc ad env
  let res@(Array sh' out) = newArray sh
      n                   = size sh'

  mallocArray out n
  d_out <- devicePtrs out
  d_in0 <- devicePtrs in0
  f_arr <- liftFun f env
  t_var <- bind mdl f_arr

  launch acc n fn (d_out ++ d_in0 ++ t_var ++ convertIx sh' ++ convertIx sh0)
  freeArray in0
  release f_arr
  return res

dispatch acc@(Replicate sliceIndex e ad) env mdl = do
  fn              <- liftIO $ CUDA.getFun mdl "replicate"
  slix            <- executeExp e env
  (Array sh0 in0) <- executeOpenAcc ad env

  let block               = extend sliceIndex (Sugar.fromElem slix) sh0
      res@(Array sh' out) = newArray (Sugar.toElem block)
      n                   = size sh'

      extend :: SliceIndex slix sl co dim -> slix -> sl -> dim
      extend (SliceNil)            ()       ()      = ()
      extend (SliceAll sliceIdx)   (slx,()) (sl,sz) = (extend sliceIdx slx sl, sz)
      extend (SliceFixed sliceIdx) (slx,sz) sl      = (extend sliceIdx slx sl, sz)

  mallocArray out n
  d_out <- devicePtrs out
  d_in0 <- devicePtrs in0
  f_arr <- liftExp e env
  t_var <- bind mdl f_arr

  launch acc n fn (d_out ++ d_in0 ++ t_var ++ convertIx sh0 ++ convertIx sh')
  freeArray in0
  release f_arr
  return res

dispatch acc@(Index sliceIndex ad e) env mdl = do
  fn              <- liftIO $ CUDA.getFun mdl "slice"
  slix            <- executeExp e env
  (Array sh0 in0) <- executeOpenAcc ad env
  let slice               = restrict sliceIndex (Sugar.fromElem slix) sh0
      slix'               = convertSliceIndex sliceIndex (Sugar.fromElem slix)
      res@(Array sh' out) = newArray (Sugar.toElem slice)
      n                   = size sh'

      restrict :: SliceIndex slix sl co dim -> slix -> dim -> sl
      restrict (SliceNil)            ()       ()      = ()
      restrict (SliceAll sliceIdx)   (slx,()) (sh,sz) = (restrict sliceIdx slx sh, sz)
      restrict (SliceFixed sliceIdx) (slx,i)  (sh,sz)
        = BOUNDS_CHECK(checkIndex) "slice" i sz $ restrict sliceIdx slx sh

  mallocArray out n
  d_out <- devicePtrs out
  d_in0 <- devicePtrs in0
  f_arr <- liftExp e env
  t_var <- bind mdl f_arr

  launch acc n fn (d_out ++ d_in0 ++ t_var ++ convertIx sh' ++ slix' ++ convertIx sh0)
  freeArray in0
  release f_arr
  return res

dispatch x _ _ =
  INTERNAL_ERROR(error) "dispatch"
  (unlines ["unsupported array primitive", render . nest 2 $ text (show x)])


-- Unified dispatch handler for left/right scan.
--
-- TLM 2010-07-02:
--   This is a little awkward. At its core we have an inclusive scan routine,
--   whereas accelerate actually returns an exclusive scan result. Multi-block
--   arrays require the inclusive scan when calculating the partial block sums,
--   which are then added to every element of an interval. Single-block arrays
--   on the other hand need to be "converted" to an exclusive result in the
--   second pass.
--
--   This optimised could be implemented by enabling some some extra code to the
--   skeleton, and dispatching accordingly.
--
dispatchScan :: OpenAcc aenv a -> Val aenv -> CUDA.Module -> CIO a
dispatchScan     (Scanr f x ad) env mdl = dispatchScan (Scanl f x ad) env mdl
dispatchScan acc@(Scanl f x ad) env mdl = do
  fscan           <- liftIO $ CUDA.getFun mdl "inclusive_scan"
  fadd            <- liftIO $ CUDA.getFun mdl "exclusive_update"
  (Array sh in0)  <- executeOpenAcc ad env
  (cta,grid,smem) <- launchConfig acc (size sh) fscan
  let a_out@(Array _ out) = newArray (Sugar.toElem sh)
      a_sum@(Array _ sum) = newArray ()
      a_bks@(Array _ bks) = newArray grid
      n                   = size sh
      interval            = (n + grid - 1) `div` grid

      unify :: Array dim e -> Array dim e -> CIO ()
      unify _ _ = return ()

  unify a_out a_bks -- TLM: *cough*

  mallocArray out n
  mallocArray sum 1
  mallocArray bks grid
  d_out <- devicePtrs out
  d_in0 <- devicePtrs in0
  d_bks <- devicePtrs bks
  d_sum <- devicePtrs sum
  f_arr <- (++) <$> liftExp x env <*> liftFun f env
  t_var <- bind mdl f_arr

  launch' (cta,grid,smem) fscan (d_out ++ d_in0 ++ d_bks ++ t_var ++ map CUDA.IArg [n,interval])
  launch' (cta,1,smem)    fscan (d_bks ++ d_bks ++ d_sum ++ map CUDA.IArg [grid,interval])
  launch' (cta,grid,smem) fadd  (d_out ++ d_bks ++ map CUDA.IArg [n,interval])

  freeArray in0
  freeArray bks
  release f_arr
  return (a_out, a_sum)

dispatchScan _ _ _ =
  error "we can never get here"


-- Initiate the device computation. The first version selects launch parameters
-- automatically, the second requires them explicitly. This tuple contains
-- threads per block, grid size, and dynamic shared memory, respectively.
--
launch :: OpenAcc aenv a -> Int -> CUDA.Fun -> [CUDA.FunParam] -> CIO ()
launch acc n fn args =
  launchConfig acc n fn >>= \cfg ->
  launch' cfg fn args

launch' :: (Int,Int,Integer) -> CUDA.Fun -> [CUDA.FunParam] -> CIO ()
launch' (cta,grid,smem) fn args =
  liftIO $ do
    CUDA.setParams     fn args
    CUDA.setSharedSize fn smem
    CUDA.setBlockShape fn (cta,1,1)
    CUDA.launch        fn (grid,1) Nothing


-- Create a new array (obviously...)
--
{-# INLINE newArray #-}
newArray :: (Sugar.Ix dim, Sugar.Elem e) => dim -> Array dim e
newArray sh = ad `seq` Array (Sugar.fromElem sh) ad
  where
    -- FIXME: small arrays are relocated by the GC
    ad = fst . runArrayData $ (,undefined) <$> newArrayData (1024 `max` Sugar.size sh)

-- Extract shape dimensions as a list of function parameters. Not that this will
-- convert to the base integer width of the device, namely, 32-bits. Singleton
-- dimensions are considered to be of unit size.
--
convertIx :: Ix dim => dim -> [CUDA.FunParam]
convertIx = post . map CUDA.IArg . shapeToList
  where post [] = [CUDA.IArg 1]
        post xs = xs

-- Convert a slice specification into storable index projection components.
-- Note: implicit conversion Int -> Int32
--
convertSliceIndex :: SliceIndex slix sl co dim -> slix -> [CUDA.FunParam]
convertSliceIndex (SliceNil)            ()     = []
convertSliceIndex (SliceAll   sliceIdx) (s,()) = convertSliceIndex sliceIdx s
convertSliceIndex (SliceFixed sliceIdx) (s,i)  = CUDA.IArg i : convertSliceIndex sliceIdx s

-- | Wait for the compilation process to finish
--
waitFor :: ProcessID -> IO ()
waitFor pid = do
  status <- getProcessStatus True True pid
  case status of
       Just (Exited ExitSuccess) -> return ()
       _                         -> error  $ "nvcc (" ++ shows pid ") terminated abnormally"