Use n-ary sums and products for NDP's generic representation
[ghc-hetmet.git] / compiler / vectorise / VectBuiltIn.hs
1 module VectBuiltIn (
2   Builtins(..), sumTyCon, prodTyCon,
3   initBuiltins, initBuiltinTyCons, initBuiltinPAs, initBuiltinPRs
4 ) where
5
6 #include "HsVersions.h"
7
8 import DsMonad
9 import IfaceEnv        ( lookupOrig )
10
11 import Module          ( Module )
12 import DataCon         ( DataCon )
13 import TyCon           ( TyCon, tyConName, tyConDataCons )
14 import Var             ( Var )
15 import Id              ( mkSysLocal )
16 import Name            ( Name )
17 import OccName         ( mkVarOccFS, mkOccNameFS, tcName )
18
19 import TypeRep         ( funTyCon )
20 import TysPrim         ( intPrimTy )
21 import TysWiredIn      ( unitTyCon, tupleTyCon, intTyConName )
22 import PrelNames
23 import BasicTypes      ( Boxity(..) )
24
25 import FastString
26 import Outputable
27
28 import Data.Array
29 import Control.Monad   ( liftM, zipWithM )
30
31 mAX_NDP_PROD :: Int
32 mAX_NDP_PROD = 3
33
34 mAX_NDP_SUM :: Int
35 mAX_NDP_SUM = 3
36
37 data Builtins = Builtins {
38                   parrayTyCon      :: TyCon
39                 , paTyCon          :: TyCon
40                 , paDataCon        :: DataCon
41                 , preprTyCon       :: TyCon
42                 , prTyCon          :: TyCon
43                 , prDataCon        :: DataCon
44                 , embedTyCon       :: TyCon
45                 , embedDataCon     :: DataCon
46                 , sumTyCons        :: Array Int TyCon
47                 , closureTyCon     :: TyCon
48                 , mkClosureVar     :: Var
49                 , applyClosureVar  :: Var
50                 , mkClosurePVar    :: Var
51                 , applyClosurePVar :: Var
52                 , lengthPAVar      :: Var
53                 , replicatePAVar   :: Var
54                 , emptyPAVar       :: Var
55                 -- , packPAVar        :: Var
56                 -- , combinePAVar     :: Var
57                 , intEqPAVar       :: Var
58                 , liftingContext   :: Var
59                 }
60
61 sumTyCon :: Int -> Builtins -> TyCon
62 sumTyCon n bi
63   | n >= 2 && n <= mAX_NDP_SUM = sumTyCons bi ! n
64   | otherwise = pprPanic "sumTyCon" (ppr n)
65
66 prodTyCon :: Int -> Builtins -> TyCon
67 prodTyCon n bi
68   | n >= 2 && n <= mAX_NDP_PROD = tupleTyCon Boxed n
69   | otherwise = pprPanic "prodTyCon" (ppr n)
70
71
72 initBuiltins :: DsM Builtins
73 initBuiltins
74   = do
75       parrayTyCon  <- dsLookupTyCon parrayTyConName
76       paTyCon      <- dsLookupTyCon paTyConName
77       let [paDataCon] = tyConDataCons paTyCon
78       preprTyCon   <- dsLookupTyCon preprTyConName
79       prTyCon      <- dsLookupTyCon prTyConName
80       let [prDataCon] = tyConDataCons prTyCon
81       embedTyCon   <- dsLookupTyCon embedTyConName
82       let [embedDataCon] = tyConDataCons embedTyCon
83       closureTyCon <- dsLookupTyCon closureTyConName
84
85       sum_tcs <- mapM (lookupExternalTyCon nDP_REPR)
86                       [mkFastString ("Sum" ++ show i) | i <- [2..mAX_NDP_SUM]]
87
88       let sumTyCons = listArray (2, mAX_NDP_SUM) sum_tcs
89
90       mkClosureVar     <- dsLookupGlobalId mkClosureName
91       applyClosureVar  <- dsLookupGlobalId applyClosureName
92       mkClosurePVar    <- dsLookupGlobalId mkClosurePName
93       applyClosurePVar <- dsLookupGlobalId applyClosurePName
94       lengthPAVar      <- dsLookupGlobalId lengthPAName
95       replicatePAVar   <- dsLookupGlobalId replicatePAName
96       emptyPAVar       <- dsLookupGlobalId emptyPAName
97       -- packPAVar        <- dsLookupGlobalId packPAName
98       -- combinePAVar     <- dsLookupGlobalId combinePAName
99       intEqPAVar       <- dsLookupGlobalId intEqPAName
100
101       liftingContext <- liftM (\u -> mkSysLocal FSLIT("lc") u intPrimTy)
102                               newUnique
103
104       return $ Builtins {
105                  parrayTyCon      = parrayTyCon
106                , paTyCon          = paTyCon
107                , paDataCon        = paDataCon
108                , preprTyCon       = preprTyCon
109                , prTyCon          = prTyCon
110                , prDataCon        = prDataCon
111                , embedTyCon       = embedTyCon
112                , embedDataCon     = embedDataCon
113                , sumTyCons        = sumTyCons
114                , closureTyCon     = closureTyCon
115                , mkClosureVar     = mkClosureVar
116                , applyClosureVar  = applyClosureVar
117                , mkClosurePVar    = mkClosurePVar
118                , applyClosurePVar = applyClosurePVar
119                , lengthPAVar      = lengthPAVar
120                , replicatePAVar   = replicatePAVar
121                , emptyPAVar       = emptyPAVar
122                -- , packPAVar        = packPAVar
123                -- , combinePAVar     = combinePAVar
124                , intEqPAVar       = intEqPAVar
125                , liftingContext   = liftingContext
126                }
127
128 initBuiltinTyCons :: DsM [(Name, TyCon)]
129 initBuiltinTyCons
130   = do
131       vects <- sequence vs
132       return (zip origs vects)
133   where
134     (origs, vs) = unzip builtinTyCons
135
136 builtinTyCons :: [(Name, DsM TyCon)]
137 builtinTyCons = [(tyConName funTyCon, dsLookupTyCon closureTyConName)]
138
139 initBuiltinDicts :: [(Name, Module, FastString)] -> DsM [(Name, Var)]
140 initBuiltinDicts ps
141   = do
142       dicts <- zipWithM lookupExternalVar mods fss
143       return $ zip tcs dicts
144   where
145     (tcs, mods, fss) = unzip3 ps
146
147 initBuiltinPAs = initBuiltinDicts builtinPAs
148
149 builtinPAs :: [(Name, Module, FastString)]
150 builtinPAs = [
151                mk closureTyConName  nDP_CLOSURE   FSLIT("dPA_Clo")
152              , mk unitTyConName     nDP_INSTANCES FSLIT("dPA_Unit")
153
154              , mk intTyConName      nDP_INSTANCES FSLIT("dPA_Int")
155              ]
156              ++ tups
157   where
158     mk name mod fs = (name, mod, fs)
159
160     tups = map mk_tup [2..3]
161     mk_tup n = mk (tyConName $ tupleTyCon Boxed n)
162                   nDP_INSTANCES
163                   (mkFastString $ "dPA_" ++ show n)
164
165 initBuiltinPRs = initBuiltinDicts . builtinPRs
166
167 builtinPRs :: Builtins -> [(Name, Module, FastString)]
168 builtinPRs bi =
169   [
170     mk (tyConName unitTyCon) nDP_REPR      FSLIT("dPR_Unit")
171   , mk embedTyConName        nDP_REPR      FSLIT("dPR_Embed")
172   , mk closureTyConName      nDP_CLOSURE   FSLIT("dPR_Clo")
173
174     -- temporary
175   , mk intTyConName          nDP_INSTANCES FSLIT("dPR_Int")
176   ]
177
178   ++ map mk_sum  [2..mAX_NDP_SUM]
179   ++ map mk_prod [2..mAX_NDP_PROD]
180   where
181     mk name mod fs = (name, mod, fs)
182
183     mk_sum n = (tyConName $ sumTyCon n bi, nDP_REPR,
184                 mkFastString ("dPR_Sum" ++ show n))
185
186     mk_prod n = (tyConName $ prodTyCon n bi, nDP_REPR,
187                  mkFastString ("dPR_" ++ show n))
188
189 lookupExternalVar :: Module -> FastString -> DsM Var
190 lookupExternalVar mod fs
191   = dsLookupGlobalId =<< lookupOrig mod (mkVarOccFS fs)
192
193 lookupExternalTyCon :: Module -> FastString -> DsM TyCon
194 lookupExternalTyCon mod fs
195   = dsLookupTyCon =<< lookupOrig mod (mkOccNameFS tcName fs)
196
197 unitTyConName = tyConName unitTyCon
198