Skip to content

Commit 07d0b83

Browse files
Zjq9409Zjq9409
authored
Add CPU and GPU eigh op implementation (#34990)
1 parent 7d9ca16 commit 07d0b83

File tree

13 files changed

+995
-4
lines changed

13 files changed

+995
-4
lines changed

‎cmake/operators.cmake‎

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -185,6 +185,7 @@ function(op_library TARGET)
185185
list(REMOVE_ITEM hip_srcs "cholesky_op.cu")
186186
list(REMOVE_ITEM hip_srcs "matrix_rank_op.cu")
187187
list(REMOVE_ITEM hip_srcs "svd_op.cu")
188+
list(REMOVE_ITEM hip_srcs "eigh_op.cu")
188189
list(REMOVE_ITEM hip_srcs "multinomial_op.cu")
189190
list(REMOVE_ITEM hip_srcs "decode_jpeg_op.cu")
190191
hip_library(${TARGET} SRCS ${cc_srcs} ${hip_cc_srcs} ${miopen_cu_cc_srcs} ${miopen_cu_srcs} ${mkldnn_cc_srcs} ${hip_srcs} DEPS ${op_library_DEPS}

‎paddle/fluid/operators/eigh_op.cc‎

Lines changed: 167 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,167 @@
1+
/* Copyright (c) 2021 PaddlePaddle Authors. All Rights Reserved.
2+
3+
Licensed under the Apache License, Version 2.0 (the "License");
4+
you may not use this file except in compliance with the License.
5+
You may obtain a copy of the License at
6+
7+
http://www.apache.org/licenses/LICENSE-2.0
8+
9+
Unless required by applicable law or agreed to in writing, software
10+
distributed under the License is distributed on an "AS IS" BASIS,
11+
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12+
See the License for the specific language governing permissions and
13+
limitations under the License. */
14+
15+
#include "paddle/fluid/operators/eigh_op.h"
16+
17+
namespace paddle {
18+
namespace operators {
19+
20+
using framework::Tensor;
21+
22+
class EighOp : public framework::OperatorWithKernel {
23+
public:
24+
using framework::OperatorWithKernel::OperatorWithKernel;
25+
26+
void InferShape(framework::InferShapeContext* ctx) const override {
27+
OP_INOUT_CHECK(ctx->HasInput("X"), "Input", "X", "Eigh");
28+
OP_INOUT_CHECK(ctx->HasOutput("Eigenvalues"), "Output", "Eigenvalues",
29+
"Eigh");
30+
OP_INOUT_CHECK(ctx->HasOutput("Eigenvectors"), "Output", "Eigenvectors",
31+
"Eigh");
32+
33+
auto input_dim = ctx->GetInputDim("X");
34+
auto rank = input_dim.size();
35+
36+
PADDLE_ENFORCE_GE(rank, 2,
37+
platform::errors::InvalidArgument(
38+
"The Input(X) should have at least 2 dimensions."
39+
"But received a %d dimension tensor.",
40+
rank));
41+
PADDLE_ENFORCE_EQ(
42+
input_dim[rank - 2], input_dim[rank - 1],
43+
platform::errors::InvalidArgument(
44+
"Eigh op is designed for square matrix, consequently"
45+
"inner-most 2 dimensions of Input(X) should be symmetric."
46+
"But received X's shape[-2] = %d and shape[-1] = %d.",
47+
input_dim[rank - 2], input_dim[rank - 1]));
48+
49+
std::vector<int64_t> values_dim;
50+
if (rank > 2) {
51+
for (auto i = 0; i < rank - 1; i++) {
52+
values_dim.emplace_back(input_dim[i]);
53+
}
54+
} else {
55+
values_dim = {input_dim[1]};
56+
}
57+
58+
ctx->SetOutputDim("Eigenvalues", framework::make_ddim(values_dim));
59+
ctx->SetOutputDim("Eigenvectors", input_dim);
60+
}
61+
};
62+
63+
class EignOpMaker : public framework::OpProtoAndCheckerMaker {
64+
public:
65+
void Make() override {
66+
AddInput("X",
67+
"(Tensor), Hermitian or real symmetric matrices."
68+
"Its shape should be [*, N, N] where * is zero or"
69+
"more batch dimensions. The data type is float32 ,"
70+
"float64, complex64, complex128.");
71+
AddOutput("Eigenvalues",
72+
"(Tensor), The eigenvalues in ascending order."
73+
"The data type is float32 or float64.");
74+
AddOutput(
75+
"Eigenvectors",
76+
"(Tensor), The column is the normalized eigenvector "
77+
"corresponding to the eigenvalue. The data type is the same as ``X``.");
78+
AddAttr<std::string>(
79+
"UPLO",
80+
"(string, default 'L'), 'L' represents the lower triangular matrix,"
81+
"'U' represents the upper triangular matrix.")
82+
.SetDefault("L");
83+
AddComment(R"DOC(
84+
Eigh Operator.
85+
86+
Computes the eigenvalues and eigenvectors of a complex Hermitian
87+
(conjugate symmetric) or a real symmetric matrix.
88+
89+
)DOC");
90+
}
91+
};
92+
93+
class EighGradOp : public framework::OperatorWithKernel {
94+
public:
95+
using framework::OperatorWithKernel::OperatorWithKernel;
96+
97+
void InferShape(framework::InferShapeContext* ctx) const override {
98+
OP_INOUT_CHECK(ctx->HasInput("Eigenvalues"), "Input", "Eigenvalues",
99+
"EighGrad");
100+
OP_INOUT_CHECK(ctx->HasInput("Eigenvectors"), "Input", "Eigenvectors",
101+
"EighGrad");
102+
OP_INOUT_CHECK(ctx->HasInputs(framework::GradVarName("Eigenvalues")),
103+
"Input", "Eigenvalues@GRAD", "EighGrad");
104+
OP_INOUT_CHECK(ctx->HasInputs(framework::GradVarName("Eigenvectors")),
105+
"Input", "Eigenvectors@GRAD", "EighGrad");
106+
auto dims = ctx->GetInputDim("Eigenvectors");
107+
auto x_grad_name = framework::GradVarName("X");
108+
if (ctx->HasOutput(x_grad_name)) {
109+
ctx->SetOutputDim(x_grad_name, dims);
110+
}
111+
}
112+
113+
protected:
114+
framework::OpKernelType GetExpectedKernelType(
115+
const framework::ExecutionContext& ctx) const override {
116+
return framework::OpKernelType(
117+
OperatorWithKernel::IndicateVarDataType(
118+
ctx, framework::GradVarName("Eigenvectors")),
119+
ctx.device_context());
120+
}
121+
};
122+
123+
template <typename T>
124+
class EighGradOpMaker : public framework::SingleGradOpMaker<T> {
125+
public:
126+
using framework::SingleGradOpMaker<T>::SingleGradOpMaker;
127+
128+
protected:
129+
void Apply(GradOpPtr<T> op) const override {
130+
op->SetType(this->ForwardOpType() + "_grad");
131+
op->SetInput("Eigenvalues", this->Output("Eigenvalues"));
132+
op->SetInput("Eigenvectors", this->Output("Eigenvectors"));
133+
op->SetInput(framework::GradVarName("Eigenvalues"),
134+
this->OutputGrad("Eigenvalues"));
135+
op->SetInput(framework::GradVarName("Eigenvectors"),
136+
this->OutputGrad("Eigenvectors"));
137+
op->SetAttrMap(this->Attrs());
138+
op->SetOutput(framework::GradVarName("X"), this->InputGrad("X"));
139+
}
140+
};
141+
142+
} // namespace operators
143+
} // namespace paddle
144+
145+
namespace ops = paddle::operators;
146+
147+
REGISTER_OPERATOR(eigh, ops::EighOp, ops::EignOpMaker,
148+
ops::EighGradOpMaker<paddle::framework::OpDesc>,
149+
ops::EighGradOpMaker<paddle::imperative::OpBase>);
150+
REGISTER_OPERATOR(eigh_grad, ops::EighGradOp);
151+
152+
REGISTER_OP_CPU_KERNEL(
153+
eigh, ops::EighKernel<paddle::platform::CPUDeviceContext, float, float>,
154+
ops::EighKernel<paddle::platform::CPUDeviceContext, double, double>,
155+
ops::EighKernel<paddle::platform::CPUDeviceContext, float,
156+
paddle::platform::complex<float>>,
157+
ops::EighKernel<paddle::platform::CPUDeviceContext, double,
158+
paddle::platform::complex<double>>);
159+
160+
REGISTER_OP_CPU_KERNEL(
161+
eigh_grad,
162+
ops::EighGradKernel<paddle::platform::CPUDeviceContext, float, float>,
163+
ops::EighGradKernel<paddle::platform::CPUDeviceContext, double, double>,
164+
ops::EighGradKernel<paddle::platform::CPUDeviceContext, float,
165+
paddle::platform::complex<float>>,
166+
ops::EighGradKernel<paddle::platform::CPUDeviceContext, double,
167+
paddle::platform::complex<double>>);

‎paddle/fluid/operators/eigh_op.cu‎

Lines changed: 53 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,53 @@
1+
/* Copyright (c) 2021 PaddlePaddle Authors. All Rights Reserved.
2+
3+
Licensed under the Apache License, Version 2.0 (the "License");
4+
you may not use this file except in compliance with the License.
5+
You may obtain a copy of the License at
6+
7+
http://www.apache.org/licenses/LICENSE-2.0
8+
9+
Unless required by applicable law or agreed to in writing, software
10+
distributed under the License is distributed on an "AS IS" BASIS,
11+
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12+
See the License for the specific language governing permissions and
13+
limitations under the License. */
14+
15+
#include "paddle/fluid/operators/eigh_op.h"
16+
17+
namespace paddle {
18+
namespace operators {
19+
20+
using Tensor = framework::Tensor;
21+
22+
template <typename ValueType, typename T>
23+
class EighGPUKernel : public framework::OpKernel<T> {
24+
public:
25+
void Compute(const framework::ExecutionContext &ctx) const override {
26+
auto input_var = ctx.Input<Tensor>("X");
27+
auto output_w_var = ctx.Output<Tensor>("Eigenvalues");
28+
auto output_v_var = ctx.Output<Tensor>("Eigenvectors");
29+
std::string lower = ctx.Attr<std::string>("UPLO");
30+
bool is_lower = (lower == "L");
31+
math::MatrixEighFunctor<ValueType, T> functor;
32+
functor(ctx, *input_var, output_w_var, output_v_var, is_lower, true);
33+
}
34+
};
35+
36+
} // namespace operators
37+
} // namespace paddle
38+
39+
namespace ops = paddle::operators;
40+
41+
REGISTER_OP_CUDA_KERNEL(
42+
eigh, ops::EighGPUKernel<float, float>, ops::EighGPUKernel<double, double>,
43+
ops::EighGPUKernel<float, paddle::platform::complex<float>>,
44+
ops::EighGPUKernel<double, paddle::platform::complex<double>>);
45+
46+
REGISTER_OP_CUDA_KERNEL(
47+
eigh_grad,
48+
ops::EighGradKernel<paddle::platform::CUDADeviceContext, float, float>,
49+
ops::EighGradKernel<paddle::platform::CUDADeviceContext, double, double>,
50+
ops::EighGradKernel<paddle::platform::CUDADeviceContext, float,
51+
paddle::platform::complex<float>>,
52+
ops::EighGradKernel<paddle::platform::CUDADeviceContext, double,
53+
paddle::platform::complex<double>>);

‎paddle/fluid/operators/eigh_op.h‎

Lines changed: 80 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,80 @@
1+
// Copyright (c) 2021 PaddlePaddle Authors. All Rights Reserved.
2+
//
3+
// Licensed under the Apache License, Version 2.0 (the "License");
4+
// you may not use this file except in compliance with the License.
5+
// You may obtain a copy of the License at
6+
//
7+
// http://www.apache.org/licenses/LICENSE-2.0
8+
//
9+
// Unless required by applicable law or agreed to in writing, software
10+
// distributed under the License is distributed on an "AS IS" BASIS,
11+
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12+
// See the License for the specific language governing permissions and
13+
// limitations under the License.
14+
15+
#pragma once
16+
17+
#include "paddle/fluid/framework/op_registry.h"
18+
#include "paddle/fluid/operators/math/eigen_values_vectors.h"
19+
20+
namespace paddle {
21+
namespace operators {
22+
23+
using Tensor = framework::Tensor;
24+
25+
template <typename T, size_t D, int MajorType = Eigen::RowMajor,
26+
typename IndexType = Eigen::DenseIndex>
27+
using EigenTensor = framework::EigenTensor<T, D, MajorType, IndexType>;
28+
template <typename T, int MajorType = Eigen::RowMajor,
29+
typename IndexType = Eigen::DenseIndex>
30+
using EigenVector = framework::EigenVector<T, MajorType, IndexType>;
31+
32+
template <typename DeviceContext, typename ValueType, typename T>
33+
class EighKernel : public framework::OpKernel<T> {
34+
public:
35+
void Compute(const framework::ExecutionContext& ctx) const override {
36+
auto input_var = ctx.Input<Tensor>("X");
37+
auto output_w_var = ctx.Output<Tensor>("Eigenvalues");
38+
auto output_v_var = ctx.Output<Tensor>("Eigenvectors");
39+
std::string lower = ctx.Attr<std::string>("UPLO");
40+
bool is_lower = (lower == "L");
41+
math::MatrixEighFunctorCPU<DeviceContext, ValueType, T> functor;
42+
functor(ctx, *input_var, output_w_var, output_v_var, is_lower, true);
43+
}
44+
};
45+
46+
template <typename DeviceContext, typename ValueType, typename T>
47+
class EighGradKernel : public framework::OpKernel<T> {
48+
public:
49+
void Compute(const framework::ExecutionContext& ctx) const override {
50+
auto& x_grad = *ctx.Output<framework::Tensor>(framework::GradVarName("X"));
51+
x_grad.mutable_data<T>(ctx.GetPlace());
52+
auto& output_w_var = *ctx.Input<Tensor>("Eigenvalues");
53+
auto& output_v_var = *ctx.Input<Tensor>("Eigenvectors");
54+
auto& output_w_grad =
55+
*ctx.Input<Tensor>(framework::GradVarName("Eigenvalues"));
56+
auto& output_v_grad =
57+
*ctx.Input<Tensor>(framework::GradVarName("Eigenvectors"));
58+
59+
auto& dims = output_v_var.dims();
60+
const int m = dims[dims.size() - 1];
61+
auto dito =
62+
math::DeviceIndependenceTensorOperations<DeviceContext, T, ValueType>(
63+
ctx);
64+
auto tV = dito.Transpose(dito.Conj(output_v_var));
65+
auto W = dito.Sub_(dito.Unsqueeze(output_w_var, -2),
66+
dito.Unsqueeze(output_w_var, -1));
67+
Tensor result = dito.Matmul(tV, output_v_grad);
68+
result.mutable_data<T>(dims, ctx.GetPlace());
69+
std::vector<int> out_shape = framework::vectorize<int>(dims);
70+
auto constant = dito.Fill(out_shape, 0.5);
71+
result = dito.Sub(result, dito.Conj(dito.Transpose(result)));
72+
result = dito.Mul(result, constant);
73+
result = dito.Div_(result, W);
74+
result = dito.DiagFill(m, m, m, 0, output_w_grad, result);
75+
x_grad = dito.Matmul(output_v_var, dito.Matmul(result, tV));
76+
}
77+
};
78+
79+
} // namespace operators
80+
} // namespace paddle

0 commit comments

Comments
 (0)