TensorRT-LLMs/examples/layer_wise_benchmarks/parse.py

import argparse
import bisect
import csv
import json
import re
import sqlite3
import subprocess
from pathlib import Path

import jinja2
import numpy as np
import pandas as pd

# Parse cmdline
parser = argparse.ArgumentParser()
parser.add_argument("--profile-dir", type=str, default="profiles")
parser.add_argument("--world-size", "--np", type=int)
parser.add_argument("--rank", type=int, default=0)
parser.add_argument("--warmup-times", type=int)
group = parser.add_mutually_exclusive_group(required=False)
group.add_argument("--error-on-unknown-kernel", action="store_true", dest="error_on_unknown_kernel")
group.add_argument(
    "--no-error-on-unknown-kernel", action="store_false", dest="error_on_unknown_kernel"
)
parser.set_defaults(error_on_unknown_kernel=None)
args = parser.parse_args()
print(args)


def lazy_convert_sqlite(nsys_rep_file_path, sqlite_file_path):
    if (
        not sqlite_file_path.is_file()
        or nsys_rep_file_path.stat().st_mtime > sqlite_file_path.stat().st_mtime
    ):
        subprocess.check_call(
            [
                "nsys",
                "export",
                "--type",
                "sqlite",
                "-o",
                sqlite_file_path,
                "--force-overwrite=true",
                nsys_rep_file_path,
            ]
        )


def shortest_common_supersequence(a, b):
    # Merge two lists into their shortest common supersequence,
    # so that both `a` and `b` are subsequences of the result.
    # Uses dynamic programming to compute the shortest common supersequence, then reconstructs it.
    m, n = len(a), len(b)
    dp = [[0] * (n + 1) for _ in range(m + 1)]
    for i in range(m + 1):
        dp[i][0] = i
    for j in range(n + 1):
        dp[0][j] = j
    for i in range(1, m + 1):
        for j in range(1, n + 1):
            if a[i - 1] == b[j - 1]:
                dp[i][j] = dp[i - 1][j - 1] + 1
            else:
                dp[i][j] = min(dp[i - 1][j] + 1, dp[i][j - 1] + 1)
    # Backtrack to build the merged sequence
    res = []
    i, j = m, n
    while i > 0 and j > 0:
        if a[i - 1] == b[j - 1]:
            res.append(a[i - 1])
            i -= 1
            j -= 1
        elif dp[i - 1][j] < dp[i][j - 1]:
            res.append(a[i - 1])
            i -= 1
        else:
            res.append(b[j - 1])
            j -= 1
    while i > 0:
        res.append(a[i - 1])
        i -= 1
    while j > 0:
        res.append(b[j - 1])
        j -= 1
    res.reverse()
    return res


profile_dir = Path(args.profile_dir)
nsys_rep_file_path = profile_dir / f"report_np{args.world_size}_rank{args.rank}.nsys-rep"
sqlite_file_path = profile_dir / f"report_np{args.world_size}_rank{args.rank}.sqlite"
csv_file_path = profile_dir / f"report_np{args.world_size}_rank{args.rank}.csv"
html_file_path = profile_dir / f"report_np{args.world_size}_rank{args.rank}.html"
lazy_convert_sqlite(nsys_rep_file_path, sqlite_file_path)

conn = sqlite3.connect(f"file:{sqlite_file_path}?mode=ro", uri=True)

query = "SELECT * FROM ENUM_NSYS_EVENT_TYPE"
df = pd.read_sql_query(query, conn)
event_id_NvtxDomainCreate = df[df["name"] == "NvtxDomainCreate"].iloc[0]["id"].tolist()
event_id_NvtxPushPopRange = df[df["name"] == "NvtxPushPopRange"].iloc[0]["id"].tolist()

query = "SELECT domainId FROM NVTX_EVENTS WHERE eventType = ? AND text = ?"
df = pd.read_sql_query(query, conn, params=(event_id_NvtxDomainCreate, "NCCL"))
nccl_domain_id = -1 if df.empty else df.iloc[0]["domainId"].tolist()

query = """SELECT T1.start, T2.value AS text
    FROM NVTX_EVENTS AS T1
    JOIN StringIds AS T2 ON T1.textId = T2.id
    WHERE eventType = ? AND T2.value LIKE ?"""
