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       pa                <- externalClass        dph_PArray      (fsLit "PA")
  50       let paTyCon     = classTyCon pa
  51           [paDataCon] = tyConDataCons paTyCon
  52           paPRSel     = classSCSelId pa 0
  53
  54       preprTyCon        <- externalTyCon        dph_PArray      (fsLit "PRepr")
  55       prTyCon           <- externalClassTyCon   dph_PArray      (fsLit "PR")
  56       let [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                , paTyCon          = paTyCon
 131                , paDataCon        = paDataCon
 132                , paPRSel          = paPRSel
 133                , preprTyCon       = preprTyCon
 134                , prTyCon          = prTyCon
 135                , prDataCon        = prDataCon
 136                , voidTyCon        = voidTyCon
 137                , wrapTyCon        = wrapTyCon
 138                , selTys           = selTys
 139                , selReplicates    = selReplicates
 140                , selPicks         = selPicks
 141                , selTagss         = selTagss
 142                , selEls           = selEls
 143                , sumTyCons        = sumTyCons
 144                , closureTyCon     = closureTyCon
 145                , voidVar          = voidVar
 146                , pvoidVar         = pvoidVar
 147                , fromVoidVar      = fromVoidVar
 148                , punitVar         = punitVar
 149                , closureVar       = closureVar
 150                , applyVar         = applyVar
 151                , liftedClosureVar = liftedClosureVar
 152                , liftedApplyVar   = liftedApplyVar
 153                , replicatePDVar   = replicatePDVar
 154                , emptyPDVar       = emptyPDVar
 155                , packByTagPDVar   = packByTagPDVar
 156                , combinePDVars    = combinePDVars
 157                , scalarClass      = scalarClass
 158                , scalarZips       = scalarZips
 159                , closureCtrFuns   = closureCtrFuns
 160                , liftingContext   = liftingContext
 161                }
 162   where
 163     mods@(Modules {
 164                dph_PArray         = dph_PArray
 165              , dph_Repr           = dph_Repr
 166              , dph_Closure        = dph_Closure
 167              , dph_Scalar         = dph_Scalar
 168              , dph_Unboxed        = dph_Unboxed
 169              })
 170       = dph_Modules pkg
 171
 172     load get_mod = dsLoadModule doc mod
 173       where
 174         mod = get_mod mods
 175         doc = ppr mod <+> ptext (sLit "is a DPH module")
 176
 177     -- Make a list of numbered strings in some range, eg foo3, foo4, foo5
 178     numbered :: String -> Int -> Int -> [FastString]
 179     numbered pfx m n = [mkFastString (pfx ++ show i) | i <- [m..n]]
 180
 181     numbered_hash :: String -> Int -> Int -> [FastString]
 182     numbered_hash pfx m n = [mkFastString (pfx ++ show i ++ "#") | i <- [m..n]]
 183
 184     mk_elements :: (Int, Int) -> DsM ((Int, Int), CoreExpr)
 185     mk_elements (i,j)
 186       = do
 187           v <- externalVar dph_Unboxed
 188              $ mkFastString ("elementsSel" ++ show i ++ "_" ++ show j ++ "#")
 189           return ((i,j), Var v)
 190
 191
 192 -- | Get the mapping of names in the Prelude to names in the DPH library.
 193 initBuiltinVars :: Builtins -> DsM [(Var, Var)]
 194 initBuiltinVars (Builtins { dphModules = mods })
 195   = do
 196       uvars <- zipWithM externalVar umods ufs
 197       vvars <- zipWithM externalVar vmods vfs
 198       cvars <- zipWithM externalVar cmods cfs
 199       return $ [(v,v) | v <- map dataConWorkId defaultDataConWorkers]
 200                ++ zip (map dataConWorkId cons) cvars
 201                ++ zip uvars vvars
 202   where
 203     (umods, ufs, vmods, vfs) = unzip4 (preludeVars mods)
 204     (cons, cmods, cfs)       = unzip3 (preludeDataCons mods)
 205
 206     defaultDataConWorkers :: [DataCon]
 207     defaultDataConWorkers = [trueDataCon, falseDataCon, unitDataCon]
 208
 209
 210 preludeDataCons :: Modules -> [(DataCon, Module, FastString)]
 211 preludeDataCons (Modules { dph_Prelude_Tuple = dph_Prelude_Tuple })
 212   = [mk_tup n dph_Prelude_Tuple (mkFastString $ "tup" ++ show n) | n <- [2..3]]
 213   where
 214     mk_tup n mod name = (tupleCon Boxed n, mod, name)
 215
 216
 217 -- | Get a list of names to `TyCon`s in the mock prelude.
 218 initBuiltinTyCons :: Builtins -> DsM [(Name, TyCon)]
 219 initBuiltinTyCons bi
 220   = do
 221       -- parr <- externalTyCon dph_Prelude_PArr (fsLit "PArr")
 222       dft_tcs <- defaultTyCons
 223       return $ (tyConName funTyCon, closureTyCon bi)
 224              : (parrTyConName,      parrayTyCon bi)
 225
 226              -- FIXME: temporary
 227              : (tyConName $ parrayTyCon bi, parrayTyCon bi)
 228
 229              : [(tyConName tc, tc) | tc <- dft_tcs]
 230
 231   where defaultTyCons :: DsM [TyCon]
 232         defaultTyCons
 233          = do   word8 <- dsLookupTyCon word8TyConName
 234                 return [intTyCon, boolTyCon, doubleTyCon, word8]
 235
 236
 237 -- | Get a list of names to `DataCon`s in the mock prelude.
 238 initBuiltinDataCons :: Builtins -> [(Name, DataCon)]
 239 initBuiltinDataCons _
 240   = [(dataConName dc, dc)| dc <- defaultDataCons]
 241   where defaultDataCons :: [DataCon]
 242         defaultDataCons = [trueDataCon, falseDataCon, unitDataCon]
 243
 244
 245 -- | Get the names of all buildin instance functions for the PA class.
 246 initBuiltinPAs :: Builtins -> (InstEnv, InstEnv) -> DsM [(Name, Var)]
 247 initBuiltinPAs (Builtins { dphModules = mods }) insts
 248   = liftM (initBuiltinDicts insts) (externalClass (dph_PArray mods) (fsLit "PA"))
 249
 250
 251 -- | Get the names of all builtin instance functions for the PR class.
 252 initBuiltinPRs :: Builtins -> (InstEnv, InstEnv) -> DsM [(Name, Var)]
 253 initBuiltinPRs (Builtins { dphModules = mods }) insts
 254   = liftM (initBuiltinDicts insts) (externalClass (dph_PArray mods) (fsLit "PR"))
 255
 256
 257 -- | Get the names of all DPH instance functions for this class.
 258 initBuiltinDicts :: (InstEnv, InstEnv) -> Class -> [(Name, Var)]
 259 initBuiltinDicts insts cls = map find $ classInstances insts cls
 260   where
 261     find i | [Just tc] <- instanceRoughTcs i    = (tc, instanceDFunId i)
 262            | otherwise                          = pprPanic "Invalid DPH instance" (ppr i)
 263
 264
 265 -- | Get a list of boxed `TyCons` in the mock prelude. This is Int only.
 266 initBuiltinBoxedTyCons :: Builtins -> DsM [(Name, TyCon)]
 267 initBuiltinBoxedTyCons
 268   = return . builtinBoxedTyCons
 269   where builtinBoxedTyCons :: Builtins -> [(Name, TyCon)]
 270         builtinBoxedTyCons _
 271                 = [(tyConName intPrimTyCon, intTyCon)]
 272
 273
 274 -- | Get a list of all scalar functions in the mock prelude.
 275 initBuiltinScalars :: Builtins -> DsM [Var]
 276 initBuiltinScalars bi
 277   = mapM (uncurry externalVar) (preludeScalars $ dphModules bi)
 278
 279
 280 -- | Lookup some variable given its name and the module that contains it.
 281 externalVar :: Module -> FastString -> DsM Var
 282 externalVar mod fs
 283   = dsLookupGlobalId =<< lookupOrig mod (mkVarOccFS fs)
 284
 285
 286 -- | Like `externalVar` but wrap the `Var` in a `CoreExpr`
 287 externalFun :: Module -> FastString -> DsM CoreExpr
 288 externalFun mod fs
 289  = do var <- externalVar mod fs
 290       return $ Var var
 291
 292
 293 -- | Lookup some `TyCon` given its name and the module that contains it.
 294 externalTyCon :: Module -> FastString -> DsM TyCon
 295 externalTyCon mod fs
 296   = dsLookupTyCon =<< lookupOrig mod (mkTcOccFS fs)
 297
 298
 299 -- | Lookup some `Type` given its name and the module that contains it.
 300 externalType :: Module -> FastString -> DsM Type
 301 externalType mod fs
 302  = do  tycon <- externalTyCon mod fs
 303        return $ mkTyConApp tycon []
 304
 305
 306 -- | Lookup some `Class` given its name and the module that contains it.
 307 externalClass :: Module -> FastString -> DsM Class
 308 externalClass mod fs
 309   = dsLookupClass =<< lookupOrig mod (mkClsOccFS fs)
 310
 311
 312 -- | Like `externalClass`, but get the TyCon of of the class.
 313 externalClassTyCon :: Module -> FastString -> DsM TyCon
 314 externalClassTyCon mod fs = liftM classTyCon (externalClass mod fs)
 315
 316