compiler/vectorise/Vectorise.hs

   1
   2 module Vectorise ( vectorise )
   3 where
   4
   5 import Vectorise.Type.Env
   6 import Vectorise.Type.Type
   7 import Vectorise.Convert
   8 import Vectorise.Utils.Hoisting
   9 import Vectorise.Exp
  10 import Vectorise.Vect
  11 import Vectorise.Env
  12 import Vectorise.Monad
  13
  14 import HscTypes hiding      ( MonadThings(..) )
  15 import CoreUnfold           ( mkInlineUnfolding )
  16 import CoreFVs
  17 import PprCore
  18 import CoreSyn
  19 import CoreMonad            ( CoreM, getHscEnv )
  20 import Type
  21 import Id
  22 import OccName
  23 import DynFlags
  24 import BasicTypes           ( isLoopBreaker )
  25 import Outputable
  26 import Util                 ( zipLazy )
  27 import MonadUtils
  28
  29 import Control.Monad
  30
  31
  32 -- | Vectorise a single module.
  33 --
  34 vectorise :: ModGuts -> CoreM ModGuts
  35 vectorise guts
  36  = do { hsc_env <- getHscEnv
  37       ; liftIO $ vectoriseIO hsc_env guts
  38       }
  39
  40 -- | Vectorise a single monad, given the dynamic compiler flags and HscEnv.
  41 --
  42 vectoriseIO :: HscEnv -> ModGuts -> IO ModGuts
  43 vectoriseIO 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 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 --
  57 vectModule :: ModGuts -> VM ModGuts
  58 vectModule guts@(ModGuts { mg_types     = types
  59                          , mg_binds     = binds
  60                          , mg_fam_insts = fam_insts
  61                          })
  62  = do { dumpOptVt Opt_D_dump_vt_trace "Before vectorisation" $
  63           pprCoreBindings binds
  64
  65           -- Vectorise the type environment.
  66           -- This may add new TyCons and DataCons.
  67       ; (types', new_fam_insts, tc_binds) <- vectTypeEnv types
  68
  69       ; (_, fam_inst_env) <- readGEnv global_fam_inst_env
  70
  71       -- dicts   <- mapM buildPADict pa_insts
  72       -- workers <- mapM vectDataConWorkers pa_insts
  73
  74           -- Vectorise all the top level bindings.
  75       ; binds'  <- mapM vectTopBind binds
  76
  77       ; return $ guts { mg_types        = types'
  78                       , mg_binds        = Rec tc_binds : binds'
  79                       , mg_fam_inst_env = fam_inst_env
  80                       , mg_fam_insts    = fam_insts ++ new_fam_insts
  81                       }
  82       }
  83
  84 -- |Try to vectorise a top-level binding.  If it doesn't vectorise then return it unharmed.
  85 --
  86 -- For example, for the binding
  87 --
  88 -- @
  89 --    foo :: Int -> Int
  90 --    foo = \x -> x + x
  91 -- @
  92 --
  93 -- we get
  94 -- @
  95 --    foo  :: Int -> Int
  96 --    foo  = \x -> vfoo $: x
  97 --
  98 --    v_foo :: Closure void vfoo lfoo
  99 --    v_foo = closure vfoo lfoo void
 100 --
 101 --    vfoo :: Void -> Int -> Int
 102 --    vfoo = ...
 103 --
 104 --    lfoo :: PData Void -> PData Int -> PData Int
 105 --    lfoo = ...
 106 -- @
 107 --
 108 -- @vfoo@ is the "vectorised", or scalar, version that does the same as the original
 109 -- function foo, but takes an explicit environment.
 110 --
 111 -- @lfoo@ is the "lifted" version that works on arrays.
 112 --
 113 -- @v_foo@ combines both of these into a `Closure` that also contains the
 114 -- environment.
 115 --
 116 -- The original binding @foo@ is rewritten to call the vectorised version
 117 -- present in the closure.
 118 --
 119 -- Vectorisation may be surpressed by annotating a binding with a 'NOVECTORISE' pragma.  If this
 120 -- pragma is used in a group of mutually recursive bindings, either all or no binding must have
 121 -- the pragma.  If only some bindings are annotated, a fatal error is being raised.
 122 -- FIXME: Once we support partial vectorisation, we may be able to vectorise parts of a group, or
 123 --   we may emit a warning and refrain from vectorising the entire group.
 124 --
 125 vectTopBind :: CoreBind -> VM CoreBind
 126 vectTopBind b@(NonRec var expr)
 127   = unlessNoVectDecl $
 128       do {   -- Vectorise the right-hand side, create an appropriate top-level binding and add it
 129              -- to the vectorisation map.
 130          ; (inline, isScalar, expr') <- vectTopRhs [] var expr
 131          ; var' <- vectTopBinder var inline expr'
 132          ; when isScalar $
 133              addGlobalScalar var
 134
 135              -- We replace the original top-level binding by a value projected from the vectorised
 136              -- closure and add any newly created hoisted top-level bindings.
 137          ; cexpr <- tryConvert var var' expr
 138          ; hs <- takeHoisted
 139          ; return . Rec $ (var, cexpr) : (var', expr') : hs
 140          }
 141      `orElseV`
 142        return b
 143   where
 144     unlessNoVectDecl vectorise
 145       = do { hasNoVectDecl <- noVectDecl var
 146            ; when hasNoVectDecl $
 147                traceVt "NOVECTORISE" $ ppr var
 148            ; if hasNoVectDecl then return b else vectorise
 149            }
 150 vectTopBind b@(Rec bs)
 151   = unlessSomeNoVectDecl $
 152       do { (vars', _, exprs', hs) <- fixV $
 153              \ ~(_, inlines, rhss, _) ->
 154                do {   -- Vectorise the right-hand sides, create an appropriate top-level bindings
 155                       -- and add them to the vectorisation map.
 156                   ; vars' <- sequence [vectTopBinder var inline rhs
 157                                       | (var, ~(inline, rhs)) <- zipLazy vars (zip inlines rhss)]
 158                   ; (inlines, areScalars, exprs') <- mapAndUnzip3M (uncurry $ vectTopRhs vars) bs
 159                   ; hs <- takeHoisted
 160                   ; if and areScalars
 161                     then      -- (1) Entire recursive group is scalar
 162                               --      => add all variables to the global set of scalars
 163                          do { mapM_ addGlobalScalar vars
 164                             ; return (vars', inlines, exprs', hs)
 165                             }
 166                     else      -- (2) At least one binding is not scalar
 167                               --     => vectorise again with empty set of local scalars
 168                          do { (inlines, _, exprs') <- mapAndUnzip3M (uncurry $ vectTopRhs []) bs
 169                             ; hs <- takeHoisted
 170                             ; return (vars', inlines, exprs', hs)
 171                             }
 172                   }
 173
 174              -- Replace the original top-level bindings by a values projected from the vectorised
 175              -- closures and add any newly created hoisted top-level bindings to the group.
 176          ; cexprs <- sequence $ zipWith3 tryConvert vars vars' exprs
 177          ; return . Rec $ zip vars cexprs ++ zip vars' exprs' ++ hs
 178          }
 179      `orElseV`
 180        return b
 181   where
 182     (vars, exprs) = unzip bs
 183
 184     unlessSomeNoVectDecl vectorise
 185       = do { hasNoVectDecls <- mapM noVectDecl vars
 186            ; when (and hasNoVectDecls) $
 187                traceVt "NOVECTORISE" $ ppr vars
 188            ; if and hasNoVectDecls
 189              then return b                              -- all bindings have 'NOVECTORISE'
 190              else if or hasNoVectDecls
 191              then cantVectorise noVectoriseErr (ppr b)  -- some (but not all) have 'NOVECTORISE'
 192              else vectorise                             -- no binding has a 'NOVECTORISE' decl
 193            }
 194     noVectoriseErr = "NOVECTORISE must be used on all or no bindings of a recursive group"
 195
 196 -- | Make the vectorised version of this top level binder, and add the mapping
 197 --   between it and the original to the state. For some binder @foo@ the vectorised
 198 --   version is @$v_foo@
 199 --
 200 --   NOTE: 'vectTopBinder' *MUST* be lazy in inline and expr because of how it is
 201 --   used inside of 'fixV' in 'vectTopBind'.
 202 --
 203 vectTopBinder :: Var      -- ^ Name of the binding.
 204               -> Inline   -- ^ Whether it should be inlined, used to annotate it.
 205               -> CoreExpr -- ^ RHS of binding, used to set the 'Unfolding' of the returned 'Var'.
 206               -> VM Var   -- ^ Name of the vectorised binding.
 207 vectTopBinder var inline expr
 208  = do {   -- Vectorise the type attached to the var.
 209       ; vty  <- vectType (idType var)
 210
 211           -- If there is a vectorisation declartion for this binding, make sure that its type
 212           --  matches
 213       ; vectDecl <- lookupVectDecl var
 214       ; case vectDecl of
 215           Nothing                 -> return ()
 216           Just (vdty, _)
 217             | eqType vty vdty -> return ()
 218             | otherwise           ->
 219               cantVectorise ("Type mismatch in vectorisation pragma for " ++ show var) $
 220                 (text "Expected type" <+> ppr vty)
 221                 $$
 222                 (text "Inferred type" <+> ppr vdty)
 223
 224           -- Make the vectorised version of binding's name, and set the unfolding used for inlining
 225       ; var' <- liftM (`setIdUnfoldingLazily` unfolding)
 226                 $  cloneId mkVectOcc var vty
 227
 228           -- Add the mapping between the plain and vectorised name to the state.
 229       ; defGlobalVar var var'
 230
 231       ; return var'
 232     }
 233   where
 234     unfolding = case inline of
 235                   Inline arity -> mkInlineUnfolding (Just arity) expr
 236                   DontInline   -> noUnfolding
 237
 238 -- | Vectorise the RHS of a top-level binding, in an empty local environment.
 239 --
 240 -- We need to distinguish three cases:
 241 --
 242 -- (1) We have a (non-scalar) vectorisation declaration for the variable (which explicitly provides
 243 --     vectorised code implemented by the user)
 244 --     => no automatic vectorisation & instead use the user-supplied code
 245 --
 246 -- (2) We have a scalar vectorisation declaration for the variable
 247 --     => generate vectorised code that uses a scalar 'map'/'zipWith' to lift the computation
 248 --
 249 -- (3) There is no vectorisation declaration for the variable
 250 --     => perform automatic vectorisation of the RHS
 251 --
 252 vectTopRhs :: [Var]           -- ^ Names of all functions in the rec block
 253            -> Var             -- ^ Name of the binding.
 254            -> CoreExpr        -- ^ Body of the binding.
 255            -> VM ( Inline     -- (1) inline specification for the binding
 256                  , Bool       -- (2) whether the right-hand side is a scalar computation
 257                  , CoreExpr)  -- (3) the vectorised right-hand side
 258 vectTopRhs recFs var expr
 259   = closedV
 260   $ do { traceVt ("vectTopRhs of " ++ show var) $ ppr expr
 261
 262        ; globalScalar <- isGlobalScalar var
 263        ; vectDecl     <- lookupVectDecl var
 264        ; rhs globalScalar vectDecl
 265        }
 266   where
 267     rhs _globalScalar (Just (_, expr'))               -- Case (1)
 268       = return (inlineMe, False, expr')
 269     rhs True          Nothing                         -- Case (2)
 270       = do { expr' <- vectScalarFun True recFs expr
 271            ; return (inlineMe, True, vectorised expr')
 272            }
 273     rhs False         Nothing                         -- Case (3)
 274       = do { let fvs = freeVars expr
 275            ; (inline, isScalar, vexpr) <- inBind var $
 276                                             vectPolyExpr (isLoopBreaker $ idOccInfo var) recFs fvs
 277            ; return (inline, isScalar, vectorised vexpr)
 278            }
 279
 280 -- | Project out the vectorised version of a binding from some closure,
 281 --   or return the original body if that doesn't work or the binding is scalar.
 282 --
 283 tryConvert :: Var       -- ^ Name of the original binding (eg @foo@)
 284            -> Var       -- ^ Name of vectorised version of binding (eg @$vfoo@)
 285            -> CoreExpr  -- ^ The original body of the binding.
 286            -> VM CoreExpr
 287 tryConvert var vect_var rhs
 288   = do { globalScalar <- isGlobalScalar var
 289        ; if globalScalar
 290          then
 291            return rhs
 292          else
 293            fromVect (idType var) (Var vect_var) `orElseV` return rhs
 294        }