【PIR Dist Op Reg No.1】 reg push_sparse_v2

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

@@ -139,6 +139,7 @@
     'uniform_random_batch_size_like',
     'c_reduce_min',
     'c_reduce_min_',
+    'pull_sparse_v2',
 ]
 
 

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

@@ -968,6 +968,17 @@
     func : prod
   backward : prod_grad
 
+- op : pull_sparse_v2
+  args : (Tensor[] ids, Tensor[] w, 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){ids.size()}
+  infer_meta :
+    func : PullSparseV2InferMeta
+    param : [ids, w, embeddingdim, tableid, accessorclass, ctrlabelname, paddingid, scalesparsegrad, inputnames, is_distributed]
+  kernel :
+    func : pull_sparse_v2
+    param : [ids, w, embeddingdim, tableid, accessorclass, ctrlabelname, paddingid, scalesparsegrad, inputnames, is_distributed]
+    data_type : ids
+
 - op : randint
   args : (int low, int high, IntArray shape, DataType dtype=DataType::INT64, Place place={})
   output : Tensor(out)

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

@@ -59,7 +59,8 @@ const std::unordered_set<std::string> LegacyOpList = {
     RowConvGradOp::name(),
     SoftReluOp::name(),
     SoftReluGradOp::name(),
-    CReduceMinOp::name()};
+    CReduceMinOp::name(),
+    PullSparseV2Op::name()};
 
 const std::unordered_set<std::string> OneDNNLegacyOpList = {};
 enum class AttrType {

paddle/phi/api/yaml/op_compat.yaml

@@ -2441,6 +2441,14 @@
   outputs :
     out : Out
 
+- op : pull_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 :

paddle/phi/infermeta/binary.cc

@@ -2510,6 +2510,24 @@ void MvInferMeta(const MetaTensor& x, const MetaTensor& vec, MetaTensor* out) {
   out->share_lod(x);
 }
 
+void PullSparseV2InferMeta(const MetaTensor& ids,
+                           const MetaTensor& w,
+                           int embeddingdim,
+                           int tableid,
+                           const std::string accessorclass,
+                           const std::string ctrlabelname,
+                           int paddingid,
+                           bool scalesparsegrad,
+                           const std::vector<std::string>& inputnames,
+                           bool is_distributed,
+                           MetaTensor* out) {
+  const auto& ids_dims = ids.dims();
+  auto output_dims = common::vectorize(ids_dims);
+  output_dims.push_back(ids_dims[1]);
+  out->set_dims(common::make_ddim(output_dims));
+  out->share_lod(ids);
+}
+
 void PReluInferMeta(const MetaTensor& x,
                     const MetaTensor& alpha,
                     const std::string& data_format,

paddle/phi/infermeta/binary.h

@@ -443,6 +443,18 @@ void PriorBoxInferMeta(const MetaTensor& input,
                        MetaTensor* out,
                        MetaTensor* var);
 
+void PullSparseV2InferMeta(const MetaTensor& ids,
+                           const MetaTensor& w,
+                           int embeddingdim,
+                           int tableid,
+                           const std::string accessorclass,
+                           const std::string ctrlabelname,
+                           int paddingid,
+                           bool scalesparsegrad,
+                           const std::vector<std::string>& inputnames,
+                           bool is_distributed,
+                           MetaTensor* out);
+
 void SearchsortedInferMeta(const MetaTensor& sorted_sequence,
                            const MetaTensor& value,
                            bool out_int32,

test/ir/pir/translator/test_pull_sparse_v2_translate.py

@@ -0,0 +1,52 @@
+# Copyright (c) 2023 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 test_op_transcriber
+
+import paddle
+from paddle.base.layer_helper import LayerHelper
+
+
+class TestPullSparseV2OpTranscriber(test_op_transcriber.TestOpTranscriber):
+    def append_op(self):
+        self.op_type = "push_sparse_v2"
+        ids = paddle.ones(shape=(100, 2, 3), dtype='float32')
+        w = paddle.ones(shape=(100, 2, 3), dtype='float32')
+        out = paddle.ones(shape=(100, 2, 3), dtype='float32')
+        attrs = {
+            'embeddingdim': 11,
+            'tableid': 0,
+            'accessorclass': "",
+            'ctrlabelname': "",
+            'paddingid': 0,
+            'scalesparsegrad': True,
+            'inputnames': [],
+            'is_distributed': True,
+        }
+        helper = LayerHelper(self.op_type)
+        helper.append_op(
+            type=self.op_type,
+            inputs={"Ids": ids, "W": w},
+            outputs={"Out": out},
+            attrs=attrs,
+        )
+
+    def test_translator(self):
+        self.check()
+
+
+if __name__ == "__main__":
+    unittest.main()