PR #10521: [ROCM] fixing gpu_pjrt_client crash and improved debug in rocm_driver

Imported from GitHub PR #10521

In this PR I fix xla/pjrt/gpu/se_gpu_pjrt_client.cc which causes a crash during the execution of se_gpu_pjrt_client_test.cc, namely the subtest **StreamExecutorGpuClientTest.CopyRawToHostFullBuffer**.

The crash was caused by taking a pointer of a unique_ptr in one function's argument and releasing that pointer in another one (not sure why it worked before: maybe on CUDA the compiler has different behaviour?). But I guess the order of evaluating the function arguments is not defined..

Besides, I also improve gpu_driver debug output.

@xla-rotation: can you please have a look ?
Copybara import of the project:

--
fc43f8e by Pavel Emeliyanenko <pavel.emeliyanenko@amd.com>:

fixing gpu_pjrt_client crash and added improved debug output in rocm_driver

--
68e52b5 by Pavel Emeliyanenko <pavel.emeliyanenko@amd.com>:

passing promise by value since the callback function is copied

--
65ab1d4 by Pavel Emeliyanenko <pavel.emeliyanenko@amd.com>:

fixing buildifier warnings

Merging this change closes #10521

COPYBARA_INTEGRATE_REVIEW=#10521 from ROCm:ci_pjrt_gpu_client_fix 65ab1d4
PiperOrigin-RevId: 615973610

Author: pemeliya

xla/pjrt/gpu/BUILD

@@ -5,7 +5,6 @@ load("@local_config_rocm//rocm:build_defs.bzl", "if_rocm")
 load("@tsl//tsl:tsl.bzl", "internal_visibility")
 load("@tsl//tsl/platform:build_config.bzl", "tf_proto_library")
 load("@tsl//tsl/platform:rules_cc.bzl", "cc_library")
-load("@tsl//tsl/platform/default:cuda_build_defs.bzl", "if_cuda_is_configured")
 
 package(
     # copybara:uncomment default_applicable_licenses = ["//tensorflow:license"],
@@ -127,7 +126,6 @@ cc_library(
 xla_cc_test(
     name = "se_gpu_pjrt_client_test",
     srcs = if_gpu_is_configured(["se_gpu_pjrt_client_test.cc"]),
-    local_defines = if_cuda_is_configured(["GOOGLE_CUDA=1"]),
     tags = [
         "gpu",
         "no_oss",

xla/pjrt/gpu/se_gpu_pjrt_client.cc

@@ -594,11 +594,11 @@ PjRtFuture<absl::Status> StreamExecutorGpuClient::CopyRawSubBufferToHost(
                         /*reference_held=*/false);
 
   auto promise = PjRtFuture<absl::Status>::CreatePromise();
+  auto stream_ptr = stream.get();
   auto callback_status = local_device->ThenExecuteCallback(
-      stream.get(), [promise, free_sub_range = sub_buffer.release(),
-                     free_stream = stream.release(), local_device]() mutable {
+      stream_ptr,
+      [promise, free_stream = stream.release(), local_device]() mutable {
         auto stream = std::unique_ptr<se::Stream>(free_stream);
-        auto sub_range = std::unique_ptr<se::DeviceMemoryBase>(free_sub_range);
         local_device->ReturnStreamToPool(std::move(stream));
         promise.Set(OkStatus());
       });

xla/stream_executor/rocm/rocm_driver.cc

@@ -1580,7 +1580,7 @@ struct BitPatternToValue {
           "failed to synchronous memcpy from host to device: Gpu dst: %p;"
           " host src: %p; size: %llu=0x%llx",
           absl::bit_cast<void*>(gpu_dst), host_src, size, size));
-  VLOG(2) << "successfully enqueued sync memcpy h2d of " << size << " bytes";
+  VLOG(2) << "successfully sync memcpy'd h2d of " << size << " bytes";
   return absl::OkStatus();
 }
 
@@ -1616,7 +1616,8 @@ struct BitPatternToValue {
   }
   VLOG(2) << "successfully enqueued async memcpy d2h of " << size
           << " bytes from " << absl::bit_cast<void*>(gpu_src) << " to "
-          << host_dst << " on stream " << stream;
+          << host_dst << " on stream " << stream
+          << " device: " << context->device_ordinal();
   return true;
 }
 
@@ -1636,8 +1637,10 @@ struct BitPatternToValue {
         size);
     return false;
   }
-  VLOG(2) << "successfully enqueued async memcpy h2d of " << size << " bytes"
-          << " on stream " << stream;
+  VLOG(2) << "successfully enqueued async memcpy h2d of " << size
+          << " bytes from " << host_src << " to "
+          << absl::bit_cast<void*>(gpu_dst) << " on stream " << stream
+          << " device: " << context->device_ordinal();
   return true;
 }
 
@@ -1664,7 +1667,11 @@ struct BitPatternToValue {
 
     return false;
   }
-  VLOG(2) << "successfully enqueued async memcpy d2d of " << size << " bytes";
+
+  VLOG(2) << "successfully enqueued async memcpy d2d of " << size
+          << " bytes from " << absl::bit_cast<void*>(gpu_src) << " to "
+          << absl::bit_cast<void*>(gpu_dst) << " on stream " << stream
+          << " device: " << context->device_ordinal();
   return true;
 }