Tidy-up marshalling interface

Tidy-up marshalling interface

Ignore-this: d2b9fb1206f4a304a9d69c6cc3f4ab65
- Have function names make it explicit that the focus is operating over
  collections of data (Array and List)
- Add several convenience functions, for temporary allocation of arrays and
  marshalling data between Haskell lists and device arrays

darcs-hash:20091221101434-dcabc-3992b11e7794a676fb45533fb224798c51925ef0.gz

{-# LANGUAGE ForeignFunctionInterface #-}
--------------------------------------------------------------------------------
-- |
-- Module    : Foreign.CUDA.Driver.Marshal
-- Copyright : (c) 2009 Trevor L. McDonell
-- License   : BSD
--
-- Memory management for low-level driver interface
--
--------------------------------------------------------------------------------

module Foreign.CUDA.Driver.Marshal
  (
    -- * Host Allocation
    HostPtr, AllocFlag(..),
    withHostPtr, mallocHostArray, freeHost,

    -- * Device Allocation
    DevicePtr,
    mallocArray, allocaArray, free,

    -- * Marshalling
    peekArray, peekArrayAsync, peekListArray,
    pokeArray, pokeArrayAsync, pokeListArray,

    -- * Combined Allocation and Marshalling
    newListArray, withListArray,

    -- * Utility
    memset, getDevicePtr
  )
  where

#include <cuda.h>
{# context lib="cuda" #}

-- Friends
import Foreign.CUDA.Internal.C2HS
import Foreign.CUDA.Driver.Error
import Foreign.CUDA.Driver.Stream               (Stream(..))

-- System
import Unsafe.Coerce
import Control.Monad                            (liftM)
import Control.Exception.Extensible

import Foreign.C
import Foreign.Ptr
import Foreign.Storable
import qualified Foreign.Marshal as F

#c
typedef enum CUmemhostalloc_option_enum {
    CU_MEMHOSTALLOC_OPTION_PORTABLE       = CU_MEMHOSTALLOC_PORTABLE,
    CU_MEMHOSTALLOC_OPTION_DEVICE_MAPPED  = CU_MEMHOSTALLOC_DEVICEMAP,
    CU_MEMHOSTALLOC_OPTION_WRITE_COMBINED = CU_MEMHOSTALLOC_WRITECOMBINED
} CUmemhostalloc_option;
#endc

--------------------------------------------------------------------------------
-- Data Types
--------------------------------------------------------------------------------

-- |
-- A reference to memory allocated on the device
--
newtype DevicePtr a = DevicePtr { useDevicePtr :: {# type CUdeviceptr #}}

instance Storable (DevicePtr a) where
  sizeOf _      = sizeOf    (undefined :: {# type CUdeviceptr #})
  alignment _   = alignment (undefined :: {# type CUdeviceptr #})
  peek p        = DevicePtr `fmap` peek (castPtr p)
  poke p v      = poke (castPtr p) (useDevicePtr v)

-- |
-- A reference to memory on the host that is page-locked and directly-accessible
-- from the device. Since the memory can be accessed directly, it can be read or
-- written at a much higher bandwidth than pageable memory from the traditional
-- malloc.
--
-- The driver automatically accelerates calls to functions such as `memcpy'
-- which reference page-locked memory.
--
newtype HostPtr a = HostPtr { useHostPtr :: Ptr a }

-- |
-- Unwrap a host pointer and execute a computation using the base pointer object
--
withHostPtr :: HostPtr a -> (Ptr a -> IO b) -> IO b
withHostPtr p f = f (useHostPtr p)

-- |
-- Options for host allocation
--
{# enum CUmemhostalloc_option as AllocFlag
    { underscoreToCase }
    with prefix="CU_MEMHOSTALLOC_OPTION" deriving (Eq, Show) #}

--------------------------------------------------------------------------------
-- Host Allocation
--------------------------------------------------------------------------------

-- |
-- Allocate a section of linear memory on the host which is page-locked and
-- directly accessible from the device. The storage is sufficient to hold the
-- given number of elements of a storable type.
--
-- Note that since the amount of pageable memory is thusly reduced, overall
-- system performance may suffer. This is best used sparingly to allocate
-- staging areas for data exchange.
--
mallocHostArray :: Storable a => [AllocFlag] -> Int -> IO (HostPtr a)
mallocHostArray flags = doMalloc undefined
  where
    doMalloc :: Storable a' => a' -> Int -> IO (HostPtr a')
    doMalloc x n = resultIfOk =<< cuMemHostAlloc (n * sizeOf x) flags

{# fun unsafe cuMemHostAlloc
  { alloca'-        `HostPtr a'   peekHP*
  ,                 `Int'
  , combineBitMasks `[AllocFlag]'         } -> `Status' cToEnum #}
  where
    alloca'  = F.alloca
    peekHP p = (HostPtr . castPtr) `fmap` peek p

-- |
-- Free a section of page-locked host memory
--
freeHost :: HostPtr a -> IO ()
freeHost p = nothingIfOk =<< cuMemFreeHost p

{# fun unsafe cuMemFreeHost
  { useHP `HostPtr a' } -> `Status' cToEnum #}
  where
    useHP = castPtr . useHostPtr

--------------------------------------------------------------------------------
-- Device Allocation
--------------------------------------------------------------------------------

-- |
-- Allocate a section of linear memory on the device, and return a reference to
-- it. The memory is sufficient to hold the given number of elements of storable
-- type. It is suitably aligned for any type, and is not cleared.
--
mallocArray :: Storable a => Int -> IO (DevicePtr a)
mallocArray = doMalloc undefined
  where
    doMalloc :: Storable a' => a' -> Int -> IO (DevicePtr a')
    doMalloc x n = resultIfOk =<< cuMemAlloc (n * sizeOf x)

{# fun unsafe cuMemAlloc
  { alloca'- `DevicePtr a' peekDP*
  ,          `Int'               } -> `Status' cToEnum #}
  where
    alloca' = F.alloca
    peekDP  = liftM DevicePtr . peek

-- |
-- Execute a computation on the device, passing a pointer to a temporarily
-- allocated block of memory sufficient to hold the given number of elements of
-- storable type. The memory is freed when the computation terminates (normally
-- or via an exception), so the pointer must not be used after this.
--
-- Note that kernel launches can be asynchronous, so you may want to add a
-- synchronisation point using 'sync' as part of the computation.
--
allocaArray :: Storable a => Int -> (DevicePtr a -> IO b) -> IO b
allocaArray n = bracket (mallocArray n) free

-- |
-- Release a section of device memory
--
free :: DevicePtr a -> IO ()
free dp = nothingIfOk =<< cuMemFree dp

{# fun unsafe cuMemFree
  { useDevicePtr `DevicePtr a' } -> `Status' cToEnum #}

--------------------------------------------------------------------------------
-- Marshalling
--------------------------------------------------------------------------------

-- |
-- Copy a number of elements from the device to host memory. This is a
-- synchronous operation
--
peekArray :: Storable a => Int -> DevicePtr a -> Ptr a -> IO ()
peekArray n dptr hptr = doPeek undefined dptr
  where
    doPeek :: Storable a' => a' -> DevicePtr a' -> IO ()
    doPeek x _ = nothingIfOk =<< cuMemcpyDtoH hptr dptr (n * sizeOf x)

{# fun unsafe cuMemcpyDtoH
  { castPtr      `Ptr a'
  , useDevicePtr `DevicePtr a'
  ,              `Int'         } -> `Status' cToEnum #}

-- |
-- Copy memory from the device asynchronously, possibly associated with a
-- particular stream. The destination host memory must be page-locked.
--
peekArrayAsync :: Storable a => Int -> DevicePtr a -> HostPtr a -> Maybe Stream -> IO ()
peekArrayAsync n dptr hptr mst = doPeek undefined dptr
  where
    doPeek :: Storable a' => a' -> DevicePtr a' -> IO ()
    doPeek x _ =
      nothingIfOk =<< case mst of
        Nothing -> cuMemcpyDtoHAsync hptr dptr (n * sizeOf x) (Stream nullPtr)
        Just st -> cuMemcpyDtoHAsync hptr dptr (n * sizeOf x) st

{# fun unsafe cuMemcpyDtoHAsync
  { useHP        `HostPtr a'
  , useDevicePtr `DevicePtr a'
  ,              `Int'
  , useStream    `Stream'  } -> `Status' cToEnum #}
  where
    useHP = castPtr . useHostPtr

-- |
-- Copy a number of elements from the device into a new Haskell list. Note that
-- this requires two memory copies: firstly from the device into a heap
-- allocated array, and from there marshalled into a list.
--
peekListArray :: Storable a => Int -> DevicePtr a -> IO [a]
peekListArray n dptr =
  F.allocaArray n $ \p -> do
    peekArray   n dptr p
    F.peekArray n p

-- |
-- Copy a number of elements onto the device. This is a synchronous operation
--
pokeArray :: Storable a => Int -> Ptr a -> DevicePtr a -> IO ()
pokeArray n hptr dptr = doPoke undefined dptr
  where
    doPoke :: Storable a' => a' -> DevicePtr a' -> IO ()
    doPoke x _ = nothingIfOk =<< cuMemcpyHtoD dptr hptr (n * sizeOf x)

{# fun unsafe cuMemcpyHtoD
  { useDevicePtr `DevicePtr a'
  , castPtr      `Ptr a'
  ,              `Int'         } -> `Status' cToEnum #}

-- |
-- Copy memory onto the device asynchronously, possibly associated with a
-- particular stream. The source host memory must be page-locked.
--
pokeArrayAsync :: Storable a => Int -> HostPtr a -> DevicePtr a -> Maybe Stream -> IO ()
pokeArrayAsync n hptr dptr mst = dopoke undefined dptr
  where
    dopoke :: Storable a' => a' -> DevicePtr a' -> IO ()
    dopoke x _ =
      nothingIfOk =<< case mst of
        Nothing -> cuMemcpyHtoDAsync dptr hptr (n * sizeOf x) (Stream nullPtr)
        Just st -> cuMemcpyHtoDAsync dptr hptr (n * sizeOf x) st

{# fun unsafe cuMemcpyHtoDAsync
  { useDevicePtr `DevicePtr a'
  , useHP        `HostPtr a'
  ,              `Int'
  , useStream    `Stream'  } -> `Status' cToEnum #}
  where
    useHP = castPtr . useHostPtr

-- |
-- Write a list of storable elements into a device array. The device array must
-- be sufficiently large to hold the entire list. This requires two marshalling
-- operations.
--
pokeListArray :: Storable a => [a] -> DevicePtr a -> IO ()
pokeListArray xs dptr = F.withArrayLen xs $ \len p -> pokeArray len p dptr

--------------------------------------------------------------------------------
-- Combined Allocation and Marshalling
--------------------------------------------------------------------------------

-- |
-- Write a list of storable elements into a newly allocated device array. Note
-- that this requires two memory copies: firstly from a Haskell list to a heap
-- allocated array, and from there onto the graphics device. The memory should
-- be 'free'd when no longer required.
--
newListArray :: Storable a => [a] -> IO (DevicePtr a)
newListArray xs =
  F.withArrayLen xs                     $ \len p ->
  bracketOnError (mallocArray len) free $ \d_xs  -> do
    pokeArray len p d_xs
    return d_xs

-- |
-- Temporarily store a list of elements into a newly allocated device array. An
-- IO action is applied to to the array, the result of which is returned.
-- Similar to 'newListArray', this requires copying the data twice.
--
-- As with 'allocaArray', the memory is freed once the action completes, so you
-- should not return the pointer from the action, and be wary of asynchronous
-- kernel execution.
--
withListArray :: Storable a => [a] -> (DevicePtr a -> IO b) -> IO b
withListArray xs = bracket (newListArray xs) free

--------------------------------------------------------------------------------
-- Utility
--------------------------------------------------------------------------------

-- |
-- Set a number of data elements to the specified value, which may be either 8-,
-- 16-, or 32-bits wide.
--
memset :: Storable a => DevicePtr a -> Int -> a -> IO ()
memset dptr n val = case sizeOf val of
    1 -> nothingIfOk =<< cuMemsetD8  dptr val n
    2 -> nothingIfOk =<< cuMemsetD16 dptr val n
    4 -> nothingIfOk =<< cuMemsetD32 dptr val n
    _ -> cudaError "can only memset 8-, 16-, and 32-bit values"

--
-- We use unsafe coerce below to reinterpret the bits of the value to memset as,
-- into the integer type required by the setting functions.
--
{# fun unsafe cuMemsetD8
  { useDevicePtr `DevicePtr a'
  , unsafeCoerce `a'
  ,              `Int'         } -> `Status' cToEnum #}

{# fun unsafe cuMemsetD16
  { useDevicePtr `DevicePtr a'
  , unsafeCoerce `a'
  ,              `Int'         } -> `Status' cToEnum #}

{# fun unsafe cuMemsetD32
  { useDevicePtr `DevicePtr a'
  , unsafeCoerce `a'
  ,              `Int'         } -> `Status' cToEnum #}

-- |
-- Return the device pointer associated with a mapped, pinned host buffer, which
-- was allocated with the 'DeviceMapped' option by 'mallocHostArray'.
--
-- Currently, no options are supported and this must be empty.
--
getDevicePtr :: [AllocFlag] -> HostPtr a -> IO (DevicePtr a)
getDevicePtr flags hp = resultIfOk =<< cuMemHostGetDevicePointer hp flags

{# fun unsafe cuMemHostGetDevicePointer
  { alloca'-        `DevicePtr a' peekDP*
  , useHP           `HostPtr a'
  , combineBitMasks `[AllocFlag]'         } -> `Status' cToEnum #}
  where
    alloca' = F.alloca
    useHP   = castPtr . useHostPtr
    peekDP  = liftM DevicePtr . peek

{-# LANGUAGE ForeignFunctionInterface #-}
--------------------------------------------------------------------------------
-- |
-- Module    : Foreign.CUDA.Driver.Marshal
-- Copyright : (c) 2009 Trevor L. McDonell
-- License   : BSD
--
-- Memory management for low-level driver interface
--
--------------------------------------------------------------------------------

module Foreign.CUDA.Driver.Marshal
  (
    DevicePtr, HostPtr, AllocFlag(..), withHostPtr,
    malloc, free, mallocHost, freeHost,
    peekArray, peekArrayAsync, pokeArray, pokeArrayAsync, memset,
    getDevicePtr
  )
  where

#include <cuda.h>
{# context lib="cuda" #}

-- Friends
import Foreign.CUDA.Internal.C2HS
import Foreign.CUDA.Driver.Error
import Foreign.CUDA.Driver.Stream               (Stream(..))

-- System
import Foreign                           hiding (peekArray, pokeArray, malloc, free)
import Foreign.C
import Control.Monad                            (liftM)
import Unsafe.Coerce

#c
typedef enum CUmemhostalloc_option_enum {
    CU_MEMHOSTALLOC_OPTION_PORTABLE       = CU_MEMHOSTALLOC_PORTABLE,
    CU_MEMHOSTALLOC_OPTION_DEVICE_MAPPED  = CU_MEMHOSTALLOC_DEVICEMAP,
    CU_MEMHOSTALLOC_OPTION_WRITE_COMBINED = CU_MEMHOSTALLOC_WRITECOMBINED
} CUmemhostalloc_option;
#endc

--------------------------------------------------------------------------------
-- Data Types
--------------------------------------------------------------------------------

-- |
-- A reference to memory allocated on the device
--
newtype DevicePtr a = DevicePtr { useDevicePtr :: {# type CUdeviceptr #}}

instance Storable (DevicePtr a) where
  sizeOf _      = sizeOf    (undefined :: {# type CUdeviceptr #})
  alignment _   = alignment (undefined :: {# type CUdeviceptr #})
  peek p        = DevicePtr `fmap` peek (castPtr p)
  poke p v      = poke (castPtr p) (useDevicePtr v)

-- |
-- A reference to memory on the host that is page-locked and directly-accessible
-- from the device. Since the memory can be accessed directly, it can be read or
-- written at a much higher bandwidth than pageable memory from the traditional
-- malloc.
--
-- The driver automatically accelerates calls to functions such as `memcpy'
-- which reference page-locked memory.
--
newtype HostPtr a = HostPtr { useHostPtr :: Ptr a }

-- |
-- Unwrap a host pointer and execute a computation using the base pointer object
--
withHostPtr :: HostPtr a -> (Ptr a -> IO b) -> IO b
withHostPtr p f = f (useHostPtr p)

-- |
-- Options for host allocation
--
{# enum CUmemhostalloc_option as AllocFlag
    { underscoreToCase }
    with prefix="CU_MEMHOSTALLOC_OPTION" deriving (Eq, Show) #}

--------------------------------------------------------------------------------
-- Host Allocation
--------------------------------------------------------------------------------

-- |
-- Allocate a section of linear memory on the host which is page-locked and
-- directly accessible from the device. The storage is sufficient to hold the
-- given number of elements of a storable type.
--
-- Note that since the amount of pageable memory is thusly reduced, overall
-- system performance may suffer. This is best used sparingly to allocate
-- staging areas for data exchange.
--
mallocHost :: Storable a => [AllocFlag] -> Int -> IO (HostPtr a)
mallocHost flags = doMalloc undefined
  where
    doMalloc :: Storable a' => a' -> Int -> IO (HostPtr a')
    doMalloc x n = resultIfOk =<< cuMemHostAlloc (n * sizeOf x) flags

{# fun unsafe cuMemHostAlloc
  { alloca'-        `HostPtr a'   peekHP*
  ,                 `Int'
  , combineBitMasks `[AllocFlag]'         } -> `Status' cToEnum #}
  where
    alloca'  = F.alloca
    peekHP p = (HostPtr . castPtr) `fmap` peek p

-- |
-- Free a section of page-locked host memory
--
freeHost :: HostPtr a -> IO ()
freeHost p = nothingIfOk =<< cuMemFreeHost p

{# fun unsafe cuMemFreeHost
  { useHP `HostPtr a' } -> `Status' cToEnum #}
  where
    useHP = castPtr . useHostPtr

--------------------------------------------------------------------------------
-- Device Allocation
--------------------------------------------------------------------------------

-- |
-- Allocate a section of linear memory on the device, and return a reference to
-- it. The memory is sufficient to hold the given number of elements of storable
-- type. It is suitably aligned for any type, and is not cleared.
--
malloc :: Storable a => Int -> IO (DevicePtr a)
malloc = doMalloc undefined
  where
    doMalloc :: Storable a' => a' -> Int -> IO (DevicePtr a')
    doMalloc x n = resultIfOk =<< cuMemAlloc (n * sizeOf x)

{# fun unsafe cuMemAlloc
  { alloca-  `DevicePtr a' dptr*
  ,          `Int'               } -> `Status' cToEnum #}
  where dptr = liftM DevicePtr . peek

-- |
-- Release a section of device memory
--
free :: DevicePtr a -> IO ()
free dp = nothingIfOk =<< cuMemFree dp

{# fun unsafe cuMemFree
  { useDevicePtr `DevicePtr a' } -> `Status' cToEnum #}

--------------------------------------------------------------------------------
-- Transfer
--------------------------------------------------------------------------------

-- |
-- Copy a number of elements from the device to host memory. This is a
-- synchronous operation
--
peekArray :: Storable a => Int -> DevicePtr a -> Ptr a -> IO ()
peekArray n dptr hptr = doPeek undefined dptr
  where
    doPeek :: Storable a' => a' -> DevicePtr a' -> IO ()
    doPeek x _ = nothingIfOk =<< cuMemcpyDtoH hptr dptr (n * sizeOf x)

{# fun unsafe cuMemcpyDtoH
  { castPtr      `Ptr a'
  , useDevicePtr `DevicePtr a'
  ,              `Int'         } -> `Status' cToEnum #}

-- |
-- Copy memory from the device asynchronously, possibly associated with a
-- particular stream. The destination host memory must be page-locked (allocated
-- by mallocHost)
--
peekArrayAsync :: Storable a => Int -> DevicePtr a -> Ptr a -> Maybe Stream -> IO ()
peekArrayAsync n dptr hptr mst = doPeek undefined dptr
  where
    doPeek :: Storable a' => a' -> DevicePtr a' -> IO ()
    doPeek x _ =
      nothingIfOk =<< case mst of
        Nothing -> cuMemcpyDtoHAsync hptr dptr (n * sizeOf x) (Stream nullPtr)
        Just st -> cuMemcpyDtoHAsync hptr dptr (n * sizeOf x) st

{# fun unsafe cuMemcpyDtoHAsync
  { castPtr      `Ptr a'
  , useDevicePtr `DevicePtr a'
  ,              `Int'
  , useStream    `Stream'  } -> `Status' cToEnum #}

-- |
-- Copy a number of elements onto the device. This is a synchronous operation
--
pokeArray :: Storable a => Int -> Ptr a -> DevicePtr a -> IO ()
pokeArray n hptr dptr = doPoke undefined dptr
  where
    doPoke :: Storable a' => a' -> DevicePtr a' -> IO ()
    doPoke x _ = nothingIfOk =<< cuMemcpyHtoD dptr hptr (n * sizeOf x)

{# fun unsafe cuMemcpyHtoD
  { useDevicePtr `DevicePtr a'
  , castPtr      `Ptr a'
  ,              `Int'         } -> `Status' cToEnum #}

-- |
-- Copy memory onto the device asynchronously, possibly associated with a
-- particular stream. The source host memory must be page-locked (allocated by
-- mallocHost)
--
pokeArrayAsync :: Storable a => Int -> Ptr a -> DevicePtr a -> Maybe Stream -> IO ()
pokeArrayAsync n hptr dptr mst = dopoke undefined dptr
  where
    dopoke :: Storable a' => a' -> DevicePtr a' -> IO ()
    dopoke x _ =
      nothingIfOk =<< case mst of
        Nothing -> cuMemcpyHtoDAsync dptr hptr (n * sizeOf x) (Stream nullPtr)
        Just st -> cuMemcpyHtoDAsync dptr hptr (n * sizeOf x) st

{# fun unsafe cuMemcpyHtoDAsync
  { useDevicePtr `DevicePtr a'
  , castPtr      `Ptr a'
  ,              `Int'
  , useStream    `Stream'  } -> `Status' cToEnum #}

--------------------------------------------------------------------------------
-- Utility
--------------------------------------------------------------------------------

-- |
-- Set a number of data elements to the specified value, which may be either 8-,
-- 16-, or 32-bits wide.
--
memset :: Storable a => DevicePtr a -> Int -> a -> IO ()
memset dptr n val = case sizeOf val of
    1 -> nothingIfOk =<< cuMemsetD8  dptr val n
    2 -> nothingIfOk =<< cuMemsetD16 dptr val n
    4 -> nothingIfOk =<< cuMemsetD32 dptr val n
    _ -> cudaError "can only memset 8-, 16-, and 32-bit values"

--
-- We use unsafe coerce below to reinterpret the bits of the value to memset as,
-- into the integer type required by the setting functions.
--
{# fun unsafe cuMemsetD8
  { useDevicePtr `DevicePtr a'
  , unsafeCoerce `a'
  ,              `Int'         } -> `Status' cToEnum #}

{# fun unsafe cuMemsetD16
  { useDevicePtr `DevicePtr a'
  , unsafeCoerce `a'
  ,              `Int'         } -> `Status' cToEnum #}

{# fun unsafe cuMemsetD32
  { useDevicePtr `DevicePtr a'
  , unsafeCoerce `a'
  ,              `Int'         } -> `Status' cToEnum #}

-- |
-- Return the device pointer associated with a mapped, pinned host buffer, which
-- was allocated with the 'DeviceMapped' option by 'mallocHost'.
--
-- Currently, no options are supported and this must be empty.
--
getDevicePtr :: [AllocFlag] -> HostPtr a -> IO (DevicePtr a)
getDevicePtr flags hp = resultIfOk =<< cuMemHostGetDevicePointer hp flags

{# fun unsafe cuMemHostGetDevicePointer
  { alloca'-        `DevicePtr a' peekDP*
  , useHP           `HostPtr a'
  , combineBitMasks `[AllocFlag]'         } -> `Status' cToEnum #}
  where
    alloca' = F.alloca
    useHP   = castPtr . useHostPtr
    peekDP  = liftM DevicePtr . peek


update with new combined alloca+memset

update with new combined alloca+memset

Ignore-this: f4560d7d75e1c6aac94be13acb58e6f8

darcs-hash:20090911082427-9241b-b8681ed812f5a0c36b5e602cb6807241f43ff0dd.gz

{-# LANGUAGE CPP, ForeignFunctionInterface #-}
--------------------------------------------------------------------------------
-- |
-- Module    : IonSeries
-- Copyright : (c) 2009 Trevor L. McDonell
-- License   : BSD
--
-- Functions to work with a series of ions, and generate theoretical spectra
-- suitable for matching against experimental results with the sequest
-- cross-correlation algorithm.
--
--------------------------------------------------------------------------------

module IonSeries
  (
    XCorrSpecThry(..),
    buildThrySpecXCorr
  )
  where

#include "kernels/kernels.h"

import Config
import Protein

import C2HS
import Foreign.CUDA (DevicePtr, withDevicePtr)
import qualified Foreign.CUDA as CUDA

--------------------------------------------------------------------------------
-- Data structures
--------------------------------------------------------------------------------

--
-- The mz/intensity spectrum array for the theoretical spectrum. Keep this as a
-- sparse array, as there are so few elements compared to the experimental data.
--
-- XXX: Changed to a dense array on the device, to facilitate eventual dot
-- product operation (and because my cuda-fu is weak...)
--
data XCorrSpecThry = XCorrSpecThry
        (Int,Int)
        (CUDA.DevicePtr Int)

--------------------------------------------------------------------------------
-- Theoretical Spectrum
--------------------------------------------------------------------------------

--
-- Generate the theoretical spectral representation of a peptide from its
-- character code sequence, and do something useful with it. The device memory
-- is deallocated once the action completes.
--
buildThrySpecXCorr :: ConfigParams
                   -> (Int,Int)                 -- ^ bounds of the output array
                   -> Int                       -- ^ precursor charge state
                   -> Peptide                   -- ^ peptide to build spectrum for
                   -> (XCorrSpecThry -> IO b)   -- ^ action to perform
                   -> IO b
buildThrySpecXCorr _cp (m,n) chrg pep action =
    CUDA.allocaBytesMemset bytes 0     $ \spec     -> do
    CUDA.withArrayLen (bIonLadder pep) $ \l b_ions -> do
    CUDA.withArray    (yIonLadder pep) $ \y_ions   -> do
    addIons chrg b_ions y_ions spec l len >> do

    action (XCorrSpecThry (m,n) spec)

    where
      len   = n - m + 1
      bytes = fromIntegral len * fromIntegral (sizeOf (undefined::Int))

{# fun unsafe addIons
    { cIntConv          `Int'             ,
      withDevicePtr*    `DevicePtr Float' ,
      withDevicePtr*    `DevicePtr Float' ,
      withDevicePtr*    `DevicePtr Int'   ,
      cIntConv          `Int'             ,
      cIntConv          `Int'             } -> `()' #}

#if 0
buildThrySpecXCorr :: ConfigParams -> Int -> Int -> Peptide -> IO XCorrSpecThry
buildThrySpecXCorr _cp len_spec charge peptide = do
    spec <- CUDA.forceEither `fmap` CUDA.malloc bytes
    rv   <- CUDA.memset spec bytes 0

    case rv of
      Just e  -> error e
      Nothing -> CUDA.withArrayLen (bIonLadder peptide) $ \len_ions b_ions -> do
                 CUDA.withArray    (yIonLadder peptide) $ \y_ions -> do
                   addIons charge b_ions y_ions spec len_ions len_spec

    return $ XCorrSpecThry (0,len_spec) spec

    where
      bytes = fromIntegral (len_spec * sizeOf (undefined::Int))
#endif
#if 0
--
-- Sequential version
--
buildThrySpecXCorr :: ConfigParams -> Float -> Peptide -> XCorrSpecThry
buildThrySpecXCorr _cp charge peptide =
    concatMap addIons $ [1 .. (max 1 (charge-1))]
    where
        addIons c = concatMap (addIonsAB c) b_ions ++ concatMap (addIonsY c) y_ions
        b_ions    = bIonLadder peptide
        y_ions    = yIonLadder peptide

--
-- Convert mass to mass/charge ratio
--
ionMZ :: Float -> Float -> Float
ionMZ m c = (m + massH*c) / c

--
-- Add a spectral peak for each fragment location, as well as the peaks
-- corresponding to neutral losses of H2O and NH3.
--
-- The factors that contribute to the collision induced dissociation process
-- used in tandem-MS experiments are not completely understood, so accurate
-- prediction of fragment ion abundances are not possible. Magnitude components
-- are assigned based on empirical knowledge.
--
addIonsAB, addIonsY :: Float -> Float -> [(Float, Float)]
addIonsAB charge mass = addIonsA : addIonsB
    where
        addIonsA = let m = ionMZ (mass - massCO) charge in (m, 10)
        addIonsB = let m = ionMZ mass charge in
          [
            (m,50), (m+1,25), (m-1,25),
            (m - massH2O/charge, 10),
            (m - massNH3/charge, 10)
          ]

addIonsY charge mass =
    let m = ionMZ (mass + massH2O) charge in
      [
        (m,50), (m+1,25), (m-1,25),
        (m - massNH3/charge, 10)
      ]
#endif

{-# LANGUAGE CPP, ForeignFunctionInterface #-}
--------------------------------------------------------------------------------
-- |
-- Module    : IonSeries
-- Copyright : (c) 2009 Trevor L. McDonell
-- License   : BSD
--
-- Functions to work with a series of ions, and generate theoretical spectra
-- suitable for matching against experimental results with the sequest
-- cross-correlation algorithm.
--
--------------------------------------------------------------------------------

module IonSeries where

#include "kernels/kernels.h"

import Config
import Protein

import C2HS
import qualified Foreign.CUDA as CUDA

--------------------------------------------------------------------------------
-- Data structures
--------------------------------------------------------------------------------

--
-- The mz/intensity spectrum array for the theoretical spectrum. Keep this as a
-- sparse array, as there are so few elements compared to the experimental data.
--
-- XXX: Changed to a dense array on the device, to facilitate eventual dot
-- product operation (and because my cuda-fu is weak...)
--
data XCorrSpecThry = XCorrSpecThry
        (Int,Int)
        (CUDA.DevicePtr Int)

--------------------------------------------------------------------------------
-- Theoretical Spectrum
--------------------------------------------------------------------------------

--
-- Generate the theoretical spectral representation of a peptide from its
-- character code sequence, and do something useful with it. The device memory
-- is deallocated once the action completes.
--
buildThrySpecXCorr :: ConfigParams
                   -> (Int,Int)                 -- ^ bounds of the output array
                   -> Int                       -- ^ precursor charge state
                   -> Peptide                   -- ^ peptide to build spectrum for
                   -> (XCorrSpecThry -> IO b)   -- ^ action to perform
                   -> IO b
buildThrySpecXCorr _cp (m,n) chrg pep fun =
    CUDA.allocaBytes bytes $ \spec ->
    CUDA.memset spec bytes 0 >>= \rv -> case rv of
      Just e  -> error e
      Nothing -> CUDA.withArrayLen (bIonLadder pep) $ \n b_ions -> do
                 CUDA.withArray    (yIonLadder pep) $ \y_ions   -> do
                 addIons chrg b_ions y_ions spec n len >> do

                 fun (XCorrSpecThry (m,n) spec)
    where
      len   = n - m + 1
      bytes = fromIntegral len * fromIntegral (sizeOf (undefined::Int))

#if 0
buildThrySpecXCorr :: ConfigParams -> Int -> Int -> Peptide -> IO XCorrSpecThry
buildThrySpecXCorr _cp len_spec charge peptide = do
    spec <- CUDA.forceEither `fmap` CUDA.malloc bytes
    rv   <- CUDA.memset spec bytes 0

    case rv of
      Just e  -> error e
      Nothing -> CUDA.withArrayLen (bIonLadder peptide) $ \len_ions b_ions -> do
                 CUDA.withArray    (yIonLadder peptide) $ \y_ions -> do
                   addIons charge b_ions y_ions spec len_ions len_spec

    return $ XCorrSpecThry (0,len_spec) spec

    where
      bytes = fromIntegral (len_spec * sizeOf (undefined::Int))
#endif

{# fun unsafe addIons
    { cIntConv          `Int'                  ,
      withDevicePtr*    `CUDA.DevicePtr Float' ,
      withDevicePtr*    `CUDA.DevicePtr Float' ,
      withDevicePtr*    `CUDA.DevicePtr Int'   ,
      cIntConv          `Int'                  ,
      cIntConv          `Int'                  } -> `()' #}
  where
    withDevicePtr = CUDA.withDevicePtr

#if 0
--
-- Sequential version
--
buildThrySpecXCorr :: ConfigParams -> Float -> Peptide -> XCorrSpecThry
buildThrySpecXCorr _cp charge peptide =
    concatMap addIons $ [1 .. (max 1 (charge-1))]
    where
        addIons c = concatMap (addIonsAB c) b_ions ++ concatMap (addIonsY c) y_ions
        b_ions    = bIonLadder peptide
        y_ions    = yIonLadder peptide

--
-- Convert mass to mass/charge ratio
--
ionMZ :: Float -> Float -> Float
ionMZ m c = (m + massH*c) / c

--
-- Add a spectral peak for each fragment location, as well as the peaks
-- corresponding to neutral losses of H2O and NH3.
--
-- The factors that contribute to the collision induced dissociation process
-- used in tandem-MS experiments are not completely understood, so accurate
-- prediction of fragment ion abundances are not possible. Magnitude components
-- are assigned based on empirical knowledge.
--
addIonsAB, addIonsY :: Float -> Float -> [(Float, Float)]
addIonsAB charge mass = addIonsA : addIonsB
    where
        addIonsA = let m = ionMZ (mass - massCO) charge in (m, 10)
        addIonsB = let m = ionMZ mass charge in
          [
            (m,50), (m+1,25), (m-1,25),
            (m - massH2O/charge, 10),
            (m - massNH3/charge, 10)
          ]

addIonsY charge mass =
    let m = ionMZ (mass + massH2O) charge in
      [
        (m,50), (m+1,25), (m-1,25),
        (m - massNH3/charge, 10)
      ]
#endif


INLINE pragma only useful inside WIN32 block

Move the INLINE pragma inside the WIN32 block. For other platforms we just have a plain foreign C import, so the INLINE pragma is not useful (I don't think).

Moreover, GHC-HEAD complains about this.

{-# LANGUAGE BangPatterns             #-}
{-# LANGUAGE ForeignFunctionInterface #-}
{-# OPTIONS_HADDOCK prune #-}
--------------------------------------------------------------------------------
-- |
-- Module    : Foreign.CUDA.Driver.Module.Base
-- Copyright : [2009..2014] Trevor L. McDonell
-- License   : BSD
--
-- Module loading for low-level driver interface
--
--------------------------------------------------------------------------------

module Foreign.CUDA.Driver.Module.Base (

  -- * Module Management
  Module(..),
  JITOption(..), JITTarget(..), JITResult(..), JITFallback(..), JITInputType(..),
  JITOptionInternal(..),

  -- ** Loading and unloading modules
  loadFile,
  loadData,   loadDataFromPtr,
  loadDataEx, loadDataFromPtrEx,
  unload,

  -- Internal
  jitOptionUnpack, jitTargetOfCompute,

) where

#include "cbits/stubs.h"
{# context lib="cuda" #}

-- Friends
import Foreign.CUDA.Analysis.Device
import Foreign.CUDA.Driver.Error
import Foreign.CUDA.Internal.C2HS

-- System
import Foreign
import Foreign.C
import Unsafe.Coerce

import Control.Monad                                    ( liftM )
import Data.ByteString.Char8                            ( ByteString )
import qualified Data.ByteString.Char8                  as B
import qualified Data.ByteString.Internal               as B

--------------------------------------------------------------------------------
-- Data Types
--------------------------------------------------------------------------------

-- |
-- A reference to a Module object, containing collections of device functions
--
newtype Module = Module { useModule :: {# type CUmodule #}}
  deriving (Eq, Show)

-- |
-- Just-in-time compilation and linking options
--
data JITOption
  = MaxRegisters       !Int             -- ^ maximum number of registers per thread
  | ThreadsPerBlock    !Int             -- ^ number of threads per block to target for
  | OptimisationLevel  !Int             -- ^ level of optimisation to apply (1-4, default 4)
  | Target             !Compute         -- ^ compilation target, otherwise determined from context
  | FallbackStrategy   !JITFallback     -- ^ fallback strategy if matching cubin not found
  | GenerateDebugInfo                   -- ^ generate debug info (-g) (requires cuda >= 5.5)
  | GenerateLineInfo                    -- ^ generate line number information (-lineinfo) (requires cuda >= 5.5)
  | Verbose                             -- ^ verbose log messages (requires cuda >= 5.5)
  deriving (Show)

-- |
-- Results of online compilation
--
data JITResult = JITResult
  {
    jitTime     :: !Float,              -- ^ milliseconds spent compiling PTX
    jitInfoLog  :: !ByteString,         -- ^ information about PTX assembly
    jitModule   :: !Module              -- ^ the compiled module
  }
  deriving (Show)

-- |
-- Online compilation target architecture
--
{# enum CUjit_target as JITTarget
    { underscoreToCase }
    with prefix="CU_TARGET" deriving (Eq, Show) #}

-- |
-- Online compilation fallback strategy
--
{# enum CUjit_fallback as JITFallback
    { underscoreToCase
    , CU_PREFER_PTX as PreferPTX }
    with prefix="CU" deriving (Eq, Show) #}

-- |
-- Device code formats that can be used for online linking
--
{# enum CUjitInputType as JITInputType
    { underscoreToCase
    , CU_JIT_INPUT_PTX as PTX }
    with prefix="CU_JIT_INPUT" deriving (Eq, Show) #}

{# enum CUjit_option as JITOptionInternal
    { }
    with prefix="CU" deriving (Eq, Show) #}

--------------------------------------------------------------------------------
-- Module management
--------------------------------------------------------------------------------

-- |
-- Load the contents of the specified file (either a ptx or cubin file) to
-- create a new module, and load that module into the current context.
--
-- <http://docs.nvidia.com/cuda/cuda-driver-api/group__CUDA__MODULE.html#group__CUDA__MODULE_1g366093bd269dafd0af21f1c7d18115d3>
--
{-# INLINEABLE loadFile #-}
loadFile :: FilePath -> IO Module
loadFile !ptx = resultIfOk =<< cuModuleLoad ptx

{-# INLINE cuModuleLoad #-}
{# fun unsafe cuModuleLoad
  { alloca-      `Module'   peekMod*
  , withCString* `FilePath'          } -> `Status' cToEnum #}

-- |
-- Load the contents of the given image into a new module, and load that module
-- into the current context. The image is (typically) the contents of a cubin or
-- PTX file.
--
-- Note that the 'ByteString' will be copied into a temporary staging area so
-- that it can be passed to C.
--
-- <http://docs.nvidia.com/cuda/cuda-driver-api/group__CUDA__MODULE.html#group__CUDA__MODULE_1g04ce266ce03720f479eab76136b90c0b>
--
{-# INLINEABLE loadData #-}
loadData :: ByteString -> IO Module
loadData !img =
  B.useAsCString img (\p -> loadDataFromPtr (castPtr p))

-- |
-- As 'loadData', but read the image data from the given pointer. The image is a
-- NULL-terminated sequence of bytes.
--
{-# INLINEABLE loadDataFromPtr #-}
loadDataFromPtr :: Ptr Word8 -> IO Module
loadDataFromPtr !img = resultIfOk =<< cuModuleLoadData img

{-# INLINE cuModuleLoadData #-}
{# fun unsafe cuModuleLoadData
  { alloca- `Module'    peekMod*
  , castPtr `Ptr Word8'          } -> ` Status' cToEnum #}

-- |
-- Load the contents of the given image into a module with online compiler
-- options, and load the module into the current context. The image is
-- (typically) the contents of a cubin or PTX file. The actual attributes of the
-- compiled kernel can be probed using 'requires'.
--
-- Note that the 'ByteString' will be copied into a temporary staging area so
-- that it can be passed to C.
--
-- <http://docs.nvidia.com/cuda/cuda-driver-api/group__CUDA__MODULE.html#group__CUDA__MODULE_1g9e8047e9dbf725f0cd7cafd18bfd4d12>
--
{-# INLINEABLE loadDataEx #-}
loadDataEx :: ByteString -> [JITOption] -> IO JITResult
loadDataEx !img !options =
  B.useAsCString img (\p -> loadDataFromPtrEx (castPtr p) options)

-- |
-- As 'loadDataEx', but read the image data from the given pointer. The image is
-- a NULL-terminated sequence of bytes.
--
{-# INLINEABLE loadDataFromPtrEx #-}
loadDataFromPtrEx :: Ptr Word8 -> [JITOption] -> IO JITResult
loadDataFromPtrEx !img !options = do
  let logSize = 2048

  fp_ilog <- B.mallocByteString logSize

  allocaArray logSize    $ \p_elog -> do
  withForeignPtr fp_ilog $ \p_ilog -> do

  let (opt,val) = unzip $
        [ (JIT_WALL_TIME, 0) -- must be first, this is extracted below
        , (JIT_INFO_LOG_BUFFER_SIZE_BYTES,  logSize)
        , (JIT_ERROR_LOG_BUFFER_SIZE_BYTES, logSize)
        , (JIT_INFO_LOG_BUFFER,  unsafeCoerce (p_ilog :: CString))
        , (JIT_ERROR_LOG_BUFFER, unsafeCoerce (p_elog :: CString))
        ]
        ++
        map jitOptionUnpack options

  withArrayLen (map cFromEnum opt)    $ \i p_opts -> do
  withArray    (map unsafeCoerce val) $ \  p_vals -> do

  (s,mdl) <- cuModuleLoadDataEx img i p_opts p_vals

  case s of
    Success     -> do
      time    <- peek (castPtr p_vals)
      infoLog <- B.fromForeignPtr (castForeignPtr fp_ilog) 0 `fmap` c_strnlen p_ilog logSize
      return  $! JITResult time infoLog mdl

    _           -> do
      errLog  <- peekCString p_elog
      cudaError (unlines [describe s, errLog])

{-# INLINE cuModuleLoadDataEx #-}
{# fun unsafe cuModuleLoadDataEx
  { alloca- `Module'       peekMod*
  , castPtr `Ptr Word8'
  ,         `Int'
  , id      `Ptr CInt'
  , id      `Ptr (Ptr ())'          } -> `Status' cToEnum #}

-- |
-- Unload a module from the current context.
--
-- <http://docs.nvidia.com/cuda/cuda-driver-api/group__CUDA__MODULE.html#group__CUDA__MODULE_1g8ea3d716524369de3763104ced4ea57b>
--
{-# INLINEABLE unload #-}
unload :: Module -> IO ()
unload !m = nothingIfOk =<< cuModuleUnload m

{-# INLINE cuModuleUnload #-}
{# fun unsafe cuModuleUnload
  { useModule `Module' } -> `Status' cToEnum #}

--------------------------------------------------------------------------------
-- Internal
--------------------------------------------------------------------------------

{-# INLINE peekMod #-}
peekMod :: Ptr {# type CUmodule #} -> IO Module
peekMod = liftM Module . peek

{-# INLINE jitOptionUnpack #-}
jitOptionUnpack :: JITOption -> (JITOptionInternal, Int)
jitOptionUnpack (MaxRegisters x)      = (JIT_MAX_REGISTERS,       x)
jitOptionUnpack (ThreadsPerBlock x)   = (JIT_THREADS_PER_BLOCK,   x)
jitOptionUnpack (OptimisationLevel x) = (JIT_OPTIMIZATION_LEVEL,  x)
jitOptionUnpack (Target x)            = (JIT_TARGET,              fromEnum (jitTargetOfCompute x))
jitOptionUnpack (FallbackStrategy x)  = (JIT_FALLBACK_STRATEGY,   fromEnum x)
#if CUDA_VERSION >= 5050
jitOptionUnpack GenerateDebugInfo     = (JIT_GENERATE_DEBUG_INFO, fromEnum True)
jitOptionUnpack GenerateLineInfo      = (JIT_GENERATE_LINE_INFO,  fromEnum True)
jitOptionUnpack Verbose               = (JIT_LOG_VERBOSE,         fromEnum True)
#else
jitOptionUnpack GenerateDebugInfo     = requireSDK 'GenerateDebugInfo 5.5
jitOptionUnpack GenerateLineInfo      = requireSDK 'GenerateLineInfo 5.5
jitOptionUnpack Verbose               = requireSDK 'Verbose 5.5
#endif

{-# INLINE jitTargetOfCompute #-}
jitTargetOfCompute :: Compute -> JITTarget
jitTargetOfCompute (Compute 1 0) = Compute10
jitTargetOfCompute (Compute 1 1) = Compute11
jitTargetOfCompute (Compute 1 2) = Compute12
jitTargetOfCompute (Compute 1 3) = Compute13
jitTargetOfCompute (Compute 2 0) = Compute20
jitTargetOfCompute (Compute 2 1) = Compute21
jitTargetOfCompute (Compute 3 0) = Compute30
jitTargetOfCompute (Compute 3 5) = Compute35
#if CUDA_VERSION >= 6000
jitTargetOfCompute (Compute 3 2) = Compute32
jitTargetOfCompute (Compute 5 0) = Compute50
#endif
#if CUDA_VERSION >= 6050
jitTargetOfCompute (Compute 3 7) = Compute37
#endif
#if CUDA_VERSION >= 7000
jitTargetOfCompute (Compute 5 2) = Compute52
#endif
jitTargetOfCompute compute       = error ("Unknown JIT Target for Compute " ++ show compute)

#if defined(WIN32)
{-# INLINE c_strnlen' #-}
c_strnlen' :: CString -> CSize -> IO CSize
c_strnlen' str size = do
  str' <- peekCStringLen (str, fromIntegral size)
  return $ stringLen 0 str'
  where
    stringLen acc []       = acc
    stringLen acc ('\0':_) = acc
    stringLen acc (_:xs)   = stringLen (acc+1) xs
#else
foreign import ccall unsafe "string.h strnlen" c_strnlen'
  :: CString -> CSize -> IO CSize
#endif

{-# INLINE c_strnlen #-}
c_strnlen :: CString -> Int -> IO Int
c_strnlen str maxlen = cIntConv `fmap` c_strnlen' str (cIntConv maxlen)

{-# LANGUAGE BangPatterns             #-}
{-# LANGUAGE ForeignFunctionInterface #-}
{-# OPTIONS_HADDOCK prune #-}
--------------------------------------------------------------------------------
-- |
-- Module    : Foreign.CUDA.Driver.Module.Base
-- Copyright : [2009..2014] Trevor L. McDonell
-- License   : BSD
--
-- Module loading for low-level driver interface
--
--------------------------------------------------------------------------------

module Foreign.CUDA.Driver.Module.Base (

  -- * Module Management
  Module(..),
  JITOption(..), JITTarget(..), JITResult(..), JITFallback(..), JITInputType(..),
  JITOptionInternal(..),

  -- ** Loading and unloading modules
  loadFile,
  loadData,   loadDataFromPtr,
  loadDataEx, loadDataFromPtrEx,
  unload,

  -- Internal
  jitOptionUnpack, jitTargetOfCompute,

) where

#include "cbits/stubs.h"
{# context lib="cuda" #}

-- Friends
import Foreign.CUDA.Analysis.Device
import Foreign.CUDA.Driver.Error
import Foreign.CUDA.Internal.C2HS

-- System
import Foreign
import Foreign.C
import Unsafe.Coerce

import Control.Monad                                    ( liftM )
import Data.ByteString.Char8                            ( ByteString )
import qualified Data.ByteString.Char8                  as B
import qualified Data.ByteString.Internal               as B

--------------------------------------------------------------------------------
-- Data Types
--------------------------------------------------------------------------------

-- |
-- A reference to a Module object, containing collections of device functions
--
newtype Module = Module { useModule :: {# type CUmodule #}}
  deriving (Eq, Show)

-- |
-- Just-in-time compilation and linking options
--
data JITOption
  = MaxRegisters       !Int             -- ^ maximum number of registers per thread
  | ThreadsPerBlock    !Int             -- ^ number of threads per block to target for
  | OptimisationLevel  !Int             -- ^ level of optimisation to apply (1-4, default 4)
  | Target             !Compute         -- ^ compilation target, otherwise determined from context
  | FallbackStrategy   !JITFallback     -- ^ fallback strategy if matching cubin not found
  | GenerateDebugInfo                   -- ^ generate debug info (-g) (requires cuda >= 5.5)
  | GenerateLineInfo                    -- ^ generate line number information (-lineinfo) (requires cuda >= 5.5)
  | Verbose                             -- ^ verbose log messages (requires cuda >= 5.5)
  deriving (Show)

-- |
-- Results of online compilation
--
data JITResult = JITResult
  {
    jitTime     :: !Float,              -- ^ milliseconds spent compiling PTX
    jitInfoLog  :: !ByteString,         -- ^ information about PTX assembly
    jitModule   :: !Module              -- ^ the compiled module
  }
  deriving (Show)

-- |
-- Online compilation target architecture
--
{# enum CUjit_target as JITTarget
    { underscoreToCase }
    with prefix="CU_TARGET" deriving (Eq, Show) #}

-- |
-- Online compilation fallback strategy
--
{# enum CUjit_fallback as JITFallback
    { underscoreToCase
    , CU_PREFER_PTX as PreferPTX }
    with prefix="CU" deriving (Eq, Show) #}

-- |
-- Device code formats that can be used for online linking
--
{# enum CUjitInputType as JITInputType
    { underscoreToCase
    , CU_JIT_INPUT_PTX as PTX }
    with prefix="CU_JIT_INPUT" deriving (Eq, Show) #}

{# enum CUjit_option as JITOptionInternal
    { }
    with prefix="CU" deriving (Eq, Show) #}

--------------------------------------------------------------------------------
-- Module management
--------------------------------------------------------------------------------

-- |
-- Load the contents of the specified file (either a ptx or cubin file) to
-- create a new module, and load that module into the current context.
--
-- <http://docs.nvidia.com/cuda/cuda-driver-api/group__CUDA__MODULE.html#group__CUDA__MODULE_1g366093bd269dafd0af21f1c7d18115d3>
--
{-# INLINEABLE loadFile #-}
loadFile :: FilePath -> IO Module
loadFile !ptx = resultIfOk =<< cuModuleLoad ptx

{-# INLINE cuModuleLoad #-}
{# fun unsafe cuModuleLoad
  { alloca-      `Module'   peekMod*
  , withCString* `FilePath'          } -> `Status' cToEnum #}

-- |
-- Load the contents of the given image into a new module, and load that module
-- into the current context. The image is (typically) the contents of a cubin or
-- PTX file.
--
-- Note that the 'ByteString' will be copied into a temporary staging area so
-- that it can be passed to C.
--
-- <http://docs.nvidia.com/cuda/cuda-driver-api/group__CUDA__MODULE.html#group__CUDA__MODULE_1g04ce266ce03720f479eab76136b90c0b>
--
{-# INLINEABLE loadData #-}
loadData :: ByteString -> IO Module
loadData !img =
  B.useAsCString img (\p -> loadDataFromPtr (castPtr p))

-- |
-- As 'loadData', but read the image data from the given pointer. The image is a
-- NULL-terminated sequence of bytes.
--
{-# INLINEABLE loadDataFromPtr #-}
loadDataFromPtr :: Ptr Word8 -> IO Module
loadDataFromPtr !img = resultIfOk =<< cuModuleLoadData img

{-# INLINE cuModuleLoadData #-}
{# fun unsafe cuModuleLoadData
  { alloca- `Module'    peekMod*
  , castPtr `Ptr Word8'          } -> ` Status' cToEnum #}

-- |
-- Load the contents of the given image into a module with online compiler
-- options, and load the module into the current context. The image is
-- (typically) the contents of a cubin or PTX file. The actual attributes of the
-- compiled kernel can be probed using 'requires'.
--
-- Note that the 'ByteString' will be copied into a temporary staging area so
-- that it can be passed to C.
--
-- <http://docs.nvidia.com/cuda/cuda-driver-api/group__CUDA__MODULE.html#group__CUDA__MODULE_1g9e8047e9dbf725f0cd7cafd18bfd4d12>
--
{-# INLINEABLE loadDataEx #-}
loadDataEx :: ByteString -> [JITOption] -> IO JITResult
loadDataEx !img !options =
  B.useAsCString img (\p -> loadDataFromPtrEx (castPtr p) options)

-- |
-- As 'loadDataEx', but read the image data from the given pointer. The image is
-- a NULL-terminated sequence of bytes.
--
{-# INLINEABLE loadDataFromPtrEx #-}
loadDataFromPtrEx :: Ptr Word8 -> [JITOption] -> IO JITResult
loadDataFromPtrEx !img !options = do
  let logSize = 2048

  fp_ilog <- B.mallocByteString logSize

  allocaArray logSize    $ \p_elog -> do
  withForeignPtr fp_ilog $ \p_ilog -> do

  let (opt,val) = unzip $
        [ (JIT_WALL_TIME, 0) -- must be first, this is extracted below
        , (JIT_INFO_LOG_BUFFER_SIZE_BYTES,  logSize)
        , (JIT_ERROR_LOG_BUFFER_SIZE_BYTES, logSize)
        , (JIT_INFO_LOG_BUFFER,  unsafeCoerce (p_ilog :: CString))
        , (JIT_ERROR_LOG_BUFFER, unsafeCoerce (p_elog :: CString))
        ]
        ++
        map jitOptionUnpack options

  withArrayLen (map cFromEnum opt)    $ \i p_opts -> do
  withArray    (map unsafeCoerce val) $ \  p_vals -> do

  (s,mdl) <- cuModuleLoadDataEx img i p_opts p_vals

  case s of
    Success     -> do
      time    <- peek (castPtr p_vals)
      infoLog <- B.fromForeignPtr (castForeignPtr fp_ilog) 0 `fmap` c_strnlen p_ilog logSize
      return  $! JITResult time infoLog mdl

    _           -> do
      errLog  <- peekCString p_elog
      cudaError (unlines [describe s, errLog])

{-# INLINE cuModuleLoadDataEx #-}
{# fun unsafe cuModuleLoadDataEx
  { alloca- `Module'       peekMod*
  , castPtr `Ptr Word8'
  ,         `Int'
  , id      `Ptr CInt'
  , id      `Ptr (Ptr ())'          } -> `Status' cToEnum #}

-- |
-- Unload a module from the current context.
--
-- <http://docs.nvidia.com/cuda/cuda-driver-api/group__CUDA__MODULE.html#group__CUDA__MODULE_1g8ea3d716524369de3763104ced4ea57b>
--
{-# INLINEABLE unload #-}
unload :: Module -> IO ()
unload !m = nothingIfOk =<< cuModuleUnload m

{-# INLINE cuModuleUnload #-}
{# fun unsafe cuModuleUnload
  { useModule `Module' } -> `Status' cToEnum #}

--------------------------------------------------------------------------------
-- Internal
--------------------------------------------------------------------------------

{-# INLINE peekMod #-}
peekMod :: Ptr {# type CUmodule #} -> IO Module
peekMod = liftM Module . peek

{-# INLINE jitOptionUnpack #-}
jitOptionUnpack :: JITOption -> (JITOptionInternal, Int)
jitOptionUnpack (MaxRegisters x)      = (JIT_MAX_REGISTERS,       x)
jitOptionUnpack (ThreadsPerBlock x)   = (JIT_THREADS_PER_BLOCK,   x)
jitOptionUnpack (OptimisationLevel x) = (JIT_OPTIMIZATION_LEVEL,  x)
jitOptionUnpack (Target x)            = (JIT_TARGET,              fromEnum (jitTargetOfCompute x))
jitOptionUnpack (FallbackStrategy x)  = (JIT_FALLBACK_STRATEGY,   fromEnum x)
#if CUDA_VERSION >= 5050
jitOptionUnpack GenerateDebugInfo     = (JIT_GENERATE_DEBUG_INFO, fromEnum True)
jitOptionUnpack GenerateLineInfo      = (JIT_GENERATE_LINE_INFO,  fromEnum True)
jitOptionUnpack Verbose               = (JIT_LOG_VERBOSE,         fromEnum True)
#else
jitOptionUnpack GenerateDebugInfo     = requireSDK 'GenerateDebugInfo 5.5
jitOptionUnpack GenerateLineInfo      = requireSDK 'GenerateLineInfo 5.5
jitOptionUnpack Verbose               = requireSDK 'Verbose 5.5
#endif

{-# INLINE jitTargetOfCompute #-}
jitTargetOfCompute :: Compute -> JITTarget
jitTargetOfCompute (Compute 1 0) = Compute10
jitTargetOfCompute (Compute 1 1) = Compute11
jitTargetOfCompute (Compute 1 2) = Compute12
jitTargetOfCompute (Compute 1 3) = Compute13
jitTargetOfCompute (Compute 2 0) = Compute20
jitTargetOfCompute (Compute 2 1) = Compute21
jitTargetOfCompute (Compute 3 0) = Compute30
jitTargetOfCompute (Compute 3 5) = Compute35
#if CUDA_VERSION >= 6000
jitTargetOfCompute (Compute 3 2) = Compute32
jitTargetOfCompute (Compute 5 0) = Compute50
#endif
#if CUDA_VERSION >= 6050
jitTargetOfCompute (Compute 3 7) = Compute37
#endif
#if CUDA_VERSION >= 7000
jitTargetOfCompute (Compute 5 2) = Compute52
#endif
jitTargetOfCompute compute       = error ("Unknown JIT Target for Compute " ++ show compute)

{-# INLINE c_strnlen' #-}
#if defined(WIN32)
c_strnlen' :: CString -> CSize -> IO CSize
c_strnlen' str size = do
  str' <- peekCStringLen (str, fromIntegral size)
  return $ stringLen 0 str'
  where
    stringLen acc []       = acc
    stringLen acc ('\0':_) = acc
    stringLen acc (_:xs)   = stringLen (acc+1) xs
#else
foreign import ccall unsafe "string.h strnlen" c_strnlen'
  :: CString -> CSize -> IO CSize
#endif

{-# INLINE c_strnlen #-}
c_strnlen :: CString -> Int -> IO Int
c_strnlen str maxlen = cIntConv `fmap` c_strnlen' str (cIntConv maxlen)