Vectorise.hs 6.54 KB
Newer Older
1
{-# OPTIONS -fno-warn-missing-signatures #-}
2

3
4
5
module Vectorise( vectorise )
where

6
7
8
import Vectorise.Type.Env
import Vectorise.Type.Type
import Vectorise.Convert
9
import Vectorise.Utils.Hoisting
10
import Vectorise.Exp
11
import Vectorise.Vect
12
import Vectorise.Env
13
import Vectorise.Monad
14

15
import HscTypes hiding      ( MonadThings(..) )
16
import Module               ( PackageId )
17
import CoreSyn
18
import CoreUnfold           ( mkInlineUnfolding )
19
import CoreFVs
Ian Lynagh's avatar
Ian Lynagh committed
20
import CoreMonad            ( CoreM, getHscEnv )
21
import FamInstEnv           ( extendFamInstEnvList )
rl@cse.unsw.edu.au's avatar
rl@cse.unsw.edu.au committed
22
import Var
rl@cse.unsw.edu.au's avatar
rl@cse.unsw.edu.au committed
23
import Id
24
import OccName
25
import BasicTypes           ( isLoopBreaker )
26
import Outputable
27
28
import Util                 ( zipLazy )
import Control.Monad
29
30
31
32
33
34

debug		= False
dtrace s x	= if debug then pprTrace "Vectorise" s x else x

-- | Vectorise a single module.
--   Takes the package containing the DPH backend we're using. Eg either dph-par or dph-seq.
35
vectorise :: PackageId -> ModGuts -> CoreM ModGuts
36
37
38
39
vectorise backend guts 
 = do hsc_env <- getHscEnv
      liftIO $ vectoriseIO backend hsc_env guts

40

41
-- | Vectorise a single monad, given its HscEnv (code gen environment).
42
43
vectoriseIO :: PackageId -> HscEnv -> ModGuts -> IO ModGuts
vectoriseIO backend hsc_env guts
44
 = do -- Get information about currently loaded external packages.
rl@cse.unsw.edu.au's avatar
rl@cse.unsw.edu.au committed
45
      eps <- hscEPS hsc_env
46
47

      -- Combine vectorisation info from the current module, and external ones.
rl@cse.unsw.edu.au's avatar
rl@cse.unsw.edu.au committed
48
      let info = hptVectInfo hsc_env `plusVectInfo` eps_vect_info eps
49
50

      -- Run the main VM computation.
51
      Just (info', guts') <- initV backend hsc_env guts info (vectModule guts)
52
      return (guts' { mg_vect_info = info' })
53

54
55

-- | Vectorise a single module, in the VM monad.
rl@cse.unsw.edu.au's avatar
rl@cse.unsw.edu.au committed
56
vectModule :: ModGuts -> VM ModGuts
57
vectModule guts
58
59
60
 = do -- Vectorise the type environment.
      -- This may add new TyCons and DataCons.
      -- TODO: What new binds do we get back here?
