diff --git a/.github/workflows/cpu-portability.yml b/.github/workflows/cpu-portability.yml index 1673a5f..5cfaad9 100644 --- a/.github/workflows/cpu-portability.yml +++ b/.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" diff --git a/CLAUDE.md b/CLAUDE.md index c3a2a30..c898bca 100644 --- a/CLAUDE.md +++ b/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/.json`. Never loosen a tolerance without saying so explicitly. - A parity claim needs numbers: `rel_absmax` and `max_abs` per key, plus device, precision diff --git a/CMakeLists.txt b/CMakeLists.txt index ed6de78..f27f9d1 100644 --- a/CMakeLists.txt +++ b/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 $<$:${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 $<$:-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) $<$:--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}) diff --git a/README.md b/README.md index 75b7b9e..4a5d1a1 100644 --- a/README.md +++ b/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 diff --git a/apps/common.h b/apps/common.h index 799aaaa..387b371 100644 --- a/apps/common.h +++ b/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) { diff --git a/apps/io/varda-forecast.cpp b/apps/io/varda-forecast.cpp index 19a5f6b..dc45120 100644 --- a/apps/io/varda-forecast.cpp +++ b/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" diff --git a/apps/metal/varda-metal-info.cpp b/apps/metal/varda-metal-info.cpp new file mode 100644 index 0000000..d384fac --- /dev/null +++ b/apps/metal/varda-metal-info.cpp @@ -0,0 +1,21 @@ +// Exit 77 explicitly distinguishes a GPU-less runner from a successful GPU test. +#include +#include +#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; + } +} diff --git a/apps/varda-parity.cpp b/apps/varda-parity.cpp index d2dd284..8b8552d 100644 --- a/apps/varda-parity.cpp +++ b/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); } diff --git a/apps/varda-rollout.cpp b/apps/varda-rollout.cpp index d5c7344..cbc978c 100644 --- a/apps/varda-rollout.cpp +++ b/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()); diff --git a/apps/varda-run.cpp b/apps/varda-run.cpp index 6d130bf..7296515 100644 --- a/apps/varda-run.cpp +++ b/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); } diff --git a/docs/apple-silicon.md b/docs/apple-silicon.md index 67f18b2..c0a1636 100644 --- a/docs/apple-silicon.md +++ b/docs/apple-silicon.md @@ -1,8 +1,9 @@ # 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 @@ -10,7 +11,7 @@ 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 diff --git a/docs/metal.md b/docs/metal.md new file mode 100644 index 0000000..7c63430 --- /dev/null +++ b/docs/metal.md @@ -0,0 +1,128 @@ +# Apple Silicon Metal backend + +`--device metal --precision fp32` runs the forecaster on a unified-memory Metal +GPU. It uses MPS matrix multiplication and native Metal compute shaders. There +is no PyTorch dependency and no silent CPU fallback. CPU and CUDA numerical +kernels, precision settings and tolerance files are unchanged. + +## Build and run + +Follow `apple-silicon.md` for OpenMP/OpenBLAS and the pinned GRIB dependencies. +`VARDA_METAL` defaults to ON on macOS and OFF elsewhere. Disable it explicitly +with `-DVARDA_METAL=OFF` for a CPU-only Mac build. Enabling it on Linux is an +error; Linux builds do not enable Objective-C++ or link any Apple framework. + +```sh +cmake -S . -B build-metal -DVARDA_CUDA=OFF -DVARDA_METAL=ON \ + -DCMAKE_PREFIX_PATH="$(brew --prefix libomp);$(brew --prefix openblas)" +cmake --build build-metal -j 4 +build-metal/varda-metal-info +build-metal/varda-run --model MODEL.safetensors --input INPUT.safetensors \ + --device metal --precision fp32 --steps 3 +``` + +The build links Apple's Foundation, Metal and MetalPerformanceShaders frameworks. +Shader source is embedded at configure time and compiled by Metal at model load; +the executable does not need the source tree or an external `.metallib` at run +time. The normal parity app still needs its reference/tolerance files. + +`varda-rollout --device metal --precision fp32` uses the existing host rollout +loop: predictions run on the GPU, while feedback and computed forcings run on +the host. `--engine device` remains CUDA-only. With an assets bundle, +`./varda forecast latest --device metal --precision fp32` runs locally and probes +the Metal GPU before downloading inputs; it does not invoke Slurm. + +## Implementation and numerics + +- Dense layers: `MPSMatrixMultiplication`, float32 inputs, weights and outputs, + with a separate bias addition and one cached multiplication object per shape. +- Graph attention: destination-sorted CSR, explicit edge projection inside the + score/value kernels, per-edge/head scores and softmax, and stored-order sums. + No edge-by-channel tensor is materialised. Empty graphs and zero-in-degree + destinations are supported. +- Normalisation, input assembly, residuals, bounding and inverse normalisation + are GPU kernels. Fraction bounding stages the products before assignment, + including when the total field is itself an output being bounded. +- Metal compilation uses safe math and precise functions, with contraction off + in the shaders. Metal has no fp64: LayerNorm uses one SIMD group per row, + queries the SIMD width at runtime and computes compensated fp32 moments after + subtracting a common row origin. Applying normalization to centered values + avoids cancellation for nearly constant rows with large offsets. + The erf form of GELU evaluates erf with A&S 7.1.26; the configured tanh form is + also supported. This is a separate numerical implementation, validated by the + existing fp32 tolerances, rather than a bitwise match to CPU or CUDA. +- Shared buffers hold the weights and reusable activations. Input/output copies + and hook reads happen only after the relevant GPU work is complete. Hooks + synchronize/read the same production calculations; they do not substitute CPU + calculations. GPU command-buffer errors propagate to the caller. Submission + batches are capped at 64 encoded operations (compute dispatches or MPS calls), + retaining a final completion wait before reading output. This bounds queued + operation count and resource retention, not elapsed GPU time. + +`VARDA_METAL_CHUNK_ROWS` controls complete destination-row chunks (default 8192). +It bounds dense/attention workspace; it does not split a destination's softmax. +All buffers are planned at load. Allocations are checked against the device's +maximum buffer size and recommended working-set budget; dimensions must fit +32-bit GPU indexing. Bounding index lists must have unique output indices. + +## Verification + +```sh +BIN="$PWD/build-metal" SYN_DIR="$PWD/local/synthetic" \ + OPENBLAS_NUM_THREADS=1 bash tests/run_synthetic.sh metal +.venv/bin/python -I python/tests/test_metal.py --bin build-metal \ + --synthetic local/synthetic --work local/metal +``` + +On Apple M1, fp32: all 12 synthetic parity runs (six configurations, default +chunking and 37-row chunks) pass unchanged `synthetic-fp32` tolerances. Worst +activation is `proc.1.out` in model d: `rel_absmax` 2.72e-6 and `max_abs` +4.18e-5, or 0.272 of its 1e-5 relative tolerance (8 host threads). Input +normalisation and assembly are bitwise equal to the reference; hook-free and +hooked outputs are bitwise equal. A small-chunk run also passes with the Metal +API validation layer enabled. + +The integration suite checks repeated inference, a four-step host rollout with +computed forcings/feedback against CPU at the existing synthetic fp32 tolerance, +completely empty graphs, FP16 rejection and invalid chunk sizes. These are +synthetic correctness checks, not meteorological forecast qualification. +It also generates three independent golden models with 7/35/96 channels, +including one hidden node, incomplete SIMD groups, 1/7/37-row chunks, and 8,201 +data nodes crossing the default chunk boundary: ten further parity checks, +including hook inertness. CI runs GPU parity and integration with Metal API +validation. The direct kernel test below also enables GPU Shader Validation. +Global shader instrumentation changes MPS GEMM results on the hosted Apple +Paravirtual GPU (an API-only control passes, while the first embedding GEMM +fails with instrumentation). Full inference therefore disables global shader +instrumentation in hosted CI; the full suite additionally passes with both +validation layers enabled on the physical M1. + +The macOS-only `test_metal_kernels` CTest target directly checks the production +LayerNorm shader against double-precision mathematical references: 384 cases +with widths 1 through 4097, random/constant/nearly constant/cancellation-prone +inputs, optional bias and output guards. All pass at 1e-5; measured worst +`rel_absmax` 3.39e-7 and `max_abs` 1.62e-6. The old serial shader fails 90 of +these cases (worst relative error 1.07e-2). This test enables Metal API/shader +validation and returns skip code 77 when no GPU is available. + +An Apple M1 fp32 microbenchmark of models e and d (two alternating baseline/new +rounds of 30 predictions, discard the first three calls, eight host threads) +measured mean warm forward times of 40.8 -> 25.9 ms and 53.2 -> 31.1 ms after +parallel LayerNorm and bounded batching. Repeats are bitwise stable; Metal has +no CUDA graph path. These synthetic timings do not establish full-model speed. + +`varda-metal-info` exits 77 if no unified-memory Metal device is present. An +explicit synthetic `metal` run also exits 77 rather than reporting a GPU pass +when execution was skipped. CI compiles the backend on macOS arm64 and runs the +GPU checks when a device is exposed; a GPU-less hosted runner reports the skip. + +## Limits + +FP32 only; no mixed precision, Metal graph capture, or GPU-resident rollout +window yet. The first call includes MPS kernel setup. Large-model performance +has not been measured; synthetic timings must not be extrapolated to Varda. + +The full Varda checkpoint and real-data GRIB forecasts have not been qualified +on Metal. Its activation storage exceeds the 16 GB test Mac's GPU working-set +budget, so a larger unified-memory machine and full reference assets are needed +for that qualification. Device limits are reported by `varda-metal-info`. diff --git a/python/tests/test_metal.py b/python/tests/test_metal.py new file mode 100644 index 0000000..5203218 --- /dev/null +++ b/python/tests/test_metal.py @@ -0,0 +1,158 @@ +"""Metal integration checks using the existing synthetic fixtures (no torch). + +Run after tests/run_synthetic.sh metal: + .venv/bin/python -I python/tests/test_metal.py --bin build-metal --synthetic local/synthetic --work local/metal +""" +import argparse +import importlib.util +import json +import os +import subprocess +import tempfile +import unittest +from pathlib import Path + +import numpy as np + +spec = importlib.util.spec_from_file_location("varda_synthetic", Path(__file__).with_name("synthetic.py")) +synthetic = importlib.util.module_from_spec(spec) +spec.loader.exec_module(synthetic) +BIN = SYNTHETIC = WORK = None + + +class MetalTests(unittest.TestCase): + @classmethod + def setUpClass(cls): + if BIN is None: + raise unittest.SkipTest("run this GPU integration suite with --bin, --synthetic and --work") + cls.model = SYNTHETIC / "a.model.safetensors" + cls.inputs = SYNTHETIC / "a.golden.safetensors" + cls.tmp = tempfile.TemporaryDirectory(prefix="metal-integration-", dir=WORK) + cls.root = Path(cls.tmp.name) + + @classmethod + def tearDownClass(cls): + cls.tmp.cleanup() + + def run_app(self, name, *args, ok=True, env=None): + result = subprocess.run([str(BIN / name), *map(str, args)], check=False, text=True, + stdout=subprocess.PIPE, stderr=subprocess.STDOUT, env=env, timeout=120) + if ok: + self.assertEqual(result.returncode, 0, result.stdout) + else: + self.assertNotEqual(result.returncode, 0, result.stdout) + return result + + def test_repeated_inference(self): + result = self.run_app("varda-run", "--model", self.model, "--input", self.inputs, + "--device", "metal", "--precision", "fp32", "--steps", 3) + self.assertIn("repeats bitwise equal", result.stdout) + + def test_irregular_shapes_and_chunk_boundaries(self): + # Fresh independent golden data exercises incomplete SIMD groups and + # workspace reuse across both small chunks and the default 8192 limit. + configs = (("odd", 35, 5, 257, 33), ("tiny", 7, 1, 9, 1), + ("cross", 96, 3, 8201, 257)) + for name, channels, heads, nodes, hidden in configs: + cfg = dict(synthetic.CONFIGS["f"], C=channels, heads=heads, + mlp=83 if channels < 40 else 192, G=nodes, H=hidden, + seed=71 + channels, n_proc=3) + synthetic.CONFIGS[name] = cfg + try: + synthetic.generate(name, str(self.root)) + tolerance = self.root / f"{name}.tolerances.json" + tolerance.write_text(json.dumps(synthetic.tolerance_file(name))) + finally: + del synthetic.CONFIGS[name] + for chunk in ((37, 8192) if name == "cross" else (1, 7, 37, 8192)): + with self.subTest(model=name, chunk=chunk): + report = self.root / f"{name}.{chunk}.report.json" + self.run_app("varda-parity", "--model", self.root / f"{name}.model.safetensors", + "--check", f"{self.root}/{name}.golden.safetensors=synthetic-fp32", + "--tolerances", tolerance, "--report", report, + "--device", "metal", "--precision", "fp32", "--threads", 2, + env=dict(os.environ, VARDA_METAL_CHUNK_ROWS=str(chunk))) + self.assertTrue(json.loads(report.read_text())["hook_inert"]) + + def test_fp16_rejected(self): + result = self.run_app("varda-run", "--model", self.model, "--input", self.inputs, + "--device", "metal", "--precision", "fp16", ok=False) + self.assertIn("Metal backend runs fp32 only", result.stdout) + + def test_invalid_chunk_rejected(self): + for value in ("0", "not-a-number"): + with self.subTest(value=value): + env = dict(os.environ, VARDA_METAL_CHUNK_ROWS=value) + result = self.run_app("varda-run", "--model", self.model, "--input", self.inputs, + "--device", "metal", "--precision", "fp32", env=env, ok=False) + self.assertIn("VARDA_METAL_CHUNK_ROWS must be", result.stdout) + + def test_host_rollout(self): + # Add rollout metadata to a separate test model: original weights and + # fixtures remain untouched. Use a real computed forcing to exercise + # feedback, constant channels, dates and CPU/GPU handover together. + tensors, metadata = synthetic.read_st(self.model) + cfg = json.loads(metadata["config"]) + cfg["source"]["uuid"] += "-metal-rollout" + prognostic = set(cfg["indices"]["residual_input"]) + kept = [i for i in range(cfg["dims"]["n_input_vars"]) if i not in prognostic] + dynamic, *kept = kept + data_index = cfg["indices"]["norm_input"][dynamic] + cfg["variables"]["input"][dynamic] = "cos_julian_day" + cfg["variables"]["data"][data_index] = "cos_julian_day" + for out, data in enumerate(cfg["indices"]["norm_output"]): + if data == data_index: + cfg["variables"]["output"][out] = "cos_julian_day" + cfg["forcings"] = {"dynamic": ["cos_julian_day"], "dynamic_input_index": [dynamic], + "constant": [cfg["variables"]["input"][i] for i in kept], + "constant_input_index": kept, "timestep_hours": 6} + nodes = cfg["dims"]["n_data_nodes"] + tensors["grid.latitudes"] = np.linspace(-45, 45, nodes, dtype=np.float64) + tensors["grid.longitudes"] = np.linspace(-180, 180, nodes, dtype=np.float64) + model = self.root / "rollout.model.safetensors" + synthetic.write_st(str(model), tensors, {"varda_format": "1", "config": json.dumps(cfg)}) + results = {} + for device in ("cpu", "metal"): + output = self.root / f"rollout.{device}.safetensors" + result = self.run_app("varda-rollout", "--model", model, "--input", self.inputs, + "--device", device, "--precision", "fp32", "--threads", 2, + "--date", "2024-06-01T00:00:00", "--steps", 4, "--out", output) + self.assertIn("engine host", result.stdout) + results[device] = synthetic.read_st(output)[0]["y"].astype(np.float64) + self.assertTrue(np.isfinite(results["metal"]).all()) + for step, (actual, reference) in enumerate(zip(results["metal"], results["cpu"])): + error = np.max(np.abs(actual - reference)) / np.max(np.abs(reference)) + self.assertLessEqual(error, synthetic.FP32_TOL, f"step {step}: rel_absmax={error}") + + def test_empty_graphs(self): + tensors, metadata = synthetic.read_st(self.model) + cfg = json.loads(metadata["config"]) + cfg["source"]["uuid"] += "-empty-graphs" + for graph in ("enc", "proc", "dec"): + tensors[f"graph.{graph}.edge_index"] = np.empty((2, 0), dtype=np.int64) + tensors[f"graph.{graph}.edge_attr"] = np.empty((0, cfg["dims"]["edge_attr_dim"]), dtype=np.float32) + model = self.root / "empty-graphs.model.safetensors" + synthetic.write_st(str(model), tensors, {"varda_format": "1", "config": json.dumps(cfg)}) + results = {} + for device in ("cpu", "metal"): + output = self.root / f"empty-graphs.{device}.safetensors" + self.run_app("varda-run", "--model", model, "--input", self.inputs, + "--device", device, "--precision", "fp32", "--threads", 2, "--out", output) + results[device] = synthetic.read_st(output)[0]["out.y"].astype(np.float64) + self.assertTrue(np.isfinite(results["metal"]).all()) + error = np.max(np.abs(results["metal"] - results["cpu"])) / np.max(np.abs(results["cpu"])) + self.assertLessEqual(error, synthetic.FP32_TOL) + + +if __name__ == "__main__": + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("--bin", required=True, type=Path) + parser.add_argument("--synthetic", required=True, type=Path) + parser.add_argument("--work", required=True, type=Path) + options, remaining = parser.parse_known_args() + BIN, SYNTHETIC, WORK = options.bin.resolve(), options.synthetic.resolve(), options.work.resolve() + WORK.mkdir(parents=True, exist_ok=True) + probe = subprocess.run([str(BIN / "varda-metal-info")], check=False) + if probe.returncode: + raise SystemExit(probe.returncode) + unittest.main(argv=[__file__, *remaining], verbosity=2) diff --git a/src/varda/backend.h b/src/varda/backend.h index 4f1f5a4..70d469e 100644 --- a/src/varda/backend.h +++ b/src/varda/backend.h @@ -25,5 +25,9 @@ std::unique_ptr make_cpu_backend(const Config& cfg, SafeTensorsFile& fi std::unique_ptr make_cuda_backend(const Config& cfg, SafeTensorsFile& file, const Options& opts); #endif +#ifdef VARDA_WITH_METAL +std::unique_ptr make_metal_backend(const Config& cfg, SafeTensorsFile& file, + const Options& opts); +#endif } // namespace varda diff --git a/src/varda/metal/backend_metal.mm b/src/varda/metal/backend_metal.mm new file mode 100644 index 0000000..94c07d5 --- /dev/null +++ b/src/varda/metal/backend_metal.mm @@ -0,0 +1,493 @@ +// Native Metal fp32 inference. MPS handles dense GEMMs; graph attention and all +// elementwise operations run as Metal compute kernels. No CPU compute fallback. +#import +#import +#import + +#include +#include +#include +#include +#include +#include +#include +#include +#include + +#include "varda/backend.h" +#include "varda/graph.h" +#include "varda/metal/metal_backend.h" +#include "varda/metal/kernels_source.h" + +namespace varda::metal { +bool available() { + @autoreleasepool { + id d = MTLCreateSystemDefaultDevice(); + return d && d.hasUnifiedMemory; + } +} +DeviceInfo device_info() { + @autoreleasepool { + id d = MTLCreateSystemDefaultDevice(); + if (!d || !d.hasUnifiedMemory) throw std::runtime_error("no unified-memory Metal GPU is available"); + return {d.name.UTF8String, d.maxBufferLength, d.recommendedMaxWorkingSetSize}; + } +} +} // namespace varda::metal + +namespace varda { +namespace { + +uint32_t u(int64_t n) { + if (n < 0 || uint64_t(n) > std::numeric_limits::max()) + throw std::runtime_error("Metal tensor dimensions exceed the 32-bit indexing limit"); + return static_cast(n); +} +struct Args { + uint32_t n = 0, a = 0, b = 0, c = 0, d = 0, e = 0, f = 0, g = 0; + float x = 0, y = 0; +}; +static_assert(sizeof(Args) == 40, "Metal parameter layout"); + +struct Tensor { + id buffer = nil; + NSUInteger offset = 0; + Tensor sub(int64_t floats) const { + Tensor t = *this; + t.offset += static_cast(floats) * sizeof(float); + return t; + } + float* data() const { return reinterpret_cast(static_cast(buffer.contents) + offset); } +}; +struct Linear { Tensor w, b; int64_t out = 0, in = 0; }; +struct Norm { Tensor w, b; }; +struct Block { + Linear q, k, v, self, edge, proj, mlp1, mlp2; + Norm src, dst, mlp; +}; +struct Graph { + CsrGraph host; + Tensor rows, sources, destinations, attrs; +}; +struct Bound { Bounding cfg; Tensor indices, totals; }; + +class MetalBackend final : public Backend { +public: + MetalBackend(const Config& cfg, SafeTensorsFile& file, const Options& opts); + void predict_step(const float* x, float* y, const ActivationHook* hook) override; + +private: + Tensor bytes(size_t n, const void* source = nullptr); + Tensor floats(int64_t n) { u(n); return bytes(static_cast(n) * sizeof(float)); } + template Tensor upload(const std::vector& v) { return bytes(v.size() * sizeof(T), v.data()); } + Tensor arena(const std::string& name, int64_t n); + Linear load_linear(SafeTensorsFile& f, const std::string& name, int64_t out, int64_t in); + Norm load_norm(SafeTensorsFile& f, const std::string& name); + Block load_block(SafeTensorsFile& f, const std::string& name, bool mapper); + Graph load_graph(SafeTensorsFile& f, const std::string& name, int64_t src, int64_t dst); + id command(); + void flush(); + void finish_op(); + void dispatch(const char* name, const Args& args, std::initializer_list buffers); + void linear(const Linear& l, const Tensor& x, int64_t rows, const Tensor& out); + void norm(const Norm& l, const Tensor& x, int64_t rows, const Tensor& out); + void copy(const Tensor& x, const Tensor& out, int64_t n); + void add(const Tensor& x, const Tensor& y, const Tensor& out, int64_t n); + void source(const Block& b, const Tensor& x, int64_t rows, int64_t dest_row); + void destination(const Block& b, const Graph& g, int64_t row, int64_t rows, + const Tensor& x, const Tensor& out); + void report(const ActivationHook* hook, const std::string& key, const Tensor& t, + std::vector shape); + void run(const float* x, float* y, const ActivationHook* hook); + + Config cfg_; + int64_t C_, G_, H_, F_, chunk_; + id device_; + id queue_; + id command_; + size_t encoded_ops_ = 0; + std::map> kernels_; + std::map, MPSMatrixMultiplication*> gemms_; + std::map arena_; + Tensor mul_in_, add_in_, mul_out_, add_out_, data_attr_, hidden_attr_, residual_in_, residual_out_; + Tensor keys_, values_, scores_, ln_, q_, self_, attention_, residual_, hidden_; + Linear enc_src_, enc_dst_, dec_dst_, out_proj_; + Norm out_norm_; + Block enc_, dec_; + std::vector processors_; + Graph enc_graph_, proc_graph_, dec_graph_; + std::vector bounds_; +}; + +Tensor MetalBackend::bytes(size_t n, const void* source) { + u(n / sizeof(float)); // every GPU buffer stores 32-bit floats or indices + const size_t size = std::max(n, sizeof(float)); + if (size > device_.maxBufferLength) + throw std::runtime_error("Metal buffer exceeds this GPU's maxBufferLength"); + const uint64_t used = device_.currentAllocatedSize, limit = device_.recommendedMaxWorkingSetSize; + if (used > limit || size > limit - used) + throw std::runtime_error("Metal model exceeds this GPU's recommended working-set budget (" + + std::to_string(limit >> 20) + " MiB); use a Mac with more unified memory"); + id buffer = [device_ newBufferWithLength:size options:MTLResourceStorageModeShared]; + if (!buffer) throw std::runtime_error("Metal buffer allocation failed (" + std::to_string(size) + " bytes)"); + if (source && n) std::memcpy(buffer.contents, source, n); + return {buffer, 0}; +} + +Tensor MetalBackend::arena(const std::string& name, int64_t n) { + Tensor& t = arena_[name]; + u(n); + if (!t.buffer || t.buffer.length < static_cast(n) * sizeof(float)) t = floats(n); + return t; +} + +Linear MetalBackend::load_linear(SafeTensorsFile& f, const std::string& name, int64_t out, int64_t in) { + const TensorView& w = f.take(name + ".weight", DType::F32, {out, in}); + Linear l{bytes(w.nbytes, w.data), {}, out, in}; + if (!cfg_.is_absent(name + ".bias")) { + const TensorView& b = f.take(name + ".bias", DType::F32, {out}); + l.b = bytes(b.nbytes, b.data); + } + return l; +} + +Norm MetalBackend::load_norm(SafeTensorsFile& f, const std::string& name) { + const TensorView& w = f.take(name + ".weight", DType::F32, {C_}); + Norm l{bytes(w.nbytes, w.data), {}}; + if (!cfg_.is_absent(name + ".bias")) { + const TensorView& b = f.take(name + ".bias", DType::F32, {C_}); + l.b = bytes(b.nbytes, b.data); + } + return l; +} + +Block MetalBackend::load_block(SafeTensorsFile& f, const std::string& name, bool mapper) { + Block b; + b.q = load_linear(f, name + ".q", C_, C_); + b.k = load_linear(f, name + ".k", C_, C_); + b.v = load_linear(f, name + ".v", C_, C_); + b.self = load_linear(f, name + ".self", C_, C_); + b.edge = load_linear(f, name + ".edge", C_, cfg_.edge_attr_dim); + b.proj = load_linear(f, name + ".proj", C_, C_); + b.mlp1 = load_linear(f, name + ".mlp1", cfg_.mlp_hidden, C_); + b.mlp2 = load_linear(f, name + ".mlp2", C_, cfg_.mlp_hidden); + b.dst = load_norm(f, name + ".ln_dst"); + b.mlp = load_norm(f, name + ".ln_mlp"); + b.src = mapper ? load_norm(f, name + ".ln_src") : b.dst; + return b; +} + +Graph MetalBackend::load_graph(SafeTensorsFile& f, const std::string& name, int64_t src, int64_t dst) { + const TensorView& ei = f.take("graph." + name + ".edge_index"); + const TensorView& ea = f.take("graph." + name + ".edge_attr"); + if (ea.shape.size() != 2 || ea.shape[1] != cfg_.edge_attr_dim) + throw std::runtime_error(ea.name + ": edge attribute shape does not match the config"); + Graph g; + g.host = build_csr_by_dst(ei, ea, src, dst); + std::vector dest(static_cast(g.host.n_edges)); + for (int64_t row = 0; row < dst; ++row) + for (int32_t e = g.host.row_ptr[row]; e < g.host.row_ptr[row + 1]; ++e) dest[e] = u(row); + g.rows = upload(g.host.row_ptr); + g.sources = upload(g.host.src); + g.destinations = upload(dest); + g.attrs = upload(g.host.attr); + return g; +} + +MetalBackend::MetalBackend(const Config& cfg, SafeTensorsFile& f, const Options& opts) + : cfg_(cfg), C_(cfg.num_channels), G_(cfg.n_data_nodes), H_(cfg.n_hidden_nodes), F_(cfg.src_features()), + chunk_(8192) { + @autoreleasepool { + if (opts.precision != Precision::FP32) throw std::runtime_error("the Metal backend runs fp32 only"); + device_ = MTLCreateSystemDefaultDevice(); + if (!device_ || !device_.hasUnifiedMemory) + throw std::runtime_error("no unified-memory Metal GPU is available; use --device cpu"); + queue_ = [device_ newCommandQueue]; + if (!queue_) throw std::runtime_error("cannot create a Metal command queue"); + MTLCompileOptions* compile = [MTLCompileOptions new]; +#if __MAC_OS_X_VERSION_MAX_ALLOWED >= 150000 + if (@available(macOS 15.0, *)) { + compile.mathMode = MTLMathModeSafe; + compile.mathFloatingPointFunctions = MTLMathFloatingPointFunctionsPrecise; + } + else { +#pragma clang diagnostic push +#pragma clang diagnostic ignored "-Wdeprecated-declarations" + compile.fastMathEnabled = NO; // compatibility with macOS before 15 +#pragma clang diagnostic pop + } +#else + compile.fastMathEnabled = NO; +#endif + compile.languageVersion = MTLLanguageVersion2_3; + NSError* error = nil; + id library = [device_ newLibraryWithSource:[NSString stringWithUTF8String:metal::kernel_source] + options:compile error:&error]; + if (!library) throw std::runtime_error("Metal shader compilation: " + std::string(error.localizedDescription.UTF8String)); + for (const char* name : {"normalise", "assemble", "bias_add", "add_arrays", "copy_array", "layernorm", + "gelu", "attention_scores", "attention_softmax", "attention_values", "residual", + "bound_clamp", "bound_product", "bound_assign", "denormalise"}) { + id function = [library newFunctionWithName:[NSString stringWithUTF8String:name]]; + id pipeline = [device_ newComputePipelineStateWithFunction:function error:&error]; + if (!pipeline) throw std::runtime_error("Metal pipeline " + std::string(name) + ": " + error.localizedDescription.UTF8String); + kernels_[name] = pipeline; + } + if (const char* value = std::getenv("VARDA_METAL_CHUNK_ROWS")) { + char* end = nullptr; + const long long n = std::strtoll(value, &end, 10); + if (!*value || *end || n < 1 || n > std::numeric_limits::max()) + throw std::runtime_error("VARDA_METAL_CHUNK_ROWS must be a positive 32-bit integer"); + chunk_ = n; + } + chunk_ = std::min(chunk_, std::max(G_, H_)); + const float* mul = f.take("norm.mul", DType::F32, {cfg.n_data_vars}).as(); + const float* add = f.take("norm.add", DType::F32, {cfg.n_data_vars}).as(); + std::vector mi, ai, mo, ao; + for (int i : cfg.norm_input) mi.push_back(mul[i]), ai.push_back(add[i]); + for (int i : cfg.norm_output) mo.push_back(mul[i]), ao.push_back(add[i]); + mul_in_ = upload(mi); add_in_ = upload(ai); mul_out_ = upload(mo); add_out_ = upload(ao); + auto attr = [&](const char* name, int64_t nodes) { + const TensorView& t = f.take(name, DType::F32, {nodes, cfg.node_attr_dim}); + return bytes(t.nbytes, t.data); + }; + data_attr_ = attr("nodes.data.attr", G_); + hidden_attr_ = attr("nodes.hidden.attr", H_); + enc_graph_ = load_graph(f, "enc", G_, H_); + proc_graph_ = load_graph(f, "proc", H_, H_); + dec_graph_ = load_graph(f, "dec", H_, G_); + enc_src_ = load_linear(f, "enc.emb_src", C_, F_); + enc_dst_ = load_linear(f, "enc.emb_dst", C_, cfg.node_attr_dim); + dec_dst_ = load_linear(f, "dec.emb_dst", C_, F_); + out_norm_ = load_norm(f, "dec.out_norm"); + out_proj_ = load_linear(f, "dec.out_proj", cfg.n_output_vars, C_); + for (const BlockInfo& bi : cfg.blocks) { + const bool mapper = bi.kind == "graph_transformer_mapper_block"; + if (bi.name == "enc.block") enc_ = load_block(f, bi.name, mapper); + else if (bi.name == "dec.block") dec_ = load_block(f, bi.name, mapper); + else processors_.push_back(load_block(f, bi.name, mapper)); + } + residual_in_ = upload(cfg.residual_input); residual_out_ = upload(cfg.residual_output); + int64_t max_bound = 0; + for (const Bounding& b : cfg.bounding) { + if (std::set(b.index.begin(), b.index.end()).size() != b.index.size()) + throw std::runtime_error("Metal bounding requires unique output indices"); + bounds_.push_back({b, upload(b.index), upload(b.total_index)}); + max_bound = std::max(max_bound, static_cast(b.index.size())); + } + const int64_t max_nodes = std::max(G_, H_); + keys_ = arena("keys", max_nodes * C_); values_ = arena("values", max_nodes * C_); + scores_ = arena("scores", std::max({enc_graph_.host.n_edges, proc_graph_.host.n_edges, + dec_graph_.host.n_edges}) * cfg.num_heads); + ln_ = arena("ln", chunk_ * C_); q_ = arena("q", chunk_ * C_); + self_ = arena("self", chunk_ * C_); attention_ = arena("attention", chunk_ * C_); + residual_ = arena("residual", chunk_ * C_); hidden_ = arena("hidden", chunk_ * cfg.mlp_hidden); + arena("input", cfg.n_input_steps * G_ * cfg.n_input_vars); + arena("xn", cfg.n_input_steps * G_ * cfg.n_input_vars); + arena("features", G_ * F_); + for (const char* name : {"z", "latent", "proc.a", "proc.b", "encoder.dest"}) arena(name, H_ * C_); + for (const char* name : {"embed", "block"}) arena(name, chunk_ * C_); + for (const char* name : {"decoder", "output"}) arena(name, G_ * cfg.n_output_vars); + arena("products", G_ * max_bound); + } +} + +id MetalBackend::command() { + if (!command_) command_ = [queue_ commandBuffer]; + if (!command_) throw std::runtime_error("cannot create a Metal command buffer"); + return command_; +} +void MetalBackend::flush() { + encoded_ops_ = 0; + if (!command_) return; + id pending = command_; + command_ = nil; + [pending commit]; + [pending waitUntilCompleted]; + if (pending.status == MTLCommandBufferStatusError) + throw std::runtime_error("Metal execution: " + std::string(pending.error.localizedDescription.UTF8String)); +} +void MetalBackend::finish_op() { + // Bound queued API operations and retained MPS resources without stalling at + // every model phase. This is an operation-count bound, not a GPU time limit. + if (++encoded_ops_ >= 64) flush(); +} +void MetalBackend::dispatch(const char* name, const Args& args, std::initializer_list buffers) { + if (!args.n) return; + id pipeline = kernels_.at(name); + id encoder = [command() computeCommandEncoder]; + if (!encoder) throw std::runtime_error("cannot create a Metal compute encoder"); + [encoder setComputePipelineState:pipeline]; + NSUInteger index = 0; + for (const Tensor& t : buffers) { + [encoder setBuffer:t.buffer offset:t.offset atIndex:index++]; + } + [encoder setBytes:&args length:sizeof(args) atIndex:8]; + const NSUInteger simd = pipeline.threadExecutionWidth; + const NSUInteger width = std::min(256, pipeline.maxTotalThreadsPerThreadgroup) / simd * simd; + const NSUInteger threads = std::strcmp(name, "layernorm") == 0 ? u(int64_t(args.n) * simd) : args.n; + [encoder dispatchThreads:MTLSizeMake(threads, 1, 1) threadsPerThreadgroup:MTLSizeMake(width, 1, 1)]; + [encoder endEncoding]; + finish_op(); +} + +void MetalBackend::linear(const Linear& l, const Tensor& x, int64_t rows, const Tensor& out) { + auto matrix = [](const Tensor& t, int64_t r, int64_t c) { + MPSMatrixDescriptor* descriptor = [MPSMatrixDescriptor matrixDescriptorWithRows:r columns:c + rowBytes:c * sizeof(float) dataType:MPSDataTypeFloat32]; + return [[MPSMatrix alloc] initWithBuffer:t.buffer offset:t.offset descriptor:descriptor]; + }; + const auto key = std::make_tuple(rows, l.out, l.in); + auto& gemm = gemms_[key]; + if (!gemm) gemm = [[MPSMatrixMultiplication alloc] initWithDevice:device_ transposeLeft:NO transposeRight:YES + resultRows:rows resultColumns:l.out interiorColumns:l.in alpha:1.0 beta:0.0]; + if (!gemm) throw std::runtime_error("cannot create an MPS matrix multiplication"); + [gemm encodeToCommandBuffer:command() leftMatrix:matrix(x, rows, l.in) + rightMatrix:matrix(l.w, l.out, l.in) resultMatrix:matrix(out, rows, l.out)]; + finish_op(); + if (l.b.buffer) dispatch("bias_add", {u(rows * l.out), u(l.out)}, {out, l.b}); +} +void MetalBackend::norm(const Norm& l, const Tensor& x, int64_t rows, const Tensor& out) { + Args p{u(rows), u(C_), l.b.buffer ? 1u : 0u}; p.x = cfg_.layer_norm_eps; + dispatch("layernorm", p, {x, l.w, l.b.buffer ? l.b : l.w, out}); +} +void MetalBackend::copy(const Tensor& x, const Tensor& out, int64_t n) { dispatch("copy_array", {u(n)}, {x, out}); } +void MetalBackend::add(const Tensor& x, const Tensor& y, const Tensor& out, int64_t n) { + dispatch("add_arrays", {u(n)}, {x, y, out}); +} + +void MetalBackend::source(const Block& b, const Tensor& x, int64_t rows, int64_t dest_row) { + for (int64_t r = 0; r < rows; r += chunk_) { + const int64_t n = std::min(chunk_, rows - r); + norm(b.src, x.sub(r * C_), n, ln_); + linear(b.k, ln_, n, keys_.sub((dest_row + r) * C_)); + linear(b.v, ln_, n, values_.sub((dest_row + r) * C_)); + } +} +void MetalBackend::destination(const Block& b, const Graph& g, int64_t row, int64_t rows, + const Tensor& x, const Tensor& out) { + norm(b.dst, x, rows, ln_); + linear(b.q, ln_, rows, q_); linear(b.self, ln_, rows, self_); + Args p{0, u(C_), u(cfg_.num_heads), u(cfg_.head_dim), u(cfg_.edge_attr_dim), + b.edge.b.buffer ? 1u : 0u, u(row), u(g.host.row_ptr[row])}; + Tensor bias = b.edge.b.buffer ? b.edge.b : b.edge.w; + p.n = u(int64_t(g.host.row_ptr[row + rows] - g.host.row_ptr[row]) * cfg_.num_heads); + dispatch("attention_scores", p, {q_, keys_, g.sources, g.destinations, g.attrs, b.edge.w, bias, scores_}); + p.n = u(rows * cfg_.num_heads); + dispatch("attention_softmax", p, {g.rows, scores_}); + p.n = u(rows * C_); + dispatch("attention_values", p, {values_, g.rows, g.sources, g.attrs, b.edge.w, bias, scores_, attention_}); + add(attention_, self_, attention_, rows * C_); + linear(b.proj, attention_, rows, residual_); + add(residual_, x, residual_, rows * C_); + norm(b.mlp, residual_, rows, attention_); + linear(b.mlp1, attention_, rows, hidden_); + dispatch("gelu", {u(rows * cfg_.mlp_hidden), cfg_.gelu_tanh ? 1u : 0u}, {hidden_}); + linear(b.mlp2, hidden_, rows, out); + add(out, residual_, out, rows * C_); +} + +void MetalBackend::report(const ActivationHook* hook, const std::string& key, const Tensor& t, + std::vector shape) { + if (!hook) return; + flush(); + (*hook)(key, t.data(), shape); +} + +void MetalBackend::run(const float* x, float* y, const ActivationHook* hook) { + const int64_t T = cfg_.n_input_steps, Vi = cfg_.n_input_vars, Vo = cfg_.n_output_vars; + Tensor input = arena_.at("input"), xn = arena_.at("xn"), features = arena_.at("features"); + std::memcpy(input.data(), x, static_cast(T * G_ * Vi) * sizeof(float)); + if (hook) (*hook)("in.x", x, {1, T, G_, Vi}); + dispatch("normalise", {u(T * G_ * Vi), u(Vi)}, {input, mul_in_, add_in_, xn}); + report(hook, "in.x_norm", xn, {1, T, 1, G_, Vi}); + dispatch("assemble", {u(G_ * F_), u(F_), u(Vi), u(T), u(G_), u(cfg_.node_attr_dim)}, + {xn, data_attr_, features}); + report(hook, "enc.x_src", features, {G_, F_}); + report(hook, "enc.x_dst", hidden_attr_, {H_, cfg_.node_attr_dim}); + + Tensor embed = arena_.at("embed"), block = arena_.at("block"), z = arena_.at("z"); + std::vector enc_full(hook ? static_cast(G_ * C_) : 0); + for (int64_t row = 0; row < G_; row += chunk_) { + const int64_t n = std::min(chunk_, G_ - row); + linear(enc_src_, features.sub(row * F_), n, embed); + if (hook) { flush(); std::memcpy(enc_full.data() + row * C_, embed.data(), n * C_ * sizeof(float)); } + source(enc_, embed, n, row); + if (hook) flush(); + } + if (hook) (*hook)("enc.emb_src", enc_full.data(), {G_, C_}); + Tensor dst = arena_.at("encoder.dest"); + linear(enc_dst_, hidden_attr_, H_, dst); + report(hook, "enc.emb_dst", dst, {H_, C_}); + for (int64_t row = 0; row < H_; row += chunk_) { + const int64_t n = std::min(chunk_, H_ - row); + destination(enc_, enc_graph_, row, n, dst.sub(row * C_), z.sub(row * C_)); + } + report(hook, "enc.out", z, {H_, C_}); + Tensor state = z; + for (int i = 0; i < cfg_.num_processor_layers; ++i) { + Tensor next = arena_.at(i % 2 ? "proc.b" : "proc.a"); + source(processors_[i], state, H_, 0); + for (int64_t row = 0; row < H_; row += chunk_) { + const int64_t n = std::min(chunk_, H_ - row); + destination(processors_[i], proc_graph_, row, n, state.sub(row * C_), next.sub(row * C_)); + } + report(hook, "proc." + std::to_string(i) + ".out", next, {H_, C_}); + state = next; + } + Tensor latent = arena_.at("latent"); + if (cfg_.latent_skip) add(state, z, latent, H_ * C_); else copy(state, latent, H_ * C_); + report(hook, "latent", latent, {H_, C_}); + source(dec_, latent, H_, 0); + Tensor decoded = arena_.at("decoder"); + std::vector dec_embed(hook ? static_cast(G_ * C_) : 0), dec_block(dec_embed.size()); + for (int64_t row = 0; row < G_; row += chunk_) { + const int64_t n = std::min(chunk_, G_ - row); + linear(dec_dst_, features.sub(row * F_), n, embed); + destination(dec_, dec_graph_, row, n, embed, block); + norm(out_norm_, block, n, ln_); + linear(out_proj_, ln_, n, decoded.sub(row * Vo)); + if (hook) flush(); + if (hook) { + std::memcpy(dec_embed.data() + row * C_, embed.data(), n * C_ * sizeof(float)); + std::memcpy(dec_block.data() + row * C_, block.data(), n * C_ * sizeof(float)); + } + } + if (hook) { + (*hook)("dec.emb_dst", dec_embed.data(), {G_, C_}); + (*hook)("dec.block_out", dec_block.data(), {G_, C_}); + } + report(hook, "dec.out", decoded, {G_, Vo}); + Tensor out = arena_.at("output"); + dispatch("residual", {u(G_), u(Vo), u(Vi), u(cfg_.residual_input.size()), u(cfg_.residual_step), u(G_)}, + {xn, decoded, residual_in_, residual_out_, out}); + report(hook, "out.pre_bound", out, {1, 1, 1, G_, Vo}); + for (const Bound& b : bounds_) { + Args p{u(G_ * b.cfg.index.size()), u(Vo), u(b.cfg.index.size()), b.cfg.type == "relu" ? 0u : 1u}; + p.x = b.cfg.min; p.y = b.cfg.max; + dispatch("bound_clamp", p, {out, b.indices}); + if (b.cfg.type == "fraction") { + p.c = u(b.cfg.total_index.size()); + Tensor products = arena_.at("products"); + dispatch("bound_product", p, {out, b.indices, b.totals, products}); + dispatch("bound_assign", p, {products, b.indices, out}); + } + } + report(hook, "out.norm", out, {1, 1, 1, G_, Vo}); + dispatch("denormalise", {u(G_ * Vo), u(Vo)}, {out, mul_out_, add_out_}); + report(hook, "out.y", out, {1, 1, 1, G_, Vo}); + flush(); + std::memcpy(y, out.data(), static_cast(G_ * Vo) * sizeof(float)); +} + +void MetalBackend::predict_step(const float* x, float* y, const ActivationHook* hook) { + @autoreleasepool { + try { run(x, y, hook); } + catch (...) { command_ = nil; encoded_ops_ = 0; throw; } + } +} +} // namespace + +std::unique_ptr make_metal_backend(const Config& cfg, SafeTensorsFile& file, const Options& opts) { + return std::make_unique(cfg, file, opts); +} +} // namespace varda diff --git a/src/varda/metal/kernels.metal b/src/varda/metal/kernels.metal new file mode 100644 index 0000000..fba9fa3 --- /dev/null +++ b/src/varda/metal/kernels.metal @@ -0,0 +1,197 @@ +#include +using namespace metal; +#pragma clang fp contract(off) + +// Matches Args in backend_metal.mm. All indexing is checked to fit uint at load. +struct Args { + uint n, a, b, c, d, e, f, g; + float x, y; +}; + +kernel void normalise(device const float* in [[buffer(0)]], device const float* mul [[buffer(1)]], + device const float* add [[buffer(2)]], device float* out [[buffer(3)]], + constant Args& p [[buffer(8)]], uint i [[thread_position_in_grid]]) { + if (i < p.n) { float v = in[i] * mul[i % p.a]; out[i] = v + add[i % p.a]; } +} + +kernel void assemble(device const float* in [[buffer(0)]], device const float* attr [[buffer(1)]], + device float* out [[buffer(2)]], constant Args& p [[buffer(8)]], + uint i [[thread_position_in_grid]]) { + if (i >= p.n) return; + uint row = i / p.a, col = i % p.a; + out[i] = col < p.b * p.c ? in[((col / p.b) * p.d + row) * p.b + col % p.b] + : attr[row * p.e + col - p.b * p.c]; +} + +kernel void bias_add(device float* out [[buffer(0)]], device const float* bias [[buffer(1)]], + constant Args& p [[buffer(8)]], uint i [[thread_position_in_grid]]) { + if (i < p.n) out[i] = out[i] + bias[i % p.a]; +} + +kernel void add_arrays(device const float* a [[buffer(0)]], device const float* b [[buffer(1)]], + device float* out [[buffer(2)]], constant Args& p [[buffer(8)]], + uint i [[thread_position_in_grid]]) { + if (i < p.n) out[i] = a[i] + b[i]; +} + +kernel void copy_array(device const float* in [[buffer(0)]], device float* out [[buffer(1)]], + constant Args& p [[buffer(8)]], uint i [[thread_position_in_grid]]) { + if (i < p.n) out[i] = in[i]; +} + +// One SIMD group per row, including widths that are not multiples of its size. +// Metal has no fp64: center on a common origin before compensated fp32 moments, +// then normalize centered values directly to avoid cancellation in x*rstd+shift. +kernel void layernorm(device const float* in [[buffer(0)]], device const float* w [[buffer(1)]], + device const float* b [[buffer(2)]], device float* out [[buffer(3)]], + constant Args& p [[buffer(8)]], uint tid [[thread_position_in_grid]], + uint lane [[thread_index_in_simdgroup]], uint simd_width [[threads_per_simdgroup]]) { + uint row = tid / simd_width; + if (row >= p.n) return; + uint base = row * p.a; + const float origin = in[base]; + float sum = 0, correction = 0; + for (uint c = lane; c < p.a; c += simd_width) { + float v = (in[base + c] - origin) - correction, next = sum + v; + correction = (next - sum) - v; sum = next; + } + float mean = simd_sum(sum) / float(p.a); + sum = 0; correction = 0; + for (uint c = lane; c < p.a; c += simd_width) { + float d = (in[base + c] - origin) - mean; + float v = d * d - correction, next = sum + v; + correction = (next - sum) - v; sum = next; + } + float rstd = 1.0f / sqrt(simd_sum(sum) / float(p.a) + p.x); + for (uint c = lane; c < p.a; c += simd_width) { + float v = ((in[base + c] - origin) - mean) * rstd; + v = v * w[c]; + out[base + c] = p.b ? v + b[c] : v; + } +} + +// erf form of GELU, not the tanh approximation. A&S 7.1.26 evaluates erf with +// maximum absolute approximation error 1.5e-7, below fp32 parity tolerances. +float erf_fp32(float x) { + float a = abs(x), t = 1.0f / (1.0f + 0.3275911f * a); + float q = ((((1.061405429f * t - 1.453152027f) * t + 1.421413741f) * t + - 0.284496736f) * t + 0.254829592f) * t; + return copysign(1.0f - q * exp(-a * a), x); +} + +kernel void gelu(device float* values [[buffer(0)]], constant Args& p [[buffer(8)]], + uint i [[thread_position_in_grid]]) { + if (i >= p.n) return; + float v = values[i]; + if (p.a) { + float cube = v * v * v; + values[i] = 0.5f * v * (1.0f + tanh(0.7978845608028654f * (v + 0.044715f * cube))); + } else { + values[i] = v * 0.5f * (1.0f + erf_fp32(v * 0.7071067811865475f)); + } +} + +float edge_value(device const float* attrs, device const float* w, device const float* b, + uint edge, uint channel, constant Args& p) { + float value = attrs[edge * p.d] * w[channel * p.d]; + for (uint a = 1; a < p.d; ++a) value += attrs[edge * p.d + a] * w[channel * p.d + a]; + return p.e ? value + b[channel] : value; +} + +// One score per (edge, head). This stores E*heads values, never an E*channels +// projected-edge tensor. d0 and nd restrict the dispatch to complete CSR rows. +kernel void attention_scores(device const float* q [[buffer(0)]], device const float* k [[buffer(1)]], + device const int* sources [[buffer(2)]], device const int* dests [[buffer(3)]], + device const float* attrs [[buffer(4)]], device const float* w [[buffer(5)]], + device const float* b [[buffer(6)]], device float* scores [[buffer(7)]], + constant Args& p [[buffer(8)]], uint i [[thread_position_in_grid]]) { + if (i >= p.n) return; + uint edge = p.g + i / p.b, head = i % p.b; + uint dst = uint(dests[edge]) - p.f, src = uint(sources[edge]); + float acc[8] = {0,0,0,0,0,0,0,0}; + uint c = 0, offset = head * p.c; + for (; c + 8 <= p.c; c += 8) + for (uint j = 0; j < 8; ++j) { + uint channel = offset + c + j; + float e = edge_value(attrs, w, b, edge, channel, p); + acc[j] += q[dst * p.a + channel] * (k[src * p.a + channel] + e); + } + float sum = ((acc[0]+acc[1])+(acc[2]+acc[3])) + ((acc[4]+acc[5])+(acc[6]+acc[7])); + for (; c < p.c; ++c) { + uint channel = offset + c; + sum += q[dst * p.a + channel] * (k[src * p.a + channel] + edge_value(attrs,w,b,edge,channel,p)); + } + scores[edge * p.b + head] = sum / sqrt(float(p.c)); +} + +kernel void attention_softmax(device const int* rows [[buffer(0)]], device float* scores [[buffer(1)]], + constant Args& p [[buffer(8)]], uint i [[thread_position_in_grid]]) { + if (i >= p.n) return; + uint dst = p.f + i / p.b, head = i % p.b; + int lo = rows[dst], hi = rows[dst + 1]; + if (lo == hi) return; + float largest = scores[uint(lo) * p.b + head]; + for (int e = lo + 1; e < hi; ++e) largest = max(largest, scores[uint(e) * p.b + head]); + float sum = 0; + for (int e = lo; e < hi; ++e) { + float v = exp(scores[uint(e) * p.b + head] - largest); + scores[uint(e) * p.b + head] = v; sum += v; + } + sum = sum + 1e-16f; + for (int e = lo; e < hi; ++e) scores[uint(e) * p.b + head] /= sum; +} + +kernel void attention_values(device const float* v [[buffer(0)]], device const int* rows [[buffer(1)]], + device const int* sources [[buffer(2)]], device const float* attrs [[buffer(3)]], + device const float* w [[buffer(4)]], device const float* b [[buffer(5)]], + device const float* scores [[buffer(6)]], device float* out [[buffer(7)]], + constant Args& p [[buffer(8)]], uint i [[thread_position_in_grid]]) { + if (i >= p.n) return; + uint dst = p.f + i / p.a, channel = i % p.a, head = channel / p.c; + float sum = 0; + for (int e = rows[dst]; e < rows[dst + 1]; ++e) { + float value = v[uint(sources[e]) * p.a + channel] + edge_value(attrs,w,b,uint(e),channel,p); + sum += value * scores[uint(e) * p.b + head]; + } + out[i] = sum; // also defines zero-in-degree destinations +} + +kernel void residual(device const float* in [[buffer(0)]], device const float* dec [[buffer(1)]], + device const int* ins [[buffer(2)]], device const int* outs [[buffer(3)]], + device float* out [[buffer(4)]], constant Args& p [[buffer(8)]], + uint row [[thread_position_in_grid]]) { + if (row >= p.n) return; + for (uint c = 0; c < p.a; ++c) out[row * p.a + c] = dec[row * p.a + c]; + for (uint j = 0; j < p.c; ++j) + out[row * p.a + uint(outs[j])] += in[(p.d * p.e + row) * p.b + uint(ins[j])]; +} + +kernel void bound_clamp(device float* out [[buffer(0)]], device const int* indices [[buffer(1)]], + constant Args& p [[buffer(8)]], uint i [[thread_position_in_grid]]) { + if (i >= p.n) return; + uint at = (i / p.b) * p.a + uint(indices[i % p.b]); + float v = out[at]; + if (p.c == 0) out[at] = v > 0.0f || isnan(v) ? v : 0.0f; + else out[at] = isnan(v) ? v : min(max(v, p.x), p.y); +} + +kernel void bound_product(device const float* in [[buffer(0)]], device const int* indices [[buffer(1)]], + device const int* totals [[buffer(2)]], device float* scratch [[buffer(3)]], + constant Args& p [[buffer(8)]], uint i [[thread_position_in_grid]]) { + if (i < p.n) { + uint row = i / p.b, j = i % p.b; + scratch[i] = in[row * p.a + uint(indices[j])] * in[row * p.a + uint(totals[p.c == 1 ? 0 : j])]; + } +} + +kernel void bound_assign(device const float* scratch [[buffer(0)]], device const int* indices [[buffer(1)]], + device float* out [[buffer(2)]], constant Args& p [[buffer(8)]], + uint i [[thread_position_in_grid]]) { + if (i < p.n) out[(i / p.b) * p.a + uint(indices[i % p.b])] = scratch[i]; +} + +kernel void denormalise(device float* out [[buffer(0)]], device const float* mul [[buffer(1)]], + device const float* add [[buffer(2)]], constant Args& p [[buffer(8)]], + uint i [[thread_position_in_grid]]) { + if (i < p.n) { float v = out[i] - add[i % p.a]; out[i] = v / mul[i % p.a]; } +} diff --git a/src/varda/metal/kernels_source.h.in b/src/varda/metal/kernels_source.h.in new file mode 100644 index 0000000..3d34218 --- /dev/null +++ b/src/varda/metal/kernels_source.h.in @@ -0,0 +1,4 @@ +#pragma once +namespace varda::metal { +inline constexpr const char* kernel_source = R"varda_msl(@VARDA_METAL_SOURCE@)varda_msl"; +} diff --git a/src/varda/metal/metal_backend.h b/src/varda/metal/metal_backend.h new file mode 100644 index 0000000..ee57bbc --- /dev/null +++ b/src/varda/metal/metal_backend.h @@ -0,0 +1,14 @@ +#pragma once +#include +#include + +namespace varda::metal { +struct DeviceInfo { + std::string name; + uint64_t max_buffer_length = 0; + uint64_t recommended_working_set = 0; +}; +// No CPU fallback: false means there is no Metal device with unified memory. +bool available(); +DeviceInfo device_info(); +} diff --git a/src/varda/model.cpp b/src/varda/model.cpp index c4bed81..db264be 100644 --- a/src/varda/model.cpp +++ b/src/varda/model.cpp @@ -8,7 +8,14 @@ namespace varda { -std::string to_string(Device d) { return d == Device::CPU ? "cpu" : "cuda"; } +std::string to_string(Device d) { + switch (d) { + case Device::CPU: return "cpu"; + case Device::CUDA: return "cuda"; + case Device::Metal: return "metal"; + } + throw std::runtime_error("unknown compute device"); +} std::string to_string(Precision p) { return p == Precision::FP32 ? "fp32" : "fp16"; } Model::~Model() = default; @@ -51,12 +58,22 @@ std::unique_ptr Model::load(const std::string& path, const Options& opts) if (opts.precision != Precision::FP32) throw std::runtime_error("the CPU backend runs fp32 only"); m->backend_ = make_cpu_backend(m->cfg_, *m->file_, opts); - } else { + } else if (opts.device == Device::CUDA) { #ifdef VARDA_WITH_CUDA m->backend_ = make_cuda_backend(m->cfg_, *m->file_, opts); #else throw std::runtime_error("varda.cpp was built without CUDA (VARDA_CUDA=OFF)"); #endif + } else if (opts.device == Device::Metal) { +#ifdef VARDA_WITH_METAL + if (opts.precision != Precision::FP32) + throw std::runtime_error("the Metal backend runs fp32 only"); + m->backend_ = make_metal_backend(m->cfg_, *m->file_, opts); +#else + throw std::runtime_error("varda.cpp was built without Metal (VARDA_METAL=OFF)"); +#endif + } else { + throw std::runtime_error("unknown compute device"); } m->file_->require_all_consumed(); return m; diff --git a/src/varda/varda.h b/src/varda/varda.h index 64a2568..a5f3a13 100644 --- a/src/varda/varda.h +++ b/src/varda/varda.h @@ -1,5 +1,5 @@ // varda.cpp public API: load a converted anemoi GraphTransformer forecaster -// (docs/format.md) and run its predict_step on the CPU or a CUDA GPU. +// (docs/format.md) and run its predict_step on CPU, CUDA or an Apple Silicon GPU. #pragma once #include @@ -12,7 +12,7 @@ namespace varda { -enum class Device { CPU, CUDA }; +enum class Device { CPU, CUDA, Metal }; // FP16 reproduces the reference's fp16 autocast (Linear in fp16, LayerNorm and attention // softmax in fp32). CUDA only. enum class Precision { FP32, FP16 }; diff --git a/tests/metal/test_metal_kernels.mm b/tests/metal/test_metal_kernels.mm new file mode 100644 index 0000000..1c74b94 --- /dev/null +++ b/tests/metal/test_metal_kernels.mm @@ -0,0 +1,160 @@ +// The production LayerNorm shader against an independent double-precision +// reference, including cancellation-prone rows and partially occupied groups. +#import +#import + +#include +#include +#include +#include +#include +#include +#include + +#include "varda/metal/kernels_source.h" + +namespace { +struct Args { + uint32_t n, a, b, c = 0, d = 0, e = 0, f = 0, g = 0; + float x = 1e-5f, y = 0; +}; +static_assert(sizeof(Args) == 40, "Metal parameter layout"); + +struct Result { double rel, abs; bool valid; }; + +Result layernorm_case(id device, id queue, + id pipeline, int channels, int rows, + int pattern, bool bias, std::mt19937& rng) { + std::normal_distribution normal(0, 1); + const size_t count = size_t(channels) * rows; + std::vector x(count), w(channels), b(channels); + for (size_t i = 0; i < count; ++i) { + switch (pattern) { + case 0: x[i] = normal(rng); break; + case 1: x[i] = 0.125f; break; + case 2: x[i] = 1000 + 0.01f * normal(rng); break; + default: x[i] = (i % 2 ? 10000 : -10000) + normal(rng); break; + } + } + for (int c = 0; c < channels; ++c) { + w[c] = 0.5f + std::abs(normal(rng)); + b[c] = normal(rng) * 0.2f; + } + const Args args{uint32_t(rows), uint32_t(channels), uint32_t(bias)}; + std::vector reference(count); + for (int r = 0; r < rows; ++r) { + double mean = 0, variance = 0; + for (int c = 0; c < channels; ++c) mean += x[r * channels + c]; + mean /= channels; + for (int c = 0; c < channels; ++c) { + const double delta = x[r * channels + c] - mean; + variance += delta * delta; + } + variance /= channels; + for (int c = 0; c < channels; ++c) + reference[r * channels + c] = (x[r * channels + c] - mean) / + std::sqrt(variance + args.x) * w[c] + (bias ? b[c] : 0); + } + + auto upload = [&](const std::vector& v) { + return [device newBufferWithBytes:v.data() length:v.size() * sizeof(float) + options:MTLResourceStorageModeShared]; + }; + id input = upload(x), weight = upload(w), offset = upload(b); + id output = [device newBufferWithLength:(count + 32) * sizeof(float) + options:MTLResourceStorageModeShared]; + if (!input || !weight || !offset || !output) throw std::runtime_error("Metal test allocation failed"); + float* values = static_cast(output.contents); + constexpr float guard = 1234567.0f; + std::fill_n(values, count + 32, guard); + id command = [queue commandBuffer]; + id encoder = [command computeCommandEncoder]; + if (!command || !encoder) throw std::runtime_error("Metal test encoding failed"); + [encoder setComputePipelineState:pipeline]; + [encoder setBuffer:input offset:0 atIndex:0]; + [encoder setBuffer:weight offset:0 atIndex:1]; + [encoder setBuffer:offset offset:0 atIndex:2]; + [encoder setBuffer:output offset:16 * sizeof(float) atIndex:3]; + [encoder setBytes:&args length:sizeof(args) atIndex:8]; + const NSUInteger simd = pipeline.threadExecutionWidth; + const NSUInteger width = std::min(256, pipeline.maxTotalThreadsPerThreadgroup) / simd * simd; + [encoder dispatchThreads:MTLSizeMake(rows * simd, 1, 1) threadsPerThreadgroup:MTLSizeMake(width, 1, 1)]; + [encoder endEncoding]; + [command commit]; + [command waitUntilCompleted]; + if (command.status == MTLCommandBufferStatusError) + throw std::runtime_error(command.error.localizedDescription.UTF8String); + + double scale = 0, error = 0; + bool valid = true; + for (size_t i = 0; i < count; ++i) { + valid &= std::isfinite(values[i + 16]); + scale = std::max(scale, std::abs(reference[i])); + error = std::max(error, std::abs(double(values[i + 16]) - reference[i])); + } + for (int i = 0; i < 16; ++i) + valid &= values[i] == guard && values[count + 16 + i] == guard; + return {scale > 0 ? error / scale : error, error, valid}; +} +} // namespace + +int main() { + @autoreleasepool { + try { + id device = MTLCreateSystemDefaultDevice(); + if (!device || !device.hasUnifiedMemory) { + std::puts("SKIP: no unified-memory Metal GPU"); + return 77; + } + id queue = [device newCommandQueue]; + if (!queue) throw std::runtime_error("cannot create test command queue"); + MTLCompileOptions* options = [MTLCompileOptions new]; +#if __MAC_OS_X_VERSION_MAX_ALLOWED >= 150000 + if (@available(macOS 15.0, *)) { + options.mathMode = MTLMathModeSafe; + options.mathFloatingPointFunctions = MTLMathFloatingPointFunctionsPrecise; + } else { +#pragma clang diagnostic push +#pragma clang diagnostic ignored "-Wdeprecated-declarations" + options.fastMathEnabled = NO; +#pragma clang diagnostic pop + } +#else + options.fastMathEnabled = NO; +#endif + options.languageVersion = MTLLanguageVersion2_3; + NSError* error = nil; + id library = [device newLibraryWithSource:[NSString stringWithUTF8String:varda::metal::kernel_source] + options:options error:&error]; + if (!library) throw std::runtime_error(error.localizedDescription.UTF8String); + id pipeline = [device newComputePipelineStateWithFunction: + [library newFunctionWithName:@"layernorm"] error:&error]; + if (!pipeline) throw std::runtime_error(error.localizedDescription.UTF8String); + + std::mt19937 rng(731); + int cases = 0, failures = 0; + double worst = 0, worst_abs = 0; + for (int channels : {1, 7, 31, 32, 33, 35, 63, 64, 65, 255, 256, 257, 768, 1024, 2048, 4097}) + for (int rows : {1, 9, 37}) + for (int pattern = 0; pattern < 4; ++pattern) + for (bool bias : {false, true}) { + @autoreleasepool { + const Result r = layernorm_case(device, queue, pipeline, channels, rows, pattern, bias, rng); + ++cases; + if (r.rel > worst) { worst = r.rel; worst_abs = r.abs; } + if (!r.valid || r.rel > 1e-5) { + ++failures; + std::printf("FAIL C=%d rows=%d pattern=%d bias=%d rel_absmax=%.3e max_abs=%.3e\n", + channels, rows, pattern, int(bias), r.rel, r.abs); + } + } + } + std::printf("%s fp32 LayerNorm: %d cases, %d failures; worst rel_absmax %.3e max_abs %.3e tol 1e-5\n", + device.name.UTF8String, cases, failures, worst, worst_abs); + return failures ? 1 : 0; + } catch (const std::exception& e) { + std::fprintf(stderr, "%s\n", e.what()); + return 1; + } + } +} diff --git a/tests/run_synthetic.sh b/tests/run_synthetic.sh index 5da8e25..69eb5c8 100755 --- a/tests/run_synthetic.sh +++ b/tests/run_synthetic.sh @@ -2,7 +2,7 @@ # Synthetic generality suite (docs/testing.md): tiny models outside the real checkpoints' # dimension envelope, checked with varda-parity against a float64 numpy reference. # -# BIN=build-cuda tests/run_synthetic.sh [cpu] [cuda] (no argument: both) +# BIN=build-cuda tests/run_synthetic.sh [cpu] [cuda] [metal] (default: cpu cuda) # # Env: # BIN directory holding varda-parity (default: /build-synth). Required to exist. @@ -52,6 +52,15 @@ fi # ---- runs have_gpu=0 if [ "${SKIP_CUDA:-0}" != 1 ] && command -v nvidia-smi >/dev/null 2>&1 && nvidia-smi -L >/dev/null 2>&1; then have_gpu=1; fi +have_metal=0 +case " $want " in *" metal "*) + [ -x "$BIN/varda-metal-info" ] || { echo "error: build has no varda-metal-info (VARDA_METAL=ON)"; exit 2; } + "$BIN/varda-metal-info" + probe=$? + if [ "$probe" = 0 ]; then have_metal=1 + elif [ "$probe" != 77 ]; then echo "error: Metal device probe failed"; exit 2 + fi +;; esac fail=0 SUMMARY="$OUT/summary.tsv" @@ -102,11 +111,23 @@ for m in $MODELS; do run "$m" cuda-fp16 synthetic-fp16 "" --device cuda --precision fp16 fi ;; esac + case " $want " in *" metal "*) + if [ "$have_metal" = 1 ]; then + run "$m" metal-fp32 synthetic-fp32 "" --device metal --precision fp32 + run "$m" metal-fp32-chunk37 synthetic-fp32 "VARDA_METAL_CHUNK_ROWS=37" --device metal --precision fp32 + fi + ;; esac done echo echo "model run status worst rel_absmax (key) ratio-to-tol" column -t -s$'\t' "$SUMMARY" 2>/dev/null || cat "$SUMMARY" case " $want " in *" cuda "*) [ "$have_gpu" = 1 ] || echo "NOTE: no GPU visible (or SKIP_CUDA=1): the CUDA runs were skipped, this is NOT a full pass." ;; esac +case " $want " in *" metal "*) + if [ "$have_metal" = 0 ]; then + echo "NOTE: no Metal GPU: Metal checks were not run (exit 77), this is NOT a GPU pass." + [ "$fail" = 0 ] && exit 77 + fi +;; esac if [ "$fail" = 0 ]; then echo "synthetic suite: PASS"; else echo "synthetic suite: FAIL (logs in $OUT)"; fi exit "$fail" diff --git a/varda b/varda index ab86bee..4c1e294 100755 --- a/varda +++ b/varda @@ -3,7 +3,7 @@ # # ./varda setup build (build/, CUDA + GRIB) and check the assets bundle (create it here) # ./varda fetch DATE|latest download the inputs of one date into the input cache -# ./varda forecast DATE|latest [--lead-time H] [--precision fp16|fp32] [--device cuda|cpu] [--out DIR] +# ./varda forecast DATE|latest [--lead-time H] [--precision fp16|fp32] [--device cuda|cpu|metal] [--out DIR] # [-- extra varda-forecast options] # # DATE is the initial time in UTC, e.g. 2026-10-22T15 (hour 00, 03, ..., 21); `latest` is the newest KENDA-CH1 @@ -70,16 +70,23 @@ cmd_setup() { done if [[ $rebuild == 1 || ! -x "$BUILD/varda-forecast" ]]; then command -v cmake >/dev/null || die "cmake not found" - say "building into $BUILD (CUDA + GRIB) ..." + say "building engine and GRIB tools into $BUILD ..." cmake -S "$REPO" -B "$BUILD" -DCMAKE_BUILD_TYPE=Release -DVARDA_GRIB=ON >"$BUILD.cmake.log" 2>&1 || { tail -20 "$BUILD.cmake.log" >&2; die "cmake configure failed (log: $BUILD.cmake.log)"; } - cmake --build "$BUILD" -j "$(nproc)" >>"$BUILD.cmake.log" 2>&1 || + local jobs + if [[ "$(uname -s)" == Darwin ]]; then jobs="$(sysctl -n hw.logicalcpu)"; else jobs="$(nproc)"; fi + cmake --build "$BUILD" -j "$jobs" >>"$BUILD.cmake.log" 2>&1 || { tail -30 "$BUILD.cmake.log" >&2; die "build failed (log: $BUILD.cmake.log)"; } mv "$BUILD.cmake.log" "$BUILD/varda-build.log" [[ -x "$BUILD/varda-forecast" ]] || die "the build has no varda-forecast: ecCodes not found? set ECCODES_HOME (log: $BUILD/varda-build.log)" - grep -q 'CMAKE_CUDA_COMPILER:' "$BUILD/CMakeCache.txt" && grep -q '^CMAKE_CUDA_COMPILER:.*nvcc' "$BUILD/CMakeCache.txt" || - say "note: built without CUDA (no nvcc found): only --device cpu will work" + if ! grep -q '^CMAKE_CUDA_COMPILER:.*nvcc' "$BUILD/CMakeCache.txt"; then + if [[ -x "$BUILD/varda-metal-info" ]]; then + say "note: built without CUDA; Metal enabled (--device metal --precision fp32)" + else + say "note: built without CUDA (no nvcc found): only --device cpu will work" + fi + fi say "built $BUILD/varda-forecast" else say "build: $BUILD/varda-forecast present (./varda setup --rebuild to rebuild)" @@ -146,7 +153,10 @@ cmd_forecast() { esac done case "$precision" in fp16|fp32) ;; *) die "--precision must be fp16 or fp32" ;; esac - case "$device" in cuda|cpu) ;; *) die "--device must be cuda or cpu" ;; esac + case "$device" in cuda|cpu|metal) ;; *) die "--device must be cuda, cpu or metal" ;; esac + if [[ "$device" == metal && "$precision" != fp32 ]]; then + die "the Metal backend runs fp32 only: pass --precision fp32" + fi if [[ -n "$lead" ]]; then case "$lead" in *d) lead=$(( ${lead%d} * 24 )) ;; @@ -161,7 +171,11 @@ cmd_forecast() { # where to run: here inside an allocation (or VARDA_SRUN=0), else through srun (decided before any download) local launch=() - if [[ -z "${SLURM_JOB_ID:-}" && "${VARDA_SRUN:-1}" != 0 ]]; then + if [[ "$device" == metal ]]; then + [[ "$(uname -s)" == Darwin ]] || die "Metal requires an Apple Silicon Mac" + [[ -x "$BUILD/varda-metal-info" ]] || die "this build has no Metal backend (VARDA_METAL=ON)" + "$BUILD/varda-metal-info" >/dev/null || die "no usable unified-memory Metal GPU" + elif [[ -z "${SLURM_JOB_ID:-}" && "${VARDA_SRUN:-1}" != 0 ]]; then command -v srun >/dev/null || die "not in a Slurm allocation and no srun: run on a GPU machine with VARDA_SRUN=0" local def="${VARDA_SLURM_GPU:-}"