Vectorisation of top-level bindings

diff --git a/compiler/vectorise/VectMonad.hs b/compiler/vectorise/VectMonad.hs
index fee294f..68966a1 100644 (file)
--- a/compiler/vectorise/VectMonad.hs
+++ b/compiler/vectorise/VectMonad.hs
@@ -263,8 +263,14 @@ newTyVar fs k
       u <- liftDs newUnique
       return $ mkTyVar (mkSysTvName u fs) k
 
-defGlobalVar :: Var -> CoreExpr -> VM ()
-defGlobalVar v e = updGEnv $ \env -> env { global_vars = extendVarEnv (global_vars env) v e }
+defGlobalVar :: Var -> Var -> VM ()
+defGlobalVar v v' = updGEnv $ \env ->
+  env { global_vars = extendVarEnv (global_vars env) v (Var v')
+      , global_exported_vars = upd (global_exported_vars env)
+      }
+  where
+    upd env | isExportedId v = extendVarEnv env v (v, v')
+            | otherwise      = env
 
 lookupVar :: Var -> VM (Scope CoreExpr (CoreExpr, CoreExpr))
 lookupVar v

diff --git a/compiler/vectorise/VectUtils.hs b/compiler/vectorise/VectUtils.hs
index 74a3405..199ef68 100644 (file)
--- a/compiler/vectorise/VectUtils.hs
+++ b/compiler/vectorise/VectUtils.hs
@@ -4,7 +4,7 @@ module VectUtils (
   mkPADictType, mkPArrayType,
   paDictArgType, paDictOfType,
   lookupPArrayFamInst,
-  hoistExpr
+  hoistExpr, takeHoisted
 ) where
 
 #include "HsVersions.h"
@@ -119,3 +119,10 @@ hoistExpr fs expr
         env { global_bindings = (var, expr) : global_bindings env }
       return var
 
+takeHoisted :: VM [(Var, CoreExpr)]
+takeHoisted
+  = do
+      env <- readGEnv id
+      setGEnv $ env { global_bindings = [] }
+      return $ global_bindings env
+

diff --git a/compiler/vectorise/Vectorise.hs b/compiler/vectorise/Vectorise.hs
index a4da858..ccb33ee 100644 (file)
--- a/compiler/vectorise/Vectorise.hs
+++ b/compiler/vectorise/Vectorise.hs
@@ -20,10 +20,11 @@ import TypeRep
 import Var
 import VarEnv
 import VarSet
-import Name                 ( mkSysTvName )
+import Name                 ( mkSysTvName, getName )
 import NameEnv
 import Id
 import MkId                 ( unwrapFamInstScrut )
+import OccName
 
 import DsMonad hiding (mapAndUnzipM)
 import DsUtils              ( mkCoreTup, mkCoreTupTy )
@@ -54,6 +55,41 @@ vectorise hsc_env guts
 vectModule :: ModGuts -> VM ModGuts
 vectModule guts = return guts
 
+vectTopBind b@(NonRec var expr)
+  = do
+      var'  <- vectTopBinder var
+      expr' <- vectTopRhs expr
+      hs    <- takeHoisted
+      return . Rec $ (var, expr) : (var', expr') : hs
+  `orElseV`
+    return b
+
+vectTopBind b@(Rec bs)
+  = do
+      vars'  <- mapM vectTopBinder vars
+      exprs' <- mapM vectTopRhs exprs
+      hs     <- takeHoisted
+      return . Rec $ bs ++ zip vars' exprs' ++ hs
+  `orElseV`
+    return b
+  where
+    (vars, exprs) = unzip bs
+
+vectTopBinder :: Var -> VM Var
+vectTopBinder var
+  = do
+      vty <- liftM (mkForAllTys tyvars) $ vectType mono_ty
+      name <- cloneName mkVectOcc (getName var)
+      let var' | isExportedId var = Id.mkExportedLocalId name vty
+               | otherwise        = Id.mkLocalId         name vty
+      defGlobalVar var var'
+      return var'
+  where
+    (tyvars, mono_ty) = splitForAllTys (idType var)
+    
+vectTopRhs :: CoreExpr -> VM CoreExpr
+vectTopRhs = liftM fst . closedV . vectPolyExpr (panic "Empty lifting context") . freeVars
+
 -- ----------------------------------------------------------------------------
 -- Bindings