Vectorise.hs 14.3 KB
Newer Older
1

2 3 4
module Vectorise( vectorise )
where

5
import VectMonad
6
import VectUtils
7
import VectType
8
import VectCore
9

10 11 12
import DynFlags
import HscTypes

13
import Module               ( dphSeqPackageId, dphParPackageId )
rl@cse.unsw.edu.au's avatar
rl@cse.unsw.edu.au committed
14
import CoreLint             ( showPass, endPass )
15
import CoreSyn
16 17
import CoreUtils
import CoreFVs
18 19
import SimplMonad           ( SimplCount, zeroSimplCount )
import Rules                ( RuleBase )
20
import DataCon
rl@cse.unsw.edu.au's avatar
rl@cse.unsw.edu.au committed
21
import TyCon
rl@cse.unsw.edu.au's avatar
rl@cse.unsw.edu.au committed
22
import Type
23
import FamInstEnv           ( extendFamInstEnvList )
rl@cse.unsw.edu.au's avatar
rl@cse.unsw.edu.au committed
24 25
import Var
import VarEnv
26
import VarSet
rl@cse.unsw.edu.au's avatar
rl@cse.unsw.edu.au committed
27
import Id
28
import OccName
rl@cse.unsw.edu.au's avatar
rl@cse.unsw.edu.au committed
29

Ian Lynagh's avatar
Ian Lynagh committed
30
import DsMonad
rl@cse.unsw.edu.au's avatar
rl@cse.unsw.edu.au committed
31

32
import Literal              ( Literal, mkMachInt )
33
import TysWiredIn
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
Ian Lynagh's avatar
Ian Lynagh committed
37
import Control.Monad        ( liftM, liftM2, zipWithM )
38
import Data.List            ( sortBy, unzip4 )
39

40
vectorise :: DPHBackend -> HscEnv -> UniqSupply -> RuleBase -> ModGuts
41
          -> IO (SimplCount, ModGuts)
42
vectorise backend hsc_env _ _ guts
rl@cse.unsw.edu.au's avatar
rl@cse.unsw.edu.au committed
43 44 45 46
  = do
      showPass dflags "Vectorisation"
      eps <- hscEPS hsc_env
      let info = hptVectInfo hsc_env `plusVectInfo` eps_vect_info eps
47 48
      Just (info', guts') <- initV (backendPackage backend) 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 (zeroSimplCount dflags, guts' { mg_vect_info = info' })
51 52 53
  where
    dflags = hsc_dflags hsc_env

54 55 56
    backendPackage DPHSeq  = dphSeqPackageId
    backendPackage DPHPar  = dphParPackageId

rl@cse.unsw.edu.au's avatar
rl@cse.unsw.edu.au committed
57
vectModule :: ModGuts -> VM ModGuts
58 59
vectModule guts
  = do
