Classify.hs 5.81 KB
Newer Older
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15
-- Extract from a list of type constructors those (1) which need to be vectorised and (2) those
-- that could be, but need not be vectorised (as a scalar representation is sufficient and more
-- efficient).  The type constructors that cannot be vectorised will be dropped.
--
-- A type constructor will only be vectorised if it is
--
-- (1) a data type constructor, with vanilla data constructors (i.e., data constructors admitted by
--     Haskell 98) and
-- (2) at least one of the type constructors that appears in its definition is also vectorised.
--
-- If (1) is met, but not (2), the type constructor may appear in vectorised code, but there is no
-- need to vectorise that type constructor itself.  This holds, for example, for all enumeration
-- types.  As '([::])' is being vectorised, any type constructor whose definition involves
-- '([::])', either directly or indirectly, will be vectorised.

16 17 18 19
module Vectorise.Type.Classify 
  ( classifyTyCons
  ) 
where
20

21
import NameSet
22 23 24 25 26 27
import UniqSet
import UniqFM
import DataCon
import TyCon
import TypeRep
import Type
28
import PrelNames
29 30 31
import Digraph


32
-- |From a list of type constructors, extract those that can be vectorised, returning them in two
33
-- sets, where the first result list /must be/ vectorised and the second result list /need not be/
34
-- vectorised. The third result list are those type constructors that we cannot convert (either
35 36
-- because they use language extensions or because they dependent on type constructors for which
-- no vectorised version is available).
37 38

-- The first argument determines the /conversion status/ of external type constructors as follows:
39
--
40 41
-- * tycons which have converted versions are mapped to 'True'
-- * tycons which are not changed by vectorisation are mapped to 'False'
42 43
-- * tycons which haven't been converted (because they can't or weren't vectorised) are not
--   elements of the map
44
--
45 46 47 48 49 50 51
classifyTyCons :: UniqFM Bool                  -- ^type constructor vectorisation status
               -> NameSet                      -- ^tycons involving parallel arrays
               -> [TyCon]                      -- ^type constructors that need to be classified
               -> ( [TyCon]                    -- to be converted
                  , [TyCon]                    -- need not be converted (but could be)
                  , [TyCon]                    -- involve parallel arrays (whether converted or not)
                  , [TyCon]                    -- can't be converted
52 53
                  )
classifyTyCons convStatus parTyCons tcs = classify [] [] [] [] convStatus parTyCons (tyConGroups tcs)
54
  where
55 56
    classify conv keep par novect _  _   []               = (conv, keep, par, novect)
    classify conv keep par novect cs pts ((tcs, ds) : rs)
57
      | can_convert && must_convert
58
      = classify (tcs ++ conv) keep (par ++ tcs_par) novect (cs `addListToUFM` [(tc, True)  | tc <- tcs]) pts' rs
59
      | can_convert
60
      = classify conv (tcs ++ keep) (par ++ tcs_par) novect (cs `addListToUFM` [(tc, False) | tc <- tcs]) pts' rs
61
      | otherwise
62
      = classify conv keep (par ++ tcs_par) (tcs ++ novect) cs pts' rs
63 64
      where
        refs = ds `delListFromUniqSet` tcs
65
        
66 67 68 69 70
          -- the tycons that directly or indirectly depend on parallel arrays
        tcs_par | any ((`elemNameSet` parTyCons) . tyConName) . eltsUFM $ refs = tcs
                | otherwise                                                    = []

        pts' = pts `addListToNameSet` map tyConName tcs_par
71

72 73
        can_convert  = (isNullUFM (refs `minusUFM` cs) && all convertable tcs)
                       || isShowClass tcs
74
        must_convert = foldUFM (||) False (intersectUFM_C const cs refs)
75
                       && (not . isShowClass $ tcs)
76

77 78 79 80 81
        -- We currently admit Haskell 2011-style data and newtype declarations as well as type
        -- constructors representing classes.
        convertable tc 
          = (isDataTyCon tc || isNewTyCon tc) && all isVanillaDataCon (tyConDataCons tc)
            || isClassTyCon tc
82 83 84 85 86
            
        -- !!!FIXME: currently we allow 'Show' in vectorised code without actually providing a
        --   vectorised definition (to be able to vectorise 'Num')
        isShowClass [tc] = tyConName tc == showClassName
        isShowClass _    = False
87

88 89 90
-- Used to group type constructors into mutually dependent groups.
--
type TyConGroup = ([TyCon], UniqSet TyCon)
91

92 93
-- Compute mutually recursive groups of tycons in topological order.
--
94 95 96 97 98 99 100 101 102 103 104
tyConGroups :: [TyCon] -> [TyConGroup]
tyConGroups tcs = map mk_grp (stronglyConnCompFromEdgedVertices edges)
  where
    edges = [((tc, ds), tc, uniqSetToList ds) | tc <- tcs
                                , let ds = tyConsOfTyCon tc]

    mk_grp (AcyclicSCC (tc, ds)) = ([tc], ds)
    mk_grp (CyclicSCC els)       = (tcs, unionManyUniqSets dss)
      where
        (tcs, dss) = unzip els

105 106
-- |Collect the set of TyCons used by the representation of some data type.
--
107
tyConsOfTyCon :: TyCon -> UniqSet TyCon
108
tyConsOfTyCon = tyConsOfTypes . concatMap dataConRepArgTys . tyConDataCons
109

110 111
-- |Collect the set of TyCons that occur in these types.
--
112 113 114
tyConsOfTypes :: [Type] -> UniqSet TyCon
tyConsOfTypes = unionManyUniqSets . map tyConsOfType

115 116
-- |Collect the set of TyCons that occur in this type.
--
117 118 119 120 121 122 123 124 125 126 127 128 129 130
tyConsOfType :: Type -> UniqSet TyCon
tyConsOfType ty
  | Just ty' <- coreView ty    = tyConsOfType ty'
tyConsOfType (TyVarTy _)       = emptyUniqSet
tyConsOfType (TyConApp tc tys) = extend (tyConsOfTypes tys)
  where
    extend |  isUnLiftedTyCon tc
           || isTupleTyCon   tc = id

           | otherwise          = (`addOneToUniqSet` tc)

tyConsOfType (AppTy a b)       = tyConsOfType a `unionUniqSets` tyConsOfType b
tyConsOfType (FunTy a b)       = (tyConsOfType a `unionUniqSets` tyConsOfType b)
                                 `addOneToUniqSet` funTyCon
131
tyConsOfType (LitTy _)         = emptyUniqSet
132
tyConsOfType (ForAllTy _ ty)   = tyConsOfType ty