Better case-of-case transformation

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

diff --git a/compiler/vectorise/Vectorise.hs b/compiler/vectorise/Vectorise.hs
index 34ca5ab..e3e9646 100644 (file)
--- a/compiler/vectorise/Vectorise.hs
+++ b/compiler/vectorise/Vectorise.hs
@@ -3,8 +3,10 @@
 module Vectorise( vectorise )
 where
 
-import VectUtils
-import VectType
+import Vectorise.Type.Env
+import Vectorise.Type.Type
+import Vectorise.Convert
+import Vectorise.Utils.Hoisting
 import Vectorise.Exp
 import Vectorise.Vect
 import Vectorise.Env
@@ -13,16 +15,17 @@ import Vectorise.Monad
 import HscTypes hiding      ( MonadThings(..) )
 import Module               ( PackageId )
 import CoreSyn
-import CoreUnfold           ( mkInlineRule )
+import CoreUnfold           ( mkInlineUnfolding )
 import CoreFVs
 import CoreMonad            ( CoreM, getHscEnv )
-import FamInstEnv           ( extendFamInstEnvList )
 import Var
 import Id
 import OccName
 import BasicTypes           ( isLoopBreaker )
 import Outputable
 import Util                 ( zipLazy )
+import MonadUtils
+
 import Control.Monad
 
 debug          = False
@@ -58,9 +61,7 @@ vectModule guts
       -- TODO: What new binds do we get back here?
       (types', fam_insts, tc_binds) <- vectTypeEnv (mg_types guts)
 
-      -- TODO: What is this?
-      let fam_inst_env' = extendFamInstEnvList (mg_fam_inst_env guts) fam_insts
-      updGEnv (setFamInstEnv fam_inst_env')
+      (_, fam_inst_env) <- readGEnv global_fam_inst_env
 
       -- dicts   <- mapM buildPADict pa_insts
       -- workers <- mapM vectDataConWorkers pa_insts
@@ -70,7 +71,7 @@ vectModule guts
 
       return $ guts { mg_types        = types'
                     , mg_binds        = Rec tc_binds : binds'
-                    , mg_fam_inst_env = fam_inst_env'
+                    , mg_fam_inst_env = fam_inst_env
                     , mg_fam_insts    = mg_fam_insts guts ++ fam_insts
                     }
 
@@ -114,7 +115,7 @@ vectModule guts
 vectTopBind :: CoreBind -> VM CoreBind
 vectTopBind b@(NonRec var expr)
  = do
-      (inline, expr')  <- vectTopRhs var expr
+      (inline, _, expr')       <- vectTopRhs [] var expr
       var'             <- vectTopBinder var inline expr'
 
       -- Vectorising the body may create other top-level bindings.
@@ -134,11 +135,16 @@ vectTopBind b@(Rec bs)
        <- fixV $ \ ~(_, inlines, rhss) ->
             do vars' <- sequence [vectTopBinder var inline rhs
                                       | (var, ~(inline, rhs)) <- zipLazy vars (zip inlines rhss)]
-               (inlines', exprs') 
-                     <- mapAndUnzipM (uncurry vectTopRhs) bs
-
-               return (vars', inlines', exprs')
-
+               (inlines', areScalars', exprs') 
+                     <- mapAndUnzip3M (uncurry $ vectTopRhs vars) bs
+               if  (and areScalars') || (length bs <= 1)
+                  then do
+                    return (vars', inlines', exprs')
+                  else do
+                    _ <- mapM deleteGlobalScalar vars
+                    (inlines'', _, exprs'')  <- mapAndUnzip3M (uncurry $ vectTopRhs []) bs
+                    return (vars', inlines'', exprs'')
+                      
       hs     <- takeHoisted
       cexprs <- sequence $ zipWith3 tryConvert vars vars' exprs
       return . Rec $ zip vars cexprs ++ zip vars' exprs' ++ hs
@@ -146,7 +152,9 @@ vectTopBind b@(Rec bs)
     return b
   where
     (vars, exprs) = unzip bs
-
+    mapAndUnzip3M f xs = do
+       ys <- mapM f xs
+       return $ unzip3 ys
 
 -- | Make the vectorised version of this top level binder, and add the mapping
 --   between it and the original to the state. For some binder @foo@ the vectorised
@@ -166,7 +174,7 @@ vectTopBinder var inline expr
       vty  <- vectType (idType var)
 
       -- Make the vectorised version of binding's name, and set the unfolding used for inlining.
-      var' <- liftM (`setIdUnfolding` unfolding) 
+      var' <- liftM (`setIdUnfoldingLazily` unfolding) 
            $  cloneId mkVectOcc var vty
 
       -- Add the mapping between the plain and vectorised name to the state.
@@ -175,23 +183,26 @@ vectTopBinder var inline expr
       return var'
   where
     unfolding = case inline of
-                  Inline arity -> mkInlineRule expr (Just arity)
+                  Inline arity -> mkInlineUnfolding (Just arity) expr
                   DontInline   -> noUnfolding
 
 
 -- | Vectorise the RHS of a top-level binding, in an empty local environment.
 vectTopRhs 
-       :: Var          -- ^ Name of the binding.
+       :: [Var]    -- ^ Names of all functions in the rec block
+       -> Var          -- ^ Name of the binding.
        -> CoreExpr     -- ^ Body of the binding.
-       -> VM (Inline, CoreExpr)
+       -> VM (Inline, Bool, CoreExpr)
 
-vectTopRhs var expr
+vectTopRhs recFs var expr
  = dtrace (vcat [text "vectTopRhs", ppr expr])
  $ closedV
- $ do (inline, vexpr) <- inBind var
-                      $ vectPolyExpr (isLoopBreaker $ idOccInfo var)
-                                      (freeVars expr)
-      return (inline, vectorised vexpr)
+ $ do (inline, isScalar, vexpr) <- 
+           inBind var $ vectPolyExpr  (isLoopBreaker $ idOccInfo var) recFs (freeVars expr)
+      if isScalar 
+         then addGlobalScalar var
+         else deleteGlobalScalar var
+      return (inline, isScalar, vectorised vexpr)
 
 
 -- | Project out the vectorised version of a binding from some closure,