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