fix(gpu): preserve PTX parameter and target facts

Fixes #6664
Fixes #6675
Fixes #6680
Fixes #6695

docs/gpu-code-objects.md

@@ -91,6 +91,12 @@ Kernel representations have different meanings:
   remain available when pointer qualifiers are understood or an unsupported
   qualifier prevents layout decoding. A pointer's `.align` describes the
   pointed-to memory and does not change parameter layout alignment.
+  Parameter names and following declarations are retained when an array
+  extent uses a parenthesized expression or when dimensions are
+  multidimensional, even though those forms do not currently yield a layout.
+  PTX target selection uses the highest-numbered module-scope `.target`
+  directive, including directives after entry declarations; unsupported or
+  conflicting target profiles leave the effective target unknown.
   Scalar register declarations count each bare or percent-prefixed name and
   each parameterized range. Unsupported declaration forms or overflowing totals
   retain an unknown count and partial decoding coverage. Unsupported parameter

src/gpu/GpuInventory.ts

@@ -45,11 +45,11 @@ import {
   readAmdElfMetadata,
   type AmdElfMetadata,
 } from "./GpuElfKernels.js";
-import { readPtxKernels, readPtxTarget } from "./PtxKernels.js";
-
-/** PTX calls compute_XX a synonym for the corresponding sm_XX target. */
-const targetIdentity = (target: string): string =>
-  target.replace(/^compute_(\d+[a-z]?)$/u, "sm_$1");
+import {
+  ptxTargetIdentity,
+  readPtxKernels,
+  readPtxTarget,
+} from "./PtxKernels.js";
 
 /** File declarations only. Every child range stays in its named parent coordinate space. */
 export class GpuInventory {
@@ -234,21 +234,28 @@ export class GpuInventory {
     }
     // PTX is positively identified by its directives, never by extension or arbitrary embedded text.
     let probe = await source.read(0, Math.min(source.size, 64 * 1024));
-    let target = readPtxTarget(probe.toString("utf8"));
-    if (
-      target === null &&
-      probe.length < Math.min(source.size, GPU_REPORT_BYTES)
-    ) {
+    let targetFacts = readPtxTarget(probe.toString("utf8"));
+    if (probe.length < Math.min(source.size, GPU_REPORT_BYTES)) {
       probe = await source.read(0, Math.min(source.size, GPU_REPORT_BYTES));
-      target = readPtxTarget(probe.toString("utf8"));
+      targetFacts = readPtxTarget(probe.toString("utf8"));
     }
-    if (target !== null) {
+    if (targetFacts.recognized) {
+      const target = probe.length === source.size ? targetFacts.target : null;
+      const limitation =
+        probe.length < source.size
+          ? "PTX target scan reached the 16 MiB text budget before all module directives were examined."
+          : targetFacts.limitation;
       const row = gpuObject(source, {
         offset: 0,
         size: source.size,
         kind: "ptx",
       });
       row.target = target;
+      if (limitation !== null) {
+        row.inspection = "partial";
+        row.limitations.push(limitation);
+        this.discovery = "partial";
+      }
       this.#add(row);
       await this.#payload(
         probe.length === source.size
@@ -357,9 +364,12 @@ export class GpuInventory {
         }
         const target = readPtxTarget(text);
         const conflict = this.#target(row, {
-          value: target,
-          key: target,
+          value: target.target,
+          key: target.target,
           scope: "exact",
+          ...(target.limitation === null
+            ? {}
+            : { limitation: target.limitation }),
         });
         if (conflict !== undefined) coverage = "partial";
         const parsed = readPtxKernels(text, row.id);
@@ -407,7 +417,7 @@ export class GpuInventory {
         const inner = /^sm_([0-9]+)$/u.exec(payload.key)?.[1];
         if (base !== undefined && inner !== undefined)
           state = base === inner ? "base-only" : "conflict";
-      } else if (targetIdentity(outer) === targetIdentity(payload.key))
+      } else if (ptxTargetIdentity(outer) === ptxTargetIdentity(payload.key))
         state = "consistent";
       else if (
         row.kind === "amdgpu" &&
@@ -432,7 +442,8 @@ export class GpuInventory {
       state === "conflict"
         ? `Conflicting target declarations: inventory ${outer}; payload ${payload.value}. Target agreement is not established.`
         : payload.limitation!;
-    this.#objectLimitations(row, [message]);
+    if (!row.limitations.includes(message))
+      this.#objectLimitations(row, [message]);
     return message;
   }
   async #amdPayload(

src/gpu/PtxKernels.test.ts

@@ -1,11 +1,16 @@
 import { expect, it } from "vitest";
 import { readPtxKernels, readPtxTarget } from "./PtxKernels.js";
 it("recognizes directives outside comments and strings and preserves UTF-8 byte coordinates", () => {
-  expect(readPtxTarget('".version 8.8\n.target sm_80\n"')).toBeNull();
+  expect(readPtxTarget('".version 8.8\n.target sm_80\n"').recognized).toBe(
+    false,
+  );
   const prefix = "// 雪 .entry fake() {}\n.version 8.8\n.target sm_80\n";
   const text =
     prefix + ".visible .entry real(.param .u64 x) { .reg .b32 %r<3>; ret; }";
-  expect(readPtxTarget(text)).toBe("sm_80");
+  expect(readPtxTarget(text)).toMatchObject({
+    recognized: true,
+    target: "sm_80",
+  });
   const rows = [...readPtxKernels(text, "object")];
   expect(rows).toHaveLength(1);
   expect(rows[0]!.coverage).toBe("complete");
@@ -16,6 +21,13 @@ it("recognizes directives outside comments and strings and preserves UTF-8 byte
     { type: ".b32", count: 3 },
   ]);
 });
+it("uses the highest module-scope target directive, including after entries", () => {
+  expect(
+    readPtxTarget(
+      ".version 8.8\n.target sm_80\n.entry first() { ret; }\n.target sm_90\n",
+    ),
+  ).toMatchObject({ recognized: true, target: "sm_90", limitation: null });
+});
 it("derives vector/array and pointer parameters while preserving their names", () => {
   const rows = [
     ...readPtxKernels(
@@ -35,6 +47,50 @@ it("derives vector/array and pointer parameters while preserving their names", (
   expect(rows[0]!.coverage).toBe("complete");
   expect(rows[0]!.kernel.parameter_count).toBe(3);
 });
+it.each([
+  [
+    "plain array control",
+    ".param .b8 payload[6], .param .u32 count",
+    ["payload", "count"],
+    [6, 4],
+    [0, 6],
+    "complete",
+  ],
+  [
+    "parenthesized extent",
+    ".param .b8 data[(4)], .param .u32 count",
+    ["data", "count"],
+    [null, 4],
+    [null, null],
+    "partial",
+  ],
+  [
+    "expression extent",
+    ".param .b8 payload[2+4], .param .u32 count",
+    ["payload", "count"],
+    [null, 4],
+    [null, null],
+    "partial",
+  ],
+  [
+    "multidimensional extent",
+    ".param .b8 payload[2][3], .param .u32 count",
+    ["payload", "count"],
+    [null, 4],
+    [null, null],
+    "partial",
+  ],
+])(
+  "preserves parameter names and trailing declarations for %s",
+  (_kind, source, names, sizes, offsets, coverage) => {
+    const [row] = [...readPtxKernels(`.entry k(${source}) { ret; }`, "object")];
+    expect(row!.kernel.parameter_count).toBe(2);
+    expect(row!.kernel.parameters.map(({ name }) => name)).toEqual(names);
+    expect(row!.kernel.parameters.map(({ size }) => size)).toEqual(sizes);
+    expect(row!.kernel.parameters.map(({ offset }) => offset)).toEqual(offsets);
+    expect(row!.coverage).toBe(coverage);
+  },
+);
 it("retains a parameter name when a qualifier is unknown", () => {
   const [row] = [
     ...readPtxKernels(
@@ -169,6 +225,6 @@ it.each([
 });
 
 it("does not label unrelated malformed quoted text as a recognized PTX artifact", () => {
-  expect(readPtxTarget('unrelated "unterminated')).toBeNull();
-  expect(readPtxTarget("/* unrelated unterminated")).toBeNull();
+  expect(readPtxTarget('unrelated "unterminated').recognized).toBe(false);
+  expect(readPtxTarget("/* unrelated unterminated").recognized).toBe(false);
 });

src/gpu/PtxKernels.ts

@@ -190,30 +190,30 @@ const declaration = (
         layoutKnown = false;
     } else layoutKnown = false;
   }
-  let nameIndex = values.length - 1;
-  let arrayStart: number | null = null;
-  if (values[nameIndex] === "]") {
-    let depth = 0;
-    for (let index = nameIndex; index >= at; index--) {
-      if (values[index] === "]") depth++;
-      else if (values[index] === "[" && --depth === 0) {
-        arrayStart = index;
-        break;
-      }
+  const name = identifier.test(values[at] ?? "") ? values[at]! : null;
+  if (name === null) layoutKnown = false;
+  let cursor = at + (name === null ? 0 : 1);
+  const dimensions: string[][] = [];
+  while (values[cursor] === "[") {
+    const start = ++cursor;
+    let parentheses = 0;
+    while (cursor < values.length) {
+      const value = values[cursor]!;
+      if (value === "(") parentheses++;
+      else if (value === ")") {
+        if (parentheses === 0) break;
+        parentheses--;
+      } else if (value === "]" && parentheses === 0) break;
+      cursor++;
     }
-    if (arrayStart !== null) nameIndex = arrayStart - 1;
-  } else {
-    const opening = values.indexOf("[", at);
-    if (opening >= 0) {
-      arrayStart = opening;
-      nameIndex = opening - 1;
+    if (values[cursor] !== "]" || parentheses !== 0) {
       layoutKnown = false;
+      break;
     }
+    dimensions.push(values.slice(start, cursor));
+    cursor++;
   }
-  const name = identifier.test(values[nameIndex] ?? "")
-    ? values[nameIndex]!
-    : null;
-  if (name === null || nameIndex !== at) layoutKnown = false;
+  if (cursor !== values.length || dimensions.length > 1) layoutKnown = false;
   let size = unit === null ? null : unit * vector;
   const type =
     scalar === undefined
@@ -223,20 +223,14 @@ const declaration = (
         : `.v${vector} ${scalar}`;
   const parameterAlignment = explicit ? alignment : size;
   const align = parameterAlignment;
-  if (arrayStart !== null) {
-    const count = integer(values[arrayStart + 1]);
-    if (
-      arrayStart !== nameIndex + 1 ||
-      arrayStart + 2 !== values.length - 1 ||
-      count === null ||
-      size === null ||
-      !Number.isSafeInteger(count * size)
-    )
+  if (dimensions.length === 1) {
+    const dimension = dimensions[0]!,
+      count = dimension.length === 1 ? integer(dimension[0]) : null;
+    if (count === null || size === null || !Number.isSafeInteger(count * size))
       layoutKnown = false;
     else size *= count;
   }
   if (
-    (arrayStart === null && values.length !== nameIndex + 1) ||
     name === null ||
     align === null ||
     align <= 0 ||
@@ -326,8 +320,21 @@ export function* readPtxKernels(
     };
     let t = next();
     if (t.value === "(") {
-      for (t = next(); t.value !== ")"; t = next()) {
-        if (t.value === ",") flush();
+      let parentheses = 0,
+        brackets = 0;
+      while (true) {
+        t = next();
+        if (t.value === "(") parentheses++;
+        else if (t.value === ")") {
+          if (parentheses > 0) parentheses--;
+          else if (brackets === 0) break;
+          else rejectGpu("Malformed PTX parameter dimension");
+        } else if (t.value === "[") brackets++;
+        else if (t.value === "]") {
+          if (brackets === 0) rejectGpu("Malformed PTX parameter dimension");
+          brackets--;
+        }
+        if (t.value === "," && parentheses === 0 && brackets === 0) flush();
         else part.push(t.value);
         if (part.length > 128)
           throw new BinaryResourceFailure(
@@ -515,36 +522,90 @@ export function* readPtxKernels(
   if (outerDepth !== 0) rejectGpu("Unbalanced PTX body");
 }
 
-/** Positive PTX header recognition uses the same lexer as kernel extraction. */
-export const readPtxTarget = (text: string): string | null => {
+export interface PtxTargetFacts {
+  readonly recognized: boolean;
+  readonly target: string | null;
+  readonly limitation: string | null;
+}
+
+/** PTX's compute spelling is an alias for the corresponding baseline SM target. */
+export const ptxTargetIdentity = (target: string): string =>
+  target.replace(/^compute_(\d+[a-z]?)$/u, "sm_$1");
+
+/** Resolve all module-scope targets through the highest numbered architecture. */
+export const readPtxTarget = (text: string): PtxTargetFacts => {
   const stream = tokens(text);
   let version = false,
-    target: string | null = null;
+    depth = 0,
+    target: string | null = null,
+    targetIdentity: string | null = null,
+    targetArchitecture = -1,
+    ambiguous = false,
+    unsupported = false,
+    malformed = false;
   try {
     for (const token of stream) {
-      if (
-        token.value === ".entry" ||
-        token.value === ".func" ||
-        token.value === "{"
-      )
-        break;
+      if (token.value === "{") {
+        depth++;
+        continue;
+      }
+      if (token.value === "}") {
+        if (depth > 0) depth--;
+        continue;
+      }
+      if (depth !== 0) continue;
       if (token.value === ".version") {
         const value = stream.next();
         version = !value.done && ptxVersion(value.value.value) !== null;
+        continue;
       }
-      if (token.value === ".target") {
-        const value = stream.next();
-        target =
-          !value.done &&
-          /^(?:sm_[0-9]+(?:a|f)?|compute_[0-9]+)$/u.test(value.value.value)
-            ? value.value.value
-            : null;
+      if (token.value !== ".target" || !version) continue;
+      const value = stream.next();
+      if (value.done) {
+        unsupported = true;
+        continue;
+      }
+      const spelling = value.value.value,
+        match = /^(?:sm_|compute_)([0-9]+)(?:a|f)?$/u.exec(spelling);
+      if (match === null) {
+        unsupported = true;
+        continue;
+      }
+      const architecture = Number(match[1]);
+      if (!Number.isSafeInteger(architecture)) {
+        unsupported = true;
+        continue;
+      }
+      const identity = ptxTargetIdentity(spelling);
+      if (architecture > targetArchitecture) {
+        target = spelling;
+        targetIdentity = identity;
+        targetArchitecture = architecture;
+        ambiguous = false;
+      } else if (
+        architecture === targetArchitecture &&
+        targetIdentity !== identity
+      ) {
+        ambiguous = true;
       }
-      if (version && target !== null) return target;
     }
   } catch (cause) {
     if (!(cause instanceof BinaryFormatFailure)) throw cause;
-    // A malformed quote/comment before both directives does not establish PTX identity.
+    malformed = true;
   }
-  return null;
+  const recognized = version && target !== null;
+  const limitation = !recognized
+    ? null
+    : unsupported
+      ? "Unsupported PTX target directives leave the effective target unknown."
+      : ambiguous
+        ? "Same-generation PTX target profiles cannot be represented by one effective target."
+        : malformed
+          ? "Malformed PTX text prevents complete target discovery."
+          : null;
+  return {
+    recognized,
+    target: limitation === null ? target : null,
+    limitation,
+  };
 };

tests/boundary/gpu/codeObjects.test.ts

@@ -496,6 +496,17 @@ it.each(["whitespace", "comment"] as const)(
     );
   },
 );
+it("uses a later module target when selecting standalone PTX", async () => {
+  const bytes = Buffer.from(
+    ".version 8.8\n.target sm_80\n.entry first() { ret; }\n.target sm_90\n",
+  );
+  const inventory = await inspect(bytes, { analysis: "inventory" });
+  const selected = await inspect(bytes, { targets: ["sm_90"] });
+  if (!inventory.ok) throw inventory.error;
+  if (!selected.ok) throw selected.error;
+  expect(inventory.value.objects[0]!.target).toBe("sm_90");
+  expect(selected.value.kernels.map(({ name }) => name)).toEqual(["first"]);
+});
 it.each(["truncated", "invalid", "extension"] as const)(
   "retains earlier AMD kernel declarations after a later %s note",
   async (kind) => {