fix(//plugins): Readding cuBLAS BUILD to allow linking of libnvinfer_…

…plugin on Jetson Signed-off-by: Naren Dasan <naren@narendasan.com>
pytorch · May 10, 2021 · a8008f4 · a8008f4
1 parent 0a3258d
commit a8008f4
Show file tree

Hide file tree

Showing 3 changed files with 66 additions and 21 deletions.
diff --git a/third_party/cublas/BUILD b/third_party/cublas/BUILD
@@ -0,0 +1,44 @@
+package(default_visibility = ["//visibility:public"])
+
+config_setting(
+    name = "aarch64_linux",
+    constraint_values = [
+        "@platforms//cpu:aarch64",
+        "@platforms//os:linux",
+    ],
+)
+
+cc_library(
+    name = "cublas_headers",
+    hdrs = ["include/cublas.h"] + glob(["include/cublas+.h"]),
+    includes = ["include/"],
+    visibility = ["//visibility:private"],
+)
+
+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",
+    }),
+    visibility = ["//visibility:private"],
+)
+
+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",
+    }),
+    visibility = ["//visibility:private"],
+)
+
+cc_library(
+    name = "cublas",
+    visibility = ["//visibility:public"],
+    deps = [
+        "cublas_headers",
+        "cublas_lib",
+        "cublas_lt_lib",
+    ],
+)
diff --git a/third_party/tensorrt/archive/BUILD b/third_party/tensorrt/archive/BUILD
@@ -165,5 +165,8 @@ cc_library(
         "nvinfer",
         "nvinferplugin_headers",
         "nvinferplugin_lib",
+        "@cublas//:cublas",
+        "@cuda//:cudart",
+        "@cudnn",
     ],
 )
diff --git a/third_party/tensorrt/local/BUILD b/third_party/tensorrt/local/BUILD
@@ -274,47 +274,45 @@ 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",
+        "@cublas//:cublas",
+        "@cuda//:cudart",
+        "@cuda//:cublas",
+        "@cudnn",
+    ],
+    alwayslink = True,
+    copts = [
+        "-pthread"
     ],
+    linkopts = [
+        "-lpthread",
+    ]
 )