"""
Reproducibility tests for Proton.
Each test should invoke one or more GPU kernels and check the validity of their profiling results.
"""

import torch
import triton
import triton.profiler as proton
import json
import pytest
from typing import NamedTuple
import pathlib
import threading

import triton.language as tl
from triton.profiler.hooks.launch import COMPUTE_METADATA_SCOPE_NAME
import triton.profiler.hooks.launch as proton_launch
import triton.profiler.viewer as viewer
from triton._internal_testing import is_hip, is_cuda, is_blackwell


@pytest.mark.parametrize("context", ["shadow", "python"])
def test_torch(context, tmp_path: pathlib.Path, device: str):
    temp_file = tmp_path / "test_torch.hatchet"
    proton.start(str(temp_file.with_suffix("")), context=context)
    proton.enter_scope("test")
    torch.ones((2, 2), device=device)
    proton.exit_scope()
    proton.finalize()
    with temp_file.open() as f:
        data = json.load(f)
    if context == "shadow":
        assert len(data[0]["children"]) == 1
        assert data[0]["children"][0]["frame"]["name"] == "test"
        assert data[0]["children"][0]["children"][0]["metrics"]["time (ns)"] > 0
    elif context == "python":
        assert len(data[0]["children"]) == 1
        # bfs search until find the "elementwise_kernel" and then check its children
        queue = [data[0]]
        import re
        while len(queue) > 0:
            parent_frame = queue.pop(0)
            for child in parent_frame["children"]:
                if "elementwise_kernel" in child["frame"]["name"]:
                    assert len(child["children"]) == 0
                    # check the regex of the parent name matches
                    # file_name:line_number@function_name
                    regex = r".+:\d+@.+"
                    assert re.match(regex, parent_frame["frame"]["name"])
                    return
                queue.append(child)


def test_triton(tmp_path: pathlib.Path, device: str):

    @triton.jit
    def foo(x, y):
        tl.store(y, tl.load(x))

    x = torch.tensor([2], device=device)
    y = torch.zeros_like(x)
    temp_file = tmp_path / "test_triton.hatchet"
    proton.start(str(temp_file.with_suffix("")))
    with proton.scope("test0"):
        with proton.scope("test1"):
            foo[(1, )](x, y)
    with proton.scope("test2"):
        foo[(1, )](x, y)
    proton.finalize()
    with temp_file.open() as f:
        data = json.load(f)
    assert len(data[0]["children"]) == 2
    assert data[0]["children"][0]["frame"]["name"] == "test0"
    assert len(data[0]["children"][0]["children"]) == 1
    assert data[0]["children"][0]["children"][0]["frame"]["name"] == "test1"
    assert data[0]["children"][1]["frame"]["name"] == "test2"


@pytest.mark.skipif(not is_cuda(), reason="HIP backend does not reliably attribute cudagraph replay launches to scopes")
def test_cudagraph(tmp_path: pathlib.Path, device: str):
    stream = torch.cuda.Stream()
    torch.cuda.set_stream(stream)

    def metadata_fn(grid: tuple, metadata: NamedTuple, args: dict):
        return {"name": "foo_test"}

    @triton.jit(launch_metadata=metadata_fn)
    def foo(x, y, z):
        tl.store(z, tl.load(y) + tl.load(x))

    def fn():
        a = torch.ones((2, 2), device=device)
        b = torch.ones((2, 2), device=device)
        c = a + b
        foo[(1, )](a, b, c)

    temp_file = tmp_path / "test_cudagraph.hatchet"
    proton.start(str(temp_file.with_suffix("")), context="shadow")

    # warmup
    # four kernels
    fn()

    # no kernels
    g = torch.cuda.CUDAGraph()
    with torch.cuda.graph(g):
        for i in range(10):
            with proton.scope(f"iter_{i}"):
                fn()

    with proton.scope("test0"):
        g.replay()

    with proton.scope("test1"):
        g.replay()

    g.reset()
    proton.finalize()

    with temp_file.open() as f:
        data = json.load(f)
    # CUDA/HIP graph may also invoke additional kernels to reset outputs
    # {torch.ones, add, foo, test}
    assert len(data[0]["children"]) >= 4
    # find the test frame
    test0_frame = None
    test1_frame = None
    for child in data[0]["children"]:
        if child["frame"]["name"] == "test0":
            test0_frame = child
        if child["frame"]["name"] == "test1":
            test1_frame = child
    assert test0_frame is not None
    assert test1_frame is not None
    # {torch.ones, add, foo}
    if is_hip():
        assert len(test0_frame["children"]) >= 2
        assert test0_frame["children"][0]["metrics"]["time (ns)"] > 0
    else:
        # cuda backend supports "<captured_at>" annotation
        for test_frame in [test0_frame, test1_frame]:
            child = test_frame["children"][0]
            assert child["frame"]["name"] == "<captured_at>"
            # 0...9 iterations
            assert len(child["children"]) == 10
            # check all iterations
            for i in range(10):
                assert child["children"][i]["frame"]["name"] == f"iter_{i}"
                assert child["children"][i]["children"][0]["metrics"]["time (ns)"] > 0


@pytest.mark.skipif(not is_cuda(), reason="Only CUDA backend supports cudagraph replay")
def test_cudagraph_not_captured_by_profiler(tmp_path: pathlib.Path, capfd, device: str):
    stream = torch.cuda.Stream()
    torch.cuda.set_stream(stream)

    @triton.jit
    def foo(x, y, z):
        tl.store(z, tl.load(y) + tl.load(x))

    def fn():
        a = torch.ones((2, 2), device=device)
        b = torch.ones((2, 2), device=device)
        c = a + b
        foo[(1, )](a, b, c)

    # Build/capture graph before profiler starts.
    fn()
    g = torch.cuda.CUDAGraph()
    with torch.cuda.graph(g):
        fn()

    temp_file = tmp_path / "test_cudagraph_not_captured_by_profiler.hatchet"
    proton.start(str(temp_file.with_suffix("")), context="shadow")
    with proton.scope("replay0"):
        g.replay()
    with proton.scope("replay1"):
        g.replay()
    proton.finalize()

    captured = capfd.readouterr()
    assert captured.err.count("Cannot find graph for graphExecId:") == 1
    assert "start profiling before the graph is created" in captured.err

    with temp_file.open() as f:
        data = json.load(f)
    replay0_frame = None
    replay1_frame = None
    for child in data[0]["children"]:
        if child["frame"]["name"] == "replay0":
            replay0_frame = child
        elif child["frame"]["name"] == "replay1":
            replay1_frame = child
    assert replay0_frame is not None
    assert replay1_frame is not None
    assert len(replay0_frame["children"]) == 3
    assert len(replay1_frame["children"]) == 3

    def has_positive_time_metric(node):
        if node["metrics"].get("time (ns)", 0) > 0:
            return True
        return any(has_positive_time_metric(child) for child in node["children"])

    assert has_positive_time_metric(replay0_frame)
    assert has_positive_time_metric(replay1_frame)


