Skip to content

Commit aa5c077

Browse files
alekstheodalekstheod
committed
Fix jax tests
1 parent eade405 commit aa5c077

File tree

2 files changed

+63
-34
lines changed

2 files changed

+63
-34
lines changed

‎build_tools/rocm/run_jax_ut.sh‎

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -22,6 +22,8 @@ bazel --bazelrc=${XLA_DIR}/build_tools/rocm/rocm_xla.bazelrc test \
2222
--test_tag_filters=cpu,gpu,-tpu,-config-cuda-only \
2323
--action_env=TF_ROCM_AMDGPU_TARGETS=gfx908,gfx90a,gfx942 \
2424
--//jax:build_jaxlib=true \
25+
--run_under=@xla//build_tools/rocm:exclusive_local_wrapper \
26+
--action_env=REMOTE_GPU_TESTING=True \
2527
"//tests/..."
2628

2729
popd

‎third_party/gpus/rocm/BUILD.tpl‎

Lines changed: 61 additions & 34 deletions
Original file line numberDiff line numberDiff line change
@@ -136,6 +136,9 @@ cc_library(
136136
deps = [":rocm_config"],
137137
)
138138

139+
# workaround to bring tensile files to the same fs layout as expected in the lib
140+
# rocblas assumes that tensile files are located in ../roblas/libraries directory
141+
# hibplatslt assumes that tensile files are located in ../hipblaslt/libraries directory
139142
cc_library(
140143
name = "rocm_rpath",
141144
linkopts = select({
@@ -163,8 +166,8 @@ cc_library(
163166

164167
cc_library(
165168
name = "rocm_hip",
166-
srcs = glob(["%{rocm_root}/lib/libamdhip*.so*"]),
167169
hdrs = glob(["%{rocm_root}/include/hip/**"]),
170+
data = glob(["%{rocm_root}/lib/libamdhip*.so*"]),
168171
include_prefix = "rocm",
169172
includes = [
170173
"%{rocm_root}/include",
@@ -173,6 +176,7 @@ cc_library(
173176
visibility = ["//visibility:public"],
174177
deps = [
175178
":rocm_config",
179+
":rocm_rpath",
176180
":rocprofiler_register",
177181
":system_libs",
178182
],
@@ -181,11 +185,13 @@ cc_library(
181185
# Used by jax_rocm_plugin to minimally link to hip runtime.
182186
cc_library(
183187
name = "hip_runtime",
188+
hdrs = glob(["%{rocm_root}/include/hip/**"]),
184189
srcs = glob([
185190
"%{rocm_root}/lib/libamdhip*.so*",
186191
"%{rocm_root}/lib/libamd_comgr.so*",
192+
"%{rocm_root}/lib/librccl*.so*",
193+
"%{rocm_root}/lib/libhipsparse*.so*",
187194
]),
188-
hdrs = glob(["%{rocm_root}/include/hip/**"]),
189195
include_prefix = "rocm",
190196
includes = [
191197
"%{rocm_root}/include",
@@ -194,7 +200,6 @@ cc_library(
194200
visibility = ["//visibility:public"],
195201
deps = [
196202
":rocm_config",
197-
":rocprofiler_register",
198203
":system_libs",
199204
],
200205
)
@@ -216,44 +221,50 @@ cc_library(
216221
includes = [
217222
"%{rocm_root}/include",
218223
],
219-
# workaround to bring tensile files to the same fs layout as expected in the lib
220-
# rocblas assumes that tensile files are located in ../roblas/libraries directory
221-
linkopts = ["-Wl,-rpath,external/local_config_rocm/rocm/%{rocm_root}/lib"],
224+
222225
strip_include_prefix = "%{rocm_root}",
223226
visibility = ["//visibility:public"],
224-
deps = [":rocm_config"],
227+
deps = [
228+
":rocm_config",
229+
":rocm_rpath",
230+
],
225231
)
226232

227233
cc_library(
228234
name = "rocfft",
229235
data = glob(["%{rocm_root}/lib/librocfft*.so*"]),
230-
linkopts = ["-Wl,-rpath,external/local_config_rocm/rocm/%{rocm_root}/lib"],
231236
include_prefix = "rocm",
232237
includes = [
233238
"%{rocm_root}/include",
234239
],
235240
linkstatic = 1,
236241
visibility = ["//visibility:public"],
237-
deps = [":rocm_config"],
242+
deps = [
243+
":rocm_config",
244+
":rocm_rpath",
245+
],
238246
)
239247

240248
cc_library(
241249
name = "hipfft",
242250
data = glob(["%{rocm_root}/lib/libhipfft*.so*"]),
243-
linkopts = ["-Wl,-rpath,external/local_config_rocm/rocm/%{rocm_root}/lib"],
244251
include_prefix = "rocm",
245252
includes = [
246253
"%{rocm_root}/include",
247254
],
255+
linkopts = ["-Wl,-rpath,external/local_config_rocm/rocm/%{rocm_root}/lib"],
248256
linkstatic = 1,
249257
visibility = ["//visibility:public"],
250-
deps = [":rocm_config"],
258+
deps = [
259+
":hipsolver",
260+
":rocm_config",
261+
],
251262
)
252263

253264
cc_library(
254265
name = "hiprand",
255-
srcs = glob(["%{rocm_root}/lib/libhiprand*.so*"]),
256266
hdrs = glob(["%{rocm_root}/include/hiprand/**"]),
267+
data = glob(["%{rocm_root}/lib/libhiprand*.so*"]),
257268
include_prefix = "rocm",
258269
includes = [
259270
"%{rocm_root}/include",
@@ -262,7 +273,10 @@ cc_library(
262273
linkstatic = 1,
263274
strip_include_prefix = "%{rocm_root}",
264275
visibility = ["//visibility:public"],
265-
deps = [":rocm_config"],
276+
deps = [
277+
":rocm_config",
278+
":rocm_rpath",
279+
],
266280
)
267281

268282
miopen_libs = glob([
@@ -281,18 +295,18 @@ cc_library(
281295
includes = [
282296
"%{rocm_root}/include",
283297
],
284-
# workaround to bring miopen db files to the same fs layout as expected in the lib
285-
# rocblas assumes that miopen db files are located in ../share/miopen/db directory
286-
linkopts = ["-Wl,-rpath,external/local_config_rocm/rocm/%{rocm_root}/lib"],
287298
strip_include_prefix = "%{rocm_root}",
288299
visibility = ["//visibility:public"],
289-
deps = [":rocm_config"],
300+
deps = [
301+
":rocm_config",
302+
":rocm_rpath",
303+
],
290304
)
291305

292306
cc_library(
293307
name = "rccl",
294-
srcs = glob(["%{rocm_root}/lib/librccl*.so*"]),
295308
hdrs = glob(["%{rocm_root}/include/rccl/**"]),
309+
data = glob(["%{rocm_root}/lib/librccl*.so*"]),
296310
include_prefix = "rocm",
297311
includes = [
298312
"%{rocm_root}/include",
@@ -303,6 +317,7 @@ cc_library(
303317
visibility = ["//visibility:public"],
304318
deps = [
305319
":rocm_config",
320+
":rocm_rpath",
306321
":system_libs",
307322
],
308323
)
@@ -338,7 +353,6 @@ cc_library(
338353

339354
cc_library(
340355
name = "hipsparse",
341-
srcs = glob(["%{rocm_root}/lib/libhipsparse*.so*"]),
342356
hdrs = glob(["%{rocm_root}/include/hipsparse/**"]),
343357
data = glob(["%{rocm_root}/lib/libhipsparse*.so*"]),
344358
include_prefix = "rocm",
@@ -347,7 +361,11 @@ cc_library(
347361
],
348362
strip_include_prefix = "%{rocm_root}",
349363
visibility = ["//visibility:public"],
350-
deps = [":rocm_config"],
364+
deps = [
365+
":rocm_config",
366+
":rocsparse",
367+
":rocm_rpath",
368+
],
351369
)
352370

353371
roctracer_libs = glob(["%{rocm_root}/lib/libroctracer*.so*"])
@@ -371,40 +389,49 @@ cc_library(
371389

372390
cc_library(
373391
name = "rocprofiler-sdk",
374-
srcs = glob(["%{rocm_root}/lib/librocprofiler-sdk*.so*"]),
375392
hdrs = glob(["%{rocm_root}/include/rocprofiler-sdk/**"]),
393+
srcs = glob(["%{rocm_root}/lib/librocprofiler-sdk*.so*"]),
376394
include_prefix = "rocm",
377395
includes = [
378396
"%{rocm_root}/include/",
379397
],
380398
strip_include_prefix = "%{rocm_root}",
381399
visibility = ["//visibility:public"],
382-
deps = [":rocm_config"],
400+
deps = [
401+
":rocm_config",
402+
":hip_runtime",
403+
],
383404
)
384405

385406
cc_library(
386407
name = "rocsolver",
387-
srcs = glob(["%{rocm_root}/lib/librocsolver*.so*"]),
388408
hdrs = glob(["%{rocm_root}/include/rocsolver/**"]),
409+
data = glob(["%{rocm_root}/lib/librocsolver*.so*"]),
389410
include_prefix = "rocm",
390411
includes = [
391412
"%{rocm_root}/include/",
392413
],
393414
strip_include_prefix = "%{rocm_root}",
394415
visibility = ["//visibility:public"],
395-
deps = [":rocm_config"],
416+
deps = [
417+
":rocm_config",
418+
":rocm_rpath",
419+
],
396420
)
397421

398422
cc_library(
399423
name = "rocsparse",
400-
srcs = glob(["%{rocm_root}/lib/librocsparse*.so*"]),
424+
data = glob(["%{rocm_root}/lib/librocsparse*.so*"]),
401425
include_prefix = "rocm",
402426
includes = [
403427
"%{rocm_root}/include/",
404428
],
405429
strip_include_prefix = "%{rocm_root}",
406430
visibility = ["//visibility:public"],
407-
deps = [":rocm_config"],
431+
deps = [
432+
":rocm_config",
433+
":rocm_rpath",
434+
],
408435
)
409436

410437
hipsolver_libs = glob([
@@ -413,7 +440,6 @@ hipsolver_libs = glob([
413440

414441
cc_library(
415442
name = "hipsolver",
416-
srcs = glob(["%{rocm_root}/lib/libhipsolver*.so*"]),
417443
hdrs = glob(["%{rocm_root}/include/hipsolver/**"]),
418444
data = select({
419445
":build_hermetic": hipsolver_libs,
@@ -439,7 +465,6 @@ cc_library(
439465
":multiple_rocm_paths": hipblas_libs,
440466
"//conditions:default": [],
441467
}),
442-
linkopts = ["-Wl,-rpath,external/local_config_rocm/rocm/%{rocm_root}/lib"],
443468
include_prefix = "rocm",
444469
includes = [
445470
"%{rocm_root}/include/",
@@ -448,7 +473,9 @@ cc_library(
448473
visibility = ["//visibility:public"],
449474
deps = [
450475
":hipblas-common",
476+
":hipsolver",
451477
":rocm_config",
478+
":rocm_rpath",
452479
],
453480
)
454481

@@ -481,22 +508,23 @@ cc_library(
481508
includes = [
482509
"%{rocm_root}/include/",
483510
],
484-
# workaround to bring tensile files to the same fs layout as expected in the lib
485-
# hibplatslt assumes that tensile files are located in ../hipblaslt/libraries directory
486-
linkopts = ["-Wl,-rpath,external/local_config_rocm/rocm/%{rocm_root}/lib"],
487511
strip_include_prefix = "%{rocm_root}",
488512
visibility = ["//visibility:public"],
489-
deps = [":rocm_config"],
513+
deps = [
514+
":rocm_config",
515+
":rocm_rpath",
516+
],
490517
)
491518

492519
cc_library(
493520
name = "rocrand",
494-
srcs = glob(["%{rocm_root}/lib/librocrand*.so*"]),
495521
hdrs = glob(["%{rocm_root}/include/rocrand/**"]),
522+
data = glob(["%{rocm_root}/lib/librocrand*.so*"]),
496523
include_prefix = "rocm",
497524
includes = [
498525
"%{rocm_root}/include/",
499526
],
527+
linkopts = ["-Wl,-rpath,external/local_config_rocm/rocm/%{rocm_root}/lib"],
500528
strip_include_prefix = "%{rocm_root}",
501529
visibility = ["//visibility:public"],
502530
deps = [":rocm_config"],
@@ -592,7 +620,6 @@ platform(
592620
"@platforms//cpu:x86_64",
593621
"@bazel_tools//tools/cpp:clang",
594622
],
595-
596623
exec_properties = {
597624
"container-image": "docker://%{rocm_rbe_docker_image}",
598625
"OSFamily": "Linux",

0 commit comments

Comments
 (0)