X-Git-Url: http://git.megacz.com/?a=blobdiff_plain;f=compiler%2Fvectorise%2FVectType.hs;h=c977e4c7dee18b5c5f7f698368bd50f5bec6b90b;hb=eaaecbaefe18da05d618942c51286cacfa1be2af;hp=6f6fca8dbadc1737c4d5910a2f3e803978d2f68d;hpb=724425265ded8958a719b3a62f43006674b506c8;p=ghc-hetmet.git diff --git a/compiler/vectorise/VectType.hs b/compiler/vectorise/VectType.hs index 6f6fca8..c977e4c 100644 --- a/compiler/vectorise/VectType.hs +++ b/compiler/vectorise/VectType.hs @@ -26,7 +26,7 @@ import Var ( Var ) import Id ( mkWildId ) import Name ( Name, getOccName ) import NameEnv -import TysWiredIn ( unitTy, intTy, intDataCon, unitDataConId ) +import TysWiredIn ( unitTy, unitTyCon, intTy, intDataCon, unitDataConId ) import TysPrim ( intPrimTy ) import Unique @@ -208,16 +208,141 @@ buildPReprTyCon orig_tc vect_tc where tyvars = tyConTyVars vect_tc +data TyConRepr = ProdRepr { + repr_prod_arg_tys :: [Type] + , repr_prod_tycon :: TyCon + , repr_prod_data_con :: DataCon + , repr_prod_arr_tycon :: TyCon + , repr_prod_arr_data_con :: DataCon + , repr_type :: Type + } + | SumRepr { + repr_tys :: [[Type]] + , repr_prod_tycons :: [Maybe TyCon] + , repr_prod_data_cons :: [Maybe DataCon] + , repr_prod_tys :: [Type] + , repr_sum_tycon :: TyCon + , repr_type :: Type + } + +arrShapeTys :: TyConRepr -> VM [Type] +arrShapeTys (ProdRepr {}) = return [intPrimTy] +arrShapeTys (SumRepr {}) + = do + uarr <- builtin uarrTyCon + return [intPrimTy, mkTyConApp uarr [intTy]] + +arrReprTys :: TyConRepr -> VM [Type] +arrReprTys (ProdRepr { repr_prod_arg_tys = tys }) + = mapM mkPArrayType tys +arrReprTys (SumRepr { repr_tys = tys }) + = concat `liftM` mapM (mapM mkPArrayType) (map mk_prod tys) + where + mk_prod [] = [unitTy] + mk_prod tys = tys + + +mkTyConRepr :: TyCon -> VM TyConRepr +mkTyConRepr vect_tc + | is_product + = let + [prod_arg_tys] = repr_tys + arity = length prod_arg_tys + in + do + prod_tycon <- builtin (prodTyCon arity) + let [prod_data_con] = tyConDataCons prod_tycon + + (arr_tycon, _) <- parrayReprTyCon + . mkTyConApp prod_tycon + $ replicate arity unitTy + + let [arr_data_con] = tyConDataCons arr_tycon + + return $ ProdRepr { + repr_prod_arg_tys = prod_arg_tys + , repr_prod_tycon = prod_tycon + , repr_prod_data_con = prod_data_con + , repr_prod_arr_tycon = arr_tycon + , repr_prod_arr_data_con = arr_data_con + , repr_type = mkTyConApp prod_tycon prod_arg_tys + } + + | otherwise + = do + uarr <- builtin uarrTyCon + prod_tycons <- mapM (mk_tycon prodTyCon) repr_tys + let prod_tys = zipWith mk_tc_app_maybe prod_tycons repr_tys + sum_tycon <- builtin (sumTyCon $ length repr_tys) + arr_repr_tys <- mapM (mapM mkPArrayType . arr_repr_elem_tys) repr_tys + + return $ SumRepr { + repr_tys = repr_tys + , repr_prod_tycons = prod_tycons + , repr_prod_data_cons = map (fmap mk_single_datacon) prod_tycons + , repr_prod_tys = prod_tys + , repr_sum_tycon = sum_tycon + , repr_type = mkTyConApp sum_tycon prod_tys + } + where + tyvars = tyConTyVars vect_tc + data_cons = tyConDataCons vect_tc + repr_tys = map dataConRepArgTys data_cons + + is_product | [_] <- data_cons = True + | otherwise = False + + mk_shape uarr = intPrimTy : mk_sel uarr + + mk_sel uarr | is_product = [] + | otherwise = [uarr_int, uarr_int] + where + uarr_int = mkTyConApp uarr [intTy] + + mk_tycon get_tc tys + | n > 1 = builtin (Just . get_tc n) + | otherwise = return Nothing + where n = length tys + + mk_single_datacon tc | [dc] <- tyConDataCons tc = dc + + mk_tc_app_maybe Nothing [] = unitTy + mk_tc_app_maybe Nothing [ty] = ty + mk_tc_app_maybe (Just tc) tys = mkTyConApp tc tys + + arr_repr_elem_tys [] = [unitTy] + arr_repr_elem_tys tys = tys + buildPReprType :: TyCon -> VM Type buildPReprType = liftM repr_type . mkTyConRepr buildToPRepr :: TyConRepr -> TyCon -> TyCon -> TyCon -> VM CoreExpr -buildToPRepr (TyConRepr { - repr_tys = repr_tys - , repr_prod_tycons = prod_tycons - , repr_prod_tys = prod_tys - , repr_sum_tycon = repr_sum_tycon - , repr_type = repr_type +buildToPRepr (ProdRepr { + repr_prod_arg_tys = prod_arg_tys + , repr_prod_data_con = prod_data_con + , repr_type = repr_type + }) + vect_tc prepr_tc _ + = do + arg <- newLocalVar FSLIT("x") arg_ty + vars <- mapM (newLocalVar FSLIT("x")) prod_arg_tys + + return . Lam arg + . wrapFamInstBody prepr_tc var_tys + $ Case (Var arg) (mkWildId arg_ty) repr_type + [(DataAlt data_con, vars, + mkConApp prod_data_con (map Type prod_arg_tys ++ map Var vars))] + where + var_tys = mkTyVarTys $ tyConTyVars vect_tc + arg_ty = mkTyConApp vect_tc var_tys + [data_con] = tyConDataCons vect_tc + +buildToPRepr (SumRepr { + repr_tys = repr_tys + , repr_prod_data_cons = prod_data_cons + , repr_prod_tys = prod_tys + , repr_sum_tycon = sum_tycon + , repr_type = repr_type }) vect_tc prepr_tc _ = do @@ -227,30 +352,143 @@ buildToPRepr (TyConRepr { return . Lam arg . wrapFamInstBody prepr_tc var_tys . Case (Var arg) (mkWildId arg_ty) repr_type - . mk_alts data_cons vars - . zipWith3 mk_prod prod_tycons repr_tys $ map (map Var) vars + . zipWith4 mk_alt data_cons vars sum_data_cons + . zipWith3 mk_prod prod_data_cons repr_tys $ map (map Var) vars where var_tys = mkTyVarTys $ tyConTyVars vect_tc arg_ty = mkTyConApp vect_tc var_tys data_cons = tyConDataCons vect_tc - Just sum_tycon = repr_sum_tycon - sum_datacons = tyConDataCons sum_tycon - - mk_alts _ _ [] = [(DEFAULT, [], Var unitDataConId)] - mk_alts [dc] [vars] [expr] = [(DataAlt dc, vars, expr)] - mk_alts dcs vars exprs = zipWith4 mk_alt dcs vars sum_datacons exprs + sum_data_cons = tyConDataCons sum_tycon mk_alt dc vars sum_dc expr = (DataAlt dc, vars, mkConApp sum_dc (map Type prod_tys ++ [expr])) mk_prod _ _ [] = Var unitDataConId mk_prod _ _ [expr] = expr - mk_prod (Just tc) tys exprs = mkConApp dc (map Type tys ++ exprs) - where - [dc] = tyConDataCons tc + mk_prod (Just dc) tys exprs = mkConApp dc (map Type tys ++ exprs) + +buildFromPRepr :: TyConRepr -> TyCon -> TyCon -> TyCon -> VM CoreExpr +buildFromPRepr (ProdRepr { + repr_prod_arg_tys = prod_arg_tys + , repr_prod_data_con = prod_data_con + , repr_type = repr_type + }) + vect_tc prepr_tc _ + = do + arg_ty <- mkPReprType res_ty + arg <- newLocalVar FSLIT("x") arg_ty + vars <- mapM (newLocalVar FSLIT("x")) prod_arg_tys + + return . Lam arg + $ Case (unwrapFamInstScrut prepr_tc var_tys (Var arg)) + (mkWildId repr_type) + res_ty + [(DataAlt prod_data_con, vars, + mkConApp data_con (map Type var_tys ++ map Var vars))] + where + var_tys = mkTyVarTys $ tyConTyVars vect_tc + ty_args = map Type var_tys + res_ty = mkTyConApp vect_tc var_tys + [data_con] = tyConDataCons vect_tc + +buildFromPRepr (SumRepr { + repr_tys = repr_tys + , repr_prod_data_cons = prod_data_cons + , repr_prod_tys = prod_tys + , repr_sum_tycon = sum_tycon + , repr_type = repr_type + }) + vect_tc prepr_tc _ + = do + arg_ty <- mkPReprType res_ty + arg <- newLocalVar FSLIT("x") arg_ty + + liftM (Lam arg + . Case (unwrapFamInstScrut prepr_tc var_tys (Var arg)) + (mkWildId repr_type) + res_ty + . zipWith mk_alt sum_data_cons) + (sequence + $ zipWith4 un_prod data_cons prod_data_cons prod_tys repr_tys) + where + var_tys = mkTyVarTys $ tyConTyVars vect_tc + ty_args = map Type var_tys + res_ty = mkTyConApp vect_tc var_tys + data_cons = tyConDataCons vect_tc + + sum_data_cons = tyConDataCons sum_tycon + + un_prod dc _ _ [] + = do + var <- newLocalVar FSLIT("u") unitTy + return (var, mkConApp dc ty_args) + un_prod dc _ _ [ty] + = do + var <- newLocalVar FSLIT("x") ty + return (var, mkConApp dc (ty_args ++ [Var var])) + + un_prod dc (Just prod_dc) prod_ty tys + = do + vars <- mapM (newLocalVar FSLIT("x")) tys + pv <- newLocalVar FSLIT("p") prod_ty + + let res = mkConApp dc (ty_args ++ map Var vars) + expr = Case (Var pv) (mkWildId prod_ty) res_ty + [(DataAlt prod_dc, vars, res)] + + return (pv, expr) + + mk_alt sum_dc (var, expr) = (DataAlt sum_dc, [var], expr) + buildToArrPRepr :: TyConRepr -> TyCon -> TyCon -> TyCon -> VM CoreExpr +buildToArrPRepr repr@(ProdRepr { + repr_prod_arg_tys = prod_arg_tys + , repr_prod_arr_tycon = prod_arr_tycon + , repr_prod_arr_data_con = prod_arr_data_con + , repr_type = repr_type + }) + vect_tc prepr_tc arr_tc + = do + arg_ty <- mkPArrayType el_ty + shape_tys <- arrShapeTys repr + arr_tys <- arrReprTys repr + res_ty <- mkPArrayType repr_type + rep_el_ty <- mkPReprType el_ty + + arg <- newLocalVar FSLIT("xs") arg_ty + shape_vars <- mapM (newLocalVar FSLIT("sh")) shape_tys + rep_vars <- mapM (newLocalVar FSLIT("ys")) arr_tys + + let vars = shape_vars ++ rep_vars + + parray_co <- mkBuiltinCo parrayTyCon + + let res = wrapFamInstBody prod_arr_tycon prod_arg_tys + . mkConApp prod_arr_data_con + $ map Type prod_arg_tys ++ map Var vars + + Just repr_co = tyConFamilyCoercion_maybe prepr_tc + co = mkAppCoercion parray_co + . mkSymCoercion + $ mkTyConApp repr_co var_tys + + return . Lam arg + . mkCoerce co + $ Case (unwrapFamInstScrut arr_tc var_tys (Var arg)) + (mkWildId (mkTyConApp arr_tc var_tys)) + res_ty + [(DataAlt arr_dc, vars, res)] + where + var_tys = mkTyVarTys $ tyConTyVars vect_tc + el_ty = mkTyConApp vect_tc var_tys + + [arr_dc] = tyConDataCons arr_tc + + +buildToArrPRepr _ _ _ _ = return (Var unitDataConId) +{- buildToArrPRepr _ vect_tc prepr_tc arr_tc = do arg_ty <- mkPArrayType el_ty @@ -285,37 +523,99 @@ buildToArrPRepr _ vect_tc prepr_tc arr_tc has_selector | [_] <- data_cons = False | otherwise = True +-} - -buildFromPRepr :: TyConRepr -> TyCon -> TyCon -> TyCon -> VM CoreExpr -buildFromPRepr _ vect_tc prepr_tc _ +buildFromArrPRepr :: TyConRepr -> TyCon -> TyCon -> TyCon -> VM CoreExpr +buildFromArrPRepr repr@(ProdRepr { + repr_prod_arg_tys = prod_arg_tys + , repr_prod_arr_tycon = prod_arr_tycon + , repr_prod_arr_data_con = prod_arr_data_con + , repr_type = repr_type + }) + vect_tc prepr_tc arr_tc = do - arg_ty <- mkPReprType res_ty - arg <- newLocalVar FSLIT("x") arg_ty - alts <- mapM mk_alt data_cons - body <- mkFromPRepr (unwrapFamInstScrut prepr_tc var_tys (Var arg)) - res_ty alts - return $ Lam arg body + rep_el_ty <- mkPReprType el_ty + arg_ty <- mkPArrayType rep_el_ty + shape_tys <- arrShapeTys repr + arr_tys <- arrReprTys repr + res_ty <- mkPArrayType el_ty + + arg <- newLocalVar FSLIT("xs") arg_ty + shape_vars <- mapM (newLocalVar FSLIT("sh")) shape_tys + rep_vars <- mapM (newLocalVar FSLIT("ys")) arr_tys + + let vars = shape_vars ++ rep_vars + + parray_co <- mkBuiltinCo parrayTyCon + + let res = wrapFamInstBody arr_tc var_tys + . mkConApp arr_dc + $ map Type var_tys ++ map Var vars + + Just repr_co = tyConFamilyCoercion_maybe prepr_tc + co = mkAppCoercion parray_co + $ mkTyConApp repr_co var_tys + + scrut = unwrapFamInstScrut prod_arr_tycon prod_arg_tys + $ mkCoerce co (Var arg) + + return . Lam arg + $ Case (scrut) + (mkWildId (mkTyConApp prod_arr_tycon prod_arg_tys)) + res_ty + [(DataAlt prod_arr_data_con, vars, res)] where - var_tys = mkTyVarTys $ tyConTyVars vect_tc - res_ty = mkTyConApp vect_tc var_tys - data_cons = tyConDataCons vect_tc + var_tys = mkTyVarTys $ tyConTyVars vect_tc + el_ty = mkTyConApp vect_tc var_tys - mk_alt dc = do - bndrs <- mapM (newLocalVar FSLIT("x")) $ dataConRepArgTys dc - return (bndrs, mkConApp dc (map Type var_tys ++ map Var bndrs)) + [arr_dc] = tyConDataCons arr_tc +buildFromArrPRepr _ _ _ _ = return (Var unitDataConId) -buildFromArrPRepr :: TyConRepr -> TyCon -> TyCon -> TyCon -> VM CoreExpr -buildFromArrPRepr _ vect_tc prepr_tc arr_tc - = mkFromArrPRepr undefined undefined undefined undefined undefined undefined +buildPRDictRepr :: TyConRepr -> VM CoreExpr +buildPRDictRepr (ProdRepr { + repr_prod_arg_tys = prod_arg_tys + , repr_prod_tycon = prod_tycon + }) + = do + prs <- mapM mkPR prod_arg_tys + dfun <- prDFunOfTyCon prod_tycon + return $ dfun `mkTyApps` prod_arg_tys `mkApps` prs + +buildPRDictRepr (SumRepr { + repr_tys = repr_tys + , repr_prod_tycons = prod_tycons + , repr_prod_tys = prod_tys + , repr_sum_tycon = sum_tycon + }) + = do + prs <- mapM (mapM mkPR) repr_tys + prod_prs <- sequence $ zipWith3 mk_prod_pr prod_tycons repr_tys prs + sum_dfun <- prDFunOfTyCon sum_tycon + return $ sum_dfun `mkTyApps` prod_tys `mkApps` prod_prs + where + mk_prod_pr _ _ [] = prDFunOfTyCon unitTyCon + mk_prod_pr _ _ [pr] = return pr + mk_prod_pr (Just tc) tys prs + = do + dfun <- prDFunOfTyCon tc + return $ dfun `mkTyApps` tys `mkApps` prs buildPRDict :: TyConRepr -> TyCon -> TyCon -> TyCon -> VM CoreExpr -buildPRDict _ vect_tc prepr_tc _ - = prCoerce prepr_tc var_tys - =<< prDictOfType (mkTyConApp prepr_tc var_tys) +buildPRDict repr vect_tc prepr_tc _ + = do + dict <- buildPRDictRepr repr + + pr_co <- mkBuiltinCo prTyCon + let co = mkAppCoercion pr_co + . mkSymCoercion + $ mkTyConApp arg_co var_tys + + return $ mkCoerce co dict where var_tys = mkTyVarTys $ tyConTyVars vect_tc + Just arg_co = tyConFamilyCoercion_maybe prepr_tc + buildPArrayTyCon :: TyCon -> TyCon -> VM TyCon buildPArrayTyCon orig_tc vect_tc = fixV $ \repr_tc -> do @@ -347,22 +647,23 @@ buildPArrayDataCon :: Name -> TyCon -> TyCon -> VM DataCon buildPArrayDataCon orig_name vect_tc repr_tc = do dc_name <- cloneName mkPArrayDataConOcc orig_name - shape <- tyConShape vect_tc - repr_tys <- mapM mkPArrayType types + repr <- mkTyConRepr vect_tc + + shape_tys <- arrShapeTys repr + repr_tys <- arrReprTys repr + + let tys = shape_tys ++ repr_tys liftDs $ buildDataCon dc_name False -- not infix - (shapeStrictness shape ++ map (const NotMarkedStrict) repr_tys) + (map (const NotMarkedStrict) tys) [] -- no field labels (tyConTyVars vect_tc) [] -- no existentials [] -- no eq spec [] -- no context - (shapeReprTys shape ++ repr_tys) + tys repr_tc - where - types = [ty | dc <- tyConDataCons vect_tc - , ty <- dataConRepArgTys dc] mkPADFun :: TyCon -> VM Var mkPADFun vect_tc