mirror of
https://github.com/vercel-labs/scriptc.git
synced 2026-11-04 08:04:37 +08:00
Map wide shape-union records to their arm through one table per union
A shape union's tag helper switched on the record's shape number. Inlined into a caller's branch on the arm, LLVM folded the two into one jump table per site, spanning the arms' shape numbers. In the native compiler, whose IR unions span about 112 shape numbers, that came to about 27,600 tables and 9 MB of read-only data, most of which a cold build pages in. A union whose arms span more than 32 shape numbers now maps the shape number to its arm through one constant table, indexed by the shape number minus the arms' lowest. An entry is a byte, or a 16-bit word for a union of 255 arms or more; the all-ones entry names no arm. Callers branch on the small arm tag. Narrower unions keep the switch, whose folded jump tables stay small and dispatch in one jump.
This commit is contained in:
Chris Tate
@@ -0,0 +1,80 @@
|
||||
import { mkdtemp, readFile, rm, writeFile } from "node:fs/promises";
|
||||
import { tmpdir } from "node:os";
|
||||
import { join } from "node:path";
|
||||
import { expect, test } from "vitest";
|
||||
import { compile, deserializeModule } from "../../index.js";
|
||||
import { validateModule } from "../../ir/validate.js";
|
||||
import { emitLlvmModule } from "./emitter.js";
|
||||
|
||||
async function llvmFor(source: string): Promise<string> {
|
||||
const dir = await mkdtemp(join(tmpdir(), "scriptc-shape-union-tags-"));
|
||||
try {
|
||||
const entry = join(dir, "main.ts"),
|
||||
output = join(dir, "main.ir.json");
|
||||
await writeFile(entry, source);
|
||||
const result = await compile(entry, { outDir: dir, outPath: output, outputKind: "ir" });
|
||||
if (!result.ok)
|
||||
throw new Error(result.diagnostics.map((item) => `${item.code}: ${item.message}`).join("\n"));
|
||||
const module = deserializeModule(await readFile(output, "utf8"));
|
||||
expect(validateModule(module)).toEqual([]);
|
||||
return emitLlvmModule(module);
|
||||
} finally {
|
||||
await rm(dir, { recursive: true, force: true });
|
||||
}
|
||||
}
|
||||
|
||||
function tagHelpers(ir: string): string[] {
|
||||
return [...ir.matchAll(/^define internal i32 @sc_stag\.[^\n]*\{[^]*?^}/gm)].map(
|
||||
(match) => match[0],
|
||||
);
|
||||
}
|
||||
|
||||
test("a union whose arms span few shape numbers switches on the shape word", async () => {
|
||||
const ir = await llvmFor(`
|
||||
type Num = { kind: "num"; value: number };
|
||||
type Neg = { kind: "neg"; operand: Node };
|
||||
type Add = { kind: "add"; left: Node; right: Node };
|
||||
type Node = Num | Neg | Add;
|
||||
function evaluate(node: Node): number {
|
||||
switch (node.kind) {
|
||||
case "num": return node.value;
|
||||
case "neg": return -evaluate(node.operand);
|
||||
case "add": return evaluate(node.left) + evaluate(node.right);
|
||||
}
|
||||
}
|
||||
const one: Node = { kind: "num", value: 1 };
|
||||
console.log(evaluate({ kind: "add", left: one, right: { kind: "neg", operand: one } }));
|
||||
`);
|
||||
expect(ir).not.toMatch(/@sc_stagtab\./);
|
||||
const helpers = tagHelpers(ir);
|
||||
expect(helpers.length).toBeGreaterThan(0);
|
||||
for (const helper of helpers) expect(helper).toMatch(/switch i32 %w/);
|
||||
});
|
||||
|
||||
test("a union whose arms span many shape numbers reads the arm from its table", async () => {
|
||||
const arms = Array.from({ length: 40 }, (_, i) => i);
|
||||
const ir = await llvmFor(`
|
||||
${arms.map((i) => `type K${i} = { kind: "k${i}"; v${i}: number };`).join("\n")}
|
||||
type Any = ${arms.map((i) => `K${i}`).join(" | ")};
|
||||
function value(item: Any): number {
|
||||
switch (item.kind) {
|
||||
${arms.map((i) => ` case "k${i}": return item.v${i};`).join("\n")}
|
||||
}
|
||||
}
|
||||
const items: Any[] = [${arms.map((i) => `{ kind: "k${i}", v${i}: ${i} }`).join(", ")}];
|
||||
let total = 0;
|
||||
for (const item of items) total += value(item);
|
||||
console.log(total);
|
||||
`);
|
||||
// One byte per shape number from the arms' lowest to highest; 255 names no arm.
|
||||
const table =
|
||||
/^(@sc_stagtab\.\S+) = private unnamed_addr constant \[(\d+) x i8\] \[([^\]]*)\]$/m.exec(ir);
|
||||
expect(table, "arm table").not.toBeNull();
|
||||
const entries = table![3]!.split(", ").map((entry) => Number(entry.replace("i8 ", "")));
|
||||
expect(entries).toHaveLength(Number(table![2]));
|
||||
expect(entries.filter((entry) => entry !== 255).sort((a, b) => a - b)).toEqual(arms);
|
||||
const helper = tagHelpers(ir).find((text) => text.includes(table![1]!));
|
||||
expect(helper, "tag helper").toBeDefined();
|
||||
expect(helper).toContain(`getelementptr inbounds [${table![2]} x i8], ptr ${table![1]}`);
|
||||
expect(helper).not.toMatch(/switch i32 %w/);
|
||||
});
|
||||
@@ -313,19 +313,46 @@ function shapeUnionRc(host: ShapeHost, union: ShapeUnion): { retain: string; rel
|
||||
};
|
||||
}
|
||||
|
||||
/** The widest span of shape numbers a union's arms may cover for its tag
|
||||
* helper to switch on the shape word (emitShapeUnionHelpers). */
|
||||
const SHAPE_TAG_SWITCH_SPAN = 32;
|
||||
|
||||
/** The tag, release and trace helpers of every shape union
|
||||
* (nullable-unions.ts). The tag helper maps a value to its arm: NULL and
|
||||
* the null sentinel are the unit arms, anything else a tagged record whose
|
||||
* shape word names its arm (the switch folds to a lookup table). */
|
||||
* shape word names its arm.
|
||||
* - Arms within SHAPE_TAG_SWITCH_SPAN shape numbers: the helper switches on
|
||||
* the shape word. Inlined, the switch folds into the caller's branch on
|
||||
* the arm, one indirect jump with a jump table of at most that many
|
||||
* entries per site.
|
||||
* - Wider unions: the shape word indexes the union's arm table, one entry
|
||||
* per shape number from the arms' lowest to highest, and callers branch
|
||||
* on the small arm tag. Folded switches would give every site of such a
|
||||
* union a jump table spanning all of its arms' shape numbers. */
|
||||
export function emitShapeUnionHelpers(host: ShapeHost): string[] {
|
||||
const out: string[] = [];
|
||||
const S = host.sizeType;
|
||||
for (const union of host.nullableUnions.allShapeUnions()) {
|
||||
const header = shapeWordHeader(S);
|
||||
const tagFn = shapeTagFn(union.unionId);
|
||||
const cases = union.arms
|
||||
.map((arm) => `i32 ${host.nullableUnions.shapeNumber(arm.shapeId)}, label %a${arm.tag}`)
|
||||
.join(" ");
|
||||
const numbers = union.arms.map((arm) => host.nullableUnions.shapeNumber(arm.shapeId));
|
||||
const lo = Math.min(...numbers);
|
||||
const span = Math.max(...numbers) - lo + 1;
|
||||
const byTable = span > SHAPE_TAG_SWITCH_SPAN;
|
||||
// Entries are bytes, or 16-bit words for a union of 255 arms or more.
|
||||
// A shape number that names no arm maps to the all-ones entry, which is
|
||||
// above every arm tag, so callers take their path for no arm as they do
|
||||
// for the -1 of a shape number outside the table.
|
||||
const wide = union.arms.some((arm) => arm.tag >= 255);
|
||||
const entry = wide ? "i16" : "i8";
|
||||
const none = wide ? 65535 : 255;
|
||||
const table = new Array<number>(span).fill(none);
|
||||
union.arms.forEach((arm, i) => (table[numbers[i]! - lo] = arm.tag));
|
||||
const tableName = tagFn.replace("@sc_stag.", "@sc_stagtab.");
|
||||
if (byTable)
|
||||
out.push(
|
||||
`${tableName} = private unnamed_addr constant [${span} x ${entry}] [${table.map((t) => `${entry} ${t}`).join(", ")}]`,
|
||||
);
|
||||
out.push(
|
||||
`define internal i32 ${tagFn}(ptr %u) alwaysinline ${FN_ATTRS} { ; shape union tag`,
|
||||
`entry:`,
|
||||
@@ -343,16 +370,35 @@ export function emitShapeUnionHelpers(host: ShapeHost): string[] {
|
||||
` ret i32 ${union.nullTag}`,
|
||||
`rec:`,
|
||||
);
|
||||
out.push(
|
||||
` %wp = getelementptr inbounds ${header}, ptr %u, i32 0, i32 ${SHAPE_WORD_INDEX}`,
|
||||
` %w = load i32, ptr %wp`,
|
||||
` switch i32 %w, label %bad [ ${cases} ]`,
|
||||
...union.arms.flatMap((arm) => [`a${arm.tag}:`, ` ret i32 ${arm.tag}`]),
|
||||
`bad:`,
|
||||
` ret i32 -1`,
|
||||
`}`,
|
||||
``,
|
||||
);
|
||||
const cases = union.arms.map((arm, i) => `i32 ${numbers[i]}, label %a${arm.tag}`).join(" ");
|
||||
if (!byTable)
|
||||
out.push(
|
||||
` %wp = getelementptr inbounds ${header}, ptr %u, i32 0, i32 ${SHAPE_WORD_INDEX}`,
|
||||
` %w = load i32, ptr %wp`,
|
||||
` switch i32 %w, label %bad [ ${cases} ]`,
|
||||
...union.arms.flatMap((arm) => [`a${arm.tag}:`, ` ret i32 ${arm.tag}`]),
|
||||
`bad:`,
|
||||
` ret i32 -1`,
|
||||
`}`,
|
||||
``,
|
||||
);
|
||||
else
|
||||
out.push(
|
||||
` %wp = getelementptr inbounds ${header}, ptr %u, i32 0, i32 ${SHAPE_WORD_INDEX}`,
|
||||
` %w = load i32, ptr %wp`,
|
||||
` %i = sub i32 %w, ${lo}`,
|
||||
` %inrange = icmp ult i32 %i, ${span}`,
|
||||
` br i1 %inrange, label %look, label %bad`,
|
||||
`look:`,
|
||||
` %tp = getelementptr inbounds [${span} x ${entry}], ptr ${tableName}, i32 0, i32 %i`,
|
||||
` %t8 = load ${entry}, ptr %tp`,
|
||||
` %tag = zext ${entry} %t8 to i32`,
|
||||
` ret i32 %tag`,
|
||||
`bad:`,
|
||||
` ret i32 -1`,
|
||||
`}`,
|
||||
``,
|
||||
);
|
||||
// Release and trace: one definition for the unions over the same
|
||||
// record arms (helperOwner). Unit encodings are NULL or the immortal
|
||||
// null sentinel; the shape word picks the arm.
|
||||
|
||||
Reference in New Issue
Block a user