Env.hs 10.6 KB
Newer Older
Ian Lynagh's avatar
Ian Lynagh committed
1
{-# OPTIONS_GHC -XNoMonoLocalBinds #-}
2

3 4 5 6 7 8
-- Vectorise a modules type environment, the structure containing all type things defined in a
-- module.
--
-- This extends the type environment with vectorised variants of data types and produces value
-- bindings for worker functions and the like.

9
module Vectorise.Type.Env ( 
10
  vectTypeEnv,
11 12
) where
  
13
import Vectorise.Env
14
import Vectorise.Vect
15 16
import Vectorise.Monad
import Vectorise.Builtins
17
import Vectorise.Type.TyConDecl
18
import Vectorise.Type.Classify
19
import Vectorise.Type.PADict
20 21 22
import Vectorise.Type.PData
import Vectorise.Type.PRepr
import Vectorise.Type.Repr
23
import Vectorise.Utils
24

25
import HscTypes
rl@cse.unsw.edu.au's avatar
rl@cse.unsw.edu.au committed
26
import CoreSyn
rl@cse.unsw.edu.au's avatar
rl@cse.unsw.edu.au committed
27
import CoreUtils
28
import CoreUnfold
29
import DataCon
30 31
import TyCon
import Type
32
import FamInstEnv
33
import Id
34
import MkId
35
import NameEnv
36
import NameSet
37

38
import Util
39
import Outputable
40
import FastString
41 42
import MonadUtils
import Control.Monad
43 44
import Data.List

45

46 47 48 49 50 51 52 53 54
-- Note [Pragmas to vectorise tycons]
-- ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
--
-- VECTORISE pragmas for type constructors cover three different flavours of vectorising data type
-- constructors:
--
-- (1) Data type constructor 'T' that may be used in vectorised code, where 'T' represents itself,
--     but the representation of 'T' is opaque in vectorised code.  
--
55
--     An example is the treatment of 'Int'.  'Int's can be used in vectorised code and remain
56 57 58
--     unchanged by vectorisation.  However, the representation of 'Int' by the 'I#' data
--     constructor wrapping an 'Int#' is not exposed in vectorised code.  Instead, computations
--     involving the representation need to be confined to scalar code.
59
--
60 61 62 63 64 65
--     'PData' and 'PRepr' instances need to be explicitly supplied for 'T' (they are not generated
--     by the vectoriser).
--
--     Type constructors declared with {-# VECTORISE SCALAR type T #-} are treated in this manner.
--     (The vectoriser never treats a type constructor automatically in this manner.)
--
66 67
-- (2) [NOT FULLY IMPLEMENTED YET]
--     Data type constructor 'T' that together with its constructors 'Cn' may be used in vectorised
68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94
--     code, where 'T' and the 'Cn' represent themselves in vectorised code.
--
--     An example is the treatment of 'Bool'.  'Bool' together with 'False' and 'True' may appear in
--     vectorised code and they remain unchanged by vectorisation.  (There is no need for a special
--     representation as the values cannot embed any arrays.)

--     'PData' and 'PRepr' instances are automatically generated by the vectoriser.
--
--     Type constructors declared with {-# VECTORISE type T #-} are treated in this manner.
--     (This is the same treatment that type constructors receive that the vectoriser deems fit for
--     use in vectorised code, but for which no special vectorised variant needs to be generated.)
--
-- (3) [NOT IMPLEMENTED YET]
--     Data type constructor 'T' that together with its constructors 'Cn' may be used in vectorised
--     code, where 'T' is represented by 'Tv' and the workers of the 'Cn' are represented 'vCn' in
--     vectorised code.
--
--     ??Example??
--
--     'PData' and 'PRepr' instances are automatically generated by the vectoriser.
--
--     ??How declared??

-- |Vectorise a type environment.
--
vectTypeEnv :: TypeEnv                  -- Original type environment
            -> [CoreVect]               -- All 'VECTORISE [SCALAR] type' declarations in this module
95 96 97
            -> VM ( TypeEnv             -- Vectorised type environment.
                  , [FamInst]           -- New type family instances.
                  , [(Var, CoreExpr)])  -- New top level bindings.
98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124
vectTypeEnv env vectTypeDecls
  = do { traceVt "** vectTypeEnv" $ ppr env

         -- Build a map containing all vectorised type constructor.  If they are scalar, they are
         -- mapped to 'False' (vectorised type constructor == original type constructor).
       ; allScalarTyConNames <- globalScalarTyCons  -- covers both current and imported modules
       ; vectTyCons          <- globalVectTyCons
       ; let vectTyConBase    = mapNameEnv (const True) vectTyCons   -- by default fully vectorised
             vectTyConFlavour = foldNameSet (\n env -> extendNameEnv env n False) vectTyConBase
                                            allScalarTyConNames

           -- Split the list of 'TyCons' into the ones (1) that we must vectorise and those (2)
           -- that we could, but don't need to vectorise.  Type constructors that are not data
           -- type constructors or use non-Haskell98 features are being dropped.  They may not
           -- appear in vectorised code.  (We also drop the local type constructors appearing in a
           -- VECTORISE SCALAR pragma, as they are being handled separately.)
       ; let localScalarTyCons      = [tycon | VectType tycon Nothing <- vectTypeDecls]
             localScalarTyConNames  = mkNameSet (map tyConName localScalarTyCons)
             notLocalScalarTyCon tc = not $ (tyConName tc) `elemNameSet` localScalarTyConNames

             maybeVectoriseTyCons   = filter notLocalScalarTyCon (typeEnvTyCons env)
             (conv_tcs, keep_tcs)   = classifyTyCons vectTyConFlavour maybeVectoriseTyCons
             orig_tcs               = keep_tcs ++ conv_tcs
             keep_dcs               = concatMap tyConDataCons keep_tcs
             
             keep_and_scalar_tcs    = keep_tcs ++ localScalarTyCons

125 126 127 128
       ; traceVt " declared SCALAR: " $ ppr localScalarTyCons
       ; traceVt " reuse          : " $ ppr keep_tcs
       ; traceVt " convert        : " $ ppr conv_tcs

129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176 177 178 179 180 181 182 183 184 185 186 187 188 189
           -- Of those type constructors that we don't need to vectorise, we use the original
           -- representation in both unvectorised and vectorised code.  For those declared VECTORISE
           -- SCALAR, we ignore their represention — see "Note [Pragmas to vectorise tycons]".
       ; zipWithM_ defTyCon   keep_and_scalar_tcs keep_and_scalar_tcs
       ; zipWithM_ defDataCon keep_dcs keep_dcs

           -- Vectorise all the data type declarations that we can and must vectorise.
       ; new_tcs <- vectTyConDecls conv_tcs

           -- We don't need new representation types for dictionary constructors. The constructors
           -- are always fully applied, and we don't need to lift them to arrays as a dictionary
           -- of a particular type always has the same value.
       ; let vect_tcs = filter (not . isClassTyCon) 
                      $ keep_tcs ++ new_tcs

           -- Build 'PRepr' and 'PData' instance type constructors and family instances for all
           -- type constructors with vectorised representations.
       ; reprs     <- mapM tyConRepr vect_tcs
       ; repr_tcs  <- zipWith3M buildPReprTyCon orig_tcs vect_tcs reprs
       ; pdata_tcs <- zipWith3M buildPDataTyCon orig_tcs vect_tcs reprs
       ; let inst_tcs  = repr_tcs ++ pdata_tcs
             fam_insts = map mkLocalFamInst inst_tcs
       ; updGEnv $ extendFamEnv fam_insts

           -- Generate dfuns for the 'PA' instances of the vectorised type constructors and
           -- associate the type constructors with their dfuns in the global environment.  We get
           -- back the dfun bindings (which we will subsequently inject into the modules toplevel).
       ; (_, binds) <- fixV $ \ ~(dfuns, _) ->
           do { defTyConPAs (zipLazy vect_tcs dfuns)
              ; dfuns <- sequence 
                      $  zipWith4 buildTyConBindings
                                  orig_tcs
                                  vect_tcs
                                  repr_tcs
                                  pdata_tcs

              ; binds <- takeHoisted
              ; return (dfuns, binds)
              }

           -- We add to the type environment: (1) the vectorised type constructors, (2) their
           -- 'PRepr' & 'PData' instance constructors, and (3) the data constructors of the fomer
           -- two.
       ; let all_new_tcs = new_tcs ++ inst_tcs
             new_env     = extendTypeEnvList env
                         $ map ATyCon all_new_tcs ++
                           [ADataCon dc | tc <- all_new_tcs
                                        , dc <- tyConDataCons tc]

       ; return (new_env, fam_insts, binds)
       }


-- Helpers -------------------

buildTyConBindings :: TyCon -> TyCon -> TyCon -> TyCon -> VM Var
buildTyConBindings orig_tc vect_tc prepr_tc pdata_tc
 = do { vectDataConWorkers orig_tc vect_tc pdata_tc
      ; repr <- tyConRepr vect_tc
      ; buildPADict vect_tc prepr_tc pdata_tc repr
      }
rl@cse.unsw.edu.au's avatar
rl@cse.unsw.edu.au committed
190

191 192
vectDataConWorkers :: TyCon -> TyCon -> TyCon -> VM ()
vectDataConWorkers orig_tc vect_tc arr_tc
193
 = do bs <- sequence
194 195 196
          . zipWith3 def_worker  (tyConDataCons orig_tc) rep_tys
          $ zipWith4 mk_data_con (tyConDataCons vect_tc)
                                 rep_tys
197 198
                                 (inits rep_tys)
                                 (tail $ tails rep_tys)
199
      mapM_ (uncurry hoistBinding) bs
200
 where
201 202 203 204 205
    tyvars   = tyConTyVars vect_tc
    var_tys  = mkTyVarTys tyvars
    ty_args  = map Type var_tys
    res_ty   = mkTyConApp vect_tc var_tys

206 207 208 209
    cons     = tyConDataCons vect_tc
    arity    = length cons
    [arr_dc] = tyConDataCons arr_tc

210 211 212 213 214
    rep_tys  = map dataConRepArgTys $ tyConDataCons vect_tc


    mk_data_con con tys pre post
      = liftM2 (,) (vect_data_con con)
rl@cse.unsw.edu.au's avatar
rl@cse.unsw.edu.au committed
215
                   (lift_data_con tys pre post (mkDataConTag con))
216

217 218 219 220 221 222 223
    sel_replicate len tag
      | arity > 1 = do
                      rep <- builtin (selReplicate arity)
                      return [rep `mkApps` [len, tag]]

      | otherwise = return []

224
    vect_data_con con = return $ mkConApp con ty_args
225
    lift_data_con tys pre_tys post_tys tag
226 227
      = do
          len  <- builtin liftingContext
Ian Lynagh's avatar
Ian Lynagh committed
228
          args <- mapM (newLocalVar (fsLit "xs"))
229
                  =<< mapM mkPDataType tys
rl@cse.unsw.edu.au's avatar
rl@cse.unsw.edu.au committed
230

231
          sel  <- sel_replicate (Var len) tag
rl@cse.unsw.edu.au's avatar
rl@cse.unsw.edu.au committed
232

233 234
          pre   <- mapM emptyPD (concat pre_tys)
          post  <- mapM emptyPD (concat post_tys)
235 236 237 238

          return . mkLams (len : args)
                 . wrapFamInstBody arr_tc var_tys
                 . mkConApp arr_dc
239
                 $ ty_args ++ sel ++ pre ++ map Var args ++ post
240 241 242

    def_worker data_con arg_tys mk_body
      = do
243
          arity <- polyArity tyvars
244 245
          body <- closedV
                . inBind orig_worker
246 247
                . polyAbstract tyvars $ \args ->
                  liftM (mkLams (tyvars ++ args) . vectorised)
248 249
                $ buildClosures tyvars [] arg_tys res_ty mk_body

250
          raw_worker <- mkVectId orig_worker (exprType body)
251
          let vect_worker = raw_worker `setIdUnfolding`
252
                              mkInlineUnfolding (Just arity) body
253 254 255 256
          defGlobalVar orig_worker vect_worker
          return (vect_worker, body)
      where
        orig_worker = dataConWorkId data_con