Reflect GCL API changes

Frederik Mellbye · Frederik Mellbye · commit bda97007905f · 2022-09-09T09:15:56.000+01:00
Reviewers: #tensorflow, #framework_ip_review_-_any_oss_or_third-party_code_use_has_been_approved, christiana Reviewed By: #tensorflow, #framework_ip_review_-_any_oss_or_third-party_code_use_has_been_approved, christiana Subscribers: danielb Maniphest Tasks: T68415 Differential Revision: https://phabricator.sourcevertex.net/D74214
diff --git a/tensorflow/compiler/plugin/poplar/driver/ops/custom_ops/popops/within_replicas.cc b/tensorflow/compiler/plugin/poplar/driver/ops/custom_ops/popops/within_replicas.cc
@@ -121,6 +121,7 @@ class AllGatherWithinReplicaOp : public PoplarOpDef {
     gcl::Chunks chunks;
 
     std::vector<poplar::Tensor> shards;
+    std::vector<gcl::Chunk> chunks_vec;
     for (auto i = 0u; i < inst->operand_count(); ++i) {
       TF_ASSIGN_OR_RETURN(
           poplar::Tensor input,
@@ -131,14 +132,16 @@ class AllGatherWithinReplicaOp : public PoplarOpDef {
       // to the subtensor after indexing the outermost dimension). Having
       // index=i means that the chunks in our gathered tensor will be in the
       // same order as the inputs.
-      chunks.chunks.push_back({input, /*index*/ i, /*offset*/ 0});
 
+      chunks_vec.emplace_back(input, i, 0);
       shards.push_back(input);
     }
 
+    chunks.setChunks(chunks_vec);
+
     auto original_input = poplar::concat(shards);
-    chunks.originalInput =
-        original_input.expand({0}).broadcast(inst->operand_count(), 0);
+    chunks.setOriginalInput(
+        original_input.expand({0}).broadcast(inst->operand_count(), 0));
     return chunks;
   }
 };
@@ -173,15 +176,15 @@ class ReduceScatterWithinReplicaOp : public PoplarOpDef {
         {debug_info, "ReduceScatterWithinReplica"},
         GetReplicatedCollectiveOptions(res));
 
-    CHECK_EQ(ipu_count, chunks.chunks.size())
+    CHECK_EQ(ipu_count, chunks.getChunks().size())
         << "Expecting to have a chunk for each IPU.";
-    TF_CHECK_OK(SetOutputs(chunks.chunks, inst, res, tensor_map));
+    TF_CHECK_OK(SetOutputs(chunks.getChunks(), inst, res, tensor_map));
 
     return seq;
   }
 
  private:
