From: Luke Lau <luke_lau@icloud.com>
Date: Sun, 2 Jun 2019 18:25:26 +0000 (+0100)
Subject: Power up the OrcJIT
X-Git-Url: https://git.lukelau.me/?p=kaleidoscope-hs.git;a=commitdiff_plain;h=2fd7e4f3686c3b977881f995ec12b8dd60e72254

Power up the OrcJIT

We're going to be derivating from the Kaleidoscope tutorial, by using
the newer OrcJIT framework rather than MCJIT.
OrcJIT allows for lazy, On Request Compliation, and is built up of
different layers for linking, compiling, etc.

We start it off by creating a new execution session which will hold all
of our layers, followed by a symbol resolver which we won't need (yet).
We then create the linking and compilation layers, and then a ModuleKey:
something that llvm-hs uses to keep track of modules across each
invocation of withModule.

There's a lot of things to keep track of so I've added a small data type
to hold them all.
---

diff --git a/Main.hs b/Main.hs
index 4dafa02..5bc3a3e 100644
--- a/Main.hs
+++ b/Main.hs
@@ -18,6 +18,8 @@ import LLVM.AST.Type as Type
 import LLVM.Context
 import LLVM.IRBuilder
 import LLVM.Module
+import LLVM.OrcJIT
+import LLVM.OrcJIT.CompileLayer
 import LLVM.PassManager
 import LLVM.Pretty
 import LLVM.Target
@@ -25,13 +27,29 @@ import System.IO
 import System.IO.Error
 import Text.Read (readMaybe)
 
+data JITEnv = JITEnv
+  { jitEnvContext :: Context
+  , jitEnvCompileLayer :: IRCompileLayer ObjectLinkingLayer
+  , jitEnvModuleKey :: ModuleKey
+  }
+
 main :: IO ()
 main = do
   withContext $ \ctx -> withHostTargetMachine $ \tm -> do
-    ast <- runReaderT (buildModuleT "main" repl) ctx
+    withExecutionSession $ \exSession ->
+      withSymbolResolver exSession (SymbolResolver symResolver) $ \symResolverPtr ->
+        withObjectLinkingLayer exSession (const $ pure symResolverPtr) $ \linkingLayer ->
+          withIRCompileLayer linkingLayer tm $ \compLayer -> do
+            withModuleKey exSession $ \mdlKey -> do
+              let env = JITEnv ctx compLayer mdlKey
+              ast <- runReaderT (buildModuleT "main" repl) env
               return ()
 
-repl :: ModuleBuilderT (ReaderT Context IO) ()
+-- This can eventually be used to resolve external functions, e.g. a stdlib call
+symResolver :: MangledSymbol -> IO (Either JITSymbolError JITSymbol)
+symResolver sym = undefined
+
+repl :: ModuleBuilderT (ReaderT JITEnv IO) ()
 repl = do
   liftIO $ hPutStr stderr "ready> "
   mline <- liftIO $ catchIOError (Just <$> getLine) eofHandler
@@ -45,14 +63,14 @@ repl = do
           def <- mostRecentDef
           
           ast <- moduleSoFar "main"
-          ctx <- lift ask
+          ctx <- lift $ asks jitEnvContext
+          env <- lift ask
           liftIO $ withModuleFromAST ctx ast $ \mdl -> do
             Text.hPutStrLn stderr $ ppll def
             let spec = defaultCuratedPassSetSpec { optLevel = Just 3 }
             -- this returns true if the module was modified
             withPassManager spec $ flip runPassManager mdl
-            Text.hPutStrLn stderr . ("\n" <>) . ppllvm =<< moduleAST mdl
-            when anon (jit mdl >>= hPrint stderr)
+            when anon (jit env mdl >>= hPrint stderr)
 
           when anon (removeDef def)
       repl
@@ -63,8 +81,10 @@ repl = do
     isAnonExpr (ConstantOperand (GlobalReference _ "__anon_expr")) = True
     isAnonExpr _ = False
 
-jit :: Module -> IO Double
-jit _mdl = putStrLn "Working on it!" >> return 0
+jit :: JITEnv -> Module -> IO Double
+jit JITEnv{jitEnvCompileLayer=compLayer, jitEnvModuleKey=mdlKey} mdl =
+  withModule compLayer mdlKey mdl $ do
+    return 0
 
 type Binds = Map.Map String Operand