@pytest.mark.skipif(not is_cuda(), reason="Only CUDA backend supports cudagraph deactivation")
def test_cudagraph_deactivate(tmp_path, device: str):
    stream = torch.cuda.Stream()
    torch.cuda.set_stream(stream)

    @triton.jit
    def foo(x, y, z):
        tl.store(z, tl.load(y) + tl.load(x))

    def fn(session):
        with proton.scope("scope_a"):
            a = torch.ones((2, 2), device=device)
        proton.deactivate(session)
        with proton.scope("scope_b"):
            b = torch.ones((2, 2), device=device)
        proton.activate(session)
        with proton.scope("scope_c"):
            c = a + b
        foo[(1, )](a, b, c)

    temp_file = tmp_path / "test_cudagraph_deactivate.hatchet"
    session = proton.start(str(temp_file.with_suffix("")), context="shadow", hook="triton")

    # warmup
    fn(session)

    # no kernels
    g = torch.cuda.CUDAGraph()
    with torch.cuda.graph(g):
        for i in range(10):
            with proton.scope(f"iter_{i}"):
                fn(session)

    with proton.scope("test0"):
        g.replay()

    proton.finalize()

    with temp_file.open() as f:
        data = json.load(f)

    # scope a and c should be recorded, b should be skipped
    children = data[0]["children"]
    test0_frame = None
    for child in children:
        if child["frame"]["name"] == "test0":
            test0_frame = child
            break
    assert test0_frame is not None
    iter_frame = test0_frame["children"][0]["children"][0]
    scope_a_frame = None
    scope_b_frame = None
    scope_c_frame = None
    for child in iter_frame["children"]:
        if child["frame"]["name"] == "scope_a":
            scope_a_frame = child
        if child["frame"]["name"] == "scope_b":
            scope_b_frame = child
        if child["frame"]["name"] == "scope_c":
            scope_c_frame = child
    assert scope_a_frame is not None
    assert scope_b_frame is None
    assert scope_c_frame is not None


def test_metrics(tmp_path: pathlib.Path, device: str):

    @triton.jit
    def foo(x, y):
        tl.store(y, tl.load(x))

    x = torch.tensor([2], device=device)
    y = torch.zeros_like(x)
    temp_file = tmp_path / "test_metrics.hatchet"
    proton.start(str(temp_file.with_suffix("")))
    with proton.scope("test0", {"foo": 1.0, "bar": [1, 2, 3], "baz": [1.0, 2.0, 3.0]}):
        foo[(1, )](x, y)
    proton.finalize()
    with temp_file.open() as f:
        data = json.load(f)
    assert len(data[0]["children"]) == 1
    assert data[0]["children"][0]["frame"]["name"] == "test0"
    assert data[0]["children"][0]["metrics"]["foo"] == 1.0
    assert data[0]["children"][0]["metrics"]["bar"] == [1, 2, 3]


def test_scope_backward(tmp_path: pathlib.Path, device: str):
    temp_file = tmp_path / "test_scope_backward.hatchet"
    proton.start(str(temp_file.with_suffix("")))
    with proton.scope("ones1"):
        a = torch.ones((100, 100), device=device, requires_grad=True)
    with proton.scope("plus"):
        a2 = a * a * a
    with proton.scope("ones2"):
        loss = torch.ones_like(a2)

    # Backward triggers two kernels in a single scope
    with proton.scope("backward"):
        a2.backward(loss)
    proton.finalize()
    with temp_file.open() as f:
        data = json.load(f)
    assert len(data[0]["children"]) == 4


def test_cpu_timed_scope(tmp_path: pathlib.Path, device: str):
    temp_file = tmp_path / "test_cpu_timed_scope.hatchet"
    proton.start(str(temp_file.with_suffix("")))
    with proton.cpu_timed_scope("test0"):
        with proton.cpu_timed_scope("test1"):
            torch.ones((100, 100), device=device)
    proton.finalize()
    with temp_file.open() as f:
        data = json.load(f)
    assert len(data[0]["children"]) == 1
    test0_frame = data[0]["children"][0]
    assert test0_frame["metrics"]["cpu_time (ns)"] > 0
    test1_frame = test0_frame["children"][0]
    assert test1_frame["metrics"]["cpu_time (ns)"] > 0


def test_get_data(tmp_path: pathlib.Path, device: str):
    temp_file = tmp_path / "test_tree_json.hatchet"
    session = proton.start(str(temp_file.with_suffix("")), context="shadow")

    @triton.jit
    def foo(x, y, size: tl.constexpr):
        offs = tl.arange(0, size)
        tl.store(y + offs, tl.load(x + offs))

    with proton.scope("test"):
        x = torch.ones((2, 2), device=device)
        foo[(1, )](x, x, 4)
        foo[(1, )](x, x, 4)

    proton.deactivate(session, flushing=True)

    database = proton.data.get(session)
    gf, _, _, _ = viewer.get_raw_metrics(database)
    foo_frame = gf.filter("MATCH ('*', c) WHERE c.'name' =~ '.*foo.*' AND c IS LEAF").dataframe
    ones_frame = gf.filter("MATCH ('*', c) WHERE c.'name' =~ '.*elementwise.*' AND c IS LEAF").dataframe

    assert len(foo_frame) == 1
    assert int(foo_frame["count"].values[0]) == 2
    assert len(ones_frame) == 1
    assert int(ones_frame["count"].values[0]) == 1

    import msgpack
    msgpack_data = proton.data.get_msgpack(session)
    database_unpacked = msgpack.loads(msgpack_data, raw=False, strict_map_key=False)
    assert database == database_unpacked

    proton.finalize()


