From 42086fdba7920786d5f097df037379fe8b31e537 Mon Sep 17 00:00:00 2001
From: morluto <76467478+morluto@users.noreply.github.com>
Date: Fri, 23 Oct 2026 19:17:42 +0800
Subject: [PATCH] fix(gpu): retain resources across entry pragmas

Fixes #6651
---
 docs/gpu-code-objects.md                   |  4 +-
 src/gpu/PtxKernels.test.ts                 |  9 +++
 src/gpu/PtxKernels.ts                      | 15 +++++
 tests/acceptance/gpuCodeObjectsCli.test.ts | 65 ++++++++++++++++++++++
 4 files changed, 92 insertions(+), 1 deletion(-)

diff --git a/docs/gpu-code-objects.md b/docs/gpu-code-objects.md
index 74226cf5ab..95ea85d67e 100644
--- a/docs/gpu-code-objects.md
+++ b/docs/gpu-code-objects.md
@@ -82,7 +82,9 @@ Kernel representations have different meanings:
   supported static declarations inside that entry; they do not establish module
   shared memory, dynamic allocation or physical device resources.
   PTX 1.0–1.3 declares entry parameters in the body; PTX 1.4 and later uses the
-  entry parameter list. A `.target compute_xx` spelling is retained as written;
+  entry parameter list. Entry-scoped `.pragma` statements before the body do
+  not end the declaration or suppress its body resources. A `.target
+  compute_xx` spelling is retained as written;
   NVIDIA PTX defines it as a synonym for `.target sm_xx` when comparing its
   target identity with the fatbin declaration. `.callprototype` formal
   parameters are excluded from entry virtual-register counts. Parameter names
diff --git a/src/gpu/PtxKernels.test.ts b/src/gpu/PtxKernels.test.ts
index 83185afe8b..08cc9190c8 100644
--- a/src/gpu/PtxKernels.test.ts
+++ b/src/gpu/PtxKernels.test.ts
@@ -47,6 +47,14 @@ it("retains a parameter name when a qualifier is unknown", () => {
   ]);
   expect(row!.coverage).toBe("partial");
 });
+it("keeps external declarations bodyless and resource values unknown", () => {
+  const [row] = [
+    ...readPtxKernels(".extern .entry k(.param .u32 input);", "object"),
+  ];
+  expect(row!.kernel.parameter_count).toBe(1);
+  expect(row!.kernel.shared_bytes).toBeNull();
+  expect(row!.coverage).toBe("partial");
+});
 it.each([".ptr .global .align 4", ".ptr .align 1", ".ptr.global.align 4"])(
   "parses PTX pointer qualifiers %s without treating pointee alignment as parameter alignment",
   (qualifiers) => {
@@ -153,6 +161,7 @@ it.each([
   ".entry k",
   ".entry k ) { ret; }",
   ".entry k() {",
+  ".entry k() .pragma; { ret; }",
   ".entry k() {.reg .u32 %x}",
   ".entry k() {.reg %x;}",
 ])("rejects incomplete PTX %s", (text) => {
diff --git a/src/gpu/PtxKernels.ts b/src/gpu/PtxKernels.ts
index 8d8c8384a1..060f401a76 100644
--- a/src/gpu/PtxKernels.ts
+++ b/src/gpu/PtxKernels.ts
@@ -345,6 +345,21 @@ export function* readPtxKernels(
     let headerTokens = 0;
     while (t.value !== "{" && t.value !== ";") {
       if (++headerTokens > 65536) rejectGpu("PTX entry header is too large");
+      if (t.value === ".pragma") {
+        t = next();
+        if (t.value !== '""') rejectGpu("Malformed PTX entry pragma");
+        while (true) {
+          if (++headerTokens > 65536)
+            rejectGpu("PTX entry header is too large");
+          t = next();
+          if (t.value === ";") break;
+          if (t.value !== ",") rejectGpu("Malformed PTX entry pragma");
+          t = next();
+          if (t.value !== '""') rejectGpu("Malformed PTX entry pragma");
+        }
+        t = next();
+        continue;
+      }
       t = next();
     }
     let end = t.end,
diff --git a/tests/acceptance/gpuCodeObjectsCli.test.ts b/tests/acceptance/gpuCodeObjectsCli.test.ts
index a0e3f02a9b..94084bd579 100644
--- a/tests/acceptance/gpuCodeObjectsCli.test.ts
+++ b/tests/acceptance/gpuCodeObjectsCli.test.ts
@@ -460,6 +460,71 @@ cliTest(
   },
 );
 
+cliTest(
+  "keeps entry-scoped PTX pragmas outside the body boundary",
+  async ({ cli }) => {
+    const root = await createTestTempDirectory("rea-ptx-entry-pragma-");
+    const preamble = ".version 8.8\n.target sm_80\n.address_size 64\n",
+      body = "{ .reg .b32 %r<3>; .shared .align 4 .b8 scratch[16]; ret; }",
+      declarations = [
+        ["plain", `.visible .entry k(.param .u32 input) ${body}`],
+        [
+          "entry-pragma",
+          `.visible .entry k(.param .u32 input) .pragma "nounroll"; ${body}`,
+        ],
+        [
+          "body-pragma",
+          `.visible .entry k(.param .u32 input) { .pragma "nounroll"; .reg .b32 %r<3>; .shared .align 4 .b8 scratch[16]; ret; }`,
+        ],
+      ] as const;
+    const reports = new Map();
+    for (const [name, declaration] of declarations) {
+      const text = `${preamble}${declaration}\n`,
+        path = join(root, `${name}.ptx`);
+      await writeFile(path, text);
+      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);
+      expect(report.ptx_artifacts[0]!.sha256).toBe(
+        createHash("sha256").update(text).digest("hex"),
+      );
+      reports.set(name, report);
+    }
+    const expected = {
+      name: "k",
+      parameter_count: 1,
+      parameter_bytes: 4,
+      registers: { virtual: [{ type: ".b32", count: 3 }] },
+      shared_bytes: 16,
+    };
+    for (const [name, report] of reports) {
+      expect(report.kernels[0]).toMatchObject(expected);
+      expect(report.coverage.payloads).toBe("complete");
+      const declaration = declarations.find(
+        ([caseName]) => caseName === name,
+      )![1];
+      expect(report.kernels[0]!.location.size).toBe(
+        Buffer.byteLength(declaration.slice(declaration.indexOf(".entry"))),
+      );
+    }
+  },
+);
+
 cliTest(
   "accepts the PTX compute_xx .target synonym verbatim",
   async ({ cli }) => {