Generate lots of __inline_me during vectorisation

3 files changed, 31 insertions, 11 deletions

diff --git a/compiler/vectorise/VectCore.hs b/compiler/vectorise/VectCore.hs
index c78f8cad16..6be1542134 100644
--- a/compiler/vectorise/VectCore.hs
+++ b/compiler/vectorise/VectCore.hs
@@ -8,12 +8,13 @@ module VectCore (
 
   vVar, vType, vNote, vLet,
   vLams, vLamsWithoutLC, vVarApps,
-  vCaseDEFAULT, vCaseProd
+  vCaseDEFAULT, vCaseProd, vInlineMe
 ) where
 
 #include "HsVersions.h"
 
 import CoreSyn
+import CoreUtils      ( mkInlineMe )
 import MkCore	      ( mkWildCase )
 import CoreUtils      ( exprType )
 import DataCon        ( DataCon )
@@ -90,3 +91,7 @@ vCaseProd (vscrut, lscrut) vty lty vdc ldc sh_bndrs bndrs
           [(DataAlt ldc, sh_bndrs ++ lbndrs, lbody)])
   where
     (vbndrs, lbndrs) = unzip bndrs
+
+vInlineMe :: VExpr -> VExpr
+vInlineMe (vexpr, lexpr) = (mkInlineMe vexpr, mkInlineMe lexpr)
+
diff --git a/compiler/vectorise/VectUtils.hs b/compiler/vectorise/VectUtils.hs
index 5c014610c5..0d5585fc04 100644
--- a/compiler/vectorise/VectUtils.hs
+++ b/compiler/vectorise/VectUtils.hs
@@ -394,12 +394,13 @@ buildClosures :: [TyVar] -> [VVar] -> [Type] -> Type -> VM VExpr -> VM VExpr
 buildClosures _   _    [] _ mk_body
   = mk_body
 buildClosures tvs vars [arg_ty] res_ty mk_body
-  = buildClosure tvs vars arg_ty res_ty mk_body
+  = liftM vInlineMe (buildClosure tvs vars arg_ty res_ty mk_body)
 buildClosures tvs vars (arg_ty : arg_tys) res_ty mk_body
   = do
       res_ty' <- mkClosureTypes arg_tys res_ty
       arg <- newLocalVVar (fsLit "x") arg_ty
-      buildClosure tvs vars arg_ty res_ty'
+      liftM vInlineMe
+        . buildClosure tvs vars arg_ty res_ty'
         . hoistPolyVExpr tvs
         $ do
             lc <- builtin liftingContext
@@ -424,7 +425,7 @@ buildClosure tvs vars arg_ty res_ty mk_body
               body  <- mk_body
               body' <- bind (vVar env_bndr)
                             (vVarApps lc body (vars ++ [arg_bndr]))
-              return (vLamsWithoutLC [env_bndr, arg_bndr] body')
+              return . vInlineMe $ vLamsWithoutLC [env_bndr, arg_bndr] body'
 
       mkClosure arg_ty res_ty env_ty fn env
 
diff --git a/compiler/vectorise/Vectorise.hs b/compiler/vectorise/Vectorise.hs
index 68b796f82e..777c195a2a 100644
--- a/compiler/vectorise/Vectorise.hs
+++ b/compiler/vectorise/Vectorise.hs
@@ -193,7 +193,7 @@ vectPolyExpr (_, AnnNote note expr)
 vectPolyExpr expr
   = polyAbstract tvs $ \abstract ->
     do
-      mono' <- vectExpr mono
+      mono' <- vectFnExpr False mono
       return $ mapVect abstract mono'
   where
     (tvs, mono) = collectAnnTypeBinders expr
@@ -263,14 +263,26 @@ vectExpr (_, AnnLet (AnnRec bs) body)
                       . inBind bndr
                       $ vectExpr rhs
 
-vectExpr e@(fvs, AnnLam bndr _)
-  | isId bndr = onlyIfV (isEmptyVarSet fvs) (vectScalarLam bs $ deAnnotate body)
-                `orElseV` vectLam fvs bs body
+vectExpr e@(_, AnnLam bndr _)
+  | isId bndr = vectFnExpr True e
+{-
+onlyIfV (isEmptyVarSet fvs) (vectScalarLam bs $ deAnnotate body)
+                `orElseV` vectLam True fvs bs body
   where
     (bs,body) = collectAnnValBinders e
+-}
 
 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
+  where
+    (bs,body) = collectAnnValBinders e
+vectFnExpr _ e = vectExpr e
+
+
 vectScalarLam :: [Var] -> CoreExpr -> VM VExpr
 vectScalarLam args body
   = do
@@ -302,8 +314,8 @@ vectScalarLam args body
     is_scalar vs (App e1 e2) = is_scalar vs e1 && is_scalar vs e2
     is_scalar _ _            = False
 
-vectLam :: VarSet -> [Var] -> CoreExprWithFVs -> VM VExpr
-vectLam fvs bs body
+vectLam :: Bool -> VarSet -> [Var] -> CoreExprWithFVs -> VM VExpr
+vectLam inline fvs bs body
   = do
       tyvars <- localTyVars
       (vs, vvs) <- readLEnv $ \env ->
@@ -319,7 +331,9 @@ vectLam fvs bs body
             lc <- builtin liftingContext
             (vbndrs, vbody) <- vectBndrsIn (vs ++ bs)
                                            (vectExpr body)
-            return $ vLams lc vbndrs vbody
+            return . maybe_inline $ vLams lc vbndrs vbody
+  where
+    maybe_inline = if inline then vInlineMe else id
 
 vectTyAppExpr :: CoreExprWithFVs -> [Type] -> VM VExpr
 vectTyAppExpr (_, AnnVar v) tys = vectPolyVar v tys