compiler/vectorise/Vectorise/Monad/Global.hs

   1 -- Operations on the global state of the vectorisation monad.
   2
   3 module Vectorise.Monad.Global (
   4   readGEnv,
   5   setGEnv,
   6   updGEnv,
   7
   8   -- * Vars
   9   defGlobalVar,
  10
  11   -- * Vectorisation declarations
  12   lookupVectDecl, noVectDecl,
  13
  14   -- * Scalars
  15   globalScalarVars, isGlobalScalar, globalScalarTyCons,
  16
  17   -- * TyCons
  18   lookupTyCon,
  19   defTyCon, globalVectTyCons,
  20
  21   -- * Datacons
  22   lookupDataCon,
  23   defDataCon,
  24
  25   -- * PA Dictionaries
  26   lookupTyConPA,
  27   defTyConPAs,
  28
  29   -- * PR Dictionaries
  30   lookupTyConPR
  31 ) where
  32
  33 import Vectorise.Monad.Base
  34 import Vectorise.Env
  35
  36 import CoreSyn
  37 import Type
  38 import TyCon
  39 import DataCon
  40 import NameEnv
  41 import NameSet
  42 import VarEnv
  43 import VarSet
  44 import Outputable
  45
  46
  47 -- Global Environment ---------------------------------------------------------
  48
  49 -- |Project something from the global environment.
  50 --
  51 readGEnv :: (GlobalEnv -> a) -> VM a
  52 readGEnv f  = VM $ \_ genv lenv -> return (Yes genv lenv (f genv))
  53
  54 -- |Set the value of the global environment.
  55 --
  56 setGEnv :: GlobalEnv -> VM ()
  57 setGEnv genv  = VM $ \_ _ lenv -> return (Yes genv lenv ())
  58
  59 -- |Update the global environment using the provided function.
  60 --
  61 updGEnv :: (GlobalEnv -> GlobalEnv) -> VM ()
  62 updGEnv f = VM $ \_ genv lenv -> return (Yes (f genv) lenv ())
  63
  64
  65 -- Vars -----------------------------------------------------------------------
  66
  67 -- |Add a mapping between a global var and its vectorised version to the state.
  68 --
  69 defGlobalVar :: Var -> Var -> VM ()
  70 defGlobalVar v v'
  71   = do { traceVt "add global var mapping:" (ppr v <+> text "-->" <+> ppr v')
  72
  73        ; updGEnv $ \env -> env { global_vars = extendVarEnv (global_vars env) v v' }
  74        }
  75
  76
  77 -- Vectorisation declarations -------------------------------------------------
  78
  79 -- |Check whether a variable has a (non-scalar) vectorisation declaration.
  80 --
  81 lookupVectDecl :: Var -> VM (Maybe (Type, CoreExpr))
  82 lookupVectDecl var = readGEnv $ \env -> lookupVarEnv (global_vect_decls env) var
  83
  84 -- |Check whether a variable has a 'NOVECTORISE' declaration.
  85 --
  86 noVectDecl :: Var -> VM Bool
  87 noVectDecl var = readGEnv $ \env -> elemVarSet var (global_novect_vars env)
  88
  89
  90 -- Scalars --------------------------------------------------------------------
  91
  92 -- |Get the set of global scalar variables.
  93 --
  94 globalScalarVars :: VM VarSet
  95 globalScalarVars = readGEnv global_scalar_vars
  96
  97 -- |Check whether a given variable is in the set of global scalar variables.
  98 --
  99 isGlobalScalar :: Var -> VM Bool
 100 isGlobalScalar var = readGEnv $ \env -> var `elemVarSet` global_scalar_vars env
 101
 102 -- |Get the set of global scalar type constructors including both those scalar type constructors
 103 -- declared in an imported module and those declared in the current module.
 104 --
 105 globalScalarTyCons :: VM NameSet
 106 globalScalarTyCons = readGEnv global_scalar_tycons
 107
 108
 109 -- TyCons ---------------------------------------------------------------------
 110
 111 -- |Lookup the vectorised version of a `TyCon` from the global environment.
 112 --
 113 lookupTyCon :: TyCon -> VM (Maybe TyCon)
 114 lookupTyCon tc
 115   | isUnLiftedTyCon tc || isTupleTyCon tc
 116   = return (Just tc)
 117   | otherwise
 118   = readGEnv $ \env -> lookupNameEnv (global_tycons env) (tyConName tc)
 119
 120 -- |Add a mapping between plain and vectorised `TyCon`s to the global environment.
 121 --
 122 defTyCon :: TyCon -> TyCon -> VM ()
 123 defTyCon tc tc' = updGEnv $ \env ->
 124   env { global_tycons = extendNameEnv (global_tycons env) (tyConName tc) tc' }
 125
 126 -- |Get the set of all vectorised type constructors.
 127 --
 128 globalVectTyCons :: VM (NameEnv TyCon)
 129 globalVectTyCons = readGEnv global_tycons
 130
 131
 132 -- DataCons -------------------------------------------------------------------
 133
 134 -- |Lookup the vectorised version of a `DataCon` from the global environment.
 135 --
 136 lookupDataCon :: DataCon -> VM (Maybe DataCon)
 137 lookupDataCon dc
 138   | isTupleTyCon (dataConTyCon dc)
 139   = return (Just dc)
 140
 141   | otherwise
 142   = readGEnv $ \env -> lookupNameEnv (global_datacons env) (dataConName dc)
 143
 144 -- |Add the mapping between plain and vectorised `DataCon`s to the global environment.
 145 --
 146 defDataCon :: DataCon -> DataCon -> VM ()
 147 defDataCon dc dc' = updGEnv $ \env ->
 148   env { global_datacons = extendNameEnv (global_datacons env) (dataConName dc) dc' }
 149
 150
 151 -- 'PA' dictionaries ------------------------------------------------------------
 152
 153 -- |Lookup the 'PA' dfun of a vectorised type constructor in the global environment.
 154 --
 155 lookupTyConPA :: TyCon -> VM (Maybe Var)
 156 lookupTyConPA tc
 157   = readGEnv $ \env -> lookupNameEnv (global_pa_funs env) (tyConName tc)
 158
 159 -- |Associate vectorised type constructors with the dfun of their 'PA' instances in the global
 160 -- environment.
 161 --
 162 defTyConPAs :: [(TyCon, Var)] -> VM ()
 163 defTyConPAs ps = updGEnv $ \env ->
 164   env { global_pa_funs = extendNameEnvList (global_pa_funs env)
 165                                            [(tyConName tc, pa) | (tc, pa) <- ps] }
 166
 167
 168 -- PR Dictionaries ------------------------------------------------------------
 169
 170 lookupTyConPR :: TyCon -> VM (Maybe Var)
 171 lookupTyConPR tc = readGEnv $ \env -> lookupNameEnv (global_pr_funs env) (tyConName tc)