Vectorise.hs 14.4 KB
Newer Older
1 2 3 4 5 6 7 8
module Vectorise( vectorise )
where

#include "HsVersions.h"

import DynFlags
import HscTypes

rl@cse.unsw.edu.au's avatar
rl@cse.unsw.edu.au committed
9
import CoreLint             ( showPass, endPass )
10
import CoreSyn
11 12
import CoreUtils
import CoreFVs
rl@cse.unsw.edu.au's avatar
rl@cse.unsw.edu.au committed
13
import TyCon
rl@cse.unsw.edu.au's avatar
rl@cse.unsw.edu.au committed
14 15
import Type
import TypeRep
rl@cse.unsw.edu.au's avatar
rl@cse.unsw.edu.au committed
16 17
import Var
import VarEnv
rl@cse.unsw.edu.au's avatar
rl@cse.unsw.edu.au committed
18
import Name                 ( mkSysTvName )
19
import NameEnv
rl@cse.unsw.edu.au's avatar
rl@cse.unsw.edu.au committed
20
import Id
rl@cse.unsw.edu.au's avatar
rl@cse.unsw.edu.au committed
21

rl@cse.unsw.edu.au's avatar
rl@cse.unsw.edu.au committed
22
import DsMonad hiding (mapAndUnzipM)
rl@cse.unsw.edu.au's avatar
rl@cse.unsw.edu.au committed
23 24 25

import PrelNames

26
import Outputable
rl@cse.unsw.edu.au's avatar
rl@cse.unsw.edu.au committed
27
import FastString
rl@cse.unsw.edu.au's avatar
rl@cse.unsw.edu.au committed
28
import Control.Monad        ( liftM, liftM2, mapAndUnzipM )
29

30 31 32
vectorise :: HscEnv -> ModGuts -> IO ModGuts
vectorise hsc_env guts
  | not (Opt_Vectorise `dopt` dflags) = return guts