-  Status SetOutputs(std::vector<gcl::Chunk>& output_chunks,
+  Status SetOutputs(const std::vector<gcl::Chunk>& output_chunks,
                     const HloInstruction* inst, CompilerResources& res,
                     TensorMap& tensor_map) {
     const auto output_tensor_shape =
@@ -191,7 +194,7 @@ class ReduceScatterWithinReplicaOp : public PoplarOpDef {
     const auto output_tensor_size = output_tensor_shape.dimensions(0);
 
     for (auto i = 0; i < output_chunks.size(); ++i) {
-      auto tensor = output_chunks[i].tensor;
+      auto tensor = output_chunks[i].getTensor();
       CHECK_EQ(tensor.rank(), 1);
 
       // Pad everything to a consistent shape. We don't know how
diff --git a/tensorflow/compiler/plugin/poplar/driver/poplar_executable.cc b/tensorflow/compiler/plugin/poplar/driver/poplar_executable.cc
@@ -72,19 +72,24 @@ void PoplarExecutableCore::PopulateCollectiveBalanceReorderHostRerrangements() {
   for (auto& param_host_rearrangement_entry :
        GetRemoteParameterHostRearrangements()) {
     int64_t id = param_host_rearrangement_entry.first;
-    auto& param_host_rearrangement = param_host_rearrangement_entry.second;
+    const auto& param_host_rearrangement =
+        param_host_rearrangement_entry.second;
     gcl::CollectiveBalancedHostRearrangement host_rearrangement;
-    host_rearrangement.replicationFactor =
-        param_host_rearrangement.replication_factor;
-    host_rearrangement.totalElementsPerReplica =
-        param_host_rearrangement.total_elements_per_replica;
-    host_rearrangement.gatheredToRefSlices.reserve(
+    host_rearrangement.setReplicationFactor(
+        param_host_rearrangement.replication_factor);
+    host_rearrangement.setTotalElementsPerReplica(
+        param_host_rearrangement.total_elements_per_replica);
+
+    std::vector<poplar::Interval> gathered_to_ref_slices(
+        host_rearrangement.getGatheredToRefSlices());
+    gathered_to_ref_slices.reserve(
         param_host_rearrangement.gathered_to_ref_slice.size());
-    for (auto& slice : param_host_rearrangement.gathered_to_ref_slice) {
-      host_rearrangement.gatheredToRefSlices.emplace_back(slice.first,
-                                                          slice.second);
+
+    for (const auto& slice : param_host_rearrangement.gathered_to_ref_slice) {
+      gathered_to_ref_slices.emplace_back(slice.first, slice.second);
     }
-    host_rearrangement.elementMap = param_host_rearrangement.element_map;
+    host_rearrangement.setGatheredToRefSlices(gathered_to_ref_slices);
+    host_rearrangement.setElementMap(param_host_rearrangement.element_map);
     cbr_host_rearrangements_[id] = std::move(host_rearrangement);
   }
 }
diff --git a/tensorflow/compiler/plugin/poplar/driver/poplar_executor.cc b/tensorflow/compiler/plugin/poplar/driver/poplar_executor.cc
@@ -307,8 +307,8 @@ PoplarExecutor::TensorControl::TensorControl(size_t size_) {
 std::size_t PoplarExecutor::TensorControl::GetRemoteBufferSize() const {
   if (host_rearrangement) {
     return ShapeUtil::ByteSizeOf(ShapeUtil::MakeShape(
-        element_type, {host_rearrangement->replicationFactor,
-                       host_rearrangement->totalElementsPerReplica}));
+        element_type, {host_rearrangement->getReplicationFactor(),
+                       host_rearrangement->getTotalElementsPerReplica()}));
   }
   return size;
 }
@@ -2848,11 +2848,12 @@ Status PoplarExecutor::MoveDeviceToHost() {
                 // code.
                 std::vector<char> temp(buffer.size());
                 tc->host_rearrangement->undoRearrangeForCollective(
-                    buffer.data(), temp.data(), bytes_per_element);
+                    buffer, temp, bytes_per_element);
                 memcpy(tc->data, temp.data(), tc->size);
               } else {
                 tc->host_rearrangement->undoRearrangeForCollective(
-                    buffer.data(), tc->data, bytes_per_element);
+                    buffer.data(), buffer.size(), tc->data, tc->size,
+                    bytes_per_element);
               }
             }
           } else {
@@ -2975,10 +2976,11 @@ Status PoplarExecutor::MoveHostToDevice() {
                 std::vector<char> temp(buffer_size);
                 memcpy(temp.data(), tc->data, tc->size);
                 tc->host_rearrangement->rearrangeForCollective(
-                    temp.data(), buffer.data(), bytes_per_element);
+                    temp, buffer, bytes_per_element);
               } else {
                 tc->host_rearrangement->rearrangeForCollective(
-                    tc->data, buffer.data(), bytes_per_element);
+                    tc->data, tc->size, buffer.data(), buffer.size(),
+                    bytes_per_element);
               }
             } else {
               buffer.resize(bytes_per_replica);
diff --git a/tensorflow/compiler/plugin/poplar/driver/visitors/partitioned_elementwise_cluster_visitor.cc b/tensorflow/compiler/plugin/poplar/driver/visitors/partitioned_elementwise_cluster_visitor.cc
@@ -138,9 +138,9 @@ Status PartitionedElementwiseClusterVisitor::ValidateShape(
   const gcl::CollectiveBalancedHostRearrangement& host_rearrangement =
       cbr->getHostRearrangement();
   const int64_t replicated_element_count =
-      host_rearrangement.totalElementsPerReplica;
+      host_rearrangement.getTotalElementsPerReplica();
   const int64_t non_replicated_element_count =
-      replicated_element_count * host_rearrangement.replicationFactor;
+      replicated_element_count * host_rearrangement.getReplicationFactor();
   const int64_t xla_element_count = ShapeUtil::ElementsIn(shape);
   VLOG(3) << "CBR slice element count: " << replicated_element_count
           << ", total collectives elements: " << non_replicated_element_count
@@ -200,7 +200,7 @@ PartitionedElementwiseClusterVisitor::MakeParameterAllocationFunction(
       TF_ASSIGN_OR_RETURN(poplar::Type type, PoplarDataType(shape));
       auto element_count = ShapeUtil::ElementsIn(shape);
       auto& host_rearrangement = cbr->getHostRearrangement();
-      if (host_rearrangement.totalElementsPerReplica == element_count) {
+      if (host_rearrangement.getTotalElementsPerReplica() == element_count) {
         tensor_like = DriverTensor(cbr->createReplicaSlice(type));
       }
     }
@@ -291,13 +291,16 @@ PartitionedElementwiseClusterVisitor::UpdateRemoteBufferInformation(
           cbr_info->host_rearrangement_id) ==
       remote_parameter_host_rearrangements.end()) {
     RemoteParameterHostRearrangement host_rearrangement;
-    host_rearrangement.replication_factor = src.replicationFactor;
-    host_rearrangement.total_elements_per_replica = src.totalElementsPerReplica;
-    for (auto& interval : src.gatheredToRefSlices) {
+    host_rearrangement.replication_factor = src.getReplicationFactor();
+    host_rearrangement.total_elements_per_replica =
+        src.getTotalElementsPerReplica();
+    for (const auto& interval : src.getGatheredToRefSlices()) {
       host_rearrangement.gathered_to_ref_slice.emplace_back(interval.begin(),
                                                             interval.end());
     }
-    host_rearrangement.element_map = src.elementMap;
+
+    host_rearrangement.element_map = src.getElementMap();
+
     remote_parameter_host_rearrangements[cbr_info->host_rearrangement_id] =
         std::move(host_rearrangement);
   }
@@ -314,7 +317,7 @@ PartitionedElementwiseClusterVisitor::UpdateRemoteBufferInformation(
   for (auto param_idx : merged_params) {
     TF_RETURN_IF_ERROR(SetRemoteBufferHostRearrangementId(
         graph, entry_comp, param_idx, cbr_info->host_rearrangement_id,
-        src.totalElementsPerReplica));
+        src.getTotalElementsPerReplica()));
   }
 
   return true;
diff --git a/tensorflow/compiler/plugin/poplar/tests/replicated_resource_update_elementwise_clustering_hw_test.cc b/tensorflow/compiler/plugin/poplar/tests/replicated_resource_update_elementwise_clustering_hw_test.cc
@@ -231,26 +231,29 @@ class ReplicatedResourceUpdateElementwiseClusteringHwTest
                 info.host_rearrangement_id);
             CHECK(hr_it !=
                   annotations.remote_parameter_host_rearrangements.end());
-            auto& host_rearrangement = hr_it->second;
+            const auto& host_rearrangement = hr_it->second;
             auto& gcl = host_rearrangements[index];
-            gcl.replicationFactor = host_rearrangement.replication_factor;
-            gcl.totalElementsPerReplica =
-                host_rearrangement.total_elements_per_replica;
-            for (auto& slice : host_rearrangement.gathered_to_ref_slice) {
-              gcl.gatheredToRefSlices.emplace_back(slice.first, slice.second);
+            gcl.setReplicationFactor(host_rearrangement.replication_factor);
+            gcl.setTotalElementsPerReplica(
+                host_rearrangement.total_elements_per_replica);
+            std::vector<poplar::Interval> gathered_to_ref_slices =
+                gcl.getGatheredToRefSlices();
+
+            for (const auto& slice : host_rearrangement.gathered_to_ref_slice) {
+              gathered_to_ref_slices.emplace_back(slice.first, slice.second);
             }
-            gcl.elementMap = host_rearrangement.element_map;
 
-            CHECK_EQ(gcl.replicationFactor, param.replication_factor);
-            per_replica_size = gcl.totalElementsPerReplica;
-            aligned_size = per_replica_size * gcl.replicationFactor;
+            gcl.setGatheredToRefSlices(gathered_to_ref_slices);
+            gcl.setElementMap(host_rearrangement.element_map);
+
+            CHECK_EQ(gcl.getReplicationFactor(), param.replication_factor);
+            per_replica_size = gcl.getTotalElementsPerReplica();
+            aligned_size = per_replica_size * gcl.getReplicationFactor();
             buffer.resize(aligned_size);
 
             std::vector<float> tmp(aligned_size);
             VLOG(1) << "Rearranging for collective...";
-            gcl.rearrangeForCollective(
-                reinterpret_cast<const char*>(buffer.data()),
-                reinterpret_cast<char*>(tmp.data()), 4);
+            gcl.rearrangeForCollective(buffer, tmp);
             buffer = std::move(tmp);
           }
 
@@ -320,9 +323,9 @@ class ReplicatedResourceUpdateElementwiseClusteringHwTest
           auto host_rearrangement_it = host_rearrangements.find(index);
           if (host_rearrangement_it != host_rearrangements.end()) {
             auto& host_rearrangement = host_rearrangement_it->second;
-            per_replica_size = host_rearrangement.totalElementsPerReplica;
+            per_replica_size = host_rearrangement.getTotalElementsPerReplica();
             aligned_size =
-                per_replica_size * host_rearrangement.replicationFactor;
+                per_replica_size * host_rearrangement.getReplicationFactor();
           }
 
           VLOG(1) << "Downloading data from " << info.buffer_name
@@ -341,9 +344,7 @@ class ReplicatedResourceUpdateElementwiseClusteringHwTest
             EXPECT_TRUE(cluster);
             std::vector<float> tmp(buffer.size());
             VLOG(1) << "Undo rearrangement for collective...";
-            host_rearrangement.undoRearrangeForCollective(
-                reinterpret_cast<const char*>(buffer.data()),
-                reinterpret_cast<char*>(tmp.data()), 4);
+            host_rearrangement.undoRearrangeForCollective(buffer, tmp);
             buffer = std::move(tmp);
           }
           buffer.resize(size);
