| load("@bazel_skylib//lib:paths.bzl", "paths") |
| load("@pybind11_bazel//:build_defs.bzl", "pybind_extension") |
| load("@rules_cc//cc:defs.bzl", "cc_binary", "cc_library", "cc_test") |
| load("@rules_python//python:defs.bzl", "py_library", "py_test") |
| load("@pytorch//third_party:substitution.bzl", "header_template_rule", "template_rule") |
| load("@pytorch//:tools/bazel.bzl", "rules") |
| load("@pytorch//tools/rules:cu.bzl", "cu_library") |
| load("@pytorch//tools/config:defs.bzl", "if_cuda") |
| load("@pytorch//:aten.bzl", "generate_aten", "intern_build_aten_ops") |
| load(":build.bzl", "GENERATED_AUTOGRAD_CPP", "GENERATED_AUTOGRAD_PYTHON", "define_targets") |
| load(":build_variables.bzl", "jit_core_sources", "lazy_tensor_ts_sources", "libtorch_core_sources", "libtorch_cuda_sources", "libtorch_distributed_sources", "libtorch_extra_sources", "libtorch_python_core_sources", "torch_cpp_srcs", "libtorch_python_cuda_sources", "libtorch_python_distributed_sources") |
| load(":ufunc_defs.bzl", "aten_ufunc_generated_cpu_kernel_sources", "aten_ufunc_generated_cpu_sources", "aten_ufunc_generated_cuda_sources") |
| load("//:tools/bazel.bzl", "rules") |
| |
| define_targets(rules = rules) |
| |
| COMMON_COPTS = [ |
| "-DHAVE_MALLOC_USABLE_SIZE=1", |
| "-DHAVE_MMAP=1", |
| "-DHAVE_SHM_OPEN=1", |
| "-DHAVE_SHM_UNLINK=1", |
| "-D_FILE_OFFSET_BITS=64", |
| "-DUSE_FBGEMM", |
| "-DUSE_DISTRIBUTED", |
| "-DAT_PER_OPERATOR_HEADERS", |
| "-DATEN_THREADING=NATIVE", |
| "-DNO_CUDNN_DESTROY_HANDLE", |
| ] + if_cuda([ |
| "-DUSE_CUDA", |
| "-DUSE_CUDNN", |
| # TODO: This should be passed only when building for CUDA-11.5 or newer |
| # use cub in a safe manner, see: |
| # https://github.com/pytorch/pytorch/pull/55292 |
| "-DCUB_WRAPPED_NAMESPACE=at_cuda_detail", |
| ]) |
| |
| aten_generation_srcs = ["aten/src/ATen/native/native_functions.yaml"] + ["aten/src/ATen/native/tags.yaml"] + glob(["aten/src/ATen/templates/**"]) |
| |
| generated_cpu_cpp = [ |
| "aten/src/ATen/RegisterBackendSelect.cpp", |
| "aten/src/ATen/RegisterCPU.cpp", |
| "aten/src/ATen/RegisterFunctionalization_0.cpp", |
| "aten/src/ATen/RegisterFunctionalization_1.cpp", |
| "aten/src/ATen/RegisterFunctionalization_2.cpp", |
| "aten/src/ATen/RegisterFunctionalization_3.cpp", |
| # "aten/src/ATen/RegisterFunctionalizationEverything.cpp", |
| "aten/src/ATen/RegisterMkldnnCPU.cpp", |
| "aten/src/ATen/RegisterNestedTensorCPU.cpp", |
| "aten/src/ATen/RegisterQuantizedCPU.cpp", |
| "aten/src/ATen/RegisterSparseCPU.cpp", |
| "aten/src/ATen/RegisterSparseCsrCPU.cpp", |
| "aten/src/ATen/RegisterZeroTensor.cpp", |
| "aten/src/ATen/RegisterCompositeImplicitAutograd.cpp", |
| "aten/src/ATen/RegisterCompositeImplicitAutogradNestedTensor.cpp", |
| "aten/src/ATen/RegisterCompositeExplicitAutograd.cpp", |
| "aten/src/ATen/RegisterCompositeExplicitAutogradNonFunctional.cpp", |
| "aten/src/ATen/RegisterMeta.cpp", |
| "aten/src/ATen/RegisterSparseMeta.cpp", |
| "aten/src/ATen/RegisterQuantizedMeta.cpp", |
| "aten/src/ATen/RegisterNestedTensorMeta.cpp", |
| "aten/src/ATen/RegisterSchema.cpp", |
| "aten/src/ATen/CPUFunctions.h", |
| "aten/src/ATen/CPUFunctions_inl.h", |
| "aten/src/ATen/CompositeExplicitAutogradFunctions.h", |
| "aten/src/ATen/CompositeExplicitAutogradFunctions_inl.h", |
| "aten/src/ATen/CompositeExplicitAutogradNonFunctionalFunctions.h", |
| "aten/src/ATen/CompositeExplicitAutogradNonFunctionalFunctions_inl.h", |
| "aten/src/ATen/CompositeImplicitAutogradFunctions.h", |
| "aten/src/ATen/CompositeImplicitAutogradFunctions_inl.h", |
| "aten/src/ATen/CompositeImplicitAutogradNestedTensorFunctions.h", |
| "aten/src/ATen/CompositeImplicitAutogradNestedTensorFunctions_inl.h", |
| "aten/src/ATen/CompositeViewCopyKernels.cpp", |
| "aten/src/ATen/FunctionalInverses.h", |
| "aten/src/ATen/Functions.h", |
| "aten/src/ATen/Functions.cpp", |
| "aten/src/ATen/RedispatchFunctions.h", |
| "aten/src/ATen/Operators.h", |
| "aten/src/ATen/Operators_0.cpp", |
| "aten/src/ATen/Operators_1.cpp", |
| "aten/src/ATen/Operators_2.cpp", |
| "aten/src/ATen/Operators_3.cpp", |
| "aten/src/ATen/Operators_4.cpp", |
| "aten/src/ATen/NativeFunctions.h", |
| "aten/src/ATen/MetaFunctions.h", |
| "aten/src/ATen/MetaFunctions_inl.h", |
| "aten/src/ATen/MethodOperators.h", |
| "aten/src/ATen/NativeMetaFunctions.h", |
| "aten/src/ATen/RegistrationDeclarations.h", |
| "aten/src/ATen/VmapGeneratedPlumbing.h", |
| "aten/src/ATen/core/aten_interned_strings.h", |
| "aten/src/ATen/core/enum_tag.h", |
| "aten/src/ATen/core/TensorBody.h", |
| "aten/src/ATen/core/TensorMethods.cpp", |
| "aten/src/ATen/core/ATenOpList.cpp", |
| ] |
| |
| generated_cuda_cpp = [ |
| "aten/src/ATen/CUDAFunctions.h", |
| "aten/src/ATen/CUDAFunctions_inl.h", |
| "aten/src/ATen/RegisterCUDA.cpp", |
| "aten/src/ATen/RegisterNestedTensorCUDA.cpp", |
| "aten/src/ATen/RegisterQuantizedCUDA.cpp", |
| "aten/src/ATen/RegisterSparseCUDA.cpp", |
| "aten/src/ATen/RegisterSparseCsrCUDA.cpp", |
| ] |
| |
| generate_aten( |
| name = "generated_aten_cpp", |
| srcs = aten_generation_srcs, |
| outs = ( |
| generated_cpu_cpp + |
| generated_cuda_cpp + |
| aten_ufunc_generated_cpu_sources("aten/src/ATen/{}") + |
| aten_ufunc_generated_cpu_kernel_sources("aten/src/ATen/{}") + |
| aten_ufunc_generated_cuda_sources("aten/src/ATen/{}") + [ |
| "aten/src/ATen/Declarations.yaml", |
| ] |
| ), |
| generator = "//torchgen:gen", |
| ) |
| |
| filegroup( |
| name = "cpp_generated_code", |
| srcs = GENERATED_AUTOGRAD_CPP, |
| data = [":generate-code"], |
| ) |
| |
| # ATen |
| filegroup( |
| name = "aten_base_cpp", |
| srcs = glob([ |
| "aten/src/ATen/*.cpp", |
| "aten/src/ATen/functorch/*.cpp", |
| "aten/src/ATen/detail/*.cpp", |
| "aten/src/ATen/cpu/*.cpp", |
| ]), |
| ) |
| |
| filegroup( |
| name = "ATen_CORE_SRCS", |
| srcs = glob( |
| [ |
| "aten/src/ATen/core/**/*.cpp", |
| ], |
| exclude = [ |
| "aten/src/ATen/core/**/*_test.cpp", |
| ], |
| ), |
| ) |
| |
| filegroup( |
| name = "aten_native_cpp", |
| srcs = glob(["aten/src/ATen/native/*.cpp"]), |
| ) |
| |
| filegroup( |
| name = "aten_native_sparse_cpp", |
| srcs = glob(["aten/src/ATen/native/sparse/*.cpp"]), |
| ) |
| |
| filegroup( |
| name = "aten_native_nested_cpp", |
| srcs = glob(["aten/src/ATen/native/nested/*.cpp"]), |
| ) |
| |
| filegroup( |
| name = "aten_native_quantized_cpp", |
| srcs = glob( |
| [ |
| "aten/src/ATen/native/quantized/*.cpp", |
| "aten/src/ATen/native/quantized/cpu/*.cpp", |
| ], |
| ), |
| ) |
| |
| filegroup( |
| name = "aten_native_transformers_cpp", |
| srcs = glob(["aten/src/ATen/native/transformers/*.cpp"]), |
| ) |
| |
| filegroup( |
| name = "aten_native_mkl_cpp", |
| srcs = glob([ |
| "aten/src/ATen/native/mkl/*.cpp", |
| "aten/src/ATen/mkl/*.cpp", |
| ]), |
| ) |
| |
| filegroup( |
| name = "aten_native_mkldnn_cpp", |
| srcs = glob(["aten/src/ATen/native/mkldnn/*.cpp"]), |
| ) |
| |
| filegroup( |
| name = "aten_native_xnnpack", |
| srcs = glob(["aten/src/ATen/native/xnnpack/*.cpp"]), |
| ) |
| |
| filegroup( |
| name = "aten_base_vulkan", |
| srcs = glob(["aten/src/ATen/vulkan/*.cpp"]), |
| ) |
| |
| filegroup( |
| name = "aten_base_metal", |
| srcs = glob(["aten/src/ATen/metal/*.cpp"]), |
| ) |
| |
| filegroup( |
| name = "ATen_QUANTIZED_SRCS", |
| srcs = glob( |
| [ |
| "aten/src/ATen/quantized/**/*.cpp", |
| ], |
| exclude = [ |
| "aten/src/ATen/quantized/**/*_test.cpp", |
| ], |
| ), |
| ) |
| |
| filegroup( |
| name = "aten_cuda_cpp_srcs", |
| srcs = glob( |
| [ |
| "aten/src/ATen/cuda/*.cpp", |
| "aten/src/ATen/cuda/detail/*.cpp", |
| "aten/src/ATen/cuda/tunable/*.cpp", |
| "aten/src/ATen/cudnn/*.cpp", |
| "aten/src/ATen/native/cuda/*.cpp", |
| "aten/src/ATen/native/cuda/linalg/*.cpp", |
| "aten/src/ATen/native/cudnn/*.cpp", |
| "aten/src/ATen/native/miopen/*.cpp", |
| "aten/src/ATen/native/nested/cuda/*.cpp", |
| "aten/src/ATen/native/quantized/cuda/*.cpp", |
| "aten/src/ATen/native/quantized/cudnn/*.cpp", |
| "aten/src/ATen/native/sparse/cuda/*.cpp", |
| "aten/src/ATen/native/transformers/cuda/*.cpp", |
| ], |
| ), |
| ) |
| |
| filegroup( |
| name = "aten_cu_srcs", |
| srcs = glob([ |
| "aten/src/ATen/cuda/*.cu", |
| "aten/src/ATen/cuda/detail/*.cu", |
| "aten/src/ATen/native/cuda/*.cu", |
| "aten/src/ATen/native/nested/cuda/*.cu", |
| "aten/src/ATen/native/quantized/cuda/*.cu", |
| "aten/src/ATen/native/sparse/cuda/*.cu", |
| "aten/src/ATen/native/transformers/cuda/*.cu", |
| ]) + aten_ufunc_generated_cuda_sources("aten/src/ATen/{}"), |
| # It's a bit puzzling to me why it's not necessary to declare the |
| # target that generates these sources... |
| ) |
| |
| header_template_rule( |
| name = "aten_src_ATen_config", |
| src = "aten/src/ATen/Config.h.in", |
| out = "aten/src/ATen/Config.h", |
| include = "aten/src", |
| substitutions = { |
| "@AT_MKLDNN_ENABLED@": "1", |
| "@AT_MKLDNN_ACL_ENABLED@": "0", |
| "@AT_MKL_ENABLED@": "1", |
| "@AT_MKL_SEQUENTIAL@": "0", |
| "@AT_POCKETFFT_ENABLED@": "0", |
| "@AT_NNPACK_ENABLED@": "0", |
| "@CAFFE2_STATIC_LINK_CUDA_INT@": "0", |
| "@AT_BUILD_WITH_BLAS@": "1", |
| "@AT_BUILD_WITH_LAPACK@": "1", |
| "@AT_PARALLEL_OPENMP@": "0", |
| "@AT_PARALLEL_NATIVE@": "1", |
| "@AT_BLAS_F2C@": "0", |
| "@AT_BLAS_USE_CBLAS_DOT@": "1", |
| }, |
| ) |
| |
| header_template_rule( |
| name = "aten_src_ATen_cuda_config", |
| src = "aten/src/ATen/cuda/CUDAConfig.h.in", |
| out = "aten/src/ATen/cuda/CUDAConfig.h", |
| include = "aten/src", |
| substitutions = { |
| "@AT_CUDNN_ENABLED@": "1", |
| "@AT_CUSPARSELT_ENABLED@": "0", |
| "@AT_ROCM_ENABLED@": "0", |
| "@AT_MAGMA_ENABLED@": "0", |
| "@NVCC_FLAGS_EXTRA@": "", |
| }, |
| ) |
| |
| cc_library( |
| name = "aten_headers", |
| hdrs = [ |
| "torch/csrc/Export.h", |
| "torch/csrc/jit/frontend/function_schema_parser.h", |
| ] + glob( |
| [ |
| "aten/src/**/*.h", |
| "aten/src/**/*.hpp", |
| "aten/src/ATen/cuda/**/*.cuh", |
| "aten/src/ATen/native/**/*.cuh", |
| "aten/src/THC/*.cuh", |
| ], |
| ) + [ |
| ":aten_src_ATen_config", |
| ":generated_aten_cpp", |
| ], |
| includes = [ |
| "aten/src", |
| ], |
| deps = [ |
| "//c10", |
| ], |
| ) |
| |
| ATEN_COPTS = COMMON_COPTS + [ |
| "-DCAFFE2_BUILD_MAIN_LIBS", |
| "-DHAVE_AVX_CPU_DEFINITION", |
| "-DHAVE_AVX2_CPU_DEFINITION", |
| "-fvisibility-inlines-hidden", |
| "-fno-math-errno", |
| "-fno-trapping-math", |
| ] |
| |
| intern_build_aten_ops( |
| copts = ATEN_COPTS, |
| extra_impls = aten_ufunc_generated_cpu_kernel_sources("aten/src/ATen/{}"), |
| deps = [ |
| ":aten_headers", |
| "@fbgemm", |
| "@mkl", |
| "@sleef", |
| "@mkl_dnn//:mkl-dnn", |
| ], |
| ) |
| |
| cc_library( |
| name = "aten", |
| srcs = [ |
| ":ATen_CORE_SRCS", |
| ":ATen_QUANTIZED_SRCS", |
| ":aten_base_cpp", |
| ":aten_base_metal", |
| ":aten_base_vulkan", |
| ":aten_native_cpp", |
| ":aten_native_mkl_cpp", |
| ":aten_native_mkldnn_cpp", |
| ":aten_native_nested_cpp", |
| ":aten_native_quantized_cpp", |
| ":aten_native_sparse_cpp", |
| ":aten_native_transformers_cpp", |
| ":aten_native_xnnpack", |
| ":aten_src_ATen_config", |
| ] + generated_cpu_cpp + aten_ufunc_generated_cpu_sources("aten/src/ATen/{}"), |
| copts = ATEN_COPTS, |
| linkopts = [ |
| "-ldl", |
| ], |
| data = if_cuda( |
| [":libcaffe2_nvrtc.so"], |
| [], |
| ), |
| visibility = ["//visibility:public"], |
| deps = [ |
| ":ATen_CPU", |
| ":aten_headers", |
| ":caffe2_for_aten_headers", |
| ":torch_headers", |
| "@fbgemm", |
| "@ideep", |
| ], |
| alwayslink = True, |
| ) |
| |
| cc_library( |
| name = "aten_nvrtc", |
| srcs = glob([ |
| "aten/src/ATen/cuda/nvrtc_stub/*.cpp", |
| ]), |
| copts = ATEN_COPTS, |
| linkstatic = True, |
| visibility = ["//visibility:public"], |
| deps = [ |
| ":aten_headers", |
| "//c10", |
| "@cuda", |
| "@cuda//:cuda_driver", |
| "@cuda//:nvrtc", |
| ], |
| alwayslink = True, |
| ) |
| |
| cc_binary( |
| name = "libcaffe2_nvrtc.so", |
| linkshared = True, |
| visibility = ["//visibility:public"], |
| deps = [ |
| ":aten_nvrtc", |
| ], |
| ) |
| |
| cc_library( |
| name = "aten_cuda_cpp", |
| srcs = [":aten_cuda_cpp_srcs"] + generated_cuda_cpp, |
| hdrs = [":aten_src_ATen_cuda_config"], |
| copts = ATEN_COPTS, |
| visibility = ["//visibility:public"], |
| deps = [ |
| ":aten", |
| "@cuda", |
| "@cuda//:cusolver", |
| "@cuda//:nvrtc", |
| "@cudnn", |
| "@cudnn_frontend", |
| ], |
| alwayslink = True, |
| ) |
| |
| torch_cuda_half_options = [ |
| "-DCUDA_HAS_FP16=1", |
| "-D__CUDA_NO_HALF_OPERATORS__", |
| "-D__CUDA_NO_HALF_CONVERSIONS__", |
| "-D__CUDA_NO_BFLOAT16_CONVERSIONS__", |
| "-D__CUDA_NO_HALF2_OPERATORS__", |
| ] |
| |
| cu_library( |
| name = "aten_cuda", |
| srcs = [":aten_cu_srcs"], |
| copts = ATEN_COPTS + torch_cuda_half_options, |
| visibility = ["//visibility:public"], |
| deps = [ |
| ":aten_cuda_cpp", |
| "//c10/util:bit_cast", |
| "@cuda//:cublas", |
| "@cuda//:cufft", |
| "@cuda//:cusparse", |
| "@cutlass", |
| ], |
| alwayslink = True, |
| ) |
| |
| # caffe2 |
| CAFFE2_COPTS = COMMON_COPTS + [ |
| "-Dcaffe2_EXPORTS", |
| "-DCAFFE2_USE_CUDNN", |
| "-DCAFFE2_BUILD_MAIN_LIB", |
| "-fvisibility-inlines-hidden", |
| "-fno-math-errno", |
| "-fno-trapping-math", |
| ] |
| |
| filegroup( |
| name = "caffe2_core_srcs", |
| srcs = [ |
| "caffe2/core/common.cc", |
| ], |
| ) |
| |
| filegroup( |
| name = "caffe2_perfkernels_srcs", |
| srcs = [ |
| "caffe2/perfkernels/embedding_lookup_idx.cc", |
| ], |
| ) |
| |
| |
| filegroup( |
| name = "caffe2_serialize_srcs", |
| srcs = [ |
| "caffe2/serialize/file_adapter.cc", |
| "caffe2/serialize/inline_container.cc", |
| "caffe2/serialize/istream_adapter.cc", |
| "caffe2/serialize/read_adapter_interface.cc", |
| ], |
| ) |
| |
| filegroup( |
| name = "caffe2_utils_srcs", |
| srcs = [ |
| "caffe2/utils/proto_wrap.cc", |
| "caffe2/utils/string_utils.cc", |
| "caffe2/utils/threadpool/ThreadPool.cc", |
| "caffe2/utils/threadpool/pthreadpool.cc", |
| "caffe2/utils/threadpool/pthreadpool_impl.cc", |
| "caffe2/utils/threadpool/thread_pool_guard.cpp", |
| ], |
| ) |
| |
| # To achieve finer granularity and make debug easier, caffe2 is split into three libraries: |
| # ATen, caffe2 and caffe2_for_aten_headers. ATen lib group up source codes under |
| # aten/ directory and caffe2 contains most files under `caffe2/` directory. Since the |
| # ATen lib and the caffe2 lib would depend on each other, `caffe2_for_aten_headers` is splitted |
| # out from `caffe2` to avoid dependency cycle. |
| cc_library( |
| name = "caffe2_for_aten_headers", |
| hdrs = [ |
| "caffe2/core/common.h", |
| "caffe2/perfkernels/common.h", |
| "caffe2/perfkernels/embedding_lookup_idx.h", |
| "caffe2/utils/fixed_divisor.h", |
| ] + glob([ |
| "caffe2/utils/threadpool/*.h", |
| ]), |
| copts = CAFFE2_COPTS, |
| visibility = ["//visibility:public"], |
| deps = [ |
| ":caffe2_core_macros", |
| "//c10", |
| ], |
| ) |
| |
| cc_library( |
| name = "caffe2_headers", |
| hdrs = glob( |
| [ |
| "caffe2/perfkernels/*.h", |
| "caffe2/serialize/*.h", |
| "caffe2/utils/*.h", |
| "caffe2/utils/threadpool/*.h", |
| "modules/**/*.h", |
| ], |
| exclude = [ |
| "caffe2/core/macros.h", |
| ], |
| ) + if_cuda(glob([ |
| "caffe2/**/*.cuh", |
| ])), |
| copts = CAFFE2_COPTS, |
| visibility = ["//visibility:public"], |
| deps = [ |
| ":caffe2_core_macros", |
| ":caffe2_for_aten_headers", |
| ], |
| ) |
| |
| cc_library( |
| name = "caffe2", |
| srcs = [ |
| ":caffe2_core_srcs", |
| ":caffe2_perfkernels_srcs", |
| ":caffe2_serialize_srcs", |
| ":caffe2_utils_srcs", |
| ], |
| copts = CAFFE2_COPTS + ["-mf16c"], |
| linkstatic = 1, |
| visibility = ["//visibility:public"], |
| deps = [ |
| ":caffe2_core_macros", |
| ":caffe2_headers", |
| ":caffe2_perfkernels_avx", |
| ":caffe2_perfkernels_avx2", |
| "//third_party/miniz-2.1.0:miniz", |
| "@com_google_protobuf//:protobuf", |
| "@eigen", |
| "@fbgemm//:fbgemm_src_headers", |
| "@fmt", |
| "@onnx", |
| ] + if_cuda( |
| [ |
| ":aten_cuda", |
| "@tensorpipe//:tensorpipe_cuda", |
| ], |
| [ |
| ":aten", |
| "@tensorpipe//:tensorpipe_cpu", |
| ], |
| ), |
| alwayslink = True, |
| ) |
| |
| cu_library( |
| name = "torch_cuda", |
| srcs = [ |
| "torch/csrc/distributed/c10d/intra_node_comm.cu", |
| "torch/csrc/distributed/c10d/NanCheck.cu", |
| "torch/csrc/distributed/c10d/quantization/quantization_gpu.cu", |
| ], |
| copts = torch_cuda_half_options, |
| visibility = ["//visibility:public"], |
| deps = [ |
| ":aten", |
| "@cuda//:cublas", |
| "@cuda//:curand", |
| "@cudnn", |
| "@eigen", |
| "@tensorpipe//:tensorpipe_cuda", |
| ], |
| alwayslink = True, |
| ) |
| |
| PERF_COPTS = [ |
| "-DHAVE_AVX_CPU_DEFINITION", |
| "-DHAVE_AVX2_CPU_DEFINITION", |
| "-DENABLE_ALIAS=1", |
| "-DHAVE_MALLOC_USABLE_SIZE=1", |
| "-DHAVE_MMAP=1", |
| "-DHAVE_SHM_OPEN=1", |
| "-DHAVE_SHM_UNLINK=1", |
| "-DSLEEF_STATIC_LIBS=1", |
| "-DTH_BALS_MKL", |
| "-D_FILE_OFFSET_BITS=64", |
| "-DUSE_FBGEMM", |
| "-fvisibility-inlines-hidden", |
| "-Wunused-parameter", |
| "-fno-math-errno", |
| "-fno-trapping-math", |
| "-mf16c", |
| ] |
| |
| PERF_HEADERS = glob([ |
| "caffe2/perfkernels/*.h", |
| "caffe2/core/*.h", |
| ]) |
| |
| cc_library( |
| name = "caffe2_perfkernels_avx", |
| srcs = glob([ |
| "caffe2/perfkernels/*_avx.cc", |
| ]), |
| hdrs = PERF_HEADERS, |
| copts = PERF_COPTS + [ |
| "-mavx", |
| ], |
| visibility = ["//visibility:public"], |
| deps = [ |
| ":caffe2_headers", |
| "//c10", |
| ], |
| alwayslink = True, |
| ) |
| |
| cc_library( |
| name = "caffe2_perfkernels_avx2", |
| srcs = glob([ |
| "caffe2/perfkernels/*_avx2.cc", |
| ]), |
| hdrs = PERF_HEADERS, |
| copts = PERF_COPTS + [ |
| "-mavx2", |
| "-mfma", |
| "-mavx", |
| ], |
| visibility = ["//visibility:public"], |
| deps = [ |
| ":caffe2_headers", |
| "//c10", |
| ], |
| alwayslink = True, |
| ) |
| |
| # torch |
| torch_cuda_headers = glob(["torch/csrc/cuda/*.h"]) |
| |
| cc_library( |
| name = "torch_headers", |
| hdrs = if_cuda( |
| torch_cuda_headers, |
| ) + glob( |
| [ |
| "torch/*.h", |
| "torch/csrc/**/*.h", |
| "torch/csrc/distributed/c10d/**/*.hpp", |
| "torch/lib/libshm/*.h", |
| ], |
| exclude = [ |
| "torch/csrc/*/generated/*.h", |
| ] + torch_cuda_headers, |
| ) + GENERATED_AUTOGRAD_CPP + [":version_h"], |
| includes = [ |
| "third_party/kineto/libkineto/include", |
| "torch/csrc", |
| "torch/csrc/api/include", |
| "torch/csrc/distributed", |
| "torch/lib", |
| "torch/lib/libshm", |
| ], |
| visibility = ["//visibility:public"], |
| deps = [ |
| ":aten_headers", |
| ":caffe2_headers", |
| "//c10", |
| "@com_github_google_flatbuffers//:flatbuffers", |
| "@local_config_python//:python_headers", |
| "@onnx", |
| ], |
| alwayslink = True, |
| ) |
| |
| TORCH_COPTS = COMMON_COPTS + [ |
| "-Dtorch_EXPORTS", |
| "-DHAVE_AVX_CPU_DEFINITION", |
| "-DHAVE_AVX2_CPU_DEFINITION", |
| "-DCAFFE2_USE_GLOO", |
| "-fvisibility-inlines-hidden", |
| "-fno-math-errno ", |
| "-fno-trapping-math", |
| "-Wno-error=unused-function", |
| ] |
| |
| torch_sources = { |
| k: "" |
| for k in ( |
| libtorch_core_sources + |
| libtorch_distributed_sources + |
| torch_cpp_srcs + |
| libtorch_extra_sources + |
| jit_core_sources + |
| lazy_tensor_ts_sources + |
| GENERATED_AUTOGRAD_CPP |
| ) |
| }.keys() |
| |
| cc_library( |
| name = "torch", |
| srcs = if_cuda(glob( |
| libtorch_cuda_sources, |
| exclude = [ |
| "torch/csrc/cuda/python_nccl.cpp", |
| "torch/csrc/cuda/nccl.cpp", |
| "torch/csrc/distributed/c10d/intra_node_comm.cu", |
| "torch/csrc/distributed/c10d/CUDASymmetricMemory.cu", |
| "torch/csrc/distributed/c10d/CUDASymmetricMemoryOps.cu", |
| "torch/csrc/distributed/c10d/NanCheck.cu", |
| "torch/csrc/distributed/c10d/quantization/quantization_gpu.cu", |
| ], |
| )) + torch_sources, |
| copts = TORCH_COPTS, |
| linkopts = [ |
| "-lrt", |
| ], |
| defines = [ |
| "CAFFE2_NIGHTLY_VERSION=20200115", |
| ], |
| visibility = ["//visibility:public"], |
| deps = [ |
| ":caffe2", |
| ":torch_headers", |
| "@kineto", |
| "@cpp-httplib", |
| "@nlohmann", |
| ] + if_cuda([ |
| "@cuda//:nvToolsExt", |
| "@cutlass", |
| ":torch_cuda", |
| ]), |
| alwayslink = True, |
| ) |
| |
| cc_library( |
| name = "shm", |
| srcs = glob(["torch/lib/libshm/*.cpp"]), |
| linkopts = [ |
| "-lrt", |
| ], |
| deps = [ |
| ":torch", |
| ], |
| ) |
| |
| cc_library( |
| name = "libtorch_headers", |
| hdrs = glob([ |
| "**/*.h", |
| "**/*.cuh", |
| ]) + [ |
| # We need the filegroup here because the raw list causes Bazel |
| # to see duplicate files. It knows how to deduplicate with the |
| # filegroup. |
| ":cpp_generated_code", |
| ], |
| includes = [ |
| "torch/csrc/api/include", |
| "torch/csrc/distributed", |
| "torch/lib", |
| "torch/lib/libshm", |
| ], |
| visibility = ["//visibility:public"], |
| deps = [ |
| ":torch_headers", |
| ], |
| ) |
| |
| cc_library( |
| name = "torch_python", |
| srcs = libtorch_python_core_sources |
| + if_cuda(libtorch_python_cuda_sources) |
| + if_cuda(libtorch_python_distributed_sources) |
| + GENERATED_AUTOGRAD_PYTHON, |
| hdrs = glob([ |
| "torch/csrc/generic/*.cpp", |
| ]), |
| copts = COMMON_COPTS + if_cuda(["-DUSE_CUDA=1"]), |
| deps = [ |
| ":torch", |
| ":shm", |
| "@pybind11", |
| ], |
| ) |
| |
| pybind_extension( |
| name = "torch/_C", |
| srcs = ["torch/csrc/stub.c"], |
| deps = [ |
| ":torch_python", |
| ":aten_nvrtc", |
| ], |
| ) |
| |
| cc_library( |
| name = "functorch", |
| hdrs = glob([ |
| "functorch/csrc/dim/*.h", |
| ]), |
| srcs = glob([ |
| "functorch/csrc/dim/*.cpp", |
| ]), |
| deps = [ |
| ":aten_nvrtc", |
| ":torch_python", |
| "@pybind11", |
| ], |
| ) |
| |
| pybind_extension( |
| name = "functorch/_C", |
| copts=[ |
| "-DTORCH_EXTENSION_NAME=_C" |
| ], |
| srcs = [ |
| "functorch/csrc/init_dim_only.cpp", |
| ], |
| deps = [ |
| ":functorch", |
| ":torch_python", |
| ":aten_nvrtc", |
| ], |
| ) |
| |
| cc_binary( |
| name = "torch/bin/torch_shm_manager", |
| srcs = [ |
| "torch/lib/libshm/manager.cpp", |
| ], |
| deps = [ |
| ":shm", |
| ], |
| linkstatic = False, |
| ) |
| |
| template_rule( |
| name = "gen_version_py", |
| src = ":torch/version.py.tpl", |
| out = "torch/version.py", |
| substitutions = if_cuda({ |
| # Set default to 11.2. Otherwise Torchvision complains about incompatibility. |
| "{{CUDA_VERSION}}": "11.2", |
| "{{VERSION}}": "2.0.0", |
| }, { |
| "{{CUDA_VERSION}}": "None", |
| "{{VERSION}}": "2.0.0", |
| }), |
| ) |
| |
| py_library( |
| name = "pytorch_py", |
| visibility = ["//visibility:public"], |
| srcs = glob(["torch/**/*.py"], exclude = ["torch/version.py"]) + [":torch/version.py"] + glob(["functorch/**/*.py"]), |
| deps = [ |
| rules.requirement("numpy"), |
| rules.requirement("pyyaml"), |
| rules.requirement("requests"), |
| rules.requirement("setuptools"), |
| rules.requirement("sympy"), |
| rules.requirement("typing_extensions"), |
| "//torchgen", |
| ], |
| data = [ |
| ":torch/_C.so", |
| ":functorch/_C.so", |
| ":torch/bin/torch_shm_manager", |
| ], |
| ) |
| |
| # cpp api tests |
| cc_library( |
| name = "test_support", |
| testonly = True, |
| srcs = [ |
| "test/cpp/api/support.cpp", |
| ], |
| hdrs = [ |
| "test/cpp/api/init_baseline.h", |
| "test/cpp/api/optim_baseline.h", |
| "test/cpp/api/support.h", |
| "test/cpp/common/support.h", |
| ], |
| deps = [ |
| ":torch", |
| "@com_google_googletest//:gtest_main", |
| ], |
| ) |
| |
| # Torch integration tests rely on a labeled data set from the MNIST database. |
| # http://yann.lecun.com/exdb/mnist/ |
| |
| cpp_api_tests = glob( |
| ["test/cpp/api/*.cpp"], |
| exclude = [ |
| "test/cpp/api/imethod.cpp", |
| "test/cpp/api/integration.cpp", |
| ], |
| ) |
| |
| cc_test( |
| name = "integration_test", |
| size = "medium", |
| srcs = ["test/cpp/api/integration.cpp"], |
| data = [ |
| ":download_mnist", |
| ], |
| tags = [ |
| "gpu-required", |
| ], |
| deps = [ |
| ":test_support", |
| "@com_google_googletest//:gtest_main", |
| ], |
| ) |
| |
| [ |
| cc_test( |
| name = paths.split_extension(paths.basename(filename))[0].replace("-", "_") + "_test", |
| size = "medium", |
| srcs = [filename], |
| deps = [ |
| ":test_support", |
| "@com_google_googletest//:gtest_main", |
| ], |
| ) |
| for filename in cpp_api_tests |
| ] |
| |
| test_suite( |
| name = "api_tests", |
| tests = [ |
| "any_test", |
| "autograd_test", |
| "dataloader_test", |
| "enum_test", |
| "expanding_array_test", |
| "functional_test", |
| "init_test", |
| "integration_test", |
| "jit_test", |
| "memory_test", |
| "misc_test", |
| "module_test", |
| "modulelist_test", |
| "modules_test", |
| "nn_utils_test", |
| "optim_test", |
| "ordered_dict_test", |
| "rnn_test", |
| "sequential_test", |
| "serialize_test", |
| "static_test", |
| "tensor_options_test", |
| "tensor_test", |
| "torch_include_test", |
| ], |
| ) |
| |
| # dist autograd tests |
| cc_test( |
| name = "torch_dist_autograd_test", |
| size = "small", |
| srcs = ["test/cpp/dist_autograd/test_dist_autograd.cpp"], |
| tags = [ |
| "exclusive", |
| "gpu-required", |
| ], |
| deps = [ |
| ":torch", |
| "@com_google_googletest//:gtest_main", |
| ], |
| ) |
| |
| # jit tests |
| # Because these individual unit tests require custom registering, |
| # it is easier to mimic the cmake build by globing together a single test. |
| cc_test( |
| name = "jit_tests", |
| size = "small", |
| srcs = glob( |
| [ |
| "test/cpp/jit/*.cpp", |
| "test/cpp/jit/*.h", |
| "test/cpp/tensorexpr/*.cpp", |
| "test/cpp/tensorexpr/*.h", |
| ], |
| exclude = [ |
| # skip this since <pybind11/embed.h> is not found in OSS build |
| "test/cpp/jit/test_exception.cpp", |
| ], |
| ), |
| linkstatic = True, |
| tags = [ |
| "exclusive", |
| "gpu-required", |
| ], |
| deps = [ |
| ":torch", |
| "@com_google_googletest//:gtest_main", |
| ], |
| ) |
| |
| cc_test( |
| name = "lazy_tests", |
| size = "small", |
| srcs = glob( |
| [ |
| "test/cpp/lazy/*.cpp", |
| "test/cpp/lazy/*.h", |
| ], |
| exclude = [ |
| # skip these since they depend on generated LazyIr.h which isn't available in bazel yet |
| "test/cpp/lazy/test_ir.cpp", |
| "test/cpp/lazy/test_lazy_ops.cpp", |
| "test/cpp/lazy/test_lazy_ops_util.cpp", |
| ], |
| ), |
| linkstatic = True, |
| tags = [ |
| "exclusive", |
| ], |
| deps = [ |
| ":torch", |
| "@com_google_googletest//:gtest_main", |
| ], |
| ) |
| |
| # python api tests |
| |
| py_test( |
| name = "test_bazel", |
| srcs = ["test/_test_bazel.py"], |
| main = "test/_test_bazel.py", |
| deps = [":pytorch_py"], |
| ) |
| |
| # all tests |
| test_suite( |
| name = "all_tests", |
| tests = [ |
| "api_tests", |
| "jit_tests", |
| "torch_dist_autograd_test", |
| "//c10/test:tests", |
| ], |
| ) |
| |
| # An internal genrule that we are converging with refers to these file |
| # as if they are from this package, so we alias them for |
| # compatibility. |
| |
| [ |
| alias( |
| name = paths.basename(path), |
| actual = path, |
| ) |
| for path in [ |
| "aten/src/ATen/templates/DispatchKeyNativeFunctions.cpp", |
| "aten/src/ATen/templates/DispatchKeyNativeFunctions.h", |
| "aten/src/ATen/templates/LazyIr.h", |
| "aten/src/ATen/templates/LazyNonNativeIr.h", |
| "aten/src/ATen/templates/RegisterDispatchKey.cpp", |
| "aten/src/ATen/templates/RegisterDispatchDefinitions.ini", |
| "aten/src/ATen/native/native_functions.yaml", |
| "aten/src/ATen/native/tags.yaml", |
| "aten/src/ATen/native/ts_native_functions.yaml", |
| "torch/csrc/lazy/core/shape_inference.h", |
| "torch/csrc/lazy/ts_backend/ts_native_functions.cpp", |
| ] |
| ] |
| |
| genrule( |
| name = "download_mnist", |
| srcs = ["//:tools/download_mnist.py"], |
| outs = [ |
| "mnist/train-images-idx3-ubyte", |
| "mnist/train-labels-idx1-ubyte", |
| "mnist/t10k-images-idx3-ubyte", |
| "mnist/t10k-labels-idx1-ubyte", |
| ], |
| cmd = "python3 tools/download_mnist.py -d $(RULEDIR)/mnist", |
| ) |
