Vectorise.hs 13.6 KB
Newer Older
1 2 3 4 5
module Vectorise( vectorise )
where

#include "HsVersions.h"

6
import VectMonad
7
import VectUtils
8

9 10 11
import DynFlags
import HscTypes

rl@cse.unsw.edu.au's avatar
rl@cse.unsw.edu.au committed
12
import CoreLint             ( showPass, endPass )
13
import CoreSyn
14 15
import CoreUtils
import CoreFVs
16
import DataCon
rl@cse.unsw.edu.au's avatar
rl@cse.unsw.edu.au committed
17
import TyCon
rl@cse.unsw.edu.au's avatar
rl@cse.unsw.edu.au committed
18 19
import Type
import TypeRep
rl@cse.unsw.edu.au's avatar
rl@cse.unsw.edu.au committed
20 21
import Var
import VarEnv
22
import VarSet
23
import Name                 ( mkSysTvName, getName )
24
import NameEnv
rl@cse.unsw.edu.au's avatar
rl@cse.unsw.edu.au committed
25
import Id
26
import MkId                 ( unwrapFamInstScrut )
27
import OccName
rl@cse.unsw.edu.au's avatar
rl@cse.unsw.edu.au committed
28

rl@cse.unsw.edu.au's avatar
rl@cse.unsw.edu.au committed
29
import DsMonad hiding (mapAndUnzipM)
30
import DsUtils              ( mkCoreTup, mkCoreTupTy )
rl@cse.unsw.edu.au's avatar
rl@cse.unsw.edu.au committed
31 32

import PrelNames
33 34
import TysWiredIn
import BasicTypes           ( Boxity(..) )
rl@cse.unsw.edu.au's avatar
rl@cse.unsw.edu.au committed
35

36
import Outputable
rl@cse.unsw.edu.au's avatar
rl@cse.unsw.edu.au committed
37
import FastString
38 39
import Control.Monad        ( liftM, liftM2, mapAndUnzipM, zipWithM_ )
import Data.Maybe           ( maybeToList )
40

41 42 43
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
44 45 46 47 48
  | otherwise
  = do
      showPass dflags "Vectorisation"
      eps <- hscEPS hsc_env
      let info = hptVectInfo hsc_env `plusVectInfo` eps_vect_info eps
49
      Just (info', guts') <- initV hsc_env guts info (vectModule guts)
rl@cse.unsw.edu.au's avatar
rl@cse.unsw.edu.au committed
50
      endPass dflags "Vectorisation" Opt_D_dump_vect (mg_binds guts')
51
      return $ guts' { mg_vect_info = info' }
52 53 54
  where
    dflags = hsc_dflags hsc_env

rl@cse.unsw.edu.au's avatar
rl@cse.unsw.edu.au committed
55 56 57
vectModule :: ModGuts -> VM ModGuts
vectModule guts = return guts

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 91 92
vectTopBind b@(NonRec var expr)
  = do
      var'  <- vectTopBinder var
      expr' <- vectTopRhs expr
      hs    <- takeHoisted
      return . Rec $ (var, expr) : (var', expr') : hs
  `orElseV`
    return b

vectTopBind b@(Rec bs)
  = do
      vars'  <- mapM vectTopBinder vars
      exprs' <- mapM vectTopRhs exprs
      hs     <- takeHoisted
      return . Rec $ bs ++ zip vars' exprs' ++ hs
  `orElseV`
    return b
  where
    (vars, exprs) = unzip bs

vectTopBinder :: Var -> VM Var
vectTopBinder var
  = do
      vty <- liftM (mkForAllTys tyvars) $ vectType mono_ty
      name <- cloneName mkVectOcc (getName var)
      let var' | isExportedId var = Id.mkExportedLocalId name vty
               | otherwise        = Id.mkLocalId         name vty
      defGlobalVar var var'
      return var'
  where
    (tyvars, mono_ty) = splitForAllTys (idType var)
    
