llvm-tf-3.0.3.2: example/Vector.hs

{-# LANGUAGE TypeOperators #-}
module Main (main) where

import Convert

import LLVM.ExecutionEngine
          (runEngineAccessWithModule, generateFunction, getPointerToFunction)
import LLVM.Util.Optimize (optimizeModule, )
import LLVM.Util.Loop (forLoop, )
import LLVM.Core

import qualified Type.Data.Num.Decimal.Number as Dec
import Type.Data.Num.Decimal.Literal (D16, )

import Control.Monad (liftM2, )
import Data.Word (Word32, )

-- Type of vector elements.
type T = Float

-- Number of vector elements.
type N = D16

retAccName, fName :: String
retAccName = "retacc"
fName = "vectest"

cgvec :: CodeGenModule (Function (T -> IO T))
cgvec = do
    -- A global variable that vectest messes with.
    acc <- createNamedGlobal False ExternalLinkage "acc" (constOf (0 :: T))

    -- Return the global variable.
    retAcc <- createNamedFunction ExternalLinkage retAccName $ do
        vacc <- load acc
        ret vacc
    let _ = retAcc :: Function (IO T)  -- Force the type of retAcc.

    -- A function that tests vector operations.
    f <- createNamedFunction ExternalLinkage fName $ \ x -> do

        let v = value (zero :: ConstValue (Vector N T))
	    n = Dec.integralFromSingleton (Dec.singleton :: Dec.Singleton N) :: Word32

        -- Fill the vector with x, x+1, x+2, ...
        (_, v1) <- forLoop (valueOf 0) (valueOf n) (x, v) $ \ i (x1, v1) -> do
            x1' <- add x1 (valueOf (1::T))
	    v1' <- insertelement v1 x1 i
	    return (x1', v1')

	-- Elementwise cubing of the vector.
	vsq <- mul v1 v1
        vcb <- mul vsq v1

        -- Sum the elements of the vector.
        s <- forLoop (valueOf 0) (valueOf n) (valueOf 0) $ \ i s -> do
            y <- extractelement vcb i
     	    add s (y :: Value T)

        -- Update the global variable.
        vacc <- load acc
        vacc' <- add vacc s
        store vacc' acc

        ret (s :: Value T)

--    liftIO $ dumpValue f
    return f

main :: IO ()
main = do
    -- Initialize jitter
    initializeNativeTarget
    -- First run standard code.
    m <- newModule
    iovec <- defineModule m $ setTarget hostTriple >> cgvec

    fptr <- runEngineAccessWithModule m $ getPointerToFunction iovec
    let fvec = convert fptr

    fvec 10 >>= print

    vec <- runEngineAccessWithModule m $ generateFunction iovec

    vec 10 >>= print

    -- And then optimize and run.
    _ <- optimizeModule 1 m

    funcs <- getModuleValues m
    print $ map fst funcs

    let iovec' :: Function (T -> IO T)
        Just iovec' = castModuleValue =<< lookup fName funcs
	ioretacc' :: Function (IO T)
        Just ioretacc' = castModuleValue =<< lookup retAccName funcs

    (vec', retacc') <-
        runEngineAccessWithModule m $
        liftM2 (,) (generateFunction iovec') (generateFunction ioretacc')

    dumpValue iovec'

    vec' 10 >>= print
    vec' 0 >>= print
    retacc' >>= print