df = pd.read_sql_query(query, conn, params=(event_id_NvtxPushPopRange, "layer_wise_benchmarks %"))
problem_start = []
problem_set = []
for start, text in df.itertuples(index=False):
    if text.startswith("layer_wise_benchmarks args {"):
        run_args = json.loads(text[len("layer_wise_benchmarks args") :])
    elif text.startswith("layer_wise_benchmarks problem_spec {"):
        problem_start.append(start)
        problem_set.append(
            {
                "spec": json.loads(text[len("layer_wise_benchmarks problem_spec") :]),
                "text": "",
                "runs": [],
                "runs_end": [],
                "ranges": [],
            }
        )

query = """SELECT T1.start, T1.end, T2.value AS text
    FROM NVTX_EVENTS AS T1
    JOIN StringIds AS T2 ON T1.textId = T2.id
    WHERE eventType = ? AND T2.value NOT LIKE ? AND T2.value NOT LIKE ? AND domainId != ?"""
df = pd.read_sql_query(
    query,
    conn,
    params=(event_id_NvtxPushPopRange, "layer_wise_benchmarks %", "[DG]%", nccl_domain_id),
)
for start, end, text in df.itertuples(index=False):
    problem_id = bisect.bisect(problem_start, start) - 1
    assert problem_id != -1
    if re.match(r"b=\d+ s=\d+ ", text):
        problem_set[problem_id]["text"] = text
        problem_set[problem_id]["runs"].append(start)
        problem_set[problem_id]["runs_end"].append(end)
    else:
        problem_set[problem_id]["ranges"].append((start, end, text))

query = """SELECT name FROM sqlite_master WHERE type = ?"""
df = pd.read_sql_query(query, conn, params=("table",))
tables = df["name"].tolist()
unified_subquery = """SELECT T1.start, T1.end, T1.demangledName, T1.correlationId, T1.graphNodeId
    FROM CUPTI_ACTIVITY_KIND_KERNEL AS T1"""
if "CUPTI_ACTIVITY_KIND_MEMCPY" in tables:
    unified_subquery += """ UNION ALL
        SELECT T2.start, T2.end, -2 AS demangledName, T2.correlationId, T2.graphNodeId
        FROM CUPTI_ACTIVITY_KIND_MEMCPY AS T2"""
if "CUPTI_ACTIVITY_KIND_MEMSET" in tables:
    unified_subquery += """ UNION ALL
        SELECT T3.start, T3.end, -3 AS demangledName, T3.correlationId, T3.graphNodeId
        FROM CUPTI_ACTIVITY_KIND_MEMSET AS T3"""
query = f"""SELECT unified.start, unified.end, unified.demangledName,
       R.start AS runtime_start, R.end AS runtime_end,
       CGE2.start AS capture_start, CGE2.end AS capture_end
FROM ({unified_subquery}) AS unified
JOIN CUPTI_ACTIVITY_KIND_RUNTIME AS R ON unified.correlationId = R.correlationId
LEFT JOIN CUDA_GRAPH_NODE_EVENTS AS CGE1 ON unified.graphNodeId = CGE1.graphNodeId AND
                                            CGE1.originalGraphNodeId IS NOT NULL
LEFT JOIN CUDA_GRAPH_NODE_EVENTS AS CGE2 ON CGE1.originalGraphNodeId = CGE2.graphNodeId"""
df = pd.read_sql_query(query, conn)
kernel_list = []
for (
    start,
    end,
    demangledName,
    runtime_start,
    runtime_end,
    capture_start,
    capture_end,
) in df.itertuples(index=False):
    problem_id = bisect.bisect(problem_start, start) - 1
    run_id = bisect.bisect(problem_set[problem_id]["runs"], runtime_start) - 1
    if (
        run_id == -1
        or run_id == len(problem_set[problem_id]["runs"])
        or runtime_start >= problem_set[problem_id]["runs_end"][run_id]
    ):
        run_id = -1
    ranges = [
        text
        for range_start, range_end, text in problem_set[problem_id]["ranges"]
        if capture_start >= range_start and capture_end <= range_end
    ]
    kernel_list.append(
        (
            problem_id,
            run_id,
            ranges,
            start,
            end,
            demangledName,
            runtime_start,
            runtime_end,
            capture_start,
            capture_end,
        )
    )
