From e59bbc18009f88c76cdc7bd5c80358513445b544 Mon Sep 17 00:00:00 2001 From: Jeremy Kun Date: Thu, 11 Jun 2026 13:37:00 -0700 Subject: [PATCH] Remove bundled openfhe from PYPI wheel This change modifies the PyPI distribution of HEIR to no longer build and ship OpenFHE. The user must now provide a self-provided OpenFHE installation (shared libs and heaers) via environment variables. As a consequence, we updated the CI to: 1. Not run frontend tests in bazel 2. After the main bazel build/test, add new steps that install openfhe, `pip install -e .`, and runs `pytest` to evaluate the frontend tests outside of bazel. I also discovered that more recent manylinux containers have a sanctioned clang, so I updated the cibuildwheel container to use 2_34. We may be able to go back to supporting 2_28 after the hermetic LLVM lands, because I believe they have hermetic glibc that you can pick explicitly. --- .github/build_openfhe.sh | 31 +++ .github/install_cibuildwheel_deps.sh | 30 +- .github/workflows/build_and_test.yml | 21 ++ .github/workflows/nightly.yml | 1 + .github/workflows/release.yml | 3 +- .gitignore | 2 + frontend/BUILD | 115 +------- frontend/README.md | 20 +- frontend/cibuildwheel.py | 51 ++++ frontend/example.py | 152 ---------- frontend/heir/backends/openfhe/__init__.py | 10 +- frontend/heir/backends/openfhe/backend.py | 5 +- frontend/heir/backends/openfhe/config.py | 260 +++++++++++++----- frontend/heir/backends/openfhe/config_test.py | 141 ++++++++++ frontend/heir/backends/util/common.py | 54 +++- frontend/heir/backends/util/cpp_compiler.py | 10 +- frontend/heir/heir_cli/heir_cli_config.py | 100 +++++-- frontend/heir/mlir/types.py | 8 +- frontend/heir/pipeline.py | 7 +- frontend/import_isolation_test.py | 81 ++++-- frontend/sbox_test.py | 156 ----------- frontend/testing.bzl | 49 ---- lib/Target/OpenFhePke/BUILD | 4 + pyproject.toml | 27 +- setup.py | 106 ++----- tests/Examples/openfhe/test.bzl | 2 +- 26 files changed, 731 insertions(+), 715 deletions(-) create mode 100755 .github/build_openfhe.sh create mode 100644 frontend/cibuildwheel.py delete mode 100644 frontend/example.py create mode 100644 frontend/heir/backends/openfhe/config_test.py delete mode 100644 frontend/sbox_test.py delete mode 100644 frontend/testing.bzl diff --git a/.github/build_openfhe.sh b/.github/build_openfhe.sh new file mode 100755 index 0000000000..a82c26dd89 --- /dev/null +++ b/.github/build_openfhe.sh @@ -0,0 +1,31 @@ +#!/usr/bin/env bash +set -e + +# Clone OpenFHE +if [ ! -d "openfhe-development" ]; then + git clone --depth 1 --branch v1.4.2 https://github.com/openfheorg/openfhe-development.git +fi + +# Build OpenFHE +mkdir -p openfhe-development/build +cd openfhe-development/build + +INSTALL_DIR="$(pwd)/../../openfhe-install" +mkdir -p "$INSTALL_DIR" + +cmake .. \ + -DCMAKE_BUILD_TYPE=Release \ + -DBUILD_UNITTESTS=OFF \ + -DBUILD_BENCHMARKS=OFF \ + -DBUILD_EXAMPLES=OFF \ + -DCMAKE_INSTALL_PREFIX="$INSTALL_DIR" + +make -j$(nproc 2>/dev/null || sysctl -n hw.ncpu 2>/dev/null || echo 2) +make install + +# Export variables to GITHUB_ENV if running in GitHub Actions +if [ -n "$GITHUB_ENV" ]; then + echo "OPENFHE_LIB_DIR=$INSTALL_DIR/lib" >> "$GITHUB_ENV" + INC_DIR="$INSTALL_DIR/include/openfhe" + echo "OPENFHE_INCLUDE_DIR=$INSTALL_DIR/include:$INC_DIR:$INC_DIR/binfhe:$INC_DIR/core:$INC_DIR/pke" >> "$GITHUB_ENV" +fi diff --git a/.github/install_cibuildwheel_deps.sh b/.github/install_cibuildwheel_deps.sh index 30a9efbbe3..933cd4142f 100644 --- a/.github/install_cibuildwheel_deps.sh +++ b/.github/install_cibuildwheel_deps.sh @@ -1,10 +1,5 @@ #!/usr/bin/env bash - -# Install clang and lld (container only has gcc by default) -# Pin to clang 20 — some HEIR deps don't yet support clang 21+ -# (cf. https://github.com/google/heir/issues/2675). -CLANG_VERSION=20 -yum install -y "clang-${CLANG_VERSION}*" "lld-${CLANG_VERSION}*" +set -e # Install bazel if ! bazel version; then @@ -16,7 +11,24 @@ if ! bazel version; then bazel_version=$(cat /project/.bazelversion) curl -L -o $HOME/bin/bazel --create-dirs "https://github.com/bazelbuild/bazel/releases/download/${bazel_version}/bazel-${bazel_version}-linux-${arch}" chmod +x $HOME/bin/bazel -else - # Bazel is installed for the correct architecture - exit 0 fi + +manylinux-install-clang -v v20.1.8.0 + +# Install OpenMP development headers using the CRB repository +dnf install -y --enablerepo=crb libomp-devel + +# Symlink OpenMP headers from system Clang to static Clang's include directory +for system_clang_dir in /usr/lib/clang/* /usr/lib64/clang/*; do + [ -d "$system_clang_dir/include" ] || continue + for f in "$system_clang_dir/include"/omp*.h; do + [ -f "$f" ] || continue + for static_clang_include in /opt/clang/lib/clang/*/include; do + [ -d "$static_clang_include" ] || continue + target="$static_clang_include/$(basename "$f")" + if [ ! -e "$target" ]; then + ln -sf "$f" "$target" + fi + done + done +done diff --git a/.github/workflows/build_and_test.yml b/.github/workflows/build_and_test.yml index 735f734551..6ef3946b8b 100644 --- a/.github/workflows/build_and_test.yml +++ b/.github/workflows/build_and_test.yml @@ -42,6 +42,7 @@ jobs: ls -al /etc/alternatives/clang bazel cquery --output=starlark --starlark:expr="str(providers(target))" @bazel_tools//tools/cpp:current_cc_toolchain | sed 's/,/\n/g' | grep -C 5 clang + # Bazel tests - name: "Run `bazel build`" run: | bazel build -c opt //... @@ -49,3 +50,23 @@ jobs: - name: "Run `bazel test`" run: | bazel test -c opt //... + + # Frontend tests + - name: Build OpenFHE + run: | + ./.github/build_openfhe.sh + + - name: Install Python 3.12 + uses: actions/setup-python@a309ff8b426b58ec0e2a45f0f869d46889d02405 # v6.2.0 + with: + python-version: "3.12" + + - name: Install python package in editable mode + run: | + python -m pip install --upgrade pip + # This should reuse the pre-built binaries from `bazel build` above. + pip install -e ".[dev,python,openfhe]" + + - name: Run Python frontend tests + run: | + pytest frontend/ diff --git a/.github/workflows/nightly.yml b/.github/workflows/nightly.yml index ef19efbb31..20fdbf3e3c 100644 --- a/.github/workflows/nightly.yml +++ b/.github/workflows/nightly.yml @@ -27,6 +27,7 @@ jobs: github.event_name == 'schedule' runs-on: ${{ matrix.runner }} strategy: + fail-fast: false matrix: include: - name: linux-x86_64 diff --git a/.github/workflows/release.yml b/.github/workflows/release.yml index 6338f84f6e..65cf613d75 100644 --- a/.github/workflows/release.yml +++ b/.github/workflows/release.yml @@ -89,6 +89,7 @@ jobs: github.event.label.name == 'ci:wheels' runs-on: ${{ matrix.runner }} strategy: + fail-fast: false matrix: include: - name: linux-x86_64 @@ -111,7 +112,7 @@ jobs: - run: pip install --upgrade pip uv - name: Build wheels on ${{ matrix.name }} using cibuildwheel - uses: pypa/cibuildwheel@5f22145df44122af0f5a201f93cf0207171beca7 # v3.0.0 + uses: pypa/cibuildwheel@294735312765b09d24a2fbec22660ce817587d55 # v4.1.0 env: SETUPTOOLS_SCM_PRETEND_VERSION: ${{ needs.setup-version.outputs.version }} CIBW_ENVIRONMENT_PASS: SETUPTOOLS_SCM_PRETEND_VERSION diff --git a/.gitignore b/.gitignore index 26bddcd05a..0f861c8694 100644 --- a/.gitignore +++ b/.gitignore @@ -58,3 +58,5 @@ MODULE.bazel.lock scripts/torch_ci/failure_logs scripts/torch_ci/models/*.mlir +openfhe-install/ +openfhe-development/ diff --git a/frontend/BUILD b/frontend/BUILD index 94ec875a8c..ec15befa6b 100644 --- a/frontend/BUILD +++ b/frontend/BUILD @@ -1,7 +1,4 @@ -load("@bazel_skylib//:bzl_library.bzl", "bzl_library") -load("@heir//frontend:testing.bzl", "frontend_test") load("@rules_python//python:py_library.bzl", "py_library") -load("@rules_python//python:py_test.bzl", "py_test") package( default_applicable_licenses = ["@heir//:license"], @@ -9,16 +6,9 @@ package( ) DATA_DEPS = [ - "@cereal//:headers", "@heir//tools:heir-opt", "@heir//tools:heir-translate", - "@openfhe//:libopenfhe", - "@openfhe//:headers", - # copybara: openfhe_binfhe_headers - # copybara: openfhe_core_headers - # copybara: openfhe_pke_headers # copybara: python_runtime_headers - "@rapidjson", ] # a single-source build dependency that gives the whole (non-test) source tree; @@ -42,103 +32,8 @@ py_library( ], ) -py_test( - name = "import_isolation_test", - srcs = ["import_isolation_test.py"], - tags = [ - # copybara: manual - "notap", - ], - deps = [ - ":frontend", - "@abseil-py//absl/testing:absltest", - "@heir_pip_deps//numba", - "@heir_pip_deps//numpy", - "@heir_pip_deps//pybind11", - ], -) - -frontend_test( - name = "e2e_test", - srcs = ["e2e_test.py"], - tags = [ - # copybara: manual - "notap", - ], -) - -frontend_test( - name = "loop_test", - srcs = ["loop_test.py"], - tags = [ - # copybara: manual - "notap", - ], -) - -frontend_test( - name = "mixed_bitwidth_test", - srcs = ["mixed_bitwidth_test.py"], - tags = [ - # copybara: manual - "notap", - ], -) - -frontend_test( - name = "tensor_test", - srcs = ["tensor_test.py"], - tags = [ - # copybara: manual - "notap", - ], -) - -frontend_test( - name = "tensor_loop_test", - srcs = ["tensor_loop_test.py"], - tags = [ - # copybara: manual - "notap", - ], -) - -frontend_test( - name = "cggi_test", - srcs = ["cggi_test.py"], - tags = [ - # copybara: manual - "notap", - ], -) - -frontend_test( - name = "cast_test", - srcs = ["cast_test.py"], - tags = [ - # copybara: manual - "notap", - ], -) - -frontend_test( - name = "unique_name_test", - srcs = ["unique_name_test.py"], - tags = [ - # copybara: manual - "notap", - ], -) - -bzl_library( - name = "testing_bzl", - srcs = ["testing.bzl"], - visibility = ["//visibility:public"], - deps = ["@rules_python//python:defs_bzl"], -) - -py_library( - name = "example", - srcs = ["example.py"], - deps = [":frontend"], -) +# Because tests require JIT-compiling generated code on the fly, and linking +# with system-provided libs, we cannot run them inside bazel's hermetic C++ +# toolchain. Instead the tests are run as a separate CI test, after installing +# the frontend package, using pytest to discover and run tests named +# frontend/**/*_test.py diff --git a/frontend/README.md b/frontend/README.md index c4d4bb4334..0ddd3a6e02 100644 --- a/frontend/README.md +++ b/frontend/README.md @@ -58,24 +58,26 @@ auto-detected resources. - Defaults to `bazel-bin/tools/heir-{opt,translate}`. - Cf. `heir/backends/openfhe/config.py` for more details. -- OpenFHE installation locations (default to where `cmake` installs them in the - OpenFHE development repo). +- OpenFHE configuration overrides. Note: If any OpenFHE environment variable is + set, BOTH `OPENFHE_LIB_DIR` and `OPENFHE_INCLUDE_DIR` must be provided + (enforcing atomic overrides), or an error will be raised. - `OPENFHE_LIB_DIR`: a string containing the directory containing the OpenFHE - .so files. Usually `/usr/local/lib` + .so files. Usually `/usr/local/lib` or `openfhe/lib`. - `OPENFHE_INCLUDE_DIR`: a colon-separated string of directories containing - OpenFHE headers. Note this usually requires four different paths due to how - OpenFHE organizes its imports, one for each of the three main subdirectories - of the project. + OpenFHE headers. For a standard installation, this usually requires four + paths: - `/usr/local/include/openfhe` - `/usr/local/include/openfhe/binfhe` - `/usr/local/include/openfhe/core` - `/usr/local/include/openfhe/pke` - - `OPENFHE_LINK_LIBS`: a colon-separated string of libraries to link against - (without `lib` or `.so`). E.g., `"OPENFHEbinfhe:OPENFHEcore:OPENFHEpke"`. + - `OPENFHE_LINK_LIBS`: a colon-separated string of libraries to link against. + Accepts short names (e.g. `"openfhe"`) OR exact library filenames prefixed + by a colon (e.g. `":libOPENFHEcore.so.1"`). Defaults to split libraries: + `":libOPENFHEcore.so.1::libOPENFHEpke.so.1::libOPENFHEbinfhe.so.1"`. - `OPENFHE_INCLUDE_TYPE`: a string indicating the include path type to use (see options on `heir-translate --emit-openfhe`). Should be - `install-relative` for a system-wide OpenFHE installation. + `install-relative` for a system-wide or PyPI wheel OpenFHE installation. ## Formatting diff --git a/frontend/cibuildwheel.py b/frontend/cibuildwheel.py new file mode 100644 index 0000000000..f4172012cf --- /dev/null +++ b/frontend/cibuildwheel.py @@ -0,0 +1,51 @@ +"""Minimal (dependency-free) test of PyPI for cibuildwheel.""" + +from heir import compile +from heir.mlir import I16, Secret +from heir.backends.cleartext import CleartextBackend + + +# Custom Pipeline Example +def custom_example(): + print("Running custom pipeline example") + + @compile( + heir_opt_options=[ + "--mlir-to-secret-arithmetic", + "--canonicalize", + "--cse", + ], + backend=CleartextBackend(), + debug=True, + ) + def custom(x: Secret[I16], y: Secret[I16]): + return (x + y) * (x - y) + (x * y) + + +# MLIR-input example +def mlir_example(): + print("Running MLIR-input example") + # Input should be a single function, just like normal MLIR input to heir-opt + src = """ + func.func @myfunc(%a : i32 {secret.secret}, %b : i32) -> i32 { + %sum = arith.addi %a, %b : i32 + return %sum : i32 + } + """ + # By passing mlir_str instead of using it to decorate a function, + # we can skip the python parsing/type inference/etc stages. + compile( + mlir_str=src, + scheme="bgv", + backend=CleartextBackend(), + debug=True, + ) + + +def main(): + custom_example() + mlir_example() + + +if __name__ == "__main__": + main() diff --git a/frontend/example.py b/frontend/example.py deleted file mode 100644 index d3c567d6f6..0000000000 --- a/frontend/example.py +++ /dev/null @@ -1,152 +0,0 @@ -"""Example of HEIR Python usage.""" - -from heir import compile -from heir.mlir import F32, I16, I64, Secret -from heir.backends.cleartext import CleartextBackend - -# TODO (#1162): Remove the need for wrapper functions around each `@compile`-d function to isolate backend pybindings - - -### Simple Example -def simple_example(): - print("Running simple example") - - @compile() # defaults to scheme="bgv", OpenFHE backend, and debug=False - def func(x: Secret[I16], y: Secret[I16]): - sum = x + y - diff = x - y - mul = x * y - expression = sum * diff + mul - deadcode = expression * mul - return expression - - print( - f"Expected result for `func`: {func.original(7,8)}, FHE result:" - f" {func(7,8)}" - ) - - -### Manual setup/enc/dec example -def manual_example(): - print("Running manual example") - - @compile() - def foo(x: Secret[I16], y: Secret[I16]): - return (x + y) * (x - y) + (x * y) - - foo.setup() # runs keygen/etc - enc_x = foo.encrypt_x(7) - enc_y = foo.encrypt_y(8) - result_enc = foo.eval(enc_x, enc_y) - result = foo.decrypt_result(result_enc) - print( - f"Expected result for `foo`: {foo.original(7,8)}, " - f"decrypted FHE result: {result}" - ) - - -### Loop Example -def loop_example(): - print("Running loop example") - - @compile() - def loop_test(a: Secret[I64]): - """An example function with a static loop.""" - result = 2 - for i in range(3): - result = a + result - return result - - print( - f"Expected result for `loop_test`: {loop_test(2)}, " - f"FHE result: {loop_test(2)}" - ) - - -### CKKS Example -def ckks_example(): - print("Running CKKS example") - - @compile(scheme="ckks") - def bar(x: Secret[F32], y: Secret[F32]): - return (x + y) * (x - y) + (x * y) - - print( - f"Expected result for `bar`: {bar.original(0.7,0.8)}, FHE result:" - f" {bar(0.7,0.8)}" - ) - - -### Ciphertext-Plaintext Example -def ctxt_ptxt_example(): - print("Running ciphertext-plaintext example") - - @compile() - def baz(x: Secret[I16], y: Secret[I16], z: I16): - ptxt_mul = x * z - ctxt_mul = x * x - ctxt_mul2 = y * y - add = ctxt_mul + ctxt_mul2 - return ptxt_mul + add - - print( - f"Expected result for `baz`: {baz.original(7,8,9)}, " - f"FHE result: {baz(7,8,9)}" - ) - - -### Custom Pipeline Example -def custom_example(): - print("Running custom pipeline example") - - @compile( - heir_opt_options=[ - "--mlir-to-secret-arithmetic", - "--canonicalize", - "--cse", - ], - backend=CleartextBackend(), # just runs the python function when `custom(...)` is called - debug=True, # so that we can see the file that contains the output of the pipeline - ) - def custom(x: Secret[I16], y: Secret[I16]): - return (x + y) * (x - y) + (x * y) - - print( - "CleartextBackend simply runs the original python function:" - f" {custom(7,8)}" - ) - - -### MLIR-input example -def mlir_example(): - print("Running MLIR-input example") - # Input should be a single function, just like normal MLIR input to heir-opt - src = """ - func.func @myfunc(%a : i32 {secret.secret}, %b : i32) -> i32 { - %sum = arith.addi %a, %b : i32 - return %sum : i32 - } - """ - # By passing mlir_str instead of using it to decorate a function, - # we can skip the python parsing/type inference/etc stages. - myfunc = compile(mlir_str=src, scheme="bgv", debug=True) - - print( - # Note that there's no 'myfunc.original' since - # there's no original Python function to call - f"Expected result for `myfunc`: {7 + 8}, FHE result: {myfunc(7,8)}" - ) - - -def main(): - simple_example() - manual_example() - loop_example() - ckks_example() - ctxt_ptxt_example() - custom_example() - mlir_example() - - -if __name__ == "__main__": - main() diff --git a/frontend/heir/backends/openfhe/__init__.py b/frontend/heir/backends/openfhe/__init__.py index 580d5e2141..fcc53de28d 100644 --- a/frontend/heir/backends/openfhe/__init__.py +++ b/frontend/heir/backends/openfhe/__init__.py @@ -1,9 +1,15 @@ from .backend import OpenFHEBackend -from .config import DEFAULT_INSTALLED_OPENFHE_CONFIG, OpenFHEConfig, from_os_env +from .config import ( + get_default_installed_config, + OpenFHEConfig, + from_os_env, + resolve_config, +) __all__ = [ "OpenFHEBackend", "OpenFHEConfig", - "DEFAULT_INSTALLED_OPENFHE_CONFIG", + "get_default_installed_config", "from_os_env", + "resolve_config", ] diff --git a/frontend/heir/backends/openfhe/backend.py b/frontend/heir/backends/openfhe/backend.py index b02af400fc..d24c4c21cf 100644 --- a/frontend/heir/backends/openfhe/backend.py +++ b/frontend/heir/backends/openfhe/backend.py @@ -260,6 +260,7 @@ def run_backend( cpp_compiler_backend = cpp_compiler.CppCompilerBackend() so_filepath = Path(workspace_dir) / f"{func_name}.so" linker_search_paths = [self.openfhe_config.lib_dir] + active_include_dirs = self.openfhe_config.include_dirs def debug_printer(args): print( @@ -271,7 +272,7 @@ def debug_printer(args): cpp_compiler_backend.compile_to_shared_object( cpp_source_filepath=cpp_filepath, shared_object_output_filepath=so_filepath, - include_paths=self.openfhe_config.include_dirs, + include_paths=active_include_dirs, linker_search_paths=linker_search_paths, link_libs=self.openfhe_config.link_libs, arg_printer=debug_printer if debug else None, @@ -282,7 +283,7 @@ def debug_printer(args): cpp_compiler_backend.compile_to_shared_object( cpp_source_filepath=pybind_filepath, shared_object_output_filepath=pybind_so_filepath, - include_paths=self.openfhe_config.include_dirs + include_paths=active_include_dirs + pybind11_includes() + [workspace_dir], linker_search_paths=linker_search_paths + pybind11_libs(), diff --git a/frontend/heir/backends/openfhe/config.py b/frontend/heir/backends/openfhe/config.py index 8c1e33b732..ddcaba2f23 100644 --- a/frontend/heir/backends/openfhe/config.py +++ b/frontend/heir/backends/openfhe/config.py @@ -2,8 +2,17 @@ import dataclasses import importlib.resources +import importlib.util import os +import pathlib +import platform +import sys +import sysconfig +from typing import Optional from heir.backends.util.common import get_repo_root, is_pip_installed +import traceback + +Path = pathlib.Path dataclass = dataclasses.dataclass @@ -28,19 +37,39 @@ class OpenFHEConfig: link_libs: list[str] -DEFAULT_INSTALLED_OPENFHE_CONFIG = OpenFHEConfig( - include_dirs=[ - "/usr/local/include/openfhe", - "/usr/local/include/openfhe/binfhe", - "/usr/local/include/openfhe/core", - "/usr/local/include/openfhe/pke", - ], - include_type="install-relative", - lib_dir="/usr/local/lib", - link_libs=[ - "openfhe", # libopenfhe.so - ], -) +def get_default_installed_config() -> Optional[OpenFHEConfig]: + prefixes = [Path("/usr/local"), Path("/usr")] + for prefix in prefixes: + include_base = prefix / "include" / "openfhe" + if include_base.is_dir(): + lib_dir = None + for lib_name in ["lib", "lib64"]: + candidate_lib = prefix / lib_name + if candidate_lib.is_dir(): + libs = _resolve_link_libs(candidate_lib) + if libs != ["openfhe"]: + lib_dir = candidate_lib + break + + if not lib_dir: + lib_dir = prefix / "lib" + + include_dirs = [ + str(include_base), + str(include_base / "binfhe"), + str(include_base / "core"), + str(include_base / "pke"), + str(prefix / "include"), + ] + include_dirs = [d for d in include_dirs if os.path.exists(d)] + + return OpenFHEConfig( + include_dirs=include_dirs, + include_type="install-relative", + lib_dir=str(lib_dir), + link_libs=_resolve_link_libs(lib_dir), + ) + return None def development_openfhe_config() -> OpenFHEConfig: @@ -86,7 +115,55 @@ def development_openfhe_config() -> OpenFHEConfig: ) -def from_os_env(debug=False) -> OpenFHEConfig: +def _resolve_link_libs(lib_dir: os.PathLike | str) -> list[str]: + lib_path = Path(lib_dir) + if not lib_path.is_dir(): + return ["openfhe"] + + core_lib = None + pke_lib = None + binfhe_lib = None + + for p in lib_path.iterdir(): + if not p.is_file(): + continue + name = p.name.lower() + if not (".so" in name or ".dylib" in name): + continue + + if "openfhe" in name: + if "core" in name: + core_lib = p + elif "pke" in name: + pke_lib = p + elif "binfhe" in name: + binfhe_lib = p + + if core_lib and pke_lib and binfhe_lib: + return [ + str(core_lib.resolve()), + str(pke_lib.resolve()), + str(binfhe_lib.resolve()), + ] + + link_libs = [] + for p in lib_path.iterdir(): + if not p.is_file(): + continue + name = p.name.lower() + if not (".so" in name or ".dylib" in name): + continue + if "openfhe" in name and not any( + x in name for x in ["core", "pke", "binfhe"] + ): + link_libs.append(str(p.resolve())) + + if link_libs: + return link_libs + return ["openfhe"] + + +def from_os_env(debug: bool = False) -> Optional[OpenFHEConfig]: """Create an OpenFHEConfig from environment variables. Note, this is required for running tests under bazel, as the openfhe @@ -109,23 +186,31 @@ def from_os_env(debug=False) -> OpenFHEConfig: debug: whether to print debug information Returns: - the OpenFHEConfig + the OpenFHEConfig or None if no OPENFHE environment variables are set. """ + trigger_keys = {"OPENFHE_INCLUDE_DIR", "OPENFHE_LIB_DIR", "OPENFHE_LINK_LIBS"} + if not any(os.environ.get(k) for k in trigger_keys): + return None + + env_keys = [k for k in os.environ.keys() if "OPENFHE" in k] + if debug: print("Env:") print(f"RUNFILES_DIR: {os.environ.get('RUNFILES_DIR', '')}") - for k, v in os.environ.items(): - if "OPENFHE" in k: - print(f"{k}: {v}") + for k in env_keys: + print(f"{k}: {os.environ.get(k)}") - include_dirs = os.environ.get("OPENFHE_INCLUDE_DIR", "").split(":") - include_type = os.environ.get("OPENFHE_INCLUDE_TYPE", "") + inc_str = os.environ.get("OPENFHE_INCLUDE_DIR", "") lib_dir = os.environ.get("OPENFHE_LIB_DIR", "") - link_libs = os.environ.get("OPENFHE_LINK_LIBS", "").split(":") - # remove empty strings from lists - include_dirs = [dir for dir in include_dirs if dir] - link_libs = [lib for lib in link_libs if lib] + if not inc_str or not lib_dir: + raise ValueError( + "Both OPENFHE_INCLUDE_DIR and OPENFHE_LIB_DIR must be set when" + f" overriding OpenFHE configuration. Found keys: {env_keys}" + ) + + include_dirs = inc_str.split(":") + include_type = os.environ.get("OPENFHE_INCLUDE_TYPE", "install-relative") # Special case for bazel, RUNFILES_DIR is in OSS, TEST_SRCDIR is # for Google-internal testing. @@ -138,59 +223,98 @@ def from_os_env(debug=False) -> OpenFHEConfig: # to $RUNFILES_DIR/openfhe/src/... include_dirs = [os.path.join(path_base, dir) for dir in include_dirs] + link_libs_str = os.environ.get("OPENFHE_LINK_LIBS", "") + if link_libs_str: + link_libs = link_libs_str.split(":") + else: + link_libs = _resolve_link_libs(Path(lib_dir)) + + extra_include_dirs = [] + for d in include_dirs: + p = Path(d) + if p.name == "openfhe": + parent_dir = str(p.parent) + if ( + parent_dir not in include_dirs + and parent_dir not in extra_include_dirs + ): + extra_include_dirs.append(parent_dir) + include_dirs.extend(extra_include_dirs) + + # remove empty strings from lists + include_dirs = [dir for dir in include_dirs if dir] + link_libs = [lib for lib in link_libs if lib] + + if not include_dirs or not link_libs: + raise ValueError( + "Invalid OpenFHE configuration: include_dirs or link_libs resolved to" + f" empty. Found keys: {env_keys}" + ) + for include_dir in include_dirs: if not os.path.exists(include_dir): print( f'Warning: OpenFHE include directory "{include_dir}" does not exist' ) - # If something has been found from the environment variables, return it - if include_dirs: - return OpenFHEConfig( - include_dirs=include_dirs, - lib_dir=lib_dir, - link_libs=link_libs, - include_type=include_type, - ) + return OpenFHEConfig( + include_dirs=include_dirs, + lib_dir=lib_dir, + link_libs=link_libs, + include_type=include_type, + ) - # if nothing is found, check the default installed config - if debug: + +def resolve_config(debug: bool = False) -> OpenFHEConfig: + """Resolve OpenFHEConfig in cascading order of preference.""" + debug_active = debug or os.environ.get("OPENFHE_DEBUG") == "1" + + if debug_active: + print("HEIRpy Debug (config): Starting resolve_config", file=sys.stderr) + + if debug_active: + print("HEIRpy Debug (config): Attempting from_os_env()", file=sys.stderr) + os_env_config = from_os_env(debug=debug_active) + if os_env_config is not None: + if debug_active: + print( + "HEIRpy Debug (config): Successfully resolved from_os_env():" + f" {os_env_config}", + file=sys.stderr, + ) + return os_env_config + else: + if debug_active: + print( + "HEIRpy Debug (config): from_os_env() returned None", file=sys.stderr + ) + + if debug_active: print( - "HEIRpy Debug (OpenFHE Backend): No valid OpenFHE config found in" - " environment variables, trying default install location." + "HEIRpy Debug (config): Attempting get_default_installed_config()", + file=sys.stderr, ) - if os.path.exists(DEFAULT_INSTALLED_OPENFHE_CONFIG.include_dirs[0]): - return DEFAULT_INSTALLED_OPENFHE_CONFIG + default_config = get_default_installed_config() + if default_config is not None: + if debug_active: + print( + "HEIRpy Debug (config): Successfully resolved" + " get_default_installed_config():" + f" {default_config}", + file=sys.stderr, + ) + return default_config + else: + if debug_active: + print( + "HEIRpy Debug (config): get_default_installed_config() returned None", + file=sys.stderr, + ) - # if nothing is found still, check the development config - if debug: + if debug_active: print( - "HEIRpy Debug (OpenFHE Backend): No valid OpenFHE config found in" - " environment variables or default install location, trying" - " development location." + "HEIRpy Debug (config): Critical fallthrough to" + " development_openfhe_config()", + file=sys.stderr, ) - return ( - development_openfhe_config() - ) # will raise a RuntimeError if repo_root not found - - -def from_pip_installation() -> OpenFHEConfig: - """ - Configure HEIR binaries from the expected pip installation structure. - """ - if not is_pip_installed(): - raise RuntimeError("HEIR is not installed via pip.") - - package_path = importlib.resources.files("heir") - return OpenFHEConfig( - include_dirs=[ - str(package_path / "openfhe+"), - str(package_path / "openfhe+" / "src" / "binfhe" / "include"), - str(package_path / "openfhe+" / "src" / "core" / "include"), - str(package_path / "openfhe+" / "src" / "pke" / "include"), - str(package_path / "cereal+" / "include"), - ], - include_type="source-relative", - lib_dir=str(package_path), - link_libs=["openfhe"], - ) + return development_openfhe_config() diff --git a/frontend/heir/backends/openfhe/config_test.py b/frontend/heir/backends/openfhe/config_test.py new file mode 100644 index 0000000000..9b7323d970 --- /dev/null +++ b/frontend/heir/backends/openfhe/config_test.py @@ -0,0 +1,141 @@ +import importlib.util +import os +import pathlib +import tempfile +import unittest +from unittest import mock + +from heir.backends.openfhe import config + + +class ConfigTest(unittest.TestCase): + + def test_from_os_env_only_debug(self): + with mock.patch.dict(os.environ, {"OPENFHE_DEBUG": "1"}, clear=True): + self.assertIsNone(config.from_os_env()) + + def test_from_os_env_valid(self): + env = { + "OPENFHE_INCLUDE_DIR": "/fake/include", + "OPENFHE_LIB_DIR": "/fake/lib", + } + with mock.patch.dict(os.environ, env, clear=True): + res = config.from_os_env() + self.assertIsNotNone(res) + self.assertEqual(res.include_dirs, ["/fake/include"]) + self.assertEqual(res.lib_dir, "/fake/lib") + + def test_from_os_env_appends_parent_of_openfhe_dir(self): + env = { + "OPENFHE_INCLUDE_DIR": "/fake/include/openfhe", + "OPENFHE_LIB_DIR": "/fake/lib", + } + with mock.patch.dict(os.environ, env, clear=True): + res = config.from_os_env() + self.assertIsNotNone(res) + self.assertEqual( + res.include_dirs, ["/fake/include/openfhe", "/fake/include"] + ) + + def test_from_os_env_missing_lib(self): + env = { + "OPENFHE_INCLUDE_DIR": "/fake/include", + } + with mock.patch.dict(os.environ, env, clear=True): + with self.assertRaises(ValueError): + config.from_os_env() + + +class ResolveLinkLibsTest(unittest.TestCase): + + def setUp(self): + super().setUp() + self.temp_dir = tempfile.TemporaryDirectory() + self.addCleanup(self.temp_dir.cleanup) + self.lib_path = pathlib.Path(self.temp_dir.name) + + def test_resolve_split_libs_lowercase(self): + core = self.lib_path / "libopenfhecore.so" + core.touch() + pke = self.lib_path / "libopenfhepke.so" + pke.touch() + binfhe = self.lib_path / "libopenfhebinfhe.so" + binfhe.touch() + + libs = config._resolve_link_libs(self.lib_path) + self.assertEqual( + libs, [str(core.resolve()), str(pke.resolve()), str(binfhe.resolve())] + ) + + def test_resolve_split_libs_uppercase(self): + core = self.lib_path / "libOPENFHEcore.so.1" + core.touch() + pke = self.lib_path / "libOPENFHEpke.so.1" + pke.touch() + binfhe = self.lib_path / "libOPENFHEbinfhe.so.1" + binfhe.touch() + + libs = config._resolve_link_libs(self.lib_path) + self.assertEqual( + libs, [str(core.resolve()), str(pke.resolve()), str(binfhe.resolve())] + ) + + def test_resolve_monolithic_lib(self): + mono = self.lib_path / "libopenfhe.so" + mono.touch() + + core = self.lib_path / "libopenfhecore.so" + core.touch() + + libs = config._resolve_link_libs(self.lib_path) + self.assertEqual(libs, [str(mono.resolve())]) + + def test_resolve_fallback(self): + libs = config._resolve_link_libs(self.lib_path) + self.assertEqual(libs, ["openfhe"]) + + +class ResolveConfigTest(unittest.TestCase): + + @mock.patch.object(config, "from_os_env") + @mock.patch.object(config, "get_default_installed_config") + @mock.patch.object(config, "development_openfhe_config") + def test_cascade_os_env(self, mock_dev, mock_default, mock_os_env): + mock_os_env.return_value = "os_env" + + res = config.resolve_config() + self.assertEqual(res, "os_env") + mock_os_env.assert_called_once() + mock_default.assert_not_called() + mock_dev.assert_not_called() + + @mock.patch.object(config, "from_os_env") + @mock.patch.object(config, "get_default_installed_config") + @mock.patch.object(config, "development_openfhe_config") + def test_cascade_default_installed(self, mock_dev, mock_default, mock_os_env): + mock_os_env.return_value = None + mock_default.return_value = "default_installed" + + res = config.resolve_config() + self.assertEqual(res, "default_installed") + mock_os_env.assert_called_once() + mock_default.assert_called_once() + mock_dev.assert_not_called() + + @mock.patch.object(config, "from_os_env") + @mock.patch.object(config, "get_default_installed_config") + @mock.patch.object(config, "development_openfhe_config") + def test_cascade_development(self, mock_dev, mock_default, mock_os_env): + mock_os_env.return_value = None + mock_default.return_value = None + mock_dev.return_value = "development" + + res = config.resolve_config() + self.assertEqual(res, "development") + mock_os_env.assert_called_once() + mock_default.assert_called_once() + mock_dev.assert_called_once() + + +if __name__ == "__main__": + unittest.main() diff --git a/frontend/heir/backends/util/common.py b/frontend/heir/backends/util/common.py index 2c41297e21..839e7e7bab 100644 --- a/frontend/heir/backends/util/common.py +++ b/frontend/heir/backends/util/common.py @@ -5,6 +5,7 @@ import os import pathlib import sysconfig +import sys from heir.interfaces import CompilationResult, EncValue @@ -66,15 +67,54 @@ def get_module_origin(module_name: str) -> Optional[str]: return spec.origin -def is_pip_installed() -> bool: +def is_pip_installed(debug: bool = False) -> bool: """Return true if heir is installed via pip.""" + debug_active = debug or os.environ.get("OPENFHE_DEBUG") == "1" try: - module_path = get_module_origin("heir") or "" - # purelib gives the environment-specific location of site-packages - # (for venv) or dist-packages (for system-wide installation) of - # non-builtin modules. - return module_path.startswith(sysconfig.get_paths()["purelib"]) - except ModuleNotFoundError: + origin = get_module_origin("heir") + if debug_active: + print( + "HEIRpy Debug (common): get_module_origin('heir') returned:" + f" {origin}", + file=sys.stderr, + ) + if not origin: + if debug_active: + print( + "HEIRpy Debug (common): is_pip_installed result: False (no origin)", + file=sys.stderr, + ) + return False + module_path = Path(origin) + purelib = Path(sysconfig.get_paths()["purelib"]) + platlib = Path(sysconfig.get_paths()["platlib"]) + + res = module_path.is_relative_to(purelib) or module_path.is_relative_to( + platlib + ) + + if debug_active: + print( + f"HEIRpy Debug (common): Module path: {module_path}", file=sys.stderr + ) + print(f"HEIRpy Debug (common): purelib: {purelib}", file=sys.stderr) + print(f"HEIRpy Debug (common): platlib: {platlib}", file=sys.stderr) + print( + f"HEIRpy Debug (common): is_pip_installed result: {res}", + file=sys.stderr, + ) + + return res + except (ModuleNotFoundError, ValueError) as e: + if debug_active: + print( + f"HEIRpy Debug (common): is_pip_installed caught exception: {e}", + file=sys.stderr, + ) + print( + "HEIRpy Debug (common): is_pip_installed result: False", + file=sys.stderr, + ) return False diff --git a/frontend/heir/backends/util/cpp_compiler.py b/frontend/heir/backends/util/cpp_compiler.py index 583ed0b577..73d573f8b3 100644 --- a/frontend/heir/backends/util/cpp_compiler.py +++ b/frontend/heir/backends/util/cpp_compiler.py @@ -16,7 +16,15 @@ def to_cpp_compiler_args(prefix, strs): - return [f"{prefix}{s}" for s in strs] if strs else [] + if not strs: + return [] + res = [] + for s in strs: + if prefix == "-l" and Path(s).is_absolute(): + res.append(s) + else: + res.append(f"{prefix}{s}") + return res class CppCompilerBackend: diff --git a/frontend/heir/heir_cli/heir_cli_config.py b/frontend/heir/heir_cli/heir_cli_config.py index baf17b245f..c368c5f314 100644 --- a/frontend/heir/heir_cli/heir_cli_config.py +++ b/frontend/heir/heir_cli/heir_cli_config.py @@ -23,7 +23,12 @@ class HEIRConfig: def development_heir_config() -> HEIRConfig: repo_root = get_repo_root() if not repo_root: - raise RuntimeError("Could not build development config. Did you run bazel?") + return HEIRConfig( + heir_opt_path="", + heir_translate_path="", + techmap_dir_path="", + abc_path="", + ) techmap_dir_path = ( repo_root @@ -47,6 +52,15 @@ def development_heir_config() -> HEIRConfig: / "abc" / "abc_bin" ) + if not abc_path.exists(): + abc_path = ( + repo_root + / "bazel-bin" + / "tools" + / "heir-opt.runfiles" + / "abc+" + / "abc_bin" + ) if not abc_path.exists(): abc_path = "" @@ -58,6 +72,40 @@ def development_heir_config() -> HEIRConfig: ) +def _resolve_binary( + env_var_name: str, + dev_path: Path | str, + binary_name: str, + required: bool = True, +) -> Path | str: + # 1. Env Var + val = os.environ.get(env_var_name) + if val: + return val + + # 2. Bazel-built dev path (if exists) + dev_path_obj = Path(dev_path) if dev_path else None + if dev_path_obj and dev_path_obj.exists(): + return dev_path_obj + + # 3. shutil.which PATH (if exists) + which_path = shutil.which(binary_name) + if which_path: + return which_path + + # 4. dev path fallback (if it is defined and we want to fall back to it without error) + if not required and dev_path: + return dev_path + + # 5. RuntimeError + if required: + raise RuntimeError( + f"Could not find binary '{binary_name}'. Please set {env_var_name} or" + " ensure it is on your PATH or built via Bazel." + ) + return "" + + def from_os_env() -> HEIRConfig: """Create a HEIRConfig from environment variables. @@ -68,33 +116,45 @@ def from_os_env() -> HEIRConfig: 1. Environment variables HEIR_OPT_PATH, HEIR_TRANSLATE_PATH, HEIR_ABC_BINARY, or HEIR_YOSYS_SCRIPTS_DIR - 2. The path to the heir-opt or heir-translate binary on the PATH - 3. The default development configuration (relative to the project root, in - bazel-bin) + 2. The default development configuration (relative to the project root, in + bazel-bin) (if they exist) + 3. The path to the heir-opt or heir-translate binary on the PATH + 4. The dev path fallback (if not found elsewhere) Returns: The HEIRConfig """ - which_heir_opt = shutil.which("heir-opt") - which_heir_translate = shutil.which("heir-translate") - which_abc = shutil.which("abc") - resolved_heir_opt_path = os.environ.get( - "HEIR_OPT_PATH", - which_heir_opt or development_heir_config().heir_opt_path, - ) - resolved_heir_translate_path = os.environ.get( - "HEIR_TRANSLATE_PATH", - which_heir_translate or development_heir_config().heir_translate_path, + dev_config = development_heir_config() + + resolved_heir_opt_path = _resolve_binary( + env_var_name="HEIR_OPT_PATH", + dev_path=dev_config.heir_opt_path, + binary_name="heir-opt", + required=True, ) - resolved_abc_path = os.environ.get( - "HEIR_ABC_BINARY", - which_abc or development_heir_config().abc_path, + resolved_heir_translate_path = _resolve_binary( + env_var_name="HEIR_TRANSLATE_PATH", + dev_path=dev_config.heir_translate_path, + binary_name="heir-translate", + required=True, ) - resolved_techmap_dir_path = os.environ.get( - "HEIR_YOSYS_SCRIPTS_DIR", - development_heir_config().techmap_dir_path, + resolved_abc_path = _resolve_binary( + env_var_name="HEIR_ABC_BINARY", + dev_path=dev_config.abc_path, + binary_name="abc", + required=False, ) + resolved_techmap_dir_path = os.environ.get("HEIR_YOSYS_SCRIPTS_DIR") + if not resolved_techmap_dir_path: + if ( + dev_config.techmap_dir_path + and Path(dev_config.techmap_dir_path).exists() + ): + resolved_techmap_dir_path = dev_config.techmap_dir_path + else: + resolved_techmap_dir_path = dev_config.techmap_dir_path + return HEIRConfig( heir_opt_path=resolved_heir_opt_path, heir_translate_path=resolved_heir_translate_path, diff --git a/frontend/heir/mlir/types.py b/frontend/heir/mlir/types.py index b0341a7752..bfb84a5368 100644 --- a/frontend/heir/mlir/types.py +++ b/frontend/heir/mlir/types.py @@ -1,14 +1,8 @@ """Defines Python type annotations for MLIR types.""" from abc import ABC, abstractmethod -import sys from typing import Generic, Optional, TypeVar, get_args, get_origin - -if sys.version_info >= (3, 11): - from typing import TypeVarTuple, Unpack -else: - # TypeVarTuple/Unpack landed in typing in 3.11; use the backport on 3.10. - from typing_extensions import TypeVarTuple, Unpack +from typing import TypeVarTuple, Unpack from numba.core.types import Type as NumbaType from numba.core.types import boolean, int8, int16, int32, int64, float32, float64 from numba.extending import type_callable diff --git a/frontend/heir/pipeline.py b/frontend/heir/pipeline.py index 6e5b22599d..a8ad93647b 100644 --- a/frontend/heir/pipeline.py +++ b/frontend/heir/pipeline.py @@ -328,10 +328,9 @@ def compile( _extras.require("pybind11", extra="openfhe") from heir.backends.openfhe import OpenFHEBackend, config as openfhe_config # pylint: disable=g-import-not-at-top - if is_pip_installed(): - backend = OpenFHEBackend(openfhe_config.from_pip_installation()) - else: - backend = OpenFHEBackend(openfhe_config.from_os_env()) + backend = OpenFHEBackend( + openfhe_config.resolve_config(debug=debug or False) + ) if debug and heir_opt_options is not None: DebugMessage(f"Overriding scheme with options {heir_opt_options}") diff --git a/frontend/import_isolation_test.py b/frontend/import_isolation_test.py index 978c98b136..e8ecabbc96 100644 --- a/frontend/import_isolation_test.py +++ b/frontend/import_isolation_test.py @@ -1,9 +1,8 @@ """Test that `import heir` does not import optional dependencies.""" -import importlib -import importlib.util +import os +import subprocess import sys - from absl.testing import absltest OPTIONAL_DEPS = ("numba", "numpy", "pybind11") @@ -12,26 +11,64 @@ class ImportIsolationTest(absltest.TestCase): def test_import_heir_does_not_load_optional_deps(self): - # The deps must be installed (else this test is vacuous) - # and not yet imported (else we couldn't attribute them to `import heir`). - for dep in OPTIONAL_DEPS: - self.assertIsNotNone( - importlib.util.find_spec(dep), - f"{dep} must be installed for this test to be meaningful", - ) - preloaded = [m for m in (*OPTIONAL_DEPS, "heir") if m in sys.modules] - self.assertEqual(preloaded, [], f"{preloaded} imported before the test ran") - - # This is normally `import heir` but this test is relative to the repository - # root. - importlib.import_module("frontend.heir") - - leaked = [dep for dep in OPTIONAL_DEPS if dep in sys.modules] + # We execute the test in a clean subprocess to ensure it runs in process isolation + # and is not polluted by other tests loaded by the test runner (like pytest). + + # Pass current sys.path via PYTHONPATH to the child process so it can find the modules. + env = os.environ.copy() + env["PYTHONPATH"] = os.pathsep.join(sys.path) + + # The code to execute in the subprocess + code = f""" +import importlib +import importlib.util +import sys + +OPTIONAL_DEPS = {OPTIONAL_DEPS} + +# Verify optional deps are installed +for dep in OPTIONAL_DEPS: + if importlib.util.find_spec(dep) is None: + print(f"Error: {{dep}} must be installed for this test to be meaningful", file=sys.stderr) + sys.exit(2) + +# Verify they are not preloaded in this clean process +preloaded = [m for m in (*OPTIONAL_DEPS, "heir") if m in sys.modules] +if preloaded: + print(f"Error: {{preloaded}} preloaded before import", file=sys.stderr) + sys.exit(3) + +# Import heir +importlib.import_module("frontend.heir") + +# Verify they were not loaded as side effects +leaked = [dep for dep in OPTIONAL_DEPS if dep in sys.modules] +if leaked: + print(f"Error: `import heir` imported {{leaked}}", file=sys.stderr) + sys.exit(4) + +print("SUCCESS") +sys.exit(0) +""" + + result = subprocess.run( + [sys.executable, "-c", code], + capture_output=True, + text=True, + env=env, + ) + + if result.returncode != 0: + print("Subprocess stdout:") + print(result.stdout) + print("Subprocess stderr:") + print(result.stderr) + self.assertEqual( - leaked, - [], - f"Error: `import heir` imported {leaked}; but the base package should" - " not import any of the optional dependencies", + result.returncode, + 0, + f"Import isolation check failed with exit code {result.returncode}. See" + " stdout/stderr for details.", ) diff --git a/frontend/sbox_test.py b/frontend/sbox_test.py deleted file mode 100644 index eef953820e..0000000000 --- a/frontend/sbox_test.py +++ /dev/null @@ -1,156 +0,0 @@ -"""Example of HEIR Python usage.""" - -from heir import compile -from heir.mlir import F32, I1, I8, I16, I64, Secret, Tensor - - -# This example does not run end to end, but is used to generate MLIR for the -# CGGI pipeline -@compile(debug=True, scheme="cggi") -def func(x: Secret[I8]): - x0: I1 = (x >> 7) & 1 - x1 = (x >> 6) & 1 - x2 = (x >> 5) & 1 - x3 = (x >> 4) & 1 - x4 = (x >> 3) & 1 - x5 = (x >> 2) & 1 - x6 = (x >> 1) & 1 - x7 = x & 1 - - y14 = x3 ^ x5 - y13 = x0 ^ x6 - y9 = x0 ^ x3 - y8 = x0 ^ x5 - t0 = x1 ^ x2 - y1 = t0 ^ x7 - y4 = y1 ^ x3 - y12 = y13 ^ y14 - y2 = y1 ^ x0 - y5 = y1 ^ x6 - y3 = y5 ^ y8 - t1 = x4 ^ y12 - y15 = t1 ^ x5 - y20 = t1 ^ x1 - y6 = y15 ^ x7 - y10 = y15 ^ t0 - y11 = y20 ^ y9 - y7 = x7 ^ y11 - y17 = y10 ^ y11 - y19 = y10 ^ y8 - y16 = t0 ^ y11 - y21 = y13 ^ y16 - y18 = x0 ^ y16 - # nonlinear - t2 = y12 & y15 - t3 = y3 & y6 - t4 = t3 ^ t2 - t5 = y4 & x7 - t6 = t5 ^ t2 - t7 = y13 & y16 - t8 = y5 & y1 - t9 = t8 ^ t7 - t10 = y2 & y7 - t11 = t10 ^ t7 - t12 = y9 & y11 - t13 = y14 & y17 - t14 = t13 ^ t12 - t15 = y8 & y10 - t16 = t15 ^ t12 - t17 = t4 ^ t14 - t18 = t6 ^ t16 - t19 = t9 ^ t14 - t20 = t11 ^ t16 - t21 = t17 ^ y20 - t22 = t18 ^ y19 - t23 = t19 ^ y21 - t24 = t20 ^ y18 - - t25 = t21 ^ t22 - t26 = t21 & t23 - t27 = t24 ^ t26 - t28 = t25 & t27 - t29 = t28 ^ t22 - t30 = t23 ^ t24 - t31 = t22 ^ t26 - t32 = t31 & t30 - t33 = t32 ^ t24 - t34 = t23 ^ t33 - t35 = t27 ^ t33 - t36 = t24 & t35 - t37 = t36 ^ t34 - t38 = t27 ^ t36 - t39 = t29 & t38 - t40 = t25 ^ t39 - - t41 = t40 ^ t37 - t42 = t29 ^ t33 - t43 = t29 ^ t40 - t44 = t33 ^ t37 - t45 = t42 ^ t41 - z0 = t44 & y15 - z1 = t37 & y6 - z2 = t33 & x7 - z3 = t43 & y16 - z4 = t40 & y1 - z5 = t29 & y7 - z6 = t42 & y11 - z7 = t45 & y17 - z8 = t41 & y10 - z9 = t44 & y12 - z10 = t37 & y3 - z11 = t33 & y4 - z12 = t43 & y13 - z13 = t40 & y5 - z14 = t29 & y2 - z15 = t42 & y9 - z16 = t45 & y14 - z17 = t41 & y8 - - t46 = z15 ^ z16 - t47 = z10 ^ z11 - t48 = z5 ^ z13 - t49 = z9 ^ z10 - t50 = z2 ^ z12 - t51 = z2 ^ z5 - t52 = z7 ^ z8 - t53 = z0 ^ z3 - t54 = z6 ^ z7 - t55 = z16 ^ z17 - t56 = z12 ^ t48 - t57 = t50 ^ t53 - t58 = z4 ^ t46 - t59 = z3 ^ t54 - t60 = t46 ^ t57 - t61 = z14 ^ t57 - t62 = t52 ^ t58 - t63 = t49 ^ t58 - t64 = z4 ^ t59 - t65 = t61 ^ t62 - t66 = z1 ^ t63 - s0 = t59 ^ t63 - s6 = t56 ^ ~t62 - s7 = t48 ^ ~t60 - t67 = t64 ^ t65 - s3 = t53 ^ t66 - s4 = t51 ^ t66 - s5 = t47 ^ t65 - s1 = t64 ^ ~s3 - s2 = t55 ^ ~t67 - - q = s0 - q = q << 1 - q = q + s1 - q = q << 1 - q = q + s2 - q = q << 1 - q = q + s3 - q = q << 1 - q = q + s4 - q = q << 1 - q = q + s5 - q = q << 1 - q = q + s6 - q = q << 1 - q = q + s7 - - return q diff --git a/frontend/testing.bzl b/frontend/testing.bzl deleted file mode 100644 index 048d03caea..0000000000 --- a/frontend/testing.bzl +++ /dev/null @@ -1,49 +0,0 @@ -"""Macros for py_test rules that use the python frontend.""" - -load("@rules_python//python:defs.bzl", "py_test") - -def frontend_test(name, srcs, deps = [], data = [], tags = []): - """A py_test replacement with an env including all backend dependencies. - """ - include_dirs = [ - "openfhe+/", - "openfhe+/src/binfhe/include", - "openfhe+/src/core/include", - "openfhe+/src/pke/include", - "cereal+/include", - "rapidjson+/include", - ] - - libs = [ - "openfhe", - ] - - py_test( - name = name, - srcs = srcs, - python_version = "PY3", - srcs_version = "PY3", - deps = deps + [ - ":frontend", - "@abseil-py//absl/testing:absltest", - "@abc//:abc_bin", - ], - imports = ["."], - data = data, - shard_count = 3, - tags = tags, - env = { - # this dir is relative to $RUNFILES_DIR, which is set by bazel at runtime - "OPENFHE_LIB_DIR": "openfhe+", - "OPENFHE_INCLUDE_TYPE": "source-relative", - "OPENFHE_LINK_LIBS": ":".join(libs), - "OPENFHE_INCLUDE_DIR": ":".join(include_dirs), - "HEIR_REPO_ROOT_MARKER": ".", - "HEIR_OPT_PATH": "tools/heir-opt", - "HEIR_TRANSLATE_PATH": "tools/heir-translate", - "PYBIND11_INCLUDE_PATH": "pybind11/include", - "HEIR_YOSYS_SCRIPTS_DIR": "lib/Transforms/YosysOptimizer/yosys", - "HEIR_ABC_BINARY": "$(rootpath @abc//:abc_bin)", - "NUMBA_USE_LEGACY_TYPE_SYSTEM": "1", - }, - ) diff --git a/lib/Target/OpenFhePke/BUILD b/lib/Target/OpenFhePke/BUILD index e47284a8ea..385f1301ee 100644 --- a/lib/Target/OpenFhePke/BUILD +++ b/lib/Target/OpenFhePke/BUILD @@ -195,6 +195,10 @@ cc_test( srcs = ["InterpreterTest.cpp"], copts = OPENMP_COPTS, linkopts = OPENMP_LINKOPTS, + tags = [ + "manual", + "notap", + ], deps = [ ":Interpreter", "@googletest//:gtest_main", diff --git a/pyproject.toml b/pyproject.toml index e7697e4572..b0dccbd29e 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -12,13 +12,14 @@ authors = [ ] description = "The HEIR compiler" readme = "README.md" -requires-python = ">=3.10" +requires-python = ">=3.11" classifiers = [ "Topic :: Security :: Cryptography", "Intended Audience :: Developers", "Topic :: Software Development :: Compilers", ] license = "MIT AND Apache-2.0" +license-files = ["LICENSE"] # The base package only ships the heir-opt/heir-translate binaries and has no deps dependencies = [] @@ -33,10 +34,13 @@ python = [ # in some of our tests as Numpy 2 introduces stricter type checking "numpy>=1.23.5,<2", "colorama>=0.4.6", - "typing_extensions>=4.6; python_version < '3.11'", ] # openfhe backend -openfhe = ["pybind11>=2.13.6", "pybind11_global>=2.13.6", "colorama>=0.4.6"] +openfhe = [ + "pybind11>=2.13.6", + "pybind11_global>=2.13.6", + "colorama>=0.4.6", +] # dev dependencies dev = [ @@ -58,7 +62,8 @@ dev = [ # For ML integration tests "torch>=2.8.0", "setuptools>=80.9.0", - "cibuildwheel>=3.0.0", + "cibuildwheel>=4.0.0", + "pytest", ] # For `uv sync` (which includes `dev` by default) or / `pip install --group dev`. @@ -89,21 +94,21 @@ skip = "*-musllinux_* pp-*" # when running `import heir` instead of the local source files. cibuildwheel # creates an isolated directory for this purpose, and cd's to it before running # this command. -test-command = "cp {package}/frontend/example.py . && python example.py" -test-extras = ["python", "openfhe"] +test-command = "cp {package}/frontend/cibuildwheel.py . && python cibuildwheel.py" +test-extras = ["python"] [tool.cibuildwheel.linux] archs = "auto64" -manylinux-x86_64-image = "manylinux_2_28" -manylinux-aarch64-image = "manylinux_2_28" +manylinux-x86_64-image = "manylinux_2_34" +manylinux-aarch64-image = "manylinux_2_34" before-all = "bash .github/install_cibuildwheel_deps.sh" # Use the rootless Bazel installation inside the container. environment = { PATH = "$PATH:$HOME/bin" } -# Force the repaired wheel to use the manylinux_2_28 platform tag. +# Force the repaired wheel to use the manylinux_2_34 platform tag. # By default, auditwheel may analyze the wheel's actual symbol requirements -# and choose or add a lower compatible manylinux_x_y tag such as 2_27 +# and choose or add a lower compatible manylinux_x_y tag such as 2_28 # which PyPi might reject (they only accept a few specific tags). -repair-wheel-command = "auditwheel repair --plat manylinux_2_28_$(uname -m) --only-plat -w {dest_dir} {wheel}" +repair-wheel-command = "auditwheel repair --plat manylinux_2_34_$(uname -m) --only-plat -w {dest_dir} {wheel}" [tool.cibuildwheel.windows] archs = "auto64" diff --git a/setup.py b/setup.py index dac20926e0..8dd3764afa 100644 --- a/setup.py +++ b/setup.py @@ -100,7 +100,6 @@ def __init__( generated_so_file: Path, target_file: str, is_binary: bool = False, - copy_include_files: bool = False, aggressive_strip: bool = False, **kwargs: Any, ): @@ -114,7 +113,6 @@ def __init__( stripped_target = bazel_target.split("//")[-1] self.relpath, self.target_name = stripped_target.split(":") self.is_binary = is_binary - self.copy_include_files = copy_include_files # Determines whether to `strip-all` when building a target for PyPI. # Critically, this should only be used for cc_binary targets that are run as @@ -133,6 +131,7 @@ def run(self): # place and use the development config to find the binaries, etc. return + self.copy_yosys_techmaps() for ext in self.extensions: self.bazel_build(ext) # explicitly call `bazel shutdown` for graceful exit @@ -145,80 +144,27 @@ def copy_extensions_to_source(self): do again in the `build_ext` base class. """ - def copy_include_files(self): - """Copy include files from bazel external directory to the libdir.""" - # Find the bazel directory that contains the source files for external - # dependencies. We copy over the include files for C++ dependencies needed - # during JIT compilation. It's a symlink and shutil.copytree doesn't work - # through symlinks, so we have to resolve() it first. - - # Normally there would be an external/ dir, but for some reason the - # containerized execution doesn't create that symlink. - - # bazel-heir -> /home/user/.cache/bazel/_bazel_j2kun//execroot/_main - bazel_out = (Path(self.build_temp) / "bazel-out").resolve() - bazel_heir = bazel_out.parent - out_root = bazel_out.parent.parent.parent - external = out_root / "external" - libdir = Path(self.build_lib) / "heir" - subdirs = [ - Path("openfhe+") / "src" / "binfhe" / "include", - Path("openfhe+") / "src" / "core" / "include", - Path("openfhe+") / "src" / "pke" / "include", - Path("cereal+") / "include", - ] - dirs_to_copy = {external / subdir: libdir / subdir for subdir in subdirs} - - # Include yosys techmaps. - dirs_to_copy[ - bazel_heir / "lib" / "Transforms" / "YosysOptimizer" / "yosys" - ] = (libdir / "techmaps") - + def copy_yosys_techmaps(self): + """Copy Yosys techmap files from source tree to the libdir.""" + src = Path("lib/Transforms/YosysOptimizer/yosys") + dst = Path(self.build_lib) / "heir" / "techmaps" patterns = ( "*.v", # techmap files - "*.h", - "*.hpp", "LICENSE", "README.md", ) - - for src, dst in dirs_to_copy.items(): - print(f"Copying {src} to {dst}") - shutil.copytree( - src=src, - dst=dst, - ignore=include_patterns(patterns), - dirs_exist_ok=True, - ) - - # Some headers (e.g. config_core.h) are generated by Bazel genrules and - # live in bazel-bin rather than the source tree. Merge those in, but skip - # files already present (the main copytree may have picked them up via - # execroot symlinks, and those copies are read-only). - def copy_if_missing(src, dst): - if not os.path.exists(dst): - shutil.copy2(src, dst) - - bazel_bin_external = Path(self.build_temp) / "bazel-bin" / "external" - for subdir in subdirs: - src = bazel_bin_external / subdir - if src.exists(): - print(f"Copying generated headers from {src} to {libdir / subdir}") - shutil.copytree( - src=src, - dst=libdir / subdir, - ignore=include_patterns(patterns), - dirs_exist_ok=True, - copy_function=copy_if_missing, - ) + print(f"Copying {src} to {dst}") + shutil.copytree( + src=src, + dst=dst, + ignore=include_patterns(patterns), + dirs_exist_ok=True, + ) def bazel_build(self, ext: BazelExtension) -> None: # noqa: C901 """Runs the bazel build to create the package.""" temp_path = Path(self.build_temp) - if ext.copy_include_files: - self.copy_include_files() - # We round to the minor version, which makes rules_python # look up the latest available patch version internally. python_version = "{}.{}".format(*sys.version_info[:2]) @@ -243,11 +189,6 @@ def bazel_build(self, ext: BazelExtension) -> None: # noqa: C901 elif IS_MAC: bazel_argv.append("--stripopt=-S") - if is_cibuildwheel() and IS_LINUX: - # TODO(#2249): OpenMP is disabled for manylinux, as libomp is not on the allowed libraries list, - # and statically bundling OpenMP is non-trivial and left as future work. - bazel_argv.append("--@openfhe//:enable_openmp=False") - if IS_WINDOWS: # Link with python*.lib. for library_dir in self.library_dirs: @@ -323,10 +264,16 @@ def bazel_build(self, ext: BazelExtension) -> None: # noqa: C901 cmdclass={ "build_ext": BuildBazelExtension, }, - package_data={"heir_py": ["py.typed", "*.pyi"]}, + package_data={ + "heir": [ + "py.typed", + "*.pyi", + "techmaps/*", + ] + }, ext_modules=[ BazelExtension( - name="heir_py._heir_opt", + name="heir._heir_opt", bazel_target="//tools:heir-opt.stripped", generated_so_file=Path("tools") / "heir-opt.stripped", target_file="heir-opt", @@ -335,7 +282,7 @@ def bazel_build(self, ext: BazelExtension) -> None: # noqa: C901 aggressive_strip=is_cibuildwheel(), ), BazelExtension( - name="heir_py._heir_translate", + name="heir._heir_translate", bazel_target="//tools:heir-translate.stripped", generated_so_file=Path("tools") / "heir-translate.stripped", target_file="heir-translate", @@ -344,7 +291,7 @@ def bazel_build(self, ext: BazelExtension) -> None: # noqa: C901 aggressive_strip=is_cibuildwheel(), ), BazelExtension( - name="heir_py._abc", + name="heir._abc", bazel_target="@abc//:abc_bin.stripped", generated_so_file=Path("external") / "abc+" / "abc_bin.stripped", target_file="abc_bin", @@ -352,15 +299,6 @@ def bazel_build(self, ext: BazelExtension) -> None: # noqa: C901 is_binary=True, aggressive_strip=is_cibuildwheel(), ), - BazelExtension( - name="heir_py._libopenfhe", - bazel_target="@openfhe//:libopenfhe", - generated_so_file=Path("external") / "openfhe+" / "libopenfhe.so", - target_file="libopenfhe.so", - py_limited_api=py_limited_api, - copy_include_files=True, - aggressive_strip=False, - ), ], options=options, ) diff --git a/tests/Examples/openfhe/test.bzl b/tests/Examples/openfhe/test.bzl index ba03f1213c..c373172db3 100644 --- a/tests/Examples/openfhe/test.bzl +++ b/tests/Examples/openfhe/test.bzl @@ -82,7 +82,7 @@ def openfhe_interpreter_test(name, mlir_src, test_src, generated_heir_opt_filena "@llvm-project//mlir:Support", ], timeout = timeout, - tags = tags, + tags = tags + ["manual", "notap"], data = data + [":" + generated_heir_opt_filename], copts = OPENMP_COPTS + copts, linkopts = OPENMP_LINKOPTS,