-- | @futhark autotune@
module Futhark.CLI.Autotune (main) where

import Control.Exception
import Control.Monad
import Data.ByteString.Char8 qualified as SBS
import Data.Function (on)
import Data.IORef
import Data.List (elemIndex, intersect, minimumBy, sort, sortOn)
import Data.Map qualified as M
import Data.Maybe
import Data.Set qualified as S
import Data.Text qualified as T
import Data.Text.IO qualified as T
import Data.Tree
import Futhark.Bench
import Futhark.Server
import Futhark.Test
import Futhark.Util (fancyTerminal, maxinum, showText)
import Futhark.Util.Options
import Futhark.Util.ProgressBar
import System.Directory
import System.Environment (getExecutablePath)
import System.Exit
import System.FilePath
import System.IO
import Text.Read (readMaybe)
import Text.Regex.TDFA

data AutotuneOptions = AutotuneOptions
  { AutotuneOptions -> String
optBackend :: String,
    AutotuneOptions -> Maybe String
optFuthark :: Maybe String,
    AutotuneOptions -> Int
optMinRuns :: Int,
    AutotuneOptions -> Maybe String
optTuning :: Maybe String,
    AutotuneOptions -> [String]
optExtraOptions :: [String],
    AutotuneOptions -> Int
optVerbose :: Int,
    AutotuneOptions -> Int
optTimeout :: Int,
    AutotuneOptions -> Bool
optSkipCompilation :: Bool,
    AutotuneOptions -> Int
optDefaultThreshold :: Int,
    AutotuneOptions -> Maybe String
optTestSpec :: Maybe FilePath
  }

initialAutotuneOptions :: AutotuneOptions
initialAutotuneOptions :: AutotuneOptions
initialAutotuneOptions =
  AutotuneOptions
    { optBackend :: String
optBackend = String
"opencl",
      optFuthark :: Maybe String
optFuthark = Maybe String
forall a. Maybe a
Nothing,
      optMinRuns :: Int
optMinRuns = Int
10,
      optTuning :: Maybe String
optTuning = String -> Maybe String
forall a. a -> Maybe a
Just String
"tuning",
      optExtraOptions :: [String]
optExtraOptions = [],
      optVerbose :: Int
optVerbose = Int
0,
      optTimeout :: Int
optTimeout = Int
600,
      optSkipCompilation :: Bool
optSkipCompilation = Bool
False,
      optDefaultThreshold :: Int
optDefaultThreshold = Int
thresholdMax,
      optTestSpec :: Maybe String
optTestSpec = Maybe String
forall a. Maybe a
Nothing
    }

compileOptions :: AutotuneOptions -> IO CompileOptions
compileOptions :: AutotuneOptions -> IO CompileOptions
compileOptions AutotuneOptions
opts = do
  futhark <- IO String -> (String -> IO String) -> Maybe String -> IO String
forall b a. b -> (a -> b) -> Maybe a -> b
maybe IO String
getExecutablePath String -> IO String
forall a. a -> IO a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (Maybe String -> IO String) -> Maybe String -> IO String
forall a b. (a -> b) -> a -> b
$ AutotuneOptions -> Maybe String
optFuthark AutotuneOptions
opts
  pure $
    CompileOptions
      { compFuthark = futhark,
        compBackend = optBackend opts,
        compOptions = mempty
      }

runOptions :: Int -> AutotuneOptions -> RunOptions
runOptions :: Int -> AutotuneOptions -> RunOptions
runOptions Int
timeout_s AutotuneOptions
opts =
  RunOptions
    { runMinRuns :: Int
runMinRuns = AutotuneOptions -> Int
optMinRuns AutotuneOptions
opts,
      runMinTime :: NominalDiffTime
runMinTime = NominalDiffTime
0.5,
      runTimeout :: Int
runTimeout = Int
timeout_s,
      runVerbose :: Int
runVerbose = AutotuneOptions -> Int
optVerbose AutotuneOptions
opts,
      runConvergencePhase :: Bool
runConvergencePhase = Bool
True,
      runConvergenceMaxTime :: NominalDiffTime
runConvergenceMaxTime = Int -> NominalDiffTime
forall a b. (Integral a, Num b) => a -> b
fromIntegral Int
timeout_s,
      runResultAction :: (Int, Maybe Double) -> IO ()
runResultAction = IO () -> (Int, Maybe Double) -> IO ()
forall a b. a -> b -> a
const (IO () -> (Int, Maybe Double) -> IO ())
-> IO () -> (Int, Maybe Double) -> IO ()
forall a b. (a -> b) -> a -> b
$ () -> IO ()
forall a. a -> IO a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (),
      runProfile :: Bool
runProfile = Bool
False
    }

type Path = [(T.Text, Int)]

regexBlocks :: Regex -> T.Text -> Maybe [T.Text]
regexBlocks :: Regex -> Text -> Maybe [Text]
regexBlocks Regex
regex Text
s = do
  (_, _, _, groups) <-
    Regex -> Text -> Maybe (Text, Text, Text, [Text])
forall regex source target (m :: * -> *).
(RegexContext regex source target, MonadFail m) =>
regex -> source -> m target
forall (m :: * -> *).
MonadFail m =>
Regex -> Text -> m (Text, Text, Text, [Text])
matchM Regex
regex Text
s :: Maybe (T.Text, T.Text, T.Text, [T.Text])
  Just groups

comparisons :: T.Text -> [(T.Text, Int)]
comparisons :: Text -> [(Text, Int)]
comparisons = (Text -> Maybe (Text, Int)) -> [Text] -> [(Text, Int)]
forall a b. (a -> Maybe b) -> [a] -> [b]
mapMaybe Text -> Maybe (Text, Int)
forall {b}. Read b => Text -> Maybe (Text, b)
isComparison ([Text] -> [(Text, Int)])
-> (Text -> [Text]) -> Text -> [(Text, Int)]
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Text -> [Text]
T.lines
  where
    regex :: Regex
regex = String -> Regex
forall regex compOpt execOpt source.
RegexMaker regex compOpt execOpt source =>
source -> regex
makeRegex (String
"Compared ([^ ]+) <= (-?[0-9]+)" :: String)
    isComparison :: Text -> Maybe (Text, b)
isComparison Text
l = do
      [thresh, val] <- Regex -> Text -> Maybe [Text]
