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

#include "HsVersions.h"

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

11
12
13
import DynFlags
import HscTypes

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
24
import FamInstEnv           ( extendFamInstEnvList )
import InstEnv              ( extendInstEnvList )
rl@cse.unsw.edu.au's avatar
rl@cse.unsw.edu.au committed
25
26
import Var
import VarEnv
27
import VarSet
28
import Name                 ( Name, mkSysTvName, getName )
29
import NameEnv
rl@cse.unsw.edu.au's avatar
rl@cse.unsw.edu.au committed
30
import Id
31
import MkId                 ( unwrapFamInstScrut )
32
import OccName
33
import Module               ( Module )
rl@cse.unsw.edu.au's avatar
rl@cse.unsw.edu.au committed
34

rl@cse.unsw.edu.au's avatar
rl@cse.unsw.edu.au committed
35
import DsMonad hiding (mapAndUnzipM)
36
import DsUtils              ( mkCoreTup, mkCoreTupTy )
rl@cse.unsw.edu.au's avatar
rl@cse.unsw.edu.au committed
37

38
import Literal              ( Literal )
rl@cse.unsw.edu.au's avatar
rl@cse.unsw.edu.au committed
39
import PrelNames
40
import TysWiredIn
41
import TysPrim              ( intPrimTy )
42
import BasicTypes           ( Boxity(..) )
rl@cse.unsw.edu.au's avatar
rl@cse.unsw.edu.au committed
43

44
import Outputable
rl@cse.unsw.edu.au's avatar
rl@cse.unsw.edu.au committed
45
import FastString
46
import Control.Monad        ( liftM, liftM2, zipWithM, mapAndUnzipM )
47

48
builtin_PAs :: [(Name, Module, FastString)]
49
builtin_PAs = [
50
                (closureTyConName, nDP_CLOSURE, FSLIT("dPA_Clo"))
51
              , mk intTyConName     FSLIT("dPA_Int")
52
53
54
              ]
              ++ tups
  where
55
    mk name fs = (name, nDP_INSTANCES, fs)
56
57

    tups = mk_tup 0 : map mk_tup [2..3]
58
59
    mk_tup n   = (getName $ tupleTyCon Boxed n, nDP_INSTANCES,
                  mkFastString $ "dPA_" ++ show n)
60

61
62
63
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
64
65
66
67
  = do
      showPass dflags "Vectorisation"
      eps <- hscEPS hsc_env
      let info = hptVectInfo hsc_env `plusVectInfo` eps_vect_info eps
68
      Just (info', guts') <- initV hsc_env guts info (vectModule guts)
rl@cse.unsw.edu.au's avatar
rl@cse.unsw.edu.au committed
69
      endPass dflags "Vectorisation" Opt_D_dump_vect (mg_binds guts')
70
      return (zeroSimplCount dflags, guts' { mg_vect_info = info' })
71
72
73
  where
    dflags = hsc_dflags hsc_env

rl@cse.unsw.edu.au's avatar
rl@cse.unsw.edu.au committed
74
vectModule :: ModGuts -> VM ModGuts
75
76
vectModule guts
  = do
77
      defTyConBuiltinPAs builtin_PAs
