X-Git-Url: http://git.megacz.com/?a=blobdiff_plain;f=compiler%2Fvectorise%2FVectorise%2FEnv.hs;h=5014fd6272f383deb64dc16e2f06f5409e2e68e3;hb=f537dd87c4a07526e2b1fc1bd1c125d652833641;hp=51d45a48c5b293516281ba6e4e4c9356972b8822;hpb=cb482d83091413830831305db007da2f088619f7;p=ghc-hetmet.git

diff --git a/compiler/vectorise/Vectorise/Env.hs b/compiler/vectorise/Vectorise/Env.hs
index 51d45a4..5014fd6 100644
--- a/compiler/vectorise/Vectorise/Env.hs
+++ b/compiler/vectorise/Vectorise/Env.hs
@@ -11,7 +11,8 @@ module Vectorise.Env (
 	initGlobalEnv,
 	extendImportedVarsEnv,
 	extendScalars,
-	setFamInstEnv,
+	setFamEnv,
+        extendFamEnv,
 	extendTyConsEnv,
 	extendDataConsEnv,
 	extendPAFunsEnv,
@@ -19,13 +20,14 @@ module Vectorise.Env (
 	setBoxedTyConsEnv,
 	updVectInfo
 ) where
+
 import HscTypes
 import InstEnv
 import FamInstEnv
 import CoreSyn
+import Type
 import TyCon
 import DataCon
-import Type
 import VarEnv
 import VarSet
 import Var
@@ -70,15 +72,23 @@ emptyLocalEnv = LocalEnv {
 
 -- GlobalEnv ------------------------------------------------------------------
 -- | The global environment.
---	These are things the exist at top-level.
+--      These are things the exist at top-level.
 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
+        = GlobalEnv {
+        -- | Mapping from global variables to their vectorised versions — aka the /vectorisation
+        --   map/.
+          global_vars           :: VarEnv Var
+
+        -- | Mapping from global variables that have a vectorisation declaration to the right-hand
+        --   side of that declaration and its type.  This mapping only applies to non-scalar
+        --   vectorisation declarations.  All variables with a scalar vectorisation declaration are
+        --   mentioned in 'global_scalars'.
+        , global_vect_decls     :: VarEnv (Type, CoreExpr)
+
+        -- | Purely scalar variables. Code which mentions only these variables doesn't have to be
+        --   lifted.  This includes variables from the current module that have a scalar
+        --   vectorisation declaration and those that the vectoriser determines to be scalar.
+        , global_scalars        :: VarSet
 
         -- | Exported variables which have a vectorised version.
         , global_exported_vars	:: VarEnv (Var, Var)
@@ -88,10 +98,10 @@ data GlobalEnv
         , global_tycons		:: NameEnv TyCon
 
         -- | Mapping from DataCons to their vectorised versions.
-        , global_datacons	:: NameEnv DataCon
+        , global_datacons       :: NameEnv DataCon
 
         -- | Mapping from TyCons to their PA dfuns.
-	, global_pa_funs	:: NameEnv Var
+        , global_pa_funs        :: NameEnv Var
 
         -- | Mapping from TyCons to their PR dfuns.
         , global_pr_funs	:: NameEnv Var
@@ -109,24 +119,26 @@ data GlobalEnv
         , global_bindings	:: [(Var, CoreExpr)]
         }
 
-
 -- | Create an initial global environment
-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        = 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      = []
-	}
-
+initGlobalEnv :: VectInfo -> [CoreVect] -> (InstEnv, InstEnv) -> FamInstEnvs -> GlobalEnv
+initGlobalEnv info vectDecls instEnvs famInstEnvs
+  = GlobalEnv 
+  { global_vars          = mapVarEnv snd $ vectInfoVar info
+  , global_vect_decls    = mkVarEnv vects
+  , global_scalars       = mkVarSet scalars
+  , global_exported_vars = emptyVarEnv
+  , 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      = []
+  }
+  where
+    vects   = [(var, (varType var, exp)) | Vect var (Just exp) <- vectDecls]
+    scalars = [var                       | Vect var Nothing    <- vectDecls]
 
 
 -- Operators on Global Environments -------------------------------------------
@@ -135,19 +147,22 @@ extendImportedVarsEnv :: [(Var, Var)] -> GlobalEnv -> GlobalEnv
 extendImportedVarsEnv ps genv
   = genv { global_vars	 = extendVarEnvList (global_vars genv) ps }
 
-
 -- | Extend the set of scalar variables in an environment.
 extendScalars :: [Var] -> GlobalEnv -> GlobalEnv
 extendScalars vs genv
   = genv { global_scalars = extendVarSetList (global_scalars genv) vs }
 
-
 -- | Set the list of type family instances in an environment.
-setFamInstEnv :: FamInstEnv -> GlobalEnv -> GlobalEnv
-setFamInstEnv l_fam_inst genv
+setFamEnv :: FamInstEnv -> GlobalEnv -> GlobalEnv
+setFamEnv l_fam_inst genv
   = genv { global_fam_inst_env = (g_fam_inst, l_fam_inst) }
   where (g_fam_inst, _) = global_fam_inst_env genv
 
+extendFamEnv :: [FamInst] -> GlobalEnv -> GlobalEnv
+extendFamEnv new genv
+  = genv { global_fam_inst_env = (g_fam_inst, extendFamInstEnvList l_fam_inst new) }
+  where (g_fam_inst, l_fam_inst) = global_fam_inst_env genv
+
 
 -- | Extend the list of type constructors in an environment.
 extendTyConsEnv :: [(Name, TyCon)] -> GlobalEnv -> GlobalEnv