def test_clear_data(tmp_path: pathlib.Path, device: str):
    temp_file = tmp_path / "test_clear_data.hatchet"
    session = proton.start(str(temp_file.with_suffix("")), context="shadow")

    with proton.scope("test0"):
        x = torch.ones((2, 2), device=device)
        x + x  # type: ignore

    proton.deactivate(session, flushing=True)
    proton.data.clear(session)
    try:
        database = proton.data.get(session)
    except RuntimeError as e:
        assert "has no data" in str(e)

    proton.activate(session)
    with proton.scope("test1"):
        x * x  # type: ignore
    proton.deactivate(session, flushing=True)
    database = proton.data.get(session)

    proton.finalize()
    assert len(database[0]["children"]) == 1
    assert database[0]["children"][0]["frame"]["name"] == "test1"
    kernel_frame = database[0]["children"][0]["children"][0]
    assert "elementwise" in kernel_frame["frame"]["name"]


def test_clear_data_up_to_phase(tmp_path: pathlib.Path, device: str):
    temp_file = tmp_path / "test_clear_data_up_to_phase.hatchet"
    session = proton.start(str(temp_file.with_suffix("")), context="shadow")

    with proton.scope("phase0"):
        x = torch.ones((2, 2), device=device)
        x + x  # type: ignore

    phase1 = proton.data.advance_phase(session)
    with proton.scope("phase1"):
        x = torch.ones((2, 2), device=device)
        x + x  # type: ignore

    proton.deactivate(session, flushing=True)

    # Clear a range of phases.
    proton.data.clear(session, phase=phase1, clear_up_to_phase=True)
    database = proton.data.get(session, phase=phase1)
    assert len(database[0]["children"]) == 0

    proton.finalize()


def test_data_is_phase_complete(tmp_path: pathlib.Path, device: str):
    temp_path = tmp_path / "test_data_is_phase_complete.hatchet"
    session = proton.start(str(temp_path.with_suffix("")), context="shadow")

    def fn():
        with proton.scope("test0"):
            x = torch.ones((2, 2), device=device)
            x + x  # type: ignore

    fn()
    assert not proton.data.is_phase_complete(session, 0)

    proton.deactivate(session)
    # likely the GPU has not completed the data yet
    assert not proton.data.is_phase_complete(session, 0)

    proton.activate(session)
    phase = proton.data.advance_phase(session)
    fn()
    proton.deactivate(session, flushing=True)
    # session 0 is a previous phase but we have called deactivate with flushing
    assert proton.data.is_phase_complete(session, 0)
    # phase 1 is the current phase so cannot be a completed phase
    assert not proton.data.is_phase_complete(session, phase)
    proton.data.advance_phase(session)
    # phase 0 should remain completed after advancing phases
    assert proton.data.is_phase_complete(session, phase - 1)

    proton.finalize()


def test_hook_launch(tmp_path: pathlib.Path, device: str):

    def metadata_fn(grid: tuple, metadata: NamedTuple, args: dict):
        # get arg's element size
        element_size = args["x"].element_size()  # non-const
        size = args["size"]  # const
        key = "flops" + str(element_size * 8)
        num_ctas = metadata.num_ctas
        # Return an extra metric key beyond the historical flops/bytes allowlist.
        return {"name": f"foo_test_{num_ctas}ctas_{size}elems", key: 1.0, "extra_metric": 7.0}

    @triton.jit(launch_metadata=metadata_fn)
    def foo(x, size: tl.constexpr, y):
        offs = tl.arange(0, size)
        tl.store(y + offs, tl.load(x + offs))

    x = torch.tensor([2], device=device, dtype=torch.float32)
    y = torch.zeros_like(x)
    temp_file = tmp_path / "test_hook_triton.hatchet"
    proton.start(str(temp_file.with_suffix("")), hook="triton")
    with proton.scope("test0"):
        foo[(1, )](x, 1, y, num_warps=4)
    proton.finalize()
    with temp_file.open() as f:
        data = json.load(f)
    assert len(data[0]["children"]) == 1
    assert data[0]["children"][0]["frame"]["name"] == "test0"
    assert data[0]["children"][0]["children"][0]["frame"]["name"] == "foo_test_1ctas_1elems"
    assert data[0]["children"][0]["children"][0]["metrics"]["flops32"] == 1.0
    assert data[0]["children"][0]["children"][0]["metrics"]["extra_metric"] == 7.0
    assert data[0]["children"][0]["children"][0]["metrics"]["time (ns)"] > 0


def test_hook_launch_filter(tmp_path: pathlib.Path, device: str):

    foo_metadata_invoked = False
    bar_metadata_invoked = False

    def foo_metadata_fn(grid: tuple, metadata: NamedTuple, args: dict):
        nonlocal foo_metadata_invoked
        foo_metadata_invoked = True
        return {"name": "foo_meta", "flops": 1.0}

    def bar_metadata_fn(grid: tuple, metadata: NamedTuple, args: dict):
        nonlocal bar_metadata_invoked
        bar_metadata_invoked = True
        return {"name": "bar_meta", "flops": 2.0}

    @triton.jit(launch_metadata=foo_metadata_fn)
    def foo(x, size: tl.constexpr, y):
        offs = tl.arange(0, size)
        tl.store(y + offs, tl.load(x + offs))

    @triton.jit(launch_metadata=bar_metadata_fn)
    def bar(x, size: tl.constexpr, y):
        offs = tl.arange(0, size)
        tl.store(y + offs, tl.load(x + offs))

    x = torch.tensor([2], device=device, dtype=torch.float32)
    y = torch.zeros_like(x)
    temp_file = tmp_path / "test_hook_triton_filter.hatchet"

    # Only allow kernels whose compiled name matches "foo" (via prefix regex).
    launch_hook = proton_launch.LaunchHook()
    launch_hook.configure(include=".*foo")
    proton.start(str(temp_file.with_suffix("")), hook=launch_hook)
    with proton.scope("test0"):
        foo[(1, )](x, 1, y, num_warps=4)
        bar[(1, )](x, 1, y, num_warps=4)
    proton.finalize()
    # Reset singleton hook state to avoid leaking filter settings across tests.
    launch_hook.configure(include=None, exclude=None)

    assert foo_metadata_invoked is True
    assert bar_metadata_invoked is False

    with temp_file.open() as f:
        data = json.load(f)

    # Ensure the "foo_meta" override exists and "bar_meta" does not.
    all_names = set()
    queue = [data[0]]
    while queue:
        node = queue.pop()
        if "frame" in node and "name" in node["frame"]:
            all_names.add(node["frame"]["name"])
        queue.extend(node.get("children", []))

    assert "foo_meta" in all_names
    assert "bar_meta" not in all_names


