Vectorise.hs 8.95 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 216 217 218 219 220
      arg_ty' <- vectType arg_ty
      res_ty' <- vectType res_ty
      fn'     <- vectExpr fn
      arg'    <- vectExpr arg
      mkClosureApp arg_ty' res_ty' fn' arg'
  where
    (arg_ty, res_ty) = splitFunTy . exprType $ deAnnotate fn
rl@cse.unsw.edu.au's avatar
rl@cse.unsw.edu.au committed
221

222 223 224 225 226 227
vectExpr (_, AnnCase scrut bndr ty alts)
  | isAlgType scrut_ty
  = vectAlgCase scrut bndr ty alts
  where
    scrut_ty = exprType (deAnnotate scrut)

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

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

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

247 248
    vect_rhs bndr rhs = localV
                      . inBind bndr
249
                      $ vectExpr rhs
250

251
vectExpr e@(fvs, AnnLam bndr _)
252
  | not (isId bndr) = pprPanic "vectExpr" (ppr $ deAnnotate e)
253
  | otherwise = vectLam fvs bs body
254 255
  where
    (bs,body) = collectAnnValBinders e
256

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

265 266 267
      arg_tys <- mapM (vectType . idType) bs
      res_ty  <- vectType (exprType $ deAnnotate body)

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