diff --git a/onnxruntime/core/providers/cpu/math/logsoftmax.cc b/onnxruntime/core/providers/cpu/math/logsoftmax.cc deleted file mode 100644 index cbf2cf4083..0000000000 --- a/onnxruntime/core/providers/cpu/math/logsoftmax.cc +++ /dev/null @@ -1,58 +0,0 @@ -// Copyright (c) Microsoft Corporation. All rights reserved. -// Licensed under the MIT License. - -#include "core/providers/cpu/math/logsoftmax.h" - -#include "core/framework/op_kernel.h" -#include "core/framework/op_kernel_context_internal.h" - -#include "core/providers/common.h" -#include "core/providers/cpu/math/softmax_shared.h" -#include "core/util/math.h" - -namespace onnxruntime { - -template <> -Status LogSoftmax::Compute(OpKernelContext* ctx) const { - auto ctx_internal = static_cast(ctx); - concurrency::ThreadPool* tp = ctx_internal->GetOperatorThreadPool(); - - const auto* tensor_pointer = ctx->Input(0); - if (tensor_pointer == nullptr) return Status(common::ONNXRUNTIME, common::FAIL, "input count mismatch"); - const Tensor& X = *tensor_pointer; - const TensorShape& input_shape{X.Shape()}; - - Tensor* Y = ctx->Output(0, input_shape); - - const int64_t axis = HandleNegativeAxis(axis_, input_shape.NumDimensions()); - - size_t N = input_shape.SizeToDimension(axis); - size_t D = input_shape.SizeFromDimension(axis); - - auto* Ydata = Y->template MutableData(); - - std::vector scale_(N); - std::vector rowmax_(N); - std::vector sum_multiplier_(D, 1.f); // initialize all multiplier values to 1.0 - - const bool logarithmic = true; - auto status = SoftmaxCPU(N, D, X.template Data(), Ydata, - scale_.data(), sum_multiplier_.data(), logarithmic, rowmax_.data(), tp); - - return status; -} - -ONNX_CPU_OPERATOR_VERSIONED_KERNEL( - LogSoftmax, - 1, - 10, - KernelDefBuilder().TypeConstraint("T", DataTypeImpl::GetTensorType()), - LogSoftmax); - -// Opset 11 starts to support Neg Axis. -ONNX_CPU_OPERATOR_KERNEL( - LogSoftmax, - 11, - KernelDefBuilder().TypeConstraint("T", DataTypeImpl::GetTensorType()), - LogSoftmax); -} // namespace onnxruntime diff --git a/onnxruntime/core/providers/cpu/math/logsoftmax.h b/onnxruntime/core/providers/cpu/math/logsoftmax.h deleted file mode 100644 index e1362d8760..0000000000 --- a/onnxruntime/core/providers/cpu/math/logsoftmax.h +++ /dev/null @@ -1,30 +0,0 @@ -// Copyright (c) Microsoft Corporation. All rights reserved. -// Licensed under the MIT License. - -#pragma once - -#include "gsl/gsl_util" - -#include "core/common/common.h" -#include "core/framework/op_kernel.h" -#include "core/util/math_cpuonly.h" - -namespace onnxruntime { -template -class LogSoftmax final : public OpKernel { - public: - LogSoftmax(const OpKernelInfo& info) : OpKernel{info}, axis_{1} { - int64_t axis; - Status status = info.GetAttr("axis", &axis); - - if (status.IsOK()) { - axis_ = gsl::narrow_cast(axis); - } - } - - Status Compute(OpKernelContext* context) const override; - - private: - int axis_; -}; -} // namespace onnxruntime diff --git a/onnxruntime/core/providers/cpu/math/softmax.cc b/onnxruntime/core/providers/cpu/math/softmax.cc index f3946c7c71..44817422ec 100644 --- a/onnxruntime/core/providers/cpu/math/softmax.cc +++ b/onnxruntime/core/providers/cpu/math/softmax.cc @@ -1,60 +1,29 @@ // Copyright (c) Microsoft Corporation. All rights reserved. // Licensed under the MIT License. - #include "core/providers/cpu/math/softmax.h" #include "core/framework/op_kernel.h" #include "core/framework/op_kernel_context_internal.h" #include "core/providers/common.h" -#include "core/providers/cpu/math/softmax_shared.h" #include "core/util/math.h" +#include "core/util/eigen_common_wrapper.h" namespace onnxruntime { -template <> -Status Softmax::Compute(OpKernelContext* ctx) const { - auto ctx_internal = static_cast(ctx); - concurrency::ThreadPool* tp = ctx_internal->GetOperatorThreadPool(); - - const auto* tensor_pointer = ctx->Input(0); - if (tensor_pointer == nullptr) return Status(common::ONNXRUNTIME, common::FAIL, "input count mismatch"); - const Tensor& X = *tensor_pointer; - const TensorShape& input_shape{X.Shape()}; - - VLOGS(ctx->Logger(), 2) << "Input tensor shape: " << input_shape; - - Tensor* Y = ctx->Output(0, input_shape); - - const int64_t axis = HandleNegativeAxis(axis_, input_shape.NumDimensions()); - - size_t N = input_shape.SizeToDimension(axis); - size_t D = input_shape.SizeFromDimension(axis); - - auto* Ydata = Y->template MutableData(); - - std::vector scale_(N); - std::vector rowmax_(N); - std::vector sum_multiplier_(D, 1.f); // initialize all multiplier values to 1.0 - - const bool logarithmic = false; - auto status = SoftmaxCPU(N, D, X.template Data(), Ydata, - scale_.data(), sum_multiplier_.data(), logarithmic, rowmax_.data(), tp); - - return status; -} - -ONNX_CPU_OPERATOR_VERSIONED_KERNEL( - Softmax, - 1, - 10, - KernelDefBuilder().TypeConstraint("T", DataTypeImpl::GetTensorType()), - Softmax); +ONNX_CPU_OPERATOR_VERSIONED_KERNEL(Softmax, 1, 10, + KernelDefBuilder().TypeConstraint("T", DataTypeImpl::GetTensorType()), + Softmax); // Opset 11 starts to support Neg Axis. -ONNX_CPU_OPERATOR_KERNEL( - Softmax, - 11, - KernelDefBuilder().TypeConstraint("T", DataTypeImpl::GetTensorType()), - Softmax); +ONNX_CPU_OPERATOR_KERNEL(Softmax, 11, KernelDefBuilder().TypeConstraint("T", DataTypeImpl::GetTensorType()), + Softmax); + +ONNX_CPU_OPERATOR_VERSIONED_KERNEL(LogSoftmax, 1, 10, + KernelDefBuilder().TypeConstraint("T", DataTypeImpl::GetTensorType()), + Softmax); + +// Opset 11 starts to support Neg Axis. +ONNX_CPU_OPERATOR_KERNEL(LogSoftmax, 11, KernelDefBuilder().TypeConstraint("T", DataTypeImpl::GetTensorType()), + Softmax); } // namespace onnxruntime diff --git a/onnxruntime/core/providers/cpu/math/softmax.h b/onnxruntime/core/providers/cpu/math/softmax.h index 2ac0a98f76..daac2eefbb 100644 --- a/onnxruntime/core/providers/cpu/math/softmax.h +++ b/onnxruntime/core/providers/cpu/math/softmax.h @@ -1,3 +1,15 @@ +/* Copyright 2015 The TensorFlow Authors. All Rights Reserved. +Licensed under the Apache License, Version 2.0 (the "License"); +you may not use this file except in compliance with the License. +You may obtain a copy of the License at + http://www.apache.org/licenses/LICENSE-2.0 +Unless required by applicable law or agreed to in writing, software +distributed under the License is distributed on an "AS IS" BASIS, +WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +See the License for the specific language governing permissions and +limitations under the License. +==============================================================================*/ + // Copyright (c) Microsoft Corporation. All rights reserved. // Licensed under the MIT License. @@ -7,10 +19,60 @@ #include "core/common/common.h" #include "core/framework/op_kernel.h" +#include "core/framework/op_kernel_context_internal.h" #include "core/util/math_cpuonly.h" +#include "core/util/eigen_common_wrapper.h" +#include "core/providers/common.h" namespace onnxruntime { -template +// copied from tensorflow/core/kernels/softmax_op.cc +template +static void ComputeSoftMax( + const Device& d, + typename Eigen::TensorMap, Eigen::Aligned> logits, + typename Eigen::TensorMap, Eigen::Aligned> softmax, + const bool log) { + const int kBatchDim = 0; + const int kClassDim = 1; + + const int batch_size = (int)logits.dimension(kBatchDim); + const int num_classes = (int)logits.dimension(kClassDim); + +// These arrays are used to reduce along the class dimension, and broadcast +// the resulting value to all classes. +#if !defined(EIGEN_HAS_INDEX_LIST) + Eigen::DSizes along_class(kClassDim); + Eigen::DSizes batch_by_one(batch_size, 1); + Eigen::DSizes one_by_class(1, num_classes); +#else + Eigen::IndexList > along_class; + Eigen::IndexList > batch_by_one; + batch_by_one.set(0, batch_size); + Eigen::IndexList, int> one_by_class; + one_by_class.set(1, num_classes); +#endif + // shifted_logits = logits - max(logits along classes); + auto shifted_logits = (logits - logits.maximum(along_class).eval().reshape(batch_by_one).broadcast(one_by_class)); + if (log) { + // Calculate the log of the softmax + // softmax = logits - max(logits along classes); + softmax.device(d) = shifted_logits; + // softmax = softmax - log(sum(exp(softmax along classes))); + softmax.device(d) = + (softmax - softmax.exp().sum(along_class).log().eval().reshape(batch_by_one).broadcast(one_by_class)); + } else { + // NOTE(touts): If you modify this implementation please run + // the BM_ImageNetSoftmaxFwd benchmark in nn_ops_test.cc. + // + // softmax = exp(logits - max(logits along classes)); + softmax.device(d) = shifted_logits.exp(); + // softmax = softmax * (1 / sum(softmax along classes)); + softmax.device(d) = + (softmax * softmax.sum(along_class).inverse().eval().reshape(batch_by_one).broadcast(one_by_class)); + } +} + +template class Softmax final : public OpKernel { public: Softmax(const OpKernelInfo& info) : OpKernel{info}, axis_{1} { @@ -22,7 +84,40 @@ class Softmax final : public OpKernel { } } - Status Compute(OpKernelContext* context) const override; + Status Compute(OpKernelContext* ctx) const override { +#ifndef USE_OPENMP + auto ctx_internal = static_cast(ctx); + concurrency::ThreadPool* tp = ctx_internal->GetOperatorThreadPool(); +#endif + const auto* tensor_pointer = ctx->Input(0); + if (tensor_pointer == nullptr) + return Status(common::ONNXRUNTIME, common::FAIL, "input count mismatch"); + const Tensor& X = *tensor_pointer; + const TensorShape& input_shape{X.Shape()}; + + VLOGS(ctx->Logger(), 2) << "Input tensor shape: " << input_shape; + + Tensor* Y = ctx->Output(0, input_shape); + + const int64_t axis = HandleNegativeAxis(axis_, input_shape.NumDimensions()); + + int N = static_cast(input_shape.SizeToDimension(axis)); + int D = static_cast(input_shape.SizeFromDimension(axis)); + + Eigen::TensorMap, Eigen::Aligned> X_tensor( + X.Data(), N, D); + Eigen::TensorMap, Eigen::Aligned> Y_tensor( + Y->MutableData(), N, D); +#ifndef USE_OPENMP + if (tp == nullptr) +#endif + ComputeSoftMax(Eigen::DefaultDevice(), X_tensor, Y_tensor, use_log); +#ifndef USE_OPENMP + else + ComputeSoftMax(Eigen::ThreadPoolDevice(&tp->GetHandler(), tp->NumThreads()), X_tensor, Y_tensor, use_log); +#endif + return Status::OK(); + } private: int axis_; diff --git a/onnxruntime/core/providers/cpu/math/softmax_shared.cc b/onnxruntime/core/providers/cpu/math/softmax_shared.cc deleted file mode 100644 index 18277f6b41..0000000000 --- a/onnxruntime/core/providers/cpu/math/softmax_shared.cc +++ /dev/null @@ -1,93 +0,0 @@ -/** -* Derived from caffe2, need copy right announcement here. -*/ - -/** -* Copyright (c) 2016-present, Facebook, Inc. -* -* Licensed under the Apache License, Version 2.0 (the "License"); -* you may not use this file except in compliance with the License. -* You may obtain a copy of the License at -* -* http://www.apache.org/licenses/LICENSE-2.0 -* -* Unless required by applicable law or agreed to in writing, software -* distributed under the License is distributed on an "AS IS" BASIS, -* WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -* See the License for the specific language governing permissions and -* limitations under the License. -*/ - -// disable noisy warning about use of std::copy_n by gsl::copy. gsl_algorithm does the same thing -// but that's too late for the disable to work -#ifdef _MSC_VER -#pragma warning(push) -#pragma warning(disable : 4996) -#endif -#include -#include -#ifdef _MSC_VER -#pragma warning(pop) -#endif - -#include "core/providers/cpu/math/softmax_shared.h" - -#include "core/util/math.h" -#include "core/util/math_cpuonly.h" - -#include "gsl/gsl_algorithm" -#include "gsl/gsl_util" - -namespace onnxruntime { - -common::Status SoftmaxCPU(const int64_t N, - const int64_t D, - const float* Xdata, - float* Ydata, - float* scale, - const float* sum_multiplier, - bool logarithmic, - float* rowmax, onnxruntime::concurrency::ThreadPool* tp) { - // the Math functions SoftmaxCPU uses only support int32_t as input, so enforce that - if (N * D > INT32_MAX || N > INT32_MAX || D > INT32_MAX) { - std::ostringstream ss; - ss << "SoftmaxCPU inputs N, D and N * D must be < " << INT32_MAX << ". N=" << N << ", D=" << D; - std::string msg = ss.str(); - - return Status(common::ONNXRUNTIME, common::INVALID_ARGUMENT, msg); - } - - const int n = gsl::narrow_cast(N); - const int d = gsl::narrow_cast(D); - const int nd = gsl::narrow_cast(N * D); - - math::RowwiseMax(n, d, Xdata, rowmax, nullptr); - - // Put the intermediate result X - max(X) into Y by first copying X to Y, and then subtracting max from each entry - gsl::copy(gsl::make_span(Xdata, nd), gsl::make_span(Ydata, nd)); - - math::Gemm(CblasNoTrans, CblasNoTrans, n, d, 1, -1, rowmax, sum_multiplier, 1, Ydata, tp); - - // Exponentiation - math::Exp(nd, Ydata, Ydata, nullptr); - math::Gemv(CblasNoTrans, n, d, 1, Ydata, sum_multiplier, 0, scale, nullptr); - - // Do division - if (!logarithmic) { - for (int i = 0; i < N; ++i) { - for (int j = 0; j < D; ++j) { - Ydata[i * D + j] /= scale[i]; - } - } - } else { - for (int i = 0; i < N; ++i) { - auto log_fmaxf_scale_i = std::log(fmaxf(scale[i], 1e-20f)); - for (int j = 0; j < D; ++j) { - Ydata[i * D + j] = Xdata[i * D + j] - rowmax[i] - log_fmaxf_scale_i; - } - } - } - - return Status::OK(); -} -} // namespace onnxruntime diff --git a/onnxruntime/core/providers/cpu/math/softmax_shared.h b/onnxruntime/core/providers/cpu/math/softmax_shared.h deleted file mode 100644 index 26ffeb193f..0000000000 --- a/onnxruntime/core/providers/cpu/math/softmax_shared.h +++ /dev/null @@ -1,25 +0,0 @@ -// Copyright (c) Microsoft Corporation. All rights reserved. -// Licensed under the MIT License. - -#pragma once - -#include "core/common/status.h" - -namespace onnxruntime { -namespace concurrency { -class ThreadPool; -} -/** -Calculate Softmax using CPU memory. -@param N Number of rows -@param D Number of elements in each row -@param Xdata Source data -@param Ydata Output data -@param scale Storage for scale calculation. Size must be >= N. -@param sum_multiplier Weights for each element. Size must be >= D. -@param logarithmic If true, compute LogSoftmax. If false compute Softmax. -@param rowmax Storage for calculation of maximum in each row. Size must be >= N. -*/ -common::Status SoftmaxCPU(int64_t N, int64_t D, const float* Xdata, float* Ydata, float* scale, - const float* sum_multiplier, bool logarithmic, float* rowmax, concurrency::ThreadPool* tp); -} // namespace onnxruntime diff --git a/onnxruntime/test/providers/cpu/math/hardmax_test.cc b/onnxruntime/test/providers/cpu/math/hardmax_test.cc index 1a0dfbd51d..44baf99860 100644 --- a/onnxruntime/test/providers/cpu/math/hardmax_test.cc +++ b/onnxruntime/test/providers/cpu/math/hardmax_test.cc @@ -1,7 +1,6 @@ // Copyright (c) Microsoft Corporation. All rights reserved. // Licensed under the MIT License. -#include "core/providers/cpu/math/softmax_shared.h" #include "gmock/gmock.h" #include "gtest/gtest.h" #include "test/providers/provider_test_utils.h" diff --git a/onnxruntime/test/providers/cpu/math/logsoftmax_test.cc b/onnxruntime/test/providers/cpu/math/logsoftmax_test.cc index 4482c7fcf7..32f44174c5 100644 --- a/onnxruntime/test/providers/cpu/math/logsoftmax_test.cc +++ b/onnxruntime/test/providers/cpu/math/logsoftmax_test.cc @@ -1,7 +1,7 @@ // Copyright (c) Microsoft Corporation. All rights reserved. // Licensed under the MIT License. -#include "core/providers/cpu/math/logsoftmax.h" +#include "core/providers/cpu/math/softmax.h" #include "gmock/gmock.h" #include "gtest/gtest.h" #include "test/providers/provider_test_utils.h" diff --git a/onnxruntime/test/providers/cpu/math/softmax_test.cc b/onnxruntime/test/providers/cpu/math/softmax_test.cc index 0273378194..48ea1bdbf1 100644 --- a/onnxruntime/test/providers/cpu/math/softmax_test.cc +++ b/onnxruntime/test/providers/cpu/math/softmax_test.cc @@ -1,7 +1,6 @@ // Copyright (c) Microsoft Corporation. All rights reserved. // Licensed under the MIT License. -#include "core/providers/cpu/math/softmax_shared.h" #include "gmock/gmock.h" #include "gtest/gtest.h" #include "test/providers/provider_test_utils.h" @@ -192,46 +191,5 @@ TEST(SoftmaxOperator, InvalidAxis) { "-10 is not in valid range [-2,1]"); } -TEST(SoftmaxOperator, TestInputTooLarge) { - float* ignored = nullptr; - concurrency::ThreadPool tp("", 1); - // N > INT32_MAX - int64_t N = int64_t(INT32_MAX) + 1; - int64_t D = 1; - auto status = SoftmaxCPU(N, D, ignored, ignored, ignored, ignored, true, ignored, &tp); - EXPECT_EQ(status.Code(), common::INVALID_ARGUMENT); - - // D > INT32_MAX - N = 1; - D = int64_t(INT32_MAX) + 1; - status = SoftmaxCPU(N, D, ignored, ignored, ignored, ignored, true, ignored, &tp); - EXPECT_EQ(status.Code(), common::INVALID_ARGUMENT); - - // N * D > INT32_MAX - N = int64_t(INT32_MAX) / 2; - D = 3; - status = SoftmaxCPU(N, D, ignored, ignored, ignored, ignored, true, ignored, &tp); - EXPECT_EQ(status.Code(), common::INVALID_ARGUMENT); - - /* - common::Status SoftmaxCPU(const int64_t N, - const int64_t D, - const float* Xdata, - float* Ydata, - float* scale, - const float* sum_multiplier, - bool logarithmic, - float* rowmax) - { - // the Math functions SoftmaxCPU uses only support int32_t as input, so enforce that - if (N * D > INT32_MAX || N > INT32_MAX || D > INT32_MAX) - { - std::ostringstream ss; - ss << "SoftmaxCPU inputs N, D and N * D must be < " << INT32_MAX << ". N=" << N << ", D=" << D; - std::string msg = ss.str(); - - return Status(common::ONNXRUNTIME, common::INVALID_ARGUMENT, msg); - }*/ -} } // namespace test } // namespace onnxruntime