hakaru-0.6.0: haskell/Language/Hakaru/CodeGen/Wrapper.hs
{-# LANGUAGE BangPatterns,
CPP,
OverloadedStrings,
DataKinds,
FlexibleContexts,
GADTs,
KindSignatures,
RankNTypes,
ScopedTypeVariables,
TypeOperators #-}
----------------------------------------------------------------
-- 2016.06.23
-- |
-- Module : Language.Hakaru.CodeGen.Wrapper
-- Copyright : Copyright (c) 2016 the Hakaru team
-- License : BSD3
-- Maintainer : zsulliva@indiana.edu
-- Stability : experimental
-- Portability : GHC-only
--
-- The purpose of the wrapper is to intelligently wrap CStatements
-- into CFunctions and CProgroms to be printed by 'hkc'
--
----------------------------------------------------------------
module Language.Hakaru.CodeGen.Wrapper
( wrapProgram
, PrintConfig(..)
) where
import Language.Hakaru.Syntax.ABT
import Language.Hakaru.Syntax.AST
import Language.Hakaru.Syntax.IClasses
import Language.Hakaru.Syntax.TypeCheck
import Language.Hakaru.Syntax.TypeOf (typeOf)
import Language.Hakaru.Types.Sing
import Language.Hakaru.CodeGen.CodeGenMonad
import Language.Hakaru.CodeGen.Flatten
import Language.Hakaru.CodeGen.Types
import Language.Hakaru.CodeGen.AST
import Language.Hakaru.CodeGen.Libs
import Language.Hakaru.Types.DataKind (Hakaru(..))
import Control.Monad.State.Strict
import Prelude as P hiding (unlines,exp)
#if __GLASGOW_HASKELL__ < 710
import Control.Applicative
#endif
-- | wrapProgram is the top level C codegen. Depending on the type a program
-- will have a different construction. It will produce an effect in the
-- CodeGenMonad that will produce a standalone C file containing the CPP
-- includes, struct declarations, functions, and sometimes a main.
wrapProgram
:: TypedAST (TrivialABT Term) -- ^ Some Hakaru ABT
-> Maybe String -- ^ Maybe an output name
-> PrintConfig -- ^ show weights?
-> CodeGen ()
wrapProgram tast@(TypedAST typ _) mn pconfig =
do sequence_ . fmap (extDeclare . CPPExt) . header $ typ
cg <- get
when (managedMem cg) $ extDeclare . CPPExt $ gcHeader
when (sharedMem cg) $ extDeclare . CPPExt $ openMpHeader
case (tast,mn) of
( TypedAST (SFun _ _) abt, Just name ) ->
flattenTopLambda abt =<< reserveIdent name
( TypedAST typ' abt, Just name ) ->
-- still buggy for measures
do mfId <- reserveIdent name
funCG (buildType typ') mfId [] $
do outE <- flattenWithName' abt "out"
putStat . CReturn . Just $ outE
( TypedAST typ' abt, Nothing ) ->
mainFunction pconfig typ' abt
header :: Sing (a :: Hakaru) -> [Preprocessor]
header (SMeasure _) = fmap PPInclude ["time.h", "stdlib.h", "stdio.h", "math.h"]
header _ = fmap PPInclude ["stdlib.h", "stdio.h", "math.h"]
--------------------------------------------------------------------------------
-- A Main Function --
--------------------------------------------------------------------------------
{-
Create standalone C program for a Hakaru ABT. This program will also print the
computed value to stdin.
-}
mainFunction
:: ABT Term abt
=> PrintConfig
-> Sing (a :: Hakaru) -- ^ type of program
-> abt '[] (a :: Hakaru) -- ^ Hakaru ABT
-> CodeGen ()
-- when measure, compile to a sampler
mainFunction pconfig typ@(SMeasure _) abt =
do mfId <- reserveIdent "measure"
mainId <- reserveIdent "main"
argVId <- reserveIdent "argv"
argCId <- reserveIdent "argc"
let (argCE:argVE:[]) = fmap CVar [argCId,argVId]
extDeclareTypes typ
-- defined a measure function that returns mdata
funCG (buildType typ) mfId [] $
(putStat . CReturn . Just) =<< flattenWithName' abt "samp"
funCG [CInt] mainId mainArgs $
do isManagedMem <- managedMem <$> get
when isManagedMem (putExprStat gcInit)
nSamples <- parseNumSamples argCE argVE
seedE <- parseSeed argCE argVE
putExprStat $ mkCallE "srand" [seedE]
putStat $ opComment "Run Hakaru Sampler"
printCG pconfig typ (CCall (CVar mfId) []) (Just nSamples)
putStat . CReturn . Just $ intE 0
mainFunction pconfig (SFun _ _) abt =
coalesceLambda abt $ \_ abt' ->
do resId <- reserveIdent "result"
mainId <- reserveIdent "main"
argVId <- reserveIdent "argv"
argCId <- reserveIdent "argc"
funId <- genIdent' "fn"
flattenTopLambda abt funId
let (resE:funE:argCE:argVE:[]) = fmap CVar [resId,funId,argCId,argVId]
typ' = typeOf abt'
funCG [CInt] mainId mainArgs $
do isManagedMem <- managedMem <$> get
when isManagedMem (putExprStat gcInit)
declare typ' resId
mns <- maybeNumSamples argCE argVE typ'
case typ' of
SMeasure _ -> do seedE <- parseSeed argCE argVE
putExprStat $ mkCallE "srand" [seedE]
_ -> return ()
withLambdaDepth' 0 abt $ \d ->
let argErr 0 = ""
argErr n = (argErr (pred n)) ++ "<arg" ++ show n ++ "> " in
ifCG (argCE .<. (intE (d+1)))
(do putExprStat $ printfE
[ stringE $ "Usage: %s " ++ argErr d ++ "\n"
, (index argVE (intE 0)) ]
putExprStat $ mkCallE "abort" [ ])
(return ())
putStat $ opComment "Parse Args"
argEs <- foldLambdaWithIndex 1 abt $ \i (Variable _ _ t) ->
do argE <- localVar' t "arg"
_ <- parseCG t (index argVE (intE i)) argE
return argE
case typ' of
SMeasure _ -> do putStat $ opComment "Run Hakaru Sampler"
printCG pconfig typ' (CCall funE argEs) mns
_ -> do putStat $ opComment "Run Hakaru Program"
putExprStat $ resE .=. (CCall funE argEs)
putStat $ opComment "Print Result"
printCG pconfig typ' resE mns
putStat . CReturn . Just $ intE 0
where withLambdaDepth'
:: ABT Term abt
=> Integer
-> abt '[] a
-> (Integer -> CodeGen ())
-> CodeGen ()
withLambdaDepth' n abt_ k =
caseVarSyn abt_
(const (k n))
(\term ->
case term of
(Lam_ :$ body :* End) ->
caseBind body $ \_ abt_' ->
withLambdaDepth' (succ n) abt_' k
_ -> k n)
maybeNumSamples
:: CExpr -> CExpr -> Sing (a :: Hakaru) -> CodeGen (Maybe CExpr)
maybeNumSamples c v (SMeasure _) = Just <$> parseNumSamples c v
maybeNumSamples _ _ _ = return Nothing
-- just a computation
mainFunction pconfig typ abt =
do resId <- reserveIdent "result"
mainId <- reserveIdent "main"
let resE = CVar resId
funCG [CInt] mainId [] $
do declare typ resId
isManagedMem <- managedMem <$> get
when isManagedMem (putExprStat gcInit)
flattenABT abt resE
printCG pconfig typ resE Nothing
putStat . CReturn . Just $ intE 0
mainArgs :: [CDecl]
mainArgs = [ CDecl [CTypeSpec CInt]
[(CDeclr Nothing (CDDeclrIdent (Ident "argc")), Nothing)]
, CDecl [CTypeSpec CChar]
[(CDeclr (Just (CPtrDeclr []))
(CDDeclrArr (CDDeclrIdent (Ident "argv")) Nothing)
, Nothing)]
]
{- the number of samples is set to -1 by default -}
parseNumSamples :: CExpr -> CExpr -> CodeGen CExpr
parseNumSamples argc argv =
do itE <- localVar SNat
outE <- localVar SNat
putStat $ opComment "Num Samples?"
putExprStat $ outE .=. (intE (-1))
forCG (itE .=. (intE 1))
(itE .<. argc)
(CUnary CPostIncOp itE)
(ifCG ((index (index argv itE) (intE 0) .==. (charE '-')) .&&.
(index (index argv itE) (intE 1) .==. (charE 'n')))
(putExprStat $
sscanfE [index argv itE,stringE "-n%d",address outE])
(return ()))
putStat $ opComment "End Num Samples?"
return outE
{- the randome seed is set to time(NULL) by default -}
parseSeed :: CExpr -> CExpr -> CodeGen CExpr
parseSeed argc argv =
do itE <- localVar SNat
outE <- localVar SNat
putStat $ opComment "Random Seed?"
putExprStat $ outE .=. (mkCallE "time" [ CVar . Ident $ "NULL"])
forCG (itE .=. (intE 1))
(itE .<. argc)
(CUnary CPostIncOp itE)
(ifCG ((index (index argv itE) (intE 0) .==. (charE '-')) .&&.
(index (index argv itE) (intE 1) .==. (charE 's')))
(putExprStat $
sscanfE [index argv itE,stringE "-s%d",address outE])
(return ()))
putStat $ opComment "End Random Seed?"
return outE
--------------------------------------------------------------------------------
-- Parsing Values --
--------------------------------------------------------------------------------
parseCG :: Sing (a :: Hakaru) -> CExpr -> CExpr -> CodeGen CExpr
parseCG (SArray t) from to =
do fpId <- genIdent' "fp"
buffId <- genIdent' "buff"
declare' $ CDecl [CTypeSpec fileT]
[(CDeclr (Just (CPtrDeclr []))
(CDDeclrIdent fpId)
, Nothing)]
declare' $ CDecl [CTypeSpec CChar]
[(CDeclr Nothing
(CDDeclrArr (CDDeclrIdent buffId) (Just (intE 1024)))
, Nothing)]
let fpE = CVar fpId
buffE = CVar buffId
putExprStat $ fpE .=. (fopenE from (stringE "r"))
itE <- localVar SNat
putExprStat $ itE .=. (intE 0)
whileCG (fgetsE buffE (intE 1024) fpE .!=. nullE)
(putExprStat $ CUnary CPostIncOp itE)
putExprStat $ arraySize to .=. itE
putMallocStat (arrayData to) itE t
putExprStat $ itE .=. (intE 0)
putExprStat $ rewindE fpE
whileCG (fgetsE buffE (intE 1024) fpE .!=. nullE)
(do checkE <- parseCG t buffE (index (arrayData to) itE)
ifCG (checkE .==. (intE 1))
(putExprStat $ CUnary CPostIncOp itE)
(putExprStat $ CUnary CPostDecOp (arraySize to)))
putExprStat $ fcloseE fpE
localVar SNat
parseCG t from to =
do checkE <- localVar SNat
putExprStat $ checkE .=. sscanfE [from,stringE . parseFormat $ t,address to]
case t of
SProb -> putExprStat $ to .=. logE to
_ -> return ()
return checkE
parseFormat :: Sing (a :: Hakaru) -> String
parseFormat SInt = "%d"
parseFormat SNat = "%u"
parseFormat SReal = "%lf"
parseFormat SProb = "%lf"
parseFormat t = error $ "parseCG{" ++ show t ++ "}: no available parsing form"
--------------------------------------------------------------------------------
-- Printing Values --
--------------------------------------------------------------------------------
{-
In HKC the printconfig is parsed from the command line. The default being that
we don't show weights and probabilities are printed as normal real values.
-}
data PrintConfig
= PrintConfig { noWeights :: Bool
, showProbInLog :: Bool
} deriving Show
printCG
:: PrintConfig
-> Sing (a :: Hakaru) -- ^ Hakaru type to be printed
-> CExpr -- ^ CExpr representing value
-> Maybe CExpr -- ^ If measure type, expr for num samples
-> CodeGen ()
printCG pconfig mtyp@(SMeasure typ) sampleFunc (Just numSamples) =
do mE <- localVar' mtyp "m"
itE <- localVar SNat
putExprStat $ itE .=. numSamples
whileCG (itE .!=. (intE 0)) $
do putExprStat $ mE .=. sampleFunc
case noWeights pconfig of
True -> printCG pconfig typ (mdataSample mE) Nothing
False -> do putExprStat $
printfE [ stringE (printFormat pconfig SProb "\t")
, expE (mdataWeight mE) ]
printCG pconfig typ (mdataSample mE) Nothing
ifCG (numSamples .>=. (intE 0))
(putExprStat $ CUnary CPostDecOp itE)
(return ())
printCG pconfig (SArray typ) arg Nothing =
do itE <- localVar' SNat "it"
putString "[ "
seqDo $
forCG (itE .=. (intE 0))
(itE .<. (arraySize arg))
(CUnary CPostIncOp itE)
(putExprStat
$ printfE [ stringE $ printFormat pconfig typ " "
, index (arrayData arg) itE ])
putString "]\n"
where putString s = putExprStat $ printfE [stringE s]
-- bool and unit
printCG _ (SData (STyCon sym) _) arg Nothing =
case ssymbolVal sym of
"Unit" -> putExprStat $ printfE [stringE "()\n"]
"Bool" -> ifCG (datumIndex arg .==. (intE 0))
(putExprStat $ printfE [stringE "true\n"])
(putExprStat $ printfE [stringE "false\n"])
_ -> error $ show sym
printCG pconfig SProb arg Nothing =
putExprStat $ printfE
[ stringE $ printFormat pconfig SProb "\n"
, if showProbInLog pconfig
then arg
else expE arg ]
printCG pconfig typ arg Nothing =
putExprStat $ printfE
[ stringE $ printFormat pconfig typ "\n"
, arg ]
-- we should only have a number of samples if it a measure
printCG _ _ _ (Just _) = error "this should not happen"
printFormat :: PrintConfig -> Sing (a :: Hakaru) -> (String -> String)
printFormat _ SInt = \s -> "%d" ++ s
printFormat _ SNat = \s -> "%d" ++ s
printFormat c SProb = \s -> if showProbInLog c
then "exp(%.15f)" ++ s
else "%.15f" ++ s
printFormat _ SReal = \s -> "%.15f" ++ s
printFormat c (SMeasure t) = if noWeights c
then printFormat c t
else \s -> if showProbInLog c
then "exp(%.15f) " ++ printFormat c t s
else "%.15f " ++ printFormat c t s
printFormat c (SArray t) = printFormat c t
printFormat _ (SFun _ _) = id
printFormat _ (SData _ _) = \s -> "TODO: printft datum" ++ s
--------------------------------------------------------------------------------
-- Wrapping Lambdas --
--------------------------------------------------------------------------------
{-
Lambdas become function in C. The Hakaru ABT only allows one arguement for each
lambda. So at the top level of a Hakaru program that is a function there may be
several nested lambdas. In C however, we can and should coalesce these into one
function with several arguements. This is what flattenTopLambda is for.
-}
flattenTopLambda
:: ABT Term abt
=> abt '[] a
-> Ident
-> CodeGen ()
flattenTopLambda abt name =
coalesceLambda abt $ \vars abt' ->
let varMs = foldMap11 (\v -> [mkVarDecl v =<< createIdent' "param" v]) vars
typ = typeOf abt'
in do argDecls <- sequence varMs
funCG (buildType typ) name argDecls $
(putStat . CReturn . Just) =<< flattenWithName' abt' "out"
-- do at top level
where mkVarDecl :: Variable (a :: Hakaru) -> Ident -> CodeGen CDecl
mkVarDecl (Variable _ _ SInt) = return . typeDeclaration SInt
mkVarDecl (Variable _ _ SNat) = return . typeDeclaration SNat
mkVarDecl (Variable _ _ SProb) = return . typeDeclaration SProb
mkVarDecl (Variable _ _ SReal) = return . typeDeclaration SReal
mkVarDecl (Variable _ _ (SArray t)) = \i ->
extDeclareTypes (SArray t) >> (return $ arrayDeclaration t i)
mkVarDecl (Variable _ _ d@(SData _ _)) = \i ->
extDeclareTypes d >> (return $ datumDeclaration d i)
mkVarDecl v = error $ "flattenSCon.Lam_.mkVarDecl cannot handle vars of type " ++ show v
coalesceLambda
:: ABT Term abt
=> abt '[] a
-> ( forall (ys :: [Hakaru]) b
. List1 Variable ys -> abt '[] b -> r)
-> r
coalesceLambda abt_ k =
caseVarSyn abt_ (const (k Nil1 abt_)) $ \term ->
case term of
(Lam_ :$ body :* End) ->
caseBind body $ \v body' ->
coalesceLambda body' $ \vars body'' -> k (Cons1 v vars) body''
_ -> k Nil1 abt_
foldLambdaWithIndex
:: ABT Term abt
=> Integer
-> abt '[] a
-> ( forall (x :: Hakaru)
. Integer
-> Variable x
-> CodeGen CExpr)
-> CodeGen [CExpr]
foldLambdaWithIndex n abt_ k =
caseVarSyn abt_
(const (return []))
(\term ->
case term of
(Lam_ :$ body :* End) ->
caseBind body $ \v abt_' ->
(do x <- k n v
xs <- foldLambdaWithIndex (succ n) abt_' k
return (x:xs))
_ -> return [])