Implement calling of Futhark Server programs from the interpreter

futhark.cabal

@@ -397,20 +397,31 @@ library
       Futhark.Transform.Rename
       Futhark.Transform.Substitute
       Futhark.Util
+      Futhark.Util.BiMap
       Futhark.Util.CMath
       Futhark.Util.IntegralExp
       Futhark.Util.Html
       Futhark.Util.Loc
       Futhark.Util.Log
+      Futhark.Util.NDArray
       Futhark.Util.Options
       Futhark.Util.Pretty
       Futhark.Util.ProgressBar
       Futhark.Util.Table
+      Futhark.Util.UID
       Futhark.Version
       Language.Futhark
       Language.Futhark.Core
       Language.Futhark.Interpreter
       Language.Futhark.Interpreter.AD
+      Language.Futhark.Interpreter.FFI
+      Language.Futhark.Interpreter.FFI.Server
+      Language.Futhark.Interpreter.FFI.Server.Explorer
+      Language.Futhark.Interpreter.FFI.Server.Interface
+      Language.Futhark.Interpreter.FFI.Server.Packer
+      Language.Futhark.Interpreter.FFI.Server.TypeLayout
+      Language.Futhark.Interpreter.FFI.UIDs
+      Language.Futhark.Interpreter.FFI.Values
       Language.Futhark.Interpreter.Values
       Language.Futhark.FreeVars
       Language.Futhark.Parser

src/Futhark/Eval.hs

@@ -9,14 +9,16 @@ module Futhark.Eval
   )
 where
 
+import Control.Arrow (Arrow(second))
 import Control.Exception (IOException, catch)
 import Control.Monad (foldM, when, (<=<))
 import Control.Monad.Except (ExceptT, runExceptT, throwError)
 import Control.Monad.Free.Church (F, runF)
 import Control.Monad.IO.Class (MonadIO, liftIO)
 import Control.Monad.Trans.Class (lift)
 import Control.Monad.Trans.Reader (ReaderT (runReaderT), ask)
