diff --git a/mindspore/lite/src/runtime/kernel/arm/fp32/reverse_fp32.cc b/mindspore/lite/src/runtime/kernel/arm/fp32/reverse_fp32.cc index ef592e1a092..4a9189979a6 100644 --- a/mindspore/lite/src/runtime/kernel/arm/fp32/reverse_fp32.cc +++ b/mindspore/lite/src/runtime/kernel/arm/fp32/reverse_fp32.cc @@ -37,6 +37,9 @@ int ReverseCPUKernel::Stride(int index) { } int ReverseCPUKernel::ReSize() { + // trans negative to positive axis + UpdateAxisInfo(); + data_size_ = in_tensors_.at(0)->ElementsNum(); thread_sz_count_ = MSMIN(op_parameter_->thread_num_, data_size_); thread_sz_stride_ = UP_DIV(data_size_, thread_sz_count_); @@ -134,6 +137,16 @@ int ReverseCPUKernel::Run() { return RET_OK; } +void ReverseCPUKernel::UpdateAxisInfo() { + auto reverse_param = reinterpret_cast(op_parameter_); + int in_shape_len = in_tensors_.front()->shape().size(); + for (int i = 0; i < reverse_param->num_axis_; ++i) { + if (reverse_param->axis_[i] < 0) { + reverse_param->axis_[i] += in_shape_len; + } + } +} + REG_KERNEL(kCPU, kNumberTypeFloat32, PrimitiveType_ReverseV2, LiteKernelCreator) REG_KERNEL(kCPU, kNumberTypeInt32, PrimitiveType_ReverseV2, LiteKernelCreator) } // namespace mindspore::kernel diff --git a/mindspore/lite/src/runtime/kernel/arm/fp32/reverse_fp32.h b/mindspore/lite/src/runtime/kernel/arm/fp32/reverse_fp32.h index 5486dd6f7ec..c78240b92aa 100644 --- a/mindspore/lite/src/runtime/kernel/arm/fp32/reverse_fp32.h +++ b/mindspore/lite/src/runtime/kernel/arm/fp32/reverse_fp32.h @@ -28,7 +28,7 @@ class ReverseCPUKernel : public LiteKernel { ReverseCPUKernel(OpParameter *parameter, const std::vector &inputs, const std::vector &outputs, const lite::InnerContext *ctx) : LiteKernel(parameter, inputs, outputs, ctx) {} - ~ReverseCPUKernel() { + ~ReverseCPUKernel() override { if (tmp_ != nullptr) { free(tmp_); tmp_ = nullptr; @@ -42,6 +42,7 @@ class ReverseCPUKernel : public LiteKernel { int DoReverse(int task_id); private: + void UpdateAxisInfo(); int thread_sz_count_ = 0; int thread_sz_stride_ = 0; int data_size_ = 0;