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