Finish breaking up VectBuiltIn and VectMonad, and add comments

1 files changed, 314 insertions, 0 deletions

diff --git a/compiler/vectorise/Vectorise/Builtins/Initialise.hs b/compiler/vectorise/Vectorise/Builtins/Initialise.hs
new file mode 100644
index 0000000000..413980a1c3
--- /dev/null
+++ b/compiler/vectorise/Vectorise/Builtins/Initialise.hs
@@ -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)
+
+