1 | /******************************************************************************* |
2 | * Copyright 2019 Intel Corporation |
3 | * |
4 | * Licensed under the Apache License, Version 2.0 (the "License"); |
5 | * you may not use this file except in compliance with the License. |
6 | * You may obtain a copy of the License at |
7 | * |
8 | * http://www.apache.org/licenses/LICENSE-2.0 |
9 | * |
10 | * Unless required by applicable law or agreed to in writing, software |
11 | * distributed under the License is distributed on an "AS IS" BASIS, |
12 | * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. |
13 | * See the License for the specific language governing permissions and |
14 | * limitations under the License. |
15 | *******************************************************************************/ |
16 | |
17 | #ifndef SUM_PD_HPP |
18 | #define SUM_PD_HPP |
19 | |
20 | #include <assert.h> |
21 | #include "mkldnn.h" |
22 | |
23 | #include "c_types_map.hpp" |
24 | #include "nstl.hpp" |
25 | #include "primitive_desc.hpp" |
26 | #include "type_helpers.hpp" |
27 | #include "utils.hpp" |
28 | |
29 | namespace mkldnn { |
30 | namespace impl { |
31 | |
32 | struct sum_pd_t: public primitive_desc_t { |
33 | sum_pd_t(engine_t *engine, const primitive_attr_t *attr, |
34 | const memory_desc_t *dst_md, int n, const float *scales, |
35 | const memory_desc_t *src_mds) |
36 | : primitive_desc_t(engine, attr, primitive_kind::sum) |
37 | , n_(n), dst_md_(*dst_md) |
38 | { |
39 | scales_.reserve(n_); |
40 | for (int i = 0; i < n_; ++i) scales_.push_back(scales[i]); |
41 | src_mds_.reserve(n_); |
42 | for (int i = 0; i < n_; ++i) src_mds_.push_back(src_mds[i]); |
43 | } |
44 | |
45 | virtual void init_info() override { impl::init_info(this, this->info_); } |
46 | |
47 | virtual arg_usage_t arg_usage(primitive_arg_index_t arg) const override { |
48 | if (arg >= MKLDNN_ARG_MULTIPLE_SRC |
49 | && arg < MKLDNN_ARG_MULTIPLE_SRC + n_inputs()) |
50 | return arg_usage_t::input; |
51 | |
52 | if (arg == MKLDNN_ARG_DST) |
53 | return arg_usage_t::output; |
54 | |
55 | return primitive_desc_t::arg_usage(arg); |
56 | } |
57 | |
58 | virtual const memory_desc_t *src_md(int index = 0) const override |
59 | { return index < n_inputs() ? &src_mds_[index] : nullptr; } |
60 | virtual const memory_desc_t *dst_md(int index = 0) const override |
61 | { return index == 0 ? &dst_md_ : nullptr; } |
62 | |
63 | virtual int n_inputs() const override { return n_; } |
64 | virtual int n_outputs() const override { return 1; } |
65 | |
66 | const float *scales() const { return &scales_[0]; } |
67 | |
68 | protected: |
69 | int n_; |
70 | nstl::vector<float> scales_; |
71 | memory_desc_t dst_md_; |
72 | nstl::vector<memory_desc_t> src_mds_; |
73 | |
74 | protected: |
75 | /* inits dst_md_ in simple cases. The call may fail. */ |
76 | status_t init() { |
77 | for (int i = 0; i < n_; ++i) { |
78 | const memory_desc_wrapper src_d(&src_mds_[i]); |
79 | if (!src_d.is_blocking_desc() || src_d.is_additional_buffer()) |
80 | return status::unimplemented; |
81 | } |
82 | bool ok = true |
83 | && set_default_params() == status::success |
84 | && attr()->has_default_values(); |
85 | return ok ? status::success : status::unimplemented; |
86 | } |
87 | |
88 | status_t set_default_params() { |
89 | if (dst_md_.format_kind != format_kind::any) |
90 | return status::success; |
91 | |
92 | /* The stupidest ever heuristics (but not the same as we had before): |
93 | * - Pick the first non-plain format; |
94 | * - If all formats are plain, pick the format of the first input |
95 | */ |
96 | for (int i = 0; i < n_; ++i) { |
97 | const memory_desc_wrapper src_d(src_mds_[i]); |
98 | if (!src_d.is_plain() && src_d.is_blocking_desc()) { |
99 | return memory_desc_init_by_blocking_desc(dst_md_, |
100 | src_d.blocking_desc()); |
101 | } |
102 | } |
103 | |
104 | if (src_mds_[0].format_kind != format_kind::blocked) |
105 | return status::unimplemented; |
106 | |
107 | dst_md_ = src_mds_[0]; |
108 | |
109 | return status::success; |
110 | } |
111 | }; |
112 | |
113 | #define DECLARE_SUM_PD_t(impl_name, ...) \ |
114 | static status_t create(sum_pd_t **sum_pd, \ |
115 | engine_t *engine, const primitive_attr_t *attr, \ |
116 | const memory_desc_t *dst_md, int n, const float *scales, \ |
117 | const memory_desc_t *src_mds) { \ |
118 | using namespace status; \ |
119 | auto _pd = new pd_t(engine, attr, dst_md, n, scales, src_mds); \ |
120 | if (_pd == nullptr) return out_of_memory; \ |
121 | if (_pd->init() != success) { delete _pd; return unimplemented; } \ |
122 | return safe_ptr_assign<sum_pd_t>(*sum_pd, _pd); \ |
123 | } \ |
124 | virtual status_t create_primitive(primitive_t **p) const override { \ |
125 | double ms = get_msec(); \ |
126 | auto ret = safe_ptr_assign<primitive_t>(*p, new (__VA_ARGS__)(this)); \ |
127 | ms = get_msec() - ms; \ |
128 | if (mkldnn_verbose()->level >= 2) { \ |
129 | printf("mkldnn_verbose,create,%s,%g\n", this->info(), ms); \ |
130 | fflush(0); \ |
131 | } \ |
132 | return ret; \ |
133 | } \ |
134 | virtual pd_t *clone() const override { return new pd_t(*this); } \ |
135 | virtual const char *name() const override { return impl_name; } \ |
136 | |
137 | #define DECLARE_SUM_PD_T(impl_name, ...) \ |
138 | DECLARE_SUM_PD_t(impl_name, __VA_ARGS__) |
139 | |
140 | } |
141 | } |
142 | |
143 | #endif |
144 | |