Update EWLayer.cpp

Author: Semyon1104

src/layers_oneDNN/EWLayer.cpp

@@ -1,7 +1,5 @@
 #include "layers_oneDNN/EWLayer.hpp"
 
-#include <chrono>
-#include <cmath>
 #include <iostream>
 #include <stdexcept>
 
@@ -22,37 +20,79 @@ void EwLayerOneDnn::run(const std::vector<Tensor>& input,
     if (data_type == Type::kFloat) {
       const std::vector<float>& input_data = *input_tensor.as<float>();
       std::vector<float> output_data(input_data.size());
+
       dnnl::memory src_mem = dnnl::memory(
           memory_desc_, *engine_, const_cast<float*>(input_data.data()));
       dnnl::memory dst_mem =
           dnnl::memory(memory_desc_, *engine_, output_data.data());
+
       eltwise_prim_->execute(
           *stream_, {{DNNL_ARG_SRC, src_mem}, {DNNL_ARG_DST, dst_mem}});
       stream_->wait();
+
       output[0] = make_tensor(output_data, input_tensor.get_shape());
     } else if (data_type == Type::kInt) {
       const std::vector<int>& input_data = *input_tensor.as<int>();
       std::vector<int> output_data(input_data.size());
 
-      std::vector<float> float_input;
-      float_input.reserve(input_data.size());
-      for (int val : input_data) {
-        float_input.push_back(static_cast<float>(val));
+      if (memory_desc_.get_data_type() != dnnl::memory::data_type::s32) {
+        std::vector<dnnl::memory::dim> dims;
+        const Shape& shape = input_tensor.get_shape();
+        for (size_t i = 0; i < shape.dims(); i++) {
+          dims.push_back(static_cast<dnnl::memory::dim>(shape.at(i)));
+        }
+
+        dnnl::memory::format_tag format;
+        switch (dims.size()) {
+          case 1:
+            format = dnnl::memory::format_tag::a;
+            break;
+          case 2:
+            format = dnnl::memory::format_tag::ab;
+            break;
+          case 3:
+            format = dnnl::memory::format_tag::abc;
+            break;
+          case 4:
+            format = dnnl::memory::format_tag::abcd;
+            break;
+          case 5:
+            format = dnnl::memory::format_tag::abcde;
+            break;
+          default:
+            throw std::invalid_argument("Unsupported tensor dimensionality: " +
+                                        std::to_string(dims.size()));
+        }
+
+        memory_desc_ =
+            dnnl::memory::desc(dims, dnnl::memory::data_type::s32, format);
+
+        float primitive_alpha = 0.0F;
+        float primitive_beta = 0.0F;
+
+        if (func_ == "relu") {
+          primitive_alpha = 0.0F;
+        } else if (func_ == "linear") {
+          primitive_alpha = alpha_;
+          primitive_beta = beta_;
+        }
+
+        auto eltwise_pd = dnnl::eltwise_forward::primitive_desc(
+            *engine_, dnnl::prop_kind::forward_inference, get_algorithm(),
+            memory_desc_, memory_desc_, primitive_alpha, primitive_beta);
+
+        eltwise_prim_ = std::make_unique<dnnl::eltwise_forward>(eltwise_pd);
       }
 
-      std::vector<float> float_output(input_data.size());
-
-      dnnl::memory src_mem =
-          dnnl::memory(memory_desc_, *engine_, float_input.data());
+      dnnl::memory src_mem = dnnl::memory(memory_desc_, *engine_,
+                                          const_cast<int*>(input_data.data()));
       dnnl::memory dst_mem =
-          dnnl::memory(memory_desc_, *engine_, float_output.data());
+          dnnl::memory(memory_desc_, *engine_, output_data.data());
+
       eltwise_prim_->execute(
           *stream_, {{DNNL_ARG_SRC, src_mem}, {DNNL_ARG_DST, dst_mem}});
       stream_->wait();
 
-      for (size_t i = 0; i < float_output.size(); ++i) {
-        output_data[i] = static_cast<int>(std::round(float_output[i]));
-      }
       output[0] = make_tensor(output_data, input_tensor.get_shape());
     } else {
       throw std::runtime_error("EwLayerOneDnn: Unsupported data type");
@@ -112,13 +152,7 @@ void EwLayerOneDnn::initialize_onednn(const Shape& shape, Type data_type) {
                                     std::to_string(dims.size()));
     }
 
-    dnnl::memory::data_type dnnl_data_type;
-    if (data_type == Type::kFloat) {
-      dnnl_data_type = dnnl::memory::data_type::f32;
-    } else {
-      throw std::invalid_argument("Unsupported data type for oneDNN EW layer");
-    }
-
+    dnnl::memory::data_type dnnl_data_type = get_dnnl_data_type(data_type);
     memory_desc_ = dnnl::memory::desc(dims, dnnl_data_type, format);
 
     dnnl::algorithm algo = get_algorithm();
@@ -148,6 +182,17 @@ void EwLayerOneDnn::initialize_onednn(const Shape& shape, Type data_type) {
   }
 }
 
+dnnl::memory::data_type EwLayerOneDnn::get_dnnl_data_type(Type type) const {
+  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::algorithm EwLayerOneDnn::get_algorithm() const {
   if (func_ == "relu") {
     return dnnl::algorithm::eltwise_relu;
@@ -170,4 +215,4 @@ bool EwLayerOneDnn::is_function_supported(const std::string& function) {
           function == "linear");
 }
 
-}  // namespace it_lab_ai
+}  // namespace it_lab_ai