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