+{-# LANGUAGE NamedFieldPuns #-}
-- | The Vectorisation monad.
module VectMonad (
- Scope(..),
VM,
noV, traceNoV, ensureV, traceEnsureV, tryV, maybeV, traceMaybeV, orElseV,
combinePDVar, scalarZip, closureCtrFun,
builtin, builtins,
- GlobalEnv(..),
setFamInstEnv,
readGEnv, setGEnv, updGEnv,
- LocalEnv(..),
readLEnv, setLEnv, updLEnv,
getBindName, inBind,
#include "HsVersions.h"
import VectBuiltIn
+import Vectorise.Env
import HscTypes hiding ( MonadThings(..) )
import Module ( PackageId )
, vectInfoPADFun = mk_env typeEnvTyCons global_pa_funs
}
where
+ mk_env :: NamedThing from =>
+ (TypeEnv -> [from])
+ -> (GlobalEnv -> NameEnv to)
+ -> NameEnv (from,to)
mk_env from_tyenv from_env = mkNameEnv [(name, (from,to))
| from <- from_tyenv tyenv
, let name = getName from
, Just to <- [lookupNameEnv (from_env env) name]]
-
-- The Vectorisation Monad ----------------------------------------------------
-- Vectorisation can either succeed with new envionment and a value,
readLEnv :: (LocalEnv -> a) -> VM a
readLEnv f = VM $ \_ genv lenv -> return (Yes genv lenv (f lenv))
+-- | Set the local environment.
setLEnv :: LocalEnv -> VM ()
setLEnv lenv = VM $ \_ genv _ -> return (Yes genv lenv ())
+-- | Update the enviroment using a provided function.
updLEnv :: (LocalEnv -> LocalEnv) -> VM ()
updLEnv f = VM $ \_ genv lenv -> return (Yes genv (f lenv) ())
u <- liftDs newUnique
return $ mkTyVar (mkSysTvName u fs) k
+
+-- | Add a mapping between a global var and its vectorised version to the state.
defGlobalVar :: Var -> Var -> VM ()
defGlobalVar v v' = updGEnv $ \env ->
env { global_vars = extendVarEnv (global_vars env) v v'
upd env | isExportedId v = extendVarEnv env v (v, v')
| otherwise = env
+-- Var ------------------------------------------------------------------------
-- | Lookup the vectorised and\/or lifted versions of this variable.
-- If it's in the global environment we get the vectorised version.
-- If it's in the local environment we get both the vectorised and lifted version.
--
lookupVar :: Var -> VM (Scope Var (Var, Var))
lookupVar v
- = do
- r <- readLEnv $ \env -> lookupVarEnv (local_vars env) v
+ = do r <- readLEnv $ \env -> lookupVarEnv (local_vars env) v
case r of
Just e -> return (Local e)
Nothing -> liftM Global
- . maybeCantVectoriseM "Variable not vectorised:" (ppr v)
+ . maybeCantVectoriseVarM v
. readGEnv $ \env -> lookupVarEnv (global_vars env) v
+maybeCantVectoriseVarM :: Monad m => Var -> m (Maybe Var) -> m Var
+maybeCantVectoriseVarM v p
+ = do r <- p
+ case r of
+ Just x -> return x
+ Nothing -> dumpVar v
+
+dumpVar :: Var -> a
+dumpVar var
+ | Just _ <- isClassOpId_maybe var
+ = cantVectorise "ClassOpId not vectorised:" (ppr var)
+
+ | otherwise
+ = cantVectorise "Variable not vectorised:" (ppr var)
+
+-------------------------------------------------------------------------------
globalScalars :: VM VarSet
globalScalars = readGEnv global_scalars
--
lookupFamInst :: TyCon -> [Type] -> VM (TyCon, [Type])
lookupFamInst tycon tys
- = ASSERT( isOpenTyCon tycon )
+ = ASSERT( isFamilyTyCon tycon )
do { instEnv <- getFamInstEnv
; case lookupFamInstEnv instEnv tycon tys of
[(fam_inst, rep_tys)] -> return (famInstTyCon fam_inst, rep_tys)
(ppr $ mkTyConApp tycon tys)
}
+
+-- | Run a vectorisation computation.
initV :: PackageId -> HscEnv -> ModGuts -> VectInfo -> VM a -> IO (Maybe (VectInfo, a))
initV pkg hsc_env guts info p
= do