Commit 60994016 authored by chak@cse.unsw.edu.au.'s avatar chak@cse.unsw.edu.au.

Added a pragma {-# NOVECTORISE f #-} that suppresses vectorisation of toplevel variable 'f'.

parent a8defd8a
......@@ -332,8 +332,9 @@ Also since rule_fn is a Name, not a Var, we have to use the grungy delUFM.
vectsFreeVars :: [CoreVect] -> VarSet
vectsFreeVars = foldr (unionVarSet . vectFreeVars) emptyVarSet
where
vectFreeVars (Vect _ Nothing) = noFVs
vectFreeVars (Vect _ (Just rhs)) = expr_fvs rhs isLocalId emptyVarSet
vectFreeVars (Vect _ Nothing) = noFVs
vectFreeVars (Vect _ (Just rhs)) = expr_fvs rhs isLocalId emptyVarSet
vectFreeVars (NoVect _) = noFVs
\end{code}
......
......@@ -714,8 +714,9 @@ 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 (Vect v Nothing) = Vect v Nothing
substVect subst (Vect v (Just rhs)) = Vect v (Just (simpleOptExprWith subst rhs))
substVect _subst (NoVect v) = NoVect v
------------------
substVarSet :: Subst -> VarSet -> VarSet
......
......@@ -417,14 +417,16 @@ Representation of desugared vectorisation declarations that are fed to the vecto
'ModGuts').
\begin{code}
data CoreVect = Vect Id (Maybe CoreExpr)
data CoreVect = Vect Id (Maybe CoreExpr)
| NoVect Id
\end{code}
%************************************************************************
%* *
Unfoldings
%* *
%* *
Unfoldings
%* *
%************************************************************************
The @Unfolding@ type is declared here to avoid numerous loops
......
......@@ -446,7 +446,7 @@ instance Outputable e => Outputable (DFunArg e) where
\end{code}
-----------------------------------------------------
-- Rules
-- Rules
-----------------------------------------------------
\begin{code}
......@@ -461,11 +461,23 @@ pprRule (BuiltinRule { ru_fn = fn, ru_name = name})
= ptext (sLit "Built in rule for") <+> ppr fn <> colon <+> doubleQuotes (ftext name)
pprRule (Rule { ru_name = name, ru_act = act, ru_fn = fn,
ru_bndrs = tpl_vars, ru_args = tpl_args,
ru_rhs = rhs })
ru_bndrs = tpl_vars, ru_args = tpl_args,
ru_rhs = rhs })
= hang (doubleQuotes (ftext name) <+> ppr act)
4 (sep [ptext (sLit "forall") <+> braces (sep (map pprTypedBinder tpl_vars)),
nest 2 (ppr fn <+> sep (map pprArg tpl_args)),
nest 2 (ptext (sLit "=") <+> pprCoreExpr rhs)
])
nest 2 (ppr fn <+> sep (map pprArg tpl_args)),
nest 2 (ptext (sLit "=") <+> pprCoreExpr rhs)
])
\end{code}
-----------------------------------------------------
-- Vectorisation declarations
-----------------------------------------------------
\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
\end{code}
......@@ -394,16 +394,11 @@ the rule is precisly to optimise them:
\begin{code}
dsVect :: LVectDecl Id -> DsM CoreVect
dsVect (L loc (HsVect v rhs))
dsVect (L loc (HsVect (L _ v) rhs))
= putSrcSpanDs loc $
do { rhs' <- fmapMaybeM dsLExpr rhs
; return $ Vect (unLoc v) rhs'
; return $ Vect v rhs'
}
-- dsVect (L loc (HsVect v Nothing))
-- = return $ Vect v Nothing
-- dsVect (L loc (HsVect v (Just rhs)))
-- = putSrcSpanDs loc $
-- do { rhs' <- dsLExpr rhs
-- ; return $ Vect v (Just rhs')
-- }
dsVect (L loc (HsNoVect (L _ v)))
= return $ NoVect v
\end{code}
......@@ -28,6 +28,7 @@ module HsDecls (
collectRuleBndrSigTys,
-- ** @VECTORISE@ declarations
VectDecl(..), LVectDecl,
lvectDeclName,
-- ** @default@ declarations
DefaultDecl(..), LDefaultDecl,
-- ** Top-level template haskell splice
......@@ -1005,10 +1006,11 @@ instance OutputableBndr name => Outputable (RuleBndr name) where
%* *
%************************************************************************
A vectorisation pragma
A vectorisation pragma, one of
{-# VECTORISE f = closure1 g (scalar_map g) #-} OR
{-# VECTORISE f = closure1 g (scalar_map g) #-}
{-# VECTORISE SCALAR f #-}
{-# NOVECTORISE f #-}
Note [Typechecked vectorisation pragmas]
~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
......@@ -1029,14 +1031,23 @@ data VectDecl name
= HsVect
(Located name)
(Maybe (LHsExpr name)) -- 'Nothing' => SCALAR declaration
| HsNoVect
(Located name)
deriving (Data, Typeable)
lvectDeclName :: LVectDecl name -> name
lvectDeclName (L _ (HsVect (L _ name) _)) = name
lvectDeclName (L _ (HsNoVect (L _ name))) = name
instance OutputableBndr name => Outputable (VectDecl name) where
ppr (HsVect v rhs)
ppr (HsVect v Nothing)
= sep [text "{-# VECTORISE SCALAR" <+> ppr v <+> text "#-}" ]
ppr (HsVect v (Just rhs))
= sep [text "{-# VECTORISE" <+> ppr v,
nest 4 (case rhs of
Nothing -> text "SCALAR #-}"
Just rhs -> pprExpr (unLoc rhs) <+> text "#-}") ]
nest 4 $
pprExpr (unLoc rhs) <+> text "#-}" ]
ppr (HsNoVect v)
= sep [text "{-# NOVECTORISE" <+> ppr v <+> text "#-}" ]
\end{code}
%************************************************************************
......
......@@ -483,6 +483,7 @@ data Token
| ITlanguage_prag
| ITvect_prag
| ITvect_scalar_prag
| ITnovect_prag
| ITdotdot -- reserved symbols
| ITcolon
......@@ -2281,7 +2282,8 @@ oneWordPrags = Map.fromList([("rules", rulePrag),
("core", token ITcore_prag),
("unpack", token ITunpack_prag),
("ann", token ITann_prag),
("vectorize", token ITvect_prag)])
("vectorize", token ITvect_prag),
("novectorize", token ITnovect_prag)])
twoWordPrags = Map.fromList([("inline conlike", token (ITinline_prag Inline ConLike)),
("notinline conlike", token (ITinline_prag NoInline ConLike)),
......@@ -2307,6 +2309,7 @@ clean_pragma prag = canon_ws (map toLower (unprefix prag))
"noinline" -> "notinline"
"specialise" -> "specialize"
"vectorise" -> "vectorize"
"novectorise" -> "novectorize"
"constructorlike" -> "conlike"
_ -> prag'
canon_ws s = unwords (map canonical (words s))
......
......@@ -252,21 +252,22 @@ incorrect.
'by' { L _ ITby } -- for list transform extension
'using' { L _ ITusing } -- for list transform extension
'{-# INLINE' { L _ (ITinline_prag _ _) }
'{-# SPECIALISE' { L _ ITspec_prag }
'{-# INLINE' { L _ (ITinline_prag _ _) }
'{-# SPECIALISE' { L _ ITspec_prag }
'{-# SPECIALISE_INLINE' { L _ (ITspec_inline_prag _) }
'{-# SOURCE' { L _ ITsource_prag }
'{-# RULES' { L _ ITrules_prag }
'{-# CORE' { L _ ITcore_prag } -- hdaume: annotated core
'{-# SCC' { L _ ITscc_prag }
'{-# GENERATED' { L _ ITgenerated_prag }
'{-# DEPRECATED' { L _ ITdeprecated_prag }
'{-# WARNING' { L _ ITwarning_prag }
'{-# UNPACK' { L _ ITunpack_prag }
'{-# ANN' { L _ ITann_prag }
'{-# SOURCE' { L _ ITsource_prag }
'{-# RULES' { L _ ITrules_prag }
'{-# CORE' { L _ ITcore_prag } -- hdaume: annotated core
'{-# SCC' { L _ ITscc_prag }
'{-# GENERATED' { L _ ITgenerated_prag }
'{-# DEPRECATED' { L _ ITdeprecated_prag }
'{-# WARNING' { L _ ITwarning_prag }
'{-# UNPACK' { L _ ITunpack_prag }
'{-# ANN' { L _ ITann_prag }
'{-# VECTORISE' { L _ ITvect_prag }
'{-# VECTORISE_SCALAR' { L _ ITvect_scalar_prag }
'#-}' { L _ ITclose_prag }
'{-# NOVECTORISE' { L _ ITnovect_prag }
'#-}' { L _ ITclose_prag }
'..' { L _ ITdotdot } -- reserved symbols
':' { L _ ITcolon }
......@@ -546,33 +547,34 @@ ops :: { Located [Located RdrName] }
-- Top-Level Declarations
topdecls :: { OrdList (LHsDecl RdrName) }
: topdecls ';' topdecl { $1 `appOL` $3 }
| topdecls ';' { $1 }
| topdecl { $1 }
: topdecls ';' topdecl { $1 `appOL` $3 }
| topdecls ';' { $1 }
| topdecl { $1 }
topdecl :: { OrdList (LHsDecl RdrName) }
: 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
unitOL (L (comb3 $1 $2 $3) (InstD (InstDecl $2 binds sigs ats)))}
: 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
unitOL (L (comb3 $1 $2 $3) (InstD (InstDecl $2 binds sigs ats)))}
| stand_alone_deriving { unitOL (LL (DerivD (unLoc $1))) }
| 'default' '(' comma_types0 ')' { unitOL (LL $ DefD (DefaultDecl $3)) }
| 'foreign' fdecl { unitOL (LL (unLoc $2)) }
| 'default' '(' comma_types0 ')' { unitOL (LL $ DefD (DefaultDecl $3)) }
| 'foreign' fdecl { unitOL (LL (unLoc $2)) }
| '{-# 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)) }
| annotation { unitOL $1 }
| decl { unLoc $1 }
-- Template Haskell Extension
-- The $(..) form is one possible form of infixexp
-- but we treat an arbitrary expression just as if
-- it had a $(..) wrapped around it
| infixexp { unitOL (LL $ mkTopSpliceDecl $1) }
| '{-# 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) }
| annotation { unitOL $1 }
| decl { unLoc $1 }
-- Template Haskell Extension
-- The $(..) form is one possible form of infixexp
-- but we treat an arbitrary expression just as if
-- it had a $(..) wrapped around it
| infixexp { unitOL (LL $ mkTopSpliceDecl $1) }
-- Type classes
--
......
......@@ -666,6 +666,10 @@ rnHsVectDecl (HsVect var (Just rhs))
; (rhs', fv_rhs) <- rnLExpr rhs
; return (HsVect var' (Just rhs'), fv_rhs `addOneFV` unLoc var')
}
rnHsVectDecl (HsNoVect var)
= do { var' <- wrapLocM lookupTopBndrRn var
; return (HsNoVect var', unitFV (unLoc var'))
}
\end{code}
%*********************************************************
......
......@@ -29,7 +29,7 @@ import FloatIn ( floatInwards )
import FloatOut ( floatOutwards )
import FamInstEnv
import Id
import BasicTypes ( CompilerPhase, isDefaultInlinePragma )
import BasicTypes
import VarSet
import VarEnv
import LiberateCase ( liberateCase )
......@@ -356,11 +356,18 @@ simplifyPgmIO pass@(CoreDoSimplify max_iterations mode)
-- space usage, especially with -O. JRS, 000620.
| let sz = coreBindsSize binds in sz == sz
= do {
-- Occurrence analysis
let { tagged_binds = {-# SCC "OccAnal" #-}
occurAnalysePgm active_rule rules [] binds } ;
Err.dumpIfSet_dyn dflags Opt_D_dump_occur_anal "Occurrence analysis"
(pprCoreBindings tagged_binds);
-- Occurrence analysis
let { -- During the 'InitialPhase' (i.e., before vectorisation), we need to make sure
-- that the right-hand sides of vectorisation declarations are taken into
-- account during occurence analysis.
maybeVects = case sm_phase mode of
InitialPhase -> mg_vect_decls guts
_ -> []
; tagged_binds = {-# SCC "OccAnal" #-}
occurAnalysePgm active_rule rules maybeVects binds
} ;
Err.dumpIfSet_dyn dflags Opt_D_dump_occur_anal "Occurrence analysis"
(pprCoreBindings tagged_binds);
-- Get any new rules, and extend the rule base
-- See Note [Overall plumbing for rules] in Rules.lhs
......
......@@ -591,7 +591,7 @@ impSpecErr name
tcVectDecls :: [LVectDecl Name] -> TcM ([LVectDecl TcId])
tcVectDecls decls
= do { decls' <- mapM (wrapLocM tcVect) decls
; let ids = [unLoc id | L _ (HsVect id _) <- decls']
; let ids = map lvectDeclName decls'
dups = findDupsEq (==) ids
; mapM_ reportVectDups dups
; traceTcConstraints "End of tcVectDecls"
......@@ -642,6 +642,11 @@ tcVect (HsVect name@(L loc _) (Just rhs))
-- to the vectoriser - see "Note [Typechecked vectorisation pragmas]" in HsDecls
; return $ HsVect (L loc id') (Just rhsWrapped)
}
tcVect (HsNoVect name)
= addErrCtxt (vectCtxt name) $
do { id <- wrapLocM tcLookupId name
; return $ HsNoVect id
}
vectCtxt :: Located Name -> SDoc
vectCtxt name = ptext (sLit "When checking the vectorisation declaration for") <+> ppr name
......
......@@ -1027,6 +1027,10 @@ zonkVect env (HsVect v (Just e))
; e' <- zonkLExpr env e
; return $ HsVect v' (Just e')
}
zonkVect env (HsNoVect v)
= do { v' <- wrapLocM (zonkIdBndr env) v
; return $ HsNoVect v'
}
\end{code}
%************************************************************************
......
{-# OPTIONS -fno-warn-missing-signatures -fno-warn-unused-do-bind #-}
module Vectorise ( vectorise )
where
......@@ -82,98 +81,124 @@ vectModule guts@(ModGuts { mg_types = types
}
}
-- | Try to vectorise a top-level binding.
-- If it doesn't vectorise then return it unharmed.
-- |Try to vectorise a top-level binding. If it doesn't vectorise then return it unharmed.
--
-- For example, for the 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 = ...
-- @
-- foo :: Int -> Int
-- foo = \x -> x + x
-- @
--
-- lfoo :: PData Void -> PData Int -> PData Int
-- lfoo = ...
-- @
-- we get
-- @
-- foo :: Int -> Int
-- foo = \x -> vfoo $: x
--
-- @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 :: Closure void vfoo lfoo
-- v_foo = closure vfoo lfoo void
--
-- vfoo :: Void -> Int -> Int
-- vfoo = ...
--
-- lfoo :: PData Void -> PData Int -> PData Int
-- lfoo = ...
-- @
--
-- @v_foo@ combines both of these into a `Closure` that also contains the
-- environment.
-- @vfoo@ is the "vectorised", or scalar, version that does the same as the original
-- function foo, but takes an explicit environment.
--
-- The original binding @foo@ is rewritten to call the vectorised version
-- present in the closure.
-- @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 surpressed 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.
-- 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 { -- Vectorise the right-hand side, create an appropriate top-level binding and add it to
-- the vectorisation map.
; (inline, isScalar, expr') <- vectTopRhs [] var expr
; var' <- vectTopBinder var inline expr'
; when isScalar $
addGlobalScalar var
-- 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
; hs <- takeHoisted
; return . Rec $ (var, cexpr) : (var', expr') : hs
}
`orElseV`
return b
= unlessNoVectDecl $
do { -- Vectorise the right-hand side, create an appropriate top-level binding and add it
-- to the vectorisation map.
; (inline, isScalar, expr') <- vectTopRhs [] var expr
; var' <- vectTopBinder var inline expr'
; when isScalar $
addGlobalScalar var
-- 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
; hs <- takeHoisted
; return . Rec $ (var, cexpr) : (var', expr') : hs
}
`orElseV`
return b
where
unlessNoVectDecl vectorise
= do { hasNoVectDecl <- noVectDecl var
; when hasNoVectDecl $
traceVt "NOVECTORISE" $ ppr var
; if hasNoVectDecl then return b else vectorise
}
vectTopBind b@(Rec bs)
= let (vars, exprs) = unzip bs
in
do { (vars', _, exprs', hs) <- fixV $
\ ~(_, inlines, rhss, _) ->
do { -- Vectorise the right-hand sides, create an appropriate top-level bindings and
-- add them to the vectorisation map.
; vars' <- sequence [vectTopBinder var inline rhs
| (var, ~(inline, rhs)) <- zipLazy vars (zip inlines rhss)]
; (inlines, areScalars, exprs') <- mapAndUnzip3M (uncurry $ vectTopRhs vars) bs
; hs <- takeHoisted
; if and areScalars
then -- (1) Entire recursive group is scalar
-- => add all variables to the global set of scalars
do { mapM addGlobalScalar vars
; return (vars', inlines, exprs', hs)
}
else -- (2) At least one binding is not scalar
-- => vectorise again with empty set of local scalars
do { (inlines, _, exprs') <- mapAndUnzip3M (uncurry $ vectTopRhs []) bs
; hs <- takeHoisted
; return (vars', inlines, exprs', hs)
}
}
-- 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.
; cexprs <- sequence $ zipWith3 tryConvert vars vars' exprs
; return . Rec $ zip vars cexprs ++ zip vars' exprs' ++ hs
}
`orElseV`
return b
= unlessSomeNoVectDecl $
do { (vars', _, exprs', hs) <- fixV $
\ ~(_, inlines, rhss, _) ->
do { -- Vectorise the right-hand sides, create an appropriate top-level bindings
-- and add them to the vectorisation map.
; vars' <- sequence [vectTopBinder var inline rhs
| (var, ~(inline, rhs)) <- zipLazy vars (zip inlines rhss)]
; (inlines, areScalars, exprs') <- mapAndUnzip3M (uncurry $ vectTopRhs vars) bs
; hs <- takeHoisted
; if and areScalars
then -- (1) Entire recursive group is scalar
-- => add all variables to the global set of scalars
do { mapM_ addGlobalScalar vars
; return (vars', inlines, exprs', hs)
}
else -- (2) At least one binding is not scalar
-- => vectorise again with empty set of local scalars
do { (inlines, _, exprs') <- mapAndUnzip3M (uncurry $ vectTopRhs []) bs
; hs <- takeHoisted
; return (vars', inlines, exprs', hs)
}
}
-- 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.
; cexprs <- sequence $ zipWith3 tryConvert vars vars' exprs
; return . Rec $ zip vars cexprs ++ zip vars' exprs' ++ hs
}
`orElseV`
return b
where
(vars, exprs) = unzip bs
unlessSomeNoVectDecl vectorise
= do { hasNoVectDecls <- mapM noVectDecl vars
; when (and hasNoVectDecls) $
traceVt "NOVECTORISE" $ ppr vars
; if and hasNoVectDecls
then return b -- all bindings have 'NOVECTORISE'
else if or hasNoVectDecls
then cantVectorise noVectoriseErr (ppr b) -- some (but not all) have 'NOVECTORISE'
else vectorise -- no binding has a 'NOVECTORISE' decl
}
noVectoriseErr = "NOVECTORISE must be used on all or no bindings of a recursive group"
-- | 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
-- 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.
......
......@@ -50,11 +50,11 @@ preludeVars (Modules { dph_Combinators = _dph_Combinators
, mk' dph_Prelude_Word8 "toInt" "toIntV"
]
++ vars_Ord dph_Prelude_Double
++ vars_Num dph_Prelude_Double
++ vars_Fractional dph_Prelude_Double
++ vars_Floating dph_Prelude_Double
++ vars_RealFrac dph_Prelude_Double
-- ++ vars_Ord dph_Prelude_Double
-- ++ vars_Num dph_Prelude_Double
-- ++ vars_Fractional dph_Prelude_Double
-- ++ vars_Floating dph_Prelude_Double
-- ++ vars_RealFrac dph_Prelude_Double
++
[ mk dph_Prelude_Bool (fsLit "andP") dph_Prelude_Bool (fsLit "andPA")
, mk dph_Prelude_Bool (fsLit "orP") dph_Prelude_Bool (fsLit "orPA")
......
......@@ -95,6 +95,10 @@ data GlobalEnv
, global_scalar_tycons :: NameSet
-- ^Type constructors whose values can only contain scalar data. Scalar code may only
-- operate on such data.
, 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)
-- ^Exported variables which have a vectorised version.
......@@ -134,6 +138,7 @@ initGlobalEnv info vectDecls instEnvs famInstEnvs
, 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
......@@ -147,6 +152,7 @@ initGlobalEnv info vectDecls instEnvs famInstEnvs
where
vects = [(var, (varType var, exp)) | Vect var (Just exp) <- vectDecls]
scalars = [var | Vect var Nothing <- vectDecls]
novects = [var | NoVect var <- vectDecls]
-- Operators on Global Environments -------------------------------------------
......
......@@ -81,6 +81,7 @@ initV hsc_env guts info thing_inside
; builtin_pas <- initBuiltinPAs builtins instEnvs
-- construct the initial global environment
; let thing_inside' = traceVt "VectDecls" (ppr (mg_vect_decls guts)) >> thing_inside
; let genv = extendImportedVarsEnv builtin_vars
. extendScalars builtin_scalars
. extendTyConsEnv builtin_tycons
......@@ -91,7 +92,7 @@ initV hsc_env guts info thing_inside
$ initGlobalEnv info (mg_vect_decls guts) instEnvs famInstEnvs
-- perform vectorisation
; r <- runVM thing_inside builtins genv emptyLocalEnv
; r <- runVM thing_inside' builtins genv emptyLocalEnv
; case r of
Yes genv _ x -> return $ Just (new_info genv, x)
No -> return Nothing
......
module Vectorise.Monad.Global (
readGEnv,
setGEnv,
updGEnv,
readGEnv,
setGEnv,
updGEnv,
-- * Vars
defGlobalVar,
-- * Vectorisation declarations
lookupVectDecl,
lookupVectDecl, noVectDecl,
-- * Scalars
globalScalars, isGlobalScalar,
-- * TyCons
lookupTyCon,
lookupBoxedTyCon,
defTyCon,
-- * Datacons
lookupDataCon,
defDataCon,
-- * PA Dictionaries
lookupTyConPA,
defTyConPA,
defTyConPAs,
-- * PR Dictionaries
lookupTyConPR
-- * TyCons
lookupTyCon,
lookupBoxedTyCon,
defTyCon,
-- * Datacons
lookupDataCon,
defDataCon,
-- * PA Dictionaries
lookupTyConPA,
defTyConPA,
defTyConPAs,
-- * PR Dictionaries
lookupTyConPR
) where
import Vectorise.Monad.Base
......@@ -45,23 +45,27 @@ import VarSet
-- Global Environment ---------------------------------------------------------
-- | Project something from the global environment.
-- |Project something from the global environment.
--
readGEnv :: (GlobalEnv -> a) -> VM a
readGEnv f = VM $ \_ genv lenv -> return (Yes genv lenv (f genv))
-- | Set the value of the global environment.
-- |Set the value of the global environment.
--
setGEnv :: GlobalEnv -> VM ()
setGEnv genv = VM $ \_ _ lenv -> return (Yes genv lenv ())
-- | Update the global environment using the provided function.
-- |Update the global environment using the provided function.
--
updGEnv :: (GlobalEnv -> GlobalEnv) -> VM ()
updGEnv f = VM $ \_ genv lenv -> return (Yes (f genv) lenv ())
-- Vars -----------------------------------------------------------------------
-- | Add a mapping between a global var and its vectorised version to the state.
-- |Add a mapping between a global var and its vectorised version to the state.
--
defGlobalVar :: Var -> Var -> VM ()
defGlobalVar v v' = updGEnv $ \env ->
env { global_vars = extendVarEnv (global_vars env) v v'
......@@ -79,6 +83,11 @@ defGlobalVar v v' = updGEnv $ \env ->
lookupVectDecl :: Var -> VM (Maybe (Type, CoreExpr))
lookupVectDecl var = readGEnv $ \env ->