Merge pull request #450 from NVIDIA/cublas_fix

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

@@ -0,0 +1,57 @@
+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 = [
+        "@platforms//cpu:aarch64",
+        "@platforms//os:linux",
+    ],
+)
+
+config_setting(
+    name = "windows",
+    constraint_values = [
+        "@platforms//os:windows",
+    ],
+)
+
+cc_library(
+    name = "cublas_headers",
+    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"],
+)
+
+cc_import(
+    name = "cublas_lib",
+    shared_library = select({
+        ":aarch64_linux": "lib/aarch64-linux-gnu/libcublas.so",
+        ":windows": "lib/x64/cublas.lib",
+        "//conditions:default": "local/cuda/targets/x86_64-linux/lib/libcublas.so",
+    }),
+    visibility = ["//visibility:private"],
+)
+
+cc_import(
+    name = "cublas_lt_lib",
+    shared_library = select({
+        ":aarch64_linux": "lib/aarch64-linux-gnu/libcublasLt.so",
+        "//conditions:default": "local/cuda/targets/x86_64-linux/lib/libcublasLt.so",
+    }),
+    visibility = ["//visibility:private"],
+)
+
+cc_library(
+    name = "cublas",
+    visibility = ["//visibility:public"],
+    deps = [
+        "cublas_headers",
+        "cublas_lib",
+        "cublas_lt_lib",
+    ],
+)

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,5 +183,11 @@ cc_library(
         "nvinfer",
         "nvinferplugin_headers",
         "nvinferplugin_lib",
-    ],
+        "@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"],
+    }),
 )
 
 ####################################################################################
@@ -274,47 +277,47 @@ cc_library(
 
 cc_library(
     name = "nvcaffeparser",
-    visibility = ["//visibility:public"],
     deps = [
         "nvcaffeparser_headers",
         "nvcaffeparser_lib",
         "nvinfer",
     ],
+    visibility = ["//visibility:public"],
 )
 
 ####################################################################################
 
-cc_import(
-    name = "nvinferplugin_lib",
-    shared_library = select({
-        ":aarch64_linux": "lib/x86_64-linux-gnu/libnvinfer_plugin.so",
-        ":windows": "lib/nvinfer_plugin.dll",
-        "//conditions:default": "lib/x86_64-linux-gnu/libnvinfer_plugin.so",
-    }),
-    visibility = ["//visibility:private"],
-)
-
 cc_library(
-    name = "nvinferplugin_headers",
+    name = "nvinferplugin",
     hdrs = select({
         ":aarch64_linux": glob(["include/aarch64-linux-gnu/NvInferPlugin*.h"]),
         ":windows": glob(["include/NvInferPlugin*.h"]),
         "//conditions:default": glob(["include/x86_64-linux-gnu/NvInferPlugin*.h"]),
     }),
+    srcs = select({
+        ":aarch64_linux": ["lib/aarch64-linux-gnu/libnvinfer_plugin.so"],
+        ":windows": ["lib/nvinfer_plugin.dll"],
+        "//conditions:default": ["lib/x86_64-linux-gnu/libnvinfer_plugin.so"],
+    }),
     includes = select({
-        ":aarch64_linux": ["include/aarch64-linux-gnu"],
+        ":aarch64_linux": ["include/aarch64-linux-gnu/"],
         ":windows": ["include/"],
         "//conditions:default": ["include/x86_64-linux-gnu/"],
     }),
-    visibility = ["//visibility:private"],
-)
-
-cc_library(
-    name = "nvinferplugin",
-    visibility = ["//visibility:public"],
     deps = [
         "nvinfer",
-        "nvinferplugin_headers",
-        "nvinferplugin_lib",
+        "@cuda//:cudart",
+        "@cudnn",
+    ] + select({
+        ":aarch64_linux": ["@cublas//:cublas"],
+        ":windows": ["@cuda//:cublas"],
+        "//conditions:default": ["@cuda//:cublas"],
+    }),
+    alwayslink = True,
+    copts = [
+        "-pthread"
     ],
+    linkopts = [
+        "-lpthread",
+    ]
 )