Fix bug in lifted environment inspection code
[ghc-hetmet.git] / compiler / vectorise / Vectorise.hs
index 993ed30..d5b78f1 100644 (file)
@@ -13,6 +13,8 @@ import CoreLint             ( showPass, endPass )
 import CoreSyn
 import CoreUtils
 import CoreFVs
+import SimplMonad           ( SimplCount, zeroSimplCount )
+import Rules                ( RuleBase )
 import DataCon
 import TyCon
 import Type
@@ -20,16 +22,18 @@ 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 )
 
 import PrelNames
 import TysWiredIn
+import TysPrim              ( intPrimTy )
 import BasicTypes           ( Boxity(..) )
 
 import Outputable
@@ -37,22 +41,58 @@ import FastString
 import Control.Monad        ( liftM, liftM2, mapAndUnzipM, zipWithM_ )
 import Data.Maybe           ( maybeToList )
 
-vectorise :: HscEnv -> ModGuts -> IO ModGuts
-vectorise hsc_env guts
-  | not (Opt_Vectorise `dopt` dflags) = return guts
-  | otherwise
+vectorise :: HscEnv -> UniqSupply -> RuleBase -> ModGuts
+          -> IO (SimplCount, ModGuts)
+vectorise hsc_env _ _ guts
   = do
       showPass dflags "Vectorisation"
       eps <- hscEPS hsc_env
       let info = hptVectInfo hsc_env `plusVectInfo` eps_vect_info eps
       Just (info', guts') <- initV hsc_env guts info (vectModule guts)
       endPass dflags "Vectorisation" Opt_D_dump_vect (mg_binds guts')
-      return $ guts' { mg_vect_info = info' }
+      return (zeroSimplCount dflags, guts' { mg_vect_info = info' })
   where
     dflags = hsc_dflags hsc_env
 
 vectModule :: ModGuts -> VM ModGuts
-vectModule guts = return guts
+vectModule guts
+  = do
+      binds' <- mapM vectTopBind (mg_binds guts)
+      return $ guts { mg_binds = binds' }
+
+vectTopBind :: CoreBind -> VM CoreBind
+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 <- vectType (idType var)
+      name <- cloneName mkVectOcc (getName var)
+      let var' | isExportedId var = Id.mkExportedLocalId name vty
+               | otherwise        = Id.mkLocalId         name vty
+      defGlobalVar var var'
+      return var'
+    
+vectTopRhs :: CoreExpr -> VM CoreExpr
+vectTopRhs = liftM fst . closedV . vectPolyExpr (panic "Empty lifting context") . freeVars
 
 -- ----------------------------------------------------------------------------
 -- Bindings
@@ -109,26 +149,25 @@ capply (vfn, lfn) (varg, larg)
     (arg_ty, res_ty) = splitClosureTy fn_ty
 
 vectVar :: CoreExpr -> Var -> VM (CoreExpr, CoreExpr)
-vectVar lc v = local v `orElseV` global v
-  where
-    local  v = maybeV (readLEnv $ \env -> lookupVarEnv (local_vars env) v)
-    global v = do
-                 vexpr <- maybeV (readGEnv $ \env -> lookupVarEnv (global_vars env) v)
-                 lexpr <- replicateP vexpr lc
-                 return (vexpr, lexpr)
+vectVar lc v
+  = do
+      r <- lookupVar v
+      case r of
+        Local es     -> return es
+        Global vexpr -> do
+                          lexpr <- replicateP vexpr lc
+                          return (vexpr, lexpr)
 
 vectPolyVar :: CoreExpr -> Var -> [Type] -> VM (CoreExpr, CoreExpr)
 vectPolyVar lc v tys
   = do
-      r <- readLEnv $ \env -> lookupVarEnv (local_vars env) v
+      r <- lookupVar v
       case r of
-        Just (vexpr, lexpr) -> liftM2 (,) (mk_app vexpr) (mk_app lexpr)
-        Nothing ->
-          do
-            poly  <- maybeV (readGEnv $ \env -> lookupVarEnv (global_vars env) v)
-            vexpr <- mk_app poly
-            lexpr <- replicateP vexpr lc
-            return (vexpr, lexpr)
+        Local (vexpr, lexpr) -> liftM2 (,) (mk_app vexpr) (mk_app lexpr)
+        Global poly          -> do
+                                  vexpr <- mk_app poly
+                                  lexpr <- replicateP vexpr lc
+                                  return (vexpr, lexpr)
   where
     mk_app e = applyToTypes e =<< mapM vectType tys
 
@@ -229,6 +268,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
+
+      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
@@ -239,8 +282,8 @@ vectExpr lc (fvs, AnnLam bndr body)
       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
@@ -332,7 +375,7 @@ mkClosureFns info tyvars arg body
 mkClosureMonoFns :: CEnvInfo -> Var -> CoreExprWithFVs -> VM (CoreExpr, CoreExpr)
 mkClosureMonoFns info arg body
   = do
-      lc_bndr <- newLocalVar FSLIT("lc") intTy
+      lc_bndr <- newLocalVar FSLIT("lc") intPrimTy
       (varg : vbndrs, larg : lbndrs, (vbody, lbody))
         <- vectBndrsIn (arg : cenv_vars info)
                        (vectExpr (Var lc_bndr) body)
@@ -361,17 +404,21 @@ mkClosureMonoFns info arg body
           return . Let (NonRec lbndr lenv)
                  $ Case (mkApps (Var lengthPA) [Type vty, (Var lbndr)])
                         lc_bndr
-                        intTy
+                        (exprType lbody)
                         [(DEFAULT, [], lbody)]
 
     bind_lenv lenv lbody lc_bndr lbndrs
-      = return
-      $ Case (unwrapFamInstScrut (cenv_repr_tycon info)
-                                 (cenv_repr_tyargs info)
-                                 lenv)
-             (mkWildId lty)
+      = let scrut = unwrapFamInstScrut (cenv_repr_tycon info)
+                                       (cenv_repr_tyargs info)
+                                       lenv
+            lbndrs' | null lbndrs = [mkWildId unitTy]
+                    | otherwise   = lbndrs
+        in
+        return
+      $ Case scrut
+             (mkWildId (exprType scrut))
              (exprType lbody)
-             [(DataAlt (cenv_repr_datacon info), lc_bndr : lbndrs, lbody)]
+             [(DataAlt (cenv_repr_datacon info), lc_bndr : lbndrs', lbody)]
           
 vectTyAppExpr :: CoreExpr -> CoreExprWithFVs -> [Type] -> VM (CoreExpr, CoreExpr)
 vectTyAppExpr lc (_, AnnVar v) tys = vectPolyVar lc v tys
@@ -394,7 +441,7 @@ vectTyCon tc
                     Nothing  -> pprTrace "ccTyCon:" (ppr tc) $ return tc
 
 vectType :: Type -> VM Type
-vectType ty | Just ty' <- coreView ty = vectType ty
+vectType ty | Just ty' <- coreView ty = vectType ty'
 vectType (TyVarTy tv) = return $ TyVarTy tv
 vectType (AppTy ty1 ty2) = liftM2 AppTy (vectType ty1) (vectType ty2)
 vectType (TyConApp tc tys) = liftM2 TyConApp (vectTyCon tc) (mapM vectType tys)