First cut at scalar vectorisation of class instances

1 files changed, 62 insertions, 24 deletions

diff --git a/compiler/vectorise/Vectorise.hs b/compiler/vectorise/Vectorise.hs
index 3ba247dfbe..7d2415caf2 100644
--- a/compiler/vectorise/Vectorise.hs
+++ b/compiler/vectorise/Vectorise.hs
@@ -81,25 +81,15 @@ vectModule guts@(ModGuts { mg_tcs        = tycons
           -- array types.
       ; (new_tycons, new_fam_insts, tc_binds) <- vectTypeEnv tycons ty_vect_decls cls_vect_decls
 
-{- TODO:
-
-instance Num Int where
-  (+) = primAdd
-{-# VECTORISE SCALAR instance Num Int #-}
-
-==> $dNumInt :: Num Int; $dNumInt = Num primAdd
-=>> $v$dNumInt :: $vNum Int
-    $v$dNumInt = $vNum (closure1 (scalar_zipWith primAdd) (scalar_zipWith primAdd))
-  $dNumInt -v> $v$dNumInt
--}
-
           -- Family instance environment for /all/ home-package modules including those instances
           -- generated by 'vectTypeEnv'.
       ; (_, fam_inst_env) <- readGEnv global_fam_inst_env
 
           -- Vectorise all the top level bindings and VECTORISE declarations on imported identifiers
+      ; let impBinds = [imp_id | Vect          imp_id _ <- vect_decls, isGlobalId imp_id] ++
+                       [imp_id | VectInst True imp_id   <- vect_decls, isGlobalId imp_id]
       ; binds_top <- mapM vectTopBind binds
-      ; binds_imp <- mapM vectImpBind [imp_id | Vect imp_id _ <- vect_decls, isGlobalId imp_id]
+      ; binds_imp <- mapM vectImpBind impBinds
 
       ; return $ guts { mg_tcs          = tycons ++ new_tycons
                         -- we produce no new classes or instances, only new class type constructors
@@ -283,21 +273,63 @@ vectTopBinder var inline expr
     unfolding = case inline of
                   Inline arity -> mkInlineUnfolding (Just arity) expr
                   DontInline   -> noUnfolding
+{-
+!!!TODO: dfuns and unfoldings:
+           -- Do not inline the dfun; instead give it a magic DFunFunfolding
+           -- See Note [ClassOp/DFun selection]
+           -- See also note [Single-method classes]
+        dfun_id_w_fun
+           | isNewTyCon class_tc
+           = dfun_id `setInlinePragma` alwaysInlinePragma { inl_sat = Just 0 }
+           | otherwise
+           = dfun_id `setIdUnfolding`  mkDFunUnfolding dfun_ty dfun_args
+                     `setInlinePragma` dfunInlinePragma
+ -}
 
 -- | Vectorise the RHS of a top-level binding, in an empty local environment.
 --
--- We need to distinguish three cases:
+-- We need to distinguish four cases:
 --
 -- (1) We have a (non-scalar) vectorisation declaration for the variable (which explicitly provides
 --     vectorised code implemented by the user)
 --     => no automatic vectorisation & instead use the user-supplied code
 -- 
--- (2) We have a scalar vectorisation declaration for the variable
+-- (2) We have a scalar vectorisation declaration for a variable that is no dfun
 --     => generate vectorised code that uses a scalar 'map'/'zipWith' to lift the computation
 -- 
--- (3) There is no vectorisation declaration for the variable
+-- (3) We have a scalar vectorisation declaration for a variable that *is* a dfun
+--     => generate vectorised code according to the the "Note [Scalar dfuns]" below
+-- 
+-- (4) There is no vectorisation declaration for the variable
 --     => perform automatic vectorisation of the RHS
 --
+-- Note [Scalar dfuns]
+-- ~~~~~~~~~~~~~~~~~~~
+--
+-- Here is the translation scheme for scalar dfuns — assume the instance declaration:
+--
+--   instance Num Int where
+--     (+) = primAdd
+--   {-# VECTORISE SCALAR instance Num Int #-}
+--
+-- It desugars to
+--
+--   $dNumInt :: Num Int
+--   $dNumInt = D:Num primAdd
+--
+-- We vectorise it to
+--
+--   $v$dNumInt :: V:Num Int
+--   $v$dNumInt = D:V:Num (closure2 ((+) $dNumInt) (scalar_zipWith ((+) $dNumInt))))
+--
+-- while adding the following entry to the vectorisation map: '$dNumInt' --> '$v$dNumInt'.
+--
+-- See "Note [Vectorising classes]" in 'Vectorise.Type.Env' for the definition of 'V:Num'.
+--
+-- NB: The outlined vectorisation scheme does not require the right-hand side of the original dfun.
+--     In fact, we definitely want to refer to the dfn variable instead of the right-hand side to 
+--     ensure that the dictionary selection rules fire.
+--
 vectTopRhs :: [Var]           -- ^ Names of all functions in the rec block
            -> Var             -- ^ Name of the binding.
            -> CoreExpr        -- ^ Body of the binding.
@@ -308,19 +340,24 @@ vectTopRhs recFs var expr
   = closedV
   $ do { globalScalar <- isGlobalScalar var
        ; vectDecl     <- lookupVectDecl var
+       ; let isDFun = isDFunId var
 
-       ; traceVt ("vectTopRhs of " ++ show var ++ info globalScalar vectDecl) $ ppr expr
+       ; traceVt ("vectTopRhs of " ++ show var ++ info globalScalar isDFun vectDecl) $ ppr expr
 
-       ; rhs globalScalar vectDecl
+       ; rhs globalScalar isDFun vectDecl
        }
   where
-    rhs _globalScalar (Just (_, expr'))               -- Case (1)
+    rhs _globalScalar _isDFun (Just (_, expr'))               -- Case (1)
       = return (inlineMe, False, expr')
-    rhs True          Nothing                         -- Case (2)
+    rhs True          False   Nothing                         -- Case (2)
       = do { expr' <- vectScalarFun True recFs expr
            ; return (inlineMe, True, vectorised expr')
            }
-    rhs False         Nothing                         -- Case (3)
+    rhs True          True    Nothing                         -- Case (3)
+      = do { expr' <- vectScalarDFun var recFs
+           ; return (DontInline, True, expr')
+           }
+    rhs False         _isDFun Nothing                         -- Case (4)
       = do { let fvs = freeVars expr
            ; (inline, isScalar, vexpr) 
                <- inBind var $
@@ -328,9 +365,10 @@ vectTopRhs recFs var expr
            ; return (inline, isScalar, vectorised vexpr)
            }
     
-    info True  _                          = " [VECTORISE SCALAR]"
-    info False vectDecl | isJust vectDecl = " [VECTORISE]"
-                        | otherwise       = " (no pragma)"
+    info True  False _                          = " [VECTORISE SCALAR]"
+    info True  True  _                          = " [VECTORISE SCALAR instance]"
+    info False _     vectDecl | isJust vectDecl = " [VECTORISE]"
+                              | otherwise       = " (no pragma)"
 
 -- |Project out the vectorised version of a binding from some closure,
 -- or return the original body if that doesn't work or the binding is scalar.