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