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