@pytest.mark.parametrize("context", ["shadow", "python"])
def test_hook_launch_context(tmp_path: pathlib.Path, context: str, device: str):

    def metadata_fn(grid: tuple, metadata: NamedTuple, args: dict):
        x = args["x"]
        # A gpu kernel, but it should be under the metadata state
        return {"name": "foo_test", "bytes": x.sum().item()}

    @triton.jit(launch_metadata=metadata_fn)
    def foo(x, size: tl.constexpr, y):
        offs = tl.arange(0, size)
        tl.store(y + offs, tl.load(x + offs))

    x = torch.tensor([2], device=device, dtype=torch.float32)
    y = torch.zeros_like(x)
    temp_file = tmp_path / "test_hook.hatchet"
    proton.start(str(temp_file.with_suffix("")), hook="triton", context=context)
    with proton.scope("test0"):
        foo[(1, )](x, 1, y, num_warps=4)
    proton.finalize()
    with temp_file.open() as f:
        data = json.load(f)
    # bfs search until find the reduce kernel and then check its parent
    queue = [data[0]]
    while len(queue) > 0:
        parent_frame = queue.pop(0)
        for child in parent_frame["children"]:
            if "reduce" in child["frame"]["name"]:
                assert parent_frame["frame"]["name"] == COMPUTE_METADATA_SCOPE_NAME
                return
            queue.append(child)


def test_hook_with_third_party(tmp_path: pathlib.Path, device: str):
    third_party_hook_invoked = False

    def third_party_hook(metadata) -> None:
        nonlocal third_party_hook_invoked
        third_party_hook_invoked = True

    triton.knobs.runtime.launch_enter_hook.add(third_party_hook)

    proton_hook_invoked = False

    def metadata_fn(grid: tuple, metadata: NamedTuple, args: dict):
        nonlocal proton_hook_invoked
        proton_hook_invoked = True
        return {"name": "foo_test"}

    @triton.jit(launch_metadata=metadata_fn)
    def foo(x, size: tl.constexpr, y):
        offs = tl.arange(0, size)
        tl.store(y + offs, tl.load(x + offs))

    x = torch.tensor([2], device=device, dtype=torch.float32)
    y = torch.zeros_like(x)
    temp_file = tmp_path / "test_hook_with_third_party.hatchet"
    proton.start(str(temp_file.with_suffix("")), hook="triton")
    foo[(1, )](x, 1, y, num_warps=4)
    proton.finalize()
    triton.knobs.runtime.launch_enter_hook.remove(third_party_hook)
    with temp_file.open() as f:
        data = json.load(f)
    assert len(data[0]["children"]) == 1
    assert data[0]["children"][0]["frame"]["name"] == "foo_test"
    assert data[0]["children"][0]["metrics"]["time (ns)"] > 0


def test_hook_multiple_threads(tmp_path: pathlib.Path, device: str):

    def metadata_fn_foo(grid: tuple, metadata: NamedTuple, args: dict):
        return {"name": "foo_test"}

    @triton.jit(launch_metadata=metadata_fn_foo)
    def foo(x, size: tl.constexpr, y):
        offs = tl.arange(0, size)
        tl.store(y + offs, tl.load(x + offs))

    def metadata_fn_bar(grid: tuple, metadata: NamedTuple, args: dict):
        return {"name": "bar_test"}

    @triton.jit(launch_metadata=metadata_fn_bar)
    def bar(x, size: tl.constexpr, y):
        offs = tl.arange(0, size)
        tl.store(y + offs, tl.load(x + offs))

    x_foo = torch.tensor([2], device=device, dtype=torch.float32)
    y_foo = torch.zeros_like(x_foo)
    x_bar = torch.tensor([2], device=device, dtype=torch.float32)
    y_bar = torch.zeros_like(x_bar)

    temp_file = tmp_path / "test_hook.hatchet"
    proton.start(str(temp_file.with_suffix("")), hook="triton")

    all_ids = set()

    # start multiple threads
    def invoke_foo():
        for _ in range(100):
            foo[(1, )](x_foo, 1, y_foo, num_warps=4)
            all_ids.add(proton_launch.id.get())

    def invoke_bar():
        for _ in range(100):
            bar[(1, )](x_bar, 1, y_bar, num_warps=4)
            all_ids.add(proton_launch.id.get())

    thread_foo = threading.Thread(target=invoke_foo)
    thread_bar = threading.Thread(target=invoke_bar)
    thread_foo.start()
    thread_bar.start()
    thread_foo.join()
    thread_bar.join()

    proton.finalize()
    assert len(all_ids) == 200

    with temp_file.open() as f:
        data = json.load(f)
    root = data[0]["children"]
    assert "foo_test" in root[0]["frame"]["name"] or root[1]["frame"]["name"]
    assert "bar_test" in root[0]["frame"]["name"] or root[1]["frame"]["name"]
    assert root[0]["metrics"]["count"] == 100
    assert root[1]["metrics"]["count"] == 100


def test_pcsampling(tmp_path: pathlib.Path, device: str):
    if not is_cuda():
        pytest.skip("Only CUDA backend supports pc sampling")

    import os

    if os.environ.get("PROTON_SKIP_PC_SAMPLING_TEST", "0") == "1":
        pytest.skip("PC sampling test is disabled")

    @triton.jit
    def foo(x, y, size: tl.constexpr):
        offs = tl.arange(0, size)
        for _ in range(1000):
            tl.store(y + offs, tl.load(x + offs))

    temp_file = tmp_path / "test_pcsampling.hatchet"
    proton.start(str(temp_file.with_suffix("")), hook="triton", backend="cupti", mode="pcsampling")
    with proton.scope("init"):
        x = torch.ones((1024, ), device=device, dtype=torch.float32)
        y = torch.zeros_like(x)
    with proton.scope("test"):
        foo[(1, )](x, y, x.size()[0], num_warps=4)
    proton.finalize()
    with temp_file.open() as f:
        data = json.load(f)
    init_frame = data[0]["children"][0]
    test_frame = data[0]["children"][1]
    # With line mapping
    assert "foo" in test_frame["children"][0]["frame"]["name"]
    assert test_frame["children"][0]["children"][0]["metrics"]["num_samples"] > 0
    assert "@" in test_frame["children"][0]["children"][0]["frame"]["name"]
    # Without line mapping
    assert "elementwise" in init_frame["children"][0]["frame"]["name"]
    assert init_frame["children"][0]["metrics"]["num_samples"] > 0


