Run LLVM JIT before all BTree tests

eclair-lang.cabal

@@ -1,10 +1,10 @@
 cabal-version: 2.2
 
--- This file has been generated from package.yaml by hpack version 0.33.0.
+-- This file has been generated from package.yaml by hpack version 0.34.4.
 --
 -- see: https://github.com/sol/hpack
 --
--- hash: a25d696986ffb9e46bdca3e3dcfa762a74976733f7c26bcf2058262f5fb2d6c9
+-- hash: b91e2e19bff2257859664e89b83efba2264daa4210b995d4bee5d8d3b806ab11
 
 name:           eclair-lang
 version:        0.1.0.0
@@ -45,12 +45,32 @@ library
       Paths_eclair_lang
   hs-source-dirs:
       lib
-  default-extensions: NoImplicitPrelude OverloadedStrings LambdaCase ViewPatterns RankNTypes TypeFamilies DataKinds KindSignatures TupleSections DeriveFunctor DeriveFoldable DeriveTraversable DeriveGeneric DeriveAnyClass DerivingStrategies DerivingVia FlexibleContexts FlexibleInstances ScopedTypeVariables
+  default-extensions:
+      NoImplicitPrelude
+      OverloadedStrings
+      LambdaCase
+      ViewPatterns
+      RankNTypes
+      TypeFamilies
+      DataKinds
+      KindSignatures
+      TupleSections
+      DeriveFunctor
+      DeriveFoldable
+      DeriveTraversable
+      DeriveGeneric
+      DeriveAnyClass
+      DerivingStrategies
+      DerivingVia
+      FlexibleContexts
+      FlexibleInstances
+      ScopedTypeVariables
   ghc-options: -fhide-source-paths -fno-show-valid-hole-fits -fno-sort-valid-hole-fits
   cxx-options: -std=c++20 -D__EMBEDDED_SOUFFLE__ -Wall
   build-depends:
       algebraic-graphs ==0.*
     , base >=4.7 && <5
+    , bytestring ==0.*
     , containers ==0.*
     , exceptions ==0.10.*
     , extra ==1.*
@@ -87,12 +107,32 @@ executable eclairc
       Paths_eclair_lang
   hs-source-dirs:
       src
-  default-extensions: NoImplicitPrelude OverloadedStrings LambdaCase ViewPatterns RankNTypes TypeFamilies DataKinds KindSignatures TupleSections DeriveFunctor DeriveFoldable DeriveTraversable DeriveGeneric DeriveAnyClass DerivingStrategies DerivingVia FlexibleContexts FlexibleInstances ScopedTypeVariables
+  default-extensions:
+      NoImplicitPrelude
+      OverloadedStrings
+      LambdaCase
+      ViewPatterns
+      RankNTypes
+      TypeFamilies
+      DataKinds
+      KindSignatures
+      TupleSections
+      DeriveFunctor
+      DeriveFoldable
+      DeriveTraversable
+      DeriveGeneric
+      DeriveAnyClass
+      DerivingStrategies
+      DerivingVia
+      FlexibleContexts
+      FlexibleInstances
+      ScopedTypeVariables
   ghc-options: -fhide-source-paths -fno-show-valid-hole-fits -fno-sort-valid-hole-fits -threaded -rtsopts -with-rtsopts=-N
   cxx-options: -std=c++20 -D__EMBEDDED_SOUFFLE__
   build-depends:
       algebraic-graphs ==0.*
     , base >=4.7 && <5
+    , bytestring ==0.*
     , containers ==0.*
     , directory ==1.*
     , eclair-lang
@@ -134,12 +174,32 @@ test-suite eclair-test
       Paths_eclair_lang
   hs-source-dirs:
       tests
