diff --git a/paddle/fluid/ir_adaptor/translator/op_compat_gen.py b/paddle/fluid/ir_adaptor/translator/op_compat_gen.py index ea844c659554d..1035674cdb3d9 100644 --- a/paddle/fluid/ir_adaptor/translator/op_compat_gen.py +++ b/paddle/fluid/ir_adaptor/translator/op_compat_gen.py @@ -164,6 +164,9 @@ def insert_new_mutable_attributes( "atol_tensor": "TolTensor", "out": "Out", } + op_arg_name_mappings['push_sparse_v2'].update( + {"out_grad_in": "Out@GRAD", "out_grad_out": "Out@GRAD"} + ) op_name_normailzer_template = env.get_template("op_compat_info.cc.j2") with open(output_source_file, 'wt') as f: diff --git a/paddle/fluid/pir/dialect/op_generator/ops_api_gen.py b/paddle/fluid/pir/dialect/op_generator/ops_api_gen.py index 04da5c4be230c..f38261bd775cc 100644 --- a/paddle/fluid/pir/dialect/op_generator/ops_api_gen.py +++ b/paddle/fluid/pir/dialect/op_generator/ops_api_gen.py @@ -145,6 +145,8 @@ 'uniform_random_batch_size_like', 'c_reduce_min', 'c_reduce_min_', + 'push_sparse_v2', + 'push_sparse_v2_', ] diff --git a/paddle/fluid/pir/dialect/operator/ir/ops.yaml b/paddle/fluid/pir/dialect/operator/ir/ops.yaml index 947e6306c3032..67f65298e0407 100644 --- a/paddle/fluid/pir/dialect/operator/ir/ops.yaml +++ b/paddle/fluid/pir/dialect/operator/ir/ops.yaml @@ -979,6 +979,17 @@ func : prod backward : prod_grad +- op : push_sparse_v2 + args : (Tensor[] ids, Tensor[] w, Tensor[] out_grad_in, int embeddingdim = 11, int tableid = 0, str accessorclass = "", str ctrlabelname = "", int paddingid = 0, bool scalesparsegrad = true, str[] inputnames = {}, bool is_distributed = true) + output : Tensor[](out_grad_out){out_grad_in.size()} + infer_meta : + func : UnchangedMultiInferMeta + param : [out_grad_in] + kernel : + func : push_sparse_v2 + data_type : out_grad_in + inplace: (out_grad_in -> out_grad_out) + - op : randint args : (int low, int high, IntArray shape, DataType dtype=DataType::INT64, Place place={}) output : Tensor(out) diff --git a/paddle/fluid/pir/dialect/operator/utils/utils.cc b/paddle/fluid/pir/dialect/operator/utils/utils.cc index eff7fecb844fb..18daeb8565580 100644 --- a/paddle/fluid/pir/dialect/operator/utils/utils.cc +++ b/paddle/fluid/pir/dialect/operator/utils/utils.cc @@ -72,7 +72,8 @@ const std::unordered_set LegacyOpList = { paddle::onednn::dialect::LrnOp::name(), paddle::onednn::dialect::LrnGradOp::name(), #endif - CReduceMinOp::name()}; + CReduceMinOp::name(), + PushSparseV2Op::name()}; enum class AttrType { UNDEFINED = 0, diff --git a/paddle/phi/api/yaml/op_compat.yaml b/paddle/phi/api/yaml/op_compat.yaml index 08172e9190b92..3c0cb80949235 100755 --- a/paddle/phi/api/yaml/op_compat.yaml +++ b/paddle/phi/api/yaml/op_compat.yaml @@ -2461,6 +2461,14 @@ outputs : out : Out +- op : push_sparse_v2 + inputs : + { x : Ids, W : w} + outputs : + out : Out + extra : + attrs : [int embeddingdim = 11, int tableid = 0, str accessorclass = "", str ctrlabelname = "", int paddingid = 0, bool scalesparsegrad = true, 'str[] inputnames = {}', bool is_distributed = true] + - op : put_along_axis backward : put_along_axis_grad inputs :