[llvm] bbf11c8 - [mlgo] Fix cmake logic detecting tf pip package location

[Mircea Trofin via llvm-commits]

Author: Mircea Trofin
Date: 2022-08-24T08:29:03-07:00
New Revision: bbf11c8d9f4bbe88c49500c26029e9ecccc0ce37

URL: https://github.com/llvm/llvm-project/commit/bbf11c8d9f4bbe88c49500c26029e9ecccc0ce37
DIFF: https://github.com/llvm/llvm-project/commit/bbf11c8d9f4bbe88c49500c26029e9ecccc0ce37.diff

LOG: [mlgo] Fix cmake logic detecting tf pip package location

New logic works for both `tensorflow` and `tf-nightly`.

Added: 
    

Modified: 
    llvm/CMakeLists.txt

Removed: 
    


################################################################################
diff  --git a/llvm/CMakeLists.txt b/llvm/CMakeLists.txt
index c31967520565..b52e2d04b646 100644
--- a/llvm/CMakeLists.txt
+++ b/llvm/CMakeLists.txt
@@ -962,15 +962,13 @@ if (TENSORFLOW_C_LIB_PATH)
   if (NOT TF_PROTO_HEADERS)
     message(STATUS "TF_PROTO_HEADERS not defined. Looking for tensorflow pip package.")
     execute_process(COMMAND
-      ${Python3_EXECUTABLE} "-m" "pip" "show" "tensorflow"
-      OUTPUT_VARIABLE TF_PIP_OUT)
-    if ("${TF_PIP_OUT}" STREQUAL "")
+      ${Python3_EXECUTABLE} "-c" "import tensorflow as tf; import os; print(os.path.dirname(tf.__file__))"
+      OUTPUT_VARIABLE TF_PIP)
+    if ("${TF_PIP}" STREQUAL "")
       message(FATAL ERROR "Tensorflow pip package is also required for 'development' mode (protobuf headers)")
     endif()
-    string(REGEX MATCH "Location: ([^\n]*\n)" TF_PIP_LOC "${TF_PIP_OUT}")
-    string(REPLACE "Location: " "" TF_PIP ${TF_PIP_LOC})
     string(STRIP ${TF_PIP} TF_PIP)
-    set(TF_PROTO_HEADERS "${TF_PIP}/tensorflow/include")
+    set(TF_PROTO_HEADERS "${TF_PIP}/include")
   endif()
   message(STATUS "Using Tensorflow headers under: ${TF_PROTO_HEADERS}")
   include_directories(${TF_PROTO_HEADERS})