Vectorise.hs 14 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
17
import SimplMonad           ( SimplCount, zeroSimplCount )
import Rules                ( RuleBase )
18
import DataCon
rl@cse.unsw.edu.au's avatar
rl@cse.unsw.edu.au committed
19
import TyCon
rl@cse.unsw.edu.au's avatar
rl@cse.unsw.edu.au committed
20
21
import Type
import TypeRep
rl@cse.unsw.edu.au's avatar
rl@cse.unsw.edu.au committed
22
23
import Var
import VarEnv
24
import VarSet
25
import Name                 ( mkSysTvName, getName )
26
import NameEnv
rl@cse.unsw.edu.au's avatar
rl@cse.unsw.edu.au committed
27
import Id
28
import MkId                 ( unwrapFamInstScrut )
29
import OccName
rl@cse.unsw.edu.au's avatar
rl@cse.unsw.edu.au committed
30

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

import PrelNames
35
import TysWiredIn
36
import TysPrim              ( intPrimTy )
37
import BasicTypes           ( Boxity(..) )
rl@cse.unsw.edu.au's avatar
rl@cse.unsw.edu.au committed
38

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

44
45
46
vectorise :: HscEnv -> UniqSupply -> RuleBase -> ModGuts
          -> IO (SimplCount, ModGuts)
vectorise hsc_env _ _ guts
rl@cse.unsw.edu.au's avatar
rl@cse.unsw.edu.au committed
47
48
49
50
  = do
      showPass dflags "Vectorisation"
      eps <- hscEPS hsc_env
      let info = hptVectInfo hsc_env `plusVectInfo` eps_vect_info eps
51
      Just (info', guts') <- initV hsc_env guts info (vectModule guts)
rl@cse.unsw.edu.au's avatar
rl@cse.unsw.edu.au committed
52
      endPass dflags "Vectorisation" Opt_D_dump_vect (mg_binds guts')
53
      return (zeroSimplCount dflags, guts' { mg_vect_info = info' })
54
55
56
  where
    dflags = hsc_dflags hsc_env