vectTopRhs :: CoreExpr -> VM CoreExpr
vectTopRhs = liftM fst . closedV . vectPolyExpr (panic "Empty lifting context") . freeVars

93 94
-- ----------------------------------------------------------------------------
-- Bindings
rl@cse.unsw.edu.au's avatar
rl@cse.unsw.edu.au committed
95 96 97 98 99

vectBndr :: Var -> VM (Var, Var)
vectBndr v
  = do
      vty <- vectType (idType v)
100
      lty <- mkPArrayType vty
rl@cse.unsw.edu.au's avatar
rl@cse.unsw.edu.au committed
101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123
      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
124
-- ----------------------------------------------------------------------------
125 126 127 128 129
-- Expressions

replicateP :: CoreExpr -> CoreExpr -> VM CoreExpr
replicateP expr len
  = do
130 131 132
      dict <- paDictOfType ty
      rep  <- builtin replicatePAVar
      return $ mkApps (Var rep) [Type ty, dict, expr, len]
133 134 135 136 137 138 139 140 141 142 143 144 145 146 147
  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)
rl@cse.unsw.edu.au's avatar
rl@cse.unsw.edu.au committed
148 149 150 151 152 153 154 155
vectVar lc v
  = do
      r <- lookupVar v
      case r of
        Local es     -> return es
        Global vexpr -> do
                          lexpr <- replicateP vexpr lc
                          return (vexpr, lexpr)
156 157 158 159

vectPolyVar :: CoreExpr -> Var -> [Type] -> VM (CoreExpr, CoreExpr)
vectPolyVar lc v tys
  = do
rl@cse.unsw.edu.au's avatar
rl@cse.unsw.edu.au committed
160
      r <- lookupVar v
161
      case r of
rl@cse.unsw.edu.au's avatar
rl@cse.unsw.edu.au committed
162 163 164 165 166
        Local (vexpr, lexpr) -> liftM2 (,) (mk_app vexpr) (mk_app lexpr)
        Global poly          -> do
                                  vexpr <- mk_app poly
                                  lexpr <- replicateP vexpr lc
                                  return (vexpr, lexpr)
167
  where
rl@cse.unsw.edu.au's avatar
rl@cse.unsw.edu.au committed
168
    mk_app e = applyToTypes e =<< mapM vectType tys
169

rl@cse.unsw.edu.au's avatar
rl@cse.unsw.edu.au committed
170 171
abstractOverTyVars :: [TyVar] -> ((CoreExpr -> CoreExpr) -> VM a) -> VM a
abstractOverTyVars tvs p
172 173
  = do
      mdicts <- mapM mk_dict_var tvs
rl@cse.unsw.edu.au's avatar
rl@cse.unsw.edu.au committed
174 175
      zipWithM_ (\tv -> maybe (deleteTyVarPA tv) (extendTyVarPA tv . Var)) tvs mdicts
      p (mk_lams mdicts)
176 177 178 179 180 181 182
  where
    mk_dict_var tv = do
                       r <- paDictArgType tv
                       case r of
                         Just ty -> liftM Just (newLocalVar FSLIT("dPA") ty)
                         Nothing -> return Nothing

rl@cse.unsw.edu.au's avatar
rl@cse.unsw.edu.au committed
183 184
    mk_lams mdicts = mkLams [arg | (tv, mdict) <- zip tvs mdicts
                                 , arg <- tv : maybeToList mdict]
rl@cse.unsw.edu.au's avatar
rl@cse.unsw.edu.au committed
185 186 187 188 189 190 191

applyToTypes :: CoreExpr -> [Type] -> VM CoreExpr
applyToTypes expr tys
  = do
      dicts <- mapM paDictOfType tys
      return $ mkApps expr [arg | (ty, dict) <- zip tys dicts
                                , arg <- [Type ty, dict]]
rl@cse.unsw.edu.au's avatar
rl@cse.unsw.edu.au committed
192 193 194 195 196 197 198 199 200 201 202 203
    

