# Licensed to the Apache Software Foundation (ASF) under one
# or more contributor license agreements.  See the NOTICE file
# distributed with this work for additional information
# regarding copyright ownership.  The ASF licenses this file
# to you under the Apache License, Version 2.0 (the
# "License"); you may not use this file except in compliance
# with the License.  You may obtain a copy of the License at
#
#   http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing,
# software distributed under the License is distributed on an
# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
# KIND, either express or implied.  See the License for the
# specific language governing permissions and limitations
# under the License.

import argparse
import os
import re
import subprocess
import sys
import time
from collections.abc import Mapping
from enum import Enum

import numpy as np
import torch
import triton.profiler as proton
import tvm_ffi

import tvm
from tvm.script import tirx as T
from tvm.support import nvcc


def is_running_under_pytest():
    """Check if the code is being executed within a pytest session."""
    return "PYTEST_CURRENT_TEST" in os.environ


def setup():
    parser = argparse.ArgumentParser()
    parser.add_argument("--dump-ptx", type=str, help="Dump PTX code to specified file")
    parser.add_argument("--dump-source", action="store_true", help="Dump source code")
    args = parser.parse_args()

    if args.dump_ptx:

        @tvm_ffi.register_global_func("tvm_callback_cuda_compile", override=True)
        def tvm_callback_cuda_compile(code, target):
            ptx = nvcc.compile_cuda(code, target_format="ptx")
            with open(args.dump_ptx, "w", encoding="utf-8") as f:
                f.write(ptx.decode())
            return ptx

    return args


_ANSI_RE = re.compile(r"\x1b\[[0-9;]*m")


# proton-viewer -m avg_time/us prints average kernel time in microseconds (see
# triton/profiler/viewer.py avg_time_factor_dict). Store microseconds as-is.
PROTON_AVG_TIME_METRIC = "avg_time/us"


def _parse_proton_tree(text, *, kernel: str = ""):
    """Parse proton-viewer tree output into {impl: time_us}.

    Accepts ALL depth-1 nodes (no KNOWN_IMPLS filter). For each depth-1 impl,
    takes the slowest depth-2 child kernel time.

    Tree numbers are microseconds when ProtonContext uses avg_time/us.

    Returns (impl_times, baseline_errors) where:
      impl_times: {str: float} — impl name to avg time in microseconds
      baseline_errors: {str: str} — impl name to error message
    """
    _ = kernel  # kept for callers; unit does not depend on workload
    impl = None
    results = {}
    baseline_errors = {}
    for raw in text.splitlines():
        line = _ANSI_RE.sub("", raw).rstrip()
        if not line:
            continue
        if line.startswith("BASELINE_ERROR:"):
            parts = line.split(":", 2)
            if len(parts) >= 3:
                baseline_errors[parts[1].strip()] = parts[2].strip()
            continue
        # Depth-1 impl header: starts with tree drawing chars
        if line and line[0] in "\u251c\u2514":  # ├ └
            parts = line.split("\u2500", 1)[-1].split()  # split on ─
            if len(parts) >= 2:
                impl = parts[1]
            else:
                impl = None
            continue
        # Depth-2 kernel: contains tree drawing chars at deeper indent
        if impl and ("\u251c\u2500" in line or "\u2514\u2500" in line):  # ├─ └─
            parts = line.split("\u2500", 1)[-1].split()
            if len(parts) >= 2:
                name = parts[1]
                if (
                    "vectorized_elementwise_kernel" in name
                    or "elementwise_kernel_with_index" in name
                ):
                    continue
                try:
                    t = float(parts[0])
                    results[impl] = max(results.get(impl, 0), t)
                except ValueError:
                    pass
    return results, baseline_errors


