Thread lifting context implicitly in the vectorisation monad

diff --git a/compiler/vectorise/VectMonad.hs b/compiler/vectorise/VectMonad.hs
index c244f0a..a991b8c 100644 (file)
--- a/compiler/vectorise/VectMonad.hs
+++ b/compiler/vectorise/VectMonad.hs
@@ -39,6 +39,7 @@ import Id
 import OccName
 import Name
 import NameEnv
+import TysPrim       ( intPrimTy )
 
 import DsMonad
 import PrelNames
@@ -69,6 +70,7 @@ data Builtins = Builtins {
                 , lengthPAVar      :: Var
                 , replicatePAVar   :: Var
                 , emptyPAVar       :: Var
+                , liftingContext   :: Var
                 }
 
 paDictTyCon :: Builtins -> TyCon
@@ -92,6 +94,9 @@ initBuiltins
       replicatePAVar   <- dsLookupGlobalId replicatePAName
       emptyPAVar       <- dsLookupGlobalId emptyPAName
 
+      liftingContext <- liftM (\u -> mkSysLocal FSLIT("lc") u intPrimTy)
+                              newUnique
+
       return $ Builtins {
                  parrayTyCon      = parrayTyCon
                , paClass          = paClass
@@ -103,6 +108,7 @@ initBuiltins
                , lengthPAVar      = lengthPAVar
                , replicatePAVar   = replicatePAVar
                , emptyPAVar       = emptyPAVar
+               , liftingContext   = liftingContext
                }
 
 data GlobalEnv = GlobalEnv {

diff --git a/compiler/vectorise/VectUtils.hs b/compiler/vectorise/VectUtils.hs
index 718db85..0727c94 100644 (file)
--- a/compiler/vectorise/VectUtils.hs
+++ b/compiler/vectorise/VectUtils.hs
@@ -4,7 +4,7 @@ module VectUtils (
   splitClosureTy,
   mkPADictType, mkPArrayType,
   paDictArgType, paDictOfType,
-  paMethod, lengthPA, replicatePA, emptyPA,
+  paMethod, lengthPA, replicatePA, emptyPA, liftPA,
   polyAbstract, polyApply, polyVApply,
   lookupPArrayFamInst,
   hoistExpr, hoistPolyVExpr, takeHoisted,
@@ -170,6 +170,12 @@ replicatePA len x = liftM (`mkApps` [len,x])
 emptyPA :: Type -> VM CoreExpr
 emptyPA = paMethod emptyPAVar
 
+liftPA :: CoreExpr -> VM CoreExpr
+liftPA x
+  = do
+      lc <- builtin liftingContext
+      replicatePA (Var lc) x
+
 newLocalVVar :: FastString -> Type -> VM VVar
 newLocalVVar fs vty
   = do
@@ -259,17 +265,18 @@ mkClosureApp (vclo, lclo) (varg, larg)
   where
     (arg_ty, res_ty) = splitClosureTy (exprType vclo)
 
-buildClosures :: [TyVar] -> Var -> [VVar] -> [Type] -> Type -> VM VExpr -> VM VExpr
-buildClosures tvs lc vars [arg_ty] res_ty mk_body
-  = buildClosure tvs lc vars arg_ty res_ty mk_body
-buildClosures tvs lc vars (arg_ty : arg_tys) res_ty mk_body
+buildClosures :: [TyVar] -> [VVar] -> [Type] -> Type -> VM VExpr -> VM VExpr
+buildClosures tvs vars [arg_ty] res_ty mk_body
+  = buildClosure tvs vars arg_ty res_ty mk_body
+buildClosures tvs vars (arg_ty : arg_tys) res_ty mk_body
   = do
       res_ty' <- mkClosureTypes arg_tys res_ty
       arg <- newLocalVVar FSLIT("x") arg_ty
-      buildClosure tvs lc vars arg_ty res_ty'
+      buildClosure tvs vars arg_ty res_ty'
         . hoistPolyVExpr tvs
         $ do
-            clo <- buildClosures tvs lc (vars ++ [arg]) arg_tys res_ty mk_body
+            lc <- builtin liftingContext
+            clo <- buildClosures tvs (vars ++ [arg]) arg_tys res_ty mk_body
             return $ vLams lc (vars ++ [arg]) clo
 
 -- (clo <x1,...,xn> <f,f^>, aclo (Arr lc xs1 ... xsn) <f,f^>)
@@ -277,27 +284,29 @@ buildClosures tvs lc vars (arg_ty : arg_tys) res_ty mk_body
 --     f  = \env v -> case env of <x1,...,xn> -> e x1 ... xn v
 --     f^ = \env v -> case env of Arr l xs1 ... xsn -> e^ l x1 ... xn v
 --
-buildClosure :: [TyVar] -> Var -> [VVar] -> Type -> Type -> VM VExpr -> VM VExpr
-buildClosure tvs lv vars arg_ty res_ty mk_body
+buildClosure :: [TyVar] -> [VVar] -> Type -> Type -> VM VExpr -> VM VExpr
+buildClosure tvs vars arg_ty res_ty mk_body
   = do
-      (env_ty, env, bind) <- buildEnv lv vars
+      (env_ty, env, bind) <- buildEnv vars
       env_bndr <- newLocalVVar FSLIT("env") env_ty
       arg_bndr <- newLocalVVar FSLIT("arg") arg_ty
 
       fn <- hoistPolyVExpr tvs
           $ do
+              lc    <- builtin liftingContext
               body  <- mk_body
               body' <- bind (vVar env_bndr)
-                            (vVarApps lv body (vars ++ [arg_bndr]))
+                            (vVarApps lc body (vars ++ [arg_bndr]))
               return (vLamsWithoutLC [env_bndr, arg_bndr] body')
 
       mkClosure arg_ty res_ty env_ty fn env
 
-buildEnv :: Var -> [VVar] -> VM (Type, VExpr, VExpr -> VExpr -> VM VExpr)
-buildEnv lv vvs
+buildEnv :: [VVar] -> VM (Type, VExpr, VExpr -> VExpr -> VM VExpr)
+buildEnv vvs
   = do
+      lc <- builtin liftingContext
       let (ty, venv, vbind) = mkVectEnv tys vs
-      (lenv, lbind) <- mkLiftEnv lv tys ls
+      (lenv, lbind) <- mkLiftEnv lc tys ls
       return (ty, (venv, lenv),
               \(venv,lenv) (vbody,lbody) ->
               do
@@ -318,28 +327,28 @@ mkVectEnv tys  vs  = (ty, mkCoreTup (map Var vs),
     ty = mkCoreTupTy tys
 
 mkLiftEnv :: Var -> [Type] -> [Var] -> VM (CoreExpr, CoreExpr -> CoreExpr -> VM CoreExpr)
-mkLiftEnv lv [ty] [v]
+mkLiftEnv lc [ty] [v]
   = return (Var v, \env body ->
                    do
                      len <- lengthPA (Var v)
                      return . Let (NonRec v env)
-                            $ Case len lv (exprType body) [(DEFAULT, [], body)])
+                            $ Case len lc (exprType body) [(DEFAULT, [], body)])
 
 -- NOTE: this transparently deals with empty environments
-mkLiftEnv lv tys vs
+mkLiftEnv lc tys vs
   = do
       (env_tc, env_tyargs) <- lookupPArrayFamInst vty
       let [env_con] = tyConDataCons env_tc
           
           env = Var (dataConWrapId env_con)
                 `mkTyApps`  env_tyargs
-                `mkVarApps` (lv : vs)
+                `mkVarApps` (lc : vs)
 
           bind env body = let scrut = unwrapFamInstScrut env_tc env_tyargs env
                           in
                           return $ Case scrut (mkWildId (exprType scrut))
                                         (exprType body)
-                                        [(DataAlt env_con, lv : bndrs, body)]
+                                        [(DataAlt env_con, lc : bndrs, body)]
       return (env, bind)
   where
     vty = mkCoreTupTy tys

diff --git a/compiler/vectorise/Vectorise.hs b/compiler/vectorise/Vectorise.hs
index 7a29a7b..e3f8008 100644 (file)
--- a/compiler/vectorise/Vectorise.hs
+++ b/compiler/vectorise/Vectorise.hs
@@ -111,10 +111,9 @@ vectTopBinder var
 vectTopRhs :: Var -> CoreExpr -> VM CoreExpr
 vectTopRhs var expr
   = do
-      lc <- newLocalVar FSLIT("lc") intPrimTy
       closedV . liftM vectorised
               . inBind var
-              $ vectPolyExpr lc (freeVars expr)
+              $ vectPolyExpr (freeVars expr)
 
 -- ----------------------------------------------------------------------------
 -- Bindings
@@ -150,19 +149,19 @@ vectBndrsIn vs p
 -- ----------------------------------------------------------------------------
 -- Expressions
 
-vectVar :: Var -> Var -> VM VExpr
-vectVar lc v
+vectVar :: Var -> VM VExpr
+vectVar v
   = do
       r <- lookupVar v
       case r of
         Local (vv,lv) -> return (Var vv, Var lv)
         Global vv     -> do
                            let vexpr = Var vv
-                           lexpr <- replicatePA (Var lc) vexpr
+                           lexpr <- liftPA vexpr
                            return (vexpr, lexpr)
 
-vectPolyVar :: Var -> Var -> [Type] -> VM VExpr
-vectPolyVar lc v tys
+vectPolyVar :: Var -> [Type] -> VM VExpr
+vectPolyVar v tys
   = do
       vtys <- mapM vectType tys
       r <- lookupVar v
@@ -171,79 +170,78 @@ vectPolyVar lc v tys
                                      (polyApply (Var lv) vtys)
         Global poly    -> do
                             vexpr <- polyApply (Var poly) vtys
-                            lexpr <- replicatePA (Var lc) vexpr
+                            lexpr <- liftPA vexpr
                             return (vexpr, lexpr)
 
-vectLiteral :: Var -> Literal -> VM VExpr
-vectLiteral lc lit
+vectLiteral :: Literal -> VM VExpr
+vectLiteral lit
   = do
-      lexpr <- replicatePA (Var lc) (Lit lit)
+      lexpr <- liftPA (Lit lit)
       return (Lit lit, lexpr)
 
-vectPolyExpr :: Var -> CoreExprWithFVs -> VM VExpr
-vectPolyExpr lc expr
+vectPolyExpr :: CoreExprWithFVs -> VM VExpr
+vectPolyExpr expr
   = polyAbstract tvs $ \abstract ->
-    -- FIXME: shadowing (tvs in lc)
     do
-      mono' <- vectExpr lc mono
+      mono' <- vectExpr mono
       return $ mapVect abstract mono'
   where
     (tvs, mono) = collectAnnTypeBinders expr  
                 
-vectExpr :: Var -> CoreExprWithFVs -> VM VExpr
-vectExpr lc (_, AnnType ty)
+vectExpr :: CoreExprWithFVs -> VM VExpr
+vectExpr (_, AnnType ty)
   = liftM vType (vectType ty)
 
-vectExpr lc (_, AnnVar v) = vectVar lc v
+vectExpr (_, AnnVar v) = vectVar v
 
-vectExpr lc (_, AnnLit lit) = vectLiteral lc lit
+vectExpr (_, AnnLit lit) = vectLiteral lit
 
-vectExpr lc (_, AnnNote note expr)
-  = liftM (vNote note) (vectExpr lc expr)
+vectExpr (_, AnnNote note expr)
+  = liftM (vNote note) (vectExpr expr)
 
-vectExpr lc e@(_, AnnApp _ arg)
+vectExpr e@(_, AnnApp _ arg)
   | isAnnTypeArg arg
-  = vectTyAppExpr lc fn tys
+  = vectTyAppExpr fn tys
   where
     (fn, tys) = collectAnnTypeArgs e
 
-vectExpr lc (_, AnnApp fn arg)
+vectExpr (_, AnnApp fn arg)
   = do
-      fn'  <- vectExpr lc fn
-      arg' <- vectExpr lc arg
+      fn'  <- vectExpr fn
+      arg' <- vectExpr arg
       mkClosureApp fn' arg'
 
-vectExpr lc (_, AnnCase expr bndr ty alts)
+vectExpr (_, AnnCase expr bndr ty alts)
   = panic "vectExpr: case"
 
-vectExpr lc (_, AnnLet (AnnNonRec bndr rhs) body)
+vectExpr (_, AnnLet (AnnNonRec bndr rhs) body)
   = do
-      vrhs <- localV . inBind bndr $ vectPolyExpr lc rhs
-      (vbndr, vbody) <- vectBndrIn bndr (vectExpr lc body)
+      vrhs <- localV . inBind bndr $ vectPolyExpr rhs
+      (vbndr, vbody) <- vectBndrIn bndr (vectExpr body)
       return $ vLet (vNonRec vbndr vrhs) vbody
 
-vectExpr lc (_, AnnLet (AnnRec bs) body)
+vectExpr (_, AnnLet (AnnRec bs) body)
   = do
       (vbndrs, (vrhss, vbody)) <- vectBndrsIn bndrs
                                 $ liftM2 (,)
                                   (zipWithM vect_rhs bndrs rhss)
-                                  (vectPolyExpr lc body)
+                                  (vectPolyExpr body)
       return $ vLet (vRec vbndrs vrhss) vbody
   where
     (bndrs, rhss) = unzip bs
 
     vect_rhs bndr rhs = localV
                       . inBind bndr
-                      $ vectExpr lc rhs
+                      $ vectExpr rhs
 
-vectExpr lc e@(fvs, AnnLam bndr _)
+vectExpr e@(fvs, AnnLam bndr _)
   | not (isId bndr) = pprPanic "vectExpr" (ppr $ deAnnotate e)
-  | otherwise = vectLam lc fvs bs body
+  | otherwise = vectLam fvs bs body
   where
     (bs,body) = collectAnnValBinders e
 
-vectLam :: Var -> VarSet -> [Var] -> CoreExprWithFVs -> VM VExpr
-vectLam lc fvs bs body
+vectLam :: VarSet -> [Var] -> CoreExprWithFVs -> VM VExpr
+vectLam fvs bs body
   = do
       tyvars <- localTyVars
       (vs, vvs) <- readLEnv $ \env ->
@@ -253,14 +251,15 @@ vectLam lc fvs bs body
       arg_tys <- mapM (vectType . idType) bs
       res_ty  <- vectType (exprType $ deAnnotate body)
 
-      buildClosures tyvars lc vvs arg_tys res_ty
+      buildClosures tyvars vvs arg_tys res_ty
         . hoistPolyVExpr tyvars
         $ do
+            lc <- builtin liftingContext
             (vbndrs, vbody) <- vectBndrsIn (vs ++ bs)
-                                           (vectExpr lc body)
+                                           (vectExpr body)
             return $ vLams lc vbndrs vbody
   
-vectTyAppExpr :: Var -> CoreExprWithFVs -> [Type] -> VM VExpr
-vectTyAppExpr lc (_, AnnVar v) tys = vectPolyVar lc v tys
-vectTyAppExpr lc e tys = pprPanic "vectTyAppExpr" (ppr $ deAnnotate e)
+vectTyAppExpr :: CoreExprWithFVs -> [Type] -> VM VExpr
+vectTyAppExpr (_, AnnVar v) tys = vectPolyVar v tys
+vectTyAppExpr e tys = pprPanic "vectTyAppExpr" (ppr $ deAnnotate e)