def test_deactivate(tmp_path: pathlib.Path, device: str):
    temp_file = tmp_path / "test_deactivate.hatchet"
    session_id = proton.start(str(temp_file.with_suffix("")), hook="triton")
    proton.deactivate(session_id)
    torch.randn((10, 10), device=device)
    proton.activate(session_id)
    torch.zeros((10, 10), device=device)
    proton.deactivate(session_id)
    proton.finalize()
    with temp_file.open() as f:
        data = json.load(f)
    # Root shouldn't have device id
    assert "device_id" not in data[0]["metrics"]
    assert len(data[0]["children"]) == 1
    assert "device_id" in data[0]["children"][0]["metrics"]


def test_multiple_sessions(tmp_path: pathlib.Path, device: str):
    temp_file0 = tmp_path / "test_multiple_sessions0.hatchet"
    temp_file1 = tmp_path / "test_multiple_sessions1.hatchet"
    session_id0 = proton.start(str(temp_file0.with_suffix("")))
    session_id1 = proton.start(str(temp_file1.with_suffix("")))
    with proton.scope("scope0"):
        torch.randn((10, 10), device=device)
        torch.randn((10, 10), device=device)
    proton.deactivate(session_id0)
    proton.finalize(session_id0)
    with proton.scope("scope1"):
        torch.randn((10, 10), device=device)
    proton.finalize(session_id1)
    # kernel has been invoked twice in session 0 and three times in session 1
    with temp_file0.open() as f:
        data = json.load(f)
    assert data[0]["children"][0]["frame"]["name"] == "scope0"
    assert int(data[0]["children"][0]["children"][0]["metrics"]["count"]) == 2
    with temp_file1.open() as f:
        data = json.load(f)
    scope0_count = int(data[0]["children"][0]["children"][0]["metrics"]["count"])
    scope1_count = int(data[0]["children"][1]["children"][0]["metrics"]["count"])
    assert scope0_count + scope1_count == 3


@pytest.mark.skipif(not is_cuda(), reason="Only CUDA backend supports metrics profiling in cudagraphs")
def test_multiple_sessions_cudagraph_metric_kernels(tmp_path: pathlib.Path, device: str):
    stream = torch.cuda.Stream()
    torch.cuda.set_stream(stream)

    foo_iters = 3
    bar_iters = 2

    def foo_metadata_fn(grid: tuple, metadata: NamedTuple, args: dict):
        x = args["x"]
        # Tensor custom metric in graph capture mode launches metric kernels.
        return {"name": "foo_with_metric", "sum_metric": x.sum()}

    def bar_metadata_fn(grid: tuple, metadata: NamedTuple, args: dict):
        # Name-only metadata (no custom metric).
        return {"name": "bar_without_metric"}

    @triton.jit(launch_metadata=foo_metadata_fn)
    def foo(x, y, z):
        tl.store(z, tl.load(y) + tl.load(x))

    @triton.jit(launch_metadata=bar_metadata_fn)
    def bar(x, y, z):
        tl.store(z, tl.load(y) - tl.load(x))

    x = torch.ones((2, 2), device=device)
    y = torch.ones((2, 2), device=device)
    z = torch.empty_like(x)

    # Compile kernels before profiling starts to reduce unrelated profile noise.
    foo[(1, )](x, y, z)
    bar[(1, )](x, y, z)

    temp_file0 = tmp_path / "test_multiple_sessions_cudagraph_metric_kernels0.hatchet"
    temp_file1 = tmp_path / "test_multiple_sessions_cudagraph_metric_kernels1.hatchet"
    session_id0 = proton.start(str(temp_file0.with_suffix("")), context="shadow", hook="triton")
    session_id1 = proton.start(str(temp_file1.with_suffix("")), context="shadow", hook="triton")

    proton.deactivate(session_id1)

    graph_foo = torch.cuda.CUDAGraph()
    with torch.cuda.graph(graph_foo):
        for _ in range(foo_iters):
            foo[(1, )](x, y, z)
    with proton.scope("session0_replay"):
        graph_foo.replay()

    proton.deactivate(session_id0)
    proton.activate(session_id1)

    graph_bar = torch.cuda.CUDAGraph()
    with torch.cuda.graph(graph_bar):
        for _ in range(bar_iters):
            bar[(1, )](x, y, z)
    with proton.scope("session1_replay"):
        graph_bar.replay()

    proton.finalize(session_id0)
    proton.finalize(session_id1)

    def get_frame_by_name(node, name: str):
        queue = [node]
        while queue:
            cur = queue.pop(0)
            if cur["frame"]["name"] == name:
                return cur
            queue.extend(cur["children"])
        return None

    def get_all_names(node):
        names = set()
        queue = [node]
        while queue:
            cur = queue.pop(0)
            names.add(cur["frame"]["name"])
            queue.extend(cur["children"])
        return names

    with temp_file0.open() as f:
        data0 = json.load(f)
    with temp_file1.open() as f:
        data1 = json.load(f)

    session0_replay_frame = get_frame_by_name(data0[0], "session0_replay")
    session1_replay_frame = get_frame_by_name(data1[0], "session1_replay")
    assert session0_replay_frame is not None
    assert session1_replay_frame is not None

    capture0 = session0_replay_frame["children"][0]
    capture1 = session1_replay_frame["children"][0]

    foo_frame0 = get_frame_by_name(capture0, "foo_with_metric")
    bar_frame0 = get_frame_by_name(capture0, "bar_without_metric")
    foo_frame1 = get_frame_by_name(capture1, "foo_with_metric")
    bar_frame1 = get_frame_by_name(capture1, "bar_without_metric")

    assert foo_frame0 is not None
    assert bar_frame0 is None
    assert foo_frame1 is None
    assert bar_frame1 is not None

    assert foo_frame0["metrics"]["sum_metric"] == float(foo_iters * x.numel())
    assert int(foo_frame0["metrics"]["count"]) == foo_iters
    assert "sum_metric" not in bar_frame1["metrics"]
    assert int(bar_frame1["metrics"]["count"]) == bar_iters


