Add macro required by CINN. (#37066)

* Add macro required by CINN. * Remove CMAKE_BUILD_TYPE form cinn.cmake.
PaddlePaddle · Nov 11, 2021 · 9a9345f · 9a9345f
1 parent a2da1ef
commit 9a9345f
Show file tree

Hide file tree

Showing 2 changed files with 16 additions and 1 deletion.
diff --git a/cmake/external/cinn.cmake b/cmake/external/cinn.cmake
@@ -28,7 +28,14 @@ include(ExternalProject)
 set(CINN_SOURCE_DIR ${THIRD_PARTY_PATH}/CINN)
 # TODO(zhhsplendid): Modify git tag after we have release tag
 set(CINN_GIT_TAG develop)
-set(CINN_OPTIONAL_ARGS -DPY_VERSION=${PY_VERSION} -DWITH_CUDA=${WITH_GPU} -DWITH_CUDNN=${WITH_GPU} -DPUBLISH_LIBS=ON -DWITH_TESTING=ON)
+set(CINN_OPTIONAL_ARGS -DPY_VERSION=${PY_VERSION}
+                       -DWITH_CUDA=${WITH_GPU}
+                       -DWITH_CUDNN=${WITH_GPU}
+                       -DWITH_MKL_CBLAS=${WITH_MKL}
+                       -DWITH_MKLDNN=${WITH_MKL}
+                       -DPUBLISH_LIBS=ON
+                       -DWITH_TESTING=ON
+)
 set(CINN_BUILD_COMMAND $(MAKE) cinnapi -j)
 ExternalProject_Add(
   external_cinn

diff --git a/cmake/third_party.cmake b/cmake/third_party.cmake
@@ -364,6 +364,14 @@ if (WITH_CINN)
     message(STATUS "Compile Paddle with CINN.")
     include(external/cinn)
     add_definitions(-DPADDLE_WITH_CINN)
+    if (WITH_GPU)
+        add_definitions(-DCINN_WITH_CUDA)
+        add_definitions(-DCINN_WITH_CUDNN)
+    endif (WITH_GPU)
+    if (WITH_MKL)
+        add_definitions(-DCINN_WITH_MKL_CBLAS)
+        add_definitions(-DCINN_WITH_MKLDNN)
+    endif (WITH_MKL)
 endif (WITH_CINN)
 
 if (WITH_CRYPTO)