Skip to content

Fix Torch-TensorRT Bazel build on JetPack #4307

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Closed
Mike-Sagiri wants to merge 1 commit into from
+85 −12
Closed

Fix Torch-TensorRT Bazel build on JetPack #4307

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
34 changes: 24 additions & 10 deletions MODULE.bazel
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -72,7 +72,7 @@ new_local_repository(
new_local_repository(
name = "cuda_l4t",
build_file = "@//third_party/cuda:BUILD",
path = "/usr/local/cuda-12.6",
path = "/usr/local/cuda-13.0",
)

# for windows
Expand Down Expand Up @@ -118,6 +118,19 @@ local_torch(name = "libtorch")
# build_file = "@//third_party/tensorrt/local:BUILD"
#)

#new_local_repository(
# name = "tensorrt_l4t",
# path = "/usr/",
# build_file = "@//third_party/tensorrt/local:BUILD"
#)

#new_local_repository(
# name = "torch_l4t",
# please modify this path according to your local python wheel path
# path = "/usr/local/lib/python3.11/dist-packages/torch",
# build_file = "third_party/libtorch/BUILD"
#)

#############################################################################################################
# Tarballs and fetched dependencies (default - use in cases when building from precompiled bin and tarballs)
#############################################################################################################
Expand Down Expand Up @@ -151,13 +164,14 @@ http_archive(
urls = ["https://download.pytorch.org/libtorch/nightly/cu130/libtorch-win-shared-with-deps-latest.zip"],
)

http_archive(
name = "torch_l4t",
build_file = "@//third_party/libtorch:BUILD",
strip_prefix = "torch",
type = "zip",
urls = ["https://pypi.jetson-ai-lab.io/jp6/cu126/+f/62a/1beee9f2f1470/torch-2.8.0-cp310-cp310-linux_aarch64.whl"],
)
# No wheel for torch2.12
#http_archive(
# name = "torch_l4t",
# build_file = "@//third_party/libtorch:BUILD",
# strip_prefix = "torch",
# type = "zip",
# urls = ["https://pypi.jetson-ai-lab.io/jp6/cu126/+f/62a/1beee9f2f1470/torch-2.8.0-cp310-cp310-linux_aarch64.whl"],
#)

# Download these tarballs manually from the NVIDIA website
# Either place them in the distdir directory in third_party and use the --distdir flag
Expand Down Expand Up @@ -194,9 +208,9 @@ http_archive(
http_archive(
name = "tensorrt_l4t",
build_file = "@//third_party/tensorrt/archive:BUILD",
strip_prefix = "TensorRT-10.3.0.26",
strip_prefix = "TensorRT-10.16.1.11",
urls = [
"https://developer.nvidia.com/downloads/compute/machine-learning/tensorrt/10.3.0/tars/TensorRT-10.3.0.26.l4t.aarch64-gnu.cuda-12.6.tar.gz",
"https://developer.nvidia.com/downloads/compute/machine-learning/tensorrt/10.16.1/tars/TensorRT-10.16.1.11.Linux.aarch64-gnu.cuda-13.2.tar.gz",
],
)

Expand Down
4 changes: 2 additions & 2 deletions setup.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -820,11 +820,11 @@ def run(self):


def get_jetpack_requirements(base_requirements):
requirements = base_requirements + ["numpy<2.0.0"]
requirements = base_requirements + ["numpy"]
if IS_DLFW_CI:
return requirements
else:
return requirements + ["torch>=2.8.0,<2.9.0", "tensorrt>=10.3.0,<10.4.0"]
return requirements + ["torch>=2.13.0.dev,<2.14.0", "tensorrt>=10.16.1,<10.17.0"]


def get_sbsa_requirements(base_requirements):
Expand Down
59 changes: 59 additions & 0 deletions third_party/tensorrt/local/BUILD
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -27,9 +27,30 @@ config_setting(
],
)

config_setting(
name = "jetpack",
constraint_values = [
"@platforms//cpu:aarch64",
"@platforms//os:linux",
],
flag_values = {
"@//toolchains/dep_collection:compute_libs": "jetpack",
},
)

cc_library(
name = "nvinfer_headers",
hdrs = select({
":jetpack": glob(
[
"include/NvInfer*.h",
],
allow_empty = True,
exclude = [
"include/NvInferPlugin.h",
"include/NvInferPluginUtils.h",
],
),
":aarch64_linux": glob(
[
"include/aarch64-linux-gnu/NvInfer*.h",
Expand Down Expand Up @@ -77,6 +98,7 @@ cc_library(
),
}),
includes = select({
":jetpack": ["include/"],
":aarch64_linux": ["include/aarch64-linux-gnu"],
":ci_rhel_x86_64_linux": ["include/"],
":windows": ["include/"],
Expand All @@ -91,6 +113,7 @@ cc_library(
cc_import(
name = "nvinfer_static_lib",
static_library = select({
":jetpack": "lib/libnvinfer_static.a",
":aarch64_linux": "lib/aarch64-linux-gnu/libnvinfer_static.a",
":ci_rhel_x86_64_linux": "lib/libnvinfer_static.a",
":windows": "lib/nvinfer_10.lib",
Expand All @@ -105,6 +128,7 @@ cc_import(
cc_import(
name = "nvinfer_lib",
shared_library = select({
":jetpack": "lib/libnvinfer.so",
":aarch64_linux": "lib/aarch64-linux-gnu/libnvinfer.so",
":ci_rhel_x86_64_linux": "lib/libnvinfer.so",
":windows": "bin/nvinfer_10.dll",
Expand All @@ -123,6 +147,7 @@ cc_library(
"nvinfer_headers",
"nvinfer_lib",
] + select({
":jetpack": ["@cuda_l4t//:cudart"],
":windows": [
"nvinfer_static_lib",
"@cuda_win//:cudart",
Expand All @@ -136,6 +161,7 @@ cc_library(
cc_import(
name = "nvparsers_lib",
shared_library = select({
":jetpack": "lib/libnvparsers.so",
":aarch64_linux": "lib/aarch64-linux-gnu/libnvparsers.so",
":ci_rhel_x86_64_linux": "lib/libnvparsers.so",
":windows": "lib/nvparsers.dll",
Expand All @@ -150,6 +176,13 @@ cc_import(
cc_library(
name = "nvparsers_headers",
hdrs = select({
":jetpack": [
"include/NvCaffeParser.h",
"include/NvOnnxConfig.h",
"include/NvOnnxParser.h",
"include/NvOnnxParserRuntime.h",
"include/NvUffParser.h",
],
":aarch64_linux": [
"include/aarch64-linux-gnu/NvCaffeParser.h",
"include/aarch64-linux-gnu/NvOnnxConfig.h",
Expand Down Expand Up @@ -185,6 +218,7 @@ cc_library(
]),
}),
includes = select({
":jetpack": ["include/"],
":aarch64_linux": ["include/aarch64-linux-gnu"],
":ci_rhel_x86_64_linux": ["include/"],
":windows": ["include/"],
Expand All @@ -211,6 +245,7 @@ cc_library(
cc_import(
name = "nvonnxparser_lib",
shared_library = select({
":jetpack": "lib/libnvonnxparser.so",
":aarch64_linux": "lib/aarch64-linux-gnu/libnvonnxparser.so",
":ci_rhel_x86_64_linux": "lib/libnvonnxparser.so",
":windows": "lib/nvonnxparser.dll",
Expand All @@ -225,6 +260,11 @@ cc_import(
cc_library(
name = "nvonnxparser_headers",
hdrs = select({
":jetpack": [
"include/NvOnnxConfig.h",
"include/NvOnnxParser.h",
"include/NvOnnxParserRuntime.h",
],
":aarch64_linux": [
"include/aarch64-linux-gnu/NvOnnxConfig.h",
"include/aarch64-linux-gnu/NvOnnxParser.h",
Expand All @@ -250,6 +290,7 @@ cc_library(
]),
}),
includes = select({
":jetpack": ["include/"],
":aarch64_linux": ["include/aarch64-linux-gnu"],
":ci_rhel_x86_64_linux": ["include/"],
":windows": ["include/"],
Expand All @@ -276,6 +317,7 @@ cc_library(
cc_import(
name = "nvonnxparser_runtime_lib",
shared_library = select({
":jetpack": "lib/libnvonnxparser_runtime.so",
":aarch64_linux": "lib/x86_64-linux-gnu/libnvonnxparser_runtime.so",
":ci_rhel_x86_64_linux": "lib/libnvonnxparser_runtime.so",
":windows": "lib/nvonnxparser_runtime.dll",
Expand All @@ -290,6 +332,9 @@ cc_import(
cc_library(
name = "nvonnxparser_runtime_header",
hdrs = select({
":jetpack": [
"include/NvOnnxParserRuntime.h",
],
":aarch64_linux": [
"include/aarch64-linux-gnu/NvOnnxParserRuntime.h",
],
Expand All @@ -305,6 +350,7 @@ cc_library(
]),
}),
includes = select({
":jetpack": ["include/"],
":aarch64_linux": ["include/aarch64-linux-gnu"],
":ci_rhel_x86_64_linux": ["include/"],
":windows": ["include/"],
Expand All @@ -331,6 +377,7 @@ cc_library(
cc_import(
name = "nvcaffeparser_lib",
shared_library = select({
":jetpack": "lib/libnvcaffe_parsers.so",
":aarch64_linux": "lib/aarch64-linux-gnu/libnvcaffe_parsers.so",
":ci_rhel_x86_64_linux": "lib/libnvcaffe_parsers.so",
":windows": "lib/nvcaffe_parsers.dll",
Expand All @@ -345,6 +392,9 @@ cc_import(
cc_library(
name = "nvcaffeparser_headers",
hdrs = select({
":jetpack": [
"include/NvCaffeParser.h",
],
":aarch64_linux": [
"include/aarch64-linux-gnu/NvCaffeParser.h",
],
Expand All @@ -360,6 +410,7 @@ cc_library(
]),
}),
includes = select({
":jetpack": ["include/"],
":aarch64_linux": ["include/aarch64-linux-gnu"],
":ci_rhel_x86_64_linux": ["include/"],
":windows": ["include/"],
Expand All @@ -386,6 +437,7 @@ cc_library(
cc_library(
name = "nvinferplugin",
srcs = select({
":jetpack": ["lib/libnvinfer_plugin.so"],
":aarch64_linux": ["lib/aarch64-linux-gnu/libnvinfer_plugin.so"],
":ci_rhel_x86_64_linux": ["lib/libnvinfer_plugin.so"],
":windows": ["lib/nvinfer_plugin_10.lib"],
Expand All @@ -395,6 +447,10 @@ cc_library(
])],
}),
hdrs = select({
":jetpack": glob(
["include/NvInferPlugin*.h"],
allow_empty = True,
),
":aarch64_linux": glob(
["include/aarch64-linux-gnu/NvInferPlugin*.h"],
allow_empty = True,
Expand All @@ -416,6 +472,7 @@ cc_library(
"-pthread",
],
includes = select({
":jetpack": ["include/"],
":aarch64_linux": ["include/aarch64-linux-gnu/"],
":ci_rhel_x86_64_linux": ["include/"],
":windows": ["include/"],
Expand All @@ -427,13 +484,15 @@ cc_library(
linkopts = [
"-lpthread",
] + select({
":jetpack": ["-Wl,--no-as-needed -ldl -lrt -Wl,--as-needed"],
":aarch64_linux": ["-Wl,--no-as-needed -ldl -lrt -Wl,--as-needed"],
"//conditions:default": [],
}),
deps = [
"nvinfer",
] + select({
":windows": ["@cuda_win//:cudart"],
":jetpack": ["@cuda_l4t//:cudart"],
"//conditions:default": ["@cuda//:cudart"],
}),
alwayslink = True,
Expand Down
Loading