Finish breaking up VectBuiltIn and VectMonad, and add comments
[ghc-hetmet.git] / compiler / vectorise / Vectorise / Builtins / Initialise.hs
diff --git a/compiler/vectorise/Vectorise/Builtins/Initialise.hs b/compiler/vectorise/Vectorise/Builtins/Initialise.hs
new file mode 100644 (file)
index 0000000..413980a
--- /dev/null
@@ -0,0 +1,314 @@
+
+
+module Vectorise.Builtins.Initialise (
+       -- * Initialisation
+       initBuiltins, initBuiltinVars, initBuiltinTyCons, initBuiltinDataCons,
+       initBuiltinPAs, initBuiltinPRs,
+       initBuiltinBoxedTyCons, initBuiltinScalars,
+) where
+import Vectorise.Builtins.Base
+import Vectorise.Builtins.Modules
+import Vectorise.Builtins.Prelude
+
+import BasicTypes
+import PrelNames
+import TysPrim
+import DsMonad
+import IfaceEnv
+import InstEnv
+import TysWiredIn
+import DataCon
+import TyCon
+import Class
+import CoreSyn
+import Type
+import OccName
+import Name
+import Module
+import Var
+import Id
+import FastString
+import Outputable
+
+import Control.Monad
+import Data.Array
+import Data.List
+
+-- | Create the initial map of builtin types and functions.
+initBuiltins 
+       :: PackageId    -- ^ package id the builtins are in, eg dph-common
+       -> DsM Builtins
+
+initBuiltins pkg
+ = do mapM_ load dph_Orphans
+
+      -- From dph-common:Data.Array.Parallel.Lifted.PArray
+      parrayTyCon      <- externalTyCon        dph_PArray      (fsLit "PArray")
+      let [parrayDataCon] = tyConDataCons parrayTyCon
+
+      pdataTyCon       <- externalTyCon        dph_PArray      (fsLit "PData")
+      paTyCon          <- externalClassTyCon   dph_PArray      (fsLit "PA")
+      let [paDataCon]  = tyConDataCons paTyCon
+
+      preprTyCon       <- externalTyCon        dph_PArray      (fsLit "PRepr")
+      prTyCon          <- externalClassTyCon   dph_PArray      (fsLit "PR")
+      let [prDataCon]  = tyConDataCons prTyCon
+
+      closureTyCon     <- externalTyCon dph_Closure            (fsLit ":->")
+
+      -- From dph-common:Data.Array.Parallel.Lifted.Repr
+      voidTyCon                <- externalTyCon        dph_Repr        (fsLit "Void")
+      wrapTyCon                <- externalTyCon        dph_Repr        (fsLit "Wrap")
+
+      -- From dph-common:Data.Array.Parallel.Lifted.Unboxed
+      sel_tys          <- mapM (externalType dph_Unboxed)
+                               (numbered "Sel" 2 mAX_DPH_SUM)
+
+      sel_replicates   <- mapM (externalFun dph_Unboxed)
+                               (numbered_hash "replicateSel" 2 mAX_DPH_SUM)
+
+      sel_picks        <- mapM (externalFun dph_Unboxed)
+                               (numbered_hash "pickSel" 2 mAX_DPH_SUM)
+
+      sel_tags         <- mapM (externalFun dph_Unboxed)
+                               (numbered "tagsSel" 2 mAX_DPH_SUM)
+
+      sel_els          <- mapM mk_elements
+                               [(i,j) | i <- [2..mAX_DPH_SUM], j <- [0..i-1]]
+
+      sum_tcs          <- mapM (externalTyCon dph_Repr)
+                               (numbered "Sum" 2 mAX_DPH_SUM)
+
+      let selTys        = listArray (2, mAX_DPH_SUM) sel_tys
+          selReplicates = listArray (2, mAX_DPH_SUM) sel_replicates
+          selPicks      = listArray (2, mAX_DPH_SUM) sel_picks
+          selTagss      = listArray (2, mAX_DPH_SUM) sel_tags
+          selEls        = array     ((2,0), (mAX_DPH_SUM, mAX_DPH_SUM)) sel_els
+          sumTyCons     = listArray (2, mAX_DPH_SUM) sum_tcs
+
+
+      voidVar          <- externalVar dph_Repr         (fsLit "void")
+      pvoidVar         <- externalVar dph_Repr         (fsLit "pvoid")
+      fromVoidVar      <- externalVar dph_Repr         (fsLit "fromVoid")
+      punitVar         <- externalVar dph_Repr         (fsLit "punit")
+      closureVar       <- externalVar dph_Closure      (fsLit "closure")
+      applyVar         <- externalVar dph_Closure      (fsLit "$:")
+      liftedClosureVar <- externalVar dph_Closure      (fsLit "liftedClosure")
+      liftedApplyVar   <- externalVar dph_Closure      (fsLit "liftedApply")
+      replicatePDVar   <- externalVar dph_PArray       (fsLit "replicatePD")
+      emptyPDVar       <- externalVar dph_PArray       (fsLit "emptyPD")
+      packByTagPDVar   <- externalVar dph_PArray       (fsLit "packByTagPD")
+
+      combines                 <- mapM (externalVar dph_PArray)
+                                       [mkFastString ("combine" ++ show i ++ "PD")
+                                       | i <- [2..mAX_DPH_COMBINE]]
+      let combinePDVars = listArray (2, mAX_DPH_COMBINE) combines
+
+      scalarClass      <- externalClass dph_PArray     (fsLit "Scalar")
+      scalar_map       <- externalVar  dph_Scalar      (fsLit "scalar_map")
+      scalar_zip2      <- externalVar  dph_Scalar      (fsLit "scalar_zipWith")
+      scalar_zips      <- mapM (externalVar dph_Scalar)
+                               (numbered "scalar_zipWith" 3 mAX_DPH_SCALAR_ARGS)
+
+      let scalarZips   = listArray (1, mAX_DPH_SCALAR_ARGS)
+                                 (scalar_map : scalar_zip2 : scalar_zips)
+
+      closures                 <- mapM (externalVar dph_Closure)
+                                       (numbered "closure" 1 mAX_DPH_SCALAR_ARGS)
+
+      let closureCtrFuns = listArray (1, mAX_DPH_COMBINE) closures
+
+      liftingContext   <- liftM (\u -> mkSysLocal (fsLit "lc") u intPrimTy)
+                               newUnique
+
+      return   $ Builtins 
+               { dphModules       = mods
+               , parrayTyCon      = parrayTyCon
+               , parrayDataCon    = parrayDataCon
+               , pdataTyCon       = pdataTyCon
+               , paTyCon          = paTyCon
+               , paDataCon        = paDataCon
+               , preprTyCon       = preprTyCon
+               , prTyCon          = prTyCon
+               , prDataCon        = prDataCon
+               , voidTyCon        = voidTyCon
+               , wrapTyCon        = wrapTyCon
+               , selTys           = selTys
+               , selReplicates    = selReplicates
+               , selPicks         = selPicks
+               , selTagss         = selTagss
+               , selEls           = selEls
+               , sumTyCons        = sumTyCons
+               , closureTyCon     = closureTyCon
+               , voidVar          = voidVar
+               , pvoidVar         = pvoidVar
+               , fromVoidVar      = fromVoidVar
+               , punitVar         = punitVar
+               , closureVar       = closureVar
+               , applyVar         = applyVar
+               , liftedClosureVar = liftedClosureVar
+               , liftedApplyVar   = liftedApplyVar
+               , replicatePDVar   = replicatePDVar
+               , emptyPDVar       = emptyPDVar
+               , packByTagPDVar   = packByTagPDVar
+               , combinePDVars    = combinePDVars
+               , scalarClass      = scalarClass
+               , scalarZips       = scalarZips
+               , closureCtrFuns   = closureCtrFuns
+               , liftingContext   = liftingContext
+               }
+  where
+    mods@(Modules {
+               dph_PArray         = dph_PArray
+             , dph_Repr           = dph_Repr
+             , dph_Closure        = dph_Closure
+             , dph_Scalar         = dph_Scalar
+             , dph_Unboxed        = dph_Unboxed
+             })
+      = dph_Modules pkg
+
+    load get_mod = dsLoadModule doc mod
+      where
+        mod = get_mod mods 
+        doc = ppr mod <+> ptext (sLit "is a DPH module")
+
+    -- Make a list of numbered strings in some range, eg foo3, foo4, foo5
+    numbered :: String -> Int -> Int -> [FastString]
+    numbered pfx m n = [mkFastString (pfx ++ show i) | i <- [m..n]]
+
+    numbered_hash :: String -> Int -> Int -> [FastString]
+    numbered_hash pfx m n = [mkFastString (pfx ++ show i ++ "#") | i <- [m..n]]
+
+    mk_elements :: (Int, Int) -> DsM ((Int, Int), CoreExpr)
+    mk_elements (i,j)
+      = do
+          v <- externalVar dph_Unboxed
+             $ mkFastString ("elementsSel" ++ show i ++ "_" ++ show j ++ "#")
+          return ((i,j), Var v)
+
+
+-- | Get the mapping of names in the Prelude to names in the DPH library.
+initBuiltinVars :: Builtins -> DsM [(Var, Var)]
+initBuiltinVars (Builtins { dphModules = mods })
+  = do
+      uvars <- zipWithM externalVar umods ufs
+      vvars <- zipWithM externalVar vmods vfs
+      cvars <- zipWithM externalVar cmods cfs
+      return $ [(v,v) | v <- map dataConWorkId defaultDataConWorkers]
+               ++ zip (map dataConWorkId cons) cvars
+               ++ zip uvars vvars
+  where
+    (umods, ufs, vmods, vfs) = unzip4 (preludeVars mods)
+    (cons, cmods, cfs)       = unzip3 (preludeDataCons mods)
+
+    defaultDataConWorkers :: [DataCon]
+    defaultDataConWorkers = [trueDataCon, falseDataCon, unitDataCon]
+
+
+preludeDataCons :: Modules -> [(DataCon, Module, FastString)]
+preludeDataCons (Modules { dph_Prelude_Tuple = dph_Prelude_Tuple })
+  = [mk_tup n dph_Prelude_Tuple (mkFastString $ "tup" ++ show n) | n <- [2..3]]
+  where
+    mk_tup n mod name = (tupleCon Boxed n, mod, name)
+
+
+-- | Get a list of names to `TyCon`s in the mock prelude.
+initBuiltinTyCons :: Builtins -> DsM [(Name, TyCon)]
+initBuiltinTyCons bi
+  = do
+      -- parr <- externalTyCon dph_Prelude_PArr (fsLit "PArr")
+      dft_tcs <- defaultTyCons
+      return $ (tyConName funTyCon, closureTyCon bi)
+             : (parrTyConName,      parrayTyCon bi)
+
+             -- FIXME: temporary
+             : (tyConName $ parrayTyCon bi, parrayTyCon bi)
+
+             : [(tyConName tc, tc) | tc <- dft_tcs]
+
+  where        defaultTyCons :: DsM [TyCon]
+       defaultTyCons
+        = do   word8 <- dsLookupTyCon word8TyConName
+               return [intTyCon, boolTyCon, doubleTyCon, word8]
+
+
+-- | Get a list of names to `DataCon`s in the mock prelude.
+initBuiltinDataCons :: Builtins -> [(Name, DataCon)]
+initBuiltinDataCons _
+  = [(dataConName dc, dc)| dc <- defaultDataCons]
+  where        defaultDataCons :: [DataCon]
+       defaultDataCons = [trueDataCon, falseDataCon, unitDataCon]
+
+
+-- | Get the names of all buildin instance functions for the PA class.
+initBuiltinPAs :: Builtins -> (InstEnv, InstEnv) -> DsM [(Name, Var)]
+initBuiltinPAs (Builtins { dphModules = mods }) insts
+  = liftM (initBuiltinDicts insts) (externalClass (dph_PArray mods) (fsLit "PA"))
+
+
+-- | Get the names of all builtin instance functions for the PR class.
+initBuiltinPRs :: Builtins -> (InstEnv, InstEnv) -> DsM [(Name, Var)]
+initBuiltinPRs (Builtins { dphModules = mods }) insts
+  = liftM (initBuiltinDicts insts) (externalClass (dph_PArray mods) (fsLit "PR"))
+
+
+-- | Get the names of all DPH instance functions for this class.
+initBuiltinDicts :: (InstEnv, InstEnv) -> Class -> [(Name, Var)]
+initBuiltinDicts insts cls = map find $ classInstances insts cls
+  where
+    find i | [Just tc] <- instanceRoughTcs i   = (tc, instanceDFunId i)
+           | otherwise                         = pprPanic "Invalid DPH instance" (ppr i)
+
+
+-- | Get a list of boxed `TyCons` in the mock prelude. This is Int only.
+initBuiltinBoxedTyCons :: Builtins -> DsM [(Name, TyCon)]
+initBuiltinBoxedTyCons 
+  = return . builtinBoxedTyCons
+  where        builtinBoxedTyCons :: Builtins -> [(Name, TyCon)]
+       builtinBoxedTyCons _ 
+               = [(tyConName intPrimTyCon, intTyCon)]
+
+
+-- | Get a list of all scalar functions in the mock prelude.
+initBuiltinScalars :: Builtins -> DsM [Var]
+initBuiltinScalars bi
+  = mapM (uncurry externalVar) (preludeScalars $ dphModules bi)
+
+
+-- | Lookup some variable given its name and the module that contains it.
+externalVar :: Module -> FastString -> DsM Var
+externalVar mod fs
+  = dsLookupGlobalId =<< lookupOrig mod (mkVarOccFS fs)
+
+
+-- | Like `externalVar` but wrap the `Var` in a `CoreExpr`
+externalFun :: Module -> FastString -> DsM CoreExpr
+externalFun mod fs
+ = do var <- externalVar mod fs
+      return $ Var var
+
+
+-- | Lookup some `TyCon` given its name and the module that contains it.
+externalTyCon :: Module -> FastString -> DsM TyCon
+externalTyCon mod fs
+  = dsLookupTyCon =<< lookupOrig mod (mkTcOccFS fs)
+
+
+-- | Lookup some `Type` given its name and the module that contains it.
+externalType :: Module -> FastString -> DsM Type
+externalType mod fs
+ = do  tycon <- externalTyCon mod fs
+       return $ mkTyConApp tycon []
+
+
+-- | Lookup some `Class` given its name and the module that contains it.
+externalClass :: Module -> FastString -> DsM Class
+externalClass mod fs
+  = dsLookupClass =<< lookupOrig mod (mkClsOccFS fs)
+
+
+-- | Like `externalClass`, but get the TyCon of of the class.
+externalClassTyCon :: Module -> FastString -> DsM TyCon
+externalClassTyCon mod fs = liftM classTyCon (externalClass mod fs)
+
+