compiler/vectorise/VectMonad.hs

   1 {-# LANGUAGE NamedFieldPuns #-}
   2
   3 -- | The Vectorisation monad.
   4 module VectMonad (
   5   VM,
   6
   7   noV, traceNoV, ensureV, traceEnsureV, tryV, maybeV, traceMaybeV, orElseV,
   8   onlyIfV, fixV, localV, closedV,
   9   initV, cantVectorise, maybeCantVectorise, maybeCantVectoriseM,
  10   liftDs,
  11   cloneName, cloneId, cloneVar,
  12   newExportedVar, newLocalVar, newLocalVars, newDummyVar, newTyVar,
  13
  14   Builtins(..), sumTyCon, prodTyCon, prodDataCon,
  15   selTy, selReplicate, selPick, selTags, selElements,
  16   combinePDVar, scalarZip, closureCtrFun,
  17   builtin, builtins,
  18
  19   setFamInstEnv,
  20   readGEnv, setGEnv, updGEnv,
  21
  22   readLEnv, setLEnv, updLEnv,
  23
  24   getBindName, inBind,
  25
  26   lookupVar, defGlobalVar, globalScalars,
  27   lookupTyCon, defTyCon,
  28   lookupDataCon, defDataCon,
  29   lookupTyConPA, defTyConPA, defTyConPAs,
  30   lookupTyConPR,
  31   lookupBoxedTyCon,
  32   lookupPrimMethod, lookupPrimPArray,
  33   lookupTyVarPA, defLocalTyVar, defLocalTyVarWithPA, localTyVars,
  34
  35   lookupInst, lookupFamInst
  36 ) where
  37
  38 #include "HsVersions.h"
  39
  40 import VectBuiltIn
  41 import Vectorise.Env
  42 import Vectorise.Vect
  43
  44 import HscTypes hiding  ( MonadThings(..) )
  45 import Module           ( PackageId )
  46 import CoreSyn
  47 import Class
  48 import TyCon
  49 import DataCon
  50 import Type
  51 import Var
  52 import VarSet
  53 import VarEnv
  54 import Id
  55 import Name
  56 import NameEnv
  57
  58 import DsMonad
  59
  60 import InstEnv
  61 import FamInstEnv
  62
  63 import Outputable
  64 import FastString
  65 import SrcLoc        ( noSrcSpan )
  66
  67 import Control.Monad
  68
  69
  70 -- The Vectorisation Monad ----------------------------------------------------
  71
  72 -- Vectorisation can either succeed with new envionment and a value,
  73 -- or return with failure.
  74 --
  75 data VResult a = Yes GlobalEnv LocalEnv a | No
  76
  77 newtype VM a = VM { runVM :: Builtins -> GlobalEnv -> LocalEnv -> DsM (VResult a) }
  78
  79 instance Monad VM where
  80   return x   = VM $ \_  genv lenv -> return (Yes genv lenv x)
  81   VM p >>= f = VM $ \bi genv lenv -> do
  82                                       r <- p bi genv lenv
  83                                       case r of
  84                                         Yes genv' lenv' x -> runVM (f x) bi genv' lenv'
  85                                         No                -> return No
  86
  87
  88 -- | Throw an error saying we can't vectorise something
  89 cantVectorise :: String -> SDoc -> a
  90 cantVectorise s d = pgmError
  91                   . showSDocDump
  92                   $ vcat [text "*** Vectorisation error ***",
  93                           nest 4 $ sep [text s, nest 4 d]]
  94
  95 maybeCantVectorise :: String -> SDoc -> Maybe a -> a
  96 maybeCantVectorise s d Nothing  = cantVectorise s d
  97 maybeCantVectorise _ _ (Just x) = x
  98
  99 maybeCantVectoriseM :: Monad m => String -> SDoc -> m (Maybe a) -> m a
 100 maybeCantVectoriseM s d p
 101   = do
 102       r <- p
 103       case r of
 104         Just x  -> return x
 105         Nothing -> cantVectorise s d
 106
 107
 108 -- Control --------------------------------------------------------------------
 109 -- | Return some result saying we've failed.
 110 noV :: VM a
 111 noV = VM $ \_ _ _ -> return No
 112
 113 traceNoV :: String -> SDoc -> VM a
 114 traceNoV s d = pprTrace s d noV
 115
 116
 117 -- | If True then carry on, otherwise fail.
 118 ensureV :: Bool -> VM ()
 119 ensureV False = noV
 120 ensureV True  = return ()
 121
 122
 123 -- | If True then return the first argument, otherwise fail.
 124 onlyIfV :: Bool -> VM a -> VM a
 125 onlyIfV b p = ensureV b >> p
 126
 127 traceEnsureV :: String -> SDoc -> Bool -> VM ()
 128 traceEnsureV s d False = traceNoV s d
 129 traceEnsureV _ _ True  = return ()
 130
 131
 132 -- | Try some vectorisation computaton.
 133 --      If it succeeds then return Just the result,
 134 --      otherwise return Nothing.
 135 tryV :: VM a -> VM (Maybe a)
 136 tryV (VM p) = VM $ \bi genv lenv ->
 137   do
 138     r <- p bi genv lenv
 139     case r of
 140       Yes genv' lenv' x -> return (Yes genv' lenv' (Just x))
 141       No                -> return (Yes genv  lenv  Nothing)
 142
 143
 144 maybeV :: VM (Maybe a) -> VM a
 145 maybeV p = maybe noV return =<< p
 146
 147 traceMaybeV :: String -> SDoc -> VM (Maybe a) -> VM a
 148 traceMaybeV s d p = maybe (traceNoV s d) return =<< p
 149
 150 orElseV :: VM a -> VM a -> VM a
 151 orElseV p q = maybe q return =<< tryV p
 152
 153 fixV :: (a -> VM a) -> VM a
 154 fixV f = VM (\bi genv lenv -> fixDs $ \r -> runVM (f (unYes r)) bi genv lenv )
 155   where
 156     -- NOTE: It is essential that we are lazy in r above so do not replace
 157     --       calls to this function by an explicit case.
 158     unYes (Yes _ _ x) = x
 159     unYes No          = panic "VectMonad.fixV: no result"
 160
 161
 162 -- Local Environments ---------------------------------------------------------
 163 -- | Perform a computation in its own local environment.
 164 --      This does not alter the environment of the current state.
 165 localV :: VM a -> VM a
 166 localV p = do
 167              env <- readLEnv id
 168              x <- p
 169              setLEnv env
 170              return x
 171
 172 -- | Perform a computation in an empty local environment.
 173 closedV :: VM a -> VM a
 174 closedV p = do
 175               env <- readLEnv id
 176               setLEnv (emptyLocalEnv { local_bind_name = local_bind_name env })
 177               x <- p
 178               setLEnv env
 179               return x
 180
 181 -- Lifting --------------------------------------------------------------------
 182 -- | Lift a desugaring computation into the vectorisation monad.
 183 liftDs :: DsM a -> VM a
 184 liftDs p = VM $ \_ genv lenv -> do { x <- p; return (Yes genv lenv x) }
 185
 186
 187
 188 -- Builtins -------------------------------------------------------------------
 189 -- Operations on Builtins
 190 liftBuiltinDs :: (Builtins -> DsM a) -> VM a
 191 liftBuiltinDs p = VM $ \bi genv lenv -> do { x <- p bi; return (Yes genv lenv x)}
 192
 193
 194 -- | Project something from the set of builtins.
 195 builtin :: (Builtins -> a) -> VM a
 196 builtin f = VM $ \bi genv lenv -> return (Yes genv lenv (f bi))
 197
 198 builtins :: (a -> Builtins -> b) -> VM (a -> b)
 199 builtins f = VM $ \bi genv lenv -> return (Yes genv lenv (`f` bi))
 200
 201
 202 -- Environments ---------------------------------------------------------------
 203 -- | Project something from the global environment.
 204 readGEnv :: (GlobalEnv -> a) -> VM a
 205 readGEnv f = VM $ \_ genv lenv -> return (Yes genv lenv (f genv))
 206
 207 setGEnv :: GlobalEnv -> VM ()
 208 setGEnv genv = VM $ \_ _ lenv -> return (Yes genv lenv ())
 209
 210 updGEnv :: (GlobalEnv -> GlobalEnv) -> VM ()
 211 updGEnv f = VM $ \_ genv lenv -> return (Yes (f genv) lenv ())
 212
 213
 214 -- | Project something from the local environment.
 215 readLEnv :: (LocalEnv -> a) -> VM a
 216 readLEnv f = VM $ \_ genv lenv -> return (Yes genv lenv (f lenv))
 217
 218 -- | Set the local environment.
 219 setLEnv :: LocalEnv -> VM ()
 220 setLEnv lenv = VM $ \_ genv _ -> return (Yes genv lenv ())
 221
 222 -- | Update the enviroment using a provided function.
 223 updLEnv :: (LocalEnv -> LocalEnv) -> VM ()
 224 updLEnv f = VM $ \_ genv lenv -> return (Yes genv (f lenv) ())
 225
 226
 227 -- InstEnv --------------------------------------------------------------------
 228 getInstEnv :: VM (InstEnv, InstEnv)
 229 getInstEnv = readGEnv global_inst_env
 230
 231 getFamInstEnv :: VM FamInstEnvs
 232 getFamInstEnv = readGEnv global_fam_inst_env
 233
 234
 235 -- Names ----------------------------------------------------------------------
 236 -- | Get the name of the local binding currently being vectorised.
 237 getBindName :: VM FastString
 238 getBindName = readLEnv local_bind_name
 239
 240 inBind :: Id -> VM a -> VM a
 241 inBind id p
 242   = do updLEnv $ \env -> env { local_bind_name = occNameFS (getOccName id) }
 243        p
 244
 245 cloneName :: (OccName -> OccName) -> Name -> VM Name
 246 cloneName mk_occ name = liftM make (liftDs newUnique)
 247   where
 248     occ_name = mk_occ (nameOccName name)
 249
 250     make u | isExternalName name = mkExternalName u (nameModule name)
 251                                                     occ_name
 252                                                     (nameSrcSpan name)
 253            | otherwise           = mkSystemName u occ_name
 254
 255 cloneId :: (OccName -> OccName) -> Id -> Type -> VM Id
 256 cloneId mk_occ id ty
 257   = do
 258       name <- cloneName mk_occ (getName id)
 259       let id' | isExportedId id = Id.mkExportedLocalId name ty
 260               | otherwise       = Id.mkLocalId         name ty
 261       return id'
 262
 263 -- Make a fresh instance of this var, with a new unique.
 264 cloneVar :: Var -> VM Var
 265 cloneVar var = liftM (setIdUnique var) (liftDs newUnique)
 266
 267 newExportedVar :: OccName -> Type -> VM Var
 268 newExportedVar occ_name ty
 269   = do
 270       mod <- liftDs getModuleDs
 271       u   <- liftDs newUnique
 272
 273       let name = mkExternalName u mod occ_name noSrcSpan
 274
 275       return $ Id.mkExportedLocalId name ty
 276
 277 newLocalVar :: FastString -> Type -> VM Var
 278 newLocalVar fs ty
 279   = do
 280       u <- liftDs newUnique
 281       return $ mkSysLocal fs u ty
 282
 283 newLocalVars :: FastString -> [Type] -> VM [Var]
 284 newLocalVars fs = mapM (newLocalVar fs)
 285
 286 newDummyVar :: Type -> VM Var
 287 newDummyVar = newLocalVar (fsLit "vv")
 288
 289 newTyVar :: FastString -> Kind -> VM Var
 290 newTyVar fs k
 291   = do
 292       u <- liftDs newUnique
 293       return $ mkTyVar (mkSysTvName u fs) k
 294
 295
 296 -- | Add a mapping between a global var and its vectorised version to the state.
 297 defGlobalVar :: Var -> Var -> VM ()
 298 defGlobalVar v v' = updGEnv $ \env ->
 299   env { global_vars = extendVarEnv (global_vars env) v v'
 300       , global_exported_vars = upd (global_exported_vars env)
 301       }
 302   where
 303     upd env | isExportedId v = extendVarEnv env v (v, v')
 304             | otherwise      = env
 305
 306 -- Var ------------------------------------------------------------------------
 307 -- | Lookup the vectorised and\/or lifted versions of this variable.
 308 --      If it's in the global environment we get the vectorised version.
 309 --      If it's in the local environment we get both the vectorised and lifted version.
 310 --
 311 lookupVar :: Var -> VM (Scope Var (Var, Var))
 312 lookupVar v
 313  = do r <- readLEnv $ \env -> lookupVarEnv (local_vars env) v
 314       case r of
 315         Just e  -> return (Local e)
 316         Nothing -> liftM Global
 317                 . maybeCantVectoriseVarM v
 318                 . readGEnv $ \env -> lookupVarEnv (global_vars env) v
 319
 320 maybeCantVectoriseVarM :: Monad m => Var -> m (Maybe Var) -> m Var
 321 maybeCantVectoriseVarM v p
 322  = do r <- p
 323       case r of
 324         Just x  -> return x
 325         Nothing -> dumpVar v
 326
 327 dumpVar :: Var -> a
 328 dumpVar var
 329         | Just _                <- isClassOpId_maybe var
 330         = cantVectorise "ClassOpId not vectorised:" (ppr var)
 331
 332         | otherwise
 333         = cantVectorise "Variable not vectorised:" (ppr var)
 334
 335 -------------------------------------------------------------------------------
 336 globalScalars :: VM VarSet
 337 globalScalars = readGEnv global_scalars
 338
 339 lookupTyCon :: TyCon -> VM (Maybe TyCon)
 340 lookupTyCon tc
 341   | isUnLiftedTyCon tc || isTupleTyCon tc = return (Just tc)
 342
 343   | otherwise = readGEnv $ \env -> lookupNameEnv (global_tycons env) (tyConName tc)
 344
 345 defTyCon :: TyCon -> TyCon -> VM ()
 346 defTyCon tc tc' = updGEnv $ \env ->
 347   env { global_tycons = extendNameEnv (global_tycons env) (tyConName tc) tc' }
 348
 349 lookupDataCon :: DataCon -> VM (Maybe DataCon)
 350 lookupDataCon dc
 351   | isTupleTyCon (dataConTyCon dc) = return (Just dc)
 352   | otherwise = readGEnv $ \env -> lookupNameEnv (global_datacons env) (dataConName dc)
 353
 354 defDataCon :: DataCon -> DataCon -> VM ()
 355 defDataCon dc dc' = updGEnv $ \env ->
 356   env { global_datacons = extendNameEnv (global_datacons env) (dataConName dc) dc' }
 357
 358 lookupPrimPArray :: TyCon -> VM (Maybe TyCon)
 359 lookupPrimPArray = liftBuiltinDs . primPArray
 360
 361 lookupPrimMethod :: TyCon -> String -> VM (Maybe Var)
 362 lookupPrimMethod tycon = liftBuiltinDs . primMethod tycon
 363
 364 lookupTyConPA :: TyCon -> VM (Maybe Var)
 365 lookupTyConPA tc = readGEnv $ \env -> lookupNameEnv (global_pa_funs env) (tyConName tc)
 366
 367 defTyConPA :: TyCon -> Var -> VM ()
 368 defTyConPA tc pa = updGEnv $ \env ->
 369   env { global_pa_funs = extendNameEnv (global_pa_funs env) (tyConName tc) pa }
 370
 371 defTyConPAs :: [(TyCon, Var)] -> VM ()
 372 defTyConPAs ps = updGEnv $ \env ->
 373   env { global_pa_funs = extendNameEnvList (global_pa_funs env)
 374                                            [(tyConName tc, pa) | (tc, pa) <- ps] }
 375
 376 lookupTyVarPA :: Var -> VM (Maybe CoreExpr)
 377 lookupTyVarPA tv = readLEnv $ \env -> lookupVarEnv (local_tyvar_pa env) tv
 378
 379 lookupTyConPR :: TyCon -> VM (Maybe Var)
 380 lookupTyConPR tc = readGEnv $ \env -> lookupNameEnv (global_pr_funs env) (tyConName tc)
 381
 382 lookupBoxedTyCon :: TyCon -> VM (Maybe TyCon)
 383 lookupBoxedTyCon tc = readGEnv $ \env -> lookupNameEnv (global_boxed_tycons env)
 384                                                        (tyConName tc)
 385
 386 defLocalTyVar :: TyVar -> VM ()
 387 defLocalTyVar tv = updLEnv $ \env ->
 388   env { local_tyvars   = tv : local_tyvars env
 389       , local_tyvar_pa = local_tyvar_pa env `delVarEnv` tv
 390       }
 391
 392 defLocalTyVarWithPA :: TyVar -> CoreExpr -> VM ()
 393 defLocalTyVarWithPA tv pa = updLEnv $ \env ->
 394   env { local_tyvars   = tv : local_tyvars env
 395       , local_tyvar_pa = extendVarEnv (local_tyvar_pa env) tv pa
 396       }
 397
 398 localTyVars :: VM [TyVar]
 399 localTyVars = readLEnv (reverse . local_tyvars)
 400
 401 -- Look up the dfun of a class instance.
 402 --
 403 -- The match must be unique - ie, match exactly one instance - but the
 404 -- type arguments used for matching may be more specific than those of
 405 -- the class instance declaration.  The found class instances must not have
 406 -- any type variables in the instance context that do not appear in the
 407 -- instances head (i.e., no flexi vars); for details for what this means,
 408 -- see the docs at InstEnv.lookupInstEnv.
 409 --
 410 lookupInst :: Class -> [Type] -> VM (DFunId, [Type])
 411 lookupInst cls tys
 412   = do { instEnv <- getInstEnv
 413        ; case lookupInstEnv instEnv cls tys of
 414            ([(inst, inst_tys)], _)
 415              | noFlexiVar -> return (instanceDFunId inst, inst_tys')
 416              | otherwise  -> pprPanic "VectMonad.lookupInst: flexi var: "
 417                                       (ppr $ mkTyConApp (classTyCon cls) tys)
 418              where
 419                inst_tys'  = [ty | Right ty <- inst_tys]
 420                noFlexiVar = all isRight inst_tys
 421            _other         ->
 422              pprPanic "VectMonad.lookupInst: not found " (ppr cls <+> ppr tys)
 423        }
 424   where
 425     isRight (Left  _) = False
 426     isRight (Right _) = True
 427
 428 -- Look up the representation tycon of a family instance.
 429 --
 430 -- The match must be unique - ie, match exactly one instance - but the
 431 -- type arguments used for matching may be more specific than those of
 432 -- the family instance declaration.
 433 --
 434 -- Return the instance tycon and its type instance.  For example, if we have
 435 --
 436 --  lookupFamInst 'T' '[Int]' yields (':R42T', 'Int')
 437 --
 438 -- then we have a coercion (ie, type instance of family instance coercion)
 439 --
 440 --  :Co:R42T Int :: T [Int] ~ :R42T Int
 441 --
 442 -- which implies that :R42T was declared as 'data instance T [a]'.
 443 --
 444 lookupFamInst :: TyCon -> [Type] -> VM (TyCon, [Type])
 445 lookupFamInst tycon tys
 446   = ASSERT( isOpenTyCon tycon )
 447     do { instEnv <- getFamInstEnv
 448        ; case lookupFamInstEnv instEnv tycon tys of
 449            [(fam_inst, rep_tys)] -> return (famInstTyCon fam_inst, rep_tys)
 450            _other                ->
 451              pprPanic "VectMonad.lookupFamInst: not found: "
 452                       (ppr $ mkTyConApp tycon tys)
 453        }
 454
 455
 456 -- | Run a vectorisation computation.
 457 initV :: PackageId -> HscEnv -> ModGuts -> VectInfo -> VM a -> IO (Maybe (VectInfo, a))
 458 initV pkg hsc_env guts info p
 459   = do
 460          -- XXX: ignores error messages and warnings, check that this is
 461          -- indeed ok (the use of "Just r" suggests so)
 462       (_,Just r) <- initDs hsc_env (mg_module guts)
 463                                (mg_rdr_env guts)
 464                                (mg_types guts)
 465                                go
 466       return r
 467   where
 468
 469     go =
 470       do
 471         builtins       <- initBuiltins pkg
 472         builtin_vars   <- initBuiltinVars builtins
 473         builtin_tycons <- initBuiltinTyCons builtins
 474         let builtin_datacons = initBuiltinDataCons builtins
 475         builtin_boxed  <- initBuiltinBoxedTyCons builtins
 476         builtin_scalars <- initBuiltinScalars builtins
 477
 478         eps <- liftIO $ hscEPS hsc_env
 479         let famInstEnvs = (eps_fam_inst_env eps, mg_fam_inst_env guts)
 480             instEnvs    = (eps_inst_env     eps, mg_inst_env     guts)
 481
 482         builtin_prs    <- initBuiltinPRs builtins instEnvs
 483         builtin_pas    <- initBuiltinPAs builtins instEnvs
 484
 485         let genv = extendImportedVarsEnv builtin_vars
 486                  . extendScalars builtin_scalars
 487                  . extendTyConsEnv builtin_tycons
 488                  . extendDataConsEnv builtin_datacons
 489                  . extendPAFunsEnv builtin_pas
 490                  . setPRFunsEnv    builtin_prs
 491                  . setBoxedTyConsEnv builtin_boxed
 492                  $ initGlobalEnv info instEnvs famInstEnvs
 493
 494         r <- runVM p builtins genv emptyLocalEnv
 495         case r of
 496           Yes genv _ x -> return $ Just (new_info genv, x)
 497           No           -> return Nothing
 498
 499     new_info genv = updVectInfo genv (mg_types guts) info
 500