text:
#include "absl/algorithm/container.h"
#include "absl/container/flat_hash_map.h"
#include "absl/container/flat_hash_set.h"
#include "absl/types/optional.h"
#include "tensorflow/compiler/xla/literal_util.h"
#include "tensorflow/compiler/xla/service/hlo_computation.h"
#include "tensorflow/compiler/xla/service/hlo_input_output_alias_config.h"
#include "tensorflow/compiler/xla/service/hlo_instruction.h"
#include "tensorflow/compiler/xla/service/hlo_instructions.h"
#include "tensorflow/compiler/xla/service/hlo_opcode.h"
#include "tensorflow/compiler/xla/service/hlo_sharding.h"
#include "tensorflow/compiler/xla/service/hlo_sharding_util.h"
#include "tensorflow/compiler/xla/service/shape_inference.h"
#include "tensorflow/compiler/xla/service/spmd/convolution_handler.h"
#include "tensorflow/compiler/xla/service/spmd/spmd_partitioner.h"
#include "tensorflow/compiler/xla/service/spmd/spmd_partitioner_util.h"
#include "tensorflow/compiler/xla/shape_util.h"
#include "tensorflow/compiler/xla/util.h"
#include "tensorflow/compiler/xla/window_util.h"
#include "tensorflow/compiler/xla/xla_data.pb.h"
#include "tensorflow/core/lib/gtl/cleanup.h"
#include "tensorflow/core/platform/numbers.h"

