Try not to avoid vectorising purely scalar functions
[ghc-hetmet.git] / compiler / vectorise / VectMonad.hs
index e578dc3..bc120cd 100644 (file)
@@ -1,20 +1,16 @@
-{-# OPTIONS -w #-}
--- The above warning supression flag is a temporary kludge.
--- While working on this module you are encouraged to remove it and fix
--- any warnings in the module. See
---     http://hackage.haskell.org/trac/ghc/wiki/Commentary/CodingStyle#Warnings
--- for details
-
 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,
+  cloneName, cloneId, cloneVar,
   newExportedVar, newLocalVar, newDummyVar, newTyVar,
   
   Builtins(..), sumTyCon, prodTyCon,
+  combinePAVar, scalarZip, closureCtrFun,
   builtin, builtins,
 
   GlobalEnv(..),
@@ -26,50 +22,47 @@ module VectMonad (
 
   getBindName, inBind,
 
-  lookupVar, defGlobalVar,
+  lookupVar, defGlobalVar, globalScalars,
   lookupTyCon, defTyCon,
   lookupDataCon, defDataCon,
   lookupTyConPA, defTyConPA, defTyConPAs,
   lookupTyConPR,
+  lookupBoxedTyCon,
   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 Class
 import Var
+import VarSet
 import VarEnv
 import Id
-import OccName
 import Name
 import NameEnv
-import TysPrim       ( intPrimTy )
-import Module
-import IfaceEnv
-import IOEnv         ( ioToIOEnv )
+import IOEnv         ( liftIO )
 
 import DsMonad
-import PrelNames
 
 import InstEnv
 import FamInstEnv
 
-import Panic
 import Outputable
 import FastString
 import SrcLoc        ( noSrcSpan )
 
-import Control.Monad ( liftM, zipWithM )
+import Control.Monad
 
 data Scope a b = Global a | Local b
 
@@ -81,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)
@@ -102,6 +99,9 @@ data GlobalEnv = GlobalEnv {
                   -- Mapping from TyCons to their PR dfuns
                 , global_pr_funs :: NameEnv Var
 
+                  -- Mapping from unboxed TyCons to their boxed versions
+                , global_boxed_tycons :: NameEnv TyCon
+
                 -- External package inst-env & home-package inst-env for class
                 -- instances
                 --
@@ -137,16 +137,26 @@ 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       = emptyVarEnv
+    , global_pr_funs       = emptyNameEnv
+    , global_boxed_tycons  = emptyNameEnv
     , global_inst_env      = instEnvs
     , global_fam_inst_env  = famInstEnvs
     , global_bindings      = []
     }
 
+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) }
@@ -157,6 +167,10 @@ extendTyConsEnv :: [(Name, TyCon)] -> GlobalEnv -> GlobalEnv
 extendTyConsEnv ps genv
   = genv { global_tycons = extendNameEnvList (global_tycons genv) ps }
 
+extendDataConsEnv :: [(Name, DataCon)] -> GlobalEnv -> GlobalEnv
+extendDataConsEnv ps genv
+  = genv { global_datacons = extendNameEnvList (global_datacons genv) ps }
+
 extendPAFunsEnv :: [(Name, Var)] -> GlobalEnv -> GlobalEnv
 extendPAFunsEnv ps genv
   = genv { global_pa_funs = extendNameEnvList (global_pa_funs genv) ps }
@@ -165,11 +179,16 @@ setPRFunsEnv :: [(Name, Var)] -> GlobalEnv -> GlobalEnv
 setPRFunsEnv ps genv
   = genv { global_pr_funs = mkNameEnv ps }
 
+setBoxedTyConsEnv :: [(Name, TyCon)] -> GlobalEnv -> GlobalEnv
+setBoxedTyConsEnv ps genv
+  = genv { global_boxed_tycons = mkNameEnv ps }
+
+emptyLocalEnv :: LocalEnv
 emptyLocalEnv = LocalEnv {
                    local_vars     = emptyVarEnv
                  , local_tyvars   = []
                  , local_tyvar_pa = emptyVarEnv
-                 , local_bind_name  = FSLIT("fn")
+                 , local_bind_name  = fsLit "fn"
                  }
 
 -- FIXME
@@ -192,19 +211,49 @@ data VResult a = Yes GlobalEnv LocalEnv a | No
 newtype VM a = VM { runVM :: Builtins -> GlobalEnv -> LocalEnv -> DsM (VResult a) }
 
 instance Monad VM where
-  return x   = VM $ \bi genv lenv -> return (Yes genv lenv x)
+  return x   = VM $ \_  genv lenv -> return (Yes genv lenv x)
   VM p >>= f = VM $ \bi genv lenv -> do
                                       r <- p bi genv lenv
                                       case r of
                                         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
