| """Explicit, reproducible native build: python -m julia.router.build.""" |
| import argparse |
| import os |
| import re |
| from pathlib import Path |
| import shutil |
| import subprocess |
| import tempfile |
|
|
|
|
| def build(output=None, bend=None, native_cpu=False): |
| source = Path(__file__).parent / 'native' / 'router.bend' |
| bend = bend or shutil.which('bend') or str(Path.home() / '.bend/bin/bend') |
| version = subprocess.check_output([bend, 'version'], text=True).strip() |
| if not re.search(r'(?<![0-9.])2\.0\.27(?![0-9.])', version): |
| raise RuntimeError(f'Native adapter requires Bend 2.0.27; found {version}') |
| subprocess.run([bend, str(source.with_name('PROOF.bend'))], check=True) |
| output = Path(output or Path(__file__).parent / 'build/libjulia_router.so').resolve() |
| output.parent.mkdir(parents=True, exist_ok=True) |
| with tempfile.TemporaryDirectory(dir=output.parent) as temp: |
| generated = Path(temp) / 'router.c' |
| library = Path(temp) / 'router.so' |
| subprocess.run([bend, str(source), '-o', str(generated)], check=True) |
| emitted = generated.read_text() |
| if not all(re.search(pattern, emitted) for pattern in ( |
| r'#define CUBE_T\s+128', r'#define LINE\s+16', |
| r'static u32\s+CUBE_LOG = 7;', r'#define RING_LOG\s+\(17 - CUBE_LOG\)')): |
| raise RuntimeError('CPU scheduling layout changed; refusing unsafe native build') |
| |
| for name, arity in [('WINNER', 1), ('SOFTMAX', 1), ('LAYERNORM', 3), ('NORM_ROWS', 3), ('DENSE_TILE', 2), ('DENSE_BATCH', 2), ('PACKED_GEMM', 9)]: |
| match = re.search(rf'#define FID_{name} (\d+)', emitted) |
| table = re.search(r'FID_ARITY_T\[\] = \{([^}]+)', emitted) |
| if not match or not table or int(table[1].split(',')[int(match[1])]) != arity: |
| raise RuntimeError(f'Unsupported generated Bend ABI for {name}') |
| winner = emitted.split('WL_CASE(FID_WINNER)', 1)[1].split('#endif', 1)[0] |
| if 'term_peek(e, _tree_0)' not in winner: |
| raise RuntimeError('WINNER ownership changed; refusing unsafe native build') |
| tile = emitted.split('WL_CASE(FID_DENSE_TILE)', 1)[1].split('#endif', 1)[0] |
| if 'term_peek(e, _weights_0)' not in tile or 'term_peek(e, _ws_0)' not in emitted: |
| raise RuntimeError('Tiled weights must be borrowed; refusing unsafe native build') |
| batch = emitted.split('WL_CASE(FID_DENSE_BATCH)', 1)[1].split('#endif', 1)[0] |
| if 'term_peek(e, _tiles_0)' not in batch: |
| raise RuntimeError('Dense input tiles must be borrowed; refusing unsafe native build') |
| packed = emitted.split('WL_CASE(FID_PACKED_DOT)', 1)[1].split('#endif', 1)[0] |
| if not all(fragment in packed for fragment in ('r0 = _x_0;', 'r1 = _w_0;', 'r2 = _o_0;', 'WL_RETN(3);')): |
| raise RuntimeError('Packed result ABI changed; refusing unsafe native build') |
| from .specialize import packed_cpu |
| generated.write_text(packed_cpu(emitted)) |
| subprocess.run([os.environ.get('CC', 'clang'), '-O3', '-ffp-contract=fast', |
| *(['-march=native'] if native_cpu else []), '-fPIC', '-shared', |
| '-pthread', str(generated), '-lm', '-ldl', '-o', str(library)], check=True) |
| os.replace(library, output) |
| return output |
|
|
|
|
| if __name__ == '__main__': |
| parser = argparse.ArgumentParser() |
| parser.add_argument('--output') |
| parser.add_argument('--bend') |
| parser.add_argument('--native-cpu', action='store_true', help='Optimize for this CPU; do not distribute to other CPUs') |
| args = parser.parse_args() |
| print(build(args.output, args.bend, args.native_cpu)) |
|
|
