Added a pragma {-# NOVECTORISE f #-} that suppresses vectorisation of toplevel variab...
[ghc-hetmet.git] / compiler / vectorise / Vectorise / Monad / Global.hs
index 4bd6c77..e471ebb 100644 (file)
@@ -1,34 +1,41 @@
 
 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
@@ -38,23 +45,27 @@ import VarSet
 
 
 -- 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'
@@ -65,15 +76,36 @@ defGlobalVar v v' = updGEnv $ \env ->
             | 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
@@ -82,14 +114,12 @@ lookupTyCon 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 ->
@@ -97,6 +127,7 @@ defTyCon tc tc' = updGEnv $ \env ->
 
 
 -- DataCons -------------------------------------------------------------------
+
 -- | Lookup the vectorised version of a `DataCon` from the global environment.
 lookupDataCon :: DataCon -> VM (Maybe DataCon)
 lookupDataCon dc