From: Luke Lau <luke_lau@icloud.com>
Date: Tue, 19 Mar 2019 23:14:39 +0000 (+0000)
Subject: Add putchard and move to OrcJIT
X-Git-Url: https://git.lukelau.me/?p=kaleidoscope-hs-old.git;a=commitdiff_plain;h=c67238064a9570e5c22566413b68906c9fcf39fe

Add putchard and move to OrcJIT
---

diff --git a/.gitignore b/.gitignore
index 1377554..fe4b493 100644
--- a/.gitignore
+++ b/.gitignore
@@ -1 +1,3 @@
 *.swp
+*.o
+*.dylib
diff --git a/Main.hs b/Main.hs
index 4d82758..f096ea6 100644
--- a/Main.hs
+++ b/Main.hs
@@ -13,7 +13,8 @@ import Foreign.Ptr
 import System.Exit
 import System.IO
 import LLVM.Context
-import LLVM.ExecutionEngine
+import LLVM.OrcJIT
+import LLVM.OrcJIT.CompileLayer
 import LLVM.Module
 import LLVM.PassManager
 import LLVM.IRBuilder
@@ -25,49 +26,73 @@ import LLVM.AST.Operand
 import LLVM.AST.Type as Type
 import LLVM.AST.Typed
 import LLVM.Pretty
+import LLVM.Linking
+import LLVM.Target
+
+import Control.Concurrent.MVar
 
 type ModuleBuilderE = ModuleBuilderT (Either String)
 
-foreign import ccall "dynamic" exprFun :: FunPtr (IO Double) -> IO Double
+foreign import ccall "dynamic" mkFun :: FunPtr (IO Double) -> IO Double
 
 main :: IO ()
 main = do
   AST.Program asts <- read <$> getContents
   let eitherMdl = buildModuleT "main" $ mapM buildAST asts
+
   case eitherMdl of
     Left err -> die err
-    Right mdl -> withContext $ \ctx -> do
-      hPutStrLn stderr "Before optimisation:"
-      Text.hPutStrLn stderr (ppllvm mdl)
-      withMCJIT ctx Nothing Nothing Nothing Nothing $ \mcjit ->
-        withModuleFromAST ctx mdl $ \mdl' -> do
-          -- withPassManager defaultCuratedPassSetSpec $ \pm -> do
-          --   runPassManager pm mdl' >>= guard
-          hPutStrLn stderr "After optimisation:"
-          Text.hPutStrLn stderr . ppllvm =<< moduleAST mdl'
-          withModuleInEngine mcjit mdl' $ \emdl -> do
-          Just f <- getFunction emdl "expr"
-          let f' = castFunPtr f :: FunPtr (IO Double)
-          exprFun f' >>= print
+    Right mdl' -> withContext $ \ctx ->
+      withHostTargetMachine $ \tm -> do
+        -- hPutStrLn stderr "Before optimisation:"
+        -- Text.hPutStrLn stderr (ppllvm mdl')
+
+        withModuleFromAST ctx mdl' $ \mdl -> do
+          let spec = defaultCuratedPassSetSpec { optLevel = Just 3 }
+          withPassManager spec $ flip runPassManager mdl
+          -- hPutStrLn stderr "After optimisation:"
+          -- Text.hPutStrLn stderr . ppllvm =<< moduleAST mdl
+          jit tm mdl >>= print
+
+jit :: TargetMachine -> Module -> IO Double
+jit tm mdl = do
+  loadLibraryPermanently (Just "stdlib.dylib") >>= guard . not
+  compLayerVar <- newEmptyMVar
+  
+  -- jit time
+  withExecutionSession $ \exSession ->
+    withSymbolResolver exSession (SymbolResolver (symResolver compLayerVar)) $ \symResolverPtr ->
+      withObjectLinkingLayer exSession (const $ pure symResolverPtr) $ \linkingLayer ->
+        withModuleKey exSession $ \mdlKey ->
+          withIRCompileLayer linkingLayer tm $ \compLayer -> do
+            putMVar compLayerVar compLayer
+
+            withModule compLayer mdlKey mdl $ do
+              mangled <- mangleSymbol compLayer "expr"
+              Right (JITSymbol fPtr _) <- findSymbolIn compLayer mdlKey mangled False
+              mkFun (castPtrToFunPtr (wordPtrToPtr fPtr))
+
+  where symResolver clv sym = do
+          cl <- readMVar clv
+          ms <- findSymbol cl sym False
+          case ms of
+            Right s -> return (return s)
+            _ -> do
+              addr <- getSymbolAddressInProcess sym
+              return $ return (JITSymbol addr (JITSymbolFlags False False True True))
 
 evalProg :: AST.Program -> IO (Maybe Double)
 evalProg (AST.Program asts) = do
   let eitherMdl = buildModuleT "main" $ mapM buildAST asts
   case eitherMdl of
     Left _ -> return Nothing
-    Right mdl -> withContext $ \ctx ->
-      withMCJIT ctx Nothing Nothing Nothing Nothing $ \mcjit ->
-        withModuleFromAST ctx mdl $ \mdl' ->
-          withModuleInEngine mcjit mdl' $ \emdl -> do
-            Just f <- getFunction emdl "expr"
-            let f' = castFunPtr f :: FunPtr (IO Double)
-            Just <$> exprFun f'
+    Right mdl' -> withContext $ \ctx ->
+      withHostTargetMachine $ \tm ->
+        withModuleFromAST ctx mdl' (fmap Just . jit tm)
 
 -- | Builds up programs at the top-level of an LLVM Module
 -- >>> evalProg (read "31 - 5")
 -- Just 26.0
--- >>> evalProg (read "extern pow(x e); pow(3,2)")
--- Just 9.0
 buildAST :: AST.AST -> ModuleBuilderE Operand
 buildAST (AST.Function nameStr paramStrs body) = do
   let n = fromString nameStr
diff --git a/Makefile b/Makefile
new file mode 100644
index 0000000..f4582a9
--- /dev/null
+++ b/Makefile
@@ -0,0 +1,2 @@
+stdlib.dylib: stdlib.c
+	clang -dynamiclib $< -o $@
diff --git a/stdlib.c b/stdlib.c
new file mode 100644
index 0000000..5948f76
--- /dev/null
+++ b/stdlib.c
@@ -0,0 +1,5 @@
+#include <stdio.h>
+extern double putchard(double x) {
+	putchar((int)x);
+	return 0;
+}
diff --git a/test.ks b/test.ks
new file mode 100644
index 0000000..c18328b
--- /dev/null
+++ b/test.ks
@@ -0,0 +1,3 @@
+def foo(x y) if x > y then x + y else x - y;
+extern putchard(x);
+putchard(42);