From b7f0c9784c29693e93286172c86da61146b7ed89 Mon Sep 17 00:00:00 2001 From: Ryan Hill Date: Thu, 29 Mar 2018 00:29:18 +0000 Subject: [PATCH] Merged PR 1135: Add Sqrt, Pow, Exp, and Log ops Another set of simple ops Plus change the ceil/floor implementations to use Eigen since I figured out a way Related work items: #152 --- .../providers/cpu/math/element_wise_ops.cc | 88 ++++++++++++++++--- .../providers/cpu/math/element_wise_ops.h | 36 ++++++++ .../cpu/math/element_wise_ops_test.cc | 65 ++++++++++++++ 3 files changed, 177 insertions(+), 12 deletions(-) diff --git a/lotus/core/providers/cpu/math/element_wise_ops.cc b/lotus/core/providers/cpu/math/element_wise_ops.cc index 25f25d2c57513..9f9acb67bb23a 100644 --- a/lotus/core/providers/cpu/math/element_wise_ops.cc +++ b/lotus/core/providers/cpu/math/element_wise_ops.cc @@ -65,6 +65,34 @@ REGISTER_KERNEL(KernelDef("Reciprocal") .TypeConstraint("T", DataTypeImpl::GetTensorType()), Reciprocal); +REGISTER_KERNEL(KernelDef("Sqrt") + .Domain(LotusIR::kOnnxDomain) + .SinceVersion(1, 2) + .Provider(LotusIR::kCpuExecutionProvider) + .TypeConstraint("T", DataTypeImpl::GetTensorType()), + Sqrt); + +REGISTER_KERNEL(KernelDef("Pow") + .Domain(LotusIR::kOnnxDomain) + .SinceVersion(1, 2) + .Provider(LotusIR::kCpuExecutionProvider) + .TypeConstraint("T", DataTypeImpl::GetTensorType()), + Pow); + +REGISTER_KERNEL(KernelDef("Exp") + .Domain(LotusIR::kOnnxDomain) + .SinceVersion(1, 2) + .Provider(LotusIR::kCpuExecutionProvider) + .TypeConstraint("T", DataTypeImpl::GetTensorType()), + Exp); + +REGISTER_KERNEL(KernelDef("Log") + .Domain(LotusIR::kOnnxDomain) + .SinceVersion(1, 2) + .Provider(LotusIR::kCpuExecutionProvider) + .TypeConstraint("T", DataTypeImpl::GetTensorType()), + Log); + REGISTER_KERNEL(KernelDef("Sum") .Domain(LotusIR::kOnnxDomain) .SinceVersion(1, 2) @@ -228,12 +256,7 @@ Status Floor::compute(OpKernelContext* ctx) const { auto& X = *ctx->input(0); auto& Y = *ctx->output(0, X.shape()); - // There is no Eigen function for ceiling, so do it ourselves - auto* pInput = X.data(); - auto* pOutput = Y.mutable_data(); - size_t count = Y.shape().Size(); - for (size_t i = 0; i < count; i++) - pOutput[i] = floor(pInput[i]); + EigenMap(Y) = EigenMap(X).array().floor(); return Status::OK(); } @@ -242,12 +265,7 @@ Status Ceil::compute(OpKernelContext* ctx) const { auto& X = *ctx->input(0); auto& Y = *ctx->output(0, X.shape()); - // There is no Eigen function for ceiling, so do it ourselves - auto* pInput = X.data(); - auto* pOutput = Y.mutable_data(); - size_t count = Y.shape().Size(); - for (size_t i = 0; i < count; i++) - pOutput[i] = ceil(pInput[i]); + EigenMap(Y) = EigenMap(X).array().ceil(); return Status::OK(); } @@ -260,6 +278,52 @@ Status Reciprocal::compute(OpKernelContext* ctx) const { return Status::OK(); } +template <> +Status Sqrt::compute(OpKernelContext* ctx) const { + auto& X = *ctx->input(0); + auto& Y = *ctx->output(0, X.shape()); + + EigenMap(Y) = EigenMap(X).cwiseSqrt(); + return Status::OK(); +} + +template <> +Status Pow::compute(OpKernelContext* ctx) const { + auto& A = *ctx->input(0); + auto& B = *ctx->input(1); + auto& C = *ctx->output(0, A.shape()); + + if (broadcast_) { + if (B.shape().NumDimensions() == 0) { + LOTUS_ENFORCE(axis_ == -1, "When broadcasting by a scalar, axis cannot be set"); + EigenMap(C) = EigenMap(A).array().pow(*B.data()); + } else + Broadcast(A, B, C, int(axis_), [](float a, float b) { return pow(a, b); }); + } else { + LOTUS_ENFORCE(A.shape() == B.shape(), "Inputs must have the same shape"); + EigenMap(C) = EigenMap(A).array().pow(EigenMap(B).array()); + } + return Status::OK(); +} + +template <> +Status Exp::compute(OpKernelContext* ctx) const { + auto& X = *ctx->input(0); + auto& Y = *ctx->output(0, X.shape()); + + EigenMap(Y) = EigenMap(X).array().exp(); + return Status::OK(); +} + +template <> +Status Log::compute(OpKernelContext* ctx) const { + auto& X = *ctx->input(0); + auto& Y = *ctx->output(0, X.shape()); + + EigenMap(Y) = EigenMap(X).array().log(); + return Status::OK(); +} + template <> Status Sum::compute(OpKernelContext* ctx) const { auto inputCount = node().InputArgCount().front(); diff --git a/lotus/core/providers/cpu/math/element_wise_ops.h b/lotus/core/providers/cpu/math/element_wise_ops.h index f3ca0c52a9969..256ce0cc6efeb 100644 --- a/lotus/core/providers/cpu/math/element_wise_ops.h +++ b/lotus/core/providers/cpu/math/element_wise_ops.h @@ -101,6 +101,42 @@ class Reciprocal final : public OpKernel { Status compute(OpKernelContext* context) const override; }; +template +class Sqrt final : public OpKernel { + public: + Sqrt(const OpKernelInfo& info) : OpKernel(info) { + } + + Status compute(OpKernelContext* context) const override; +}; + +template +class Pow final : public BroadcastAxisKernel { + public: + Pow(const OpKernelInfo& info) : BroadcastAxisKernel(info) { + } + + Status compute(OpKernelContext* context) const override; +}; + +template +class Exp final : public BroadcastAxisKernel { + public: + Exp(const OpKernelInfo& info) : BroadcastAxisKernel(info) { + } + + Status compute(OpKernelContext* context) const override; +}; + +template +class Log final : public BroadcastAxisKernel { + public: + Log(const OpKernelInfo& info) : BroadcastAxisKernel(info) { + } + + Status compute(OpKernelContext* context) const override; +}; + template class Sum final : public OpKernel { public: diff --git a/lotus/test/providers/cpu/math/element_wise_ops_test.cc b/lotus/test/providers/cpu/math/element_wise_ops_test.cc index d3d14e6da32aa..35ca37de734e3 100644 --- a/lotus/test/providers/cpu/math/element_wise_ops_test.cc +++ b/lotus/test/providers/cpu/math/element_wise_ops_test.cc @@ -167,6 +167,71 @@ TEST(MathOpTest, Reciprocal) { test.Run(dims, expected_vals); } +TEST(MathOpTest, Sqrt) { + LotusIR::NodeArg input_def("X", &s_typeProto_float), output_def("Y", &s_typeProto_float); + TestModel model("Sqrt", {&input_def}, {&output_def}); + SimpleFloatTest test(model); + + std::vector dims{2, 2}; + test.AddInput(dims, {1.0f, 4.0f, 0.0f, 9.0f}); + test.AddOutput(dims); + float expected_vals[]{1.0f, 2.0f, 0.0f, 3.0f}; + test.Run(dims, expected_vals); +} + +TEST(MathOpTest, Pow) { + LotusIR::NodeArg input1_def("X", &s_typeProto_float), input2_def("Y", &s_typeProto_float), output_def("Z", &s_typeProto_float); + TestModel model("Pow", {&input1_def, &input2_def}, {&output_def}); + SimpleFloatTest test(model); + + std::vector dims{2, 2}; + test.AddInput(dims, {2.0f, 2.0f, sqrt(2.0f), 1.0f}); + test.AddInput(dims, {0.0f, 8.0f, 2.0f, 9.0f}); + test.AddOutput(dims); + float expected_vals[]{1.0f, 256.0f, 2.0f, 1.0f}; + test.Run(dims, expected_vals); +} + +TEST(MathOpTest, Pow_Broadcast_Scalar) { + LotusIR::NodeArg input1_def("X", &s_typeProto_float), input2_def("Y", &s_typeProto_float), output_def("Z", &s_typeProto_float); + TestModel model("Pow", {&input1_def, &input2_def}, {&output_def}); + + EXPECT_TRUE(model.Node().AddAttribute("broadcast", int64_t{1})); + + SimpleFloatTest test(model); + + std::vector dims{3}; + test.AddInput(dims, {1.0f, 2.0f, 3.0f}); + test.AddInput({}, {2.0f}); + test.AddOutput(dims); + float expected_vals[]{1.0f, 4.0f, 9.0f}; + test.Run(dims, expected_vals); +} + +TEST(MathOpTest, Exp) { + LotusIR::NodeArg input_def("X", &s_typeProto_float), output_def("Y", &s_typeProto_float); + TestModel model("Exp", {&input_def}, {&output_def}); + SimpleFloatTest test(model); + + std::vector dims{2, 2}; + test.AddInput(dims, {0.0f, 1.0f, 2.0f, 10.0f}); + test.AddOutput(dims); + float expected_vals[]{1.0f, exp(1.0f), exp(2.0f), exp(10.0f)}; + test.Run(dims, expected_vals); +} + +TEST(MathOpTest, Log) { + LotusIR::NodeArg input_def("X", &s_typeProto_float), output_def("Y", &s_typeProto_float); + TestModel model("Log", {&input_def}, {&output_def}); + SimpleFloatTest test(model); + + std::vector dims{2, 2}; + test.AddInput(dims, {1.0f, 2.0f, 5.0f, 10.0f}); + test.AddOutput(dims); + float expected_vals[]{0.0f, log(2.0f), log(5.0f), log(10.0f)}; + test.Run(dims, expected_vals); +} + TEST(MathOpTest, Sum) { LotusIR::NodeArg input1_def("data_0", &s_typeProto_float), input2_def("data_1", &s_typeProto_float), input3_def("data_3", &s_typeProto_float), output_def("sum", &s_typeProto_float); TestModel model("Sum", {&input1_def, &input2_def, &input3_def}, {&output_def});