302 lines
9.6 KiB
TypeScript
302 lines
9.6 KiB
TypeScript
import { createHash } from "node:crypto";
|
|
import { deflateSync, zstdCompressSync } from "node:zlib";
|
|
import { encode } from "@msgpack/msgpack";
|
|
export const ptxFixture = `// Producer-shaped static PTX; portable tests do not invoke a GPU compiler.
|
|
.version 8.8
|
|
.target sm_80
|
|
.address_size 64
|
|
.visible .entry add_one(.param .u64 output, .param .u32 value) {
|
|
.reg .b32 %r<9>;
|
|
.reg .b64 %rd<5>;
|
|
.shared .align 4 .b8 scratch[32];
|
|
ret;
|
|
}
|
|
`;
|
|
export interface ElfFixtureSection {
|
|
name: string;
|
|
type?: number;
|
|
flags?: bigint;
|
|
bytes?: Buffer;
|
|
size?: number;
|
|
info?: number;
|
|
link?: number;
|
|
entrySize?: number;
|
|
}
|
|
/** Format boundary fixtures, including numeric records checked against nvcc 12.9 cuobjdump. */
|
|
export const elfFixture = (
|
|
machine: number,
|
|
sections: ElfFixtureSection[],
|
|
version = 1,
|
|
): Buffer => {
|
|
const names = Buffer.from(
|
|
`\0.shstrtab\0${sections.map((row) => row.name).join("\0")}\0`,
|
|
);
|
|
const rows = [
|
|
{ name: "", type: 0, bytes: Buffer.alloc(0) },
|
|
{ name: ".shstrtab", type: 3, bytes: names },
|
|
...sections,
|
|
];
|
|
let position = 64;
|
|
const entries = rows.map((row) => {
|
|
position = Math.ceil(position / 8) * 8;
|
|
const offset = position;
|
|
if (row.type !== 8) position += row.bytes?.length ?? 0;
|
|
return { row, offset };
|
|
});
|
|
const table = Math.ceil(position / 8) * 8;
|
|
const result = Buffer.alloc(table + rows.length * 64);
|
|
result.set([
|
|
0x7f,
|
|
0x45,
|
|
0x4c,
|
|
0x46,
|
|
2,
|
|
1,
|
|
1,
|
|
machine === 190 ? 0x33 : machine === 224 ? 0x40 : 0,
|
|
]);
|
|
result.writeUInt16LE(2, 16);
|
|
result.writeUInt16LE(machine, 18);
|
|
result.writeUInt32LE(version, 20);
|
|
result.writeBigUInt64LE(BigInt(table), 40);
|
|
result.writeUInt32LE(machine === 190 ? 0x500550 : 0x12c, 48);
|
|
result.writeUInt16LE(64, 52);
|
|
result.writeUInt16LE(64, 58);
|
|
result.writeUInt16LE(rows.length, 60);
|
|
result.writeUInt16LE(1, 62);
|
|
for (let index = 1; index < entries.length; index++) {
|
|
const { row, offset } = entries[index]!,
|
|
h = table + index * 64;
|
|
result.writeUInt32LE(names.indexOf(Buffer.from(`${row.name}\0`)), h);
|
|
result.writeUInt32LE(row.type ?? 1, h + 4);
|
|
result.writeBigUInt64LE(row.flags ?? 0n, h + 8);
|
|
result.writeBigUInt64LE(BigInt(offset), h + 24);
|
|
result.writeBigUInt64LE(BigInt(row.size ?? row.bytes?.length ?? 0), h + 32);
|
|
result.writeUInt32LE(row.link ?? 0, h + 40);
|
|
result.writeUInt32LE(row.info ?? 0, h + 44);
|
|
result.writeBigUInt64LE(BigInt(row.entrySize ?? 0), h + 56);
|
|
if (row.type !== 8) row.bytes?.copy(result, offset);
|
|
}
|
|
return result;
|
|
};
|
|
export const cubinFixture = (
|
|
options: {
|
|
shared?: number;
|
|
textRegisters?: number;
|
|
globalRegisters?: number;
|
|
name?: string;
|
|
globalRecords?: number;
|
|
globalAliases?: number;
|
|
architecture?: number;
|
|
osAbi?: number;
|
|
flags?: number;
|
|
} = {},
|
|
): Buffer => {
|
|
const name = options.name ?? "add_one";
|
|
const symbols = Buffer.alloc(48);
|
|
symbols.writeUInt32LE(1, 24);
|
|
symbols[28] = 0x12;
|
|
symbols[29] = 0x10;
|
|
symbols.writeUInt16LE(4, 30);
|
|
symbols.writeBigUInt64LE(16n, 40);
|
|
const record = Buffer.from("042f0800010000000a000000", "hex");
|
|
record.writeUInt32LE(options.globalRegisters ?? 10, 8);
|
|
const global = Buffer.alloc(record.length * (options.globalRecords ?? 1));
|
|
for (let offset = 0; offset < global.length; offset += record.length)
|
|
record.copy(global, offset);
|
|
const info = Buffer.from(
|
|
"03190c0004170c00000000000100080000f0110004170c00000000000000000000f02100",
|
|
"hex",
|
|
);
|
|
const result = elfFixture(
|
|
190,
|
|
[
|
|
{ name: ".strtab", type: 3, bytes: Buffer.from(`\0${name}\0`) },
|
|
{ name: ".symtab", type: 2, link: 2, entrySize: 24, bytes: symbols },
|
|
{
|
|
name: `.text.${name}`,
|
|
flags: 0x100006n,
|
|
info: ((options.textRegisters ?? 0) << 24) | 1,
|
|
bytes: Buffer.alloc(16),
|
|
},
|
|
{ name: ".nv.info", type: 0x70000000, link: 3, bytes: global },
|
|
{
|
|
name: `.nv.info.${name}`,
|
|
type: 0x70000000,
|
|
link: 3,
|
|
info: 4,
|
|
bytes: info,
|
|
},
|
|
{
|
|
name: `.nv.shared.${name}`,
|
|
type: 8,
|
|
info: 4,
|
|
size: options.shared ?? 1056,
|
|
},
|
|
...Array.from({ length: (options.globalAliases ?? 1) - 1 }, () => ({
|
|
name: ".nv.info",
|
|
type: 0x70000000,
|
|
link: 3,
|
|
bytes: Buffer.alloc(0),
|
|
})),
|
|
],
|
|
0x81,
|
|
);
|
|
if (options.architecture !== undefined)
|
|
result.writeUInt32LE(
|
|
(result.readUInt32LE(48) & ~0xff) | options.architecture,
|
|
48,
|
|
);
|
|
const table = Number(result.readBigUInt64LE(40));
|
|
if (options.osAbi !== undefined) {
|
|
result[7] = options.osAbi;
|
|
result.writeUInt32LE(1, 20);
|
|
}
|
|
if (options.flags !== undefined) result.writeUInt32LE(options.flags, 48);
|
|
for (let index = 8; index < result.readUInt16LE(60); index++) {
|
|
result.writeBigUInt64LE(
|
|
result.readBigUInt64LE(table + 5 * 64 + 24),
|
|
table + index * 64 + 24,
|
|
);
|
|
result.writeBigUInt64LE(BigInt(global.length), table + index * 64 + 32);
|
|
}
|
|
return result;
|
|
};
|
|
export const amdFixture = (hidden = false, trailingNote?: Buffer): Buffer => {
|
|
const args = [
|
|
{ ".offset": 0, ".size": 8, ".value_kind": "global_buffer" },
|
|
{
|
|
".offset": 8,
|
|
".size": 4,
|
|
".value_kind": hidden ? "hidden_global_offset_x" : "by_value",
|
|
},
|
|
];
|
|
const metadata = Buffer.from(
|
|
encode({
|
|
"amdhsa.target": "amdgcn-amd-amdhsa--gfx900",
|
|
"amdhsa.version": [1, 1],
|
|
"amdhsa.kernels": [
|
|
{
|
|
".name": "add_one",
|
|
".symbol": "add_one.kd",
|
|
".args": args,
|
|
".kernarg_segment_size": 12,
|
|
".group_segment_fixed_size": 0,
|
|
".private_segment_fixed_size": 0,
|
|
".sgpr_count": 6,
|
|
".vgpr_count": 2,
|
|
},
|
|
],
|
|
}),
|
|
);
|
|
const note = Buffer.alloc(20 + Math.ceil(metadata.length / 4) * 4);
|
|
note.writeUInt32LE(7, 0);
|
|
note.writeUInt32LE(metadata.length, 4);
|
|
note.writeUInt32LE(32, 8);
|
|
note.write("AMDGPU\0", 12);
|
|
metadata.copy(note, 20);
|
|
return elfFixture(224, [
|
|
{
|
|
name: ".note",
|
|
type: 7,
|
|
bytes:
|
|
trailingNote === undefined ? note : Buffer.concat([note, trailingNote]),
|
|
},
|
|
]);
|
|
};
|
|
export const cudaFixture = (
|
|
compression: "none" | "zstd" = "none",
|
|
entries = [
|
|
{ kind: 1, target: 80, bytes: Buffer.from(ptxFixture) },
|
|
{ kind: 2, target: 90, bytes: cubinFixture({ architecture: 90 }) },
|
|
],
|
|
): Buffer => {
|
|
const records = entries.map(({ kind, target, bytes }) => {
|
|
const compressed = compression === "zstd",
|
|
data = compressed ? zstdCompressSync(bytes) : bytes;
|
|
const result = Buffer.alloc(64 + Math.ceil(data.length / 8) * 8);
|
|
result.writeUInt16LE(kind, 0);
|
|
result.writeUInt16LE(0x101, 2);
|
|
result.writeUInt32LE(64, 4);
|
|
result.writeBigUInt64LE(BigInt(result.length - 64), 8);
|
|
result.writeUInt32LE(compressed ? data.length : 0, 16);
|
|
result.writeUInt32LE(target, 28);
|
|
result.writeBigUInt64LE(compressed ? 0x8011n : 0x11n, 40);
|
|
result.writeBigUInt64LE(compressed ? BigInt(bytes.length) : 0n, 56);
|
|
data.copy(result, 64);
|
|
return result;
|
|
});
|
|
const header = Buffer.alloc(16);
|
|
header.writeUInt32LE(0xba55ed50);
|
|
header.writeUInt16LE(1, 4);
|
|
header.writeUInt16LE(16, 6);
|
|
header.writeBigUInt64LE(
|
|
BigInt(records.reduce((n, row) => n + row.length, 0)),
|
|
8,
|
|
);
|
|
return Buffer.concat([header, ...records]);
|
|
};
|
|
export const hipFixture = (
|
|
payload = amdFixture(),
|
|
targetId = "hipv4-amdgcn-amd-amdhsa--gfx900",
|
|
): Buffer => {
|
|
const ids = ["host-x86_64-unknown-linux", targetId];
|
|
const headers = ids.map((id) => Buffer.alloc(24 + Buffer.byteLength(id)));
|
|
const start = 32 + headers.reduce((total, row) => total + row.length, 0);
|
|
for (let index = 0; index < headers.length; index++) {
|
|
const row = headers[index]!;
|
|
row.writeBigUInt64LE(BigInt(start));
|
|
row.writeBigUInt64LE(index === 0 ? 0n : BigInt(payload.length), 8);
|
|
row.writeBigUInt64LE(BigInt(Buffer.byteLength(ids[index]!)), 16);
|
|
row.write(ids[index]!, 24);
|
|
}
|
|
const head = Buffer.alloc(32);
|
|
head.write("__CLANG_OFFLOAD_BUNDLE__");
|
|
head.writeBigUInt64LE(2n, 24);
|
|
return Buffer.concat([head, ...headers, payload]);
|
|
};
|
|
export const ccobFixture = (
|
|
version: 1 | 2 | 3 = 2,
|
|
method: "zlib" | "zstd" = "zstd",
|
|
bundle = hipFixture(),
|
|
): Buffer => {
|
|
const compressed =
|
|
method === "zstd" ? zstdCompressSync(bundle) : deflateSync(bundle);
|
|
const width = version === 1 ? 20 : version === 2 ? 24 : 32,
|
|
head = Buffer.alloc(width);
|
|
head.write("CCOB");
|
|
head.writeUInt16LE(version, 4);
|
|
head.writeUInt16LE(method === "zstd" ? 1 : 0, 6);
|
|
if (version === 3) {
|
|
head.writeBigUInt64LE(BigInt(width + compressed.length), 8);
|
|
head.writeBigUInt64LE(BigInt(bundle.length), 16);
|
|
} else if (version === 2) {
|
|
head.writeUInt32LE(width + compressed.length, 8);
|
|
head.writeUInt32LE(bundle.length, 12);
|
|
} else head.writeUInt32LE(bundle.length, 8);
|
|
createHash("md5")
|
|
.update(bundle)
|
|
.digest()
|
|
.copy(head, width - 8, 0, 8);
|
|
return Buffer.concat([head, compressed]);
|
|
};
|
|
export const peGpuFixture = (payload = cudaFixture()): Buffer => {
|
|
const result = Buffer.alloc(512 + payload.length),
|
|
pe = 64,
|
|
optionalSize = 240,
|
|
section = pe + 24 + optionalSize;
|
|
result.write("MZ");
|
|
result.writeUInt32LE(pe, 60);
|
|
result.write("PE\0\0", pe);
|
|
result.writeUInt16LE(0x8664, pe + 4);
|
|
result.writeUInt16LE(1, pe + 6);
|
|
result.writeUInt16LE(optionalSize, pe + 20);
|
|
result.writeUInt16LE(0x20b, pe + 24);
|
|
result.write(".nv_fatb", section);
|
|
result.writeUInt32LE(payload.length, section + 8);
|
|
result.writeUInt32LE(payload.length, section + 16);
|
|
result.writeUInt32LE(512, section + 20);
|
|
payload.copy(result, 512);
|
|
return result;
|
|
};
|