# TODO: Parse CTX phases

query = "SELECT * FROM StringIds"
df = pd.read_sql_query(query, conn)
string_ids = dict(zip(df["id"], df["value"]))

conn.close()

kernel_list.sort(key=lambda t: (t[6], t[8]))
kernels = [[[] for _ in problem["runs"]] for problem in problem_set]
for (
    problem_id,
    run_id,
    ranges,
    start,
    end,
    demangledName,
    runtime_start,
    runtime_end,
    capture_start,
    capture_end,
) in kernel_list:
    if run_id != -1:
        kernels[problem_id][run_id].append((demangledName, start, end, ranges))
for problem_id in range(len(kernels)):
    required_seq = [demangledName for demangledName, _, _, _ in kernels[problem_id][0]]
    for run_id in range(len(kernels[problem_id])):
        seq = [demangledName for demangledName, _, _, _ in kernels[problem_id][run_id]]
        assert seq == required_seq


parser_keywords = [
    ("cuBLASGemm", "nvjet"),
    ("splitKreduce", "splitKreduce_kernel"),
    ("fusedAGemm", "fused_a_gemm_kernel"),
    ("RMSNorm", "RMSNormKernel"),
    ("torchCat", "CatArrayBatchedCopy"),
    ("applyMLARope", "applyMLARope"),
    ("fmhaSm100f", "fmhaSm100fKernel_Qkv"),
    ("fmhaReduction", "fmhaReductionKernel"),
    ("quant", "quantize_with_block_size"),
    ("AllGather", "ncclDevKernel_AllGather_"),
    ("ReduceScatter", "ncclDevKernel_ReduceScatter_"),
    ("allreduce_oneshot", "allreduce_fusion_kernel_oneshot_lamport"),
    ("allreduce_twoshot", "allreduce_fusion_kernel_twoshot_sync"),
    ("expandInput", "expandInputRowsKernel"),
    ("computeStrides", "computeStridesTmaWarpSpecializedKernel"),
    ("cutlassGroupGemm", "cutlass::device_kernel<cutlass::gemm::kernel::GemmUniversal"),
    ("doActivation", "doActivationKernel"),
    ("cutlassGemm", "GemmUniversal"),
    ("deepseek_v3_topk", "deepseek_v3_topk_kernel"),
    ("CountAndIndice", "computeCountAndIndiceDevice"),
    ("Cumsum", "computeCumsumDevice"),
    ("moveIndice", "moveIndiceDevice"),
    ("moeAllToAll", "moeAllToAllKernel"),
    ("moeA2APrepareDispatch", "moe_comm::moeA2APrepareDispatchKernel"),
    ("moeA2ADispatch", "moe_comm::moeA2ADispatchKernel"),
    ("moeA2ASanitizeExpertIds", "moe_comm::moeA2ASanitizeExpertIdsKernel"),
    ("moeA2APrepareCombine", "moe_comm::moeA2APrepareCombineKernel"),
    ("moeA2ACombine", "moe_comm::moeA2ACombineKernel"),
    ("memsetExpertIds", "memsetExpertIdsDevice"),
    ("blockSum", "blockExpertPrefixSumKernel"),
    ("globalSum", "globalExpertPrefixSumKernel"),
    ("mergePrefix", "mergeExpertPrefixSumKernel"),
    ("fusedBuildExpertMaps", "fusedBuildExpertMapsSortFirstTokenKernel"),
    ("swiglu", "silu_and_mul_kernel"),
    ("torchAdd", "CUDAFunctor_add"),
    ("torchFill", "at::native::FillFunctor"),
    ("triton_fused_add_sum", "triton_red_fused_add_sum_0"),
    ("torchCopy", "at::native::bfloat16_copy_kernel_cuda"),
    ("torchDistribution", "distribution_elementwise_grid_stride_kernel"),
    ("torchArange", "at::native::arange_cuda_out"),
    ("torchDirectCopy", "at::native::direct_copy_kernel_cuda"),
    ("torchBitonicSort", "at::native::bitonicSortKVInPlace"),
    ("routingInitExpertCounts", "routingInitExpertCounts"),
    ("routingIndicesCluster", "routingIndicesClusterKernel"),
    ("routingIndicesCoop", "routingIndicesCoopKernel"),
    ("bmm_4_44_32", "bmm_E2m1_E2m1E2m1_Fp32_t"),
    ("finalize", "finalize::finalizeKernel"),
    ("bmm_16_44_32", "bmm_Bfloat16_E2m1E2m1_Fp32_"),
    ("deep_gemm_gemm", "deep_gemm::sm100_fp8_gemm_1d1d_impl<"),
    ("per_token_quant", "_per_token_quant_and_transform_kernel"),
    ("triton_fused_layer_norm", "triton_per_fused__to_copy_native_layer_norm_0"),
    ("flashinferRoPE", "flashinfer::BatchQKApplyRotaryPosIdsCosSinCacheHeadParallelismKernel<"),
    ("fp8_blockscale_gemm", "tensorrt_llm::kernels::fp8_blockscale_gemm"),
    ("triton_fused_mul_squeeze", "triton_poi_fused_mul_squeeze_0"),
    ("indexerKCacheScatter", "tensorrt_llm::kernels::indexerKCacheScatterUnifiedKernel"),
    ("deep_gemm_mqa_logits", "deep_gemm::sm100_fp8_paged_mqa_logits<"),
    ("topKPerRowDecode", "tensorrt_llm::kernels::topKPerRowDecode<"),
    ("torchAdd<int>", "at::native::CUDAFunctorOnSelf_add"),
    ("convert_req_index", "_convert_req_index_to_global_index_kernel_with_stride_factor"),
    ("preprocess_after_permute", "_preprocess_after_permute_kernel"),
    ("masked_index_copy_quant", "_masked_index_copy_group_quant_fp8"),
    ("swiglu_quant", "_silu_and_mul_post_quant_kernel"),
    ("masked_index_gather", "masked_index_gather_kernel"),
    ("finalizeMoeRouting", "tensorrt_llm::kernels::cutlass_kernels::finalizeMoeRoutingKernel<"),
    ("fused_qkvzba_split", "fused_qkvzba_split_reshape_cat_kernel"),
    ("causal_conv1d_update", "tensorrt_llm::kernels::causal_conv1d::causal_conv1d_update_kernel<"),
    ("fused_delta_rule_update", "fused_sigmoid_gating_delta_rule_update_kernel"),
    ("layer_norm_fwd_1pass", "_layer_norm_fwd_1pass_kernel"),
    ("torchGatherTopK", "at::native::sbtopk::gatherTopK<"),
    ("softmax_warp_forward", "softmax_warp_forward<"),
    ("torchSigmoid", "at::native::sigmoid_kernel_cuda"),
    ("torchMul", "at::native::binary_internal::MulFunctor<"),
    ("applyBiasRopeUpdateKVCache", "tensorrt_llm::kernels::applyBiasRopeUpdateKVCacheV2<"),
    ("routingIndicesHistogramScores", "routingRenormalize::routingIndicesHistogramScoresKernel<"),
    ("routingIndicesHistogram", "routingIndicesHistogramKernel<"),
    ("routingIndicesOffsets", "routingIndicesOffsetsKernel<"),
    ("torchReduceSum", ["at::native::reduce_kernel<", "at::native::sum_functor<"]),
    ("CuteDSLMoePermute", "cute_dsl::moePermuteKernel"),
    (
        "CuteDSLGroupedGemmSwiglu",
        ["cute_dsl_kernels", "blockscaled_contiguous_grouped_gemm_swiglu_fusion"],
    ),
    (
        "CuteDSLGroupedGemmFinalize",
        ["cute_dsl_kernels", "blockscaled_contiguous_grouped_gemm_finalize_fusion"],
    ),
]
warned_names = set()


