From 649ebb90938609d194b5aeeb34d9a010686441d9 Mon Sep 17 00:00:00 2001
From: Alexander Lyashuk <crem@google.com>
Date: Fri, 4 Oct 2024 12:48:28 -0700
Subject: [PATCH] [XLA] Delete unused code intended to handle auto layout in
 HLO.

The code is related to frontend HLO attributes that were supposed to mirror MLIR attributes, but in the end they didn't end up being used and HLO paths (in places where AUTO layout is handled at all) use different method.

This is preparation for more changes that will fix AUTO layout handling in non-MLIR invocations of XLA.

PiperOrigin-RevId: 682418818
---
 xla/pjrt/utils.cc | 52 -----------------------------------------------
 xla/pjrt/utils.h  | 12 -----------
 2 files changed, 64 deletions(-)

diff --git a/xla/pjrt/utils.cc b/xla/pjrt/utils.cc
index 3bc88b6a36dfe..b50fa31b648db 100644
--- a/xla/pjrt/utils.cc
+++ b/xla/pjrt/utils.cc
@@ -384,58 +384,6 @@ absl::StatusOr<std::vector<MemorySpaceColor>> GetOutputMemoryKinds(
 // Make sure to choose delimiter that will never show up in Layout strings.
 static const char* kDelimiter = ";";
 
-static std::string GetFrontendAttr(absl::Span<const LayoutMode> layout_modes) {
-  return absl::StrJoin(layout_modes, kDelimiter,
-                       [](std::string* out, const LayoutMode& mode) {
-                         absl::StrAppend(out, mode.ToString());
-                       });
-}
-
-absl::Status AddLayoutModesToFrontendAttrs(mlir::ModuleOp module,
-                                           XlaComputation& xla_computation) {
-  TF_ASSIGN_OR_RETURN(std::vector<LayoutMode> arg_layout_modes,
-                      GetArgLayoutModes(module));
-  TF_ASSIGN_OR_RETURN(std::vector<LayoutMode> out_layout_modes,
-                      GetOutputLayoutModes(module));
-
-  // Type is string->string proto map. Using auto here to deal with different
-  // build environments.
-  auto& frontend_attrs = *xla_computation.mutable_proto()
-                              ->mutable_frontend_attributes()
-                              ->mutable_map();
-  frontend_attrs["arg_layout_modes"] = GetFrontendAttr(arg_layout_modes);
-  frontend_attrs["out_layout_modes"] = GetFrontendAttr(out_layout_modes);
-  return absl::OkStatus();
-}
-
-static std::string GetFrontendAttrForMemorySpace(
-    const std::vector<MemorySpaceColor>& memory_spaces) {
-  return absl::StrJoin(
-      memory_spaces, kDelimiter,
-      [](std::string* out, const MemorySpaceColor memory_kind) {
-        absl::StrAppend(out, memory_kind);
-      });
-}
-
-absl::Status AddMemoryKindsToFrontendAttrs(mlir::ModuleOp module,
-                                           XlaComputation& xla_computation) {
-  TF_ASSIGN_OR_RETURN(std::vector<MemorySpaceColor> arg_memory_spaces,
-                      GetArgMemoryKinds(module));
-  TF_ASSIGN_OR_RETURN(std::vector<MemorySpaceColor> out_memory_spaces,
-                      GetOutputMemoryKinds(module));
-
-  // Type is string->string proto map. Using auto here to deal with different
-  // build environments.
-  auto& frontend_attrs = *xla_computation.mutable_proto()
-                              ->mutable_frontend_attributes()
-                              ->mutable_map();
-  frontend_attrs["arg_memory_spaces"] =
-      GetFrontendAttrForMemorySpace(arg_memory_spaces);
-  frontend_attrs["out_memory_spaces"] =
-      GetFrontendAttrForMemorySpace(out_memory_spaces);
-  return absl::OkStatus();
-}
-
 static absl::StatusOr<std::vector<LayoutMode>> GetLayoutModesFromFrontendAttr(
     absl::string_view attr) {
   // SkipEmpty() needed to avoid returning the empty string when attr is empty.
diff --git a/xla/pjrt/utils.h b/xla/pjrt/utils.h
index d6aa862876fec..3470bd164d72a 100644
--- a/xla/pjrt/utils.h
+++ b/xla/pjrt/utils.h
@@ -70,18 +70,6 @@ absl::StatusOr<std::vector<MemorySpaceColor>> GetArgMemoryKinds(
 absl::StatusOr<std::vector<MemorySpaceColor>> GetOutputMemoryKinds(
     mlir::ModuleOp module);
 
-// Populates the frontend attributes "arg_layout_mode" and "out_layout_mode" in
-// xla_computation based on `module`. This function must be called before the
-// LayoutMode getters below work correctly on `computation`.
-absl::Status AddLayoutModesToFrontendAttrs(mlir::ModuleOp module,
-                                           XlaComputation& xla_computation);
-
-// Populates the frontend attributes "arg_memory_kinds" and "out_memory_kinds"
-// in xla_computation based on `module`. This function must be called before the
-// LayoutMode getters below work correctly on `computation`.
-absl::Status AddMemoryKindsToFrontendAttrs(mlir::ModuleOp module,
-                                           XlaComputation& xla_computation);
-
 // Returns the LayoutMode for each argument of the computations. Checks for the
 // "arg_layout_mode" frontend attribute, and if not present, assumes
 // LayoutMode::Mode::kDefault.