mirror of
https://github.com/saymrwulf/pytorch.git
synced 2026-05-14 20:57:59 +00:00
129 lines
2.6 KiB
Text
129 lines
2.6 KiB
Text
cc_library(
|
|
name = "core",
|
|
srcs = Glob(["core/*.cc", "utils/*.cc"],
|
|
excludes=["*gpu*", "*cudnn*", "*test*"]),
|
|
hdrs=Glob(["core/*.h", "utils/*.h"]),
|
|
deps = [
|
|
"//caffe2/proto:caffe2_proto",
|
|
"//third_party:glog",
|
|
"//third_party:gflags",
|
|
"//third_party:eigen",
|
|
],
|
|
whole_archive = True,
|
|
)
|
|
|
|
cuda_library(
|
|
name="core_gpu_cu",
|
|
srcs=Glob(["core/*.cu", "utils/*.cu"]),
|
|
deps=[
|
|
":core",
|
|
"//third_party:cuda",
|
|
],
|
|
)
|
|
|
|
cc_library(
|
|
name = "core_gpu",
|
|
srcs = Glob(["core/*_gpu.cc", "utils/*_gpu.cc"]),
|
|
deps = [
|
|
":core",
|
|
":core_gpu_cu",
|
|
"//third_party:cnmem",
|
|
"//third_party:cub",
|
|
"//third_party:cuda",
|
|
],
|
|
whole_archive = True,
|
|
)
|
|
|
|
cc_library(
|
|
name = "core_cudnn",
|
|
srcs = Glob(["core/*_cudnn.cc"]),
|
|
deps = [
|
|
":core_gpu",
|
|
"//third_party:cudnn",
|
|
],
|
|
whole_archive = True,
|
|
)
|
|
|
|
cc_test(
|
|
name = "core_test",
|
|
srcs = Glob(["core/*_test.cc", "utils/*_test.cc"],
|
|
excludes=["*gpu_test*", "*cudnn_test*"]),
|
|
deps = [
|
|
":core",
|
|
"//caffe2/operators:core_ops",
|
|
"//third_party:gtest",
|
|
"//caffe2/test:caffe2_gtest_main",
|
|
],
|
|
)
|
|
|
|
cc_test(
|
|
name = "core_gpu_test",
|
|
srcs = Glob([
|
|
"core/*_gpu_test.cc",
|
|
"core/*_cudnn_test.cc",
|
|
"utils/*_gpu_test.cc",
|
|
]),
|
|
deps = [
|
|
":core_gpu",
|
|
":core_cudnn",
|
|
"//third_party:gtest",
|
|
"//caffe2/test:caffe2_gtest_main",
|
|
],
|
|
test_flags = [
|
|
"--caffe2_cuda_memory_pool cnmem",
|
|
"--caffe2_cuda_memory_pool cub",
|
|
],
|
|
)
|
|
|
|
cc_library(
|
|
name = "all_available_cpu_ops",
|
|
srcs = [],
|
|
# core_ops should always be available. Otherwise, we have a problem.
|
|
deps = [
|
|
"//caffe2/db:db",
|
|
"//caffe2/operators:core_ops",
|
|
],
|
|
optional_deps = [
|
|
"//caffe2/image:image_ops",
|
|
"//caffe2/mpi:mpi_ops",
|
|
"//caffe2/queue:queue_ops",
|
|
"//caffe2/sgd:sgd_ops",
|
|
],
|
|
)
|
|
|
|
cc_library(
|
|
name = "all_available_ops",
|
|
srcs = [],
|
|
deps = [
|
|
":all_available_cpu_ops",
|
|
],
|
|
optional_deps = [
|
|
"//caffe2/operators:core_ops_gpu",
|
|
"//caffe2/operators:core_ops_cudnn",
|
|
"//caffe2/contrib/nccl:nccl_ops",
|
|
"//caffe2/cuda_rtc:rtc_ops",
|
|
"//caffe2/db:db_gpu",
|
|
"//caffe2/image:image_ops_gpu",
|
|
"//caffe2/mpi:mpi_ops_gpu",
|
|
"//caffe2/queue:queue_ops_gpu",
|
|
"//caffe2/sgd:sgd_ops_gpu",
|
|
],
|
|
)
|
|
|
|
cc_library(
|
|
name = "caffe2_core",
|
|
srcs = [],
|
|
deps = [
|
|
"//caffe2/operators:core_ops",
|
|
],
|
|
build_shared = True,
|
|
)
|
|
|
|
cc_library(
|
|
name = "caffe2",
|
|
srcs = [],
|
|
deps = [
|
|
"all_available_ops",
|
|
],
|
|
build_shared=True,
|
|
)
|