From 48dea84bf03971fafeb59eccf08d3237dc209690 Mon Sep 17 00:00:00 2001
From: dongzhihong <dzhwinter@gmail.com>
Date: Wed, 6 Sep 2017 21:12:27 -0700
Subject: [PATCH] "nccl multigpu init"

---
 paddle/operators/nccl/nccl_gpu_common.h | 39 ++++++++++++++++++++
 paddle/operators/nccl/nccl_ops.cc       | 48 +++++++++++++++++++++++++
 paddle/operators/nccl/nccl_ops.h        |  7 ++++
 3 files changed, 94 insertions(+)
 create mode 100644 paddle/operators/nccl/nccl_gpu_common.h
 create mode 100644 paddle/operators/nccl/nccl_ops.cc
 create mode 100644 paddle/operators/nccl/nccl_ops.h

diff --git a/paddle/operators/nccl/nccl_gpu_common.h b/paddle/operators/nccl/nccl_gpu_common.h
new file mode 100644
index 00000000000..017492a0d8f
--- /dev/null
+++ b/paddle/operators/nccl/nccl_gpu_common.h
@@ -0,0 +1,39 @@
+#pragma once
+#include <nccl.h>
+
+#include "paddle/platform/device_context.h"
+
+namespace paddle {
+namespace platform {
+
+class NCCLManager {
+ public:
+  static NCCLManager* Get() {
+    static NCCLManager m;
+    return &m;
+  }
+
+  NCCLManager() { _comms.resize(_gpu_worlds.size()); }
+  ~NCCLManager() {}
+
+ private:
+  // clang-format off
+  std::vector<ncclComm_t> _comms;
+  std::vector<int> _gpu_worlds;
+  // clang-format on
+};
+
+class NCCLContext : public DeviceContext {
+ public:
+  explicit NCCLContext(GPUPlace place);
+  virtual ~NCCLContext();
+
+ private:
+  // clang-format off
+  std::vector<int> _gpu_ids;
+  std::vector<cudaStream_t> _streams;
+  int root_gpu;
+  // clang-format on
+};
+}
+}
diff --git a/paddle/operators/nccl/nccl_ops.cc b/paddle/operators/nccl/nccl_ops.cc
new file mode 100644
index 00000000000..a4bd8b9c0f6
--- /dev/null
+++ b/paddle/operators/nccl/nccl_ops.cc
@@ -0,0 +1,48 @@
+#include "paddle/framework/op_registry.h"
+#include "paddle/operators/nccl/nccl_gpu_common.h"
+
+namespace paddle {
+namespace operators {
+
+// AllreduceOp
+class NCCLAllreduceOp : public framework::OperatorWithKernel {
+ public:
+  using framework::OperatorWithKernel::OperatorWithKernel;
+
+ protected:
+  // allreduce do nothing in infershape
+  void InferShape(const framework::InferShapeContext &ctx) const override {}
+};
+
+template <typename T>
+class NCCLAllreduceOp : public framework::OpKernel {
+ public:
+  void Compute(const framework::ExecutionContext &context) const override {
+    auto *ctx = static_cast<NCCLContext *>(context.device_context());
+    // auto *comm = ;
+    // auto *src = ;
+    // ncclAllReduce(src, dest, )
+  }
+};
+
+// BcastSendOp
+template <typename T>
+class NCCLBroadcastSendOp final : public framework::OperatorWithKernel {
+ public:
+  using framework::OperatorWithKernel::OperatorWithKernel;
+
+ protected:
+  void InferShape(const framework::InferShapeContext &ctx) const override {}
+};
+
+// BcastRecvOp
+template <typename T>
+class NCCLBroadcastRecvOp final : public framework::OperatorWithKernel {
+ public:
+  using framework::OperatorWithKernel::OperatorWithKernel;
+
+ protected:
+  void InferShape(const framework::InferShapeContext &ctx) const override {}
+};
+}
+}
diff --git a/paddle/operators/nccl/nccl_ops.h b/paddle/operators/nccl/nccl_ops.h
new file mode 100644
index 00000000000..0d78c606395
--- /dev/null
+++ b/paddle/operators/nccl/nccl_ops.h
@@ -0,0 +1,7 @@
+#pragma once
+#include "paddle/framework/op_registry.h"
+#include "paddle/operators/nccl/nccl_gpu_common.h"
+
+namespace paddle {
+namespace operators {}
+}
-- 
GitLab