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,
12 collectAnnTypeBinders,
17 replicatePD, emptyPD, packByTagPD,
21 zipScalars, scalarClosure,
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
42 -- Annotated Exprs ------------------------------------------------------------
43 collectAnnTypeArgs :: AnnExpr b ann -> (AnnExpr b ann, [Type])
44 collectAnnTypeArgs expr = go expr []
46 go (_, AnnApp f (_, AnnType ty)) tys = go f (ty : tys)
49 collectAnnTypeBinders :: AnnExpr Var ann -> ([Var], AnnExpr Var ann)
50 collectAnnTypeBinders expr = go [] expr
52 go bs (_, AnnLam b e) | isTyCoVar b = go (b:bs) e
53 go bs e = (reverse bs, e)
55 collectAnnValBinders :: AnnExpr Var ann -> ([Var], AnnExpr Var ann)
56 collectAnnValBinders expr = go [] expr
58 go bs (_, AnnLam b e) | isId b = go (b:bs) e
59 go bs e = (reverse bs, e)
61 isAnnTypeArg :: AnnExpr b ann -> Bool
62 isAnnTypeArg (_, AnnType _) = True
63 isAnnTypeArg _ = False
66 -- PD "Parallel Data" Functions -----------------------------------------------
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.
71 -- In the DPH backend, the types of these functions are defined
72 -- in dph-common/D.A.P.Lifted/PArray.hs
75 -- | An empty array of the given type.
76 emptyPD :: Type -> VM CoreExpr
77 emptyPD = paMethod emptyPDVar "emptyPD"
80 -- | Produce an array containing copies of a given element.
82 :: CoreExpr -- ^ Number of copies in the resulting array.
83 -> CoreExpr -- ^ Value to replicate.
87 = liftM (`mkApps` [len,x])
88 $ paMethod replicatePDVar "replicatePD" (exprType x)
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
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.
104 packByTagPD ty xs len tags t
105 = liftM (`mkApps` [xs, len, tags, t])
106 (paMethod packByTagPDVar "packByTagPD" ty)
109 -- | Combine some arrays based on a selector.
110 -- The selector says which source array to choose for each element of the
113 :: Type -- ^ Element type
114 -> CoreExpr -- ^ Length of resulting array
115 -> CoreExpr -- ^ Selector.
116 -> [CoreExpr] -- ^ Arrays to combine.
119 combinePD ty len sel xs
120 = liftM (`mkApps` (len : sel : xs))
121 (paMethod (combinePDVar n) ("combine" ++ show n ++ "PD") ty)
126 -- | Like `replicatePD` but use the lifting context in the vectoriser state.
127 liftPD :: CoreExpr -> VM CoreExpr
130 lc <- builtin liftingContext
131 replicatePD (Var lc) x
134 -- Scalars --------------------------------------------------------------------
135 zipScalars :: [Type] -> Type -> VM CoreExpr
136 zipScalars arg_tys res_ty
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
143 ty_args = arg_tys ++ [res_ty]
146 scalarClosure :: [Type] -> Type -> CoreExpr -> CoreExpr -> VM CoreExpr
147 scalarClosure arg_tys res_ty scalar_fun array_fun
149 ctr <- builtin (closureCtrFun $ length arg_tys)
150 pas <- mapM paDictOfType (init arg_tys)
151 return $ Var ctr `mkTyApps` (arg_tys ++ [res_ty])
152 `mkApps` (pas ++ [scalar_fun, array_fun])
157 boxExpr :: Type -> VExpr -> VM VExpr
158 boxExpr ty (vexpr, lexpr)
159 | Just (tycon, []) <- splitTyConApp_maybe ty
160 , isUnLiftedTyCon tycon
162 r <- lookupBoxedTyCon tycon
164 Just tycon' -> let [dc] = tyConDataCons tycon'
166 return (mkConApp dc [vexpr], lexpr)
167 Nothing -> return (vexpr, lexpr)