cmake: allow USE_SYSTEM_ZSTD (pytorch#104611)

Fixes pytorch#44255.

This is part of larger work I'm doing to allow for more `USE_SYSTEM_*` options to allow Nix to have faster re-builds of PyTorch: NixOS/nixpkgs#239291.

Pull Request resolved: pytorch#104611
Approved by: https://github.com/ezyang, https://github.com/malfet

CMakeLists.txt

@@ -413,6 +413,7 @@ option(USE_SYSTEM_FXDIV "Use system-provided fxdiv." OFF)
 option(USE_SYSTEM_BENCHMARK "Use system-provided google benchmark." OFF)
 option(USE_SYSTEM_ONNX "Use system-provided onnx." OFF)
 option(USE_SYSTEM_XNNPACK "Use system-provided xnnpack." OFF)
+option(USE_SYSTEM_ZSTD "Use system-provided zstd." OFF)
 option(USE_GOLD_LINKER "Use ld.gold to link" OFF)
 if(USE_SYSTEM_LIBS)
   set(USE_SYSTEM_CPUINFO ON)
@@ -434,6 +435,9 @@ if(USE_SYSTEM_LIBS)
   if(USE_TBB)
     set(USE_SYSTEM_TBB ON)
   endif()
+  if(USE_ZSTD)
+    set(USE_SYSTEM_ZSTD ON)
+  endif()
 endif()
 
 # Used when building Caffe2 through setup.py

cmake/Dependencies.cmake

@@ -1505,10 +1505,22 @@ if(NOT INTERN_BUILD_MOBILE AND BUILD_CAFFE2_OPS)
 endif()
 
 if(USE_ZSTD)
-  list(APPEND Caffe2_DEPENDENCY_LIBS libzstd_static)
-  include_directories(SYSTEM ${CMAKE_CURRENT_LIST_DIR}/../third_party/zstd/lib)
-  add_subdirectory(${CMAKE_CURRENT_LIST_DIR}/../third_party/zstd/build/cmake)
-  set_property(TARGET libzstd_static PROPERTY POSITION_INDEPENDENT_CODE ON)
+  if(USE_SYSTEM_ZSTD)
+    find_package(zstd REQUIRED)
+    if(TARGET zstd::libzstd_shared)
+      set(ZSTD_TARGET zstd::libzstd_shared)
+    else()
+      set(ZSTD_TARGET zstd::libzstd_static)
+    endif()
+    list(APPEND Caffe2_DEPENDENCY_LIBS ${ZSTD_TARGET})
+    get_property(ZSTD_INCLUDE_DIR TARGET ${ZSTD_TARGET} PROPERTY INTERFACE_INCLUDE_DIRECTORIES)
+    include_directories(SYSTEM ${ZSTD_INCLUDE_DIR})
+  else()
+    list(APPEND Caffe2_DEPENDENCY_LIBS libzstd_static)
+    include_directories(SYSTEM ${CMAKE_CURRENT_LIST_DIR}/../third_party/zstd/lib)
+    add_subdirectory(${CMAKE_CURRENT_LIST_DIR}/../third_party/zstd/build/cmake)
+    set_property(TARGET libzstd_static PROPERTY POSITION_INDEPENDENT_CODE ON)
+  endif()
 endif()
 
 # ---[ Onnx