[IFRT] Add custom_options in ifrt::ExecuteOptions

The users can provide any runtime-specific execution options at
`ifrt::LoadedExecutable::Execute()` using
`ifrt::ExecuteOptions::custom_options`. These options are intended to convey
small per-execution metadata that are to be used by implementation internals
with little to (ideally) no interaction with IFRT API semantics. Example use
cases are propagating a profiling key, which currently goes out of scope of
IFRT APIs.

PiperOrigin-RevId: 702546163

xla/python/ifrt/BUILD

@@ -173,6 +173,20 @@ xla_cc_test(
 tf_proto_library(
     name = "execute_options_proto",
     srcs = ["execute_options.proto"],
+    protodeps = [":attribute_map_proto"],
+)
+
+xla_cc_test(
+    name = "executable_test",
+    size = "small",
+    srcs = ["executable_test.cc"],
+    deps = [
+        ":attribute_map",
+        ":ifrt",
+        "@com_google_googletest//:gtest_main",
+        "@tsl//tsl/platform:statusor",
+        "@tsl//tsl/platform:test",
+    ],
 )
 
 xla_cc_test(

xla/python/ifrt/executable.cc

@@ -16,35 +16,43 @@ limitations under the License.
 #include "xla/python/ifrt/executable.h"
 
 #include "absl/status/statusor.h"
+#include "xla/python/ifrt/attribute_map.h"
 #include "xla/python/ifrt/execute_options.pb.h"
+#include "tsl/platform/statusor.h"
 
 namespace xla {
 namespace ifrt {
 
 char Executable::ID = 0;
 char LoadedExecutable::ID = 0;
 
-absl::StatusOr<xla::ifrt::ExecuteOptionsProto> ExecuteOptions::ToProto() const {
+absl::StatusOr<ExecuteOptionsProto> ExecuteOptions::ToProto() const {
   ExecuteOptionsProto proto;
 
   proto.set_launch_id(launch_id);
   proto.mutable_non_donatable_input_indices()->Add(
       non_donatable_input_indices.begin(), non_donatable_input_indices.end());
   proto.set_fill_status(fill_status);
+  if (custom_options.has_value()) {
+    *proto.mutable_custom_options() = custom_options->ToProto();
+  }
 
   return proto;
 }
 
-absl::StatusOr<xla::ifrt::ExecuteOptions> ExecuteOptions::FromProto(
-    const xla::ifrt::ExecuteOptionsProto& proto) {
+absl::StatusOr<ExecuteOptions> ExecuteOptions::FromProto(
+    const ExecuteOptionsProto& proto) {
   ExecuteOptions options;
 
   options.launch_id = proto.launch_id();
   options.non_donatable_input_indices.insert(
       proto.non_donatable_input_indices().begin(),
       proto.non_donatable_input_indices().end());
   options.fill_status = proto.fill_status();
-
+  if (proto.has_custom_options()) {
+    TF_ASSIGN_OR_RETURN(options.custom_options,
+                        AttributeMap::FromProto(proto.custom_options()));
+  }
   return options;
 }
 

xla/python/ifrt/executable.h

@@ -97,7 +97,7 @@ class Executable : public llvm::RTTIExtends<Executable, llvm::RTTIRoot> {
   // Returns named values for cost properties of this executable (such as
   // operations, size of input/outputs, and run time estimate). Properties may
   // differ for different implementations and platforms.
-  virtual absl::StatusOr<xla::ifrt::AttributeMap> GetCostAnalysis() const = 0;
+  virtual absl::StatusOr<AttributeMap> GetCostAnalysis() const = 0;
 
   // Returns the compile options used to compile this executable.
   // TODO(phawkins): consider removing this API and having the client remember
@@ -127,6 +127,10 @@ struct ExecuteOptions {
   // an invalid future.
   bool fill_status = false;
 
+  // Custom execution options specific to the runtime. The user and the runtime
+  // are responsible for ensuring version compatibility.
+  std::optional<AttributeMap> custom_options;
+
   absl::StatusOr<ExecuteOptionsProto> ToProto() const;
 
   static absl::StatusOr<ExecuteOptions> FromProto(
@@ -201,11 +205,11 @@ class LoadedExecutable
   // Returns named values for cost properties of this executable (such as
   // operations, size of input/outputs, and run time estimate). Properties may
   // differ for different implementations and platforms.
-  virtual absl::StatusOr<xla::ifrt::AttributeMap> GetCostAnalysis() const = 0;
+  virtual absl::StatusOr<AttributeMap> GetCostAnalysis() const = 0;
 
   // `LoadedExecutable` methods.
 
-  using ExecuteOptions = xla::ifrt::ExecuteOptions;
+  using ExecuteOptions = ::xla::ifrt::ExecuteOptions;
 
   // Result from an execution.
   struct ExecuteResult {

xla/python/ifrt/executable_test.cc

@@ -0,0 +1,53 @@
+/* Copyright 2024 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/executable.h"
+
+#include <gmock/gmock.h>
+#include <gtest/gtest.h>
+#include "xla/python/ifrt/attribute_map.h"
+#include "tsl/platform/statusor.h"
+#include "tsl/platform/test.h"
+
+namespace xla {
+namespace ifrt {
+
+using ::testing::Pair;
+using ::testing::UnorderedElementsAre;
+
+TEST(ExecuteOptionsTest, RoundTrip) {
+  LoadedExecutable::ExecuteOptions options;
+  options.launch_id = 1234;
+  options.non_donatable_input_indices.insert(0);
+  options.non_donatable_input_indices.insert(3);
+  options.fill_status = true;
+  options.custom_options = AttributeMap(
+      AttributeMap::Map({{"foo", AttributeMap::StringValue("bar")}}));
+  TF_ASSERT_OK_AND_ASSIGN(ExecuteOptionsProto serialized, options.ToProto());
+  TF_ASSERT_OK_AND_ASSIGN(
+      auto deserialized,
+      LoadedExecutable::ExecuteOptions::FromProto(serialized));
+  EXPECT_EQ(deserialized.launch_id, 1234);
+  EXPECT_THAT(deserialized.non_donatable_input_indices,
+              UnorderedElementsAre(0, 3));
+  EXPECT_TRUE(deserialized.fill_status);
+  ASSERT_TRUE(deserialized.custom_options.has_value());
+  EXPECT_THAT(
+      deserialized.custom_options->map(),
+      UnorderedElementsAre(Pair("foo", AttributeMap::StringValue("bar"))));
+}
+
+}  // namespace ifrt
+}  // namespace xla

xla/python/ifrt/execute_options.proto

@@ -17,11 +17,14 @@ syntax = "proto3";
 
 package xla.ifrt;
 
+import "xla/python/ifrt/attribute_map.proto";
+
 message ExecuteOptionsProto {
   bool untuple_result = 2;
   int32 launch_id = 3;
   repeated int32 non_donatable_input_indices = 7;
   bool fill_status = 9;
+  AttributeMapProto custom_options = 10;
 
   reserved 1, 4 to 6, 8;
 }