Skip to content
GitLab
Projects
Groups
Snippets
/
Help
Help
Support
Community forum
Keyboard shortcuts
?
Submit feedback
Sign in / Register
Toggle navigation
Menu
Open sidebar
jberryman
GHC
Commits
112780e0
Commit
112780e0
authored
Aug 30, 2010
by
benl@ouroborus.net
Browse files
Comments and formatting to vectoriser
parent
5ee7f0e6
Changes
5
Hide whitespace changes
Inline
Side-by-side
compiler/iface/BuildTyCl.lhs
View file @
112780e0
...
...
@@ -36,8 +36,8 @@ import Outputable
------------------------------------------------------
buildSynTyCon :: Name -> [TyVar]
-> SynTyConRhs
-> Kind -- Kind of the RHS
-> Maybe (TyCon, [Type]) -- family instance if applicable
-> Kind --
^
Kind of the RHS
-> Maybe (TyCon, [Type]) --
^
family instance if applicable
-> TcRnIf m n TyCon
buildSynTyCon tc_name tvs rhs@(OpenSynTyCon {}) rhs_kind _
...
...
@@ -61,12 +61,12 @@ buildSynTyCon tc_name tvs rhs@(SynonymTyCon {}) rhs_kind mb_family
------------------------------------------------------
buildAlgTyCon :: Name -> [TyVar]
-> ThetaType -- Stupid theta
-> ThetaType --
^
Stupid theta
-> AlgTyConRhs
-> RecFlag
-> Bool -- True <=> want generics functions
-> Bool -- True <=> was declared in GADT syntax
-> Maybe (TyCon, [Type]) -- family instance if applicable
-> Bool --
^
True <=> want generics functions
-> Bool --
^
True <=> was declared in GADT syntax
-> Maybe (TyCon, [Type]) --
^
family instance if applicable
-> TcRnIf m n TyCon
buildAlgTyCon tc_name tvs stupid_theta rhs is_rec want_generics gadt_syn
...
...
@@ -84,7 +84,7 @@ buildAlgTyCon tc_name tvs stupid_theta rhs is_rec want_generics gadt_syn
; return tycon
}
-- If a family tycon with instance types is given, the current tycon is an
--
|
If a family tycon with instance types is given, the current tycon is an
-- instance of that family and we need to
--
-- (1) create a coercion that identifies the family instance type and the
...
...
@@ -132,9 +132,9 @@ mkDataTyConRhs cons
}
mkNewTyConRhs :: Name -> TyCon -> DataCon -> TcRnIf m n AlgTyConRhs
-- Monadic because it makes a Name for the coercion TyCon
-- We pass the Name of the parent TyCon, as well as the TyCon itself,
-- because the latter is part of a knot, whereas the former is not.
--
^
Monadic because it makes a Name for the coercion TyCon
--
We pass the Name of the parent TyCon, as well as the TyCon itself,
--
because the latter is part of a knot, whereas the former is not.
mkNewTyConRhs tycon_name tycon con
= do { co_tycon_name <- newImplicitBinder tycon_name mkNewTyCoOcc
; let co_tycon = mkNewTypeCoercion co_tycon_name tycon etad_tvs etad_rhs
...
...
compiler/vectorise/VectMonad.hs
View file @
112780e0
...
...
@@ -439,6 +439,8 @@ newTyVar fs k
u
<-
liftDs
newUnique
return
$
mkTyVar
(
mkSysTvName
u
fs
)
k
-- | 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'
...
...
@@ -448,14 +450,14 @@ defGlobalVar v v' = updGEnv $ \env ->
upd
env
|
isExportedId
v
=
extendVarEnv
env
v
(
v
,
v'
)
|
otherwise
=
env
-- Var ------------------------------------------------------------------------
-- | Lookup the vectorised and\/or lifted versions of this variable.
-- If it's in the global environment we get the vectorised version.
-- If it's in the local environment we get both the vectorised and lifted version.
--
lookupVar
::
Var
->
VM
(
Scope
Var
(
Var
,
Var
))
lookupVar
v
=
do
r
<-
readLEnv
$
\
env
->
lookupVarEnv
(
local_vars
env
)
v
=
do
r
<-
readLEnv
$
\
env
->
lookupVarEnv
(
local_vars
env
)
v
case
r
of
Just
e
->
return
(
Local
e
)
Nothing
->
liftM
Global
...
...
@@ -581,6 +583,8 @@ lookupFamInst tycon tys
(
ppr
$
mkTyConApp
tycon
tys
)
}
-- | Run a vectorisation computation.
initV
::
PackageId
->
HscEnv
->
ModGuts
->
VectInfo
->
VM
a
->
IO
(
Maybe
(
VectInfo
,
a
))
initV
pkg
hsc_env
guts
info
p
=
do
...
...
compiler/vectorise/VectType.hs
View file @
112780e0
...
...
@@ -45,13 +45,16 @@ import Data.List ( inits, tails, zipWith4, zipWith5 )
-- ----------------------------------------------------------------------------
-- Types
-- | Vectorise a type constructor.
vectTyCon
::
TyCon
->
VM
TyCon
vectTyCon
tc
|
isFunTyCon
tc
=
builtin
closureTyCon
|
isBoxedTupleTyCon
tc
=
return
tc
|
isUnLiftedTyCon
tc
=
return
tc
|
otherwise
=
maybeCantVectoriseM
"Tycon not vectorised:"
(
ppr
tc
)
$
lookupTyCon
tc
|
otherwise
=
maybeCantVectoriseM
"Tycon not vectorised: "
(
ppr
tc
)
$
lookupTyCon
tc
vectAndLiftType
::
Type
->
VM
(
Type
,
Type
)
vectAndLiftType
ty
|
Just
ty'
<-
coreView
ty
=
vectAndLiftType
ty'
...
...
@@ -67,6 +70,7 @@ vectAndLiftType ty
(
tyvars
,
mono_ty
)
=
splitForAllTys
ty
-- | Vectorise a type.
vectType
::
Type
->
VM
Type
vectType
ty
|
Just
ty'
<-
coreView
ty
=
vectType
ty'
vectType
(
TyVarTy
tv
)
=
return
$
TyVarTy
tv
...
...
@@ -87,6 +91,7 @@ vectType ty = cantVectorise "Can't vectorise type" (ppr ty)
vectAndBoxType
::
Type
->
VM
Type
vectAndBoxType
ty
=
vectType
ty
>>=
boxType
-- | Add quantified vars and dictionary parameters to the front of a type.
abstractType
::
[
TyVar
]
->
[
Type
]
->
Type
->
Type
abstractType
tyvars
dicts
=
mkForAllTys
tyvars
.
mkFunTys
dicts
...
...
@@ -102,6 +107,7 @@ boxType ty
case
r
of
Just
tycon'
->
return
$
mkTyConApp
tycon'
[]
Nothing
->
return
ty
boxType
ty
=
return
ty
-- ----------------------------------------------------------------------------
...
...
@@ -109,14 +115,21 @@ boxType ty = return ty
type
TyConGroup
=
([
TyCon
],
UniqSet
TyCon
)
-- | Vectorise a type environment.
-- The type environment contains all the type things defined in a module.
vectTypeEnv
::
TypeEnv
->
VM
(
TypeEnv
,
[
FamInst
],
[(
Var
,
CoreExpr
)])
vectTypeEnv
env
=
do
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
(
conv_tcs
,
keep_tcs
)
=
classifyTyCons
cs
groups
keep_dcs
=
concatMap
tyConDataCons
keep_tcs
zipWithM_
defTyCon
keep_tcs
keep_tcs
zipWithM_
defDataCon
keep_dcs
keep_dcs
new_tcs
<-
vectTyConDecls
conv_tcs
let
orig_tcs
=
keep_tcs
++
conv_tcs
...
...
@@ -151,6 +164,7 @@ vectTypeEnv env
mk_map
env
=
listToUFM_Directly
[(
u
,
getUnique
n
/=
u
)
|
(
u
,
n
)
<-
nameEnvUniqueElts
env
]
-- | Vectorise some (possibly recursively defined) type constructors.
vectTyConDecls
::
[
TyCon
]
->
VM
[
TyCon
]
vectTyConDecls
tcs
=
fixV
$
\
tcs'
->
do
...
...
@@ -848,8 +862,8 @@ paMethods = [("dictPRepr", buildPRDict),
(
"fromArrPRepr"
,
buildFromArrPRepr
)]
-- | 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:
--
converted (first list) or not (second list). The first argument contains
--
information about the conversion status of external tycons:
--
-- * tycons which have converted versions are mapped to True
-- * tycons which are not changed by vectorisation are mapped to False
...
...
compiler/vectorise/VectUtils.hs
View file @
112780e0
...
...
@@ -281,6 +281,7 @@ combinePD ty len sel xs
where
n
=
length
xs
-- | Like `replicatePD` but use the lifting context in the vectoriser state.
liftPD
::
CoreExpr
->
VM
CoreExpr
liftPD
x
=
do
...
...
@@ -358,6 +359,8 @@ addInlineArity DontInline _ = DontInline
inlineMe
::
Inline
inlineMe
=
Inline
0
-- Hoising --------------------------------------------------------------------
hoistBinding
::
Var
->
CoreExpr
->
VM
()
hoistBinding
v
e
=
updGEnv
$
\
env
->
env
{
global_bindings
=
(
v
,
e
)
:
global_bindings
env
}
...
...
@@ -412,24 +415,24 @@ boxExpr ty (vexpr, lexpr)
Nothing -> return (vexpr, lexpr)
-}
-- Closures -------------------------------------------------------------------
mkClosure
::
Type
->
Type
->
Type
->
VExpr
->
VExpr
->
VM
VExpr
mkClosure
arg_ty
res_ty
env_ty
(
vfn
,
lfn
)
(
venv
,
lenv
)
=
do
dict
<-
paDictOfType
env_ty
mkv
<-
builtin
closureVar
mkl
<-
builtin
liftedClosureVar
=
do
Just
dict
<-
paDictOfType
env_ty
mkv
<-
builtin
closureVar
mkl
<-
builtin
liftedClosureVar
return
(
Var
mkv
`
mkTyApps
`
[
arg_ty
,
res_ty
,
env_ty
]
`
mkApps
`
[
dict
,
vfn
,
lfn
,
venv
],
Var
mkl
`
mkTyApps
`
[
arg_ty
,
res_ty
,
env_ty
]
`
mkApps
`
[
dict
,
vfn
,
lfn
,
lenv
])
mkClosureApp
::
Type
->
Type
->
VExpr
->
VExpr
->
VM
VExpr
mkClosureApp
arg_ty
res_ty
(
vclo
,
lclo
)
(
varg
,
larg
)
=
do
vapply
<-
builtin
applyVar
=
do
vapply
<-
builtin
applyVar
lapply
<-
builtin
liftedApplyVar
lc
<-
builtin
liftingContext
return
(
Var
vapply
`
mkTyApps
`
[
arg_ty
,
res_ty
]
`
mkApps
`
[
vclo
,
varg
],
Var
lapply
`
mkTyApps
`
[
arg_ty
,
res_ty
]
`
mkApps
`
[
Var
lc
,
lclo
,
larg
])
buildClosures
::
[
TyVar
]
->
[
VVar
]
->
[
Type
]
->
Type
->
VM
VExpr
->
VM
VExpr
buildClosures
_
_
[]
_
mk_body
=
mk_body
...
...
@@ -471,6 +474,8 @@ buildClosure tvs vars arg_ty res_ty mk_body
mkClosure
arg_ty
res_ty
env_ty
fn
env
-- Environments ---------------------------------------------------------------
buildEnv
::
[
VVar
]
->
VM
(
Type
,
VExpr
,
VExpr
->
VExpr
->
VExpr
)
buildEnv
[]
=
do
ty
<-
voidType
...
...
compiler/vectorise/Vectorise.hs
View file @
112780e0
...
...
@@ -42,51 +42,109 @@ vectorise backend guts = do
hsc_env
<-
getHscEnv
liftIO
$
vectoriseIO
backend
hsc_env
guts
-- | Vectorise a single monad, given its HscEnv (code gen environment).
vectoriseIO
::
PackageId
->
HscEnv
->
ModGuts
->
IO
ModGuts
vectoriseIO
backend
hsc_env
guts
=
do
=
do
-- Get information about currently loaded external packages.
eps
<-
hscEPS
hsc_env
-- Combine vectorisation info from the current module, and external ones.
let
info
=
hptVectInfo
hsc_env
`
plusVectInfo
`
eps_vect_info
eps
-- Run the main VM computation.
Just
(
info'
,
guts'
)
<-
initV
backend
hsc_env
guts
info
(
vectModule
guts
)
return
(
guts'
{
mg_vect_info
=
info'
})
-- | Vectorise a single module, in the VM monad.
vectModule
::
ModGuts
->
VM
ModGuts
vectModule
guts
=
do
=
do
-- Vectorise the type environment.
-- This may add new TyCons and DataCons.
-- TODO: What new binds do we get back here?
(
types'
,
fam_insts
,
tc_binds
)
<-
vectTypeEnv
(
mg_types
guts
)
-- TODO: What is this?
let
fam_inst_env'
=
extendFamInstEnvList
(
mg_fam_inst_env
guts
)
fam_insts
updGEnv
(
setFamInstEnv
fam_inst_env'
)
-- dicts <- mapM buildPADict pa_insts
-- workers <- mapM vectDataConWorkers pa_insts
-- Vectorise all the top level bindings.
binds'
<-
mapM
vectTopBind
(
mg_binds
guts
)
return
$
guts
{
mg_types
=
types'
,
mg_binds
=
Rec
tc_binds
:
binds'
,
mg_fam_inst_env
=
fam_inst_env'
,
mg_fam_insts
=
mg_fam_insts
guts
++
fam_insts
}
-- | Try to vectorise a top-level binding.
-- If it doesn't vectorise then return it unharmed.
--
-- For example, for the binding
--
-- @
-- foo :: Int -> Int
-- foo = \x -> x + x
-- @
--
-- we get
-- @
-- foo :: Int -> Int
-- foo = \x -> vfoo $: x
--
-- v_foo :: Closure void vfoo lfoo
-- v_foo = closure vfoo lfoo void
--
-- vfoo :: Void -> Int -> Int
-- vfoo = ...
--
-- lfoo :: PData Void -> PData Int -> PData Int
-- lfoo = ...
-- @
--
-- @vfoo@ is the "vectorised", or scalar, version that does the same as the original
-- function foo, but takes an explicit environment.
--
-- @lfoo@ is the "lifted" version that works on arrays.
--
-- @v_foo@ combines both of these into a `Closure` that also contains the
-- environment.
--
-- The original binding @foo@ is rewritten to call the vectorised version
-- present in the closure.
--
vectTopBind
::
CoreBind
->
VM
CoreBind
vectTopBind
b
@
(
NonRec
var
expr
)
=
do
(
inline
,
expr'
)
<-
vectTopRhs
var
expr
var'
<-
vectTopBinder
var
inline
expr'
hs
<-
takeHoisted
cexpr
<-
tryConvert
var
var'
expr
=
do
(
inline
,
expr'
)
<-
vectTopRhs
var
expr
var'
<-
vectTopBinder
var
inline
expr'
-- Vectorising the body may create other top-level bindings.
hs
<-
takeHoisted
-- To get the same functionality as the original body we project
-- out its vectorised version from the closure.
cexpr
<-
tryConvert
var
var'
expr
return
.
Rec
$
(
var
,
cexpr
)
:
(
var'
,
expr'
)
:
hs
`
orElseV
`
return
b
vectTopBind
b
@
(
Rec
bs
)
=
do
(
vars'
,
_
,
exprs'
)
<-
fixV
$
\
~
(
_
,
inlines
,
rhss
)
->
do
vars'
<-
sequence
[
vectTopBinder
var
inline
rhs
|
(
var
,
~
(
inline
,
rhs
))
<-
zipLazy
vars
(
zip
inlines
rhss
)]
(
inlines'
,
exprs'
)
<-
mapAndUnzipM
(
uncurry
vectTopRhs
)
bs
return
(
vars'
,
inlines'
,
exprs'
)
=
do
(
vars'
,
_
,
exprs'
)
<-
fixV
$
\
~
(
_
,
inlines
,
rhss
)
->
do
vars'
<-
sequence
[
vectTopBinder
var
inline
rhs
|
(
var
,
~
(
inline
,
rhs
))
<-
zipLazy
vars
(
zip
inlines
rhss
)]
(
inlines'
,
exprs'
)
<-
mapAndUnzipM
(
uncurry
vectTopRhs
)
bs
return
(
vars'
,
inlines'
,
exprs'
)
hs
<-
takeHoisted
cexprs
<-
sequence
$
zipWith3
tryConvert
vars
vars'
exprs
return
.
Rec
$
zip
vars
cexprs
++
zip
vars'
exprs'
++
hs
...
...
@@ -95,11 +153,22 @@ vectTopBind b@(Rec bs)
where
(
vars
,
exprs
)
=
unzip
bs
-- NOTE: vectTopBinder *MUST* be lazy in inline and expr because of how it is
-- used inside of fixV in vectTopBind
vectTopBinder
::
Var
->
Inline
->
CoreExpr
->
VM
Var
-- | Make the vectorised version of this top level binder, and add the mapping
-- between it and the original to the state. For some binder @foo@ the vectorised
-- version is @$v_foo@
--
-- NOTE: vectTopBinder *MUST* be lazy in inline and expr because of how it is
-- used inside of fixV in vectTopBind
vectTopBinder
::
Var
-- ^ Name of the binding.
->
Inline
-- ^ Whether it should be inlined, used to annotate it.
->
CoreExpr
-- ^ RHS of the binding, used to set the `Unfolding` of the returned `Var`.
->
VM
Var
-- ^ Name of the vectorised binding.
vectTopBinder
var
inline
expr
=
do
=
do
-- Vectorise the type attached to the var.
vty
<-
vectType
(
idType
var
)
var'
<-
liftM
(`
setIdUnfolding
`
unfolding
)
$
cloneId
mkVectOcc
var
vty
defGlobalVar
var
var'
...
...
@@ -109,22 +178,37 @@ vectTopBinder var inline expr
Inline
arity
->
mkInlineRule
expr
(
Just
arity
)
DontInline
->
noUnfolding
vectTopRhs
::
Var
->
CoreExpr
->
VM
(
Inline
,
CoreExpr
)
-- | Vectorise the RHS of a top-level binding, in an empty local environment.
vectTopRhs
::
Var
-- ^ Name of the binding.
->
CoreExpr
-- ^ Body of the binding.
->
VM
(
Inline
,
CoreExpr
)
vectTopRhs
var
expr
=
closedV
$
do
(
inline
,
vexpr
)
<-
inBind
var
$
vectPolyExpr
(
isLoopBreaker
$
idOccInfo
var
)
=
dtrace
(
vcat
[
text
"vectTopRhs"
,
ppr
expr
])
$
closedV
$
do
(
inline
,
vexpr
)
<-
inBind
var
$
vectPolyExpr
(
isLoopBreaker
$
idOccInfo
var
)
(
freeVars
expr
)
return
(
inline
,
vectorised
vexpr
)
tryConvert
::
Var
->
Var
->
CoreExpr
->
VM
CoreExpr
-- | Project out the vectorised version of a binding from some closure,
-- or return the original body if that doesn't work.
tryConvert
::
Var
-- ^ Name of the original binding (eg @foo@)
->
Var
-- ^ Name of vectorised version of binding (eg @$vfoo@)
->
CoreExpr
-- ^ The original body of the binding.
->
VM
CoreExpr
tryConvert
var
vect_var
rhs
=
fromVect
(
idType
var
)
(
Var
vect_var
)
`
orElseV
`
return
rhs
-- ----------------------------------------------------------------------------
-- Bindings
-- | Vectorise a binder variable, along with its attached type.
vectBndr
::
Var
->
VM
VVar
vectBndr
v
=
do
...
...
@@ -136,6 +220,9 @@ vectBndr v
where
mapTo
vv
lv
env
=
env
{
local_vars
=
extendVarEnv
(
local_vars
env
)
v
(
vv
,
lv
)
}
-- | Vectorise a binder variable, along with its attached type,
-- but give the result a new name.
vectBndrNew
::
Var
->
FastString
->
VM
VVar
vectBndrNew
v
fs
=
do
...
...
@@ -146,6 +233,8 @@ vectBndrNew v fs
where
upd
vv
env
=
env
{
local_vars
=
extendVarEnv
(
local_vars
env
)
v
vv
}
-- | Vectorise a binder then run a computation with that binder in scope.
vectBndrIn
::
Var
->
VM
a
->
VM
(
VVar
,
a
)
vectBndrIn
v
p
=
localV
...
...
@@ -154,6 +243,8 @@ vectBndrIn v p
x
<-
p
return
(
vv
,
x
)
-- | Vectorise a binder, give it a new name, then run a computation with that binder in scope.
vectBndrNewIn
::
Var
->
FastString
->
VM
a
->
VM
(
VVar
,
a
)
vectBndrNewIn
v
fs
p
=
localV
...
...
@@ -162,6 +253,7 @@ vectBndrNewIn v fs p
x
<-
p
return
(
vv
,
x
)
-- | Vectorise some binders, then run a computation with them in scope.
vectBndrsIn
::
[
Var
]
->
VM
a
->
VM
([
VVar
],
a
)
vectBndrsIn
vs
p
=
localV
...
...
@@ -170,13 +262,17 @@ vectBndrsIn vs p
x
<-
p
return
(
vvs
,
x
)
-- ----------------------------------------------------------------------------
-- Expressions
-- | Vectorise a variable, producing the vectorised and lifted versions.
vectVar
::
Var
->
VM
VExpr
vectVar
v
=
do
=
do
-- lookup the variable from the environment.
r
<-
lookupVar
v
case
r
of
Local
(
vv
,
lv
)
->
return
(
Var
vv
,
Var
lv
)
Global
vv
->
do
...
...
@@ -184,30 +280,42 @@ vectVar v
lexpr
<-
liftPD
vexpr
return
(
vexpr
,
lexpr
)
-- | Like `vectVar` but also add type applications to the variables.
vectPolyVar
::
Var
->
[
Type
]
->
VM
VExpr
vectPolyVar
v
tys
=
do
vtys
<-
mapM
vectType
tys
r
<-
lookupVar
v
vtys
<-
mapM
vectType
tys
r
<-
lookupVar
v
case
r
of
Local
(
vv
,
lv
)
->
liftM2
(,)
(
polyApply
(
Var
vv
)
vtys
)
(
polyApply
(
Var
lv
)
vtys
)
Global
poly
->
do
vexpr
<-
polyApply
(
Var
poly
)
vtys
lexpr
<-
liftPD
vexpr
return
(
vexpr
,
lexpr
)
Local
(
vv
,
lv
)
->
liftM2
(,)
(
polyApply
(
Var
vv
)
vtys
)
(
polyApply
(
Var
lv
)
vtys
)
Global
poly
->
do
vexpr
<-
polyApply
(
Var
poly
)
vtys
lexpr
<-
liftPD
vexpr
return
(
vexpr
,
lexpr
)
-- | Lifted literals are created by replicating them.
vectLiteral
::
Literal
->
VM
VExpr
vectLiteral
lit
=
do
lexpr
<-
liftPD
(
Lit
lit
)
return
(
Lit
lit
,
lexpr
)
vectPolyExpr
::
Bool
->
CoreExprWithFVs
->
VM
(
Inline
,
VExpr
)
-- | Vectorise a polymorphic expression
vectPolyExpr
::
Bool
-- ^ When vectorising the RHS of a binding, whether that
-- binding is a loop breaker.
->
CoreExprWithFVs
->
VM
(
Inline
,
VExpr
)
vectPolyExpr
loop_breaker
(
_
,
AnnNote
note
expr
)
=
do
(
inline
,
expr'
)
<-
vectPolyExpr
loop_breaker
expr
=
do
(
inline
,
expr'
)
<-
vectPolyExpr
loop_breaker
expr
return
(
inline
,
vNote
note
expr'
)
vectPolyExpr
loop_breaker
expr
=
do
arity
<-
polyArity
tvs
...
...
@@ -219,13 +327,17 @@ vectPolyExpr loop_breaker expr
where
(
tvs
,
mono
)
=
collectAnnTypeBinders
expr
-- | Vectorise a core expression.
vectExpr
::
CoreExprWithFVs
->
VM
VExpr
vectExpr
(
_
,
AnnType
ty
)
=
liftM
vType
(
vectType
ty
)
vectExpr
(
_
,
AnnVar
v
)
=
vectVar
v
vectExpr
(
_
,
AnnVar
v
)
=
vectVar
v
vectExpr
(
_
,
AnnLit
lit
)
=
vectLiteral
lit
vectExpr
(
_
,
AnnLit
lit
)
=
vectLiteral
lit
vectExpr
(
_
,
AnnNote
note
expr
)
=
liftM
(
vNote
note
)
(
vectExpr
expr
)
...
...
@@ -247,12 +359,22 @@ vectExpr (_, AnnApp (_, AnnVar v) (_, AnnLit lit))
is_special_con
con
=
con
`
elem
`
[
intDataCon
,
floatDataCon
,
doubleDataCon
]
-- TODO: Avoid using closure application for dictionaries.
-- vectExpr (_, AnnApp fn arg)
-- | if is application of dictionary
-- just use regular app instead of closure app.
-- for lifted version.
-- do liftPD (sub a dNumber)
-- lift the result of the selection, not sub and dNumber seprately.
vectExpr
(
_
,
AnnApp
fn
arg
)
=
do
arg_ty'
<-
vectType
arg_ty
res_ty'
<-
vectType
res_ty
fn'
<-
vectExpr
fn
arg'
<-
vectExpr
arg
mkClosureApp
arg_ty'
res_ty'
fn'
arg'
where
(
arg_ty
,
res_ty
)
=
splitFunTy
.
exprType
$
deAnnotate
fn
...
...
@@ -296,7 +418,14 @@ onlyIfV (isEmptyVarSet fvs) (vectScalarLam bs $ deAnnotate body)
vectExpr
e
=
cantVectorise
"Can't vectorise expression"
(
ppr
$
deAnnotate
e
)
vectFnExpr
::
Bool
->
Bool
->
CoreExprWithFVs
->
VM
(
Inline
,
VExpr
)
-- | Vectorise an expression with an outer lambda abstraction.
vectFnExpr
::
Bool
-- ^ When the RHS of a binding, whether that binding should be inlined.
->
Bool
-- ^ Whether the binding is a loop breaker.
->
CoreExprWithFVs
-- ^ Expression to vectorise. Must have an outer `AnnLam`.
->
VM
(
Inline
,
VExpr
)
vectFnExpr
inline
loop_breaker
e
@
(
fvs
,
AnnLam
bndr
_
)
|
isId
bndr
=
onlyIfV
(
isEmptyVarSet
fvs
)
(
mark
DontInline
.
vectScalarLam
bs
$
deAnnotate
body
)
...
...
@@ -308,7 +437,12 @@ vectFnExpr _ _ e = mark DontInline $ vectExpr e
mark
::
Inline
->
VM
a
->
VM
(
Inline
,
a
)
mark
b
p
=
do
{
x
<-
p
;
return
(
b
,
x
)
}
vectScalarLam
::
[
Var
]
->
CoreExpr
->
VM
VExpr
-- | Vectorise a function where are the args have scalar type, that is Int, Float or Double.
vectScalarLam
::
[
Var
]
-- ^ Bound variables of function.
->
CoreExpr
-- ^ Function body.
->
VM
VExpr
vectScalarLam
args
body
=
do
scalars
<-
globalScalars
...
...
@@ -317,23 +451,24 @@ vectScalarLam args body
&&
is_scalar
(
extendVarSetList
scalars
args
)
body
&&
uses
scalars
body
)
$
do
fn_var
<-
hoistExpr
(
fsLit
"fn"
)
(
mkLams
args
body
)
DontInline
zipf
<-
zipScalars
arg_tys
res_ty
clo
<-
scalarClosure
arg_tys
res_ty
(
Var
fn_var
)
fn_var
<-
hoistExpr
(
fsLit
"fn"
)
(
mkLams
args
body
)
DontInline
zipf
<-
zipScalars
arg_tys
res_ty
clo
<-
scalarClosure
arg_tys
res_ty
(
Var
fn_var
)
(
zipf
`
App
`
Var
fn_var
)
clo_var
<-
hoistExpr
(
fsLit
"clo"
)
clo
DontInline
lclo
<-
liftPD
(
Var
clo_var
)
lclo
<-
liftPD
(
Var
clo_var
)
return
(
Var
clo_var
,
lclo
)
where
arg_tys
=
map
idType
args
res_ty
=
exprType
body
is_scalar_ty
ty
|
Just
(
tycon
,
[]
)
<-
splitTyConApp_maybe
ty
=
tycon
==
intTyCon
||
tycon
==
floatTyCon
||
tycon
==
doubleTyCon
is_scalar_ty
ty
|
Just
(
tycon
,
[]
)
<-
splitTyConApp_maybe
ty
=
tycon
==
intTyCon
||
tycon
==
floatTyCon
||
tycon
==
doubleTyCon