compiler/vectorise/Vectorise/Convert.hs

   1
   2 module Vectorise.Convert
   3         (fromVect)
   4 where
   5 import Vectorise.Monad
   6 import Vectorise.Builtins
   7 import Vectorise.Type.Type
   8
   9 import CoreSyn
  10 import TyCon
  11 import Type
  12 import TypeRep
  13 import FastString
  14
  15
  16 -- | Build an expression that calls the vectorised version of some
  17 --   function from a `Closure`.
  18 --
  19 --   For example
  20 --   @
  21 --      \(x :: Double) ->
  22 --      \(y :: Double) ->
  23 --      ($v_foo $: x) $: y
  24 --   @
  25 --
  26 --   We use the type of the original binding to work out how many
  27 --   outer lambdas to add.
  28 --
  29 fromVect
  30         :: Type         -- ^ The type of the original binding.
  31         -> CoreExpr     -- ^ Expression giving the closure to use, eg @$v_foo@.
  32         -> VM CoreExpr
  33
  34 -- Convert the type to the core view if it isn't already.
  35 fromVect ty expr
  36         | Just ty' <- coreView ty
  37         = fromVect ty' expr
  38
  39 -- For each function constructor in the original type we add an outer
  40 -- lambda to bind the parameter variable, and an inner application of it.
  41 fromVect (FunTy arg_ty res_ty) expr
  42   = do
  43       arg     <- newLocalVar (fsLit "x") arg_ty
  44       varg    <- toVect arg_ty (Var arg)
  45       varg_ty <- vectType arg_ty
  46       vres_ty <- vectType res_ty
  47       apply   <- builtin applyVar
  48       body    <- fromVect res_ty
  49                $ Var apply `mkTyApps` [varg_ty, vres_ty] `mkApps` [expr, varg]
  50       return $ Lam arg body
  51
  52 -- If the type isn't a function then it's time to call on the closure.
  53 fromVect ty expr
  54   = identityConv ty >> return expr
  55
  56
  57 -- TODO: What is this really doing?
  58 toVect :: Type -> CoreExpr -> VM CoreExpr
  59 toVect ty expr = identityConv ty >> return expr
  60
  61
  62 -- | Check that we have the vectorised versions of all the
  63 --   type constructors in this type.
  64 identityConv :: Type -> VM ()
  65 identityConv ty
  66   | Just ty' <- coreView ty
  67   = identityConv ty'
  68
  69 identityConv (TyConApp tycon tys)
  70  = do mapM_ identityConv tys
  71       identityConvTyCon tycon
  72
  73 identityConv _ = noV
  74
  75
  76 -- | Check that we have the vectorised version of this type constructor.
  77 identityConvTyCon :: TyCon -> VM ()
  78 identityConvTyCon tc
  79   | isBoxedTupleTyCon tc = return ()
  80   | isUnLiftedTyCon tc   = return ()
  81   | otherwise
  82   = do tc' <- maybeV (lookupTyCon tc)
  83        if tc == tc' then return () else noV