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 #include "HsVersions.h"
  12
  13 import DsMonad
  14 import IfaceEnv        ( lookupOrig )
  15
  16 import Module
  17 import DataCon         ( DataCon, dataConName, dataConWorkId )
  18 import TyCon           ( TyCon, tyConName, tyConDataCons )
  19 import Var             ( Var )
  20 import Id              ( mkSysLocal )
  21 import Name            ( Name, getOccString )
  22 import NameEnv
  23 import OccName
  24
  25 import TypeRep         ( funTyCon )
  26 import Type            ( Type, mkTyConApp )
  27 import TysPrim
  28 import TysWiredIn      ( unitTyCon, unitDataCon,
  29                          tupleTyCon, tupleCon,
  30                          intTyCon, intTyConName, intTy,
  31                          doubleTyCon, doubleTyConName,
  32                          boolTyCon, boolTyConName, trueDataCon, falseDataCon,
  33                          parrTyConName )
  34 import PrelNames       ( gHC_PARR )
  35 import BasicTypes      ( Boxity(..) )
  36
  37 import FastString
  38 import Outputable
  39
  40 import Data.Array
  41 import Control.Monad   ( liftM, zipWithM )
  42 import Data.List       ( unzip4 )
  43
  44 mAX_NDP_PROD :: Int
  45 mAX_NDP_PROD = 5
  46
  47 mAX_NDP_SUM :: Int
  48 mAX_NDP_SUM = 3
  49
  50 mAX_NDP_COMBINE :: Int
  51 mAX_NDP_COMBINE = 2
  52
  53 mkNDPModule :: FastString -> Module
  54 mkNDPModule m = mkModule ndpPackageId (mkModuleNameFS m)
  55
  56 nDP_UARR,
  57   nDP_PARRAY,
  58   nDP_REPR,
  59   nDP_CLOSURE,
  60   nDP_UNBOXED,
  61   nDP_INSTANCES,
  62   nDP_COMBINATORS,
  63   nDP_PRELUDE_PARR,
  64   nDP_PRELUDE_INT,
  65   nDP_PRELUDE_DOUBLE,
  66   nDP_PRELUDE_TUPLE :: Module
  67
  68 nDP_UARR        = mkNDPModule FSLIT("Data.Array.Parallel.Unlifted.Flat.UArr")
  69 nDP_PARRAY      = mkNDPModule FSLIT("Data.Array.Parallel.Lifted.PArray")
  70 nDP_REPR        = mkNDPModule FSLIT("Data.Array.Parallel.Lifted.Repr")
  71 nDP_CLOSURE     = mkNDPModule FSLIT("Data.Array.Parallel.Lifted.Closure")
  72 nDP_UNBOXED     = mkNDPModule FSLIT("Data.Array.Parallel.Lifted.Unboxed")
  73 nDP_INSTANCES   = mkNDPModule FSLIT("Data.Array.Parallel.Lifted.Instances")
  74 nDP_COMBINATORS = mkNDPModule FSLIT("Data.Array.Parallel.Lifted.Combinators")
  75
  76 nDP_PRELUDE_PARR = mkNDPModule FSLIT("Data.Array.Parallel.Prelude.Base.PArr")
  77 nDP_PRELUDE_INT  = mkNDPModule FSLIT("Data.Array.Parallel.Prelude.Base.Int")
  78 nDP_PRELUDE_DOUBLE = mkNDPModule FSLIT("Data.Array.Parallel.Prelude.Base.Double")
  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
 247     , mk nDP_PRELUDE_INT  FSLIT("plus") nDP_PRELUDE_INT FSLIT("plusV")
 248     , mk nDP_PRELUDE_INT  FSLIT("minus") nDP_PRELUDE_INT FSLIT("minusV")
 249     , mk nDP_PRELUDE_INT  FSLIT("mult")  nDP_PRELUDE_INT FSLIT("multV")
 250     , mk nDP_PRELUDE_INT  FSLIT("intDiv")  nDP_PRELUDE_INT FSLIT("intDivV")
 251     , mk nDP_PRELUDE_INT  FSLIT("sumP")  nDP_PRELUDE_INT FSLIT("sumPA")
 252     , mk nDP_PRELUDE_INT  FSLIT("upToP") nDP_PRELUDE_INT FSLIT("upToPA")
 253
 254     , mk nDP_PRELUDE_INT  FSLIT("eq") nDP_PRELUDE_INT FSLIT("eqV")
 255     , mk nDP_PRELUDE_INT  FSLIT("neq") nDP_PRELUDE_INT FSLIT("neqV")
 256     , mk nDP_PRELUDE_INT  FSLIT("le")  nDP_PRELUDE_INT FSLIT("leV")
 257     , mk nDP_PRELUDE_INT  FSLIT("lt") nDP_PRELUDE_INT FSLIT("ltV")
 258     , mk nDP_PRELUDE_INT  FSLIT("ge") nDP_PRELUDE_INT FSLIT("geV")
 259     , mk nDP_PRELUDE_INT  FSLIT("gt")  nDP_PRELUDE_INT FSLIT("gtV")
 260
 261     , mk nDP_PRELUDE_DOUBLE  FSLIT("plus") nDP_PRELUDE_DOUBLE FSLIT("plusV")
 262     , mk nDP_PRELUDE_DOUBLE  FSLIT("minus") nDP_PRELUDE_DOUBLE FSLIT("minusV")
 263     , mk nDP_PRELUDE_DOUBLE  FSLIT("mult")  nDP_PRELUDE_DOUBLE FSLIT("multV")
 264     , mk nDP_PRELUDE_DOUBLE  FSLIT("divide")  nDP_PRELUDE_DOUBLE FSLIT("divideV")
 265     , mk nDP_PRELUDE_DOUBLE  FSLIT("sumP")  nDP_PRELUDE_DOUBLE FSLIT("sumPA")
 266     , mk nDP_PRELUDE_DOUBLE  FSLIT("minIndexP")
 267          nDP_PRELUDE_DOUBLE  FSLIT("minIndexPA")
 268     , mk nDP_PRELUDE_DOUBLE  FSLIT("maxIndexP")
 269          nDP_PRELUDE_DOUBLE  FSLIT("maxIndexPA")
 270
 271     , mk nDP_PRELUDE_DOUBLE  FSLIT("eq") nDP_PRELUDE_DOUBLE FSLIT("eqV")
 272     , mk nDP_PRELUDE_DOUBLE  FSLIT("neq") nDP_PRELUDE_DOUBLE FSLIT("neqV")
 273     , mk nDP_PRELUDE_DOUBLE  FSLIT("le")  nDP_PRELUDE_DOUBLE FSLIT("leV")
 274     , mk nDP_PRELUDE_DOUBLE  FSLIT("lt") nDP_PRELUDE_DOUBLE FSLIT("ltV")
 275     , mk nDP_PRELUDE_DOUBLE  FSLIT("ge") nDP_PRELUDE_DOUBLE FSLIT("geV")
 276     , mk nDP_PRELUDE_DOUBLE  FSLIT("gt")  nDP_PRELUDE_DOUBLE FSLIT("gtV")
 277
 278     -- FIXME: temporary
 279     , mk nDP_PRELUDE_PARR FSLIT("fromPArrayP") nDP_PRELUDE_PARR FSLIT("fromPArrayPA")
 280     , mk nDP_PRELUDE_PARR FSLIT("toPArrayP") nDP_PRELUDE_PARR FSLIT("toPArrayPA")
 281     , mk nDP_PRELUDE_PARR FSLIT("fromNestedPArrayP") nDP_PRELUDE_PARR FSLIT("fromNestedPArrayPA")
 282     ]
 283   where
 284     mk = (,,,)
 285
 286 initBuiltinTyCons :: Builtins -> DsM [(Name, TyCon)]
 287 initBuiltinTyCons bi
 288   = do
 289       -- parr <- externalTyCon nDP_PRELUDE_PARR FSLIT("PArr")
 290       return $ (tyConName funTyCon, closureTyCon bi)
 291              : (parrTyConName,      parrayTyCon bi)
 292
 293              -- FIXME: temporary
 294              : (tyConName $ parrayTyCon bi, parrayTyCon bi)
 295
 296              : [(tyConName tc, tc) | tc <- defaultTyCons]
 297
 298 defaultTyCons :: [TyCon]
 299 defaultTyCons = [intTyCon, boolTyCon, doubleTyCon]
 300
 301 initBuiltinDataCons :: Builtins -> [(Name, DataCon)]
 302 initBuiltinDataCons _ = [(dataConName dc, dc)| dc <- defaultDataCons]
 303
 304 defaultDataCons :: [DataCon]
 305 defaultDataCons = [trueDataCon, falseDataCon, unitDataCon]
 306
 307 initBuiltinDicts :: [(Name, Module, FastString)] -> DsM [(Name, Var)]
 308 initBuiltinDicts ps
 309   = do
 310       dicts <- zipWithM externalVar mods fss
 311       return $ zip tcs dicts
 312   where
 313     (tcs, mods, fss) = unzip3 ps
 314
 315 initBuiltinPAs :: Builtins -> DsM [(Name, Var)]
 316 initBuiltinPAs = initBuiltinDicts . builtinPAs
 317
 318 builtinPAs :: Builtins -> [(Name, Module, FastString)]
 319 builtinPAs bi
 320   = [
 321       mk (tyConName $ closureTyCon bi)  nDP_CLOSURE     FSLIT("dPA_Clo")
 322     , mk (tyConName $ voidTyCon bi)     nDP_REPR        FSLIT("dPA_Void")
 323     , mk (tyConName $ parrayTyCon bi)   nDP_INSTANCES   FSLIT("dPA_PArray")
 324     , mk unitTyConName                  nDP_INSTANCES   FSLIT("dPA_Unit")
 325
 326     , mk intTyConName                   nDP_INSTANCES   FSLIT("dPA_Int")
 327     , mk doubleTyConName                nDP_INSTANCES   FSLIT("dPA_Double")
 328     , mk boolTyConName                  nDP_INSTANCES   FSLIT("dPA_Bool")
 329     ]
 330     ++ tups
 331   where
 332     mk name mod fs = (name, mod, fs)
 333
 334     tups = map mk_tup [2..mAX_NDP_PROD]
 335     mk_tup n = mk (tyConName $ tupleTyCon Boxed n)
 336                   nDP_INSTANCES
 337                   (mkFastString $ "dPA_" ++ show n)
 338
 339 initBuiltinPRs :: Builtins -> DsM [(Name, Var)]
 340 initBuiltinPRs = initBuiltinDicts . builtinPRs
 341
 342 builtinPRs :: Builtins -> [(Name, Module, FastString)]
 343 builtinPRs bi =
 344   [
 345     mk (tyConName unitTyCon)          nDP_REPR      FSLIT("dPR_Unit")
 346   , mk (tyConName $ voidTyCon bi)     nDP_REPR      FSLIT("dPR_Void")
 347   , mk (tyConName $ wrapTyCon bi)     nDP_REPR      FSLIT("dPR_Wrap")
 348   , mk (tyConName $ enumerationTyCon bi) nDP_REPR   FSLIT("dPR_Enumeration")
 349   , mk (tyConName $ closureTyCon bi)  nDP_CLOSURE   FSLIT("dPR_Clo")
 350
 351     -- temporary
 352   , mk intTyConName          nDP_INSTANCES FSLIT("dPR_Int")
 353   , mk doubleTyConName       nDP_INSTANCES FSLIT("dPR_Double")
 354   ]
 355
 356   ++ map mk_sum  [2..mAX_NDP_SUM]
 357   ++ map mk_prod [2..mAX_NDP_PROD]
 358   where
 359     mk name mod fs = (name, mod, fs)
 360
 361     mk_sum n = (tyConName $ sumTyCon n bi, nDP_REPR,
 362                 mkFastString ("dPR_Sum" ++ show n))
 363
 364     mk_prod n = (tyConName $ prodTyCon n bi, nDP_REPR,
 365                  mkFastString ("dPR_" ++ show n))
 366
 367 initBuiltinBoxedTyCons :: Builtins -> DsM [(Name, TyCon)]
 368 initBuiltinBoxedTyCons = return . builtinBoxedTyCons
 369
 370 builtinBoxedTyCons :: Builtins -> [(Name, TyCon)]
 371 builtinBoxedTyCons _ =
 372   [(tyConName intPrimTyCon, intTyCon)]
 373
 374 externalVar :: Module -> FastString -> DsM Var
 375 externalVar mod fs
 376   = dsLookupGlobalId =<< lookupOrig mod (mkVarOccFS fs)
 377
 378 externalTyCon :: Module -> FastString -> DsM TyCon
 379 externalTyCon mod fs
 380   = dsLookupTyCon =<< lookupOrig mod (mkOccNameFS tcName fs)
 381
 382 unitTyConName :: Name
 383 unitTyConName = tyConName unitTyCon
 384
 385
 386 primMethod :: TyCon -> String -> DsM (Maybe Var)
 387 primMethod tycon method
 388   | Just suffix <- lookupNameEnv prim_ty_cons (tyConName tycon)
 389   = liftM Just
 390   $ dsLookupGlobalId =<< lookupOrig nDP_UNBOXED (mkVarOcc $ method ++ suffix)
 391
 392   | otherwise = return Nothing
 393
 394 primPArray :: TyCon -> DsM (Maybe TyCon)
 395 primPArray tycon
 396   | Just suffix <- lookupNameEnv prim_ty_cons (tyConName tycon)
 397   = liftM Just
 398   $ dsLookupTyCon =<< lookupOrig nDP_UNBOXED (mkOccName tcName $ "PArray" ++ suffix)
 399
 400   | otherwise = return Nothing
 401
 402 prim_ty_cons :: NameEnv String
 403 prim_ty_cons = mkNameEnv [mk_prim intPrimTyCon]
 404   where
 405     mk_prim tycon = (tyConName tycon, '_' : getOccString tycon)
 406