Vectorisation-specific imports

[ghc-hetmet.git] / compiler / vectorise / Vectorise.hs

diff --git a/compiler/vectorise/Vectorise.hs b/compiler/vectorise/Vectorise.hs
index 648f0ab..6f92724 100644 (file)
--- a/compiler/vectorise/Vectorise.hs
+++ b/compiler/vectorise/Vectorise.hs
@@ -7,6 +7,9 @@ import DynFlags
 import HscTypes
 
 import CoreLint             ( showPass, endPass )
+import CoreSyn
+import CoreUtils
+import CoreFVs
 import TyCon
 import Type
 import TypeRep
@@ -85,9 +88,14 @@ initBuiltins
                }
 
 data VEnv = VEnv {
-              -- Mapping from variables to their vectorised versions
+              -- Mapping from global variables to their vectorised versions.
+              -- 
+              vect_global_vars :: VarEnv CoreExpr
+
+              -- Mapping from local variables to their vectorised and lifted
+              -- versions.
               --
-              vect_vars :: VarEnv Var
+            , vect_local_vars :: VarEnv (CoreExpr, CoreExpr)
 
               -- Exported variables which have a vectorised version
               --
@@ -96,15 +104,27 @@ data VEnv = VEnv {
               -- Mapping from TyCons to their vectorised versions.
               -- TyCons which do not have to be vectorised are mapped to
               -- themselves.
+              --
             , vect_tycons :: NameEnv TyCon
+
+              -- Mapping from TyCons to their PA dictionaries
+              --
+            , vect_tycon_pa :: NameEnv CoreExpr
+
+              -- Mapping from tyvars to their PA dictionaries
+              --
+            , vect_tyvar_pa :: VarEnv CoreExpr
             }
 
 initVEnv :: VectInfo -> DsM VEnv
 initVEnv info
   = return $ VEnv {
-               vect_vars          = mapVarEnv  snd $ vectInfoCCVar   info
+               vect_global_vars   = mapVarEnv  (Var . snd) $ vectInfoCCVar   info
+             , vect_local_vars    = emptyVarEnv
              , vect_exported_vars = emptyVarEnv
              , vect_tycons        = mapNameEnv snd $ vectInfoCCTyCon info
+             , vect_tycon_pa      = emptyNameEnv
+             , vect_tyvar_pa      = emptyVarEnv
              }
 
 -- FIXME
@@ -123,28 +143,48 @@ updVectInfo env guts = guts { mg_vect_info = info' }
                                             , let tc_name = tyConName tc
                                             , Just tc' <- [lookupNameEnv (vect_tycons env) tc_name]]
 
-newtype VM a = VM { runVM :: Builtins -> VEnv -> DsM (VEnv, a) }
+data VResult a = Yes VEnv a | No
+
+newtype VM a = VM { runVM :: Builtins -> VEnv -> DsM (VResult a) }
 
 instance Monad VM where
-  return x   = VM $ \bi env -> return (env, x)
+  return x   = VM $ \bi env -> return (Yes env x)
   VM p >>= f = VM $ \bi env -> do
-                                 (env', x) <- p bi env
-                                 runVM (f x) bi env'
+                                 r <- p bi env
+                                 case r of
+                                   Yes env' x -> runVM (f x) bi env'
+                                   No         -> return No
+
+noV :: VM a
+noV = VM $ \bi env -> return No
+
+tryV :: VM a -> VM (Maybe a)
+tryV (VM p) = VM $ \bi env -> do
+                                r <- p bi env
+                                case r of
+                                  Yes env' x -> return (Yes env' (Just x))
+                                  No         -> return (Yes env Nothing)
+
+maybeV :: VM (Maybe a) -> VM a
+maybeV p = maybe noV return =<< p
+
+orElseV :: VM a -> VM a -> VM a
+orElseV p q = maybe q return =<< tryV p
 
 liftDs :: DsM a -> VM a
-liftDs p = VM $ \bi env -> do { x <- p; return (env, x) }
+liftDs p = VM $ \bi env -> do { x <- p; return (Yes env x) }
 
 builtin :: (Builtins -> a) -> VM a
-builtin f = VM $ \bi env -> return (env, f bi)
+builtin f = VM $ \bi env -> return (Yes env (f bi))
 
 readEnv :: (VEnv -> a) -> VM a
-readEnv f = VM $ \bi env -> return (env, f env)
+readEnv f = VM $ \bi env -> return (Yes env (f env))
 
 setEnv :: VEnv -> VM ()
-setEnv env = VM $ \_ _ -> return (env, ())
+setEnv env = VM $ \_ _ -> return (Yes env ())
 
 updEnv :: (VEnv -> VEnv) -> VM ()
-updEnv f = VM $ \_ env -> return (f env, ())
+updEnv f = VM $ \_ env -> return (Yes (f env) ())
 
 newTyVar :: FastString -> Kind -> VM Var
 newTyVar fs k
@@ -152,6 +192,9 @@ newTyVar fs k
       u <- liftDs newUnique
       return $ mkTyVar (mkSysTvName u fs) k
 
+lookupVar :: Var -> VM CoreExpr
+lookupVar v = maybeV . readEnv $ \env -> lookupVarEnv (vect_vars env) v
+
 lookupTyCon :: TyCon -> VM (Maybe TyCon)
 lookupTyCon tc = readEnv $ \env -> lookupNameEnv (vect_tycons env) (tyConName tc)
 
@@ -163,8 +206,10 @@ vectoriseModule info guts
   = do
       builtins <- initBuiltins
       env <- initVEnv info
-      (env', guts') <- runVM (vectModule guts) builtins env
-      return $ updVectInfo env' guts'
+      r <- runVM (vectModule guts) builtins env
+      case r of
+        Yes env' guts' -> return $ updVectInfo env' guts'
+        No             -> return guts
 
 vectModule :: ModGuts -> VM ModGuts
 vectModule guts = return guts
@@ -204,7 +249,9 @@ paArgType' ty k
 
 vectTyCon :: TyCon -> VM TyCon
 vectTyCon tc
-  | isFunTyCon tc = builtin closureTyCon
+  | isFunTyCon tc        = builtin closureTyCon
+  | isBoxedTupleTyCon tc = return tc
+  | isUnLiftedTyCon tc   = return tc
   | otherwise = do
                   r <- lookupTyCon tc
                   case r of
@@ -228,3 +275,14 @@ vectType (ForAllTy tv ty)
 
 vectType ty = pprPanic "vectType:" (ppr ty)
 
+isClosureTyCon :: TyCon -> Bool
+isClosureTyCon tc = tyConUnique tc == closureTyConKey
+
+splitClosureTy :: Type -> (Type, Type)
+splitClosureTy ty
+  | Just (tc, [arg_ty, res_ty]) <- splitTyConApp_maybe ty
+  , isClosureTyCon tc
+  = (arg_ty, res_ty)
+
+  | otherwise = pprPanic "splitClosureTy" (ppr ty)
+