feat(pipeline): support configurable GPU architectures

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.

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
 
 

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"]

gpu_test/test_vast_session.py

@@ -31,14 +31,16 @@ 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:
             self.hide_next_listing = False
             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:

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 <memory>
 #include <string>
@@ -19,13 +20,22 @@ class OpPassManager;
 
 namespace warpforth {
 
+/// Command-line parser for a supported NVVM chip.
+class NVVMChipParser : public llvm::cl::parser<std::string> {
+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> {
   WarpForthPipelineOptions();
 
   /// NVVM target chip, such as `sm_70`.
-  PassOptions::Option<std::string> chip;
+  PassOptions::Option<std::string, NVVMChipParser> chip;
   /// NVVM target features, such as `+ptx60`.
   PassOptions::Option<std::string> features;
   /// Path to the CUDA libdevice bitcode library.

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")),

test/Pipeline/target-options.forth

@@ -1,10 +1,15 @@
 \ 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
 
 \! kernel main
 42
 
 \ CHECK: #nvvm.target<O = 3, chip = "sm_80"
+\ ARCH: #nvvm.target<chip = "sm_80"
+\ BAD-ARCH: unsupported GPU architecture 'sm_99'; supported architectures: sm_70, sm_75, sm_80, sm_86, sm_89, sm_90
 \ BAD-OPT: for the --opt-level option: Cannot find option named '4'!
 \ BAD-FORMAT: for the --compilation-target option: Cannot find option named 'bogus'!

tools/warpforthc/warpforthc.cpp

@@ -43,6 +43,10 @@ static llvm::cl::opt<std::string>
     outputFilename("o", llvm::cl::desc("Output filename"),
                    llvm::cl::value_desc("filename"), llvm::cl::init("-"));
 
+static llvm::cl::opt<std::string, false, warpforth::NVVMChipParser>
+    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;