ghc/compiler/ndpFlatten/FlattenMonad.hs

   1 --  $Id$
   2 --
   3 --  Copyright (c) [2001..2002] Manuel M T Chakravarty & Gabriele Keller
   4 --
   5 --  Monad maintaining parallel contexts and substitutions for flattening.
   6 --
   7 --- DESCRIPTION ---------------------------------------------------------------
   8 --
   9 --  The flattening transformation needs to perform a fair amount of plumbing.
  10 --  It needs to mainatin a set of variables, called the parallel context for
  11 --  lifting, variable substitutions in case alternatives, and so on.
  12 --  Moreover, we need to manage uniques to create new variables.  The monad
  13 --  defined in this module takes care of maintaining this state.
  14 --
  15 --- DOCU ----------------------------------------------------------------------
  16 --
  17 --  Language: Haskell 98
  18 --
  19 --  * a parallel context is a set of variables that get vectorised during a
  20 --    lifting transformations (ie, their type changes from `t' to `[:t:]')
  21 --
  22 --  * all vectorised variables in a parallel context have the same size; we
  23 --    call this also the size of the parallel context
  24 --
  25 --  * we represent contexts by maps that give the lifted version of a variable
  26 --    (remember that in GHC, variables contain type information that changes
  27 --    during lifting)
  28 --
  29 --- TODO ----------------------------------------------------------------------
  30 --
  31 --  * Assumptions currently made that should (if they turn out to be true) be
  32 --    documented in The Commentary:
  33 --
  34 --    - Local bindings can be copied without any need to alpha-rename bound
  35 --      variables (or their uniques).  Such renaming is only necessary when
  36 --      bindings in a recursive group are replicated; implying that this is
  37 --      required in the case of top-level bindings).  (Note: The CoreTidy path
  38 --      generates global uniques before code generation.)
  39 --
  40 --  * One FIXME left to resolve.
  41 --
  42
  43 module FlattenMonad (
  44
  45   -- monad definition
  46   --
  47   Flatten, runFlatten,
  48
  49   -- variable generation
  50   --
  51   newVar, mkBind,
  52
  53   -- context management & query operations
  54   --
  55   extendContext, packContext, liftVar, liftConst, intersectWithContext,
  56
  57   -- construction of prelude functions
  58   --
  59   mk'fst, mk'eq, mk'neq, mk'and, mk'or, mk'lengthP, mk'replicateP, mk'mapP,
  60   mk'bpermuteP, mk'bpermuteDftP, mk'indexOfP
  61 ) where
  62
  63 -- standard
  64 import Monad        (mplus)
  65
  66 -- GHC
  67 import CmdLineOpts  (opt_Flatten)
  68 import Panic        (panic)
  69 import Outputable   (Outputable(ppr), pprPanic)
  70 import UniqSupply   (UniqSupply, splitUniqSupply, uniqFromSupply)
  71 import OccName      (UserFS)
  72 import Var          (Var(..))
  73 import Id           (Id, mkSysLocal)
  74 import Name         (Name)
  75 import VarSet       (VarSet, emptyVarSet, unitVarSet, extendVarSet,
  76                      varSetElems, unionVarSet)
  77 import VarEnv       (VarEnv, emptyVarEnv, unitVarEnv, zipVarEnv, plusVarEnv,
  78                      elemVarEnv, lookupVarEnv, lookupVarEnv_NF, delVarEnvList)
  79 import TyCon        (tyConName)
  80 import Type         (Type, tyConAppTyCon)
  81 import HscTypes     (HomeSymbolTable, PersistentCompilerState(..),
  82                      TyThing(..), lookupType)
  83 import PrelNames    (charPrimTyConName, intPrimTyConName, floatPrimTyConName,
  84                      doublePrimTyConName, fstName, andName, orName,
  85                      lengthPName, replicatePName, mapPName, bpermutePName,
  86                      bpermuteDftPName, indexOfPName)
  87 import PrimOp       (eqCharName, eqIntName, eqFloatName, eqDoubleName,
  88                      neqIntName)
  89                      -- neqCharName, neqFloatName,neqDoubleName,
  90 import CoreSyn      (Expr(..), Bind(..), CoreBndr, CoreExpr, CoreBind, mkApps,
  91                      bindersOfBinds)
  92 import CoreUtils    (exprType)
  93
  94 -- friends
  95 import NDPCoreUtils (parrElemTy)
  96
  97
  98 -- definition of the monad
  99 -- -----------------------
 100
 101 -- state maintained by the flattening monad
 102 --
 103 data FlattenState = FlattenState {
 104
 105                       -- our source for uniques
 106                       --
 107                       us       :: UniqSupply,
 108
 109                       -- environment containing all known names (including all
 110                       -- Prelude functions)
 111                       --
 112                       env      :: Name -> Id,
 113
 114                       -- this variable determines the parallel context; if
 115                       -- `Nothing', we are in pure vectorisation mode, no
 116                       -- lifting going on
 117                       --
 118                       ctxtVar  :: Maybe Var,
 119
 120                       -- environment that maps each variable that is
 121                       -- vectorised in the current parallel context to the
 122                       -- vectorised version of that variable
 123                       --
 124                       ctxtEnv :: VarEnv Var,
 125
 126                       -- those variables from the *domain* of `ctxtEnv' that
 127                       -- have been used since the last context restriction (cf.
 128                       -- `restrictContext')
 129                       --
 130                       usedVars :: VarSet
 131                     }
 132
 133 -- initial value of the flattening state
 134 --
 135 initialFlattenState :: PersistentCompilerState
 136                     -> HomeSymbolTable
 137                     -> UniqSupply
 138                     -> FlattenState
 139 initialFlattenState pcs hst us =
 140   FlattenState {
 141     us       = us,
 142     env      = lookup,
 143     ctxtVar  = Nothing,
 144     ctxtEnv  = emptyVarEnv,
 145     usedVars = emptyVarSet
 146   }
 147   where
 148     lookup n =
 149       case lookupType hst (pcs_PTE pcs) n of
 150         Just (AnId v) -> v
 151         _             -> pprPanic "FlattenMonad: unknown name:" (ppr n)
 152
 153 -- the monad representation (EXPORTED ABSTRACTLY)
 154 --
 155 newtype Flatten a = Flatten {
 156                       unFlatten :: (FlattenState -> (a, FlattenState))
 157                     }
 158
 159 instance Monad Flatten where
 160   return x = Flatten $ \s -> (x, s)
 161   m >>= n  = Flatten $ \s -> let
 162                                (r, s') = unFlatten m s
 163                              in
 164                              unFlatten (n r) s'
 165
 166 -- execute the given flattening computation (EXPORTED)
 167 --
 168 runFlatten :: PersistentCompilerState
 169            -> HomeSymbolTable
 170            -> UniqSupply
 171            -> Flatten a
 172            -> a
 173 runFlatten pcs hst us m = fst $ unFlatten m (initialFlattenState pcs hst us)
 174
 175
 176 -- variable generation
 177 -- -------------------
 178
 179 -- generate a new local variable whose name is based on the given lexeme and
 180 -- whose type is as specified in the second argument (EXPORTED)
 181 --
 182 newVar           :: UserFS -> Type -> Flatten Var
 183 newVar lexeme ty  = Flatten $ \state ->
 184   let
 185     (us1, us2) = splitUniqSupply (us state)
 186     state'     = state {us = us2}
 187   in
 188   (mkSysLocal lexeme (uniqFromSupply us1) ty, state')
 189
 190 -- generate a non-recursive binding using a new binder whose name is derived
 191 -- from the given lexeme (EXPORTED)
 192 --
 193 mkBind          :: UserFS -> CoreExpr -> Flatten (CoreBndr, CoreBind)
 194 mkBind lexeme e  =
 195   do
 196     v <- newVar lexeme (exprType e)
 197     return (v, NonRec v e)
 198
 199
 200 -- context management
 201 -- ------------------
 202
 203 -- extend the parallel context by the given set of variables (EXPORTED)
 204 --
 205 -- * if there is no parallel context at the moment, the first element of the
 206 --   variable list will be used to determine the new parallel context
 207 --
 208 -- * the second argument is executed in the current context extended with the
 209 --   given variables
 210 --
 211 -- * the variables must already have been lifted by transforming their type,
 212 --   but they *must* have retained their original name (or, at least, their
 213 --   unique); this is needed so that they match the original variable in
 214 --   variable environments
 215 --
 216 -- * any trace of the given set of variables has to be removed from the state
 217 --   at the end of this operation
 218 --
 219 extendContext      :: [Var] -> Flatten a -> Flatten a
 220 extendContext [] m  = m
 221 extendContext vs m  = Flatten $ \state ->
 222   let
 223     extState       = state {
 224                        ctxtVar = ctxtVar state `mplus` Just (head vs),
 225                        ctxtEnv = ctxtEnv state `plusVarEnv` zipVarEnv vs vs
 226                      }
 227     (r, extState') = unFlatten m extState
 228     resState       = extState' { -- remove `vs' from the result state
 229                        ctxtVar  = ctxtVar state,
 230                        ctxtEnv  = ctxtEnv state,
 231                        usedVars = usedVars extState' `delVarEnvList` vs
 232                      }
 233   in
 234   (r, resState)
 235
 236 -- execute the second argument in a restricted context (EXPORTED)
 237 --
 238 -- * all variables in the current parallel context are packed according to
 239 --   the permutation vector associated with the variable passed as the first
 240 --   argument (ie, all elements of vectorised context variables that are
 241 --   invalid in the restricted context are dropped)
 242 --
 243 -- * the returned list of core binders contains the operations that perform
 244 --   the restriction on all variables in the parallel context that *do* occur
 245 --   during the execution of the second argument (ie, `liftVar' is executed at
 246 --   least once on any such variable)
 247 --
 248 packContext        :: Var -> Flatten a -> Flatten (a, [CoreBind])
 249 packContext perm m  = Flatten $ \state ->
 250   let
 251     -- FIXME: To set the packed environment to the unpacked on is a hack of
 252     --   which I am not sure yet (a) whether it works and (b) whether it's
 253     --   really worth it.  The one advantages is that, we can use a var set,
 254     --   after all, instead of a var environment.
 255     --
 256     --   The idea is the following: If we have to pack a variable `x', we
 257     --   generate `let{-NonRec-} x = bpermuteP perm x in ...'.  As this is a
 258     --   non-recursive binding, the lhs `x' overshadows the rhs `x' in the
 259     --   body of the let.
 260     --
 261     --   NB: If we leave it like this, `mkCoreBind' can be simplified.
 262     packedCtxtEnv     = ctxtEnv state
 263     packedState       = state {
 264                           ctxtVar  = fmap
 265                                        (lookupVarEnv_NF packedCtxtEnv)
 266                                        (ctxtVar state),
 267                           ctxtEnv  = packedCtxtEnv,
 268                           usedVars = emptyVarSet
 269                         }
 270     (r, packedState') = unFlatten m packedState
 271     resState          = state {    -- revert to the unpacked context
 272                           ctxtVar  = ctxtVar state,
 273                           ctxtEnv  = ctxtEnv state,
 274                         }
 275     bndrs             = map mkCoreBind . varSetElems . usedVars $ packedState'
 276
 277     -- generate a binding for the packed variant of a context variable
 278     --
 279     mkCoreBind var    = let
 280                           rhs = fst $ unFlatten (mk'bpermuteP (varType var)
 281                                                               (Var perm)
 282                                                               (Var var)
 283                                                 ) state
 284                         in
 285                         NonRec (lookupVarEnv_NF packedCtxtEnv var) $ rhs
 286
 287   in
 288   ((r, bndrs), resState)
 289
 290 -- lift a single variable in the current context (EXPORTED)
 291 --
 292 -- * if the variable does not occur in the context, it's value is vectorised to
 293 --   match the size of the current context
 294 --
 295 -- * otherwise, the variable is replaced by whatever the context environment
 296 --   maps it to (this may either be simply the lifted version of the original
 297 --   variable or a packed variant of that variable)
 298 --
 299 -- * the monad keeps track of all lifted variables that occur in the parallel
 300 --   context, so that `packContext' can determine the correct set of core
 301 --   bindings
 302 --
 303 liftVar     :: Var -> Flatten CoreExpr
 304 liftVar var  = Flatten $ \s ->
 305   let
 306     v          = ctxtVarErr s
 307     v'elemType = parrElemTy . varType $ v
 308     len        = fst $ unFlatten (mk'lengthP v'elemType (Var v)) s
 309     replicated = fst $ unFlatten (mk'replicateP (varType var) len (Var var)) s
 310   in case lookupVarEnv (ctxtEnv s) var of
 311     Just liftedVar -> (Var liftedVar,
 312                        s {usedVars = usedVars s `extendVarSet` var})
 313     Nothing        -> (replicated, s)
 314
 315 -- lift a constant expression in the current context (EXPORTED)
 316 --
 317 -- * the value of the constant expression is vectorised to match the current
 318 --   parallel context
 319 --
 320 liftConst   :: CoreExpr -> Flatten CoreExpr
 321 liftConst e  = Flatten $ \s ->
 322   let
 323      v          = ctxtVarErr s
 324      v'elemType = parrElemTy . varType $ v
 325      len        = fst $ unFlatten (mk'lengthP v'elemType (Var v)) s
 326   in
 327   (fst $ unFlatten (mk'replicateP (exprType e) len e ) s, s)
 328
 329 -- pick those variables of the given set that occur (if albeit in lifted form)
 330 -- in the current parallel context (EXPORTED)
 331 --
 332 -- * the variables returned are from the given set and *not* the corresponding
 333 --   context variables
 334 --
 335 intersectWithContext    :: VarSet -> Flatten [Var]
 336 intersectWithContext vs  = Flatten $ \s ->
 337   let
 338     vs' = filter (`elemVarEnv` ctxtEnv s) (varSetElems vs)
 339   in
 340   (vs', s)
 341
 342
 343 -- construct applications of prelude functions
 344 -- -------------------------------------------
 345
 346 -- NB: keep all the used names listed in `FlattenInfo.namesNeededForFlattening'
 347
 348 -- generate an application of `fst' (EXPORTED)
 349 --
 350 mk'fst           :: Type -> Type -> CoreExpr -> Flatten CoreExpr
 351 mk'fst ty1 ty2 a  = mkFunApp fstName [Type ty1, Type ty2, a]
 352
 353 -- generate an application of `&&' (EXPORTED)
 354 --
 355 mk'and       :: CoreExpr -> CoreExpr -> Flatten CoreExpr
 356 mk'and a1 a2  = mkFunApp andName [a1, a2]
 357
 358 -- generate an application of `||' (EXPORTED)
 359 --
 360 mk'or       :: CoreExpr -> CoreExpr -> Flatten CoreExpr
 361 mk'or a1 a2  = mkFunApp orName [a1, a2]
 362
 363 -- generate an application of `==' where the arguments may only be literals
 364 -- that may occur in a Core case expression (i.e., `Char', `Int', `Float', and
 365 -- `Double') (EXPORTED)
 366 --
 367 mk'eq          :: Type -> CoreExpr -> CoreExpr -> Flatten CoreExpr
 368 mk'eq ty a1 a2  = mkFunApp eqName [a1, a2]
 369                   where
 370                     name = tyConName . tyConAppTyCon $ ty
 371                     --
 372                     eqName | name == charPrimTyConName   = eqCharName
 373                            | name == intPrimTyConName    = eqIntName
 374                            | name == floatPrimTyConName  = eqFloatName
 375                            | name == doublePrimTyConName = eqDoubleName
 376                            | otherwise                   =
 377                              pprPanic "FlattenMonad.mk'eq: " (ppr ty)
 378
 379 -- generate an application of `==' where the arguments may only be literals
 380 -- that may occur in a Core case expression (i.e., `Char', `Int', `Float', and
 381 -- `Double') (EXPORTED)
 382 --
 383 mk'neq          :: Type -> CoreExpr -> CoreExpr -> Flatten CoreExpr
 384 mk'neq ty a1 a2  = mkFunApp neqName [a1, a2]
 385                    where
 386                      name = tyConName . tyConAppTyCon $ ty
 387                      --
 388                      neqName {- | name == charPrimTyConName   = neqCharName -}
 389                              | name == intPrimTyConName    = neqIntName
 390                              {- | name == floatPrimTyConName  = neqFloatName -}
 391                              {- | name == doublePrimTyConName = neqDoubleName -}
 392                              | otherwise                   =
 393                                pprPanic "FlattenMonad.mk'neq: " (ppr ty)
 394
 395 -- generate an application of `lengthP' (EXPORTED)
 396 --
 397 mk'lengthP      :: Type -> CoreExpr -> Flatten CoreExpr
 398 mk'lengthP ty a  = mkFunApp lengthPName [Type ty, a]
 399
 400 -- generate an application of `replicateP' (EXPORTED)
 401 --
 402 mk'replicateP          :: Type -> CoreExpr -> CoreExpr -> Flatten CoreExpr
 403 mk'replicateP ty a1 a2  = mkFunApp replicatePName [Type ty, a1, a2]
 404
 405 -- generate an application of `replicateP' (EXPORTED)
 406 --
 407 mk'mapP :: Type -> Type -> CoreExpr -> CoreExpr -> Flatten CoreExpr
 408 mk'mapP ty1 ty2 a1 a2  = mkFunApp mapPName [Type ty1, Type ty2, a1, a2]
 409
 410 -- generate an application of `bpermuteP' (EXPORTED)
 411 --
 412 mk'bpermuteP          :: Type -> CoreExpr -> CoreExpr -> Flatten CoreExpr
 413 mk'bpermuteP ty a1 a2  = mkFunApp bpermutePName [Type ty, a1, a2]
 414
 415 -- generate an application of `bpermuteDftP' (EXPORTED)
 416 --
 417 mk'bpermuteDftP :: Type -> CoreExpr -> CoreExpr -> CoreExpr -> Flatten CoreExpr
 418 mk'bpermuteDftP ty a1 a2 a3 = mkFunApp bpermuteDftPName [Type ty, a1, a2, a3]
 419
 420 -- generate an application of `indexOfP' (EXPORTED)
 421 --
 422 mk'indexOfP          :: Type -> CoreExpr -> CoreExpr -> Flatten CoreExpr
 423 mk'indexOfP ty a1 a2  = mkFunApp indexOfPName [Type ty, a1, a2]
 424
 425
 426 -- auxilliary functions
 427 -- --------------------
 428
 429 -- obtain the context variable, aborting if it is not available (as this
 430 -- signals an internal error in the usage of the `Flatten' monad)
 431 --
 432 ctxtVarErr   :: FlattenState -> Var
 433 ctxtVarErr s  = case ctxtVar s of
 434                   Nothing -> panic "FlattenMonad.ctxtVarErr: No context \
 435                                    \variable available!"
 436                   Just v  -> v
 437
 438 -- given the name of a known function and a set of arguments (needs to include
 439 -- all needed type arguments), build a Core expression that applies the named
 440 -- function to those arguments
 441 --
 442 mkFunApp           :: Name -> [CoreExpr] -> Flatten CoreExpr
 443 mkFunApp name args  =
 444   do
 445     fun <- lookupName name
 446     return $ mkApps (Var fun) args
 447
 448 -- get the `Id' of a known `Name'
 449 --
 450 -- * this can be the `Name' of any function that's visible on the toplevel of
 451 --   the current compilation unit
 452 --
 453 lookupName      :: Name -> Flatten Id
 454 lookupName name  = Flatten $ \s ->
 455   (env s name, s)