From 72f4971c5966f51db847e72d9b73cdcc5bdea657 Mon Sep 17 00:00:00 2001 From: Ivy Zheng Date: Thu, 14 Nov 2024 13:00:25 -0800 Subject: [PATCH] Setup the flaxlib in C++, using Meson and Nanobind. Not enabling GitHub auto-build of flaxlib yet - gonna do that later. PiperOrigin-RevId: 696623140 --- .github/workflows/flax_test.yml | 3 -- .gitignore | 4 +++ flaxlib/README.md | 1 - flaxlib/flaxlib/__init__.py | 15 --------- {flaxlib => flaxlib_src}/.gitignore | 0 {flaxlib => flaxlib_src}/Cargo.lock | 0 {flaxlib => flaxlib_src}/Cargo.toml | 0 {flaxlib => flaxlib_src}/LICENSE | 0 flaxlib_src/README.md | 34 ++++++++++++++++++++ {flaxlib/flaxlib => flaxlib_src}/flaxlib.pyi | 0 flaxlib_src/meson.build | 14 ++++++++ {flaxlib => flaxlib_src}/pyproject.toml | 8 ++--- flaxlib_src/src/lib.cc | 14 ++++++++ {flaxlib => flaxlib_src}/src/lib.rs | 0 {flaxlib => flaxlib_src}/uv.lock | 0 tests/flaxlib_test.py | 13 +++++--- 16 files changed, 77 insertions(+), 29 deletions(-) delete mode 100644 flaxlib/README.md delete mode 100644 flaxlib/flaxlib/__init__.py rename {flaxlib => flaxlib_src}/.gitignore (100%) rename {flaxlib => flaxlib_src}/Cargo.lock (100%) rename {flaxlib => flaxlib_src}/Cargo.toml (100%) rename {flaxlib => flaxlib_src}/LICENSE (100%) create mode 100644 flaxlib_src/README.md rename {flaxlib/flaxlib => flaxlib_src}/flaxlib.pyi (100%) create mode 100644 flaxlib_src/meson.build rename {flaxlib => flaxlib_src}/pyproject.toml (67%) create mode 100644 flaxlib_src/src/lib.cc rename {flaxlib => flaxlib_src}/src/lib.rs (100%) rename {flaxlib => flaxlib_src}/uv.lock (100%) diff --git a/.github/workflows/flax_test.yml b/.github/workflows/flax_test.yml index db737ed27d..4bed8d8179 100644 --- a/.github/workflows/flax_test.yml +++ b/.github/workflows/flax_test.yml @@ -108,13 +108,10 @@ jobs: uses: astral-sh/setup-uv@v2 with: version: "0.3.0" - - name: Setup Rust (flaxlib) - uses: actions-rust-lang/setup-rust-toolchain@v1 - name: Install dependencies run: | uv sync --extra all --extra testing --extra docs - uv pip install ./flaxlib - name: Install JAX run: | if [[ "${{ matrix.jax-version }}" == "newest" ]]; then diff --git a/.gitignore b/.gitignore index 2d436c7105..0bc7f3cbe6 100644 --- a/.gitignore +++ b/.gitignore @@ -13,6 +13,10 @@ build/ docs*/**/_autosummary docs*/_build docs*/**/tmp +flaxlib_src/build +flaxlib_src/builddir +flaxlib_src/dist +flaxlib_src/subprojects # used by direnv .envrc diff --git a/flaxlib/README.md b/flaxlib/README.md deleted file mode 100644 index 66910f7e27..0000000000 --- a/flaxlib/README.md +++ /dev/null @@ -1 +0,0 @@ -# flaxlib \ No newline at end of file diff --git a/flaxlib/flaxlib/__init__.py b/flaxlib/flaxlib/__init__.py deleted file mode 100644 index 435dad41b5..0000000000 --- a/flaxlib/flaxlib/__init__.py +++ /dev/null @@ -1,15 +0,0 @@ -# Copyright 2024 The Flax Authors. -# -# Licensed under the Apache License, Version 2.0 (the "License"); -# you may not use this file except in compliance with the License. -# You may obtain a copy of the License at -# -# http://www.apache.org/licenses/LICENSE-2.0 -# -# Unless required by applicable law or agreed to in writing, software -# distributed under the License is distributed on an "AS IS" BASIS, -# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -# See the License for the specific language governing permissions and -# limitations under the License. - -from flaxlib.flaxlib import sum_as_string as sum_as_string diff --git a/flaxlib/.gitignore b/flaxlib_src/.gitignore similarity index 100% rename from flaxlib/.gitignore rename to flaxlib_src/.gitignore diff --git a/flaxlib/Cargo.lock b/flaxlib_src/Cargo.lock similarity index 100% rename from flaxlib/Cargo.lock rename to flaxlib_src/Cargo.lock diff --git a/flaxlib/Cargo.toml b/flaxlib_src/Cargo.toml similarity index 100% rename from flaxlib/Cargo.toml rename to flaxlib_src/Cargo.toml diff --git a/flaxlib/LICENSE b/flaxlib_src/LICENSE similarity index 100% rename from flaxlib/LICENSE rename to flaxlib_src/LICENSE diff --git a/flaxlib_src/README.md b/flaxlib_src/README.md new file mode 100644 index 0000000000..29b4a837c8 --- /dev/null +++ b/flaxlib_src/README.md @@ -0,0 +1,34 @@ +# flaxlib + +## Build flaxlib from source + +Install necessary dependencies to build the C++ based package. + +```shell +pip install meson-python ninja build +``` + +Clone the Flax repository, navigate to the flaxlib source directory. + +```shell +git clone git@github.com:google/flax.git +cd flax/flaxlib_src +``` + +Configure the build. + +```shell +mkdir -p subprojects +meson wrap install robin-map +meson wrap install nanobind +meson setup builddir +``` + +Compile the code. You'll need to run this repeatedly if you modify the source +code. Note that the actual wheel name will differ depending on your system. + +```shell +meson compile -C builddir +python -m build . -w +pip install dist/flaxlib-0.0.1-cp311-cp311-macosx_14_0_arm64.whl --force-reinstall +``` diff --git a/flaxlib/flaxlib/flaxlib.pyi b/flaxlib_src/flaxlib.pyi similarity index 100% rename from flaxlib/flaxlib/flaxlib.pyi rename to flaxlib_src/flaxlib.pyi diff --git a/flaxlib_src/meson.build b/flaxlib_src/meson.build new file mode 100644 index 0000000000..0d78d9436b --- /dev/null +++ b/flaxlib_src/meson.build @@ -0,0 +1,14 @@ +project( + 'flaxlib', + 'cpp', + version: '0.0.1', + default_options: ['cpp_std=c++17'], +) +py = import('python').find_installation() +nanobind_dep = dependency('nanobind', static: true) +py.extension_module( + 'flaxlib', + sources: ['src/lib.cc'], + dependencies: [nanobind_dep], + install: true, +) \ No newline at end of file diff --git a/flaxlib/pyproject.toml b/flaxlib_src/pyproject.toml similarity index 67% rename from flaxlib/pyproject.toml rename to flaxlib_src/pyproject.toml index 993b9703a6..0afc7699a5 100644 --- a/flaxlib/pyproject.toml +++ b/flaxlib_src/pyproject.toml @@ -1,12 +1,12 @@ [build-system] -requires = ["maturin>=1.7,<2.0"] -build-backend = "maturin" +requires = ['meson-python'] +build-backend = 'mesonpy' [project] name = "flaxlib" requires-python = ">=3.10" classifiers = [ - "Programming Language :: Rust", + "Programming Language :: C++", "Programming Language :: Python :: Implementation :: CPython", "Programming Language :: Python :: Implementation :: PyPy", ] @@ -15,5 +15,3 @@ dynamic = ["version"] tests = [ "pytest", ] -[tool.maturin] -features = ["pyo3/extension-module"] diff --git a/flaxlib_src/src/lib.cc b/flaxlib_src/src/lib.cc new file mode 100644 index 0000000000..c714588118 --- /dev/null +++ b/flaxlib_src/src/lib.cc @@ -0,0 +1,14 @@ +#include + +#include "nanobind/nanobind.h" +#include "nanobind/stl/string.h" + +namespace flaxlib { +std::string sum_as_string(int a, int b) { + return std::to_string(a + b); +} + +NB_MODULE(flaxlib, m) { + m.def("sum_as_string", &sum_as_string); +} +} // namespace flaxlib \ No newline at end of file diff --git a/flaxlib/src/lib.rs b/flaxlib_src/src/lib.rs similarity index 100% rename from flaxlib/src/lib.rs rename to flaxlib_src/src/lib.rs diff --git a/flaxlib/uv.lock b/flaxlib_src/uv.lock similarity index 100% rename from flaxlib/uv.lock rename to flaxlib_src/uv.lock diff --git a/tests/flaxlib_test.py b/tests/flaxlib_test.py index dc36f6a21a..c23f70baa7 100644 --- a/tests/flaxlib_test.py +++ b/tests/flaxlib_test.py @@ -12,11 +12,14 @@ # See the License for the specific language governing permissions and # limitations under the License. -from absl.testing import absltest -import flaxlib +# TODO: Re-enable this test after setting up CI build for flaxlib CC. -class TestFlaxlib(absltest.TestCase): +# from absl.testing import absltest +# import flaxlib - def test_flaxlib(self): - self.assertEqual(flaxlib.sum_as_string(1, 2), '3') + +# class TestFlaxlib(absltest.TestCase): + +# def test_flaxlib(self): +# self.assertEqual(flaxlib.sum_as_string(1, 2), '3')