diff --git a/TraceLens/TreePerf/tree_perf.py b/TraceLens/TreePerf/tree_perf.py index e65a92254..47068609e 100644 --- a/TraceLens/TreePerf/tree_perf.py +++ b/TraceLens/TreePerf/tree_perf.py @@ -130,6 +130,69 @@ def get_max_achievable_tflops(perf_model, arch): return maf_specs.get(compute_spec) +_MEM_LATENCY_MISSING_WARNED = False + + +def get_mem_latency_us(arch): + """Return global memory latency floor from arch JSON, or None if unset.""" + if arch is None: + return None + mem_latency_us = arch.get("mem_latency_us") + if mem_latency_us is None: + return None + return float(mem_latency_us) + + +def warn_if_missing_mem_latency(arch): + """ + Warn once when roofline arch has bandwidth but no memory latency floor. + + Without mem_latency_us, roofline classification falls back to the legacy + bandwidth-only model (COMPUTE_BOUND / MEMORY_BOUND only). + """ + global _MEM_LATENCY_MISSING_WARNED + if arch is None or _MEM_LATENCY_MISSING_WARNED: + return + if arch.get("mem_bw_gbps") is not None and arch.get("mem_latency_us") is None: + warnings.warn( + "GPU arch JSON has mem_bw_gbps but no mem_latency_us; roofline memory " + "leg uses bandwidth-only classification (COMPUTE_BOUND / MEMORY_BOUND). " + "Add mem_latency_us (global HBM access latency in microseconds) to enable " + "LATENCY_BOUND for small global-memory transfers.", + UserWarning, + stacklevel=2, + ) + _MEM_LATENCY_MISSING_WARNED = True + + +def compute_roofline_bound( + compute_time_us, bytes_moved, mem_bw_gbps, mem_latency_us=None +): + """ + Classify roofline bound using compute vs global-memory limits. + + When mem_latency_us is provided, memory time is max(latency floor, bytes/bw). + Otherwise falls back to bytes/bw only (legacy behavior). + """ + transfer_time_us = (bytes_moved / (mem_bw_gbps * 1e9)) * 1e6 + if mem_latency_us is not None: + memory_time_us = max(mem_latency_us, transfer_time_us) + if compute_time_us >= memory_time_us: + roofline_bound = "COMPUTE_BOUND" + elif transfer_time_us >= mem_latency_us: + roofline_bound = "MEMORY_BOUND" + else: + roofline_bound = "LATENCY_BOUND" + else: + memory_time_us = transfer_time_us + if compute_time_us >= memory_time_us: + roofline_bound = "COMPUTE_BOUND" + else: + roofline_bound = "MEMORY_BOUND" + roofline_time_us = max(compute_time_us, memory_time_us) + return roofline_time_us, roofline_bound + + def _perf_model_init_kwargs( perf_model_class, event, arch, python_path, enable_origami, inductor_cache_dir=None ): @@ -275,6 +338,7 @@ def __init__( add_python_func = True self.add_python_func = add_python_func self.arch = arch + warn_if_missing_mem_latency(arch) self.python_path = python_path self.enable_origami = enable_origami self.inductor_cache_dir = inductor_cache_dir @@ -492,13 +556,13 @@ def compute_perf_metrics( ): # Compute time: flops / (peak_tflops * 1e12) gives seconds, convert to µs compute_time_us = (gflops * 1e9 / (peak_tflops * 1e12)) * 1e6 - # Memory time: bytes / (bandwidth_gbps * 1e9) gives seconds, convert to µs - memory_time_us = (bytes_moved / (mem_bw_gbps * 1e9)) * 1e6 - roofline_time_us = max(compute_time_us, memory_time_us) - if compute_time_us >= memory_time_us: - roofline_bound = "COMPUTE_BOUND" - else: - roofline_bound = "MEMORY_BOUND" + mem_latency_us = get_mem_latency_us(self.arch) + roofline_time_us, roofline_bound = compute_roofline_bound( + compute_time_us, + bytes_moved, + mem_bw_gbps, + mem_latency_us=mem_latency_us, + ) dict_metrics["Roofline Time (µs)"] = roofline_time_us dict_metrics["Roofline Bound"] = roofline_bound dict_metrics["Pct Roofline"] = ( diff --git a/examples/gpu_arch_example.md b/examples/gpu_arch_example.md index 0eda77caa..0e8ff2871 100644 --- a/examples/gpu_arch_example.md +++ b/examples/gpu_arch_example.md @@ -14,6 +14,7 @@ GPU architecture JSON files define specifications used for roofline analysis in { "name": "MI300X", "mem_bw_gbps": 5300, + "mem_latency_us": 0.3, "max_achievable_tflops": { "matrix_fp16": 654, "matrix_bf16": 708, @@ -36,6 +37,7 @@ GPU architecture JSON files define specifications used for roofline analysis in |-------|-------------| | `name` | GPU model name | | `mem_bw_gbps` | Memory bandwidth in GB/s | +| `mem_latency_us` | (Optional) Global HBM access latency floor in microseconds. Enables `LATENCY_BOUND` roofline classification for small transfers. If omitted, roofline falls back to bandwidth-only (`COMPUTE_BOUND` / `MEMORY_BOUND`) and emits a warning. | | `max_achievable_tflops` | Max achievable TFLOPS by compute type and precision | | `_reference` | (Optional) Source reference for the values | diff --git a/tests/test_roofline_bound.py b/tests/test_roofline_bound.py index 0b9f8f18c..277b991de 100644 --- a/tests/test_roofline_bound.py +++ b/tests/test_roofline_bound.py @@ -13,6 +13,8 @@ from TraceLens.Reporting.generate_perf_report_pytorch import ( generate_perf_report_pytorch, ) +from TraceLens.TreePerf import tree_perf as tree_perf_module + from conftest import list_perf_report_csv_sheets _ORIGAMI_AVAILABLE = importlib.util.find_spec("origami") is not None @@ -44,6 +46,14 @@ ) VALID_BOUND_VALUES = {"COMPUTE_BOUND", "MEMORY_BOUND"} +VALID_BOUND_VALUES_WITH_LATENCY = VALID_BOUND_VALUES | {"LATENCY_BOUND"} + + +@pytest.fixture(autouse=True) +def reset_mem_latency_warning(): + tree_perf_module._MEM_LATENCY_MISSING_WARNED = False + yield + tree_perf_module._MEM_LATENCY_MISSING_WARNED = False @pytest.fixture(scope="module") @@ -67,6 +77,28 @@ def perf_report(tmp_path_factory): return csv_dir +@pytest.fixture(scope="module") +def perf_report_with_mem_latency(tmp_path_factory): + """Generate a perf report with mem_latency_us in the arch JSON.""" + output_dir = tmp_path_factory.mktemp("roofline_bound_latency") + arch_path = str(output_dir / "mi300x.json") + csv_dir = str(output_dir / "perf_report_csvs") + arch = dict(MI300X_ARCH, mem_latency_us=0.3) + + with open(arch_path, "w") as f: + json.dump(arch, f) + + generate_perf_report_pytorch( + profile_json_path=TRACE_PATH, + output_xlsx_path=None, + output_csvs_dir=csv_dir, + gpu_arch_json_path=arch_path, + enable_origami=True, + ) + + return csv_dir + + def _find_col(df, prefix): """Find the Roofline Bound column, which may have an aggregation suffix.""" for col in df.columns: @@ -87,6 +119,23 @@ def test_roofline_bound_in_unified_perf_summary(perf_report): assert bound_vals <= VALID_BOUND_VALUES, f"Unexpected values: {bound_vals}" +def test_roofline_bound_legacy_arch_warns(perf_report): + """Legacy arch JSON without mem_latency_us should warn once during report gen.""" + with pytest.warns(UserWarning, match="mem_latency_us"): + tree_perf_module.warn_if_missing_mem_latency(MI300X_ARCH) + + +def test_roofline_bound_with_mem_latency(perf_report_with_mem_latency): + """With mem_latency_us, LATENCY_BOUND may appear in unified_perf_summary.""" + df = pd.read_csv( + os.path.join(perf_report_with_mem_latency, "unified_perf_summary.csv") + ) + bound_col = _find_col(df, "Roofline Bound") + assert bound_col is not None + bound_vals = set(df[bound_col].dropna().unique()) + assert bound_vals <= VALID_BOUND_VALUES_WITH_LATENCY, f"Unexpected values: {bound_vals}" + + @pytest.mark.skipif(not _ORIGAMI_AVAILABLE, reason="requires origami (rocm-origami)") def test_origami_time_in_unified_perf_summary(perf_report): """Origami Time (µs) column must appear in unified_perf_summary.""" diff --git a/tests/test_roofline_latency_bound.py b/tests/test_roofline_latency_bound.py new file mode 100644 index 000000000..53b94216e --- /dev/null +++ b/tests/test_roofline_latency_bound.py @@ -0,0 +1,96 @@ +############################################################################### +# Copyright (c) 2024 - 2025 Advanced Micro Devices, Inc. All rights reserved. +# +# See LICENSE for license information. +############################################################################### + +"""Unit tests for three-way roofline bound classification.""" + +import pytest + +from TraceLens.TreePerf import tree_perf as tree_perf_module +from TraceLens.TreePerf.tree_perf import ( + compute_roofline_bound, + warn_if_missing_mem_latency, +) + +MI300X_MEM_BW_GBPS = 5300 +MI300X_MEM_LATENCY_US = 0.3 + + +@pytest.fixture(autouse=True) +def reset_mem_latency_warning(): + tree_perf_module._MEM_LATENCY_MISSING_WARNED = False + yield + tree_perf_module._MEM_LATENCY_MISSING_WARNED = False + + +def test_compute_roofline_bound_latency_regime(): + compute_time_us = 0.01 + bytes_moved = 1024 + roofline_time_us, roofline_bound = compute_roofline_bound( + compute_time_us, + bytes_moved, + MI300X_MEM_BW_GBPS, + mem_latency_us=MI300X_MEM_LATENCY_US, + ) + assert roofline_bound == "LATENCY_BOUND" + assert roofline_time_us == pytest.approx(MI300X_MEM_LATENCY_US) + + +def test_compute_roofline_bound_bandwidth_regime(): + compute_time_us = 1.0 + bytes_moved = 16 * 1024 * 1024 + roofline_time_us, roofline_bound = compute_roofline_bound( + compute_time_us, + bytes_moved, + MI300X_MEM_BW_GBPS, + mem_latency_us=MI300X_MEM_LATENCY_US, + ) + assert roofline_bound == "MEMORY_BOUND" + assert roofline_time_us > MI300X_MEM_LATENCY_US + + +def test_compute_roofline_bound_compute_regime(): + compute_time_us = 100.0 + bytes_moved = 1024 + roofline_time_us, roofline_bound = compute_roofline_bound( + compute_time_us, + bytes_moved, + MI300X_MEM_BW_GBPS, + mem_latency_us=MI300X_MEM_LATENCY_US, + ) + assert roofline_bound == "COMPUTE_BOUND" + assert roofline_time_us == pytest.approx(compute_time_us) + + +def test_compute_roofline_bound_fallback_without_mem_latency(): + roofline_time_us, roofline_bound = compute_roofline_bound( + 0.01, + 1024, + MI300X_MEM_BW_GBPS, + mem_latency_us=None, + ) + assert roofline_bound in {"COMPUTE_BOUND", "MEMORY_BOUND"} + assert roofline_bound != "LATENCY_BOUND" + assert roofline_time_us < MI300X_MEM_LATENCY_US + + +def test_warn_if_missing_mem_latency_emits_once(): + arch = {"mem_bw_gbps": MI300X_MEM_BW_GBPS} + with pytest.warns(UserWarning, match="mem_latency_us"): + warn_if_missing_mem_latency(arch) + warn_if_missing_mem_latency(arch) + + +def test_warn_if_missing_mem_latency_skips_when_present(): + import warnings as py_warnings + + arch = { + "mem_bw_gbps": MI300X_MEM_BW_GBPS, + "mem_latency_us": MI300X_MEM_LATENCY_US, + } + with py_warnings.catch_warnings(record=True) as caught: + py_warnings.simplefilter("always") + warn_if_missing_mem_latency(arch) + assert len(caught) == 0