diff --git a/mlir/test/perf-scripts/runtime/tuningRunner-gemm.py b/mlir/test/perf-scripts/runtime/tuningRunner-gemm.py index f92f25a96a44..204cd53464a6 100644 --- a/mlir/test/perf-scripts/runtime/tuningRunner-gemm.py +++ b/mlir/test/perf-scripts/runtime/tuningRunner-gemm.py @@ -1,8 +1,8 @@ # Using a tiny GEMM with --debug-quick-tune-data. This emits a `.debug` # TSV of the per-config table entries (PerfConfig + TFlops) but, unlike the -# full ``--debug`` flag, omits the heavy per-iteration ``MeasurementsMs`` -# arrays. Verify the debug file is produced, has the expected header and -# per-config rows, and that the measurements column is absent. +# full ``--debug`` flag, omits the per-config timing statistics in ``Stats``. +# Verify the debug file is produced, has the expected header and per-config +# rows, and that the statistics column is absent. # # tuningRunner.py drives real GPU tuning, so it needs the ROCm runner / GPU # runtime. @@ -11,7 +11,7 @@ # RUN: tuningRunner.py --op gemm --tuning-space=quick --debug-quick-tune-data \ # RUN: --config='-g 1 -m 64 -n 64 -k 64 -t f32 -out_datatype f32 -transA 0 -transB 0' \ # RUN: -q -o %t2.tsv -# RUN: FileCheck %s --check-prefix=DEBUG --implicit-check-not=MeasurementsMs < %t2.tsv.debug +# RUN: FileCheck %s --check-prefix=DEBUG --implicit-check-not=Stats < %t2.tsv.debug # # DEBUG: PerfConfig{{.*}}TFlops # DEBUG: {{v[0-9]+:}} diff --git a/mlir/tools/rocmlir-tuning-driver/rocmlir-tuning-driver.cpp b/mlir/tools/rocmlir-tuning-driver/rocmlir-tuning-driver.cpp index 9e0111a5a2a1..77a5e4f85dd9 100644 --- a/mlir/tools/rocmlir-tuning-driver/rocmlir-tuning-driver.cpp +++ b/mlir/tools/rocmlir-tuning-driver/rocmlir-tuning-driver.cpp @@ -46,8 +46,10 @@ #include "llvm/Support/CommandLine.h" #include "llvm/Support/ErrorHandling.h" #include "llvm/Support/InitLLVM.h" +#include "llvm/Support/Signals.h" #include "llvm/Support/SourceMgr.h" #include "llvm/Support/ThreadPool.h" +#include "llvm/Support/raw_ostream.h" #include #include @@ -57,6 +59,7 @@ #include #include #include +#include #include "CacheFlush.h" #include "ConcurrentQueue.h" @@ -161,16 +164,16 @@ static llvm::cl::opt sleepUs( static llvm::cl::opt showStats( "show-stats", llvm::cl::desc( - "Print detailed stats (min, max, median, stddev, cv) in JSON format. " - "In case of small kernels print total_cpu_time and number of " + "Print detailed stats (min, max, median, stddev, cv) in nanoseconds as " + "JSON. In case of small kernels print total_cpu_time and number of " "iterations."), llvm::cl::init(false)); static llvm::cl::opt showAllMeasurements( "show-all-measurements", llvm::cl::desc( - "Print all individual timing measurements in JSON format. In case of " - "small kernels print total_cpu_time and number of iterations."), + "Print all individual timing measurements in nanoseconds as JSON. In " + "case of small kernels print total_cpu_time and number of iterations."), llvm::cl::init(false)); static llvm::cl::opt benchmarkConfig( @@ -514,7 +517,7 @@ benchmarkKernels(ArrayRef binaries, // Load all modules once to reduce overhead std::vector modules; std::vector functions; - auto moduleCleanup = llvm::make_scope_exit([&]() { + llvm::scope_exit moduleCleanup([&]() { for (hipModule_t mod : modules) { if (!mod) continue; @@ -537,7 +540,7 @@ benchmarkKernels(ArrayRef binaries, } // Sleep guard to avoid GPU throttling - auto sleepGuard = llvm::make_scope_exit([¶ms] { + llvm::scope_exit sleepGuard([¶ms] { if (params.sleepUs > 0) { std::this_thread::sleep_for(std::chrono::microseconds(params.sleepUs)); } @@ -630,9 +633,16 @@ benchmarkKernels(ArrayRef binaries, return failure(); } + constexpr auto msToNs = [](double ms) { return 1e6 * ms; }; + + // Convert measurements from milliseconds to nanoseconds + for (double &measurement : measurements) { + measurement = msToNs(measurement); + } + if (params.showAllMeasurements) { if (isSmallKernel) { - llvm::outs() << "{\"total_cpu_time\":" << smallKernelCpuMs + llvm::outs() << "{\"total_cpu_time\":" << msToNs(smallKernelCpuMs) << ",\"iterations\":" << iterations << "}\t"; } else { llvm::outs() << "["; @@ -651,7 +661,7 @@ benchmarkKernels(ArrayRef binaries, // We cannot show the rest of the stats because the small kernel case uses // one timer only, so we cannot actually compute the min, max, etc. if (isSmallKernel) { - llvm::outs() << "{\"total_cpu_time\":" << smallKernelCpuMs + llvm::outs() << "{\"total_cpu_time\":" << msToNs(smallKernelCpuMs) << ",\"iterations\":" << iterations << "}\t"; } if (measurements.size() > 1) { @@ -667,11 +677,10 @@ benchmarkKernels(ArrayRef binaries, } } - auto msToNs = [](double ms) { return 1e6 * ms; }; if (params.useMedian) - return msToNs(computeMedian(measurements)); + return computeMedian(measurements); else - return msToNs(computeMean(trimValues(measurements, params.trimPercent))); + return computeMean(trimValues(measurements, params.trimPercent)); } static int toKernelOrder(Attribute attr) { @@ -702,6 +711,46 @@ static LogicalResult extractFuncOps(ModuleOp op, return success(); } +/// Perf config the calling thread is compiling, or null when it is between +/// configs. Points into the config list, which outlives the workers. +static thread_local const SmallString<64> *compilingConfig = nullptr; + +/// Writes to stderr without llvm::errs(), which is neither async-signal-safe +/// nor safe to reenter from a fatal error handler. Give up silently on a failed +/// write: there is nothing left to report it with. +static void writeToStderr(StringRef message) { + ssize_t written = ::write(STDERR_FILENO, message.data(), message.size()); + (void)written; +} + +/// Names the perf config being compiled on the crashing thread. Clears the +/// pointer so that the abort() following report_fatal_error() does not report +/// the same config twice through the signal handler. +static void reportCompilingConfig() { + if (!compilingConfig) + return; + SmallString<256> message; + { + llvm::raw_svector_ostream os(message); + os << "Offending perf config: " << *compilingConfig << "\n" + << "Reproduce with `--benchmark-config=" << *compilingConfig << "`\n"; + } + writeToStderr(message); + compilingConfig = nullptr; +} + +static void compilationFatalErrorHandler(void *, const char *reason, bool) { + SmallString<256> message; + { + llvm::raw_svector_ostream os(message); + os << "LLVM ERROR: " << reason << "\n"; + } + writeToStderr(message); + reportCompilingConfig(); +} + +static void compilationSignalHandler(void *) { reportCompilingConfig(); } + static bool doesModuleHaveFusions(ModuleOp module) { WalkResult result = module.walk([](Operation *op) { // Check for linalg.generic or rock.reduce (standalone fusion ops) @@ -767,7 +816,7 @@ static LogicalResult runTuningLoop(ModuleOp source) { // 3. Create HIP stream and allocate device buffers hipStream_t stream; HIPCHECK(hipStreamCreate(&stream)); - auto streamCleanup = llvm::make_scope_exit([&]() { + llvm::scope_exit streamCleanup([&]() { hipError_t status = hipStreamDestroy(stream); if (status != hipSuccess) { llvm::errs() << "HIP error in hipStreamDestroy: " @@ -776,7 +825,7 @@ static LogicalResult runTuningLoop(ModuleOp source) { }); std::vector gpuBuffers; - auto bufferCleanup = llvm::make_scope_exit([&]() { + llvm::scope_exit bufferCleanup([&]() { for (void *buffer : gpuBuffers) { // hipFree does not allow nullptrs, so make sure to check for it first if (!buffer) @@ -1020,7 +1069,11 @@ static LogicalResult runTuningLoop(ModuleOp source) { if (idx >= configs.size()) break; - if (!compilationResults.push(compileConfig(idx, myRes))) + compilingConfig = &configs[idx]; + CompilationResult result = compileConfig(idx, myRes); + compilingConfig = nullptr; + + if (!compilationResults.push(std::move(result))) break; // Queue terminated } @@ -1036,7 +1089,7 @@ static LogicalResult runTuningLoop(ModuleOp source) { threads.emplace_back(worker); } - auto threadCleanup = llvm::make_scope_exit([&] { + llvm::scope_exit threadCleanup([&] { // In case of early termination, signal all threads to stop compilationResults.terminate(); for (auto &t : threads) { @@ -1104,6 +1157,11 @@ static LogicalResult runTuningLoop(ModuleOp source) { int main(int argc, char **argv) { llvm::InitLLVM y(argc, argv); + // Name the perf config under compilation if a backend pass in the pipeline + // dies, either through report_fatal_error or a crash signal. + llvm::install_fatal_error_handler(compilationFatalErrorHandler); + llvm::sys::AddSignalHandler(compilationSignalHandler, nullptr); + mlir::registerMLIRCLOptions(); llvm::cl::ParseCommandLineOptions(argc, argv, "rocMLIR tuning driver"); diff --git a/mlir/utils/performance/perfRunner.py b/mlir/utils/performance/perfRunner.py index 060f87b33480..c2bdd10d9b04 100644 --- a/mlir/utils/performance/perfRunner.py +++ b/mlir/utils/performance/perfRunner.py @@ -259,6 +259,16 @@ def chip_has_mfma(): return has_feature(lookup_arch_info(get_chip()).default_features, GemmFeatures.MFMA) +def chip_uses_wgp_mode(chip: Optional[str] = None) -> bool: + """Whether the chip is built in workgroup processor (WGP) mode. + + In WGP mode HIP counts one multiprocessor per two-CU workgroup processor, so the CU + count it reports is half of the physical one that rocminfo prints. Wave size is the + same proxy for WGP mode that fixNaviProperties() uses in AmdArchDb.cpp. + """ + return lookup_arch_info(chip or get_chip()).wave_size == 32 + + DATA_TYPES_ATTENTION = None diff --git a/mlir/utils/performance/tests/test_perfRunner.py b/mlir/utils/performance/tests/test_perfRunner.py index 66898c6dbe23..3696bae61cd7 100644 --- a/mlir/utils/performance/tests/test_perfRunner.py +++ b/mlir/utils/performance/tests/test_perfRunner.py @@ -168,6 +168,31 @@ def fake_lookup(arch): assert captured["arch"] == "native:1" +class TestChipUsesWgpMode: + """Tests for chip_uses_wgp_mode (wave size proxy for WGP mode).""" + + def test_wave64_chip_is_not_wgp_mode(self, monkeypatch): + monkeypatch.setattr(perfRunner, "lookup_arch_info", + lambda arch: types.SimpleNamespace(wave_size=64)) + assert not perfRunner.chip_uses_wgp_mode("gfx942") + + def test_wave32_chip_is_wgp_mode(self, monkeypatch): + monkeypatch.setattr(perfRunner, "lookup_arch_info", + lambda arch: types.SimpleNamespace(wave_size=32)) + assert perfRunner.chip_uses_wgp_mode("gfx1100") + + def test_defaults_to_current_chip(self, monkeypatch): + captured = {} + + def fake_lookup(arch): + captured["arch"] = arch + return types.SimpleNamespace(wave_size=32) + + monkeypatch.setattr(perfRunner, "lookup_arch_info", fake_lookup) + assert perfRunner.chip_uses_wgp_mode() + assert captured["arch"] == "gfx900" + + class TestParseDataTypes: """Tests for parse_data_types (gemm data types).""" diff --git a/mlir/utils/performance/tests/test_tuningRunner.py b/mlir/utils/performance/tests/test_tuningRunner.py index 5ed9f03029da..25a3fcaecc48 100644 --- a/mlir/utils/performance/tests/test_tuningRunner.py +++ b/mlir/utils/performance/tests/test_tuningRunner.py @@ -11,6 +11,7 @@ import tempfile import threading import time +import types from pathlib import Path import pytest @@ -214,6 +215,58 @@ def test_running_becomes_crashed_on_load(self): if os.path.exists(path): os.unlink(path) + def _ctx_key(self, num_cu=None): + return f"{self._ARCH}/{num_cu or self._NUM_CU}/{self._NUM_CHIPLETS}/full" + + def _write_state_file(self, contexts): + with tempfile.NamedTemporaryFile(mode="w", suffix=".state", delete=False) as f: + f.write(json.dumps({"contexts": contexts})) + return f.name + + def test_legacy_wgp_context_is_merged_and_migrated(self, monkeypatch): + """A context keyed by the doubled CU count belongs to this device on WGP-mode chips.""" + monkeypatch.setattr(perfRunner, "lookup_arch_info", + lambda arch: types.SimpleNamespace(wave_size=32)) + legacy_key = self._ctx_key(2 * self._NUM_CU) + path = self._write_state_file({ + legacy_key: { + self._TV_A: "crashed", + self._TV_B: "failed" + }, + self._ctx_key(): { + self._TV_A: "timed_out" + }, + }) + try: + sf = self._make_state_file(path) + # The active context wins where the two overlap + assert sf.state.configs.get(self._TV_A) == ConfigState.TIMED_OUT + assert sf.state.configs.get(self._TV_B) == ConfigState.FAILED + + with open(path, 'r') as f: + contexts = json.load(f)["contexts"] + assert legacy_key not in contexts + assert contexts[self._ctx_key()] == {self._TV_A: "timed_out", self._TV_B: "failed"} + finally: + if os.path.exists(path): + os.unlink(path) + + def test_doubled_cu_context_untouched_in_cu_mode(self, monkeypatch): + monkeypatch.setattr(perfRunner, "lookup_arch_info", + lambda arch: types.SimpleNamespace(wave_size=64)) + other_key = self._ctx_key(2 * self._NUM_CU) + path = self._write_state_file({other_key: {self._TV_A: "failed"}}) + try: + sf = self._make_state_file(path) + assert sf.state.is_empty() + + with open(path, 'r') as f: + contexts = json.load(f)["contexts"] + assert contexts == {other_key: {self._TV_A: "failed"}} + finally: + if os.path.exists(path): + os.unlink(path) + def test_old_state_file_configs_are_canonicalized(self): """Non-canonical test vectors in state file are canonicalized on load.""" non_canonical = "-g 1 -m 1024 -k 769 -n 512 -t f32 -out_datatype f32 -transA false -transB false" @@ -293,6 +346,53 @@ def test_parse_new_format_tsv(self): finally: os.unlink(path) + def _write_gemm_tsv(self, tv, num_cu, arch="gfx1100"): + with tempfile.NamedTemporaryFile(mode="w", suffix=".tsv", delete=False) as f: + f.write( + "# arch\tnumCUs\tnumChiplets\ttestVector\tperfConfig\tTFlops\ttuningSpace\tcommitId\ttimestamp\tdurationSec\n" + ) + f.write( + f"{arch}\t{num_cu}\t1\t{tv}\tperf_best\t1.5\tfull\tabc123\t2025-01-01T00:00:00Z\t10.0\n" + ) + return f.name + + def test_legacy_doubled_num_cu_loaded_in_wgp_mode(self, monkeypatch): + tv = "-t f32 -out_datatype f32 -transA false -transB false -g 1 -m 1024 -n 512 -k 769" + path = self._write_gemm_tsv(tv, num_cu=128) + monkeypatch.setattr(perfRunner, "lookup_arch_info", + lambda arch: types.SimpleNamespace(wave_size=32)) + try: + opts = self._options(path, arch="gfx1100", num_cu=64) + cache = TunedConfigsCache.from_output_file(opts, GemmConfiguration) + assert cache.count() == 1 + assert cache.contains(tv) + finally: + os.unlink(path) + + def test_doubled_num_cu_not_loaded_in_cu_mode(self, monkeypatch): + tv = "-t f32 -out_datatype f32 -transA false -transB false -g 1 -m 1024 -n 512 -k 769" + path = self._write_gemm_tsv(tv, num_cu=128, arch="gfx942") + monkeypatch.setattr(perfRunner, "lookup_arch_info", + lambda arch: types.SimpleNamespace(wave_size=64)) + try: + opts = self._options(path, arch="gfx942", num_cu=64) + cache = TunedConfigsCache.from_output_file(opts, GemmConfiguration) + assert cache.count() == 0 + finally: + os.unlink(path) + + def test_unrelated_num_cu_not_loaded_in_wgp_mode(self, monkeypatch): + tv = "-t f32 -out_datatype f32 -transA false -transB false -g 1 -m 1024 -n 512 -k 769" + path = self._write_gemm_tsv(tv, num_cu=96) + monkeypatch.setattr(perfRunner, "lookup_arch_info", + lambda arch: types.SimpleNamespace(wave_size=32)) + try: + opts = self._options(path, arch="gfx1100", num_cu=64) + cache = TunedConfigsCache.from_output_file(opts, GemmConfiguration) + assert cache.count() == 0 + finally: + os.unlink(path) + def test_arch_mismatch_not_loaded(self): with tempfile.NamedTemporaryFile(mode="w", suffix=".tsv", delete=False) as f: f.write( @@ -552,6 +652,24 @@ def test_tuning_space_choices(self): ) assert parsed.tuning_space == "quick" + def _parse_output(self, config_args): + topology = _make_mock_gpu_topology([(0, "gfx900")]) + parsed = tuningRunner.parse_arguments(topology, [0], ["--op", "gemm"] + config_args) + return parsed.output + + def test_output_defaults_to_configs_file_plus_tsv(self): + derived = self._parse_output(["-c", "configs/tier1-gemm-configs"]) + assert derived == "configs/tier1-gemm-configs.tsv" + + def test_explicit_output_overrides_configs_file_name(self): + explicit = self._parse_output(["-c", "configs/gemm.txt", "-o", "/tmp/out.tsv"]) + assert explicit == "/tmp/out.tsv" + + def test_output_falls_back_without_a_configs_file(self): + single_config = self._parse_output(["--config", "-g 1 -m 1024 -k 769 -n 512"]) + assert single_config == tuningRunner.DEFAULT_OUTPUT_FILE + assert self._parse_output(["-c", "-"]) == tuningRunner.DEFAULT_OUTPUT_FILE + def test_negative_gpu_run_timeout_rejected(self, capsys): topology = _make_mock_gpu_topology([(0, "gfx900")]) available = [0] diff --git a/mlir/utils/performance/tuningRunner.py b/mlir/utils/performance/tuningRunner.py index 7b91e6071ba3..e55b1d8b7fc5 100755 --- a/mlir/utils/performance/tuningRunner.py +++ b/mlir/utils/performance/tuningRunner.py @@ -81,6 +81,9 @@ 'commitId', 'timestamp', 'durationSec' ] +# Used when there is no configs file to name the results after +DEFAULT_OUTPUT_FILE = "tuning_results_local.tsv" + # Only these operation types support GPU validation # Keep in sync with isGpuValidationSupported() in rocmlir-gen.cpp # ConvConfiguration covers both fwd and bwd @@ -450,6 +453,11 @@ def __init__(self, filepath: Optional[str], chip: str, arch: str, num_cu: int, num_chiplets: int, tuning_space: str, conf_class: type): self.filepath = filepath self.context_key = f"{chip}/{num_cu}/{num_chiplets}/{tuning_space}" + # State written before the CU count came from HIP holds the doubled, rocminfo-reported + # count on WGP-mode chips. Such a context describes this device, so adopt it. + self.legacy_context_key = None + if num_cu and perfRunner.chip_uses_wgp_mode(chip): + self.legacy_context_key = f"{chip}/{2 * num_cu}/{num_chiplets}/{tuning_space}" self._conf_class = conf_class self._arch = arch self._num_cu = num_cu @@ -464,6 +472,9 @@ def __init__(self, filepath: Optional[str], chip: str, arch: str, num_cu: int, def _load(self) -> None: """Load state from file. + A legacy context (see legacy_context_key) is merged into the active one and dropped + from the file, so the next save leaves a single context for this device. + For the active context only: - INTERRUPTED configs are removed (will be retried) - RUNNING configs become CRASHED (stale = crash) @@ -476,28 +487,33 @@ def _load(self) -> None: data = json.load(f) self._all_contexts = data['contexts'] - # Process configs for active context with state transitions - if self.context_key in self._all_contexts: - for tv, state_str in self._all_contexts[self.context_key].items(): - try: - state = ConfigState(state_str) - except ValueError: - logger.warning(f"Unknown state '{state_str}' for config '{tv}' in state file") - continue + # Legacy configs go first so the active context wins where the two overlap + if self.legacy_context_key: + self._load_context(self._all_contexts.pop(self.legacy_context_key, {})) + self._load_context(self._all_contexts.get(self.context_key, {})) - if state == ConfigState.INTERRUPTED: - continue - if state == ConfigState.RUNNING: - state = ConfigState.CRASHED + def _load_context(self, configs: Dict[str, str]) -> None: + """Apply the state transitions of a single context to the in-memory state.""" + for tv, state_str in configs.items(): + try: + state = ConfigState(state_str) + except ValueError: + logger.warning(f"Unknown state '{state_str}' for config '{tv}' in state file") + continue - try: - canonical_tv = canonicalize_test_vector(tv, self._conf_class, self._arch, - self._num_cu, self._num_chiplets) - except ValueError as e: - logger.debug(f"Failed to canonicalize config in state file: {e}") - canonical_tv = tv # Keep the raw key so it survives a save/load round-trip + if state == ConfigState.INTERRUPTED: + continue + if state == ConfigState.RUNNING: + state = ConfigState.CRASHED + + try: + canonical_tv = canonicalize_test_vector(tv, self._conf_class, self._arch, + self._num_cu, self._num_chiplets) + except ValueError as e: + logger.debug(f"Failed to canonicalize config in state file: {e}") + canonical_tv = tv # Keep the raw key so it survives a save/load round-trip - self._state.configs[canonical_tv] = state + self._state.configs[canonical_tv] = state @property def state(self) -> TuningState: @@ -575,6 +591,18 @@ def get_state_filepath(output_filepath: str) -> Optional[str]: # ============================================================================= +def matches_current_num_cu(file_num_cu: str, options: Options) -> bool: + """Check a CU count recorded in an output file against the current device. + + On chips running in WGP mode the recorded count used to be the physical CU count read + from rocminfo, which is twice the count HIP reports and that we write today. Accept + that doubled value so results tuned before the switch are still reused. + """ + if file_num_cu == str(options.num_cu): + return True + return (perfRunner.chip_uses_wgp_mode(options.chip) and file_num_cu == str(2 * options.num_cu)) + + @dataclass(frozen=True) class TunedConfigsCache: """Cache for previously tuned configurations loaded from output file.""" @@ -684,7 +712,7 @@ def _parse_data_line(fields: List[str], column_indices: Dict[str, int], options: A line is valid if: - arch matches current system (chip or arch for backwards compatibility) - - numCUs and numChiplets match current system + - numCUs (see matches_current_num_cu) and numChiplets match current system - tuning space matches (from column or header) - testVector is present, parseable, and belongs to the expected operation - perfConfig is present and not 'None' @@ -703,7 +731,7 @@ def get_field(name: str) -> Optional[str]: # Check numCUs match file_num_cu = get_field('numCUs') - if file_num_cu and file_num_cu != str(options.num_cu): + if file_num_cu and not matches_current_num_cu(file_num_cu, options): return None # Check numChiplets match @@ -1120,45 +1148,6 @@ def raise_if_terminated(returncode: int) -> None: raise KeyboardInterrupt() -class TuningArgumentParser(argparse.ArgumentParser): - """ArgumentParser with custom validation for tuning arguments.""" - - def __init__(self, *args, gpu_topology: Optional[GpuTopology] = None, **kwargs): - super().__init__(*args, **kwargs) - self._gpu_topology = gpu_topology - - def parse_args(self, args=None, namespace=None): - parsed = super().parse_args(args, namespace) - - op_type = Operation.from_name(parsed.op) - - if op_type == Operation.FUSION and not parsed.test_dir: - self.error("argument --op=fusion: requires --test-dir to be specified") - - if parsed.test_dir and op_type != Operation.FUSION: - self.error("argument --test-dir: only allowed with --op=fusion") - - if parsed.verify_perf_configs and parsed.verify_mode == "none": - self.error("argument --verify-perf-configs: not allowed with --verify-mode=none") - - if self._gpu_topology and not self._gpu_topology.validate_homogeneity(parsed.gpus): - details = ", ".join(f"GPU {g}: {self._gpu_topology.gpus[g].sku}" for g in parsed.gpus) - self.error(f"argument --gpus: mixed GPU models not supported. Found: {details}") - - return parsed - - -class UniqueChoicesAction(argparse.Action): - """Argparse action that ensures no duplicate values.""" - - def __call__(self, parser, namespace, values, option_string=None): - if len(values) != len(set(values)): - duplicates = [v for v in values if values.count(v) > 1] - parser.error( - f"argument {option_string}: duplicate values not allowed: {set(duplicates)}") - setattr(namespace, self.dest, values) - - @functools.lru_cache(maxsize=1) def get_git_commit_hash() -> str: """Get the current git commit hash.""" @@ -1416,10 +1405,10 @@ def find_best_perfconfig( try: if time == "N/A": nano_seconds = np.nan - measurements = None + stats = None else: nano_seconds = float(time) - measurements = json.loads(parts[1]) if len(parts) == 3 else None + stats = json.loads(parts[1]) if len(parts) == 3 else None except (ValueError, json.JSONDecodeError): gpu_logger.debug(f"Skipping malformed tuning output line: '{result}'") continue @@ -1427,7 +1416,7 @@ def find_best_perfconfig( config.set_perfconfig(perfconfig) entry = config.table_entry(nano_seconds) if options.debug: - entry["MeasurementsMs"] = measurements + entry["Stats"] = stats entries.append(entry) if options.verify_perfconfigs and not np.isnan(nano_seconds): @@ -1454,7 +1443,7 @@ def tune_config(test_vector: str, conf_class: type, paths: Paths, options: Optio f"--warmup-iterations={WARMUP_ITERATIONS}", "--use-median", f"--sleep-us={SLEEP_US}", - f"--show-all-measurements={options.debug}", + f"--show-stats={options.debug}", f"--num-compile-threads={num_compile_threads}", f"--wait-for-compiles={options.wait_for_compiles}", f"--gpu-run-timeout={options.gpu_run_timeout}", @@ -1930,10 +1919,65 @@ def canonicalize_test_vector(tv: str, conf_class: type, arch: str, num_cu: int, # ============================================================================= -# Entry Point +# Argument Parsing # ============================================================================= +class TuningArgumentParser(argparse.ArgumentParser): + """ArgumentParser with custom validation for tuning arguments.""" + + def __init__(self, *args, gpu_topology: Optional[GpuTopology] = None, **kwargs): + super().__init__(*args, **kwargs) + self._gpu_topology = gpu_topology + + def parse_args(self, args=None, namespace=None): + parsed = super().parse_args(args, namespace) + + op_type = Operation.from_name(parsed.op) + + if parsed.gpu_run_timeout < 0: + self.error("argument --gpu-run-timeout: must be non-negative") + + if op_type == Operation.FUSION and not parsed.test_dir: + self.error("argument --op=fusion: requires --test-dir to be specified") + + if parsed.test_dir and op_type != Operation.FUSION: + self.error("argument --test-dir: only allowed with --op=fusion") + + if parsed.verify_perf_configs and parsed.verify_mode == "none": + self.error("argument --verify-perf-configs: not allowed with --verify-mode=none") + + if self._gpu_topology and not self._gpu_topology.validate_homogeneity(parsed.gpus): + details = ", ".join(f"GPU {g}: {self._gpu_topology.gpus[g].sku}" for g in parsed.gpus) + self.error(f"argument --gpus: mixed GPU models not supported. Found: {details}") + + if parsed.output is None: + parsed.output = default_output_path(parsed.configs_file) + + return parsed + + +def default_output_path(configs_file: Optional[str]) -> str: + """Name the results after the configs file, e.g. 'tier1-gemm-configs' -> 'tier1-gemm-configs.tsv'. + + Falls back to DEFAULT_OUTPUT_FILE when there is no name to derive from, i.e. for stdin, --config and --test-dir. + """ + if not configs_file or configs_file == '-': + return DEFAULT_OUTPUT_FILE + return f"{configs_file}.tsv" + + +class UniqueChoicesAction(argparse.Action): + """Argparse action that ensures no duplicate values.""" + + def __call__(self, parser, namespace, values, option_string=None): + if len(values) != len(set(values)): + duplicates = [v for v in values if values.count(v) > 1] + parser.error( + f"argument {option_string}: duplicate values not allowed: {set(duplicates)}") + setattr(namespace, self.dest, values) + + def parse_arguments(gpu_topology: GpuTopology, available_gpus: List[int], args=None) -> argparse.Namespace: @@ -1979,10 +2023,11 @@ def parse_arguments(gpu_topology: GpuTopology, "-o", "--output", type=str, - default="tuning_results_local.tsv", + default=None, metavar='FILE', help= - "Output file path for tuning results in TSV format. Results will be appended if file exists. Use '-' for stdout." + f"Output file path for tuning results in TSV format. Results will be appended if file exists. Use '-' for stdout. " + f"Defaults to the --configs-file path with '.tsv' appended, or '{DEFAULT_OUTPUT_FILE}' if there is no configs file." ) parser.add_argument( @@ -2006,13 +2051,15 @@ def parse_arguments(gpu_topology: GpuTopology, "--debug", action='store_true', default=False, - help="Enable debug output including detailed per-iteration measurements") + help="Enable debug output including per-config timing statistics") - parser.add_argument("--debug-quick-tune-data", - action='store_true', - default=False, - help="Enable debug output for quick tuning data generation without the " - "detailed per-iteration measurement arrays") + parser.add_argument( + "--debug-quick-tune-data", + action='store_true', + default=False, + help= + "Enable debug output for quick tuning data generation without the per-config timing statistics" + ) parser.add_argument("--tuning-space", default="full", @@ -2140,10 +2187,12 @@ def parse_arguments(gpu_topology: GpuTopology, default=False, help="Only show tuning status without performing any tuning") - parsed_args = parser.parse_args(args) - if parsed_args.gpu_run_timeout < 0: - parser.error("argument --gpu-run-timeout: must be non-negative") - return parsed_args + return parser.parse_args(args) + + +# ============================================================================= +# Main +# ============================================================================= def main(args=None):