Scope(..),
VM,
- noV, tryV, maybeV, orElseV, fixV, localV, closedV, initV,
- cloneName, cloneId,
+ noV, traceNoV, ensureV, traceEnsureV, tryV, maybeV, traceMaybeV, orElseV,
+ onlyIfV, fixV, localV, closedV,
+ initV, cantVectorise, maybeCantVectorise, maybeCantVectoriseM,
+ liftDs,
+ cloneName, cloneId, cloneVar,
newExportedVar, newLocalVar, newDummyVar, newTyVar,
- Builtins(..), paDictTyCon, paDictDataCon,
- builtin,
+ Builtins(..), sumTyCon, prodTyCon,
+ combinePAVar, scalarZip, closureCtrFun,
+ builtin, builtins,
GlobalEnv(..),
- setInstEnvs,
+ setFamInstEnv,
readGEnv, setGEnv, updGEnv,
LocalEnv(..),
getBindName, inBind,
- lookupVar, defGlobalVar,
+ lookupVar, defGlobalVar, globalScalars,
lookupTyCon, defTyCon,
lookupDataCon, defDataCon,
- lookupTyConPA, defTyConPA,
+ lookupTyConPA, defTyConPA, defTyConPAs,
+ lookupTyConPR,
+ lookupBoxedTyCon,
+ lookupPrimMethod, lookupPrimPArray,
lookupTyVarPA, defLocalTyVar, defLocalTyVarWithPA, localTyVars,
lookupInst, lookupFamInst
#include "HsVersions.h"
-import HscTypes
+import VectBuiltIn
+
+import HscTypes hiding ( MonadThings(..) )
+import Module ( PackageId )
import CoreSyn
import Class
import TyCon
import DataCon
import Type
import Var
+import VarSet
import VarEnv
import Id
-import OccName
import Name
import NameEnv
-import TysPrim ( intPrimTy )
-import RdrName
+import IOEnv ( liftIO )
import DsMonad
-import PrelNames
import InstEnv
import FamInstEnv
-import Panic
import Outputable
import FastString
import SrcLoc ( noSrcSpan )
-import Control.Monad ( liftM )
+import Control.Monad
data Scope a b = Global a | Local b
-- ----------------------------------------------------------------------------
-- Vectorisation monad
-data Builtins = Builtins {
- parrayTyCon :: TyCon
- , paClass :: Class
- , closureTyCon :: TyCon
- , mkClosureVar :: Var
- , applyClosureVar :: Var
- , mkClosurePVar :: Var
- , applyClosurePVar :: Var
- , lengthPAVar :: Var
- , replicatePAVar :: Var
- , emptyPAVar :: Var
- , liftingContext :: Var
- }
-
-paDictTyCon :: Builtins -> TyCon
-paDictTyCon = classTyCon . paClass
-
-paDictDataCon :: Builtins -> DataCon
-paDictDataCon = classDataCon . paClass
-
-initBuiltins :: DsM Builtins
-initBuiltins
- = do
- parrayTyCon <- dsLookupTyCon parrayTyConName
- paClass <- dsLookupClass paClassName
- closureTyCon <- dsLookupTyCon closureTyConName
-
- mkClosureVar <- dsLookupGlobalId mkClosureName
- applyClosureVar <- dsLookupGlobalId applyClosureName
- mkClosurePVar <- dsLookupGlobalId mkClosurePName
- applyClosurePVar <- dsLookupGlobalId applyClosurePName
- lengthPAVar <- dsLookupGlobalId lengthPAName
- replicatePAVar <- dsLookupGlobalId replicatePAName
- emptyPAVar <- dsLookupGlobalId emptyPAName
-
- liftingContext <- liftM (\u -> mkSysLocal FSLIT("lc") u intPrimTy)
- newUnique
-
- return $ Builtins {
- parrayTyCon = parrayTyCon
- , paClass = paClass
- , closureTyCon = closureTyCon
- , mkClosureVar = mkClosureVar
- , applyClosureVar = applyClosureVar
- , mkClosurePVar = mkClosurePVar
- , applyClosurePVar = applyClosurePVar
- , lengthPAVar = lengthPAVar
- , replicatePAVar = replicatePAVar
- , emptyPAVar = emptyPAVar
- , liftingContext = liftingContext
- }
-
data GlobalEnv = GlobalEnv {
-- Mapping from global variables to their vectorised versions.
--
global_vars :: VarEnv Var
+ -- Purely scalar variables. Code which mentions only these
+ -- variables doesn't have to be lifted.
+ , global_scalars :: VarSet
+
-- Exported variables which have a vectorised version
--
, global_exported_vars :: VarEnv (Var, Var)
--
, global_pa_funs :: NameEnv Var
+ -- Mapping from TyCons to their PR dfuns
+ , global_pr_funs :: NameEnv Var
+
+ -- Mapping from unboxed TyCons to their boxed versions
+ , global_boxed_tycons :: NameEnv TyCon
+
-- External package inst-env & home-package inst-env for class
-- instances
--
-- Hoisted bindings
, global_bindings :: [(Var, CoreExpr)]
-
- -- Global Rdr environment (from ModGuts)
- --
- , global_rdr_env :: GlobalRdrEnv
}
data LocalEnv = LocalEnv {
-- Local binding name
, local_bind_name :: FastString
}
-
-initGlobalEnv :: VectInfo -> (InstEnv, InstEnv) -> FamInstEnvs -> Builtins -> GlobalRdrEnv
- -> GlobalEnv
-initGlobalEnv info instEnvs famInstEnvs bi rdr_env
+initGlobalEnv :: VectInfo -> (InstEnv, InstEnv) -> FamInstEnvs -> GlobalEnv
+initGlobalEnv info instEnvs famInstEnvs
= GlobalEnv {
global_vars = mapVarEnv snd $ vectInfoVar info
+ , global_scalars = emptyVarSet
, global_exported_vars = emptyVarEnv
- , global_tycons = extendNameEnv (mapNameEnv snd (vectInfoTyCon info))
- (tyConName funTyCon) (closureTyCon bi)
-
+ , global_tycons = mapNameEnv snd $ vectInfoTyCon info
, global_datacons = mapNameEnv snd $ vectInfoDataCon info
, global_pa_funs = mapNameEnv snd $ vectInfoPADFun info
+ , global_pr_funs = emptyNameEnv
+ , global_boxed_tycons = emptyNameEnv
, global_inst_env = instEnvs
, global_fam_inst_env = famInstEnvs
, global_bindings = []
- , global_rdr_env = rdr_env
}
-setInstEnvs :: InstEnv -> FamInstEnv -> GlobalEnv -> GlobalEnv
-setInstEnvs l_inst l_fam_inst genv
- | (g_inst, _) <- global_inst_env genv
- , (g_fam_inst, _) <- global_fam_inst_env genv
- = genv { global_inst_env = (g_inst, l_inst)
- , global_fam_inst_env = (g_fam_inst, l_fam_inst) }
+extendImportedVarsEnv :: [(Var, Var)] -> GlobalEnv -> GlobalEnv
+extendImportedVarsEnv ps genv
+ = genv { global_vars = extendVarEnvList (global_vars genv) ps }
+
+extendScalars :: [Var] -> GlobalEnv -> GlobalEnv
+extendScalars vs genv
+ = genv { global_scalars = extendVarSetList (global_scalars genv) vs }
+
+setFamInstEnv :: FamInstEnv -> GlobalEnv -> GlobalEnv
+setFamInstEnv l_fam_inst genv
+ = genv { global_fam_inst_env = (g_fam_inst, l_fam_inst) }
+ where
+ (g_fam_inst, _) = global_fam_inst_env genv
+
+extendTyConsEnv :: [(Name, TyCon)] -> GlobalEnv -> GlobalEnv
+extendTyConsEnv ps genv
+ = genv { global_tycons = extendNameEnvList (global_tycons genv) ps }
+
+extendDataConsEnv :: [(Name, DataCon)] -> GlobalEnv -> GlobalEnv
+extendDataConsEnv ps genv
+ = genv { global_datacons = extendNameEnvList (global_datacons genv) ps }
+
+extendPAFunsEnv :: [(Name, Var)] -> GlobalEnv -> GlobalEnv
+extendPAFunsEnv ps genv
+ = genv { global_pa_funs = extendNameEnvList (global_pa_funs genv) ps }
+
+setPRFunsEnv :: [(Name, Var)] -> GlobalEnv -> GlobalEnv
+setPRFunsEnv ps genv
+ = genv { global_pr_funs = mkNameEnv ps }
+setBoxedTyConsEnv :: [(Name, TyCon)] -> GlobalEnv -> GlobalEnv
+setBoxedTyConsEnv ps genv
+ = genv { global_boxed_tycons = mkNameEnv ps }
+
+emptyLocalEnv :: LocalEnv
emptyLocalEnv = LocalEnv {
local_vars = emptyVarEnv
, local_tyvars = []
, local_tyvar_pa = emptyVarEnv
- , local_bind_name = FSLIT("fn")
+ , local_bind_name = fsLit "fn"
}
-- FIXME
newtype VM a = VM { runVM :: Builtins -> GlobalEnv -> LocalEnv -> DsM (VResult a) }
instance Monad VM where
- return x = VM $ \bi genv lenv -> return (Yes genv lenv x)
+ return x = VM $ \_ genv lenv -> return (Yes genv lenv x)
VM p >>= f = VM $ \bi genv lenv -> do
r <- p bi genv lenv
case r of
Yes genv' lenv' x -> runVM (f x) bi genv' lenv'
No -> return No
+
+cantVectorise :: String -> SDoc -> a
+cantVectorise s d = pgmError
+ . showSDocDump
+ $ vcat [text "*** Vectorisation error ***",
+ nest 4 $ sep [text s, nest 4 d]]
+
+maybeCantVectorise :: String -> SDoc -> Maybe a -> a
+maybeCantVectorise s d Nothing = cantVectorise s d
+maybeCantVectorise _ _ (Just x) = x
+
+maybeCantVectoriseM :: Monad m => String -> SDoc -> m (Maybe a) -> m a
+maybeCantVectoriseM s d p
+ = do
+ r <- p
+ case r of
+ Just x -> return x
+ Nothing -> cantVectorise s d
+
noV :: VM a
noV = VM $ \_ _ _ -> return No
traceNoV :: String -> SDoc -> VM a
traceNoV s d = pprTrace s d noV
+ensureV :: Bool -> VM ()
+ensureV False = noV
+ensureV True = return ()
+
+onlyIfV :: Bool -> VM a -> VM a
+onlyIfV b p = ensureV b >> p
+
+traceEnsureV :: String -> SDoc -> Bool -> VM ()
+traceEnsureV s d False = traceNoV s d
+traceEnsureV _ _ True = return ()
+
tryV :: VM a -> VM (Maybe a)
tryV (VM p) = VM $ \bi genv lenv ->
do
fixV :: (a -> VM a) -> VM a
fixV f = VM (\bi genv lenv -> fixDs $ \r -> runVM (f (unYes r)) bi genv lenv )
where
+ -- NOTE: It is essential that we are lazy in r above so do not replace
+ -- calls to this function by an explicit case.
unYes (Yes _ _ x) = x
+ unYes No = panic "VectMonad.fixV: no result"
localV :: VM a -> VM a
localV p = do
return x
liftDs :: DsM a -> VM a
-liftDs p = VM $ \bi genv lenv -> do { x <- p; return (Yes genv lenv x) }
+liftDs p = VM $ \_ genv lenv -> do { x <- p; return (Yes genv lenv x) }
+
+liftBuiltinDs :: (Builtins -> DsM a) -> VM a
+liftBuiltinDs p = VM $ \bi genv lenv -> do { x <- p bi; return (Yes genv lenv x)}
builtin :: (Builtins -> a) -> VM a
builtin f = VM $ \bi genv lenv -> return (Yes genv lenv (f bi))
+builtins :: (a -> Builtins -> b) -> VM (a -> b)
+builtins f = VM $ \bi genv lenv -> return (Yes genv lenv (`f` bi))
+
readGEnv :: (GlobalEnv -> a) -> VM a
-readGEnv f = VM $ \bi genv lenv -> return (Yes genv lenv (f genv))
+readGEnv f = VM $ \_ genv lenv -> return (Yes genv lenv (f genv))
setGEnv :: GlobalEnv -> VM ()
setGEnv genv = VM $ \_ _ lenv -> return (Yes genv lenv ())
updGEnv f = VM $ \_ genv lenv -> return (Yes (f genv) lenv ())
readLEnv :: (LocalEnv -> a) -> VM a
-readLEnv f = VM $ \bi genv lenv -> return (Yes genv lenv (f lenv))
+readLEnv f = VM $ \_ genv lenv -> return (Yes genv lenv (f lenv))
setLEnv :: LocalEnv -> VM ()
setLEnv lenv = VM $ \_ genv _ -> return (Yes genv lenv ())
| otherwise = Id.mkLocalId name ty
return id'
+cloneVar :: Var -> VM Var
+cloneVar var = liftM (setIdUnique var) (liftDs newUnique)
+
newExportedVar :: OccName -> Type -> VM Var
newExportedVar occ_name ty
= do
return $ mkSysLocal fs u ty
newDummyVar :: Type -> VM Var
-newDummyVar = newLocalVar FSLIT("ds")
+newDummyVar = newLocalVar (fsLit "vv")
newTyVar :: FastString -> Kind -> VM Var
newTyVar fs k
case r of
Just e -> return (Local e)
Nothing -> liftM Global
- $ traceMaybeV "lookupVar" (ppr v)
- (readGEnv $ \env -> lookupVarEnv (global_vars env) v)
+ . maybeCantVectoriseM "Variable not vectorised:" (ppr v)
+ . readGEnv $ \env -> lookupVarEnv (global_vars env) v
+
+globalScalars :: VM VarSet
+globalScalars = readGEnv global_scalars
lookupTyCon :: TyCon -> VM (Maybe TyCon)
lookupTyCon tc
env { global_tycons = extendNameEnv (global_tycons env) (tyConName tc) tc' }
lookupDataCon :: DataCon -> VM (Maybe DataCon)
-lookupDataCon dc = readGEnv $ \env -> lookupNameEnv (global_datacons env) (dataConName dc)
+lookupDataCon dc
+ | isTupleTyCon (dataConTyCon dc) = return (Just dc)
+ | otherwise = readGEnv $ \env -> lookupNameEnv (global_datacons env) (dataConName dc)
defDataCon :: DataCon -> DataCon -> VM ()
defDataCon dc dc' = updGEnv $ \env ->
env { global_datacons = extendNameEnv (global_datacons env) (dataConName dc) dc' }
+lookupPrimPArray :: TyCon -> VM (Maybe TyCon)
+lookupPrimPArray = liftBuiltinDs . primPArray
+
+lookupPrimMethod :: TyCon -> String -> VM (Maybe Var)
+lookupPrimMethod tycon = liftBuiltinDs . primMethod tycon
+
lookupTyConPA :: TyCon -> VM (Maybe Var)
lookupTyConPA tc = readGEnv $ \env -> lookupNameEnv (global_pa_funs env) (tyConName tc)
defTyConPA tc pa = updGEnv $ \env ->
env { global_pa_funs = extendNameEnv (global_pa_funs env) (tyConName tc) pa }
+defTyConPAs :: [(TyCon, Var)] -> VM ()
+defTyConPAs ps = updGEnv $ \env ->
+ env { global_pa_funs = extendNameEnvList (global_pa_funs env)
+ [(tyConName tc, pa) | (tc, pa) <- ps] }
+
lookupTyVarPA :: Var -> VM (Maybe CoreExpr)
-lookupTyVarPA tv = readLEnv $ \env -> lookupVarEnv (local_tyvar_pa env) tv
+lookupTyVarPA tv = readLEnv $ \env -> lookupVarEnv (local_tyvar_pa env) tv
+
+lookupTyConPR :: TyCon -> VM (Maybe Var)
+lookupTyConPR tc = readGEnv $ \env -> lookupNameEnv (global_pr_funs env) (tyConName tc)
+
+lookupBoxedTyCon :: TyCon -> VM (Maybe TyCon)
+lookupBoxedTyCon tc = readGEnv $ \env -> lookupNameEnv (global_boxed_tycons env)
+ (tyConName tc)
defLocalTyVar :: TyVar -> VM ()
defLocalTyVar tv = updLEnv $ \env ->
where
inst_tys' = [ty | Right ty <- inst_tys]
noFlexiVar = all isRight inst_tys
- _other -> traceNoV "lookupInst" (ppr cls <+> ppr tys)
+ _other ->
+ pprPanic "VectMonad.lookupInst: not found " (ppr cls <+> ppr tys)
}
where
isRight (Left _) = False
(ppr $ mkTyConApp tycon tys)
}
-initV :: HscEnv -> ModGuts -> VectInfo -> VM a -> IO (Maybe (VectInfo, a))
-initV hsc_env guts info p
+initV :: PackageId -> HscEnv -> ModGuts -> VectInfo -> VM a -> IO (Maybe (VectInfo, a))
+initV pkg hsc_env guts info p
= do
- eps <- hscEPS hsc_env
- let famInstEnvs = (eps_fam_inst_env eps, mg_fam_inst_env guts)
- let instEnvs = (eps_inst_env eps, mg_inst_env guts)
-
- Just r <- initDs hsc_env (mg_module guts)
+ -- XXX: ignores error messages and warnings, check that this is
+ -- indeed ok (the use of "Just r" suggests so)
+ (_,Just r) <- initDs hsc_env (mg_module guts)
(mg_rdr_env guts)
(mg_types guts)
- (go instEnvs famInstEnvs)
+ go
return r
where
- go instEnvs famInstEnvs =
+ go =
do
- builtins <- initBuiltins
- r <- runVM p builtins (initGlobalEnv info
- instEnvs
- famInstEnvs
- builtins
- (mg_rdr_env guts))
- emptyLocalEnv
+ builtins <- initBuiltins pkg
+ builtin_vars <- initBuiltinVars builtins
+ builtin_tycons <- initBuiltinTyCons builtins
+ let builtin_datacons = initBuiltinDataCons builtins
+ builtin_pas <- initBuiltinPAs builtins
+ builtin_prs <- initBuiltinPRs builtins
+ builtin_boxed <- initBuiltinBoxedTyCons builtins
+ builtin_scalars <- initBuiltinScalars builtins
+
+ eps <- liftIO $ hscEPS hsc_env
+ let famInstEnvs = (eps_fam_inst_env eps, mg_fam_inst_env guts)
+ instEnvs = (eps_inst_env eps, mg_inst_env guts)
+
+ let genv = extendImportedVarsEnv builtin_vars
+ . extendScalars builtin_scalars
+ . extendTyConsEnv builtin_tycons
+ . extendDataConsEnv builtin_datacons
+ . extendPAFunsEnv builtin_pas
+ . setPRFunsEnv builtin_prs
+ . setBoxedTyConsEnv builtin_boxed
+ $ initGlobalEnv info instEnvs famInstEnvs
+
+ r <- runVM p builtins genv emptyLocalEnv
case r of
Yes genv _ x -> return $ Just (new_info genv, x)
No -> return Nothing