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