compiler/vectorise/Vectorise/Builtins/Initialise.hs

   1
   2
   3 module Vectorise.Builtins.Initialise (
   4         -- * Initialisation
   5         initBuiltins, initBuiltinVars, initBuiltinTyCons, initBuiltinDataCons,
   6         initBuiltinPAs, initBuiltinPRs,
   7         initBuiltinBoxedTyCons, initBuiltinScalars,
   8 ) where
   9 import Vectorise.Builtins.Base
  10 import Vectorise.Builtins.Modules
  11 import Vectorise.Builtins.Prelude
  12
  13 import BasicTypes
  14 import PrelNames
  15 import TysPrim
  16 import DsMonad
  17 import IfaceEnv
  18 import InstEnv
  19 import TysWiredIn
  20 import DataCon
  21 import TyCon
  22 import Class
  23 import CoreSyn
  24 import Type
  25 import Name
  26 import Module
  27 import Var
  28 import Id
  29 import FastString
  30 import Outputable
  31
  32 import Control.Monad
  33 import Data.Array
  34 import Data.List
  35
  36 -- | Create the initial map of builtin types and functions.
  37 initBuiltins
  38         :: PackageId    -- ^ package id the builtins are in, eg dph-common
  39         -> DsM Builtins
  40
  41 initBuiltins pkg
  42  = do mapM_ load dph_Orphans
  43
  44       -- From dph-common:Data.Array.Parallel.Lifted.PArray
  45       parrayTyCon       <- externalTyCon        dph_PArray      (fsLit "PArray")
  46       let [parrayDataCon] = tyConDataCons parrayTyCon
  47
  48       pdataTyCon        <- externalTyCon        dph_PArray      (fsLit "PData")
  49       paClass           <- externalClass        dph_PArray      (fsLit "PA")
  50       let paTyCon     = classTyCon paClass
  51           [paDataCon] = tyConDataCons paTyCon
  52           paPRSel     = classSCSelId paClass 0
  53
  54       preprTyCon        <- externalTyCon        dph_PArray      (fsLit "PRepr")
  55       prClass           <- externalClass        dph_PArray      (fsLit "PR")
  56       let prTyCon     = classTyCon prClass
  57           [prDataCon] = tyConDataCons prTyCon
  58
  59       closureTyCon      <- externalTyCon dph_Closure            (fsLit ":->")
  60
  61       -- From dph-common:Data.Array.Parallel.Lifted.Repr
  62       voidTyCon         <- externalTyCon        dph_Repr        (fsLit "Void")
  63       wrapTyCon         <- externalTyCon        dph_Repr        (fsLit "Wrap")
  64
  65       -- From dph-common:Data.Array.Parallel.Lifted.Unboxed
  66       sel_tys           <- mapM (externalType dph_Unboxed)
  67                                 (numbered "Sel" 2 mAX_DPH_SUM)
  68
  69       sel_replicates    <- mapM (externalFun dph_Unboxed)
  70                                 (numbered_hash "replicateSel" 2 mAX_DPH_SUM)
  71
  72       sel_picks         <- mapM (externalFun dph_Unboxed)
  73                                 (numbered_hash "pickSel" 2 mAX_DPH_SUM)
  74
  75       sel_tags          <- mapM (externalFun dph_Unboxed)
  76                                 (numbered "tagsSel" 2 mAX_DPH_SUM)
  77
  78       sel_els           <- mapM mk_elements
  79                                 [(i,j) | i <- [2..mAX_DPH_SUM], j <- [0..i-1]]
  80
  81       sum_tcs           <- mapM (externalTyCon dph_Repr)
  82                                 (numbered "Sum" 2 mAX_DPH_SUM)
  83
  84       let selTys        = listArray (2, mAX_DPH_SUM) sel_tys
  85           selReplicates = listArray (2, mAX_DPH_SUM) sel_replicates
  86           selPicks      = listArray (2, mAX_DPH_SUM) sel_picks
  87           selTagss      = listArray (2, mAX_DPH_SUM) sel_tags
  88           selEls        = array     ((2,0), (mAX_DPH_SUM, mAX_DPH_SUM)) sel_els
  89           sumTyCons     = listArray (2, mAX_DPH_SUM) sum_tcs
  90
  91
  92       voidVar          <- externalVar dph_Repr          (fsLit "void")
  93       pvoidVar         <- externalVar dph_Repr          (fsLit "pvoid")
  94       fromVoidVar      <- externalVar dph_Repr          (fsLit "fromVoid")
  95       punitVar         <- externalVar dph_Repr          (fsLit "punit")
  96       closureVar       <- externalVar dph_Closure       (fsLit "closure")
  97       applyVar         <- externalVar dph_Closure       (fsLit "$:")
  98       liftedClosureVar <- externalVar dph_Closure       (fsLit "liftedClosure")
  99       liftedApplyVar   <- externalVar dph_Closure       (fsLit "liftedApply")
 100       replicatePDVar   <- externalVar dph_PArray        (fsLit "replicatePD")
 101       emptyPDVar       <- externalVar dph_PArray        (fsLit "emptyPD")
 102       packByTagPDVar   <- externalVar dph_PArray        (fsLit "packByTagPD")
 103
 104       combines          <- mapM (externalVar dph_PArray)
 105                                 [mkFastString ("combine" ++ show i ++ "PD")
 106                                         | i <- [2..mAX_DPH_COMBINE]]
 107       let combinePDVars = listArray (2, mAX_DPH_COMBINE) combines
 108
 109       scalarClass       <- externalClass dph_PArray     (fsLit "Scalar")
 110       scalar_map        <- externalVar  dph_Scalar      (fsLit "scalar_map")
 111       scalar_zip2       <- externalVar  dph_Scalar      (fsLit "scalar_zipWith")
 112       scalar_zips       <- mapM (externalVar dph_Scalar)
 113                                 (numbered "scalar_zipWith" 3 mAX_DPH_SCALAR_ARGS)
 114
 115       let scalarZips    = listArray (1, mAX_DPH_SCALAR_ARGS)
 116                                  (scalar_map : scalar_zip2 : scalar_zips)
 117
 118       closures          <- mapM (externalVar dph_Closure)
 119                                 (numbered "closure" 1 mAX_DPH_SCALAR_ARGS)
 120
 121       let closureCtrFuns = listArray (1, mAX_DPH_COMBINE) closures
 122
 123       liftingContext    <- liftM (\u -> mkSysLocal (fsLit "lc") u intPrimTy)
 124                                 newUnique
 125
 126       return   $ Builtins
 127                { dphModules       = mods
 128                , parrayTyCon      = parrayTyCon
 129                , parrayDataCon    = parrayDataCon
 130                , pdataTyCon       = pdataTyCon
 131                , paClass          = paClass
 132                , paTyCon          = paTyCon
 133                , paDataCon        = paDataCon
 134                , paPRSel          = paPRSel
 135                , preprTyCon       = preprTyCon
 136                , prClass          = prClass
 137                , prTyCon          = prTyCon
 138                , prDataCon        = prDataCon
 139                , voidTyCon        = voidTyCon
 140                , wrapTyCon        = wrapTyCon
 141                , selTys           = selTys
 142                , selReplicates    = selReplicates
 143                , selPicks         = selPicks
 144                , selTagss         = selTagss
 145                , selEls           = selEls
 146                , sumTyCons        = sumTyCons
 147                , closureTyCon     = closureTyCon
 148                , voidVar          = voidVar
 149                , pvoidVar         = pvoidVar
 150                , fromVoidVar      = fromVoidVar
 151                , punitVar         = punitVar
 152                , closureVar       = closureVar
 153                , applyVar         = applyVar
 154                , liftedClosureVar = liftedClosureVar
 155                , liftedApplyVar   = liftedApplyVar
 156                , replicatePDVar   = replicatePDVar
 157                , emptyPDVar       = emptyPDVar
 158                , packByTagPDVar   = packByTagPDVar
 159                , combinePDVars    = combinePDVars
 160                , scalarClass      = scalarClass
 161                , scalarZips       = scalarZips
 162                , closureCtrFuns   = closureCtrFuns
 163                , liftingContext   = liftingContext
 164                }
 165   where
 166     mods@(Modules {
 167                dph_PArray         = dph_PArray
 168              , dph_Repr           = dph_Repr
 169              , dph_Closure        = dph_Closure
 170              , dph_Scalar         = dph_Scalar
 171              , dph_Unboxed        = dph_Unboxed
 172              })
 173       = dph_Modules pkg
 174
 175     load get_mod = dsLoadModule doc mod
 176       where
 177         mod = get_mod mods
 178         doc = ppr mod <+> ptext (sLit "is a DPH module")
 179
 180     -- Make a list of numbered strings in some range, eg foo3, foo4, foo5
 181     numbered :: String -> Int -> Int -> [FastString]
 182     numbered pfx m n = [mkFastString (pfx ++ show i) | i <- [m..n]]
 183
 184     numbered_hash :: String -> Int -> Int -> [FastString]
 185     numbered_hash pfx m n = [mkFastString (pfx ++ show i ++ "#") | i <- [m..n]]
 186
 187     mk_elements :: (Int, Int) -> DsM ((Int, Int), CoreExpr)
 188     mk_elements (i,j)
 189       = do
 190           v <- externalVar dph_Unboxed
 191              $ mkFastString ("elementsSel" ++ show i ++ "_" ++ show j ++ "#")
 192           return ((i,j), Var v)
 193
 194 -- | Get the mapping of names in the Prelude to names in the DPH library.
 195 --
 196 initBuiltinVars :: Bool   -- FIXME
 197                 -> Builtins -> DsM [(Var, Var)]
 198 initBuiltinVars compilingDPH (Builtins { dphModules = mods })
 199   = do
 200       uvars <- zipWithM externalVar umods ufs
 201       vvars <- zipWithM externalVar vmods vfs
 202       cvars <- zipWithM externalVar cmods cfs
 203       return $ [(v,v) | v <- map dataConWorkId defaultDataConWorkers]
 204                ++ zip (map dataConWorkId cons) cvars
 205                ++ zip uvars vvars
 206   where
 207     (umods, ufs, vmods, vfs) = if compilingDPH then ([], [], [], []) else unzip4 (preludeVars mods)
 208     (cons, cmods, cfs)       = unzip3 (preludeDataCons mods)
 209
 210     defaultDataConWorkers :: [DataCon]
 211     defaultDataConWorkers = [trueDataCon, falseDataCon, unitDataCon]
 212
 213
 214 preludeDataCons :: Modules -> [(DataCon, Module, FastString)]
 215 preludeDataCons (Modules { dph_Prelude_Tuple = dph_Prelude_Tuple })
 216   = [mk_tup n dph_Prelude_Tuple (mkFastString $ "tup" ++ show n) | n <- [2..3]]
 217   where
 218     mk_tup n mod name = (tupleCon Boxed n, mod, name)
 219
 220
 221 -- | Get a list of names to `TyCon`s in the mock prelude.
 222 initBuiltinTyCons :: Builtins -> DsM [(Name, TyCon)]
 223 initBuiltinTyCons bi
 224   = do
 225       -- parr <- externalTyCon dph_Prelude_PArr (fsLit "PArr")
 226       dft_tcs <- defaultTyCons
 227       return $ (tyConName funTyCon, closureTyCon bi)
 228              : (parrTyConName,      parrayTyCon bi)
 229
 230              -- FIXME: temporary
 231              : (tyConName $ parrayTyCon bi, parrayTyCon bi)
 232
 233              : [(tyConName tc, tc) | tc <- dft_tcs]
 234
 235   where defaultTyCons :: DsM [TyCon]
 236         defaultTyCons
 237          = do   word8 <- dsLookupTyCon word8TyConName
 238                 return [intTyCon, boolTyCon, doubleTyCon, word8]
 239
 240
 241 -- | Get a list of names to `DataCon`s in the mock prelude.
 242 initBuiltinDataCons :: Builtins -> [(Name, DataCon)]
 243 initBuiltinDataCons _
 244   = [(dataConName dc, dc)| dc <- defaultDataCons]
 245   where defaultDataCons :: [DataCon]
 246         defaultDataCons = [trueDataCon, falseDataCon, unitDataCon]
 247
 248
 249 -- | Get the names of all buildin instance functions for the PA class.
 250 initBuiltinPAs :: Builtins -> (InstEnv, InstEnv) -> DsM [(Name, Var)]
 251 initBuiltinPAs (Builtins { dphModules = mods }) insts
 252   = liftM (initBuiltinDicts insts) (externalClass (dph_PArray mods) (fsLit "PA"))
 253
 254
 255 -- | Get the names of all builtin instance functions for the PR class.
 256 initBuiltinPRs :: Builtins -> (InstEnv, InstEnv) -> DsM [(Name, Var)]
 257 initBuiltinPRs (Builtins { dphModules = mods }) insts
 258   = liftM (initBuiltinDicts insts) (externalClass (dph_PArray mods) (fsLit "PR"))
 259
 260
 261 -- | Get the names of all DPH instance functions for this class.
 262 initBuiltinDicts :: (InstEnv, InstEnv) -> Class -> [(Name, Var)]
 263 initBuiltinDicts insts cls = map find $ classInstances insts cls
 264   where
 265     find i | [Just tc] <- instanceRoughTcs i    = (tc, instanceDFunId i)
 266            | otherwise                          = pprPanic "Invalid DPH instance" (ppr i)
 267
 268
 269 -- | Get a list of boxed `TyCons` in the mock prelude. This is Int only.
 270 initBuiltinBoxedTyCons :: Builtins -> DsM [(Name, TyCon)]
 271 initBuiltinBoxedTyCons
 272   = return . builtinBoxedTyCons
 273   where builtinBoxedTyCons :: Builtins -> [(Name, TyCon)]
 274         builtinBoxedTyCons _
 275                 = [(tyConName intPrimTyCon, intTyCon)]
 276
 277 -- | Get a list of all scalar functions in the mock prelude.
 278 --
 279 initBuiltinScalars :: Bool
 280                    -> Builtins -> DsM [Var]
 281 initBuiltinScalars True  _bi = return []
 282 initBuiltinScalars False bi  = mapM (uncurry externalVar) (preludeScalars $ dphModules bi)
 283
 284 -- | Lookup some variable given its name and the module that contains it.
 285 externalVar :: Module -> FastString -> DsM Var
 286 externalVar mod fs
 287   = dsLookupGlobalId =<< lookupOrig mod (mkVarOccFS fs)
 288
 289
 290 -- | Like `externalVar` but wrap the `Var` in a `CoreExpr`
 291 externalFun :: Module -> FastString -> DsM CoreExpr
 292 externalFun mod fs
 293  = do var <- externalVar mod fs
 294       return $ Var var
 295
 296
 297 -- | Lookup some `TyCon` given its name and the module that contains it.
 298 externalTyCon :: Module -> FastString -> DsM TyCon
 299 externalTyCon mod fs
 300   = dsLookupTyCon =<< lookupOrig mod (mkTcOccFS fs)
 301
 302
 303 -- | Lookup some `Type` given its name and the module that contains it.
 304 externalType :: Module -> FastString -> DsM Type
 305 externalType mod fs
 306  = do  tycon <- externalTyCon mod fs
 307        return $ mkTyConApp tycon []
 308
 309
 310 -- | Lookup some `Class` given its name and the module that contains it.
 311 externalClass :: Module -> FastString -> DsM Class
 312 externalClass mod fs
 313   = dsLookupClass =<< lookupOrig mod (mkClsOccFS fs)
 314