392 lines
13 KiB
TypeScript
392 lines
13 KiB
TypeScript
import { z } from "zod";
|
|
import { createHash } from "node:crypto";
|
|
import { digestSchema } from "../digests.js";
|
|
import {
|
|
absoluteLocalPathStringSchema,
|
|
localPathStringSchema,
|
|
} from "../localPath.js";
|
|
const count = z.number().int().nonnegative().max(Number.MAX_SAFE_INTEGER);
|
|
export const gpuLocationSchema = z.strictObject({
|
|
container_id: z.string().nullable(),
|
|
offset: count,
|
|
size: count,
|
|
section_index: count.nullable(),
|
|
section_name: z.string().nullable(),
|
|
});
|
|
export const gpuObjectSchema = z.strictObject({
|
|
id: z.string(),
|
|
parent_id: z.string().nullable(),
|
|
kind: z.enum([
|
|
"cuda-fatbin",
|
|
"hip-bundle",
|
|
"native-section",
|
|
"ptx",
|
|
"cubin",
|
|
"amdgpu",
|
|
"host",
|
|
"unknown",
|
|
]),
|
|
target: z.string().nullable(),
|
|
payload_target: z.string().nullable(),
|
|
cuda_target_header: z
|
|
.strictObject({
|
|
os_abi: z.number().int().min(0).max(255),
|
|
flags: z.string().regex(/^0x[0-9a-f]{8}$/u),
|
|
location: gpuLocationSchema,
|
|
})
|
|
.nullable(),
|
|
target_consistency: z.enum([
|
|
"unexamined",
|
|
"single-source",
|
|
"consistent",
|
|
"base-only",
|
|
"conflict",
|
|
"unknown",
|
|
]),
|
|
bundle_id: z.string().nullable(),
|
|
compression: z.enum(["none", "zstd", "zlib", "lz4", "unsupported"]),
|
|
compression_id: z.string().nullable(),
|
|
encoded_size: count,
|
|
decoded_size: count.nullable(),
|
|
sha256: digestSchema.nullable(),
|
|
raw_flags: z.string().nullable(),
|
|
location: gpuLocationSchema,
|
|
metadata_location: gpuLocationSchema,
|
|
inspection: z.enum([
|
|
"inventory",
|
|
"complete",
|
|
"partial",
|
|
"not-selected",
|
|
"unsupported",
|
|
"malformed",
|
|
"capacity",
|
|
]),
|
|
limitations: z.array(z.string()),
|
|
});
|
|
export const gpuParameterSchema = z.strictObject({
|
|
ordinal: count,
|
|
name: z.string().nullable(),
|
|
offset: count.nullable(),
|
|
size: count.nullable(),
|
|
type: z.string().nullable(),
|
|
hidden: z.boolean().nullable(),
|
|
});
|
|
export const gpuKernelSchema = z.strictObject({
|
|
id: z.string(),
|
|
object_id: z.string(),
|
|
representation: z.enum([
|
|
"ptx-declarations",
|
|
"cuda-elf-metadata",
|
|
"amdgpu-metadata",
|
|
]),
|
|
name: z.string(),
|
|
symbol: z.string().nullable(),
|
|
parameter_count: count.nullable(),
|
|
parameter_bytes: count.nullable(),
|
|
parameters: z.array(gpuParameterSchema),
|
|
registers: z.strictObject({
|
|
physical: count.nullable(),
|
|
sgpr: count.nullable(),
|
|
vgpr: count.nullable(),
|
|
agpr: count.nullable(),
|
|
virtual: z.array(
|
|
z.strictObject({ type: z.string(), count: count.nullable() }),
|
|
),
|
|
}),
|
|
shared_bytes: count
|
|
.nullable()
|
|
.describe(
|
|
"PTX: entry-local declared static bytes; CUDA ELF: allocated section bytes; AMDGPU: group_segment_fixed_size. Runtime dynamic/module shared allocations are not established by PTX.",
|
|
),
|
|
private_bytes: count.nullable(),
|
|
location: gpuLocationSchema,
|
|
sources: z.strictObject({
|
|
parameters: gpuLocationSchema.nullable(),
|
|
registers: gpuLocationSchema.nullable(),
|
|
shared: gpuLocationSchema.nullable(),
|
|
private: gpuLocationSchema.nullable(),
|
|
}),
|
|
limitations: z.array(z.string()),
|
|
});
|
|
export const gpuPtxArtifactSchema = z
|
|
.strictObject({
|
|
id: z.string(),
|
|
object_id: z.string(),
|
|
sha256: digestSchema,
|
|
media_type: z.literal("text/x-ptx"),
|
|
encoding: z.literal("utf-8"),
|
|
byte_length: count,
|
|
text: z.string(),
|
|
location: gpuLocationSchema,
|
|
})
|
|
.superRefine((artifact, context) => {
|
|
if (
|
|
/[\uD800-\uDFFF]/u.test(artifact.text) ||
|
|
Buffer.byteLength(artifact.text, "utf8") !== artifact.byte_length ||
|
|
artifact.location.size !== artifact.byte_length
|
|
)
|
|
context.addIssue({
|
|
code: "custom",
|
|
path: ["byte_length"],
|
|
message: "PTX text and its UTF-8 source extent disagree",
|
|
});
|
|
if (
|
|
createHash("sha256").update(artifact.text, "utf8").digest("hex") !==
|
|
artifact.sha256
|
|
)
|
|
context.addIssue({
|
|
code: "custom",
|
|
path: ["sha256"],
|
|
message: "PTX digest does not match retained UTF-8 text",
|
|
});
|
|
});
|
|
export const gpuCodeObjectsObjectSchema = z.strictObject({
|
|
artifact: z.strictObject({
|
|
path: localPathStringSchema,
|
|
file_size: count,
|
|
sha256: digestSchema,
|
|
bytes_read: count,
|
|
}),
|
|
format: z.enum(["elf", "pe", "cuda-fatbin", "hip-bundle", "ptx"]),
|
|
analysis: z.enum(["inventory", "kernels"]),
|
|
objects: z.array(gpuObjectSchema),
|
|
kernels: z.array(gpuKernelSchema),
|
|
ptx_artifacts: z.array(gpuPtxArtifactSchema),
|
|
coverage: z.strictObject({
|
|
discovery_scope: z.literal(
|
|
"supported-section-names-and-standalone-formats",
|
|
),
|
|
discovery: z.enum(["complete", "partial"]),
|
|
payloads: z.enum(["not-requested", "complete", "partial"]),
|
|
examined_sections: count,
|
|
unknown_section_names: count,
|
|
}),
|
|
limitations: z.array(z.string()),
|
|
});
|
|
/** Validate local identities and coordinate spaces once before consumers build lookup indexes. */
|
|
export const gpuCodeObjectsSchema = gpuCodeObjectsObjectSchema.superRefine(
|
|
(report, context) => {
|
|
const objects = new Map<string, GpuObject>();
|
|
const issue = (path: (string | number)[], message: string): void =>
|
|
context.addIssue({ code: "custom", path, message });
|
|
for (let index = 0; index < report.objects.length; index++) {
|
|
const row = report.objects[index]!;
|
|
if (objects.has(row.id))
|
|
issue(["objects", index, "id"], "GPU object IDs must be unique");
|
|
if (row.parent_id !== null && !objects.has(row.parent_id))
|
|
issue(
|
|
["objects", index, "parent_id"],
|
|
"GPU parents must precede their children",
|
|
);
|
|
if (
|
|
row.location.container_id !== row.parent_id ||
|
|
row.metadata_location.container_id !== row.parent_id
|
|
)
|
|
issue(
|
|
["objects", index, "location"],
|
|
"GPU object ranges must use their parent's coordinate space",
|
|
);
|
|
objects.set(row.id, row);
|
|
}
|
|
const range = (value: GpuLocation, path: (string | number)[]): void => {
|
|
const size =
|
|
value.container_id === null
|
|
? report.artifact.file_size
|
|
: objects.get(value.container_id)?.decoded_size;
|
|
if (
|
|
size === undefined ||
|
|
(size !== null &&
|
|
(value.offset > size || value.size > size - value.offset))
|
|
)
|
|
issue(
|
|
path,
|
|
"GPU source range is outside its declared coordinate space",
|
|
);
|
|
};
|
|
for (let index = 0; index < report.objects.length; index++) {
|
|
const row = report.objects[index]!;
|
|
range(row.location, ["objects", index, "location"]);
|
|
range(row.metadata_location, ["objects", index, "metadata_location"]);
|
|
const header = row.cuda_target_header;
|
|
if (header !== null) {
|
|
range(header.location, [
|
|
"objects",
|
|
index,
|
|
"cuda_target_header",
|
|
"location",
|
|
]);
|
|
if (
|
|
row.kind !== "cubin" ||
|
|
![52, 64].includes(header.location.size) ||
|
|
(header.location.container_id === row.id
|
|
? header.location.offset !== 0
|
|
: header.location.container_id !== row.parent_id ||
|
|
header.location.offset !== row.location.offset)
|
|
)
|
|
issue(
|
|
["objects", index, "cuda_target_header"],
|
|
"CUDA header provenance must start at its CUBIN object",
|
|
);
|
|
}
|
|
}
|
|
const kernelIds = new Set<string>();
|
|
for (let index = 0; index < report.kernels.length; index++) {
|
|
const row = report.kernels[index]!;
|
|
if (kernelIds.has(row.id))
|
|
issue(["kernels", index, "id"], "GPU kernel IDs must be unique");
|
|
kernelIds.add(row.id);
|
|
if (
|
|
!objects.has(row.object_id) ||
|
|
row.location.container_id !== row.object_id
|
|
)
|
|
issue(
|
|
["kernels", index, "object_id"],
|
|
"GPU kernels require an inventoried object and its coordinate space",
|
|
);
|
|
range(row.location, ["kernels", index, "location"]);
|
|
for (const [field, source] of Object.entries(row.sources))
|
|
if (source !== null)
|
|
range(source, ["kernels", index, "sources", field]);
|
|
}
|
|
const artifactIds = new Set<string>();
|
|
for (let index = 0; index < report.ptx_artifacts.length; index++) {
|
|
const row = report.ptx_artifacts[index]!;
|
|
if (artifactIds.has(row.id))
|
|
issue(
|
|
["ptx_artifacts", index, "id"],
|
|
"PTX artifact IDs must be unique",
|
|
);
|
|
artifactIds.add(row.id);
|
|
if (
|
|
objects.get(row.object_id)?.kind !== "ptx" ||
|
|
row.location.container_id !== row.object_id
|
|
)
|
|
issue(
|
|
["ptx_artifacts", index, "object_id"],
|
|
"PTX artifacts require an inventoried PTX object",
|
|
);
|
|
range(row.location, ["ptx_artifacts", index, "location"]);
|
|
}
|
|
},
|
|
);
|
|
/** Interrupted observations retain facts without claiming final source identity or completeness. */
|
|
export const gpuCodeObjectsPartialSchema = z.strictObject({
|
|
kind: z.literal("gpu-code-objects-partial"),
|
|
artifact: z.strictObject({
|
|
path: localPathStringSchema,
|
|
observed_file_size: count,
|
|
sha256: z.null(),
|
|
file_identity: z.literal("unverified"),
|
|
}),
|
|
analysis: z.enum(["inventory", "kernels"]),
|
|
objects: z.array(gpuObjectSchema),
|
|
kernels: z.array(gpuKernelSchema),
|
|
ptx_artifacts: z.array(gpuPtxArtifactSchema),
|
|
coverage: z.strictObject({
|
|
discovery: z.literal("partial"),
|
|
payloads: z.literal("partial"),
|
|
}),
|
|
limitations: z.array(z.string()),
|
|
});
|
|
export type GpuCodeObjectsPartial = z.infer<typeof gpuCodeObjectsPartialSchema>;
|
|
export const gpuTargetSchema = z.strictObject({
|
|
id: z.string(),
|
|
target: z.string().nullable(),
|
|
confidence: z.literal("derived"),
|
|
object_count: count,
|
|
kernel_count: count,
|
|
encoded_bytes: z.string().regex(/^(?:0|[1-9][0-9]*)$/u),
|
|
decoded_bytes: z
|
|
.string()
|
|
.regex(/^(?:0|[1-9][0-9]*)$/u)
|
|
.nullable(),
|
|
kernel_coverage: z.enum(["not-requested", "complete", "partial"]),
|
|
by_kind: z.strictObject({
|
|
ptx: count,
|
|
cubin: count,
|
|
amdgpu: count,
|
|
unknown: count,
|
|
}),
|
|
});
|
|
export type GpuTarget = z.infer<typeof gpuTargetSchema>;
|
|
const selection = {
|
|
ids: z.array(z.string()).optional(),
|
|
targets: z.array(z.string()).optional(),
|
|
offset: count.default(0),
|
|
limit: z.number().int().positive().default(32),
|
|
};
|
|
export const gpuCodeObjectsSelectionSchema = z.discriminatedUnion(
|
|
"collection",
|
|
[
|
|
z.strictObject({
|
|
...selection,
|
|
collection: z.literal("objects"),
|
|
fields: z.array(gpuObjectSchema.keyof()).optional(),
|
|
}),
|
|
z.strictObject({
|
|
...selection,
|
|
collection: z.literal("kernels"),
|
|
object_ids: z.array(z.string()).optional(),
|
|
names: z.array(z.string()).optional(),
|
|
fields: z.array(gpuKernelSchema.keyof()).optional(),
|
|
}),
|
|
z.strictObject({
|
|
collection: z.literal("ptx"),
|
|
object_id: z.string(),
|
|
offset: count.default(0),
|
|
limit: z.number().int().positive().default(4096),
|
|
}),
|
|
z.strictObject({
|
|
...selection,
|
|
collection: z.literal("targets"),
|
|
fields: z.array(gpuTargetSchema.keyof()).optional(),
|
|
}),
|
|
],
|
|
);
|
|
export const inspectGpuCodeObjectsInputSchema = z
|
|
.strictObject({
|
|
path: absoluteLocalPathStringSchema,
|
|
analysis: z.enum(["inventory", "kernels"]).default("inventory"),
|
|
expected_sha256: digestSchema
|
|
.optional()
|
|
.describe(
|
|
"Require the current full-file digest; mandatory when selecting inventory object_ids",
|
|
),
|
|
object_ids: z
|
|
.array(z.string())
|
|
.optional()
|
|
.describe(
|
|
"Decode only these inventory IDs; IDs are scoped to the artifact SHA-256",
|
|
),
|
|
targets: z
|
|
.array(z.string())
|
|
.optional()
|
|
.describe(
|
|
"Decode only exact target declarations, preserving all inventory rows",
|
|
),
|
|
detail: z.enum(["summary", "complete"]).default("summary"),
|
|
view: gpuCodeObjectsSelectionSchema.optional(),
|
|
})
|
|
.superRefine((input, context) => {
|
|
if (input.object_ids !== undefined && input.expected_sha256 === undefined)
|
|
context.addIssue({
|
|
code: "custom",
|
|
path: ["expected_sha256"],
|
|
message: "object_ids require the artifact SHA-256 from their inventory",
|
|
});
|
|
if (input.view !== undefined && input.detail === "complete")
|
|
context.addIssue({
|
|
code: "custom",
|
|
path: ["view"],
|
|
message: "view requires summary detail",
|
|
});
|
|
});
|
|
export type GpuCodeObjects = z.infer<typeof gpuCodeObjectsSchema>;
|
|
export type GpuObject = z.infer<typeof gpuObjectSchema>;
|
|
export type GpuKernel = z.infer<typeof gpuKernelSchema>;
|
|
export type GpuLocation = z.infer<typeof gpuLocationSchema>;
|
|
export type GpuCodeObjectsInput = Pick<
|
|
z.infer<typeof inspectGpuCodeObjectsInputSchema>,
|
|
"path" | "analysis" | "object_ids" | "targets" | "expected_sha256"
|
|
>;
|