Add native Apple Silicon GPU inference with Metal and MPS

.github/workflows/cpu-portability.yml

@@ -71,3 +71,27 @@ jobs:
           build-ci/varda-run --model "$RUNNER_TEMP/varda-synthetic/e.model.safetensors" \
             --input "$RUNNER_TEMP/varda-synthetic/e.golden.safetensors" \
             --device cpu --precision fp32 --threads 2 --steps 3
+      - name: Metal availability and numerical checks
+        if: runner.os == 'macOS'
+        run: |
+          set +e
+          build-ci/varda-metal-info
+          probe=$?
+          set -e
+          if [[ "$probe" == 77 ]]; then
+            echo '::notice::No Metal GPU on this runner; GPU execution checks are SKIPPED, not a GPU pass.'
+            exit 0
+          fi
+          if [[ "$probe" != 0 ]]; then exit "$probe"; fi
+          # Global shader instrumentation corrupts MPS GEMM results on the
+          # hosted Paravirtual GPU. Keep API validation here; the direct shader
+          # CTest above independently enables both API and shader validation.
+          if ! MTL_DEBUG_LAYER=1 MTL_SHADER_VALIDATION=0 \
+            BIN="$PWD/build-ci" SYN_DIR="$RUNNER_TEMP/varda-synthetic" \
+            OUT="$RUNNER_TEMP/varda-metal-parity" bash tests/run_synthetic.sh metal; then
+            cat "$RUNNER_TEMP/varda-metal-parity/a.metal-fp32.log"
+            exit 1
+          fi
+          MTL_DEBUG_LAYER=1 MTL_SHADER_VALIDATION=0 \
+            .venv/bin/python -I python/tests/test_metal.py --bin build-ci \
+            --synthetic "$RUNNER_TEMP/varda-synthetic" --work "$RUNNER_TEMP/varda-metal"

CLAUDE.md

@@ -7,7 +7,8 @@
 - `docs/format.md`: the contracts (converted file, activation names, golden files,
   tolerances). Code that produces or consumes any of them follows it exactly.
 - `docs/design.md`: the map: pipeline, decisions, verification, methodology.
-- `docs/cpu.md`, `docs/cuda.md`: the two backends. `docs/rollout-engine.md`: the C++ rollout.
+- `docs/cpu.md`, `docs/cuda.md`: CPU and CUDA backends. `docs/rollout-engine.md`: the C++ rollout.
+- `docs/metal.md`: the optional Apple Silicon Metal/MPS fp32 backend and its qualification limits.
   `docs/forcings.md`: computed forcings. `docs/rollout.md`: reference rollouts and their
   tolerances. `docs/tools.md`: converter. `docs/testing.md`: every test suite and how to run it.
 - `docs/forecast.md`: running a forecast from real data (GRIB in, GRIB out), overlap, timings, limits.
@@ -23,14 +24,15 @@
 
 ## Hard rules
 
-- The engine (`src/`, `apps/`, `tests/*.cpp|cu`) is C++17/CUDA only: standard library,
+- The engine (`src/`, `apps/`, `tests/*.cpp|cu`) uses C++17/CUDA (Objective-C++ for Metal): standard library,
   OpenMP, BLAS (cblas), CUDA runtime, cuBLAS/cuBLASLt, and the vendored
-  `third_party/nlohmann/json.hpp`. No Python, no libtorch, no other dependency without
+  `third_party/nlohmann/json.hpp`, plus Apple's Foundation/Metal/MetalPerformanceShaders on macOS.
+  No Python, no libtorch, no other dependency without
   asking the integrator.
 - `python/` holds tools only (converter, golden export, reference drivers). They run in
   `.venv` (torch 2.10 CPU, numpy, pyyaml). Never unpickle the checkpoint into real anemoi
   classes in `.venv`: use the restricted unpickler (stubs every non-torch global).
-- The CPU fp32 path is the readable oracle. Optimisations go in the CUDA backend and are
+- The CPU fp32 path is the readable oracle. Optimisations go in the CUDA or Metal backend and are
   verified against golden data at `tolerances/<model>.json`. Never loosen a tolerance
   without saying so explicitly.
 - A parity claim needs numbers: `rel_absmax` and `max_abs` per key, plus device, precision

