Improve error message for unsupported RTS config

Summary: This reverts commit df5bdd3.

Test Plan: revert-hammer

Reviewers:

Subscribers:

Author: georgepaw

tensorflow/compiler/plugin/poplar/driver/passes/replicated_resource_update_elementwise_clustering.cc

@@ -280,6 +280,19 @@ ReplicatedResourceUpdateElementwiseClustering::AddClusterInput(
     // replica 3: 0 1 1 1                       1 2                     3 4
     //            -------                   -------                 -------
 
+    if (ipu_link_domain_replication_factor_ < global_replication_factor_ &&
+        partition_replication_factor_ < ipu_link_domain_replication_factor_) {
+      // TODO(T45704): GCL does not support this and will report an error, but
+      // the message might be hard to understand for the user. So we catch it
+      // here to attempt to report a more understandable error message.
+      return Unimplemented(
+          "Replicated partitioning is not supported when there are multiple "
+          "instances per IPU-link domain. The number of local replicas per "
+          "partitioning domain (%u) is less than the number of replicas "
+          "per IPU-link domain (%u).",
+          partition_replication_factor_, ipu_link_domain_replication_factor_);
+    }
+
     CHECK_NE(partition_replication_factor_, 0);
     CHECK_EQ(global_replication_factor_ % partition_replication_factor_, 0);
     const uint64 orthogonal_group_size =
@@ -549,7 +562,7 @@ ReplicatedResourceUpdateElementwiseClustering::UpdateClusterBackendConfig(
           cluster, backend_config));
   auto* call_config = backend_config.mutable_call_config();
   auto* function_config = call_config->mutable_function_config();
-  // Setting parititoned_elementwise_cluster attribute indicates that we will
+  // Setting partitioned_elementwise_cluster attribute indicates that we will
   // process those clusters differently:
   // - Remote buffer outlining pass will outline load/stores regardless if it's
   // unique or not.

tensorflow/compiler/plugin/poplar/driver/passes/replicated_resource_update_elementwise_clustering.h

@@ -33,15 +33,19 @@ class ReplicatedResourceUpdateElementwiseClustering final
  public:
   ReplicatedResourceUpdateElementwiseClustering(
       CompilerAnnotations& annotations, uint32 partition_replication_factor,
-      uint32 global_replication_factor)
+      uint32 global_replication_factor,
+      uint32 ipu_link_domain_replication_factor)
       : annotations_(annotations),
         partition_replication_factor_(partition_replication_factor),
-        global_replication_factor_(global_replication_factor) {}
+        global_replication_factor_(global_replication_factor),
+        ipu_link_domain_replication_factor_(
+            ipu_link_domain_replication_factor) {}
 
   explicit ReplicatedResourceUpdateElementwiseClustering(
       CompilerAnnotations& annotations, uint32 replication_factor)
       : ReplicatedResourceUpdateElementwiseClustering(
-            annotations, replication_factor, replication_factor) {}
+            annotations, replication_factor, replication_factor,
+            replication_factor) {}
 
   absl::string_view name() const override {
     return "replicated-resource-update-elementwise-clustering";
@@ -86,6 +90,7 @@ class ReplicatedResourceUpdateElementwiseClustering final
   CompilerAnnotations& annotations_;
   uint32 partition_replication_factor_;
   uint32 global_replication_factor_;
+  uint32 ipu_link_domain_replication_factor_;
 };
 
 }  // namespace poplarplugin

tensorflow/compiler/plugin/poplar/driver/poplar_compiler.cc

@@ -1228,17 +1228,35 @@ StatusOr<std::unique_ptr<PoplarExecutableCore>> CompileEngine(
         " The number of shards needs to divide the number of local IPUs.");
   }
 
-  CHECK_LE(local_replication_factor, replication_factor);
-
   // Currently we only support performing replica partitioning across the local
   // replicas in each process, as this allows access to all the parts of a
   // partitioned remote buffer locally. This means that copying to/from all the
   // parts of the partitioned remote buffer can be done without any additional
   // inter-process collective communication.
   const auto partition_replication_factor = local_replication_factor;
 
