message: [XLA] Add fields in HloModule for SPMD shardings of input/output, so that we do not need to keep them in the HloInstructions after partitioning.
diff: @@ -300,6 +300,38 @@ class HloModule {
  return &fusion_config_;
  }
  + // Checks if this config has a list of entry parameters' HLO shardings for
  + // SPMD.
  + bool has_spmd_parameters_shardings() const {
  + return spmd_parameters_shardings_.has_value();
  + }
  +
  + // Getter and setter for the list of entry parameters' HLO shardings for SPMD.
  + const std::vector<HloSharding>& spmd_parameters_shardings() const {
  + CHECK(spmd_parameters_shardings_.has_value());
  + return *spmd_parameters_shardings_;
  + }
  + void set_spmd_parameters_shardings(
  + const std::vector<HloSharding>& shardings) {
  + spmd_parameters_shardings_ = shardings;
  + }
  +
  + // Checks if this config has the entry computation output's HLO sharding for
  + // SPMD.
  + bool has_spmd_output_sharding() const {
  + return spmd_output_sharding_.has_value();
  + }
  +
  + // Getter and setter for the entry computation output's HLO shardings for
  + // SPMD.
  + const HloSharding& spmd_output_sharding() const {
  + CHECK(spmd_output_sharding_.has_value());
  + return *spmd_output_sharding_;
  + }
  + void set_spmd_output_sharding(const HloSharding& sharding) {
  + spmd_output_sharding_ = sharding;
  + }
  +
  private:
  HloComputation* AddComputationInternal(
  std::unique_ptr<HloComputation> computation, bool is_entry,
  @@ -342,6 +374,14 @@ class HloModule {
  // Fusion configuration.
  std::vector<std::vector<bool>> fusion_config_;
  +
  + // The HLO shardings of the entry computation's parameters for
  + // SPMD-partitioned programs.
  + absl::optional<std::vector<HloSharding>> spmd_parameters_shardings_;
  +
  + // The HLO sharding of the entry computation's output (root) for
  + // SPMD-partitioned programs.
  + absl::optional<HloSharding> spmd_output_sharding_;
  };
  } // namespace xla
  