Concat&&OneDNN

app/Graph/build.cpp

@@ -519,7 +519,7 @@ ParseResult parse_json_model(RuntimeOptions options,
             concat_connections[layer_name].push_back(base_input_name);
           }
         }
-        auto concat_layer = std::make_shared<it_lab_ai::ConcatLayer>(axis);
+        auto concat_layer = LayerFactory::createConcatLayer(axis, options);
         layer = concat_layer;
         concat_connected_inputs[layer_name] = std::unordered_set<std::string>();
       } else if (layer_type == "Split") {

app/Graph/build.hpp

@@ -33,6 +33,7 @@
 #include "layers/Tensor.hpp"
 #include "layers/TransposeLayer.hpp"
 #include "layers_oneDNN/BinaryOpLayer.hpp"
+#include "layers_oneDNN/ConcatLayer.hpp"
 #include "layers_oneDNN/ConvLayer.hpp"
 #include "layers_oneDNN/EWLayer.hpp"
 #include "layers_oneDNN/PoolingLayer.hpp"
@@ -134,6 +135,14 @@ class LayerFactory {
     return std::make_shared<PoolingLayer>(shape, strides, pads, dilations,
                                           ceil_mode, PoolType);
   }
+
+  static std::shared_ptr<Layer> createConcatLayer(
+      int64_t axis, const RuntimeOptions& options) {
+    if (options.backend == Backend::kOneDnn) {
+      return std::make_shared<ConcatLayerOneDnn>(axis);
+    }
+    return std::make_shared<ConcatLayer>(axis);
+  }
 };
 
 }  // namespace it_lab_ai

include/layers_oneDNN/ConcatLayer.hpp

@@ -0,0 +1,69 @@
+#pragma once
+
+#include <dnnl.hpp>
+#include <memory>
+#include <vector>
+
+#include "layers/Layer.hpp"
+#include "layers/Tensor.hpp"
+
+namespace it_lab_ai {
+
+class ConcatLayerOneDnn : public Layer {
+ public:
+  explicit ConcatLayerOneDnn(int64_t axis = 0) : Layer(kConcat), axis_(axis) {}
+
+  ConcatLayerOneDnn(const ConcatLayerOneDnn& c)
+      : Layer(kConcat), axis_(c.axis_) {}
+
+  void run(const std::vector<Tensor>& input,
+           std::vector<Tensor>& output) override;
+
+#ifdef ENABLE_STATISTIC_WEIGHTS
+  Tensor get_weights() override {
+    std::vector<int> v = {0};
+    return make_tensor(v);
+  }
+#endif
+
+ private:
+  int64_t axis_;
+
+  bool initialized_ = false;
+  Type last_type_;
+  std::vector<Shape> last_shapes_;
+
+  std::unique_ptr<dnnl::engine> engine_;
+  std::unique_ptr<dnnl::stream> stream_;
+  std::unique_ptr<dnnl::concat> concat_prim_;
+
+  std::vector<dnnl::memory::desc> src_mds_;
+  dnnl::memory::desc dst_md_;
+
+  Shape output_shape_;
+
+  std::vector<dnnl::memory> src_mems_;
+  dnnl::memory dst_mem_;
+  std::unordered_map<int, dnnl::memory> args_;
+
+  std::vector<float> dst_buffer_f32_;
+  std::vector<int> dst_buffer_s32_;
+
+  void initialize_onednn(const std::vector<Tensor>& input);
+
+  static void validate_input(const std::vector<Tensor>& input);
+
+  [[nodiscard]] static dnnl::memory::data_type get_dnnl_data_type(Type type);
+
+  [[nodiscard]] static dnnl::memory::format_tag pick_format(size_t ndims);
+
+  [[nodiscard]] static std::vector<dnnl::memory::dim> shape_to_dims(
+      const Shape& shape);
+
+  [[nodiscard]] static Shape calculate_output_shape(
+      const std::vector<Tensor>& inputs, int64_t axis);
+
+  [[nodiscard]] static int64_t normalize_axis(int64_t axis, size_t rank);
+};
+
+}  // namespace it_lab_ai

