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
29namespace mkldnn {
30namespace impl {
31
32struct 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
68protected:
69 int n_;
70 nstl::vector<float> scales_;
71 memory_desc_t dst_md_;
72 nstl::vector<memory_desc_t> src_mds_;
73
74protected:
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