<commit_before>{-# LANGUAGE 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 TypeRep
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 MonadUtils

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.
--  
encapsulateScalars :: CoreExprWithVectInfo -> VM CoreExprWithVectInfo
encapsulateScalars ce@(_, AnnType _ty)
  = return ce
encapsulateScalars ce@((_, VISimple), AnnVar v)
      -- NB: diverts from the paper: encapsulate variables with scalar type (includes 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 
    { varsS <- allScalarVarTypeSet fvs 
    ; case (vi, varsS) of
        (VISimple, True) -> liftSimpleAndCase ce
        _                -> do 
                            { encExpr <- encapsulateScalars expr
                            ; return ((fvs, vi), AnnLam bndr encExpr)
                            }
    }
encapsulateScalars ce@((fvs, vi), AnnApp ce1 ce2)
  = do
    { varsS <- allScalarVarTypeSet fvs
    ; case (vi, varsS) of
        (VISimple, True) -> liftSimpleAndCase ce
        _                -> do 
                            { encCe1 <- encapsulateScalars ce1
                            ; encCe2 <- encapsulateScalars ce2
                            ; return ((fvs, vi), AnnApp encCe1 encCe2)
                            }
    }
encapsulateScalars ce@((fvs, vi), AnnCase scrut bndr ty alts) 
  = do 
    { varsS <- allScalarVarTypeSet fvs 
    ; case (vi, 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 
    { varsS <- allScalarVarTypeSet fvs 
    ; case (vi, varsS) of
        (VISimple, True) -> liftSimpleAndCase ce<commit_msg>Vectoriser: binders of encapsulated lambdas need to be scalar<commit_after>{-# LANGUAGE 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 TypeRep
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 MonadUtils

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.
--  
encapsulateScalars :: CoreExprWithVectInfo -> VM CoreExprWithVectInfo
encapsulateScalars ce@(_, AnnType _ty)
  = return ce
encapsulateScalars ce@((_, VISimple), AnnVar _v)
      -- NB: diverts from the paper: encapsulate variables with scalar type (includes 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 
    { 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, 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
    { varsS <- allScalarVarTypeSet fvs
    ; case (vi, varsS) of
        (VISimple, True) -> liftSimpleAndCase ce
        _                -> do 
                            { encCe1 <- encapsulateScalars ce1
                            ; encCe2 <- encapsulateScalars ce2
                            ; return ((fvs, vi), AnnApp encCe1 encCe2)
                            }
    }
encapsulateScalars ce@((fvs, vi), AnnCase scrut bndr ty alts) 
  = do 
    { varsS <- allScalarVarTypeSet fvs 
    ; case (vi, 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 
    { varsS <- allScalarVarTypeSet fvs 
    ; case (vi, varsS) of
        (VISimple, True) -> liftSimpleAndCase ce
{'fraction_non_alphanumeric': 0.08239947264337508, 'fraction_numerical': 0.0013183915622940012, 'mean_word_length': 2.462117754450023, 'pattern_counts': {'":': 0, '<': 58, '<?xml version=': 0, '>': 49, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 0, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '5610534', 'n_tokens_mistral': 4558, 'n_tokens_neox': 4349, 'n_words': 1525}