Special case dictionary abstraction and application during vectorisation
authorManuel M T Chakravarty <chak@cse.unsw.edu.au>
Thu, 17 Nov 2011 11:33:49 +0000 (22:33 +1100)
committerManuel M T Chakravarty <chak@cse.unsw.edu.au>
Thu, 17 Nov 2011 11:40:31 +0000 (22:40 +1100)
compiler/vectorise/Vectorise.hs
compiler/vectorise/Vectorise/Exp.hs
compiler/vectorise/Vectorise/Monad.hs
compiler/vectorise/Vectorise/Monad/Base.hs
compiler/vectorise/Vectorise/Type/Env.hs
compiler/vectorise/Vectorise/Type/Type.hs
compiler/vectorise/Vectorise/Utils.hs
compiler/vectorise/Vectorise/Utils/Closure.hs
compiler/vectorise/Vectorise/Utils/Hoisting.hs
compiler/vectorise/Vectorise/Utils/Poly.hs
compiler/vectorise/Vectorise/Var.hs

index 7d2415c..dc467f5 100644 (file)
@@ -86,10 +86,11 @@ vectModule guts@(ModGuts { mg_tcs        = tycons
       ; (_, fam_inst_env) <- readGEnv global_fam_inst_env
 
           -- Vectorise all the top level bindings and VECTORISE declarations on imported identifiers
+          -- NB: Need to vectorise the imported bindings first (local bindings may depend on them).
       ; let impBinds = [imp_id | Vect          imp_id _ <- vect_decls, isGlobalId imp_id] ++
                        [imp_id | VectInst True imp_id   <- vect_decls, isGlobalId imp_id]
-      ; binds_top <- mapM vectTopBind binds
       ; binds_imp <- mapM vectImpBind impBinds
+      ; binds_top <- mapM vectTopBind binds
 
       ; return $ guts { mg_tcs          = tycons ++ new_tycons
                         -- we produce no new classes or instances, only new class type constructors
@@ -301,7 +302,8 @@ vectTopBinder var inline expr
 --     => generate vectorised code according to the the "Note [Scalar dfuns]" below
 -- 
 -- (4) There is no vectorisation declaration for the variable
---     => perform automatic vectorisation of the RHS
+--     => perform automatic vectorisation of the RHS (the definition may or may not be a dfun;
+--        vectorisation proceeds differently depending on which it is)
 --
 -- Note [Scalar dfuns]
 -- ~~~~~~~~~~~~~~~~~~~
@@ -342,7 +344,8 @@ vectTopRhs recFs var expr
        ; vectDecl     <- lookupVectDecl var
        ; let isDFun = isDFunId var
 
-       ; traceVt ("vectTopRhs of " ++ show var ++ info globalScalar isDFun vectDecl) $ ppr expr
+       ; traceVt ("vectTopRhs of " ++ show var ++ info globalScalar isDFun vectDecl ++ ":") $ 
+           ppr expr
 
        ; rhs globalScalar isDFun vectDecl
        }
@@ -357,14 +360,18 @@ vectTopRhs recFs var expr
       = do { expr' <- vectScalarDFun var recFs
            ; return (DontInline, True, expr')
            }
-    rhs False         _isDFun Nothing                         -- Case (4)
-      = do { let fvs = freeVars expr
+    rhs False         False   Nothing                         -- Case (4) — not a dfun
+      = do { let exprFvs = freeVars expr
            ; (inline, isScalar, vexpr) 
                <- inBind var $
-                    vectPolyExpr (isStrongLoopBreaker $ idOccInfo var) recFs fvs
+                    vectPolyExpr (isStrongLoopBreaker $ idOccInfo var) recFs exprFvs
            ; return (inline, isScalar, vectorised vexpr)
            }
-    
+    rhs False         True    Nothing                         -- Case (4) — is a dfun
+      = do { expr' <- vectDictExpr expr
+           ; return  (DontInline, True, expr')
+           }
+
     info True  False _                          = " [VECTORISE SCALAR]"
     info True  True  _                          = " [VECTORISE SCALAR instance]"
     info False _     vectDecl | isJust vectDecl = " [VECTORISE]"
index bf6fe31..ca7b13f 100644 (file)
@@ -1,9 +1,12 @@
+{-# LANGUAGE TupleSections #-}
+
 -- |Vectorisation of expressions.
 
 module Vectorise.Exp
   (   -- * Vectorise polymorphic expressions with special cases for right-hand sides of particular 
       --   variable bindings
     vectPolyExpr
+  , vectDictExpr
   , vectScalarFun
   , vectScalarDFun
   ) 
@@ -42,50 +45,45 @@ import Outputable
 import FastString
 import Control.Monad
 import Control.Applicative
+import Data.Maybe
 import Data.List
 
 
--- | Vectorise a polymorphic expression.
+-- |Vectorise a polymorphic expression.
 --
-vectPolyExpr :: Bool            -- ^ When vectorising the RHS of a
-                                -- binding, whether that binding is a
-                                -- loop breaker.
+vectPolyExpr :: Bool    -- ^ When vectorising the RHS of a binding: is that binding a loop breaker?
              -> [Var]                     
              -> CoreExprWithFVs
              -> VM (Inline, Bool, VExpr)
 vectPolyExpr loop_breaker recFns (_, AnnTick tickish expr)
- = do (inline, isScalarFn, expr') <- vectPolyExpr loop_breaker recFns expr
-      return (inline, isScalarFn, vTick tickish expr')
+ = do { (inline, isScalarFn, expr') <- vectPolyExpr loop_breaker recFns expr
+      ; return (inline, isScalarFn, vTick tickish expr')
+      }
 vectPolyExpr loop_breaker recFns expr
- = do
-      arity <- polyArity tvs
-      polyAbstract tvs $ \args ->
-        do
-          (inline, isScalarFn, mono') <- vectFnExpr False loop_breaker recFns mono
-          return (addInlineArity inline arity, isScalarFn, 
-                  mapVect (mkLams $ tvs ++ args) mono')
+ = do { arity <- polyArity tvs
+      ; polyAbstract tvs $ \args -> do
+      { (inline, isScalarFn, mono') <- vectFnExpr False loop_breaker recFns mono
+      ; return (addInlineArity inline arity, isScalarFn, mapVect (mkLams $ tvs ++ args) mono')
+      } }
   where
     (tvs, mono) = collectAnnTypeBinders expr
 
-
 -- |Vectorise an expression.
 --
 vectExpr :: CoreExprWithFVs -> VM VExpr
-vectExpr (_, AnnType ty)
-  = liftM vType (vectType ty)
 
 vectExpr (_, AnnVar v) 
   = vectVar v
 
 vectExpr (_, AnnLit lit) 
-  = vectLiteral lit
+  = vectConst $ Lit lit
 
-vectExpr (_, AnnTick tickish expr)
-  = liftM (vTick tickish) (vectExpr expr)
+vectExpr e@(_, AnnLam bndr _)
+  | isId bndr = (\(_, _, ve) -> ve) <$> vectFnExpr True False [] e
 
--- SPECIAL CASE: Vectorise/lift 'patError @ ty err' by only vectorising/lifting the type 'ty';
---   its only purpose is to abort the program, but we need to adjust the type to keep CoreLint
---   happy.
+  -- SPECIAL CASE: Vectorise/lift 'patError @ ty err' by only vectorising/lifting the type 'ty';
+  --   its only purpose is to abort the program, but we need to adjust the type to keep CoreLint
+  --   happy.
 -- FIXME: can't be do this with a VECTORISE pragma on 'pAT_ERROR_ID' now?
 vectExpr (_, AnnApp (_, AnnApp (_, AnnVar v) (_, AnnType ty)) err)
   | v == pAT_ERROR_ID
@@ -95,12 +93,14 @@ vectExpr (_, AnnApp (_, AnnApp (_, AnnVar v) (_, AnnType ty)) err)
   where
     err' = deAnnotate err
 
+  -- type application (handle multiple consecutive type applications simultaneously to ensure the
+  -- PA dictionaries are put at the right places)
 vectExpr e@(_, AnnApp _ arg)
   | isAnnTypeArg arg
-  = vectTyAppExpr fn tys
-  where
-    (fn, tys) = collectAnnTypeArgs e
-
+  = vectPolyApp e
+    
+  -- 'Int', 'Float', or 'Double' literal
+  -- FIXME: this needs to be generalised
 vectExpr (_, AnnApp (_, AnnVar v) (_, AnnLit lit))
   | Just con <- isDataConId_maybe v
   , is_special_con con
@@ -111,25 +111,22 @@ vectExpr (_, AnnApp (_, AnnVar v) (_, AnnLit lit))
   where
     is_special_con con = con `elem` [intDataCon, floatDataCon, doubleDataCon]
 
-
--- TODO: Avoid using closure application for dictionaries.
--- vectExpr (_, AnnApp fn arg)
---  | if is application of dictionary 
---    just use regular app instead of closure app.
-
--- for lifted version. 
---      do liftPD (sub a dNumber)
---      lift the result of the selection, not sub and dNumber seprately. 
-
-vectExpr (_, AnnApp fn arg)
- = do
-      arg_ty' <- vectType arg_ty
-      res_ty' <- vectType res_ty
-
-      fn'     <- vectExpr fn
-      arg'    <- vectExpr arg
-
-      mkClosureApp arg_ty' res_ty' fn' arg'
+  -- value application (dictionary or user value)
+vectExpr e@(_, AnnApp fn arg)
+  | isPredTy arg_ty   -- dictionary application (whose result is not a dictionary)
+  = vectPolyApp e
+  | otherwise         -- user value
+  = do {   -- vectorise the types
+       ; varg_ty <- vectType arg_ty
+       ; vres_ty <- vectType res_ty
+
+           -- vectorise the function and argument expression
+       ; vfn  <- vectExpr fn
+       ; varg <- vectExpr arg
+
+           -- the vectorised function is a closure; apply it to the vectorised argument
+       ; mkClosureApp varg_ty vres_ty vfn varg
+       }
   where
     (arg_ty, res_ty) = splitFunTy . exprType $ deAnnotate fn
 
@@ -162,18 +159,19 @@ vectExpr (_, AnnLet (AnnRec bs) body)
                       . liftM (\(_,_,z)->z)
                       $ vectPolyExpr (isStrongLoopBreaker $ idOccInfo bndr) [] rhs
 
-vectExpr e@(_, AnnLam bndr _)
-  | isId bndr = liftM (\(_,_,z) ->z) $ vectFnExpr True False [] e
-{-
-onlyIfV (isEmptyVarSet fvs) (vectScalarLam bs $ deAnnotate body)
-                `orElseV` vectLam True fvs bs body
-  where
-    (bs,body) = collectAnnValBinders e
--}
+vectExpr (_, AnnTick tickish expr)
+  = liftM (vTick tickish) (vectExpr expr)
+
+vectExpr (_, AnnType ty)
+  = liftM vType (vectType ty)
 
 vectExpr e = cantVectorise "Can't vectorise expression (vectExpr)" (ppr $ deAnnotate e)
 
--- |Vectorise an expression with an outer lambda abstraction.
+-- |Vectorise an expression that *may* have an outer lambda abstraction.
+--
+-- We do not handle type variables at this point, as they will already have been stripped off by
+-- 'vectPolyExpr'.  We also only have to worry about one set of dictionary arguments as we (1) only
+-- deal with Haskell 2011 and (2) class selectors are vectorised elsewhere.
 --
 vectFnExpr :: Bool             -- ^ If we process the RHS of a binding, whether that binding should
                                --   be inlined
@@ -181,15 +179,138 @@ vectFnExpr :: Bool             -- ^ If we process the RHS of a binding, whether
            -> [Var]            -- ^ Names of function in same recursive binding group
            -> CoreExprWithFVs  -- ^ Expression to vectorise; must have an outer `AnnLam`
            -> VM (Inline, Bool, VExpr)
-vectFnExpr inline loop_breaker recFns expr@(_fvs, AnnLam bndr _)
-  | isId bndr = mark DontInline True (vectScalarFun False recFns (deAnnotate expr))
-                `orElseV` 
-                mark inlineMe False (vectLam inline loop_breaker expr)
-vectFnExpr _ _ _  e = mark DontInline False $ vectExpr e
+vectFnExpr inline loop_breaker recFns expr@(_fvs, AnnLam bndr body)
+      -- predicate abstraction: leave as a normal abstraction, but vectorise the predicate type
+  | isId bndr
+    && isPredTy (idType bndr)
+  = do { vBndr <- vectBndr bndr
+       ; (inline, isScalarFn, vbody) <- vectFnExpr inline loop_breaker recFns body
+       ; return (inline, isScalarFn, mapVect (mkLams [vectorised vBndr]) vbody)
+       }
+      -- non-predicate abstraction: vectorise (try to vectorise as a scalar computation)
+  | isId bndr
+  = mark DontInline True (vectScalarFun False recFns (deAnnotate expr))
+    `orElseV` 
+    mark inlineMe False (vectLam inline loop_breaker expr)
+vectFnExpr _ _ _  e 
+      -- not an abstraction: vectorise as a vanilla expression
+  = mark DontInline False $ vectExpr e
 
 mark :: Inline -> Bool -> VM a -> VM (Inline, Bool, a)
 mark b isScalarFn p = do { x <- p; return (b, isScalarFn, x) }
 
+-- |Vectorise type and dictionary applications.
+--
+-- These are always headed by a variable (as we don't support higher-rank polymorphism), but may
+-- involve two sets of type variables and dictionaries.  Consider,
+--
+-- > class C a where
+-- >   m :: D b => b -> a
+--
+-- The type of 'm' is 'm :: forall a. C a => forall b. D b => b -> a'.
+--
+vectPolyApp :: CoreExprWithFVs -> VM VExpr
+vectPolyApp e0
+  = case e4 of
+      (_, AnnVar var)
+        -> do {   -- get the vectorised form of the variable
+              ; vVar <- lookupVar var
+              ; traceVt "vectPolyApp of" (ppr var)
+
+                  -- vectorise type and dictionary arguments
+              ; vDictsOuter <- mapM vectDictExpr (map deAnnotate dictsOuter)
+              ; vDictsInner <- mapM vectDictExpr (map deAnnotate dictsInner)
+              ; vTysOuter   <- mapM vectType     tysOuter
+              ; vTysInner   <- mapM vectType     tysInner
+              
+              ; let reconstructOuter v = (`mkApps` vDictsOuter) <$> polyApply v vTysOuter
+
+              ; case vVar of
+                  Local (vv, lv)
+                    -> do { MASSERT( null dictsInner )    -- local vars cannot be class selectors
+                          ; traceVt "  LOCAL" (text "")
+                          ; (,) <$> reconstructOuter (Var vv) <*> reconstructOuter (Var lv)
+                          }
+                  Global vv
+                    | isDictComp var                      -- dictionary computation
+                    -> do {   -- in a dictionary computation, the innermost, non-empty set of
+                              -- arguments are non-vectorised arguments, where no 'PA'dictionaries
+                              -- are needed for the type variables
+                          ; ve <- if null dictsInner
+                                  then 
+                                    return $ Var vv `mkTyApps` vTysOuter `mkApps` vDictsOuter
+                                  else 
+                                    reconstructOuter 
+                                      (Var vv `mkTyApps` vTysInner `mkApps` vDictsInner)
+                          ; traceVt "  GLOBAL (dict):" (ppr ve)
+                          ; vectConst ve
+                          }
+                    | otherwise                           -- non-dictionary computation
+                    -> do { MASSERT( null dictsInner )
+                          ; ve <- reconstructOuter (Var vv)
+                          ; traceVt "  GLOBAL (non-dict):" (ppr ve)
+                          ; vectConst ve
+                          }
+              }
+      _ -> pprSorry "Cannot vectorise programs with higher-rank types:" (ppr . deAnnotate $ e0)
+  where
+    -- if there is only one set of variables or dictionaries, it will be the outer set
+    (e1, dictsOuter) = collectAnnDictArgs e0
+    (e2, tysOuter)   = collectAnnTypeArgs e1
+    (e3, dictsInner) = collectAnnDictArgs e2
+    (e4, tysInner)   = collectAnnTypeArgs e3
+    --
+    isDictComp var = (isJust . isClassOpId_maybe $ var) || isDFunId var
+    
+-- |Vectorise the body of a dfun.  
+--
+-- Dictionary computations are special for the following reasons.  The application of dictionary
+-- functions are always saturated, so there is no need to create closures.  Dictionary computations
+-- don't depend on array values, so they are always scalar computations whose result we can
+-- replicate (instead of executing them in parallel).
+--
+-- NB: To keep things simple, we are not rewriting any of the bindings introduced in a dictionary
+--     computation.  Consequently, the variable case needs to deal with cases where binders are
+--     in the vectoriser environments and where that is not the case.
+--
+vectDictExpr :: CoreExpr -> VM CoreExpr
+vectDictExpr (Var var)
+  = do { mb_scope <- lookupVar_maybe var
+       ; case mb_scope of
+           Nothing                -> return $ Var var   -- binder from within the dict. computation
+           Just (Local (vVar, _)) -> return $ Var vVar  -- local vectorised variable
+           Just (Global vVar)     -> return $ Var vVar  -- global vectorised variable
+       }
+vectDictExpr (Lit lit)
+  = pprPanic "Vectorise.Exp.vectDictExpr: literal in dictionary computation" (ppr lit)
+vectDictExpr (Lam bndr e)
+  = Lam bndr <$> vectDictExpr e
+vectDictExpr (App fn arg)
+  = App <$> vectDictExpr fn <*> vectDictExpr arg
+vectDictExpr (Case e bndr ty alts)
+  = Case <$> vectDictExpr e <*> pure bndr <*> vectType ty <*> mapM vectDictAlt alts
+  where
+    vectDictAlt (con, bs, e) = (,,) <$> vectDictAltCon con <*> pure bs <*> vectDictExpr e
+    --
+    vectDictAltCon (DataAlt datacon) = DataAlt <$> maybeV dataConErr (lookupDataCon datacon)
+      where
+        dataConErr = ptext (sLit "Cannot vectorise data constructor:") <+> ppr datacon
+    vectDictAltCon (LitAlt lit)      = return $ LitAlt lit
+    vectDictAltCon DEFAULT           = return DEFAULT
+vectDictExpr (Let bnd body)
+  = Let <$> vectDictBind bnd <*> vectDictExpr body
+  where
+    vectDictBind (NonRec bndr e) = NonRec bndr <$> vectDictExpr e
+    vectDictBind (Rec bnds)      = Rec <$> mapM (\(bndr, e) -> (bndr,) <$> vectDictExpr e) bnds
+vectDictExpr e@(Cast _e _coe)
+  = pprSorry "Vectorise.Exp.vectDictExpr: cast" (ppr e)
+vectDictExpr (Tick tickish e)
+  = Tick tickish <$> vectDictExpr e
+vectDictExpr (Type ty)
+  = Type <$> vectType ty
+vectDictExpr (Coercion coe)
+  = pprSorry "Vectorise.Exp.vectDictExpr: coercion" (ppr coe)
+
 -- |Vectorise an expression of functional type, where all arguments and the result are of scalar
 -- type (i.e., 'Int', 'Float', 'Double' etc.) and which does not contain any subcomputations that
 -- involve parallel arrays.  Such functionals do not requires the full blown vectorisation
@@ -398,53 +519,68 @@ unVectDict ty e
                                        Nothing  -> panic "Vectorise.Exp.unVectDict: no class"
     selIds                         = classAllSelIds cls
 
--- | Vectorise a lambda abstraction.
+-- |Vectorise an 'n'-ary lambda abstraction by building a set of 'n' explicit closures.
+--
+-- All non-dictionary free variables go into the closure's environment, whereas the dictionary
+-- variables are passed explicit (as conventional arguments) into the body during closure construction.
 --
 vectLam :: Bool             -- ^ When the RHS of a binding, whether that binding should be inlined.
         -> Bool             -- ^ Whether the binding is a loop breaker.
         -> CoreExprWithFVs  -- ^ Body of abstraction.
         -> VM VExpr
 vectLam inline loop_breaker expr@(fvs, AnnLam _ _)
- = do let (bs, body) = collectAnnValBinders expr
-      tyvars    <- localTyVars
-      (vs, vvs) <- readLEnv $ \env ->
-                   unzip [(var, vv) | var <- varSetElems fvs
-                                    , Just vv <- [lookupVarEnv (local_vars env) var]]
-
-      arg_tys   <- mapM (vectType . idType) bs
-      res_ty    <- vectType (exprType $ deAnnotate body)
-
-      buildClosures tyvars vvs arg_tys res_ty
-        . hoistPolyVExpr tyvars (maybe_inline (length vs + length bs))
-        $ do
-            lc              <- builtin liftingContext
-            (vbndrs, vbody) <- vectBndrsIn (vs ++ bs) (vectExpr body)
-
-            vbody' <- break_loop lc res_ty vbody
-            return $ vLams lc vbndrs vbody'
+ = do { let (bndrs, body) = collectAnnValBinders expr
+
+          -- grab the in-scope type variables
+      ; tyvars <- localTyVars
+
+          -- collect and vectorise all /local/ free variables
+      ; vfvs <- readLEnv $ \env ->
+                  [ (var, fromJust mb_vv) 
+                  | var <- varSetElems fvs
+                  , let mb_vv = lookupVarEnv (local_vars env) var
+                  , isJust mb_vv         -- its local == is in local var env
+                  ]
+          -- separate dictionary from non-dictionary variables in the free variable set
+      ; let (vvs_dict, vvs_nondict)     = partition (isPredTy . varType . fst) vfvs
+            (_fvs_dict, vfvs_dict)      = unzip vvs_dict
+            (fvs_nondict, vfvs_nondict) = unzip vvs_nondict
+
+          -- compute the type of the vectorised closure
+      ; arg_tys <- mapM (vectType . idType) bndrs
+      ; res_ty  <- vectType (exprType $ deAnnotate body)
+
+      ; let arity      = length fvs_nondict + length bndrs
+            vfvs_dict' = map vectorised vfvs_dict
+      ; buildClosures tyvars vfvs_dict' vfvs_nondict arg_tys res_ty
+        . hoistPolyVExpr tyvars vfvs_dict' (maybe_inline arity)
+        $ do {   -- generate the vectorised body of the lambda abstraction
+             ; lc              <- builtin liftingContext
+             ; (vbndrs, vbody) <- vectBndrsIn (fvs_nondict ++ bndrs) (vectExpr body)
+
+             ; vbody' <- break_loop lc res_ty vbody
+             ; return $ vLams lc vbndrs vbody'
+             }
+      }
   where
     maybe_inline n | inline    = Inline n
                    | otherwise = DontInline
 
+    -- If this is the body of a binding marked as a loop breaker, add a recursion termination test
+    -- to the /lifted/ version of the function body.  The termination tests checks if the lifting
+    -- context is empty.  If so, it returns an empty array of the (lifted) result type instead of
+    -- executing the function body.  This is the test from the last line (defining \mathcal{L}')
+    -- in Figure 6 of HtM.
     break_loop lc ty (ve, le)
       | loop_breaker
-      = do
-          empty <- emptyPD ty
-          lty <- mkPDataType ty
-          return (ve, mkWildCase (Var lc) intPrimTy lty
-                        [(DEFAULT, [], le),
-                         (LitAlt (mkMachInt 0), [], empty)])
-
+      = do { empty <- emptyPD ty
+           ; lty   <- mkPDataType ty
+           ; return (ve, mkWildCase (Var lc) intPrimTy lty
+                           [(DEFAULT, [], le),
+                            (LitAlt (mkMachInt 0), [], empty)])
+           }
       | otherwise = return (ve, le)
 vectLam _ _ _ = panic "vectLam"
-
-vectTyAppExpr :: CoreExprWithFVs -> [Type] -> VM VExpr
-vectTyAppExpr (_, AnnVar v) tys = vectPolyVar v tys
-vectTyAppExpr e tys = cantVectorise "Can't vectorise expression (vectTyExpr)"
-                        (ppr $ deAnnotate e `mkTyApps` tys)
-
 
 -- | Vectorise an algebraic case expression.
 --   We convert
index 6fbdb4e..b9a1fdf 100644 (file)
@@ -13,8 +13,7 @@ module Vectorise.Monad (
   
   -- * Variables
   lookupVar,
-  maybeCantVectoriseVarM,
-  dumpVar,
+  lookupVar_maybe,
   addGlobalScalar, 
 ) where
 
@@ -41,7 +40,6 @@ import Name
 import ErrUtils
 import Outputable
 
-import Control.Monad
 import System.IO
 
 
@@ -142,32 +140,31 @@ builtins f = VM $ \bi genv lenv -> return (Yes genv lenv (`f` bi))
 
 -- Var ------------------------------------------------------------------------
 
--- |Lookup the vectorised, and if local, also the lifted versions of a variable.
+-- |Lookup the vectorised, and if local, also the lifted version of a variable.
 --
 -- * If it's in the global environment we get the vectorised version.
 -- * If it's in the local environment we get both the vectorised and lifted version.
 --
 lookupVar :: Var -> VM (Scope Var (Var, Var))
 lookupVar v
= do r <- readLEnv $ \env -> lookupVarEnv (local_vars env) v
-      case r of
-        Just e  -> return (Local e)
-        Nothing -> liftM Global
-                . maybeCantVectoriseVarM v
-                . readGEnv $ \env -> lookupVarEnv (global_vars env) v
-
-maybeCantVectoriseVarM :: Monad m => Var -> m (Maybe Var) -> m Var
-maybeCantVectoriseVarM v p
- = do r <- p
-      case r of
-        Just x  -> return x
-        Nothing -> dumpVar v
 = do { mb_res <- lookupVar_maybe v
+       ; case mb_res of
+           Just x  -> return x
+           Nothing -> dumpVar v
+       }
+
+lookupVar_maybe :: Var -> VM (Maybe (Scope Var (Var, Var)))
+lookupVar_maybe v
+ = do { r <- readLEnv $ \env -> lookupVarEnv (local_vars env) v
+      ; case r of
+          Just e  -> return $ Just (Local e)
+          Nothing -> fmap Global <$> (readGEnv $ \env -> lookupVarEnv (global_vars env) v)
+      }
 
 dumpVar :: Var -> a
 dumpVar var
   | Just _    <- isClassOpId_maybe var
   = cantVectorise "ClassOpId not vectorised:" (ppr var)
-
   | otherwise
   = cantVectorise "Variable not vectorised:" (ppr var)
 
index 01fb6a5..b8a18c3 100644 (file)
@@ -37,6 +37,7 @@ import DynFlags
 import StaticFlags
 
 import Control.Monad
+import Control.Applicative
 import System.IO (stderr)
 
 
@@ -60,6 +61,10 @@ instance Monad VM where
                                          Yes genv' lenv' x -> runVM (f x) bi genv' lenv'
                                          No reason         -> return $ No reason
 
+instance Applicative VM where
+  pure  = return
+  (<*>) = ap
+  
 instance Functor VM where
   fmap = liftM
   
index e1efc96..87d0717 100644 (file)
@@ -308,7 +308,7 @@ vectDataConWorkers orig_tc vect_tc arr_tc
                 . inBind orig_worker
                 . polyAbstract tyvars $ \args ->
                   liftM (mkLams (tyvars ++ args) . vectorised)
-                $ buildClosures tyvars [] arg_tys res_ty mk_body
+                $ buildClosures tyvars [] [] arg_tys res_ty mk_body
 
           raw_worker <- mkVectId orig_worker (exprType body)
           let vect_worker = raw_worker `setIdUnfolding`
index cdd7bed..db724ad 100644 (file)
@@ -1,20 +1,22 @@
 -- Apply the vectorisation transformation to types. This is the \mathcal{L}_t scheme in HtM.
 
-module Vectorise.Type.Type (
-  vectTyCon,
-  vectAndLiftType,
-  vectType
-) where
+module Vectorise.Type.Type
+  ( vectTyCon
+  , vectAndLiftType
+  , vectType
+  ) 
+where
 
 import Vectorise.Utils
 import Vectorise.Monad
 import Vectorise.Builtins
-import TypeRep
+import TcType
 import Type
+import TypeRep
 import TyCon
 import Outputable
 import Control.Monad
-import Data.List
+import Control.Applicative
 import Data.Maybe
 
 -- | Vectorise a type constructor.
@@ -30,55 +32,53 @@ vectTyCon tc
 
 -- |Produce the vectorised and lifted versions of a type.
 --
+-- NB: Here we are limited to properly handle predicates at the toplevel only.  Anything embedded
+--     in what is called the 'body_ty' below will end up as an argument to the type family 'PData'.
+--
 vectAndLiftType :: Type -> VM (Type, Type)
 vectAndLiftType ty | Just ty' <- coreView ty = vectAndLiftType ty'
 vectAndLiftType ty
-  = do
-      mdicts   <- mapM paDictArgType (reverse tyvars)
-      let dicts = [dict | Just dict <- mdicts]
-      vmono_ty <- vectType mono_ty
-      lmono_ty <- mkPDataType vmono_ty
-      return (abstractType tyvars dicts vmono_ty,
-              abstractType tyvars dicts lmono_ty)
+  = do { padicts  <- liftM catMaybes $ mapM paDictArgType tyvars
+       ; vmono_ty <- vectType mono_ty
+       ; lmono_ty <- mkPDataType vmono_ty
+       ; return (abstractType tyvars (padicts ++ theta) vmono_ty,
+                 abstractType tyvars (padicts ++ theta) lmono_ty)
+       }
   where
-    (tyvars, mono_ty) = splitForAllTys ty
+    (tyvars, phiTy)  = splitForAllTys ty
+    (theta, mono_ty) = tcSplitPhiTy phiTy 
 
 -- |Vectorise a type.
 --
+-- For each quantified var we need to add a PA dictionary out the front of the type.
+-- So          forall a.         C  a => a -> a   
+-- turns into  forall a. PA a => Cv a => a :-> a
+--
 vectType :: Type -> VM Type
 vectType ty
   | Just ty'  <- coreView ty
   = vectType ty'
-  
 vectType (TyVarTy tv)      = return $ TyVarTy tv
-vectType (AppTy ty1 ty2)   = liftM2 AppTy    (vectType ty1) (vectType ty2)
-vectType (TyConApp tc tys) = liftM2 TyConApp (vectTyCon tc) (mapM vectType tys)
-vectType (FunTy ty1 ty2)   = liftM2 TyConApp (builtin closureTyCon) (mapM vectType [ty1,ty2])
-
--- For each quantified var we need to add a PA dictionary out the front of the type.
--- So          forall a. C  a => a -> a   
--- turns into  forall a. Cv a => PA a => a :-> a
+vectType (AppTy ty1 ty2)   = AppTy <$> vectType ty1 <*> vectType ty2
+vectType (TyConApp tc tys) = TyConApp <$> vectTyCon tc <*> mapM vectType tys
+vectType (FunTy ty1 ty2)   
+  | isPredTy ty1
+  = FunTy <$> vectType ty1 <*> vectType ty2   -- don't build a closure for dictionary abstraction
+  | otherwise
+  = TyConApp <$> builtin closureTyCon <*> mapM vectType [ty1, ty2]
 vectType ty@(ForAllTy _ _)
- = do
-      -- split the type into the quantified vars, its dictionaries and the body.
-      let (tyvars, tyBody)   = splitForAllTys ty
-      let (tyArgs, tyResult) = splitFunTys    tyBody
-
-      let (tyArgs_dict, tyArgs_regular) 
-                  = partition isDictTy tyArgs
-
-      -- vectorise the body.
-      let tyBody' = mkFunTys tyArgs_regular tyResult
-      tyBody''    <- vectType tyBody'
+ = do {   -- strip off consecutive foralls
+      ; let (tyvars, tyBody) = splitForAllTys ty
 
-      -- vectorise the dictionary parameters.
-      dictsVect   <- mapM vectType tyArgs_dict
+          -- vectorise the body
+      ; vtyBody <- vectType tyBody
 
-      -- make a PA dictionary for each of the type variables.
-      dictsPA     <- liftM catMaybes $ mapM paDictArgType tyvars
+          -- make a PA dictionary for each of the type variables
+      ; dictsPA <- liftM catMaybes $ mapM paDictArgType tyvars
 
-      -- pack it all back together.
-      return $ abstractType tyvars (dictsPA ++ dictsVect) tyBody''
+          -- add the PA dictionaries after the foralls
+      ; return $ abstractType tyvars dictsPA vtyBody
+      }
 
 -- |Add quantified vars and dictionary parameters to the front of a type.
 --
index 255a6c5..c5f1cb7 100644 (file)
@@ -7,6 +7,7 @@ module Vectorise.Utils (
 
   -- * Annotated Exprs
   collectAnnTypeArgs,
+  collectAnnDictArgs,
   collectAnnTypeBinders,
   collectAnnValBinders,
   isAnnTypeArg,
@@ -31,6 +32,7 @@ import Vectorise.Monad
 import Vectorise.Builtins
 import CoreSyn
 import CoreUtils
+import Id
 import Type
 import Control.Monad
 
@@ -43,17 +45,28 @@ collectAnnTypeArgs expr = go expr []
     go (_, AnnApp f (_, AnnType ty)) tys = go f (ty : tys)
     go e                             tys = (e, tys)
 
+collectAnnDictArgs :: AnnExpr Var ann -> (AnnExpr Var ann, [AnnExpr Var ann])
+collectAnnDictArgs expr = go expr []
+  where
+    go e@(_, AnnApp f arg) dicts 
+      | isPredTy . exprType . deAnnotate $ arg = go f (arg : dicts)
+      | otherwise                              = (e, dicts)
+    go e                        dicts          = (e, dicts)
+
 collectAnnTypeBinders :: AnnExpr Var ann -> ([Var], AnnExpr Var ann)
 collectAnnTypeBinders expr = go [] expr
   where
-    go bs (_, AnnLam b e) | isTyVar b = go (b:bs) e
+    go bs (_, AnnLam b e) | isTyVar b = go (b : bs) e
     go bs e                           = (reverse bs, e)
 
+-- |Collect all consecutive value binders that are not dictionaries.
+--
 collectAnnValBinders :: AnnExpr Var ann -> ([Var], AnnExpr Var ann)
 collectAnnValBinders expr = go [] expr
   where
-    go bs (_, AnnLam b e) | isId b = go (b:bs) e
-    go bs e                        = (reverse bs, e)
+    go bs (_, AnnLam b e) | isId b 
+                          && (not . isPredTy . idType $ b) = go (b : bs) e
+    go bs e                                                = (reverse bs, e)
 
 isAnnTypeArg :: AnnExpr b ann -> Bool
 isAnnTypeArg (_, AnnType _) = True
index 1f99ee5..0a918f8 100644 (file)
@@ -1,12 +1,11 @@
 -- |Utils concerning closure construction and application.
 
-module Vectorise.Utils.Closure (
-  mkClosure,
-  mkClosureApp,
-  buildClosure,
-  buildClosures,
-  buildEnv
-) where
+module Vectorise.Utils.Closure
+  ( mkClosure
+  , mkClosureApp
+  , buildClosures
+  )
+where
 
 import Vectorise.Builtins
 import Vectorise.Vect
@@ -56,51 +55,72 @@ mkClosureApp arg_ty res_ty (vclo, lclo) (varg, larg)
       return (Var vapply `mkTyApps` [arg_ty, res_ty] `mkApps` [vclo, varg],
               Var lapply `mkTyApps` [arg_ty, res_ty] `mkApps` [Var lc, lclo, larg])
 
-buildClosures :: [TyVar]
-              -> [VVar]
-              -> [Type]   -- ^ Type of the arguments.
-              -> Type     -- ^ Type of result.
+-- |Build a set of 'n' closures corresponding to an 'n'-ary vectorised function.  The length of
+-- the list of types of arguments determines the arity.
+--
+-- In addition to a set of type variables, a set of value variables is passed during closure
+-- /construction/.  In contrast, the closure environment and the arguments are passed during closure
+-- application.
+--
+buildClosures :: [TyVar]    -- ^ Type variables passed during closure construction.
+              -> [Var]      -- ^ Variables passed during closure construction.
+              -> [VVar]     -- ^ Variables in the environment.
+              -> [Type]     -- ^ Type of the arguments.
+              -> Type       -- ^ Type of result.
               -> VM VExpr
               -> VM VExpr
-buildClosures _   _    [] _ mk_body
+buildClosures _tvs _vars _env [] _res_ty mk_body
  = mk_body
-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 vars arg_ty res_ty'
-        . hoistPolyVExpr tvs (Inline (length vars + 1))
-        $ do
-            lc     <- builtin liftingContext
-            clo    <- buildClosures tvs (vars ++ [arg]) arg_tys res_ty mk_body
-            return $ vLams lc (vars ++ [arg]) clo
-
+buildClosures tvs vars env [arg_ty] res_ty mk_body
+ =  buildClosure tvs vars env arg_ty res_ty mk_body
+buildClosures tvs vars env (arg_ty : arg_tys) res_ty mk_body
+ = do { res_ty' <- mkClosureTypes arg_tys res_ty
+      ; arg     <- newLocalVVar (fsLit "x") arg_ty
+      ; buildClosure tvs vars env arg_ty res_ty'
+          . hoistPolyVExpr tvs vars (Inline (length env + 1))
+          $ do { lc     <- builtin liftingContext
+               ; clo    <- buildClosures tvs vars (env ++ [arg]) arg_tys res_ty mk_body
+               ; return $ vLams lc (env ++ [arg]) clo
+               }
+      }
+
+-- Build a closure taking one extra argument during closure application.
+--
 -- (clo <x1,...,xn> <f,f^>, aclo (Arr lc xs1 ... xsn) <f,f^>)
 --   where
 --     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] -> [VVar] -> Type -> Type -> VM VExpr -> VM VExpr
-buildClosure tvs vars arg_ty res_ty mk_body
-  = do
-      (env_ty, env, bind) <- buildEnv vars
-      env_bndr <- newLocalVVar (fsLit "env") env_ty
-      arg_bndr <- newLocalVVar (fsLit "arg") arg_ty
-
-      fn <- hoistPolyVExpr tvs (Inline 2)
-          $ do
-              lc     <- builtin liftingContext
-              body   <- mk_body
-              return .  vLams lc [env_bndr, arg_bndr]
-                     $  bind (vVar env_bndr)
-                             (vVarApps lc body (vars ++ [arg_bndr]))
-
-      mkClosure arg_ty res_ty env_ty fn env
-
-
--- Environments ---------------------------------------------------------------
-
+-- In addition to a set of type variables, a set of value variables is passed during closure
+-- /construction/.  In contrast, the closure environment and the closure argument are passed during
+-- closure application.
+--
+buildClosure :: [TyVar]         -- ^Type variables passed during closure construction.
+             -> [Var]           -- ^Variables passed during closure construction.
+             -> [VVar]          -- ^Variables in the environment.
+             -> Type            -- ^Type of the closure argument.
+             -> Type            -- ^Type of the result.
+             -> VM VExpr 
+             -> VM VExpr
+buildClosure tvs vars vvars arg_ty res_ty mk_body
+  = do { (env_ty, env, bind) <- buildEnv vvars
+       ; env_bndr <- newLocalVVar (fsLit "env") env_ty
+       ; arg_bndr <- newLocalVVar (fsLit "arg") arg_ty
+
+           -- generate the closure function as a hoisted binding
+       ; fn <- hoistPolyVExpr tvs vars (Inline 2) $
+                 do { lc     <- builtin liftingContext
+                    ; body   <- mk_body
+                    ; return . vLams lc [env_bndr, arg_bndr]
+                             $ bind (vVar env_bndr)
+                                    (vVarApps lc body (vvars ++ [arg_bndr]))
+                    }
+
+       ; mkClosure arg_ty res_ty env_ty fn env
+       }
+
+-- Build the environment for a single closure.
+--
 buildEnv :: [VVar] -> VM (Type, VExpr, VExpr -> VExpr -> VExpr)
 buildEnv [] 
  = do
index 3625508..7275bb2 100644 (file)
@@ -1,23 +1,16 @@
-
-{-# OPTIONS -fno-warn-tabs #-}
--- The above warning supression flag is a temporary kludge.
--- While working on this module you are encouraged to remove it and
--- detab the module (please do the detabbing in a separate patch). See
---     http://hackage.haskell.org/trac/ghc/wiki/Commentary/CodingStyle#TabsvsSpaces
--- for details
-
-module Vectorise.Utils.Hoisting (
-       Inline(..),
-       addInlineArity,
-       inlineMe,
-       
-       hoistBinding,
-       hoistExpr,
-       hoistVExpr,
-       hoistPolyVExpr,
-       takeHoisted
-)
+module Vectorise.Utils.Hoisting
+  ( Inline(..)
+  , addInlineArity
+  , inlineMe
+  
+  , hoistBinding
+  , hoistExpr
+  , hoistVExpr
+  , hoistPolyVExpr
+  , takeHoisted
+  )
 where
+
 import Vectorise.Monad
 import Vectorise.Env
 import Vectorise.Vect
@@ -28,33 +21,38 @@ import CoreUtils
 import CoreUnfold
 import Type
 import Id
-import BasicTypes( Arity )
+import BasicTypes  (Arity)
 import FastString
 import Control.Monad
+import Control.Applicative
 
 
 -- Inline ---------------------------------------------------------------------
--- | Records whether we should inline a particular binding.
+
+-- |Records whether we should inline a particular binding.
+--
 data Inline 
         = Inline Arity
         | DontInline
 
--- | Add to the arity contained within an `Inline`, if any.
+-- |Add to the arity contained within an `Inline`, if any.
+--
 addInlineArity :: Inline -> Int -> Inline
 addInlineArity (Inline m) n = Inline (m+n)
 addInlineArity DontInline _ = DontInline
 
--- | Says to always inline a binding.
+-- |Says to always inline a binding.
+--
 inlineMe :: Inline
 inlineMe = Inline 0
 
 
--- Hoising --------------------------------------------------------------------
+-- Hoisting --------------------------------------------------------------------
+
 hoistBinding :: Var -> CoreExpr -> VM ()
 hoistBinding v e = updGEnv $ \env ->
   env { global_bindings = (v,e) : global_bindings env }
 
-
 hoistExpr :: FastString -> CoreExpr -> Inline -> VM Var
 hoistExpr fs expr inl
   = do
@@ -67,7 +65,6 @@ hoistExpr fs expr inl
                                       mkInlineUnfolding (Just arity) expr
                       DontInline   -> var
 
-
 hoistVExpr :: VExpr -> Inline -> VM VVar
 hoistVExpr (ve, le) inl
   = do
@@ -76,16 +73,22 @@ hoistVExpr (ve, le) inl
       lv <- hoistExpr ('l' `consFS` fs) le (addInlineArity inl 1)
       return (vv, lv)
 
-
-hoistPolyVExpr :: [TyVar] -> Inline -> VM VExpr -> VM VExpr
-hoistPolyVExpr tvs inline p
-  = do
-      inline' <- liftM (addInlineArity inline) (polyArity tvs)
-      expr <- closedV . polyAbstract tvs $ \args ->
-              liftM (mapVect (mkLams $ tvs ++ args)) p
-      fn   <- hoistVExpr expr inline'
-      polyVApply (vVar fn) (mkTyVarTys tvs)
-
+-- |Hoist a polymorphic vectorised expression into a new top-level binding (representing a closure
+-- function).
+--
+-- The hoisted expression is parameterised by (1) a set of type variables and (2) a set of value
+-- variables that are passed as conventional type and value arguments.  The latter is implicitly
+-- extended by the set of 'PA' dictionaries required for the type variables.
+--
+hoistPolyVExpr :: [TyVar] -> [Var] -> Inline -> VM VExpr -> VM VExpr
+hoistPolyVExpr tvs vars inline p
+  = do { inline' <- addInlineArity inline . (+ length vars) <$> polyArity tvs
+       ; expr <- closedV . polyAbstract tvs $ \args ->
+                   mapVect (mkLams $ tvs ++ args ++ vars) <$> p
+       ; fn   <- hoistVExpr expr inline'
+       ; let varArgs = varsToCoreExprs vars
+       ; mapVect (\e -> e `mkApps` varArgs) <$> polyVApply (vVar fn) (mkTyVarTys tvs)
+       }
 
 takeHoisted :: VM [(Var, CoreExpr)]
 takeHoisted
index f33fef3..e943313 100644 (file)
@@ -1,18 +1,12 @@
 -- |Auxiliary functions to vectorise type abstractions.
 
-{-# OPTIONS -fno-warn-tabs #-}
--- The above warning supression flag is a temporary kludge.
--- While working on this module you are encouraged to remove it and
--- detab the module (please do the detabbing in a separate patch). See
---     http://hackage.haskell.org/trac/ghc/wiki/Commentary/CodingStyle#TabsvsSpaces
--- for details
-
-module Vectorise.Utils.Poly (
-  polyAbstract, 
-  polyApply,
-  polyVApply,
-  polyArity
-) where
+module Vectorise.Utils.Poly
+  ( polyAbstract
+  , polyApply
+  , polyVApply
+  , polyArity
+  ) 
+where
 
 import Vectorise.Vect
 import Vectorise.Monad
@@ -60,7 +54,7 @@ polyArity tvs
        ; return $ length [() | Just _ <- tys]
        }
 
--- |Apply a variable to its type arguments as well as 'PA' dictionaries for these type arguments.
+-- |Apply a expression to its type arguments as well as 'PA' dictionaries for these type arguments.
 --
 polyApply :: CoreExpr -> [Type] -> VM CoreExpr
 polyApply expr tys
index de04f4d..a299731 100644 (file)
@@ -1,22 +1,17 @@
+{-# LANGUAGE TupleSections #-}
 
-{-# OPTIONS -fno-warn-tabs #-}
--- The above warning supression flag is a temporary kludge.
--- While working on this module you are encouraged to remove it and
--- detab the module (please do the detabbing in a separate patch). See
---     http://hackage.haskell.org/trac/ghc/wiki/Commentary/CodingStyle#TabsvsSpaces
--- for details
-
--- | Vectorise variables and literals.
-module Vectorise.Var (
-       vectBndr,
-       vectBndrNew,
-       vectBndrIn,
-       vectBndrNewIn,
-       vectBndrsIn,
-       vectVar,
-       vectPolyVar,
-       vectLiteral
-) where
+-- |Vectorise variables and literals.
+
+module Vectorise.Var 
+  ( vectBndr
+  , vectBndrNew
+  , vectBndrIn
+  , vectBndrNewIn
+  , vectBndrsIn
+  , vectVar
+  , vectConst
+  )
+where
 
 import Vectorise.Utils
 import Vectorise.Monad
@@ -26,15 +21,15 @@ import Vectorise.Type.Type
 import CoreSyn
 import Type
 import VarEnv
-import Literal
 import Id
 import FastString
-import Control.Monad
+import Control.Applicative
 
 
 -- Binders ----------------------------------------------------------------------------------------
 
--- | Vectorise a binder variable, along with its attached type.
+-- |Vectorise a binder variable, along with its attached type.
+--
 vectBndr :: Var -> VM VVar
 vectBndr v
  = do (vty, lty) <- vectAndLiftType (idType v)
@@ -47,8 +42,8 @@ vectBndr v
   where
     mapTo vv lv env = env { local_vars = extendVarEnv (local_vars env) v (vv, lv) }
 
--- | Vectorise a binder variable, along with its attached type, 
---   but give the result a new name.
+-- |Vectorise a binder variable, along with its attached type, but give the result a new name.
+--
 vectBndrNew :: Var -> FastString -> VM VVar
 vectBndrNew v fs
  = do vty <- vectType (idType v)
@@ -58,7 +53,8 @@ vectBndrNew v fs
   where
     upd vv env = env { local_vars = extendVarEnv (local_vars env) v vv }
 
--- | Vectorise a binder then run a computation with that binder in scope.
+-- |Vectorise a binder then run a computation with that binder in scope.
+--
 vectBndrIn :: Var -> VM a -> VM (VVar, a)
 vectBndrIn v p
  = localV
@@ -66,7 +62,8 @@ vectBndrIn v p
       x <- p
       return (vv, x)
 
--- | Vectorise a binder, give it a new name, then run a computation with that binder in scope.
+-- |Vectorise a binder, give it a new name, then run a computation with that binder in scope.
+--
 vectBndrNewIn :: Var -> FastString -> VM a -> VM (VVar, a)
 vectBndrNewIn v fs p
  = localV
@@ -74,60 +71,33 @@ vectBndrNewIn v fs p
       x  <- p
       return (vv, x)
 
--- | Vectorise some binders, then run a computation with them in scope.
+-- |Vectorise some binders, then run a computation with them in scope.
+--
 vectBndrsIn :: [Var] -> VM a -> VM ([VVar], a)
 vectBndrsIn vs p
  = localV
  $ do vvs <- mapM vectBndr vs
-      x          <- p
+      x   <- p
       return (vvs, x)
 
 
 -- Variables --------------------------------------------------------------------------------------
 
--- | Vectorise a variable, producing the vectorised and lifted versions.
+-- |Vectorise a variable, producing the vectorised and lifted versions.
+--
 vectVar :: Var -> VM VExpr
-vectVar v
- = do 
-      -- lookup the variable from the environment.
-      r        <- lookupVar v
-
-      case r of
-        -- If it's been locally bound then we'll already have both versions available.
-        Local (vv,lv) 
-         -> return (Var vv, Var lv)
-
-        -- To create the lifted version of a global variable we replicate it
-       -- using the integer context in the VM state for the number of elements.
-        Global vv     
-         -> do let vexpr = Var vv
-               lexpr <- liftPD vexpr
-               return (vexpr, lexpr)
-
--- | Like `vectVar` but also add type applications to the variables.
--- FIXME: 'vectVar' is really just a special case, which 'vectPolyVar' should handle fine as well —
---        MERGE the two functions!
-vectPolyVar :: Var -> [Type] -> VM VExpr
-vectPolyVar v tys
- = do vtys     <- mapM vectType tys
-      r                <- lookupVar v
-      case r of
-        Local (vv, lv) 
-         -> liftM2 (,) (polyApply (Var vv) vtys)
-                       (polyApply (Var lv) vtys)
-
-        Global poly    
-         -> do vexpr <- polyApply (Var poly) vtys
-               lexpr <- liftPD vexpr
-               return (vexpr, lexpr)
-
-
--- Literals ---------------------------------------------------------------------------------------
-
--- | Lifted literals are created by replicating them
---   We use the the integer context in the `VM` state for the number
---   of elements in the output array.
-vectLiteral :: Literal -> VM VExpr
-vectLiteral lit
- = do lexpr <- liftPD (Lit lit)
-      return (Lit lit, lexpr)
+vectVar var
+  = do { vVar <- lookupVar var
+       ; case vVar of
+           Local (vv, lv) -> return (Var vv, Var lv) -- local variables have a vect & lifted version
+           Global vv      -> vectConst (Var vv)      -- global variables get replicated
+       }
+
+
+-- Constants --------------------------------------------------------------------------------------
+
+-- |Constants are lifted by replication along the integer context in the `VM` state for the number
+-- of elements in the result array.
+--
+vectConst :: CoreExpr -> VM VExpr
+vectConst c = (c,) <$> liftPD c