Try not to avoid vectorising purely scalar functions
[ghc-hetmet.git] / compiler / vectorise / VectMonad.hs
index 6838338..bc120cd 100644 (file)
@@ -2,12 +2,16 @@ module VectMonad (
   Scope(..),
   VM,
 
-  noV, tryV, maybeV, orElseV, fixV, localV, closedV, initV,
-  cloneName, cloneId,
+  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(..),
-  builtin,
+  Builtins(..), sumTyCon, prodTyCon,
+  combinePAVar, scalarZip, closureCtrFun,
+  builtin, builtins,
 
   GlobalEnv(..),
   setFamInstEnv,
@@ -18,103 +22,62 @@ module VectMonad (
 
   getBindName, inBind,
 
-  lookupVar, defGlobalVar,
+  lookupVar, defGlobalVar, globalScalars,
   lookupTyCon, defTyCon,
   lookupDataCon, defDataCon,
-  lookupTyConPA, defTyConPA, defTyConPAs, defTyConRdrPAs,
+  lookupTyConPA, defTyConPA, defTyConPAs,
+  lookupTyConPR,
+  lookupBoxedTyCon,
+  lookupPrimMethod, lookupPrimPArray,
   lookupTyVarPA, defLocalTyVar, defLocalTyVarWithPA, localTyVars,
 
-  {-lookupInst,-} lookupFamInst
+  lookupInst, lookupFamInst
 ) where
 
 #include "HsVersions.h"
 
-import HscTypes
+import VectBuiltIn
+
+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 OccName
 import Name
 import NameEnv
-import TysPrim       ( intPrimTy )
-import RdrName
+import IOEnv         ( liftIO )
 
 import DsMonad
-import PrelNames
 
 import InstEnv
 import FamInstEnv
 
-import Panic
 import Outputable
 import FastString
 import SrcLoc        ( noSrcSpan )
 
-import Control.Monad ( liftM )
+import Control.Monad
 
 data Scope a b = Global a | Local b
 
 -- ----------------------------------------------------------------------------
 -- Vectorisation monad
 
-data Builtins = Builtins {
-                  parrayTyCon      :: TyCon
-                , paTyCon          :: TyCon
-                , paDataCon        :: DataCon
-                , closureTyCon     :: TyCon
-                , mkClosureVar     :: Var
-                , applyClosureVar  :: Var
-                , mkClosurePVar    :: Var
-                , applyClosurePVar :: Var
-                , lengthPAVar      :: Var
-                , replicatePAVar   :: Var
-                , emptyPAVar       :: Var
-                , liftingContext   :: Var
-                }
-
-initBuiltins :: DsM Builtins
-initBuiltins
-  = do
-      parrayTyCon  <- dsLookupTyCon parrayTyConName
-      paTyCon      <- dsLookupTyCon paTyConName
-      let paDataCon = case tyConDataCons paTyCon of [dc] -> dc
-      closureTyCon <- dsLookupTyCon closureTyConName
-
-      mkClosureVar     <- dsLookupGlobalId mkClosureName
-      applyClosureVar  <- dsLookupGlobalId applyClosureName
-      mkClosurePVar    <- dsLookupGlobalId mkClosurePName
-      applyClosurePVar <- dsLookupGlobalId applyClosurePName
-      lengthPAVar      <- dsLookupGlobalId lengthPAName
-      replicatePAVar   <- dsLookupGlobalId replicatePAName
-      emptyPAVar       <- dsLookupGlobalId emptyPAName
-
-      liftingContext <- liftM (\u -> mkSysLocal FSLIT("lc") u intPrimTy)
-                              newUnique
-
-      return $ Builtins {
-                 parrayTyCon      = parrayTyCon
-               , paTyCon          = paTyCon
-               , paDataCon        = paDataCon
-               , closureTyCon     = closureTyCon
-               , mkClosureVar     = mkClosureVar
-               , applyClosureVar  = applyClosureVar
-               , mkClosurePVar    = mkClosurePVar
-               , applyClosurePVar = applyClosurePVar
-               , lengthPAVar      = lengthPAVar
-               , replicatePAVar   = replicatePAVar
-               , emptyPAVar       = emptyPAVar
-               , liftingContext   = liftingContext
-               }
-
 data GlobalEnv = GlobalEnv {
                   -- Mapping from global variables to their vectorised versions.
                   -- 
                   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)
@@ -133,6 +96,12 @@ data GlobalEnv = GlobalEnv {
                   --
                 , global_pa_funs :: NameEnv Var
 
+                  -- 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
                 --
@@ -145,10 +114,6 @@ data GlobalEnv = GlobalEnv {
 
                 -- Hoisted bindings
                 , global_bindings :: [(Var, CoreExpr)]
-
-                  -- Global Rdr environment (from ModGuts)
-                  --
-                , global_rdr_env :: GlobalRdrEnv
                 }
 
 data LocalEnv = LocalEnv {
@@ -167,36 +132,63 @@ data LocalEnv = LocalEnv {
                  -- Local binding name
                , local_bind_name :: FastString
                }
-              
 
-initGlobalEnv :: VectInfo -> (InstEnv, InstEnv) -> FamInstEnvs -> Builtins -> GlobalRdrEnv
-              -> GlobalEnv
-initGlobalEnv info instEnvs famInstEnvs bi rdr_env
+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        = extendNameEnv (mapNameEnv snd (vectInfoTyCon info))
-                                           (tyConName funTyCon) (closureTyCon bi)
-                              
+    , global_tycons        = mapNameEnv snd $ vectInfoTyCon info
     , global_datacons      = mapNameEnv snd $ vectInfoDataCon info
     , global_pa_funs       = mapNameEnv snd $ vectInfoPADFun info
+    , global_pr_funs       = emptyNameEnv
+    , global_boxed_tycons  = emptyNameEnv
     , global_inst_env      = instEnvs
     , global_fam_inst_env  = famInstEnvs
     , global_bindings      = []
-    , global_rdr_env       = rdr_env
     }
 
+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) }
   where
     (g_fam_inst, _) = global_fam_inst_env genv
 
