Added a VECTORISE pragma

- Added a pragma {-# VECTORISE var = exp #-} that prevents
  the vectoriser from vectorising the definition of 'var'.
  Instead it uses the binding '$v_var = exp' to vectorise
  'var'.  The vectoriser checks that the Core type of 'exp'
  matches the vectorised Core type of 'var'.  (It would be
  quite complicated to perform that check in the type checker
  as the vectorisation of a type needs the state of the VM
  monad.)
- Added parts of a related VECTORISE SCALAR pragma
- Documented -ddump-vect
- Added -ddump-vt-trace
- Some clean up

1 files changed, 137 insertions, 88 deletions

diff --git a/compiler/vectorise/Vectorise.hs b/compiler/vectorise/Vectorise.hs
index e3e9646a19..72cca6e1c6 100644
--- a/compiler/vectorise/Vectorise.hs
+++ b/compiler/vectorise/Vectorise.hs
@@ -1,6 +1,6 @@
 {-# OPTIONS -fno-warn-missing-signatures #-}
 
-module Vectorise( vectorise )
+module Vectorise ( vectorise )
 where
 
 import Vectorise.Type.Env
@@ -13,14 +13,16 @@ import Vectorise.Env
 import Vectorise.Monad
 
 import HscTypes hiding      ( MonadThings(..) )
-import Module               ( PackageId )
-import CoreSyn
 import CoreUnfold           ( mkInlineUnfolding )
 import CoreFVs
+import PprCore
+import CoreSyn
 import CoreMonad            ( CoreM, getHscEnv )
+import Type
 import Var
 import Id
 import OccName
+import DynFlags
 import BasicTypes           ( isLoopBreaker )
 import Outputable
 import Util                 ( zipLazy )
@@ -28,53 +30,58 @@ import MonadUtils
 
 import Control.Monad
 
-debug		= False
-dtrace s x	= if debug then pprTrace "Vectorise" s x else x
 
 -- | Vectorise a single module.
---   Takes the package containing the DPH backend we're using. Eg either dph-par or dph-seq.
-vectorise :: PackageId -> ModGuts -> CoreM ModGuts
-vectorise backend guts 
- = do hsc_env <- getHscEnv
-      liftIO $ vectoriseIO backend hsc_env guts
-
-
--- | Vectorise a single monad, given its HscEnv (code gen environment).
-vectoriseIO :: PackageId -> HscEnv -> ModGuts -> IO ModGuts
-vectoriseIO backend hsc_env guts
- = do -- Get information about currently loaded external packages.
-      eps <- hscEPS hsc_env
+--
+vectorise :: ModGuts -> CoreM ModGuts
+vectorise guts
+ = do { hsc_env <- getHscEnv
+      ; liftIO $ vectoriseIO hsc_env guts
+      }
 
-      -- Combine vectorisation info from the current module, and external ones.
-      let info = hptVectInfo hsc_env `plusVectInfo` eps_vect_info eps
+-- | Vectorise a single monad, given the dynamic compiler flags and HscEnv.
+--
+vectoriseIO :: HscEnv -> ModGuts -> IO ModGuts
+vectoriseIO hsc_env guts
+ = do {   -- Get information about currently loaded external packages.
+      ; eps <- hscEPS hsc_env
 
-      -- Run the main VM computation.
-      Just (info', guts') <- initV backend hsc_env guts info (vectModule guts)
-      return (guts' { mg_vect_info = info' })
+          -- Combine vectorisation info from the current module, and external ones.
+      ; let info = hptVectInfo hsc_env `plusVectInfo` eps_vect_info eps
 
+          -- Run the main VM computation.
+      ; Just (info', guts') <- initV hsc_env guts info (vectModule guts)
+      ; return (guts' { mg_vect_info = info' })
+      }
 
 -- | Vectorise a single module, in the VM monad.
+--
 vectModule :: ModGuts -> VM ModGuts
