1 | /******************************************************************************* |
2 | * Copyright 2016-2018 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 | #include <assert.h> |
18 | #include "mkldnn.h" |
19 | |
20 | #include "c_types_map.hpp" |
21 | #include "memory_desc_wrapper.hpp" |
22 | #include "utils.hpp" |
23 | |
24 | using namespace mkldnn::impl; |
25 | using namespace mkldnn::impl::utils; |
26 | using namespace mkldnn::impl::status; |
27 | using namespace mkldnn::impl::prop_kind; |
28 | using namespace mkldnn::impl::alg_kind; |
29 | using namespace mkldnn::impl::types; |
30 | |
31 | namespace { |
32 | status_t softmax_desc_init(softmax_desc_t *softmax_desc, prop_kind_t prop_kind, |
33 | const memory_desc_t *data_desc, const memory_desc_t *diff_desc, int softmax_axis) { |
34 | bool args_ok = true |
35 | && !any_null(softmax_desc, data_desc) |
36 | && 0 <= softmax_axis |
37 | && softmax_axis < data_desc->ndims; |
38 | if (!args_ok) return invalid_arguments; |
39 | |
40 | auto sd = softmax_desc_t(); |
41 | sd.primitive_kind = primitive_kind::softmax; |
42 | sd.prop_kind = prop_kind; |
43 | |
44 | bool is_bwd = (sd.prop_kind == backward_data); |
45 | sd.data_desc = *data_desc; |
46 | sd.diff_desc = is_bwd ? *diff_desc : zero_md(); |
47 | sd.softmax_axis = softmax_axis; |
48 | |
49 | *softmax_desc = sd; |
50 | return success; |
51 | } |
52 | } |
53 | |
54 | status_t mkldnn_softmax_forward_desc_init(softmax_desc_t *softmax_desc, |
55 | prop_kind_t prop_kind, const memory_desc_t *data_desc, |
56 | int softmax_axis) { |
57 | if (!one_of(prop_kind, forward_inference, forward_training)) |
58 | return invalid_arguments; |
59 | return softmax_desc_init(softmax_desc, prop_kind, data_desc, nullptr, softmax_axis); |
60 | } |
61 | |
62 | status_t mkldnn_softmax_backward_desc_init(softmax_desc_t *softmax_desc, |
63 | const memory_desc_t *diff_desc, const mkldnn_memory_desc_t *data_desc, |
64 | int softmax_axis) { |
65 | return softmax_desc_init(softmax_desc, prop_kind::backward_data, |
66 | data_desc, diff_desc, softmax_axis); |
67 | } |
68 | // vim: et ts=4 sw=4 cindent cino^=l0,\:0,N-s |
69 | |