Try not to avoid vectorising purely scalar functions
[ghc-hetmet.git] / compiler / vectorise / VectMonad.hs
index 4ae6c17..bc120cd 100644 (file)
@@ -2,13 +2,15 @@ module VectMonad (
   Scope(..),
   VM,
 
-  noV, tryV, maybeV, traceMaybeV, orElseV, fixV, localV, closedV, initV,
+  noV, traceNoV, ensureV, traceEnsureV, tryV, maybeV, traceMaybeV, orElseV,
+  onlyIfV, fixV, localV, closedV,
+  initV, cantVectorise, maybeCantVectorise, maybeCantVectoriseM,
   liftDs,
   cloneName, cloneId, cloneVar,
   newExportedVar, newLocalVar, newDummyVar, newTyVar,
   
-  Builtins(..), sumTyCon, prodTyCon, uarrTy, intPrimArrayTy,
-  combinePAVar,
+  Builtins(..), sumTyCon, prodTyCon,
+  combinePAVar, scalarZip, closureCtrFun,
   builtin, builtins,
 
   GlobalEnv(..),
@@ -20,7 +22,7 @@ module VectMonad (
 
   getBindName, inBind,
 
-  lookupVar, defGlobalVar,
+  lookupVar, defGlobalVar, globalScalars,
   lookupTyCon, defTyCon,
   lookupDataCon, defDataCon,
   lookupTyConPA, defTyConPA, defTyConPAs,
@@ -29,19 +31,22 @@ module VectMonad (
   lookupPrimMethod, lookupPrimPArray,
   lookupTyVarPA, defLocalTyVar, defLocalTyVarWithPA, localTyVars,
 
-  {-lookupInst,-} lookupFamInst
+  lookupInst, lookupFamInst
 ) where
 
 #include "HsVersions.h"
 
 import VectBuiltIn
 
-import HscTypes
+import HscTypes hiding  ( MonadThings(..) )
+import Module           ( PackageId )
 import CoreSyn
+import Class
 import TyCon
 import DataCon
 import Type
 import Var
+import VarSet
 import VarEnv
 import Id
 import Name
@@ -69,6 +74,10 @@ data GlobalEnv = GlobalEnv {
                   -- 
                   global_vars :: VarEnv Var
 
+                  -- Purely scalar variables. Code which mentions only these
+                  -- variables doesn't have to be lifted.
+                , global_scalars :: VarSet
+
                   -- Exported variables which have a vectorised version
                   --
                 , global_exported_vars :: VarEnv (Var, Var)
@@ -128,6 +137,7 @@ 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
@@ -143,6 +153,10 @@ extendImportedVarsEnv :: [(Var, Var)] -> GlobalEnv -> GlobalEnv
 extendImportedVarsEnv ps genv
   = genv { global_vars = extendVarEnvList (global_vars genv) ps }
 
+extendScalars :: [Var] -> GlobalEnv -> GlobalEnv
+extendScalars vs genv
+  = genv { global_scalars = extendVarSetList (global_scalars genv) vs }
+
 setFamInstEnv :: FamInstEnv -> GlobalEnv -> GlobalEnv
 setFamInstEnv l_fam_inst genv
   = genv { global_fam_inst_env = (g_fam_inst, l_fam_inst) }
@@ -204,12 +218,42 @@ instance Monad VM where
                                         Yes genv' lenv' x -> runVM (f x) bi genv' lenv'
                                         No                -> return No
 
+
+cantVectorise :: String -> SDoc -> a
+cantVectorise s d = pgmError
+                  . showSDocDump
+                  $ vcat [text "*** Vectorisation error ***",
+                          nest 4 $ sep [text s, nest 4 d]]
+
+maybeCantVectorise :: String -> SDoc -> Maybe a -> a
+maybeCantVectorise s d Nothing  = cantVectorise s d
+maybeCantVectorise _ _ (Just x) = x
+
+maybeCantVectoriseM :: Monad m => String -> SDoc -> m (Maybe a) -> m a
+maybeCantVectoriseM s d p
+  = do
+      r <- p
+      case r of
+        Just x  -> return x
+        Nothing -> cantVectorise s d
+
 noV :: VM a
 noV = VM $ \_ _ _ -> return No
 
 traceNoV :: String -> SDoc -> VM a
 traceNoV s d = pprTrace s d noV
 
+ensureV :: Bool -> VM ()
+ensureV False = noV
+ensureV True  = return ()
+
+onlyIfV :: Bool -> VM a -> VM a
+onlyIfV b p = ensureV b >> p
+
+traceEnsureV :: String -> SDoc -> Bool -> VM ()
+traceEnsureV s d False = traceNoV s d
+traceEnsureV s d True  = return ()
+
 tryV :: VM a -> VM (Maybe a)
 tryV (VM p) = VM $ \bi genv lenv ->
   do
@@ -253,6 +297,9 @@ closedV p = do
 liftDs :: DsM a -> VM a
 liftDs p = VM $ \_ genv lenv -> do { x <- p; return (Yes genv lenv x) }
 
+liftBuiltinDs :: (Builtins -> DsM a) -> VM a
+liftBuiltinDs p = VM $ \bi genv lenv -> do { x <- p bi; return (Yes genv lenv x)}
+
 builtin :: (Builtins -> a) -> VM a
 builtin f = VM $ \bi genv lenv -> return (Yes genv lenv (f bi))
 
@@ -277,10 +324,8 @@ setLEnv lenv = VM $ \_ genv _ -> return (Yes genv lenv ())
 updLEnv :: (LocalEnv -> LocalEnv) -> VM ()
 updLEnv f = VM $ \_ genv lenv -> return (Yes genv (f lenv) ())
 
-{-
 getInstEnv :: VM (InstEnv, InstEnv)
 getInstEnv = readGEnv global_inst_env
--}
 
 getFamInstEnv :: VM FamInstEnvs
 getFamInstEnv = readGEnv global_fam_inst_env
@@ -331,7 +376,7 @@ newLocalVar fs ty
       return $ mkSysLocal fs u ty
 
 newDummyVar :: Type -> VM Var
-newDummyVar = newLocalVar (fsLit "ds")
+newDummyVar = newLocalVar (fsLit "vv")
 
 newTyVar :: FastString -> Kind -> VM Var
 newTyVar fs k
@@ -355,8 +400,11 @@ lookupVar v
       case r of
         Just e  -> return (Local e)
         Nothing -> liftM Global
-                 $  traceMaybeV "lookupVar" (ppr v)
-                                (readGEnv $ \env -> lookupVarEnv (global_vars env) v)
+                . maybeCantVectoriseM "Variable not vectorised:" (ppr v)
+                . readGEnv $ \env -> lookupVarEnv (global_vars env) v
+
+globalScalars :: VM VarSet
+globalScalars = readGEnv global_scalars
 
 lookupTyCon :: TyCon -> VM (Maybe TyCon)
 lookupTyCon tc
@@ -378,10 +426,10 @@ defDataCon dc dc' = updGEnv $ \env ->
   env { global_datacons = extendNameEnv (global_datacons env) (dataConName dc) dc' }
 
 lookupPrimPArray :: TyCon -> VM (Maybe TyCon)
-lookupPrimPArray = liftDs . primPArray
+lookupPrimPArray = liftBuiltinDs . primPArray
 
 lookupPrimMethod :: TyCon -> String -> VM (Maybe Var)
-lookupPrimMethod tycon = liftDs . primMethod tycon
+lookupPrimMethod tycon = liftBuiltinDs . primMethod tycon
 
 lookupTyConPA :: TyCon -> VM (Maybe Var)
 lookupTyConPA tc = readGEnv $ \env -> lookupNameEnv (global_pa_funs env) (tyConName tc)
@@ -429,7 +477,6 @@ localTyVars = readLEnv (reverse . local_tyvars)
 -- instances head (i.e., no flexi vars); for details for what this means,
 -- see the docs at InstEnv.lookupInstEnv.
 --
-{-
 lookupInst :: Class -> [Type] -> VM (DFunId, [Type])
 lookupInst cls tys
   = do { instEnv <- getInstEnv
@@ -441,12 +488,12 @@ lookupInst cls tys
              where
                inst_tys'  = [ty | Right ty <- inst_tys]
                noFlexiVar = all isRight inst_tys
-          _other         -> traceNoV "lookupInst" (ppr cls <+> ppr tys)
+          _other         ->
+             pprPanic "VectMonad.lookupInst: not found " (ppr cls <+> ppr tys)
        }
   where
     isRight (Left  _) = False
     isRight (Right _) = True
--}
 
 -- Look up the representation tycon of a family instance.
 --
@@ -475,10 +522,12 @@ lookupFamInst tycon tys
                       (ppr $ mkTyConApp tycon tys)
        }
 
-initV :: HscEnv -> ModGuts -> VectInfo -> VM a -> IO (Maybe (VectInfo, a))
-initV hsc_env guts info p
+initV :: PackageId -> HscEnv -> ModGuts -> VectInfo -> VM a -> IO (Maybe (VectInfo, a))
+initV pkg hsc_env guts info p
   = do
-      Just r <- initDs hsc_env (mg_module guts)
+         -- XXX: ignores error messages and warnings, check that this is
+         -- indeed ok (the use of "Just r" suggests so)
+      (_,Just r) <- initDs hsc_env (mg_module guts)
                                (mg_rdr_env guts)
                                (mg_types guts)
                                go
@@ -487,19 +536,21 @@ initV hsc_env guts info p
 
     go =
       do
-        builtins       <- initBuiltins
+        builtins       <- initBuiltins pkg
         builtin_vars   <- initBuiltinVars builtins
         builtin_tycons <- initBuiltinTyCons builtins
         let builtin_datacons = initBuiltinDataCons builtins
         builtin_pas    <- initBuiltinPAs builtins
         builtin_prs    <- initBuiltinPRs builtins
         builtin_boxed  <- initBuiltinBoxedTyCons builtins
+        builtin_scalars <- initBuiltinScalars builtins
 
         eps <- liftIO $ hscEPS hsc_env
         let famInstEnvs = (eps_fam_inst_env eps, mg_fam_inst_env guts)
             instEnvs    = (eps_inst_env     eps, mg_inst_env     guts)
 
         let genv = extendImportedVarsEnv builtin_vars
+                 . extendScalars builtin_scalars
                  . extendTyConsEnv builtin_tycons
                  . extendDataConsEnv builtin_datacons
                  . extendPAFunsEnv builtin_pas