78
      (types', fam_insts, tc_binds) <- vectTypeEnv (mg_types guts)
79
      
80
81
      let fam_inst_env' = extendFamInstEnvList (mg_fam_inst_env guts) fam_insts
      updGEnv (setFamInstEnv fam_inst_env')
82
     
83
84
      -- dicts   <- mapM buildPADict pa_insts
      -- workers <- mapM vectDataConWorkers pa_insts
rl@cse.unsw.edu.au's avatar
rl@cse.unsw.edu.au committed
85
      binds'  <- mapM vectTopBind (mg_binds guts)
86
      return $ guts { mg_types        = types'
87
                    , mg_binds        = Rec tc_binds : binds'
88
89
90
                    , 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
91

92
vectTopBind :: CoreBind -> VM CoreBind
93
94
95
vectTopBind b@(NonRec var expr)
  = do
      var'  <- vectTopBinder var
96
      expr' <- vectTopRhs var expr
97
98
99
100
101
102
103
104
      hs    <- takeHoisted
      return . Rec $ (var, expr) : (var', expr') : hs
  `orElseV`
    return b

vectTopBind b@(Rec bs)
  = do
      vars'  <- mapM vectTopBinder vars
105
      exprs' <- zipWithM vectTopRhs vars exprs
106
107
108
109
110
111
112
113
114
115
      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
116
117
      vty  <- vectType (idType var)
      var' <- cloneId mkVectOcc var vty
118
119
120
      defGlobalVar var var'
      return var'
    
121
122
vectTopRhs :: Var -> CoreExpr -> VM CoreExpr
vectTopRhs var expr
123
124
  = do
      closedV . liftM vectorised
125
              . inBind var
126
              $ vectPolyExpr (freeVars expr)
127

128
129
-- ----------------------------------------------------------------------------
-- Bindings
rl@cse.unsw.edu.au's avatar
rl@cse.unsw.edu.au committed
130

131
vectBndr :: Var -> VM VVar
rl@cse.unsw.edu.au's avatar
rl@cse.unsw.edu.au committed
132
133
134
vectBndr v
  = do
      vty <- vectType (idType v)
135
      lty <- mkPArrayType vty
rl@cse.unsw.edu.au's avatar
rl@cse.unsw.edu.au committed
136
137
138
139
140
      let vv = v `Id.setIdType` vty
          lv = v `Id.setIdType` lty
      updLEnv (mapTo vv lv)
      return (vv, lv)
  where
141
    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
142

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

151
152
153
154
155
156
157
158
vectBndrIn' :: Var -> (VVar -> VM a) -> VM (VVar, a)
vectBndrIn' v p
  = localV
  $ do
      vv <- vectBndr v
      x  <- p vv
      return (vv, x)

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

rl@cse.unsw.edu.au's avatar
rl@cse.unsw.edu.au committed
167
-- ----------------------------------------------------------------------------
168
169
-- Expressions

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

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

194
195
vectLiteral :: Literal -> VM VExpr
vectLiteral lit
196
  = do
197
      lexpr <- liftPA (Lit lit)
198
199
      return (Lit lit, lexpr)

200
201
vectPolyExpr :: CoreExprWithFVs -> VM VExpr
vectPolyExpr expr
202
  = polyAbstract tvs $ \abstract ->
rl@cse.unsw.edu.au's avatar
rl@cse.unsw.edu.au committed
203
    do
204
      mono' <- vectExpr mono
205
      return $ mapVect abstract mono'
rl@cse.unsw.edu.au's avatar
rl@cse.unsw.edu.au committed
206
207
  where
    (tvs, mono) = collectAnnTypeBinders expr  
208
                
209
210
vectExpr :: CoreExprWithFVs -> VM VExpr
vectExpr (_, AnnType ty)
211
  = liftM vType (vectType ty)
rl@cse.unsw.edu.au's avatar
rl@cse.unsw.edu.au committed
212

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

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

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

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

226
vectExpr (_, AnnApp fn arg)
227
  = do
228
229
      fn'  <- vectExpr fn
      arg' <- vectExpr arg
rl@cse.unsw.edu.au's avatar
rl@cse.unsw.edu.au committed
230
      mkClosureApp fn' arg'
rl@cse.unsw.edu.au's avatar
rl@cse.unsw.edu.au committed
231

232
233
234
235
236
237
vectExpr (_, AnnCase scrut bndr ty alts)
  | isAlgType scrut_ty
  = vectAlgCase scrut bndr ty alts
  where
    scrut_ty = exprType (deAnnotate scrut)

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

241
vectExpr (_, AnnLet (AnnNonRec bndr rhs) body)
rl@cse.unsw.edu.au's avatar
rl@cse.unsw.edu.au committed
242
  = do
243
244
      vrhs <- localV . inBind bndr $ vectPolyExpr rhs
      (vbndr, vbody) <- vectBndrIn bndr (vectExpr body)
245
      return $ vLet (vNonRec vbndr vrhs) vbody
rl@cse.unsw.edu.au's avatar
rl@cse.unsw.edu.au committed
246

247
vectExpr (_, AnnLet (AnnRec bs) body)
rl@cse.unsw.edu.au's avatar
rl@cse.unsw.edu.au committed
248
  = do
249
250
      (vbndrs, (vrhss, vbody)) <- vectBndrsIn bndrs
                                $ liftM2 (,)
251
                                  (zipWithM vect_rhs bndrs rhss)
252
                                  (vectPolyExpr body)
253
      return $ vLet (vRec vbndrs vrhss) vbody
rl@cse.unsw.edu.au's avatar
rl@cse.unsw.edu.au committed
254
  where
255
    (bndrs, rhss) = unzip bs
rl@cse.unsw.edu.au's avatar
rl@cse.unsw.edu.au committed
256

257
258
    vect_rhs bndr rhs = localV
                      . inBind bndr
259
                      $ vectExpr rhs
260

261
vectExpr e@(fvs, AnnLam bndr _)
262
  | not (isId bndr) = pprPanic "vectExpr" (ppr $ deAnnotate e)
263
  | otherwise = vectLam fvs bs body
264
265
  where
    (bs,body) = collectAnnValBinders e
266

267
268
vectLam :: VarSet -> [Var] -> CoreExprWithFVs -> VM VExpr
vectLam fvs bs body
269
  = do
270
      tyvars <- localTyVars
271
272
273
      (vs, vvs) <- readLEnv $ \env ->
                   unzip [(var, vv) | var <- varSetElems fvs
                                    , Just vv <- [lookupVarEnv (local_vars env) var]]
274

275
276
277
      arg_tys <- mapM (vectType . idType) bs
      res_ty  <- vectType (exprType $ deAnnotate body)

278
      buildClosures tyvars vvs arg_tys res_ty
279
        . hoistPolyVExpr tyvars
280
        $ do
281
            lc <- builtin liftingContext
282
            (vbndrs, vbody) <- vectBndrsIn (vs ++ bs)
283
                                           (vectExpr body)
284
            return $ vLams lc vbndrs vbody
285
  
286
287
288
vectTyAppExpr :: CoreExprWithFVs -> [Type] -> VM VExpr
vectTyAppExpr (_, AnnVar v) tys = vectPolyVar v tys
vectTyAppExpr e tys = pprPanic "vectTyAppExpr" (ppr $ deAnnotate e)
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
type CoreAltWithFVs = AnnAlt Id VarSet

-- We convert
--
--   case e :: t of v { ... }
--
-- to
--
--   V:    let v = e in case v of _ { ... }
--   L:    let v = e in case v `cast` ... of _ { ... }
--
-- When lifting, we have to do it this way because v must have the type
-- [:V(T):] but the scrutinee must be cast to the representation type.
--   

-- FIXME: this is too lazy
vectAlgCase scrut bndr ty [(DEFAULT, [], body)]
  = do
      vscrut <- vectExpr scrut
      vty    <- vectType ty
      lty    <- mkPArrayType vty
      (vbndr, vbody) <- vectBndrIn bndr (vectExpr body)
      return $ vCaseDEFAULT vscrut vbndr vty lty vbody

vectAlgCase scrut bndr ty [(DataAlt dc, bndrs, body)]
  = do
      vty <- vectType ty
      lty <- mkPArrayType vty
      vexpr <- vectExpr scrut
      (vbndr, (vbndrs, vbody)) <- vectBndrIn bndr
                                . vectBndrsIn bndrs
                                $ vectExpr body

      (vscrut, arr_tc, arg_tys) <- mkVScrut (vVar vbndr)
      vect_dc <- maybeV (lookupDataCon dc)
      let [arr_dc] = tyConDataCons arr_tc
      let shape_tys = take (dataConRepArity arr_dc - length bndrs)
                           (dataConRepArgTys arr_dc)
      shape_bndrs <- mapM (newLocalVar FSLIT("s")) shape_tys
      return . vLet (vNonRec vbndr vexpr)
             $ vCaseProd vscrut vty lty vect_dc arr_dc shape_bndrs vbndrs vbody