Vectorise.hs 7.48 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
34
import RdrName              ( RdrName, mkRdrQual )
import Module               ( mkModuleNameFS )
rl@cse.unsw.edu.au's avatar
rl@cse.unsw.edu.au committed
35

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

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

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

49
50
51
52
53
54
mkNDPVar :: FastString -> RdrName
mkNDPVar fs = mkRdrQual nDP_BUILTIN (mkVarOccFS fs)

builtin_PAs :: [(Name, RdrName)]
builtin_PAs = [(intTyConName, mkNDPVar FSLIT("dPA_Int"))]

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

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

87
vectTopBind :: CoreBind -> VM CoreBind
88
89
90
vectTopBind b@(NonRec var expr)
  = do
      var'  <- vectTopBinder var
91
      expr' <- vectTopRhs var expr
92
93
94
95
96
97
98
99
      hs    <- takeHoisted
      return . Rec $ (var, expr) : (var', expr') : hs
  `orElseV`
    return b

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

123
124
-- ----------------------------------------------------------------------------
-- Bindings
rl@cse.unsw.edu.au's avatar
rl@cse.unsw.edu.au committed
125

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

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

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

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

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

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

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

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

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

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

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

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

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

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

222
vectExpr (_, AnnLet (AnnNonRec bndr rhs) body)
rl@cse.unsw.edu.au's avatar
rl@cse.unsw.edu.au committed
223
  = do
224
225
      vrhs <- localV . inBind bndr $ vectPolyExpr rhs
      (vbndr, vbody) <- vectBndrIn bndr (vectExpr body)
226
      return $ vLet (vNonRec vbndr vrhs) vbody
rl@cse.unsw.edu.au's avatar
rl@cse.unsw.edu.au committed
227

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

238
239
    vect_rhs bndr rhs = localV
                      . inBind bndr
240
                      $ vectExpr rhs
241

242
vectExpr e@(fvs, AnnLam bndr _)
243
  | not (isId bndr) = pprPanic "vectExpr" (ppr $ deAnnotate e)
244
  | otherwise = vectLam fvs bs body
245
246
  where
    (bs,body) = collectAnnValBinders e
247

248
249
vectLam :: VarSet -> [Var] -> CoreExprWithFVs -> VM VExpr
vectLam fvs bs body
250
  = do
251
      tyvars <- localTyVars
252
253
254
      (vs, vvs) <- readLEnv $ \env ->
                   unzip [(var, vv) | var <- varSetElems fvs
                                    , Just vv <- [lookupVarEnv (local_vars env) var]]
255

256
257
258
      arg_tys <- mapM (vectType . idType) bs
      res_ty  <- vectType (exprType $ deAnnotate body)

259
      buildClosures tyvars vvs arg_tys res_ty
260
        . hoistPolyVExpr tyvars
261
        $ do
262
            lc <- builtin liftingContext
263
            (vbndrs, vbody) <- vectBndrsIn (vs ++ bs)
264
                                           (vectExpr body)
265
            return $ vLams lc vbndrs vbody
266
  
267
268
269
vectTyAppExpr :: CoreExprWithFVs -> [Type] -> VM VExpr
vectTyAppExpr (_, AnnVar v) tys = vectPolyVar v tys
vectTyAppExpr e tys = pprPanic "vectTyAppExpr" (ppr $ deAnnotate e)
270