Skip to main content

Apple Metal native runtime for hash-grid + MLP neural-field training (Python binding ships in v1.0).

Project description

tiny-metal-nn

tiny-metal-nn (tmnn) is a Metal-native, tiny-cuda-nn-inspired C++ library for hash-grid + MLP neural-field training and inference on Apple GPUs.

tiny-cuda-nn is widely used in the CUDA ecosystem (instant-ngp, NeuS, Plenoxels, NerfAcc, and others) but does not run on Apple Silicon. tmnn is an attempt to provide a runtime in that style on Metal — fully fused MLPs, hash-grid encodings, and JSON config compatible with tcnn — without trying to be a general ML framework, a PyTorch / MLX replacement, or a generic GPU compute substrate.

See STATUS.md for what is and is not in the box today, and what we explicitly do not yet claim.

When tmnn might fit

  • You are on Apple Silicon and want a small-network runtime with an API surface close to tcnn (JSON config, fused MLP families, an image-fitting sample), so that a tcnn-based codebase can be ported with relatively small diffs.
  • You want explicit family selection (FullyFusedMetal / TiledMetal / SafeDebugMetal), inspectable planner reasons, manifest-backed kernel prewarm, numerics telemetry, and a frozen optimizer-checkpoint contract — operational surfaces we have intentionally exposed.

For a general Apple-native tensor framework, MLX is broader and more mature. On NVIDIA, tiny-cuda-nn is the natural choice. A side-by-side comparison, including what we have not yet measured, is in docs/VS-MLX-AND-TCNN.md.

Status

  • ~17K LOC C++ source; ~10.5K LOC tests
  • samples/mlp_learning_an_image.cpp — flagship runnable sample, deliberately matches tiny-cuda-nn's filename and target shape so migration diffs stay small
  • Pre-1.0. The v0.1.0 tag is pending external comparison benchmarks (vs MLX, vs tcnn) and at least two more flagship samples (NeRF synthetic, hash-grid SDF fitting). See STATUS.md for the current honest scope and the "Roadmap" section below.

Architecture

JSON config → factory_json::canonicalize_model_config
  → NetworkPlan (planner) → kernel selection + prewarm
  → MetalContext (Metal device + pipeline registry + batch pool)
  → fully-fused or tiled MLP training/inference loop
  → checkpoint contract (frozen, optimizer state portable across versions)

Targets:

Target Purpose
tiny_metal_nn_core header-only public API
tiny_metal_nn_runtime Metal runtime (MetalContext, pipeline registry, batch pool, training-step lifecycle)
tiny_metal_nn_kernels MSL kernel generation (KernelSpec, KernelCompiler, MLPKernelEmitter)
tiny_metal_nn_extensions built-in adapters (DNL, RMHE, 4D, standard SDF, multi-output MLP)

See docs/ARCHITECTURE.md for the design.

Build (C++)

tmnn uses vcpkg for nlohmann_json and gtest. Either place a vcpkg checkout at ./.deps/vcpkg, or point VCPKG_ROOT at an existing one:

git clone https://github.com/microsoft/vcpkg.git ~/vcpkg
~/vcpkg/bootstrap-vcpkg.sh
export VCPKG_ROOT=~/vcpkg

Then build:

git clone <repo-url> tiny-metal-nn
cd tiny-metal-nn
cmake -S . -B build
cmake --build build -j
./build/samples/mlp_learning_an_image                           # flagship sample (built-in config)
./build/samples/mlp_learning_an_image my_config.json            # ...with a custom JSON config
ctest --test-dir build -V                                       # tests

Python binding (optional)

A pybind11 binding ships alongside the C++ library. Install via the project's build backend (scikit-build-core, which invokes the same CMake project with -DTMNN_BUILD_PYTHON_MODULE=ON):

python3.13 -m venv .venv
.venv/bin/pip install scikit-build-core pybind11
VCPKG_ROOT=~/vcpkg .venv/bin/pip install -e ".[dev]" --no-build-isolation
.venv/bin/python -c "import tiny_metal_nn as tmnn; print(tmnn.__version__)"

