未验证 提交 848ae7dc 编写于 作者: H hong 提交者: GitHub

Move digamma to pten (#39240)

* move digamma to pten; test=develop

* fix mutable_data bugs; test=develop

* remove useless code; test=develop

* remove kernel compute; test=develop

* fix bug; test=develop
上级 91dd0f0d
...@@ -64,6 +64,13 @@ class DigammaGradOp : public framework::OperatorWithKernel { ...@@ -64,6 +64,13 @@ class DigammaGradOp : public framework::OperatorWithKernel {
ctx->SetOutputDim(framework::GradVarName("X"), dout_dims); ctx->SetOutputDim(framework::GradVarName("X"), dout_dims);
ctx->ShareLoD(framework::GradVarName("Out"), framework::GradVarName("X")); ctx->ShareLoD(framework::GradVarName("Out"), framework::GradVarName("X"));
} }
framework::KernelSignature GetExpectedPtenKernelArgs(
const framework::ExecutionContext &ctx) const override {
return framework::KernelSignature("digamma_grad",
{framework::GradVarName("Out"), "X"}, {},
{framework::GradVarName("X")});
}
}; };
template <typename T> template <typename T>
...@@ -89,12 +96,3 @@ REGISTER_OPERATOR(digamma, ops::DigammaOp, ops::DigammaOpMaker, ...@@ -89,12 +96,3 @@ REGISTER_OPERATOR(digamma, ops::DigammaOp, ops::DigammaOpMaker,
ops::DigammaGradOpMaker<paddle::framework::OpDesc>, ops::DigammaGradOpMaker<paddle::framework::OpDesc>,
ops::DigammaGradOpMaker<paddle::imperative::OpBase>); ops::DigammaGradOpMaker<paddle::imperative::OpBase>);
REGISTER_OPERATOR(digamma_grad, ops::DigammaGradOp); REGISTER_OPERATOR(digamma_grad, ops::DigammaGradOp);
REGISTER_OP_CPU_KERNEL(
digamma, ops::DigammaKernel<paddle::platform::CPUDeviceContext, float>,
ops::DigammaKernel<paddle::platform::CPUDeviceContext, double>);
REGISTER_OP_CPU_KERNEL(
digamma_grad,
ops::DigammaGradKernel<paddle::platform::CPUDeviceContext, float>,
ops::DigammaGradKernel<paddle::platform::CPUDeviceContext, double>);
/* Copyright (c) 2021 PaddlePaddle 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. */
#include "paddle/fluid/operators/digamma_op.h"
namespace ops = paddle::operators;
REGISTER_OP_CUDA_KERNEL(
digamma, ops::DigammaKernel<paddle::platform::CUDADeviceContext, float>,
ops::DigammaKernel<paddle::platform::CUDADeviceContext, double>);
REGISTER_OP_CUDA_KERNEL(
digamma_grad,
ops::DigammaGradKernel<paddle::platform::CUDADeviceContext, float>,
ops::DigammaGradKernel<paddle::platform::CUDADeviceContext, double>);
...@@ -14,86 +14,5 @@ limitations under the License. */ ...@@ -14,86 +14,5 @@ limitations under the License. */
#pragma once #pragma once
#include <unsupported/Eigen/SpecialFunctions>
#include "paddle/fluid/framework/op_registry.h" #include "paddle/fluid/framework/op_registry.h"
#include "paddle/fluid/framework/operator.h" #include "paddle/fluid/framework/operator.h"
#include "paddle/fluid/platform/for_range.h"
namespace paddle {
namespace operators {
template <typename T>
struct DigammaFunctor {
DigammaFunctor(const T* input, T* output, int64_t numel)
: input_(input), output_(output), numel_(numel) {}
HOSTDEVICE void operator()(int64_t idx) const {
output_[idx] = Eigen::numext::digamma(input_[idx]);
}
private:
const T* input_;
T* output_;
int64_t numel_;
};
template <typename T>
struct DigammaGradFunctor {
DigammaGradFunctor(const T* dout, const T* x, T* output, int64_t numel)
: dout_(dout), x_(x), output_(output), numel_(numel) {}
HOSTDEVICE void operator()(int64_t idx) const {
output_[idx] = dout_[idx] * Eigen::numext::polygamma(T(1), x_[idx]);
}
private:
const T* dout_;
const T* x_;
T* output_;
int64_t numel_;
};
using Tensor = framework::Tensor;
template <typename DeviceContext, typename T>
class DigammaKernel : public framework::OpKernel<T> {
public:
void Compute(const framework::ExecutionContext& context) const override {
const Tensor* x = context.Input<Tensor>("X");
Tensor* out = context.Output<Tensor>("Out");
auto numel = x->numel();
auto* x_data = x->data<T>();
auto* out_data = out->mutable_data<T>(context.GetPlace(),
size_t(x->numel() * sizeof(T)));
auto& dev_ctx = context.template device_context<DeviceContext>();
platform::ForRange<DeviceContext> for_range(dev_ctx, numel);
DigammaFunctor<T> functor(x_data, out_data, numel);
for_range(functor);
}
};
template <typename DeviceContext, typename T>
class DigammaGradKernel : public framework::OpKernel<T> {
public:
void Compute(const framework::ExecutionContext& context) const override {
const Tensor* d_out = context.Input<Tensor>(framework::GradVarName("Out"));
const Tensor* x = context.Input<Tensor>("X");
auto* d_x = context.Output<Tensor>(framework::GradVarName("X"));
auto numel = d_out->numel();
auto* dout_data = d_out->data<T>();
auto* x_data = x->data<T>();
auto* dx_data = d_x->mutable_data<T>(
context.GetPlace(), static_cast<size_t>(numel * sizeof(T)));
auto& dev_ctx = context.template device_context<DeviceContext>();
platform::ForRange<DeviceContext> for_range(dev_ctx, numel);
DigammaGradFunctor<T> functor(dout_data, x_data, dx_data, numel);
for_range(functor);
}
};
} // namespace operators
} // namespace paddle
// Copyright (c) 2022 PaddlePaddle 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.
#include "paddle/pten/kernels/digamma_grad_kernel.h"
#include "paddle/pten/backends/cpu/cpu_context.h"
#include "paddle/pten/common/scalar.h"
#include "paddle/pten/core/dense_tensor.h"
#include "paddle/pten/core/kernel_registry.h"
#include "paddle/pten/kernels/impl/digamma_grad_kernel_impl.h"
PT_REGISTER_KERNEL(
digamma_grad, CPU, ALL_LAYOUT, pten::DigammaGradKernel, float, double) {}
// Copyright (c) 2022 PaddlePaddle 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.
#include "paddle/pten/kernels/digamma_kernel.h"
#include "paddle/pten/backends/cpu/cpu_context.h"
#include "paddle/pten/common/scalar.h"
#include "paddle/pten/core/dense_tensor.h"
#include "paddle/pten/core/kernel_registry.h"
#include "paddle/pten/kernels/impl/digamma_kernel_impl.h"
PT_REGISTER_KERNEL(
digamma, CPU, ALL_LAYOUT, pten::DigammaKernel, float, double) {}
// Copyright (c) 2022 PaddlePaddle 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.
#pragma once
#include "paddle/pten/core/dense_tensor.h"
namespace pten {
template <typename T, typename Context>
void DigammaGradKernel(const Context& ctx,
const DenseTensor& out_grad,
const DenseTensor& x,
DenseTensor* x_grad);
} // namepsace pten
// Copyright (c) 2022 PaddlePaddle 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.
#pragma once
#include "paddle/pten/core/dense_tensor.h"
namespace pten {
template <typename T, typename Context>
void DigammaKernel(const Context& ctx, const DenseTensor& x, DenseTensor* out);
} // namepsace pten
// Copyright (c) 2022 PaddlePaddle 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.
#include "paddle/pten/backends/gpu/gpu_context.h"
#include "paddle/pten/core/dense_tensor.h"
#include "paddle/pten/core/kernel_registry.h"
#include "paddle/pten/kernels/digamma_grad_kernel.h"
#include "paddle/pten/kernels/impl/digamma_grad_kernel_impl.h"
PT_REGISTER_KERNEL(
digamma_grad, GPU, ALL_LAYOUT, pten::DigammaGradKernel, float, double) {}
// Copyright (c) 2022 PaddlePaddle 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.
#include "paddle/pten/backends/gpu/gpu_context.h"
#include "paddle/pten/common/scalar.h"
#include "paddle/pten/core/dense_tensor.h"
#include "paddle/pten/core/kernel_registry.h"
#include "paddle/pten/kernels/digamma_kernel.h"
#include "paddle/pten/kernels/impl/digamma_kernel_impl.h"
PT_REGISTER_KERNEL(
digamma, GPU, ALL_LAYOUT, pten::DigammaKernel, float, double) {}
// Copyright (c) 2022 PaddlePaddle 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.
#pragma once
#include <unsupported/Eigen/SpecialFunctions>
#include "paddle/fluid/platform/for_range.h"
#include "paddle/pten/core/dense_tensor.h"
namespace pten {
template <typename T>
struct DigammaGradFunctor {
DigammaGradFunctor(const T* dout, const T* x, T* output, int64_t numel)
: dout_(dout), x_(x), output_(output), numel_(numel) {}
HOSTDEVICE void operator()(int64_t idx) const {
output_[idx] = dout_[idx] * Eigen::numext::polygamma(T(1), x_[idx]);
}
private:
const T* dout_;
const T* x_;
T* output_;
int64_t numel_;
};
template <typename T, typename Context>
void DigammaGradKernel(const Context& ctx,
const DenseTensor& out_grad,
const DenseTensor& x,
DenseTensor* x_grad) {
x_grad->mutable_data<T>(ctx.GetPlace());
auto* dout_data = out_grad.data<T>();
auto* x_data = x.data<T>();
auto* dx_data = x_grad->data<T>();
auto numel = out_grad.numel();
platform::ForRange<Context> for_range(ctx, numel);
DigammaGradFunctor<T> functor(dout_data, x_data, dx_data, numel);
for_range(functor);
}
} // namespace pten
// Copyright (c) 2022 PaddlePaddle 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.
#pragma once
#include <unsupported/Eigen/SpecialFunctions>
#include "paddle/fluid/platform/for_range.h"
#include "paddle/pten/core/dense_tensor.h"
namespace pten {
template <typename T>
struct DigammaFunctor {
DigammaFunctor(const T* input, T* output, int64_t numel)
: input_(input), output_(output), numel_(numel) {}
HOSTDEVICE void operator()(int64_t idx) const {
output_[idx] = Eigen::numext::digamma(input_[idx]);
}
private:
const T* input_;
T* output_;
int64_t numel_;
};
template <typename T, typename Context>
void DigammaKernel(const Context& ctx, const DenseTensor& x, DenseTensor* out) {
out->mutable_data<T>(ctx.GetPlace());
auto* x_data = x.data<T>();
auto* out_data = out->data<T>();
auto numel = x.numel();
platform::ForRange<Context> for_range(ctx, numel);
DigammaFunctor<T> functor(x_data, out_data, numel);
for_range(functor);
}
} // namespace pten
Markdown is supported
0% .
You are about to add 0 people to the discussion. Proceed with caution.
先完成此消息的编辑!
想要评论请 注册