1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
|
{-# LANGUAGE LambdaCase #-}
{-# LANGUAGE TupleSections #-}
module Compile.Exec (
KernelLib,
buildKernel,
callKernelFun,
) where
import Control.Monad (when)
import Data.IORef
import qualified Data.Map.Strict as Map
import Data.Map.Strict (Map)
import Foreign (Ptr)
import Foreign.Ptr (FunPtr)
import System.Directory (removeDirectoryRecursive)
import System.Environment (lookupEnv)
import System.IO (hPutStrLn, stderr)
import System.IO.Unsafe (unsafePerformIO)
import System.Posix.DynamicLinker
import System.Posix.Temp (mkdtemp)
import System.Process (readProcess)
debug :: Bool
debug = False
-- The IORef wrapper is required for the finalizer to attach properly (see the 'Weak' docs)
data KernelLib = KernelLib !(IORef (Map String (FunPtr (Ptr () -> IO ()))))
buildKernel :: String -> [String] -> IO KernelLib
buildKernel csource funnames = do
template <- (++ "/tmp.chad.") <$> getTempDir
path <- mkdtemp template
let outso = path ++ "/out.so"
let args = ["-O3", "-march=native"
,"-shared", "-fPIC"
,"-std=c99", "-x", "c"
,"-o", outso, "-"
,"-Wall", "-Wextra", "-Wno-unused-variable", "-Wno-unused-parameter"]
_ <- readProcess "gcc" args csource
numLoaded <- atomicModifyIORef' numLoadedCounter (\n -> (n+1, n+1))
when debug $ hPutStrLn stderr $ "[chad] loading kernel " ++ path ++ " (" ++ show numLoaded ++ " total)"
dl <- dlopen outso [RTLD_LAZY, RTLD_LOCAL]
removeDirectoryRecursive path -- we keep a reference anyway because we have the file open now
ptrs <- Map.fromList <$> sequence [(name,) <$> dlsym dl name | name <- funnames]
ref <- newIORef ptrs
_ <- mkWeakIORef ref (do numLeft <- atomicModifyIORef' numLoadedCounter (\n -> (n-1, n-1))
when debug $ hPutStrLn stderr $ "[chad] unloading kernel " ++ path ++ " (" ++ show numLeft ++ " left)"
dlclose dl)
return (KernelLib ref)
foreign import ccall "dynamic"
wrapKernelFun :: FunPtr (Ptr () -> IO ()) -> Ptr () -> IO ()
-- Ensure that keeping a reference to the returned function also keeps the 'KernelLib' alive
{-# NOINLINE callKernelFun #-}
callKernelFun :: String -> KernelLib -> Ptr () -> IO ()
callKernelFun key (KernelLib ref) arg = do
mp <- readIORef ref
wrapKernelFun (mp Map.! key) arg
getTempDir :: IO FilePath
getTempDir =
lookupEnv "TMPDIR" >>= \case
Just s | not (null s) -> return s
_ -> return "/tmp"
{-# NOINLINE numLoadedCounter #-}
numLoadedCounter :: IORef Int
numLoadedCounter = unsafePerformIO $ newIORef 0
|
