Teach vectorisation about tuple datacons

[ghc-hetmet.git] / compiler / vectorise / Vectorise.hs

diff --git a/compiler/vectorise/Vectorise.hs b/compiler/vectorise/Vectorise.hs
index 96ed01e..d85ef6a 100644 (file)
--- a/compiler/vectorise/Vectorise.hs
+++ b/compiler/vectorise/Vectorise.hs
@@ -89,7 +89,8 @@ vectTopBind b@(NonRec var expr)
       var'  <- vectTopBinder var
       expr' <- vectTopRhs var expr
       hs    <- takeHoisted
-      return . Rec $ (var, expr) : (var', expr') : hs
+      cexpr <- tryConvert var var' expr
+      return . Rec $ (var, cexpr) : (var', expr') : hs
   `orElseV`
     return b
 
@@ -98,7 +99,8 @@ vectTopBind b@(Rec bs)
       vars'  <- mapM vectTopBinder vars
       exprs' <- zipWithM vectTopRhs vars exprs
       hs     <- takeHoisted
-      return . Rec $ bs ++ zip vars' exprs' ++ hs
+      cexprs <- sequence $ zipWith3 tryConvert vars vars' exprs
+      return . Rec $ zip vars cexprs ++ zip vars' exprs' ++ hs
   `orElseV`
     return b
   where
@@ -119,6 +121,10 @@ vectTopRhs var expr
               . inBind var
               $ vectPolyExpr (freeVars expr)
 
+tryConvert :: Var -> Var -> CoreExpr -> VM CoreExpr
+tryConvert var vect_var rhs
+  = fromVect (idType var) (Var vect_var) `orElseV` return rhs
+
 -- ----------------------------------------------------------------------------
 -- Bindings
 
@@ -210,6 +216,8 @@ vectLiteral lit
       return (Lit lit, lexpr)
 
 vectPolyExpr :: CoreExprWithFVs -> VM VExpr
+vectPolyExpr (_, AnnNote note expr)
+  = liftM (vNote note) $ vectPolyExpr expr
 vectPolyExpr expr
   = polyAbstract tvs $ \abstract ->
     do
@@ -435,9 +443,12 @@ packLiftingContext len shape tag fvs res_ty p
       lc_var <- builtin liftingContext
       localV $
         do
-          bnds <- mapM (packFreeVar (Var lc_var) (Var sel_var)) (varSetElems fvs)
+          bnds <- mapM (packFreeVar (Var lc_var) (Var sel_var))
+                . filter isLocalId
+                $ varSetElems fvs
           (vexpr, lexpr) <- p
           return (vexpr, Let (NonRec sel_var sel_expr)
+                         . mkLets (concat bnds)
                          $ Case len lc_var res_ty [(DEFAULT, [], lexpr)])
 
 packFreeVar :: CoreExpr -> CoreExpr -> Var -> VM [CoreBind]