From 2857bea216c05d628a8cb89044ddf440fd1e94df Mon Sep 17 00:00:00 2001 From: lucianopaz Date: Thu, 23 Nov 2023 16:33:18 +0100 Subject: [PATCH 1/3] Include conda Library bin and lib paths to default blas ldflags search dirs --- pytensor/link/c/cmodule.py | 11 ++++++++- tests/link/c/test_cmodule.py | 45 ++++++++++++++++++++++++++++++++++++ 2 files changed, 55 insertions(+), 1 deletion(-) diff --git a/pytensor/link/c/cmodule.py b/pytensor/link/c/cmodule.py index a47fd978b9..35a37f1303 100644 --- a/pytensor/link/c/cmodule.py +++ b/pytensor/link/c/cmodule.py @@ -2744,7 +2744,9 @@ def get_cxx_library_dirs(): [pathlib.Path(p).resolve() for p in line[len("libraries: =") :].split(":")] for line in stdout.decode(sys.stdout.encoding).splitlines() if line.startswith("libraries: =") - ][0] + ] + if len(maybe_lib_dirs) > 0: + maybe_lib_dirs = maybe_lib_dirs[0] return [str(d) for d in maybe_lib_dirs if d.exists() and d.is_dir()] def check_libs( @@ -2793,6 +2795,13 @@ def check_libs( cxx_library_dirs = get_cxx_library_dirs() searched_library_dirs = cxx_library_dirs + _std_lib_dirs + if sys.platform == "win32": + # Conda on Windows saves MKL libraries under CONDA_PREFIX\Library\bin + # From the conda manual (https://docs.conda.io/projects/conda-build/en/stable/user-guide/environment-variables.html) + # it seems like conda could also save some libraries into the CONDA_PREFIX\Library\lib + # directory. We will include both in our searched library dirs + searched_library_dirs.append(os.path.join(sys.prefix, "Library", "bin")) + searched_library_dirs.append(os.path.join(sys.prefix, "Library", "lib")) all_libs = [ l for path in [ diff --git a/tests/link/c/test_cmodule.py b/tests/link/c/test_cmodule.py index 176779971a..574e224368 100644 --- a/tests/link/c/test_cmodule.py +++ b/tests/link/c/test_cmodule.py @@ -260,6 +260,51 @@ def test_default_blas_ldflags_no_cxx(): assert default_blas_ldflags() == "" +@pytest.fixture() +def windows_conda_libs(blas_libs): + libtemplate = "{lib}.dll" + libraries = [] + with tempfile.TemporaryDirectory() as d: + subdir = os.path.join(d, "Library", "bin") + os.makedirs(subdir, exist_ok=True) + flags = f'-L"{subdir}"' + for lib in blas_libs: + lib_path = os.path.join(subdir, libtemplate.format(lib=lib)) + with open(lib_path, "wb") as f: + f.write(b"1") + libraries.append(lib_path) + 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 d, flags + + +@patch("pytensor.link.c.cmodule.std_lib_dirs", return_value=[]) +@patch("pytensor.link.c.cmodule.check_mkl_openmp", return_value=None) +def test_default_blas_ldflags_conda_windows( + mock_std_lib_dirs, mock_check_mkl_openmp, windows_conda_libs +): + mock_sys_prefix, expected_blas_ldflags = windows_conda_libs + mock_process = MagicMock() + mock_process.communicate = lambda *args, **kwargs: (b"", b"") + mock_process.returncode = 0 + with patch("sys.platform", "win32"): + with patch("sys.prefix", mock_sys_prefix): + 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), + ): + assert set(default_blas_ldflags().split(" ")) == set( + expected_blas_ldflags.split(" ") + ) + + @patch( "os.listdir", return_value=["mkl_core.1.dll", "mkl_rt.1.0.dll", "mkl_rt.1.1.lib"] ) From 07fa4e59fcef0ace4970a89f31ba812ea3b8e05c Mon Sep 17 00:00:00 2001 From: lucianopaz Date: Thu, 23 Nov 2023 22:59:57 +0100 Subject: [PATCH 2/3] Ignore import warnings in test_ExternalCOp_c_code_cache_version --- tests/link/c/test_op.py | 13 +++++++------ 1 file changed, 7 insertions(+), 6 deletions(-) diff --git a/tests/link/c/test_op.py b/tests/link/c/test_op.py index 95b65feed0..c9f40bbb72 100644 --- a/tests/link/c/test_op.py +++ b/tests/link/c/test_op.py @@ -207,11 +207,11 @@ def get_hash(modname, seed=None): cmd_line, stdin=subprocess.PIPE, stdout=subprocess.PIPE, - stderr=subprocess.STDOUT, + stderr=subprocess.PIPE, env=env, ) out, err = p.communicate() - return out, err + return out, err, p.returncode def test_ExternalCOp_c_code_cache_version(): @@ -222,10 +222,11 @@ def test_ExternalCOp_c_code_cache_version(): tmp.seek(0) # modname = os.path.splitext(tmp.name)[0] modname = tmp.name - out_1, err = get_hash(modname, seed=428) - assert err is None - out_2, err = get_hash(modname, seed=3849) - assert err is None + out_1, err1, returncode1 = get_hash(modname, seed=428) + out_2, err2, returncode2 = get_hash(modname, seed=3849) + assert returncode1 == 0 + assert returncode2 == 0 + assert err1 == err2 hash_1, msg, _ = out_1.decode().split("\n") assert msg == "__success__" From 7943c3d2a51c261134f185b84ae07ef7b25db8b6 Mon Sep 17 00:00:00 2001 From: lucianopaz Date: Thu, 23 Nov 2023 23:08:17 +0100 Subject: [PATCH 3/3] Make test_debugprint robust to environments without blas__ldflags --- tests/test_printing.py | 31 +++++++++++++++++-------------- 1 file changed, 17 insertions(+), 14 deletions(-) diff --git a/tests/test_printing.py b/tests/test_printing.py index 7e4fa2064b..9f949362ad 100644 --- a/tests/test_printing.py +++ b/tests/test_printing.py @@ -272,29 +272,32 @@ def test_debugprint(): print_view_map=True, ) s = s.getvalue() + Gemv_op_name = "CGemv" if pytensor.config.blas__ldflags else "Gemv" exp_res = dedent( r""" Composite{(i2 + (i0 - i1))} 4 - ├─ ExpandDims{axis=0} v={0: [0]} 3 - │ └─ CGemv{inplace} d={0: [0]} 2 - │ ├─ AllocEmpty{dtype='float64'} 1 - │ │ └─ Shape_i{0} 0 - │ │ └─ B - │ ├─ 1.0 - │ ├─ B - │ ├─ - │ └─ 0.0 - ├─ D - └─ A + ├─ ExpandDims{axis=0} v={0: [0]} 3 + """ + f" │ └─ {Gemv_op_name}{{inplace}} d={{0: [0]}} 2" + r""" + │ ├─ AllocEmpty{dtype='float64'} 1 + │ │ └─ Shape_i{0} 0 + │ │ └─ B + │ ├─ 1.0 + │ ├─ B + │ ├─ + │ └─ 0.0 + ├─ D + └─ A Inner graphs: Composite{(i2 + (i0 - i1))} - ← add 'o0' + ← add 'o0' ├─ i2 └─ sub - ├─ i0 - └─ i1 + ├─ i0 + └─ i1 """ ).lstrip()