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