compiler/vectorise/Vectorise.hs

   1 {-# OPTIONS -fno-warn-missing-signatures #-}
   2
   3 module Vectorise( vectorise )
   4 where
   5
   6 import VectMonad
   7 import VectUtils
   8 import VectType
   9 import VectCore
  10
  11 import HscTypes hiding      ( MonadThings(..) )
  12
  13 import Module               ( PackageId )
  14 import CoreSyn
  15 import CoreUtils
  16 import CoreUnfold           ( mkInlineRule )
  17 import MkCore               ( mkWildCase )
  18 import CoreFVs
  19 import CoreMonad            ( CoreM, getHscEnv )
  20 import DataCon
  21 import TyCon
  22 import Type
  23 import FamInstEnv           ( extendFamInstEnvList )
  24 import Var
  25 import VarEnv
  26 import VarSet
  27 import Id
  28 import OccName
  29 import BasicTypes           ( isLoopBreaker )
  30
  31 import Literal              ( Literal, mkMachInt )
  32 import TysWiredIn
  33 import TysPrim              ( intPrimTy )
  34
  35 import Outputable
  36 import FastString
  37 import Util                 ( zipLazy )
  38 import Control.Monad
  39 import Data.List            ( sortBy, unzip4 )
  40
  41
  42 debug           = False
  43 dtrace s x      = if debug then pprTrace "Vectorise" s x else x
  44
  45 -- | Vectorise a single module.
  46 --   Takes the package containing the DPH backend we're using. Eg either dph-par or dph-seq.
  47 vectorise :: PackageId -> ModGuts -> CoreM ModGuts
  48 vectorise backend guts
  49  = do hsc_env <- getHscEnv
  50       liftIO $ vectoriseIO backend hsc_env guts
  51
  52
  53 -- | Vectorise a single monad, given its HscEnv (code gen environment).
  54 vectoriseIO :: PackageId -> HscEnv -> ModGuts -> IO ModGuts
  55 vectoriseIO backend hsc_env guts
  56  = do -- Get information about currently loaded external packages.
  57       eps <- hscEPS hsc_env
  58
  59       -- Combine vectorisation info from the current module, and external ones.
  60       let info = hptVectInfo hsc_env `plusVectInfo` eps_vect_info eps
  61
  62       -- Run the main VM computation.
  63       Just (info', guts') <- initV backend hsc_env guts info (vectModule guts)
  64       return (guts' { mg_vect_info = info' })
  65
  66
  67 -- | Vectorise a single module, in the VM monad.
  68 vectModule :: ModGuts -> VM ModGuts
  69 vectModule guts
  70  = do -- Vectorise the type environment.
  71       -- This may add new TyCons and DataCons.
  72       -- TODO: What new binds do we get back here?
  73       (types', fam_insts, tc_binds) <- vectTypeEnv (mg_types guts)
  74
  75       -- TODO: What is this?
  76       let fam_inst_env' = extendFamInstEnvList (mg_fam_inst_env guts) fam_insts
  77       updGEnv (setFamInstEnv fam_inst_env')
  78
  79       -- dicts   <- mapM buildPADict pa_insts
  80       -- workers <- mapM vectDataConWorkers pa_insts
  81
  82       -- Vectorise all the top level bindings.
  83       binds'  <- mapM vectTopBind (mg_binds guts)
  84
  85       return $ guts { mg_types        = types'
  86                     , mg_binds        = Rec tc_binds : binds'
  87                     , mg_fam_inst_env = fam_inst_env'
  88                     , mg_fam_insts    = mg_fam_insts guts ++ fam_insts
  89                     }
  90
  91
  92 -- | Try to vectorise a top-level binding.
  93 --   If it doesn't vectorise then return it unharmed.
  94 --
  95 --   For example, for the binding
  96 --
  97 --   @
  98 --      foo :: Int -> Int
  99 --      foo = \x -> x + x
 100 --   @
 101 --
 102 --   we get
 103 --   @
 104 --      foo  :: Int -> Int
 105 --      foo  = \x -> vfoo $: x
 106 --
 107 --      v_foo :: Closure void vfoo lfoo
 108 --      v_foo = closure vfoo lfoo void
 109 --
 110 --      vfoo :: Void -> Int -> Int
 111 --      vfoo = ...
 112 --
 113 --      lfoo :: PData Void -> PData Int -> PData Int
 114 --      lfoo = ...
 115 --   @
 116 --
 117 --   @vfoo@ is the "vectorised", or scalar, version that does the same as the original
 118 --   function foo, but takes an explicit environment.
 119 --
 120 --   @lfoo@ is the "lifted" version that works on arrays.
 121 --
 122 --   @v_foo@ combines both of these into a `Closure` that also contains the
 123 --   environment.
 124 --
 125 --   The original binding @foo@ is rewritten to call the vectorised version
 126 --   present in the closure.
 127 --
 128 vectTopBind :: CoreBind -> VM CoreBind
 129 vectTopBind b@(NonRec var expr)
 130  = do
 131       (inline, expr')   <- vectTopRhs var expr
 132       var'              <- vectTopBinder var inline expr'
 133
 134       -- Vectorising the body may create other top-level bindings.
 135       hs        <- takeHoisted
 136
 137       -- To get the same functionality as the original body we project
 138       -- out its vectorised version from the closure.
 139       cexpr     <- tryConvert var var' expr
 140
 141       return . Rec $ (var, cexpr) : (var', expr') : hs
 142   `orElseV`
 143     return b
 144
 145 vectTopBind b@(Rec bs)
 146  = do
 147       (vars', _, exprs')
 148         <- fixV $ \ ~(_, inlines, rhss) ->
 149             do vars' <- sequence [vectTopBinder var inline rhs
 150                                       | (var, ~(inline, rhs)) <- zipLazy vars (zip inlines rhss)]
 151                (inlines', exprs')
 152                      <- mapAndUnzipM (uncurry vectTopRhs) bs
 153
 154                return (vars', inlines', exprs')
 155
 156       hs     <- takeHoisted
 157       cexprs <- sequence $ zipWith3 tryConvert vars vars' exprs
 158       return . Rec $ zip vars cexprs ++ zip vars' exprs' ++ hs
 159   `orElseV`
 160     return b
 161   where
 162     (vars, exprs) = unzip bs
 163
 164
 165 -- | Make the vectorised version of this top level binder, and add the mapping
 166 --   between it and the original to the state. For some binder @foo@ the vectorised
 167 --   version is @$v_foo@
 168 --
 169 --   NOTE: vectTopBinder *MUST* be lazy in inline and expr because of how it is
 170 --   used inside of fixV in vectTopBind
 171 vectTopBinder
 172         :: Var          -- ^ Name of the binding.
 173         -> Inline       -- ^ Whether it should be inlined, used to annotate it.
 174         -> CoreExpr     -- ^ RHS of the binding, used to set the `Unfolding` of the returned `Var`.
 175         -> VM Var       -- ^ Name of the vectorised binding.
 176
 177 vectTopBinder var inline expr
 178  = do
 179       -- Vectorise the type attached to the var.
 180       vty  <- vectType (idType var)
 181
 182       -- Make the vectorised version of binding's name, and set the unfolding used for inlining.
 183       var' <- liftM (`setIdUnfolding` unfolding)
 184            $  cloneId mkVectOcc var vty
 185
 186       -- Add the mapping between the plain and vectorised name to the state.
 187       defGlobalVar var var'
 188
 189       return var'
 190   where
 191     unfolding = case inline of
 192                   Inline arity -> mkInlineRule expr (Just arity)
 193                   DontInline   -> noUnfolding
 194
 195
 196 -- | Vectorise the RHS of a top-level binding, in an empty local environment.
 197 vectTopRhs
 198         :: Var          -- ^ Name of the binding.
 199         -> CoreExpr     -- ^ Body of the binding.
 200         -> VM (Inline, CoreExpr)
 201
 202 vectTopRhs var expr
 203  = dtrace (vcat [text "vectTopRhs", ppr expr])
 204  $ closedV
 205  $ do (inline, vexpr) <- inBind var
 206                       $ vectPolyExpr (isLoopBreaker $ idOccInfo var)
 207                                       (freeVars expr)
 208       return (inline, vectorised vexpr)
 209
 210
 211 -- | Project out the vectorised version of a binding from some closure,
 212 --      or return the original body if that doesn't work.
 213 tryConvert
 214         :: Var          -- ^ Name of the original binding (eg @foo@)
 215         -> Var          -- ^ Name of vectorised version of binding (eg @$vfoo@)
 216         -> CoreExpr     -- ^ The original body of the binding.
 217         -> VM CoreExpr
 218
 219 tryConvert var vect_var rhs
 220   = fromVect (idType var) (Var vect_var) `orElseV` return rhs
 221
 222
 223 -- ----------------------------------------------------------------------------
 224 -- Bindings
 225
 226 -- | Vectorise a binder variable, along with its attached type.
 227 vectBndr :: Var -> VM VVar
 228 vectBndr v
 229   = do
 230       (vty, lty) <- vectAndLiftType (idType v)
 231       let vv = v `Id.setIdType` vty
 232           lv = v `Id.setIdType` lty
 233       updLEnv (mapTo vv lv)
 234       return (vv, lv)
 235   where
 236     mapTo vv lv env = env { local_vars = extendVarEnv (local_vars env) v (vv, lv) }
 237
 238
 239 -- | Vectorise a binder variable, along with its attached type,
 240 --   but give the result a new name.
 241 vectBndrNew :: Var -> FastString -> VM VVar
 242 vectBndrNew v fs
 243   = do
 244       vty <- vectType (idType v)
 245       vv  <- newLocalVVar fs vty
 246       updLEnv (upd vv)
 247       return vv
 248   where
 249     upd vv env = env { local_vars = extendVarEnv (local_vars env) v vv }
 250
 251
 252 -- | Vectorise a binder then run a computation with that binder in scope.
 253 vectBndrIn :: Var -> VM a -> VM (VVar, a)
 254 vectBndrIn v p
 255   = localV
 256   $ do
 257       vv <- vectBndr v
 258       x <- p
 259       return (vv, x)
 260
 261
 262 -- | Vectorise a binder, give it a new name, then run a computation with that binder in scope.
 263 vectBndrNewIn :: Var -> FastString -> VM a -> VM (VVar, a)
 264 vectBndrNewIn v fs p
 265   = localV
 266   $ do
 267       vv <- vectBndrNew v fs
 268       x  <- p
 269       return (vv, x)
 270
 271 -- | Vectorise some binders, then run a computation with them in scope.
 272 vectBndrsIn :: [Var] -> VM a -> VM ([VVar], a)
 273 vectBndrsIn vs p
 274   = localV
 275   $ do
 276       vvs <- mapM vectBndr vs
 277       x <- p
 278       return (vvs, x)
 279
 280
 281 -- ----------------------------------------------------------------------------
 282 -- Expressions
 283
 284 -- | Vectorise a variable, producing the vectorised and lifted versions.
 285 vectVar :: Var -> VM VExpr
 286 vectVar v
 287  = do
 288       -- lookup the variable from the environment.
 289       r <- lookupVar v
 290
 291       case r of
 292         -- If it's been locally bound then we'll already have both versions available.
 293         Local (vv,lv)
 294          -> return (Var vv, Var lv)
 295
 296         -- To create the lifted version of a global variable we replicate it.
 297         Global vv
 298          -> do let vexpr = Var vv
 299                lexpr <- liftPD vexpr
 300                return (vexpr, lexpr)
 301
 302
 303 -- | Like `vectVar` but also add type applications to the variables.
 304 vectPolyVar :: Var -> [Type] -> VM VExpr
 305 vectPolyVar v tys
 306   = do
 307       vtys      <- mapM vectType tys
 308       r         <- lookupVar v
 309       case r of
 310         Local (vv, lv)
 311          -> liftM2 (,) (polyApply (Var vv) vtys)
 312                        (polyApply (Var lv) vtys)
 313
 314         Global poly
 315          -> do vexpr <- polyApply (Var poly) vtys
 316                lexpr <- liftPD vexpr
 317                return (vexpr, lexpr)
 318
 319
 320 -- | Lifted literals are created by replicating them.
 321 vectLiteral :: Literal -> VM VExpr
 322 vectLiteral lit
 323   = do
 324       lexpr <- liftPD (Lit lit)
 325       return (Lit lit, lexpr)
 326
 327
 328 -- | Vectorise a polymorphic expression
 329 vectPolyExpr
 330         :: Bool                 -- ^ When vectorising the RHS of a binding, whether that
 331                                 --   binding is a loop breaker.
 332         -> CoreExprWithFVs
 333         -> VM (Inline, VExpr)
 334
 335 vectPolyExpr loop_breaker (_, AnnNote note expr)
 336  = do (inline, expr') <- vectPolyExpr loop_breaker expr
 337       return (inline, vNote note expr')
 338
 339 vectPolyExpr loop_breaker expr
 340  = dtrace (vcat [text "vectPolyExpr", ppr (deAnnotate expr)])
 341  $ do
 342       arity <- polyArity tvs
 343       polyAbstract tvs $ \args ->
 344         do
 345           (inline, mono') <- vectFnExpr False loop_breaker mono
 346           return (addInlineArity inline arity,
 347                   mapVect (mkLams $ tvs ++ args) mono')
 348   where
 349     (tvs, mono) = collectAnnTypeBinders expr
 350
 351
 352 -- | Vectorise a core expression.
 353 vectExpr :: CoreExprWithFVs -> VM VExpr
 354 vectExpr (_, AnnType ty)
 355   = liftM vType (vectType ty)
 356
 357 vectExpr (_, AnnVar v)
 358   = vectVar v
 359
 360 vectExpr (_, AnnLit lit)
 361   = vectLiteral lit
 362
 363 vectExpr (_, AnnNote note expr)
 364   = liftM (vNote note) (vectExpr expr)
 365
 366 vectExpr e@(_, AnnApp _ arg)
 367   | isAnnTypeArg arg
 368   = vectTyAppExpr fn tys
 369   where
 370     (fn, tys) = collectAnnTypeArgs e
 371
 372 vectExpr (_, AnnApp (_, AnnVar v) (_, AnnLit lit))
 373   | Just con <- isDataConId_maybe v
 374   , is_special_con con
 375   = do
 376       let vexpr = App (Var v) (Lit lit)
 377       lexpr <- liftPD vexpr
 378       return (vexpr, lexpr)
 379   where
 380     is_special_con con = con `elem` [intDataCon, floatDataCon, doubleDataCon]
 381
 382
 383 -- TODO: Avoid using closure application for dictionaries.
 384 -- vectExpr (_, AnnApp fn arg)
 385 --  | if is application of dictionary
 386 --    just use regular app instead of closure app.
 387
 388 -- for lifted version.
 389 --      do liftPD (sub a dNumber)
 390 --      lift the result of the selection, not sub and dNumber seprately.
 391
 392 vectExpr (_, AnnApp fn arg)
 393  = dtrace (text "AnnApp" <+> ppr (deAnnotate fn) <+> ppr (deAnnotate arg))
 394  $ do
 395       arg_ty' <- vectType arg_ty
 396       res_ty' <- vectType res_ty
 397
 398       dtrace (text "vectorising fn " <> ppr (deAnnotate fn))  $ return ()
 399       fn'     <- vectExpr fn
 400       dtrace (text "fn' = "       <> ppr fn') $ return ()
 401
 402       arg'    <- vectExpr arg
 403
 404       mkClosureApp arg_ty' res_ty' fn' arg'
 405   where
 406     (arg_ty, res_ty) = splitFunTy . exprType $ deAnnotate fn
 407
 408 vectExpr (_, AnnCase scrut bndr ty alts)
 409   | Just (tycon, ty_args) <- splitTyConApp_maybe scrut_ty
 410   , isAlgTyCon tycon
 411   = vectAlgCase tycon ty_args scrut bndr ty alts
 412   where
 413     scrut_ty = exprType (deAnnotate scrut)
 414
 415 vectExpr (_, AnnLet (AnnNonRec bndr rhs) body)
 416   = do
 417       vrhs <- localV . inBind bndr . liftM snd $ vectPolyExpr False rhs
 418       (vbndr, vbody) <- vectBndrIn bndr (vectExpr body)
 419       return $ vLet (vNonRec vbndr vrhs) vbody
 420
 421 vectExpr (_, AnnLet (AnnRec bs) body)
 422   = do
 423       (vbndrs, (vrhss, vbody)) <- vectBndrsIn bndrs
 424                                 $ liftM2 (,)
 425                                   (zipWithM vect_rhs bndrs rhss)
 426                                   (vectExpr body)
 427       return $ vLet (vRec vbndrs vrhss) vbody
 428   where
 429     (bndrs, rhss) = unzip bs
 430
 431     vect_rhs bndr rhs = localV
 432                       . inBind bndr
 433                       . liftM snd
 434                       $ vectPolyExpr (isLoopBreaker $ idOccInfo bndr) rhs
 435
 436 vectExpr e@(_, AnnLam bndr _)
 437   | isId bndr = liftM snd $ vectFnExpr True False e
 438 {-
 439 onlyIfV (isEmptyVarSet fvs) (vectScalarLam bs $ deAnnotate body)
 440                 `orElseV` vectLam True fvs bs body
 441   where
 442     (bs,body) = collectAnnValBinders e
 443 -}
 444
 445 vectExpr e = cantVectorise "Can't vectorise expression" (ppr $ deAnnotate e)
 446
 447
 448 -- | Vectorise an expression with an outer lambda abstraction.
 449 vectFnExpr
 450         :: Bool                 -- ^ When the RHS of a binding, whether that binding should be inlined.
 451         -> Bool                 -- ^ Whether the binding is a loop breaker.
 452         -> CoreExprWithFVs      -- ^ Expression to vectorise. Must have an outer `AnnLam`.
 453         -> VM (Inline, VExpr)
 454
 455 vectFnExpr inline loop_breaker e@(fvs, AnnLam bndr _)
 456   | isId bndr = onlyIfV (isEmptyVarSet fvs)
 457                         (mark DontInline . vectScalarLam bs $ deAnnotate body)
 458                 `orElseV` mark inlineMe (vectLam inline loop_breaker fvs bs body)
 459   where
 460     (bs,body) = collectAnnValBinders e
 461
 462 vectFnExpr _ _ e = mark DontInline $ vectExpr e
 463
 464 mark :: Inline -> VM a -> VM (Inline, a)
 465 mark b p = do { x <- p; return (b,x) }
 466
 467
 468 -- | Vectorise a function where are the args have scalar type, that is Int, Float or Double.
 469 vectScalarLam
 470         :: [Var]        -- ^ Bound variables of function.
 471         -> CoreExpr     -- ^ Function body.
 472         -> VM VExpr
 473 vectScalarLam args body
 474  = dtrace (vcat [text "vectScalarLam ", ppr args, ppr body])
 475  $ do scalars <- globalScalars
 476       onlyIfV (all is_scalar_ty arg_tys
 477                && is_scalar_ty res_ty
 478                && is_scalar (extendVarSetList scalars args) body
 479                && uses scalars body)
 480         $ do
 481             fn_var  <- hoistExpr (fsLit "fn") (mkLams args body) DontInline
 482             zipf    <- zipScalars arg_tys res_ty
 483             clo     <- scalarClosure arg_tys res_ty (Var fn_var)
 484                                                 (zipf `App` Var fn_var)
 485             clo_var <- hoistExpr (fsLit "clo") clo DontInline
 486             lclo    <- liftPD (Var clo_var)
 487             return (Var clo_var, lclo)
 488   where
 489     arg_tys = map idType args
 490     res_ty  = exprType body
 491
 492     is_scalar_ty ty
 493         | Just (tycon, [])   <- splitTyConApp_maybe ty
 494         =    tycon == intTyCon
 495           || tycon == floatTyCon
 496           || tycon == doubleTyCon
 497
 498         | otherwise = False
 499
 500     is_scalar vs (Var v)     = v `elemVarSet` vs
 501     is_scalar _ e@(Lit _)    = is_scalar_ty $ exprType e
 502     is_scalar vs (App e1 e2) = is_scalar vs e1 && is_scalar vs e2
 503     is_scalar _ _            = False
 504
 505     -- A scalar function has to actually compute something. Without the check,
 506     -- we would treat (\(x :: Int) -> x) as a scalar function and lift it to
 507     -- (map (\x -> x)) which is very bad. Normal lifting transforms it to
 508     -- (\n# x -> x) which is what we want.
 509     uses funs (Var v)     = v `elemVarSet` funs
 510     uses funs (App e1 e2) = uses funs e1 || uses funs e2
 511     uses _ _              = False
 512
 513
 514 vectLam
 515         :: Bool                 -- ^ When the RHS of a binding, whether that binding should be inlined.
 516         -> Bool                 -- ^ Whether the binding is a loop breaker.
 517         -> VarSet               -- ^ The free variables in the body.
 518         -> [Var]                --
 519         -> CoreExprWithFVs
 520         -> VM VExpr
 521
 522 vectLam inline loop_breaker fvs bs body
 523  = dtrace (vcat [ text "vectLam "
 524                 , text "free vars    = " <> ppr fvs
 525                 , text "binding vars = " <> ppr bs
 526                 , text "body         = " <> ppr (deAnnotate body)])
 527
 528  $ do tyvars    <- localTyVars
 529       (vs, vvs) <- readLEnv $ \env ->
 530                    unzip [(var, vv) | var <- varSetElems fvs
 531                                     , Just vv <- [lookupVarEnv (local_vars env) var]]
 532
 533       arg_tys   <- mapM (vectType . idType) bs
 534
 535       dtrace (text "arg_tys = " <> ppr arg_tys) $ return ()
 536
 537       res_ty    <- vectType (exprType $ deAnnotate body)
 538
 539       dtrace (text "res_ty = " <> ppr res_ty) $ return ()
 540
 541       buildClosures tyvars vvs arg_tys res_ty
 542         . hoistPolyVExpr tyvars (maybe_inline (length vs + length bs))
 543         $ do
 544             lc              <- builtin liftingContext
 545             (vbndrs, vbody) <- vectBndrsIn (vs ++ bs) (vectExpr body)
 546
 547             dtrace (text "vbody = " <> ppr vbody) $ return ()
 548
 549             vbody' <- break_loop lc res_ty vbody
 550             return $ vLams lc vbndrs vbody'
 551   where
 552     maybe_inline n | inline    = Inline n
 553                    | otherwise = DontInline
 554
 555     break_loop lc ty (ve, le)
 556       | loop_breaker
 557       = do
 558           empty <- emptyPD ty
 559           lty <- mkPDataType ty
 560           return (ve, mkWildCase (Var lc) intPrimTy lty
 561                         [(DEFAULT, [], le),
 562                          (LitAlt (mkMachInt 0), [], empty)])
 563
 564       | otherwise = return (ve, le)
 565
 566
 567 vectTyAppExpr :: CoreExprWithFVs -> [Type] -> VM VExpr
 568 vectTyAppExpr (_, AnnVar v) tys = vectPolyVar v tys
 569 vectTyAppExpr e tys = cantVectorise "Can't vectorise expression"
 570                         (ppr $ deAnnotate e `mkTyApps` tys)
 571
 572 -- We convert
 573 --
 574 --   case e :: t of v { ... }
 575 --
 576 -- to
 577 --
 578 --   V:    let v' = e in case v' of _ { ... }
 579 --   L:    let v' = e in case v' `cast` ... of _ { ... }
 580 --
 581 -- When lifting, we have to do it this way because v must have the type
 582 -- [:V(T):] but the scrutinee must be cast to the representation type. We also
 583 -- have to handle the case where v is a wild var correctly.
 584 --
 585
 586 -- FIXME: this is too lazy
 587 vectAlgCase :: TyCon -> [Type] -> CoreExprWithFVs -> Var -> Type
 588             -> [(AltCon, [Var], CoreExprWithFVs)]
 589             -> VM VExpr
 590 vectAlgCase _tycon _ty_args scrut bndr ty [(DEFAULT, [], body)]
 591   = do
 592       vscrut         <- vectExpr scrut
 593       (vty, lty)     <- vectAndLiftType ty
 594       (vbndr, vbody) <- vectBndrIn bndr (vectExpr body)
 595       return $ vCaseDEFAULT vscrut vbndr vty lty vbody
 596
 597 vectAlgCase _tycon _ty_args scrut bndr ty [(DataAlt _, [], body)]
 598   = do
 599       vscrut         <- vectExpr scrut
 600       (vty, lty)     <- vectAndLiftType ty
 601       (vbndr, vbody) <- vectBndrIn bndr (vectExpr body)
 602       return $ vCaseDEFAULT vscrut vbndr vty lty vbody
 603
 604 vectAlgCase _tycon _ty_args scrut bndr ty [(DataAlt dc, bndrs, body)]
 605   = do
 606       (vty, lty) <- vectAndLiftType ty
 607       vexpr      <- vectExpr scrut
 608       (vbndr, (vbndrs, (vect_body, lift_body)))
 609          <- vect_scrut_bndr
 610           . vectBndrsIn bndrs
 611           $ vectExpr body
 612       let (vect_bndrs, lift_bndrs) = unzip vbndrs
 613       (vscrut, lscrut, pdata_tc, _arg_tys) <- mkVScrut (vVar vbndr)
 614       vect_dc <- maybeV (lookupDataCon dc)
 615       let [pdata_dc] = tyConDataCons pdata_tc
 616
 617       let vcase = mk_wild_case vscrut vty vect_dc  vect_bndrs vect_body
 618           lcase = mk_wild_case lscrut lty pdata_dc lift_bndrs lift_body
 619
 620       return $ vLet (vNonRec vbndr vexpr) (vcase, lcase)
 621   where
 622     vect_scrut_bndr | isDeadBinder bndr = vectBndrNewIn bndr (fsLit "scrut")
 623                     | otherwise         = vectBndrIn bndr
 624
 625     mk_wild_case expr ty dc bndrs body
 626       = mkWildCase expr (exprType expr) ty [(DataAlt dc, bndrs, body)]
 627
 628 vectAlgCase tycon _ty_args scrut bndr ty alts
 629   = do
 630       vect_tc     <- maybeV (lookupTyCon tycon)
 631       (vty, lty)  <- vectAndLiftType ty
 632
 633       let arity = length (tyConDataCons vect_tc)
 634       sel_ty <- builtin (selTy arity)
 635       sel_bndr <- newLocalVar (fsLit "sel") sel_ty
 636       let sel = Var sel_bndr
 637
 638       (vbndr, valts) <- vect_scrut_bndr
 639                       $ mapM (proc_alt arity sel vty lty) alts'
 640       let (vect_dcs, vect_bndrss, lift_bndrss, vbodies) = unzip4 valts
 641
 642       vexpr <- vectExpr scrut
 643       (vect_scrut, lift_scrut, pdata_tc, _arg_tys) <- mkVScrut (vVar vbndr)
 644       let [pdata_dc] = tyConDataCons pdata_tc
 645
 646       let (vect_bodies, lift_bodies) = unzip vbodies
 647
 648       vdummy <- newDummyVar (exprType vect_scrut)
 649       ldummy <- newDummyVar (exprType lift_scrut)
 650       let vect_case = Case vect_scrut vdummy vty
 651                            (zipWith3 mk_vect_alt vect_dcs vect_bndrss vect_bodies)
 652
 653       lc <- builtin liftingContext
 654       lbody <- combinePD vty (Var lc) sel lift_bodies
 655       let lift_case = Case lift_scrut ldummy lty
 656                            [(DataAlt pdata_dc, sel_bndr : concat lift_bndrss,
 657                              lbody)]
 658
 659       return . vLet (vNonRec vbndr vexpr)
 660              $ (vect_case, lift_case)
 661   where
 662     vect_scrut_bndr | isDeadBinder bndr = vectBndrNewIn bndr (fsLit "scrut")
 663                     | otherwise         = vectBndrIn bndr
 664
 665     alts' = sortBy (\(alt1, _, _) (alt2, _, _) -> cmp alt1 alt2) alts
 666
 667     cmp (DataAlt dc1) (DataAlt dc2) = dataConTag dc1 `compare` dataConTag dc2
 668     cmp DEFAULT       DEFAULT       = EQ
 669     cmp DEFAULT       _             = LT
 670     cmp _             DEFAULT       = GT
 671     cmp _             _             = panic "vectAlgCase/cmp"
 672
 673     proc_alt arity sel _ lty (DataAlt dc, bndrs, body)
 674       = do
 675           vect_dc <- maybeV (lookupDataCon dc)
 676           let ntag = dataConTagZ vect_dc
 677               tag  = mkDataConTag vect_dc
 678               fvs  = freeVarsOf body `delVarSetList` bndrs
 679
 680           sel_tags  <- liftM (`App` sel) (builtin (selTags arity))
 681           lc        <- builtin liftingContext
 682           elems     <- builtin (selElements arity ntag)
 683
 684           (vbndrs, vbody)
 685             <- vectBndrsIn bndrs
 686              . localV
 687              $ do
 688                  binds    <- mapM (pack_var (Var lc) sel_tags tag)
 689                            . filter isLocalId
 690                            $ varSetElems fvs
 691                  (ve, le) <- vectExpr body
 692                  return (ve, Case (elems `App` sel) lc lty
 693                              [(DEFAULT, [], (mkLets (concat binds) le))])
 694                  -- empty    <- emptyPD vty
 695                  -- return (ve, Case (elems `App` sel) lc lty
 696                  --             [(DEFAULT, [], Let (NonRec flags_var flags_expr)
 697                  --                             $ mkLets (concat binds) le),
 698                  --               (LitAlt (mkMachInt 0), [], empty)])
 699           let (vect_bndrs, lift_bndrs) = unzip vbndrs
 700           return (vect_dc, vect_bndrs, lift_bndrs, vbody)
 701
 702     proc_alt _ _ _ _ _ = panic "vectAlgCase/proc_alt"
 703
 704     mk_vect_alt vect_dc bndrs body = (DataAlt vect_dc, bndrs, body)
 705
 706     pack_var len tags t v
 707       = do
 708           r <- lookupVar v
 709           case r of
 710             Local (vv, lv) ->
 711               do
 712                 lv'  <- cloneVar lv
 713                 expr <- packByTagPD (idType vv) (Var lv) len tags t
 714                 updLEnv (\env -> env { local_vars = extendVarEnv
 715                                                 (local_vars env) v (vv, lv') })
 716                 return [(NonRec lv' expr)]
 717
 718             _ -> return []
 719