feat(extraction): add CUDA language support (.cu/.cuh) (#387, #648) (#1172)

CUDA rides the C++ grammar via the Metal (#1121) dialect pattern:
blankCudaConstructs (offset-preserving) blanks execution-space specifiers
(__global__ family), __launch_bounds__(...), and <<<grid, block>>> launch
configs — which otherwise lex as shift operators and destroy the
host→kernel call edge entirely. Gated by .cu/.cuh extension OR by content
(looksLikeCudaSource), because much real CUDA lives in .h/.hpp headers:
cutlass launches most kernels from headers and flash-attention's launch
templates are .h. Safe by construction — no CUDA marker is valid C++
anywhere, and the launch blank is bounded + brace-balance-checked so a
stray <<< (committed merge-conflict markers) can never blank real code.

All real-world launch styles connect: plain, templated
(k<T, 256><<<...>>>), function-pointer (auto kernel = &fn<...>; with
branch reassignments each linked), dim3{...} brace-init configs, and
kernels defined through name-in-first-argument macros
(DEFINE_FLASH_FORWARD_KERNEL style — gtest TEST_F / PYBIND11_MODULE
shapes deliberately excluded by the two-lone-identifiers rule).

Two general C++ resolution wins the flow validation forced out:
- namespace blocks now prefix contained symbols' qualifiedNames
  (prefix-only — no namespace nodes, avoiding #1093-style crowd-out), so
  ns::fn(...) calls resolve; previously every namespace-qualified C++
  call was a permanently dead edge. cutlass: +30,864 edges (~10%), node
  count byte-identical.
- templated callees (fn<T, 256>(args)) strip template args at extraction
  (mirroring #1043 for base classes), so they match their definitions.

Validated on llm.c (165 host→kernel launch edges, was 0),
flash-attention (run_flash_fwd → flash_fwd_kernel → compute_attn traces
in one codegraph_explore call), and NVIDIA CUTLASS; fmt as the plain-C++
control (unchanged). A/B n=2/arm: Read/Grep displacement decisive on all
three repos (flash-attention Reads 29,13 → 5,2).

Co-authored-by: Claude Fable 5 <noreply@anthropic.com>

src/extraction/grammars.ts

@@ -125,6 +125,13 @@ export const EXTENSION_MAP: Record<string, Language> = {
   // structs, and calls. MSL-specific `[[attribute]]` annotations are blanked
   // pre-parse for `.metal` files (see blankMetalAttributes in c-cpp.ts). (#1121)
   '.metal': 'cpp',
+  // CUDA ≈ C++ plus execution-space specifiers (`__global__` …) and
+  // `<<<grid, block>>>` kernel-launch syntax: the C++ grammar extracts its
+  // functions/structs/classes/calls once blankCudaConstructs (pre-parse; gated
+  // by these extensions OR by content for CUDA living in `.h`/`.hpp` headers —
+  // see c-cpp.ts) blanks the CUDA-only tokens. (#387)
+  '.cu': 'cpp',
+  '.cuh': 'cpp',
   // XML: file-level tracking; the MyBatis extractor matches `<mapper namespace="...">`
   // shape and emits SQL-statement nodes (other XML returns empty).
   '.xml': 'xml',

src/extraction/languages/c-cpp.ts

@@ -27,7 +27,54 @@ function findDeclaratorQualifiedId(declarator: SyntaxNode): SyntaxNode | undefin
   return undefined;
 }
 
+/**
+ * Recover the real function name from the macro-definition idiom
+ * `MACRO_NAME(real_name, typed args…) { body }` — flash-attention's
+ * `DEFINE_FLASH_FORWARD_KERNEL(flash_fwd_kernel, bool Is_dropout, …) { … }`
+ * being the motivating case: tree-sitter parses the invocation as a
+ * function_definition NAMED after the macro, so every such kernel shared one
+ * name (`DEFINE_FLASH_FORWARD_KERNEL`) and the launch sites' calls to the real
+ * names (`flash_fwd_kernel<…><<<…>>>`) could never resolve.
+ *
+ * Deliberately narrow so name-in-first-arg is unambiguous — ALL of:
+ *  - the parsed name is macro-shaped: ALL-CAPS with at least one underscore
+ *    (`TEST` never matches; K&R C definitions have lowercase names);
+ *  - the first "parameter" is a LONE identifier (no type, no declarator)
+ *    containing a lowercase letter — the name being defined;
+ *  - at least one more parameter follows and NONE of them is another lone
+ *    identifier — a second bare arg means the first isn't the name (gtest's
+ *    `TEST_F(Fixture, Name)`, `PYBIND11_MODULE(ext, m)`,
+ *    google-benchmark's `BENCHMARK_DEFINE_F(Fix, name)` all bail here).
+ */
+function recoverCppMacroDefinedName(node: SyntaxNode, source: string): string | undefined {
+  if (node.type !== 'function_definition') return undefined;
+  const declarator = getChildByField(node, 'declarator');
+  if (declarator?.type !== 'function_declarator') return undefined;
+  const inner = getChildByField(declarator, 'declarator');
+  if (inner?.type !== 'identifier') return undefined;
+  const macroName = getNodeText(inner, source);
+  if (!/^[A-Z][A-Z0-9]*(?:_[A-Z0-9]+)+$/.test(macroName)) return undefined;
+  const params = getChildByField(declarator, 'parameters');
+  if (!params || params.namedChildCount < 2) return undefined;
+  const loneIdentText = (p: SyntaxNode): string | null =>
+    p.type === 'parameter_declaration' &&
+    p.namedChildCount === 1 &&
+    p.namedChild(0)?.type === 'type_identifier'
+      ? getNodeText(p.namedChild(0)!, source)
+      : null;
+  const first = params.namedChild(0);
+  const name = first ? loneIdentText(first) : null;
+  if (!name || !/[a-z]/.test(name)) return undefined;
+  for (let i = 1; i < params.namedChildCount; i++) {
+    const p = params.namedChild(i);
+    if (p && loneIdentText(p) !== null) return undefined;
+  }
+  return name;
+}
+
 function extractCppQualifiedMethodName(node: SyntaxNode, source: string): string | undefined {
+  const macroDefined = recoverCppMacroDefinedName(node, source);
+  if (macroDefined) return macroDefined;
   const declarator = getChildByField(node, 'declarator');
   if (!declarator) return undefined;
   const qid = findDeclaratorQualifiedId(declarator);
@@ -123,6 +170,8 @@ function extractCppReturnType(node: SyntaxNode, source: string): string | undefi
 }
 
 export const cExtractor: LanguageExtractor = {
+  // CUDA in C-detected headers (content-gated blank; see preParseCSource).
+  preParse: preParseCSource,
   // Universal net: recover a real name from any macro-mangled function name.
   recoverMangledName: recoverMangledCppName,
   functionTypes: ['function_definition'],
@@ -384,14 +433,108 @@ export function blankMetalAttributes(source: string): string {
   return source.replace(METAL_ATTRIBUTE_RE, (m) => ' '.repeat(m.length));
 }
 
+/**
+ * Blank CUDA-specific constructs before parsing `.cu`/`.cuh` files (parsed with
+ * the C++ grammar). Three shapes tree-sitter-cpp can't reconcile, each replaced
+ * with equal-length whitespace so every byte offset survives (#387):
+ *
+ * 1. Execution-space / storage specifiers: in `__global__ void step(…)` or
+ *    `__shared__ float tile[256]` the specifier parses as the declaration's
+ *    TYPE and shunts the real return/value type into an ERROR node — mangling
+ *    signatures and, for `__shared__` arrays, the declared name itself. Blanked
+ *    unconditionally (no following-token lookahead) so extended lambdas
+ *    (`[=] __device__ (int i) { … }`) recover too. `__restrict__` is deliberately
+ *    absent: the grammar already parses it natively as a type_qualifier.
+ * 2. `__launch_bounds__(…)` between specifier and declarator — same misparse.
+ *    The parenthesized form is blanked first; a bare leftover token is caught
+ *    by the specifier list.
+ * 3. Kernel-launch configs `step<<<grid, block, smem, stream>>>(args)`: the
+ *    chevrons lex as shift operators around an empty-named template, so no
+ *    call_expression exists and the host→kernel call edge — the main reason to
+ *    index CUDA at all — is lost. Blanking the `<<<…>>>` span leaves
+ *    `step                              (args)`, a plain call the grammar
+ *    parses natively (templated launches `k<T, 256><<<…>>>(…)` included).
+ *
+ * The launch-config match is deliberately bounded — statement/brace characters
+ * excluded, span capped, newlines preserved by the replacer — so a stray `<<<`
+ * (a committed merge-conflict marker, a string literal) can never blank a run
+ * of real code: an unmatched launch degrades to the status quo for that call
+ * site (no call edge), never to corruption. Applied to `.cu`/`.cuh` files and —
+ * because much real CUDA lives in extension-less headers (cutlass launches the
+ * majority of its kernels from `.h`; flash-attention's launch templates are
+ * `.h`; llm.c keeps device helpers in C-detected `.h`) — to any C/C++-family
+ * file whose CONTENT carries a strong CUDA marker (`looksLikeCudaSource`).
+ * Unlike Metal's `[[attribute]]` (legal C++ syntax elsewhere, hence Metal's
+ * strict extension gate), no CUDA marker is valid C++ anywhere: `<<<` isn't
+ * legal syntax and the dunder specifiers are implementation-reserved names no
+ * real codebase defines — so a content-triggered blank on a non-CUDA file can
+ * only ever whitespace tokens inside comments or strings, which parse the same.
+ */
+const CUDA_LAUNCH_BOUNDS_RE = /\b__launch_bounds__\s*\([^()\n]*\)/g;
+const CUDA_SPECIFIER_RE =
+  /\b__(?:global|device|host|constant|shared|managed|grid_constant|forceinline|noinline|launch_bounds)__\b/g;
+// `;` stays excluded (launch configs are expressions; a stray `<<<` spanning
+// real statements always crosses one) and the span is capped. Braces are
+// allowed through the regex — `k<<<dim3{1,1,1}, dim3{256,1,1}>>>(…)` is a real
+// launch shape — but the replacer only blanks a BALANCED match: a merge
+// conflict's `<<<<<<< … >>>>>>>` region that dodges every `;` still opens
+// braces it never closes, so it fails the balance check and stays untouched.
+const CUDA_LAUNCH_CONFIG_RE = /<<<[^;]{0,400}?>>>/g;
+export function blankCudaConstructs(source: string): string {
+  let out = source;
+  if (out.indexOf('__') !== -1) {
+    out = out
+      .replace(CUDA_LAUNCH_BOUNDS_RE, (m) => ' '.repeat(m.length))
+      .replace(CUDA_SPECIFIER_RE, (m) => ' '.repeat(m.length));
+  }
+  if (out.indexOf('<<<') !== -1) {
+    out = out.replace(CUDA_LAUNCH_CONFIG_RE, (m) => {
+      let depth = 0;
+      for (let i = 0; i < m.length; i++) {
+        const ch = m.charCodeAt(i);
+        if (ch === 0x7b /* { */) depth++;
+        else if (ch === 0x7d /* } */ && --depth < 0) return m;
+      }
+      return depth === 0 ? m.replace(/[^\n]/g, ' ') : m;
+    });
+  }
+  return out;
+}
+
+/** Strong content markers for CUDA source in files without a CUDA extension
+ * (headers). The dunders are execution-space specifiers that only nvcc defines;
+ * `cudaStream_t` is the runtime's stream handle, pervasive in launcher headers
+ * that themselves declare no kernel. Deliberately excludes weak markers (`dim3`,
+ * `<<<`) that could plausibly appear in non-CUDA text. */
+function looksLikeCudaSource(source: string): boolean {
+  return (
+    source.indexOf('__global__') !== -1 ||
+    source.indexOf('__device__') !== -1 ||
+    source.indexOf('__constant__') !== -1 ||
+    source.indexOf('cudaStream_t') !== -1
+  );
+}
+
 /** C/C++ source pre-processing before tree-sitter: recover both macro-annotated
- * class definitions and macro-prefixed function definitions — plus, for `.metal`
- * shaders (parsed with the C++ grammar), MSL attribute annotations. Offset-preserving. */
+ * class definitions and macro-prefixed function definitions — plus the non-C++
+ * surface of the dialects parsed with the C++ grammar: `.metal` MSL attribute
+ * annotations, and CUDA specifiers + launch syntax (by `.cu`/`.cuh` extension
+ * or by content, for CUDA living in `.h`/`.hpp` headers). Offset-preserving. */
 function preParseCppSource(source: string, filePath?: string): string {
   const blanked = blankCppInlineMacros(blankCppExportMacros(source));
-  return filePath && filePath.toLowerCase().endsWith('.metal')
-    ? blankMetalAttributes(blanked)
-    : blanked;
+  const lower = filePath ? filePath.toLowerCase() : '';
+  if (lower.endsWith('.metal')) return blankMetalAttributes(blanked);
+  if (lower.endsWith('.cu') || lower.endsWith('.cuh') || looksLikeCudaSource(source)) {
+    return blankCudaConstructs(blanked);
+  }
+  return blanked;
+}
+
+/** C source pre-processing: C-detected headers in CUDA projects (llm.c keeps
+ * `__device__` helpers and kernel prototypes in plain `.h`) get the same
+ * content-gated CUDA blank as C++. */
+function preParseCSource(source: string): string {
+  return looksLikeCudaSource(source) ? blankCudaConstructs(source) : source;
 }
 
 export const cppExtractor: LanguageExtractor = {

src/extraction/tree-sitter.ts

@@ -382,6 +382,20 @@ export class TreeSitterExtractor {
   private errors: ExtractionError[] = [];
   private extractor: LanguageExtractor | null = null;
   private nodeStack: string[] = []; // Stack of parent node IDs
+  // C/C++ enclosing `namespace ns { … }` names, prepended to every contained
+  // symbol's qualifiedName (see visitNode). Prefix-only by design — no
+  // namespace NODE is created: `namespace cutlass {` opens in thousands of
+  // files, and a node per block would flood search with same-named symbols
+  // (the #1093 crowd-out failure mode). Always empty outside C/C++.
+  private namespacePrefix: string[] = [];
+  // C++ local function-pointer bindings, per enclosing symbol:
+  // `auto kernel = &flash_fwd_kernel<…>;` recorded as callerId → kernel →
+  // {flash_fwd_kernel}, so a later `kernel<<<grid, block>>>(params)` (or plain
+  // `kernel(args)`) in the same body emits calls refs to the real target(s)
+  // instead of an unresolvable local name. Branch reassignments accumulate —
+  // each assigned target is a genuine possible callee. Same-body locality is
+  // the precision guard (the #932 table-dispatch philosophy scoped to locals).
+  private cppLocalFnPtrs = new Map<string, Map<string, Set<string>>>();
   private methodIndex: Map<string, string> | null = null; // lookup key → node ID for Pascal defProc lookup
   // Function-as-value capture (#756): per-language spec + candidates collected
   // during the walk, gated & flushed into unresolvedReferences at end-of-file
@@ -909,6 +923,31 @@ export class TreeSitterExtractor {
       if (skipChildren) return;
     }
 
+    // C++ namespace blocks: carry the namespace name as a qualifiedName prefix
+    // while walking the body, so `namespace flash { void compute_attn(); }`
+    // indexes compute_attn with qualifiedName `flash::compute_attn` and a
+    // namespace-qualified call (`flash::compute_attn(...)`) resolves by exact
+    // qualified match instead of never resolving — C++ namespaces previously
+    // left no trace in qualifiedNames at all, so every `ns::fn()` call site
+    // was a permanently dead edge (surfaced by #387 flow validation on
+    // flash-attention/cutlass, whose kernel dispatch is namespace-qualified).
+    // C++17 nested forms (`namespace a::b {`) prefix as written. An anonymous
+    // namespace falls through to the generic walk — its contents stay bare,
+    // matching how call sites spell them.
+    if (this.language === 'cpp' && nodeType === 'namespace_definition') {
+      const nameNode = getChildByField(node, 'name');
+      const nsName = nameNode ? getNodeText(nameNode, this.source) : '';
+      if (nsName) {
+        this.namespacePrefix.push(nsName);
+        for (let i = 0; i < node.namedChildCount; i++) {
+          const child = node.namedChild(i);
+          if (child) this.visitNode(child);
+        }
+        this.namespacePrefix.pop();
+        return;
+      }
+    }
+
     // Function-as-value capture (#756) — independent of the dispatch ladder
     // below (the captured container types have no other handler there), so it
     // can never shadow or be shadowed by an extraction branch.
@@ -1345,7 +1384,8 @@ export class TreeSitterExtractor {
   private buildQualifiedName(name: string): string {
     // Build a qualified name from the semantic hierarchy only (no file path).
     // The file path is stored separately in filePath and pollutes FTS if included here.
-    const parts: string[] = [];
+    // C/C++ enclosing namespaces prefix first (empty for every other language).
+    const parts: string[] = [...this.namespacePrefix];
     for (const nodeId of this.nodeStack) {
       const node = this.nodes.find((n) => n.id === nodeId);
       if (node && node.kind !== 'file') {
@@ -4183,6 +4223,44 @@ export class TreeSitterExtractor {
       if (conv && conv[1]) calleeName = conv[1];
     }
 
+    // C/C++ templated callees — a direct templated call (`fn<T, 256>(args)`,
+    // the shape every CUDA kernel-launch site takes once its `<<<…>>>` config
+    // is blanked) or a qualified one (`ns::fn<T>(args)`) — carry template
+    // arguments in the callee text, which can never match the bare name the
+    // function was DEFINED as, so the call edge silently never resolves. Strip
+    // them: the same normalization base-class `extends` refs already get
+    // (#1043). `operator<`/`operator<<` callees are excluded — their `<` is the
+    // operator itself, not a template-argument list.
+    if (
+      calleeName &&
+      calleeName.includes('<') &&
+      (this.language === 'cpp' || this.language === 'c') &&
+      !calleeName.includes('operator')
+    ) {
+      calleeName = stripCppTemplateArgs(calleeName);
+    }
+
+    // C++ call/launch through a local function pointer: `auto kernel =
+    // &flash_fwd_kernel<…>; … kernel<<<grid, block>>>(params);` — the callee
+    // is an unresolvable local name. When the same enclosing symbol bound the
+    // local from `&fn` (each branch assignment counts), emit the call against
+    // every recorded target instead of the local.
+    if (calleeName && this.language === 'cpp' && /^[A-Za-z_]\w*$/.test(calleeName)) {
+      const targets = this.cppLocalFnPtrs.get(callerId)?.get(calleeName);
+      if (targets && targets.size > 0) {
+        for (const target of targets) {
+          this.unresolvedReferences.push({
+            fromNodeId: callerId,
+            referenceName: target,
+            referenceKind: 'calls',
+            line: node.startPosition.row + 1,
+            column: node.startPosition.column,
+          });
+        }
+        return;
+      }
+    }
+
     if (calleeName) {
       this.unresolvedReferences.push({
         fromNodeId: callerId,
@@ -4704,6 +4782,42 @@ export class TreeSitterExtractor {
     flush();
   }
 
+  /**
+   * Record a C++ local function-pointer binding (`local = &fn` / `&fn<…>` /
+   * `&ns::fn<…>`) for the CURRENT enclosing symbol, so calls through the local
+   * resolve to the real target (see cppLocalFnPtrs). Only the address-of shape
+   * is accepted — a bare-identifier RHS (`auto x = y;`) is any value copy, and
+   * linking through it would guess.
+   */
+  private recordCppFnPtrBinding(localName: string, value: SyntaxNode | null): void {
+    if (!value || value.type !== 'pointer_expression') return;
+    if (value.child(0)?.type !== '&') return; // `*p` dereference, not address-of
+    const arg = getChildByField(value, 'argument') ?? value.namedChild(0);
+    if (
+      !arg ||
+      (arg.type !== 'identifier' &&
+        arg.type !== 'template_function' &&
+        arg.type !== 'qualified_identifier')
+    ) {
+      return;
+    }
+    const callerId = this.nodeStack[this.nodeStack.length - 1];
+    if (!callerId) return;
+    const target = stripCppTemplateArgs(getNodeText(arg, this.source));
+    if (!target || target === localName) return;
+    let locals = this.cppLocalFnPtrs.get(callerId);
+    if (!locals) {
+      locals = new Map();
+      this.cppLocalFnPtrs.set(callerId, locals);
+    }
+    let targets = locals.get(localName);
+    if (!targets) {
+      targets = new Set();
+      locals.set(localName, targets);
+    }
+    targets.add(target);
+  }
+
   private visitFunctionBody(body: SyntaxNode, _functionId: string): void {
     if (!this.extractor) return;
 
@@ -4763,6 +4877,35 @@ export class TreeSitterExtractor {
         this.extractInstantiation(node);
       }
 
+      // C++ local function-pointer bindings (see cppLocalFnPtrs): record
+      // `auto kernel = &fn<…>;` declarations and `kernel = &other_fn<…>;`
+      // branch reassignments so a call/launch through the local links to the
+      // real target(s). The body walker sees these in source order, and C++
+      // requires declaration-before-use, so the map is always populated before
+      // the call that consumes it.
+      if (this.language === 'cpp' && this.nodeStack.length > 0) {
+        if (nodeType === 'declaration') {
+          for (let i = 0; i < node.namedChildCount; i++) {
+            const child = node.namedChild(i);
+            if (child?.type !== 'init_declarator') continue;
+            const decl = getChildByField(child, 'declarator');
+            if (decl?.type !== 'identifier') continue;
+            this.recordCppFnPtrBinding(
+              getNodeText(decl, this.source),
+              getChildByField(child, 'value')
+            );
+          }
+        } else if (nodeType === 'assignment_expression') {
+          const left = getChildByField(node, 'left');
+          if (left?.type === 'identifier') {
+            this.recordCppFnPtrBinding(
+              getNodeText(left, this.source),
+              getChildByField(node, 'right')
+            );
+          }
+        }
+      }
+
       // Static-member / value-read: `Enum.value`, `Type.CONST`, `Foo::BAR`.
       this.extractStaticMemberRef(node);