Fix mkl-static issue for Windows. (#130697)

xuhancn · pytorchmergebot · commit f1456c74a09e · 2024-07-15T19:28:11.000Z
Background: We found the pytorch Windows release/2.4 performance regression: #130619 After some debug works, I found the pytorch Windows static mkl build options are wrong: <img width="1049" alt="image" src="https://github.com/user-attachments/assets/38692142-bfca-4c98-8092-6e105c82bb13"> 1. Thread lib is wrong. 2. Miss `openmp` lib and config. > Debug history: #130619 (comment) and #130619 (comment) This PR will fix `mkl-static` build options issue. <img width="863" alt="image" src="https://github.com/user-attachments/assets/834f6cee-7e6d-4d74-b2bc-8a270f05e429"> Reference: <img width="482" alt="image" src="https://github.com/user-attachments/assets/8184dadb-f230-4062-a49f-51df1d7285f5"> https://www.intel.com/content/www/us/en/developer/tools/oneapi/onemkl-link-line-advisor.html#gs.c6izlg Pull Request resolved: #130697 Approved by: https://github.com/jgong5, https://github.com/atalman
diff --git a/cmake/Modules/FindMKL.cmake b/cmake/Modules/FindMKL.cmake
@@ -88,25 +88,41 @@ ELSE ("${SIZE_OF_VOIDP}" EQUAL 8)
   SET(iccvers "ia32")
   SET(mkl64s)
 ENDIF ("${SIZE_OF_VOIDP}" EQUAL 8)
-IF(CMAKE_COMPILER_IS_GNUCC)
-  IF ("${MKL_THREADING}" STREQUAL "TBB")
-    SET(mklthreads "mkl_tbb_thread")
-    SET(mklrtls "tbb")
-  ELSE()
-    SET(mklthreads "mkl_gnu_thread" "mkl_intel_thread")
-    SET(mklrtls "gomp" "iomp5")
-  ENDIF()
-  SET(mklifaces  "intel" "gf")
-ELSE(CMAKE_COMPILER_IS_GNUCC)
+
+IF(WIN32)
   IF ("${MKL_THREADING}" STREQUAL "TBB")
     SET(mklthreads "mkl_tbb_thread")
-    SET(mklrtls "tbb")
+    IF (CMAKE_BUILD_TYPE STREQUAL Debug)
+      SET(mklrtls "tbb12_debug")
+    ELSE ()
+      SET(mklrtls "tbb12")
+    ENDIF ()
   ELSE()
     SET(mklthreads "mkl_intel_thread")
-    SET(mklrtls "iomp5" "guide")
+    SET(mklrtls "libiomp5md")
   ENDIF()
-  SET(mklifaces  "intel")
-ENDIF (CMAKE_COMPILER_IS_GNUCC)
+    SET(mklifaces  "intel")
+ELSE(WIN32)
+  IF(CMAKE_COMPILER_IS_GNUCC)
+    IF ("${MKL_THREADING}" STREQUAL "TBB")
+      SET(mklthreads "mkl_tbb_thread")
+      SET(mklrtls "tbb")
+    ELSE()
+      SET(mklthreads "mkl_gnu_thread" "mkl_intel_thread")
+      SET(mklrtls "gomp" "iomp5")
+    ENDIF()
+    SET(mklifaces  "intel" "gf")
+  ELSE(CMAKE_COMPILER_IS_GNUCC)
+    IF ("${MKL_THREADING}" STREQUAL "TBB")
+      SET(mklthreads "mkl_tbb_thread")
+      SET(mklrtls "tbb")
+    ELSE()
+      SET(mklthreads "mkl_intel_thread")
+      SET(mklrtls "iomp5" "guide")
+    ENDIF()
+    SET(mklifaces  "intel")
+  ENDIF (CMAKE_COMPILER_IS_GNUCC)
+ENDIF(WIN32)
 
 # Kernel libraries dynamically loaded
 SET(mklkerlibs "mc" "mc3" "nc" "p4n" "p4m" "p4m3" "p4p" "def")
