optimize slice op and slice grad op (#32266)

* optimize slice op and slice grad op, test=develop

* optimize variable name and annotation information, test=develop

paddle/fluid/operators/slice_op.h

@@ -259,7 +259,20 @@ class SliceKernel : public framework::OpKernel<T> {
     auto out_t =
         framework::EigenTensor<T, D, Eigen::RowMajor, Eigen::DenseIndex>::From(
             *out, new_out_dims);
-    out_t.device(place) = in_t.slice(offsets, extents);
+
+    if (in->numel() <= Eigen::NumTraits<int>::highest()) {
+      // similar to tf.slice:
+      // if element number less than INT_MAX, change the type of index to int
+      Eigen::DSizes<int, D> offsets_32bit, extents_32bit;
+      for (size_t i = 0; i < D; i++) {
+        offsets_32bit[i] = offsets[i];
+        extents_32bit[i] = extents[i];
+      }
+      framework::To32BitIndex(out_t).device(place) =
+          framework::To32BitIndex(in_t).slice(offsets_32bit, extents_32bit);
+    } else {
+      out_t.device(place) = in_t.slice(offsets, extents);
+    }
 
     out->Resize(out_dims);
   }
@@ -300,8 +313,6 @@ class SliceGradKernel : public framework::OpKernel<T> {
  private:
   template <size_t D>
   void SliceCompute(const framework::ExecutionContext& context) const {
-    auto& place =
-        *context.template device_context<DeviceContext>().eigen_device();
     auto axes = context.Attr<std::vector<int>>("axes");
 
     auto starts_int = context.Attr<std::vector<int>>("starts");
@@ -435,13 +446,189 @@ class SliceGradKernel : public framework::OpKernel<T> {
       paddings[i].first = offsets[i];
       paddings[i].second = (in_dims[i] - out_dims[i]) - offsets[i];
     }
+    EigenPaddingCompute(context, d_input, in_dims, d_out, out_dims, paddings);
+  }
+
+  template <size_t D>
+  void EigenPaddingCompute(
+      const framework::ExecutionContext& context, framework::Tensor* d_input,
+      const framework::DDim& in_dims, const framework::Tensor* d_out,
+      const framework::DDim& out_dims,
+      const Eigen::array<std::pair<int64_t, int64_t>, D>& paddings) const {
+    if (D <= 3) {
+      // if dimension less than 3, cannot reduce dimension
+      LaunchEigenPadding(context, d_input, in_dims, d_out, out_dims, paddings);
+    } else {  // else we can reduce dimension
+      // count not-zero padding number, and record the dimension
+      int need_pad_num = 0, pad_dim = -1;
+      for (size_t i = 0; i < D; i++) {
+        if (paddings[i].first != 0 || paddings[i].second != 0) {
+          need_pad_num++;
+          pad_dim = i;
+        }
+      }
+
+      if (need_pad_num == 0) {
+        // do not need padding, pass if data address same, else copy
+        if (d_input->mutable_data<T>(context.GetPlace()) == d_out->data<T>()) {
+          // inplace, do not any operator, pass
+        } else {
+          framework::TensorCopy(
+              *d_out, context.GetPlace(),
+              context.template device_context<platform::DeviceContext>(),
+              d_input);
+        }
+      } else if (need_pad_num == 1) {
+        // only need padding one dimension, we can reduce dimension.
+        // only the padding dimension is available for us.
+        // How to reduce dimension(5 to 3 for example):
+        // before(D=5):
+        // in_dims:        [x1,  x2,  x3,  x4,  x5]
+        // padding.first:  [0,   0,   a,   0,  0]
+        // padding.second: [0,   0,   b,   0,  0]
+        //                     | |
+        //                     V V
+        // after(D=3):
+        // reshaped_in_dims:        [x1*x2,  x3,  x4*x5]
+        // reshaped_padding.first:  [0,      a,     0]
+        // reshaped_padding.second: [0,      b,     0]
+
+        if (pad_dim == D - 1) {
+          // only last dimension need padding,
+          // reshape the dimension of tensor in 2: [preceding, padding]
+          std::vector<int64_t> in_tore_shape(2, 1), out_tore_shape(2, 1);
+          Eigen::array<std::pair<int64_t, int64_t>, 2> reshaped_padding;
+
+          // first dimension is the accumulate of preceding dimension
+          for (int i = 0; i < pad_dim; i++) {
+            in_tore_shape[0] *= in_dims[i];
+            out_tore_shape[0] *= out_dims[i];
+          }
+          // second dimension is the padding dimension
+          in_tore_shape[1] = in_dims[pad_dim];
+          out_tore_shape[1] = out_dims[pad_dim];
+
+          // convert array from std::vector to DDim
+          framework::DDim reshaped_in_dims =
+              framework::make_ddim(in_tore_shape);
+          framework::DDim reshaped_out_dims =
+              framework::make_ddim(out_tore_shape);
+
+          // after reshape: the first dimension do not need padding,
+          // set padding[0] zero
+          reshaped_padding[0].first = reshaped_padding[0].second = 0;
+          // the second dimension is the previous padding dimension
+          reshaped_padding[1].first = paddings[pad_dim].first;
+          reshaped_padding[1].second = paddings[pad_dim].second;
+
+          LaunchEigenPadding(context, d_input, reshaped_in_dims, d_out,
+                             reshaped_out_dims, reshaped_padding);
+        } else if (pad_dim == 0) {
+          // only first dimension need padding,
+          // reshape the dimension of tensor in 2: [padding, succeeding]
+          // similar to (D - 1)
+          std::vector<int64_t> in_tore_shape(2, 1), out_tore_shape(2, 1);
+          Eigen::array<std::pair<int64_t, int64_t>, 2> reshaped_padding;
+
+          // first dimension is the padding dimension
+          in_tore_shape[0] = in_dims[pad_dim];
+          out_tore_shape[0] = out_dims[pad_dim];
+          // sencond dimension is the accumulate of succeeding dimension
+          for (size_t i = pad_dim + 1; i < D; i++) {
+            in_tore_shape[1] *= in_dims[i];
+            out_tore_shape[1] *= out_dims[i];
+          }
+
+          // convert array from std::vector to DDim
+          framework::DDim reshaped_in_dims =
+              framework::make_ddim(in_tore_shape);
+          framework::DDim reshaped_out_dims =
+              framework::make_ddim(out_tore_shape);
+
+          // after reshape:
+          // the first dimension is the previous padding dimension
+          reshaped_padding[0].first = paddings[pad_dim].first;
+          reshaped_padding[0].second = paddings[pad_dim].second;
+          // the second dimension do not need padding, set padding[1] zero
+          reshaped_padding[1].first = reshaped_padding[1].second = 0;
+
+          LaunchEigenPadding(context, d_input, reshaped_in_dims, d_out,
+                             reshaped_out_dims, reshaped_padding);
+        } else {
+          // other dimension need padding
+          // reshape the dimension of tensor in 3:
+          // [preceding, padding, succeeding]
+          std::vector<int64_t> in_tore_shape(3, 1), out_tore_shape(3, 1);
+          Eigen::array<std::pair<int64_t, int64_t>, 3> reshaped_padding;
+
+          // first dimension is the accumulate of preceding dimension
+          for (int i = 0; i < pad_dim; i++) {
+            in_tore_shape[0] *= in_dims[i];
+            out_tore_shape[0] *= out_dims[i];
+          }
+          // second dimension is the padding dimension
+          in_tore_shape[1] = in_dims[pad_dim];
+          out_tore_shape[1] = out_dims[pad_dim];
+          // third dimension is the accumulate of succeeding dimension
+          for (size_t i = pad_dim + 1; i < D; i++) {
+            in_tore_shape[2] *= in_dims[i];
+            out_tore_shape[2] *= out_dims[i];
+          }
+
+          // convert array from std::vector to DDim
+          framework::DDim reshaped_in_dims =
+              framework::make_ddim(in_tore_shape);
+          framework::DDim reshaped_out_dims =
+              framework::make_ddim(out_tore_shape);
+
+          // after reshape:
+          // the first dimension do not need padding, set padding[0] zero
+          reshaped_padding[0].first = reshaped_padding[2].second = 0;
+          // the second dimension is the previous padding dimension
+          reshaped_padding[1].first = paddings[pad_dim].first;
+          reshaped_padding[1].second = paddings[pad_dim].second;
+          // the third dimension do not need padding, set padding[2] zero
+          reshaped_padding[2].first = reshaped_padding[2].second = 0;
+
+          LaunchEigenPadding(context, d_input, reshaped_in_dims, d_out,
+                             reshaped_out_dims, reshaped_padding);
+        }
+      } else {
+        // need padding at many dimension, cannot reduce dimension
+        LaunchEigenPadding(context, d_input, in_dims, d_out, out_dims,
+                           paddings);
+      }
+    }
+  }
+
+  template <size_t D>
+  void LaunchEigenPadding(
+      const framework::ExecutionContext& context, framework::Tensor* d_input,
+      const framework::DDim& in_dims, const framework::Tensor* d_out,
+      const framework::DDim& out_dims,
+      const Eigen::array<std::pair<int64_t, int64_t>, D>& paddings) const {
+    auto& place =
+        *context.template device_context<DeviceContext>().eigen_device();
     auto d_in_t =
         framework::EigenTensor<T, D, Eigen::RowMajor, Eigen::DenseIndex>::From(
-            *d_input);
+            *d_input, in_dims);
     auto d_out_t =
         framework::EigenTensor<T, D, Eigen::RowMajor, Eigen::DenseIndex>::From(
             *d_out, out_dims);
-    d_in_t.device(place) = d_out_t.pad(paddings, T(0));
+
+    if (d_input->numel() <= Eigen::NumTraits<int>::highest()) {
+      // similar to tf.pad:
+      // if element number less than INT_MAX, change the type of index to int
+      Eigen::array<std::pair<int, int>, D> paddings_32bit;
+      for (size_t i = 0; i < D; i++) {
+        paddings_32bit[i] =
+            std::make_pair(paddings[i].first, paddings[i].second);
+      }
+      framework::To32BitIndex(d_in_t).device(place) =
+          framework::To32BitIndex(d_out_t).pad(paddings_32bit, T(0));
+    } else {
+      d_in_t.device(place) = d_out_t.pad(paddings, T(0));
+    }
   }
 };
 }  // namespace operators