Fix scalar vectorisation of superclasses and recursive dfuns

2 files changed, 35 insertions, 7 deletions

diff --git a/compiler/vectorise/Vectorise.hs b/compiler/vectorise/Vectorise.hs
index cd87868081..88fc947242 100644
--- a/compiler/vectorise/Vectorise.hs
+++ b/compiler/vectorise/Vectorise.hs
@@ -218,15 +218,23 @@ vectTopBind b@(Rec bs)
 -- Add a vectorised binding to an imported top-level variable that has a VECTORISE [SCALAR] pragma
 -- in this module.
 --
+-- RESTIRCTION: Currently, we cannot use the pragma vor mutually recursive definitions.
+--
 vectImpBind :: Id -> VM CoreBind
 vectImpBind var
   = do {   -- Vectorise the right-hand side, create an appropriate top-level binding and add it
            -- to the vectorisation map.  For the non-lifted version, we refer to the original
            -- definition — i.e., 'Var var'.
-       ; (inline, isScalar, expr') <- vectTopRhs [] var (Var var)
-       ; var' <- vectTopBinder var inline expr'
-       ; when isScalar $ 
-           addGlobalScalarVar var
+           -- NB: To support recursive definitions, we tie a lazy knot.
+       ; (var', _, expr') <- fixV $
+           \ ~(_, inline, rhs) ->
+             do { var' <- vectTopBinder var inline rhs
+                ; (inline, isScalar, expr') <- vectTopRhs [] var (Var var)
+
+                ; when isScalar $ 
+                    addGlobalScalarVar var
+                ; return (var', inline, expr')
+                }
 
            -- We add any newly created hoisted top-level bindings.
        ; hs <- takeHoisted
diff --git a/compiler/vectorise/Vectorise/Exp.hs b/compiler/vectorise/Vectorise/Exp.hs
index 3970549f75..778a3a5d19 100644
--- a/compiler/vectorise/Vectorise/Exp.hs
+++ b/compiler/vectorise/Vectorise/Exp.hs
@@ -318,6 +318,10 @@ vectDictExpr (Coercion coe)
 -- 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.)
 --
+-- Dictionary functions are also scalar functions (as dictionaries themselves are not vectorised,
+-- instead they become dictionaries of vectorised methods).  We treat them differently, though see
+-- "Note [Scalar dfuns]" in 'Vectorise'.
+--
 vectScalarFun :: Bool       -- ^ Was the function marked as scalar by the user?
               -> [Var]      -- ^ Functions names in same recursive binding group
               -> CoreExpr   -- ^ Expression to be vectorised
@@ -344,14 +348,20 @@ vectScalarFun forceScalar recFns expr
     --           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
+      | isPredTy ty               -- dictionaries never get into the environment
+      = True
       | Just (tycon, _) <- splitTyConApp_maybe ty
       = tyConName tycon `elem` [boolTyConName, intTyConName, word8TyConName, doubleTyConName]
-      | otherwise = False
+      | otherwise 
+      = False
 
     is_scalar_ty scalarTyCons ty 
+      | isPredTy ty               -- dictionaries never get into the environment
+      = True
       | Just (tycon, _) <- splitTyConApp_maybe ty
       = tyConName tycon `elemNameSet` scalarTyCons
-      | otherwise = False
+      | otherwise 
+      = False
 
     -- Checks whether an expression contain a non-scalar subexpression. 
     --
@@ -427,9 +437,17 @@ vectScalarFun forceScalar recFns expr
 
     uses_alt funs (_, _bs, e) = uses funs e 
 
+-- Generate code for a scalar function by generating a scalar closure.  If the function is a
+-- dictionary function, vectorise it as dictionary code.
+-- 
 mkScalarFun :: [Type] -> Type -> CoreExpr -> VM VExpr
 mkScalarFun arg_tys res_ty expr
-  = do { traceVt "mkScalarFun: " $ ppr expr
+  | isPredTy res_ty
+  = do { vExpr <- vectDictExpr expr
+       ; return (vExpr, unused)
+       }
+  | otherwise
+  = do { traceVt "mkScalarFun: " $ ppr expr $$ ptext (sLit "  ::") <+> ppr (mkFunTys arg_tys res_ty)
 
        ; fn_var  <- hoistExpr (fsLit "fn") expr DontInline
        ; zipf    <- zipScalars arg_tys res_ty
@@ -438,6 +456,8 @@ mkScalarFun arg_tys res_ty expr
        ; lclo    <- liftPD (Var clo_var)
        ; return (Var clo_var, lclo)
        }
+  where
+    unused = error "Vectorise.Exp.mkScalarFun: we don't lift dictionary expressions"
 
 -- |Vectorise a dictionary function that has a 'VECTORISE SCALAR instance' pragma.
 --