From 4a57c62e9e15e09f9f5f4775cf8e411b279d741f Mon Sep 17 00:00:00 2001 From: Alex Cameron Date: Sat, 29 Aug 2026 16:37:32 +0000 Subject: [PATCH] feat(runner): replace CLI with Python JSON protocol --- AGENTS.md | 4 +- README.md | 23 +- gpu_test/conftest.py | 141 ++++---- gpu_test/runner.py | 327 ++++++++++++++++++ gpu_test/test_runner.py | 178 ++++++++++ pyproject.toml | 1 + tools/CMakeLists.txt | 1 - tools/warpforth-runner/CMakeLists.txt | 9 - tools/warpforth-runner/warpforth-runner.cpp | 347 -------------------- uv.lock | 49 +++ 10 files changed, 659 insertions(+), 421 deletions(-) create mode 100644 gpu_test/runner.py create mode 100644 gpu_test/test_runner.py delete mode 100644 tools/warpforth-runner/CMakeLists.txt delete mode 100644 tools/warpforth-runner/warpforth-runner.cpp diff --git a/AGENTS.md b/AGENTS.md index fdf579f..cba4eb4 100644 --- a/AGENTS.md +++ b/AGENTS.md @@ -38,7 +38,7 @@ CMake. **Tools:** - `tools/warpforth-translate/warpforth-translate.cpp` - Translation tool entry point - `tools/warpforth-opt/warpforth-opt.cpp` - Optimization tool entry point -- `tools/warpforth-runner/warpforth-runner.cpp` - PTX execution tool for GPU kernels +- `gpu_test/runner.py` - JSON PTX execution tool for GPU kernels ## Tools Usage @@ -56,7 +56,7 @@ CMake. ./build/bin/warpforth-translate --mlir-to-ptx > kernel.ptx # Execute PTX on GPU -./warpforth-runner kernel.ptx --param i64[]:1,2,3 --param i64:42 --output-param 0 --output-count 3 +cat request.json | python3 gpu_test/runner.py ``` ## Adding New Operations diff --git a/README.md b/README.md index 846b739..b0d62fc 100644 --- a/README.md +++ b/README.md @@ -59,12 +59,21 @@ Compile to PTX: Test on a GPU (A is 2x4 row-major, B is 4x3 row-major, C is 2x3 output): ```bash -./build/bin/warpforth-runner matmul.ptx \ - --param 'i64[]:1,2,3,4,5,6,7,8' \ - --param 'i64[]:1,2,3,4,5,6,7,8,9,10,11,12' \ - --param 'i64[]:0,0,0,0,0,0' \ - --grid 6,1,1 --block 1,1,1 \ - --output-param 2 --output-count 6 +cat > request.json < None: ) self._attach_ssh_key() self._wait_for_sshd() - self._compile_runner() + self._install_runner() return time.sleep(POLL_INTERVAL_S) @@ -356,14 +358,15 @@ def _wait_for_sshd(self) -> None: ) raise TimeoutError(msg) - def _compile_runner(self) -> None: - """Upload warpforth-runner.cpp and compile it on the remote host.""" - self.scp_upload(RUNNER_SRC, f"{REMOTE_TMP}/warpforth-runner.cpp") - nvcc_cmd = ( - f"nvcc -o {REMOTE_TMP}/warpforth-runner" - f" {REMOTE_TMP}/warpforth-runner.cpp -lcuda -std=c++17" + def _install_runner(self) -> None: + """Install cuda-python and upload the runner to the remote host.""" + install_cmd = ( + "apt-get update -qq && " + "DEBIAN_FRONTEND=noninteractive apt-get install -y -qq python3-pip && " + "python3 -m pip install -q 'cuda-python>=12.4,<13'" ) - self.ssh_run(nvcc_cmd, timeout=60) + self.ssh_run(install_cmd, timeout=180) + self.scp_upload(RUNNER_SRC, f"{REMOTE_TMP}/warpforth-runner.py") def _record_labeled_instance_ids( self, @@ -484,10 +487,11 @@ def _ssh_cmd(self) -> list[str]: f"root@{self.ssh_host}", ] - def ssh_run(self, cmd: str, *, timeout: int = 120) -> str: + def ssh_run(self, cmd: str, *, input_text: str | None = None, timeout: int = 120) -> str: """Execute a command on the remote instance via SSH.""" result = subprocess.run( [*self._ssh_cmd(), cmd], + input=input_text, capture_output=True, text=True, timeout=timeout, @@ -620,7 +624,22 @@ def run( params = params or {} - # Validate output_param + self._validate_output(decls, output_param) + ptx = self.compiler.compile_source(forth_source) + request = self._build_request( + ptx, + kernel_name, + decls, + params, + grid, + block, + output_param, + output_count, + ) + return self._run_request(request) + + @staticmethod + def _validate_output(decls: list[ParamDecl], output_param: int) -> None: if output_param < 0 or output_param >= len(decls): msg = f"output_param {output_param} out of range (have {len(decls)} params)" raise ValueError(msg) @@ -629,27 +648,31 @@ def run( msg = f"output_param {output_param} ('{name}') is a scalar and cannot be read back" raise ValueError(msg) - # Compile locally - ptx = self.compiler.compile_source(forth_source) - - # Write PTX to temp file and upload - with tempfile.NamedTemporaryFile(mode="w", suffix=".ptx", delete=False) as f: - f.write(ptx) - ptx_path = Path(f.name) - - try: - self.session.scp_upload(ptx_path, f"{REMOTE_TMP}/kernel.ptx") - finally: - ptx_path.unlink() - - # Build remote command - cmd_parts = [ - f"{REMOTE_TMP}/warpforth-runner", - f"{REMOTE_TMP}/kernel.ptx", - "--kernel", - kernel_name, - ] - + @staticmethod + def _build_request( + ptx: str, + kernel_name: str, + decls: list[ParamDecl], + params: dict[str, list[int] | list[float] | int | float], + grid: tuple[int, int, int], + block: tuple[int, int, int], + output_param: int, + output_count: int | None, + ) -> dict[str, object]: + request: dict[str, object] = { + "ptx_base64": base64.b64encode(ptx.encode()).decode(), + "kernel": kernel_name, + "grid": list(grid), + "block": list(block), + "params": [], + "outputs": [ + { + "param": output_param, + **({"count": output_count} if output_count is not None else {}), + } + ], + } + request_params: list[dict[str, object]] = [] for decl in decls: if decl.is_array: values = params.get(decl.name, []) @@ -660,35 +683,43 @@ def run( buf = [zero] * decl.size for i, v in enumerate(values): buf[i] = v - cmd_parts.extend(["--param", f"{decl.base_type}[]:{','.join(str(v) for v in buf)}"]) + request_params.append({"type": f"{decl.base_type}[]", "values": buf}) else: value = params.get(decl.name, 0.0 if decl.base_type == "f64" else 0) if isinstance(value, list): msg = f"Scalar param '{decl.name}' expects a scalar, got list" raise TypeError(msg) - cmd_parts.extend(["--param", f"{decl.base_type}:{value}"]) - - cmd_parts.extend( - [ - "--grid", - f"{grid[0]},{grid[1]},{grid[2]}", - "--block", - f"{block[0]},{block[1]},{block[2]}", - "--output-param", - str(output_param), - ] + request_params.append({"type": decl.base_type, "value": value}) + request["params"] = request_params + return request + + def _run_request(self, request: dict[str, object]) -> list[int] | list[float]: + stdout = self.session.ssh_run( + f"python3 {REMOTE_TMP}/warpforth-runner.py", + input_text=json.dumps(request), + timeout=120, ) - - if output_count is not None: - cmd_parts.extend(["--output-count", str(output_count)]) - - cmd = " ".join(cmd_parts) - stdout = self.session.ssh_run(cmd, timeout=120) - - # Parse CSV output — type depends on the output param - out_type = decls[output_param].base_type - parse = float if out_type == "f64" else int - return [parse(v) for v in stdout.strip().split(",")] + try: + response = json.loads(stdout) + except json.JSONDecodeError as exc: + msg = f"warpforth-runner returned invalid JSON: {stdout!r}" + raise RuntimeError(msg) from exc + if response.get("status") != "ok": + msg = f"warpforth-runner failed: {response.get('error', 'unknown error')}" + raise RuntimeError(msg) + outputs = response.get("outputs") + if not isinstance(outputs, list) or len(outputs) != 1: + msg = "warpforth-runner returned an invalid outputs array" + raise RuntimeError(msg) + output = outputs[0] + if not isinstance(output, dict): + msg = "warpforth-runner output must be an object" + raise TypeError(msg) + values = output.get("values") + if not isinstance(values, list): + msg = "warpforth-runner output has no values array" + raise TypeError(msg) + return values # --- Fixtures --- diff --git a/gpu_test/runner.py b/gpu_test/runner.py new file mode 100644 index 0000000..89187ad --- /dev/null +++ b/gpu_test/runner.py @@ -0,0 +1,327 @@ +"""Execute a PTX kernel from a JSON request on stdin.""" + +from __future__ import annotations + +import base64 +import binascii +import ctypes +import importlib +import json +import sys +from dataclasses import dataclass +from typing import TYPE_CHECKING, cast + +if TYPE_CHECKING: + from collections.abc import Callable + from types import SimpleNamespace + +DIMENSION_COUNT = 3 +SINGLE_VALUE_RESULT_LENGTH = 2 +SUPPORTED_TYPES = ("i64", "f64", "i64[]", "f64[]") + + +class RequestError(ValueError): + """Raised when the input request does not match the runner protocol.""" + + +class CudaError(RuntimeError): + """Raised when a CUDA Driver API call fails.""" + + +@dataclass +class Param: + """A kernel parameter and any associated host/device storage.""" + + type_name: str + host_value: object + device_ptr: object | None = None + + +@dataclass +class RunnerRequest: + """Validated runner input.""" + + ptx: bytes + kernel: str + grid: tuple[int, int, int] + block: tuple[int, int, int] + params: list[Param] + outputs: list[tuple[int, int]] + + +def _require_object(value: object, name: str) -> dict[str, object]: + if not isinstance(value, dict): + message = f"{name} must be an object" + raise RequestError(message) + return value + + +def _parse_dims(request: dict[str, object], name: str) -> tuple[int, int, int]: + value = request.get(name, [1, 1, 1]) + if ( + not isinstance(value, list) + or len(value) != DIMENSION_COUNT + or any(not isinstance(dim, int) or isinstance(dim, bool) or dim <= 0 for dim in value) + ): + message = f"{name} must contain three positive integers" + raise RequestError(message) + return value[0], value[1], value[2] + + +def _parse_params(request: dict[str, object]) -> list[Param]: + specs = request.get("params") + if not isinstance(specs, list) or not specs: + message = "params must be a non-empty array" + raise RequestError(message) + return [_parse_param(raw_spec, index) for index, raw_spec in enumerate(specs)] + + +def _parse_param(raw_spec: object, index: int) -> Param: + spec = _require_object(raw_spec, f"params[{index}]") + type_name = spec.get("type") + if type_name not in SUPPORTED_TYPES: + message = f"params[{index}].type is unsupported: {type_name!r}" + raise RequestError(message) + + type_name = cast("str", type_name) + ctype = ctypes.c_double if type_name.startswith("f64") else ctypes.c_int64 + if type_name.endswith("[]"): + values = spec.get("values") + if not isinstance(values, list) or not values: + message = f"params[{index}].values must be a non-empty array" + raise RequestError(message) + try: + return Param(type_name, (ctype * len(values))(*values)) + except (OverflowError, TypeError, ValueError) as exc: + message = f"params[{index}].values do not match {type_name}" + raise RequestError(message) from exc + + if "value" not in spec: + message = f"params[{index}].value is required" + raise RequestError(message) + try: + return Param(type_name, ctype(spec["value"])) + except (OverflowError, TypeError, ValueError) as exc: + message = f"params[{index}].value does not match {type_name}" + raise RequestError(message) from exc + + +def _parse_outputs(request: dict[str, object], params: list[Param]) -> list[tuple[int, int]]: + specs = request.get("outputs") + if not isinstance(specs, list) or not specs: + message = "outputs must be a non-empty array" + raise RequestError(message) + return [_parse_output(raw_spec, index, params) for index, raw_spec in enumerate(specs)] + + +def _parse_output(raw_spec: object, index: int, params: list[Param]) -> tuple[int, int]: + spec = _require_object(raw_spec, f"outputs[{index}]") + param_index = spec.get("param") + if ( + not isinstance(param_index, int) + or isinstance(param_index, bool) + or param_index < 0 + or param_index >= len(params) + ): + message = f"outputs[{index}].param is out of range" + raise RequestError(message) + + param = params[param_index] + if not param.type_name.endswith("[]"): + message = f"outputs[{index}].param refers to a scalar" + raise RequestError(message) + size = len(cast("ctypes.Array[object]", param.host_value)) + count = spec.get("count", size) + if not isinstance(count, int) or isinstance(count, bool) or count < 0 or count > size: + message = f"outputs[{index}].count must be between 0 and {size}" + raise RequestError(message) + return param_index, count + + +def _parse_request(value: object) -> RunnerRequest: + request = _require_object(value, "request") + encoded_ptx = request.get("ptx_base64") + if not isinstance(encoded_ptx, str) or not encoded_ptx: + message = "ptx_base64 must be a non-empty string" + raise RequestError(message) + try: + ptx = base64.b64decode(encoded_ptx, validate=True) + except (binascii.Error, ValueError) as exc: + message = "ptx_base64 is not valid base64" + raise RequestError(message) from exc + if not ptx: + message = "ptx_base64 decodes to empty PTX" + raise RequestError(message) + + kernel = request.get("kernel") + if not isinstance(kernel, str) or not kernel: + message = "kernel must be a non-empty string" + raise RequestError(message) + params = _parse_params(request) + return RunnerRequest( + ptx=ptx, + kernel=kernel, + grid=_parse_dims(request, "grid"), + block=_parse_dims(request, "block"), + params=params, + outputs=_parse_outputs(request, params), + ) + + +def _load_cuda() -> object: + try: + return importlib.import_module("cuda.bindings.driver") + except ImportError: + try: + return importlib.import_module("cuda.cuda") + except ImportError as exc: + message = "cuda-python is not installed" + raise RuntimeError(message) from exc + + +def _cuda_check(cuda: object, result: tuple[object, ...]) -> object: + cuda_module = cast("SimpleNamespace", cuda) + success = cuda_module.CUresult.CUDA_SUCCESS + error = result[0] + if error != success: + _, name = _cuda_result(cuda, "cuGetErrorName", error) + _, description = _cuda_result(cuda, "cuGetErrorString", error) + message = f"{_decode(name)}: {_decode(description)}" + raise CudaError(message) + if len(result) == 1: + return None + if len(result) == SINGLE_VALUE_RESULT_LENGTH: + return result[1] + return result[1:] + + +def _cuda_result(cuda: object, name: str, *args: object) -> tuple[object, ...]: + function = cast("Callable[..., tuple[object, ...]]", getattr(cuda, name)) + return function(*args) + + +def _cuda_call(cuda: object, name: str, *args: object) -> object: + return _cuda_check(cuda, _cuda_result(cuda, name, *args)) + + +def _decode(value: object) -> str: + return value.decode("utf-8", "replace") if isinstance(value, bytes) else str(value) + + +class CudaSession: + """Own the CUDA resources used to execute one request.""" + + def __init__(self, cuda: object, request: RunnerRequest) -> None: + self.cuda = cuda + self.request = request + self.context: object | None = None + self.module: object | None = None + + def load_kernel(self) -> object: + _cuda_call(self.cuda, "cuInit", 0) + device = _cuda_call(self.cuda, "cuDeviceGet", 0) + try: + self.context = _cuda_call(self.cuda, "cuCtxCreate", None, 0, device) + except TypeError: + self.context = _cuda_call(self.cuda, "cuCtxCreate", 0, device) + ptx_buffer = ctypes.create_string_buffer(self.request.ptx) + self.module = _cuda_call(self.cuda, "cuModuleLoadData", ptx_buffer) + return _cuda_call( + self.cuda, + "cuModuleGetFunction", + self.module, + self.request.kernel.encode(), + ) + + def allocate_params(self) -> tuple[tuple[object, ...], tuple[object, ...]]: + values: list[object] = [] + types: list[object] = [] + for param in self.request.params: + value, value_type = self._allocate_param(param) + values.append(value) + types.append(value_type) + return tuple(values), tuple(types) + + def _allocate_param(self, param: Param) -> tuple[object, object]: + if not param.type_name.endswith("[]"): + scalar = cast("ctypes.c_int64 | ctypes.c_double", param.host_value) + return scalar.value, type(scalar) + size_bytes = ctypes.sizeof(param.host_value) + param.device_ptr = _cuda_call(self.cuda, "cuMemAlloc", size_bytes) + _cuda_call( + self.cuda, + "cuMemcpyHtoD", + param.device_ptr, + ctypes.addressof(param.host_value), + size_bytes, + ) + return param.device_ptr, ctypes.c_void_p + + def read_outputs(self) -> list[dict[str, object]]: + outputs: list[dict[str, object]] = [] + for param_index, count in self.request.outputs: + param = self.request.params[param_index] + output_buffer = type(param.host_value)() + _cuda_call( + self.cuda, + "cuMemcpyDtoH", + ctypes.addressof(output_buffer), + param.device_ptr, + ctypes.sizeof(output_buffer), + ) + outputs.append( + { + "param": param_index, + "type": param.type_name, + "values": list(output_buffer)[:count], + } + ) + return outputs + + def cleanup(self) -> None: + for param in self.request.params: + if param.device_ptr is not None: + _cuda_call(self.cuda, "cuMemFree", param.device_ptr) + if self.module is not None: + _cuda_call(self.cuda, "cuModuleUnload", self.module) + if self.context is not None: + _cuda_call(self.cuda, "cuCtxDestroy", self.context) + + +def execute(value: object, cuda: object | None = None) -> list[dict[str, object]]: + """Validate and execute one runner request.""" + request = _parse_request(value) + session = CudaSession(cuda or _load_cuda(), request) + try: + function = session.load_kernel() + kernel_params = session.allocate_params() + _cuda_call( + session.cuda, + "cuLaunchKernel", + function, + *request.grid, + *request.block, + 0, + 0, + kernel_params, + 0, + ) + _cuda_call(session.cuda, "cuCtxSynchronize") + return session.read_outputs() + finally: + session.cleanup() + + +def main() -> int: + """Read one request and write one response.""" + try: + response = {"status": "ok", "outputs": execute(json.load(sys.stdin))} + except (ImportError, OSError, RuntimeError, TypeError, ValueError) as exc: + response = {"status": "error", "error": str(exc)} + json.dump(response, sys.stdout, allow_nan=False, separators=(",", ":")) + sys.stdout.write("\n") + return 0 + + +if __name__ == "__main__": + raise SystemExit(main()) diff --git a/gpu_test/test_runner.py b/gpu_test/test_runner.py new file mode 100644 index 0000000..ad3a95d --- /dev/null +++ b/gpu_test/test_runner.py @@ -0,0 +1,178 @@ +"""Unit tests for the JSON GPU runner protocol.""" + +from __future__ import annotations + +import base64 +import ctypes +import io +import json +import sys +from types import SimpleNamespace +from typing import TYPE_CHECKING, cast + +from gpu_test import runner +from gpu_test.conftest import KernelRunner + +if TYPE_CHECKING: + import pytest + + +class FakeCuda: + """Minimal successful CUDA Driver API implementation backed by host bytes.""" + + class CUresult: + CUDA_SUCCESS = 0 + + def __init__(self) -> None: + self.next_pointer = 100 + self.memory: dict[int, bytes] = {} + self.launch_params: object = None + self.CUresult = SimpleNamespace(CUDA_SUCCESS=0) + methods = { + "cuInit": self._init, + "cuDeviceGet": self._device_get, + "cuCtxCreate": self._context_create, + "cuModuleLoadData": self._module_load_data, + "cuModuleGetFunction": self._module_get_function, + "cuMemAlloc": self._mem_alloc, + "cuMemcpyHtoD": self._memcpy_host_to_device, + "cuLaunchKernel": self._launch_kernel, + "cuCtxSynchronize": self._context_synchronize, + "cuMemcpyDtoH": self._memcpy_device_to_host, + "cuMemFree": self._mem_free, + "cuModuleUnload": self._module_unload, + "cuCtxDestroy": self._context_destroy, + } + for name, method in methods.items(): + setattr(self, name, method) + + def _init(self, _flags: int) -> tuple[int]: + return (0,) + + def _device_get(self, _ordinal: int) -> tuple[int, int]: + return 0, 1 + + def _context_create(self, *_args: object) -> tuple[int, int]: + return 0, 2 + + def _module_load_data(self, _ptx: object) -> tuple[int, int]: + return 0, 3 + + def _module_get_function(self, _module: int, _name: bytes) -> tuple[int, int]: + return 0, 4 + + def _mem_alloc(self, size: int) -> tuple[int, int]: + pointer = self.next_pointer + self.next_pointer += 1 + self.memory[pointer] = bytes(size) + return 0, pointer + + def _memcpy_host_to_device(self, pointer: int, source: int, size: int) -> tuple[int]: + self.memory[pointer] = ctypes.string_at(source, size) + return (0,) + + def _launch_kernel(self, *_args: object) -> tuple[int]: + self.launch_params = _args[-2] + return (0,) + + def _context_synchronize(self) -> tuple[int]: + return (0,) + + def _memcpy_device_to_host(self, destination: int, pointer: int, size: int) -> tuple[int]: + ctypes.memmove(destination, self.memory[pointer], size) + return (0,) + + def _mem_free(self, _pointer: int) -> tuple[int]: + return (0,) + + def _module_unload(self, _module: int) -> tuple[int]: + return (0,) + + def _context_destroy(self, _context: int) -> tuple[int]: + return (0,) + + +def test_runner_supports_typed_params_and_multiple_outputs() -> None: + cuda = FakeCuda() + request = { + "ptx_base64": base64.b64encode(b"// PTX").decode(), + "kernel": "main", + "grid": [2, 1, 1], + "block": [32, 1, 1], + "params": [ + {"type": "i64[]", "values": [1, 2, 3]}, + {"type": "f64", "value": 3.5}, + {"type": "f64[]", "values": [4.25, 5.5]}, + ], + "outputs": [{"param": 0, "count": 2}, {"param": 2}], + } + + assert runner.execute(request, cuda) == [ + {"param": 0, "type": "i64[]", "values": [1, 2]}, + {"param": 2, "type": "f64[]", "values": [4.25, 5.5]}, + ] + launch_params = cast("tuple[tuple[object, ...], tuple[object, ...]]", cuda.launch_params) + assert launch_params[0][1] == 3.5 + assert launch_params[1][1] is ctypes.c_double + + +def test_runner_reports_invalid_requests_as_json(monkeypatch: pytest.MonkeyPatch) -> None: + output = io.StringIO() + monkeypatch.setattr(sys, "stdin", io.StringIO("{}")) + monkeypatch.setattr(sys, "stdout", output) + + assert runner.main() == 0 + assert json.loads(output.getvalue()) == { + "status": "error", + "error": "ptx_base64 must be a non-empty string", + } + + +class FakeCompiler: + def compile_source(self, _source: str) -> str: + return "// compiled PTX" + + +class FakeSession: + def __init__(self) -> None: + self.request: dict[str, object] | None = None + + def ssh_run(self, _command: str, *, input_text: str, timeout: int) -> str: + assert timeout == 120 + self.request = json.loads(input_text) + return json.dumps( + { + "status": "ok", + "outputs": [{"param": 0, "type": "i64[]", "values": [7, 8]}], + } + ) + + +def test_kernel_runner_sends_json_request() -> None: + session = FakeSession() + runner = KernelRunner(session, FakeCompiler()) + source = """\\! kernel main +\\! param output i64[3] +\\! param scale f64 +""" + + result = runner.run( + source, + params={"output": [7, 8], "scale": 2.5}, + grid=(3, 2, 1), + block=(16, 1, 1), + output_count=2, + ) + + assert result == [7, 8] + assert session.request == { + "ptx_base64": base64.b64encode(b"// compiled PTX").decode(), + "kernel": "main", + "grid": [3, 2, 1], + "block": [16, 1, 1], + "params": [ + {"type": "i64[]", "values": [7, 8, 0]}, + {"type": "f64", "value": 2.5}, + ], + "outputs": [{"param": 0, "count": 2}], + } diff --git a/pyproject.toml b/pyproject.toml index f416b7e..823a434 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -3,6 +3,7 @@ name = "warpforth" version = "0.1.0" requires-python = ">=3.11" dependencies = [ + "cuda-python>=12.4,<13", "lit>=18.1.0", "numpy", "pytest", diff --git a/tools/CMakeLists.txt b/tools/CMakeLists.txt index 7d4740f..f9e8857 100644 --- a/tools/CMakeLists.txt +++ b/tools/CMakeLists.txt @@ -1,4 +1,3 @@ add_subdirectory(warpforth-translate) add_subdirectory(warpforth-opt) add_subdirectory(warpforthc) -add_subdirectory(warpforth-runner) diff --git a/tools/warpforth-runner/CMakeLists.txt b/tools/warpforth-runner/CMakeLists.txt deleted file mode 100644 index a2786f2..0000000 --- a/tools/warpforth-runner/CMakeLists.txt +++ /dev/null @@ -1,9 +0,0 @@ -find_package(CUDAToolkit QUIET) - -if(CUDAToolkit_FOUND) - add_executable(warpforth-runner warpforth-runner.cpp) - target_link_libraries(warpforth-runner PRIVATE CUDA::cuda_driver) - target_compile_features(warpforth-runner PRIVATE cxx_std_17) -else() - message(STATUS "CUDAToolkit not found — skipping warpforth-runner") -endif() diff --git a/tools/warpforth-runner/warpforth-runner.cpp b/tools/warpforth-runner/warpforth-runner.cpp deleted file mode 100644 index 67e1d3f..0000000 --- a/tools/warpforth-runner/warpforth-runner.cpp +++ /dev/null @@ -1,347 +0,0 @@ -/// warpforth-runner: Execute PTX kernels via the CUDA Driver API. -/// -/// Single-file C++ program designed to be uploaded and compiled on a remote -/// GPU host with `nvcc -o warpforth-runner warpforth-runner.cpp -lcuda -/// -std=c++17`. -/// -/// Usage: -/// warpforth-runner kernel.ptx --param i64[]:1,2,3 --param f64:3.14 \ -/// --grid 4,1,1 --block 64,1,1 --kernel main \ -/// --output-param 0 --output-count 3 - -#include - -#include -#include -#include -#include -#include -#include -#include -#include -#include -#include -#include -#include -#include - -#define CHECK_CU(call) \ - do { \ - CUresult err = (call); \ - if (err != CUDA_SUCCESS) { \ - const char *errStr = nullptr; \ - cuGetErrorString(err, &errStr); \ - std::cerr << "CUDA error at " << __FILE__ << ":" << __LINE__ << ": " \ - << (errStr ? errStr : "unknown") << "\n"; \ - exit(1); \ - } \ - } while (0) - -template struct ArrayParam { - std::vector values; - CUdeviceptr devicePtr = 0; -}; - -template struct ScalarParam { - T value; -}; - -using Param = std::variant, ArrayParam, - ScalarParam, ScalarParam>; - -template static void allocDevice(ArrayParam &arr) { - size_t bytes = arr.values.size() * sizeof(T); - CHECK_CU(cuMemAlloc(&arr.devicePtr, bytes)); - CHECK_CU(cuMemcpyHtoD(arr.devicePtr, arr.values.data(), bytes)); -} - -template -static void printOutput(ArrayParam &arr, size_t count) { - std::vector output(arr.values.size()); - CHECK_CU(cuMemcpyDtoH(output.data(), arr.devicePtr, - arr.values.size() * sizeof(T))); - for (size_t i = 0; i < count; ++i) { - if (i > 0) - std::cout << ","; - if constexpr (std::is_floating_point_v) - std::cout << std::setprecision(17) << output[i]; - else - std::cout << output[i]; - } - std::cout << "\n"; -} - -static void *kernelArgPtr(Param &p) { - if (auto *a = std::get_if>(&p)) - return &a->devicePtr; - if (auto *a = std::get_if>(&p)) - return &a->devicePtr; - if (auto *s = std::get_if>(&p)) - return &s->value; - return &std::get>(p).value; -} - -static bool isScalar(const Param &p) { - return std::holds_alternative>(p) || - std::holds_alternative>(p); -} - -struct Dims { - unsigned x = 1, y = 1, z = 1; -}; - -static int parseIntArg(std::string_view s, std::string_view optName) { - int value = 0; - auto [ptr, ec] = std::from_chars(s.data(), s.data() + s.size(), value); - if (ec != std::errc{} || ptr != s.data() + s.size()) { - std::cerr << "Error: " << optName << " expects an integer, got: " << s - << "\n"; - exit(1); - } - return value; -} - -static Dims parseDims(std::string_view s) { - Dims d; - const char *p = s.data(); - const char *end = s.data() + s.size(); - - auto dimsErr = [&]() { - std::cerr << "Error: expected 3 comma-separated values, got: " << s << "\n"; - exit(1); - }; - - auto [p1, ec1] = std::from_chars(p, end, d.x); - if (ec1 != std::errc{} || p1 == end || *p1 != ',') - dimsErr(); - - auto [p2, ec2] = std::from_chars(p1 + 1, end, d.y); - if (ec2 != std::errc{} || p2 == end || *p2 != ',') - dimsErr(); - - auto [p3, ec3] = std::from_chars(p2 + 1, end, d.z); - if (ec3 != std::errc{} || p3 != end) - dimsErr(); - - return d; -} - -static Param parseParam(std::string_view s) { - std::string input(s); - - auto colonPos = input.find(':'); - if (colonPos == std::string::npos) { - std::cerr << "Error: --param requires type prefix (e.g. i64:42 or " - "f64[]:1.0,2.0), got: " - << s << "\n"; - exit(1); - } - - std::string typePrefix = input.substr(0, colonPos); - std::string valueStr = input.substr(colonPos + 1); - - if (valueStr.empty()) { - std::cerr << "Error: --param requires at least one value, got: " << s - << "\n"; - exit(1); - } - - // Parse comma-separated values into a typed vector - auto parseValues = [&](auto convert) { - using T = decltype(convert(std::string{})); - std::vector vals; - std::istringstream iss(valueStr); - std::string token; - while (std::getline(iss, token, ',')) - vals.push_back(convert(token)); - return vals; - }; - - auto toI64 = [&](const std::string &tok) -> int64_t { - try { - return std::stoll(tok); - } catch (const std::exception &) { - std::cerr << "Error: invalid integer value '" << tok << "' in --param " - << s << "\n"; - exit(1); - } - }; - auto toF64 = [&](const std::string &tok) -> double { - try { - return std::stod(tok); - } catch (const std::exception &) { - std::cerr << "Error: invalid float value '" << tok << "' in --param " << s - << "\n"; - exit(1); - } - }; - - if (typePrefix == "i64[]") - return Param{ArrayParam{parseValues(toI64)}}; - if (typePrefix == "f64[]") - return Param{ArrayParam{parseValues(toF64)}}; - - // Scalars — must be exactly one value - if (valueStr.find(',') != std::string::npos) { - std::cerr << "Error: scalar param expects exactly one value, got: " << s - << "\n"; - exit(1); - } - - if (typePrefix == "i64") - return Param{ScalarParam{toI64(valueStr)}}; - if (typePrefix == "f64") - return Param{ScalarParam{toF64(valueStr)}}; - - std::cerr << "Error: unsupported param type '" << typePrefix - << "' (expected i64, i64[], f64, or f64[]), got: " << s << "\n"; - exit(1); -} - -static std::string readFile(std::string_view path) { - std::ifstream f(std::string(path), std::ios::binary); - if (!f) { - std::cerr << "Error: cannot open " << path << "\n"; - exit(1); - } - std::ostringstream ss; - ss << f.rdbuf(); - return ss.str(); -} - -int main(int argc, char **argv) { - const char *ptxFile = nullptr; - const char *kernelName = nullptr; - std::vector params; - Dims grid, block; - int outputParam = 0; - int outputCount = -1; // -1 = all - - // Parse arguments - for (int i = 1; i < argc; ++i) { - std::string_view arg = argv[i]; - auto needsValue = [&](std::string_view opt) { - if (++i >= argc) { - std::cerr << "Error: " << opt << " requires a value\n"; - exit(1); - } - }; - if (arg == "--param") { - needsValue("--param"); - params.push_back(parseParam(argv[i])); - } else if (arg == "--grid") { - needsValue("--grid"); - grid = parseDims(argv[i]); - } else if (arg == "--block") { - needsValue("--block"); - block = parseDims(argv[i]); - } else if (arg == "--output-param") { - needsValue("--output-param"); - outputParam = parseIntArg(argv[i], "--output-param"); - } else if (arg == "--output-count") { - needsValue("--output-count"); - outputCount = parseIntArg(argv[i], "--output-count"); - } else if (arg == "--kernel") { - needsValue("--kernel"); - kernelName = argv[i]; - } else if (arg[0] == '-') { - std::cerr << "Error: unknown option " << arg << "\n"; - exit(1); - } else { - ptxFile = argv[i]; - } - } - - if (!ptxFile) { - std::cerr << "Usage: warpforth-runner kernel.ptx --kernel NAME " - "[--param i64[]:V,...] [--param f64[]:V,...] " - "[--param i64:V] [--param f64:V] [--grid X,Y,Z] " - "[--block X,Y,Z] [--output-param N] [--output-count N]\n"; - return 1; - } - - if (!kernelName) { - std::cerr << "Error: --kernel NAME is required\n"; - return 1; - } - - if (params.empty()) { - std::cerr << "Error: at least one --param is required\n"; - return 1; - } - - if (outputParam < 0 || outputParam >= static_cast(params.size())) { - std::cerr << "Error: output-param " << outputParam << " out of range (have " - << params.size() << " params)\n"; - return 1; - } - - if (isScalar(params[outputParam])) { - std::cerr << "Error: output-param " << outputParam - << " is a scalar (cannot read back)\n"; - return 1; - } - - // Read PTX - std::string ptx = readFile(ptxFile); - - // Initialize CUDA - CHECK_CU(cuInit(0)); - - CUdevice device; - CHECK_CU(cuDeviceGet(&device, 0)); - - CUcontext ctx; - CHECK_CU(cuCtxCreate(&ctx, 0, device)); - - // Load PTX module - CUmodule module; - CHECK_CU(cuModuleLoadData(&module, ptx.c_str())); - - CUfunction func; - CHECK_CU(cuModuleGetFunction(&func, module, kernelName)); - - // Allocate device buffers for array params - for (auto &p : params) { - if (auto *a = std::get_if>(&p)) - allocDevice(*a); - else if (auto *a = std::get_if>(&p)) - allocDevice(*a); - } - - // Set up kernel parameters — Driver API expects array of pointers to args - std::vector kernelArgs(params.size()); - for (size_t i = 0; i < params.size(); ++i) - kernelArgs[i] = kernelArgPtr(params[i]); - - // Launch kernel - CHECK_CU(cuLaunchKernel(func, grid.x, grid.y, grid.z, block.x, block.y, - block.z, 0, nullptr, kernelArgs.data(), nullptr)); - - CHECK_CU(cuCtxSynchronize()); - - // Copy back and print output param - size_t count = outputCount >= 0 ? static_cast(outputCount) : 0; - if (auto *iArr = std::get_if>(¶ms[outputParam])) { - if (outputCount < 0) - count = iArr->values.size(); - printOutput(*iArr, count); - } else { - auto &fArr = std::get>(params[outputParam]); - if (outputCount < 0) - count = fArr.values.size(); - printOutput(fArr, count); - } - - // Cleanup — only free device memory for array params - for (auto &p : params) { - if (auto *a = std::get_if>(&p)) - cuMemFree(a->devicePtr); - else if (auto *a = std::get_if>(&p)) - cuMemFree(a->devicePtr); - } - cuModuleUnload(module); - cuCtxDestroy(ctx); - - return 0; -} diff --git a/uv.lock b/uv.lock index f2bb18e..c9df123 100644 --- a/uv.lock +++ b/uv.lock @@ -469,6 +469,53 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/aa/50/a9caea39ad19c431c1a3f8a31114df65b260cdfe67786b6c7e7c040c4c44/cryptography-49.0.0-pp311-pypy311_pp73-win_amd64.whl", hash = "sha256:be9fcb48a55f023493482827d4f459bd263cc20efde64f204b97c123201850c6", size = 3783731, upload-time = "2026-06-12T20:02:43.319Z" }, ] +[[package]] +name = "cuda-bindings" +version = "12.9.7" +source = { registry = "https://pypi.org/simple" } +dependencies = [ + { name = "cuda-pathfinder" }, +] +wheels = [ + { url = "https://files.pythonhosted.org/packages/40/f3/f9d1095f90d2a4df24cfcafe7487fd9444c6dacb94e3722be6fedd8ac26c/cuda_bindings-12.9.7-cp311-cp311-manylinux_2_24_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:16043ef5b15ab88fe9954c5c2061b1d8007591b27f2c916331056de0ebc6187e", size = 7114834, upload-time = "2026-05-27T18:44:07.746Z" }, + { url = "https://files.pythonhosted.org/packages/3a/8a/1251e1794b69865aacd5629936006b18ea0816a495de4ecea9a825556eb3/cuda_bindings-12.9.7-cp311-cp311-manylinux_2_24_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:c6496a88d84b1209d6651b0370c19c26319e157c22f6d018bf9a358cd8049041", size = 7647147, upload-time = "2026-05-27T18:44:09.4Z" }, + { url = "https://files.pythonhosted.org/packages/1e/39/158392f6572e6e0def70ca39029c46b75e02ea4a43c63ff7320b3d180a29/cuda_bindings-12.9.7-cp311-cp311-win_amd64.whl", hash = "sha256:c392ffa5010ef4073bfd9dfff4d1ae56032094ed52d3d732014f8e41a73e6b59", size = 7218081, upload-time = "2026-05-27T18:44:11.104Z" }, + { url = "https://files.pythonhosted.org/packages/32/45/557d4ed1fa54f0c7db8aee083229f624990d69f7d00f55477eed5c7e169a/cuda_bindings-12.9.7-cp312-cp312-manylinux_2_24_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:0666d3c082ef8f4b2d670950589373550e9f3bf564d635dd883f24a0b40402ff", size = 7071026, upload-time = "2026-05-27T18:44:13.356Z" }, + { url = "https://files.pythonhosted.org/packages/91/97/e3c6e58ece26a053419ba0a18444b5443cfc64451bbf37f84e8143b8bdca/cuda_bindings-12.9.7-cp312-cp312-manylinux_2_24_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:4c7ef48c5e13ae90f3b2ecfb72f8e99ac43c8f4c43e67e1325b8aae331453687", size = 7611059, upload-time = "2026-05-27T18:44:15.252Z" }, + { url = "https://files.pythonhosted.org/packages/6d/39/afaa3de4d491a55af8961081e0b69c08d51bfbe471c359a7bddb4a28ca41/cuda_bindings-12.9.7-cp312-cp312-win_amd64.whl", hash = "sha256:3c089aaf4f5f570ec50244c68f5a2b00a2c9a8e01e04219fd2e36e340be0d88b", size = 7400841, upload-time = "2026-05-27T18:44:17.164Z" }, + { url = "https://files.pythonhosted.org/packages/eb/7b/f1575e41e1a17dc2f2a408b2e8e864c9324e41e3e23f6401e5efc54c152a/cuda_bindings-12.9.7-cp313-cp313-manylinux_2_24_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:266379e4942051f544a8e7ea1a30ead8d7e8199b6b30fcdc8917cae2bf614e61", size = 6978549, upload-time = "2026-05-27T18:44:18.839Z" }, + { url = "https://files.pythonhosted.org/packages/9d/dc/62d62eb4f91eb721bcf46da51b13e9872ccd8fa7e60eb8ba7b7baeac72c6/cuda_bindings-12.9.7-cp313-cp313-manylinux_2_24_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:59cf4a37b0d662ba15037c9ceebe1a306ebf2c01a8235a09be13cd07094fdb74", size = 7457675, upload-time = "2026-05-27T18:44:20.637Z" }, + { url = "https://files.pythonhosted.org/packages/43/b2/753fe88151001d0dc23f56a8e119fe06b991b0d1a885fa02f9852b12f523/cuda_bindings-12.9.7-cp313-cp313-win_amd64.whl", hash = "sha256:5bd89dcb78475a6d8a4620ea94b74edf0cbbeacee6d1622d8f94452c1e8d3f15", size = 7360097, upload-time = "2026-05-27T18:44:22.405Z" }, + { url = "https://files.pythonhosted.org/packages/f9/77/94d9b85f26add6fe9c9cb7c4ec3b96bc598f7ea5cfbd7490cc0a36adf5be/cuda_bindings-12.9.7-cp313-cp313t-manylinux_2_24_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:2dbcd4801954eb3508f4dc2fa0d0c8eb93eb3f45326fd61be2731418c371e7a0", size = 6870886, upload-time = "2026-05-27T18:44:24.164Z" }, + { url = "https://files.pythonhosted.org/packages/04/dd/3ec34b569e1b990b11276feba306bf8f446656cc38e8ed0f49b5facfeffa/cuda_bindings-12.9.7-cp313-cp313t-manylinux_2_24_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:3747ea132642416786a8e31bf229032df3a7856911ae5426a7be53d032df183d", size = 7345663, upload-time = "2026-05-27T18:44:26.333Z" }, + { url = "https://files.pythonhosted.org/packages/b8/c8/d79a20ba396e7ab2dfdd4b72b62356972b25b88aee2ded49a70c797ddea1/cuda_bindings-12.9.7-cp313-cp313t-win_amd64.whl", hash = "sha256:64f7ade7a7a3b69001489753acc21706d9dbda32db8deb68a767a0a0aab30b68", size = 7780136, upload-time = "2026-05-27T18:44:28.121Z" }, + { url = "https://files.pythonhosted.org/packages/68/e4/075052d42872cf8162da53f14447a4b8abc004c3750e4b724ee502428da0/cuda_bindings-12.9.7-cp314-cp314-manylinux_2_24_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:775960ac9e530717f3b48e165cc6f68684fa9a4141764fd923e4c1a9820acc73", size = 7060090, upload-time = "2026-05-27T18:44:30.281Z" }, + { url = "https://files.pythonhosted.org/packages/ec/cd/3289c810a4d45e5364a3387a74b4c9b6f6f57ee96ae0e5b537cc61dec242/cuda_bindings-12.9.7-cp314-cp314-manylinux_2_24_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:3c47ec1a7a441d91aab32339951df7a1be53451121a12c094bba51467717a35a", size = 7504419, upload-time = "2026-05-27T18:44:31.992Z" }, + { url = "https://files.pythonhosted.org/packages/5b/a0/c429fdcfa5aae181415504c5085ea5944f782b417dd16a7f2a14be0da80d/cuda_bindings-12.9.7-cp314-cp314-win_amd64.whl", hash = "sha256:1e2a4f2ec5b67408c04bb4fbed45d214b66de1f00ee2e972865cacb8708d4e1e", size = 7493876, upload-time = "2026-05-27T18:44:33.618Z" }, + { url = "https://files.pythonhosted.org/packages/11/43/472a6281c3d94e71687e27c657a8f60718d3579b4d94c41deea503165f8a/cuda_bindings-12.9.7-cp314-cp314t-manylinux_2_24_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:00a833d399b31071fab4cf3de2929840ae462dc4848116eeff033d09219e7116", size = 6899146, upload-time = "2026-05-27T18:44:35.556Z" }, + { url = "https://files.pythonhosted.org/packages/2b/13/10c1d0b32a9da65142d213e0733d748457fb3fd066aee4317335266f15c6/cuda_bindings-12.9.7-cp314-cp314t-manylinux_2_24_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:11aeafa2b33995f890086b3fb0f062075176d956e9b6a6fe1a699dddc413f6ad", size = 7369087, upload-time = "2026-05-27T18:44:37.359Z" }, + { url = "https://files.pythonhosted.org/packages/33/10/c71a07cd2a1d4db119bada1848b4752a874ccfe4927d419bfdd05f250920/cuda_bindings-12.9.7-cp314-cp314t-win_amd64.whl", hash = "sha256:ece8dfbc22e6de96a26940ab9887eb3cfe1fc1bc3966169391cdb866bb82bb64", size = 8208198, upload-time = "2026-05-27T18:44:39.053Z" }, +] + +[[package]] +name = "cuda-pathfinder" +version = "1.8.0" +source = { registry = "https://pypi.org/simple" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/a1/b1/ef21259ec74fe0b265ed201379de1d0ef7c14178313ee03705952f1b7093/cuda_pathfinder-1.8.0-py3-none-any.whl", hash = "sha256:c44e574dc997fae2814721d1ae97d0fd6db76db82decbe9b753bf75de53f515e", size = 62539, upload-time = "2026-08-27T21:33:03.229Z" }, +] + +[[package]] +name = "cuda-python" +version = "12.9.7" +source = { registry = "https://pypi.org/simple" } +dependencies = [ + { name = "cuda-bindings" }, +] +wheels = [ + { url = "https://files.pythonhosted.org/packages/c1/9d/05e753afbaac3f92691059b3ba875589c98a425d69e5808cec32b31b580c/cuda_python-12.9.7-py3-none-any.whl", hash = "sha256:23a1fc406d491eef7a7e985095725cb7b20a04a7bd9b7a66400e5c86e082e0aa", size = 7597, upload-time = "2026-05-27T19:50:32.605Z" }, +] + [[package]] name = "curlify" version = "3.0.0" @@ -1895,6 +1942,7 @@ name = "warpforth" version = "0.1.0" source = { virtual = "." } dependencies = [ + { name = "cuda-python" }, { name = "lit" }, { name = "numpy" }, { name = "pytest" }, @@ -1916,6 +1964,7 @@ docs = [ [package.metadata] requires-dist = [ + { name = "cuda-python", specifier = ">=12.4,<13" }, { name = "lit", specifier = ">=18.1.0" }, { name = "numpy" }, { name = "pytest" },