Skip to content

Commit a1e5c11

Browse files
narendasannarendasan
committed
infra: bazel aarch64 colocation
1 parent a4c2093 commit a1e5c11

File tree

40 files changed

+917
-178
lines changed

40 files changed

+917
-178
lines changed

‎MODULE.bazel

Lines changed: 37 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -27,7 +27,7 @@ local_repository = use_repo_rule("@bazel_tools//tools/build_defs/repo:local.bzl"
2727
# External dependency for torch_tensorrt if you already have precompiled binaries.
2828
local_repository(
2929
name = "torch_tensorrt",
30-
path = "/opt/conda/lib/python3.8/site-packages/torch_tensorrt",
30+
path = "/opt/conda/lib/python3.10/site-packages/torch_tensorrt",
3131
)
3232

3333

@@ -37,7 +37,7 @@ new_local_repository = use_repo_rule("@bazel_tools//tools/build_defs/repo:local.
3737
new_local_repository(
3838
name = "cuda",
3939
build_file = "@//third_party/cuda:BUILD",
40-
path = "",
40+
path = "/usr/local/cuda-12.8",
4141
)
4242

4343
new_local_repository(
@@ -57,7 +57,19 @@ http_archive(
5757
name = "libtorch",
5858
build_file = "@//third_party/libtorch:BUILD",
5959
strip_prefix = "libtorch",
60+
<<<<<<< HEAD
6061
urls = ["https://download.pytorch.org/libtorch///libtorch-cxx11-abi-shared-with-deps-latest.zip"],
62+
=======
63+
urls = ["https://download.pytorch.org/libtorch/nightly/cu128/libtorch-cxx11-abi-shared-with-deps-latest.zip"],
64+
)
65+
66+
http_archive(
67+
name = "torch_whl",
68+
build_file = "@//third_party/libtorch:BUILD",
69+
strip_prefix = "torch",
70+
type = "zip",
71+
urls = ["https://download.pytorch.org/whl/nightly/cu128/torch-2.8.0.dev20250414%2Bcu128-cp39-cp39-manylinux_2_28_aarch64.whl"],
72+
>>>>>>> 75173f897 (infra: bazel aarch64 colocation)
6173
)
6274

6375
# Download these tarballs manually from the NVIDIA website
@@ -73,6 +85,15 @@ http_archive(
7385
],
7486
)
7587

88+
http_archive(
89+
name = "tensorrt_sbsa",
90+
build_file = "@//third_party/tensorrt/archive:BUILD",
91+
strip_prefix = "TensorRT-10.9.0.34",
92+
urls = [
93+
"https://developer.nvidia.com/downloads/compute/machine-learning/tensorrt/10.9.0/tars/TensorRT-10.9.0.34.Linux.aarch64-gnu.cuda-12.8.tar.gz",
94+
],
95+
)
96+
7697
http_archive(
7798
name = "tensorrt_win",
7899
build_file = "@//third_party/tensorrt/archive:BUILD",
@@ -94,6 +115,7 @@ http_archive(
94115
# x86_64 python distribution. If using NVIDIA's version just point to the root of the package
95116
# for both versions here and do not use --config=pre-cxx11-abi
96117

118+
<<<<<<< HEAD
97119
new_local_repository(
98120
name = "libtorch_win",
99121
path = "",
@@ -105,6 +127,19 @@ new_local_repository(
105127
path = "",
106128
build_file = "third_party/libtorch/BUILD"
107129
)
130+
=======
131+
# new_local_repository(
132+
# name = "libtorch",
133+
# path = "/workspace/tensorrt/.venv/lib/python3.9/site-packages/torch",
134+
# build_file = "third_party/libtorch/BUILD"
135+
# )
136+
137+
# new_local_repository(
138+
# name = "libtorch_pre_cxx11_abi",
139+
# path = "",
140+
# build_file = "third_party/libtorch/BUILD"
141+
# )
142+
>>>>>>> 75173f897 (infra: bazel aarch64 colocation)
108143

109144
#new_local_repository(
110145
# name = "tensorrt",

‎core/BUILD

Lines changed: 34 additions & 9 deletions
Original file line numberDiff line numberDiff line change
@@ -4,16 +4,29 @@ load("@rules_pkg//:pkg.bzl", "pkg_tar")
44
package(default_visibility = ["//visibility:public"])
55

66
config_setting(
7-
name = "use_pre_cxx11_abi",
8-
values = {
9-
"define": "abi=pre_cxx11_abi",
7+
name = "use_torch_whl",
8+
flag_values = {
9+
"//toolchains/build_srcs:torch": "whl"
1010
},
1111
)
1212

1313
config_setting(
14-
name = "python_core",
15-
values = {
16-
"define": "target_lang=python",
14+
name = "sbsa",
15+
constraint_values = [
16+
"@platforms//cpu:aarch64",
17+
],
18+
flag_values = {
19+
"//toolchains/dep_collection:compute_lib_collection": "datacenter"
20+
},
21+
)
22+
23+
config_setting(
24+
name = "jetpack",
25+
constraint_values = [
26+
"@platforms//cpu:aarch64",
27+
],
28+
flag_values = {
29+
"//toolchains/dep_collection:compute_lib_collection": "jetpack"
1730
},
1831
)
1932

@@ -24,6 +37,13 @@ config_setting(
2437
],
2538
)
2639

40+
config_setting(
41+
name = "python_core",
42+
values = {
43+
"define": "target_lang=python",
44+
},
45+
)
46+
2747
cc_library(
2848
name = "core",
2949
srcs = [
@@ -39,9 +59,14 @@ cc_library(
3959
"//core/runtime",
4060
"//core/util/logging",
4161
] + select({
42-
":windows": ["@tensorrt_win//:nvinfer", "@libtorch_win//:libtorch"],
43-
":use_pre_cxx11_abi": ["@tensorrt//:nvinfer", "@libtorch_pre_cxx11_abi//:libtorch"],
44-
"//conditions:default": ["@tensorrt//:nvinfer", "@libtorch"],
62+
":windows": ["@tensorrt_win//:nvinfer"],
63+
":sbsa": ["@tensorrt_sbsa//:nvinfer"],
64+
":jetpack": ["@tensorrt_jetpack//:nvinfer"],
65+
"//conditions:default": ["@tensorrt//:nvinfer"],
66+
}) + select({
67+
":windows": ["@libtorch_win//:libtorch"],
68+
":use_torch_whl": ["@torch_whl//:libtorch"],
69+
"//conditions:default": ["@libtorch"],
4570
}),
4671
alwayslink = True,
4772
)

‎core/conversion/BUILD

Lines changed: 31 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -4,9 +4,29 @@ load("@rules_pkg//:pkg.bzl", "pkg_tar")
44
package(default_visibility = ["//visibility:public"])
55

66
config_setting(
7-
name = "use_pre_cxx11_abi",
8-
values = {
9-
"define": "abi=pre_cxx11_abi",
7+
name = "use_torch_whl",
8+
flag_values = {
9+
"//toolchains/build_srcs:torch": "whl"
10+
},
11+
)
12+
13+
config_setting(
14+
name = "sbsa",
15+
constraint_values = [
16+
"@platforms//cpu:aarch64",
17+
],
18+
flag_values = {
19+
"//toolchains/dep_collection:compute_lib_collection": "datacenter"
20+
},
21+
)
22+
23+
config_setting(
24+
name = "jetpack",
25+
constraint_values = [
26+
"@platforms//cpu:aarch64",
27+
],
28+
flag_values = {
29+
"//toolchains/dep_collection:compute_lib_collection": "jetpack"
1030
},
1131
)
1232

@@ -34,9 +54,14 @@ cc_library(
3454
"//core/ir",
3555
"//core/util:prelude",
3656
] + select({
37-
":windows": ["@tensorrt_win//:nvinfer", "@libtorch_win//:libtorch"],
38-
":use_pre_cxx11_abi": ["@tensorrt//:nvinfer", "@libtorch_pre_cxx11_abi//:libtorch"],
39-
"//conditions:default": ["@tensorrt//:nvinfer", "@libtorch"],
57+
":windows": ["@tensorrt_win//:nvinfer"],
58+
":sbsa": ["@tensorrt_sbsa//:nvinfer"],
59+
":jetpack": ["@tensorrt_jetpack//:nvinfer"],
60+
"//conditions:default": ["@tensorrt//:nvinfer"],
61+
}) + select({
62+
":windows": ["@libtorch_win//:libtorch"],
63+
":use_torch_whl": ["@torch_whl//:libtorch"],
64+
"//conditions:default": ["@libtorch"],
4065
}),
4166
alwayslink = True,
4267
)

‎core/conversion/conversionctx/BUILD

Lines changed: 31 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -4,9 +4,29 @@ load("@rules_pkg//:pkg.bzl", "pkg_tar")
44
package(default_visibility = ["//visibility:public"])
55

66
config_setting(
7-
name = "use_pre_cxx11_abi",
8-
values = {
9-
"define": "abi=pre_cxx11_abi",
7+
name = "use_torch_whl",
8+
flag_values = {
9+
"//toolchains/build_srcs:torch": "whl"
10+
},
11+
)
12+
13+
config_setting(
14+
name = "sbsa",
15+
constraint_values = [
16+
"@platforms//cpu:aarch64",
17+
],
18+
flag_values = {
19+
"//toolchains/dep_collection:compute_lib_collection": "datacenter"
20+
},
21+
)
22+
23+
config_setting(
24+
name = "jetpack",
25+
constraint_values = [
26+
"@platforms//cpu:aarch64",
27+
],
28+
flag_values = {
29+
"//toolchains/dep_collection:compute_lib_collection": "jetpack"
1030
},
1131
)
1232

@@ -29,9 +49,14 @@ cc_library(
2949
"//core/ir",
3050
"//core/util:prelude",
3151
] + select({
32-
":windows": ["@tensorrt_win//:nvinfer", "@libtorch_win//:libtorch"],
33-
":use_pre_cxx11_abi": ["@tensorrt//:nvinfer", "@libtorch_pre_cxx11_abi//:libtorch"],
34-
"//conditions:default": ["@tensorrt//:nvinfer", "@libtorch"],
52+
":windows": ["@tensorrt_win//:nvinfer"],
53+
":sbsa": ["@tensorrt_sbsa//:nvinfer"],
54+
":jetpack": ["@tensorrt_jetpack//:nvinfer"],
55+
"//conditions:default": ["@tensorrt//:nvinfer"],
56+
}) + select({
57+
":windows": ["@libtorch_win//:libtorch"],
58+
":use_torch_whl": ["@torch_whl//:libtorch"],
59+
"//conditions:default": ["@libtorch"],
3560
}),
3661
alwayslink = True,
3762
)

‎core/conversion/converters/BUILD

Lines changed: 47 additions & 12 deletions
Original file line numberDiff line numberDiff line change
@@ -4,9 +4,29 @@ load("@rules_pkg//:pkg.bzl", "pkg_tar")
44
package(default_visibility = ["//visibility:public"])
55

66
config_setting(
7-
name = "use_pre_cxx11_abi",
8-
values = {
9-
"define": "abi=pre_cxx11_abi",
7+
name = "use_torch_whl",
8+
flag_values = {
9+
"//toolchains/build_srcs:torch": "whl"
10+
},
11+
)
12+
13+
config_setting(
14+
name = "sbsa",
15+
constraint_values = [
16+
"@platforms//cpu:aarch64",
17+
],
18+
flag_values = {
19+
"//toolchains/dep_collection:compute_lib_collection": "datacenter"
20+
},
21+
)
22+
23+
config_setting(
24+
name = "jetpack",
25+
constraint_values = [
26+
"@platforms//cpu:aarch64",
27+
],
28+
flag_values = {
29+
"//toolchains/dep_collection:compute_lib_collection": "jetpack"
1030
},
1131
)
1232

@@ -29,9 +49,14 @@ cc_library(
2949
"//core/conversion/conversionctx",
3050
"//core/util:prelude",
3151
] + select({
32-
":windows": ["@tensorrt_win//:nvinfer", "@libtorch_win//:libtorch"],
33-
":use_pre_cxx11_abi": ["@tensorrt//:nvinfer", "@libtorch_pre_cxx11_abi//:libtorch"],
34-
"//conditions:default": ["@tensorrt//:nvinfer", "@libtorch"],
52+
":windows": ["@tensorrt_win//:nvinfer"],
53+
":sbsa": ["@tensorrt_sbsa//:nvinfer"],
54+
":jetpack": ["@tensorrt_jetpack//:nvinfer"],
55+
"//conditions:default": ["@tensorrt//:nvinfer"],
56+
}) + select({
57+
":windows": ["@libtorch_win//:libtorch"],
58+
":use_torch_whl": ["@torch_whl//:libtorch"],
59+
"//conditions:default": ["@libtorch"],
3560
}),
3661
alwayslink = True,
3762
)
@@ -49,9 +74,14 @@ cc_library(
4974
"//core/conversion/conversionctx",
5075
"//core/util:prelude",
5176
] + select({
52-
":windows": ["@tensorrt_win//:nvinfer", "@libtorch_win//:libtorch"],
53-
":use_pre_cxx11_abi": ["@tensorrt//:nvinfer", "@libtorch_pre_cxx11_abi//:libtorch"],
54-
"//conditions:default": ["@tensorrt//:nvinfer", "@libtorch"],
77+
":windows": ["@tensorrt_win//:nvinfer"],
78+
":sbsa": ["@tensorrt_sbsa//:nvinfer"],
79+
":jetpack": ["@tensorrt_jetpack//:nvinfer"],
80+
"//conditions:default": ["@tensorrt//:nvinfer"],
81+
}) + select({
82+
":windows": ["@libtorch_win//:libtorch"],
83+
":use_torch_whl": ["@torch_whl//:libtorch"],
84+
"//conditions:default": ["@libtorch"],
5585
}),
5686
alwayslink = True,
5787
)
@@ -106,9 +136,14 @@ cc_library(
106136
"//core/plugins:torch_tensorrt_plugins",
107137
"//core/util:prelude",
108138
] + select({
109-
":windows": ["@tensorrt_win//:nvinfer", "@libtorch_win//:libtorch"],
110-
":use_pre_cxx11_abi": ["@tensorrt//:nvinfer", "@libtorch_pre_cxx11_abi//:libtorch"],
111-
"//conditions:default": ["@tensorrt//:nvinfer", "@libtorch"],
139+
":windows": ["@tensorrt_win//:nvinfer"],
140+
":sbsa": ["@tensorrt_sbsa//:nvinfer"],
141+
":jetpack": ["@tensorrt_jetpack//:nvinfer"],
142+
"//conditions:default": ["@tensorrt//:nvinfer"],
143+
}) + select({
144+
":windows": ["@libtorch_win//:libtorch"],
145+
":use_torch_whl": ["@torch_whl//:libtorch"],
146+
"//conditions:default": ["@libtorch"],
112147
}),
113148
alwayslink = True,
114149
)

‎core/conversion/evaluators/BUILD

Lines changed: 29 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -4,9 +4,29 @@ load("@rules_pkg//:pkg.bzl", "pkg_tar")
44
package(default_visibility = ["//visibility:public"])
55

66
config_setting(
7-
name = "use_pre_cxx11_abi",
8-
values = {
9-
"define": "abi=pre_cxx11_abi",
7+
name = "use_torch_whl",
8+
flag_values = {
9+
"//toolchains/build_srcs:torch": "whl"
10+
},
11+
)
12+
13+
config_setting(
14+
name = "sbsa",
15+
constraint_values = [
16+
"@platforms//cpu:aarch64",
17+
],
18+
flag_values = {
19+
"//toolchains/dep_collection:compute_lib_collection": "datacenter"
20+
},
21+
)
22+
23+
config_setting(
24+
name = "jetpack",
25+
constraint_values = [
26+
"@platforms//cpu:aarch64",
27+
],
28+
flag_values = {
29+
"//toolchains/dep_collection:compute_lib_collection": "jetpack"
1030
},
1131
)
1232

@@ -35,8 +55,13 @@ cc_library(
3555
"//core/conversion/var",
3656
"//core/util:prelude",
3757
] + select({
58+
":windows": ["@tensorrt_win//:nvinfer"],
59+
":sbsa": ["@tensorrt_sbsa//:nvinfer"],
60+
":jetpack": ["@tensorrt_jetpack//:nvinfer"],
61+
"//conditions:default": ["@tensorrt//:nvinfer"],
62+
}) + select({
3863
":windows": ["@libtorch_win//:libtorch"],
39-
":use_pre_cxx11_abi": ["@libtorch_pre_cxx11_abi//:libtorch"],
64+
":use_torch_whl": ["@torch_whl//:libtorch"],
4065
"//conditions:default": ["@libtorch"],
4166
}),
4267
alwayslink = True,

0 commit comments

Comments
 (0)