Vectorisation of: enumFromTo, div, intSqrt
[ghc-hetmet.git] / compiler / vectorise / VectBuiltIn.hs
1 module VectBuiltIn (
2   Builtins(..), sumTyCon, prodTyCon, uarrTy, intPrimArrayTy,
3   combinePAVar,
4   initBuiltins, initBuiltinVars, initBuiltinTyCons, initBuiltinDataCons,
5   initBuiltinPAs, initBuiltinPRs,
6   initBuiltinBoxedTyCons,
7
8   primMethod, primPArray
9 ) where
10
11 import DsMonad
12 import IfaceEnv        ( lookupOrig )
13
14 import Module
15 import DataCon         ( DataCon, dataConName, dataConWorkId )
16 import TyCon           ( TyCon, tyConName, tyConDataCons )
17 import Var             ( Var )
18 import Id              ( mkSysLocal )
19 import Name            ( Name, getOccString )
20 import NameEnv
21 import OccName
22
23 import TypeRep         ( funTyCon )
24 import Type            ( Type, mkTyConApp )
25 import TysPrim
26 import TysWiredIn      ( unitTyCon, unitDataCon,
27                          tupleTyCon, tupleCon,
28                          intTyCon, intTyConName, intTy,
29                          doubleTyCon, doubleTyConName,
30                          boolTyCon, boolTyConName, trueDataCon, falseDataCon,
31                          parrTyConName )
32 import PrelNames       ( gHC_PARR )
33 import BasicTypes      ( Boxity(..) )
34
35 import FastString
36 import Outputable
37
38 import Data.Array
39 import Control.Monad   ( liftM, zipWithM )
40 import Data.List       ( unzip4 )
41
42 mAX_NDP_PROD :: Int
43 mAX_NDP_PROD = 5
44
45 mAX_NDP_SUM :: Int
46 mAX_NDP_SUM = 3
47
48 mAX_NDP_COMBINE :: Int
49 mAX_NDP_COMBINE = 2
50
51 mkNDPModule :: FastString -> Module
52 mkNDPModule m = mkModule ndpPackageId (mkModuleNameFS m)
53
54 nDP_UARR,
55   nDP_PARRAY,
56   nDP_REPR,
57   nDP_CLOSURE,
58   nDP_UNBOXED,
59   nDP_INSTANCES,
60   nDP_COMBINATORS,
61   nDP_PRELUDE_PARR,
62   nDP_PRELUDE_INT,
63   nDP_PRELUDE_DOUBLE,
64   nDP_PRELUDE_BOOL,
65   nDP_PRELUDE_TUPLE :: Module
66
67 nDP_UARR        = mkNDPModule (fsLit "Data.Array.Parallel.Unlifted.Flat.UArr")
68 nDP_PARRAY      = mkNDPModule (fsLit "Data.Array.Parallel.Lifted.PArray")
69 nDP_REPR        = mkNDPModule (fsLit "Data.Array.Parallel.Lifted.Repr")
70 nDP_CLOSURE     = mkNDPModule (fsLit "Data.Array.Parallel.Lifted.Closure")
71 nDP_UNBOXED     = mkNDPModule (fsLit "Data.Array.Parallel.Lifted.Unboxed")
72 nDP_INSTANCES   = mkNDPModule (fsLit "Data.Array.Parallel.Lifted.Instances")
73 nDP_COMBINATORS = mkNDPModule (fsLit "Data.Array.Parallel.Lifted.Combinators")
74
75 nDP_PRELUDE_PARR = mkNDPModule (fsLit "Data.Array.Parallel.Prelude.Base.PArr")
76 nDP_PRELUDE_INT  = mkNDPModule (fsLit "Data.Array.Parallel.Prelude.Base.Int")
77 nDP_PRELUDE_DOUBLE = mkNDPModule (fsLit "Data.Array.Parallel.Prelude.Base.Double")
78 nDP_PRELUDE_BOOL = mkNDPModule (fsLit "Data.Array.Parallel.Prelude.Base.Bool")
79 nDP_PRELUDE_TUPLE  = mkNDPModule (fsLit "Data.Array.Parallel.Prelude.Base.Tuple")
80
81 data Builtins = Builtins {
82                   parrayTyCon      :: TyCon
83                 , paTyCon          :: TyCon
84                 , paDataCon        :: DataCon
85                 , preprTyCon       :: TyCon
86                 , prTyCon          :: TyCon
87                 , prDataCon        :: DataCon
88                 , uarrTyCon        :: TyCon
89                 , voidTyCon        :: TyCon
90                 , wrapTyCon        :: TyCon
91                 , enumerationTyCon :: TyCon
92                 , sumTyCons        :: Array Int TyCon
93                 , closureTyCon     :: TyCon
94                 , voidVar          :: Var
95                 , mkPRVar          :: Var
96                 , mkClosureVar     :: Var
97                 , applyClosureVar  :: Var
98                 , mkClosurePVar    :: Var
99                 , applyClosurePVar :: Var
100                 , replicatePAIntPrimVar :: Var
101                 , upToPAIntPrimVar :: Var
102                 , selectPAIntPrimVar :: Var
103                 , truesPABoolPrimVar :: Var
104                 , lengthPAVar      :: Var
105                 , replicatePAVar   :: Var
106                 , emptyPAVar       :: Var
107                 , packPAVar        :: Var
108                 , combinePAVars    :: Array Int Var
109                 , liftingContext   :: Var
110                 }
111
112 uarrTy :: Type -> Builtins -> Type
113 uarrTy ty bi = mkTyConApp (uarrTyCon bi) [ty]
114
115 intPrimArrayTy :: Builtins -> Type
116 intPrimArrayTy = uarrTy intTy
117
118 sumTyCon :: Int -> Builtins -> TyCon
119 sumTyCon n bi
120   | n >= 2 && n <= mAX_NDP_SUM = sumTyCons bi ! n
121   | otherwise = pprPanic "sumTyCon" (ppr n)
122
123 prodTyCon :: Int -> Builtins -> TyCon
124 prodTyCon n bi
125   | n == 1                      = wrapTyCon bi
126   | n >= 0 && n <= mAX_NDP_PROD = tupleTyCon Boxed n
127   | otherwise = pprPanic "prodTyCon" (ppr n)
128
129 combinePAVar :: Int -> Builtins -> Var
130 combinePAVar n bi
131   | n >= 2 && n <= mAX_NDP_COMBINE = combinePAVars bi ! n
132   | otherwise = pprPanic "combinePAVar" (ppr n)
133
134 initBuiltins :: DsM Builtins
135 initBuiltins
136   = do
137       parrayTyCon  <- externalTyCon nDP_PARRAY (fsLit "PArray")
138       paTyCon      <- externalTyCon nDP_PARRAY (fsLit "PA")
139       let [paDataCon] = tyConDataCons paTyCon
140       preprTyCon   <- externalTyCon nDP_PARRAY (fsLit "PRepr")
141       prTyCon      <- externalTyCon nDP_PARRAY (fsLit "PR")
142       let [prDataCon] = tyConDataCons prTyCon
143       uarrTyCon    <- externalTyCon nDP_UARR   (fsLit "UArr")
144       closureTyCon <- externalTyCon nDP_CLOSURE (fsLit ":->")
145
146       voidTyCon    <- externalTyCon nDP_REPR (fsLit "Void")
147       wrapTyCon    <- externalTyCon nDP_REPR (fsLit "Wrap")
148       enumerationTyCon <- externalTyCon nDP_REPR (fsLit "Enumeration")
149       sum_tcs <- mapM (externalTyCon nDP_REPR)
150                       [mkFastString ("Sum" ++ show i) | i <- [2..mAX_NDP_SUM]]
151
152       let sumTyCons = listArray (2, mAX_NDP_SUM) sum_tcs
153
154       voidVar          <- externalVar nDP_REPR (fsLit "void")
155       mkPRVar          <- externalVar nDP_PARRAY (fsLit "mkPR")
156       mkClosureVar     <- externalVar nDP_CLOSURE (fsLit "mkClosure")
157       applyClosureVar  <- externalVar nDP_CLOSURE (fsLit "$:")
158       mkClosurePVar    <- externalVar nDP_CLOSURE (fsLit "mkClosureP")
159       applyClosurePVar <- externalVar nDP_CLOSURE (fsLit "$:^")
160       replicatePAIntPrimVar <- externalVar nDP_UNBOXED (fsLit "replicatePA_Int#")
161       upToPAIntPrimVar <- externalVar nDP_UNBOXED (fsLit "upToPA_Int#")
162       selectPAIntPrimVar <- externalVar nDP_UNBOXED (fsLit "selectPA_Int#")
163       truesPABoolPrimVar <- externalVar nDP_UNBOXED (fsLit "truesPA_Bool#")
164       lengthPAVar      <- externalVar nDP_PARRAY (fsLit "lengthPA#")
165       replicatePAVar   <- externalVar nDP_PARRAY (fsLit "replicatePA#")
166       emptyPAVar       <- externalVar nDP_PARRAY (fsLit "emptyPA")
167       packPAVar        <- externalVar nDP_PARRAY (fsLit "packPA#")
168
169       combines <- mapM (externalVar nDP_PARRAY)
170                        [mkFastString ("combine" ++ show i ++ "PA#")
171                           | i <- [2..mAX_NDP_COMBINE]]
172       let combinePAVars = listArray (2, mAX_NDP_COMBINE) combines
173
174       liftingContext <- liftM (\u -> mkSysLocal (fsLit "lc") u intPrimTy)
175                               newUnique
176
177       return $ Builtins {
178                  parrayTyCon      = parrayTyCon
179                , paTyCon          = paTyCon
180                , paDataCon        = paDataCon
181                , preprTyCon       = preprTyCon
182                , prTyCon          = prTyCon
183                , prDataCon        = prDataCon
184                , uarrTyCon        = uarrTyCon
185                , voidTyCon        = voidTyCon
186                , wrapTyCon        = wrapTyCon
187                , enumerationTyCon = enumerationTyCon
188                , sumTyCons        = sumTyCons
189                , closureTyCon     = closureTyCon
190                , voidVar          = voidVar
191                , mkPRVar          = mkPRVar
192                , mkClosureVar     = mkClosureVar
193                , applyClosureVar  = applyClosureVar
194                , mkClosurePVar    = mkClosurePVar
195                , applyClosurePVar = applyClosurePVar
196                , replicatePAIntPrimVar = replicatePAIntPrimVar
197                , upToPAIntPrimVar = upToPAIntPrimVar
198                , selectPAIntPrimVar = selectPAIntPrimVar
199                , truesPABoolPrimVar = truesPABoolPrimVar
200                , lengthPAVar      = lengthPAVar
201                , replicatePAVar   = replicatePAVar
202                , emptyPAVar       = emptyPAVar
203                , packPAVar        = packPAVar
204                , combinePAVars    = combinePAVars
205                , liftingContext   = liftingContext
206                }
207
208 initBuiltinVars :: Builtins -> DsM [(Var, Var)]
209 initBuiltinVars _
210   = do
211       uvars <- zipWithM externalVar umods ufs
212       vvars <- zipWithM externalVar vmods vfs
213       cvars <- zipWithM externalVar cmods cfs
214       return $ [(v,v) | v <- map dataConWorkId defaultDataConWorkers]
215                ++ zip (map dataConWorkId cons) cvars
216                ++ zip uvars vvars
217   where
218     (umods, ufs, vmods, vfs) = unzip4 preludeVars
219
220     (cons, cmods, cfs) = unzip3 preludeDataCons
221
222 defaultDataConWorkers :: [DataCon]
223 defaultDataConWorkers = [trueDataCon, falseDataCon, unitDataCon]
224
225 preludeDataCons :: [(DataCon, Module, FastString)]
226 preludeDataCons
227   = [mk_tup n nDP_PRELUDE_TUPLE (mkFastString $ "tup" ++ show n) | n <- [2..3]]
228   where
229     mk_tup n mod name = (tupleCon Boxed n, mod, name)
230
231 preludeVars :: [(Module, FastString, Module, FastString)]
232 preludeVars
233   = [
234       mk gHC_PARR (fsLit "mapP")       nDP_COMBINATORS (fsLit "mapPA")
235     , mk gHC_PARR (fsLit "zipWithP")   nDP_COMBINATORS (fsLit "zipWithPA")
236     , mk gHC_PARR (fsLit "zipP")       nDP_COMBINATORS (fsLit "zipPA")
237     , mk gHC_PARR (fsLit "unzipP")     nDP_COMBINATORS (fsLit "unzipPA")
238     , mk gHC_PARR (fsLit "filterP")    nDP_COMBINATORS (fsLit "filterPA")
239     , mk gHC_PARR (fsLit "lengthP")    nDP_COMBINATORS (fsLit "lengthPA")
240     , mk gHC_PARR (fsLit "replicateP") nDP_COMBINATORS (fsLit "replicatePA")
241     , mk gHC_PARR (fsLit "!:")         nDP_COMBINATORS (fsLit "indexPA")
242     , mk gHC_PARR (fsLit "crossMapP")  nDP_COMBINATORS (fsLit "crossMapPA")
243     , mk gHC_PARR (fsLit "singletonP") nDP_COMBINATORS (fsLit "singletonPA")
244     , mk gHC_PARR (fsLit "concatP")    nDP_COMBINATORS (fsLit "concatPA")
245     , mk gHC_PARR (fsLit "+:+")        nDP_COMBINATORS (fsLit "appPA")
246     , mk gHC_PARR (fsLit "emptyP")     nDP_PARRAY (fsLit "emptyPA")
247
248     , mk nDP_PRELUDE_INT  (fsLit "plus") nDP_PRELUDE_INT (fsLit "plusV")
249     , mk nDP_PRELUDE_INT  (fsLit "minus") nDP_PRELUDE_INT (fsLit "minusV")
250     , mk nDP_PRELUDE_INT  (fsLit "mult")  nDP_PRELUDE_INT (fsLit "multV")
251     , mk nDP_PRELUDE_INT  (fsLit "intDiv")  nDP_PRELUDE_INT (fsLit "intDivV")
252     , mk nDP_PRELUDE_INT  (fsLit "intMod")  nDP_PRELUDE_INT (fsLit "intModV")
253     , mk nDP_PRELUDE_INT  (fsLit "intSquareRoot")  nDP_PRELUDE_INT (fsLit "intSquareRootV")
254     , mk nDP_PRELUDE_INT  (fsLit "sumP")  nDP_PRELUDE_INT (fsLit "sumPA")
255     , mk nDP_PRELUDE_INT  (fsLit "enumFromToP")  nDP_PRELUDE_INT (fsLit "enumFromToPA")
256     , mk nDP_PRELUDE_INT  (fsLit "upToP") nDP_PRELUDE_INT (fsLit "upToPA")
257
258     , mk nDP_PRELUDE_INT  (fsLit "eq") nDP_PRELUDE_INT (fsLit "eqV")
259     , mk nDP_PRELUDE_INT  (fsLit "neq") nDP_PRELUDE_INT (fsLit "neqV")
260     , mk nDP_PRELUDE_INT  (fsLit "le")  nDP_PRELUDE_INT (fsLit "leV")
261     , mk nDP_PRELUDE_INT  (fsLit "lt") nDP_PRELUDE_INT (fsLit "ltV")
262     , mk nDP_PRELUDE_INT  (fsLit "ge") nDP_PRELUDE_INT (fsLit "geV")
263     , mk nDP_PRELUDE_INT  (fsLit "gt")  nDP_PRELUDE_INT (fsLit "gtV")
264
265     , mk nDP_PRELUDE_DOUBLE  (fsLit "plus") nDP_PRELUDE_DOUBLE (fsLit "plusV")
266     , mk nDP_PRELUDE_DOUBLE  (fsLit "minus") nDP_PRELUDE_DOUBLE (fsLit "minusV")
267     , mk nDP_PRELUDE_DOUBLE  (fsLit "mult")  nDP_PRELUDE_DOUBLE (fsLit "multV")
268     , mk nDP_PRELUDE_DOUBLE  (fsLit "divide")  nDP_PRELUDE_DOUBLE (fsLit "divideV")
269     , mk nDP_PRELUDE_DOUBLE  (fsLit  "squareRoot")  nDP_PRELUDE_DOUBLE (fsLit "squareRootV")    
270     , mk nDP_PRELUDE_DOUBLE  (fsLit "sumP")  nDP_PRELUDE_DOUBLE (fsLit "sumPA")
271     , mk nDP_PRELUDE_DOUBLE  (fsLit "minIndexP") 
272          nDP_PRELUDE_DOUBLE  (fsLit "minIndexPA")
273     , mk nDP_PRELUDE_DOUBLE  (fsLit "maxIndexP")
274          nDP_PRELUDE_DOUBLE  (fsLit "maxIndexPA")
275
276     , mk nDP_PRELUDE_DOUBLE  (fsLit "eq") nDP_PRELUDE_DOUBLE (fsLit "eqV")
277     , mk nDP_PRELUDE_DOUBLE  (fsLit "neq") nDP_PRELUDE_DOUBLE (fsLit "neqV")
278     , mk nDP_PRELUDE_DOUBLE  (fsLit "le")  nDP_PRELUDE_DOUBLE (fsLit "leV")
279     , mk nDP_PRELUDE_DOUBLE  (fsLit "lt") nDP_PRELUDE_DOUBLE (fsLit "ltV")
280     , mk nDP_PRELUDE_DOUBLE  (fsLit "ge") nDP_PRELUDE_DOUBLE (fsLit "geV")
281     , mk nDP_PRELUDE_DOUBLE  (fsLit "gt")  nDP_PRELUDE_DOUBLE (fsLit "gtV")
282
283     , mk nDP_PRELUDE_BOOL  (fsLit "andP")  nDP_PRELUDE_BOOL (fsLit "andPA")
284     , mk nDP_PRELUDE_BOOL  (fsLit "orP")  nDP_PRELUDE_BOOL (fsLit "orPA")
285
286     -- FIXME: temporary
287     , mk nDP_PRELUDE_PARR (fsLit "fromPArrayP") nDP_PRELUDE_PARR (fsLit "fromPArrayPA")
288     , mk nDP_PRELUDE_PARR (fsLit "toPArrayP") nDP_PRELUDE_PARR (fsLit "toPArrayPA")
289     , mk nDP_PRELUDE_PARR (fsLit "fromNestedPArrayP") nDP_PRELUDE_PARR (fsLit "fromNestedPArrayPA")
290     , mk nDP_PRELUDE_PARR (fsLit "combineP")    nDP_COMBINATORS (fsLit "combine2PA")
291     ]
292   where
293     mk = (,,,)
294
295 initBuiltinTyCons :: Builtins -> DsM [(Name, TyCon)]
296 initBuiltinTyCons bi
297   = do
298       -- parr <- externalTyCon nDP_PRELUDE_PARR (fsLit "PArr")
299       return $ (tyConName funTyCon, closureTyCon bi)
300              : (parrTyConName,      parrayTyCon bi)
301
302              -- FIXME: temporary
303              : (tyConName $ parrayTyCon bi, parrayTyCon bi)
304
305              : [(tyConName tc, tc) | tc <- defaultTyCons]
306
307 defaultTyCons :: [TyCon]
308 defaultTyCons = [intTyCon, boolTyCon, doubleTyCon]
309
310 initBuiltinDataCons :: Builtins -> [(Name, DataCon)]
311 initBuiltinDataCons _ = [(dataConName dc, dc)| dc <- defaultDataCons]
312
313 defaultDataCons :: [DataCon]
314 defaultDataCons = [trueDataCon, falseDataCon, unitDataCon]
315
316 initBuiltinDicts :: [(Name, Module, FastString)] -> DsM [(Name, Var)]
317 initBuiltinDicts ps
318   = do
319       dicts <- zipWithM externalVar mods fss
320       return $ zip tcs dicts
321   where
322     (tcs, mods, fss) = unzip3 ps
323
324 initBuiltinPAs :: Builtins -> DsM [(Name, Var)]
325 initBuiltinPAs = initBuiltinDicts . builtinPAs
326
327 builtinPAs :: Builtins -> [(Name, Module, FastString)]
328 builtinPAs bi
329   = [
330       mk (tyConName $ closureTyCon bi)  nDP_CLOSURE     (fsLit "dPA_Clo")
331     , mk (tyConName $ voidTyCon bi)     nDP_REPR        (fsLit "dPA_Void")
332     , mk (tyConName $ parrayTyCon bi)   nDP_INSTANCES   (fsLit "dPA_PArray")
333     , mk unitTyConName                  nDP_INSTANCES   (fsLit "dPA_Unit")
334
335     , mk intTyConName                   nDP_INSTANCES   (fsLit "dPA_Int")
336     , mk doubleTyConName                nDP_INSTANCES   (fsLit "dPA_Double")
337     , mk boolTyConName                  nDP_INSTANCES   (fsLit "dPA_Bool")
338     ]
339     ++ tups
340   where
341     mk name mod fs = (name, mod, fs)
342
343     tups = map mk_tup [2..mAX_NDP_PROD]
344     mk_tup n = mk (tyConName $ tupleTyCon Boxed n)
345                   nDP_INSTANCES
346                   (mkFastString $ "dPA_" ++ show n)
347
348 initBuiltinPRs :: Builtins -> DsM [(Name, Var)]
349 initBuiltinPRs = initBuiltinDicts . builtinPRs
350
351 builtinPRs :: Builtins -> [(Name, Module, FastString)]
352 builtinPRs bi =
353   [
354     mk (tyConName unitTyCon)          nDP_REPR      (fsLit "dPR_Unit")
355   , mk (tyConName $ voidTyCon bi)     nDP_REPR      (fsLit "dPR_Void")
356   , mk (tyConName $ wrapTyCon bi)     nDP_REPR      (fsLit "dPR_Wrap")
357   , mk (tyConName $ enumerationTyCon bi) nDP_REPR   (fsLit "dPR_Enumeration")
358   , mk (tyConName $ closureTyCon bi)  nDP_CLOSURE   (fsLit "dPR_Clo")
359
360     -- temporary
361   , mk intTyConName          nDP_INSTANCES (fsLit "dPR_Int")
362   , mk doubleTyConName       nDP_INSTANCES (fsLit "dPR_Double")
363   ]
364
365   ++ map mk_sum  [2..mAX_NDP_SUM]
366   ++ map mk_prod [2..mAX_NDP_PROD]
367   where
368     mk name mod fs = (name, mod, fs)
369
370     mk_sum n = (tyConName $ sumTyCon n bi, nDP_REPR,
371                 mkFastString ("dPR_Sum" ++ show n))
372
373     mk_prod n = (tyConName $ prodTyCon n bi, nDP_REPR,
374                  mkFastString ("dPR_" ++ show n))
375
376 initBuiltinBoxedTyCons :: Builtins -> DsM [(Name, TyCon)]
377 initBuiltinBoxedTyCons = return . builtinBoxedTyCons
378
379 builtinBoxedTyCons :: Builtins -> [(Name, TyCon)]
380 builtinBoxedTyCons _ =
381   [(tyConName intPrimTyCon, intTyCon)]
382
383 externalVar :: Module -> FastString -> DsM Var
384 externalVar mod fs
385   = dsLookupGlobalId =<< lookupOrig mod (mkVarOccFS fs)
386
387 externalTyCon :: Module -> FastString -> DsM TyCon
388 externalTyCon mod fs
389   = dsLookupTyCon =<< lookupOrig mod (mkOccNameFS tcName fs)
390
391 unitTyConName :: Name
392 unitTyConName = tyConName unitTyCon
393
394
395 primMethod :: TyCon -> String -> DsM (Maybe Var)
396 primMethod tycon method
397   | Just suffix <- lookupNameEnv prim_ty_cons (tyConName tycon)
398   = liftM Just
399   $ dsLookupGlobalId =<< lookupOrig nDP_UNBOXED (mkVarOcc $ method ++ suffix)
400
401   | otherwise = return Nothing
402
403 primPArray :: TyCon -> DsM (Maybe TyCon)
404 primPArray tycon
405   | Just suffix <- lookupNameEnv prim_ty_cons (tyConName tycon)
406   = liftM Just
407   $ dsLookupTyCon =<< lookupOrig nDP_UNBOXED (mkOccName tcName $ "PArray" ++ suffix)
408
409   | otherwise = return Nothing
410
411 prim_ty_cons :: NameEnv String
412 prim_ty_cons = mkNameEnv [mk_prim intPrimTyCon]
413   where
414     mk_prim tycon = (tyConName tycon, '_' : getOccString tycon)
415