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 CONCAT_PD_HPP
18#define CONCAT_PD_HPP
19
20#include <assert.h>
21
22#include "c_types_map.hpp"
23#include "nstl.hpp"
24#include "primitive_desc.hpp"
25#include "type_helpers.hpp"
26#include "utils.hpp"
27
28namespace mkldnn {
29namespace impl {
30
31struct concat_pd_t: public primitive_desc_t {
32 concat_pd_t(engine_t *engine, const primitive_attr_t *attr,
33 const memory_desc_t *dst_md, int n, int concat_dim,
34 const memory_desc_t *src_mds)
35 : primitive_desc_t(engine, attr, primitive_kind::concat)
36 , n_(n), concat_dim_(concat_dim), dst_md_(*dst_md)
37 {
38 src_mds_.reserve(n_);
39 for (int i = 0; i < n_; ++i) src_mds_.push_back(src_mds[i]);
40 }
41
42 concat_pd_t(const concat_pd_t &rhs) = default;
43
44 virtual void init_info() override { impl::init_info(this, this->info_); }
45
46 virtual arg_usage_t arg_usage(primitive_arg_index_t arg) const override {
47 if (arg >= MKLDNN_ARG_MULTIPLE_SRC
48 && arg < MKLDNN_ARG_MULTIPLE_SRC + n_inputs())
49 return arg_usage_t::input;
50
51 if (arg == MKLDNN_ARG_DST)
52 return arg_usage_t::output;
53
54 return primitive_desc_t::arg_usage(arg);
55 }
56
57 virtual const memory_desc_t *src_md(int index = 0) const override
58 { return index < n_inputs() ? &src_mds_[index] : nullptr; }
59 virtual const memory_desc_t *dst_md(int index = 0) const override
60 { return index == 0 ? &dst_md_ : nullptr; }
61
62 virtual int n_inputs() const override { return n_; }
63 virtual int n_outputs() const override { return 1; }
64
65 int concat_dim() const { return concat_dim_; }
66
67 const memory_desc_t *src_image_md(int index = 0) const
68 { return index < n_inputs() ? &src_image_mds_[index] : nullptr; }
69
70protected:
71 int n_, concat_dim_;
72 memory_desc_t dst_md_;
73 nstl::vector<memory_desc_t> src_mds_;
74
75 /* contains images of srcs in the dst memory (if possible)
76 * Lives here to simplify some implementations. An implementation might
77 * use this auxiliary array iff init() returned success */
78 nstl::vector<memory_desc_t> src_image_mds_;
79
80protected:
81 /* inits src_image_mds_ and dst_md_ in simple cases. The call may fail */
82 status_t init() {
83 bool ok = true
84 && set_default_params() == status::success
85 && attr()->has_default_values();
86 if (!ok) return status::unimplemented;
87
88 for (int i = 0; i < n_; ++i) {
89 const memory_desc_wrapper i_d(&src_mds_[i]);
90 if (!i_d.is_blocking_desc() || i_d.is_additional_buffer())
91 return status::unimplemented;
92 }
93
94 const int ndims = dst_md_.ndims;
95 int current_concat_dim_offset = 0;
96 for (int i = 0; i < n_; ++i) {
97 const int dim = src_mds_[i].dims[concat_dim_];
98 dims_t dims, offsets = {};
99 utils::array_copy(dims, dst_md_.dims, ndims);
100 dims[concat_dim_] = dim;
101 offsets[concat_dim_] = current_concat_dim_offset;
102
103 memory_desc_t src_img_d;
104 status_t status = mkldnn_memory_desc_init_submemory(&src_img_d,
105 &dst_md_, dims, offsets);
106 if (status != status::success) return status;
107 src_image_mds_.push_back(src_img_d);
108 current_concat_dim_offset += dim;
109 }
110
111 return status::success;
112 }
113
114 status_t set_default_params() {
115 if (dst_md_.format_kind != format_kind::any)
116 return status::success;
117
118 const int ndims = dst_md_.ndims;
119
120 /* The stupidest ever heuristics (but not the same as we had before):
121 * - Pick the first non-plain format;
122 * - If all formats are plain or it is not possible to create a
123 * blocked format for the output, pick the format of the plain input
124 * - If this fails as well, use plain layout (abcd...)
125 */
126 status_t status = status::unimplemented;
127 for (int i = 0; i < n_; ++i) {
128 const memory_desc_wrapper src_d(src_mds_[i]);
129 if (src_d.is_blocking_desc() && !src_d.is_plain()) {
130 status = memory_desc_init_by_blocking_desc(dst_md_,
131 src_d.blocking_desc());
132 if (status == status::success) break;
133 }
134 }
135
136 if (status == status::success) {
137 /* check if we can create a sub-memory for the dst */
138 bool desired_format_ok = true;
139 int current_concat_dim_offset = 0;
140 for (int i = 0; i < n_; ++i) {
141 const int dim = src_mds_[i].dims[concat_dim_];
142 dims_t dims, offsets = {};
143 utils::array_copy(dims, dst_md_.dims, ndims);
144 dims[concat_dim_] = dim;
145 offsets[concat_dim_] = current_concat_dim_offset;
146
147 memory_desc_t src_img_d;
148 status_t status = mkldnn_memory_desc_init_submemory(&src_img_d,
149 &dst_md_, dims, offsets);
150 if (status != status::success) {
151 desired_format_ok = false;
152 break;
153 }
154 current_concat_dim_offset += dim;
155 }
156
157 if (!desired_format_ok)
158 status = status::unimplemented;
159 }
160
161 /* if no success so far, try using the format of the first plain input */
162 if (status != status::success) {
163 for (int i = 0; i < n_; ++i) {
164 const memory_desc_wrapper src_d(src_mds_[i]);
165 if (src_d.is_blocking_desc() && src_d.is_plain()) {
166 status = memory_desc_init_by_blocking_desc(dst_md_,
167 memory_desc_wrapper(src_mds_[0]).blocking_desc());
168 if (status == status::success) return status;
169 }
170 }
171 }
172
173 /* the last line of defense: use plain abcd... format */
174 if (status != status::success)
175 status = memory_desc_init_by_strides(dst_md_, nullptr);
176
177 return status;
178 }
179};
180
181#define DECLARE_CONCAT_PD_t(impl_name, ...) \
182 static status_t create(concat_pd_t **concat_pd, \
183 engine_t *engine, const primitive_attr_t *attr, \
184 const memory_desc_t *dst_md, int n, int concat_dim, \
185 const memory_desc_t *src_mds) { \
186 using namespace status; \
187 auto _pd = new pd_t(engine, attr, dst_md, n, concat_dim, src_mds); \
188 if (_pd == nullptr) return out_of_memory; \
189 if (_pd->init() != success) { delete _pd; return unimplemented; } \
190 return safe_ptr_assign<concat_pd_t>(*concat_pd, _pd); \
191 } \
192 virtual status_t create_primitive(primitive_t **p) const override { \
193 double ms = get_msec(); \
194 auto ret = safe_ptr_assign<primitive_t>(*p, new (__VA_ARGS__)(this)); \
195 ms = get_msec() - ms; \
196 if (mkldnn_verbose()->level >= 2) { \
197 printf("mkldnn_verbose,create,%s,%g\n", this->info(), ms); \
198 fflush(0); \
199 } \
200 return ret; \
201 } \
202 virtual pd_t *clone() const override { return new pd_t(*this); } \
203 virtual const char *name() const override { return impl_name; } \
204
205#define DECLARE_CONCAT_PD_T(impl_name, ...) \
206 DECLARE_CONCAT_PD_t(impl_name, __VA_ARGS__)
207
208}
209}
210
211#endif
212