Vectorise.hs 12.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
rl@cse.unsw.edu.au's avatar
rl@cse.unsw.edu.au committed
23
import Name                 ( mkSysTvName )
24
import NameEnv
rl@cse.unsw.edu.au's avatar
rl@cse.unsw.edu.au committed
25
import Id
26
import MkId                 ( unwrapFamInstScrut )
rl@cse.unsw.edu.au's avatar
rl@cse.unsw.edu.au committed
27

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

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

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

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

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

57 58
-- ----------------------------------------------------------------------------
-- Bindings
rl@cse.unsw.edu.au's avatar
rl@cse.unsw.edu.au committed
59 60 61 62 63

vectBndr :: Var -> VM (Var, Var)
vectBndr v
  = do
      vty <- vectType (idType v)
64
      lty <- mkPArrayType vty
rl@cse.unsw.edu.au's avatar
rl@cse.unsw.edu.au committed
65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87
      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
88
-- ----------------------------------------------------------------------------
89 90 91 92 93
-- Expressions

replicateP :: CoreExpr -> CoreExpr -> VM CoreExpr
replicateP expr len
  = do
94 95 96
      dict <- paDictOfType ty
      rep  <- builtin replicatePAVar
      return $ mkApps (Var rep) [Type ty, dict, expr, len]
97 98 99 100 101 102 103 104 105 106 107 108 109 110 111
  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
112 113 114 115 116 117 118 119
vectVar lc v
  = do
      r <- lookupVar v
      case r of
        Local es     -> return es
        Global vexpr -> do
                          lexpr <- replicateP vexpr lc
                          return (vexpr, lexpr)
120 121 122 123

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
124
      r <- lookupVar v
125
      case r of
rl@cse.unsw.edu.au's avatar
rl@cse.unsw.edu.au committed
126 127 128 129 130
        Local (vexpr, lexpr) -> liftM2 (,) (mk_app vexpr) (mk_app lexpr)
        Global poly          -> do
                                  vexpr <- mk_app poly
                                  lexpr <- replicateP vexpr lc
                                  return (vexpr, lexpr)
131
  where
rl@cse.unsw.edu.au's avatar
rl@cse.unsw.edu.au committed
132
    mk_app e = applyToTypes e =<< mapM vectType tys
133

rl@cse.unsw.edu.au's avatar
rl@cse.unsw.edu.au committed
134 135
abstractOverTyVars :: [TyVar] -> ((CoreExpr -> CoreExpr) -> VM a) -> VM a
abstractOverTyVars tvs p
136 137
  = do
      mdicts <- mapM mk_dict_var tvs
rl@cse.unsw.edu.au's avatar
rl@cse.unsw.edu.au committed
138 139
      zipWithM_ (\tv -> maybe (deleteTyVarPA tv) (extendTyVarPA tv . Var)) tvs mdicts
      p (mk_lams mdicts)
140 141 142 143 144 145 146
  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
147 148
    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
149 150 151 152 153 154 155

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
156 157 158 159 160 161 162 163 164 165 166 167
    

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  
168 169 170 171 172 173
                
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
174

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

177 178 179 180 181
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
182

183 184 185 186
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
187

188 189 190 191 192
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
193

194 195 196 197 198
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
199

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

rl@cse.unsw.edu.au's avatar
rl@cse.unsw.edu.au committed
203 204
vectExpr lc (_, AnnLet (AnnNonRec bndr rhs) body)
  = do
205
      (vrhs, lrhs) <- vectPolyExpr lc rhs
rl@cse.unsw.edu.au's avatar
rl@cse.unsw.edu.au committed
206 207 208
      (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
209

rl@cse.unsw.edu.au's avatar
rl@cse.unsw.edu.au committed
210 211 212 213 214 215 216 217 218 219
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
220
             (vbody, lbody) <- vectPolyExpr lc body
rl@cse.unsw.edu.au's avatar
rl@cse.unsw.edu.au committed
221
             return (vrhss, vbody, lrhss, lbody)
rl@cse.unsw.edu.au's avatar
rl@cse.unsw.edu.au committed
222

223 224 225
vectExpr lc e@(_, AnnLam bndr body)
  | isTyVar bndr = pprPanic "vectExpr" (ppr $ deAnnotate e)

226 227 228 229 230
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
231 232 233 234

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

235 236 237 238 239 240 241 242 243 244
      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
245 246
      mono_vfn <- applyToTypes (Var vfn_var) (map TyVarTy tyvars)
      mono_lfn <- applyToTypes (Var lfn_var) (map TyVarTy tyvars)
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 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 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

      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)]
          
379 380 381
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)
382

383 384 385
-- ----------------------------------------------------------------------------
-- Types

rl@cse.unsw.edu.au's avatar
rl@cse.unsw.edu.au committed
386 387
vectTyCon :: TyCon -> VM TyCon
vectTyCon tc
388 389 390
  | 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
391 392 393 394 395 396 397 398 399 400 401 402 403 404 405 406 407
  | 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
408
      r   <- paDictArgType tv
rl@cse.unsw.edu.au's avatar
rl@cse.unsw.edu.au committed
409
      ty' <- vectType ty
410 411 412 413
      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
414 415 416

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