[XLA] Add an autotune result wrapper that can be used for externally …

…caching autotune results. PiperOrigin-RevId: 647114148
tensorflow · Jun 27, 2024 · 437d4ac · 437d4ac
1 parent 4acdaaa
commit 437d4ac
Show file tree

Hide file tree

Showing 4 changed files with 307 additions and 0 deletions.
diff --git a/third_party/xla/xla/BUILD b/third_party/xla/xla/BUILD
@@ -1182,6 +1182,34 @@ tf_proto_library(
     ]),
 )
 
+cc_library(
+    name = "autotune_result_wrapper",
+    srcs = ["autotune_result_wrapper.cc"],
+    hdrs = ["autotune_result_wrapper.h"],
+    visibility = ["//visibility:public"],
+    deps = [
+        ":autotune_results_proto_cc",
+        ":autotuning_proto_cc",
+        "@com_google_absl//absl/log:check",
+        "@com_google_absl//absl/status",
+        "@com_google_absl//absl/status:statusor",
+        "@local_tsl//tsl/lib/strings:proto_serialization",
+    ],
+)
+
+xla_cc_test(
+    name = "autotune_result_wrapper_test",
+    srcs = ["autotune_result_wrapper_test.cc"],
+    deps = [
+        ":autotune_result_wrapper",
+        ":autotuning_proto_cc",
+        ":test",
+        ":test_helpers",
+        "@local_tsl//tsl/platform:statusor",
+        "@local_tsl//tsl/platform:test_main",
+    ],
+)
+
 cc_library(
     name = "printer",
     srcs = ["printer.cc"],

diff --git a/third_party/xla/xla/autotune_result_wrapper.cc b/third_party/xla/xla/autotune_result_wrapper.cc
@@ -0,0 +1,96 @@
+/* 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/autotune_result_wrapper.h"
+
+#include <vector>
+
+#include "absl/log/check.h"
+#include "absl/status/status.h"
+#include "xla/autotuning.pb.h"
+#include "tsl/lib/strings/proto_serialization.h"
+
+namespace xla {
+
+/*static*/ absl::StatusOr<AutotuneResultWrapper>
+AutotuneResultWrapper::FromKeyAndValue(OpaqueKey key, OpaqueValue value) {
+  AutotuneResults key_proto;
+  if (!key_proto.ParseFromString(key)) {
+    return absl::Status(absl::StatusCode::kInvalidArgument,
+                        "Could not parse the provided key");
+  }
+
+  AutotuneResults::Entry value_entry;
+  if (!value_entry.ParseFromString(value)) {
+    return absl::Status(absl::StatusCode::kInvalidArgument,
+                        "Could not parse the provided value");
+  }
+
+  AutotuneResults::Entry full_entry;
+  full_entry.set_device(key_proto.results(0).device());
+  full_entry.set_hlo(key_proto.results(0).hlo());
+  *full_entry.mutable_result() = value_entry.result();
+  return AutotuneResultWrapper(full_entry, key_proto.version());
+}
+
+AutotuneResultWrapper::OpaqueKey AutotuneResultWrapper::Key() const {
+  AutotuneResults key_proto;
+  key_proto.set_version(version_);
+  auto entry = key_proto.add_results();
+  entry->set_device(autotune_result_.device());
+  entry->set_hlo(autotune_result_.hlo());
+  OpaqueKey serialized;
+  CHECK(tsl::SerializeToStringDeterministic(key_proto, &serialized));
+  return serialized;
+}
+
+AutotuneResultWrapper::OpaqueValue AutotuneResultWrapper::Value() const {
+  AutotuneResults::Entry entry;
+  *entry.mutable_result() = autotune_result_.result();
+  OpaqueValue serialized;
+  CHECK(tsl::SerializeToStringDeterministic(entry, &serialized));
+  return serialized;
+}
+
+/*static*/ std::vector<AutotuneResultWrapper>
+AutotuneResultWrapper::AutotuneResultsToWrappers(
+    const AutotuneResults& autotune_results) {
+  std::vector<AutotuneResultWrapper> wrappers;
+  wrappers.reserve(autotune_results.results_size());
+  for (const auto& result : autotune_results.results()) {
+    wrappers.push_back(
+        AutotuneResultWrapper(result, autotune_results.version()));
+  }
+  return wrappers;
+}
+
+/*static*/ absl::StatusOr<AutotuneResults>
+AutotuneResultWrapper::AutotuneResultsFromWrappers(
+    const std::vector<AutotuneResultWrapper>& wrappers) {
+  AutotuneResults autotune_results;
+  for (const auto& wrapper : wrappers) {
+    if (autotune_results.results_size() > 0 &&
+        autotune_results.version() != wrapper.version_) {
+      return absl::Status(absl::StatusCode::kInvalidArgument,
+                          "All wrappers must have the same version number");
+    }
+
+    *autotune_results.add_results() = wrapper.autotune_result_;
+    autotune_results.set_version(wrapper.version_);
+  }
+  return autotune_results;
+}
+
+}  // namespace xla
diff --git a/third_party/xla/xla/autotune_result_wrapper.h b/third_party/xla/xla/autotune_result_wrapper.h
@@ -0,0 +1,69 @@
+/* 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.
+==============================================================================*/
+
+#ifndef XLA_AUTOTUNE_RESULT_WRAPPER_H_
+#define XLA_AUTOTUNE_RESULT_WRAPPER_H_
+
+#include <cstdint>
+#include <string>
+#include <vector>
+
+#include "absl/status/statusor.h"
+#include "xla/autotune_results.pb.h"
+
+namespace xla {
+
+// This class is a thin wrapper around AutotuneResults::Entry. It is used to
+// provide opaque accessors to an entry's key and value without exposing the
+// internal structure of the entry.
+class AutotuneResultWrapper {
+ public:
+  using OpaqueKey = std::string;
+  using OpaqueValue = std::string;
+
+  // Creates an AutotuneResultWrapper from a key and value. The provided key and
+  // value must be ones that were previously returned by calls to Key() and
+  // Value().
+  static absl::StatusOr<AutotuneResultWrapper> FromKeyAndValue(
+      OpaqueKey key, OpaqueValue value);
+
+  // An opaque string that can be used as a key for this Autotuning result.
+  // Do not rely on the format of this string.
+  OpaqueKey Key() const;
+
+  // An opaque string that encodes the autotuning result.
+  // Do not rely on the format of this string.
+  OpaqueValue Value() const;
+
+  static std::vector<AutotuneResultWrapper> AutotuneResultsToWrappers(
+      const AutotuneResults& autotune_results);
+
+  // Returns the AutotuneResults proto that corresponds to the provided
+  // wrappers. This function will return an error if the provided wrappers have
+  // inconsistent versions.
+  static absl::StatusOr<AutotuneResults> AutotuneResultsFromWrappers(
+      const std::vector<AutotuneResultWrapper>& wrappers);
+
+ private:
+  AutotuneResultWrapper(const AutotuneResults::Entry& result, int32_t version)
+      : autotune_result_(result), version_(version) {}
+
+  AutotuneResults::Entry autotune_result_;
+  int32_t version_;
+};
+
+}  // namespace xla
+
+#endif  // XLA_AUTOTUNE_RESULT_WRAPPER_H_
diff --git a/third_party/xla/xla/autotune_result_wrapper_test.cc b/third_party/xla/xla/autotune_result_wrapper_test.cc
@@ -0,0 +1,114 @@
+/* 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/autotune_result_wrapper.h"
+
+#include <cstdint>
+#include <utility>
+#include <vector>
+
+#include "xla/autotuning.pb.h"
+#include "xla/test.h"
+#include "xla/test_helpers.h"
+#include "tsl/platform/statusor.h"
+
+namespace xla {
+namespace {
+
+AutotuneResults ThreeAutotuneEntries(int32_t version) {
+  AutotuneResults results;
+  results.set_version(version);
+  auto r1 = results.add_results();
+  r1->set_device("dev1");
+  r1->set_hlo("hlo1");
+  r1->mutable_result()->set_scratch_bytes(1);
+
+  auto r2 = results.add_results();
+  r2->set_device("dev2");
+  r2->set_hlo("hlo2");
+  r2->mutable_result()->set_scratch_bytes(2);
+
+  auto r3 = results.add_results();
+  r3->set_device("dev3");
+  r3->set_hlo("hlo3");
+  r3->mutable_result()->set_scratch_bytes(3);
+
+  return results;
+}
+
+TEST(AutotuneResultWrapperTest, FullRoundTrip) {
+  std::vector<AutotuneResultWrapper> wrappers =
+      AutotuneResultWrapper::AutotuneResultsToWrappers(
+          ThreeAutotuneEntries(/*version=*/42));
+
+  std::vector<std::pair<AutotuneResultWrapper::OpaqueKey,
+                        AutotuneResultWrapper::OpaqueValue>>
+      key_value_pairs;
+  for (const auto& wrapper : wrappers) {
+    key_value_pairs.push_back(std::make_pair(wrapper.Key(), wrapper.Value()));
+  }
+
+  std::vector<AutotuneResultWrapper> new_wrappers;
+  for (const auto& [key, value] : key_value_pairs) {
+    TF_ASSERT_OK_AND_ASSIGN(AutotuneResultWrapper wrapper,
+                            AutotuneResultWrapper::FromKeyAndValue(key, value));
+    new_wrappers.push_back(std::move(wrapper));
+  }
+
+  TF_ASSERT_OK_AND_ASSIGN(
+      AutotuneResults round_tripped,
+      AutotuneResultWrapper::AutotuneResultsFromWrappers(new_wrappers));
+  EXPECT_EQ(round_tripped.results_size(), 3);
+  EXPECT_EQ(round_tripped.version(), 42);
+  EXPECT_EQ(round_tripped.results(0).device(), "dev1");
+  EXPECT_EQ(round_tripped.results(0).hlo(), "hlo1");
+  EXPECT_EQ(round_tripped.results(0).result().scratch_bytes(), 1);
+  EXPECT_EQ(round_tripped.results(1).device(), "dev2");
+  EXPECT_EQ(round_tripped.results(1).hlo(), "hlo2");
+  EXPECT_EQ(round_tripped.results(1).result().scratch_bytes(), 2);
+  EXPECT_EQ(round_tripped.results(2).device(), "dev3");
+  EXPECT_EQ(round_tripped.results(2).hlo(), "hlo3");
+  EXPECT_EQ(round_tripped.results(2).result().scratch_bytes(), 3);
+}
+
+TEST(AutotuneResultWrapperTest, InconsistentVersions) {
+  std::vector<AutotuneResultWrapper> wrappers =
+      AutotuneResultWrapper::AutotuneResultsToWrappers(
+          ThreeAutotuneEntries(/*version=*/42));
+  auto inconsistent_wrappers = AutotuneResultWrapper::AutotuneResultsToWrappers(
+      ThreeAutotuneEntries(/*version=*/43));
+  wrappers.insert(wrappers.end(), inconsistent_wrappers.begin(),
+                  inconsistent_wrappers.end());
+
+  std::vector<std::pair<AutotuneResultWrapper::OpaqueKey,
+                        AutotuneResultWrapper::OpaqueValue>>
+      key_value_pairs;
+  for (const auto& wrapper : wrappers) {
+    key_value_pairs.push_back(std::make_pair(wrapper.Key(), wrapper.Value()));
+  }
+
+  std::vector<AutotuneResultWrapper> decoded_wrappers;
+  for (const auto& [key, value] : key_value_pairs) {
+    TF_ASSERT_OK_AND_ASSIGN(AutotuneResultWrapper wrapper,
+                            AutotuneResultWrapper::FromKeyAndValue(key, value));
+    decoded_wrappers.push_back(std::move(wrapper));
+  }
+
+  EXPECT_IS_NOT_OK(
+      AutotuneResultWrapper::AutotuneResultsFromWrappers(decoded_wrappers));
+}
+
+}  // namespace
+}  // namespace xla