content: -- | @futhark autotune@
  module Futhark.CLI.Autotune (main) where
  
  import Control.Monad
  import Data.ByteString.Char8 qualified as SBS
  import Data.Function (on)
  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 (maxinum, showText)
  import Futhark.Util.Options
  import System.Directory
  import System.Environment (getExecutablePath)
  import System.Exit
  import System.FilePath
  import Text.Read (readMaybe)
  import Text.Regex.TDFA
  
  data AutotuneOptions = AutotuneOptions
    { optBackend :: String,
      optFuthark :: Maybe String,
      optMinRuns :: Int,
      optTuning :: Maybe String,
      optExtraOptions :: [String],
      optVerbose :: Int,
      optTimeout :: Int,
      optSkipCompilation :: Bool,
      optDefaultThreshold :: Int,
      optTestSpec :: Maybe FilePath
    }
  
  initialAutotuneOptions :: AutotuneOptions
  initialAutotuneOptions =
    AutotuneOptions
      { optBackend = "opencl",
        optFuthark = Nothing,
        optMinRuns = 10,
        optTuning = Just "tuning",
        optExtraOptions = [],
        optVerbose = 0,
        optTimeout = 600,
        optSkipCompilation = False,
        optDefaultThreshold = thresholdMax,
        optTestSpec = Nothing
      }
  
  compileOptions :: AutotuneOptions -> IO CompileOptions
  compileOptions opts = do
    futhark <- maybe getExecutablePath pure $ optFuthark opts
    pure $
      CompileOptions
        { compFuthark = futhark,
          compBackend = optBackend opts,
          compOptions = mempty
        }
  
  runOptions :: Int -> AutotuneOptions -> RunOptions
  runOptions timeout_s opts =
    RunOptions
      { runMinRuns = optMinRuns opts,
        runMinTime = 0.5,
        runTimeout = timeout_s,
        runVerbose = optVerbose opts,
        runConvergencePhase = True,
        runConvergenceMaxTime = fromIntegral timeout_s,
        runResultAction = const $ pure (),
        runProfile = False
      }
  
  type Path = [(T.Text, Int)]
  
  regexBlocks :: Regex -> T.Text -> Maybe [T.Text]
  regexBlocks regex s = do
    (_, _, _, groups) <-
      matchM regex s :: Maybe (T.Text, T.Text, T.Text, [T.Text])
    Just groups
  
  comparisons :: T.Text -> [(T.Text, Int)]
  comparisons = mapMaybe isComparison . T.lines
    where
      regex = makeRegex ("Compared ([^ ]+) <= (-?[0-9]+)" :: String)
      isComparison l = do
        [thresh, val] <- regexBlocks regex 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 opts =
    "--default-threshold"
      : show (optDefaultThreshold opts)
      : "-L"
      : optExtraOptions opts
  
  checkCmd :: Either CmdFailure a -> IO a
  checkCmd = either (error . T.unpack . T.unlines . failureMsg) pure
  
  setTuningParam :: Server -> T.Text -> Int -> IO ()
  setTuningParam server name val =
    void $ checkCmd =<< cmdSetTuningParam server name (showText val)
  
  setTuningParams :: Server -> Path -> IO ()
  setTuningParams server = mapM_ (uncurry $ setTuningParam server)
  
  restoreTuningParams :: AutotuneOptions -> Server -> Path -> IO ()
  restoreTuningParams opts server = mapM_ opt
    where
      opt (name, _) = setTuningParam server name (optDefaultThreshold opts)
  
  prepare :: AutotuneOptions -> FutharkExe -> FilePath -> IO [(DatasetName, RunDataset, T.Text)]
  prepare opts futhark prog = do
    spec <-
      maybe (testSpecFromProgramOrDie prog) testSpecFromFileOrDie $
        optTestSpec opts
    copts <- compileOptions opts
  
    truns <-
      case testAction spec of
        RunCases ios _ _ | not $ null ios -> do
          when (optVerbose opts > 1) $
            putStrLn $
              unwords ("Entry points:" : map (T.unpack . iosEntryPoint) ios)
  
          if optSkipCompilation opts
            then do
              exists <- doesFileExist $ binaryName prog
              if exists
                then pure ios
                else do
                  putStrLn $ binaryName prog ++ " does not exist, but --skip-compilation passed."
                  exitFailure
            else do
              res <- prepareBenchmarkProgram Nothing copts prog ios
              case res of
                Left (err, errstr) -> do
                  putStrLn err
                  maybe (pure ()) SBS.putStrLn errstr
                  exitFailure
                Right () ->
                  pure ios
        _ ->
          fail "Unsupported test spec."
  
    let runnableDataset entry_point trun =
          case runExpectedResult trun of
            Succeeds expected
              | null (runTags trun `intersect` ["notune", "disable"]) ->
                  Just
                    ( runDescription trun,
                      \server -> run server entry_point trun expected
                    )
            _ -> Nothing
  
    fmap concat . forM truns $ \ios -> do
      let cases =
            mapMaybe (runnableDataset $ iosEntryPoint ios) (iosTestRuns ios)
      forM cases $ \(dataset, do_run) ->
        pure (dataset, do_run, iosEntryPoint ios)
    where
      run server entry_point trun expected timeout path = do
        let bestRuntime (runres, errout, _) =
              ( comparisons errout,
                minimum $ map runMicroseconds runres
              )
  
            ropts = runOptions timeout opts
  
        when (optVerbose opts > 1) $
          putStrLn ("Trying path: " ++ show 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.
        setTuningParams server path
  
        either (Left . T.unpack) (Right . bestRuntime)
          <$> benchmarkDataset
            server
            ropts
            futhark
            prog
            entry_point
            (runInput trun)
            expected
            (testRunReferenceOutput prog entry_point trun)
          <* restoreTuningParams opts server path
  
  --- Benchmarking a program
  
  data DatasetResult = DatasetResult [(T.Text, Int)] Double
    deriving (Show)
  
  --- Finding initial comparisons.
  
  --- Extracting threshold hierarchy.
  
  type ThresholdForest = Forest (T.Text, Bool)
  
  thresholdMin, thresholdMax :: Int
  thresholdMin = 1
  thresholdMax = 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 = concatMap (treePaths [])
    where
      treePaths ancestors (Node (v, _) children) =
        concatMap (onChild ancestors v) children ++ [(v, ancestors)]
  
      onChild ancestors v child@(Node (_, cmp) _) =
        treePaths (ancestors ++ [(v, t cmp)]) child
  
      t False = thresholdMax
      t True = thresholdMin
  
  allTuningParams :: Server -> IO [(T.Text, T.Text)]
  allTuningParams server = do
    entry_points <- checkCmd =<< cmdEntryPoints 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 = do
    thresholds <- mapMaybe findThreshold <$> allTuningParams server
    let root (v, _) = ((v, False), [])
    pure $
      unfoldForest (unfold thresholds) $
        map root $
          filter (null . snd) thresholds
    where
      regex = makeRegex ("threshold\\(([^ ]+,)(.*)\\)" :: T.Text)
  
      findThreshold :: (T.Text, T.Text) -> Maybe (T.Text, [(T.Text, Bool)])
      findThreshold (name, param_class) = do
        [_, grp] <- regexBlocks regex param_class
        pure
          ( name,
            filter (not . T.null . fst)
              $ map
                ( \x ->
                    if "!" `T.isPrefixOf` x
                      then (T.drop 1 x, False)
                      else (x, True)
                )
              $ T.words grp
          )
  
      unfold thresholds ((parent, parent_cmp), ancestors) =
        let ancestors' = parent : ancestors
  
            isChild (v, v_ancestors) = do
              cmp <- lookup parent v_ancestors
              guard $
                sort (map fst v_ancestors)
                  == sort (parent : ancestors)
              pure ((v, cmp), ancestors')
         in ((parent, parent_cmp), mapMaybe isChild 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 = 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 opts server datasets (already_tuned, best_runtimes0) (v, _v_path) = do
    (tune_result, best_runtimes) <-
      foldM tuneDataset (Nothing, best_runtimes0) datasets
    case tune_result of
      Nothing ->
        pure ((v, thresholdMin) : already_tuned, best_runtimes)
      Just (_, threshold) ->
        pure ((v, threshold) : already_tuned, best_runtimes)
    where
      tuneDataset ::
        (Maybe (Int, Int), M.Map DatasetName Int) ->
        (DatasetName, RunDataset, T.Text) ->
        IO (Maybe (Int, Int), M.Map DatasetName Int)
      tuneDataset (thresholds, best_runtimes) (dataset_name, run, entry_point) = do
        relevant <- checkCmd =<< cmdTuningParams server 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
            T.putStrLn $
              T.unwords
                [ "Tuning",
                  v,
                  "on entry point",
                  entry_point,
                  "and dataset",
                  dataset_name
                ]
  
            sample_run <-
              run
                server
                (optTimeout opts)
                ((v, maybe thresholdMax snd thresholds) : already_tuned)
  
            case sample_run of
              Left 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.
                when (optVerbose opts > 0) $
                  putStrLn $
                    "Sampling run failed:\n" ++ err
                pure (thresholds, best_runtimes)
              Right (cmps, t) -> do
                let (tMin, tMax) = fromMaybe (thresholdMin, thresholdMax) thresholds
                let ePars =
                      S.toAscList $
                        S.map snd $
                          S.filter (candidateEPar (tMin, tMax)) $
                            S.fromList cmps
  
                    runner :: Int -> Int -> IO (Maybe Int)
                    runner timeout' threshold = do
                      res <- run server timeout' ((v, threshold) : already_tuned)
                      case res of
                        Right (_, runTime) ->
                          pure $ Just runTime
                        _ ->
                          pure Nothing
  
                when (optVerbose opts > 1) $
                  putStrLn $
                    unwords ("Got ePars: " : map show ePars)
  
                (best_t, newMax) <- binarySearch runner (t, tMax) ePars
                let newMinIdx = do
                      i <- pred <$> elemIndex newMax ePars
                      if i < 0 then fail "Invalid lower index" else pure i
                let newMin = maxinum $ catMaybes [Just tMin, fmap (ePars !!) newMinIdx]
                best_runtimes' <-
                  case dataset_name `M.lookup` best_runtimes of
                    Just rt
                      | fromIntegral rt * epsilon < fromIntegral best_t -> do
                          T.putStrLn $
                            T.unwords
                              [ "WARNING! Possible non-monotonicity detected. Previous best run-time for dataset",
                                dataset_name,
                                " was",
                                showText rt,
                                "but after tuning threshold",
                                v,
                                "it is",
                                showText best_t
                              ]
                          pure best_runtimes
                    _ ->
                      pure $ M.insertWith min dataset_name best_t best_runtimes
                pure (Just (newMin, newMax), best_runtimes')
  
      bestPair :: [(Int, Int)] -> (Int, Int)
      bestPair = minimumBy (compare `on` fst)
  
      timeout :: Int -> Int
      -- We wish to let datasets run for the untuned time + 20% + 1 second.
      timeout elapsed = ceiling (fromIntegral elapsed * 1.2 :: Double) + 1
  
      candidateEPar :: (Int, Int) -> (T.Text, Int) -> Bool
      candidateEPar (tMin, tMax) (threshold, ePar) =
        ePar > tMin && ePar < tMax && threshold == v
  
      binarySearch :: (Int -> Int -> IO (Maybe Int)) -> (Int, Int) -> [Int] -> IO (Int, Int)
      binarySearch runner best@(best_t, best_e_par) xs =
        case splitAt (length xs `div` 2) xs of
          (lower, middle : middle' : upper) -> do
            when (optVerbose opts > 0) $
              putStrLn $
                unwords
                  [ "Trying e_par",
                    show middle,
                    "and",
                    show middle'
                  ]
            candidate <- runner (timeout best_t) middle
            candidate' <- runner (timeout best_t) middle'
            case (candidate, candidate') of
              (Just new_t, Just new_t') ->
                if new_t < new_t'
                  then -- recurse into lower half
                    binarySearch runner (bestPair [(new_t, middle), best]) lower
                  else -- recurse into upper half
                    binarySearch runner (bestPair [(new_t', middle'), best]) upper
              (Just new_t, Nothing) ->
                -- recurse into lower half
                binarySearch runner (bestPair [(new_t, middle), best]) lower
              (Nothing, Just new_t') ->
                -- recurse into upper half
                binarySearch runner (bestPair [(new_t', middle'), best]) upper
              (Nothing, Nothing) -> do
                when (optVerbose opts > 2) $
                  putStrLn $
                    unwords
                      [ "Timing failed for candidates",
                        show middle,
                        "and",
                        show middle'
                      ]
                pure (best_t, best_e_par)
          (_, _) -> do
            when (optVerbose opts > 0) $
              putStrLn $
                unwords ["Trying e_pars", show xs]
            candidates <-
              catMaybes . zipWith (fmap . flip (,)) xs
                <$> mapM (runner $ timeout best_t) xs
            pure $ bestPair $ best : candidates
  
  --- CLI
  
  tune :: AutotuneOptions -> FilePath -> IO Path
  tune opts prog = do
    futhark <- fmap FutharkExe $ maybe getExecutablePath pure $ optFuthark opts
  
    putStrLn $ "Compiling " ++ prog ++ "..."
    datasets <- prepare opts futhark prog
  
    putStrLn $ "Running with options: " ++ unwords (serverOptions opts)
    let progbin = "." </> dropExtension prog
    withServer (futharkServerCfg progbin (serverOptions opts)) $ \server -> do
      forest <- thresholdForest server
      when (optVerbose opts > 0) $
        putStrLn $
          ("Threshold forest:\n" <>) $
            drawForest (map (fmap show) forest)
  
      fmap fst . foldM (tuneThreshold opts server datasets) ([], mempty) $
        tuningPaths forest
  
  runAutotuner :: AutotuneOptions -> FilePath -> IO ()
  runAutotuner opts prog = do
    best <- tune opts prog
  
    let tuning = T.unlines $ do
          (s, n) <- sortOn fst best
          pure $ s <> "=" <> showText n
  
    case optTuning opts of
      Nothing -> pure ()
      Just suffix -> do
        T.writeFile (prog <.> suffix) tuning
        putStrLn $ "Wrote " ++ prog <.> suffix
  
    T.putStrLn $ "Result of autotuning:\n" <> tuning
  
  supportedBackends :: [String]
  supportedBackends = ["opencl", "cuda", "hip"]
  
  commandLineOptions :: [FunOptDescr AutotuneOptions]
  commandLineOptions =
    [ Option
        "r"
        ["runs"]
        ( ReqArg
            ( \n ->
                case reads n of
                  [(n', "")] | n' >= 0 ->
                    Right $ \config -> config {optMinRuns = n'}
                  _ ->
                    Left $ optionsError $ "'" ++ n ++ "' is not a non-negative integer."
            )
            "RUNS"
        )
        "Run each test case this many times.",
      Option
        []
        ["backend"]
        ( ReqArg
            ( \backend ->
                if backend `elem` supportedBackends
                  then Right $ \config -> config {optBackend = backend}
                  else Left $ optionsError $ "autotuning is only supported for these backends: " <> unwords supportedBackends
            )
            "BACKEND"
        )
        "The backend used (defaults to 'opencl').",
      Option
        []
        ["futhark"]
        ( ReqArg
            (\prog -> Right $ \config -> config {optFuthark = Just prog})
            "PROGRAM"
        )
        "The binary used for operations (defaults to 'futhark').",
      Option
        []
        ["pass-option"]
        ( ReqArg
            ( \opt ->
                Right $ \config ->
                  config {optExtraOptions = opt : optExtraOptions config}
            )
            "OPT"
        )
        "Pass this option to programs being run.",
      Option
        []
        ["tuning"]
        ( ReqArg
            (\s -> Right $ \config -> config {optTuning = Just s})
            "EXTENSION"
        )
        "Write tuning files with this extension (default: .tuning).",
      Option
        []
        ["timeout"]
        ( ReqArg
            ( \n ->
                case reads n of
                  [(n', "")] ->
                    Right $ \config -> config {optTimeout = n'}
                  _ ->
                    Left $ optionsError $ "'" ++ n ++ "' is not a non-negative integer."
            )
            "SECONDS"
        )
        "Initial tuning timeout for each dataset. Later tuning runs are based off of the runtime of the first run.",
      Option
        []
        ["skip-compilation"]
        (NoArg $ Right $ \config -> config {optSkipCompilation = True})
        "Use already compiled program.",
      Option
        "v"
        ["verbose"]
        (NoArg $ Right $ \config -> config {optVerbose = optVerbose config + 1})
        "Enable logging.  Pass multiple times for more.",
      Option
        []
        ["spec-file"]
        (ReqArg (\s -> Right $ \config -> config {optTestSpec = Just s}) "FILE")
        "Use test specification from this file."
    ]
  
  -- | Run @futhark autotune@
  main :: String -> [String] -> IO ()
  main = mainWithOptions initialAutotuneOptions commandLineOptions "options... program" $ \progs config ->
    case progs of
      [prog] -> Just $ runAutotuner config prog
      _ -> Nothing
  