/* Copyright 2020 The TensorFlow Authors. All Rights Reserved.

Licensed under the Apache License, Version 2.0 (the "License");
you may not use this file except in compliance with the License.
You may obtain a copy of the License at

    http://www.apache.org/licenses/LICENSE-2.0

Unless required by applicable law or agreed to in writing, software
distributed under the License is distributed on an "AS IS" BASIS,
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
See the License for the specific language governing permissions and
limitations under the License.
==============================================================================*/

#include "tensorflow/compiler/xla/service/spmd/spmd_partitioner.h"

#include "tensorflow/compiler/xla/service/hlo_casting_utils.h"
#include "tensorflow/compiler/xla/service/hlo_instructions.h"
#include "tensorflow/compiler/xla/service/hlo_matchers.h"
#include "tensorflow/compiler/xla/service/hlo_parser.h"
#include "tensorflow/compiler/xla/service/hlo_pass_pipeline.h"
#include "tensorflow/compiler/xla/service/hlo_verifier.h"
#include "tensorflow/compiler/xla/tests/hlo_test_base.h"
#include "tensorflow/compiler/xla/util.h"
#include "tensorflow/compiler/xla/xla_data.pb.h"
#include "tensorflow/core/lib/core/status_test_util.h"

namespace xla {
namespace spmd {
namespace {

using ::testing::_;
using ::testing::AllOf;
namespace op = xla::testing::opcode_matchers;

class SpmdPartitioningTest : public HloTestBase {
 public:
  StatusOr<std::unique_ptr<HloModule>> PartitionComputation(
      absl::string_view hlo_module, int64 num_devices,
      bool conv_halo_exchange_always_on_lhs = true,
      bool choose_faster_windowed_einsum = false) {
    // Some tests (BackpropFilter convs) set this flag false to test two
    // different paths of the implementation.
    SpmdPartitionerOptions options;
    options.conv_halo_exchange_always_on_lhs = conv_halo_exchange_always_on_lhs;
    options.allow_module_signature_change = true;
    options.choose_faster_windowed_einsum_over_mem =
        choose_faster_windowed_einsum;
    auto collective_ops_creator =
        GetDefaultCollectiveOpsCreator(num_devices, /*num_replicas=*/1);
    // Do not use all-gather for pattern-matching purpose, as the partitioner
    // might create reshape/transposes around it.
    collective_ops_creator.create_cross_partition_all_gather = nullptr;

    TF_ASSIGN_OR_RETURN(auto module, ParseAndReturnVerifiedModule(
                                         hlo_module, GetModuleConfigForTest()));
    HloPassPipeline pass("spmd-partitioning");
    pass.AddPass<HloVerifier>(/*layout_sensitive=*/false,
                              /*allow_mixed_precision=*/false);
    pass.AddPass<SpmdPartitioner>(num_devices, /*num_replicas=*/1, options,
                                  collective_ops_creator);
    pass.AddPass<HloVerifier>(/*layout_sensitive=*/false,
                              /*allow_mixed_precision=*/false);
    TF_RETURN_IF_ERROR(pass.Run(module.get()).status());
    return StatusOr<std::unique_ptr<HloModule>>(std::move(module));
  }
};

TEST_F(SpmdPartitioningTest, InvalidSharding) {
  absl::string_view hlo_string = R"(
HloModule module

ENTRY entry {
  token0 = token[] after-all(), sharding={maximal device=0}
  infeed = (f32[8,2]{1,0}, token[]) infeed(token0),
    sharding={{devices=[2,1]0,1}, {maximal device=0}}
  ROOT infeed.data = f32[8,2]{1,0} get-tuple-element(infeed), index=0,
    sharding={maximal device=0}
})";
  auto module_status = PartitionComputation(hlo_string, /*num_devices=*/4);
  EXPECT_FALSE(module_status.status().ok());
  EXPECT_THAT(module_status.status().ToString(),
              ::testing::HasSubstr(
                  "only supports tile sharding that includes all partitions"));
}

TEST_F(SpmdPartitioningTest, SingleDeviceToReplicated) {
  absl::string_view hlo_string = R"(
HloModule module

ENTRY entry {
  %constant = s32[2,3]{1,0} constant({{1,1,1},{1,1,1}}),
    sharding={maximal device=0}
  ROOT %copy = s32[2,3]{1,0} copy(%constant), sharding={replicated}
})";
  TF_ASSERT_OK_AND_ASSIGN(auto module,
                          PartitionComputation(hlo_string, /*num_devices=*/2));
  VLOG(1) << module->ToString();
  HloInstruction* root = module->entry_computation()->root_instruction();
  EXPECT_THAT(root, AllOf(op::Copy(op::AllReduce(
                              op::Select(op::Broadcast(op::Compare()),
                                         op::Constant(), op::Broadcast()))),
                          op::Shape("s32[2,3]")));
}

TEST_F(SpmdPartitioningTest, SingleDeviceToSingleDevice) {
  absl::string_view hlo_string = R"(
HloModule module

ENTRY entry {
  %constant = s32[2,3]{1,0} constant({{1,1,1},{1,1,1}}),
    sharding={maximal device=0}
  ROOT %copy = s32[2,3]{1,0} copy(%constant), sharding={maximal device=1}
})";
  TF_ASSERT_OK_AND_ASSIGN(auto module,
                          PartitionComputation(hlo_string, /*num_devices=*/2));
  HloInstruction* root = module->entry_computation()->root_instruction();
  VLOG(1) << module->ToString();
  EXPECT_THAT(root, op::Copy(AllOf(op::Copy(op::AllReduce(op::Select(
                                       op::Broadcast(op::Compare()),
                                       op::Constant(), op::Broadcast()))),
                                   op::Shape("s32[2,3]"))));
}

TEST_F(SpmdPartitioningTest, SingleDeviceToTiled) {
  absl::string_view hlo_string = R"(
HloModule module

ENTRY entry {
  %constant = s32[2,3]{1,0} constant({{1,1,1},{1,1,1}}),
    sharding={maximal device=0}
  ROOT %copy = s32[2,3]{1,0} copy(%constant),
    sharding={devices=[2,1]1,0}
})";
  TF_ASSERT_OK_AND_ASSIGN(auto module,
                          PartitionComputation(hlo_string, /*num_devices=*/2));
  VLOG(1) << module->ToString();
  HloInstruction* root = module->entry_computation()->root_instruction();
  EXPECT_THAT(
      root,
      AllOf(
          op::Copy(op::DynamicSlice(
              op::AllReduce(op::Select(
                  op::Broadcast(op::Compare(op::PartitionId(), op::Constant())),
                  op::Constant(), op::Broadcast())),
              op::Reshape(op::DynamicSlice(op::Constant(), op::PartitionId())),
              op::Constant())),
          op::Shape("s32[1,3]")));
}

TEST_F(SpmdPartitioningTest, TiledToReplicated) {
  absl::string_view hlo_string = R"(
HloModule module

ENTRY entry {
  %constant = s32[2,3]{1,0} constant({{1,1,1},{1,1,1}}),
    sharding={devices=[2,1]0,1}
  ROOT %copy = s32[2,3]{1,0} copy(%constant), sharding={replicated}
})";
  TF_ASSERT_OK_AND_ASSIGN(auto module,
                          PartitionComputation(hlo_string, /*num_devices=*/2));
  HloInstruction* root = module->entry_computation()->root_instruction();
  EXPECT_THAT(
      root,
      op::Copy(op::AllReduce(AllOf(
          op::DynamicUpdateSlice(
              op::Broadcast(), AllOf(op::Constant(), op::Shape("s32[1,3]")),
              op::Reshape(op::DynamicSlice(op::Constant(), op::PartitionId())),
              op::Constant()),
          op::Shape("s32[2,3]")))));
}

TEST_F(SpmdPartitioningTest, TiledToSingleDevice) {
  absl::string_view hlo_string = R"(
HloModule module

ENTRY entry {
  %constant = s32[2,3]{1,0} constant({{1,1,1},{1,1,1}}),
    sharding={devices=[2,1]0,1}
  ROOT %copy = s32[2,3]{1,0} copy(%constant), sharding={maximal device=0}
})";
  TF_ASSERT_OK_AND_ASSIGN(auto module,
                          PartitionComputation(hlo_string, /*num_devices=*/2));
  HloInstruction* root = module->entry_computation()->root_instruction();
  EXPECT_THAT(
      root,
      op::Copy(op::Copy(op::AllReduce(AllOf(
          op::DynamicUpdateSlice(
              op::Broadcast(), AllOf(op::Constant(), op::Shape("s32[1,3]")),
              op::Reshape(op::DynamicSlice(op::Constant(), op::PartitionId())),
              op::Constant()),
          op::Shape("s32[2,3]"))))));
}

TEST_F(SpmdPartitioningTest, TiledToTiledEven) {
  absl::string_view hlo_string = R"(
HloModule module

ENTRY entry {
  %param= s32[8,2]{1,0} parameter(0), sharding={devices=[2,1]0,1}
  ROOT %copy = s32[8,2]{1,0} copy(%param), sharding={devices=[1,2]0,1}
})";
  TF_ASSERT_OK_AND_ASSIGN(auto module,
                          PartitionComputation(hlo_string, /*num_devices=*/2));
  VLOG(1) << module->ToString();

  HloInstruction* root = module->entry_computation()->root_instruction();
  EXPECT_THAT(
      root,
      AllOf(op::Copy(op::Reshape(op::Transpose(op::AllToAll(AllOf(
                op::Reshape(op::Parameter()), op::Shape("s32[4,2,1]")))))),
            op::Shape("s32[8,1]")));
}

TEST_F(SpmdPartitioningTest, TiledToTiledUneven) {
  absl::string_view hlo_string = R"(
HloModule module

ENTRY entry {
  %param= f32[7,31,128]{2,1,0} parameter(0), sharding={devices=[1,2,1]0,1}
  ROOT %copy = f32[7,31,128]{2,1,0} copy(%param), sharding={devices=[2,1,1]0,1}
})";
  TF_ASSERT_OK_AND_ASSIGN(auto module,
                          PartitionComputation(hlo_string, /*num_devices=*/2));
  VLOG(1) << module->ToString();

  HloInstruction* root = module->entry_computation()->root_instruction();
  EXPECT_THAT(
      root,
      AllOf(op::Copy(op::Slice(op::Reshape(AllOf(op::Transpose(op::AllToAll(
          op::Reshape(AllOf(op::Pad(), op::Shape("f32[8,16,128]")))))))))));
}

TEST_F(SpmdPartitioningTest, GetTupleElementSwapDevice) {
  absl::string_view hlo_string = R"(
HloModule module

ENTRY entry {
  %param.0 = (f32[2,3]{1,0}, u32[]) parameter(0),
    sharding={{maximal device=1}, {maximal device=1}}
  %gte.0 = f32[2,3]{1,0} get-tuple-element(%param.0), index=0,
    sharding={maximal device=0}
  %gte.1 = u32[] get-tuple-element(%param.0), index=1,
    sharding={maximal device=0}
  ROOT %tuple = (f32[2,3]{1,0}, u32[]) tuple(%gte.0, %gte.1),
    sharding={{maximal device=0},{maximal device=0}}
})";
  TF_ASSERT_OK_AND_ASSIGN(auto module,
                          PartitionComputation(hlo_string, /*num_devices=*/2));
  VLOG(1) << module->ToString();
  HloInstruction* root = module->entry_computation()->root_instruction();
  ASSERT_THAT(root, op::Tuple());

  EXPECT_THAT(root->operand(0),
              op::Copy(op::AllReduce(op::Select(
                  op::Broadcast(op::Compare(op::PartitionId(), op::Constant())),
                  op::GetTupleElement(op::Parameter()), op::Broadcast()))));
  EXPECT_THAT(root->operand(1),
              op::Copy(op::AllReduce(op::Select(
                  op::Broadcast(op::Compare(op::PartitionId(), op::Constant())),
                  op::GetTupleElement(op::Parameter()), op::Broadcast()))));
}

TEST_F(SpmdPartitioningTest, GetTupleElementTiled) {
  absl::string_view hlo_string = R"(
HloModule module

ENTRY entry {
  param.0 = (f32[2,3]{1,0}, u32[2,3]{1,0}) parameter(0),
    sharding={{replicated}, {replicated}}
  gte.0 = f32[2,3]{1,0} get-tuple-element(param.0), index=0,
    sharding={devices=[2,1]0,1}
  gte.1 = u32[2,3]{1,0} get-tuple-element(param.0), index=1,
    sharding={devices=[2,1]0,1}
  ROOT %tuple = (f32[2,3]{1,0}, u32[2,3]{1,0}) tuple(gte.0, gte.1),
    sharding={{devices=[2,1]0,1},{devices=[2,1]0,1}}
})";
  TF_ASSERT_OK_AND_ASSIGN(auto module,
                          PartitionComputation(hlo_string, /*num_devices=*/2));
  VLOG(1) << module->ToString();
  HloInstruction* root = module->entry_computation()->root_instruction();
  ASSERT_THAT(root, op::Tuple());

  auto offset =
      op::Reshape(op::DynamicSlice(op::Constant(), op::PartitionId()));

  EXPECT_THAT(root->operand(0),
              op::DynamicSlice(op::GetTupleElement(op::Parameter()), offset,
                               op::Constant()));
  EXPECT_THAT(root->operand(1),
              op::DynamicSlice(op::GetTupleElement(op::Parameter()), offset,
                               op::Constant()));
}

TEST_F(SpmdPartitioningTest, TiledInfeed) {
  absl::string_view hlo_string = R"(
HloModule module

ENTRY entry {
  token0 = token[] after-all(), sharding={maximal device=0}
  infeed = (f32[8,2]{1,0}, token[]) infeed(token0),
    sharding={{devices=[2,1]0,1}, {maximal device=0}}
  ROOT infeed.data = f32[8,2]{1,0} get-tuple-element(infeed), index=0,
    sharding={maximal device=0}
})";
  TF_ASSERT_OK_AND_ASSIGN(auto module,
                          PartitionComputation(hlo_string, /*num_devices=*/2));
  HloInstruction* root = module->entry_computation()->root_instruction();
  EXPECT_THAT(
      root,
      op::Copy(op::AllReduce(op::DynamicUpdateSlice(
          op::Broadcast(),
          op::GetTupleElement(
              AllOf(op::Infeed(), op::Shape("(f32[4,2]{1,0}, token[])"))),
          op::Reshape(op::DynamicSlice(op::Constant(), op::PartitionId())),
          op::Constant()))));
}

TEST_F(SpmdPartitioningTest, UnevenTiledInfeed) {
  absl::string_view hlo_string = R"(
HloModule module

ENTRY entry {
  token0 = token[] after-all(), sharding={maximal device=0}
  infeed = (f32[9,2]{1,0}, token[]) infeed(token0),
    sharding={{devices=[2,1]0,1}, {maximal device=0}}
  ROOT infeed.data = f32[9,2]{1,0} get-tuple-element(infeed), index=0,
    sharding={devices=[2,1]0,1}
})";
  TF_ASSERT_OK_AND_ASSIGN(auto module,
                          PartitionComputation(hlo_string, /*num_devices=*/2));
  VLOG(1) << module->ToString();
  HloInstruction* root = module->entry_computation()->root_instruction();
  EXPECT_THAT(
      root, AllOf(op::Shape("f32[5,2]"), op::GetTupleElement(op::Conditional(
                                             op::Convert(op::PartitionId()),
                                             op::AfterAll(), op::AfterAll()))));
  EXPECT_THAT(
      root->operand(0)->called_computations()[0]->root_instruction(),
      AllOf(op::Shape("(f32[5,2], token[])"), op::Infeed(op::Parameter())));
  auto second_infeed =
      AllOf(op::Shape("(f32[4,2], token[])"), op::Infeed(op::Parameter()));
  EXPECT_THAT(root->operand(0)->called_computations()[1]->root_instruction(),
              AllOf(op::Shape("(f32[5,2], token[])"),
                    op::Tuple(op::Pad(op::GetTupleElement(second_infeed),
                                      op::Constant()),
                              op::GetTupleElement(second_infeed))));
}

TEST_F(SpmdPartitioningTest, UnevenTiledTupleInfeed) {
  absl::string_view hlo_string = R"(
HloModule module

ENTRY entry {
  token0 = token[] after-all(), sharding={maximal device=0}
  infeed = ((f32[9,2]{1,0}, f32[2]{0}), token[]) infeed(token0),
    sharding={{devices=[2,1]0,1}, {replicated}, {maximal device=0}}
  ROOT infeed.data = (f32[9,2]{1,0}, f32[2]{0}) get-tuple-element(infeed),
    index=0, sharding={{devices=[2,1]0,1}, {replicated}}
})";
  TF_ASSERT_OK_AND_ASSIGN(auto module,
                          PartitionComputation(hlo_string, /*num_devices=*/2));
  VLOG(1) << module->ToString();
  HloInstruction* root = module->entry_computation()->root_instruction();
  EXPECT_THAT(root, AllOf(op::Shape("(f32[5,2], f32[2])"),
                          op::GetTupleElement(op::Conditional(
                              op::Convert(op::PartitionId()), op::AfterAll(),
                              op::AfterAll()))));
  EXPECT_THAT(root->operand(0)->called_computations()[0]->root_instruction(),
              AllOf(op::Shape("((f32[5,2], f32[2]), token[])"),
                    op::Infeed(op::Parameter())));
  auto second_infeed = AllOf(op::Shape("((f32[4,2], f32[2]), token[])"),
                             op::Infeed(op::Parameter()));
  EXPECT_THAT(
      root->operand(0)->called_computations()[1]->root_instruction(),
      AllOf(op::Shape("((f32[5,2], f32[2]), token[])"),
            op::Tuple(op::Tuple(op::Pad(op::GetTupleElement(
                                            op::GetTupleElement(second_infeed)),
                                        op::Constant()),
                                op::GetTupleElement(
                                    op::GetTupleElement(second_infeed))),
                      op::GetTupleElement(second_infeed))));
}

TEST_F(SpmdPartitioningTest, MixedTupleInfeed) {
  absl::string_view hlo_string = R"(
HloModule module

ENTRY entry {
  token0 = token[] after-all(), sharding={maximal device=0}
  infeed = ((f32[9,2]{1,0}, f32[2]{0}), token[]) infeed(token0),
    sharding={{maximal device=0}, {maximal device=1}, {maximal device=0}}
  ROOT infeed.data = (f32[9,2]{1,0}, f32[2]{0}) get-tuple-element(infeed),
    index=0, sharding={{maximal device=0}, {maximal device=1}}
})";
  TF_ASSERT_OK_AND_ASSIGN(auto module,
                          PartitionComputation(hlo_string, /*num_devices=*/2));
  VLOG(1) << module->ToString();
  HloInstruction* root = module->entry_computation()->root_instruction();
  EXPECT_THAT(root, AllOf(op::Shape("(f32[9,2], f32[2])"),
                          op::GetTupleElement(op::Conditional(
                              op::Convert(op::PartitionId()), op::AfterAll(),
                              op::AfterAll()))));
  auto first_infeed = AllOf(op::Shape("((f32[9,2], ()), token[])"),
                            op::Infeed(op::Parameter()));
  EXPECT_THAT(root->operand(0)->called_computations()[0]->root_instruction(),
              AllOf(op::Shape("((f32[9,2], f32[2]), token[])"),
                    op::Tuple(op::Tuple(op::GetTupleElement(
                                            op::GetTupleElement(first_infeed)),
                                        op::Broadcast(op::Constant())),
                              op::GetTupleElement(first_infeed))));
  auto second_infeed =
      AllOf(op::Shape("(((), f32[2]), token[])"), op::Infeed(op::Parameter()));
  EXPECT_THAT(root->operand(0)->called_computations()[1]->root_instruction(),
              AllOf(op::Shape("((f32[9,2], f32[2]), token[])"),
                    op::Tuple(op::Tuple(op::Broadcast(op::Constant()),
                                        op::GetTupleElement(op::GetTupleElement(
                                            second_infeed))),
                              op::GetTupleElement(second_infeed))));
}

TEST_F(SpmdPartitioningTest, TiledToReplicatedReduce) {
  absl::string_view hlo_string = R"(
HloModule module

sum {
  a = f32[] parameter(0)
  b = f32[] parameter(1)
  ROOT add = f32[] add(a, b)
}

ENTRY entry {
  constant = f32[3,3]{1,0} constant({{1,1,1},{1,1,1},{1,1,1}}),
    sharding={devices=[2,1]0,1}
  constant.1 = f32[] constant(0), sharding={replicated}
  ROOT reduce = f32[] reduce(constant, constant.1), dimensions={0,1},
    to_apply=sum, sharding={replicated}
})";
  TF_ASSERT_OK_AND_ASSIGN(auto module,
                          PartitionComputation(hlo_string, /*num_devices=*/2));
  VLOG(1) << module->ToString();
  HloInstruction* root = module->entry_computation()->root_instruction();
  EXPECT_THAT(
      root,
      op::AllReduce(op::Reduce(
          op::Select(
              op::Compare(op::Add(op::Iota(), op::Broadcast(op::Reshape())),
                          op::Broadcast(op::Constant())),
              AllOf(op::Shape("f32[2,3]{1,0}"),
                    op::DynamicSlice(op::Pad(op::Constant(), op::Constant()),
                                     op::Reshape(), op::Constant())),
              op::Broadcast(op::Constant())),
          op::Constant())));
}

TEST_F(SpmdPartitioningTest, TiledElementwise) {
  absl::string_view hlo_string = R"(
HloModule module

ENTRY entry {
  constant = f32[3,3]{1,0} constant({{1,1,1},{1,1,1},{1,1,1}}),
    sharding={devices=[2,1]0,1}
  constant.1 = f32[3,3]{1,0} constant({{2,2,2},{2,2,2},{2,2,2}}),
    sharding={replicated}
  multiply = f32[3,3]{1,0} multiply(constant, constant.1),
    sharding={devices=[2,1]0,1}
  ROOT add = f32[3,3]{1,0} add(multiply, constant.1),
    sharding={devices=[2,1]0,1}
})";
  TF_ASSERT_OK_AND_ASSIGN(auto module,
                          PartitionComputation(hlo_string, /*num_devices=*/2));
  VLOG(1) << module->ToString();
  HloInstruction* root = module->entry_computation()->root_instruction();
  EXPECT_THAT(
      root,
      AllOf(
          op::Shape("f32[2,3]{1,0}"),
          op::Add(op::Multiply(
                      op::DynamicSlice(op::Pad(op::Constant(), op::Constant()),
                                       op::Reshape(), op::Constant()),
                      op::DynamicSlice(op::Pad(op::Constant(), op::Constant()),
                                       op::Reshape(), op::Constant())),
                  op::DynamicSlice(op::Pad(op::Constant(), op::Constant()),
                                   op::Reshape(), op::Constant()))));
}

TEST_F(SpmdPartitioningTest, TiledAllReduce) {
  absl::string_view hlo_string = R"(
HloModule module

sum {
  a = f32[] parameter(0)
  b = f32[] parameter(1)
  ROOT add = f32[] add(a, b)
}

ENTRY entry {
  parameter = f32[3,3]{1,0} parameter(0), sharding={devices=[2,1]0,1}
  ROOT all-reduce = f32[3,3]{1,0} all-reduce(parameter), to_apply=sum,
    replica_groups={}, sharding={devices=[2,1]0,1}
})";
  TF_ASSERT_OK_AND_ASSIGN(auto module,
                          PartitionComputation(hlo_string, /*num_devices=*/2));
  VLOG(1) << module->ToString();
  HloInstruction* root = module->entry_computation()->root_instruction();
  EXPECT_THAT(
      root, AllOf(op::Shape("f32[2,3]{1,0}"), op::AllReduce(op::Parameter(0))));
}

TEST_F(SpmdPartitioningTest, BroadcastOnlyNewDimsSharded) {
  absl::string_view hlo_string = R"(
HloModule module

ENTRY entry {
  constant = f32[4,3]{1,0} constant({{1,1,1},{1,1,1},{1,1,1},{1,1,1}}),
    sharding={replicated}
  ROOT broadcast = f32[3,4,3]{2,1,0} broadcast(constant), dimensions={1,2},
    sharding={devices=[2,1,1]0,1}
})";
  TF_ASSERT_OK_AND_ASSIGN(auto module,
                          PartitionComputation(hlo_string, /*num_devices=*/2));
  VLOG(1) << module->ToString();
  HloInstruction* root = module->entry_computation()->root_instruction();
  EXPECT_THAT(root, AllOf(op::Shape("f32[2,4,3]{2,1,0}"),
                          op::Broadcast(op::Constant())));
}

TEST_F(SpmdPartitioningTest, BroadcastOnlyOldDimsSharded) {
  absl::string_view hlo_string = R"(
HloModule module

ENTRY entry {
  constant = f32[4,3]{1,0} constant({{1,1,1},{1,1,1},{1,1,1},{1,1,1}}),
    sharding={replicated}
  ROOT broadcast = f32[4,4,3]{2,1,0} broadcast(constant), dimensions={1,2},
    sharding={devices=[1,2,1]0,1}
})";
  TF_ASSERT_OK_AND_ASSIGN(auto module,
                          PartitionComputation(hlo_string, /*num_devices=*/2));
  VLOG(1) << module->ToString();
  HloInstruction* root = module->entry_computation()->root_instruction();
  EXPECT_THAT(root, AllOf(op::Shape("f32[4,2,3]{2,1,0}"),
                          op::Broadcast(op::DynamicSlice(
                              op::Constant(), op::Reshape(), op::Constant()))));
}

TEST_F(SpmdPartitioningTest, BroadcastBothOldAndNewDimsSharded) {
  absl::string_view hlo_string = R"(
HloModule module

ENTRY entry {
  constant = f32[4,3]{1,0} constant({{1,1,1},{1,1,1},{1,1,1},{1,1,1}}),
    sharding={replicated}
  ROOT broadcast = f32[4,4,3]{2,1,0} broadcast(constant), dimensions={1,2},
    sharding={devices=[2,2,1]0,1,2,3}
})";
  TF_ASSERT_OK_AND_ASSIGN(auto module,
                          PartitionComputation(hlo_string, /*num_devices=*/4));
  VLOG(1) << module->ToString();
  HloInstruction* root = module->entry_computation()->root_instruction();
  EXPECT_THAT(
      root,
      AllOf(op::Shape("f32[2,2,3]{2,1,0}"),
            op::Broadcast(AllOf(op::Shape("f32[2,3]{1,0}"),
                                op::DynamicSlice(op::Constant(), op::Reshape(),
                                                 op::Constant())))));
}

TEST_F(SpmdPartitioningTest,
       BroadcastBothOldAndNewDimsShardedPartiallySharded) {
  absl::string_view hlo_string = R"(
HloModule module

ENTRY entry {
  param = f32[4,3] parameter(0),
    sharding={devices=[1,2,4]0,1,4,5,2,3,6,7 last_tile_dim_replicate}
  ROOT broadcast = f32[4,4,3] broadcast(param), dimensions={1,2},
    sharding={devices=[2,1,2,2]0,1,2,3,4,5,6,7 last_tile_dim_replicate}
})";
  TF_ASSERT_OK_AND_ASSIGN(auto module,
                          PartitionComputation(hlo_string, /*num_devices=*/8));
  VLOG(1) << module->ToString();
  HloInstruction* root = module->entry_computation()->root_instruction();
  EXPECT_THAT(
      root,
      AllOf(op::Shape("f32[2,4,2]"),
            op::Broadcast(AllOf(op::Shape("f32[4,2]"), op::Parameter(0)))));
}

TEST_F(SpmdPartitioningTest,
       ConvWithParallelDimAndNonParallelSpatialDimPartitioned) {
  absl::string_view hlo_string = R"(
HloModule module

ENTRY entry {
  %lhs = f32[32,12,12,24,32] parameter(0)
  %lhs.copy = f32[32,12,12,24,32] copy(%lhs),
    sharding={devices=[2,2,1,1,1]0,1,2,3}
  %rhs = f32[32,6,6,16,32] parameter(1)
  %rhs.copy = f32[32,6,6,16,32] copy(%rhs),
    sharding={devices=[2,2,1,1,1]0,1,2,3}
  ROOT %conv = f32[32,7,7,24,16] convolution(%lhs.copy, %rhs.copy),
    dim_labels=012bf_012oi->012bf,
    window={size=32x6x6 stride=31x1x1 lhs_dilate=32x1x1},
    sharding={devices=[2,2,1,1,1]0,1,2,3}
})";

  TF_ASSERT_OK_AND_ASSIGN(auto module,
                          PartitionComputation(hlo_string, /*num_devices=*/4));
  VLOG(1) << module->ToString();
  auto root = module->entry_computation()->root_instruction();
  auto lhs = AllOf(op::Copy(op::DynamicSlice(op::Parameter(), op::Reshape(),
                                             op::Reshape(), op::Constant(),
                                             op::Constant(), op::Constant())),
                   op::Shape("f32[16,6,12,24,32]"));
  auto rhs = AllOf(op::Copy(op::DynamicSlice(op::Parameter(), op::Reshape(),
                                             op::Reshape(), op::Constant(),
                                             op::Constant(), op::Constant())),
                   op::Shape("f32[16,3,6,16,32]"));
  auto resharded_rhs =
      AllOf(op::Shape("f32[16,6,6,16,32]"),
            op::AllReduce(op::DynamicUpdateSlice(
                op::Broadcast(), rhs, op::Constant(), op::Reshape(),
                op::Constant(), op::Constant(), op::Constant())));

  auto left_halo = AllOf(op::CollectivePermute(op::Slice(lhs)),
                         op::Shape("f32[16,2,12,24,32]"));
  auto right_halo = AllOf(op::CollectivePermute(op::Slice(lhs)),
                          op::Shape("f32[16,3,12,24,32]"));
  EXPECT_THAT(
      root,
      AllOf(op::Convolution(
                op::Select(op::Compare(),
                           op::DynamicSlice(
                               op::Concatenate(left_halo, lhs, right_halo),
                               op::Constant(), op::Add(), op::Constant(),
                               op::Constant(), op::Constant()),
                           op::Broadcast()),
                resharded_rhs),
            op::Shape("f32[16,4,7,24,16]")));
}

TEST_F(SpmdPartitioningTest, BroadcastPropagateTiledSharding) {
  absl::string_view hlo_string = R"(
HloModule module

ENTRY entry {
  constant = f32[4,3]{1,0} constant({{1,1,1},{1,4,1},{1,3,1},{1,2,1}}),
    sharding={devices=[2,1]0,1}
  ROOT broadcast = f32[4,4,3]{2,1,0} broadcast(constant), dimensions={1,2},
    sharding={devices=[1,2,1]0,1}
})";
  TF_ASSERT_OK_AND_ASSIGN(auto module,
                          PartitionComputation(hlo_string, /*num_devices=*/2));
  VLOG(1) << module->ToString();
  HloInstruction* root = module->entry_computation()->root_instruction();
  EXPECT_THAT(root, AllOf(op::Shape("f32[4,2,3]{2,1,0}"),
                          op::Broadcast(op::DynamicSlice(
                              op::Constant(), op::Reshape(), op::Constant()))));
}

TEST_F(SpmdPartitioningTest, OutfeedSingleDevice) {
  absl::string_view hlo_string = R"(
HloModule module

ENTRY entry {
  token.0 = token[] after-all()
  data = f32[1024]{0} parameter(0), sharding={maximal device=0}
  outfeed = token[] outfeed(data, token.0), sharding={maximal device=0}
})";
  TF_ASSERT_OK_AND_ASSIGN(auto module,
                          PartitionComputation(hlo_string, /*num_devices=*/2));
  VLOG(1) << module->ToString();
  HloInstruction* root = module->entry_computation()->root_instruction();
  EXPECT_THAT(root, AllOf(op::Shape("token[]"),
                          op::Conditional(
                              op::Compare(op::PartitionId(), op::Constant()),
                              op::Tuple(op::Parameter(0), op::AfterAll()),
                              op::Tuple(op::Parameter(0), op::AfterAll()))));

  HloInstruction* root_b0 = root->branch_computation(0)->root_instruction();
  EXPECT_THAT(root_b0,
              AllOf(op::Shape("token[]"),
                    op::Outfeed(op::GetTupleElement(op::Parameter(), 0),
                                op::GetTupleElement(op::Parameter(), 1))));

  HloInstruction* root_b1 = root->branch_computation(1)->root_instruction();
  EXPECT_THAT(root_b1, AllOf(op::Shape("token[]"), op::AfterAll()));
}

TEST_F(SpmdPartitioningTest, OutfeedEvenlyTiled) {
  absl::string_view hlo_string = R"(
HloModule module

ENTRY entry {
  token.0 = token[] after-all()
  data = f32[1024]{0} parameter(0), sharding={devices=[2]0,1}
  ROOT outfeed = token[] outfeed(data, token.0), sharding={devices=[2]0,1}
})";
  TF_ASSERT_OK_AND_ASSIGN(auto module,
                          PartitionComputation(hlo_string, /*num_devices=*/2));
  VLOG(1) << module->ToString();
  HloInstruction* root = module->entry_computation()->root_instruction();
  EXPECT_THAT(root, AllOf(op::Shape("token[]"),
                          op::Outfeed(op::Parameter(), op::AfterAll())));
}

TEST_F(SpmdPartitioningTest, OutfeedTupleEvenlyTiled) {
  absl::string_view hlo_string = R"(
HloModule module

ENTRY entry {
  token.0 = token[] after-all()
  data = (f32[1024,2]{1,0}, f32[2]{0}) parameter(0), sharding={{devices=[2,1]0,1},
    {devices=[2]0,1}}
  ROOT outfeed = token[] outfeed(data, token.0),
    outfeed_shape=(f32[1024,2]{0,1}, f32[2]{0}), sharding={{devices=[2,1]0,1},
    {devices=[2]0,1}}
})";
  TF_ASSERT_OK_AND_ASSIGN(auto module,
                          PartitionComputation(hlo_string, /*num_devices=*/2));
  VLOG(1) << module->ToString();
  HloInstruction* root = module->entry_computation()->root_instruction();
  EXPECT_THAT(root, AllOf(op::Shape("token[]"),
                          op::Outfeed(op::Parameter(), op::AfterAll())));
  auto expected_layout0 = LayoutUtil::MakeLayout({0, 1});
  auto expected_layout1 = LayoutUtil::MakeLayout({0});
  EXPECT_TRUE(LayoutUtil::Equal(root->outfeed_shape().tuple_shapes(0).layout(),
                                expected_layout0));
  EXPECT_TRUE(LayoutUtil::Equal(root->outfeed_shape().tuple_shapes(1).layout(),
                                expected_layout1));
}

TEST_F(SpmdPartitioningTest, OutfeedReplicated) {
  absl::string_view hlo_string = R"(
HloModule module

ENTRY entry {
  token.0 = token[] after-all()
  data = (f32[1024,2]{1,0}, f32[2]{0}) parameter(0), sharding={{devices=[2,1]0,1},
    {replicated}}
  ROOT outfeed = token[] outfeed(data, token.0), sharding={{devices=[2,1]0,1},
    {replicated}}
})";
  TF_ASSERT_OK_AND_ASSIGN(auto module,
                          PartitionComputation(hlo_string, /*num_devices=*/2));
  VLOG(1) << module->ToString();
  HloInstruction* root = module->entry_computation()->root_instruction();
  EXPECT_THAT(root, AllOf(op::Shape("token[]"),
                          op::Outfeed(op::Parameter(), op::AfterAll())));
}

TEST_F(SpmdPartitioningTest, OutfeedUnevenlyTiled) {
  absl::string_view hlo_string = R"(
HloModule module

ENTRY entry {
  token.0 = token[] after-all()
  data = (f32[1023,2]{1,0}, f32[3]{0}) parameter(0), sharding={{devices=[2,1]0,1},
    {devices=[2]0,1}}
  outfeed = token[] outfeed(data, token.0),
    outfeed_shape=(f32[1023,2]{0,1}, f32[3]{0}), sharding={{devices=[2,1]0,1},
    {devices=[2]0,1}}
})";
  TF_ASSERT_OK_AND_ASSIGN(auto module,
                          PartitionComputation(hlo_string, /*num_devices=*/2));
  VLOG(1) << module->ToString();

  HloInstruction* root = module->entry_computation()->root_instruction();
  EXPECT_THAT(
      root, AllOf(op::Shape("token[]"),
                  op::Conditional(op::Convert(),
                                  op::Tuple(op::Parameter(), op::AfterAll()),
                                  op::Tuple(op::Parameter(), op::AfterAll()))));

  auto first_outfeed =
      AllOf(op::Shape("(f32[512,2], f32[2])"), op::GetTupleElement());
  EXPECT_THAT(root->called_computations()[0]->root_instruction(),
              AllOf(op::Shape("token[]"),
                    op::Outfeed(first_outfeed, op::GetTupleElement())));

  auto second_outfeed = AllOf(op::Shape("(f32[511,2], f32[1])"), op::Tuple());
  EXPECT_THAT(root->called_computations()[1]->root_instruction(),
              AllOf(op::Shape("token[]"),
                    op::Outfeed(second_outfeed, op::GetTupleElement())));

  auto expected_layout0 = LayoutUtil::MakeLayout({0, 1});
  auto expected_layout1 = LayoutUtil::MakeLayout({0});
  auto first_outfeed_instr = root->called_computations()[0]->root_instruction();
  auto second_outfeed_instr =
      root->called_computations()[1]->root_instruction();
  EXPECT_TRUE(LayoutUtil::Equal(
      first_outfeed_instr->outfeed_shape().tuple_shapes(0).layout(),
      expected_layout0));
  EXPECT_TRUE(LayoutUtil::Equal(
      first_outfeed_instr->outfeed_shape().tuple_shapes(1).layout(),
      expected_layout1));
  EXPECT_TRUE(LayoutUtil::Equal(
      second_outfeed_instr->outfeed_shape().tuple_shapes(0).layout(),
      expected_layout0));
  EXPECT_TRUE(LayoutUtil::Equal(
      second_outfeed_instr->outfeed_shape().tuple_shapes(1).layout(),
      expected_layout1));
}

TEST_F(SpmdPartitioningTest, ReduceWindowReplicatedInput) {
  absl::string_view hlo_string = R"(
HloModule module

sum {
  a = f32[] parameter(0)
  b = f32[] parameter(1)
  ROOT add = f32[] add(a, b)
}

ENTRY entry {
  constant = f32[6,2]{1,0} constant({{1,1},{1,4},{2,1},{3,1},{1,2},{2,2}}),
    sharding={replicated}
  constant.1 = f32[] constant(0), sharding={replicated}
  ROOT reduce-window = f32[3,2]{1,0} reduce-window(constant, constant.1),
    window={size=3x1 stride=2x1 pad=1_0x0_0}, to_apply=sum,
    sharding={devices=[2,1]0,1}
})";
  TF_ASSERT_OK_AND_ASSIGN(auto module,
                          PartitionComputation(hlo_string, /*num_devices=*/2));
  VLOG(1) << module->ToString();
  HloInstruction* root = module->entry_computation()->root_instruction();
  EXPECT_THAT(
      root,
      AllOf(op::Shape("f32[2,2]{1,0}"),
            op::ReduceWindow(
                op::DynamicSlice(AllOf(op::Shape("f32[9,2]{1,0}"),
                                       op::Pad(op::Constant(), op::Constant())),
                                 op::Multiply(op::Reshape(), op::Constant()),
                                 op::Constant()),
                op::Constant())));
}

TEST_F(SpmdPartitioningTest, ReduceWindowTiledNegativeLeftHalo) {
  absl::string_view hlo_string = R"(
HloModule module

sum {
  a = f32[] parameter(0)
  b = f32[] parameter(1)
  ROOT add = f32[] add(a, b)
}

ENTRY entry {
  constant = f32[6,2]{1,0} constant({{1,1},{1,4},{2,1},{3,1},{1,2},{2,2}}),
    sharding={devices=[2,1]0,1}
  constant.1 = f32[] constant(0), sharding={replicated}
  ROOT %reduce-window = f32[3,2]{1,0} reduce-window(%constant, %constant.1),
    window={size=3x1 stride=2x1 pad=0_1x0_0}, to_apply=sum,
    sharding={devices=[2,1]0,1}
})";
  TF_ASSERT_OK_AND_ASSIGN(auto module,
                          PartitionComputation(hlo_string, /*num_devices=*/2));
  VLOG(1) << module->ToString();
  HloInstruction* root = module->entry_computation()->root_instruction();

  auto sharded_input =
      op::DynamicSlice(op::Constant(), op::Reshape(), op::Constant());
  auto right_halo = AllOf(op::Shape("f32[2,2]{1,0}"),
                          op::CollectivePermute(op::Slice(sharded_input)));
  auto pre_masking = op::DynamicSlice(
      AllOf(
          op::Shape("f32[6,2]{1,0}"),
          op::Pad(op::Concatenate(sharded_input, right_halo), op::Constant())),
      op::Reshape(), op::Constant());
  auto index_in_padded = op::Add(
      op::Iota(), op::Broadcast(op::Multiply(op::Reshape(), op::Constant())));
  auto masked =
      op::Select(op::Compare(index_in_padded, op::Broadcast(op::Constant())),
                 pre_masking, op::Broadcast(op::Constant()));
  EXPECT_THAT(root, AllOf(op::Shape("f32[2,2]{1,0}"),
                          op::ReduceWindow(masked, op::Constant())));
}

TEST_F(SpmdPartitioningTest, ReduceWindowTiledOneSideHaloBeyondNeighbor) {
  absl::string_view hlo_string = R"(
HloModule module

sum {
  a = f32[] parameter(0)
  b = f32[] parameter(1)
  ROOT add = f32[] add(a, b)
}

ENTRY entry {
  param = f32[9,2] parameter(0), sharding={devices=[5,1]0,1,2,3,4}
  constant.1 = f32[] constant(0), sharding={replicated}
  ROOT reduce-window = f32[5,2]{1,0} reduce-window(param, constant.1),
    window={size=4x1 stride=2x1 pad=3_0x0_0}, to_apply=sum,
    sharding={devices=[5,1]0,1,2,3,4}
})";
  TF_ASSERT_OK_AND_ASSIGN(auto module,
                          PartitionComputation(hlo_string, /*num_devices=*/5));
  VLOG(1) << module->ToString();
  auto halo0 = AllOf(op::Shape("f32[1,2]"),
                     op::CollectivePermute(op::Slice(op::Parameter(0))));
  auto halo1 =
      AllOf(op::Shape("f32[2,2]"), op::CollectivePermute(op::Parameter(0)));
  auto pre_mask =
      AllOf(op::Shape("f32[4,2]"),
            op::Slice(AllOf(op::Shape("f32[5,2]"),
                            op::Concatenate(halo0, halo1, op::Parameter(0)))));
  auto masked =
      op::Select(op::Compare(op::Add(op::Iota(), op::Broadcast(op::Multiply())),
                             op::Broadcast(op::Constant())),
                 pre_mask, op::Broadcast(op::Constant()));
  HloInstruction* root = module->entry_computation()->root_instruction();
  EXPECT_THAT(root, AllOf(op::Shape("f32[1,2]{1,0}"),
                          op::ReduceWindow(masked, op::Constant())));
}

TEST_F(SpmdPartitioningTest, ReduceWindowTiledOneSideUnequalHalo) {
  absl::string_view hlo_string = R"(
HloModule module

sum {
  a = f32[] parameter(0)
  b = f32[] parameter(1)
  ROOT add = f32[] add(a, b)
}

ENTRY entry {
  constant = f32[9,2]{1,0} constant(
    {{1,1},{1,4},{2,1},{3,1},{1,2},{2,2},{4,1},{1,2},{2,1}}),
    sharding={devices=[3,1]0,1,2}
  constant.1 = f32[] constant(0), sharding={replicated}
  ROOT reduce-window = f32[5,2]{1,0} reduce-window(constant, constant.1),
    window={size=3x1 stride=2x1 pad=1_1x0_0}, to_apply=sum,
    sharding={devices=[3,1]0,1,2}
})";
  TF_ASSERT_OK_AND_ASSIGN(auto module,
                          PartitionComputation(hlo_string, /*num_devices=*/3));
  VLOG(1) << module->ToString();
  HloInstruction* root = module->entry_computation()->root_instruction();

  auto sharded_input =
      op::DynamicSlice(op::Constant(), op::Reshape(), op::Constant());
  auto right_halo = AllOf(op::Shape("f32[2,2]{1,0}"),
                          op::CollectivePermute(op::Slice(sharded_input)));
  auto pre_masking = op::DynamicSlice(
      AllOf(
          op::Shape("f32[7,2]{1,0}"),
          op::Pad(op::Concatenate(sharded_input, right_halo), op::Constant())),
      op::Reshape(), op::Constant());
  auto index_in_padded = op::Add(
      op::Iota(), op::Broadcast(op::Multiply(op::Reshape(), op::Constant())));
  auto masked = op::Select(
      op::And(op::Compare(index_in_padded, op::Broadcast(op::Constant())),
              op::Compare(index_in_padded, op::Broadcast(op::Constant()))),
      pre_masking, op::Broadcast(op::Constant()));
  EXPECT_THAT(root, AllOf(op::Shape("f32[2,2]{1,0}"),
                          op::ReduceWindow(masked, op::Constant())));
}

TEST_F(SpmdPartitioningTest, ReduceWindowTiledTwoSideHalo) {
  absl::string_view hlo_string = R"(
HloModule module

sum {
  a = f32[] parameter(0)
  b = f32[] parameter(1)
  ROOT add = f32[] add(a, b)
}

ENTRY entry {
  constant = f32[4,2]{1,0} constant({{1,1},{1,4},{2,1},{3,1}}),
    sharding={devices=[2,1]0,1}
  constant.1 = f32[] constant(0), sharding={replicated}
  ROOT reduce-window = f32[2,2]{1,0} reduce-window(constant, constant.1),
    window={size=5x1 stride=3x1 pad=2_2x0_0}, to_apply=sum,
    sharding={devices=[2,1]0,1}
})";
  TF_ASSERT_OK_AND_ASSIGN(auto module,
                          PartitionComputation(hlo_string, /*num_devices=*/2));
  VLOG(1) << module->ToString();
  HloInstruction* root = module->entry_computation()->root_instruction();

  auto sharded_input =
      op::DynamicSlice(op::Constant(), op::Reshape(), op::Constant());
  auto left_halo = AllOf(op::Shape("f32[1,2]{1,0}"),
                         op::CollectivePermute(op::Slice(sharded_input)));
  auto right_halo = AllOf(op::Shape("f32[1,2]{1,0}"),
                          op::CollectivePermute(op::Slice(sharded_input)));
  auto pre_masking = AllOf(
      op::Shape("f32[5,2]{1,0}"),
      op::DynamicSlice(
          AllOf(op::Shape("f32[6,2]{1,0}"),
                op::Pad(op::Concatenate(left_halo, sharded_input, right_halo),
                        op::Constant())),
          op::Reshape(), op::Constant()));
  auto index_in_padded = op::Add(
      op::Iota(), op::Broadcast(op::Multiply(op::Reshape(), op::Constant())));
  auto masked = op::Select(
      op::And(op::Compare(index_in_padded, op::Broadcast(op::Constant())),
              op::Compare(index_in_padded, op::Broadcast(op::Constant()))),
      pre_masking, op::Broadcast(op::Constant()));
  EXPECT_THAT(root, AllOf(op::Shape("f32[1,2]{1,0}"),
                          op::ReduceWindow(masked, op::Constant())));
}

TEST_F(SpmdPartitioningTest, ReduceWindowTiled2D) {
  absl::string_view hlo_string = R"(
HloModule module

sum {
  a = f32[] parameter(0)
  b = f32[] parameter(1)
  ROOT add = f32[] add(a, b)
}

ENTRY entry {
  token0 = token[] after-all(), sharding={maximal device=0}
  infeed = (f32[4,4,2,2]{3,2,1,0}, token[]) infeed(token0),
    sharding={{devices=[2,2,1,1]0,1,2,3}, {maximal device=0}}
  infeed.data = f32[4,4,2,2]{3,2,1,0} get-tuple-element(infeed), index=0,
    sharding={devices=[2,2,1,1]0,1,2,3}
  constant = f32[] constant(0), sharding={replicated}
  ROOT reduce-window = f32[2,2,2,2]{3,2,1,0} reduce-window(infeed.data, constant),
    window={size=5x5x1x1 stride=3x3x1x1 pad=2_2x2_2x0_0x0_0}, to_apply=sum,
    sharding={devices=[2,2,1,1]0,1,2,3}
})";
  TF_ASSERT_OK_AND_ASSIGN(auto module,
                          PartitionComputation(hlo_string, /*num_devices=*/4));
  VLOG(1) << module->ToString();
  HloInstruction* root = module->entry_computation()->root_instruction();

  auto sharded_input = AllOf(op::Shape("f32[2,2,2,2]{3,2,1,0}"),
                             op::GetTupleElement(op::Infeed()));
  auto dim0_left_halo = AllOf(op::Shape("f32[1,2,2,2]{3,2,1,0}"),
                              op::CollectivePermute(op::Slice(sharded_input)));
  auto dim0_right_halo = AllOf(op::Shape("f32[1,2,2,2]{3,2,1,0}"),
                               op::CollectivePermute(op::Slice(sharded_input)));
  auto dim0_pre_masking = op::DynamicSlice(
      AllOf(op::Shape("f32[6,2,2,2]{3,2,1,0}"),
            op::Pad(
                op::Concatenate(dim0_left_halo, sharded_input, dim0_right_halo),
                op::Constant())),
      op::Reshape(), op::Constant(), op::Constant(), op::Constant());
  auto dim0_index_in_padded = op::Add(
      op::Iota(), op::Broadcast(op::Multiply(op::Reshape(), op::Constant())));
  auto dim0_masked = op::Select(
      op::And(op::Compare(dim0_index_in_padded, op::Broadcast(op::Constant())),
              op::Compare(dim0_index_in_padded, op::Broadcast(op::Constant()))),
      dim0_pre_masking, op::Broadcast(op::Constant()));
  auto dim0_resharded = AllOf(op::Shape("f32[5,2,2,2]{3,2,1,0}"), dim0_masked);
  auto dim1_left_halo = AllOf(op::Shape("f32[5,1,2,2]{3,2,1,0}"),
                              op::CollectivePermute(op::Slice(dim0_resharded)));
  auto dim1_right_halo =
      AllOf(op::Shape("f32[5,1,2,2]{3,2,1,0}"),
            op::CollectivePermute(op::Slice(dim0_resharded)));
  auto dim1_pre_masking = op::DynamicSlice(
      AllOf(op::Shape("f32[5,6,2,2]{3,2,1,0}"),
            op::Pad(op::Concatenate(dim1_left_halo, dim0_resharded,
                                    dim1_right_halo),
                    op::Constant())),
      op::Constant(), op::Reshape(), op::Constant(), op::Constant());
  auto dim1_index_in_padded = op::Add(
      op::Iota(), op::Broadcast(op::Multiply(op::Reshape(), op::Constant())));
  auto dim1_masked = op::Select(
      op::And(op::Compare(dim1_index_in_padded, op::Broadcast(op::Constant())),
              op::Compare(dim1_index_in_padded, op::Broadcast(op::Constant()))),
      dim1_pre_masking, op::Broadcast(op::Constant()));
  auto dim1_resharded = AllOf(op::Shape("f32[5,5,2,2]{3,2,1,0}"), dim1_masked);
  EXPECT_THAT(root, AllOf(op::Shape("f32[1,1,2,2]{3,2,1,0}"),
                          op::ReduceWindow(dim1_resharded, op::Constant())));
}

TEST_F(SpmdPartitioningTest, ConvolutionLhsTiledRhsReplicated) {
  absl::string_view hlo_string = R"(
HloModule module

ENTRY entry {
  %lhs = f32[128,224,224,3] parameter(0)
  %lhs.copy = f32[128,224,224,3] copy(f32[128,224,224,3] %lhs),
    sharding={devices=[1,2,1,1]0,1}
  %rhs = f32[7,7,3,64] parameter(1)
  %rhs.copy = f32[7,7,3,64] copy(f32[7,7,3,64] %rhs),
    sharding={replicated}
  ROOT %conv = f32[128,112,112,64] convolution(
    f32[128,224,224,3] %lhs.copy,
    f32[7,7,3,64] %rhs.copy),
    window={size=7x7 stride=2x2 pad=3_3x3_3},
    dim_labels=b01f_01io->b01f,
    sharding={devices=[1,2,1,1]0,1}
})";

  TF_ASSERT_OK_AND_ASSIGN(auto module,
                          PartitionComputation(hlo_string, /*num_devices=*/2));
  VLOG(1) << module->ToString();

  auto root = module->entry_computation()->root_instruction();
  auto lhs = AllOf(
      op::Copy(op::DynamicSlice(op::Parameter(), op::Constant(), op::Reshape(),
                                op::Constant(), op::Constant())),
      op::Shape("f32[128,112,224,3]"));
  auto rhs = AllOf(op::Copy(op::Parameter()), op::Shape("f32[7,7,3,64]"));

  auto left_halo = AllOf(op::CollectivePermute(op::Slice(lhs)),
                         op::Shape("f32[128,3,224,3]"));
  auto right_halo = AllOf(op::CollectivePermute(op::Slice(lhs)),
                          op::Shape("f32[128,2,224,3]"));
  EXPECT_THAT(root,
              AllOf(op::Convolution(
                        op::Select(op::And(),
                                   op::Concatenate(left_halo, lhs, right_halo),
                                   op::Broadcast()),
                        rhs),
                    op::Shape("f32[128,56,112,64]")));
}

TEST_F(SpmdPartitioningTest, ConvolutionLhsTiledRhsReplicatedNeedReshard) {
  absl::string_view hlo_string = R"(
HloModule module

ENTRY entry {
  %lhs = f32[128,224,224,3] parameter(0)
  %lhs.copy = f32[128,224,224,3] copy(f32[128,224,224,3] %lhs),
    sharding={devices=[2,1,1,1]0,1}
  %rhs = f32[7,7,3,64] parameter(1)
  %rhs.copy = f32[7,7,3,64] copy(f32[7,7,3,64] %rhs),
    sharding={replicated}
  ROOT %conv = f32[128,112,112,64] convolution(
    f32[128,224,224,3] %lhs.copy,
    f32[7,7,3,64] %rhs.copy),
    window={size=7x7 stride=2x2 pad=3_3x3_3},
    dim_labels=b01f_01io->b01f,
    sharding={devices=[1,2,1,1]0,1}
})";

  TF_ASSERT_OK_AND_ASSIGN(auto module,
                          PartitionComputation(hlo_string, /*num_devices=*/2));
  VLOG(1) << module->ToString();

  auto root = module->entry_computation()->root_instruction();
  auto lhs = AllOf(
      op::Copy(op::DynamicSlice(op::Parameter(), op::Reshape(), op::Constant(),
                                op::Constant(), op::Constant())),
      op::Shape("f32[64,224,224,3]"));
  auto all_to_all =
      AllOf(op::AllToAll(op::Reshape(lhs)), op::Shape("f32[64,2,112,224,3]"));
  auto reshard_lhs = AllOf(op::Reshape(op::Transpose(all_to_all)),
                           op::Shape("f32[128,112,224,3]"));

  auto rhs = AllOf(op::Copy(op::Parameter()), op::Shape("f32[7,7,3,64]"));

  auto left_halo = AllOf(op::CollectivePermute(op::Slice(reshard_lhs)),
                         op::Shape("f32[128,3,224,3]"));
  auto right_halo = AllOf(op::CollectivePermute(op::Slice(reshard_lhs)),
                          op::Shape("f32[128,2,224,3]"));
  EXPECT_THAT(
      root,
      AllOf(op::Convolution(
                op::Select(op::And(),
                           op::Concatenate(left_halo, reshard_lhs, right_halo),
                           op::Broadcast()),
                rhs),
            op::Shape("f32[128,56,112,64]")));
}

TEST_F(SpmdPartitioningTest, ConvolutionLhsTiledRhsReplicatedReordered) {
  absl::string_view hlo_string = R"(
HloModule module

ENTRY entry {
  %lhs = f32[224,224,3,128] parameter(0)
  %lhs.copy = f32[224,224,3,128] copy(%lhs), sharding={devices=[2,1,1,1]0,1}
  %rhs = f32[7,7,3,64] parameter(1)
  %rhs.copy = f32[7,7,3,64] copy(%rhs), sharding={replicated}
  ROOT %conv = f32[128,112,112,64] convolution(%lhs.copy, %rhs.copy),
    window={size=7x7 stride=2x2 pad=3_3x3_3},
    dim_labels=01fb_01io->b01f,
    sharding={devices=[1,2,1,1]0,1}
})";

  TF_ASSERT_OK_AND_ASSIGN(auto module,
                          PartitionComputation(hlo_string, /*num_devices=*/2));
  VLOG(1) << module->ToString();

  auto root = module->entry_computation()->root_instruction();
  auto lhs = AllOf(
      op::Copy(op::DynamicSlice(op::Parameter(), op::Reshape(), op::Constant(),
                                op::Constant(), op::Constant())),
      op::Shape("f32[112,224,3,128]"));
  auto rhs = AllOf(op::Copy(op::Parameter()), op::Shape("f32[7,7,3,64]"));

  auto left_halo = AllOf(op::CollectivePermute(op::Slice(lhs)),
                         op::Shape("f32[3,224,3,128]"));
  auto right_halo = AllOf(op::CollectivePermute(op::Slice(lhs)),
                          op::Shape("f32[2,224,3,128]"));
  EXPECT_THAT(root,
              AllOf(op::Convolution(
                        op::Select(op::And(),
                                   op::Concatenate(left_halo, lhs, right_halo),
                                   op::Broadcast()),
                        rhs),
                    op::Shape("f32[128,56,112,64]")));
}

// (stride * per_shard_window_count) % dilation == 0
TEST_F(SpmdPartitioningTest,
       ConvolutionBaseDilationSameStartPatternLhsTiledRhsReplicated) {
  absl::string_view hlo_string = R"(
HloModule module

ENTRY entry {
  %lhs = f32[128,7,7,512] parameter(0)
  %lhs.copy = f32[128,7,7,512] copy(%lhs),
    sharding={devices=[1,2,1,1]0,1}
  %rhs = f32[3,3,512,512] parameter(1)
  %rhs.copy = f32[3,3,512,512] copy(%rhs),
    sharding={replicated}
  ROOT %conv = f32[128,4,4,512] convolution(%lhs.copy, %rhs.copy),
    window={size=3x3 stride=4x4 pad=1_1x1_1 lhs_dilate=2x2 rhs_reversal=1x1},
    dim_labels=b01f_01io->b01f,
    sharding={devices=[1,2,1,1]0,1}
})";

  TF_ASSERT_OK_AND_ASSIGN(auto module,
                          PartitionComputation(hlo_string, /*num_devices=*/2));
  VLOG(1) << module->ToString();

  auto root = module->entry_computation()->root_instruction();
  // There is no halo exchange, and because the last element in the shard is not
  // needed (stride == 4), the LHS will be just a slice.
  auto sliced_lhs =
      AllOf(op::Slice(op::Copy(op::DynamicSlice(
                op::Pad(op::Parameter(), op::Constant()), op::Constant(),
                op::Reshape(), op::Constant(), op::Constant()))),
            op::Shape("f32[128,3,7,512]"));
  auto rhs = AllOf(op::Copy(op::Parameter()), op::Shape("f32[3,3,512,512]"));
  EXPECT_THAT(root, AllOf(op::Convolution(sliced_lhs, rhs),
                          op::Shape("f32[128,2,4,512]")));
  EXPECT_EQ(root->window().dimensions(0).padding_low(), 1);
  EXPECT_EQ(root->window().dimensions(0).padding_high(), 1);
}

// (stride * per_shard_window_count) % dilation != 0 but stride == 1
TEST_F(SpmdPartitioningTest,
       ConvolutionBaseDilationStride1LhsTiledRhsReplicated) {
  absl::string_view hlo_string = R"(
HloModule module

ENTRY entry {
  %lhs = f32[128,7,7,512] parameter(0)
  %lhs.copy = f32[128,7,7,512] copy(%lhs),
    sharding={devices=[1,2,1,1]0,1}
  %rhs = f32[3,3,512,512] parameter(1)
  %rhs.copy = f32[3,3,512,512] copy(%rhs),
    sharding={replicated}
  ROOT %conv = f32[128,14,14,512] convolution(%lhs.copy, %rhs.copy),
    window={size=3x3 pad=1_2x1_2 lhs_dilate=2x2 rhs_reversal=1x1},
    dim_labels=b01f_01io->b01f,
    sharding={devices=[1,2,1,1]0,1}
})";

  TF_ASSERT_OK_AND_ASSIGN(auto module,
                          PartitionComputation(hlo_string, /*num_devices=*/2));
  VLOG(1) << module->ToString();

  auto root = module->entry_computation()->root_instruction();
  auto lhs = AllOf(op::Copy(op::DynamicSlice(
                       op::Pad(op::Parameter(), op::Constant()), op::Constant(),
                       op::Reshape(), op::Constant(), op::Constant())),
                   op::Shape("f32[128,4,7,512]"));
  auto rhs = AllOf(op::Copy(op::Parameter()), op::Shape("f32[3,3,512,512]"));

  auto left_halo = AllOf(op::CollectivePermute(op::Slice(lhs)),
                         op::Shape("f32[128,1,7,512]"));
  auto start_window = op::Multiply(op::Reshape(), op::Constant());
  auto start_input_element = op::Divide(start_window, op::Constant());
  auto dynamic_offset_for_padded_concat = op::Subtract(
      op::Constant(), op::Subtract(op::Multiply(op::Reshape(), op::Constant()),
                                   start_input_element));
  auto pre_masking =
      AllOf(op::Shape("f32[128,5,7,512]"),
            op::DynamicSlice(
                AllOf(op::Shape("f32[128,6,7,512]"),
                      op::Pad(op::Concatenate(left_halo, lhs), op::Constant())),
                op::Constant(), dynamic_offset_for_padded_concat,
                op::Constant(), op::Constant()));
  auto masked = op::Select(
      op::Compare(op::Add(op::Iota(), op::Broadcast(start_input_element)),
                  op::Broadcast(op::Constant())),
      pre_masking, op::Broadcast(op::Constant()));
  auto dynamic_offset_on_output = op::Subtract(
      start_window, op::Multiply(start_input_element, op::Constant()));
  EXPECT_THAT(root,
              AllOf(op::DynamicSlice(AllOf(op::Convolution(masked, rhs),
                                           op::Shape("f32[128,8,14,512]")),
                                     op::Constant(), dynamic_offset_on_output,
                                     op::Constant(), op::Constant()),
                    op::Shape("f32[128,7,14,512]")));
  EXPECT_EQ(root->operand(0)->window().dimensions(0).padding_low(), 1);
  EXPECT_EQ(root->operand(0)->window().dimensions(0).padding_high(), 0);
}

TEST_F(SpmdPartitioningTest, SelectAndScatterNoOverlap) {
  absl::string_view hlo_string = R"(
HloModule module

ge {
  a = f32[] parameter(0)
  b = f32[] parameter(1)
  ROOT compare = pred[] compare(a, b), direction=GE
}

sum {
  c = f32[] parameter(0)
  d = f32[] parameter(1)
  ROOT add = f32[] add(c, d)
}

ENTRY entry {
  %param = f32[11,4]{1,0} parameter(0)
  %param.copy = f32[11,4] copy(%param),
    sharding={devices=[4,1]0,1,2,3}
  constant = f32[4,2]{1,0} constant({{1,2},{3,4},{1,0},{2,8}}),
    sharding={devices=[4,1]0,1,2,3}
  constant.1 = f32[] constant(0), sharding={replicated}
  ROOT select-and-scatter = f32[11,4]{1,0} select-and-scatter(param.copy,
    constant, constant.1), window={size=3x2 stride=3x2 pad=0_1x0_0},
    select=ge, scatter=sum, sharding={devices=[4,1]0,1,2,3}
})";
  TF_ASSERT_OK_AND_ASSIGN(auto module,
                          PartitionComputation(hlo_string, /*num_devices=*/4));
  VLOG(1) << module->ToString();
  auto root = module->entry_computation()->root_instruction();
  auto source =
      AllOf(op::Shape("f32[1,2]{1,0}"),
            op::DynamicSlice(op::Constant(), op::Reshape(), op::Constant()));
  auto masked_data = AllOf(
      op::Shape("f32[3,4]{1,0}"),
      op::Select(
          op::Compare(op::Add(op::Iota(), op::Broadcast(op::Multiply(
                                              op::Reshape(), op::Constant()))),
                      op::Broadcast(op::Constant())),
          op::Copy(op::DynamicSlice(op::Pad(op::Parameter(), op::Constant()),
                                    op::Reshape(), op::Constant())),
          op::Broadcast(op::Constant())));

  EXPECT_THAT(root,
              AllOf(op::SelectAndScatter(masked_data, source, op::Constant()),
                    op::Shape("f32[3,4]{1,0}")));
  EXPECT_EQ(root->window().dimensions(0).padding_low(), 0);
  EXPECT_EQ(root->window().dimensions(0).padding_high(), 0);
}

TEST_F(SpmdPartitioningTest, SelectAndScatterNoOverlapReshard) {
  absl::string_view hlo_string = R"(
HloModule module

ge {
  a = f32[] parameter(0)
  b = f32[] parameter(1)
  ROOT compare = pred[] compare(a, b), direction=GE
}

sum {
  c = f32[] parameter(0)
  d = f32[] parameter(1)
  ROOT add = f32[] add(c, d)
}

ENTRY entry {
  %param = f32[11,4]{1,0} parameter(0)
  %param.copy = f32[11,4] copy(%param),
    sharding={devices=[1,4]0,1,2,3}
  constant = f32[4,2]{1,0} constant({{1,2},{3,4},{1,0},{2,8}}),
    sharding={devices=[4,1]0,1,2,3}
  constant.1 = f32[] constant(0), sharding={replicated}
  ROOT select-and-scatter = f32[11,4]{1,0} select-and-scatter(param.copy,
    constant, constant.1), window={size=3x2 stride=3x2 pad=0_1x0_0},
    select=ge, scatter=sum, sharding={devices=[4,1]0,1,2,3}
})";
  TF_ASSERT_OK_AND_ASSIGN(auto module,
                          PartitionComputation(hlo_string, /*num_devices=*/4));
  VLOG(1) << module->ToString();
  auto root = module->entry_computation()->root_instruction();
  auto source =
      AllOf(op::Shape("f32[1,2]{1,0}"),
            op::DynamicSlice(op::Constant(), op::Reshape(), op::Constant()));
  auto operand = AllOf(op::Copy(op::DynamicSlice(
                           op::Parameter(0), op::Constant(), op::Reshape())),
                       op::Shape("f32[11,1]"));
  auto reshard_operand = op::Reshape(op::Transpose(
      op::AllToAll(op::Reshape(op::Pad(operand, op::Constant())))));
  auto masked_data = AllOf(
      op::Shape("f32[3,4]{1,0}"),
      op::Select(
          op::Compare(op::Add(op::Iota(), op::Broadcast(op::Multiply(
                                              op::Reshape(), op::Constant()))),
                      op::Broadcast(op::Constant())),
          reshard_operand, op::Broadcast(op::Constant())));

  EXPECT_THAT(root,
              AllOf(op::SelectAndScatter(masked_data, source, op::Constant()),
                    op::Shape("f32[3,4]{1,0}")));
  EXPECT_EQ(root->window().dimensions(0).padding_low(), 0);
  EXPECT_EQ(root->window().dimensions(0).padding_high(), 0);
}

TEST_F(SpmdPartitioningTest, SelectAndScatterWithOverlap) {
  absl::string_view hlo_string = R"(
HloModule module

ge {
  a = f32[] parameter(0)
  b = f32[] parameter(1)
  ROOT compare = pred[] compare(a, b), direction=GE
}

sum {
  c = f32[] parameter(0)
  d = f32[] parameter(1)
  ROOT add = f32[] add(c, d)
}

ENTRY entry {
  %param = f32[11,4]{1,0} parameter(0)
  %param.copy = f32[11,4] copy(%param),
    sharding={devices=[4,1]0,1,2,3}
  constant = f32[6,2]{1,0} constant({{1,2},{3,4},{1,0},{2,8},{6,6},{1,9}}),
    sharding={devices=[4,1]0,1,2,3}
  constant.1 = f32[] constant(0), sharding={replicated}
  ROOT select-and-scatter = f32[11,4]{1,0} select-and-scatter(param.copy,
    constant, constant.1), window={size=3x2 stride=2x2 pad=1_1x0_0},
    select=ge, scatter=sum, sharding={devices=[4,1]0,1,2,3}
})";
  TF_ASSERT_OK_AND_ASSIGN(auto module,
                          PartitionComputation(hlo_string, /*num_devices=*/4));
  VLOG(1) << module->ToString();
  auto root = module->entry_computation()->root_instruction();

  auto source_shard =
      AllOf(op::Shape("f32[2,2]{1,0}"),
            op::DynamicSlice(op::Pad(), op::Reshape(), op::Constant()));
  // Max halo size is the same as the shard size, so slice is not needed.
  auto source_left_halo = op::CollectivePermute(source_shard);
  auto required_source_shard_start =
      op::Divide(op::Multiply(op::Reshape(), op::Constant()), op::Constant());
  auto source_with_halo = op::DynamicSlice(
      AllOf(op::Shape("f32[5,2]{1,0}"),
            op::Pad(op::Concatenate(source_left_halo, source_shard),
                    op::Constant())),
      op::Subtract(op::Constant(),
                   op::Subtract(op::Multiply(op::Reshape(), op::Constant()),
                                required_source_shard_start)),
      op::Constant());
  auto masked_source_with_halo = AllOf(
      AllOf(op::Shape("f32[3,2]{1,0}")),
      op::Select(
          op::Compare(
              op::Add(op::Iota(), op::Broadcast(required_source_shard_start)),
              op::Broadcast(op::Constant())),
          source_with_halo, op::Broadcast(op::Constant())));

  auto data_shard =
      AllOf(op::Shape("f32[3,4]{1,0}"),
            op::Copy(op::DynamicSlice(op::Pad(op::Parameter(), op::Constant()),
                                      op::Reshape(), op::Constant())));
  auto data_left_halo = AllOf(op::Shape("f32[2,4]{1,0}"),
                              op::CollectivePermute(op::Slice(data_shard)));
  auto data_right_halo = AllOf(op::Shape("f32[2,4]{1,0}"),
                               op::CollectivePermute(op::Slice(data_shard)));
  auto required_data_start_on_padded =
      op::Multiply(required_source_shard_start, op::Constant());
  auto left_halo_size = op::Subtract(
      op::Add(op::Multiply(op::Reshape(), op::Constant()), op::Constant()),
      required_data_start_on_padded);
  auto data_with_halo =
      AllOf(op::Shape("f32[7,4]{1,0}"),
            op::DynamicSlice(
                AllOf(op::Shape("f32[8,4]{1,0}"),
                      op::Pad(op::Concatenate(data_left_halo, data_shard,
                                              data_right_halo),
                              op::Constant())),
                op::Subtract(op::Constant(), left_halo_size), op::Constant()));
  auto index_on_padded =
      op::Add(op::Iota(), op::Broadcast(required_data_start_on_padded));
  auto masked_data_with_halo = op::Select(
      op::And(op::Compare(index_on_padded, op::Broadcast(op::Constant())),
              op::Compare(index_on_padded, op::Broadcast(op::Constant()))),
      data_with_halo, op::Broadcast(op::Constant()));

  EXPECT_THAT(
      root, AllOf(op::DynamicSlice(op::SelectAndScatter(masked_data_with_halo,
                                                        masked_source_with_halo,
                                                        op::Constant()),
                                   left_halo_size, op::Constant()),
                  op::Shape("f32[3,4]{1,0}")));
  EXPECT_EQ(root->operand(0)->window().dimensions(0).padding_low(), 0);
  EXPECT_EQ(root->operand(0)->window().dimensions(0).padding_high(), 0);
}

TEST_F(SpmdPartitioningTest, ConvolutionLhsTiledRhsTiled) {
  absl::string_view hlo_string = R"(
HloModule module

ENTRY entry {
  %lhs = f32[128,56,56,64] parameter(0)
  %lhs.copy = f32[128,56,56,64] copy(%lhs), sharding={devices=[1,2,1,1]0,1}
  %rhs = f32[128,56,56,256] parameter(1)
  %rhs.copy = f32[128,56,56,256] copy(%rhs), sharding={devices=[1,2,1,1]0,1}
  ROOT %conv = f32[1,1,64,256] convolution(%lhs.copy, %rhs.copy),
    window={size=56x56}, dim_labels=f01b_i01o->01bf, sharding={replicated}
})";

  TF_ASSERT_OK_AND_ASSIGN(auto module,
                          PartitionComputation(hlo_string, /*num_devices=*/2));
  VLOG(1) << module->ToString();

  auto root = module->entry_computation()->root_instruction();
  auto lhs = AllOf(
      op::Copy(op::DynamicSlice(op::Parameter(), op::Constant(), op::Reshape(),
                                op::Constant(), op::Constant())),
      op::Shape("f32[128,28,56,64]"));
  auto rhs = AllOf(
      op::Copy(op::DynamicSlice(op::Parameter(), op::Constant(), op::Reshape(),
                                op::Constant(), op::Constant())),
      op::Shape("f32[128,28,56,256]"));

  EXPECT_THAT(root, AllOf(op::AllReduce(op::Convolution(lhs, rhs)),
                          op::Shape("f32[1,1,64,256]")));
}

TEST_F(SpmdPartitioningTest, ConvolutionLhsTiledRhsTiledWindowReversal) {
  absl::string_view hlo_string = R"(
HloModule module

ENTRY entry {
  %lhs = f32[5,128,64] parameter(0), sharding={devices=[2,1,1]0,1}
  %rhs = f32[5,128,256] parameter(1), sharding={devices=[2,1,1]1,0}
  ROOT %conv = f32[1,64,256] convolution(%lhs, %rhs),
    window={size=5 rhs_reversal=1}, dim_labels=0fb_0io->0bf,
    sharding={replicated}
})";

  TF_ASSERT_OK_AND_ASSIGN(auto module,
                          PartitionComputation(hlo_string, /*num_devices=*/2));
  VLOG(1) << module->ToString();

  auto lhs_masked =
      AllOf(op::Shape("f32[3,128,64]"), op::Select(_, op::Parameter(0), _));
  auto rhs_left_padded = op::Slice(op::Concatenate(
      op::CollectivePermute(op::Slice(op::Parameter(1))), op::Parameter(1)));
  auto rhs_masked =
      AllOf(op::Shape("f32[3,128,256]"), op::Select(_, rhs_left_padded, _));

  auto root = module->entry_computation()->root_instruction();
  EXPECT_THAT(root,
              AllOf(op::AllReduce(op::Convolution(lhs_masked, rhs_masked)),
                    op::Shape("f32[1,64,256]")));
}

TEST_F(SpmdPartitioningTest, DotLhsTiledRhsTiledWithReshard) {
  absl::string_view hlo_string = R"(
HloModule module

ENTRY entry {
  %lhs = f32[128,56,56,64] parameter(0)
  %lhs.copy = f32[128,56,56,64] copy(%lhs), sharding={devices=[1,2,1,1]0,1}
  %rhs = f32[128,56,56,256] parameter(1)
  %rhs.copy = f32[128,56,56,256] copy(%rhs), sharding={devices=[2,1,1,1]0,1}
  ROOT %conv = f32[1,1,64,256] convolution(%lhs.copy, %rhs.copy),
    window={size=56x56}, dim_labels=f01b_i01o->01bf, sharding={replicated}
})";

  TF_ASSERT_OK_AND_ASSIGN(auto module,
                          PartitionComputation(hlo_string, /*num_devices=*/2));
  VLOG(1) << module->ToString();

  auto root = module->entry_computation()->root_instruction();
  auto lhs = AllOf(
      op::Copy(op::DynamicSlice(op::Parameter(), op::Constant(), op::Reshape(),
                                op::Constant(), op::Constant())),
      op::Shape("f32[128,28,56,64]"));
  auto rhs = AllOf(
      op::Copy(op::DynamicSlice(op::Parameter(), op::Reshape(), op::Constant(),
                                op::Constant(), op::Constant())),
      op::Shape("f32[64,56,56,256]"));
  auto all_to_all =
      AllOf(op::AllToAll(op::Reshape(lhs)), op::Shape("f32[2,64,28,56,64]"));
  auto reshard = AllOf(op::Reshape(op::Transpose(all_to_all)));

  EXPECT_THAT(root, AllOf(op::AllReduce(op::Convolution(reshard, rhs)),
                          op::Shape("f32[1,1,64,256]")));
}

TEST_F(SpmdPartitioningTest, ConvolutionLhsTiledRhsTiledWithReshard) {
  absl::string_view hlo_string = R"(
HloModule module

ENTRY entry {
  %lhs = f32[128,56,56,512] parameter(0)
  %lhs.copy = f32[128,56,56,512] copy(%lhs), sharding={devices=[1,2,1,1]0,1}
  %rhs = f32[128,28,28,64] parameter(1)
  %rhs.copy = f32[128,28,28,64] copy(%rhs), sharding={devices=[2,1,1,1]0,1}
  ROOT %conv = f32[1,1,512,64] convolution(%lhs.copy, %rhs.copy),
    window={size=28x28 pad=0_-1x0_-1 rhs_dilate=2x2},
    dim_labels=f01b_i01o->01bf, sharding={replicated}
})";

  TF_ASSERT_OK_AND_ASSIGN(auto module,
                          PartitionComputation(hlo_string, /*num_devices=*/2));
  VLOG(1) << module->ToString();

  auto root = module->entry_computation()->root_instruction();
  auto lhs = AllOf(
      op::Copy(op::DynamicSlice(op::Parameter(), op::Constant(), op::Reshape(),
                                op::Constant(), op::Constant())),
      op::Shape("f32[128,28,56,512]"));
  auto rhs = AllOf(
      op::Copy(op::DynamicSlice(op::Parameter(), op::Reshape(), op::Constant(),
                                op::Constant(), op::Constant())),
      op::Shape("f32[64,28,28,64]"));
  auto all_to_all =
      AllOf(op::AllToAll(op::Reshape(rhs)), op::Shape("f32[64,2,14,28,64]"));
  auto reshard = op::Reshape(op::Transpose(all_to_all));

  EXPECT_THAT(root,
              AllOf(op::AllReduce(op::Convolution(op::Slice(lhs), reshard)),
                    op::Shape("f32[1,1,512,64]")));
}

TEST_F(SpmdPartitioningTest,
       ConvolutionLhsTiledRhsTiled_UnevenDilatedRHSPartitioned) {
  absl::string_view hlo_string = R"(
HloModule module

ENTRY entry {
  %lhs = f32[8,28,28,8] parameter(0)
  %lhs.copy = f32[8,28,28,8] copy(%lhs), sharding={devices=[1,4,1,1]0,1,2,3}
  %rhs = f32[8,14,14,64] parameter(1)
  %rhs.copy = f32[8,14,14,64] copy(%rhs), sharding={devices=[1,4,1,1]0,1,2,3}
  ROOT %conv = f32[1,1,8,64] convolution(%lhs.copy, %rhs.copy),
    window={size=14x14 pad=0_-1x0_-1 rhs_dilate=2x2},
    dim_labels=f01b_i01o->01bf, sharding={replicated}
})";

  TF_ASSERT_OK_AND_ASSIGN(auto module,
                          PartitionComputation(hlo_string, /*num_devices=*/4));
  VLOG(1) << module->ToString();

  auto root = module->entry_computation()->root_instruction();
  auto lhs = AllOf(
      op::Copy(op::DynamicSlice(op::Parameter(), op::Constant(), op::Reshape(),
                                op::Constant(), op::Constant())),
      op::Shape("f32[8,7,28,8]"));
  auto rhs = AllOf(op::Pad(op::Parameter(), op::Constant()),
                   op::Shape("f32[8,16,14,64]"));
  auto selected_rhs = AllOf(
      op::Select(op::Compare(),
                 op::Copy(op::DynamicSlice(rhs, op::Constant(), op::Reshape(),
                                           op::Constant(), op::Constant())),
                 op::Broadcast()),
      op::Shape("f32[8,4,14,64]"));
  auto right_halo =
      AllOf(op::CollectivePermute(op::Slice(lhs)), op::Shape("f32[8,2,28,8]"));
  auto selected_lhs =
      AllOf(op::DynamicSlice(
                op::Pad(op::Concatenate(lhs, right_halo), op::Constant()),
                op::Constant(), op::Reshape(), op::Constant(), op::Constant()),
            op::Shape("f32[8,7,28,8]"));
  EXPECT_THAT(root,
              AllOf(op::AllReduce(op::Convolution(selected_lhs, selected_rhs)),
                    op::Shape("f32[1,1,8,64]")));
}

TEST_F(SpmdPartitioningTest, ConvolutionLhsTiledRhsTiledWithPadding) {
  absl::string_view hlo_string = R"(
HloModule module

ENTRY entry {
  %lhs = f32[32,28,28,128] parameter(0)
  %lhs.copy = f32[32,28,28,128] copy(%lhs), sharding={devices=[1,2,1,1]0,1}
  %rhs = f32[32,28,28,64] parameter(1)
  %rhs.copy = f32[32,28,28,64] copy(%rhs), sharding={devices=[1,2,1,1]0,1}
  ROOT %conv = f32[3,3,128,64] convolution(%lhs.copy, %rhs.copy),
    window={size=28x28 pad=1_1x1_1}, dim_labels=f01b_i01o->01bf, sharding={replicated}
})";

  TF_ASSERT_OK_AND_ASSIGN(
      auto module,
      PartitionComputation(hlo_string, /*num_devices=*/2,
                           /*conv_halo_exchange_always_on_lhs=*/false));
  VLOG(1) << module->ToString();

  auto root = module->entry_computation()->root_instruction();
  auto lhs = AllOf(
      op::Copy(op::DynamicSlice(op::Parameter(), op::Constant(), op::Reshape(),
                                op::Constant(), op::Constant())),
      op::Shape("f32[32,14,28,128]"));
  auto rhs = AllOf(
      op::Copy(op::DynamicSlice(op::Parameter(), op::Constant(), op::Reshape(),
                                op::Constant(), op::Constant())),
      op::Shape("f32[32,14,28,64]"));

  auto left_halo = AllOf(op::CollectivePermute(op::Slice(rhs)),
                         op::Shape("f32[32,1,28,64]"));
  auto right_halo = AllOf(op::CollectivePermute(op::Slice(rhs)),
                          op::Shape("f32[32,1,28,64]"));
  EXPECT_THAT(root,
              AllOf(op::AllReduce(op::Convolution(
                        lhs, AllOf(op::Concatenate(left_halo, rhs, right_halo),
                                   op::Shape("f32[32,16,28,64]")))),
                    op::Shape("f32[3,3,128,64]")));
}

TEST_F(SpmdPartitioningTest, ConvolutionLhsTiledRhsTiledWindowDilate) {
  absl::string_view hlo_string = R"(
HloModule module

ENTRY entry {
  %lhs = f32[128,224,224,3] parameter(0)
  %lhs.copy = f32[128,224,224,3] copy(%lhs), sharding={devices=[1,2,1,1]0,1}
  %rhs = f32[128,112,112,64] parameter(1)
  %rhs.copy = f32[128,112,112,64] copy(%rhs), sharding={devices=[1,2,1,1]0,1}
  ROOT %conv = f32[7,7,3,64] convolution(%lhs.copy, %rhs.copy),
    window={size=112x112 pad=3_2x3_2 rhs_dilate=2x2}, dim_labels=f01b_i01o->01bf, sharding={replicated}
})";

  TF_ASSERT_OK_AND_ASSIGN(
      auto module,
      PartitionComputation(hlo_string, /*num_devices=*/2,
                           /*conv_halo_exchange_always_on_lhs=*/false));
  VLOG(1) << module->ToString();

  auto root = module->entry_computation()->root_instruction();
  auto lhs = AllOf(
      op::Copy(op::DynamicSlice(op::Parameter(), op::Constant(), op::Reshape(),
                                op::Constant(), op::Constant())),
      op::Shape("f32[128,112,224,3]"));
  auto rhs = AllOf(
      op::Copy(op::DynamicSlice(op::Parameter(), op::Constant(), op::Reshape(),
                                op::Constant(), op::Constant())),
      op::Shape("f32[128,56,112,64]"));

  auto left_halo = AllOf(op::CollectivePermute(op::Slice(rhs)),
                         op::Shape("f32[128,2,112,64]"));
  auto right_halo = AllOf(op::CollectivePermute(op::Slice(rhs)),
                          op::Shape("f32[128,2,112,64]"));
  EXPECT_THAT(root,
              AllOf(op::AllReduce(op::Convolution(
                        lhs, AllOf(op::Concatenate(left_halo, rhs, right_halo),
                                   op::Shape("f32[128,60,112,64]")))),
                    op::Shape("f32[7,7,3,64]")));
}

TEST_F(SpmdPartitioningTest,
       ConvolutionLhsTiledRhsTiledWindowDilateNegativeRhsPadding) {
  absl::string_view hlo_string = R"(
HloModule module

ENTRY entry {
  %lhs = f32[128,56,56,256] parameter(0)
  %lhs.copy = f32[128,56,56,256] copy(%lhs), sharding={devices=[1,2,1,1]0,1}
  %rhs = f32[128,28,28,512] parameter(1)
  %rhs.copy = f32[128,28,28,512] copy(%rhs), sharding={devices=[1,2,1,1]0,1}
  ROOT %conv = f32[1,1,256,512] convolution(%lhs.copy, %rhs.copy),
    window={size=28x28 pad=0_-1x0_-1 rhs_dilate=2x2}, dim_labels=f01b_i01o->01bf, sharding={replicated}
})";

  TF_ASSERT_OK_AND_ASSIGN(
      auto module,
      PartitionComputation(hlo_string, /*num_devices=*/2,
                           /*conv_halo_exchange_always_on_lhs=*/false));
  VLOG(1) << module->ToString();

  auto root = module->entry_computation()->root_instruction();
  auto lhs = AllOf(
      op::Copy(op::DynamicSlice(op::Parameter(), op::Constant(), op::Reshape(),
                                op::Constant(), op::Constant())),
      op::Shape("f32[128,28,56,256]"));
  auto rhs = AllOf(
      op::Copy(op::DynamicSlice(op::Parameter(), op::Constant(), op::Reshape(),
                                op::Constant(), op::Constant())),
      op::Shape("f32[128,14,28,512]"));

  EXPECT_THAT(root, AllOf(op::AllReduce(op::Convolution(lhs, rhs)),
                          op::Shape("f32[1,1,256,512]")));
}

TEST_F(SpmdPartitioningTest, ConvolutionLhsTiledRhsTiledWindowDilateUneven) {
  absl::string_view hlo_string = R"(
HloModule module

ENTRY entry {
  %lhs = f32[128,14,14,512] parameter(0)
  %lhs.copy = f32[128,14,14,512] copy(%lhs), sharding={devices=[1,2,1,1]0,1}
  %rhs = f32[128,7,7,512] parameter(1)
  %rhs.copy = f32[128,7,7,512] copy(%rhs), sharding={devices=[1,2,1,1]0,1}
  ROOT %conv = f32[3,3,512,512] convolution(%lhs.copy, %rhs.copy),
    window={size=7x7 pad=1_0x1_0 rhs_dilate=2x2}, dim_labels=f01b_i01o->01bf, sharding={replicated}
})";

  TF_ASSERT_OK_AND_ASSIGN(
      auto module,
      PartitionComputation(hlo_string, /*num_devices=*/2,
                           /*conv_halo_exchange_always_on_lhs=*/false));
  VLOG(1) << module->ToString();

  auto root = module->entry_computation()->root_instruction();
  auto lhs = AllOf(
      op::Copy(op::DynamicSlice(op::Parameter(), op::Constant(), op::Reshape(),
                                op::Constant(), op::Constant())),
      op::Shape("f32[128,7,14,512]"));
  auto rhs = AllOf(
      op::Select(op::Compare(),
                 op::Copy(op::DynamicSlice(
                     op::Pad(op::Parameter(), op::Constant()), op::Constant(),
                     op::Reshape(), op::Constant(), op::Constant())),
                 op::Broadcast()),
      op::Shape("f32[128,4,7,512]"));

  auto left_halo = AllOf(op::CollectivePermute(op::Slice(rhs)),
                         op::Shape("f32[128,1,7,512]"));
  EXPECT_THAT(root,
              AllOf(op::AllReduce(op::Convolution(
                        AllOf(op::DynamicSlice(op::Pad(lhs, op::Constant()),
                                               op::Constant(), op::Subtract(),
                                               op::Constant(), op::Constant()),
                              op::Shape("f32[128,10,14,512]")),
                        AllOf(op::Concatenate(left_halo, rhs),
                              op::Shape("f32[128,5,7,512]")))),
                    op::Shape("f32[3,3,512,512]")));
}

TEST_F(SpmdPartitioningTest, ConvolutionLhsTiledRhsTiledWithPadding_HaloOnLhs) {
  absl::string_view hlo_string = R"(
HloModule module

ENTRY entry {
  %lhs = f32[32,28,28,128] parameter(0)
  %lhs.copy = f32[32,28,28,128] copy(%lhs), sharding={devices=[1,2,1,1]0,1}
  %rhs = f32[32,28,28,64] parameter(1)
  %rhs.copy = f32[32,28,28,64] copy(%rhs), sharding={devices=[1,2,1,1]0,1}
  ROOT %conv = f32[3,3,128,64] convolution(%lhs.copy, %rhs.copy),
    window={size=28x28 pad=1_1x1_1}, dim_labels=f01b_i01o->01bf, sharding={replicated}
})";

  TF_ASSERT_OK_AND_ASSIGN(auto module,
                          PartitionComputation(hlo_string, /*num_devices=*/2));
  VLOG(1) << module->ToString();

  auto root = module->entry_computation()->root_instruction();
  auto lhs = AllOf(
      op::Copy(op::DynamicSlice(op::Parameter(), op::Constant(), op::Reshape(),
                                op::Constant(), op::Constant())),
      op::Shape("f32[32,14,28,128]"));
  auto rhs = AllOf(
      op::Copy(op::DynamicSlice(op::Parameter(), op::Constant(), op::Reshape(),
                                op::Constant(), op::Constant())),
      op::Shape("f32[32,14,28,64]"));

  auto left_halo = AllOf(op::CollectivePermute(op::Slice(lhs)),
                         op::Shape("f32[32,1,28,128]"));
  auto right_halo = AllOf(op::CollectivePermute(op::Slice(lhs)),
                          op::Shape("f32[32,1,28,128]"));
  EXPECT_THAT(root, AllOf(op::AllReduce(op::Convolution(
                              AllOf(op::Concatenate(left_halo, lhs, right_halo),
                                    op::Shape("f32[32,16,28,128]")),
                              rhs)),
                          op::Shape("f32[3,3,128,64]")));
}

TEST_F(SpmdPartitioningTest,
       ConvolutionLhsTiledRhsTiledWindowDilate_HaloOnLhs) {
  absl::string_view hlo_string = R"(
HloModule module

ENTRY entry {
  %lhs = f32[128,224,224,3] parameter(0)
  %lhs.copy = f32[128,224,224,3] copy(%lhs), sharding={devices=[1,2,1,1]0,1}
  %rhs = f32[128,112,112,64] parameter(1)
  %rhs.copy = f32[128,112,112,64] copy(%rhs), sharding={devices=[1,2,1,1]0,1}
  ROOT %conv = f32[7,7,3,64] convolution(%lhs.copy, %rhs.copy),
    window={size=112x112 pad=3_2x3_2 rhs_dilate=2x2}, dim_labels=f01b_i01o->01bf, sharding={replicated}
})";

  TF_ASSERT_OK_AND_ASSIGN(auto module,
                          PartitionComputation(hlo_string, /*num_devices=*/2));
  VLOG(1) << module->ToString();

  auto root = module->entry_computation()->root_instruction();
  auto lhs = AllOf(
      op::Copy(op::DynamicSlice(op::Parameter(), op::Constant(), op::Reshape(),
                                op::Constant(), op::Constant())),
      op::Shape("f32[128,112,224,3]"));
  auto rhs = AllOf(
      op::Copy(op::DynamicSlice(op::Parameter(), op::Constant(), op::Reshape(),
                                op::Constant(), op::Constant())),
      op::Shape("f32[128,56,112,64]"));

  auto left_halo = AllOf(op::CollectivePermute(op::Slice(lhs)),
                         op::Shape("f32[128,3,224,3]"));
  auto right_halo = AllOf(op::CollectivePermute(op::Slice(lhs)),
                          op::Shape("f32[128,2,224,3]"));
  EXPECT_THAT(root, AllOf(op::AllReduce(op::Convolution(
                              AllOf(op::Concatenate(left_halo, lhs, right_halo),
                                    op::Shape("f32[128,117,224,3]")),
                              rhs)),
                          op::Shape("f32[7,7,3,64]")));
}

TEST_F(SpmdPartitioningTest,
       ConvolutionLhsTiledRhsTiledWindowDilateNegativeRhsPadding_HaloOnLhs) {
  absl::string_view hlo_string = R"(
HloModule module

ENTRY entry {
  %lhs = f32[128,56,56,256] parameter(0)
  %lhs.copy = f32[128,56,56,256] copy(%lhs), sharding={devices=[1,2,1,1]0,1}
  %rhs = f32[128,28,28,512] parameter(1)
  %rhs.copy = f32[128,28,28,512] copy(%rhs), sharding={devices=[1,2,1,1]0,1}
  ROOT %conv = f32[1,1,256,512] convolution(%lhs.copy, %rhs.copy),
    window={size=28x28 pad=0_-1x0_-1 rhs_dilate=2x2}, dim_labels=f01b_i01o->01bf, sharding={replicated}
})";

  TF_ASSERT_OK_AND_ASSIGN(auto module,
                          PartitionComputation(hlo_string, /*num_devices=*/2));
  VLOG(1) << module->ToString();

  auto root = module->entry_computation()->root_instruction();
  auto lhs = AllOf(
      op::Copy(op::DynamicSlice(op::Parameter(), op::Constant(), op::Reshape(),
                                op::Constant(), op::Constant())),
      op::Shape("f32[128,28,56,256]"));
  auto rhs = AllOf(
      op::Copy(op::DynamicSlice(op::Parameter(), op::Constant(), op::Reshape(),
                                op::Constant(), op::Constant())),
      op::Shape("f32[128,14,28,512]"));

  EXPECT_THAT(root, AllOf(op::AllReduce(op::Convolution(op::Slice(lhs), rhs)),
                          op::Shape("f32[1,1,256,512]")));
}

TEST_F(SpmdPartitioningTest,
       ConvolutionLhsTiledRhsTiledWindowDilateUneven_HaloOnLhs) {
  absl::string_view hlo_string = R"(
HloModule module

ENTRY entry {
  %lhs = f32[128,14,14,512] parameter(0)
  %lhs.copy = f32[128,14,14,512] copy(%lhs), sharding={devices=[1,2,1,1]0,1}
  %rhs = f32[128,7,7,512] parameter(1)
  %rhs.copy = f32[128,7,7,512] copy(%rhs), sharding={devices=[1,2,1,1]0,1}
  ROOT %conv = f32[3,3,512,512] convolution(%lhs.copy, %rhs.copy),
    window={size=7x7 pad=1_0x1_0 rhs_dilate=2x2}, dim_labels=f01b_i01o->01bf, sharding={replicated}
})";

  TF_ASSERT_OK_AND_ASSIGN(auto module,
                          PartitionComputation(hlo_string, /*num_devices=*/2));
  VLOG(1) << module->ToString();

  auto root = module->entry_computation()->root_instruction();
  auto lhs = AllOf(
      op::Copy(op::DynamicSlice(op::Parameter(), op::Constant(), op::Reshape(),
                                op::Constant(), op::Constant())),
      op::Shape("f32[128,7,14,512]"));
  auto rhs = AllOf(
      op::Select(op::Compare(),
                 op::Copy(op::DynamicSlice(
                     op::Pad(op::Parameter(), op::Constant()), op::Constant(),
                     op::Reshape(), op::Constant(), op::Constant())),
                 op::Broadcast()),
      op::Shape("f32[128,4,7,512]"));

  auto right_halo = AllOf(op::CollectivePermute(op::Slice(lhs)),
                          op::Shape("f32[128,1,14,512]"));
  EXPECT_THAT(
      root, AllOf(op::AllReduce(op::Convolution(
                      AllOf(op::DynamicSlice(
                                AllOf(op::Pad(op::Concatenate(lhs, right_halo),
                                              op::Constant()),
                                      op::Shape("f32[128,10,14,512]")),
                                op::Constant(), op::Reshape(), op::Constant(),
                                op::Constant()),
                            op::Shape("f32[128,9,14,512]")),
                      rhs)),
                  op::Shape("f32[3,3,512,512]")));
}

TEST_F(SpmdPartitioningTest, ConcatenateAlongNonPartitionedDimension) {
  absl::string_view hlo_string = R"(
HloModule module

ENTRY entry {
  %param0 = f32[14,257] parameter(0)
  %param0.copy = f32[14,257] copy(%param0), sharding={devices=[2,1]0,1}
  %param1 = f32[14,116] parameter(1)
  %param1.copy = f32[14,116] copy(%param1), sharding={devices=[2,1]0,1}
  ROOT %concatenate = f32[14,373] concatenate(%param0.copy, %param1.copy),
    dimensions={1}, sharding={devices=[2,1]0,1}
})";

  TF_ASSERT_OK_AND_ASSIGN(auto module,
                          PartitionComputation(hlo_string, /*num_devices=*/2));
  VLOG(1) << module->ToString();

  auto root = module->entry_computation()->root_instruction();
  auto param0 = AllOf(op::Copy(op::DynamicSlice(op::Parameter(), op::Reshape(),
                                                op::Constant())),
                      op::Shape("f32[7,257]"));
  auto param1 = AllOf(op::Copy(op::DynamicSlice(op::Parameter(), op::Reshape(),
                                                op::Constant())),
                      op::Shape("f32[7,116]"));
  EXPECT_THAT(root,
              AllOf(op::Concatenate(param0, param1), op::Shape("f32[7,373]")));
}

TEST_F(SpmdPartitioningTest, ConcatenateAlongPartitionedDimension) {
  absl::string_view hlo_string = R"(
HloModule module

ENTRY entry {
  %param0 = f32[14,257] parameter(0)
  %param0.copy = f32[14,257] copy(%param0), sharding={devices=[1,2]0,1}
  %param1 = f32[14,116] parameter(1)
  %param1.copy = f32[14,116] copy(%param1), sharding={devices=[1,2]0,1}
  ROOT %concatenate = f32[14,373] concatenate(%param0.copy, %param1.copy),
    dimensions={1}, sharding={devices=[1,2]0,1}
})";

  TF_ASSERT_OK_AND_ASSIGN(auto module,
                          PartitionComputation(hlo_string, /*num_devices=*/2));
  VLOG(1) << module->ToString();

  auto root = module->entry_computation()->root_instruction();
  auto param0 =
      AllOf(op::Copy(op::DynamicSlice(op::Pad(op::Parameter(), op::Constant()),
                                      op::Constant(), op::Reshape())),
            op::Shape("f32[14,129]"));
  auto param1 = AllOf(op::Copy(op::DynamicSlice(op::Parameter(), op::Constant(),
                                                op::Reshape())),
                      op::Shape("f32[14,58]"));
  EXPECT_THAT(root, AllOf(op::DynamicSlice(
                              AllOf(op::AllReduce(op::DynamicUpdateSlice(
                                        op::DynamicUpdateSlice(
                                            op::Broadcast(), param0,
                                            op::Constant(), op::Multiply()),
                                        param1, op::Constant(), op::Add())),
                                    op::Shape("f32[14,374]")),
                              op::Constant(), op::Multiply()),
                          op::Shape("f32[14,187]")));
}

TEST_F(SpmdPartitioningTest, ConcatenateAlongBothDimensions) {
  const char* const hlo_string = R"(
HloModule module

ENTRY entry {
  %param0 = f32[14,257] parameter(0), sharding={devices=[2,2]0,1,2,3}
  %param1 = f32[14,116] parameter(1), sharding={devices=[2,2]0,1,2,3}
  ROOT %concatenate = f32[14,373] concatenate(%param0, %param1),
    dimensions={1}, sharding={devices=[2,2]0,1,2,3}
})";

  TF_ASSERT_OK_AND_ASSIGN(auto module,
                          PartitionComputation(hlo_string, /*num_devices=*/4));
  VLOG(1) << module->ToString();

  auto root = module->entry_computation()->root_instruction();
  auto param0 = AllOf(op::Parameter(0), op::Shape("f32[7,129]"));
  auto param1 = AllOf(op::Parameter(1), op::Shape("f32[7,58]"));
  EXPECT_THAT(root, AllOf(op::DynamicSlice(
                              AllOf(op::AllReduce(op::DynamicUpdateSlice(
                                        op::DynamicUpdateSlice(
                                            op::Broadcast(), param0,
                                            op::Constant(), op::Multiply()),
                                        param1, op::Constant(), op::Add())),
                                    op::Shape("f32[7,374]")),
                              op::Constant(), op::Multiply()),
                          op::Shape("f32[7,187]")));
}

TEST_F(SpmdPartitioningTest, PadAlongNonPartitionedDimension) {
  absl::string_view hlo_string = R"(
HloModule module

ENTRY entry {
  %param0 = f32[128,14,257] parameter(0), sharding={devices=[1,1,2]0,1}
  %const = f32[] constant(0)
  ROOT %pad = f32[128,17,257] pad(%param0, %const), padding=0_0x1_2x0_0,
    sharding={devices=[1,1,2]0,1}
})";

  TF_ASSERT_OK_AND_ASSIGN(auto module,
                          PartitionComputation(hlo_string, /*num_devices=*/2));
  VLOG(1) << module->ToString();

  auto root = module->entry_computation()->root_instruction();
  auto param0 = AllOf(op::Parameter(), op::Shape("f32[128,14,129]"));
  EXPECT_THAT(root, AllOf(op::Pad(param0, op::Constant()),
                          op::Shape("f32[128,17,129]")));
}

TEST_F(SpmdPartitioningTest, PadAlongPartitionedDimension) {
  absl::string_view hlo_string = R"(
HloModule module

ENTRY entry {
  %param0 = f32[14,257] parameter(0), sharding={devices=[1,2]0,1}
  %const = f32[] constant(0)
  ROOT %pad = f32[14,259] pad(%param0, %const), padding=0_0x0_2,
    sharding={devices=[1,2]0,1}
})";

  TF_ASSERT_OK_AND_ASSIGN(auto module,
                          PartitionComputation(hlo_string, /*num_devices=*/2));
  VLOG(1) << module->ToString();

  auto root = module->entry_computation()->root_instruction();
  auto param0 = AllOf(op::Parameter(), op::Shape("f32[14,129]"));
  auto after_halo_exchange =
      AllOf(op::Shape("f32[14,130]"),
            op::Concatenate(param0, op::CollectivePermute(op::Slice(param0))));
  auto pad = AllOf(op::Shape("f32[14,131]"),
                   op::Pad(after_halo_exchange, op::Constant()));
  EXPECT_THAT(root, op::DynamicSlice(pad, op::Constant(), _));
}

TEST_F(SpmdPartitioningTest, PadAlongPartitionedDimensionWithInteriorPadding) {
  absl::string_view hlo_string = R"(
HloModule module

ENTRY entry {
  %param0 = f32[7] parameter(0), sharding={devices=[2]0,1}
  %param1 = f32[] parameter(1), sharding={replicated}
  ROOT %pad = f32[22] pad(%param0, %param1), padding=2_1_2,
    sharding={devices=[2]0,1}
})";

  TF_ASSERT_OK_AND_ASSIGN(auto module,
                          PartitionComputation(hlo_string, /*num_devices=*/2));
  VLOG(1) << module->ToString();
  auto root = module->entry_computation()->root_instruction();

  auto param0 = AllOf(op::Parameter(), op::Shape("f32[4]"));
  auto after_halo_exchange =
      AllOf(op::Shape("f32[4]"),
            op::DynamicSlice(
                AllOf(op::Shape("f32[5]"),
                      op::Concatenate(op::CollectivePermute(op::Slice(param0)),
                                      param0)),
                _));
  auto pad = op::Pad(after_halo_exchange, op::Parameter(1));
  EXPECT_THAT(root, op::DynamicSlice(pad, _));
}

TEST_F(SpmdPartitioningTest, PartialReplicatePad) {
  absl::string_view hlo_string = R"(
HloModule module

ENTRY entry {
  %param0 = f32[11,7] parameter(0),
    sharding={devices=[1,2,2]0,1,2,3 last_tile_dim_replicate}
  %param1 = f32[] parameter(1), sharding={replicated}
  ROOT %pad = f32[27,22] pad(%param0, %param1), padding=2_4_1x2_1_2,
    sharding={devices=[1,2,2]0,1,2,3 last_tile_dim_replicate}
})";

  TF_ASSERT_OK_AND_ASSIGN(auto module,
                          PartitionComputation(hlo_string, /*num_devices=*/4));
  VLOG(1) << module->ToString();
  auto root = module->entry_computation()->root_instruction();

  auto param0 = AllOf(op::Parameter(), op::Shape("f32[11,4]"));
  auto after_halo_exchange =
      AllOf(op::Shape("f32[11,4]"),
            op::DynamicSlice(
                AllOf(op::Shape("f32[11,5]"),
                      op::Concatenate(op::CollectivePermute(op::Slice(param0)),
                                      param0)),
                op::Constant(), _));
  auto pad = op::Pad(after_halo_exchange, op::Parameter(1));
  EXPECT_THAT(root, AllOf(op::DynamicSlice(pad, op::Constant(), _),
                          op::Shape("f32[27,11]")));
}

TEST_F(SpmdPartitioningTest, SliceAlongNonPartitionedDimension) {
  absl::string_view hlo_string = R"(
HloModule module

ENTRY entry {
  %param0 = f32[128,14,257] parameter(0)
  %param0.copy = f32[128,14,257] copy(%param0), sharding={devices=[1,1,2]0,1}
  ROOT %slice = f32[128,11,257] slice(%param0.copy),
    slice={[0:128:1], [2:13:1], [0:257:1]}, sharding={devices=[1,1,2]0,1}
})";

  TF_ASSERT_OK_AND_ASSIGN(auto module,
                          PartitionComputation(hlo_string, /*num_devices=*/2));
  VLOG(1) << module->ToString();

  auto root = module->entry_computation()->root_instruction();
  auto param0 = AllOf(
      op::Copy(op::DynamicSlice(op::Pad(op::Parameter(), op::Constant()),
                                op::Constant(), op::Constant(), op::Reshape())),
      op::Shape("f32[128,14,129]"));
  EXPECT_THAT(root, AllOf(op::Slice(param0), op::Shape("f32[128,11,129]")));
}

TEST_F(SpmdPartitioningTest, SliceAlongPartitionedDimension) {
  absl::string_view hlo_string = R"(
HloModule module

ENTRY entry {
  %param0 = f32[128,14,257] parameter(0)
  %param0.copy = f32[128,14,257] copy(%param0), sharding={devices=[1,1,2]0,1}
  ROOT %slice = f32[63,14,251] slice(%param0.copy),
    slice={[2:128:2], [0:14:1], [5:256:1]}, sharding={devices=[1,1,2]0,1}
})";

  TF_ASSERT_OK_AND_ASSIGN(auto module,
                          PartitionComputation(hlo_string, /*num_devices=*/2));
  VLOG(1) << module->ToString();

  auto root = module->entry_computation()->root_instruction();
  auto param0 = AllOf(
      op::Copy(op::DynamicSlice(op::Pad(op::Parameter(), op::Constant()),
                                op::Constant(), op::Constant(), op::Reshape())),
      op::Shape("f32[128,14,129]"));
  EXPECT_THAT(
      root,
      AllOf(op::Slice(AllOf(
                op::DynamicSlice(
                    AllOf(op::Concatenate(
                              param0,
                              AllOf(op::CollectivePermute(op::Slice(param0)),
                                    op::Shape("f32[128,14,2]"))),
                          op::Shape("f32[128,14,131]")),
                    op::Constant(), op::Constant(), op::Add()),
                op::Shape("f32[128,14,126]"))),
            op::Shape("f32[63,14,126]")));
}

TEST_F(SpmdPartitioningTest,
       PartialReplicateSliceAlongNonPartitionedDimension) {
  absl::string_view hlo_string = R"(
HloModule module

ENTRY entry {
  %param0 = f32[128,14,257] parameter(0), sharding={devices=[1,1,2,2]0,1,2,3 last_tile_dim_replicate}
  ROOT %slice = f32[128,11,257] slice(%param0),
    slice={[0:128:1], [2:13:1], [0:257:1]}, sharding={devices=[1,1,2,2]0,1,2,3 last_tile_dim_replicate}
})";

  TF_ASSERT_OK_AND_ASSIGN(auto module,
                          PartitionComputation(hlo_string, /*num_devices=*/4));
  VLOG(1) << module->ToString();

  auto root = module->entry_computation()->root_instruction();
  auto param0 = AllOf(op::Parameter(), op::Shape("f32[128,14,129]"));
  EXPECT_THAT(root, AllOf(op::Slice(param0), op::Shape("f32[128,11,129]")));
}

TEST_F(SpmdPartitioningTest, PartialReplicateSliceAlongPartitionedDimension) {
  absl::string_view hlo_string = R"(
HloModule module

ENTRY entry {
  %param0 = f32[128,14,257] parameter(0), sharding={devices=[1,1,2,2]0,1,2,3 last_tile_dim_replicate}
  ROOT %slice = f32[63,14,251] slice(%param0),
    slice={[2:128:2], [0:14:1], [5:256:1]}, sharding={devices=[1,1,2,2]0,1,2,3 last_tile_dim_replicate}
})";

  TF_ASSERT_OK_AND_ASSIGN(auto module,
                          PartitionComputation(hlo_string, /*num_devices=*/4));
  VLOG(1) << module->ToString();

  auto root = module->entry_computation()->root_instruction();
  auto param0 = AllOf(op::Parameter(), op::Shape("f32[128,14,129]"));
  EXPECT_THAT(
      root,
      AllOf(
          op::Slice(AllOf(
              op::DynamicSlice(
                  AllOf(op::Concatenate(
                            param0,
                            AllOf(op::CollectivePermute(op::Slice(param0)),
                                  op::Shape("f32[128,14,2]"))),
                        op::Shape("f32[128,14,131]")),
                  op::Constant(), op::Constant(),
                  op::Add(op::Multiply(op::Reshape(op::DynamicSlice(
                                           op::Constant(), op::PartitionId())),
                                       op::Constant()),
                          op::Constant())),
              op::Shape("f32[128,14,126]"))),
          op::Shape("f32[63,14,126]")));
}

TEST_F(SpmdPartitioningTest, SortAlongNonPartitionedDimension) {
  absl::string_view hlo_string = R"(
HloModule module

ge {
  p.0.lhs.1247 = f32[]{:T(256)} parameter(0), sharding={replicated}
  bitcast-convert = s32[]{:T(256)} bitcast-convert(p.0.lhs.1247), sharding={replicated}
  constant = s32[]{:T(256)} constant(0), sharding={replicated}
  compare = pred[]{:T(256)E(32)} compare(bitcast-convert, constant), direction=LT, sharding={replicated}
  constant.1 = u32[]{:T(256)} constant(2147483647), sharding={replicated}
  bitcast-convert.1 = u32[]{:T(256)} bitcast-convert(p.0.lhs.1247), sharding={replicated}
  subtract = u32[]{:T(256)} subtract(constant.1, bitcast-convert.1), sharding={replicated}
  bitcast-convert.2 = s32[]{:T(256)} bitcast-convert(subtract), sharding={replicated}
  select = s32[]{:T(256)} select(compare, bitcast-convert.2, bitcast-convert), sharding={replicated}
  p.0.rhs.1248 = f32[]{:T(256)} parameter(1), sharding={replicated}
  bitcast-convert.3 = s32[]{:T(256)} bitcast-convert(p.0.rhs.1248), sharding={replicated}
  compare.1 = pred[]{:T(256)E(32)} compare(bitcast-convert.3, constant), direction=LT, sharding={replicated}
  bitcast-convert.4 = u32[]{:T(256)} bitcast-convert(p.0.rhs.1248), sharding={replicated}
  subtract.1 = u32[]{:T(256)} subtract(constant.1, bitcast-convert.4), sharding={replicated}
  bitcast-convert.5 = s32[]{:T(256)} bitcast-convert(subtract.1), sharding={replicated}
  select.1 = s32[]{:T(256)} select(compare.1, bitcast-convert.5, bitcast-convert.3), sharding={replicated}
  compare.2 = pred[]{:T(256)E(32)} compare(select, select.1), direction=GT, sharding={replicated}
  compare.258 = pred[]{:T(256)E(32)} compare(select.1, select), direction=GT, sharding={replicated}
  compare.259 = pred[]{:T(256)E(32)} compare(compare.2, compare.258), direction=EQ, sharding={replicated}
  p.1.lhs.1249 = s32[]{:T(256)} parameter(2), sharding={replicated}
  p.1.rhs.1250 = s32[]{:T(256)} parameter(3), sharding={replicated}
  compare.260 = pred[]{:T(256)E(32)} compare(p.1.lhs.1249, p.1.rhs.1250), direction=LT, sharding={replicated}
  ROOT select.86 = pred[]{:T(256)E(32)} select(compare.259, compare.260, compare.2), sharding={replicated}
}

ENTRY entry {
  %param0 = f32[128,14,257] parameter(0)
  %param0.copy = f32[128,14,257] copy(%param0), sharding={devices=[1,2,1]0,1}
  %param1 = s32[128,14,257] parameter(1)
  %param1.copy = s32[128,14,257] copy(%param1), sharding={devices=[1,2,1]0,1}
  ROOT %sort.6 = (f32[128,14,257]{2,1,0:T(8,128)}, s32[128,14,257]{2,1,0:T(8,128)})
    sort(%param0.copy, %param1.copy), dimensions={2}, is_stable=true,
    to_apply=%ge, sharding={{devices=[1,2,1]0,1},{devices=[1,2,1]0,1}}
})";

  TF_ASSERT_OK_AND_ASSIGN(auto module,
                          PartitionComputation(hlo_string, /*num_devices=*/2));
  VLOG(1) << module->ToString();

  auto root = module->entry_computation()->root_instruction();
  auto param0 =
      AllOf(op::Copy(op::DynamicSlice(op::Parameter(0), op::Constant(),
                                      op::Reshape(), op::Constant())),
            op::Shape("f32[128,7,257]"));
  auto param1 =
      AllOf(op::Copy(op::DynamicSlice(op::Parameter(1), op::Constant(),
                                      op::Reshape(), op::Constant())),
            op::Shape("s32[128,7,257]"));
  EXPECT_THAT(root, AllOf(op::Sort(param0, param1),
                          op::Shape("(f32[128,7,257], s32[128,7,257])")));
}

TEST_F(SpmdPartitioningTest, PartitionCustomCall) {
  absl::string_view hlo_string = R"(
HloModule cluster_2013453984438090939__.47

ENTRY %cluster_2013453984438090939__.47
  (arg_tuple.1: ()) -> (bf16[2,2000], s32[2,2000]) {
  %arg_tuple.1 = bf16[2,209664] parameter(0)
  %copy.arg_tuple.1 = bf16[2,209664] copy(%arg_tuple.1), sharding={devices=[1,2]0,1}
  %custom-call = (bf16[2,2000]{1,0}, s32[2,2000]{1,0})
    custom-call(bf16[2,209664]{1,0} %copy.arg_tuple.1), custom_call_target="TopK"
  %get-tuple-element = bf16[2,2000]{1,0}
    get-tuple-element((bf16[2,2000]{1,0}, s32[2,2000]{1,0}) %custom-call),
    index=0, sharding={replicated}
  %get-tuple-element.1 = s32[2,2000]{1,0} get-tuple-element((bf16[2,2000]{1,0},
    s32[2,2000]{1,0}) %custom-call), index=1, sharding={replicated}
  ROOT %tuple.46 = (bf16[2,2000]{1,0}, s32[2,2000]{1,0})
    tuple(bf16[2,2000]{1,0} %get-tuple-element, s32[2,2000]{1,0}
    %get-tuple-element.1), sharding={{replicated}, {replicated}},
    metadata={op_name="XLA_Retvals"}
})";

  TF_ASSERT_OK_AND_ASSIGN(auto module,
                          PartitionComputation(hlo_string, /*num_devices=*/2));
  VLOG(1) << module->ToString();
  auto custom_call = FindInstruction(module.get(), "custom-call.1");
  EXPECT_EQ(custom_call->operand(0)->shape().dimensions(1), 104832);
  auto sort = FindInstruction(module.get(), "sort");
  EXPECT_EQ(sort->operand(0)->shape().dimensions(1), 4000);
  EXPECT_EQ(sort->operand(1)->shape().dimensions(1), 4000);
}

TEST_F(SpmdPartitioningTest, PartitionCustomCall_TwoPartitionedDims) {
  absl::string_view hlo_string = R"(
HloModule module

ENTRY entry {
  %param0 = f32[8,32128] parameter(0)
  %copy.0 = f32[8,32128] copy(%param0),
    sharding={devices=[4,2]0,1,2,3,4,5,6,7}
  %custom-call = (f32[8,2]{1,0}, s32[8,2]{1,0})
    custom-call(%copy.0), custom_call_target="TopK"
  %get-tuple-element = f32[8,2]{1,0}
    get-tuple-element((f32[8,2]{1,0}, s32[8,2]{1,0}) %custom-call), index=0,
    sharding={devices=[4,1,2]0,1,2,3,4,5,6,7 last_tile_dim_replicate}
  %get-tuple-element.1 = s32[8,2]{1,0}
    get-tuple-element((f32[8,2]{1,0}, s32[8,2]{1,0}) %custom-call), index=1,
    sharding={devices=[4,1,2]0,1,2,3,4,5,6,7 last_tile_dim_replicate}
  ROOT %tuple = (f32[8,2]{1,0}, s32[8,2]{1,0})
    tuple(%get-tuple-element, %get-tuple-element.1),
    sharding={{replicated}, {replicated}}
})";

  TF_ASSERT_OK_AND_ASSIGN(auto module,
                          PartitionComputation(hlo_string, /*num_devices=*/8));
  VLOG(1) << module->ToString();
  auto custom_call = FindInstruction(module.get(), "custom-call.1");
  EXPECT_EQ(custom_call->operand(0)->shape().dimensions(1), 16064);
  auto sort = FindInstruction(module.get(), "sort");
  EXPECT_EQ(sort->operand(0)->shape().dimensions(0), 2);
  EXPECT_EQ(sort->operand(0)->shape().dimensions(1), 4);
  EXPECT_EQ(sort->operand(1)->shape().dimensions(0), 2);
  EXPECT_EQ(sort->operand(1)->shape().dimensions(1), 4);
}

TEST_F(SpmdPartitioningTest, PartitionSortInTopK) {
  absl::string_view hlo_string = R"(
HloModule module

%compare-greater-than.8 (p.0.lhs.9: bf16[], p.0.rhs.10: bf16[], p.1.lhs.11:
   s32[], p.1.rhs.12: s32[]) -> pred[] {
  %p.1.lhs.11 = s32[] parameter(2)
  %p.1.rhs.12 = s32[] parameter(3)
  %p.0.lhs.9 = bf16[] parameter(0)
  %convert.13 = f32[] convert(bf16[] %p.0.lhs.9)
  %bitcast-convert.16 = s32[] bitcast-convert(f32[] %convert.13)
  %constant.20 = s32[] constant(0)
  %compare.21 = pred[] compare(s32[] %bitcast-convert.16, s32[] %constant.20),
    direction=LT
  %constant.15 = u32[] constant(2147483647)
  %bitcast-convert.17 = u32[] bitcast-convert(f32[] %convert.13)
  %subtract.18 = u32[] subtract(u32[] %constant.15, u32[] %bitcast-convert.17)
  %bitcast-convert.19 = s32[] bitcast-convert(u32[] %subtract.18)
  %select.22 = s32[] select(pred[] %compare.21, s32[] %bitcast-convert.19, s32[]
    %bitcast-convert.16)
  %p.0.rhs.10 = bf16[] parameter(1)
  %convert.14 = f32[] convert(bf16[] %p.0.rhs.10)
  %bitcast-convert.24 = s32[] bitcast-convert(f32[] %convert.14)
  %constant.28 = s32[] constant(0)
  %compare.29 = pred[] compare(s32[] %bitcast-convert.24, s32[] %constant.28),
    direction=LT
  %constant.23 = u32[] constant(2147483647)
  %bitcast-convert.25 = u32[] bitcast-convert(f32[] %convert.14)
  %subtract.26 = u32[] subtract(u32[] %constant.23, u32[] %bitcast-convert.25)
  %bitcast-convert.27 = s32[] bitcast-convert(u32[] %subtract.26)
  %select.30 = s32[] select(pred[] %compare.29, s32[] %bitcast-convert.27, s32[]
    %bitcast-convert.24)
  ROOT %compare.31 = pred[] compare(s32[] %select.22, s32[] %select.30),
    direction=GT
}

ENTRY entry
  (arg_tuple.1: ()) -> (bf16[2,2000], s32[2,2000]) {
  %arg_tuple.1 = bf16[2,209664] parameter(0)
  %copy.arg_tuple.1 = bf16[2,209664] copy(%arg_tuple.1), sharding={devices=[1,2]0,1}
  %iota.7 = s32[2,209664] iota(), iota_dimension=1,
    metadata={op_type="TopKV2" op_name="TopKV2"}
  %sort.32 = (bf16[2,209664], s32[2,209664])
    sort(bf16[2,209664] %copy.arg_tuple.1, s32[2,209664] %iota.7),
    dimensions={1}, is_stable=true, to_apply=%compare-greater-than.8,
    metadata={op_type="TopKV2" op_name="TopKV2"}
  %get-tuple-element.33 = bf16[2,209664]
    get-tuple-element((bf16[2,209664], s32[2,209664]) %sort.32),
    index=0, metadata={op_type="TopKV2" op_name="TopKV2"}
  %slice.34 = bf16[2,2000] slice(bf16[2,209664]
    %get-tuple-element.33), slice={[0:2], [0:2000]},
    metadata={op_type="TopKV2" op_name="TopKV2"}
  %get-tuple-element.35 = s32[2,209664]
    get-tuple-element((bf16[2,209664], s32[2,209664]) %sort.32),
    index=1, metadata={op_type="TopKV2" op_name="TopKV2"}
  %slice.36 = s32[2,2000] slice(s32[2,209664]
    %get-tuple-element.35), slice={[0:2], [0:2000]},
    metadata={op_type="TopKV2" op_name="TopKV2"}
  ROOT %tuple.46 = (bf16[2,2000], s32[2,2000])
    tuple(bf16[2,2000] %slice.34, s32[2,2000]
    %slice.36), sharding={{replicated}, {replicated}},
    metadata={op_name="XLA_Retvals"}
})";

  TF_ASSERT_OK_AND_ASSIGN(auto module,
                          PartitionComputation(hlo_string, /*num_devices=*/2));
  VLOG(1) << module->ToString();
  auto sort = FindInstruction(module.get(), "sort");
  EXPECT_EQ(sort->operand(0)->shape().dimensions(1), 104832);
  EXPECT_EQ(sort->operand(1)->shape().dimensions(1), 104832);
  auto final_sort = FindInstruction(module.get(), "sort.1");
  EXPECT_EQ(final_sort->operand(0)->shape().dimensions(1), 4000);
  EXPECT_EQ(final_sort->operand(1)->shape().dimensions(1), 4000);
}

TEST_F(SpmdPartitioningTest, PartitionSortInTopKWhenComparisonWithSelect) {
  absl::string_view hlo_string = R"(
HloModule module

%compare-greater-than.8 (p.0.lhs.2566: bf16[],
  p.0.rhs.2567: bf16[], p.1.lhs.2586: s32[], p.1.rhs.2587: s32[]) -> pred[] {
  %p.0.lhs.2566 = bf16[] parameter(0)
  %convert.164 = f32[] convert(bf16[] %p.0.lhs.2566)
  %bitcast-convert.48 = s32[] bitcast-convert(f32[] %convert.164)
  %constant.285 = s32[] constant(0)
  %compare.84 = pred[] compare(s32[] %bitcast-convert.48, s32[] %constant.285),
    direction=LT
  %constant.286 = u32[] constant(2147483647)
  %bitcast-convert.49 = u32[] bitcast-convert(f32[] %convert.164)
  %subtract.84 = u32[] subtract(u32[] %constant.286, u32[] %bitcast-convert.49)
  %bitcast-convert.50 = s32[] bitcast-convert(u32[] %subtract.84)
  %select.40 = s32[] select(pred[] %compare.84, s32[] %bitcast-convert.50,
    s32[] %bitcast-convert.48)
  %p.0.rhs.2567 = bf16[] parameter(1)
  %convert.165 = f32[] convert(bf16[] %p.0.rhs.2567)
  %bitcast-convert.51 = s32[] bitcast-convert(f32[] %convert.165)
  %compare.85 = pred[] compare(s32[] %bitcast-convert.51, s32[] %constant.285),
    direction=LT
  %bitcast-convert.52 = u32[] bitcast-convert(f32[] %convert.165)
  %subtract.85 = u32[] subtract(u32[] %constant.286, u32[] %bitcast-convert.52)
  %bitcast-convert.53 = s32[] bitcast-convert(u32[] %subtract.85)
  %select.41 = s32[] select(pred[] %compare.85, s32[] %bitcast-convert.53,
    s32[] %bitcast-convert.51)
  %compare.86 = pred[] compare(s32[] %select.40, s32[] %select.41), direction=GT
  %compare.1645 = pred[] compare(s32[] %select.41, s32[] %select.40), direction=GT
  %compare.1646 = pred[] compare(pred[] %compare.86, pred[] %compare.1645),
    direction=EQ
  %p.1.lhs.2586 = s32[] parameter(2)
  %p.1.rhs.2587 = s32[] parameter(3)
  %compare.1647 = pred[] compare(s32[] %p.1.lhs.2586, s32[] %p.1.rhs.2587),
    direction=LT
  ROOT %select.1054 = pred[] select(pred[] %compare.1646, pred[] %compare.1647,
    pred[] %compare.86)
}

ENTRY entry
  (arg_tuple.1: ()) -> (bf16[2,2000], s32[2,2000]) {
  %arg_tuple.1 = bf16[2,209664] parameter(0)
  %copy.arg_tuple.1 = bf16[2,209664] copy(%arg_tuple.1), sharding={devices=[1,2]0,1}
  %iota.7 = s32[2,209664] iota(), iota_dimension=1,
    metadata={op_type="TopKV2" op_name="TopKV2"}
  %sort.32 = (bf16[2,209664], s32[2,209664])
    sort(bf16[2,209664] %copy.arg_tuple.1, s32[2,209664] %iota.7),
    dimensions={1}, is_stable=true, to_apply=%compare-greater-than.8,
    metadata={op_type="TopKV2" op_name="TopKV2"}
  %get-tuple-element.33 = bf16[2,209664]
    get-tuple-element((bf16[2,209664], s32[2,209664]) %sort.32),
    index=0, metadata={op_type="TopKV2" op_name="TopKV2"}
  %slice.34 = bf16[2,2000] slice(bf16[2,209664]
    %get-tuple-element.33), slice={[0:2], [0:2000]},
    metadata={op_type="TopKV2" op_name="TopKV2"}
  %get-tuple-element.35 = s32[2,209664]
    get-tuple-element((bf16[2,209664], s32[2,209664]) %sort.32),
    index=1, metadata={op_type="TopKV2" op_name="TopKV2"}
  %slice.36 = s32[2,2000] slice(s32[2,209664]
    %get-tuple-element.35), slice={[0:2], [0:2000]},
    metadata={op_type="TopKV2" op_name="TopKV2"}
  ROOT %tuple.46 = (bf16[2,2000], s32[2,2000])
    tuple(bf16[2,2000] %slice.34, s32[2,2000]
    %slice.36), sharding={{replicated}, {replicated}},
    metadata={op_name="XLA_Retvals"}
})";

  TF_ASSERT_OK_AND_ASSIGN(auto module,
                          PartitionComputation(hlo_string, /*num_devices=*/2));
  VLOG(1) << module->ToString();
  auto sort = FindInstruction(module.get(), "sort");
  EXPECT_EQ(sort->operand(0)->shape().dimensions(1), 104832);
  EXPECT_EQ(sort->operand(1)->shape().dimensions(1), 104832);
  auto final_sort = FindInstruction(module.get(), "sort.1");
  EXPECT_EQ(final_sort->operand(0)->shape().dimensions(1), 4000);
  EXPECT_EQ(final_sort->operand(1)->shape().dimensions(1), 4000);
}

TEST_F(SpmdPartitioningTest, NoPartitionSortInTopKWhenSecondOperandIsNotIota) {
  absl::string_view hlo_string = R"(
HloModule module

%compare-greater-than.8 (p.0.lhs.2566: bf16[],
  p.0.rhs.2567: bf16[], p.1.lhs.2586: s32[], p.1.rhs.2587: s32[]) -> pred[] {
  %p.0.lhs.2566 = bf16[] parameter(0)
  %convert.164 = f32[] convert(bf16[] %p.0.lhs.2566)
  %bitcast-convert.48 = s32[] bitcast-convert(f32[] %convert.164)
  %constant.285 = s32[] constant(0)
  %compare.84 = pred[] compare(s32[] %bitcast-convert.48, s32[] %constant.285),
    direction=LT
  %constant.286 = u32[] constant(2147483647)
  %bitcast-convert.49 = u32[] bitcast-convert(f32[] %convert.164)
  %subtract.84 = u32[] subtract(u32[] %constant.286, u32[] %bitcast-convert.49)
  %bitcast-convert.50 = s32[] bitcast-convert(u32[] %subtract.84)
  %select.40 = s32[] select(pred[] %compare.84, s32[] %bitcast-convert.50,
    s32[] %bitcast-convert.48)
  %p.0.rhs.2567 = bf16[] parameter(1)
  %convert.165 = f32[] convert(bf16[] %p.0.rhs.2567)
  %bitcast-convert.51 = s32[] bitcast-convert(f32[] %convert.165)
  %compare.85 = pred[] compare(s32[] %bitcast-convert.51, s32[] %constant.285),
    direction=LT
  %bitcast-convert.52 = u32[] bitcast-convert(f32[] %convert.165)
  %subtract.85 = u32[] subtract(u32[] %constant.286, u32[] %bitcast-convert.52)
  %bitcast-convert.53 = s32[] bitcast-convert(u32[] %subtract.85)
  %select.41 = s32[] select(pred[] %compare.85, s32[] %bitcast-convert.53,
    s32[] %bitcast-convert.51)
  %compare.86 = pred[] compare(s32[] %select.40, s32[] %select.41), direction=GT
  %compare.1645 = pred[] compare(s32[] %select.41, s32[] %select.40), direction=GT
  %compare.1646 = pred[] compare(pred[] %compare.86, pred[] %compare.1645),
    direction=EQ
  %p.1.lhs.2586 = s32[] parameter(2)
  %p.1.rhs.2587 = s32[] parameter(3)
  %compare.1647 = pred[] compare(s32[] %p.1.lhs.2586, s32[] %p.1.rhs.2587),
    direction=LT
  ROOT %select.1054 = pred[] select(pred[] %compare.1646, pred[] %compare.1647,
    pred[] %compare.86)
}

ENTRY entry {
  %arg_tuple.1 = bf16[2,209664] parameter(0)
  %arg_tuple.2 = s32[2,209664] parameter(1)
  %copy.arg_tuple.1 = bf16[2,209664] copy(%arg_tuple.1), sharding={devices=[1,2]0,1}
  %sort.32 = (bf16[2,209664], s32[2,209664])
    sort(bf16[2,209664] %copy.arg_tuple.1, s32[2,209664] %arg_tuple.2),
    dimensions={1}, is_stable=true, to_apply=%compare-greater-than.8,
    metadata={op_type="TopKV2" op_name="TopKV2"}
  %get-tuple-element.33 = bf16[2,209664]
    get-tuple-element((bf16[2,209664], s32[2,209664]) %sort.32),
    index=0, metadata={op_type="TopKV2" op_name="TopKV2"}
  %slice.34 = bf16[2,2000] slice(bf16[2,209664]
    %get-tuple-element.33), slice={[0:2], [0:2000]},
    metadata={op_type="TopKV2" op_name="TopKV2"}
  %get-tuple-element.35 = s32[2,209664]
    get-tuple-element((bf16[2,209664], s32[2,209664]) %sort.32),
    index=1, metadata={op_type="TopKV2" op_name="TopKV2"}
  %slice.36 = s32[2,2000] slice(s32[2,209664]
    %get-tuple-element.35), slice={[0:2], [0:2000]},
    metadata={op_type="TopKV2" op_name="TopKV2"}
  ROOT %tuple.46 = (bf16[2,2000], s32[2,2000])
    tuple(bf16[2,2000] %slice.34, s32[2,2000]
    %slice.36), sharding={{replicated}, {replicated}},
    metadata={op_name="XLA_Retvals"}
})";

  TF_ASSERT_OK_AND_ASSIGN(auto module,
                          PartitionComputation(hlo_string, /*num_devices=*/2));
  VLOG(1) << module->ToString();
  auto sort = FindInstruction(module.get(), "sort");
  EXPECT_EQ(sort->operand(0)->shape().dimensions(1), 209664);
  EXPECT_EQ(sort->operand(1)->shape().dimensions(1), 209664);
}

TEST_F(SpmdPartitioningTest, NoPartitionSortInTopKWhenNoPartitionInSortDim) {
  absl::string_view hlo_string = R"(
HloModule module

%compare-greater-than.8 (p.0.lhs.2566: bf16[],
  p.0.rhs.2567: bf16[], p.1.lhs.2586: s32[], p.1.rhs.2587: s32[]) -> pred[] {
  %p.0.lhs.2566 = bf16[] parameter(0)
  %convert.164 = f32[] convert(bf16[] %p.0.lhs.2566)
  %bitcast-convert.48 = s32[] bitcast-convert(f32[] %convert.164)
  %constant.285 = s32[] constant(0)
  %compare.84 = pred[] compare(s32[] %bitcast-convert.48, s32[] %constant.285),
    direction=LT
  %constant.286 = u32[] constant(2147483647)
  %bitcast-convert.49 = u32[] bitcast-convert(f32[] %convert.164)
  %subtract.84 = u32[] subtract(u32[] %constant.286, u32[] %bitcast-convert.49)
  %bitcast-convert.50 = s32[] bitcast-convert(u32[] %subtract.84)
  %select.40 = s32[] select(pred[] %compare.84, s32[] %bitcast-convert.50,
    s32[] %bitcast-convert.48)
  %p.0.rhs.2567 = bf16[] parameter(1)
  %convert.165 = f32[] convert(bf16[] %p.0.rhs.2567)
  %bitcast-convert.51 = s32[] bitcast-convert(f32[] %convert.165)
  %compare.85 = pred[] compare(s32[] %bitcast-convert.51, s32[] %constant.285),
    direction=LT
  %bitcast-convert.52 = u32[] bitcast-convert(f32[] %convert.165)
  %subtract.85 = u32[] subtract(u32[] %constant.286, u32[] %bitcast-convert.52)
  %bitcast-convert.53 = s32[] bitcast-convert(u32[] %subtract.85)
  %select.41 = s32[] select(pred[] %compare.85, s32[] %bitcast-convert.53,
    s32[] %bitcast-convert.51)
  %compare.86 = pred[] compare(s32[] %select.40, s32[] %select.41), direction=GT
  %compare.1645 = pred[] compare(s32[] %select.41, s32[] %select.40), direction=GT
  %compare.1646 = pred[] compare(pred[] %compare.86, pred[] %compare.1645),
    direction=EQ
  %p.1.lhs.2586 = s32[] parameter(2)
  %p.1.rhs.2587 = s32[] parameter(3)
  %compare.1647 = pred[] compare(s32[] %p.1.lhs.2586, s32[] %p.1.rhs.2587),
    direction=LT
  ROOT %select.1054 = pred[] select(pred[] %compare.1646, pred[] %compare.1647,
    pred[] %compare.86)
}

ENTRY entry
  (arg_tuple.1: ()) -> (bf16[2,2000], s32[2,2000]) {
  %arg_tuple.1 = bf16[2,209664] parameter(0)
  %copy.arg_tuple.1 = bf16[2,209664] copy(%arg_tuple.1), sharding={devices=[2,1]0,1}
  %iota.7 = s32[2,209664] iota(), iota_dimension=1,
    metadata={op_type="TopKV2" op_name="TopKV2"}
  %sort.32 = (bf16[2,209664], s32[2,209664])
    sort(bf16[2,209664] %copy.arg_tuple.1, s32[2,209664] %iota.7),
    dimensions={1}, is_stable=true, to_apply=%compare-greater-than.8,
    metadata={op_type="TopKV2" op_name="TopKV2"}
  %get-tuple-element.33 = bf16[2,209664]
    get-tuple-element((bf16[2,209664], s32[2,209664]) %sort.32),
    index=0, metadata={op_type="TopKV2" op_name="TopKV2"}
  %slice.34 = bf16[2,2000] slice(bf16[2,209664]
    %get-tuple-element.33), slice={[0:2], [0:2000]},
    metadata={op_type="TopKV2" op_name="TopKV2"}
  %get-tuple-element.35 = s32[2,209664]
    get-tuple-element((bf16[2,209664], s32[2,209664]) %sort.32),
    index=1, metadata={op_type="TopKV2" op_name="TopKV2"}
  %slice.36 = s32[2,2000] slice(s32[2,209664]
    %get-tuple-element.35), slice={[0:2], [0:2000]},
    metadata={op_type="TopKV2" op_name="TopKV2"}
  ROOT %tuple.46 = (bf16[2,2000], s32[2,2000])
    tuple(bf16[2,2000] %slice.34, s32[2,2000]
    %slice.36), sharding={{replicated}, {replicated}},
    metadata={op_name="XLA_Retvals"}
})";

  TF_ASSERT_OK_AND_ASSIGN(auto module,
                          PartitionComputation(hlo_string, /*num_devices=*/2));
  VLOG(1) << module->ToString();
  auto sort = FindInstruction(module.get(), "sort");
  EXPECT_EQ(sort->operand(0)->shape().dimensions(1), 209664);
  EXPECT_EQ(sort->operand(1)->shape().dimensions(1), 209664);
}

TEST_F(SpmdPartitioningTest, NoPartitionSortInTopKWhenSliceInOtherDim) {
  absl::string_view hlo_string = R"(
HloModule module

%compare-greater-than.8 (p.0.lhs.2566: bf16[],
  p.0.rhs.2567: bf16[], p.1.lhs.2586: s32[], p.1.rhs.2587: s32[]) -> pred[] {
  %p.0.lhs.2566 = bf16[] parameter(0)
  %convert.164 = f32[] convert(bf16[] %p.0.lhs.2566)
  %bitcast-convert.48 = s32[] bitcast-convert(f32[] %convert.164)
  %constant.285 = s32[] constant(0)
  %compare.84 = pred[] compare(s32[] %bitcast-convert.48, s32[] %constant.285),
    direction=LT
  %constant.286 = u32[] constant(2147483647)
  %bitcast-convert.49 = u32[] bitcast-convert(f32[] %convert.164)
  %subtract.84 = u32[] subtract(u32[] %constant.286, u32[] %bitcast-convert.49)
  %bitcast-convert.50 = s32[] bitcast-convert(u32[] %subtract.84)
  %select.40 = s32[] select(pred[] %compare.84, s32[] %bitcast-convert.50,
    s32[] %bitcast-convert.48)
  %p.0.rhs.2567 = bf16[] parameter(1)
  %convert.165 = f32[] convert(bf16[] %p.0.rhs.2567)
  %bitcast-convert.51 = s32[] bitcast-convert(f32[] %convert.165)
  %compare.85 = pred[] compare(s32[] %bitcast-convert.51, s32[] %constant.285),
    direction=LT
  %bitcast-convert.52 = u32[] bitcast-convert(f32[] %convert.165)
  %subtract.85 = u32[] subtract(u32[] %constant.286, u32[] %bitcast-convert.52)
  %bitcast-convert.53 = s32[] bitcast-convert(u32[] %subtract.85)
  %select.41 = s32[] select(pred[] %compare.85, s32[] %bitcast-convert.53,
    s32[] %bitcast-convert.51)
  %compare.86 = pred[] compare(s32[] %select.40, s32[] %select.41), direction=GT
  %compare.1645 = pred[] compare(s32[] %select.41, s32[] %select.40), direction=GT
  %compare.1646 = pred[] compare(pred[] %compare.86, pred[] %compare.1645),
    direction=EQ
  %p.1.lhs.2586 = s32[] parameter(2)
  %p.1.rhs.2587 = s32[] parameter(3)
  %compare.1647 = pred[] compare(s32[] %p.1.lhs.2586, s32[] %p.1.rhs.2587),
    direction=LT
  ROOT %select.1054 = pred[] select(pred[] %compare.1646, pred[] %compare.1647,
    pred[] %compare.86)
}

ENTRY entry {
  %arg_tuple.1 = bf16[2,209664] parameter(0)
  %copy.arg_tuple.1 = bf16[2,209664] copy(%arg_tuple.1), sharding={devices=[1,2]0,1}
  %iota.7 = s32[2,209664] iota(), iota_dimension=1,
    metadata={op_type="TopKV2" op_name="TopKV2"}
  %sort.32 = (bf16[2,209664], s32[2,209664])
    sort(bf16[2,209664] %copy.arg_tuple.1, s32[2,209664] %iota.7),
    dimensions={1}, is_stable=true, to_apply=%compare-greater-than.8,
    metadata={op_type="TopKV2" op_name="TopKV2"}
  %get-tuple-element.33 = bf16[2,209664]
    get-tuple-element((bf16[2,209664], s32[2,209664]) %sort.32),
    index=0, metadata={op_type="TopKV2" op_name="TopKV2"}
  %slice.34 = bf16[1,209664] slice(bf16[2,209664]
    %get-tuple-element.33), slice={[0:1], [0:209664]},
    metadata={op_type="TopKV2" op_name="TopKV2"}
  %get-tuple-element.35 = s32[2,209664]
    get-tuple-element((bf16[2,209664], s32[2,209664]) %sort.32),
    index=1, metadata={op_type="TopKV2" op_name="TopKV2"}
  %slice.36 = s32[1,209664] slice(s32[2,209664]
    %get-tuple-element.35), slice={[0:1], [0:209664]},
    metadata={op_type="TopKV2" op_name="TopKV2"}
  ROOT %tuple.46 = (bf16[1,209664], s32[1,209664])
    tuple(bf16[1,209664] %slice.34, s32[1,209664]
    %slice.36), sharding={{replicated}, {replicated}},
    metadata={op_name="XLA_Retvals"}
})";

  TF_ASSERT_OK_AND_ASSIGN(auto module,
                          PartitionComputation(hlo_string, /*num_devices=*/2));
  VLOG(1) << module->ToString();
  auto sort = FindInstruction(module.get(), "sort");
  EXPECT_EQ(sort->operand(0)->shape().dimensions(1), 209664);
  EXPECT_EQ(sort->operand(1)->shape().dimensions(1), 209664);
}

TEST_F(SpmdPartitioningTest, ShardableTranspose) {
  absl::string_view hlo_string = R"(
HloModule module

ENTRY entry {
  %param0 = f32[16,38,38,4] parameter(0)
  %param0.copy = f32[16,38,38,4] copy(%param0), sharding={devices=[1,2,1,1]0,1}
  ROOT %transpose = f32[16,4,38,38] transpose(%param0.copy),
    dimensions={0,3,1,2}, sharding={devices=[1,1,2,1]0,1}
})";

  TF_ASSERT_OK_AND_ASSIGN(auto module,
                          PartitionComputation(hlo_string, /*num_devices=*/2));
  VLOG(1) << module->ToString();

  auto root = module->entry_computation()->root_instruction();
  auto param0 = AllOf(
      op::Copy(op::DynamicSlice(op::Parameter(), op::Constant(), op::Reshape(),
                                op::Constant(), op::Constant())),
      op::Shape("f32[16,19,38,4]"));
  EXPECT_THAT(root, AllOf(op::Transpose(param0), op::Shape("f32[16,4,19,38]")));
}

TEST_F(SpmdPartitioningTest, MultiDimensionShardedTranspose) {
  absl::string_view hlo_string = R"(
HloModule module

ENTRY entry {
  %param0 = f32[16,38,38,4] parameter(0)
  %param0.copy = f32[16,38,38,4] copy(%param0),
    sharding={devices=[4,2,1,1]0,1,2,3,4,5,6,7}
  ROOT %transpose = f32[38,4,16,38] transpose(%param0.copy),
    dimensions={1,3,0,2}, sharding={devices=[2,1,4,1]0,2,4,6,1,3,5,7}
})";

  TF_ASSERT_OK_AND_ASSIGN(auto module,
                          PartitionComputation(hlo_string, /*num_devices=*/8));
  VLOG(1) << module->ToString();

  auto root = module->entry_computation()->root_instruction();
  auto param0 = AllOf(
      op::Copy(op::DynamicSlice(op::Parameter(), op::Reshape(), op::Reshape(),
                                op::Constant(), op::Constant())),
      op::Shape("f32[4,19,38,4]"));
  EXPECT_THAT(root, AllOf(op::Transpose(param0), op::Shape("f32[19,4,4,38]")));
}

TEST_F(SpmdPartitioningTest, NonShardableTranspose) {
  absl::string_view hlo_string = R"(
HloModule module

ENTRY entry {
  %param0 = f32[16,38,38,4] parameter(0)
  %param0.copy = f32[16,38,38,4] copy(%param0), sharding={devices=[1,2,1,1]0,1}
  ROOT %transpose = f32[16,4,38,38] transpose(%param0.copy),
    dimensions={0,3,1,2}, sharding={devices=[1,2,1,1]0,1}
})";

  TF_ASSERT_OK_AND_ASSIGN(auto module,
                          PartitionComputation(hlo_string, /*num_devices=*/2));
  VLOG(1) << module->ToString();

  auto root = module->entry_computation()->root_instruction();
  auto resahrd = AllOf(op::Reshape(op::Transpose(op::Reshape(op::AllToAll()))),
                       op::Shape("f32[16,38,38,2]"));
  EXPECT_THAT(root, AllOf(op::Transpose(), op::Shape("f32[16,2,38,38]")));
}

TEST_F(SpmdPartitioningTest, PartialReplicateShardableTranspose) {
  absl::string_view hlo_string = R"(
HloModule module

ENTRY entry {
  %param0 = f32[16,38,38,4] parameter(0)
  %param0.copy = f32[16,38,38,4] copy(%param0),
    sharding={devices=[1,2,1,1,2]0,1,2,3 last_tile_dim_replicate}
  ROOT %transpose = f32[16,4,38,38] transpose(%param0.copy),
    dimensions={0,3,1,2},
    sharding={devices=[1,1,2,1,2]0,1,2,3 last_tile_dim_replicate}
})";

  TF_ASSERT_OK_AND_ASSIGN(auto module,
                          PartitionComputation(hlo_string, /*num_devices=*/4));
  VLOG(1) << module->ToString();

  auto root = module->entry_computation()->root_instruction();
  auto param0 = AllOf(
      op::Copy(op::DynamicSlice(op::Parameter(), op::Constant(), op::Reshape(),
                                op::Constant(), op::Constant())),
      op::Shape("f32[16,19,38,4]"));
  EXPECT_THAT(root, AllOf(op::Transpose(param0), op::Shape("f32[16,4,19,38]")));
}

TEST_F(SpmdPartitioningTest, PartialReplicateNonShardableTranspose) {
  absl::string_view hlo_string = R"(
HloModule module

ENTRY entry {
  %param0 = f32[16,38,38,4] parameter(0)
  %param0.copy = f32[16,38,38,4] copy(%param0),
    sharding={devices=[1,2,1,1,2]0,1,2,3 last_tile_dim_replicate}
  ROOT %transpose = f32[16,4,38,38] transpose(%param0.copy),
    dimensions={0,3,1,2},
    sharding={devices=[1,2,1,1,2]0,1,2,3 last_tile_dim_replicate}
})";

  TF_ASSERT_OK_AND_ASSIGN(auto module,
                          PartitionComputation(hlo_string, /*num_devices=*/4));
  VLOG(1) << module->ToString();

  auto root = module->entry_computation()->root_instruction();
  auto resahrd = AllOf(op::Reshape(op::Transpose(op::Reshape(op::AllToAll()))),
                       op::Shape("f32[16,38,38,2]"));
  EXPECT_THAT(root, AllOf(op::Transpose(), op::Shape("f32[16,2,38,38]")));
}

TEST_F(SpmdPartitioningTest, PartialReplicateMultiDimensionShardedTranspose) {
  absl::string_view hlo_string = R"(
HloModule module

ENTRY entry {
  %param0 = f32[16,38,38,4] parameter(0)
  %param0.copy = f32[16,38,38,4] copy(%param0),
    sharding={devices=[2,2,1,1,2]0,1,2,3,4,5,6,7 last_tile_dim_replicate}
  ROOT %transpose = f32[38,4,16,38] transpose(%param0.copy),
    dimensions={1,3,0,2},
    sharding={devices=[2,1,2,1,2]0,1,4,5,2,3,6,7 last_tile_dim_replicate}
})";

  TF_ASSERT_OK_AND_ASSIGN(auto module,
                          PartitionComputation(hlo_string, /*num_devices=*/8));
  VLOG(1) << module->ToString();

  auto root = module->entry_computation()->root_instruction();
  auto param0 = AllOf(
      op::Copy(op::DynamicSlice(op::Parameter(), op::Reshape(), op::Reshape(),
                                op::Constant(), op::Constant())),
      op::Shape("f32[8,19,38,4]"));
  EXPECT_THAT(root, AllOf(op::Transpose(param0), op::Shape("f32[19,4,8,38]")));
}

TEST_F(SpmdPartitioningTest, ShardableReshape) {
  absl::string_view hlo_string = R"(
HloModule module

ENTRY entry {
  %param0 = f32[38,38,324] parameter(0)
  %param0.copy = f32[38,38,324] copy(%param0), sharding={devices=[2,1,1]0,1}
  ROOT %reshape = f32[38,38,4,81] reshape(%param0.copy),
    sharding={devices=[2,1,1,1]0,1}
})";

  TF_ASSERT_OK_AND_ASSIGN(auto module,
                          PartitionComputation(hlo_string, /*num_devices=*/2));
  VLOG(1) << module->ToString();

  auto root = module->entry_computation()->root_instruction();
  auto param0 =
      AllOf(op::Copy(op::DynamicSlice(op::Parameter(), op::Reshape(),
                                      op::Constant(), op::Constant())),
            op::Shape("f32[19,38,324]"));
  EXPECT_THAT(root, AllOf(op::Reshape(param0), op::Shape("f32[19,38,4,81]")));
}

TEST_F(SpmdPartitioningTest, ReshapeWithReshard) {
  absl::string_view hlo_string = R"(
HloModule module

ENTRY entry {
  %param0 = f32[38,38,324] parameter(0), sharding={devices=[2,1,1]0,1}
  ROOT %reshape = f32[38,38,4,81] reshape(%param0),
    sharding={devices=[1,2,1,1]0,1}
})";

  TF_ASSERT_OK_AND_ASSIGN(auto module,
                          PartitionComputation(hlo_string, /*num_devices=*/2));
  VLOG(1) << module->ToString();

  auto root = module->entry_computation()->root_instruction();
  auto input_reshard =
      op::Reshape(op::Transpose(op::AllToAll(op::Reshape(op::Parameter(0)))));
  EXPECT_THAT(root,
              AllOf(op::Reshape(input_reshard), op::Shape("f32[38,19,4,81]")));
}

TEST_F(SpmdPartitioningTest, ReshapeWithReshard2) {
  absl::string_view hlo_string = R"(
HloModule module

ENTRY entry {
  %param0 = f32[38,38,324] parameter(0), sharding={devices=[2,1,1]0,1}
  ROOT %reshape = f32[38,38,2,162] reshape(%param0),
    sharding={devices=[1,1,1,2]0,1}
})";

  TF_ASSERT_OK_AND_ASSIGN(auto module,
                          PartitionComputation(hlo_string, /*num_devices=*/2));
  VLOG(1) << module->ToString();

  auto root = module->entry_computation()->root_instruction();
  auto local_reshape =
      AllOf(op::Reshape(op::Parameter(0)), op::Shape("f32[19,38,2,162]"));
  EXPECT_THAT(root, AllOf(op::Shape("f32[38,38,2,81]"),
                          op::Reshape(op::Transpose(
                              op::AllToAll(op::Reshape(local_reshape))))));
}

TEST_F(SpmdPartitioningTest, PartialReplicateShardableReshape) {
  absl::string_view hlo_string = R"(
HloModule module

ENTRY entry {
  %param0 = f32[38,38,324] parameter(0)
  %param0.copy = f32[38,38,324] copy(%param0),
    sharding={devices=[2,1,1,2]0,1,2,3 last_tile_dim_replicate}
  ROOT %reshape = f32[38,38,4,81] reshape(%param0.copy),
    sharding={devices=[2,1,1,1,2]0,1,2,3 last_tile_dim_replicate}
})";

  TF_ASSERT_OK_AND_ASSIGN(auto module,
                          PartitionComputation(hlo_string, /*num_devices=*/4));
  VLOG(1) << module->ToString();

  auto root = module->entry_computation()->root_instruction();
  auto param0 =
      AllOf(op::Copy(op::DynamicSlice(op::Parameter(), op::Reshape(),
                                      op::Constant(), op::Constant())),
            op::Shape("f32[19,38,324]"));
  EXPECT_THAT(root, AllOf(op::Reshape(param0), op::Shape("f32[19,38,4,81]")));
}

TEST_F(SpmdPartitioningTest, ReshapeMergeDimsWithHaloExchange) {
  absl::string_view hlo_string = R"(
HloModule module

ENTRY entry {
  %input = s32[2,3,7,10] parameter(0), sharding={devices=[1,1,2,1]0,1}
  ROOT %reshape = s32[3,2,1,14,5] reshape(%input),
    sharding={devices=[1,1,1,2,1]0,1}
})";

  TF_ASSERT_OK_AND_ASSIGN(auto module,
                          PartitionComputation(hlo_string, /*num_devices=*/2));
  VLOG(1) << module->ToString();

  auto reshape =
      AllOf(op::Reshape(op::Parameter(0)), op::Shape("s32[3,2,1,8,5]"));
  auto halo = op::CollectivePermute(op::Slice(reshape));
  auto exchanged =
      op::DynamicSlice(op::Concatenate(halo, reshape), _, _, _, _, _);
  auto root = module->entry_computation()->root_instruction();
  EXPECT_THAT(root, AllOf(exchanged, op::Shape("s32[3,2,1,7,5]")));
}

TEST_F(SpmdPartitioningTest, PartialReplicateReshapeMergeDimsWithHaloExchange) {
  absl::string_view hlo_string = R"(
HloModule module

ENTRY entry {
  %input = s32[2,3,7,10] parameter(0),
    sharding={devices=[1,1,2,1,2]0,1,2,3 last_tile_dim_replicate}
  ROOT %reshape = s32[3,2,1,14,5] reshape(%input),
    sharding={devices=[1,1,1,2,1,2]0,1,2,3 last_tile_dim_replicate}
})";

  TF_ASSERT_OK_AND_ASSIGN(auto module,
                          PartitionComputation(hlo_string, /*num_devices=*/4));
  VLOG(1) << module->ToString();

  auto reshape =
      AllOf(op::Reshape(op::Parameter(0)), op::Shape("s32[3,2,1,8,5]"));
  auto halo = op::CollectivePermute(op::Slice(reshape));
  auto exchanged =
      op::DynamicSlice(op::Concatenate(halo, reshape), _, _, _, _, _);
  auto root = module->entry_computation()->root_instruction();
  EXPECT_THAT(root, AllOf(exchanged, op::Shape("s32[3,2,1,7,5]")));
}

// Produces an invalid module after transformation.
TEST_F(SpmdPartitioningTest, InceptionV3_4_way_ReduceWindowDilated) {
  absl::string_view hlo_string = R"(
HloModule module

sum {
  a = f32[] parameter(0)
  b = f32[] parameter(1)
  ROOT add = f32[] add(a, b)
}

ENTRY entry {
  %param0 = f32[128,5,5,768] parameter(0)
  %param0.copy = f32[128,5,5,768] copy(%param0),
    sharding={devices=[1,4,1,1]0,1,2,3}
  %constant.1 = f32[] constant(0), sharding={replicated}
  ROOT %rw = f32[128,17,17,768] reduce-window(%param0.copy, %constant.1),
    window={size=1x5x5x1 pad=0_0x4_4x4_4x0_0 lhs_dilate=1x3x3x1},
    to_apply=sum, sharding={devices=[1,4,1,1]0,1,2,3}
})";

  TF_ASSERT_OK_AND_ASSIGN(auto module,
                          PartitionComputation(hlo_string, /*num_devices=*/4));
  VLOG(1) << module->ToString();

  auto input_shard = op::Copy(op::DynamicSlice(
      op::Pad(op::Parameter(0), op::Constant()), op::Constant(), op::Reshape(),
      op::Constant(), op::Constant()));
  auto id_mul4_add1 =
      op::Add(op::Multiply(op::Reshape(), op::Constant()), op::Constant());
  auto id_mul5 = op::Multiply(op::Reshape(), op::Constant());
  auto id_mul5_add1_div3 =
      op::Divide(op::Add(id_mul5, op::Constant()), op::Constant());
  auto before_masking = AllOf(
      op::Shape("f32[128,3,5,768]"),
      op::DynamicSlice(
          AllOf(
              op::Shape("f32[128,4,5,768]"),
              op::Concatenate(op::CollectivePermute(input_shard), input_shard)),
          op::Constant(),
          op::Subtract(op::Constant(),
                       op::Subtract(id_mul4_add1, id_mul5_add1_div3)),
          op::Constant(), op::Constant()));
  auto masked = op::Select(
      op::And(op::Compare(op::Add(op::Iota(), op::Broadcast(id_mul5_add1_div3)),
                          op::Broadcast(op::Constant())),
              op::Compare(op::Add(op::Iota(), op::Broadcast(id_mul5_add1_div3)),
                          op::Broadcast(op::Constant()))),
      before_masking, op::Broadcast(op::Constant()));
  auto rw = AllOf(op::Shape("f32[128,7,17,768]"),
                  op::ReduceWindow(masked, op::Constant()));
  auto final_slice_index = op::Subtract(
      id_mul5,
      op::Add(op::Multiply(id_mul5_add1_div3, op::Constant()), op::Constant()));
  auto root = module->entry_computation()->root_instruction();
  EXPECT_THAT(root,
              AllOf(op::Shape("f32[128,5,17,768]"),
                    op::DynamicSlice(rw, op::Constant(), final_slice_index,
                                     op::Constant(), op::Constant())));
}

TEST_F(SpmdPartitioningTest, TiledToTiledReduce) {
  absl::string_view hlo_string = R"(
HloModule module

sum {
  a = f32[] parameter(0)
  b = f32[] parameter(1)
  ROOT add = f32[] add(a, b)
}

ENTRY entry {
  %param0 = f32[4,32,32,128] parameter(0)
  %param0.copy = f32[4,32,32,128] copy(%param0),
    sharding={devices=[1,1,1,2]0,1}
  %constant.1 = f32[] constant(0), sharding={replicated}
  %reduce = f32[128] reduce(%param0.copy, %constant.1), dimensions={0,1,2},
    to_apply=%sum, sharding={devices=[2]0,1}
})";

  TF_ASSERT_OK_AND_ASSIGN(auto module,
                          PartitionComputation(hlo_string, /*num_devices=*/2));
  VLOG(1) << module->ToString();

  auto root = module->entry_computation()->root_instruction();
  auto param0 = AllOf(
      op::Copy(op::DynamicSlice(op::Parameter(), op::Constant(), op::Constant(),
                                op::Constant(), op::Reshape())),
      op::Shape("f32[4,32,32,64]"));

  EXPECT_THAT(root,
              AllOf(op::Reduce(param0, op::Constant()), op::Shape("f32[64]")));
}

TEST_F(SpmdPartitioningTest, PartialTiledToPartialTiledReduce) {
  absl::string_view hlo_string = R"(
HloModule module

sum {
  a = f32[] parameter(0)
  b = f32[] parameter(1)
  ROOT add = f32[] add(a, b)
}

ENTRY entry {
  %param0 = f32[4,4] parameter(0),
    sharding={devices=[2,2,2]0,1,2,3,4,5,6,7 last_tile_dim_replicate}
  %constant.1 = f32[] constant(0), sharding={replicated}
  ROOT %reduce = f32[4] reduce(%param0, %constant.1), dimensions={0},
    to_apply=%sum,
    sharding={devices=[2,4]0,1,4,5,2,3,6,7 last_tile_dim_replicate}
})";

  TF_ASSERT_OK_AND_ASSIGN(auto module,
                          PartitionComputation(hlo_string, /*num_devices=*/8));
  VLOG(1) << module->ToString();

  auto root = module->entry_computation()->root_instruction();
  EXPECT_THAT(root,
              AllOf(op::AllReduce(op::Reduce(op::Parameter(0), op::Constant())),
                    op::Shape("f32[2]")));
}

TEST_F(SpmdPartitioningTest, TiledToTiledTupleReduce) {
  absl::string_view hlo_string = R"(
HloModule module

%minmax_func {
  %lhs_value = f32[] parameter(0)
  %rhs_value = f32[] parameter(2)
  %compare.2 = pred[] compare(%lhs_value, %rhs_value), direction=GT
  %select.4 = f32[] select(%compare.2, %lhs_value, %rhs_value)
  %lhs_index = s32[] parameter(1)
  %rhs_index = s32[] parameter(3)
  %select.5 = s32[] select(%compare.2, %lhs_index, %rhs_index)
  ROOT %tuple.2 = (f32[], s32[]) tuple(%select.4, %select.5)
}

ENTRY %main {
  %param0 = f32[28,10] parameter(0), sharding={devices=[2,1]0,1}
  %param1 = s32[28,10] parameter(1), sharding={devices=[2,1]0,1}
  %init0 = f32[] parameter(2)
  %init1 = s32[] parameter(3)
  ROOT %reduce = (f32[28], s32[28]) reduce(%param0, %param1, %init0, %init1),
    dimensions={1}, to_apply=%minmax_func,
    sharding={{devices=[2]0,1}, {devices=[2]0,1}}
})";

  TF_ASSERT_OK_AND_ASSIGN(auto module,
                          PartitionComputation(hlo_string, /*num_devices=*/2));
  VLOG(1) << module->ToString();

  auto root = module->entry_computation()->root_instruction();
  EXPECT_THAT(root, AllOf(op::Reduce(op::Parameter(0), op::Parameter(1),
                                     op::Parameter(2), op::Parameter(3)),
                          op::Shape("(f32[14], s32[14])")));
}

TEST_F(SpmdPartitioningTest, TiledToPartiallyTiledTupleReduce) {
  absl::string_view hlo_string = R"(
HloModule module

%minmax_func {
  %lhs_value = f32[] parameter(0)
  %rhs_value = f32[] parameter(2)
  %compare.2 = pred[] compare(%lhs_value, %rhs_value), direction=GT
  %select.4 = f32[] select(%compare.2, %lhs_value, %rhs_value)
  %lhs_index = s32[] parameter(1)
  %rhs_index = s32[] parameter(3)
  %select.5 = s32[] select(%compare.2, %lhs_index, %rhs_index)
  ROOT %tuple.2 = (f32[], s32[]) tuple(%select.4, %select.5)
}

ENTRY %main {
  %param0 = f32[28,12] parameter(0), sharding={devices=[2,4]0,1,2,3,4,5,6,7}
  %param1 = s32[28,12] parameter(1), sharding={devices=[2,4]0,1,2,3,4,5,6,7}
  %init0 = f32[] parameter(2)
  %init1 = s32[] parameter(3)
  ROOT %reduce = (f32[28], s32[28]) reduce(%param0, %param1, %init0, %init1),
    dimensions={1}, to_apply=%minmax_func,
    sharding={{devices=[2,4]0,1,2,3,4,5,6,7 last_tile_dim_replicate},
              {devices=[2,4]0,1,2,3,4,5,6,7 last_tile_dim_replicate}}
})";

  TF_ASSERT_OK_AND_ASSIGN(auto module,
                          PartitionComputation(hlo_string, /*num_devices=*/8));
  VLOG(1) << module->ToString();

  auto lhs = AllOf(op::Shape("f32[14,3]"), op::Parameter(0));
  auto rhs = AllOf(op::Shape("s32[14,3]"), op::Parameter(1));
  auto local_reduce =
      AllOf(op::Reduce(lhs, rhs, op::Parameter(2), op::Parameter(3)),
            op::Shape("(f32[14], s32[14])"));
  auto reshape_l = AllOf(op::Reshape(op::GetTupleElement(local_reduce)),
                         op::Shape("f32[14,1]"));
  auto reshape_r = AllOf(op::Reshape(op::GetTupleElement(local_reduce)),
                         op::Shape("s32[14,1]"));
  auto broadcast_l =
      AllOf(op::AllReduce(op::DynamicUpdateSlice(_, reshape_l, _, _)),
            op::Shape("f32[14,4]"));
  auto broadcast_r =
      AllOf(op::AllReduce(op::DynamicUpdateSlice(_, reshape_r, _, _)),
            op::Shape("s32[14,4]"));
  auto root = module->entry_computation()->root_instruction();
  EXPECT_THAT(root, AllOf(op::Reduce(broadcast_l, broadcast_r, op::Parameter(2),
                                     op::Parameter(3)),
                          op::Shape("(f32[14], s32[14])")));
}

TEST_F(SpmdPartitioningTest, TiledToTiledReduceOutputReshard) {
  absl::string_view hlo_string = R"(
HloModule module

sum {
  a = f32[] parameter(0)
  b = f32[] parameter(1)
  ROOT add = f32[] add(a, b)
}

ENTRY entry {
  %param0 = f32[4,32,32,128] parameter(0)
  %param0.copy = f32[4,32,32,128] copy(%param0),
    sharding={devices=[1,2,1,1]0,1}
  %constant.1 = f32[] constant(0), sharding={replicated}
  %reduce = f32[128] reduce(%param0.copy, %constant.1), dimensions={0,1,2},
    to_apply=%sum, sharding={devices=[2]0,1}
})";

  TF_ASSERT_OK_AND_ASSIGN(auto module,
                          PartitionComputation(hlo_string, /*num_devices=*/2));
  VLOG(1) << module->ToString();

  auto root = module->entry_computation()->root_instruction();
  auto param0 = AllOf(
      op::Copy(op::DynamicSlice(op::Parameter(), op::Constant(), op::Reshape(),
                                op::Constant(), op::Constant())),
      op::Shape("f32[4,16,32,128]"));

  EXPECT_THAT(root,
              AllOf(op::DynamicSlice(
                        AllOf(op::AllReduce(op::Reduce(param0, op::Constant())),
                              op::Shape("f32[128]")),
                        op::Reshape()),
                    op::Shape("f32[64]")));
}

TEST_F(SpmdPartitioningTest, IotaAlongNonTileDimension) {
  absl::string_view hlo_string = R"(
HloModule module

ENTRY entry {
  ROOT %iota = s32[16,80,91] iota(), iota_dimension=1,
    sharding={devices=[1,1,2]0,1}
})";

  TF_ASSERT_OK_AND_ASSIGN(auto module,
                          PartitionComputation(hlo_string, /*num_devices=*/2));
  VLOG(1) << module->ToString();

  auto root = module->entry_computation()->root_instruction();
  EXPECT_THAT(root, AllOf(op::Iota(), op::Shape("s32[16,80,46]")));
}

TEST_F(SpmdPartitioningTest, IotaAlongTileDimension) {
  absl::string_view hlo_string = R"(
HloModule module

ENTRY entry {
  ROOT %iota = s32[16,80,91] iota(), iota_dimension=2,
    sharding={devices=[1,1,2]0,1}
})";

  TF_ASSERT_OK_AND_ASSIGN(auto module,
                          PartitionComputation(hlo_string, /*num_devices=*/2));
  VLOG(1) << module->ToString();

  auto root = module->entry_computation()->root_instruction();
  EXPECT_THAT(root, AllOf(op::Add(op::Iota(), op::Broadcast()),
                          op::Shape("s32[16,80,46]")));
}

TEST_F(SpmdPartitioningTest, U32IotaAlongTileDimension) {
  absl::string_view hlo_string = R"(
HloModule module

ENTRY entry {
  ROOT %iota = u32[16,80,91] iota(), iota_dimension=2,
    sharding={devices=[1,1,2]0,1}
})";

  TF_ASSERT_OK_AND_ASSIGN(auto module,
                          PartitionComputation(hlo_string, /*num_devices=*/2));
  VLOG(1) << module->ToString();

  auto root = module->entry_computation()->root_instruction();
  EXPECT_THAT(root, AllOf(op::Add(op::Iota(), op::Broadcast()),
                          op::Shape("u32[16,80,46]")));
}

TEST_F(SpmdPartitioningTest, Conditional) {
  absl::string_view hlo_string = R"(
HloModule module

Negate {
  x = f32[4,5] parameter(0), sharding={replicated}
  ROOT negate = f32[4,5] negate(x), sharding={replicated}
}

Identity {
  y = f32[4,5] parameter(0), sharding={devices=[2,1]0,1}
  ROOT copy = f32[4,5] copy(y), sharding={devices=[2,1]0,1}
}

ENTRY entry {
  %param.0 = pred[] parameter(0)
  %param.0.copy = pred[] copy(%param.0), sharding={maximal device=0}
  %param.1 = f32[4,5] parameter(1)
  %param.1.copy = f32[4,5] copy(%param.1), sharding={replicated}
  %param.2 = f32[4,5] parameter(2)
  %param.2.copy = f32[4,5] copy(%param.2), sharding={devices=[2,1]0,1}
  ROOT cond = f32[4,5] conditional(%param.0.copy, %param.1.copy, %param.2.copy),
    true_computation=Negate, false_computation=Identity,
    sharding={devices=[2,1]0,1}
})";

  TF_ASSERT_OK_AND_ASSIGN(auto module,
                          PartitionComputation(hlo_string, /*num_devices=*/2));
  VLOG(1) << module->ToString();

  auto param0 = AllOf(op::Copy(op::Copy(op::Parameter()), op::Shape("pred[]")));
  auto param1 = AllOf(op::Copy(op::Parameter()), op::Shape("f32[4,5]"));
  auto param2 = AllOf(op::Copy(op::DynamicSlice(op::Parameter(), op::Reshape(),
                                                op::Constant())),
                      op::Shape("f32[2,5]"));

  auto root = module->entry_computation()->root_instruction();
  EXPECT_THAT(root, AllOf(op::Conditional(op::AllReduce(), param1, param2),
                          op::Shape("f32[2,5]")));

  auto then_branch_root = root->branch_computation(0)->root_instruction();
  EXPECT_THAT(then_branch_root,
              AllOf(op::DynamicSlice(op::Negate(op::Parameter()), op::Reshape(),
                                     op::Constant()),
                    op::Shape("f32[2,5]")));

  auto else_branch_root = root->branch_computation(1)->root_instruction();
  EXPECT_THAT(else_branch_root,
              AllOf(op::Copy(op::Parameter()), op::Shape("f32[2,5]")));
}

TEST_F(SpmdPartitioningTest, SelectAndScatter_RetinaNet) {
  absl::string_view hlo_string = R"(
HloModule module

ge {
  a = f32[] parameter(0)
  b = f32[] parameter(1)
  ROOT compare = pred[] compare(a, b), direction=GE
}

sum {
  c = f32[] parameter(0)
  d = f32[] parameter(1)
  ROOT add = f32[] add(c, d)
}

ENTRY entry {
  %param.0 = f32[32,128,384,64] parameter(0)
  %param.0.copy = f32[32,128,384,64] copy(%param.0),
    sharding={devices=[1,8,1,1]0,1,2,3,4,5,6,7}
  %param.1 = f32[32,64,192,64] parameter(1)
  %param.1.copy = f32[32,64,192,64] copy(%param.1),
    sharding={devices=[1,8,1,1]0,1,2,3,4,5,6,7}
  constant.1 = f32[] constant(0), sharding={replicated}
  ROOT select-and-scatter = f32[32,128,384,64] select-and-scatter(param.0.copy,
    %param.1.copy, constant.1), window={size=1x1x1x1 stride=1x2x2x1},
    select=ge, scatter=sum, sharding={devices=[1,8,1,1]0,1,2,3,4,5,6,7}
})";
  TF_ASSERT_OK_AND_ASSIGN(auto module,
                          PartitionComputation(hlo_string, /*num_devices=*/8));
  VLOG(1) << module->ToString();

  auto root = module->entry_computation()->root_instruction();
  auto source = AllOf(
      op::Shape("f32[32,8,192,64]"),
      op::Copy(op::DynamicSlice(op::Parameter(1), op::Constant(), op::Reshape(),
                                op::Constant(), op::Constant())));
  auto data = AllOf(
      op::Shape("f32[32,16,384,64]"),
      op::Copy(op::DynamicSlice(op::Parameter(0), op::Constant(), op::Reshape(),
                                op::Constant(), op::Constant())));

  EXPECT_THAT(root, op::SelectAndScatter(data, source, op::Constant()));
  EXPECT_EQ(root->window().dimensions(0).padding_low(), 0);
  EXPECT_EQ(root->window().dimensions(0).padding_high(), 0);
}

TEST_F(SpmdPartitioningTest, TiledDot) {
  absl::string_view hlo_string = R"(
HloModule module

ENTRY entry {
  %lhs = f32[128,64] parameter(0)
  %lhs.copy = f32[128,64] copy(%lhs), sharding={devices=[1,2]0,1}
  %rhs = f32[64,256] parameter(1)
  %rhs.copy = f32[64,256] copy(%rhs), sharding={devices=[2,1]0,1}
  ROOT %conv = f32[128,256] convolution(%lhs.copy, %rhs.copy),
    dim_labels=bf_io->bf, sharding={replicated}
})";

  TF_ASSERT_OK_AND_ASSIGN(
      auto module,
      PartitionComputation(hlo_string, /*num_devices=*/2,
                           /*conv_halo_exchange_always_on_lhs=*/false));
  VLOG(1) << module->ToString();

  auto root = module->entry_computation()->root_instruction();
  auto lhs = AllOf(op::Copy(op::DynamicSlice(op::Parameter(), op::Constant(),
                                             op::Reshape())),
                   op::Shape("f32[128,32]"));
  auto rhs = AllOf(op::Copy(op::DynamicSlice(op::Parameter(), op::Reshape(),
                                             op::Constant())),
                   op::Shape("f32[32,256]"));
  EXPECT_THAT(root, AllOf(op::AllReduce(op::Convolution(lhs, rhs)),
                          op::Shape("f32[128,256]")));
}

TEST_F(SpmdPartitioningTest, TiledDotOutputTiled) {
  absl::string_view hlo_string = R"(
HloModule module

ENTRY entry {
  %lhs = f32[128,64] parameter(0)
  %lhs.copy = f32[128,64] copy(%lhs), sharding={devices=[1,2]0,1}
  %rhs = f32[64,256] parameter(1)
  %rhs.copy = f32[64,256] copy(%rhs), sharding={devices=[2,1]0,1}
  ROOT %conv = f32[128,256] convolution(%lhs.copy, %rhs.copy),
    dim_labels=bf_io->bf, sharding={devices=[1,2]0,1}
})";

  TF_ASSERT_OK_AND_ASSIGN(auto module,
                          PartitionComputation(hlo_string, /*num_devices=*/2));
  VLOG(1) << module->ToString();

  auto root = module->entry_computation()->root_instruction();
  auto lhs = AllOf(op::Copy(op::DynamicSlice(op::Parameter(), op::Constant(),
                                             op::Reshape())),
                   op::Shape("f32[128,32]"));
  auto rhs = AllOf(op::Copy(op::DynamicSlice(op::Parameter(), op::Reshape(),
                                             op::Constant())),
                   op::Shape("f32[32,256]"));
  EXPECT_THAT(root, AllOf(op::DynamicSlice(
                              AllOf(op::AllReduce(op::Convolution(lhs, rhs)),
                                    op::Shape("f32[128,256]")),
                              op::Constant(), op::Reshape()),
                          op::Shape("f32[128,128]")));
}

TEST_F(SpmdPartitioningTest, BatchPartitionedConvolution) {
  absl::string_view hlo_string = R"(
HloModule module

ENTRY entry {
  %lhs = f32[128,256,256] parameter(0)
  %lhs.copy = f32[128,256,256] copy(%lhs), sharding={devices=[1,2,1]0,1}
  %rhs = f32[256,8,1] parameter(1)
  %rhs.copy = f32[256,8,1] copy(%rhs), sharding={replicated}
  ROOT %conv = f32[128,256,8] convolution(%lhs.copy, %rhs.copy),
    window={size=1}, dim_labels=0bf_io0->0bf, sharding={devices=[1,2,1]0,1}
})";

  TF_ASSERT_OK_AND_ASSIGN(auto module,
                          PartitionComputation(hlo_string, /*num_devices=*/2));
  VLOG(1) << module->ToString();

  auto root = module->entry_computation()->root_instruction();
  auto lhs = AllOf(op::Copy(op::DynamicSlice(op::Parameter(0), op::Constant(),
                                             op::Reshape(), op::Constant())),
                   op::Shape("f32[128,128,256]"));
  auto rhs = AllOf(op::Copy(op::Parameter(1)), op::Shape("f32[256,8,1]"));
  EXPECT_THAT(root,
              AllOf(op::Convolution(lhs, rhs), op::Shape("f32[128,128,8]")));
}

TEST_F(SpmdPartitioningTest, DotOutputFeaturePartitioned) {
  absl::string_view hlo_string = R"(
HloModule module

ENTRY entry {
  %lhs = f32[24,64] parameter(0)
  %lhs.copy = f32[24,64] copy(%lhs), sharding={replicated}
  %rhs = f32[39296,64] parameter(1)
  %rhs.copy = f32[39296,64] copy(%rhs), sharding={devices=[2,1]0,1}
  ROOT %dot = f32[24,39296] dot(%lhs.copy, %rhs.copy),
    lhs_batch_dims={}, rhs_batch_dims={},
    lhs_contracting_dims={1}, rhs_contracting_dims={1},
    sharding={devices=[1,2]0,1}
})";

  TF_ASSERT_OK_AND_ASSIGN(auto module,
                          PartitionComputation(hlo_string, /*num_devices=*/2));
  VLOG(1) << module->ToString();

  auto root = module->entry_computation()->root_instruction();
  auto lhs = AllOf(op::Copy(op::Parameter(0)), op::Shape("f32[24,64]"));
  auto rhs = AllOf(op::Copy(op::DynamicSlice(op::Parameter(1), op::Reshape(),
                                             op::Constant())),
                   op::Shape("f32[19648,64]"));
  EXPECT_THAT(root, AllOf(op::Dot(lhs, rhs), op::Shape("f32[24,19648]")));
}

TEST_F(SpmdPartitioningTest, DotPartialDeviceOrder) {
  absl::string_view hlo_string = R"(
HloModule module

ENTRY entry {
  %lhs = f32[16,256,4096] parameter(0), sharding={devices=[1,1,2,2]1,3,0,2 last_tile_dim_replicate}
  %rhs = f32[4096,2048] parameter(1), sharding={devices=[2,2]3,1,2,0}
  ROOT %dot = f32[16,256,2048] dot(%lhs, %rhs),
    lhs_batch_dims={}, rhs_batch_dims={},
    lhs_contracting_dims={2}, rhs_contracting_dims={0},
    sharding={devices=[1,1,2,2]2,3,0,1 last_tile_dim_replicate}
})";

  TF_ASSERT_OK_AND_ASSIGN(auto module,
                          PartitionComputation(hlo_string, /*num_devices=*/4));
  VLOG(1) << module->ToString();

  auto root = module->entry_computation()->root_instruction();
  auto lhs = AllOf(op::Parameter(0), op::Shape("f32[16,256,2048]"));
  auto rhs = AllOf(op::Parameter(1), op::Shape("f32[2048,1024]"));
  EXPECT_THAT(root, AllOf(op::AllReduce(op::Dot(lhs, rhs)),
                          op::Shape("f32[16,256,1024]")));
}

TEST_F(SpmdPartitioningTest, EinsumBatchPartitioned) {
  absl::string_view hlo_string = R"(
HloModule module

ENTRY entry {
  %lhs = f32[32,24,64] parameter(0)
  %lhs.copy = f32[32,24,64] copy(%lhs), sharding={devices=[2,1,1]0,1}
  %rhs = f32[32,39296,64] parameter(1)
  %rhs.copy = f32[32,39296,64] copy(%rhs), sharding={devices=[2,1,1]0,1}
  ROOT %dot = f32[32,24,39296] dot(%lhs.copy, %rhs.copy),
    lhs_batch_dims={0}, rhs_batch_dims={0},
    lhs_contracting_dims={2}, rhs_contracting_dims={2},
    sharding={devices=[2,1,1]0,1}
})";

  TF_ASSERT_OK_AND_ASSIGN(auto module,
                          PartitionComputation(hlo_string, /*num_devices=*/2));
  VLOG(1) << module->ToString();

  auto root = module->entry_computation()->root_instruction();
  auto lhs = AllOf(op::Copy(op::DynamicSlice(op::Parameter(0), op::Reshape(),
                                             op::Constant(), op::Constant())),
                   op::Shape("f32[16,24,64]"));
  auto rhs = AllOf(op::Copy(op::DynamicSlice(op::Parameter(1), op::Reshape(),
                                             op::Constant(), op::Constant())),
                   op::Shape("f32[16,39296,64]"));
  EXPECT_THAT(root, AllOf(op::Dot(lhs, rhs), op::Shape("f32[16,24,39296]")));
}

TEST_F(SpmdPartitioningTest, EinsumLHSandOutputBatchPartitioned) {
  absl::string_view hlo_string = R"(
HloModule module

ENTRY entry {
  %lhs = f32[32,24,64] parameter(0)
  %lhs.copy = f32[32,24,64] copy(%lhs), sharding={devices=[2,1,1]0,1}
  %rhs = f32[32,39296,64] parameter(1)
  %rhs.copy = f32[32,39296,64] copy(%rhs), sharding={replicated}
  ROOT %dot = f32[32,24,39296] dot(%lhs.copy, %rhs.copy),
    lhs_batch_dims={0}, rhs_batch_dims={0},
    lhs_contracting_dims={2}, rhs_contracting_dims={2},
    sharding={devices=[2,1,1]0,1}
})";

  TF_ASSERT_OK_AND_ASSIGN(auto module,
                          PartitionComputation(hlo_string, /*num_devices=*/2));
  VLOG(1) << module->ToString();

  auto root = module->entry_computation()->root_instruction();
  auto lhs = AllOf(op::Copy(op::DynamicSlice(op::Parameter(0), op::Reshape(),
                                             op::Constant(), op::Constant())),
                   op::Shape("f32[16,24,64]"));
  auto rhs = AllOf(op::Copy(op::Parameter(1)), op::Shape("f32[32,39296,64]"));
  EXPECT_THAT(root, AllOf(op::Dot(lhs, op::DynamicSlice(rhs, op::Reshape(),
                                                        op::Constant(),
                                                        op::Constant())),
                          op::Shape("f32[16,24,39296]")));
}

TEST_F(SpmdPartitioningTest, EinsumRHSandOutputBatchPartitioned) {
  absl::string_view hlo_string = R"(
HloModule module

ENTRY entry {
  %lhs = f32[32,24,64] parameter(0)
  %lhs.copy = f32[32,24,64] copy(%lhs), sharding={devices=[1,2,1]0,1}
  %rhs = f32[32,39296,64] parameter(1)
  %rhs.copy = f32[32,39296,64] copy(%rhs), sharding={devices=[2,1,1]0,1}
  ROOT %dot = f32[32,24,39296] dot(%lhs.copy, %rhs.copy),
    lhs_batch_dims={0}, rhs_batch_dims={0},
    lhs_contracting_dims={2}, rhs_contracting_dims={2},
    sharding={devices=[2,1,1]0,1}
})";

  TF_ASSERT_OK_AND_ASSIGN(auto module,
                          PartitionComputation(hlo_string, /*num_devices=*/2));
  VLOG(1) << module->ToString();

  auto root = module->entry_computation()->root_instruction();
  auto lhs = AllOf(op::Copy(op::DynamicSlice(op::Parameter(0), op::Constant(),
                                             op::Reshape(), op::Constant())),
                   op::Shape("f32[32,12,64]"));
  auto rhs = AllOf(op::Copy(op::DynamicSlice(op::Parameter(1), op::Reshape(),
                                             op::Constant(), op::Constant())),
                   op::Shape("f32[16,39296,64]"));
  auto lhs_reshard = op::Reshape(op::Transpose(op::AllToAll(op::Reshape(lhs))));
  EXPECT_THAT(root,
              AllOf(op::Dot(lhs_reshard, rhs), op::Shape("f32[16,24,39296]")));
}

TEST_F(SpmdPartitioningTest, EinsumOutputBatchPartitioned) {
  absl::string_view hlo_string = R"(
HloModule module

ENTRY entry {
  %lhs = f32[32,24,64] parameter(0)
  %lhs.copy = f32[32,24,64] copy(%lhs), sharding={replicated}
  %rhs = f32[32,39296,64] parameter(1)
  %rhs.copy = f32[32,39296,64] copy(%rhs), sharding={replicated}
  ROOT %dot = f32[32,24,39296] dot(%lhs.copy, %rhs.copy),
    lhs_batch_dims={0}, rhs_batch_dims={0},
    lhs_contracting_dims={2}, rhs_contracting_dims={2},
    sharding={devices=[2,1,1]0,1}
})";

  TF_ASSERT_OK_AND_ASSIGN(auto module,
                          PartitionComputation(hlo_string, /*num_devices=*/2));
  VLOG(1) << module->ToString();

  auto root = module->entry_computation()->root_instruction();
  auto lhs_slice =
      AllOf(op::DynamicSlice(op::Copy(op::Parameter(0)), op::Reshape(),
                             op::Constant(), op::Constant()),
            op::Shape("f32[16,24,64]"));
  auto rhs_slice =
      AllOf(op::DynamicSlice(op::Copy(op::Parameter(1)), op::Reshape(),
                             op::Constant(), op::Constant()),
            op::Shape("f32[16,39296,64]"));
  EXPECT_THAT(root, AllOf(op::Dot(lhs_slice, rhs_slice),
                          op::Shape("f32[16,24,39296]")));
}

TEST_F(SpmdPartitioningTest, EinsumContractingDimsPartitioned) {
  absl::string_view hlo_string = R"(
HloModule module

ENTRY entry {
  %lhs = f32[32,24,64,128] parameter(0)
  %lhs.copy = f32[32,24,64,128] copy(%lhs), sharding={devices=[1,1,2,2]0,1,2,3}
  %rhs = f32[32,39296,64,128] parameter(1)
  %rhs.copy = f32[32,39296,64,128] copy(%rhs), sharding={devices=[1,1,2,2]0,1,2,3}
  ROOT %dot = f32[32,24,39296] dot(%lhs.copy, %rhs.copy),
    lhs_batch_dims={0}, rhs_batch_dims={0},
    lhs_contracting_dims={2,3}, rhs_contracting_dims={2,3},
    sharding={replicated}
})";

  TF_ASSERT_OK_AND_ASSIGN(auto module,
                          PartitionComputation(hlo_string, /*num_devices=*/4));
  VLOG(1) << module->ToString();

  auto root = module->entry_computation()->root_instruction();
  auto lhs = AllOf(
      op::Copy(op::DynamicSlice(op::Parameter(0), op::Constant(),
                                op::Constant(), op::Reshape(), op::Reshape())),
      op::Shape("f32[32,24,32,64]"));
  auto rhs = AllOf(
      op::Copy(op::DynamicSlice(op::Parameter(1), op::Constant(),
                                op::Constant(), op::Reshape(), op::Reshape())),
      op::Shape("f32[32,39296,32,64]"));
  EXPECT_THAT(root, AllOf(op::AllReduce(op::AllReduce(op::Dot(lhs, rhs))),
                          op::Shape("f32[32,24,39296]")));
}

TEST_F(SpmdPartitioningTest, EinsumLHSNonContractingDimsPartitioned) {
  absl::string_view hlo_string = R"(
HloModule module

ENTRY entry {
  %lhs = f32[32,24,64,128] parameter(0)
  %lhs.copy = f32[32,24,64,128] copy(%lhs), sharding={devices=[1,2,1,2]0,1,2,3}
  %rhs = f32[32,39296,64] parameter(1)
  %rhs.copy = f32[32,39296,64] copy(%rhs), sharding={replicated}
  ROOT %dot = f32[32,24,128,39296] dot(%lhs.copy, %rhs.copy),
    lhs_batch_dims={0}, rhs_batch_dims={0},
    lhs_contracting_dims={2}, rhs_contracting_dims={2},
    sharding={devices=[1,2,2,1]0,1,2,3}
})";

  TF_ASSERT_OK_AND_ASSIGN(auto module,
                          PartitionComputation(hlo_string, /*num_devices=*/4));
  VLOG(1) << module->ToString();

  auto root = module->entry_computation()->root_instruction();
  auto lhs = AllOf(
      op::Copy(op::DynamicSlice(op::Parameter(0), op::Constant(), op::Reshape(),
                                op::Constant(), op::Reshape())),
      op::Shape("f32[32,12,64,64]"));
  auto rhs = AllOf(op::Copy(op::Parameter(1)), op::Shape("f32[32,39296,64]"));
  EXPECT_THAT(root, AllOf(op::Dot(lhs, rhs), op::Shape("f32[32,12,64,39296]")));
}

TEST_F(SpmdPartitioningTest, EinsumRHSNonContractingDimsPartitioned) {
  absl::string_view hlo_string = R"(
HloModule module

ENTRY entry {
  %lhs = f32[32,24,64] parameter(0)
  %lhs.copy = f32[32,24,64] copy(%lhs), sharding={replicated}
  %rhs = f32[32,39296,64,128] parameter(1)
  %rhs.copy = f32[32,39296,64,128] copy(%rhs), sharding={devices=[1,2,1,2]0,1,2,3}
  ROOT %dot = f32[32,24,39296,128] dot(%lhs.copy, %rhs.copy),
    lhs_batch_dims={0}, rhs_batch_dims={0},
    lhs_contracting_dims={2}, rhs_contracting_dims={2},
    sharding={devices=[1,1,2,2]0,1,2,3}
})";

  TF_ASSERT_OK_AND_ASSIGN(auto module,
                          PartitionComputation(hlo_string, /*num_devices=*/4));
  VLOG(1) << module->ToString();

  auto root = module->entry_computation()->root_instruction();
  auto lhs = AllOf(op::Copy(op::Parameter(0)), op::Shape("f32[32,24,64]"));
  auto rhs = AllOf(
      op::Copy(op::DynamicSlice(op::Parameter(1), op::Constant(), op::Reshape(),
                                op::Constant(), op::Reshape())),
      op::Shape("f32[32,19648,64,64]"));
  EXPECT_THAT(root, AllOf(op::Dot(lhs, rhs), op::Shape("f32[32,24,19648,64]")));
}

TEST_F(SpmdPartitioningTest, EinsumOutputLHSNonContractingDimPartitioned) {
  absl::string_view hlo_string = R"(
HloModule module

ENTRY entry {
  %lhs = f32[32,24,64,128] parameter(0)
  %lhs.copy = f32[32,24,64,128] copy(%lhs), sharding={replicated}
  %rhs = f32[32,39296,64,128] parameter(1)
  %rhs.copy = f32[32,39296,64,128] copy(%rhs), sharding={replicated}
  ROOT %dot = f32[32,24,39296] dot(%lhs.copy, %rhs.copy),
    lhs_batch_dims={0}, rhs_batch_dims={0},
    lhs_contracting_dims={2,3}, rhs_contracting_dims={2,3},
    sharding={devices=[1,2,1]0,1}
})";

  TF_ASSERT_OK_AND_ASSIGN(auto module,
                          PartitionComputation(hlo_string, /*num_devices=*/2));
  VLOG(1) << module->ToString();

  auto root = module->entry_computation()->root_instruction();
  auto lhs = AllOf(op::Copy(op::Parameter(0)), op::Shape("f32[32,24,64,128]"));
  auto rhs =
      AllOf(op::Copy(op::Parameter(1)), op::Shape("f32[32,39296,64,128]"));
  EXPECT_THAT(
      root,
      AllOf(op::Dot(AllOf(op::DynamicSlice(lhs, op::Constant(), op::Reshape(),
                                           op::Constant(), op::Constant()),
                          op::Shape("f32[32,12,64,128]")),
                    rhs),
            op::Shape("f32[32,12,39296]")));
}

TEST_F(SpmdPartitioningTest, EinsumOutputRHSNonContractingDimPartitioned) {
  absl::string_view hlo_string = R"(
HloModule module

ENTRY entry {
  %lhs = f32[32,24,64,128] parameter(0)
  %lhs.copy = f32[32,24,64,128] copy(%lhs), sharding={replicated}
  %rhs = f32[32,39296,64,128] parameter(1)
  %rhs.copy = f32[32,39296,64,128] copy(%rhs), sharding={replicated}
  ROOT %dot = f32[32,24,39296] dot(%lhs.copy, %rhs.copy),
    lhs_batch_dims={0}, rhs_batch_dims={0},
    lhs_contracting_dims={2,3}, rhs_contracting_dims={2,3},
    sharding={devices=[1,1,2]0,1}
})";

  TF_ASSERT_OK_AND_ASSIGN(auto module,
                          PartitionComputation(hlo_string, /*num_devices=*/2));
  VLOG(1) << module->ToString();

  auto root = module->entry_computation()->root_instruction();
  auto lhs = AllOf(op::Copy(op::Parameter(0)), op::Shape("f32[32,24,64,128]"));
  auto rhs =
      AllOf(op::Copy(op::Parameter(1)), op::Shape("f32[32,39296,64,128]"));
  EXPECT_THAT(root,
              AllOf(op::Dot(lhs, AllOf(op::DynamicSlice(
                                           rhs, op::Constant(), op::Reshape(),
                                           op::Constant(), op::Constant()),
                                       op::Shape("f32[32,19648,64,128]"))),
                    op::Shape("f32[32,24,19648]")));
}

TEST_F(SpmdPartitioningTest,
       EinsumRHSWindowedInContractingOutNonContractingPartitioned) {
  absl::string_view hlo_string = R"(
HloModule module

ENTRY entry {
  %lhs = f32[320,25,64,128] parameter(0)
  %lhs.copy = f32[320,25,64,128] copy(%lhs), sharding={devices=[1,1,4,1]0,1,2,3}
  %rhs = f32[320,39296,64,128] parameter(1)
  %rhs.copy = f32[320,39296,64,128] copy(%rhs),
    sharding={devices=[1,1,4,1]0,1,2,3}
  ROOT %dot = f32[320,25,39296] dot(%lhs.copy, %rhs.copy),
    lhs_batch_dims={0}, rhs_batch_dims={0},
    lhs_contracting_dims={2,3}, rhs_contracting_dims={2,3},
    sharding={devices=[1,4,1]0,1,2,3}
})";

  TF_ASSERT_OK_AND_ASSIGN(auto module,
                          PartitionComputation(hlo_string, /*num_devices=*/4));
  VLOG(1) << module->ToString();

  auto root = module->entry_computation()->root_instruction();
  auto lhs = AllOf(
      op::Copy(op::DynamicSlice(op::Parameter(0), op::Constant(),
                                op::Constant(), op::Reshape(), op::Constant())),
      op::Shape("f32[320,25,16,128]"));
  auto rhs = AllOf(
      op::Copy(op::DynamicSlice(op::Parameter(1), op::Constant(),
                                op::Constant(), op::Reshape(), op::Constant())),
      op::Shape("f32[320,39296,16,128]"));
  EXPECT_THAT(
      root,
      AllOf(op::GetTupleElement(op::While(op::Tuple(
                lhs, rhs, op::Broadcast(), op::Broadcast(), op::Constant()))),
            op::Shape("f32[320,7,39296]")));

  auto while_loop = root->operand(0);
  // Check loop condition.
  EXPECT_THAT(
      while_loop->while_condition()->root_instruction(),
      op::Compare(op::GetTupleElement(op::Parameter(0)), op::Constant()));

  // Check loop body.
  auto next_i = op::Add(op::GetTupleElement(op::Parameter(0)), op::Constant());
  auto ds =
      AllOf(op::DynamicSlice(
                op::Pad(op::GetTupleElement(op::Parameter(0)), op::Constant()),
                op::Constant(), op::Reshape(), op::Constant(), op::Constant()),
            op::Shape("f32[320,7,16,128]"));
  auto partial_output =
      AllOf(op::Add(op::GetTupleElement(op::Parameter(0)),
                    op::Dot(ds, op::GetTupleElement(op::Parameter(0)))),
            op::Shape("f32[320,7,39296]"));
  auto window = op::Conditional(op::Compare(next_i, op::Constant()),
                                partial_output, partial_output);
  EXPECT_THAT(while_loop->while_body()->root_instruction(),
              op::Tuple(op::GetTupleElement(op::Parameter(0)),
                        op::GetTupleElement(op::Parameter(0)), window,
                        op::GetTupleElement(op::Parameter(0)), next_i));

  // Check the conditional that contains the collective permute.
  auto cp_conditional =
      while_loop->while_body()->root_instruction()->operand(2);
  EXPECT_THAT(cp_conditional->true_computation()->root_instruction(),
              op::CollectivePermute(op::Parameter(0)));
  EXPECT_THAT(cp_conditional->false_computation()->root_instruction(),
              op::Parameter(0));
}

TEST_F(SpmdPartitioningTest,
       EinsumRHSWindowedInContractingOutNonContractingFromBroadcast) {
  absl::string_view hlo_string = R"(
HloModule module

ENTRY entry {
  %constant.1 = f32[] constant(2)
  %broadcast = f32[32,25,64,128] broadcast(%constant.1), dimensions={},
    sharding={devices=[1,1,4,1]0,1,2,3}
  %add = f32[32,25,64,128] add(%broadcast, %broadcast),
    sharding={devices=[1,1,4,1]0,1,2,3}
  %rhs = f32[32,39296,64,128] parameter(0)
  %rhs.copy = f32[32,39296,64,128] copy(%rhs),
    sharding={devices=[1,1,4,1]0,1,2,3}
  ROOT %dot = f32[32,25,39296] dot(%add, %rhs.copy),
    lhs_batch_dims={0}, rhs_batch_dims={0},
    lhs_contracting_dims={2,3}, rhs_contracting_dims={2,3},
    sharding={devices=[1,4,1]0,1,2,3}
})";

  TF_ASSERT_OK_AND_ASSIGN(auto module, PartitionComputation(hlo_string,
                                                            /*num_devices=*/4));
  VLOG(1) << module->ToString();
  // Involves loop code motion, skips pattern matching.
}

TEST_F(SpmdPartitioningTest,
       EinsumLHSWindowedInContractingOutNonContractingPartitioned) {
  absl::string_view hlo_string = R"(
HloModule module

ENTRY entry {
  %lhs = f32[16,1024,16384] parameter(0)
  %lhs.copy = f32[16,1024,16384] copy(%lhs),
    sharding={devices=[2,1,4]0,1,2,3,4,5,6,7}
  %rhs = f32[16384,67,128] parameter(1)
  %rhs.copy = f32[16384,67,128] copy(%rhs),
    sharding={devices=[4,1,1,2]0,4,1,5,2,6,3,7 last_tile_dim_replicate}
  ROOT %dot = f32[16,1024,67,128] dot(%lhs.copy, %rhs.copy),
    lhs_batch_dims={}, rhs_batch_dims={},
    lhs_contracting_dims={2}, rhs_contracting_dims={0},
    sharding={devices=[2,1,4,1]0,1,2,3,4,5,6,7}
})";

  TF_ASSERT_OK_AND_ASSIGN(auto module,
                          PartitionComputation(hlo_string, /*num_devices=*/8));
  VLOG(1) << module->ToString();

  auto root = module->entry_computation()->root_instruction();
  auto lhs = AllOf(op::Copy(op::DynamicSlice(op::Parameter(0), op::Reshape(),
                                             op::Constant(), op::Reshape())),
                   op::Shape("f32[8,1024,4096]"));
  auto rhs = AllOf(op::Copy(op::DynamicSlice(op::Parameter(1), op::Reshape(),
                                             op::Constant(), op::Constant())),
                   op::Shape("f32[4096,67,128]"));
  EXPECT_THAT(
      root,
      AllOf(op::GetTupleElement(op::While(op::Tuple(
                lhs, rhs, op::Broadcast(), op::Broadcast(), op::Constant()))),
            op::Shape("f32[8,1024,17,128]")));

  auto while_loop = root->operand(0);
  // Check loop condition.
  EXPECT_THAT(
      while_loop->while_condition()->root_instruction(),
      op::Compare(op::GetTupleElement(op::Parameter(0)), op::Constant()));

  // Check loop body.
  auto next_i = op::Add(op::GetTupleElement(op::Parameter(0)), op::Constant());
  auto ds =
      AllOf(op::DynamicSlice(
                op::Pad(op::GetTupleElement(op::Parameter(0)), op::Constant()),
                op::Constant(), op::Reshape(), op::Constant()),
            op::Shape("f32[4096,17,128]"));
  auto partial_output =
      AllOf(op::Add(op::GetTupleElement(op::Parameter(0)),
                    op::Dot(op::GetTupleElement(op::Parameter(0)), ds)),
            op::Shape("f32[8,1024,17,128]"));
  auto window = op::Conditional(op::Compare(next_i, op::Constant()),
                                partial_output, partial_output);
  EXPECT_THAT(while_loop->while_body()->root_instruction(),
              op::Tuple(op::GetTupleElement(op::Parameter(0)),
                        op::GetTupleElement(op::Parameter(0)), window,
                        op::GetTupleElement(op::Parameter(0)), next_i));

  // Check the conditional that contains the collective permute.
  auto cp_conditional =
      while_loop->while_body()->root_instruction()->operand(2);
  EXPECT_THAT(cp_conditional->true_computation()->root_instruction(),
              op::CollectivePermute(op::Parameter(0)));
  EXPECT_THAT(cp_conditional->false_computation()->root_instruction(),
              op::Parameter(0));
}

TEST_F(SpmdPartitioningTest,
       EinsumLHSWindowedInContractingOutNonContractingPartitioned2) {
  absl::string_view hlo_string = R"(
HloModule module

ENTRY entry {
  %lhs = f32[16,1024,16384] parameter(0)
  %lhs.copy = f32[16,1024,16384] copy(%lhs),
    sharding={devices=[2,1,4]0,1,2,3,4,5,6,7}
  %rhs = f32[16384,2,33,128] parameter(1)
  %rhs.copy = f32[16384,2,33,128] copy(%rhs),
    sharding={devices=[4,1,1,1,2]0,4,1,5,2,6,3,7 last_tile_dim_replicate}
  ROOT %dot = f32[16,1024,2,33,128] dot(%lhs.copy, %rhs.copy),
    lhs_batch_dims={}, rhs_batch_dims={},
    lhs_contracting_dims={2}, rhs_contracting_dims={0},
    sharding={devices=[2,1,2,2,1]0,1,2,3,4,5,6,7}
})";

  TF_ASSERT_OK_AND_ASSIGN(auto module,
                          PartitionComputation(hlo_string, /*num_devices=*/8));
  VLOG(1) << module->ToString();

  auto root = module->entry_computation()->root_instruction();
  auto lhs = AllOf(op::Copy(op::DynamicSlice(op::Parameter(0), op::Reshape(),
                                             op::Constant(), op::Reshape())),
                   op::Shape("f32[8,1024,4096]"));
  auto rhs = AllOf(
      op::Copy(op::DynamicSlice(op::Parameter(1), op::Reshape(), op::Constant(),
                                op::Constant(), op::Constant())),
      op::Shape("f32[4096,2,33,128]"));
  EXPECT_THAT(
      root,
      AllOf(op::GetTupleElement(op::While(op::Tuple(
                lhs, rhs, op::Broadcast(), op::Broadcast(), op::Constant()))),
            op::Shape("f32[8,1024,1,17,128]")));

  auto while_loop = root->operand(0);
  // Check loop condition.
  EXPECT_THAT(
      while_loop->while_condition()->root_instruction(),
      op::Compare(op::GetTupleElement(op::Parameter(0)), op::Constant()));

  // Check loop body.
  auto next_i = op::Add(op::GetTupleElement(op::Parameter(0)), op::Constant());
  auto ds =
      AllOf(op::DynamicSlice(
                op::Pad(op::GetTupleElement(op::Parameter(0)), op::Constant()),
                op::Constant(), op::Reshape(), op::Reshape(), op::Constant()),
            op::Shape("f32[4096,1,17,128]"));
  auto partial_output =
      AllOf(op::Add(op::GetTupleElement(op::Parameter(0)),
                    op::Dot(op::GetTupleElement(op::Parameter(0)), ds)),
            op::Shape("f32[8,1024,1,17,128]"));
  auto window = op::Conditional(op::Compare(next_i, op::Constant()),
                                partial_output, partial_output);
  EXPECT_THAT(while_loop->while_body()->root_instruction(),
              op::Tuple(op::GetTupleElement(op::Parameter(0)),
                        op::GetTupleElement(op::Parameter(0)), window,
                        op::GetTupleElement(op::Parameter(0)), next_i));

  // Check the conditional that contains the collective permute.
  auto cp_conditional =
      while_loop->while_body()->root_instruction()->operand(2);
  EXPECT_THAT(cp_conditional->true_computation()->root_instruction(),
              op::CollectivePermute(op::Parameter(0)));
  EXPECT_THAT(cp_conditional->false_computation()->root_instruction(),
              op::Parameter(0));
}

TEST_F(SpmdPartitioningTest, EinsumRHSWindowedNonContracting) {
  absl::string_view hlo_string = R"(
HloModule module

ENTRY entry {
  %lhs = f32[32,24,64,128] parameter(0)
  %lhs.copy = f32[32,24,64,128] copy(%lhs), sharding={devices=[1,2,1,1]0,1}
  %rhs = f32[32,39295,64,128] parameter(1)
  %rhs.copy = f32[32,39295,64,128] copy(%rhs), sharding={devices=[1,2,1,1]0,1}
  ROOT %dot = f32[32,24,39295] dot(%lhs.copy, %rhs.copy),
    lhs_batch_dims={0}, rhs_batch_dims={0},
    lhs_contracting_dims={2,3}, rhs_contracting_dims={2,3},
    sharding={devices=[1,2,1]0,1}
})";

  TF_ASSERT_OK_AND_ASSIGN(auto module, PartitionComputation(hlo_string,
                                                            /*num_devices=*/2));
  VLOG(1) << module->ToString();
  auto root = module->entry_computation()->root_instruction();
  auto lhs = AllOf(
      op::Copy(op::DynamicSlice(op::Parameter(0), op::Constant(), op::Reshape(),
                                op::Constant(), op::Constant())),
      op::Shape("f32[32,12,64,128]"));
  auto rhs =
      AllOf(op::Copy(op::DynamicSlice(op::Pad(op::Parameter(1), op::Constant()),
                                      op::Constant(), op::Reshape(),
                                      op::Constant(), op::Constant())),
            op::Shape("f32[32,19648,64,128]"));
  EXPECT_THAT(root,
              AllOf(op::Slice(AllOf(op::GetTupleElement(op::While(op::Tuple(
                                        lhs, rhs, op::Broadcast(),
                                        op::Broadcast(), op::Constant()))),
                                    op::Shape("f32[32,12,39296]"))),
                    op::Shape("f32[32,12,39295]")));
  auto while_loop = root->operand(0)->operand(0);
  // Check loop condition.
  EXPECT_THAT(
      while_loop->while_condition()->root_instruction(),
      op::Compare(op::GetTupleElement(op::Parameter(0)), op::Constant()));

  // Check loop body.
  auto next_i = op::Add(op::GetTupleElement(op::Parameter(0)), op::Constant());
  auto window = op::Conditional(op::Compare(next_i, op::Constant()),
                                op::GetTupleElement(op::Parameter(0)),
                                op::GetTupleElement(op::Parameter(0)));
  auto partial_output = op::Dot(op::GetTupleElement(op::Parameter(0)),
                                op::GetTupleElement(op::Parameter(0)));
  EXPECT_THAT(
      while_loop->while_body()->root_instruction(),
      op::Tuple(op::GetTupleElement(op::Parameter(0)), window,
                op::DynamicUpdateSlice(op::GetTupleElement(op::Parameter(0)),
                                       partial_output, op::Constant(),
                                       op::Constant(), op::Reshape()),
                op::GetTupleElement(op::Parameter(0)), next_i));

  // Check the conditional that contains the collective permute.
  auto cp_conditional =
      while_loop->while_body()->root_instruction()->operand(1);
  EXPECT_THAT(cp_conditional->true_computation()->root_instruction(),
              op::CollectivePermute(op::Parameter(0)));
  EXPECT_THAT(cp_conditional->false_computation()->root_instruction(),
              op::Parameter(0));
}

TEST_F(SpmdPartitioningTest, EinsumRHSWindowedContracting) {
  absl::string_view hlo_string = R"(
HloModule module

ENTRY entry {
  %lhs = f32[32,24,63,128] parameter(0)
  %lhs.copy = f32[32,24,63,128] copy(%lhs), sharding={devices=[1,2,1,1]0,1}
  %rhs = f32[32,39296,63,128] parameter(1)
  %rhs.copy = f32[32,39296,63,128] copy(%rhs), sharding={devices=[1,1,2,1]0,1}
  ROOT %dot = f32[32,24,39296] dot(%lhs.copy, %rhs.copy),
    lhs_batch_dims={0}, rhs_batch_dims={0},
    lhs_contracting_dims={2,3}, rhs_contracting_dims={2,3},
    sharding={devices=[1,2,1]0,1}
})";

  TF_ASSERT_OK_AND_ASSIGN(auto module, PartitionComputation(hlo_string,
                                                            /*num_devices=*/2));
  VLOG(1) << module->ToString();
  auto root = module->entry_computation()->root_instruction();
  auto lhs = AllOf(
      op::Copy(op::DynamicSlice(op::Parameter(0), op::Constant(), op::Reshape(),
                                op::Constant(), op::Constant())),
      op::Shape("f32[32,12,63,128]"));
  auto rhs =
      AllOf(op::Copy(op::DynamicSlice(op::Pad(op::Parameter(1), op::Constant()),
                                      op::Constant(), op::Constant(),
                                      op::Reshape(), op::Constant())),
            op::Shape("f32[32,39296,32,128]"));
  auto masked_rhs =
      op::Select(op::Compare(), rhs, op::Broadcast(op::Constant()));
  EXPECT_THAT(root, AllOf(op::GetTupleElement(op::While(
                              op::Tuple(lhs, masked_rhs, op::Broadcast(),
                                        op::Broadcast(), op::Constant()))),
                          op::Shape("f32[32,12,39296]")));
  auto while_loop = root->operand(0);
  // Check loop condition.
  EXPECT_THAT(
      while_loop->while_condition()->root_instruction(),
      op::Compare(op::GetTupleElement(op::Parameter(0)), op::Constant()));

  // Check loop body.
  auto next_i = op::Add(op::GetTupleElement(op::Parameter(0)), op::Constant());
  auto window = op::Conditional(op::Compare(next_i, op::Constant()),
                                op::GetTupleElement(op::Parameter(0)),
                                op::GetTupleElement(op::Parameter(0)));
  auto partial_output = op::Dot(
      op::DynamicSlice(
          op::Pad(op::GetTupleElement(op::Parameter(0)), op::Constant()),
          op::Constant(), op::Constant(), op::Reshape(), op::Constant()),
      op::GetTupleElement(op::Parameter(0)));
  EXPECT_THAT(
      while_loop->while_body()->root_instruction(),
      op::Tuple(op::GetTupleElement(op::Parameter(0)), window,
                op::Add(op::GetTupleElement(op::Parameter(0)), partial_output),
                op::GetTupleElement(op::Parameter(0)), next_i));

  // Check the conditional that contains the collective permute.
  auto cp_conditional =
      while_loop->while_body()->root_instruction()->operand(1);
  EXPECT_THAT(cp_conditional->true_computation()->root_instruction(),
              op::CollectivePermute(op::Parameter(0)));
  EXPECT_THAT(cp_conditional->false_computation()->root_instruction(),
              op::Parameter(0));
}

TEST_F(SpmdPartitioningTest, EinsumRHSWindowedNonContractingReduce1) {
  absl::string_view hlo_string = R"(
HloModule module

sum {
  a = f32[] parameter(0)
  b = f32[] parameter(1)
  ROOT add = f32[] add(a, b)
}

ENTRY entry {
  %lhs = f32[32,24,64,128] parameter(0)
  %lhs.copy = f32[32,24,64,128] copy(%lhs), sharding={devices=[1,2,1,1]0,1}
  %rhs = f32[32,39295,64,128] parameter(1)
  %rhs.copy = f32[32,39295,64,128] copy(%rhs), sharding={devices=[1,2,1,1]0,1}
  %dot = f32[32,24,39295] dot(%lhs.copy, %rhs.copy),
    lhs_batch_dims={0}, rhs_batch_dims={0},
    lhs_contracting_dims={2,3}, rhs_contracting_dims={2,3},
    sharding={devices=[1,2,1]0,1}
  %constant = f32[] constant(0)
  %constant.1 = f32[] constant(2)
  %broadcast = f32[32,24,39295] broadcast(%constant.1), dimensions={},
    sharding={devices=[1,2,1]0,1}
  %multiply = f32[32,24,39295] multiply(%dot, %broadcast),
  sharding={devices=[1,2,1]0,1}
  ROOT %reduce = f32[32,24] reduce(%multiply, %constant), dimensions={2},
    to_apply=sum, sharding={devices=[1,2]0,1}
})";

  TF_ASSERT_OK_AND_ASSIGN(auto module, PartitionComputation(hlo_string,
                                                            /*num_devices=*/2));
  VLOG(1) << module->ToString();
  // Involves loop code motion, skips pattern matching.
}

TEST_F(SpmdPartitioningTest, EinsumRHSWindowedNonContractingReduce2) {
  absl::string_view hlo_string = R"(
HloModule module

sum {
  a = f32[] parameter(0)
  b = f32[] parameter(1)
  ROOT add = f32[] add(a, b)
}

ENTRY entry {
  %lhs = f32[32,24,64,128] parameter(0)
  %lhs.copy = f32[32,24,64,128] copy(%lhs), sharding={devices=[1,2,1,1]0,1}
  %rhs = f32[32,39295,64,128] parameter(1)
  %rhs.copy = f32[32,39295,64,128] copy(%rhs), sharding={devices=[1,2,1,1]0,1}
  %dot = f32[32,24,39295] dot(%lhs.copy, %rhs.copy),
    lhs_batch_dims={0}, rhs_batch_dims={0},
    lhs_contracting_dims={2,3}, rhs_contracting_dims={2,3},
    sharding={devices=[1,2,1]0,1}
  %constant = f32[] constant(0)
  %constant.1 = f32[] constant(2)
  %broadcast = f32[32,24,39295] broadcast(%constant.1), dimensions={},
    sharding={devices=[1,2,1]0,1}
  %multiply = f32[32,24,39295] multiply(%dot, %broadcast),
    sharding={devices=[1,2,1]0,1}
  ROOT %reduce = f32[32,39295] reduce(%multiply, %constant), dimensions={1},
    to_apply=sum, sharding={replicated}
})";

  TF_ASSERT_OK_AND_ASSIGN(auto module, PartitionComputation(hlo_string,
                                                            /*num_devices=*/2));
  VLOG(1) << module->ToString();
  // Involves loop code motion, skips pattern matching.
}

TEST_F(SpmdPartitioningTest, EinsumRHSWindowedContractingFromBroadcast) {
  absl::string_view hlo_string = R"(
HloModule module

ENTRY entry {
  %rhs = f32[32,39296,63,128] parameter(0)
  %rhs.copy = f32[32,39296,63,128] copy(%rhs), sharding={devices=[1,1,2,1]0,1}
  %constant.1 = f32[] constant(2)
  %broadcast = f32[32,24,63,128] broadcast(%constant.1), dimensions={},
    sharding={devices=[1,2,1,1]0,1}
  %add = f32[32,24,63,128] add(%broadcast, %broadcast),
    sharding={devices=[1,2,1,1]0,1}
  ROOT %dot = f32[32,24,39296] dot(%add, %rhs.copy),
    lhs_batch_dims={0}, rhs_batch_dims={0},
    lhs_contracting_dims={2,3}, rhs_contracting_dims={2,3},
    sharding={devices=[1,2,1]0,1}
})";

  TF_ASSERT_OK_AND_ASSIGN(auto module, PartitionComputation(hlo_string,
                                                            /*num_devices=*/2));
  VLOG(1) << module->ToString();
  // Involves loop code motion, skips pattern matching.
}

TEST_F(SpmdPartitioningTest, ReplicatedRng) {
  absl::string_view hlo_string = R"(
HloModule module

ENTRY entry {
  %lhs = s32[] parameter(0)
  %lhs.copy = s32[] copy(%lhs), sharding={replicated}
  %rhs = s32[] parameter(1)
  %rhs.copy = s32[] copy(%rhs), sharding={replicated}
  ROOT %rng = s32[4]{0} rng(%lhs.copy, %rhs.copy),
      distribution=rng_uniform, sharding={replicated}
})";

  TF_ASSERT_OK_AND_ASSIGN(auto module,
                          PartitionComputation(hlo_string, /*num_devices=*/2));
  VLOG(1) << module->ToString();

  auto root = module->entry_computation()->root_instruction();
  auto lhs = AllOf(op::Copy(op::Parameter(0)), op::Shape("s32[]"));
  auto rhs = AllOf(op::Copy(op::Parameter(1)), op::Shape("s32[]"));
  EXPECT_THAT(
      root,
      AllOf(op::AllReduce(op::Select(
                op::Broadcast(op::Compare(op::PartitionId(), op::Constant())),
                op::Rng(), op::Broadcast(op::Constant()))),
            op::Shape("s32[4]")));
}

TEST_F(SpmdPartitioningTest, PartitionedRng) {
  absl::string_view hlo_string = R"(
HloModule module

ENTRY entry {
  %lhs = s32[] parameter(0)
  %lhs.copy = s32[] copy(%lhs), sharding={replicated}
  %rhs = s32[] parameter(1)
  %rhs.copy = s32[] copy(%rhs), sharding={maximal device=1}
  ROOT %rng = s32[4]{0} rng(%lhs.copy, %rhs.copy),
      distribution=rng_uniform, sharding={devices=[2]0,1}
})";

  TF_ASSERT_OK_AND_ASSIGN(auto module,
                          PartitionComputation(hlo_string, /*num_devices=*/2));
  VLOG(1) << module->ToString();

  auto root = module->entry_computation()->root_instruction();
  auto lhs = AllOf(op::Copy(op::Parameter(0)), op::Shape("s32[]"));
  auto rhs = AllOf(op::Copy(op::Copy(op::Parameter(1))), op::Shape("s32[]"));
  EXPECT_THAT(root, AllOf(op::Rng(lhs, op::AllReduce(op::Select(
                                           op::Broadcast(op::Compare()), rhs,
                                           op::Broadcast(op::Constant())))),
                          op::Shape("s32[2]")));
}

TEST_F(SpmdPartitioningTest, PartialReplicatedRng) {
  absl::string_view hlo_string = R"(
HloModule module

ENTRY entry {
  %lhs = s32[] parameter(0), sharding={replicated}
  %rhs = s32[] parameter(1), sharding={replicated}
  ROOT %rng = s32[8]{0} rng(%lhs, %rhs),
      distribution=rng_uniform,
      sharding={devices=[2,4]0,1,2,3,4,5,6,7 last_tile_dim_replicate}
})";

  TF_ASSERT_OK_AND_ASSIGN(auto module,
                          PartitionComputation(hlo_string, /*num_devices=*/8));
  VLOG(1) << module->ToString();

  auto root = module->entry_computation()->root_instruction();
  auto lhs = AllOf(op::Parameter(0), op::Shape("s32[]"));
  auto rhs = AllOf(op::Parameter(1), op::Shape("s32[]"));
  auto partition_id =
      AllOf(op::Reshape(op::DynamicSlice(op::Constant(), op::PartitionId())),
            op::Shape("u32[]"));
  EXPECT_THAT(
      root, AllOf(op::AllReduce(op::Select(
                      op::Broadcast(op::Compare(partition_id, op::Constant())),
                      op::Rng(lhs, rhs), op::Broadcast(op::Constant()))),
                  op::Shape("s32[4]")));
}

TEST_F(SpmdPartitioningTest, DynamicSliceAlongNonPartitionedDimension) {
  absl::string_view hlo_string = R"(
HloModule module

ENTRY entry {
  %input = s32[128,64] parameter(0)
  %input.copy = s32[128,64] copy(%input), sharding={devices=[2,1]0,1}
  %index = s32[] parameter(1)
  %constant = s32[] constant(0)
  ROOT %dynamic-slice = s32[128,2] dynamic-slice(%input.copy, %constant, %index),
    dynamic_slice_sizes={128,2}, sharding={devices=[2,1]0,1}
})";

  TF_ASSERT_OK_AND_ASSIGN(auto module,
                          PartitionComputation(hlo_string, /*num_devices=*/2));
  VLOG(1) << module->ToString();

  auto root = module->entry_computation()->root_instruction();
  auto input = AllOf(op::Copy(op::DynamicSlice(op::Parameter(0), op::Reshape(),
                                               op::Constant())),
                     op::Shape("s32[64,64]"));
  EXPECT_THAT(root,
              AllOf(op::DynamicSlice(input, op::Constant(), op::Parameter(1)),
                    op::Shape("s32[64,2]")));
}

TEST_F(SpmdPartitioningTest, DynamicUpdateSliceAlongNonPartitionedDimension) {
  absl::string_view hlo_string = R"(
HloModule module

ENTRY entry {
  %input = s32[128,64] parameter(0)
  %input.copy = s32[128,64] copy(%input), sharding={devices=[2,1]0,1}
  %index = s32[] parameter(1)
  %constant = s32[] constant(0)
  %update = s32[128,2] parameter(2)
  %update.copy = s32[128,2] copy(%update), sharding={devices=[2,1]0,1}
  ROOT %dynamic-update-slice = s32[128,64]
    dynamic-update-slice(%input.copy, %update.copy, %constant, %index),
    sharding={devices=[2,1]0,1}
})";

  TF_ASSERT_OK_AND_ASSIGN(auto module,
                          PartitionComputation(hlo_string, /*num_devices=*/2));
  VLOG(1) << module->ToString();

  auto root = module->entry_computation()->root_instruction();
  auto input = AllOf(op::Copy(op::DynamicSlice(op::Parameter(0), op::Reshape(),
                                               op::Constant())),
                     op::Shape("s32[64,64]"));
  auto update = AllOf(op::Copy(op::DynamicSlice(op::Parameter(2), op::Reshape(),
                                                op::Constant())),
                      op::Shape("s32[64,2]"));
  EXPECT_THAT(root, AllOf(op::DynamicUpdateSlice(input, update, op::Constant(),
                                                 op::Parameter(1)),
                          op::Shape("s32[64,64]")));
}

TEST_F(SpmdPartitioningTest, DynamicUpdateSliceAlongPartitionedDimension) {
  absl::string_view hlo_string = R"(
HloModule module

ENTRY entry {
  %input = s32[128,64] parameter(0)
  %input.copy = s32[128,64] copy(%input), sharding={devices=[1,2]0,1}
  %index = s32[] parameter(1)
  %constant = s32[] constant(60)
  %update = s32[128,2] parameter(2)
  %update.copy = s32[128,2] copy(%update), sharding={devices=[1,2]0,1}
  ROOT %dynamic-update-slice = s32[128,64]
    dynamic-update-slice(%input.copy, %update.copy, %index, %constant),
    sharding={devices=[1,2]0,1}
})";

  TF_ASSERT_OK_AND_ASSIGN(auto module,
                          PartitionComputation(hlo_string, /*num_devices=*/2));
  VLOG(1) << module->ToString();

  auto root = module->entry_computation()->root_instruction();
  auto input = AllOf(op::Copy(op::DynamicSlice(op::Parameter(0), op::Constant(),
                                               op::Reshape())),
                     op::Shape("s32[128,32]"));
  auto update = AllOf(op::AllReduce(op::DynamicUpdateSlice(
                          op::Broadcast(),
                          op::Copy(op::DynamicSlice(
                              op::Parameter(2), op::Constant(), op::Reshape())),
                          op::Constant(), op::Reshape())),
                      op::Shape("s32[128,2]"));

  EXPECT_THAT(
      root, AllOf(op::Select(op::Broadcast(),
                             op::DynamicUpdateSlice(
                                 input, update, op::Parameter(1), op::Select()),
                             input),
                  op::Shape("s32[128,32]")));
}

TEST_F(SpmdPartitioningTest, DynamicUpdateSlicePartitionSliceAndNonSliceDims) {
  absl::string_view hlo_string = R"(
HloModule module

ENTRY entry {
  %input = s32[128,64] parameter(0)
  %input.copy = s32[128,64] copy(%input), sharding={devices=[2,2]0,1,2,3}
  %constant.0 = s32[] constant(0)
  %constant.1 = s32[] constant(60)
  %update = s32[128,2] parameter(1)
  %update.copy = s32[128,2] copy(%update), sharding={devices=[2,2]0,1,2,3}
  ROOT %dynamic-update-slice = s32[128,64]
    dynamic-update-slice(%input.copy, %update.copy, %constant.0, %constant.1),
    sharding={devices=[2,2]0,1,2,3}
})";

  TF_ASSERT_OK_AND_ASSIGN(auto module,
                          PartitionComputation(hlo_string, /*num_devices=*/4));
  VLOG(1) << module->ToString();

  auto root = module->entry_computation()->root_instruction();
  auto input = AllOf(op::Copy(op::DynamicSlice(op::Parameter(0), op::Reshape(),
                                               op::Reshape())),
                     op::Shape("s32[64,32]"));
  auto update = AllOf(op::AllReduce(op::DynamicUpdateSlice(
                          op::Broadcast(),
                          op::Copy(op::DynamicSlice(
                              op::Parameter(1), op::Reshape(), op::Reshape())),
                          op::Constant(), op::Reshape())),
                      op::Shape("s32[64,2]"));

  EXPECT_THAT(root,
              AllOf(op::Select(op::Broadcast(),
                               op::DynamicUpdateSlice(
                                   input, update, op::Constant(), op::Select()),
                               input),
                    op::Shape("s32[64,32]")));
}

TEST_F(SpmdPartitioningTest, PassthroughGather) {
  absl::string_view hlo_string = R"(
HloModule module

ENTRY entry {
  %input = f32[2,9] parameter(0), sharding={devices=[1,2]0,1}
  %indices = s32[3] parameter(1), sharding={replicated}
  ROOT %gather = f32[3,9] gather(%input, %indices), offset_dims={1},
    collapsed_slice_dims={0}, start_index_map={0}, index_vector_dim=1,
    slice_sizes={1,9}, sharding={devices=[1,2]0,1}
})";
  TF_ASSERT_OK_AND_ASSIGN(auto module,
                          PartitionComputation(hlo_string, /*num_devices=*/2));
  VLOG(1) << module->ToString();
  HloInstruction* root = module->entry_computation()->root_instruction();
  EXPECT_THAT(root, AllOf(op::Gather(op::Parameter(0), op::Parameter(1)),
                          op::Shape("f32[3,5]")));
}

TEST_F(SpmdPartitioningTest, PassthroughGather_PartialReplicate) {
  absl::string_view hlo_string = R"(
HloModule module

ENTRY entry {
  %input = f32[2,9] parameter(0),
    sharding={devices=[1,2,2]0,1,2,3 last_tile_dim_replicate}
  %indices = s32[3] parameter(1), sharding={replicated}
  ROOT %gather = f32[3,9] gather(%input, %indices), offset_dims={1},
    collapsed_slice_dims={0}, start_index_map={0}, index_vector_dim=1,
    slice_sizes={1,9}, sharding={devices=[1,2,2]0,1,2,3 last_tile_dim_replicate}
})";
  TF_ASSERT_OK_AND_ASSIGN(auto module,
                          PartitionComputation(hlo_string, /*num_devices=*/4));
  VLOG(1) << module->ToString();
  HloInstruction* root = module->entry_computation()->root_instruction();
  EXPECT_THAT(root, AllOf(op::Gather(op::Parameter(0), op::Parameter(1)),
                          op::Shape("f32[3,5]")));
}

TEST_F(SpmdPartitioningTest, IndexPassthroughGather) {
  absl::string_view hlo_string = R"(
HloModule module

ENTRY entry {
  %input = f32[2,9,8] parameter(0), sharding={replicated}
  %indices = s32[4,2,4] parameter(1), sharding={devices=[2,1,2]0,1,2,3}
  ROOT %gather = f32[8,4,4] gather(%input, %indices), offset_dims={0},
    collapsed_slice_dims={0,1}, start_index_map={0,1}, index_vector_dim=1,
    slice_sizes={1,1,8}, sharding={devices=[1,2,2]0,1,2,3}
})";
  TF_ASSERT_OK_AND_ASSIGN(auto module,
                          PartitionComputation(hlo_string, /*num_devices=*/4));
  VLOG(1) << module->ToString();
  HloInstruction* root = module->entry_computation()->root_instruction();
  EXPECT_THAT(root, AllOf(op::Gather(op::Parameter(0), op::Parameter(1)),
                          op::Shape("f32[8,2,2]")));
}

TEST_F(SpmdPartitioningTest, IndexPassthroughGather_PartialReplicate) {
  absl::string_view hlo_string = R"(
HloModule module

ENTRY entry {
  %input = f32[2,9,8] parameter(0), sharding={replicated}
  %indices = s32[4,2,4] parameter(1),
    sharding={devices=[2,1,2,2]0,1,2,3,4,5,6,7 last_tile_dim_replicate}
  ROOT %gather = f32[8,4,4] gather(%input, %indices), offset_dims={0},
    collapsed_slice_dims={0,1}, start_index_map={0,1}, index_vector_dim=1,
    slice_sizes={1,1,8},
    sharding={devices=[1,2,2,2]0,1,2,3,4,5,6,7 last_tile_dim_replicate}
})";
  TF_ASSERT_OK_AND_ASSIGN(auto module,
                          PartitionComputation(hlo_string, /*num_devices=*/8));
  VLOG(1) << module->ToString();
  HloInstruction* root = module->entry_computation()->root_instruction();
  EXPECT_THAT(root, AllOf(op::Gather(op::Parameter(0), op::Parameter(1)),
                          op::Shape("f32[8,2,2]")));
}

TEST_F(SpmdPartitioningTest, GatherPartitionedOnTrivialSliceDims) {
  absl::string_view hlo_string = R"(
HloModule module

ENTRY entry {
  %input = f32[17,9] parameter(0), sharding={devices=[2,1]0,1}
  %indices = s32[2,3] parameter(1), sharding={replicated}
  ROOT %gather = f32[2,3,9] gather(%input, %indices), offset_dims={2},
    collapsed_slice_dims={0}, start_index_map={0}, index_vector_dim=2,
    slice_sizes={1,9}, sharding={replicated}
})";
  TF_ASSERT_OK_AND_ASSIGN(auto module,
                          PartitionComputation(hlo_string, /*num_devices=*/2));
  VLOG(1) << module->ToString();
  auto offset =
      op::Reshape(op::DynamicSlice(op::Constant(), op::PartitionId()));
  auto min = AllOf(op::Broadcast(offset), op::Shape("s32[2,3]"));
  auto max = AllOf(op::Broadcast(op::Add(offset, op::Constant())),
                   op::Shape("s32[2,3]"));
  auto clamp = op::Clamp(min, op::Parameter(1), max);
  auto gather = op::Gather(op::Parameter(0), op::Subtract(clamp, min));
  auto mask =
      op::Or(op::Lt(op::Parameter(1), min), op::Gt(op::Parameter(1), max));
  auto masked =
      op::Select(op::Broadcast(mask), op::Broadcast(op::Constant()), gather);
  HloInstruction* root = module->entry_computation()->root_instruction();
  EXPECT_THAT(root, AllOf(op::AllReduce(masked), op::Shape("f32[2,3,9]")));
}

TEST_F(SpmdPartitioningTest,
       GatherPartitionedOnTrivialSliceDims_PartialReplicate) {
  absl::string_view hlo_string = R"(
HloModule module

ENTRY entry {
  %input = f32[17,9] parameter(0),
    sharding={devices=[2,1,2]0,1,2,3 last_tile_dim_replicate}
  %indices = s32[2,3] parameter(1), sharding={replicated}
  ROOT %gather = f32[2,3,9] gather(%input, %indices), offset_dims={2},
    collapsed_slice_dims={0}, start_index_map={0}, index_vector_dim=2,
    slice_sizes={1,9}, sharding={replicated}
})";
  TF_ASSERT_OK_AND_ASSIGN(auto module,
                          PartitionComputation(hlo_string, /*num_devices=*/4));
  VLOG(1) << module->ToString();
  auto offset =
      op::Reshape(op::DynamicSlice(op::Constant(), op::PartitionId()));
  auto min = AllOf(op::Broadcast(offset), op::Shape("s32[2,3]"));
  auto max = AllOf(op::Broadcast(op::Add(offset, op::Constant())),
                   op::Shape("s32[2,3]"));
  auto clamp = op::Clamp(min, op::Parameter(1), max);
  auto gather = op::Gather(op::Parameter(0), op::Subtract(clamp, min));
  auto mask =
      op::Or(op::Lt(op::Parameter(1), min), op::Gt(op::Parameter(1), max));
  auto masked =
      op::Select(op::Broadcast(mask), op::Broadcast(op::Constant()), gather);
  HloInstruction* root = module->entry_computation()->root_instruction();
  EXPECT_THAT(root, AllOf(op::AllReduce(masked), op::Shape("f32[2,3,9]")));
}

TEST_F(SpmdPartitioningTest, PassthroughScatter) {
  absl::string_view hlo_string = R"(
HloModule module

add (lhs: f32[], rhs: f32[]) -> f32[] {
  lhs = f32[] parameter(0)
  rhs = f32[] parameter(1)
  ROOT sum = f32[] add(lhs, rhs)
}

ENTRY entry {
  %input = f32[2,9] parameter(0), sharding={devices=[1,2]0,1}
  %indices = s32[3] parameter(1), sharding={replicated}
  %updates = f32[3,9] parameter(2), sharding={devices=[1,2]0,1}
  ROOT %scatter = f32[2,9] scatter(%input, %indices, %updates),
      to_apply=add,
      update_window_dims={1},
      inserted_window_dims={0},
      scatter_dims_to_operand_dims={0},
      index_vector_dim=1, sharding={devices=[1,2]0,1}
})";
  TF_ASSERT_OK_AND_ASSIGN(auto module,
                          PartitionComputation(hlo_string, /*num_devices=*/2));
  VLOG(1) << module->ToString();
  HloInstruction* root = module->entry_computation()->root_instruction();
  EXPECT_THAT(root, AllOf(op::Scatter(op::Parameter(0), op::Parameter(1),
                                      op::Parameter(2)),
                          op::Shape("f32[2,5]")));
}

TEST_F(SpmdPartitioningTest, PassthroughScatter_PartialReplicate) {
  absl::string_view hlo_string = R"(
HloModule module

add (lhs: f32[], rhs: f32[]) -> f32[] {
  lhs = f32[] parameter(0)
  rhs = f32[] parameter(1)
  ROOT sum = f32[] add(lhs, rhs)
}

ENTRY entry {
  %input = f32[2,9] parameter(0),
    sharding={devices=[1,2,2]0,1,2,3 last_tile_dim_replicate}
  %indices = s32[3] parameter(1), sharding={replicated}
  %updates = f32[3,9] parameter(2),
    sharding={devices=[1,2,2]0,1,2,3 last_tile_dim_replicate}
  ROOT %scatter = f32[2,9] scatter(%input, %indices, %updates),
      to_apply=add,
      update_window_dims={1},
      inserted_window_dims={0},
      scatter_dims_to_operand_dims={0},
      index_vector_dim=1,
      sharding={devices=[1,2,2]0,1,2,3 last_tile_dim_replicate}
})";
  TF_ASSERT_OK_AND_ASSIGN(auto module,
                          PartitionComputation(hlo_string, /*num_devices=*/4));
  VLOG(1) << module->ToString();
  HloInstruction* root = module->entry_computation()->root_instruction();
  EXPECT_THAT(root, AllOf(op::Scatter(op::Parameter(0), op::Parameter(1),
                                      op::Parameter(2)),
                          op::Shape("f32[2,5]")));
}

TEST_F(SpmdPartitioningTest, IndexPassthroughScatter) {
  absl::string_view hlo_string = R"(
HloModule module

add (lhs: f32[], rhs: f32[]) -> f32[] {
  lhs = f32[] parameter(0)
  rhs = f32[] parameter(1)
  ROOT sum = f32[] add(lhs, rhs)
}

ENTRY entry {
  %input = f32[2,9,8] parameter(0), sharding={replicated}
  %indices = s32[4,2,4] parameter(1), sharding={devices=[2,1,2]0,1,2,3}
  %updates = f32[4,4,8] parameter(2), sharding={devices=[2,2,1]0,1,2,3}
  ROOT %scatter = f32[2,9,8] scatter(%input, %indices, %updates),
      to_apply=add,
      update_window_dims={2},
      inserted_window_dims={0,1},
      scatter_dims_to_operand_dims={0,1},
      index_vector_dim=1, sharding={replicated}
})";
  TF_ASSERT_OK_AND_ASSIGN(auto module,
                          PartitionComputation(hlo_string, /*num_devices=*/4));
  VLOG(1) << module->ToString();
  HloInstruction* root = module->entry_computation()->root_instruction();
  EXPECT_THAT(
      root,
      AllOf(op::AllReduce(op::AllReduce(op::Scatter(
                op::Select(op::Broadcast(op::Convert(op::PartitionId())),
                           op::Broadcast(op::Constant()), op::Parameter(0)),
                op::Parameter(1), op::Parameter(2)))),
            op::Shape("f32[2,9,8]")));
}

TEST_F(SpmdPartitioningTest, IndexPassthroughScatter_PartialReplicate) {
  absl::string_view hlo_string = R"(
HloModule module

add (lhs: f32[], rhs: f32[]) -> f32[] {
  lhs = f32[] parameter(0)
  rhs = f32[] parameter(1)
  ROOT sum = f32[] add(lhs, rhs)
}

ENTRY entry {
  %input = f32[2,9,8] parameter(0), sharding={replicated}
  %indices = s32[4,2,4] parameter(1),
    sharding={devices=[2,1,2,2]0,1,2,3,4,5,6,7 last_tile_dim_replicate}
  %updates = f32[4,4,8] parameter(2),
    sharding={devices=[2,2,1,2]0,1,2,3,4,5,6,7 last_tile_dim_replicate}
  ROOT %scatter = f32[2,9,8] scatter(%input, %indices, %updates),
      to_apply=add,
      update_window_dims={2},
      inserted_window_dims={0,1},
      scatter_dims_to_operand_dims={0,1},
      index_vector_dim=1, sharding={replicated}
})";
  TF_ASSERT_OK_AND_ASSIGN(auto module,
                          PartitionComputation(hlo_string, /*num_devices=*/8));
  VLOG(1) << module->ToString();
  HloInstruction* root = module->entry_computation()->root_instruction();
  EXPECT_THAT(
      root,
      AllOf(op::AllReduce(op::AllReduce(op::Scatter(
                op::Select(op::Broadcast(op::Convert(op::Reshape())),
                           op::Broadcast(op::Constant()), op::Parameter(0)),
                op::Parameter(1), op::Parameter(2)))),
            op::Shape("f32[2,9,8]")));
}

TEST_F(SpmdPartitioningTest, IndexPassthroughScatter_Min) {
  absl::string_view hlo_string = R"(
HloModule module

min (lhs: f32[], rhs: f32[]) -> f32[] {
  lhs = f32[] parameter(0)
  rhs = f32[] parameter(1)
  ROOT min = f32[] minimum(lhs, rhs)
}

ENTRY entry {
  %input = f32[2,9,8] parameter(0), sharding={replicated}
  %indices = s32[4,2,4] parameter(1), sharding={devices=[2,1,2]0,1,2,3}
  %updates = f32[4,4,8] parameter(2), sharding={devices=[2,2,1]0,1,2,3}
  ROOT %scatter = f32[2,9,8] scatter(%input, %indices, %updates),
      to_apply=min,
      update_window_dims={2},
      inserted_window_dims={0,1},
      scatter_dims_to_operand_dims={0,1},
      index_vector_dim=1, sharding={replicated}
})";
  TF_ASSERT_OK_AND_ASSIGN(auto module,
                          PartitionComputation(hlo_string, /*num_devices=*/4));
  VLOG(1) << module->ToString();
  HloInstruction* root = module->entry_computation()->root_instruction();
  EXPECT_THAT(
      root,
      AllOf(op::AllReduce(op::AllReduce(op::Scatter(
                op::Select(op::Broadcast(op::Convert(op::PartitionId())),
                           op::Broadcast(op::Constant()), op::Parameter(0)),
                op::Parameter(1), op::Parameter(2)))),
            op::Shape("f32[2,9,8]")));
}

TEST_F(SpmdPartitioningTest, ScatterPartitionedOnTrivialSliceDims) {
  absl::string_view hlo_string = R"(
HloModule module

add (lhs: f32[], rhs: f32[]) -> f32[] {
  lhs = f32[] parameter(0)
  rhs = f32[] parameter(1)
  ROOT sum = f32[] add(lhs, rhs)
}

ENTRY entry {
  %input = f32[17,9] parameter(0), sharding={devices=[2,1]0,1}
  %indices = s32[2,3] parameter(1), sharding={replicated}
  %updates = f32[2,3,9] parameter(2), sharding={replicated}
  ROOT %scatter = f32[17,9] scatter(%input, %indices, %updates),
      to_apply=add,
      update_window_dims={2},
      inserted_window_dims={0},
      scatter_dims_to_operand_dims={0},
      index_vector_dim=2, sharding={devices=[2,1]0,1}
})";
  TF_ASSERT_OK_AND_ASSIGN(auto module,
                          PartitionComputation(hlo_string, /*num_devices=*/2));
  VLOG(1) << module->ToString();
  auto offset =
      op::Reshape(op::DynamicSlice(op::Constant(), op::PartitionId()));
  auto indices = op::Subtract(
      op::Parameter(1), AllOf(op::Broadcast(offset), op::Shape("s32[2,3]")));
  HloInstruction* root = module->entry_computation()->root_instruction();
  EXPECT_THAT(root,
              AllOf(op::Scatter(op::Parameter(0), indices, op::Parameter(2)),
                    op::Shape("f32[9,9]")));
}

TEST_F(SpmdPartitioningTest,
       ScatterPartitionedOnTrivialSliceDims_PartialReplicate) {
  absl::string_view hlo_string = R"(
HloModule module

add (lhs: f32[], rhs: f32[]) -> f32[] {
  lhs = f32[] parameter(0)
  rhs = f32[] parameter(1)
  ROOT sum = f32[] add(lhs, rhs)
}

ENTRY entry {
  %input = f32[17,9] parameter(0),
    sharding={devices=[2,1,2]0,1,2,3 last_tile_dim_replicate}
  %indices = s32[2,3] parameter(1), sharding={replicated}
  %updates = f32[2,3,9] parameter(2), sharding={replicated}
  ROOT %scatter = f32[17,9] scatter(%input, %indices, %updates),
      to_apply=add,
      update_window_dims={2},
      inserted_window_dims={0},
      scatter_dims_to_operand_dims={0},
      index_vector_dim=2,
      sharding={devices=[2,1,2]0,1,2,3 last_tile_dim_replicate}
})";
  TF_ASSERT_OK_AND_ASSIGN(auto module,
                          PartitionComputation(hlo_string, /*num_devices=*/4));
  VLOG(1) << module->ToString();
  auto offset =
      op::Reshape(op::DynamicSlice(op::Constant(), op::PartitionId()));
  auto indices = op::Subtract(
      op::Parameter(1), AllOf(op::Broadcast(offset), op::Shape("s32[2,3]")));
  HloInstruction* root = module->entry_computation()->root_instruction();
  EXPECT_THAT(root,
              AllOf(op::Scatter(op::Parameter(0), indices, op::Parameter(2)),
                    op::Shape("f32[9,9]")));
}

TEST_F(SpmdPartitioningTest, TiledReversePassthrough) {
  absl::string_view hlo_string = R"(
HloModule module

ENTRY entry {
  constant = f32[3,3]{1,0} constant({{1,1,1},{1,1,1},{1,1,1}}),
    sharding={devices=[2,1]0,1}
  ROOT reverse = f32[3,3]{1,0} reverse(constant), dimensions={1},
    sharding={devices=[2,1]0,1}
})";
  TF_ASSERT_OK_AND_ASSIGN(auto module,
                          PartitionComputation(hlo_string, /*num_devices=*/2));
  VLOG(1) << module->ToString();
  HloInstruction* root = module->entry_computation()->root_instruction();
  EXPECT_THAT(root, AllOf(op::Shape("f32[2,3]{1,0}"),
                          op::Reverse(op::DynamicSlice(
                              op::Pad(op::Constant(), op::Constant()),
                              op::Reshape(), op::Constant()))));
}

TEST_F(SpmdPartitioningTest, TiledReversePassthroughViaReversedSharding) {
  absl::string_view hlo_string = R"(
HloModule module

ENTRY entry {
  param = f32[4] parameter(0), sharding={devices=[2]0,1}
  ROOT reverse = f32[4] reverse(param), dimensions={0},
    sharding={devices=[2]1,0}
})";
  TF_ASSERT_OK_AND_ASSIGN(auto module,
                          PartitionComputation(hlo_string, /*num_devices=*/2));
  VLOG(1) << module->ToString();
  HloInstruction* root = module->entry_computation()->root_instruction();
  EXPECT_THAT(root, AllOf(op::Shape("f32[2]"), op::Reverse(op::Parameter(0))));
}

TEST_F(SpmdPartitioningTest, TiledReverseSwapShards) {
  absl::string_view hlo_string = R"(
HloModule module

ENTRY entry {
  param = f32[4] parameter(0), sharding={devices=[2]0,1}
  ROOT reverse = f32[4] reverse(param), dimensions={0},
    sharding={devices=[2]0,1}
})";
  TF_ASSERT_OK_AND_ASSIGN(auto module,
                          PartitionComputation(hlo_string, /*num_devices=*/2));
  VLOG(1) << module->ToString();
  HloInstruction* root = module->entry_computation()->root_instruction();
  EXPECT_THAT(root,
              AllOf(op::Shape("f32[2]"),
                    op::Reverse(op::CollectivePermute(op::Parameter(0)))));
}

TEST_F(SpmdPartitioningTest, TiledReverseHaloExchange) {
  absl::string_view hlo_string = R"(
HloModule module

ENTRY entry {
  param = f32[3] parameter(0), sharding={devices=[2]0,1}
  ROOT reverse = f32[3] reverse(param), dimensions={0},
    sharding={devices=[2]1,0}
})";
  TF_ASSERT_OK_AND_ASSIGN(auto module,
                          PartitionComputation(hlo_string, /*num_devices=*/2));
  VLOG(1) << module->ToString();
  HloInstruction* root = module->entry_computation()->root_instruction();
  auto halo_exchange_concat =
      op::Concatenate(AllOf(op::Shape("f32[1]"),
                            op::CollectivePermute(op::Slice(op::Parameter(0)))),
                      op::Parameter(0));
  auto after_halo_exchange = op::Slice(halo_exchange_concat);
  EXPECT_THAT(root,
              AllOf(op::Shape("f32[2]"), op::Reverse(after_halo_exchange)));
}

TEST_F(SpmdPartitioningTest, MixWithManualPartitioning) {
  absl::string_view hlo_string = R"(
HloModule module

ENTRY entry {
  param = (f32[8,2], f32[4,2]) parameter(0), sharding={{devices=[2,1]0,1},{manual}}
  param0 = f32[8,2] get-tuple-element(param), index=0, sharding={devices=[2,1]0,1}
  param1 = f32[4,2] get-tuple-element(param), index=1, sharding={manual}
  to_shard = f32[4,2] custom-call(param0), custom_call_target="SPMDFullToShardShape", sharding={manual}
  add = f32[4,2] add(to_shard, param1), sharding={manual}
  to_full = f32[8,2] custom-call(add), custom_call_target="SPMDShardToFullShape", sharding={devices=[2,1]0,1}
  mul = f32[8,2] multiply(to_full, param0), sharding={devices=[2,1]0,1}
  to_shard2 = f32[4,2] custom-call(mul), custom_call_target="SPMDFullToShardShape", sharding={manual}
  ROOT tuple = (f32[4,2]) tuple(to_shard2), sharding={{manual}}
})";
  TF_ASSERT_OK_AND_ASSIGN(auto module,
                          PartitionComputation(hlo_string, /*num_devices=*/2));
  VLOG(1) << module->ToString();
  HloInstruction* root = module->entry_computation()->root_instruction();
  auto p0 = op::GetTupleElement(op::Parameter(0));
  auto to_shard = op::Copy(p0);
  auto p1 = op::GetTupleElement(op::Parameter(0));
  auto mul = AllOf(op::Shape("f32[4,2]"),
                   op::Multiply(op::Copy(op::Add(to_shard, p1)), p0));
  EXPECT_THAT(root, op::Tuple(op::Copy(mul)));
}

TEST_F(SpmdPartitioningTest, SubgroupAllToAllReshard) {
  absl::string_view hlo_string = R"(
HloModule module

ENTRY entry {
  %param0 = f32[8,8,8,8] parameter(0),
    sharding={devices=[2,2,1,2]0,1,2,3,4,5,6,7}
  ROOT %copy = f32[8,8,8,8] copy(%param0),
    sharding={devices=[1,2,2,2]0,1,4,5,2,3,6,7}
})";

  TF_ASSERT_OK_AND_ASSIGN(auto module,
                          PartitionComputation(hlo_string, /*num_devices=*/8));
  VLOG(1) << module->ToString();

  auto root = module->entry_computation()->root_instruction();
  auto reshape =
      AllOf(op::Shape("f32[4,4,2,4,4]"), op::Reshape(op::Parameter(0)));
  auto all_to_all = AllOf(op::Shape("f32[4,4,2,4,4]"), op::AllToAll(reshape));
  auto xpose = AllOf(op::Shape("f32[2,4,4,4,4]"), op::Transpose(all_to_all));
  EXPECT_THAT(root,
              op::Copy(AllOf(op::Reshape(xpose), op::Shape("f32[8,4,4,4]"))));
  EXPECT_EQ(root->operand(0)->operand(0)->operand(0)->replica_groups().size(),
            4);
}

TEST_F(SpmdPartitioningTest, SubgroupAllToAllReshard2) {
  absl::string_view hlo_string = R"(
HloModule module

ENTRY entry {
  %param0 = f32[8,8] parameter(0),
    sharding={devices=[2,4]0,1,2,3,4,5,6,7}
  ROOT %copy = f32[8,8] copy(%param0),
    sharding={devices=[4,2]0,1,4,5,2,3,6,7}
})";

  TF_ASSERT_OK_AND_ASSIGN(auto module,
                          PartitionComputation(hlo_string, /*num_devices=*/8));
  VLOG(1) << module->ToString();

  auto root = module->entry_computation()->root_instruction();
  auto all_to_all = op::AllToAll(
      AllOf(op::Shape("f32[2,2,2]"), op::Reshape(op::Parameter(0))));
  auto reshape =
      AllOf(op::Shape("f32[2,4]"), op::Reshape(op::Transpose(all_to_all)));
  EXPECT_THAT(root, op::Copy(op::CollectivePermute(reshape)));
}

TEST_F(SpmdPartitioningTest, SubgroupAllToAllReshard3) {
  absl::string_view hlo_string = R"(
HloModule module

ENTRY entry {
  %param0 = f32[8,8,8] parameter(0),
    sharding={devices=[2,4,1]0,1,2,3,4,5,6,7}
  ROOT %copy = f32[8,8,8] copy(%param0),
    sharding={devices=[1,2,4]0,1,4,5,2,3,6,7}
})";

  TF_ASSERT_OK_AND_ASSIGN(auto module,
                          PartitionComputation(hlo_string, /*num_devices=*/8));
  VLOG(1) << module->ToString();

  auto root = module->entry_computation()->root_instruction();
  auto all_to_all = op::AllToAll(
      AllOf(op::Shape("f32[4,2,4,2]"), op::Reshape(op::Parameter(0))));
  auto reshape =
      AllOf(op::Shape("f32[4,8,2]"), op::Reshape(op::Transpose(all_to_all)));
  auto all_to_all2 =
      op::AllToAll(AllOf(op::Shape("f32[4,2,4,2]"), op::Reshape(reshape)));
  auto reshape2 =
      AllOf(op::Shape("f32[8,4,2]"), op::Reshape(op::Transpose(all_to_all2)));
  EXPECT_THAT(root, op::Copy(op::CollectivePermute(reshape2)));
}

TEST_F(SpmdPartitioningTest, Dot2DPartitionedNonContractingAndContracting0) {
  absl::string_view hlo_string = R"(
HloModule module

ENTRY entry {
  %lhs = f32[48,12] parameter(0), sharding={devices=[2,2]0,1,2,3}
  %rhs = f32[32,12] parameter(1), sharding={devices=[2,2]0,2,1,3}
  ROOT %dot = f32[48,32] dot(%lhs, %rhs),
    lhs_batch_dims={}, rhs_batch_dims={},
    lhs_contracting_dims={1}, rhs_contracting_dims={1},
    sharding={devices=[2,2]0,1,2,3}
})";

  TF_ASSERT_OK_AND_ASSIGN(auto module,
                          PartitionComputation(hlo_string, /*num_devices=*/4));
  VLOG(1) << module->ToString();

  auto lhs = AllOf(op::Shape("f32[24,6]"), op::Parameter(0));
  auto partial_replicated_lhs =
      AllOf(op::Shape("f32[24,12]"),
            op::AllReduce(op::DynamicUpdateSlice(_, lhs, _, _)));
  auto rhs = AllOf(op::Shape("f32[16,6]"), op::Parameter(1));
  auto partial_replicated_rhs =
      AllOf(op::Shape("f32[16,12]"),
            op::AllReduce(op::DynamicUpdateSlice(_, rhs, _, _)));
  auto root = module->entry_computation()->root_instruction();
  EXPECT_THAT(root,
              AllOf(op::Dot(partial_replicated_lhs, partial_replicated_rhs),
                    op::Shape("f32[24,16]")));
}

TEST_F(SpmdPartitioningTest, Dot2DPartitionedNonContractingAndContracting1) {
  absl::string_view hlo_string = R"(
HloModule module

ENTRY entry {
  %lhs = f32[48,100] parameter(0), sharding={devices=[2,2]0,1,2,3}
  %rhs = f32[32,100] parameter(1), sharding={devices=[2,2]0,1,2,3}
  ROOT %dot = f32[48,32] dot(%lhs, %rhs),
    lhs_batch_dims={}, rhs_batch_dims={},
    lhs_contracting_dims={1}, rhs_contracting_dims={1},
    sharding={devices=[2,2]0,1,2,3}
})";

  TF_ASSERT_OK_AND_ASSIGN(auto module,
                          PartitionComputation(hlo_string, /*num_devices=*/4));
  VLOG(1) << module->ToString();

  auto lhs = AllOf(op::Shape("f32[24,50]"), op::Parameter(0));
  auto rhs = AllOf(op::Shape("f32[16,50]"), op::Parameter(1));
  auto partial_replicated_rhs =
      AllOf(op::Shape("f32[32,50]"),
            op::AllReduce(op::DynamicUpdateSlice(_, rhs, _, _)));
  auto root = module->entry_computation()->root_instruction();
  EXPECT_THAT(
      root, AllOf(op::Shape("f32[24,16]"),
                  op::DynamicSlice(
                      op::AllReduce(AllOf(op::Dot(lhs, partial_replicated_rhs),
                                          op::Shape("f32[24,32]"))),
                      _, _)));
}

TEST_F(SpmdPartitioningTest, Dot2DPartitionedNonContractingAndContracting2) {
  absl::string_view hlo_string = R"(
HloModule module

ENTRY entry {
  %lhs = f32[48,100] parameter(0), sharding={replicated}
  %rhs = f32[32,100] parameter(1), sharding={devices=[2,2]0,1,2,3}
  ROOT %dot = f32[48,32] dot(%lhs, %rhs),
    lhs_batch_dims={}, rhs_batch_dims={},
    lhs_contracting_dims={1}, rhs_contracting_dims={1},
    sharding={devices=[2,2]0,1,2,3}
})";

  TF_ASSERT_OK_AND_ASSIGN(auto module,
                          PartitionComputation(hlo_string, /*num_devices=*/4));
  VLOG(1) << module->ToString();

  auto lhs = AllOf(op::Shape("f32[48,100]"), op::Parameter(0));
  auto lhs_slice = AllOf(op::Shape("f32[24,100]"), op::DynamicSlice(lhs, _, _));
  auto rhs = AllOf(op::Shape("f32[16,50]"), op::Parameter(1));
  auto partial_replicated_rhs = AllOf(
      op::Shape("f32[16,100]"), op::AllReduce(op::DynamicUpdateSlice(
                                    _, op::CollectivePermute(rhs), _, _)));
  auto root = module->entry_computation()->root_instruction();
  EXPECT_THAT(root, AllOf(op::Shape("f32[24,16]"),
                          op::Dot(lhs_slice, partial_replicated_rhs)));
}

TEST_F(SpmdPartitioningTest, Dot2DPartitionedNoncontractingAndContracting3) {
  absl::string_view hlo_string = R"(
HloModule module

ENTRY entry {
  %lhs = f32[23,24] parameter(0), sharding={devices=[2,1,2]0,1,2,3 last_tile_dim_replicate}
  %rhs = f32[23,32] parameter(1), sharding={devices=[2,2]0,1,2,3}
  ROOT %dot = f32[24,32] dot(%lhs, %rhs),
    lhs_contracting_dims={0}, rhs_contracting_dims={0},
    sharding={devices=[2,2]1,0,3,2}
})";

  TF_ASSERT_OK_AND_ASSIGN(auto module,
                          PartitionComputation(hlo_string, /*num_devices=*/4));
  VLOG(1) << module->ToString();

  auto lhs = AllOf(op::Shape("f32[12,24]"), op::Parameter(0));
  auto masked_lhs = op::Select(_, lhs, op::Broadcast(op::Constant()));
  auto rhs = AllOf(op::Shape("f32[12,16]"), op::Parameter(1));
  auto masked_rhs = op::Select(_, rhs, op::Broadcast(op::Constant()));
  auto root = module->entry_computation()->root_instruction();
  EXPECT_THAT(
      root,
      AllOf(op::Shape("f32[12,16]"),
            op::DynamicSlice(
                AllOf(op::Shape("f32[24,16]"),
                      op::AllReduce(op::Dot(
                          masked_lhs, op::CollectivePermute(masked_rhs)))),
                _, _)));
}

TEST_F(SpmdPartitioningTest, Dot2DPartitionedBatchAndNonContracting) {
  absl::string_view hlo_string = R"(
HloModule module

ENTRY entry {
  %lhs = f32[4,24,100] parameter(0), sharding={devices=[2,2,1]0,1,2,3}
  %rhs = f32[4,32,100] parameter(1), sharding={devices=[2,2,1]0,1,2,3}
  ROOT %dot = f32[4,24,32] dot(%lhs, %rhs),
    lhs_batch_dims={0}, rhs_batch_dims={0},
    lhs_contracting_dims={2}, rhs_contracting_dims={2},
    sharding={devices=[2,2,1]0,1,2,3}
})";

  TF_ASSERT_OK_AND_ASSIGN(auto module,
                          PartitionComputation(hlo_string, /*num_devices=*/4));
  VLOG(1) << module->ToString();

  auto lhs = AllOf(op::Shape("f32[2,12,100]"), op::Parameter(0));
  auto rhs = AllOf(op::Shape("f32[2,16,100]"), op::Parameter(1));
  auto partial_replicated_rhs =
      AllOf(op::Shape("f32[2,32,100]"),
            op::AllReduce(op::DynamicUpdateSlice(_, rhs, _, _, _)));
  auto root = module->entry_computation()->root_instruction();
  EXPECT_THAT(root, AllOf(op::Shape("f32[2,12,32]"),
                          op::Dot(lhs, partial_replicated_rhs)));
}

TEST_F(SpmdPartitioningTest, Dot2DPartitionedBatchAndContracting) {
  absl::string_view hlo_string = R"(
HloModule module

ENTRY entry {
  %lhs = f32[4,24,100] parameter(0), sharding={devices=[2,1,2]0,1,2,3}
  %rhs = f32[4,32,100] parameter(1), sharding={devices=[1,2,2]0,1,2,3}
  ROOT %dot = f32[4,24,32] dot(%lhs, %rhs),
    lhs_batch_dims={0}, rhs_batch_dims={0},
    lhs_contracting_dims={2}, rhs_contracting_dims={2},
    sharding={devices=[2,2,1]0,1,2,3}
})";

  TF_ASSERT_OK_AND_ASSIGN(auto module,
                          PartitionComputation(hlo_string, /*num_devices=*/4));
  VLOG(1) << module->ToString();

  auto lhs = AllOf(op::Shape("f32[2,24,50]"), op::Parameter(0));
  auto rhs = AllOf(op::Shape("f32[4,16,50]"), op::Parameter(1));
  auto resharded_rhs =
      AllOf(op::Shape("f32[2,32,50]"),
            op::Reshape(op::Transpose(op::AllToAll(op::Reshape(rhs)))));
  auto root = module->entry_computation()->root_instruction();
  EXPECT_THAT(root, AllOf(op::Shape("f32[2,12,32]"),
                          op::DynamicSlice(
                              AllOf(op::Shape("f32[2,24,32]"),
                                    op::AllReduce(op::Dot(lhs, resharded_rhs))),
                              _, _, _)));
}

TEST_F(SpmdPartitioningTest, Dot2DPartitionedBatchAndContracting2) {
  absl::string_view hlo_string = R"(
HloModule module

ENTRY entry {
  %lhs = f32[4,24,100] parameter(0), sharding={devices=[2,1,2]0,1,2,3}
  %rhs = f32[4,32,100] parameter(1), sharding={replicated}
  ROOT %dot = f32[4,24,32] dot(%lhs, %rhs),
    lhs_batch_dims={0}, rhs_batch_dims={0},
    lhs_contracting_dims={2}, rhs_contracting_dims={2},
    sharding={devices=[2,2,1]0,1,2,3}
})";

  TF_ASSERT_OK_AND_ASSIGN(auto module,
                          PartitionComputation(hlo_string, /*num_devices=*/4));
  VLOG(1) << module->ToString();

  auto lhs = AllOf(op::Shape("f32[2,24,50]"), op::Parameter(0));
  auto resharded_lhs =
      AllOf(op::Shape("f32[2,12,100]"),
            op::Reshape(op::Transpose(op::AllToAll(op::Reshape(lhs)))));
  auto rhs = AllOf(op::Shape("f32[4,32,100]"), op::Parameter(1));
  auto rhs_slice =
      AllOf(op::Shape("f32[2,32,100]"), op::DynamicSlice(rhs, _, _, _));
  auto root = module->entry_computation()->root_instruction();
  EXPECT_THAT(root, AllOf(op::Shape("f32[2,12,32]"),
                          op::Dot(resharded_lhs, rhs_slice)));
}

TEST_F(SpmdPartitioningTest,
       Dot2DPartitionedBatchNonContractingAndContracting) {
  absl::string_view hlo_string = R"(
HloModule module

ENTRY entry {
  %lhs = f32[4,24,100] parameter(0), sharding={devices=[2,1,2]0,1,2,3}
  %rhs = f32[4,32,100] parameter(1), sharding={devices=[2,2,1]0,1,2,3}
  ROOT %dot = f32[4,24,32] dot(%lhs, %rhs),
    lhs_batch_dims={0}, rhs_batch_dims={0},
    lhs_contracting_dims={2}, rhs_contracting_dims={2},
    sharding={devices=[2,1,2]0,1,2,3}
})";

  TF_ASSERT_OK_AND_ASSIGN(auto module,
                          PartitionComputation(hlo_string, /*num_devices=*/4));
  VLOG(1) << module->ToString();

  auto lhs = AllOf(op::Shape("f32[2,24,50]"), op::Parameter(0));
  auto rhs = AllOf(op::Shape("f32[2,16,100]"), op::Parameter(1));
  auto partial_replicated_lhs =
      AllOf(op::Shape("f32[2,24,100]"),
            op::AllReduce(op::DynamicUpdateSlice(_, lhs, _, _, _)));
  auto root = module->entry_computation()->root_instruction();
  EXPECT_THAT(root, AllOf(op::Shape("f32[2,24,16]"),
                          op::Dot(partial_replicated_lhs, rhs)));
}

TEST_F(SpmdPartitioningTest, Dot2DPartitionedBatchAndReshard) {
  absl::string_view hlo_string = R"(
HloModule module

ENTRY entry {
  %lhs = f32[4,8,24,100] parameter(0), sharding={devices=[2,1,2,1]0,1,2,3}
  %rhs = f32[4,8,32,100] parameter(1), sharding={devices=[2,1,2,1]0,1,2,3}
  ROOT %dot = f32[4,8,24,32] dot(%lhs, %rhs),
    lhs_batch_dims={0,1}, rhs_batch_dims={0,1},
    lhs_contracting_dims={3}, rhs_contracting_dims={3},
    sharding={devices=[1,2,2,1]0,1,2,3}
})";

  TF_ASSERT_OK_AND_ASSIGN(auto module,
                          PartitionComputation(hlo_string, /*num_devices=*/4));
  VLOG(1) << module->ToString();

  auto lhs = AllOf(op::Shape("f32[2,8,12,100]"), op::Parameter(0));
  auto rhs = AllOf(op::Shape("f32[2,8,16,100]"), op::Parameter(1));
  auto partial_replicated_rhs =
      AllOf(op::Shape("f32[2,8,32,100]"),
            op::AllReduce(op::DynamicUpdateSlice(_, rhs, _, _, _, _)));
  auto dot =
      AllOf(op::Shape("f32[2,8,12,32]"), op::Dot(lhs, partial_replicated_rhs));
  auto reshape = AllOf(op::Shape("f32[2,2,4,12,32]"), op::Reshape(dot));
  auto all_to_all = AllOf(op::Shape("f32[2,2,4,12,32]"), op::AllToAll(reshape));
  auto xpose = AllOf(op::Shape("f32[2,2,4,12,32]"), op::Transpose(all_to_all));
  auto root = module->entry_computation()->root_instruction();
  EXPECT_THAT(root, AllOf(op::Shape("f32[4,4,12,32]"), op::Reshape(xpose)));
}

TEST_F(SpmdPartitioningTest, SimpleDotPartial) {
  absl::string_view hlo_string = R"(
HloModule module

ENTRY entry {
  %lhs = f32[2,24,100] parameter(0),
    sharding={devices=[2,1,1,2]0,1,2,3 last_tile_dim_replicate}
  %rhs = f32[2,32,100] parameter(1),
    sharding={devices=[2,1,1,2]0,1,2,3 last_tile_dim_replicate}
  ROOT %dot = f32[2,24,32] dot(%lhs, %rhs),
    lhs_batch_dims={0}, rhs_batch_dims={0},
    lhs_contracting_dims={2}, rhs_contracting_dims={2},
    sharding={devices=[2,1,1,2]0,1,2,3 last_tile_dim_replicate}
})";

  TF_ASSERT_OK_AND_ASSIGN(auto module,
                          PartitionComputation(hlo_string, /*num_devices=*/4));
  VLOG(1) << module->ToString();

  auto lhs = AllOf(op::Shape("f32[1,24,100]"), op::Parameter(0));
  auto rhs = AllOf(op::Shape("f32[1,32,100]"), op::Parameter(1));
  auto dot = AllOf(op::Shape("f32[1,24,32]"), op::Dot(lhs, rhs));
  auto root = module->entry_computation()->root_instruction();
  EXPECT_THAT(root, dot);
}

TEST_F(SpmdPartitioningTest, DotPartialContracting) {
  absl::string_view hlo_string = R"(
HloModule module

ENTRY entry {
  %lhs = f32[24,100] parameter(0),
    sharding={devices=[1,2,2]0,1,2,3 last_tile_dim_replicate}
  %rhs = f32[32,100] parameter(1),
    sharding={devices=[1,2,2]0,1,2,3 last_tile_dim_replicate}
  ROOT %dot = f32[24,32] dot(%lhs, %rhs),
    lhs_batch_dims={}, rhs_batch_dims={},
    lhs_contracting_dims={1}, rhs_contracting_dims={1},
    sharding={replicated}
})";

  TF_ASSERT_OK_AND_ASSIGN(auto module,
                          PartitionComputation(hlo_string, /*num_devices=*/4));
  VLOG(1) << module->ToString();

  auto lhs = AllOf(op::Shape("f32[24,50]"), op::Parameter(0));
  auto rhs = AllOf(op::Shape("f32[32,50]"), op::Parameter(1));
  auto dot = AllOf(op::Shape("f32[24,32]"), op::Dot(lhs, rhs));
  auto root = module->entry_computation()->root_instruction();
  EXPECT_THAT(root, op::AllReduce(dot));
}

TEST_F(SpmdPartitioningTest, DotPartialContracting2) {
  absl::string_view hlo_string = R"(
HloModule module

ENTRY entry {
  %lhs = f32[24,100] parameter(0),
    sharding={devices=[1,2,2]0,1,2,3 last_tile_dim_replicate}
  %rhs = f32[32,100] parameter(1),
    sharding={devices=[1,2,2]0,1,2,3 last_tile_dim_replicate}
  ROOT %dot = f32[24,32] dot(%lhs, %rhs),
    lhs_batch_dims={}, rhs_batch_dims={},
    lhs_contracting_dims={1}, rhs_contracting_dims={1},
    sharding={devices=[2,1,2]0,2,1,3 last_tile_dim_replicate}
})";

  TF_ASSERT_OK_AND_ASSIGN(auto module,
                          PartitionComputation(hlo_string, /*num_devices=*/4));
  VLOG(1) << module->ToString();

  auto lhs = AllOf(op::Shape("f32[24,50]"), op::Parameter(0));
  auto rhs = AllOf(op::Shape("f32[32,50]"), op::Parameter(1));
  auto dot =
      AllOf(op::Shape("f32[12,32]"),
            op::Dot(AllOf(op::Shape("f32[12,50]"), op::DynamicSlice(lhs, _, _)),
                    rhs));
  auto root = module->entry_computation()->root_instruction();
  EXPECT_THAT(root, op::AllReduce(dot));
}

TEST_F(SpmdPartitioningTest, DotPartialContracting3) {
  absl::string_view hlo_string = R"(
HloModule module

ENTRY entry {
  %lhs = f32[24,100] parameter(0),
    sharding={devices=[1,2,4]0,1,2,3,4,5,6,7 last_tile_dim_replicate}
  %rhs = f32[32,100] parameter(1),
    sharding={devices=[1,2,4]0,1,2,3,4,5,6,7 last_tile_dim_replicate}
  ROOT %dot = f32[24,32] dot(%lhs, %rhs),
    lhs_batch_dims={}, rhs_batch_dims={},
    lhs_contracting_dims={1}, rhs_contracting_dims={1},
    sharding={devices=[1,2,4]0,1,2,3,4,5,6,7 last_tile_dim_replicate}
})";

  TF_ASSERT_OK_AND_ASSIGN(auto module,
                          PartitionComputation(hlo_string, /*num_devices=*/8));
  VLOG(1) << module->ToString();

  auto lhs = AllOf(op::Shape("f32[24,50]"), op::Parameter(0));
  auto rhs =
      AllOf(op::Shape("f32[16,50]"), op::DynamicSlice(op::Parameter(1), _, _));
  auto dot = AllOf(op::Shape("f32[24,16]"), op::Dot(lhs, rhs));
  auto root = module->entry_computation()->root_instruction();
  EXPECT_THAT(root, op::CollectivePermute(op::AllReduce(dot)));
}

TEST_F(SpmdPartitioningTest, DotBatchAndPartialContracting) {
  absl::string_view hlo_string = R"(
HloModule module

ENTRY entry {
  %lhs = f32[4,24,100] parameter(0),
    sharding={devices=[2,2,2]0,1,2,3,4,5,6,7}
  %rhs = f32[4,32,100] parameter(1),
    sharding={devices=[2,1,2,2]0,2,1,3,4,6,5,7 last_tile_dim_replicate}
  ROOT %dot = f32[4,24,32] dot(%lhs, %rhs),
    lhs_batch_dims={0}, rhs_batch_dims={0},
    lhs_contracting_dims={2}, rhs_contracting_dims={2},
    sharding={devices=[2,2,1,2]0,1,2,3,4,5,6,7 last_tile_dim_replicate}
})";

  TF_ASSERT_OK_AND_ASSIGN(auto module,
                          PartitionComputation(hlo_string, /*num_devices=*/8));
  VLOG(1) << module->ToString();

  auto lhs = AllOf(op::Shape("f32[2,12,50]"), op::Parameter(0));
  auto rhs = AllOf(op::Shape("f32[2,32,50]"), op::Parameter(1));
  auto dot = AllOf(op::Shape("f32[2,12,32]"), op::Dot(lhs, rhs));
  auto root = module->entry_computation()->root_instruction();
  EXPECT_THAT(root, op::AllReduce(dot));
}

TEST_F(SpmdPartitioningTest, DotPartialNonContracting) {
  absl::string_view hlo_string = R"(
HloModule module

ENTRY entry {
  %lhs = f32[24,8,100] parameter(0),
    sharding={devices=[2,1,1,2]0,1,2,3 last_tile_dim_replicate}
  %rhs = f32[32,100] parameter(1), sharding={devices=[2,2]0,2,1,3}
  ROOT %dot = f32[24,8,32] dot(%lhs, %rhs),
    lhs_batch_dims={}, rhs_batch_dims={},
    lhs_contracting_dims={2}, rhs_contracting_dims={1},
    sharding={devices=[2,1,2]0,1,2,3}
})";

  TF_ASSERT_OK_AND_ASSIGN(auto module,
                          PartitionComputation(hlo_string, /*num_devices=*/4));
  VLOG(1) << module->ToString();

  auto lhs = AllOf(op::Shape("f32[12,8,100]"), op::Parameter(0));
  auto rhs = AllOf(op::Shape("f32[16,50]"), op::Parameter(1));
  auto partially_replicated_rhs =
      AllOf(op::Shape("f32[16,100]"),
            op::AllReduce(op::DynamicUpdateSlice(op::Broadcast(_), rhs, _, _)));
  auto dot =
      AllOf(op::Shape("f32[12,8,16]"), op::Dot(lhs, partially_replicated_rhs));
  auto root = module->entry_computation()->root_instruction();
  EXPECT_THAT(root, dot);
}

TEST_F(SpmdPartitioningTest, DotPartialNonContractingPartialMatch) {
  absl::string_view hlo_string = R"(
HloModule module

ENTRY entry {
  %lhs = f32[24,8,100] parameter(0), sharding={devices=[2,2,1]0,1,2,3}
  %rhs = f32[32,100] parameter(1),
    sharding={devices=[2,1,2]0,2,1,3 last_tile_dim_replicate}
  ROOT %dot = f32[24,8,32] dot(%lhs, %rhs),
    lhs_batch_dims={}, rhs_batch_dims={},
    lhs_contracting_dims={2}, rhs_contracting_dims={1},
    sharding={devices=[2,1,2]0,1,2,3}
})";

  TF_ASSERT_OK_AND_ASSIGN(auto module,
                          PartitionComputation(hlo_string, /*num_devices=*/4));
  VLOG(1) << module->ToString();

  auto lhs = AllOf(op::Shape("f32[12,4,100]"), op::Parameter(0));
  auto rhs = AllOf(op::Shape("f32[16,100]"), op::Parameter(1));
  auto partially_replicated_lhs = AllOf(
      op::Shape("f32[12,8,100]"),
      op::AllReduce(op::DynamicUpdateSlice(op::Broadcast(_), lhs, _, _, _)));
  auto dot =
      AllOf(op::Shape("f32[12,8,16]"), op::Dot(partially_replicated_lhs, rhs));
  auto root = module->entry_computation()->root_instruction();
  EXPECT_THAT(root, dot);
}

TEST_F(SpmdPartitioningTest, DotPartialContractingPartialMatch) {
  absl::string_view hlo_string = R"(
HloModule module

ENTRY entry {
  %lhs = f32[24,8,100] parameter(0), sharding={devices=[1,2,2]0,1,2,3}
  %rhs = f32[32,8,100] parameter(1),
    sharding={devices=[1,1,2,2]0,2,1,3 last_tile_dim_replicate}
  ROOT %dot = f32[24,32] dot(%lhs, %rhs),
    lhs_batch_dims={}, rhs_batch_dims={},
    lhs_contracting_dims={1,2}, rhs_contracting_dims={1,2},
    sharding={replicated}
})";

  TF_ASSERT_OK_AND_ASSIGN(auto module,
                          PartitionComputation(hlo_string, /*num_devices=*/4));
  VLOG(1) << module->ToString();

  auto lhs = AllOf(op::Shape("f32[24,4,50]"), op::Parameter(0));
  auto rhs = AllOf(op::Shape("f32[32,8,50]"), op::Parameter(1));
  auto dot = AllOf(op::Shape("f32[24,32]"),
                   op::Dot(lhs, AllOf(op::Shape("f32[32,4,50]"),
                                      op::DynamicSlice(rhs, _, _, _))));
  auto root = module->entry_computation()->root_instruction();
  EXPECT_THAT(root, op::AllReduce(op::AllReduce(dot)));
}

TEST_F(SpmdPartitioningTest, DotNonContractingPartialMatchContractingMatch) {
  absl::string_view hlo_string = R"(
HloModule module

ENTRY entry {
  %lhs = f32[24,8,100] parameter(0), sharding={devices=[2,1,2]0,1,2,3}
  %rhs = f32[100,50] parameter(1), sharding={devices=[2,2]0,2,1,3}
  ROOT %dot = f32[24,8,50] dot(%lhs, %rhs),
    lhs_batch_dims={}, rhs_batch_dims={},
    lhs_contracting_dims={2}, rhs_contracting_dims={0},
    sharding={devices=[2,2,1]0,1,2,3}
})";

  TF_ASSERT_OK_AND_ASSIGN(auto module,
                          PartitionComputation(hlo_string, /*num_devices=*/4));
  VLOG(1) << module->ToString();

  auto lhs = AllOf(op::Shape("f32[12,8,50]"), op::Parameter(0));
  auto rhs = AllOf(op::Shape("f32[50,25]"), op::Parameter(1));
  auto dot = AllOf(
      op::Shape("f32[12,8,50]"),
      op::Dot(lhs, AllOf(op::Shape("f32[50,50]"),
                         op::AllReduce(op::DynamicUpdateSlice(_, rhs, _, _)))));
  auto root = module->entry_computation()->root_instruction();
  EXPECT_THAT(root, AllOf(op::Shape("f32[12,4,50]"),
                          op::DynamicSlice(op::AllReduce(dot), _, _, _)))
      << module->ToString();
}

TEST_F(SpmdPartitioningTest, DotLHSMutiNonContractingRHSNotMatch) {
  absl::string_view hlo_string = R"(
HloModule module

ENTRY entry {
  %lhs = f32[24,8,10] parameter(0), sharding={devices=[2,2,1]0,1,2,3}
  %rhs = f32[10,50] parameter(1),
    sharding={devices=[2,1,2]0,2,1,3 last_tile_dim_replicate}
  ROOT %dot = f32[24,8,50] dot(%lhs, %rhs),
    lhs_batch_dims={}, rhs_batch_dims={},
    lhs_contracting_dims={2}, rhs_contracting_dims={0},
    sharding={devices=[2,2,1]0,1,2,3}
})";

  TF_ASSERT_OK_AND_ASSIGN(auto module,
                          PartitionComputation(hlo_string, /*num_devices=*/4));
  VLOG(1) << module->ToString();

  auto lhs = AllOf(op::Shape("f32[12,4,10]"), op::Parameter(0));
  auto rhs = AllOf(op::Shape("f32[5,50]"), op::Parameter(1));
  auto dot = AllOf(
      op::Shape("f32[12,4,50]"),
      op::Dot(lhs, AllOf(op::Shape("f32[10,50]"),
                         op::AllReduce(op::DynamicUpdateSlice(_, rhs, _, _)))));
  auto root = module->entry_computation()->root_instruction();
  EXPECT_THAT(root, dot) << module->ToString();
}

TEST_F(SpmdPartitioningTest,
       ElementwiseTest_PartialReplicateToTiledHaloExchange) {
  absl::string_view hlo_string = R"(
HloModule module

ENTRY entry {
  constant = f32[6,3]{1,0}
    constant({{1,3,7},{5,1,4},{1,2,8},{2,3,7},{5,2,4},{2,2,8}}),
    sharding={replicated}
  constant.1 = f32[6,3]{1,0}
    constant({{2,7,2},{2,9,2},{2,6,2},{3,7,2},{2,9,3},{2,3,2}}),
    sharding={replicated}
  multiply = f32[6,3]{1,0} multiply(constant, constant.1),
    sharding={devices=[2,1,2]0,1,2,3 last_tile_dim_replicate}
  ROOT add = f32[6,3]{1,0} add(multiply, constant.1),
    sharding={devices=[4,1]0,1,2,3}
}
)";

  TF_ASSERT_OK_AND_ASSIGN(auto module,
                          PartitionComputation(hlo_string, /*num_devices=*/4));
  VLOG(1) << module->ToString();
  auto partial_replicate_lhs =
      AllOf(op::Shape("f32[3,3]"),
            op::DynamicSlice(op::Constant(), op::Reshape(), op::Constant()));
  auto partial_replicate_rhs =
      AllOf(op::Shape("f32[3,3]"),
            op::DynamicSlice(op::Constant(), op::Reshape(), op::Constant()));
  auto multiply =
      AllOf(op::Shape("f32[3,3]"),
            op::Multiply(partial_replicate_lhs, partial_replicate_rhs));
  auto right_halo =
      AllOf(op::Shape("f32[1,3]"), op::CollectivePermute(op::Slice(multiply)));
  auto add_lhs = AllOf(
      op::Shape("f32[2,3]"),
      op::DynamicSlice(
          op::DynamicSlice(
              op::Pad(op::Concatenate(multiply, right_halo), op::Constant()),
              op::Reshape(), op::Constant()),
          op::Subtract(), op::Subtract()));
  auto add_rhs = AllOf(op::Shape("f32[2,3]"),
                       op::DynamicSlice(op::Pad(op::Constant(), op::Constant()),
                                        op::Reshape(), op::Constant()));
  auto root = module->entry_computation()->root_instruction();
  EXPECT_THAT(root, AllOf(op::Shape("f32[2,3]"), op::Add(add_lhs, add_rhs)));
}

TEST_F(SpmdPartitioningTest, TileToPartialReplicateReshard) {
  absl::string_view hlo_string = R"(
HloModule module

ENTRY entry {
  %param0 = f32[8,8] parameter(0)
  %copy = f32[8,8] copy(%param0),
    sharding={devices=[2,2]0,1,2,3}
  ROOT %copy0 = f32[8,8] copy(%copy),
    sharding={devices=[2,1,2]0,1,2,3 last_tile_dim_replicate}
})";

  TF_ASSERT_OK_AND_ASSIGN(auto module,
                          PartitionComputation(hlo_string, /*num_devices=*/4));
  VLOG(1) << module->ToString();
  auto tiled = AllOf(op::Shape("f32[4,4]"),
                     op::Copy(op::DynamicSlice(op::Parameter(0), op::Reshape(),
                                               op::Reshape())));
  auto partially_replicated = AllOf(
      op::Shape("f32[4,8]"), op::Copy(op::AllReduce(op::DynamicUpdateSlice(
                                 op::Broadcast(_), tiled, _, _))));
  auto root = module->entry_computation()->root_instruction();
  EXPECT_THAT(root, partially_replicated);
}

TEST_F(SpmdPartitioningTest, TileToPartialReplicateReshardUnevenPartition) {
  absl::string_view hlo_string = R"(
HloModule module

ENTRY entry {
  %param0 = f32[8,8] parameter(0),
    sharding={devices=[2,3]0,1,2,3,4,5}
  ROOT %copy0 = f32[8,8] copy(%param0),
    sharding={devices=[1,2,3]0,1,2,3,4,5 last_tile_dim_replicate}
})";

  TF_ASSERT_OK_AND_ASSIGN(auto module,
                          PartitionComputation(hlo_string, /*num_devices=*/6));
  VLOG(1) << module->ToString();
  auto tiled = AllOf(op::Shape("f32[4,3]"), op::Parameter(0));
  auto partially_replicated = AllOf(
      op::Shape("f32[8,4]"),
      op::Copy(op::Reshape(
          op::Transpose(op::AllToAll(op::Reshape(op::Slice(op::AllReduce(
              op::DynamicUpdateSlice(op::Broadcast(), tiled, _, _)))))))));
  auto root = module->entry_computation()->root_instruction();
  EXPECT_THAT(root, partially_replicated);
}

TEST_F(SpmdPartitioningTest, PartialReplicateToTileReshardUnevenPartition) {
  absl::string_view hlo_string = R"(
HloModule module

ENTRY entry {
  %param0 = f32[8,8] parameter(0),
    sharding={devices=[1,2,3]0,1,2,3,4,5 last_tile_dim_replicate}
  ROOT %copy0 = f32[8,8] copy(%param0),
    sharding={devices=[2,3]0,1,2,3,4,5}
})";

  TF_ASSERT_OK_AND_ASSIGN(auto module,
                          PartitionComputation(hlo_string, /*num_devices=*/6));
  VLOG(1) << module->ToString();
  auto partial_replicated = AllOf(op::Shape("f32[8,4]"), op::Parameter(0));
  auto tiled = AllOf(
      op::Shape("f32[4,3]"),
      op::Copy(op::DynamicSlice(op::Pad(op::Reshape(op::Transpose(op::AllToAll(
                                            op::Reshape(partial_replicated)))),
                                        _),
                                _, _)));
  auto root = module->entry_computation()->root_instruction();
  EXPECT_THAT(root, tiled);
}

TEST_F(SpmdPartitioningTest, PartialReplicateToTileReshard) {
  absl::string_view hlo_string = R"(
HloModule module

ENTRY entry {
  %param0 = f32[8,8] parameter(0)
  %copy = f32[8,8] copy(%param0),
    sharding={devices=[2,1,2]0,1,2,3 last_tile_dim_replicate}
  ROOT %copy0 = f32[8,8] copy(%copy),
    sharding={devices=[2,2]0,1,2,3}
})";

  TF_ASSERT_OK_AND_ASSIGN(auto module,
                          PartitionComputation(hlo_string, /*num_devices=*/4));
  VLOG(1) << module->ToString();
  auto partially_replicated =
      AllOf(op::Shape("f32[4,8]"),
            op::Copy(op::DynamicSlice(op::Parameter(0), op::Reshape(),
                                      op::Constant())));
  auto tiled =
      AllOf(op::Shape("f32[4,4]"),
            op::Copy(op::DynamicSlice(partially_replicated, op::Subtract(),
                                      op::Subtract())));
  auto root = module->entry_computation()->root_instruction();
  EXPECT_THAT(root, tiled);
}

TEST_F(SpmdPartitioningTest,
       PartialReplicateToPartialReplicateReshard_AllReduce) {
  absl::string_view hlo_string = R"(
HloModule module

ENTRY entry {
  %param0 = f32[8,8] parameter(0)
  %copy = f32[8,8] copy(param0),
    sharding={devices=[2,2,2]0,1,2,3,4,5,6,7 last_tile_dim_replicate}
  ROOT %copy0 = f32[8,8] copy(%copy),
    sharding={devices=[2,1,4]0,1,2,3,4,5,6,7 last_tile_dim_replicate}
})";

  TF_ASSERT_OK_AND_ASSIGN(auto module,
                          PartitionComputation(hlo_string, /*num_devices=*/8));

  VLOG(1) << module->ToString();
  auto partially_replicated_init =
      AllOf(op::Shape("f32[4,4]"),
            op::Copy(op::DynamicSlice(op::Parameter(0), op::Reshape(),
                                      op::Reshape())));
  auto partially_replicated =
      AllOf(op::Shape("f32[4,8]"),
            op::Copy(op::AllReduce(op::DynamicUpdateSlice(
                op::Broadcast(_), partially_replicated_init, _, _))));
  auto root = module->entry_computation()->root_instruction();
  EXPECT_THAT(root, partially_replicated);
}

TEST_F(SpmdPartitioningTest,
       PartialReplicateToPartialReplicateReshard_DynamicSlice) {
  absl::string_view hlo_string = R"(
HloModule module

ENTRY entry {
  %param0 = f32[8,8] parameter(0)
  %copy = f32[8,8] copy(%param0),
    sharding={devices=[2,1,4]0,1,2,3,4,5,6,7 last_tile_dim_replicate}
  ROOT %copy0 = f32[8,8] copy(%copy),
    sharding={devices=[2,2,2]0,1,2,3,4,5,6,7 last_tile_dim_replicate}
})";

  TF_ASSERT_OK_AND_ASSIGN(auto module,
                          PartitionComputation(hlo_string, /*num_devices=*/8));
  VLOG(1) << module->ToString();
  auto partially_replicated =
      AllOf(op::Shape("f32[4,8]"),
            op::Copy(op::DynamicSlice(op::Parameter(0), op::Reshape(),
                                      op::Constant())));
  auto tiled =
      AllOf(op::Shape("f32[4,4]"),
            op::Copy(op::DynamicSlice(partially_replicated, op::Subtract(),
                                      op::Subtract())));
  auto root = module->entry_computation()->root_instruction();
  EXPECT_THAT(root, tiled);
}

TEST_F(SpmdPartitioningTest,
       PartialReplicateToPartialReplicateReshardWithCollectivePermute) {
  absl::string_view hlo_string = R"(
HloModule module

ENTRY entry {
  %param0 = f32[8,8] parameter(0)
  %copy = f32[8,8] copy(param0),
    sharding={devices=[2,2,2]0,1,2,3,4,5,6,7 last_tile_dim_replicate}
  ROOT %copy0 = f32[8,8] copy(%copy),
    sharding={devices=[1,2,4]0,1,2,3,4,5,6,7 last_tile_dim_replicate}
})";

  TF_ASSERT_OK_AND_ASSIGN(auto module,
                          PartitionComputation(hlo_string, /*num_devices=*/8));

  VLOG(1) << module->ToString();
  auto partially_replicated_init =
      AllOf(op::Shape("f32[4,4]"),
            op::CollectivePermute(op::Copy(op::DynamicSlice(
                op::Parameter(0), op::Reshape(), op::Reshape()))));
  auto partially_replicated =
      AllOf(op::Shape("f32[8,4]"),
            op::Copy(op::AllReduce(op::DynamicUpdateSlice(
                op::Broadcast(_), partially_replicated_init, _, _))));
  auto root = module->entry_computation()->root_instruction();
  EXPECT_THAT(root, partially_replicated);
}

TEST_F(SpmdPartitioningTest,
       PartialReplicateToPartialReplicateReshardCollectivePermute1) {
  absl::string_view hlo_string = R"(
HloModule module

ENTRY entry {
  %param0 = f32[8,8] parameter(0)
  %copy = f32[8,8] copy(%param0),
    sharding={devices=[1,2,4]0,1,2,3,4,5,6,7 last_tile_dim_replicate}
  ROOT %copy0 = f32[8,8] copy(%copy),
    sharding={devices=[2,2,2]0,1,2,3,4,5,6,7 last_tile_dim_replicate}
})";

  TF_ASSERT_OK_AND_ASSIGN(auto module,
                          PartitionComputation(hlo_string, /*num_devices=*/8));
  VLOG(1) << module->ToString();
  auto partially_replicated =
      AllOf(op::Shape("f32[8,4]"),
            op::Copy(op::DynamicSlice(op::Parameter(0), op::Constant(),
                                      op::Reshape())));
  auto tiled =
      AllOf(op::Shape("f32[4,4]"),
            op::Copy(op::CollectivePermute(op::DynamicSlice(
                partially_replicated, op::Subtract(), op::Subtract()))));
  auto root = module->entry_computation()->root_instruction();
  EXPECT_THAT(root, tiled);
}

TEST_F(SpmdPartitioningTest,
       PartialReplicateToPartialReplicateReshardHaloExchange) {
  absl::string_view hlo_string = R"(
HloModule module

ENTRY entry {
  %param0 = f32[6,3] parameter(0)
  %copy = f32[6,3] copy(param0),
    sharding={devices=[4,1,2]0,1,2,3,4,5,6,7 last_tile_dim_replicate}
  ROOT %copy0 = f32[6,3] copy(%copy),
    sharding={devices=[2,1,4]0,1,2,3,4,5,6,7 last_tile_dim_replicate}
})";

  TF_ASSERT_OK_AND_ASSIGN(auto module,
                          PartitionComputation(hlo_string, /*num_devices=*/8));

  VLOG(1) << module->ToString();
  auto partially_replicated_init =
      AllOf(op::Shape("f32[2,3]"),
            op::Copy(op::DynamicSlice(op::Pad(op::Parameter(0), op::Constant()),
                                      op::Reshape(), op::Constant())));
  auto slice =
      AllOf(op::Shape("f32[2,3]"),
            op::DynamicSlice(op::Concatenate(op::CollectivePermute(op::Slice(
                                                 partially_replicated_init)),
                                             partially_replicated_init),
                             _, _));
  auto partially_replicated =
      AllOf(op::Shape("f32[3,3]"),
            op::Copy(op::Slice(op::AllReduce(
                op::DynamicUpdateSlice(op::Broadcast(_), slice, _, _)))));
  auto root = module->entry_computation()->root_instruction();
  EXPECT_THAT(root, partially_replicated);
}

TEST_F(SpmdPartitioningTest,
       PartialReplicateToPartialReplicateReshardHaloExchange1) {
  absl::string_view hlo_string = R"(
HloModule module

ENTRY entry {
  %param0 = f32[6,3] parameter(0)
  %copy = f32[6,3] copy(param0),
    sharding={devices=[2,1,4]0,1,2,3,4,5,6,7 last_tile_dim_replicate}
  ROOT %copy0 = f32[6,3] copy(%copy),
    sharding={devices=[4,1,2]0,1,2,3,4,5,6,7 last_tile_dim_replicate}
})";

  TF_ASSERT_OK_AND_ASSIGN(auto module,
                          PartitionComputation(hlo_string, /*num_devices=*/8));

  VLOG(1) << module->ToString();
  auto partially_replicated_init =
      AllOf(op::Shape("f32[3,3]"),
            op::Copy(op::DynamicSlice(op::Parameter(0), op::Reshape(),
                                      op::Constant())));
  auto slice = AllOf(
      op::Shape("f32[4,3]"),
      op::DynamicSlice(op::Pad(op::Concatenate(partially_replicated_init,
                                               op::CollectivePermute(op::Slice(
                                                   partially_replicated_init))),
                               op::Constant()),
                       _, _));
  auto partially_replicated =
      AllOf(op::Shape("f32[2,3]"), op::Copy(op::DynamicSlice(slice, _, _)));
  auto root = module->entry_computation()->root_instruction();
  EXPECT_THAT(root, partially_replicated);
}

TEST_F(SpmdPartitioningTest, PartitionConvWithBathGroupCount) {
  absl::string_view hlo_string = R"(
HloModule module

ENTRY entry {
  %lhs = f32[16,801,1,1024] parameter(0)
  %lhs.copy = f32[16,801,1,1024] copy(%lhs),
    sharding={devices=[1,1,1,2]0,1}
  %rhs = f32[16,801,1,1024] parameter(1)
  %rhs.copy = f32[16,801,1,1024] copy(%rhs),
    sharding={devices=[1,1,1,2]0,1}
  ROOT %conv = f32[5,1,1,1024] convolution(%lhs.copy, %rhs.copy),
    dim_labels=f01b_i01o->01bf,batch_group_count=1024,
    window={size=801x1 pad=2_2x0_0},
    sharding={devices=[1,1,1,2]0,1}
})";

  TF_ASSERT_OK_AND_ASSIGN(auto module,
                          PartitionComputation(hlo_string, /*num_devices=*/2));

  VLOG(1) << module->ToString();
  auto root = module->entry_computation()->root_instruction();
  auto lhs = AllOf(
      op::Copy(op::DynamicSlice(op::Parameter(), op::Constant(), op::Constant(),
                                op::Constant(), op::Reshape())),
      op::Shape("f32[16,801,1,512]"));
  auto rhs = AllOf(
      op::Copy(op::DynamicSlice(op::Parameter(), op::Constant(), op::Constant(),
                                op::Constant(), op::Reshape())),
      op::Shape("f32[16,801,1,512]"));
  EXPECT_THAT(root,
              AllOf(op::Convolution(lhs, rhs), op::Shape("f32[5,1,1,512]")));
}

TEST_F(SpmdPartitioningTest, PartitionConvWithBathGroupCountRHSAlignWithLHS) {
  absl::string_view hlo_string = R"(
HloModule module

ENTRY entry {
  %lhs = f32[16,801,1,1024] parameter(0)
  %lhs.copy = f32[16,801,1,1024] copy(%lhs),
    sharding={devices=[1,1,1,2]0,1}
  %rhs = f32[16,801,1,1024] parameter(1)
  %rhs.copy = f32[16,801,1,1024] copy(%rhs),
    sharding={devices=[1,2,1,1]0,1}
  ROOT %conv = f32[5,1,1,1024] convolution(%lhs.copy, %rhs.copy),
    dim_labels=f01b_i01o->01bf,batch_group_count=1024,
    window={size=801x1 pad=2_2x0_0},
    sharding={devices=[1,1,1,2]0,1}
})";

  TF_ASSERT_OK_AND_ASSIGN(auto module,
                          PartitionComputation(hlo_string, /*num_devices=*/2));
  VLOG(1) << module->ToString();
  auto root = module->entry_computation()->root_instruction();
  auto lhs = AllOf(
      op::Copy(op::DynamicSlice(op::Parameter(), op::Constant(), op::Constant(),
                                op::Constant(), op::Reshape())),
      op::Shape("f32[16,801,1,512]"));
  auto rhs = AllOf(op::Copy(op::DynamicSlice(
                       op::Pad(op::Parameter(), op::Constant()), op::Constant(),
                       op::Reshape(), op::Constant(), op::Constant())),
                   op::Shape("f32[16,401,1,1024]"));
  auto resharded_rhs = AllOf(
      op::Slice(op::Reshape(op::Transpose(op::AllToAll(op::Reshape(rhs))))),
      op::Shape("f32[16,801,1,512]"));
  EXPECT_THAT(root, AllOf(op::Convolution(lhs, resharded_rhs),
                          op::Shape("f32[5,1,1,512]")));
}

TEST_F(SpmdPartitioningTest, PartitionConvWithBathGroupCountLHSAlignWithRHS) {
  absl::string_view hlo_string = R"(
HloModule module

ENTRY entry {
  %lhs = f32[16,801,1,1024] parameter(0)
  %lhs.copy = f32[16,801,1,1024] copy(%lhs),
    sharding={devices=[1,2,1,1]0,1}
  %rhs = f32[16,801,1,1024] parameter(1)
  %rhs.copy = f32[16,801,1,1024] copy(%rhs),
    sharding={devices=[1,1,1,2]0,1}
  ROOT %conv = f32[5,1,1,1024] convolution(%lhs.copy, %rhs.copy),
    dim_labels=f01b_i01o->01bf,batch_group_count=1024,
    window={size=801x1 pad=2_2x0_0},
    sharding={devices=[1,1,1,2]0,1}
})";

  TF_ASSERT_OK_AND_ASSIGN(auto module,
                          PartitionComputation(hlo_string, /*num_devices=*/2));
  VLOG(1) << module->ToString();
  auto root = module->entry_computation()->root_instruction();
  auto rhs = AllOf(
      op::Copy(op::DynamicSlice(op::Parameter(), op::Constant(), op::Constant(),
                                op::Constant(), op::Reshape())),
      op::Shape("f32[16,801,1,512]"));
  auto lhs = AllOf(op::Copy(op::DynamicSlice(
                       op::Pad(op::Parameter(), op::Constant()), op::Constant(),
                       op::Reshape(), op::Constant(), op::Constant())),
                   op::Shape("f32[16,401,1,1024]"));
  auto resharded_lhs = AllOf(
      op::Slice(op::Reshape(op::Transpose(op::AllToAll(op::Reshape(lhs))))),
      op::Shape("f32[16,801,1,512]"));
  EXPECT_THAT(root, AllOf(op::Convolution(resharded_lhs, rhs),
                          op::Shape("f32[5,1,1,512]")));
}

TEST_F(SpmdPartitioningTest,
       PartitionConvWithBathGroupCountOutputAlignWithLHS) {
  absl::string_view hlo_string = R"(
HloModule module

ENTRY entry {
  %lhs = f32[16,801,1,1024] parameter(0)
  %lhs.copy = f32[16,801,1,1024] copy(%lhs),
    sharding={devices=[1,1,1,2]0,1}
  %rhs = f32[16,801,1,1024] parameter(1)
  %rhs.copy = f32[16,801,1,1024] copy(%rhs),
    sharding={devices=[1,1,1,2]0,1}
  ROOT %conv = f32[5,1,1,1024] convolution(%lhs.copy, %rhs.copy),
    dim_labels=f01b_i01o->01bf,batch_group_count=1024,
    window={size=801x1 pad=2_2x0_0},
    sharding={devices=[2,1,1,1]0,1}
})";

  TF_ASSERT_OK_AND_ASSIGN(auto module,
                          PartitionComputation(hlo_string, /*num_devices=*/2));
  VLOG(1) << module->ToString();
  auto root = module->entry_computation()->root_instruction();
  auto lhs = AllOf(
      op::Copy(op::DynamicSlice(op::Parameter(), op::Constant(), op::Constant(),
                                op::Constant(), op::Reshape())),
      op::Shape("f32[16,801,1,512]"));
  auto rhs = AllOf(
      op::Copy(op::DynamicSlice(op::Parameter(), op::Constant(), op::Constant(),
                                op::Constant(), op::Reshape())),
      op::Shape("f32[16,801,1,512]"));
  auto conv = AllOf(op::Convolution(lhs, rhs), op::Shape("f32[5,1,1,512]"));
  EXPECT_THAT(root, AllOf(op::Reshape(op::Transpose(op::AllToAll(
                              op::Reshape(op::Pad(conv, op::Constant()))))),
                          op::Shape("f32[3,1,1,1024]")));
}

TEST_F(SpmdPartitioningTest,
       PartitionConvWithBathGroupCountOutputAlignWithRHS) {
  absl::string_view hlo_string = R"(
HloModule module

ENTRY entry {
  %lhs = f32[16,801,1,1024] parameter(0)
  %lhs.copy = f32[16,801,1,1024] copy(%lhs),
    sharding={devices=[1,2,1,1]0,1}
  %rhs = f32[16,801,1,1024] parameter(1)
  %rhs.copy = f32[16,801,1,1024] copy(%rhs),
    sharding={devices=[1,1,1,2]0,1}
  ROOT %conv = f32[5,1,1,1024] convolution(%lhs.copy, %rhs.copy),
    dim_labels=f01b_i01o->01bf,batch_group_count=1024,
    window={size=801x1 pad=2_2x0_0},
    sharding={devices=[2,1,1,1]0,1}
})";

  TF_ASSERT_OK_AND_ASSIGN(auto module,
                          PartitionComputation(hlo_string, /*num_devices=*/2));
  VLOG(1) << module->ToString();
  auto root = module->entry_computation()->root_instruction();
  auto rhs = AllOf(
      op::Copy(op::DynamicSlice(op::Parameter(), op::Constant(), op::Constant(),
                                op::Constant(), op::Reshape())),
      op::Shape("f32[16,801,1,512]"));
  auto lhs = AllOf(op::Copy(op::DynamicSlice(
                       op::Pad(op::Parameter(), op::Constant()), op::Constant(),
                       op::Reshape(), op::Constant(), op::Constant())),
                   op::Shape("f32[16,401,1,1024]"));
  auto resharded_lhs = AllOf(
      op::Slice(op::Reshape(op::Transpose(op::AllToAll(op::Reshape(lhs))))),
      op::Shape("f32[16,801,1,512]"));
  auto conv =
      AllOf(op::Convolution(resharded_lhs, rhs), op::Shape("f32[5,1,1,512]"));
  EXPECT_THAT(root, AllOf(op::Reshape(op::Transpose(op::AllToAll(
                              op::Reshape(op::Pad(conv, op::Constant()))))),
                          op::Shape("f32[3,1,1,1024]")));
}

TEST_F(SpmdPartitioningTest, PartitionConvWithFeatureGroupCount) {
  absl::string_view hlo_string = R"(
HloModule module

ENTRY entry {
  %lhs = f32[16,801,1,1024] parameter(0)
  %lhs.copy = f32[16,801,1,1024] copy(%lhs),
    sharding={devices=[1,1,1,2]0,1}
  %rhs = f32[5,1,1,1024] parameter(1)
  %rhs.copy = f32[5,1,1,1024] copy(%rhs),
    sharding={devices=[1,1,1,2]0,1}
  ROOT %conv = f32[16,801,1,1024] convolution(%lhs.copy, %rhs.copy),
    dim_labels=b01f_01io->b01f,feature_group_count=1024,
    window={size=5x1 pad=2_2x0_0},
    sharding={devices=[1,1,1,2]0,1}
})";

  TF_ASSERT_OK_AND_ASSIGN(auto module,
                          PartitionComputation(hlo_string, /*num_devices=*/2));
  VLOG(1) << module->ToString();
  auto root = module->entry_computation()->root_instruction();
  auto lhs = AllOf(
      op::Copy(op::DynamicSlice(op::Parameter(), op::Constant(), op::Constant(),
                                op::Constant(), op::Reshape())),
      op::Shape("f32[16,801,1,512]"));
  auto rhs = AllOf(
      op::Copy(op::DynamicSlice(op::Parameter(), op::Constant(), op::Constant(),
                                op::Constant(), op::Reshape())),
      op::Shape("f32[5,1,1,512]"));
  EXPECT_THAT(root,
              AllOf(op::Convolution(lhs, rhs), op::Shape("f32[16,801,1,512]")));
}

TEST_F(SpmdPartitioningTest,
       PartitionConvWithFeatureGroupCountRHSAlignWithLHS) {
  absl::string_view hlo_string = R"(
HloModule module

ENTRY entry {
  %lhs = f32[16,801,1,1024] parameter(0)
  %lhs.copy = f32[16,801,1,1024] copy(%lhs),
    sharding={devices=[1,1,1,2]0,1}
  %rhs = f32[5,1,1,1024] parameter(1)
  %rhs.copy = f32[5,1,1,1024] copy(%rhs),
    sharding={devices=[2,1,1,1]0,1}
  ROOT %conv = f32[16,801,1,1024] convolution(%lhs.copy, %rhs.copy),
    dim_labels=b01f_01io->b01f,feature_group_count=1024,
    window={size=5x1 pad=2_2x0_0},
    sharding={devices=[1,1,1,2]0,1}
})";

  TF_ASSERT_OK_AND_ASSIGN(auto module,
                          PartitionComputation(hlo_string, /*num_devices=*/2));
  VLOG(1) << module->ToString();
  auto root = module->entry_computation()->root_instruction();
  auto lhs = AllOf(
      op::Copy(op::DynamicSlice(op::Parameter(), op::Constant(), op::Constant(),
                                op::Constant(), op::Reshape())),
      op::Shape("f32[16,801,1,512]"));
  auto rhs = AllOf(op::Copy(op::DynamicSlice(
                       op::Pad(op::Parameter(), op::Constant()), op::Reshape(),
                       op::Constant(), op::Constant(), op::Constant())),
                   op::Shape("f32[3,1,1,1024]"));
  auto resharded_rhs = AllOf(
      op::Slice(op::Reshape(op::Transpose(op::AllToAll(op::Reshape(rhs))))),
      op::Shape("f32[5,1,1,512]"));
  EXPECT_THAT(root, AllOf(op::Convolution(lhs, resharded_rhs),
                          op::Shape("f32[16,801,1,512]")));
}

TEST_F(SpmdPartitioningTest,
       PartitionConvWithFeatureGroupCountLHSAlignWithRHS) {
  absl::string_view hlo_string = R"(
HloModule module

ENTRY entry {
  %lhs = f32[16,801,1,1024] parameter(0)
  %lhs.copy = f32[16,801,1,1024] copy(%lhs),
    sharding={devices=[1,2,1,1]0,1}
  %rhs = f32[5,1,1,1024] parameter(1)
  %rhs.copy = f32[5,1,1,1024] copy(%rhs),
    sharding={devices=[1,1,1,2]0,1}
  ROOT %conv = f32[16,801,1,1024] convolution(%lhs.copy, %rhs.copy),
    dim_labels=b01f_01io->b01f,feature_group_count=1024,
    window={size=5x1 pad=2_2x0_0},
    sharding={devices=[1,1,1,2]0,1}
})";

  TF_ASSERT_OK_AND_ASSIGN(auto module,
                          PartitionComputation(hlo_string, /*num_devices=*/2));
  VLOG(1) << module->ToString();
  auto root = module->entry_computation()->root_instruction();
  auto lhs = AllOf(op::Copy(op::DynamicSlice(
                       op::Pad(op::Parameter(), op::Constant()), op::Constant(),
                       op::Reshape(), op::Constant(), op::Constant())),
                   op::Shape("f32[16,401,1,1024]"));
  auto rhs = AllOf(
      op::Copy(op::DynamicSlice(op::Parameter(), op::Constant(), op::Constant(),
                                op::Constant(), op::Reshape())),
      op::Shape("f32[5,1,1,512]"));
  auto resharded_lhs = AllOf(
      op::Slice(op::Reshape(op::Transpose(op::AllToAll(op::Reshape(lhs))))),
      op::Shape("f32[16,801,1,512]"));
  EXPECT_THAT(root, AllOf(op::Convolution(resharded_lhs, rhs),
                          op::Shape("f32[16,801,1,512]")));
}

TEST_F(SpmdPartitioningTest,
       PartitionConvWithFeatureGroupCountAlignOuputWithLHS) {
  absl::string_view hlo_string = R"(
HloModule module

ENTRY entry {
  %lhs = f32[16,801,1,1024] parameter(0)
  %lhs.copy = f32[16,801,1,1024] copy(%lhs),
    sharding={devices=[1,1,1,2]0,1}
  %rhs = f32[5,1,1,1024] parameter(1)
  %rhs.copy = f32[5,1,1,1024] copy(%rhs),
    sharding={devices=[1,1,1,2]0,1}
  ROOT %conv = f32[16,801,1,1024] convolution(%lhs.copy, %rhs.copy),
    dim_labels=b01f_01io->b01f,feature_group_count=1024,
    window={size=5x1 pad=2_2x0_0},
    sharding={devices=[2,1,1,1]0,1}
})";

  TF_ASSERT_OK_AND_ASSIGN(auto module,
                          PartitionComputation(hlo_string, /*num_devices=*/2));
  VLOG(1) << module->ToString();
  auto root = module->entry_computation()->root_instruction();
  auto lhs = AllOf(
      op::Copy(op::DynamicSlice(op::Parameter(), op::Constant(), op::Constant(),
                                op::Constant(), op::Reshape())),
      op::Shape("f32[16,801,1,512]"));
  auto rhs = AllOf(
      op::Copy(op::DynamicSlice(op::Parameter(), op::Constant(), op::Constant(),
                                op::Constant(), op::Reshape())),
      op::Shape("f32[5,1,1,512]"));
  auto conv = AllOf(op::Convolution(lhs, rhs), op::Shape("f32[16,801,1,512]"));
  EXPECT_THAT(root,
              AllOf(op::Reshape(op::Transpose(op::AllToAll(op::Reshape(conv)))),
                    op::Shape("f32[8,801,1,1024]")));
}

TEST_F(SpmdPartitioningTest,
       PartitionConvGroupOnFeatureGroupCount_RHSPartialReplicate) {
  absl::string_view hlo_string = R"(
HloModule module

ENTRY entry {
  %lhs = f32[16,801,1,1024] parameter(0)
  %lhs.copy = f32[16,801,1,1024] copy(%lhs),
    sharding={devices=[1,2,1,2]0,1,2,3}
  %rhs = f32[5,1,1,1024] parameter(1)
  %rhs.copy = f32[5,1,1,1024] copy(%rhs),
    sharding={devices=[1,1,1,2,2]0,2,1,3 last_tile_dim_replicate}
  ROOT %conv = f32[16,801,1,1024] convolution(%lhs.copy, %rhs.copy),
    dim_labels=b01f_01io->b01f,feature_group_count=1024,
    window={size=5x1 pad=2_2x0_0},
    sharding={devices=[1,2,1,2]0,1,2,3}
})";

  TF_ASSERT_OK_AND_ASSIGN(auto module,
                          PartitionComputation(hlo_string, /*num_devices=*/4));
  VLOG(1) << module->ToString();
  auto root = module->entry_computation()->root_instruction();
  auto lhs = AllOf(op::Copy(op::DynamicSlice(
                       op::Pad(op::Parameter(), op::Constant()), op::Constant(),
                       op::Reshape(), op::Constant(), op::Reshape())),
                   op::Shape("f32[16,401,1,512]"));
  auto left_halo = AllOf(op::Shape("f32[16,2, 1, 512]"),
                         op::CollectivePermute(op::Slice(lhs)));
  auto right_halo = AllOf(op::Shape("f32[16,2, 1, 512]"),
                          op::CollectivePermute(op::Slice(lhs)));
  auto rhs = AllOf(
      op::Copy(op::DynamicSlice(op::Parameter(), op::Constant(), op::Constant(),
                                op::Constant(), op::Reshape())),
      op::Shape("f32[5,1,1,512]"));
  EXPECT_THAT(
      root,
      AllOf(op::Convolution(
                op::Select(_, op::Concatenate(left_halo, lhs, right_halo), _),
                rhs),
            op::Shape("f32[16, 401, 1, 512]")));
}

TEST_F(SpmdPartitioningTest,
       PartitionConvGroupOnFeatureGroupCount_RHSAlignWithOutput) {
  absl::string_view hlo_string = R"(
HloModule module

ENTRY entry {
  %lhs = f32[16,801,1,1024] parameter(0)
  %lhs.copy = f32[16,801,1,1024] copy(%lhs),
    sharding={devices=[1,2,1,2]0,1,2,3}
  %rhs = f32[5,1,1,1024] parameter(1), sharding={replicated}
  ROOT %conv = f32[16,801,1,1024] convolution(%lhs.copy, %rhs),
    dim_labels=b01f_01io->b01f,feature_group_count=1024,
    window={size=5x1 pad=2_2x0_0},
    sharding={devices=[1,2,1,2]0,1,2,3}
})";
  TF_ASSERT_OK_AND_ASSIGN(auto module,
                          PartitionComputation(hlo_string, /*num_devices=*/4));
  VLOG(1) << module->ToString();
  auto root = module->entry_computation()->root_instruction();
  auto lhs = AllOf(op::Copy(op::DynamicSlice(
                       op::Pad(op::Parameter(), op::Constant()), op::Constant(),
                       op::Reshape(), op::Constant(), op::Reshape())),
                   op::Shape("f32[16,401,1,512]"));
  auto left_halo = AllOf(op::Shape("f32[16,2, 1, 512]"),
                         op::CollectivePermute(op::Slice(lhs)));
  auto right_halo = AllOf(op::Shape("f32[16,2, 1, 512]"),
                          op::CollectivePermute(op::Slice(lhs)));
  auto rhs =
      AllOf(op::DynamicSlice(op::Parameter(), op::Constant(), op::Constant(),
                             op::Constant(), op::Reshape()),
            op::Shape("f32[5,1,1,512]"));
  EXPECT_THAT(
      root,
      AllOf(op::Convolution(
                op::Select(_, op::Concatenate(left_halo, lhs, right_halo), _),
                rhs),
            op::Shape("f32[16, 401, 1, 512]")));
}

TEST_F(SpmdPartitioningTest,
       PartitionConvGroupOnFeatureGroupCount_LHSAlignWithOutput) {
  absl::string_view hlo_string = R"(
HloModule module

ENTRY entry {
  %lhs = f32[16,801,1,1024] parameter(0)
  %lhs.copy = f32[16,801,1,1024] copy(%lhs),
    sharding={devices=[2,1,1,1,2]0,1,2,3 last_tile_dim_replicate}
  %rhs = f32[5,1,1,1024] parameter(1)
  %rhs.copy = f32[5,1,1,1024] copy(%rhs),
    sharding={devices=[1,1,1,2,2]0,2,1,3 last_tile_dim_replicate}
  ROOT %conv = f32[16,801,1,1024] convolution(%lhs.copy, %rhs.copy),
    dim_labels=b01f_01io->b01f,feature_group_count=1024,
    window={size=5x1 pad=2_2x0_0},
    sharding={devices=[1,2,1,2]0,1,2,3}
})";
  TF_ASSERT_OK_AND_ASSIGN(auto module,
                          PartitionComputation(hlo_string, /*num_devices=*/4));
  VLOG(1) << module->ToString();
  auto root = module->entry_computation()->root_instruction();
  auto lhs = AllOf(
      op::Copy(op::DynamicSlice(op::Parameter(), op::Reshape(), op::Constant(),
                                op::Constant(), op::Constant())),
      op::Shape("f32[8,801,1,1024]"));
  auto resharded_lhs =
      AllOf(op::Reshape(op::Transpose(op::AllToAll(op::Reshape(
                op::Pad(op::DynamicSlice(lhs, op::Subtract(), op::Subtract(),
                                         op::Subtract(), op::Subtract()),
                        op::Constant()))))),
            op::Shape("f32[16,401,1,512]"));
  auto left_halo = AllOf(op::Shape("f32[16,2, 1, 512]"),
                         op::CollectivePermute(op::Slice(resharded_lhs)));
  auto right_halo = AllOf(op::Shape("f32[16,2, 1, 512]"),
                          op::CollectivePermute(op::Slice(resharded_lhs)));
  auto rhs = AllOf(
      op::Copy(op::DynamicSlice(op::Parameter(), op::Constant(), op::Constant(),
                                op::Constant(), op::Reshape())),
      op::Shape("f32[5,1,1,512]"));
  EXPECT_THAT(
      root,
      AllOf(
          op::Convolution(
              op::Select(
                  _, op::Concatenate(left_halo, resharded_lhs, right_halo), _),
              rhs),
          op::Shape("f32[16, 401, 1, 512]")));
}

TEST_F(SpmdPartitioningTest, PartitionConvGroupOnBatchGroupCount) {
  absl::string_view hlo_string = R"(
HloModule module

ENTRY entry {
  %lhs = f32[16,801,1,1024] parameter(0)
  %lhs.copy = f32[16,801,1,1024] copy(%lhs),
    sharding={devices=[1,2,1,2]0,1,2,3}
  %rhs = f32[16,801,1,1024] parameter(1)
  %rhs.copy = f32[16,801,1,1024] copy(%rhs),
    sharding={devices=[1,2,1,2]0,1,2,3}
  ROOT %conv = f32[5,1,1,1024] convolution(%lhs.copy, %rhs.copy),
    dim_labels=f01b_i01o->01bf,batch_group_count=1024,
    window={size=801x1 pad=2_2x0_0},
    sharding={devices=[1,1,1,2,2]0,1,2,3 last_tile_dim_replicate}
})";

  TF_ASSERT_OK_AND_ASSIGN(auto module,
                          PartitionComputation(hlo_string, /*num_devices=*/4));
  VLOG(1) << module->ToString();
  auto root = module->entry_computation()->root_instruction();
  auto lhs = AllOf(
      op::Select(_,
                 op::Copy(op::DynamicSlice(
                     op::Pad(op::Parameter(), op::Constant()), op::Constant(),
                     op::Reshape(), op::Constant(), op::Reshape())),
                 _),
      op::Shape("f32[16,401,1,512]"));
  auto left_halo = AllOf(op::Shape("f32[16,2, 1, 512]"),
                         op::CollectivePermute(op::Slice(lhs)));
  auto right_halo = AllOf(op::Shape("f32[16,2, 1, 512]"),
                          op::CollectivePermute(op::Slice(lhs)));
  auto rhs = AllOf(op::Copy(op::DynamicSlice(
                       op::Pad(op::Parameter(), op::Constant()), op::Constant(),
                       op::Reshape(), op::Constant(), op::Reshape())),
                   op::Shape("f32[16,401,1,512]"));
  auto conv = AllOf(op::Convolution(op::Concatenate(left_halo, lhs, right_halo),
                                    op::Select(_, rhs, _)),
                    op::Shape("f32[5,1,1,512]"));
  EXPECT_THAT(root, AllOf(op::CollectivePermute(op::AllReduce(conv)),
                          op::Shape("f32[5,1,1,512]")));
}

TEST_F(SpmdPartitioningTest,
       PartitionConvWithFeatureGroupCountAlignOuputWithRHS) {
  absl::string_view hlo_string = R"(
HloModule module

ENTRY entry {
  %lhs = f32[16,801,1,1024] parameter(0)
  %lhs.copy = f32[16,801,1,1024] copy(%lhs),
    sharding={devices=[1,2,1,1]0,1}
  %rhs = f32[5,1,1,1024] parameter(1)
  %rhs.copy = f32[5,1,1,1024] copy(%rhs),
    sharding={devices=[1,1,1,2]0,1}
  ROOT %conv = f32[16,801,1,1024] convolution(%lhs.copy, %rhs.copy),
    dim_labels=b01f_01io->b01f,feature_group_count=1024,
    window={size=5x1 pad=2_2x0_0},
    sharding={devices=[2,1,1,1]0,1}
})";

  TF_ASSERT_OK_AND_ASSIGN(auto module,
                          PartitionComputation(hlo_string, /*num_devices=*/2));
  VLOG(1) << module->ToString();
  auto root = module->entry_computation()->root_instruction();
  auto lhs = AllOf(op::Copy(op::DynamicSlice(
                       op::Pad(op::Parameter(), op::Constant()), op::Constant(),
                       op::Reshape(), op::Constant(), op::Constant())),
                   op::Shape("f32[16,401,1,1024]"));
  auto rhs = AllOf(
      op::Copy(op::DynamicSlice(op::Parameter(), op::Constant(), op::Constant(),
                                op::Constant(), op::Reshape())),
      op::Shape("f32[5,1,1,512]"));
  auto resharded_lhs = AllOf(
      op::Slice(op::Reshape(op::Transpose(op::AllToAll(op::Reshape(lhs))))),
      op::Shape("f32[16,801,1,512]"));
  auto conv = AllOf(op::Convolution(resharded_lhs, rhs),
                    op::Shape("f32[16,801,1,512]"));
  EXPECT_THAT(root,
              AllOf(op::Reshape(op::Transpose(op::AllToAll(op::Reshape(conv)))),
                    op::Shape("f32[8,801,1,1024]")));
}

TEST_F(SpmdPartitioningTest, PartitionConvWithFeatureGroupCountBackProp) {
  absl::string_view hlo_string = R"(
HloModule module

ENTRY entry {
  %lhs = f32[16,801,1,1024] parameter(0)
  %lhs.copy = f32[16,801,1,1024] copy(%lhs),
    sharding={devices=[1,1,1,2]0,1}
  %rhs = f32[5,1,1024,1] parameter(1)
  %rhs.copy = f32[5,1,1024,1] copy(%rhs),
    sharding={devices=[1,1,2,1]0,1}
  ROOT %conv = f32[16,801,1,1024] convolution(%lhs.copy, %rhs.copy),
    dim_labels=b01f_01oi->b01f,feature_group_count=1024,
    window={size=5x1 pad=2_2x0_0 rhs_reversal=1x1},
    sharding={devices=[1,1,1,2]0,1}
})";

  TF_ASSERT_OK_AND_ASSIGN(auto module,
                          PartitionComputation(hlo_string, /*num_devices=*/2));
  VLOG(1) << module->ToString();
  auto root = module->entry_computation()->root_instruction();
  auto lhs = AllOf(
      op::Copy(op::DynamicSlice(op::Parameter(), op::Constant(), op::Constant(),
                                op::Constant(), op::Reshape())),
      op::Shape("f32[16,801,1,512]"));
  auto rhs = AllOf(
      op::Copy(op::DynamicSlice(op::Parameter(), op::Constant(), op::Constant(),
                                op::Reshape(), op::Constant())),
      op::Shape("f32[5,1,512,1]"));
  EXPECT_THAT(root,
              AllOf(op::Convolution(lhs, rhs), op::Shape("f32[16,801,1,512]")));
}

TEST_F(SpmdPartitioningTest, NoReshardOnBroadcastDims) {
  absl::string_view hlo_string = R"(
HloModule module

ENTRY entry {
  %param0 = f32[2,3] parameter(0)
  %param1 = f32[2,3,20] parameter(1)
  %br0 = f32[20,2,20,3,20] broadcast(%param0), dimensions={1,3}, sharding={devices=[2,1,2,1,2]0,1,2,3,4,5,6,7}
  %br1 = f32[20,2,20,3,20] broadcast(%param1), dimensions={1,3,4}, sharding={devices=[2,1,2,1,2]0,1,2,3,4,5,6,7}
  %add = f32[20,2,20,3,20] add(%br0, %br1), sharding={devices=[2,1,2,1,2]0,1,2,3,4,5,6,7}
  %reshape = f32[10,4,10,6,20] reshape(%br0), sharding={devices=[2,1,2,1,2]0,1,2,3,4,5,6,7}
  %transpose = f32[2,3,20,20,20] transpose(%br0), dimensions={1,3,0,2,4}, sharding={devices=[1,1,2,2,2]0,1,2,3,4,5,6,7}
  %copy_add0 = f32[20,2,20,3,20] copy(%add), sharding={devices=[2,1,2,1,2]6,7,2,3,4,5,0,1}
  %copy_add1 = f32[20,2,20,3,20] copy(%add), sharding={devices=[2,1,2,1,2]7,6,3,2,5,4,0,1}
  %copy_reshape = f32[10,4,10,6,20] copy(%reshape), sharding={devices=[2,1,2,1,2]7,6,3,2,5,4,0,1}
  %copy_transpose = f32[2,3,20,20,20] copy(%transpose), sharding={devices=[1,1,2,2,2]7,6,3,2,5,4,0,1}
  ROOT %tuple = (f32[20,2,20,3,20], f32[20,2,20,3,20], f32[10,4,10,6,20], f32[2,3,20,20,20])
    tuple(%copy_add0, %copy_add1, %copy_reshape, %copy_transpose),
    sharding={{devices=[2,1,2,1,2]6,7,2,3,4,5,0,1},{devices=[2,1,2,1,2]7,6,3,2,5,4,0,1},{devices=[2,1,2,1,2]7,6,3,2,5,4,0,1},{devices=[1,1,2,2,2]7,6,3,2,5,4,0,1}}
})";

  TF_ASSERT_OK_AND_ASSIGN(auto module,
                          PartitionComputation(hlo_string, /*num_devices=*/8));
  VLOG(1) << module->ToString();
  auto root = module->entry_computation()->root_instruction();
  // Reshard on copy_add0 only happens on broadcast dims, can be skipped.
  auto copy_add0 =
      op::Copy(op::Copy(op::Add(op::Broadcast(_), op::Broadcast(_))));
  // Reshard on copy_add1 also happens on non-broadcast dims.
  auto copy_add1 = op::Copy(
      op::CollectivePermute(op::Add(op::Broadcast(_), op::Broadcast(_))));
  // Reshard on copy_reshape only happens on broadcast dims, can be skipped.
  auto copy_reshape = op::Copy(op::Copy(op::Reshape(op::Broadcast(_))));
  // Reshard on copy_transpose only happens on broadcast dims, can be skipped.
  auto copy_transpose = op::Copy(op::Copy(op::Transpose(op::Broadcast(_))));
  EXPECT_THAT(root,
              op::Tuple(copy_add0, copy_add1, copy_reshape, copy_transpose));
}

TEST_F(SpmdPartitioningTest,
       ConvolutionFilterIFOFPartitionedInputPartialReplicate) {
  absl::string_view hlo_string = R"(
HloModule module

ENTRY entry {
  %lhs = f32[128,112,112,12] parameter(0)
  %lhs.copy = f32[128,112,112,12] copy(f32[128,112,112,12] %lhs),
    sharding={devices=[1,1,1,2,2]0,1,2,3 last_tile_dim_replicate}
  %rhs = f32[7,7,12,64] parameter(1)
  %rhs.copy = f32[7,7,12,64] copy(f32[7,7,12,64] %rhs),
    sharding={devices=[1,1,2,2]0,1,2,3}
  ROOT %conv = f32[128,56,56,64] convolution(
    f32[128,112,112,12] %lhs.copy,
    f32[7,7,12,64] %rhs.copy),
    window={size=7x7 stride=2x2 pad=3_3x3_3},
    dim_labels=b01f_01io->b01f,
    sharding={devices=[1,1,1,2,2]0,1,2,3 last_tile_dim_replicate}
})";

  TF_ASSERT_OK_AND_ASSIGN(auto module,
                          PartitionComputation(hlo_string, /*num_devices=*/4));
  VLOG(1) << module->ToString();
  auto root = module->entry_computation()->root_instruction();
  auto lhs = AllOf(
      op::Copy(op::DynamicSlice(op::Parameter(), op::Constant(), op::Constant(),
                                op::Constant(), op::Reshape())),
      op::Shape("f32[128,112,112,6]"));
  auto rhs = AllOf(
      op::Copy(op::DynamicSlice(op::Parameter(), op::Constant(), op::Constant(),
                                op::Reshape(), op::Reshape())),
      op::Shape("f32[7,7,6,32]"));

  EXPECT_THAT(
      root,
      AllOf(op::CollectivePermute(op::AllReduce(op::Convolution(lhs, rhs))),
            op::Shape("f32[128,56,56,32]")));
}

TEST_F(SpmdPartitioningTest,
       ConvolutionInputKernelNonContractingDimPartialReplicate) {
  absl::string_view hlo_string = R"(
HloModule module

ENTRY entry {
  %lhs = f32[128,56,56,256] parameter(0)
  %lhs.copy = f32[128,56,56,256] copy(%lhs),
  sharding={devices=[1,1,1,2,2]0,1,2,3 last_tile_dim_replicate}
  %rhs = f32[128,28,28,512] parameter(1)
  %rhs.copy = f32[128,28,28,512] copy(%rhs),
  sharding={devices=[1,1,1,2,2]0,1,2,3 last_tile_dim_replicate}
  ROOT %conv = f32[1,1,256,512] convolution(%lhs.copy, %rhs.copy),
    window={size=28x28 pad=0_-1x0_-1 rhs_dilate=2x2}, dim_labels=f01b_i01o->01bf,
    sharding={devices=[1,1,2,2]0,1,2,3}
})";

  TF_ASSERT_OK_AND_ASSIGN(auto module,
                          PartitionComputation(hlo_string, /*num_devices=*/4));
  VLOG(1) << module->ToString();
  auto root = module->entry_computation()->root_instruction();
  auto lhs = AllOf(
      op::Copy(op::DynamicSlice(op::Parameter(), op::Constant(), op::Constant(),
                                op::Constant(), op::Reshape())),
      op::Shape("f32[128,56,56,128]"));
  auto rhs = AllOf(
      op::Copy(op::DynamicSlice(op::Parameter(), op::Constant(), op::Constant(),
                                op::Constant(), op::Reshape())),
      op::Shape("f32[128,28,28,256]"));

  EXPECT_THAT(root, AllOf(op::Convolution(lhs, op::CollectivePermute(rhs)),
                          op::Shape("f32[1,1,128,256]")));
}

TEST_F(SpmdPartitioningTest,
       ConvolutionInputSpatialDimAndFeatureDimParttiioned) {
  absl::string_view hlo_string = R"(
HloModule module

ENTRY entry {
  %lhs = f32[8,210,210,12] parameter(0)
  %lhs.copy = f32[8,210,210,12] copy(f32[8,210,210,12] %lhs),
    sharding={devices=[1,2,1,2]0,1,2,3}
  %rhs = f32[3,3,12,32] parameter(1)
  %rhs.copy = f32[3,3,12,32] copy(f32[3,3,12,32] %rhs),
    sharding={devices=[1,1,2,1,2]0,1,2,3 last_tile_dim_replicate}
  ROOT %conv = f32[8,210,210,32] convolution(
    f32[8,210,210,12] %lhs.copy,
    f32[3,3,12,32] %rhs.copy),
    window={size=3x3 pad=1_1x1_1},
    dim_labels=b01f_01io->b01f,
    sharding={devices=[1,2,1,1,2]0,1,2,3 last_tile_dim_replicate}
})";
  TF_ASSERT_OK_AND_ASSIGN(auto module,
                          PartitionComputation(hlo_string, /*num_devices=*/4));
  VLOG(1) << module->ToString();
  auto root = module->entry_computation()->root_instruction();
  auto lhs = AllOf(
      op::Copy(op::DynamicSlice(op::Parameter(), op::Constant(), op::Reshape(),
                                op::Constant(), op::Reshape())),
      op::Shape("f32[8,105,210,6]"));
  auto left_halo =
      AllOf(op::CollectivePermute(op::Slice(lhs)), op::Shape("f32[8,1,210,6]"));
  auto right_halo =
      AllOf(op::CollectivePermute(op::Slice(lhs)), op::Shape("f32[8,1,210,6]"));
  auto exchanged_lhs = AllOf(
      op::Select(op::And(_, _), op::Concatenate(left_halo, lhs, right_halo),
                 op::Broadcast(_)),
      op::Shape("f32[8,107,210,6]"));
  auto rhs = AllOf(
      op::Copy(op::DynamicSlice(op::Parameter(), op::Constant(), op::Constant(),
                                op::Reshape(), op::Constant())),
      op::Shape("f32[3,3,6,32]"));
  EXPECT_THAT(root, AllOf(op::AllReduce(op::Convolution(
                              exchanged_lhs, op::CollectivePermute(rhs))),
                          op::Shape("f32[8,105,210,32]")));
}

TEST_F(SpmdPartitioningTest, Fft3D) {
  absl::string_view hlo_string = R"(
HloModule module

ENTRY entry {
  constant = c64[1,1,6]
    constant({{{(0,0),(1,1),(2,2),(3,3),(4,4),(5,5)}}}),
    sharding={devices=[1,1,2]0,1}
  ROOT fft = c64[1,1,6] fft(c64[1,1,6] constant), fft_type=FFT, fft_length={6},
    sharding={devices=[1,1,2]0,1}
}
)";

  TF_ASSERT_OK_AND_ASSIGN(auto module,
                          PartitionComputation(hlo_string, /*num_devices=*/2));
  VLOG(1) << module->ToString();
  auto root = module->entry_computation()->root_instruction();
  auto input = AllOf(op::DynamicSlice(op::Constant(), op::Constant(),
                                      op::Constant(), op::Reshape()),
                     op::Shape("c64[1,1,3]"));
  auto padded_input =
      AllOf(op::DynamicSlice(
                op::Concatenate(input, op::CollectivePermute(op::Slice())),
                op::Constant(), op::Constant(), op::Reshape()),
            op::Shape("c64[1,1,4]"));

  auto shuffled_input =
      AllOf(op::Slice(op::AllToAll(op::Dot(padded_input, op::Convert()))),
            op::Shape("c64[1,1,3]"));

  auto local_fft = AllOf(op::Fft(shuffled_input), op::Shape("c64[1,1,3]"));

  EXPECT_THAT(root, AllOf(op::GetTupleElement(op::While(op::Tuple(
                              _, op::Multiply(local_fft, op::Exp()), _, _, _))),
                          op::Shape("c64[1,1,3]")));
}

TEST_F(SpmdPartitioningTest, DotInputsAreIdentical) {
  absl::string_view hlo_string = R"(
HloModule module

ENTRY entry {
  %parameter.1 = f32[4000,4000]{1,0} parameter(0),
    sharding={devices=[2,4]0,1,2,3,4,5,6,7}
  ROOT %convolution = f32[4000,4000]{1,0} convolution(
    f32[4000,4000]{1,0} %parameter.1, f32[4000,4000]{1,0} %parameter.1),
    dim_labels=bf_io->bf, sharding={devices=[2,4]0,1,2,3,4,5,6,7}
}

)";

  TF_ASSERT_OK_AND_ASSIGN(auto module,
                          PartitionComputation(hlo_string, /*num_devices=*/8));
  VLOG(1) << module->ToString();
  auto root = module->entry_computation()->root_instruction();
  auto param = AllOf(op::Parameter(), op::Shape("f32[2000, 1000]"));
  auto resharded_lhs =
      AllOf(op::AllReduce(op::DynamicUpdateSlice(_, param, _, _)),
            op::Shape("f32[2000, 4000]"));
  auto resharded_rhs =
      AllOf(op::AllReduce(op::DynamicUpdateSlice(_, op::Copy(param), _, _)),
            op::Shape("f32[4000, 1000]"));
  EXPECT_THAT(root, AllOf(op::Convolution(resharded_lhs, resharded_rhs),
                          op::Shape("f32[2000, 1000]")));
}

TEST_F(SpmdPartitioningTest, ConstantSliceReshard) {
  absl::string_view hlo_string = R"(
HloModule module

ENTRY entry {
  %constant.785 = f32[1,8] constant({{0,1,2,3,4,5,6,7}}),
    sharding={devices=[1,8]0,1,2,3,4,5,6,7}
  %slice.62 = f32[1,1] slice(%constant.785), slice={[0:1], [0:1]},
    sharding={devices=[1,8]0,1,2,3,4,5,6,7}
  ROOT %reshape.779 = f32[] reshape(%slice.62), sharding={replicated}
})";
  TF_ASSERT_OK_AND_ASSIGN(auto module,
                          PartitionComputation(hlo_string, /*num_devices=*/8));
  auto root = module->entry_computation()->root_instruction();
  VLOG(1) << module->ToString();
  auto slice = AllOf(op::Shape("f32[1,1]"),
                     op::Copy(op::DynamicSlice(op::Constant(), _, _)));
  EXPECT_THAT(root, op::Reshape(op::Slice(op::AllReduce(op::DynamicUpdateSlice(
                        op::Broadcast(), slice, _, _)))));
}

TEST_F(SpmdPartitioningTest, GatherParallelDimRedistributionOperand) {
  absl::string_view hlo_string = R"(
HloModule module

ENTRY %module {
  %parameter.0 = s32[8,4,2,2]{3,2,1,0} parameter(0),
    sharding={devices=[4,2,1,1]0,1,2,3,4,5,6,7}
  %iota = s32[1,8,4]{2,1,0} iota(), iota_dimension=2,
    sharding={devices=[1,8,1]0,1,2,3,4,5,6,7}
  %iota2 = s32[1,8,4]{2,1,0} iota(), iota_dimension=1,
    sharding={devices=[1,8,1]0,1,2,3,4,5,6,7}
  %concatenate.19 = s32[2,8,4]{2,1,0} concatenate(s32[1,8,4]{2,1,0} %iota,
    s32[1,8,4]{2,1,0} %iota2), dimensions={0},
    sharding={devices=[1,8,1]0,1,2,3,4,5,6,7}
  ROOT %gather.20 = s32[8,4,2,2]{3,2,1,0} gather(
    s32[8,4,2,2]{3,2,1,0} %parameter.0,
    s32[2,8,4]{2,1,0} %concatenate.19), offset_dims={2,3},
    collapsed_slice_dims={0,1}, start_index_map={1,0}, index_vector_dim=0,
    slice_sizes={1,1,2,2}, sharding={replicated}
})";
  TF_ASSERT_OK_AND_ASSIGN(auto module,
                          PartitionComputation(hlo_string, /*num_devices=*/8));
  auto root = module->entry_computation()->root_instruction();
  VLOG(1) << module->ToString();
  auto operand = AllOf(op::Shape("s32[1,4,2,2]"), op::DynamicSlice());
  auto indices = AllOf(op::Shape("s32[2,1,4]"), op::Subtract());
  auto gather = AllOf(op::Shape("s32[1,4,2,2]"), op::Gather(operand, indices));
  EXPECT_THAT(root,
              op::AllReduce(op::DynamicUpdateSlice(_, gather, _, _, _, _)));
}

TEST_F(SpmdPartitioningTest, GatherParallelDimRedistributionIndices) {
  absl::string_view hlo_string = R"(
HloModule module

ENTRY %module {
  %parameter.0 = s32[8,4,2,2]{3,2,1,0} parameter(0),
    sharding={devices=[8,1,1,1]0,1,2,3,4,5,6,7}
  %iota = s32[1,8,4]{2,1,0} iota(), iota_dimension=1,
    sharding={devices=[1,4,2]0,1,2,3,4,5,6,7}
  %iota2 = s32[1,8,4]{2,1,0} iota(), iota_dimension=2,
    sharding={devices=[1,4,2]0,1,2,3,4,5,6,7}
  %concatenate.19 = s32[2,8,4]{2,1,0} concatenate(s32[1,8,4]{2,1,0} %iota,
    s32[1,8,4]{2,1,0} %iota2), dimensions={0},
    sharding={devices=[1,4,2]0,1,2,3,4,5,6,7}
  ROOT %gather.20 = s32[8,4,2,2]{3,2,1,0} gather(s32[8,4,2,2]{3,2,1,0} %parameter.0,
    s32[2,8,4]{2,1,0} %concatenate.19), offset_dims={2,3},
    collapsed_slice_dims={0,1}, start_index_map={0,1}, index_vector_dim=0,
    slice_sizes={1,1,2,2}, sharding={replicated}
})";
  TF_ASSERT_OK_AND_ASSIGN(auto module,
                          PartitionComputation(hlo_string, /*num_devices=*/8));
  auto root = module->entry_computation()->root_instruction();
  VLOG(1) << module->ToString();
  auto operand = AllOf(op::Shape("s32[2,2,2,2]"), op::DynamicSlice());
  auto indices = AllOf(op::Shape("s32[2,2,2]"), op::Subtract());
  auto gather = AllOf(op::Shape("s32[2,2,2,2]"), op::Gather(operand, indices));
  EXPECT_THAT(root, op::AllReduce(op::AllReduce(
                        op::DynamicUpdateSlice(_, gather, _, _, _, _))));
}

TEST_F(SpmdPartitioningTest, GatherParallelDimReplicatedIndices) {
  absl::string_view hlo_string = R"(
HloModule module

ENTRY %module {
  %parameter.0 = s32[8,4,2,2]{3,2,1,0} parameter(0),
    sharding={devices=[8,1,1,1]0,1,2,3,4,5,6,7}
  %iota = s32[1,8,4]{2,1,0} iota(), iota_dimension=1,
    sharding={replicated}
  %iota2 = s32[1,8,4]{2,1,0} iota(), iota_dimension=2,
    sharding={replicated}
  %concatenate.19 = s32[2,8,4]{2,1,0} concatenate(s32[1,8,4]{2,1,0} %iota,
    s32[1,8,4]{2,1,0} %iota2), dimensions={0},
    sharding={replicated}
  ROOT %gather.20 = s32[8,4,2,2]{3,2,1,0} gather(
    s32[8,4,2,2]{3,2,1,0} %parameter.0,
    s32[2,8,4]{2,1,0} %concatenate.19), offset_dims={2,3},
    collapsed_slice_dims={0,1}, start_index_map={0,1}, index_vector_dim=0,
    slice_sizes={1,1,2,2}, sharding={replicated}
})";
  TF_ASSERT_OK_AND_ASSIGN(auto module,
                          PartitionComputation(hlo_string, /*num_devices=*/8));
  auto root = module->entry_computation()->root_instruction();
  VLOG(1) << module->ToString();
  auto operand = AllOf(op::Shape("s32[1,4,2,2]"), op::Parameter());
  auto indices = AllOf(op::Shape("s32[2,1,4]"), op::Subtract());
  auto gather = AllOf(op::Shape("s32[1,4,2,2]"), op::Gather(operand, indices));
  EXPECT_THAT(root,
              op::AllReduce(op::DynamicUpdateSlice(_, gather, _, _, _, _)));
}

TEST_F(SpmdPartitioningTest, GatherParallelDimReplicatedOperand) {
  absl::string_view hlo_string = R"(
HloModule module

ENTRY %module {
  %parameter.0 = s32[8,4,2,2]{3,2,1,0} parameter(0), sharding={replicated}
  %iota = s32[1,8,4]{2,1,0} iota(), iota_dimension=1,
    sharding={devices=[1,8,1]0,1,2,3,4,5,6,7}
  %iota2 = s32[1,8,4]{2,1,0} iota(), iota_dimension=2,
    sharding={devices=[1,8,1]0,1,2,3,4,5,6,7}
  %concatenate.19 = s32[2,8,4]{2,1,0} concatenate(s32[1,8,4]{2,1,0} %iota,
    s32[1,8,4]{2,1,0} %iota2), dimensions={0},
    sharding={devices=[1,8,1]0,1,2,3,4,5,6,7}
  ROOT %gather.20 = s32[8,4,2,2]{3,2,1,0} gather(
    s32[8,4,2,2]{3,2,1,0} %parameter.0,
    s32[2,8,4]{2,1,0} %concatenate.19), offset_dims={2,3},
    collapsed_slice_dims={0,1}, start_index_map={0,1}, index_vector_dim=0,
    slice_sizes={1,1,2,2}, sharding={replicated}
})";
  TF_ASSERT_OK_AND_ASSIGN(auto module,
                          PartitionComputation(hlo_string, /*num_devices=*/8));
  auto root = module->entry_computation()->root_instruction();
  VLOG(1) << module->ToString();
  auto operand = AllOf(op::Shape("s32[1,4,2,2]"), op::DynamicSlice());
  auto indices = AllOf(op::Shape("s32[2,1,4]"), op::Subtract());
  auto gather = AllOf(op::Shape("s32[1,4,2,2]"), op::Gather(operand, indices));
  EXPECT_THAT(root,
              op::AllReduce(op::DynamicUpdateSlice(_, gather, _, _, _, _)));
}

TEST_F(SpmdPartitioningTest, GatherParallelDimPartialReplicatedIndices) {
  absl::string_view hlo_string = R"(
HloModule module

ENTRY %module {
  %parameter.0 = s32[8,4,2,2]{3,2,1,0} parameter(0),
    sharding={devices=[8,1,1,1]0,1,2,3,4,5,6,7}
  %iota = s32[1,8,4]{2,1,0} iota(), iota_dimension=1,
    sharding={devices=[1,2,1,4]0,1,2,3,4,5,6,7 last_tile_dim_replicate}
  %iota2 = s32[1,8,4]{2,1,0} iota(), iota_dimension=2,
    sharding={devices=[1,2,1,4]0,1,2,3,4,5,6,7 last_tile_dim_replicate}
  %concatenate.19 = s32[2,8,4]{2,1,0} concatenate(s32[1,8,4]{2,1,0} %iota,
    s32[1,8,4]{2,1,0} %iota2), dimensions={0},
    sharding={devices=[1,2,1,4]0,1,2,3,4,5,6,7 last_tile_dim_replicate}
  ROOT %gather.20 = s32[8,4,2,2]{3,2,1,0} gather(
    s32[8,4,2,2]{3,2,1,0} %parameter.0,
    s32[2,8,4]{2,1,0} %concatenate.19), offset_dims={2,3},
    collapsed_slice_dims={0,1}, start_index_map={0,1}, index_vector_dim=0,
    slice_sizes={1,1,2,2}, sharding={replicated}
})";
  TF_ASSERT_OK_AND_ASSIGN(auto module,
                          PartitionComputation(hlo_string, /*num_devices=*/8));
  auto root = module->entry_computation()->root_instruction();
  VLOG(1) << module->ToString();
  auto operand = AllOf(op::Shape("s32[1,4,2,2]"), op::Parameter());
  auto indices = AllOf(op::Shape("s32[2,1,4]"), op::Subtract());
  auto gather = AllOf(op::Shape("s32[1,4,2,2]"), op::Gather(operand, indices));
  EXPECT_THAT(root,
              op::AllReduce(op::DynamicUpdateSlice(_, gather, _, _, _, _)));
}

TEST_F(SpmdPartitioningTest, GatherParallelDimPartialReplicatedOperand) {
  absl::string_view hlo_string = R"(
HloModule module

ENTRY %module {
  %parameter.0 = s32[8,4,2,2]{3,2,1,0} parameter(0), sharding={
    devices=[2,1,1,1,4]0,1,2,3,4,5,6,7 last_tile_dim_replicate}
  %iota = s32[1,8,4]{2,1,0} iota(), iota_dimension=1,
    sharding={devices=[1,8,1]0,1,2,3,4,5,6,7}
  %iota2 = s32[1,8,4]{2,1,0} iota(), iota_dimension=2,
    sharding={devices=[1,8,1]0,1,2,3,4,5,6,7}
  %concatenate.19 = s32[2,8,4]{2,1,0} concatenate(s32[1,8,4]{2,1,0} %iota,
    s32[1,8,4]{2,1,0} %iota2), dimensions={0},
    sharding={devices=[1,8,1]0,1,2,3,4,5,6,7}
  ROOT %gather.20 = s32[8,4,2,2]{3,2,1,0} gather(
    s32[8,4,2,2]{3,2,1,0} %parameter.0,
    s32[2,8,4]{2,1,0} %concatenate.19), offset_dims={2,3},
    collapsed_slice_dims={0,1}, start_index_map={0,1}, index_vector_dim=0,
    slice_sizes={1,1,2,2}, sharding={replicated}
})";
  TF_ASSERT_OK_AND_ASSIGN(auto module,
                          PartitionComputation(hlo_string, /*num_devices=*/8));
  auto root = module->entry_computation()->root_instruction();
  VLOG(1) << module->ToString();
  auto operand = AllOf(op::Shape("s32[1,4,2,2]"), op::DynamicSlice());
  auto indices = AllOf(op::Shape("s32[2,1,4]"), op::Subtract());
  auto gather = AllOf(op::Shape("s32[1,4,2,2]"), op::Gather(operand, indices));
  EXPECT_THAT(root,
              op::AllReduce(op::DynamicUpdateSlice(_, gather, _, _, _, _)));
}

TEST_F(SpmdPartitioningTest, GatherParallelDimSwappedDimensions) {
  absl::string_view hlo_string = R"(
HloModule module

ENTRY %module {
  %parameter.0 = s32[8,4,2,2]{3,2,1,0} parameter(0), sharding={
    devices=[4,2,1,1]0,1,2,3,4,5,6,7}
  %iota = s32[1,8,4]{2,1,0} iota(), iota_dimension=1,
    sharding={devices=[1,2,4]0,1,2,3,4,5,6,7}
  %iota2 = s32[1,8,4]{2,1,0} iota(), iota_dimension=2,
    sharding={devices=[1,2,4]0,1,2,3,4,5,6,7}
  %concatenate.19 = s32[2,8,4]{2,1,0} concatenate(s32[1,8,4]{2,1,0} %iota,
    s32[1,8,4]{2,1,0} %iota2), dimensions={0},
    sharding={devices=[1,2,4]0,1,2,3,4,5,6,7}
  ROOT %gather.20 = s32[8,4,2,2]{3,2,1,0} gather(
    s32[8,4,2,2]{3,2,1,0} %parameter.0,
    s32[2,8,4]{2,1,0} %concatenate.19), offset_dims={2,3},
    collapsed_slice_dims={0,1}, start_index_map={0,1}, index_vector_dim=0,
    slice_sizes={1,1,2,2}, sharding={replicated}
})";
  TF_ASSERT_OK_AND_ASSIGN(auto module,
                          PartitionComputation(hlo_string, /*num_devices=*/8));
  auto root = module->entry_computation()->root_instruction();
  VLOG(1) << module->ToString();
  auto operand = AllOf(op::Shape("s32[4,1,2,2]"), op::CollectivePermute());
  auto indices = AllOf(op::Shape("s32[2,4,1]"), op::Subtract());
  auto gather = AllOf(op::Shape("s32[4,1,2,2]"), op::Gather(operand, indices));
  EXPECT_THAT(root, op::AllReduce(op::AllReduce(
                        op::DynamicUpdateSlice(_, gather, _, _, _, _))));
}

TEST_F(SpmdPartitioningTest, GatherMergedParalleIndexPassthrough) {
  absl::string_view hlo_string = R"(
HloModule module

ENTRY %module {
  %parameter.0 = s32[8,4,2,2]{3,2,1,0} parameter(0),
    sharding={devices=[2,2,2,1]0,1,2,3,4,5,6,7}
  %iota = s32[1,8,4]{2,1,0} iota(), iota_dimension=2,
    sharding={devices=[1,4,1,2]0,1,2,3,4,5,6,7 last_tile_dim_replicate}
  %iota2 = s32[1,8,4]{2,1,0} iota(), iota_dimension=1,
    sharding={devices=[1,4,1,2]0,1,2,3,4,5,6,7 last_tile_dim_replicate}
  %concatenate.19 = s32[2,8,4]{2,1,0} concatenate(s32[1,8,4]{2,1,0} %iota,
    s32[1,8,4]{2,1,0} %iota2), dimensions={0},
    sharding={devices=[1,4,1,2]0,1,2,3,4,5,6,7 last_tile_dim_replicate}
  ROOT %gather.20 = s32[8,4,2,2]{3,2,1,0} gather(
    s32[8,4,2,2]{3,2,1,0} %parameter.0,
    s32[2,8,4]{2,1,0} %concatenate.19), offset_dims={2,3},
    collapsed_slice_dims={0,1}, start_index_map={1,0}, index_vector_dim=0,
    slice_sizes={1,1,2,2}, sharding={replicated}
})";
  TF_ASSERT_OK_AND_ASSIGN(auto module,
                          PartitionComputation(hlo_string, /*num_devices=*/8));
  auto root = module->entry_computation()->root_instruction();
  auto operand = AllOf(op::Shape("s32[2,4,1,2]"), op::DynamicSlice());
  auto indices = AllOf(op::Shape("s32[2,2,4]"), op::Subtract());
  auto gather = AllOf(op::Shape("s32[2,4,1,2]"), op::Gather(operand, indices));
  EXPECT_THAT(
      root, op::AllReduce(op::DynamicUpdateSlice(
                _, op::AllReduce(op::DynamicUpdateSlice(_, gather, _, _, _, _)),
                _, _, _, _)));
}

TEST_F(SpmdPartitioningTest, GatherMergedParallelIndexTrivialSlice) {
  absl::string_view hlo_string = R"(
HloModule module

ENTRY %module {
  %parameter.0 = s32[8,4,2,2]{3,2,1,0} parameter(0),
    sharding={devices=[4,2,1,1]0,1,2,3,4,5,6,7}
  %parameter.1 = s32[1,8,1]{2,1,0} parameter(1),
    sharding={devices=[1,4,1,2]0,1,2,3,4,5,6,7 last_tile_dim_replicate}
  %iota = s32[1,8,1]{2,1,0} iota(), iota_dimension=1,
    sharding={devices=[1,4,1,2]0,1,2,3,4,5,6,7 last_tile_dim_replicate}
  %concatenate.19 = s32[2,8,1]{2,1,0} concatenate(
    s32[1,8,1]{2,1,0} %parameter.1, s32[1,8,1]{2,1,0} %iota), dimensions={0},
    sharding={devices=[1,4,1,2]0,1,2,3,4,5,6,7 last_tile_dim_replicate}
  ROOT %gather.20 = s32[8,1,2,2]{3,2,1,0} gather(
    s32[8,4,2,2]{3,2,1,0} %parameter.0,
    s32[2,8,1]{2,1,0} %concatenate.19), offset_dims={2,3},
    collapsed_slice_dims={0,1}, start_index_map={1,0}, index_vector_dim=0,
    slice_sizes={1,1,2,2}, sharding={replicated}
})";
  TF_ASSERT_OK_AND_ASSIGN(auto module,
                          PartitionComputation(hlo_string, /*num_devices=*/8));
  auto root = module->entry_computation()->root_instruction();
  auto operand = AllOf(op::Shape("s32[2,2,2,2]"), op::Parameter());
  auto indices = AllOf(op::Shape("s32[2,2,1]"), op::Subtract());
  auto gather = AllOf(op::Shape("s32[2,1,2,2]"), op::Gather(operand, indices));
  EXPECT_THAT(root,
              op::AllReduce(op::DynamicUpdateSlice(
                  _, op::AllReduce(op::Select(_, _, gather)), _, _, _, _)));
}

TEST_F(SpmdPartitioningTest, SortTopKNonSortDimension) {
  absl::string_view hlo_string = R"(
HloModule module

%compare-greater-than.42077 (p.0.lhs.42078: f32[],
  p.0.rhs.42079: f32[], p.1.lhs.42080: s32[], p.1.rhs.42081: s32[]) -> pred[] {
  %p.0.lhs.42078 = f32[] parameter(0)
  %bitcast-convert.135 = s32[] bitcast-convert(f32[] %p.0.lhs.42078)
  %constant.45054 = s32[] constant(0)
  %compare.133 = pred[] compare(s32[] %bitcast-convert.135,
    s32[] %constant.45054), direction=LT
  %constant.45278 = u32[] constant(2147483647)
  %bitcast-convert.136 = u32[] bitcast-convert(f32[] %p.0.lhs.42078)
  %subtract.337 = u32[] subtract(u32[] %constant.45278,
    u32[] %bitcast-convert.136)
  %bitcast-convert.137 = s32[] bitcast-convert(u32[] %subtract.337)
  %select.282 = s32[] select(pred[] %compare.133, s32[] %bitcast-convert.137,
    s32[] %bitcast-convert.135)
  %p.0.rhs.42079 = f32[] parameter(1)
  %bitcast-convert.138 = s32[] bitcast-convert(f32[] %p.0.rhs.42079)
  %compare.134 = pred[] compare(s32[] %bitcast-convert.138,
    s32[] %constant.45054), direction=LT
  %bitcast-convert.139 = u32[] bitcast-convert(f32[] %p.0.rhs.42079)
  %subtract.338 = u32[] subtract(u32[] %constant.45278,
    u32[] %bitcast-convert.139)
  %bitcast-convert.140 = s32[] bitcast-convert(u32[] %subtract.338)
  %select.283 = s32[] select(pred[] %compare.134, s32[] %bitcast-convert.140,
    s32[] %bitcast-convert.138)
  %compare.135 = pred[] compare(s32[] %select.282,
    s32[] %select.283), direction=GT
  %compare.428 = pred[] compare(s32[] %select.283,
    s32[] %select.282), direction=GT
  %compare.429 = pred[] compare(pred[] %compare.135,
    pred[] %compare.428), direction=EQ
  %p.1.lhs.42080 = s32[] parameter(2)
  %p.1.rhs.42081 = s32[] parameter(3)
  %compare.430 = pred[] compare(s32[] %p.1.lhs.42080,
    s32[] %p.1.rhs.42081), direction=LT
  ROOT %select.579 = pred[] select(pred[] %compare.429,
    pred[] %compare.430, pred[] %compare.135)
}

ENTRY %module {
  %parameter.0 = f32[2,64,32128]{2,1,0} parameter(0),
     sharding={devices=[2,1,4]0,1,2,3,4,5,6,7}
  %iota = s32[2,64,32128]{2,1,0} iota(), iota_dimension=2,
    sharding={devices=[2,1,4]0,1,2,3,4,5,6,7}
  %sort.18 = (f32[2,64,32128]{2,1,0}, s32[2,64,32128]{2,1,0}) sort(
    f32[2,64,32128]{2,1,0} %parameter.0, s32[2,64,32128]{2,1,0} %iota),
    dimensions={2}, is_stable=true, to_apply=%compare-greater-than.42077,
    sharding={{devices=[2,1,4]0,1,2,3,4,5,6,7},
    {devices=[2,1,4]0,1,2,3,4,5,6,7}}
  output = f32[2,64,32128]{2,1,0} get-tuple-element(%sort.18), index=0,
    sharding={devices=[2,1,4]0,1,2,3,4,5,6,7}
  %slice.0 = f32[2,64,2]{2,1,0} slice(f32[2,64,32128]{2,1,0} output),
    slice={[0:2], [0:64], [0:2]},
    sharding={devices=[2,1,4]0,1,2,3,4,5,6,7}
  output2 = s32[2,64,32128]{2,1,0} get-tuple-element(%sort.18), index=1,
    sharding={replicated}
  %slice.1 = s32[2,64,2]{2,1,0} slice(s32[2,64,32128]{2,1,0} output2),
    slice={[0:2], [0:64], [0:2]},
    sharding={devices=[2,1,4]0,1,2,3,4,5,6,7}
  ROOT output.t = (f32[2,64,2]{2,1,0},
    s32[2,64,2]{2,1,0}) tuple(slice.0, slice.1),
    sharding={{replicated}, {replicated}}
})";
  TF_ASSERT_OK_AND_ASSIGN(auto module,
                          PartitionComputation(hlo_string, /*num_devices=*/8));

  const HloInstruction* sort = FindInstruction(module.get(), "sort");
  EXPECT_NE(sort, nullptr);
  auto sort_match =
      AllOf(op::Shape("(f32[2,64,32128], s32[2,64,32128])"), op::Sort(_, _));
  EXPECT_THAT(sort, sort_match);
}

TEST_F(SpmdPartitioningTest, SortTopKPropagateBaseShape) {
  absl::string_view hlo_string = R"(
HloModule module

%compare-greater-than.42077 (p.0.lhs.42078: f32[],
  p.0.rhs.42079: f32[], p.1.lhs.42080: s32[], p.1.rhs.42081: s32[]) -> pred[] {
  %p.0.lhs.42078 = f32[] parameter(0)
  %bitcast-convert.135 = s32[] bitcast-convert(f32[] %p.0.lhs.42078)
  %constant.45054 = s32[] constant(0)
  %compare.133 = pred[] compare(s32[] %bitcast-convert.135,
    s32[] %constant.45054), direction=LT
  %constant.45278 = u32[] constant(2147483647)
  %bitcast-convert.136 = u32[] bitcast-convert(f32[] %p.0.lhs.42078)
  %subtract.337 = u32[] subtract(u32[] %constant.45278,
    u32[] %bitcast-convert.136)
  %bitcast-convert.137 = s32[] bitcast-convert(u32[] %subtract.337)
  %select.282 = s32[] select(pred[] %compare.133, s32[] %bitcast-convert.137,
    s32[] %bitcast-convert.135)
  %p.0.rhs.42079 = f32[] parameter(1)
  %bitcast-convert.138 = s32[] bitcast-convert(f32[] %p.0.rhs.42079)
  %compare.134 = pred[] compare(s32[] %bitcast-convert.138,
    s32[] %constant.45054), direction=LT
  %bitcast-convert.139 = u32[] bitcast-convert(f32[] %p.0.rhs.42079)
  %subtract.338 = u32[] subtract(u32[] %constant.45278,
    u32[] %bitcast-convert.139)
  %bitcast-convert.140 = s32[] bitcast-convert(u32[] %subtract.338)
  %select.283 = s32[] select(pred[] %compare.134, s32[] %bitcast-convert.140,
    s32[] %bitcast-convert.138)
  %compare.135 = pred[] compare(s32[] %select.282,
    s32[] %select.283), direction=GT
  %compare.428 = pred[] compare(s32[] %select.283,
    s32[] %select.282), direction=GT
  %compare.429 = pred[] compare(pred[] %compare.135,
    pred[] %compare.428), direction=EQ
  %p.1.lhs.42080 = s32[] parameter(2)
  %p.1.rhs.42081 = s32[] parameter(3)
  %compare.430 = pred[] compare(s32[] %p.1.lhs.42080,
    s32[] %p.1.rhs.42081), direction=LT
  ROOT %select.579 = pred[] select(pred[] %compare.429,
    pred[] %compare.430, pred[] %compare.135)
}

ENTRY %module {
  %parameter.0 = f32[2,64,32128]{2,1,0} parameter(0),
     sharding={devices=[1,1,8]0,1,2,3,4,5,6,7}
  %iota = s32[2,64,32128]{2,1,0} iota(), iota_dimension=2,
    sharding={devices=[1,1,8]0,1,2,3,4,5,6,7}
  %sort.18 = (f32[2,64,32128]{2,1,0}, s32[2,64,32128]{2,1,0}) sort(
    f32[2,64,32128]{2,1,0} %parameter.0, s32[2,64,32128]{2,1,0} %iota),
    dimensions={2}, is_stable=true, to_apply=%compare-greater-than.42077,
    sharding={{devices=[1,1,8]0,1,2,3,4,5,6,7},
    {devices=[1,1,8]0,1,2,3,4,5,6,7}}
  output = f32[2,64,32128]{2,1,0} get-tuple-element(%sort.18), index=0,
    sharding={devices=[1,1,8]0,1,2,3,4,5,6,7}
  %slice.0 = f32[2,64,2]{2,1,0} slice(f32[2,64,32128]{2,1,0} output),
    slice={[0:2], [0:64], [0:2]},
    sharding={devices=[1,1,8]0,1,2,3,4,5,6,7}
  output2 = s32[2,64,32128]{2,1,0} get-tuple-element(%sort.18), index=1,
    sharding={replicated}
  %slice.1 = s32[2,64,2]{2,1,0} slice(s32[2,64,32128]{2,1,0} output2),
    slice={[0:2], [0:64], [0:2]},
    sharding={devices=[1,1,8]0,1,2,3,4,5,6,7}
  ROOT output.t = (f32[2,64,2]{2,1,0},
    s32[2,64,2]{2,1,0}) tuple(slice.0, slice.1),
    sharding={{replicated}, {replicated}}
})";
  TF_ASSERT_OK_AND_ASSIGN(auto module,
                          PartitionComputation(hlo_string, /*num_devices=*/8));

  const HloInstruction* root = module->entry_computation()->root_instruction();
  auto all_reduce_val =
      AllOf(op::Shape("f32[2,64,2]"),
            op::Slice(op::AllReduce(op::DynamicUpdateSlice(_, _, _, _, _))));
  auto all_reduce_idx =
      AllOf(op::Shape("s32[2,64,2]"),
            op::Slice(op::AllReduce(op::DynamicUpdateSlice(_, _, _, _, _))));
  auto tuple = AllOf(op::Shape("(f32[2,64,2], s32[2,64,2])"),
                     op::Tuple(all_reduce_val, all_reduce_idx));
  EXPECT_THAT(root, tuple);
}

TEST_F(SpmdPartitioningTest, GatherIndexOnlyCorrectReplacement) {
  absl::string_view hlo_string = R"(
HloModule module

ENTRY %module {
  %parameter.0 = bf16[1,8,6,6]{3,2,1,0} parameter(0),
    sharding={replicated}
  %parameter.1 = s32[2,4]{1,0} parameter(1),
     sharding={devices=[2,1,4]0,1,2,3,4,5,6,7 last_tile_dim_replicate}
  %gather.100 = bf16[2,1,8,1,6]{4,3,2,1,0} gather(
    bf16[1,8,6,6]{3,2,1,0} %parameter.0, s32[2,4]{1,0} %parameter.1),
    offset_dims={1,2,3,4}, collapsed_slice_dims={}, start_index_map={0,1,2,3},
    index_vector_dim=1, slice_sizes={1,8,1,6},
    sharding={devices=[2,1,4,1,1]0,1,2,3,4,5,6,7}
  %constant.45590 = s32[] constant(0), sharding={replicated}
  %broadcast.54515 = s32[2,64,1,1]{3,2,1,0} broadcast(s32[] %constant.45590),
    dimensions={},
    sharding={devices=[2,1,1,1,4]0,1,2,3,4,5,6,7 last_tile_dim_replicate}
  ROOT %reshape.4243 = bf16[2,8,6]{2,1,0} reshape(
    bf16[2,1,8,1,6]{4,3,2,1,0} %gather.100),
    sharding={devices=[2,4,1]0,1,2,3,4,5,6,7}
})";
  TF_ASSERT_OK_AND_ASSIGN(auto module,
                          PartitionComputation(hlo_string, /*num_devices=*/8));

  const HloInstruction* root = module->entry_computation()->root_instruction();
  auto param0 = AllOf(op::Shape("bf16[1,8,6,6]"), op::Parameter());
  auto param1 = AllOf(op::Shape("s32[1,4]"), op::Parameter());
  auto reshape = AllOf(
      op::Shape("bf16[1,2,6]"),
      op::Reshape(op::DynamicSlice(op::Gather(param0, param1), _, _, _, _, _)));
  EXPECT_THAT(root, reshape);
}

TEST_F(SpmdPartitioningTest, WindowedEinsumPreferMemoryFootprint) {
  absl::string_view hlo_string = R"(
HloModule module

ENTRY %module {
  %parameter.0 = bf16[128,1024,4,4,1152,1,1]{6,5,4,3,2,1,0} parameter(0),
    sharding={devices=[4,1,2,1,1,1,1]0,1,2,3,4,5,6,7}
  %parameter.1 = bf16[4,4,1152,4,176,256,1]{6,5,4,3,2,1,0} parameter(1),
    sharding={devices=[2,2,1,2,1,1,1]0,1,2,3,4,5,6,7}
  %convolution.3 = bf16[128,1024,4,176,256,1,1]{6,5,4,3,2,1,0}
    convolution(bf16[128,1024,4,4,1152,1,1]{6,5,4,3,2,1,0} %parameter.0,
    bf16[4,4,1152,4,176,256,1]{6,5,4,3,2,1,0} %parameter.1),
    window={size=1x4x176x4x4 pad=0_0x3_3x175_175x0_0x0_0
    rhs_reversal=0x1x1x0x0}, dim_labels=0b34f12_34i12o0->0b12f34,
    sharding={devices=[4,1,2,1,1,1,1]0,1,2,3,4,5,6,7}
  ROOT %reshape.3973 = bf16[128,1024,4,176,256]{4,3,2,1,0}
    reshape(bf16[128,1024,4,176,256,1,1]{6,5,4,3,2,1,0} %convolution.3),
    sharding={replicated}
})";
  TF_ASSERT_OK_AND_ASSIGN(
      auto module,
      PartitionComputation(hlo_string, /*num_devices=*/8,
                           /*conv_halo_exchange_always_on_lhs =*/true,
                           /*choose_faster_windowed_einsum =*/false));
  const HloInstruction* while_inst = FindInstruction(module.get(), "while");
  EXPECT_NE(while_inst, nullptr);
  const HloComputation* cond_comp = while_inst->while_condition();
  const HloInstruction* root = cond_comp->root_instruction();
  EXPECT_THAT(root, op::Compare(_, op::Constant()));
  const HloConstantInstruction* iterations =
      Cast<HloConstantInstruction>(root->operand(1));
  EXPECT_TRUE(iterations->literal().GetFirstInteger());
  EXPECT_EQ(*iterations->literal().GetFirstInteger(), 4);
}

TEST_F(SpmdPartitioningTest, WindowedEinsumPreferNumberIterations) {
  absl::string_view hlo_string = R"(
HloModule module

ENTRY %module {
  %parameter.0 = bf16[128,1024,4,4,1152,1,1]{6,5,4,3,2,1,0} parameter(0),
    sharding={devices=[4,1,2,1,1,1,1]0,1,2,3,4,5,6,7}
  %parameter.1 = bf16[4,4,1152,4,176,256,1]{6,5,4,3,2,1,0} parameter(1),
    sharding={devices=[2,2,1,2,1,1,1]0,1,2,3,4,5,6,7}
  %convolution.3 = bf16[128,1024,4,176,256,1,1]{6,5,4,3,2,1,0}
    convolution(bf16[128,1024,4,4,1152,1,1]{6,5,4,3,2,1,0} %parameter.0,
    bf16[4,4,1152,4,176,256,1]{6,5,4,3,2,1,0} %parameter.1),
    window={size=1x4x176x4x4 pad=0_0x3_3x175_175x0_0x0_0
    rhs_reversal=0x1x1x0x0}, dim_labels=0b34f12_34i12o0->0b12f34,
    sharding={devices=[4,1,2,1,1,1,1]0,1,2,3,4,5,6,7}
  ROOT %reshape.3973 = bf16[128,1024,4,176,256]{4,3,2,1,0}
    reshape(bf16[128,1024,4,176,256,1,1]{6,5,4,3,2,1,0} %convolution.3),
    sharding={replicated}
})";
  TF_ASSERT_OK_AND_ASSIGN(
      auto module,
      PartitionComputation(hlo_string, /*num_devices=*/8,
                           /*conv_halo_exchange_always_on_lhs =*/true,
                           /*choose_faster_windowed_einsum =*/true));
  const HloInstruction* while_inst = FindInstruction(module.get(), "while");
  EXPECT_NE(while_inst, nullptr);
  const HloComputation* cond_comp = while_inst->while_condition();
  const HloInstruction* root = cond_comp->root_instruction();
  EXPECT_THAT(root, op::Compare(_, op::Constant()));
  const HloConstantInstruction* iterations =
      Cast<HloConstantInstruction>(root->operand(1));
  EXPECT_TRUE(iterations->literal().GetFirstInteger());
  EXPECT_EQ(*iterations->literal().GetFirstInteger(), 2);
}

TEST_F(SpmdPartitioningTest, WindowedEinsumPreferNumberIterations2) {
  const char* const hlo_string = R"(
HloModule module

ENTRY entry {
  %lhs = bf16[512,1024,16,36,256]{4,3,2,1,0} parameter(0)
  %lhs.copy = bf16[512,1024,16,36,256]{4,3,2,1,0} copy(%lhs),
  sharding={devices=[8,1,4,1,1,1,1]0,1,2,3,4,5,6,7,8,9,10,11,12,13,14,15,16,17,
            18,19,20,21,22,23,24,25,26,27,28,29,30,31}
  %rhs = bf16[512,1024,16,4,288]{4,3,2,1,0} parameter(1)
  %rhs.copy = bf16[512,1024,16,4,288]{4,3,2,1,0} copy(%rhs),
    sharding={devices=[8,1,4,1,1,1,1]0,1,2,3,4,5,6,7,8,9,10,11,12,13,14,15,16,
              17,18,19,20,21,22,23,24,25,26,27,28,29,30,31}
  %reshape.2556 = bf16[512,1024,16,4,288,1,1]{6,5,4,3,2,1,0} reshape(
    bf16[512,1024,16,4,288]{4,3,2,1,0} %rhs.copy), sharding={
      devices=[8,1,4,1,1,1,1]0,1,2,3,4,5,6,7,8,9,10,11,12,13,14,15,16,17,18,19,
        20,21,22,23,24,25,26,27,28,29,30,31}
  %reshape.2570 = bf16[512,1024,16,36,256,1,1]{6,5,4,3,2,1,0}
    reshape(bf16[512,1024,16,36,256]{4,3,2,1,0} %lhs.copy), sharding={
    devices=[8,1,4,1,1,1,1]0,1,2,3,4,5,6,7,8,9,10,11,12,13,14,15,16,17,18,19,
             20,21,22,23,24,25,26,27,28,29,30,31}
  %convolution.10 = bf16[16,36,256,16,4,288,1]{6,5,4,3,2,1,0}
    convolution(bf16[512,1024,16,36,256,1,1]{6,5,4,3,2,1,0} %reshape.2570,
    bf16[512,1024,16,4,288,1,1]{6,5,4,3,2,1,0} %reshape.2556),
    window={size=1x1x16x4x512 pad=0_0x0_0x15_15x3_3x0_0 rhs_reversal=0x0x1x1x0},
    dim_labels=4f01b23_4i23o01->01b23f4, sharding={devices=[4,1,1,4,2,1,1]0,4,8,
    12,16,20,24,28,1,5,9,13,17,21,25,29,2,6,10,14,18,22,26,30,3,7,11,15,19,23,
    27,31}
  ROOT %output = bf16[16,36,256,16,4,288,1]{6,5,4,3,2,1,0}
   copy(%convolution.10), sharding={replicated}
})";
  TF_ASSERT_OK_AND_ASSIGN(
      auto module,
      PartitionComputation(hlo_string, /*num_devices=*/32,
                           /*conv_halo_exchange_always_on_lhs =*/true,
                           /*choose_faster_windowed_einsum =*/true));
  const HloInstruction* while_inst = FindInstruction(module.get(), "while");
  EXPECT_NE(while_inst, nullptr);
  const HloComputation* cond_comp = while_inst->while_condition();
  const HloInstruction* root = cond_comp->root_instruction();
  EXPECT_THAT(root, op::Compare(_, op::Constant()));
  const HloConstantInstruction* iterations =
      Cast<HloConstantInstruction>(root->operand(1));
  EXPECT_TRUE(iterations->literal().GetFirstInteger());
  EXPECT_EQ(*iterations->literal().GetFirstInteger(), 4);
}

TEST_F(SpmdPartitioningTest, WindowedEinsumPreferMemoryFootprint2) {
  const char* const hlo_string = R"(
HloModule module

ENTRY entry {
  %lhs = bf16[512,1024,16,36,256]{4,3,2,1,0} parameter(0)
  %lhs.copy = bf16[512,1024,16,36,256]{4,3,2,1,0} copy(%lhs),
  sharding={devices=[8,1,4,1,1,1,1]0,1,2,3,4,5,6,7,8,9,10,11,12,13,14,15,16,17,
            18,19,20,21,22,23,24,25,26,27,28,29,30,31}
  %rhs = bf16[512,1024,16,4,288]{4,3,2,1,0} parameter(1)
  %rhs.copy = bf16[512,1024,16,4,288]{4,3,2,1,0} copy(%rhs),
    sharding={devices=[8,1,4,1,1,1,1]0,1,2,3,4,5,6,7,8,9,10,11,12,13,14,15,16,
              17,18,19,20,21,22,23,24,25,26,27,28,29,30,31}
  %reshape.2556 = bf16[512,1024,16,4,288,1,1]{6,5,4,3,2,1,0} reshape(
    bf16[512,1024,16,4,288]{4,3,2,1,0} %rhs.copy), sharding={
      devices=[8,1,4,1,1,1,1]0,1,2,3,4,5,6,7,8,9,10,11,12,13,14,15,16,17,18,19,
        20,21,22,23,24,25,26,27,28,29,30,31}
  %reshape.2570 = bf16[512,1024,16,36,256,1,1]{6,5,4,3,2,1,0}
    reshape(bf16[512,1024,16,36,256]{4,3,2,1,0} %lhs.copy), sharding={
    devices=[8,1,4,1,1,1,1]0,1,2,3,4,5,6,7,8,9,10,11,12,13,14,15,16,17,18,19,
             20,21,22,23,24,25,26,27,28,29,30,31}
  %convolution.10 = bf16[16,36,256,16,4,288,1]{6,5,4,3,2,1,0}
    convolution(bf16[512,1024,16,36,256,1,1]{6,5,4,3,2,1,0} %reshape.2570,
    bf16[512,1024,16,4,288,1,1]{6,5,4,3,2,1,0} %reshape.2556),
    window={size=1x1x16x4x512 pad=0_0x0_0x15_15x3_3x0_0 rhs_reversal=0x0x1x1x0},
    dim_labels=4f01b23_4i23o01->01b23f4, sharding={devices=[4,1,1,4,2,1,1]0,4,8,
    12,16,20,24,28,1,5,9,13,17,21,25,29,2,6,10,14,18,22,26,30,3,7,11,15,19,23,
    27,31}
  ROOT %output = bf16[16,36,256,16,4,288,1]{6,5,4,3,2,1,0}
   copy(%convolution.10), sharding={replicated}
})";
  TF_ASSERT_OK_AND_ASSIGN(
      auto module,
      PartitionComputation(hlo_string, /*num_devices=*/32,
                           /*conv_halo_exchange_always_on_lhs =*/true,
                           /*choose_faster_windowed_einsum =*/false));
  const HloInstruction* while_inst = FindInstruction(module.get(), "while");
  EXPECT_NE(while_inst, nullptr);
  const HloComputation* cond_comp = while_inst->while_condition();
  const HloInstruction* root = cond_comp->root_instruction();
  EXPECT_THAT(root, op::Compare(_, op::Constant()));
  const HloConstantInstruction* iterations =
      Cast<HloConstantInstruction>(root->operand(1));
  EXPECT_TRUE(iterations->literal().GetFirstInteger());
  EXPECT_EQ(*iterations->literal().GetFirstInteger(), 8);
}

TEST_F(SpmdPartitioningTest, ContractingPartitionDotOperandsSlicedWrong) {
  const char* const hlo_string = R"(
HloModule module

ENTRY entry {
  %lhs = f32[8,2,15,4] parameter(0)
  %lhs.copy = f32[8,2,15,4] copy(%lhs),
    sharding={devices=[1,2,4,1]0,1,2,3,4,5,6,7}
  %rhs = f32[2,15,4] parameter(1)
  %rhs.copy = f32[2,15,4] copy(%rhs),
    sharding={devices=[2,4,1]0,1,2,3,4,5,6,7}
  %dot = f32[8,2,2] dot(%lhs.copy, %rhs.copy),
    lhs_batch_dims={}, rhs_batch_dims={},
    lhs_contracting_dims={2,3}, rhs_contracting_dims={1,2},
    operand_precision={HIGH,HIGH},
    sharding={devices=[2,2,2]0,1,2,3,4,5,6,7}
  ROOT %output = f32[8,2,2] copy(%dot), sharding={replicated}
})";
  TF_ASSERT_OK_AND_ASSIGN(
      auto module,
      PartitionComputation(hlo_string, /*num_devices=*/8,
                           /*conv_halo_exchange_always_on_lhs =*/true,
                           /*choose_faster_windowed_einsum =*/true));

  const HloInstruction* dot_op = FindInstruction(module.get(), "dot.1");
  auto op1 = op::Shape("f32[4,2,4,4]");
  auto op2 = op::Shape("f32[2,4,4]");
  EXPECT_THAT(dot_op, op::Dot(op1, op2));
}

TEST_F(SpmdPartitioningTest, PartitionDotGroupOnBatchContractingReshard) {
  absl::string_view hlo_string = R"(
HloModule module

ENTRY entry {
  %lhs = f32[32,32,24,4096] parameter(0),
    sharding={devices=[2,1,1,2]0,1,2,3}
  %rhs = f32[32,4096,1024] parameter(1),
    sharding={devices=[2,2,1]0,1,2,3}
  ROOT %dot = f32[32,32,24,1024] dot(%lhs, %rhs),
    lhs_batch_dims={0}, rhs_batch_dims={0},
    lhs_contracting_dims={3}, rhs_contracting_dims={1},
    sharding={devices=[1,2,1,2]0,1,2,3}
})";

  TF_ASSERT_OK_AND_ASSIGN(auto module,
                          PartitionComputation(hlo_string, /*num_devices=*/4));
  VLOG(1) << module->ToString();
  auto root = module->entry_computation()->root_instruction();
  auto dot = AllOf(op::Shape("f32[16,32,24,1024]"),
                   op::Dot(op::Parameter(0), op::Parameter(1)));
  auto reduce_scatter = AllOf(op::Shape("f32[16,32,24,512]"),
                              op::DynamicSlice(op::AllReduce(dot), _, _, _, _));
  EXPECT_THAT(root, AllOf(op::Reshape(op::Transpose(
                              op::AllToAll(op::Reshape(reduce_scatter)))),
                          op::Shape("f32[32,16,24,512]")));
}

TEST_F(SpmdPartitioningTest, PartitionPassthroughScatterCorrectOutputSharding) {
  absl::string_view hlo_string = R"(
HloModule module

%scatter_add (parameter.0: bf16[], parameter.1: bf16[]) -> bf16[] {
  %parameter.0 = bf16[] parameter(0)
  %parameter.1 = bf16[] parameter(1)
  ROOT %add = bf16[] add(bf16[] %parameter.0, bf16[] %parameter.1)
}

ENTRY entry {
  %operand = bf16[2,1024]{1,0} parameter(0),
    sharding={devices=[1,2,2]0,1,2,3 last_tile_dim_replicate}
  %indices = s32[8,512,1]{2,1,0} parameter(1),
    sharding={devices=[2,1,1,2]0,2,1,3 last_tile_dim_replicate}
  %updates = bf16[8,512,1024]{2,1,0} parameter(2),
    sharding={devices=[2,1,2]0,2,1,3}
  ROOT %scatter = bf16[2,1024]{1,0} scatter(bf16[2,1024]{1,0} %operand,
    s32[8,512,1]{2,1,0} %indices,
    bf16[8,512,1024]{2,1,0} %updates), update_window_dims={2},
    inserted_window_dims={0}, scatter_dims_to_operand_dims={0},
    index_vector_dim=2, to_apply=%scatter_add,
    sharding={devices=[1,2,2]0,1,2,3 last_tile_dim_replicate}
})";

  TF_ASSERT_OK_AND_ASSIGN(auto module,
                          PartitionComputation(hlo_string, /*num_devices=*/4));
  VLOG(1) << module->ToString();
  auto root = module->entry_computation()->root_instruction();
  auto scatter = AllOf(op::Shape("bf16[2,512]"), op::Scatter(_, _, _));
  EXPECT_THAT(root, scatter);
}

}  // namespace
}  // namespace spmd
}  // namespace xla
