compiler/vectorise/Vectorise.hs

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