def parse_kernel_name(demangledName):
    if demangledName == -2:
        return "Memcpy"
    if demangledName == -3:
        return "Memset"
    name = string_ids[demangledName]
    for dst, src in parser_keywords:
        if not isinstance(src, (tuple, list)):
            src = [src]
        if all(keyword in name for keyword in src):
            return dst
    if name not in warned_names:
        print(f"Unknown kernel name: {name}")
        warned_names.add(name)
        if args.error_on_unknown_kernel:
            raise NotImplementedError(f"Unknown kernel name: {name}")
    return name[:30]


converted_seqs = []
for runs in kernels:
    warmup_times = run_args["warmup_times"] if args.warmup_times is None else args.warmup_times
    converted_seq = []
    # Kernel time
    for i, (demangledName, _, _, ranges) in enumerate(runs[0]):
        name = parse_kernel_name(demangledName)
        category = (*ranges, name)
        time_list = [run[i][2] - run[i][1] for run in runs]
        t = np.mean(time_list[warmup_times:]).tolist()
        converted_seq.append((category, t))
    # Space and Overlap
    overlap_list = []
    space_list = []
    for run in runs:
        sorted_run = sorted(run, key=lambda op: op[1])
        last_end = sorted_run[0][1]
        overlap_time = 0
        space_time = 0
        for _, start, end, _ in sorted_run:
            if start > last_end:
                space_time += start - last_end
            else:
                overlap_time += min(last_end, end) - start
            last_end = max(last_end, end)
        overlap_list.append(-overlap_time)
        space_list.append(space_time)
    converted_seq.append((("Overlap",), np.mean(overlap_list[warmup_times:]).tolist()))
    converted_seq.append((("Space",), np.mean(space_list[warmup_times:]).tolist()))
    converted_seq.append((("Total",), sum(t for _, t in converted_seq)))
    converted_seqs.append(converted_seq)

