compiler/vectorise/Vectorise/Utils.hs

   1
   2 module Vectorise.Utils (
   3   module Vectorise.Utils.Base,
   4   module Vectorise.Utils.Closure,
   5   module Vectorise.Utils.Hoisting,
   6   module Vectorise.Utils.PADict,
   7   module Vectorise.Utils.PRDict,
   8   module Vectorise.Utils.Poly,
   9
  10   -- * Annotated Exprs
  11   collectAnnTypeArgs,
  12   collectAnnTypeBinders,
  13   collectAnnValBinders,
  14   isAnnTypeArg,
  15
  16   -- * PD Functions
  17   replicatePD, emptyPD, packByTagPD,
  18   combinePD, liftPD,
  19
  20   -- * Scalars
  21   zipScalars, scalarClosure,
  22
  23   -- * Naming
  24   newLocalVar
  25 )
  26 where
  27 import Vectorise.Utils.Base
  28 import Vectorise.Utils.Closure
  29 import Vectorise.Utils.Hoisting
  30 import Vectorise.Utils.PADict
  31 import Vectorise.Utils.PRDict
  32 import Vectorise.Utils.Poly
  33 import Vectorise.Monad
  34 import Vectorise.Builtins
  35 import CoreSyn
  36 import CoreUtils
  37 import Type
  38 import Var
  39 import Control.Monad
  40
  41
  42 -- Annotated Exprs ------------------------------------------------------------
  43 collectAnnTypeArgs :: AnnExpr b ann -> (AnnExpr b ann, [Type])
  44 collectAnnTypeArgs expr = go expr []
  45   where
  46     go (_, AnnApp f (_, AnnType ty)) tys = go f (ty : tys)
  47     go e                             tys = (e, tys)
  48
  49 collectAnnTypeBinders :: AnnExpr Var ann -> ([Var], AnnExpr Var ann)
  50 collectAnnTypeBinders expr = go [] expr
  51   where
  52     go bs (_, AnnLam b e) | isTyCoVar b = go (b:bs) e
  53     go bs e                           = (reverse bs, e)
  54
  55 collectAnnValBinders :: AnnExpr Var ann -> ([Var], AnnExpr Var ann)
  56 collectAnnValBinders expr = go [] expr
  57   where
  58     go bs (_, AnnLam b e) | isId b = go (b:bs) e
  59     go bs e                        = (reverse bs, e)
  60
  61 isAnnTypeArg :: AnnExpr b ann -> Bool
  62 isAnnTypeArg (_, AnnType _) = True
  63 isAnnTypeArg _              = False
  64
  65
  66 -- PD "Parallel Data" Functions -----------------------------------------------
  67 --
  68 --   Given some data that has a PA dictionary, we can convert it to its
  69 --   representation type, perform some operation on the data, then convert it back.
  70 --
  71 --   In the DPH backend, the types of these functions are defined
  72 --   in dph-common/D.A.P.Lifted/PArray.hs
  73 --
  74
  75 -- | An empty array of the given type.
  76 emptyPD :: Type -> VM CoreExpr
  77 emptyPD = paMethod emptyPDVar "emptyPD"
  78
  79
  80 -- | Produce an array containing copies of a given element.
  81 replicatePD
  82         :: CoreExpr     -- ^ Number of copies in the resulting array.
  83         -> CoreExpr     -- ^ Value to replicate.
  84         -> VM CoreExpr
  85
  86 replicatePD len x
  87         = liftM (`mkApps` [len,x])
  88         $ paMethod replicatePDVar "replicatePD" (exprType x)
  89
  90
  91 -- | Select some elements from an array that correspond to a particular tag value
  92 ---  and pack them into a new array.
  93 --   eg  packByTagPD Int# [:23, 42, 95, 50, 27, 49:]  3 [:1, 2, 1, 2, 3, 2:] 2
  94 --          ==> [:42, 50, 49:]
  95 --
  96 packByTagPD
  97         :: Type         -- ^ Element type.
  98         -> CoreExpr     -- ^ Source array.
  99         -> CoreExpr     -- ^ Length of resulting array.
 100         -> CoreExpr     -- ^ Tag values of elements in source array.
 101         -> CoreExpr     -- ^ The tag value for the elements to select.
 102         -> VM CoreExpr
 103
 104 packByTagPD ty xs len tags t
 105   = liftM (`mkApps` [xs, len, tags, t])
 106           (paMethod packByTagPDVar "packByTagPD" ty)
 107
 108
 109 -- | Combine some arrays based on a selector.
 110 --     The selector says which source array to choose for each element of the
 111 --     resulting array.
 112 combinePD
 113         :: Type         -- ^ Element type
 114         -> CoreExpr     -- ^ Length of resulting array
 115         -> CoreExpr     -- ^ Selector.
 116         -> [CoreExpr]   -- ^ Arrays to combine.
 117         -> VM CoreExpr
 118
 119 combinePD ty len sel xs
 120   = liftM (`mkApps` (len : sel : xs))
 121           (paMethod (combinePDVar n) ("combine" ++ show n ++ "PD") ty)
 122   where
 123     n = length xs
 124
 125
 126 -- | Like `replicatePD` but use the lifting context in the vectoriser state.
 127 liftPD :: CoreExpr -> VM CoreExpr
 128 liftPD x
 129   = do
 130       lc <- builtin liftingContext
 131       replicatePD (Var lc) x
 132
 133
 134 -- Scalars --------------------------------------------------------------------
 135 zipScalars :: [Type] -> Type -> VM CoreExpr
 136 zipScalars arg_tys res_ty
 137   = do
 138       scalar <- builtin scalarClass
 139       (dfuns, _) <- mapAndUnzipM (\ty -> lookupInst scalar [ty]) ty_args
 140       zipf <- builtin (scalarZip $ length arg_tys)
 141       return $ Var zipf `mkTyApps` ty_args `mkApps` map Var dfuns
 142     where
 143       ty_args = arg_tys ++ [res_ty]
 144
 145
 146 scalarClosure :: [Type] -> Type -> CoreExpr -> CoreExpr -> VM CoreExpr
 147 scalarClosure arg_tys res_ty scalar_fun array_fun
 148   = do
 149       ctr      <- builtin (closureCtrFun $ length arg_tys)
 150       Just pas <- liftM sequence $ mapM paDictOfType (init arg_tys)
 151       return $ Var ctr `mkTyApps` (arg_tys ++ [res_ty])
 152                        `mkApps`   (pas ++ [scalar_fun, array_fun])
 153
 154
 155
 156 {-
 157 boxExpr :: Type -> VExpr -> VM VExpr
 158 boxExpr ty (vexpr, lexpr)
 159   | Just (tycon, []) <- splitTyConApp_maybe ty
 160   , isUnLiftedTyCon tycon
 161   = do
 162       r <- lookupBoxedTyCon tycon
 163       case r of
 164         Just tycon' -> let [dc] = tyConDataCons tycon'
 165                        in
 166                        return (mkConApp dc [vexpr], lexpr)
 167         Nothing     -> return (vexpr, lexpr)
 168 -}