From 3149e399d31045c1357d07e0f5de6dda1e2798d8 Mon Sep 17 00:00:00 2001 From: zyfncg Date: Mon, 14 Mar 2022 20:01:10 +0800 Subject: [PATCH] [PHI] Move set_value_grad kernel form fluid to phi (#40478) * move set_value_grad kernel form fluid to phi * add unittest for passing coverage ci --- paddle/fluid/operators/set_value_op.cc | 8 - paddle/fluid/operators/set_value_op.cu | 25 -- paddle/fluid/operators/set_value_op.h | 268 -------------- .../phi/kernels/cpu/set_value_grad_kernel.cc | 29 ++ .../phi/kernels/gpu/set_value_grad_kernel.cu | 29 ++ .../kernels/impl/set_value_grad_kernel_impl.h | 344 ++++++++++++++++++ .../phi/kernels/impl/set_value_kernel_impl.h | 1 - paddle/phi/kernels/set_value_grad_kernel.h | 34 ++ paddle/phi/ops/compat/set_value_sig.cc | 102 ++++++ paddle/phi/tests/ops/test_op_signature.cc | 65 ++++ 10 files changed, 603 insertions(+), 302 deletions(-) delete mode 100644 paddle/fluid/operators/set_value_op.cu create mode 100644 paddle/phi/kernels/cpu/set_value_grad_kernel.cc create mode 100644 paddle/phi/kernels/gpu/set_value_grad_kernel.cu create mode 100644 paddle/phi/kernels/impl/set_value_grad_kernel_impl.h create mode 100644 paddle/phi/kernels/set_value_grad_kernel.h diff --git a/paddle/fluid/operators/set_value_op.cc b/paddle/fluid/operators/set_value_op.cc index 7d0d782b83..513ab46e9b 100644 --- a/paddle/fluid/operators/set_value_op.cc +++ b/paddle/fluid/operators/set_value_op.cc @@ -243,14 +243,6 @@ REGISTER_OPERATOR(set_value, ops::SetValue, ops::SetValueMaker, REGISTER_OPERATOR(set_value_grad, ops::SetValueGrad); -REGISTER_OP_CPU_KERNEL( - set_value_grad, - ops::SetValueGradKernel, - ops::SetValueGradKernel, - ops::SetValueGradKernel, - ops::SetValueGradKernel, - ops::SetValueGradKernel); - REGISTER_OP_VERSION(set_value) .AddCheckpoint( R"ROC( diff --git a/paddle/fluid/operators/set_value_op.cu b/paddle/fluid/operators/set_value_op.cu deleted file mode 100644 index 9f291a863c..0000000000 --- a/paddle/fluid/operators/set_value_op.cu +++ /dev/null @@ -1,25 +0,0 @@ -// Copyright (c) 2020 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/set_value_op.h" - -namespace ops = paddle::operators; - -REGISTER_OP_CUDA_KERNEL( - set_value_grad, - ops::SetValueGradKernel, - ops::SetValueGradKernel, - ops::SetValueGradKernel, - ops::SetValueGradKernel, - ops::SetValueGradKernel); diff --git a/paddle/fluid/operators/set_value_op.h b/paddle/fluid/operators/set_value_op.h index 4d459f8c01..4696907f32 100644 --- a/paddle/fluid/operators/set_value_op.h +++ b/paddle/fluid/operators/set_value_op.h @@ -19,14 +19,10 @@ #include #include "paddle/fluid/framework/convert_utils.h" -#include "paddle/fluid/framework/eigen.h" #include "paddle/fluid/framework/op_registry.h" #include "paddle/fluid/framework/tensor_util.h" #include "paddle/fluid/operators/assign_value_op.h" -#include "paddle/fluid/operators/eigen/eigen_function.h" -#include "paddle/fluid/operators/elementwise/elementwise_op_function.h" #include "paddle/fluid/operators/slice_utils.h" -#include "paddle/fluid/operators/strided_slice_op.h" #include "paddle/fluid/operators/utils.h" #include "paddle/fluid/platform/enforce.h" @@ -36,23 +32,6 @@ namespace operators { using Tensor = framework::Tensor; using DDim = framework::DDim; -inline void GetOffsets(const DDim& big_dim, const DDim& small_dim, - DDim start_offset, int cur_dim, - std::vector* offsets) { - if (cur_dim == big_dim.size()) { - offsets->push_back(start_offset); - return; - } - if (small_dim[cur_dim] == big_dim[cur_dim]) { - GetOffsets(big_dim, small_dim, start_offset, cur_dim + 1, offsets); - } else { - for (int i = 0; i < big_dim[cur_dim]; i++) { - GetOffsets(big_dim, small_dim, start_offset, cur_dim + 1, offsets); - start_offset[cur_dim] += 1; - } - } -} - inline std::string GetValueName(framework::proto::VarType::Type data_type) { std::string value_name; switch (data_type) { @@ -121,253 +100,6 @@ inline void CheckIsDimsMatch(const framework::DDim first, "of target shape: %d, but now shape is %d.", second.to_str(), first.to_str())); } -template -class SetValueGradKernel : public framework::OpKernel { - public: - void Compute(const framework::ExecutionContext& ctx) const override { - int rank = ctx.Input(framework::GradVarName("Out"))->dims().size(); - - switch (rank) { - case 1: - SetValueGradCompute<1>(ctx); - break; - case 2: - SetValueGradCompute<2>(ctx); - break; - case 3: - SetValueGradCompute<3>(ctx); - break; - case 4: - SetValueGradCompute<4>(ctx); - break; - case 5: - SetValueGradCompute<5>(ctx); - break; - case 6: - SetValueGradCompute<6>(ctx); - break; - default: - PADDLE_THROW(platform::errors::InvalidArgument( - "The rank of set_value_grad's input should be less than 7, but " - "received %d.", - rank)); - } - } - - private: - template - void SetValueGradCompute(const framework::ExecutionContext& context) const { - auto starts = context.Attr>("starts"); - auto ends = context.Attr>("ends"); - auto steps = context.Attr>("steps"); - - auto axes_int64 = context.Attr>("axes"); - std::vector axes(axes_int64.begin(), axes_int64.end()); - - auto starts_indices = Eigen::DSizes(); - auto ends_indices = Eigen::DSizes(); - auto steps_indices = Eigen::DSizes(); - auto reverse_axis = Eigen::array(); - - auto list_new_ends_tensor = - context.MultiInput("EndsTensorList"); - auto list_new_starts_tensor = - context.MultiInput("StartsTensorList"); - auto list_new_steps_tensor = - context.MultiInput("StepsTensorList"); - - if (list_new_starts_tensor.size() > 0) { - starts = GetDataFromTensorList(list_new_starts_tensor); - } - - if (list_new_ends_tensor.size() > 0) { - ends = GetDataFromTensorList(list_new_ends_tensor); - } - - if (list_new_steps_tensor.size() > 0) { - steps = GetDataFromTensorList(list_new_steps_tensor); - } - - auto in = context.Input(framework::GradVarName("Out")); - PADDLE_ENFORCE_EQ( - in->IsInitialized(), true, - platform::errors::PermissionDenied( - "The input of `set_value_grad`(%s) has not been initialized", - framework::GradVarName("Out"))); - auto grad_value = context.Output( - framework::GradVarName("ValueTensor")); - auto grad_input = - context.Output(framework::GradVarName("Input")); - auto in_dims = in->dims(); - - auto decrease_axis_int64 = - context.Attr>("decrease_axes"); - std::vector decrease_axis(decrease_axis_int64.begin(), - decrease_axis_int64.end()); - std::vector infer_flags(axes.size(), 1); - std::vector out_dims_vector(in_dims.size(), -1); - StridedSliceOutDims(starts, ends, steps, axes, infer_flags, in_dims, - decrease_axis, out_dims_vector.data(), axes.size(), - false); - - framework::DDim out_dims(phi::make_ddim(out_dims_vector)); - - std::vector reverse_vector(starts.size(), 0); - StridedSliceFunctor(starts.data(), ends.data(), steps.data(), axes.data(), - reverse_vector.data(), in_dims, infer_flags, - decrease_axis, starts.size()); - - for (size_t axis = 0; axis < D; axis++) { - starts_indices[axis] = 0; - ends_indices[axis] = out_dims[axis]; - steps_indices[axis] = 1; - reverse_axis[axis] = false; - } - - for (size_t axis = 0; axis < axes.size(); axis++) { - int axis_index = axes[axis]; - starts_indices[axis_index] = starts[axis]; - ends_indices[axis_index] = ends[axis]; - steps_indices[axis_index] = steps[axis]; - reverse_axis[axis_index] = (reverse_vector[axis] == 1) ? true : false; - } - - bool need_reverse = false; - for (size_t axis = 0; axis < axes.size(); axis++) { - if (reverse_vector[axis] == 1) { - need_reverse = true; - break; - } - } - - auto& dev_ctx = context.template device_context(); - auto& place = - *context.template device_context().eigen_device(); - phi::funcs::SetConstant set_zero; - - if (grad_input) { - // Set gradient of `Input` - paddle::framework::TensorCopy(*in, context.GetPlace(), grad_input); - - auto grad_input_t = - framework::EigenTensor::From(*grad_input); - - framework::Tensor tmp(grad_input->dtype()); - tmp.mutable_data(out_dims, context.GetPlace()); - set_zero(dev_ctx, &tmp, static_cast(0)); - auto tmp_t = framework::EigenTensor::From(tmp); - - grad_input_t.stridedSlice(starts_indices, ends_indices, steps_indices) - .device(place) = tmp_t; - } - if (grad_value) { - grad_value->mutable_data(context.GetPlace()); - set_zero(dev_ctx, grad_value, static_cast(0)); - - auto in_t = framework::EigenTensor::From(*in); - - if (grad_value->dims() == out_dims) { - auto grad_value_t = - framework::EigenTensor::From(*grad_value); - if (need_reverse) { - framework::Tensor tmp(grad_value->dtype()); - tmp.mutable_data(out_dims, context.GetPlace()); - set_zero(dev_ctx, &tmp, static_cast(0)); - auto tmp_t = framework::EigenTensor::From(tmp); - - tmp_t.device(place) = - in_t.stridedSlice(starts_indices, ends_indices, steps_indices); - grad_value_t.device(place) = tmp_t.reverse(reverse_axis); - } else { - grad_value_t.device(place) = - in_t.stridedSlice(starts_indices, ends_indices, steps_indices); - } - } else { - int out_dims_size = out_dims.size(); - auto grad_value_dims = grad_value->dims(); - auto fake_grad_value_dims = out_dims; - - // Create an extented shape according to the rules of broadcast. - auto grad_value_dims_size = grad_value_dims.size(); - - int num_decrease = 0; - - int decrease_axis_size = decrease_axis.size(); - for (int i = 0; i < out_dims_size; i++) { - if (decrease_axis.end() != - std::find(decrease_axis.begin(), decrease_axis.end(), i)) { - fake_grad_value_dims[i] = 1; - num_decrease++; - } else if (i < out_dims_size - (grad_value_dims_size + - decrease_axis_size - num_decrease)) { - fake_grad_value_dims[i] = 1; - } else { - auto index_grad = - i - (out_dims_size - (grad_value_dims_size + - decrease_axis_size - num_decrease)); - fake_grad_value_dims[i] = grad_value_dims[index_grad]; - - PADDLE_ENFORCE_EQ((out_dims[i] == grad_value_dims[index_grad]) || - (grad_value_dims[index_grad] == 1), - true, - platform::errors::InvalidArgument( - "An error occurred while calculating %s: " - "[%s] can not be accumulated into [%s].", - framework::GradVarName("ValueTensor"), - out_dims, grad_value_dims)); - } - } - - VLOG(3) << "Dimensions of " << framework::GradVarName("ValueTensor") - << "([" << grad_value_dims << "])is broadcasted into [" - << fake_grad_value_dims << "]."; - - auto extent = Eigen::DSizes(); - auto offset = out_dims; - for (int i = 0; i < out_dims_size; i++) { - offset[i] = 0; - extent[i] = fake_grad_value_dims[i]; - } - std::vector offsets; - GetOffsets(out_dims, fake_grad_value_dims, offset, 0, &offsets); - - auto grad_value_t = - framework::EigenTensor:: - From(*grad_value, fake_grad_value_dims); - - framework::Tensor tmp(grad_value->dtype()); - tmp.mutable_data(out_dims, context.GetPlace()); - set_zero(dev_ctx, &tmp, static_cast(0)); - auto tmp_t = framework::EigenTensor::From(tmp); - - tmp_t.device(place) = - in_t.stridedSlice(starts_indices, ends_indices, steps_indices); - - // accumulate gradient - for (auto offset : offsets) { - grad_value_t.device(place) = - grad_value_t + - tmp_t.slice(framework::EigenDim::From(offset), extent); - } - if (need_reverse) { - framework::Tensor tmp_value(grad_value->dtype()); - tmp_value.mutable_data(fake_grad_value_dims, context.GetPlace()); - auto tmp_value_t = - framework::EigenTensor::From(tmp_value); - tmp_value_t.device(place) = grad_value_t.reverse(reverse_axis); - grad_value_t.device(place) = tmp_value_t; - } - } - } - } -}; } // namespace operators } // namespace paddle diff --git a/paddle/phi/kernels/cpu/set_value_grad_kernel.cc b/paddle/phi/kernels/cpu/set_value_grad_kernel.cc new file mode 100644 index 0000000000..44df36bb9f --- /dev/null +++ b/paddle/phi/kernels/cpu/set_value_grad_kernel.cc @@ -0,0 +1,29 @@ +// 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/phi/kernels/set_value_grad_kernel.h" + +#include "paddle/phi/backends/cpu/cpu_context.h" +#include "paddle/phi/core/kernel_registry.h" +#include "paddle/phi/kernels/impl/set_value_grad_kernel_impl.h" + +PD_REGISTER_KERNEL(set_value_grad, + CPU, + ALL_LAYOUT, + phi::SetValueGradKernel, + float, + double, + int, + int64_t, + bool) {} diff --git a/paddle/phi/kernels/gpu/set_value_grad_kernel.cu b/paddle/phi/kernels/gpu/set_value_grad_kernel.cu new file mode 100644 index 0000000000..7eed96699e --- /dev/null +++ b/paddle/phi/kernels/gpu/set_value_grad_kernel.cu @@ -0,0 +1,29 @@ +// 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/phi/kernels/set_value_grad_kernel.h" + +#include "paddle/phi/backends/gpu/gpu_context.h" +#include "paddle/phi/core/kernel_registry.h" +#include "paddle/phi/kernels/impl/set_value_grad_kernel_impl.h" + +PD_REGISTER_KERNEL(set_value_grad, + GPU, + ALL_LAYOUT, + phi::SetValueGradKernel, + float, + double, + int, + int64_t, + bool) {} diff --git a/paddle/phi/kernels/impl/set_value_grad_kernel_impl.h b/paddle/phi/kernels/impl/set_value_grad_kernel_impl.h new file mode 100644 index 0000000000..4947170088 --- /dev/null +++ b/paddle/phi/kernels/impl/set_value_grad_kernel_impl.h @@ -0,0 +1,344 @@ +// 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/phi/common/scalar_array.h" +#include "paddle/phi/core/dense_tensor.h" + +#include "paddle/phi/kernels/copy_kernel.h" +#include "paddle/phi/kernels/full_kernel.h" +#include "paddle/phi/kernels/funcs/eigen/common.h" +#include "paddle/phi/kernels/funcs/eigen/eigen_function.h" +#include "paddle/phi/kernels/funcs/math_function.h" + +#include "paddle/fluid/operators/strided_slice_op.h" + +namespace phi { + +inline void GetOffsets(const DDim& big_dim, + const DDim& small_dim, + DDim start_offset, + int cur_dim, + std::vector* offsets) { + if (cur_dim == big_dim.size()) { + offsets->push_back(start_offset); + return; + } + if (small_dim[cur_dim] == big_dim[cur_dim]) { + GetOffsets(big_dim, small_dim, start_offset, cur_dim + 1, offsets); + } else { + for (int i = 0; i < big_dim[cur_dim]; i++) { + GetOffsets(big_dim, small_dim, start_offset, cur_dim + 1, offsets); + start_offset[cur_dim] += 1; + } + } +} + +template +void SetValueGradImpl(const Context& dev_ctx, + const DenseTensor& out_grad, + const ScalarArray& starts, + const ScalarArray& ends, + const ScalarArray& steps, + const std::vector& axes, + const std::vector& decrease_axes, + const std::vector& none_axes, + DenseTensor* x_grad, + DenseTensor* value_grad) { + PADDLE_ENFORCE_EQ( + out_grad.IsInitialized(), + true, + errors::PermissionDenied( + "The input of `set_value_grad`(out_grad) has not been initialized")); + + auto in_dims = out_grad.dims(); + + std::vector decrease_axis_int32(decrease_axes.begin(), + decrease_axes.end()); + std::vector axes_int32(axes.begin(), axes.end()); + std::vector infer_flags(axes.size(), 1); + std::vector out_dims_vector(in_dims.size(), -1); + std::vector starts_local = starts.GetData(); + std::vector ends_local = ends.GetData(); + std::vector steps_local = steps.GetData(); + paddle::operators::StridedSliceOutDims(starts_local, + ends_local, + steps_local, + axes_int32, + infer_flags, + in_dims, + decrease_axis_int32, + out_dims_vector.data(), + axes.size(), + false); + + DDim out_dims(phi::make_ddim(out_dims_vector)); + + std::vector reverse_vector(starts_local.size(), 0); + paddle::operators::StridedSliceFunctor(starts_local.data(), + ends_local.data(), + steps_local.data(), + axes_int32.data(), + reverse_vector.data(), + in_dims, + infer_flags, + decrease_axis_int32, + starts_local.size()); + + auto starts_indices = Eigen::DSizes(); + auto ends_indices = Eigen::DSizes(); + auto steps_indices = Eigen::DSizes(); + auto reverse_axis = Eigen::array(); + + for (size_t axis = 0; axis < RANK; axis++) { + starts_indices[axis] = 0; + ends_indices[axis] = out_dims[axis]; + steps_indices[axis] = 1; + reverse_axis[axis] = false; + } + + for (size_t axis = 0; axis < axes.size(); axis++) { + int axis_index = axes[axis]; + starts_indices[axis_index] = starts_local[axis]; + ends_indices[axis_index] = ends_local[axis]; + steps_indices[axis_index] = steps_local[axis]; + reverse_axis[axis_index] = (reverse_vector[axis] == 1) ? true : false; + } + + bool need_reverse = false; + for (size_t axis = 0; axis < axes.size(); axis++) { + if (reverse_vector[axis] == 1) { + need_reverse = true; + break; + } + } + + auto& place = *dev_ctx.eigen_device(); + phi::funcs::SetConstant set_zero; + + if (x_grad) { + // Set gradient of `Input` + Copy(dev_ctx, out_grad, dev_ctx.GetPlace(), false, x_grad); + + auto x_grad_t = + EigenTensor::From(*x_grad); + + DenseTensor tmp = Full(dev_ctx, out_dims_vector, static_cast(0)); + auto tmp_t = + EigenTensor::From(tmp); + + x_grad_t.stridedSlice(starts_indices, ends_indices, steps_indices) + .device(place) = tmp_t; + } + if (value_grad) { + dev_ctx.template Alloc(value_grad); + set_zero(dev_ctx, value_grad, static_cast(0)); + + auto in_t = EigenTensor::From( + out_grad); + + if (value_grad->dims() == out_dims) { + auto value_grad_t = + EigenTensor::From( + *value_grad); + if (need_reverse) { + DenseTensor tmp = Full(dev_ctx, out_dims_vector, static_cast(0)); + auto tmp_t = + EigenTensor::From(tmp); + + tmp_t.device(place) = + in_t.stridedSlice(starts_indices, ends_indices, steps_indices); + value_grad_t.device(place) = tmp_t.reverse(reverse_axis); + } else { + value_grad_t.device(place) = + in_t.stridedSlice(starts_indices, ends_indices, steps_indices); + } + } else { + int out_dims_size = out_dims.size(); + auto value_grad_dims = value_grad->dims(); + auto fake_value_grad_dims = out_dims; + + // Create an extented shape according to the rules of broadcast. + auto value_grad_dims_size = value_grad_dims.size(); + + int num_decrease = 0; + + int decrease_axis_size = decrease_axes.size(); + for (int i = 0; i < out_dims_size; i++) { + if (decrease_axes.end() != + std::find(decrease_axes.begin(), decrease_axes.end(), i)) { + fake_value_grad_dims[i] = 1; + num_decrease++; + } else if (i < out_dims_size - (value_grad_dims_size + + decrease_axis_size - num_decrease)) { + fake_value_grad_dims[i] = 1; + } else { + auto index_grad = + i - (out_dims_size - + (value_grad_dims_size + decrease_axis_size - num_decrease)); + fake_value_grad_dims[i] = value_grad_dims[index_grad]; + + PADDLE_ENFORCE_EQ((out_dims[i] == value_grad_dims[index_grad]) || + (value_grad_dims[index_grad] == 1), + true, + errors::InvalidArgument( + "An error occurred while calculating %s: " + "[%s] can not be accumulated into [%s].", + paddle::framework::GradVarName("ValueTensor"), + out_dims, + value_grad_dims)); + } + } + + VLOG(3) << "Dimensions of " + << paddle::framework::GradVarName("ValueTensor") << "([" + << value_grad_dims << "])is broadcasted into [" + << fake_value_grad_dims << "]."; + + auto extent = Eigen::DSizes(); + auto offset = out_dims; + for (int i = 0; i < out_dims_size; i++) { + offset[i] = 0; + extent[i] = fake_value_grad_dims[i]; + } + std::vector offsets; + GetOffsets(out_dims, fake_value_grad_dims, offset, 0, &offsets); + + auto value_grad_t = + EigenTensor::From( + *value_grad, fake_value_grad_dims); + + DenseTensor tmp = Full(dev_ctx, out_dims_vector, static_cast(0)); + auto tmp_t = + EigenTensor::From(tmp); + + tmp_t.device(place) = + in_t.stridedSlice(starts_indices, ends_indices, steps_indices); + + // accumulate gradient + for (auto offset : offsets) { + value_grad_t.device(place) = + value_grad_t + tmp_t.slice(EigenDim::From(offset), extent); + } + if (need_reverse) { + DenseTensor tmp_value = + Full(dev_ctx, + {fake_value_grad_dims.Get(), fake_value_grad_dims.size()}, + static_cast(0)); + auto tmp_value_t = + EigenTensor::From( + tmp_value); + tmp_value_t.device(place) = value_grad_t.reverse(reverse_axis); + value_grad_t.device(place) = tmp_value_t; + } + } + } +} + +template +void SetValueGradKernel(const Context& dev_ctx, + const DenseTensor& out_grad, + const ScalarArray& starts, + const ScalarArray& ends, + const ScalarArray& steps, + const std::vector& axes, + const std::vector& decrease_axes, + const std::vector& none_axes, + DenseTensor* x_grad, + DenseTensor* value_grad) { + const int rank = out_grad.dims().size(); + + switch (rank) { + case 1: + SetValueGradImpl(dev_ctx, + out_grad, + starts, + ends, + steps, + axes, + decrease_axes, + none_axes, + x_grad, + value_grad); + break; + case 2: + SetValueGradImpl(dev_ctx, + out_grad, + starts, + ends, + steps, + axes, + decrease_axes, + none_axes, + x_grad, + value_grad); + break; + case 3: + SetValueGradImpl(dev_ctx, + out_grad, + starts, + ends, + steps, + axes, + decrease_axes, + none_axes, + x_grad, + value_grad); + break; + case 4: + SetValueGradImpl(dev_ctx, + out_grad, + starts, + ends, + steps, + axes, + decrease_axes, + none_axes, + x_grad, + value_grad); + break; + case 5: + SetValueGradImpl(dev_ctx, + out_grad, + starts, + ends, + steps, + axes, + decrease_axes, + none_axes, + x_grad, + value_grad); + break; + case 6: + SetValueGradImpl(dev_ctx, + out_grad, + starts, + ends, + steps, + axes, + decrease_axes, + none_axes, + x_grad, + value_grad); + break; + default: + PADDLE_THROW(phi::errors::InvalidArgument( + "The rank of set_value_grad's input should be less than 7, but " + "received %d.", + rank)); + } +} + +} // namespace phi diff --git a/paddle/phi/kernels/impl/set_value_kernel_impl.h b/paddle/phi/kernels/impl/set_value_kernel_impl.h index 5aebffe51b..99db559f3b 100644 --- a/paddle/phi/kernels/impl/set_value_kernel_impl.h +++ b/paddle/phi/kernels/impl/set_value_kernel_impl.h @@ -25,7 +25,6 @@ #include "paddle/phi/kernels/funcs/eigen/eigen_function.h" #include "paddle/phi/kernels/funcs/elementwise_functor.h" -#include "paddle/fluid/framework/tensor_util.h" #include "paddle/fluid/operators/slice_utils.h" namespace phi { diff --git a/paddle/phi/kernels/set_value_grad_kernel.h b/paddle/phi/kernels/set_value_grad_kernel.h new file mode 100644 index 0000000000..6a028b0c8d --- /dev/null +++ b/paddle/phi/kernels/set_value_grad_kernel.h @@ -0,0 +1,34 @@ +// 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/phi/common/scalar_array.h" +#include "paddle/phi/core/dense_tensor.h" + +namespace phi { + +template +void SetValueGradKernel(const Context& dev_ctx, + const DenseTensor& out_grad, + const ScalarArray& starts, + const ScalarArray& ends, + const ScalarArray& steps, + const std::vector& axes, + const std::vector& decrease_axes, + const std::vector& none_axes, + DenseTensor* x_grad, + DenseTensor* value_grad); + +} // namespace phi diff --git a/paddle/phi/ops/compat/set_value_sig.cc b/paddle/phi/ops/compat/set_value_sig.cc index eacfff26d5..9653250bde 100644 --- a/paddle/phi/ops/compat/set_value_sig.cc +++ b/paddle/phi/ops/compat/set_value_sig.cc @@ -731,6 +731,108 @@ KernelSignature SetValueOpArgumentMapping(const ArgumentMappingContext& ctx) { } return KernelSignature("unregistered", {}, {}, {}); } + +KernelSignature SetValueGradOpArgumentMapping( + const ArgumentMappingContext& ctx) { + if (ctx.HasInput("StartsTensorList")) { + if (ctx.HasInput("EndsTensorList")) { + if (ctx.HasInput("StepsTensorList")) { + return KernelSignature( + "set_value_grad", + {GradVarName("Out")}, + {"StartsTensorList", + "EndsTensorList", + "StepsTensorList", + "axes", + "decrease_axes", + "none_axes"}, + {GradVarName("Input"), GradVarName("ValueTensor")}); + } else { + return KernelSignature( + "set_value_grad", + {GradVarName("Out")}, + {"StartsTensorList", + "EndsTensorList", + "steps", + "axes", + "decrease_axes", + "none_axes"}, + {GradVarName("Input"), GradVarName("ValueTensor")}); + } + } else { + if (ctx.HasInput("StepsTensorList")) { + return KernelSignature( + "set_value_grad", + {GradVarName("Out")}, + {"StartsTensorList", + "ends", + "StepsTensorList", + "axes", + "decrease_axes", + "none_axes"}, + {GradVarName("Input"), GradVarName("ValueTensor")}); + } else { + return KernelSignature( + "set_value_grad", + {GradVarName("Out")}, + {"StartsTensorList", + "ends", + "steps", + "axes", + "decrease_axes", + "none_axes"}, + {GradVarName("Input"), GradVarName("ValueTensor")}); + } + } + } else { + if (ctx.HasInput("EndsTensorList")) { + if (ctx.HasInput("StepsTensorList")) { + return KernelSignature( + "set_value_grad", + {GradVarName("Out")}, + {"starts", + "EndsTensorList", + "StepsTensorList", + "axes", + "decrease_axes", + "none_axes"}, + {GradVarName("Input"), GradVarName("ValueTensor")}); + } else { + return KernelSignature( + "set_value_grad", + {GradVarName("Out")}, + {"starts", + "EndsTensorList", + "steps", + "axes", + "decrease_axes", + "none_axes"}, + {GradVarName("Input"), GradVarName("ValueTensor")}); + } + } else { + if (ctx.HasInput("StepsTensorList")) { + return KernelSignature( + "set_value_grad", + {GradVarName("Out")}, + {"starts", + "ends", + "StepsTensorList", + "axes", + "decrease_axes", + "none_axes"}, + {GradVarName("Input"), GradVarName("ValueTensor")}); + } else { + return KernelSignature( + "set_value_grad", + {GradVarName("Out")}, + {"starts", "ends", "steps", "axes", "decrease_axes", "none_axes"}, + {GradVarName("Input"), GradVarName("ValueTensor")}); + } + } + } +} + } // namespace phi PD_REGISTER_ARG_MAPPING_FN(set_value, phi::SetValueOpArgumentMapping); +PD_REGISTER_ARG_MAPPING_FN(set_value_grad, phi::SetValueGradOpArgumentMapping); diff --git a/paddle/phi/tests/ops/test_op_signature.cc b/paddle/phi/tests/ops/test_op_signature.cc index c74049e0f0..36923972ea 100644 --- a/paddle/phi/tests/ops/test_op_signature.cc +++ b/paddle/phi/tests/ops/test_op_signature.cc @@ -484,6 +484,71 @@ TEST(ARG_MAP, set_value) { "set_value"); } +TEST(ARG_MAP, set_value_grad) { + TestArgumentMappingContext arg_case( + {"Out@GRAD", "StartsTensorList", "EndsTensorList"}, + {}, + {}, + {"Input@GRAD", "ValueTensor@GRAD"}, + {}); + ASSERT_EQ(OpUtilsMap::Instance() + .GetArgumentMappingFn("set_value_grad")(arg_case) + .name, + "set_value_grad"); + + TestArgumentMappingContext arg_case1( + {"Out@GRAD", "StartsTensorList", "StepsTensorList"}, + {}, + {}, + {"Input@GRAD", "ValueTensor@GRAD"}, + {}); + ASSERT_EQ(OpUtilsMap::Instance() + .GetArgumentMappingFn("set_value_grad")(arg_case1) + .name, + "set_value_grad"); + + TestArgumentMappingContext arg_case2({"Out@GRAD", "StartsTensorList"}, + {}, + {}, + {"Input@GRAD", "ValueTensor@GRAD"}, + {}); + ASSERT_EQ(OpUtilsMap::Instance() + .GetArgumentMappingFn("set_value_grad")(arg_case2) + .name, + "set_value_grad"); + + TestArgumentMappingContext arg_case3( + {"Out@GRAD", "EndsTensorList", "StepsTensorList"}, + {}, + {}, + {"Input@GRAD", "ValueTensor@GRAD"}, + {}); + ASSERT_EQ(OpUtilsMap::Instance() + .GetArgumentMappingFn("set_value_grad")(arg_case3) + .name, + "set_value_grad"); + + TestArgumentMappingContext arg_case4({"Out@GRAD", "EndsTensorList"}, + {}, + {}, + {"Input@GRAD", "ValueTensor@GRAD"}, + {}); + ASSERT_EQ(OpUtilsMap::Instance() + .GetArgumentMappingFn("set_value_grad")(arg_case4) + .name, + "set_value_grad"); + + TestArgumentMappingContext arg_case5({"Out@GRAD", "StepsTensorList"}, + {}, + {}, + {"Input@GRAD", "ValueTensor@GRAD"}, + {}); + ASSERT_EQ(OpUtilsMap::Instance() + .GetArgumentMappingFn("set_value_grad")(arg_case5) + .name, + "set_value_grad"); +} + TEST(ARG_MAP, allclose) { TestArgumentMappingContext arg_case1( {"Input", "Other", "Rtol"}, -- GitLab