aarch64 python build

Author: narendasan

.github/scripts/filter-matrix.py

@@ -26,8 +26,12 @@ def main(args: List[str]) -> None:
     includes = matrix_dict["include"]
     filtered_includes = []
     for item in includes:
-        if all([item["python_version"] not in disabled_python_versions,
-                item["desired_cuda"] != cpu_tag]):
+        if all(
+            [
+                item["python_version"] not in disabled_python_versions,
+                item["desired_cuda"] != cpu_tag
+            ]
+        ):
             filtered_includes.append(item)
     filtered_matrix_dict = {}
     filtered_matrix_dict["include"] = filtered_includes

.github/scripts/generate_binary_build_matrix.py

@@ -50,7 +50,7 @@
     "release": "12.4",
 }
 
-CUDA_AARCH64_ARCHES = ["12.8-aarch64", "12.6-aarch64", "11.8-aarch64"]
+CUDA_AARCH64_ARCHES = ["12.8-aarch64", "12.6-aarch64"]
 
 PACKAGE_TYPES = ["wheel", "conda", "libtorch"]
 CXX11_ABI = "cxx11-abi"

MODULE.bazel

@@ -80,13 +80,13 @@ http_archive(
 # It is possible to specify a wheel file to use as the libtorch source by providing the URL below and
 # using the build flag `--//toolchains/dep_src:torch="whl"`
 
-# http_archive(
-#     name = "torch_whl",
-#     build_file = "@//third_party/libtorch:BUILD",
-#     strip_prefix = "torch",
-#     type = "zip",
-#     urls = ["https://pypi.jetson-ai-lab.dev/jp6/cu126/+f/52c/2cbdd62b78f32/torch-2.7.0-cp310-cp310-linux_aarch64.whl#sha256=52c2cbdd62b78f32c51fa178212e4721241a2ba9e0c4d7d690dd808bd890d51b"],
-# )
+http_archive(
+    name = "torch_whl",
+    build_file = "@//third_party/libtorch:BUILD",
+    strip_prefix = "torch",
+    type = "zip",
+    urls = ["https://download.pytorch.org/whl/nightly/cu128/torch-2.8.0.dev20250415%2Bcu128-cp310-cp310-manylinux_2_28_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

setup.py

@@ -188,6 +188,10 @@ def build_libtorchtrt_cxx11_abi(
     if IS_JETPACK:
         cmd.append("--config=jetpack")
 
+    if IS_SBSA:
+        if CI_BUILD:
+            cmd.append("--//toolchains/dep_src:torch=whl")
+
     if CI_BUILD:
         cmd.append("--platforms=//toolchains:ci_rhel_x86_64_linux")
         print("CI based build")
@@ -456,14 +460,6 @@ def run(self):
 package_data = {}
 
 if not (PY_ONLY or NO_TS):
-    tensorrt_windows_external_dir = (
-        lambda: subprocess.check_output(
-            [BAZEL_EXE, "query", "@tensorrt_win//:nvinfer", "--output", "location"]
-        )
-        .decode("ascii")
-        .strip()
-        .split("/BUILD.bazel")[0]
-    )
 
     tensorrt_x86_64_external_dir = (
         lambda: subprocess.check_output(
@@ -493,15 +489,15 @@ def run(self):
     )
 
     if IS_SBSA:
-        tensorrt_linux_external_dir = tensorrt_sbsa_external_dir()
+        tensorrt_linux_external_dir = tensorrt_sbsa_external_dir
     elif IS_JETPACK:
-        tensorrt_linux_external_dir = tensorrt_jetpack_external_dir()
+        tensorrt_linux_external_dir = tensorrt_jetpack_external_dir
     else:
-        tensorrt_linux_external_dir = tensorrt_x86_64_external_dir()
+        tensorrt_linux_external_dir = tensorrt_x86_64_external_dir
 
     tensorrt_windows_external_dir = (
         lambda: subprocess.check_output(
-            [BAZEL_EXE, "query", "@tensorrt_windows//:nvinfer", "--output", "location"]
+            [BAZEL_EXE, "query", "@tensorrt_win//:nvinfer", "--output", "location"]
         )
         .decode("ascii")
         .strip()