Fix bug in vectorisation of case expressions

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

diff --git a/compiler/vectorise/Vectorise.hs b/compiler/vectorise/Vectorise.hs
index 9f2e2b7..562e46d 100644 (file)
--- a/compiler/vectorise/Vectorise.hs
+++ b/compiler/vectorise/Vectorise.hs
@@ -42,7 +42,7 @@ import Module               ( Module )
 import DsMonad hiding (mapAndUnzipM)
 import DsUtils              ( mkCoreTup, mkCoreTupTy )
 
-import Literal              ( Literal )
+import Literal              ( Literal, mkMachInt )
 import PrelNames
 import TysWiredIn
 import TysPrim              ( intPrimTy )
@@ -70,10 +70,10 @@ vectModule :: ModGuts -> VM ModGuts
 vectModule guts
   = do
       (types', fam_insts, tc_binds) <- vectTypeEnv (mg_types guts)
-      
+
       let fam_inst_env' = extendFamInstEnvList (mg_fam_inst_env guts) fam_insts
       updGEnv (setFamInstEnv fam_inst_env')
-     
+
       -- dicts   <- mapM buildPADict pa_insts
       -- workers <- mapM vectDataConWorkers pa_insts
       binds'  <- mapM vectTopBind (mg_binds guts)
@@ -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
@@ -111,7 +113,7 @@ vectTopBinder var
       var' <- cloneId mkVectOcc var vty
       defGlobalVar var var'
       return var'
-    
+
 vectTopRhs :: Var -> CoreExpr -> VM CoreExpr
 vectTopRhs var expr
   = do
@@ -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,14 +216,16 @@ 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
       mono' <- vectExpr mono
       return $ mapVect abstract mono'
   where
-    (tvs, mono) = collectAnnTypeBinders expr  
-                
+    (tvs, mono) = collectAnnTypeBinders expr
+
 vectExpr :: CoreExprWithFVs -> VM VExpr
 vectExpr (_, AnnType ty)
   = liftM vType (vectType ty)
@@ -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
@@ -281,6 +300,8 @@ vectExpr e@(fvs, AnnLam bndr _)
   where
     (bs,body) = collectAnnValBinders e
 
+vectExpr e = pprPanic "vectExpr" (ppr $ deAnnotate e)
+
 vectLam :: VarSet -> [Var] -> CoreExprWithFVs -> VM VExpr
 vectLam fvs bs body
   = do
@@ -299,7 +320,7 @@ vectLam fvs bs body
             (vbndrs, vbody) <- vectBndrsIn (vs ++ bs)
                                            (vectExpr body)
             return $ vLams lc vbndrs vbody
-  
+
 vectTyAppExpr :: CoreExprWithFVs -> [Type] -> VM VExpr
 vectTyAppExpr (_, AnnVar v) tys = vectPolyVar v tys
 vectTyAppExpr e tys = pprPanic "vectTyAppExpr" (ppr $ deAnnotate e)
@@ -318,7 +339,7 @@ type CoreAltWithFVs = AnnAlt Id VarSet
 -- When lifting, we have to do it this way because v must have the type
 -- [:V(T):] but the scrutinee must be cast to the representation type. We also
 -- have to handle the case where v is a wild var correctly.
---   
+--
 
 -- FIXME: this is too lazy
 vectAlgCase tycon ty_args scrut bndr ty [(DEFAULT, [], body)]
@@ -329,6 +350,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)
@@ -360,7 +389,7 @@ vectAlgCase tycon ty_args scrut bndr ty alts
       shape_bndrs <- arrShapeVars repr
       (len, sel, indices) <- arrSelector repr (map Var shape_bndrs)
 
-      (vbndr, valts) <- vect_scrut_bndr $ mapM (proc_alt sel lty) alts'
+      (vbndr, valts) <- vect_scrut_bndr $ mapM (proc_alt sel vty lty) alts'
       let (vect_dcs, vect_bndrss, lift_bndrss, vbodies) = unzip4 valts
 
       vexpr <- vectExpr scrut
@@ -391,14 +420,14 @@ vectAlgCase tycon ty_args scrut bndr ty alts
     cmp DEFAULT       _             = LT
     cmp _             DEFAULT       = GT
 
-    proc_alt sel lty (DataAlt dc, bndrs, body)
+    proc_alt sel vty lty (DataAlt dc, bndrs, body)
       = do
           vect_dc <- maybeV (lookupDataCon dc)
           let tag = mkDataConTag vect_dc
               fvs = freeVarsOf body `delVarSetList` bndrs
           (vect_bndrs, lift_bndrs, vbody)
             <- vect_alt_bndrs bndrs
-             $ \len -> packLiftingContext len sel tag fvs lty
+             $ \len -> packLiftingContext len sel tag fvs vty lty
              $ vectExpr body
 
           return (vect_dc, vect_bndrs, lift_bndrs, vbody)
@@ -426,8 +455,9 @@ vectAlgCase tycon ty_args scrut bndr ty alts
 
     mk_vect_alt vect_dc bndrs body = (DataAlt vect_dc, bndrs, body)
 
-packLiftingContext :: CoreExpr -> CoreExpr -> CoreExpr -> VarSet -> Type -> VM VExpr -> VM VExpr
-packLiftingContext len shape tag fvs res_ty p
+packLiftingContext :: CoreExpr -> CoreExpr -> CoreExpr -> VarSet
+                   -> Type -> Type -> VM VExpr -> VM VExpr
+packLiftingContext len shape tag fvs vty lty p
   = do
       select <- builtin selectPAIntPrimVar
       let sel_expr = mkApps (Var select) [shape, tag]
@@ -439,8 +469,11 @@ packLiftingContext len shape tag fvs res_ty p
                 . filter isLocalId
                 $ varSetElems fvs
           (vexpr, lexpr) <- p
+          empty <- emptyPA vty
           return (vexpr, Let (NonRec sel_var sel_expr)
-                         $ Case len lc_var res_ty [(DEFAULT, [], lexpr)])
+                         $ Case len lc_var lty
+                             [(DEFAULT, [], mkLets (concat bnds) lexpr),
+                              (LitAlt (mkMachInt 0), [], empty)])
 
 packFreeVar :: CoreExpr -> CoreExpr -> Var -> VM [CoreBind]
 packFreeVar len sel v