refactor(//third_party/cublas): Adding paths for x85 in the cuBLAS BUILD

to satisfy Bazel

Adding the paths for cuBLAS in cuda to the BUILD for cuBLAS just to let
bazel track them. We will still use the cuBLAS from the CUDA build on
x86 since it will let us version better.

Signed-off-by: Naren Dasan <[email protected]>
Signed-off-by: Naren Dasan <[email protected]>

Author: narendasan

WORKSPACE

@@ -38,6 +38,11 @@ new_local_repository(
     path = "/usr/local/cuda-11.1/",
 )
 
+new_local_repository(
+    name = "cublas",
+    build_file = "@//third_party/cublas:BUILD",
+    path = "/usr",
+)
 #############################################################################################################
 # Tarballs and fetched dependencies (default - use in cases when building from precompiled bin and tarballs)
 #############################################################################################################

third_party/cublas/BUILD

@@ -1,5 +1,7 @@
 package(default_visibility = ["//visibility:public"])
 
+# NOTE: This BUILD file is only really targeted at aarch64, the rest of the configuration is just to satisfy bazel, x86 uses the cublas source from the CUDA build file since it will be versioned with CUDA.
+
 config_setting(
     name = "aarch64_linux",
     constraint_values = [
@@ -8,9 +10,19 @@ config_setting(
     ],
 )
 
+config_setting(
+    name = "windows",
+    constraint_values = [
+        "@platforms//os:windows",
+    ],
+)
+
 cc_library(
     name = "cublas_headers",
-    hdrs = ["include/cublas.h"] + glob(["include/cublas+.h"]),
+    hdrs = select({
+        ":aarch64_linux": ["include/cublas.h"] + glob(["usr/include/cublas+.h"]),
+        "//conditions:default": ["local/cuda/include/cublas.h"] + glob(["usr/cuda/include/cublas+.h"]),
+    }),
     includes = ["include/"],
     visibility = ["//visibility:private"],
 )
@@ -19,7 +31,8 @@ cc_import(
     name = "cublas_lib",
     shared_library = select({
         ":aarch64_linux": "lib/aarch64-linux-gnu/libcublas.so",
-        "//conditions:default": "lib/x86_64-linux-gnu/libcublas.so",
+        ":windows": "lib/x64/cublas.lib",
+        "//conditions:default": "local/cuda/targets/x86_64-linux/lib/libcublas.so",
     }),
     visibility = ["//visibility:private"],
 )
@@ -28,7 +41,7 @@ cc_import(
     name = "cublas_lt_lib",
     shared_library = select({
         ":aarch64_linux": "lib/aarch64-linux-gnu/libcublasLt.so",
-        "//conditions:default": "lib/x86_64-linux-gnu/libcublasLt.so",
+        "//conditions:default": "local/cuda/targets/x86_64-linux/lib/libcublasLt.so",
     }),
     visibility = ["//visibility:private"],
 )

third_party/tensorrt/archive/BUILD

@@ -1,5 +1,20 @@
 package(default_visibility = ["//visibility:public"])
 
+config_setting(
+    name = "aarch64_linux",
+    constraint_values = [
+        "@platforms//cpu:aarch64",
+        "@platforms//os:linux",
+    ],
+)
+
+config_setting(
+    name = "windows",
+    constraint_values = [
+        "@platforms//os:windows",
+    ],
+)
+
 cc_library(
     name = "nvinfer_headers",
     hdrs = glob(
@@ -27,10 +42,13 @@ cc_library(
     deps = [
         "nvinfer_headers",
         "nvinfer_lib",
-        "@cuda//:cublas",
         "@cuda//:cudart",
         "@cudnn",
-    ],
+    ] + select({
+        ":aarch64_linux": ["@cublas//:cublas"],
+        ":windows": ["@cuda//:cublas"],
+        "//conditions:default": ["@cuda//:cublas"],
+    }),
 )
 
 ####################################################################################
@@ -165,8 +183,11 @@ cc_library(
         "nvinfer",
         "nvinferplugin_headers",
         "nvinferplugin_lib",
-        "@cublas//:cublas",
         "@cuda//:cudart",
         "@cudnn",
-    ],
+    ] + select({
+        ":aarch64_linux": ["@cublas//:cublas"],
+        ":windows": ["@cuda//:cublas"],
+        "//conditions:default": ["@cuda//:cublas"],
+    }),
 )

third_party/tensorrt/local/BUILD

@@ -86,10 +86,13 @@ cc_library(
     deps = [
         "nvinfer_headers",
         "nvinfer_lib",
-        "@cuda//:cublas",
         "@cuda//:cudart",
         "@cudnn",
-    ],
+    ] + select({
+        ":aarch64_linux": ["@cublas//:cublas"],
+        ":windows": ["@cuda//:cublas"],
+        "//conditions:default": ["@cuda//:cublas"],
+    }),
 )
 
 ####################################################################################
@@ -305,9 +308,12 @@ cc_library(
         "nvinfer",
         "@cublas//:cublas",
         "@cuda//:cudart",
-        "@cuda//:cublas",
         "@cudnn",
-    ],
+    ] + select({
+        ":aarch64_linux": ["@cublas//:cublas"],
+        ":windows": ["@cuda//:cublas"],
+        "//conditions:default": ["@cuda//:cublas"],
+    }),
     alwayslink = True,
     copts = [
         "-pthread"