Vectorise.hs 7.79 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                 ( 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
rl@cse.unsw.edu.au's avatar
rl@cse.unsw.edu.au committed
33

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

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

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

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

rl@cse.unsw.edu.au's avatar
rl@cse.unsw.edu.au committed
60
vectModule :: ModGuts -> VM ModGuts
61
62
vectModule guts
  = do
63
64
65
66
      (types', fam_insts, pa_insts) <- vectTypeEnv (mg_types guts)
      
      let insts         = map painstInstance pa_insts
          fam_inst_env' = extendFamInstEnvList (mg_fam_inst_env guts) fam_insts
67
68
          inst_env'     = extendInstEnvList (mg_inst_env guts) insts
      updGEnv (setInstEnvs inst_env' fam_inst_env')
69
70
     
      dicts  <- mapM buildPADict pa_insts 
71
      binds' <- mapM vectTopBind (mg_binds guts)
72
      return $ guts { mg_types        = types'
73
                    , mg_binds        = Rec (concat dicts) : binds'
74
75
76
77
78
                    , mg_inst_env     = inst_env'
                    , mg_fam_inst_env = fam_inst_env'
                    , mg_insts        = mg_insts guts ++ insts
                    , mg_fam_insts    = mg_fam_insts guts ++ fam_insts
                    }
rl@cse.unsw.edu.au's avatar
rl@cse.unsw.edu.au committed
79

80
vectTopBind :: CoreBind -> VM CoreBind
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
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
104
      vty <- vectType (idType var)
105
106
107
108
109
110
111
      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
112
113
114
115
116
vectTopRhs expr
  = do
      lc <- newLocalVar FSLIT("lc") intPrimTy
      closedV . liftM vectorised
              $ vectPolyExpr lc (freeVars expr)
117

118
119
-- ----------------------------------------------------------------------------
-- Bindings
rl@cse.unsw.edu.au's avatar
rl@cse.unsw.edu.au committed
120

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

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

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

rl@cse.unsw.edu.au's avatar
rl@cse.unsw.edu.au committed
149
-- ----------------------------------------------------------------------------
150
151
-- Expressions

152
capply :: VExpr -> VExpr -> VM VExpr
153
154
155
156
157
158
159
160
161
162
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

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

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

187
188
189
190
191
192
vectLiteral :: Var -> Literal -> VM VExpr
vectLiteral lc lit
  = do
      lexpr <- replicatePA (Lit lit) (Var lc)
      return (Lit lit, lexpr)

193
vectPolyExpr :: Var -> CoreExprWithFVs -> VM VExpr
rl@cse.unsw.edu.au's avatar
rl@cse.unsw.edu.au committed
194
vectPolyExpr lc expr
195
  = polyAbstract tvs $ \abstract ->
rl@cse.unsw.edu.au's avatar
rl@cse.unsw.edu.au committed
196
197
    -- FIXME: shadowing (tvs in lc)
    do
198
199
      mono' <- vectExpr lc mono
      return $ mapVect abstract mono'
rl@cse.unsw.edu.au's avatar
rl@cse.unsw.edu.au committed
200
201
  where
    (tvs, mono) = collectAnnTypeBinders expr  
202
                
203
vectExpr :: Var -> CoreExprWithFVs -> VM VExpr
204
vectExpr lc (_, AnnType ty)
205
  = liftM vType (vectType ty)
rl@cse.unsw.edu.au's avatar
rl@cse.unsw.edu.au committed
206

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

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

211
vectExpr lc (_, AnnNote note expr)
212
  = liftM (vNote note) (vectExpr lc expr)
rl@cse.unsw.edu.au's avatar
rl@cse.unsw.edu.au committed
213

214
215
216
217
218
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
219

220
221
222
223
224
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
225

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

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

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

245
246
247
vectExpr lc e@(_, AnnLam bndr body)
  | isTyVar bndr = pprPanic "vectExpr" (ppr $ deAnnotate e)

248
249
vectExpr lc (fvs, AnnLam bndr body)
  = do
250
      tyvars <- localTyVars
251
252
253
      (vs, vvs) <- readLEnv $ \env ->
                   unzip [(var, vv) | var <- varSetElems fvs
                                    , Just vv <- [lookupVarEnv (local_vars env) var]]
254

255
      arg_ty <- vectType (idType bndr)
256
      res_ty <- vectType (exprType $ deAnnotate body)
257
258
259
260
261
262
263
      buildClosure tyvars lc vvs arg_ty res_ty
        . hoistPolyVExpr FSLIT("fn") tyvars
        $ do
            new_lc <- newLocalVar FSLIT("lc") intPrimTy
            (vbndrs, vbody) <- vectBndrsIn (vs ++ [bndr])
                                           (vectExpr new_lc body)
            return $ vLams new_lc vbndrs vbody
264

265
vectTyAppExpr :: Var -> CoreExprWithFVs -> [Type] -> VM (CoreExpr, CoreExpr)
266
267
vectTyAppExpr lc (_, AnnVar v) tys = vectPolyVar lc v tys
vectTyAppExpr lc e tys = pprPanic "vectTyAppExpr" (ppr $ deAnnotate e)
268