fix(gpu): preserve PTX and AMD target facts

Fixes #6469

Fixes #6483

src/gpu/GpuElfKernels.ts

@@ -351,12 +351,14 @@ export const readAmdElfMetadata = async (
   let retainedBytes = 0;
   const observed = (): AmdElfMetadata => {
     const target = targets.size === 1 ? [...targets][0]! : null;
-    const target_key =
+    // LLVM notes use both the empty and explicit `unknown` environment forms.
+    const targetMatch =
       target === null
         ? null
-        : (/^amdgcn-amd-amdhsa--(gfx[a-zA-Z0-9]+(?::[^\s]+)?)$/u.exec(
+        : /^amdgcn-amd-amdhsa-(?:-(gfx[a-zA-Z0-9]+(?::[^\s]+)?)|unknown-(gfx[a-zA-Z0-9]+(?::[^\s]+)?))$/u.exec(
             target,
-          )?.[1] ?? null);
+          );
+    const target_key = targetMatch?.[1] ?? targetMatch?.[2] ?? null;
     return { target, target_key, limitations, notes };
   };
   try {

src/gpu/PtxKernels.test.ts

@@ -66,6 +66,7 @@ it("does not round overflowing register totals or accept truncated ranges", () =
   for (const declaration of [
     "%r<9007199254740991>, %s",
     "%r<4",
+    "%r<2><3>",
     "%r<0>",
     "%r,",
   ]) {
@@ -80,6 +81,15 @@ it("does not round overflowing register totals or accept truncated ranges", () =
     ]);
   }
 });
