(//docker): Updated cublas changes for arm64

Signed-off-by: Anurag Dixit <[email protected]>

Author: Anurag Dixit

third_party/cublas/BUILD

@@ -16,7 +16,6 @@ config_setting(
     constraint_values = [
         "@platforms//cpu:aarch64",
         "@platforms//os:linux",
-        "@//toolchains/jetpack:4.6"
     ]
 )
 
@@ -67,4 +66,4 @@ cc_library(
         "cublas_lib",
         "cublas_lt_lib",
     ],
-)
+)

third_party/cuda/BUILD

@@ -1,13 +1,5 @@
 package(default_visibility = ["//visibility:public"])
 
-config_setting(
-    name = "aarch64_linux",
-    constraint_values = [
-        "@platforms//cpu:aarch64",
-        "@platforms//os:linux",
-    ],
-)
-
 config_setting(
     name = "windows",
     constraint_values = [
@@ -18,9 +10,6 @@ config_setting(
 cc_library(
     name = "cudart",
     srcs = select({
-        ":aarch64_linux": [
-            "targets/aarch64-linux/lib/libcudart.so",
-        ],
         ":windows": [
             "lib/x64/cudart.lib",
         ],
@@ -39,24 +28,18 @@ cc_library(
 cc_library(
     name = "nvToolsExt",
     srcs = select({
-        ":aarch64_linux": [
-            "targets/aarch64-linux/lib/libnvToolsExt.so.1",
-        ],
         ":windows": [
             "bin/nvToolsExt64_1.dll",
         ],
         "//conditions:default": [
-            "targets/x86_64-linux/lib/libnvToolsExt.so.1",
+            "ltargets/x86_64-linux/lib/libnvToolsExt.so.1",
         ],
     }),
 )
 
 cc_library(
     name = "cuda",
     srcs = select({
-        ":aarch64_linux": glob([
-            "targets/aarch64-linux/lib/**/lib*.so",
-        ]),
         ":windows": [
             "bin/*.dll",
         ],
@@ -76,14 +59,11 @@ cc_library(
 cc_library(
     name = "cublas",
     srcs = select({
-        ":aarch64_linux": glob([
-            "lib/**/*libcublas.so",
-        ]),
         ":windows": glob([
             "lib/x64/cublas.lib",
         ]),
         "//conditions:default": glob([
-            "lib/**/*libcublas.so",
+            "targets/x86_64-linux/lib/**/*libcublas.so",
         ]),
     }),
     hdrs = glob([

third_party/tensorrt/archive/BUILD

@@ -45,7 +45,6 @@ cc_library(
         "@cuda//:cudart",
         "@cudnn",
     ] + select({
-        ":aarch64_linux": ["@cublas//:cublas"],
         ":windows": ["@cuda//:cublas"],
         "//conditions:default": ["@cuda//:cublas"],
     }),
@@ -186,7 +185,6 @@ cc_library(
         "@cuda//:cudart",
         "@cudnn",
     ] + select({
-        ":aarch64_linux": ["@cublas//:cublas"],
         ":windows": ["@cuda//:cublas"],
         "//conditions:default": ["@cuda//:cublas"],
     }),

third_party/tensorrt/local/BUILD

@@ -89,7 +89,6 @@ cc_library(
         "@cuda//:cudart",
         "@cudnn",
     ] + select({
-        ":aarch64_linux": ["@cublas//:cublas"],
         ":windows": ["@cuda//:cublas"],
         "//conditions:default": ["@cuda//:cublas"],
     }),
@@ -309,7 +308,6 @@ cc_library(
         "@cuda//:cudart",
         "@cudnn",
     ] + select({
-        ":aarch64_linux": ["@cublas//:cublas"],
         ":windows": ["@cuda//:cublas"],
         "//conditions:default": ["@cuda//:cublas"],
     }),