60
      (types', fam_insts, tc_binds) <- vectTypeEnv (mg_types guts)
Ian Lynagh's avatar
Ian Lynagh committed
61

62 63
      let fam_inst_env' = extendFamInstEnvList (mg_fam_inst_env guts) fam_insts
      updGEnv (setFamInstEnv fam_inst_env')
Ian Lynagh's avatar
Ian Lynagh committed
64

65 66
      -- dicts   <- mapM buildPADict pa_insts
      -- workers <- mapM vectDataConWorkers pa_insts
rl@cse.unsw.edu.au's avatar
rl@cse.unsw.edu.au committed
67
      binds'  <- mapM vectTopBind (mg_binds guts)
68
      return $ guts { mg_types        = types'
69
                    , mg_binds        = Rec tc_binds : binds'
70 71 72
                    , mg_fam_inst_env = fam_inst_env'
                    , mg_fam_insts    = mg_fam_insts guts ++ fam_insts
                    }
rl@cse.unsw.edu.au's avatar
rl@cse.unsw.edu.au committed
73

74
vectTopBind :: CoreBind -> VM CoreBind
75 76 77
vectTopBind b@(NonRec var expr)
  = do
      var'  <- vectTopBinder var
78
      expr' <- vectTopRhs var expr
79
      hs    <- takeHoisted
80 81
      cexpr <- tryConvert var var' expr
      return . Rec $ (var, cexpr) : (var', expr') : hs
82 83 84 85 86 87
  `orElseV`
    return b

vectTopBind b@(Rec bs)
  = do
      vars'  <- mapM vectTopBinder vars
88
      exprs' <- zipWithM vectTopRhs vars exprs
89
      hs     <- takeHoisted
90 91
      cexprs <- sequence $ zipWith3 tryConvert vars vars' exprs
      return . Rec $ zip vars cexprs ++ zip vars' exprs' ++ hs
92 93 94 95 96 97 98 99
  `orElseV`
    return b
  where
    (vars, exprs) = unzip bs

vectTopBinder :: Var -> VM Var
vectTopBinder var
  = do
rl@cse.unsw.edu.au's avatar
rl@cse.unsw.edu.au committed
100 101
      vty  <- vectType (idType var)
      var' <- cloneId mkVectOcc var vty
102 103
      defGlobalVar var var'
      return var'
Ian Lynagh's avatar
Ian Lynagh committed
104

105 106
vectTopRhs :: Var -> CoreExpr -> VM CoreExpr
vectTopRhs var expr
107 108
  = do
      closedV . liftM vectorised
109
              . inBind var
110
              $ vectPolyExpr (freeVars expr)
111

112 113 114 115
tryConvert :: Var -> Var -> CoreExpr -> VM CoreExpr
tryConvert var vect_var rhs
  = fromVect (idType var) (Var vect_var) `orElseV` return rhs

116 117
-- ----------------------------------------------------------------------------
-- Bindings
rl@cse.unsw.edu.au's avatar
rl@cse.unsw.edu.au committed
118

119
vectBndr :: Var -> VM VVar
rl@cse.unsw.edu.au's avatar
rl@cse.unsw.edu.au committed
120 121
vectBndr v
  = do
122
      (vty, lty) <- vectAndLiftType (idType v)
rl@cse.unsw.edu.au's avatar
rl@cse.unsw.edu.au committed
123 124 125 126 127
      let vv = v `Id.setIdType` vty
          lv = v `Id.setIdType` lty
      updLEnv (mapTo vv lv)
      return (vv, lv)
  where
128
    mapTo vv lv env = env { local_vars = extendVarEnv (local_vars env) v (vv, lv) }
rl@cse.unsw.edu.au's avatar
rl@cse.unsw.edu.au committed
129

130 131 132 133 134 135 136 137 138 139
vectBndrNew :: Var -> FastString -> VM VVar
vectBndrNew v fs
  = do
      vty <- vectType (idType v)
      vv  <- newLocalVVar fs vty
      updLEnv (upd vv)
      return vv
  where
    upd vv env = env { local_vars = extendVarEnv (local_vars env) v vv }

140
vectBndrIn :: Var -> VM a -> VM (VVar, a)
rl@cse.unsw.edu.au's avatar
rl@cse.unsw.edu.au committed
141 142 143
vectBndrIn v p
  = localV
  $ do
144
      vv <- vectBndr v
rl@cse.unsw.edu.au's avatar
rl@cse.unsw.edu.au committed
145
      x <- p
146
      return (vv, x)
rl@cse.unsw.edu.au's avatar
rl@cse.unsw.edu.au committed
147

148 149 150 151 152 153 154 155
vectBndrNewIn :: Var -> FastString -> VM a -> VM (VVar, a)
vectBndrNewIn v fs p
  = localV
  $ do
      vv <- vectBndrNew v fs
      x  <- p
      return (vv, x)

156
vectBndrsIn :: [Var] -> VM a -> VM ([VVar], a)
rl@cse.unsw.edu.au's avatar
rl@cse.unsw.edu.au committed
157 158 159
vectBndrsIn vs p
  = localV
  $ do
160
      vvs <- mapM vectBndr vs
rl@cse.unsw.edu.au's avatar
rl@cse.unsw.edu.au committed
161
      x <- p
162
      return (vvs, x)
rl@cse.unsw.edu.au's avatar
rl@cse.unsw.edu.au committed
163

rl@cse.unsw.edu.au's avatar
rl@cse.unsw.edu.au committed
164
-- ----------------------------------------------------------------------------
165 166
-- Expressions

167 168
vectVar :: Var -> VM VExpr
vectVar v
rl@cse.unsw.edu.au's avatar
rl@cse.unsw.edu.au committed
169 170 171
  = do
      r <- lookupVar v
      case r of
172 173 174
        Local (vv,lv) -> return (Var vv, Var lv)
        Global vv     -> do
                           let vexpr = Var vv
175
                           lexpr <- liftPA vexpr
176
                           return (vexpr, lexpr)
177

178 179
vectPolyVar :: Var -> [Type] -> VM VExpr
vectPolyVar v tys
180
  = do
181
      vtys <- mapM vectType tys
rl@cse.unsw.edu.au's avatar
rl@cse.unsw.edu.au committed
182
      r <- lookupVar v
183
      case r of
184 185 186 187
        Local (vv, lv) -> liftM2 (,) (polyApply (Var vv) vtys)
                                     (polyApply (Var lv) vtys)
        Global poly    -> do
                            vexpr <- polyApply (Var poly) vtys
188
                            lexpr <- liftPA vexpr
189
                            return (vexpr, lexpr)
190

191 192
vectLiteral :: Literal -> VM VExpr
vectLiteral lit
193
  = do
194
      lexpr <- liftPA (Lit lit)
195 196
      return (Lit lit, lexpr)

197
vectPolyExpr :: CoreExprWithFVs -> VM VExpr
198 199
vectPolyExpr (_, AnnNote note expr)
  = liftM (vNote note) $ vectPolyExpr expr
200
vectPolyExpr expr
201
  = polyAbstract tvs $ \abstract ->
rl@cse.unsw.edu.au's avatar
rl@cse.unsw.edu.au committed
202
    do
203
      mono' <- vectExpr mono
204
      return $ mapVect abstract mono'
rl@cse.unsw.edu.au's avatar
rl@cse.unsw.edu.au committed
205
  where
Ian Lynagh's avatar
Ian Lynagh committed
206 207
    (tvs, mono) = collectAnnTypeBinders expr

208 209
vectExpr :: CoreExprWithFVs -> VM VExpr
vectExpr (_, AnnType ty)
210
  = liftM vType (vectType ty)
rl@cse.unsw.edu.au's avatar
rl@cse.unsw.edu.au committed
211

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

214
vectExpr (_, AnnLit lit) = vectLiteral lit
rl@cse.unsw.edu.au's avatar
rl@cse.unsw.edu.au committed
215

216 217
vectExpr (_, AnnNote note expr)
  = liftM (vNote note) (vectExpr expr)
rl@cse.unsw.edu.au's avatar
rl@cse.unsw.edu.au committed
218

219
vectExpr e@(_, AnnApp _ arg)
220
  | isAnnTypeArg arg
221
  = vectTyAppExpr fn tys
222 223
  where
    (fn, tys) = collectAnnTypeArgs e
rl@cse.unsw.edu.au's avatar
rl@cse.unsw.edu.au committed
224

225 226 227 228 229 230 231 232 233
vectExpr (_, AnnApp (_, AnnVar v) (_, AnnLit lit))
  | Just con <- isDataConId_maybe v
  , is_special_con con
  = do
      let vexpr = App (Var v) (Lit lit)
      lexpr <- liftPA vexpr
      return (vexpr, lexpr)
  where
    is_special_con con = con `elem` [intDataCon, floatDataCon, doubleDataCon]
Ian Lynagh's avatar
Ian Lynagh committed
234

235

236
vectExpr (_, AnnApp fn arg)
237
  = do
238 239 240 241 242 243 244
      arg_ty' <- vectType arg_ty
      res_ty' <- vectType res_ty
      fn'     <- vectExpr fn
      arg'    <- vectExpr arg
      mkClosureApp arg_ty' res_ty' fn' arg'
  where
    (arg_ty, res_ty) = splitFunTy . exprType $ deAnnotate fn
rl@cse.unsw.edu.au's avatar
rl@cse.unsw.edu.au committed
245

246
vectExpr (_, AnnCase scrut bndr ty alts)
247 248 249
  | Just (tycon, ty_args) <- splitTyConApp_maybe scrut_ty
  , isAlgTyCon tycon
  = vectAlgCase tycon ty_args scrut bndr ty alts
250 251 252
  where
    scrut_ty = exprType (deAnnotate scrut)

Ian Lynagh's avatar
Ian Lynagh committed
253
vectExpr (_, AnnCase _ _ _ _)
rl@cse.unsw.edu.au's avatar
rl@cse.unsw.edu.au committed
254
  = panic "vectExpr: case"
rl@cse.unsw.edu.au's avatar
rl@cse.unsw.edu.au committed
255

256
vectExpr (_, AnnLet (AnnNonRec bndr rhs) body)
rl@cse.unsw.edu.au's avatar
rl@cse.unsw.edu.au committed
257
  = do
258 259
      vrhs <- localV . inBind bndr $ vectPolyExpr rhs
      (vbndr, vbody) <- vectBndrIn bndr (vectExpr body)
260
      return $ vLet (vNonRec vbndr vrhs) vbody
rl@cse.unsw.edu.au's avatar
rl@cse.unsw.edu.au committed
261

262
vectExpr (_, AnnLet (AnnRec bs) body)
rl@cse.unsw.edu.au's avatar
rl@cse.unsw.edu.au committed
263
  = do
264 265
      (vbndrs, (vrhss, vbody)) <- vectBndrsIn bndrs
                                $ liftM2 (,)
266
                                  (zipWithM vect_rhs bndrs rhss)
267
                                  (vectPolyExpr body)
268
      return $ vLet (vRec vbndrs vrhss) vbody
rl@cse.unsw.edu.au's avatar
rl@cse.unsw.edu.au committed
269
  where
270
    (bndrs, rhss) = unzip bs
rl@cse.unsw.edu.au's avatar
rl@cse.unsw.edu.au committed
271

272 273
    vect_rhs bndr rhs = localV
                      . inBind bndr
274
                      $ vectExpr rhs
275

276
vectExpr e@(fvs, AnnLam bndr _)
277
  | not (isId bndr) = pprPanic "vectExpr" (ppr $ deAnnotate e)
278
  | otherwise = vectLam fvs bs body
279 280
  where
    (bs,body) = collectAnnValBinders e
281

282 283
vectExpr e = pprPanic "vectExpr" (ppr $ deAnnotate e)

284 285
vectLam :: VarSet -> [Var] -> CoreExprWithFVs -> VM VExpr
vectLam fvs bs body
286
  = do
287
      tyvars <- localTyVars
288 289 290
      (vs, vvs) <- readLEnv $ \env ->
                   unzip [(var, vv) | var <- varSetElems fvs
                                    , Just vv <- [lookupVarEnv (local_vars env) var]]
291

292 293 294
      arg_tys <- mapM (vectType . idType) bs
      res_ty  <- vectType (exprType $ deAnnotate body)

295
      buildClosures tyvars vvs arg_tys res_ty
296
        . hoistPolyVExpr tyvars
297
        $ do
298
            lc <- builtin liftingContext
299
            (vbndrs, vbody) <- vectBndrsIn (vs ++ bs)
300
                                           (vectExpr body)
301
            return $ vLams lc vbndrs vbody
Ian Lynagh's avatar
Ian Lynagh committed
302

303 304
vectTyAppExpr :: CoreExprWithFVs -> [Type] -> VM VExpr
vectTyAppExpr (_, AnnVar v) tys = vectPolyVar v tys
Ian Lynagh's avatar
Ian Lynagh committed
305
vectTyAppExpr e _ = pprPanic "vectTyAppExpr" (ppr $ deAnnotate e)
306 307 308 309 310 311 312

-- We convert
--
--   case e :: t of v { ... }
--
-- to
--
313 314
--   V:    let v' = e in case v' of _ { ... }
--   L:    let v' = e in case v' `cast` ... of _ { ... }
315 316
--
-- When lifting, we have to do it this way because v must have the type
317 318
-- [:V(T):] but the scrutinee must be cast to the representation type. We also
-- have to handle the case where v is a wild var correctly.
Ian Lynagh's avatar
Ian Lynagh committed
319
--
320 321

-- FIXME: this is too lazy
Ian Lynagh's avatar
Ian Lynagh committed
322 323 324 325
vectAlgCase :: TyCon -> [Type] -> CoreExprWithFVs -> Var -> Type
            -> [(AltCon, [Var], CoreExprWithFVs)]
            -> VM VExpr
vectAlgCase _tycon _ty_args scrut bndr ty [(DEFAULT, [], body)]
326
  = do
327 328
      vscrut         <- vectExpr scrut
      (vty, lty)     <- vectAndLiftType ty
329 330 331
      (vbndr, vbody) <- vectBndrIn bndr (vectExpr body)
      return $ vCaseDEFAULT vscrut vbndr vty lty vbody

Ian Lynagh's avatar
Ian Lynagh committed
332
vectAlgCase _tycon _ty_args scrut bndr ty [(DataAlt _, [], body)]
333
  = do
334 335
      vscrut         <- vectExpr scrut
      (vty, lty)     <- vectAndLiftType ty
336 337 338
      (vbndr, vbody) <- vectBndrIn bndr (vectExpr body)
      return $ vCaseDEFAULT vscrut vbndr vty lty vbody

Ian Lynagh's avatar
Ian Lynagh committed
339
vectAlgCase tycon _ty_args scrut bndr ty [(DataAlt dc, bndrs, body)]
340
  = do
341 342 343
      vect_tc    <- maybeV (lookupTyCon tycon)
      (vty, lty) <- vectAndLiftType ty
      vexpr      <- vectExpr scrut
344
      (vbndr, (vbndrs, vbody)) <- vect_scrut_bndr
345 346 347
                                . vectBndrsIn bndrs
                                $ vectExpr body

Ian Lynagh's avatar
Ian Lynagh committed
348
      (vscrut, arr_tc, _arg_tys) <- mkVScrut (vVar vbndr)
349 350
      vect_dc <- maybeV (lookupDataCon dc)
      let [arr_dc] = tyConDataCons arr_tc
351 352
      repr <- mkRepr vect_tc
      shape_bndrs <- arrShapeVars repr
353 354
      return . vLet (vNonRec vbndr vexpr)
             $ vCaseProd vscrut vty lty vect_dc arr_dc shape_bndrs vbndrs vbody
355
  where
Ian Lynagh's avatar
Ian Lynagh committed
356
    vect_scrut_bndr | isDeadBinder bndr = vectBndrNewIn bndr (fsLit "scrut")
357 358
                    | otherwise         = vectBndrIn bndr

Ian Lynagh's avatar
Ian Lynagh committed
359
vectAlgCase tycon _ty_args scrut bndr ty alts
360
  = do
361 362
      vect_tc     <- maybeV (lookupTyCon tycon)
      (vty, lty)  <- vectAndLiftType ty
363 364 365 366
      repr        <- mkRepr vect_tc
      shape_bndrs <- arrShapeVars repr
      (len, sel, indices) <- arrSelector repr (map Var shape_bndrs)

367
      (vbndr, valts) <- vect_scrut_bndr $ mapM (proc_alt sel vty lty) alts'
368 369 370
      let (vect_dcs, vect_bndrss, lift_bndrss, vbodies) = unzip4 valts

      vexpr <- vectExpr scrut
Ian Lynagh's avatar
Ian Lynagh committed
371
      (vscrut, arr_tc, _arg_tys) <- mkVScrut (vVar vbndr)
372 373 374 375 376 377 378 379 380 381 382 383 384 385 386 387
      let [arr_dc] = tyConDataCons arr_tc

      let (vect_scrut,  lift_scrut)  = vscrut
          (vect_bodies, lift_bodies) = unzip vbodies

      let vect_case = Case vect_scrut (mkWildId (exprType vect_scrut)) vty
                           (zipWith3 mk_vect_alt vect_dcs vect_bndrss vect_bodies)

      lbody <- combinePA vty len sel indices lift_bodies
      let lift_case = Case lift_scrut (mkWildId (exprType lift_scrut)) lty
                           [(DataAlt arr_dc, shape_bndrs ++ concat lift_bndrss,
                             lbody)]

      return . vLet (vNonRec vbndr vexpr)
             $ (vect_case, lift_case)
  where
Ian Lynagh's avatar
Ian Lynagh committed
388
    vect_scrut_bndr | isDeadBinder bndr = vectBndrNewIn bndr (fsLit "scrut")
389 390 391 392 393 394 395 396
                    | otherwise         = vectBndrIn bndr

    alts' = sortBy (\(alt1, _, _) (alt2, _, _) -> cmp alt1 alt2) alts

    cmp (DataAlt dc1) (DataAlt dc2) = dataConTag dc1 `compare` dataConTag dc2
    cmp DEFAULT       DEFAULT       = EQ
    cmp DEFAULT       _             = LT
    cmp _             DEFAULT       = GT
Ian Lynagh's avatar
Ian Lynagh committed
397
    cmp _             _             = panic "vectAlgCase/cmp"
398

399
    proc_alt sel vty lty (DataAlt dc, bndrs, body)
400 401 402 403 404 405
      = do
          vect_dc <- maybeV (lookupDataCon dc)
          let tag = mkDataConTag vect_dc
              fvs = freeVarsOf body `delVarSetList` bndrs
          (vect_bndrs, lift_bndrs, vbody)
            <- vect_alt_bndrs bndrs
406
             $ \len -> packLiftingContext len sel tag fvs vty lty
407 408 409
             $ vectExpr body

          return (vect_dc, vect_bndrs, lift_bndrs, vbody)
Ian Lynagh's avatar
Ian Lynagh committed
410
    proc_alt _ _ _ _ = panic "vectAlgCase/proc_alt"
411 412 413 414 415 416

    vect_alt_bndrs [] p
      = do
          void_tc <- builtin voidTyCon
          let void_ty = mkTyConApp void_tc []
          arr_ty <- mkPArrayType void_ty
Ian Lynagh's avatar
Ian Lynagh committed
417
          bndr   <- newLocalVar (fsLit "voids") arr_ty
418 419 420 421 422 423 424 425 426 427 428 429 430 431 432 433 434
          len    <- lengthPA void_ty (Var bndr)
          e      <- p len
          return ([], [bndr], e)

    vect_alt_bndrs bndrs p
       = localV
       $ do
           vbndrs <- mapM vectBndr bndrs
           let (vect_bndrs, lift_bndrs) = unzip vbndrs
               vv : _ = vect_bndrs
               lv : _ = lift_bndrs
           len <- lengthPA (idType vv) (Var lv)
           e   <- p len
           return (vect_bndrs, lift_bndrs, e)

    mk_vect_alt vect_dc bndrs body = (DataAlt vect_dc, bndrs, body)

435 436 437
packLiftingContext :: CoreExpr -> CoreExpr -> CoreExpr -> VarSet
                   -> Type -> Type -> VM VExpr -> VM VExpr
packLiftingContext len shape tag fvs vty lty p
438 439 440
  = do
      select <- builtin selectPAIntPrimVar
      let sel_expr = mkApps (Var select) [shape, tag]
Ian Lynagh's avatar
Ian Lynagh committed
441
      sel_var <- newLocalVar (fsLit "sel#") (exprType sel_expr)
442 443 444
      lc_var <- builtin liftingContext
      localV $
        do
445 446 447
          bnds <- mapM (packFreeVar (Var lc_var) (Var sel_var))
                . filter isLocalId
                $ varSetElems fvs
448
          (vexpr, lexpr) <- p
449
          empty <- emptyPA vty
450
          return (vexpr, Let (NonRec sel_var sel_expr)
451
                         $ Case len lc_var lty
452
                             [(DEFAULT, [], mkLets (concat bnds) lexpr),
453
                              (LitAlt (mkMachInt 0), [], empty)])
454 455 456 457 458 459 460 461 462 463 464 465 466 467 468 469 470

packFreeVar :: CoreExpr -> CoreExpr -> Var -> VM [CoreBind]
packFreeVar len sel v
  = do
      r <- lookupVar v
      case r of
        Local (vv,lv) ->
          do
            lv' <- cloneVar lv
            expr <- packPA (idType vv) (Var lv) len sel
            updLEnv (upd vv lv')
            return [(NonRec lv' expr)]

        _  -> return []
  where
    upd vv lv' env = env { local_vars = extendVarEnv (local_vars env) v (vv, lv') }