Env.hs 12.5 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 14
#include "HsVersions.h"

15
import Vectorise.Env
16
import Vectorise.Vect
17 18
import Vectorise.Monad
import Vectorise.Builtins
19
import Vectorise.Type.TyConDecl
20
import Vectorise.Type.Classify
21
import Vectorise.Type.PADict
22 23 24
import Vectorise.Type.PData
import Vectorise.Type.PRepr
import Vectorise.Type.Repr
25
import Vectorise.Utils
26

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

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

46

47 48 49 50 51 52 53 54 55
-- 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.  
--
56
--     An example is the treatment of 'Int'.  'Int's can be used in vectorised code and remain
57 58 59
--     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.
60
--
61 62 63 64 65 66
--     '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.)
--
67 68 69 70 71 72
-- (2) Data type constructor 'T' that together with its constructors 'Cn' may be used in vectorised
--     code, where 'T' and the 'Cn' are automatically vectorised in the same manner as data types
--     declared in a vectorised module.  This includes the case where the vectoriser determines that
--     the original representation of 'T' may be used in vectorised code (as it does not embed any
--     parallel arrays.)  This case is for type constructors that are *imported* from a non-
--     vectorised module, but that we want to use with full vectorisation support.
73
--
74 75
--     An example is the treatment of 'Ordering' and '[]'.  The former remains unchanged by
--     vectorisation, whereas the latter is fully vectorised.
76 77 78 79 80

