Skip to content

Commit

Permalink
Merged PR 1135: Add Sqrt, Pow, Exp, and Log ops
Browse files Browse the repository at this point in the history 
Another set of simple ops

Plus change the ceil/floor implementations to use Eigen since I figured out a way

Related work items: #152
  • Loading branch information
@RyanUnderhill
RyanUnderhill committed Mar 29, 2018
1 parent d064446 commit b7f0c97
Show file tree
Hide file tree
Showing 3 changed files with 177 additions and 12 deletions.
88 changes: 76 additions & 12 deletions lotus/core/providers/cpu/math/element_wise_ops.cc
View file
Original file line number Diff line number Diff line change
Expand Up @@ -65,6 +65,34 @@ REGISTER_KERNEL(KernelDef("Reciprocal")
.TypeConstraint("T", DataTypeImpl::GetTensorType<float>()),
Reciprocal<float>);

REGISTER_KERNEL(KernelDef("Sqrt")
.Domain(LotusIR::kOnnxDomain)
.SinceVersion(1, 2)
.Provider(LotusIR::kCpuExecutionProvider)
.TypeConstraint("T", DataTypeImpl::GetTensorType<float>()),
Sqrt<float>);

REGISTER_KERNEL(KernelDef("Pow")
.Domain(LotusIR::kOnnxDomain)
.SinceVersion(1, 2)
.Provider(LotusIR::kCpuExecutionProvider)
.TypeConstraint("T", DataTypeImpl::GetTensorType<float>()),
Pow<float>);

REGISTER_KERNEL(KernelDef("Exp")
.Domain(LotusIR::kOnnxDomain)
.SinceVersion(1, 2)
.Provider(LotusIR::kCpuExecutionProvider)
.TypeConstraint("T", DataTypeImpl::GetTensorType<float>()),
Exp<float>);

REGISTER_KERNEL(KernelDef("Log")
.Domain(LotusIR::kOnnxDomain)
.SinceVersion(1, 2)
.Provider(LotusIR::kCpuExecutionProvider)
.TypeConstraint("T", DataTypeImpl::GetTensorType<float>()),
Log<float>);

REGISTER_KERNEL(KernelDef("Sum")
.Domain(LotusIR::kOnnxDomain)
.SinceVersion(1, 2)
Expand Down Expand Up @@ -228,12 +256,7 @@ Status Floor<float>::compute(OpKernelContext* ctx) const {
auto& X = *ctx->input<Tensor>(0);
auto& Y = *ctx->output(0, X.shape());

// There is no Eigen function for ceiling, so do it ourselves
auto* pInput = X.data<float>();
auto* pOutput = Y.mutable_data<float>();
size_t count = Y.shape().Size();
for (size_t i = 0; i < count; i++)
pOutput[i] = floor(pInput[i]);
EigenMap<float>(Y) = EigenMap<float>(X).array().floor();
return Status::OK();
}

Expand All @@ -242,12 +265,7 @@ Status Ceil<float>::compute(OpKernelContext* ctx) const {
auto& X = *ctx->input<Tensor>(0);
auto& Y = *ctx->output(0, X.shape());

// There is no Eigen function for ceiling, so do it ourselves
auto* pInput = X.data<float>();
auto* pOutput = Y.mutable_data<float>();
size_t count = Y.shape().Size();
for (size_t i = 0; i < count; i++)
pOutput[i] = ceil(pInput[i]);
EigenMap<float>(Y) = EigenMap<float>(X).array().ceil();
return Status::OK();
}

Expand All @@ -260,6 +278,52 @@ Status Reciprocal<float>::compute(OpKernelContext* ctx) const {
return Status::OK();
}

template <>
Status Sqrt<float>::compute(OpKernelContext* ctx) const {
auto& X = *ctx->input<Tensor>(0);
auto& Y = *ctx->output(0, X.shape());

EigenMap<float>(Y) = EigenMap<float>(X).cwiseSqrt();
return Status::OK();
}

template <>
Status Pow<float>::compute(OpKernelContext* ctx) const {
auto& A = *ctx->input<Tensor>(0);
auto& B = *ctx->input<Tensor>(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<float>(C) = EigenMap<float>(A).array().pow(*B.data<float>());
} else
Broadcast<float>(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<float>(C) = EigenMap<float>(A).array().pow(EigenMap<float>(B).array());
}
return Status::OK();
}

template <>
Status Exp<float>::compute(OpKernelContext* ctx) const {
auto& X = *ctx->input<Tensor>(0);
auto& Y = *ctx->output(0, X.shape());

EigenMap<float>(Y) = EigenMap<float>(X).array().exp();
return Status::OK();
}

template <>
Status Log<float>::compute(OpKernelContext* ctx) const {
auto& X = *ctx->input<Tensor>(0);
auto& Y = *ctx->output(0, X.shape());

EigenMap<float>(Y) = EigenMap<float>(X).array().log();
return Status::OK();
}

template <>
Status Sum<float>::compute(OpKernelContext* ctx) const {
auto inputCount = node().InputArgCount().front();
Expand Down
36 changes: 36 additions & 0 deletions lotus/core/providers/cpu/math/element_wise_ops.h
View file
Original file line number Diff line number Diff line change
Expand Up @@ -101,6 +101,42 @@ class Reciprocal final : public OpKernel {
Status compute(OpKernelContext* context) const override;
};

template <typename T>
class Sqrt final : public OpKernel {
public:
Sqrt(const OpKernelInfo& info) : OpKernel(info) {
}

Status compute(OpKernelContext* context) const override;
};

template <typename T>
class Pow final : public BroadcastAxisKernel {
public:
Pow(const OpKernelInfo& info) : BroadcastAxisKernel(info) {
}

Status compute(OpKernelContext* context) const override;
};

template <typename T>
class Exp final : public BroadcastAxisKernel {
public:
Exp(const OpKernelInfo& info) : BroadcastAxisKernel(info) {
}

Status compute(OpKernelContext* context) const override;
};

template <typename T>
class Log final : public BroadcastAxisKernel {
public:
Log(const OpKernelInfo& info) : BroadcastAxisKernel(info) {
}

Status compute(OpKernelContext* context) const override;
};

template <typename T>
class Sum final : public OpKernel {
public:
Expand Down
65 changes: 65 additions & 0 deletions lotus/test/providers/cpu/math/element_wise_ops_test.cc
View file
Original file line number Diff line number Diff line change
Expand Up @@ -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<Sqrt> test(model);

std::vector<int64_t> 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<Pow> test(model);

std::vector<int64_t> 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<Pow> test(model);

std::vector<int64_t> 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<Exp> test(model);

std::vector<int64_t> 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<Log> test(model);

std::vector<int64_t> 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});
Expand Down

0 comments on commit b7f0c97

Please sign in to comment.