Vectorise.hs 6.97 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 )
rl@cse.unsw.edu.au's avatar
rl@cse.unsw.edu.au committed
21
import Var
rl@cse.unsw.edu.au's avatar
rl@cse.unsw.edu.au committed
22
import Id
23
import OccName
24
import BasicTypes           ( isLoopBreaker )
25
import Outputable
26
import Util                 ( zipLazy )
27 28
import MonadUtils

29
import Control.Monad
30 31 32 33 34 35

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.
36
vectorise :: PackageId -> ModGuts -> CoreM ModGuts
37 38 39 40
vectorise backend guts 
 = do hsc_env <- getHscEnv
      liftIO $ vectoriseIO backend hsc_env guts

41

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

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

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

55 56

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

64
      (_, fam_inst_env) <- readGEnv global_fam_inst_env
Ian Lynagh's avatar
Ian Lynagh committed
65

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

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

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

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

-- | 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.
--
115
vectTopBind :: CoreBind -> VM CoreBind
116
vectTopBind b@(NonRec var expr)
117
 = do
118
      (inline, _, expr') 	<- vectTopRhs [] var expr
119 120 121 122 123 124 125 126 127
      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

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

vectTopBind b@(Rec bs)
133 134 135 136 137
 = do
      (vars', _, exprs') 
	<- fixV $ \ ~(_, inlines, rhss) ->
            do vars' <- sequence [vectTopBinder var inline rhs
                                      | (var, ~(inline, rhs)) <- zipLazy vars (zip inlines rhss)]
138 139 140 141 142 143
               (inlines', areScalars', exprs') 
                     <- mapAndUnzip3M (uncurry $ vectTopRhs vars) bs
               if  (and areScalars') || (length bs <= 1)
                  then do
                    return (vars', inlines', exprs')
                  else do
144
                    _ <- mapM deleteGlobalScalar vars
145 146 147
                    (inlines'', _, exprs'')  <- mapAndUnzip3M (uncurry $ vectTopRhs []) bs
                    return (vars', inlines'', exprs'')
                      
148
      hs     <- takeHoisted
149 150
      cexprs <- sequence $ zipWith3 tryConvert vars vars' exprs
      return . Rec $ zip vars cexprs ++ zip vars' exprs' ++ hs
151 152 153 154
  `orElseV`
    return b
  where
    (vars, exprs) = unzip bs
155 156 157
    mapAndUnzip3M f xs = do
       ys <- mapM f xs
       return $ unzip3 ys
158 159 160 161 162 163 164 165 166 167 168 169 170

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

171
vectTopBinder var inline expr
172 173
 = do
      -- Vectorise the type attached to the var.
rl@cse.unsw.edu.au's avatar
rl@cse.unsw.edu.au committed
174
      vty  <- vectType (idType var)
175 176

      -- Make the vectorised version of binding's name, and set the unfolding used for inlining.
177
      var' <- liftM (`setIdUnfoldingLazily` unfolding) 
178 179 180
           $  cloneId mkVectOcc var vty

      -- Add the mapping between the plain and vectorised name to the state.
181
      defGlobalVar var var'
182

183
      return var'
184 185
  where
    unfolding = case inline of
186
                  Inline arity -> mkInlineUnfolding (Just arity) expr
187
                  DontInline   -> noUnfolding
Ian Lynagh's avatar
Ian Lynagh committed
188

189 190 191

-- | Vectorise the RHS of a top-level binding, in an empty local environment.
vectTopRhs 
192 193
	:: [Var]    -- ^ Names of all functions in the rec block
	-> Var 		-- ^ Name of the binding.
194
	-> CoreExpr	-- ^ Body of the binding.
195
	-> VM (Inline, Bool, CoreExpr)
196

197
vectTopRhs recFs var expr
198 199
 = dtrace (vcat [text "vectTopRhs", ppr expr])
 $ closedV
200 201
 $ do (inline, isScalar, vexpr) <- 
           inBind var $ vectPolyExpr  (isLoopBreaker $ idOccInfo var) recFs (freeVars expr)
202 203
      if isScalar 
         then addGlobalScalar var
204 205
         else deleteGlobalScalar var
      return (inline, isScalar, vectorised vexpr)
206

207 208 209 210 211 212 213 214 215

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

216 217 218
tryConvert var vect_var rhs
  = fromVect (idType var) (Var vect_var) `orElseV` return rhs