Use packByTag instead of pack in the vectoriser

4 files changed, 31 insertions, 13 deletions

diff --git a/compiler/vectorise/VectBuiltIn.hs b/compiler/vectorise/VectBuiltIn.hs
index 77b4243b00..d417898776 100644
--- a/compiler/vectorise/VectBuiltIn.hs
+++ b/compiler/vectorise/VectBuiltIn.hs
@@ -1,6 +1,6 @@
 module VectBuiltIn (
   Builtins(..), sumTyCon, prodTyCon, prodDataCon,
-  selTy, selReplicate, selPick, selElements,
+  selTy, selReplicate, selPick, selTags, selElements,
   combinePDVar, scalarZip, closureCtrFun,
   initBuiltins, initBuiltinVars, initBuiltinTyCons, initBuiltinDataCons,
   initBuiltinPAs, initBuiltinPRs,
@@ -111,6 +111,7 @@ data Builtins = Builtins {
                 , selTys           :: Array Int Type
                 , selReplicates    :: Array Int CoreExpr
                 , selPicks         :: Array Int CoreExpr
+                , selTagss         :: Array Int CoreExpr
                 , selEls           :: Array (Int, Int) CoreExpr
                 , sumTyCons        :: Array Int TyCon
                 , closureTyCon     :: TyCon
@@ -125,6 +126,7 @@ data Builtins = Builtins {
                 , replicatePDVar   :: Var
                 , emptyPDVar       :: Var
                 , packPDVar        :: Var
+                , packByTagPDVar   :: Var
                 , combinePDVars    :: Array Int Var
                 , scalarClass      :: Class
                 , scalarZips       :: Array Int Var
@@ -149,6 +151,9 @@ selReplicate = indexBuiltin "selReplicate" selReplicates
 selPick :: Int -> Builtins -> CoreExpr
 selPick = indexBuiltin "selPick" selPicks
 
+selTags :: Int -> Builtins -> CoreExpr
+selTags = indexBuiltin "selTags" selTagss
+
 selElements :: Int -> Int -> Builtins -> CoreExpr
 selElements i j = indexBuiltin "selElements" selEls (i,j)
 
@@ -196,6 +201,8 @@ initBuiltins pkg
                              (numbered "replicate" 2 mAX_DPH_SUM)
       sel_picks    <- mapM (externalFun dph_Selector)
                            (numbered "pick" 2 mAX_DPH_SUM)
+      sel_tags     <- mapM (externalFun dph_Selector)
+                           (numbered "tagsSel" 2 mAX_DPH_SUM)
       sel_els      <- mapM mk_elements
                            [(i,j) | i <- [2..mAX_DPH_SUM], j <- [0..i-1]]
       sum_tcs      <- mapM (externalTyCon dph_Repr)
@@ -204,6 +211,7 @@ initBuiltins pkg
       let selTys        = listArray (2, mAX_DPH_SUM) sel_tys
           selReplicates = listArray (2, mAX_DPH_SUM) sel_replicates
           selPicks      = listArray (2, mAX_DPH_SUM) sel_picks
+          selTagss      = listArray (2, mAX_DPH_SUM) sel_tags
           selEls        = array ((2,0), (mAX_DPH_SUM, mAX_DPH_SUM)) sel_els
           sumTyCons     = listArray (2, mAX_DPH_SUM) sum_tcs
 
@@ -218,6 +226,7 @@ initBuiltins pkg
       replicatePDVar   <- externalVar dph_PArray (fsLit "replicatePD")
       emptyPDVar       <- externalVar dph_PArray (fsLit "emptyPD")
       packPDVar        <- externalVar dph_PArray (fsLit "packPD")
+      packByTagPDVar   <- externalVar dph_PArray (fsLit "packByTagPD")
 
       combines <- mapM (externalVar dph_PArray)
                        [mkFastString ("combine" ++ show i ++ "PD")
@@ -253,6 +262,7 @@ initBuiltins pkg
                , selTys           = selTys
                , selReplicates    = selReplicates
                , selPicks         = selPicks
+               , selTagss         = selTagss
                , selEls           = selEls
                , sumTyCons        = sumTyCons
                , closureTyCon     = closureTyCon
@@ -267,6 +277,7 @@ initBuiltins pkg
                , replicatePDVar   = replicatePDVar
                , emptyPDVar       = emptyPDVar
                , packPDVar        = packPDVar
+               , packByTagPDVar   = packByTagPDVar
                , combinePDVars    = combinePDVars
                , scalarClass      = scalarClass
                , scalarZips       = scalarZips
diff --git a/compiler/vectorise/VectMonad.hs b/compiler/vectorise/VectMonad.hs
index a8c84acc65..98701f0b44 100644
--- a/compiler/vectorise/VectMonad.hs
+++ b/compiler/vectorise/VectMonad.hs
@@ -10,7 +10,7 @@ module VectMonad (
   newExportedVar, newLocalVar, newLocalVars, newDummyVar, newTyVar,
   
   Builtins(..), sumTyCon, prodTyCon, prodDataCon,
-  selTy, selReplicate, selPick, selElements,
+  selTy, selReplicate, selPick, selTags, selElements,
   combinePDVar, scalarZip, closureCtrFun,
   builtin, builtins,
 
diff --git a/compiler/vectorise/VectUtils.hs b/compiler/vectorise/VectUtils.hs
index 9ff5b5a820..e5084241d6 100644
--- a/compiler/vectorise/VectUtils.hs
+++ b/compiler/vectorise/VectUtils.hs
@@ -11,7 +11,7 @@ module VectUtils (
   pdataReprTyCon, pdataReprDataCon, mkVScrut,
   prDictOfType, prDFunOfTyCon,
   paDictArgType, paDictOfType, paDFunType,
-  paMethod, wrapPR, replicatePD, emptyPD, packPD,
+  paMethod, wrapPR, replicatePD, emptyPD, packPD, packByTagPD,
   combinePD,
   liftPD,
   zipScalars, scalarClosure,
@@ -269,6 +269,12 @@ packPD :: Type -> CoreExpr -> CoreExpr -> CoreExpr -> VM CoreExpr
 packPD ty xs len sel = liftM (`mkApps` [xs, len, sel])
                              (paMethod packPDVar "packPD" ty)
 
+packByTagPD :: Type -> CoreExpr -> CoreExpr -> CoreExpr -> CoreExpr
+                 -> VM CoreExpr
+packByTagPD ty xs len tags t
+  = liftM (`mkApps` [xs, len, tags, t])
+          (paMethod packByTagPDVar "packByTagPD" ty)
+
 combinePD :: Type -> CoreExpr -> CoreExpr -> [CoreExpr]
           -> VM CoreExpr
 combinePD ty len sel xs
diff --git a/compiler/vectorise/Vectorise.hs b/compiler/vectorise/Vectorise.hs
index 36ee7b7bbb..2bce391a8f 100644
--- a/compiler/vectorise/Vectorise.hs
+++ b/compiler/vectorise/Vectorise.hs
@@ -27,6 +27,7 @@ import OccName
 
 import Literal              ( Literal, mkMachInt )
 import TysWiredIn
+import TysPrim              ( intPrimTy )
 
 import Outputable
 import FastString
@@ -447,9 +448,7 @@ vectAlgCase tycon _ty_args scrut bndr ty alts
               tag  = mkDataConTag vect_dc
               fvs  = freeVarsOf body `delVarSetList` bndrs
 
-          pick <- builtin (selPick arity)
-          let flags_expr = mkApps pick [sel, tag]
-          flags_var <- newLocalVar (fsLit "flags") (exprType flags_expr)
+          sel_tags  <- liftM (`App` sel) (builtin (selTags arity))
           lc        <- builtin liftingContext
           elems     <- builtin (selElements arity ntag)
 
@@ -457,15 +456,17 @@ vectAlgCase tycon _ty_args scrut bndr ty alts
             <- vectBndrsIn bndrs
              . localV
              $ do
-                 binds    <- mapM (pack_var (Var lc) (Var flags_var))
+                 binds    <- mapM (pack_var (Var lc) sel_tags tag)
                            . filter isLocalId
                            $ varSetElems fvs
                  (ve, le) <- vectExpr body
-                 empty    <- emptyPD vty
                  return (ve, Case (elems `App` sel) lc lty
-                               [(DEFAULT, [], Let (NonRec flags_var flags_expr)
-                                              $ mkLets (concat binds) le),
-                                (LitAlt (mkMachInt 0), [], empty)])
+                             [(DEFAULT, [], (mkLets (concat binds) le))])
+                 -- empty    <- emptyPD vty
+                 -- return (ve, Case (elems `App` sel) lc lty
+                 --             [(DEFAULT, [], Let (NonRec flags_var flags_expr)
+                 --                             $ mkLets (concat binds) le),
+                 --               (LitAlt (mkMachInt 0), [], empty)])
           let (vect_bndrs, lift_bndrs) = unzip vbndrs
           return (vect_dc, vect_bndrs, lift_bndrs, vbody)
 
@@ -473,14 +474,14 @@ vectAlgCase tycon _ty_args scrut bndr ty alts
 
     mk_vect_alt vect_dc bndrs body = (DataAlt vect_dc, bndrs, body)
 
-    pack_var len flags v
+    pack_var len tags t v
       = do
           r <- lookupVar v
           case r of
             Local (vv, lv) ->
               do
                 lv'  <- cloneVar lv
-                expr <- packPD (idType vv) (Var lv) len flags
+                expr <- packByTagPD (idType vv) (Var lv) len tags t
                 updLEnv (\env -> env { local_vars = extendVarEnv
                                                 (local_vars env) v (vv, lv') })
                 return [(NonRec lv' expr)]