Skip to content

Commit f540558

Browse files
Semyon1104Semyon1104
committed
fix realization
1 parent f940135 commit f540558

2 files changed

Lines changed: 82 additions & 151 deletions

File tree

‎include/layers_oneDNN/EWLayer.hpp‎

Lines changed: 3 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -38,10 +38,13 @@ class EwLayerOneDnn : public Layer {
3838
[[nodiscard]] dnnl::algorithm get_algorithm() const;
3939
void validate_input(const std::vector<Tensor>& input) const;
4040
[[nodiscard]] static dnnl::memory::data_type get_dnnl_data_type(Type type);
41+
static dnnl::memory::format_tag pick_format(size_t ndims);
4142

4243
std::string func_;
4344
float alpha_;
4445
float beta_;
46+
Shape last_shape_;
47+
Type last_type_;
4548

4649
std::unique_ptr<dnnl::engine> engine_;
4750
std::unique_ptr<dnnl::stream> stream_;

‎src/layers_oneDNN/EWLayer.cpp‎

Lines changed: 79 additions & 151 deletions
Original file line numberDiff line numberDiff line change
@@ -9,96 +9,41 @@ void EwLayerOneDnn::run(const std::vector<Tensor>& input,
99
std::vector<Tensor>& output) {
1010
validate_input(input);
1111

12-
const Tensor& input_tensor = input[0];
13-
Type data_type = input_tensor.get_type();
12+
const Tensor& in = input[0];
13+
Type type = in.get_type();
1414

15-
if (!initialized_) {
16-
initialize_onednn(input_tensor.get_shape(), data_type);
15+
bool need_reinit =
16+
!initialized_ || last_type_ != type || last_shape_ != in.get_shape();
17+
18+
if (need_reinit) {
19+
initialize_onednn(in.get_shape(), type);
1720
}
1821

19-
try {
20-
if (data_type == Type::kFloat) {
21-
const std::vector<float>& input_data = *input_tensor.as<float>();
22-
std::vector<float> output_data(input_data.size());
23-
24-
dnnl::memory src_mem = dnnl::memory(
25-
memory_desc_, *engine_, const_cast<float*>(input_data.data()));
26-
dnnl::memory dst_mem =
27-
dnnl::memory(memory_desc_, *engine_, output_data.data());
28-
29-
eltwise_prim_->execute(
30-
*stream_, {{DNNL_ARG_SRC, src_mem}, {DNNL_ARG_DST, dst_mem}});
31-
stream_->wait();
32-
33-
output[0] = make_tensor(output_data, input_tensor.get_shape());
34-
} else if (data_type == Type::kInt) {
35-
const std::vector<int>& input_data = *input_tensor.as<int>();
36-
std::vector<int> output_data(input_data.size());
37-
38-
if (memory_desc_.get_data_type() != dnnl::memory::data_type::s32) {
39-
std::vector<dnnl::memory::dim> dims;
40-
const Shape& shape = input_tensor.get_shape();
41-
for (size_t i = 0; i < shape.dims(); i++) {
42-
dims.push_back(static_cast<dnnl::memory::dim>(shape.at(i)));
43-
}
44-
45-
dnnl::memory::format_tag format;
46-
switch (dims.size()) {
47-
case 1:
48-
format = dnnl::memory::format_tag::a;
49-
break;
50-
case 2:
51-
format = dnnl::memory::format_tag::ab;
52-
break;
53-
case 3:
54-
format = dnnl::memory::format_tag::abc;
55-
break;
56-
case 4:
57-
format = dnnl::memory::format_tag::abcd;
58-
break;
59-
case 5:
60-
format = dnnl::memory::format_tag::abcde;
61-
break;
62-
default:
63-
throw std::invalid_argument("Unsupported tensor dimensionality: " +
64-
std::to_string(dims.size()));
65-
}
66-
67-
memory_desc_ =
68-
dnnl::memory::desc(dims, dnnl::memory::data_type::s32, format);
69-
70-
float primitive_alpha = 0.0F;
71-
float primitive_beta = 0.0F;
72-
73-
if (func_ == "relu") {
74-
primitive_alpha = 0.0F;
75-
} else if (func_ == "linear") {
76-
primitive_alpha = alpha_;
77-
primitive_beta = beta_;
78-
}
79-
80-
auto eltwise_pd = dnnl::eltwise_forward::primitive_desc(
81-
*engine_, dnnl::prop_kind::forward_inference, get_algorithm(),
82-
memory_desc_, memory_desc_, primitive_alpha, primitive_beta);
83-
84-
eltwise_prim_ = std::make_unique<dnnl::eltwise_forward>(eltwise_pd);
85-
}
86-
87-
dnnl::memory src_mem = dnnl::memory(memory_desc_, *engine_,
88-
const_cast<int*>(input_data.data()));
89-
dnnl::memory dst_mem =
90-
dnnl::memory(memory_desc_, *engine_, output_data.data());
91-
92-
eltwise_prim_->execute(
93-
*stream_, {{DNNL_ARG_SRC, src_mem}, {DNNL_ARG_DST, dst_mem}});
94-
stream_->wait();
95-
96-
output[0] = make_tensor(output_data, input_tensor.get_shape());
97-
}
98-
99-
} catch (const std::exception& e) {
100-
std::cerr << "oneDNN execution failed: " << e.what() << '\n';
101-
throw;
22+
if (type == Type::kFloat) {
23+
const auto& src = *in.as<float>();
24+
std::vector<float> dst(src.size());
25+
26+
dnnl::memory src_mem(memory_desc_, *engine_,
27+
const_cast<float*>(src.data()));
28+
dnnl::memory dst_mem(memory_desc_, *engine_, dst.data());
29+
30+
eltwise_prim_->execute(*stream_,
31+
{{DNNL_ARG_SRC, src_mem}, {DNNL_ARG_DST, dst_mem}});
32+
33+
stream_->wait();
34+
output[0] = make_tensor(dst, in.get_shape());
35+
} else if (type == Type::kInt) {
36+
const auto& src = *in.as<int>();
37+
std::vector<int> dst(src.size());
38+
39+
dnnl::memory src_mem(memory_desc_, *engine_, const_cast<int*>(src.data()));
40+
dnnl::memory dst_mem(memory_desc_, *engine_, dst.data());
41+
42+
eltwise_prim_->execute(*stream_,
43+
{{DNNL_ARG_SRC, src_mem}, {DNNL_ARG_DST, dst_mem}});
44+
45+
stream_->wait();
46+
output[0] = make_tensor(dst, in.get_shape());
10247
}
10348
}
10449

@@ -110,74 +55,39 @@ void EwLayerOneDnn::validate_input(const std::vector<Tensor>& input) const {
11055
if (!is_function_supported(func_)) {
11156
throw std::invalid_argument("Unsupported function for oneDNN: " + func_);
11257
}
113-
114-
Type data_type = input[0].get_type();
115-
if (data_type != Type::kFloat && data_type != Type::kInt) {
116-
throw std::runtime_error(
117-
"EwLayerOneDnn supports only float and int tensors");
118-
}
11958
}
12059

12160
void EwLayerOneDnn::initialize_onednn(const Shape& shape, Type data_type) {
122-
try {
123-
engine_ = std::make_unique<dnnl::engine>(dnnl::engine::kind::cpu, 0);
124-
stream_ = std::make_unique<dnnl::stream>(*engine_);
125-
126-
std::vector<dnnl::memory::dim> dims;
127-
for (size_t i = 0; i < shape.dims(); i++) {
128-
dims.push_back(static_cast<dnnl::memory::dim>(shape.at(i)));
129-
}
130-
131-
dnnl::memory::format_tag format;
132-
switch (dims.size()) {
133-
case 1:
134-
format = dnnl::memory::format_tag::a;
135-
break;
136-
case 2:
137-
format = dnnl::memory::format_tag::ab;
138-
break;
139-
case 3:
140-
format = dnnl::memory::format_tag::abc;
141-
break;
142-
case 4:
143-
format = dnnl::memory::format_tag::abcd;
144-
break;
145-
case 5:
146-
format = dnnl::memory::format_tag::abcde;
147-
break;
148-
default:
149-
throw std::invalid_argument("Unsupported tensor dimensionality: " +
150-
std::to_string(dims.size()));
151-
}
152-
153-
dnnl::memory::data_type dnnl_data_type = get_dnnl_data_type(data_type);
154-
memory_desc_ = dnnl::memory::desc(dims, dnnl_data_type, format);
155-
156-
dnnl::algorithm algo = get_algorithm();
157-
158-
float primitive_alpha = 0.0F;
159-
float primitive_beta = 0.0F;
160-
161-
if (func_ == "relu") {
162-
primitive_alpha = 0.0F;
163-
} else if (func_ == "linear") {
164-
primitive_alpha = alpha_;
165-
primitive_beta = beta_;
166-
}
167-
168-
auto eltwise_pd = dnnl::eltwise_forward::primitive_desc(
169-
*engine_, dnnl::prop_kind::forward_inference, algo, memory_desc_,
170-
memory_desc_, primitive_alpha, primitive_beta);
171-
172-
eltwise_prim_ = std::make_unique<dnnl::eltwise_forward>(eltwise_pd);
173-
174-
initialized_ = true;
175-
176-
} catch (const std::exception& e) {
177-
std::cerr << "oneDNN initialization failed for function '" << func_
178-
<< "': " << e.what() << '\n';
179-
throw;
61+
engine_ = std::make_unique<dnnl::engine>(dnnl::engine::kind::cpu, 0);
62+
stream_ = std::make_unique<dnnl::stream>(*engine_);
63+
64+
std::vector<dnnl::memory::dim> dims;
65+
for (size_t i = 0; i < shape.dims(); ++i) {
66+
dims.push_back(static_cast<dnnl::memory::dim>(shape.at(i)));
18067
}
68+
69+
auto format = pick_format(dims.size());
70+
auto dnnl_type = get_dnnl_data_type(data_type);
71+
72+
memory_desc_ = dnnl::memory::desc(dims, dnnl_type, format);
73+
74+
float alpha = 0.0f;
75+
float beta = 0.0f;
76+
77+
if (func_ == "linear") {
78+
alpha = alpha_;
79+
beta = beta_;
80+
}
81+
82+
auto eltwise_pd = dnnl::eltwise_forward::primitive_desc(
83+
*engine_, dnnl::prop_kind::forward_inference, get_algorithm(),
84+
memory_desc_, memory_desc_, alpha, beta);
85+
86+
eltwise_prim_ = std::make_unique<dnnl::eltwise_forward>(eltwise_pd);
87+
88+
last_shape_ = shape;
89+
last_type_ = data_type;
90+
initialized_ = true;
18191
}
18292

18393
dnnl::memory::data_type EwLayerOneDnn::get_dnnl_data_type(Type type) {
@@ -213,4 +123,22 @@ bool EwLayerOneDnn::is_function_supported(const std::string& function) {
213123
function == "linear");
214124
}
215125

126+
dnnl::memory::format_tag EwLayerOneDnn::pick_format(size_t ndims) {
127+
switch (ndims) {
128+
case 1:
129+
return dnnl::memory::format_tag::a;
130+
case 2:
131+
return dnnl::memory::format_tag::ab;
132+
case 3:
133+
return dnnl::memory::format_tag::abc;
134+
case 4:
135+
return dnnl::memory::format_tag::abcd;
136+
case 5:
137+
return dnnl::memory::format_tag::abcde;
138+
default:
139+
throw std::invalid_argument("Unsupported tensor dimensionality: " +
140+
std::to_string(ndims));
141+
}
142+
}
143+
216144
} // namespace it_lab_ai

0 commit comments

Comments
 (0)