File size: 6,033 Bytes
3ce1a40 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 | /**
* Host runtime for compiled kernels.
*
* Owns one WebAssembly.Memory shared by every kernel it loads, hands out typed
* device buffers from it with a bump allocator, and validates arguments before a
* launch so a type mismatch fails in JavaScript instead of corrupting memory.
*/
import { compile } from '../../kernel/compiler/compile.js';
import { Type } from '../../kernel/ir/ir.js';
import { MEMORY_IMPORT, ELEMENT, elementTypeOf } from '../abi/abi.js';
const PAGE_BYTES = 65536;
const MAX_I32 = 0x7fffffff;
/** Offset 0 is never handed out, so a zero pointer is always a bug. */
const HEAP_BASE = 16;
export class RuntimeError extends Error {
constructor(message) {
super(message);
this.name = 'RuntimeError';
}
}
export class DeviceBuffer {
constructor(runtime, ptr, length, type) {
this.runtime = runtime;
this.ptr = ptr;
this.length = length;
this.type = type;
this.generation = runtime.generation;
}
get byteLength() {
return this.length * ELEMENT[this.type].bytes;
}
/** A live view over device memory. Invalidated when memory grows; prefer read(). */
view() {
this.runtime.checkBuffer(this);
return new ELEMENT[this.type].array(this.runtime.memory.buffer, this.ptr, this.length);
}
read() {
return this.view().slice();
}
write(data) {
if (data.length !== this.length) {
throw new RuntimeError(`Expected ${this.length} elements but got ${data.length}`);
}
this.view().set(data);
}
}
export class KernelRuntime {
constructor({ initialPages = 1, maximumPages = 16384 } = {}) {
this.memory = new WebAssembly.Memory({ initial: initialPages, maximum: maximumPages });
this.top = HEAP_BASE;
this.generation = 0;
}
/** Allocates `length` elements of IR type `type` ('i32', 'u32', 'f32', 'f64'), zeroed. */
alloc(type, length) {
const element = ELEMENT[type];
if (!element) throw new RuntimeError(`Unsupported buffer type '${type}'`);
if (!Number.isInteger(length) || length < 0) throw new RuntimeError(`Invalid buffer length ${length}`);
const ptr = Math.ceil(this.top / 8) * 8;
const end = ptr + length * element.bytes;
if (end > MAX_I32) throw new RuntimeError('Out of device memory');
const needed = Math.ceil(end / PAGE_BYTES) - this.memory.buffer.byteLength / PAGE_BYTES;
if (needed > 0) {
try {
this.memory.grow(needed);
} catch {
throw new RuntimeError('Out of device memory');
}
}
this.top = end;
const buffer = new DeviceBuffer(this, ptr, length, type);
// Memory is reused after reset(), so clear it rather than trusting it is zero.
new Uint8Array(this.memory.buffer, ptr, buffer.byteLength).fill(0);
return buffer;
}
/** Allocates a buffer matching a typed array and copies it in. */
upload(array) {
const type = elementTypeOf(array);
if (!type) throw new RuntimeError(`Unsupported array type ${array?.constructor?.name}`);
const buffer = this.alloc(type, array.length);
buffer.write(array);
return buffer;
}
/** Frees every buffer. Existing DeviceBuffers become invalid. */
reset() {
this.top = HEAP_BASE;
this.generation += 1;
}
checkBuffer(buffer) {
if (buffer.runtime !== this) throw new RuntimeError('Buffer belongs to a different runtime');
if (buffer.generation !== this.generation) throw new RuntimeError('Buffer was freed by reset()');
}
/** Compiles kernel source and instantiates it against this runtime's memory. */
async load(source) {
const { ir, bytes } = compile(source);
const imports = { [MEMORY_IMPORT.module]: { [MEMORY_IMPORT.name]: this.memory } };
const { instance } = await WebAssembly.instantiate(bytes, imports);
return new Kernel(this, ir, bytes, instance.exports[ir.name]);
}
}
export class Kernel {
constructor(runtime, ir, bytes, fn) {
this.runtime = runtime;
this.ir = ir;
this.bytes = bytes;
this.fn = fn;
}
get name() {
return this.ir.name;
}
/** Runs a void kernel on grid × block threads: launch({ grid, block }, ...args). */
launch({ grid, block }, ...args) {
if (this.ir.returnType !== null) throw new RuntimeError(`'${this.name}' returns a value; use call()`);
for (const [label, n] of [['grid', grid], ['block', block]]) {
if (!Number.isInteger(n) || n < 0 || n > MAX_I32) throw new RuntimeError(`Invalid ${label} size ${n}`);
}
if (grid * block > Number.MAX_SAFE_INTEGER) throw new RuntimeError('Launch is too large');
this.fn(...this.marshal(args), grid, block);
}
/** Calls a value-returning kernel once and returns its result. */
call(...args) {
if (this.ir.returnType === null) throw new RuntimeError(`'${this.name}' is a void kernel; use launch()`);
const result = this.fn(...this.marshal(args));
return this.ir.returnType === Type.U32 ? result >>> 0 : result;
}
marshal(args) {
const { params } = this.ir;
if (args.length !== params.length) {
throw new RuntimeError(`'${this.name}' takes ${params.length} arguments but got ${args.length}`);
}
return params.map((param, i) => {
const arg = args[i];
if (param.kind === 'buffer') {
if (!(arg instanceof DeviceBuffer)) throw new RuntimeError(`Argument '${param.name}' must be a DeviceBuffer`);
this.runtime.checkBuffer(arg);
if (arg.type !== param.type) {
throw new RuntimeError(`Argument '${param.name}' must be a ${param.type} buffer, not ${arg.type}`);
}
return arg.ptr;
}
if (typeof arg !== 'number') throw new RuntimeError(`Argument '${param.name}' must be a number`);
if (param.type === Type.I32 || param.type === Type.U32) {
const [min, max] = param.type === Type.I32 ? [-0x80000000, MAX_I32] : [0, 0xffffffff];
if (!Number.isInteger(arg) || arg < min || arg > max) {
throw new RuntimeError(`Argument '${param.name}' must be a ${param.type} but got ${arg}`);
}
}
return arg;
});
}
}
|