SNAPKITTYWEST's picture
October 2026 main drop: mirror from GitHub
3ce1a40 verified
Raw History Blame Contribute Delete
11.8 kB
/**
* Kernel IR -> WebAssembly binary.
*
* Every IR virtual register becomes a typed wasm local, which keeps the stack
* discipline trivial: each instruction pushes its operands with local.get and
* stores its result with local.set. Structured IR control flow maps directly onto
* wasm block / loop / if. The module layout follows ../abi/abi.js.
*/
import { Op, Type, Builtin, isFloat } from '../../kernel/ir/ir.js';
import { MEMORY_IMPORT, THREAD_PARAMS, LAUNCH_PARAMS, ELEMENT } from '../abi/abi.js';
import {
ValType, BLOCK_EMPTY, Section, ExternalKind, OpCode, SatOpCode, MAGIC_AND_VERSION,
uleb, sleb, f32, f64, name, vec, section, funcType, funcBody,
} from './encoder.js';
export class CodegenError extends Error {
constructor(message) {
super(message);
this.name = 'CodegenError';
}
}
const VAL_TYPE = {
[Type.I32]: ValType.I32, [Type.U32]: ValType.I32, [Type.BOOL]: ValType.I32,
[Type.F32]: ValType.F32, [Type.F64]: ValType.F64,
};
/** The wasm instruction family an IR type computes in: i32 (signed), u32, f32 or f64. */
const family = (type) => (type === Type.BOOL ? Type.I32 : type);
const ARITH = {
[Op.ADD]: { i32: 'i32_add', u32: 'i32_add', f32: 'f32_add', f64: 'f64_add' },
[Op.SUB]: { i32: 'i32_sub', u32: 'i32_sub', f32: 'f32_sub', f64: 'f64_sub' },
[Op.MUL]: { i32: 'i32_mul', u32: 'i32_mul', f32: 'f32_mul', f64: 'f64_mul' },
[Op.DIV]: { i32: 'i32_div_s', u32: 'i32_div_u', f32: 'f32_div', f64: 'f64_div' },
[Op.REM]: { i32: 'i32_rem_s', u32: 'i32_rem_u' },
[Op.BIT_AND]: { i32: 'i32_and', u32: 'i32_and' },
[Op.BIT_OR]: { i32: 'i32_or', u32: 'i32_or' },
[Op.BIT_XOR]: { i32: 'i32_xor', u32: 'i32_xor' },
[Op.SHL]: { i32: 'i32_shl', u32: 'i32_shl' },
[Op.SHR]: { i32: 'i32_shr_s', u32: 'i32_shr_u' },
[Op.CMP_LT]: { i32: 'i32_lt_s', u32: 'i32_lt_u', f32: 'f32_lt', f64: 'f64_lt' },
[Op.CMP_GT]: { i32: 'i32_gt_s', u32: 'i32_gt_u', f32: 'f32_gt', f64: 'f64_gt' },
[Op.CMP_LE]: { i32: 'i32_le_s', u32: 'i32_le_u', f32: 'f32_le', f64: 'f64_le' },
[Op.CMP_GE]: { i32: 'i32_ge_s', u32: 'i32_ge_u', f32: 'f32_ge', f64: 'f64_ge' },
[Op.CMP_EQ]: { i32: 'i32_eq', u32: 'i32_eq', f32: 'f32_eq', f64: 'f64_eq' },
[Op.CMP_NE]: { i32: 'i32_ne', u32: 'i32_ne', f32: 'f32_ne', f64: 'f64_ne' },
};
/** CONVERT opcodes keyed by `${from}>${to}`; an empty list is a reinterpretation. */
const CONVERSIONS = {
'i32>u32': [], 'u32>i32': [], 'bool>i32': [], 'bool>u32': [],
'i32>f32': ['f32_convert_i32_s'], 'u32>f32': ['f32_convert_i32_u'], 'bool>f32': ['f32_convert_i32_s'],
'i32>f64': ['f64_convert_i32_s'], 'u32>f64': ['f64_convert_i32_u'], 'bool>f64': ['f64_convert_i32_s'],
'f32>f64': ['f64_promote_f32'], 'f64>f32': ['f32_demote_f64'],
'f32>i32': [{ sat: 'i32_trunc_sat_f32_s' }], 'f32>u32': [{ sat: 'i32_trunc_sat_f32_u' }],
'f64>i32': [{ sat: 'i32_trunc_sat_f64_s' }], 'f64>u32': [{ sat: 'i32_trunc_sat_f64_u' }],
};
const LOAD = { i32: 'i32_load', u32: 'i32_load', f32: 'f32_load', f64: 'f64_load' };
const STORE = { i32: 'i32_store', u32: 'i32_store', f32: 'f32_store', f64: 'f64_store' };
/** Emits the body function of a kernel. Returns { params, results, locals, code }. */
function emitBody(ir) {
const isVoid = ir.returnType === null;
const code = [];
const localIndex = new Map();
const params = [];
const locals = [];
for (const p of ir.params) {
localIndex.set(p.name, params.length);
params.push(p.kind === 'buffer' ? ValType.I32 : VAL_TYPE[p.type]);
}
if (isVoid) {
for (const t of THREAD_PARAMS) {
localIndex.set(`@${t}`, params.length);
params.push(ValType.I32);
}
}
const addLocal = (key, valType) => {
localIndex.set(key, params.length + locals.length);
locals.push(valType);
};
for (const local of ir.locals) addLocal(local.name, VAL_TYPE[local.type]);
const vreg = (val) => {
const key = `%${val.id}`;
if (!localIndex.has(key)) addLocal(key, VAL_TYPE[val.type]);
return localIndex.get(key);
};
const buffers = new Map(ir.params.filter((p) => p.kind === 'buffer').map((p) => [p.name, p]));
const op = (mnemonic) => code.push(OpCode[mnemonic]);
const get = (val) => code.push(OpCode.local_get, ...uleb(vreg(val)));
const set = (val) => code.push(OpCode.local_set, ...uleb(vreg(val)));
const getNamed = (key) => code.push(OpCode.local_get, ...uleb(localIndex.get(key)));
const setNamed = (key) => code.push(OpCode.local_set, ...uleb(localIndex.get(key)));
function pushConst(type, imm) {
if (type === Type.F32) code.push(OpCode.f32_const, ...f32(imm));
else if (type === Type.F64) code.push(OpCode.f64_const, ...f64(imm));
else code.push(OpCode.i32_const, ...sleb(imm));
}
/** Pushes the byte address of buffer[index]. */
function address(bufferName, index) {
const buffer = buffers.get(bufferName);
getNamed(buffer.name);
get(index);
code.push(OpCode.i32_const, ...sleb(ELEMENT[buffer.type].log2));
op('i32_shl');
op('i32_add');
return buffer;
}
/** Memory immediate: alignment hint (log2 bytes) and a zero offset. */
const memarg = (type) => code.push(...uleb(ELEMENT[type].log2), ...uleb(0));
// Control stack of labels so BREAK / CONTINUE can compute branch depths.
const labels = [];
const depthOf = (kind) => {
for (let i = labels.length - 1; i >= 0; i -= 1) {
if (labels[i] === kind) return labels.length - 1 - i;
}
throw new CodegenError(`'${kind}' has no enclosing loop`);
};
const open = (opcode, label) => {
code.push(opcode, BLOCK_EMPTY);
labels.push(label);
};
const close = () => {
op('end');
labels.pop();
};
function emitInstr(ins) {
switch (ins.op) {
case Op.CONST:
pushConst(ins.out.type, ins.imm);
set(ins.out);
return;
case Op.LOAD:
getNamed(ins.name);
set(ins.out);
return;
case Op.STORE:
get(ins.value);
setNamed(ins.name);
return;
case Op.BUILTIN:
if (ins.name === Builtin.GLOBAL_ID_X) {
getNamed('@block_id_x');
getNamed('@block_dim_x');
op('i32_mul');
getNamed('@local_id_x');
op('i32_add');
} else {
getNamed(`@${ins.name}`);
}
set(ins.out);
return;
case Op.NOT:
get(ins.operand);
op('i32_eqz');
set(ins.out);
return;
case Op.NEG:
if (isFloat(ins.out.type)) {
get(ins.operand);
op(ins.out.type === Type.F32 ? 'f32_neg' : 'f64_neg');
} else {
code.push(OpCode.i32_const, 0);
get(ins.operand);
op('i32_sub');
}
set(ins.out);
return;
case Op.BIT_NOT:
get(ins.operand);
code.push(OpCode.i32_const, ...sleb(-1));
op('i32_xor');
set(ins.out);
return;
case Op.CONVERT: {
const steps = CONVERSIONS[`${ins.value.type}>${ins.out.type}`];
if (!steps) throw new CodegenError(`No conversion from ${ins.value.type} to ${ins.out.type}`);
get(ins.value);
for (const step of steps) {
if (typeof step === 'string') op(step);
else code.push(0xfc, ...uleb(SatOpCode[step.sat]));
}
set(ins.out);
return;
}
case Op.BUFFER_LOAD: {
const buffer = address(ins.buffer, ins.index);
op(LOAD[buffer.type]);
memarg(buffer.type);
set(ins.out);
return;
}
case Op.BUFFER_STORE: {
const buffer = address(ins.buffer, ins.index);
get(ins.value);
op(STORE[buffer.type]);
memarg(buffer.type);
return;
}
case Op.IF:
get(ins.cond);
open(OpCode.if, 'if');
emitList(ins.then);
if (ins.otherwise.length) {
op('else');
emitList(ins.otherwise);
}
close();
return;
case Op.LOOP:
// block $break { loop $top { head; br_if $break (!cond); block $continue { body } step; br $top } }
open(OpCode.block, 'break');
open(OpCode.loop, 'top');
emitList(ins.head);
if (ins.cond) {
get(ins.cond);
op('i32_eqz');
code.push(OpCode.br_if, ...uleb(depthOf('break')));
}
open(OpCode.block, 'continue');
emitList(ins.body);
close();
emitList(ins.step);
code.push(OpCode.br, ...uleb(depthOf('top')));
close();
close();
return;
case Op.BREAK:
code.push(OpCode.br, ...uleb(depthOf('break')));
return;
case Op.CONTINUE:
code.push(OpCode.br, ...uleb(depthOf('continue')));
return;
case Op.RET:
if (ins.value) get(ins.value);
op('return');
return;
default: {
const table = ARITH[ins.op];
if (!table) throw new CodegenError(`Unsupported IR op '${ins.op}'`);
const mnemonic = table[family(ins.left.type)];
if (!mnemonic) throw new CodegenError(`'${ins.op}' is not defined for ${ins.left.type}`);
get(ins.left);
get(ins.right);
op(mnemonic);
set(ins.out);
}
}
}
function emitList(list) {
for (const ins of list) emitInstr(ins);
}
emitList(ir.body);
// A value-returning body must not fall off the end; wasm validation needs the stack typed.
if (!isVoid) op('unreachable');
return { params, results: isVoid ? [] : [VAL_TYPE[ir.returnType]], locals, code };
}
/**
* The exported launcher for a void kernel: loops over blocks then threads and calls
* the body (function `bodyIndex`) once per thread.
*/
function emitLauncher(ir, bodyIndex) {
const n = ir.params.length;
const params = [
...ir.params.map((p) => (p.kind === 'buffer' ? ValType.I32 : VAL_TYPE[p.type])),
...LAUNCH_PARAMS.map(() => ValType.I32),
];
const grid = n;
const blockDim = n + 1;
const block = n + 2;
const thread = n + 3;
const get = (i) => [OpCode.local_get, ...uleb(i)];
const set = (i) => [OpCode.local_set, ...uleb(i)];
const increment = (i) => [...get(i), OpCode.i32_const, 1, OpCode.i32_add, ...set(i)];
const call = [];
for (let i = 0; i < n; i += 1) call.push(...get(i));
// THREAD_PARAMS order: block_id_x, local_id_x, block_dim_x, grid_dim_x
call.push(...get(block), ...get(thread), ...get(blockDim), ...get(grid), OpCode.call, ...uleb(bodyIndex));
const code = [
OpCode.block, BLOCK_EMPTY, OpCode.loop, BLOCK_EMPTY,
...get(block), ...get(grid), OpCode.i32_ge_u, OpCode.br_if, 1,
OpCode.i32_const, 0, ...set(thread),
OpCode.block, BLOCK_EMPTY, OpCode.loop, BLOCK_EMPTY,
...get(thread), ...get(blockDim), OpCode.i32_ge_u, OpCode.br_if, 1,
...call,
...increment(thread),
OpCode.br, 0,
OpCode.end, OpCode.end,
...increment(block),
OpCode.br, 0,
OpCode.end, OpCode.end,
];
return { params, results: [], locals: [ValType.I32, ValType.I32], code };
}
/** Compiles one kernel's IR to a complete wasm module. */
export function emitModule(ir) {
const body = emitBody(ir);
const functions = [body];
if (ir.returnType === null) functions.push(emitLauncher(ir, 0));
const exported = functions.length - 1;
const bytes = [
...MAGIC_AND_VERSION,
...section(Section.TYPE, vec(functions.map((f) => funcType(f.params, f.results)))),
...section(Section.IMPORT, vec([[
...name(MEMORY_IMPORT.module), ...name(MEMORY_IMPORT.name), ExternalKind.MEMORY, 0x00, ...uleb(0),
]])),
...section(Section.FUNCTION, vec(functions.map((_, i) => uleb(i)))),
...section(Section.EXPORT, vec([[...name(ir.name), ExternalKind.FUNC, ...uleb(exported)]])),
...section(Section.CODE, vec(functions.map((f) => funcBody(f.locals, f.code)))),
];
return new Uint8Array(bytes);
}