[PTen] Support SelectedRows in execution and remove scale OpKernel an…

…d InferShape (#39351)

* adapt selectedrows in execution

* impl selected rows branch

* support selectedrow in infershape utils

* fix device compile failed

* fix new exe test failed

* revert some changes

paddle/fluid/eager/tests/performance_tests/benchmark_eager_cpu.cc

@@ -175,7 +175,7 @@ TEST(Benchmark, EagerIntermediateMLPCPU) {
   }
 }
 
-USE_OP(scale);
+USE_OP_ITSELF(scale);
 USE_OP(elementwise_add);
 USE_OP(matmul_v2);
 USE_OP(reduce_sum);

paddle/fluid/eager/tests/performance_tests/benchmark_eager_cuda.cc

@@ -185,7 +185,7 @@ TEST(Benchmark, EagerIntermediateMLPCUDA) {
   }
 }
 
-USE_OP(scale);
+USE_OP_ITSELF(scale);
 USE_OP(matmul_v2);
 USE_OP(reduce_sum);
 USE_OP(reduce_sum_grad);

paddle/fluid/eager/tests/performance_tests/benchmark_fluid_cpu.cc

@@ -211,7 +211,7 @@ TEST(Benchmark, FluidMLPCPU) {
 }  // namespace imperative
 }  // namespace paddle
 
-USE_OP(scale);
+USE_OP_ITSELF(scale);
 USE_OP(elementwise_add);
 USE_OP(matmul_v2);
 USE_OP(reduce_sum);

paddle/fluid/eager/tests/performance_tests/benchmark_fluid_cuda.cc

@@ -245,7 +245,7 @@ TEST(Benchmark, FluidMLPCUDA) {
 }  // namespace imperative
 }  // namespace paddle
 
-USE_OP(scale);
+USE_OP_ITSELF(scale);
 USE_OP(matmul_v2);
 USE_OP(reduce_sum);
 USE_OP(reduce_sum_grad);

paddle/fluid/framework/heter_pipeline_trainer_test.cc

@@ -26,7 +26,7 @@
 #define _LINUX
 #endif
 