CMakeLists.txt

@@ -2,6 +2,7 @@ cmake_minimum_required(VERSION 3.24)
 project(varda LANGUAGES CXX)
 
 option(VARDA_CUDA "Build the CUDA backend" ON)
+option(VARDA_METAL "Build the Apple Silicon Metal fp32 backend" ${APPLE})
 option(VARDA_NATIVE "Tune CPU code for the build host (-march=native)" OFF)
 option(VARDA_GRIB "Build the GRIB input/output layer (needs ecCodes, ECCODES_HOME)" ON)
 
@@ -41,6 +42,27 @@ target_include_directories(varda PUBLIC src third_party ${CBLAS_INCLUDE_DIR})
 target_compile_options(varda PRIVATE $<$<COMPILE_LANGUAGE:CXX>:${VARDA_CXX_FLAGS}>)
 target_link_libraries(varda PUBLIC OpenMP::OpenMP_CXX ${BLAS_LIBRARIES})
 
+if(VARDA_METAL)
+  if(NOT APPLE)
+    message(FATAL_ERROR "VARDA_METAL requires macOS and the Apple Metal frameworks")
+  endif()
+  enable_language(OBJCXX)
+  set(CMAKE_OBJCXX_STANDARD 17)
+  set(CMAKE_OBJCXX_STANDARD_REQUIRED ON)
+  find_library(VARDA_FOUNDATION_FRAMEWORK Foundation REQUIRED)
+  find_library(VARDA_METAL_FRAMEWORK Metal REQUIRED)
+  find_library(VARDA_MPS_FRAMEWORK MetalPerformanceShaders REQUIRED)
+  file(READ src/varda/metal/kernels.metal VARDA_METAL_SOURCE)
+  set_property(DIRECTORY APPEND PROPERTY CMAKE_CONFIGURE_DEPENDS src/varda/metal/kernels.metal)
+  configure_file(src/varda/metal/kernels_source.h.in generated/varda/metal/kernels_source.h @ONLY)
+  target_sources(varda PRIVATE src/varda/metal/backend_metal.mm)
+  set_target_properties(varda PROPERTIES OBJCXX_STANDARD 17 OBJCXX_STANDARD_REQUIRED ON)
+  target_include_directories(varda PRIVATE "${CMAKE_CURRENT_BINARY_DIR}/generated")
+  target_compile_definitions(varda PUBLIC VARDA_WITH_METAL)
+  target_compile_options(varda PRIVATE $<$<COMPILE_LANGUAGE:OBJCXX>:-fobjc-arc;-ffp-contract=off;-Wall;-Wextra>)
+  target_link_libraries(varda PUBLIC ${VARDA_FOUNDATION_FRAMEWORK} ${VARDA_METAL_FRAMEWORK} ${VARDA_MPS_FRAMEWORK})
+endif()
+
 if(VARDA_CUDA)
   include(CheckLanguage)
   check_language(CUDA)
@@ -61,7 +83,7 @@ if(VARDA_CUDA)
       $<$<COMPILE_LANGUAGE:CUDA>:--expt-relaxed-constexpr -lineinfo -Xcompiler=-Wall>)
     target_link_libraries(varda PUBLIC CUDA::cudart CUDA::cublas CUDA::cublasLt)
   else()
-    message(WARNING "VARDA_CUDA=ON but no CUDA compiler found: building the CPU backend only")
+    message(WARNING "VARDA_CUDA=ON but no CUDA compiler found: building without the CUDA backend")
   endif()
 endif()
 