vectPolyExpr :: CoreExpr -> CoreExprWithFVs -> VM (CoreExpr, CoreExpr)
vectPolyExpr lc expr
  = localV
  . abstractOverTyVars tvs $ \mk_lams ->
    -- FIXME: shadowing (tvs in lc)
    do
      (vmono, lmono) <- vectExpr lc mono
      return $ (mk_lams vmono, mk_lams lmono)
  where
    (tvs, mono) = collectAnnTypeBinders expr  
204 205 206 207 208 209
                
vectExpr :: CoreExpr -> CoreExprWithFVs -> VM (CoreExpr, CoreExpr)
vectExpr lc (_, AnnType ty)
  = do
      vty <- vectType ty
      return (Type vty, Type vty)
rl@cse.unsw.edu.au's avatar
rl@cse.unsw.edu.au committed
210

211
vectExpr lc (_, AnnVar v)   = vectVar lc v
rl@cse.unsw.edu.au's avatar
rl@cse.unsw.edu.au committed
212

213 214 215 216 217
vectExpr lc (_, AnnLit lit)
  = do
      let vexpr = Lit lit
      lexpr <- replicateP vexpr lc
      return (vexpr, lexpr)
rl@cse.unsw.edu.au's avatar
rl@cse.unsw.edu.au committed
218

219 220 221 222
vectExpr lc (_, AnnNote note expr)
  = do
      (vexpr, lexpr) <- vectExpr lc expr
      return (Note note vexpr, Note note lexpr)
rl@cse.unsw.edu.au's avatar
rl@cse.unsw.edu.au committed
223

224 225 226 227 228
vectExpr lc e@(_, AnnApp _ arg)
  | isAnnTypeArg arg
  = vectTyAppExpr lc fn tys
  where
    (fn, tys) = collectAnnTypeArgs e
rl@cse.unsw.edu.au's avatar
rl@cse.unsw.edu.au committed
229

230 231 232 233 234
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
235

rl@cse.unsw.edu.au's avatar
rl@cse.unsw.edu.au committed
236 237
vectExpr lc (_, AnnCase expr bndr ty alts)
  = panic "vectExpr: case"
rl@cse.unsw.edu.au's avatar
rl@cse.unsw.edu.au committed
238

rl@cse.unsw.edu.au's avatar
rl@cse.unsw.edu.au committed
239 240
vectExpr lc (_, AnnLet (AnnNonRec bndr rhs) body)
  = do
241
      (vrhs, lrhs) <- vectPolyExpr lc rhs
rl@cse.unsw.edu.au's avatar
rl@cse.unsw.edu.au committed
242 243 244
      (vbndr, lbndr, (vbody, lbody)) <- vectBndrIn bndr (vectExpr lc body)
      return (Let (NonRec vbndr vrhs) vbody,
              Let (NonRec lbndr lrhs) lbody)
rl@cse.unsw.edu.au's avatar
rl@cse.unsw.edu.au committed
245

rl@cse.unsw.edu.au's avatar
rl@cse.unsw.edu.au committed
246 247 248 249 250 251 252 253 254 255
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
256
             (vbody, lbody) <- vectPolyExpr lc body
rl@cse.unsw.edu.au's avatar
rl@cse.unsw.edu.au committed
257
             return (vrhss, vbody, lrhss, lbody)
rl@cse.unsw.edu.au's avatar
rl@cse.unsw.edu.au committed
258

259 260 261
vectExpr lc e@(_, AnnLam bndr body)
  | isTyVar bndr = pprPanic "vectExpr" (ppr $ deAnnotate e)

262 263 264 265 266
vectExpr lc (fvs, AnnLam bndr body)
  = do
      let tyvars = filter isTyVar (varSetElems fvs)
      info <- mkCEnvInfo fvs bndr body
      (poly_vfn, poly_lfn) <- mkClosureFns info tyvars bndr body