--     'PData' and 'PRepr' instances are automatically generated by the vectoriser.
--
--     Type constructors declared with {-# VECTORISE type T #-} are treated in this manner.
--
81 82 83
-- (3) Data type constructor 'T' that together with its constructors 'Cn' may be used in vectorised
--     code, where 'T' is represented by an explicitly given 'Tv' whose constructors 'Cvn' represent
--     the original constructors in vectorised code.  As a special case, we can have 'Tv = T'
84
--
85 86 87
--     An example is the treatment of 'Bool', which is represented by itself in vectorised code
--     (as it cannot embed any parallel arrays).  However, we do not want any automatic generation
--     of class and family instances, which is why Case (2) does not apply.
88
--
89 90
--     'PData' and 'PRepr' instances need to be explicitly supplied for 'T' (they are not generated
--     by the vectoriser).
91
--
92
--     Type constructors declared with {-# VECTORISE type T = T' #-} are treated in this manner.
93 94 95 96 97
--
-- In addition, we have also got a single pragma form for type classes: {-# VECTORISE class C #-}.  It
-- implies that the class type constructor may be used in vectorised code together with its data
-- constructor.  We generally produce a vectorised version of the data type and data constructor.
-- We do not generate 'PData' and 'PRepr' instances for class type constructors.
98 99 100

-- |Vectorise a type environment.
--
101
vectTypeEnv :: [TyCon]                  -- TyCons defined in this module
102
            -> [CoreVect]               -- All 'VECTORISE [SCALAR] type' declarations in this module
103
            -> VM ( [TyCon]             -- old TyCons ++ new TyCons
104 105
                  , [FamInst]           -- New type family instances.
                  , [(Var, CoreExpr)])  -- New top level bindings.
106 107
vectTypeEnv tycons vectTypeDecls
  = do { traceVt "** vectTypeEnv" $ ppr tycons
108 109 110 111 112 113 114 115 116

         -- 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

117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132
       ; let   -- {-# VECTORISE SCALAR type T -#} (imported and local tycons)
             localScalarTyCons      = [tycon | VectType True  tycon Nothing <- vectTypeDecls]

               -- {-# VECTORISE type T -#} (ONLY the imported tycons)
             impVectTyCons          = [tycon | VectType False tycon Nothing <- vectTypeDecls]
                                      \\ tycons

               -- {-# VECTORISE type T = ty -#} (imported and local tycons)
             vectTyConsWithRHS      = [ (tycon, rhs) 
                                      | VectType False tycon (Just rhs) <- vectTypeDecls]

               -- filter VECTORISE SCALAR tycons and VECTORISE tycons with explicit rhses
             vectSpecialTyConNames  = mkNameSet . map tyConName $
                                        localScalarTyCons ++ map fst vectTyConsWithRHS
             notLocalScalarTyCon tc = not $ (tyConName tc) `elemNameSet` vectSpecialTyConNames

133 134 135 136
           -- 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
137 138 139
           -- VECTORISE SCALAR pragma or a VECTORISE pragma with an explicit right-hand side, as
           -- these are being handled separately.)
       ; let maybeVectoriseTyCons   = filter notLocalScalarTyCon tycons ++ impVectTyCons
140 141 142
             (conv_tcs, keep_tcs)   = classifyTyCons vectTyConFlavour maybeVectoriseTyCons
             orig_tcs               = keep_tcs ++ conv_tcs
             
143 144
       ; traceVt " VECT SCALAR    : " $ ppr localScalarTyCons
       ; traceVt " VECT with rhs  : " $ ppr (map fst vectTyConsWithRHS)
145 146 147
       ; traceVt " reuse          : " $ ppr keep_tcs
       ; traceVt " convert        : " $ ppr conv_tcs

148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165
       ; let defTyConDataCons origTyCon vectTyCon
               = do { defTyCon origTyCon vectTyCon
                    ; MASSERT(length (tyConDataCons origTyCon) == length (tyConDataCons vectTyCon))
                    ; zipWithM_ defDataCon (tyConDataCons origTyCon) (tyConDataCons vectTyCon)
                    }

           -- For the type constructors that we don't need to vectorise, we use the original
           -- representation in both unvectorised and vectorised code.
       ; zipWithM_ defTyConDataCons keep_tcs keep_tcs

           -- We do the same for type constructors declared VECTORISE SCALAR, while ignoring their
           -- representation (data constructors) — see "Note [Pragmas to vectorise tycons]".
       ; zipWithM_ defTyCon localScalarTyCons localScalarTyCons

           -- For type constructors declared VECTORISE with an explicit vectorised type, we use the
           -- explicitly given type in vectorised code and map data constructors one for one — see
           -- "Note [Pragmas to vectorise tycons]".
       ; mapM_ (uncurry defTyConDataCons) vectTyConsWithRHS
166 167 168 169 170 171 172 173 174 175 176 177 178 179 180 181 182 183 184 185 186 187 188 189 190 191 192 193 194 195 196 197 198 199 200

           -- 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)
              }

201 202 203
           -- Return the vectorised variants of type constructors as well as the generated instance type
           -- constructors, family instances, and dfun bindings.
       ; return (new_tcs ++ inst_tcs, fam_insts, binds)
204 205 206 207 208 209 210 211 212 213 214
       }


-- 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
215

216 217
vectDataConWorkers :: TyCon -> TyCon -> TyCon -> VM ()
vectDataConWorkers orig_tc vect_tc arr_tc
218
 = do bs <- sequence
219 220 221
          . zipWith3 def_worker  (tyConDataCons orig_tc) rep_tys
          $ zipWith4 mk_data_con (tyConDataCons vect_tc)
                                 rep_tys
222 223
                                 (inits rep_tys)
                                 (tail $ tails rep_tys)
224
      mapM_ (uncurry hoistBinding) bs
225
 where
226 227 228 229 230
    tyvars   = tyConTyVars vect_tc
    var_tys  = mkTyVarTys tyvars
    ty_args  = map Type var_tys
    res_ty   = mkTyConApp vect_tc var_tys

231 232 233 234
    cons     = tyConDataCons vect_tc
    arity    = length cons
    [arr_dc] = tyConDataCons arr_tc

235 236 237 238 239
    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
240
                   (lift_data_con tys pre post (mkDataConTag con))
241

242 243 244 245 246 247 248
    sel_replicate len tag
      | arity > 1 = do
                      rep <- builtin (selReplicate arity)
                      return [rep `mkApps` [len, tag]]

      | otherwise = return []

249
    vect_data_con con = return $ mkConApp con ty_args
250
    lift_data_con tys pre_tys post_tys tag
251 252
      = do
          len  <- builtin liftingContext
Ian Lynagh's avatar
Ian Lynagh committed
253
          args <- mapM (newLocalVar (fsLit "xs"))
254
                  =<< mapM mkPDataType tys
rl@cse.unsw.edu.au's avatar
rl@cse.unsw.edu.au committed
255

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

258 259
          pre   <- mapM emptyPD (concat pre_tys)
          post  <- mapM emptyPD (concat post_tys)
260 261 262 263

          return . mkLams (len : args)
                 . wrapFamInstBody arr_tc var_tys
                 . mkConApp arr_dc
264
                 $ ty_args ++ sel ++ pre ++ map Var args ++ post
265 266 267

    def_worker data_con arg_tys mk_body
      = do
268
          arity <- polyArity tyvars
269 270
          body <- closedV
                . inBind orig_worker
271 272
                . polyAbstract tyvars $ \args ->
                  liftM (mkLams (tyvars ++ args) . vectorised)
273 274
                $ buildClosures tyvars [] arg_tys res_ty mk_body

275
          raw_worker <- mkVectId orig_worker (exprType body)
276
          let vect_worker = raw_worker `setIdUnfolding`
277
                              mkInlineUnfolding (Just arity) body
278 279 280 281
          defGlobalVar orig_worker vect_worker
          return (vect_worker, body)
      where
        orig_worker = dataConWorkId data_con