The editable install reuses a persistent CMake binary directory, so subsequent edits to src/python/ (or any C++ source the binding pulls in) auto-rebuild on the next import tiny_metal_nn. See docs/QUICKSTART.md for a Python "hello world" and docs/ARCHITECTURE.md § 6 for the binding's design.

For users moving from tinycudann, a migration tool (tools/migrate_tcnn.py) handles the mechanical conversion (imports + the canonical 5-line training-loop body); the harder cases (custom losses, non-canonical shapes) are flagged with actionable diagnostics. A worked example pair lives at examples/sphere_sdf/. See docs/TCNN-MIGRATION-GUIDE.md § 10.

For an instrumented build with AddressSanitizer + UndefinedBehaviorSanitizer (off by default; consumers must rebuild any downstream code with the same sanitizer flags):

cmake -S . -B build-asan -DTMNN_ENABLE_SANITIZERS=ON
cmake --build build-asan -j
ctest --test-dir build-asan -j

macOS Apple Silicon is the primary target. Linux / non-Apple builds use a Metal stub (compiles, but GPU-runtime tests skip).

Performance

A standalone GPU-measured benchmark binary lives at tests/benchmarks/tmnn_runtime_benchmarks.cpp. Build and run it on your hardware to get real numbers:

cmake --build build --target tmnn_runtime_benchmarks
./build/tests/tmnn_runtime_benchmarks --smoke    # quick check
./build/tests/tmnn_runtime_benchmarks            # full run

The binary emits planner / autotune-search / hot-step / morton-sort medians on the local Metal device. We are intentionally not publishing pre-measured "reference numbers" in this README until a same-hardware comparison story against MLX (architecture-equivalent) and tiny-cuda-nn (cross-platform) lands; see docs/VS-MLX-AND-TCNN.md for the positioning we can honestly defend today, and STATUS.md for what is still on the roadmap.

Dependencies

External (vcpkg manifest):

  • nlohmann_json ≥ 3.11.3 for JSON config + checkpoint serialization
  • gtest ≥ 1.15.0 (test-only)

Apple frameworks (system, macOS): Metal, Foundation.

Documentation

Roadmap to v0.1.0

The following work is on the path to a stable v0.1.0 tag. Until those land, expect breaking changes on main without deprecation notice:

  1. Comparison benchmark vs MLX on the same hardware / same workload, with architecture-equivalent models on both sides
  2. At least two additional flagship samples (NeRF synthetic, hash-grid SDF fitting)
  3. CI matrix workflows shipped (.github/workflows/test.yml, .github/workflows/release.yml); pending first run + minutes-budget review on a public-repo runner
  4. PyPI listing for the Python wheel (workflow gated off until a maintainer configures trusted publishing)
  5. MAINTAINERS.md with response SLA

For the current honest scope ("what works today" vs "what does not work yet"), see STATUS.md.

Related projects

  • tiny-cuda-nn — the upstream model for tmnn's API surface and sample shape; the natural choice on NVIDIA hardware
  • MLX — Apple's official ML framework; broader scope, the natural choice for general Apple-native tensor work
  • slangcsg (currently private) — sibling project consuming tmnn_core + tmnn_runtime for differentiable CSG / SDF compute

License

Apache-2.0. See LICENSE.

Project details


Download files

Download the file for your platform. If you're not sure which to choose, learn more about installing packages.

Source Distributions

No source distribution files available for this release.See tutorial on generating distribution archives.

Built Distributions

If you're not sure about the file name format, learn more about wheel file names.

tiny_metal_nn-0.1.0a2-cp313-cp313-macosx_14_0_arm64.whl (758.8 kB view details)

Uploaded CPython 3.13macOS 14.0+ ARM64

tiny_metal_nn-0.1.0a2-cp312-cp312-macosx_14_0_arm64.whl (758.7 kB view details)

Uploaded CPython 3.12macOS 14.0+ ARM64

tiny_metal_nn-0.1.0a2-cp311-cp311-macosx_14_0_arm64.whl (755.2 kB view details)

Uploaded CPython 3.11macOS 14.0+ ARM64

