module Vectorise.Monad.Global (
- readGEnv,
- setGEnv,
- updGEnv,
-
- -- * Vars
- defGlobalVar,
-
- -- * Scalars
- globalScalars,
-
- -- * TyCons
- lookupTyCon,
- lookupBoxedTyCon,
- defTyCon,
-
- -- * Datacons
- lookupDataCon,
- defDataCon,
-
- -- * PA Dictionaries
- lookupTyConPA,
- defTyConPA,
- defTyConPAs,
-
- -- * PR Dictionaries
- lookupTyConPR
+ readGEnv,
+ setGEnv,
+ updGEnv,
+
+ -- * Vars
+ defGlobalVar,
+
+ -- * Vectorisation declarations
+ lookupVectDecl, noVectDecl,
+
+ -- * Scalars
+ globalScalars, isGlobalScalar,
+
+ -- * TyCons
+ lookupTyCon,
+ lookupBoxedTyCon,
+ defTyCon,
+
+ -- * Datacons
+ lookupDataCon,
+ defDataCon,
+
+ -- * PA Dictionaries
+ lookupTyConPA,
+ defTyConPA,
+ defTyConPAs,
+
+ -- * PR Dictionaries
+ lookupTyConPR
) where
+
import Vectorise.Monad.Base
import Vectorise.Env
+
+import CoreSyn
+import Type
import TyCon
import DataCon
import NameEnv
-- Global Environment ---------------------------------------------------------
--- | Project something from the global environment.
+
+-- |Project something from the global environment.
+--
readGEnv :: (GlobalEnv -> a) -> VM a
readGEnv f = VM $ \_ genv lenv -> return (Yes genv lenv (f genv))
-
--- | Set the value of the global environment.
+-- |Set the value of the global environment.
+--
setGEnv :: GlobalEnv -> VM ()
setGEnv genv = VM $ \_ _ lenv -> return (Yes genv lenv ())
-
--- | Update the global environment using the provided function.
+-- |Update the global environment using the provided function.
+--
updGEnv :: (GlobalEnv -> GlobalEnv) -> VM ()
updGEnv f = VM $ \_ genv lenv -> return (Yes (f genv) lenv ())
-- Vars -----------------------------------------------------------------------
--- | Add a mapping between a global var and its vectorised version to the state.
+
+-- |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'
| otherwise = env
+-- Vectorisation declarations -------------------------------------------------
+
+-- |Check whether a variable has a (non-scalar) vectorisation declaration.
+--
+lookupVectDecl :: Var -> VM (Maybe (Type, CoreExpr))
+lookupVectDecl var = readGEnv $ \env -> lookupVarEnv (global_vect_decls env) var
+
+-- |Check whether a variable has a 'NOVECTORISE' declaration.
+--
+noVectDecl :: Var -> VM Bool
+noVectDecl var = readGEnv $ \env -> elemVarSet var (global_novect_vars env)
+
+
-- Scalars --------------------------------------------------------------------
--- | Get the set of global scalar variables.
+
+-- |Get the set of global scalar variables.
+--
globalScalars :: VM VarSet
-globalScalars
- = readGEnv global_scalars
+globalScalars = readGEnv global_scalar_vars
+
+-- |Check whether a given variable is in the set of global scalar variables.
+--
+isGlobalScalar :: Var -> VM Bool
+isGlobalScalar var = readGEnv $ \env -> elemVarSet var (global_scalar_vars env)
-- TyCons ---------------------------------------------------------------------
--- | Lookup the vectorised version of a `TyCon` from the global environment.
+
+-- |Lookup the vectorised version of a `TyCon` from the global environment.
+--
lookupTyCon :: TyCon -> VM (Maybe TyCon)
lookupTyCon tc
| isUnLiftedTyCon tc || isTupleTyCon tc
| otherwise
= readGEnv $ \env -> lookupNameEnv (global_tycons env) (tyConName tc)
-
-- | Lookup the vectorised version of a boxed `TyCon` from the global environment.
lookupBoxedTyCon :: TyCon -> VM (Maybe TyCon)
lookupBoxedTyCon tc
= readGEnv $ \env -> lookupNameEnv (global_boxed_tycons env)
(tyConName tc)
-
-- | Add a mapping between plain and vectorised `TyCon`s to the global environment.
defTyCon :: TyCon -> TyCon -> VM ()
defTyCon tc tc' = updGEnv $ \env ->
-- DataCons -------------------------------------------------------------------
+
-- | Lookup the vectorised version of a `DataCon` from the global environment.
lookupDataCon :: DataCon -> VM (Maybe DataCon)
lookupDataCon dc