class ProtonContext:
    """Context manager for Proton profiling sessions.

    Always captures proton-viewer output and parses impl times so that
    get_impl_times() / get_baseline_errors() work after exiting the context.

    The proton tree is printed to **stdout** by default (visible on screen
    when running kernels interactively).  When the environment variable
    ``TIRX_BENCH_JSON=1`` is set (done automatically by ``--json`` mode),
    the tree goes to **stderr** instead so it does not corrupt the JSON on
    stdout.
    """

    def __init__(
        self,
        name="kernel",
        hook="triton",
        debug=False,
        nsight=False,
        metric=PROTON_AVG_TIME_METRIC,
        kernel="",
    ):
        self.name = name
        self.hook = hook
        self.debug = debug
        self.nsight = nsight
        self.metric = metric
        self.kernel = kernel
        self._impl_times = {}
        self._baseline_errors = {}

    def __enter__(self):
        if not is_running_under_pytest() and not self.debug and not self.nsight:
            proton.start(self.name, hook=self.hook)
            proton.deactivate()
        return self

    def __exit__(self, exc_type, exc_val, exc_tb):
        if not is_running_under_pytest() and not self.debug and not self.nsight:
            proton.finalize()

            hatchet = f"{self.name}.hatchet"
            result = subprocess.run(
                ["proton-viewer", "-m", self.metric, hatchet],
                capture_output=True,
                text=True,
                check=False,
            )
            if result.returncode == 0:
                self._impl_times, self._baseline_errors = _parse_proton_tree(
                    result.stdout, kernel=self.kernel
                )
                out = sys.stderr if os.environ.get("TIRX_BENCH_JSON") else sys.stdout
                print(f"# proton {PROTON_AVG_TIME_METRIC} (microseconds)\n", file=out, end="")
                print(result.stdout, file=out, end="")
            else:
                print(
                    f"proton-viewer failed (rc={result.returncode}): {result.stderr}",
                    file=sys.stderr,
                )

            if os.path.exists(hatchet):
                os.remove(hatchet)

    def get_impl_times(self):
        """Return {impl_name: avg_time_us} parsed from proton-viewer output."""
        return dict(self._impl_times)

    def get_baseline_errors(self):
        """Return {impl_name: error_message} from BASELINE_ERROR lines."""
        return dict(self._baseline_errors)


def _get_l2_cache_bytes():
    """Query L2 cache size from the current CUDA device, fallback to 128MB."""
    try:
        props = torch.cuda.get_device_properties(torch.cuda.current_device())
        if hasattr(props, "l2_cache_size") and props.l2_cache_size > 0:
            return props.l2_cache_size
    except Exception:
        pass
    return 128 * 1024 * 1024  # 128MB default (B200)


def _tensor_bytes(args, _seen=None):
    """Sum the byte size of all torch/tvm tensors in a nested value."""
    if _seen is None:
        _seen = set()
    total = 0
    if isinstance(args, list | tuple):
        for a in args:
            total += _tensor_bytes(a, _seen)
    elif isinstance(args, Mapping):
        for a in args.values():
            total += _tensor_bytes(a, _seen)
    elif isinstance(args, torch.Tensor):
        key = ("torch", args.device.type, args.device.index, int(args.data_ptr()))
        if key not in _seen:
            _seen.add(key)
            total += args.nelement() * args.element_size()
    elif hasattr(args, "numpy"):  # tvm.runtime.NDArray
        try:
            key = ("tvm", int(args.handle.value))
        except Exception:
            key = ("tvm", id(args))
        if key not in _seen:
            _seen.add(key)
            try:
                total += int(np.prod(args.shape)) * np.dtype(str(args.dtype)).itemsize
            except Exception:
                total += args.numpy().nbytes
    return total


def tensor_bytes(*values):
    """Return unique torch/tvm tensor bytes for kernel-owned byte accounting.

    The benchmark driver does not use this implicitly.  Kernel benchmark
    factories may call it when their invocation footprint is exactly the set of
    tensors in ``values``.
    """
    if len(values) == 1:
        return _tensor_bytes(values[0])
    return _tensor_bytes(values)


