compiler/vectorise/Vectorise.hs

   1 {-# OPTIONS -fno-warn-missing-signatures #-}
   2
   3 module Vectorise( vectorise )
   4 where
   5
   6 import Vectorise.Type.Env
   7 import Vectorise.Type.Type
   8 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
  20 import CoreMonad            ( CoreM, getHscEnv )
  21 import Var
  22 import Id
  23 import OccName
  24 import BasicTypes           ( isLoopBreaker )
  25 import Outputable
  26 import Util                 ( zipLazy )
  27 import MonadUtils
  28
  29 import Control.Monad
  30
  31 debug           = False
  32 dtrace s x      = if debug then pprTrace "Vectorise" s x else x
  33
  34 -- | Vectorise a single module.
  35 --   Takes the package containing the DPH backend we're using. Eg either dph-par or dph-seq.
  36 vectorise :: PackageId -> ModGuts -> CoreM ModGuts
  37 vectorise backend guts
  38  = do hsc_env <- getHscEnv
  39       liftIO $ vectoriseIO backend hsc_env guts
  40
  41
  42 -- | Vectorise a single monad, given its HscEnv (code gen environment).
  43 vectoriseIO :: PackageId -> HscEnv -> ModGuts -> IO ModGuts
  44 vectoriseIO backend hsc_env guts
  45  = do -- Get information about currently loaded external packages.
  46       eps <- hscEPS hsc_env
  47
  48       -- Combine vectorisation info from the current module, and external ones.
  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.
  57 vectModule :: ModGuts -> VM ModGuts
  58 vectModule guts
  59  = do -- Vectorise the type environment.
  60       -- This may add new TyCons and DataCons.
  61       -- TODO: What new binds do we get back here?
  62       (types', fam_insts, tc_binds) <- vectTypeEnv (mg_types guts)
  63
  64       (_, fam_inst_env) <- readGEnv global_fam_inst_env
  65
  66       -- dicts   <- mapM buildPADict pa_insts
  67       -- workers <- mapM vectDataConWorkers pa_insts
  68
  69       -- Vectorise all the top level bindings.
  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                     , mg_fam_insts    = mg_fam_insts guts ++ fam_insts
  76                     }
  77
  78
  79 -- | Try to vectorise a top-level binding.
  80 --   If it doesn't vectorise then return it unharmed.
  81 --
  82 --   For example, for the binding
  83 --
  84 --   @
  85 --      foo :: Int -> Int
  86 --      foo = \x -> x + x
  87 --   @
  88 --
  89 --   we get
  90 --   @
  91 --      foo  :: Int -> Int
  92 --      foo  = \x -> vfoo $: x
  93 --
  94 --      v_foo :: Closure void vfoo lfoo
  95 --      v_foo = closure vfoo lfoo void
  96 --
  97 --      vfoo :: Void -> Int -> Int
  98 --      vfoo = ...
  99 --
 100 --      lfoo :: PData Void -> PData Int -> PData Int
 101 --      lfoo = ...
 102 --   @
 103 --
 104 --   @vfoo@ is the "vectorised", or scalar, version that does the same as the original
 105 --   function foo, but takes an explicit environment.
 106 --
 107 --   @lfoo@ is the "lifted" version that works on arrays.
 108 --
 109 --   @v_foo@ combines both of these into a `Closure` that also contains the
 110 --   environment.
 111 --
 112 --   The original binding @foo@ is rewritten to call the vectorised version
 113 --   present in the closure.
 114 --
 115 vectTopBind :: CoreBind -> VM CoreBind
 116 vectTopBind b@(NonRec var expr)
 117  = do
 118       (inline, expr')   <- vectTopRhs var expr
 119       var'              <- vectTopBinder var inline expr'
 120
 121       -- Vectorising the body may create other top-level bindings.
 122       hs        <- takeHoisted
 123
 124       -- To get the same functionality as the original body we project
 125       -- out its vectorised version from the closure.
 126       cexpr     <- tryConvert var var' expr
 127
 128       return . Rec $ (var, cexpr) : (var', expr') : hs
 129   `orElseV`
 130     return b
 131
 132 vectTopBind b@(Rec bs)
 133  = do
 134       (vars', _, exprs')
 135         <- fixV $ \ ~(_, inlines, rhss) ->
 136             do vars' <- sequence [vectTopBinder var inline rhs
 137                                       | (var, ~(inline, rhs)) <- zipLazy vars (zip inlines rhss)]
 138                (inlines', exprs')
 139                      <- mapAndUnzipM (uncurry vectTopRhs) bs
 140
 141                return (vars', inlines', exprs')
 142
 143       hs     <- takeHoisted
 144       cexprs <- sequence $ zipWith3 tryConvert vars vars' exprs
 145       return . Rec $ zip vars cexprs ++ zip vars' exprs' ++ hs
 146   `orElseV`
 147     return b
 148   where
 149     (vars, exprs) = unzip bs
 150
 151
 152 -- | Make the vectorised version of this top level binder, and add the mapping
 153 --   between it and the original to the state. For some binder @foo@ the vectorised
 154 --   version is @$v_foo@
 155 --
 156 --   NOTE: vectTopBinder *MUST* be lazy in inline and expr because of how it is
 157 --   used inside of fixV in vectTopBind
 158 vectTopBinder
 159         :: Var          -- ^ Name of the binding.
 160         -> Inline       -- ^ Whether it should be inlined, used to annotate it.
 161         -> CoreExpr     -- ^ RHS of the binding, used to set the `Unfolding` of the returned `Var`.
 162         -> VM Var       -- ^ Name of the vectorised binding.
 163
 164 vectTopBinder var inline expr
 165  = do
 166       -- Vectorise the type attached to the var.
 167       vty  <- vectType (idType var)
 168
 169       -- Make the vectorised version of binding's name, and set the unfolding used for inlining.
 170       var' <- liftM (`setIdUnfoldingLazily` unfolding)
 171            $  cloneId mkVectOcc var vty
 172
 173       -- Add the mapping between the plain and vectorised name to the state.
 174       defGlobalVar var var'
 175
 176       return var'
 177   where
 178     unfolding = case inline of
 179                   Inline arity -> mkInlineUnfolding (Just arity) expr
 180                   DontInline   -> noUnfolding
 181
 182
 183 -- | Vectorise the RHS of a top-level binding, in an empty local environment.
 184 vectTopRhs
 185         :: Var          -- ^ Name of the binding.
 186         -> CoreExpr     -- ^ Body of the binding.
 187         -> VM (Inline, CoreExpr)
 188
 189 vectTopRhs var expr
 190  = dtrace (vcat [text "vectTopRhs", ppr expr])
 191  $ closedV
 192  $ do (inline, vexpr) <- inBind var
 193                       $ vectPolyExpr (isLoopBreaker $ idOccInfo var)
 194                                       (freeVars expr)
 195       return (inline, vectorised vexpr)
 196
 197
 198 -- | Project out the vectorised version of a binding from some closure,
 199 --      or return the original body if that doesn't work.
 200 tryConvert
 201         :: Var          -- ^ Name of the original binding (eg @foo@)
 202         -> Var          -- ^ Name of vectorised version of binding (eg @$vfoo@)
 203         -> CoreExpr     -- ^ The original body of the binding.
 204         -> VM CoreExpr
 205
 206 tryConvert var vect_var rhs
 207   = fromVect (idType var) (Var vect_var) `orElseV` return rhs
 208