compiler/vectorise/VectType.hs

   1 module VectType ( vectTyCon, vectType, vectTypeEnv )
   2 where
   3
   4 #include "HsVersions.h"
   5
   6 import VectMonad
   7 import VectUtils
   8
   9 import HscTypes          ( TypeEnv, extendTypeEnvList, typeEnvTyCons )
  10 import CoreSyn
  11 import DataCon
  12 import TyCon
  13 import Type
  14 import TypeRep
  15 import Coercion
  16 import FamInstEnv        ( FamInst, mkLocalFamInst )
  17 import InstEnv           ( Instance )
  18 import OccName
  19 import MkId
  20 import BasicTypes        ( StrictnessMark(..), boolToRecFlag )
  21 import Id                ( mkWildId )
  22 import Name              ( Name )
  23 import NameEnv
  24 import TysWiredIn        ( intTy )
  25 import TysPrim           ( intPrimTy )
  26
  27 import Unique
  28 import UniqFM
  29 import UniqSet
  30 import Digraph           ( SCC(..), stronglyConnComp )
  31
  32 import Outputable
  33
  34 import Control.Monad  ( liftM2, zipWithM, zipWithM_ )
  35
  36 -- ----------------------------------------------------------------------------
  37 -- Types
  38
  39 vectTyCon :: TyCon -> VM TyCon
  40 vectTyCon tc
  41   | isFunTyCon tc        = builtin closureTyCon
  42   | isBoxedTupleTyCon tc = return tc
  43   | isUnLiftedTyCon tc   = return tc
  44   | otherwise = do
  45                   r <- lookupTyCon tc
  46                   case r of
  47                     Just tc' -> return tc'
  48
  49                     -- FIXME: just for now
  50                     Nothing  -> pprTrace "ccTyCon:" (ppr tc) $ return tc
  51
  52 vectType :: Type -> VM Type
  53 vectType ty | Just ty' <- coreView ty = vectType ty'
  54 vectType (TyVarTy tv) = return $ TyVarTy tv
  55 vectType (AppTy ty1 ty2) = liftM2 AppTy (vectType ty1) (vectType ty2)
  56 vectType (TyConApp tc tys) = liftM2 TyConApp (vectTyCon tc) (mapM vectType tys)
  57 vectType (FunTy ty1 ty2)   = liftM2 TyConApp (builtin closureTyCon)
  58                                              (mapM vectType [ty1,ty2])
  59 vectType ty@(ForAllTy _ _)
  60   = do
  61       mdicts   <- mapM paDictArgType tyvars
  62       mono_ty' <- vectType mono_ty
  63       return $ tyvars `mkForAllTys` ([dict | Just dict <- mdicts] `mkFunTys` mono_ty')
  64   where
  65     (tyvars, mono_ty) = splitForAllTys ty
  66
  67 vectType ty = pprPanic "vectType:" (ppr ty)
  68
  69 -- ----------------------------------------------------------------------------
  70 -- Type definitions
  71
  72 type TyConGroup = ([TyCon], UniqSet TyCon)
  73
  74 vectTypeEnv :: TypeEnv -> VM (TypeEnv, [FamInst], [Instance])
  75 vectTypeEnv env
  76   = do
  77       cs <- readGEnv $ mk_map . global_tycons
  78       let (conv_tcs, keep_tcs) = classifyTyCons cs groups
  79           keep_dcs             = concatMap tyConDataCons keep_tcs
  80       zipWithM_ defTyCon   keep_tcs keep_tcs
  81       zipWithM_ defDataCon keep_dcs keep_dcs
  82       vect_tcs <- vectTyConDecls conv_tcs
  83       parr_tcs1 <- zipWithM buildPArrayTyCon keep_tcs keep_tcs
  84       parr_tcs2 <- zipWithM buildPArrayTyCon conv_tcs vect_tcs
  85       let new_tcs = vect_tcs ++ parr_tcs1 ++ parr_tcs2
  86
  87       let new_env = extendTypeEnvList env
  88                        (map ATyCon new_tcs
  89                         ++ [ADataCon dc | tc <- new_tcs
  90                                         , dc <- tyConDataCons tc])
  91
  92       return (new_env, map mkLocalFamInst (parr_tcs1 ++ parr_tcs2), [])
  93   where
  94     tycons = typeEnvTyCons env
  95     groups = tyConGroups tycons
  96
  97     mk_map env = listToUFM_Directly [(u, getUnique n /= u) | (u,n) <- nameEnvUniqueElts env]
  98
  99     keep_tc tc = let dcs = tyConDataCons tc
 100                  in
 101                  defTyCon tc tc >> zipWithM_ defDataCon dcs dcs
 102
 103
 104 vectTyConDecls :: [TyCon] -> VM [TyCon]
 105 vectTyConDecls tcs = fixV $ \tcs' ->
 106   do
 107     mapM_ (uncurry defTyCon) (lazy_zip tcs tcs')
 108     mapM vectTyConDecl tcs
 109   where
 110     lazy_zip [] _ = []
 111     lazy_zip (x:xs) ~(y:ys) = (x,y) : lazy_zip xs ys
 112
 113 vectTyConDecl :: TyCon -> VM TyCon
 114 vectTyConDecl tc
 115   = do
 116       name' <- cloneName mkVectTyConOcc name
 117       rhs'  <- vectAlgTyConRhs (algTyConRhs tc)
 118
 119       return $ mkAlgTyCon name'
 120                           kind
 121                           tyvars
 122                           []              -- no stupid theta
 123                           rhs'
 124                           []              -- no selector ids
 125                           NoParentTyCon   -- FIXME
 126                           rec_flag        -- FIXME: is this ok?
 127                           False           -- FIXME: no generics
 128                           False           -- not GADT syntax
 129   where
 130     name   = tyConName tc
 131     kind   = tyConKind tc
 132     tyvars = tyConTyVars tc
 133     rec_flag = boolToRecFlag (isRecursiveTyCon tc)
 134
 135 vectAlgTyConRhs :: AlgTyConRhs -> VM AlgTyConRhs
 136 vectAlgTyConRhs (DataTyCon { data_cons = data_cons
 137                            , is_enum   = is_enum
 138                            })
 139   = do
 140       data_cons' <- mapM vectDataCon data_cons
 141       zipWithM_ defDataCon data_cons data_cons'
 142       return $ DataTyCon { data_cons = data_cons'
 143                          , is_enum   = is_enum
 144                          }
 145
 146 vectDataCon :: DataCon -> VM DataCon
 147 vectDataCon dc
 148   | not . null $ dataConExTyVars dc = pprPanic "vectDataCon: existentials" (ppr dc)
 149   | not . null $ dataConEqSpec   dc = pprPanic "vectDataCon: eq spec" (ppr dc)
 150   | otherwise
 151   = do
 152       name'    <- cloneName mkVectDataConOcc name
 153       tycon'   <- vectTyCon tycon
 154       arg_tys  <- mapM vectType rep_arg_tys
 155       wrk_name <- cloneName mkDataConWorkerOcc name'
 156
 157       let ids      = mkDataConIds (panic "vectDataCon: wrapper id")
 158                                   wrk_name
 159                                   data_con
 160           data_con = mkDataCon name'
 161                                False           -- not infix
 162                                (map (const NotMarkedStrict) arg_tys)
 163                                []              -- no labelled fields
 164                                univ_tvs
 165                                []              -- no existential tvs for now
 166                                []              -- no eq spec for now
 167                                []              -- no theta
 168                                arg_tys
 169                                tycon'
 170                                []              -- no stupid theta
 171                                ids
 172       return data_con
 173   where
 174     name        = dataConName dc
 175     univ_tvs    = dataConUnivTyVars dc
 176     rep_arg_tys = dataConRepArgTys dc
 177     tycon       = dataConTyCon dc
 178
 179 buildPArrayTyCon :: TyCon -> TyCon -> VM TyCon
 180 buildPArrayTyCon orig_tc vect_tc = fixV $ \repr_tc ->
 181   do
 182     name'  <- cloneName mkPArrayTyConOcc orig_name
 183     parent <- buildPArrayParentInfo orig_name vect_tc repr_tc
 184     rhs    <- buildPArrayTyConRhs orig_name vect_tc repr_tc
 185
 186     return $ mkAlgTyCon name'
 187                         kind
 188                         tyvars
 189                         []              -- no stupid theta
 190                         rhs
 191                         []              -- no selector ids
 192                         parent
 193                         rec_flag        -- FIXME: is this ok?
 194                         False           -- FIXME: no generics
 195                         False           -- not GADT syntax
 196   where
 197     orig_name = tyConName orig_tc
 198     name   = tyConName vect_tc
 199     kind   = tyConKind vect_tc
 200     tyvars = tyConTyVars vect_tc
 201     rec_flag = boolToRecFlag (isRecursiveTyCon vect_tc)
 202
 203
 204 buildPArrayParentInfo :: Name -> TyCon -> TyCon -> VM TyConParent
 205 buildPArrayParentInfo orig_name vect_tc repr_tc
 206   = do
 207       parray_tc <- builtin parrayTyCon
 208       co_name <- cloneName mkInstTyCoOcc (tyConName repr_tc)
 209
 210       let inst_tys = [mkTyConApp vect_tc (map mkTyVarTy tyvars)]
 211
 212       return . FamilyTyCon parray_tc inst_tys
 213              $ mkFamInstCoercion co_name
 214                                  tyvars
 215                                  parray_tc
 216                                  inst_tys
 217                                  repr_tc
 218   where
 219     tyvars = tyConTyVars vect_tc
 220
 221 buildPArrayTyConRhs :: Name -> TyCon -> TyCon -> VM AlgTyConRhs
 222 buildPArrayTyConRhs orig_name vect_tc repr_tc
 223   = do
 224       data_con <- buildPArrayDataCon orig_name vect_tc repr_tc
 225       return $ DataTyCon { data_cons = [data_con], is_enum = False }
 226
 227 buildPArrayDataCon :: Name -> TyCon -> TyCon -> VM DataCon
 228 buildPArrayDataCon orig_name vect_tc repr_tc
 229   = do
 230       dc_name  <- cloneName mkPArrayDataConOcc orig_name
 231       shape_ty <- mkPArrayType intTy   -- FIXME: we want to unbox this!
 232       repr_tys <- mapM mkPArrayType types
 233       wrk_name <- cloneName mkDataConWorkerOcc  dc_name
 234       wrp_name <- cloneName mkDataConWrapperOcc dc_name
 235
 236       let ids      = mkDataConIds wrp_name wrk_name data_con
 237           data_con = mkDataCon dc_name
 238                                False
 239                                (MarkedStrict : map (const NotMarkedStrict) repr_tys)
 240                                []
 241                                (tyConTyVars vect_tc)
 242                                []
 243                                []
 244                                []
 245                                (shape_ty : repr_tys)
 246                                repr_tc
 247                                []
 248                                ids
 249
 250       return data_con
 251   where
 252     types = [ty | dc <- tyConDataCons vect_tc
 253                 , ty <- dataConRepArgTys dc]
 254
 255 buildLengthPA :: TyCon -> VM CoreExpr
 256 buildLengthPA repr_tc
 257   = do
 258       arg   <- newLocalVar FSLIT("xs") arg_ty
 259       shape <- newLocalVar FSLIT("sel") shape_ty
 260       body  <- lengthPA (Var shape)
 261       return . Lam arg
 262              $ Case (Var arg) (mkWildId arg_ty) intPrimTy
 263                     [(DataAlt repr_dc, shape : map mkWildId repr_tys, body)]
 264   where
 265     arg_ty = mkTyConApp repr_tc . mkTyVarTys $ tyConTyVars repr_tc
 266     [repr_dc] = tyConDataCons repr_tc
 267     shape_ty : repr_tys = dataConRepArgTys repr_dc
 268
 269 -- | Split the given tycons into two sets depending on whether they have to be
 270 -- converted (first list) or not (second list). The first argument contains
 271 -- information about the conversion status of external tycons:
 272 --
 273 --   * tycons which have converted versions are mapped to True
 274 --   * tycons which are not changed by vectorisation are mapped to False
 275 --   * tycons which can't be converted are not elements of the map
 276 --
 277 classifyTyCons :: UniqFM Bool -> [TyConGroup] -> ([TyCon], [TyCon])
 278 classifyTyCons = classify [] []
 279   where
 280     classify conv keep cs [] = (conv, keep)
 281     classify conv keep cs ((tcs, ds) : rs)
 282       | can_convert && must_convert
 283         = classify (tcs ++ conv) keep (cs `addListToUFM` [(tc,True) | tc <- tcs]) rs
 284       | can_convert
 285         = classify conv (tcs ++ keep) (cs `addListToUFM` [(tc,False) | tc <- tcs]) rs
 286       | otherwise
 287         = classify conv keep cs rs
 288       where
 289         refs = ds `delListFromUniqSet` tcs
 290
 291         can_convert  = isNullUFM (refs `minusUFM` cs) && all convertable tcs
 292         must_convert = foldUFM (||) False (intersectUFM_C const cs refs)
 293
 294         convertable tc = isDataTyCon tc && all isVanillaDataCon (tyConDataCons tc)
 295
 296 -- | Compute mutually recursive groups of tycons in topological order
 297 --
 298 tyConGroups :: [TyCon] -> [TyConGroup]
 299 tyConGroups tcs = map mk_grp (stronglyConnComp edges)
 300   where
 301     edges = [((tc, ds), tc, uniqSetToList ds) | tc <- tcs
 302                                 , let ds = tyConsOfTyCon tc]
 303
 304     mk_grp (AcyclicSCC (tc, ds)) = ([tc], ds)
 305     mk_grp (CyclicSCC els)       = (tcs, unionManyUniqSets dss)
 306       where
 307         (tcs, dss) = unzip els
 308
 309 tyConsOfTyCon :: TyCon -> UniqSet TyCon
 310 tyConsOfTyCon
 311   = tyConsOfTypes . concatMap dataConRepArgTys . tyConDataCons
 312
 313 tyConsOfType :: Type -> UniqSet TyCon
 314 tyConsOfType ty
 315   | Just ty' <- coreView ty    = tyConsOfType ty'
 316 tyConsOfType (TyVarTy v)       = emptyUniqSet
 317 tyConsOfType (TyConApp tc tys) = extend (tyConsOfTypes tys)
 318   where
 319     extend | isUnLiftedTyCon tc
 320            || isTupleTyCon   tc = id
 321
 322            | otherwise          = (`addOneToUniqSet` tc)
 323
 324 tyConsOfType (AppTy a b)       = tyConsOfType a `unionUniqSets` tyConsOfType b
 325 tyConsOfType (FunTy a b)       = (tyConsOfType a `unionUniqSets` tyConsOfType b)
 326                                  `addOneToUniqSet` funTyCon
 327 tyConsOfType (ForAllTy _ ty)   = tyConsOfType ty
 328 tyConsOfType other             = pprPanic "ClosureConv.tyConsOfType" $ ppr other
 329
 330 tyConsOfTypes :: [Type] -> UniqSet TyCon
 331 tyConsOfTypes = unionManyUniqSets . map tyConsOfType
 332