X-Git-Url: http://git.megacz.com/?a=blobdiff_plain;f=compiler%2Fvectorise%2FVectorise.hs;h=59fded3c4f01c552c081869fd34c7698b2e285d6;hb=222415a5b658e737a0a1f2c980c6f80635289f75;hp=27cdde35a089d1f9724c741bddb3ce2f3abb6018;hpb=3736e30f683990ee94055b60905cce208a467e8b;p=ghc-hetmet.git diff --git a/compiler/vectorise/Vectorise.hs b/compiler/vectorise/Vectorise.hs index 27cdde3..59fded3 100644 --- a/compiler/vectorise/Vectorise.hs +++ b/compiler/vectorise/Vectorise.hs @@ -12,6 +12,7 @@ import HscTypes hiding ( MonadThings(..) ) import Module ( PackageId ) import CoreSyn import CoreUtils +import CoreUnfold ( mkInlineRule ) import MkCore ( mkWildCase ) import CoreFVs import CoreMonad ( CoreM, getHscEnv ) @@ -24,13 +25,16 @@ import VarEnv import VarSet import Id import OccName +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 @@ -66,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 @@ -76,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 @@ -86,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 @@ -186,14 +203,19 @@ vectLiteral 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 @@ -244,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 @@ -253,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 @@ -273,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 @@ -290,11 +316,11 @@ 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 + clo_var <- hoistExpr (fsLit "clo") clo DontInline lclo <- liftPD (Var clo_var) return (Var clo_var, lclo) where @@ -313,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 -> @@ -325,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 @@ -371,9 +411,8 @@ 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, (vect_body, lift_body))) @@ -441,16 +480,14 @@ vectAlgCase tycon _ty_args scrut bndr ty alts cmp _ DEFAULT = GT cmp _ _ = panic "vectAlgCase/cmp" - proc_alt arity sel vty lty (DataAlt dc, bndrs, body) + proc_alt arity sel _ lty (DataAlt dc, bndrs, body) = do vect_dc <- maybeV (lookupDataCon dc) let ntag = dataConTagZ vect_dc tag = mkDataConTag vect_dc fvs = freeVarsOf body `delVarSetList` bndrs - pick <- builtin (selPick arity) - let flags_expr = mkApps pick [sel, tag] - flags_var <- newLocalVar (fsLit "flags") (exprType flags_expr) + sel_tags <- liftM (`App` sel) (builtin (selTags arity)) lc <- builtin liftingContext elems <- builtin (selElements arity ntag) @@ -458,15 +495,17 @@ vectAlgCase tycon _ty_args scrut bndr ty alts <- vectBndrsIn bndrs . localV $ do - binds <- mapM (pack_var (Var lc) (Var flags_var)) + binds <- mapM (pack_var (Var lc) sel_tags tag) . filter isLocalId $ varSetElems fvs (ve, le) <- vectExpr body - 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)]) + [(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) @@ -474,14 +513,14 @@ vectAlgCase tycon _ty_args scrut bndr ty alts mk_vect_alt vect_dc bndrs body = (DataAlt vect_dc, bndrs, body) - pack_var len flags v + pack_var len tags t v = do r <- lookupVar v case r of Local (vv, lv) -> do lv' <- cloneVar lv - expr <- packPD (idType vv) (Var lv) len flags + 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)]