Support selected rows new ir

paddle/fluid/ir/dialect/kernel_dialect.cc

@@ -41,6 +41,7 @@ PaddleKernelDialect::PaddleKernelDialect(ir::IrContext *context)
 
 void PaddleKernelDialect::initialize() {
   RegisterTypes<paddle::dialect::AllocatedDenseTensorType>();
+  RegisterTypes<paddle::dialect::AllocatedSelectedRowsType>();
   RegisterOps<dialect::PhiKernelOp>();
 
   RegisterAttributes<paddle::dialect::KernelAttribute>();

paddle/fluid/ir/dialect/kernel_type.cc

@@ -41,7 +41,32 @@ const size_t& AllocatedDenseTensorType::offset() const {
   return storage()->dense_tensor_type_.offset();
 }
 
+const phi::Place& AllocatedSelectedRowsType::place() const {
+  return storage()->place_;
+}
+
+const ir::Type& AllocatedSelectedRowsType::dtype() const {
+  return storage()->selected_rows_type_.dtype();
+}
+
+const phi::DDim& AllocatedSelectedRowsType::dims() const {
+  return storage()->selected_rows_type_.dims();
+}
+
+const phi::DataLayout& AllocatedSelectedRowsType::data_layout() const {
+  return storage()->selected_rows_type_.data_layout();
+}
+
+const phi::LoD& AllocatedSelectedRowsType::lod() const {
+  return storage()->selected_rows_type_.lod();
+}
+
+const size_t& AllocatedSelectedRowsType::offset() const {
+  return storage()->selected_rows_type_.offset();
+}
+
 }  // namespace dialect
 }  // namespace paddle
 
 IR_DEFINE_EXPLICIT_TYPE_ID(paddle::dialect::AllocatedDenseTensorType)
+IR_DEFINE_EXPLICIT_TYPE_ID(paddle::dialect::AllocatedSelectedRowsType)

paddle/fluid/ir/dialect/kernel_type.h