regexBlocks Regex
regex Text
l
      val' <- readMaybe $ T.unpack val
      pure (thresh, val')

type RunDataset = Server -> Int -> Path -> IO (Either String ([(T.Text, Int)], Int))

type DatasetName = T.Text

serverOptions :: AutotuneOptions -> [String]
serverOptions :: AutotuneOptions -> [String]
serverOptions AutotuneOptions
opts =
  String
"--default-threshold"
    String -> [String] -> [String]
forall a. a -> [a] -> [a]
: Int -> String
forall a. Show a => a -> String
show (AutotuneOptions -> Int
optDefaultThreshold AutotuneOptions
opts)
    String -> [String] -> [String]
forall a. a -> [a] -> [a]
: String
"-L"
    String -> [String] -> [String]
forall a. a -> [a] -> [a]
: AutotuneOptions -> [String]
optExtraOptions AutotuneOptions
opts

checkCmd :: Either CmdFailure a -> IO a
checkCmd :: forall a. Either CmdFailure a -> IO a
checkCmd = (CmdFailure -> IO a) -> (a -> IO a) -> Either CmdFailure a -> IO a
forall a c b. (a -> c) -> (b -> c) -> Either a b -> c
either (String -> IO a
forall a. HasCallStack => String -> a
error (String -> IO a) -> (CmdFailure -> String) -> CmdFailure -> IO a
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Text -> String
T.unpack (Text -> String) -> (CmdFailure -> Text) -> CmdFailure -> String
forall b c a. (b -> c) -> (a -> b) -> a -> c
. [Text] -> Text
T.unlines ([Text] -> Text) -> (CmdFailure -> [Text]) -> CmdFailure -> Text
forall b c a. (b -> c) -> (a -> b) -> a -> c
. CmdFailure -> [Text]
failureMsg) a -> IO a
forall a. a -> IO a
forall (f :: * -> *) a. Applicative f => a -> f a
pure

setTuningParam :: Server -> T.Text -> Int -> IO ()
setTuningParam :: Server -> Text -> Int -> IO ()
setTuningParam Server
server Text
name Int
val =
  IO () -> IO ()
forall (f :: * -> *) a. Functor f => f a -> f ()
void (IO () -> IO ()) -> IO () -> IO ()
forall a b. (a -> b) -> a -> b
$ Either CmdFailure () -> IO ()
forall a. Either CmdFailure a -> IO a
checkCmd (Either CmdFailure () -> IO ())
-> (Maybe CmdFailure -> Either CmdFailure ())
-> Maybe CmdFailure
-> IO ()
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Either CmdFailure ()
-> (CmdFailure -> Either CmdFailure ())
-> Maybe CmdFailure
-> Either CmdFailure ()
forall b a. b -> (a -> b) -> Maybe a -> b
maybe (() -> Either CmdFailure ()
forall a b. b -> Either a b
Right ()) CmdFailure -> Either CmdFailure ()
forall a b. a -> Either a b
Left (Maybe CmdFailure -> IO ()) -> IO (Maybe CmdFailure) -> IO ()
forall (m :: * -> *) a b. Monad m => (a -> m b) -> m a -> m b
=<< Server -> Text -> Int -> IO (Maybe CmdFailure)
cmdSetTuningParam Server
server Text
name Int
val

setTuningParams :: Server -> Path -> IO ()
setTuningParams :: Server -> [(Text, Int)] -> IO ()
setTuningParams Server
server = ((Text, Int) -> IO ()) -> [(Text, Int)] -> IO ()
forall (t :: * -> *) (m :: * -> *) a b.
(Foldable t, Monad m) =>
(a -> m b) -> t a -> m ()
mapM_ ((Text -> Int -> IO ()) -> (Text, Int) -> IO ()
forall a b c. (a -> b -> c) -> (a, b) -> c
uncurry ((Text -> Int -> IO ()) -> (Text, Int) -> IO ())
-> (Text -> Int -> IO ()) -> (Text, Int) -> IO ()
forall a b. (a -> b) -> a -> b
$ Server -> Text -> Int -> IO ()
setTuningParam Server
server)

restoreTuningParams :: AutotuneOptions -> Server -> Path -> IO ()
restoreTuningParams :: AutotuneOptions -> Server -> [(Text, Int)] -> IO ()
restoreTuningParams AutotuneOptions
opts Server
server = ((Text, Int) -> IO ()) -> [(Text, Int)] -> IO ()
forall (t :: * -> *) (m :: * -> *) a b.
(Foldable t, Monad m) =>
(a -> m b) -> t a -> m ()
mapM_ (Text, Int) -> IO ()
forall {b}. (Text, b) -> IO ()
opt
  where
    opt :: (Text, b) -> IO ()
opt (Text
name, b
_) = Server -> Text -> Int -> IO ()
setTuningParam Server
server Text
name (AutotuneOptions -> Int
optDefaultThreshold AutotuneOptions
opts)

prepare :: AutotuneOptions -> FutharkExe -> FilePath -> IO [(DatasetName, RunDataset, T.Text)]
prepare :: AutotuneOptions
-> FutharkExe -> String -> IO [(Text, RunDataset, Text)]
prepare AutotuneOptions
opts FutharkExe
futhark String
prog = do
  spec <-
    IO ProgramTest
-> (String -> IO ProgramTest) -> Maybe String -> IO ProgramTest
forall b a. b -> (a -> b) -> Maybe a -> b
maybe (String -> IO ProgramTest
testSpecFromProgramOrDie String
prog) String -> IO ProgramTest
testSpecFromFileOrDie (Maybe String -> IO ProgramTest) -> Maybe String -> IO ProgramTest
forall a b. (a -> b) -> a -> b
$
      AutotuneOptions -> Maybe String
optTestSpec AutotuneOptions
opts
  copts <- compileOptions opts

  truns <-
    case testAction spec of
      RunCases [InputOutputs]
ios [StructureTest]
_ [WarningTest]
_ [PropertyCase]
_ | Bool -> Bool
not (Bool -> Bool) -> Bool -> Bool
forall a b. (a -> b) -> a -> b
$ [InputOutputs] -> Bool
forall a. [a] -> Bool
forall (t :: * -> *) a. Foldable t => t a -> Bool
null [InputOutputs]
ios -> do
        Bool -> IO () -> IO ()
forall (f :: * -> *). Applicative f => Bool -> f () -> f ()
when (AutotuneOptions -> Int
optVerbose AutotuneOptions
opts Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
> Int
1) (IO () -> IO ()) -> IO () -> IO ()
forall a b. (a -> b) -> a -> b
$
          String -> IO ()
putStrLn (String -> IO ()) -> String -> IO ()
forall a b. (a -> b) -> a -> b
$
            [String] -> String
unwords (String
"Entry points:" String -> [String] -> [String]
forall a. a -> [a] -> [a]
: (InputOutputs -> String) -> [InputOutputs] -> [String]
forall a b. (a -> b) -> [a] -> [b]
map (Text -> String
T.unpack (Text -> String)
-> (InputOutputs -> Text) -> InputOutputs -> String
forall b c a. (b -> c) -> (a -> b) -> a -> c
. InputOutputs -> Text
iosEntryPoint) [InputOutputs]
ios)

        if AutotuneOptions -> Bool
optSkipCompilation AutotuneOptions
opts
          then do
            exists <- String -> IO Bool
doesFileExist (String -> IO Bool) -> String -> IO Bool
forall a b. (a -> b) -> a -> b
$ String -> String
binaryName String
prog
            if exists
              then pure ios
              else do
                putStrLn $ binaryName prog ++ " does not exist, but --skip-compilation passed."
                exitFailure
          else do
            res <- Maybe Int
-> CompileOptions
-> String
-> [InputOutputs]
-> IO (Either (String, Maybe ByteString) ())
forall (m :: * -> *).
MonadIO m =>
Maybe Int
-> CompileOptions
-> String
-> [InputOutputs]
-> m (Either (String, Maybe ByteString) ())
prepareBenchmarkProgram Maybe Int
forall a. Maybe a
Nothing CompileOptions
copts String
prog [InputOutputs]
ios
            case res of
              Left (String
err, Maybe ByteString
errstr) -> do
                String -> IO ()
putStrLn String
err
                IO () -> (ByteString -> IO ()) -> Maybe ByteString -> IO ()
forall b a. b -> (a -> b) -> Maybe a -> b
maybe (() -> IO ()
forall a. a -> IO a
forall (f :: * -> *) a. Applicative f => a -> f a
pure ()) ByteString -> IO ()
SBS.putStrLn Maybe ByteString
errstr
                IO [InputOutputs]
forall a. IO a
exitFailure
              Right () ->
                [InputOutputs] -> IO [InputOutputs]
forall a. a -> IO a
forall (f :: * -> *) a. Applicative f => a -> f a
pure [InputOutputs]
ios
      TestAction
_ ->
        String -> IO [InputOutputs]
forall a. String -> IO a
forall (m :: * -> *) a. MonadFail m => String -> m a
fail String
"Unsupported test spec."

  let runnableDataset Text
entry_point TestRun
trun =
        case TestRun -> ExpectedResult Success
runExpectedResult TestRun
trun of
          Succeeds Maybe Success
expected
            | [Text] -> Bool
forall a. [a] -> Bool
forall (t :: * -> *) a. Foldable t => t a -> Bool
null (TestRun -> [Text]
runTags TestRun
trun [Text] -> [Text] -> [Text]
forall a. Eq a => [a] -> [a] -> [a]
`intersect` [Text
"notune", Text
"disable"]) ->
                (Text, RunDataset) -> Maybe (Text, RunDataset)
forall a. a -> Maybe a
Just
                  ( TestRun -> Text
runDescription TestRun
trun,
                    \Server
server -> Server
-> Text
-> TestRun
-> Maybe Success
-> Int
-> [(Text, Int)]
-> IO (Either String ([(Text, Int)], Int))
run Server
server Text
entry_point TestRun
trun Maybe Success
expected
                  )
          ExpectedResult Success
_ -> Maybe (Text, RunDataset)
forall a. Maybe a
Nothing

  fmap concat . forM truns $ \InputOutputs
ios -> do
    let cases :: [(Text, RunDataset)]
cases =
          (TestRun -> Maybe (Text, RunDataset))
-> [TestRun] -> [(Text, RunDataset)]
forall a b. (a -> Maybe b) -> [a] -> [b]
mapMaybe (Text -> TestRun -> Maybe (Text, RunDataset)
runnableDataset (Text -> TestRun -> Maybe (Text, RunDataset))
-> Text -> TestRun -> Maybe (Text, RunDataset)
forall a b. (a -> b) -> a -> b
$ InputOutputs -> Text
iosEntryPoint InputOutputs
ios) (InputOutputs -> [TestRun]
iosTestRuns InputOutputs
ios)
    [(Text, RunDataset)]
-> ((Text, RunDataset) -> IO (Text, RunDataset, Text))
-> IO [(Text, RunDataset, Text)]
forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
t a -> (a -> m b) -> m (t b)
forM [(Text, RunDataset)]
cases (((Text, RunDataset) -> IO (Text, RunDataset, Text))
 -> IO [(Text, RunDataset, Text)])
-> ((Text, RunDataset) -> IO (Text, RunDataset, Text))
-> IO [(Text, RunDataset, Text)]
forall a b. (a -> b) -> a -> b
$ \(Text
dataset, RunDataset
do_run) ->
      (Text, RunDataset, Text) -> IO (Text, RunDataset, Text)
forall a. a -> IO a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (Text
dataset, RunDataset
do_run, InputOutputs -> Text
iosEntryPoint InputOutputs
ios)
  where
    run :: Server
-> Text
-> TestRun
-> Maybe Success
-> Int
-> [(Text, Int)]
-> IO (Either String ([(Text, Int)], Int))
run Server
server Text
entry_point TestRun
trun Maybe Success
expected Int
timeout [(Text, Int)]
path = do
      let bestRuntime :: ([RunResult], Text, c) -> ([(Text, Int)], Int)
bestRuntime ([RunResult]
runres, Text
errout, c
_) =
            ( Text -> [(Text, Int)]
comparisons Text
errout,
              [Int] -> Int
forall a. Ord a => [a] -> a
forall (t :: * -> *) a. (Foldable t, Ord a) => t a -> a
minimum ([Int] -> Int) -> [Int] -> Int
forall a b. (a -> b) -> a -> b
$ (RunResult -> Int) -> [RunResult] -> [Int]
forall a b. (a -> b) -> [a] -> [b]
map RunResult -> Int
runMicroseconds [RunResult]
runres
            )

          ropts :: RunOptions
ropts = Int -> AutotuneOptions -> RunOptions
runOptions Int
timeout AutotuneOptions
opts

      Bool -> IO () -> IO ()
forall (f :: * -> *). Applicative f => Bool -> f () -> f ()
when (AutotuneOptions -> Int
optVerbose AutotuneOptions
opts Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
> Int
1) (IO () -> IO ()) -> IO () -> IO ()
forall a b. (a -> b) -> a -> b
$
        String -> IO ()
putStrLn (String
"Trying path: " String -> String -> String
forall a. [a] -> [a] -> [a]
++ [(Text, Int)] -> String
forall a. Show a => a -> String
show [(Text, Int)]
path)

      -- Setting the tuning parameters is a stateful action, so we
      -- must be careful to restore the defaults below.  This is
      -- because we rely on parameters not in 'path' to have their
      -- default value.
      Server -> [(Text, Int)] -> IO ()
setTuningParams Server
server [(Text, Int)]
path

      (Text -> Either String ([(Text, Int)], Int))
-> (([RunResult], Text, ProfilingReport)
    -> Either String ([(Text, Int)], Int))
-> Either Text ([RunResult], Text, ProfilingReport)
-> Either String ([(Text, Int)], Int)
forall a c b. (a -> c) -> (b -> c) -> Either a b -> c
either (String -> Either String ([(Text, Int)], Int)
forall a b. a -> Either a b
Left (String -> Either String ([(Text, Int)], Int))
-> (Text -> String) -> Text -> Either String ([(Text, Int)], Int)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Text -> String
T.unpack) (([(Text, Int)], Int) -> Either String ([(Text, Int)], Int)
forall a b. b -> Either a b
Right (([(Text, Int)], Int) -> Either String ([(Text, Int)], Int))
-> (([RunResult], Text, ProfilingReport) -> ([(Text, Int)], Int))
-> ([RunResult], Text, ProfilingReport)
-> Either String ([(Text, Int)], Int)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. ([RunResult], Text, ProfilingReport) -> ([(Text, Int)], Int)
forall {c}. ([RunResult], Text, c) -> ([(Text, Int)], Int)
bestRuntime)
        (Either Text ([RunResult], Text, ProfilingReport)
 -> Either String ([(Text, Int)], Int))
-> IO (Either Text ([RunResult], Text, ProfilingReport))
-> IO (Either String ([(Text, Int)], Int))
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Server
-> RunOptions
-> FutharkExe
-> String
-> Text
-> Values
-> Maybe Success
-> String
-> IO (Either Text ([RunResult], Text, ProfilingReport))
benchmarkDataset
          Server
server
          RunOptions
ropts
          FutharkExe
futhark
          String
prog
          Text
entry_point
          (TestRun -> Values
runInput TestRun
trun)
          Maybe Success
expected
          (String -> Text -> TestRun -> String
testRunReferenceOutput String
prog Text
entry_point TestRun
trun)
        IO (Either String ([(Text, Int)], Int))
-> IO () -> IO (Either String ([(Text, Int)], Int))
forall a b. IO a -> IO b -> IO a
forall (f :: * -> *) a b. Applicative f => f a -> f b -> f a
<* AutotuneOptions -> Server -> [(Text, Int)] -> IO ()
restoreTuningParams AutotuneOptions
opts Server
server [(Text, Int)]
path

--- Benchmarking a program

data DatasetResult = DatasetResult [(T.Text, Int)] Double
  deriving (Int -> DatasetResult -> String -> String
[DatasetResult] -> String -> String
DatasetResult -> String
(Int -> DatasetResult -> String -> String)
-> (DatasetResult -> String)
-> ([DatasetResult] -> String -> String)
-> Show DatasetResult
forall a.
(Int -> a -> String -> String)
-> (a -> String) -> ([a] -> String -> String) -> Show a
$cshowsPrec :: Int -> DatasetResult -> String -> String
showsPrec :: Int -> DatasetResult -> String -> String
$cshow :: DatasetResult -> String
show :: DatasetResult -> String
$cshowList :: [DatasetResult] -> String -> String
showList :: [DatasetResult] -> String -> String
Show)

--- Finding initial comparisons.

--- Extracting threshold hierarchy.

type ThresholdForest = Forest (T.Text, Bool)

thresholdMin, thresholdMax :: Int
thresholdMin :: Int
thresholdMin = Int
1
thresholdMax :: Int
thresholdMax = Int
2000000000

-- | Depth-first list of thresholds to tune in order, and a
-- corresponding assignment of ancestor thresholds to ensure that they
-- are used.
tuningPaths :: ThresholdForest -> [(T.Text, Path)]
tuningPaths :: ThresholdForest -> [(Text, [(Text, Int)])]
tuningPaths = (Tree (Text, Bool) -> [(Text, [(Text, Int)])])
-> ThresholdForest -> [(Text, [(Text, Int)])]
forall (t :: * -> *) a b. Foldable t => (a -> [b]) -> t a -> [b]
concatMap ([(Text, Int)] -> Tree (Text, Bool) -> [(Text, [(Text, Int)])]
forall {a}. [(a, Int)] -> Tree (a, Bool) -> [(a, [(a, Int)])]
treePaths [])
  where
    treePaths :: [(a, Int)] -> Tree (a, Bool) -> [(a, [(a, Int)])]
treePaths [(a, Int)]
ancestors (Node (a
v, Bool
_) [Tree (a, Bool)]
children) =
      (Tree (a, Bool) -> [(a, [(a, Int)])])
-> [Tree (a, Bool)] -> [(a, [(a, Int)])]
forall (t :: * -> *) a b. Foldable t => (a -> [b]) -> t a -> [b]
concatMap ([(a, Int)] -> a -> Tree (a, Bool) -> [(a, [(a, Int)])]
onChild [(a, Int)]
ancestors a
v) [Tree (a, Bool)]
children [(a, [(a, Int)])] -> [(a, [(a, Int)])] -> [(a, [(a, Int)])]
forall a. [a] -> [a] -> [a]
++ [(a
v, [(a, Int)]
ancestors)]

    onChild :: [(a, Int)] -> a -> Tree (a, Bool) -> [(a, [(a, Int)])]
onChild [(a, Int)]
ancestors a
v child :: Tree (a, Bool)
child@(Node (a
_, Bool
cmp) [Tree (a, Bool)]
_) =
      [(a, Int)] -> Tree (a, Bool) -> [(a, [(a, Int)])]
treePaths ([(a, Int)]
ancestors [(a, Int)] -> [(a, Int)] -> [(a, Int)]
forall a. [a] -> [a] -> [a]
++ [(a
v, Bool -> Int
t Bool
cmp)]) Tree (a, Bool)
child

    t :: Bool -> Int
t Bool
False = Int
thresholdMax
    t Bool
True = Int
thresholdMin

allTuningParams :: Server -> IO [(T.Text, T.Text)]
allTuningParams :: Server -> IO [(Text, Text)]
allTuningParams Server
server = do
  entry_points <- Either CmdFailure [Text] -> IO [Text]
forall a. Either CmdFailure a -> IO a
checkCmd (Either CmdFailure [Text] -> IO [Text])
-> IO (Either CmdFailure [Text]) -> IO [Text]
forall (m :: * -> *) a b. Monad m => (a -> m b) -> m a -> m b
=<< Server -> IO (Either CmdFailure [Text])
cmdEntryPoints Server
server
  param_names <- concat <$> mapM (checkCmd <=< cmdTuningParams server) entry_points
  param_classes <- mapM (checkCmd <=< cmdTuningParamClass server) param_names
  pure $ zip param_names param_classes

thresholdForest :: Server -> IO ThresholdForest
thresholdForest :: Server -> IO ThresholdForest
thresholdForest Server
server = do
  thresholds <- ((Text, Text) -> Maybe (Text, [(Text, Bool)]))
-> [(Text, Text)] -> [(Text, [(Text, Bool)])]
forall a b. (a -> Maybe b) -> [a] -> [b]
mapMaybe (Text, Text) -> Maybe (Text, [(Text, Bool)])
findThreshold ([(Text, Text)] -> [(Text, [(Text, Bool)])])
-> IO [(Text, Text)] -> IO [(Text, [(Text, Bool)])]
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Server -> IO [(Text, Text)]
allTuningParams Server
server
  let root (a
v, b
_) = ((a
v, Bool
False), [])
  pure $
    unfoldForest (unfold thresholds) $
      map root $
        filter (null . snd) thresholds
  where
    regex :: Regex
regex = Text -> Regex
forall regex compOpt execOpt source.
RegexMaker regex compOpt execOpt source =>
source -> regex
makeRegex (Text
"threshold\\(([^ ]+,)(.*)\\)" :: T.Text)

    findThreshold :: (T.Text, T.Text) -> Maybe (T.Text, [(T.Text, Bool)])
    findThreshold :: (Text, Text) -> Maybe (Text, [(Text, Bool)])
findThreshold (Text
name, Text
param_class) = do
      [_, grp] <- Regex -> Text -> Maybe [Text]
regexBlocks Regex
regex Text
param_class
      pure
        ( name,
          filter (not . T.null . fst)
            $ map
              ( \Text
x ->
                  if Text
"!" Text -> Text -> Bool
`T.isPrefixOf` Text
x
                    then (Int -> Text -> Text
T.drop Int
1 Text
x, Bool
False)
                    else (Text
x, Bool
True)
              )
            $ T.words grp
        )

    unfold :: [(a, [(a, b)])] -> ((a, b), [a]) -> ((a, b), [((a, b), [a])])
unfold [(a, [(a, b)])]
thresholds ((a
parent, b
parent_cmp), [a]
ancestors) =
      let ancestors' :: [a]
ancestors' = a
parent a -> [a] -> [a]
forall a. a -> [a] -> [a]
: [a]
ancestors

          isChild :: (a, [(a, b)]) -> Maybe ((a, b), [a])
isChild (a
v, [(a, b)]
v_ancestors) = do
            cmp <- a -> [(a, b)] -> Maybe b
forall a b. Eq a => a -> [(a, b)] -> Maybe b
lookup a
parent [(a, b)]
v_ancestors
            guard $
              sort (map fst v_ancestors)
                == sort (parent : ancestors)
            pure ((v, cmp), ancestors')
       in ((a
parent, b
parent_cmp), ((a, [(a, b)]) -> Maybe ((a, b), [a]))
-> [(a, [(a, b)])] -> [((a, b), [a])]
forall a b. (a -> Maybe b) -> [a] -> [b]
mapMaybe (a, [(a, b)]) -> Maybe ((a, b), [a])
forall {a} {b}. (a, [(a, b)]) -> Maybe ((a, b), [a])
isChild [(a, [(a, b)])]
thresholds)

-- | The performance difference in percentage that triggers a non-monotonicity
-- warning. This is to account for slight variantions in run-time.
epsilon :: Double
epsilon :: Double
epsilon = Double
1.02

--- Doing the atual tuning

tuneThreshold ::
  AutotuneOptions ->
  Server ->
  [(DatasetName, RunDataset, T.Text)] ->
  (Path, M.Map DatasetName Int) ->
  (T.Text, Path) ->
  IO (Path, M.Map DatasetName Int)
tuneThreshold :: AutotuneOptions
-> Server
-> [(Text, RunDataset, Text)]
-> ([(Text, Int)], Map Text Int)
-> (Text, [(Text, Int)])
-> IO ([(Text, Int)], Map Text Int)
tuneThreshold AutotuneOptions
opts Server
server [(Text, RunDataset, Text)]
datasets ([(Text, Int)]
already_tuned, Map Text Int
best_runtimes0) (Text
v, [(Text, Int)]
_v_path) = do
  (tune_result, best_runtimes) <-
    ((Maybe (Int, Int), Map Text Int)
 -> (Text, RunDataset, Text) -> IO (Maybe (Int, Int), Map Text Int))
-> (Maybe (Int, Int), Map Text Int)
-> [(Text, RunDataset, Text)]
-> IO (Maybe (Int, Int), Map Text Int)
forall (t :: * -> *) (m :: * -> *) b a.
(Foldable t, Monad m) =>
(b -> a -> m b) -> b -> t a -> m b
foldM (Maybe (Int, Int), Map Text Int)
-> (Text, RunDataset, Text) -> IO (Maybe (Int, Int), Map Text Int)
tuneDataset (Maybe (Int, Int)
forall a. Maybe a
Nothing, Map Text Int
best_runtimes0) [(Text, RunDataset, Text)]
datasets
  case tune_result of
    Maybe (Int, Int)
Nothing ->
      ([(Text, Int)], Map Text Int) -> IO ([(Text, Int)], Map Text Int)
forall a. a -> IO a
forall (f :: * -> *) a. Applicative f => a -> f a
pure ((Text
v, Int
thresholdMin) (Text, Int) -> [(Text, Int)] -> [(Text, Int)]
forall a. a -> [a] -> [a]
: [(Text, Int)]
already_tuned, Map Text Int
best_runtimes)
    Just (Int
_, Int
threshold) ->
      ([(Text, Int)], Map Text Int) -> IO ([(Text, Int)], Map Text Int)
forall a. a -> IO a
forall (f :: * -> *) a. Applicative f => a -> f a
pure ((Text
v, Int
threshold) (Text, Int) -> [(Text, Int)] -> [(Text, Int)]
forall a. a -> [a] -> [a]
: [(Text, Int)]
already_tuned, Map Text Int
best_runtimes)
  where
    tuneDataset ::
      (Maybe (Int, Int), M.Map DatasetName Int) ->
      (DatasetName, RunDataset, T.Text) ->
      IO (Maybe (Int, Int), M.Map DatasetName Int)
    tuneDataset :: (Maybe (Int, Int), Map Text Int)
-> (Text, RunDataset, Text) -> IO (Maybe (Int, Int), Map Text Int)
tuneDataset (Maybe (Int, Int)
thresholds, Map Text Int
best_runtimes) (Text
dataset_name, RunDataset
run, Text
entry_point) = do
      relevant <- Either CmdFailure [Text] -> IO [Text]
forall a. Either CmdFailure a -> IO a
checkCmd (Either CmdFailure [Text] -> IO [Text])
-> IO (Either CmdFailure [Text]) -> IO [Text]
forall (m :: * -> *) a b. Monad m => (a -> m b) -> m a -> m b
=<< Server -> Text -> IO (Either CmdFailure [Text])
cmdTuningParams Server
server Text
entry_point
      if v `notElem` relevant
        then do
          when (optVerbose opts > 0) $
            T.putStrLn $
              T.unwords [v, "is irrelevant for", entry_point]
          pure (thresholds, best_runtimes)
        else do
          sample_run <-
            run
              server
              (optTimeout opts)
              ((v, maybe thresholdMax snd thresholds) : already_tuned)

          case sample_run of
            Left String
err -> do
              -- If the sampling run fails, we treat it as zero information.
              -- One of our ancestor thresholds will have be set such that
              -- this path is never taken.
              Bool -> IO () -> IO ()
forall (f :: * -> *). Applicative f => Bool -> f () -> f ()
when (AutotuneOptions -> Int
optVerbose AutotuneOptions
opts Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
> Int
0) (IO () -> IO ()) -> IO () -> IO ()
forall a b. (a -> b) -> a -> b
$
                String -> IO ()
putStrLn (String -> IO ()) -> String -> IO ()
forall a b. (a -> b) -> a -> b
$
                  String
"Sampling run failed:\n" String -> String -> String
forall a. [a] -> [a] -> [a]
++ String
err
              (Maybe (Int, Int), Map Text Int)
-> IO (Maybe (Int, Int), Map Text Int)
forall a. a -> IO a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (Maybe (Int, Int)
thresholds, Map Text Int
best_runtimes)
            Right ([(Text, Int)]
cmps, Int
t) -> do
              let (Int
tMin, Int
tMax) = (Int, Int) -> Maybe (Int, Int) -> (Int, Int)
forall a. a -> Maybe a -> a
fromMaybe (Int
thresholdMin, Int
thresholdMax) Maybe (Int, Int)
thresholds
              let ePars :: [Int]
ePars =
                    Set Int -> [Int]
forall a. Set a -> [a]
S.toAscList (Set Int -> [Int]) -> Set Int -> [Int]
forall a b. (a -> b) -> a -> b
$
                      ((Text, Int) -> Int) -> Set (Text, Int) -> Set Int
forall b a. Ord b => (a -> b) -> Set a -> Set b
S.map (Text, Int) -> Int
forall a b. (a, b) -> b
snd (Set (Text, Int) -> Set Int) -> Set (Text, Int) -> Set Int
forall a b. (a -> b) -> a -> b
$
                        ((Text, Int) -> Bool) -> Set (Text, Int) -> Set (Text, Int)
forall a. (a -> Bool) -> Set a -> Set a
S.filter ((Int, Int) -> (Text, Int) -> Bool
candidateEPar (Int
tMin, Int
tMax)) (Set (Text, Int) -> Set (Text, Int))
-> Set (Text, Int) -> Set (Text, Int)
forall a b. (a -> b) -> a -> b
$
                          [(Text, Int)] -> Set (Text, Int)
forall a. Ord a => [a] -> Set a
S.fromList [(Text, Int)]
cmps

                  runner :: Int -> Int -> IO (Maybe Int)
                  runner :: Int -> Int -> IO (Maybe Int)
runner Int
timeout' Int
threshold = do
                    res <- RunDataset
run Server
server Int
timeout' ((Text
v, Int
threshold) (Text, Int) -> [(Text, Int)] -> [(Text, Int)]
forall a. a -> [a] -> [a]
: [(Text, Int)]
already_tuned)
                    case res of
                      Right ([(Text, Int)]
_, Int
runTime) ->
                        Maybe Int -> IO (Maybe Int)
