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