tiny_metal_nn-0.1.0a2-cp310-cp310-macosx_14_0_arm64.whl (753.7 kB view details)

Uploaded CPython 3.10macOS 14.0+ ARM64

File details

Details for the file tiny_metal_nn-0.1.0a2-cp313-cp313-macosx_14_0_arm64.whl.

File metadata

File hashes

Hashes for tiny_metal_nn-0.1.0a2-cp313-cp313-macosx_14_0_arm64.whl
Algorithm Hash digest
SHA256 6e540da471aa796a9671fc0e76a9e5529ea5c68e1300d8cb2dcd6bded144f460
MD5 dcb3ee3ef30427da03fc176296872d9e
BLAKE2b-256 002b453f5e94ec00277057f5be9a451cbcf2fede5249f1bc1d6dfb457fd31446

See more details on using hashes here.

Provenance

The following attestation bundles were made for tiny_metal_nn-0.1.0a2-cp313-cp313-macosx_14_0_arm64.whl:

Publisher: release.yml on randallyanh/tiny-metal-nn

Attestations: Values shown here reflect the state when the release was signed and may no longer be current.

File details

Details for the file tiny_metal_nn-0.1.0a2-cp312-cp312-macosx_14_0_arm64.whl.

File metadata

File hashes

Hashes for tiny_metal_nn-0.1.0a2-cp312-cp312-macosx_14_0_arm64.whl
Algorithm Hash digest
SHA256 dfd2bf402417e2544cf915e2b9b9b489cb913e548ef1a3da77f42c73499243bc
MD5 e4b37f018a4dd7ac90f931df857a9322
BLAKE2b-256 ac7955cf3a09d37a11b2eb30029a15d6801e126c5a325bcf4772f1cf614060f1

See more details on using hashes here.

Provenance

The following attestation bundles were made for tiny_metal_nn-0.1.0a2-cp312-cp312-macosx_14_0_arm64.whl:

Publisher: release.yml on randallyanh/tiny-metal-nn

Attestations: Values shown here reflect the state when the release was signed and may no longer be current.

File details

Details for the file tiny_metal_nn-0.1.0a2-cp311-cp311-macosx_14_0_arm64.whl.

File metadata

File hashes

Hashes for tiny_metal_nn-0.1.0a2-cp311-cp311-macosx_14_0_arm64.whl
Algorithm Hash digest
SHA256 20b267c43a397fffdd7270321ad0f16be904227a67624a94de12e246afb6e17c
MD5 7b1d81a9021e3760fe3582d69bf52e55
BLAKE2b-256 0b3aadaaaa77dc82de3923db2d43f89edecc4f7255494ab371f5d3bcb21075f0

See more details on using hashes here.

Provenance

The following attestation bundles were made for tiny_metal_nn-0.1.0a2-cp311-cp311-macosx_14_0_arm64.whl:

Publisher: release.yml on randallyanh/tiny-metal-nn

Attestations: Values shown here reflect the state when the release was signed and may no longer be current.

File details

Details for the file tiny_metal_nn-0.1.0a2-cp310-cp310-macosx_14_0_arm64.whl.

File metadata

File hashes

Hashes for tiny_metal_nn-0.1.0a2-cp310-cp310-macosx_14_0_arm64.whl
Algorithm Hash digest
SHA256 afdb8875e210d0653ae4b805f86afb88b87ef7aa267595b11960ddace78e242d
MD5 b362709a0b17bcecc2d4d47c60b46d84
BLAKE2b-256 5f92153b08c1f77d8433144092dd1772dd634f468e515108f0b4dd8da173aa23

See more details on using hashes here.

Provenance

The following attestation bundles were made for tiny_metal_nn-0.1.0a2-cp310-cp310-macosx_14_0_arm64.whl:

Publisher: release.yml on randallyanh/tiny-metal-nn

Attestations: Values shown here reflect the state when the release was signed and may no longer be current.

Supported by

AWS Cloud computing and Security Sponsor Datadog Monitoring Depot Continuous Integration Fastly CDN Google Download Analytics Pingdom Monitoring Sentry Error logging StatusPage Status page