From b1bec33c80343e3c2afd5c3bafe78e96619417b4 Mon Sep 17 00:00:00 2001
From: pizzud <pizzud@google.com>
Date: Fri, 27 Sep 2024 14:03:27 -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: 679721068
---
 xla/service/hlo_runner_pjrt.cc | 8 +++-----
 1 file changed, 3 insertions(+), 5 deletions(-)

diff --git a/xla/service/hlo_runner_pjrt.cc b/xla/service/hlo_runner_pjrt.cc
index 3965bf61870f3a..66b64492d32c35 100644
--- a/xla/service/hlo_runner_pjrt.cc
+++ b/xla/service/hlo_runner_pjrt.cc
@@ -125,11 +125,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 +373,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(