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 )
68 -- | Indicates what scope something (a variable) is in.
69 data Scope a b = Global a | Local b
72 -- | The global environment.
73 data GlobalEnv = GlobalEnv {
74 -- | Mapping from global variables to their vectorised versions.
76 global_vars :: VarEnv Var
78 -- | Purely scalar variables. Code which mentions only these
79 -- variables doesn't have to be lifted.
80 , global_scalars :: VarSet
82 -- | Exported variables which have a vectorised version
84 , global_exported_vars :: VarEnv (Var, Var)
86 -- | Mapping from TyCons to their vectorised versions.
87 -- TyCons which do not have to be vectorised are mapped to
90 , global_tycons :: NameEnv TyCon
92 -- | Mapping from DataCons to their vectorised versions
94 , global_datacons :: NameEnv DataCon
96 -- | Mapping from TyCons to their PA dfuns
98 , global_pa_funs :: NameEnv Var
100 -- | Mapping from TyCons to their PR dfuns
101 , global_pr_funs :: NameEnv Var
103 -- | Mapping from unboxed TyCons to their boxed versions
104 , global_boxed_tycons :: NameEnv TyCon
106 -- | External package inst-env & home-package inst-env for class
109 , global_inst_env :: (InstEnv, InstEnv)
111 -- | External package inst-env & home-package inst-env for family
114 , global_fam_inst_env :: FamInstEnvs
116 -- | Hoisted bindings
117 , global_bindings :: [(Var, CoreExpr)]
120 -- | The local environment.
121 data LocalEnv = LocalEnv {
122 -- Mapping from local variables to their vectorised and
125 local_vars :: VarEnv (Var, Var)
127 -- In-scope type variables
129 , local_tyvars :: [TyVar]
131 -- Mapping from tyvars to their PA dictionaries
132 , local_tyvar_pa :: VarEnv CoreExpr
134 -- Local binding name
135 , local_bind_name :: FastString
139 -- | Create an initial global environment
140 initGlobalEnv :: VectInfo -> (InstEnv, InstEnv) -> FamInstEnvs -> GlobalEnv
141 initGlobalEnv info instEnvs famInstEnvs
143 global_vars = mapVarEnv snd $ vectInfoVar info
144 , global_scalars = emptyVarSet
145 , global_exported_vars = emptyVarEnv
146 , global_tycons = mapNameEnv snd $ vectInfoTyCon info
147 , global_datacons = mapNameEnv snd $ vectInfoDataCon info
148 , global_pa_funs = mapNameEnv snd $ vectInfoPADFun info
149 , global_pr_funs = emptyNameEnv
150 , global_boxed_tycons = emptyNameEnv
151 , global_inst_env = instEnvs
152 , global_fam_inst_env = famInstEnvs
153 , global_bindings = []
157 -- Operators on Global Environments -------------------------------------------
158 extendImportedVarsEnv :: [(Var, Var)] -> GlobalEnv -> GlobalEnv
159 extendImportedVarsEnv ps genv
160 = genv { global_vars = extendVarEnvList (global_vars genv) ps }
162 extendScalars :: [Var] -> GlobalEnv -> GlobalEnv
163 extendScalars vs genv
164 = genv { global_scalars = extendVarSetList (global_scalars genv) vs }
166 setFamInstEnv :: FamInstEnv -> GlobalEnv -> GlobalEnv
167 setFamInstEnv l_fam_inst genv
168 = genv { global_fam_inst_env = (g_fam_inst, l_fam_inst) }
170 (g_fam_inst, _) = global_fam_inst_env genv
172 extendTyConsEnv :: [(Name, TyCon)] -> GlobalEnv -> GlobalEnv
173 extendTyConsEnv ps genv
174 = genv { global_tycons = extendNameEnvList (global_tycons genv) ps }
176 extendDataConsEnv :: [(Name, DataCon)] -> GlobalEnv -> GlobalEnv
177 extendDataConsEnv ps genv
178 = genv { global_datacons = extendNameEnvList (global_datacons genv) ps }
180 extendPAFunsEnv :: [(Name, Var)] -> GlobalEnv -> GlobalEnv
181 extendPAFunsEnv ps genv
182 = genv { global_pa_funs = extendNameEnvList (global_pa_funs genv) ps }
184 setPRFunsEnv :: [(Name, Var)] -> GlobalEnv -> GlobalEnv
186 = genv { global_pr_funs = mkNameEnv ps }
188 setBoxedTyConsEnv :: [(Name, TyCon)] -> GlobalEnv -> GlobalEnv
189 setBoxedTyConsEnv ps genv
190 = genv { global_boxed_tycons = mkNameEnv ps }
193 -- | Create an empty local environment.
194 emptyLocalEnv :: LocalEnv
195 emptyLocalEnv = LocalEnv {
196 local_vars = emptyVarEnv
198 , local_tyvar_pa = emptyVarEnv
199 , local_bind_name = fsLit "fn"
203 updVectInfo :: GlobalEnv -> TypeEnv -> VectInfo -> VectInfo
204 updVectInfo env tyenv info
206 vectInfoVar = global_exported_vars env
207 , vectInfoTyCon = mk_env typeEnvTyCons global_tycons
208 , vectInfoDataCon = mk_env typeEnvDataCons global_datacons
209 , vectInfoPADFun = mk_env typeEnvTyCons global_pa_funs
212 mk_env :: NamedThing from =>
214 -> (GlobalEnv -> NameEnv to)
216 mk_env from_tyenv from_env = mkNameEnv [(name, (from,to))
217 | from <- from_tyenv tyenv
218 , let name = getName from
219 , Just to <- [lookupNameEnv (from_env env) name]]
221 -- The Vectorisation Monad ----------------------------------------------------
223 -- Vectorisation can either succeed with new envionment and a value,
224 -- or return with failure.
226 data VResult a = Yes GlobalEnv LocalEnv a | No
228 newtype VM a = VM { runVM :: Builtins -> GlobalEnv -> LocalEnv -> DsM (VResult a) }
230 instance Monad VM where
231 return x = VM $ \_ genv lenv -> return (Yes genv lenv x)
232 VM p >>= f = VM $ \bi genv lenv -> do
235 Yes genv' lenv' x -> runVM (f x) bi genv' lenv'
239 -- | Throw an error saying we can't vectorise something
240 cantVectorise :: String -> SDoc -> a
241 cantVectorise s d = pgmError
243 $ vcat [text "*** Vectorisation error ***",
244 nest 4 $ sep [text s, nest 4 d]]
246 maybeCantVectorise :: String -> SDoc -> Maybe a -> a
247 maybeCantVectorise s d Nothing = cantVectorise s d
248 maybeCantVectorise _ _ (Just x) = x
250 maybeCantVectoriseM :: Monad m => String -> SDoc -> m (Maybe a) -> m a
251 maybeCantVectoriseM s d p
256 Nothing -> cantVectorise s d
259 -- Control --------------------------------------------------------------------
260 -- | Return some result saying we've failed.
262 noV = VM $ \_ _ _ -> return No
264 traceNoV :: String -> SDoc -> VM a
265 traceNoV s d = pprTrace s d noV
268 -- | If True then carry on, otherwise fail.
269 ensureV :: Bool -> VM ()
271 ensureV True = return ()
274 -- | If True then return the first argument, otherwise fail.
275 onlyIfV :: Bool -> VM a -> VM a
276 onlyIfV b p = ensureV b >> p
278 traceEnsureV :: String -> SDoc -> Bool -> VM ()
279 traceEnsureV s d False = traceNoV s d
280 traceEnsureV _ _ True = return ()
283 -- | Try some vectorisation computaton.
284 -- If it succeeds then return Just the result,
285 -- otherwise return Nothing.
286 tryV :: VM a -> VM (Maybe a)
287 tryV (VM p) = VM $ \bi genv lenv ->
291 Yes genv' lenv' x -> return (Yes genv' lenv' (Just x))
292 No -> return (Yes genv lenv Nothing)
295 maybeV :: VM (Maybe a) -> VM a
296 maybeV p = maybe noV return =<< p
298 traceMaybeV :: String -> SDoc -> VM (Maybe a) -> VM a
299 traceMaybeV s d p = maybe (traceNoV s d) return =<< p
301 orElseV :: VM a -> VM a -> VM a
302 orElseV p q = maybe q return =<< tryV p
304 fixV :: (a -> VM a) -> VM a
305 fixV f = VM (\bi genv lenv -> fixDs $ \r -> runVM (f (unYes r)) bi genv lenv )
307 -- NOTE: It is essential that we are lazy in r above so do not replace
308 -- calls to this function by an explicit case.
309 unYes (Yes _ _ x) = x
310 unYes No = panic "VectMonad.fixV: no result"
313 -- Local Environments ---------------------------------------------------------
314 -- | Perform a computation in its own local environment.
315 -- This does not alter the environment of the current state.
316 localV :: VM a -> VM a
323 -- | Perform a computation in an empty local environment.
324 closedV :: VM a -> VM a
327 setLEnv (emptyLocalEnv { local_bind_name = local_bind_name env })
332 -- Lifting --------------------------------------------------------------------
333 -- | Lift a desugaring computation into the vectorisation monad.
334 liftDs :: DsM a -> VM a
335 liftDs p = VM $ \_ genv lenv -> do { x <- p; return (Yes genv lenv x) }
339 -- Builtins -------------------------------------------------------------------
340 -- Operations on Builtins
341 liftBuiltinDs :: (Builtins -> DsM a) -> VM a
342 liftBuiltinDs p = VM $ \bi genv lenv -> do { x <- p bi; return (Yes genv lenv x)}
345 -- | Project something from the set of builtins.
346 builtin :: (Builtins -> a) -> VM a
347 builtin f = VM $ \bi genv lenv -> return (Yes genv lenv (f bi))
349 builtins :: (a -> Builtins -> b) -> VM (a -> b)
350 builtins f = VM $ \bi genv lenv -> return (Yes genv lenv (`f` bi))
353 -- Environments ---------------------------------------------------------------
354 -- | Project something from the global environment.
355 readGEnv :: (GlobalEnv -> a) -> VM a
356 readGEnv f = VM $ \_ genv lenv -> return (Yes genv lenv (f genv))
358 setGEnv :: GlobalEnv -> VM ()
359 setGEnv genv = VM $ \_ _ lenv -> return (Yes genv lenv ())
361 updGEnv :: (GlobalEnv -> GlobalEnv) -> VM ()
362 updGEnv f = VM $ \_ genv lenv -> return (Yes (f genv) lenv ())
365 -- | Project something from the local environment.
366 readLEnv :: (LocalEnv -> a) -> VM a
367 readLEnv f = VM $ \_ genv lenv -> return (Yes genv lenv (f lenv))
369 -- | Set the local environment.
370 setLEnv :: LocalEnv -> VM ()
371 setLEnv lenv = VM $ \_ genv _ -> return (Yes genv lenv ())
373 -- | Update the enviroment using a provided function.
374 updLEnv :: (LocalEnv -> LocalEnv) -> VM ()
375 updLEnv f = VM $ \_ genv lenv -> return (Yes genv (f lenv) ())
378 -- InstEnv --------------------------------------------------------------------
379 getInstEnv :: VM (InstEnv, InstEnv)
380 getInstEnv = readGEnv global_inst_env
382 getFamInstEnv :: VM FamInstEnvs
383 getFamInstEnv = readGEnv global_fam_inst_env
386 -- Names ----------------------------------------------------------------------
387 -- | Get the name of the local binding currently being vectorised.
388 getBindName :: VM FastString
389 getBindName = readLEnv local_bind_name
391 inBind :: Id -> VM a -> VM a
393 = do updLEnv $ \env -> env { local_bind_name = occNameFS (getOccName id) }
396 cloneName :: (OccName -> OccName) -> Name -> VM Name
397 cloneName mk_occ name = liftM make (liftDs newUnique)
399 occ_name = mk_occ (nameOccName name)
401 make u | isExternalName name = mkExternalName u (nameModule name)
404 | otherwise = mkSystemName u occ_name
406 cloneId :: (OccName -> OccName) -> Id -> Type -> VM Id
409 name <- cloneName mk_occ (getName id)
410 let id' | isExportedId id = Id.mkExportedLocalId name ty
411 | otherwise = Id.mkLocalId name ty
414 -- Make a fresh instance of this var, with a new unique.
415 cloneVar :: Var -> VM Var
416 cloneVar var = liftM (setIdUnique var) (liftDs newUnique)
418 newExportedVar :: OccName -> Type -> VM Var
419 newExportedVar occ_name ty
421 mod <- liftDs getModuleDs
422 u <- liftDs newUnique
424 let name = mkExternalName u mod occ_name noSrcSpan
426 return $ Id.mkExportedLocalId name ty
428 newLocalVar :: FastString -> Type -> VM Var
431 u <- liftDs newUnique
432 return $ mkSysLocal fs u ty
434 newLocalVars :: FastString -> [Type] -> VM [Var]
435 newLocalVars fs = mapM (newLocalVar fs)
437 newDummyVar :: Type -> VM Var
438 newDummyVar = newLocalVar (fsLit "vv")
440 newTyVar :: FastString -> Kind -> VM Var
443 u <- liftDs newUnique
444 return $ mkTyVar (mkSysTvName u fs) k
447 -- | Add a mapping between a global var and its vectorised version to the state.
448 defGlobalVar :: Var -> Var -> VM ()
449 defGlobalVar v v' = updGEnv $ \env ->
450 env { global_vars = extendVarEnv (global_vars env) v v'
451 , global_exported_vars = upd (global_exported_vars env)
454 upd env | isExportedId v = extendVarEnv env v (v, v')
457 -- Var ------------------------------------------------------------------------
458 -- | Lookup the vectorised and\/or lifted versions of this variable.
459 -- If it's in the global environment we get the vectorised version.
460 -- If it's in the local environment we get both the vectorised and lifted version.
462 lookupVar :: Var -> VM (Scope Var (Var, Var))
464 = do r <- readLEnv $ \env -> lookupVarEnv (local_vars env) v
466 Just e -> return (Local e)
467 Nothing -> liftM Global
468 . maybeCantVectoriseVarM v
469 . readGEnv $ \env -> lookupVarEnv (global_vars env) v
471 maybeCantVectoriseVarM :: Monad m => Var -> m (Maybe Var) -> m Var
472 maybeCantVectoriseVarM v p
480 | Just _ <- isClassOpId_maybe var
481 = cantVectorise "ClassOpId not vectorised:" (ppr var)
484 = cantVectorise "Variable not vectorised:" (ppr var)
486 -------------------------------------------------------------------------------
487 globalScalars :: VM VarSet
488 globalScalars = readGEnv global_scalars
490 lookupTyCon :: TyCon -> VM (Maybe TyCon)
492 | isUnLiftedTyCon tc || isTupleTyCon tc = return (Just tc)
494 | otherwise = readGEnv $ \env -> lookupNameEnv (global_tycons env) (tyConName tc)
496 defTyCon :: TyCon -> TyCon -> VM ()
497 defTyCon tc tc' = updGEnv $ \env ->
498 env { global_tycons = extendNameEnv (global_tycons env) (tyConName tc) tc' }
500 lookupDataCon :: DataCon -> VM (Maybe DataCon)
502 | isTupleTyCon (dataConTyCon dc) = return (Just dc)
503 | otherwise = readGEnv $ \env -> lookupNameEnv (global_datacons env) (dataConName dc)
505 defDataCon :: DataCon -> DataCon -> VM ()
506 defDataCon dc dc' = updGEnv $ \env ->
507 env { global_datacons = extendNameEnv (global_datacons env) (dataConName dc) dc' }
509 lookupPrimPArray :: TyCon -> VM (Maybe TyCon)
510 lookupPrimPArray = liftBuiltinDs . primPArray
512 lookupPrimMethod :: TyCon -> String -> VM (Maybe Var)
513 lookupPrimMethod tycon = liftBuiltinDs . primMethod tycon
515 lookupTyConPA :: TyCon -> VM (Maybe Var)
516 lookupTyConPA tc = readGEnv $ \env -> lookupNameEnv (global_pa_funs env) (tyConName tc)
518 defTyConPA :: TyCon -> Var -> VM ()
519 defTyConPA tc pa = updGEnv $ \env ->
520 env { global_pa_funs = extendNameEnv (global_pa_funs env) (tyConName tc) pa }
522 defTyConPAs :: [(TyCon, Var)] -> VM ()
523 defTyConPAs ps = updGEnv $ \env ->
524 env { global_pa_funs = extendNameEnvList (global_pa_funs env)
525 [(tyConName tc, pa) | (tc, pa) <- ps] }
527 lookupTyVarPA :: Var -> VM (Maybe CoreExpr)
528 lookupTyVarPA tv = readLEnv $ \env -> lookupVarEnv (local_tyvar_pa env) tv
530 lookupTyConPR :: TyCon -> VM (Maybe Var)
531 lookupTyConPR tc = readGEnv $ \env -> lookupNameEnv (global_pr_funs env) (tyConName tc)
533 lookupBoxedTyCon :: TyCon -> VM (Maybe TyCon)
534 lookupBoxedTyCon tc = readGEnv $ \env -> lookupNameEnv (global_boxed_tycons env)
537 defLocalTyVar :: TyVar -> VM ()
538 defLocalTyVar tv = updLEnv $ \env ->
539 env { local_tyvars = tv : local_tyvars env
540 , local_tyvar_pa = local_tyvar_pa env `delVarEnv` tv
543 defLocalTyVarWithPA :: TyVar -> CoreExpr -> VM ()
544 defLocalTyVarWithPA tv pa = updLEnv $ \env ->
545 env { local_tyvars = tv : local_tyvars env
546 , local_tyvar_pa = extendVarEnv (local_tyvar_pa env) tv pa
549 localTyVars :: VM [TyVar]
550 localTyVars = readLEnv (reverse . local_tyvars)
552 -- Look up the dfun of a class instance.
554 -- The match must be unique - ie, match exactly one instance - but the
555 -- type arguments used for matching may be more specific than those of
556 -- the class instance declaration. The found class instances must not have
557 -- any type variables in the instance context that do not appear in the
558 -- instances head (i.e., no flexi vars); for details for what this means,
559 -- see the docs at InstEnv.lookupInstEnv.
561 lookupInst :: Class -> [Type] -> VM (DFunId, [Type])
563 = do { instEnv <- getInstEnv
564 ; case lookupInstEnv instEnv cls tys of
565 ([(inst, inst_tys)], _)
566 | noFlexiVar -> return (instanceDFunId inst, inst_tys')
567 | otherwise -> pprPanic "VectMonad.lookupInst: flexi var: "
568 (ppr $ mkTyConApp (classTyCon cls) tys)
570 inst_tys' = [ty | Right ty <- inst_tys]
571 noFlexiVar = all isRight inst_tys
573 pprPanic "VectMonad.lookupInst: not found " (ppr cls <+> ppr tys)
576 isRight (Left _) = False
577 isRight (Right _) = True
579 -- Look up the representation tycon of a family instance.
581 -- The match must be unique - ie, match exactly one instance - but the
582 -- type arguments used for matching may be more specific than those of
583 -- the family instance declaration.
585 -- Return the instance tycon and its type instance. For example, if we have
587 -- lookupFamInst 'T' '[Int]' yields (':R42T', 'Int')
589 -- then we have a coercion (ie, type instance of family instance coercion)
591 -- :Co:R42T Int :: T [Int] ~ :R42T Int
593 -- which implies that :R42T was declared as 'data instance T [a]'.
595 lookupFamInst :: TyCon -> [Type] -> VM (TyCon, [Type])
596 lookupFamInst tycon tys
597 = ASSERT( isFamilyTyCon tycon )
598 do { instEnv <- getFamInstEnv
599 ; case lookupFamInstEnv instEnv tycon tys of
600 [(fam_inst, rep_tys)] -> return (famInstTyCon fam_inst, rep_tys)
602 pprPanic "VectMonad.lookupFamInst: not found: "
603 (ppr $ mkTyConApp tycon tys)
607 -- | Run a vectorisation computation.
608 initV :: PackageId -> HscEnv -> ModGuts -> VectInfo -> VM a -> IO (Maybe (VectInfo, a))
609 initV pkg hsc_env guts info p
611 -- XXX: ignores error messages and warnings, check that this is
612 -- indeed ok (the use of "Just r" suggests so)
613 (_,Just r) <- initDs hsc_env (mg_module guts)
622 builtins <- initBuiltins pkg
623 builtin_vars <- initBuiltinVars builtins
624 builtin_tycons <- initBuiltinTyCons builtins
625 let builtin_datacons = initBuiltinDataCons builtins
626 builtin_boxed <- initBuiltinBoxedTyCons builtins
627 builtin_scalars <- initBuiltinScalars builtins
629 eps <- liftIO $ hscEPS hsc_env
630 let famInstEnvs = (eps_fam_inst_env eps, mg_fam_inst_env guts)
631 instEnvs = (eps_inst_env eps, mg_inst_env guts)
633 builtin_prs <- initBuiltinPRs builtins instEnvs
634 builtin_pas <- initBuiltinPAs builtins instEnvs
636 let genv = extendImportedVarsEnv builtin_vars
637 . extendScalars builtin_scalars
638 . extendTyConsEnv builtin_tycons
639 . extendDataConsEnv builtin_datacons
640 . extendPAFunsEnv builtin_pas
641 . setPRFunsEnv builtin_prs
642 . setBoxedTyConsEnv builtin_boxed
643 $ initGlobalEnv info instEnvs famInstEnvs
645 r <- runVM p builtins genv emptyLocalEnv
647 Yes genv _ x -> return $ Just (new_info genv, x)
650 new_info genv = updVectInfo genv (mg_types guts) info