Remove '-favoid-vect' and add '-fvectorisation-avoidance'

* By default '-fvectorisation-avoidance' is enabled at all optimisation levels (but it only matters in combination with '-fvectorise').
* The new vectoriser always uses vectorisation avoidance, but with '-fno-vectorisation-avoidance' it restricts it to simple scalar applications (and dictionary computations)

4 files changed, 60 insertions, 17 deletions

diff --git a/compiler/vectorise/Vectorise/Env.hs b/compiler/vectorise/Vectorise/Env.hs
index efda22a1c5..3358ceafab 100644
--- a/compiler/vectorise/Vectorise/Env.hs
+++ b/compiler/vectorise/Vectorise/Env.hs
@@ -79,7 +79,10 @@ emptyLocalEnv = LocalEnv
 --
 data GlobalEnv 
         = GlobalEnv
-        { global_vars                 :: VarEnv Var
+        { global_vect_avoid           :: Bool
+          -- ^'True' implies to avoid vectorisation as far as possible.
+
+        , global_vars                 :: VarEnv Var
           -- ^Mapping from global variables to their vectorised versions — aka the /vectorisation
           -- map/.
 
@@ -136,10 +139,11 @@ data GlobalEnv
 -- to the global table, so that we can query scalarness during vectorisation, and especially, when
 -- vectorising the scalar entities' definitions themselves.
 --
-initGlobalEnv :: VectInfo -> [CoreVect] -> (InstEnv, InstEnv) -> FamInstEnvs -> GlobalEnv
-initGlobalEnv info vectDecls instEnvs famInstEnvs
+initGlobalEnv :: Bool -> VectInfo -> [CoreVect] -> (InstEnv, InstEnv) -> FamInstEnvs -> GlobalEnv
+initGlobalEnv vectAvoid info vectDecls instEnvs famInstEnvs
   = GlobalEnv 
-  { global_vars                 = mapVarEnv snd $ vectInfoVar info
+  { global_vect_avoid           = vectAvoid
+  , global_vars                 = mapVarEnv snd $ vectInfoVar info
   , global_vect_decls           = mkVarEnv vects
   , global_parallel_vars        = vectInfoParallelVars info
   , global_parallel_tycons      = vectInfoParallelTyCons info
diff --git a/compiler/vectorise/Vectorise/Exp.hs b/compiler/vectorise/Vectorise/Exp.hs
index d4eee26553..2fdd223975 100644
--- a/compiler/vectorise/Vectorise/Exp.hs
+++ b/compiler/vectorise/Vectorise/Exp.hs
@@ -157,7 +157,12 @@ vectAnnPolyExpr loop_breaker expr
 -- * All free variables and the result type must be /simple/ types.
 -- * The expression is sufficiently complex (to warrant special treatment).  For now, that is
 --   every expression that is not constant and contains at least one operation.
---  
+--
+--
+-- The user has an option to choose between aggressive and minimal vectorisation avoidance. With
+-- minimal vectorisation avoidance, we only encapsulate individual scalar operations. With
+-- aggressive vectorisation avoidance, we encapsulate subexpression that are as big as possible.
+--
 encapsulateScalars :: CoreExprWithVectInfo -> VM CoreExprWithVectInfo
 encapsulateScalars ce@(_, AnnType _ty)
   = return ce
@@ -175,12 +180,13 @@ encapsulateScalars ((fvs, vi), AnnTick tck expr)
     }
 encapsulateScalars ce@((fvs, vi), AnnLam bndr expr) 
   = do 
-    { varsS  <- allScalarVarTypeSet fvs 
+    { vectAvoid <- isVectAvoidanceAggressive
+    ; varsS     <- allScalarVarTypeSet fvs 
         -- NB: diverts from the paper: we need to check the scalarness of bound variables as well,
         --     as 'vectScalarFun' will handle them just the same as those introduced for the 'fvs'
         --     by encapsulation.
-    ; bndrsS <- allScalarVarType bndrs
-    ; case (vi, varsS && bndrsS) of
+    ; bndrsS    <- allScalarVarType bndrs
+    ; case (vi, vectAvoid && varsS && bndrsS) of
         (VISimple, True) -> liftSimpleAndCase ce
         _                -> do 
                             { encExpr <- encapsulateScalars expr
@@ -191,8 +197,9 @@ encapsulateScalars ce@((fvs, vi), AnnLam bndr expr)
     (bndrs, _) = collectAnnBndrs ce
 encapsulateScalars ce@((fvs, vi), AnnApp ce1 ce2)
   = do
-    { varsS <- allScalarVarTypeSet fvs
-    ; case (vi, varsS) of
+    { vectAvoid <- isVectAvoidanceAggressive
+    ; varsS     <- allScalarVarTypeSet fvs
+    ; case (vi, (vectAvoid || isSimpleApplication ce) && varsS) of
         (VISimple, True) -> liftSimpleAndCase ce
         _                -> do 
                             { encCe1 <- encapsulateScalars ce1
@@ -200,10 +207,26 @@ encapsulateScalars ce@((fvs, vi), AnnApp ce1 ce2)
                             ; return ((fvs, vi), AnnApp encCe1 encCe2)
                             }
     }
+  where
+    isSimpleApplication :: CoreExprWithVectInfo -> Bool
+    isSimpleApplication (_, AnnTick _ ce)                 = isSimpleApplication ce
+    isSimpleApplication (_, AnnCast ce _)                 = isSimpleApplication ce
+    isSimpleApplication ce                  | isSimple ce = True
+    isSimpleApplication (_, AnnApp ce1 ce2)               = isSimple ce1 && isSimpleApplication ce2
+    isSimpleApplication _                                 = False
+    --
+    isSimple :: CoreExprWithVectInfo -> Bool
+    isSimple (_, AnnType {})   = True
+    isSimple (_, AnnVar  {})   = True
+    isSimple (_, AnnLit  {})   = True
+    isSimple (_, AnnTick _ ce) = isSimple ce
+    isSimple (_, AnnCast ce _) = isSimple ce
+    isSimple _                 = False
 encapsulateScalars ce@((fvs, vi), AnnCase scrut bndr ty alts) 
   = do 
-    { varsS <- allScalarVarTypeSet fvs 
-    ; case (vi, varsS) of
+    { vectAvoid <- isVectAvoidanceAggressive
+    ; varsS     <- allScalarVarTypeSet fvs 
+    ; case (vi, vectAvoid && varsS) of
         (VISimple, True) -> liftSimpleAndCase ce
         _                -> do 
                             { encScrut <- encapsulateScalars scrut
@@ -215,8 +238,9 @@ encapsulateScalars ce@((fvs, vi), AnnCase scrut bndr ty alts)
     encAlt (con, bndrs, expr) = (con, bndrs,) <$> encapsulateScalars expr
 encapsulateScalars ce@((fvs, vi), AnnLet (AnnNonRec bndr expr1) expr2) 
   = do 
-    { varsS <- allScalarVarTypeSet fvs 
-    ; case (vi, varsS) of
+    { vectAvoid <- isVectAvoidanceAggressive
+    ; varsS     <- allScalarVarTypeSet fvs 
+    ; case (vi, vectAvoid && varsS) of
         (VISimple, True) -> liftSimpleAndCase ce
         _                -> do 
                             { encExpr1 <- encapsulateScalars expr1
@@ -226,8 +250,9 @@ encapsulateScalars ce@((fvs, vi), AnnLet (AnnNonRec bndr expr1) expr2)
     }
 encapsulateScalars ce@((fvs, vi), AnnLet (AnnRec binds) expr) 
   = do 
-    { varsS <- allScalarVarTypeSet fvs 
-    ; case (vi, varsS) of 
+    { vectAvoid <- isVectAvoidanceAggressive
+    ; varsS     <- allScalarVarTypeSet fvs 
+    ; case (vi, vectAvoid && varsS) of 
         (VISimple, True) -> liftSimpleAndCase ce
         _                -> do 
                             { encBinds <- mapM encBind binds
diff --git a/compiler/vectorise/Vectorise/Monad.hs b/compiler/vectorise/Vectorise/Monad.hs
index 6b5e9cc354..e9078ab850 100644
--- a/compiler/vectorise/Vectorise/Monad.hs
+++ b/compiler/vectorise/Vectorise/Monad.hs
@@ -92,7 +92,8 @@ initV hsc_env guts info thing_inside
            ; let genv = extendImportedVarsEnv builtin_vars
                         . setPAFunsEnv        builtin_pas
                         . setPRFunsEnv        builtin_prs
-                        $ initGlobalEnv info (mg_vect_decls guts) instEnvs famInstEnvs
+                        $ initGlobalEnv (dopt Opt_VectorisationAvoidance dflags) 
+                                        info (mg_vect_decls guts) instEnvs famInstEnvs
  
                -- perform vectorisation
            ; r <- runVM thing_inside builtins genv emptyLocalEnv
diff --git a/compiler/vectorise/Vectorise/Monad/Global.hs b/compiler/vectorise/Vectorise/Monad/Global.hs
index 0fe460ad73..143330554f 100644
--- a/compiler/vectorise/Vectorise/Monad/Global.hs
+++ b/compiler/vectorise/Vectorise/Monad/Global.hs
@@ -5,6 +5,9 @@ module Vectorise.Monad.Global (
   setGEnv,
   updGEnv,
   
+  -- * Configuration
+  isVectAvoidanceAggressive,
+  
   -- * Vars
   defGlobalVar, undefGlobalVar,
   
@@ -66,6 +69,16 @@ updGEnv :: (GlobalEnv -> GlobalEnv) -> VM ()
 updGEnv f = VM $ \_ genv lenv -> return (Yes (f genv) lenv ())
 
 
+-- Configuration --------------------------------------------------------------
+
+-- |Should we avoid as much vectorisation as possible?
+--
+-- Set by '-f[no]-vectorisation-avoidance'
+--
+isVectAvoidanceAggressive :: VM Bool
+isVectAvoidanceAggressive = readGEnv global_vect_avoid
+
+
 -- Vars -----------------------------------------------------------------------
 
 -- |Add a mapping between a global var and its vectorised version to the state.