From aa52ff6f4bffb47c8a3dcf60e1800e7dce5fad02 Mon Sep 17 00:00:00 2001 From: morluto <76467478+morluto@users.noreply.github.com> Date: Fri, 23 Oct 2026 21:01:48 +0800 Subject: [PATCH] fix(gpu): preserve PTX parameter and target facts Fixes #6664 Fixes #6675 Fixes #6680 Fixes #6695 --- docs/gpu-code-objects.md | 6 + src/gpu/GpuInventory.ts | 43 ++++--- src/gpu/PtxKernels.test.ts | 64 +++++++++- src/gpu/PtxKernels.ts | 163 +++++++++++++++++-------- tests/boundary/gpu/codeObjects.test.ts | 11 ++ 5 files changed, 216 insertions(+), 71 deletions(-) diff --git a/docs/gpu-code-objects.md b/docs/gpu-code-objects.md index 95ea85d67e..512b4a264a 100644 --- a/docs/gpu-code-objects.md +++ b/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 diff --git a/src/gpu/GpuInventory.ts b/src/gpu/GpuInventory.ts index eb5c8d2de3..63946c5818 100644 --- a/src/gpu/GpuInventory.ts +++ b/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( diff --git a/src/gpu/PtxKernels.test.ts b/src/gpu/PtxKernels.test.ts index 08cc9190c8..01d1ba3ca3 100644 --- a/src/gpu/PtxKernels.test.ts +++ b/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); }); diff --git a/src/gpu/PtxKernels.ts b/src/gpu/PtxKernels.ts index 060f401a76..89b4846adf 100644 --- a/src/gpu/PtxKernels.ts +++ b/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, + }; }; diff --git a/tests/boundary/gpu/codeObjects.test.ts b/tests/boundary/gpu/codeObjects.test.ts index cbe4ced8bd..8536a7574f 100644 --- a/tests/boundary/gpu/codeObjects.test.ts +++ b/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) => {