xref: /aosp_15_r20/external/pytorch/caffe2/perfkernels/CMakeLists.txt (revision da0073e96a02ea20f0ac840b70461e3646d07c45)
1if(INTERN_BUILD_MOBILE)
2  list(APPEND Caffe2_CPU_SRCS
3    "${CMAKE_CURRENT_SOURCE_DIR}/embedding_lookup_idx.cc"
4  )
5  set(Caffe2_CPU_SRCS ${Caffe2_CPU_SRCS} PARENT_SCOPE)
6  return()
7endif()
8
9# ---[ CPU files.
10file(GLOB common_srcs *.cc)
11file(GLOB avx_srcs *_avx.cc)
12file(GLOB avx2_srcs *_avx2.cc)
13# exclude avx and avx2 srcs from common_srcs
14exclude(common_srcs "${common_srcs}" ${avx_srcs})
15exclude(common_srcs "${common_srcs}" ${avx2_srcs})
16
17# We will always build common srcs.
18set(Caffe2_CPU_SRCS ${Caffe2_CPU_SRCS} ${common_srcs})
19
20# We will only build the perf kernel files if the compiler supports avx2
21# extensions.
22if(CXX_AVX2_FOUND)
23  add_library(Caffe2_perfkernels_avx2 STATIC ${avx2_srcs})
24  target_link_libraries(Caffe2_perfkernels_avx2 PRIVATE c10)
25
26  if(MSVC AND NOT "${CMAKE_CXX_COMPILER_ID}" STREQUAL "Clang")
27    target_compile_options(Caffe2_perfkernels_avx2
28        PRIVATE "/arch:AVX2"
29        PRIVATE "/D__FMA__"
30        PRIVATE "/D__F16C__")
31  else()
32    target_compile_options(Caffe2_perfkernels_avx2
33        PRIVATE "-mavx2"
34        PRIVATE "-mfma"
35        PRIVATE "-mavx"
36        PRIVATE "-mf16c")
37  endif()
38  caffe2_interface_library(
39      Caffe2_perfkernels_avx2 Caffe2_perfkernels_avx2_interface)
40  list(APPEND
41       Caffe2_DEPENDENCY_WHOLE_LINK_LIBS
42       "Caffe2_perfkernels_avx2_interface")
43endif()
44
45# TODO(jiayq): currently, we only implement the very base files for the
46# perfkernels. This is because to implement avx and avx2 files, we actually
47# need to set up different compilation units and this is a bit more involving
48# in terms of CMakefile changes. This is a stop-gap solution until we get a
49# more proper implementation.
50
51set(Caffe2_CPU_SRCS ${Caffe2_CPU_SRCS} PARENT_SCOPE)
52set(Caffe2_DEPENDENCY_WHOLE_LINK_LIBS
53    ${Caffe2_DEPENDENCY_WHOLE_LINK_LIBS}
54    PARENT_SCOPE)
55