Implement graph deep copy function, switch to pointers to weights/bias (#269)

Author: NeiroYT

include/graph/graph.hpp

@@ -23,6 +23,9 @@ struct BranchState {
   std::vector<std::pair<int, int>> distribution;
 };
 
+std::shared_ptr<Layer> layer_based_shared_copy(
+    const std::shared_ptr<Layer>& layer, const RuntimeOptions& options);
+
 class Graph {
   int BiggestSize_;
   int V_;  // amount of ids
@@ -72,6 +75,9 @@ class Graph {
   Graph& operator=(Graph&&) noexcept = default;
   ~Graph() = default;
 
+  void clone(Graph& result, Tensor& out,
+             const RuntimeOptions& options = RuntimeOptions()) const;
+
   void setSplitDistribution(
       std::vector<std::vector<std::pair<int, int>>> split_dist) {
     split_distribution_ = std::move(split_dist);

include/layers/ConvLayer.hpp

@@ -22,35 +22,48 @@ class ConvolutionalLayer : public Layer {
   size_t stride_;
   size_t pads_;
   size_t dilations_;
-  Tensor kernel_;
-  Tensor bias_;
+  std::shared_ptr<Tensor> kernel_;
+  std::shared_ptr<Tensor> bias_;
   size_t group_;
   bool useLegacyImpl_;
 
  public:
-  ConvolutionalLayer() : Layer(kConvolution) {
+  ConvolutionalLayer() : Layer(kConvolution), kernel_(nullptr), bias_(nullptr) {
     stride_ = 0;
     pads_ = 0;
     dilations_ = 0;
   }
   ConvolutionalLayer(size_t step, size_t pads, size_t dilations,
                      const Tensor& kernel, const Tensor& bias = Tensor(),
                      size_t group = 1, bool useLegacyImpl = false)
-      : Layer(kConvolution) {
+      : Layer(kConvolution),
+        kernel_(std::make_shared<Tensor>(kernel)),
+        bias_(std::make_shared<Tensor>(bias)) {
+    stride_ = step;
+    pads_ = pads;
+    group_ = group;
+    dilations_ = dilations;
+    useLegacyImpl_ = useLegacyImpl;
+  }
+  ConvolutionalLayer(size_t step, size_t pads, size_t dilations,
+                     std::shared_ptr<Tensor> kernel,
+                     std::shared_ptr<Tensor> bias = std::make_shared<Tensor>(),
+                     size_t group = 1, bool useLegacyImpl = false)
+      : Layer(kConvolution),
+        kernel_(std::move(kernel)),
+        bias_(std::move(bias)) {
     stride_ = step;
     pads_ = pads;
     group_ = group;
     dilations_ = dilations;
-    kernel_ = kernel;
-    bias_ = bias;
     useLegacyImpl_ = useLegacyImpl;
   }
   void run(const std::vector<Tensor>& input,
            std::vector<Tensor>& output) override;
   void run(const std::vector<Tensor>& input, std::vector<Tensor>& output,
            const RuntimeOptions& options) override;
 #ifdef ENABLE_STATISTIC_WEIGHTS
-  Tensor get_weights() override { return kernel_; }
+  Tensor get_weights() override { return *kernel_; }
 #endif
 };
 

include/layers/FCLayer.hpp

@@ -11,17 +11,23 @@ namespace it_lab_ai {
 
 class FCLayer : public Layer {
  private:
-  Tensor weights_;
-  Tensor bias_;
+  std::shared_ptr<Tensor> weights_;
+  std::shared_ptr<Tensor> bias_;
 
  public:
-  FCLayer() : Layer(kFullyConnected) {}
-  FCLayer(Tensor weights, const Tensor& bias)
-      : Layer(kFullyConnected), weights_(std::move(weights)), bias_(bias) {}
+  FCLayer() : Layer(kFullyConnected), weights_(nullptr), bias_(nullptr) {}
+  FCLayer(const Tensor& weights, const Tensor& bias)
+      : Layer(kFullyConnected),
+        weights_(std::make_shared<Tensor>(weights)),
+        bias_(std::make_shared<Tensor>(bias)) {}
+  FCLayer(std::shared_ptr<Tensor> weights, std::shared_ptr<Tensor> bias)
+      : Layer(kFullyConnected),
+        weights_(std::move(weights)),
+        bias_(std::move(bias)) {}
   void run(const std::vector<Tensor>& input,
            std::vector<Tensor>& output) override;
 #ifdef ENABLE_STATISTIC_WEIGHTS
-  Tensor get_weights() override { return weights_; }
+  Tensor get_weights() override { return *weights_; }
 #endif
 };
 

include/layers_oneDNN/BinaryOpLayer.hpp

@@ -17,6 +17,10 @@ class BinaryOpLayerOneDnn : public Layer {
   explicit BinaryOpLayerOneDnn(BinaryOpLayer::Operation op)
       : Layer(kBinaryOp), op_(op) {}
 
+  BinaryOpLayerOneDnn(const BinaryOpLayerOneDnn& c) : Layer(kBinaryOp) {
+    this->op_ = c.op_;
+  }
+
   void run(const std::vector<Tensor>& input,
            std::vector<Tensor>& output) override;
 

include/layers_oneDNN/ConvLayer.hpp

@@ -16,6 +16,8 @@ class ConvLayerOneDnn : public Layer {
         stride_(1),
         pads_(0),
         dilations_(1),
+        kernel_(nullptr),
+        bias_(nullptr),
         group_(1),
         use_legacy_(false) {}
 
@@ -26,16 +28,39 @@ class ConvLayerOneDnn : public Layer {
         stride_(stride),
         pads_(pads),
         dilations_(dilations),
-        kernel_(kernel),
-        bias_(bias),
+        kernel_(std::make_shared<Tensor>(kernel)),
+        bias_(std::make_shared<Tensor>(bias)),
         group_(group),
         use_legacy_(use_legacy) {}
 
+  ConvLayerOneDnn(size_t stride, size_t pads, size_t dilations,
+                  std::shared_ptr<Tensor> kernel,
+                  std::shared_ptr<Tensor> bias = std::make_shared<Tensor>(),
+                  size_t group = 1, bool use_legacy = false)
+      : Layer(kConvolution),
+        stride_(stride),
+        pads_(pads),
+        dilations_(dilations),
+        kernel_(std::move(kernel)),
+        bias_(std::move(bias)),
+        group_(group),
+        use_legacy_(use_legacy) {}
+
+  ConvLayerOneDnn(const ConvLayerOneDnn& c) : Layer(kConvolution) {
+    this->stride_ = c.stride_;
+    this->pads_ = c.pads_;
+    this->dilations_ = c.dilations_;
+    this->kernel_ = c.kernel_;
+    this->bias_ = c.bias_;
+    this->group_ = c.group_;
+    this->use_legacy_ = c.use_legacy_;
+  }
+
   void run(const std::vector<Tensor>& input,
            std::vector<Tensor>& output) override;
 
 #ifdef ENABLE_STATISTIC_WEIGHTS
-  Tensor get_weights() override { return kernel_; }
+  Tensor get_weights() override { return *kernel_; }
 #endif
 
  private:
@@ -85,8 +110,8 @@ class ConvLayerOneDnn : public Layer {
   size_t stride_;
   size_t pads_;
   size_t dilations_;
-  Tensor kernel_;
-  Tensor bias_;
+  std::shared_ptr<Tensor> kernel_;
+  std::shared_ptr<Tensor> bias_;
   size_t group_;
   bool use_legacy_;
 

include/layers_oneDNN/EWLayer.hpp

@@ -21,6 +21,12 @@ class EwLayerOneDnn : public Layer {
         alpha_(alpha),
         beta_(beta) {}
 
+  EwLayerOneDnn(const EwLayerOneDnn& c) : Layer(kElementWise) {
+    this->func_ = c.func_;
+    this->alpha_ = c.alpha_;
+    this->beta_ = c.beta_;
+  }
+
   void run(const std::vector<Tensor>& input,
            std::vector<Tensor>& output) override;
   static bool is_function_supported(const std::string& function);

include/layers_oneDNN/PoolingLayer.hpp

@@ -27,6 +27,18 @@ class PoolingLayerOneDnn : public Layer {
         engine_(std::make_unique<dnnl::engine>(dnnl::engine::kind::cpu, 0)),
         stream_(std::make_unique<dnnl::stream>(*engine_)) {}
 
+  PoolingLayerOneDnn(const PoolingLayerOneDnn& c)
+      : Layer(kPooling),
+        engine_(std::make_unique<dnnl::engine>(dnnl::engine::kind::cpu, 0)),
+        stream_(std::make_unique<dnnl::stream>(*engine_)) {
+    this->poolingShape_ = c.poolingShape_;
+    this->strides_ = c.strides_;
+    this->pads_ = c.pads_;
+    this->dilations_ = c.dilations_;
+    this->ceil_mode_ = c.ceil_mode_;
+    this->poolingType_ = c.poolingType_;
+  }
+
   void run(const std::vector<Tensor>& input,
            std::vector<Tensor>& output) override;
 

include/layers_oneDNN/ReduceLayer.hpp

@@ -20,6 +20,12 @@ class ReduceLayerOneDnn : public Layer {
                              const std::vector<int64_t>& axes = {})
       : ReduceLayerOneDnn(ReduceLayer::Operation::kSum, keepdims, axes) {}
 
+  ReduceLayerOneDnn(const ReduceLayerOneDnn& c) : Layer(kReduce) {
+    this->op_ = c.op_;
+    this->keepdims_ = c.keepdims_;
+    this->axes_ = c.axes_;
+  }
+
   void run(const std::vector<Tensor>& input,
            std::vector<Tensor>& output) override;
 

src/graph/CMakeLists.txt

@@ -1,4 +1,4 @@
 file(GLOB_RECURSE graph_src *.cpp)
 add_library(graph_lib STATIC "${GRAPH_HEADERS}" "${graph_src}")
-target_link_libraries(graph_lib PUBLIC TBB_unified)
+target_link_libraries(graph_lib PUBLIC dnnl TBB_unified)
 add_dependencies(graph_lib kokkos_external)