diff options
Diffstat (limited to 'compiler/vectorise/Vectorise.hs')
-rw-r--r-- | compiler/vectorise/Vectorise.hs | 356 |
1 files changed, 0 insertions, 356 deletions
diff --git a/compiler/vectorise/Vectorise.hs b/compiler/vectorise/Vectorise.hs deleted file mode 100644 index 2e09adbbbe..0000000000 --- a/compiler/vectorise/Vectorise.hs +++ /dev/null @@ -1,356 +0,0 @@ --- 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 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 - } |