@@ -225,7 +274,10 @@ orElseV p q = maybe q return =<< tryV p
 fixV :: (a -> VM a) -> VM a
 fixV f = VM (\bi genv lenv -> fixDs $ \r -> runVM (f (unYes r)) bi genv lenv )
   where
+    -- NOTE: It is essential that we are lazy in r above so do not replace
+    --       calls to this function by an explicit case.
     unYes (Yes _ _ x) = x
+    unYes No          = panic "VectMonad.fixV: no result"
 
 localV :: VM a -> VM a
 localV p = do
@@ -243,7 +295,10 @@ closedV p = do
               return x
 
 liftDs :: DsM a -> VM a
-liftDs p = VM $ \bi genv lenv -> do { x <- p; return (Yes genv lenv x) }
+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))
@@ -252,7 +307,7 @@ builtins :: (a -> Builtins -> b) -> VM (a -> b)
 builtins f = VM $ \bi genv lenv -> return (Yes genv lenv (`f` bi))
 
 readGEnv :: (GlobalEnv -> a) -> VM a
-readGEnv f = VM $ \bi genv lenv -> return (Yes genv lenv (f genv))
+readGEnv f = VM $ \_ genv lenv -> return (Yes genv lenv (f genv))
 
 setGEnv :: GlobalEnv -> VM ()
 setGEnv genv = VM $ \_ _ lenv -> return (Yes genv lenv ())
@@ -261,7 +316,7 @@ updGEnv :: (GlobalEnv -> GlobalEnv) -> VM ()
 updGEnv f = VM $ \_ genv lenv -> return (Yes (f genv) lenv ())
 
 readLEnv :: (LocalEnv -> a) -> VM a
-readLEnv f = VM $ \bi genv lenv -> return (Yes genv lenv (f lenv))
+readLEnv f = VM $ \_ genv lenv -> return (Yes genv lenv (f lenv))
 
 setLEnv :: LocalEnv -> VM ()
 setLEnv lenv = VM $ \_ genv _ -> return (Yes genv lenv ())
@@ -301,6 +356,9 @@ cloneId mk_occ id ty
               | otherwise       = Id.mkLocalId         name ty
       return id'
 
+cloneVar :: Var -> VM Var
+cloneVar var = liftM (setIdUnique var) (liftDs newUnique)
+
 newExportedVar :: OccName -> Type -> VM Var
 newExportedVar occ_name ty 
   = do
@@ -318,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
@@ -342,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
@@ -356,17 +417,19 @@ defTyCon tc tc' = updGEnv $ \env ->
   env { global_tycons = extendNameEnv (global_tycons env) (tyConName tc) tc' }
 
 lookupDataCon :: DataCon -> VM (Maybe DataCon)
-lookupDataCon dc = readGEnv $ \env -> lookupNameEnv (global_datacons env) (dataConName dc)
+lookupDataCon dc
+  | isTupleTyCon (dataConTyCon dc) = return (Just dc)
+  | otherwise = readGEnv $ \env -> lookupNameEnv (global_datacons env) (dataConName dc)
 
 defDataCon :: DataCon -> DataCon -> VM ()
 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)
@@ -386,6 +449,10 @@ lookupTyVarPA tv = readLEnv $ \env -> lookupVarEnv (local_tyvar_pa env) tv
 lookupTyConPR :: TyCon -> VM (Maybe Var)
 lookupTyConPR tc = readGEnv $ \env -> lookupNameEnv (global_pr_funs env) (tyConName tc)
 
+lookupBoxedTyCon :: TyCon -> VM (Maybe TyCon)
+lookupBoxedTyCon tc = readGEnv $ \env -> lookupNameEnv (global_boxed_tycons env)
+                                                       (tyConName tc)
+
 defLocalTyVar :: TyVar -> VM ()
 defLocalTyVar tv = updLEnv $ \env ->
   env { local_tyvars   = tv : local_tyvars env
@@ -410,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
@@ -422,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.
 --
@@ -456,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
@@ -468,18 +536,26 @@ initV hsc_env guts info p
 
     go =
       do
-        builtins       <- initBuiltins
-        let builtin_tycons = initBuiltinTyCons builtins
+        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 <- ioToIOEnv $ hscEPS hsc_env
+        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 = extendTyConsEnv builtin_tycons
+        let genv = extendImportedVarsEnv builtin_vars
+                 . extendScalars builtin_scalars
+                 . extendTyConsEnv builtin_tycons
+                 . extendDataConsEnv builtin_datacons
                  . extendPAFunsEnv builtin_pas
                  . setPRFunsEnv    builtin_prs
+                 . setBoxedTyConsEnv builtin_boxed
                  $ initGlobalEnv info instEnvs famInstEnvs
 
         r <- runVM p builtins genv emptyLocalEnv