X-Git-Url: http://git.megacz.com/?a=blobdiff_plain;f=compiler%2Fvectorise%2FVectorise.hs;h=777c195a2a37d8b322524943e763739157fc711b;hb=7bb3d1fc79521d591cd9f824893963141a7997b6;hp=92a19d562c9c6271cf247e77bf812026e6716280;hpb=84dc1adf521191e1c171e684c7cd28f03475125e;p=ghc-hetmet.git

diff --git a/compiler/vectorise/Vectorise.hs b/compiler/vectorise/Vectorise.hs
index 92a19d5..777c195 100644
--- a/compiler/vectorise/Vectorise.hs
+++ b/compiler/vectorise/Vectorise.hs
@@ -7,16 +7,13 @@ import VectUtils
 import VectType
 import VectCore
 
-import DynFlags
 import HscTypes hiding      ( MonadThings(..) )
 
 import Module               ( PackageId )
-import CoreLint             ( showPass, endPass )
 import CoreSyn
 import CoreUtils
 import CoreFVs
-import SimplMonad           ( SimplCount, zeroSimplCount )
-import Rules                ( RuleBase )
+import CoreMonad            ( CoreM, getHscEnv, liftIO )
 import DataCon
 import TyCon
 import Type
@@ -37,18 +34,18 @@ import FastString
 import Control.Monad        ( liftM, liftM2, zipWithM )
 import Data.List            ( sortBy, unzip4 )
 
-vectorise :: PackageId -> HscEnv -> UniqSupply -> RuleBase -> ModGuts
-          -> IO (SimplCount, ModGuts)
-vectorise backend hsc_env _ _ guts
+vectorise :: PackageId -> ModGuts -> CoreM ModGuts
+vectorise backend guts = do
+    hsc_env <- getHscEnv
+    liftIO $ vectoriseIO backend hsc_env guts
+
+vectoriseIO :: PackageId -> HscEnv -> ModGuts -> IO ModGuts
+vectoriseIO backend hsc_env guts
   = do
-      showPass dflags "Vectorisation"
       eps <- hscEPS hsc_env
       let info = hptVectInfo hsc_env `plusVectInfo` eps_vect_info eps
       Just (info', guts') <- initV backend hsc_env guts info (vectModule guts)
-      endPass dflags "Vectorisation" Opt_D_dump_vect (mg_binds guts')
-      return (zeroSimplCount dflags, guts' { mg_vect_info = info' })
-  where
-    dflags = hsc_dflags hsc_env
+      return (guts' { mg_vect_info = info' })
 
 vectModule :: ModGuts -> VM ModGuts
 vectModule guts
@@ -196,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
@@ -266,15 +263,59 @@ vectExpr (_, AnnLet (AnnRec bs) body)
                       . inBind bndr
                       $ vectExpr rhs
 
-vectExpr e@(fvs, AnnLam bndr _)
-  | isId bndr = 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)
 
-vectLam :: VarSet -> [Var] -> CoreExprWithFVs -> VM VExpr
-vectLam fvs bs body
+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
+      scalars <- globalScalars
+      onlyIfV (all is_scalar_ty arg_tys
+               && is_scalar_ty res_ty
+               && is_scalar (extendVarSetList scalars args) body)
+        $ do
+            fn_var <- hoistExpr (fsLit "fn") (mkLams args body)
+            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)
+            return (Var clo_var, lclo)
+  where
+    arg_tys = map idType args
+    res_ty  = exprType body
+
+    is_scalar_ty ty | Just (tycon, []) <- splitTyConApp_maybe ty
+                    = tycon == intTyCon
+                      || tycon == floatTyCon
+                      || tycon == doubleTyCon
+
+                    | otherwise = False
+
+    is_scalar vs (Var v)     = v `elemVarSet` vs
+    is_scalar _ e@(Lit _)    = is_scalar_ty $ exprType e
+    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
   = do
       tyvars <- localTyVars
       (vs, vvs) <- readLEnv $ \env ->
@@ -290,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