From 79a0d535b113dd0f007da842e7566e06342b7807 Mon Sep 17 00:00:00 2001 From: Luciano Paz Date: Wed, 30 Oct 2024 22:44:09 +0100 Subject: [PATCH 1/2] Add Accelerate framework blas__ldflags tests --- pytensor/link/c/cmodule.py | 38 ++++++++++++++++-- pytensor/tensor/blas.py | 34 ++++++++++++++-- tests/link/c/test_cmodule.py | 75 ++++++++++++++++++++++++++++++------ 3 files changed, 127 insertions(+), 20 deletions(-) diff --git a/pytensor/link/c/cmodule.py b/pytensor/link/c/cmodule.py index 70be23a0a5..62f5adea01 100644 --- a/pytensor/link/c/cmodule.py +++ b/pytensor/link/c/cmodule.py @@ -2458,7 +2458,23 @@ def patch_ldflags(flag_list: list[str]) -> list[str]: @staticmethod def linking_patch(lib_dirs: list[str], libs: list[str]) -> list[str]: if sys.platform != "win32": - return [f"-l{l}" for l in libs] + patched_libs = [] + framework = False + for lib in libs: + # The clang framework flag is handled differently. + # The flag will have the format -framework framework_name + # If we find a lib that is called -framework, we keep it and the following + # entry in the lib list unchanged. Anything else, we add the standard + # -l library prefix. + if lib == "-framework": + framework = True + patched_libs.append(lib) + elif framework: + framework = False + patched_libs.append(lib) + else: + patched_libs.append(f"-l{lib}") + return patched_libs else: # In explicit else because of https://github.com/python/mypy/issues/10773 def sort_key(lib): @@ -2466,6 +2482,8 @@ def sort_key(lib): return (extension == "dll", tuple(map(int, numbers))) patched_lib_ldflags = [] + # Should we also add a framework possibility on windows? I didn't do so because + # clang is not intended to be used there at the moment. for lib in libs: ldflag = f"-l{lib}" for lib_dir in lib_dirs: @@ -2873,9 +2891,21 @@ def check_libs( ) except Exception as e: _logger.debug(e) + try: + # 3. Mac Accelerate framework + _logger.debug("Checking Accelerate framework") + flags = ["-framework", "Accelerate"] + if rpath: + flags = [*flags, f"-Wl,-rpath,{rpath}"] + validated_flags = try_blas_flag(flags) + if validated_flags == "": + raise Exception("Accelerate framework flag failed ") + return validated_flags + except Exception as e: + _logger.debug(e) try: _logger.debug("Checking Lapack + blas") - # 3. Try to use LAPACK + BLAS + # 4. Try to use LAPACK + BLAS return check_libs( all_libs, required_libs=["lapack", "blas", "cblas", "m"], @@ -2885,7 +2915,7 @@ def check_libs( except Exception as e: _logger.debug(e) try: - # 4. Try to use BLAS alone + # 5. Try to use BLAS alone _logger.debug("Checking blas alone") return check_libs( all_libs, @@ -2896,7 +2926,7 @@ def check_libs( except Exception as e: _logger.debug(e) try: - # 5. Try to use openblas + # 6. Try to use openblas _logger.debug("Checking openblas") return check_libs( all_libs, diff --git a/pytensor/tensor/blas.py b/pytensor/tensor/blas.py index b3cf96cbd4..6170a02a98 100644 --- a/pytensor/tensor/blas.py +++ b/pytensor/tensor/blas.py @@ -78,7 +78,9 @@ import functools import logging import os +import shlex import time +from pathlib import Path import numpy as np @@ -396,7 +398,7 @@ def _ldflags( rval = [] if libs_dir: found_dyn = False - dirs = [x[2:] for x in ldflags_str.split() if x.startswith("-L")] + dirs = [x[2:] for x in shlex.split(ldflags_str) if x.startswith("-L")] l = _ldflags( ldflags_str=ldflags_str, libs=True, @@ -409,6 +411,9 @@ def _ldflags( if f.endswith(".so") or f.endswith(".dylib") or f.endswith(".dll"): if any(f.find(ll) >= 0 for ll in l): found_dyn = True + # Special treatment of clang framework. Specifically for MacOS Accelerate + if "-framework" in l and "Accelerate" in l: + found_dyn = True if not found_dyn and dirs: _logger.warning( "We did not find a dynamic library in the " @@ -416,7 +421,12 @@ def _ldflags( "ATLAS, make sure to compile it with dynamics library." ) - for t in ldflags_str.split(): + split_flags = shlex.split(ldflags_str) + skip = False + for pos, t in enumerate(split_flags): + if skip: + skip = False + continue # Remove extra quote. if (t.startswith("'") and t.endswith("'")) or ( t.startswith('"') and t.endswith('"') @@ -425,10 +435,26 @@ def _ldflags( try: t0, t1 = t[0], t[1] - assert t0 == "-" + assert t0 == "-" or Path(t).exists() except Exception: raise ValueError(f'invalid token "{t}" in ldflags_str: "{ldflags_str}"') - if libs_dir and t1 == "L": + if t == "-framework": + skip = True + # Special treatment of clang framework. Specifically for MacOS Accelerate + # The clang framework implicitly adds: header dirs, libraries, and library dirs. + # If we choose to always return these flags, we run into a huge deal amount of + # incompatibilities. For this reason, we only return the framework if libs are + # requested. + if ( + libs + and len(split_flags) >= pos + and split_flags[pos + 1] == "Accelerate" + ): + # We only add the Accelerate framework, but in the future we could extend it to + # other frameworks + rval.append(t) + rval.append(split_flags[pos + 1]) + elif libs_dir and t1 == "L": rval.append(t[2:]) elif include_dir and t1 == "I": raise ValueError( diff --git a/tests/link/c/test_cmodule.py b/tests/link/c/test_cmodule.py index 0eae1db68e..2242bc12e9 100644 --- a/tests/link/c/test_cmodule.py +++ b/tests/link/c/test_cmodule.py @@ -165,13 +165,22 @@ def test_flag_detection(): @pytest.fixture( scope="module", - params=["mkl_intel", "mkl_gnu", "openblas", "lapack", "blas", "no_blas"], + params=[ + "mkl_intel", + "mkl_gnu", + "accelerate", + "openblas", + "lapack", + "blas", + "no_blas", + ], ) def blas_libs(request): key = request.param libs = { "mkl_intel": ["mkl_core", "mkl_rt", "mkl_intel_thread", "iomp5", "pthread"], "mkl_gnu": ["mkl_core", "mkl_rt", "mkl_gnu_thread", "gomp", "pthread"], + "accelerate": ["vecLib_placeholder"], "openblas": ["openblas", "gfortran", "gomp", "m"], "lapack": ["lapack", "blas", "cblas", "m"], "blas": ["blas", "cblas"], @@ -190,25 +199,37 @@ def mock_system(request): def cxx_search_dirs(blas_libs, mock_system): libext = {"Linux": "so", "Windows": "dll", "Darwin": "dylib"} libraries = [] + enabled_accelerate_framework = False with tempfile.TemporaryDirectory() as d: flags = None for lib in blas_libs: - lib_path = Path(d) / f"{lib}.{libext[mock_system]}" - lib_path.write_bytes(b"1") - libraries.append(lib_path) - if flags is None: - flags = f"-l{lib}" + if lib == "vecLib_placeholder": + if mock_system != "Darwin": + flags = "" + else: + flags = "-framework Accelerate" + enabled_accelerate_framework = True else: - flags += f" -l{lib}" + lib_path = Path(d) / f"{lib}.{libext[mock_system]}" + lib_path.write_bytes(b"1") + libraries.append(lib_path) + if flags is None: + flags = f"-l{lib}" + else: + flags += f" -l{lib}" if "gomp" in blas_libs and "mkl_gnu_thread" not in blas_libs: flags += " -fopenmp" if len(blas_libs) == 0: flags = "" - yield f"libraries: ={d}".encode(sys.stdout.encoding), flags + yield ( + f"libraries: ={d}".encode(sys.stdout.encoding), + flags, + enabled_accelerate_framework, + ) @pytest.fixture( - scope="function", params=[False, True], ids=["Working_CXX", "Broken_CXX"] + scope="function", params=[True, False], ids=["Working_CXX", "Broken_CXX"] ) def cxx_search_dirs_status(request): return request.param @@ -219,22 +240,39 @@ def cxx_search_dirs_status(request): def test_default_blas_ldflags( mock_std_lib_dirs, mock_check_mkl_openmp, cxx_search_dirs, cxx_search_dirs_status ): - cxx_search_dirs, expected_blas_ldflags = cxx_search_dirs + cxx_search_dirs, expected_blas_ldflags, enabled_accelerate_framework = ( + cxx_search_dirs + ) mock_process = MagicMock() if cxx_search_dirs_status: error_message = "" mock_process.communicate = lambda *args, **kwargs: (cxx_search_dirs, b"") mock_process.returncode = 0 else: + enabled_accelerate_framework = False error_message = "Unsupported argument -print-search-dirs" error_message_bytes = error_message.encode(sys.stderr.encoding) mock_process.communicate = lambda *args, **kwargs: (b"", error_message_bytes) mock_process.returncode = 1 + + def patched_compile_tmp(*args, **kwargs): + def wrapped(test_code, tmp_prefix, flags, try_run, output): + if len(flags) >= 2 and flags[:2] == ["-framework", "Accelerate"]: + print(enabled_accelerate_framework) + if enabled_accelerate_framework: + return (True, True) + else: + return (False, False, "", "Invalid flags -framework Accelerate") + else: + return (True, True) + + return wrapped + with patch("pytensor.link.c.cmodule.subprocess_Popen", return_value=mock_process): with patch.object( pytensor.link.c.cmodule.GCC_compiler, "try_compile_tmp", - return_value=(True, True), + new_callable=patched_compile_tmp, ): if cxx_search_dirs_status: assert set(default_blas_ldflags().split(" ")) == set( @@ -267,6 +305,9 @@ def windows_conda_libs(blas_libs): subdir.mkdir(exist_ok=True, parents=True) flags = f'-L"{subdir}"' for lib in blas_libs: + if lib == "vecLib_placeholder": + flags = "" + break lib_path = subdir / f"{lib}.dll" lib_path.write_bytes(b"1") libraries.append(lib_path) @@ -287,6 +328,16 @@ def test_default_blas_ldflags_conda_windows( mock_process = MagicMock() mock_process.communicate = lambda *args, **kwargs: (b"", b"") mock_process.returncode = 0 + + def patched_compile_tmp(*args, **kwargs): + def wrapped(test_code, tmp_prefix, flags, try_run, output): + if len(flags) >= 2 and flags[:2] == ["-framework", "Accelerate"]: + return (False, False, "", "Invalid flags -framework Accelerate") + else: + return (True, True) + + return wrapped + with patch("sys.platform", "win32"): with patch("sys.prefix", mock_sys_prefix): with patch( @@ -295,7 +346,7 @@ def test_default_blas_ldflags_conda_windows( with patch.object( pytensor.link.c.cmodule.GCC_compiler, "try_compile_tmp", - return_value=(True, True), + new_callable=patched_compile_tmp, ): assert set(default_blas_ldflags().split(" ")) == set( expected_blas_ldflags.split(" ") From 8307af15ffc437009a44b2991506aff642087b40 Mon Sep 17 00:00:00 2001 From: Luciano Paz Date: Fri, 1 Nov 2024 07:37:56 +0100 Subject: [PATCH 2/2] Add workflow job on mac to check blas --- .github/workflows/test.yml | 40 ++++++++++++++++++++++++++++++-------- 1 file changed, 32 insertions(+), 8 deletions(-) diff --git a/.github/workflows/test.yml b/.github/workflows/test.yml index 7298d5df61..a8456c8292 100644 --- a/.github/workflows/test.yml +++ b/.github/workflows/test.yml @@ -61,16 +61,17 @@ jobs: python-version: ${{ matrix.python-version }} - uses: pre-commit/action@v3.0.1 - test_ubuntu: - name: "Test py${{ matrix.python-version }} : fast-compile ${{ matrix.fast-compile }} : float32 ${{ matrix.float32 }} : ${{ matrix.part }}" + test: + name: "${{ matrix.os }} test py${{ matrix.python-version }} : fast-compile ${{ matrix.fast-compile }} : float32 ${{ matrix.float32 }} : ${{ matrix.part }}" needs: - changes - style - runs-on: ubuntu-latest + runs-on: ${{ matrix.os }} if: ${{ needs.changes.outputs.changes == 'true' && needs.style.result == 'success' }} strategy: fail-fast: false matrix: + os: ["ubuntu-latest"] python-version: ["3.10", "3.12"] fast-compile: [0, 1] float32: [0, 1] @@ -103,30 +104,44 @@ jobs: fast-compile: 1 include: - install-numba: 1 + os: "ubuntu-latest" python-version: "3.10" fast-compile: 0 float32: 0 part: "tests/link/numba" - install-numba: 1 + os: "ubuntu-latest" python-version: "3.12" fast-compile: 0 float32: 0 part: "tests/link/numba" - install-jax: 1 + os: "ubuntu-latest" python-version: "3.10" fast-compile: 0 float32: 0 part: "tests/link/jax" - install-jax: 1 + os: "ubuntu-latest" python-version: "3.12" fast-compile: 0 float32: 0 part: "tests/link/jax" - install-torch: 1 + os: "ubuntu-latest" python-version: "3.10" fast-compile: 0 float32: 0 part: "tests/link/pytorch" + - os: macos-latest + python-version: "3.12" + fast-compile: 0 + float32: 0 + install-numba: 0 + install-jax: 0 + install-torch: 0 + part: "tests/tensor/test_blas.py tests/tensor/test_elemwise.py tests/tensor/test_math_scipy.py" + steps: - uses: actions/checkout@v4 with: @@ -146,7 +161,7 @@ jobs: MATRIX_CONTEXT: ${{ toJson(matrix) }} run: | echo $MATRIX_CONTEXT - export MATRIX_ID=`echo $MATRIX_CONTEXT | md5sum | cut -c 1-32` + export MATRIX_ID=`echo $MATRIX_CONTEXT | sha256sum | cut -c 1-32` echo $MATRIX_ID echo "id=$MATRIX_ID" >> $GITHUB_OUTPUT @@ -154,7 +169,11 @@ jobs: shell: micromamba-shell {0} run: | - micromamba install --yes -q "python~=${PYTHON_VERSION}=*_cpython" mkl numpy scipy pip mkl-service graphviz cython pytest coverage pytest-cov pytest-benchmark pytest-mock + if [[ $OS == "macos-latest" ]]; then + micromamba install --yes -q "python~=${PYTHON_VERSION}=*_cpython" numpy scipy pip graphviz cython pytest coverage pytest-cov pytest-benchmark pytest-mock libblas=*=*accelerate; + else + micromamba install --yes -q "python~=${PYTHON_VERSION}=*_cpython" mkl numpy scipy pip mkl-service graphviz cython pytest coverage pytest-cov pytest-benchmark pytest-mock; + fi if [[ $INSTALL_NUMBA == "1" ]]; then micromamba install --yes -q -c conda-forge "python~=${PYTHON_VERSION}=*_cpython" "numba>=0.57"; fi if [[ $INSTALL_JAX == "1" ]]; then micromamba install --yes -q -c conda-forge "python~=${PYTHON_VERSION}=*_cpython" jax jaxlib numpyro && pip install tensorflow-probability; fi if [[ $INSTALL_TORCH == "1" ]]; then micromamba install --yes -q -c conda-forge "python~=${PYTHON_VERSION}=*_cpython" pytorch pytorch-cuda=12.1 "mkl<=2024.0" -c pytorch -c nvidia; fi @@ -163,12 +182,17 @@ jobs: pip install -e ./ micromamba list && pip freeze python -c 'import pytensor; print(pytensor.config.__str__(print_doc=False))' - python -c 'import pytensor; assert pytensor.config.blas__ldflags != "", "Blas flags are empty"' + if [[ $OS == "macos-latest" ]]; then + python -c 'import pytensor; assert pytensor.config.blas__ldflags.startswith("-framework Accelerate"), "Blas flags are not set to MacOS Accelerate"'; + else + python -c 'import pytensor; assert pytensor.config.blas__ldflags != "", "Blas flags are empty"'; + fi env: PYTHON_VERSION: ${{ matrix.python-version }} INSTALL_NUMBA: ${{ matrix.install-numba }} INSTALL_JAX: ${{ matrix.install-jax }} INSTALL_TORCH: ${{ matrix.install-torch}} + OS: ${{ matrix.os}} - name: Run tests shell: micromamba-shell {0} @@ -249,10 +273,10 @@ jobs: if: ${{ always() }} runs-on: ubuntu-latest name: "All tests" - needs: [changes, style, test_ubuntu] + needs: [changes, style, test] steps: - name: Check build matrix status - if: ${{ needs.changes.outputs.changes == 'true' && (needs.style.result != 'success' || needs.test_ubuntu.result != 'success') }} + if: ${{ needs.changes.outputs.changes == 'true' && (needs.style.result != 'success' || needs.test.result != 'success') }} run: exit 1 upload-coverage: