Preserve C++ template arguments in NCU kernel matches

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:

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*, double const*)"
+19.2,"1,920,000",50,38400,38000,35000,42000,2000,"void sbd::MultUnified<double, (int)0, (int)1>(double*, double const*)"
+17.0,"1,700,000",50,34000,33800,30000,39000,1800,"void sbd::MultUnified<double, (int)1, (int)1>(double*, double const*)"
+2.7,"270,000",50,5400,5300,5000,6500,400,"void sbd::MultUnified<double, (int)0, (int)0>(double*, double const*)"
+2.5,"250,000",50,5000,4900,4500,6000,350,"void sbd::MultUnified<double, (int)1, (int)0>(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<double>.*" and
+  .profiles[1].kernel_match.pattern == "regex:.*sbd::MultUnified<double,[[:space:]]*\\(int\\)0,[[:space:]]*\\(int\\)1>.*" and
+  .profiles[2].kernel_match.pattern == "regex:.*sbd::MultUnified<double,[[:space:]]*\\(int\\)1,[[:space:]]*\\(int\\)1>.*" and
+  .profiles[3].kernel_match.pattern == "regex:.*sbd::MultUnified<double,[[:space:]]*\\(int\\)0,[[:space:]]*\\(int\\)0>.*" and
+  .profiles[4].kernel_match.pattern == "regex:.*sbd::MultUnified<double,[[:space:]]*\\(int\\)1,[[:space:]]*\\(int\\)0>.*" 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