Add popdist::run to prevent HST in TF2

Reviewers: christiana, #tensorflow, #framework_ip_review_-_any_oss_or_third-party_code_use_has_been_approved, georgep

Reviewed By: #tensorflow, #framework_ip_review_-_any_oss_or_third-party_code_use_has_been_approved, georgep

Maniphest Tasks: T63781

Differential Revision: https://phabricator.sourcevertex.net/D72533

Author: Frederik Mellbye

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

@@ -25,6 +25,7 @@ limitations under the License.
 #include <utility>
 #include <vector>
 
+#include <popdist/backend.hpp>
 #include <poplar/DeviceManager.hpp>
 #include <poplar/IPUModel.hpp>
 #include <poplar/StreamCallback.hpp>
@@ -3788,7 +3789,7 @@ Status PoplarExecutor::ExecuteEngineImpl(se::DeviceMemoryBase* result_buffer,
 
       // Run the main engine
       current_engine_->enableExecutionProfiling();
-      current_engine_->run(PoplarProgramType::MAIN_SEQUENCE);
+      popdist::run(*current_engine_, PoplarProgramType::MAIN_SEQUENCE);
 
       StopIOThreads();
 

tensorflow/compiler/plugin/poplar/tests/poprun_basic_test.py

@@ -28,6 +28,10 @@
 
 
 class PoprunBasicTest(test_util.TensorFlowTestCase):  # pylint: disable=abstract-method
+  @classmethod
+  def setUpClass(cls):
+    popdist.init()
+
   @tu.test_uses_ipus(num_ipus=4)
   @test_util.deprecated_graph_mode_only
   def test_cross_replica_sum(self):

tensorflow/python/ipu/tests/poprun_replica_partitioning_test.py

@@ -71,6 +71,10 @@ def cross_replica_reduce(self, grad):
 
 
 class PoprunReplicaPartitioningTest(test.TestCase):
+  @classmethod
+  def setUpClass(cls):
+    popdist.init()
+
   def _compare_partitioned_to_non_partitioned(self, stages, repeat_count,
                                               gradient_accumulation_count,
                                               dataset_fn, optimizer_fn):