compiler/vectorise/Vectorise/Type/Env.hs

   1 {-# OPTIONS_GHC -XNoMonoLocalBinds #-}
   2
   3 module Vectorise.Type.Env (
   4         vectTypeEnv,
   5 ) where
   6
   7 import Vectorise.Env
   8 import Vectorise.Vect
   9 import Vectorise.Monad
  10 import Vectorise.Builtins
  11 import Vectorise.Type.TyConDecl
  12 import Vectorise.Type.Classify
  13 import Vectorise.Type.PADict
  14 import Vectorise.Type.PData
  15 import Vectorise.Type.PRepr
  16 import Vectorise.Type.Repr
  17 import Vectorise.Utils
  18
  19 import HscTypes
  20 import CoreSyn
  21 import CoreUtils
  22 import CoreUnfold
  23 import DataCon
  24 import TyCon
  25 import Type
  26 import FamInstEnv
  27 import OccName
  28 import Id
  29 import MkId
  30 import NameEnv
  31
  32 import Unique
  33 import UniqFM
  34 import Util
  35 import Outputable
  36 import FastString
  37 import MonadUtils
  38 import Control.Monad
  39 import Data.List
  40
  41
  42 -- | Vectorise a type environment.
  43 --   The type environment contains all the type things defined in a module.
  44 --
  45 vectTypeEnv :: TypeEnv
  46             -> VM ( TypeEnv             -- Vectorised type environment.
  47                   , [FamInst]           -- New type family instances.
  48                   , [(Var, CoreExpr)])  -- New top level bindings.
  49 vectTypeEnv env
  50   = do
  51       traceVt "** vectTypeEnv" $ ppr env
  52
  53       cs <- readGEnv $ mk_map . global_tycons
  54
  55       -- Split the list of TyCons into the ones we have to vectorise vs the
  56       -- ones we can pass through unchanged. We also pass through algebraic
  57       -- types that use non Haskell98 features, as we don't handle those.
  58       let tycons               = typeEnvTyCons env
  59           groups               = tyConGroups tycons
  60
  61       let (conv_tcs, keep_tcs) = classifyTyCons cs groups
  62           orig_tcs             = keep_tcs ++ conv_tcs
  63           keep_dcs             = concatMap tyConDataCons keep_tcs
  64
  65       -- Just use the unvectorised versions of these constructors in vectorised code.
  66       zipWithM_ defTyCon   keep_tcs keep_tcs
  67       zipWithM_ defDataCon keep_dcs keep_dcs
  68
  69       -- Vectorise all the declarations.
  70       new_tcs      <- vectTyConDecls conv_tcs
  71
  72       -- We don't need to make new representation types for dictionary
  73       -- constructors. The constructors are always fully applied, and we don't
  74       -- need to lift them to arrays as a dictionary of a particular type
  75       -- always has the same value.
  76       let vect_tcs  = filter (not . isClassTyCon)
  77                     $ keep_tcs ++ new_tcs
  78
  79       reprs <- mapM tyConRepr vect_tcs
  80       repr_tcs  <- zipWith3M buildPReprTyCon orig_tcs vect_tcs reprs
  81       pdata_tcs <- zipWith3M buildPDataTyCon orig_tcs vect_tcs reprs
  82       updGEnv $ extendFamEnv
  83               $ map mkLocalFamInst
  84               $ repr_tcs ++ pdata_tcs
  85
  86       -- Create PRepr and PData instances for the vectorised types.
  87       -- We get back the binds for the instance functions,
  88       -- and some new type constructors for the representation types.
  89       (_, binds, inst_tcs) <- fixV $ \ ~(dfuns', _, _) ->
  90         do
  91           defTyConPAs (zipLazy vect_tcs dfuns')
  92           reprs     <- mapM tyConRepr vect_tcs
  93
  94           dfuns     <- sequence
  95                     $  zipWith5 buildTyConBindings
  96                                orig_tcs
  97                                vect_tcs
  98                                repr_tcs
  99                                pdata_tcs
 100                                reprs
 101
 102           binds     <- takeHoisted
 103           return (dfuns, binds, repr_tcs ++ pdata_tcs)
 104
 105       -- The new type constructors are the vectorised versions of the originals,
 106       -- plus the new type constructors that we use for the representations.
 107       let all_new_tcs = new_tcs ++ inst_tcs
 108
 109       let new_env     =  extendTypeEnvList env
 110                       $  map ATyCon all_new_tcs
 111                       ++ [ADataCon dc | tc <- all_new_tcs
 112                                       , dc <- tyConDataCons tc]
 113
 114       return (new_env, map mkLocalFamInst inst_tcs, binds)
 115
 116    where
 117     mk_map env = listToUFM_Directly [(u, getUnique n /= u) | (u,n) <- nameEnvUniqueElts env]
 118
 119 buildTyConBindings :: TyCon -> TyCon -> TyCon -> TyCon -> SumRepr -> VM Var
 120 buildTyConBindings orig_tc vect_tc prepr_tc pdata_tc repr
 121  = do vectDataConWorkers orig_tc vect_tc pdata_tc
 122       buildPADict vect_tc prepr_tc pdata_tc repr
 123
 124 vectDataConWorkers :: TyCon -> TyCon -> TyCon -> VM ()
 125 vectDataConWorkers orig_tc vect_tc arr_tc
 126  = do bs <- sequence
 127           . zipWith3 def_worker  (tyConDataCons orig_tc) rep_tys
 128           $ zipWith4 mk_data_con (tyConDataCons vect_tc)
 129                                  rep_tys
 130                                  (inits rep_tys)
 131                                  (tail $ tails rep_tys)
 132       mapM_ (uncurry hoistBinding) bs
 133  where
 134     tyvars   = tyConTyVars vect_tc
 135     var_tys  = mkTyVarTys tyvars
 136     ty_args  = map Type var_tys
 137     res_ty   = mkTyConApp vect_tc var_tys
 138
 139     cons     = tyConDataCons vect_tc
 140     arity    = length cons
 141     [arr_dc] = tyConDataCons arr_tc
 142
 143     rep_tys  = map dataConRepArgTys $ tyConDataCons vect_tc
 144
 145
 146     mk_data_con con tys pre post
 147       = liftM2 (,) (vect_data_con con)
 148                    (lift_data_con tys pre post (mkDataConTag con))
 149
 150     sel_replicate len tag
 151       | arity > 1 = do
 152                       rep <- builtin (selReplicate arity)
 153                       return [rep `mkApps` [len, tag]]
 154
 155       | otherwise = return []
 156
 157     vect_data_con con = return $ mkConApp con ty_args
 158     lift_data_con tys pre_tys post_tys tag
 159       = do
 160           len  <- builtin liftingContext
 161           args <- mapM (newLocalVar (fsLit "xs"))
 162                   =<< mapM mkPDataType tys
 163
 164           sel  <- sel_replicate (Var len) tag
 165
 166           pre   <- mapM emptyPD (concat pre_tys)
 167           post  <- mapM emptyPD (concat post_tys)
 168
 169           return . mkLams (len : args)
 170                  . wrapFamInstBody arr_tc var_tys
 171                  . mkConApp arr_dc
 172                  $ ty_args ++ sel ++ pre ++ map Var args ++ post
 173
 174     def_worker data_con arg_tys mk_body
 175       = do
 176           arity <- polyArity tyvars
 177           body <- closedV
 178                 . inBind orig_worker
 179                 . polyAbstract tyvars $ \args ->
 180                   liftM (mkLams (tyvars ++ args) . vectorised)
 181                 $ buildClosures tyvars [] arg_tys res_ty mk_body
 182
 183           raw_worker <- cloneId mkVectOcc orig_worker (exprType body)
 184           let vect_worker = raw_worker `setIdUnfolding`
 185                               mkInlineUnfolding (Just arity) body
 186           defGlobalVar orig_worker vect_worker
 187           return (vect_worker, body)
 188       where
 189         orig_worker = dataConWorkId data_con
 190