Autogen embedding static graph code (#52460)

* autogen embedding

* deal

* fix bug in CompatMetaTensor::share_lod

paddle/fluid/framework/infershape_utils.cc

@@ -392,8 +392,14 @@ void CompatMetaTensor::share_lod(const MetaTensor& meta_tensor) {
     }
   } else {
     auto* var = PADDLE_GET(VarDesc*, var_);
-    if (!meta_tensor.is_dense() && !meta_tensor.is_tensor_array()) {
-      VLOG(3) << "input metatensor is not phi::DenseTensor or LoDTensorArray.";
+    // NOTE(lizhiyu): If var is select_rows and meta_tensor is dense,
+    // 'var->SetLodLevel' will fail. This case will happen when execute
+    // 'test_hsigmoid_op.py'. So it is needed to assert 'var' type.
+    if ((var && (var->GetType() != proto::VarType::LOD_TENSOR &&
+                 var->GetType() != proto::VarType::LOD_TENSOR_ARRAY)) ||
+        (!meta_tensor.is_dense() && !meta_tensor.is_tensor_array())) {
+      VLOG(3) << "this tensor or input metatensor is not phi::DenseTensor or "
+                 "LoDTensorArray.";
       return;
     }
     if (var) {
@@ -410,7 +416,9 @@ void CompatMetaTensor::share_dims(const MetaTensor& meta_tensor) {
   if (is_runtime_) {
     auto* var = PADDLE_GET(Variable*, var_);
     if (var == nullptr) return;
-    if (var->IsType<phi::SelectedRows>()) {
+    // NOTE(lizhiyu): If var is select_rows and meta_tensor is dense,
+    // `var->GetMutable<phi::SelectedRows>()` will failed.
+    if (var->IsType<phi::SelectedRows>() && meta_tensor.is_selected_rows()) {
       auto* selected_rows = var->GetMutable<phi::SelectedRows>();
       auto& input_selected_rows =
           static_cast<const CompatMetaTensor&>(meta_tensor).GetSelectedRows();

paddle/fluid/operators/generator/filters.py

@@ -39,6 +39,28 @@ def get_infer_var_type_func(op_name):
     ctx->SyncTypeAndDataType("X", "Out");
   }}
 }};
+"""
+    elif op_name == "lookup_table_v2_grad":
+        return f"""
+class {to_pascal_case(op_name)}InferVarType : public framework::VarTypeInference {{
+ public:
+  void operator()(framework::InferVarTypeContext* ctx) const override {{
+    auto out_var_name = framework::GradVarName("W");
+    auto attr = ctx->GetAttr("is_sparse");
+    bool is_sparse = PADDLE_GET(bool, attr);
+    if (is_sparse) {{
+      VLOG(3) << "lookup_table_v2_grad op " << framework::GradVarName("W")
+              << " is set to SelectedRows";
+      ctx->SetOutputType(out_var_name,
+                         framework::proto::VarType::SELECTED_ROWS);
+    }} else {{
+      VLOG(3) << "lookup_table_v2_grad op " << framework::GradVarName("W")
+              << " is set to phi::DenseTensor";
+      ctx->SetOutputType(out_var_name, framework::proto::VarType::LOD_TENSOR);
+    }}
+    ctx->SetOutputDataType(out_var_name, ctx->GetInputDataType("W"));
+  }}
+}};
 """
     elif op_name == "merge_selected_rows":
         return f"""

paddle/fluid/operators/lookup_table_v2_op.cc

-/* Copyright (c) 2019 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. */
-
-#include "paddle/fluid/operators/lookup_table_v2_op.h"
-
-#include <memory>
-
-#include "paddle/fluid/framework/no_need_buffer_vars_inference.h"
-#include "paddle/fluid/framework/op_version_registry.h"
-#include "paddle/fluid/framework/var_type_inference.h"
-
-namespace paddle {
-namespace operators {
-
-class LookupTableV2Op : public framework::OperatorWithKernel {
- public:
-  using framework::OperatorWithKernel::OperatorWithKernel;
-
-  void InferShape(framework::InferShapeContext* ctx) const override {
-    PADDLE_ENFORCE_EQ(ctx->HasInput("W"),
-                      true,
-                      platform::errors::InvalidArgument(
-                          "Input(W) of LookupTableV2Op should not be null."));
-    PADDLE_ENFORCE_EQ(ctx->HasInput("Ids"),
-                      true,
-                      platform::errors::InvalidArgument(
-                          "Input(Ids) of LookupTableV2Op should not be null."));
-    PADDLE_ENFORCE_EQ(
-        ctx->HasOutput("Out"),
-        true,
-        platform::errors::InvalidArgument(
-            "Output(Out) of LookupTableV2Op should not be null."));
-
-    auto table_dims = ctx->GetInputDim("W");
-    auto ids_dims = ctx->GetInputDim("Ids");
-    int ids_rank = ids_dims.size();
-    VLOG(5) << "ids rank is " << ids_rank << std::endl;
-    PADDLE_ENFORCE_EQ(
-        table_dims.size(),
-        2,
-        platform::errors::InvalidArgument(
-            "ShapeError: The dimensions of the 'lookup table' must be 2. "
-            "But received lookup table's dimensions = %d, "
-            "lookup table's shape = [%s].",
-            table_dims.size(),
-            table_dims));
-
-    auto output_dims = phi::vectorize(ids_dims);
-    output_dims.push_back(table_dims[1]);
-    ctx->SetOutputDim("Out", phi::make_ddim(output_dims));
-
-    if (ctx->GetOutputsVarType("Out")[0] ==
-        framework::proto::VarType::LOD_TENSOR) {
-      ctx->ShareLoD("Ids", /*->*/ "Out");
-    }
-  }
-
- protected:
-  phi::KernelKey GetExpectedKernelType(
-      const framework::ExecutionContext& ctx) const override {
-    auto data_type = OperatorWithKernel::IndicateVarDataType(ctx, "W");
-    return phi::KernelKey(data_type, ctx.device_context().GetPlace());
-  }
-};
-
-class LookupTableV2OpMaker : public framework::OpProtoAndCheckerMaker {
- public:
-  void Make() override {
-    AddInput("W",
-             "(Tensor) The input represents embedding tensors, "
-             "which is a learnable parameter.");
-    AddInput("Ids",
-             "An input with type int64 "
-             "contains the ids to be looked up in W.");
-    AddOutput("Out", "The lookup results, which have the same type as W.");
-    AddAttr<int64_t>("padding_idx",
-                     "(int64, default -1) "
-                     "If the value is -1, it makes no effect to lookup. "
-                     "Otherwise the given value indicates padding the output "
-                     "with zeros whenever lookup encounters it in Ids.")
-        .SetDefault(kNoPadding);
-    AddComment(R"DOC(
-Lookup Table V2 Operator.
-
-This operator is used to perform lookups on the parameter W,
-then concatenated into a dense tensor.
-
-The input Ids can carry the LoD (Level of Details) information,
-or not. And the output only shares the LoD information with input Ids.
-
-)DOC");
-  }
-};
-
-DECLARE_NO_NEED_BUFFER_VARS_INFERER(LookupTableV2GradOpNoBufferVarsInferer,
-                                    "W");
-
-template <typename T>
-class LookupTableV2GradOpMaker : public framework::SingleGradOpMaker<T> {
- public:
-  using framework::SingleGradOpMaker<T>::SingleGradOpMaker;
-
- protected:
-  void Apply(GradOpPtr<T> op) const override {
-    op->SetType("lookup_table_v2_grad");
-
-    op->SetInput("W", this->Input("W"));
-    op->SetInput("Ids", this->Input("Ids"));
-    op->SetInput(framework::GradVarName("Out"), this->OutputGrad("Out"));
-
-    op->SetOutput(framework::GradVarName("W"), this->InputGrad("W"));
-
-    op->SetAttrMap(this->Attrs());
-  }
-};
-
-class LookupTableV2OpGrad : public framework::OperatorWithKernel {
- public:
-  using framework::OperatorWithKernel::OperatorWithKernel;
-
-  void InferShape(framework::InferShapeContext* ctx) const override {
-    auto table_dims = ctx->GetInputDim("W");
-    ctx->SetOutputDim(framework::GradVarName("W"), table_dims);
-  }
-
- protected:
-  phi::KernelKey GetExpectedKernelType(
-      const framework::ExecutionContext& ctx) const override {
-    auto data_type = OperatorWithKernel::IndicateVarDataType(
-        ctx, framework::GradVarName("Out"));
-    return phi::KernelKey(data_type, ctx.device_context().GetPlace());
-  }
-};
-
-class LookupTableV2OpGradVarTypeInference : public framework::VarTypeInference {
- public:
-  void operator()(framework::InferVarTypeContext* ctx) const override {
-    auto out_var_name = framework::GradVarName("W");
-    auto attr = ctx->GetAttr("is_sparse");
-    bool is_sparse = PADDLE_GET(bool, attr);
-    if (is_sparse) {
-      VLOG(3) << "lookup_table_v2_grad op " << framework::GradVarName("W")
-              << " is set to SelectedRows";
-      ctx->SetOutputType(out_var_name,
-                         framework::proto::VarType::SELECTED_ROWS);
-    } else {
-      VLOG(3) << "lookup_table_v2_grad op " << framework::GradVarName("W")
-              << " is set to phi::DenseTensor";
-      ctx->SetOutputType(out_var_name, framework::proto::VarType::LOD_TENSOR);
-    }
-    ctx->SetOutputDataType(out_var_name, ctx->GetInputDataType("W"));
-  }
-};
-
-}  // namespace operators
-}  // namespace paddle
-
-namespace ops = paddle::operators;
-REGISTER_OPERATOR(lookup_table_v2,
-                  ops::LookupTableV2Op,
-                  ops::LookupTableV2OpMaker,
-                  ops::LookupTableV2GradOpMaker<paddle::framework::OpDesc>,
-                  ops::LookupTableV2GradOpMaker<paddle::imperative::OpBase>);
-
-REGISTER_OPERATOR(lookup_table_v2_grad,
-                  ops::LookupTableV2OpGrad,
-                  ops::LookupTableV2GradOpNoBufferVarsInferer,
-                  ops::LookupTableV2OpGradVarTypeInference);
-
-/* ==========================  register checkpoint ===========================*/
-REGISTER_OP_VERSION(lookup_table_v2)
-    .AddCheckpoint(
-        R"ROC(fix lookup_table_v2, add input type `int32`)ROC",
-        paddle::framework::compatible::OpVersionDesc()
-            .BugfixWithBehaviorChanged("lookup_table_v2 support input type "
-                                       "`int64`; after support input type "
-                                       "`int32/int64`"));
-
-/* ========================================================================== */

paddle/phi/api/yaml/op_compat.yaml

@@ -592,6 +592,11 @@
 
 - op : embedding (lookup_table_v2)
   backward : embedding_grad (lookup_table_v2_grad)
+  inputs :
+    {x : Ids, weight : W}
+  outputs :
+    out : Out
+  manual_signature : [embedding_grad]
   extra :
     attrs : [bool is_sparse = false, bool is_distributed = false, bool remote_prefetch = false,
              int trainer_id = 0, int slot = 0, 'int64_t[] height_sections = {}', 'str[] epmap = {}',

paddle/phi/api/yaml/op_version.yaml

@@ -43,6 +43,13 @@
         - add_input : Max
           comment : Pass the mix, min value as input, not attribute. Max is dispensable.
 
+- op : embedding
+  version :
+    - checkpoint : Upgrade flip, add new attr [axis] and delete attr [dims]
+      action :
+        - fix_bug : fix_bug
+          comment : lookup_table_v2 support input type `int64`; after support input type `int32/int64`
+
 - op : equal
   version :
     - checkpoint : Upgrade compare ops, add a new attribute [force_cpu]

paddle/phi/api/yaml/static_backward.yaml

@@ -7,6 +7,21 @@
   composite: assign_grad(out_grad, x_grad)
   invoke : assign(out_grad)
 
+- backward_op : embedding_grad
+  forward : embedding (Tensor x, Tensor weight, int64_t padding_idx=-1) -> Tensor(out)
+  args : (Tensor x, Tensor weight, Tensor out_grad, int64_t padding_idx=-1)
+  output : Tensor(weight_grad)
+  infer_meta :
+    func : EmbeddingGradInferMeta
+    param : [x,weght]
+  kernel :
+    func : embedding_grad {dense, dense, dense -> dense}
+           embedding_sparse_grad {dense, dense, dense -> selected_rows}
+           sparse_weight_embedding_grad {selected_rows, dense, dense -> dense}
+           sparse_weight_embedding_sparse_grad {selected_rows, dense, dense -> selected_rows}
+    data_type : out_grad
+  no_need_buffer : weight
+
 - backward_op : frobenius_norm_grad
   forward: frobenius_norm (Tensor x, IntArray axis={0}, bool keepdim=false, bool reduce_all=false, int in_dtype=-1, int out_dtype=-1) -> Tensor(out)
   args : (Tensor x, Tensor out, Tensor out_grad, IntArray axis={0}, bool keepdim=false, bool reduce_all=false, int in_dtype=-1, int out_dtype=-1)

paddle/phi/api/yaml/static_ops.yaml

@@ -57,6 +57,19 @@
     func : broadcast
     param: [x, root]
 
+- op : embedding
+  args : (Tensor x, Tensor weight, int64_t padding_idx=-1)
+  output : Tensor
+  infer_meta :
+    func : EmbeddingInferMeta
+    param : [x, weight, padding_idx]
+  kernel :
+    func : embedding {dense, dense -> dense}
+           sparse_weight_embedding {dense, selected_rows -> dense}
+    param : [x, weight, padding_idx]
+    data_type : weight
+  backward : embedding_grad
+
 - op : equal
   args : (Tensor x, Tensor y, int axis = -1, bool force_cpu=false)
   output : Tensor(out)

paddle/phi/infermeta/backward.cc

@@ -323,6 +323,15 @@ void EigvalshGradInferMeta(const MetaTensor& out_v,
   }
 }
 
+void EmbeddingGradInferMeta(const MetaTensor& x,
+                            const MetaTensor& weight,
+                            MetaTensor* out) {
+  (void)x;
+  if (weight) {
+    out->share_dims(weight);
+  }
+}
+
 void FFTC2RGradInferMeta(const MetaTensor& x,
                          const std::vector<int64_t>& axes,
                          const std::string& normalization,

paddle/phi/infermeta/backward.h

@@ -151,6 +151,10 @@ void EigvalshGradInferMeta(const MetaTensor& out_v,
                            bool is_test,
                            MetaTensor* x_grad);
 
+void EmbeddingGradInferMeta(const MetaTensor& x,
+                            const MetaTensor& weight,
+                            MetaTensor* out);
+
 void FFTC2RGradInferMeta(const MetaTensor& x,
                          const std::vector<int64_t>& axes,
                          const std::string& normalization,

paddle/phi/ops/compat/embedding_sig.cc

@@ -16,15 +16,6 @@
 
 namespace phi {
 
-KernelSignature EmbeddingOpArgumentMapping(const ArgumentMappingContext& ctx) {
-  if (ctx.IsDenseTensorInput("W")) {
-    return KernelSignature("embedding", {"Ids", "W"}, {"padding_idx"}, {"Out"});
-  } else {
-    return KernelSignature(
-        "sparse_weight_embedding", {"Ids", "W"}, {"padding_idx"}, {"Out"});
-  }
-}
-
 KernelSignature EmbeddingGradOpArgumentMapping(
     const ArgumentMappingContext& ctx) {
   if (ctx.IsDenseTensorInput("W")) {
@@ -56,7 +47,6 @@ KernelSignature EmbeddingGradOpArgumentMapping(
 
 }  // namespace phi
 
-PD_REGISTER_BASE_KERNEL_NAME(lookup_table_v2, embedding);
 PD_REGISTER_BASE_KERNEL_NAME(lookup_table_v2_grad, embedding_grad);
 PD_REGISTER_BASE_KERNEL_NAME(lookup_table_v2_grad, embedding_sparse_grad);
 PD_REGISTER_BASE_KERNEL_NAME(lookup_table_v2_grad,
@@ -64,6 +54,5 @@ PD_REGISTER_BASE_KERNEL_NAME(lookup_table_v2_grad,
 PD_REGISTER_BASE_KERNEL_NAME(lookup_table_v2_grad,
                              sparse_weight_embedding_sparse_grad);
 
-PD_REGISTER_ARG_MAPPING_FN(lookup_table_v2, phi::EmbeddingOpArgumentMapping);
 PD_REGISTER_ARG_MAPPING_FN(lookup_table_v2_grad,
                            phi::EmbeddingGradOpArgumentMapping);