Vectorise.hs 6.52 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           ( mkInlineRule )
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 171 172 173 174

      -- Make the vectorised version of binding's name, and set the unfolding used for inlining.
      var' <- liftM (`setIdUnfolding` unfolding) 
           $  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 -> mkInlineRule expr (Just arity)
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