{-# LANGUAGE BangPatterns             #-}
{-# LANGUAGE CPP                      #-}
{-# LANGUAGE ForeignFunctionInterface #-}
{-# LANGUAGE MagicHash                #-}
{-# LANGUAGE TemplateHaskell          #-}
{-# LANGUAGE UnboxedTuples            #-}
{-# OPTIONS_GHC -funbox-strict-fields #-}
--------------------------------------------------------------------------------
-- |
-- Module    : Foreign.NVVM.Compile
-- Copyright : [2016] Trevor L. McDonell
-- License   : BSD
--
-- Program compilation
--
--------------------------------------------------------------------------------

module Foreign.NVVM.Compile (

  Program,
  Result(..),
  CompileOption(..),

  compileModule, compileModules,

  create,
  destroy,
  addModule,     addModuleFromPtr,
  addModuleLazy, addModuleLazyFromPtr,
  compile,
  verify

) where

import Foreign.CUDA.Analysis
import Foreign.NVVM.Error
import Foreign.NVVM.Internal.C2HS

import Foreign.C
import Foreign.Marshal
import Foreign.Ptr
import Foreign.ForeignPtr
import Foreign.Storable

import Control.Exception
import Data.Word
import Text.Printf
import Data.ByteString                                              ( ByteString )
import Data.ByteString.Short                                        ( ShortByteString )
import qualified Data.ByteString.Char8                              as B
import qualified Data.ByteString.Unsafe                             as B
import qualified Data.ByteString.Internal                           as B
import qualified Data.ByteString.Short                              as BS
import qualified Data.ByteString.Short.Internal                     as BSI

import GHC.Exts
import GHC.Base                                                     ( IO(..) )


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


-- | An NVVM program
--
newtype Program = Program { useProgram :: {# type nvvmProgram #} }
  deriving ( Eq, Show )

-- | The result of compiling an NVVM program.
--
data Result = Result
  { compileResult :: !ByteString  -- ^ The compiled kernel, which can be loaded into the current program using 'Foreign.CUDA.Driver.loadData*'
  , compileLog    :: !ByteString  -- ^ Warning messages generated by the compiler/verifier
  }

-- | Program compilation options
--
data CompileOption
  = OptimisationLevel !Int        -- ^ optimisation level, from 0 (disable optimisations) to 3 (default)
  | Target !Compute               -- ^ target architecture to compile for (default: compute 2.0)
  | FlushToZero                   -- ^ flush denormal values to zero when performing single-precision floating-point operations (default: no)
  | NoFMA                         -- ^ disable fused-multiply-add instructions (default: enabled)
  | FastSqrt                      -- ^ use a fast approximation for single-precision floating-point square root (default: no)
  | FastDiv                       -- ^ use a fast approximation for single-precision floating-point division and reciprocal (default: no)
  | GenerateDebugInfo             -- ^ generate debugging information (-g) (default: no)
  deriving ( Eq, Show )


-- | Compile an NVVM IR module, in either bitcode or textual representation,
-- into PTX code.
--
{-# INLINEABLE compileModule #-}
compileModule
    :: ShortByteString            -- ^ name of the module
    -> ByteString                 -- ^ NVVM IR in either textual or bitcode representation
    -> [CompileOption]            -- ^ compiler options
    -> IO Result
compileModule !name !bs !opts =
  compileModules [(name,bs)] opts


-- | Compile a collection of NVVM IR modules into PTX code
--
{-# INLINEABLE compileModules #-}
compileModules
    :: [(ShortByteString, ByteString)]  -- ^ (module name, module NVVM IR) pairs to compile
    -> [CompileOption]                  -- ^ compiler options
    -> IO Result
compileModules !bss !opts =
  bracket create destroy $ \prg -> do
    mapM_ (uncurry (addModule prg)) bss
    (messages, result) <- compile prg opts
    case result of
      Nothing  -> nvvmErrorIO (B.unpack messages)
      Just ptx -> return $ Result ptx messages


-- | Create an empty 'Program'
--
-- <http://docs.nvidia.com/cuda/libnvvm-api/group__compilation.html#group__compilation_1g46a0ab04a063cba28bfbb41a1939e3f4>
--
{-# INLINEABLE create #-}
{# fun unsafe nvvmCreateProgram as create
    { alloca- `Program' peekProgram*
    }
    -> `()' checkStatus*-
#}


-- | Destroy a 'Program'
--
-- <http://docs.nvidia.com/cuda/libnvvm-api/group__compilation.html#group__compilation_1gfba94cab1224c0152841b80690d366aa>
--
{-# INLINEABLE destroy #-}
{#
  fun unsafe nvvmDestroyProgram as destroy
    { withProgram* `Program'
    }
    -> `()' checkStatus*-
#}


-- | Add a module level NVVM IR to a program
--
-- <http://docs.nvidia.com/cuda/libnvvm-api/group__compilation.html#group__compilation_1g0c22d2b9be033c165bc37b16f3ed75c6>
--
{-# INLINEABLE addModule #-}
addModule
    :: Program              -- ^ NVVM program to add to
    -> ShortByteString      -- ^ Name of the module (defaults to \"@\<unnamed\>@\" if empty)
    -> ByteString           -- ^ NVVM IR module in either bitcode or textual representation
    -> IO ()
addModule !prg !name !bs =
  B.unsafeUseAsCStringLen bs $ \(ptr,len) ->
  addModuleFromPtr prg name len (castPtr ptr)


-- | As with 'addModule', but read the specified number of bytes from the given
-- pointer.
--
{-# INLINEABLE addModuleFromPtr #-}
addModuleFromPtr
    :: Program              -- ^ NVVM program to add to
    -> ShortByteString      -- ^ Name of the module (defaults to \"@\<unnamed\>@\" if empty)
    -> Int                  -- ^ Number of bytes in the module
    -> Ptr Word8            -- ^ NVVM IR module in bitcode or textual representation
    -> IO ()
addModuleFromPtr !prg !name !size !buffer =
  nvvmAddModuleToProgram prg buffer size name
  where
    {#
      fun unsafe nvvmAddModuleToProgram
        { useProgram    `Program'
        , castPtr       `Ptr Word8'
        , cIntConv      `Int'
        , useAsCString* `ShortByteString'
        }
        -> `()' checkStatus*-
    #}


-- | Add a module level NVVM IR to a program.
--
-- The module is loaded lazily: only symbols required by modules loaded using
-- 'addModule' or 'addModuleFromPtr' will be loaded.
--
-- Requires CUDA-10.0
--
-- <https://docs.nvidia.com/cuda/libnvvm-api/group__compilation.html#group__compilation_1g5356ce5063db232cd4330b666c62219b>
--
-- @since 0.9.0.0
--
{-# INLINEABLE addModuleLazy #-}
addModuleLazy
    :: Program              -- ^ NVVM program to add to
    -> ShortByteString      -- ^ Name of the module (defaults to \"@\<unnamed\>@\" if empty)
    -> ByteString           -- ^ NVVM IR module in either bitcode or textual representation
    -> IO ()
#if CUDA_VERSION < 10000
addModuleLazy = requireSDK 'addModuleLazy 10.0
#else
addModuleLazy !prg !name !bs =
  B.unsafeUseAsCStringLen bs $ \(buffer, size) ->
  addModuleLazyFromPtr prg name size (castPtr buffer)
#endif


-- | As with 'addModuleLazy', but read the specified number of bytes from the
-- given pointer (the symbols are loaded lazily, the data in the buffer will be
-- read immediately).
--
-- Requires CUDA-10.0
--
-- @since 0.9.0.0
--
{-# INLINEABLE addModuleLazyFromPtr #-}
addModuleLazyFromPtr
    :: Program              -- ^ NVVM program to add to
    -> ShortByteString      -- ^ Name of the module (defaults to \"@\<unnamed\>@\" if empty)
    -> Int                  -- ^ Number of bytes in the module
    -> Ptr Word8            -- ^ NVVM IR in bitcode or textual representation
    -> IO ()
#if CUDA_VERSION < 10000
addModuleLazyFromPtr = requireSDK 'addModuleLazyFromPtr 10.0
#else
addModuleLazyFromPtr !prg !name !size !buffer =
  nvvmLazyAddModuleToProgram prg buffer size name
  where
    {#
      fun unsafe nvvmLazyAddModuleToProgram
        { useProgram    `Program'
        , castPtr       `Ptr Word8'
        , cIntConv      `Int'
        , useAsCString* `ShortByteString'
        }
        -> `()' checkStatus*-
    #}
#endif


-- | Compile the NVVM program. Returns the log from the compiler/verifier and,
-- if successful, the compiled program.
--
-- <http://docs.nvidia.com/cuda/libnvvm-api/group__compilation.html#group__compilation_1g76ac1e23f5d0e2240e78be0e63450346>
--
{-# INLINEABLE compile #-}
compile :: Program -> [CompileOption] -> IO (ByteString, Maybe ByteString)
compile !prg !opts = do
  status    <- withCompileOptions opts (nvvmCompileProgram prg)
  messages  <- retrieve (nvvmGetProgramLogSize prg) (nvvmGetProgramLog prg)
  case status of
    Success -> do ptx <- retrieve (nvvmGetCompiledResultSize prg) (nvvmGetCompiledResult prg)
                  return (messages, Just ptx)
    _       ->    return (messages, Nothing)
  where
    {# fun unsafe nvvmCompileProgram
        { useProgram `Program'
        , cIntConv   `Int'
        , id         `Ptr CString'
        }
        -> `Status' cToEnum
    #}

    {# fun unsafe nvvmGetCompiledResultSize
        { useProgram `Program'
        , alloca-    `Int'     peekIntConv*
        }
        -> `()' checkStatus*-
    #}

    {# fun unsafe nvvmGetCompiledResult
        { useProgram       `Program'
        , withForeignPtr'* `ForeignPtr Word8'
        }
        -> `()' checkStatus*-
    #}


-- | Verify the NVVM program. Returns whether compilation will succeed, together
-- with any error or warning messages.
--
{-# INLINEABLE verify #-}
verify :: Program -> [CompileOption] -> IO (Status, ByteString)
verify !prg !opts = do
  status   <- withCompileOptions opts (nvvmVerifyProgram prg)
  messages <- retrieve (nvvmGetProgramLogSize prg) (nvvmGetProgramLog prg)
  return (status, messages)
  where
    {#
      fun unsafe nvvmVerifyProgram
        { useProgram `Program'
        , cIntConv   `Int'
        , id         `Ptr CString'
        }
        -> `Status' cToEnum
    #}


{# fun unsafe nvvmGetProgramLogSize
    { useProgram `Program'
    , alloca-    `Int'     peekIntConv*
    }
    -> `()' checkStatus*-
#}

{# fun unsafe nvvmGetProgramLog
    { useProgram       `Program'
    , withForeignPtr'* `ForeignPtr Word8'
    }
    -> `()' checkStatus*-
#}


-- Utilities
-- ---------

{-# INLINEABLE withForeignPtr' #-}
withForeignPtr' :: ForeignPtr Word8 -> (Ptr CChar -> IO a) -> IO a
withForeignPtr' fp f = withForeignPtr fp (f . castPtr)


{-# INLINEABLE withCompileOptions #-}
withCompileOptions :: [CompileOption] -> (Int -> Ptr CString -> IO a) -> IO a
withCompileOptions opts next =
  withMany withCString (map toStr opts) $ \cs -> withArrayLen cs next
  where
    toStr :: CompileOption -> String
    toStr (OptimisationLevel n)  = printf "-opt=%d" n
    toStr (Target (Compute n m)) = printf "-arch=compute_%d%d" n m
    toStr FlushToZero            = "-ftz=1"
    toStr NoFMA                  = "-fma=0"
    toStr FastSqrt               = "-prec-sqrt=0"
    toStr FastDiv                = "-prec-div=0"
    toStr GenerateDebugInfo      = "-g"

{-# INLINEABLE retrieve #-}
retrieve :: IO Int -> (ForeignPtr Word8 -> IO ()) -> IO ByteString
retrieve size fill = do
  bytes <- size
  if bytes <= 1             -- size includes NULL terminator
    then return B.empty
    else do fp <- mallocForeignPtrBytes bytes
            _  <- fill fp
            return (B.fromForeignPtr fp 0 bytes)

{-# INLINEABLE peekProgram #-}
peekProgram :: Ptr {# type nvvmProgram #} -> IO Program
peekProgram p = Program `fmap` peek p

{-# INLINEABLE withProgram #-}
withProgram :: Program -> (Ptr {# type nvvmProgram #} -> IO a) -> IO a
withProgram p = with (useProgram p)


-- [Short]ByteStrings are not null-terminated, so can't be passed directly to C.
--
-- unsafeUseAsCString :: ShortByteString -> CString
-- unsafeUseAsCString (BI.SBS ba#) = Ptr (byteArrayContents# ba#)

{-# INLINE useAsCString #-}
useAsCString :: ShortByteString -> (CString -> IO a) -> IO a
useAsCString (BSI.SBS ba#) action = IO $ \s0 ->
  case sizeofByteArray# ba#                    of { n# ->
  case newPinnedByteArray# (n# +# 1#) s0       of { (# s1, mba# #) ->
  case byteArrayContents# (unsafeCoerce# mba#) of { addr# ->
  case copyByteArrayToAddr# ba# 0# addr# n# s1 of { s2 ->
  case writeWord8OffAddr# addr# n# 0## s2      of { s3 ->
  case action (Ptr addr#)                      of { IO action' ->
  case action' s3                              of { (# s4, r  #) ->
  case touch# mba# s4                          of { s5 ->
  (# s5, r #)
 }}}}}}}}