From 46fa261eee74c1c1a1be52f9394ff131183024da Mon Sep 17 00:00:00 2001 From: Manuel M T Chakravarty Date: Wed, 17 Aug 2011 14:41:59 +1000 Subject: [PATCH] Add VECTORISE [SCALAR] type pragma - Pragma to determine how a given type is vectorised - At this stage only the VECTORISE SCALAR variant is used by the vectoriser. - '{-# VECTORISE SCALAR type t #-}' implies that 't' cannot contain parallel arrays and may be used in vectorised code. However, its constructors can only be used in scalar code. We use this, e.g., for 'Int'. - May be used on imported types See also http://hackage.haskell.org/trac/ghc/wiki/DataParallel/VectPragma --- compiler/coreSyn/CoreFVs.lhs | 2 + compiler/coreSyn/CoreSubst.lhs | 3 +- compiler/coreSyn/CoreSyn.lhs | 9 +- compiler/coreSyn/PprCore.lhs | 11 +- compiler/deSugar/Desugar.lhs | 6 +- compiler/hsSyn/HsDecls.lhs | 32 ++- compiler/parser/Parser.y.pp | 14 +- compiler/rename/RnSource.lhs | 25 +- compiler/typecheck/TcBinds.lhs | 14 ++ compiler/typecheck/TcHsSyn.lhs | 17 +- compiler/vectorise/Vectorise.hs | 24 +- compiler/vectorise/Vectorise/Builtins.hs | 13 +- .../Vectorise/Builtins/Initialise.hs | 7 +- compiler/vectorise/Vectorise/Env.hs | 72 +++--- compiler/vectorise/Vectorise/Exp.hs | 88 +++---- compiler/vectorise/Vectorise/Monad/Global.hs | 71 +++--- compiler/vectorise/Vectorise/Type/Classify.hs | 76 +++--- compiler/vectorise/Vectorise/Type/Env.hs | 224 +++++++++++------- compiler/vectorise/Vectorise/Type/Repr.hs | 30 +-- 19 files changed, 449 insertions(+), 289 deletions(-) diff --git a/compiler/coreSyn/CoreFVs.lhs b/compiler/coreSyn/CoreFVs.lhs index f5cd76254d..71ddc8c8cc 100644 --- a/compiler/coreSyn/CoreFVs.lhs +++ b/compiler/coreSyn/CoreFVs.lhs @@ -334,6 +334,8 @@ vectsFreeVars = foldr (unionVarSet . vectFreeVars) emptyVarSet vectFreeVars (Vect _ Nothing) = noFVs vectFreeVars (Vect _ (Just rhs)) = expr_fvs rhs isLocalId emptyVarSet vectFreeVars (NoVect _) = noFVs + vectFreeVars (VectType _ _) = noFVs + -- this function is only concerned with values, not types \end{code} diff --git a/compiler/coreSyn/CoreSubst.lhs b/compiler/coreSyn/CoreSubst.lhs index ca0fbd5a52..effc5f8459 100644 --- a/compiler/coreSyn/CoreSubst.lhs +++ b/compiler/coreSyn/CoreSubst.lhs @@ -735,7 +735,8 @@ substVects subst = map (substVect subst) substVect :: Subst -> CoreVect -> CoreVect substVect _subst (Vect v Nothing) = Vect v Nothing substVect subst (Vect v (Just rhs)) = Vect v (Just (simpleOptExprWith subst rhs)) -substVect _subst (NoVect v) = NoVect v +substVect _subst vd@(NoVect _) = vd +substVect _subst vd@(VectType _ _) = vd ------------------ substVarSet :: Subst -> VarSet -> VarSet diff --git a/compiler/coreSyn/CoreSyn.lhs b/compiler/coreSyn/CoreSyn.lhs index ccb87e7782..f91a8f6b23 100644 --- a/compiler/coreSyn/CoreSyn.lhs +++ b/compiler/coreSyn/CoreSyn.lhs @@ -87,12 +87,13 @@ import Coercion import Name import Literal import DataCon +import TyCon import BasicTypes import FastString import Outputable import Util -import Data.Data +import Data.Data hiding (TyCon) import Data.Word infixl 4 `mkApps`, `mkTyApps`, `mkVarApps`, `App`, `mkCoApps` @@ -428,9 +429,9 @@ Representation of desugared vectorisation declarations that are fed to the vecto 'ModGuts'). \begin{code} -data CoreVect = Vect Id (Maybe CoreExpr) - | NoVect Id - +data CoreVect = Vect Id (Maybe CoreExpr) + | NoVect Id + | VectType TyCon (Maybe Type) \end{code} diff --git a/compiler/coreSyn/PprCore.lhs b/compiler/coreSyn/PprCore.lhs index 58a940c72a..cf9292408f 100644 --- a/compiler/coreSyn/PprCore.lhs +++ b/compiler/coreSyn/PprCore.lhs @@ -473,8 +473,11 @@ pprRule (Rule { ru_name = name, ru_act = act, ru_fn = fn, \begin{code} instance Outputable CoreVect where - ppr (Vect var Nothing) = ptext (sLit "VECTORISE SCALAR") <+> ppr var - ppr (Vect var (Just e)) = hang (ptext (sLit "VECTORISE") <+> ppr var <+> char '=') - 4 (pprCoreExpr e) - ppr (NoVect var) = ptext (sLit "NOVECTORISE") <+> ppr var + ppr (Vect var Nothing) = ptext (sLit "VECTORISE SCALAR") <+> ppr var + ppr (Vect var (Just e)) = hang (ptext (sLit "VECTORISE") <+> ppr var <+> char '=') + 4 (pprCoreExpr e) + ppr (NoVect var) = ptext (sLit "NOVECTORISE") <+> ppr var + ppr (VectType var Nothing) = ptext (sLit "VECTORISE SCALAR type") <+> ppr var + ppr (VectType var (Just ty)) = hang (ptext (sLit "VECTORISE type") <+> ppr var <+> char '=') + 4 (ppr ty) \end{code} diff --git a/compiler/deSugar/Desugar.lhs b/compiler/deSugar/Desugar.lhs index 2f265221e8..2c5a3c820b 100644 --- a/compiler/deSugar/Desugar.lhs +++ b/compiler/deSugar/Desugar.lhs @@ -403,7 +403,11 @@ dsVect (L loc (HsVect (L _ v) rhs)) = putSrcSpanDs loc $ do { rhs' <- fmapMaybeM dsLExpr rhs ; return $ Vect v rhs' - } + } dsVect (L _loc (HsNoVect (L _ v))) = return $ NoVect v +dsVect (L _loc (HsVectTypeOut tycon ty)) + = return $ VectType tycon ty +dsVect vd@(L _ (HsVectTypeIn _ _ty)) + = pprPanic "Desugar.dsVect: unexpected 'HsVectTypeIn'" (ppr vd) \end{code} diff --git a/compiler/hsSyn/HsDecls.lhs b/compiler/hsSyn/HsDecls.lhs index 9d3382fd8a..c1b06809d7 100644 --- a/compiler/hsSyn/HsDecls.lhs +++ b/compiler/hsSyn/HsDecls.lhs @@ -59,6 +59,7 @@ import HsBinds import HsPat import HsTypes import HsDoc +import TyCon import NameSet import {- Kind parts of -} Type import BasicTypes @@ -72,7 +73,7 @@ import SrcLoc import FastString import Control.Monad ( liftM ) -import Data.Data +import Data.Data hiding (TyCon) import Data.Maybe ( isJust ) \end{code} @@ -1014,6 +1015,9 @@ A vectorisation pragma, one of {-# VECTORISE f = closure1 g (scalar_map g) #-} {-# VECTORISE SCALAR f #-} {-# NOVECTORISE f #-} + + {-# VECTORISE type T = ty #-} + {-# VECTORISE SCALAR type T #-} Note [Typechecked vectorisation pragmas] ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ @@ -1036,11 +1040,19 @@ data VectDecl name (Maybe (LHsExpr name)) -- 'Nothing' => SCALAR declaration | HsNoVect (Located name) + | HsVectTypeIn -- pre type-checking + (Located name) + (Maybe (LHsType name)) -- 'Nothing' => SCALAR declaration + | HsVectTypeOut -- post type-checking + TyCon + (Maybe Type) -- 'Nothing' => SCALAR declaration deriving (Data, Typeable) -lvectDeclName :: LVectDecl name -> name -lvectDeclName (L _ (HsVect (L _ name) _)) = name -lvectDeclName (L _ (HsNoVect (L _ name))) = name +lvectDeclName :: Outputable name => LVectDecl name -> name +lvectDeclName (L _ (HsVect (L _ name) _)) = name +lvectDeclName (L _ (HsNoVect (L _ name))) = name +lvectDeclName (L _ (HsVectTypeIn (L _ name) _)) = name +lvectDeclName (L _ (HsVectTypeOut name _)) = pprPanic "HsDecls.HsVectTypeOut" (ppr name) instance OutputableBndr name => Outputable (VectDecl name) where ppr (HsVect v Nothing) @@ -1051,6 +1063,18 @@ instance OutputableBndr name => Outputable (VectDecl name) where pprExpr (unLoc rhs) <+> text "#-}" ] ppr (HsNoVect v) = sep [text "{-# NOVECTORISE" <+> ppr v <+> text "#-}" ] + ppr (HsVectTypeIn t Nothing) + = sep [text "{-# VECTORISE SCALAR type" <+> ppr t <+> text "#-}" ] + ppr (HsVectTypeIn t (Just ty)) + = sep [text "{-# VECTORISE type" <+> ppr t, + nest 4 $ + ppr (unLoc ty) <+> text "#-}" ] + ppr (HsVectTypeOut t Nothing) + = sep [text "{-# VECTORISE SCALAR type" <+> ppr t <+> text "#-}" ] + ppr (HsVectTypeOut t (Just ty)) + = sep [text "{-# VECTORISE type" <+> ppr t, + nest 4 $ + ppr ty <+> text "#-}" ] \end{code} %************************************************************************ diff --git a/compiler/parser/Parser.y.pp b/compiler/parser/Parser.y.pp index 42988feeeb..c1e1d8810a 100644 --- a/compiler/parser/Parser.y.pp +++ b/compiler/parser/Parser.y.pp @@ -563,8 +563,8 @@ topdecls :: { OrdList (LHsDecl RdrName) } | topdecl { $1 } topdecl :: { OrdList (LHsDecl RdrName) } - : cl_decl { unitOL (L1 (TyClD (unLoc $1))) } - | ty_decl { unitOL (L1 (TyClD (unLoc $1))) } + : cl_decl { unitOL (L1 (TyClD (unLoc $1))) } + | ty_decl { unitOL (L1 (TyClD (unLoc $1))) } | 'instance' inst_type where_inst { let (binds, sigs, ats, _) = cvBindsAndSigs (unLoc $3) in @@ -575,9 +575,13 @@ topdecl :: { OrdList (LHsDecl RdrName) } | '{-# DEPRECATED' deprecations '#-}' { $2 } | '{-# WARNING' warnings '#-}' { $2 } | '{-# RULES' rules '#-}' { $2 } - | '{-# VECTORISE_SCALAR' qvar '#-}' { unitOL $ LL $ VectD (HsVect $2 Nothing) } - | '{-# VECTORISE' qvar '=' exp '#-}' { unitOL $ LL $ VectD (HsVect $2 (Just $4)) } - | '{-# NOVECTORISE' qvar '#-}' { unitOL $ LL $ VectD (HsNoVect $2) } + | '{-# VECTORISE_SCALAR' qvar '#-}' { unitOL $ LL $ VectD (HsVect $2 Nothing) } + | '{-# VECTORISE' qvar '=' exp '#-}' { unitOL $ LL $ VectD (HsVect $2 (Just $4)) } + | '{-# NOVECTORISE' qvar '#-}' { unitOL $ LL $ VectD (HsNoVect $2) } + | '{-# VECTORISE_SCALAR' 'type' qtycon '#-}' + { unitOL $ LL $ VectD (HsVectTypeIn $3 Nothing) } + | '{-# VECTORISE' 'type' qtycon '=' ctype '#-}' + { unitOL $ LL $ VectD (HsVectTypeIn $3 (Just $5)) } | annotation { unitOL $1 } | decl { unLoc $1 } diff --git a/compiler/rename/RnSource.lhs b/compiler/rename/RnSource.lhs index 3d73e4b7bc..64feaed8e4 100644 --- a/compiler/rename/RnSource.lhs +++ b/compiler/rename/RnSource.lhs @@ -659,24 +659,37 @@ badRuleLhsErr name lhs bad_e \begin{code} rnHsVectDecl :: VectDecl RdrName -> RnM (VectDecl Name, FreeVars) rnHsVectDecl (HsVect var Nothing) - = do { var' <- wrapLocM lookupTopBndrRn var + = do { var' <- lookupLocatedTopBndrRn var ; return (HsVect var' Nothing, unitFV (unLoc var')) } rnHsVectDecl (HsVect var (Just rhs)) - = do { var' <- wrapLocM lookupTopBndrRn var + = do { var' <- lookupLocatedTopBndrRn var ; (rhs', fv_rhs) <- rnLExpr rhs ; return (HsVect var' (Just rhs'), fv_rhs `addOneFV` unLoc var') } rnHsVectDecl (HsNoVect var) - = do { var' <- wrapLocM lookupTopBndrRn var + = do { var' <- lookupLocatedTopBndrRn var ; return (HsNoVect var', unitFV (unLoc var')) } +rnHsVectDecl (HsVectTypeIn tycon Nothing) + = do { tycon' <- lookupLocatedOccRn tycon + ; return (HsVectTypeIn tycon' Nothing, unitFV (unLoc tycon')) + } +rnHsVectDecl (HsVectTypeIn tycon (Just ty)) + = do { tycon' <- lookupLocatedOccRn tycon + ; (ty', fv_ty) <- rnHsTypeFVs vect_doc ty + ; return (HsVectTypeIn tycon' (Just ty'), fv_ty `addOneFV` unLoc tycon') + } + where + vect_doc = text "In the VECTORISE pragma for type constructor" <+> quotes (ppr tycon) +rnHsVectDecl (HsVectTypeOut _ _) + = panic "RnSource.rnHsVectDecl: Unexpected 'HsVectTypeOut'" \end{code} %********************************************************* -%* * +%* * \subsection{Type, class and iface sig declarations} -%* * +%* * %********************************************************* @rnTyDecl@ uses the `global name function' to create a new type @@ -711,7 +724,7 @@ rnTyClDecl (ForeignType {tcdLName = name, tcdExtName = ext_name}) return (ForeignType {tcdLName = name', tcdExtName = ext_name}, emptyFVs) --- all flavours of type family declarations ("type family", "newtype fanily", +-- all flavours of type family declarations ("type family", "newtype family", -- and "data family") rnTyClDecl tydecl@TyFamily {} = rnFamily tydecl bindTyVarsFV diff --git a/compiler/typecheck/TcBinds.lhs b/compiler/typecheck/TcBinds.lhs index 0f404c6923..599b5334a1 100644 --- a/compiler/typecheck/TcBinds.lhs +++ b/compiler/typecheck/TcBinds.lhs @@ -24,6 +24,7 @@ import TcSimplify import TcHsType import TcPat import TcMType +import TyCon import TcType -- import Coercion import TysPrim @@ -682,10 +683,23 @@ tcVect (HsNoVect name) do { id <- wrapLocM tcLookupId name ; return $ HsNoVect id } +tcVect (HsVectTypeIn lname@(L _ name) ty) + = addErrCtxt (vectCtxt lname) $ + do { tycon <- tcLookupTyCon name + ; checkTc (tyConArity tycon /= 0) scalarTyConMustBeNullary + + ; ty' <- fmapMaybeM dsHsType ty + ; return $ HsVectTypeOut tycon ty' + } +tcVect (HsVectTypeOut _ _) + = panic "TcBinds.tcVect: Unexpected 'HsVectTypeOut'" vectCtxt :: Located Name -> SDoc vectCtxt name = ptext (sLit "When checking the vectorisation declaration for") <+> ppr name +scalarTyConMustBeNullary :: Message +scalarTyConMustBeNullary = ptext (sLit "VECTORISE SCALAR type constructor must be nullary") + -------------- -- If typechecking the binds fails, then return with each -- signature-less binder given type (forall a.a), to minimise diff --git a/compiler/typecheck/TcHsSyn.lhs b/compiler/typecheck/TcHsSyn.lhs index 699869c824..65bd79c204 100644 --- a/compiler/typecheck/TcHsSyn.lhs +++ b/compiler/typecheck/TcHsSyn.lhs @@ -1,4 +1,4 @@ -1% +% % (c) The University of Glasgow 2006 % (c) The AQUA Project, Glasgow University, 1996-1998 % @@ -1022,19 +1022,20 @@ zonkVects :: ZonkEnv -> [LVectDecl TcId] -> TcM [LVectDecl Id] zonkVects env = mappM (wrapLocM (zonkVect env)) zonkVect :: ZonkEnv -> VectDecl TcId -> TcM (VectDecl Id) -zonkVect env (HsVect v Nothing) - = do { v' <- wrapLocM (zonkIdBndr env) v - ; return $ HsVect v' Nothing - } -zonkVect env (HsVect v (Just e)) +zonkVect env (HsVect v e) = do { v' <- wrapLocM (zonkIdBndr env) v - ; e' <- zonkLExpr env e - ; return $ HsVect v' (Just e') + ; e' <- fmapMaybeM (zonkLExpr env) e + ; return $ HsVect v' e' } zonkVect env (HsNoVect v) = do { v' <- wrapLocM (zonkIdBndr env) v ; return $ HsNoVect v' } +zonkVect _env (HsVectTypeOut t ty) + = do { ty' <- fmapMaybeM zonkTypeZapping ty + ; return $ HsVectTypeOut t ty' + } +zonkVect _ (HsVectTypeIn _ _) = panic "TcHsSyn.zonkVect: HsVectTypeIn" \end{code} %************************************************************************ diff --git a/compiler/vectorise/Vectorise.hs b/compiler/vectorise/Vectorise.hs index f5795424da..1d54b3803d 100644 --- a/compiler/vectorise/Vectorise.hs +++ b/compiler/vectorise/Vectorise.hs @@ -1,3 +1,9 @@ +-- 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 @@ -55,22 +61,22 @@ vectoriseIO hsc_env guts -- | Vectorise a single module, in the VM monad. -- vectModule :: ModGuts -> VM ModGuts -vectModule guts@(ModGuts { mg_types = types - , mg_binds = binds - , mg_fam_insts = fam_insts +vectModule guts@(ModGuts { mg_types = types + , mg_binds = binds + , mg_fam_insts = fam_insts + , mg_vect_decls = vect_decls }) = 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 + -- Vectorise the type environment. This will add vectorised type constructors, their + -- representaions, and the conrresponding 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. + ; (types', new_fam_insts, tc_binds) <- vectTypeEnv types [vd | vd@(VectType _ _) <- vect_decls] ; (_, 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 binds diff --git a/compiler/vectorise/Vectorise/Builtins.hs b/compiler/vectorise/Vectorise/Builtins.hs index 125d26482e..46da134fba 100644 --- a/compiler/vectorise/Vectorise/Builtins.hs +++ b/compiler/vectorise/Vectorise/Builtins.hs @@ -1,12 +1,9 @@ - --- | Builtin types and functions used by the vectoriser. --- The source program uses functions from Data.Array.Parallel, which the vectoriser rewrites --- to use equivalent vectorised versions in the DPH backend packages. --- --- The `Builtins` structure holds the name of all the things in the DPH packages --- we will need. We can get specific things using the selectors, which print a --- civilized panic message if the specified thing cannot be found. +-- Types and functions declared in the DPH packages and used by the vectoriser. -- +-- The @Builtins@ structure holds the name of all the things in the DPH packages that appear in +-- code generated by the vectoriser. We can get specific things using the selectors, which print a +-- civilized panic message if the specified thing cannot be found. + module Vectorise.Builtins ( -- * Builtins Builtins(..), diff --git a/compiler/vectorise/Vectorise/Builtins/Initialise.hs b/compiler/vectorise/Vectorise/Builtins/Initialise.hs index 9fdf3ba8f5..9c21eef6f9 100644 --- a/compiler/vectorise/Vectorise/Builtins/Initialise.hs +++ b/compiler/vectorise/Vectorise/Builtins/Initialise.hs @@ -1,3 +1,4 @@ +-- Set up the data structures provided by 'Vectorise.Builtins'. module Vectorise.Builtins.Initialise ( -- * Initialisation @@ -81,10 +82,10 @@ initBuiltins pkg -- From dph-common:Data.Array.Parallel.PArray.Types voidTyCon <- externalTyCon dph_PArray_Types (fsLit "Void") - voidVar <- externalVar dph_PArray_Types (fsLit "void") - fromVoidVar <- externalVar dph_PArray_Types (fsLit "fromVoid") + voidVar <- externalVar dph_PArray_Types (fsLit "void") + fromVoidVar <- externalVar dph_PArray_Types (fsLit "fromVoid") wrapTyCon <- externalTyCon dph_PArray_Types (fsLit "Wrap") - sum_tcs <- mapM (externalTyCon dph_PArray_Types) (numbered "Sum" 2 mAX_DPH_SUM) + sum_tcs <- mapM (externalTyCon dph_PArray_Types) (numbered "Sum" 2 mAX_DPH_SUM) -- from dph-common:Data.Array.Parallel.PArray.PDataInstances pvoidVar <- externalVar dph_PArray_PDataInstances (fsLit "pvoid") diff --git a/compiler/vectorise/Vectorise/Env.hs b/compiler/vectorise/Vectorise/Env.hs index d70f09affd..3fbfb924d5 100644 --- a/compiler/vectorise/Vectorise/Env.hs +++ b/compiler/vectorise/Vectorise/Env.hs @@ -76,55 +76,56 @@ emptyLocalEnv = LocalEnv { -- data GlobalEnv = GlobalEnv - { global_vars :: VarEnv Var + { global_vars :: VarEnv Var -- ^Mapping from global variables to their vectorised versions — aka the /vectorisation -- map/. - , global_vect_decls :: VarEnv (Type, CoreExpr) + , global_vect_decls :: VarEnv (Type, CoreExpr) -- ^Mapping from global variables that have a vectorisation declaration to the right-hand -- side of that declaration and its type. This mapping only applies to non-scalar -- vectorisation declarations. All variables with a scalar vectorisation declaration are -- mentioned in 'global_scalars_vars'. - , global_scalar_vars :: VarSet + , global_scalar_vars :: VarSet -- ^Purely scalar variables. Code which mentions only these variables doesn't have to be -- lifted. This includes variables from the current module that have a scalar -- vectorisation declaration and those that the vectoriser determines to be scalar. - , global_scalar_tycons :: NameSet - -- ^Type constructors whose values can only contain scalar data. Scalar code may only - -- operate on such data. + , global_scalar_tycons :: NameSet + -- ^Type constructors whose values can only contain scalar data and that appear in a + -- 'VECTORISE SCALAR type' pragma in the current or an imported module. Scalar code may + -- only operate on such data. - , global_novect_vars :: VarSet + , global_novect_vars :: VarSet -- ^Variables that are not vectorised. (They may be referenced in the right-hand sides -- of vectorisation declarations, though.) - , global_exported_vars :: VarEnv (Var, Var) + , global_exported_vars :: VarEnv (Var, Var) -- ^Exported variables which have a vectorised version. - , global_tycons :: NameEnv TyCon + , global_tycons :: NameEnv TyCon -- ^Mapping from TyCons to their vectorised versions. -- TyCons which do not have to be vectorised are mapped to themselves. - , global_datacons :: NameEnv DataCon + , global_datacons :: NameEnv DataCon -- ^Mapping from DataCons to their vectorised versions. - , global_pa_funs :: NameEnv Var + , global_pa_funs :: NameEnv Var -- ^Mapping from TyCons to their PA dfuns. - , global_pr_funs :: NameEnv Var + , global_pr_funs :: NameEnv Var -- ^Mapping from TyCons to their PR dfuns. - , global_boxed_tycons :: NameEnv TyCon + , global_boxed_tycons :: NameEnv TyCon -- ^Mapping from unboxed TyCons to their boxed versions. - , global_inst_env :: (InstEnv, InstEnv) + , global_inst_env :: (InstEnv, InstEnv) -- ^External package inst-env & home-package inst-env for class instances. - , global_fam_inst_env :: FamInstEnvs + , global_fam_inst_env :: FamInstEnvs -- ^External package inst-env & home-package inst-env for family instances. - , global_bindings :: [(Var, CoreExpr)] + , global_bindings :: [(Var, CoreExpr)] -- ^Hoisted bindings. } @@ -133,25 +134,26 @@ data GlobalEnv initGlobalEnv :: VectInfo -> [CoreVect] -> (InstEnv, InstEnv) -> FamInstEnvs -> GlobalEnv initGlobalEnv info vectDecls instEnvs famInstEnvs = GlobalEnv - { global_vars = mapVarEnv snd $ vectInfoVar info - , global_vect_decls = mkVarEnv vects - , global_scalar_vars = vectInfoScalarVars info `extendVarSetList` scalars - , global_scalar_tycons = vectInfoScalarTyCons info - , global_novect_vars = mkVarSet novects - , global_exported_vars = emptyVarEnv - , global_tycons = mapNameEnv snd $ vectInfoTyCon info - , global_datacons = mapNameEnv snd $ vectInfoDataCon info - , global_pa_funs = mapNameEnv snd $ vectInfoPADFun info - , global_pr_funs = emptyNameEnv - , global_boxed_tycons = emptyNameEnv - , global_inst_env = instEnvs - , global_fam_inst_env = famInstEnvs - , global_bindings = [] + { global_vars = mapVarEnv snd $ vectInfoVar info + , global_vect_decls = mkVarEnv vects + , global_scalar_vars = vectInfoScalarVars info `extendVarSetList` scalar_vars + , global_scalar_tycons = vectInfoScalarTyCons info `addListToNameSet` scalar_tycons + , global_novect_vars = mkVarSet novects + , global_exported_vars = emptyVarEnv + , global_tycons = mapNameEnv snd $ vectInfoTyCon info + , global_datacons = mapNameEnv snd $ vectInfoDataCon info + , global_pa_funs = mapNameEnv snd $ vectInfoPADFun info + , global_pr_funs = emptyNameEnv + , global_boxed_tycons = emptyNameEnv + , global_inst_env = instEnvs + , global_fam_inst_env = famInstEnvs + , global_bindings = [] } where - vects = [(var, (varType var, exp)) | Vect var (Just exp) <- vectDecls] - scalars = [var | Vect var Nothing <- vectDecls] - novects = [var | NoVect var <- vectDecls] + vects = [(var, (varType var, exp)) | Vect var (Just exp) <- vectDecls] + scalar_vars = [var | Vect var Nothing <- vectDecls] + novects = [var | NoVect var <- vectDecls] + scalar_tycons = [tyConName tycon | VectType tycon Nothing <- vectDecls] -- Operators on Global Environments ------------------------------------------- @@ -214,9 +216,9 @@ modVectInfo :: GlobalEnv -> TypeEnv -> VectInfo -> VectInfo modVectInfo env tyenv info = info { vectInfoVar = global_exported_vars env - , vectInfoTyCon = mk_env typeEnvTyCons global_tycons + , vectInfoTyCon = mk_env typeEnvTyCons global_tycons , vectInfoDataCon = mk_env typeEnvDataCons global_datacons - , vectInfoPADFun = mk_env typeEnvTyCons global_pa_funs + , vectInfoPADFun = mk_env typeEnvTyCons global_pa_funs , vectInfoScalarVars = global_scalar_vars env `minusVarSet` vectInfoScalarVars info , vectInfoScalarTyCons = global_scalar_tycons env `minusNameSet` vectInfoScalarTyCons info } diff --git a/compiler/vectorise/Vectorise/Exp.hs b/compiler/vectorise/Vectorise/Exp.hs index 6d6a473b44..2b7accc646 100644 --- a/compiler/vectorise/Vectorise/Exp.hs +++ b/compiler/vectorise/Vectorise/Exp.hs @@ -26,6 +26,7 @@ import CoreFVs import DataCon import TyCon import Type +import NameSet import Var import VarEnv import VarSet @@ -42,11 +43,11 @@ import Data.List -- | Vectorise a polymorphic expression. -- -vectPolyExpr :: Bool -- ^ When vectorising the RHS of a binding, whether that - -- binding is a loop breaker. - -> [Var] - -> CoreExprWithFVs - -> VM (Inline, Bool, VExpr) +vectPolyExpr :: Bool -- ^ When vectorising the RHS of a binding, whether that + -- binding is a loop breaker. + -> [Var] + -> CoreExprWithFVs + -> VM (Inline, Bool, VExpr) vectPolyExpr loop_breaker recFns (_, AnnNote note expr) = do (inline, isScalarFn, expr') <- vectPolyExpr loop_breaker recFns expr return (inline, isScalarFn, vNote note expr') @@ -194,26 +195,24 @@ vectScalarFun :: Bool -- ^ Was the function marked as scalar by the user? -> CoreExpr -- ^ Expression to be vectorised -> VM VExpr vectScalarFun forceScalar recFns expr - = do { gscalars <- globalScalars - ; let scalars = gscalars `extendVarSetList` recFns + = do { gscalarVars <- globalScalarVars + ; scalarTyCons <- globalScalarTyCons + ; let scalarVars = gscalarVars `extendVarSetList` recFns (arg_tys, res_ty) = splitFunTys (exprType expr) ; MASSERT( not $ null arg_tys ) - ; onlyIfV (forceScalar -- user asserts the functions is scalar + ; onlyIfV (forceScalar -- user asserts the functions is scalar || - all is_prim_ty arg_tys -- check whether the function is scalar - && is_prim_ty res_ty - && is_scalar scalars expr - && uses scalars expr) + all (is_scalar_ty scalarTyCons) arg_tys -- check whether the function is scalar + && is_scalar_ty scalarTyCons res_ty + && is_scalar scalarVars (is_scalar_ty scalarTyCons) expr + && uses scalarVars expr) $ mkScalarFun arg_tys res_ty expr } where - -- FIXME: This is woefully insufficient!!! We need a scalar pragma for types!!! - is_prim_ty ty - | Just (tycon, []) <- splitTyConApp_maybe ty - = tycon == intTyCon - || tycon == floatTyCon - || tycon == doubleTyCon - | otherwise = False + is_scalar_ty scalarTyCons ty + | Just (tycon, _) <- splitTyConApp_maybe ty + = tyConName tycon `elemNameSet` scalarTyCons + | otherwise = False -- Checks whether an expression contain a non-scalar subexpression. -- @@ -223,40 +222,45 @@ vectScalarFun forceScalar recFns expr -- them to the list of scalar variables) and then check them. If one of them turns out not to -- be scalar, the entire group is regarded as not being scalar. -- - -- FIXME: Currently, doesn't regard external (non-data constructor) variable and anonymous - -- data constructor as scalar. Should be changed once scalar types are passed - -- through VectInfo. + -- The second argument is a predicate that checks whether a type is scalar. -- - is_scalar :: VarSet -> CoreExpr -> Bool - is_scalar scalars (Var v) = v `elemVarSet` scalars - is_scalar _scalars (Lit _) = True - is_scalar scalars e@(App e1 e2) - | maybe_parr_ty (exprType e) = False - | otherwise = is_scalar scalars e1 && is_scalar scalars e2 - is_scalar scalars (Lam var body) - | maybe_parr_ty (varType var) = False - | otherwise = is_scalar (scalars `extendVarSet` var) body - is_scalar scalars (Let bind body) = bindsAreScalar && is_scalar scalars' body + is_scalar :: VarSet -> (Type -> Bool) -> CoreExpr -> Bool + is_scalar scalars _isScalarTC (Var v) = v `elemVarSet` scalars + is_scalar _scalars _isScalarTC (Lit _) = True + is_scalar scalars isScalarTC e@(App e1 e2) + | maybe_parr_ty (exprType e) = False + | otherwise = is_scalar scalars isScalarTC e1 && + is_scalar scalars isScalarTC e2 + is_scalar scalars isScalarTC (Lam var body) + | maybe_parr_ty (varType var) = False + | otherwise = is_scalar (scalars `extendVarSet` var) + isScalarTC body + is_scalar scalars isScalarTC (Let bind body) = bindsAreScalar && + is_scalar scalars' isScalarTC body where - (bindsAreScalar, scalars') = is_scalar_bind scalars bind - is_scalar scalars (Case e var ty alts) - | is_prim_ty ty = is_scalar scalars' e && all (is_scalar_alt scalars') alts + (bindsAreScalar, scalars') = is_scalar_bind scalars isScalarTC bind + is_scalar scalars isScalarTC (Case e var ty alts) + | isScalarTC ty = is_scalar scalars' isScalarTC e && + all (is_scalar_alt scalars' isScalarTC) alts | otherwise = False where scalars' = scalars `extendVarSet` var - is_scalar scalars (Cast e _coe) = is_scalar scalars e - is_scalar scalars (Note _ e ) = is_scalar scalars e - is_scalar _scalars (Type {}) = True - is_scalar _scalars (Coercion {}) = True + is_scalar scalars isScalarTC (Cast e _coe) = is_scalar scalars isScalarTC e + is_scalar scalars isScalarTC (Note _ e ) = is_scalar scalars isScalarTC e + is_scalar _scalars _isScalarTC (Type {}) = True + is_scalar _scalars _isScalarTC (Coercion {}) = True -- Result: (, scalars ++ variables bound in this group) - is_scalar_bind scalars (NonRec var e) = (is_scalar scalars e, scalars `extendVarSet` var) - is_scalar_bind scalars (Rec bnds) = (all (is_scalar scalars') es, scalars') + is_scalar_bind scalars isScalarTCs (NonRec var e) = (is_scalar scalars isScalarTCs e, + scalars `extendVarSet` var) + is_scalar_bind scalars isScalarTCs (Rec bnds) = (all (is_scalar scalars' isScalarTCs) es, + scalars') where (vars, es) = unzip bnds scalars' = scalars `extendVarSetList` vars - is_scalar_alt scalars (_, vars, e) = is_scalar (scalars `extendVarSetList ` vars) e + is_scalar_alt scalars isScalarTCs (_, vars, e) = is_scalar (scalars `extendVarSetList ` vars) + isScalarTCs e -- Checks whether the type might be a parallel array type. In particular, if the outermost -- constructor is a type family, we conservatively assume that it may be a parallel array type. diff --git a/compiler/vectorise/Vectorise/Monad/Global.hs b/compiler/vectorise/Vectorise/Monad/Global.hs index e471ebbc03..96448fb26a 100644 --- a/compiler/vectorise/Vectorise/Monad/Global.hs +++ b/compiler/vectorise/Vectorise/Monad/Global.hs @@ -1,3 +1,4 @@ +-- Operations on the global state of the vectorisation monad. module Vectorise.Monad.Global ( readGEnv, @@ -11,12 +12,11 @@ module Vectorise.Monad.Global ( lookupVectDecl, noVectDecl, -- * Scalars - globalScalars, isGlobalScalar, + globalScalarVars, isGlobalScalar, globalScalarTyCons, -- * TyCons - lookupTyCon, - lookupBoxedTyCon, - defTyCon, + lookupTyCon, lookupBoxedTyCon, + defTyCon, globalVectTyCons, -- * Datacons lookupDataCon, @@ -24,7 +24,6 @@ module Vectorise.Monad.Global ( -- * PA Dictionaries lookupTyConPA, - defTyConPA, defTyConPAs, -- * PR Dictionaries @@ -39,6 +38,7 @@ import Type import TyCon import DataCon import NameEnv +import NameSet import Var import VarEnv import VarSet @@ -49,17 +49,17 @@ import VarSet -- |Project something from the global environment. -- readGEnv :: (GlobalEnv -> a) -> VM a -readGEnv f = VM $ \_ genv lenv -> return (Yes genv lenv (f genv)) +readGEnv f = VM $ \_ genv lenv -> return (Yes genv lenv (f genv)) -- |Set the value of the global environment. -- setGEnv :: GlobalEnv -> VM () -setGEnv genv = VM $ \_ _ lenv -> return (Yes genv lenv ()) +setGEnv genv = VM $ \_ _ lenv -> return (Yes genv lenv ()) -- |Update the global environment using the provided function. -- updGEnv :: (GlobalEnv -> GlobalEnv) -> VM () -updGEnv f = VM $ \_ genv lenv -> return (Yes (f genv) lenv ()) +updGEnv f = VM $ \_ genv lenv -> return (Yes (f genv) lenv ()) -- Vars ----------------------------------------------------------------------- @@ -93,13 +93,19 @@ noVectDecl var = readGEnv $ \env -> elemVarSet var (global_novect_vars env) -- |Get the set of global scalar variables. -- -globalScalars :: VM VarSet -globalScalars = readGEnv global_scalar_vars +globalScalarVars :: VM VarSet +globalScalarVars = readGEnv global_scalar_vars -- |Check whether a given variable is in the set of global scalar variables. -- isGlobalScalar :: Var -> VM Bool -isGlobalScalar var = readGEnv $ \env -> elemVarSet var (global_scalar_vars env) +isGlobalScalar var = readGEnv $ \env -> var `elemVarSet` global_scalar_vars env + +-- |Get the set of global scalar type constructors including both those scalar type constructors +-- declared in an imported module and those declared in the current module. +-- +globalScalarTyCons :: VM NameSet +globalScalarTyCons = readGEnv global_scalar_tycons -- TyCons --------------------------------------------------------------------- @@ -110,25 +116,32 @@ lookupTyCon :: TyCon -> VM (Maybe TyCon) lookupTyCon tc | isUnLiftedTyCon tc || isTupleTyCon tc = return (Just tc) - | otherwise = readGEnv $ \env -> lookupNameEnv (global_tycons env) (tyConName tc) --- | Lookup the vectorised version of a boxed `TyCon` from the global environment. +-- |Lookup the vectorised version of a boxed `TyCon` from the global environment. +-- lookupBoxedTyCon :: TyCon -> VM (Maybe TyCon) lookupBoxedTyCon tc - = readGEnv $ \env -> lookupNameEnv (global_boxed_tycons env) - (tyConName tc) + = readGEnv $ \env -> lookupNameEnv (global_boxed_tycons env) + (tyConName tc) --- | Add a mapping between plain and vectorised `TyCon`s to the global environment. +-- |Add a mapping between plain and vectorised `TyCon`s to the global environment. +-- defTyCon :: TyCon -> TyCon -> VM () defTyCon tc tc' = updGEnv $ \env -> env { global_tycons = extendNameEnv (global_tycons env) (tyConName tc) tc' } +-- |Get the set of all vectorised type constructors. +-- +globalVectTyCons :: VM (NameEnv TyCon) +globalVectTyCons = readGEnv global_tycons + -- DataCons ------------------------------------------------------------------- --- | Lookup the vectorised version of a `DataCon` from the global environment. +-- |Lookup the vectorised version of a `DataCon` from the global environment. +-- lookupDataCon :: DataCon -> VM (Maybe DataCon) lookupDataCon dc | isTupleTyCon (dataConTyCon dc) @@ -137,27 +150,24 @@ lookupDataCon dc | otherwise = readGEnv $ \env -> lookupNameEnv (global_datacons env) (dataConName dc) - --- | Add the mapping between plain and vectorised `DataCon`s to the global environment. +-- |Add the mapping between plain and vectorised `DataCon`s to the global environment. +-- defDataCon :: DataCon -> DataCon -> VM () defDataCon dc dc' = updGEnv $ \env -> env { global_datacons = extendNameEnv (global_datacons env) (dataConName dc) dc' } --- PA dictionaries ------------------------------------------------------------ --- | Lookup a PA `TyCon` from the global environment. +-- 'PA' dictionaries ------------------------------------------------------------ + +-- |Lookup the 'PA' dfun of a vectorised type constructor in the global environment. +-- lookupTyConPA :: TyCon -> VM (Maybe Var) lookupTyConPA tc - = readGEnv $ \env -> lookupNameEnv (global_pa_funs env) (tyConName tc) - + = readGEnv $ \env -> lookupNameEnv (global_pa_funs env) (tyConName tc) --- | Add a mapping between a PA TyCon and is vectorised version to the global environment. -defTyConPA :: TyCon -> Var -> VM () -defTyConPA tc pa = updGEnv $ \env -> - env { global_pa_funs = extendNameEnv (global_pa_funs env) (tyConName tc) pa } - - --- | Add several mapping between PA TyCons and their vectorised versions to the global environment. +-- |Associate vectorised type constructors with the dfun of their 'PA' instances in the global +-- environment. +-- defTyConPAs :: [(TyCon, Var)] -> VM () defTyConPAs ps = updGEnv $ \env -> env { global_pa_funs = extendNameEnvList (global_pa_funs env) @@ -165,6 +175,7 @@ defTyConPAs ps = updGEnv $ \env -> -- PR Dictionaries ------------------------------------------------------------ + lookupTyConPR :: TyCon -> VM (Maybe Var) lookupTyConPR tc = readGEnv $ \env -> lookupNameEnv (global_pr_funs env) (tyConName tc) diff --git a/compiler/vectorise/Vectorise/Type/Classify.hs b/compiler/vectorise/Vectorise/Type/Classify.hs index 79cd0357c4..283af8175d 100644 --- a/compiler/vectorise/Vectorise/Type/Classify.hs +++ b/compiler/vectorise/Vectorise/Type/Classify.hs @@ -1,9 +1,22 @@ +-- Extract from a list of type constructors those (1) which need to be vectorised and (2) those +-- that could be, but need not be vectorised (as a scalar representation is sufficient and more +-- efficient). The type constructors that cannot be vectorised will be dropped. +-- +-- A type constructor will only be vectorised if it is +-- +-- (1) a data type constructor, with vanilla data constructors (i.e., data constructors admitted by +-- Haskell 98) and +-- (2) at least one of the type constructors that appears in its definition is also vectorised. +-- +-- If (1) is met, but not (2), the type constructor may appear in vectorised code, but there is no +-- need to vectorise that type constructor itself. This holds, for example, for all enumeration +-- types. As '([::])' is being vectorised, any type constructor whose definition involves +-- '([::])', either directly or indirectly, will be vectorised. + +module Vectorise.Type.Classify ( + classifyTyCons +) where -module Vectorise.Type.Classify - ( TyConGroup - , classifyTyCons - , tyConGroups) -where import UniqSet import UniqFM import DataCon @@ -13,31 +26,30 @@ import Type import Digraph import Outputable -type TyConGroup = ([TyCon], UniqSet TyCon) --- | Split the given tycons into two sets depending on whether they have to be --- converted (first list) or not (second list). The first argument contains --- information about the conversion status of external tycons: +-- |From a list of type constructors, extract those thatcan be vectorised, returning them in two +-- sets, where the first result list /must be/ vectorised and the second result list /need not be/ +-- vectroised. + +-- The first argument determines the /conversion status/ of external type constructors as follows: -- --- * tycons which have converted versions are mapped to True --- * tycons which are not changed by vectorisation are mapped to False --- * tycons which can't be converted are not elements of the map +-- * tycons which have converted versions are mapped to 'True' +-- * tycons which are not changed by vectorisation are mapped to 'False' +-- * tycons which can't be converted are not elements of the map -- -classifyTyCons - :: UniqFM Bool - -> [TyConGroup] - -> ([TyCon], [TyCon]) - -classifyTyCons = classify [] [] +classifyTyCons :: UniqFM Bool -- ^type constructor conversion status + -> [TyCon] -- ^type constructors that need to be classified + -> ([TyCon], [TyCon]) -- ^tycons to be converted & not to be converted +classifyTyCons convStatus tcs = classify [] [] convStatus (tyConGroups tcs) where classify conv keep _ [] = (conv, keep) classify conv keep cs ((tcs, ds) : rs) | can_convert && must_convert - = classify (tcs ++ conv) keep (cs `addListToUFM` [(tc,True) | tc <- tcs]) rs + = classify (tcs ++ conv) keep (cs `addListToUFM` [(tc, True) | tc <- tcs]) rs | can_convert - = classify conv (tcs ++ keep) (cs `addListToUFM` [(tc,False) | tc <- tcs]) rs + = classify conv (tcs ++ keep) (cs `addListToUFM` [(tc, False) | tc <- tcs]) rs | otherwise - = classify conv keep cs rs + = classify conv keep cs rs where refs = ds `delListFromUniqSet` tcs @@ -46,8 +58,12 @@ classifyTyCons = classify [] [] convertable tc = isDataTyCon tc && all isVanillaDataCon (tyConDataCons tc) +-- Used to group type constructors into mutually dependent groups. +-- +type TyConGroup = ([TyCon], UniqSet TyCon) --- | Compute mutually recursive groups of tycons in topological order +-- Compute mutually recursive groups of tycons in topological order. +-- tyConGroups :: [TyCon] -> [TyConGroup] tyConGroups tcs = map mk_grp (stronglyConnCompFromEdgedVertices edges) where @@ -59,19 +75,18 @@ tyConGroups tcs = map mk_grp (stronglyConnCompFromEdgedVertices edges) where (tcs, dss) = unzip els - --- | Collect the set of TyCons used by the representation of some data type. +-- |Collect the set of TyCons used by the representation of some data type. +-- tyConsOfTyCon :: TyCon -> UniqSet TyCon -tyConsOfTyCon - = tyConsOfTypes . concatMap dataConRepArgTys . tyConDataCons - +tyConsOfTyCon = tyConsOfTypes . concatMap dataConRepArgTys . tyConDataCons --- | Collect the set of TyCons that occur in these types. +-- |Collect the set of TyCons that occur in these types. +-- tyConsOfTypes :: [Type] -> UniqSet TyCon tyConsOfTypes = unionManyUniqSets . map tyConsOfType - --- | Collect the set of TyCons that occur in this type. +-- |Collect the set of TyCons that occur in this type. +-- tyConsOfType :: Type -> UniqSet TyCon tyConsOfType ty | Just ty' <- coreView ty = tyConsOfType ty' @@ -88,4 +103,3 @@ tyConsOfType (FunTy a b) = (tyConsOfType a `unionUniqSets` tyConsOfType b) `addOneToUniqSet` funTyCon tyConsOfType (ForAllTy _ ty) = tyConsOfType ty tyConsOfType other = pprPanic "ClosureConv.tyConsOfType" $ ppr other - diff --git a/compiler/vectorise/Vectorise/Type/Env.hs b/compiler/vectorise/Vectorise/Type/Env.hs index 4910464709..fcc6300022 100644 --- a/compiler/vectorise/Vectorise/Type/Env.hs +++ b/compiler/vectorise/Vectorise/Type/Env.hs @@ -1,7 +1,13 @@ {-# OPTIONS_GHC -XNoMonoLocalBinds #-} +-- Vectorise a modules type environment, the structure containing all type things defined in a +-- module. +-- +-- This extends the type environment with vectorised variants of data types and produces value +-- bindings for worker functions and the like. + module Vectorise.Type.Env ( - vectTypeEnv, + vectTypeEnv, ) where import Vectorise.Env @@ -28,9 +34,8 @@ import OccName import Id import MkId import NameEnv +import NameSet -import Unique -import UniqFM import Util import Outputable import FastString @@ -39,87 +44,145 @@ import Control.Monad import Data.List --- | Vectorise a type environment. --- The type environment contains all the type things defined in a module. +-- Note [Pragmas to vectorise tycons] +-- ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ +-- +-- VECTORISE pragmas for type constructors cover three different flavours of vectorising data type +-- constructors: +-- +-- (1) Data type constructor 'T' that may be used in vectorised code, where 'T' represents itself, +-- but the representation of 'T' is opaque in vectorised code. +-- +-- An example is the treatment of Int'. 'Int's can be used in vectorised code and remain +-- unchanged by vectorisation. However, the representation of 'Int' by the 'I#' data +-- constructor wrapping an 'Int#' is not exposed in vectorised code. Instead, computations +-- involving the representation need to be confined to scalar code. -- -vectTypeEnv :: TypeEnv +-- 'PData' and 'PRepr' instances need to be explicitly supplied for 'T' (they are not generated +-- by the vectoriser). +-- +-- Type constructors declared with {-# VECTORISE SCALAR type T #-} are treated in this manner. +-- (The vectoriser never treats a type constructor automatically in this manner.) +-- +-- (2) Data type constructor 'T' that together with its constructors 'Cn' may be used in vectorised +-- code, where 'T' and the 'Cn' represent themselves in vectorised code. +-- +-- An example is the treatment of 'Bool'. 'Bool' together with 'False' and 'True' may appear in +-- vectorised code and they remain unchanged by vectorisation. (There is no need for a special +-- representation as the values cannot embed any arrays.) + +-- 'PData' and 'PRepr' instances are automatically generated by the vectoriser. +-- +-- Type constructors declared with {-# VECTORISE type T #-} are treated in this manner. +-- (This is the same treatment that type constructors receive that the vectoriser deems fit for +-- use in vectorised code, but for which no special vectorised variant needs to be generated.) +-- +-- (3) [NOT IMPLEMENTED YET] +-- Data type constructor 'T' that together with its constructors 'Cn' may be used in vectorised +-- code, where 'T' is represented by 'Tv' and the workers of the 'Cn' are represented 'vCn' in +-- vectorised code. +-- +-- ??Example?? +-- +-- 'PData' and 'PRepr' instances are automatically generated by the vectoriser. +-- +-- ??How declared?? + +-- |Vectorise a type environment. +-- +vectTypeEnv :: TypeEnv -- Original type environment + -> [CoreVect] -- All 'VECTORISE [SCALAR] type' declarations in this module -> VM ( TypeEnv -- Vectorised type environment. , [FamInst] -- New type family instances. , [(Var, CoreExpr)]) -- New top level bindings. -vectTypeEnv env - = do - traceVt "** vectTypeEnv" $ ppr env - - cs <- readGEnv $ mk_map . global_tycons - - -- Split the list of TyCons into the ones we have to vectorise vs the - -- ones we can pass through unchanged. We also pass through algebraic - -- types that use non Haskell98 features, as we don't handle those. - let tycons = typeEnvTyCons env - groups = tyConGroups tycons - - let (conv_tcs, keep_tcs) = classifyTyCons cs groups - orig_tcs = keep_tcs ++ conv_tcs - keep_dcs = concatMap tyConDataCons keep_tcs - - -- Just use the unvectorised versions of these constructors in vectorised code. - zipWithM_ defTyCon keep_tcs keep_tcs - zipWithM_ defDataCon keep_dcs keep_dcs - - -- Vectorise all the declarations. - new_tcs <- vectTyConDecls conv_tcs - - -- We don't need to make new representation types for dictionary - -- constructors. The constructors are always fully applied, and we don't - -- need to lift them to arrays as a dictionary of a particular type - -- always has the same value. - let vect_tcs = filter (not . isClassTyCon) - $ keep_tcs ++ new_tcs - - reprs <- mapM tyConRepr vect_tcs - repr_tcs <- zipWith3M buildPReprTyCon orig_tcs vect_tcs reprs - pdata_tcs <- zipWith3M buildPDataTyCon orig_tcs vect_tcs reprs - updGEnv $ extendFamEnv - $ map mkLocalFamInst - $ repr_tcs ++ pdata_tcs - - -- Create PRepr and PData instances for the vectorised types. - -- We get back the binds for the instance functions, - -- and some new type constructors for the representation types. - (_, binds, inst_tcs) <- fixV $ \ ~(dfuns', _, _) -> - do - defTyConPAs (zipLazy vect_tcs dfuns') - reprs <- mapM tyConRepr vect_tcs - - dfuns <- sequence - $ zipWith5 buildTyConBindings - orig_tcs - vect_tcs - repr_tcs - pdata_tcs - reprs - - binds <- takeHoisted - return (dfuns, binds, repr_tcs ++ pdata_tcs) - - -- The new type constructors are the vectorised versions of the originals, - -- plus the new type constructors that we use for the representations. - let all_new_tcs = new_tcs ++ inst_tcs - - let new_env = extendTypeEnvList env - $ map ATyCon all_new_tcs - ++ [ADataCon dc | tc <- all_new_tcs - , dc <- tyConDataCons tc] - - return (new_env, map mkLocalFamInst inst_tcs, binds) - - where - mk_map env = listToUFM_Directly [(u, getUnique n /= u) | (u,n) <- nameEnvUniqueElts env] - -buildTyConBindings :: TyCon -> TyCon -> TyCon -> TyCon -> SumRepr -> VM Var -buildTyConBindings orig_tc vect_tc prepr_tc pdata_tc repr - = do vectDataConWorkers orig_tc vect_tc pdata_tc - buildPADict vect_tc prepr_tc pdata_tc repr +vectTypeEnv env vectTypeDecls + = do { traceVt "** vectTypeEnv" $ ppr env + + -- Build a map containing all vectorised type constructor. If they are scalar, they are + -- mapped to 'False' (vectorised type constructor == original type constructor). + ; allScalarTyConNames <- globalScalarTyCons -- covers both current and imported modules + ; vectTyCons <- globalVectTyCons + ; let vectTyConBase = mapNameEnv (const True) vectTyCons -- by default fully vectorised + vectTyConFlavour = foldNameSet (\n env -> extendNameEnv env n False) vectTyConBase + allScalarTyConNames + + -- Split the list of 'TyCons' into the ones (1) that we must vectorise and those (2) + -- that we could, but don't need to vectorise. Type constructors that are not data + -- type constructors or use non-Haskell98 features are being dropped. They may not + -- appear in vectorised code. (We also drop the local type constructors appearing in a + -- VECTORISE SCALAR pragma, as they are being handled separately.) + ; let localScalarTyCons = [tycon | VectType tycon Nothing <- vectTypeDecls] + localScalarTyConNames = mkNameSet (map tyConName localScalarTyCons) + notLocalScalarTyCon tc = not $ (tyConName tc) `elemNameSet` localScalarTyConNames + + maybeVectoriseTyCons = filter notLocalScalarTyCon (typeEnvTyCons env) + (conv_tcs, keep_tcs) = classifyTyCons vectTyConFlavour maybeVectoriseTyCons + orig_tcs = keep_tcs ++ conv_tcs + keep_dcs = concatMap tyConDataCons keep_tcs + + keep_and_scalar_tcs = keep_tcs ++ localScalarTyCons + + -- Of those type constructors that we don't need to vectorise, we use the original + -- representation in both unvectorised and vectorised code. For those declared VECTORISE + -- SCALAR, we ignore their represention — see "Note [Pragmas to vectorise tycons]". + ; zipWithM_ defTyCon keep_and_scalar_tcs keep_and_scalar_tcs + ; zipWithM_ defDataCon keep_dcs keep_dcs + + -- Vectorise all the data type declarations that we can and must vectorise. + ; new_tcs <- vectTyConDecls conv_tcs + + -- We don't need new representation types for dictionary constructors. The constructors + -- are always fully applied, and we don't need to lift them to arrays as a dictionary + -- of a particular type always has the same value. + ; let vect_tcs = filter (not . isClassTyCon) + $ keep_tcs ++ new_tcs + + -- Build 'PRepr' and 'PData' instance type constructors and family instances for all + -- type constructors with vectorised representations. + ; reprs <- mapM tyConRepr vect_tcs + ; repr_tcs <- zipWith3M buildPReprTyCon orig_tcs vect_tcs reprs + ; pdata_tcs <- zipWith3M buildPDataTyCon orig_tcs vect_tcs reprs + ; let inst_tcs = repr_tcs ++ pdata_tcs + fam_insts = map mkLocalFamInst inst_tcs + ; updGEnv $ extendFamEnv fam_insts + + -- Generate dfuns for the 'PA' instances of the vectorised type constructors and + -- associate the type constructors with their dfuns in the global environment. We get + -- back the dfun bindings (which we will subsequently inject into the modules toplevel). + ; (_, binds) <- fixV $ \ ~(dfuns, _) -> + do { defTyConPAs (zipLazy vect_tcs dfuns) + ; dfuns <- sequence + $ zipWith4 buildTyConBindings + orig_tcs + vect_tcs + repr_tcs + pdata_tcs + + ; binds <- takeHoisted + ; return (dfuns, binds) + } + + -- We add to the type environment: (1) the vectorised type constructors, (2) their + -- 'PRepr' & 'PData' instance constructors, and (3) the data constructors of the fomer + -- two. + ; let all_new_tcs = new_tcs ++ inst_tcs + new_env = extendTypeEnvList env + $ map ATyCon all_new_tcs ++ + [ADataCon dc | tc <- all_new_tcs + , dc <- tyConDataCons tc] + + ; return (new_env, fam_insts, binds) + } + + +-- Helpers ------------------- + +buildTyConBindings :: TyCon -> TyCon -> TyCon -> TyCon -> VM Var +buildTyConBindings orig_tc vect_tc prepr_tc pdata_tc + = do { vectDataConWorkers orig_tc vect_tc pdata_tc + ; repr <- tyConRepr vect_tc + ; buildPADict vect_tc prepr_tc pdata_tc repr + } vectDataConWorkers :: TyCon -> TyCon -> TyCon -> VM () vectDataConWorkers orig_tc vect_tc arr_tc @@ -187,4 +250,3 @@ vectDataConWorkers orig_tc vect_tc arr_tc return (vect_worker, body) where orig_worker = dataConWorkId data_con - diff --git a/compiler/vectorise/Vectorise/Type/Repr.hs b/compiler/vectorise/Vectorise/Type/Repr.hs index bb300ca863..2fd788432c 100644 --- a/compiler/vectorise/Vectorise/Type/Repr.hs +++ b/compiler/vectorise/Vectorise/Type/Repr.hs @@ -1,17 +1,10 @@ +-- |Compute the representation type for data type constructors. + +module Vectorise.Type.Repr ( + CompRepr (..), ProdRepr (..), ConRepr (..), SumRepr (..), + tyConRepr, sumReprType, conReprType, prodReprType, compReprType, compOrigType +) where --- | Representation of Algebraic Data Types. -module Vectorise.Type.Repr - ( CompRepr (..) - , ProdRepr (..) - , ConRepr (..) - , SumRepr (..) - , tyConRepr - , sumReprType - , conReprType - , prodReprType - , compReprType - , compOrigType) -where import Vectorise.Utils import Vectorise.Monad import Vectorise.Builtins @@ -41,10 +34,12 @@ data SumRepr = EmptySum | Sum { repr_sum_tc :: TyCon -- representation sum tycon , repr_psum_tc :: TyCon -- PData representation tycon , repr_sel_ty :: Type -- type of selector - , repr_con_tys :: [Type] -- representation types of + , repr_con_tys :: [Type] -- representation types of , repr_cons :: [ConRepr] -- components } +-- |Determine the representation type of a data type constructor. +-- tyConRepr :: TyCon -> VM SumRepr tyConRepr tc = sum_repr (tyConDataCons tc) where @@ -102,9 +97,10 @@ prodReprType (Prod { repr_tup_tc = tup_tc, repr_comp_tys = tys }) compReprType :: CompRepr -> VM Type compReprType (Keep ty _) = return ty -compReprType (Wrap ty) = do - wrap_tc <- builtin wrapTyCon - return $ mkTyConApp wrap_tc [ty] +compReprType (Wrap ty) + = do { wrap_tc <- builtin wrapTyCon + ; return $ mkTyConApp wrap_tc [ty] + } compOrigType :: CompRepr -> Type compOrigType (Keep ty _) = ty -- GitLab