diff --git a/src/gpu/PtxKernels.test.ts b/src/gpu/PtxKernels.test.ts index cbb380533d..379c8ccb6c 100644 --- a/src/gpu/PtxKernels.test.ts +++ b/src/gpu/PtxKernels.test.ts @@ -183,6 +183,20 @@ it("sums bare names and every register range while preserving unsupported totals ]); expect(row!.kernel.limitations.join(" ")).toContain("count unknown"); }); +it("counts numeric percent-prefixed PTX register names", () => { + const [row] = [ + ...readPtxKernels( + ".entry kernel() { .reg .b32 %0; mov.u32 %0, 0; ret; }", + "object", + ), + ]; + expect(row!.kernel.registers.virtual).toEqual([{ type: ".b32", count: 1 }]); +}); +it.each(["0", "%"])("rejects malformed PTX register identifier %s", (name) => { + expect(() => [ + ...readPtxKernels(`.entry kernel() { .reg .b32 ${name}; ret; }`, "object"), + ]).toThrow(); +}); it("does not round overflowing register totals or accept truncated ranges", () => { for (const declaration of [ "%r<9007199254740991>, %s", diff --git a/src/gpu/PtxKernels.ts b/src/gpu/PtxKernels.ts index 6a0363437f..b1b1c7ccb6 100644 --- a/src/gpu/PtxKernels.ts +++ b/src/gpu/PtxKernels.ts @@ -85,7 +85,7 @@ 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; }; -const identifier = /^%?[A-Za-z_$][\w$@.]*$/u; +const identifier = /^(?:%?[A-Za-z_$][\w$@.]*|%[0-9][A-Za-z0-9_$]*)$/u; /** Stream register names/ranges without retaining the declaration token list. */ const registerDeclaration = ( next: () => Token, diff --git a/tests/acceptance/gpuCodeObjectsCli.test.ts b/tests/acceptance/gpuCodeObjectsCli.test.ts index 94084bd579..0c219b7dc0 100644 --- a/tests/acceptance/gpuCodeObjectsCli.test.ts +++ b/tests/acceptance/gpuCodeObjectsCli.test.ts @@ -412,6 +412,44 @@ cliTest( }, ); +cliTest( + "counts numeric percent-prefixed PTX register declarations", + async ({ cli }) => { + const root = await createTestTempDirectory("rea-ptx-numeric-register-"); + for (const register of ["%0", "%r0"] as const) { + const text = `.version 7.0\n.target sm_80\n.address_size 64\n.visible .entry kernel() {\n.reg .b32 ${register};\nmov.u32 ${register}, 0;\nret;\n}\n`; + const path = join(root, `${register.slice(1)}.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 report = gpuCodeObjectsSchema.parse( + toolContract("inspect_gpu_code_objects").outputSchema.parse( + response.json, + ).normalized_result, + ); + expect(report.kernels).toMatchObject([ + { + name: "kernel", + registers: { virtual: [{ type: ".b32", count: 1 }] }, + }, + ]); + expect(report.ptx_artifacts[0]!.sha256).toBe( + createHash("sha256").update(text).digest("hex"), + ); + } + }, +); + cliTest( "preserves PTX target spelling and percent-prefixed entry names", async ({ cli }) => {