X-Git-Url: http://git.megacz.com/?a=blobdiff_plain;f=compiler%2Fvectorise%2FVectorise%2FExp.hs;h=9cd34e3ac3f54d0a2e0403b5f025a048660ea8da;hb=6815209779aeeedc5d9b79e7c16238c4c658230b;hp=d792fd681973a330709bc9dd741e12c79a9337f9;hpb=1c4593fd4e45a1ea5a32551408866e73840251f1;p=ghc-hetmet.git diff --git a/compiler/vectorise/Vectorise/Exp.hs b/compiler/vectorise/Vectorise/Exp.hs index d792fd6..9cd34e3 100644 --- a/compiler/vectorise/Vectorise/Exp.hs +++ b/compiler/vectorise/Vectorise/Exp.hs @@ -35,21 +35,22 @@ import Data.List -- | Vectorise a polymorphic expression. vectPolyExpr :: Bool -- ^ When vectorising the RHS of a binding, whether that - -- binding is a loop breaker. + -- binding is a loop breaker. + -> [Var] -> CoreExprWithFVs - -> VM (Inline, VExpr) + -> VM (Inline, Bool, VExpr) -vectPolyExpr loop_breaker (_, AnnNote note expr) - = do (inline, expr') <- vectPolyExpr loop_breaker expr - return (inline, vNote note expr') +vectPolyExpr loop_breaker recFns (_, AnnNote note expr) + = do (inline, isScalarFn, expr') <- vectPolyExpr loop_breaker recFns expr + return (inline, isScalarFn, vNote note expr') -vectPolyExpr loop_breaker expr +vectPolyExpr loop_breaker recFns expr = do arity <- polyArity tvs polyAbstract tvs $ \args -> do - (inline, mono') <- vectFnExpr False loop_breaker mono - return (addInlineArity inline arity, + (inline, isScalarFn, mono') <- vectFnExpr False loop_breaker recFns mono + return (addInlineArity inline arity, isScalarFn, mapVect (mkLams $ tvs ++ args) mono') where (tvs, mono) = collectAnnTypeBinders expr @@ -111,12 +112,13 @@ vectExpr (_, AnnCase scrut bndr ty alts) | Just (tycon, ty_args) <- splitTyConApp_maybe scrut_ty , isAlgTyCon tycon = vectAlgCase tycon ty_args scrut bndr ty alts + | otherwise = cantVectorise "Can't vectorise expression" (ppr scrut_ty) where scrut_ty = exprType (deAnnotate scrut) vectExpr (_, AnnLet (AnnNonRec bndr rhs) body) = do - vrhs <- localV . inBind bndr . liftM snd $ vectPolyExpr False rhs + vrhs <- localV . inBind bndr . liftM (\(_,_,z)->z) $ vectPolyExpr False [] rhs (vbndr, vbody) <- vectBndrIn bndr (vectExpr body) return $ vLet (vNonRec vbndr vrhs) vbody @@ -132,11 +134,11 @@ vectExpr (_, AnnLet (AnnRec bs) body) vect_rhs bndr rhs = localV . inBind bndr - . liftM snd - $ vectPolyExpr (isLoopBreaker $ idOccInfo bndr) rhs + . liftM (\(_,_,z)->z) + $ vectPolyExpr (isLoopBreaker $ idOccInfo bndr) [] rhs vectExpr e@(_, AnnLam bndr _) - | isId bndr = liftM snd $ vectFnExpr True False e + | isId bndr = liftM (\(_,_,z) ->z) $ vectFnExpr True False [] e {- onlyIfV (isEmptyVarSet fvs) (vectScalarLam bs $ deAnnotate body) `orElseV` vectLam True fvs bs body @@ -144,40 +146,43 @@ onlyIfV (isEmptyVarSet fvs) (vectScalarLam bs $ deAnnotate body) (bs,body) = collectAnnValBinders e -} -vectExpr e = cantVectorise "Can't vectorise expression" (ppr $ deAnnotate e) +vectExpr e = cantVectorise "Can't vectorise expression (vectExpr)" (ppr $ deAnnotate e) -- | Vectorise an expression with an outer lambda abstraction. vectFnExpr :: Bool -- ^ When the RHS of a binding, whether that binding should be inlined. -> Bool -- ^ Whether the binding is a loop breaker. + -> [Var] -> CoreExprWithFVs -- ^ Expression to vectorise. Must have an outer `AnnLam`. - -> VM (Inline, VExpr) + -> VM (Inline, Bool, 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) +vectFnExpr inline loop_breaker recFns e@(fvs, AnnLam bndr _) + | isId bndr = onlyIfV True -- (isEmptyVarSet fvs) -- we check for free variables later. TODO: clean up + (mark DontInline True . vectScalarLam bs recFns $ deAnnotate body) + `orElseV` mark inlineMe False (vectLam inline loop_breaker fvs bs body) where (bs,body) = collectAnnValBinders e -vectFnExpr _ _ e = mark DontInline $ vectExpr e +vectFnExpr _ _ _ e = mark DontInline False $ vectExpr e -mark :: Inline -> VM a -> VM (Inline, a) -mark b p = do { x <- p; return (b,x) } +mark :: Inline -> Bool -> VM a -> VM (Inline, Bool, a) +mark b isScalarFn p = do { x <- p; return (b, isScalarFn, x) } -- | Vectorise a function where are the args have scalar type, -- that is Int, Float, Double etc. vectScalarLam - :: [Var] -- ^ Bound variables of function. + :: [Var] -- ^ Bound variables of function + -> [Var] -> CoreExpr -- ^ Function body. -> VM VExpr -vectScalarLam args body - = do scalars <- globalScalars - onlyIfV (all is_scalar_ty arg_tys - && is_scalar_ty res_ty +vectScalarLam args recFns body + = do scalars' <- globalScalars + let scalars = unionVarSet (mkVarSet recFns) scalars' + onlyIfV (all is_prim_ty arg_tys + && is_prim_ty res_ty && is_scalar (extendVarSetList scalars args) body && uses scalars body) $ do @@ -202,14 +207,15 @@ vectScalarLam args body cantbe_parr_expr expr = not $ maybe_parr_ty $ exprType expr - maybe_parr_ty ty = maybe_parr_ty' [] ty + maybe_parr_ty ty = maybe_parr_ty' [] ty + + maybe_parr_ty' _ ty | Nothing <- splitTyConApp_maybe ty = False -- TODO: is this really what we want to do with polym. types? maybe_parr_ty' alreadySeen ty | isPArrTyCon tycon = True | isPrimTyCon tycon = False | isAbstractTyCon tycon = True | isFunTyCon tycon || isProductTyCon tycon || isTupleTyCon tycon = any (maybe_parr_ty' alreadySeen) args - | isDataTyCon tycon = pprTrace "isDataTyCon" (ppr tycon) $ - any (maybe_parr_ty' alreadySeen) args || + | isDataTyCon tycon = any (maybe_parr_ty' alreadySeen) args || hasParrDataCon alreadySeen tycon | otherwise = True where @@ -226,31 +232,25 @@ vectScalarLam args body is_scalar vs e@(Var v) | Just _ <- isDataConId_maybe v = cantbe_parr_expr e | otherwise = cantbe_parr_expr e && (v `elemVarSet` vs) - is_scalar _ e@(Lit _) = -- pprTrace "is_scalar Lit" (ppr e) $ - cantbe_parr_expr e + is_scalar _ e@(Lit _) = cantbe_parr_expr e - is_scalar vs e@(App e1 e2) = -- pprTrace "is_scalar App" (ppr e) $ - cantbe_parr_expr e && + is_scalar vs e@(App e1 e2) = cantbe_parr_expr e && is_scalar vs e1 && is_scalar vs e2 is_scalar vs e@(Let (NonRec b letExpr) body) - = -- pprTrace "is_scalar Let" (ppr e) $ - cantbe_parr_expr e && + = cantbe_parr_expr e && is_scalar vs letExpr && is_scalar (extendVarSet vs b) body - is_scalar vs e@(Let (Rec bnds) body) + is_scalar vs e@(Let (Rec bnds) body) = let vs' = extendVarSetList vs (map fst bnds) - in -- pprTrace "is_scalar Rec" (ppr e) $ - cantbe_parr_expr e && + in cantbe_parr_expr e && all (is_scalar vs') (map snd bnds) && is_scalar vs' body is_scalar vs e@(Case eC eId ty alts) = let vs' = extendVarSet vs eId - in -- pprTrace "is_scalar Case" (ppr e) $ - cantbe_parr_expr e && + in cantbe_parr_expr e && is_prim_ty ty && is_scalar vs' eC && (all (is_scalar_alt vs') alts) - is_scalar _ e = -- pprTrace "is_scalar other" (ppr e) $ - False + is_scalar _ _ = False is_scalar_alt vs (_, bs, e) = is_scalar (extendVarSetList vs bs) e @@ -314,7 +314,7 @@ vectLam inline loop_breaker fvs bs body vectTyAppExpr :: CoreExprWithFVs -> [Type] -> VM VExpr vectTyAppExpr (_, AnnVar v) tys = vectPolyVar v tys -vectTyAppExpr e tys = cantVectorise "Can't vectorise expression" +vectTyAppExpr e tys = cantVectorise "Can't vectorise expression (vectTyExpr)" (ppr $ deAnnotate e `mkTyApps` tys)