diff options
Diffstat (limited to 'compiler/vectorise')
| -rw-r--r-- | compiler/vectorise/Vectorise.hs | 24 | ||||
| -rw-r--r-- | compiler/vectorise/Vectorise/Builtins.hs | 13 | ||||
| -rw-r--r-- | compiler/vectorise/Vectorise/Builtins/Initialise.hs | 7 | ||||
| -rw-r--r-- | compiler/vectorise/Vectorise/Env.hs | 72 | ||||
| -rw-r--r-- | compiler/vectorise/Vectorise/Exp.hs | 88 | ||||
| -rw-r--r-- | compiler/vectorise/Vectorise/Monad/Global.hs | 71 | ||||
| -rw-r--r-- | compiler/vectorise/Vectorise/Type/Classify.hs | 76 | ||||
| -rw-r--r-- | compiler/vectorise/Vectorise/Type/Env.hs | 224 | ||||
| -rw-r--r-- | compiler/vectorise/Vectorise/Type/Repr.hs | 30 |
9 files changed, 349 insertions, 256 deletions
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: (<is this binding group scalar>, 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 |