Add ext_tensor.slice() API

paddle/fluid/extension/include/ext_tensor.h

@@ -88,10 +88,20 @@ class PD_DLL_DECL Tensor {
   /// It's usually used to set the input tensor data.
   /// \param PlaceType of target place, of which
   /// the tensor will copy to.
-
   template <typename T>
   Tensor copy_to(const PlaceType& place) const;
 
+  /// \brief Return a sub-tensor of the given tensor.
+  /// It is usually used to extract a sub-tensor (which supports
+  /// modifying the data of the original tensor) to perform further
+  /// operations.
+  /// \param begin_idx The index of the start row (inclusive) to slice.
+  ///                  The index number begins from 0.
+  /// \param end_idx  The index of the end row (exclusive) to slice.
+  ///                 The index number begins from begin_idx + 1.
+  /// \return The sliced tensor.
+  Tensor slice(const int64_t begin_idx, const int64_t end_idx) const;
+
   /// \brief Return the shape of the Tensor.
   std::vector<int64_t> shape() const;
 

paddle/fluid/extension/src/ext_tensor.cc

@@ -99,6 +99,21 @@ void GpuCopy(T *src, T *dst, PlaceType src_plc, PlaceType dst_plc,
   }                                                     \
   auto *tensor = static_cast<framework::LoDTensor *>(tensor_.get());
 
+#define GET_INNER_PLACE                               \
+  platform::Place place;                              \
+  switch (place_) {                                   \
+    case PlaceType::kCPU:                             \
+      place = platform::CPUPlace();                   \
+      break;                                          \
+    case PlaceType::kGPU:                             \
+      place = platform::CUDAPlace();                  \
+      break;                                          \
+    default:                                          \
+      PADDLE_THROW(platform::errors::Unavailable(     \
+          "Custom operator unsupported place id(%d)", \
+          static_cast<int>(place_)));                 \
+  }
+
 void Tensor::reshape(const std::vector<int64_t> &shape) {
   GET_CASTED_TENSOR
   auto new_dim = framework::make_ddim(shape);
@@ -221,6 +236,63 @@ Tensor Tensor::copy_to(const PlaceType &target_place) const {
   return target;
 }
 
+Tensor Tensor::slice(const int64_t begin_idx, const int64_t end_idx) const {
+  GET_CASTED_TENSOR
+  GET_INNER_PLACE
+  framework::Tensor intermediate = tensor->Slice(begin_idx, end_idx);
+  std::vector<int64_t> shape_vec =
+      framework::vectorize<int64_t>(tensor->dims());
+  shape_vec[0] = end_idx - begin_idx;
+  Tensor target = Tensor(place_, shape_vec);
+  auto *target_tensor =
+      static_cast<framework::LoDTensor *>(target.tensor_.get());
+  auto dtype = tensor->type();
+  switch (dtype) {
+    case framework::proto::VarType::FP32:
+      target_tensor->mutable_data<float>(place);
+      break;
+    case framework::proto::VarType::FP64:
+      target_tensor->mutable_data<double>(place);
+      break;
+    case framework::proto::VarType::INT64:
+      target_tensor->mutable_data<int64_t>(place);
+      break;
+    case framework::proto::VarType::INT32:
+      target_tensor->mutable_data<int32_t>(place);
+      break;
+    case framework::proto::VarType::INT16:
+      target_tensor->mutable_data<int16_t>(place);
+      break;
+    case framework::proto::VarType::INT8:
+      target_tensor->mutable_data<int8_t>(place);
+      break;
+    case framework::proto::VarType::UINT8:
+      target_tensor->mutable_data<uint8_t>(place);
+      break;
+    case framework::proto::VarType::BOOL:
+      target_tensor->mutable_data<bool>(place);
+      break;
+    case framework::proto::VarType::COMPLEX64:
+      target_tensor->mutable_data<complex64>(place);
+      break;
+    case framework::proto::VarType::COMPLEX128:
+      target_tensor->mutable_data<complex128>(place);
+      break;
+    case framework::proto::VarType::FP16:
+      target_tensor->mutable_data<float16>(place);
+      break;
+    // TODO(JiabinYang) support more data types if needed.
+    default:
+      PADDLE_THROW(platform::errors::Unavailable(
+          "Not supported VarType(%s) for slicing.",
+          ToString(framework::CustomTensorUtils::ConvertInnerDTypeToEnumDType(
+              dtype))));
+      break;
+  }
+  target_tensor->ShareDataWith(intermediate);
+  return target;
+}
+
 template PD_DLL_DECL Tensor
 Tensor::copy_to<float>(const PlaceType &target_place) const;
 template PD_DLL_DECL Tensor

paddle/fluid/framework/custom_tensor_test.cc

@@ -75,6 +75,36 @@ void TestAPISizeAndShape() {
   CHECK(t1.shape() == tensor_shape);
 }
 
+void TestAPISlice() {
+  std::vector<int64_t> tensor_shape_origin = {5, 5};
+  std::vector<int64_t> tensor_shape_sub1 = {3, 5};
+  std::vector<int64_t> tensor_shape_sub2 = {1, 5};
+#ifdef PADDLE_WITH_CUDA
+  auto t1 = paddle::Tensor(paddle::PlaceType::kGPU, tensor_shape_origin);
+  t1.mutable_data<float>();
+  CHECK(t1.slice(0, 5).shape() == tensor_shape_origin);
+  CHECK(t1.slice(0, 3).shape() == tensor_shape_sub1);
+  CHECK(t1.slice(4, 5).shape() == tensor_shape_sub2);
+#endif
+  auto t2 = paddle::Tensor(paddle::PlaceType::kCPU, tensor_shape_origin);
+  t2.mutable_data<float>();
+  CHECK(t2.slice(0, 5).shape() == tensor_shape_origin);
+  CHECK(t2.slice(0, 3).shape() == tensor_shape_sub1);
+  CHECK(t2.slice(4, 5).shape() == tensor_shape_sub2);
+
+  // Test writing function for sliced tensor
+  auto t = InitCPUTensorForTest<float>();
+  auto t_sliced = t.slice(0, 1);
+  auto* t_sliced_data_ptr = t_sliced.mutable_data<float>();
+  for (int64_t i = 0; i < t_sliced.size(); i++) {
+    t_sliced_data_ptr[i] += float(5);
+  }
+  auto* t_data_ptr = t.mutable_data<float>();
+  for (int64_t i = 0; i < t_sliced.size(); i++) {
+    CHECK_EQ(t_data_ptr[i], float(10));
+  }
+}
+
 template <typename T>
 paddle::DataType TestDtype() {
   std::vector<int64_t> tensor_shape = {5, 5};
@@ -244,6 +274,8 @@ TEST(CustomTensor, copyTest) {
   TestAPISizeAndShape();
   VLOG(2) << "TestPlace";
   TestAPIPlace();
+  VLOG(2) << "TestSlice";
+  TestAPISlice();
   VLOG(2) << "TestCast";
   GroupTestCast();
   VLOG(2) << "TestDtypeConvert";