Treat some standard data cons specially during vectorisation

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

diff --git a/compiler/vectorise/Vectorise.hs b/compiler/vectorise/Vectorise.hs
index 96ed01e..80c896c 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
@@ -235,6 +243,17 @@ vectExpr e@(_, AnnApp _ arg)
   where
     (fn, tys) = collectAnnTypeArgs e
 
+vectExpr (_, AnnApp (_, AnnVar v) (_, AnnLit lit))
+  | Just con <- isDataConId_maybe v
+  , is_special_con con
+  = do
+      let vexpr = App (Var v) (Lit lit)
+      lexpr <- liftPA vexpr
+      return (vexpr, lexpr)
+  where
+    is_special_con con = con `elem` [intDataCon, floatDataCon, doubleDataCon]
+                
+
 vectExpr (_, AnnApp fn arg)
   = do
       arg_ty' <- vectType arg_ty
@@ -329,6 +348,14 @@ vectAlgCase tycon ty_args scrut bndr ty [(DEFAULT, [], body)]
       (vbndr, vbody) <- vectBndrIn bndr (vectExpr body)
       return $ vCaseDEFAULT vscrut vbndr vty lty vbody
 
+vectAlgCase tycon ty_args scrut bndr ty [(DataAlt dc, [], body)]
+  = do
+      vscrut <- vectExpr scrut
+      vty    <- vectType ty
+      lty    <- mkPArrayType vty
+      (vbndr, vbody) <- vectBndrIn bndr (vectExpr body)
+      return $ vCaseDEFAULT vscrut vbndr vty lty vbody
+
 vectAlgCase tycon ty_args scrut bndr ty [(DataAlt dc, bndrs, body)]
   = do
       vect_tc <- maybeV (lookupTyCon tycon)
@@ -435,9 +462,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]