xref: /aosp_15_r20/external/pytorch/defs.bzl (revision da0073e96a02ea20f0ac840b70461e3646d07c45)
1def get_blas_gomp_arch_deps():
2    return [
3        ("x86_64", [
4            "fbsource//third-party/mkl:{}".format(native.read_config("fbcode", "mkl_lp64", "mkl_lp64_omp")),
5        ]),
6        ("aarch64", [
7            "third-party//OpenBLAS:OpenBLAS",
8            "third-party//openmp:omp",
9        ]),
10    ]
11
12default_compiler_flags = [
13    "-Wall",
14    "-Wextra",
15    "-Wno-unused-function",
16    "-Wno-unused-parameter",
17    "-Wno-error=strict-aliasing",
18    "-Wno-shadow-compatible-local",
19    "-Wno-maybe-uninitialized",  # aten is built with gcc as part of HHVM
20    "-Wno-unknown-pragmas",
21    "-Wno-strict-overflow",
22    # See https://fb.facebook.com/groups/fbcode/permalink/1813348245368673/
23    # These trigger on platform007
24    "-Wno-stringop-overflow",
25    "-Wno-class-memaccess",
26    "-DHAVE_MMAP",
27    "-DUSE_GCC_ATOMICS=1",
28    "-D_FILE_OFFSET_BITS=64",
29    "-DHAVE_SHM_OPEN=1",
30    "-DHAVE_SHM_UNLINK=1",
31    "-DHAVE_MALLOC_USABLE_SIZE=1",
32    "-DCPU_CAPABILITY_DEFAULT",
33    "-DTH_INDEX_BASE=0",
34    "-DMAGMA_V2",
35    "-DNO_CUDNN_DESTROY_HANDLE",
36    "-DUSE_FBGEMM",
37    "-DUSE_PYTORCH_QNNPACK",
38    # The dynamically loaded NVRTC trick doesn't work in fbcode,
39    # and it's not necessary anyway, because we have a stub
40    # nvrtc library which we load canonically anyway
41    "-DUSE_DIRECT_NVRTC",
42    "-DUSE_RUY_QMATMUL",
43] + select({
44    # XNNPACK depends on an updated version of pthreadpool interface, whose implementation
45    # includes <pthread.h> - a header not available on Windows.
46    "DEFAULT": ["-DUSE_XNNPACK"],
47    "ovr_config//os:windows": [],
48}) + (["-O1"] if native.read_config("fbcode", "build_mode_test_label", "") == "dev-nosan" else [])
49
50compiler_specific_flags = {
51    "clang": [
52        "-Wno-absolute-value",
53        "-Wno-pass-failed",
54        "-Wno-braced-scalar-init",
55    ],
56    "gcc": [
57        "-Wno-error=array-bounds",
58    ],
59}
60
61def get_cpu_parallel_backend_flags():
62    parallel_backend = native.read_config("pytorch", "parallel_backend", "openmp")
63    defs = []
64    if parallel_backend == "openmp":
65        defs.append("-DAT_PARALLEL_OPENMP_FBCODE=1")
66    elif parallel_backend == "native":
67        defs.append("-DAT_PARALLEL_NATIVE_FBCODE=1")
68    else:
69        fail("Unsupported parallel backend: " + parallel_backend)
70    if native.read_config("pytorch", "exp_single_thread_pool", "0") == "1":
71        defs.append("-DAT_EXPERIMENTAL_SINGLE_THREAD_POOL=1")
72    mkl_ver = native.read_config("fbcode", "mkl_lp64", "mkl_lp64_omp")
73    if mkl_ver == "mkl_lp64_seq":
74        defs.append("-DATEN_MKL_SEQUENTIAL_FBCODE=1")
75    return defs
76
77def is_cpu_static_dispatch_build():
78    mode = native.read_config("fbcode", "caffe2_static_dispatch_mode", "none")
79    return mode == "cpu"
80