-vectModule guts
- = do -- Vectorise the type environment.
-      -- This may add new TyCons and DataCons.
-      -- TODO: What new binds do we get back here?
-      (types', fam_insts, tc_binds) <- vectTypeEnv (mg_types guts)
-
-      (_, fam_inst_env) <- readGEnv global_fam_inst_env
+vectModule guts@(ModGuts { mg_types     = types
+                         , mg_binds     = binds
+                         , mg_fam_insts = fam_insts
+                         })
+ = do { dumpOptVt Opt_D_dump_vt_trace "Before vectorisation" $ 
+          pprCoreBindings binds
+ 
+          -- Vectorise the type environment.
+          -- This may add new TyCons and DataCons.
+      ; (types', new_fam_insts, tc_binds) <- vectTypeEnv types
+
+      ; (_, fam_inst_env) <- readGEnv global_fam_inst_env
 
       -- dicts   <- mapM buildPADict pa_insts
       -- workers <- mapM vectDataConWorkers pa_insts
 
-      -- Vectorise all the top level bindings.
-      binds'  <- mapM vectTopBind (mg_binds guts)
-
-      return $ guts { mg_types        = types'
-                    , mg_binds        = Rec tc_binds : binds'
-                    , mg_fam_inst_env = fam_inst_env
-                    , mg_fam_insts    = mg_fam_insts guts ++ fam_insts
-                    }
+          -- Vectorise all the top level bindings.
+      ; binds'  <- mapM vectTopBind binds
 
+      ; return $ guts { mg_types        = types'
+                      , mg_binds        = Rec tc_binds : binds'
+                      , mg_fam_inst_env = fam_inst_env
+                      , mg_fam_insts    = fam_insts ++ new_fam_insts
+                      }
+      }
 
 -- | Try to vectorise a top-level binding.
 --   If it doesn't vectorise then return it unharmed.
@@ -116,14 +123,14 @@ vectTopBind :: CoreBind -> VM CoreBind
 vectTopBind b@(NonRec var expr)
  = do
       (inline, _, expr') 	<- vectTopRhs [] var expr
-      var' 		<- vectTopBinder var inline expr'
+      var' <- vectTopBinder var inline expr'
 
       -- Vectorising the body may create other top-level bindings.
-      hs	<- takeHoisted
+      hs <- takeHoisted
 
       -- To get the same functionality as the original body we project
       -- out its vectorised version from the closure.
-      cexpr	<- tryConvert var var' expr
+      cexpr <- tryConvert var var' expr
 
       return . Rec $ (var, cexpr) : (var', expr') : hs
   `orElseV`
@@ -132,7 +139,7 @@ vectTopBind b@(NonRec var expr)
 vectTopBind b@(Rec bs)
  = do
       (vars', _, exprs') 
-	<- fixV $ \ ~(_, inlines, rhss) ->
+        <- fixV $ \ ~(_, inlines, rhss) ->
             do vars' <- sequence [vectTopBinder var inline rhs
                                       | (var, ~(inline, rhs)) <- zipLazy vars (zip inlines rhss)]
                (inlines', areScalars', exprs') 
@@ -152,67 +159,109 @@ 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
 --   version is @$v_foo@
 --
 --   NOTE: vectTopBinder *MUST* be lazy in inline and expr because of how it is
 --   used inside of fixV in vectTopBind
-vectTopBinder 
-	:: Var 		-- ^ Name of the binding.
-	-> Inline 	-- ^ Whether it should be inlined, used to annotate it.
-	-> CoreExpr 	-- ^ RHS of the binding, used to set the `Unfolding` of the returned `Var`.
-	-> VM Var	-- ^ Name of the vectorised binding.
-
+--
+vectTopBinder :: Var      -- ^ Name of the binding.
+              -> Inline   -- ^ Whether it should be inlined, used to annotate it.
+              -> CoreExpr -- ^ RHS of binding, used to set the 'Unfolding' of the returned 'Var'.
+              -> VM Var   -- ^ Name of the vectorised binding.
 vectTopBinder var inline expr
- = do
-      -- Vectorise the type attached to the var.
-      vty  <- vectType (idType var)
-
-      -- Make the vectorised version of binding's name, and set the unfolding used for inlining.
-      var' <- liftM (`setIdUnfoldingLazily` unfolding) 
-           $  cloneId mkVectOcc var vty
-
-      -- Add the mapping between the plain and vectorised name to the state.
-      defGlobalVar var var'
-
-      return var'
+ = do {   -- Vectorise the type attached to the var.
+      ; vty  <- vectType (idType var)
+      
+          -- If there is a vectorisation declartion for this binding, make sure that its type
+          --  matches
+      ; vectDecl <- lookupVectDecl var
+      ; case vectDecl of
+          Nothing                 -> return ()
+          Just (vdty, _) 
+            | coreEqType vty vdty -> return ()
+            | otherwise           -> 
+              cantVectorise ("Type mismatch in vectorisation pragma for " ++ show var) $
+                (text "Expected type" <+> ppr vty)
+                $$
+                (text "Inferred type" <+> ppr vdty)
+
+          -- Make the vectorised version of binding's name, and set the unfolding used for inlining
+      ; var' <- liftM (`setIdUnfoldingLazily` unfolding) 
+                $  cloneId mkVectOcc var vty
+
+          -- Add the mapping between the plain and vectorised name to the state.
+      ; defGlobalVar var var'
+
+      ; return var'
+    }
   where
     unfolding = case inline of
                   Inline arity -> mkInlineUnfolding (Just arity) expr
                   DontInline   -> noUnfolding
 
-
 -- | Vectorise the RHS of a top-level binding, in an empty local environment.
-vectTopRhs 
-	:: [Var]    -- ^ Names of all functions in the rec block
-	-> Var 		-- ^ Name of the binding.
-	-> CoreExpr	-- ^ Body of the binding.
-	-> VM (Inline, Bool, CoreExpr)
-
+--
+-- We need to distinguish three 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
+--     => generate vectorised code that uses a scalar 'map'/'zipWith' to lift the computation
+-- 
+-- (3) There is no vectorisation declaration for the variable
+--     => perform automatic vectorisation of the RHS
+--
+vectTopRhs :: [Var]           -- ^ Names of all functions in the rec block
+           -> Var             -- ^ Name of the binding.
+           -> CoreExpr        -- ^ Body of the binding.
+           -> VM ( Inline     -- (1) inline specification for the binding
+                 , Bool       -- (2) whether the right-hand side is a scalar computation
+                 , CoreExpr)  -- (3) the vectorised right-hand side
 vectTopRhs recFs var expr
- = dtrace (vcat [text "vectTopRhs", ppr expr])
- $ closedV
- $ 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)
-
+  = closedV
+  $ do { traceVt ("vectTopRhs of " ++ show var) $ ppr expr
+  
+       ; globalScalar <- isGlobalScalar var
+       ; vectDecl     <- lookupVectDecl var
+       ; rhs globalScalar vectDecl
+       }
+  where
+    rhs _globalScalar (Just (_, expr'))               -- Case (1)
+      = return (inlineMe, False, expr')
+    rhs True          _vectDecl                       -- Case (2)
+      = return (inlineMe, True, scalarRHS)
+                          -- FIXME: that True is not enough to register scalarness
+    rhs False         _vectDecl                       -- Case (3)
+      = do { let fvs = freeVars expr
+           ; (inline, isScalar, vexpr) <- inBind var $
+                                            vectPolyExpr (isLoopBreaker $ idOccInfo var) recFs fvs
+           ; if isScalar 
+             then addGlobalScalar var
+             else deleteGlobalScalar var
+           ; return (inline, isScalar, vectorised vexpr)
+           }
+      
+    -- For scalar right-hand sides, we know that the original binding will remain unaltered
+    -- (hence, we can refer to it without risk of cycles) - cf, 'tryConvert'.
+    scalarRHS = panic "Vectorise.scalarRHS: not implemented yet"
 
 -- | Project out the vectorised version of a binding from some closure,
---	or return the original body if that doesn't work.	
-tryConvert 
-	:: Var	 	-- ^ Name of the original binding (eg @foo@)
-	-> Var 		-- ^ Name of vectorised version of binding (eg @$vfoo@)
-	-> CoreExpr	-- ^ The original body of the binding.
-	-> VM CoreExpr
-
+--   or return the original body if that doesn't work or the binding is scalar. 
+--
+tryConvert :: Var       -- ^ Name of the original binding (eg @foo@)
+           -> Var       -- ^ Name of vectorised version of binding (eg @$vfoo@)
+           -> CoreExpr  -- ^ The original body of the binding.
+           -> VM CoreExpr
 tryConvert var vect_var rhs
-  = fromVect (idType var) (Var vect_var) `orElseV` return rhs
-
+  = do { globalScalar <- isGlobalScalar var
+       ; if globalScalar
+         then
+           return rhs
+         else
+           fromVect (idType var) (Var vect_var) `orElseV` return rhs
+       }