Merge pull request #3 from narendasan/anuragd/aarch64-jetpack-support

Small changes to the aarch64 PR

Author: Anurag Dixit

WORKSPACE

@@ -99,16 +99,12 @@ http_archive(
 #new_local_repository(
 #    name = "libtorch",
 #    path = "/usr/local/lib/python3.6/dist-packages/torch",
-#    # Use below for aarch64
-#    #path = "/home/nvidia/.local/lib/python3.6/site-packages/torch",
 #    build_file = "third_party/libtorch/BUILD"
 #)
 
 #new_local_repository(
 #    name = "libtorch_pre_cxx11_abi",
 #    path = "/usr/local/lib/python3.6/dist-packages/torch",
-#    # Use below for aarch64
-#    #path = "/home/nvidia/.local/lib/python3.6/site-packages/torch",
 #    build_file = "third_party/libtorch/BUILD"
 #)
 

tests/core/converters/BUILD

@@ -86,5 +86,3 @@ test_suite(
     ":test_stack"
   ]
 )
-
-

tests/util/BUILD

@@ -23,6 +23,7 @@ cc_library(
         "//core/conversion",
         "//core/util:prelude",
         "//cpp/api:trtorch",
+	"@tensorrt//:nvinfer"
     ] + select({
         ":use_pre_cxx11_abi":  [
             "@libtorch_pre_cxx11_abi//:libtorch",

third_party/cuda/BUILD

@@ -12,7 +12,7 @@ cc_library(
     name = "cudart",
     srcs = select({
         ":aarch64_linux": [
-            "targets/aarch64-linux-gnu/lib/libcudart.so",
+            "targets/aarch64-linux/lib/libcudart.so",
         ],
         "//conditions:default": [
             "targets/x86_64-linux/lib/libcudart.so",
@@ -30,7 +30,7 @@ cc_library(
     name = "nvToolsExt",
     srcs = select({
         ":aarch64_linux": [
-            "targets/aarch64-linux-gnu/lib/libnvToolsExt.so.1",
+            "targets/aarch64-linux/lib/libnvToolsExt.so.1",
         ],
         "//conditions:default": [
             "targets/x86_64-linux/lib/libnvToolsExt.so.1",
@@ -42,7 +42,7 @@ cc_library(
     name = "cuda",
     srcs = select({
         ":aarch64_linux": glob([
-            "targets/aarch64-linux-gnu/lib/**/lib*.so",
+            "targets/aarch64-linux/lib/**/lib*.so",
         ]),
         "//conditions:default": glob([
             "targets/x86_64-linux/lib/**/lib*.so",

toolchains/BUILD

@@ -0,0 +1,9 @@
+package(default_visibility = ["//visibility:public"])
+
+platform(
+    name = "aarch64_linux",
+    constraint_values = [
+        "@platforms//os:linux",
+        "@platforms//cpu:aarch64",
+    ],
+)