267 268 269 270

      vfn_var <- hoistExpr FSLIT("vfn") poly_vfn
      lfn_var <- hoistExpr FSLIT("lfn") poly_lfn

271 272 273 274 275 276 277 278 279 280
      let (venv, lenv) = mkClosureEnvs info lc

      let env_ty = cenv_vty info

      pa_dict <- paDictOfType env_ty

      arg_ty <- vectType (varType bndr)
      res_ty <- vectType (exprType $ deAnnotate body)

      -- FIXME: move the functions to the top level
281 282
      mono_vfn <- applyToTypes (Var vfn_var) (map TyVarTy tyvars)
      mono_lfn <- applyToTypes (Var lfn_var) (map TyVarTy tyvars)
283 284 285 286 287 288 289 290 291 292 293 294 295 296 297 298 299 300 301 302 303 304 305 306 307 308 309 310 311 312 313 314 315 316 317 318 319 320 321 322 323 324 325 326 327 328 329 330 331 332 333 334 335 336 337 338 339 340 341 342 343 344 345 346 347 348 349 350 351 352 353 354 355 356 357 358 359 360 361 362 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 394 395 396 397 398 399 400 401 402 403 404 405 406 407 408 409 410 411 412 413 414

      mk_clo <- builtin mkClosureVar
      mk_cloP <- builtin mkClosurePVar

      let vclo = Var mk_clo  `mkTyApps` [arg_ty, res_ty, env_ty]
                             `mkApps`   [pa_dict, mono_vfn, mono_lfn, venv]
          
          lclo = Var mk_cloP `mkTyApps` [arg_ty, res_ty, env_ty]
                             `mkApps`   [pa_dict, mono_vfn, mono_lfn, lenv]

      return (vclo, lclo)
       

data CEnvInfo = CEnvInfo {
               cenv_vars         :: [Var]
             , cenv_values       :: [(CoreExpr, CoreExpr)]
             , cenv_vty          :: Type
             , cenv_lty          :: Type
             , cenv_repr_tycon   :: TyCon
             , cenv_repr_tyargs  :: [Type]
             , cenv_repr_datacon :: DataCon
             }

mkCEnvInfo :: VarSet -> Var -> CoreExprWithFVs -> VM CEnvInfo
mkCEnvInfo fvs arg body
  = do
      locals <- readLEnv local_vars
      let
          (vars, vals) = unzip
                 [(var, val) | var      <- varSetElems fvs
                             , Just val <- [lookupVarEnv locals var]]
      vtys <- mapM (vectType . varType) vars

      (vty, repr_tycon, repr_tyargs, repr_datacon) <- mk_env_ty vtys
      lty <- mkPArrayType vty
      
      return $ CEnvInfo {
                 cenv_vars         = vars
               , cenv_values       = vals
               , cenv_vty          = vty
               , cenv_lty          = lty
               , cenv_repr_tycon   = repr_tycon
               , cenv_repr_tyargs  = repr_tyargs
               , cenv_repr_datacon = repr_datacon
               }
  where
    mk_env_ty [vty]
      = return (vty, error "absent cinfo_repr_tycon"
                   , error "absent cinfo_repr_tyargs"
                   , error "absent cinfo_repr_datacon")

    mk_env_ty vtys
      = do
          let ty = mkCoreTupTy vtys
          (repr_tc, repr_tyargs) <- lookupPArrayFamInst ty
          let [repr_con] = tyConDataCons repr_tc
          return (ty, repr_tc, repr_tyargs, repr_con)

    

mkClosureEnvs :: CEnvInfo -> CoreExpr -> (CoreExpr, CoreExpr)
mkClosureEnvs info lc
  | [] <- vals
  = (Var unitDataConId, mkApps (Var $ dataConWrapId (cenv_repr_datacon info))
                               [lc, Var unitDataConId])

  | [(vval, lval)] <- vals
  = (vval, lval)

  | otherwise
  = (mkCoreTup vvals, Var (dataConWrapId $ cenv_repr_datacon info)
                      `mkTyApps` cenv_repr_tyargs info
                      `mkApps`   (lc : lvals))

  where
    vals = cenv_values info
    (vvals, lvals) = unzip vals

