[cherry-pick][code-gen] Support code-gen for opmaker of sparse op (#46993) (#47417)

* support generating code of opmaker for backward op invoke forward op (#46912)

* [code-gen] Support code-gen for opmaker of sparse op (#46993)

* support generating code of opmaker for backward op invoke forward op

* gsupport code-gen of opmaker for sparse op

* refind logic of choose phi kernrel

* fix complie budg

* fix code_gen bug

* fix bug

* fix kernel signature code-gen

* fix complie bug of VarType

* fix complie bug of VarType

* fix test_sparse_conv_op

* fix test_sparse_norm_op

* [Phi] Refactor logic of judging whether having a phi kernrel (#46920)

* refind logic of choose phi kernrel

* fix complie budg

* update cmake

Author: zyfncg

.gitignore

@@ -71,7 +71,9 @@ paddle/fluid/pybind/eager_op_function.cc
 
 # these files (directories) are generated before build system generation
 paddle/fluid/operators/generated_op.cc
+paddle/fluid/operators/generated_sparse_op.cc
 paddle/phi/ops/compat/generated_sig.cc
+paddle/phi/ops/compat/generated_sparse_sig.cc
 paddle/phi/api/yaml/parsed_apis/
 python/paddle/utils/code_gen/
 paddle/fluid/pybind/tmp_eager_op_function_impl.h

paddle/fluid/eager/auto_code_generator/eager_generator.cc

@@ -55,7 +55,9 @@ static std::unordered_set<std::string> black_ops_list = {"run_program",
                                                          "fused_gate_attention",
                                                          "fused_feedforward",
                                                          "fused_attention",
-                                                         "fused_gemm_epilogue"};
+                                                         "fused_gemm_epilogue",
+                                                         "sparse_divide_scalar",
+                                                         "sparse_scale"};
 
 static std::string LegalizeVariableName(const std::string& var_name) {
   std::string ret = var_name;
@@ -3161,6 +3163,12 @@ static void DygraphCodeGeneration(const std::string& output_dir,
       continue;
     }
 
+    // Skip the sparse op
+    if (op_type.compare(0, 7, "sparse_") == 0 && op_type != "sparse_momentum" &&
+        op_type != "sparse_attention") {
+      continue;
+    }
+
     GradNodeGenerationInfo bwd_info;
 
     bool is_available = CollectGradInformationFromOpInfo(op_info, &bwd_info);

paddle/fluid/framework/CMakeLists.txt

@@ -190,7 +190,7 @@ cc_test(
 cc_library(
   var_type_traits
   SRCS var_type_traits.cc
-  DEPS framework_proto scope tensor_array sparse_coo_tensor)
+  DEPS framework_proto scope tensor_array sparse_coo_tensor sparse_csr_tensor)
 if(WITH_GPU)
   target_link_libraries(var_type_traits dynload_cuda)
 endif()

paddle/fluid/framework/framework.proto

@@ -156,6 +156,8 @@ message VarType {
     PSTRING = 29;
     // the data type of phi::SparseCooTensor
     SPARSE_COO = 30;
+    // the data type of phi::SparseCsrTensor
+    SPARSE_CSR = 31;
   }
 
   required Type type = 1;
@@ -189,6 +191,7 @@ message VarType {
   optional TensorDesc strings = 9;
   optional TensorDesc vocab = 10;
   optional TensorDesc sparse_coo = 11;
+  optional TensorDesc sparse_csr = 12;
 }
 
 message VarDesc {

paddle/fluid/framework/infershape_utils.cc

@@ -106,6 +106,11 @@ class InferShapeArgumentMappingContext : public phi::ArgumentMappingContext {
     return var_type == proto::VarType::SPARSE_COO;
   }
 
+  bool IsSparseCsrTensorInput(const std::string& name) const override {
+    auto var_type = ctx_.GetInputVarType(name);
+    return var_type == proto::VarType::SPARSE_CSR;
+  }
+
   bool IsDenseTensorOutput(const std::string& name) const override {
     auto var_types = ctx_.GetOutputsVarType(name);
     return std::all_of(var_types.begin(),

paddle/fluid/framework/operator.h

@@ -529,6 +529,11 @@ class ExecutionArgumentMappingContext : public phi::ArgumentMappingContext {
     return var->IsType<phi::SparseCooTensor>();
   }
 
+  bool IsSparseCsrTensorInput(const std::string& name) const override {
+    const auto* var = ctx_.InputVar(name);
+    return var->IsType<phi::SparseCsrTensor>();
+  }
+
   bool IsDenseTensorOutput(const std::string& name) const override {
     auto vars = ctx_.MultiOutputVar(name);
     return std::all_of(vars.begin(), vars.end(), [](const Variable* var) {

paddle/fluid/framework/tensor.h

@@ -18,6 +18,7 @@ limitations under the License. */
 #include "paddle/fluid/framework/mixed_vector.h"
 #include "paddle/phi/core/dense_tensor.h"
 #include "paddle/phi/core/sparse_coo_tensor.h"
+#include "paddle/phi/core/sparse_csr_tensor.h"
 
 namespace paddle {
 namespace framework {

paddle/fluid/framework/var_type_traits.h

@@ -55,6 +55,7 @@ namespace phi {
 class DenseTensor;
 class SelectedRows;
 class SparseCooTensor;
+class SparseCsrTensor;
 }  // namespace phi
 
 // Users should add forward declarations here
@@ -182,6 +183,7 @@ using VarTypeRegistry = detail::VarTypeRegistryImpl<
     Tensor,
     phi::SelectedRows,
     phi::SparseCooTensor,
+    phi::SparseCsrTensor,
     std::vector<Scope *>,
     LoDRankTable,
     Strings,

paddle/fluid/inference/tensorrt/plugin_arg_mapping_context.cc

@@ -108,6 +108,10 @@ bool PluginArgumentMappingContext::IsSparseCooTensorInput(
     const std::string& name) const {
   return false;
 }
+bool PluginArgumentMappingContext::IsSparseCsrTensorInput(
+    const std::string& name) const {
+  return false;
+}
 bool PluginArgumentMappingContext::IsDenseTensorVectorInput(
     const std::string& name) const {
   return false;

paddle/fluid/inference/tensorrt/plugin_arg_mapping_context.h

@@ -48,6 +48,8 @@ class PluginArgumentMappingContext : public ::phi::ArgumentMappingContext {
 
   bool IsSparseCooTensorInput(const std::string& name) const override;
 
+  bool IsSparseCsrTensorInput(const std::string& name) const override;
+
   bool IsDenseTensorVectorInput(const std::string& name) const override;
 
   bool IsDenseTensorOutput(const std::string& name) const override;