[pull] master from tensorflow:master

third_party/xla/xla/python/ifrt/ir/BUILD

@@ -331,26 +331,17 @@ cc_library(
     ],
 )
 
-exports_files(
-    ["basic_atom_program_compiler.h"],
-    visibility = if_google(
-        [
-            "//xla/python/ifrt/ir/google:__pkg__",
-            "//xla/python/ifrt:friends",
-        ],
-        ["//visibility:private"],
-    ),
-)
-
 cc_library(
     name = "basic_atom_program_compiler",
+    srcs = ["basic_atom_program_compiler.cc"],
     hdrs = ["basic_atom_program_compiler.h"],
     compatible_with = get_compatible_with_libtpu_portable(),
     visibility = internal_visibility(["//xla/python/ifrt:friends"]),
     deps = [
-        if_google("//xla/python/ifrt/ir/google:basic_atom_program_compiler", ":basic_atom_program_compiler_oss"),
         ":atom_program_compiler",
+        ":compilation_utils",
         ":ir",
+        "//xla:status_macros",
         "//xla/pjrt:pjrt_executable",
         "//xla/python/ifrt",
         "//xla/python/ifrt/hlo:hlo_program",
@@ -364,20 +355,28 @@ cc_library(
         "@com_google_absl//absl/status",
         "@com_google_absl//absl/status:statusor",
         "@com_google_absl//absl/strings:str_format",
+        "@com_google_absl//absl/strings:string_view",
         "@com_google_absl//absl/types:span",
     ],
 )
 
+exports_files(
+    ["compilation_utils.h"],
+    visibility = if_google(
+        [
+            "//xla/python/ifrt/ir/google:__pkg__",
+        ],
+        ["//visibility:private"],
+    ),
+)
+
 cc_library(
-    name = "basic_atom_program_compiler_oss",
-    srcs = [
-        "basic_atom_program_compiler.cc",
-        "basic_atom_program_compiler.h",
-    ],
-    visibility = ["//visibility:private"],
+    name = "compilation_utils",
+    hdrs = ["compilation_utils.h"],
+    compatible_with = get_compatible_with_libtpu_portable(),
     deps = [
-        ":atom_program_compiler",
-        ":ir",
+        if_google("//xla/python/ifrt/ir/google:compilation_utils", ":compilation_utils_oss"),
+        "//xla/client:executable_build_options",
         "//xla/pjrt:pjrt_executable",
         "//xla/python/ifrt",
         "//xla/python/ifrt/hlo:hlo_program",
@@ -391,8 +390,23 @@ cc_library(
         "@com_google_absl//absl/status",
         "@com_google_absl//absl/status:statusor",
         "@com_google_absl//absl/strings:str_format",
+        "@com_google_absl//absl/strings:string_view",
         "@com_google_absl//absl/types:span",
     ],
+)
+
+cc_library(
+    name = "compilation_utils_oss",
+    srcs = [
+        "compilation_utils.cc",
+        "compilation_utils.h",
+    ],
+    compatible_with = get_compatible_with_portable(),
+    deps = [
+        "//xla/pjrt:pjrt_executable",
+        "@com_google_absl//absl/status",
+        "@com_google_absl//absl/strings:string_view",
+    ],
     alwayslink = True,
 )
 

third_party/xla/xla/python/ifrt/ir/basic_atom_program_compiler.cc

@@ -25,6 +25,7 @@ limitations under the License.
 #include "absl/status/status.h"
 #include "absl/status/statusor.h"
 #include "absl/strings/str_format.h"
+#include "absl/strings/string_view.h"
 #include "absl/types/span.h"
 #include "xla/pjrt/pjrt_executable.h"
 #include "xla/python/ifrt/client.h"
@@ -34,10 +35,12 @@ limitations under the License.
 #include "xla/python/ifrt/executable.h"
 #include "xla/python/ifrt/hlo/hlo_program.h"
 #include "xla/python/ifrt/ir/atom_program_compiler.h"
+#include "xla/python/ifrt/ir/compilation_utils.h"
 #include "xla/python/ifrt/ir/ifrt_dialect.h"
 #include "xla/python/ifrt/shape.h"
 #include "xla/python/pjrt_ifrt/xla_compiler.h"
 #include "xla/service/computation_placer.h"
+#include "xla/status_macros.h"
 #include "xla/tsl/concurrency/future.h"
 #include "xla/tsl/platform/errors.h"
 #include "xla/tsl/platform/statusor.h"
@@ -83,6 +86,28 @@ tsl::Future<LoadedExecutableRef> BasicAtomProgramCompiler::CompileXla(
 
   TF_ASSIGN_OR_RETURN(DeviceListRef devices,
                       GetDeviceListFromXlaCompileOptions(client_, options));
+
+  TF_RET_CHECK(!devices->devices().empty())
+      << "CompileXla was called with empty device assignment.";
+  xla::ifrt::Device* first_device = devices->devices().front();
+  absl::string_view platform_name = first_device->PlatformName();
+  TF_RET_CHECK(platform_name == "tpu" || platform_name == "cpu" ||
+               platform_name == "cuda" || platform_name == "rocm")
+      << "Unsupported platform: " << platform_name;
+
+  if (strict_memory_reservation_) {
+    TF_RET_CHECK(first_device->PlatformName() == "tpu")
+        << "IFRT IR `strict_memory_reservation` option is only supported for "
+           "TPU devices. Got device platform: "
+        << first_device->PlatformName();
+    TF_ASSIGN_OR_RETURN(
+        int64_t device_memory,
+        first_device->Attributes().Get<int64_t>("device_memory_bytes_limit"));
+
+    TF_RETURN_IF_ERROR(SetStrictMemoryReservation(hlo_program->name(),
+                                                  device_memory, options));
+  }
+
   return client_->GetDefaultCompiler()->CompileAndLoad(
       std::move(hlo_program), std::make_unique<XlaCompileOptions>(
                                   std::move(options), std::move(devices)));

third_party/xla/xla/python/ifrt/ir/compilation_utils.cc

@@ -0,0 +1,34 @@
+/* Copyright 2026 The OpenXLA Authors.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+==============================================================================*/
+
+#include "xla/python/ifrt/ir/compilation_utils.h"
+
+#include <cstdint>
+
+#include "absl/status/status.h"
+#include "absl/strings/string_view.h"
+#include "xla/pjrt/pjrt_executable.h"
+
+namespace xla {
+namespace ifrt {
+
+absl::Status SetStrictMemoryReservation(absl::string_view program_name,
+                                        int64_t device_memory,
+                                        xla::CompileOptions& options) {
+  return absl::OkStatus();
+}
+
+}  // namespace ifrt
+}  // namespace xla

third_party/xla/xla/python/ifrt/ir/compilation_utils.h

@@ -0,0 +1,36 @@
+/* Copyright 2026 The OpenXLA Authors.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+==============================================================================*/
+
+#ifndef XLA_PYTHON_IFRT_IR_COMPILATION_UTILS_H_
+#define XLA_PYTHON_IFRT_IR_COMPILATION_UTILS_H_
+
+#include <cstdint>
+
+#include "absl/status/status.h"
+#include "absl/strings/string_view.h"
+#include "xla/pjrt/pjrt_executable.h"
+
+namespace xla {
+namespace ifrt {
+
+// Sets the strict memory reservation in the compile options.
+absl::Status SetStrictMemoryReservation(absl::string_view program_name,
+                                        int64_t device_memory,
+                                        xla::CompileOptions& options);
+
+}  // namespace ifrt
+}  // namespace xla
+
+#endif  // XLA_PYTHON_IFRT_IR_COMPILATION_UTILS_H_