vectorise/Vectorise/Exp.hs

    {-# LANGUAGE CPP, TupleSections #-}

    -- |Vectorisation of expressions.

    module Vectorise.Exp
      (   -- * Vectorise right-hand sides of toplevel bindings
        vectTopExpr
      , vectTopExprs
      , vectScalarFun
      , vectScalarDFun
      )
    where

    #include "HsVersions.h"

    import Vectorise.Type.Type
    import Vectorise.Var
    import Vectorise.Convert
    import Vectorise.Vect
    import Vectorise.Env
    import Vectorise.Monad
    import Vectorise.Builtins
    import Vectorise.Utils

    import CoreUtils
    import MkCore
    import CoreSyn
    import CoreFVs
    import Class
    import DataCon
    import TyCon
    import TcType
    import Type
    import TyCoRep
    import Var
    import VarEnv
    import VarSet
    import NameSet
    import Id
    import BasicTypes( isStrongLoopBreaker )
    import Literal
    import TysPrim
    import Outputable
    import FastString
    import DynFlags
    import Util

    import Control.Monad
    import Data.Maybe
    import Data.List

    -- Main entry point to vectorise expressions -----------------------------------

    -- |Vectorise a polymorphic expression that forms a *non-recursive* binding.
    --
    -- Return 'Nothing' if the expression is scalar; otherwise, the first component of the result
    -- (which is of type 'Bool') indicates whether the expression is parallel (i.e., whether it is
    -- tagged as 'VIParr').
    --
    -- We have got the non-recursive case as a special case as it doesn't require to compute
    -- vectorisation information twice.
    --
    vectTopExpr :: Var -> CoreExpr -> VM (Maybe (Bool, Inline, CoreExpr))
    vectTopExpr var expr
      = do
        { exprVI <- encapsulateScalars <=< vectAvoidInfo emptyVarSet . freeVars $ expr
        ; if isVIEncaps exprVI
          then
            return Nothing
          else do
          { vExpr <- closedV $
                       inBind var $
                         vectAnnPolyExpr False exprVI
          ; inline <- computeInline exprVI
          ; return $ Just (isVIParr exprVI, inline, vectorised vExpr)
          }
        }

    -- Compute the inlining hint for the right-hand side of a top-level binding.
    --
    computeInline :: CoreExprWithVectInfo -> VM Inline
    computeInline ((_, VIDict), _)     = return $ DontInline
    computeInline (_, AnnTick _ expr)  = computeInline expr
    computeInline expr@(_, AnnLam _ _) = Inline <$> polyArity tvs
      where
        (tvs, _) = collectAnnTypeBinders expr
    computeInline _expr                = return $ DontInline

    -- |Vectorise a recursive group of top-level polymorphic expressions.
    --
    -- Return 'Nothing' if the expression group is scalar; otherwise, the first component of the result
    -- (which is of type 'Bool') indicates whether the expressions are parallel (i.e., whether they are
    -- tagged as 'VIParr').
    --
    vectTopExprs :: [(Var, CoreExpr)] -> VM (Maybe (Bool, [(Inline, CoreExpr)]))
    vectTopExprs binds
      = do
        { exprVIs <- mapM (vectAvoidAndEncapsulate emptyVarSet) exprs
        ; if all isVIEncaps exprVIs
            -- if all bindings are scalar => don't vectorise this group of bindings
          then return Nothing
          else do
          {   -- non-scalar bindings need to be vectorised
          ; let areVIParr = any isVIParr exprVIs
          ; revised_exprVIs <- if not areVIParr
                                 -- if no binding is parallel => 'exprVIs' is ready for vectorisation
                               then return exprVIs
                                 -- if any binding is parallel => recompute the vectorisation info
                               else mapM (vectAvoidAndEncapsulate (mkVarSet vars)) exprs

          ; vExprs <- zipWithM vect vars revised_exprVIs
          ; return $ Just (areVIParr, vExprs)
          }
        }
      where
        (vars, exprs) = unzip binds

        vectAvoidAndEncapsulate pvs = encapsulateScalars <=< vectAvoidInfo pvs . freeVars

        vect var exprVI
          = do
            { vExpr  <- closedV $
                          inBind var $
                            vectAnnPolyExpr (isStrongLoopBreaker $ idOccInfo var) exprVI
            ; inline <- computeInline exprVI
            ; return (inline, vectorised vExpr)
            }

    -- |Vectorise a polymorphic expression annotated with vectorisation information.
    --
    -- The special case of dictionary functions is currently handled separately. (Would be neater to
    -- integrate them, though!)
    --
    vectAnnPolyExpr :: Bool -> CoreExprWithVectInfo -> VM VExpr
    vectAnnPolyExpr loop_breaker (_, AnnTick tickish expr)
        -- traverse through ticks
      = vTick tickish <$> vectAnnPolyExpr loop_breaker expr
    vectAnnPolyExpr loop_breaker expr
      | isVIDict expr
        -- special case the right-hand side of dictionary functions
      = (, undefined) <$> vectDictExpr (deAnnotate expr)
      | otherwise
        -- collect and vectorise type abstractions; then, descent into the body
      = polyAbstract tvs $ \args ->
          mapVect (mkLams $ tvs ++ args) <$> vectFnExpr False loop_breaker mono
      where
        (tvs, mono) = collectAnnTypeBinders expr

    -- Encapsulate every purely sequential subexpression of a (potentially) parallel expression into a
    -- lambda abstraction over all its free variables followed by the corresponding application to those
    -- variables.  We can, then, avoid the vectorisation of the ensapsulated subexpressions.
    --
    -- Preconditions:
    --
    -- * All free variables and the result type must be /simple/ types.
    -- * The expression is sufficiently complex (to warrant special treatment).  For now, that is
    --   every expression that is not constant and contains at least one operation.
    --
    --
    -- The user has an option to choose between aggressive and minimal vectorisation avoidance. With
    -- minimal vectorisation avoidance, we only encapsulate individual scalar operations. With
    -- aggressive vectorisation avoidance, we encapsulate subexpression that are as big as possible.
    --
    encapsulateScalars :: CoreExprWithVectInfo -> VM CoreExprWithVectInfo
    encapsulateScalars ce@(_, AnnType _ty)
      = return ce
    encapsulateScalars ce@((_, VISimple), AnnVar _v)
          -- NB: diverts from the paper: encapsulate scalar variables (including functions)
      = liftSimpleAndCase ce
    encapsulateScalars ce@(_, AnnVar _v)
      = return ce
    encapsulateScalars ce@(_, AnnLit _)
      = return ce
    encapsulateScalars ((fvs, vi), AnnTick tck expr)
      = do
        { encExpr <- encapsulateScalars expr
        ; return ((fvs, vi), AnnTick tck encExpr)
        }
    encapsulateScalars ce@((fvs, vi), AnnLam bndr expr)
      = do
        { vectAvoid <- isVectAvoidanceAggressive
        ; varsS     <- allScalarVarTypeSet fvs
            -- NB: diverts from the paper: we need to check the scalarness of bound variables as well,
            --     as 'vectScalarFun' will handle them just the same as those introduced for the 'fvs'
            --     by encapsulation.
        ; bndrsS    <- allScalarVarType bndrs
        ; case (vi, vectAvoid && varsS && bndrsS) of
            (VISimple, True) -> liftSimpleAndCase ce
            _                -> do
                                { encExpr <- encapsulateScalars expr
                                ; return ((fvs, vi), AnnLam bndr encExpr)
                                }
        }
      where
        (bndrs, _) = collectAnnBndrs ce
    encapsulateScalars ce@((fvs, vi), AnnApp ce1 ce2)
      = do
        { vectAvoid <- isVectAvoidanceAggressive
        ; varsS     <- allScalarVarTypeSet fvs
        ; case (vi, (vectAvoid || isSimpleApplication ce) && varsS) of
            (VISimple, True) -> liftSimpleAndCase ce
            _                -> do
                                { encCe1 <- encapsulateScalars ce1
                                ; encCe2 <- encapsulateScalars ce2
                                ; return ((fvs, vi), AnnApp encCe1 encCe2)
                                }
        }
      where
        isSimpleApplication :: CoreExprWithVectInfo -> Bool
        isSimpleApplication (_, AnnTick _ ce)                 = isSimpleApplication ce
        isSimpleApplication (_, AnnCast ce _)                 = isSimpleApplication ce
        isSimpleApplication ce                  | isSimple ce = True
        isSimpleApplication (_, AnnApp ce1 ce2)               = isSimple ce1 && isSimpleApplication ce2
        isSimpleApplication _                                 = False
        --
        isSimple :: CoreExprWithVectInfo -> Bool
        isSimple (_, AnnType {})   = True
        isSimple (_, AnnVar  {})   = True
        isSimple (_, AnnLit  {})   = True
        isSimple (_, AnnTick _ ce) = isSimple ce
        isSimple (_, AnnCast ce _) = isSimple ce
        isSimple _                 = False
    encapsulateScalars ce@((fvs, vi), AnnCase scrut bndr ty alts)
      = do
        { vectAvoid <- isVectAvoidanceAggressive
        ; varsS     <- allScalarVarTypeSet fvs
        ; case (vi, vectAvoid && varsS) of
            (VISimple, True) -> liftSimpleAndCase ce
            _                -> do
                                { encScrut <- encapsulateScalars scrut
                                ; encAlts  <- mapM encAlt alts
                                ; return ((fvs, vi), AnnCase encScrut bndr ty encAlts)
                                }
        }
      where
        encAlt (con, bndrs, expr) = (con, bndrs,) <$> encapsulateScalars expr
    encapsulateScalars ce@((fvs, vi), AnnLet (AnnNonRec bndr expr1) expr2)
      = do
        { vectAvoid <- isVectAvoidanceAggressive
        ; varsS     <- allScalarVarTypeSet fvs
        ; case (vi, vectAvoid && varsS) of
            (VISimple, True) -> liftSimpleAndCase ce
            _                -> do
                                { encExpr1 <- encapsulateScalars expr1
                                ; encExpr2 <- encapsulateScalars expr2
                                ; return ((fvs, vi), AnnLet (AnnNonRec bndr encExpr1) encExpr2)
                                }
        }
    encapsulateScalars ce@((fvs, vi), AnnLet (AnnRec binds) expr)
      = do
        { vectAvoid <- isVectAvoidanceAggressive
        ; varsS     <- allScalarVarTypeSet fvs
        ; case (vi, vectAvoid && varsS) of
            (VISimple, True) -> liftSimpleAndCase ce
            _                -> do
                                { encBinds <- mapM encBind binds
                                ; encExpr  <- encapsulateScalars expr
                                ; return ((fvs, vi), AnnLet (AnnRec encBinds) encExpr)
                                }
        }
     where
       encBind (bndr, expr) = (bndr,) <$> encapsulateScalars expr
    encapsulateScalars ((fvs, vi), AnnCast expr coercion)
      = do
        { encExpr <- encapsulateScalars expr
        ; return ((fvs, vi), AnnCast encExpr coercion)
        }
    encapsulateScalars _
      = panic "Vectorise.Exp.encapsulateScalars: unknown constructor"

    -- Lambda-lift the given simple expression and apply it to the abstracted free variables.
    --
    -- If the expression is a case expression scrutinising anything, but a scalar type, then lift
    -- each alternative individually.
    --
    liftSimpleAndCase :: CoreExprWithVectInfo -> VM CoreExprWithVectInfo
    liftSimpleAndCase aexpr@((fvs, _vi), AnnCase expr bndr t alts)
      = do
        { vi <- vectAvoidInfoTypeOf expr
        ; if (vi == VISimple)
          then
            liftSimple aexpr  -- if the scrutinee is scalar, we need no special treatment
          else do
          { alts' <- mapM (\(ac, bndrs, aexpr) -> (ac, bndrs,) <$> liftSimpleAndCase aexpr) alts
          ; return ((fvs, vi), AnnCase expr bndr t alts')
          }
        }
    liftSimpleAndCase aexpr = liftSimple aexpr

    liftSimple :: CoreExprWithVectInfo -> VM CoreExprWithVectInfo
    liftSimple ((fvs, vi), AnnVar v)
      | v `elemDVarSet` fvs               -- special case to avoid producing: (\v -> v) v
      && not (isToplevel v)               --   NB: if 'v' not free or is toplevel, we must get the 'VIEncaps'
      = return $ ((fvs, vi), AnnVar v)
    liftSimple aexpr@((fvs_orig, VISimple), expr)
      = do
        { let liftedExpr = mkAnnApps (mkAnnLams (reverse vars) fvs expr) vars

        ; traceVt "encapsulate:" $ ppr (deAnnotate aexpr) $$ text "==>" $$ ppr (deAnnotate liftedExpr)

        ; return $ liftedExpr
        }
      where
        vars = dVarSetElems fvs
        fvs  = filterDVarSet (not . isToplevel) fvs_orig -- only include 'Id's that are not toplevel

        mkAnnLams :: [Var] -> DVarSet -> AnnExpr' Var (DVarSet, VectAvoidInfo) -> CoreExprWithVectInfo
        mkAnnLams []     fvs expr = ASSERT(isEmptyDVarSet fvs)
                                    ((emptyDVarSet, VIEncaps), expr)
        mkAnnLams (v:vs) fvs expr = mkAnnLams vs (fvs `delDVarSet` v) (AnnLam v ((fvs, VIEncaps), expr))

        mkAnnApps :: CoreExprWithVectInfo -> [Var] -> CoreExprWithVectInfo
        mkAnnApps aexpr []     = aexpr
        mkAnnApps aexpr (v:vs) = mkAnnApps (mkAnnApp aexpr v) vs

        mkAnnApp :: CoreExprWithVectInfo -> Var -> CoreExprWithVectInfo
        mkAnnApp aexpr@((fvs, _vi), _expr) v
          = ((fvs `extendDVarSet` v, VISimple), AnnApp aexpr ((unitDVarSet v, VISimple), AnnVar v))
    liftSimple aexpr
      = pprPanic "Vectorise.Exp.liftSimple: not simple" $ ppr (deAnnotate aexpr)

    isToplevel :: Var -> Bool
    isToplevel v | isId v    = case realIdUnfolding v of
                                 NoUnfolding                     -> False
                                 BootUnfolding                   -> False
                                 OtherCon      {}                -> True
                                 DFunUnfolding {}                -> True
                                 CoreUnfolding {uf_is_top = top} -> top
                 | otherwise = False

    -- |Vectorise an expression.
    --
    vectExpr :: CoreExprWithVectInfo -> VM VExpr

    vectExpr aexpr
        -- encapsulated expression of functional type => try to vectorise as a scalar subcomputation
      | (isFunTy . annExprType $ aexpr) && isVIEncaps aexpr
      = vectFnExpr True False aexpr
        -- encapsulated constant => vectorise as a scalar constant
      | isVIEncaps aexpr
      = traceVt "vectExpr (encapsulated constant):" (ppr . deAnnotate $ aexpr) >>
        vectConst (deAnnotate aexpr)

    vectExpr (_, AnnVar v)
      = vectVar v

    vectExpr (_, AnnLit lit)
      = vectConst $ Lit lit

    vectExpr aexpr@(_, AnnLam _ _)
      = traceVt "vectExpr [AnnLam]:" (ppr . deAnnotate $ aexpr) >>
        vectFnExpr True False aexpr

      -- 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
      = do
        { (vty, lty) <- vectAndLiftType ty
        ; return (mkCoreApps (Var v) [Type (getRuntimeRep vty), Type vty, err'],
                  mkCoreApps (Var v) [Type lty, 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
      = vectPolyApp e

      -- Lifted literal
    vectExpr (_, AnnApp (_, AnnVar v) (_, AnnLit lit))
      | Just _con <- isDataConId_maybe v
      = do
        { let vexpr = App (Var v) (Lit lit)
        ; lexpr <- liftPD vexpr
        ; return (vexpr, lexpr)
        }

      -- 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

    vectExpr (_, AnnCase scrut bndr ty alts)
      | Just (tycon, ty_args) <- splitTyConApp_maybe scrut_ty
      , isAlgTyCon tycon
      = vectAlgCase tycon ty_args scrut bndr ty alts
      | otherwise
      = do
        { dflags <- getDynFlags
        ; cantVectorise dflags "Can't vectorise expression (no algebraic type constructor)" $
            ppr scrut_ty
        }
      where
        scrut_ty = exprType (deAnnotate scrut)

    vectExpr (_, AnnLet (AnnNonRec bndr rhs) body)
      = do
        { traceVt "let binding (non-recursive)" Outputable.empty
        ; vrhs <- localV $
                    inBind bndr $
                      vectAnnPolyExpr False rhs
        ; traceVt "let body (non-recursive)" Outputable.empty
        ; (vbndr, vbody) <- vectBndrIn bndr (vectExpr body)
        ; return $ vLet (vNonRec vbndr vrhs) vbody
        }

    vectExpr (_, AnnLet (AnnRec bs) body)
      = do
        { (vbndrs, (vrhss, vbody)) <- vectBndrsIn bndrs $ do
                                      { traceVt "let bindings (recursive)" Outputable.empty
                                      ; vrhss <- zipWithM vect_rhs bndrs rhss
                                      ; traceVt "let body (recursive)" Outputable.empty
                                      ; vbody <- vectExpr body
                                      ; return (vrhss, vbody)
                                      }
        ; return $ vLet (vRec vbndrs vrhss) vbody
        }
      where
        (bndrs, rhss) = unzip bs

        vect_rhs bndr rhs = localV $
                              inBind bndr $
                                vectAnnPolyExpr (isStrongLoopBreaker $ idOccInfo bndr) rhs

    vectExpr (_, AnnTick tickish expr)
      = vTick tickish <$> vectExpr expr

    vectExpr (_, AnnType ty)
      = vType <$> vectType ty

    vectExpr e
      = do
        { dflags <- getDynFlags
        ; cantVectorise dflags "Can't vectorise expression (vectExpr)" $ ppr (deAnnotate e)
        }

    -- |Vectorise an expression that *may* have an outer lambda abstraction. If the expression is marked
    -- as encapsulated ('VIEncaps'), vectorise it as a scalar computation (using a generalised scalar
    -- zip).
    --
    -- 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
               -> Bool                  -- ^Whether the binding is a loop breaker
               -> CoreExprWithVectInfo  -- ^Expression to vectorise; must have an outer `AnnLam`
               -> VM VExpr
    vectFnExpr inline loop_breaker aexpr@(_ann, AnnLam bndr body)
        -- predicate abstraction: leave as a normal abstraction, but vectorise the predicate type
      | isId bndr
        && isPredTy (idType bndr)
      = do
        { vBndr <- vectBndr bndr
        ; vbody <- vectFnExpr inline loop_breaker body
        ; return $ mapVect (mkLams [vectorised vBndr]) vbody
        }
        -- encapsulated non-predicate abstraction: vectorise as a scalar computation
      | isId bndr && isVIEncaps aexpr
      = vectScalarFun . deAnnotate $ aexpr
        -- non-predicate abstraction: vectorise as a non-scalar computation
      | isId bndr
      = vectLam inline loop_breaker aexpr
      | otherwise
      = do
        { dflags <- getDynFlags
        ; cantVectorise dflags "Vectorise.Exp.vectFnExpr: Unexpected type lambda" $
            ppr (deAnnotate aexpr)
        }
    vectFnExpr _ _ aexpr
        -- encapsulated function: vectorise as a scalar computation
      | (isFunTy . annExprType $ aexpr) && isVIEncaps aexpr
      = vectScalarFun . deAnnotate $ aexpr
      | otherwise
        -- not an abstraction: vectorise as a non-scalar vanilla expression
        -- NB: we can get here due to the recursion in the first case above and from 'vectAnnPolyExpr'
      = vectExpr aexpr

    -- |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 :: CoreExprWithVectInfo -> 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 = text "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 primitive
    -- types (i.e., 'Int', 'Float', 'Double' etc., which have instances of the 'Scalar' type class) and
    -- which does not contain any subcomputations that involve parallel arrays.  Such functionals do not
    -- require the full blown vectorisation transformation; instead, they can be lifted by application
    -- of a member of the zipWith family (i.e., 'map', 'zipWith', zipWith3', etc.)
    --
    -- Dictionary functions are also scalar functions (as dictionaries themselves are not vectorised,
    -- instead they become dictionaries of vectorised methods).  We treat them differently, though see
    -- "Note [Scalar dfuns]" in 'Vectorise'.
    --
    vectScalarFun :: CoreExpr -> VM VExpr
    vectScalarFun expr
      = do
        { traceVt "vectScalarFun:" (ppr expr)
        ; let (arg_tys, res_ty) = splitFunTys (exprType expr)
        ; mkScalarFun arg_tys res_ty expr
        }

    -- Generate code for a scalar function by generating a scalar closure.  If the function is a
    -- dictionary function, vectorise it as dictionary code.
    --
    mkScalarFun :: [Type] -> Type -> CoreExpr -> VM VExpr
    mkScalarFun arg_tys res_ty expr
      | isPredTy res_ty
      = do { vExpr <- vectDictExpr expr
           ; return (vExpr, unused)
           }
      | otherwise
      = do { traceVt "mkScalarFun: " $ ppr expr $$ text "  ::" <+>
                                       ppr (mkFunTys arg_tys res_ty)

           ; fn_var  <- hoistExpr (fsLit "fn") expr DontInline
           ; zipf    <- zipScalars arg_tys res_ty
           ; clo     <- scalarClosure arg_tys res_ty (Var fn_var) (zipf `App` Var fn_var)
           ; clo_var <- hoistExpr (fsLit "clo") clo DontInline
           ; lclo    <- liftPD (Var clo_var)
           ; return (Var clo_var, lclo)
           }
      where
        unused = error "Vectorise.Exp.mkScalarFun: we don't lift dictionary expressions"

    -- |Vectorise a dictionary function that has a 'VECTORISE SCALAR instance' pragma.
    --
    -- In other words, all methods in that dictionary are scalar functions — to be vectorised with
    -- 'vectScalarFun'.  The dictionary "function" itself may be a constant, though.
    --
    -- NB: You may think that we could implement this function guided by the struture of the Core
    --     expression of the right-hand side of the dictionary function.  We cannot proceed like this as
    --     'vectScalarDFun' must also work for *imported* dfuns, where we don't necessarily have access
    --     to the Core code of the unvectorised dfun.
    --
    -- Here an example — assume,
    --
    -- > class Eq a where { (==) :: a -> a -> Bool }
    -- > instance (Eq a, Eq b) => Eq (a, b) where { (==) = ... }
    -- > {-# VECTORISE SCALAR instance Eq (a, b) }
    --
    -- The unvectorised dfun for the above instance has the following signature:
    --
    -- > $dEqPair :: forall a b. Eq a -> Eq b -> Eq (a, b)
    --
    -- We generate the following (scalar) vectorised dfun (liberally using TH notation):
    --
    -- > $v$dEqPair :: forall a b. V:Eq a -> V:Eq b -> V:Eq (a, b)
    -- > $v$dEqPair = /\a b -> \dEqa :: V:Eq a -> \dEqb :: V:Eq b ->
    -- >                D:V:Eq $(vectScalarFun True recFns
    -- >                         [| (==) @(a, b) ($dEqPair @a @b $(unVect dEqa) $(unVect dEqb)) |])
    --
    -- NB:
    -- * '(,)' vectorises to '(,)' — hence, the type constructor in the result type remains the same.
    -- * We share the '$(unVect di)' sub-expressions between the different selectors, but duplicate
    --   the application of the unvectorised dfun, to enable the dictionary selection rules to fire.
    --
    vectScalarDFun :: Var        -- ^ Original dfun
                   -> VM CoreExpr
    vectScalarDFun var
      = do {   -- bring the type variables into scope
           ; mapM_ defLocalTyVar tvs

               -- vectorise dictionary argument types and generate variables for them
           ; vTheta     <- mapM vectType theta
           ; vThetaBndr <- mapM (newLocalVar (fsLit "vd")) vTheta
           ; let vThetaVars = varsToCoreExprs vThetaBndr

               -- vectorise superclass dictionaries and methods as scalar expressions
           ; thetaVars  <- mapM (newLocalVar (fsLit "d")) theta
           ; thetaExprs <- zipWithM unVectDict theta vThetaVars
           ; let thetaDictBinds = zipWith NonRec thetaVars thetaExprs
                 dict           = Var var `mkTyApps` (mkTyVarTys tvs) `mkVarApps` thetaVars
                 scsOps         = map (\selId -> varToCoreExpr selId `mkTyApps` tys `mkApps` [dict])
                                      selIds
           ; vScsOps <- mapM (\e -> vectorised <$> vectScalarFun e) scsOps

               -- vectorised applications of the class-dictionary data constructor
           ; Just vDataCon <- lookupDataCon dataCon
           ; vTys          <- mapM vectType tys
           ; let vBody = thetaDictBinds `mkLets` mkCoreConApps vDataCon (map Type vTys ++ vScsOps)

           ; return $ mkLams (tvs ++ vThetaBndr) vBody
           }
      where
        ty                   = varType var
        (tvs, theta, pty)    = tcSplitSigmaTy  ty        -- 'theta' is the instance context
        (cls, tys)           = tcSplitDFunHead pty       -- 'pty' is the instance head
        selIds               = classAllSelIds cls
        dataCon              = classDataCon cls

    -- Build a value of the dictionary before vectorisation from original, unvectorised type and an
    -- expression computing the vectorised dictionary.
    --
    -- Given the vectorised version of a dictionary 'vd :: V:C vt1..vtn', generate code that computes
    -- the unvectorised version, thus:
    --
    -- > D:C op1 .. opm
    -- > where
    -- >   opi = $(fromVect opTyi [| vSeli @vt1..vtk vd |])
    --
    -- where 'opTyi' is the type of the i-th superclass or op of the unvectorised dictionary.
    --
    unVectDict :: Type -> CoreExpr -> VM CoreExpr
    unVectDict ty e
      = do { vTys <- mapM vectType tys
           ; let meths = map (\sel -> Var sel `mkTyApps` vTys `mkApps` [e]) selIds
           ; scOps <- zipWithM fromVect methTys meths
           ; return $ mkCoreConApps dataCon (map Type tys ++ scOps)
           }
      where
        (tycon, tys) = splitTyConApp ty
        Just dataCon = isDataProductTyCon_maybe tycon
        Just cls     = tyConClass_maybe tycon
        methTys      = dataConInstArgTys dataCon tys
        selIds       = classAllSelIds cls

    -- 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                 -- ^ Should the RHS of a binding be inlined?
            -> Bool                 -- ^ Whether the binding is a loop breaker.
            -> CoreExprWithVectInfo -- ^ Body of abstraction.
            -> VM VExpr
    vectLam inline loop_breaker expr@((fvs, _vi), AnnLam _ _)
     = do { traceVt "fully vectorise a lambda expression" (ppr . deAnnotate $ expr)

          ; 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 <- dVarSetElems 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 { dflags <- getDynFlags
               ; empty <- emptyPD ty
               ; lty   <- mkPDataType ty
               ; return (ve, mkWildCase (Var lc) intPrimTy lty
                               [(DEFAULT, [], le),
                                (LitAlt (mkMachInt dflags 0), [], empty)])
               }
          | otherwise = return (ve, le)
    vectLam _ _ _ = panic "Vectorise.Exp.vectLam: not a lambda"

    -- Vectorise an algebraic case expression.
    --
    -- We convert
    --
    --   case e :: t of v { ... }
    --
    -- to
    --
    --   V:    let v' = e in case v' of _ { ... }
    --   L:    let v' = e in case v' `cast` ... of _ { ... }
    --
    --   When lifting, we have to do it this way because v must have the type
    --   [:V(T):] but the scrutinee must be cast to the representation type. We also
    --   have to handle the case where v is a wild var correctly.
    --

    -- FIXME: this is too lazy...is it?
    vectAlgCase :: TyCon -> [Type] -> CoreExprWithVectInfo -> Var -> Type
                -> [(AltCon, [Var], CoreExprWithVectInfo)]
                -> VM VExpr
    vectAlgCase _tycon _ty_args scrut bndr ty [(DEFAULT, [], body)]
      = do
        { traceVt "scrutinee (DEFAULT only)" Outputable.empty
        ; vscrut         <- vectExpr scrut
        ; (vty, lty)     <- vectAndLiftType ty
        ; traceVt "alternative body (DEFAULT only)" Outputable.empty
        ; (vbndr, vbody) <- vectBndrIn bndr (vectExpr body)
        ; return $ vCaseDEFAULT vscrut vbndr vty lty vbody
        }
    vectAlgCase _tycon _ty_args scrut bndr ty [(DataAlt _, [], body)]
      = do
        { traceVt "scrutinee (one shot w/o binders)" Outputable.empty
        ; vscrut         <- vectExpr scrut
        ; (vty, lty)     <- vectAndLiftType ty
        ; traceVt "alternative body (one shot w/o binders)" Outputable.empty
        ; (vbndr, vbody) <- vectBndrIn bndr (vectExpr body)
        ; return $ vCaseDEFAULT vscrut vbndr vty lty vbody
        }
    vectAlgCase _tycon _ty_args scrut bndr ty [(DataAlt dc, bndrs, body)]
      = do
        { traceVt "scrutinee (one shot w/ binders)" Outputable.empty
        ; vexpr      <- vectExpr scrut
        ; (vty, lty) <- vectAndLiftType ty
        ; traceVt "alternative body (one shot w/ binders)" Outputable.empty
        ; (vbndr, (vbndrs, (vect_body, lift_body)))
            <- vect_scrut_bndr
             . vectBndrsIn bndrs
             $ vectExpr body
        ; let (vect_bndrs, lift_bndrs) = unzip vbndrs
        ; (vscrut, lscrut, pdata_dc) <- pdataUnwrapScrut (vVar vbndr)
        ; vect_dc <- maybeV dataConErr (lookupDataCon dc)

        ; let vcase = mk_wild_case vscrut vty vect_dc  vect_bndrs vect_body
              lcase = mk_wild_case lscrut lty pdata_dc lift_bndrs lift_body

        ; return $ vLet (vNonRec vbndr vexpr) (vcase, lcase)
        }
      where
        vect_scrut_bndr | isDeadBinder bndr = vectBndrNewIn bndr (fsLit "scrut")
                        | otherwise         = vectBndrIn bndr

        mk_wild_case expr ty dc bndrs body
          = mkWildCase expr (exprType expr) ty [(DataAlt dc, bndrs, body)]

        dataConErr = (text "vectAlgCase: data constructor not vectorised" <+> ppr dc)

    vectAlgCase tycon _ty_args scrut bndr ty alts
      = do
        { traceVt "scrutinee (general case)" Outputable.empty
        ; vexpr <- vectExpr scrut

        ; vect_tc     <- vectTyCon tycon
        ; (vty, lty)  <- vectAndLiftType ty

        ; let arity = length (tyConDataCons vect_tc)
        ; sel_ty <- builtin (selTy arity)
        ; sel_bndr <- newLocalVar (fsLit "sel") sel_ty
        ; let sel = Var sel_bndr

        ; traceVt "alternatives' body (general case)" Outputable.empty
        ; (vbndr, valts) <- vect_scrut_bndr
                          $ mapM (proc_alt arity sel vty lty) alts'
        ; let (vect_dcs, vect_bndrss, lift_bndrss, vbodies) = unzip4 valts

        ; (vect_scrut, lift_scrut, pdata_dc) <- pdataUnwrapScrut (vVar vbndr)

        ; let (vect_bodies, lift_bodies) = unzip vbodies

        ; vdummy <- newDummyVar (exprType vect_scrut)
        ; ldummy <- newDummyVar (exprType lift_scrut)
        ; let vect_case = Case vect_scrut vdummy vty
                               (zipWith3 mk_vect_alt vect_dcs vect_bndrss vect_bodies)

        ; lc <- builtin liftingContext
        ; lbody <- combinePD vty (Var lc) sel lift_bodies
        ; let lift_case = Case lift_scrut ldummy lty
                               [(DataAlt pdata_dc, sel_bndr : concat lift_bndrss,
                                 lbody)]

        ; return . vLet (vNonRec vbndr vexpr)
                 $ (vect_case, lift_case)
        }
      where
        vect_scrut_bndr | isDeadBinder bndr = vectBndrNewIn bndr (fsLit "scrut")
                        | otherwise         = vectBndrIn bndr

        alts' = sortBy (\(alt1, _, _) (alt2, _, _) -> cmp alt1 alt2) alts

        cmp (DataAlt dc1) (DataAlt dc2) = dataConTag dc1 `compare` dataConTag dc2
        cmp DEFAULT       DEFAULT       = EQ
        cmp DEFAULT       _             = LT
        cmp _             DEFAULT       = GT
        cmp _             _             = panic "vectAlgCase/cmp"

        proc_alt arity sel _ lty (DataAlt dc, bndrs, body@((fvs_body, _), _))
          = do
              dflags <- getDynFlags
              vect_dc <- maybeV dataConErr (lookupDataCon dc)
              let ntag = dataConTagZ vect_dc
                  tag  = mkDataConTag dflags vect_dc
                  fvs  = fvs_body `delDVarSetList` bndrs

              sel_tags  <- liftM (`App` sel) (builtin (selTags arity))
              lc        <- builtin liftingContext
              elems     <- builtin (selElements arity ntag)

              (vbndrs, vbody)
                <- vectBndrsIn bndrs
                 . localV
                 $ do
                   { binds    <- mapM (pack_var (Var lc) sel_tags tag)
                               . filter isLocalId
                               $ dVarSetElems fvs
                   ; traceVt "case alternative:" (ppr . deAnnotate $ body)
                   ; (ve, le) <- vectExpr body
                   ; return (ve, Case (elems `App` sel) lc lty
                                 [(DEFAULT, [], (mkLets (concat binds) le))])
                   }
                     -- empty    <- emptyPD vty
                     -- return (ve, Case (elems `App` sel) lc lty
                     --             [(DEFAULT, [], Let (NonRec flags_var flags_expr)
                     --                             $ mkLets (concat binds) le),
                     --               (LitAlt (mkMachInt 0), [], empty)])
              let (vect_bndrs, lift_bndrs) = unzip vbndrs
              return (vect_dc, vect_bndrs, lift_bndrs, vbody)
          where
            dataConErr = (text "vectAlgCase: data constructor not vectorised" <+> ppr dc)

        proc_alt _ _ _ _ _ = panic "vectAlgCase/proc_alt"

        mk_vect_alt vect_dc bndrs body = (DataAlt vect_dc, bndrs, body)

          -- Pack a variable for a case alternative context *if* the variable is vectorised. If it
          -- isn't, ignore it as scalar variables don't need to be packed.
        pack_var len tags t v
          = do
            { r <- lookupVar_maybe v
            ; case r of
                Just (Local (vv, lv)) ->
                  do
                  { lv'  <- cloneVar lv
                  ; expr <- packByTagPD (idType vv) (Var lv) len tags t
                  ; updLEnv (\env -> env { local_vars = extendVarEnv (local_vars env) v (vv, lv') })
                  ; return [(NonRec lv' expr)]
                  }
                _ -> return []
            }

    -- Support to compute information for vectorisation avoidance ------------------

    -- Annotation for Core AST nodes that describes how they should be handled during vectorisation
    -- and especially if vectorisation of the corresponding computation can be avoided.
    --
    data VectAvoidInfo = VIParr       -- tree contains parallel computations
                       | VISimple     -- result type is scalar & no parallel subcomputation
                       | VIComplex    -- any result type, no parallel subcomputation
                       | VIEncaps     -- tree encapsulated by 'liftSimple'
                       | VIDict       -- dictionary computation (never parallel)
                       deriving (Eq, Show)

    -- Core expression annotated with free variables and vectorisation-specific information.
    --
    type CoreExprWithVectInfo = AnnExpr Id (DVarSet, VectAvoidInfo)

    -- Yield the type of an annotated core expression.
    --
    annExprType :: AnnExpr Var ann -> Type
    annExprType = exprType . deAnnotate

    -- Project the vectorisation information from an annotated Core expression.
    --
    vectAvoidInfoOf :: CoreExprWithVectInfo -> VectAvoidInfo
    vectAvoidInfoOf ((_, vi), _) = vi

    -- Is this a 'VIParr' node?
    --
    isVIParr :: CoreExprWithVectInfo -> Bool
    isVIParr = (== VIParr) . vectAvoidInfoOf

    -- Is this a 'VIEncaps' node?
    --
    isVIEncaps :: CoreExprWithVectInfo -> Bool
    isVIEncaps = (== VIEncaps) . vectAvoidInfoOf

    -- Is this a 'VIDict' node?
    --
    isVIDict :: CoreExprWithVectInfo -> Bool
    isVIDict = (== VIDict) . vectAvoidInfoOf

    -- 'VIParr' if either argument is 'VIParr'; otherwise, the first argument.
    --
    unlessVIParr :: VectAvoidInfo -> VectAvoidInfo -> VectAvoidInfo
    unlessVIParr _  VIParr = VIParr
    unlessVIParr vi _      = vi

    -- 'VIParr' if either arguments vectorisation information is 'VIParr'; otherwise, the vectorisation
    -- information of the first argument is produced.
    --
    unlessVIParrExpr :: VectAvoidInfo -> CoreExprWithVectInfo -> VectAvoidInfo
    infixl `unlessVIParrExpr`
    unlessVIParrExpr e1 e2 = e1 `unlessVIParr` vectAvoidInfoOf e2

    -- Compute Core annotations to determine for which subexpressions we can avoid vectorisation.
    --
    -- * The first argument is the set of free, local variables whose evaluation may entail parallelism.
    --
    vectAvoidInfo :: VarSet -> CoreExprWithFVs -> VM CoreExprWithVectInfo
    vectAvoidInfo pvs ce@(_, AnnVar v)
      = do
        { gpvs <- globalParallelVars
        ; vi <- if v `elemVarSet` pvs || v `elemDVarSet` gpvs
                then return VIParr
                else vectAvoidInfoTypeOf ce
        ; viTrace ce vi []
        ; when (vi == VIParr) $
            traceVt "  reason:" $ if v `elemVarSet` pvs  then text "local"  else
                                  if v `elemDVarSet` gpvs then text "global" else text "parallel type"

        ; return ((fvs, vi), AnnVar v)
        }
      where
        fvs = freeVarsOf ce

    vectAvoidInfo _pvs ce@(_, AnnLit lit)
      = do
        { vi <- vectAvoidInfoTypeOf ce
        ; viTrace ce vi []
        ; return ((fvs, vi), AnnLit lit)
        }
      where
        fvs = freeVarsOf ce

    vectAvoidInfo pvs ce@(_, AnnApp e1 e2)
      = do
        { ceVI <- vectAvoidInfoTypeOf ce
        ; eVI1 <- vectAvoidInfo pvs e1
        ; eVI2 <- vectAvoidInfo pvs e2
        ; let vi = ceVI `unlessVIParrExpr` eVI1 `unlessVIParrExpr` eVI2
        -- ; viTrace ce vi [eVI1, eVI2]
        ; return ((fvs, vi), AnnApp eVI1 eVI2)
        }
      where
        fvs = freeVarsOf ce

    vectAvoidInfo pvs ce@(_, AnnLam var body)
      = do
        { bodyVI <- vectAvoidInfo pvs body
        ; varVI  <- vectAvoidInfoType $ varType var
        ; let vi = vectAvoidInfoOf bodyVI `unlessVIParr` varVI
        -- ; viTrace ce vi [bodyVI]
        ; return ((fvs, vi), AnnLam var bodyVI)
        }
      where
        fvs = freeVarsOf ce

    vectAvoidInfo pvs ce@(_, AnnLet (AnnNonRec var e) body)
      = do
        { ceVI       <- vectAvoidInfoTypeOf ce
        ; eVI        <- vectAvoidInfo pvs e
        ; isScalarTy <- isScalar $ varType var
        ; (bodyVI, vi) <- if isVIParr eVI && not isScalarTy
            then do -- binding is parallel
            { bodyVI <- vectAvoidInfo (pvs `extendVarSet` var) body
            ; return (bodyVI, VIParr)
            }
            else do -- binding doesn't affect parallelism
            { bodyVI <- vectAvoidInfo pvs body
            ; return (bodyVI, ceVI `unlessVIParrExpr` bodyVI)
            }
        -- ; viTrace ce vi [eVI, bodyVI]
        ; return ((fvs, vi), AnnLet (AnnNonRec var eVI) bodyVI)
        }
      where
        fvs = freeVarsOf ce

    vectAvoidInfo pvs ce@(_, AnnLet (AnnRec bnds) body)
      = do
        { ceVI         <- vectAvoidInfoTypeOf ce
        ; bndsVI       <- mapM (vectAvoidInfoBnd pvs) bnds
        ; parrBndrs    <- map fst <$> filterM isVIParrBnd bndsVI
        ; if not . null $ parrBndrs
          then do         -- body may trigger parallelism via at least one binding
            { new_pvs <- filterM ((not <$>) . isScalar . varType) parrBndrs
            ; let extendedPvs = pvs `extendVarSetList` new_pvs
            ; bndsVI <- mapM (vectAvoidInfoBnd extendedPvs) bnds
            ; bodyVI <- vectAvoidInfo extendedPvs body
            -- ; viTrace ce VIParr (map snd bndsVI ++ [bodyVI])
            ; return ((fvs, VIParr), AnnLet (AnnRec bndsVI) bodyVI)
            }
          else do         -- demanded bindings cannot trigger parallelism
            { bodyVI <- vectAvoidInfo pvs body
            ; let vi = ceVI `unlessVIParrExpr` bodyVI
            -- ; viTrace ce vi (map snd bndsVI ++ [bodyVI])
            ; return ((fvs, vi), AnnLet (AnnRec bndsVI) bodyVI)
            }
        }
      where
        fvs = freeVarsOf ce
        vectAvoidInfoBnd pvs (var, e) = (var,) <$> vectAvoidInfo pvs e

        isVIParrBnd (var, eVI)
          = do
            { isScalarTy <- isScalar (varType var)
            ; return $ isVIParr eVI && not isScalarTy
            }

    vectAvoidInfo pvs ce@(_, AnnCase e var ty alts)
      = do
        { ceVI           <- vectAvoidInfoTypeOf ce
        ; eVI            <- vectAvoidInfo pvs e
        ; altsVI         <- mapM (vectAvoidInfoAlt (isVIParr eVI)) alts
        ; let alteVIs = [eVI | (_, _, eVI) <- altsVI]
              vi      =  foldl unlessVIParrExpr ceVI (eVI:alteVIs)  -- NB: same effect as in the paper
        -- ; viTrace ce vi (eVI : alteVIs)
        ; return ((fvs, vi), AnnCase eVI var ty altsVI)
        }
      where
        fvs = freeVarsOf ce
        vectAvoidInfoAlt scrutIsPar (con, bndrs, e)
          = do
            { allScalar <- allScalarVarType bndrs
            ; let altPvs | scrutIsPar && not allScalar = pvs `extendVarSetList` bndrs
                         | otherwise                   = pvs
            ; (con, bndrs,) <$> vectAvoidInfo altPvs e
            }

    vectAvoidInfo pvs ce@(_, AnnCast e (fvs_ann, ann))
      = do
        { eVI <- vectAvoidInfo pvs e
        ; return ((fvs, vectAvoidInfoOf eVI), AnnCast eVI ((freeVarsOfAnn fvs_ann, VISimple), ann))
        }
      where
        fvs = freeVarsOf ce

    vectAvoidInfo pvs ce@(_, AnnTick tick e)
      = do
        { eVI <- vectAvoidInfo pvs e
        ; return ((fvs, vectAvoidInfoOf eVI), AnnTick tick eVI)
        }
      where
        fvs = freeVarsOf ce

    vectAvoidInfo _pvs ce@(_, AnnType ty)
      = return ((fvs, VISimple), AnnType ty)
      where
        fvs = freeVarsOf ce

    vectAvoidInfo _pvs ce@(_, AnnCoercion coe)
      = return ((fvs, VISimple), AnnCoercion coe)
      where
        fvs = freeVarsOf ce

    -- Compute vectorisation avoidance information for a type.
    --
    vectAvoidInfoType :: Type -> VM VectAvoidInfo
    vectAvoidInfoType ty
      | isPredTy ty
      = return VIDict
      | Just (arg, res) <- splitFunTy_maybe ty
      = do
        { argVI <- vectAvoidInfoType arg
        ; resVI <- vectAvoidInfoType res
        ; case (argVI, resVI) of
            (VISimple, VISimple) -> return VISimple   -- NB: diverts from the paper: scalar functions
            (_       , VIDict)   -> return VIDict
            _                    -> return $ VIComplex `unlessVIParr` argVI `unlessVIParr` resVI
        }
      | otherwise
      = do
        { parr <- maybeParrTy ty
        ; if parr
          then return VIParr
          else do
        { scalar <- isScalar ty
        ; if scalar
          then return VISimple
          else return VIComplex
        } }

    -- Compute vectorisation avoidance information for the type of a Core expression (with FVs).
    --
    vectAvoidInfoTypeOf :: AnnExpr Var ann -> VM VectAvoidInfo
    vectAvoidInfoTypeOf = vectAvoidInfoType . annExprType

    -- Checks whether the type might be a parallel array type.
    --
    maybeParrTy :: Type -> VM Bool
    maybeParrTy ty
        -- looking through newtypes
      | Just ty'      <- coreView ty
      = (== VIParr) <$> vectAvoidInfoType ty'
        -- decompose constructor applications
      | Just (tc, ts) <- splitTyConApp_maybe ty
      = do
        { isParallel <- (tyConName tc `elemNameSet`) <$> globalParallelTyCons
        ; if isParallel
          then return True
          else or <$> mapM maybeParrTy ts
        }
      -- must be a Named ForAllTy because anon ones respond to splitTyConApp_maybe
    maybeParrTy (ForAllTy _ ty) = maybeParrTy ty
    maybeParrTy _               = return False

    -- Are the types of all variables in the 'Scalar' class or toplevel variables?
    --
    -- NB: 'liftSimple' does not abstract over toplevel variables.
    --
    allScalarVarType :: [Var] -> VM Bool
    allScalarVarType vs = and <$> mapM isScalarOrToplevel vs
      where
        isScalarOrToplevel v | isToplevel v = return True
                             | otherwise    = isScalar (varType v)

    -- Are the types of all variables in the set in the 'Scalar' class or toplevel variables?
    --
    allScalarVarTypeSet :: DVarSet -> VM Bool
    allScalarVarTypeSet = allScalarVarType . dVarSetElems

    -- Debugging support
    --
    viTrace :: CoreExprWithFVs -> VectAvoidInfo -> [CoreExprWithVectInfo] -> VM ()
    viTrace ce vi vTs
      = traceVt ("vect info: " ++ show vi ++ "[" ++
                 (concat $ map ((++ " ") . show . vectAvoidInfoOf) vTs) ++ "]")
                (ppr $ deAnnotate ce)