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', areScalars', exprs')
 139                      <- mapAndUnzip3M (uncurry $ vectTopRhs vars) bs
 140                if  (and areScalars') || (length bs <= 1)
 141                   then do
 142                     return (vars', inlines', exprs')
 143                   else do
 144                     _ <- mapM deleteGlobalScalar vars
 145                     (inlines'', _, exprs'')  <- mapAndUnzip3M (uncurry $ vectTopRhs []) bs
 146                     return (vars', inlines'', exprs'')
 147
 148       hs     <- takeHoisted
 149       cexprs <- sequence $ zipWith3 tryConvert vars vars' exprs
 150       return . Rec $ zip vars cexprs ++ zip vars' exprs' ++ hs
 151   `orElseV`
 152     return b
 153   where
 154     (vars, exprs) = unzip bs
 155     mapAndUnzip3M f xs = do
 156        ys <- mapM f xs
 157        return $ unzip3 ys
 158
 159 -- | Make the vectorised version of this top level binder, and add the mapping
 160 --   between it and the original to the state. For some binder @foo@ the vectorised
 161 --   version is @$v_foo@
 162 --
 163 --   NOTE: vectTopBinder *MUST* be lazy in inline and expr because of how it is
 164 --   used inside of fixV in vectTopBind
 165 vectTopBinder
 166         :: Var          -- ^ Name of the binding.
 167         -> Inline       -- ^ Whether it should be inlined, used to annotate it.
 168         -> CoreExpr     -- ^ RHS of the binding, used to set the `Unfolding` of the returned `Var`.
 169         -> VM Var       -- ^ Name of the vectorised binding.
 170
 171 vectTopBinder var inline expr
 172  = do
 173       -- Vectorise the type attached to the var.
 174       vty  <- vectType (idType var)
 175
 176       -- Make the vectorised version of binding's name, and set the unfolding used for inlining.
 177       var' <- liftM (`setIdUnfoldingLazily` unfolding)
 178            $  cloneId mkVectOcc var vty
 179
 180       -- Add the mapping between the plain and vectorised name to the state.
 181       defGlobalVar var var'
 182
 183       return var'
 184   where
 185     unfolding = case inline of
 186                   Inline arity -> mkInlineUnfolding (Just arity) expr
 187                   DontInline   -> noUnfolding
 188
 189
 190 -- | Vectorise the RHS of a top-level binding, in an empty local environment.
 191 vectTopRhs
 192         :: [Var]    -- ^ Names of all functions in the rec block
 193         -> Var          -- ^ Name of the binding.
 194         -> CoreExpr     -- ^ Body of the binding.
 195         -> VM (Inline, Bool, CoreExpr)
 196
 197 vectTopRhs recFs var expr
 198  = dtrace (vcat [text "vectTopRhs", ppr expr])
 199  $ closedV
 200  $ do (inline, isScalar, vexpr) <-
 201            inBind var $ vectPolyExpr  (isLoopBreaker $ idOccInfo var) recFs (freeVars expr)
 202       if isScalar
 203          then addGlobalScalar var
 204          else deleteGlobalScalar var
 205       return (inline, isScalar, vectorised vexpr)
 206
 207
 208 -- | Project out the vectorised version of a binding from some closure,
 209 --      or return the original body if that doesn't work.
 210 tryConvert
 211         :: Var          -- ^ Name of the original binding (eg @foo@)
 212         -> Var          -- ^ Name of vectorised version of binding (eg @$vfoo@)
 213         -> CoreExpr     -- ^ The original body of the binding.
 214         -> VM CoreExpr
 215
 216 tryConvert var vect_var rhs
 217   = fromVect (idType var) (Var vect_var) `orElseV` return rhs
 218