Skip to content

Commit a7376d0

Browse files
zhiboniuzhiboniu
committed
add api fill_diagonal_inplace
1 parent 6151ccd commit a7376d0

5 files changed

Lines changed: 586 additions & 0 deletions

File tree

‎paddle/fluid/operators/fill_diagonal_op.cc‎

Lines changed: 217 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,217 @@
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/fill_diagonal_op.h"
16+
17+
namespace paddle {
18+
namespace operators {
19+
20+
int64_t CalStride(framework::DDim dim) {
21+
int rank = dim.size();
22+
int64_t dimsum = 1;
23+
int64_t strides = 0;
24+
for (int i = rank - 1; i >= 0; i--) {
25+
strides += dimsum;
26+
dimsum *= dim[i];
27+
}
28+
return strides;
29+
}
30+
31+
class FillIDiagonalOpMaker : public framework::OpProtoAndCheckerMaker {
32+
public:
33+
void Make() override {
34+
AddComment(R"DOC(Fill replace operator
35+
Fill the diagonal of an tensor with 'value'.
36+
)DOC");
37+
AddInput("X", "(Tensor) The input tensor.");
38+
AddOutput("Out",
39+
"Tensor, the output tensor, with the same shape and data type "
40+
"as input(x)");
41+
AddAttr<float>(
42+
"value",
43+
"The float values of tensor, whose dim is one, and no need of grad")
44+
.SetDefault(0);
45+
AddAttr<bool>("wrap",
46+
"the diagonal 'wrapped' after N columns for tall matrices")
47+
.SetDefault(false);
48+
AddAttr<int>("offset",
49+
"offset of diagonal, zero means no offset, positive means "
50+
"offset to up-right corner; negtive means offset to "
51+
"bottom-left corner")
52+
.SetDefault(0);
53+
}
54+
};
55+
56+
class FillIDiagonalOp : public framework::OperatorWithKernel {
57+
public:
58+
using framework::OperatorWithKernel::OperatorWithKernel;
59+
60+
void InferShape(framework::InferShapeContext *context) const override {
61+
OP_INOUT_CHECK(context->HasInput("X"), "Input", "X", "FillIDiagonal");
62+
OP_INOUT_CHECK(context->HasOutput("Out"), "Output", "Out", "FillIDiagonal");
63+
auto x_dims = context->GetInputDim("X");
64+
context->SetOutputDim("Out", x_dims);
65+
}
66+
67+
protected:
68+
framework::OpKernelType GetExpectedKernelType(
69+
const framework::ExecutionContext &ctx) const override {
70+
return framework::OpKernelType(
71+
OperatorWithKernel::IndicateVarDataType(ctx, "X"), ctx.GetPlace());
72+
}
73+
};
74+
75+
class FillIDiagonalOpVarTypeInference : public framework::VarTypeInference {
76+
public:
77+
void operator()(framework::InferVarTypeContext *ctx) const override {
78+
auto var_type = ctx->GetInputType("X", 0);
79+
auto data_type = ctx->GetInputDataType("X", 0);
80+
ctx->SetOutputType("Out", var_type, framework::ALL_ELEMENTS);
81+
ctx->SetOutputDataType("Out", data_type, framework::ALL_ELEMENTS);
82+
}
83+
};
84+
85+
template <typename T>
86+
class FillIDiagonalKernel : public framework::OpKernel<T> {
87+
public:
88+
void Compute(const paddle::framework::ExecutionContext &ctx) const override {
89+
auto fill_val = ctx.template Attr<float>("value");
90+
auto *out = ctx.Output<framework::Tensor>("Out");
91+
auto offset = ctx.Attr<int>("offset");
92+
auto wrap = ctx.Attr<bool>("wrap");
93+
94+
auto *xin = ctx.Input<framework::Tensor>("X");
95+
96+
T temp_var = static_cast<T>(fill_val);
97+
98+
T *out_data = out->mutable_data<T>(ctx.GetPlace());
99+
framework::TensorCopy(*xin, ctx.GetPlace(), out);
100+
101+
auto out_dims = out->dims();
102+
auto strides = CalStride(out_dims);
103+
auto size = out->numel();
104+
105+
// The wrap mode supported only the dims equels to 2; In wrap mode, the
106+
// value will be filled in cycles
107+
if (!wrap) {
108+
size = std::min(size, out_dims[1] * out_dims[1]);
109+
}
110+
111+
for (int64_t i = offset; i < size; i += strides) {
112+
out_data[i] = temp_var;
113+
}
114+
}
115+
};
116+
117+
class FillIDiagonalGradOp : public framework::OperatorWithKernel {
118+
public:
119+
using framework::OperatorWithKernel::OperatorWithKernel;
120+
121+
void InferShape(framework::InferShapeContext *ctx) const override {
122+
OP_INOUT_CHECK(ctx->HasInput(framework::GradVarName("Out")), "Input",
123+
"Out@GRAD", "mul");
124+
auto x_dims = ctx->GetInputDim(framework::GradVarName("Out"));
125+
auto x_grad_name = framework::GradVarName("X");
126+
if (ctx->HasOutput(x_grad_name)) {
127+
ctx->SetOutputDim(x_grad_name, x_dims);
128+
}
129+
}
130+
131+
framework::OpKernelType GetExpectedKernelType(
132+
const framework::ExecutionContext &ctx) const override {
133+
// Note: don't get data type from ctx.Input<framework::Tensor>("Input");
134+
auto dtype =
135+
ctx.Input<framework::Tensor>(framework::GradVarName("Out"))->type();
136+
return framework::OpKernelType(dtype, ctx.GetPlace());
137+
}
138+
};
139+
140+
template <typename T>
141+
class FillIDiagonalGradOpMaker : public framework::SingleGradOpMaker<T> {
142+
public:
143+
using framework::SingleGradOpMaker<T>::SingleGradOpMaker;
144+
145+
protected:
146+
void Apply(GradOpPtr<T> retv) const override {
147+
retv->SetType("fill_diagonal_grad");
148+
retv->SetInput(framework::GradVarName("Out"), this->OutputGrad("Out"));
149+
retv->SetOutput(framework::GradVarName("X"), this->InputGrad("X"));
150+
retv->SetAttrMap(this->Attrs());
151+
}
152+
};
153+
154+
template <typename T>
155+
class FillIDiagonalGradKernel : public framework::OpKernel<T> {
156+
public:
157+
void Compute(const paddle::framework::ExecutionContext &ctx) const override {
158+
auto *dx = ctx.Output<framework::Tensor>(framework::GradVarName("X"));
159+
auto *dout = ctx.Input<framework::Tensor>(framework::GradVarName("Out"));
160+
161+
auto offset = ctx.Attr<int>("offset");
162+
auto wrap = ctx.Attr<bool>("wrap");
163+
164+
if (dx) {
165+
auto *data = dx->mutable_data<T>(ctx.GetPlace());
166+
framework::TensorCopy(*dout, ctx.GetPlace(), dx);
167+
168+
auto dx_dims = dx->dims();
169+
auto strides = CalStride(dx_dims);
170+
auto size = dx->numel();
171+
auto wrapsize = std::min(size, dx_dims[1] * dx_dims[1]);
172+
173+
// The wrap mode supported only the dims equels to 2; In wrap mode, the
174+
// value will be filled in cycles
175+
if (wrap) {
176+
wrapsize = size;
177+
}
178+
179+
for (int64_t i = offset; i < wrapsize; i += strides) {
180+
data[i] = T(0);
181+
}
182+
}
183+
}
184+
};
185+
186+
DECLARE_INPLACE_OP_INFERER(FillIDiagonalOpInplaceInferer, {"X", "Out"});
187+
DECLARE_INPLACE_OP_INFERER(FillIDiagonalGradOpInplaceInferer,
188+
{framework::GradVarName("Out"),
189+
framework::GradVarName("X")});
190+
191+
} // namespace operators
192+
} // namespace paddle
193+
namespace ops = paddle::operators;
194+
195+
REGISTER_OPERATOR(fill_diagonal, ops::FillIDiagonalOp,
196+
ops::FillIDiagonalOpMaker,
197+
ops::FillIDiagonalOpVarTypeInference,
198+
ops::FillIDiagonalGradOpMaker<paddle::framework::OpDesc>,
199+
ops::FillIDiagonalGradOpMaker<paddle::imperative::OpBase>,
200+
ops::FillIDiagonalOpInplaceInferer);
201+
202+
REGISTER_OPERATOR(fill_diagonal_grad, ops::FillIDiagonalGradOp,
203+
ops::FillIDiagonalGradOpInplaceInferer);
204+
205+
REGISTER_OP_CPU_KERNEL(fill_diagonal, ops::FillIDiagonalKernel<float>,
206+
ops::FillIDiagonalKernel<double>,
207+
ops::FillIDiagonalKernel<int64_t>,
208+
ops::FillIDiagonalKernel<int>,
209+
ops::FillIDiagonalKernel<paddle::platform::float16>,
210+
ops::FillIDiagonalKernel<bool>);
211+
212+
REGISTER_OP_CPU_KERNEL(fill_diagonal_grad, ops::FillIDiagonalGradKernel<float>,
213+
ops::FillIDiagonalGradKernel<double>,
214+
ops::FillIDiagonalGradKernel<int64_t>,
215+
ops::FillIDiagonalGradKernel<int>,
216+
ops::FillIDiagonalGradKernel<paddle::platform::float16>,
217+
ops::FillIDiagonalGradKernel<bool>);

