module Vectorise.Builtins.Initialise (
initBuiltins, initBuiltinVars, initBuiltinTyCons, initBuiltinDataCons,
initBuiltinPAs, initBuiltinPRs,
initBuiltinBoxedTyCons
) where
import Vectorise.Builtins.Base
import Vectorise.Builtins.Modules
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 Name
import Module
import Id
import FastString
import Outputable
import Control.Monad
import Data.Array
initBuiltins :: PackageId
-> DsM Builtins
initBuiltins pkg
= do mapM_ load dph_Orphans
pdataTyCon <- externalTyCon dph_PArray_PData (fsLit "PData")
prClass <- externalClass dph_PArray_PData (fsLit "PR")
let prTyCon = classTyCon prClass
[prDataCon] = tyConDataCons prTyCon
preprTyCon <- externalTyCon dph_PArray_PRepr (fsLit "PRepr")
paClass <- externalClass dph_PArray_PRepr (fsLit "PA")
let paTyCon = classTyCon paClass
[paDataCon] = tyConDataCons paTyCon
paPRSel = classSCSelId paClass 0
replicatePDVar <- externalVar dph_PArray_PRepr (fsLit "replicatePD")
emptyPDVar <- externalVar dph_PArray_PRepr (fsLit "emptyPD")
packByTagPDVar <- externalVar dph_PArray_PRepr (fsLit "packByTagPD")
combines <- mapM (externalVar dph_PArray_PRepr)
[mkFastString ("combine" ++ show i ++ "PD")
| i <- [2..mAX_DPH_COMBINE]]
let combinePDVars = listArray (2, mAX_DPH_COMBINE) combines
scalarClass <- externalClass dph_PArray_Scalar (fsLit "Scalar")
parrayTyCon <- externalTyCon dph_PArray_Base (fsLit "PArray")
let [parrayDataCon] = tyConDataCons parrayTyCon
voidTyCon <- externalTyCon dph_PArray_Types (fsLit "Void")
voidVar <- externalVar dph_PArray_Types (fsLit "void")
fromVoidVar <- externalVar dph_PArray_Types (fsLit "fromVoid")
wrapTyCon <- externalTyCon dph_PArray_Types (fsLit "Wrap")
sum_tcs <- mapM (externalTyCon dph_PArray_Types) (numbered "Sum" 2 mAX_DPH_SUM)
pvoidVar <- externalVar dph_PArray_PDataInstances (fsLit "pvoid")
punitVar <- externalVar dph_PArray_PDataInstances (fsLit "punit")
closureTyCon <- externalTyCon dph_Closure (fsLit ":->")
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..i1]]
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
closureVar <- externalVar dph_Closure (fsLit "closure")
applyVar <- externalVar dph_Closure (fsLit "$:")
liftedClosureVar <- externalVar dph_Closure (fsLit "liftedClosure")
liftedApplyVar <- externalVar dph_Closure (fsLit "liftedApply")
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
, paClass = paClass
, paTyCon = paTyCon
, paDataCon = paDataCon
, paPRSel = paPRSel
, preprTyCon = preprTyCon
, prClass = prClass
, 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_Base = dph_PArray_Base
, dph_PArray_Scalar = dph_PArray_Scalar
, dph_PArray_PRepr = dph_PArray_PRepr
, dph_PArray_PData = dph_PArray_PData
, dph_PArray_PDataInstances = dph_PArray_PDataInstances
, dph_PArray_Types = dph_PArray_Types
, 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")
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)
initBuiltinVars :: Builtins -> DsM [(Var, Var)]
initBuiltinVars (Builtins { dphModules = mods })
= do
cvars <- zipWithM externalVar cmods cfs
return $ [(v,v) | v <- map dataConWorkId defaultDataConWorkers]
++ zip (map dataConWorkId cons) cvars
where
(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)
initBuiltinTyCons :: Builtins -> DsM [(Name, TyCon)]
initBuiltinTyCons bi
= do
dft_tcs <- defaultTyCons
return $ (tyConName funTyCon, closureTyCon bi)
: (parrTyConName, parrayTyCon bi)
: (tyConName $ parrayTyCon bi, parrayTyCon bi)
: [(tyConName tc, tc) | tc <- dft_tcs]
where
defaultTyCons :: DsM [TyCon]
defaultTyCons
= do word8 <- dsLookupTyCon word8TyConName
return [intTyCon, boolTyCon, floatTyCon, doubleTyCon, word8]
initBuiltinDataCons :: Builtins -> [(Name, DataCon)]
initBuiltinDataCons _
= [(dataConName dc, dc)| dc <- defaultDataCons]
where
defaultDataCons :: [DataCon]
defaultDataCons = [trueDataCon, falseDataCon, unitDataCon]
initBuiltinPAs :: Builtins -> (InstEnv, InstEnv) -> DsM [(Name, Var)]
initBuiltinPAs (Builtins { dphModules = mods }) insts
= liftM (initBuiltinDicts insts) (externalClass (dph_PArray_PRepr mods) (fsLit "PA"))
initBuiltinPRs :: Builtins -> (InstEnv, InstEnv) -> DsM [(Name, Var)]
initBuiltinPRs (Builtins { dphModules = mods }) insts
= liftM (initBuiltinDicts insts) (externalClass (dph_PArray_PData mods) (fsLit "PR"))
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)
initBuiltinBoxedTyCons :: Builtins -> DsM [(Name, TyCon)]
initBuiltinBoxedTyCons
= return . builtinBoxedTyCons
where
builtinBoxedTyCons :: Builtins -> [(Name, TyCon)]
builtinBoxedTyCons _
= [(tyConName intPrimTyCon, intTyCon)]
externalVar :: Module -> FastString -> DsM Var
externalVar mod fs
= dsLookupGlobalId =<< lookupOrig mod (mkVarOccFS fs)
externalFun :: Module -> FastString -> DsM CoreExpr
externalFun mod fs
= do var <- externalVar mod fs
return $ Var var
externalTyCon :: Module -> FastString -> DsM TyCon
externalTyCon mod fs
= dsLookupTyCon =<< lookupOrig mod (mkTcOccFS fs)
externalType :: Module -> FastString -> DsM Type
externalType mod fs
= do tycon <- externalTyCon mod fs
return $ mkTyConApp tycon []
externalClass :: Module -> FastString -> DsM Class
externalClass mod fs
= dsLookupClass =<< lookupOrig mod (mkClsOccFS fs)