@@ -20,9 +20,7 @@
 
 namespace paddle {
 namespace dialect {
-///
-/// \brief Define built-in parametric types.
-///
+
 class AllocatedDenseTensorType : public ir::Type {
  public:
   using Type::Type;
@@ -64,7 +62,49 @@ class AllocatedDenseTensorType : public ir::Type {
   const size_t &offset() const;
 };
 
+class AllocatedSelectedRowsType : public ir::Type {
+ public:
+  using Type::Type;
+
+  DECLARE_TYPE_UTILITY_FUNCTOR(AllocatedSelectedRowsType,
+                               AllocatedSelectedRowsTypeStorage);
+
+  static AllocatedSelectedRowsType get(ir::IrContext *ctx,
+                                       const phi::Place &place,
+                                       dialect::SelectedRowsType type) {
+    return ir::TypeManager::template get<AllocatedSelectedRowsType>(
+        ctx, place, type);
+  }
+
+  static AllocatedSelectedRowsType get(ir::IrContext *ctx,
+                                       const phi::Place &place,
+                                       const ir::Type &dtype,
+                                       const phi::DDim &dims,
+                                       const phi::DataLayout &layout,
+                                       const phi::LoD &lod,
+                                       size_t offset) {
+    dialect::SelectedRowsType dense_tensor_type =
+        dialect::SelectedRowsType::get(ctx, dtype, dims, layout, lod, offset);
+
+    return ir::TypeManager::template get<AllocatedSelectedRowsType>(
+        ctx, place, dense_tensor_type);
+  }
+
+  const phi::Place &place() const;
+
+  const ir::Type &dtype() const;
+
+  const phi::DDim &dims() const;
+
+  const phi::DataLayout &data_layout() const;
+
+  const phi::LoD &lod() const;
+
+  const size_t &offset() const;
+};
+
 }  // namespace dialect
 }  // namespace paddle
 
 IR_DECLARE_EXPLICIT_TYPE_ID(paddle::dialect::AllocatedDenseTensorType)
+IR_DECLARE_EXPLICIT_TYPE_ID(paddle::dialect::AllocatedSelectedRowsType)

paddle/fluid/ir/dialect/kernel_type_storage.h

@@ -88,5 +88,67 @@ struct AllocatedDenseTensorTypeStorage : public ir::TypeStorage {
   dialect::DenseTensorType dense_tensor_type_;
 };
 
+///
+/// \brief Define Parametric TypeStorage for AllocatedSelectedRowsTypeStorage.
+///
+///
+struct AllocatedSelectedRowsTypeStorage : public ir::TypeStorage {
+  using Place = phi::Place;
+  ///
+  /// \brief Declare ParamKey according to parameter type.
+  ///
+  using ParamKey = std::tuple<phi::Place, dialect::SelectedRowsType>;
+
+  AllocatedSelectedRowsTypeStorage(const phi::Place& place,
+                                   const dialect::SelectedRowsType& type)
+      : place_(place), selected_rows_type_(type) {}
+
+  ///
+  /// \brief Each derived TypeStorage must define a Construct method, which
+  /// StorageManager uses to construct a derived TypeStorage.
+  ///
+  static AllocatedSelectedRowsTypeStorage* Construct(const ParamKey& key) {
+    return new AllocatedSelectedRowsTypeStorage(std::get<0>(key),
+                                                std::get<1>(key));
+  }
+
+  ///
+  /// \brief Each derived TypeStorage must provide a HashValue method.
+  ///
+  static std::size_t HashValue(const ParamKey& key) {
+    std::size_t hash_value = 791;
+    // hash place
+    hash_value = ir::hash_combine(hash_value, std::get<0>(key).HashValue());
+
+    // hash dtype
+    auto selected_rows_type = std::get<1>(key);
+    hash_value = ir::hash_combine(hash_value,
+                                  dialect::DenseTensorTypeStorage::HashValue(
+                                      dialect::DenseTensorTypeStorage::ParamKey(
+                                          selected_rows_type.dtype(),
+                                          selected_rows_type.dims(),
+                                          selected_rows_type.data_layout(),
+                                          selected_rows_type.lod(),
+                                          selected_rows_type.offset())));
+    return hash_value;
+  }
+
+  ///
+  /// \brief Each derived TypeStorage needs to overload operator==.
+  ///
+  bool operator==(const ParamKey& key) const {
+    return ParamKey(place_, selected_rows_type_) == key;
+  }
+
+  ParamKey GetAsKey() const { return ParamKey(place_, selected_rows_type_); }
+
+  ///
+  /// \brief AllocatedSelectedRowsTypeStorage include five parameters: place,
+  /// SelectedRowsType
+  ///
+  phi::Place place_;
+  dialect::SelectedRowsType selected_rows_type_;
+};
+
 }  // namespace dialect
 }  // namespace paddle

paddle/fluid/ir/dialect/pd_dialect.cc