src/layers_oneDNN/ConcatLayer.cpp

@@ -0,0 +1,222 @@
+#include "layers_oneDNN/ConcatLayer.hpp"
+
+#include <stdexcept>
+
+namespace it_lab_ai {
+
+void ConcatLayerOneDnn::run(const std::vector<Tensor>& input,
+                            std::vector<Tensor>& output) {
+  validate_input(input);
+
+  if (input.size() == 1) {
+    output = input;
+    return;
+  }
+
+  Type type = input[0].get_type();
+
+  bool need_reinit = !initialized_ || last_type_ != type ||
+                     last_shapes_.size() != input.size();
+
+  if (!need_reinit) {
+    for (size_t i = 0; i < input.size(); ++i) {
+      if (last_shapes_[i] != input[i].get_shape()) {
+        need_reinit = true;
+        break;
+      }
+    }
+  }
+
+  if (need_reinit) {
+    initialize_onednn(input);
+  }
+
+  output.resize(1);
+
+  if (type == Type::kFloat) {
+    for (size_t i = 0; i < input.size(); ++i) {
+      if (last_type_ == Type::kFloat) {
+        src_mems_[i].set_data_handle(
+            const_cast<float*>(input[i].as<float>()->data()));
+      } else {
+        src_mems_[i].set_data_handle(
+            const_cast<int*>(input[i].as<int>()->data()));
+      }
+
+      args_[DNNL_ARG_MULTIPLE_SRC + i] = src_mems_[i];
+    }
+
+    args_[DNNL_ARG_DST] = dst_mem_;
+
+    concat_prim_->execute(*stream_, args_);
+    stream_->wait();
+
+    output[0] = make_tensor(dst_buffer_f32_, output_shape_);
+  } else if (type == Type::kInt) {
+    for (size_t i = 0; i < input.size(); ++i) {
+      src_mems_[i].set_data_handle(
+          const_cast<int*>(input[i].as<int>()->data()));
+      args_[DNNL_ARG_MULTIPLE_SRC + i] = src_mems_[i];
+    }
+
+    args_[DNNL_ARG_DST] = dst_mem_;
+
+    concat_prim_->execute(*stream_, args_);
+    stream_->wait();
+
+    output[0] = make_tensor(dst_buffer_s32_, output_shape_);
+  }
+}
+
+void ConcatLayerOneDnn::validate_input(const std::vector<Tensor>& input) {
+  Type type = input[0].get_type();
+  const Shape& base = input[0].get_shape();
+
+  for (size_t i = 1; i < input.size(); ++i) {
+    if (input[i].get_type() != type) {
+      throw std::runtime_error(
+          "ConcatLayerOneDnn: All tensors must have same type");
+    }
+
+    if (input[i].get_shape().dims() != base.dims()) {
+      throw std::runtime_error(
+          "ConcatLayerOneDnn: All tensors must have same rank");
+    }
+  }
+}
+
+void ConcatLayerOneDnn::initialize_onednn(const std::vector<Tensor>& input) {
+  if (!engine_) {
+    engine_ = std::make_unique<dnnl::engine>(dnnl::engine::kind::cpu, 0);
+  }
+  if (!stream_) {
+    stream_ = std::make_unique<dnnl::stream>(*engine_);
+  }
+
+  size_t rank = input[0].get_shape().dims();
+  int64_t axis = normalize_axis(axis_, rank);
+
+  last_type_ = input[0].get_type();
+  auto type = get_dnnl_data_type(last_type_);
+
+  auto layout = pick_format(rank);
+
+  src_mds_.clear();
+  for (const auto& t : input) {
+    src_mds_.emplace_back(shape_to_dims(t.get_shape()), type, layout);
+  }
+
+  output_shape_ = calculate_output_shape(input, axis);
+
+  dst_md_ = dnnl::memory::desc(shape_to_dims(output_shape_), type, layout);
+
+  auto concat_pd =
+      dnnl::concat::primitive_desc(*engine_, dst_md_, axis, src_mds_);
+  concat_prim_ = std::make_unique<dnnl::concat>(concat_pd);
+
+  dst_md_ = concat_pd.dst_desc();
+  src_mds_.clear();
+  for (size_t i = 0; i < input.size(); ++i) {
+    src_mds_.push_back(concat_pd.src_desc(i));
+  }
+
+  size_t n = input.size();
+  src_mems_.resize(n);
+  for (size_t i = 0; i < n; ++i) {
+    src_mems_[i] = dnnl::memory(src_mds_[i], *engine_, nullptr);
+  }
+
+  size_t out_size = output_shape_.count();
+  if (last_type_ == Type::kFloat) {
+    dst_buffer_f32_.resize(out_size);
+    dst_mem_ = dnnl::memory(dst_md_, *engine_, dst_buffer_f32_.data());
+  } else {
+    dst_buffer_s32_.resize(out_size);
+    dst_mem_ = dnnl::memory(dst_md_, *engine_, dst_buffer_s32_.data());
+  }
+
+  args_.clear();
+  for (size_t i = 0; i < n; ++i) {
+    args_[DNNL_ARG_MULTIPLE_SRC + i] = src_mems_[i];
+  }
+  args_[DNNL_ARG_DST] = dst_mem_;
+
+  last_shapes_.clear();
+  for (const auto& t : input) {
+    last_shapes_.push_back(t.get_shape());
+  }
+
+  initialized_ = true;
+}
+
+dnnl::memory::data_type ConcatLayerOneDnn::get_dnnl_data_type(Type type) {
+  switch (type) {
+    case Type::kFloat:
+      return dnnl::memory::data_type::f32;
+    case Type::kInt:
+      return dnnl::memory::data_type::s32;
+    default:
+      throw std::runtime_error("Unsupported data type for oneDNN");
+  }
+}
+
+dnnl::memory::format_tag ConcatLayerOneDnn::pick_format(size_t ndims) {
+  switch (ndims) {
+    case 1:
+      return dnnl::memory::format_tag::a;
+    case 2:
+      return dnnl::memory::format_tag::ab;
+    case 3:
+      return dnnl::memory::format_tag::abc;
+    case 4:
+      return dnnl::memory::format_tag::abcd;
+    case 5:
+      return dnnl::memory::format_tag::abcde;
+    default:
+      return dnnl::memory::format_tag::any;
+  }
+}
+
+std::vector<dnnl::memory::dim> ConcatLayerOneDnn::shape_to_dims(
+    const Shape& shape) {
+  std::vector<dnnl::memory::dim> dims;
+
+  for (size_t i = 0; i < shape.dims(); ++i) {
+    dims.push_back(static_cast<dnnl::memory::dim>(shape.at(i)));
+  }
+
+  return dims;
+}
+
+Shape ConcatLayerOneDnn::calculate_output_shape(
+    const std::vector<Tensor>& inputs, int64_t axis) {
+  const Shape& base = inputs[0].get_shape();
+
+  std::vector<size_t> dims(base.dims());
+
+  for (size_t i = 0; i < base.dims(); ++i) {
+    dims[i] = base[i];
+  }
+
+  dims[axis] = 0;
+
+  for (const auto& t : inputs) {
+    dims[axis] += t.get_shape()[axis];
+  }
+
+  return Shape(dims);
+}
+
+int64_t ConcatLayerOneDnn::normalize_axis(int64_t axis, size_t rank) {
+  if (axis < 0) {
+    axis += rank;
+  }
+
+  if (axis < 0 || axis >= static_cast<int64_t>(rank)) {
+    throw std::runtime_error("ConcatLayerOneDnn: axis out of range");
+  }
+
+  return axis;
+}
+
+}  // namespace it_lab_ai