forall a. a -> IO a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (Maybe Int -> IO (Maybe Int)) -> Maybe Int -> IO (Maybe Int)
forall a b. (a -> b) -> a -> b
$ Int -> Maybe Int
forall a. a -> Maybe a
Just Int
runTime
                      Either String ([(Text, Int)], Int)
_ ->
                        Maybe Int -> IO (Maybe Int)
forall a. a -> IO a
forall (f :: * -> *) a. Applicative f => a -> f a
pure Maybe Int
forall a. Maybe a
Nothing

              Bool -> IO () -> IO ()
forall (f :: * -> *). Applicative f => Bool -> f () -> f ()
when (AutotuneOptions -> Int
optVerbose AutotuneOptions
opts Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
> Int
1) (IO () -> IO ()) -> IO () -> IO ()
forall a b. (a -> b) -> a -> b
$
                String -> IO ()
putStrLn (String -> IO ()) -> String -> IO ()
forall a b. (a -> b) -> a -> b
$
                  [String] -> String
unwords (String
"Got ePars: " String -> [String] -> [String]
forall a. a -> [a] -> [a]
: (Int -> String) -> [Int] -> [String]
forall a b. (a -> b) -> [a] -> [b]
map Int -> String
forall a. Show a => a -> String
show [Int]
ePars)

              (best_t, newMax) <- (Int -> Int -> IO (Maybe Int))
