2 module VectType ( vectTyCon, vectAndLiftType, vectType, vectTypeEnv,
3 mkRepr, arrShapeTys, arrShapeVars, arrSelector,
12 import HscTypes ( TypeEnv, extendTypeEnvList, typeEnvTyCons )
21 import FamInstEnv ( FamInst, mkLocalFamInst )
24 import BasicTypes ( StrictnessMark(..), boolToRecFlag )
25 import Var ( Var, TyVar )
26 import Id ( mkWildId )
27 import Name ( Name, getOccName )
30 import TysPrim ( intPrimTy )
36 import Digraph ( SCC(..), stronglyConnCompFromEdgedVertices )
41 import Control.Monad ( liftM, liftM2, zipWithM, zipWithM_, mapAndUnzipM )
42 import Data.List ( inits, tails, zipWith4, zipWith5 )
44 -- ----------------------------------------------------------------------------
47 vectTyCon :: TyCon -> VM TyCon
49 | isFunTyCon tc = builtin closureTyCon
50 | isBoxedTupleTyCon tc = return tc
51 | isUnLiftedTyCon tc = return tc
52 | otherwise = maybeCantVectoriseM "Tycon not vectorised:" (ppr tc)
55 vectAndLiftType :: Type -> VM (Type, Type)
56 vectAndLiftType ty | Just ty' <- coreView ty = vectAndLiftType ty'
59 mdicts <- mapM paDictArgType tyvars
60 let dicts = [dict | Just dict <- mdicts]
61 vmono_ty <- vectType mono_ty
62 lmono_ty <- mkPArrayType vmono_ty
63 return (abstractType tyvars dicts vmono_ty,
64 abstractType tyvars dicts lmono_ty)
66 (tyvars, mono_ty) = splitForAllTys ty
69 vectType :: Type -> VM Type
70 vectType ty | Just ty' <- coreView ty = vectType ty'
71 vectType (TyVarTy tv) = return $ TyVarTy tv
72 vectType (AppTy ty1 ty2) = liftM2 AppTy (vectType ty1) (vectType ty2)
73 vectType (TyConApp tc tys) = liftM2 TyConApp (vectTyCon tc) (mapM vectType tys)
74 vectType (FunTy ty1 ty2) = liftM2 TyConApp (builtin closureTyCon)
75 (mapM vectAndBoxType [ty1,ty2])
76 vectType ty@(ForAllTy _ _)
78 mdicts <- mapM paDictArgType tyvars
79 mono_ty' <- vectType mono_ty
80 return $ abstractType tyvars [dict | Just dict <- mdicts] mono_ty'
82 (tyvars, mono_ty) = splitForAllTys ty
84 vectType ty = cantVectorise "Can't vectorise type" (ppr ty)
86 vectAndBoxType :: Type -> VM Type
87 vectAndBoxType ty = vectType ty >>= boxType
89 abstractType :: [TyVar] -> [Type] -> Type -> Type
90 abstractType tyvars dicts = mkForAllTys tyvars . mkFunTys dicts
92 -- ----------------------------------------------------------------------------
95 boxType :: Type -> VM Type
97 | Just (tycon, []) <- splitTyConApp_maybe ty
98 , isUnLiftedTyCon tycon
100 r <- lookupBoxedTyCon tycon
102 Just tycon' -> return $ mkTyConApp tycon' []
104 boxType ty = return ty
106 -- ----------------------------------------------------------------------------
109 type TyConGroup = ([TyCon], UniqSet TyCon)
111 vectTypeEnv :: TypeEnv -> VM (TypeEnv, [FamInst], [(Var, CoreExpr)])
114 cs <- readGEnv $ mk_map . global_tycons
115 let (conv_tcs, keep_tcs) = classifyTyCons cs groups
116 keep_dcs = concatMap tyConDataCons keep_tcs
117 zipWithM_ defTyCon keep_tcs keep_tcs
118 zipWithM_ defDataCon keep_dcs keep_dcs
119 new_tcs <- vectTyConDecls conv_tcs
121 let orig_tcs = keep_tcs ++ conv_tcs
122 vect_tcs = keep_tcs ++ new_tcs
124 repr_tcs <- zipWithM buildPReprTyCon orig_tcs vect_tcs
125 parr_tcs <- zipWithM buildPArrayTyCon orig_tcs vect_tcs
126 dfuns <- mapM mkPADFun vect_tcs
127 defTyConPAs (zip vect_tcs dfuns)
128 binds <- sequence (zipWith5 buildTyConBindings orig_tcs
134 let all_new_tcs = new_tcs ++ repr_tcs ++ parr_tcs
136 let new_env = extendTypeEnvList env
137 (map ATyCon all_new_tcs
138 ++ [ADataCon dc | tc <- all_new_tcs
139 , dc <- tyConDataCons tc])
141 return (new_env, map mkLocalFamInst (repr_tcs ++ parr_tcs), concat binds)
143 tycons = typeEnvTyCons env
144 groups = tyConGroups tycons
146 mk_map env = listToUFM_Directly [(u, getUnique n /= u) | (u,n) <- nameEnvUniqueElts env]
149 vectTyConDecls :: [TyCon] -> VM [TyCon]
150 vectTyConDecls tcs = fixV $ \tcs' ->
152 mapM_ (uncurry defTyCon) (zipLazy tcs tcs')
153 mapM vectTyConDecl tcs
155 vectTyConDecl :: TyCon -> VM TyCon
158 name' <- cloneName mkVectTyConOcc name
159 rhs' <- vectAlgTyConRhs tc (algTyConRhs tc)
161 liftDs $ buildAlgTyCon name'
163 [] -- no stupid theta
165 rec_flag -- FIXME: is this ok?
166 False -- FIXME: no generics
167 False -- not GADT syntax
168 Nothing -- not a family instance
171 tyvars = tyConTyVars tc
172 rec_flag = boolToRecFlag (isRecursiveTyCon tc)
174 vectAlgTyConRhs :: TyCon -> AlgTyConRhs -> VM AlgTyConRhs
175 vectAlgTyConRhs _ (DataTyCon { data_cons = data_cons
179 data_cons' <- mapM vectDataCon data_cons
180 zipWithM_ defDataCon data_cons data_cons'
181 return $ DataTyCon { data_cons = data_cons'
184 vectAlgTyConRhs tc _ = cantVectorise "Can't vectorise type definition:" (ppr tc)
186 vectDataCon :: DataCon -> VM DataCon
188 | not . null $ dataConExTyVars dc
189 = cantVectorise "Can't vectorise constructor (existentials):" (ppr dc)
190 | not . null $ dataConEqSpec dc
191 = cantVectorise "Can't vectorise constructor (eq spec):" (ppr dc)
194 name' <- cloneName mkVectDataConOcc name
195 tycon' <- vectTyCon tycon
196 arg_tys <- mapM vectType rep_arg_tys
198 liftDs $ buildDataCon name'
200 (map (const NotMarkedStrict) arg_tys)
201 [] -- no labelled fields
203 [] -- no existential tvs for now
204 [] -- no eq spec for now
209 name = dataConName dc
210 univ_tvs = dataConUnivTyVars dc
211 rep_arg_tys = dataConRepArgTys dc
212 tycon = dataConTyCon dc
214 mk_fam_inst :: TyCon -> TyCon -> (TyCon, [Type])
215 mk_fam_inst fam_tc arg_tc
216 = (fam_tc, [mkTyConApp arg_tc . mkTyVarTys $ tyConTyVars arg_tc])
218 buildPReprTyCon :: TyCon -> TyCon -> VM TyCon
219 buildPReprTyCon orig_tc vect_tc
221 name <- cloneName mkPReprTyConOcc (tyConName orig_tc)
222 rhs_ty <- buildPReprType vect_tc
223 prepr_tc <- builtin preprTyCon
224 liftDs $ buildSynTyCon name
226 (SynonymTyCon rhs_ty)
228 (Just $ mk_fam_inst prepr_tc vect_tc)
230 tyvars = tyConTyVars vect_tc
233 data Repr = ProdRepr {
234 prod_components :: [Type]
235 , prod_tycon :: TyCon
236 , prod_data_con :: DataCon
237 , prod_arr_tycon :: TyCon
238 , prod_arr_data_con :: DataCon
242 sum_components :: [Repr]
244 , sum_arr_tycon :: TyCon
245 , sum_arr_data_con :: DataCon
252 , void_bottom :: CoreExpr
257 , enum_data_con :: DataCon
258 , enum_arr_tycon :: TyCon
259 , enum_arr_data_con :: DataCon
265 tycon <- builtin voidTyCon
266 var <- builtin voidVar
269 , void_bottom = Var var
276 tycon <- builtin enumerationTyCon
277 let [data_con] = tyConDataCons tycon
278 (arr_tycon, _) <- parrayReprTyCon (mkTyConApp tycon [])
279 let [arr_data_con] = tyConDataCons arr_tycon
283 , enum_data_con = data_con
284 , enum_arr_tycon = arr_tycon
285 , enum_arr_data_con = arr_data_con
289 unboxedProductRepr :: [Type] -> VM Repr
290 unboxedProductRepr [] = voidRepr
291 unboxedProductRepr [ty] = return $ IdRepr ty
292 unboxedProductRepr tys = boxedProductRepr tys
294 boxedProductRepr :: [Type] -> VM Repr
297 tycon <- builtin (prodTyCon arity)
298 let [data_con] = tyConDataCons tycon
300 tys' <- mapM boxType tys
301 (arr_tycon, _) <- parrayReprTyCon $ mkTyConApp tycon tys'
302 let [arr_data_con] = tyConDataCons arr_tycon
305 prod_components = tys
307 , prod_data_con = data_con
308 , prod_arr_tycon = arr_tycon
309 , prod_arr_data_con = arr_data_con
314 sumRepr :: [Repr] -> VM Repr
315 sumRepr [] = voidRepr
316 sumRepr [repr] = boxRepr repr
319 tycon <- builtin (sumTyCon arity)
320 (arr_tycon, _) <- parrayReprTyCon
324 let [arr_data_con] = tyConDataCons arr_tycon
327 sum_components = reprs
329 , sum_arr_tycon = arr_tycon
330 , sum_arr_data_con = arr_data_con
335 splitSumRepr :: Repr -> [Repr]
336 splitSumRepr (SumRepr { sum_components = reprs }) = reprs
337 splitSumRepr repr = [repr]
339 boxRepr :: Repr -> VM Repr
340 boxRepr (VoidRepr {}) = boxedProductRepr []
341 boxRepr (IdRepr ty) = boxedProductRepr [ty]
342 boxRepr repr = return repr
344 reprType :: Repr -> Type
345 reprType (ProdRepr { prod_tycon = tycon, prod_components = tys })
346 = mkTyConApp tycon tys
347 reprType (SumRepr { sum_tycon = tycon, sum_components = reprs })
348 = mkTyConApp tycon (map reprType reprs)
349 reprType (IdRepr ty) = ty
350 reprType (VoidRepr { void_tycon = tycon }) = mkTyConApp tycon []
351 reprType (EnumRepr { enum_tycon = tycon }) = mkTyConApp tycon []
353 arrReprType :: Repr -> VM Type
354 arrReprType = mkPArrayType . reprType
356 arrShapeTys :: Repr -> VM [Type]
357 arrShapeTys (SumRepr {}) = sumShapeTys
358 arrShapeTys (ProdRepr {}) = return [intPrimTy]
359 arrShapeTys (IdRepr _) = return []
360 arrShapeTys (VoidRepr {}) = return [intPrimTy]
361 arrShapeTys (EnumRepr {}) = sumShapeTys
363 sumShapeTys :: VM [Type]
365 int_arr <- builtin intPrimArrayTy
366 return [intPrimTy, int_arr, int_arr]
369 arrShapeVars :: Repr -> VM [Var]
370 arrShapeVars repr = mapM (newLocalVar (fsLit "sh")) =<< arrShapeTys repr
372 replicateShape :: Repr -> CoreExpr -> CoreExpr -> VM [CoreExpr]
373 replicateShape (ProdRepr {}) len _ = return [len]
374 replicateShape (SumRepr {}) len tag = replicateSumShape len tag
375 replicateShape (IdRepr _) _ _ = return []
376 replicateShape (VoidRepr {}) len _ = return [len]
377 replicateShape (EnumRepr {}) len tag = replicateSumShape len tag
379 replicateSumShape :: CoreExpr -> CoreExpr -> VM [CoreExpr]
380 replicateSumShape len tag
382 rep <- builtin replicatePAIntPrimVar
383 up <- builtin upToPAIntPrimVar
384 return [len, Var rep `mkApps` [len, tag], Var up `App` len]
386 arrSelector :: Repr -> [CoreExpr] -> VM (CoreExpr, CoreExpr, CoreExpr)
387 arrSelector (SumRepr {}) [len, sel, is] = return (len, sel, is)
388 arrSelector (EnumRepr {}) [len, sel, is] = return (len, sel, is)
389 arrSelector _ _ = panic "arrSelector"
391 emptyArrRepr :: Repr -> VM [CoreExpr]
392 emptyArrRepr (SumRepr { sum_components = prods })
393 = liftM concat $ mapM emptyArrRepr prods
394 emptyArrRepr (ProdRepr { prod_components = [] })
395 = return [Var unitDataConId]
396 emptyArrRepr (ProdRepr { prod_components = tys })
398 emptyArrRepr (IdRepr ty)
399 = liftM singleton $ emptyPA ty
400 emptyArrRepr (VoidRepr { void_tycon = tycon })
401 = liftM singleton $ emptyPA (mkTyConApp tycon [])
402 emptyArrRepr (EnumRepr {})
405 arrReprTys :: Repr -> VM [Type]
406 arrReprTys (SumRepr { sum_components = reprs })
407 = liftM concat $ mapM arrReprTys reprs
408 arrReprTys (ProdRepr { prod_components = [] })
410 arrReprTys (ProdRepr { prod_components = tys })
411 = mapM mkPArrayType tys
412 arrReprTys (IdRepr ty)
413 = liftM singleton $ mkPArrayType ty
414 arrReprTys (VoidRepr { void_tycon = tycon })
415 = liftM singleton $ mkPArrayType (mkTyConApp tycon [])
416 arrReprTys (EnumRepr {})
419 arrReprTys' :: Repr -> VM [[Type]]
420 arrReprTys' (SumRepr { sum_components = reprs })
421 = mapM arrReprTys reprs
422 arrReprTys' repr = liftM singleton $ arrReprTys repr
424 arrReprVars :: Repr -> VM [[Var]]
426 = mapM (mapM (newLocalVar (fsLit "rs"))) =<< arrReprTys' repr
428 mkRepr :: TyCon -> VM Repr
430 | [tys] <- rep_tys = boxedProductRepr tys
431 -- removed: | all null rep_tys = enumRepr
432 | otherwise = sumRepr =<< mapM unboxedProductRepr rep_tys
434 rep_tys = map dataConRepArgTys $ tyConDataCons vect_tc
436 buildPReprType :: TyCon -> VM Type
437 buildPReprType = liftM reprType . mkRepr
439 buildToPRepr :: Repr -> TyCon -> TyCon -> TyCon -> VM CoreExpr
440 buildToPRepr repr vect_tc prepr_tc _
442 arg <- newLocalVar (fsLit "x") arg_ty
443 result <- to_repr repr (Var arg)
446 . wrapFamInstBody prepr_tc var_tys
449 var_tys = mkTyVarTys $ tyConTyVars vect_tc
450 arg_ty = mkTyConApp vect_tc var_tys
451 res_ty = reprType repr
453 cons = tyConDataCons vect_tc
456 to_repr (SumRepr { sum_components = prods
457 , sum_tycon = tycon })
460 (vars, bodies) <- mapAndUnzipM to_unboxed prods
461 return . Case expr (mkWildId (exprType expr)) res_ty
462 $ zipWith4 mk_alt cons vars (tyConDataCons tycon) bodies
464 mk_alt con vars sum_con body
465 = (DataAlt con, vars, mkConApp sum_con (ty_args ++ [body]))
467 ty_args = map (Type . reprType) prods
469 to_repr (EnumRepr { enum_data_con = data_con }) expr
470 = return . Case expr (mkWildId (exprType expr)) res_ty
473 mk_alt con = (DataAlt con, [], mkConApp data_con [mkDataConTag con])
477 (vars, body) <- to_unboxed prod
478 return $ Case expr (mkWildId (exprType expr)) res_ty
479 [(DataAlt con, vars, body)]
481 to_unboxed (ProdRepr { prod_components = tys
482 , prod_data_con = data_con })
484 vars <- mapM (newLocalVar (fsLit "r")) tys
485 return (vars, mkConApp data_con (map Type tys ++ map Var vars))
487 to_unboxed (IdRepr ty)
489 var <- newLocalVar (fsLit "y") ty
490 return ([var], Var var)
492 to_unboxed (VoidRepr { void_bottom = bottom })
493 = return ([], bottom)
495 to_unboxed _ = panic "buildToPRepr/to_unboxed"
498 buildFromPRepr :: Repr -> TyCon -> TyCon -> TyCon -> VM CoreExpr
499 buildFromPRepr repr vect_tc prepr_tc _
501 arg_ty <- mkPReprType res_ty
502 arg <- newLocalVar (fsLit "x") arg_ty
506 $ unwrapFamInstScrut prepr_tc var_tys (Var arg)
508 var_tys = mkTyVarTys $ tyConTyVars vect_tc
509 res_ty = mkTyConApp vect_tc var_tys
511 cons = map (`mkConApp` map Type var_tys) (tyConDataCons vect_tc)
514 from_repr repr@(SumRepr { sum_components = prods
515 , sum_tycon = tycon })
518 vars <- mapM (newLocalVar (fsLit "x")) (map reprType prods)
519 bodies <- sequence . zipWith3 from_unboxed prods cons
521 return . Case expr (mkWildId (reprType repr)) res_ty
522 $ zipWith3 sum_alt (tyConDataCons tycon) vars bodies
524 sum_alt data_con var body = (DataAlt data_con, [var], body)
526 from_repr repr@(EnumRepr { enum_data_con = data_con }) expr
528 var <- newLocalVar (fsLit "n") intPrimTy
530 let res = Case (Var var) (mkWildId intPrimTy) res_ty
531 $ (DEFAULT, [], error_expr)
532 : zipWith mk_alt (tyConDataCons vect_tc) cons
534 return $ Case expr (mkWildId (reprType repr)) res_ty
535 [(DataAlt data_con, [var], res)]
537 mk_alt data_con con = (LitAlt (mkDataConTagLit data_con), [], con)
539 error_expr = mkRuntimeErrorApp rUNTIME_ERROR_ID res_ty
541 $ sep [text "Invalid NDP representation of", ppr vect_tc]
543 from_repr repr expr = from_unboxed repr con expr
545 from_unboxed prod@(ProdRepr { prod_components = tys
546 , prod_data_con = data_con })
550 vars <- mapM (newLocalVar (fsLit "y")) tys
551 return $ Case expr (mkWildId (reprType prod)) res_ty
552 [(DataAlt data_con, vars, con `mkVarApps` vars)]
554 from_unboxed (IdRepr _) con expr
555 = return $ con `App` expr
557 from_unboxed (VoidRepr {}) con _
560 from_unboxed _ _ _ = panic "buildFromPRepr/from_unboxed"
562 buildToArrPRepr :: Repr -> TyCon -> TyCon -> TyCon -> VM CoreExpr
563 buildToArrPRepr repr vect_tc prepr_tc arr_tc
565 arg_ty <- mkPArrayType el_ty
566 arg <- newLocalVar (fsLit "xs") arg_ty
568 res_ty <- mkPArrayType (reprType repr)
570 shape_vars <- arrShapeVars repr
571 repr_vars <- arrReprVars repr
573 parray_co <- mkBuiltinCo parrayTyCon
575 let Just repr_co = tyConFamilyCoercion_maybe prepr_tc
576 co = mkAppCoercion parray_co
578 $ mkTyConApp repr_co var_tys
580 scrut = unwrapFamInstScrut arr_tc var_tys (Var arg)
582 result <- to_repr shape_vars repr_vars repr
586 $ Case scrut (mkWildId (mkTyConApp arr_tc var_tys)) res_ty
587 [(DataAlt arr_dc, shape_vars ++ concat repr_vars, result)]
589 var_tys = mkTyVarTys $ tyConTyVars vect_tc
590 el_ty = mkTyConApp vect_tc var_tys
592 [arr_dc] = tyConDataCons arr_tc
594 to_repr shape_vars@(_ : _)
596 (SumRepr { sum_components = prods
597 , sum_arr_tycon = tycon
598 , sum_arr_data_con = data_con })
600 exprs <- zipWithM to_prod repr_vars prods
602 return . wrapFamInstBody tycon tys
604 $ map Type tys ++ map Var shape_vars ++ exprs
606 tys = map reprType prods
610 (ProdRepr { prod_components = tys
611 , prod_arr_tycon = tycon
612 , prod_arr_data_con = data_con })
613 = return . wrapFamInstBody tycon tys
615 $ map Type tys ++ map Var (len_var : repr_vars)
619 (EnumRepr { enum_arr_tycon = tycon
620 , enum_arr_data_con = data_con })
621 = return . wrapFamInstBody tycon []
625 to_repr _ _ _ = panic "buildToArrPRepr/to_repr"
627 to_prod repr_vars@(r : _)
628 (ProdRepr { prod_components = tys@(ty : _)
629 , prod_arr_tycon = tycon
630 , prod_arr_data_con = data_con })
632 len <- lengthPA ty (Var r)
633 return . wrapFamInstBody tycon tys
635 $ map Type tys ++ len : map Var repr_vars
637 to_prod [var] (IdRepr _) = return (Var var)
638 to_prod [var] (VoidRepr {}) = return (Var var)
639 to_prod _ _ = panic "buildToArrPRepr/to_prod"
642 buildFromArrPRepr :: Repr -> TyCon -> TyCon -> TyCon -> VM CoreExpr
643 buildFromArrPRepr repr vect_tc prepr_tc arr_tc
645 arg_ty <- mkPArrayType =<< mkPReprType el_ty
646 arg <- newLocalVar (fsLit "xs") arg_ty
648 res_ty <- mkPArrayType el_ty
650 shape_vars <- arrShapeVars repr
651 repr_vars <- arrReprVars repr
653 parray_co <- mkBuiltinCo parrayTyCon
655 let Just repr_co = tyConFamilyCoercion_maybe prepr_tc
656 co = mkAppCoercion parray_co
657 $ mkTyConApp repr_co var_tys
659 scrut = mkCoerce co (Var arg)
661 result = wrapFamInstBody arr_tc var_tys
663 $ map Type var_tys ++ map Var (shape_vars ++ concat repr_vars)
666 (from_repr repr scrut shape_vars repr_vars res_ty result)
668 var_tys = mkTyVarTys $ tyConTyVars vect_tc
669 el_ty = mkTyConApp vect_tc var_tys
671 [arr_dc] = tyConDataCons arr_tc
673 from_repr (SumRepr { sum_components = prods
674 , sum_arr_tycon = tycon
675 , sum_arr_data_con = data_con })
682 vars <- mapM (newLocalVar (fsLit "xs")) =<< mapM arrReprType prods
683 result <- go prods repr_vars vars body
685 let scrut = unwrapFamInstScrut tycon ty_args expr
686 return . Case scrut (mkWildId scrut_ty) res_ty
687 $ [(DataAlt data_con, shape_vars ++ vars, result)]
689 ty_args = map reprType prods
690 scrut_ty = mkTyConApp tycon ty_args
692 go [] [] [] body = return body
693 go (prod : prods) (repr_vars : rss) (var : vars) body
695 shape_vars <- mapM (newLocalVar (fsLit "s")) =<< arrShapeTys prod
697 from_prod prod (Var var) shape_vars repr_vars res_ty
698 =<< go prods rss vars body
699 go _ _ _ _ = panic "buildFromArrPRepr/go"
701 from_repr repr expr shape_vars [repr_vars] res_ty body
702 = from_prod repr expr shape_vars repr_vars res_ty body
704 from_repr _ _ _ _ _ _ = panic "buildFromArrPRepr/from_repr"
706 from_prod (ProdRepr { prod_components = tys
707 , prod_arr_tycon = tycon
708 , prod_arr_data_con = data_con })
715 let scrut = unwrapFamInstScrut tycon tys expr
716 scrut_ty = mkTyConApp tycon tys
718 return $ Case scrut (mkWildId scrut_ty) res_ty
719 [(DataAlt data_con, shape_vars ++ repr_vars, body)]
721 from_prod (EnumRepr { enum_arr_tycon = tycon
722 , enum_arr_data_con = data_con })
728 = let scrut = unwrapFamInstScrut tycon [] expr
729 scrut_ty = mkTyConApp tycon []
731 return $ Case scrut (mkWildId scrut_ty) res_ty
732 [(DataAlt data_con, shape_vars, body)]
740 = return $ Let (NonRec repr_var expr) body
742 from_prod (VoidRepr {})
748 = return $ Let (NonRec repr_var expr) body
750 from_prod _ _ _ _ _ _ = panic "buildFromArrPRepr/from_prod"
752 buildPRDictRepr :: Repr -> VM CoreExpr
753 buildPRDictRepr (VoidRepr { void_tycon = tycon })
754 = prDFunOfTyCon tycon
755 buildPRDictRepr (IdRepr ty) = mkPR ty
756 buildPRDictRepr (ProdRepr {
757 prod_components = tys
762 dfun <- prDFunOfTyCon tycon
763 return $ dfun `mkTyApps` tys `mkApps` prs
765 buildPRDictRepr (SumRepr {
766 sum_components = prods
767 , sum_tycon = tycon })
769 prs <- mapM buildPRDictRepr prods
770 dfun <- prDFunOfTyCon tycon
771 return $ dfun `mkTyApps` map reprType prods `mkApps` prs
773 buildPRDictRepr (EnumRepr { enum_tycon = tycon })
774 = prDFunOfTyCon tycon
776 buildPRDict :: Repr -> TyCon -> TyCon -> TyCon -> VM CoreExpr
777 buildPRDict repr vect_tc prepr_tc _
779 dict <- buildPRDictRepr repr
781 pr_co <- mkBuiltinCo prTyCon
782 let co = mkAppCoercion pr_co
784 $ mkTyConApp arg_co var_tys
786 return $ mkCoerce co dict
788 var_tys = mkTyVarTys $ tyConTyVars vect_tc
790 Just arg_co = tyConFamilyCoercion_maybe prepr_tc
792 buildPArrayTyCon :: TyCon -> TyCon -> VM TyCon
793 buildPArrayTyCon orig_tc vect_tc = fixV $ \repr_tc ->
795 name' <- cloneName mkPArrayTyConOcc orig_name
796 rhs <- buildPArrayTyConRhs orig_name vect_tc repr_tc
797 parray <- builtin parrayTyCon
799 liftDs $ buildAlgTyCon name'
801 [] -- no stupid theta
803 rec_flag -- FIXME: is this ok?
804 False -- FIXME: no generics
805 False -- not GADT syntax
806 (Just $ mk_fam_inst parray vect_tc)
808 orig_name = tyConName orig_tc
809 tyvars = tyConTyVars vect_tc
810 rec_flag = boolToRecFlag (isRecursiveTyCon vect_tc)
813 buildPArrayTyConRhs :: Name -> TyCon -> TyCon -> VM AlgTyConRhs
814 buildPArrayTyConRhs orig_name vect_tc repr_tc
816 data_con <- buildPArrayDataCon orig_name vect_tc repr_tc
817 return $ DataTyCon { data_cons = [data_con], is_enum = False }
819 buildPArrayDataCon :: Name -> TyCon -> TyCon -> VM DataCon
820 buildPArrayDataCon orig_name vect_tc repr_tc
822 dc_name <- cloneName mkPArrayDataConOcc orig_name
823 repr <- mkRepr vect_tc
825 shape_tys <- arrShapeTys repr
826 repr_tys <- arrReprTys repr
828 let tys = shape_tys ++ repr_tys
830 liftDs $ buildDataCon dc_name
832 (map (const NotMarkedStrict) tys)
833 [] -- no field labels
834 (tyConTyVars vect_tc)
835 [] -- no existentials
841 mkPADFun :: TyCon -> VM Var
843 = newExportedVar (mkPADFunOcc $ getOccName vect_tc) =<< paDFunType vect_tc
845 buildTyConBindings :: TyCon -> TyCon -> TyCon -> TyCon -> Var
846 -> VM [(Var, CoreExpr)]
847 buildTyConBindings orig_tc vect_tc prepr_tc arr_tc dfun
849 repr <- mkRepr vect_tc
850 vectDataConWorkers repr orig_tc vect_tc arr_tc
851 dict <- buildPADict repr vect_tc prepr_tc arr_tc dfun
853 return $ (dfun, dict) : binds
855 vectDataConWorkers :: Repr -> TyCon -> TyCon -> TyCon
857 vectDataConWorkers repr orig_tc vect_tc arr_tc
860 . zipWith3 def_worker (tyConDataCons orig_tc) rep_tys
861 $ zipWith4 mk_data_con (tyConDataCons vect_tc)
865 mapM_ (uncurry hoistBinding) bs
867 tyvars = tyConTyVars vect_tc
868 var_tys = mkTyVarTys tyvars
869 ty_args = map Type var_tys
871 res_ty = mkTyConApp vect_tc var_tys
873 rep_tys = map dataConRepArgTys $ tyConDataCons vect_tc
874 reprs = splitSumRepr repr
876 [arr_dc] = tyConDataCons arr_tc
878 mk_data_con con tys pre post
879 = liftM2 (,) (vect_data_con con)
880 (lift_data_con tys pre post (mkDataConTag con))
882 vect_data_con con = return $ mkConApp con ty_args
883 lift_data_con tys pre_reprs post_reprs tag
885 len <- builtin liftingContext
886 args <- mapM (newLocalVar (fsLit "xs"))
887 =<< mapM mkPArrayType tys
889 shape <- replicateShape repr (Var len) tag
890 repr <- mk_arr_repr (Var len) (map Var args)
892 pre <- liftM concat $ mapM emptyArrRepr pre_reprs
893 post <- liftM concat $ mapM emptyArrRepr post_reprs
895 return . mkLams (len : args)
896 . wrapFamInstBody arr_tc var_tys
898 $ ty_args ++ shape ++ pre ++ repr ++ post
902 units <- replicatePA len (Var unitDataConId)
905 mk_arr_repr _ arrs = return arrs
907 def_worker data_con arg_tys mk_body
911 . polyAbstract tyvars $ \abstract ->
912 liftM (abstract . vectorised)
913 $ buildClosures tyvars [] arg_tys res_ty mk_body
915 vect_worker <- cloneId mkVectOcc orig_worker (exprType body)
916 defGlobalVar orig_worker vect_worker
917 return (vect_worker, body)
919 orig_worker = dataConWorkId data_con
921 buildPADict :: Repr -> TyCon -> TyCon -> TyCon -> Var -> VM CoreExpr
922 buildPADict repr vect_tc prepr_tc arr_tc _
923 = polyAbstract tvs $ \abstract ->
925 meth_binds <- mapM (mk_method repr) paMethods
926 let meth_exprs = map (Var . fst) meth_binds
928 pa_dc <- builtin paDataCon
929 let dict = mkConApp pa_dc (Type (mkTyConApp vect_tc arg_tys) : meth_exprs)
930 body = Let (Rec meth_binds) dict
931 return . mkInlineMe $ abstract body
933 tvs = tyConTyVars arr_tc
934 arg_tys = mkTyVarTys tvs
936 mk_method repr (name, build)
939 body <- build repr vect_tc prepr_tc arr_tc
940 var <- newLocalVar name (exprType body)
941 return (var, mkInlineMe body)
943 paMethods :: [(FastString, Repr -> TyCon -> TyCon -> TyCon -> VM CoreExpr)]
944 paMethods = [(fsLit "toPRepr", buildToPRepr),
945 (fsLit "fromPRepr", buildFromPRepr),
946 (fsLit "toArrPRepr", buildToArrPRepr),
947 (fsLit "fromArrPRepr", buildFromArrPRepr),
948 (fsLit "dictPRepr", buildPRDict)]
950 -- | Split the given tycons into two sets depending on whether they have to be
951 -- converted (first list) or not (second list). The first argument contains
952 -- information about the conversion status of external tycons:
954 -- * tycons which have converted versions are mapped to True
955 -- * tycons which are not changed by vectorisation are mapped to False
956 -- * tycons which can't be converted are not elements of the map
958 classifyTyCons :: UniqFM Bool -> [TyConGroup] -> ([TyCon], [TyCon])
959 classifyTyCons = classify [] []
961 classify conv keep _ [] = (conv, keep)
962 classify conv keep cs ((tcs, ds) : rs)
963 | can_convert && must_convert
964 = classify (tcs ++ conv) keep (cs `addListToUFM` [(tc,True) | tc <- tcs]) rs
966 = classify conv (tcs ++ keep) (cs `addListToUFM` [(tc,False) | tc <- tcs]) rs
968 = classify conv keep cs rs
970 refs = ds `delListFromUniqSet` tcs
972 can_convert = isNullUFM (refs `minusUFM` cs) && all convertable tcs
973 must_convert = foldUFM (||) False (intersectUFM_C const cs refs)
975 convertable tc = isDataTyCon tc && all isVanillaDataCon (tyConDataCons tc)
977 -- | Compute mutually recursive groups of tycons in topological order
979 tyConGroups :: [TyCon] -> [TyConGroup]
980 tyConGroups tcs = map mk_grp (stronglyConnCompFromEdgedVertices edges)
982 edges = [((tc, ds), tc, uniqSetToList ds) | tc <- tcs
983 , let ds = tyConsOfTyCon tc]
985 mk_grp (AcyclicSCC (tc, ds)) = ([tc], ds)
986 mk_grp (CyclicSCC els) = (tcs, unionManyUniqSets dss)
988 (tcs, dss) = unzip els
990 tyConsOfTyCon :: TyCon -> UniqSet TyCon
992 = tyConsOfTypes . concatMap dataConRepArgTys . tyConDataCons
994 tyConsOfType :: Type -> UniqSet TyCon
996 | Just ty' <- coreView ty = tyConsOfType ty'
997 tyConsOfType (TyVarTy _) = emptyUniqSet
998 tyConsOfType (TyConApp tc tys) = extend (tyConsOfTypes tys)
1000 extend | isUnLiftedTyCon tc
1001 || isTupleTyCon tc = id
1003 | otherwise = (`addOneToUniqSet` tc)
1005 tyConsOfType (AppTy a b) = tyConsOfType a `unionUniqSets` tyConsOfType b
1006 tyConsOfType (FunTy a b) = (tyConsOfType a `unionUniqSets` tyConsOfType b)
1007 `addOneToUniqSet` funTyCon
1008 tyConsOfType (ForAllTy _ ty) = tyConsOfType ty
1009 tyConsOfType other = pprPanic "ClosureConv.tyConsOfType" $ ppr other
1011 tyConsOfTypes :: [Type] -> UniqSet TyCon
1012 tyConsOfTypes = unionManyUniqSets . map tyConsOfType
1015 -- ----------------------------------------------------------------------------
1018 fromVect :: Type -> CoreExpr -> VM CoreExpr
1019 fromVect ty expr | Just ty' <- coreView ty = fromVect ty' expr
1020 fromVect (FunTy arg_ty res_ty) expr
1022 arg <- newLocalVar (fsLit "x") arg_ty
1023 varg <- toVect arg_ty (Var arg)
1024 varg_ty <- vectType arg_ty
1025 vres_ty <- vectType res_ty
1026 apply <- builtin applyClosureVar
1027 body <- fromVect res_ty
1028 $ Var apply `mkTyApps` [varg_ty, vres_ty] `mkApps` [expr, varg]
1029 return $ Lam arg body
1031 = identityConv ty >> return expr
1033 toVect :: Type -> CoreExpr -> VM CoreExpr
1034 toVect ty expr = identityConv ty >> return expr
1036 identityConv :: Type -> VM ()
1037 identityConv ty | Just ty' <- coreView ty = identityConv ty'
1038 identityConv (TyConApp tycon tys)
1040 mapM_ identityConv tys
1041 identityConvTyCon tycon
1042 identityConv _ = noV
1044 identityConvTyCon :: TyCon -> VM ()
1045 identityConvTyCon tc
1046 | isBoxedTupleTyCon tc = return ()
1047 | isUnLiftedTyCon tc = return ()
1049 tc' <- maybeV (lookupTyCon tc)
1050 if tc == tc' then return () else noV