BUILD.pytorch

config_setting(
    name = "using_cuda",
    values = {"define": "using_cuda=true"},
)

config_setting(
    name = "using_rocm",
    values = {"define": "using_rocm=true"},
)

cc_library(
    name = "torch",
    srcs = [
        "torch/lib/libtorch.so",
        "torch/lib/libtorch_python.so",
        "torch/lib/libtorch_cpu.so",
        "torch/lib/libshm.so",
        "torch/lib/libc10.so",
    ] + glob([
        "torch/lib/libgomp*.so*",
    ]) + select({
        "//:using_cuda": [
            "torch/lib/libtorch_cuda.so",
            "torch/lib/libc10_cuda.so",
        ] + glob([
            "torch/lib/libcudart*.so*",
            "torch/lib/libnvToolsExt*.so*",
        ]),
        "//:using_rocm": [
            "torch/lib/libtorch_hip.so",
            "torch/lib/libc10_hip.so",
        ] +  glob([
            "torch/lib/libmagma*.so*",
            # "torch/lib/libtinfo.so*",
            # "torch/lib/libsuitesparseconfig.so*",
            # "torch/lib/libcholmod.so*",
            # "torch/lib/libsatlas.so*",
            # "torch/lib/libgfortran.so*",
            # "torch/lib/libquadmath.so*",
            # "torch/lib/libgomp*.so*",
            # "torch/lib/libhip*.so*",
            # "torch/lib/lib*amd*.so*",
        ]),
        "//conditions:default": [],
    }),
    hdrs = glob([
        "torch/include/*.h",
        "torch/include/**/*.h",
        "torch/include/**/*.cuh",
        "torch/include/**/*.hpp",
    ]),
    deps = [
        "@local_config_python//:python_headers",
        "@local_config_python//:python_lib",
    ],
    strip_include_prefix = "torch/include",
    visibility = ["//visibility:public"],
)

cc_library(
    name = "torch_api",
    hdrs = glob([
        "torch/include/torch/csrc/api/include/**/*.h",
    ]),
    deps = [
        "@local_config_python//:python_headers",
    ],
    strip_include_prefix = "torch/include/torch/csrc/api/include",
    visibility = ["//visibility:public"],
)

cc_library(
    name = "torch_libs",
    srcs = glob([
        "torch.libs/libarm_compute*.so",
        "torch.libs/libgfortran*.so*",
        "torch.libs/libgomp*.so*",
        "torch.libs/libopenblasp*.so",
    ]),
    deps = [
        "@local_config_python//:python_headers",
    ],
    visibility = ["//visibility:public"],
)