@@ -111,15 +111,23 @@ void PaddleDialect::initialize() {
 }
 
 void PaddleDialect::PrintType(ir::Type type, std::ostream &os) const {
-  DenseTensorType tensor_type = type.dyn_cast<DenseTensorType>();
-
-  os << "tensor<";
-  for (auto d : phi::vectorize(tensor_type.dims())) {
-    os << d;
-    os << "x";
+  if (auto tensor_type = type.dyn_cast<DenseTensorType>()) {
+    os << "tensor<";
+    for (auto d : phi::vectorize(tensor_type.dims())) {
+      os << d;
+      os << "x";
+    }
+    tensor_type.dtype().Print(os);
+    os << ">";
+  } else if (auto selected_rows_type = type.dyn_cast<SelectedRowsType>()) {
+    os << "selectedrows<";
+    for (auto d : phi::vectorize(selected_rows_type.dims())) {
+      os << d;
+      os << "x";
+    }
+    selected_rows_type.dtype().Print(os);
+    os << ">";
   }
-  tensor_type.dtype().Print(os);
-  os << ">";
 }
 
 void PaddleDialect::PrintAttribute(ir::Attribute attr, std::ostream &os) const {

paddle/fluid/ir/dialect/pd_op.yaml

@@ -187,18 +187,18 @@
   no_need_buffer: null
   data_transform: null
   infer_meta:
-    func: UnchangedInferMeta
+    func: EmbeddingGradSparseInferMeta
     param: [weight]
   kernel:
-    func: [embedding_grad_sparse]
+    func: [embedding_sparse_grad]
     param: [x, weight, out_grad, padding_idx, sparse]
     backend: null
     layout: null
     data_type:
       ordered: false
       candidates: [weight]
       to_complex_flag: [false]
-    dispatch: {embedding_grad_sparse: null}
+    dispatch: {embedding_sparse_grad: null}
     force_backend: null
   inplace: null
   view: null

paddle/fluid/ir/phi_kernel_adaptor/phi_kernel_util.cc

@@ -104,6 +104,8 @@ void BuildValue(ir::Value value,
     var->GetMutable<phi::DenseTensor>();
   } else if (value.type().isa<paddle::dialect::AllocatedDenseTensorType>()) {
     var->GetMutable<phi::DenseTensor>();
+  } else if (value.type().isa<paddle::dialect::AllocatedSelectedRowsType>()) {
+    var->GetMutable<phi::SelectedRows>();
   } else if (value.type().isa<ir::VectorType>()) {
     auto tensor_array = var->GetMutable<paddle::framework::TensorRefArray>();
     for (size_t i = 0; i < value.type().dyn_cast<ir::VectorType>().size();

paddle/fluid/ir/phi_kernel_adaptor/phi_kernel_util.h

@@ -33,6 +33,7 @@
 #include "paddle/phi/core/kernel_context.h"
 
 #include "paddle/fluid/ir/dialect/kernel_attribute.h"
+#include "paddle/fluid/ir/dialect/kernel_type.h"
 #include "paddle/fluid/ir/dialect/pd_attribute.h"
 #include "paddle/fluid/ir/interface/op_yaml_info_parser.h"
 #include "paddle/phi/core/infermeta_utils.h"
@@ -238,6 +239,23 @@ void BuildPhiContext(
     } else if (attr_type_name == "ir::ArrayAttribute<ir::Int64Attribute>") {
       auto array_list = attr_map[t].dyn_cast<ir::ArrayAttribute>().data();
 
+      std::vector<int64_t> vec_res;
+      if (array_list.size() > 0) {
+        PADDLE_ENFORCE_EQ(
+            array_list[0].isa<ir::Int64Attribute>(),
+            true,
+            phi::errors::PreconditionNotMet(
+                "Element in array list MUST be ir::Int64Attribute "));
+
+        for (size_t i = 0; i < array_list.size(); ++i) {
+          vec_res.push_back(
+              array_list[i].dyn_cast<ir::Int64Attribute>().data());
+        }
+      }
+      ctx->EmplaceBackAttr(vec_res);
+    } else if (attr_type_name == "ir::ArrayAttribute<ir::Int64Attribute>") {
+      auto array_list = attr_map[t].dyn_cast<ir::ArrayAttribute>().data();
+
       std::vector<int64_t> vec_res;
       if (array_list.size() > 0) {
         PADDLE_ENFORCE_EQ(
@@ -280,13 +298,19 @@ void BuildPhiContext(
     for (size_t i = 0; i < op->num_results(); ++i) {
       ir::Value out_ptr = op->result(i);
       auto name = name_map.at(out_ptr);
-      if (out_ptr.type()) {
-        ctx->EmplaceBackOutput(OutType(const_cast<phi::DenseTensor*>(
-            &(inner_scope->FindVar(name)->Get<phi::DenseTensor>()))));
-      } else {
+      auto out_type = out_ptr.type();
+      if (!out_type) {
         phi::DenseTensor* ptr = nullptr;
         OutType out_ptr(ptr);
         ctx->EmplaceBackOutput(out_ptr);
+      } else if (out_type.isa<paddle::dialect::AllocatedDenseTensorType>()) {
+        ctx->EmplaceBackOutput(OutType(const_cast<phi::DenseTensor*>(
+            &(scope->Var(name)->Get<phi::DenseTensor>()))));
+      } else if (out_type.isa<paddle::dialect::AllocatedSelectedRowsType>()) {
+        ctx->EmplaceBackOutput(OutType(const_cast<phi::SelectedRows*>(
+            &(scope->Var(name)->Get<phi::SelectedRows>()))));
+      } else {
+        PADDLE_THROW("not support type");
       }
 
       if (output_map != nullptr) {

paddle/fluid/ir/transforms/pd_op_to_kernel_pass.cc

@@ -252,6 +252,13 @@ std::unique_ptr<ir::Program> PdOpLowerToKernelPass(ir::Program* prog) {
 
           ir::Type t1 = ir::VectorType::get(ctx, vec_inner_types);
           op_output_types.push_back(t1);
+        } else if (result_type.isa<dialect::SelectedRowsType>()) {
+          auto allocated_selected_rows_dtype =
+              paddle::dialect::AllocatedSelectedRowsType::get(
+                  ctx,
+                  phi::TransToPhiPlace(kernel_key.backend()),
+                  result_type.dyn_cast<dialect::SelectedRowsType>());
+          op_output_types.push_back(allocated_selected_rows_dtype);
         } else {
           PADDLE_THROW(phi::errors::Unimplemented(
               "Result type only support DenseTensorType and VectorType"));
@@ -322,6 +329,8 @@ std::unique_ptr<ir::Program> PdOpLowerToKernelPass(ir::Program* prog) {
             }
           } else if (new_in_type.isa<ir::VectorType>()) {
             // [ todo need update here, support combine data transfomer]
+          } else if (new_in_type.isa<dialect::AllocatedSelectedRowsType>()) {
+            // do nothing here
           } else {
             PADDLE_THROW(phi::errors::Unimplemented(
                 "only support allocated dense tensor type for now"));

paddle/fluid/ir_adaptor/translator/op_translator.cc

@@ -893,6 +893,41 @@ struct RnnOpTranscriber : public OpTranscriber {
   };
 };
 
+struct EmbeddingGradOpTranscriber : public OpTranscriber {
+  void HandleNonexistentAttribute(ir::IrContext* ctx,
+                                  ir::AttributeMap* attribute_map,
+                                  const OpAttributeInfo& info) override {
+    if (info.name == "padding_idx") {
+      (*attribute_map)[info.name] = ir::Int64Attribute::get(ctx, -1);
+    } else if (info.name == "sparse") {
+      (*attribute_map)[info.name] = ir::BoolAttribute::get(ctx, false);
+    }
+  }
+
+  ir::OpInfo LoopkUpOpInfo(ir::IrContext* ctx, const OpDesc& op_desc) override {
+    std::string target_op_name =
+        kTargetDialectPrefix + OpNameCompatibleMapping(op_desc.Type());
+
+    bool is_sparse = paddle::get<bool>(op_desc.GetAttr("is_sparse"));
+
+    if (is_sparse) {
+      target_op_name = "pd.embedding_grad_sparse";
+    } else {
+      target_op_name = "pd.embedding_grad_dense";
+    }
+    VLOG(6) << "[op name normalizing: " << op_desc.Type() << " to "
+            << target_op_name;
+    auto op_info = ctx->GetRegisteredOpInfo(target_op_name);
+    if (!op_info) {
+      IR_THROW("Op %d should have corresponding OpInfo %d",
+               op_desc.Type(),
+               target_op_name);
+    }
+
+    return op_info;
+  }
+};
+
 struct FeedOpTranscriber : public OpTranscriber {
   ir::AttributeMap TranslateOpAttribute(
       ir::IrContext* ctx,
@@ -959,6 +994,7 @@ OpTranslator::OpTranslator() {
   special_handlers["fetch_v2"] = FetchOpTranscriber();
   special_handlers["cast"] = CastOpTranscriber();
   special_handlers["lookup_table_v2"] = EmbeddingOpTranscriber();
+  special_handlers["lookup_table_v2_grad"] = EmbeddingGradOpTranscriber();
   special_handlers["assign_value"] = AssignValueOpTranscriber();
   special_handlers["increment"] = IncrementOpTranscriber();
   special_handlers["rnn"] = RnnOpTranscriber();

paddle/fluid/ir_adaptor/translator/program_translator.cc

@@ -59,7 +59,6 @@ void ProgramTranslator::Translate() {
       platform::errors::PreconditionNotMet(
           "Not support multi block ProgramDesc translated, now has %d blocks",
           legacy_program_->Size()));
-
   for (size_t block_idx = 0; block_idx < legacy_program_->Size(); block_idx++) {
     const BlockDesc& block = legacy_program_->Block(block_idx);
     GetParameterForSingleBlock(block);

paddle/fluid/ir_adaptor/translator/translate.cc

@@ -32,7 +32,6 @@ std::unique_ptr<Program> TranslateLegacyProgramToProgram(
   ir::IrContext* ctx = ir::IrContext::Instance();
   ctx->GetOrRegisterDialect<dialect::PaddleDialect>();
   auto program = std::make_unique<Program>(ctx);
-
   translator::ProgramTranslator program_translator(&legacy_program,
                                                    program.get());
   program_translator.Translate();