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
 195 -- | Get the mapping of names in the Prelude to names in the DPH library.
 196 initBuiltinVars :: Builtins -> DsM [(Var, Var)]
 197 initBuiltinVars (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) = 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
 277 -- | Get a list of all scalar functions in the mock prelude.
 278 initBuiltinScalars :: Builtins -> DsM [Var]
 279 initBuiltinScalars bi
 280   = mapM (uncurry externalVar) (preludeScalars $ dphModules bi)
 281
 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