‎paddle/fluid/operators/fill_diagonal_op.cu‎

Lines changed: 122 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,122 @@
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/fill_diagonal_op.h"
16+
17+
namespace paddle {
18+
namespace operators {
19+
20+
using Tensor = framework::Tensor;
21+
using CUDADeviceContext = paddle::platform::CUDADeviceContext;
22+
23+
template <typename T>
24+
__global__ void fill_constant_kernel(const int64_t featuresize, T* in_data,
25+
int64_t strides, int offset, T fillvar) {
26+
for (int64_t idx = blockIdx.x * featuresize + threadIdx.x;
27+
idx * strides + offset < (blockIdx.x + 1) * featuresize;
28+
idx += blockDim.x) {
29+
in_data[idx * strides + offset] = fillvar;
30+
}
31+
}
32+
33+
template <typename T>
34+
class FillIDiagonalCUDAKernel : public framework::OpKernel<T> {
35+
public:
36+
void Compute(const framework::ExecutionContext& ctx) const override {
37+
#ifdef __HIPCC__
38+
const int64_t kMaxBlockDim = 256;
39+
#else
40+
const int64_t kMaxBlockDim = 512;
41+
#endif
42+
auto* out = ctx.Output<Tensor>("Out");
43+
auto offset = ctx.Attr<int>("offset");
44+
auto wrap = ctx.Attr<bool>("wrap");
45+
46+
auto* xin = ctx.Input<framework::Tensor>("X");
47+
framework::TensorCopy(*xin, ctx.GetPlace(), out);
48+
49+
T* out_data = out->mutable_data<T>(ctx.GetPlace());
50+
auto fill_val = static_cast<T>(ctx.template Attr<float>("value"));
51+
T temp_var = static_cast<T>(fill_val);
52+
53+
auto size = out->numel();
54+
auto out_dims = out->dims();
55+
auto strides = CalStride(out_dims);
56+
57+
// The wrap mode supported only the dims equels to 2; In wrap mode, the
58+
// value will be filled in cycles
59+
if (!wrap) {
60+
size = std::min(size, out_dims[1] * out_dims[1]);
61+
}
62+
63+
int64_t kBlockDim = std::min(int64_t(size / strides), kMaxBlockDim);
64+
fill_constant_kernel<T><<<1, kBlockDim, 0>>>(size, out_data, strides,
65+
offset, temp_var);
66+
}
67+
};
68+
69+
template <typename T>
70+
class FillIDiagonalGradCUDAKernel : public framework::OpKernel<T> {
71+
public:
72+
void Compute(const framework::ExecutionContext& ctx) const override {
73+
#ifdef __HIPCC__
74+
const int64_t kMaxBlockDim = 256;
75+
#else
76+
const int64_t kMaxBlockDim = 512;
77+
#endif
78+
auto* dx = ctx.Output<framework::Tensor>(framework::GradVarName("X"));
79+
auto* in_data = dx->mutable_data<T>(ctx.GetPlace());
80+
auto* dout = ctx.Input<framework::Tensor>(framework::GradVarName("Out"));
81+
auto offset = ctx.Attr<int>("offset");
82+
auto wrap = ctx.Attr<bool>("wrap");
83+
84+
framework::TensorCopy(*dout, ctx.GetPlace(), dx);
85+
86+
auto size = dx->numel();
87+
auto out_dims = dx->dims();
88+
auto strides = CalStride(out_dims);
89+
90+
auto wrapsize = std::min(size, out_dims[1] * out_dims[1]);
91+
// The wrap mode supported only the dims equels to 2; In wrap mode, the
92+
// value will be filled in cycles
93+
if (wrap) {
94+
wrapsize = size;
95+
}
96+
97+
int64_t kBlockDim = std::min(int64_t(size), kMaxBlockDim);
98+
fill_constant_kernel<T><<<1, kBlockDim, 0>>>(wrapsize, in_data, strides,
99+
offset, T(0));
100+
}
101+
};
102+
103+
} // namespace operators
104+
} // namespace paddle
105+
106+
namespace ops = paddle::operators;
107+
namespace plat = paddle::platform;
108+
109+
REGISTER_OP_CUDA_KERNEL(fill_diagonal, ops::FillIDiagonalCUDAKernel<float>,
110+
ops::FillIDiagonalCUDAKernel<double>,
111+
ops::FillIDiagonalCUDAKernel<plat::float16>,
112+
ops::FillIDiagonalCUDAKernel<int>,
113+
ops::FillIDiagonalCUDAKernel<int64_t>,
114+
ops::FillIDiagonalCUDAKernel<bool>);
115+
116+
REGISTER_OP_CUDA_KERNEL(fill_diagonal_grad,
117+
ops::FillIDiagonalGradCUDAKernel<float>,
118+
ops::FillIDiagonalGradCUDAKernel<double>,
119+
ops::FillIDiagonalGradCUDAKernel<int>,
120+
ops::FillIDiagonalGradCUDAKernel<int64_t>,
121+
ops::FillIDiagonalGradCUDAKernel<plat::float16>,
122+
ops::FillIDiagonalGradCUDAKernel<bool>);

‎paddle/fluid/operators/fill_diagonal_op.h‎

Lines changed: 25 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,25 @@
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+
19+
namespace paddle {
20+
namespace operators {
21+
22+
int64_t CalStride(framework::DDim dim);
23+
24+
} // namespace operators
25+
} // namespace paddle

0 commit comments

Comments
 (0)