ghc-8.4.1: vectorise/Vectorise.hs
-- Main entry point to the vectoriser. It is invoked iff the option '-fvectorise' is passed.
--
-- This module provides the function 'vectorise', which vectorises an entire (desugared) module.
-- It vectorises all type declarations and value bindings. It also processes all VECTORISE pragmas
-- (aka vectorisation declarations), which can lead to the vectorisation of imported data types
-- and the enrichment of imported functions with vectorised versions.
module Vectorise ( vectorise )
where
import GhcPrelude
import Vectorise.Type.Env
import Vectorise.Type.Type
import Vectorise.Convert
import Vectorise.Utils.Hoisting
import Vectorise.Exp
import Vectorise.Env
import Vectorise.Monad
import HscTypes hiding ( MonadThings(..) )
import CoreUnfold ( mkInlineUnfoldingWithArity )
import PprCore
import CoreSyn
import CoreMonad ( CoreM, getHscEnv )
import Type
import Id
import DynFlags
import Outputable
import Util ( zipLazy )
import MonadUtils
import Control.Monad
-- |Vectorise a single module.
--
vectorise :: ModGuts -> CoreM ModGuts
vectorise guts
= do { hsc_env <- getHscEnv
; liftIO $ vectoriseIO hsc_env guts
}
-- 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
-- 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@(ModGuts { mg_tcs = tycons
, mg_binds = binds
, mg_fam_insts = fam_insts
, mg_vect_decls = vect_decls
})
= do { dumpOptVt Opt_D_dump_vt_trace "Before vectorisation" $
pprCoreBindings binds
-- Pick out all 'VECTORISE [SCALAR] type' and 'VECTORISE class' pragmas
; let ty_vect_decls = [vd | vd@(VectType _ _ _) <- vect_decls]
cls_vect_decls = [vd | vd@(VectClass _) <- vect_decls]
-- Vectorise the type environment. This will add vectorised
-- type constructors, their representations, and the
-- corresponding data constructors. Moreover, we produce
-- bindings for dfuns and family instances of the classes
-- and type families used in the DPH library to represent
-- array types.
; (new_tycons, new_fam_insts, tc_binds) <- vectTypeEnv tycons ty_vect_decls cls_vect_decls
-- 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
-- NB: Need to vectorise the imported bindings first (local bindings may depend on them).
; let impBinds = [(imp_id, expr) | Vect imp_id expr <- vect_decls, isGlobalId imp_id]
; binds_imp <- mapM vectImpBind impBinds
; binds_top <- mapM vectTopBind binds
; return $ guts { mg_tcs = tycons ++ new_tycons
-- we produce no new classes or instances, only new class type constructors
-- and dfuns
, mg_binds = Rec tc_binds : (binds_top ++ binds_imp)
, 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, or if it is entirely scalar, then
-- omit vectorisation of that binding.
--
-- For example, for the binding
--
-- @
-- foo :: Int -> Int
-- foo = \x -> x + x
-- @
--
-- we get
-- @
-- foo :: Int -> Int
-- foo = \x -> vfoo $: x
--
-- v_foo :: Closure void vfoo lfoo
-- v_foo = closure vfoo lfoo void
--
-- vfoo :: Void -> Int -> Int
-- vfoo = ...
--
-- lfoo :: PData Void -> PData Int -> PData Int
-- lfoo = ...
-- @
--
-- @vfoo@ is the "vectorised", or scalar, version that does the same as the original function foo,
-- but takes an explicit environment.
--
-- @lfoo@ is the "lifted" version that works on arrays.
--
-- @v_foo@ combines both of these into a `Closure` that also contains the environment.
--
-- The original binding @foo@ is rewritten to call the vectorised version present in the closure.
--
-- Vectorisation may be suppressed by annotating a binding with a 'NOVECTORISE' pragma. If this
-- pragma is used in a group of mutually recursive bindings, either all or no binding must have
-- the pragma. If only some bindings are annotated, a fatal error is being raised. (In the case of
-- scalar bindings, we only omit vectorisation if all bindings in a group are scalar.)
--
-- FIXME: Once we support partial vectorisation, we may be able to vectorise parts of a group, or
-- we may emit a warning and refrain from vectorising the entire group.
--
vectTopBind :: CoreBind -> VM CoreBind
vectTopBind b@(NonRec var expr)
= do
{ traceVt "= Vectorise non-recursive top-level variable" (ppr var)
; (hasNoVect, vectDecl) <- lookupVectDecl var
; if hasNoVect
then do
{ -- 'NOVECTORISE' pragma => leave this binding as it is
; traceVt "NOVECTORISE" $ ppr var
; return b
}
else do
{ vectRhs <- case vectDecl of
Just (_, expr') ->
-- 'VECTORISE' pragma => just use the provided vectorised rhs
do
{ traceVt "VECTORISE" $ ppr var
; addGlobalParallelVar var
; return $ Just (False, inlineMe, expr')
}
Nothing ->
-- no pragma => standard vectorisation of rhs
do
{ traceVt "[Vanilla]" $ ppr var <+> char '=' <+> ppr expr
; vectTopExpr var expr
}
; hs <- takeHoisted -- make sure we clean those out (even if we skip)
; case vectRhs of
{ Nothing ->
-- scalar binding => leave this binding as it is
do
{ traceVt "scalar binding [skip]" $ ppr var
; return b
}
; Just (parBind, inline, expr') -> do
{
-- vanilla case => create an appropriate top-level binding & add it to the vectorisation map
; when parBind $
addGlobalParallelVar var
; var' <- vectTopBinder var inline expr'
-- We replace the original top-level binding by a value projected from the vectorised
-- closure and add any newly created hoisted top-level bindings.
; cexpr <- tryConvert var var' expr
; return . Rec $ (var, cexpr) : (var', expr') : hs
} } } }
`orElseErrV`
do
{ emitVt " Could NOT vectorise top-level binding" $ ppr var
; return b
}
vectTopBind b@(Rec binds)
= do
{ traceVt "= Vectorise recursive top-level variables" $ ppr vars
; vectDecls <- mapM lookupVectDecl vars
; let hasNoVects = map fst vectDecls
; if and hasNoVects
then do
{ -- 'NOVECTORISE' pragmas => leave this entire binding group as it is
; traceVt "NOVECTORISE" $ ppr vars
; return b
}
else do
{ if or hasNoVects
then do
{ -- Inconsistent 'NOVECTORISE' pragmas => bail out
; dflags <- getDynFlags
; cantVectorise dflags noVectoriseErr (ppr b)
}
else do
{ traceVt "[Vanilla]" $ vcat [ppr var <+> char '=' <+> ppr expr | (var, expr) <- binds]
-- For all bindings *with* a pragma, just use the pragma-supplied vectorised expression
; newBindsWPragma <- concat <$>
sequence [ vectTopBindAndConvert bind inlineMe expr'
| (bind, (_, Just (_, expr'))) <- zip binds vectDecls]
-- Standard vectorisation of all rhses that are *without* a pragma.
-- NB: The reason for 'fixV' is rather subtle: 'vectTopBindAndConvert' adds entries for
-- the bound variables in the recursive group to the vectorisation map, which in turn
-- are needed by 'vectPolyExprs' (unless it returns 'Nothing').
; let bindsWOPragma = [bind | (bind, (_, Nothing)) <- zip binds vectDecls]
; (newBinds, _) <- fixV $
\ ~(_, exprs') ->
do
{ -- Create appropriate top-level bindings, enter them into the vectorisation map, and
-- vectorise the right-hand sides
; newBindsWOPragma <- concat <$>
sequence [vectTopBindAndConvert bind inline expr
| (bind, ~(inline, expr)) <- zipLazy bindsWOPragma exprs']
-- irrefutable pattern and 'zipLazy' to tie the knot;
-- hence, can't use 'zipWithM'
; vectRhses <- vectTopExprs bindsWOPragma
; hs <- takeHoisted -- make sure we clean those out (even if we skip)
; case vectRhses of
Nothing ->
-- scalar bindings => skip all bindings except those with pragmas and retract the
-- entries into the vectorisation map for the scalar bindings
do
{ traceVt "scalar bindings [skip]" $ ppr vars
; mapM_ (undefGlobalVar . fst) bindsWOPragma
; return (bindsWOPragma ++ newBindsWPragma, exprs')
}
Just (parBind, exprs') ->
-- vanilla case => record parallel variables and return the final bindings
do
{ when parBind $
mapM_ addGlobalParallelVar vars
; return (newBindsWOPragma ++ newBindsWPragma ++ hs, exprs')
}
}
; return $ Rec newBinds
} } }
`orElseErrV`
do
{ emitVt " Could NOT vectorise top-level bindings" $ ppr vars
; return b
}
where
vars = map fst binds
noVectoriseErr = "NOVECTORISE must be used on all or no bindings of a recursive group"
-- Replace the original top-level bindings by a values projected from the vectorised
-- closures and add any newly created hoisted top-level bindings to the group.
vectTopBindAndConvert (var, expr) inline expr'
= do
{ var' <- vectTopBinder var inline expr'
; cexpr <- tryConvert var var' expr
; return [(var, cexpr), (var', expr')]
}
-- Add a vectorised binding to an imported top-level variable that has a VECTORISE pragma
-- in this module.
--
-- RESTRICTION: Currently, we cannot use the pragma for mutually recursive definitions.
--
vectImpBind :: (Id, CoreExpr) -> VM CoreBind
vectImpBind (var, expr)
= do
{ traceVt "= Add vectorised binding to imported variable" (ppr var)
; var' <- vectTopBinder var inlineMe expr
; return $ NonRec var' expr
}
-- |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 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)
-- If there is a vectorisation declaration for this binding, make sure its type matches
; (_, vectDecl) <- lookupVectDecl var
; case vectDecl of
Nothing -> return ()
Just (vdty, _)
| eqType vty vdty -> return ()
| otherwise ->
do
{ dflags <- getDynFlags
; cantVectorise dflags ("Type mismatch in vectorisation pragma for " ++ showPpr dflags 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 (`setIdUnfolding` unfolding)
$ mkVectId 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 -> mkInlineUnfoldingWithArity 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
-}
-- |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
tryConvert var vect_var rhs
= fromVect (idType var) (Var vect_var)
`orElseErrV`
do
{ emitVt " Could NOT call vectorised from original version" $ ppr var <+> dcolon <+> ppr (idType var)
; return rhs
}