fix(gpu): resolve extended CUBIN symbol sections

Fixes #6600

src/gpu/GpuElfKernels.ts

@@ -266,6 +266,8 @@ const cudaKernels = async (
       rejectGpu("CUBIN symbol string table is invalid");
     const names = await sectionBytes(source, namesSection!, index.budget);
     const bytes = await sectionBytes(source, section, index.budget);
+    const symbolCount = bytes.length / width;
+    let extendedIndices: Buffer | null = null;
     for (let offset = 0; offset < bytes.length; offset += width) {
       index.budget.visit();
       const symbol = readElfSymbol(
@@ -275,7 +277,31 @@ const cudaKernels = async (
       );
       // STO_CUDA_ENTRY distinguishes launchable kernels from ordinary device functions.
       if (symbol.type !== 2 || (symbol.other & 0x10) === 0) continue;
-      const text = directory.sections[symbol.section_index];
+      let sectionIndex = symbol.section_index;
+      if (sectionIndex === 0xffff) {
+        if (extendedIndices === null) {
+          const matches = directory.sections.filter(
+            (row) => row.type === 18 && row.link === section.index,
+          );
+          if (matches.length !== 1)
+            rejectGpu(
+              "CUBIN extended symbol section-index table is unavailable",
+            );
+          const extended = matches[0]!;
+          if (
+            extended.entry_size !== 4n ||
+            extended.size !== BigInt(symbolCount) * 4n
+          )
+            rejectGpu(
+              "CUBIN extended symbol section-index table has an invalid extent",
+            );
+          extendedIndices = await sectionBytes(source, extended, index.budget);
+        }
+        sectionIndex = new BinaryReader(extendedIndices, directory.little).u32(
+          (offset / width) * 4,
+        );
+      }
+      const text = directory.sections[sectionIndex];
       if (text === undefined || text.type !== 1 || (text.flags & 4n) === 0n)
         rejectGpu("CUBIN kernel symbol has no executable section");
       const kernel = blankKernel(

tests/acceptance/gpuCodeObjectsCli.test.ts

@@ -63,6 +63,84 @@ cliTest(
   },
 );
 
+cliTest(
+  "resolves CUBIN SHN_XINDEX symbols through their ELF extension table",
+  async ({ cli }) => {
+    const root = await createTestTempDirectory("rea-cubin-shndx-cli-");
+    const bytes = cubinFixture({ extendedSectionIndex: true });
+    const path = join(root, "extended.cubin");
+    await writeFile(path, bytes);
+    const response = await cli.run({
+      arguments: [
+        "inspect-gpu-code-objects",
+        path,
+        "--analysis",
+        "kernels",
+        "--detail",
+        "complete",
+        "--json",
+      ],
+    });
+    expect(response.exitCode, response.stderr).toBe(0);
+    const evidence = parseEvidence(
+      toolContract("inspect_gpu_code_objects").outputSchema.parse(
+        response.json,
+      ),
+    );
+    const report = gpuCodeObjectsSchema.parse(evidence.normalized_result);
+    const sha256 = createHash("sha256").update(bytes).digest("hex");
+    expect(report.artifact.sha256).toBe(sha256);
+    expect(report.objects[0]!.sha256).toBe(sha256);
+    expect(report.kernels).toHaveLength(1);
+    expect(report.kernels[0]).toMatchObject({
+      name: "add_one",
+      parameter_count: 2,
+      parameter_bytes: 12,
+      shared_bytes: 1056,
+      sources: {
+        parameters: expect.objectContaining({ section_index: 65537 }),
+        shared: expect.objectContaining({ section_index: 65538 }),
+      },
+    });
+
+    const malformedBytes = Buffer.from(bytes);
+    const sectionTable = Number(malformedBytes.readBigUInt64LE(40));
+    const sectionCount = Number(
+      malformedBytes.readBigUInt64LE(sectionTable + 32),
+    );
+    let indexTable: number | undefined;
+    for (let index = 0; index < sectionCount; index++)
+      if (malformedBytes.readUInt32LE(sectionTable + index * 64 + 4) === 18) {
+        indexTable = index;
+        break;
+      }
+    expect(indexTable).toBeDefined();
+    malformedBytes.writeBigUInt64LE(4n, sectionTable + indexTable! * 64 + 32);
+    const malformedPath = join(root, "short-shndx.cubin");
+    await writeFile(malformedPath, malformedBytes);
+    const malformed = await cli.run({
+      arguments: [
+        "inspect-gpu-code-objects",
+        malformedPath,
+        "--analysis",
+        "kernels",
+        "--json",
+      ],
+    });
+    expect(malformed.exitCode, malformed.stderr).toBe(0);
+    const malformedEvidence = parseEvidence(
+      toolContract("inspect_gpu_code_objects").outputSchema.parse(
+        malformed.json,
+      ),
+    );
+    const malformedReport = gpuCodeObjectsSchema.parse(
+      malformedEvidence.normalized_result,
+    );
+    expect(malformedReport.objects[0]!.inspection).toBe("malformed");
+    expect(malformedReport.coverage.payloads).toBe("partial");
+  },
+);
+
 cliTest(
   "preserves PTX target spelling and percent-prefixed entry names",
   async ({ cli }) => {

tests/fixtures/gpu/codeObjects.ts

@@ -36,6 +36,7 @@ export const elfFixture = (
     { name: ".shstrtab", type: 3, bytes: names },
     ...sections,
   ];
+  const extendedSectionCount = rows.length >= 0xff00;
   let position = 64;
   const entries = rows.map((row) => {
     position = Math.ceil(position / 8) * 8;
@@ -62,8 +63,10 @@ export const elfFixture = (
   result.writeUInt32LE(machine === 190 ? 0x500550 : 0x12c, 48);
   result.writeUInt16LE(64, 52);
   result.writeUInt16LE(64, 58);
-  result.writeUInt16LE(rows.length, 60);
+  result.writeUInt16LE(extendedSectionCount ? 0 : rows.length, 60);
   result.writeUInt16LE(1, 62);
+  if (extendedSectionCount)
+    result.writeBigUInt64LE(BigInt(rows.length), table + 32);
   for (let index = 1; index < entries.length; index++) {
     const { row, offset } = entries[index]!,
       h = table + index * 64;
@@ -90,6 +93,7 @@ export const cubinFixture = (
     architecture?: number;
     osAbi?: number;
     flags?: number;
+    extendedSectionIndex?: boolean;
   } = {},
 ): Buffer => {
   const name = options.name ?? "add_one";
@@ -97,7 +101,7 @@ export const cubinFixture = (
   symbols.writeUInt32LE(1, 24);
   symbols[28] = 0x12;
   symbols[29] = 0x10;
-  symbols.writeUInt16LE(4, 30);
+  symbols.writeUInt16LE(options.extendedSectionIndex ? 0xffff : 4, 30);
   symbols.writeBigUInt64LE(16n, 40);
   const record = Buffer.from("042f0800010000000a000000", "hex");
   record.writeUInt32LE(options.globalRegisters ?? 10, 8);
@@ -108,11 +112,20 @@ export const cubinFixture = (
     "03190c0004170c00000000000100080000f0110004170c00000000000000000000f02100",
     "hex",
   );
+  const paddingSections = options.extendedSectionIndex
+    ? Array.from({ length: 65531 }, (_, index) => ({
+        name: `.padding.${index}`,
+      }))
+    : [];
+  const textSectionIndex = options.extendedSectionIndex ? 65535 : 4;
+  const extendedIndices = Buffer.alloc(8);
+  extendedIndices.writeUInt32LE(textSectionIndex, 4);
   const result = elfFixture(
     190,
     [
       { name: ".strtab", type: 3, bytes: Buffer.from(`\0${name}\0`) },
       { name: ".symtab", type: 2, link: 2, entrySize: 24, bytes: symbols },
+      ...paddingSections,
       {
         name: `.text.${name}`,
         flags: 0x100006n,
@@ -124,15 +137,26 @@ export const cubinFixture = (
         name: `.nv.info.${name}`,
         type: 0x70000000,
         link: 3,
-        info: 4,
+        info: textSectionIndex,
         bytes: info,
       },
       {
         name: `.nv.shared.${name}`,
         type: 8,
-        info: 4,
+        info: textSectionIndex,
         size: options.shared ?? 1056,
       },
+      ...(options.extendedSectionIndex
+        ? [
+            {
+              name: ".symtab_shndx",
+              type: 18,
+              link: 3,
+              entrySize: 4,
+              bytes: extendedIndices,
+            },
+          ]
+        : []),
       ...Array.from({ length: (options.globalAliases ?? 1) - 1 }, () => ({
         name: ".nv.info",
         type: 0x70000000,
@@ -148,12 +172,21 @@ export const cubinFixture = (
       48,
     );
   const table = Number(result.readBigUInt64LE(40));
+  const sectionCount =
+    result.readUInt16LE(60) === 0
+      ? Number(result.readBigUInt64LE(table + 32))
+      : result.readUInt16LE(60);
   if (options.osAbi !== undefined) {
     result[7] = options.osAbi;
     result.writeUInt32LE(1, 20);
   }
   if (options.flags !== undefined) result.writeUInt32LE(options.flags, 48);
-  for (let index = 8; index < result.readUInt16LE(60); index++) {
+  for (
+    let index =
+      8 + paddingSections.length + (options.extendedSectionIndex ? 1 : 0);
+    index < sectionCount;
+    index++
+  ) {
     result.writeBigUInt64LE(
       result.readBigUInt64LE(table + 5 * 64 + 24),
       table + index * 64 + 24,