Collect hoisted vectorised functions

diff --git a/compiler/vectorise/VectMonad.hs b/compiler/vectorise/VectMonad.hs
index a658253..dc26b4b 100644 (file)
--- a/compiler/vectorise/VectMonad.hs
+++ b/compiler/vectorise/VectMonad.hs
@@ -124,6 +124,9 @@ data LocalEnv = LocalEnv {
 
                  -- Mapping from tyvars to their PA dictionaries
                , local_tyvar_pa :: VarEnv CoreExpr
 
                  -- Mapping from tyvars to their PA dictionaries
                , local_tyvar_pa :: VarEnv CoreExpr

+
+                 -- Hoisted bindings
+               , local_bindings :: [(Var, CoreExpr)]

                }
               
 
                }
               
 

@@ -141,6 +144,7 @@ initGlobalEnv info instEnvs famInstEnvs
 emptyLocalEnv = LocalEnv {
                    local_vars     = emptyVarEnv
                  , local_tyvar_pa = emptyVarEnv
 emptyLocalEnv = LocalEnv {
                    local_vars     = emptyVarEnv
                  , local_tyvar_pa = emptyVarEnv

+                 , local_bindings = []

                  }
 
 -- FIXME
                  }
 
 -- FIXME

diff --git a/compiler/vectorise/VectUtils.hs b/compiler/vectorise/VectUtils.hs
index 630c425..5b70bf4 100644 (file)
--- a/compiler/vectorise/VectUtils.hs
+++ b/compiler/vectorise/VectUtils.hs
@@ -3,7 +3,8 @@ module VectUtils (
   splitClosureTy,
   mkPADictType, mkPArrayType,
   paDictArgType, paDictOfType,
   splitClosureTy,
   mkPADictType, mkPArrayType,
   paDictArgType, paDictOfType,

-  lookupPArrayFamInst
+  lookupPArrayFamInst,
+  hoistExpr

 ) where
 
 #include "HsVersions.h"
 ) where
 
 #include "HsVersions.h"

@@ -11,6 +12,7 @@ module VectUtils (
 import VectMonad
 
 import CoreSyn
 import VectMonad
 
 import CoreSyn

+import CoreUtils

 import Type
 import TypeRep
 import TyCon
 import Type
 import TypeRep
 import TyCon

@@ -18,6 +20,7 @@ import Var
 import PrelNames
 
 import Outputable
 import PrelNames
 
 import Outputable

+import FastString

 
 import Control.Monad         ( liftM )
 
 
 import Control.Monad         ( liftM )
 

@@ -108,3 +111,11 @@ paDFunApply dfun tys
 lookupPArrayFamInst :: Type -> VM (TyCon, [Type])
 lookupPArrayFamInst ty = builtin parrayTyCon >>= (`lookupFamInst` [ty])
 
 lookupPArrayFamInst :: Type -> VM (TyCon, [Type])
 lookupPArrayFamInst ty = builtin parrayTyCon >>= (`lookupFamInst` [ty])
 

+hoistExpr :: FastString -> CoreExpr -> VM Var
+hoistExpr fs expr
+  = do
+      var <- newLocalVar fs (exprType expr)
+      updLEnv $ \env ->
+        env { local_bindings = (var, expr) : local_bindings env }
+      return var
+

diff --git a/compiler/vectorise/Vectorise.hs b/compiler/vectorise/Vectorise.hs
index 993ed30..c9df41b 100644 (file)
--- a/compiler/vectorise/Vectorise.hs
+++ b/compiler/vectorise/Vectorise.hs
@@ -229,6 +229,10 @@ vectExpr lc (fvs, AnnLam bndr body)
       let tyvars = filter isTyVar (varSetElems fvs)
       info <- mkCEnvInfo fvs bndr body
       (poly_vfn, poly_lfn) <- mkClosureFns info tyvars bndr body
       let tyvars = filter isTyVar (varSetElems fvs)
       info <- mkCEnvInfo fvs bndr body
       (poly_vfn, poly_lfn) <- mkClosureFns info tyvars bndr body

+
+      vfn_var <- hoistExpr FSLIT("vfn") poly_vfn
+      lfn_var <- hoistExpr FSLIT("lfn") poly_lfn
+

       let (venv, lenv) = mkClosureEnvs info lc
 
       let env_ty = cenv_vty info
       let (venv, lenv) = mkClosureEnvs info lc
 
       let env_ty = cenv_vty info

@@ -239,8 +243,8 @@ vectExpr lc (fvs, AnnLam bndr body)
       res_ty <- vectType (exprType $ deAnnotate body)
 
       -- FIXME: move the functions to the top level
       res_ty <- vectType (exprType $ deAnnotate body)
 
       -- FIXME: move the functions to the top level

-      mono_vfn <- applyToTypes poly_vfn (map TyVarTy tyvars)
-      mono_lfn <- applyToTypes poly_lfn (map TyVarTy tyvars)
+      mono_vfn <- applyToTypes (Var vfn_var) (map TyVarTy tyvars)
+      mono_lfn <- applyToTypes (Var lfn_var) (map TyVarTy tyvars)

 
       mk_clo <- builtin mkClosureVar
       mk_cloP <- builtin mkClosurePVar
 
       mk_clo <- builtin mkClosureVar
       mk_cloP <- builtin mkClosurePVar