Arith.hs

{-# LANGUAGE DeriveFoldable #-}
{-# LANGUAGE DeriveFunctor #-}
{-# LANGUAGE DeriveTraversable #-}
{-# LANGUAGE FlexibleContexts #-}
{-# LANGUAGE FlexibleInstances #-}
{-# LANGUAGE ForeignFunctionInterface #-}
{-# LANGUAGE OverloadedStrings #-}
{-# LANGUAGE TypeSynonymInstances #-}

module Main where

import Control.DeepSeq (NFData, force)
import Control.Exception
import Control.Monad
import Control.Monad.IO.Class
import Data.ByteString (ByteString)
import qualified Data.ByteString.Char8 as BS
import Data.Foldable
import Data.Functor.Foldable hiding (fold)
import Data.IORef
import Data.Map.Strict (Map)
import qualified Data.Map.Strict as Map
import Data.Monoid
import qualified Data.Set as Set
import Data.Text.Lazy (Text)
import qualified Data.Text.Lazy.IO as Text
import Foreign.Ptr
import qualified LLVM.AST as LLVM
import qualified LLVM.AST.Constant as LLVM
import qualified LLVM.AST.Float as LLVM
import qualified LLVM.AST.Type as LLVM
import qualified LLVM.CodeGenOpt as CodeGenOpt
import qualified LLVM.CodeModel as CodeModel
import qualified LLVM.Context as JIT
import qualified LLVM.IRBuilder.Instruction as LLVMIR
import qualified LLVM.IRBuilder.Module as LLVMIR
import qualified LLVM.IRBuilder.Monad as LLVMIR
import qualified LLVM.Internal.OrcJIT.CompileLayer as JIT
import qualified LLVM.Linking as JIT
import qualified LLVM.Module as JIT
import qualified LLVM.OrcJIT as JIT
import qualified LLVM.Pretty as LLVMPretty
import qualified LLVM.Relocation as Reloc
import qualified LLVM.Target as JIT

-- * Core expression type

-- | An expression will be any value of type @'Fix' 'ExprF'@, which
--   has as values arbitrarily nested applications of constructors from
--   'ExprF'. This is equivalent to just having an 'Expr type with no type
--   parameter and all @a@s replaced by 'Expr', but the 'Functor' and 'Foldable'
--   instances are quite handy, especially combined with the /recursion-schemes/
--   library.
--
--   This type allows us to express the body of a @Double -> Double@ function,
--   where 'Var' allows us to refer to the (only) argument of the function.
data ExprF a
  = -- | a 'Double' literal
    Lit Double
  | -- | @a+b@
    Add a a
  | -- | @a-b@
    Sub a a
  | -- | @a*b@
    Mul a a
  | -- | @a/b@
    Div a a
  | -- | @-a@
    Neg a
  | -- | @'exp' a@
    Exp a
  | -- | @'log' a@
    Log a
  | -- | @'sqrt' a@
    Sqrt a
  | -- | @'sin' a@
    Sin a
  | -- | @'cos' a@
    Cos a
  | -- | @'x'@
    Var
  deriving (Functor, Foldable, Traversable)

type Expr = Fix ExprF

-- * Helpers for building expressions

x :: Expr
x = Fix Var

lit :: Double -> Expr
lit d = Fix (Lit d)

add :: Expr -> Expr -> Expr
add a b = Fix (Add a b)

sub :: Expr -> Expr -> Expr
sub a b = Fix (Sub a b)

mul :: Expr -> Expr -> Expr
mul a b = Fix (Mul a b)

neg :: Expr -> Expr
neg a = Fix (Neg a)

instance Num Expr where
  fromInteger = lit . fromInteger
  (+) = add
  (-) = sub
  (*) = mul
  negate = neg
  abs = notImplemented "Expr.abs"
  signum = notImplemented "Expr.signum"

divide :: Expr -> Expr -> Expr
divide a b = Fix (Div a b)

instance Fractional Expr where
  (/) = divide
  recip = divide 1
  fromRational = lit . fromRational

instance Floating Expr where
  pi = lit pi
  exp = Fix . Exp
  log = Fix . Log
  sqrt = Fix . Sqrt
  sin = Fix . Sin
  cos = Fix . Cos
  asin = notImplemented "Expr.asin"
  acos = notImplemented "Expr.acos"
  atan = notImplemented "Expr.atan"
  sinh = notImplemented "Expr.sinh"
  cosh = notImplemented "Expr.cosh"
  asinh = notImplemented "Expr.asinh"
  acosh = notImplemented "Expr.acosh"
  atanh = notImplemented "Expr.atanh"

notImplemented :: String -> a
notImplemented = error . (++ " is not implemented")

-- * Pretty printing

-- | Pretty print an 'Expr'
pp :: Expr -> String
pp e = funprefix ++ para ppExpAlg e
  where
    funprefix = "\\x -> "

printExpr :: MonadIO m => Expr -> m ()
printExpr expr = liftIO $ do
  putStrLn "*** Expression ***\n"
  putStrLn (pp expr)

-- | Core pretty printing function. For each
--   constructor that contains sub expressions,
--   we get the string for the sub expression as
--   well as the original 'Expr' value, to help us
--   decide when to use parens.
ppExpAlg :: ExprF (Expr, String) -> String
ppExpAlg (Lit d) = show d
ppExpAlg (Add (_, a) (_, b)) = a ++ " + " ++ b
ppExpAlg (Sub (_, a) (e2, b)) =
  a ++ " - " ++ paren (isAdd e2 || isSub e2) b
ppExpAlg (Mul (e1, a) (e2, b)) =
  paren (isAdd e1 || isSub e1) a ++ " * " ++ paren (isAdd e2 || isSub e2) b
ppExpAlg (Div (e1, a) (e2, b)) =
  paren (isAdd e1 || isSub e1) a ++ " / " ++ paren (isComplex e2) b
  where
    isComplex (Fix (Add _ _)) = True
    isComplex (Fix (Sub _ _)) = True
    isComplex (Fix (Mul _ _)) = True
    isComplex (Fix (Div _ _)) = True
    isComplex _ = False
ppExpAlg (Neg (_, a)) = function "negate" a
ppExpAlg (Exp (_, a)) = function "exp" a
ppExpAlg (Log (_, a)) = function "log" a
ppExpAlg (Sqrt (_, a)) = function "sqrt" a
ppExpAlg (Sin (_, a)) = function "sin" a
ppExpAlg (Cos (_, a)) = function "cos" a
ppExpAlg Var = "x"

paren :: Bool -> String -> String
paren b x
  | b = "(" ++ x ++ ")"
  | otherwise = x

function name arg =
  name ++ paren True arg

isAdd :: Expr -> Bool
isAdd (Fix (Add _ _)) = True
isAdd _ = False

isSub :: Expr -> Bool
isSub (Fix (Sub _ _)) = True
isSub _ = False

isLit :: Expr -> Bool
isLit (Fix (Lit _)) = True
isLit _ = False

isVar :: Expr -> Bool
isVar (Fix Var) = True
isVar _ = False

-- * Simple evaluator

-- | Evaluate an 'Expr'ession using standard
--   'Num', 'Fractional' and 'Floating' operations.
eval :: Expr -> (Double -> Double)
eval fexpr x = cata alg fexpr
  where
    alg e = case e of
      Var -> x
      Lit d -> d
      Add a b -> a + b
      Sub a b -> a - b
      Mul a b -> a * b
      Div a b -> a / b
      Neg a -> negate a
      Exp a -> exp a
      Log a -> log a
      Sqrt a -> sqrt a
      Sin a -> sin a
      Cos a -> cos a

-- * Code generation

-- | Helper for calling intrinsics for 'exp', 'log' and friends.
callDblfun ::
  LLVMIR.MonadIRBuilder m => LLVM.Operand -> LLVM.Operand -> m LLVM.Operand
callDblfun fun arg = LLVMIR.call fun [(arg, [])]

xparam :: LLVMIR.ParameterName
xparam = LLVMIR.ParameterName "x"

-- | Generate @declare@ statements for all the intrinsics required for
--   executing the given expression and return a mapping from function
--   name to 'Operand' so that we can very easily refer to those functions
--   for calling them, when generating the code for the expression itself.
declarePrimitives ::
  LLVMIR.MonadModuleBuilder m => Expr -> m (Map String LLVM.Operand)
declarePrimitives expr = fmap Map.fromList
  $ forM primitives
  $ \primName -> do
    f <-
      LLVMIR.extern
        (LLVM.mkName ("llvm." <> primName <> ".f64"))
        [LLVM.double]
        LLVM.double
    return (primName, f)
  where
    primitives = Set.toList (cata alg expr)
    alg (Exp ps) = Set.insert "exp" ps
    alg (Log ps) = Set.insert "log" ps
    alg (Sqrt ps) = Set.insert "sqrt" ps
    alg (Sin ps) = Set.insert "sin" ps
    alg (Cos ps) = Set.insert "cos" ps
    alg e = fold e

-- | Generate an LLVM IR module for the given expression,
--   including @declare@ statements for the intrinsics and
--   a function, always called @f@, that will perform the copoutations
--   described by the 'Expr'ession.
codegen :: Expr -> LLVM.Module
codegen fexpr = LLVMIR.buildModule "arith.ll" $ do
  prims <- declarePrimitives fexpr
  _ <- LLVMIR.function "f" [(LLVM.double, xparam)] LLVM.double $ \[arg] -> do
    res <- cataM (alg arg prims) fexpr
    LLVMIR.ret res
  return ()
  where
    alg arg _ (Lit d) =
      return (LLVM.ConstantOperand $ LLVM.Float $ LLVM.Double d)
    alg arg _ Var = return arg
    alg arg _ (Add a b) = LLVMIR.fadd a b `LLVMIR.named` "x"
    alg arg _ (Sub a b) = LLVMIR.fsub a b `LLVMIR.named` "x"
    alg arg _ (Mul a b) = LLVMIR.fmul a b `LLVMIR.named` "x"
    alg arg _ (Div a b) = LLVMIR.fdiv a b `LLVMIR.named` "x"
    alg arg ps (Neg a) = do
      z <- alg arg ps (Lit 0)
      LLVMIR.fsub z a `LLVMIR.named` "x"
    alg arg ps (Exp a) = callDblfun (ps Map.! "exp") a `LLVMIR.named` "x"
    alg arg ps (Log a) = callDblfun (ps Map.! "log") a `LLVMIR.named` "x"
    alg arg ps (Sqrt a) = callDblfun (ps Map.! "sqrt") a `LLVMIR.named` "x"
    alg arg ps (Sin a) = callDblfun (ps Map.! "sin") a `LLVMIR.named` "x"
    alg arg ps (Cos a) = callDblfun (ps Map.! "cos") a `LLVMIR.named` "x"

codegenText :: Expr -> Text
codegenText = LLVMPretty.ppllvm . codegen

printCodegen :: Expr -> IO ()
printCodegen = Text.putStrLn . codegenText

-- * JIT compilation & loading

-- | This allows us to call dynamically loaded functions
foreign import ccall "dynamic"
  mkDoubleFun :: FunPtr (Double -> Double) -> (Double -> Double)

resolver ::
  JIT.IRCompileLayer l ->
  JIT.MangledSymbol ->
  IO (Either JIT.JITSymbolError JIT.JITSymbol)
resolver compileLayer symbol =
  JIT.findSymbol compileLayer symbol True

symbolFromProcess :: JIT.MangledSymbol -> IO JIT.JITSymbol
symbolFromProcess sym =
  (\addr -> JIT.JITSymbol addr JIT.defaultJITSymbolFlags)
    <$> JIT.getSymbolAddressInProcess sym

resolv :: JIT.IRCompileLayer l -> JIT.SymbolResolver
resolv cl = JIT.SymbolResolver (\sym -> JIT.findSymbol cl sym True)

printIR :: MonadIO m => ByteString -> m ()
printIR = liftIO . BS.putStrLn . ("\n*** LLVM IR ***\n\n" <>)

-- | JIT-compile the given 'Expr'ession and use the resulting function.
withSimpleJIT ::
  NFData a =>
  Expr ->
  -- | what to do with the generated functiion
  ((Double -> Double) -> a) ->
  IO a
withSimpleJIT expr doFun = do
  resolvers <- newIORef Map.empty
  JIT.withContext $ \context -> (>>) (JIT.loadLibraryPermanently Nothing)
    $ JIT.withModuleFromAST context (codegen expr)
    $ \mod' ->
      JIT.withHostTargetMachine Reloc.PIC CodeModel.Default CodeGenOpt.Default $ \tm ->
        JIT.withExecutionSession $ \es ->
          JIT.withObjectLinkingLayer es (\k -> fmap (\rs -> rs Map.! k) (readIORef resolvers)) $ \objectLayer ->
            JIT.withIRCompileLayer objectLayer tm $ \compileLayer -> do
              asm <- JIT.moduleLLVMAssembly mod'
              printExpr expr
              printIR asm
              JIT.withModuleKey es $ \k ->
                JIT.withModule compileLayer k mod' $ do
                  fSymbol <- JIT.mangleSymbol compileLayer "f"
                  Right (JIT.JITSymbol fnAddr _) <- JIT.findSymbol compileLayer fSymbol True
                  let f = mkDoubleFun . castPtrToFunPtr $ wordPtrToPtr fnAddr
                  liftIO (putStrLn "*** Result ***\n")
                  evaluate $ force (doFun f)

-- * Utilities

cataM ::
  (Monad m, Traversable (Base t), Recursive t) =>
  (Base t a -> m a) ->
  t ->
  m a
cataM alg = c
  where
    c = alg <=< traverse c . project

-- * Main

f :: Floating a => a -> a
f t = sin (pi * t / 2) * (1 + sqrt t) ^ 2

main :: IO ()
main = do
  let res1 = map f [0 .. 10] :: [Double]
  res2 <- withSimpleJIT (f x) (\fopt -> map fopt [0 .. 10])
  if res1 == res2
    then putStrLn "results match" >> print res1
    else print res1 >> print res2 >> putStrLn "results don't match"