-  default-extensions: NoImplicitPrelude OverloadedStrings LambdaCase ViewPatterns RankNTypes TypeFamilies DataKinds KindSignatures TupleSections DeriveFunctor DeriveFoldable DeriveTraversable DeriveGeneric DeriveAnyClass DerivingStrategies DerivingVia FlexibleContexts FlexibleInstances ScopedTypeVariables
+  default-extensions:
+      NoImplicitPrelude
+      OverloadedStrings
+      LambdaCase
+      ViewPatterns
+      RankNTypes
+      TypeFamilies
+      DataKinds
+      KindSignatures
+      TupleSections
+      DeriveFunctor
+      DeriveFoldable
+      DeriveTraversable
+      DeriveGeneric
+      DeriveAnyClass
+      DerivingStrategies
+      DerivingVia
+      FlexibleContexts
+      FlexibleInstances
+      ScopedTypeVariables
   ghc-options: -fhide-source-paths -fno-show-valid-hole-fits -fno-sort-valid-hole-fits
   cxx-options: -std=c++20 -D__EMBEDDED_SOUFFLE__
   build-depends:
       algebraic-graphs ==0.*
     , base >=4.7 && <5
+    , bytestring ==0.*
     , containers ==0.*
     , eclair-lang
     , exceptions ==0.10.*

lib/Eclair/Lowering/RA.hs

@@ -98,7 +98,8 @@ generateFnsForRelations indexMap typeInfo = do
     -- TODO: avoid codegen collisions between relations
     for (toList idxs) $ \idx -> do
       let meta = mkMeta idx (fromJust $ Map.lookup r typeInfo)
-      (idx,) <$> BTree.codegen meta
+      (fns, _sizes) <- BTree.codegen meta
+      pure (idx, fns)
 
   pure $ map Map.fromList results
 

lib/Eclair/Runtime/BTree.hs

