Skip to content

Commit a8008f4

Browse files
committed
fix(//plugins): Readding cuBLAS BUILD to allow linking of libnvinfer_plugin on Jetson
Signed-off-by: Naren Dasan <naren@narendasan.com>
1 parent 0a3258d commit a8008f4

File tree

3 files changed

+66
-21
lines changed

3 files changed

+66
-21
lines changed

third_party/cublas/BUILD

+44
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,44 @@
1+
package(default_visibility = ["//visibility:public"])
2+
3+
config_setting(
4+
name = "aarch64_linux",
5+
constraint_values = [
6+
"@platforms//cpu:aarch64",
7+
"@platforms//os:linux",
8+
],
9+
)
10+
11+
cc_library(
12+
name = "cublas_headers",
13+
hdrs = ["include/cublas.h"] + glob(["include/cublas+.h"]),
14+
includes = ["include/"],
15+
visibility = ["//visibility:private"],
16+
)
17+
18+
cc_import(
19+
name = "cublas_lib",
20+
shared_library = select({
21+
":aarch64_linux": "lib/aarch64-linux-gnu/libcublas.so",
22+
"//conditions:default": "lib/x86_64-linux-gnu/libcublas.so",
23+
}),
24+
visibility = ["//visibility:private"],
25+
)
26+
27+
cc_import(
28+
name = "cublas_lt_lib",
29+
shared_library = select({
30+
":aarch64_linux": "lib/aarch64-linux-gnu/libcublasLt.so",
31+
"//conditions:default": "lib/x86_64-linux-gnu/libcublasLt.so",
32+
}),
33+
visibility = ["//visibility:private"],
34+
)
35+
36+
cc_library(
37+
name = "cublas",
38+
visibility = ["//visibility:public"],
39+
deps = [
40+
"cublas_headers",
41+
"cublas_lib",
42+
"cublas_lt_lib",
43+
],
44+
)

third_party/tensorrt/archive/BUILD

+3
Original file line numberDiff line numberDiff line change
@@ -165,5 +165,8 @@ cc_library(
165165
"nvinfer",
166166
"nvinferplugin_headers",
167167
"nvinferplugin_lib",
168+
"@cublas//:cublas",
169+
"@cuda//:cudart",
170+
"@cudnn",
168171
],
169172
)

third_party/tensorrt/local/BUILD

+19-21
Original file line numberDiff line numberDiff line change
@@ -274,47 +274,45 @@ cc_library(
274274

275275
cc_library(
276276
name = "nvcaffeparser",
277-
visibility = ["//visibility:public"],
278277
deps = [
279278
"nvcaffeparser_headers",
280279
"nvcaffeparser_lib",
281280
"nvinfer",
282281
],
282+
visibility = ["//visibility:public"],
283283
)
284284

285285
####################################################################################
286286

287-
cc_import(
288-
name = "nvinferplugin_lib",
289-
shared_library = select({
290-
":aarch64_linux": "lib/x86_64-linux-gnu/libnvinfer_plugin.so",
291-
":windows": "lib/nvinfer_plugin.dll",
292-
"//conditions:default": "lib/x86_64-linux-gnu/libnvinfer_plugin.so",
293-
}),
294-
visibility = ["//visibility:private"],
295-
)
296-
297287
cc_library(
298-
name = "nvinferplugin_headers",
288+
name = "nvinferplugin",
299289
hdrs = select({
300290
":aarch64_linux": glob(["include/aarch64-linux-gnu/NvInferPlugin*.h"]),
301291
":windows": glob(["include/NvInferPlugin*.h"]),
302292
"//conditions:default": glob(["include/x86_64-linux-gnu/NvInferPlugin*.h"]),
303293
}),
294+
srcs = select({
295+
":aarch64_linux": ["lib/aarch64-linux-gnu/libnvinfer_plugin.so"],
296+
":windows": ["lib/nvinfer_plugin.dll"],
297+
"//conditions:default": ["lib/x86_64-linux-gnu/libnvinfer_plugin.so"],
298+
}),
304299
includes = select({
305-
":aarch64_linux": ["include/aarch64-linux-gnu"],
300+
":aarch64_linux": ["include/aarch64-linux-gnu/"],
306301
":windows": ["include/"],
307302
"//conditions:default": ["include/x86_64-linux-gnu/"],
308303
}),
309-
visibility = ["//visibility:private"],
310-
)
311-
312-
cc_library(
313-
name = "nvinferplugin",
314-
visibility = ["//visibility:public"],
315304
deps = [
316305
"nvinfer",
317-
"nvinferplugin_headers",
318-
"nvinferplugin_lib",
306+
"@cublas//:cublas",
307+
"@cuda//:cudart",
308+
"@cuda//:cublas",
309+
"@cudnn",
310+
],
311+
alwayslink = True,
312+
copts = [
313+
"-pthread"
319314
],
315+
linkopts = [
316+
"-lpthread",
317+
]
320318
)

0 commit comments

Comments
 (0)