Roofline: classify LATENCY_BOUND for small global-memory transfers

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"] = (

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 |
 

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."""

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