+it("retains overlong unsupported register type prefixes within the PTX text", () => {
+  const [row] = [
+    ...readPtxKernels(".entry k() { .reg .v2 .u32 .extra %r; ret; }", "object"),
+  ];
+  expect(row!.kernel.registers.virtual).toEqual([
+    { type: ".v2 .u32 .extra", count: null },
+  ]);
+  expect(row!.coverage).toBe("partial");
+});
 it.each([
   "/* unfinished",
   ".entry broken(",

src/gpu/PtxKernels.ts

@@ -75,36 +75,57 @@ const sizeOf = (type: string | undefined): number | null => {
   const match = /^\.(?:b|u|s|f)(8|16|32|64|128)$/u.exec(type ?? "");
   return match === null ? null : Number(match[1]) / 8;
 };
-/** Parse every declared register name/range; unsupported forms poison the type total. */
+/** Stream register names/ranges without retaining the declaration token list. */
 const registerDeclaration = (
-  values: string[],
+  next: () => Token,
 ): { type: string; count: number | null } => {
   const identifier = /^%?[A-Za-z_$][\w$@.]*$/u;
-  const firstName = values.findIndex((value) => identifier.test(value));
-  if (firstName <= 0)
-    return rejectGpu("PTX register declaration has no type or name");
-  const type = values.slice(0, firstName).join(" ");
-  if (firstName !== 1 || !/^\.(?:pred|[busf](?:8|16|32|64|128))$/u.test(type))
-    return { type, count: null };
-  let at = firstName,
-    count = 0;
-  while (at < values.length) {
-    if (!identifier.test(values[at++]!)) return { type, count: null };
-    let declared = 1;
-    if (values[at] === "<") {
-      const range = integer(values[at + 1]);
-      if (range === null || range === 0 || values[at + 2] !== ">")
-        return { type, count: null };
-      declared = range;
-      at += 3;
-    }
-    if (!Number.isSafeInteger(count + declared)) return { type, count: null };
-    count += declared;
-    if (at === values.length) return { type, count };
-    if (values[at++] !== "," || at === values.length)
-      return { type, count: null };
+  let type = "";
+  let t = next();
+  while (!identifier.test(t.value)) {
+    if (t.value === ";" || t.value === "}")
+      rejectGpu("PTX register declaration has no type or name");
+    type += `${type.length === 0 ? "" : " "}${t.value}`;
+    t = next();
+  }
+  if (type.length === 0)
+    rejectGpu("PTX register declaration has no type or name");
+  const supported = /^\.(?:pred|[busf](?:8|16|32|64|128))$/u.test(type);
+  let count = 0;
+  const discard = (token: Token): { type: string; count: null } => {
+    while (token.value !== ";") {
+      if (token.value === "}") rejectGpu("Malformed PTX resource declaration");
+      token = next();
+    }
+    return { type, count: null };
+  };
+  while (true) {
+    let declared = 1;
+    let separator = next();
+    if (separator.value === "<") {
+      const range = next();
+      if (range.value === ";" || range.value === "}") return discard(range);
+      const close = next();
+      if (close.value !== ">") return discard(close);
+      const value = integer(range.value);
+      if (value === null || value === 0) return discard(next());
+      declared = value;
+      separator = next();
+    }
+    if (separator.value === ";") {
+      const total = count + declared;
+      return {
+        type,
+        count: supported && Number.isSafeInteger(total) ? total : null,
+      };
+    }
+    if (separator.value !== ",") return discard(separator);
+    count += declared;
+    if (!Number.isSafeInteger(count) || !supported) return discard(separator);
+    const name = next();
+    if (name.value === ";" || !identifier.test(name.value))
+      return discard(name);
   }
-  return { type, count: null };
 };
 const declaration = (
   source: string[],
@@ -262,15 +283,8 @@ export function* readPtxKernels(
           continue;
         }
         if (t.value !== ".reg" && t.value !== ".shared") continue;
-        const kind = t.value;
-        const declarationTokens: string[] = [];
-        for (t = next(); t.value !== ";"; t = next()) {
-          if (t.value === "}" || declarationTokens.length > 128)
-            rejectGpu("Malformed PTX resource declaration");
-          declarationTokens.push(t.value);
-        }
-        if (kind === ".reg") {
-          const { type, count } = registerDeclaration(declarationTokens);
+        if (t.value === ".reg") {
+          const { type, count } = registerDeclaration(next);
           const previous = virtual.get(type);
           const total =
             previous === null || count === null
@@ -280,20 +294,26 @@ export function* readPtxKernels(
             type,
             total === null || !Number.isSafeInteger(total) ? null : total,
           );
-        } else {
-          const value = declaration(declarationTokens);
-          if (
-            shared === null ||
-            value.size === null ||
-            value.align === null ||
-            value.align === 0
-          )
-            shared = null;
-          else {
-            const total =
-              Math.ceil(shared / value.align) * value.align + value.size;
-            shared = Number.isSafeInteger(total) ? total : null;
-          }
+          continue;
+        }
+        const declarationTokens: string[] = [];
+        for (t = next(); t.value !== ";"; t = next()) {
+          if (t.value === "}" || declarationTokens.length > 128)
+            rejectGpu("Malformed PTX resource declaration");
+          declarationTokens.push(t.value);
+        }
+        const value = declaration(declarationTokens);
+        if (
+          shared === null ||
+          value.size === null ||
+          value.align === null ||
+          value.align === 0
+        )
+          shared = null;
+        else {
+          const total =
+            Math.ceil(shared / value.align) * value.align + value.size;
+          shared = Number.isSafeInteger(total) ? total : null;
         }
       }
     }

tests/boundary/gpu/codeObjects.test.ts

@@ -230,6 +230,46 @@ it("discovers standalone AMD targets before exact selection", async () => {
   expect(skipped.value.kernels).toEqual([]);
   expect(skipped.value.coverage.payloads).toBe("partial");
 });
+it.each(["match", "conflict", "feature-unknown"] as const)(
+  "compares AMDGPU targets with an explicit unknown environment (%s)",
+  async (comparison) => {
+    const payloadTarget =
+      comparison === "feature-unknown"
+        ? "amdgcn-amd-amdhsa-unknown-gfx900:xnack+"
+        : "amdgcn-amd-amdhsa-unknown-gfx900";
+    const inventoryTarget = comparison === "conflict" ? "gfx1100" : "gfx900";
+    const bytes = hipFixture(
+      amdFixture(false, undefined, payloadTarget),
+      `hipv4-amdgcn-amd-amdhsa--${inventoryTarget}`,
+    );
+    const result = await inspect(bytes);
+    if (!result.ok) throw result.error;
+    const row = result.value.objects.find((value) => value.kind === "amdgpu")!;
+    expect(row.payload_target).toBe(payloadTarget);
+    expect(result.value.kernels).toHaveLength(1);
+    if (comparison === "conflict") {
+      expect(row).toMatchObject({
+        target: "gfx1100",
+        target_consistency: "conflict",
+        inspection: "partial",
+      });
+      expect(result.value.coverage.payloads).toBe("partial");
+      expect(result.value.kernels[0]!.limitations.join(" ")).toContain(
+        "Conflicting target declarations",
+      );
+    } else {
+      expect(row).toMatchObject({
+        target: "gfx900",
+        target_consistency: comparison === "match" ? "consistent" : "unknown",
+        inspection: "complete",
+      });
+      expect(result.value.coverage.payloads).toBe("complete");
+      expect(result.value.kernels[0]!.limitations.join(" ")).not.toContain(
+        "Conflicting target declarations",
+      );
+    }
+  },
+);
 it.each([
   "hipv4-amdgcn-amd-amdhsa--gfx900",
   "hipv4-amdgcn-amd-amdhsa-gfx900",
@@ -314,6 +354,65 @@ it("retains unknown PTX virtual counts and partial payload coverage", async () =
     count: null,
   });
 });
+it.each([64, 65])(
+  "counts all %d scalar names in one PTX register list",
+  async (count) => {
+    const names = Array.from(
+      { length: count },
+      (_, index) => `%r${index}`,
+    ).join(", ");
+    const result = await inspect(
+      Buffer.from(
+        `.version 8.8\n.target sm_80\n.address_size 64\n.entry k() { .reg .b32 ${names}; ret; }\n`,
+      ),
+    );
+    if (!result.ok) throw result.error;
+    expect(result.value.objects[0]!.inspection).toBe("complete");
+    expect(result.value.coverage.payloads).toBe("complete");
+    expect(result.value.kernels[0]!.registers.virtual).toEqual([
+      { type: ".b32", count },
+    ]);
+  },
+);
+it("retains PTX kernels before and after a long scalar register list", async () => {
+  const names = Array.from({ length: 65 }, (_, index) => `%r${index}`).join(
+    ", ",
+  );
+  const result = await inspect(
+    Buffer.from(
+      `.version 8.8\n.target sm_80\n.address_size 64\n.entry before() { ret; }\n.entry affected() { .reg .b32 ${names}; ret; }\n.entry after() { ret; }\n`,
+    ),
+  );
+  if (!result.ok) throw result.error;
+  expect(result.value.objects[0]!.inspection).toBe("complete");
+  expect(result.value.coverage.payloads).toBe("complete");
+  expect(result.value.kernels.map((kernel) => kernel.name)).toEqual([
+    "before",
+    "affected",
+    "after",
+  ]);
+  expect(result.value.kernels[1]!.registers.virtual).toEqual([
+    { type: ".b32", count: 65 },
+  ]);
+});
+it("retains later PTX kernels when a virtual register total overflows", async () => {
+  const result = await inspect(
+    Buffer.from(
+      ".version 8.8\n.target sm_80\n.address_size 64\n.entry before() { ret; }\n.entry affected() { .reg .b32 %r<9007199254740991>, %s; ret; }\n.entry after() { ret; }\n",
+    ),
+  );
+  if (!result.ok) throw result.error;
+  expect(result.value.objects[0]!.inspection).toBe("partial");
+  expect(result.value.coverage.payloads).toBe("partial");
+  expect(result.value.kernels.map((kernel) => kernel.name)).toEqual([
+    "before",
+    "affected",
+    "after",
+  ]);
+  expect(result.value.kernels[1]!.registers.virtual).toEqual([
+    { type: ".b32", count: null },
+  ]);
+});
 it.each([
   ".entry k { ret; }",
   ".entry k() { ret; }",

tests/fixtures/gpu/codeObjects.ts

@@ -162,7 +162,11 @@ export const cubinFixture = (
   }
   return result;
 };
-export const amdFixture = (hidden = false, trailingNote?: Buffer): Buffer => {
+export const amdFixture = (
+  hidden = false,
+  trailingNote?: Buffer,
+  target = "amdgcn-amd-amdhsa--gfx900",
+): Buffer => {
   const args = [
     { ".offset": 0, ".size": 8, ".value_kind": "global_buffer" },
     {
@@ -173,7 +177,7 @@ export const amdFixture = (hidden = false, trailingNote?: Buffer): Buffer => {
   ];
   const metadata = Buffer.from(
     encode({
-      "amdhsa.target": "amdgcn-amd-amdhsa--gfx900",
+      "amdhsa.target": target,
       "amdhsa.version": [1, 1],
       "amdhsa.kernels": [
         {