61
      (types', fam_insts, tc_binds) <- vectTypeEnv (mg_types guts)
Ian Lynagh's avatar
Ian Lynagh committed
62

63
      -- TODO: What is this?
64
65
      let fam_inst_env' = extendFamInstEnvList (mg_fam_inst_env guts) fam_insts
      updGEnv (setFamInstEnv fam_inst_env')
Ian Lynagh's avatar
Ian Lynagh committed
66

67
68
      -- dicts   <- mapM buildPADict pa_insts
      -- workers <- mapM vectDataConWorkers pa_insts
69
70

      -- Vectorise all the top level bindings.
rl@cse.unsw.edu.au's avatar
rl@cse.unsw.edu.au committed
71
      binds'  <- mapM vectTopBind (mg_binds guts)
72

73
      return $ guts { mg_types        = types'
74
                    , mg_binds        = Rec tc_binds : binds'
75
76
77
                    , mg_fam_inst_env = fam_inst_env'
                    , mg_fam_insts    = mg_fam_insts guts ++ fam_insts
                    }
rl@cse.unsw.edu.au's avatar
rl@cse.unsw.edu.au committed
78

79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115

-- | Try to vectorise a top-level binding.
--   If it doesn't vectorise then return it unharmed.
--
--   For example, for the binding 
--
--   @  
--      foo :: Int -> Int
--      foo = \x -> x + x
--   @
--  
--   we get
--   @
--      foo  :: Int -> Int
--      foo  = \x -> vfoo $: x                  
-- 
--      v_foo :: Closure void vfoo lfoo
--      v_foo = closure vfoo lfoo void        
-- 
--      vfoo :: Void -> Int -> Int
--      vfoo = ...
--
--      lfoo :: PData Void -> PData Int -> PData Int
--      lfoo = ...
--   @ 
--
--   @vfoo@ is the "vectorised", or scalar, version that does the same as the original
--   function foo, but takes an explicit environment.
-- 
--   @lfoo@ is the "lifted" version that works on arrays.
--
--   @v_foo@ combines both of these into a `Closure` that also contains the
--   environment.
--
--   The original binding @foo@ is rewritten to call the vectorised version
--   present in the closure.
--
116
vectTopBind :: CoreBind -> VM CoreBind
117
vectTopBind b@(NonRec var expr)
118
119
120
121
122
123
124
125
126
127
128
 = do
      (inline, expr') 	<- vectTopRhs var expr
      var' 		<- vectTopBinder var inline expr'

      -- Vectorising the body may create other top-level bindings.
      hs	<- takeHoisted

      -- To get the same functionality as the original body we project
      -- out its vectorised version from the closure.
      cexpr	<- tryConvert var var' expr

129
      return . Rec $ (var, cexpr) : (var', expr') : hs
130
131
132
133
  `orElseV`
    return b

vectTopBind b@(Rec bs)
134
135
136
137
138
139
140
141
142
143
 = do
      (vars', _, exprs') 
	<- fixV $ \ ~(_, inlines, rhss) ->
            do vars' <- sequence [vectTopBinder var inline rhs
                                      | (var, ~(inline, rhs)) <- zipLazy vars (zip inlines rhss)]
               (inlines', exprs') 
                     <- mapAndUnzipM (uncurry vectTopRhs) bs

               return (vars', inlines', exprs')

144
      hs     <- takeHoisted
145
146
      cexprs <- sequence $ zipWith3 tryConvert vars vars' exprs
      return . Rec $ zip vars cexprs ++ zip vars' exprs' ++ hs
147
148
149
150
151
  `orElseV`
    return b
  where
    (vars, exprs) = unzip bs

152
153
154
155
156
157
158
159
160
161
162
163
164

-- | Make the vectorised version of this top level binder, and add the mapping
--   between it and the original to the state. For some binder @foo@ the vectorised
--   version is @$v_foo@
--
--   NOTE: vectTopBinder *MUST* be lazy in inline and expr because of how it is
--   used inside of fixV in vectTopBind
vectTopBinder 
	:: Var 		-- ^ Name of the binding.
	-> Inline 	-- ^ Whether it should be inlined, used to annotate it.
	-> CoreExpr 	-- ^ RHS of the binding, used to set the `Unfolding` of the returned `Var`.
	-> VM Var	-- ^ Name of the vectorised binding.

165
vectTopBinder var inline expr
166
167
 = do
      -- Vectorise the type attached to the var.
rl@cse.unsw.edu.au's avatar
rl@cse.unsw.edu.au committed
168
      vty  <- vectType (idType var)
169
170

      -- Make the vectorised version of binding's name, and set the unfolding used for inlining.
171
      var' <- liftM (`setIdUnfoldingLazily` unfolding) 
172
173
174
           $  cloneId mkVectOcc var vty

      -- Add the mapping between the plain and vectorised name to the state.
175
      defGlobalVar var var'
176

177
      return var'
178
179
  where
    unfolding = case inline of
180
                  Inline arity -> mkInlineUnfolding (Just arity) expr
181
                  DontInline   -> noUnfolding
Ian Lynagh's avatar
Ian Lynagh committed
182

183
184
185
186
187
188
189

-- | Vectorise the RHS of a top-level binding, in an empty local environment.
vectTopRhs 
	:: Var 		-- ^ Name of the binding.
	-> CoreExpr	-- ^ Body of the binding.
	-> VM (Inline, CoreExpr)

190
vectTopRhs var expr
191
192
193
194
 = dtrace (vcat [text "vectTopRhs", ppr expr])
 $ closedV
 $ do (inline, vexpr) <- inBind var
                      $ vectPolyExpr (isLoopBreaker $ idOccInfo var)
195
196
                                      (freeVars expr)
      return (inline, vectorised vexpr)
197

198
199
200
201
202
203
204
205
206

-- | Project out the vectorised version of a binding from some closure,
--	or return the original body if that doesn't work.	
tryConvert 
	:: Var	 	-- ^ Name of the original binding (eg @foo@)
	-> Var 		-- ^ Name of vectorised version of binding (eg @$vfoo@)
	-> CoreExpr	-- ^ The original body of the binding.
	-> VM CoreExpr

207
208
209
tryConvert var vect_var rhs
  = fromVect (idType var) (Var vect_var) `orElseV` return rhs