+  // The IPU-link domain size is the number of IPUs per IPU-link domain.
+  // A CPU target cannot be trusted to report a sensible value.
+  const auto num_ipu_link_domain_ipus =
+      target.getTargetType() == poplar::TargetType::CPU
+          ? num_ipus
+          : std::min(num_ipus, target.getIpuLinkDomainSize());
+  CHECK_GE(num_ipu_link_domain_ipus, num_shards);
+  CHECK_EQ(num_ipu_link_domain_ipus % num_shards, 0);
+  const auto ipu_link_domain_replication_factor =
+      num_ipu_link_domain_ipus / num_shards;
+
+  // Replication factor invariant: local <= ipu_link_domain <= global.
+  CHECK_LE(local_replication_factor, ipu_link_domain_replication_factor);
+  CHECK_LE(ipu_link_domain_replication_factor, replication_factor);
+
+  CHECK_EQ(replication_factor % ipu_link_domain_replication_factor, 0);
+  CHECK_EQ(ipu_link_domain_replication_factor % local_replication_factor, 0);
+
   VLOG(1) << "Local replication factor " << local_replication_factor
           << ", global replication factor " << replication_factor
+          << ", IPU-link domain replication factor "
+          << ipu_link_domain_replication_factor
           << ", partition replication factor " << partition_replication_factor;
 
   auto config = module->config();
