VectMonad.hs 16 KB
Newer Older
1
{-# LANGUAGE NamedFieldPuns #-}
benl@ouroborus.net's avatar
benl@ouroborus.net committed
2
3

-- | The Vectorisation monad.
4
5
6
module VectMonad (
  VM,

7
8
  noV, traceNoV, ensureV, traceEnsureV, tryV, maybeV, traceMaybeV, orElseV,
  onlyIfV, fixV, localV, closedV,
9
  initV, cantVectorise, maybeCantVectorise, maybeCantVectoriseM,
10
  liftDs,
rl@cse.unsw.edu.au's avatar
rl@cse.unsw.edu.au committed
11
  cloneName, cloneId, cloneVar,
12
  newExportedVar, newLocalVar, newLocalVars, newDummyVar, newTyVar,
13
  
14
  Builtins(..), sumTyCon, prodTyCon, prodDataCon,
15
  selTy, selReplicate, selPick, selTags, selElements,
16
  combinePDVar, scalarZip, closureCtrFun,
17
  builtin, builtins,
18

19
  setFamInstEnv,
20
21
22
23
  readGEnv, setGEnv, updGEnv,

  readLEnv, setLEnv, updLEnv,

24
25
  getBindName, inBind,

26
  lookupVar, defGlobalVar, globalScalars,
27
28
  lookupTyCon, defTyCon,
  lookupDataCon, defDataCon,
29
  lookupTyConPA, defTyConPA, defTyConPAs,
30
  lookupTyConPR,
31
  lookupBoxedTyCon,
32
  lookupPrimMethod, lookupPrimPArray,
33
  lookupTyVarPA, defLocalTyVar, defLocalTyVarWithPA, localTyVars,
34

35
  lookupInst, lookupFamInst
36
37
38
39
) where

#include "HsVersions.h"

40
import VectBuiltIn
41
import Vectorise.Env
42
import Vectorise.Vect
43

44
import HscTypes hiding  ( MonadThings(..) )
45
import Module           ( PackageId )
46
import CoreSyn
47
import Class
48
import TyCon
49
import DataCon
50
51
import Type
import Var
52
import VarSet
53
54
55
56
57
58
59
import VarEnv
import Id
import Name
import NameEnv

import DsMonad

60
61
62
63
import InstEnv
import FamInstEnv

import Outputable
64
import FastString
65
import SrcLoc        ( noSrcSpan )
66

twanvl's avatar
twanvl committed
67
import Control.Monad
rl@cse.unsw.edu.au's avatar
rl@cse.unsw.edu.au committed
68

benl@ouroborus.net's avatar
benl@ouroborus.net committed
69
70
71
72
73
74

-- The Vectorisation Monad ----------------------------------------------------

-- Vectorisation can either succeed with new envionment and a value,
-- or return with failure.
--
75
76
77
78
79
data VResult a = Yes GlobalEnv LocalEnv a | No

newtype VM a = VM { runVM :: Builtins -> GlobalEnv -> LocalEnv -> DsM (VResult a) }

instance Monad VM where
twanvl's avatar
twanvl committed
80
  return x   = VM $ \_  genv lenv -> return (Yes genv lenv x)
81
82
83
84
85
86
  VM p >>= f = VM $ \bi genv lenv -> do
                                      r <- p bi genv lenv
                                      case r of
                                        Yes genv' lenv' x -> runVM (f x) bi genv' lenv'
                                        No                -> return No

87

benl@ouroborus.net's avatar
benl@ouroborus.net committed
88
-- | Throw an error saying we can't vectorise something
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
cantVectorise :: String -> SDoc -> a
cantVectorise s d = pgmError
                  . showSDocDump
                  $ vcat [text "*** Vectorisation error ***",
                          nest 4 $ sep [text s, nest 4 d]]

maybeCantVectorise :: String -> SDoc -> Maybe a -> a
maybeCantVectorise s d Nothing  = cantVectorise s d
maybeCantVectorise _ _ (Just x) = x

maybeCantVectoriseM :: Monad m => String -> SDoc -> m (Maybe a) -> m a
maybeCantVectoriseM s d p
  = do
      r <- p
      case r of
        Just x  -> return x
        Nothing -> cantVectorise s d

benl@ouroborus.net's avatar
benl@ouroborus.net committed
107
108
109

-- Control --------------------------------------------------------------------
-- | Return some result saying we've failed.
110
111
112
noV :: VM a
noV = VM $ \_ _ _ -> return No

113
114
115
traceNoV :: String -> SDoc -> VM a
traceNoV s d = pprTrace s d noV

benl@ouroborus.net's avatar
benl@ouroborus.net committed
116
117

-- | If True then carry on, otherwise fail.
118
119
120
121
ensureV :: Bool -> VM ()
ensureV False = noV
ensureV True  = return ()

benl@ouroborus.net's avatar
benl@ouroborus.net committed
122
123

-- | If True then return the first argument, otherwise fail.
124
125
126
127
128
onlyIfV :: Bool -> VM a -> VM a
onlyIfV b p = ensureV b >> p

traceEnsureV :: String -> SDoc -> Bool -> VM ()
traceEnsureV s d False = traceNoV s d
rl@cse.unsw.edu.au's avatar
rl@cse.unsw.edu.au committed
129
traceEnsureV _ _ True  = return ()
130

benl@ouroborus.net's avatar
benl@ouroborus.net committed
131
132
133
134

-- | Try some vectorisation computaton.
--	If it succeeds then return Just the result,
--	otherwise return Nothing.
135
136
137
138
139
140
141
142
tryV :: VM a -> VM (Maybe a)
tryV (VM p) = VM $ \bi genv lenv ->
  do
    r <- p bi genv lenv
    case r of
      Yes genv' lenv' x -> return (Yes genv' lenv' (Just x))
      No                -> return (Yes genv  lenv  Nothing)

benl@ouroborus.net's avatar
benl@ouroborus.net committed
143

144
145
146
maybeV :: VM (Maybe a) -> VM a
maybeV p = maybe noV return =<< p

147
148
149
traceMaybeV :: String -> SDoc -> VM (Maybe a) -> VM a
traceMaybeV s d p = maybe (traceNoV s d) return =<< p

150
151
152
orElseV :: VM a -> VM a -> VM a
orElseV p q = maybe q return =<< tryV p

rl@cse.unsw.edu.au's avatar
rl@cse.unsw.edu.au committed
153
fixV :: (a -> VM a) -> VM a
rl@cse.unsw.edu.au's avatar
rl@cse.unsw.edu.au committed
154
155
156
157
158
159
fixV f = VM (\bi genv lenv -> fixDs $ \r -> runVM (f (unYes r)) bi genv lenv )
  where
    -- NOTE: It is essential that we are lazy in r above so do not replace
    --       calls to this function by an explicit case.
    unYes (Yes _ _ x) = x
    unYes No          = panic "VectMonad.fixV: no result"
rl@cse.unsw.edu.au's avatar
rl@cse.unsw.edu.au committed
160

benl@ouroborus.net's avatar
benl@ouroborus.net committed
161
162
163
164

-- Local Environments ---------------------------------------------------------
-- | Perform a computation in its own local environment.
--	This does not alter the environment of the current state.
165
166
167
168
169
170
171
localV :: VM a -> VM a
localV p = do
             env <- readLEnv id
             x <- p
             setLEnv env
             return x

benl@ouroborus.net's avatar
benl@ouroborus.net committed
172
-- | Perform a computation in an empty local environment.
rl@cse.unsw.edu.au's avatar
rl@cse.unsw.edu.au committed
173
174
175
closedV :: VM a -> VM a
closedV p = do
              env <- readLEnv id
176
              setLEnv (emptyLocalEnv { local_bind_name = local_bind_name env })
rl@cse.unsw.edu.au's avatar
rl@cse.unsw.edu.au committed
177
178
179
180
              x <- p
              setLEnv env
              return x

benl@ouroborus.net's avatar
benl@ouroborus.net committed
181
182
-- Lifting --------------------------------------------------------------------
-- | Lift a desugaring computation into the vectorisation monad.
183
liftDs :: DsM a -> VM a
twanvl's avatar
twanvl committed
184
liftDs p = VM $ \_ genv lenv -> do { x <- p; return (Yes genv lenv x) }
185

benl@ouroborus.net's avatar
benl@ouroborus.net committed
186
187
188
189


-- Builtins -------------------------------------------------------------------
-- Operations on Builtins
190
191
192
liftBuiltinDs :: (Builtins -> DsM a) -> VM a
liftBuiltinDs p = VM $ \bi genv lenv -> do { x <- p bi; return (Yes genv lenv x)}

benl@ouroborus.net's avatar
benl@ouroborus.net committed
193
194

-- | Project something from the set of builtins.
195
196
197
builtin :: (Builtins -> a) -> VM a
builtin f = VM $ \bi genv lenv -> return (Yes genv lenv (f bi))

198
199
200
builtins :: (a -> Builtins -> b) -> VM (a -> b)
builtins f = VM $ \bi genv lenv -> return (Yes genv lenv (`f` bi))

benl@ouroborus.net's avatar
benl@ouroborus.net committed
201
202
203

-- Environments ---------------------------------------------------------------
-- | Project something from the global environment.
204
readGEnv :: (GlobalEnv -> a) -> VM a
twanvl's avatar
twanvl committed
205
readGEnv f = VM $ \_ genv lenv -> return (Yes genv lenv (f genv))
206
207
208
209
210
211
212

setGEnv :: GlobalEnv -> VM ()
setGEnv genv = VM $ \_ _ lenv -> return (Yes genv lenv ())

updGEnv :: (GlobalEnv -> GlobalEnv) -> VM ()
updGEnv f = VM $ \_ genv lenv -> return (Yes (f genv) lenv ())

benl@ouroborus.net's avatar
benl@ouroborus.net committed
213
214

-- | Project something from the local environment.
215
readLEnv :: (LocalEnv -> a) -> VM a
twanvl's avatar
twanvl committed
216
readLEnv f = VM $ \_ genv lenv -> return (Yes genv lenv (f lenv))
217

218
-- | Set the local environment.
219
220
221
setLEnv :: LocalEnv -> VM ()
setLEnv lenv = VM $ \_ genv _ -> return (Yes genv lenv ())

222
-- | Update the enviroment using a provided function.
223
224
225
updLEnv :: (LocalEnv -> LocalEnv) -> VM ()
updLEnv f = VM $ \_ genv lenv -> return (Yes genv (f lenv) ())

benl@ouroborus.net's avatar
benl@ouroborus.net committed
226
227

-- InstEnv --------------------------------------------------------------------
228
229
230
231
232
233
getInstEnv :: VM (InstEnv, InstEnv)
getInstEnv = readGEnv global_inst_env

getFamInstEnv :: VM FamInstEnvs
getFamInstEnv = readGEnv global_fam_inst_env

benl@ouroborus.net's avatar
benl@ouroborus.net committed
234
235
236

-- Names ----------------------------------------------------------------------
-- | Get the name of the local binding currently being vectorised.
237
238
239
240
241
242
243
244
getBindName :: VM FastString
getBindName = readLEnv local_bind_name

inBind :: Id -> VM a -> VM a
inBind id p
  = do updLEnv $ \env -> env { local_bind_name = occNameFS (getOccName id) }
       p

245
246
247
248
249
250
251
252
253
254
cloneName :: (OccName -> OccName) -> Name -> VM Name
cloneName mk_occ name = liftM make (liftDs newUnique)
  where
    occ_name = mk_occ (nameOccName name)

    make u | isExternalName name = mkExternalName u (nameModule name)
                                                    occ_name
                                                    (nameSrcSpan name)
           | otherwise           = mkSystemName u occ_name

rl@cse.unsw.edu.au's avatar
rl@cse.unsw.edu.au committed
255
256
257
258
259
260
261
262
cloneId :: (OccName -> OccName) -> Id -> Type -> VM Id
cloneId mk_occ id ty
  = do
      name <- cloneName mk_occ (getName id)
      let id' | isExportedId id = Id.mkExportedLocalId name ty
              | otherwise       = Id.mkLocalId         name ty
      return id'

benl@ouroborus.net's avatar
benl@ouroborus.net committed
263
-- Make a fresh instance of this var, with a new unique.
rl@cse.unsw.edu.au's avatar
rl@cse.unsw.edu.au committed
264
265
266
cloneVar :: Var -> VM Var
cloneVar var = liftM (setIdUnique var) (liftDs newUnique)

267
268
269
270
271
272
273
274
275
276
newExportedVar :: OccName -> Type -> VM Var
newExportedVar occ_name ty 
  = do
      mod <- liftDs getModuleDs
      u   <- liftDs newUnique

      let name = mkExternalName u mod occ_name noSrcSpan
      
      return $ Id.mkExportedLocalId name ty

277
278
279
280
281
282
newLocalVar :: FastString -> Type -> VM Var
newLocalVar fs ty
  = do
      u <- liftDs newUnique
      return $ mkSysLocal fs u ty

283
284
285
newLocalVars :: FastString -> [Type] -> VM [Var]
newLocalVars fs = mapM (newLocalVar fs)

286
newDummyVar :: Type -> VM Var
rl@cse.unsw.edu.au's avatar
rl@cse.unsw.edu.au committed
287
newDummyVar = newLocalVar (fsLit "vv")
288

289
290
291
292
293
294
newTyVar :: FastString -> Kind -> VM Var
newTyVar fs k
  = do
      u <- liftDs newUnique
      return $ mkTyVar (mkSysTvName u fs) k

295
296

-- | Add a mapping between a global var and its vectorised version to the state.
297
298
defGlobalVar :: Var -> Var -> VM ()
defGlobalVar v v' = updGEnv $ \env ->
299
  env { global_vars = extendVarEnv (global_vars env) v v'
300
301
302
303
304
      , global_exported_vars = upd (global_exported_vars env)
      }
  where
    upd env | isExportedId v = extendVarEnv env v (v, v')
            | otherwise      = env
rl@cse.unsw.edu.au's avatar
rl@cse.unsw.edu.au committed
305

306
-- Var ------------------------------------------------------------------------
benl@ouroborus.net's avatar
benl@ouroborus.net committed
307
308
309
310
-- | Lookup the vectorised and\/or lifted versions of this variable.
--	If it's in the global environment we get the vectorised version.
--      If it's in the local environment we get both the vectorised and lifted version.
--	
311
lookupVar :: Var -> VM (Scope Var (Var, Var))
rl@cse.unsw.edu.au's avatar
rl@cse.unsw.edu.au committed
312
lookupVar v
313
 = do r <- readLEnv $ \env -> lookupVarEnv (local_vars env) v
rl@cse.unsw.edu.au's avatar
rl@cse.unsw.edu.au committed
314
315
316
      case r of
        Just e  -> return (Local e)
        Nothing -> liftM Global
317
                . maybeCantVectoriseVarM v
318
                . readGEnv $ \env -> lookupVarEnv (global_vars env) v
rl@cse.unsw.edu.au's avatar
rl@cse.unsw.edu.au committed
319

320
321
322
323
324
325
326
327
328
maybeCantVectoriseVarM :: Monad m => Var -> m (Maybe Var) -> m Var
maybeCantVectoriseVarM v p
 = do r <- p
      case r of
        Just x  -> return x
        Nothing -> dumpVar v

dumpVar :: Var -> a
dumpVar var
benl@ouroborus.net's avatar
benl@ouroborus.net committed
329
	| Just _		<- isClassOpId_maybe var
330
331
332
333
334
335
	= cantVectorise "ClassOpId not vectorised:" (ppr var)

	| otherwise
	= cantVectorise "Variable not vectorised:" (ppr var)

-------------------------------------------------------------------------------
336
337
338
globalScalars :: VM VarSet
globalScalars = readGEnv global_scalars

339
lookupTyCon :: TyCon -> VM (Maybe TyCon)
340
341
342
343
lookupTyCon tc
  | isUnLiftedTyCon tc || isTupleTyCon tc = return (Just tc)

  | otherwise = readGEnv $ \env -> lookupNameEnv (global_tycons env) (tyConName tc)
344

345
346
347
348
349
defTyCon :: TyCon -> TyCon -> VM ()
defTyCon tc tc' = updGEnv $ \env ->
  env { global_tycons = extendNameEnv (global_tycons env) (tyConName tc) tc' }

lookupDataCon :: DataCon -> VM (Maybe DataCon)
350
351
352
lookupDataCon dc
  | isTupleTyCon (dataConTyCon dc) = return (Just dc)
  | otherwise = readGEnv $ \env -> lookupNameEnv (global_datacons env) (dataConName dc)
353
354
355
356
357

defDataCon :: DataCon -> DataCon -> VM ()
defDataCon dc dc' = updGEnv $ \env ->
  env { global_datacons = extendNameEnv (global_datacons env) (dataConName dc) dc' }

358
lookupPrimPArray :: TyCon -> VM (Maybe TyCon)
359
lookupPrimPArray = liftBuiltinDs . primPArray
360

361
lookupPrimMethod :: TyCon -> String -> VM (Maybe Var)
362
lookupPrimMethod tycon = liftBuiltinDs . primMethod tycon
363

364
365
366
367
368
369
370
lookupTyConPA :: TyCon -> VM (Maybe Var)
lookupTyConPA tc = readGEnv $ \env -> lookupNameEnv (global_pa_funs env) (tyConName tc)

defTyConPA :: TyCon -> Var -> VM ()
defTyConPA tc pa = updGEnv $ \env ->
  env { global_pa_funs = extendNameEnv (global_pa_funs env) (tyConName tc) pa }

371
372
373
374
375
defTyConPAs :: [(TyCon, Var)] -> VM ()
defTyConPAs ps = updGEnv $ \env ->
  env { global_pa_funs = extendNameEnvList (global_pa_funs env)
                                           [(tyConName tc, pa) | (tc, pa) <- ps] }

376
lookupTyVarPA :: Var -> VM (Maybe CoreExpr)
377
378
379
380
lookupTyVarPA tv = readLEnv $ \env -> lookupVarEnv (local_tyvar_pa env) tv

lookupTyConPR :: TyCon -> VM (Maybe Var)
lookupTyConPR tc = readGEnv $ \env -> lookupNameEnv (global_pr_funs env) (tyConName tc)
381

382
383
384
385
lookupBoxedTyCon :: TyCon -> VM (Maybe TyCon)
lookupBoxedTyCon tc = readGEnv $ \env -> lookupNameEnv (global_boxed_tycons env)
                                                       (tyConName tc)

386
387
388
389
390
391
392
393
394
395
396
defLocalTyVar :: TyVar -> VM ()
defLocalTyVar tv = updLEnv $ \env ->
  env { local_tyvars   = tv : local_tyvars env
      , local_tyvar_pa = local_tyvar_pa env `delVarEnv` tv
      }

defLocalTyVarWithPA :: TyVar -> CoreExpr -> VM ()
defLocalTyVarWithPA tv pa = updLEnv $ \env ->
  env { local_tyvars   = tv : local_tyvars env
      , local_tyvar_pa = extendVarEnv (local_tyvar_pa env) tv pa
      }
397

398
399
localTyVars :: VM [TyVar]
localTyVars = readLEnv (reverse . local_tyvars)
400

401
402
403
404
405
406
407
408
409
410
411
412
413
414
415
416
417
418
419
420
-- Look up the dfun of a class instance.
--
-- The match must be unique - ie, match exactly one instance - but the 
-- type arguments used for matching may be more specific than those of 
-- the class instance declaration.  The found class instances must not have
-- any type variables in the instance context that do not appear in the
-- instances head (i.e., no flexi vars); for details for what this means,
-- see the docs at InstEnv.lookupInstEnv.
--
lookupInst :: Class -> [Type] -> VM (DFunId, [Type])
lookupInst cls tys
  = do { instEnv <- getInstEnv
       ; case lookupInstEnv instEnv cls tys of
	   ([(inst, inst_tys)], _) 
             | noFlexiVar -> return (instanceDFunId inst, inst_tys')
             | otherwise  -> pprPanic "VectMonad.lookupInst: flexi var: " 
                                      (ppr $ mkTyConApp (classTyCon cls) tys)
             where
               inst_tys'  = [ty | Right ty <- inst_tys]
               noFlexiVar = all isRight inst_tys
421
422
	   _other         ->
             pprPanic "VectMonad.lookupInst: not found " (ppr cls <+> ppr tys)
423
424
425
426
427
428
429
430
431
432
433
434
435
436
437
438
439
440
441
442
443
444
445
446
447
448
449
450
451
452
453
454
       }
  where
    isRight (Left  _) = False
    isRight (Right _) = True

-- Look up the representation tycon of a family instance.
--
-- The match must be unique - ie, match exactly one instance - but the 
-- type arguments used for matching may be more specific than those of 
-- the family instance declaration.
--
-- Return the instance tycon and its type instance.  For example, if we have
--
--  lookupFamInst 'T' '[Int]' yields (':R42T', 'Int')
--
-- then we have a coercion (ie, type instance of family instance coercion)
--
--  :Co:R42T Int :: T [Int] ~ :R42T Int
--
-- which implies that :R42T was declared as 'data instance T [a]'.
--
lookupFamInst :: TyCon -> [Type] -> VM (TyCon, [Type])
lookupFamInst tycon tys
  = ASSERT( isOpenTyCon tycon )
    do { instEnv <- getFamInstEnv
       ; case lookupFamInstEnv instEnv tycon tys of
	   [(fam_inst, rep_tys)] -> return (famInstTyCon fam_inst, rep_tys)
	   _other                -> 
             pprPanic "VectMonad.lookupFamInst: not found: " 
                      (ppr $ mkTyConApp tycon tys)
       }

455
456

-- | Run a vectorisation computation.
457
458
initV :: PackageId -> HscEnv -> ModGuts -> VectInfo -> VM a -> IO (Maybe (VectInfo, a))
initV pkg hsc_env guts info p
459
  = do
460
461
462
         -- XXX: ignores error messages and warnings, check that this is
         -- indeed ok (the use of "Just r" suggests so)
      (_,Just r) <- initDs hsc_env (mg_module guts)
463
464
                               (mg_rdr_env guts)
                               (mg_types guts)
465
                               go
466
467
      return r
  where
468

469
    go =
470
      do
471
        builtins       <- initBuiltins pkg
472
473
474
        builtin_vars   <- initBuiltinVars builtins
        builtin_tycons <- initBuiltinTyCons builtins
        let builtin_datacons = initBuiltinDataCons builtins
475
        builtin_boxed  <- initBuiltinBoxedTyCons builtins
476
        builtin_scalars <- initBuiltinScalars builtins
477

478
        eps <- liftIO $ hscEPS hsc_env
479
480
481
        let famInstEnvs = (eps_fam_inst_env eps, mg_fam_inst_env guts)
            instEnvs    = (eps_inst_env     eps, mg_inst_env     guts)

482
483
484
        builtin_prs    <- initBuiltinPRs builtins instEnvs
        builtin_pas    <- initBuiltinPAs builtins instEnvs

485
        let genv = extendImportedVarsEnv builtin_vars
486
                 . extendScalars builtin_scalars
487
                 . extendTyConsEnv builtin_tycons
488
                 . extendDataConsEnv builtin_datacons
489
                 . extendPAFunsEnv builtin_pas
490
                 . setPRFunsEnv    builtin_prs
491
                 . setBoxedTyConsEnv builtin_boxed
492
493
494
                 $ initGlobalEnv info instEnvs famInstEnvs

        r <- runVM p builtins genv emptyLocalEnv
495
496
497
        case r of
          Yes genv _ x -> return $ Just (new_info genv, x)
          No           -> return Nothing
498
499
500

    new_info genv = updVectInfo genv (mg_types guts) info