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