-> (Int, Int) -> [Int] -> IO (Int, Int)
binarySearch Int -> Int -> IO (Maybe Int)
runner (Int
t, Int
tMax) [Int]
ePars
              let newMinIdx = do
                    i <- Int -> Int
forall a. Enum a => a -> a
pred (Int -> Int) -> Maybe Int -> Maybe Int
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Int -> [Int] -> Maybe Int
forall a. Eq a => a -> [a] -> Maybe Int
elemIndex Int
newMax [Int]
ePars
                    if i < 0 then fail "Invalid lower index" else pure i
              let newMin = [Int] -> Int
forall a (f :: * -> *). (Num a, Ord a, Foldable f) => f a -> a
maxinum ([Int] -> Int) -> [Int] -> Int
forall a b. (a -> b) -> a -> b
$ [Maybe Int] -> [Int]
forall a. [Maybe a] -> [a]
catMaybes [Int -> Maybe Int
forall a. a -> Maybe a
Just Int
tMin, (Int -> Int) -> Maybe Int -> Maybe Int
forall a b. (a -> b) -> Maybe a -> Maybe b
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap ([Int]
ePars [Int] -> Int -> Int
forall a. HasCallStack => [a] -> Int -> a
!!) Maybe Int
newMinIdx]
              best_runtimes' <-
                case dataset_name `M.lookup` best_runtimes of
                  Just Int
