Download unified-engine/kernel/ir/lowering.js from Snapkitty/repoverse: direct link, hf CLI and curl.
- Browser
- Download file 13.6 kB
-
https://huggingface.co/Snapkitty/repoverse/resolve/main/unified-engine/kernel/ir/lowering.js
- Command line
-
hf download hf://Snapkitty/repoverse/unified-engine/kernel/ir/lowering.js
-
curl -L -o lowering.js https://huggingface.co/Snapkitty/repoverse/resolve/main/unified-engine/kernel/ir/lowering.js
13.6 kB
| /** | |
| * AST -> IR lowering. | |
| * | |
| * Resolves names against C block scopes, assigns every expression a type using C's | |
| * usual arithmetic conversions, inserts explicit CONVERTs, and turns && / || into | |
| * short-circuit control flow. Every failure raises LowerError; the IR it returns is | |
| * fully typed so a backend never has to infer anything. | |
| */ | |
| import { | |
| Op, Type, SCALAR_TYPES, Builtin, isFloat, | |
| makeIdCounter, value, instr, kernel, | |
| } from './ir.js'; | |
| export class LowerError extends Error { | |
| constructor(message) { | |
| super(message); | |
| this.name = 'LowerError'; | |
| } | |
| } | |
| const ARITHMETIC = { '+': Op.ADD, '-': Op.SUB, '*': Op.MUL, '/': Op.DIV, '%': Op.REM }; | |
| const BITWISE = { '&': Op.BIT_AND, '|': Op.BIT_OR, '^': Op.BIT_XOR }; | |
| const SHIFTS = { '<<': Op.SHL, '>>': Op.SHR }; | |
| const COMPARISONS = { | |
| '<': Op.CMP_LT, '>': Op.CMP_GT, '<=': Op.CMP_LE, | |
| '>=': Op.CMP_GE, '==': Op.CMP_EQ, '!=': Op.CMP_NE, | |
| }; | |
| /** Builtins the lowering can read. Only the x dimension is launched today. */ | |
| const SUPPORTED_BUILTINS = new Set(Object.values(Builtin)); | |
| /** Integer promotion: bool arithmetic happens in int. */ | |
| const promote = (type) => (type === Type.BOOL ? Type.I32 : type); | |
| /** C's usual arithmetic conversions, restricted to the supported scalars. */ | |
| function commonType(a, b) { | |
| const pa = promote(a); | |
| const pb = promote(b); | |
| for (const type of [Type.F64, Type.F32, Type.U32]) { | |
| if (pa === type || pb === type) return type; | |
| } | |
| return Type.I32; | |
| } | |
| function scalarType(name) { | |
| const type = SCALAR_TYPES[name]; | |
| if (!type) throw new LowerError(`'${name}' is not a supported value type`); | |
| return type; | |
| } | |
| export function lower(ast) { | |
| const ids = makeIdCounter(); | |
| const isVoid = ast.returnType === 'void'; | |
| const returnType = isVoid ? null : scalarType(ast.returnType); | |
| const params = []; | |
| const usedNames = new Set(); | |
| const paramScope = new Map(); | |
| for (const p of ast.params) { | |
| if (usedNames.has(p.name)) throw new LowerError(`Duplicate parameter '${p.name}'`); | |
| usedNames.add(p.name); | |
| const param = { name: p.name, kind: p.type.kind, type: scalarType(p.type.scalar), isConst: p.type.isConst }; | |
| params.push(param); | |
| paramScope.set(p.name, param); | |
| } | |
| const locals = []; | |
| const builtins = new Set(); | |
| const scopes = [paramScope]; | |
| let out = []; | |
| let loopDepth = 0; | |
| // ---- helpers ---- | |
| const fresh = (type) => value(ids.next(), type); | |
| const emit = (ins) => { | |
| out.push(ins); | |
| return ins.out; | |
| }; | |
| /** Runs `fn` with a fresh instruction list and returns what it emitted. */ | |
| const collect = (fn) => { | |
| const saved = out; | |
| out = []; | |
| try { | |
| fn(); | |
| return out; | |
| } finally { | |
| out = saved; | |
| } | |
| }; | |
| const scoped = (fn) => { | |
| scopes.push(new Map()); | |
| try { | |
| return fn(); | |
| } finally { | |
| scopes.pop(); | |
| } | |
| }; | |
| function declareLocal(sourceName, type) { | |
| let name = sourceName; | |
| for (let n = 1; usedNames.has(name); n += 1) name = `${sourceName}.${n}`; | |
| usedNames.add(name); | |
| const local = { name, kind: 'scalar', type, isConst: false }; | |
| locals.push({ name, type }); | |
| return local; | |
| } | |
| function lookup(name) { | |
| for (let i = scopes.length - 1; i >= 0; i -= 1) { | |
| const binding = scopes[i].get(name); | |
| if (binding) return binding; | |
| } | |
| throw new LowerError(`Unknown variable '${name}'`); | |
| } | |
| const constant = (type, imm) => emit(instr(Op.CONST, { imm }, fresh(type))); | |
| function convert(val, to) { | |
| if (val.type === to) return val; | |
| if (to === Type.BOOL) return toBool(val); | |
| return emit(instr(Op.CONVERT, { value: val }, fresh(to))); | |
| } | |
| function toBool(val) { | |
| if (val.type === Type.BOOL) return val; | |
| const zero = constant(val.type, 0); | |
| return emit(instr(Op.CMP_NE, { left: val, right: zero }, fresh(Type.BOOL))); | |
| } | |
| function requireInteger(val, what) { | |
| if (isFloat(val.type)) throw new LowerError(`${what} requires integer operands`); | |
| } | |
| function readBuiltin(name) { | |
| if (!SUPPORTED_BUILTINS.has(name)) { | |
| throw new LowerError(`Builtin '${name}' is not supported; only the x dimension is available`); | |
| } | |
| if (!isVoid) throw new LowerError(`Builtin '${name}' is only available in void kernels`); | |
| builtins.add(name); | |
| return emit(instr(Op.BUILTIN, { name }, fresh(Type.I32))); | |
| } | |
| function bufferBinding(expr) { | |
| if (expr.kind !== 'var') throw new LowerError('Only a buffer parameter can be indexed'); | |
| const binding = lookup(expr.name); | |
| if (binding.kind !== 'buffer') throw new LowerError(`'${expr.name}' is not a buffer`); | |
| return binding; | |
| } | |
| function lowerIndex(expr) { | |
| const index = lowerExpr(expr); | |
| requireInteger(index, 'A buffer index'); | |
| return convert(index, Type.I32); | |
| } | |
| /** Applies a binary operator to two already-lowered values. */ | |
| function binaryOp(op, left, right) { | |
| if (op in ARITHMETIC) { | |
| const type = commonType(left.type, right.type); | |
| if (op === '%' && isFloat(type)) throw new LowerError("'%' requires integer operands"); | |
| const l = convert(left, type); | |
| const r = convert(right, type); | |
| return emit(instr(ARITHMETIC[op], { left: l, right: r }, fresh(type))); | |
| } | |
| if (op in BITWISE) { | |
| requireInteger(left, `'${op}'`); | |
| requireInteger(right, `'${op}'`); | |
| const type = commonType(left.type, right.type); | |
| return emit(instr(BITWISE[op], { left: convert(left, type), right: convert(right, type) }, fresh(type))); | |
| } | |
| if (op in SHIFTS) { | |
| requireInteger(left, `'${op}'`); | |
| requireInteger(right, `'${op}'`); | |
| // The result has the promoted type of the left operand, not the common type. | |
| const type = promote(left.type); | |
| return emit(instr(SHIFTS[op], { left: convert(left, type), right: convert(right, type) }, fresh(type))); | |
| } | |
| if (op in COMPARISONS) { | |
| const type = commonType(left.type, right.type); | |
| const l = convert(left, type); | |
| const r = convert(right, type); | |
| return emit(instr(COMPARISONS[op], { left: l, right: r }, fresh(Type.BOOL))); | |
| } | |
| throw new LowerError(`Unsupported operator '${op}'`); | |
| } | |
| /** a && b / a || b: the right side runs only when it can change the result. */ | |
| function shortCircuit(expr) { | |
| const result = declareLocal('$sc', Type.BOOL); | |
| const left = toBool(lowerExpr(expr.left)); | |
| emit(instr(Op.STORE, { name: result.name, value: left })); | |
| const evalRight = collect(() => { | |
| const right = toBool(lowerExpr(expr.right)); | |
| emit(instr(Op.STORE, { name: result.name, value: right })); | |
| }); | |
| const isAnd = expr.op === '&&'; | |
| emit(instr(Op.IF, { cond: left, then: isAnd ? evalRight : [], otherwise: isAnd ? [] : evalRight })); | |
| return emit(instr(Op.LOAD, { name: result.name }, fresh(Type.BOOL))); | |
| } | |
| function lowerAssign(expr) { | |
| const compound = expr.op === '=' ? null : expr.op.slice(0, -1); | |
| if (expr.target.kind === 'var') { | |
| const binding = lookup(expr.target.name); | |
| if (binding.kind === 'buffer') throw new LowerError(`Cannot assign to buffer '${expr.target.name}'`); | |
| if (binding.isConst) throw new LowerError(`Cannot assign to const '${expr.target.name}'`); | |
| let result = lowerExpr(expr.value); | |
| if (compound) { | |
| const current = emit(instr(Op.LOAD, { name: binding.name }, fresh(binding.type))); | |
| result = binaryOp(compound, current, result); | |
| } | |
| result = convert(result, binding.type); | |
| emit(instr(Op.STORE, { name: binding.name, value: result })); | |
| return result; | |
| } | |
| if (expr.target.kind === 'index') { | |
| const binding = bufferBinding(expr.target.base); | |
| if (binding.isConst) throw new LowerError(`Cannot write to const buffer '${binding.name}'`); | |
| const index = lowerIndex(expr.target.index); | |
| let result = lowerExpr(expr.value); | |
| if (compound) { | |
| const current = emit(instr(Op.BUFFER_LOAD, { buffer: binding.name, index }, fresh(binding.type))); | |
| result = binaryOp(compound, current, result); | |
| } | |
| result = convert(result, binding.type); | |
| emit(instr(Op.BUFFER_STORE, { buffer: binding.name, index, value: result })); | |
| return result; | |
| } | |
| throw new LowerError('The left side of an assignment must be a variable or a buffer element'); | |
| } | |
| function lowerExpr(expr) { | |
| switch (expr.kind) { | |
| case 'int': { | |
| if (expr.value <= 0x7fffffff) return constant(Type.I32, expr.value); | |
| if (expr.value <= 0xffffffff) return constant(Type.U32, expr.value); | |
| throw new LowerError(`Integer literal ${expr.value} does not fit in 32 bits`); | |
| } | |
| case 'float': | |
| return constant(expr.single ? Type.F32 : Type.F64, expr.value); | |
| case 'bool': | |
| return constant(Type.BOOL, expr.value ? 1 : 0); | |
| case 'var': { | |
| const binding = lookup(expr.name); | |
| if (binding.kind === 'buffer') throw new LowerError(`Buffer '${expr.name}' must be indexed`); | |
| return emit(instr(Op.LOAD, { name: binding.name }, fresh(binding.type))); | |
| } | |
| case 'builtin': | |
| return readBuiltin(expr.name); | |
| case 'call': | |
| // The parser maps member builtins such as threadIdx.x onto call nodes. | |
| if (expr.args.length !== 0) throw new LowerError(`Unknown function '${expr.callee}'`); | |
| return readBuiltin(expr.callee); | |
| case 'index': { | |
| const binding = bufferBinding(expr.base); | |
| const index = lowerIndex(expr.index); | |
| return emit(instr(Op.BUFFER_LOAD, { buffer: binding.name, index }, fresh(binding.type))); | |
| } | |
| case 'unary': { | |
| const operand = lowerExpr(expr.operand); | |
| if (expr.op === '!') { | |
| return emit(instr(Op.NOT, { operand: toBool(operand) }, fresh(Type.BOOL))); | |
| } | |
| if (expr.op === '-') { | |
| const type = promote(operand.type); | |
| return emit(instr(Op.NEG, { operand: convert(operand, type) }, fresh(type))); | |
| } | |
| if (expr.op === '~') { | |
| requireInteger(operand, "'~'"); | |
| const type = promote(operand.type); | |
| return emit(instr(Op.BIT_NOT, { operand: convert(operand, type) }, fresh(type))); | |
| } | |
| throw new LowerError(`Unsupported unary operator '${expr.op}'`); | |
| } | |
| case 'binary': | |
| if (expr.op === '&&' || expr.op === '||') return shortCircuit(expr); | |
| return binaryOp(expr.op, lowerExpr(expr.left), lowerExpr(expr.right)); | |
| case 'assign': | |
| return lowerAssign(expr); | |
| case 'cast': | |
| return convert(lowerExpr(expr.operand), scalarType(expr.to)); | |
| default: | |
| throw new LowerError(`Unsupported expression '${expr.kind}'`); | |
| } | |
| } | |
| function lowerLoopBody(stmts) { | |
| loopDepth += 1; | |
| try { | |
| return collect(() => scoped(() => lowerStmts(stmts))); | |
| } finally { | |
| loopDepth -= 1; | |
| } | |
| } | |
| function lowerStmt(stmt) { | |
| switch (stmt.kind) { | |
| case 'decl': { | |
| const type = scalarType(stmt.type); | |
| if (scopes.at(-1).has(stmt.name)) throw new LowerError(`'${stmt.name}' is already declared in this scope`); | |
| // Evaluated before the name is in scope, so `int x = x;` reads an outer x. | |
| // Without an initializer the local is zeroed so loop iterations stay deterministic. | |
| const init = stmt.init ? convert(lowerExpr(stmt.init), type) : constant(type, 0); | |
| const local = declareLocal(stmt.name, type); | |
| scopes.at(-1).set(stmt.name, local); | |
| emit(instr(Op.STORE, { name: local.name, value: init })); | |
| return; | |
| } | |
| case 'assign': | |
| lowerAssign(stmt); | |
| return; | |
| case 'expr': | |
| lowerExpr(stmt.value); | |
| return; | |
| case 'block': | |
| scoped(() => lowerStmts(stmt.body)); | |
| return; | |
| case 'if': { | |
| const cond = toBool(lowerExpr(stmt.cond)); | |
| const then = collect(() => scoped(() => lowerStmts(stmt.then))); | |
| const otherwise = collect(() => scoped(() => lowerStmts(stmt.otherwise))); | |
| emit(instr(Op.IF, { cond, then, otherwise })); | |
| return; | |
| } | |
| case 'while': { | |
| let cond = null; | |
| const head = collect(() => { cond = toBool(lowerExpr(stmt.cond)); }); | |
| const body = lowerLoopBody(stmt.body); | |
| emit(instr(Op.LOOP, { head, cond, body, step: [] })); | |
| return; | |
| } | |
| case 'for': | |
| scoped(() => { | |
| if (stmt.init) lowerStmt(stmt.init); | |
| let cond = null; | |
| const head = collect(() => { if (stmt.cond) cond = toBool(lowerExpr(stmt.cond)); }); | |
| const body = lowerLoopBody(stmt.body); | |
| const step = collect(() => { if (stmt.step) lowerExpr(stmt.step); }); | |
| emit(instr(Op.LOOP, { head, cond, body, step })); | |
| }); | |
| return; | |
| case 'break': | |
| case 'continue': | |
| if (loopDepth === 0) throw new LowerError(`'${stmt.kind}' outside of a loop`); | |
| emit(instr(stmt.kind === 'break' ? Op.BREAK : Op.CONTINUE)); | |
| return; | |
| case 'return': { | |
| if (isVoid) { | |
| if (stmt.value) throw new LowerError('A void kernel cannot return a value'); | |
| emit(instr(Op.RET, { value: null })); | |
| return; | |
| } | |
| if (!stmt.value) throw new LowerError(`Kernel '${ast.name}' must return a ${ast.returnType}`); | |
| emit(instr(Op.RET, { value: convert(lowerExpr(stmt.value), returnType) })); | |
| return; | |
| } | |
| default: | |
| throw new LowerError(`Unsupported statement '${stmt.kind}'`); | |
| } | |
| } | |
| function lowerStmts(stmts) { | |
| for (const stmt of stmts) lowerStmt(stmt); | |
| } | |
| const body = collect(() => scoped(() => lowerStmts(ast.body))); | |
| return kernel(ast.name, returnType, params, locals, body, [...builtins]); | |
| } | |