@@ -2,6 +2,7 @@
 
 module Eclair.Runtime.BTree
   ( Meta(..)
+  , Sizes(..)
   , SearchIndex
   , SearchType(..)
   , codegen
@@ -69,6 +70,8 @@ data Externals
 data Sizes
   = Sizes
   { pointerSize :: Word64
+  , treeSize :: Word64
+  , iterSize :: Word64
   , valueSize :: Word64
   , nodeDataSize :: Word64
   , leafNodeSize :: Word64
@@ -89,13 +92,14 @@ type IRCodegen = IRBuilderT ModuleCodegen
 type ModuleCodegen = ReaderT CGState ModuleBuilder
 
 
-codegen :: Meta -> ModuleBuilderT IO Functions
+codegen :: Meta -> ModuleBuilderT IO (Functions, Sizes)
 codegen meta = do
   sizes <- computeSizes meta
   hoist intoIO $ do
     tys <- runReaderT (generateTypes sizes) meta
     exts <- mkExternals
-    runReaderT generateFunctions $ CGState meta tys sizes exts
+    fns <- runReaderT generateFunctions $ CGState meta tys sizes exts
+    pure (fns, sizes)
   where intoIO = pure . runIdentity
 
 mkExternals :: ModuleBuilder Externals
@@ -125,7 +129,10 @@ computeSizes meta = do
       innerNodeTy = wrap [nodeTy, ArrayType (numKeys' + 1) (ptr nodeTy)]
   leafNodeSize <- sizeOfType ("leaf_node_t", nodeTy)
   innerNodeSize <- sizeOfType ("inner_node_t", innerNodeTy)
-  pure $ Sizes ptrSize valueSize nodeDataSize leafNodeSize innerNodeSize
+  let positionTy = i16
+  iterSize <- sizeOfType ("btree_iterator_t", wrap [ptr nodeTy, positionTy])
+  treeSize <- sizeOfType ("btree_t", wrap [ptr nodeTy, ptr nodeTy])
+  pure $ Sizes ptrSize treeSize iterSize valueSize nodeDataSize leafNodeSize innerNodeSize
   where
     wrap = StructureType False
 

lib/Eclair/Runtime/LLVM.hs

@@ -4,7 +4,7 @@ module Eclair.Runtime.LLVM
   ( module Eclair.Runtime.LLVM
   ) where
 
-import Protolude hiding ( Type, (.), bit )
+import Protolude hiding ( Type, (.), bit, moduleName )
 import Control.Category
 import Control.Monad.Morph
 import Control.Monad.Fix
@@ -17,6 +17,7 @@ import LLVM.IRBuilder.Instruction
 import qualified LLVM.AST.IntegerPredicate as IP
 import qualified LLVM.AST.Constant as Constant
 import LLVM.AST.Operand ( Operand(..) )
+import LLVM.AST (Module, moduleName, moduleDefinitions, defaultModule)
 import LLVM.AST.Type
 import LLVM.AST.Name
 import Eclair.Runtime.Hash
@@ -30,6 +31,7 @@ import qualified LLVM.CodeGenOpt as CG
 import qualified LLVM.CodeModel as CM
 import qualified LLVM.Relocation as Rel
 import LLVM.Target
+import Data.ByteString.Short
 
 
 -- TODO: remove, import directly from llvm-hs
@@ -72,6 +74,12 @@ sizeOfType (n, ty) = do
       defineType n n' t'
       setNamedType t' ty
 
+codegenModule :: ShortByteString -> ModuleBuilderT IO a -> IO (a, Module)
+codegenModule name mod =
+  map mkModule <$> runModuleBuilderT emptyModuleBuilder mod
+  where
+    mkModule defs = defaultModule { moduleName = name, moduleDefinitions = defs }
+
 -- NOTE: Orphan instance, but should give no conflicts.
 instance MFunctor ModuleBuilderT where
   hoist nat = ModuleBuilderT . hoist nat . unModuleBuilderT

package.yaml

@@ -22,6 +22,7 @@ dependencies:
   - protolude == 0.3.0
   - extra == 1.*
   - text == 1.*
+  - bytestring == 0.*
   - vector == 0.12.*
   - containers == 0.*
   - transformers == 0.*

tests/Test/Eclair/Runtime/BTreeSpec.hs

@@ -4,11 +4,169 @@ module Test.Eclair.Runtime.BTreeSpec
 
 import Protolude
 import Test.Hspec
+import Control.Monad.Cont
+import Data.ByteString.Short hiding (index)
+import Eclair.Runtime.BTree
+import Eclair.Runtime.Store
+import Eclair.Runtime.LLVM
+import LLVM.IRBuilder.Module
+import LLVM.Context
+import LLVM.Module
+import LLVM.Target
+import LLVM.OrcJIT
+import qualified LLVM.Relocation as Relocation
+import qualified LLVM.CodeModel as CodeModel
+import qualified LLVM.CodeGenOpt as CodeGenOpt
+import qualified LLVM.Internal.OrcJIT.CompileLayer as CL
+import Foreign
+
+
+jit :: ModuleBuilderT IO a -> (forall l. CompileLayer l => l -> a -> IO b) -> IO b
+jit code f = flip runContT pure $ do
+  ctx <- ContT withContext
+  (a, ffiCode) <- lift $ codegenModule "test.ll" code
+  mod <- ContT $ withModuleFromAST ctx ffiCode
+  tm <- ContT $ withHostTargetMachine Relocation.PIC CodeModel.JITDefault CodeGenOpt.None
+  exeSession <- ContT withExecutionSession
+  resolver <- ContT $ withSymbolResolver exeSession (SymbolResolver symbolResolver)
+  objectLayer <- ContT $ withObjectLinkingLayer exeSession (\_key -> pure resolver)
+  compileLayer <- ContT $ withIRCompileLayer objectLayer tm
+  key <- ContT $ withModuleKey exeSession
+  lift $ withModule compileLayer key mod $ f compileLayer a
+
+symbolResolver :: MangledSymbol -> IO (Either JITSymbolError JITSymbol)
+symbolResolver = panic "Not implemented!"
+
+data Tree
+data Value
+data Iter
+
+data FFI
+  = FFI
+  { ffiWithEmptyTree :: forall a. (Ptr Tree -> IO a) -> IO a
+  , ffiPurge :: Ptr Tree -> IO ()
+  , ffiSwap :: Ptr Tree -> Ptr Tree -> IO ()
+  , ffiBegin :: Ptr Tree -> IO (ForeignPtr Value)
+  , ffiEnd :: Ptr Tree -> IO (ForeignPtr Value)
+  , ffiInsert :: Ptr Tree -> Ptr Value -> IO Bool
+  , ffiInsertRange :: Ptr Tree -> Ptr Iter -> Ptr Iter -> IO ()
+  , ffiIsEmpty :: Ptr Tree -> IO Bool
+  , ffiLowerBound :: Ptr Tree -> Ptr Value -> IO (ForeignPtr Iter)
+  , ffiUpperBound :: Ptr Tree -> Ptr Value -> IO (ForeignPtr Iter)
+  , ffiContains :: Ptr Tree -> Ptr Value -> IO Bool
+  , ffiIterIsEqual :: Ptr Iter -> Ptr Iter -> IO Bool
+  , ffiIterCurrent :: Ptr Iter -> IO (Ptr Value)
+  , ffiIterNext :: Ptr Iter -> IO ()
+  }
+
+jitCompile :: IO FFI
+jitCompile = do
+  let meta = Meta { numColumns = 4
+                  , index = [1, 3]
+                  , blockSize = 256
+                  , searchType = Linear
+                  }
+  jit (codegen meta) $ \compileLayer (_, sizes) -> do
+    fnInitEmpty <- importSymbol compileLayer "btree_init_empty" foreignInitEmpty
+    fnDestroy <- importSymbol compileLayer "btree_destroy" foreignDestroy
+    fnPurge <- importSymbol compileLayer "btree_clear" foreignPurge
+    fnSwap <- importSymbol compileLayer "btree_swap" foreignSwap
+    fnBegin <- importSymbol compileLayer "btree_begin" foreignBegin
+    fnEnd <- importSymbol compileLayer "btree_end" foreignEnd
+    fnInsert <- importSymbol compileLayer "btree_insert_value" foreignInsert
+    fnInsertRange <- importSymbol compileLayer "btree_insert_range" foreignInsertRange
+    fnIsEmpty <- importSymbol compileLayer "btree_is_empty" foreignIsEmpty
+    fnLowerBound <- importSymbol compileLayer "btree_lower_bound" foreignLowerBound
+    fnUpperBound <- importSymbol compileLayer "btree_upper_bound" foreignUpperBound
+    fnContains <- importSymbol compileLayer "btree_contains" foreignContains
+    fnIterIsEqual <- importSymbol compileLayer "iterator_is_equal" foreignIterIsEqual
+    fnIterCurrent <- importSymbol compileLayer "iterator_current" foreignIterCurrent
+    fnIterNext <- importSymbol compileLayer "iterator_next" foreignIterNext
+    pure FFI { ffiWithEmptyTree = withResource (treeSize sizes) fnInitEmpty fnDestroy
+             , ffiPurge = fnPurge
+             , ffiSwap = fnSwap
+             , ffiBegin = \tree -> allocateAndApply (valueSize sizes) (fnBegin tree)
+             , ffiEnd = \tree -> allocateAndApply (valueSize sizes) (fnEnd tree)
+             , ffiInsert = fnInsert
+             , ffiInsertRange = fnInsertRange
+             , ffiIsEmpty = fnIsEmpty
+             , ffiLowerBound = \tree val -> allocateAndApply (iterSize sizes) (fnLowerBound tree val)
+             , ffiUpperBound = \tree val -> allocateAndApply (iterSize sizes) (fnUpperBound tree val)
+             , ffiContains = fnContains
+             , ffiIterIsEqual = fnIterIsEqual
+             , ffiIterCurrent = fnIterCurrent
+             , ffiIterNext = fnIterNext
+             }
+
+importSymbol :: CompileLayer l => l -> ShortByteString -> (FunPtr a -> a) -> IO a
+importSymbol compileLayer symbol fn = do
+  mangled <- mangleSymbol compileLayer symbol
+  Right (JITSymbol symbolPtr _) <- CL.findSymbol compileLayer mangled True
+  pure $ fn $ castPtrToFunPtr $ wordPtrToPtr symbolPtr
+
+foreign import ccall "dynamic" foreignInitEmpty
+  :: FunPtr (Ptr Tree -> IO ()) -> Ptr Tree -> IO ()
+foreign import ccall "dynamic" foreignDestroy
+  :: FunPtr (Ptr Tree -> IO ()) -> Ptr Tree -> IO ()
+foreign import ccall "dynamic" foreignPurge
+  :: FunPtr (Ptr Tree -> IO ()) -> Ptr Tree -> IO ()
+foreign import ccall "dynamic" foreignSwap
+  :: FunPtr (Ptr Tree -> Ptr Tree -> IO ())
+  -> Ptr Tree -> Ptr Tree -> IO ()
+foreign import ccall "dynamic" foreignBegin
+  :: FunPtr (Ptr Tree -> Ptr Value -> IO ())
+  -> Ptr Tree -> Ptr Value -> IO ()
+foreign import ccall "dynamic" foreignEnd
+  :: FunPtr (Ptr Tree -> Ptr Value -> IO ())
+  -> Ptr Tree -> Ptr Value -> IO ()
+foreign import ccall "dynamic" foreignInsert
+  :: FunPtr (Ptr Tree -> Ptr Value -> IO Bool)
+  -> Ptr Tree -> Ptr Value -> IO Bool
+foreign import ccall "dynamic" foreignInsertRange
+  :: FunPtr (Ptr Tree -> Ptr Iter -> Ptr Iter -> IO ())
+  -> Ptr Tree -> Ptr Iter -> Ptr Iter -> IO ()
+foreign import ccall "dynamic" foreignIsEmpty
+  :: FunPtr (Ptr Tree -> IO Bool)
+  -> Ptr Tree -> IO Bool
+foreign import ccall "dynamic" foreignLowerBound
+  :: FunPtr (Ptr Tree -> Ptr Value -> Ptr Iter -> IO ())
+  -> Ptr Tree -> Ptr Value -> Ptr Iter -> IO ()
+foreign import ccall "dynamic" foreignUpperBound
+  :: FunPtr (Ptr Tree -> Ptr Value -> Ptr Iter -> IO ())
+  -> Ptr Tree -> Ptr Value -> Ptr Iter -> IO ()
+foreign import ccall "dynamic" foreignContains
+  :: FunPtr (Ptr Tree -> Ptr Value -> IO Bool)
+  -> Ptr Tree -> Ptr Value -> IO Bool
+foreign import ccall "dynamic" foreignIterIsEqual
+  :: FunPtr (Ptr Iter -> Ptr Iter -> IO Bool)
+  -> Ptr Iter -> Ptr Iter -> IO Bool
+foreign import ccall "dynamic" foreignIterCurrent
+  :: FunPtr (Ptr Iter -> IO (Ptr Value))
+  -> Ptr Iter -> IO (Ptr Value)
+foreign import ccall "dynamic" foreignIterNext
+  :: FunPtr (Ptr Iter -> IO ())
+  -> Ptr Iter -> IO ()
+
+withResource :: Word64            -- number of bytes to allocate
+             -> (Ptr a -> IO ())  -- "constructor"
+             -> (Ptr a -> IO ())  -- "destructor"
+             -> (Ptr a -> IO b)
+             -> IO b
+withResource numBytes construct destruct =
+  bracket (mask_ $ do
+            ptr <- mallocBytes (fromIntegral numBytes)
+            construct ptr
+            pure ptr)
+          (\ptr -> destruct ptr *> free ptr)
+
+allocateAndApply :: Word64 -> (Ptr a -> IO b) -> IO (ForeignPtr a)
+allocateAndApply size f = do
+  ptr <- mallocForeignPtrBytes (fromIntegral size)
+  withForeignPtr ptr f
+  pure ptr
 
 {-
 TODO:
-specialize for 1 specific situation
-JIT compile using LLVM
 import foreign functions into haskell
 make ADT for all operations
 generate list of ops using hedgehog, run program with that