49104647097c983857688a2b38df73aa6b973aaa
[ghc.git] / compiler / vectorise / Vectorise / Type / Env.hs
1 {-# OPTIONS_GHC -XNoMonoLocalBinds #-}
2
3 module Vectorise.Type.Env (
4 vectTypeEnv,
5 ) where
6
7 import Vectorise.Env
8 import Vectorise.Vect
9 import Vectorise.Monad
10 import Vectorise.Builtins
11 import Vectorise.Type.TyConDecl
12 import Vectorise.Type.Classify
13 import Vectorise.Type.PADict
14 import Vectorise.Type.PData
15 import Vectorise.Type.PRepr
16 import Vectorise.Type.Repr
17 import Vectorise.Utils
18
19 import HscTypes
20 import CoreSyn
21 import CoreUtils
22 import CoreUnfold
23 import DataCon
24 import TyCon
25 import Type
26 import FamInstEnv
27 import OccName
28 import Id
29 import MkId
30 import NameEnv
31
32 import Unique
33 import UniqFM
34 import Util
35 import Outputable
36 import FastString
37 import MonadUtils
38 import Control.Monad
39 import Data.List
40
41
42 -- | Vectorise a type environment.
43 -- The type environment contains all the type things defined in a module.
44 --
45 vectTypeEnv :: TypeEnv
46 -> VM ( TypeEnv -- Vectorised type environment.
47 , [FamInst] -- New type family instances.
48 , [(Var, CoreExpr)]) -- New top level bindings.
49 vectTypeEnv env
50 = do
51 traceVt "** vectTypeEnv" $ ppr env
52
53 cs <- readGEnv $ mk_map . global_tycons
54
55 -- Split the list of TyCons into the ones we have to vectorise vs the
56 -- ones we can pass through unchanged. We also pass through algebraic
57 -- types that use non Haskell98 features, as we don't handle those.
58 let tycons = typeEnvTyCons env
59 groups = tyConGroups tycons
60
61 let (conv_tcs, keep_tcs) = classifyTyCons cs groups
62 orig_tcs = keep_tcs ++ conv_tcs
63 keep_dcs = concatMap tyConDataCons keep_tcs
64
65 -- Just use the unvectorised versions of these constructors in vectorised code.
66 zipWithM_ defTyCon keep_tcs keep_tcs
67 zipWithM_ defDataCon keep_dcs keep_dcs
68
69 -- Vectorise all the declarations.
70 new_tcs <- vectTyConDecls conv_tcs
71
72 -- We don't need to make new representation types for dictionary
73 -- constructors. The constructors are always fully applied, and we don't
74 -- need to lift them to arrays as a dictionary of a particular type
75 -- always has the same value.
76 let vect_tcs = filter (not . isClassTyCon)
77 $ keep_tcs ++ new_tcs
78
79 reprs <- mapM tyConRepr vect_tcs
80 repr_tcs <- zipWith3M buildPReprTyCon orig_tcs vect_tcs reprs
81 pdata_tcs <- zipWith3M buildPDataTyCon orig_tcs vect_tcs reprs
82 updGEnv $ extendFamEnv
83 $ map mkLocalFamInst
84 $ repr_tcs ++ pdata_tcs
85
86 -- Create PRepr and PData instances for the vectorised types.
87 -- We get back the binds for the instance functions,
88 -- and some new type constructors for the representation types.
89 (_, binds, inst_tcs) <- fixV $ \ ~(dfuns', _, _) ->
90 do
91 defTyConPAs (zipLazy vect_tcs dfuns')
92 reprs <- mapM tyConRepr vect_tcs
93
94 dfuns <- sequence
95 $ zipWith5 buildTyConBindings
96 orig_tcs
97 vect_tcs
98 repr_tcs
99 pdata_tcs
100 reprs
101
102 binds <- takeHoisted
103 return (dfuns, binds, repr_tcs ++ pdata_tcs)
104
105 -- The new type constructors are the vectorised versions of the originals,
106 -- plus the new type constructors that we use for the representations.
107 let all_new_tcs = new_tcs ++ inst_tcs
108
109 let new_env = extendTypeEnvList env
110 $ map ATyCon all_new_tcs
111 ++ [ADataCon dc | tc <- all_new_tcs
112 , dc <- tyConDataCons tc]
113
114 return (new_env, map mkLocalFamInst inst_tcs, binds)
115
116 where
117 mk_map env = listToUFM_Directly [(u, getUnique n /= u) | (u,n) <- nameEnvUniqueElts env]
118
119 buildTyConBindings :: TyCon -> TyCon -> TyCon -> TyCon -> SumRepr -> VM Var
120 buildTyConBindings orig_tc vect_tc prepr_tc pdata_tc repr
121 = do vectDataConWorkers orig_tc vect_tc pdata_tc
122 buildPADict vect_tc prepr_tc pdata_tc repr
123
124 vectDataConWorkers :: TyCon -> TyCon -> TyCon -> VM ()
125 vectDataConWorkers orig_tc vect_tc arr_tc
126 = do bs <- sequence
127 . zipWith3 def_worker (tyConDataCons orig_tc) rep_tys
128 $ zipWith4 mk_data_con (tyConDataCons vect_tc)
129 rep_tys
130 (inits rep_tys)
131 (tail $ tails rep_tys)
132 mapM_ (uncurry hoistBinding) bs
133 where
134 tyvars = tyConTyVars vect_tc
135 var_tys = mkTyVarTys tyvars
136 ty_args = map Type var_tys
137 res_ty = mkTyConApp vect_tc var_tys
138
139 cons = tyConDataCons vect_tc
140 arity = length cons
141 [arr_dc] = tyConDataCons arr_tc
142
143 rep_tys = map dataConRepArgTys $ tyConDataCons vect_tc
144
145
146 mk_data_con con tys pre post
147 = liftM2 (,) (vect_data_con con)
148 (lift_data_con tys pre post (mkDataConTag con))
149
150 sel_replicate len tag
151 | arity > 1 = do
152 rep <- builtin (selReplicate arity)
153 return [rep `mkApps` [len, tag]]
154
155 | otherwise = return []
156
157 vect_data_con con = return $ mkConApp con ty_args
158 lift_data_con tys pre_tys post_tys tag
159 = do
160 len <- builtin liftingContext
161 args <- mapM (newLocalVar (fsLit "xs"))
162 =<< mapM mkPDataType tys
163
164 sel <- sel_replicate (Var len) tag
165
166 pre <- mapM emptyPD (concat pre_tys)
167 post <- mapM emptyPD (concat post_tys)
168
169 return . mkLams (len : args)
170 . wrapFamInstBody arr_tc var_tys
171 . mkConApp arr_dc
172 $ ty_args ++ sel ++ pre ++ map Var args ++ post
173
174 def_worker data_con arg_tys mk_body
175 = do
176 arity <- polyArity tyvars
177 body <- closedV
178 . inBind orig_worker
179 . polyAbstract tyvars $ \args ->
180 liftM (mkLams (tyvars ++ args) . vectorised)
181 $ buildClosures tyvars [] arg_tys res_ty mk_body
182
183 raw_worker <- cloneId mkVectOcc orig_worker (exprType body)
184 let vect_worker = raw_worker `setIdUnfolding`
185 mkInlineUnfolding (Just arity) body
186 defGlobalVar orig_worker vect_worker
187 return (vect_worker, body)
188 where
189 orig_worker = dataConWorkId data_con
190