Download unified-engine/wasm/modules/codegen.js from Snapkitty/repoverse: direct link, hf CLI and curl.
- Browser
- Download file 11.8 kB
-
https://huggingface.co/Snapkitty/repoverse/resolve/main/unified-engine/wasm/modules/codegen.js
- Command line
-
hf download hf://Snapkitty/repoverse/unified-engine/wasm/modules/codegen.js
-
curl -L -o codegen.js https://huggingface.co/Snapkitty/repoverse/resolve/main/unified-engine/wasm/modules/codegen.js
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); | |
| } | |