merged_title = []
for converted_seq in converted_seqs:
    title = [name for name, _ in converted_seq]
    merged_title = shortest_common_supersequence(merged_title, title)

merged_data = [[0.0] * len(problem_set) for _ in merged_title]
for problem_id, converted_seq in enumerate(converted_seqs):
    cur = 0
    for category, t in converted_seq:
        cur = merged_title.index(category, cur)
        merged_data[cur][problem_id] = t
        cur += 1

print("Problem set:")
for problem in problem_set:
    print(
        f'- "{problem["text"]}"    {len(problem["runs"])} runs'
        f"    Ranges: [{', '.join(text for _, _, text in problem['ranges'])}]"
    )

stack = []
csv_data = [["", *[problem["text"] for problem in problem_set]]]
js_data = []
js_stack = [js_data]
max_title_len = max((len(title) - 1) * 3 + len(title[-1]) for title in merged_title)
for title, time_data in zip(merged_title, merged_data):
    while stack != list(title[: len(stack)]):
        level_title = stack[-1]
        stack.pop()
        js_stack[-2].append(
            {
                "name": level_title,
                "children": js_stack[-1],
            }
        )
        js_stack.pop()
    while len(stack) != len(title) - 1:
        level_title = title[len(stack)]
        stack.append(level_title)
        level = len(stack)
        print("|--" * (level - 1) + level_title)
        csv_data.append(["|--" * (level - 1) + level_title])
        js_stack.append([])
    level = len(stack) + 1
    print(
        "|--" * (level - 1) + title[-1] + " " * (max_title_len - (level - 1) * 3 - len(title[-1])),
        *[f"{x / 1000:-6.1f}" for x in time_data],
    )
    csv_data.append(["|--" * (level - 1) + title[-1], *[f"{x / 1000:.1f}" for x in time_data]])
    if title != ("Total",):
        js_stack[-1].append(
            {
                "name": title[-1],
                "time": [x / 1000 for x in time_data],
            }
        )
# TODO: Group repeated modules
with csv_file_path.open("w", newline="") as f:
    csv_writer = csv.writer(f, quoting=csv.QUOTE_MINIMAL)
    for row in csv_data:
        csv_writer.writerow(row)
js_header_config = [{"name": problem["text"]} for problem in problem_set]
loader = jinja2.FileSystemLoader(Path(__file__).parent)
template = jinja2.Environment(loader=loader).get_template("template.html")
with html_file_path.open("w") as f:
    f.write(
        template.render(
            headerConfig=js_header_config, rawData=js_data, runArgs=json.dumps(run_args, indent=4)
        )
    )