rt
                    | Int -> Double
forall a b. (Integral a, Num b) => a -> b
fromIntegral Int
rt Double -> Double -> Double
forall a. Num a => a -> a -> a
* Double
epsilon Double -> Double -> Bool
forall a. Ord a => a -> a -> Bool
< Int -> Double
forall a b. (Integral a, Num b) => a -> b
fromIntegral Int
best_t -> do
                        Text -> IO ()
T.putStrLn (Text -> IO ()) -> Text -> IO ()
forall a b. (a -> b) -> a -> b
$
                          [Text] -> Text
T.unwords
                            [ (if Bool
fancyTerminal then Text
"\r\ESC[K" else Text
"")
                                Text -> Text -> Text
forall a. Semigroup a => a -> a -> a
<> Text
"WARNING! Possible non-monotonicity detected. Previous best run-time for dataset",
                              Text
dataset_name,
                              Text
" was",
                              Int -> Text
forall a. Show a => a -> Text
showText Int
rt,
                              Text
"but after tuning threshold",
                              Text
v,
                              Text
"it is",
                              Int -> Text
forall a. Show a => a -> Text
showText Int
best_t
                            ]
                        Map Text Int -> IO (Map Text Int)
forall a. a -> IO a
forall (f :: * -> *) a. Applicative f => a -> f a
pure Map Text Int
best_runtimes
                  Maybe Int
_ ->
                    Map Text Int -> IO (Map Text Int)