def test_trace(tmp_path: pathlib.Path, device: str):
    temp_file = tmp_path / "test_trace.chrome_trace"
    proton.start(str(temp_file.with_suffix("")), data="trace")

    @triton.jit
    def foo(x, y, size: tl.constexpr):
        offs = tl.arange(0, size)
        tl.store(y + offs, tl.load(x + offs))

    with proton.scope("init"):
        x = torch.ones((1024, ), device=device, dtype=torch.float32)
        y = torch.zeros_like(x)

    with proton.scope("test"):
        foo[(1, )](x, y, x.size()[0], num_warps=4)

    proton.finalize()

    with temp_file.open() as f:
        data = json.load(f)
        trace_events = data["traceEvents"]
        assert len(trace_events) == 3
        assert trace_events[-1]["name"] == "foo"
        assert trace_events[-1]["args"]["call_stack"] == ["ROOT", "test", "foo"]


def test_scope_multiple_threads(tmp_path: pathlib.Path, device: str):
    temp_file = tmp_path / "test_scope_threads.hatchet"
    proton.start(str(temp_file.with_suffix("")))

    N = 50
    thread_names = ["threadA", "threadB"]

    def worker(prefix: str):
        for i in range(N):
            name = f"{prefix}_{i}"
            proton.enter_scope(name)
            torch.ones((1, ), device=device)
            proton.exit_scope()

    threads = [threading.Thread(target=worker, args=(tname, )) for tname in thread_names]
    for t in threads:
        t.start()
    for t in threads:
        t.join()

    proton.finalize()

    with temp_file.open() as f:
        data = json.load(f)

    children = data[0]["children"]
    assert len(children) == N * len(thread_names)
    names = {c["frame"]["name"] for c in children}
    expected = {f"{t}_{i}" for t in thread_names for i in range(N)}
    assert names == expected


@pytest.mark.skipif(not is_cuda() and not is_hip(), reason="Only CUDA/HIP backend supports NVTX profiling")
@pytest.mark.parametrize("enable_nvtx", [None, True, False])
def test_nvtx_range_push_pop(enable_nvtx, fresh_knobs, tmp_path: pathlib.Path, device: str):
    if enable_nvtx is not None:
        fresh_knobs.proton.enable_nvtx = enable_nvtx
    temp_file = tmp_path / "test_nvtx_range_push_pop.hatchet"
    proton.start(str(temp_file.with_suffix("")))

    with proton.scope("proton_scope"):
        torch.cuda.nvtx.range_push("nvtx_range0")
        torch.cuda.nvtx.range_push("nvtx_range1")
        torch.ones((1, ), device=device)
        torch.cuda.nvtx.range_pop()
        torch.cuda.nvtx.range_pop()

    proton.finalize()

    with temp_file.open() as f:
        data = json.load(f)

    children = data[0]["children"]
    assert len(children) == 1
    proton_scope = children[0]
    assert proton_scope["frame"]["name"] == "proton_scope"
    assert len(proton_scope["children"]) == 1
    if enable_nvtx or enable_nvtx is None:
        nvtx_range0 = proton_scope["children"][0]
        assert nvtx_range0["frame"]["name"] == "nvtx_range0"
        assert len(nvtx_range0["children"]) == 1
        nvtx_range1 = nvtx_range0["children"][0]
        assert nvtx_range1["frame"]["name"] == "nvtx_range1"
        assert len(nvtx_range1["children"]) == 1
        kernel = nvtx_range1["children"][0]
    else:
        kernel = proton_scope["children"][0]
    assert "elementwise" in kernel["frame"]["name"]
    assert kernel["metrics"]["count"] == 1


def test_tensor_metrics_scope(tmp_path: pathlib.Path, device: str):
    temp_file = tmp_path / "test_tensor_metrics_scope.hatchet"
    proton.start(str(temp_file.with_suffix("")))

    x = torch.ones((10, 10), device=device, dtype=torch.float32)
    x_mean = x.mean()
    x_std = x.std()
    with proton.scope("test", metrics={"x_mean": x_mean, "x_std": x_std}):
        torch.randn((10, 10), device=device)
        torch.zeros_like(x)

    proton.finalize()

    with temp_file.open() as f:
        data = json.load(f)

    children = data[0]["children"]
    assert len(children) == 4
    # get the test frame
    test_frame = None
    for child in children:
        if child["frame"]["name"] == "test":
            test_frame = child
            break
    assert test_frame is not None
    assert test_frame["metrics"]["x_mean"] == 1.0
    assert test_frame["metrics"]["x_std"] == 0.0


def test_tensor_metrics_hook(tmp_path: pathlib.Path, device: str):
    temp_file = tmp_path / "test_tensor_metrics_hook.hatchet"

    def metadata_fn(grid: tuple, metadata: NamedTuple, args: dict):
        metric_value = torch.tensor(8.0, device=device)
        return {"name": "foo_test", "flops": metric_value}

    @triton.jit(launch_metadata=metadata_fn)
    def foo(x, size: tl.constexpr, y):
        offs = tl.arange(0, size)
        tl.store(y + offs, tl.load(x + offs))

    x = torch.ones((8, ), device=device, dtype=torch.float32)
    y = torch.zeros_like(x)

    proton.start(str(temp_file.with_suffix("")), hook="triton")
    foo[(1, )](x, x.numel(), y, num_warps=4)
    proton.finalize()

    with temp_file.open() as f:
        data = json.load(f)

    children = data[0]["children"]
    # metadata scope + foo_test
    assert len(children) == 2
    foo_test_frame = None
    for child in children:
        if child["frame"]["name"] == "foo_test":
            foo_test_frame = child
            break
    assert foo_test_frame is not None
    assert foo_test_frame["metrics"]["flops"] == 8.0


@pytest.mark.skipif(not is_cuda(), reason="Only CUDA backend supports metrics profiling in cudagraphs")
def test_tensor_metrics_cudagraph(tmp_path: pathlib.Path, device: str):
    stream = torch.cuda.Stream()
    torch.cuda.set_stream(stream)

    def metadata_fn(grid: tuple, metadata: NamedTuple, args: dict):
        x = args["x"]
        x_sum = x.sum()
        return {"name": "foo_test", "bytes": x.numel() * x.element_size(), "flops": x_sum}

    @triton.jit(launch_metadata=metadata_fn)
    def foo(x, y, z):
        tl.store(z, tl.load(y) + tl.load(x))

    def fn():
        with proton.scope("scope_a", metrics={"bytes": 4 * 4}):
            a = torch.ones((2, 2), device=device)
        with proton.metadata_state():
            a_sum = a.sum()
        with proton.scope("scope_b", metrics={"sum": a_sum}):
            b = torch.ones((2, 2), device=device)
        c = a + b
        foo[(1, )](a, b, c)
        with proton.metadata_state():
            d = torch.arange(4, device="cuda")
        with proton.scope("scope_d", metrics={"vec": d}):
            e = d * 2  # noqa: F841

    temp_file = tmp_path / "test_tensor_metrics_cudagraph.hatchet"
    proton.start(str(temp_file.with_suffix("")), context="shadow", hook="triton")

    # warmup
    # four kernels
    fn()

    # no kernels
    g = torch.cuda.CUDAGraph()
    with torch.cuda.graph(g):
        for _ in range(10):
            fn()

    with proton.scope("test0"):
        g.replay()

    proton.finalize()

    with temp_file.open() as f:
        data = json.load(f)

    children = data[0]["children"]
    # metadata scope + kernels + scope_a + scope_b + test0 + scope_d
    assert len(children) == 8
    test0_frame = None
    for child in children:
        if child["frame"]["name"] == "test0":
            test0_frame = child
            break
    assert test0_frame is not None
    capture_at_frame = test0_frame["children"][0]

    foo_test_frame = None
    scope_a_frame = None
    scope_b_frame = None
    scope_d_frame = None
    for child in capture_at_frame["children"]:
        if child["frame"]["name"] == "foo_test":
            foo_test_frame = child
        if child["frame"]["name"] == "scope_a":
            scope_a_frame = child
        if child["frame"]["name"] == "scope_b":
            scope_b_frame = child
        if child["frame"]["name"] == "scope_d":
            scope_d_frame = child
    assert foo_test_frame is not None
    assert foo_test_frame["metrics"]["bytes"] == 160
    assert foo_test_frame["metrics"]["flops"] == 40
    assert foo_test_frame["metrics"]["count"] == 10
    assert scope_a_frame is not None
    assert scope_a_frame["metrics"]["bytes"] == 160
    assert "count" not in scope_a_frame["metrics"]
    assert scope_b_frame is not None
    assert scope_b_frame["metrics"]["sum"] == 40.0
    assert "count" not in scope_b_frame["metrics"]
    assert scope_d_frame is not None
    assert scope_d_frame["metrics"]["vec"] == [0, 10, 20, 30]


@pytest.mark.skipif(not is_cuda(), reason="Only CUDA backend supports metrics profiling in cudagraphs")
def test_tensor_metrics_cudagraph_deactivate(tmp_path: pathlib.Path, device: str):
    stream = torch.cuda.Stream()
    torch.cuda.set_stream(stream)

    def fn(session):
        proton.deactivate(session)
        with proton.scope("scope_b", metrics={"sum": 4}):
            b = torch.ones((2, 2), device=device)
        proton.activate(session)
        c = b * 2  # noqa: F841

    temp_file = tmp_path / "test_tensor_metrics_cudagraph_deactivate.hatchet"
    session = proton.start(str(temp_file.with_suffix("")), context="shadow", hook="triton")

    # warmup
    fn(session)

    # no kernels
    g = torch.cuda.CUDAGraph()
    with torch.cuda.graph(g):
        for _ in range(10):
            fn(session)

    with proton.scope("test0"):
        g.replay()

    proton.finalize()

    # only a single kernel b * 2
    with temp_file.open() as f:
        data = json.load(f)
        children = data[0]["children"]
        test0_frame = None
        for child in children:
            if child["frame"]["name"] == "test0":
                test0_frame = child
                break
        assert test0_frame is not None
        capture_at_frame = test0_frame["children"][0]
        scope_b_frame = None
        c_frame = None
        for child in capture_at_frame["children"]:
            if child["frame"]["name"] == "scope_b":
                scope_b_frame = child
            if "elementwise" in child["frame"]["name"]:
                c_frame = child
        assert scope_b_frame is None
        assert c_frame is not None
        assert c_frame["metrics"]["count"] == 10


@pytest.mark.skipif(not is_cuda(), reason="Only CUDA backend supports metrics profiling in cudagraphs")
def test_tensor_metrics_multi_device_cudagraph(tmp_path: pathlib.Path):
    if torch.cuda.device_count() < 2:
        pytest.skip("Requires at least two CUDA devices")

    devices = [torch.device(f"cuda:{i}") for i in range(2)]
    streams = []
    for device in devices:
        with torch.cuda.device(device):
            streams.append(torch.cuda.Stream(device=device))

    def metadata_fn(grid: tuple, metadata: NamedTuple, args: dict):
        x = args["x"]
        x_sum = x.sum()
        device_idx = x.device.index
        return {"name": f"foo_test_{device_idx}", "bytes": x.numel() * x.element_size(), "flops": x_sum}

    @triton.jit(launch_metadata=metadata_fn)
    def foo(x, y, z):
        tl.store(z, tl.load(y) + tl.load(x))

    def run_on_device(device_id):
        with proton.scope(f"scope_a_{device_id}", metrics={"bytes": 4 * 4}):
            a = torch.ones((2, 2), device=f"cuda:{device_id}")
        with proton.metadata_state():
            a_sum = a.sum()
        with proton.scope(f"scope_b_{device_id}", metrics={"sum": a_sum}):
            b = torch.ones((2, 2), device=f"cuda:{device_id}")
        c = a + b
        foo[(1, )](a, b, c)

    temp_file = tmp_path / "test_tensor_metrics_multi_device_cudagraph.hatchet"
    proton.start(str(temp_file.with_suffix("")), context="shadow", hook="triton")

    graphs = []
    for device, stream in zip(devices, streams):
        with torch.cuda.device(device):
            torch.cuda.set_stream(stream)
            # warmup
            run_on_device(device.index)
            # graph capture
            g = torch.cuda.CUDAGraph()
            with torch.cuda.graph(g, stream=stream):
                for _ in range(10):
                    run_on_device(device.index)
        graphs.append((device, stream, g))

    for device, stream, graph in graphs:
        with torch.cuda.device(device):
            torch.cuda.set_stream(stream)
            with proton.scope(f"test_device_{device.index}"):
                graph.replay()

    proton.finalize()

    with temp_file.open() as f:
        data = json.load(f)

    children = data[0]["children"]
    for device in devices:
        device_name = f"test_device_{device.index}"
        launch_frame = next((child for child in children if child["frame"]["name"] == device_name), None)
        assert launch_frame is not None
        capture_at_frame = launch_frame["children"][0]
        assert capture_at_frame["frame"]["name"] == "<captured_at>"

        foo_frame = None
        scope_a_frame = None
        scope_b_frame = None
        for child in capture_at_frame["children"]:
            if child["frame"]["name"] == f"foo_test_{device.index}":
                foo_frame = child
            if child["frame"]["name"] == f"scope_a_{device.index}":
                scope_a_frame = child
            if child["frame"]["name"] == f"scope_b_{device.index}":
                scope_b_frame = child

        assert foo_frame is not None
        assert scope_a_frame is not None
        assert scope_b_frame is not None
        assert foo_frame["metrics"]["bytes"] == 160
        assert foo_frame["metrics"]["flops"] == 40
        assert foo_frame["metrics"]["device_id"] == str(device.index)
        assert scope_a_frame["metrics"]["bytes"] == 160
        assert scope_b_frame["metrics"]["sum"] == 40.0

    assert len(data) > 1
    cuda_devices = data[1].get("CUDA", {})
    assert len(cuda_devices) >= 2