@@ -1459,7 +1477,7 @@ StatusOr<std::unique_ptr<PoplarExecutableCore>> CompileEngine(
       pipeline.AddPass<PipelineFIFOInserter>(resources.remote_memory_supported);
       pipeline.AddPass<ReplicatedResourceUpdateElementwiseClustering>(
           resources.annotations, resources.partition_replication_factor,
-          resources.replication_factor);
+          resources.replication_factor, ipu_link_domain_replication_factor);
       {
         auto inline_fusion = [](const HloInstruction* inst) {
           return IsReplicatedParameterLoadFusion(inst) ||

tensorflow/compiler/plugin/poplar/tests/replicated_resource_update_elementwise_clustering_test.cc

@@ -1436,7 +1436,8 @@ TEST_F(TestPartitionReplicationFactor, TestCollectiveGroups) {
   EXPECT_TRUE(offloaded);
 
   ReplicatedResourceUpdateElementwiseClustering pass(
-      annotations, partition_replication_factor, global_replication_factor);
+      annotations, partition_replication_factor, global_replication_factor,
+      global_replication_factor);
   auto elementwise_comps =
       ElementwiseCluster::GetElementwiseClusterableComputations(module.get());
   TF_ASSERT_OK_AND_ASSIGN(auto clusters,
@@ -1497,6 +1498,96 @@ TEST_F(TestPartitionReplicationFactor, TestCollectiveGroups) {
                                               partition_replication_factor));
 }
 
+TEST_F(TestPartitionReplicationFactor, TestUnsupportedPartitioning) {
+  const std::string hlo = R"(
+  HloModule main
+
+  sum {
+    y = f16[] parameter(1)
+    x = f16[] parameter(0), control-predecessors={y}
+    ROOT add = f16[] add(x, y), backend_config="{\"isInplace\":true}"
+  }
+
+  resource_update {
+    arg0 = f16[128] parameter(0)
+    arg1 = f16[128] parameter(1)
+    arg2 = f16[128] parameter(2)
+
+    arg0_r = f16[128] all-reduce(arg0), to_apply=sum
+    arg2_new = f16[128] add(arg0_r, arg2)
+    arg1_new = f16[128] add(arg1, arg2_new)
+
+    ROOT t = (f16[128],f16[128]) tuple(arg1_new, arg2_new)
+    counter_0 = s32[] constant(4)
+    gac = () custom-call(s32[] counter_0), custom_call_target="GradientAccumulationCount"
+  }
+
+  loop {
+    after-all = token[] after-all()
+    infeed = (f16[128], token[]) infeed(after-all), infeed_config="140121807314576"
+    input = f16[128] get-tuple-element(infeed), index=0
+
+    arg0 = f16[128] parameter(0)
+    arg1 = f16[128] parameter(1)
+
+    add.1 = f16[128] add(input, arg0)
+    call = (f16[128],f16[128]) call(add.1, arg0, arg1), to_apply=resource_update, frontend_attributes={CALL_CONFIG_TYPE="ResourceUpdate"}, backend_config="{\"callConfig\":{\"type\":\"ResourceUpdate\",\"resourceUpdateConfig\":{\"offloadVariables\":\"THREESTATE_ON\", \"partitionOffloadedVariables\":\"THREESTATE_ON\"}}}"
+    gte0 = f16[128] get-tuple-element(call), index=0
+    gte1 = f16[128] get-tuple-element(call), index=1
+    ROOT r = (f16[128],f16[128]) tuple(gte0, gte1)
+  }
+
+  ENTRY e {
+    e.in0 = f16[128] parameter(0)
+    e.in1 = f16[128] parameter(1)
+    loop_call = (f16[128],f16[128]) call(e.in0, e.in1), to_apply=loop, backend_config="{\"callConfig\":{\"type\":\"RepeatLoop\",\"repeatConfig\":{\"repeatCount\":\"100\"}}}"
+    gte0 = f16[128] get-tuple-element(loop_call), index=0
+    gte1 = f16[128] get-tuple-element(loop_call), index=1
+    ROOT r = (f16[128],f16[128]) tuple(gte0, gte1)
+  }
+  )";
+
+  auto config = GetModuleConfigForTest();
+  config.set_argument_input_indices({});
+  config.set_resource_input_indices({0, 1});
+  config.set_resource_input_initialized({true, true});
+  config.set_resource_update_to_input_index({0, 1});
+  TF_ASSERT_OK_AND_ASSIGN(auto module,
+                          ParseAndReturnVerifiedModule(hlo, config));
+
+  HloInstruction* loop =
+      CHECK_NOTNULL(FindInstruction(module.get(), "loop_call"));
+
+  const uint64 partition_replication_factor = 2;
+  const uint64 ipu_link_domain_replication_factor = 4;
+  const uint64 global_replication_factor = 8;
+
+  CompilerAnnotations annotations(module.get());
+  TF_ASSERT_OK_AND_ASSIGN(
+      bool offloaded,
+      VariablesOffloadAndPartition(
+          annotations, /*remote_memory_supported=*/true,
+          /*minimum_remote_tensor_size=*/4, partition_replication_factor)
+          .Run(module.get()));
+  EXPECT_TRUE(offloaded);
+
+  ReplicatedResourceUpdateElementwiseClustering pass(
+      annotations, partition_replication_factor, global_replication_factor,
+      ipu_link_domain_replication_factor);
+  auto elementwise_comps =
+      ElementwiseCluster::GetElementwiseClusterableComputations(module.get());
+  TF_ASSERT_OK_AND_ASSIGN(auto clusters,
+                          pass.GetClustersIn(loop, elementwise_comps));
+  ASSERT_THAT(clusters.size(), 1);
+  auto& cluster = clusters.front();
+  const auto status = pass.OutlineCluster(cluster).status();
+  ASSERT_THAT(status.code(), tensorflow::error::UNIMPLEMENTED);
+  EXPECT_THAT(status.error_message(),
+              ::testing::StartsWith(
+                  "Replicated partitioning is not supported when there are "
+                  "multiple instances per IPU-link domain."));
+}
+
 TEST_F(TestPartitionReplicationFactor, TestNonGlobalAllReduce) {
   const std::string hlo = R"(
   HloModule main
@@ -1569,7 +1660,8 @@ TEST_F(TestPartitionReplicationFactor, TestNonGlobalAllReduce) {
   EXPECT_TRUE(offloaded);
 
   ReplicatedResourceUpdateElementwiseClustering pass(
-      annotations, partition_replication_factor, global_replication_factor);
+      annotations, partition_replication_factor, global_replication_factor,
+      global_replication_factor);
   auto elementwise_comps =
       ElementwiseCluster::GetElementwiseClusterableComputations(module.get());
   TF_ASSERT_OK_AND_ASSIGN(auto clusters,