compiler/vectorise/Vectorise.hs

   1 module Vectorise( vectorise )
   2 where
   3
   4 #include "HsVersions.h"
   5
   6 import DynFlags
   7 import HscTypes
   8
   9 import CoreLint       ( showPass, endPass )
  10 import TyCon
  11 import Var
  12 import VarEnv
  13 import NameEnv
  14
  15 import DsMonad
  16
  17 import PrelNames
  18
  19 import Outputable
  20
  21 vectorise :: HscEnv -> ModGuts -> IO ModGuts
  22 vectorise hsc_env guts
  23   | not (Opt_Vectorise `dopt` dflags) = return guts
  24   | otherwise
  25   = do
  26       showPass dflags "Vectorisation"
  27       eps <- hscEPS hsc_env
  28       let info = hptVectInfo hsc_env `plusVectInfo` eps_vect_info eps
  29       Just guts' <- initDs hsc_env (mg_module guts)
  30                                    (mg_rdr_env guts)
  31                                    (mg_types guts)
  32                                    (vectoriseModule info guts)
  33       endPass dflags "Vectorisation" Opt_D_dump_vect (mg_binds guts')
  34       return guts'
  35   where
  36     dflags = hsc_dflags hsc_env
  37
  38 -- ----------------------------------------------------------------------------
  39 -- Vectorisation monad
  40
  41 data Builtins = Builtins {
  42                   parrayTyCon      :: TyCon
  43                 , paTyCon          :: TyCon
  44                 , closureTyCon     :: TyCon
  45                 , mkClosureVar     :: Var
  46                 , applyClosureVar  :: Var
  47                 , mkClosurePVar    :: Var
  48                 , applyClosurePVar :: Var
  49                 , closurePAVar     :: Var
  50                 , lengthPAVar      :: Var
  51                 , replicatePAVar   :: Var
  52                 }
  53
  54 initBuiltins :: DsM Builtins
  55 initBuiltins
  56   = do
  57       parrayTyCon  <- dsLookupTyCon parrayTyConName
  58       paTyCon      <- dsLookupTyCon paTyConName
  59       closureTyCon <- dsLookupTyCon closureTyConName
  60
  61       mkClosureVar     <- dsLookupGlobalId mkClosureName
  62       applyClosureVar  <- dsLookupGlobalId applyClosureName
  63       mkClosurePVar    <- dsLookupGlobalId mkClosurePName
  64       applyClosurePVar <- dsLookupGlobalId applyClosurePName
  65       closurePAVar     <- dsLookupGlobalId closurePAName
  66       lengthPAVar      <- dsLookupGlobalId lengthPAName
  67       replicatePAVar   <- dsLookupGlobalId replicatePAName
  68
  69       return $ Builtins {
  70                  parrayTyCon      = parrayTyCon
  71                , paTyCon          = paTyCon
  72                , closureTyCon     = closureTyCon
  73                , mkClosureVar     = mkClosureVar
  74                , applyClosureVar  = applyClosureVar
  75                , mkClosurePVar    = mkClosurePVar
  76                , applyClosurePVar = applyClosurePVar
  77                , closurePAVar     = closurePAVar
  78                , lengthPAVar      = lengthPAVar
  79                , replicatePAVar   = replicatePAVar
  80                }
  81
  82 data VEnv = VEnv {
  83               -- Mapping from variables to their vectorised versions
  84               --
  85               vect_vars :: VarEnv Var
  86
  87               -- Exported variables which have a vectorised version
  88               --
  89             , vect_exported_vars :: VarEnv (Var, Var)
  90
  91               -- Mapping from TyCons to their vectorised versions.
  92               -- TyCons which do not have to be vectorised are mapped to
  93               -- themselves.
  94             , vect_tycons :: NameEnv TyCon
  95             }
  96
  97 initVEnv :: VectInfo -> DsM VEnv
  98 initVEnv info
  99   = return $ VEnv {
 100                vect_vars          = mapVarEnv  snd $ vectInfoCCVar   info
 101              , vect_exported_vars = emptyVarEnv
 102              , vect_tycons        = mapNameEnv snd $ vectInfoCCTyCon info
 103              }
 104
 105 -- FIXME
 106 updVectInfo :: VEnv -> ModGuts -> ModGuts
 107 updVectInfo env guts = guts { mg_vect_info = info' }
 108   where
 109     info' = info {
 110               vectInfoCCVar   = vect_exported_vars env
 111             , vectInfoCCTyCon = tc_env
 112             }
 113
 114     info  = mg_vect_info guts
 115     tyenv = mg_types guts
 116
 117     tc_env = mkNameEnv [(tc_name, (tc,tc')) | tc <- typeEnvTyCons tyenv
 118                                             , let tc_name = tyConName tc
 119                                             , Just tc' <- [lookupNameEnv (vect_tycons env) tc_name]]
 120
 121 newtype VM a = VM { runVM :: Builtins -> VEnv -> DsM (VEnv, a) }
 122
 123 instance Monad VM where
 124   return x   = VM $ \bi env -> return (env, x)
 125   VM p >>= f = VM $ \bi env -> do
 126                                  (env', x) <- p bi env
 127                                  runVM (f x) bi env'
 128
 129 builtin :: (Builtins -> a) -> VM a
 130 builtin f = VM $ \bi env -> return (env, f bi)
 131
 132 readEnv :: (VEnv -> a) -> VM a
 133 readEnv f = VM $ \bi env -> return (env, f env)
 134
 135 setEnv :: VEnv -> VM ()
 136 setEnv env = VM $ \_ _ -> return (env, ())
 137
 138 updEnv :: (VEnv -> VEnv) -> VM ()
 139 updEnv f = VM $ \_ env -> return (f env, ())
 140
 141
 142 lookupTyCon :: TyCon -> VM (Maybe TyCon)
 143 lookupTyCon tc = readEnv $ \env -> lookupNameEnv (vect_tycons env) (tyConName tc)
 144
 145 -- ----------------------------------------------------------------------------
 146 -- Bindings
 147
 148 vectoriseModule :: VectInfo -> ModGuts -> DsM ModGuts
 149 vectoriseModule info guts
 150   = do
 151       builtins <- initBuiltins
 152       env <- initVEnv info
 153       (env', guts') <- runVM (vectModule guts) builtins env
 154       return $ updVectInfo env' guts'
 155
 156 vectModule :: ModGuts -> VM ModGuts
 157 vectModule guts = return guts
 158