Find our JIT'ed function and run it

[kaleidoscope-hs.git] / Main.hs

diff --git a/Main.hs b/Main.hs
index c3f929214099d738794ea2b6a76abe7be0d3124f..468573d865ab4fccb350bc78bf446d4e9d400563 100644 (file)
--- a/Main.hs
+++ b/Main.hs
@@ -2,12 +2,14 @@
 
 import AST as K -- K for Kaleidoscope
 import Utils
 
 import AST as K -- K for Kaleidoscope
 import Utils

+import Control.Monad

 import Control.Monad.Trans.Class
 import Control.Monad.Trans.Reader
 import Control.Monad.IO.Class
 import Data.String
 import qualified Data.Map as Map
 import qualified Data.Text.Lazy.IO as Text
 import Control.Monad.Trans.Class
 import Control.Monad.Trans.Reader
 import Control.Monad.IO.Class
 import Data.String
 import qualified Data.Map as Map
 import qualified Data.Text.Lazy.IO as Text

+import Foreign.Ptr

 import LLVM.AST.AddrSpace
 import LLVM.AST.Constant
 import LLVM.AST.Float
 import LLVM.AST.AddrSpace
 import LLVM.AST.Constant
 import LLVM.AST.Float


@@ -17,6 +19,8 @@ import LLVM.AST.Type as Type
 import LLVM.Context
 import LLVM.IRBuilder
 import LLVM.Module
 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
 import LLVM.PassManager
 import LLVM.Pretty
 import LLVM.Target


@@ -24,13 +28,31 @@ import System.IO
 import System.IO.Error
 import Text.Read (readMaybe)
 
 import System.IO.Error
 import Text.Read (readMaybe)
 

+foreign import ccall "dynamic" mkFun :: FunPtr (IO Double) -> IO Double
+
+data JITEnv = JITEnv
+  { jitEnvContext :: Context
+  , jitEnvCompileLayer :: IRCompileLayer ObjectLinkingLayer
+  , jitEnvModuleKey :: ModuleKey
+  }
+

 main :: IO ()
 main :: IO ()

-main = do
-  withContext $ \ctx -> withHostTargetMachineDefault $ \tm -> do
-    ast <- runReaderT (buildModuleT "main" repl) ctx
+main =
+  withContext $ \ctx -> withHostTargetMachineDefault $ \tm ->
+    withExecutionSession $ \exSession ->
+      withSymbolResolver exSession (SymbolResolver symResolver) $ \symResolverPtr ->
+        withObjectLinkingLayer exSession (const $ pure symResolverPtr) $ \linkingLayer ->
+          withIRCompileLayer linkingLayer tm $ \compLayer ->
+            withModuleKey exSession $ \mdlKey -> do
+              let env = JITEnv ctx compLayer mdlKey
+              _ast <- runReaderT (buildModuleT "main" repl) env

               return ()
 
               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
 repl = do
   liftIO $ hPutStr stderr "ready> "
   mline <- liftIO $ catchIOError (Just <$> getLine) eofHandler


@@ -40,21 +62,34 @@ repl = do
       case readMaybe l of
         Nothing ->  liftIO $ hPutStrLn stderr "Couldn't parse"
         Just ast -> do
       case readMaybe l of
         Nothing ->  liftIO $ hPutStrLn stderr "Couldn't parse"
         Just ast -> do

-          hoist $ buildAST ast
-          mostRecentDef >>= liftIO . Text.hPutStrLn stderr . ppll
+          anon <- isAnonExpr <$> hoist (buildAST ast)
+          def <- mostRecentDef

           
           

-          ast <- moduleSoFar "main"
-          ctx <- lift ask
-          liftIO $ withModuleFromAST ctx ast $ \mdl -> do
+          llvmAst <- moduleSoFar "main"
+          ctx <- lift $ asks jitEnvContext
+          env <- lift ask
+          liftIO $ withModuleFromAST ctx llvmAst $ \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
             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 env mdl >>= hPrint stderr)
+
+          when anon (removeDef def)

       repl
   where
     eofHandler e
       | isEOFError e = return Nothing
       | otherwise = ioError e
       repl
   where
     eofHandler e
       | isEOFError e = return Nothing
       | otherwise = ioError e

+    isAnonExpr (ConstantOperand (GlobalReference _ "__anon_expr")) = True
+    isAnonExpr _ = False
+
+jit :: JITEnv -> Module -> IO Double
+jit JITEnv{jitEnvCompileLayer=compLayer, jitEnvModuleKey=mdlKey} mdl =
+  withModule compLayer mdlKey mdl $ do
+    mangled <- mangleSymbol compLayer "__anon_expr"
+    Right (JITSymbol fPtr _) <- findSymbolIn compLayer mdlKey mangled False
+    mkFun (castPtrToFunPtr (wordPtrToPtr fPtr))

 
 type Binds = Map.Map String Operand
 
 
 type Binds = Map.Map String Operand