[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

@@ -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.

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.