@@ -92,6 +114,9 @@ if(VARDA_GRIB)
 endif()
 
 file(GLOB VARDA_APPS CONFIGURE_DEPENDS apps/*.cpp)
+if(VARDA_METAL)
+  list(APPEND VARDA_APPS "${CMAKE_CURRENT_SOURCE_DIR}/apps/metal/varda-metal-info.cpp")
+endif()
 if(VARDA_HAVE_GRIB)
   file(GLOB VARDA_IO_APPS CONFIGURE_DEPENDS apps/io/*.cpp)
 endif()
@@ -124,6 +149,16 @@ foreach(t ${VARDA_TESTS})
   target_link_libraries(${name} PRIVATE varda)
   add_test(NAME ${name} COMMAND ${name})
 endforeach()
+if(VARDA_METAL)
+  add_executable(test_metal_kernels tests/metal/test_metal_kernels.mm)
+  set_target_properties(test_metal_kernels PROPERTIES OBJCXX_STANDARD 17 OBJCXX_STANDARD_REQUIRED ON)
+  target_compile_options(test_metal_kernels PRIVATE -fobjc-arc -ffp-contract=off -Wall -Wextra)
+  target_include_directories(test_metal_kernels PRIVATE "${CMAKE_CURRENT_BINARY_DIR}/generated")
+  target_link_libraries(test_metal_kernels PRIVATE varda)
+  add_test(NAME test_metal_kernels COMMAND test_metal_kernels)
+  set_tests_properties(test_metal_kernels PROPERTIES SKIP_RETURN_CODE 77
+    ENVIRONMENT "MTL_DEBUG_LAYER=1;MTL_SHADER_VALIDATION=1;MTL_SHADER_VALIDATION_REPORT_TO_STDERR=1")
+endif()
 if(VARDA_HAVE_GRIB)
   file(GLOB VARDA_IO_TESTS CONFIGURE_DEPENDS tests/io/*.cpp)
   foreach(t ${VARDA_IO_TESTS})

README.md

@@ -1,6 +1,6 @@
 # varda.cpp
 
-A from-scratch C++/CUDA inference engine for the **forecaster** of
+A from-scratch C++/CUDA/Metal inference engine for the **forecaster** of
 [MeteoSwiss Varda-single-1.0](https://huggingface.co/MeteoSwiss/Varda-single-1.0)
 (`varda-forecaster-det-sgm/model.ckpt`): an anemoi GraphTransformer
 encoder–processor–decoder on a stretched grid (1 688 650 data nodes: 1 km over the Alps,
@@ -120,7 +120,8 @@ anemoi-inference does); it needs ecCodes (`-DVARDA_GRIB=ON`, `tools/env.sh`): `d
 
 Apple Silicon: the CPU fp32 backend builds natively on macOS arm64 with OpenMP and OpenBLAS.
 Build, GRIB dependencies and verification: [docs/apple-silicon.md](docs/apple-silicon.md).
-There is no Apple GPU backend; CUDA requires an NVIDIA GPU.
+The optional [Metal/MPS backend](docs/metal.md) runs on the Apple GPU with
+`--device metal --precision fp32`; CUDA requires an NVIDIA GPU.
 
 ## Licence and disclaimer
 

apps/common.h

@@ -129,7 +129,8 @@ inline Metric compare(const float* a, const float* r, int64_t n) {
 inline varda::Device parse_device(const std::string& s) {
     if (s == "cpu") return varda::Device::CPU;
     if (s == "cuda") return varda::Device::CUDA;
-    throw std::runtime_error("unknown --device '" + s + "' (cpu | cuda)");
+    if (s == "metal") return varda::Device::Metal;
+    throw std::runtime_error("unknown --device '" + s + "' (cpu | cuda | metal)");
 }
 
 inline varda::Precision parse_precision(const std::string& s) {

apps/io/varda-forecast.cpp

@@ -7,7 +7,7 @@
 //                    --regrid-matrix R.safetensors --templates HF_TEMPLATES_DIR [--definitions D] )
 //                  --out DIR
 //                  ( --realdata DIR | --date T --kenda DIR --ifs-prev F.. --ifs F.. --ifs-constants F.. )
-//                  [--lead-time-hours H] [--device cuda|cpu] [--precision fp16|fp32] [--no-graph]
+//                  [--lead-time-hours H] [--device cuda|cpu|metal] [--precision fp16|fp32] [--no-graph]
 //                  [--host-loop] [--allow-nans] [--threads N] [--dump-y] [--write-queue N] [--quiet]
 //
 // --assets BUNDLE fills the five model-release files from one directory (docs/quickstart.md):
@@ -61,7 +61,7 @@ using namespace varda::io;
     std::fprintf(stderr,
                  "usage: varda-forecast (--assets BUNDLE | --model M --io-config J --regrid-matrix R --templates DIR)\n"
                  "         --out DIR (--realdata DIR | --date T --kenda DIR --ifs-prev F.. --ifs F.. --ifs-constants F..)\n"
-                 "         [--lead-time-hours H] [--device cuda|cpu] [--precision fp16|fp32] [--no-graph]\n"
+                 "         [--lead-time-hours H] [--device cuda|cpu|metal] [--precision fp16|fp32] [--no-graph]\n"
                  "         [--host-loop] [--definitions D] [--allow-nans] [--threads N] [--dump-y]\n"
                  "         [--write-queue N] [--quiet]\n"
                  "  --assets BUNDLE: model.safetensors, varda-io.json, regrid-ifs0p25-n320.safetensors, templates/,\n"

apps/metal/varda-metal-info.cpp

@@ -0,0 +1,21 @@
+// Exit 77 explicitly distinguishes a GPU-less runner from a successful GPU test.
+#include <cstdio>
+#include <exception>
+#include "varda/metal/metal_backend.h"
+
+int main() {
+    if (!varda::metal::available()) {
+        std::fprintf(stderr, "no unified-memory Metal GPU is available\n");
+        return 77;
+    }
+    try {
+        const auto d = varda::metal::device_info();
+        std::printf("Metal GPU: %s\nmax buffer: %.2f GiB\nrecommended working set: %.2f GiB\n",
+                    d.name.c_str(), d.max_buffer_length / 1073741824.0,
+                    d.recommended_working_set / 1073741824.0);
+        return 0;
+    } catch (const std::exception& e) {
+        std::fprintf(stderr, "%s\n", e.what());
+        return 1;
+    }
+}

apps/varda-parity.cpp

@@ -1,7 +1,7 @@
 // Parity check of the engine against golden activations (docs/format.md §2-4).
 //
 //   varda-parity --model M.safetensors --check GOLDEN.safetensors=TABLE [--check ...]
-//                [--report-only GOLDEN.safetensors] [--device cpu|cuda] [--precision fp32|fp16]
+//                [--report-only GOLDEN.safetensors] [--device cpu|cuda|metal] [--precision fp32|fp16]
 //                [--threads N] [--no-cuda-graph] [--tolerances FILE] [--report OUT.json]
 //
 // The input `in.x` comes from the first golden file (all must hold the same bytes). One
@@ -34,7 +34,7 @@ struct Check {
 [[noreturn]] void usage() {
     std::fprintf(stderr,
                  "usage: varda-parity --model M --check GOLDEN=TABLE [--check ...] [--report-only GOLDEN]\n"
-                 "                    [--device cpu|cuda] [--precision fp32|fp16] [--threads N] [--no-cuda-graph]\n"
+                 "                    [--device cpu|cuda|metal] [--precision fp32|fp16] [--threads N] [--no-cuda-graph]\n"
                  "                    [--tolerances FILE] [--report OUT.json]\n");
     std::exit(2);
 }

apps/varda-rollout.cpp

@@ -1,7 +1,7 @@
 // Autoregressive rollout (src/varda/rollout.h, src/varda/cuda/rollout_cuda.h, docs/rollout-engine.md).
 //
 //   varda-rollout --model M.safetensors --input IN [--key K] [--date ISO] [--steps N]
-//                 [--device cpu|cuda] [--precision fp32|fp16] [--threads N] [--cuda-device I]
+//                 [--device cpu|cuda|metal] [--precision fp32|fp16] [--threads N] [--cuda-device I]
 //                 [--engine host|device] [--no-cuda-graph]
 //                 [--forcings computed|replay] [--replay-from ROLLOUT]
 //                 [--compare REF]... [--tolerances FILE] [--check-forcings]
@@ -76,7 +76,7 @@ namespace {
 
 [[noreturn]] void usage() {
     std::fprintf(stderr,
-                 "usage: varda-rollout --model M --input IN [--key K] [--date ISO] [--steps N] [--device cpu|cuda]\n"
+                 "usage: varda-rollout --model M --input IN [--key K] [--date ISO] [--steps N] [--device cpu|cuda|metal]\n"
                  "         [--precision fp32|fp16] [--threads N] [--cuda-device I] [--engine host|device]\n"
                  "         [--no-cuda-graph] [--forcings computed|replay] [--replay-from ROLLOUT] [--compare REF]...\n"
                  "         [--tolerances FILE] [--check-forcings] [--verify-replay] [--verify-host] [--verify-constants]\n"
@@ -453,7 +453,11 @@ int main(int argc, char** argv) {
         else if (a == "--key") key = next();
         else if (a == "--date") date_s = next();
         else if (a == "--steps") steps = std::stoi(next());
-        else if (a == "--device") { std::string d = next(); if (d != "cpu" && d != "cuda") usage(); opts.device = d == "cuda" ? varda::Device::CUDA : varda::Device::CPU; }
+        else if (a == "--device") {
+            const std::string d = next();
+            if (d != "cpu" && d != "cuda" && d != "metal") usage();
+            opts.device = d == "cuda" ? varda::Device::CUDA : d == "metal" ? varda::Device::Metal : varda::Device::CPU;
+        }
         else if (a == "--precision") { std::string p = next(); if (p != "fp32" && p != "fp16") usage(); opts.precision = p == "fp16" ? varda::Precision::FP16 : varda::Precision::FP32; }
         else if (a == "--threads") opts.num_threads = std::stoi(next());
         else if (a == "--cuda-device") opts.cuda_device = std::stoi(next());

apps/varda-run.cpp

@@ -1,7 +1,7 @@
 // Runs predict_step on an input read from a safetensors file and reports timings.
 //
 //   varda-run --model M.safetensors --input X.safetensors [--key in.x] [--steps N]
-//             [--device cpu|cuda] [--precision fp32|fp16] [--threads N] [--no-cuda-graph]
+//             [--device cpu|cuda|metal] [--precision fp32|fp16] [--threads N] [--no-cuda-graph]
 //             [--out Y.safetensors]
 //
 // The input tensor (default key `in.x`, e.g. from a golden file) must have the contract shape
@@ -25,7 +25,7 @@ namespace {
 
 [[noreturn]] void usage() {
     std::fprintf(stderr,
-                 "usage: varda-run --model M --input X [--key in.x] [--steps N] [--device cpu|cuda]\n"
+                 "usage: varda-run --model M --input X [--key in.x] [--steps N] [--device cpu|cuda|metal]\n"
                  "                 [--precision fp32|fp16] [--threads N] [--no-cuda-graph] [--out Y]\n");
     std::exit(2);
 }

docs/apple-silicon.md

@@ -1,16 +1,17 @@
 # Apple Silicon (macOS arm64)
 
 The existing CPU fp32 backend runs natively on Apple Silicon using Apple Clang,
-OpenMP and OpenBLAS. CUDA is NVIDIA-only; there is no Metal or MPS backend, so the
-Apple GPU is not used. No model arithmetic or numerical tolerances are changed.
+OpenMP and OpenBLAS. The optional native Metal/MPS backend runs fp32 inference on
+the Apple GPU; see [metal.md](metal.md). CUDA remains NVIDIA-only. CPU arithmetic
+and numerical tolerances are unchanged.
 
 ## CPU build
 
 Install Xcode Command Line Tools and the build dependencies:
 
 ```sh
 brew install cmake libomp openblas
-cmake -S . -B build-apple -DVARDA_CUDA=OFF -DVARDA_GRIB=OFF \
+cmake -S . -B build-apple -DVARDA_CUDA=OFF -DVARDA_METAL=OFF -DVARDA_GRIB=OFF \
   -DCMAKE_PREFIX_PATH="$(brew --prefix libomp);$(brew --prefix openblas)"
 cmake --build build-apple -j 4
 OMP_NUM_THREADS=8 OPENBLAS_NUM_THREADS=1 ctest --test-dir build-apple --output-on-failure