Function add_proton_records

benchmarking/insert_proton_records.py:34–74

Add proton.record statements to the FMMS kernel TTGIR file. Uses a two-pass approach: first scan for insertion points, then insert from bottom to top so line indices remain stable.

(input_file: str)

def add_proton_records(input_file: str) -> None:
    """Add proton.record statements to the FMMS kernel TTGIR file.

    Uses a two-pass approach: first scan for insertion points, then insert
    from bottom to top so line indices remain stable.
    """
    with open(input_file) as f:
        lines = f.readlines()

    if any("proton.record" in line for line in lines):
        raise AssertionError(
            "File already contains `proton.record` statements! "
            "Re-dump the TTGIR without proton scopes first."
        )

    insertions = _find_insertion_points(lines)
    _validate_insertions(insertions, input_file)

    # Insert from bottom to top so earlier indices stay valid.
    # At the same line index, the last-processed item ends up first in the file.
    # Secondary key: "end" before "start" in the output (so end=0, start=1).
    # Tertiary key: "kernel" is the outermost scope, so its start must appear
    # first (processed last → nesting=0 for start) and its end must appear
    # last (processed first → nesting=0 for end).
    def _sort_key(item):
        line_idx, text = item
        is_start = 1 if text.startswith("start") else 0
        is_kernel = '"kernel"' in text
        # For starts at same line: kernel (nesting=0) processed last → first in file
        # For ends at same line: kernel (nesting=1) processed first → last in file
        nesting = (0 if is_kernel else 1) if is_start else (1 if is_kernel else 0)
        return (line_idx, is_start, nesting)

    for line_idx, text in sorted(insertions, key=_sort_key, reverse=True):
        lines.insert(line_idx, f"      proton.record {text} loc(#loc1)\n")

    with open(input_file, "w") as f:
        f.writelines(lines)

    scopes = sorted({t.split('"')[1] for _, t in insertions})
    print(f"Added proton records to {input_file}: {scopes}")


def _find_insertion_points(lines: list[str]) -> list[tuple[int, str]]: