Shared selected rows (#39608)

* merge legacy to fluid

* Remove legacy code

* Remove legacy code

* Remove DataType test

* Using Tensor directly instead of using EagerTensor

* support gradient_accumulation

* make test_imperative_lod_tensor_to_selected_rows longer

* make test_imperative_lod_tensor_to_selected_rows longer

* refine code

* Rename all EagerTensor to Tensor

* Rename some EagerTensor to Tensor

* rename EagerTensor to EagerVariable

* add more test

* Support copiable selected rows and merge develop

paddle/fluid/eager/eager_tensor.h

@@ -52,9 +52,9 @@ class EagerVariable final {
       : name_(tensor.name()) {
     if (tensor.defined()) {
       if (tensor.is_dense_tensor()) {
-        ConstructVariableFromTensor(tensor);
+        ConstructVariableFromTensor<pten::DenseTensor>(tensor);
       } else if (tensor.is_selected_rows()) {
-        ConstructVariableFromSelectedRows(tensor);
+        ConstructVariableFromTensor<pten::SelectedRows>(tensor);
       } else {
         PADDLE_THROW(paddle::platform::errors::Fatal(
             "Unrecognized egr::EagerVariable type, only "
@@ -71,9 +71,9 @@ class EagerVariable final {
     if (var_.IsInitialized()) {
       if (var_.IsType<paddle::framework::LoDTensor>() ||
           var_.IsType<paddle::framework::Tensor>()) {
-        return SetImplWithLegacyTensor();
+        return SetImplWithLegacyTensor<pten::DenseTensor>();
       } else if (var_.IsType<pten::SelectedRows>()) {
-        return SetImplWithLegacySelectedRows();
+        return SetImplWithLegacyTensor<pten::SelectedRows>();
       } else {
         PADDLE_THROW(paddle::platform::errors::Fatal(
             "Unable to fetch underlying tensor "
@@ -98,26 +98,18 @@ class EagerVariable final {
   void set_name(const std::string& name) { name_ = name; }
 
  private:
+  template <typename VarType>
   std::shared_ptr<pten::TensorBase> SetImplWithLegacyTensor() {
-    const auto& framework_tensor = var_.Get<pten::DenseTensor>();
+    const auto& framework_tensor = var_.Get<VarType>();
     VLOG(8) << "Sync Var to tensor for: " << name();
-    return std::make_shared<pten::DenseTensor>(framework_tensor);
-  }
-
-  std::shared_ptr<pten::TensorBase> SetImplWithLegacySelectedRows() {
-    auto* framework_tensor = var_.GetMutable<pten::SelectedRows>();
-    VLOG(8) << "Sync SelectedRows to tensor for: " << name();
-    auto res =
-        std::make_shared<pten::SelectedRows>(std::move(*framework_tensor));
-    var_.Clear();
-    return res;
+    return std::make_shared<VarType>(framework_tensor);
   }
 
+  template <typename VarType>
   void ConstructVariableFromTensor(const paddle::experimental::Tensor& tensor) {
-    auto* framework_tensor = var_.GetMutable<pten::DenseTensor>();
+    auto* framework_tensor = var_.GetMutable<VarType>();
     // Contruct framework::Tensor from egr::EagerVariable
-    auto tensor_dense =
-        std::dynamic_pointer_cast<pten::DenseTensor>(tensor.impl());
+    auto tensor_dense = std::dynamic_pointer_cast<VarType>(tensor.impl());
     PADDLE_ENFORCE_EQ(
         (tensor_dense.get() && tensor_dense), true,
         paddle::platform::errors::Fatal(
@@ -128,22 +120,6 @@ class EagerVariable final {
     *framework_tensor = *tensor_dense;
   }
 
-  void ConstructVariableFromSelectedRows(
-      const paddle::experimental::Tensor& tensor) {
-    auto* framework_tensor = var_.GetMutable<pten::SelectedRows>();
-    // Contruct framework::Tensor from egr::EagerVariable
-    auto tensor_dense =
-        std::dynamic_pointer_cast<pten::SelectedRows>(tensor.impl());
-    PADDLE_ENFORCE_EQ(
-        (tensor_dense.get() && tensor_dense), true,
-        paddle::platform::errors::Fatal(
-            "Tensor %s does not hold pten::SelectedRows or pten::DenseTensor. "
-            "Or it holds empty impl, this should not happend since we should "
-            "treat all kinds of tensor as what they are.",
-            tensor.name()));
-    *framework_tensor = std::move(*tensor_dense);
-  }
-
  private:
   std::string name_{""};
   paddle::framework::Variable var_;

paddle/pten/core/CMakeLists.txt

@@ -23,7 +23,7 @@ cc_library(sparse_csr_tensor SRCS sparse_csr_tensor.cc DEPS dense_tensor tensor_
 
 cc_library(meta_tensor SRCS meta_tensor.cc DEPS tensor_base tensor_meta dense_tensor)
 cc_library(infermeta_utils SRCS infermeta_utils.cc DEPS meta_tensor)
-cc_library(selected_rows SRCS selected_rows.cc DEPS dense_tensor mixed_vector pten_enforce ddim)
+cc_library(selected_rows SRCS selected_rows_impl.cc DEPS dense_tensor mixed_vector pten_enforce ddim)
 
 # Will remove once we implemented MKLDNN_Tensor
 if(WITH_MKLDNN)

paddle/pten/core/selected_rows.h

@@ -21,14 +21,8 @@ limitations under the License. */
 #include <utility>
 #include <vector>
 
-#include "paddle/pten/common/place.h"
-#include "paddle/pten/core/ddim.h"
-#include "paddle/pten/core/dense_tensor.h"
-#include "paddle/pten/core/enforce.h"
-#include "paddle/pten/core/utils/rw_lock.h"
-
-// See Note [ Why still include the fluid headers? ]
-#include "paddle/fluid/framework/mixed_vector.h"
+#include "paddle/pten/core/selected_rows_impl.h"
+
 namespace pten {
 class SelectedRows : public TensorBase,
                      public TypeInfoTraits<TensorBase, SelectedRows> {
@@ -49,31 +43,28 @@ class SelectedRows : public TensorBase,
    */
  public:
   SelectedRows(const std::vector<int64_t>& rows, const int64_t& height)
-      : rows_(rows), height_(height) {
-    value_.reset(new DenseTensor());
-    rwlock_.reset(new RWLock);
-  }
+      : impl_(std::make_shared<pten::SelectedRowsImpl>(rows, height)) {}
 
-  SelectedRows() {
-    height_ = 0;
-    value_.reset(new DenseTensor());
-    rwlock_.reset(new RWLock);
-  }
+  SelectedRows() : impl_(std::make_shared<pten::SelectedRowsImpl>()) {}
 
-  const DenseTensor& value() const { return *value_; }
+  const DenseTensor& value() const { return impl_->value(); }
 
-  DenseTensor* mutable_value() { return value_.get(); }
+  DenseTensor* mutable_value() { return impl_->mutable_value(); }
 
-  int64_t height() const { return height_; }
+  int64_t height() const { return impl_->height(); }
 
-  void set_height(int64_t height) { height_ = height; }
+  void set_height(int64_t height) { impl_->set_height(height); }
 
-  const paddle::framework::Vector<int64_t>& rows() const { return rows_; }
+  const paddle::framework::Vector<int64_t>& rows() const {
+    return impl_->rows();
+  }
 
-  paddle::framework::Vector<int64_t>* mutable_rows() { return &rows_; }
+  paddle::framework::Vector<int64_t>* mutable_rows() {
+    return impl_->mutable_rows();
+  }
 
   void set_rows(const paddle::framework::Vector<int64_t>& rows) {
-    rows_ = rows;
+    impl_->set_rows(rows);
   }
 
   /*
@@ -81,21 +72,14 @@ class SelectedRows : public TensorBase,
    *
    * @return -1 if the key does not exists.
    */
-  int64_t Index(int64_t key) const {
-    auto it = std::find(rows_.begin(), rows_.end(), key);
-    if (it == rows_.end()) {
-      PADDLE_THROW(paddle::platform::errors::NotFound(
-          "Input id (%lld) is not in current rows table.", key));
-    }
-    return static_cast<int64_t>(std::distance(rows_.begin(), it));
-  }
+  int64_t Index(int64_t key) const { return impl_->Index(key); }
 
   /*
    * @brief whether has the specified key in the table.
    *
    * @return true if the key is exists.
    */
-  bool HasKey(int64_t key) const;
+  bool HasKey(int64_t key) const { return impl_->HasKey(key); }
 
   /*
    * @brief Get value by the key list.
@@ -109,11 +93,15 @@ class SelectedRows : public TensorBase,
   void Get(const DenseTensor& ids,
            DenseTensor* value,
            bool auto_grown = false,
-           bool is_test = false);
+           bool is_test = false) {
+    impl_->Get(ids, value, auto_grown, is_test);
+  }
 
   void* AllocateFrom(Allocator* allocator,
                      DataType dtype,
-                     size_t requested_size = 0) override;
+                     size_t requested_size = 0) override {
+    return impl_->AllocateFrom(allocator, dtype, requested_size);
+  }
 
   /*
    * @brief Get the index of the key from id_to_index_ map. If the key not
@@ -126,28 +114,23 @@ class SelectedRows : public TensorBase,
    *
    * @return index of the key.
    */
-  int64_t AutoGrownIndex(int64_t key, bool auto_grown, bool is_test = false);
+  int64_t AutoGrownIndex(int64_t key, bool auto_grown, bool is_test = false) {
+    return impl_->AutoGrownIndex(key, auto_grown, is_test);
+  }
 
   /*
    * @brief Get the index of the key from id_to_index_ map.
    */
   inline int64_t GetIndexFromId(int64_t key) const {
-    auto iter = id_to_index_.find(key);
-    if (iter == id_to_index_.end()) {
-      return -1;
-    } else {
-      return iter->second;
-    }
+    return impl_->GetIndexFromId(key);
   }
 
-  void SyncIndex();
+  void SyncIndex() { impl_->SyncIndex(); }
   /*
    * @brief Get complete Dims before
    */
   pten::framework::DDim GetCompleteDims() const {
-    std::vector<int64_t> dims = vectorize(value_->dims());
-    dims[0] = height_;
-    return pten::framework::make_ddim(dims);
+    return impl_->GetCompleteDims();
   }
 
   /// \brief Returns the name of the class for type traits.
@@ -156,45 +139,37 @@ class SelectedRows : public TensorBase,
 
   /// \brief Returns the number of elements contained in tensor.
   /// \return The number of elements contained in tensor.
-  int64_t numel() const override { return value_->numel(); };
+  int64_t numel() const override { return impl_->numel(); };
 
   /// \brief Returns the dims of the tensor.
   /// \return The dims of the tensor.
   const DDim& dims() const noexcept override {
-    return value_->dims();
+    return impl_->dims();
     // return paddle::framework::make_ddim(dims);
   }
 
   /// \brief Returns the data type of the tensor.
   /// \return The data type of the tensor.
-  DataType dtype() const noexcept override { return value_->dtype(); }
+  DataType dtype() const noexcept override { return impl_->dtype(); }
 
   /// \brief Returns the data layout of the tensor.
   /// \return The data layout of the tensor.
-  DataLayout layout() const noexcept override { return value_->layout(); }
+  DataLayout layout() const noexcept override { return impl_->layout(); }
 
   /// \brief Returns the data place of the tensor.
   /// \return The data place of the tensor.
-  const Place& place() const override { return value_->place(); };
+  const Place& place() const override { return impl_->place(); };
 
   /// \brief Test whether the metadata is valid.
   /// \return Whether the metadata is valid.
-  bool valid() const noexcept override { return value_->valid(); }
+  bool valid() const noexcept override { return impl_->valid(); }
 
   /// \brief Test whether the storage is allocated.
   /// return Whether the storage is allocated.
-  bool initialized() const override { return value_->initialized(); }
+  bool initialized() const override { return impl_->initialized(); }
 
  private:
-  // Notice: rows can be duplicate. We can have {0, 4, 7, 0, 5, 7, 9} here.
-  // SelectedRows are simply concated when adding together. Until a
-  // SelectedRows add a Tensor, will the duplicate rows be handled.
-  paddle::framework::Vector<int64_t> rows_;
-  std::unordered_map<int64_t, int64_t>
-      id_to_index_;  // should not be used when rows_ has duplicate member
-  std::unique_ptr<DenseTensor> value_{nullptr};
-  int64_t height_;  // height indicates the underline tensor's height
-  std::unique_ptr<RWLock> rwlock_{nullptr};
+  std::shared_ptr<pten::SelectedRowsImpl> impl_{nullptr};
 };
 
 }  // namespace pten

paddle/pten/core/selected_rows.cc → paddle/pten/core/selected_rows_impl.cc

@@ -12,7 +12,7 @@ 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/pten/core/selected_rows.h"
+#include "paddle/pten/core/selected_rows_impl.h"
 
 #include "paddle/pten/core/utils/data_type.h"
 
@@ -95,20 +95,20 @@ struct TensorFillVisitor {
   int64_t size_;
 };
 
-void* SelectedRows::AllocateFrom(Allocator* allocator,
-                                 DataType dtype,
-                                 size_t requested_size) {
+void* SelectedRowsImpl::AllocateFrom(Allocator* allocator,
+                                     DataType dtype,
+                                     size_t requested_size) {
   return value_->AllocateFrom(allocator, dtype, requested_size);
 }
 
-bool SelectedRows::HasKey(int64_t key) const {
+bool SelectedRowsImpl::HasKey(int64_t key) const {
   return std::find(rows_.begin(), rows_.end(), key) == rows_.end() ? false
                                                                    : true;
 }
 
-int64_t SelectedRows::AutoGrownIndex(int64_t key,
-                                     bool auto_grown,
-                                     bool is_test) {
+int64_t SelectedRowsImpl::AutoGrownIndex(int64_t key,
+                                         bool auto_grown,
+                                         bool is_test) {
   if (is_test) {
     auto iter = id_to_index_.find(key);
     if (iter == id_to_index_.end()) {
@@ -164,7 +164,7 @@ int64_t SelectedRows::AutoGrownIndex(int64_t key,
   }
 }
 
-void SelectedRows::SyncIndex() {
+void SelectedRowsImpl::SyncIndex() {
   rwlock_->WRLock();
   id_to_index_.clear();
   for (size_t i = 0; i < rows_.size(); ++i) {
@@ -173,10 +173,10 @@ void SelectedRows::SyncIndex() {
   rwlock_->UNLock();
 }
 
-void SelectedRows::Get(const pten::DenseTensor& ids,
-                       pten::DenseTensor* value,
-                       bool auto_grown,
-                       bool is_test) {
+void SelectedRowsImpl::Get(const pten::DenseTensor& ids,
+                           pten::DenseTensor* value,
+                           bool auto_grown,
+                           bool is_test) {
   PADDLE_ENFORCE_EQ(value->IsInitialized(),
                     true,
                     paddle::platform::errors::InvalidArgument(

paddle/pten/core/selected_rows_impl.h

+/* 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 <algorithm>
+#include <memory>
+#include <mutex>  // NOLINT
+#include <unordered_map>
+#include <utility>
+#include <vector>
+
+#include "paddle/pten/common/place.h"
+#include "paddle/pten/core/ddim.h"
+#include "paddle/pten/core/dense_tensor.h"
+#include "paddle/pten/core/enforce.h"
+#include "paddle/pten/core/utils/rw_lock.h"
+
+// See Note [ Why still include the fluid headers? ]
+#include "paddle/fluid/framework/mixed_vector.h"
+namespace pten {
+class SelectedRowsImpl {
+  /*
+   * @brief We can use the SelectedRowsImpl structure to reproduce a sparse
+   * table.
+   *  A sparse table is a key-value structure that the key is an `int64_t`,
+   *  and the value is a Tensor which the first dimension is 0.
+   *  You can use the following interface to operate the sparse table, and you
+   * can find
+   *  some detail information from the comments of each interface:
+   *
+   *  HasKey(key), whether the sparse table has the specified key.
+   *  Set(key, value), set a key-value pair into the sparse table.
+   *  Get(keys, value*), get value by given key list and apply it to the given
+   * value pointer
+   *    with the specified offset.
+   *
+   */
+ public:
+  SelectedRowsImpl(const std::vector<int64_t>& rows, const int64_t& height)
+      : rows_(rows), height_(height) {
+    value_.reset(new DenseTensor());
+    rwlock_.reset(new RWLock);
+  }
+
+  SelectedRowsImpl() {
+    height_ = 0;
+    value_.reset(new DenseTensor());
+    rwlock_.reset(new RWLock);
+  }
+
+  const DenseTensor& value() const { return *value_; }
+
+  DenseTensor* mutable_value() { return value_.get(); }
+
+  int64_t height() const { return height_; }
+
+  void set_height(int64_t height) { height_ = height; }
+
+  const paddle::framework::Vector<int64_t>& rows() const { return rows_; }
+
+  paddle::framework::Vector<int64_t>* mutable_rows() { return &rows_; }
+
+  void set_rows(const paddle::framework::Vector<int64_t>& rows) {
+    rows_ = rows;
+  }
+
+  /*
+   * @brief Get the index of key in rows
+   *
+   * @return -1 if the key does not exists.
+   */
+  int64_t Index(int64_t key) const {
+    auto it = std::find(rows_.begin(), rows_.end(), key);
+    if (it == rows_.end()) {
+      PADDLE_THROW(paddle::platform::errors::NotFound(
+          "Input id (%lld) is not in current rows table.", key));
+    }
+    return static_cast<int64_t>(std::distance(rows_.begin(), it));
+  }
+
+  /*
+   * @brief whether has the specified key in the table.
+   *
+   * @return true if the key is exists.
+   */
+  bool HasKey(int64_t key) const;
+
+  /*
+   * @brief Get value by the key list.
+   * Note!!! this interface is only used when selected_rows is used as
+   * parameters
+   * for distribute lookup table.
+   *
+   * @return a list of pair which contains the non-exists key and the index in
+   * the value
+   */
+  void Get(const DenseTensor& ids,
+           DenseTensor* value,
+           bool auto_grown = false,
+           bool is_test = false);
+
+  void* AllocateFrom(Allocator* allocator,
+                     DataType dtype,
+                     size_t requested_size = 0);
+
+  /*
+   * @brief Get the index of the key from id_to_index_ map. If the key not
+   * exist,
+   * add the key into id_to_index_.
+   *
+   * Note!!! this interface is only used when selected_rows is used as
+   * parameters
+   * for distribute lookup table.
+   *
+   * @return index of the key.
+   */
+  int64_t AutoGrownIndex(int64_t key, bool auto_grown, bool is_test = false);
+
+  /*
+   * @brief Get the index of the key from id_to_index_ map.
+   */
+  inline int64_t GetIndexFromId(int64_t key) const {
+    auto iter = id_to_index_.find(key);
+    if (iter == id_to_index_.end()) {
+      return -1;
+    } else {
+      return iter->second;
+    }
+  }
+
+  void SyncIndex();
+  /*
+   * @brief Get complete Dims before
+   */
+  pten::framework::DDim GetCompleteDims() const {
+    std::vector<int64_t> dims = vectorize(value_->dims());
+    dims[0] = height_;
+    return pten::framework::make_ddim(dims);
+  }
+
+  /// \brief Returns the number of elements contained in tensor.
+  /// \return The number of elements contained in tensor.
+  int64_t numel() const { return value_->numel(); }
+
+  /// \brief Returns the dims of the tensor.
+  /// \return The dims of the tensor.
+  const DDim& dims() const noexcept {
+    return value_->dims();
+    // return paddle::framework::make_ddim(dims);
+  }
+
+  /// \brief Returns the data type of the tensor.
+  /// \return The data type of the tensor.
+  DataType dtype() const noexcept { return value_->dtype(); }
+
+  /// \brief Returns the data layout of the tensor.
+  /// \return The data layout of the tensor.
+  DataLayout layout() const noexcept { return value_->layout(); }
+
+  /// \brief Returns the data place of the tensor.
+  /// \return The data place of the tensor.
+  const Place& place() const { return value_->place(); }
+
+  /// \brief Test whether the metadata is valid.
+  /// \return Whether the metadata is valid.
+  bool valid() const noexcept { return value_->valid(); }
+
+  /// \brief Test whether the storage is allocated.
+  /// return Whether the storage is allocated.
+  bool initialized() const { return value_->initialized(); }
+
+ private:
+  // Notice: rows can be duplicate. We can have {0, 4, 7, 0, 5, 7, 9} here.
+  // SelectedRowsImpl are simply concated when adding together. Until a
+  // SelectedRowsImpl add a Tensor, will the duplicate rows be handled.
+  paddle::framework::Vector<int64_t> rows_;
+  std::unordered_map<int64_t, int64_t>
+      id_to_index_;  // should not be used when rows_ has duplicate member
+  std::unique_ptr<DenseTensor> value_{nullptr};
+  int64_t height_;  // height indicates the underline tensor's height
+  std::unique_ptr<RWLock> rwlock_{nullptr};
+};
+
+}  // namespace pten