rl@cse.unsw.edu.au's avatar
rl@cse.unsw.edu.au committed
33 34 35 36 37 38 39 40 41 42 43
  | otherwise
  = do
      showPass dflags "Vectorisation"
      eps <- hscEPS hsc_env
      let info = hptVectInfo hsc_env `plusVectInfo` eps_vect_info eps
      Just guts' <- initDs hsc_env (mg_module guts)
                                   (mg_rdr_env guts)
                                   (mg_types guts)
                                   (vectoriseModule info guts)
      endPass dflags "Vectorisation" Opt_D_dump_vect (mg_binds guts')
      return guts'
44 45 46
  where
    dflags = hsc_dflags hsc_env

rl@cse.unsw.edu.au's avatar
rl@cse.unsw.edu.au committed
47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90
-- ----------------------------------------------------------------------------
-- Vectorisation monad

data Builtins = Builtins {
                  parrayTyCon      :: TyCon
                , paTyCon          :: TyCon
                , closureTyCon     :: TyCon
                , mkClosureVar     :: Var
                , applyClosureVar  :: Var
                , mkClosurePVar    :: Var
                , applyClosurePVar :: Var
                , closurePAVar     :: Var
                , lengthPAVar      :: Var
                , replicatePAVar   :: Var
                }

initBuiltins :: DsM Builtins
initBuiltins
  = do
      parrayTyCon  <- dsLookupTyCon parrayTyConName
      paTyCon      <- dsLookupTyCon paTyConName
      closureTyCon <- dsLookupTyCon closureTyConName

      mkClosureVar     <- dsLookupGlobalId mkClosureName
      applyClosureVar  <- dsLookupGlobalId applyClosureName
      mkClosurePVar    <- dsLookupGlobalId mkClosurePName
      applyClosurePVar <- dsLookupGlobalId applyClosurePName
      closurePAVar     <- dsLookupGlobalId closurePAName
      lengthPAVar      <- dsLookupGlobalId lengthPAName
      replicatePAVar   <- dsLookupGlobalId replicatePAName

      return $ Builtins {
                 parrayTyCon      = parrayTyCon
               , paTyCon          = paTyCon
               , closureTyCon     = closureTyCon
               , mkClosureVar     = mkClosureVar
               , applyClosureVar  = applyClosureVar
               , mkClosurePVar    = mkClosurePVar
               , applyClosurePVar = applyClosurePVar
               , closurePAVar     = closurePAVar
               , lengthPAVar      = lengthPAVar
               , replicatePAVar   = replicatePAVar
               }

91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109
data GlobalEnv = GlobalEnv {
                  -- Mapping from global variables to their vectorised versions.
                  -- 
                  global_vars :: VarEnv CoreExpr

                  -- Exported variables which have a vectorised version
                  --
                , global_exported_vars :: VarEnv (Var, Var)

                  -- Mapping from TyCons to their vectorised versions.
                  -- TyCons which do not have to be vectorised are mapped to
                  -- themselves.
                  --
                , global_tycons :: NameEnv TyCon

                  -- Mapping from TyCons to their PA dictionaries
                  --
                , global_tycon_pa :: NameEnv CoreExpr
                }
rl@cse.unsw.edu.au's avatar
rl@cse.unsw.edu.au committed
110

111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134
data LocalEnv = LocalEnv {
                 -- Mapping from local variables to their vectorised and
                 -- lifted versions
                 --
                 local_vars :: VarEnv (CoreExpr, CoreExpr)

                 -- Mapping from tyvars to their PA dictionaries
               , local_tyvar_pa :: VarEnv CoreExpr
               }
              

initGlobalEnv :: VectInfo -> GlobalEnv
initGlobalEnv info
  = GlobalEnv {
      global_vars          = mapVarEnv  (Var . snd) $ vectInfoCCVar   info
    , global_exported_vars = emptyVarEnv
    , global_tycons        = mapNameEnv snd $ vectInfoCCTyCon info
    , global_tycon_pa      = emptyNameEnv
    }

emptyLocalEnv = LocalEnv {
                   local_vars     = emptyVarEnv
                 , local_tyvar_pa = emptyVarEnv
                 }
rl@cse.unsw.edu.au's avatar
rl@cse.unsw.edu.au committed
135 136

-- FIXME
137
updVectInfo :: GlobalEnv -> ModGuts -> ModGuts
138 139 140
updVectInfo env guts = guts { mg_vect_info = info' }
  where
    info' = info {
141
              vectInfoCCVar   = global_exported_vars env
142 143 144 145 146 147 148 149
            , vectInfoCCTyCon = tc_env
            }

    info  = mg_vect_info guts
    tyenv = mg_types guts

    tc_env = mkNameEnv [(tc_name, (tc,tc')) | tc <- typeEnvTyCons tyenv
                                            , let tc_name = tyConName tc
150
                                            , Just tc' <- [lookupNameEnv (global_tycons env) tc_name]]
rl@cse.unsw.edu.au's avatar
rl@cse.unsw.edu.au committed
151

152
data VResult a = Yes GlobalEnv LocalEnv a | No
153

154
newtype VM a = VM { runVM :: Builtins -> GlobalEnv -> LocalEnv -> DsM (VResult a) }
rl@cse.unsw.edu.au's avatar
rl@cse.unsw.edu.au committed
155 156

instance Monad VM where
157 158 159 160 161 162
  return x   = VM $ \bi genv lenv -> return (Yes genv lenv x)
  VM p >>= f = VM $ \bi genv lenv -> do
                                      r <- p bi genv lenv
                                      case r of
                                        Yes genv' lenv' x -> runVM (f x) bi genv' lenv'
                                        No                -> return No
163 164

noV :: VM a
165
noV = VM $ \_ _ _ -> return No
166 167

tryV :: VM a -> VM (Maybe a)
168 169 170 171 172 173
tryV (VM p) = VM $ \bi genv lenv ->
  do
    r <- p bi genv lenv
    case r of
      Yes genv' lenv' x -> return (Yes genv' lenv' (Just x))
      No                -> return (Yes genv  lenv  Nothing)
rl@cse.unsw.edu.au's avatar
rl@cse.unsw.edu.au committed
174

175 176 177
maybeV :: VM (Maybe a) -> VM a
maybeV p = maybe noV return =<< p

178 179 180
orElseV :: VM a -> VM a -> VM a
orElseV p q = maybe q return =<< tryV p

rl@cse.unsw.edu.au's avatar
rl@cse.unsw.edu.au committed
181 182 183 184 185 186 187
localV :: VM a -> VM a
localV p = do
             env <- readLEnv id
             x <- p
             setLEnv env
             return x

rl@cse.unsw.edu.au's avatar
rl@cse.unsw.edu.au committed
188
liftDs :: DsM a -> VM a
189
liftDs p = VM $ \bi genv lenv -> do { x <- p; return (Yes genv lenv x) }
rl@cse.unsw.edu.au's avatar
rl@cse.unsw.edu.au committed
190

191
builtin :: (Builtins -> a) -> VM a
192 193 194 195 196 197 198 199 200 201
builtin f = VM $ \bi genv lenv -> return (Yes genv lenv (f bi))

readGEnv :: (GlobalEnv -> a) -> VM a
readGEnv f = VM $ \bi genv lenv -> return (Yes genv lenv (f genv))

setGEnv :: GlobalEnv -> VM ()
setGEnv genv = VM $ \_ _ lenv -> return (Yes genv lenv ())

updGEnv :: (GlobalEnv -> GlobalEnv) -> VM ()
updGEnv f = VM $ \_ genv lenv -> return (Yes (f genv) lenv ())
202

203 204
readLEnv :: (LocalEnv -> a) -> VM a
readLEnv f = VM $ \bi genv lenv -> return (Yes genv lenv (f lenv))
205

206 207
setLEnv :: LocalEnv -> VM ()
setLEnv lenv = VM $ \_ genv _ -> return (Yes genv lenv ())
208

209 210
updLEnv :: (LocalEnv -> LocalEnv) -> VM ()
updLEnv f = VM $ \_ genv lenv -> return (Yes genv (f lenv) ())
211

212 213 214 215 216 217
newLocalVar :: FastString -> Type -> VM Var
newLocalVar fs ty
  = do
      u <- liftDs newUnique
      return $ mkSysLocal fs u ty

rl@cse.unsw.edu.au's avatar
rl@cse.unsw.edu.au committed
218 219 220 221 222
newTyVar :: FastString -> Kind -> VM Var
newTyVar fs k
  = do
      u <- liftDs newUnique
      return $ mkTyVar (mkSysTvName u fs) k
223 224

lookupTyCon :: TyCon -> VM (Maybe TyCon)
225
lookupTyCon tc = readGEnv $ \env -> lookupNameEnv (global_tycons env) (tyConName tc)
226

227 228 229 230

extendTyVarPA :: Var -> CoreExpr -> VM ()
extendTyVarPA tv pa = updLEnv $ \env -> env { local_tyvar_pa = extendVarEnv (local_tyvar_pa env) tv pa }

231 232 233
-- ----------------------------------------------------------------------------
-- Bindings

rl@cse.unsw.edu.au's avatar
rl@cse.unsw.edu.au committed
234 235 236 237
vectoriseModule :: VectInfo -> ModGuts -> DsM ModGuts
vectoriseModule info guts
  = do
      builtins <- initBuiltins
238
      r <- runVM (vectModule guts) builtins (initGlobalEnv info) emptyLocalEnv
239
      case r of
240 241
        Yes genv _ guts' -> return $ updVectInfo genv guts'
        No               -> return guts
rl@cse.unsw.edu.au's avatar
rl@cse.unsw.edu.au committed
242 243 244 245

vectModule :: ModGuts -> VM ModGuts
vectModule guts = return guts

rl@cse.unsw.edu.au's avatar
rl@cse.unsw.edu.au committed
246 247 248 249 250 251 252 253 254 255 256 257 258 259 260 261 262 263 264 265 266 267 268 269 270 271 272 273 274 275


vectBndr :: Var -> VM (Var, Var)
vectBndr v
  = do
      vty <- vectType (idType v)
      lty <- mkPArrayTy vty
      let vv = v `Id.setIdType` vty
          lv = v `Id.setIdType` lty
      updLEnv (mapTo vv lv)
      return (vv, lv)
  where
    mapTo vv lv env = env { local_vars = extendVarEnv (local_vars env) v (Var vv, Var lv) }

vectBndrIn :: Var -> VM a -> VM (Var, Var, a)
vectBndrIn v p
  = localV
  $ do
      (vv, lv) <- vectBndr v
      x <- p
      return (vv, lv, x)

vectBndrsIn :: [Var] -> VM a -> VM ([Var], [Var], a)
vectBndrsIn vs p
  = localV
  $ do
      (vvs, lvs) <- mapAndUnzipM vectBndr vs
      x <- p
      return (vvs, lvs, x)

rl@cse.unsw.edu.au's avatar
rl@cse.unsw.edu.au committed
276
-- ----------------------------------------------------------------------------
277 278 279 280 281 282 283 284 285 286 287 288 289 290 291 292 293 294 295 296 297 298 299 300 301
-- Expressions

replicateP :: CoreExpr -> CoreExpr -> VM CoreExpr
replicateP expr len
  = do
      pa  <- paOfType ty
      rep <- builtin replicatePAVar
      return $ mkApps (Var rep) [Type ty, pa, expr, len]
  where
    ty = exprType expr

capply :: (CoreExpr, CoreExpr) -> (CoreExpr, CoreExpr) -> VM (CoreExpr, CoreExpr)
capply (vfn, lfn) (varg, larg)
  = do
      apply  <- builtin applyClosureVar
      applyP <- builtin applyClosurePVar
      return (mkApps (Var apply)  [Type arg_ty, Type res_ty, vfn, varg],
              mkApps (Var applyP) [Type arg_ty, Type res_ty, lfn, larg])
  where
    fn_ty            = exprType vfn
    (arg_ty, res_ty) = splitClosureTy fn_ty

vectVar :: CoreExpr -> Var -> VM (CoreExpr, CoreExpr)
vectVar lc v = local v `orElseV` global v
  where
302
    local  v = maybeV (readLEnv $ \env -> lookupVarEnv (local_vars env) v)
303
    global v = do
304
                 vexpr <- maybeV (readGEnv $ \env -> lookupVarEnv (global_vars env) v)
305 306 307 308 309 310 311 312 313 314 315 316 317 318 319 320 321 322 323 324 325 326 327
                 lexpr <- replicateP vexpr lc
                 return (vexpr, lexpr)
                
vectExpr :: CoreExpr -> CoreExprWithFVs -> VM (CoreExpr, CoreExpr)
vectExpr lc (_, AnnType ty)
  = do
      vty <- vectType ty
      return (Type vty, Type vty)
vectExpr lc (_, AnnVar v)   = vectVar lc v
vectExpr lc (_, AnnLit lit)
  = do
      let vexpr = Lit lit
      lexpr <- replicateP vexpr lc
      return (vexpr, lexpr)
vectExpr lc (_, AnnNote note expr)
  = do
      (vexpr, lexpr) <- vectExpr lc expr
      return (Note note vexpr, Note note lexpr)
vectExpr lc (_, AnnApp fn arg)
  = do
      fn'  <- vectExpr lc fn
      arg' <- vectExpr lc arg
      capply fn' arg'
rl@cse.unsw.edu.au's avatar
rl@cse.unsw.edu.au committed
328 329 330 331 332 333 334 335 336 337 338 339 340 341 342 343 344 345 346 347
vectExpr lc (_, AnnCase expr bndr ty alts)
  = panic "vectExpr: case"
vectExpr lc (_, AnnLet (AnnNonRec bndr rhs) body)
  = do
      (vrhs, lrhs) <- vectExpr lc rhs
      (vbndr, lbndr, (vbody, lbody)) <- vectBndrIn bndr (vectExpr lc body)
      return (Let (NonRec vbndr vrhs) vbody,
              Let (NonRec lbndr lrhs) lbody)
vectExpr lc (_, AnnLet (AnnRec prs) body)
  = do
      (vbndrs, lbndrs, (vrhss, vbody, lrhss, lbody)) <- vectBndrsIn bndrs vect
      return (Let (Rec (zip vbndrs vrhss)) vbody,
              Let (Rec (zip lbndrs lrhss)) lbody)
  where
    (bndrs, rhss) = unzip prs
    
    vect = do
             (vrhss, lrhss) <- mapAndUnzipM (vectExpr lc) rhss
             (vbody, lbody) <- vectExpr lc body
             return (vrhss, vbody, lrhss, lbody)
348 349 350 351 352 353 354 355 356 357 358 359
vectExpr lc (_, AnnLam bndr body)
  | isTyVar bndr
  = do
      pa_ty          <- paArgType' (TyVarTy bndr) (tyVarKind bndr)
      pa_var         <- newLocalVar FSLIT("dPA") pa_ty
      (vbody, lbody) <- localV
                      $ do
                          extendTyVarPA bndr (Var pa_var)
                          -- FIXME: what about shadowing here (bndr in lc)?
                          vectExpr lc body
      return (mkLams [bndr, pa_var] vbody,
              mkLams [bndr, pa_var] lbody)
360 361 362

-- ----------------------------------------------------------------------------
-- PA dictionaries
rl@cse.unsw.edu.au's avatar
rl@cse.unsw.edu.au committed
363 364 365 366 367 368 369 370 371 372 373 374 375 376 377 378 379 380 381 382 383 384 385 386 387 388 389 390 391 392 393

paArgType :: Type -> Kind -> VM (Maybe Type)
paArgType ty k
  | Just k' <- kindView k = paArgType ty k'

-- Here, we assume that for a kind (k1 -> k2) to be valid, k1 and k2 can only
-- be made up of * and (->), i.e., they can't be coercion kinds or #.
paArgType ty (FunTy k1 k2)
  = do
      tv  <- newTyVar FSLIT("a") k1
      ty1 <- paArgType' (TyVarTy tv) k1
      ty2 <- paArgType' (AppTy ty (TyVarTy tv)) k2
      return . Just $ ForAllTy tv (FunTy ty1 ty2)

paArgType ty k
  | isLiftedTypeKind k
  = do
      tc <- builtin paTyCon
      return . Just $ TyConApp tc [ty]

  | otherwise
  = return Nothing 

paArgType' :: Type -> Kind -> VM Type
paArgType' ty k
  = do
      r <- paArgType ty k
      case r of
        Just ty' -> return ty'
        Nothing  -> pprPanic "paArgType'" (ppr ty)

394 395
paOfTyCon :: TyCon -> VM CoreExpr
-- FIXME: just for now
396
paOfTyCon tc = maybeV (readGEnv $ \env -> lookupNameEnv (global_tycon_pa env) (tyConName tc))
397 398 399 400

paOfType :: Type -> VM CoreExpr
paOfType ty | Just ty' <- coreView ty = paOfType ty'

401
paOfType (TyVarTy tv) = maybeV (readLEnv $ \env -> lookupVarEnv (local_tyvar_pa env) tv)
402 403 404 405 406 407 408 409 410 411 412 413 414 415 416 417 418 419 420
paOfType (AppTy ty1 ty2)
  = do
      e1 <- paOfType ty1
      e2 <- paOfType ty2
      return $ mkApps e1 [Type ty2, e2]
paOfType (TyConApp tc tys)
  = do
      e  <- paOfTyCon tc
      es <- mapM paOfType tys
      return $ mkApps e [arg | (t,e) <- zip tys es, arg <- [Type t, e]]
paOfType (FunTy ty1 ty2) = paOfType (TyConApp funTyCon [ty1,ty2])
paOfType t@(ForAllTy tv ty) = pprPanic "paOfType:" (ppr t)
paOfType ty = pprPanic "paOfType:" (ppr ty)
        


-- ----------------------------------------------------------------------------
-- Types

rl@cse.unsw.edu.au's avatar
rl@cse.unsw.edu.au committed
421 422
vectTyCon :: TyCon -> VM TyCon
vectTyCon tc
423 424 425
  | isFunTyCon tc        = builtin closureTyCon
  | isBoxedTupleTyCon tc = return tc
  | isUnLiftedTyCon tc   = return tc
rl@cse.unsw.edu.au's avatar
rl@cse.unsw.edu.au committed
426 427 428 429 430 431 432 433 434 435 436 437 438 439 440 441 442 443 444 445 446 447 448
  | otherwise = do
                  r <- lookupTyCon tc
                  case r of
                    Just tc' -> return tc'

                    -- FIXME: just for now
                    Nothing  -> pprTrace "ccTyCon:" (ppr tc) $ return tc

vectType :: Type -> VM Type
vectType ty | Just ty' <- coreView ty = vectType ty
vectType (TyVarTy tv) = return $ TyVarTy tv
vectType (AppTy ty1 ty2) = liftM2 AppTy (vectType ty1) (vectType ty2)
vectType (TyConApp tc tys) = liftM2 TyConApp (vectTyCon tc) (mapM vectType tys)
vectType (FunTy ty1 ty2)   = liftM2 TyConApp (builtin closureTyCon)
                                             (mapM vectType [ty1,ty2])
vectType (ForAllTy tv ty)
  = do
      r   <- paArgType (TyVarTy tv) (tyVarKind tv)
      ty' <- vectType ty
      return . ForAllTy tv $ case r of { Just paty -> FunTy paty ty'; Nothing -> ty' }

vectType ty = pprPanic "vectType:" (ppr ty)

449 450 451 452 453 454 455 456 457 458 459
isClosureTyCon :: TyCon -> Bool
isClosureTyCon tc = tyConUnique tc == closureTyConKey

splitClosureTy :: Type -> (Type, Type)
splitClosureTy ty
  | Just (tc, [arg_ty, res_ty]) <- splitTyConApp_maybe ty
  , isClosureTyCon tc
  = (arg_ty, res_ty)

  | otherwise = pprPanic "splitClosureTy" (ppr ty)

rl@cse.unsw.edu.au's avatar
rl@cse.unsw.edu.au committed
460 461 462 463 464
mkPArrayTy :: Type -> VM Type
mkPArrayTy ty = do
                  tc <- builtin parrayTyCon
                  return $ TyConApp tc [ty]