From 048af2b19489a6e246b401fd05ff6e3be5c49eca Mon Sep 17 00:00:00 2001 From: Alex Cameron Date: Fri, 21 Aug 2026 02:26:42 +0000 Subject: [PATCH] feat(pipeline): support configurable GPU architectures --- README.md | 10 ++++++--- gpu_test/conftest.py | 30 ++++++++++++++++++++------- gpu_test/test_compiler.py | 15 ++++++++++++++ gpu_test/test_vast_session.py | 16 +++++++++++++- include/warpforth/Conversion/Passes.h | 12 ++++++++++- lib/Conversion/Passes.cpp | 19 +++++++++++++++++ test/Pipeline/target-options.forth | 5 +++++ tools/warpforthc/warpforthc.cpp | 8 ++++++- 8 files changed, 101 insertions(+), 14 deletions(-) create mode 100644 gpu_test/test_compiler.py diff --git a/README.md b/README.md index c783eff..846b739 100644 --- a/README.md +++ b/README.md @@ -53,7 +53,7 @@ CELLS C + ! Compile to PTX: ```bash -./build/bin/warpforthc matmul.forth -o matmul.ptx +./build/bin/warpforthc matmul.forth -o matmul.ptx --arch sm_80 ``` Test on a GPU (A is 2x4 row-major, B is 4x3 row-major, C is 2x3 output): @@ -97,7 +97,8 @@ output: Available options are `chip` (default `sm_70`), `features` (default `+ptx60`), `libdevice` (the configured libdevice path by default), `opt-level` (`0`, `1`, `2`, or `3`; default `2`), and `compilation-target` (`llvm`, `isa`, `bin`, or -`fatbin`; default `isa`, textual PTX). +`fatbin`; default `isa`, textual PTX). The pipeline's `chip` option accepts the +same architectures as `warpforthc --arch`. ## Language Reference @@ -128,5 +129,8 @@ The `demo/` directory contains a GPT-2 text generation demo that routes scaled d cmake --build build --target check-warpforth # Run end-to-end GPU tests (requires Vast.ai API key) -VASTAI_API_KEY=xxx uv run pytest -v -m gpu +VASTAI_API_KEY=xxx uv run pytest -v -m gpu --arch sm_80 ``` + +The GPU tests pass `--arch` to `warpforthc` and only consider Vast.ai GPUs with +a compatible compute capability. diff --git a/gpu_test/conftest.py b/gpu_test/conftest.py index a74ef18..f7fdf6a 100644 --- a/gpu_test/conftest.py +++ b/gpu_test/conftest.py @@ -33,6 +33,8 @@ POLL_TIMEOUT_S = 300 INSTANCE_LABEL_PREFIX = "warpforth-test-" REMOTE_TMP = "/tmp" # noqa: S108 +DEFAULT_ARCHITECTURE = "sm_70" +SUPPORTED_ARCHITECTURES = ("sm_70", "sm_75", "sm_80", "sm_86", "sm_89", "sm_90") @dataclass @@ -52,11 +54,12 @@ class CompileError(Exception): class Compiler: """Wraps the local warpforthc binary to compile Forth source to PTX.""" - def __init__(self, binary: Path = WARPFORTHC) -> None: + def __init__(self, binary: Path = WARPFORTHC, arch: str = DEFAULT_ARCHITECTURE) -> None: if not binary.exists(): msg = f"warpforthc not found at {binary}; run: cmake --build build" raise FileNotFoundError(msg) self.binary = binary + self.arch = arch def compile_source(self, forth_src: str) -> str: """Compile Forth source code to PTX, returning PTX as a string.""" @@ -66,7 +69,7 @@ def compile_source(self, forth_src: str) -> str: try: result = subprocess.run( - [self.binary, src_path], + [self.binary, src_path, "--arch", self.arch], capture_output=True, text=True, timeout=60, @@ -87,8 +90,9 @@ class VastSession: readiness, and unconditional cleanup on exit. """ - def __init__(self, api_key: str) -> None: + def __init__(self, api_key: str, arch: str = DEFAULT_ARCHITECTURE) -> None: self.sdk = VastAI(api_key, retry=1) + self.arch = arch timestamp = datetime.now(tz=UTC).strftime("%Y%m%dT%H%M%SZ") self.instance_label = f"{INSTANCE_LABEL_PREFIX}{timestamp}-{uuid4().hex}" self.instance_id: int | None = None @@ -221,8 +225,9 @@ def _log_existing_instances(self) -> None: def _launch(self) -> None: """Find the cheapest suitable offer and launch an instance.""" + minimum_compute_capability = int(self.arch.removeprefix("sm_")) * 10 query = ( - f"num_gpus=1 rentable=True rented=False compute_cap>=700" + f"num_gpus=1 rentable=True rented=False compute_cap>={minimum_compute_capability}" f" reliability>=0.95 inet_up>=100 dph<={MAX_COST_PER_HOUR}" ) offers = self.sdk.search_offers(query=query, order="dph", limit=5) @@ -689,18 +694,27 @@ def run( # --- Fixtures --- +def pytest_addoption(parser: pytest.Parser) -> None: + parser.addoption( + "--arch", + choices=SUPPORTED_ARCHITECTURES, + default=DEFAULT_ARCHITECTURE, + help="GPU architecture used for PTX compilation and Vast.ai selection", + ) + + @pytest.fixture(scope="session") -def compiler() -> Compiler: - return Compiler() +def compiler(pytestconfig: pytest.Config) -> Compiler: + return Compiler(arch=pytestconfig.getoption("arch")) @pytest.fixture(scope="session") -def gpu_session() -> Generator[VastSession]: +def gpu_session(pytestconfig: pytest.Config) -> Generator[VastSession]: api_key = os.environ.get("VASTAI_API_KEY") if not api_key: pytest.skip("VASTAI_API_KEY not set") - with VastSession(api_key) as session: + with VastSession(api_key, arch=pytestconfig.getoption("arch")) as session: yield session diff --git a/gpu_test/test_compiler.py b/gpu_test/test_compiler.py new file mode 100644 index 0000000..202ac86 --- /dev/null +++ b/gpu_test/test_compiler.py @@ -0,0 +1,15 @@ +"""Unit tests for local WarpForth compiler invocation.""" + +from pathlib import Path + +from gpu_test.conftest import Compiler + + +def test_compiler_passes_architecture_to_warpforthc(tmp_path: Path) -> None: + binary = tmp_path / "warpforthc" + binary.write_text("#!/bin/sh\nprintf '%s\\n' \"$@\"\n") + binary.chmod(0o755) + + output = Compiler(binary=binary, arch="sm_80").compile_source("\\! kernel main\n42\n") + + assert output.splitlines()[-2:] == ["--arch", "sm_80"] diff --git a/gpu_test/test_vast_session.py b/gpu_test/test_vast_session.py index 3e1ad39..3df0128 100644 --- a/gpu_test/test_vast_session.py +++ b/gpu_test/test_vast_session.py @@ -31,6 +31,7 @@ def __init__(self, instances: list[dict[str, object]] | None = None) -> None: self.hide_next_listing = False self.attached_keys: list[tuple[int, str]] = [] self.attach_result: dict[str, bool] = {"success": True} + self.search_calls: list[dict[str, object]] = [] def show_instances(self) -> list[dict[str, object]]: if self.hide_next_listing: @@ -38,7 +39,8 @@ def show_instances(self) -> list[dict[str, object]]: return [] return [dict(instance) for instance in self.instances] - def search_offers(self, **_kwargs: object) -> list[dict[str, int]]: + def search_offers(self, **kwargs: object) -> list[dict[str, int]]: + self.search_calls.append(kwargs) return [{"id": 1}, {"id": 2}] def create_instance(self, *, label: str, **kwargs: object) -> dict[str, object]: @@ -160,6 +162,18 @@ def test_startup_logs_existing_instances_without_destroying_them( assert "will not be destroyed" in caplog.text +def test_architecture_constrains_gpu_offer_compute_capability( + monkeypatch: pytest.MonkeyPatch, +) -> None: + sdk = FakeVastAI() + configure_session_test(monkeypatch, sdk) + + with VastSession("key", arch="sm_89"): + pass + + assert "compute_cap>=890" in str(sdk.search_calls[0]["query"]) + + def test_rejected_response_reconciles_without_creating_again( monkeypatch: pytest.MonkeyPatch, ) -> None: diff --git a/include/warpforth/Conversion/Passes.h b/include/warpforth/Conversion/Passes.h index 68a9dac..683401f 100644 --- a/include/warpforth/Conversion/Passes.h +++ b/include/warpforth/Conversion/Passes.h @@ -10,6 +10,7 @@ #include "mlir/Dialect/GPU/IR/CompilationInterfaces.h" #include "mlir/Pass/Pass.h" #include "mlir/Pass/PassOptions.h" +#include "llvm/ADT/StringRef.h" #include "llvm/Support/CodeGen.h" #include #include @@ -19,13 +20,22 @@ class OpPassManager; namespace warpforth { +/// Command-line parser for a supported NVVM chip. +class NVVMChipParser : public llvm::cl::parser { +public: + using parser::parser; + + bool parse(llvm::cl::Option &option, StringRef argName, StringRef arg, + std::string &value); +}; + /// Target configuration for the WarpForth compilation pipeline. struct WarpForthPipelineOptions : public PassPipelineOptions { WarpForthPipelineOptions(); /// NVVM target chip, such as `sm_70`. - PassOptions::Option chip; + PassOptions::Option chip; /// NVVM target features, such as `+ptx60`. PassOptions::Option features; /// Path to the CUDA libdevice bitcode library. diff --git a/lib/Conversion/Passes.cpp b/lib/Conversion/Passes.cpp index c28b7cd..096e4c8 100644 --- a/lib/Conversion/Passes.cpp +++ b/lib/Conversion/Passes.cpp @@ -18,10 +18,29 @@ #include "mlir/Transforms/Passes.h" #include "warpforth/Conversion/ForthToGPU/ForthToGPU.h" #include "warpforth/Conversion/ForthToMemRef/ForthToMemRef.h" +#include "llvm/ADT/STLExtras.h" +#include "llvm/Support/raw_ostream.h" namespace mlir { namespace warpforth { +static constexpr llvm::StringLiteral supportedNVVMChips[] = { + "sm_70", "sm_75", "sm_80", "sm_86", "sm_89", "sm_90"}; + +bool NVVMChipParser::parse(llvm::cl::Option &option, StringRef argName, + StringRef arg, std::string &value) { + if (!llvm::is_contained(supportedNVVMChips, arg)) { + std::string supported; + llvm::raw_string_ostream os(supported); + llvm::interleaveComma(supportedNVVMChips, os); + return option.error("unsupported GPU architecture '" + arg + + "'; supported architectures: " + os.str(), + argName); + } + value = arg.str(); + return false; +} + WarpForthPipelineOptions::WarpForthPipelineOptions() : chip(*this, "chip", llvm::cl::desc("NVVM target chip"), llvm::cl::init("sm_70")), diff --git a/test/Pipeline/target-options.forth b/test/Pipeline/target-options.forth index 845fda8..cbc7cd2 100644 --- a/test/Pipeline/target-options.forth +++ b/test/Pipeline/target-options.forth @@ -1,4 +1,7 @@ \ RUN: %warpforth-translate --forth-to-mlir %s | %warpforth-opt '--pass-pipeline=builtin.module(warpforth-pipeline{chip=sm_80 opt-level=3 compilation-target=isa})' --mlir-print-ir-after=nvvm-attach-target --mlir-disable-threading 2>&1 | %FileCheck %s +\ RUN: %warpforthc --arch sm_80 --mlir-print-ir-after-all --mlir-disable-threading %s -o %t.ptx 2>&1 | %FileCheck %s --check-prefix=ARCH +\ RUN: %not %warpforthc --arch sm_99 %s -o %t.ptx 2>&1 | %FileCheck %s --check-prefix=BAD-ARCH +\ RUN: %warpforth-translate --forth-to-mlir %s | %not %warpforth-opt '--pass-pipeline=builtin.module(warpforth-pipeline{chip=sm_99})' 2>&1 | %FileCheck %s --check-prefix=BAD-ARCH \ RUN: %warpforth-translate --forth-to-mlir %s | %not %warpforth-opt '--pass-pipeline=builtin.module(warpforth-pipeline{opt-level=4})' 2>&1 | %FileCheck %s --check-prefix=BAD-OPT \ RUN: %warpforth-translate --forth-to-mlir %s | %not %warpforth-opt '--pass-pipeline=builtin.module(warpforth-pipeline{compilation-target=bogus})' 2>&1 | %FileCheck %s --check-prefix=BAD-FORMAT @@ -6,5 +9,7 @@ 42 \ CHECK: #nvvm.target outputFilename("o", llvm::cl::desc("Output filename"), llvm::cl::value_desc("filename"), llvm::cl::init("-")); +static llvm::cl::opt + targetArchitecture("arch", llvm::cl::desc("GPU architecture"), + llvm::cl::value_desc("sm_XX"), llvm::cl::init("sm_70")); + int main(int argc, char **argv) { llvm::InitLLVM y(argc, argv); warpforth::registerConversionPasses(); @@ -93,7 +97,9 @@ int main(int argc, char **argv) { PassManager pm(&context); if (failed(applyPassManagerCLOptions(pm))) return 1; - warpforth::buildWarpForthPipeline(pm); + warpforth::WarpForthPipelineOptions pipelineOptions; + pipelineOptions.chip = targetArchitecture.getValue(); + warpforth::buildWarpForthPipeline(pm, pipelineOptions); if (failed(pm.run(*module))) { llvm::errs() << "error: compilation pipeline failed\n"; return 1;