vectorise/Vectorise/Convert.hs

    module Vectorise.Convert
      ( fromVect
      )
    where

    import Vectorise.Monad
    import Vectorise.Builtins
    import Vectorise.Type.Type

    import CoreSyn
    import TyCon
    import Type
    import TyCoRep
    import NameSet
    import FastString
    import Outputable

    import Control.Applicative
    import Prelude -- avoid redundant import warning due to AMP

    -- |Convert a vectorised expression such that it computes the non-vectorised equivalent of its
    -- value.
    --
    -- For functions, we eta expand the function and convert the arguments and result:

    -- For example
    -- @
    --    \(x :: Double) ->
    --    \(y :: Double) ->
    --    ($v_foo $: x) $: y
    -- @
    --
    -- We use the type of the original binding to work out how many outer lambdas to add.
    --
    fromVect :: Type        -- ^ The type of the original binding.
             -> CoreExpr    -- ^ Expression giving the closure to use, eg @$v_foo@.
             -> VM CoreExpr

    -- Convert the type to the core view if it isn't already.
    --
    fromVect ty expr
      | Just ty' <- coreView ty
      = fromVect ty' expr

    -- For each function constructor in the original type we add an outer
    -- lambda to bind the parameter variable, and an inner application of it.
    fromVect (FunTy arg_ty res_ty) expr
      = do
          arg     <- newLocalVar (fsLit "x") arg_ty
          varg    <- toVect arg_ty (Var arg)
          varg_ty <- vectType arg_ty
          vres_ty <- vectType res_ty
          apply   <- builtin applyVar
          body    <- fromVect res_ty
                   $ Var apply `mkTyApps` [varg_ty, vres_ty] `mkApps` [expr, varg]
          return $ Lam arg body

    -- If the type isn't a function, then we can't current convert it unless the type is scalar (i.e.,
    -- is identical to the non-vectorised version).
    --
    fromVect ty expr
      = identityConv ty >> return expr

    -- Convert an expression such that it evaluates to the vectorised equivalent of the value of the
    -- original expression.
    --
    -- WARNING: Currently only works for the scalar types, where the vectorised value coincides with the
    --          original one.
    --
    toVect :: Type -> CoreExpr -> VM CoreExpr
    toVect ty expr = identityConv ty >> return expr

    -- |Check that the type is neutral under type vectorisation — i.e., all involved type constructor
    -- are not altered by vectorisation as they contain no parallel arrays.
    --
    identityConv :: Type -> VM ()
    identityConv ty
      | Just ty' <- coreView ty
      = identityConv ty'
    identityConv (TyConApp tycon tys)
      = do { mapM_ identityConv tys
           ; identityConvTyCon tycon
           }
    identityConv (LitTy {})      = noV $ text "identityConv: not sure about literal types under vectorisation"
    identityConv (TyVarTy {})    = noV $ text "identityConv: type variable changes under vectorisation"
    identityConv (AppTy {})      = noV $ text "identityConv: type appl. changes under vectorisation"
    identityConv (FunTy {})      = noV $ text "identityConv: function type changes under vectorisation"
    identityConv (ForAllTy {})   = noV $ text "identityConv: quantified type changes under vectorisation"
    identityConv (CastTy {})     = noV $ text "identityConv: not sure about casted types under vectorisation"
    identityConv (CoercionTy {}) = noV $ text "identityConv: not sure about coercions under vectorisation"

    -- |Check that this type constructor is not changed by vectorisation — i.e., it does not embed any
    -- parallel arrays.
    --
    identityConvTyCon :: TyCon -> VM ()
    identityConvTyCon tc
      = do
        { isParallel <- (tyConName tc `elemNameSet`) <$> globalParallelTyCons
        ; parray     <- builtin parrayTyCon
        ; if isParallel && not (tc == parray)
          then noV idErr
          else return ()
        }
      where
        idErr = text "identityConvTyCon: type constructor contains parallel arrays" <+> ppr tc