<commit_before>{-# LANGUAGE GeneralizedNewtypeDeriving #-}
{-# LANGUAGE FlexibleInstances #-}
{-# LANGUAGE MultiParamTypeClasses #-}
module Futhark.Optimise.Fusion ( fuseSOACs )
  where

import Control.Monad.State
import Control.Applicative
import Control.Monad.Reader
import Control.Monad.Except

import Data.Hashable
import Data.Maybe
import Data.Monoid

import qualified Data.HashMap.Lazy as HM
import qualified Data.HashSet      as HS
import qualified Data.List         as L

import Prelude

import Futhark.Representation.SOACS hiding (SOAC(..))
import qualified Futhark.Representation.SOACS as Futhark
import Futhark.MonadFreshNames
import Futhark.Representation.SOACS.Simplify
import Futhark.Optimise.Fusion.LoopKernel
import Futhark.Binder
import qualified Futhark.Analysis.HORepresentation.SOAC as SOAC
import Futhark.Transform.Rename
import Futhark.Pass

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

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

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

arrsInScope :: FusionGEnv -> HM.HashMap VName (VName, Type, SOAC.ArrayTransforms)
arrsInScope = HM.fromList . mapMaybe asArray . HM.toList . varsInScope
  where asArray (name, IsArray srcname ts attr) =
          Just (name, (srcname, typeOf attr, ts))
        asArray (_, IsNotArray _ _) =
          Nothing

data 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 = toScope <$> asks varsInScope
    where toScope = HM.map varEntryType

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

arrayTransforms :: VName -> Type -> FusionGM (VName, Type, SOAC.ArrayTransforms)
arrayTransforms name t = do
  v' <- asks $ HM.lookup name . arrsInScope
  case v' of Just res -> return res
             Nothing  -> return (name, t, SOAC.noTransforms)

-- | Binds an array name to the set of used-array vars
bindVar :: FusionGEnv -> Ident -> FusionGEnv
bindVar env (Ident name t) =
  env { varsInScope = HM.insert name entry $ varsInScope env }
  where entry = case t of
          Array {} -> IsArray name mempty $ LetInfo t
          _        -> IsNotArray name $ LetInfo t

bindVars :: FusionGEnv -> [Ident] -> FusionGEnv
bindVars = foldl bindVar

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

gatherBindingPattern :: Pattern -> FusionGM FusedRes -> FusionGM FusedRes
gatherBindingPattern pat m = do
  res <- binding (patternIdents pat) m
  checkForUpdates pat res

bindingPat :: Pattern -> FusionGM a -> FusionGM a
bindingPat = binding . patternIdents

bindingFParams :: [FParam] -> FusionGM a -> FusionGM a
bindingFParams = binding  . 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       = HM.insert nm faml $ soacs env
      , varsInScope = HM.insert nm (IsArray nm mempty $ LetInfo t) $
                      varsInScope env
      }

checkForUpdates :: Pattern -> FusedRes -> FusionGM FusedRes
checkForUpdates pat res = foldM checkForUpdate res $ patternElements pat
  where checkForUpdate res' (PatElem _ BindVar _) =
          return res'
        checkForUpdate res' (PatElem _ (BindInPlace _ src is) _) = do
          res'' <- foldM fusionGatherSubExp res' (Var src : is)
          let aliases = [src]
              inspectKer k =
                let inplace' = foldl (flip HS.insert) (inplace k) aliases
                in  k { inplace = inplace' }
          return $ res'' { kernels = HM.map inspectKer $ kernels res'' }

-- | Updates the environment: (i) the @soacs@ (map) by binding each pattern
--   element identifier to all pattern elements (identifiers) and (ii) the
--   @arrsInScope@ (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 m = do
  res <- local bind m
  checkForUpdates pat res
  where idents = patternIdents pat
        family = patternNames pat
        bind env = foldl (bindingFamilyVar family) env idents

bindingTransform :: Ident -> VName -> SOAC.ArrayTransform -> FusionGM a -> FusionGM a
bindingTransform v srcname trns = local $ \env ->
  case HM.lookup srcname $ varsInScope env of
    Just (IsArray src' ts attr) ->
      env { varsInScope =
              HM.insert vname (IsArray src' (ts SOAC.|> trns) attr) $
              varsInScope env
          }
    _ -> bindVar env v
  where vname = identName v

-- | Binds the fusion result to the environment.
bindRes :: FusedRes -> FusionGM a -> FusionGM a
bindRes rrr = local (\x -> x { fusedRes = rrr })

-- | The fusion transformation runs in this monad.  The mutable
-- state refers to the fresh-names engine.
-- The reader hides the vtable that associates ... to ... (fill in, please).
-- The 'Either' monad is used for error handling.
runFusionGatherM :: MonadFreshNames m =>
                    FusionGM a -> FusionGEnv -> m (Either Error a)
runFusionGatherM (FusionGM a) env =
  modifyNameSource $ \src -> runReader (runStateT (runExceptT a) src) env

badFusionGM :: Error -> FusionGM a
badFusionGM = throwError

------------------------------------------------------------------------
--- Fusion Entry Points: gather the to-be-fused kernels@pgm level    ---
---    and fuse them in a second pass!                               ---
------------------------------------------------------------------------

fuseSOACs :: Pass SOACS SOACS
fuseSOACs =
  Pass { passName = "Fuse SOACs"
       , passDescription = "Perform higher-order optimisation, i.e., fusion."
       , passFunction = fuseProg
       }

fuseProg :: Prog -> PassM Prog
fuseProg prog = do
  let env  = FusionGEnv { soacs = HM.empty
                        , varsInScope = HM.empty
                        , fusedRes = mkFreshFusionRes
                        }
      funs = progFunctions prog
  ks <- liftEitherM $ runFusionGatherM (mapM fusionGatherFun funs) env
  let ks'    = map cleanFusionResult ks
  let succc = any rsucc ks'
  if not succc
  then return prog
  else do funs' <- liftEitherM $ runFusionGatherM (zipWithM fuseInFun ks' funs) env
          renameProg $ Prog funs'

fusionGatherFun :: FunDec -> FusionGM FusedRes
fusionGatherFun fundec =
  bindingFParams (funDecParams fundec) $
  fusionGatherBody mkFreshFusionRes $ funDecBody fundec

fuseInFun :: FusedRes -> FunDec -> FusionGM FunDec
fuseInFun res fundec = do<commit_msg>Run dead code elimination inside the fusion module.<commit_after>{-# LANGUAGE GeneralizedNewtypeDeriving #-}
{-# LANGUAGE FlexibleInstances #-}
{-# LANGUAGE MultiParamTypeClasses #-}
module Futhark.Optimise.Fusion ( fuseSOACs )
  where

import Control.Monad.State
import Control.Applicative
import Control.Monad.Reader
import Control.Monad.Except

import Data.Hashable
import Data.Maybe
import Data.Monoid

import qualified Data.HashMap.Lazy as HM
import qualified Data.HashSet      as HS
import qualified Data.List         as L

import Prelude

import Futhark.Optimise.DeadVarElim (deadCodeElim)
import Futhark.Representation.SOACS hiding (SOAC(..))
import qualified Futhark.Representation.SOACS as Futhark
import Futhark.MonadFreshNames
import Futhark.Representation.SOACS.Simplify
import Futhark.Optimise.Fusion.LoopKernel
import Futhark.Binder
import qualified Futhark.Analysis.HORepresentation.SOAC as SOAC
import Futhark.Transform.Rename
import Futhark.Pass

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

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

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

arrsInScope :: FusionGEnv -> HM.HashMap VName (VName, Type, SOAC.ArrayTransforms)
arrsInScope = HM.fromList . mapMaybe asArray . HM.toList . varsInScope
  where asArray (name, IsArray srcname ts attr) =
          Just (name, (srcname, typeOf attr, ts))
        asArray (_, IsNotArray _ _) =
          Nothing

data 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 = toScope <$> asks varsInScope
    where toScope = HM.map varEntryType

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

arrayTransforms :: VName -> Type -> FusionGM (VName, Type, SOAC.ArrayTransforms)
arrayTransforms name t = do
  v' <- asks $ HM.lookup name . arrsInScope
  case v' of Just res -> return res
             Nothing  -> return (name, t, SOAC.noTransforms)

-- | Binds an array name to the set of used-array vars
bindVar :: FusionGEnv -> Ident -> FusionGEnv
bindVar env (Ident name t) =
  env { varsInScope = HM.insert name entry $ varsInScope env }
  where entry = case t of
          Array {} -> IsArray name mempty $ LetInfo t
          _        -> IsNotArray name $ LetInfo t

bindVars :: FusionGEnv -> [Ident] -> FusionGEnv
bindVars = foldl bindVar

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

gatherBindingPattern :: Pattern -> FusionGM FusedRes -> FusionGM FusedRes
gatherBindingPattern pat m = do
  res <- binding (patternIdents pat) m
  checkForUpdates pat res

bindingPat :: Pattern -> FusionGM a -> FusionGM a
bindingPat = binding . patternIdents

bindingFParams :: [FParam] -> FusionGM a -> FusionGM a
bindingFParams = binding  . 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       = HM.insert nm faml $ soacs env
      , varsInScope = HM.insert nm (IsArray nm mempty $ LetInfo t) $
                      varsInScope env
      }

checkForUpdates :: Pattern -> FusedRes -> FusionGM FusedRes
checkForUpdates pat res = foldM checkForUpdate res $ patternElements pat
  where checkForUpdate res' (PatElem _ BindVar _) =
          return res'
        checkForUpdate res' (PatElem _ (BindInPlace _ src is) _) = do
          res'' <- foldM fusionGatherSubExp res' (Var src : is)
          let aliases = [src]
              inspectKer k =
                let inplace' = foldl (flip HS.insert) (inplace k) aliases
                in  k { inplace = inplace' }
          return $ res'' { kernels = HM.map inspectKer $ kernels res'' }

-- | Updates the environment: (i) the @soacs@ (map) by binding each pattern
--   element identifier to all pattern elements (identifiers) and (ii) the
--   @arrsInScope@ (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 m = do
  res <- local bind m
  checkForUpdates pat res
  where idents = patternIdents pat
        family = patternNames pat
        bind env = foldl (bindingFamilyVar family) env idents

bindingTransform :: Ident -> VName -> SOAC.ArrayTransform -> FusionGM a -> FusionGM a
bindingTransform v srcname trns = local $ \env ->
  case HM.lookup srcname $ varsInScope env of
    Just (IsArray src' ts attr) ->
      env { varsInScope =
              HM.insert vname (IsArray src' (ts SOAC.|> trns) attr) $
              varsInScope env
          }
    _ -> bindVar env v
  where vname = identName v

-- | Binds the fusion result to the environment.
bindRes :: FusedRes -> FusionGM a -> FusionGM a
bindRes rrr = local (\x -> x { fusedRes = rrr })

-- | The fusion transformation runs in this monad.  The mutable
-- state refers to the fresh-names engine.
-- The reader hides the vtable that associates ... to ... (fill in, please).
-- The 'Either' monad is used for error handling.
runFusionGatherM :: MonadFreshNames m =>
                    FusionGM a -> FusionGEnv -> m (Either Error a)
runFusionGatherM (FusionGM a) env =
  modifyNameSource $ \src -> runReader (runStateT (runExceptT a) src) env

badFusionGM :: Error -> FusionGM a
badFusionGM = throwError

------------------------------------------------------------------------
--- Fusion Entry Points: gather the to-be-fused kernels@pgm level    ---
---    and fuse them in a second pass!                               ---
------------------------------------------------------------------------

fuseSOACs :: Pass SOACS SOACS
fuseSOACs =
  Pass { passName = "Fuse SOACs"
       , passDescription = "Perform higher-order optimisation, i.e., fusion."
       , passFunction = fuseProg
       }

fuseProg :: Prog -> PassM Prog
fuseProg prog = do
  let env  = FusionGEnv { soacs = HM.empty
                        , varsInScope = HM.empty
                        , fusedRes = mkFreshFusionRes
                        }
      funs = progFunctions prog
  ks <- liftEitherM $ runFusionGatherM (mapM fusionGatherFun funs) env
  let ks'    = map cleanFusionResult ks
  let succc = any rsucc ks'
  if not succc
  then return prog
  else do funs' <- liftEitherM $ runFusionGatherM (zipWithM fuseInFun ks' funs) env
          deadCodeElim <$> renameProg (Prog funs')

fusionGatherFun :: FunDec -> FusionGM FusedRes
fusionGatherFun fundec =
  bindingFParams (funDecParams fundec) $
  fusionGatherBody mkFreshFusionRes $ funDecBody fundec

fuseInFun :: FusedRes -> FunDec -> FusionGM FunDec
fuseInFun res fundec = do
{'fraction_non_alphanumeric': 0.12072663895964607, 'fraction_numerical': 0.0, 'mean_word_length': 2.501290776812955, 'pattern_counts': {'":': 0, '<': 18, '<?xml version=': 0, '>': 100, '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': '7224391', 'n_tokens_mistral': 4622, 'n_tokens_neox': 4387, 'n_words': 1486}