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

Fix the vectorisation of workers of data constructors

parent 96daec08
......@@ -1943,6 +1943,9 @@ data VectInfo
-- |Vectorisation information for 'ModIface'; i.e, the vectorisation information propagated
-- across module boundaries.
--
-- NB: The field 'ifaceVectInfoVar' explicitly contains the workers of data constructors as well as
-- class selectors — i.e., their mappings are /not/ implicitly generated from the data types.
--
data IfaceVectInfo
= IfaceVectInfo
{ ifaceVectInfoVar :: [Name] -- ^ All variables in here have a vectorised variant
......
......@@ -211,9 +211,10 @@ modVectInfo env mg_ids mg_tyCons vectDecls info
vectTypeTyCons = [tycon | VectType _ tycon _ <- vectDecls] ++
[tycon | VectClass tycon <- vectDecls]
vectDataCons = concatMap tyConDataCons vectTypeTyCons
ids = mg_ids ++ vectIds ++ selIds
ids = mg_ids ++ vectIds ++ dataConIds ++ selIds
tyCons = mg_tyCons ++ vectTypeTyCons
dataCons = concatMap tyConDataCons mg_tyCons ++ vectDataCons
dataConIds = map dataConWorkId dataCons
selIds = concat [ classAllSelIds cls
| tycon <- tyCons
, cls <- maybeToList . tyConClass_maybe $ tycon]
......
......@@ -522,7 +522,8 @@ unVectDict ty e
-- |Vectorise an 'n'-ary lambda abstraction by building a set of 'n' explicit closures.
--
-- All non-dictionary free variables go into the closure's environment, whereas the dictionary
-- variables are passed explicit (as conventional arguments) into the body during closure construction.
-- variables are passed explicit (as conventional arguments) into the body during closure
-- construction.
--
vectLam :: Bool -- ^ When the RHS of a binding, whether that binding should be inlined.
-> Bool -- ^ Whether the binding is a loop breaker.
......
......@@ -287,7 +287,7 @@ buildToArrPRepr vect_tc prepr_tc pdata_tc _ r
to_comp expr (Wrap ty)
= do
wrap_tc <- builtin wrapTyCon
(pwrap_tc, _) <- pdataReprTyCon (mkTyConApp wrap_tc [ty])
pwrap_tc <- pdataReprTyConExact (mkTyConApp wrap_tc [ty])
return $ wrapNewTypeBody pwrap_tc [ty] expr
......@@ -355,8 +355,8 @@ buildFromArrPRepr vect_tc prepr_tc pdata_tc _ r
from_comp _ res expr (Keep _ _) = return (res, [expr])
from_comp _ res expr (Wrap ty)
= do wrap_tc <- builtin wrapTyCon
(pwrap_tc, _) <- pdataReprTyCon (mkTyConApp wrap_tc [ty])
= do wrap_tc <- builtin wrapTyCon
pwrap_tc <- pdataReprTyConExact (mkTyConApp wrap_tc [ty])
return (res, [unwrapNewTypeBody pwrap_tc [ty]
$ unwrapFamInstScrut pwrap_tc [ty] expr])
......
......@@ -164,7 +164,6 @@ vectTypeEnv tycons vectTypeDecls vectClassDecls
-- Furthermore, 'drop_tcs' are those type constructors that we cannot vectorise.
; let maybeVectoriseTyCons = filter notLocalScalarTyCon tycons ++ impVectTyCons
(conv_tcs, keep_tcs, drop_tcs) = classifyTyCons vectTyConFlavour maybeVectoriseTyCons
orig_tcs = keep_tcs ++ conv_tcs
; traceVt " VECT SCALAR : " $ ppr localScalarTyCons
; traceVt " VECT [class] : " $ ppr impVectTyCons
......@@ -206,8 +205,8 @@ vectTypeEnv tycons vectTypeDecls vectClassDecls
-- 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
; let orig_tcs = filter (not . isClassTyCon) $ keep_tcs ++ conv_tcs
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.
......@@ -220,18 +219,36 @@ vectTypeEnv tycons vectTypeDecls vectClassDecls
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).
-- Generate workers for the vectorised data constructors, 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
$ zipWith5 buildTyConBindings
orig_tcs
vect_tcs
repr_tcs
pdata_tcs
pdatas_tcs
-- query the 'PData' instance type constructors for type constructors that have a
-- VECTORISE pragma with an explicit right-hand side (this is Item (3) of
-- "Note [Pragmas to vectorise tycons]" above)
; pdata_withRHS_tcs <- mapM pdataReprTyConExact
[ mkTyConApp tycon tys
| (tycon, _) <- vectTyConsWithRHS
, let tys = mkTyVarTys (tyConTyVars tycon)
]
-- build workers for all vectorised data constructors (except scalar ones)
; sequence_ $
zipWith3 vectDataConWorkers (orig_tcs ++ map fst vectTyConsWithRHS)
(vect_tcs ++ map snd vectTyConsWithRHS)
(pdata_tcs ++ pdata_withRHS_tcs)
-- build a 'PA' dictionary for all type constructors (except scalar ones and those
-- defined with an explicit right-hand side where the dictionary is user-supplied)
; dfuns <- sequence $
zipWith4 buildTyConPADict
vect_tcs
repr_tcs
pdata_tcs
pdatas_tcs
; binds <- takeHoisted
; return (dfuns, binds)
......@@ -244,23 +261,32 @@ vectTypeEnv tycons vectTypeDecls vectClassDecls
-- Helpers --------------------------------------------------------------------
buildTyConBindings :: TyCon -> TyCon -> TyCon -> TyCon -> TyCon -> VM Var
buildTyConBindings orig_tc vect_tc prepr_tc pdata_tc pdatas_tc
= do vectDataConWorkers orig_tc vect_tc pdata_tc
repr <- tyConRepr vect_tc
buildPADict vect_tc prepr_tc pdata_tc pdatas_tc repr
buildTyConPADict :: TyCon -> TyCon -> TyCon -> TyCon -> VM Var
buildTyConPADict vect_tc prepr_tc pdata_tc pdatas_tc
= tyConRepr vect_tc >>= buildPADict vect_tc prepr_tc pdata_tc pdatas_tc
-- Produce a custom-made worker for the data constructors of a vectorised data type. This includes
-- all data constructors that may be used in vetcorised code — i.e., all data constructors of data
-- types other than scalar ones. Also adds a mapping from the original to vectorised worker into
-- the vectorisation map.
--
-- FIXME: It's not nice that we need create a special worker after the data constructors has
-- already been constructed. Also, I don't think the worker is properly added to the data
-- constructor. Seems messy.
vectDataConWorkers :: TyCon -> TyCon -> TyCon -> VM ()
vectDataConWorkers orig_tc vect_tc arr_tc
= do bs <- sequence
. zipWith3 def_worker (tyConDataCons orig_tc) rep_tys
$ zipWith4 mk_data_con (tyConDataCons vect_tc)
rep_tys
(inits rep_tys)
(tail $ tails rep_tys)
mapM_ (uncurry hoistBinding) bs
where
= do { traceVt "Building vectorised worker for datatype" (ppr orig_tc)
; bs <- sequence
. zipWith3 def_worker (tyConDataCons orig_tc) rep_tys
$ zipWith4 mk_data_con (tyConDataCons vect_tc)
rep_tys
(inits rep_tys)
(tail $ tails rep_tys)
; mapM_ (uncurry hoistBinding) bs
}
where
tyvars = tyConTyVars vect_tc
var_tys = mkTyVarTys tyvars
ty_args = map Type var_tys
......@@ -272,7 +298,6 @@ vectDataConWorkers orig_tc vect_tc arr_tc
rep_tys = map dataConRepArgTys $ tyConDataCons vect_tc
mk_data_con con tys pre post
= liftM2 (,) (vect_data_con con)
(lift_data_con tys pre post (mkDataConTag con))
......
......@@ -276,8 +276,8 @@ buildToArrPRepr vect_tc prepr_tc pdata_tc r
to_comp expr (Wrap ty)
= do
wrap_tc <- builtin wrapTyCon
(pwrap_tc, _) <- pdataReprTyCon (mkTyConApp wrap_tc [ty])
return $ wrapNewTypeBody pwrap_tc [ty] expr
pwrap_tc <- pdataReprTyConExact (mkTyConApp wrap_tc [ty])
return $ wrapNewTypeBody pwrap_tc [ty] expr
buildFromArrPRepr :: TyCon -> TyCon -> TyCon -> SumRepr -> VM CoreExpr
......@@ -358,7 +358,7 @@ buildFromArrPRepr vect_tc prepr_tc pdata_tc r
from_comp _ res expr (Wrap ty)
= do
wrap_tc <- builtin wrapTyCon
(pwrap_tc, _) <- pdataReprTyCon (mkTyConApp wrap_tc [ty])
pwrap_tc <- pdataReprTyConExact (mkTyConApp wrap_tc [ty])
return (res, [unwrapNewTypeBody pwrap_tc [ty]
$ unwrapFamInstScrut pwrap_tc [ty] expr])
......
......@@ -15,7 +15,7 @@ module Vectorise.Utils.Base (
mkBuiltinCo,
mkVScrut,
pdataReprTyCon, pdatasReprTyCon,
pdataReprTyCon, pdataReprTyConExact, pdatasReprTyCon,
pdataReprDataCon, pdatasReprDataCon,
prDFunOfTyCon
) where
......@@ -28,12 +28,14 @@ import CoreSyn
import CoreUtils
import Coercion
import Type
import TypeRep
import TyCon
import DataCon
import MkId
import Literal
import Outputable
import FastString
import ListSetOps
import Control.Monad (liftM)
......@@ -150,6 +152,7 @@ mkBuiltinCo get_tc
-------------------------------------------------------------------------------
mkVScrut :: VExpr -> VM (CoreExpr, CoreExpr, TyCon, [Type])
mkVScrut (ve, le)
= do
......@@ -158,37 +161,57 @@ mkVScrut (ve, le)
where
ty = exprType ve
-- | Get the PData tycon that represents this type.
-- This tycon does not appear explicitly in the source program.
-- See Note [PData TyCons] in Vectorise.PRepr
-- |Get the representation tycon of the 'PData' data family for a given type.
--
-- This tycon does not appear explicitly in the source program — see Note [PData TyCons] in
-- 'Vectorise.Generic.Description':
--
-- @pdataReprTyCon {Sum2} = {PDataSum2}@
--
-- The type for which we look up a 'PData' instance may be more specific than the type in the
-- instance declaration. In that case the second component of the result will be more specific than
-- a set of distinct type variables.
--
pdataReprTyCon :: Type -> VM (TyCon, [Type])
pdataReprTyCon ty
= builtin pdataTyCon >>= (`lookupFamInst` [ty])
pdataReprTyCon ty = builtin pdataTyCon >>= (`lookupFamInst` [ty])
-- |Get the representation tycon of the 'PData' data family for a given type which must match the
-- type index in the looked up 'PData' instance exactly.
--
pdataReprTyConExact :: Type -> VM TyCon
pdataReprTyConExact ty
= do { (tycon, tys) <- pdataReprTyCon ty
; if uniqueTyVars tys
then
return tycon
else
cantVectorise "No exact 'PData' family instance for" (ppr ty)
}
where
uniqueTyVars tys = all isTyVarTy tys && hasNoDups (map extractTyVar tys)
where
extractTyVar (TyVarTy tv) = tv
extractTyVar _ = panic "Vectorise.Utils.Base: extractTyVar"
pdataReprDataCon :: Type -> VM (DataCon, [Type])
pdataReprDataCon ty
= do (tc, arg_tys) <- pdataReprTyCon ty
let [dc] = tyConDataCons tc
return (dc, arg_tys)
= do { (tc, arg_tys) <- pdataReprTyCon ty
; let [dc] = tyConDataCons tc
; return (dc, arg_tys)
}
pdatasReprTyCon :: Type -> VM (TyCon, [Type])
pdatasReprTyCon ty
= builtin pdatasTyCon >>= (`lookupFamInst` [ty])
pdatasReprTyCon ty = builtin pdatasTyCon >>= (`lookupFamInst` [ty])
pdatasReprDataCon :: Type -> VM (DataCon, [Type])
pdatasReprDataCon ty
= do (tc, arg_tys) <- pdatasReprTyCon ty
let [dc] = tyConDataCons tc
return (dc, arg_tys)
= do { (tc, arg_tys) <- pdatasReprTyCon ty
; let [dc] = tyConDataCons tc
; return (dc, arg_tys)
}
prDFunOfTyCon :: TyCon -> VM CoreExpr
prDFunOfTyCon tycon
= liftM Var
. maybeCantVectoriseM "No PR dictionary for tycon" (ppr tycon)
$ lookupTyConPR tycon
Markdown is supported
0% or
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment