Vectorise.hs 8.8 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
49
50
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
51
52
53
54
  = do
      showPass dflags "Vectorisation"
      eps <- hscEPS hsc_env
      let info = hptVectInfo hsc_env `plusVectInfo` eps_vect_info eps
55
      Just (info', guts') <- initV hsc_env guts info (vectModule guts)
rl@cse.unsw.edu.au's avatar
rl@cse.unsw.edu.au committed
56
      endPass dflags "Vectorisation" Opt_D_dump_vect (mg_binds guts')
57
      return (zeroSimplCount dflags, guts' { mg_vect_info = info' })
58
59
60
  where
    dflags = hsc_dflags hsc_env

rl@cse.unsw.edu.au's avatar
rl@cse.unsw.edu.au committed
61
vectModule :: ModGuts -> VM ModGuts
62
63
vectModule guts
  = do
64
      (types', fam_insts, tc_binds) <- vectTypeEnv (mg_types guts)
65
      
66
67
      let fam_inst_env' = extendFamInstEnvList (mg_fam_inst_env guts) fam_insts
      updGEnv (setFamInstEnv fam_inst_env')
68
     
69
70
      -- dicts   <- mapM buildPADict pa_insts
      -- workers <- mapM vectDataConWorkers pa_insts
rl@cse.unsw.edu.au's avatar
rl@cse.unsw.edu.au committed
71
      binds'  <- mapM vectTopBind (mg_binds guts)
72
      return $ guts { mg_types        = types'
73
                    , mg_binds        = Rec tc_binds : binds'
74
75
76
                    , 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
77

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

vectTopBind b@(Rec bs)
  = do
      vars'  <- mapM vectTopBinder vars
91
      exprs' <- zipWithM vectTopRhs vars exprs
92
93
94
95
96
97
98
99
100
101
      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
102
103
      vty  <- vectType (idType var)
      var' <- cloneId mkVectOcc var vty
104
105
106
      defGlobalVar var var'
      return var'
    
107
108
vectTopRhs :: Var -> CoreExpr -> VM CoreExpr
vectTopRhs var expr
109
110
  = do
      closedV . liftM vectorised
111
              . inBind var
112
              $ vectPolyExpr (freeVars expr)
113

114
115
-- ----------------------------------------------------------------------------
-- Bindings
rl@cse.unsw.edu.au's avatar
rl@cse.unsw.edu.au committed
116

117
vectBndr :: Var -> VM VVar
rl@cse.unsw.edu.au's avatar
rl@cse.unsw.edu.au committed
118
119
120
vectBndr v
  = do
      vty <- vectType (idType v)
121
      lty <- mkPArrayType vty
rl@cse.unsw.edu.au's avatar
rl@cse.unsw.edu.au committed
122
123
124
125
126
      let vv = v `Id.setIdType` vty
          lv = v `Id.setIdType` lty
      updLEnv (mapTo vv lv)
      return (vv, lv)
  where
127
    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
128

129
vectBndrIn :: Var -> VM a -> VM (VVar, a)
rl@cse.unsw.edu.au's avatar
rl@cse.unsw.edu.au committed
130
131
132
vectBndrIn v p
  = localV
  $ do
133
      vv <- vectBndr v
rl@cse.unsw.edu.au's avatar
rl@cse.unsw.edu.au committed
134
      x <- p
135
      return (vv, x)
rl@cse.unsw.edu.au's avatar
rl@cse.unsw.edu.au committed
136

137
138
139
140
141
142
143
144
vectBndrIn' :: Var -> (VVar -> VM a) -> VM (VVar, a)
vectBndrIn' v p
  = localV
  $ do
      vv <- vectBndr v
      x  <- p vv
      return (vv, x)

145
vectBndrsIn :: [Var] -> VM a -> VM ([VVar], a)
rl@cse.unsw.edu.au's avatar
rl@cse.unsw.edu.au committed
146
147
148
vectBndrsIn vs p
  = localV
  $ do
149
      vvs <- mapM vectBndr vs
rl@cse.unsw.edu.au's avatar
rl@cse.unsw.edu.au committed
150
      x <- p
151
      return (vvs, x)
rl@cse.unsw.edu.au's avatar
rl@cse.unsw.edu.au committed
152

rl@cse.unsw.edu.au's avatar
rl@cse.unsw.edu.au committed
153
-- ----------------------------------------------------------------------------
154
155
-- Expressions

156
157
vectVar :: Var -> VM VExpr
vectVar v
rl@cse.unsw.edu.au's avatar
rl@cse.unsw.edu.au committed
158
159
160
  = do
      r <- lookupVar v
      case r of
161
162
163
        Local (vv,lv) -> return (Var vv, Var lv)
        Global vv     -> do
                           let vexpr = Var vv
164
                           lexpr <- liftPA vexpr
165
                           return (vexpr, lexpr)
166

167
168
vectPolyVar :: Var -> [Type] -> VM VExpr
vectPolyVar v tys
169
  = do
170
      vtys <- mapM vectType tys
rl@cse.unsw.edu.au's avatar
rl@cse.unsw.edu.au committed
171
      r <- lookupVar v
172
      case r of
173
174
175
176
        Local (vv, lv) -> liftM2 (,) (polyApply (Var vv) vtys)
                                     (polyApply (Var lv) vtys)
        Global poly    -> do
                            vexpr <- polyApply (Var poly) vtys
177
                            lexpr <- liftPA vexpr
178
                            return (vexpr, lexpr)
179

180
181
vectLiteral :: Literal -> VM VExpr
vectLiteral lit
182
  = do
183
      lexpr <- liftPA (Lit lit)
184
185
      return (Lit lit, lexpr)

186
187
vectPolyExpr :: CoreExprWithFVs -> VM VExpr
vectPolyExpr expr
188
  = polyAbstract tvs $ \abstract ->
rl@cse.unsw.edu.au's avatar
rl@cse.unsw.edu.au committed
189
    do
190
      mono' <- vectExpr mono
191
      return $ mapVect abstract mono'
rl@cse.unsw.edu.au's avatar
rl@cse.unsw.edu.au committed
192
193
  where
    (tvs, mono) = collectAnnTypeBinders expr  
194
                
195
196
vectExpr :: CoreExprWithFVs -> VM VExpr
vectExpr (_, AnnType ty)
197
  = liftM vType (vectType ty)
rl@cse.unsw.edu.au's avatar
rl@cse.unsw.edu.au committed
198

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

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

203
204
vectExpr (_, AnnNote note expr)
  = liftM (vNote note) (vectExpr expr)
rl@cse.unsw.edu.au's avatar
rl@cse.unsw.edu.au committed
205

206
vectExpr e@(_, AnnApp _ arg)
207
  | isAnnTypeArg arg
208
  = vectTyAppExpr fn tys
209
210
  where
    (fn, tys) = collectAnnTypeArgs e
rl@cse.unsw.edu.au's avatar
rl@cse.unsw.edu.au committed
211

212
vectExpr (_, AnnApp fn arg)
213
  = do
214
215
      fn'  <- vectExpr fn
      arg' <- vectExpr arg
rl@cse.unsw.edu.au's avatar
rl@cse.unsw.edu.au committed
216
      mkClosureApp fn' arg'
rl@cse.unsw.edu.au's avatar
rl@cse.unsw.edu.au committed
217

218
219
220
221
222
223
vectExpr (_, AnnCase scrut bndr ty alts)
  | isAlgType scrut_ty
  = vectAlgCase scrut bndr ty alts
  where
    scrut_ty = exprType (deAnnotate scrut)

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

227
vectExpr (_, AnnLet (AnnNonRec bndr rhs) body)
rl@cse.unsw.edu.au's avatar
rl@cse.unsw.edu.au committed
228
  = do
229
230
      vrhs <- localV . inBind bndr $ vectPolyExpr rhs
      (vbndr, vbody) <- vectBndrIn bndr (vectExpr body)
231
      return $ vLet (vNonRec vbndr vrhs) vbody
rl@cse.unsw.edu.au's avatar
rl@cse.unsw.edu.au committed
232

233
vectExpr (_, AnnLet (AnnRec bs) body)
rl@cse.unsw.edu.au's avatar
rl@cse.unsw.edu.au committed
234
  = do
235
236
      (vbndrs, (vrhss, vbody)) <- vectBndrsIn bndrs
                                $ liftM2 (,)
237
                                  (zipWithM vect_rhs bndrs rhss)
238
                                  (vectPolyExpr body)
239
      return $ vLet (vRec vbndrs vrhss) vbody
rl@cse.unsw.edu.au's avatar
rl@cse.unsw.edu.au committed
240
  where
241
    (bndrs, rhss) = unzip bs
rl@cse.unsw.edu.au's avatar
rl@cse.unsw.edu.au committed
242

243
244
    vect_rhs bndr rhs = localV
                      . inBind bndr
245
                      $ vectExpr rhs
246

247
vectExpr e@(fvs, AnnLam bndr _)
248
  | not (isId bndr) = pprPanic "vectExpr" (ppr $ deAnnotate e)
249
  | otherwise = vectLam fvs bs body
250
251
  where
    (bs,body) = collectAnnValBinders e
252

253
254
vectLam :: VarSet -> [Var] -> CoreExprWithFVs -> VM VExpr
vectLam fvs bs body
255
  = do
256
      tyvars <- localTyVars
257
258
259
      (vs, vvs) <- readLEnv $ \env ->
                   unzip [(var, vv) | var <- varSetElems fvs
                                    , Just vv <- [lookupVarEnv (local_vars env) var]]
260

261
262
263
      arg_tys <- mapM (vectType . idType) bs
      res_ty  <- vectType (exprType $ deAnnotate body)

264
      buildClosures tyvars vvs arg_tys res_ty
265
        . hoistPolyVExpr tyvars
266
        $ do
267
            lc <- builtin liftingContext
268
            (vbndrs, vbody) <- vectBndrsIn (vs ++ bs)
269
                                           (vectExpr body)
270
            return $ vLams lc vbndrs vbody
271
  
272
273
274
vectTyAppExpr :: CoreExprWithFVs -> [Type] -> VM VExpr
vectTyAppExpr (_, AnnVar v) tys = vectPolyVar v tys
vectTyAppExpr e tys = pprPanic "vectTyAppExpr" (ppr $ deAnnotate e)
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
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