[commit: ghc] ghc-7.4: Be more careful when deciding which functions are scalar (946c947)

[Ian Lynagh]

Repository : ssh://darcs.haskell.org//srv/darcs/ghc

On branch  : ghc-7.4

http://hackage.haskell.org/trac/ghc/changeset/946c947c1a27d7e0a98055ca02823a9b42bcf2e4

>---------------------------------------------------------------

commit 946c947c1a27d7e0a98055ca02823a9b42bcf2e4
Author: Manuel M T Chakravarty <chak at cse.unsw.edu.au>
Date:   Wed Dec 14 15:37:56 2011 +1100

    Be more careful when deciding which functions are scalar
    
    Although scalar functions can use any scalar data type, their arguments and functions may only involve primitive types at the moment.

>---------------------------------------------------------------

 compiler/vectorise/Vectorise/Builtins.hs |    7 ++++-
 compiler/vectorise/Vectorise/Env.hs      |    3 ++
 compiler/vectorise/Vectorise/Exp.hs      |   33 ++++++++++++++++++++---------
 3 files changed, 31 insertions(+), 12 deletions(-)

diff --git a/compiler/vectorise/Vectorise/Builtins.hs b/compiler/vectorise/Vectorise/Builtins.hs
index d194135..a897ad2 100644
--- a/compiler/vectorise/Vectorise/Builtins.hs
+++ b/compiler/vectorise/Vectorise/Builtins.hs
@@ -1,9 +1,12 @@
 -- Types and functions declared in 'Data.Array.Parallel.Prim' and used by the vectoriser.
 --
--- The @Builtins@ structure holds the name of all the things in 'Data.Array.Parallel.Prim' that appear in
--- code generated by the vectoriser.
+-- The @Builtins@ structure holds the name of all the things in 'Data.Array.Parallel.Prim' that
+-- appear in code generated by the vectoriser.
 
 module Vectorise.Builtins (
+  -- * Restrictions
+  mAX_DPH_SCALAR_ARGS,
+  
   -- * Builtins
   Builtins(..),
   
diff --git a/compiler/vectorise/Vectorise/Env.hs b/compiler/vectorise/Vectorise/Env.hs
index 166262f..cf5bf96 100644
--- a/compiler/vectorise/Vectorise/Env.hs
+++ b/compiler/vectorise/Vectorise/Env.hs
@@ -98,6 +98,9 @@ data GlobalEnv
           -- *without* a right-hand side in the current or an imported module as well as type
           -- constructors that are automatically identified as scalar by the vectoriser (in
           -- 'Vectorise.Type.Env').  Scalar code may only operate on such data.
+          --
+          -- NB: Not all type constructors in that set are members of the 'Scalar' type class
+          --     (which can be trivially marshalled across scalar code boundaries).
         
         , global_novect_vars          :: VarSet
           -- ^Variables that are not vectorised.  (They may be referenced in the right-hand sides
diff --git a/compiler/vectorise/Vectorise/Exp.hs b/compiler/vectorise/Vectorise/Exp.hs
index d695fcb..3970549 100644
--- a/compiler/vectorise/Vectorise/Exp.hs
+++ b/compiler/vectorise/Vectorise/Exp.hs
@@ -32,6 +32,7 @@ import DataCon
 import TyCon
 import TcType
 import Type
+import PrelNames
 import NameSet
 import Var
 import VarEnv
@@ -311,11 +312,11 @@ vectDictExpr (Type ty)
 vectDictExpr (Coercion coe)
   = pprSorry "Vectorise.Exp.vectDictExpr: coercion" (ppr coe)
 
--- |Vectorise an expression of functional type, where all arguments and the result are of scalar
--- type (i.e., 'Int', 'Float', 'Double' etc.) and which does not contain any subcomputations that
--- involve parallel arrays.  Such functionals do not requires the full blown vectorisation
--- transformation; instead, they can be lifted by application of a member of the zipWith family
--- (i.e., 'map', 'zipWith', zipWith3', etc.)
+-- |Vectorise an expression of functional type, where all arguments and the result are of primitive
+-- types (i.e., 'Int', 'Float', 'Double' etc., which have instances of the 'Scalar' type class) and
+-- which does not contain any subcomputations that involve parallel arrays.  Such functionals do not
+-- requires the full blown vectorisation transformation; instead, they can be lifted by application
+-- of a member of the zipWith family (i.e., 'map', 'zipWith', zipWith3', etc.)
 --
 vectScalarFun :: Bool       -- ^ Was the function marked as scalar by the user?
               -> [Var]      -- ^ Functions names in same recursive binding group
@@ -328,15 +329,25 @@ vectScalarFun forceScalar recFns expr
             (arg_tys, res_ty) = splitFunTys (exprType expr)
       ; MASSERT( not $ null arg_tys )
       ; onlyIfV (ptext (sLit "not a scalar function"))
-                (forceScalar                              -- user asserts the functions is scalar
+                (forceScalar                                 -- user asserts the functions is scalar
                  ||
-                 all (is_scalar_ty scalarTyCons) arg_tys  -- check whether the function is scalar
-                  && is_scalar_ty scalarTyCons res_ty
+                 all is_primitive_ty arg_tys                 -- check whether the function is scalar
+                  && is_primitive_ty res_ty
                   && is_scalar scalarVars (is_scalar_ty scalarTyCons) expr
-                  && uses scalarVars expr)
+                  && uses scalarVars expr
+                  && length arg_tys <= mAX_DPH_SCALAR_ARGS)
         $ mkScalarFun arg_tys res_ty expr
       }
   where
+    -- !!!FIXME: We would like to allow scalar functions with arguments and results that can be
+    --           any 'scalarTyCons', but can't at the moment, as those argument and result types
+    --           need to be members of the 'Scalar' class (that in its current form would better
+    --           be called 'Primitive'). *ALSO* the hardcoded list of types is ugly!
+    is_primitive_ty ty
+      | Just (tycon, _) <- splitTyConApp_maybe ty
+      = tyConName tycon `elem` [boolTyConName, intTyConName, word8TyConName, doubleTyConName]
+      | otherwise = False
+
     is_scalar_ty scalarTyCons ty 
       | Just (tycon, _) <- splitTyConApp_maybe ty
       = tyConName tycon `elemNameSet` scalarTyCons
@@ -418,7 +429,9 @@ vectScalarFun forceScalar recFns expr
 
 mkScalarFun :: [Type] -> Type -> CoreExpr -> VM VExpr
 mkScalarFun arg_tys res_ty expr
-  = do { fn_var  <- hoistExpr (fsLit "fn") expr DontInline
+  = do { traceVt "mkScalarFun: " $ ppr expr
+
+       ; fn_var  <- hoistExpr (fsLit "fn") expr DontInline
        ; zipf    <- zipScalars arg_tys res_ty
        ; clo     <- scalarClosure arg_tys res_ty (Var fn_var) (zipf `App` Var fn_var)
        ; clo_var <- hoistExpr (fsLit "clo") clo DontInline