diff --git a/HerbiePlugin.cabal b/HerbiePlugin.cabal new file mode 100644 --- /dev/null +++ b/HerbiePlugin.cabal @@ -0,0 +1,114 @@ +-- Initial herbie-haskell.cabal generated by cabal init. For further +-- documentation, see http://haskell.org/cabal/users-guide/ + +-- The name of the package. +name: HerbiePlugin + +-- The package version. See the Haskell package versioning policy (PVP) +-- for standards guiding when and how versions should be incremented. +-- http://www.haskell.org/haskellwiki/Package_versioning_policy +-- PVP summary: +-+------- breaking API changes +-- | | +----- non-breaking API additions +-- | | | +--- code changes with no API change +version: 0.1.0.0 + +-- A short (one-line) description of the package. +synopsis: automatically improve your code's numeric stability + +-- A longer description of the package. +description: + This package contains a GHC plugin that automatically improves the numerical stability of your Haskell code. + See for details on how it works and how to use it. + +-- URL for the project homepage or repository. +homepage: github.com/mikeizbicki/herbie-haskell + +-- The license under which the package is released. +license: BSD3 + +-- The file containing the license text. +license-file: LICENSE + +-- The package author(s). +author: Mike Izbicki + +-- An email address to which users can send suggestions, bug reports, and +-- patches. +maintainer: mike@izbicki.me + +-- A copyright notice. +-- copyright: + +category: Math + +build-type: Simple + +-- Extra files to be distributed with the package, such as examples or a +-- README. +-- extra-source-files: + +data-files: + Herbie.db + +data-dir: + data + +-- Constraint on the version of Cabal needed to build this package. +cabal-version: >=1.10 + +source-repository head + type: git + location: http://github.com/mikeizbicki/HerbiePlugin + + +library + -- Modules exported by the library. + exposed-modules: + Herbie + + -- Modules included in this library but not exported. + other-modules: + Herbie.CoreManip + Herbie.ForeignInterface + Herbie.MathExpr + Herbie.MathInfo + Show + Paths_HerbiePlugin + + -- LANGUAGE extensions used by modules in this package. + default-extensions: + MultiWayIf + ScopedTypeVariables + DeriveGeneric + DeriveAnyClass + + -- Other library packages from which modules are imported. + build-depends: base >=4.8 && <4.9 + , ghc + , template-haskell + , process >= 1.1.0.0 + , sqlite-simple + , text + , directory + , deepseq + , mtl + + -- Directories containing source files. + hs-source-dirs: src + + -- Base language which the package is written in. + default-language: Haskell2010 + +Test-Suite Tests + default-language: Haskell2010 + type: exitcode-stdio-1.0 + hs-source-dirs: test + main-is: Tests.hs + + ghc-options: + -fplugin=Herbie +-- -dcore-lint + + build-depends: + subhask, + HerbiePlugin diff --git a/LICENSE b/LICENSE new file mode 100644 --- /dev/null +++ b/LICENSE @@ -0,0 +1,29 @@ +Copyright (c) 2015, Mike Izbicki +All rights reserved. + +Redistribution and use in source and binary forms, with or without +modification, are permitted provided that the following conditions are met: + + * Redistributions of source code must retain the above copyright + notice, this list of conditions and the following disclaimer. + + * Redistributions in binary form must reproduce the above + copyright notice, this list of conditions and the following + disclaimer in the documentation and/or other materials provided + with the distribution. + + * Neither the name of Mike Izbicki nor the names of other + contributors may be used to endorse or promote products derived + from this software without specific prior written permission. + +THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS +"AS IS" AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT +LIMITED TO, THE IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR +A PARTICULAR PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT +OWNER OR CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, +SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT +LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, +DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY +THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT +(INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE +OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. diff --git a/Setup.hs b/Setup.hs new file mode 100644 --- /dev/null +++ b/Setup.hs @@ -0,0 +1,2 @@ +import Distribution.Simple +main = defaultMain diff --git a/data/Herbie.db b/data/Herbie.db new file mode 100644 Binary files /dev/null and b/data/Herbie.db differ diff --git a/src/Herbie.hs b/src/Herbie.hs new file mode 100644 --- /dev/null +++ b/src/Herbie.hs @@ -0,0 +1,185 @@ +module Herbie + ( plugin + , pass + ) + where + +import Class +import DsBinds +import DsMonad +import ErrUtils +import GhcPlugins +import Id +import Unique +import MkId +import PrelNames +import TcRnMonad +import TcSimplify + +import Control.Monad +import Control.Monad.Except +import Data.Maybe + +import Herbie.CoreManip +import Herbie.ForeignInterface +import Herbie.MathExpr +import Herbie.MathInfo + +import Debug.Trace + +import Prelude +import Show +import Data.IORef + +plugin :: Plugin +plugin = defaultPlugin + { installCoreToDos = install + } + +install :: [CommandLineOption] -> [CoreToDo] -> CoreM [CoreToDo] +install opts todo = do + putMsgS "Compiling with Herbie floating point stabilization" + reinitializeGlobals + return (CoreDoPluginPass "MathInfo" (pass opts) : todo) + +pass :: [CommandLineOption] -> ModGuts -> CoreM ModGuts +pass opts guts = do + dflags <- getDynFlags + liftIO $ writeIORef dynFlags_ref dflags + bindsOnlyPass (mapM (modBind opts guts)) guts + +-- | This function gets run on each binding on the Haskell source file. +modBind :: [CommandLineOption] -> ModGuts -> CoreBind -> CoreM CoreBind +modBind opts guts bndr@(Rec _) = return bndr +modBind opts guts bndr@(NonRec b e) = do +-- dflags <- getDynFlags +-- putMsgS "" +-- putMsgS $ showSDoc dflags (ppr b) +-- ++ "::" +-- ++ showSDoc dflags (ppr $ varType b) +-- putMsgS $ myshow dflags e +-- return bndr + e' <- go [] e + return $ NonRec b e' + where + -- Recursively descend into the expression e. + -- For each math expression we find, run Herbie on it. + -- We need to save each dictionary we find because + -- it might be needed to create the replacement expressions. + go dicts e = do + dflags <- getDynFlags + case mkMathInfo dflags dicts (varType b) e of + + -- not a math expression, so recurse into subexpressions + Nothing -> case e of + + -- Lambda expression: + -- If the variable is a dictionary, add it to the list; + -- Always recurse into the subexpression + -- + -- FIXME: + -- Currently, we're removing deadness annotations from any dead variables. + -- This is so that we can use all the dictionaries that the type signatures allow. + -- Core lint complains about using dead variables if we don't. + -- This causes us to remove ALL deadness annotations in the entire program. + -- I'm not sure the drawback of this. + -- This could be fixed by having a second pass through the code + -- to remove only the appropriate deadness annotations. + Lam a b -> do + let a' = undeadenId a + b' <- go (extractDicts a'++dicts) b + return $ Lam a' b' + + -- Let binding: + -- If the variable is a dictionary, add it to the list; + -- Always recurse into the subexpression + Let (NonRec a e) b -> do + let a' = undeadenId a + e' <- go dicts e + b' <- go (extractDicts a'++dicts) b + return $ Let (NonRec a' e') b' + + Let (Rec bndrs) expr -> do + bndrs' <- forM bndrs $ \(a,e) -> do + let a' = undeadenId a + e' <- go dicts e + return (a',e') + expr' <- go dicts expr + return $ Let (Rec bndrs') expr' + + -- Function application: + -- Math expressions may appear on either side, so recurse on both + App a b -> do + a' <- go dicts a + b' <- go dicts b + return $ App a' b' + + -- Case statement: + -- Math expressions may appear in the condition or in any of the branches + Case cond w t es -> do + cond' <- go dicts cond + es' <- forM es $ \ (altcon, xs, expr) -> do + expr' <- go dicts expr + return $ (altcon, xs, expr') + return $ Case cond' w t es' + + -- Ticks and Casts are just annotating extra information on an expression. + -- We ignore the extra information and recurse into the expression. + Tick a b -> do + b' <- go dicts b + return $ Tick a b' + + Cast a b -> do + a' <- go dicts a + return $ Cast a' b + + -- There's nothing to do for these statements. + -- They form the recursion's base case. + Var v -> return $ Var v + Lit l -> return $ Lit l + Type t -> return $ Type t + Coercion c -> return $ Coercion c + + -- We found a math expression, so process it + Just mathInfo -> do + putMsgS $ "Found math expression within binding " + ++ showSDoc dflags (ppr b) + ++ " :: " + ++ showSDoc dflags (ppr $ varType b) + putMsgS $ " original expression = "++pprMathInfo mathInfo + let dbgInfo = DbgInfo + { dbgComments = concat opts + , modName = showSDoc dflags (ppr $ moduleName $ mg_module guts) + , functionName = showSDoc dflags (ppr b) + , functionType = showSDoc dflags (ppr $ varType b) + } + res <- liftIO $ stabilizeMathExpr dbgInfo $ getMathExpr mathInfo + let mathInfo' = mathInfo { getMathExpr = cmdout res } + putMsgS $ " improved expression = "++pprMathInfo mathInfo' + putMsgS $ " original error = "++show (errin res)++" bits" + putMsgS $ " improved error = "++show (errout res)++" bits" + ret <- runExceptT $ mathInfo2expr guts mathInfo' + case ret of + Left str -> do + putMsgS " WARNING: Not substituting the improved expression into your code" + putMsgS str + return e + Right e' -> do +-- putMsgS $ " before = " ++ myshow dflags e +-- putMsgS $ " after = " ++ myshow dflags e' + return e' + +-- | Return a list with the given variable if the variable is a dictionary or tuple of dictionaries, +-- otherwise return []. +extractDicts :: Var -> [Var] +extractDicts v = case classifyPredType (varType v) of + ClassPred _ _ -> [v] + EqPred _ _ _ -> [v] + TuplePred _ -> [v] + IrredPred _ -> [] + +-- | If a variable is marked as dead, remove the marking +undeadenId :: Var -> Var +undeadenId a = if isDeadBinder a + then setIdOccInfo a NoOccInfo + else a diff --git a/src/Herbie/CoreManip.hs b/src/Herbie/CoreManip.hs new file mode 100644 --- /dev/null +++ b/src/Herbie/CoreManip.hs @@ -0,0 +1,622 @@ +{-# LANGUAGE CPP #-} +module Herbie.CoreManip + where + +import Class +import DsBinds +import DsMonad +import ErrUtils +import GhcPlugins hiding (trace) +import Unique +import MkId +import PrelNames +import UniqSupply +import TcRnMonad +import TcSimplify +import Type + +import Control.Monad +import Control.Monad.Except +import Control.Monad.Trans +import Data.Char +import Data.List +import Data.Maybe +import Data.Ratio + +import Herbie.MathExpr + +import Prelude +import Show + +-- import Debug.Trace hiding (traceM) +trace a b = b +traceM a = return () + +-------------------------------------------------------------------------------- + +instance MonadUnique m => MonadUnique (ExceptT e m) where + getUniqueSupplyM = lift getUniqueSupplyM + +instance (Monad m, HasDynFlags m) => HasDynFlags (ExceptT e m) where + getDynFlags = lift getDynFlags + +instance MonadThings m => MonadThings (ExceptT e m) where + lookupThing name = lift $ lookupThing name + +---------------------------------------- +-- core manipulation + +-- | Converts a string into a Core variable +getVar :: ModGuts -> String -> ExceptT String CoreM Var +getVar guts opstr = do + let opname = getName guts opstr + hscenv <- lift getHscEnv + dflags <- getDynFlags + eps <- liftIO $ hscEPS hscenv + optype <- case lookupNameEnv (eps_PTE eps) opname of + Just (AnId i) -> return $ varType i + _ -> throwError $ " WARNING: variable \""++opstr++"\" not in scope" + return $ mkGlobalVar VanillaId opname optype vanillaIdInfo + + where + getName :: ModGuts -> String -> Name + getName guts str = case filter isCorrectVar (concat $ occEnvElts (mg_rdr_env guts)) of + xs -> if length xs>0 + then gre_name $ head $ xs + else error $ "getName: '"++str++"'" + where + isCorrectVar x = (getString $ gre_name x) == str + && (str == "abs" || case gre_par x of NoParent -> False; _ -> True) + +-- | Like "decorateFunction", but first finds the function variable given a string. +getDecoratedFunction :: ModGuts -> String -> Type -> [CoreExpr] -> ExceptT String CoreM CoreExpr +getDecoratedFunction guts str t preds = do + f <- getVar guts str + decorateFunction guts f t preds + +-- | Given a variable that contains a function, +-- the type the function is being applied to, +-- and all in scope predicates, +-- apply the type and any needed dictionaries to the function. +decorateFunction :: ModGuts -> Var -> Type -> [CoreExpr] -> ExceptT String CoreM CoreExpr +decorateFunction guts f t preds = do + let ([v],unquantified) = extractQuantifiers $ varType f + (cxt,_) = extractContext unquantified + cxt' = substTysWith [v] [t] cxt + + cxt'' <- mapM getDict cxt' + + return $ mkApps (App (Var f) (Type t)) cxt'' + where + getDict :: PredType -> ExceptT String CoreM CoreExpr + getDict pred = do + catchError + (getDictionary guts pred) + (\_ -> getPredEvidence guts pred preds) + +-- | Given a non-polymorphic PredType (e.g. `Num Float`), +-- return the corresponding dictionary. +getDictionary :: ModGuts -> Type -> ExceptT String CoreM CoreExpr +getDictionary guts dictTy = do + let dictVar = mkGlobalVar + VanillaId + (mkSystemName (mkUnique 'z' 1337) (mkVarOcc $ "magicDictionaryName")) + dictTy + vanillaIdInfo + + bnds <- lift $ runTcM guts $ do + loc <- getCtLoc $ GivenOrigin UnkSkol + let nonC = mkNonCanonical $ CtWanted + { ctev_pred = varType dictVar + , ctev_evar = dictVar + , ctev_loc = loc + } + wCs = mkSimpleWC [nonC] + (x, evBinds) <- solveWantedsTcM wCs + bnds <- initDsTc $ dsEvBinds evBinds + +-- liftIO $ do +-- putStrLn $ "dictType="++showSDoc dflags (ppr dictType) +-- putStrLn $ "dictVar="++showSDoc dflags (ppr dictVar) +-- +-- putStrLn $ "nonC="++showSDoc dflags (ppr nonC) +-- putStrLn $ "wCs="++showSDoc dflags (ppr wCs) +-- putStrLn $ "bnds="++showSDoc dflags (ppr bnds) +-- putStrLn $ "x="++showSDoc dflags (ppr x) + + return bnds + + case bnds of + [NonRec _ dict] -> return dict + otherwise -> throwError $ + " WARNING: Cannot satisfy the constraint: "++dbg dictTy + +-- | Given a predicate for which we don't have evidence +-- and a list of expressions that contain evidence for predicates, +-- construct an expression that contains evidence for the given predicate. +getPredEvidence :: ModGuts -> PredType -> [CoreExpr] -> ExceptT String CoreM CoreExpr +getPredEvidence guts pred evidenceExprs = go $ prepEvidence evidenceExprs + where + + go :: [(CoreExpr,Type)] -> ExceptT String CoreM CoreExpr + + -- We've looked at all the evidence, but didn't find anything + go [] = throwError $ + " WARNING: Cannot satisfy the constraint: "++dbg pred + + -- Recursively descend into all the available predicates. + -- The list tracks both the evidence expression (this will change in recursive descent), + -- and the baseTy that gave rise to the expression (this stays constant). + go ((expr,baseTy):exprs) = if exprType expr == pred + + -- The expression we've found matches the predicate. + -- We're done! + then return expr + + -- The expression doesn't match the predicate, + -- so we recurse by searching for sub-predicates within expr + -- and adding them to the list. + else case classifyPredType (exprType expr) of + + -- What we've found contains no more predicates to recurse into, + -- so we don't add anything to the list of exprs to search. + IrredPred _ -> go exprs + + EqPred _ t1 t2 -> trace ("getPredEvidence.go.EP: pred="++dbg pred + ++"; origType="++dbg (baseTy) + ++"; exprType="++dbg (exprType expr) + ) $ case splitAppTy_maybe pred of + Nothing -> trace " A" $ go exprs +-- Just (tyCon,tyApp) -> if baseTy/=tyApp + Just (tyCon,tyApp) -> trace " A'" $ if t1/=tyApp && t2 /=tyApp + then trace (" B: baseTy="++dbg baseTy++"; tyApp="++dbg tyApp) + $ go exprs + else do + let pred' = mkAppTy tyCon $ if t1==tyApp + then t2 + else t1 + getDictionary guts pred' >>= castToType evidenceExprs pred + + -- We've found a class dictionary. + -- Recurse into each field (selId) of the dictionary. + -- Some (but not all) of these may be more dictionaries. + -- + -- FIXME: Multiparamter classes broken + ClassPred c' [ct] -> trace ("getPredEvidence.go.CP: pred="++dbg pred + ++"; origType="++dbg (baseTy) + ++"; exprType="++dbg (exprType expr) + ) $ + go $ + exprs++ + [ ( App (App (Var selId) (Type baseTy)) expr + , baseTy + ) + | selId <- classAllSelIds c' + ] + + ClassPred _ _ -> go exprs + + -- We've found a tuple of evidence. + -- For each field of the tuple we extract it with a case statement, then recurse. + TuplePred preds -> do + trace ("getPredEvidence.go.TP: pred="++dbg pred + ++"; origType="++dbg (baseTy) + ++"; exprType="++dbg (exprType expr) + ) $ return () + + uniqs <- getUniquesM + + traceM $ " tupelems: baseTy="++dbg baseTy++"; preds="++dbg preds + let tupelems = + [ mkLocalVar + VanillaId + (mkSystemName uniq (mkVarOcc $ "a"++show j)) + t' +-- (mkAppTy (fst $ splitAppTys t') baseTy) + vanillaIdInfo + | (j,t',uniq) <- zip3 [0..] preds uniqs + ] + + uniq <- getUniqueM + let wildName = mkSystemName uniq (mkVarOcc $ "wild") + wildVar = mkLocalVar VanillaId wildName (exprType expr) vanillaIdInfo + + let ret = + [ ( Case expr wildVar (varType $ tupelems!!i) + [ ( DataAlt $ tupleCon ConstraintTuple $ length preds + , tupelems + , Var $ tupelems!!i + ) + ] + , baseTy + ) + | (i,t) <- zip [0..] preds + ] + + sequence_ [ traceM $ " ret!!"++show i++"="++myshow dynFlags (fst $ ret!!i) | i<-[0..length ret-1]] + + go $ ret++exprs + +-- | Given some evidence, an expression, and a type: +-- try to prove that the expression can be cast to the type. +-- If it can, return the cast expression. +castToType :: [CoreExpr] -> Type -> CoreExpr -> ExceptT String CoreM CoreExpr +castToType xs castTy inputExpr = if exprType inputExpr == castTy + then return inputExpr + else go $ prepEvidence xs +-- else go $ catMaybes [ (x, extractBaseTy $ exprType x) | x <- xs ] + where + + + go :: [(CoreExpr,Type)] -> ExceptT String CoreM CoreExpr + + -- base case: we've searched through all the evidence, but couldn't create a cast + go [] = throwError $ + " WARNING: Could not cast expression of type "++dbg (exprType inputExpr)++" to "++dbg castTy + + -- recursively try each evidence expression looking for a cast + go ((expr,baseTy):exprs) = case classifyPredType $ exprType expr of + + IrredPred _ -> go exprs + + EqPred _ t1 t2 -> trace ("castToType.go.EP: castTy="++dbg castTy + ++"; origType="++dbg (baseTy) + ++"; exprType="++dbg (exprType expr) + ) $ goEqPred [] castTy (exprType inputExpr) + where + -- Check if a cast is possible. + -- We need to recursively peel off all the type constructors + -- on the inputTyRHS and castTyRHS types. + -- As long as the type constructors match, + -- we might be able to do a cast at any level of the peeling + goEqPred :: [TyCon] -> Type -> Type -> ExceptT String CoreM CoreExpr + goEqPred tyCons castTyRHS inputTyRHS = if + | t1==castTyRHS && t2==inputTyRHS -> mkCast True + | t2==castTyRHS && t1==inputTyRHS -> mkCast False + | otherwise -> case ( splitTyConApp_maybe castTyRHS + , splitTyConApp_maybe inputTyRHS + ) of + (Just (castTyCon, [castTyRHS']), Just (inputTyCon,[inputTyRHS'])) -> + if castTyCon == inputTyCon + then goEqPred (castTyCon:tyCons) castTyRHS' inputTyRHS' + else go exprs + _ -> go exprs + where + + -- Constructs the actual cast from one variable type to another. + -- + -- There's some subtle voodoo in here involving GHC's Roles. + -- Basically, everything gets created as a Nominal role, + -- but the final Coercion needs to be Representational. + -- mkSubCo converts from Nominal into Representational. + -- See https://ghc.haskell.org/trac/ghc/wiki/RolesImplementation + mkCast :: Bool -> ExceptT String CoreM CoreExpr + mkCast isFlipped = do + coboxUniq <- getUniqueM + let coboxName = mkSystemName coboxUniq (mkVarOcc $ "cobox") + coboxType = if isFlipped + then mkCoercionType Nominal castTyRHS inputTyRHS + else mkCoercionType Nominal inputTyRHS castTyRHS + coboxVar = mkLocalVar VanillaId coboxName coboxType vanillaIdInfo + + -- Reapplies the list of tyCons that we peeled off during the recursion. + let mkCoercion [] = if isFlipped + then mkSymCo $ mkCoVarCo coboxVar + else mkCoVarCo coboxVar + mkCoercion (x:xs) = mkTyConAppCo Nominal x [mkCoercion xs] + + wildUniq <- getUniqueM + let wildName = mkSystemName wildUniq (mkVarOcc $ "wild") + wildType = exprType expr + wildVar = mkLocalVar VanillaId wildName wildType vanillaIdInfo + + return $ Case + expr + wildVar + castTy + [ ( DataAlt eqBoxDataCon + , [coboxVar] + , Cast inputExpr $ mkSubCo $ mkCoercion tyCons + ) ] + + -- | FIXME: ClassPred and TuplePred are both handled the same + -- within castToPred and getPredEvidence. + -- They should be factored out? + ClassPred c' [ct] -> go $ + exprs++ + [ ( App (App (Var selId) (Type baseTy)) expr + , baseTy + ) + | selId <- classAllSelIds c' + ] + + ClassPred _ _ -> go exprs + + TuplePred preds -> do + uniqs <- getUniquesM + let tupelems = + [ mkLocalVar + VanillaId + (mkSystemName uniq (mkVarOcc $ "a"++show j)) +-- (mkAppTy (fst $ splitAppTys t') baseTy) + t' + vanillaIdInfo + | (j,t',uniq) <- zip3 [0..] preds uniqs + ] + + uniq <- getUniqueM + let wildName = mkSystemName uniq (mkVarOcc $ "wild") + wildVar = mkLocalVar VanillaId wildName (exprType expr) vanillaIdInfo + + let ret = + [ ( Case expr wildVar (varType $ tupelems!!i) + [ ( DataAlt $ tupleCon ConstraintTuple $ length preds + , tupelems + , Var $ tupelems!!i + ) + ] + , baseTy + ) + | (i,t) <- zip [0..] preds + ] + + go $ ret++exprs + +-- | Each element in the input list must contain evidence of a predicate. +-- The output list contains evidence of a predicate along with a type that will be used for casting. +prepEvidence :: [CoreExpr] -> [(CoreExpr,Type)] +prepEvidence exprs = catMaybes + [ case extractBaseTy $ exprType x of + Just t -> Just (x,t) + Nothing -> Nothing --(x, extractBaseTy $ exprType x) + | x <- exprs + ] + + where + -- Extracts the type that each of our pieces of evidence is applied to + extractBaseTy :: Type -> Maybe Type + extractBaseTy t = case classifyPredType t of + + ClassPred _ [x] -> Just x + + EqPred rel t1 t2 -> if + | t1 == boolTy -> Just t2 + | t2 == boolTy -> Just t1 + | otherwise -> Nothing + + _ -> Nothing + +-- | Return all the TyVars that occur anywhere in the Type +extractTyVars :: Type -> [TyVar] +extractTyVars t = case getTyVar_maybe t of + Just x -> [x] + Nothing -> case tyConAppArgs_maybe t of + Just xs -> concatMap extractTyVars xs + Nothing -> concatMap extractTyVars $ snd $ splitAppTys t + +-- | Given a quantified type of the form: +-- +-- > forall a. (Num a, Ord a) => a -> a +-- +-- The first element of the returned tuple is the list of quantified variables, +-- and the seecond element is the unquantified type. +extractQuantifiers :: Type -> ([Var],Type) +extractQuantifiers t = case splitForAllTy_maybe t of + Nothing -> ([],t) + Just (a,b) -> (a:as,b') + where + (as,b') = extractQuantifiers b + +-- | Given unquantified types of the form: +-- +-- > (Num a, Ord a) => a -> a +-- +-- The first element of the returned tuple contains everything to the left of "=>"; +-- and the second element contains everything to the right. +extractContext :: Type -> ([Type],Type) +extractContext t = case splitTyConApp_maybe t of + Nothing -> ([],t) + Just (tycon,xs) -> if (occNameString $ nameOccName $ tyConName tycon)/="(->)" + || not hasCxt + then ([],t) + else (head xs:cxt',t') + where + (cxt',t') = extractContext $ head $ tail xs + + hasCxt = case classifyPredType $ head xs of + IrredPred _ -> False + _ -> True + +-- | given a function, get the type of the parameters +-- +-- FIXME: this should be deleted +extractParam :: Type -> Maybe Type +extractParam t = case splitTyConApp_maybe t of + Nothing -> Nothing + Just (tycon,xs) -> if (occNameString $ nameOccName $ tyConName tycon)/="(->)" + then Just t -- Nothing + else Just (head xs) + + +-- | Given a type of the form +-- +-- > A -> ... -> C +-- +-- returns C +getReturnType :: Type -> Type +getReturnType t = case splitForAllTys t of + (_,t') -> go t' + where + go t = case splitTyConApp_maybe t of + Just (tycon,[_,t']) -> if getString tycon=="(->)" + then go t' + else t + _ -> t + + +-------------------------------------------------------------------------------- +-- + +runTcM :: ModGuts -> TcM a -> CoreM a +runTcM guts tcm = do + env <- getHscEnv + dflags <- getDynFlags +#if __GLASGOW_HASKELL__ < 710 || (__GLASGOW_HASKELL__ == 710 && __GLASGOW_HASKELL_PATCHLEVEL1__ < 2) + (msgs, mr) <- liftIO $ initTc env HsSrcFile False (mg_module guts) tcm +#else + let realSrcSpan = mkRealSrcSpan + (mkRealSrcLoc (mkFastString "a") 0 1) + (mkRealSrcLoc (mkFastString "b") 2 3) + (msgs, mr) <- liftIO $ initTc env HsSrcFile False (mg_module guts) realSrcSpan tcm +#endif + let showMsgs (warns, errs) = showSDoc dflags $ vcat + $ text "Errors:" : pprErrMsgBag errs + ++ text "Warnings:" : pprErrMsgBag warns + maybe (fail $ showMsgs msgs) return mr + where + pprErrMsgBag = pprErrMsgBagWithLoc + +-------------------------------------------------------------------------------- +-- utils + +getString :: NamedThing a => a -> String +getString = occNameString . getOccName + +expr2str :: DynFlags -> Expr Var -> String +expr2str dflags (Var v) = {-"var_" ++-} var2str v +expr2str dflags e = "expr_" ++ (decorate $ showSDoc dflags (ppr e)) + where + decorate :: String -> String + decorate = map go + where + go x = if not (isAlphaNum x) + then '_' + else x + +lit2rational :: Literal -> Rational +lit2rational l = case l of + MachInt i -> toRational i + MachInt64 i -> toRational i + MachWord i -> toRational i + MachWord64 i -> toRational i + MachFloat r -> r + MachDouble r -> r + LitInteger i _ -> toRational i + +var2str :: Var -> String +var2str = occNameString . occName . varName + +maybeHead :: [a] -> Maybe a +maybeHead (a:_) = Just a +maybeHead _ = Nothing + +myshow :: DynFlags -> Expr Var -> String +myshow dflags = go 1 + where + go i (Var v) = "Var "++showSDoc dflags (ppr v) + ++"_"++showSDoc dflags (ppr $ getUnique v) + ++"::"++showSDoc dflags (ppr $ varType v) + go i (Lit (MachFloat l )) = "FloatLiteral " ++show (fromRational l :: Double) + go i (Lit (MachDouble l )) = "DoubleLiteral " ++show (fromRational l :: Double) + go i (Lit (MachInt l )) = "IntLiteral " ++show (fromIntegral l :: Double) + go i (Lit (MachInt64 l )) = "Int64Literal " ++show (fromIntegral l :: Double) + go i (Lit (MachWord l )) = "WordLiteral " ++show (fromIntegral l :: Double) + go i (Lit (MachWord64 l )) = "Word64Literal " ++show (fromIntegral l :: Double) + go i (Lit (LitInteger l t)) = "IntegerLiteral "++show (fromIntegral l :: Double)++ + "::"++showSDoc dflags (ppr t) + go i (Lit l) = "Lit" + go i (Type t) = "Type "++showSDoc dflags (ppr t) + go i (Tick a b) = "Tick (" ++ show a ++ ") ("++go (i+1) b++")" + go i (Coercion l) = "Coercion "++myCoercionShow dflags l + go i (Cast a b) + = "Cast \n" + ++white++"(" ++ go (i+1) a ++ ")\n" + ++white++"("++myshow dflags (Coercion b)++")\n" + ++drop 4 white + where + white=replicate (4*i) ' ' + go i (Let (NonRec a e) b) + = "Let "++getString a++"_"++showSDoc dflags (ppr $ getUnique a) + ++"::"++showSDoc dflags (ppr $ varType a)++"\n" + ++white++"("++go (i+1) e++")\n" + ++white++"("++go (i+1) b++")\n" + ++drop 4 white + where + white=replicate (4*i) ' ' + go i (Let _ _) = error "myshow: recursive let" + go i (Lam a b) + = "Lam "++getString a++"_"++showSDoc dflags (ppr $ getUnique a) + ++"::"++showSDoc dflags (ppr $ varType a) + ++"; coercion="++show (isCoVar a)++"\n" + ++white++"("++go (i+1) b++")\n" + ++drop 4 white + where + white=replicate (4*i) ' ' + go i (App a b) + = "App\n" + ++white++"(" ++ go (i+1) a ++ ")\n" + ++white++"("++go (i+1) b++")\n" + ++drop 4 white + where + white=replicate (4*i) ' ' + go i (Case a b c d) + = "Case\n" + ++white++"("++go (i+1) a++")\n" + ++white++"("++getString b++"_"++showSDoc dflags (ppr $ getUnique b) + ++"::"++showSDoc dflags (ppr $ varType b)++")\n" + ++white++"("++showSDoc dflags (ppr c)++"; "++show (fmap (myshow dflags . Var) $ getTyVar_maybe c)++")\n" + ++white++"["++concatMap altShow d++"]\n" + ++drop 4 white + where + white=replicate (4*i) ' ' + + altShow :: Alt Var -> String + altShow (con,xs,expr) = "("++con'++", "++xs'++", "++go (i+1) expr++")\n"++white + where + con' = case con of + DataAlt x -> showSDoc dflags (ppr x) + LitAlt x -> showSDoc dflags (ppr x) + DEFAULT -> "DEFAULT" + + xs' = show $ map (myshow dflags . Var) xs + +myCoercionShow :: DynFlags -> Coercion -> String +myCoercionShow dflags c = go c + where + go (Refl _ _ ) = "Refl" + go (TyConAppCo a b c ) = "TyConAppCo "++showSDoc dflags (ppr a)++" " + ++showSDoc dflags (ppr b)++" " + ++showSDoc dflags (ppr c) + go (AppCo _ _ ) = "AppCo" + go (ForAllCo _ _ ) = "ForAllCo" + go (CoVarCo v ) = "CoVarCo ("++myshow dflags (Var v)++")" + go (AxiomInstCo _ _ _ ) = "AxiomInstCo" + go (UnivCo _ _ _ _ ) = "UnivCo" + go (SymCo c' ) = "SymCo ("++myCoercionShow dflags c'++")" + go (TransCo _ _ ) = "TransCo" + go (AxiomRuleCo _ _ _ ) = "AxiomRuleCo" + go (NthCo _ _ ) = "NthCo" + go (LRCo _ _ ) = "LRCo" + go (InstCo _ _ ) = "InstCo" + go (SubCo c' ) = "SubCo ("++myCoercionShow dflags c'++")" + + +-- instance Show (Coercion) where +-- show _ = "Coercion" +-- +-- instance Show b => Show (Bind b) where +-- show _ = "Bind" +-- +-- instance Show (Tickish Id) where +-- show _ = "(Tickish Id)" +-- +-- instance Show Type where +-- show _ = "Type" +-- +-- instance Show AltCon where +-- show _ = "AltCon" +-- +-- instance Show Var where +-- show v = getString v + + diff --git a/src/Herbie/ForeignInterface.hs b/src/Herbie/ForeignInterface.hs new file mode 100644 --- /dev/null +++ b/src/Herbie/ForeignInterface.hs @@ -0,0 +1,247 @@ +{-# LANGUAGE OverloadedStrings #-} + +module Herbie.ForeignInterface + where + +import Control.Applicative +import Control.Exception +import Control.DeepSeq +import Data.List +import Data.String +import qualified Data.Text as T +import Database.SQLite.Simple +import Database.SQLite.Simple.FromRow +import Database.SQLite.Simple.FromField +import Database.SQLite.Simple.ToField +import GHC.Generics hiding (modName) +import System.Directory +import System.Process +import System.Timeout + +import Paths_HerbiePlugin +import Herbie.MathInfo +import Herbie.MathExpr + +import Prelude + +-- | Given a MathExpr, return a numerically stable version. +stabilizeMathExpr :: DbgInfo -> MathExpr -> IO (StabilizerResult MathExpr) +stabilizeMathExpr dbgInfo cmdin = do + let (cmdinLisp,varmap) = getCanonicalLispCmd $ haskellOpsToHerbieOps cmdin + res <- stabilizeLisp dbgInfo cmdinLisp + cmdout <- do + -- FIXME: + -- Due to a bug in Herbie, fromCanonicalLispCmd sometimes throws an exception. + ret <- try $ do + let ret = herbieOpsToHaskellOps $ fromCanonicalLispCmd (cmdout res,varmap) + deepseq ret $ return ret + case ret of + Left (SomeException e) -> do + putStrLn $ "WARNING in stabilizeMathExpr: "++show e + return cmdin + Right x -> return x + let res' = res + { cmdin = cmdin + , cmdout = cmdout + } +-- putStrLn $ "cmdin: "++cmdinLisp +-- putStrLn $ "cmdout: "++cmdout res +-- putStrLn $ "stabilizeLisp': "++mathExpr2lisp (fromCanonicalLispCmd (cmdout res,varmap)) + return res' + +-- | Given a Lisp command, return a numerically stable version. +-- It first checks if the command is in the global database; +-- if it's not, then it runs "execHerbie". +stabilizeLisp :: DbgInfo -> String -> IO (StabilizerResult String) +stabilizeLisp dbgInfo cmd = do + dbResult <- lookupDatabase cmd + ret <- case dbResult of + Just x -> do + return x + Nothing -> do + putStrLn " Not found in database. Running Herbie..." + res <- execHerbie cmd + insertDatabase res + return res + insertDatabaseDbgInfo dbgInfo ret + + -- FIXME: + -- Herbie has a bug where it sometimes outputs a less numerically stable version. + -- So we need to check to make sure we return the more stable output. + return $ if errin ret > errout ret + then ret + else ret { errout = errin ret, cmdout = cmdin ret } + +-- | Run the `herbie` command and return the result +execHerbie :: String -> IO (StabilizerResult String) +execHerbie lisp = do + + -- build the command string we will pass to Herbie + let varstr = "("++(intercalate " " $ lisp2vars lisp)++")" + stdin = "(herbie-test "++varstr++" \"cmd\" "++lisp++") \n" + + -- Herbie can take a long time to run. + -- Here we limit it to 2 minutes. + -- + -- FIXME: + -- This should be a parameter the user can pass to the plugin + ret <- timeout 120000000 $ do + + -- launch Herbie with a fixed seed to ensure reproducible builds + (_,stdout,stderr) <- readProcessWithExitCode + "herbie-exec" + [ "-r", "#(1461197085 2376054483 1553562171 1611329376 2497620867 2308122621)" ] + stdin + + -- try to parse Herbie's output; + -- if we can't parse it, that means Herbie had an error and we should abort gracefully + ret <- try $ do + let (line1:line2:line3:_) = lines stdout + let ret = StabilizerResult + { errin + = read + $ drop 1 + $ dropWhile (/=':') + $ line1 + , errout + = read + $ drop 1 + $ dropWhile (/=':') + $ line2 + , cmdin + = lisp + , cmdout + = (!!2) + $ groupByParens + $ init + $ tail + $ line3 + } + deepseq ret $ return ret + + case ret of + Left (SomeException e) -> do + putStrLn $ "WARNING in execHerbie: "++show e + putStrLn $ "WARNING in execHerbie: stdin="++stdin + putStrLn $ "WARNING in execHerbie: stdout="++stdout + return $ StabilizerResult + { errin = 0/0 + , errout = 0/0 + , cmdin = lisp + , cmdout = lisp + } + Right x -> return x + + case ret of + Just x -> return x + Nothing -> do + putStrLn $ "WARNING: Call to Herbie timed out after 2 minutes." + return $ StabilizerResult + { errin = 0/0 + , errout = 0/0 + , cmdin = lisp + , cmdout = lisp + } + + +-- | The result of running Herbie +data StabilizerResult a = StabilizerResult + { cmdin :: !a + , cmdout :: !a + , errin :: !Double + , errout :: !Double + } + deriving (Show,Generic,NFData) + +instance FromField a => FromRow (StabilizerResult a) where + fromRow = StabilizerResult <$> field <*> field <*> field <*> field + +instance ToField a => ToRow (StabilizerResult a) where + toRow (StabilizerResult cmdin cmdout errin errout) = toRow (cmdin, cmdout, errin, errout) + +-- | Returns a connection to the sqlite3 database +mkConn = do + path <- getDataFileName "Herbie.db" + open path + +-- | Check the database to see if we already know the answer for running Herbie +-- +-- FIXME: +-- When Herbie times out, NULL gets inserted into the database for errin and errout. +-- The Sqlite3 bindings don't support putting NULL into Double's as NaNs, +-- so the query below raises an exception. +-- This isn't so bad, except a nasty error message gets printed, +-- and the plugin attempts to run Herbie again (wasting a lot of time). +lookupDatabase :: String -> IO (Maybe (StabilizerResult String)) +lookupDatabase cmdin = do + ret <- try $ do + conn <- mkConn + res <- queryNamed + conn + "SELECT cmdin,cmdout,errin,errout from StabilizerResults where cmdin = :cmdin" + [":cmdin" := cmdin] + :: IO [StabilizerResult String] + close conn + return $ case res of + [x] -> Just x + [] -> Nothing + case ret of + Left (SomeException e) -> do + putStrLn $ "WARNING in lookupDatabase: "++show e + return Nothing + Right x -> return x + +-- | Inserts a "StabilizerResult" into the global database of commands +insertDatabase :: StabilizerResult String -> IO () +insertDatabase res = do + ret <- try $ do + conn <- mkConn + execute_ conn $ fromString $ + "CREATE TABLE IF NOT EXISTS StabilizerResults " + ++"( id INTEGER PRIMARY KEY" + ++", cmdin TEXT UNIQUE NOT NULL" + ++", cmdout TEXT NOT NULL" + ++", errin DOUBLE " + ++", errout DOUBLE " + ++")" + execute_ conn "CREATE INDEX IF NOT EXISTS StabilizerResultsIndex ON StabilizerResults(cmdin)" + execute conn "INSERT INTO StabilizerResults (cmdin,cmdout,errin,errout) VALUES (?,?,?,?)" res + close conn + case ret of + Left (SomeException e) -> putStrLn $ "WARNING in insertDatabase: "++show e + Right _ -> return () + return () + + +-- | This information gets stored in a separate db table for debugging purposes +data DbgInfo = DbgInfo + { dbgComments :: String + , modName :: String + , functionName :: String + , functionType :: String + } + +insertDatabaseDbgInfo :: DbgInfo -> StabilizerResult String -> IO () +insertDatabaseDbgInfo dbgInfo res = do + ret <- try $ do + conn <- mkConn + execute_ conn $ fromString $ + "CREATE TABLE IF NOT EXISTS DbgInfo " + ++"( id INTEGER PRIMARY KEY" + ++", resid INTEGER NOT NULL" + ++", dbgComments TEXT" + ++", modName TEXT" + ++", functionName TEXT" + ++", functionType TEXT" + ++")" + res <- queryNamed + conn + "SELECT id,cmdout from StabilizerResults where cmdin = :cmdin" + [":cmdin" := (cmdin res)] + :: IO [(Int,String)] + execute conn "INSERT INTO DbgInfo (resid,dbgComments,modName,functionName,functionType) VALUES (?,?,?,?,?)" (fst $ head res,dbgComments dbgInfo,modName dbgInfo,functionName dbgInfo,functionType dbgInfo) + close conn + case ret of + Left (SomeException e) -> putStrLn $ "WARNING in insertDatabaseDbgInfo: "++show e + Right _ -> return () + return () diff --git a/src/Herbie/MathExpr.hs b/src/Herbie/MathExpr.hs new file mode 100644 --- /dev/null +++ b/src/Herbie/MathExpr.hs @@ -0,0 +1,249 @@ +{-# LANGUAGE DeriveAnyClass,DeriveGeneric #-} +module Herbie.MathExpr + where + +import Control.DeepSeq +import Data.List +import Data.Maybe +import GHC.Generics + +import Debug.Trace +import Prelude +ifThenElse True t f = t +ifThenElse False t f = f + +------------------------------------------------------------------------------- +-- constants that define valid math expressions + +monOpList = + [ "cos" + , "sin" + , "tan" + , "acos" + , "asin" + , "atan" + , "cosh" + , "sinh" + , "tanh" + , "exp" + , "log" + , "sqrt" + , "abs" + , "size" + ] + +binOpList = [ "^", "**", "^^", "/", "-", "expt" ] ++ commutativeOpList +commutativeOpList = [ "*", "+"] -- , "max", "min" ] + +-------------------------------------------------------------------------------- + +-- | Stores the AST for a math expression in a generic form that requires no knowledge of Core syntax. +data MathExpr + = EBinOp String MathExpr MathExpr + | EMonOp String MathExpr + | EIf MathExpr MathExpr MathExpr + | ELit Rational + | ELeaf String + deriving (Show,Eq,Generic,NFData) + +instance Ord MathExpr where + compare (ELeaf _) (ELeaf _) = EQ + compare (ELeaf _) _ = LT + + compare (ELit r1) (ELit r2) = compare r1 r2 + compare (ELit _ ) (ELeaf _) = GT + compare (ELit _ ) _ = LT + + compare (EMonOp op1 e1) (EMonOp op2 e2) = case compare op1 op2 of + EQ -> compare e1 e2 + x -> x + compare (EMonOp _ _) (ELeaf _) = GT + compare (EMonOp _ _) (ELit _) = GT + compare (EMonOp _ _) _ = LT + + compare (EBinOp op1 e1a e1b) (EBinOp op2 e2a e2b) = case compare op1 op2 of + EQ -> case compare e1a e2a of + EQ -> compare e1b e2b + _ -> EQ + _ -> EQ + compare (EBinOp _ _ _) _ = LT + +-- | Converts all Haskell operators in the MathExpr into Herbie operators +haskellOpsToHerbieOps :: MathExpr -> MathExpr +haskellOpsToHerbieOps = go + where + go (EBinOp op e1 e2) = EBinOp op' (go e1) (go e2) + where + op' = case op of + "**" -> "expt" + "^^" -> "expt" + "^" -> "expt" + x -> x + + go (EMonOp op e1) = EMonOp op' (go e1) + where + op' = case op of + "size" -> "abs" + x -> x + + go (EIf cond e1 e2) = EIf (go cond) (go e1) (go e2) + go x = x + +-- | Converts all Herbie operators in the MathExpr into Haskell operators +herbieOpsToHaskellOps :: MathExpr -> MathExpr +herbieOpsToHaskellOps = go + where + go (EBinOp op e1 e2) = EBinOp op' (go e1) (go e2) + where + op' = case op of + "^" -> "**" + "expt" -> "**" + x -> x + + go (EMonOp "sqr" e1) = EBinOp "*" (go e1) (go e1) + go (EMonOp op e1) = EMonOp op' (go e1) + where + op' = case op of + "-" -> "negate" + "abs" -> "size" + x -> x + + go (EIf cond e1 e2) = EIf (go cond) (go e1) (go e2) + go x = x + +-- | Replace all the variables in the MathExpr with canonical names (x0,x1,x2...) +-- and reorder commutative binary operations. +-- This lets us more easily compare MathExpr's based on their structure. +-- The returned map lets us convert the canoncial MathExpr back into the original. +toCanonicalMathExpr :: MathExpr -> (MathExpr,[(String,String)]) +toCanonicalMathExpr e = go [] e + where + go :: [(String,String)] -> MathExpr -> (MathExpr,[(String,String)]) + go acc (EBinOp op e1 e2) = (EBinOp op e1' e2',acc2') + where + (e1_,e2_) = if op `elem` commutativeOpList + then (min e1 e2,max e1 e2) + else (e1,e2) + + (e1',acc1') = go acc e1_ + (e2',acc2') = go acc1' e2_ + + go acc (EMonOp op e1) = (EMonOp op e1', acc1') + where + (e1',acc1') = go acc e1 + go acc (ELit r) = (ELit r,acc) + go acc (ELeaf str) = (ELeaf str',acc') + where + (acc',str') = case lookup str acc of + Nothing -> ((str,"herbie"++show (length acc)):acc, "herbie"++show (length acc)) + Just x -> (acc,x) + +-- | Convert a canonical MathExpr into its original form. +-- +-- FIXME: +-- A bug in Herbie causes it to sometimes output infinities, +-- which break this function and cause it to error. +fromCanonicalMathExpr :: (MathExpr,[(String,String)]) -> MathExpr +fromCanonicalMathExpr (e,xs) = go e + where + xs' = map (\(a,b) -> (b,a)) xs + + go (EMonOp op e1) = EMonOp op (go e1) + go (EBinOp op e1 e2) = EBinOp op (go e1) (go e2) + go (EIf (EBinOp "<" _ (ELeaf "-inf.0")) e1 e2) = go e2 -- FIXME: added due to bug above + go (EIf cond e1 e2) = EIf (go cond) (go e1) (go e2) + go (ELit r) = ELit r + go (ELeaf str) = case lookup str xs' of + Just x -> ELeaf x + Nothing -> error $ "fromCanonicalMathExpr: str="++str++"; xs="++show xs' + +-- | Calculates the maximum depth of the AST. +mathExprDepth :: MathExpr -> Int +mathExprDepth (EBinOp _ e1 e2) = 1+max (mathExprDepth e1) (mathExprDepth e2) +mathExprDepth (EMonOp _ e1 ) = 1+mathExprDepth e1 +mathExprDepth _ = 0 + +-------------------------------------------------------------------------------- +-- functions for manipulating math expressions in lisp form + +getCanonicalLispCmd :: MathExpr -> (String,[(String,String)]) +getCanonicalLispCmd me = (mathExpr2lisp me',varmap) + where + (me',varmap) = toCanonicalMathExpr me + +fromCanonicalLispCmd :: (String,[(String,String)]) -> MathExpr +fromCanonicalLispCmd (lisp,varmap) = fromCanonicalMathExpr (lisp2mathExpr lisp,varmap) + +-- | Converts MathExpr into a lisp command suitable for passing to Herbie +mathExpr2lisp :: MathExpr -> String +mathExpr2lisp = go + where + go (EBinOp op a1 a2) = "("++op++" "++go a1++" "++go a2++")" + go (EMonOp op a) = "("++op++" "++go a++")" + go (EIf cond e1 e2) = "(if "++go cond++" "++go e1++" "++go e2++")" + go (ELeaf e) = e + go (ELit r) = if (toRational (floor r::Integer) == r) + then show (floor r :: Integer) + else show (fromRational r :: Double) + +-- | Converts a lisp command into a MathExpr +lisp2mathExpr :: String -> MathExpr +lisp2mathExpr ('-':xs) = EMonOp "negate" (lisp2mathExpr xs) +lisp2mathExpr ('(':xs) = if length xs > 1 && last xs==')' + then case groupByParens $ init xs of + [op,e1] -> EMonOp op (lisp2mathExpr e1) + [op,e1,e2] -> EBinOp op (lisp2mathExpr e1) (lisp2mathExpr e2) + ["if",cond,e1,e2] -> EIf (lisp2mathExpr cond) (lisp2mathExpr e1) (lisp2mathExpr e2) + _ -> error $ "lisp2mathExpr: "++xs + else error $ "lisp2mathExpr: malformed input '("++xs++"'" +lisp2mathExpr xs = case readMaybe xs :: Maybe Double of + Just x -> ELit $ toRational x + Nothing -> ELeaf xs + +-- | Extracts all the variables from the lisp commands with no duplicates. +lisp2vars :: String -> [String] +lisp2vars = nub . lisp2varsNoNub + +-- | Extracts all the variables from the lisp commands. +-- Each variable occurs once in the output for each time it occurs in the input. +lisp2varsNoNub :: String -> [String] +lisp2varsNoNub lisp + = sort + $ filter (\x -> x/="(" + && x/=")" + && not (x `elem` binOpList) + && not (x `elem` monOpList) + && not (head x `elem` ("1234567890"::String)) + ) + $ tokenize lisp :: [String] + where + -- We just need to add spaces around the parens before calling "words" + tokenize :: String -> [String] + tokenize = words . concat . map go + where + go '(' = " ( " + go ')' = " ) " + go x = [x] + +lispHasRepeatVars :: String -> Bool +lispHasRepeatVars lisp = length (lisp2vars lisp) /= length (lisp2varsNoNub lisp) + +------------------------------------------------------------------------------- +-- utilities + +readMaybe :: Read a => String -> Maybe a +readMaybe = fmap fst . listToMaybe . reads + +-- | Given an expression, break it into tokens only outside parentheses +groupByParens :: String -> [String] +groupByParens str = go 0 str [] [] + where + go 0 (' ':xs) [] ret = go 0 xs [] ret + go 0 (' ':xs) acc ret = go 0 xs [] (ret++[acc]) + go 0 (')':xs) acc ret = go 0 xs [] (ret++[acc]) + go i (')':xs) acc ret = go (i-1) xs (acc++")") ret + go i ('(':xs) acc ret = go (i+1) xs (acc++"(") ret + go i (x :xs) acc ret = go i xs (acc++[x]) ret + go _ [] acc ret = ret++[acc] + + diff --git a/src/Herbie/MathInfo.hs b/src/Herbie/MathInfo.hs new file mode 100644 --- /dev/null +++ b/src/Herbie/MathInfo.hs @@ -0,0 +1,278 @@ +{-# LANGUAGE FlexibleInstances,FlexibleContexts,MultiWayIf,CPP #-} +module Herbie.MathInfo + where + +import Class +import DsBinds +import DsMonad +import ErrUtils +import GhcPlugins hiding (trace) +import Unique +import MkId +import PrelNames +import UniqSupply +import TcRnMonad +import TcSimplify +import Type + +import Control.Monad +import Control.Monad.Except +import Control.Monad.Trans +import Data.Char +import Data.List +import Data.Maybe +import Data.Ratio + +import Herbie.CoreManip +import Herbie.MathExpr + +import Prelude +import Show + +-- import Debug.Trace hiding (traceM) +trace a b = b +traceM a = return () + +-------------------------------------------------------------------------------- + +-- | The fields of this type correspond to the sections of a function type. +-- +-- Must satisfy the invariant that every class in "getCxt" has an associated dictionary in "getDicts". +data ParamType = ParamType + { getQuantifier :: [Var] + , getCxt :: [Type] + , getDicts :: [CoreExpr] + , getParam :: Type + } + +-- | This type is a simplified version of the CoreExpr type. +-- It only supports math expressions. +-- We first convert a CoreExpr into a MathInfo, +-- perform all the manipulation on the MathExpr within the MathInfo, +-- then use the information in MathInfo to convert the MathExpr back into a CoreExpr. +data MathInfo = MathInfo + { getMathExpr :: MathExpr + , getParamType :: ParamType + , getExprs :: [(String,Expr Var)] + -- ^ the fst value is the unique name assigned to non-mathematical expressions + -- the snd value is the expression + } + +-- | Pretty print a math expression +pprMathInfo :: MathInfo -> String +pprMathInfo mathInfo = go 1 False $ getMathExpr mathInfo + where + isLitOrLeaf :: MathExpr -> Bool + isLitOrLeaf (ELit _ ) = True + isLitOrLeaf (ELeaf _) = True + isLitOrLeaf _ = False + + go :: Int -> Bool -> MathExpr -> String + go i b e = if b && not (isLitOrLeaf e) + then "("++str++")" + else str + where + str = case e of + EMonOp op e1 -> op++" "++(go i True e1) + + EBinOp op e1 e2 -> go i parens1 e1++" "++op++" "++go i parens2 e2 + where + parens1 = case e1 of + (EBinOp op' _ _) -> op/=op' + _ -> True + + parens2 = case e2 of + (EBinOp op' _ _) -> op/=op' + _ -> True + + ELit l -> if toRational (floor l) == l + then if length (show (floor l :: Integer)) < 10 + then show (floor l :: Integer) + else show (fromRational l :: Double) + else show (fromRational l :: Double) + + ELeaf l -> case lookup l $ getExprs mathInfo of + Just (Var _) -> l + _ -> "???" + + EIf cond e1 e2 -> "if "++go i False cond++"\n" + ++white++"then "++go (i+1) False e1++"\n" + ++white++"else "++go (i+1) False e2 + where + white = replicate (4*i) ' ' + +-- If the given expression is a math expression, +-- returns the type of the variable that the math expression operates on. +varTypeIfValidExpr :: CoreExpr -> Maybe Type +varTypeIfValidExpr e = case e of + + -- might be a binary math operation + (App (App (App (App (Var v) (Type t)) _) _) _) -> if var2str v `elem` binOpList + then if isValidType t + then Just t + else Nothing + else Nothing + + -- might be a unary math operation + (App (App (App (Var v) (Type t)) _) _) -> if var2str v `elem` monOpList + then if isValidType t + then Just t + else Nothing + else Nothing + + -- first function is anything else means that we're not a math expression + _ -> Nothing + + where + isValidType :: Type -> Bool + isValidType t = isTyVarTy t || case splitTyConApp_maybe t of + Nothing -> True + Just (tyCon,_) -> tyCon == floatTyCon || tyCon == doubleTyCon + +-- | Converts a CoreExpr into a MathInfo +mkMathInfo :: DynFlags -> [Var] -> Type -> Expr Var -> Maybe MathInfo +mkMathInfo dflags dicts bndType e = case varTypeIfValidExpr e of + Nothing -> Nothing + Just t -> if mathExprDepth getMathExpr>1 && lispHasRepeatVars (mathExpr2lisp getMathExpr) + then Just $ MathInfo + getMathExpr + ( ParamType + { getQuantifier = quantifier + , getCxt = cxt + , getDicts = map Var dicts + , getParam = t + } + ) exprs + else Nothing + + where + (getMathExpr,exprs) = go e [] + + -- this should never return Nothing if validExpr is not Nothing + (quantifier,unquantified) = extractQuantifiers bndType + (cxt,uncxt) = extractContext unquantified + + -- recursively converts the `Expr Var` into a MathExpr and a dictionary + go :: Expr Var + -> [(String,Expr Var)] + -> (MathExpr + ,[(String,Expr Var)] + ) + + -- we need to special case the $ operator for when MathExpr is run before any rewrite rules + go e@(App (App (App (App (Var v) (Type _)) (Type _)) a1) a2) exprs + = if var2str v == "$" + then go (App a1 a2) exprs + else (ELeaf $ expr2str dflags e,[(expr2str dflags e,e)]) + + -- polymorphic literals created via fromInteger + go e@(App (App (App (Var v) (Type _)) dict) (Lit l)) exprs + = (ELit $ lit2rational l, exprs) + + -- polymorphic literals created via fromRational + go e@(App (App (App (Var v) (Type _)) dict) + (App (App (App (Var _) (Type _)) (Lit l1)) (Lit l2))) exprs + = (ELit $ lit2rational l1 / lit2rational l2, exprs) + + -- non-polymorphic literals + go e@(App (Var _) (Lit l)) exprs + = (ELit $ lit2rational l, exprs) + + -- binary operators + go e@(App (App (App (App (Var v) (Type _)) dict) a1) a2) exprs + = if var2str v `elem` binOpList + then let (a1',exprs1) = go a1 [] + (a2',exprs2) = go a2 [] + in ( EBinOp (var2str v) a1' a2' + , exprs++exprs1++exprs2 + ) + else (ELeaf $ expr2str dflags e,[(expr2str dflags e,e)]) + + -- unary operators + go e@(App (App (App (Var v) (Type _)) dict) a) exprs + = if var2str v `elem` monOpList + then let (a',exprs') = go a [] + in ( EMonOp (var2str v) a' + , exprs++exprs' + ) + else (ELeaf $ expr2str dflags e,(expr2str dflags e,e):exprs) + + -- everything else + go e exprs = (ELeaf $ expr2str dflags e,[(expr2str dflags e,e)]) + +-- | Converts a MathInfo back into a CoreExpr +mathInfo2expr :: ModGuts -> MathInfo -> ExceptT String CoreM CoreExpr +mathInfo2expr guts herbie = go (getMathExpr herbie) + where + pt = getParamType herbie + + -- binary operators + go (EBinOp opstr a1 a2) = do + a1' <- go a1 + a2' <- go a2 + f <- getDecoratedFunction guts opstr (getParam pt) (getDicts pt) + return $ App (App f a1') a2' + + -- unary operators + go (EMonOp opstr a) = do + a' <- go a + f <- getDecoratedFunction guts opstr (getParam pt) (getDicts pt) + castToType + (getDicts pt) + (getParam pt) + $ App f a' + + -- if statements + go (EIf cond a1 a2) = do + cond' <- go cond >>= castToType (getDicts pt) boolTy + a1' <- go a1 + a2' <- go a2 + + wildUniq <- getUniqueM + let wildName = mkSystemName wildUniq (mkVarOcc $ "wild") + wildVar = mkLocalVar VanillaId wildName boolTy vanillaIdInfo + + return $ Case + cond' + wildVar + (getParam pt) + [ (DataAlt falseDataCon, [], a2') + , (DataAlt trueDataCon, [], a1') + ] + + -- leaf is a numeric literal + go (ELit r) = do + fromRationalExpr <- getDecoratedFunction guts "fromRational" (getParam pt) (getDicts pt) + + integerTyCon <- lookupTyCon integerTyConName + let integerTy = mkTyConTy integerTyCon + + ratioTyCon <- lookupTyCon ratioTyConName + tmpUniq <- getUniqueM + let tmpName = mkSystemName tmpUniq (mkVarOcc $ "a") + tmpVar = mkTyVar tmpName liftedTypeKind + tmpVarT = mkTyVarTy tmpVar + ratioConTy = mkForAllTy tmpVar $ mkFunTys [tmpVarT,tmpVarT] $ mkAppTy (mkTyConTy ratioTyCon) tmpVarT + ratioConVar = mkGlobalVar VanillaId ratioDataConName ratioConTy vanillaIdInfo + + return $ App + fromRationalExpr + (App + (App + (App + (Var ratioConVar ) + (Type integerTy) + ) + (Lit $ LitInteger (numerator r) integerTy) + ) + (Lit $ LitInteger (denominator r) integerTy) + ) + + -- leaf is any other expression + go (ELeaf str) = do + dflags <- getDynFlags + return $ case lookup str (getExprs herbie) of + Just x -> x + Nothing -> error $ "mathInfo2expr: var " ++ str ++ " not in scope" + ++"; in scope vars="++show (nub $ map fst $ getExprs herbie) + diff --git a/src/Show.hs b/src/Show.hs new file mode 100644 --- /dev/null +++ b/src/Show.hs @@ -0,0 +1,160 @@ +{-# LANGUAGE FlexibleInstances, MultiWayIf, StandaloneDeriving, + TypeSynonymInstances #-} + +{-# OPTIONS_GHC -fno-warn-orphans #-} + +-- | We define lots of orphan Show instances here, for debugging and learning +-- purposes. +-- +-- Most of the time while trying to figure out when a constructor is used or how +-- is a term compiled, it's easiest to just create an example and run the plugin +-- on it. +-- +-- Without Show instances though, we can't easily inspect compiled outputs. +-- Outputable generated strings hide lots of details(especially constructors), +-- but we still export a `showOutputable` here, for similar reasons. +-- +module Show where + +import Data.IORef +import Data.List (intercalate) +import System.IO.Unsafe (unsafePerformIO) + +import Class +import CostCentre +import ForeignCall +import Demand +import GhcPlugins +import IdInfo +import PrimOp +import TypeRep + +import Prelude + +-------------------------------------------------------------------------------- + +dbg :: Outputable a => a -> String +dbg a = showSDoc dynFlags (ppr a) + +{-# NOINLINE dynFlags_ref #-} +dynFlags_ref :: IORef DynFlags +dynFlags_ref = unsafePerformIO (newIORef undefined) + +{-# NOINLINE dynFlags #-} +dynFlags :: DynFlags +dynFlags = unsafePerformIO (readIORef dynFlags_ref) + +showOutputable :: Outputable a => a -> String +showOutputable = showSDoc dynFlags . ppr + +-------------------------------------------------------------------------------- +-- Orphan Show instances + +deriving instance Show a => Show (Expr a) +deriving instance Show Type +deriving instance Show Literal +deriving instance Show a => Show (Tickish a) +deriving instance Show a => Show (Bind a) +deriving instance Show AltCon +deriving instance Show TyLit +deriving instance Show FunctionOrData +deriving instance Show Module +deriving instance Show CostCentre +deriving instance Show Role +deriving instance Show LeftOrRight +deriving instance Show IsCafCC + +instance Show Class where + show _ = "" + +deriving instance Show IdDetails +deriving instance Show PrimOp +deriving instance Show ForeignCall +deriving instance Show TickBoxOp +deriving instance Show PrimOpVecCat +deriving instance Show CCallSpec +deriving instance Show CCallTarget +deriving instance Show CCallConv +deriving instance Show SpecInfo +deriving instance Show OccInfo +deriving instance Show InlinePragma +deriving instance Show OneShotInfo +deriving instance Show CafInfo +deriving instance Show Unfolding +deriving instance Show UnfoldingSource +deriving instance Show UnfoldingGuidance +deriving instance Show Activation +deriving instance Show CoreRule +-- deriving instance Show IsOrphan +deriving instance Show StrictSig +deriving instance Show DmdType + +instance Show RuleFun where + show _ = "" + +instance Show (UniqFM a) where + show _ = "" + +instance Show IdInfo where + show info = + "Info{" ++ intercalate "," [show arityInfo_, show specInfo_, show unfoldingInfo_, + show cafInfo_, show oneShotInfo_, show inlinePragInfo_, + show occInfo_, show strictnessInfo_, show demandInfo_, + show callArityInfo_] ++ "}" + where + arityInfo_ = arityInfo info + specInfo_ = specInfo info + unfoldingInfo_ = unfoldingInfo info + cafInfo_ = cafInfo info + oneShotInfo_ = oneShotInfo info + inlinePragInfo_ = inlinePragInfo info + occInfo_ = occInfo info + strictnessInfo_ = strictnessInfo info + demandInfo_ = demandInfo info + callArityInfo_ = callArityInfo info + +instance Show Var where + show v = + if | isId v -> + let details = idDetails v + info = idInfo v + in "Id{" ++ intercalate "," [show name, show uniq, show ty, show details, show info] ++ "}" + | isTyVar v -> "TyVar{" ++ show name ++ "}" + | otherwise -> "TcTyVar{" ++ show name ++ "}" + where + name = varName v + uniq = varUnique v + ty = varType v + +instance Show DataCon where + show = show . dataConName + +instance Show TyCon where + show = show . tyConName + +instance Show ModuleName where + show = show . moduleNameString + +instance Show PackageKey where + show = show . packageKeyString + +instance Show Name where + show = showOutputable . nameOccName + +-- deriving instance Show Name +instance Show OccName where + show = showOutputable + +instance Show Coercion where + show _ = "" + + +-- Instance for non-terms related stuff. + +deriving instance Show CoreToDo +deriving instance Show SimplifierMode +deriving instance Show CompilerPhase +deriving instance Show FloatOutSwitches + +instance Show PluginPass where + show _ = "PluginPass" diff --git a/test/Tests.hs b/test/Tests.hs new file mode 100644 --- /dev/null +++ b/test/Tests.hs @@ -0,0 +1,281 @@ +{-# LANGUAGE GADTs,RebindableSyntax,CPP,FlexibleContexts,FlexibleInstances,ConstraintKinds #-} +{- + - The idea of this test suite is that it should be compiled + - with the -fplugin=Herbie and -dcore-lint flags. + - Then we check to make sure GHC didn't throw any errors during + - the core type checking process. + -} +module Main + where + +import SubHask + +-------------------------------------------------------------------------------- + +-- This section tests that Herbie gets run on the correct types. +-- Herbie should be run on all the functions below. + +#define f1(x) (sqrt ((x)+1) - sqrt (x)) + +herbie1 :: Real a => a -> a +herbie1 x = f1(x) + +herbie2 :: Real a => a -> a -> a -> a -> a +herbie2 a b c d = f1(a)+f1(b)+f1(c)+f1(d) + +herbie3 :: Float -> Float +herbie3 x = f1(x) + +herbie4 :: String -> String +herbie4 str = show $ f1(x1) + where + x1 = fromIntegral (length str) :: Float + +herbie5 :: (Show a, Real a) => String -> a -> String +herbie5 str x1 = show $ f1(x1) + +herbie6 :: (Show a, Real a) => a -> String -> String +herbie6 x1 str = show $ f1(x1) + +herbie7 :: Semigroup a => a -> a +herbie7 x1 = x1+x1+x1+x1+x1 + +herbie8 :: Float -> Float +herbie8 x1 = case x1 of + 1.0 -> f1(x1) + 2.0 -> x1 + +herbie9 :: Float -> Float +herbie9 x1 = go 4 x1 + where + go :: Float -> Float -> Float + go 0 b = b + go a b = go (a-1) (sqrt (b-1)) + +-- Herbie should not get run on any of the functions in this section. + +#define f2(a,b) a+b*(a+b*a)+a*b + +noherbie1 :: String -> String +noherbie1 x = x++"hello world" + +noherbie2 :: Rational -> Rational -> Rational +noherbie2 a b = f2(a,b) + +noherbie3 :: Int -> Int -> Int +noherbie3 a b = f2(a,b) + +noherbie4 :: x -> Int -> Int -> Int +noherbie4 x a b = f2(a,b) + +-------------------------------------------------------------------------------- + +-- Herbie shouldn't process these because the expression size is too small. +-- We're unlikely to get any benefit, and it might take a long time. + +toosmall1 :: Float -> Float +toosmall1 a = a+a + +toosmall2 :: Float -> Float -> Float -> Float +toosmall2 a b c = a+b*c + +-- These are big enough and should get processed + +bigenough1 :: Float -> Float +bigenough1 a = a+a*a + +bigenough2 :: Float -> Float -> Float -> Float +bigenough2 a b c = a+b*(c+a) + +bigenough3 :: Float -> Float -> Float -> Float +bigenough3 a b c = f1(c) + +-------------------------------------------------------------------------------- + +-- This section contains lots of examples of expressions that the Herbie plugin can parse +-- and find improved versions. + +example1 x1 x2 = sqrt (x1*x1 + x2*x2) + +example2 x = exp(log(x)+8) + +example3 x = sqrt(x*x +1) -1 + +example4 x = exp(x)-1 + +example5 x = log(1+x) + +example6 x y = sqrt(x+ y) - sqrt(y) + +example7 k r a = k*(r-a)^3 + +example8 k r a = k*(r-a)^2 + +example9 x y = sin(x - y) + +example10 p1x p2x p1y p2y = sqrt((p1x - p2x) * (p1x - p2x) + (p1y - p2y) * (p1y - p2y)) + +example11 x = sin(x)-x + +example12 x = 1-cos(x) + +example13 x1 x2 = sqrt((x1 - x2) * (x1 - x2)) + +example14 x y z = sqrt(x*x + y*y + z*z) + +example15 x y z c = sqrt(x*x + y*y + z*z)/c + +example16 tdx dx tdy dy = (tdx * dx + tdy * dy) / (dx * dx + dy * dy) + +example17 tdx dx tdy dy sl2 = (tdx * dx + tdy * dy) / sl2 + +example18 x = (x + 0.1)-x + +example19 x = log(x) - sin(x+1) + +example20 a b = exp(1+log(a) + log(b)) + +example21 x = (1+sqrt(x-1))/(x-1)^2 + +example22 x = (1+sqrt(x))/(x-1)^2 + +example23 a b c d e f = a+b+(((d-c)*(d-c))*e*f/(e+f)) + +example24 q = sqrt(q*(q-1)) + +example25 a = sqrt(a^2-1) + +example26 a b c d = ((a*b)+(c*d))/(a+c) + +example27 x = sqrt(x^2) + +example28 x y = sqrt(x) * y * y + +example29 x y z = sqrt(x*x+y*y+z*z) + +example30 x y = 1.75 * x * y*y + sqrt(x/y) + +example31 x = exp(3*log(x)+2) + +example32 x = exp(2*log(x)) + +example33 x = sqrt(1/x + 1) - sqrt(1/x) + +example34 left i right count = left + i * ((left - right) / count) + +example35 left right count = left + count * ((left - right) / count) + +example36 x y = sqrt(x*x) - sqrt(y*y) + +example37 x = log(x+1)-log(x) + +example38 x = log(x+1)^x + +example39 minval minstep val = (minval/minstep + val) * minstep + +example40 x = x*x*cos(x/2 - sqrt(x)) + +example41 x = sqrt(4+x^2+x) + +example42 x y z = x / sqrt(x*x + y*y + z*z) + +example43 x = sin(sqrt(x+1)) + +example44 x = sqrt(x-2)-sqrt(x*x-3) + +example45 x = (sin(x) - tan(x)) / x + +example46 x y = 1 / sqrt(x^2 - y^2) + +example47 x1 x2 = sqrt((x1 - x2)^2) + +example48 x = x - sin(x) + +example49 x = sqrt(x + 1) - 1 + x + +example50 a b c = (a*a - c*c)/b + +example51 x y = sin(x+y)-cos(x+y) + +example52 x = (x + 1)^2 - 1 + +example53 x = sqrt(1+x) - sqrt(x) + +example54 x = sqrt(x + 1) / (x*x) + +example55 x = sqrt(x^2 / 3) + +example56 a b = 100*(a-b)/a + +example57 x = abs(x^3)-x^3 + +example58 x = log(x) - log(x+1) + +example59 x = 1/x - 1/(x+1) + +example60 a b c = -b + sqrt(b*b-4*a*c)/(2*a) + +example61 a c an cn = log(exp(a)*an + exp(c)*cn) - log(an+cn) + +example62 x = sqrt(sin(x)) - sqrt(x) + +example63 x = log(1+x) + +example64 a b = a * b / (1 - b + a * b) + +example65 a b = b*sqrt(a * a + 1.0) + +example65' a = sqrt(a * a + 1.0) + +example66 x y = x * y * x*pi/y + +example67 x = sqrt(x + 1) - sqrt(x - 1) + +example68 x = cos(x + 1) * x^2 + +example69 a b = b*(a/b - log(1 + a/b)) + +example70 a b = b*(a/b - 1 - log(a/b)) + +example71 x = (6/(x^99))*(x^101) + +example72 x = (1/(x^99))*(x^101) + +example73 x = (1/(x^100))*(x^100) + +example74 x y z = cos(sqrt(x*x+y*y+z*z)) + +example75 x = sqrt(sqrt(x*x+1)+1) + +example76 a k = a + sqrt(a*a-k) + +example77 a k = -a - sqrt(a*a-k) + +example78 a b x = x*x*a+x*(a+b) +x*b + +example79 x = (x + x) ^ 3 / x + +example80 x = sqrt(x+1)-sqrt 1 + +example81 x = (x+1)-x + +example82 x = sqrt(x+100)-sqrt(x) + +example83 x = 1-cos(x) + +example84 u v = sqrt(sqrt(u^2 + v^2) - u) + +example85 x = exp(log(x)) + +example86 x = sqrt(x + 1) - sqrt x + sin(x - 1) + +example87 x = exp x / sqrt(exp x - 1) * sqrt x + +example88 x = (exp(x) - 1) / x + +example89 x = sqrt(x + 2) - sqrt(x) + +-------------------------------------------------------------------------------- + +-- The main function does nothing +main = return ()