1 /*
2 * Copyright (c) 2017-2020 Arm Limited.
3 *
4 * SPDX-License-Identifier: MIT
5 *
6 * Permission is hereby granted, free of charge, to any person obtaining a copy
7 * of this software and associated documentation files (the "Software"), to
8 * deal in the Software without restriction, including without limitation the
9 * rights to use, copy, modify, merge, publish, distribute, sublicense, and/or
10 * sell copies of the Software, and to permit persons to whom the Software is
11 * furnished to do so, subject to the following conditions:
12 *
13 * The above copyright notice and this permission notice shall be included in all
14 * copies or substantial portions of the Software.
15 *
16 * THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
17 * IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
18 * FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
19 * AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
20 * LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
21 * OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
22 * SOFTWARE.
23 */
24 #include "arm_compute/runtime/NEON/functions/NESoftmaxLayer.h"
25
26 #include "arm_compute/core/Helpers.h"
27 #include "arm_compute/core/utils/misc/ShapeCalculator.h"
28 #include "arm_compute/runtime/NEON/NEScheduler.h"
29 #include "src/core/NEON/kernels/NEFillBorderKernel.h"
30 #include "src/core/NEON/kernels/NESoftmaxLayerKernel.h"
31 #include "src/core/NEON/kernels/NESoftmaxLayerKernel.h"
32 #include "src/core/helpers/SoftmaxHelpers.h"
33 #include "support/MemorySupport.h"
34
35 namespace arm_compute
36 {
37 template <bool IS_LOG>
38 NESoftmaxLayerGeneric<IS_LOG>::~NESoftmaxLayerGeneric() = default;
39
40 template <bool IS_LOG>
NESoftmaxLayerGeneric(std::shared_ptr<IMemoryManager> memory_manager)41 NESoftmaxLayerGeneric<IS_LOG>::NESoftmaxLayerGeneric(std::shared_ptr<IMemoryManager> memory_manager)
42 : _memory_group(std::move(memory_manager)), _permute_input(), _permute_output(), _max_kernel(), _softmax_kernel(), _fill_border_kernel(), _max(), _tmp(), _input_permuted(), _output_permuted(),
43 _needs_permute(false)
44 {
45 }
46
47 template <bool IS_LOG>
configure(ITensor * input,ITensor * output,float beta,int32_t axis)48 void NESoftmaxLayerGeneric<IS_LOG>::configure(ITensor *input, ITensor *output, float beta, int32_t axis)
49 {
50 // Perform validation step
51 ARM_COMPUTE_ERROR_ON_NULLPTR(input, output);
52 ARM_COMPUTE_ERROR_THROW_ON(NESoftmaxLayerGeneric::validate(input->info(), output->info(), beta, axis));
53
54 const unsigned int actual_axis = static_cast<unsigned int>(wrap_around(axis, static_cast<int32_t>(input->info()->num_dimensions())));
55
56 _needs_permute = actual_axis > 0;
57
58 if(_needs_permute)
59 {
60 // Add to the memory manager _input_permuted
61 _memory_group.manage(&_input_permuted);
62
63 _permute_input.configure(input, &_input_permuted, softmax_helpers::get_permutation_vector_from_softmax_axis(actual_axis));
64 }
65
66 // We want to deal with a 2D input. Either it is the permuted version of the original input (4D case)
67 // or it is the original input case (2D case)
68 ITensor *tmp_input = (_needs_permute ? &_input_permuted : input);
69
70 // Create intermediate tensors shapes
71 const TensorInfo input_info = tmp_input->info()->clone()->reset_padding().set_is_resizable(true);
72 DataType tmp_data_type = is_data_type_quantized_asymmetric(tmp_input->info()->data_type()) ? DataType::F32 : tmp_input->info()->data_type();
73 TensorInfo tensor_info_tmp(input_info.clone()->set_data_type(tmp_data_type));
74
75 // Init intermediate tensors
76 TensorShape max_sum_shape = tmp_input->info()->tensor_shape();
77 max_sum_shape.set(0, 1);
78 _max.allocator()->init(input_info.clone()->set_tensor_shape(max_sum_shape));
79 _tmp.allocator()->init(tensor_info_tmp);
80
81 // Manage intermediate buffers
82 _memory_group.manage(&_max);
83 _memory_group.manage(&_tmp);
84
85 // Configure kernels
86 _max_kernel = arm_compute::support::cpp14::make_unique<NELogits1DMaxKernel>();
87 _softmax_kernel = arm_compute::support::cpp14::make_unique<NELogits1DSoftmaxKernel<IS_LOG>>();
88 _max_kernel->configure(tmp_input, &_max);
89 if(_needs_permute)
90 {
91 // Add to the memory manager _output_permuted
92 _memory_group.manage(&_output_permuted);
93
94 // The normalization kernel stores the result in a permuted output tensor
95 _softmax_kernel->configure(tmp_input, &_max, &_output_permuted, beta, &_tmp);
96 _input_permuted.allocator()->allocate();
97
98 // Re-permute the permuted output into the requested (4D) output
99 _permute_output.configure(&_output_permuted, output, softmax_helpers::get_permutation_vector_from_softmax_axis(actual_axis));
100
101 // Allocate the intermediate permuted tensors
102 _output_permuted.allocator()->allocate();
103 }
104 else
105 {
106 // Softmax 2D case
107 _fill_border_kernel = arm_compute::support::cpp14::make_unique<NEFillBorderKernel>();
108 _fill_border_kernel->configure(tmp_input, _max_kernel->border_size(), BorderMode::REPLICATE);
109 _softmax_kernel->configure(tmp_input, &_max, output, beta, &_tmp);
110 }
111
112 // Allocate intermediate buffers
113 _max.allocator()->allocate();
114 _tmp.allocator()->allocate();
115 }
116
117 template <bool IS_LOG>
validate(const ITensorInfo * input,const ITensorInfo * output,float beta,int32_t axis)118 Status NESoftmaxLayerGeneric<IS_LOG>::validate(const ITensorInfo *input, const ITensorInfo *output, float beta, int32_t axis)
119 {
120 // Perform validation step
121 ARM_COMPUTE_RETURN_ERROR_ON_NULLPTR(input, output);
122 ARM_COMPUTE_RETURN_ERROR_ON_MSG(input->num_dimensions() > 4, "Only up to 4 dimensions are supported");
123 ARM_COMPUTE_UNUSED(beta);
124 ARM_COMPUTE_RETURN_ERROR_ON(axis < static_cast<int32_t>(-input->num_dimensions()) || static_cast<int32_t>(input->num_dimensions()) <= axis);
125
126 // Create intermediate tensor info
127 DataType tmp_data_type = input->data_type();
128 const TensorInfo tensor_info_tmp(input->clone()->set_data_type(tmp_data_type).set_is_resizable(true));
129
130 TensorShape max_sum_shape = input->tensor_shape();
131 max_sum_shape.set(0, 1);
132 const TensorInfo tensor_info_max_sum(input->clone()->set_tensor_shape(max_sum_shape).set_data_type(tmp_data_type).set_quantization_info(input->quantization_info()).set_is_resizable(true));
133 const TensorInfo dont_care;
134
135 const unsigned int actual_axis = static_cast<unsigned int>(wrap_around(axis, static_cast<int32_t>(input->num_dimensions())));
136
137 const bool needs_permute = actual_axis > 0;
138
139 if(needs_permute)
140 {
141 const PermutationVector permutation_vector = softmax_helpers::get_permutation_vector_from_softmax_axis(actual_axis);
142 const TensorShape permuted_shape = misc::shape_calculator::compute_permutation_output_shape(*input, permutation_vector);
143 TensorInfo input_permuted(input->clone()->set_tensor_shape(permuted_shape));
144 ARM_COMPUTE_RETURN_ON_ERROR(NEPermute::validate(input, &input_permuted, permutation_vector));
145 TensorInfo output_permuted(output->clone()->set_tensor_shape(permuted_shape));
146 ARM_COMPUTE_RETURN_ON_ERROR(NEPermute::validate(&output_permuted, output, permutation_vector));
147 }
148
149 ARM_COMPUTE_RETURN_ON_ERROR(NELogits1DMaxKernel::validate(input, &tensor_info_max_sum));
150 ARM_COMPUTE_RETURN_ON_ERROR(NELogits1DSoftmaxKernel<IS_LOG>::validate(&tensor_info_tmp, &tensor_info_max_sum, output, beta, &dont_care));
151
152 return Status{};
153 }
154
155 template <bool IS_LOG>
run()156 void NESoftmaxLayerGeneric<IS_LOG>::run()
157 {
158 MemoryGroupResourceScope scope_mg(_memory_group);
159
160 if(_needs_permute)
161 {
162 _permute_input.run();
163 }
164 else
165 {
166 NEScheduler::get().schedule(_fill_border_kernel.get(), Window::DimY);
167 }
168
169 NEScheduler::get().schedule(_max_kernel.get(), Window::DimY);
170 NEScheduler::get().schedule(_softmax_kernel.get(), Window::DimY);
171
172 if(_needs_permute)
173 {
174 _permute_output.run();
175 }
176 }
177
178 template class NESoftmaxLayerGeneric<false>;
179 template class NESoftmaxLayerGeneric<true>;
180
181 } // namespace arm_compute
182