Convert.hs 3.36 KB
Newer Older
1
module Vectorise.Convert
2 3
  ( fromVect
  )
4
where
5

6 7 8 9 10 11 12 13
import Vectorise.Monad
import Vectorise.Builtins
import Vectorise.Type.Type

import CoreSyn
import TyCon
import Type
import TypeRep
14
import NameSet
15
import FastString
16
import Outputable
17

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

21 22
-- |Convert a vectorised expression such that it computes the non-vectorised equivalent of its
-- value.
23
--
24 25 26 27 28 29 30 31
-- For functions, we eta expand the function and convert the arguments and result:

-- For example
-- @   
--    \(x :: Double) -> 
--    \(y :: Double) -> 
--    ($v_foo $: x) $: y
-- @
32
--
33
-- We use the type of the original binding to work out how many outer lambdas to add.
34
--
35 36 37 38
fromVect :: Type        -- ^ The type of the original binding.
         -> CoreExpr    -- ^ Expression giving the closure to use, eg @$v_foo@.
         -> VM CoreExpr
  
39
-- Convert the type to the core view if it isn't already.
40
--
41
fromVect ty expr 
42 43
  | Just ty' <- coreView ty 
  = fromVect ty' expr
44 45 46 47 48 49 50 51 52 53 54 55 56 57

-- 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

58 59 60
-- 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).
--
61 62 63
fromVect ty expr
  = identityConv ty >> return expr

64 65 66 67 68 69
-- 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.
--
70 71 72
toVect :: Type -> CoreExpr -> VM CoreExpr
toVect ty expr = identityConv ty >> return expr

73 74 75
-- |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.
--
76 77 78 79 80
identityConv :: Type -> VM ()
identityConv ty 
  | Just ty' <- coreView ty 
  = identityConv ty'
identityConv (TyConApp tycon tys)
81 82 83
  = do { mapM_ identityConv tys
       ; identityConvTyCon tycon
       }
84 85 86 87 88
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"
89

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