llvm-3.2.0.1: examples/Array.hs

module Array where
import Data.Word

import LLVM.Core
--import LLVM.ExecutionEngine
import LLVM.Util.Loop
import LLVM.Util.Optimize

cg :: CodeGenModule (Function (Double -> IO (Ptr Double)))
cg = do
    dotProd <- createFunction InternalLinkage $ \ size aPtr aStride bPtr bStride -> do
        r <- forLoop (valueOf 0) size (valueOf 0) $ \ i s -> do
            ai <- mul aStride i
            bi <- mul bStride i
            ap <- getElementPtr aPtr (ai, ())
            bp <- getElementPtr bPtr (bi, ())
            a <- load ap
            b <- load bp
            ab <- mul a b
            add (s :: Value Double) ab
        ret r
    let _ = dotProd :: Function (Word32 -> Ptr Double -> Word32 -> Ptr Double -> Word32 -> IO Double)

    -- multiply a:[n x m], b:[m x l]
    matMul <- createFunction InternalLinkage $ \ n m l aPtr bPtr cPtr -> do
        forLoop (valueOf 0) n () $ \ ni () -> do
            forLoop (valueOf 0) l () $ \ li () -> do
                ni' <- mul ni m
                row <- getElementPtr aPtr (ni', ())
                col <- getElementPtr bPtr (li, ())
                x <- call dotProd m row (valueOf 1) col m
                j <- add ni' li
                p <- getElementPtr cPtr (j, ())
                store x p
                return ()
        ret ()
    let _ = matMul :: Function (Word32 -> Word32 -> Word32 -> Ptr Double -> Ptr Double -> Ptr Double -> IO ())

    let fillArray _ [] = return ()
        fillArray ptr (x:xs) = do store x ptr; ptr' <- getElementPtr ptr (1::Word32,()); fillArray ptr' xs

    test <- createNamedFunction ExternalLinkage "test" $ \ x -> do
        a <- arrayMalloc (4 :: Word32)
        fillArray a $ map valueOf [1,2,3,4]
        b <- arrayMalloc (4 :: Word32)
        fillArray b [x,x,x,x]
        c <- arrayMalloc (4 :: Word32)
        _ <- call matMul (valueOf 2) (valueOf 2) (valueOf 2) a b c
        ret c
    let _ = test :: Function (Double -> IO (Ptr Double))

    return test

main :: IO ()
main = do
    -- Initialize jitter
    initializeNativeTarget
    m <- newModule
    _f <- defineModule m cg
    writeBitcodeToFile "Arr.bc" m
    _ <- optimizeModule 3 m
    writeBitcodeToFile "Arr-opt.bc" m