compiler/vectorise/Vectorise.hs

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