mkClosureFns :: CEnvInfo -> [TyVar] -> Var -> CoreExprWithFVs
             -> VM (CoreExpr, CoreExpr)
mkClosureFns info tyvars arg body
  = closedV
  . abstractOverTyVars tyvars
  $ \mk_tlams ->
  do
    (vfn, lfn) <- mkClosureMonoFns info arg body
    return (mk_tlams vfn, mk_tlams lfn)

mkClosureMonoFns :: CEnvInfo -> Var -> CoreExprWithFVs -> VM (CoreExpr, CoreExpr)
mkClosureMonoFns info arg body
  = do
      lc_bndr <- newLocalVar FSLIT("lc") intTy
      (varg : vbndrs, larg : lbndrs, (vbody, lbody))
        <- vectBndrsIn (arg : cenv_vars info)
                       (vectExpr (Var lc_bndr) body)

      venv_bndr <- newLocalVar FSLIT("env") vty
      lenv_bndr <- newLocalVar FSLIT("env") lty

      let vcase = bind_venv (Var venv_bndr) vbody vbndrs
      lcase <- bind_lenv (Var lenv_bndr) lbody lc_bndr lbndrs
      return (mkLams [venv_bndr, varg] vcase, mkLams [lenv_bndr, larg] lcase)
  where
    vty = cenv_vty info
    lty = cenv_lty info

    arity = length (cenv_vars info)

    bind_venv venv vbody []      = vbody
    bind_venv venv vbody [vbndr] = Let (NonRec vbndr venv) vbody
    bind_venv venv vbody vbndrs
      = Case venv (mkWildId vty) (exprType vbody)
             [(DataAlt (tupleCon Boxed arity), vbndrs, vbody)]

    bind_lenv lenv lbody lc_bndr [lbndr]
      = do
          lengthPA <- builtin lengthPAVar
          return . Let (NonRec lbndr lenv)
                 $ Case (mkApps (Var lengthPA) [Type vty, (Var lbndr)])
                        lc_bndr
                        intTy
                        [(DEFAULT, [], lbody)]

    bind_lenv lenv lbody lc_bndr lbndrs
      = return
      $ Case (unwrapFamInstScrut (cenv_repr_tycon info)
                                 (cenv_repr_tyargs info)
                                 lenv)
             (mkWildId lty)
             (exprType lbody)
             [(DataAlt (cenv_repr_datacon info), lc_bndr : lbndrs, lbody)]
          
415 416 417
vectTyAppExpr :: CoreExpr -> CoreExprWithFVs -> [Type] -> VM (CoreExpr, CoreExpr)
vectTyAppExpr lc (_, AnnVar v) tys = vectPolyVar lc v tys
vectTyAppExpr lc e tys = pprPanic "vectTyAppExpr" (ppr $ deAnnotate e)
418

419 420 421
-- ----------------------------------------------------------------------------
-- Types

rl@cse.unsw.edu.au's avatar
rl@cse.unsw.edu.au committed
422 423
vectTyCon :: TyCon -> VM TyCon
vectTyCon tc
424 425 426
  | 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
427 428 429 430 431 432 433 434 435 436 437 438 439 440 441 442 443
  | 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
444
      r   <- paDictArgType tv
rl@cse.unsw.edu.au's avatar
rl@cse.unsw.edu.au committed
445
      ty' <- vectType ty
446 447 448 449
      return $ ForAllTy tv (wrap r ty')
  where
    wrap Nothing      = id
    wrap (Just pa_ty) = FunTy pa_ty
rl@cse.unsw.edu.au's avatar
rl@cse.unsw.edu.au committed
450 451 452

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