+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 }
+
+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
@@ -219,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
@@ -252,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
@@ -270,13 +295,19 @@ 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))
 
+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 ())
@@ -285,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 ())
@@ -307,21 +338,6 @@ inBind id p
   = do updLEnv $ \env -> env { local_bind_name = occNameFS (getOccName id) }
        p
 
-lookupRdrName :: RdrName -> VM Name
-lookupRdrName rdr_name
-  = do
-      rdr_env <- readGEnv global_rdr_env
-      case lookupGRE_RdrName rdr_name rdr_env of
-        [gre] -> return (gre_name gre)
-        []    -> pprPanic "VectMonad.lookupRdrName: not found" (ppr rdr_name)
-        _     -> pprPanic "VectMonad.lookupRdrName: ambiguous" (ppr rdr_name)
-
-lookupRdrVar :: RdrName -> VM Var
-lookupRdrVar rdr_name
-  = do
-      name <- lookupRdrName rdr_name
-      liftDs (dsLookupGlobalId name)
-
 cloneName :: (OccName -> OccName) -> Name -> VM Name
 cloneName mk_occ name = liftM make (liftDs newUnique)
   where
@@ -340,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
@@ -357,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
@@ -381,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
@@ -395,12 +417,20 @@ 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 = liftBuiltinDs . primPArray
+
+lookupPrimMethod :: TyCon -> String -> VM (Maybe Var)
+lookupPrimMethod tycon = liftBuiltinDs . primMethod tycon
+
 lookupTyConPA :: TyCon -> VM (Maybe Var)
 lookupTyConPA tc = readGEnv $ \env -> lookupNameEnv (global_pa_funs env) (tyConName tc)
 
@@ -413,18 +443,15 @@ defTyConPAs ps = updGEnv $ \env ->
   env { global_pa_funs = extendNameEnvList (global_pa_funs env)
                                            [(tyConName tc, pa) | (tc, pa) <- ps] }
 
-defTyConRdrPAs :: [(Name, RdrName)] -> VM ()
-defTyConRdrPAs ps
-  = do
-      pas <- mapM lookupRdrVar rdr_names
-      updGEnv $ \env ->
-        env { global_pa_funs = extendNameEnvList (global_pa_funs env)
-                                                 (zip tcs pas) }
-  where
-    (tcs, rdr_names) = unzip ps
-
 lookupTyVarPA :: Var -> VM (Maybe CoreExpr)
-lookupTyVarPA tv = readLEnv $ \env -> lookupVarEnv (local_tyvar_pa env) tv 
+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 ->
@@ -450,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
@@ -462,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.
 --
@@ -496,29 +522,43 @@ 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
-      eps <- hscEPS hsc_env
-      let famInstEnvs = (eps_fam_inst_env eps, mg_fam_inst_env guts)
-      let instEnvs    = (eps_inst_env     eps, mg_inst_env     guts)
-
-      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 instEnvs famInstEnvs)
+                               go
       return r
   where
 
-    go instEnvs famInstEnvs = 
+    go =
       do
-        builtins <- initBuiltins
-        r <- runVM p builtins (initGlobalEnv info
-                                             instEnvs
-                                             famInstEnvs
-                                             builtins
-                                             (mg_rdr_env guts))
-                   emptyLocalEnv
+        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
+                 . setPRFunsEnv    builtin_prs
+                 . setBoxedTyConsEnv builtin_boxed
+                 $ initGlobalEnv info instEnvs famInstEnvs
+
+        r <- runVM p builtins genv emptyLocalEnv
         case r of
           Yes genv _ x -> return $ Just (new_info genv, x)
           No           -> return Nothing