diff --git a/mlir/CMakeLists.txt b/mlir/CMakeLists.txt --- a/mlir/CMakeLists.txt +++ b/mlir/CMakeLists.txt @@ -88,6 +88,8 @@ message(STATUS "Python prefix = '${PYTHON_MODULE_PREFIX}', " "suffix = '${PYTHON_MODULE_SUFFIX}', " "extension = '${PYTHON_MODULE_EXTENSION}") + mlir_detect_numpy_install() + message(STATUS "Found numpy v${numpy_VERSION}: ${numpy_INCLUDE_DIR}") endif() include_directories( "include") diff --git a/mlir/cmake/modules/MLIRDetectPythonEnv.cmake b/mlir/cmake/modules/MLIRDetectPythonEnv.cmake --- a/mlir/cmake/modules/MLIRDetectPythonEnv.cmake +++ b/mlir/cmake/modules/MLIRDetectPythonEnv.cmake @@ -24,3 +24,26 @@ set(pybind11_DIR "${PACKAGE_DIR}" PARENT_SCOPE) endif() endfunction() + +# Detects a numpy package installed in the current python environment +# and sets variables to indicate numpy version and directory. +function(mlir_detect_numpy_install) + message(STATUS "Checking for numpy in python path...") + execute_process( + COMMAND "${Python3_EXECUTABLE}" + -c "import numpy as np; print(np.__version__); print(np.get_include());" + WORKING_DIRECTORY ${CMAKE_CURRENT_SOURCE_DIR} + RESULT_VARIABLE STATUS + OUTPUT_VARIABLE numpy_VALUES_OUTPUT + ERROR_QUIET) + if(NOT STATUS EQUAL "0") + message(WARNING "not found (install via 'pip install numpy')") + return() + endif() + string(REGEX REPLACE "\n" ";" numpy_VALUES ${numpy_VALUES_OUTPUT}) + list(GET numpy_VALUES 0 numpy_VERSION) + list(GET numpy_VALUES 1 numpy_INCLUDE_DIR) + message(STATUS "found (${numpy_INCLUDE_DIR})") + set(numpy_VERSION "${numpy_VERSION}" PARENT_SCOPE) + set(numpy_INCLUDE_DIR "${numpy_INCLUDE_DIR}" PARENT_SCOPE) +endfunction()