1 {-# LANGUAGE NamedFieldPuns #-}
3 -- | The Vectorisation monad.
7 noV, traceNoV, ensureV, traceEnsureV, tryV, maybeV, traceMaybeV, orElseV,
8 onlyIfV, fixV, localV, closedV,
9 initV, cantVectorise, maybeCantVectorise, maybeCantVectoriseM,
11 cloneName, cloneId, cloneVar,
12 newExportedVar, newLocalVar, newLocalVars, newDummyVar, newTyVar,
14 Builtins(..), sumTyCon, prodTyCon, prodDataCon,
15 selTy, selReplicate, selPick, selTags, selElements,
16 combinePDVar, scalarZip, closureCtrFun,
20 readGEnv, setGEnv, updGEnv,
22 readLEnv, setLEnv, updLEnv,
26 lookupVar, defGlobalVar, globalScalars,
27 lookupTyCon, defTyCon,
28 lookupDataCon, defDataCon,
29 lookupTyConPA, defTyConPA, defTyConPAs,
32 lookupPrimMethod, lookupPrimPArray,
33 lookupTyVarPA, defLocalTyVar, defLocalTyVarWithPA, localTyVars,
35 lookupInst, lookupFamInst
38 #include "HsVersions.h"
43 import HscTypes hiding ( MonadThings(..) )
44 import Module ( PackageId )
64 import SrcLoc ( noSrcSpan )
69 -- The Vectorisation Monad ----------------------------------------------------
71 -- Vectorisation can either succeed with new envionment and a value,
72 -- or return with failure.
74 data VResult a = Yes GlobalEnv LocalEnv a | No
76 newtype VM a = VM { runVM :: Builtins -> GlobalEnv -> LocalEnv -> DsM (VResult a) }
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
83 Yes genv' lenv' x -> runVM (f x) bi genv' lenv'
87 -- | Throw an error saying we can't vectorise something
88 cantVectorise :: String -> SDoc -> a
89 cantVectorise s d = pgmError
91 $ vcat [text "*** Vectorisation error ***",
92 nest 4 $ sep [text s, nest 4 d]]
94 maybeCantVectorise :: String -> SDoc -> Maybe a -> a
95 maybeCantVectorise s d Nothing = cantVectorise s d
96 maybeCantVectorise _ _ (Just x) = x
98 maybeCantVectoriseM :: Monad m => String -> SDoc -> m (Maybe a) -> m a
99 maybeCantVectoriseM s d p
104 Nothing -> cantVectorise s d
107 -- Control --------------------------------------------------------------------
108 -- | Return some result saying we've failed.
110 noV = VM $ \_ _ _ -> return No
112 traceNoV :: String -> SDoc -> VM a
113 traceNoV s d = pprTrace s d noV
116 -- | If True then carry on, otherwise fail.
117 ensureV :: Bool -> VM ()
119 ensureV True = return ()
122 -- | If True then return the first argument, otherwise fail.
123 onlyIfV :: Bool -> VM a -> VM a
124 onlyIfV b p = ensureV b >> p
126 traceEnsureV :: String -> SDoc -> Bool -> VM ()
127 traceEnsureV s d False = traceNoV s d
128 traceEnsureV _ _ True = return ()
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 ->
139 Yes genv' lenv' x -> return (Yes genv' lenv' (Just x))
140 No -> return (Yes genv lenv Nothing)
143 maybeV :: VM (Maybe a) -> VM a
144 maybeV p = maybe noV return =<< p
146 traceMaybeV :: String -> SDoc -> VM (Maybe a) -> VM a
147 traceMaybeV s d p = maybe (traceNoV s d) return =<< p
149 orElseV :: VM a -> VM a -> VM a
150 orElseV p q = maybe q return =<< tryV p
152 fixV :: (a -> VM a) -> VM a
153 fixV f = VM (\bi genv lenv -> fixDs $ \r -> runVM (f (unYes r)) bi genv lenv )
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"
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
171 -- | Perform a computation in an empty local environment.
172 closedV :: VM a -> VM a
175 setLEnv (emptyLocalEnv { local_bind_name = local_bind_name env })
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) }
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)}
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))
197 builtins :: (a -> Builtins -> b) -> VM (a -> b)
198 builtins f = VM $ \bi genv lenv -> return (Yes genv lenv (`f` bi))
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))
206 setGEnv :: GlobalEnv -> VM ()
207 setGEnv genv = VM $ \_ _ lenv -> return (Yes genv lenv ())
209 updGEnv :: (GlobalEnv -> GlobalEnv) -> VM ()
210 updGEnv f = VM $ \_ genv lenv -> return (Yes (f genv) lenv ())
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))
217 -- | Set the local environment.
218 setLEnv :: LocalEnv -> VM ()
219 setLEnv lenv = VM $ \_ genv _ -> return (Yes genv lenv ())
221 -- | Update the enviroment using a provided function.
222 updLEnv :: (LocalEnv -> LocalEnv) -> VM ()
223 updLEnv f = VM $ \_ genv lenv -> return (Yes genv (f lenv) ())
226 -- InstEnv --------------------------------------------------------------------
227 getInstEnv :: VM (InstEnv, InstEnv)
228 getInstEnv = readGEnv global_inst_env
230 getFamInstEnv :: VM FamInstEnvs
231 getFamInstEnv = readGEnv global_fam_inst_env
234 -- Names ----------------------------------------------------------------------
235 -- | Get the name of the local binding currently being vectorised.
236 getBindName :: VM FastString
237 getBindName = readLEnv local_bind_name
239 inBind :: Id -> VM a -> VM a
241 = do updLEnv $ \env -> env { local_bind_name = occNameFS (getOccName id) }
244 cloneName :: (OccName -> OccName) -> Name -> VM Name
245 cloneName mk_occ name = liftM make (liftDs newUnique)
247 occ_name = mk_occ (nameOccName name)
249 make u | isExternalName name = mkExternalName u (nameModule name)
252 | otherwise = mkSystemName u occ_name
254 cloneId :: (OccName -> OccName) -> Id -> Type -> VM Id
257 name <- cloneName mk_occ (getName id)
258 let id' | isExportedId id = Id.mkExportedLocalId name ty
259 | otherwise = Id.mkLocalId name ty
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)
266 newExportedVar :: OccName -> Type -> VM Var
267 newExportedVar occ_name ty
269 mod <- liftDs getModuleDs
270 u <- liftDs newUnique
272 let name = mkExternalName u mod occ_name noSrcSpan
274 return $ Id.mkExportedLocalId name ty
276 newLocalVar :: FastString -> Type -> VM Var
279 u <- liftDs newUnique
280 return $ mkSysLocal fs u ty
282 newLocalVars :: FastString -> [Type] -> VM [Var]
283 newLocalVars fs = mapM (newLocalVar fs)
285 newDummyVar :: Type -> VM Var
286 newDummyVar = newLocalVar (fsLit "vv")
288 newTyVar :: FastString -> Kind -> VM Var
291 u <- liftDs newUnique
292 return $ mkTyVar (mkSysTvName u fs) k
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)
302 upd env | isExportedId v = extendVarEnv env v (v, v')
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.
310 lookupVar :: Var -> VM (Scope Var (Var, Var))
312 = do r <- readLEnv $ \env -> lookupVarEnv (local_vars env) v
314 Just e -> return (Local e)
315 Nothing -> liftM Global
316 . maybeCantVectoriseVarM v
317 . readGEnv $ \env -> lookupVarEnv (global_vars env) v
319 maybeCantVectoriseVarM :: Monad m => Var -> m (Maybe Var) -> m Var
320 maybeCantVectoriseVarM v p
328 | Just _ <- isClassOpId_maybe var
329 = cantVectorise "ClassOpId not vectorised:" (ppr var)
332 = cantVectorise "Variable not vectorised:" (ppr var)
334 -------------------------------------------------------------------------------
335 globalScalars :: VM VarSet
336 globalScalars = readGEnv global_scalars
338 lookupTyCon :: TyCon -> VM (Maybe TyCon)
340 | isUnLiftedTyCon tc || isTupleTyCon tc = return (Just tc)
342 | otherwise = readGEnv $ \env -> lookupNameEnv (global_tycons env) (tyConName tc)
344 defTyCon :: TyCon -> TyCon -> VM ()
345 defTyCon tc tc' = updGEnv $ \env ->
346 env { global_tycons = extendNameEnv (global_tycons env) (tyConName tc) tc' }
348 lookupDataCon :: DataCon -> VM (Maybe DataCon)
350 | isTupleTyCon (dataConTyCon dc) = return (Just dc)
351 | otherwise = readGEnv $ \env -> lookupNameEnv (global_datacons env) (dataConName dc)
353 defDataCon :: DataCon -> DataCon -> VM ()
354 defDataCon dc dc' = updGEnv $ \env ->
355 env { global_datacons = extendNameEnv (global_datacons env) (dataConName dc) dc' }
357 lookupPrimPArray :: TyCon -> VM (Maybe TyCon)
358 lookupPrimPArray = liftBuiltinDs . primPArray
360 lookupPrimMethod :: TyCon -> String -> VM (Maybe Var)
361 lookupPrimMethod tycon = liftBuiltinDs . primMethod tycon
363 lookupTyConPA :: TyCon -> VM (Maybe Var)
364 lookupTyConPA tc = readGEnv $ \env -> lookupNameEnv (global_pa_funs env) (tyConName tc)
366 defTyConPA :: TyCon -> Var -> VM ()
367 defTyConPA tc pa = updGEnv $ \env ->
368 env { global_pa_funs = extendNameEnv (global_pa_funs env) (tyConName tc) pa }
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] }
375 lookupTyVarPA :: Var -> VM (Maybe CoreExpr)
376 lookupTyVarPA tv = readLEnv $ \env -> lookupVarEnv (local_tyvar_pa env) tv
378 lookupTyConPR :: TyCon -> VM (Maybe Var)
379 lookupTyConPR tc = readGEnv $ \env -> lookupNameEnv (global_pr_funs env) (tyConName tc)
381 lookupBoxedTyCon :: TyCon -> VM (Maybe TyCon)
382 lookupBoxedTyCon tc = readGEnv $ \env -> lookupNameEnv (global_boxed_tycons env)
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
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
397 localTyVars :: VM [TyVar]
398 localTyVars = readLEnv (reverse . local_tyvars)
400 -- Look up the dfun of a class instance.
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.
409 lookupInst :: Class -> [Type] -> VM (DFunId, [Type])
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)
418 inst_tys' = [ty | Right ty <- inst_tys]
419 noFlexiVar = all isRight inst_tys
421 pprPanic "VectMonad.lookupInst: not found " (ppr cls <+> ppr tys)
424 isRight (Left _) = False
425 isRight (Right _) = True
427 -- Look up the representation tycon of a family instance.
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.
433 -- Return the instance tycon and its type instance. For example, if we have
435 -- lookupFamInst 'T' '[Int]' yields (':R42T', 'Int')
437 -- then we have a coercion (ie, type instance of family instance coercion)
439 -- :Co:R42T Int :: T [Int] ~ :R42T Int
441 -- which implies that :R42T was declared as 'data instance T [a]'.
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)
450 pprPanic "VectMonad.lookupFamInst: not found: "
451 (ppr $ mkTyConApp tycon tys)
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
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)
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
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)
481 builtin_prs <- initBuiltinPRs builtins instEnvs
482 builtin_pas <- initBuiltinPAs builtins instEnvs
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
493 r <- runVM p builtins genv emptyLocalEnv
495 Yes genv _ x -> return $ Just (new_info genv, x)
498 new_info genv = updVectInfo genv (mg_types guts) info