[Cherry-pick] Fix indexing shape bug and Optimize (#62117)

* tensor_array slice in PIR (#60503) * use slice_array, now will meet error of destory opresult still in use * disable the pir test until the bug fixed * Optimize advanced setting by remove the last set_value (#60771) * pure-advanced setitem will not set_value back * fix multi output in tensor_array_pir * only in dynamic mode * add only advanced-setting case to fix coverage * fast pass for bool setitem (#61021) * fast pass for bool setitem * fix 0-size value case * remove final set_value OP in combined-indexing setting (#60983) * remove setvalue in combined indexing set * using combined-setting case to test * Optimize index put preprocess (#61060) * reduce vector operations when no bool index * reduce vector in index_put * reduce vector operations * no change for value * fix shape error in combine-getitem (#61922)

Author: zoooo0820

paddle/fluid/pybind/eager_method.cc

@@ -54,6 +54,7 @@ typedef SSIZE_T ssize_t;
 #pragma GCC diagnostic ignored "-Wmissing-field-initializers"
 #include "paddle/common/ddim.h"
 #include "paddle/fluid/eager/amp_utils.h"
+#include "paddle/fluid/eager/api/generated/eager_generated/backwards/nodes.h"
 #include "paddle/fluid/eager/api/generated/eager_generated/forwards/dygraph_functions.h"
 #include "paddle/fluid/eager/eager_amp_auto_cast.h"
 #include "paddle/fluid/framework/python_headers.h"
@@ -1359,6 +1360,7 @@ static PyObject* tensor__getitem_dygraph(TensorObject* self,
  &use_strided_slice);
 
  // step2: Dealing with basic indexing
+ bool out_is_view = false;
  auto out = getTensorWithBasicIndexing(tensor,
  &slice_axes,
  &slice_starts,
@@ -1367,7 +1369,8 @@ static PyObject* tensor__getitem_dygraph(TensorObject* self,
  &decrease_axis,
  &none_axes,
  &infer_flags,
- &use_strided_slice);
+ &use_strided_slice,
+ &out_is_view);
 
  if (!has_advanced_index) {
  return ToPyObject(out);
@@ -1386,7 +1389,8 @@ static PyObject* tensor__getitem_dygraph(TensorObject* self,
  &trans_back_dim,
  &pos_of_new_dim,
  &rank_of_new_dim,
- &trans_dim);
+ &trans_dim,
+ &out_is_view);
 
  if (transed_index.size() == 1 &&
  transed_index[0].dtype() == phi::DataType::BOOL) {
@@ -1416,14 +1420,14 @@ static PyObject* tensor__getitem_dygraph(TensorObject* self,
 
  if (pos_of_new_dim != 0) {
  std::vector<int> perm(out.shape().size(), 0);
- int tmp1 = pos_of_new_dim, tmp2 = 0,
+ int tmp1 = rank_of_new_dim, tmp2 = 0,
  tmp3 = pos_of_new_dim + rank_of_new_dim;
  for (int i = 0; i < static_cast<int>(out.shape().size()); ++i) {
- if (i < rank_of_new_dim) {
+ if (i < pos_of_new_dim) {
  perm[i] =
- tmp1++; // range(pos_of_new_dim, pos_of_new_dim + rank_of_new_dim)
- } else if (i >= rank_of_new_dim && i < pos_of_new_dim + rank_of_new_dim) {
- perm[i] = tmp2++; // range(0, pos_of_new_dim)
+ tmp1++; // range(rank_of_new_dim, pos_of_new_dim + rank_of_new_dim)
+ } else if (i >= pos_of_new_dim && i < pos_of_new_dim + rank_of_new_dim) {
+ perm[i] = tmp2++; // range(0, rank_of_new_dim)
  } else {
  perm[i] = tmp3++; // range(pos_of_new_dim + rank_of_new_dim, out.ndim)
  }
@@ -1681,6 +1685,7 @@ static PyObject* tensor__setitem_dygraph(TensorObject* self,
  // 3. assign values to the sliced result by index_put OP;
  // 4. transpose back and assign the result to original tensor by set_value
  // OP.
+ bool out_is_view = false;
  paddle::Tensor sub_tensor = getTensorWithBasicIndexing(tensor,
  &slice_axes,
  &slice_starts,
@@ -1689,7 +1694,8 @@ static PyObject* tensor__setitem_dygraph(TensorObject* self,
  &decrease_axis,
  &none_axes,
  &infer_flags,
- &use_strided_slice);
+ &use_strided_slice,
+ &out_is_view);
 
  std::vector<paddle::Tensor> transed_index;
  std::vector<int> trans_back_dim, trans_dim;
@@ -1705,65 +1711,126 @@ static PyObject* tensor__setitem_dygraph(TensorObject* self,
  &trans_back_dim,
  &pos_of_new_dim,
  &rank_of_new_dim,
- &trans_dim);
+ &trans_dim,
+ &out_is_view);
 
  // Release gil and do tracing
  py::gil_scoped_release release;
-
- if (value_tensor.initialized() &&
- (self->tensor.dtype() != value_tensor.dtype())) {
- if (egr::Controller::Instance().GetAMPLevel() !=
- paddle::imperative::AmpLevel::O0) {
- paddle::small_vector<std::vector<paddle::Tensor>,
- egr::kSlotSmallVectorSize>
- tmps = {{self->tensor}, {value_tensor}};
- auto amp_dtype = egr::GetAmpDestDtype("index_put", tmps);
- self->tensor = egr::EagerAmpAutoCast(
- self->tensor.name(), self->tensor, amp_dtype, "index_put");
- value_tensor = egr::EagerAmpAutoCast(
- value_tensor.name(), value_tensor, amp_dtype, "index_put");
- }
+ if (value_tensor.initialized()) {
  if (self->tensor.dtype() != value_tensor.dtype()) {
- value_tensor = cast_ad_func(value_tensor, self->tensor.dtype());
+ if (egr::Controller::Instance().GetAMPLevel() !=
+ paddle::imperative::AmpLevel::O0) {
+ paddle::small_vector<std::vector<paddle::Tensor>,
+ egr::kSlotSmallVectorSize>
+ tmps = {{self->tensor}, {value_tensor}};
+ auto amp_dtype = egr::GetAmpDestDtype("index_put", tmps);
+ self->tensor = egr::EagerAmpAutoCast(
+ self->tensor.name(), self->tensor, amp_dtype, "index_put");
+ value_tensor = egr::EagerAmpAutoCast(
+ value_tensor.name(), value_tensor, amp_dtype, "index_put");
+ }
+ if (self->tensor.dtype() != value_tensor.dtype()) {
+ value_tensor = cast_ad_func(value_tensor, self->tensor.dtype());
+ }
  }
- }
 
- if (value_tensor.dims().size() > 1 && pos_of_new_dim != 0) {
- value_tensor = transpose_ad_func(value_tensor, trans_dim);
- }
+  if (value_tensor.dims().size() > 1 && pos_of_new_dim != 0) {
+  value_tensor = transpose_ad_func(value_tensor, trans_dim);
+  }
 
- // TODO(zoooo0820) 1.Using inplace version index_put
- // 2.Remove following code after backward bug fixed.
- transed_sub_tensor = assign_ad_func(transed_sub_tensor);
+ const phi::distributed::ProcessMesh* mesh = nullptr;
+ if (InputsContainDistTensor(
+ &mesh, self->tensor, transed_sub_tensor, value_tensor)) {
+ ConvertAllInputsToDistTensor(
+ mesh, self->tensor, transed_sub_tensor, value_tensor);
+ }
 
- const phi::distributed::ProcessMesh* mesh = nullptr;
- if (InputsContainDistTensor(
- &mesh, self->tensor, transed_sub_tensor, value_tensor)) {
- ConvertAllInputsToDistTensor(
- mesh, self->tensor, transed_sub_tensor, value_tensor);
- }
+ if (transed_index.size() == 1 &&
+ transed_index[0].dtype() == phi::DataType::BOOL &&
+ transed_index[0].shape().size() == self->tensor.shape().size()) {
+ if (value_tensor.shape() != self->tensor.shape()) {
+ value_tensor = expand_ad_func(value_tensor, self->tensor.shape());
+ }
+ transed_sub_tensor =
+ where__ad_func(logical_not_ad_func(transed_index[0]),
+ transed_sub_tensor,
+ value_tensor);
+ } else {
+ transed_sub_tensor =
+ index_put__ad_func(transed_sub_tensor, transed_index, value_tensor);
+ }
 
- transed_sub_tensor =
- index_put_ad_func(transed_sub_tensor, transed_index, value_tensor);
-
- paddle::Tensor transback_sub_tensor =
- transpose_ad_func(transed_sub_tensor, trans_back_dim);
-
- self->tensor = set_value_with_tensor__ad_func(self->tensor,
- transback_sub_tensor,
- slice_starts,
- slice_ends,
- slice_strides,
- slice_axes,
- decrease_axis,
- none_axes);
- if (PyCheckTensor(value_obj)) {
- // pass the stop_gradient from value to tensor.
- // pass stop gradient should be done after CheckInplace in
- // set_value__dygraph_function.
- if (!egr::EagerUtils::autograd_meta(&value_tensor)->StopGradient() &&
- egr::EagerUtils::autograd_meta(&self->tensor)->StopGradient()) {
- egr::EagerUtils::autograd_meta(&self->tensor)->SetStopGradient(false);
+ if (out_is_view) {
+ // NOTE(zoooo0820): if out_is_view is true, it is a case of
+ // combined-indexing setitem, i.e. firstly we get a view of
+ // self->tensor, then modified it with inplace api index_put_ For now,
+ // in design of Paddle, the forward result is right. But the backward
+ // edge can not be established because the Base Tensor cannot sense
+ // whether it has been modified by other operations. Following codes are
+ // to add a new node (set_value_with_tensor_grad) to record the backward
+ // edge, with out ad_function which needs to do the forward calculation.
+
+ egr::AutogradMeta* x_autograd_meta =
+ egr::EagerUtils::nullable_autograd_meta(self->tensor);
+ egr::AutogradMeta* values_autograd_meta =
+ egr::EagerUtils::nullable_autograd_meta(transed_sub_tensor);
+ bool trace_backward = egr::Controller::Instance().HasGrad();
+ bool require_any_grad = egr::EagerUtils::ComputeRequireGrad(
+ trace_backward, x_autograd_meta, values_autograd_meta);
+ // Node Declaration
+ std::shared_ptr<SetValueWithTensorGradNode> grad_node;
+ // Set grad_node before API Call
+ if (require_any_grad) {
+ paddle::Tensor transback_sub_tensor =
+ transpose_ad_func(transed_sub_tensor, trans_back_dim);
+ const auto& values_tmp =
+ (require_any_grad && transback_sub_tensor.is_dense_tensor() &&
+ !std::dynamic_pointer_cast<phi::DenseTensor>(
+ transback_sub_tensor.impl())
+ ->meta()
+ .is_contiguous())
+ ? paddle::Tensor(
+ std::make_shared<phi::DenseTensor>(
+ std::move(paddle::experimental::Trans2Contiguous(
+ *(std::dynamic_pointer_cast<phi::DenseTensor>(
+ transback_sub_tensor.impl()))))),
+ transback_sub_tensor.mutable_autograd_meta())
+ : transback_sub_tensor;
+
+ grad_node = std::shared_ptr<SetValueWithTensorGradNode>(
+ new SetValueWithTensorGradNode(1, 2)); // NOLINT
+ grad_node->SetAttributestarts(slice_starts);
+ grad_node->SetAttributeends(slice_ends);
+ grad_node->SetAttributesteps(slice_strides);
+ grad_node->SetAttributeaxes(slice_axes);
+ grad_node->SetAttributedecrease_axes(decrease_axis);
+ grad_node->SetAttributenone_axes(none_axes);
+ grad_node->SetTensorWrappervalues(values_tmp);
+
+ paddle::memory::LogDeviceMemoryStats(
+ egr::Controller::Instance().GetExpectedPlace(),
+ "set_value_with_tensor");
+ egr::EagerUtils::CheckInplace(
+ self->tensor, x_autograd_meta, require_any_grad);
+ egr::EagerUtils::PassStopGradient(false, x_autograd_meta);
+ // SetGradOutMeta & SetEdges
+ grad_node->SetGradOutMeta(self->tensor, 0);
+ grad_node->SetGradOutMeta(transback_sub_tensor, 1);
+ if (x_autograd_meta) {
+ egr::EagerUtils::SetOutRankWithSlot(x_autograd_meta, 0);
+ egr::EagerUtils::SetHistory(x_autograd_meta, grad_node);
+ }
+ grad_node->SetGradInMeta(self->tensor, 0);
+ }
+ }
+ if (PyCheckTensor(value_obj)) {
+ // pass the stop_gradient from value to tensor.
+ // pass stop gradient should be done after CheckInplace in
+ // set_value__dygraph_function.
+ if (!egr::EagerUtils::autograd_meta(&value_tensor)->StopGradient() &&
+ egr::EagerUtils::autograd_meta(&self->tensor)->StopGradient()) {
+ egr::EagerUtils::autograd_meta(&self->tensor)->SetStopGradient(false);
+ }
  }
  }
  }

paddle/fluid/pybind/slice_utils.h

@@ -347,11 +347,13 @@ static paddle::Tensor getTensorWithBasicIndexing(
  std::vector<int64_t>* decrease_axis,
  std::vector<int64_t>* none_axes,
  std::vector<int64_t>* infer_flags,
- bool* use_strided_slice) {
+ bool* use_strided_slice,
+ bool* out_is_view) {
  paddle::Tensor out;
  if (slice_axes->empty()) {
  out = tensor;
  } else {
+ *out_is_view = true;
  if (!(*use_strided_slice)) {
  eager_gil_scoped_release guard;
  out = slice_ad_func(tensor,
@@ -372,6 +374,7 @@ static paddle::Tensor getTensorWithBasicIndexing(
  }
  }
  if (!none_axes->empty()) {
+ *out_is_view = true;
  eager_gil_scoped_release guard;
  // Deal with cases that decrease_axes is not empty
  // For example:
@@ -400,7 +403,8 @@ static paddle::Tensor dealWithAdvancedIndex(
  std::vector<int>* trans_back_dim,
  int* pos_of_new_dim,
  int* rank_of_new_dim,
- std::vector<int>* trans_dim) {
+ std::vector<int>* trans_dim,
+ bool* out_is_view) {
  int p = 0;
  for (size_t i = 0; i < advanced_index_dim->size(); ++i) {
  auto index_dim = (*advanced_index_dim)[i];
@@ -443,6 +447,7 @@ static paddle::Tensor dealWithAdvancedIndex(
  if (original_dim_order == *trans_dim) {
  transed_tensor = tensor;
  } else {
+ *out_is_view = true;
  transed_tensor = transpose_ad_func(tensor, *trans_dim);
  }