【PIR Dist Op Reg No.2】 reg distributed_push_sparse

paddle/fluid/pir/dialect/op_generator/ops_api_gen.py

@@ -136,6 +136,7 @@
     'c_softmax_with_cross_entropy',
     'c_split',
     'decayed_adagrad',
+    'distributed_push_sparse',
     'distributed_lookup_table',
     'dpsgd',
     'embedding_grad_sparse',

paddle/fluid/pir/dialect/operator/ir/ops.yaml

@@ -473,6 +473,15 @@
     func : distributed_lookup_table
     data_type : dtype
 
+- op : distributed_push_sparse
+  args : (Tensor[] ids, Tensor[] shows, Tensor[] clicks, int table_id = 0, int size = 8, bool is_distributed = false, str push_sparse_version = "push_sparse", int64_t padding_idx = -1, DataType dtype=DataType::FLOAT32, bool is_test = false, bool use_cvm_op = false)
+  output : Tensor[](output){ids.size()}
+  infer_meta :
+    func : DistributedPushSparseInferMeta
+  kernel :
+    func: distributed_push_sparse
+    data_type : dtype
+
 - op : divide
   args : (Tensor x, Tensor y)
   output : Tensor(out)

paddle/fluid/pir/dialect/operator/utils/utils.cc

@@ -41,6 +41,7 @@ const std::unordered_set<std::string> LegacyOpList = {
     CBroadcast_Op::name(),
     CSyncCalcStream_Op::name(),
     CSyncCommStream_Op::name(),
+    DistributedPushSparseOp::name(),
     FtrlOp::name(),
     FusedElemwiseAddActivationOp::name(),
     FusedElemwiseAddActivationGradOp::name(),

paddle/phi/api/yaml/op_compat.yaml

@@ -843,6 +843,10 @@
     out : Out
 
 - op : distributed_push_sparse
+  inputs :
+    {ids : Ids, shows : Shows, clicks: Clicks}
+  outputs :
+    output : Outputs
   extra :
     attrs : ['int[] slots = {}']
 

paddle/phi/infermeta/ternary.cc

@@ -255,6 +255,37 @@ void BoxCoderInferMeta(const MetaTensor& prior_box,
   output_box->set_dtype(target_box.dtype());
 }
 
+void DistributedPushSparseInferMeta(
+    const std::vector<const MetaTensor*>& ids,
+    const std::vector<const MetaTensor*>& shows,
+    const std::vector<const MetaTensor*>& clicks,
+    int table_id,
+    int size,
+    bool is_distributed,
+    const std::string& push_sparse_version,
+    int64_t padding_idx,
+    DataType dtype,
+    bool is_test,
+    bool use_cvm_op,
+    std::vector<MetaTensor*> output) {
+  auto ids_size = ids.size();
+  std::vector<DDim> ids_dims;
+  ids_dims.reserve(ids.size());
+  for (size_t i = 1; i < ids_size; ++i) {
+    PADDLE_ENFORCE_EQ(ids_dims[i].size(),
+                      2,
+                      phi::errors::InvalidArgument(
+                          "The dimension of the 'Ids' tensor must be 2."));
+  }
+
+  for (auto& out : output) {
+    if (out == nullptr) {
+      continue;
+    }
+    out->set_dtype(ids[0]->dtype());
+  }
+}
+
 void DpsgdInferMeta(const MetaTensor& param,
                     const MetaTensor& grad,
                     const MetaTensor& learning_rate,

paddle/phi/infermeta/ternary.h

@@ -63,6 +63,20 @@ void BoxCoderInferMeta(const MetaTensor& prior_box,
                        MetaTensor* output_box,
                        MetaConfig config = MetaConfig());
 
+void DistributedPushSparseInferMeta(
+    const std::vector<const MetaTensor*>& ids,
+    const std::vector<const MetaTensor*>& shows,
+    const std::vector<const MetaTensor*>& clicks,
+    int table_id,
+    int size,
+    bool is_distributed,
+    const std::string& push_sparse_version,
+    int64_t padding_idx,
+    DataType dtype,
+    bool is_test,
+    bool use_cvm_op,
+    std::vector<MetaTensor*> output);
+
 void DpsgdInferMeta(const MetaTensor& param,
                     const MetaTensor& grad,
                     const MetaTensor& learning_rate,

test/ir/pir/translator/CMakeLists.txt

@@ -13,6 +13,9 @@ list(APPEND DISTRIBUTED_OP_TRANSLATOR_TEST test_c_split_translator)
 list(APPEND DISTRIBUTED_OP_TRANSLATOR_TEST test_distributed_fused_lamb_init)
 list(APPEND DISTRIBUTED_OP_TRANSLATOR_TEST
      test_distributed_lookup_table_translate)
+list(APPEND DISTRIBUTED_OP_TRANSLATOR_TEST
+     test_distributed_push_sparse_translator)
+list(APPEND DISTRIBUTED_OP_TRANSLATOR_TEST test_distributed_fused_lamb_init)
 list(APPEND DISTRIBUTED_OP_TRANSLATOR_TEST test_nop_translator)
 list(APPEND DISTRIBUTED_OP_TRANSLATOR_TEST test_partial_send_translator)
 list(APPEND DISTRIBUTED_OP_TRANSLATOR_TEST test_partial_recv_translator)

test/ir/pir/translator/test_distributed_push_sparse_translator.py

@@ -0,0 +1,59 @@
+# Copyright (c) 2024 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.
+
+import unittest
+
+import numpy as np
+import test_op_translator
+
+import paddle
+from paddle.base.framework import (
+    convert_np_dtype_to_dtype_,
+)
+from paddle.base.layer_helper import LayerHelper
+
+
+class TestDistributedPushSparseOpTranslator(
+    test_op_translator.TestOpTranslator
+):
+    def append_op(self):
+        self.op_type = "distributed_push_sparse"
+        ids = paddle.ones(shape=(1, 1), dtype='float32')
+        shows = paddle.ones(shape=(1, 1), dtype='float32')
+        clicks = paddle.ones(shape=(1, 1), dtype='float32')
+        output = paddle.ones(shape=(1, 1), dtype='float32')
+        attrs = {
+            'table_id': 0,
+            'size': 8,
+            'is_distributed': False,
+            'push_sparse_version': 'push_sparse',
+            'padding_idx': -1,
+            'dtype': convert_np_dtype_to_dtype_(np.float32),
+            'is_test': False,
+            'use_cvm_op': False,
+        }
+        helper = LayerHelper(self.op_type)
+        helper.append_op(
+            type=self.op_type,
+            inputs={"Ids": [ids], "Shows": shows, "Clicks": clicks},
+            outputs={"Outputs": [output]},
+            attrs=attrs,
+        )
+
+    def test_translator(self):
+        self.check()
+
+
+if __name__ == "__main__":
+    unittest.main()