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

3 4 5
module Vectorise( vectorise )
where

6
import VectType
7
import Vectorise.Utils.Hoisting
8
import Vectorise.Exp
9
import Vectorise.Vect
10
import Vectorise.Env
11
import Vectorise.Monad
12

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

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

38

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

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

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

52 53

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

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

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

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

71
      return $ guts { mg_types        = types'
72
                    , mg_binds        = Rec tc_binds : binds'
73 74 75
                    , 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
76

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

-- | 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.
--
114
vectTopBind :: CoreBind -> VM CoreBind
115
vectTopBind b@(NonRec var expr)
116 117 118 119 120 121 122 123 124 125 126
 = 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

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

vectTopBind b@(Rec bs)
132 133 134 135 136 137 138 139 140 141
 = 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')

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

150 151 152 153 154 155 156 157 158 159 160 161 162

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

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

      -- 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.
173
      defGlobalVar var var'
174

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

181 182 183 184 185 186 187

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

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

196 197 198 199 200 201 202 203 204

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

205 206 207
tryConvert var vect_var rhs
  = fromVect (idType var) (Var vect_var) `orElseV` return rhs