compiler/vectorise/VectMonad.hs

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