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