diff --git a/src/gpu/GpuElfKernels.ts b/src/gpu/GpuElfKernels.ts index 6b1caacbcb..0ce4bec885 100644 --- a/src/gpu/GpuElfKernels.ts +++ b/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 { diff --git a/src/gpu/PtxKernels.test.ts b/src/gpu/PtxKernels.test.ts index 48dd48c324..548d23275a 100644 --- a/src/gpu/PtxKernels.test.ts +++ b/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(", diff --git a/src/gpu/PtxKernels.ts b/src/gpu/PtxKernels.ts index d15ac6bd60..0acac9a033 100644 --- a/src/gpu/PtxKernels.ts +++ b/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; } } } diff --git a/tests/boundary/gpu/codeObjects.test.ts b/tests/boundary/gpu/codeObjects.test.ts index b5f6d067fd..3b5c01a208 100644 --- a/tests/boundary/gpu/codeObjects.test.ts +++ b/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; }", diff --git a/tests/fixtures/gpu/codeObjects.ts b/tests/fixtures/gpu/codeObjects.ts index 5eb8165065..b425895e68 100644 --- a/tests/fixtures/gpu/codeObjects.ts +++ b/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": [ {