diff --git a/scripts/profiling/generate_ncu_plan.py b/scripts/profiling/generate_ncu_plan.py index e67b612d..8e635392 100644 --- a/scripts/profiling/generate_ncu_plan.py +++ b/scripts/profiling/generate_ncu_plan.py @@ -193,26 +193,78 @@ def time_or_none(item: tuple[str, str] | None) -> float | None: return kernels +def _strip_function_arguments(kernel_name: str) -> str: + """Drop the top-level function argument list from a demangled kernel name.""" + + template_depth = 0 + paren_depth = 0 + last_top_level_open = -1 + for index, char in enumerate(kernel_name): + if char == "<": + template_depth += 1 + elif char == ">" and template_depth > 0: + template_depth -= 1 + elif template_depth == 0 and char == "(": + if paren_depth == 0: + last_top_level_open = index + paren_depth += 1 + elif template_depth == 0 and char == ")" and paren_depth > 0: + paren_depth -= 1 + if last_top_level_open >= 0: + return kernel_name[:last_top_level_open].strip() + return kernel_name.strip() + + +def _last_top_level_token(stem: str) -> str: + """Return the final whitespace-delimited token outside templates/parens.""" + + template_depth = 0 + paren_depth = 0 + token_start = 0 + for index, char in enumerate(stem): + if char == "<": + template_depth += 1 + elif char == ">" and template_depth > 0: + template_depth -= 1 + elif template_depth == 0 and char == "(": + paren_depth += 1 + elif template_depth == 0 and char == ")" and paren_depth > 0: + paren_depth -= 1 + elif template_depth == 0 and paren_depth == 0 and char.isspace(): + token_start = index + 1 + return stem[token_start:].strip() or stem.strip() + + +def _escape_regex_with_flexible_whitespace(token: str) -> str: + parts: list[str] = [] + in_whitespace = False + for char in token: + if char.isspace(): + if not in_whitespace: + parts.append("[[:space:]]*") + in_whitespace = True + continue + parts.append(re.escape(char)) + in_whitespace = False + return "".join(parts) + + def _kernel_match_token(kernel_name: str) -> str: - """Return a shell-safe demangled-name substring for NCU kernel matching. + """Return a demangled-name substring for NCU kernel matching. - Nsight Systems summaries often include the full demangled signature. Passing - that signature through shell variables is fragile because argument lists and - template parameters contain spaces. NCU only needs a stable substring, so - use the demangled function stem instead of the full signature. + Keep C++ template arguments because they may be the semantic distinction + between CUDA kernels. Drop only the top-level function argument list and + return type. """ - stem = kernel_name.split("(", 1)[0].strip() - stem = re.sub(r"<.*>$", "", stem).strip() - parts = stem.split() - if parts: - stem = parts[-1] + stem = _strip_function_arguments(kernel_name) + stem = _last_top_level_token(stem) return stem or kernel_name.strip() def _regex_kernel_name(kernel_name: str) -> str: token = _kernel_match_token(kernel_name) - return f"regex:.*{re.escape(token)}.*" + return f"regex:.*{_escape_regex_with_flexible_whitespace(token)}.*" def _slugify_kernel_name(kernel_name: str, index: int) -> str: diff --git a/scripts/tests/test_ncu_plan_generation.sh b/scripts/tests/test_ncu_plan_generation.sh index 49db0e73..b3499174 100644 --- a/scripts/tests/test_ncu_plan_generation.sh +++ b/scripts/tests/test_ncu_plan_generation.sh @@ -55,7 +55,7 @@ jq -e ' .profiles[0].section == null and .profiles[0].kernel_match.pattern == "regex:.*kern_compute_force_nonbond_table_linear_univ__inter_cell.*" and .profiles[1].kernel_match.pattern == "regex:.*kern_compute_force_nonbond_table_linear_univ__intra_cell.*" and - .profiles[2].kernel_match.pattern == "regex:.*kern_build_pairlist.*" and + .profiles[2].kernel_match.pattern == "regex:.*kern_build_pairlist<4,[[:space:]]*256>.*" and (.profiles[2].kernel_name | contains("build_pairlist")) ' "${TMP_DIR}/ncu_plan.json" >/dev/null @@ -93,7 +93,7 @@ jq -e ' .profiles[0].selection.source_gpu_duration_ns == 9750000 and .profiles[0].selection.discovery_gpu_time_pct == 97.5 and .profiles[1].kernel_match.pattern == "regex:.*kern_compute_force_nonbond_table_linear_univ__force_intra_cell.*" and - .profiles[2].kernel_match.pattern == "regex:.*kern_build_pairlist.*" + .profiles[2].kernel_match.pattern == "regex:.*kern_build_pairlist<4,[[:space:]]*256>.*" ' "${TMP_DIR}/dominant_ncu_plan.json" >/dev/null "${PYTHON_BIN}" "${REPO_DIR}/scripts/profiling/generate_ncu_plan.py" \ @@ -107,6 +107,32 @@ jq -e ' .profiles[3].kernel_match.pattern == "regex:.*kern_compute_energy_nonbond_table_linear_univ__energyforce_inter_cell.*" ' "${TMP_DIR}/all_ncu_plan.json" >/dev/null +cat > "${TMP_DIR}/sbd_nsys.csv" <<'CSV' +CUDA Kernel Summary +"Time (%)","Total Time (ns)","Instances","Avg (ns)","Med (ns)","Min (ns)","Max (ns)","StdDev (ns)","Name" +58.5,"5,850,000",50,117000,116000,100000,130000,4000,"void sbd::MultAlphaBeta(double*, double const*)" +19.2,"1,920,000",50,38400,38000,35000,42000,2000,"void sbd::MultUnified(double*, double const*)" +17.0,"1,700,000",50,34000,33800,30000,39000,1800,"void sbd::MultUnified(double*, double const*)" +2.7,"270,000",50,5400,5300,5000,6500,400,"void sbd::MultUnified(double*, double const*)" +2.5,"250,000",50,5000,4900,4500,6000,350,"void sbd::MultUnified(double*, double const*)" +CSV + +"${PYTHON_BIN}" "${REPO_DIR}/scripts/profiling/generate_ncu_plan.py" \ + --nsys-csv "${TMP_DIR}/sbd_nsys.csv" \ + --out-plan "${TMP_DIR}/sbd_ncu_plan.json" \ + --top-k 5 \ + --launch-count 10 >/dev/null + +jq -e ' + (.profiles | length) == 5 and + .profiles[0].kernel_match.pattern == "regex:.*sbd::MultAlphaBeta.*" and + .profiles[1].kernel_match.pattern == "regex:.*sbd::MultUnified.*" and + .profiles[2].kernel_match.pattern == "regex:.*sbd::MultUnified.*" and + .profiles[3].kernel_match.pattern == "regex:.*sbd::MultUnified.*" and + .profiles[4].kernel_match.pattern == "regex:.*sbd::MultUnified.*" and + ([.profiles[].kernel_match.pattern] | unique | length) == 5 +' "${TMP_DIR}/sbd_ncu_plan.json" >/dev/null + jq -e ' .schema_version == 1 and .execution.profiler == "ncu" and