From 035c917d779f1e13f1cb5864d23a8af529ba2996 Mon Sep 17 00:00:00 2001
From: pizzud <pizzud@google.com>
Date: Sun, 29 Sep 2024 16:43:09 -0700
Subject: [PATCH] hlo_runner_pjrt: Have PjRtWrappedExecutable own the
 underlying executable.

Avoids a memory leak in various unit tests currently masked by not actually
looking for leaks.

PiperOrigin-RevId: 680342833
---
 xla/service/hlo_runner_pjrt.cc | 9 ++++-----
 1 file changed, 4 insertions(+), 5 deletions(-)

diff --git a/xla/service/hlo_runner_pjrt.cc b/xla/service/hlo_runner_pjrt.cc
index 3965bf61870f3..0047ccd428819 100644
--- a/xla/service/hlo_runner_pjrt.cc
+++ b/xla/service/hlo_runner_pjrt.cc
@@ -114,6 +114,7 @@ absl::StatusOr<ExecuteOptions> GenerateExecuteOptions(const HloModule& module) {
 // TODO(b/245550554): Remove the use of PjRtWrappedExecutable.
 class PjRtWrappedExecutable : public Executable {
  public:
+  // Takes ownership of the provided executable.
   explicit PjRtWrappedExecutable(std::shared_ptr<HloModule> hlo_module,
                                  PjRtLoadedExecutable* pjrt_loaded_executable)
       : Executable(hlo_module),
@@ -125,11 +126,11 @@ class PjRtWrappedExecutable : public Executable {
       HloExecutionProfile* hlo_execution_profile) override;
 
   PjRtLoadedExecutable* GetPjRtLoadedExecutable() const {
-    return pjrt_loaded_executable_;
+    return pjrt_loaded_executable_.get();
   }
 
  private:
-  PjRtLoadedExecutable* pjrt_loaded_executable_;
+  std::unique_ptr<PjRtLoadedExecutable> pjrt_loaded_executable_;
 };
 
 absl::StatusOr<ExecutionOutput> PjRtWrappedExecutable::ExecuteAsyncOnStream(
@@ -373,9 +374,7 @@ absl::StatusOr<std::unique_ptr<Executable>> HloRunnerPjRt::CreateExecutable(
           std::move(pjrt_executable->GetHloModules().value()[0])),
       pjrt_executable.release());
 
-  std::unique_ptr<Executable> exec =
-      static_cast<std::unique_ptr<Executable>>(executable.release());
-  return exec;
+  return executable;
 }
 
 absl::StatusOr<std::vector<Literal>> HloRunnerPjRt::ExecuteReplicated(