def _compute_group_count(input_bytes, l2_bytes=None):
    """Return TK-style input-group count from one invocation's byte footprint."""
    if input_bytes <= 0:
        return 1
    if l2_bytes is None:
        l2_bytes = _get_l2_cache_bytes()
    threshold = l2_bytes * 3
    if input_bytes >= threshold:
        return 1
    return int(threshold // input_bytes) + 1


def _make_bench_input(input_factory):
    value = input_factory()
    if not isinstance(value, tuple) or len(value) != 2:
        raise TypeError("input_factory must return (case, input_bytes)")

    case, input_bytes = value
    try:
        input_bytes = int(input_bytes)
    except (TypeError, ValueError) as err:
        raise TypeError("input_factory input_bytes must be an integer") from err
    if input_bytes < 0:
        raise ValueError("input_factory input_bytes must be non-negative")
    return case, input_bytes


def prepare_input_groups(input_factory, l2_bytes=None):
    """Materialize TK-style input groups from a single-group factory.

    ``input_factory`` must return ``(case, input_bytes)``.  ``case`` is passed
    back to every benchmark function unchanged.  ``input_bytes`` defines one
    invocation's L2-eviction footprint and is intentionally owned by the kernel
    benchmark harness instead of inferred here.
    """
    if not callable(input_factory):
        raise TypeError("input_factory must be callable")
    if l2_bytes is None:
        l2_bytes = _get_l2_cache_bytes()

    sample, input_bytes = _make_bench_input(input_factory)
    num_groups = _compute_group_count(input_bytes, l2_bytes)
    groups = [sample]
    for _ in range(num_groups - 1):
        case, _ = _make_bench_input(input_factory)
        groups.append(case)

    return groups, {
        "num_groups": num_groups,
        "input_bytes": input_bytes,
        "l2_bytes": l2_bytes,
        "l2_eviction_factor": 3,
        "flush_l2": False,
    }


def _bench_event_groups(funcs, groups, warmup, repeat, cooldown_s):
    num_groups = len(groups)
    results = {}

    for idx, (name, func) in enumerate(funcs.items()):
        if idx > 0:
            time.sleep(cooldown_s)

        for i in range(warmup):
            func(groups[i % num_groups])

        start_event = torch.cuda.Event(enable_timing=True)
        end_event = torch.cuda.Event(enable_timing=True)
        torch.cuda.synchronize()

        start_event.record()
        for i in range(repeat):
            func(groups[i % num_groups])
        end_event.record()

        torch.cuda.synchronize()
        results[name] = start_event.elapsed_time(end_event) / repeat * 1000.0

        time.sleep(cooldown_s)

    return results


def _bench_proton_groups(
    funcs, groups, warmup, repeat, cooldown_s, proton_name, debug, nsight, *, kernel=""
):
    num_groups = len(groups)
    with ProtonContext(proton_name, debug=debug, nsight=nsight, kernel=kernel) as ctx:
        for idx, (name, func) in enumerate(funcs.items()):
            if idx > 0:
                time.sleep(cooldown_s)

            for i in range(warmup):
                func(groups[i % num_groups])
            torch.cuda.synchronize()

            if not is_running_under_pytest() and not debug and not nsight:
                proton.activate()
                with proton.scope(name, metrics={}):
                    for i in range(repeat):
                        func(groups[i % num_groups])
                proton.deactivate()
            else:
                for i in range(repeat):
                    func(groups[i % num_groups])
            torch.cuda.synchronize()

            time.sleep(cooldown_s)

    return ctx.get_impl_times(), ctx.get_baseline_errors()


def _flush_l2_legacy(flush_l2_size):
    if flush_l2_size > 0:
        torch.empty(flush_l2_size, dtype=torch.int, device="cuda").zero_()


def _bench_legacy_callable(func, warmup, repeat, proton_name, debug, nsight, flush_l2_size):
    start_event = torch.cuda.Event(enable_timing=True)
    end_event = torch.cuda.Event(enable_timing=True)

    def timed_loop():
        start_event.record()
        for _ in range(repeat):
            _flush_l2_legacy(flush_l2_size)
            func()
        end_event.record()

    for _ in range(warmup):
        _flush_l2_legacy(flush_l2_size)
        func()
    torch.cuda.synchronize()
    if not is_running_under_pytest() and not debug and not nsight:
        proton.activate()
        with proton.scope(proton_name, metrics={}):
            timed_loop()
        proton.deactivate()
    else:
        timed_loop()
    torch.cuda.synchronize()

    return start_event.elapsed_time(end_event) / repeat * 1000.0


# Labels identifying our own kernel (vs external reference impls). Must match
# OUR_IMPLS in bench_suite's ratio_diff.py. Used by the TIRX_BENCH_IMPLS filter.
OURS_IMPLS = frozenset({"tir", "tirx"})


def bench_impls_mode():
    """Current impl-selection mode: 'all' (default), 'ours', or 'baseline'.

    Set via the ``TIRX_BENCH_IMPLS`` env var (by bench_suite ``run.py --impls``).
    A kernel's ``run_bench`` can use this to skip *building/warming* reference
    impls (e.g. flashinfer autotune, deepgemm/cublas ext setup) that ``bench``'s
    filter alone cannot avoid, since that setup runs before ``bench`` is called.
    """
    mode = os.environ.get("TIRX_BENCH_IMPLS", "all").lower()
    if mode not in {"all", "ours", "baseline"}:
        raise ValueError(f"TIRX_BENCH_IMPLS must be 'all', 'ours', or 'baseline', got {mode!r}")
    return mode


def bench_only_ours():
    """True when only our own kernel should be benched (reference setup skippable)."""
    return bench_impls_mode() == "ours"


def filter_impls(funcs):
    """Filter a ``{label: callable}`` impl map per the current ``bench_impls_mode``.

    Call this right after building ``funcs`` so any subsequent
    ``if "<ref>" in funcs:`` reference-setup blocks are skipped in 'ours' mode.
    """
    mode = bench_impls_mode()
    if mode == "ours":
        return {n: f for n, f in funcs.items() if n in OURS_IMPLS}
    if mode == "baseline":
        return {n: f for n, f in funcs.items() if n not in OURS_IMPLS}
    return funcs


def bench(
    funcs,
    input_factory=None,
    warmup=500,
    repeat=100,
    cooldown_s=1.0,
    timer="proton",
    proton_name="kernel",
    l2_bytes=None,
    debug=False,
    nsight=False,
    flush_l2_size=int(8e8 // 4),
    references=None,
    rounds=1,
    round_cooldown_s=1.0,
    validate_case=None,
):
    """Benchmark implementations with a factory-owned input footprint.

    This is the single TIRx benchmark API.  It follows the ThunderKittens-style
    multi-input protocol for L2 eviction and supports either Proton/CUPTI or
    CUDA-event timing.  The benchmark driver never infers which tensors belong
    to a workload; ``input_factory`` owns that definition by returning
    ``(case, input_bytes)``.

    Parameters
    ----------
    funcs : dict[str, callable]
        Map of implementation name to callable.  Each callable receives one
        ``case`` returned by ``input_factory``.  This should hold only *our*
        kernel(s); external baselines go in ``references``.
    references : dict[str, callable], optional
        Map of reference-impl name to a no-arg *builder* that does the heavy
        import/setup and returns the run callable.  Builders run lazily and only
        when that impl will actually be benched (skipped entirely under
        ``--impls ours``); a builder that raises is recorded as a
        ``BASELINE_ERROR`` instead of failing the workload.
    input_factory : callable
        Factory returning ``(case, input_bytes)`` for one benchmark group.
    warmup : int
        Number of untimed warmup iterations per implementation.
    repeat : int
        Number of timed iterations per round.
    cooldown_s : float
        Seconds to sleep between impls for thermal cooldown.
    rounds : int
        Independent benchmark rounds (compile + inputs once; each round runs
        warmup + repeat for every selected impl).
    round_cooldown_s : float
        Seconds to sleep between rounds (ignored when ``rounds == 1``).
    validate_case : callable, optional
        Called once on the first prepared ``case`` (after ``prepare_input_groups``,
        before warmup/repeat rounds). Under bench_suite, ``run_kernel_bench`` holds
        the per-GPU lock for the whole ``run_bench()`` call.
    timer : {"event", "proton"}
        Timing backend.

    Returns
    -------
    dict
        ``{"impls": {name: us}, "round_samples": {name: [us, ...]}, ...}``.
        Times are stored in microseconds (same unit as pinned bench_suite baselines).
    """
    if repeat <= 0:
        raise ValueError("repeat must be positive")
    if warmup < 0:
        raise ValueError("warmup must be non-negative")
    if rounds < 1:
        raise ValueError("rounds must be >= 1")
    if round_cooldown_s < 0:
        raise ValueError("round_cooldown_s must be non-negative")
    if timer not in {"event", "proton"}:
        raise ValueError(f"unsupported timer {timer!r}; expected event or proton")

    if callable(funcs) and input_factory is None:
        return _bench_legacy_callable(
            funcs,
            warmup=warmup,
            repeat=repeat,
            proton_name=proton_name,
            debug=debug,
            nsight=nsight,
            flush_l2_size=flush_l2_size,
        )

    if input_factory is None:
        raise TypeError("input_factory is required when funcs is a mapping")
    if not isinstance(funcs, Mapping) or not funcs:
        raise TypeError("funcs must be a non-empty mapping of name to callable")
    for name, func in funcs.items():
        if not isinstance(name, str):
            raise TypeError("func names must be strings")
        if not callable(func):
            raise TypeError(f"funcs[{name!r}] must be callable")

    # Select impls for this run. ``funcs`` holds our own kernel(s); external
    # baselines are passed as ``references`` (name -> no-arg builder). A builder
    # is invoked here ONLY when its impl will actually be benched, so --impls
    # ours skips all reference setup, --impls baseline skips ours, and a builder
    # that fails is recorded as a BASELINE_ERROR rather than failing the
    # workload. Legacy kernels that put references directly in ``funcs`` are
    # still handled by the label filter (filter_impls) below.
    funcs = filter_impls(funcs)
    build_errors: dict[str, str] = {}
    if bench_impls_mode() != "ours":
        for ref_name, builder in (references or {}).items():
            if not isinstance(ref_name, str) or not callable(builder):
                raise TypeError("references must map a name to a no-arg builder callable")
            try:
                ref_fn = builder()
            except Exception as e:
                build_errors[ref_name] = str(e)
                print(f"BASELINE_ERROR: {ref_name}: {e}", file=sys.stderr)
                continue
            if ref_fn is None:
                continue
            if not callable(ref_fn):
                raise TypeError(f"references[{ref_name!r}] builder must return a callable")
            funcs = {**funcs, ref_name: ref_fn}

    if not funcs:
        # Nothing to bench in this mode (e.g. 'ours' on a kernel with no tir
        # impl, or 'baseline' with no references). Return an empty-but-valid
        # result so the workload is a no-op.
        return {
            "impls": {},
            "round_samples": {},
            "errors": build_errors,
            "timer": timer,
            "benchmark_protocol": {
                "warmup": warmup,
                "repeat": repeat,
                "cooldown_s": cooldown_s,
                "rounds": rounds,
                "round_cooldown_s": round_cooldown_s,
                "order": [],
            },
        }

    inputs, protocol = prepare_input_groups(input_factory, l2_bytes=l2_bytes)
    num_groups = len(inputs)
    if num_groups == 0:
        return {
            "impls": {},
            "round_samples": {},
            "errors": build_errors,
            "timer": timer,
            "benchmark_protocol": {
                **protocol,
                "warmup": warmup,
                "repeat": repeat,
                "cooldown_s": cooldown_s,
                "rounds": rounds,
                "round_cooldown_s": round_cooldown_s,
                "order": list(funcs.keys()),
            },
        }

    if validate_case is not None:
        validate_case(inputs[0])

    errors = dict(build_errors)
    round_samples: dict[str, list[float]] = {}
    for round_idx in range(rounds):
        if round_idx > 0:
            time.sleep(round_cooldown_s)
        if timer == "event":
            impls = _bench_event_groups(funcs, inputs, warmup, repeat, cooldown_s)
            proton_errors = {}
        else:
            impls, proton_errors = _bench_proton_groups(
                funcs,
                inputs,
                warmup,
                repeat,
                cooldown_s,
                proton_name,
                debug,
                nsight,
                kernel=proton_name,
            )
        errors.update(proton_errors)
        for impl, sec in impls.items():
            round_samples.setdefault(impl, []).append(sec)

    if not round_samples:
        aggregated = {}
    else:
        import statistics

        aggregated = {impl: statistics.mean(samples) for impl, samples in round_samples.items()}

    return {
        "impls": aggregated,
        "round_samples": round_samples,
        "errors": errors,
        "timer": timer,
        "benchmark_protocol": {
            **protocol,
            "warmup": warmup,
            "repeat": repeat,
            "cooldown_s": cooldown_s,
            "rounds": rounds,
            "round_cooldown_s": round_cooldown_s,
            "order": list(funcs.keys()),
        },
    }


# utils for tg4perfetto profiler, adapted from https://github.com/flashinfer-ai/flashinfer


class EventType(Enum):
    kBegin = 0
    kEnd = 1
    kInstant = 2
    kFinalize = 3


def decode_tag(tag, num_groups):
    block_group_tag = tag >> 12
    event_idx = (tag >> 2) & 0x3FF
    event_type = tag & 0x3
    return (block_group_tag // num_groups, block_group_tag % num_groups, event_idx, event_type)


def export_to_perfetto_trace(
    profiler_buffer: np.ndarray, file_name: str, event_type_names: list[str]
) -> None:
    if is_running_under_pytest():
        return

    import torch

    # pip install git+https://github.com/ihavnoid/tg4perfetto.git
    from tg4perfetto import TraceGenerator

    profiler_buffer_host = torch.tensor(profiler_buffer)
    num_blocks, num_groups = profiler_buffer_host[:1].view(dtype=torch.int32)
    num_blocks = int(num_blocks)
    num_groups = int(num_groups)
    tgen = TraceGenerator(file_name)

    tid_map = {}
    track_map = {}
    finish_idx = set()
    for block_idx in range(num_blocks):
        pid = tgen.create_group(f"block_{block_idx}")
        for group_idx in range(num_groups):
            tid = pid.create_group(f"group_{group_idx}")
            tid_map[(block_idx, group_idx)] = tid

    for i in range(1, len(profiler_buffer_host)):
        if profiler_buffer_host[i] == 0:
            continue
        tag, timestamp = profiler_buffer_host[i : i + 1].view(dtype=torch.uint32)
        tag = int(tag)
        timestamp = int(timestamp)
        block_idx, group_idx, event_idx, event_type = decode_tag(tag, num_groups)

        if event_type == EventType.kFinalize.value:
            finish_idx.add((block_idx, group_idx))
            if len(finish_idx) == num_blocks * num_groups:
                break
        else:
            if (block_idx, group_idx) in finish_idx:
                continue

        event = event_type_names[event_idx]
        tid = tid_map[(block_idx, group_idx)]

        if (block_idx, group_idx, event_idx) in track_map:
            track = track_map[(block_idx, group_idx, event_idx)]
        else:
            track = tid.create_track()
            track_map[(block_idx, group_idx, event_idx)] = track

        if event_type == EventType.kBegin.value:
            track.open(timestamp, event)
        elif event_type == EventType.kEnd.value:
            track.close(timestamp)
        elif event_type == EventType.kInstant.value:
            track.instant(timestamp, event)

    tgen.flush()


@T.meta_class
class CudaProfiler:
    """A lightweight wrapper around T.timer_* CUDA intrinsics.

    Stores repeated arguments used by timer_init/start/end/finalize so users can
    call concise methods in kernels. Intended to mirror Pipeline/TileScheduler helpers.

    When ``profiler_enabled`` is False (or a false-y Expr), calls to
    ``init/start/end/finalize`` become no-ops. This allows constructing a
    profiler unconditionally and eliminating external ``if PROFILER_ON:`` guards.
    """

    def __init__(
        self,
        profiler_buffer: T.Buffer,
        write_stride: int,
        num_groups: int,
        default_leader: None | tvm.tirx.Expr | bool = None,
        profiler_enabled: bool | tvm.tirx.Expr = True,
    ):
        self.buffer = profiler_buffer
        self.write_stride = write_stride
        self.num_groups = num_groups
        self.default_leader = default_leader
        # Accept either a Python bool or a Expr; normalize simple bools to T.bool
        # so we can use it uniformly inside macros for conditional emission.
        if isinstance(profiler_enabled, bool | np.bool_):
            self.profiler_enabled = T.bool(bool(profiler_enabled))
        else:
            # Assume Expr-like input; use as-is
            self.profiler_enabled = profiler_enabled  # type: ignore[assignment]

        self.profiler_tag = T.alloc_buffer([1], "uint64", scope="local", align=8)
        self.profiler_write_offset = T.alloc_buffer([1], "uint32", scope="local", align=8)

    def _leader(self, leader: None | tvm.tirx.Expr | bool):
        if leader is not None:
            if isinstance(leader, bool | np.bool_):
                return T.bool(bool(leader))
            return leader
        if self.default_leader is not None:
            return self.default_leader
        return T.bool(True)

    @T.inline
    def init(self, group_id: tvm.tirx.Expr):
        if self.profiler_enabled:
            T.cuda.timer_init(
                self.buffer.data,
                self.profiler_tag.data,
                self.profiler_write_offset.data,
                self.num_groups,
                group_id,
            )

    @T.inline
    def start(self, event_type: Enum, leader: None | tvm.tirx.Expr | bool = None):
        if self.profiler_enabled:
            T.cuda.timer_start(
                event_type,
                self.buffer.data,
                self.profiler_tag.data,
                self.profiler_write_offset.data,
                self.write_stride,
                self._leader(leader),
            )

    @T.inline
    def end(self, event_type: Enum, leader: None | tvm.tirx.Expr | bool = None):
        if self.profiler_enabled:
            T.cuda.timer_end(
                event_type,
                self.buffer.data,
                self.profiler_tag.data,
                self.profiler_write_offset.data,
                self.write_stride,
                self._leader(leader),
            )

    @T.inline
    def finalize(self, leader: None | tvm.tirx.Expr | bool = None):
        if self.profiler_enabled:
            T.cuda.timer_finalize(
                self.buffer.data,
                self.profiler_tag.data,
                self.profiler_write_offset.data,
                self.write_stride,
                self._leader(leader),
            )