-import Data.IORef (IORef, modifyIORef')
+import Data.Bifunctor (first)
+import Data.IORef (IORef, modifyIORef', readIORef, writeIORef, newIORef)
 import Data.Map qualified as M
 import Data.Maybe (maybeToList)
 import Data.Sequence (Seq, (|>))
@@ -27,12 +29,17 @@ import Futhark.Error (externalErrorS, prettyCompilerError)
 import Futhark.FreshNames (VNameSource)
 import Futhark.Util.Pretty (commasep, hPutDoc, hPutDocLn, hardline, putDocLn)
 import Language.Futhark.Interpreter qualified as I
+import Language.Futhark.Interpreter.FFI qualified as S
+import Language.Futhark.Interpreter.FFI.Server (FutharkServer)
+import Language.Futhark.Interpreter.FFI.Server qualified as S
+import Language.Futhark.Interpreter.FFI.Server.Packer qualified as SP
+import Language.Futhark.Interpreter.FFI.Values (Location (Location), Value (Atom))
 import Language.Futhark.Parser (parseExp)
 import Language.Futhark.Parser.Monad (SyntaxError (SyntaxError))
 import Language.Futhark.Pretty (toName)
 import Language.Futhark.Prop (typeOf)
 import Language.Futhark.Semantic qualified as T
-import Language.Futhark.Syntax (nameToText, typeParamName)
+import Language.Futhark.Syntax (nameToText, typeParamName, VName (VName), ProgBase (progDecs), DecBase (ValDec), ValBindBase (ValBind), nameToString, nameFromString)
 import Language.Futhark.TypeChecker qualified as T
 import Prettyprinter (Doc, align, pretty, unAnnotate, vcat, (<+>))
 import Prettyprinter.Render.Terminal (AnsiStyle)
@@ -78,17 +85,36 @@ runEvalRecordRef ::
 runEvalRecordRef msgRef (EvalRecordRef action) =
   flip runReaderT msgRef $ runExceptT action
 
-newtype InterpreterState = InterpreterState (VNameSource, T.Env, I.Ctx)
+newtype InterpreterState = InterpreterState (VNameSource, T.Env, I.Ctx, Maybe FutharkServer)
+
+-- TODO: Should NOT be IORef. This is temporary, for testing
+call :: IORef (Maybe FutharkServer) -> VName -> [I.Value] -> IO I.Value
+call s (VName n _) p = do
+  let p' = map S.fromInterpreterValue p
+  (Just s') <- readIORef s
+  (r, s'') <- first S.toInterpreterValue <$> S.runFutharkServerM (SP.call (nameToText n) p') s'
+  writeIORef s $ Just s''
+  pure r
+
+realize' l = SP.realize' l
+
+-- TODO: Should NOT be IORef. This is temporary, for testing
+realize :: IORef (Maybe FutharkServer) -> Location -> IO I.Value
+realize s l = do
+  (Just s') <- readIORef s
+  (r, s'') <- first S.toInterpreterValue <$> S.runFutharkServerM (realize' l) s'
+  writeIORef s $ Just s''
+  pure r
 
 -- | Run an expression in the given interpreter state. The expression is parsed,
 -- type checked, and then run. Returns a prettyprinted result. Must be run in a
 -- monad that supports aborting and traces.
 runExpr ::
-  (Evaluation m) =>
+  (Evaluation m, MonadIO m) =>
   InterpreterState ->
   T.Text ->
   m (Doc AnsiStyle)
-runExpr (InterpreterState (src, env, ctx)) str = do
+runExpr (InterpreterState (src, env, ctx, s)) str = do
   uexp <- case parseExp "" str of
     Left (SyntaxError _ serr) -> abort $ pretty serr
     Right e -> pure e
@@ -103,7 +129,8 @@ runExpr (InterpreterState (src, env, ctx)) str = do
             "The following types are ambiguous: "
               <> commasep (map (pretty . nameToText . toName . typeParamName) tparams)
           ]
-  pval <- runInterpreterNoBreak $ I.interpretExp ctx fexp
+  is <- liftIO $ newIORef s
+  pval <- runInterpreterNoBreak call realize is $ I.interpretExp ctx fexp
   case pval of
     Left err -> do
       abort $ I.prettyInterpreterError err
@@ -136,34 +163,59 @@ newFutharkiState cfg maybe_file vfs = runExceptT $ do
       hPutDoc stderr $
         prettyWarnings ws
 
+  let modifyLast _ []     = []
+      modifyLast f [x]    = [f x]
+      modifyLast f (x:xs) = x : modifyLast f xs
+
+  (imports', s) <- case maybe_file of
+        Just file -> liftIO $
+          -- TODO: This relies quite heavily on a bunch of imports, but hey, it's a start
+          let mdec (ValDec (ValBind (Just a) (VName vn vi) c d e f g h i j)) = ValDec (ValBind (Just a) (VName (nameFromString $ "$" ++ nameToString vn) vi) c d e f g h i j)
+              mdec o = o
+              (_, m) = last imports
+              m' = m { fileProg = (fileProg m) { progDecs = map mdec $ progDecs $ fileProg m} }
+          in (modifyLast (second $ const m') imports,) . Just <$> S.startServer (take (length file - 4) file)
+        Nothing -> pure (imports, Nothing)
+
+  is <- liftIO $ newIORef s
   ictx <-
     let foldFile ctx =
           badOnLeft I.prettyInterpreterError
-            <=< runInterpreterNoBreak
+            <=< runInterpreterNoBreak call realize is
               . I.interpretImport ctx
      in foldM foldFile I.initialCtx $
-          map (fmap fileProg) imports
+          map (fmap fileProg) imports'
+  s' <- liftIO $ readIORef is
 
   let (tenv, ienv) =
-        let (iname, fm) = last imports
+        let (iname, fm) = last imports'
          in ( fileScope fm,
               ictx {I.ctxEnv = I.ctxImports ictx M.! iname}
             )
 
-  pure $ InterpreterState (src, tenv, ienv)
+  pure $ InterpreterState (src, tenv, ienv, s')
   where
     badOnLeft :: (Monad m) => (err -> err') -> Either err a -> ExceptT err' m a
     badOnLeft _ (Right x) = pure x
     badOnLeft p (Left err) = throwError $ p err
 
 runInterpreterNoBreak ::
-  (Evaluation m) =>
+  (Evaluation m, MonadIO m) =>
+  (IORef (Maybe FutharkServer) -> VName -> [I.Value] -> IO I.Value) ->
+  (IORef (Maybe FutharkServer) -> Location -> IO I.Value) ->
+  IORef (Maybe FutharkServer) ->
   F I.ExtOp a ->
   m (Either I.InterpreterError a)
-runInterpreterNoBreak m = runF m (pure . Right) intOp
+runInterpreterNoBreak call' realize' s m = runF m (pure . Right) intOp
   where
     intOp (I.ExtOpError err) = pure $ Left err
     intOp (I.ExtOpTrace w v c) = do
       trace $ pretty w <> ":" <+> align (unAnnotate v)
       c
     intOp (I.ExtOpBreak _ _ _ c) = c
+    intOp (I.ExtOpCall n p c) = do
+      r <- liftIO $ call' s n p
+      c r
+    intOp (I.ExtOpRealize l c) = do
+      r <- liftIO $ realize' s l
+      c r

src/Futhark/Util/BiMap.hs

@@ -0,0 +1,27 @@
+module Futhark.Util.BiMap
+  ( BiMap,
+    insert,
+    lookupRight,
+    lookupLeft
+  )
+where
+
+import Data.Map qualified as M
+
+data BiMap a b = BiMap (M.Map a b) (M.Map b a)
+  deriving (Eq, Ord, Show)
+
+instance (Ord a, Ord b) => Monoid (BiMap a b) where
+  mempty = BiMap mempty mempty
+
+instance (Ord a, Ord b) => Semigroup (BiMap a b) where
+  BiMap r1 l1 <> BiMap r2 l2 = BiMap (r1 <> r2) (l1 <> l2)
+
+insert :: (Ord a, Ord b) => a -> b -> BiMap a b -> BiMap a b
+insert l r (BiMap mr ml) = BiMap (M.insert l r mr) (M.insert r l ml)
+
+lookupRight :: Ord a => a -> BiMap a b -> Maybe b
+lookupRight l (BiMap mr _) = M.lookup l mr
+
+lookupLeft :: Ord b => b -> BiMap a b -> Maybe a
+lookupLeft r (BiMap _ ml) = M.lookup r ml

src/Futhark/Util/NDArray.hs

@@ -0,0 +1,59 @@
+module Futhark.Util.NDArray
+  ( NDArray,
+    fromList,
+    shape,
+    size,
+    rank,
+    (!),
+    elems,
+    flatten,
+    mapWithIndex,
+    mapMWithIndex,
+    mapMWithIndex_
+  )
+where
+
+import Data.Array qualified as A
+import Control.Monad (zipWithM, zipWithM_)
+
+data NDArray a = NDArray [Int] (A.Array Int a)
+  deriving (Eq, Ord, Show, Functor, Foldable, Traversable)
+
+fromList :: [Int] -> [a] -> NDArray a
+fromList s l = NDArray s $ A.array (0, length l - 1) (zip [0..] l)
+
+shape :: NDArray a -> [Int]
+shape (NDArray s _) = s
+
+size :: NDArray a -> Int
+size = foldl (*) 1 . shape
+
+rank :: NDArray a -> Int
+rank = length . shape
+
+(!) :: NDArray a -> [Int] -> a
+(!) (NDArray s a) idx =
+  let i = foldl (+) 0 $ zipWith (*) (reverse idx) $ scanl (*) 1 s
+   in a A.! i
+
+elems :: NDArray a -> [a]
+elems (NDArray _ a) = A.elems a
+
+flatten :: NDArray a -> A.Array Int a
+flatten (NDArray _ a) = a
+
+indexOf :: [Int] -> Int -> [Int]
+indexOf (d : ds) i = (i `mod` d) : indexOf ds (i `div` d)
+indexOf [] _ = []
+
+mapWithIndex :: ([Int] -> a -> b) -> NDArray a -> NDArray b
+mapWithIndex f nd =
+  fromList (shape nd) $ zipWith f (map (indexOf $ shape nd) [0..size nd]) $ elems nd
+
+mapMWithIndex :: Monad m => ([Int] -> a -> m b) -> NDArray a -> m (NDArray b)
+mapMWithIndex f nd =
+  fromList (shape nd) <$> zipWithM f (map (indexOf $ shape nd) [0..size nd]) (elems nd)
+
+mapMWithIndex_ :: Monad m => ([Int] -> a -> m b) -> NDArray a -> m ()
+mapMWithIndex_ f nd =
+  zipWithM_ f (map (indexOf $ shape nd) [0..size nd]) (elems nd)

src/Futhark/Util/UID.hs

@@ -0,0 +1,66 @@
+{-# LANGUAGE UndecidableInstances #-}
+module Futhark.Util.UID
+  ( -- Unique IDs
+    UID (uid),
+
+    -- Unique ID source
+    UIDSource,
+    nextUID,
+
+    -- Unique ID source monad transformer
+    UIDSourceT,
+    runUIDSourceT,
+    UIDSourceM,
+    runUIDSourceM,
+    getUID,
+    getUIDs,
+  )
+where
+
+import Control.Monad.Trans.Class
+import Control.Monad.State (StateT (runStateT), MonadState (state, get, put), MonadIO)
+import Data.Functor.Identity (Identity (runIdentity))
+import Control.Monad.RWS (MonadReader (ask, local))
+
+-- External IDs
+newtype UID p r = UID { uid :: r }
+  deriving (Show, Eq, Ord, Functor)
+
+-- External ID source
+newtype UIDSource r = UIDSource r
+
+instance Ord r => Semigroup (UIDSource r) where
+  UIDSource i1 <> UIDSource i2 = UIDSource $ max i1 i2
+
+instance (Ord r, Bounded r) => Monoid (UIDSource r) where
+  mempty = UIDSource minBound
+
+nextUID :: Enum r => UIDSource r -> (UID p r, UIDSource r)
+nextUID (UIDSource i) = (UID i, UIDSource $ succ i)
+
+-- External ID source monad transformer
+newtype UIDSourceT r m a = UIDSourceT (StateT (UIDSource r) m a)
+  deriving (Functor, Applicative, Monad, MonadTrans, MonadIO)
+
+type UIDSourceM r = UIDSourceT r Identity
+
+instance MonadState s m => MonadState s (UIDSourceT r m) where
+  get = lift get
+  put = lift . put
+  state = lift . state
+
+instance MonadReader r m => MonadReader r (UIDSourceT r m) where
+  ask = lift ask
+  local f (UIDSourceT m) = UIDSourceT (local f m)
+
+runUIDSourceT :: UIDSourceT r m a -> UIDSource r -> m (a, UIDSource r)
+runUIDSourceT (UIDSourceT m) s = runStateT m s
+
+runUIDSourceM :: UIDSourceT r Identity a -> UIDSource r -> (a, UIDSource r)
+runUIDSourceM m = runIdentity . runUIDSourceT m
+
+getUID :: (Monad m, Enum r) => UIDSourceT r m (UID p r)
+getUID = UIDSourceT $ state nextUID
+
+getUIDs :: (Monad m, Bounded r, Enum r) => r -> UIDSourceT r m [UID p r]
+getUIDs n = mapM (const getUID) [minBound..n]