<commit_before>{-# LANGUAGE FlexibleContexts #-}
{-# LANGUAGE GeneralizedNewtypeDeriving #-}
{-# LANGUAGE MultiParamTypeClasses #-}
{-# LANGUAGE TypeFamilies #-}

-- | Perform horizontal and vertical fusion of SOACs.  See the paper
-- /A T2 Graph-Reduction Approach To Fusion/ for the basic idea (some
-- extensions discussed in /Design and GPGPU Performance of Futhark’s
-- Redomap Construct/).
module Futhark.Optimise.Fusion (fuseSOACs) where

import Control.Monad.Except
import Control.Monad.Reader
import Control.Monad.State
import qualified Data.List as L
import qualified Data.Map.Strict as M
import Data.Maybe
import qualified Data.Set as S
import qualified Futhark.Analysis.Alias as Alias
import qualified Futhark.Analysis.HORep.SOAC as SOAC
import Futhark.Construct
import qualified Futhark.IR.Aliases as Aliases
import Futhark.IR.Prop.Aliases
import Futhark.IR.SOACS hiding (SOAC (..))
import qualified Futhark.IR.SOACS as Futhark
import Futhark.IR.SOACS.Simplify
import Futhark.Optimise.Fusion.LoopKernel
import Futhark.Pass
import Futhark.Transform.Rename
import Futhark.Transform.Substitute
import Futhark.Util (maxinum)

data VarEntry
  = IsArray VName (NameInfo SOACS) Names SOAC.Input
  | IsNotArray (NameInfo SOACS)

varEntryType :: VarEntry -> NameInfo SOACS
varEntryType (IsArray _ dec _ _) =
  dec
varEntryType (IsNotArray dec) =
  dec

varEntryAliases :: VarEntry -> Names
varEntryAliases (IsArray _ _ x _) = x
varEntryAliases _ = mempty

data FusionGEnv = FusionGEnv
  { -- | Mapping from variable name to its entire family.
    soacs :: M.Map VName [VName],
    varsInScope :: M.Map VName VarEntry,
    fusedRes :: FusedRes
  }

lookupArr :: VName -> FusionGEnv -> Maybe SOAC.Input
lookupArr v env = asArray =<< M.lookup v (varsInScope env)
  where
    asArray (IsArray _ _ _ input) = Just input
    asArray IsNotArray {} = Nothing

newtype Error = Error String

instance Show Error where
  show (Error msg) = "Fusion error:\n" ++ msg

newtype FusionGM a = FusionGM (ExceptT Error (StateT VNameSource (Reader FusionGEnv)) a)
  deriving
    ( Monad,
      Applicative,
      Functor,
      MonadError Error,
      MonadState VNameSource,
      MonadReader FusionGEnv
    )

instance MonadFreshNames FusionGM where
  getNameSource = get
  putNameSource = put

instance HasScope SOACS FusionGM where
  askScope = asks $ toScope . varsInScope
    where
      toScope = M.map varEntryType

------------------------------------------------------------------------
--- Monadic Helpers: bind/new/runFusionGatherM, etc
------------------------------------------------------------------------

-- | Binds an array name to the set of used-array vars
bindVar :: FusionGEnv -> (Ident, Names) -> FusionGEnv
bindVar env (Ident name t, aliases) =
  env {varsInScope = M.insert name entry $ varsInScope env}
  where
    entry = case t of
      Array {} -> IsArray name (LetName t) aliases' $ SOAC.identInput $ Ident name t
      _ -> IsNotArray $ LetName t
    expand = maybe mempty varEntryAliases . flip M.lookup (varsInScope env)
    aliases' = aliases <> mconcat (map expand $ namesToList aliases)

bindVars :: FusionGEnv -> [(Ident, Names)] -> FusionGEnv
bindVars = foldl bindVar

binding :: [(Ident, Names)] -> FusionGM a -> FusionGM a
binding vs = local (`bindVars` vs)

gatherStmPattern :: Pattern -> Exp -> FusionGM FusedRes -> FusionGM FusedRes
gatherStmPattern pat e = binding $ zip idents aliases
  where
    idents = patternIdents pat
    aliases =
      replicate (length (patternContextNames pat)) mempty
        ++ expAliases (Alias.analyseExp mempty e)

bindingPat :: Pattern -> FusionGM a -> FusionGM a
bindingPat = binding . (`zip` repeat mempty) . patternIdents

bindingParams :: Typed t => [Param t] -> FusionGM a -> FusionGM a
bindingParams = binding . (`zip` repeat mempty) . map paramIdent

-- | Binds an array name to the set of soac-produced vars
bindingFamilyVar :: [VName] -> FusionGEnv -> Ident -> FusionGEnv
bindingFamilyVar faml env (Ident nm t) =
  env
    { soacs = M.insert nm faml $ soacs env,
      varsInScope =
        M.insert
          nm
          ( IsArray nm (LetName t) mempty $
              SOAC.identInput $ Ident nm t
          )
          $ varsInScope env
    }

varAliases :: VName -> FusionGM Names
varAliases v =
  asks $
    (oneName v <>) . maybe mempty varEntryAliases
      . M.lookup v
      . varsInScope

varsAliases :: Names -> FusionGM Names
varsAliases = fmap mconcat . mapM varAliases . namesToList

updateKerInPlaces :: FusedRes -> ([VName], [VName]) -> FusionGM FusedRes
updateKerInPlaces res (ip_vs, other_infuse_vs) = do
  res' <- foldM addVarToInfusible res (ip_vs ++ other_infuse_vs)
  aliases <- mconcat <$> mapM varAliases ip_vs
  let inspectKer k = k {inplace = aliases <> inplace k}
  return res' {kernels = M.map inspectKer $ kernels res'}

checkForUpdates :: FusedRes -> Exp -> FusionGM FusedRes
checkForUpdates res (BasicOp (Update src is _)) = do
  let ifvs = namesToList $ mconcat $ map freeIn is
  updateKerInPlaces res ([src], ifvs)
checkForUpdates res (Op (Futhark.Scatter _ _ _ written_info)) = do
  let updt_arrs = map (\(_,_,x)->x) written_info
  updateKerInPlaces res (updt_arrs, [])
checkForUpdates res _ = return res

-- | Updates the environment: (i) the @soacs@ (map) by binding each pattern
--   element identifier to all pattern elements (identifiers) and (ii) the
--   variables in scope (map) by inserting each (pattern-array) name.
--   Finally, if the binding is an in-place update, then the @inplace@ field
--   of each (result) kernel is updated with the new in-place updates.
bindingFamily :: Pattern -> FusionGM FusedRes -> FusionGM FusedRes
bindingFamily pat = local bind
  where
    idents = patternIdents pat
    family = patternNames pat
    bind env = foldl (bindingFamilyVar family) env idents

bindingTransform :: PatElem -> VName -> SOAC.ArrayTransform -> FusionGM a -> FusionGM a
bindingTransform pe srcname trns = local $ \env -><commit_msg>Style fix.<commit_after>{-# LANGUAGE FlexibleContexts #-}
{-# LANGUAGE GeneralizedNewtypeDeriving #-}
{-# LANGUAGE MultiParamTypeClasses #-}
{-# LANGUAGE TypeFamilies #-}

-- | Perform horizontal and vertical fusion of SOACs.  See the paper
-- /A T2 Graph-Reduction Approach To Fusion/ for the basic idea (some
-- extensions discussed in /Design and GPGPU Performance of Futhark’s
-- Redomap Construct/).
module Futhark.Optimise.Fusion (fuseSOACs) where

import Control.Monad.Except
import Control.Monad.Reader
import Control.Monad.State
import qualified Data.List as L
import qualified Data.Map.Strict as M
import Data.Maybe
import qualified Data.Set as S
import qualified Futhark.Analysis.Alias as Alias
import qualified Futhark.Analysis.HORep.SOAC as SOAC
import Futhark.Construct
import qualified Futhark.IR.Aliases as Aliases
import Futhark.IR.Prop.Aliases
import Futhark.IR.SOACS hiding (SOAC (..))
import qualified Futhark.IR.SOACS as Futhark
import Futhark.IR.SOACS.Simplify
import Futhark.Optimise.Fusion.LoopKernel
import Futhark.Pass
import Futhark.Transform.Rename
import Futhark.Transform.Substitute
import Futhark.Util (maxinum)

data VarEntry
  = IsArray VName (NameInfo SOACS) Names SOAC.Input
  | IsNotArray (NameInfo SOACS)

varEntryType :: VarEntry -> NameInfo SOACS
varEntryType (IsArray _ dec _ _) =
  dec
varEntryType (IsNotArray dec) =
  dec

varEntryAliases :: VarEntry -> Names
varEntryAliases (IsArray _ _ x _) = x
varEntryAliases _ = mempty

data FusionGEnv = FusionGEnv
  { -- | Mapping from variable name to its entire family.
    soacs :: M.Map VName [VName],
    varsInScope :: M.Map VName VarEntry,
    fusedRes :: FusedRes
  }

lookupArr :: VName -> FusionGEnv -> Maybe SOAC.Input
lookupArr v env = asArray =<< M.lookup v (varsInScope env)
  where
    asArray (IsArray _ _ _ input) = Just input
    asArray IsNotArray {} = Nothing

newtype Error = Error String

instance Show Error where
  show (Error msg) = "Fusion error:\n" ++ msg

newtype FusionGM a = FusionGM (ExceptT Error (StateT VNameSource (Reader FusionGEnv)) a)
  deriving
    ( Monad,
      Applicative,
      Functor,
      MonadError Error,
      MonadState VNameSource,
      MonadReader FusionGEnv
    )

instance MonadFreshNames FusionGM where
  getNameSource = get
  putNameSource = put

instance HasScope SOACS FusionGM where
  askScope = asks $ toScope . varsInScope
    where
      toScope = M.map varEntryType

------------------------------------------------------------------------
--- Monadic Helpers: bind/new/runFusionGatherM, etc
------------------------------------------------------------------------

-- | Binds an array name to the set of used-array vars
bindVar :: FusionGEnv -> (Ident, Names) -> FusionGEnv
bindVar env (Ident name t, aliases) =
  env {varsInScope = M.insert name entry $ varsInScope env}
  where
    entry = case t of
      Array {} -> IsArray name (LetName t) aliases' $ SOAC.identInput $ Ident name t
      _ -> IsNotArray $ LetName t
    expand = maybe mempty varEntryAliases . flip M.lookup (varsInScope env)
    aliases' = aliases <> mconcat (map expand $ namesToList aliases)

bindVars :: FusionGEnv -> [(Ident, Names)] -> FusionGEnv
bindVars = foldl bindVar

binding :: [(Ident, Names)] -> FusionGM a -> FusionGM a
binding vs = local (`bindVars` vs)

gatherStmPattern :: Pattern -> Exp -> FusionGM FusedRes -> FusionGM FusedRes
gatherStmPattern pat e = binding $ zip idents aliases
  where
    idents = patternIdents pat
    aliases =
      replicate (length (patternContextNames pat)) mempty
        ++ expAliases (Alias.analyseExp mempty e)

bindingPat :: Pattern -> FusionGM a -> FusionGM a
bindingPat = binding . (`zip` repeat mempty) . patternIdents

bindingParams :: Typed t => [Param t] -> FusionGM a -> FusionGM a
bindingParams = binding . (`zip` repeat mempty) . map paramIdent

-- | Binds an array name to the set of soac-produced vars
bindingFamilyVar :: [VName] -> FusionGEnv -> Ident -> FusionGEnv
bindingFamilyVar faml env (Ident nm t) =
  env
    { soacs = M.insert nm faml $ soacs env,
      varsInScope =
        M.insert
          nm
          ( IsArray nm (LetName t) mempty $
              SOAC.identInput $ Ident nm t
          )
          $ varsInScope env
    }

varAliases :: VName -> FusionGM Names
varAliases v =
  asks $
    (oneName v <>) . maybe mempty varEntryAliases
      . M.lookup v
      . varsInScope

varsAliases :: Names -> FusionGM Names
varsAliases = fmap mconcat . mapM varAliases . namesToList

updateKerInPlaces :: FusedRes -> ([VName], [VName]) -> FusionGM FusedRes
updateKerInPlaces res (ip_vs, other_infuse_vs) = do
  res' <- foldM addVarToInfusible res (ip_vs ++ other_infuse_vs)
  aliases <- mconcat <$> mapM varAliases ip_vs
  let inspectKer k = k {inplace = aliases <> inplace k}
  return res' {kernels = M.map inspectKer $ kernels res'}

checkForUpdates :: FusedRes -> Exp -> FusionGM FusedRes
checkForUpdates res (BasicOp (Update src is _)) = do
  let ifvs = namesToList $ mconcat $ map freeIn is
  updateKerInPlaces res ([src], ifvs)
checkForUpdates res (Op (Futhark.Scatter _ _ _ written_info)) = do
  let updt_arrs = map (\(_, _, x) -> x) written_info
  updateKerInPlaces res (updt_arrs, [])
checkForUpdates res _ = return res

-- | Updates the environment: (i) the @soacs@ (map) by binding each pattern
--   element identifier to all pattern elements (identifiers) and (ii) the
--   variables in scope (map) by inserting each (pattern-array) name.
--   Finally, if the binding is an in-place update, then the @inplace@ field
--   of each (result) kernel is updated with the new in-place updates.
bindingFamily :: Pattern -> FusionGM FusedRes -> FusionGM FusedRes
bindingFamily pat = local bind
  where
    idents = patternIdents pat
    family = patternNames pat
    bind env = foldl (bindingFamilyVar family) env idents

bindingTransform :: PatElem -> VName -> SOAC.ArrayTransform -> FusionGM a -> FusionGM a
bindingTransform pe srcname trns = local $ \env ->
{'fraction_non_alphanumeric': 0.1127150650440621, 'fraction_numerical': 0.0001678556441460344, 'mean_word_length': 3.176656151419558, 'pattern_counts': {'":': 0, '<': 19, '<?xml version=': 0, '>': 85, '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': '3927036', 'n_tokens_mistral': 3825, 'n_tokens_neox': 3694, 'n_words': 1236}