LeelaChessZero · zamfofex · Oct 5, 2024
diff --git a/cross-files/wasm32-emscripten b/cross-files/wasm32-emscripten
@@ -0,0 +1,24 @@
+
+[host_machine]
+system = 'emscripten'
+cpu_family = 'wasm32'
+cpu = 'wasm32'
+endian = 'little'
+
+[binaries]
+c = 'emcc'
+cpp = 'em++'
+ar = 'emar'
+strip = 'emstrip'
+
+[built-in options]
+cpp_args = ['-fexceptions', '-msse', '-msse2', '-msse3', '-msimd128']
+cpp_link_args = [
+		'-fexceptions',
+		'-sASYNCIFY', '-sASYNCIFY_STACK_SIZE=65536',
+		'-sMODULARIZE', '-sEXPORT_ES6',
+		'-sDEFAULT_LIBRARY_FUNCS_TO_INCLUDE=$stringToNewUTF8',
+		'-sALLOW_MEMORY_GROWTH',
+		'-sWASM_BIGINT',
+		'-sENVIRONMENT=worker',
+	]
diff --git a/js/.gitignore b/js/.gitignore
@@ -0,0 +1,5 @@
+node_modules
+package-lock.json
+dist
+build
+lc0.tar
diff --git a/js/README.md b/js/README.md
@@ -0,0 +1,21 @@
+<!-- TODO: improve documentation! -->
+
+Leela Chess Zero (Wasm)
+===
+
+Lc0 compiled to WebAssembly (running on WebGPU when supported).
+
+compiling
+---
+
+- Install [Emscripten].
+- Install [esbuild].
+- Install [npm].
+- Install [Meson].
+- Run `npm install`
+- Run `npm run build` (or `./build.sh`)
+
+[Emscripten]: <https://emscripten.org/docs/getting_started/downloads.html>
+[esbuild]: <https://esbuild.github.io/getting-started/>
+[npm]: <https://docs.npmjs.com/cli/configuring-npm/install>
+[Meson]: <https://mesonbuild.com/Getting-meson.html>
diff --git a/js/build.sh b/js/build.sh
@@ -0,0 +1,29 @@
+#!/usr/bin/env sh
+set -ex
+meson setup --buildtype=release -Ddefault_library=static --prefer-static --cross-file=../cross-files/wasm32-emscripten -Dblas=false build .. || :
+meson compile -C build lc0
+esbuild --minify --outdir=dist --format=esm main.js worker.js build/lc0.js build/lc0.worker.mjs
+mv dist/build/lc0.worker.js dist/build/lc0.worker.mjs
+cp build/lc0.wasm dist/build
+cat > dist/package.json << END
+{
+	"name": "lc0",
+	"description": "Leela Chess Zero",
+	"version": "0.0.0.1",
+	"license": "GPL",
+	"homepage": "https://lczero.org",
+	"repository": {
+		"type": "git",
+		"url": "https://github.com/LeelaChessZero/lc0"
+	},
+	"main": "./main.js",
+	"exports": {
+		".": {
+			"import": "./main.js"
+		}
+	},
+	"dependencies": {
+		"onnxruntime-web": "1.19.2"
+	}
+}
+END
diff --git a/js/example/.gitignore b/js/example/.gitignore
@@ -0,0 +1,5 @@
+node_modules
+package-lock.json
+dist
+net.onnx
+.parcel-cache
diff --git a/js/example/README.md b/js/example/README.md
@@ -0,0 +1,20 @@
+<!-- TODO: improve documentation! -->
+
+Lc0 Web Example (Wasm)
+===
+
+First, convert an Lc0 network to ONNX format by using `lc0 leela2onnx` appropriately. Name the file `net.onnx` and place it on this directory.
+
+Then, run the following commands:
+
+- Install [Emscripten].
+- Install [esbuild].
+- Install [npm].
+- Install [Meson].
+- Run `npm install`
+- Run `npm run dev`
+
+[Emscripten]: <https://emscripten.org/docs/getting_started/downloads.html>
+[esbuild]: <https://esbuild.github.io/getting-started/>
+[npm]: <https://docs.npmjs.com/cli/configuring-npm/install>
+[Meson]: <https://mesonbuild.com/Getting-meson.html>
diff --git a/js/example/index.html b/js/example/index.html
@@ -0,0 +1,44 @@
+
+<!doctype html>
+<meta charset="utf-8">
+
+<title> Lc0 Web Test </title>
+
+<style> @import "@xterm/xterm/css/xterm.css"; </style>
+
+<script type="module">
+
+import {Lc0} from "lc0"
+import {Terminal} from "@xterm/xterm"
+
+let response = await fetch("net.onnx")
+let lc0 = Lc0(response.body)
+
+let terminal = new Terminal()
+terminal.open(document.querySelector("div"))
+
+let show = async out =>
+{
+	for await (let line of out) terminal.writeln(line)
+}
+
+show(lc0)
+show(lc0.stderr)
+
+let line = ""
+terminal.onData(data =>
+{
+	if (data === "\x7F") return
+	if (data === "\r") data = "\r\n"
+	terminal.write(data)
+	if (data === "\r\n") {
+		lc0.post(line)
+		line = ""
+		return
+	}
+	line += data
+})
+
+</script>
+
+<div></div>
diff --git a/js/example/package.json b/js/example/package.json
@@ -0,0 +1,11 @@
+{
+	"scripts": {
+		"pack": "cd .. && npm run pack",
+		"dev": "npm run pack && npm install && vite"
+	},
+	"dependencies": {
+		"@xterm/xterm": "5.5.0",
+		"lc0": "../lc0.tar",
+		"vite": "5.4.8"
+	}
+}
diff --git a/js/example/vite.config.js b/js/example/vite.config.js
@@ -0,0 +1,11 @@
+export default {
+	server: {
+		headers: {
+			"Cross-Origin-Embedder-Policy": "require-corp",
+			"Cross-Origin-Opener-Policy": "same-origin",
+		}
+	},
+	optimizeDeps: {
+		exclude: ["lc0"],
+	},
+}
diff --git a/js/main.js b/js/main.js
@@ -0,0 +1,69 @@
+let Stream = (worker, type, finished) =>
+{
+	let gotLine
+	let lines = []
+
+	worker.addEventListener("message", ({data}) =>
+	{
+		if (data.type !== type) return
+		if (!gotLine) {
+			lines.push(data.text)
+			return
+		}
+		gotLine(data.text)
+		gotLine = undefined
+	})
+
+	let next = () =>
+	{
+		let value
+		if (lines.length !== 0) return lines.shift()
+		else return new Promise(resolve => gotLine = resolve)
+	}
+
+	let it = {next: async () => finished() && lines.length === 0 ? {done: true} : {done: false, value: await next()}}
+	Object.freeze(it)
+
+	let peek = () => lines[0]
+
+	return {next, peek, [Symbol.asyncIterator]: () => it}
+}
+
+export let Lc0 = network =>
+{
+	let worker = new Worker(new URL("worker.js", import.meta.url), {type: "module"})
+
+	let commands = []
+	let post0 = command => commands.push(command)
+
+	worker.addEventListener("message", () =>
+	{
+		worker.postMessage({network}, [network])
+		for (let command of commands) worker.postMessage(command)
+		commands = undefined
+		post0 = command => worker.postMessage(command)
+	}, {once: true})
+
+	let post = command =>
+	{
+		if (finished) throw new Error("Cannot post command to finished Lc0")
+		post0(String(command))
+	}
+
+	let finished = false
+
+	// todo: this should send a message to the worker instead
+	// so that it can end its pthread workers too
+	let finish = () =>
+	{
+		finished = true
+		worker.terminate()
+	}
+
+	let stdout = Stream(worker, "stdout", () => finished)
+	let stderr = Stream(worker, "stderr", () => finished)
+
+	let lc0 = {post, finish, ...stdout, stderr, get finished() { return finished }}
+	Object.freeze(lc0)
+	return lc0
+}
diff --git a/js/package.json b/js/package.json
@@ -0,0 +1,7 @@
+{
+	"scripts": {
+		"build": "npm install && ./build.sh",
+		"pack": "npm run build && tar cf lc0.tar dist"
+	}
+}
+
diff --git a/js/worker.js b/js/worker.js
@@ -0,0 +1,105 @@
+import Module from "./build/lc0.js"
+import {InferenceSession, Tensor} from "onnxruntime-web/all"
+
+let gotLine
+let lines = []
+
+addEventListener("message", ({data}) =>
+{
+	if (typeof data !== "string") return
+	if (!gotLine) {
+		lines.push(data)
+		return
+	}
+	gotLine(data)
+	gotLine = undefined
+})
+
+let {data: {network}} = await new Promise(resolve =>
+{
+	postMessage({type: "ready"})
+	addEventListener("message", resolve, {once: true})
+})
+
+let session = await InferenceSession.create(await new Response(network).arrayBuffer(), {executionProviders: ["webgpu", "wasm"]})
+let map = new Map()
+
+let lc0web_get_line = async () =>
+{
+	if (lines.length !== 0) return lines.shift()
+	return new Promise(resolve => gotLine = resolve)
+}
+
+let lc0web_is_cpu = () => true
+
+let id = 0
+let lc0web_id = () =>
+{
+	let i = id++
+	map.set(i, {input: []})
+	return i
+}
+
+let lc0web_batch_size = id => map.get(id).input.length
+
+let lc0web_remove = id => map.delete(id)
+
+let lc0web_q_val = (id, sample) =>
+{
+	let [w, d, l] = map.get(id).output["/output/wdl"].cpuData
+	return w - l
+}
+
+let lc0web_d_val = (id, sample) =>
+{
+	let [w, d, l] = map.get(id).output["/output/wdl"].cpuData
+	return d
+}
+
+let lc0web_p_val = (id, sample, moveID) =>
+	map.get(id).output["/output/policy"].cpuData[sample * 1858 + moveID]
+
+let lc0web_m_val = (id, sample) =>
+	map.get(id).output["/output/mlh"].cpuData[sample]
+
+let lc0web_add_input = id => map.get(id).input.push([])
+
+let lc0web_add_plane = (id, index, mask, value) =>
+{
+	let array = map.get(id).input[index]
+	for (let i = 0 ; i < 64 ; i++) {
+		if (mask & 1n) array.push(value)
+		else array.push(0)
+		mask >>= 1n
+	}
+}
+
+let lc0web_compute = async id =>
+{
+	let value = map.get(id)
+	let {input} = value
+	let array = new Float32Array(input.flat(Infinity))
+	let tensor = new Tensor("float32", array, [input.length, 112, 8, 8])
+	value.output = await session.run({"/input/planes": tensor})
+}
+
+Object.assign(globalThis, {
+	lc0web_get_line,
+	lc0web_is_cpu,
+	lc0web_id,
+	lc0web_batch_size,
+	lc0web_remove,
+	lc0web_q_val,
+	lc0web_d_val,
+	lc0web_p_val,
+	lc0web_m_val,
+	lc0web_add_input,
+	lc0web_add_plane,
+	lc0web_compute,
+})
+
+Module({
+	arguments: ["--preload"],
+	print: text => postMessage({type: "stdout", text}),
+	printErr: text => postMessage({type: "stderr", text}),
+})
diff --git a/meson.build b/meson.build
@@ -632,6 +632,11 @@ if get_option('build_backends')
       has_backends = true
   endif
 
+  if host_machine.system() == 'emscripten'
+      files += 'src/neural/js/network_js.cc'
+      has_backends = true
+  endif
+
 endif # if get_option('build_backends')
 
 if not has_backends and get_option('lc0') and get_option('build_backends')