<commit_before>{-# LANGUAGE ExistentialQuantification #-}
{-# LANGUAGE BangPatterns #-}

module Flow.Run
  ( dumpStrategy
  , dumpStrategyDOT
  , dumpSteps
  , execStrategy
  ) where

import Control.Monad
import Control.Applicative
import Control.Concurrent

import Data.List
import Data.Maybe ( fromMaybe, fromJust, mapMaybe )
import qualified Data.IntMap as IM
import qualified Data.IntSet as IS
import qualified Data.Map.Strict as Map
import Data.IORef
import Data.Ord
import Data.Time.Clock
import Data.Typeable

import System.IO

import Flow.Internal
import Flow.Builder
import Flow.Vector

dumpStrategy :: Strategy a -> IO ()
dumpStrategy strat = do

  -- Construct strategy map
  let kerns = runStrategy (void strat)

  -- Generate sorted kernel
  let kerns' = sortBy (comparing kernId) (stepsToKernels kerns)

  let dumpKern (KernelBind kid kfl kname repr deps _ _) = do
        putStrLn $ concat
          [ "Kernel ", show kid, ":", kname, " implementing ", flName kfl
          , " producing ", show repr, " using ", show deps ]
  forM_ kerns' dumpKern

dumpStrategyDOT :: FilePath -> Strategy a -> IO ()
dumpStrategyDOT file strat = do

  -- Construct strategy map
  let kerns = runStrategy (void strat)

  -- Generate sorted kernel
  let kerns' = sortBy (comparing kernId) (stepsToKernels kerns)

  -- Open output file
  h <- openFile file WriteMode
  hPutStrLn h "digraph strategy {"
  let kidName kid = "kernel" ++ show kid
  let dumpKern (KernelBind kid kfl kname repr deps _ _) = do
        hPutStrLn h $ concat
          [ kidName kid, " [label=\"" ++ kname, " implementing ", flName kfl,
            " producing ", show repr, "\"]"]
        forM_ (deps) $ \kid' ->
          hPutStrLn h $ concat
            [ kidName kid', " -> ", kidName kid]
  forM_ kerns' dumpKern
  hPutStrLn h "}"
  hClose h

dumpSteps :: Strategy a -> IO ()
dumpSteps strat = do

  let dump ind (DomainStep m_kid dh)
        = putStrLn $ ind ++ "Domain " ++ show dh ++ maybe "" (\kid -> " from kernel " ++ show kid) m_kid
      dump ind (SplitStep dh steps)
        = do putStrLn $ ind ++ "Split Domain " ++ show (dhId (fromJust $ dhParent dh)) ++ " into " ++ show dh
             forM_ steps (dump ("  " ++ ind))
      dump ind (KernelStep kb@KernelBind{kernRepr=ReprI rep})
        = putStrLn $ ind ++ "Over " ++ show (reprDomain rep) ++ " run " ++ show kb
      dump ind step@(DistributeStep did sched steps)
        = do putStrLn $ ind ++ "Distribute " ++ show did ++ " using " ++ show sched ++
                        " deps " ++ show (stepKernDeps step)
             forM_ steps (dump ("  " ++ ind))

  forM_ (runStrategy (void strat)) (dump "")

data AnyDH = forall a. Typeable a => AnyDH (Domain a)

adhFilterBox :: AnyDH -> RegionBox -> Region -> Maybe Region
adhFilterBox (AnyDH dom) = dhFilterBox dom

type DataMap = IM.IntMap (ReprI, RegionData)
type DomainMap = IM.IntMap (AnyDH, [Region])

dataMapInsert :: KernelId -> ReprI -> RegionData -> DataMap -> DataMap
dataMapInsert kid repr rdata = IM.insertWith update kid (repr, rdata)
  where update _ (repr', m) = (repr', Map.union rdata m)

dataMapUnion :: DataMap -> DataMap -> DataMap
dataMapUnion = IM.unionWith combine
  where combine (repr, bufs) (_, bufs') = (repr, bufs `Map.union` bufs')

dataMapUnions :: [DataMap] -> DataMap
dataMapUnions = IM.unionsWith combine
  where combine (repr, bufs) (_, bufs') = (repr, bufs `Map.union` bufs')

dataMapDifference :: DataMap -> DataMap -> DataMap
dataMapDifference = IM.differenceWith remove
  where remove (repr, bufs) (_, bufs')
           = if Map.null diff then Nothing else Just (repr, diff)
          where diff = bufs `Map.difference` bufs'

type KernelSet = IS.IntSet
type DomainSet = IS.IntSet

-- | Kernel dependencies of a step
stepKernDeps :: Step -> KernelSet
stepKernDeps (DomainStep (Just kid) _)  = IS.singleton kid
stepKernDeps (KernelStep kbind)         = IS.fromList $ map kdepId $ kernDeps kbind
stepKernDeps (SplitStep _ steps)        = stepsKernDeps steps
stepKernDeps (DistributeStep _ _ steps) = stepsKernDeps steps
stepKernDeps _                          = IS.empty

-- | Kernel dependencies of a series of steps
stepsKernDeps :: [Step] -> KernelSet
stepsKernDeps (step@(KernelStep kbind) : steps)
  = IS.delete (kernId kbind) $ stepsKernDeps steps `IS.union` stepKernDeps step
stepsKernDeps (step : steps)
  = stepKernDeps step `IS.union` stepsKernDeps steps
stepsKernDeps []
  = IS.empty

-- | Domain dependencies of a step
stepDomainDeps :: Step -> DomainSet
stepDomainDeps (KernelStep kbind)
  = IS.fromList $ concatMap kdepDomain $ kernDeps kbind
stepDomainDeps (SplitStep dh steps)
  = IS.insert (dhId (fromJust (dhParent dh))) $ IS.delete (dhId dh) $ stepsDomainDeps steps
stepDomainDeps (DistributeStep _ _ steps)
  = stepsDomainDeps steps
stepDomainDeps _
  = IS.empty

-- | Domain dependencies of a series of steps
stepsDomainDeps :: [Step] -> DomainSet
stepsDomainDeps (step@(DomainStep _ dh) : steps)
  = IS.delete (dhId dh) $ stepDomainDeps step `IS.union` stepsDomainDeps steps
stepsDomainDeps (step : steps)
  = stepDomainDeps step `IS.union` stepsDomainDeps steps
stepsDomainDeps []
  = IS.empty

execStrategy :: Strategy () -> IO ()
execStrategy strat = do

  -- Initialise maps
  dataMapRef <- newIORef IM.empty
  domainMapRef <- newIORef IM.empty

  -- Run steps given by strategy
  execSteps dataMapRef domainMapRef IS.empty $
    runStrategy strat
<commit_msg>Count return domains for domain dependencies<commit_after>{-# LANGUAGE ExistentialQuantification #-}
{-# LANGUAGE BangPatterns #-}

module Flow.Run
  ( dumpStrategy
  , dumpStrategyDOT
  , dumpSteps
  , execStrategy
  ) where

import Control.Monad
import Control.Applicative
import Control.Concurrent

import Data.List
import Data.Maybe ( fromMaybe, fromJust, mapMaybe )
import qualified Data.IntMap as IM
import qualified Data.IntSet as IS
import qualified Data.Map.Strict as Map
import Data.IORef
import Data.Ord
import Data.Time.Clock
import Data.Typeable

import System.IO

import Flow.Internal
import Flow.Builder
import Flow.Vector

dumpStrategy :: Strategy a -> IO ()
dumpStrategy strat = do

  -- Construct strategy map
  let kerns = runStrategy (void strat)

  -- Generate sorted kernel
  let kerns' = sortBy (comparing kernId) (stepsToKernels kerns)

  let dumpKern (KernelBind kid kfl kname repr deps _ _) = do
        putStrLn $ concat
          [ "Kernel ", show kid, ":", kname, " implementing ", flName kfl
          , " producing ", show repr, " using ", show deps ]
  forM_ kerns' dumpKern

dumpStrategyDOT :: FilePath -> Strategy a -> IO ()
dumpStrategyDOT file strat = do

  -- Construct strategy map
  let kerns = runStrategy (void strat)

  -- Generate sorted kernel
  let kerns' = sortBy (comparing kernId) (stepsToKernels kerns)

  -- Open output file
  h <- openFile file WriteMode
  hPutStrLn h "digraph strategy {"
  let kidName kid = "kernel" ++ show kid
  let dumpKern (KernelBind kid kfl kname repr deps _ _) = do
        hPutStrLn h $ concat
          [ kidName kid, " [label=\"" ++ kname, " implementing ", flName kfl,
            " producing ", show repr, "\"]"]
        forM_ (deps) $ \kid' ->
          hPutStrLn h $ concat
            [ kidName kid', " -> ", kidName kid]
  forM_ kerns' dumpKern
  hPutStrLn h "}"
  hClose h

dumpSteps :: Strategy a -> IO ()
dumpSteps strat = do

  let dump ind (DomainStep m_kid dh)
        = putStrLn $ ind ++ "Domain " ++ show dh ++ maybe "" (\kid -> " from kernel " ++ show kid) m_kid
      dump ind (SplitStep dh steps)
        = do putStrLn $ ind ++ "Split Domain " ++ show (dhId (fromJust $ dhParent dh)) ++ " into " ++ show dh
             forM_ steps (dump ("  " ++ ind))
      dump ind (KernelStep kb@KernelBind{kernRepr=ReprI rep})
        = putStrLn $ ind ++ "Over " ++ show (reprDomain rep) ++ " run " ++ show kb
      dump ind step@(DistributeStep did sched steps)
        = do putStrLn $ ind ++ "Distribute " ++ show did ++ " using " ++ show sched ++
                        " deps " ++ show (stepKernDeps step)
             forM_ steps (dump ("  " ++ ind))

  forM_ (runStrategy (void strat)) (dump "")

data AnyDH = forall a. Typeable a => AnyDH (Domain a)

adhFilterBox :: AnyDH -> RegionBox -> Region -> Maybe Region
adhFilterBox (AnyDH dom) = dhFilterBox dom

type DataMap = IM.IntMap (ReprI, RegionData)
type DomainMap = IM.IntMap (AnyDH, [Region])

dataMapInsert :: KernelId -> ReprI -> RegionData -> DataMap -> DataMap
dataMapInsert kid repr rdata = IM.insertWith update kid (repr, rdata)
  where update _ (repr', m) = (repr', Map.union rdata m)

dataMapUnion :: DataMap -> DataMap -> DataMap
dataMapUnion = IM.unionWith combine
  where combine (repr, bufs) (_, bufs') = (repr, bufs `Map.union` bufs')

dataMapUnions :: [DataMap] -> DataMap
dataMapUnions = IM.unionsWith combine
  where combine (repr, bufs) (_, bufs') = (repr, bufs `Map.union` bufs')

dataMapDifference :: DataMap -> DataMap -> DataMap
dataMapDifference = IM.differenceWith remove
  where remove (repr, bufs) (_, bufs')
           = if Map.null diff then Nothing else Just (repr, diff)
          where diff = bufs `Map.difference` bufs'

type KernelSet = IS.IntSet
type DomainSet = IS.IntSet

-- | Kernel dependencies of a step
stepKernDeps :: Step -> KernelSet
stepKernDeps (DomainStep (Just kid) _)  = IS.singleton kid
stepKernDeps (KernelStep kbind)         = IS.fromList $ map kdepId $ kernDeps kbind
stepKernDeps (SplitStep _ steps)        = stepsKernDeps steps
stepKernDeps (DistributeStep _ _ steps) = stepsKernDeps steps
stepKernDeps _                          = IS.empty

-- | Kernel dependencies of a series of steps
stepsKernDeps :: [Step] -> KernelSet
stepsKernDeps (step@(KernelStep kbind) : steps)
  = IS.delete (kernId kbind) $ stepsKernDeps steps `IS.union` stepKernDeps step
stepsKernDeps (step : steps)
  = stepKernDeps step `IS.union` stepsKernDeps steps
stepsKernDeps []
  = IS.empty

-- | Domain dependencies of a step
stepDomainDeps :: Step -> DomainSet
stepDomainDeps (KernelStep kbind)
  = (IS.fromList $ kernDomain kbind) `IS.union`
    (IS.fromList $ concatMap kdepDomain $ kernDeps kbind)
stepDomainDeps (SplitStep dh steps)
  = IS.insert (dhId (fromJust (dhParent dh))) $ IS.delete (dhId dh) $ stepsDomainDeps steps
stepDomainDeps (DistributeStep _ _ steps)
  = stepsDomainDeps steps
stepDomainDeps _
  = IS.empty

-- | Domain dependencies of a series of steps
stepsDomainDeps :: [Step] -> DomainSet
stepsDomainDeps (step@(DomainStep _ dh) : steps)
  = IS.delete (dhId dh) $ stepDomainDeps step `IS.union` stepsDomainDeps steps
stepsDomainDeps (step : steps)
  = stepDomainDeps step `IS.union` stepsDomainDeps steps
stepsDomainDeps []
  = IS.empty

execStrategy :: Strategy () -> IO ()
execStrategy strat = do

  -- Initialise maps
  dataMapRef <- newIORef IM.empty
  domainMapRef <- newIORef IM.empty

  -- Run steps given by strategy
  execSteps dataMapRef domainMapRef IS.empty $
    runStrategy strat

{'fraction_non_alphanumeric': 0.10609812012838148, 'fraction_numerical': 0.0, 'mean_word_length': 3.2953918865695155, 'pattern_counts': {'":': 2, '<': 9, '<?xml version=': 0, '>': 53, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 2, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '3676554', 'n_tokens_mistral': 3647, 'n_tokens_neox': 3509, 'n_words': 1121}