rl@cse.unsw.edu.au's avatar
rl@cse.unsw.edu.au committed
57
vectModule :: ModGuts -> VM ModGuts
58
59
60
61
vectModule guts
  = do
      binds' <- mapM vectTopBind (mg_binds guts)
      return $ guts { mg_binds = binds' }
rl@cse.unsw.edu.au's avatar
rl@cse.unsw.edu.au committed
62

63
vectTopBind :: CoreBind -> VM CoreBind
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
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
rl@cse.unsw.edu.au's avatar
rl@cse.unsw.edu.au committed
87
      vty <- vectType (idType var)
88
89
90
91
92
93
94
95
96
      name <- cloneName mkVectOcc (getName var)
      let var' | isExportedId var = Id.mkExportedLocalId name vty
               | otherwise        = Id.mkLocalId         name vty
      defGlobalVar var var'
      return var'
    
vectTopRhs :: CoreExpr -> VM CoreExpr
vectTopRhs = liftM fst . closedV . vectPolyExpr (panic "Empty lifting context") . freeVars

97
98
-- ----------------------------------------------------------------------------
-- Bindings
rl@cse.unsw.edu.au's avatar
rl@cse.unsw.edu.au committed
99
100
101
102
103

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

replicateP :: CoreExpr -> CoreExpr -> VM CoreExpr
replicateP expr len
  = do
134
135
136
      dict <- paDictOfType ty
      rep  <- builtin replicatePAVar
      return $ mkApps (Var rep) [Type ty, dict, expr, len]
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
  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
152
153
154
155
156
157
158
159
vectVar lc v
  = do
      r <- lookupVar v
      case r of
        Local es     -> return es
        Global vexpr -> do
                          lexpr <- replicateP vexpr lc
                          return (vexpr, lexpr)
160
161
162
163

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
164
      r <- lookupVar v
165
      case r of
rl@cse.unsw.edu.au's avatar
rl@cse.unsw.edu.au committed
166
167
168
169
170
        Local (vexpr, lexpr) -> liftM2 (,) (mk_app vexpr) (mk_app lexpr)
        Global poly          -> do
                                  vexpr <- mk_app poly
                                  lexpr <- replicateP vexpr lc
                                  return (vexpr, lexpr)
171
  where
rl@cse.unsw.edu.au's avatar
rl@cse.unsw.edu.au committed
172
    mk_app e = applyToTypes e =<< mapM vectType tys
173

rl@cse.unsw.edu.au's avatar
rl@cse.unsw.edu.au committed
174
175
abstractOverTyVars :: [TyVar] -> ((CoreExpr -> CoreExpr) -> VM a) -> VM a
abstractOverTyVars tvs p
176
177
  = do
      mdicts <- mapM mk_dict_var tvs
rl@cse.unsw.edu.au's avatar
rl@cse.unsw.edu.au committed
178
179
      zipWithM_ (\tv -> maybe (deleteTyVarPA tv) (extendTyVarPA tv . Var)) tvs mdicts
      p (mk_lams mdicts)
180
181
182
183
184
185
186
  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
187
188
    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
189
190
191
192
193
194
195

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
196
197
198
199
200
201
202
203
204
205
206
207
    

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  
208
209
210
211
212
213
                
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
214

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

217
218
219
220
221
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
222

223
224
225
226
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
227

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

234
235
236
237
238
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
239

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

rl@cse.unsw.edu.au's avatar
rl@cse.unsw.edu.au committed
243
244
vectExpr lc (_, AnnLet (AnnNonRec bndr rhs) body)
  = do
245
      (vrhs, lrhs) <- vectPolyExpr lc rhs
rl@cse.unsw.edu.au's avatar
rl@cse.unsw.edu.au committed
246
247
248
      (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
249

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

263
264
265
vectExpr lc e@(_, AnnLam bndr body)
  | isTyVar bndr = pprPanic "vectExpr" (ppr $ deAnnotate e)

266
267
268
269
270
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
271
272
273
274

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

275
276
277
278
279
280
281
282
283
284
      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
285
286
      mono_vfn <- applyToTypes (Var vfn_var) (map TyVarTy tyvars)
      mono_lfn <- applyToTypes (Var lfn_var) (map TyVarTy tyvars)
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

      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
378
      lc_bndr <- newLocalVar FSLIT("lc") intPrimTy
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
      (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
407
                        (exprType lbody)
408
409
410
                        [(DEFAULT, [], lbody)]

    bind_lenv lenv lbody lc_bndr lbndrs
411
412
413
414
415
416
417
418
419
      = let scrut = unwrapFamInstScrut (cenv_repr_tycon info)
                                       (cenv_repr_tyargs info)
                                       lenv
            lbndrs' | null lbndrs = [mkWildId unitTy]
                    | otherwise   = lbndrs
        in
        return
      $ Case scrut
             (mkWildId (exprType scrut))
420
             (exprType lbody)
421
             [(DataAlt (cenv_repr_datacon info), lc_bndr : lbndrs', lbody)]
422
          
423
424
425
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)
426

427
428
429
-- ----------------------------------------------------------------------------
-- Types

rl@cse.unsw.edu.au's avatar
rl@cse.unsw.edu.au committed
430
431
vectTyCon :: TyCon -> VM TyCon
vectTyCon tc
432
433
434
  | 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
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
444
vectType ty | Just ty' <- coreView ty = vectType ty'
rl@cse.unsw.edu.au's avatar
rl@cse.unsw.edu.au committed
445
446
447
448
449
450
451
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
452
      r   <- paDictArgType tv
rl@cse.unsw.edu.au's avatar
rl@cse.unsw.edu.au committed
453
      ty' <- vectType ty
454
455
456
457
      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
458
459
460

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