@pytest.mark.parametrize("buffer_size", [256 * 1024, 64 * 1024 * 1024])
@pytest.mark.parametrize("data_format", ["hatchet_msgpack", "hatchet"])
def test_periodic_flushing(tmp_path, fresh_knobs, data_format, buffer_size, device: str):
    fresh_knobs.proton.profile_buffer_size = buffer_size
    temp_file = tmp_path / f"test_periodic_flushing.{data_format}"
    session = proton.start(str(temp_file.with_suffix("")), mode=f"periodic_flushing:format={data_format}")

    for i in range(10000):
        if i != 0 and i % 1000 == 0:
            proton.data.advance_phase(session=session)
        with proton.scope(f"test_{i}", metrics={"count": 1}):
            torch.zeros((100), device=device)

    proton.finalize(output_format=data_format)

    # Find all *.hatchet files under the directory `tmp_path`
    import glob
    import msgpack
    hatchet_files = glob.glob(str(tmp_path / f"*.{data_format}"))
    assert len(hatchet_files) == 10
    num_scopes = 0
    for hatchet_file in hatchet_files:
        if data_format == "hatchet_msgpack":
            with open(hatchet_file, "rb") as f:
                data = msgpack.load(f, raw=False, strict_map_key=False)
        else:
            with open(hatchet_file, "r", encoding="utf-8") as f:
                data = json.load(f)
        assert len(data[0]["children"]) == 1000
        assert data[0]["children"][0]["metrics"]["count"] == 1
        assert data[0]["children"][0]["frame"]["name"].startswith("test_")
        assert data[0]["children"][0]["children"][0]["metrics"]["time (ns)"] > 0
        num_scopes += len(data[0]["children"])
    assert num_scopes == 10000


@pytest.mark.skipif(not is_cuda(), reason="Only CUDA backend supports metrics profiling in cudagraphs")
@pytest.mark.parametrize("buffer_size", [256 * 1024, 64 * 1024 * 1024])
@pytest.mark.parametrize("data_format", ["hatchet_msgpack", "hatchet"])
def test_periodic_flushing_cudagraph(tmp_path, fresh_knobs, data_format, buffer_size, device: str):
    fresh_knobs.proton.profile_buffer_size = buffer_size
    temp_file = tmp_path / f"test_periodic_flushing.{data_format}"
    session = proton.start(str(temp_file.with_suffix("")), mode=f"periodic_flushing:format={data_format}",
                           hook="triton")

    def metadata_fn(grid: tuple, metadata: NamedTuple, args: dict):
        x = args["x"]
        x_sum = x.sum()
        return {"name": "foo_test", "bytes": x.numel() * x.element_size(), "flops": x_sum}

    @triton.jit(launch_metadata=metadata_fn)
    def foo(x, y, z):
        tl.store(z, tl.load(y) + tl.load(x))

    def fn():
        with proton.scope("scope_a", metrics={"bytes": 4 * 4}):
            a = torch.ones((2, 2), device=device)
        c = a + a
        foo[(1, )](a, a, c)

    # warmup
    fn()

    # no kernels
    g = torch.cuda.CUDAGraph()
    with torch.cuda.graph(g):
        fn()

    with proton.scope("test0"):
        for i in range(10000):
            if i != 0 and i % 1000 == 0:
                proton.data.advance_phase(session=session)
            g.replay()

    proton.finalize(output_format=data_format)

    # Find all *.hatchet files under the directory `tmp_path`
    import glob
    import msgpack
    hatchet_files = glob.glob(str(tmp_path / f"*.{data_format}"))
    assert len(hatchet_files) == 10
    for hatchet_file in hatchet_files:
        if data_format == "hatchet_msgpack":
            with open(hatchet_file, "rb") as f:
                data = msgpack.load(f, raw=False, strict_map_key=False)
        else:
            with open(hatchet_file, "r", encoding="utf-8") as f:
                data = json.load(f)
        capture_frame = None
        for child in data[0]["children"]:
            if child["frame"]["name"] == "test0":
                capture_frame = child["children"][0]
                break
        assert capture_frame is not None
        scope_a_frame = None
        foo_test_frame = None
        for child in capture_frame["children"]:
            if child["frame"]["name"] == "scope_a":
                scope_a_frame = child
            if child["frame"]["name"] == "foo_test":
                foo_test_frame = child
        assert scope_a_frame is not None
        assert foo_test_frame is not None
        assert scope_a_frame["metrics"]["bytes"] == 16000
        assert foo_test_frame["metrics"]["bytes"] == 16000
        assert foo_test_frame["metrics"]["flops"] == 4000


@pytest.mark.skipif(not is_blackwell(), reason="HW trace is only supported on Blackwell GPUs")
def test_hw_trace(fresh_knobs, tmp_path: pathlib.Path, device: str):
    fresh_knobs.proton.enable_hw_trace = True
    temp_file = tmp_path / "test_hw_trace.hatchet"
    proton.start(str(temp_file.with_suffix("")), hook="triton")

    with proton.scope("init"):
        x = torch.ones((1024, ), device=device, dtype=torch.float32)  # noqa: F841

    proton.finalize()

    with temp_file.open() as f:
        data = json.load(f)
    kernel_frame = data[0]["children"][0]["children"][0]
    assert "elementwise" in kernel_frame["frame"]["name"]
    assert kernel_frame["metrics"]["time (ns)"] > 0
