futhark-0.26.3: src/Futhark/CLI/Autotune.hs
-- | @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
{ 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 . maybe (Right ()) Left =<< cmdSetTuningParam server name 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
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
[ (if fancyTerminal then "\r\ESC[K" else "")
<> "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
[ (if fancyTerminal then "\r\ESC[K" else "") ++ "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 [(if fancyTerminal then "\r\ESC[K" else "") ++ "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
let paths = tuningPaths forest
let total = length paths
when (optVerbose opts > 0) $
putStrLn $
("Threshold forest:\n" <>) $
drawForest (map (fmap show) forest)
counter <- newIORef (0 :: Int)
let dataset_names =
T.intercalate "," $ map (\(name, _, ep) -> ep <> ":" <> name) datasets
result <-
foldM
( \acc tp@(v, _) -> do
modifyIORef' counter (+ 1)
n <- readIORef counter
let bar =
"\rTuning: "
++ show n
++ "/"
++ show total
++ " "
++ T.unpack
( progressBar
ProgressBar
{ progressBarSteps = 10,
progressBarBound = 1,
progressBarElapsed = fromIntegral n / fromIntegral total
}
)
++ T.unpack v
++ " on "
++ T.unpack 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 (ServerException s) = do
T.hPutStrLn stderr s
exitFailure
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
`catch` onServerException
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