PR #33505: [XLA:GPU] Dump command buffer contents to folder specified by --xla-dump-to through dump.h

Imported from GitHub PR #33505

📝 Summary of Changes
This PR migrates the command buffer content's dump through dump.h

Copybara import of the project:

--
0b870f2 by Shawn Wang <[email protected]>:

Dump command buffer contents to folder specified by --xla-dump-to through dump.h

--
123c254 by Shawn Wang <[email protected]>:

clang format fix

Merging this change closes #33505

COPYBARA_INTEGRATE_REVIEW=#33505 from shawnwang18:shawnw/stringfy_cmd_buffer 123c254
PiperOrigin-RevId: 828357446

Author: shawnwang18

xla/stream_executor/command_buffer.h

@@ -311,6 +311,8 @@ class CommandBuffer {
   // Returns command buffer state.
   virtual State state() const = 0;
 
+  virtual std::string ToString() const = 0;
+
   //--------------------------------------------------------------------------//
   // Command buffer tracing API
   //--------------------------------------------------------------------------//

xla/stream_executor/cuda/BUILD

@@ -1559,6 +1559,7 @@ cc_library(
         "@com_google_absl//absl/types:span",
         "@local_config_cuda//cuda:cuda_headers",
         "@tsl//tsl/platform:casts",
+        "@tsl//tsl/platform:path",
     ],
 )
 

xla/stream_executor/cuda/cuda_command_buffer.cc

@@ -53,6 +53,7 @@ limitations under the License.
 #include "xla/tsl/platform/errors.h"
 #include "xla/tsl/platform/statusor.h"
 #include "tsl/platform/casts.h"
+#include "tsl/platform/path.h"
 
 namespace stream_executor::gpu {
 namespace {
@@ -831,4 +832,25 @@ absl::Status CudaCommandBuffer::CheckCanBeUpdated() {
   return absl::OkStatus();
 }
 
+std::string CudaCommandBuffer::ToString() const {
+  std::string path = tsl::io::GetTempFilename(/*extension=*/"dot");
+#if CUDA_VERSION >= 12000
+  int flags = CU_GRAPH_DEBUG_DOT_FLAGS_VERBOSE;
+  auto dot_print_status =
+      cuda::ToStatus(cuGraphDebugDotPrint(graph_, path.c_str(), flags),
+                     "Failed to print gpu graph debug file");
+  if (!dot_print_status.ok()) {
+    return std::string(dot_print_status.message());
+  }
+  std::string dot_file_contents;
+  auto read_status =
+      tsl::ReadFileToString(tsl::Env::Default(), path, &dot_file_contents);
+  if (!read_status.ok()) {
+    return std::string(read_status.message());
+  }
+  return dot_file_contents;
+#endif  // CUDA_VERSION >= 12000
+  return "CUDA graph debug dot print is not supported.";
+}
+
 }  // namespace stream_executor::gpu

xla/stream_executor/cuda/cuda_command_buffer.h

@@ -53,6 +53,8 @@ class CudaCommandBuffer final : public GpuCommandBuffer {
   static absl::StatusOr<std::unique_ptr<CudaCommandBuffer>> Create(
       Mode mode, StreamExecutor* executor, CudaContext* cuda_context);
 
+  std::string ToString() const override;
+
   ~CudaCommandBuffer() override;
 
  private:

xla/stream_executor/gpu/BUILD

@@ -137,6 +137,8 @@ cc_library(
         "gpu_command_buffer.h",
     ],
     deps = [
+        "//xla:debug_options_flags",
+        "//xla/service:dump",
         "//xla/stream_executor:bit_pattern",
         "//xla/stream_executor:command_buffer",
         "//xla/stream_executor:device_memory",

xla/stream_executor/gpu/gpu_command_buffer.cc

@@ -33,6 +33,8 @@ limitations under the License.
 #include "absl/strings/str_format.h"
 #include "absl/strings/string_view.h"
 #include "absl/types/span.h"
+#include "xla/debug_options_flags.h"
+#include "xla/service/dump.h"
 #include "xla/stream_executor/bit_pattern.h"
 #include "xla/stream_executor/command_buffer.h"
 #include "xla/stream_executor/device_memory.h"
@@ -571,14 +573,10 @@ absl::Status GpuCommandBuffer::Finalize() {
   // Maybe dump created GPU graph to a dot file for debugging.
   if (state_ == State::kCreate &&
       (VLOG_IS_ON(10) || (VLOG_IS_ON(9) && mode_ == Mode::kPrimary))) {
-    std::string path = tsl::io::GetTempFilename(/*extension=*/"dot");
-    TF_RETURN_IF_ERROR(WriteGraphToDotFile(path));
-    if (VLOG_IS_ON(100) || (VLOG_IS_ON(90) && mode_ == Mode::kPrimary)) {
-      std::string dot_file_contents;
-      TF_RETURN_IF_ERROR(
-          tsl::ReadFileToString(tsl::Env::Default(), path, &dot_file_contents));
-      VLOG(90) << "Contents of " << path << " is:\n" << dot_file_contents;
-    }
+    xla::DebugOptions debug_options = xla::GetDebugOptionsFromFlags();
+    std::string contents = ToString();
+    std::string filename = absl::StrFormat("gpu_command_buffer_%p.dot", this);
+    xla::DumpToFileInDir(debug_options, filename, contents);
   }
 
   size_t num_commands = commands_.size();

xla/stream_executor/rocm/rocm_command_buffer.cc

@@ -463,4 +463,8 @@ absl::Status RocmCommandBuffer::CheckCanBeUpdated() {
   return absl::OkStatus();
 }
 
+std::string RocmCommandBuffer::ToString() const {
+  return "ROCM graph debug dot print is not supported.";
+}
+
 }  // namespace stream_executor::gpu

xla/stream_executor/rocm/rocm_command_buffer.h

@@ -46,6 +46,8 @@ class RocmCommandBuffer : public GpuCommandBuffer {
   static absl::StatusOr<std::unique_ptr<RocmCommandBuffer>> Create(
       Mode mode, StreamExecutor* executor);
 
+  std::string ToString() const override;
+
   ~RocmCommandBuffer() override;
 
  private: