Add BroadcastDistributed op (#2243)

* Add MPI_Bcast to nGraph core

* Add mpi_broadcast op and pass the unit test on CPU backend

* Modify API for MLSL

* Remove unused variable

* Fix the format

* Add the op call in python/pyngraph/ops/regmodule_pyngraph_op.cpp.

* Fix emitter function for DistBroadcast op

* Fix format

* Fix the emitter function for distbroadcast op

* Fix format

* Delete the reference code for distbroadcast op

* Include mlsl.h from ngraph build for genernal external usage

* Change back for build

* Fix alphabetical order issue and line space

* Rename disbroadcast to broadcastdistributed

* Change back continue

* Update the changes

* Fix style

* Delete unused variables

* Change some text format

* Add support for interpreter backend for broadcastdistributed

* Delete some redundant code

* Fix format

* rename

* remove unused file

* More renames

* renaming

* Remove one unused file

* Fix a bug

* add back redundant

* One more commit

* More commits

* Unify distributed_send and distributed_recv functions

* Fix format

CODEOWNERS

@@ -53,7 +53,7 @@
 /src/ngraph/type/                   @diyessi
 /src/ngraph/serializer.*pp          @rkimballn1
 
-/test/distributed.in.cpp            @wenzhe-nrv @diyessi @avijit-nervana
+/test/distributed.in.cpp            @wenzhe-nrv @diyessi @avijit-nervana @jianyinglang
 
 # Putting this last so it's not overridden by directory rules
 CMakeLists.txt                      @rkimballn1 @silee2

doc/sphinx/source/ops/broadcast_distributed.rst

+.. broadcastdistributed.rst:
+
+#####################
+BroadcastDistributed
+#####################
+
+.. code-block:: cpp
+
+   BroadcastDistributed // Collective operation
+
+
+Description
+===========
+
+Broadcast values from a primary root process or device to other processes or 
+devices within the op communicator.
+
+
+Inputs
+------
+
++-----------------+-------------------------+--------------------------------+
+| Name            | Element Type            | Shape                          |
++=================+=========================+================================+
+| ``arg``         | ``element::f32``        | Any                            |
+|                 | ``element::f64``        |                                |
++-----------------+-------------------------+--------------------------------+
+
+
+Outputs (in place)
+-------
+
++-----------------+-------------------------+--------------------------------+
+| Name            | Element Type            | Shape                          |
++=================+=========================+================================+
+| ``arg``         | ``element::f32``        | Same as ``arg``                |
+|                 | ``element::f64``        |                                |
++-----------------+-------------------------+--------------------------------+
+
+
+C++ Interface
+=============
+
+.. doxygenclass:: ngraph::op::BroadcastDistributed
+   :project: ngraph
+   :members:

doc/sphinx/source/ops/index.rst

@@ -26,6 +26,7 @@ Not currently a comprehensive list.
    * :doc:`batch_norm_training`
    * :doc:`batch_norm_training_backprop`
    * :doc:`broadcast`
+   * :doc:`broadcastdistributed`
    * :doc:`ceiling`
    * :doc:`concat`
    * :doc:`constant`
@@ -97,6 +98,7 @@ Not currently a comprehensive list.
    batch_norm_training.rst
    batch_norm_training_backprop.rst
    broadcast.rst
+   broadcast_distributed.rst
    ceiling.rst
    concat.rst
    constant.rst
@@ -194,4 +196,4 @@ that must be performed are:
 - Transformer handlers:
 
   * Interpreter (reference) implementation of behavior. The
-    implementation should favor clarity over efficiency.
\ No newline at end of file
+    implementation should favor clarity over efficiency.

python/pyngraph/ops/broadcast_distributed.cpp

+//*****************************************************************************
+// Copyright 2017-2018 Intel Corporation
+//
+// 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 <pybind11/pybind11.h>
+#include <pybind11/stl.h>
+
+#include "ngraph/op/broadcast_distributed.hpp"
+#include "pyngraph/ops/broadcast_distributed.hpp"
+
+namespace py = pybind11;
+
+void regclass_pyngraph_op_BroadcastDistributed(py::module m)
+{
+    py::class_<ngraph::op::BroadcastDistributed,
+               std::shared_ptr<ngraph::op::BroadcastDistributed>,
+               ngraph::op::Op>
+        broadcastdistributed(m, "BroadcastDistributed");
+    broadcastdistributed.doc() =
+        "ngraph.impl.op.BroadcastDistributed wraps ngraph::op::BroadcastDistributed";
+    broadcastdistributed.def(py::init<const std::shared_ptr<ngraph::Node>&>());
+}

python/pyngraph/ops/broadcast_distributed.hpp

+//*****************************************************************************
+// Copyright 2017-2019 Intel Corporation
+//
+// 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.
+//*****************************************************************************
+
+#pragma once
+
+#include <pybind11/pybind11.h>
+
+namespace py = pybind11;
+
+void regclass_pyngraph_op_BroadcastDistributed(py::module m);

python/pyngraph/ops/regmodule_pyngraph_op.cpp

@@ -37,6 +37,7 @@ void regmodule_pyngraph_op(py::module m_op)
     regclass_pyngraph_op_BatchNormTraining(m_op);
     regclass_pyngraph_op_BatchNormTrainingBackprop(m_op);
     regclass_pyngraph_op_Broadcast(m_op);
+    regclass_pyngraph_op_BroadcastDistributed(m_op);
     regclass_pyngraph_op_Ceiling(m_op);
     regclass_pyngraph_op_Concat(m_op);
     regclass_pyngraph_op_Constant(m_op);

python/pyngraph/ops/regmodule_pyngraph_op.hpp

@@ -29,6 +29,7 @@
 #include "pyngraph/ops/avg_pool.hpp"
 #include "pyngraph/ops/batch_norm.hpp"
 #include "pyngraph/ops/broadcast.hpp"
+#include "pyngraph/ops/broadcast_distributed.hpp"
 #include "pyngraph/ops/ceiling.hpp"
 #include "pyngraph/ops/concat.hpp"
 #include "pyngraph/ops/constant.hpp"

python/setup.py

@@ -172,6 +172,7 @@ sources = [
     'pyngraph/ops/atan.cpp',
     'pyngraph/ops/avg_pool.cpp',
     'pyngraph/ops/broadcast.cpp',
+    'pyngraph/ops/broadcast_distributed.cpp',
     'pyngraph/ops/concat.cpp',
     'pyngraph/ops/constant.cpp',
     'pyngraph/ops/convert.cpp',

src/ngraph/CMakeLists.txt

@@ -101,7 +101,7 @@ set (SRC
     op/batch_norm.cpp
     op/batch_norm.hpp
     op/broadcast.cpp
-    op/broadcast.hpp
+    op/broadcast_distributed.cpp
     op/ceiling.cpp
     op/ceiling.hpp
     op/concat.cpp

src/ngraph/ngraph.hpp

@@ -73,6 +73,7 @@
 #include "ngraph/op/avg_pool.hpp"
 #include "ngraph/op/batch_norm.hpp"
 #include "ngraph/op/broadcast.hpp"
+#include "ngraph/op/broadcast_distributed.hpp"
 #include "ngraph/op/ceiling.hpp"
 #include "ngraph/op/concat.hpp"
 #include "ngraph/op/constant.hpp"

src/ngraph/op/broadcast_distributed.cpp

+//*****************************************************************************
+// Copyright 2017-2018 Intel Corporation
+//
+// 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 "ngraph/op/broadcast_distributed.hpp"
+
+using namespace std;
+using namespace ngraph;
+
+op::BroadcastDistributed::BroadcastDistributed(const shared_ptr<Node>& arg)
+    : Op("BroadcastDistributed", check_single_output_args({arg}))
+{
+    constructor_validate_and_infer_types();
+}
+
+void op::BroadcastDistributed::validate_and_infer_types()
+{
+    NODE_VALIDATION_CHECK(this,
+                          get_input_element_type(0).is_dynamic() ||
+                              get_input_element_type(0) == element::f32 ||
+                              get_input_element_type(0) == element::f64,
+                          "Only element types f32 and f64 are supported (argument element type: ",
+                          get_input_element_type(0),
+                          ").");
+
+    set_output_type(0, get_input_element_type(0), get_input_partial_shape(0));
+}
+
+shared_ptr<Node> op::BroadcastDistributed::copy_with_new_args(const NodeVector& new_args) const
+{
+    check_new_args_count(this, new_args);
+    return make_shared<BroadcastDistributed>(new_args.at(0));
+}

src/ngraph/op/broadcast_distributed.hpp

+//*****************************************************************************
+// Copyright 2017-2018 Intel Corporation
+//
+// 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.
+//*****************************************************************************
+
+#pragma once
+
+#include <memory>
+
+#include "ngraph/op/op.hpp"
+
+namespace ngraph
+{
+    namespace op
+    {
+        class BroadcastDistributed : public Op
+        {
+        public:
+            BroadcastDistributed(const std::shared_ptr<Node>& arg);
+
+            void validate_and_infer_types() override;
+
+            virtual std::shared_ptr<Node>
+                copy_with_new_args(const NodeVector& new_args) const override;
+        };
+    }
+}

src/ngraph/op/op_tbl.hpp

@@ -67,6 +67,7 @@ NGRAPH_OP(BatchNormInference, ngraph::op)
 NGRAPH_OP(BatchNormTraining, ngraph::op)
 NGRAPH_OP(BatchNormTrainingBackprop, ngraph::op)
 NGRAPH_OP(Broadcast, ngraph::op)
+NGRAPH_OP(BroadcastDistributed, ngraph::op)
 NGRAPH_OP(BroadcastLike, ngraph::op)
 NGRAPH_OP(Ceiling, ngraph::op)
 NGRAPH_OP(Concat, ngraph::op)

src/ngraph/runtime/cpu/CMakeLists.txt

@@ -38,6 +38,7 @@ set(SRC
     builder/argmax.cpp
     builder/batch_norm.cpp
     builder/broadcast.cpp
+    builder/broadcast_distributed.cpp
     builder/bounded_relu.cpp
     builder/concat.cpp
     builder/convert.cpp

src/ngraph/runtime/cpu/builder/broadcast_distributed.cpp

+//*****************************************************************************
+// Copyright 2017-2019 Intel Corporation
+//
+// 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.
+//*****************************************************************************
+#ifdef NGRAPH_DISTRIBUTED_ENABLE
+
+#ifdef NGRAPH_DISTRIBUTED_MLSL_ENABLE
+#include <mlsl.hpp>
+#elif NGRAPH_DISTRIBUTED_OMPI_ENABLE
+#include <mpi.h>
+#endif
+
+#include "ngraph/op/broadcast_distributed.hpp"
+#include "ngraph/runtime/cpu/cpu_builder.hpp"
+
+using namespace std;
+using namespace ngraph;
+
+namespace ngraph
+{
+    namespace runtime
+    {
+        namespace cpu
+        {
+            template <>
+            void Builder::BUILDER_DECL(ngraph::op::BroadcastDistributed)
+            {
+                auto& functors = external_function->get_functors();
+
+                auto& arg_tensor = external_function->get_tensor_data(args[0].get_name());
+                auto count = static_cast<int>(args[0].get_size());
+
+#ifdef NGRAPH_DISTRIBUTED_MLSL_ENABLE
+                auto data_type = MLSL::DT_FLOAT;
+
+                if (args[0].get_element_type() == element::f32)
+                {
+                    data_type = MLSL::DT_FLOAT;
+                }
+                else if (args[0].get_element_type() == element::f64)
+                {
+                    data_type = MLSL::DT_DOUBLE;
+                }
+
+                auto functor = [&, count, data_type](CPURuntimeContext* ctx,
+                                                     CPUExecutionContext* ectx) {
+                    MLSL::CommReq* req =
+                        ctx->mlsl_dist->Bcast(arg_tensor, count, data_type, 0, MLSL::GT_DATA);
+                    ctx->mlsl_env->Wait(req);
+                };
+#elif NGRAPH_DISTRIBUTED_OMPI_ENABLE
+                auto data_type = MPI_FLOAT;
+
+                if (args[0].get_element_type() == element::f32)
+                {
+                    data_type = MPI_FLOAT;
+                }
+                else if (args[0].get_element_type() == element::f64)
+                {
+                    data_type = MPI_DOUBLE;
+                }
+
+                auto functor = [&, count, data_type](CPURuntimeContext* ctx,
+                                                     CPUExecutionContext* ectx) {
+                    MPI_Bcast(arg_tensor, count, data_type, 0, MPI_COMM_WORLD);
+                };
+#else
+                throw ngraph_error("Distributed Library not supported/mentioned");
+#endif
+                functors.emplace_back(functor);
+            }
+
+            REGISTER_OP_BUILDER(BroadcastDistributed);
+        }
+    }
+}
+#endif

src/ngraph/runtime/cpu/cpu_emitter.cpp

@@ -37,6 +37,7 @@
 #include "ngraph/op/avg_pool.hpp"
 #include "ngraph/op/batch_norm.hpp"
 #include "ngraph/op/broadcast.hpp"
+#include "ngraph/op/broadcast_distributed.hpp"
 #include "ngraph/op/ceiling.hpp"
 #include "ngraph/op/concat.hpp"
 #include "ngraph/op/constant.hpp"
@@ -244,6 +245,49 @@ namespace ngraph
                 writer.block_end();
 #else
                 throw ngraph_error("Distributed Library not supported/mentioned");
+#endif
+            }
+
+            template <>
+            void CPU_Emitter::EMITTER_DECL(ngraph::op::BroadcastDistributed)
+            {
+                const element::Type& element_type = args[0].get_element_type();
+#ifdef NGRAPH_DISTRIBUTED_MLSL_ENABLE
+                auto data_type = "MLSL::DT_FLOAT";
+
+                if (element_type == element::f32)
+                {
+                    data_type = "MLSL::DT_FLOAT";
+                }
+                else if (element_type == element::f64)
+                {
+                    data_type = "MLSL::DT_DOUBLE";
+                }
+
+                writer.block_begin();
+                writer << "MLSL::CommReq* req = ctx->mlsl_dist->Bcast(" << args[0].get_name()
+                       << ", " << args[0].get_size() << ", " << data_type
+                       << ", 0, MLSL::GT_DATA);\n";
+                writer << "ctx->mlsl_env->Wait(req);\n";
+                writer.block_end();
+#elif NGRAPH_DISTRIBUTED_OMPI_ENABLE
+                auto data_type = "MPI_FLOAT";
+
+                if (element_type == element::f32)
+                {
+                    data_type = "MPI_FLOAT";
+                }
+                else if (element_type == element::f64)
+                {
+                    data_type = "MPI_DOUBLE";
+                }
+
+                writer.block_begin();
+                writer << "MPI_Bcast(" << args[0].get_name() << ", " << args[0].get_size() << ", "
+                       << data_type << ", 0, MPI_COMM_WORLD);\n";
+                writer.block_end();
+#else
+                throw ngraph_error("Distributed Library not supported/mentioned");
 #endif
             }
 #endif

src/ngraph/runtime/cpu/cpu_external_function.cpp

@@ -180,6 +180,7 @@
 
 #ifdef NGRAPH_DISTRIBUTED_ENABLE
 #include "ngraph/op/allreduce.hpp"
+#include "ngraph/op/broadcast_distributed.hpp"
 #endif
 
 using namespace std;
@@ -290,6 +291,8 @@ static const runtime::cpu::OpMap dispatcher{
     {TI(ngraph::op::Add), &runtime::cpu::CPU_Emitter::emit<op::Add>},
 #ifdef NGRAPH_DISTRIBUTED_ENABLE
     {TI(ngraph::op::AllReduce), &runtime::cpu::CPU_Emitter::emit<op::AllReduce>},
+    {TI(ngraph::op::BroadcastDistributed),
+     &runtime::cpu::CPU_Emitter::emit<op::BroadcastDistributed>},
 #endif
     {TI(ngraph::op::MatmulBias), &runtime::cpu::CPU_Emitter::emit<op::MatmulBias>},
     {TI(ngraph::op::Dot), &runtime::cpu::CPU_Emitter::emit<op::Dot>},

src/ngraph/runtime/generic_cpu/gcpu_executable.hpp

@@ -135,6 +135,7 @@
 
 #ifdef NGRAPH_DISTRIBUTED_ENABLE
 #include "ngraph/runtime/reference/allreduce.hpp"
+#include "ngraph/runtime/reference/broadcast_distributed.hpp"
 #endif
 
 namespace ngraph
@@ -437,6 +438,32 @@ private:
                                        broadcast_axes);
             break;
         }
+        case OP_TYPEID::BroadcastDistributed: {
+#ifdef NGRAPH_DISTRIBUTED_ENABLE
+            Distributed dist;
+            int Rank_ID;
+            Rank_ID = dist.get_rank();
+            if (Rank_ID == 0)
+            {
+                reference::broadcastdistributed<T>(
+                    static_cast<T*>(args[0]),
+                    node.get_input_element_type(0),
+                    static_cast<int>(shape_size(node.get_input_shape(0))));
+                auto memSize = static_cast<int>(shape_size(node.get_input_shape(0))) *
+                               sizeof(node.get_input_element_type(0));
+                memcpy(out[0], args[0], memSize);
+            }
+            else
+            {
+                reference::broadcastdistributed<T>(
+                    static_cast<T*>(out[0]),
+                    node.get_input_element_type(0),
+                    static_cast<int>(shape_size(node.get_input_shape(0))));
+            }
+            break;
+#endif
+            break;
+        }
         case OP_TYPEID::BroadcastLike: break;
         case OP_TYPEID::Ceiling:
         {

src/ngraph/runtime/gpu/gpu_emitter.cpp

@@ -44,6 +44,7 @@
 #include "ngraph/op/avg_pool.hpp"
 #include "ngraph/op/batch_norm.hpp"
 #include "ngraph/op/broadcast.hpp"
+#include "ngraph/op/broadcast_distributed.hpp"
 #include "ngraph/op/ceiling.hpp"
 #include "ngraph/op/concat.hpp"
 #include "ngraph/op/constant.hpp"
@@ -548,6 +549,11 @@ std::string runtime::gpu::GPU_Emitter::emit_Cosh(EMIT_ARGS)
     return emit_elementwise<ngraph::op::Cosh>(compiled_function, function_name, node, args, out);
 }
 
+std::string runtime::gpu::GPU_Emitter::emit_BroadcastDistributed(EMIT_ARGS)
+{
+    throw unsupported_op("Unsupported op '" + node->description() + "'");
+}
+
 std::string runtime::gpu::GPU_Emitter::emit_Divide(EMIT_ARGS)
 {
     return emit_elementwise<ngraph::op::Divide>(compiled_function, function_name, node, args, out);

src/ngraph/runtime/intelgpu/intelgpu_backend.cpp

@@ -1995,6 +1995,7 @@ shared_ptr<runtime::Executable>
             break;
         }
         case OP_TYPEID::AllReduce:
+        case OP_TYPEID::BroadcastDistributed:
         case OP_TYPEID::BroadcastLike:
         case OP_TYPEID::DynReshape:
         case OP_TYPEID::DynSlice:

src/ngraph/runtime/interpreter/int_backend.hpp

@@ -24,10 +24,6 @@
 
 #include "ngraph/runtime/tensor.hpp"
 
-#ifdef NGRAPH_DISTRIBUTED_ENABLE
-#include "ngraph/runtime/reference/allreduce.hpp"
-#endif
-
 namespace ngraph
 {
     namespace runtime

src/ngraph/runtime/interpreter/int_executable.hpp

@@ -140,6 +140,7 @@
 
 #ifdef NGRAPH_DISTRIBUTED_ENABLE
 #include "ngraph/runtime/reference/allreduce.hpp"
+#include "ngraph/runtime/reference/broadcast_distributed.hpp"
 #endif
 
 namespace ngraph
@@ -440,6 +441,32 @@ private:
                                     broadcast_axes);
             break;
         }
+        case OP_TYPEID::BroadcastDistributed: {
+#ifdef NGRAPH_DISTRIBUTED_ENABLE
+            Distributed dist;
+            int Rank_ID;
+            Rank_ID = dist.get_rank();
+            if (Rank_ID == 0)
+            {
+                reference::broadcastdistributed<T>(
+                    args[0]->get_data_ptr<T>(),
+                    node.get_input_element_type(0),
+                    static_cast<int>(shape_size(node.get_input_shape(0))));
+                auto memSize = static_cast<int>(shape_size(node.get_input_shape(0))) *
+                               sizeof(node.get_input_element_type(0));
+                memcpy(out[0]->get_data_ptr<T>(), args[0]->get_data_ptr<T>(), memSize);
+            }
+            else
+            {
+                reference::broadcastdistributed<T>(
+                    out[0]->get_data_ptr<T>(),
+                    node.get_input_element_type(0),
+                    static_cast<int>(shape_size(node.get_input_shape(0))));
+            }
+            break;
+#endif
+            break;
+        }
         case OP_TYPEID::BroadcastLike: break;
         case OP_TYPEID::Ceiling:
         {

src/ngraph/runtime/reference/broadcast_distributed.hpp

+//*****************************************************************************
+// Copyright 2017-2019 Intel Corporation
+//
+// 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.
+//*****************************************************************************
+
+#pragma once
+
+#ifdef NGRAPH_DISTRIBUTED_ENABLE
+#ifdef NGRAPH_DISTRIBUTED_MLSL_ENABLE
+#include <mlsl.hpp>
+#elif NGRAPH_DISTRIBUTED_OMPI_ENABLE
+#include <mpi.h>
+#endif
+#include "ngraph/type/element_type.hpp"
+
+namespace ngraph
+{
+    namespace runtime
+    {
+        namespace reference
+        {
+            template <typename T>
+            void broadcastdistributed(T* arg, const element::Type element_type, int count)
+            {
+#ifdef NGRAPH_DISTRIBUTED_MLSL_ENABLE
+                auto data_type = MLSL::DT_FLOAT;
+
+                if (element_type == element::f64)
+                {
+                    data_type = MLSL::DT_DOUBLE;
+                }
+                else if (element_type != element::f32)
+                {
+                    throw std::runtime_error(
+                        "BroadcastDistributed op supports only f32 and f64 types");
+                }
+
+                MLSL::Environment& env = MLSL::Environment::GetEnv();
+                MLSL::Distribution* distribution = env.CreateDistribution(env.GetProcessCount(), 1);
+                MLSL::CommReq* req = distribution->Bcast(arg, count, data_type, 0, MLSL::GT_DATA);
+                env.Wait(req);
+                env.DeleteDistribution(distribution);
+#elif NGRAPH_DISTRIBUTED_OMPI_ENABLE
+                auto data_type = MPI_FLOAT;
+
+                if (element_type == element::f64)
+                {
+                    data_type = MPI_DOUBLE;
+                }
+                else if (element_type != element::f32)
+                {
+                    throw std::runtime_error(
+                        "BroadcastDistributed op supports only f32 and f64 types");
+                }
+                MPI_Bcast(arg, count, data_type, 0, MPI_COMM_WORLD);
+#else
+                throw ngraph_error("Distributed Library not supported/mentioned");
+#endif
+            }
+        }
+    }
+}
+
+#endif

src/ngraph/serializer.cpp

@@ -34,6 +34,7 @@
 #include "ngraph/op/avg_pool.hpp"
 #include "ngraph/op/batch_norm.hpp"
 #include "ngraph/op/broadcast.hpp"
+#include "ngraph/op/broadcast_distributed.hpp"
 #include "ngraph/op/ceiling.hpp"
 #include "ngraph/op/concat.hpp"
 #include "ngraph/op/constant.hpp"
@@ -589,6 +590,11 @@ static shared_ptr<ngraph::Function>
                 node = make_shared<op::Broadcast>(args[0], shape, axes);
                 break;
             }
+            case OP_TYPEID::BroadcastDistributed:
+            {
+                node = make_shared<op::BroadcastDistributed>(args[0]);
+                break;
+            }
             case OP_TYPEID::BroadcastLike:
             {
                 auto initial_axes = node_js.at("initial_axes").get<set<size_t>>();
@@ -1448,6 +1454,8 @@ static json write(const Node& n, bool binary_constant_data)
         node["shape"] = tmp->get_broadcast_shape();
         break;
     }
+    case OP_TYPEID::BroadcastDistributed: { break;
+    }
     case OP_TYPEID::BroadcastLike:
     {
         auto tmp = dynamic_cast<const op::BroadcastLike*>(&n);

test/distributed.in.cpp

@@ -56,3 +56,27 @@ TEST(distributed_${BACKEND_NAME}, allreduce)
         EXPECT_TRUE(test::all_close_f(v, read_vector<float>(result)));
     }
 }
+
+TEST(distributed_${BACKEND_NAME}, broadcastdistributed)
+{
+    auto shape = Shape{2, 2};
+    auto A = make_shared<op::Parameter>(element::f32, shape);
+    auto f = make_shared<Function>(make_shared<op::BroadcastDistributed>(A), ParameterVector{A});
+
+    auto backend = runtime::Backend::create("${BACKEND_NAME}");
+
+    auto v = vector<float>{1, 2, 3, 4};
+    auto result = backend->create_tensor(element::f32, shape);
+    copy_data(result, vector<float>(4, 0));
+
+    DistributedSetup distsetup;
+    auto processIdx = distsetup.get_comm_rank();
+    if (processIdx == 0)
+    {
+        copy_data(result, v);
+    }
+
+    auto handle = backend->compile(f);
+    handle->call_with_validate({result}, {result});
+    EXPECT_EQ(v, read_vector<float>(result));
+}