X-Git-Url: http://git.megacz.com/?a=blobdiff_plain;f=compiler%2Fvectorise%2FVectorise.hs;h=59fded3c4f01c552c081869fd34c7698b2e285d6;hb=222415a5b658e737a0a1f2c980c6f80635289f75;hp=777c195a2a37d8b322524943e763739157fc711b;hpb=bee06bad431d372bd862b5c6e921d8fc87eaffc9;p=ghc-hetmet.git diff --git a/compiler/vectorise/Vectorise.hs b/compiler/vectorise/Vectorise.hs index 777c195..59fded3 100644 --- a/compiler/vectorise/Vectorise.hs +++ b/compiler/vectorise/Vectorise.hs @@ -12,8 +12,10 @@ import HscTypes hiding ( MonadThings(..) ) import Module ( PackageId ) import CoreSyn import CoreUtils +import CoreUnfold ( mkInlineRule ) +import MkCore ( mkWildCase ) import CoreFVs -import CoreMonad ( CoreM, getHscEnv, liftIO ) +import CoreMonad ( CoreM, getHscEnv ) import DataCon import TyCon import Type @@ -23,15 +25,16 @@ import VarEnv import VarSet import Id import OccName - -import DsMonad +import BasicTypes ( isLoopBreaker ) import Literal ( Literal, mkMachInt ) import TysWiredIn +import TysPrim ( intPrimTy ) import Outputable import FastString -import Control.Monad ( liftM, liftM2, zipWithM ) +import Util ( zipLazy ) +import Control.Monad import Data.List ( sortBy, unzip4 ) vectorise :: PackageId -> ModGuts -> CoreM ModGuts @@ -67,8 +70,8 @@ vectModule guts vectTopBind :: CoreBind -> VM CoreBind vectTopBind b@(NonRec var expr) = do - var' <- vectTopBinder var - expr' <- vectTopRhs var expr + (inline, expr') <- vectTopRhs var expr + var' <- vectTopBinder var inline expr' hs <- takeHoisted cexpr <- tryConvert var var' expr return . Rec $ (var, cexpr) : (var', expr') : hs @@ -77,8 +80,13 @@ vectTopBind b@(NonRec var expr) vectTopBind b@(Rec bs) = do - vars' <- mapM vectTopBinder vars - exprs' <- zipWithM vectTopRhs vars exprs + (vars', _, exprs') <- fixV $ \ ~(_, inlines, rhss) -> + do + vars' <- sequence [vectTopBinder var inline rhs + | (var, ~(inline, rhs)) + <- zipLazy vars (zip inlines rhss)] + (inlines', exprs') <- mapAndUnzipM (uncurry vectTopRhs) bs + return (vars', inlines', exprs') hs <- takeHoisted cexprs <- sequence $ zipWith3 tryConvert vars vars' exprs return . Rec $ zip vars cexprs ++ zip vars' exprs' ++ hs @@ -87,20 +95,28 @@ vectTopBind b@(Rec bs) where (vars, exprs) = unzip bs -vectTopBinder :: Var -> VM Var -vectTopBinder var +-- NOTE: vectTopBinder *MUST* be lazy in inline and expr because of how it is +-- used inside of fixV in vectTopBind +vectTopBinder :: Var -> Inline -> CoreExpr -> VM Var +vectTopBinder var inline expr = do vty <- vectType (idType var) - var' <- cloneId mkVectOcc var vty + var' <- liftM (`setIdUnfolding` unfolding) $ cloneId mkVectOcc var vty defGlobalVar var var' return var' + where + unfolding = case inline of + Inline arity -> mkInlineRule InlSat expr arity + DontInline -> noUnfolding -vectTopRhs :: Var -> CoreExpr -> VM CoreExpr +vectTopRhs :: Var -> CoreExpr -> VM (Inline, CoreExpr) vectTopRhs var expr - = do - closedV . liftM vectorised - . inBind var - $ vectPolyExpr (freeVars expr) + = closedV + $ do + (inline, vexpr) <- inBind var + $ vectPolyExpr (isLoopBreaker $ idOccInfo var) + (freeVars expr) + return (inline, vectorised vexpr) tryConvert :: Var -> Var -> CoreExpr -> VM CoreExpr tryConvert var vect_var rhs @@ -165,7 +181,7 @@ vectVar v Local (vv,lv) -> return (Var vv, Var lv) Global vv -> do let vexpr = Var vv - lexpr <- liftPA vexpr + lexpr <- liftPD vexpr return (vexpr, lexpr) vectPolyVar :: Var -> [Type] -> VM VExpr @@ -178,23 +194,28 @@ vectPolyVar v tys (polyApply (Var lv) vtys) Global poly -> do vexpr <- polyApply (Var poly) vtys - lexpr <- liftPA vexpr + lexpr <- liftPD vexpr return (vexpr, lexpr) vectLiteral :: Literal -> VM VExpr vectLiteral lit = do - lexpr <- liftPA (Lit lit) + lexpr <- liftPD (Lit lit) return (Lit lit, lexpr) -vectPolyExpr :: CoreExprWithFVs -> VM VExpr -vectPolyExpr (_, AnnNote note expr) - = liftM (vNote note) $ vectPolyExpr expr -vectPolyExpr expr - = polyAbstract tvs $ \abstract -> - do - mono' <- vectFnExpr False mono - return $ mapVect abstract mono' +vectPolyExpr :: Bool -> CoreExprWithFVs -> VM (Inline, VExpr) +vectPolyExpr loop_breaker (_, AnnNote note expr) + = do + (inline, expr') <- vectPolyExpr loop_breaker expr + return (inline, vNote note expr') +vectPolyExpr loop_breaker expr + = do + arity <- polyArity tvs + polyAbstract tvs $ \args -> + do + (inline, mono') <- vectFnExpr False loop_breaker mono + return (addInlineArity inline arity, + mapVect (mkLams $ tvs ++ args) mono') where (tvs, mono) = collectAnnTypeBinders expr @@ -220,7 +241,7 @@ vectExpr (_, AnnApp (_, AnnVar v) (_, AnnLit lit)) , is_special_con con = do let vexpr = App (Var v) (Lit lit) - lexpr <- liftPA vexpr + lexpr <- liftPD vexpr return (vexpr, lexpr) where is_special_con con = con `elem` [intDataCon, floatDataCon, doubleDataCon] @@ -245,7 +266,7 @@ vectExpr (_, AnnCase scrut bndr ty alts) vectExpr (_, AnnLet (AnnNonRec bndr rhs) body) = do - vrhs <- localV . inBind bndr $ vectPolyExpr rhs + vrhs <- localV . inBind bndr . liftM snd $ vectPolyExpr False rhs (vbndr, vbody) <- vectBndrIn bndr (vectExpr body) return $ vLet (vNonRec vbndr vrhs) vbody @@ -254,17 +275,18 @@ vectExpr (_, AnnLet (AnnRec bs) body) (vbndrs, (vrhss, vbody)) <- vectBndrsIn bndrs $ liftM2 (,) (zipWithM vect_rhs bndrs rhss) - (vectPolyExpr body) + (vectExpr body) return $ vLet (vRec vbndrs vrhss) vbody where (bndrs, rhss) = unzip bs vect_rhs bndr rhs = localV . inBind bndr - $ vectExpr rhs + . liftM snd + $ vectPolyExpr (isLoopBreaker $ idOccInfo bndr) rhs vectExpr e@(_, AnnLam bndr _) - | isId bndr = vectFnExpr True e + | isId bndr = liftM snd $ vectFnExpr True False e {- onlyIfV (isEmptyVarSet fvs) (vectScalarLam bs $ deAnnotate body) `orElseV` vectLam True fvs bs body @@ -274,14 +296,17 @@ onlyIfV (isEmptyVarSet fvs) (vectScalarLam bs $ deAnnotate body) vectExpr e = cantVectorise "Can't vectorise expression" (ppr $ deAnnotate e) -vectFnExpr :: Bool -> CoreExprWithFVs -> VM VExpr -vectFnExpr inline e@(fvs, AnnLam bndr _) - | isId bndr = onlyIfV (isEmptyVarSet fvs) (vectScalarLam bs $ deAnnotate body) - `orElseV` vectLam inline fvs bs body +vectFnExpr :: Bool -> Bool -> CoreExprWithFVs -> VM (Inline, VExpr) +vectFnExpr inline loop_breaker e@(fvs, AnnLam bndr _) + | isId bndr = onlyIfV (isEmptyVarSet fvs) + (mark DontInline . vectScalarLam bs $ deAnnotate body) + `orElseV` mark inlineMe (vectLam inline loop_breaker fvs bs body) where (bs,body) = collectAnnValBinders e -vectFnExpr _ e = vectExpr e +vectFnExpr _ _ e = mark DontInline $ vectExpr e +mark :: Inline -> VM a -> VM (Inline, a) +mark b p = do { x <- p; return (b,x) } vectScalarLam :: [Var] -> CoreExpr -> VM VExpr vectScalarLam args body @@ -291,12 +316,12 @@ vectScalarLam args body && is_scalar_ty res_ty && is_scalar (extendVarSetList scalars args) body) $ do - fn_var <- hoistExpr (fsLit "fn") (mkLams args body) + fn_var <- hoistExpr (fsLit "fn") (mkLams args body) DontInline zipf <- zipScalars arg_tys res_ty clo <- scalarClosure arg_tys res_ty (Var fn_var) (zipf `App` Var fn_var) - clo_var <- hoistExpr (fsLit "clo") clo - lclo <- liftPA (Var clo_var) + clo_var <- hoistExpr (fsLit "clo") clo DontInline + lclo <- liftPD (Var clo_var) return (Var clo_var, lclo) where arg_tys = map idType args @@ -314,8 +339,8 @@ vectScalarLam args body is_scalar vs (App e1 e2) = is_scalar vs e1 && is_scalar vs e2 is_scalar _ _ = False -vectLam :: Bool -> VarSet -> [Var] -> CoreExprWithFVs -> VM VExpr -vectLam inline fvs bs body +vectLam :: Bool -> Bool -> VarSet -> [Var] -> CoreExprWithFVs -> VM VExpr +vectLam inline loop_breaker fvs bs body = do tyvars <- localTyVars (vs, vvs) <- readLEnv $ \env -> @@ -326,14 +351,28 @@ vectLam inline fvs bs body res_ty <- vectType (exprType $ deAnnotate body) buildClosures tyvars vvs arg_tys res_ty - . hoistPolyVExpr tyvars + . hoistPolyVExpr tyvars (maybe_inline (length vs + length bs)) $ do lc <- builtin liftingContext (vbndrs, vbody) <- vectBndrsIn (vs ++ bs) (vectExpr body) - return . maybe_inline $ vLams lc vbndrs vbody + vbody' <- break_loop lc res_ty vbody + return $ vLams lc vbndrs vbody' where - maybe_inline = if inline then vInlineMe else id + maybe_inline n | inline = Inline n + | otherwise = DontInline + + break_loop lc ty (ve, le) + | loop_breaker + = do + empty <- emptyPD ty + lty <- mkPDataType ty + return (ve, mkWildCase (Var lc) intPrimTy lty + [(DEFAULT, [], le), + (LitAlt (mkMachInt 0), [], empty)]) + + | otherwise = return (ve, le) + vectTyAppExpr :: CoreExprWithFVs -> [Type] -> VM VExpr vectTyAppExpr (_, AnnVar v) tys = vectPolyVar v tys @@ -372,52 +411,59 @@ vectAlgCase _tycon _ty_args scrut bndr ty [(DataAlt _, [], body)] (vbndr, vbody) <- vectBndrIn bndr (vectExpr body) return $ vCaseDEFAULT vscrut vbndr vty lty vbody -vectAlgCase tycon _ty_args scrut bndr ty [(DataAlt dc, bndrs, body)] +vectAlgCase _tycon _ty_args scrut bndr ty [(DataAlt dc, bndrs, body)] = do - vect_tc <- maybeV (lookupTyCon tycon) (vty, lty) <- vectAndLiftType ty vexpr <- vectExpr scrut - (vbndr, (vbndrs, vbody)) <- vect_scrut_bndr - . vectBndrsIn bndrs - $ vectExpr body - - (vscrut, arr_tc, _arg_tys) <- mkVScrut (vVar vbndr) + (vbndr, (vbndrs, (vect_body, lift_body))) + <- vect_scrut_bndr + . vectBndrsIn bndrs + $ vectExpr body + let (vect_bndrs, lift_bndrs) = unzip vbndrs + (vscrut, lscrut, pdata_tc, _arg_tys) <- mkVScrut (vVar vbndr) vect_dc <- maybeV (lookupDataCon dc) - let [arr_dc] = tyConDataCons arr_tc - repr <- mkRepr vect_tc - shape_bndrs <- arrShapeVars repr - return . vLet (vNonRec vbndr vexpr) - $ vCaseProd vscrut vty lty vect_dc arr_dc shape_bndrs vbndrs vbody + let [pdata_dc] = tyConDataCons pdata_tc + + let vcase = mk_wild_case vscrut vty vect_dc vect_bndrs vect_body + lcase = mk_wild_case lscrut lty pdata_dc lift_bndrs lift_body + + return $ vLet (vNonRec vbndr vexpr) (vcase, lcase) where vect_scrut_bndr | isDeadBinder bndr = vectBndrNewIn bndr (fsLit "scrut") | otherwise = vectBndrIn bndr + mk_wild_case expr ty dc bndrs body + = mkWildCase expr (exprType expr) ty [(DataAlt dc, bndrs, body)] + vectAlgCase tycon _ty_args scrut bndr ty alts = do vect_tc <- maybeV (lookupTyCon tycon) (vty, lty) <- vectAndLiftType ty - repr <- mkRepr vect_tc - shape_bndrs <- arrShapeVars repr - (len, sel, indices) <- arrSelector repr (map Var shape_bndrs) - (vbndr, valts) <- vect_scrut_bndr $ mapM (proc_alt sel vty lty) alts' + let arity = length (tyConDataCons vect_tc) + sel_ty <- builtin (selTy arity) + sel_bndr <- newLocalVar (fsLit "sel") sel_ty + let sel = Var sel_bndr + + (vbndr, valts) <- vect_scrut_bndr + $ mapM (proc_alt arity sel vty lty) alts' let (vect_dcs, vect_bndrss, lift_bndrss, vbodies) = unzip4 valts vexpr <- vectExpr scrut - (vscrut, arr_tc, _arg_tys) <- mkVScrut (vVar vbndr) - let [arr_dc] = tyConDataCons arr_tc + (vect_scrut, lift_scrut, pdata_tc, _arg_tys) <- mkVScrut (vVar vbndr) + let [pdata_dc] = tyConDataCons pdata_tc - let (vect_scrut, lift_scrut) = vscrut - (vect_bodies, lift_bodies) = unzip vbodies + let (vect_bodies, lift_bodies) = unzip vbodies vdummy <- newDummyVar (exprType vect_scrut) ldummy <- newDummyVar (exprType lift_scrut) let vect_case = Case vect_scrut vdummy vty (zipWith3 mk_vect_alt vect_dcs vect_bndrss vect_bodies) - lbody <- combinePA vty len sel indices lift_bodies + lc <- builtin liftingContext + lbody <- combinePD vty (Var lc) sel lift_bodies let lift_case = Case lift_scrut ldummy lty - [(DataAlt arr_dc, shape_bndrs ++ concat lift_bndrss, + [(DataAlt pdata_dc, sel_bndr : concat lift_bndrss, lbody)] return . vLet (vNonRec vbndr vexpr) @@ -434,75 +480,50 @@ vectAlgCase tycon _ty_args scrut bndr ty alts cmp _ DEFAULT = GT cmp _ _ = panic "vectAlgCase/cmp" - proc_alt sel vty lty (DataAlt dc, bndrs, body) + proc_alt arity sel _ lty (DataAlt dc, bndrs, body) = do vect_dc <- maybeV (lookupDataCon dc) - let tag = mkDataConTag vect_dc - fvs = freeVarsOf body `delVarSetList` bndrs - (vect_bndrs, lift_bndrs, vbody) - <- vect_alt_bndrs bndrs - $ \len -> packLiftingContext len sel tag fvs vty lty - $ vectExpr body - + let ntag = dataConTagZ vect_dc + tag = mkDataConTag vect_dc + fvs = freeVarsOf body `delVarSetList` bndrs + + sel_tags <- liftM (`App` sel) (builtin (selTags arity)) + lc <- builtin liftingContext + elems <- builtin (selElements arity ntag) + + (vbndrs, vbody) + <- vectBndrsIn bndrs + . localV + $ do + binds <- mapM (pack_var (Var lc) sel_tags tag) + . filter isLocalId + $ varSetElems fvs + (ve, le) <- vectExpr body + return (ve, Case (elems `App` sel) lc lty + [(DEFAULT, [], (mkLets (concat binds) le))]) + -- empty <- emptyPD vty + -- return (ve, Case (elems `App` sel) lc lty + -- [(DEFAULT, [], Let (NonRec flags_var flags_expr) + -- $ mkLets (concat binds) le), + -- (LitAlt (mkMachInt 0), [], empty)]) + let (vect_bndrs, lift_bndrs) = unzip vbndrs return (vect_dc, vect_bndrs, lift_bndrs, vbody) - proc_alt _ _ _ _ = panic "vectAlgCase/proc_alt" - vect_alt_bndrs [] p - = do - void_tc <- builtin voidTyCon - let void_ty = mkTyConApp void_tc [] - arr_ty <- mkPArrayType void_ty - bndr <- newLocalVar (fsLit "voids") arr_ty - len <- lengthPA void_ty (Var bndr) - e <- p len - return ([], [bndr], e) - - vect_alt_bndrs bndrs p - = localV - $ do - vbndrs <- mapM vectBndr bndrs - let (vect_bndrs, lift_bndrs) = unzip vbndrs - vv : _ = vect_bndrs - lv : _ = lift_bndrs - len <- lengthPA (idType vv) (Var lv) - e <- p len - return (vect_bndrs, lift_bndrs, e) + proc_alt _ _ _ _ _ = panic "vectAlgCase/proc_alt" mk_vect_alt vect_dc bndrs body = (DataAlt vect_dc, bndrs, body) -packLiftingContext :: CoreExpr -> CoreExpr -> CoreExpr -> VarSet - -> Type -> Type -> VM VExpr -> VM VExpr -packLiftingContext len shape tag fvs vty lty p - = do - select <- builtin selectPAIntPrimVar - let sel_expr = mkApps (Var select) [shape, tag] - sel_var <- newLocalVar (fsLit "sel#") (exprType sel_expr) - lc_var <- builtin liftingContext - localV $ - do - bnds <- mapM (packFreeVar (Var lc_var) (Var sel_var)) - . filter isLocalId - $ varSetElems fvs - (vexpr, lexpr) <- p - empty <- emptyPA vty - return (vexpr, Let (NonRec sel_var sel_expr) - $ Case len lc_var lty - [(DEFAULT, [], mkLets (concat bnds) lexpr), - (LitAlt (mkMachInt 0), [], empty)]) - -packFreeVar :: CoreExpr -> CoreExpr -> Var -> VM [CoreBind] -packFreeVar len sel v - = do - r <- lookupVar v - case r of - Local (vv,lv) -> - do - lv' <- cloneVar lv - expr <- packPA (idType vv) (Var lv) len sel - updLEnv (upd vv lv') - return [(NonRec lv' expr)] - - _ -> return [] - where - upd vv lv' env = env { local_vars = extendVarEnv (local_vars env) v (vv, lv') } + pack_var len tags t v + = do + r <- lookupVar v + case r of + Local (vv, lv) -> + do + lv' <- cloneVar lv + expr <- packByTagPD (idType vv) (Var lv) len tags t + updLEnv (\env -> env { local_vars = extendVarEnv + (local_vars env) v (vv, lv') }) + return [(NonRec lv' expr)] + + _ -> return []