<commit_before>{-# LANGUAGE GeneralizedNewtypeDeriving, TypeSynonymInstances, FlexibleInstances, MultiParamTypeClasses #-}
-- |
--
-- Perform a range of loosely connected low-level transformations
-- based on data dependency information.  This module will:
--
--    * Perform common-subexpression elimination (CSE).
--
--    * Rewrite expressions such that the dependecy path from the root
--    variables (the arguments to the function containing the
--    expression) is as short as possible.  For example, @size(0,b)@
--    will be rewritten to @size(0,a)@ if @b@ is the result of a @map@
--    on @a@, as @a@ and @b@ will in that case have the same number of
--    rows.
--
--    * Hoist expressions out of loops (including lambdas) and
--    branches.  This is done as aggressively as possible.
--
-- For this module to work properly, the input program should be fully
-- normalised; this can be accomplished through use of
-- "L0C.FullNormalizer".
--
-- CSE (and other transformations) may also create many bindings of
-- the form @let a=b@, so it is recommended to run copy propagation
-- after the rebinder.
--
module L0C.Rebinder
  ( transformProg
  , transformProgAggr
  )
  where

import Control.Applicative
import Control.Monad.Writer
import Control.Monad.Reader
import Control.Monad.State
import Control.Monad.RWS

import Data.Graph
import Data.Hashable
import Data.List
import Data.Loc
import Data.Maybe
import qualified Data.HashMap.Strict as HM
import Data.Ord
import qualified Data.HashSet as HS
import qualified Data.Set as S

import L0C.L0
import L0C.MonadFreshNames
import qualified L0C.HORepresentation.SOACNest as Nest

import L0C.Rebinder.CSE
import qualified L0C.Rebinder.SizeTracking as SZ

data BindNeed = LoopBind TupIdent Exp Ident Exp Exp
              | LetBind TupIdent Exp [Exp]
              | LetWithBind Certificates Ident Ident (Maybe Certificates) [Exp] Exp
                deriving (Show, Eq)

type NeedSet = [BindNeed]

asTail :: BindNeed -> Exp
asTail (LoopBind mergepat mergeexp i bound loopbody) =
  DoLoop mergepat mergeexp i bound loopbody (TupLit [] loc) loc
    where loc = srclocOf mergepat
asTail (LetBind pat e _) =
  LetPat pat e (TupLit [] loc) loc
    where loc = srclocOf pat
asTail (LetWithBind cs dest src idxcs idxs ve) =
  LetWith cs dest src idxcs idxs ve (Var dest) $ srclocOf dest

requires :: BindNeed -> HS.HashSet VName
requires (LetBind pat e (alte:alts)) =
  requires (LetBind pat e []) <> requires (LetBind pat alte alts)
requires bnd = HS.map identName $ freeInExp $ asTail bnd

provides :: BindNeed -> HS.HashSet VName
provides (LoopBind mergepat _ _ _ _)  = patNameSet mergepat
provides (LetBind pat _ _)            = patNameSet pat
provides (LetWithBind _ dest _ _ _ _) = HS.singleton $ identName dest

data Need = Need { needBindings :: NeedSet
                 , freeInBound  :: HS.HashSet VName
                 }

instance Monoid Need where
  Need b1 f1 `mappend` Need b2 f2 = Need (b1 <> b2) (f1 <> f2)
  mempty = Need [] HS.empty

data Env = Env { envBindings :: SZ.ShapeMap
               , envDupeState :: DupeState
               , envAggressive :: Bool
               }

emptyEnv :: Env
emptyEnv = Env {
             envBindings = HM.empty
           , envDupeState = newDupeState
           , envAggressive = False
           }

isArrayIdent :: Ident -> Bool
isArrayIdent idd = case identType idd of
                     Array {} -> True
                     _        -> False



newtype HoistM a = HoistM (RWS
                           Env                -- Reader
                           Need               -- Writer
                           (NameSource VName) -- State
                           a)
  deriving (Applicative, Functor, Monad,
            MonadWriter Need, MonadReader Env, MonadState (NameSource VName))

instance MonadFreshNames VName HoistM where
  getNameSource = get
  putNameSource = put

runHoistM :: HoistM a -> NameSource VName -> Env -> a
runHoistM (HoistM m) src env = let (x, _, _) = runRWS m env src
                               in x

needThis :: BindNeed -> HoistM ()
needThis need = tell $ Need [need] HS.empty

boundFree :: HS.HashSet VName -> HoistM ()
boundFree fs = tell $ Need [] fs

withNewBinding :: String -> Exp -> (Ident -> HoistM a) -> HoistM a
withNewBinding k e m = do
  k' <- newVName k
  let ident = Ident { identName = k'
                    , identType = typeOf e
                    , identSrcLoc = srclocOf e }
  withBinding (Id ident) e $ m ident

withBinding :: TupIdent -> Exp -> HoistM a -> HoistM a
withBinding pat e@(Size _ i (Var x) _) m = do
  let mkAlt es = case drop i es of
                   des:_ -> S.toList des
                   _     -> []
  alts <- mkAlt <$> asks (SZ.lookup x . envBindings)
  withSeveralBindings pat e alts m

withBinding pat e m = withSingleBinding pat e m

withSingleBinding :: TupIdent -> Exp -> HoistM a -> HoistM a
withSingleBinding pat e = withSeveralBindings pat e []

withSeveralBindings :: TupIdent -> Exp -> [Exp]
                    -> HoistM a -> HoistM a
withSeveralBindings pat e alts m = do
  ds <- asks envDupeState
  let (e', ds') = performCSE ds pat e
      (es, ds'') = performMultipleCSE ds pat alts
  needThis $ LetBind pat e' es
  local (\env -> env { envDupeState = ds' <> ds''}) m

bindLet :: TupIdent -> Exp -> HoistM a -> HoistM a
bindLet (Id dest) (Var src) m =
  withBinding (Id dest) (Var src) $
  case identType src of Array {} -> withShape dest (slice [] 0 src) m
                        _        -> m

bindLet pat@(Id dest) e@(Iota (Var x) _) m =
  withBinding pat e $
  withShape dest [Var x] m

bindLet pat@(Id dest) e@(Replicate (Var x) (Var y) _) m =
  withBinding pat e $
  withShape dest (Var x:slice [] 0 y) m

bindLet pat@(TupId [Id dest1, Id dest2] _) e@(Split cs (Var n) (Var src) _ loc) m =
  withBinding pat e $
  withNewBinding "split_src_sz" (Size cs 0 (Var src) loc) $ \src_sz ->
  withNewBinding "split_sz" (BinOp Minus (Var src_sz) (Var n) (Elem Int) loc) $ \split_sz ->
  withShapes [(dest1, Var n : rest),
              (dest2, Var split_sz : rest)] m
    where rest = [ Size cs i (Var src) loc
                   | i <- [1.. arrayRank (identType src) - 1]]

bindLet pat@(Id dest) e@(Concat cs (Var x) (Var y) loc) m =
  withBinding pat e $
  withNewBinding "concat_x" (Size cs 0 (Var x) loc) $ \concat_x ->
  withNewBinding "concat_y" (Size cs 0 (Var y) loc) $ \concat_y ->
  withNewBinding "concat_sz" (BinOp Plus (Var concat_x) (Var concat_y) (Elem Int) loc) $ \concat_sz ->
  withShape dest (Var concat_sz :
                  [Size cs i (Var x) loc
                     | i <- [1..arrayRank (identType x) - 1]])
  m
<commit_msg>Make Rebinder accept slightly less normalised programs.<commit_after>{-# LANGUAGE GeneralizedNewtypeDeriving, TypeSynonymInstances, FlexibleInstances, MultiParamTypeClasses #-}
-- |
--
-- Perform a range of loosely connected low-level transformations
-- based on data dependency information.  This module will:
--
--    * Perform common-subexpression elimination (CSE).
--
--    * Rewrite expressions such that the dependecy path from the root
--    variables (the arguments to the function containing the
--    expression) is as short as possible.  For example, @size(0,b)@
--    will be rewritten to @size(0,a)@ if @b@ is the result of a @map@
--    on @a@, as @a@ and @b@ will in that case have the same number of
--    rows.
--
--    * Hoist expressions out of loops (including lambdas) and
--    branches.  This is done as aggressively as possible.
--
-- For this module to work properly, the input program should be fully
-- normalised; this can be accomplished through use of
-- "L0C.FullNormalizer".
--
-- CSE (and other transformations) may also create many bindings of
-- the form @let a=b@, so it is recommended to run copy propagation
-- after the rebinder.
--
module L0C.Rebinder
  ( transformProg
  , transformProgAggr
  )
  where

import Control.Applicative
import Control.Monad.Writer
import Control.Monad.Reader
import Control.Monad.State
import Control.Monad.RWS

import Data.Graph
import Data.Hashable
import Data.List
import Data.Loc
import Data.Maybe
import qualified Data.HashMap.Strict as HM
import Data.Ord
import qualified Data.HashSet as HS
import qualified Data.Set as S

import L0C.L0
import L0C.MonadFreshNames
import qualified L0C.HORepresentation.SOACNest as Nest

import L0C.Rebinder.CSE
import qualified L0C.Rebinder.SizeTracking as SZ

data BindNeed = LoopBind TupIdent Exp Ident Exp Exp
              | LetBind TupIdent Exp [Exp]
              | LetWithBind Certificates Ident Ident (Maybe Certificates) [Exp] Exp
                deriving (Show, Eq)

type NeedSet = [BindNeed]

asTail :: BindNeed -> Exp
asTail (LoopBind mergepat mergeexp i bound loopbody) =
  DoLoop mergepat mergeexp i bound loopbody (TupLit [] loc) loc
    where loc = srclocOf mergepat
asTail (LetBind pat e _) =
  LetPat pat e (TupLit [] loc) loc
    where loc = srclocOf pat
asTail (LetWithBind cs dest src idxcs idxs ve) =
  LetWith cs dest src idxcs idxs ve (Var dest) $ srclocOf dest

requires :: BindNeed -> HS.HashSet VName
requires (LetBind pat e (alte:alts)) =
  requires (LetBind pat e []) <> requires (LetBind pat alte alts)
requires bnd = HS.map identName $ freeInExp $ asTail bnd

provides :: BindNeed -> HS.HashSet VName
provides (LoopBind mergepat _ _ _ _)  = patNameSet mergepat
provides (LetBind pat _ _)            = patNameSet pat
provides (LetWithBind _ dest _ _ _ _) = HS.singleton $ identName dest

data Need = Need { needBindings :: NeedSet
                 , freeInBound  :: HS.HashSet VName
                 }

instance Monoid Need where
  Need b1 f1 `mappend` Need b2 f2 = Need (b1 <> b2) (f1 <> f2)
  mempty = Need [] HS.empty

data Env = Env { envBindings :: SZ.ShapeMap
               , envDupeState :: DupeState
               , envAggressive :: Bool
               }

emptyEnv :: Env
emptyEnv = Env {
             envBindings = HM.empty
           , envDupeState = newDupeState
           , envAggressive = False
           }

isArrayIdent :: Ident -> Bool
isArrayIdent idd = case identType idd of
                     Array {} -> True
                     _        -> False



newtype HoistM a = HoistM (RWS
                           Env                -- Reader
                           Need               -- Writer
                           (NameSource VName) -- State
                           a)
  deriving (Applicative, Functor, Monad,
            MonadWriter Need, MonadReader Env, MonadState (NameSource VName))

instance MonadFreshNames VName HoistM where
  getNameSource = get
  putNameSource = put

runHoistM :: HoistM a -> NameSource VName -> Env -> a
runHoistM (HoistM m) src env = let (x, _, _) = runRWS m env src
                               in x

needThis :: BindNeed -> HoistM ()
needThis need = tell $ Need [need] HS.empty

boundFree :: HS.HashSet VName -> HoistM ()
boundFree fs = tell $ Need [] fs

withNewBinding :: String -> Exp -> (Ident -> HoistM a) -> HoistM a
withNewBinding k e m = do
  k' <- newVName k
  let ident = Ident { identName = k'
                    , identType = typeOf e
                    , identSrcLoc = srclocOf e }
  withBinding (Id ident) e $ m ident

withBinding :: TupIdent -> Exp -> HoistM a -> HoistM a
withBinding pat e@(Size _ i (Var x) _) m = do
  let mkAlt es = case drop i es of
                   des:_ -> S.toList des
                   _     -> []
  alts <- mkAlt <$> asks (SZ.lookup x . envBindings)
  withSeveralBindings pat e alts m

withBinding pat e m = withSingleBinding pat e m

withSingleBinding :: TupIdent -> Exp -> HoistM a -> HoistM a
withSingleBinding pat e = withSeveralBindings pat e []

withSeveralBindings :: TupIdent -> Exp -> [Exp]
                    -> HoistM a -> HoistM a
withSeveralBindings pat e alts m = do
  ds <- asks envDupeState
  let (e', ds') = performCSE ds pat e
      (es, ds'') = performMultipleCSE ds pat alts
  needThis $ LetBind pat e' es
  local (\env -> env { envDupeState = ds' <> ds''}) m

bindLet :: TupIdent -> Exp -> HoistM a -> HoistM a
bindLet (Id dest) (Var src) m =
  withBinding (Id dest) (Var src) $
  case identType src of Array {} -> withShape dest (slice [] 0 src) m
                        _        -> m

bindLet pat@(Id dest) e@(Iota ne _) m =
  withBinding pat e $
  withShape dest [ne] m

bindLet pat@(Id dest) e@(Replicate ne (Var y) _) m =
  withBinding pat e $
  withShape dest (ne:slice [] 0 y) m

bindLet pat@(TupId [Id dest1, Id dest2] _) e@(Split cs ne srce _ loc) m =
  withBinding pat e $
  withNewBinding "split_src_sz" (Size cs 0 srce loc) $ \src_sz ->
  withNewBinding "split_sz" (BinOp Minus (Var src_sz) ne (Elem Int) loc) $ \split_sz ->
  withShapes [(dest1, ne : rest),
              (dest2, Var split_sz : rest)] m
    where rest = [ Size cs i srce loc
                   | i <- [1.. arrayRank (typeOf srce) - 1]]

bindLet pat@(Id dest) e@(Concat cs (Var x) (Var y) loc) m =
  withBinding pat e $
  withNewBinding "concat_x" (Size cs 0 (Var x) loc) $ \concat_x ->
  withNewBinding "concat_y" (Size cs 0 (Var y) loc) $ \concat_y ->
  withNewBinding "concat_sz" (BinOp Plus (Var concat_x) (Var concat_y) (Elem Int) loc) $ \concat_sz ->
  withShape dest (Var concat_sz :
                  [Size cs i (Var x) loc
                     | i <- [1..arrayRank (identType x) - 1]])
  m

{'fraction_non_alphanumeric': 0.0957782451923077, 'fraction_numerical': 0.004657451923076923, 'mean_word_length': 2.424125514403292, 'pattern_counts': {'":': 0, '<': 23, '<?xml version=': 0, '>': 89, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 8, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '5749776', 'n_tokens_mistral': 4393, 'n_tokens_neox': 4233, 'n_words': 1550}