-- Copyright 2022 Google LLC
--
-- Use of this source code is governed by a BSD-style
-- license that can be found in the LICENSE file or at
-- https://developers.google.com/open-source/licenses/bsd

module LLVM.CUDA (
  LLVMKernel (..), compileCUDAKernel, ptxTargetTriple, ptxDataLayout) where

import qualified LLVM.AST as L
import qualified LLVM.AST.Global as L
import qualified LLVM.AST.Global as GL
import qualified LLVM.AST.AddrSpace as L
import qualified LLVM.AST.Constant as C
import qualified LLVM.AST.DataLayout as L
import qualified LLVM.AST.FunctionAttribute as FA
import qualified LLVM.Relocation as R
import qualified LLVM.CodeModel as CM
import qualified LLVM.CodeGenOpt as CGO
import qualified LLVM.Module as Mod
#if MIN_VERSION_llvm_hs(15,0,0)
import qualified LLVM.Passes as P
#else
import qualified LLVM.Transforms as P
#endif
import qualified LLVM.Target as T
import LLVM.Context
import System.IO
import System.IO.Unsafe
import System.IO.Temp
import System.Directory (listDirectory)
import System.Process
import qualified System.Environment as E
import System.Exit

import Data.ByteString.Short (ShortByteString)
import qualified Data.ByteString.Char8 as B
import qualified Data.Map as M
import qualified Data.Set as S

import LLVM.Compile
import Types.Imp
import Types.Source


data LLVMKernel = LLVMKernel L.Module

cudaPath :: IO String
cudaPath = maybe "/usr/local/cuda" id <$> E.lookupEnv "CUDA_PATH"

compileCUDAKernel :: PassLogger -> LLVMKernel -> String -> IO CUDAKernel
compileCUDAKernel logger (LLVMKernel ast) arch = undefined
  T.initializeAllTargets
  withContext \ctx ->
    Mod.withModuleFromAST ctx ast \m -> do
      withGPUTargetMachine (B.pack arch) \tm -> do
        linkLibdevice m
        standardCompilationPipeline OptAggressively logger ["kernel"] tm m
        ptx <- Mod.moduleTargetAssembly tm m
        usePTXAS <- maybe False (=="1") <$> E.lookupEnv "DEX_USE_PTXAS"
        if usePTXAS
          then do
            ptxasPath <- (++"/bin/ptxas") <$> cudaPath
            withSystemTempFile "kernel.ptx" \ptxPath ptxH -> do
              B.hPut ptxH ptx
              hClose ptxH
              withSystemTempFile "kernel.sass" \sassPath sassH -> do
                let cmd = proc ptxasPath [ptxPath, "-o", sassPath, "-arch=" ++ arch, "-O3"]
                withCreateProcess cmd \_ _ _ ptxas -> do
                  code <- waitForProcess ptxas
                  case code of
                    ExitSuccess   -> return ()
                    ExitFailure _ -> error "ptxas invocation failed"
                -- TODO: B.readFile might be faster, but withSystemTempFile seems to lock the file...
                CUDAKernel <$> B.hGetContents sassH
          else return $ CUDAKernel ptx
{-# SCC compileCUDAKernel #-}

{-# NOINLINE libdevice #-}
libdevice :: L.Module
libdevice = unsafePerformIO $ do
  withContext \ctx -> do
    libdeviceDirectory <- (flip (++) $ "/nvvm/libdevice") <$> cudaPath
    [libdeviceFileName] <- listDirectory libdeviceDirectory
    let libdevicePath = libdeviceDirectory ++ "/" ++ libdeviceFileName
    libdeviceBC <- B.readFile libdevicePath
    m <- Mod.withModuleFromBitcode ctx (libdevicePath, libdeviceBC) Mod.moduleAST
    -- Override the data layout and target triple to avoid warnings when linking
    return $ m { L.moduleDataLayout = Just ptxDataLayout
               , L.moduleTargetTriple = Just ptxTargetTriple }

linkLibdevice :: Mod.Module -> IO ()
linkLibdevice m = do
  ctx <- Mod.moduleContext m
  Mod.withModuleFromAST ctx zeroNVVMReflect \reflectm ->
    Mod.withModuleFromAST ctx libdevice \ldm -> do
      Mod.linkModules m ldm
      Mod.linkModules m reflectm
      runPasses [P.AlwaysInline True] Nothing m

-- llvm-hs does not expose the NVVM reflect pass, so we have to eliminate all calls to
-- __nvvm_reflect by ourselves. Since we aren't really interested in setting any reflection
-- flags to a value other than 0 for now, we eliminate the function by linking with this
-- trivial module and forcing the definition to get inlined.
zeroNVVMReflect :: L.Module
zeroNVVMReflect =
  L.Module { L.moduleName = "zero-nvvm-reflect"
           , L.moduleSourceFileName = ""
           , L.moduleDataLayout = Just ptxDataLayout
           , L.moduleTargetTriple = Just ptxTargetTriple
           , L.moduleDefinitions =
               [ L.GlobalDefinition $ L.functionDefaults
                   { L.name = "__nvvm_reflect"
                   , L.returnType = L.IntegerType 32
                   , L.parameters =
                       ([ L.Parameter i8p "name" [] ], False)
                   , GL.functionAttributes = [ Right FA.AlwaysInline ]
                   , L.basicBlocks = [
                       L.BasicBlock "entry" [] $ L.Do $ L.Ret (Just $ L.ConstantOperand $ C.Int 32 0) []
                     ]
                   }
               ]
           }
  where
#if MIN_VERSION_llvm_hs(15,0,0)
    i8p = L.PointerType (L.AddrSpace 0)
#else
    i8p = L.PointerType (L.IntegerType 8) (L.AddrSpace 0)
#endif

ptxTargetTriple :: ShortByteString
ptxTargetTriple = "nvptx64-nvidia-cuda"

ptxDataLayout :: L.DataLayout
ptxDataLayout = (L.defaultDataLayout L.LittleEndian)
    { L.endianness     = L.LittleEndian
    , L.pointerLayouts = M.fromList [(L.AddrSpace 0, (64, L.AlignmentInfo 64 64))]
    , L.typeLayouts    = M.fromList $
        [ ((L.IntegerAlign, 1), L.AlignmentInfo 8 8) ] ++
        [ ((L.IntegerAlign, w), L.AlignmentInfo w w) | w <- [8, 16, 32, 64] ] ++
        [ ((L.FloatAlign,   w), L.AlignmentInfo w w) | w <- [32, 64] ] ++
        [ ((L.VectorAlign,  w), L.AlignmentInfo w w) | w <- [16, 32, 64, 128] ]
    , L.nativeSizes    = Just $ S.fromList [16, 32, 64]
    }

-- === supported target machines ===

withGPUTargetMachine :: B.ByteString -> (T.TargetMachine -> IO a) -> IO a
withGPUTargetMachine computeCapability next = do
  (tripleTarget, _) <- T.lookupTarget Nothing ptxTargetTriple
  T.withTargetOptions \topt ->
    T.withTargetMachine
      tripleTarget
      ptxTargetTriple
      computeCapability
      (M.singleton (T.CPUFeature "ptx64") True)
      topt
      R.Default
      CM.Default
      CGO.Aggressive
      next