namespace xla {
namespace spmd {

Status SpmdPartitioningVisitor::HandleDot(HloInstruction* hlo) {
  DotConvDimsMapping mapping;
  const auto& dnums = hlo->dot_dimension_numbers();
  int64 next_output_dim = 0;
  for (int64 i = 0; i < dnums.lhs_batch_dimensions_size(); ++i) {
    mapping.batch_dims.emplace_back();
    mapping.batch_dims.back().lhs = dnums.lhs_batch_dimensions(i);
    mapping.batch_dims.back().rhs = dnums.rhs_batch_dimensions(i);
    mapping.batch_dims.back().output = next_output_dim++;
  }
  for (int64 i = 0; i < dnums.lhs_contracting_dimensions_size(); ++i) {
    mapping.contracting_dims.emplace_back();
    mapping.contracting_dims.back().lhs = dnums.lhs_contracting_dimensions(i);
    mapping.contracting_dims.back().rhs = dnums.rhs_contracting_dimensions(i);
    mapping.contracting_dims.back().output = -1;
  }
  for (int64 i = 0; i < hlo->operand(0)->shape().rank(); ++i) {
    if (absl::c_linear_search(dnums.lhs_batch_dimensions(), i) ||
        absl::c_linear_search(dnums.lhs_contracting_dimensions(), i)) {
      continue;
    }
    mapping.lhs_non_contracting_dims.emplace_back();
    mapping.lhs_non_contracting_dims.back().lhs = i;
    mapping.lhs_non_contracting_dims.back().rhs = -1;
    mapping.lhs_non_contracting_dims.back().output = next_output_dim++;
  }
  for (int64 i = 0; i < hlo->operand(1)->shape().rank(); ++i) {
    if (absl::c_linear_search(dnums.rhs_batch_dimensions(), i) ||
        absl::c_linear_search(dnums.rhs_contracting_dimensions(), i)) {
      continue;
    }
    mapping.rhs_non_contracting_dims.emplace_back();
    mapping.rhs_non_contracting_dims.back().lhs = -1;
    mapping.rhs_non_contracting_dims.back().rhs = i;
    mapping.rhs_non_contracting_dims.back().output = next_output_dim++;
  }
  auto create_sharded_dot =
      [&](HloInstruction* l, HloInstruction* r, SpmdBuilder* b,
          const Window& conv_window) -> StatusOr<HloInstruction*> {
    TF_ASSIGN_OR_RETURN(
        auto sharded_dot_shape,
        ShapeInference::InferDotOpShape(
            l->shape(), r->shape(), hlo->dot_dimension_numbers(),
            /*preferred_element_type=*/hlo->shape().element_type()));
    return b->AddInstruction(HloInstruction::CreateDot(
        sharded_dot_shape, l, r, hlo->dot_dimension_numbers(),
        hlo->precision_config()));
  };
  return HandleDotHelper(hlo, mapping, create_sharded_dot);
}

namespace {

enum class WindowedEinsumOperand { LHS, RHS };

struct WindowedEinsumConfig {
  WindowedEinsumOperand windowed_op;
  bool windowed_at_contracting_dims;
  bool windowed_at_batch_dims;
  bool operands_sharded_at_contracting_dims;
};

struct DotDimensionIndexMapping {
  std::vector<int64> lhs_to_rhs_indices;
  std::vector<int64> lhs_to_output_indices;
  std::vector<int64> rhs_to_lhs_indices;
  std::vector<int64> rhs_to_output_indices;
  std::vector<int64> output_to_lhs_indices;
  std::vector<int64> output_to_rhs_indices;
};

void UpdateDDNums(DotDimensionNumbers* new_ddnums, int64 reshaped_dim,
                  bool lhs) {
  auto update_dims =
      [&reshaped_dim](tensorflow::protobuf::RepeatedField<int64>* dims) {
        for (int64 i = 0; i < dims->size(); ++i) {
          auto dim = dims->at(i);
          if (reshaped_dim <= dim) {
            dims->Set(i, dim + 1);
          }
        }
        if (absl::c_linear_search(*dims, reshaped_dim)) {
          dims->Add(reshaped_dim);
        }
      };

  if (lhs) {
    update_dims(new_ddnums->mutable_lhs_contracting_dimensions());
    update_dims(new_ddnums->mutable_lhs_batch_dimensions());
  } else {  // rhs
    update_dims(new_ddnums->mutable_rhs_contracting_dimensions());
    update_dims(new_ddnums->mutable_rhs_batch_dimensions());
  }
}

Window GenNewWindow(const HloInstruction* original_dot,
                    const HloInstruction* dot_lhs,
                    const HloInstruction* dot_rhs, int64 lhs_concat_dim,
                    int64 rhs_concat_dim, bool windowed_at_contracting_dims,
                    bool windowed_at_batch_dims) {
  auto new_window = original_dot->window();
  const ConvolutionDimensionNumbers& conv_dnums =
      original_dot->convolution_dimension_numbers();
  if (lhs_concat_dim != -1) {
    for (int64 i = 0; i < conv_dnums.input_spatial_dimensions_size(); ++i) {
      if (conv_dnums.input_spatial_dimensions(i) == lhs_concat_dim) {
        auto wd = new_window.mutable_dimensions(i);
        auto lhs_size = dot_lhs->shape().dimensions(lhs_concat_dim + 1);
        if (windowed_at_contracting_dims) {
          wd->set_size(lhs_size);
        }
        if (windowed_at_batch_dims) {
          wd->set_size(lhs_size);
          wd->set_padding_low(0);
          wd->set_padding_high(0);
          wd->set_stride(std::max<int64>(1, lhs_size - 1));
          wd->set_window_dilation(1);
          wd->set_base_dilation(lhs_size);
          wd->set_window_reversal(false);
        }
      }
    }
  }
  if (rhs_concat_dim != -1) {
    for (int64 i = 0; i < conv_dnums.kernel_spatial_dimensions_size(); ++i) {
      if (conv_dnums.kernel_spatial_dimensions(i) == rhs_concat_dim &&
          !windowed_at_contracting_dims && !windowed_at_batch_dims &&
          lhs_concat_dim == -1) {
        auto wd = new_window.mutable_dimensions(i);
        auto rhs_size = dot_rhs->shape().dimensions(rhs_concat_dim + 1);
        wd->set_size(rhs_size);
        wd->set_padding_low(rhs_size - 1);
        wd->set_padding_high(rhs_size - 1);
      }
    }
  }
  // Add the extra dimension to window.
  WindowDimension* new_dim = new_window.add_dimensions();
  if (windowed_at_contracting_dims) {
    new_dim->set_size(2);
    new_dim->set_padding_low(0);
    new_dim->set_padding_high(0);
    new_dim->set_stride(1);
    new_dim->set_window_dilation(1);
    new_dim->set_base_dilation(1);
    new_dim->set_window_reversal(false);
  } else if (windowed_at_batch_dims) {
    new_dim->set_size(2);
    new_dim->set_padding_low(0);
    new_dim->set_padding_high(0);
    new_dim->set_stride(1);  // std::max<int64>(1, 2 - 1)
    new_dim->set_window_dilation(1);
    new_dim->set_base_dilation(2);
    new_dim->set_window_reversal(false);
  } else {
    if (lhs_concat_dim != -1) {
      new_dim->set_size(1);
      new_dim->set_padding_low(0);
      new_dim->set_padding_high(0);
      new_dim->set_stride(1);
      new_dim->set_window_dilation(1);
      new_dim->set_base_dilation(1);
      new_dim->set_window_reversal(false);
    }
    if (rhs_concat_dim != -1) {
      new_dim->set_size(2);          // rhs_size
      new_dim->set_padding_low(1);   // rhs_size - 1
      new_dim->set_padding_high(1);  // rhs_size - 1
      new_dim->set_stride(1);
      new_dim->set_window_dilation(1);
      new_dim->set_base_dilation(1);
      new_dim->set_window_reversal(true);
    }
  }

  VLOG(2) << "new_window: " << new_window.ShortDebugString();
  return new_window;
}

ConvolutionDimensionNumbers GenNewConvDNums(
    const HloInstruction* original_dot, const HloInstruction* dot_lhs,
    const HloInstruction* dot_rhs, int64 lhs_concat_dim, int64 rhs_concat_dim,
    bool windowed_at_contracting_dims, bool windowed_at_batch_dims,
    const std::vector<int64>& lhs_to_output_indices,
    const std::vector<int64>& rhs_to_output_indices,
    const Shape& new_dot_shape) {
  // Generate the new conv dimension numbers.
  const ConvolutionDimensionNumbers& dnums =
      original_dot->convolution_dimension_numbers();
  // Handle the LHS dimension numbers.
  int64 input_batch_dimension = dnums.input_batch_dimension();
  int64 input_feature_dimension = dnums.input_feature_dimension();
  std::vector<int64> input_spatial_dimensions(
      dnums.input_spatial_dimensions().begin(),
      dnums.input_spatial_dimensions().end());
  if (lhs_concat_dim != -1) {
    if (lhs_concat_dim <= input_batch_dimension) {
      input_batch_dimension++;
    }
    if (lhs_concat_dim <= input_feature_dimension) {
      input_feature_dimension++;
    }
    for (int64 i = 0; i < input_spatial_dimensions.size(); ++i) {
      if (lhs_concat_dim <= input_spatial_dimensions[i]) {
        input_spatial_dimensions[i]++;
      }
    }
    input_spatial_dimensions.push_back(lhs_concat_dim);
  }
  if (rhs_concat_dim != -1 && !windowed_at_contracting_dims &&
      !windowed_at_batch_dims) {
    input_spatial_dimensions.push_back(dot_lhs->shape().dimensions_size() - 1);
  }
  // Handle the RHS dimension numbers.
  int64 kernel_input_feature_dimension = dnums.kernel_input_feature_dimension();
  int64 kernel_output_feature_dimension =
      dnums.kernel_output_feature_dimension();
  std::vector<int64> kernel_spatial_dimensions(
      dnums.kernel_spatial_dimensions().begin(),
      dnums.kernel_spatial_dimensions().end());
  if (rhs_concat_dim != -1) {
    if (rhs_concat_dim <= kernel_input_feature_dimension) {
      kernel_input_feature_dimension++;
    }
    if (rhs_concat_dim <= kernel_output_feature_dimension) {
      kernel_output_feature_dimension++;
    }
    for (int64 i = 0; i < kernel_spatial_dimensions.size(); ++i) {
      if (rhs_concat_dim <= kernel_spatial_dimensions[i]) {
        kernel_spatial_dimensions[i]++;
      }
    }
    kernel_spatial_dimensions.push_back(rhs_concat_dim);
  }
  if (lhs_concat_dim != -1 && !windowed_at_contracting_dims &&
      !windowed_at_batch_dims) {
    kernel_spatial_dimensions.push_back(dot_rhs->shape().dimensions_size() - 1);
  }
  // Handle the Output dimension numbers.
  int64 output_batch_dimension = dnums.output_batch_dimension();
  int64 output_feature_dimension = dnums.output_feature_dimension();
  std::vector<int64> output_spatial_dimensions(
      dnums.output_spatial_dimensions().begin(),
      dnums.output_spatial_dimensions().end());
  if (!windowed_at_contracting_dims) {
    auto output_slice_dim = lhs_concat_dim != -1
                                ? lhs_to_output_indices[lhs_concat_dim]
                                : rhs_to_output_indices[rhs_concat_dim];
    if (output_slice_dim <= output_batch_dimension) {
      output_batch_dimension++;
    }
    if (output_slice_dim <= output_feature_dimension) {
      output_feature_dimension++;
    }
    for (int64 i = 0; i < output_spatial_dimensions.size(); ++i) {
      if (output_slice_dim <= output_spatial_dimensions[i]) {
        output_spatial_dimensions[i]++;
      }
    }
    output_spatial_dimensions.push_back(output_slice_dim);
  } else {
    output_spatial_dimensions.push_back(new_dot_shape.dimensions_size() - 1);
  }
  // Construct the new dot dimension numbers.
  ConvolutionDimensionNumbers new_dnums;
  new_dnums.set_input_batch_dimension(input_batch_dimension);
  new_dnums.set_input_feature_dimension(input_feature_dimension);
  for (auto dim : input_spatial_dimensions) {
    new_dnums.add_input_spatial_dimensions(dim);
  }
  new_dnums.set_kernel_input_feature_dimension(kernel_input_feature_dimension);
  new_dnums.set_kernel_output_feature_dimension(
      kernel_output_feature_dimension);
  for (auto dim : kernel_spatial_dimensions) {
    new_dnums.add_kernel_spatial_dimensions(dim);
  }
  new_dnums.set_output_batch_dimension(output_batch_dimension);
  new_dnums.set_output_feature_dimension(output_feature_dimension);
  for (auto dim : output_spatial_dimensions) {
    new_dnums.add_output_spatial_dimensions(dim);
  }

  return new_dnums;
}

DotDimensionIndexMapping ComputeDimensionIndexMapping(
    const DotConvDimsMapping& dims_mapping, int64 lhs_rank, int64 rhs_rank,
    int64 output_rank) {
  std::vector<int64> lhs_to_rhs_indices(lhs_rank, -1);
  std::vector<int64> lhs_to_output_indices(lhs_rank, -1);
  std::vector<int64> rhs_to_lhs_indices(rhs_rank, -1);
  std::vector<int64> rhs_to_output_indices(rhs_rank, -1);
  std::vector<int64> output_to_lhs_indices(output_rank, -1);
  std::vector<int64> output_to_rhs_indices(output_rank, -1);
  auto populate_indices_mapping =
      [&](const DotConvDimsMapping::DimsMapping& mapping) {
        if (mapping.lhs >= 0) {
          lhs_to_rhs_indices[mapping.lhs] = mapping.rhs;
          lhs_to_output_indices[mapping.lhs] = mapping.output;
        }
        if (mapping.rhs >= 0) {
          rhs_to_lhs_indices[mapping.rhs] = mapping.lhs;
          rhs_to_output_indices[mapping.rhs] = mapping.output;
        }
        if (mapping.output >= 0) {
          output_to_lhs_indices[mapping.output] = mapping.lhs;
          output_to_rhs_indices[mapping.output] = mapping.rhs;
        }
      };
  for (const auto& mapping : dims_mapping.batch_dims) {
    populate_indices_mapping(mapping);
  }
  for (const auto& mapping : dims_mapping.contracting_dims) {
    populate_indices_mapping(mapping);
  }
  for (const auto& mapping : dims_mapping.lhs_non_contracting_dims) {
    populate_indices_mapping(mapping);
  }
  for (const auto& mapping : dims_mapping.rhs_non_contracting_dims) {
    populate_indices_mapping(mapping);
  }
  for (const auto& mapping : dims_mapping.conv_spatial_dims) {
    populate_indices_mapping(mapping);
  }
  return DotDimensionIndexMapping{lhs_to_rhs_indices,    lhs_to_output_indices,
                                  rhs_to_lhs_indices,    rhs_to_output_indices,
                                  output_to_lhs_indices, output_to_rhs_indices};
}

absl::optional<WindowedEinsumConfig> GetWindowedEinsumConfiguration(
    int64 num_partitions, int64 output_lhs_non_contracting_partitions,
    int64 output_rhs_non_contracting_partitions,
    int64 rhs_contracting_partitions, int64 rhs_non_contracting_partitions,
    int64 rhs_batch_partitions, int64 lhs_contracting_partitions,
    int64 lhs_non_contracting_partitions, int64 lhs_batch_partitions,
    int64 rhs_shape_size, int64 lhs_shape_size, int64 output_shape_size,
    int64 einsum_threshold_mib,
    const absl::optional<HloSharding>& output_sharding_transposed_to_match_lhs,
    const absl::optional<HloSharding>& output_sharding_transposed_to_match_rhs,
    const HloSharding& lhs_sharding, const HloSharding& rhs_sharding) {
  if (output_lhs_non_contracting_partitions == num_partitions &&
      output_sharding_transposed_to_match_lhs == lhs_sharding &&
      rhs_shape_size >= einsum_threshold_mib * 1024 * 1024) {
    if (rhs_contracting_partitions == num_partitions) {
      return WindowedEinsumConfig{
          /*windowed_op=*/WindowedEinsumOperand::RHS,
          /*windowed_at_contracting_dims*/ true,
          /*windowed_at_batch_dims=*/false,
          /*operands_sharded_at_contracting_dims=*/false};
    }
    if (rhs_non_contracting_partitions == num_partitions) {
      return WindowedEinsumConfig{
          /*windowed_op=*/WindowedEinsumOperand::RHS,
          /*windowed_at_contracting_dims*/ false,
          /*windowed_at_batch_dims=*/false,
          /*operands_sharded_at_contracting_dims=*/false};
    }
    if (rhs_batch_partitions == num_partitions) {
      return WindowedEinsumConfig{
          /*windowed_op=*/WindowedEinsumOperand::RHS,
          /*windowed_at_contracting_dims*/ false,
          /*windowed_at_batch_dims=*/true,
          /*operands_sharded_at_contracting_dims=*/false};
    }
  }
  if (output_rhs_non_contracting_partitions == num_partitions &&
      output_sharding_transposed_to_match_rhs == rhs_sharding &&
      lhs_shape_size >= einsum_threshold_mib * 1024 * 1024) {
    if (lhs_contracting_partitions == num_partitions) {
      return WindowedEinsumConfig{
          /*windowed_op=*/WindowedEinsumOperand::LHS,
          /*windowed_at_contracting_dims*/ true,
          /*windowed_at_batch_dims=*/false,
          /*operands_sharded_at_contracting_dims=*/false};
    }
    if (lhs_non_contracting_partitions == num_partitions) {
      return WindowedEinsumConfig{
          /*windowed_op=*/WindowedEinsumOperand::LHS,
          /*windowed_at_contracting_dims*/ false,
          /*windowed_at_batch_dims=*/false,
          /*operands_sharded_at_contracting_dims=*/false};
    }
    if (lhs_batch_partitions == num_partitions) {
      return WindowedEinsumConfig{
          /*windowed_op=*/WindowedEinsumOperand::LHS,
          /*windowed_at_contracting_dims*/ false,
          /*windowed_at_batch_dims=*/true,
          /*operands_sharded_at_contracting_dims=*/false};
    }
  }
  if (lhs_contracting_partitions == rhs_contracting_partitions &&
      lhs_contracting_partitions == num_partitions &&
      (output_lhs_non_contracting_partitions == num_partitions ||
       output_rhs_non_contracting_partitions == num_partitions) &&
      output_shape_size >= einsum_threshold_mib * 1024 * 1024) {
    if (output_lhs_non_contracting_partitions == num_partitions) {
      return WindowedEinsumConfig{
          /*windowed_op=*/WindowedEinsumOperand::RHS,
          /*windowed_at_contracting_dims*/ false,
          /*windowed_at_batch_dims=*/false,
          /*operands_sharded_at_contracting_dims=*/true};
    }
    if (output_rhs_non_contracting_partitions == num_partitions) {
      return WindowedEinsumConfig{
          /*windowed_op=*/WindowedEinsumOperand::LHS,
          /*windowed_at_contracting_dims*/ false,
          /*windowed_at_batch_dims=*/false,
          /*operands_sharded_at_contracting_dims=*/true};
    }
  }
  return absl::nullopt;
}

// We use a recursive approach where sets of matching dimensions are recognized
// one at a time. The base shapes and shardings can be changed during the
// recursion as we group devices together. So refer to the passed in shapes and
// shardings for inputs and output, and do not use shape inference.

StatusOr<HloInstruction*> PartitionBaseCase(
    PartitionedHlo lhs, PartitionedHlo rhs, const Shape& output_base_shape,
    const HloSharding& output_sharding, const DotConvDimsMapping& dims_mapping,
    int64 num_partitions,
    const std::function<StatusOr<HloInstruction*>(
        HloInstruction*, HloInstruction*, SpmdBuilder*,
        const Window& conv_window)>& create_sharded_dot,
    const Window& conv_window, HloModule* module, HloInstruction* original_hlo,
    int64 lhs_batch_partitions, int64 rhs_batch_partitions,
    int64 output_batch_partitions, int64 lhs_contracting_partitions,
    int64 rhs_contracting_partitions, int64 lhs_non_contracting_partitions,
    int64 rhs_non_contracting_partitions,
    int64 output_lhs_non_contracting_partitions,
    int64 output_rhs_non_contracting_partitions,
    const SpmdPartitionerOptions& options, SpmdBuilder* b,
    std::vector<SpmdPartitioningVisitor::WindowedDotGeneralLoop>*
        windowed_dot_general_loops,
    bool may_reshard_without_detecting_match) {
  const HloSharding& lhs_sharding = lhs.sharding();
  const HloSharding& rhs_sharding = rhs.sharding();
  if (lhs_sharding.ReplicateOnLastTileDim() ||
      rhs_sharding.ReplicateOnLastTileDim() ||
      output_sharding.ReplicateOnLastTileDim()) {
    return nullptr;
  }
  DotDimensionIndexMapping indices_map = ComputeDimensionIndexMapping(
      dims_mapping, lhs.base_shape().rank(), rhs.base_shape().rank(),
      output_base_shape.rank());
  auto lhs_sharding_transposed_to_match_rhs =
      hlo_sharding_util::TransposeShardingWithCollapsedDims(
          lhs_sharding, indices_map.lhs_to_rhs_indices,
          indices_map.rhs_to_lhs_indices);
  auto rhs_sharding_transposed_to_match_lhs =
      hlo_sharding_util::TransposeShardingWithCollapsedDims(
          rhs_sharding, indices_map.rhs_to_lhs_indices,
          indices_map.lhs_to_rhs_indices);
  auto lhs_sharding_transposed_to_match_output =
      hlo_sharding_util::TransposeShardingWithCollapsedDims(
          lhs_sharding, indices_map.lhs_to_output_indices,
          indices_map.output_to_lhs_indices);
  auto rhs_sharding_transposed_to_match_output =
      hlo_sharding_util::TransposeShardingWithCollapsedDims(
          rhs_sharding, indices_map.rhs_to_output_indices,
          indices_map.output_to_rhs_indices);
  auto output_sharding_transposed_to_match_lhs =
      hlo_sharding_util::TransposeShardingWithCollapsedDims(
          output_sharding, indices_map.output_to_lhs_indices,
          indices_map.lhs_to_output_indices);
  auto output_sharding_transposed_to_match_rhs =
      hlo_sharding_util::TransposeShardingWithCollapsedDims(
          output_sharding, indices_map.output_to_rhs_indices,
          indices_map.rhs_to_output_indices);

  // LHS and RHS are partitioned the same way and only partitioned in batch
  // dimensions.
  if (lhs_batch_partitions == rhs_batch_partitions &&
      rhs_batch_partitions == num_partitions &&
      lhs_sharding_transposed_to_match_rhs == rhs_sharding) {
    TF_ASSIGN_OR_RETURN(
        auto dot, create_sharded_dot(lhs.hlo(), rhs.hlo(), b, conv_window));
    dot->set_sharding(*lhs_sharding_transposed_to_match_output);
    return PartitionedHlo(dot, output_base_shape, lhs.state())
        .Reshard(output_sharding)
        .hlo();
  }

  // Try emit batch-partitioned einsum with one operand resharded. Returns
  // partitioned HLO or nullptr if the attempt fails. If
  // may_reshard_with_allreduce is false, reshard must be done using
  // all-to-all/collective-permute; otherwise this attempt fails.
  auto try_emit_output_batch_partitioned_einsum_with_reshard =
      [&](bool may_reshard_with_allreduce) -> StatusOr<HloInstruction*> {
    // LHS and output are batch partitioned in the same way.
    if (lhs_batch_partitions == num_partitions &&
        output_batch_partitions == num_partitions &&
        lhs_sharding_transposed_to_match_output == output_sharding) {
      if (!may_reshard_with_allreduce &&
          !CanReshardWithCollectivePermute(
              rhs.sharding(), *lhs_sharding_transposed_to_match_rhs) &&
          !GetReshardAllToAllSourceTargetDims(
              rhs.sharding(), *lhs_sharding_transposed_to_match_rhs)) {
        return nullptr;
      }
      auto resharded_rhs = rhs.Reshard(*lhs_sharding_transposed_to_match_rhs);
      TF_ASSIGN_OR_RETURN(
          auto dot,
          create_sharded_dot(lhs.hlo(), resharded_rhs.hlo(), b, conv_window));
      return dot;
    }
    // RHS and output are batch partitioned in the same way.
    if (rhs_batch_partitions == num_partitions &&
        output_batch_partitions == num_partitions &&
        rhs_sharding_transposed_to_match_output == output_sharding) {
      if (!may_reshard_with_allreduce &&
          !CanReshardWithCollectivePermute(
              lhs.sharding(), *rhs_sharding_transposed_to_match_lhs) &&
          !GetReshardAllToAllSourceTargetDims(
              lhs.sharding(), *rhs_sharding_transposed_to_match_lhs)) {
        return nullptr;
      }
      auto resharded_lhs = lhs.Reshard(*rhs_sharding_transposed_to_match_lhs);
      TF_ASSIGN_OR_RETURN(
          auto dot,
          create_sharded_dot(resharded_lhs.hlo(), rhs.hlo(), b, conv_window));
      return dot;
    }
    return nullptr;
  };

  {
    // Try batch-parallel by resharding one operand, and not using all-reduce.
    TF_ASSIGN_OR_RETURN(
        HloInstruction * partitioned_dot,
        try_emit_output_batch_partitioned_einsum_with_reshard(false));
    if (partitioned_dot) {
      return partitioned_dot;
    }
  }

  // Try to emit windowed DotGeneral when one operand is partitioned in the same
  // way as the output along non-contracting dimensions, but the other operand
  // is tiled in other dimensions. Or both operands are partitioned in the same
  // way along contracting dimensions, but the output is partitioned along
  // non-contracting dimensions.
  auto emit_windowed_dot_general =
      [&](const WindowedEinsumConfig& einsum_config)
      -> StatusOr<HloInstruction*> {
    CHECK(!einsum_config.windowed_at_batch_dims ||
          !einsum_config.windowed_at_contracting_dims);
    const bool windowed_at_batch_dims = einsum_config.windowed_at_batch_dims;
    const bool windowed_at_contracting_dims =
        einsum_config.windowed_at_contracting_dims;
    const bool operands_sharded_at_contracting_dims =
        einsum_config.operands_sharded_at_contracting_dims;
    auto unpadded_result_buffer_shape =
        MakePartitionedShape(output_base_shape, output_sharding);
    auto padded_result_buffer_shape = unpadded_result_buffer_shape;
    const bool windowed_op_is_lhs =
        einsum_config.windowed_op == WindowedEinsumOperand::LHS;
    // For windowing at batch/non-contracting dims, we produce the result one
    // partition at a time, so we need to pad the shape in case of uneven
    // partitioning in order to make dynamic-update-slice in-bound.
    if (!windowed_at_contracting_dims &&
        !operands_sharded_at_contracting_dims) {
      padded_result_buffer_shape = GetPaddedShapeForUnevenPartitioning(
          padded_result_buffer_shape,
          windowed_op_is_lhs ? *lhs_sharding_transposed_to_match_output
                             : *rhs_sharding_transposed_to_match_output);
    }
    // Mask the padding area of the windowed operand with zero if there is
    // uneven partitioning.
    if (windowed_at_contracting_dims) {
      auto& to_mask = windowed_op_is_lhs ? lhs : rhs;
      to_mask =
          to_mask.PadWithValue(b->AddInstruction(HloInstruction::CreateConstant(
              LiteralUtil::Zero(output_base_shape.element_type()))));
    }
    if (operands_sharded_at_contracting_dims) {
      auto zero = b->AddInstruction(HloInstruction::CreateConstant(
          LiteralUtil::Zero(output_base_shape.element_type())));
      lhs = lhs.PadWithValue(zero);
      rhs = rhs.PadWithValue(zero);
    }
    auto result_buffer = CreateZero(padded_result_buffer_shape, b);
    auto extra_buffer =
        (!(options.bidirectional_windowed_einsum && num_partitions % 4 == 0) ||
         operands_sharded_at_contracting_dims)
            ? CreateZero(padded_result_buffer_shape, b)
        : windowed_op_is_lhs ? lhs.hlo()
                             : rhs.hlo();

    if (options.bidirectional_windowed_einsum && num_partitions % 4 == 0 &&
        !operands_sharded_at_contracting_dims) {
      std::vector<std::pair<int64, int64>> pre_sd_pairs(num_partitions);
      for (int64 source = 0; source < num_partitions; ++source) {
        // 0 -> 1, 1 -> 2, 2 -> 3, ...
        pre_sd_pairs[source] = {source, (source + 1) % num_partitions};
      }
      extra_buffer =
          lhs.state()
              .collective_ops_creator.create_cross_partition_collective_permute(
                  b, extra_buffer, pre_sd_pairs,
                  (*lhs.state().next_channel_id)++);
    }

    auto iteration = b->AddInstruction(
        HloInstruction::CreateConstant(LiteralUtil::CreateR0<uint32>(0)));

    // Create a while loop that computes one window per iteration. During each
    // iteration, each partition sends its input window to its neighbor using
    // collective-permute for the next iteration.
    SpmdBuilder body_b("windowed_dot_general_body", original_hlo);

    // Generate partial results used by bidirectional algorithm.
    auto get_partial_bid_results =
        [&](HloInstruction* l, HloInstruction* r, HloInstruction* o,
            HloInstruction* extra_inout, HloInstruction* cw_cp_output,
            HloInstruction* i) -> StatusOr<std::vector<HloInstruction*>> {
      auto partition_id =
          lhs.state().collective_ops_creator.create_partition_id(&body_b);
      auto partition_count =
          body_b.AddInstruction(HloInstruction::CreateConstant(
              LiteralUtil::CreateR0<uint32>(num_partitions)));
      auto ccw_data_partition_id =
          body_b.AddInstruction(HloInstruction::CreateBinary(
              i->shape(), HloOpcode::kAdd, i, partition_id));
      auto cw_data_partition_id =
          body_b.AddInstruction(HloInstruction::CreateBinary(
              i->shape(), HloOpcode::kAdd, partition_count, partition_id));
      if (operands_sharded_at_contracting_dims) {
        ccw_data_partition_id =
            body_b.AddInstruction(HloInstruction::CreateBinary(
                i->shape(), HloOpcode::kAdd, ccw_data_partition_id,
                body_b.AddInstruction(HloInstruction::CreateConstant(
                    LiteralUtil::CreateR0<uint32>(num_partitions / 2 + 1)))));
        cw_data_partition_id =
            body_b.AddInstruction(HloInstruction::CreateBinary(
                i->shape(), HloOpcode::kSubtract, cw_data_partition_id,
                body_b.AddInstruction(HloInstruction::CreateConstant(
                    LiteralUtil::CreateR0<uint32>(num_partitions / 2)))));
      } else {
        cw_data_partition_id =
            body_b.AddInstruction(HloInstruction::CreateBinary(
                i->shape(), HloOpcode::kSubtract, cw_data_partition_id,
                CreateOne(cw_data_partition_id->shape(), &body_b)));
      }
      ccw_data_partition_id = body_b.AddInstruction(
          HloInstruction::CreateBinary(i->shape(), HloOpcode::kRemainder,
                                       ccw_data_partition_id, partition_count));
      cw_data_partition_id = body_b.AddInstruction(HloInstruction::CreateBinary(
          i->shape(), HloOpcode::kSubtract, cw_data_partition_id, i));
      cw_data_partition_id = body_b.AddInstruction(
          HloInstruction::CreateBinary(i->shape(), HloOpcode::kRemainder,
                                       cw_data_partition_id, partition_count));
      // Calculate concat dim.
      const HloSharding* slice_sharding;
      if (operands_sharded_at_contracting_dims) {
        slice_sharding = windowed_op_is_lhs
                             ? &*output_sharding_transposed_to_match_rhs
                             : &*output_sharding_transposed_to_match_lhs;
      } else if (windowed_at_contracting_dims || windowed_at_batch_dims) {
        slice_sharding = windowed_op_is_lhs
                             ? &*lhs_sharding_transposed_to_match_rhs
                             : &*rhs_sharding_transposed_to_match_lhs;
      } else {
        slice_sharding = windowed_op_is_lhs
                             ? &*lhs_sharding_transposed_to_match_output
                             : &*rhs_sharding_transposed_to_match_output;
      }
      CHECK_EQ(Product(slice_sharding->tile_assignment().dimensions()),
               num_partitions);
      int64 slice_sharding_dim = -1;
      for (int64 i = 0; i < slice_sharding->tile_assignment().num_dimensions();
           ++i) {
        if (slice_sharding->tile_assignment().dim(i) > 1) {
          slice_sharding_dim = i;
          break;
        }
      }
      int64 lhs_concat_dim = -1;
      int64 rhs_concat_dim = -1;
      if (operands_sharded_at_contracting_dims) {
        if (windowed_op_is_lhs) {
          rhs_concat_dim = slice_sharding_dim;
        } else {
          lhs_concat_dim = slice_sharding_dim;
        }
      } else if (windowed_at_contracting_dims || windowed_at_batch_dims) {
        lhs_concat_dim =
            windowed_op_is_lhs
                ? indices_map.rhs_to_lhs_indices[slice_sharding_dim]
                : slice_sharding_dim;
        rhs_concat_dim =
            windowed_op_is_lhs
                ? slice_sharding_dim
                : indices_map.lhs_to_rhs_indices[slice_sharding_dim];
      } else {
        if (windowed_op_is_lhs) {
          lhs_concat_dim =
              indices_map.output_to_lhs_indices[slice_sharding_dim];
        } else {
          rhs_concat_dim =
              indices_map.output_to_rhs_indices[slice_sharding_dim];
        }
      }

      DotDimensionNumbers new_ddnums;
      if (original_hlo->opcode() == HloOpcode::kDot) {
        new_ddnums = original_hlo->dot_dimension_numbers();
      }

      auto dot_lhs = l;
      auto dot_rhs = r;
      auto original_dot_lhs = l;
      auto original_dot_rhs = r;
      if (windowed_at_contracting_dims || windowed_at_batch_dims ||
          operands_sharded_at_contracting_dims) {
        // Slice the matching operand according to the partitioned dimensions
        // on the windowed operand or the output.
        auto slice_operand = !windowed_op_is_lhs ? l : r;

        // Pad the sharding dim first (then the concat dim) for correctness.
        auto sharding_dim_size =
            slice_operand->shape().dimensions(slice_sharding_dim);
        if (sharding_dim_size % num_partitions != 0) {
          slice_operand = PadBaseShapeBeforeUnevenTiledSharding(
              slice_operand, *slice_sharding, &body_b);
        }

        // We do this by treating the matching operand as replicated, and
        // resharding it to match the windowed operand or the output.
        auto gen_slice = [&](HloInstruction* data_partition_id,
                             bool ccw) -> HloInstruction* {
          std::vector<int64> new_dims;
          for (int64 i = 0; i < slice_operand->shape().dimensions_size(); ++i) {
            if (i == slice_sharding_dim) {
              new_dims.push_back(1);
            }
            new_dims.push_back(slice_operand->shape().dimensions(i));
          }
          auto reshaped_slice_operand =
              body_b.AddInstruction(HloInstruction::CreateReshape(
                  ShapeUtil::MakeShape(slice_operand->shape().element_type(),
                                       new_dims),
                  slice_operand));
          auto min = body_b.AddInstruction(
              HloInstruction::CreateConstant(LiteralUtil::MinValue(
                  reshaped_slice_operand->shape().element_type())));
          std::vector<int64> min_padding(
              reshaped_slice_operand->shape().rank());
          auto padded_slice_operand = reshaped_slice_operand;
          auto padded_shape = padded_slice_operand->shape();
          int64 padding_dim = slice_sharding_dim;
          padded_shape.set_dimensions(padding_dim, 2);
          if (ccw) {
            // ccw pad high
            PaddingConfig ccw_pad_config =
                window_util::MakeSymmetricPadding(min_padding);
            ccw_pad_config.mutable_dimensions(padding_dim)
                ->set_edge_padding_low(0);
            ccw_pad_config.mutable_dimensions(padding_dim)
                ->set_edge_padding_high(1);
            padded_slice_operand =
                body_b.AddInstruction(HloInstruction::CreatePad(
                    padded_shape, padded_slice_operand, min, ccw_pad_config));
          } else {
            // cw pad low
            PaddingConfig cw_pad_config =
                window_util::MakeSymmetricPadding(min_padding);
            cw_pad_config.mutable_dimensions(padding_dim)
                ->set_edge_padding_low(1);
            cw_pad_config.mutable_dimensions(padding_dim)
                ->set_edge_padding_high(0);
            padded_slice_operand =
                body_b.AddInstruction(HloInstruction::CreatePad(
                    padded_shape, padded_slice_operand, min, cw_pad_config));
          }

          padded_slice_operand->set_sharding(HloSharding::Replicate());
          auto state = lhs.state();
          state.b = &body_b;
          state.partition_id = data_partition_id;
          state.reshard_cache->per_hlo_cache.erase(padded_slice_operand);
          auto padded_slice_sharding = hlo_sharding_util::ReshapeSharding(
              slice_operand->shape(), reshaped_slice_operand->shape(),
              *slice_sharding);
          auto padded_slice =
              PartitionedHlo(padded_slice_operand,
                             padded_slice_operand->shape(), state)
                  .Reshard(*padded_slice_sharding)
                  .hlo();
          padded_slice_operand->clear_sharding();
          return padded_slice;
        };

        auto ccw_slice = gen_slice(ccw_data_partition_id, true);
        auto cw_slice = gen_slice(cw_data_partition_id, false);
        auto slice = body_b.AddInstruction(HloInstruction::CreateBinary(
            ccw_slice->shape(), HloOpcode::kMaximum, ccw_slice, cw_slice));
        // Reshape. The reshaped slice will not be used to produce the final
        // result, but used as a hint for the shape inference.
        std::vector<int64> reshaped_slice_dims;
        for (int64 i = 0; i < slice->shape().dimensions_size(); ++i) {
          auto dim_size = slice->shape().dimensions(i);
          if (i == (slice_sharding_dim + 1)) {
            reshaped_slice_dims.push_back(dim_size * 2);
          } else if (i != slice_sharding_dim) {
            reshaped_slice_dims.push_back(dim_size);
          }
        }
        auto reshaped_slice =
            body_b.AddInstruction(HloInstruction::CreateReshape(
                ShapeUtil::MakeShape(slice->shape().element_type(),
                                     reshaped_slice_dims),
                slice));

        if (!windowed_op_is_lhs) {
          dot_lhs = slice;
          original_dot_lhs = reshaped_slice;
          if (original_hlo->opcode() == HloOpcode::kDot) {
            UpdateDDNums(&new_ddnums, slice_sharding_dim, true);
          }
        } else {
          dot_rhs = slice;
          original_dot_rhs = reshaped_slice;
          if (original_hlo->opcode() == HloOpcode::kDot) {
            UpdateDDNums(&new_ddnums, slice_sharding_dim, false);
          }
        }
      }

      auto ccw_dot_lhs = l;
      auto ccw_dot_rhs = r;
      auto cw_dot_lhs = windowed_op_is_lhs ? extra_inout : l;
      auto cw_dot_rhs = windowed_op_is_lhs ? r : extra_inout;
      if (lhs_concat_dim != -1 && windowed_op_is_lhs) {
        auto lhs_concat_shape = ccw_dot_lhs->shape();
        lhs_concat_shape.set_dimensions(
            lhs_concat_dim,
            ccw_dot_lhs->shape().dimensions(lhs_concat_dim) * 2);
        dot_lhs = body_b.AddInstruction(HloInstruction::CreateConcatenate(
            lhs_concat_shape, {ccw_dot_lhs, cw_dot_lhs}, lhs_concat_dim));
        original_dot_lhs = dot_lhs;

        // Reshape
        std::vector<int64> reshaped_dims(dot_lhs->shape().dimensions().begin(),
                                         dot_lhs->shape().dimensions().end());
        reshaped_dims[lhs_concat_dim] /= 2;
        reshaped_dims.insert(reshaped_dims.begin() + lhs_concat_dim, 2);
        dot_lhs = body_b.AddInstruction(HloInstruction::CreateReshape(
            ShapeUtil::MakeShape(dot_lhs->shape().element_type(),
                                 reshaped_dims),
            dot_lhs));

        if (original_hlo->opcode() == HloOpcode::kDot) {
          UpdateDDNums(&new_ddnums, lhs_concat_dim, true);
        }
      }
      if (rhs_concat_dim != -1 && !windowed_op_is_lhs) {
        auto rhs_concat_shape = ccw_dot_rhs->shape();
        rhs_concat_shape.set_dimensions(
            rhs_concat_dim,
            ccw_dot_rhs->shape().dimensions(rhs_concat_dim) * 2);
        dot_rhs = body_b.AddInstruction(HloInstruction::CreateConcatenate(
            rhs_concat_shape, {ccw_dot_rhs, cw_dot_rhs}, rhs_concat_dim));
        original_dot_rhs = dot_rhs;

        // Reshape
        std::vector<int64> reshaped_dims(dot_rhs->shape().dimensions().begin(),
                                         dot_rhs->shape().dimensions().end());
        reshaped_dims[rhs_concat_dim] /= 2;
        reshaped_dims.insert(reshaped_dims.begin() + rhs_concat_dim, 2);
        dot_rhs = body_b.AddInstruction(HloInstruction::CreateReshape(
            ShapeUtil::MakeShape(dot_rhs->shape().element_type(),
                                 reshaped_dims),
            dot_rhs));

        if (original_hlo->opcode() == HloOpcode::kDot) {
          UpdateDDNums(&new_ddnums, rhs_concat_dim, false);
        }
      }

      // The generated original dot will not be used.
      TF_ASSIGN_OR_RETURN(auto original_dot,
                          create_sharded_dot(original_dot_lhs, original_dot_rhs,
                                             &body_b, conv_window));
      VLOG(2) << original_dot->ToString();

      // Generate the correct shape of the new dot/conv.
      auto original_sharded_dot_shape = original_dot->shape();
      auto new_dot_shape = original_sharded_dot_shape;
      std::vector<int64> new_dims(new_dot_shape.dimensions().begin(),
                                  new_dot_shape.dimensions().end());
      if (!windowed_at_contracting_dims) {
        auto slice_dim =
            lhs_concat_dim != -1
                ? indices_map.lhs_to_output_indices[lhs_concat_dim]
                : indices_map.rhs_to_output_indices[rhs_concat_dim];
        new_dims[slice_dim] /= 2;
        new_dims.insert(new_dims.begin() + slice_dim, 2);
      } else {
        new_dims.push_back(1);
      }
      new_dot_shape =
          ShapeUtil::MakeShape(original_hlo->shape().element_type(), new_dims);

      HloInstruction* dot;
      if (original_hlo->opcode() == HloOpcode::kDot) {
        dot = body_b.AddInstruction(HloInstruction::CreateDot(
            new_dot_shape, dot_lhs, dot_rhs, new_ddnums,
            original_hlo->precision_config()));
      } else {
        if (!windowed_at_contracting_dims && !windowed_at_batch_dims) {
          if (lhs_concat_dim != -1) {
            std::vector<int64> new_dims(dot_rhs->shape().dimensions().begin(),
                                        dot_rhs->shape().dimensions().end());
            new_dims.push_back(1);
            dot_rhs = body_b.AddInstruction(HloInstruction::CreateReshape(
                ShapeUtil::MakeShape(dot_rhs->shape().element_type(), new_dims),
                dot_rhs));
          }
          if (rhs_concat_dim != -1) {
            std::vector<int64> new_dims(dot_lhs->shape().dimensions().begin(),
                                        dot_lhs->shape().dimensions().end());
            new_dims.push_back(1);
            dot_lhs = body_b.AddInstruction(HloInstruction::CreateReshape(
                ShapeUtil::MakeShape(dot_lhs->shape().element_type(), new_dims),
                dot_lhs));
          }
        }

        dot = body_b.AddInstruction(HloInstruction::CreateConvolve(
            new_dot_shape, dot_lhs, dot_rhs,
            original_dot->feature_group_count(),
            original_dot->batch_group_count(),
            GenNewWindow(original_dot, dot_lhs, dot_rhs, lhs_concat_dim,
                         rhs_concat_dim, windowed_at_contracting_dims,
                         windowed_at_batch_dims),
            GenNewConvDNums(original_dot, dot_lhs, dot_rhs, lhs_concat_dim,
                            rhs_concat_dim, windowed_at_contracting_dims,
                            windowed_at_batch_dims,
                            indices_map.lhs_to_output_indices,
                            indices_map.rhs_to_output_indices, new_dot_shape),
            original_dot->precision_config()));
      }
      VLOG(2) << dot->ToString();

      // Reshape to the original sharded dot shape.
      dot = body_b.AddInstruction(
          HloInstruction::CreateReshape(original_sharded_dot_shape, dot));

      if (windowed_at_contracting_dims) {
        // Accumulate the partial output to the result buffer.
        o = body_b.AddInstruction(
            HloInstruction::CreateBinary(o->shape(), HloOpcode::kAdd, o, dot));
      } else {
        // The windowing operand is partitioned along batch/non-contracting
        // dimensions, so we need a dynamic-update-slice to save the partial
        // output in the result buffer.
        auto slice_shape = dot->shape();
        auto slice_dim =
            lhs_concat_dim != -1
                ? indices_map.lhs_to_output_indices[lhs_concat_dim]
                : indices_map.rhs_to_output_indices[rhs_concat_dim];
        slice_shape.set_dimensions(slice_dim,
                                   dot->shape().dimensions(slice_dim) / 2);
        std::vector<int64> ccw_start_indices(dot->shape().rank(), 0);
        std::vector<int64> cw_start_indices(dot->shape().rank(), 0);
        cw_start_indices[slice_dim] = dot->shape().dimensions(slice_dim) / 2;
        auto ccw_dot = body_b.AddInstruction(HloInstruction::CreateSlice(
            slice_shape, dot, ccw_start_indices, slice_shape.dimensions(),
            std::vector<int64>(dot->shape().rank(), 1)));
        auto cw_dot = body_b.AddInstruction(HloInstruction::CreateSlice(
            slice_shape, dot, cw_start_indices, dot->shape().dimensions(),
            std::vector<int64>(dot->shape().rank(), 1)));

        if (operands_sharded_at_contracting_dims) {
          // Accumulate the partial output to the result buffer.
          o = body_b.AddInstruction(HloInstruction::CreateBinary(
              o->shape(), HloOpcode::kAdd, o, ccw_dot));
          cw_cp_output = body_b.AddInstruction(HloInstruction::CreateBinary(
              o->shape(), HloOpcode::kAdd, cw_cp_output, cw_dot));
        } else {
          auto ccw_offsets = MakePartitionOffsets(
              o->shape(),
              windowed_op_is_lhs ? *lhs_sharding_transposed_to_match_output
                                 : *rhs_sharding_transposed_to_match_output,
              ccw_data_partition_id, &body_b);
          auto cw_offsets = MakePartitionOffsets(
              o->shape(),
              windowed_op_is_lhs ? *lhs_sharding_transposed_to_match_output
                                 : *rhs_sharding_transposed_to_match_output,
              cw_data_partition_id, &body_b);
          o = body_b.AddInstruction(HloInstruction::CreateDynamicUpdateSlice(
              o->shape(), o, ccw_dot, ccw_offsets));
          o = body_b.AddInstruction(HloInstruction::CreateDynamicUpdateSlice(
              o->shape(), o, cw_dot, cw_offsets));
        }
      }

      std::vector<HloInstruction*> partial_results;
      partial_results.push_back(o);
      partial_results.push_back(cw_cp_output);
      return partial_results;
    };

    // Generate partial result used by unidirectional algorithm.
    auto get_partial_unid_result =
        [&](HloInstruction* l, HloInstruction* r, HloInstruction* o,
            HloInstruction* i) -> StatusOr<HloInstruction*> {
      auto partition_id =
          lhs.state().collective_ops_creator.create_partition_id(&body_b);
      auto data_partition_id =
          body_b.AddInstruction(HloInstruction::CreateBinary(
              i->shape(), HloOpcode::kAdd, i, partition_id));
      auto partition_count =
          body_b.AddInstruction(HloInstruction::CreateConstant(
              LiteralUtil::CreateR0<uint32>(num_partitions)));
      data_partition_id = body_b.AddInstruction(
          HloInstruction::CreateBinary(i->shape(), HloOpcode::kRemainder,
                                       data_partition_id, partition_count));
      auto dot_lhs = l;
      auto dot_rhs = r;
      if (windowed_at_contracting_dims || windowed_at_batch_dims ||
          operands_sharded_at_contracting_dims) {
        // Slice the matching operand according to the partitioned dimensions on
        // the windowed operand or the output.
        auto slice_operand = !windowed_op_is_lhs ? l : r;
        // We do this by treating the matching operand as replicated, and
        // resharding it to match the windowed operand or the output.
        slice_operand->set_sharding(HloSharding::Replicate());
        auto state = lhs.state();
        state.b = &body_b;
        state.partition_id = data_partition_id;
        state.reshard_cache->per_hlo_cache.erase(slice_operand);
        const HloSharding* slice_sharding;
        if (operands_sharded_at_contracting_dims) {
          slice_sharding = windowed_op_is_lhs
                               ? &*output_sharding_transposed_to_match_rhs
                               : &*output_sharding_transposed_to_match_lhs;
        } else {
          slice_sharding = windowed_op_is_lhs
                               ? &*lhs_sharding_transposed_to_match_rhs
                               : &*rhs_sharding_transposed_to_match_lhs;
        }
        auto slice =
            PartitionedHlo(slice_operand, slice_operand->shape(), state)
                .Reshard(*slice_sharding)
                .hlo();
        slice_operand->clear_sharding();
        if (!windowed_op_is_lhs) {
          dot_lhs = slice;
        } else {
          dot_rhs = slice;
        }
      }
      TF_ASSIGN_OR_RETURN(
          auto dot, create_sharded_dot(dot_lhs, dot_rhs, &body_b, conv_window));
      if (windowed_at_contracting_dims ||
          operands_sharded_at_contracting_dims) {
        // Accumulate the partial output to the result buffer.
        o = body_b.AddInstruction(
            HloInstruction::CreateBinary(o->shape(), HloOpcode::kAdd, o, dot));
      } else {
        // The windowing operand is partitioned along batch/non-contracting
        // dimensions, so we need a dynamic-update-slice to save the partial
        // output in the result buffer.
        auto offsets = MakePartitionOffsets(
            o->shape(),
            windowed_op_is_lhs ? *lhs_sharding_transposed_to_match_output
                               : *rhs_sharding_transposed_to_match_output,
            data_partition_id, &body_b);
        o = body_b.AddInstruction(HloInstruction::CreateDynamicUpdateSlice(
            o->shape(), o, dot, offsets));
      }
      return o;
    };

    auto param = body_b.AddInstruction(HloInstruction::CreateParameter(
        /*parameter_number=*/0,
        ShapeUtil::MakeTupleShape({lhs.hlo()->shape(), rhs.hlo()->shape(),
                                   result_buffer->shape(),
                                   extra_buffer->shape(), iteration->shape()}),
        "param"));
    auto l = body_b.AddInstruction(
        HloInstruction::CreateGetTupleElement(lhs.hlo()->shape(), param, 0));
    auto r = body_b.AddInstruction(
        HloInstruction::CreateGetTupleElement(rhs.hlo()->shape(), param, 1));
    auto o = body_b.AddInstruction(HloInstruction::CreateGetTupleElement(
        result_buffer->shape(), param, 2));
    auto extra_inout = body_b.AddInstruction(
        HloInstruction::CreateGetTupleElement(extra_buffer->shape(), param, 3));
    auto i = body_b.AddInstruction(
        HloInstruction::CreateGetTupleElement(iteration->shape(), param, 4));

    // The bidirectional collective permute implementation has loop unrolling
    // of degree 2, so num_partitions is required to be a multiple of 4.
    if (options.bidirectional_windowed_einsum && num_partitions % 4 == 0) {
      std::vector<std::pair<int64, int64>> ccw_sd_pairs(num_partitions);
      for (int64 source = 0; source < num_partitions; ++source) {
        // 0 -> n-1, 1 -> 0, 2 -> 1, ...
        ccw_sd_pairs[source] = {source,
                                (source - 1 + num_partitions) % num_partitions};
      }
      std::vector<std::pair<int64, int64>> cw_sd_pairs(num_partitions);
      for (int64 source = 0; source < num_partitions; ++source) {
        // 0 -> 1, 1 -> 2, 2 -> 3, ...
        cw_sd_pairs[source] = {source, (source + 1) % num_partitions};
      }

      // Even number iteration.
      auto next_l = l;
      auto next_r = r;
      auto ccw_cp_input = operands_sharded_at_contracting_dims ? o
                          : windowed_op_is_lhs                 ? l
                                                               : r;
      auto ccw_cp_output =
          lhs.state()
              .collective_ops_creator.create_cross_partition_collective_permute(
                  &body_b, ccw_cp_input, ccw_sd_pairs,
                  (*lhs.state().next_channel_id)++);
      if (operands_sharded_at_contracting_dims) {
        o = ccw_cp_output;
      } else if (windowed_op_is_lhs) {
        next_l = ccw_cp_output;
      } else {
        next_r = ccw_cp_output;
      }
      auto cw_cp_input = extra_inout;
      auto cw_cp_output =
          lhs.state()
              .collective_ops_creator.create_cross_partition_collective_permute(
                  &body_b, cw_cp_input, cw_sd_pairs,
                  (*lhs.state().next_channel_id)++);

      TF_ASSIGN_OR_RETURN(
          auto outputs,
          get_partial_bid_results(l, r, o, extra_inout, cw_cp_output, i));
      o = outputs[0];
      cw_cp_output = outputs[1];

      // ++i
      i = body_b.AddInstruction(HloInstruction::CreateBinary(
          i->shape(), HloOpcode::kAdd, i, CreateOne(i->shape(), &body_b)));

      // Odd number iteration.
      auto second_next_l = next_l;
      auto second_next_r = next_r;
      ccw_cp_input = operands_sharded_at_contracting_dims ? o
                     : windowed_op_is_lhs                 ? next_l
                                                          : next_r;
      ccw_cp_output =
          lhs.state()
              .collective_ops_creator.create_cross_partition_collective_permute(
                  &body_b, ccw_cp_input, ccw_sd_pairs,
                  (*lhs.state().next_channel_id)++);
      if (operands_sharded_at_contracting_dims) {
        o = ccw_cp_output;
      } else if (windowed_op_is_lhs) {
        second_next_l = ccw_cp_output;
      } else {
        second_next_r = ccw_cp_output;
      }
      auto next_cw_cp_input = cw_cp_output;
      auto next_cw_cp_output =
          lhs.state()
              .collective_ops_creator.create_cross_partition_collective_permute(
                  &body_b, next_cw_cp_input, cw_sd_pairs,
                  (*lhs.state().next_channel_id)++);

      TF_ASSIGN_OR_RETURN(
          outputs, get_partial_bid_results(next_l, next_r, o, cw_cp_output,
                                           next_cw_cp_output, i));
      o = outputs[0];
      next_cw_cp_output = outputs[1];

      // ++i
      i = body_b.AddInstruction(HloInstruction::CreateBinary(
          i->shape(), HloOpcode::kAdd, i, CreateOne(i->shape(), &body_b)));

      body_b.AddInstruction(HloInstruction::CreateTuple(
          {second_next_l, second_next_r, o, next_cw_cp_output, i}));

    } else if (options.unroll_windowed_einsum && num_partitions % 2 == 0) {
      if (operands_sharded_at_contracting_dims) {
        std::vector<std::pair<int64, int64>> output_sd_pairs(num_partitions);
        for (int64 source = 0; source < num_partitions; ++source) {
          // 0 -> n-2, 1 -> n-1, 2 -> 0, ...
          output_sd_pairs[source] = {
              source, (source - 2 + num_partitions) % num_partitions};
        }

        o = lhs.state()
                .collective_ops_creator
                .create_cross_partition_collective_permute(
                    &body_b, o, output_sd_pairs,
                    (*lhs.state().next_channel_id)++);

        TF_ASSIGN_OR_RETURN(extra_inout,
                            get_partial_unid_result(l, r, extra_inout, i));

        extra_inout = lhs.state()
                          .collective_ops_creator
                          .create_cross_partition_collective_permute(
                              &body_b, extra_inout, output_sd_pairs,
                              (*lhs.state().next_channel_id)++);

        // i+2
        i = body_b.AddInstruction(HloInstruction::CreateBinary(
            i->shape(), HloOpcode::kAdd, i,
            body_b.AddInstruction(HloInstruction::CreateConstant(
                LiteralUtil::CreateR0<uint32>(2)))));
        auto real_i = body_b.AddInstruction(HloInstruction::CreateBinary(
            i->shape(), HloOpcode::kAdd, i,
            body_b.AddInstruction(HloInstruction::CreateConstant(
                LiteralUtil::CreateR0<uint32>(1)))));

        TF_ASSIGN_OR_RETURN(o, get_partial_unid_result(l, r, o, real_i));
        body_b.AddInstruction(
            HloInstruction::CreateTuple({l, r, o, extra_inout, i}));
      } else {
        std::vector<std::pair<int64, int64>> sd_pairs(num_partitions);
        for (int64 source = 0; source < num_partitions; ++source) {
          // 0 -> n-1, 1 -> 0, 2 -> 1, ...
          sd_pairs[source] = {source,
                              (source - 1 + num_partitions) % num_partitions};
        }

        // Even number iteration.
        auto next_l = l;
        auto next_r = r;
        auto cp_input = windowed_op_is_lhs ? l : r;
        auto cp_output = lhs.state()
                             .collective_ops_creator
                             .create_cross_partition_collective_permute(
                                 &body_b, cp_input, sd_pairs,
                                 (*lhs.state().next_channel_id)++);
        if (windowed_op_is_lhs) {
          next_l = cp_output;
        } else {
          next_r = cp_output;
        }
        TF_ASSIGN_OR_RETURN(o, get_partial_unid_result(l, r, o, i));

        // ++i
        i = body_b.AddInstruction(HloInstruction::CreateBinary(
            i->shape(), HloOpcode::kAdd, i,
            body_b.AddInstruction(HloInstruction::CreateConstant(
                LiteralUtil::CreateR0<uint32>(1)))));

        // Odd number iteration.
        auto second_next_l = next_l;
        auto second_next_r = next_r;
        cp_input = windowed_op_is_lhs ? next_l : next_r;
        cp_output = lhs.state()
                        .collective_ops_creator
                        .create_cross_partition_collective_permute(
                            &body_b, cp_input, sd_pairs,
                            (*lhs.state().next_channel_id)++);
        if (windowed_op_is_lhs) {
          second_next_l = cp_output;
        } else {
          second_next_r = cp_output;
        }
        TF_ASSIGN_OR_RETURN(o, get_partial_unid_result(next_l, next_r, o, i));

        // ++i
        i = body_b.AddInstruction(HloInstruction::CreateBinary(
            i->shape(), HloOpcode::kAdd, i,
            body_b.AddInstruction(HloInstruction::CreateConstant(
                LiteralUtil::CreateR0<uint32>(1)))));

        body_b.AddInstruction(HloInstruction::CreateTuple(
            {second_next_l, second_next_r, o, extra_inout, i}));
      }
    } else {
      auto real_i = i;
      if (operands_sharded_at_contracting_dims) {
        // For reduce-scatter case, start from the data_partition_id + 1 to make
        // the data_partition_id of the final data shard in each partition the
        // same as the corresponding partition_id.
        real_i = body_b.AddInstruction(HloInstruction::CreateBinary(
            real_i->shape(), HloOpcode::kAdd, real_i,
            CreateOne(real_i->shape(), &body_b)));
      }
      TF_ASSIGN_OR_RETURN(o, get_partial_unid_result(l, r, o, real_i));

      // ++i
      i = body_b.AddInstruction(HloInstruction::CreateBinary(
          i->shape(), HloOpcode::kAdd, i,
          body_b.AddInstruction(HloInstruction::CreateConstant(
              LiteralUtil::CreateR0<uint32>(1)))));
      auto has_more = body_b.AddInstruction(HloInstruction::CreateCompare(
          ShapeUtil::MakeShape(PRED, {}), i,
          body_b.AddInstruction(HloInstruction::CreateConstant(
              LiteralUtil::CreateR0<uint32>(num_partitions))),
          ComparisonDirection::kLt));
      // Collective-permute for the next window. We don't need it for the last
      // iteration, so we use a conditional around the collective-permute.
      HloInstruction* conditional;
      {
        SpmdBuilder cp_b("window_collective_permute", original_hlo);
        {
          auto p = cp_b.AddInstruction(HloInstruction::CreateParameter(
              0,
              operands_sharded_at_contracting_dims ? o->shape()
              : windowed_op_is_lhs                 ? l->shape()
                                                   : r->shape(),
              "window"));
          std::vector<std::pair<int64, int64>> sd_pairs(num_partitions);
          for (int64 source = 0; source < num_partitions; ++source) {
            // 0 -> n-1, 1 -> 0, 2 -> 1, ...
            sd_pairs[source] = {source,
                                (source - 1 + num_partitions) % num_partitions};
          }
          lhs.state()
              .collective_ops_creator.create_cross_partition_collective_permute(
                  &cp_b, p, sd_pairs, (*lhs.state().next_channel_id)++);
        }
        SpmdBuilder ncp_b("last_iteration_noop", original_hlo);
        {
          ncp_b.AddInstruction(HloInstruction::CreateParameter(
              0,
              operands_sharded_at_contracting_dims ? o->shape()
              : windowed_op_is_lhs                 ? l->shape()
                                                   : r->shape(),
              "window"));
        }
        conditional = body_b.AddInstruction(HloInstruction::CreateConditional(
            operands_sharded_at_contracting_dims ? o->shape()
            : windowed_op_is_lhs                 ? l->shape()
                                                 : r->shape(),
            has_more,
            operands_sharded_at_contracting_dims ? o
            : windowed_op_is_lhs                 ? l
                                                 : r,
            module->AddEmbeddedComputation(cp_b.Build()),
            operands_sharded_at_contracting_dims ? o
            : windowed_op_is_lhs                 ? l
                                                 : r,
            module->AddEmbeddedComputation(ncp_b.Build())));
      }
      if (operands_sharded_at_contracting_dims) {
        o = conditional;
      } else if (windowed_op_is_lhs) {
        l = conditional;
      } else {
        r = conditional;
      }
      body_b.AddInstruction(
          HloInstruction::CreateTuple({l, r, o, extra_inout, i}));
    }

    SpmdBuilder cond_b("windowed_dot_general_cond", original_hlo);
    auto cond_param = cond_b.AddInstruction(HloInstruction::CreateParameter(
        /*parameter_number=*/0,
        ShapeUtil::MakeTupleShape({lhs.hlo()->shape(), rhs.hlo()->shape(),
                                   result_buffer->shape(),
                                   extra_buffer->shape(), iteration->shape()}),
        "param"));
    auto cond_i = cond_b.AddInstruction(HloInstruction::CreateGetTupleElement(
        iteration->shape(), cond_param, 4));
    int64 adapted_num_partitions =
        (options.bidirectional_windowed_einsum && num_partitions % 4 == 0)
            ? num_partitions / 2
            : num_partitions;
    cond_b.AddInstruction(HloInstruction::CreateCompare(
        ShapeUtil::MakeShape(PRED, {}), cond_i,
        cond_b.AddInstruction(HloInstruction::CreateConstant(
            LiteralUtil::CreateR0<uint32>(adapted_num_partitions))),
        ComparisonDirection::kLt));
    auto while_loop = b->AddInstruction(HloInstruction::CreateWhile(
        cond_param->shape(), module->AddEmbeddedComputation(cond_b.Build()),
        module->AddEmbeddedComputation(body_b.Build()),
        b->AddInstruction(HloInstruction::CreateTuple(
            {lhs.hlo(), rhs.hlo(), result_buffer, extra_buffer, iteration}))));
    windowed_dot_general_loops->push_back(
        {while_loop, windowed_op_is_lhs ? 0 : 1, windowed_at_contracting_dims,
         windowed_at_batch_dims, operands_sharded_at_contracting_dims});
    auto result = b->AddInstruction(HloInstruction::CreateGetTupleElement(
        result_buffer->shape(), while_loop, 2));
    if (((options.bidirectional_windowed_einsum && num_partitions % 4 == 0) ||
         (options.unroll_windowed_einsum && num_partitions % 2 == 0)) &&
        operands_sharded_at_contracting_dims) {
      std::vector<std::pair<int64, int64>> extra_sd_pairs(num_partitions);
      for (int64 source = 0; source < num_partitions; ++source) {
        // 0 -> 1, 1 -> 2, 2 -> 3, ...
        extra_sd_pairs[source] = {source, (source + 1) % num_partitions};
      }
      auto extra_result =
          b->AddInstruction(HloInstruction::CreateGetTupleElement(
              extra_buffer->shape(), while_loop, 3));
      if (options.bidirectional_windowed_einsum && num_partitions % 4 == 0) {
        extra_result = lhs.state()
                           .collective_ops_creator
                           .create_cross_partition_collective_permute(
                               b, extra_result, extra_sd_pairs,
                               (*lhs.state().next_channel_id)++);
      }
      if (options.unroll_windowed_einsum && num_partitions % 2 == 0) {
        result = lhs.state()
                     .collective_ops_creator
                     .create_cross_partition_collective_permute(
                         b, result, extra_sd_pairs,
                         (*lhs.state().next_channel_id)++);
      }
      result = b->AddInstruction(HloInstruction::CreateBinary(
          result->shape(), HloOpcode::kAdd, result, extra_result));
    }
    if (!ShapeUtil::Compatible(padded_result_buffer_shape,
                               unpadded_result_buffer_shape)) {
      result = b->AddInstruction(HloInstruction::CreateSlice(
          unpadded_result_buffer_shape, result,
          std::vector<int64>(padded_result_buffer_shape.rank(), 0),
          unpadded_result_buffer_shape.dimensions(),
          std::vector<int64>(padded_result_buffer_shape.rank(), 1)));
    }
    return result;
  };
  absl::optional<WindowedEinsumConfig> e_config =
      GetWindowedEinsumConfiguration(
          num_partitions, output_lhs_non_contracting_partitions,
          output_rhs_non_contracting_partitions, rhs_contracting_partitions,
          rhs_non_contracting_partitions, rhs_batch_partitions,
          lhs_contracting_partitions, lhs_non_contracting_partitions,
          lhs_batch_partitions, ShapeSizeInBytes(rhs.base_shape()),
          ShapeSizeInBytes(lhs.base_shape()),
          ShapeSizeInBytes(output_base_shape),
          options.threshold_for_windowed_einsum_mib,
          output_sharding_transposed_to_match_lhs,
          output_sharding_transposed_to_match_rhs, lhs_sharding, rhs_sharding);
  if (e_config) {
    return emit_windowed_dot_general(*e_config);
  }

  {
    // Try batch-parallel by resharding one operand, and allowing all-reduce.
    TF_ASSIGN_OR_RETURN(
        HloInstruction * partitioned_dot,
        try_emit_output_batch_partitioned_einsum_with_reshard(true));
    if (partitioned_dot) {
      return partitioned_dot;
    }
  }

  // LHS and RHS have the same partitioned contracting dimensions.
  if (lhs_contracting_partitions == rhs_contracting_partitions &&
      lhs_contracting_partitions == num_partitions) {
    auto zero = b->AddInstruction(HloInstruction::CreateConstant(
        LiteralUtil::Zero(output_base_shape.element_type())));
    // Pad both sides with zero, since NaN at one side cannot be masked by zero
    // on the other side.
    if (ShapeSizeInBytes(lhs.base_shape()) <
        ShapeSizeInBytes(rhs.base_shape())) {
      lhs =
          lhs.Reshard(*rhs_sharding_transposed_to_match_lhs).PadWithValue(zero);
      rhs = rhs.PadWithValue(zero);
    } else {
      lhs = lhs.PadWithValue(zero);
      rhs =
          rhs.Reshard(*lhs_sharding_transposed_to_match_rhs).PadWithValue(zero);
    }
    TF_ASSIGN_OR_RETURN(
        auto dot, create_sharded_dot(lhs.hlo(), rhs.hlo(), b, conv_window));
    std::vector<int64> lhs_contracting_dims;
    lhs_contracting_dims.reserve(lhs.base_shape().rank());
    for (const auto& cd : dims_mapping.contracting_dims) {
      lhs_contracting_dims.push_back(cd.lhs);
    }
    auto ar = lhs.state().partitioner->AllReduceAlongShardingDims(
        b, dot, lhs.sharding(), lhs.state().next_channel_id,
        lhs_contracting_dims, lhs.state().collective_ops_creator,
        MakeBinaryAdd(output_base_shape.element_type(), module));
    ar->set_sharding(HloSharding::Replicate());
    return PartitionedHlo(ar, output_base_shape, lhs.state())
        .Reshard(output_sharding)
        .hlo();
  }

  // LHS and output have the same partitioned non-contracting dimensions.
  if (lhs_non_contracting_partitions == num_partitions &&
      output_lhs_non_contracting_partitions == num_partitions &&
      lhs_sharding_transposed_to_match_output == output_sharding) {
    auto rhs_replicated = rhs.Reshard(HloSharding::Replicate()).hlo();
    TF_ASSIGN_OR_RETURN(auto dot, create_sharded_dot(lhs.hlo(), rhs_replicated,
                                                     b, conv_window));
    return dot;
  }

  // RHS and output have the same partitioned non-contracting dimensions.
  if (rhs_non_contracting_partitions == num_partitions &&
      output_rhs_non_contracting_partitions == num_partitions &&
      rhs_sharding_transposed_to_match_output == output_sharding) {
    auto lhs_replicated = lhs.Reshard(HloSharding::Replicate()).hlo();
    TF_ASSIGN_OR_RETURN(auto dot, create_sharded_dot(lhs_replicated, rhs.hlo(),
                                                     b, conv_window));
    return dot;
  }

  if (may_reshard_without_detecting_match) {
    // Output is batch partitioned.
    if (output_batch_partitions == num_partitions) {
      auto resharded_lhs =
          lhs.Reshard(*output_sharding_transposed_to_match_lhs);
      auto resharded_rhs =
          rhs.Reshard(*output_sharding_transposed_to_match_rhs);
      TF_ASSIGN_OR_RETURN(
          auto dot, create_sharded_dot(resharded_lhs.hlo(), resharded_rhs.hlo(),
                                       b, conv_window));
      return dot;
    }
    // Output is partitioned along LHS non-contracting dimensions.
    if (output_lhs_non_contracting_partitions == num_partitions) {
      auto resharded_lhs =
          lhs.Reshard(*output_sharding_transposed_to_match_lhs);
      auto replicated_rhs = rhs.Reshard(HloSharding::Replicate());
      TF_ASSIGN_OR_RETURN(
          auto dot, create_sharded_dot(resharded_lhs.hlo(),
                                       replicated_rhs.hlo(), b, conv_window));
      return dot;
    }
    // Output is partitioned along RHS non-contracting dimensions.
    if (output_rhs_non_contracting_partitions == num_partitions) {
      auto replicated_lhs = lhs.Reshard(HloSharding::Replicate());
      auto resharded_rhs =
          rhs.Reshard(*output_sharding_transposed_to_match_rhs);
      TF_ASSIGN_OR_RETURN(
          auto dot, create_sharded_dot(replicated_lhs.hlo(),
                                       resharded_rhs.hlo(), b, conv_window));
      return dot;
    }
  }

  // Returns true if it is beneficial to reshard the operand at `operand_idx`
  // across the contracting dimension.
  const auto should_partition_contracting_dim = [&](int64 operand_idx) {
    if (!output_sharding.IsReplicated()) {
      return false;
    }

    if (operand_idx == 0) {
      // If LHS and output are replicated, we compare the cost of all-gather
      // on RHS vs all-reduce on the output.
      return (rhs_contracting_partitions == num_partitions) &&
             lhs.sharding().IsReplicated() &&
             ShapeUtil::ElementsIn(rhs.base_shape()) >
                 ShapeUtil::ElementsIn(output_base_shape);
    } else {
      return (lhs_contracting_partitions == num_partitions) &&
             rhs.sharding().IsReplicated() &&
             ShapeUtil::ElementsIn(lhs.base_shape()) >
                 ShapeUtil::ElementsIn(output_base_shape);
    }
  };

  // When the output is replicated and one of the operands is partitioned along
  // contracting dimension, align the other operand to be partitioned along
  // the contracting dimensions.
  if (output_sharding.IsReplicated() && (should_partition_contracting_dim(0) ||
                                         should_partition_contracting_dim(1))) {
    auto zero = b->AddInstruction(HloInstruction::CreateConstant(
        LiteralUtil::Zero(output_base_shape.element_type())));
    if (should_partition_contracting_dim(0)) {
      lhs =
          lhs.Reshard(*rhs_sharding_transposed_to_match_lhs).PadWithValue(zero);
      rhs = rhs.PadWithValue(zero);
    } else {
      lhs = lhs.PadWithValue(zero);
      rhs =
          rhs.Reshard(*lhs_sharding_transposed_to_match_rhs).PadWithValue(zero);
    }
    TF_ASSIGN_OR_RETURN(
        auto dot, create_sharded_dot(lhs.hlo(), rhs.hlo(), b, conv_window));

    std::vector<int64> lhs_contracting_dims;
    lhs_contracting_dims.reserve(lhs.base_shape().rank());
    for (const auto& cd : dims_mapping.contracting_dims) {
      lhs_contracting_dims.push_back(cd.lhs);
    }
    return lhs.state().partitioner->AllReduceAlongShardingDims(
        b, dot, lhs.sharding(), lhs.state().next_channel_id,
        lhs_contracting_dims, lhs.state().collective_ops_creator,
        MakeBinaryAdd(output_base_shape.element_type(), module));
  }
  return nullptr;
}

StatusOr<HloInstruction*> PartitionDot(
    PartitionedHlo lhs, PartitionedHlo rhs, const Shape& output_base_shape,
    const HloSharding& output_sharding, const DotConvDimsMapping& dims_mapping,
    int64 num_partitions,
    const std::function<StatusOr<HloInstruction*>(
        HloInstruction*, HloInstruction*, SpmdBuilder*,
        const Window& conv_window)>& create_sharded_dot,
    const Window& conv_window, HloModule* module, HloInstruction* original_hlo,
    const SpmdPartitionerOptions& options, SpmdBuilder* b,
    std::vector<SpmdPartitioningVisitor::WindowedDotGeneralLoop>*
        windowed_dot_general_loops);

StatusOr<HloInstruction*> PartitionDotGroupOnBatch(
    PartitionedHlo lhs, PartitionedHlo rhs, const Shape& output_base_shape,
    const HloSharding& output_sharding, const DotConvDimsMapping& dims_mapping,
    int64 num_partitions, int64 lhs_contracting_partitions,
    int64 rhs_contracting_partitions, int64 lhs_non_contracting_partitions,
    int64 rhs_non_contracting_partitions,
    const std::function<StatusOr<HloInstruction*>(
        HloInstruction*, HloInstruction*, SpmdBuilder*,
        const Window& conv_window)>& create_sharded_dot,
    const Window& conv_window, HloModule* module, HloInstruction* original_hlo,
    bool require_matching_devices_to_group,
    const SpmdPartitionerOptions& options, SpmdBuilder* b,
    std::vector<SpmdPartitioningVisitor::WindowedDotGeneralLoop>*
        windowed_dot_general_loops) {
  std::vector<std::pair<HloInstruction*, HloSharding>>
      top_level_sharding_to_reset;
  auto cleaner = tensorflow::gtl::MakeCleanup([&] {
    for (auto& to_reset : top_level_sharding_to_reset) {
      to_reset.first->set_sharding(to_reset.second);
    }
  });
  std::vector<int64> lhs_dims;
  std::vector<int64> rhs_dims;
  std::vector<int64> output_dims;
  auto lhs_sharding_dims_adjusted_to_output =
      lhs.sharding().IsReplicated()
          ? std::vector<int64>(lhs.base_shape().rank(), 1)
          : lhs.sharding().tile_assignment().dimensions();
  auto rhs_sharding_dims_adjusted_to_output =
      rhs.sharding().IsReplicated()
          ? std::vector<int64>(rhs.base_shape().rank(), 1)
          : rhs.sharding().tile_assignment().dimensions();
  auto output_sharding_dims_adjusted_to_lhs =
      output_sharding.tile_assignment().dimensions();
  bool lhs_rhs_dims_matching = true;
  for (const auto& dim : dims_mapping.batch_dims) {
    lhs_dims.push_back(dim.lhs);
    rhs_dims.push_back(dim.rhs);
    output_dims.push_back(dim.output);
    if (lhs_sharding_dims_adjusted_to_output[dim.lhs] !=
        rhs_sharding_dims_adjusted_to_output[dim.rhs]) {
      lhs_rhs_dims_matching = false;
    }
    lhs_sharding_dims_adjusted_to_output[dim.lhs] =
        output_sharding.tile_assignment().dim(dim.output);
    rhs_sharding_dims_adjusted_to_output[dim.rhs] =
        output_sharding.tile_assignment().dim(dim.output);
    output_sharding_dims_adjusted_to_lhs[dim.output] =
        lhs.sharding().tile_assignment().dim(dim.lhs);
  }
  if (require_matching_devices_to_group && lhs_rhs_dims_matching) {
    lhs_rhs_dims_matching =
        rhs.sharding() == UngroupSharding(AlignGroupsWith(
                              GroupShardingOnDims(rhs.sharding(), rhs_dims),
                              GroupShardingOnDims(lhs.sharding(), lhs_dims)));
  }
  auto output_grouped = GroupShardingOnDims(output_sharding, output_dims);
  PartitionedHlo per_group_lhs = lhs;
  PartitionedHlo per_group_rhs = rhs;
  if (lhs_rhs_dims_matching) {
    auto lhs_grouped = GroupShardingOnDims(lhs.sharding(), lhs_dims);
    auto rhs_grouped = GroupShardingOnDims(rhs.sharding(), rhs_dims);
    if (ShapeUtil::ByteSizeOf(lhs.hlo()->shape()) >
        ShapeUtil::ByteSizeOf(rhs.hlo()->shape())) {
      rhs_grouped = AlignGroupsWith(std::move(rhs_grouped), lhs_grouped);
      rhs = rhs.Reshard(UngroupSharding(rhs_grouped));
    } else {
      lhs_grouped = AlignGroupsWith(std::move(lhs_grouped), rhs_grouped);
      lhs = lhs.Reshard(UngroupSharding(lhs_grouped));
    }
    auto reshaped_output_tiling = output_sharding.tile_assignment();
    reshaped_output_tiling.Reshape(output_sharding_dims_adjusted_to_lhs);
    output_grouped = AlignGroupsWith(
        GroupShardingOnDims(
            output_sharding.ReplicateOnLastTileDim()
                ? HloSharding::PartialTile(reshaped_output_tiling)
                : HloSharding::Tile(reshaped_output_tiling),
            output_dims),
        lhs_grouped);
    auto per_group_partitioner_state = CreatePerGroupPartitioningState(
        lhs.state(), lhs_grouped.device_groups, b);
    top_level_sharding_to_reset.emplace_back(lhs.hlo(), lhs.sharding());
    lhs.hlo()->set_sharding(lhs_grouped.sharding);
    top_level_sharding_to_reset.emplace_back(rhs.hlo(), rhs.sharding());
    rhs.hlo()->set_sharding(rhs_grouped.sharding);
    CHECK(lhs.hlo() != rhs.hlo() ||
          lhs_grouped.sharding == rhs_grouped.sharding);
    per_group_lhs = PartitionedHlo(
        lhs.hlo(), GetPerGroupBaseShape(lhs_grouped, lhs.base_shape()),
        per_group_partitioner_state);
    per_group_rhs = PartitionedHlo(
        rhs.hlo(), GetPerGroupBaseShape(rhs_grouped, rhs.base_shape()),
        per_group_partitioner_state);
  } else {
    auto per_group_partitioner_state = CreatePerGroupPartitioningState(
        lhs.state(), output_grouped.device_groups, b);
    auto reshard_to_output_batch =
        [&](PartitionedHlo operand, absl::Span<const int64> batch_dims,
            absl::Span<const int64> contracting_dims,
            absl::Span<const int64> non_contracting_dims,
            int64 contracting_dim_partitions,
            int64 non_contracting_dim_partitions,
            int64 other_contracting_dim_partitions,
            std::vector<int64>* sharding_dims_adjusted_to_output)
        -> absl::optional<PartitionedHlo> {
      if (operand.sharding().IsTileMaximal()) {
        auto partially_sharded = PerGroupSliceFromReplicated(
            operand.Replicate().hlo(), operand.state().partition_id,
            output_grouped.device_groups, batch_dims,
            output_grouped.group_dim_sizes, b);
        partially_sharded->set_sharding(HloSharding::Replicate());
        return PartitionedHlo(partially_sharded, partially_sharded->shape(),
                              per_group_partitioner_state);
      }
      auto reshaped_tiling = operand.sharding().tile_assignment();
      // It's possible that the operand is not initially sharded on batch
      // dimensions in the same way as the output, although being tiled. In that
      // case, the current sharding_dims_adjusted_to_output may contain more
      // partitions than available devices. We remove partitioning on other
      // dimensions.
      if (Product(*sharding_dims_adjusted_to_output) >
          reshaped_tiling.num_elements()) {
        if (Product(*sharding_dims_adjusted_to_output) %
                reshaped_tiling.num_elements() !=
            0) {
          return absl::nullopt;
        }
        int64 ratio = Product(*sharding_dims_adjusted_to_output) /
                      reshaped_tiling.num_elements();
        if (operand.sharding().ReplicateOnLastTileDim() &&
            reshaped_tiling.dimensions().back() % ratio == 0) {
          sharding_dims_adjusted_to_output->back() /= ratio;
          if (sharding_dims_adjusted_to_output->back() == 1) {
            sharding_dims_adjusted_to_output->pop_back();
          }
        } else if (ratio == non_contracting_dim_partitions &&
                   (ratio != contracting_dim_partitions ||
                    contracting_dim_partitions ==
                        other_contracting_dim_partitions)) {
          for (int64 dim : non_contracting_dims) {
            (*sharding_dims_adjusted_to_output)[dim] = 1;
          }
        } else if (ratio == contracting_dim_partitions) {
          for (int64 dim : contracting_dims) {
            (*sharding_dims_adjusted_to_output)[dim] = 1;
          }
        } else {
          return absl::nullopt;
        }
      }
      // If the operand is initially sharded more ways than the output in the
      // batch dimensions, sharding_dims_adjusted_to_output currently contains
      // fewer partitions than available devices. We do not handle this case.
      if (Product(*sharding_dims_adjusted_to_output) <
          reshaped_tiling.num_elements()) {
        return absl::nullopt;
      }
      reshaped_tiling.Reshape(*sharding_dims_adjusted_to_output);
      auto grouped = AlignGroupsWith(
          GroupShardingOnDims(operand.base_shape().rank() <
                                      sharding_dims_adjusted_to_output->size()
                                  ? HloSharding::PartialTile(reshaped_tiling)
                                  : HloSharding::Tile(reshaped_tiling),
                              batch_dims),
          output_grouped);
      if (require_matching_devices_to_group &&
          operand.sharding() != UngroupSharding(grouped)) {
        return absl::nullopt;
      }
      auto resharded = operand.Reshard(UngroupSharding(grouped));
      top_level_sharding_to_reset.emplace_back(resharded.hlo(),
                                               resharded.sharding());
      resharded.hlo()->set_sharding(grouped.sharding);
      return PartitionedHlo(resharded.hlo(),
                            GetPerGroupBaseShape(grouped, operand.base_shape()),
                            per_group_partitioner_state);
    };
    std::vector<int64> lhs_contracting_dims;
    std::vector<int64> rhs_contracting_dims;
    lhs_contracting_dims.reserve(dims_mapping.contracting_dims.size());
    rhs_contracting_dims.reserve(dims_mapping.contracting_dims.size());
    for (const auto& dim : dims_mapping.contracting_dims) {
      lhs_contracting_dims.push_back(dim.lhs);
      rhs_contracting_dims.push_back(dim.rhs);
    }
    std::vector<int64> lhs_non_contracting_dims;
    std::vector<int64> rhs_non_contracting_dims;
    lhs_non_contracting_dims.reserve(
        dims_mapping.lhs_non_contracting_dims.size());
    rhs_non_contracting_dims.reserve(
        dims_mapping.rhs_non_contracting_dims.size());
    for (const auto& dim : dims_mapping.lhs_non_contracting_dims) {
      lhs_non_contracting_dims.push_back(dim.lhs);
    }
    for (const auto& dim : dims_mapping.rhs_non_contracting_dims) {
      rhs_non_contracting_dims.push_back(dim.rhs);
    }
    if (auto resharded = reshard_to_output_batch(
            lhs, lhs_dims, lhs_contracting_dims, lhs_non_contracting_dims,
            lhs_contracting_partitions, lhs_non_contracting_partitions,
            rhs_contracting_partitions,
            &lhs_sharding_dims_adjusted_to_output)) {
      per_group_lhs = *resharded;
    } else {
      return nullptr;
    }
    if (auto resharded = reshard_to_output_batch(
            rhs, rhs_dims, rhs_contracting_dims, rhs_non_contracting_dims,
            rhs_contracting_partitions, rhs_non_contracting_partitions,
            lhs_contracting_partitions,
            &rhs_sharding_dims_adjusted_to_output)) {
      per_group_rhs = *resharded;
    } else {
      return nullptr;
    }
    CHECK(lhs.hlo() != rhs.hlo() ||
          per_group_lhs.sharding() == per_group_rhs.sharding());
  }
  TF_ASSIGN_OR_RETURN(
      auto dot,
      PartitionDot(per_group_lhs, per_group_rhs,
                   GetPerGroupBaseShape(output_grouped, output_base_shape),
                   output_grouped.sharding, dims_mapping,
                   num_partitions / output_grouped.device_groups.size(),
                   create_sharded_dot, conv_window, module, original_hlo,
                   options, b, windowed_dot_general_loops));
  dot->set_sharding(UngroupSharding(output_grouped));
  return PartitionedHlo(dot, output_base_shape, lhs.state())
      .Reshard(output_sharding)
      .hlo();
}

GroupedSharding GetNonContractingPartitionGroupedShardingForMatchedOperand(
    bool lhs_matching, const HloSharding& matching_sharding,
    const HloSharding& output_sharding,
    absl::Span<const DotConvDimsMapping::DimsMapping> partitioned_dims) {
  std::vector<int64> matching_sharding_dims =
      matching_sharding.tile_assignment().dimensions();
  std::vector<int64> matching_dims;
  std::vector<int64> output_dims;
  // Make sure the partitioning on matching's non-contracting dimensions
  // defines the same device groups for both matching and output.
  for (const auto& dim : partitioned_dims) {
    int64 md = lhs_matching ? dim.lhs : dim.rhs;
    matching_sharding_dims[md] =
        output_sharding.tile_assignment().dim(dim.output);
    matching_dims.push_back(md);
    output_dims.push_back(dim.output);
  }
  GroupedSharding output_grouped =
      GroupShardingOnDims(output_sharding, output_dims);
  Array<int64> reshaped_matching_tiling = matching_sharding.tile_assignment();
  reshaped_matching_tiling.Reshape(matching_sharding_dims);
  return AlignGroupsWith(
      GroupShardingOnDims(
          matching_sharding.ReplicateOnLastTileDim()
              ? HloSharding::PartialTile(reshaped_matching_tiling)
              : HloSharding::Tile(reshaped_matching_tiling),
          matching_dims),
      output_grouped);
}

absl::optional<GroupedSharding>
GetNonContractingPartitionGroupedShardingForOtherOperand(
    bool lhs_matching, const Shape& output_base_shape, const Shape& other_shape,
    int64 other_contracting_partitions, int64 other_non_contracting_partitions,
    int64 matching_contracting_partitions,
    int64 output_other_non_contracting_partitions,
    const HloSharding& other_sharding, const HloSharding& output_sharding,
    absl::Span<const DotConvDimsMapping::DimsMapping> matching_partitioned_dims,
    absl::Span<const DotConvDimsMapping::DimsMapping>
        other_non_contracting_dims,
    absl::Span<const DotConvDimsMapping::DimsMapping> other_contracting_dims) {
  int64 group_count = 1;
  std::vector<int64> output_dims;
  for (const auto& dim : matching_partitioned_dims) {
    output_dims.push_back(dim.output);
    group_count *= output_sharding.tile_assignment().dim(dim.output);
  }
  GroupedSharding output_grouped =
      GroupShardingOnDims(output_sharding, output_dims);
  std::vector<int64> other_group_dims;
  if (other_sharding.ReplicateOnLastTileDim() &&
      other_sharding.tile_assignment().dimensions().back() % group_count == 0) {
    other_group_dims.push_back(
        other_sharding.tile_assignment().num_dimensions() - 1);
  } else {
    const bool may_replicate_other_contracting_dims =
        (other_contracting_partitions == group_count &&
         other_non_contracting_partitions ==
             output_other_non_contracting_partitions);
    const bool may_replicate_other_non_contracting_dims =
        group_count == other_non_contracting_partitions &&
        matching_contracting_partitions == other_contracting_partitions;
    if (auto found_dims = FindMatchingPartitionedDimsForGrouping(
            other_sharding, output_grouped.device_groups)) {
      other_group_dims = std::move(*found_dims);
    } else if (may_replicate_other_contracting_dims &&
               (!may_replicate_other_non_contracting_dims ||
                ShapeUtil::ByteSizeOf(other_shape)) <=
                   ShapeUtil::ByteSizeOf(MakePartitionedShape(
                       output_base_shape, output_sharding))) {
      for (const auto& dim : other_contracting_dims) {
        other_group_dims.push_back(lhs_matching ? dim.rhs : dim.lhs);
      }
    } else if (may_replicate_other_non_contracting_dims) {
      for (const auto& dim : other_non_contracting_dims) {
        other_group_dims.push_back(lhs_matching ? dim.rhs : dim.lhs);
      }
    } else {
      return absl::nullopt;
    }
  }
  if (other_group_dims.size() == 1 &&
      other_group_dims[0] ==
          other_sharding.tile_assignment().num_dimensions() - 1) {
    return AlignGroupsWith(
        GroupShardingOnDims(
            other_sharding, {other_group_dims[0]},
            {other_sharding.tile_assignment().dimensions().back() /
             group_count}),
        output_grouped, /*ignore_group_order=*/true);

  } else if (!other_sharding.IsReplicated()) {
    return AlignGroupsWith(
        GroupShardingOnDims(other_sharding, other_group_dims), output_grouped,
        /*ignore_group_order=*/true);
  }
  return absl::nullopt;
}

StatusOr<HloInstruction*> PartitionDotGroupOnNonContracting(
    bool lhs_matching, PartitionedHlo matching, PartitionedHlo other,
    int64 matching_contracting_partitions, int64 other_contracting_partitions,
    absl::Span<const DotConvDimsMapping::DimsMapping>
        partitioned_non_contracting_dims,
    int64 other_non_contracting_partitions,
    int64 output_other_non_contracting_partitions,
    const Shape& output_base_shape, const HloSharding& output_sharding,
    const DotConvDimsMapping& dims_mapping, int64 num_partitions,
    const std::function<StatusOr<HloInstruction*>(
        HloInstruction*, HloInstruction*, SpmdBuilder*,
        const Window& conv_window)>& create_sharded_dot,
    const Window& conv_window, HloModule* module, HloInstruction* original_hlo,
    bool require_matching_devices_to_group,
    const SpmdPartitionerOptions& options, SpmdBuilder* b,
    std::vector<SpmdPartitioningVisitor::WindowedDotGeneralLoop>*
        windowed_dot_general_loops) {
  std::vector<std::pair<HloInstruction*, HloSharding>>
      top_level_sharding_to_reset;
  auto cleaner = tensorflow::gtl::MakeCleanup([&] {
    for (auto& to_reset : top_level_sharding_to_reset) {
      to_reset.first->set_sharding(to_reset.second);
    }
  });

  std::vector<int64> output_dims;
  for (const auto& dim : partitioned_non_contracting_dims) {
    output_dims.push_back(dim.output);
  }
  GroupedSharding output_grouped =
      GroupShardingOnDims(output_sharding, output_dims);
  GroupedSharding matching_grouped =
      GetNonContractingPartitionGroupedShardingForMatchedOperand(
          lhs_matching, matching.sharding(), output_sharding,
          partitioned_non_contracting_dims);
  if (require_matching_devices_to_group &&
      matching.sharding() != UngroupSharding(matching_grouped)) {
    return nullptr;
  }
  absl::optional<GroupedSharding> other_grouped =
      GetNonContractingPartitionGroupedShardingForOtherOperand(
          lhs_matching, output_base_shape, other.hlo()->shape(),
          other_contracting_partitions, other_non_contracting_partitions,
          matching_contracting_partitions,
          output_other_non_contracting_partitions, other.sharding(),
          output_sharding, partitioned_non_contracting_dims,
          lhs_matching ? dims_mapping.rhs_non_contracting_dims
                       : dims_mapping.lhs_non_contracting_dims,
          dims_mapping.contracting_dims);

  if (!other_grouped) {
    other = other.Replicate();
  }
  matching = matching.Reshard(UngroupSharding(matching_grouped));
  auto per_group_partitioner_state = CreatePerGroupPartitioningState(
      matching.state(), matching_grouped.device_groups, b);
  top_level_sharding_to_reset.emplace_back(matching.hlo(), matching.sharding());
  matching.hlo()->set_sharding(matching_grouped.sharding);
  auto matching_p = PartitionedHlo(
      matching.hlo(),
      GetPerGroupBaseShape(matching_grouped, matching.base_shape()),
      per_group_partitioner_state);

  auto partially_replicated_other = other.hlo();
  if (other_grouped && other_grouped->group_dims.size() == 1 &&
      other_grouped->group_dims[0] == other.base_shape().rank()) {
    // Group on replication dim.
    other = other.Reshard(UngroupSharding(*other_grouped));
    partially_replicated_other = other.hlo();
    top_level_sharding_to_reset.emplace_back(other.hlo(), other.sharding());
    partially_replicated_other->set_sharding(other_grouped->sharding);
  } else if (!other.sharding().IsReplicated()) {
    other = other.Reshard(UngroupSharding(*other_grouped));
    partially_replicated_other =
        other
            .Reshard(hlo_sharding_util::PartiallyReplicateTiledShardingOnDims(
                other.sharding(), other_grouped->group_dims))
            .hlo();
    top_level_sharding_to_reset.emplace_back(
        partially_replicated_other, partially_replicated_other->sharding());
    partially_replicated_other->set_sharding(other_grouped->sharding);
  }
  auto other_p = PartitionedHlo(partially_replicated_other, other.base_shape(),
                                per_group_partitioner_state);
  TF_ASSIGN_OR_RETURN(
      auto dot,
      PartitionDot(lhs_matching ? matching_p : other_p,
                   lhs_matching ? other_p : matching_p,
                   GetPerGroupBaseShape(output_grouped, output_base_shape),
                   output_grouped.sharding, dims_mapping,
                   num_partitions / matching_grouped.device_groups.size(),
                   create_sharded_dot, conv_window, module, original_hlo,
                   options, b, windowed_dot_general_loops));
  return dot;
}

std::pair<HloSharding, HloSharding>
GetDotGroupPartitionContractingOutputShardings(
    const DotConvDimsMapping& dims_mapping, const GroupedSharding& lhs_grouped,
    const Shape& output_base_shape, const HloSharding& output_sharding,
    int64 group_count, int64 output_lhs_non_contracting_partitions,
    int64 output_rhs_non_contracting_partitions, int64 output_batch_partitions,
    std::vector<int64>* output_slice_dims_out) {
  HloSharding inner_output_sharding = HloSharding::Replicate();
  HloSharding outer_output_tmp_sharding = HloSharding::Replicate();
  std::vector<int64> output_slice_dims;
  if (output_sharding.ReplicateOnLastTileDim() &&
      output_sharding.tile_assignment().dimensions().back() % group_count ==
          0) {
    auto grouped = AlignGroupsWith(
        GroupShardingOnDims(
            output_sharding,
            {output_sharding.tile_assignment().num_dimensions() - 1},
            {output_sharding.tile_assignment().dimensions().back() /
             group_count}),
        lhs_grouped,
        /*ignore_group_order=*/true);
    outer_output_tmp_sharding = UngroupSharding(grouped);
    inner_output_sharding = std::move(grouped.sharding);
  } else {
    if (auto found_dims = FindMatchingPartitionedDimsForGrouping(
            output_sharding, lhs_grouped.device_groups)) {
      output_slice_dims = std::move(*found_dims);
    } else if (output_lhs_non_contracting_partitions == group_count ||
               output_rhs_non_contracting_partitions == group_count ||
               output_batch_partitions == group_count) {
      if (output_lhs_non_contracting_partitions == group_count) {
        for (const auto& dim : dims_mapping.lhs_non_contracting_dims) {
          output_slice_dims.push_back(dim.output);
        }
      } else if (output_rhs_non_contracting_partitions == group_count) {
        for (const auto& dim : dims_mapping.rhs_non_contracting_dims) {
          output_slice_dims.push_back(dim.output);
        }
      } else {
        for (const auto& dim : dims_mapping.batch_dims) {
          output_slice_dims.push_back(dim.output);
        }
      }
    }
    if (!output_slice_dims.empty()) {
      auto grouped = AlignGroupsWith(
          GroupShardingOnDims(output_sharding, output_slice_dims), lhs_grouped);
      inner_output_sharding = grouped.sharding;
      outer_output_tmp_sharding = UngroupSharding(grouped);
    }
  }
  if (output_slice_dims_out) {
    (*output_slice_dims_out) = std::move(output_slice_dims);
  }
  return std::make_pair(inner_output_sharding, outer_output_tmp_sharding);
}

std::pair<HloSharding, HloSharding>
GetDotGroupPartitionContractingLhsRhsShardings(
    const PartitionedHlo& lhs, const PartitionedHlo& rhs,
    absl::Span<const DotConvDimsMapping::DimsMapping>
        partitioned_contracting_dims) {
  HloSharding lhs_sharding = lhs.sharding();
  HloSharding rhs_sharding = rhs.sharding();
  std::vector<int64> lhs_tile_shape =
      lhs_sharding.tile_assignment().dimensions();
  std::vector<int64> rhs_tile_shape =
      rhs_sharding.tile_assignment().dimensions();
  if (ShapeUtil::ByteSizeOf(lhs.hlo()->shape()) >
      ShapeUtil::ByteSizeOf(rhs.hlo()->shape())) {
    for (const auto& dim : partitioned_contracting_dims) {
      rhs_tile_shape[dim.rhs] = lhs_tile_shape[dim.lhs];
    }
    auto new_tile = rhs.sharding().tile_assignment();
    new_tile.Reshape(rhs_tile_shape);
    rhs_sharding = rhs_sharding.ReplicateOnLastTileDim()
                       ? HloSharding::PartialTile(new_tile)
                       : HloSharding::Tile(new_tile);
  } else {
    for (const auto& dim : partitioned_contracting_dims) {
      lhs_tile_shape[dim.lhs] = rhs_tile_shape[dim.rhs];
    }
    auto new_tile = lhs.sharding().tile_assignment();
    new_tile.Reshape(lhs_tile_shape);
    lhs_sharding = lhs_sharding.ReplicateOnLastTileDim()
                       ? HloSharding::PartialTile(new_tile)
                       : HloSharding::Tile(new_tile);
  }
  return std::make_pair(lhs_sharding, rhs_sharding);
}

StatusOr<HloInstruction*> PartitionDotGroupOnContracting(
    PartitionedHlo lhs, PartitionedHlo rhs,
    absl::Span<const DotConvDimsMapping::DimsMapping>
        partitioned_contracting_dims,
    int64 output_batch_partitions, int64 output_lhs_non_contracting_partitions,
    int64 output_rhs_non_contracting_partitions, const Shape& output_base_shape,
    const HloSharding& output_sharding, const DotConvDimsMapping& dims_mapping,
    int64 num_partitions,
    const std::function<StatusOr<HloInstruction*>(
        HloInstruction*, HloInstruction*, SpmdBuilder*,
        const Window& conv_window)>& create_sharded_dot,
    const Window& conv_window, HloModule* module, HloInstruction* original_hlo,
    bool require_matching_devices_to_group,
    const SpmdPartitionerOptions& options, SpmdBuilder* b,
    std::vector<SpmdPartitioningVisitor::WindowedDotGeneralLoop>*
        windowed_dot_general_loops) {
  std::vector<std::pair<HloInstruction*, HloSharding>>
      top_level_sharding_to_reset;
  auto cleaner = tensorflow::gtl::MakeCleanup([&] {
    for (auto& to_reset : top_level_sharding_to_reset) {
      to_reset.first->set_sharding(to_reset.second);
    }
  });
  std::vector<int64> lhs_dims;
  std::vector<int64> rhs_dims;
  int64 group_count = 1;
  for (const auto& dim : partitioned_contracting_dims) {
    lhs_dims.push_back(dim.lhs);
    rhs_dims.push_back(dim.rhs);
    group_count *= lhs.sharding().tile_assignment().dim(dim.lhs);
  }
  HloSharding lhs_sharding = HloSharding::Replicate();
  HloSharding rhs_sharding = HloSharding::Replicate();
  std::tie(lhs_sharding, rhs_sharding) =
      GetDotGroupPartitionContractingLhsRhsShardings(
          lhs, rhs, partitioned_contracting_dims);
  auto lhs_grouped = GroupShardingOnDims(lhs_sharding, lhs_dims);
  auto rhs_grouped = GroupShardingOnDims(rhs_sharding, rhs_dims);
  if (ShapeUtil::ByteSizeOf(lhs.hlo()->shape()) >
      ShapeUtil::ByteSizeOf(rhs.hlo()->shape())) {
    rhs_grouped = AlignGroupsWith(rhs_grouped, lhs_grouped);
    rhs_sharding = UngroupSharding(rhs_grouped);
    if (require_matching_devices_to_group && rhs.sharding() != rhs_sharding) {
      return nullptr;
    }
    rhs = rhs.Reshard(rhs_sharding);
  } else {
    lhs_grouped = AlignGroupsWith(lhs_grouped, rhs_grouped);
    lhs_sharding = UngroupSharding(lhs_grouped);
    if (require_matching_devices_to_group && lhs.sharding() != lhs_sharding) {
      return nullptr;
    }
    lhs = lhs.Reshard(lhs_sharding);
  }
  // Mask out invalid data.
  std::vector<int64> lhs_skipped_dims;
  for (int64 i = 0; i < lhs.base_shape().rank(); ++i) {
    if (absl::c_linear_search(lhs_dims, i)) {
      continue;
    }
    lhs_skipped_dims.push_back(i);
  }
  lhs = lhs.PadWithValue(
      CreateZero(ShapeUtil::MakeShape(lhs.base_shape().element_type(), {}), b),
      /*left_padded_dims=*/{}, lhs_skipped_dims);
  std::vector<int64> rhs_skipped_dims;
  for (int64 i = 0; i < rhs.base_shape().rank(); ++i) {
    if (absl::c_linear_search(rhs_dims, i)) {
      continue;
    }
    rhs_skipped_dims.push_back(i);
  }
  rhs = rhs.PadWithValue(
      CreateZero(ShapeUtil::MakeShape(rhs.base_shape().element_type(), {}), b),
      /*left_padded_dims=*/{}, rhs_skipped_dims);
  top_level_sharding_to_reset.emplace_back(lhs.hlo(), lhs_sharding);
  lhs.hlo()->set_sharding(lhs_grouped.sharding);
  top_level_sharding_to_reset.emplace_back(rhs.hlo(), rhs_sharding);
  rhs.hlo()->set_sharding(rhs_grouped.sharding);

  HloSharding inner_output_sharding = HloSharding::Replicate();
  HloSharding outer_output_tmp_sharding = HloSharding::Replicate();
  std::vector<int64> output_slice_dims;
  std::tie(inner_output_sharding, outer_output_tmp_sharding) =
      GetDotGroupPartitionContractingOutputShardings(
          dims_mapping, lhs_grouped, output_base_shape, output_sharding,
          group_count, output_lhs_non_contracting_partitions,
          output_rhs_non_contracting_partitions, output_batch_partitions,
          &output_slice_dims);
  Shape inner_output_base_shape = output_base_shape;
  auto get_non_slice_dims = [&] {
    std::vector<int64> non_group_dims;
    for (int64 i = 0; i < output_base_shape.rank(); ++i) {
      if (!absl::c_linear_search(output_slice_dims, i)) {
        non_group_dims.push_back(i);
      }
    }
    return non_group_dims;
  };
  if (!output_slice_dims.empty()) {
    inner_output_base_shape = MakePartitionedShape(
        output_base_shape,
        hlo_sharding_util::PartiallyReplicateTiledShardingOnDims(
            output_sharding, get_non_slice_dims()));
  }
  std::function<StatusOr<HloInstruction*>(HloInstruction*, HloInstruction*,
                                          SpmdBuilder*, const Window&)>
      inner_creator =
          [&](HloInstruction* l, HloInstruction* r, SpmdBuilder* b,
              const Window& conv_window) -> StatusOr<HloInstruction*> {
    TF_ASSIGN_OR_RETURN(auto inner_dot,
                        create_sharded_dot(l, r, b, conv_window));
    auto ar = lhs.state().partitioner->AllReduceAlongShardingDims(
        b, inner_dot, lhs_sharding, lhs.state().next_channel_id, lhs_dims,
        lhs.state().collective_ops_creator,
        MakeBinaryAdd(output_base_shape.element_type(), module));
    if (output_slice_dims.empty()) {
      return ar;
    }
    // Use resharding to slice the output. Use a temporary reshard cache since
    // we are faking with replicated sharding.
    PartitionedHlo::PartitioningState new_state = lhs.state();
    new_state.b = b;
    new_state.partition_id =
        lhs.state().collective_ops_creator.create_partition_id(b);
    PartitionedHlo::ReshardCache tmp_cache;
    new_state.reshard_cache = &tmp_cache;
    ar->set_sharding(HloSharding::Replicate());
    return PartitionedHlo(ar, ar->shape(), new_state)
        .Reshard(hlo_sharding_util::PartiallyReplicateTiledShardingOnDims(
            output_sharding, get_non_slice_dims()))
        .hlo();
  };
  // Disable doing the inner reshard when the "faster windowed einsum" flag is
  // enabled, because the windowed einsum implementation is currently slow with
  // this kind of reshard happening.
  if (options.choose_faster_windowed_einsum_over_mem) {
    inner_output_base_shape = output_base_shape;
    inner_creator = create_sharded_dot;
    outer_output_tmp_sharding =
        hlo_sharding_util::PartiallyReplicateTiledShardingOnDims(
            outer_output_tmp_sharding, output_slice_dims);
  }
  PartitionedHlo::PartitioningState inner_state =
      CreatePerGroupPartitioningState(lhs.state(), lhs_grouped.device_groups,
                                      b);
  TF_ASSIGN_OR_RETURN(
      auto dot,
      PartitionDot(
          PartitionedHlo(lhs.hlo(),
                         GetPerGroupBaseShape(lhs_grouped, lhs.base_shape()),
                         inner_state),
          PartitionedHlo(rhs.hlo(),
                         GetPerGroupBaseShape(rhs_grouped, rhs.base_shape()),
                         inner_state),
          inner_output_base_shape, inner_output_sharding, dims_mapping,
          num_partitions / group_count, inner_creator, conv_window, module,
          original_hlo, options, b, windowed_dot_general_loops));
  if (!dot) {
    return nullptr;
  }

  if (options.choose_faster_windowed_einsum_over_mem) {
    HloInstruction* ar = lhs.state().partitioner->AllReduceAlongShardingDims(
        b, dot, lhs_sharding, lhs.state().next_channel_id, lhs_dims,
        lhs.state().collective_ops_creator,
        MakeBinaryAdd(output_base_shape.element_type(), module));
    dot = ar;
  }

  dot->set_sharding(outer_output_tmp_sharding);
  auto d = PartitionedHlo(dot, output_base_shape, lhs.state())
               .Reshard(output_sharding)
               .hlo();
  return d;
}

DotConvDimsMapping ConvertDimsMappingWithFeatureGroupCount(
    const DotConvDimsMapping& dims_mapping, HloInstruction* original_hlo) {
  const auto& dnums = original_hlo->convolution_dimension_numbers();
  DotConvDimsMapping new_dims_mapping;
  new_dims_mapping.batch_dims = dims_mapping.batch_dims;
  new_dims_mapping.conv_spatial_dims = dims_mapping.conv_spatial_dims;
  // Append batch dims.
  new_dims_mapping.batch_dims.emplace_back();
  new_dims_mapping.batch_dims.back().lhs = dnums.input_feature_dimension();
  new_dims_mapping.batch_dims.back().rhs =
      dnums.kernel_output_feature_dimension();
  new_dims_mapping.batch_dims.back().output = dnums.output_feature_dimension();
  new_dims_mapping.batch_dims.back().spatial = -1;
  // Setup non contracting dims.
  new_dims_mapping.lhs_non_contracting_dims.emplace_back();
  new_dims_mapping.lhs_non_contracting_dims.back().lhs =
      dnums.input_batch_dimension();
  new_dims_mapping.rhs_non_contracting_dims.emplace_back();
  new_dims_mapping.rhs_non_contracting_dims.back().rhs =
      dnums.kernel_input_feature_dimension();
  return new_dims_mapping;
}

DotConvDimsMapping ConvertDimsMappingWithBatchGroupCount(
    const DotConvDimsMapping& dims_mapping, HloInstruction* original_hlo) {
  const auto& dnums = original_hlo->convolution_dimension_numbers();
  DotConvDimsMapping new_dims_mapping;
  new_dims_mapping.batch_dims = dims_mapping.batch_dims;
  new_dims_mapping.conv_spatial_dims = dims_mapping.conv_spatial_dims;
  new_dims_mapping.contracting_dims = dims_mapping.contracting_dims;
  // Append batch dims.
  new_dims_mapping.batch_dims.emplace_back();
  new_dims_mapping.batch_dims.back().lhs = dnums.input_batch_dimension();
  new_dims_mapping.batch_dims.back().rhs =
      dnums.kernel_output_feature_dimension();
  new_dims_mapping.batch_dims.back().output = dnums.output_feature_dimension();
  new_dims_mapping.batch_dims.back().spatial = -1;
  return new_dims_mapping;
}

// Estimate the number of iterations of a subsequent windowed einsum
// partitioning if its partitioned in the non-contracting dimensions.
// First value returned is the estimate of the number of iterations if LHS is
// matched while the second is the number of iterations if RHS is matched.
std::pair<absl::optional<int64>, absl::optional<int64>>
EstimateWindowedEinsumIterationsForNonContractingPartitioning(
    const DotConvDimsMapping& dims_mapping, const PartitionedHlo& lhs,
    const PartitionedHlo& rhs, const Shape& output_base_shape,
    const HloSharding& output_sharding, const SpmdPartitionerOptions& options,
    int64 num_partitions, int64 lhs_non_contracting_partitions,
    int64 rhs_non_contracting_partitions, int64 lhs_matching_partitions,
    int64 rhs_matching_partitions, int64 lhs_contracting_partitions,
    int64 rhs_contracting_partitions,
    int64 output_lhs_non_contracting_partitions,
    int64 output_rhs_non_contracting_partitions, int64 lhs_batch_partitions,
    int64 rhs_batch_partitions) {
  const DotDimensionIndexMapping indices_map = ComputeDimensionIndexMapping(
      dims_mapping, lhs.base_shape().rank(), rhs.base_shape().rank(),
      output_base_shape.rank());
  auto subsequent_einsum_iterations_estimate =
      [&](bool assume_lhs_match) -> absl::optional<int64> {
    const std::vector<DotConvDimsMapping::DimsMapping>&
        matching_non_contracting_dims =
            assume_lhs_match ? dims_mapping.lhs_non_contracting_dims
                             : dims_mapping.rhs_non_contracting_dims;
    const std::vector<DotConvDimsMapping::DimsMapping>&
        other_non_contracting_dims =
            assume_lhs_match ? dims_mapping.rhs_non_contracting_dims
                             : dims_mapping.lhs_non_contracting_dims;
    const std::vector<int64>& output_to_matching_indices =
        assume_lhs_match ? indices_map.output_to_lhs_indices
                         : indices_map.output_to_rhs_indices;
    const std::vector<int64>& output_to_other_indices =
        assume_lhs_match ? indices_map.output_to_rhs_indices
                         : indices_map.output_to_lhs_indices;
    const std::vector<int64>& matching_to_output_indices =
        assume_lhs_match ? indices_map.lhs_to_output_indices
                         : indices_map.rhs_to_output_indices;
    const std::vector<int64>& other_to_output_indices =
        assume_lhs_match ? indices_map.rhs_to_output_indices
                         : indices_map.lhs_to_output_indices;
    const HloSharding& matching_sharding =
        assume_lhs_match ? lhs.sharding() : rhs.sharding();
    const HloSharding& other_sharding =
        assume_lhs_match ? rhs.sharding() : lhs.sharding();
    const PartitionedHlo& matching_partitioned = assume_lhs_match ? lhs : rhs;
    const PartitionedHlo& other_partitioned = assume_lhs_match ? rhs : lhs;
    const int64 matching_non_contracting_partitions =
        assume_lhs_match ? lhs_non_contracting_partitions
                         : rhs_non_contracting_partitions;
    const int64 other_non_contracting_partitions =
        assume_lhs_match ? rhs_non_contracting_partitions
                         : lhs_non_contracting_partitions;
    const int64 matching_contracting_partitions =
        assume_lhs_match ? lhs_contracting_partitions
                         : rhs_contracting_partitions;
    const int64 other_contracting_partitions = assume_lhs_match
                                                   ? rhs_contracting_partitions
                                                   : lhs_contracting_partitions;
    const int64 output_matching_non_contracting_partitions =
        assume_lhs_match ? output_lhs_non_contracting_partitions
                         : output_rhs_non_contracting_partitions;
    const int64 output_other_non_contracting_partitions =
        assume_lhs_match ? output_rhs_non_contracting_partitions
                         : output_lhs_non_contracting_partitions;
    const int64 matching_batch_partitions =
        assume_lhs_match ? lhs_batch_partitions : rhs_batch_partitions;
    const int64 other_batch_partitions =
        assume_lhs_match ? rhs_batch_partitions : lhs_batch_partitions;
    const int64 matching_matched_non_contracting_partitions =
        assume_lhs_match ? lhs_non_contracting_partitions
                         : rhs_non_contracting_partitions;
    std::vector<int64> output_dims;
    output_dims.reserve(matching_non_contracting_dims.size());
    for (const DotConvDimsMapping::DimsMapping& dim :
         matching_non_contracting_dims) {
      output_dims.push_back(dim.output);
    }
    GroupedSharding output_grouped =
        GroupShardingOnDims(output_sharding, output_dims);
    GroupedSharding matching_grouped =
        GetNonContractingPartitionGroupedShardingForMatchedOperand(
            assume_lhs_match, matching_sharding, output_sharding,
            matching_non_contracting_dims);
    absl::optional<GroupedSharding> other_grouped =
        GetNonContractingPartitionGroupedShardingForOtherOperand(
            assume_lhs_match, output_base_shape,
            other_partitioned.hlo()->shape(), other_contracting_partitions,
            other_non_contracting_partitions, matching_contracting_partitions,
            output_other_non_contracting_partitions, other_sharding,
            output_sharding, matching_non_contracting_dims,
            other_non_contracting_dims, dims_mapping.contracting_dims);
    if (!other_grouped) {
      return absl::nullopt;
    }
    absl::optional<HloSharding> output_sharding_transposed_to_match_matching =
        hlo_sharding_util::TransposeShardingWithCollapsedDims(
            output_grouped.sharding, output_to_matching_indices,
            matching_to_output_indices);
    absl::optional<HloSharding> output_sharding_transposed_to_match_other =
        hlo_sharding_util::TransposeShardingWithCollapsedDims(
            output_grouped.sharding, output_to_other_indices,
            other_to_output_indices);
    const int64 new_num_partitions =
        num_partitions / matching_non_contracting_partitions;
    absl::optional<WindowedEinsumConfig> e_config =
        GetWindowedEinsumConfiguration(
            new_num_partitions, output_matching_non_contracting_partitions,
            output_other_non_contracting_partitions,
            other_contracting_partitions, other_non_contracting_partitions,
            other_batch_partitions, matching_contracting_partitions,
            matching_non_contracting_partitions /
                matching_matched_non_contracting_partitions,
            matching_batch_partitions,
            ShapeSizeInBytes(other_partitioned.base_shape()),
            ShapeSizeInBytes(matching_partitioned.base_shape()) /
                matching_non_contracting_partitions,
            ShapeSizeInBytes(
                GetPerGroupBaseShape(output_grouped, output_base_shape)),
            options.threshold_for_windowed_einsum_mib,
            output_sharding_transposed_to_match_matching,
            output_sharding_transposed_to_match_other,
            matching_grouped.sharding, other_grouped->sharding);
    return e_config ? new_num_partitions : absl::optional<int64>(absl::nullopt);
  };
  absl::optional<int64> lhs_matching_iterations;
  if (lhs_matching_partitions != 0) {
    lhs_matching_iterations = subsequent_einsum_iterations_estimate(true);
  }
  absl::optional<int64> rhs_matching_iterations;
  if (rhs_matching_partitions != 0) {
    rhs_matching_iterations = subsequent_einsum_iterations_estimate(false);
  }
  return std::make_pair(lhs_matching_iterations, rhs_matching_iterations);
}

// Return if we should prioritize partitioning in the contracting dimensions
// first then non-contracting dimensions if we estimate that would allow
// for a fewer number of iterations of the windowed einsum.
bool PrioritizeContractingDimensionsPartitioning(
    const DotConvDimsMapping& dims_mapping, const PartitionedHlo& lhs,
    const PartitionedHlo& rhs, const Shape& output_base_shape,
    const HloSharding& output_sharding, const SpmdPartitionerOptions& options,
    int64 num_partitions, int64 lhs_non_contracting_partitions,
    int64 rhs_non_contracting_partitions, int64 lhs_contracting_partitions,
    int64 rhs_contracting_partitions,
    int64 output_lhs_non_contracting_partitions,
    int64 output_rhs_non_contracting_partitions, int64 lhs_batch_partitions,
    int64 rhs_batch_partitions, int64 output_batch_partitions,
    bool require_matching_devices_to_group) {
  const bool may_group_on_lhs_non_contracting =
      lhs_non_contracting_partitions == output_lhs_non_contracting_partitions &&
      lhs_non_contracting_partitions > 1;
  const bool may_group_on_rhs_non_contracting =
      rhs_non_contracting_partitions == output_rhs_non_contracting_partitions &&
      rhs_non_contracting_partitions > 1;
  if (!options.choose_faster_windowed_einsum_over_mem) {
    return false;
  }
  // Check only for perfect dimensions match for now.
  if (!may_group_on_lhs_non_contracting && !may_group_on_rhs_non_contracting) {
    return false;
  }
  absl::optional<int64> lhs_matching_iterations;
  absl::optional<int64> rhs_matching_iterations;
  const int64 lhs_matching_non_contracting_partitions =
      may_group_on_lhs_non_contracting ? lhs_non_contracting_partitions : 0;
  const int64 rhs_matching_non_contracting_partitions =
      may_group_on_rhs_non_contracting ? rhs_non_contracting_partitions : 0;
  std::tie(lhs_matching_iterations, rhs_matching_iterations) =
      EstimateWindowedEinsumIterationsForNonContractingPartitioning(
          dims_mapping, lhs, rhs, output_base_shape, output_sharding, options,
          num_partitions, lhs_non_contracting_partitions,
          rhs_non_contracting_partitions,
          lhs_matching_non_contracting_partitions,
          rhs_matching_non_contracting_partitions, lhs_contracting_partitions,
          rhs_contracting_partitions, output_lhs_non_contracting_partitions,
          output_rhs_non_contracting_partitions, lhs_batch_partitions,
          rhs_batch_partitions);
  if (!lhs_matching_iterations && !rhs_matching_iterations) {
    return false;
  }
  // Be conservative and handle only case where the two partitions in rhs and
  // lhs match
  if (!(lhs_contracting_partitions == rhs_contracting_partitions &&
        lhs_contracting_partitions > 1)) {
    return false;
  }
  // Estimate the iterations in the case we perform the partitioning on the
  // contracting dimensions instead.
  std::vector<int64> lhs_dims;
  std::vector<int64> rhs_dims;
  int64 group_count = 1;
  for (const auto& dim : dims_mapping.contracting_dims) {
    lhs_dims.push_back(dim.lhs);
    rhs_dims.push_back(dim.rhs);
    group_count *= lhs.sharding().tile_assignment().dim(dim.lhs);
  }
  HloSharding lhs_sharding = HloSharding::Replicate();
  HloSharding rhs_sharding = HloSharding::Replicate();
  std::tie(lhs_sharding, rhs_sharding) =
      GetDotGroupPartitionContractingLhsRhsShardings(
          lhs, rhs, dims_mapping.contracting_dims);
  auto lhs_grouped = GroupShardingOnDims(lhs_sharding, lhs_dims);
  auto rhs_grouped = GroupShardingOnDims(rhs_sharding, rhs_dims);
  rhs_grouped = AlignGroupsWith(rhs_grouped, lhs_grouped);
  rhs_sharding = UngroupSharding(rhs_grouped);

  if (require_matching_devices_to_group && rhs.sharding() != rhs_sharding) {
    return false;
  }
  const int64 new_num_partitions = num_partitions / lhs_contracting_partitions;

  HloSharding inner_output_sharding = HloSharding::Replicate();
  HloSharding outer_output_tmp_sharding = HloSharding::Replicate();
  std::vector<int64> output_slice_dims;
  std::tie(inner_output_sharding, outer_output_tmp_sharding) =
      GetDotGroupPartitionContractingOutputShardings(
          dims_mapping, lhs_grouped, output_base_shape, output_sharding,
          group_count, output_lhs_non_contracting_partitions,
          output_rhs_non_contracting_partitions, output_batch_partitions,
          &output_slice_dims);
  Shape inner_output_base_shape = output_base_shape;
  if (!output_slice_dims.empty()) {
    std::vector<int64> non_group_dims;
    for (int64 i = 0; i < output_base_shape.rank(); ++i) {
      if (!absl::c_linear_search(output_slice_dims, i)) {
        non_group_dims.push_back(i);
      }
    }
    inner_output_base_shape = MakePartitionedShape(
        output_base_shape,
        hlo_sharding_util::PartiallyReplicateTiledShardingOnDims(
            output_sharding, non_group_dims));
  }
  int64 new_output_lhs_non_contracting_partitions = 1;
  int64 new_output_rhs_non_contracting_partitions = 1;
  if (!inner_output_sharding.IsTileMaximal()) {
    for (const auto& dim : dims_mapping.lhs_non_contracting_dims) {
      new_output_lhs_non_contracting_partitions *=
          inner_output_sharding.tile_assignment().dim(dim.output);
    }
    for (const auto& dim : dims_mapping.rhs_non_contracting_dims) {
      if (dim.output != -1) {
        new_output_rhs_non_contracting_partitions *=
            inner_output_sharding.tile_assignment().dim(dim.output);
      }
    }
  }

  const DotDimensionIndexMapping indices_map = ComputeDimensionIndexMapping(
      dims_mapping, lhs.base_shape().rank(), rhs.base_shape().rank(),
      inner_output_base_shape.rank());
  absl::optional<HloSharding> output_sharding_transposed_to_match_lhs =
      hlo_sharding_util::TransposeShardingWithCollapsedDims(
          inner_output_sharding, indices_map.output_to_lhs_indices,
          indices_map.lhs_to_output_indices);
  absl::optional<HloSharding> output_sharding_transposed_to_match_rhs =
      hlo_sharding_util::TransposeShardingWithCollapsedDims(
          inner_output_sharding, indices_map.output_to_rhs_indices,
          indices_map.rhs_to_output_indices);
  absl::optional<WindowedEinsumConfig> e_config =
      GetWindowedEinsumConfiguration(
          new_num_partitions, new_output_lhs_non_contracting_partitions,
          new_output_rhs_non_contracting_partitions, 1,
          rhs_non_contracting_partitions, rhs_batch_partitions, 1,
          lhs_non_contracting_partitions, lhs_batch_partitions,
          ShapeSizeInBytes(GetPerGroupBaseShape(rhs_grouped, rhs.base_shape())),
          ShapeSizeInBytes(GetPerGroupBaseShape(lhs_grouped, lhs.base_shape())),
          ShapeSizeInBytes(inner_output_base_shape),
          options.threshold_for_windowed_einsum_mib,
          output_sharding_transposed_to_match_lhs,
          output_sharding_transposed_to_match_rhs, lhs_grouped.sharding,
          rhs_grouped.sharding);
  if (!e_config) {
    return false;
  }
  const int64 min_nc_iterations =
      std::min(lhs_matching_iterations ? *lhs_matching_iterations : INT64_MAX,
               rhs_matching_iterations ? *rhs_matching_iterations : INT64_MAX);
  return min_nc_iterations > new_num_partitions;
}

// Return if it would be better to match the LHS operand or RHS operand
// of a dot for non-contracting partitioning.
bool LhsIsBestMatchForNonContractingPartitioning(
    const DotConvDimsMapping& dims_mapping, const PartitionedHlo& lhs,
    const PartitionedHlo& rhs, const Shape& output_base_shape,
    const HloSharding& output_sharding, const SpmdPartitionerOptions& options,
    int64 num_partitions, int64 lhs_non_contracting_partitions,
    int64 rhs_non_contracting_partitions, int64 lhs_matching_partitions,
    int64 rhs_matching_partitions, int64 lhs_contracting_partitions,
    int64 rhs_contracting_partitions,
    int64 output_lhs_non_contracting_partitions,
    int64 output_rhs_non_contracting_partitions, int64 lhs_batch_partitions,
    int64 rhs_batch_partitions) {
  const bool may_group_on_lhs_non_contracting =
      lhs_non_contracting_partitions == output_lhs_non_contracting_partitions &&
      lhs_non_contracting_partitions > 1;
  const bool may_group_on_rhs_non_contracting =
      rhs_non_contracting_partitions == output_rhs_non_contracting_partitions &&
      rhs_non_contracting_partitions > 1;
  // If both match output non-contracting dimensions, choose the one which
  // will result in smaller replication of the other operand.
  bool lhs_matching = may_group_on_lhs_non_contracting &&
                      (!may_group_on_rhs_non_contracting ||
                       lhs_non_contracting_partitions *
                               ShapeUtil::ByteSizeOf(rhs.hlo()->shape()) <
                           rhs_non_contracting_partitions *
                               ShapeUtil::ByteSizeOf(lhs.hlo()->shape()));
  // If both grouping are available and the option to choose faster windowed
  // einsums vs saving memory is enabled then try to determine which of the
  // operands will generate the least amount of iterations for the windowed
  // einsum when matched (if a windowed einsum is gonna be generated at
  // all).
  if (may_group_on_lhs_non_contracting && may_group_on_rhs_non_contracting &&
      options.choose_faster_windowed_einsum_over_mem) {
    const DotDimensionIndexMapping indices_map = ComputeDimensionIndexMapping(
        dims_mapping, lhs.base_shape().rank(), rhs.base_shape().rank(),
        output_base_shape.rank());
    absl::optional<int64> lhs_matching_iterations;
    absl::optional<int64> rhs_matching_iterations;
    std::tie(lhs_matching_iterations, rhs_matching_iterations) =
        EstimateWindowedEinsumIterationsForNonContractingPartitioning(
            dims_mapping, lhs, rhs, output_base_shape, output_sharding, options,
            num_partitions, lhs_non_contracting_partitions,
            rhs_non_contracting_partitions, lhs_matching_partitions,
            rhs_matching_partitions, lhs_contracting_partitions,
            rhs_contracting_partitions, output_lhs_non_contracting_partitions,
            output_rhs_non_contracting_partitions, lhs_batch_partitions,
            rhs_batch_partitions);
    if (lhs_matching_iterations && rhs_matching_iterations &&
        *lhs_matching_iterations != *rhs_matching_iterations) {
      lhs_matching = *lhs_matching_iterations < *rhs_matching_iterations;
    }
  }
  return lhs_matching;
}

// Recursive partitioning function. If there are partial dimensions matching
// in the operands and output, group the devices and recursively partition
// the in-group dot.
StatusOr<HloInstruction*> PartitionDot(
    PartitionedHlo lhs, PartitionedHlo rhs, const Shape& output_base_shape,
    const HloSharding& output_sharding, const DotConvDimsMapping& dims_mapping,
    int64 num_partitions,
    const std::function<StatusOr<HloInstruction*>(
        HloInstruction*, HloInstruction*, SpmdBuilder*,
        const Window& conv_window)>& create_sharded_dot,
    const Window& conv_window, HloModule* module, HloInstruction* original_hlo,
    bool require_matching_devices_to_group,
    const SpmdPartitionerOptions& options, SpmdBuilder* b,
    std::vector<SpmdPartitioningVisitor::WindowedDotGeneralLoop>*
        windowed_dot_general_loops) {
  // If lhs' hlo and rhs' hlo are identical, make a copy for rhs.
  if (lhs.hlo() == rhs.hlo()) {
    auto copy_hlo = b->AddInstruction(HloInstruction::CreateUnary(
        rhs.hlo()->shape(), HloOpcode::kCopy, rhs.hlo()));
    copy_hlo->set_sharding(rhs.sharding());
    rhs = PartitionedHlo(copy_hlo, rhs.base_shape(), rhs.state());
  }

  // lhs_rhs_or_output: 0 lhs, 1 rhs, 2 output.
  auto get_partitions_for_dims =
      [&](const HloSharding& sharding,
          absl::Span<const DotConvDimsMapping::DimsMapping> dims,
          int lhs_rhs_or_output) {
        int64 partitions = 1;
        if (sharding.IsTileMaximal()) {
          return partitions;
        }
        for (const auto& dim : dims) {
          if (lhs_rhs_or_output == 0) {
            partitions *= sharding.tile_assignment().dim(dim.lhs);
          } else if (lhs_rhs_or_output == 1) {
            partitions *= sharding.tile_assignment().dim(dim.rhs);
          } else {
            CHECK_EQ(lhs_rhs_or_output, 2);
            partitions *= sharding.tile_assignment().dim(dim.output);
          }
        }
        return partitions;
      };
  const int64 lhs_batch_partitions =
      get_partitions_for_dims(lhs.sharding(), dims_mapping.batch_dims, 0);
  const int64 rhs_batch_partitions =
      get_partitions_for_dims(rhs.sharding(), dims_mapping.batch_dims, 1);
  const int64 output_batch_partitions =
      get_partitions_for_dims(output_sharding, dims_mapping.batch_dims, 2);
  const int64 lhs_contracting_partitions =
      get_partitions_for_dims(lhs.sharding(), dims_mapping.contracting_dims, 0);
  const int64 rhs_contracting_partitions =
      get_partitions_for_dims(rhs.sharding(), dims_mapping.contracting_dims, 1);
  const int64 lhs_non_contracting_partitions = get_partitions_for_dims(
      lhs.sharding(), dims_mapping.lhs_non_contracting_dims, 0);
  const int64 rhs_non_contracting_partitions = get_partitions_for_dims(
      rhs.sharding(), dims_mapping.rhs_non_contracting_dims, 1);
  const int64 output_lhs_non_contracting_partitions = get_partitions_for_dims(
      output_sharding, dims_mapping.lhs_non_contracting_dims, 2);
  const int64 output_rhs_non_contracting_partitions = get_partitions_for_dims(
      output_sharding, dims_mapping.rhs_non_contracting_dims, 2);
  const int64 lhs_conv_spatial_partitions = get_partitions_for_dims(
      lhs.sharding(), dims_mapping.conv_spatial_dims, 0);
  const int64 rhs_conv_spatial_partitions = get_partitions_for_dims(
      rhs.sharding(), dims_mapping.conv_spatial_dims, 1);
  const int64 output_conv_spatial_partitions = get_partitions_for_dims(
      output_sharding, dims_mapping.conv_spatial_dims, 2);
  // Before we find partial matches along the dimensions, invoke base case
  // again without may_reshard_without_detecting_match.

  // Try partition the purely spatially-partitioned convolution with
  // convolution spatial dimension partitioned or depthwise parallel
  // dimension partitioned.
  bool is_conv_spatial_dim_partitioned =
      (lhs_conv_spatial_partitions > 1 || rhs_conv_spatial_partitions > 1 ||
       output_conv_spatial_partitions > 1);
  bool is_conv_batch_or_contracting_dim_partitioned =
      (lhs_batch_partitions > 1 || rhs_batch_partitions > 1 ||
       output_batch_partitions > 1 ||
       (lhs_contracting_partitions > 1 && rhs_contracting_partitions > 1));
  if ((!dims_mapping.conv_spatial_dims.empty() &&
       is_conv_spatial_dim_partitioned &&
       !is_conv_batch_or_contracting_dim_partitioned) ||
      (original_hlo->opcode() == HloOpcode::kConvolution &&
       (original_hlo->batch_group_count() > 1 ||
        original_hlo->feature_group_count() > 1))) {
    // Partition with kernel_input_feature_dim > 1 and feature_group_count >
    // 1 is not supported.
    const auto& dnums = original_hlo->convolution_dimension_numbers();
    if (original_hlo->feature_group_count() > 1 &&
        rhs.hlo()->shape().dimensions(dnums.kernel_input_feature_dimension()) >
            1) {
      return nullptr;
    }

    TF_ASSIGN_OR_RETURN(
        auto partitioned_conv,
        PartitionConvolution(lhs, rhs, output_base_shape, output_sharding,
                             dims_mapping, create_sharded_dot, conv_window,
                             original_hlo, num_partitions, options,
                             lhs.state().partition_id, module, b));

    if (partitioned_conv) {
      return partitioned_conv;
    }

    // Recursively partition on different types of dimensions for
    // convolution. Case 0.a: Group partitions by feature group count.
    if (original_hlo->feature_group_count() > 1 ||
        original_hlo->batch_group_count() > 1) {
      DotConvDimsMapping new_dims_mapping;
      if (original_hlo->feature_group_count() > 1) {
        new_dims_mapping =
            ConvertDimsMappingWithFeatureGroupCount(dims_mapping, original_hlo);
      }

      if (original_hlo->batch_group_count() > 1) {
        new_dims_mapping =
            ConvertDimsMappingWithBatchGroupCount(dims_mapping, original_hlo);
      }

      const int64 conv_lhs_contracting_partitions = get_partitions_for_dims(
          lhs.sharding(), new_dims_mapping.contracting_dims, 0);
      const int64 conv_rhs_contracting_partitions = get_partitions_for_dims(
          rhs.sharding(), new_dims_mapping.contracting_dims, 1);
      const int64 conv_lhs_non_contracting_partitions = get_partitions_for_dims(
          lhs.sharding(), new_dims_mapping.lhs_non_contracting_dims, 0);
      const int64 conv_rhs_non_contracting_partitions = get_partitions_for_dims(
          rhs.sharding(), new_dims_mapping.rhs_non_contracting_dims, 1);
      const int64 conv_lhs_batch_partitions = get_partitions_for_dims(
          lhs.sharding(), new_dims_mapping.batch_dims, 0);
      const int64 conv_rhs_batch_partitions = get_partitions_for_dims(
          rhs.sharding(), new_dims_mapping.batch_dims, 1);
      const int64 conv_output_batch_partitions = get_partitions_for_dims(
          output_sharding, new_dims_mapping.batch_dims, 2);
      if ((conv_lhs_batch_partitions == conv_output_batch_partitions ||
           conv_rhs_batch_partitions == conv_output_batch_partitions) &&
          conv_output_batch_partitions > 1) {
        TF_ASSIGN_OR_RETURN(
            auto try_partitioned_conv,
            PartitionDotGroupOnBatch(
                lhs, rhs, output_base_shape, output_sharding, new_dims_mapping,
                num_partitions, conv_lhs_contracting_partitions,
                conv_rhs_contracting_partitions,
                conv_lhs_non_contracting_partitions,
                conv_rhs_non_contracting_partitions, create_sharded_dot,
                conv_window, module, original_hlo,
                require_matching_devices_to_group, options, b,
                windowed_dot_general_loops));
        if (try_partitioned_conv) {
          return try_partitioned_conv;
        }
      }
      return nullptr;
    }
  }

  TF_ASSIGN_OR_RETURN(
      auto try_partitioned_dot,
      PartitionBaseCase(
          lhs, rhs, output_base_shape, output_sharding, dims_mapping,
          num_partitions, create_sharded_dot, conv_window, module, original_hlo,
          lhs_batch_partitions, rhs_batch_partitions, output_batch_partitions,
          lhs_contracting_partitions, rhs_contracting_partitions,
          lhs_non_contracting_partitions, rhs_non_contracting_partitions,
          output_lhs_non_contracting_partitions,
          output_rhs_non_contracting_partitions, options, b,
          windowed_dot_general_loops,
          /*may_reshard_without_detecting_match=*/false));
  if (try_partitioned_dot) {
    return try_partitioned_dot;
  }

  // Recursively partition on different types of dimensions.
  //
  // Case 1: Group partitions by batch.
  if ((lhs_batch_partitions == output_batch_partitions ||
       rhs_batch_partitions == output_batch_partitions) &&
      output_batch_partitions > 1) {
    TF_ASSIGN_OR_RETURN(
        auto dot,
        PartitionDotGroupOnBatch(
            lhs, rhs, output_base_shape, output_sharding, dims_mapping,
            num_partitions, lhs_contracting_partitions,
            rhs_contracting_partitions, lhs_non_contracting_partitions,
            rhs_non_contracting_partitions, create_sharded_dot, conv_window,
            module, original_hlo, require_matching_devices_to_group, options, b,
            windowed_dot_general_loops));
    if (dot) {
      return dot;
    }
  }

  // Case 2: Group partitions by non-contracting dimensions.
  const bool may_group_on_lhs_non_contracting =
      lhs_non_contracting_partitions == output_lhs_non_contracting_partitions &&
      lhs_non_contracting_partitions > 1;
  const bool may_group_on_rhs_non_contracting =
      rhs_non_contracting_partitions == output_rhs_non_contracting_partitions &&
      rhs_non_contracting_partitions > 1;
  bool lhs_matching = false;
  std::vector<DotConvDimsMapping::DimsMapping> matching_dims;
  if (may_group_on_lhs_non_contracting || may_group_on_rhs_non_contracting) {
    lhs_matching = LhsIsBestMatchForNonContractingPartitioning(
        dims_mapping, lhs, rhs, output_base_shape, output_sharding, options,
        num_partitions, lhs_non_contracting_partitions,
        rhs_non_contracting_partitions, lhs_non_contracting_partitions,
        rhs_non_contracting_partitions, lhs_contracting_partitions,
        rhs_contracting_partitions, output_lhs_non_contracting_partitions,
        output_rhs_non_contracting_partitions, lhs_batch_partitions,
        rhs_batch_partitions);
    matching_dims = lhs_matching ? dims_mapping.lhs_non_contracting_dims
                                 : dims_mapping.rhs_non_contracting_dims;
  } else if (lhs_non_contracting_partitions > 1 &&
             output_lhs_non_contracting_partitions > 1) {
    lhs_matching = true;
    for (const auto& dim : dims_mapping.lhs_non_contracting_dims) {
      int64 lhs_partitions = lhs.sharding().tile_assignment().dim(dim.lhs);
      if (lhs_partitions > 1 &&
          lhs_partitions == output_sharding.tile_assignment().dim(dim.output)) {
        matching_dims.push_back(dim);
      }
    }
  } else if (rhs_non_contracting_partitions > 1 &&
             output_rhs_non_contracting_partitions > 1) {
    lhs_matching = false;
    for (const auto& dim : dims_mapping.rhs_non_contracting_dims) {
      int64 rhs_partitions = rhs.sharding().tile_assignment().dim(dim.rhs);
      if (rhs_partitions > 1 &&
          rhs_partitions == output_sharding.tile_assignment().dim(dim.output)) {
        matching_dims.push_back(dim);
      }
    }
  }
  const bool prioritize_contracting_for_faster_windowed_einsum =
      PrioritizeContractingDimensionsPartitioning(
          dims_mapping, lhs, rhs, output_base_shape, output_sharding, options,
          num_partitions, lhs_non_contracting_partitions,
          rhs_non_contracting_partitions, lhs_contracting_partitions,
          rhs_contracting_partitions, output_lhs_non_contracting_partitions,
          output_rhs_non_contracting_partitions, lhs_batch_partitions,
          rhs_batch_partitions, output_batch_partitions,
          require_matching_devices_to_group);
  if (!(matching_dims.empty() ||
        prioritize_contracting_for_faster_windowed_einsum)) {
    TF_ASSIGN_OR_RETURN(
        auto dot,
        PartitionDotGroupOnNonContracting(
            lhs_matching, lhs_matching ? lhs : rhs, lhs_matching ? rhs : lhs,
            lhs_matching ? lhs_contracting_partitions
                         : rhs_contracting_partitions,
            lhs_matching ? rhs_contracting_partitions
                         : lhs_contracting_partitions,
            matching_dims,
            lhs_matching ? rhs_non_contracting_partitions
                         : lhs_non_contracting_partitions,
            lhs_matching ? output_rhs_non_contracting_partitions
                         : output_lhs_non_contracting_partitions,
            output_base_shape, output_sharding, dims_mapping, num_partitions,
            create_sharded_dot, conv_window, module, original_hlo,
            require_matching_devices_to_group, options, b,
            windowed_dot_general_loops));
    if (dot) {
      return dot;
    }
  }

  // Case 3: Group partitions by contracting dimensions.
  if (lhs_contracting_partitions == rhs_contracting_partitions &&
      lhs_contracting_partitions > 1) {
    TF_ASSIGN_OR_RETURN(
        auto dot,
        PartitionDotGroupOnContracting(
            lhs, rhs, dims_mapping.contracting_dims, output_batch_partitions,
            output_lhs_non_contracting_partitions,
            output_rhs_non_contracting_partitions, output_base_shape,
            output_sharding, dims_mapping, num_partitions, create_sharded_dot,
            conv_window, module, original_hlo,
            require_matching_devices_to_group, options, b,
            windowed_dot_general_loops));
    if (dot) {
      return dot;
    }
  }
  if (lhs_contracting_partitions > 1 && rhs_contracting_partitions > 1) {
    // If part of contracting dims match, try them.
    std::vector<DotConvDimsMapping::DimsMapping> matching_dims;
    for (const auto& dim : dims_mapping.contracting_dims) {
      int64 lhs_partitions = lhs.sharding().tile_assignment().dim(dim.lhs);
      if (lhs_partitions > 1 &&
          lhs_partitions == rhs.sharding().tile_assignment().dim(dim.rhs)) {
        matching_dims.push_back(dim);
      }
    }
    if (!matching_dims.empty()) {
      TF_ASSIGN_OR_RETURN(
          auto dot, PartitionDotGroupOnContracting(
                        lhs, rhs, matching_dims, output_batch_partitions,
                        output_lhs_non_contracting_partitions,
                        output_rhs_non_contracting_partitions,
                        output_base_shape, output_sharding, dims_mapping,
                        num_partitions, create_sharded_dot, conv_window, module,
                        original_hlo, require_matching_devices_to_group,
                        options, b, windowed_dot_general_loops));
      if (dot) {
        return dot;
      }
    }
  }

  // Case 4: If operands are replicated but output is partially replicated,
  // recursive call with partial replication removed.
  if (lhs.sharding().IsReplicated() && rhs.sharding().IsReplicated() &&
      output_sharding.ReplicateOnLastTileDim()) {
    auto grouped_output =
        GroupShardingOnDims(output_sharding, {output_base_shape.rank()});
    auto inner_state = CreatePerGroupPartitioningState(
        lhs.state(), grouped_output.device_groups, b);
    TF_ASSIGN_OR_RETURN(
        auto dot,
        PartitionDot(PartitionedHlo(lhs.hlo(), lhs.base_shape(), inner_state),
                     PartitionedHlo(rhs.hlo(), rhs.base_shape(), inner_state),
                     output_base_shape, grouped_output.sharding, dims_mapping,
                     output_sharding.NumTiles(), create_sharded_dot,
                     conv_window, module, original_hlo, options, b,
                     windowed_dot_general_loops));
    if (dot) {
      return dot;
    }
  }

  // We failed to find partial matches, invoke base case again with
  // may_reshard_without_detecting_match.
  TF_ASSIGN_OR_RETURN(
      auto dot,
      PartitionBaseCase(
          lhs, rhs, output_base_shape, output_sharding, dims_mapping,
          num_partitions, create_sharded_dot, conv_window, module, original_hlo,
          lhs_batch_partitions, rhs_batch_partitions, output_batch_partitions,
          lhs_contracting_partitions, rhs_contracting_partitions,
          lhs_non_contracting_partitions, rhs_non_contracting_partitions,
          output_lhs_non_contracting_partitions,
          output_rhs_non_contracting_partitions, options, b,
          windowed_dot_general_loops,
          /*may_reshard_without_detecting_match=*/true));
  if (dot) {
    return dot;
  }
  return nullptr;
}

StatusOr<HloInstruction*> PartitionDot(
    PartitionedHlo lhs, PartitionedHlo rhs, const Shape& output_base_shape,
    const HloSharding& output_sharding, const DotConvDimsMapping& dims_mapping,
    int64 num_partitions,
    const std::function<StatusOr<HloInstruction*>(
        HloInstruction*, HloInstruction*, SpmdBuilder*,
        const Window& conv_window)>& create_sharded_dot,
    const Window& conv_window, HloModule* module, HloInstruction* original_hlo,
    const SpmdPartitionerOptions& options, SpmdBuilder* b,
    std::vector<SpmdPartitioningVisitor::WindowedDotGeneralLoop>*
        windowed_dot_general_loops) {
  // First try partitioning without resharding the groups, then try allow
  // resharding the groups.
  for (bool require_matching_devices_to_group : {true, false}) {
    TF_ASSIGN_OR_RETURN(
        auto try_partition,
        PartitionDot(lhs, rhs, output_base_shape, output_sharding, dims_mapping,
                     num_partitions, create_sharded_dot, conv_window, module,
                     original_hlo, require_matching_devices_to_group, options,
                     b, windowed_dot_general_loops));
    if (try_partition) {
      return try_partition;
    }
  }

  // Default action.
  TF_ASSIGN_OR_RETURN(
      auto dot, create_sharded_dot(lhs.Replicate().hlo(), rhs.Replicate().hlo(),
                                   b, conv_window));
  dot->set_sharding(HloSharding::Replicate());
  return PartitionedHlo(dot, output_base_shape, lhs.state())
      .Reshard(output_sharding)
      .hlo();
}

}  // namespace

Status SpmdPartitioningVisitor::HandleDotHelper(
    HloInstruction* hlo, const DotConvDimsMapping& dims_mapping,
    const std::function<StatusOr<HloInstruction*>(
        HloInstruction*, HloInstruction*, SpmdBuilder*,
        const Window& conv_window)>& create_sharded_dot) {
  auto& lhs = GetPartitionedHlo(hlo->operand(0));
  auto& rhs = GetPartitionedHlo(hlo->operand(1));
  Window conv_window;
  if (hlo->opcode() == HloOpcode::kConvolution) {
    conv_window = hlo->window();
  }

  TF_ASSIGN_OR_RETURN(
      auto partitioned_dot,
      PartitionDot(lhs, rhs, hlo->shape(), hlo->sharding(), dims_mapping,
                   num_partitions_, create_sharded_dot, conv_window, module_,
                   hlo, options_, &b_, &windowed_dot_general_loops_));
  SetPartitionedHlo(hlo, [&] { return partitioned_dot; });
  return Status::OK();
}

namespace {

// Finds a cluster of nodes that produce the inputs for `hlo` which only
// depend on small operands, which means the cluster should start with
// broadcasts, constants and iotas. All other internal nodes must be
// non-side-effecting elemntwise ops. Returns the set of nodes, and the small
// operands. E.g., for the following graph,
//
//     a -> broadcast -> multiply
//     iota  ---> add--/
//     constant/
//
// FindInputNodesIfOnlyDependOnSmallOperands(multiply) will return
//    <{broadcast, iota, constant, add, multiply}, [a]>.
std::pair<absl::flat_hash_set<HloInstruction*>, std::vector<HloInstruction*>>
FindInputNodesIfOnlyDependOnSmallOperands(HloInstruction* hlo) {
  absl::flat_hash_set<HloInstruction*> nodes_found;
  std::vector<HloInstruction*> new_operands;
  absl::flat_hash_set<const HloInstruction*> new_operands_set;
  std::vector<HloInstruction*> worklist;
  worklist.push_back(hlo);
  while (!worklist.empty()) {
    auto inst = worklist.back();
    worklist.pop_back();
    if (nodes_found.count(inst) > 0) {
      continue;
    }
    if (inst->opcode() == HloOpcode::kBroadcast ||
        inst->opcode() == HloOpcode::kConstant ||
        inst->opcode() == HloOpcode::kIota) {
      nodes_found.insert(inst);
      for (auto o : inst->operands()) {
        auto res = new_operands_set.emplace(o);
        if (res.second) {
          new_operands.push_back(o);
        }
      }
    } else if (inst->IsElementwise() && !inst->HasSideEffectNoRecurse() &&
               inst->opcode() != HloOpcode::kAllReduce &&
               absl::c_all_of(inst->operands(),
                              [inst](const HloInstruction* o) {
                                return ShapeUtil::CompatibleIgnoringElementType(
                                    o->shape(), inst->shape());
                              })) {
      nodes_found.insert(inst);
      for (auto o : inst->operands()) {
        worklist.push_back(o);
      }
    } else {
      nodes_found.clear();
      new_operands.clear();
      break;
    }
  }
  return {std::move(nodes_found), std::move(new_operands)};
}

// Moves a cluster of memory-reducing nodes into the windowed dot-general loop
// on contracting dimensions. Such a loop has a dynamic slice on the
// non-windowed operand. If we move the input nodes into the loop, the
// dynamic-slice could be merged with them by later optimization passes, which
// reduces memory.
//
// small_operands             small_operands
//        |                          |
// input_nodes                loop { |
//        |          =>         input_nodes
// loop { |                          |
//    dynamic-slice             dynamic-slice
//    ...                       ...
// }                          }
//
// Later optimization passes (TpuPadSliceMover) will merge the dynamic slice
// with the input nodes.
Status SinkInputNodesIntoWindowedDotGeneralLoopOnContractingDimensions(
    HloInstruction* loop, int64 non_windowed_operand_index) {
  auto input_tuple = loop->mutable_operand(0);
  auto old_operand = input_tuple->mutable_operand(non_windowed_operand_index);
  auto input_nodes = FindInputNodesIfOnlyDependOnSmallOperands(old_operand);
  auto to_sink = std::move(input_nodes.first);
  auto new_operands = std::move(input_nodes.second);
  if (to_sink.empty()) {
    return Status::OK();
  }
  auto computation = loop->parent();
  // Replace the old operand with a tuple of the found small operands.
  auto new_input_subtuple =
      computation->AddInstruction(HloInstruction::CreateTuple(new_operands));
  TF_RETURN_IF_ERROR(input_tuple->ReplaceOperandWithDifferentShape(
      non_windowed_operand_index, new_input_subtuple));

  auto body = loop->while_body();
  auto body_param = body->parameter_instruction(0);
  auto old_body_param_users = body_param->users();
  // Update all tuple shapes.
  for (auto tuple : std::vector<HloInstruction*>{
           input_tuple, loop, loop->while_condition()->parameter_instruction(0),
           body_param, body->root_instruction()}) {
    *ShapeUtil::GetMutableSubshape(tuple->mutable_shape(),
                                   {non_windowed_operand_index}) =
        new_input_subtuple->shape();
  }
  // Now update the loop body.
  auto new_operand_tuple_inside =
      body->AddInstruction(HloInstruction::CreateGetTupleElement(
          new_input_subtuple->shape(), body_param, non_windowed_operand_index));
  TF_RETURN_IF_ERROR(body->root_instruction()->ReplaceOperandWithDifferentShape(
      non_windowed_operand_index, new_operand_tuple_inside));

  // Create nodes inside the loop body.
  std::vector<HloInstruction*> worklist;
  absl::flat_hash_map<const HloInstruction*, HloInstruction*> outside_to_inside;
  auto add_users_if_available = [&](HloInstruction* inst) {
    for (auto u : inst->users()) {
      if (outside_to_inside.count(u) == 0 && to_sink.count(u) > 0 &&
          absl::c_all_of(u->operands(), [&](const HloInstruction* o) {
            return outside_to_inside.count(o) > 0;
          })) {
        worklist.push_back(u);
      }
    }
  };
  for (int64 i = 0; i < new_operands.size(); ++i) {
    outside_to_inside[new_operands[i]] =
        body->AddInstruction(HloInstruction::CreateGetTupleElement(
            new_operands[i]->shape(), new_operand_tuple_inside, i));
    add_users_if_available(new_operands[i]);
  }
  // HLOs to sink without operands.
  std::vector<HloInstruction*> nullaries_to_sink;
  for (auto inst : to_sink) {
    if (inst->operand_count() == 0) {
      nullaries_to_sink.push_back(inst);
    }
  }
  // Sort nullaries_to_sink to make it deterministic.
  absl::c_sort(nullaries_to_sink,
               [](const HloInstruction* a, const HloInstruction* b) {
                 return a->unique_id() < b->unique_id();
               });
  worklist.reserve(nullaries_to_sink.size());
  for (auto inst : nullaries_to_sink) {
    worklist.push_back(inst);
  }
  while (!worklist.empty()) {
    auto inst = worklist.back();
    worklist.pop_back();
    std::vector<HloInstruction*> inst_new_operands(inst->operand_count());
    for (int64 i = 0; i < inst->operand_count(); ++i) {
      inst_new_operands[i] = outside_to_inside[inst->operand(i)];
    }
    outside_to_inside[inst] = body->AddInstruction(
        inst->CloneWithNewOperands(inst->shape(), inst_new_operands));
    add_users_if_available(inst);
  }
  TF_RET_CHECK(outside_to_inside.count(old_operand) > 0);
  for (auto ou : old_body_param_users) {
    if (ou->opcode() == HloOpcode::kGetTupleElement &&
        ou->tuple_index() == non_windowed_operand_index) {
      TF_RETURN_IF_ERROR(
          ou->ReplaceAllUsesWith(outside_to_inside[old_operand]));
      TF_RETURN_IF_ERROR(body->RemoveInstruction(ou));
    }
  }
  return Status::OK();
}

// Moves a cluster of memory-reducing nodes (with reduce nodes at the end)
// into the windowed dot-general loop on non-contracting dimensions. Such a
// loop has a dynamic-update-slice at the output. If we move the user nodes
// into the loop and before the dynamic-update-slice, the user nodes can
// operate on smaller shapes, which reduces memory.
//
// small_operands                   small_operands
//  | |                 =>                  | |
//  | |  loop {                     loop {  | |
//  | |    conv                             | broadcast      conv
//  | |      |                              |     |           /
//  | | dynamic-update-slice                |  dynamic-slice /
//  | |         |                           |     |         /
//  | |  }      |                           |  multiply-----
//  |broadcast  /                           |    /
//  | |        /                            reduce
//  |multiply--                             |
//  \ |                                dynamic-update-slice
//   reduce                         }
//
// Later optimization passes (TpuPadSliceMover) will merge the dynamic slice
// with the input nodes (broadcast).
Status MoveUsersIntoWindowedDotGeneralLoopOnNonContractingDimensions(
    HloInstruction* loop) {
  CHECK_EQ(loop->user_count(), 1);
  // There should be a single direct user of the while loop, which is the
  // gte for element 2, i.e., the dot output.
  auto user_gte = loop->users().front();
  CHECK_EQ(user_gte->opcode(), HloOpcode::kGetTupleElement);
  CHECK_EQ(user_gte->tuple_index(), 2);
  auto computation = loop->parent();

  // Find the reduce outputs and the input nodes they depend on, if input
  // nodes only have small operands.
  absl::flat_hash_set<HloInstruction*> to_move;
  std::vector<HloInstruction*> new_operands;
  absl::flat_hash_set<const HloInstruction*> new_operands_set;
  std::vector<HloInstruction*> reduce_outputs;
  std::vector<HloInstruction*> worklist;
  Shape padded_shape = user_gte->shape();
  Shape unpadded_shape = user_gte->shape();
  auto original_output = user_gte;

  if (user_gte->user_count() == 1 &&
      user_gte->users().back()->opcode() == HloOpcode::kSlice) {
    original_output = user_gte->users().back();
    unpadded_shape = original_output->shape();
  }
  for (auto u : original_output->users()) {
    worklist.push_back(u);
  }
  to_move.insert(original_output);
  while (!worklist.empty()) {
    auto inst = worklist.back();
    worklist.pop_back();
    if (to_move.count(inst) > 0) {
      continue;
    }
    // We only support reduces with simple reduction function, since we may
    // need to accumulate across iterations manually.
    if (inst->opcode() == HloOpcode::kReduce &&
        inst->to_apply()->instruction_count() == 3 &&
        inst->to_apply()->num_parameters() == 2 &&
        inst->to_apply()->root_instruction()->IsElementwise()) {
      to_move.insert(inst);
      auto other_operand = inst->mutable_operand(1);
      auto res = new_operands_set.emplace(other_operand);
      if (res.second) {
        new_operands.push_back(other_operand);
      }
      reduce_outputs.push_back(inst);
    } else if (inst != computation->root_instruction() &&
               inst->user_count() > 0 && inst->IsElementwise() &&
               !inst->HasSideEffectNoRecurse() &&
               inst->opcode() != HloOpcode::kAllReduce &&
               absl::c_all_of(inst->operands(),
                              [inst](const HloInstruction* o) {
                                return ShapeUtil::CompatibleIgnoringElementType(
                                    o->shape(), inst->shape());
                              })) {
      // For an elementwise op, we need to make sure that they depend on only
      // nodes already in to_move and nodes with small operands.
      bool can_include = true;
      for (auto operand : inst->operands()) {
        if (to_move.count(operand) > 0) {
          continue;
        }
        auto find_result = FindInputNodesIfOnlyDependOnSmallOperands(operand);
        if (find_result.first.empty()) {
          can_include = false;
          break;
        }
        for (auto n : find_result.first) {
          to_move.insert(n);
        }
        for (auto new_operand : find_result.second) {
          auto res = new_operands_set.insert(new_operand);
          if (res.second) {
            new_operands.push_back(new_operand);
          }
        }
      }
      if (!can_include) {
        to_move.clear();
        break;
      }
      to_move.insert(inst);
      for (auto u : inst->users()) {
        worklist.push_back(u);
      }
    } else {
      to_move.clear();
      break;
    }
  }
  // If nothing is found, to_move could contain only original_output, or
  // cleared by the above code.
  if (to_move.size() <= 1) {
    return Status::OK();
  }

  // We will replace the original loop output with reduce-shape outputs.
  // Create the initial buffers before the loop.
  for (auto out : reduce_outputs) {
    auto padded_out_shape = out->shape();
    int64 operand_dim = 0;
    int64 output_dim = 0;
    while (output_dim < padded_out_shape.rank()) {
      if (absl::c_linear_search(out->dimensions(), operand_dim)) {
        // Dimension colapsed.
        ++operand_dim;
        continue;
      }
      // Kept dimensions have the same size of the padded shape.
      padded_out_shape.set_dimensions(output_dim,
                                      padded_shape.dimensions(operand_dim));
      ++operand_dim;
      ++output_dim;
    }
    auto broadcast =
        computation->AddInstruction(HloInstruction::CreateBroadcast(
            padded_out_shape,
            computation->AddInstruction(HloInstruction::CreateConstant(
                LiteralUtil::Zero(out->shape().element_type()))),
            {}));
    new_operands.push_back(broadcast);
  }

  auto input_tuple = loop->mutable_operand(0);
  // Create the new input subtuple that contains the small operands and the
  // reduce-shape result buffers.
  auto new_input_subtuple =
      computation->AddInstruction(HloInstruction::CreateTuple(new_operands));
  TF_RETURN_IF_ERROR(
      input_tuple->ReplaceOperandWithDifferentShape(2, new_input_subtuple));
  auto body = loop->while_body();
  auto body_param = body->parameter_instruction(0);
  auto body_root = body->root_instruction();
  CHECK_EQ(body_root->opcode(), HloOpcode::kTuple);
  // Update tuple shapes.
  for (auto tuple : std::vector<HloInstruction*>{
           input_tuple, loop, loop->while_condition()->parameter_instruction(0),
           body_param, body_root}) {
    *ShapeUtil::GetMutableSubshape(tuple->mutable_shape(), {2}) =
        new_input_subtuple->shape();
  }
  auto new_loop_input =
      body->AddInstruction(HloInstruction::CreateGetTupleElement(
          new_input_subtuple->shape(), body_param, 2));

  // Now create the moved nodes inside the loop body.
  absl::flat_hash_map<const HloInstruction*, HloInstruction*> outside_to_inside;
  worklist.clear();
  auto add_users_if_available = [&](HloInstruction* inst) {
    for (auto u : inst->users()) {
      if (outside_to_inside.count(u) == 0 && to_move.count(u) > 0 &&
          absl::c_all_of(u->operands(), [&](const HloInstruction* o) {
            return outside_to_inside.count(o) > 0;
          })) {
        worklist.push_back(u);
      }
    }
  };
  for (int64 i = 0; i < new_operands.size(); ++i) {
    outside_to_inside[new_operands[i]] =
        body->AddInstruction(HloInstruction::CreateGetTupleElement(
            new_operands[i]->shape(), new_loop_input, i));
    add_users_if_available(new_operands[i]);
  }
  // The elementwise nodes will be created with sliced shape. The original
  // loop output corresponds to the dynamic-update-slice's update slice.
  auto dus = body_root->mutable_operand(2);
  CHECK_EQ(dus->opcode(), HloOpcode::kDynamicUpdateSlice);
  outside_to_inside[original_output] = dus->mutable_operand(1);
  add_users_if_available(original_output);
  std::vector<HloInstruction*> slice_offsets(padded_shape.rank());
  for (int64 i = 0; i < slice_offsets.size(); ++i) {
    slice_offsets[i] = dus->mutable_operand(i + 2);
  }
  auto get_slice = [&](HloInstruction* padded) {
    return body->AddInstruction(HloInstruction::CreateDynamicSlice(
        ShapeUtil::ChangeElementType(dus->operand(1)->shape(),
                                     padded->shape().element_type()),
        padded, slice_offsets, dus->operand(1)->shape().dimensions()));
  };
  // Helper functions to create nodes with small operands.
  auto add_broadcast = [&](const HloInstruction* broadcast) {
    auto padded_operand_shape = broadcast->operand(0)->shape();
    for (int64 i = 0; i < broadcast->dimensions().size(); ++i) {
      padded_operand_shape.set_dimensions(
          i, padded_shape.dimensions(broadcast->dimensions(i)));
    }
    auto padded_operand = PadToShape(outside_to_inside[broadcast->operand(0)],
                                     padded_operand_shape, nullptr, body);
    outside_to_inside[broadcast] =
        get_slice(body->AddInstruction(broadcast->CloneWithNewOperands(
            ShapeUtil::ChangeElementType(padded_shape,
                                         padded_operand_shape.element_type()),
            {padded_operand})));
  };
  auto add_iota = [&](const HloInstruction* iota) {
    outside_to_inside[iota] =
        get_slice(body->AddInstruction(iota->CloneWithNewOperands(
            ShapeUtil::ChangeElementType(padded_shape,
                                         iota->shape().element_type()),
            {})));
  };
  auto add_constant = [&](const HloInstruction* constant) {
    outside_to_inside[constant] = body->AddInstruction(constant->Clone());
    outside_to_inside[constant] = get_slice(
        PadToShape(outside_to_inside[constant],
                   ShapeUtil::ChangeElementType(
                       padded_shape, constant->shape().element_type()),
                   nullptr, body));
  };
  while (!worklist.empty()) {
    auto inst = worklist.back();
    worklist.pop_back();
    if (outside_to_inside.count(inst) > 0) {
      continue;
    }
    if (inst->opcode() == HloOpcode::kBroadcast) {
      add_broadcast(inst);
    } else if (inst->opcode() == HloOpcode::kIota) {
      add_iota(inst);
    } else if (inst->opcode() == HloOpcode::kConstant) {
      add_constant(inst);
    } else if (inst->opcode() == HloOpcode::kReduce) {
      // This is an output, for which we has special handling later.
    } else {
      std::vector<HloInstruction*> operands_inside(inst->operand_count());
      for (int64 i = 0; i < operands_inside.size(); ++i) {
        operands_inside[i] = outside_to_inside[inst->operand(i)];
      }
      outside_to_inside[inst] = body->AddInstruction(inst->CloneWithNewOperands(
          ShapeUtil::ChangeElementType(dus->operand(1)->shape(),
                                       inst->shape().element_type()),
          operands_inside));
    }
    add_users_if_available(inst);
  }
  std::vector<HloInstruction*> new_outputs_inside(new_operands.size());
  for (int64 i = 0; i < new_outputs_inside.size(); ++i) {
    new_outputs_inside[i] = outside_to_inside[new_operands[i]];
  }
  // Now create the reduce outpus inside of the loop.
  for (int64 i = 0; i < reduce_outputs.size(); ++i) {
    auto reduce_outside = reduce_outputs[i];
    CHECK_EQ(reduce_outside->opcode(), HloOpcode::kReduce);
    int64 index_in_operand = new_operands.size() - reduce_outputs.size() + i;
    auto last_iter_result = outside_to_inside[new_operands[index_in_operand]];
    auto operand0 = outside_to_inside[reduce_outside->operand(0)];
    auto operand1 = outside_to_inside[reduce_outside->operand(1)];
    TF_ASSIGN_OR_RETURN(auto reduce_shape,
                        ShapeInference::InferReduceShape(
                            {&operand0->shape(), &operand1->shape()},
                            reduce_outside->dimensions(),
                            reduce_outside->to_apply()->ComputeProgramShape()));
    *reduce_shape.mutable_layout() = reduce_outside->shape().layout();
    std::vector<HloInstruction*> reduce_dus_offsets;
    // If any collapsed dimension is windowed, we need to accumulate with last
    // iteration's result. If such a dimension has padding, we also need to
    // mask off invalid data.
    bool needs_accumulate = false;
    std::vector<int64> dims_to_mask;
    for (int64 i = 0; i < slice_offsets.size(); ++i) {
      if (absl::c_linear_search(reduce_outside->dimensions(), i)) {
        if (reduce_outside->operand(0)->shape().dimensions(i) !=
            operand0->shape().dimensions(i)) {
          needs_accumulate = true;
          if (unpadded_shape.dimensions(i) != padded_shape.dimensions(i)) {
            dims_to_mask.push_back(i);
          }
        }
        continue;
      }
      reduce_dus_offsets.push_back(slice_offsets[i]);
    }
    // Mask off invalid data in collapsed dimensions.
    for (int64 dim : dims_to_mask) {
      auto iota = body->AddInstruction(HloInstruction::CreateIota(
          ShapeUtil::ChangeElementType(operand0->shape(), S32), dim));
      auto add = body->AddInstruction(HloInstruction::CreateBinary(
          iota->shape(), HloOpcode::kAdd, iota,
          body->AddInstruction(HloInstruction::CreateBroadcast(
              iota->shape(), slice_offsets[dim], {}))));
      auto limit = body->AddInstruction(HloInstruction::CreateBroadcast(
          iota->shape(),
          body->AddInstruction(
              HloInstruction::CreateConstant(LiteralUtil::CreateR0<int32>(
                  reduce_outside->operand(0)->shape().dimensions(dim)))),
          {}));
      auto compare = body->AddInstruction(HloInstruction::CreateCompare(
          ShapeUtil::ChangeElementType(iota->shape(), PRED), add, limit,
          ComparisonDirection::kLt));
      operand0 = body->AddInstruction(HloInstruction::CreateTernary(
          operand0->shape(), HloOpcode::kSelect, compare, operand0,
          body->AddInstruction(HloInstruction::CreateBroadcast(
              operand0->shape(), operand1, {}))));
    }
    auto output_inside =
        body->AddInstruction(reduce_outside->CloneWithNewOperands(
            reduce_shape, {operand0, operand1}));
    // Accumulate with previous results if needed.
    if (needs_accumulate) {
      auto input_slice =
          body->AddInstruction(HloInstruction::CreateDynamicSlice(
              output_inside->shape(), last_iter_result, reduce_dus_offsets,
              output_inside->shape().dimensions()));
      output_inside = body->AddInstruction(HloInstruction::CreateBinary(
          output_inside->shape(),
          reduce_outside->to_apply()->root_instruction()->opcode(),
          output_inside, input_slice));
    }
    // Dynamic-update-slice if needed.
    if (!ShapeUtil::Compatible(output_inside->shape(),
                               last_iter_result->shape())) {
      output_inside =
          body->AddInstruction(HloInstruction::CreateDynamicUpdateSlice(
              last_iter_result->shape(), last_iter_result, output_inside,
              reduce_dus_offsets));
    }
    new_outputs_inside[index_in_operand] = output_inside;
  }
  // Body output.
  auto new_output_inside =
      body->AddInstruction(HloInstruction::CreateTuple(new_outputs_inside));
  TF_RETURN_IF_ERROR(
      body_root->ReplaceOperandWithDifferentShape(2, new_output_inside));
  TF_RETURN_IF_ERROR(body->RemoveInstructionAndUnusedOperands(dus));
  // Replace uses of the reduces outside the loop.
  auto new_output_gte =
      computation->AddInstruction(HloInstruction::CreateGetTupleElement(
          new_output_inside->shape(), loop, 2));
  for (int64 i = 0; i < reduce_outputs.size(); ++i) {
    int64 index_in_operand = new_operands.size() - reduce_outputs.size() + i;
    auto new_output =
        computation->AddInstruction(HloInstruction::CreateGetTupleElement(
            new_outputs_inside[index_in_operand]->shape(), new_output_gte,
            index_in_operand));
    if (!ShapeUtil::Compatible(new_output->shape(),
                               reduce_outputs[i]->shape())) {
      new_output = computation->AddInstruction(HloInstruction::CreateSlice(
          reduce_outputs[i]->shape(), new_output,
          std::vector<int64>(new_output->shape().rank(), 0),
          reduce_outputs[i]->shape().dimensions(),
          std::vector<int64>(new_output->shape().rank(), 1)));
    }
    TF_RETURN_IF_ERROR(reduce_outputs[i]->ReplaceAllUsesWith(new_output));
    TF_RETURN_IF_ERROR(
        computation->RemoveInstructionAndUnusedOperands(reduce_outputs[i]));
  }
  return Status::OK();
}

}  // namespace

Status SpmdPartitioningVisitor::DoCodeMotionForWindowedDotGeneralLoops(
    HloComputation* computation, const SpmdPartitionerOptions& options) {
  for (auto& loop : windowed_dot_general_loops_) {
    if (loop.windowed_in_contracting_dims || loop.windowed_in_batch_dims ||
        loop.operands_sharded_at_contracting_dims) {
      // We have a dynamic-slice for the non-windowed operand in
      // batch/contracting-dim/noncontracting-dim windowed dot-general. So
      // moving the broadcast/iota/elementwise ops into the loop could help
      // reduce memory via fusion.
      TF_RETURN_IF_ERROR(
          SinkInputNodesIntoWindowedDotGeneralLoopOnContractingDimensions(
              loop.while_loop, 1 - loop.windowed_operand));
    }
    // Currently unrolled loop does not support this optimization.
    if (!options.bidirectional_windowed_einsum &&
        !options.unroll_windowed_einsum && !loop.windowed_in_contracting_dims &&
        !loop.operands_sharded_at_contracting_dims) {
      // We have a dynamic-update-slice for the output in
      // batch/non-contracting-dim windowed dot-general. So moving reduce ops
      // into the loop could help reduce memory.
      TF_RETURN_IF_ERROR(
          MoveUsersIntoWindowedDotGeneralLoopOnNonContractingDimensions(
              loop.while_loop));
    }
  }
  return Status::OK();
}

}  // namespace spmd
}  // namespace xla

redpajama_set_name: RedPajamaGithub