forall a. a -> IO a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (Map Text Int -> IO (Map Text Int))
-> Map Text Int -> IO (Map Text Int)
forall a b. (a -> b) -> a -> b
$ (Int -> Int -> Int) -> Text -> Int -> Map Text Int -> Map Text Int
forall k a. Ord k => (a -> a -> a) -> k -> a -> Map k a -> Map k a
M.insertWith Int -> Int -> Int
forall a. Ord a => a -> a -> a
min Text
dataset_name Int
best_t Map Text Int
best_runtimes
              pure (Just (newMin, newMax), best_runtimes')

    bestPair :: [(Int, Int)] -> (Int, Int)
    bestPair :: [(Int, Int)] -> (Int, Int)
bestPair = ((Int, Int) -> (Int, Int) -> Ordering)
-> [(Int, Int)] -> (Int, Int)
forall (t :: * -> *) a.
Foldable t =>
(a -> a -> Ordering) -> t a -> a
minimumBy (Int -> Int -> Ordering
forall a. Ord a => a -> a -> Ordering
compare (Int -> Int -> Ordering)
-> ((Int, Int) -> Int) -> (Int, Int) -> (Int, Int) -> Ordering
forall b c a. (b -> b -> c) -> (a -> b) -> a -> a -> c
`on` (Int, Int) -> Int
forall a b. (a, b) -> a
fst)

    timeout :: Int -> Int
    -- We wish to let datasets run for the untuned time + 20% + 1 second.
    timeout :: Int -> Int
timeout Int
elapsed = Double -> Int
forall b. Integral b => Double -> b
forall a b. (RealFrac a, Integral b) => a -> b
ceiling (Int -> Double
forall a b. (Integral a, Num b) => a -> b
fromIntegral Int
elapsed Double -> Double -> Double
forall a. Num a => a -> a -> a
* Double
1.2 :: Double) Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
1

    candidateEPar :: (Int, Int) -> (T.Text, Int) -> Bool
    candidateEPar :: (Int, Int) -> (Text, Int) -> Bool
candidateEPar (Int
tMin, Int
tMax) (Text
threshold, Int
ePar) =
      Int
ePar Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
> Int
tMin Bool -> Bool -> Bool
&& Int
ePar Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
< Int
tMax Bool -> Bool -> Bool
&& Text
threshold Text -> Text -> Bool
forall a. Eq a => a -> a -> Bool
== Text
v

    binarySearch :: (Int -> Int -> IO (Maybe Int)) -> (Int, Int) -> [Int] -> IO (Int, Int)
    binarySearch :: (Int -> Int -> IO (Maybe Int))
-> (Int, Int) -> [Int] -> IO (Int, Int)
binarySearch Int -> Int -> IO (Maybe Int)
runner best :: (Int, Int)
best@(Int
best_t, Int
best_e_par) [Int]
xs =
      case Int -> [Int] -> ([Int], [Int])
forall a. Int -> [a] -> ([a], [a])
splitAt ([Int] -> Int
forall a. [a] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length [Int]
xs Int -> Int -> Int
forall a. Integral a => a -> a -> a
`div` Int
2) [Int]
xs of
        ([Int]
lower, Int
middle : Int
middle' : [Int]
upper) -> do
          Bool -> IO () -> IO ()
forall (f :: * -> *). Applicative f => Bool -> f () -> f ()
when (AutotuneOptions -> Int
optVerbose AutotuneOptions
opts Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
> Int
0) (IO () -> IO ()) -> IO () -> IO ()
forall a b. (a -> b) -> a -> b
$
            String -> IO ()
putStrLn (String -> IO ()) -> String -> IO ()
forall a b. (a -> b) -> a -> b
$
              [String] -> String
unwords
                [ (if Bool
fancyTerminal then String
"\r\ESC[K" else String
"") String -> String -> String
forall a. [a] -> [a] -> [a]
++ String
"Trying e_par",
                  Int -> String
forall a. Show a => a -> String
show Int
middle,
                  String
"and",
                  Int -> String
forall a. Show a => a -> String
show Int
middle'
                ]
          candidate <- Int -> Int -> IO (Maybe Int)
runner (Int -> Int
timeout Int
best_t) Int
middle
          candidate' <- runner (timeout best_t) middle'
          case (candidate, candidate') of
            (Just Int
new_t, Just Int
new_t') ->
              if Int
new_t Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
< Int
new_t'
                then -- recurse into lower half
                  (Int -> Int -> IO (Maybe Int))
-> (Int, Int) -> [Int] -> IO (Int, Int)
binarySearch Int -> Int -> IO (Maybe Int)
runner ([(Int, Int)] -> (Int, Int)
bestPair [(Int
new_t, Int
middle), (Int, Int)
best]) [Int]
lower
                else -- recurse into upper half
                  (Int -> Int -> IO (Maybe Int))
-> (Int, Int) -> [Int] -> IO (Int, Int)
binarySearch Int -> Int -> IO (Maybe Int)
runner ([(Int, Int)] -> (Int, Int)
bestPair [(Int
new_t', Int
middle'), (Int, Int)
best]) [Int]
upper
            (Just Int
new_t, Maybe Int
Nothing) ->
              -- recurse into lower half
              (Int -> Int -> IO (Maybe Int))
-> (Int, Int) -> [Int] -> IO (Int, Int)
binarySearch Int -> Int -> IO (Maybe Int)
runner ([(Int, Int)] -> (Int, Int)
bestPair [(Int
new_t, Int
middle), (Int, Int)
best]) [Int]
lower
            (Maybe Int
Nothing, Just Int
new_t') ->
              -- recurse into upper half
              (Int -> Int -> IO (Maybe Int))
-> (Int, Int) -> [Int] -> IO (Int, Int)
binarySearch Int -> Int -> IO (Maybe Int)
runner ([(Int, Int)] -> (Int, Int)
bestPair [(Int
new_t', Int
middle'), (Int, Int)
best]) [Int]
upper
            (Maybe Int
Nothing, Maybe Int
Nothing) -> do
              Bool -> IO () -> IO ()
forall (f :: * -> *). Applicative f => Bool -> f () -> f ()
when (AutotuneOptions -> Int
optVerbose AutotuneOptions
opts Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
> Int
2) (IO () -> IO ()) -> IO () -> IO ()
forall a b. (a -> b) -> a -> b
$
                String -> IO ()
putStrLn (String -> IO ()) -> String -> IO ()
forall a b. (a -> b) -> a -> b
$
                  [String] -> String
unwords
                    [ String
"Timing failed for candidates",
                      Int -> String
forall a. Show a => a -> String
show Int
middle,
                      String
"and",
                      Int -> String
forall a. Show a => a -> String
show Int
middle'
                    ]
              (Int, Int) -> IO (Int, Int)
forall a. a -> IO a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (Int
best_t, Int
best_e_par)
        ([Int]
_, [Int]
_) -> do
          Bool -> IO () -> IO ()
forall (f :: * -> *). Applicative f => Bool -> f () -> f ()
when (AutotuneOptions -> Int
optVerbose AutotuneOptions
opts Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
> Int
0) (IO () -> IO ()) -> IO () -> IO ()
forall a b. (a -> b) -> a -> b
$
            String -> IO ()
putStrLn (String -> IO ()) -> String -> IO ()
forall a b. (a -> b) -> a -> b
$
              [String] -> String
unwords [(if Bool
fancyTerminal then String
"\r\ESC[K" else String
"") String -> String -> String
forall a. [a] -> [a] -> [a]
++ String
"Trying e_pars", [Int] -> String
forall a. Show a => a -> String
show [Int]
xs]
          candidates <-
            [Maybe (Int, Int)] -> [(Int, Int)]
forall a. [Maybe a] -> [a]
catMaybes ([Maybe (Int, Int)] -> [(Int, Int)])
-> ([Maybe Int] -> [Maybe (Int, Int)])
-> [Maybe Int]
-> [(Int, Int)]
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (Int -> Maybe Int -> Maybe (Int, Int))
-> [Int] -> [Maybe Int] -> [Maybe (Int, Int)]
forall a b c. (a -> b -> c) -> [a] -> [b] -> [c]
zipWith ((Int -> (Int, Int)) -> Maybe Int -> Maybe (Int, Int)
forall a b. (a -> b) -> Maybe a -> Maybe b
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap ((Int -> (Int, Int)) -> Maybe Int -> Maybe (Int, Int))
-> (Int -> Int -> (Int, Int))
-> Int
-> Maybe Int
-> Maybe (Int, Int)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (Int -> Int -> (Int, Int)) -> Int -> Int -> (Int, Int)
forall a b c. (a -> b -> c) -> b -> a -> c
flip (,)) [Int]
xs
              ([Maybe Int] -> [(Int, Int)]) -> IO [Maybe Int] -> IO [(Int, Int)]
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> (Int -> IO (Maybe Int)) -> [Int] -> IO [Maybe Int]
forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
(a -> m b) -> t a -> m (t b)
forall (m :: * -> *) a b. Monad m => (a -> m b) -> [a] -> m [b]
mapM (Int -> Int -> IO (Maybe Int)
runner (Int -> Int -> IO (Maybe Int)) -> Int -> Int -> IO (Maybe Int)
forall a b. (a -> b) -> a -> b
$ Int -> Int
timeout Int
best_t) [Int]
xs
          pure $ bestPair $ best : candidates

--- CLI

tune :: AutotuneOptions -> FilePath -> IO Path
tune :: AutotuneOptions -> String -> IO [(Text, Int)]
tune AutotuneOptions
opts String
prog = do
  futhark <- (String -> FutharkExe) -> IO String -> IO FutharkExe
forall a b. (a -> b) -> IO a -> IO b
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap String -> FutharkExe
FutharkExe (IO String -> IO FutharkExe) -> IO String -> IO FutharkExe
forall a b. (a -> b) -> a -> b
$ IO String -> (String -> IO String) -> Maybe String -> IO String
forall b a. b -> (a -> b) -> Maybe a -> b
maybe IO String
getExecutablePath String -> IO String
forall a. a -> IO a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (Maybe String -> IO String) -> Maybe String -> IO String
forall a b. (a -> b) -> a -> b
$ AutotuneOptions -> Maybe String
optFuthark AutotuneOptions
opts

  putStrLn $ "Compiling " ++ prog ++ "..."
  datasets <- prepare opts futhark prog

  putStrLn $ "Running with options: " ++ unwords (serverOptions opts)
  let progbin = String
"." String -> String -> String
</> String -> String
dropExtension String
prog
  withServer (futharkServerCfg progbin (serverOptions opts)) $ \Server
server -> do
    forest <- Server -> IO ThresholdForest
thresholdForest Server
server
    let paths = ThresholdForest -> [(Text, [(Text, Int)])]
tuningPaths ThresholdForest
forest
    let total = [(Text, [(Text, Int)])] -> Int
forall a. [a] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length [(Text, [(Text, Int)])]
paths
    when (optVerbose opts > 0) $
      putStrLn $
        ("Threshold forest:\n" <>) $
          drawForest (map (fmap show) forest)

    counter <- newIORef (0 :: Int)
    let dataset_names =
          Text -> [Text] -> Text
T.intercalate Text
"," ([Text] -> Text) -> [Text] -> Text
forall a b. (a -> b) -> a -> b
$ ((Text, RunDataset, Text) -> Text)
-> [(Text, RunDataset, Text)] -> [Text]
forall a b. (a -> b) -> [a] -> [b]
map (\(Text
name, RunDataset
_, Text
ep) -> Text
ep Text -> Text -> Text
forall a. Semigroup a => a -> a -> a
<> Text
":" Text -> Text -> Text
forall a. Semigroup a => a -> a -> a
<> Text
name) [(Text, RunDataset, Text)]
datasets
    result <-
      foldM
        ( \([(Text, Int)], Map Text Int)
acc tp :: (Text, [(Text, Int)])
tp@(Text
v, [(Text, Int)]
_) -> do
            IORef Int -> (Int -> Int) -> IO ()
forall a. IORef a -> (a -> a) -> IO ()
modifyIORef' IORef Int
counter (Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
1)
            n <- IORef Int -> IO Int
forall a. IORef a -> IO a
readIORef IORef Int
counter
            let bar =
                  String
"\rTuning: "
                    String -> String -> String
forall a. [a] -> [a] -> [a]
++ Int -> String
forall a. Show a => a -> String
show Int
n
                    String -> String -> String
forall a. [a] -> [a] -> [a]
++ String
"/"
                    String -> String -> String
forall a. [a] -> [a] -> [a]
++ Int -> String
forall a. Show a => a -> String
show Int
total
                    String -> String -> String
forall a. [a] -> [a] -> [a]
++ String
" "
                    String -> String -> String
forall a. [a] -> [a] -> [a]
++ Text -> String
T.unpack
                      ( ProgressBar -> Text
progressBar
                          ProgressBar
                            { progressBarSteps :: Int
progressBarSteps = Int
10,
                              progressBarBound :: Double
progressBarBound = Double
1,
                              progressBarElapsed :: Double
progressBarElapsed = Int -> Double
forall a b. (Integral a, Num b) => a -> b
fromIntegral Int
n Double -> Double -> Double
forall a. Fractional a => a -> a -> a
/ Int -> Double
forall a b. (Integral a, Num b) => a -> b
fromIntegral Int
total
                            }
                      )
                    String -> String -> String
forall a. [a] -> [a] -> [a]
++ Text -> String
T.unpack Text
v
                    String -> String -> String
forall a. [a] -> [a] -> [a]
++ String
" on "
                    String -> String -> String
forall a. [a] -> [a] -> [a]
++ Text -> String
T.unpack Text
dataset_names
            when fancyTerminal $ putStr bar
            hFlush stdout
            r <- tuneThreshold opts server datasets acc tp
            when fancyTerminal $ putStr bar
            hFlush stdout
            pure r
        )
        ([], mempty)
        paths
    when fancyTerminal $ do
      putStr "\r\ESC[K"
      hFlush stdout
    pure (fst result)

onServerException :: ServerException -> IO a
onServerException :: forall a. ServerException -> IO a
onServerException (ServerException Text
s) = do
  Handle -> Text -> IO ()
T.hPutStrLn Handle
stderr Text
s
  IO a
forall a. IO a
exitFailure

runAutotuner :: AutotuneOptions -> FilePath -> IO ()
runAutotuner :: AutotuneOptions -> String -> IO ()
runAutotuner AutotuneOptions
opts String
prog =
  do
    best <- AutotuneOptions -> String -> IO [(Text, Int)]
tune AutotuneOptions
opts String
prog

    let tuning = [Text] -> Text
T.unlines ([Text] -> Text) -> [Text] -> Text
forall a b. (a -> b) -> a -> b
$ do
          (s, n) <- ((Text, Int) -> Text) -> [(Text, Int)] -> [(Text, Int)]
forall b a. Ord b => (a -> b) -> [a] -> [a]
sortOn (Text, Int) -> Text
forall a b. (a, b) -> a
fst [(Text, Int)]
best
          pure $ s <> "=" <> showText n

    case optTuning opts of
      Maybe String
Nothing -> () -> IO ()
forall a. a -> IO a
forall (f :: * -> *) a. Applicative f => a -> f a
pure ()
      Just String
suffix -> do
        String -> Text -> IO ()
T.writeFile (String
prog String -> String -> String
<.> String
suffix) Text
tuning
        String -> IO ()
putStrLn (String -> IO ()) -> String -> IO ()
forall a b. (a -> b) -> a -> b
$ String
"Wrote " String -> String -> String
forall a. [a] -> [a] -> [a]
++ String
prog String -> String -> String
<.> String
suffix

    T.putStrLn $ "Result of autotuning:\n" <> tuning
    IO () -> (ServerException -> IO ()) -> IO ()
forall e a. Exception e => IO a -> (e -> IO a) -> IO a
`catch` ServerException -> IO ()
forall a. ServerException -> IO a
onServerException

supportedBackends :: [String]
supportedBackends :: [String]
supportedBackends = [String
"opencl", String
"cuda", String
"hip"]

commandLineOptions :: [FunOptDescr AutotuneOptions]
commandLineOptions :: [FunOptDescr AutotuneOptions]
commandLineOptions =
  [ String
-> [String]
-> ArgDescr (Either (IO ()) (AutotuneOptions -> AutotuneOptions))
-> String
-> FunOptDescr AutotuneOptions
forall a. String -> [String] -> ArgDescr a -> String -> OptDescr a
Option
      String
"r"
      [String
"runs"]
      ( (String -> Either (IO ()) (AutotuneOptions -> AutotuneOptions))
-> String
-> ArgDescr (Either (IO ()) (AutotuneOptions -> AutotuneOptions))
forall a. (String -> a) -> String -> ArgDescr a
ReqArg
          ( \String
n ->
              case ReadS Int
forall a. Read a => ReadS a
reads String
n of
                [(Int
n', String
"")] | Int
n' Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
>= Int
0 ->
                  (AutotuneOptions -> AutotuneOptions)
-> Either (IO ()) (AutotuneOptions -> AutotuneOptions)
forall a b. b -> Either a b
Right ((AutotuneOptions -> AutotuneOptions)
 -> Either (IO ()) (AutotuneOptions -> AutotuneOptions))
-> (AutotuneOptions -> AutotuneOptions)
-> Either (IO ()) (AutotuneOptions -> AutotuneOptions)
forall a b. (a -> b) -> a -> b
$ \AutotuneOptions
config -> AutotuneOptions
config {optMinRuns = n'}
                [(Int, String)]
_ ->
                  IO () -> Either (IO ()) (AutotuneOptions -> AutotuneOptions)
forall a b. a -> Either a b
Left (IO () -> Either (IO ()) (AutotuneOptions -> AutotuneOptions))
-> IO () -> Either (IO ()) (AutotuneOptions -> AutotuneOptions)
forall a b. (a -> b) -> a -> b
$ String -> IO ()
optionsError (String -> IO ()) -> String -> IO ()
forall a b. (a -> b) -> a -> b
$ String
"'" String -> String -> String
forall a. [a] -> [a] -> [a]
++ String
n String -> String -> String
forall a. [a] -> [a] -> [a]
++ String
"' is not a non-negative integer."
          )
          String
"RUNS"
      )
      String
"Run each test case this many times.",
    String
-> [String]
-> ArgDescr (Either (IO ()) (AutotuneOptions -> AutotuneOptions))
-> String
-> FunOptDescr AutotuneOptions
forall a. String -> [String] -> ArgDescr a -> String -> OptDescr a
Option
      []
      [String
"backend"]
      ( (String -> Either (IO ()) (AutotuneOptions -> AutotuneOptions))
-> String
-> ArgDescr (Either (IO ()) (AutotuneOptions -> AutotuneOptions))
forall a. (String -> a) -> String -> ArgDescr a
ReqArg
          ( \String
backend ->
              if String
backend String -> [String] -> Bool
forall a. Eq a => a -> [a] -> Bool
forall (t :: * -> *) a. (Foldable t, Eq a) => a -> t a -> Bool
`elem` [String]
supportedBackends
                then (AutotuneOptions -> AutotuneOptions)
-> Either (IO ()) (AutotuneOptions -> AutotuneOptions)
forall a b. b -> Either a b
Right ((AutotuneOptions -> AutotuneOptions)
 -> Either (IO ()) (AutotuneOptions -> AutotuneOptions))
-> (AutotuneOptions -> AutotuneOptions)
-> Either (IO ()) (AutotuneOptions -> AutotuneOptions)
forall a b. (a -> b) -> a -> b
$ \AutotuneOptions
config -> AutotuneOptions
config {optBackend = backend}
                else IO () -> Either (IO ()) (AutotuneOptions -> AutotuneOptions)
forall a b. a -> Either a b
Left (IO () -> Either (IO ()) (AutotuneOptions -> AutotuneOptions))
-> IO () -> Either (IO ()) (AutotuneOptions -> AutotuneOptions)
forall a b. (a -> b) -> a -> b
$ String -> IO ()
optionsError (String -> IO ()) -> String -> IO ()
forall a b. (a -> b) -> a -> b
$ String
"autotuning is only supported for these backends: " String -> String -> String
forall a. Semigroup a => a -> a -> a
<> [String] -> String
unwords [String]
supportedBackends
          )
          String
"BACKEND"
      )
      String
"The backend used (defaults to 'opencl').",
    String
-> [String]
-> ArgDescr (Either (IO ()) (AutotuneOptions -> AutotuneOptions))
-> String
-> FunOptDescr AutotuneOptions
forall a. String -> [String] -> ArgDescr a -> String -> OptDescr a
Option
      []
      [String
"futhark"]
      ( (String -> Either (IO ()) (AutotuneOptions -> AutotuneOptions))
-> String
-> ArgDescr (Either (IO ()) (AutotuneOptions -> AutotuneOptions))
forall a. (String -> a) -> String -> ArgDescr a
ReqArg
          (\String
prog -> (AutotuneOptions -> AutotuneOptions)
-> Either (IO ()) (AutotuneOptions -> AutotuneOptions)
forall a b. b -> Either a b
Right ((AutotuneOptions -> AutotuneOptions)
 -> Either (IO ()) (AutotuneOptions -> AutotuneOptions))
-> (AutotuneOptions -> AutotuneOptions)
-> Either (IO ()) (AutotuneOptions -> AutotuneOptions)
forall a b. (a -> b) -> a -> b
$ \AutotuneOptions
config -> AutotuneOptions
config {optFuthark = Just prog})
          String
"PROGRAM"
      )
      String
"The binary used for operations (defaults to 'futhark').",
    String
-> [String]
-> ArgDescr (Either (IO ()) (AutotuneOptions -> AutotuneOptions))
-> String
-> FunOptDescr AutotuneOptions
forall a. String -> [String] -> ArgDescr a -> String -> OptDescr a
Option
      []
      [String
"pass-option"]
      ( (String -> Either (IO ()) (AutotuneOptions -> AutotuneOptions))
-> String
-> ArgDescr (Either (IO ()) (AutotuneOptions -> AutotuneOptions))
forall a. (String -> a) -> String -> ArgDescr a
ReqArg
          ( \String
opt ->
              (AutotuneOptions -> AutotuneOptions)
-> Either (IO ()) (AutotuneOptions -> AutotuneOptions)
forall a b. b -> Either a b
Right ((AutotuneOptions -> AutotuneOptions)
 -> Either (IO ()) (AutotuneOptions -> AutotuneOptions))
-> (AutotuneOptions -> AutotuneOptions)
-> Either (IO ()) (AutotuneOptions -> AutotuneOptions)
forall a b. (a -> b) -> a -> b
$ \AutotuneOptions
config ->
                AutotuneOptions
config {optExtraOptions = opt : optExtraOptions config}
          )
          String
"OPT"
      )
      String
"Pass this option to programs being run.",
    String
-> [String]
-> ArgDescr (Either (IO ()) (AutotuneOptions -> AutotuneOptions))
-> String
-> FunOptDescr AutotuneOptions
forall a. String -> [String] -> ArgDescr a -> String -> OptDescr a
Option
      []
      [String
"tuning"]
      ( (String -> Either (IO ()) (AutotuneOptions -> AutotuneOptions))
-> String
-> ArgDescr (Either (IO ()) (AutotuneOptions -> AutotuneOptions))
forall a. (String -> a) -> String -> ArgDescr a
ReqArg
          (\String
s -> (AutotuneOptions -> AutotuneOptions)
-> Either (IO ()) (AutotuneOptions -> AutotuneOptions)
forall a b. b -> Either a b
Right ((AutotuneOptions -> AutotuneOptions)
 -> Either (IO ()) (AutotuneOptions -> AutotuneOptions))
-> (AutotuneOptions -> AutotuneOptions)
-> Either (IO ()) (AutotuneOptions -> AutotuneOptions)
forall a b. (a -> b) -> a -> b
$ \AutotuneOptions
config -> AutotuneOptions
config {optTuning = Just s})
          String
"EXTENSION"
      )
      String
"Write tuning files with this extension (default: .tuning).",
    String
-> [String]
-> ArgDescr (Either (IO ()) (AutotuneOptions -> AutotuneOptions))
-> String
-> FunOptDescr AutotuneOptions
forall a. String -> [String] -> ArgDescr a -> String -> OptDescr a
Option
      []
      [String
"timeout"]
      ( (String -> Either (IO ()) (AutotuneOptions -> AutotuneOptions))
-> String
-> ArgDescr (Either (IO ()) (AutotuneOptions -> AutotuneOptions))
forall a. (String -> a) -> String -> ArgDescr a
ReqArg
          ( \String
n ->
              case ReadS Int
forall a. Read a => ReadS a
reads String
n of
                [(Int
n', String
"")] ->
                  (AutotuneOptions -> AutotuneOptions)
-> Either (IO ()) (AutotuneOptions -> AutotuneOptions)
forall a b. b -> Either a b
Right ((AutotuneOptions -> AutotuneOptions)
 -> Either (IO ()) (AutotuneOptions -> AutotuneOptions))
-> (AutotuneOptions -> AutotuneOptions)
-> Either (IO ()) (AutotuneOptions -> AutotuneOptions)
forall a b. (a -> b) -> a -> b
$ \AutotuneOptions
config -> AutotuneOptions
config {optTimeout = n'}
                [(Int, String)]
_ ->
                  IO () -> Either (IO ()) (AutotuneOptions -> AutotuneOptions)
forall a b. a -> Either a b
Left (IO () -> Either (IO ()) (AutotuneOptions -> AutotuneOptions))
-> IO () -> Either (IO ()) (AutotuneOptions -> AutotuneOptions)
forall a b. (a -> b) -> a -> b
$ String -> IO ()
optionsError (String -> IO ()) -> String -> IO ()
forall a b. (a -> b) -> a -> b
$ String
"'" String -> String -> String
forall a. [a] -> [a] -> [a]
++ String
n String -> String -> String
forall a. [a] -> [a] -> [a]
++ String
"' is not a non-negative integer."
          )
          String
"SECONDS"
      )
      String
"Initial tuning timeout for each dataset. Later tuning runs are based off of the runtime of the first run.",
    String
-> [String]
-> ArgDescr (Either (IO ()) (AutotuneOptions -> AutotuneOptions))
-> String
-> FunOptDescr AutotuneOptions
forall a. String -> [String] -> ArgDescr a -> String -> OptDescr a
Option
      []
      [String
"skip-compilation"]
      (Either (IO ()) (AutotuneOptions -> AutotuneOptions)
-> ArgDescr (Either (IO ()) (AutotuneOptions -> AutotuneOptions))
forall a. a -> ArgDescr a
NoArg (Either (IO ()) (AutotuneOptions -> AutotuneOptions)
 -> ArgDescr (Either (IO ()) (AutotuneOptions -> AutotuneOptions)))
-> Either (IO ()) (AutotuneOptions -> AutotuneOptions)
-> ArgDescr (Either (IO ()) (AutotuneOptions -> AutotuneOptions))
forall a b. (a -> b) -> a -> b
$ (AutotuneOptions -> AutotuneOptions)
-> Either (IO ()) (AutotuneOptions -> AutotuneOptions)
forall a b. b -> Either a b
Right ((AutotuneOptions -> AutotuneOptions)
 -> Either (IO ()) (AutotuneOptions -> AutotuneOptions))
-> (AutotuneOptions -> AutotuneOptions)
-> Either (IO ()) (AutotuneOptions -> AutotuneOptions)
forall a b. (a -> b) -> a -> b
$ \AutotuneOptions
config -> AutotuneOptions
config {optSkipCompilation = True})
      String
"Use already compiled program.",
    String
-> [String]
-> ArgDescr (Either (IO ()) (AutotuneOptions -> AutotuneOptions))
-> String
-> FunOptDescr AutotuneOptions
forall a. String -> [String] -> ArgDescr a -> String -> OptDescr a
Option
      String
"v"
      [String
"verbose"]
      (Either (IO ()) (AutotuneOptions -> AutotuneOptions)
-> ArgDescr (Either (IO ()) (AutotuneOptions -> AutotuneOptions))
forall a. a -> ArgDescr a
NoArg (Either (IO ()) (AutotuneOptions -> AutotuneOptions)
 -> ArgDescr (Either (IO ()) (AutotuneOptions -> AutotuneOptions)))
-> Either (IO ()) (AutotuneOptions -> AutotuneOptions)
-> ArgDescr (Either (IO ()) (AutotuneOptions -> AutotuneOptions))
forall a b. (a -> b) -> a -> b
$ (AutotuneOptions -> AutotuneOptions)
-> Either (IO ()) (AutotuneOptions -> AutotuneOptions)
forall a b. b -> Either a b
Right ((AutotuneOptions -> AutotuneOptions)
 -> Either (IO ()) (AutotuneOptions -> AutotuneOptions))
-> (AutotuneOptions -> AutotuneOptions)
-> Either (IO ()) (AutotuneOptions -> AutotuneOptions)
forall a b. (a -> b) -> a -> b
$ \AutotuneOptions
config -> AutotuneOptions
config {optVerbose = optVerbose config + 1})
      String
"Enable logging.  Pass multiple times for more.",
    String
-> [String]
-> ArgDescr (Either (IO ()) (AutotuneOptions -> AutotuneOptions))
-> String
-> FunOptDescr AutotuneOptions
forall a. String -> [String] -> ArgDescr a -> String -> OptDescr a
Option
      []
      [String
"spec-file"]
      ((String -> Either (IO ()) (AutotuneOptions -> AutotuneOptions))
-> String
-> ArgDescr (Either (IO ()) (AutotuneOptions -> AutotuneOptions))
forall a. (String -> a) -> String -> ArgDescr a
ReqArg (\String
s -> (AutotuneOptions -> AutotuneOptions)
-> Either (IO ()) (AutotuneOptions -> AutotuneOptions)
forall a b. b -> Either a b
Right ((AutotuneOptions -> AutotuneOptions)
 -> Either (IO ()) (AutotuneOptions -> AutotuneOptions))
-> (AutotuneOptions -> AutotuneOptions)
-> Either (IO ()) (AutotuneOptions -> AutotuneOptions)
forall a b. (a -> b) -> a -> b
$ \AutotuneOptions
config -> AutotuneOptions
config {optTestSpec = Just s}) String
"FILE")
      String
"Use test specification from this file."
  ]

-- | Run @futhark autotune@
main :: String -> [String] -> IO ()
main :: String -> [String] -> IO ()
main = AutotuneOptions
-> [FunOptDescr AutotuneOptions]
-> String
-> ([String] -> AutotuneOptions -> Maybe (IO ()))
-> String
-> [String]
-> IO ()
forall cfg.
cfg
-> [FunOptDescr cfg]
-> String
-> ([String] -> cfg -> Maybe (IO ()))
-> String
-> [String]
-> IO ()
mainWithOptions AutotuneOptions
initialAutotuneOptions [FunOptDescr AutotuneOptions]
commandLineOptions String
"options... program" (([String] -> AutotuneOptions -> Maybe (IO ()))
 -> String -> [String] -> IO ())
-> ([String] -> AutotuneOptions -> Maybe (IO ()))
-> String
-> [String]
-> IO ()
forall a b. (a -> b) -> a -> b
$ \[String]
progs AutotuneOptions
config ->
  case [String]
progs of
    [String
prog] -> IO () -> Maybe (IO ())
forall a. a -> Maybe a
Just (IO () -> Maybe (IO ())) -> IO () -> Maybe (IO ())
forall a b. (a -> b) -> a -> b
$ AutotuneOptions -> String -> IO ()
runAutotuner AutotuneOptions
config String
prog
    [String]
_ -> Maybe (IO ())
forall a. Maybe a
Nothing