-USE_OP(scale);
+USE_OP_ITSELF(scale);
 USE_NO_KERNEL_OP(heter_listen_and_serv);
 namespace paddle {
 namespace framework {

paddle/fluid/framework/infershape_utils.cc

@@ -78,7 +78,6 @@ class InferShapeArgumentMappingContext : public pten::ArgumentMappingContext {
   const InferShapeContext& ctx_;
 };
 
-// TODO(chenweihang): Support SelectedRows later
 // TODO(chenweihang): Support TensorArray later
 class CompatMetaTensor : public pten::MetaTensor {
  public:
@@ -104,7 +103,14 @@ class CompatMetaTensor : public pten::MetaTensor {
   DDim dims() const override {
     if (is_runtime_) {
       auto* var = BOOST_GET_CONST(Variable*, var_);
-      return var->Get<LoDTensor>().dims();
+      if (var->IsType<pten::DenseTensor>()) {
+        return var->Get<pten::DenseTensor>().dims();
+      } else if (var->IsType<pten::SelectedRows>()) {
+        return var->Get<pten::SelectedRows>().dims();
+      } else {
+        PADDLE_THROW(platform::errors::Unimplemented(
+            "Currently, only can get dims from DenseTensor or SelectedRows."));
+      }
     } else {
       auto* var = BOOST_GET_CONST(VarDesc*, var_);
       return make_ddim(var->GetShape());
@@ -114,7 +120,14 @@ class CompatMetaTensor : public pten::MetaTensor {
   pten::DataType dtype() const override {
     if (is_runtime_) {
       auto* var = BOOST_GET_CONST(Variable*, var_);
-      return var->Get<LoDTensor>().dtype();
+      if (var->IsType<pten::DenseTensor>()) {
+        return var->Get<pten::DenseTensor>().dtype();
+      } else if (var->IsType<pten::SelectedRows>()) {
+        return var->Get<pten::SelectedRows>().dtype();
+      } else {
+        PADDLE_THROW(platform::errors::Unimplemented(
+            "Currently, only can get dtype from DenseTensor or SelectedRows."));
+      }
     } else {
       auto* var = BOOST_GET_CONST(VarDesc*, var_);
       return pten::TransToPtenDataType(var->GetDataType());
@@ -135,10 +148,16 @@ class CompatMetaTensor : public pten::MetaTensor {
   void set_dims(const DDim& dims) override {
     if (is_runtime_) {
       auto* var = BOOST_GET(Variable*, var_);
-      LoDTensor* tensor = var->GetMutable<LoDTensor>();
-      pten::DenseTensorUtils::GetMutableMeta(
-          static_cast<pten::DenseTensor*>(tensor))
-          ->dims = dims;
+      if (var->IsType<pten::DenseTensor>()) {
+        auto* tensor = var->GetMutable<pten::DenseTensor>();
+        pten::DenseTensorUtils::GetMutableMeta(tensor)->dims = dims;
+      } else if (var->IsType<pten::SelectedRows>()) {
+        auto* tensor = var->GetMutable<pten::SelectedRows>()->mutable_value();
+        pten::DenseTensorUtils::GetMutableMeta(tensor)->dims = dims;
+      } else {
+        PADDLE_THROW(platform::errors::Unimplemented(
+            "Currently, only can set dims from DenseTensor or SelectedRows."));
+      }
     } else {
       auto* var = BOOST_GET(VarDesc*, var_);
       var->SetShape(vectorize(dims));
@@ -148,10 +167,16 @@ class CompatMetaTensor : public pten::MetaTensor {
   void set_dtype(pten::DataType dtype) override {
     if (is_runtime_) {
       auto* var = BOOST_GET(Variable*, var_);
-      LoDTensor* tensor = var->GetMutable<LoDTensor>();
-      pten::DenseTensorUtils::GetMutableMeta(
-          static_cast<pten::DenseTensor*>(tensor))
-          ->dtype = dtype;
+      if (var->IsType<pten::DenseTensor>()) {
+        auto* tensor = var->GetMutable<pten::DenseTensor>();
+        pten::DenseTensorUtils::GetMutableMeta(tensor)->dtype = dtype;
+      } else if (var->IsType<pten::SelectedRows>()) {
+        auto* tensor = var->GetMutable<pten::SelectedRows>()->mutable_value();
+        pten::DenseTensorUtils::GetMutableMeta(tensor)->dtype = dtype;
+      } else {
+        PADDLE_THROW(platform::errors::Unimplemented(
+            "Currently, only can set dtype from DenseTensor or SelectedRows."));
+      }
     } else {
       auto* var = BOOST_GET(VarDesc*, var_);
       var->SetDataType(pten::TransToProtoVarType(dtype));
@@ -174,11 +199,14 @@ class CompatMetaTensor : public pten::MetaTensor {
   void share_lod(const MetaTensor& meta_tensor) override {
     if (is_runtime_) {
       auto* var = BOOST_GET(Variable*, var_);
-      LoDTensor* tensor = var->GetMutable<LoDTensor>();
-      pten::DenseTensorUtils::GetMutableMeta(
-          static_cast<pten::DenseTensor*>(tensor))
-          ->lod =
-          static_cast<const CompatMetaTensor&>(meta_tensor).GetRuntimeLoD();
+      if (var->IsType<pten::DenseTensor>()) {
+        auto* tensor = var->GetMutable<pten::DenseTensor>();
+        pten::DenseTensorUtils::GetMutableMeta(tensor)->lod =
+            static_cast<const CompatMetaTensor&>(meta_tensor).GetRuntimeLoD();
+      } else {
+        // NOTE(chenweihang): do nothing
+        // only LoDTensor need to share lod
+      }
     } else {
       auto* var = BOOST_GET(VarDesc*, var_);
       var->SetLoDLevel(static_cast<const CompatMetaTensor&>(meta_tensor)
@@ -191,19 +219,45 @@ class CompatMetaTensor : public pten::MetaTensor {
     set_dtype(meta_tensor.dtype());
     // VarDesc doesn't contains layout, so we cannot share layout
     // set_layout(meta_tensor.layout());
+
+    // special case 1: share lod of LoDTensor
     share_lod(meta_tensor);
+
+    // special case 2: share height and rows of SelectedRows in runtime
+    if (is_runtime_) {
+      auto* var = BOOST_GET(Variable*, var_);
+      if (var->IsType<pten::SelectedRows>()) {
+        auto* selected_rows = var->GetMutable<pten::SelectedRows>();
+        auto& input_selected_rows =
+            static_cast<const CompatMetaTensor&>(meta_tensor).GetSelectedRows();
+        selected_rows->set_rows(input_selected_rows.rows());
+        selected_rows->set_height(input_selected_rows.height());
+      }
+    }
   }
 
  private:
   const LoD& GetRuntimeLoD() const {
     auto* var = BOOST_GET_CONST(Variable*, var_);
     return var->Get<LoDTensor>().lod();
   }
+
   int32_t GetCompileTimeLoD() const {
     auto* var = BOOST_GET_CONST(VarDesc*, var_);
     return var->GetLoDLevel();
   }
 
+  const pten::SelectedRows& GetSelectedRows() const {
+    PADDLE_ENFORCE_EQ(is_runtime_, true,
+                      platform::errors::Unavailable(
+                          "Only can get Tensor from MetaTensor in rumtime."));
+    auto* var = BOOST_GET_CONST(Variable*, var_);
+    PADDLE_ENFORCE_EQ(var->IsType<pten::SelectedRows>(), true,
+                      platform::errors::Unavailable(
+                          "The Tensor in MetaTensor is not SelectedRows."));
+    return var->Get<pten::SelectedRows>();
+  }
+
   InferShapeVarPtr var_;
   bool is_runtime_;
 };

paddle/fluid/framework/ir/memory_optimize_pass/test_reference_count_pass_last_lived_ops.cc

@@ -21,7 +21,7 @@
 #include "paddle/fluid/framework/parallel_executor.h"
 #include "paddle/fluid/framework/program_desc.h"
 
-USE_OP(scale);
+USE_OP_ITSELF(scale);
 USE_OP(elementwise_mul);
 USE_OP(elementwise_add);
 USE_OP(elementwise_add_grad);

paddle/fluid/framework/new_executor/interpretercore.cc

@@ -393,7 +393,8 @@ void InterpreterCore::RunInstruction(const Instruction& instr_node) {
     platform::RecordEvent infershape_event("InferShape");
     // If it is OperatorBase, InferShape do nothing.
     if (op_with_kernel != nullptr)
-      op_with_kernel->InferShape(instr_node.InnerInferShapeContext().get());
+      op_with_kernel->Info().infer_shape_(
+          instr_node.InnerInferShapeContext().get());
   }
 
   if (op_with_kernel != nullptr &&

paddle/fluid/framework/operator.cc

@@ -1998,16 +1998,17 @@ void OperatorWithKernel::BuildPtenKernelContext(
     size_t end_idx = start_idx + ins_vector.size();
 
     for (size_t offset = 0; offset < ins_vector.size(); ++offset) {
-      const framework::Tensor* tensor_in = nullptr;
+      const pten::TensorBase* tensor_in = nullptr;
       auto* var = ins_vector[offset];
-      if (var->IsType<framework::LoDTensor>()) {
-        tensor_in = &(var->Get<framework::LoDTensor>());
+      if (var->IsType<pten::DenseTensor>()) {
+        tensor_in = &(var->Get<pten::DenseTensor>());
+      } else if (var->IsType<pten::SelectedRows>()) {
+        tensor_in = &(var->Get<pten::SelectedRows>());
       } else {
         PADDLE_THROW(platform::errors::Unimplemented(
             "Unsupported input `%s` type when call pt kernel.",
             framework::ToTypeName(var->Type())));
-      }  // TODO(zyfncg): Add support for SelectedRows
-
+      }
       pt_kernel_context->EmplaceBackInputWithoutSetRange(tensor_in);
     }
     pt_kernel_context->AssignInputRange(std::make_pair(start_idx, end_idx), i);
@@ -2021,17 +2022,20 @@ void OperatorWithKernel::BuildPtenKernelContext(
     size_t end_idx = start_idx + outs_vector.size();
 
     for (size_t offset = 0; offset < outs_vector.size(); ++offset) {
-      framework::Tensor* tensor_out = nullptr;
+      pten::TensorBase* tensor_out = nullptr;
       auto* var = outs_vector[offset];
-      if (var->template IsType<framework::LoDTensor>()) {
-        tensor_out = var->template GetMutable<framework::LoDTensor>();
+      if (var->template IsType<pten::DenseTensor>()) {
+        tensor_out = var->template GetMutable<pten::DenseTensor>();
+      } else if (var->template IsType<pten::SelectedRows>()) {
+        tensor_out = var->template GetMutable<pten::SelectedRows>();
       } else {
         PADDLE_THROW(platform::errors::Unimplemented(
             "Unsupported output `%s` type when call pt kernel.",
             framework::ToTypeName(var->Type())));
-      }  // TODO(zyfncg): Add support for SelectedRows
+      }
 
-      experimental::ResetTensorByArgDef(tensor_out, output_defs.at(i));
+      experimental::ResetTensorDtypeAndLayoutByArgDef(tensor_out,
+                                                      output_defs.at(i));
       SetAllocationForOutputTenosr(
           tensor_out, pten::TransToFluidPlace(output_defs.at(i).backend));
 

paddle/fluid/framework/pten_utils.cc

@@ -207,21 +207,40 @@ void InitDefaultKernelSignatureMap() {
   });
 }
 
-void SetAllocationForOutputTenosr(pten::DenseTensor* tensor,
+static void SetAllocationForUninitializedDenseTensor(
+    pten::DenseTensor* dense_tensor, const platform::Place& place) {
+  int dtype_size = dense_tensor->dtype() == DataType::UNDEFINED
+                       ? 0
+                       : experimental::SizeOf(dense_tensor->dtype());
+  int64_t numels = product(dense_tensor->dims());
+  numels = numels < 0 ? 0 : numels;
+  auto tmp_allocation_ptr = memory::Alloc(place, numels * dtype_size);
+  auto& deleter = tmp_allocation_ptr.get_deleter();
+  auto* allocation_ptr = tmp_allocation_ptr.release();
+  auto shared_allocation =
+      std::shared_ptr<pten::Allocation>(allocation_ptr, deleter);
+
+  dense_tensor->ResetHolder(shared_allocation);
+}
+
+void SetAllocationForOutputTenosr(pten::TensorBase* tensor,
                                   const platform::Place& place) {
-  if (!tensor->IsInitialized() || !(tensor->place() == place)) {
-    int dtype_size = tensor->dtype() == DataType::UNDEFINED
-                         ? 0
-                         : experimental::SizeOf(tensor->dtype());
-    int64_t numels = product(tensor->dims());
-    numels = numels < 0 ? 0 : numels;
-    auto tmp_allocation_ptr = memory::Alloc(place, numels * dtype_size);
-    auto& deleter = tmp_allocation_ptr.get_deleter();
-    auto* allocation_ptr = tmp_allocation_ptr.release();
-    auto shared_allocation =
-        std::shared_ptr<pten::Allocation>(allocation_ptr, deleter);
-
-    tensor->ResetHolder(shared_allocation);
+  if (pten::DenseTensor::classof(tensor)) {
+    auto* dense_tensor = static_cast<pten::DenseTensor*>(tensor);
+    if (!dense_tensor->IsInitialized() || !(dense_tensor->place() == place)) {
+      SetAllocationForUninitializedDenseTensor(dense_tensor, place);
+    }
+  } else if (pten::SelectedRows::classof(tensor)) {
+    auto* selected_rows = static_cast<pten::SelectedRows*>(tensor);
+    if (!selected_rows->value().IsInitialized() ||
+        !(selected_rows->place() == place)) {
+      SetAllocationForUninitializedDenseTensor(selected_rows->mutable_value(),
+                                               place);
+    }
+  } else {
+    PADDLE_THROW(platform::errors::Unimplemented(
+        "Unsupported tensor type is received when setting allocation for "
+        "output tensor."));
   }
 }
 

paddle/fluid/framework/pten_utils.h

@@ -63,7 +63,7 @@ class KernelArgsNameMaker {
 
 void InitDefaultKernelSignatureMap();
 
-void SetAllocationForOutputTenosr(pten::DenseTensor* tensor,
+void SetAllocationForOutputTenosr(pten::TensorBase* tensor,
                                   const platform::Place& place);
 
 // TODO(Wilber): support others device context.

paddle/fluid/imperative/prepared_operator.h

@@ -29,6 +29,9 @@
 #include "paddle/fluid/imperative/type_defs.h"
 #include "paddle/fluid/imperative/var_helper.h"
 
+#include "paddle/pten/core/dense_tensor.h"
+#include "paddle/pten/core/selected_rows.h"
+
 DECLARE_bool(use_mkldnn);
 
 namespace paddle {
@@ -262,7 +265,17 @@ void BuildDygraphPtenKernelContext(
     size_t end_idx = start_idx + ins_vector.size();
 
     for (size_t offset = 0; offset < ins_vector.size(); ++offset) {
-      const auto* tensor_in = GetTensorFromVar(ins_vector[offset]->Var());
+      const pten::TensorBase* tensor_in = nullptr;
+      auto& var = ins_vector[offset]->Var();
+      if (var.template IsType<pten::DenseTensor>()) {
+        tensor_in = &(var.template Get<pten::DenseTensor>());
+      } else if (var.template IsType<pten::SelectedRows>()) {
+        tensor_in = &(var.template Get<pten::SelectedRows>());
+      } else {
+        PADDLE_THROW(platform::errors::Unimplemented(
+            "Unsupported input `%s` type when call pt kernel.",
+            framework::ToTypeName(var.Type())));
+      }
       kernel_ctx->EmplaceBackInputWithoutSetRange(tensor_in);
     }
     kernel_ctx->AssignInputRange(std::make_pair(start_idx, end_idx), i);
@@ -287,17 +300,21 @@ void BuildDygraphPtenKernelContext(
         kernel_ctx->EmplaceBackOutputWithoutSetRange({nullptr});
         continue;
       }
+
+      pten::TensorBase* tensor_out = nullptr;
       auto* var = outs_vector[offset]->MutableVar();
-      framework::Tensor* tensor_out = nullptr;
-      if (var->template IsType<framework::LoDTensor>()) {
-        tensor_out = var->template GetMutable<framework::LoDTensor>();
+      if (var->template IsType<pten::DenseTensor>()) {
+        tensor_out = var->template GetMutable<pten::DenseTensor>();
+      } else if (var->template IsType<pten::SelectedRows>()) {
+        tensor_out = var->template GetMutable<pten::SelectedRows>();
       } else {
         PADDLE_THROW(platform::errors::Unimplemented(
             "Unsupported output `%s` type when call pt kernel.",
             framework::ToTypeName(var->Type())));
-      }  // TODO(zyfncg): Add support for SelectedRows
+      }
 
-      experimental::ResetTensorByArgDef(tensor_out, output_defs.at(i));
+      experimental::ResetTensorDtypeAndLayoutByArgDef(tensor_out,
+                                                      output_defs.at(i));
       framework::SetAllocationForOutputTenosr(
           tensor_out, pten::TransToFluidPlace(output_defs.at(i).backend));
 

paddle/fluid/operators/pscore/heter_listen_and_server_test.cc

@@ -33,7 +33,7 @@ using MultiVarMsg = ::paddle::distributed::MultiVariableMessage;
 using VarMsg = ::paddle::distributed::VariableMessage;
 DECLARE_double(eager_delete_tensor_gb);
 
-USE_OP(scale);
+USE_OP_ITSELF(scale);
 USE_NO_KERNEL_OP(heter_listen_and_serv);
 
 framework::BlockDesc* AppendSendAndRecvBlock(framework::ProgramDesc* program) {

paddle/fluid/operators/pscore/heter_server_test.cc

@@ -29,7 +29,7 @@ namespace distributed = paddle::distributed;
 using MultiVarMsg = ::paddle::distributed::MultiVariableMessage;
 using VarMsg = ::paddle::distributed::VariableMessage;
 
-USE_OP(scale);
+USE_OP_ITSELF(scale);
 
 std::shared_ptr<distributed::HeterServer> b_rpc_service;