compiler/vectorise/Vectorise/Utils/Closure.hs

   1
   2 module Vectorise.Utils.Closure (
   3         mkClosure,
   4         mkClosureApp,
   5         buildClosure,
   6         buildClosures,
   7         buildEnv
   8 )
   9 where
  10 import VectUtils
  11 import Vectorise.Builtins
  12 import Vectorise.Vect
  13 import Vectorise.Monad
  14
  15 import CoreSyn
  16 import Type
  17 import Var
  18 import MkCore
  19 import CoreUtils
  20 import TyCon
  21 import DataCon
  22 import MkId
  23 import TysWiredIn
  24 import BasicTypes
  25 import FastString
  26
  27
  28 mkClosure :: Type -> Type -> Type -> VExpr -> VExpr -> VM VExpr
  29 mkClosure arg_ty res_ty env_ty (vfn,lfn) (venv,lenv)
  30  = do Just dict <- paDictOfType env_ty
  31       mkv       <- builtin closureVar
  32       mkl       <- builtin liftedClosureVar
  33       return (Var mkv `mkTyApps` [arg_ty, res_ty, env_ty] `mkApps` [dict, vfn, lfn, venv],
  34               Var mkl `mkTyApps` [arg_ty, res_ty, env_ty] `mkApps` [dict, vfn, lfn, lenv])
  35
  36
  37 mkClosureApp :: Type -> Type -> VExpr -> VExpr -> VM VExpr
  38 mkClosureApp arg_ty res_ty (vclo, lclo) (varg, larg)
  39  = do vapply <- builtin applyVar
  40       lapply <- builtin liftedApplyVar
  41       lc     <- builtin liftingContext
  42       return (Var vapply `mkTyApps` [arg_ty, res_ty] `mkApps` [vclo, varg],
  43               Var lapply `mkTyApps` [arg_ty, res_ty] `mkApps` [Var lc, lclo, larg])
  44
  45
  46 buildClosures :: [TyVar] -> [VVar] -> [Type] -> Type -> VM VExpr -> VM VExpr
  47 buildClosures _   _    [] _ mk_body
  48   = mk_body
  49 buildClosures tvs vars [arg_ty] res_ty mk_body
  50   = -- liftM vInlineMe $
  51       buildClosure tvs vars arg_ty res_ty mk_body
  52 buildClosures tvs vars (arg_ty : arg_tys) res_ty mk_body
  53   = do
  54       res_ty' <- mkClosureTypes arg_tys res_ty
  55       arg <- newLocalVVar (fsLit "x") arg_ty
  56       -- liftM vInlineMe
  57       buildClosure tvs vars arg_ty res_ty'
  58         . hoistPolyVExpr tvs (Inline (length vars + 1))
  59         $ do
  60             lc <- builtin liftingContext
  61             clo <- buildClosures tvs (vars ++ [arg]) arg_tys res_ty mk_body
  62             return $ vLams lc (vars ++ [arg]) clo
  63
  64
  65 -- (clo <x1,...,xn> <f,f^>, aclo (Arr lc xs1 ... xsn) <f,f^>)
  66 --   where
  67 --     f  = \env v -> case env of <x1,...,xn> -> e x1 ... xn v
  68 --     f^ = \env v -> case env of Arr l xs1 ... xsn -> e^ l x1 ... xn v
  69 --
  70 buildClosure :: [TyVar] -> [VVar] -> Type -> Type -> VM VExpr -> VM VExpr
  71 buildClosure tvs vars arg_ty res_ty mk_body
  72   = do
  73       (env_ty, env, bind) <- buildEnv vars
  74       env_bndr <- newLocalVVar (fsLit "env") env_ty
  75       arg_bndr <- newLocalVVar (fsLit "arg") arg_ty
  76
  77       fn <- hoistPolyVExpr tvs (Inline 2)
  78           $ do
  79               lc    <- builtin liftingContext
  80               body  <- mk_body
  81               return -- . vInlineMe
  82                      . vLams lc [env_bndr, arg_bndr]
  83                      $ bind (vVar env_bndr)
  84                             (vVarApps lc body (vars ++ [arg_bndr]))
  85
  86       mkClosure arg_ty res_ty env_ty fn env
  87
  88
  89 -- Environments ---------------------------------------------------------------
  90 buildEnv :: [VVar] -> VM (Type, VExpr, VExpr -> VExpr -> VExpr)
  91 buildEnv [] = do
  92              ty    <- voidType
  93              void  <- builtin voidVar
  94              pvoid <- builtin pvoidVar
  95              return (ty, vVar (void, pvoid), \_ body -> body)
  96
  97 buildEnv [v] = return (vVarType v, vVar v,
  98                     \env body -> vLet (vNonRec v env) body)
  99
 100 buildEnv vs
 101   = do
 102
 103       (lenv_tc, lenv_tyargs) <- pdataReprTyCon ty
 104
 105       let venv_con   = tupleCon Boxed (length vs)
 106           [lenv_con] = tyConDataCons lenv_tc
 107
 108           venv       = mkCoreTup (map Var vvs)
 109           lenv       = Var (dataConWrapId lenv_con)
 110                        `mkTyApps` lenv_tyargs
 111                        `mkApps`   map Var lvs
 112
 113           vbind env body = mkWildCase env ty (exprType body)
 114                            [(DataAlt venv_con, vvs, body)]
 115
 116           lbind env body =
 117             let scrut = unwrapFamInstScrut lenv_tc lenv_tyargs env
 118             in
 119             mkWildCase scrut (exprType scrut) (exprType body)
 120               [(DataAlt lenv_con, lvs, body)]
 121
 122           bind (venv, lenv) (vbody, lbody) = (vbind venv vbody,
 123                                               lbind lenv lbody)
 124
 125       return (ty, (venv, lenv), bind)
 126   where
 127     (vvs, lvs) = unzip vs
 128     tys        = map vVarType vs
 129     ty         = mkBoxedTupleTy tys