[PIR slice] Optimize bool index logic for setitem and getitem (#72644)

changeyoung98 · Eddie-Wang1120 · web-flow · commit d0dc9c5be5cb · 2025-05-12T14:23:53.000+08:00
* add single bool index logic &amp;&amp; optimize set_tensor_value_op

* add single bool branch for getitem

---------

Co-authored-by: Eddie-Wang1120 &lt;wangjinheng1120@163.com&gt;
diff --git a/paddle/fluid/pybind/eager_method.cc b/paddle/fluid/pybind/eager_method.cc
@@ -1452,6 +1452,15 @@ static PyObject* tensor__getitem_dygraph(TensorObject* self,
   PyObject* _index = PyTuple_GET_ITEM(args, 0);
   VLOG(4) << "Call new indexing strategy _getitem_dygraph";
 
+  PyObject* index_ptr =
+      !PyTuple_Check(_index) ? PyTuple_Pack(1, _index) : _index;
+  DEFINE_PADDLE_SCOPE_GUARD([index_ptr, &_index]() {
+    if (!PyTuple_Check(_index)) {
+      Py_DECREF(index_ptr);
+      VLOG(4) << "Call Py_DECREF";
+    }
+  });
+
   // Note(0x45f): Using defined() instead of initialized()
   // to support slice tensor which shape like [0, 0, 0].
   PADDLE_ENFORCE_EQ(
@@ -1476,7 +1485,7 @@ static PyObject* tensor__getitem_dygraph(TensorObject* self,
 
   // step1: parsing the index and recording them
   ParseIndex(tensor,
-             _index,
+             index_ptr,
              &slice_axes,
              &slice_starts,
              &slice_ends,
@@ -1489,6 +1498,23 @@ static PyObject* tensor__getitem_dygraph(TensorObject* self,
              &has_advanced_index,
              &use_strided_slice);
 
+  // Special: Check if the index is single bool
+  if (PyTuple_GET_SIZE(_index) == 1 &&
+      PyBool_Check(PyTuple_GetItem(_index, 0))) {
+    if (PyTuple_GetItem(_index, 0) == Py_True) {
+      // unsqueeze the tensor to a new tensor with shape (1,)
+      paddle::Tensor out;
+      out.copy_(unsqueeze_ad_func(tensor, {0}), tensor.place(), false);
+      return ToPyObject(out);
+    } else {
+      // create a new tensor with shape (0,)
+      auto shape = tensor.shape();
+      shape.insert(shape.begin(), 0);
+      auto out = paddle::empty(shape, tensor.dtype(), tensor.place());
+      return ToPyObject(out);
+    }
+  }
+
   // step2: Dealing with basic indexing
   bool out_is_view = false;
   auto out = getTensorWithBasicIndexing(tensor,
@@ -1748,6 +1774,7 @@ static PyObject* tensor__setitem_dygraph(TensorObject* self,
             tensor.name()));
   }
   const int rank = tensor.shape().size();
+  const int size = PyTuple_GET_SIZE(index_ptr);
   std::vector<int> slice_starts, slice_ends, slice_strides;
   std::vector<int64_t> slice_axes, decrease_axis, infer_flags, none_axes;
 
@@ -1760,7 +1787,7 @@ static PyObject* tensor__setitem_dygraph(TensorObject* self,
 
   // step1: parsing the index and recording them
   ParseIndex(tensor,
-             _index,
+             index_ptr,
              &slice_axes,
              &slice_starts,
              &slice_ends,
@@ -1808,14 +1835,18 @@ static PyObject* tensor__setitem_dygraph(TensorObject* self,
       if (InputsContainDistTensor(&mesh, self->tensor, value_tensor)) {
         ConvertAllInputsToDistTensor(mesh, self->tensor, value_tensor);
       }
-      self->tensor = set_value_with_tensor__ad_func(self->tensor,
-                                                    value_tensor,
-                                                    slice_starts,
-                                                    slice_ends,
-                                                    slice_strides,
-                                                    slice_axes,
-                                                    decrease_axis,
-                                                    none_axes);
+      if (size == 1 && PyTuple_GetItem(index_ptr, 0) == Py_False) {
+        // do nothing
+      } else {
+        self->tensor = set_value_with_tensor__ad_func(self->tensor,
+                                                      value_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
@@ -1830,15 +1861,19 @@ static PyObject* tensor__setitem_dygraph(TensorObject* self,
       if (InputsContainDistTensor(&mesh, self->tensor)) {
         ConvertAllInputsToDistTensor(mesh, self->tensor);
       }
-      self->tensor = set_value__ad_func(self->tensor,
-                                        slice_starts,
-                                        slice_ends,
-                                        slice_strides,
-                                        slice_axes,
-                                        decrease_axis,
-                                        none_axes,
-                                        {1},
-                                        values);
+      if (size == 1 && PyTuple_GetItem(index_ptr, 0) == Py_False) {
+        // do nothing
+      } else {
+        self->tensor = set_value__ad_func(self->tensor,
+                                          slice_starts,
+                                          slice_ends,
+                                          slice_strides,
+                                          slice_axes,
+                                          decrease_axis,
+                                          none_axes,
+                                          {1},
+                                          values);
+      }
     }
   } else {
     // step3.2: Case for there are advanced indexing.
diff --git a/paddle/fluid/pybind/slice_utils.h b/paddle/fluid/pybind/slice_utils.h
@@ -171,7 +171,7 @@ static int _PySlice_GetIndices(PySliceObject* r,
 }
 
 static void ParseIndex(const paddle::Tensor& tensor,
-                       PyObject* _index,
+                       PyObject* index,
                        std::vector<int64_t>* slice_axes,
                        std::vector<int>* slice_starts,
                        std::vector<int>* slice_ends,
@@ -183,14 +183,6 @@ static void ParseIndex(const paddle::Tensor& tensor,
                        std::vector<paddle::Tensor>* advanced_index,
                        bool* has_advanced_index,
                        bool* use_strided_slice) {
-  // NOTE(zhiqiu): PyTuple_Pack increases refcount.
-  PyObject* index = !PyTuple_Check(_index) ? PyTuple_Pack(1, _index) : _index;
-  DEFINE_PADDLE_SCOPE_GUARD([index, _index]() {
-    if (!PyTuple_Check(_index)) {
-      Py_DECREF(index);
-      VLOG(4) << "Call Py_DECREF";
-    }
-  });
   // for case 0-size tensor in slice
   PADDLE_ENFORCE_EQ(
       tensor.defined(),
@@ -199,7 +191,11 @@ static void ParseIndex(const paddle::Tensor& tensor,
   const auto& shape = tensor.dims();
   const int rank = shape.size();
   const int size = PyTuple_GET_SIZE(index);
-
+  if (size == 1 && PyBool_Check(PyTuple_GetItem(index, 0))) {
+    // true and none using set_value full_set branch
+    // false do nothing
+    return;
+  }
   // Check Ellipsis is valid
   int specified_dims = 0;
   int ell_count = 0;
diff --git a/paddle/phi/kernels/impl/set_value_kernel_impl.h b/paddle/phi/kernels/impl/set_value_kernel_impl.h
@@ -86,6 +86,13 @@ void SetValueImpl(const Context& dev_ctx,
   std::vector<int64_t> starts_local = starts.GetData();
   std::vector<int64_t> ends_local = ends.GetData();
   std::vector<int64_t> steps_local = steps.GetData();
+  if (starts_local.empty() && ends_local.empty() && steps_local.empty() &&
+      axes.empty() && decrease_axes.empty() && none_axes.empty() &&
+      value.numel() == 1) {
+    ExpandKernel<T, Context>(
+        dev_ctx, value, IntArray{phi::vectorize<int64_t>(in.dims())}, out);
+    return;
+  }
   phi::funcs::CheckAndUpdateSliceAttrs(
       in_dims, axes, &starts_local, &ends_local, &steps_local);
   auto slice_dims = phi::funcs::GetSliceDims(
