X-Git-Url: http://git.megacz.com/?a=blobdiff_plain;f=compiler%2Fvectorise%2FVectorise.hs;h=59e5264aeac72fabee34fe905f73dd12ddce5b8d;hb=ce39c447ab47ac1616cea079210c7651f486f425;hp=055137abe41c5abb102e22d1aababaa1760b72a8;hpb=0a21de62e274acc8e8e260298da4f6c1ee18ecc2;p=ghc-hetmet.git

diff --git a/compiler/vectorise/Vectorise.hs b/compiler/vectorise/Vectorise.hs
index 055137a..59e5264 100644
--- a/compiler/vectorise/Vectorise.hs
+++ b/compiler/vectorise/Vectorise.hs
@@ -42,7 +42,7 @@ import BasicTypes           ( Boxity(..) )
 
 import Outputable
 import FastString
-import Control.Monad        ( liftM, liftM2, mapAndUnzipM )
+import Control.Monad        ( liftM, liftM2, zipWithM, mapAndUnzipM )
 
 vectorise :: HscEnv -> UniqSupply -> RuleBase -> ModGuts
           -> IO (SimplCount, ModGuts)
@@ -81,7 +81,7 @@ vectTopBind :: CoreBind -> VM CoreBind
 vectTopBind b@(NonRec var expr)
   = do
       var'  <- vectTopBinder var
-      expr' <- vectTopRhs expr
+      expr' <- vectTopRhs var expr
       hs    <- takeHoisted
       return . Rec $ (var, expr) : (var', expr') : hs
   `orElseV`
@@ -90,7 +90,7 @@ vectTopBind b@(NonRec var expr)
 vectTopBind b@(Rec bs)
   = do
       vars'  <- mapM vectTopBinder vars
-      exprs' <- mapM vectTopRhs exprs
+      exprs' <- zipWithM vectTopRhs vars exprs
       hs     <- takeHoisted
       return . Rec $ bs ++ zip vars' exprs' ++ hs
   `orElseV`
@@ -108,11 +108,12 @@ vectTopBinder var
       defGlobalVar var var'
       return var'
     
-vectTopRhs :: CoreExpr -> VM CoreExpr
-vectTopRhs expr
+vectTopRhs :: Var -> CoreExpr -> VM CoreExpr
+vectTopRhs var expr
   = do
       lc <- newLocalVar FSLIT("lc") intPrimTy
       closedV . liftM vectorised
+              . inBind var
               $ vectPolyExpr lc (freeVars expr)
 
 -- ----------------------------------------------------------------------------
@@ -228,7 +229,7 @@ vectExpr lc (_, AnnCase expr bndr ty alts)
 
 vectExpr lc (_, AnnLet (AnnNonRec bndr rhs) body)
   = do
-      vrhs <- vectPolyExpr lc rhs
+      vrhs <- localV . inBind bndr $ vectPolyExpr lc rhs
       (vbndr, vbody) <- vectBndrIn bndr (vectExpr lc body)
       return $ vLet (vNonRec vbndr vrhs) vbody
 
@@ -236,32 +237,41 @@ vectExpr lc (_, AnnLet (AnnRec bs) body)
   = do
       (vbndrs, (vrhss, vbody)) <- vectBndrsIn bndrs
                                 $ liftM2 (,)
-                                  (mapM (vectExpr lc) rhss)
+                                  (zipWithM vect_rhs bndrs rhss)
                                   (vectPolyExpr lc body)
       return $ vLet (vRec vbndrs vrhss) vbody
   where
     (bndrs, rhss) = unzip bs
 
-vectExpr lc e@(_, AnnLam bndr body)
-  | isTyVar bndr = pprPanic "vectExpr" (ppr $ deAnnotate e)
+    vect_rhs bndr rhs = localV
+                      . inBind bndr
+                      $ vectExpr lc rhs
 
-vectExpr lc (fvs, AnnLam bndr body)
+vectExpr lc e@(fvs, AnnLam bndr _)
+  | not (isId bndr) = pprPanic "vectExpr" (ppr $ deAnnotate e)
+  | otherwise = vectLam lc fvs bs body
+  where
+    (bs,body) = collectAnnValBinders e
+
+vectLam :: Var -> VarSet -> [Var] -> CoreExprWithFVs -> VM VExpr
+vectLam lc fvs bs body
   = do
       tyvars <- localTyVars
       (vs, vvs) <- readLEnv $ \env ->
                    unzip [(var, vv) | var <- varSetElems fvs
                                     , Just vv <- [lookupVarEnv (local_vars env) var]]
 
-      arg_ty <- vectType (idType bndr)
-      res_ty <- vectType (exprType $ deAnnotate body)
-      buildClosure tyvars lc vvs arg_ty res_ty
-        . hoistPolyVExpr FSLIT("fn") tyvars
+      arg_tys <- mapM (vectType . idType) bs
+      res_ty  <- vectType (exprType $ deAnnotate body)
+
+      buildClosures tyvars lc vvs arg_tys res_ty
+        . hoistPolyVExpr tyvars
         $ do
             new_lc <- newLocalVar FSLIT("lc") intPrimTy
-            (vbndrs, vbody) <- vectBndrsIn (vs ++ [bndr])
+            (vbndrs, vbody) <- vectBndrsIn (vs ++ bs)
                                            (vectExpr new_lc body)
             return $ vLams new_lc vbndrs vbody
-
+  
 vectTyAppExpr :: Var -> CoreExprWithFVs -> [Type] -> VM (CoreExpr, CoreExpr)
 vectTyAppExpr lc (_, AnnVar v) tys = vectPolyVar lc v tys
 vectTyAppExpr lc e tys = pprPanic "vectTyAppExpr" (ppr $ deAnnotate e)