【PIR Dist Op Reg No.10】 reg pull_gpups_sparse (#62935)

xiaoyewww · web-flow · commit cf0f65b3a8d3 · 2024-04-25T10:51:00.000+08:00
* feat(pir): reg pull_gpups_sparse

* feat(pir): reg pull_gpups_sparse

* feat(pir): reg pull_gpups_sparse

* feat(pir): reg pull_gpups_sparse

* feat(pir): reg pull_gpups_sparse

* feat(pir): reg pull_gpups_sparse

* feat(pir): reg pull_gpups_sparse

* feat(pir): reg pull_gpups_sparse

* feat(pir): reg pull_gpups_sparse

* feat(pir): reg pull_gpups_sparse

* feat(pir): reg pull_gpups_sparse

* feat(pir): reg pull_gpups_sparse

* feat(pir): reg pull_gpups_sparse

* feat(pir): reg pull_gpups_sparse

* feat(pir): reg pull_gpups_sparse

* feat(pir): reg pull_gpups_sparse

* feat(pir): reg pull_gpups_sparse

* feat(pir): reg pull_gpups_sparse

* feat(pir): reg pull_gpups_sparse

* feat(pir): reg pull_gpups_sparse

* feat(pir): reg pull_gpups_sparse

* feat(pir): reg pull_gpups_sparse

* feat(pir): reg pull_gpups_sparse

* feat(pir): reg pull_gpups_sparse
diff --git a/paddle/fluid/ir_adaptor/translator/op_compat_gen.py b/paddle/fluid/ir_adaptor/translator/op_compat_gen.py
@@ -169,6 +169,10 @@ def insert_new_mutable_attributes(
         {"out_grad_in": "Out@GRAD", "out_grad_out": "Out@GRAD"}
     )
 
+    op_arg_name_mappings['push_gpups_sparse'].update(
+        {"out_grad": "Out@GRAD", "out_grad_grad": "Out@GRAD"}
+    )
+
     op_name_normalizer_template = env.get_template("op_compat_info.cc.j2")
     with open(output_source_file, 'wt') as f:
         op_compat_definition = op_name_normalizer_template.render(
diff --git a/paddle/fluid/pir/dialect/op_generator/ops_api_gen.py b/paddle/fluid/pir/dialect/op_generator/ops_api_gen.py
@@ -166,6 +166,10 @@
     'lars_momentum_',
     'max_pool2d_v2',
     'partial_sum',
+    'pull_gpups_sparse',
+    'pull_gpups_sparse_',
+    'push_gpups_sparse',
+    'push_gpups_sparse_',
     'random_routing',
     'rank_attention',
     'rnn_',
diff --git a/paddle/fluid/pir/dialect/operator/interface/parse_kernel_key.cc b/paddle/fluid/pir/dialect/operator/interface/parse_kernel_key.cc
@@ -40,6 +40,10 @@ KernelKeyTuple Nop_OpParseKernelKey(pir::Operation* op) {
   return {phi::DataType::FLOAT32, phi::Backend::UNDEFINED};
 }
 
+KernelKeyTuple PullGpupsSparseKernelKey(pir::Operation* op) {
+  return {phi::DataType::FLOAT32, phi::Backend::UNDEFINED};
+}
+
 }  // namespace paddle::dialect
 
 IR_DEFINE_EXPLICIT_TYPE_ID(paddle::dialect::ParseKernelKeyInterface)
diff --git a/paddle/fluid/pir/dialect/operator/interface/parse_kernel_key.h b/paddle/fluid/pir/dialect/operator/interface/parse_kernel_key.h
@@ -63,6 +63,8 @@ KernelKeyTuple NopOpParseKernelKey(pir::Operation *op);
 
 KernelKeyTuple Nop_OpParseKernelKey(pir::Operation *op);
 
+KernelKeyTuple PullGpupsSparseKernelKey(pir::Operation *op);
+
 }  // namespace dialect
 }  // namespace paddle
 
diff --git a/paddle/fluid/pir/dialect/operator/ir/ops.yaml b/paddle/fluid/pir/dialect/operator/ir/ops.yaml
@@ -1351,6 +1351,17 @@
     func : prune_gate_by_capacity
     data_type : gate_idx
 
+- op : pull_gpups_sparse
+  args : (Tensor w, Tensor[] ids, int[] size={}, bool is_sparse=false, bool is_distributed=false)
+  output : Tensor[](out){ids.size()}
+  infer_meta :
+    func : PullGpupsSparseInferMeta
+  kernel :
+    func : pull_gpups_sparse
+    data_type : PullGpupsSparseKernelKey
+  optional : w
+  backward: push_gpups_sparse
+
 - op : push_dense
   args : (Tensor[] ids, int table_id = -1, float scale_data_norm = -1.0f, str[] input_names = {})
   output :
diff --git a/paddle/phi/api/yaml/backward.yaml b/paddle/phi/api/yaml/backward.yaml
@@ -1870,6 +1870,16 @@
     data_type : x
   optional : boxes_num
 
+- backward_op : push_gpups_sparse
+  forward : pull_gpups_sparse (Tensor w, Tensor[] ids, int[] size={}, bool is_sparse=false, bool is_distributed=false) -> Tensor[](out)
+  args : (Tensor[] ids, Tensor[] out_grad, int[] size, bool is_sparse, bool is_distributed)
+  output : Tensor[](out_grad_grad){out_grad.size()}
+  infer_meta :
+    func : PushGpupsSparseInferMeta
+  kernel :
+    func : push_gpups_sparse
+  inplace : (out_grad -> out_grad_grad)
+
 - backward_op : put_along_axis_grad
   forward : put_along_axis (Tensor arr, Tensor indices, Tensor values, int axis, str reduce = "assign", bool include_self = true) -> Tensor(out)
   args : (Tensor arr, Tensor indices, Tensor values, Tensor out, Tensor out_grad, int axis, str reduce, bool include_self)
diff --git a/paddle/phi/api/yaml/op_compat.yaml b/paddle/phi/api/yaml/op_compat.yaml
@@ -2671,6 +2671,13 @@
   outputs :
     out : Out
 
+- op : pull_gpups_sparse
+  backward : push_gpups_sparse
+  inputs :
+    {w : W, ids : Ids}
+  outputs :
+    out : Out
+
 - op : push_dense
   inputs :
     ids : Ids
diff --git a/paddle/phi/infermeta/binary.cc b/paddle/phi/infermeta/binary.cc
@@ -2738,6 +2738,55 @@ void PReluInferMeta(const MetaTensor& x,
   out->share_lod(x);
 }
 
+void PullGpupsSparseInferMeta(const MetaTensor& w,
+                              const std::vector<const MetaTensor*>& ids,
+                              const std::vector<int>& size,
+                              bool is_sparse,
+                              bool is_distributed,
+                              std::vector<MetaTensor*> out) {
+  PADDLE_ENFORCE_GE(
+      ids.size(),
+      1UL,
+      phi::errors::InvalidArgument(
+          "Inputs(Ids) of PullGpuPSSparseOp should not be empty."));
+  PADDLE_ENFORCE_GE(
+      out.size(),
+      1UL,
+      phi::errors::InvalidArgument(
+          "Outputs(Out) of PullGpuPSSparseOp should not be empty."));
+  PADDLE_ENFORCE_EQ(
+      ids.size(),
+      size.size(),
+      phi::errors::InvalidArgument("The ids size: %lu must be equal to "
+                                   "the length of embedding size: %lu.",
+                                   ids.size(),
+                                   size.size()));
+  const size_t n_ids = ids.size();
+  std::vector<phi::DDim> outs_dims;
+  outs_dims.resize(n_ids);
+  for (size_t i = 0; i < n_ids; ++i) {
+    int embedding_size = size[i];
+    const auto ids_dims = ids[i]->dims();
+    int ids_rank = ids_dims.size();
+    PADDLE_ENFORCE_EQ(ids_dims[ids_rank - 1],
+                      1,
+                      phi::errors::InvalidArgument(
+                          "Shape error in %lu id, the last dimension of the "
+                          "'Ids' tensor must be 1.",
+                          i));
+    auto out_dim =
+        common::vectorize(common::slice_ddim(ids_dims, 0, ids_rank - 1));
+    out_dim.push_back(embedding_size);
+    outs_dims[i] = common::make_ddim(out_dim);
+  }
+
+  for (size_t i = 0; i < n_ids; ++i) {
+    out[i]->set_dims(outs_dims[i]);
+    out[i]->share_lod(*ids[i], i);
+    out[i]->set_dtype(w.dtype());
+  }
+}
+
 void ApplyPerChannelScaleInferMeta(const MetaTensor& x,
                                    const MetaTensor& scales,
                                    MetaTensor* out) {
diff --git a/paddle/phi/infermeta/binary.h b/paddle/phi/infermeta/binary.h
@@ -473,6 +473,13 @@ void PReluInferMeta(const MetaTensor& x,
                     MetaTensor* out,
                     MetaConfig config = MetaConfig());
 
+void PullGpupsSparseInferMeta(const MetaTensor& w,
+                              const std::vector<const MetaTensor*>& ids,
+                              const std::vector<int>& size,
+                              bool is_sparse,
+                              bool is_distributed,
+                              std::vector<MetaTensor*> out);
+
 void RepeatInterleaveWithTensorIndexInferMeta(const MetaTensor& x,
                                               const MetaTensor& repeats,
                                               int dim,
diff --git a/paddle/phi/infermeta/ternary.cc b/paddle/phi/infermeta/ternary.cc
@@ -1124,6 +1124,13 @@ void PutAlongAxisInferMeta(const MetaTensor& x,
   out->set_dtype(x.dtype());
 }
 
+void PushGpupsSparseInferMeta(const std::vector<const MetaTensor*>& ids,
+                              const std::vector<const MetaTensor*>& out,
+                              const std::vector<int>& size,
+                              bool is_sparse,
+                              bool is_distributed,
+                              std::vector<MetaTensor*> out_grad) {}
+
 void RandomRoutingInferMeta(const MetaTensor& prob,
                             const MetaTensor& topk_value,
                             const MetaTensor& topk_idx,
diff --git a/paddle/phi/infermeta/ternary.h b/paddle/phi/infermeta/ternary.h
@@ -203,6 +203,13 @@ void NllLossRawInferMeta(const MetaTensor& input,
                          MetaTensor* total_weight,
                          MetaConfig config = MetaConfig());
 
+void PushGpupsSparseInferMeta(const std::vector<const MetaTensor*>& ids,
+                              const std::vector<const MetaTensor*>& out,
+                              const std::vector<int>& size,
+                              bool is_sparse,
+                              bool is_distributed,
+                              std::vector<MetaTensor*> out_grad);
+
 void PutAlongAxisInferMeta(const MetaTensor& x,
                            const MetaTensor& index,
                            const MetaTensor& value,
diff --git a/test/ir/pir/translator/CMakeLists.txt b/test/ir/pir/translator/CMakeLists.txt
@@ -27,6 +27,7 @@ list(APPEND DISTRIBUTED_OP_TRANSLATOR_TEST test_push_dense_translator)
 list(APPEND DISTRIBUTED_OP_TRANSLATOR_TEST test_partial_recv_translator)
 list(APPEND DISTRIBUTED_OP_TRANSLATOR_TEST
      test_prune_gate_by_capacity_translator)
+list(APPEND DISTRIBUTED_OP_TRANSLATOR_TEST test_pull_gpups_sparse_translator)
 list(APPEND DISTRIBUTED_OP_TRANSLATOR_TEST test_random_routing_translator)
 list(APPEND DISTRIBUTED_OP_TRANSLATOR_TEST test_limit_by_capacity_translator)
 list(APPEND DISTRIBUTED_OP_TRANSLATOR_TEST test_global_scatter_translator)
diff --git a/test/ir/pir/translator/test_op_translator.py b/test/ir/pir/translator/test_op_translator.py
@@ -63,6 +63,7 @@ def build_model(self):
         with paddle.static.scope_guard(self.new_scope):
             with paddle.static.program_guard(self.main_program):
                 out = self.append_op()
+                out.stop_gradient = False
                 append_backward(out)
 
     def check(self):
diff --git a/test/ir/pir/translator/test_pull_gpups_sparse_translator.py b/test/ir/pir/translator/test_pull_gpups_sparse_translator.py
@@ -0,0 +1,54 @@
+# 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
+#
+#     https://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_translator
+
+import paddle
+from paddle.base import core
+from paddle.base.layer_helper import LayerHelper
+
+
+class TestPullGpupsSparseOpTranslator(
+    test_op_translator.TestOpWithBackwardTranslator
+):
+    def setUp(self):
+        self.place = core.Place()
+        self.place.set_place(paddle.CPUPlace())
+        self.new_scope = paddle.static.Scope()
+        self.main_program = paddle.static.Program()
+        self.forward_op_type = "pull_gpups_sparse"
+        self.backward_op_type = "push_gpups_sparse"
+
+    def append_op(self):
+        self.op_type = "pull_gpups_sparse"
+        ids = paddle.ones(shape=(1,), dtype='int64')
+        out = paddle.ones(shape=(1,), dtype='int64')
+        attrs = {'size': [1], 'is_sparse': False, 'is_distributed': False}
+        helper = LayerHelper(self.op_type)
+        helper.append_op(
+            type=self.op_type,
+            inputs={"Ids": ids},
+            outputs={"Out": out},
+            attrs=attrs,
+        )
+        return out
+
+    def test_translator(self):
+        self.check()
+
+
+if __name__ == "__main__":
+    unittest.main()

Original file line number	Diff line number	Diff line change
`@@ -40,6 +40,10 @@ KernelKeyTuple Nop_OpParseKernelKey(pir::Operation* op) {`
`40`	`40`	`return {phi::DataType::FLOAT32, phi::Backend::UNDEFINED};`
`41`	`41`	`}`
`42`	`42`
	`43`	`+KernelKeyTuple PullGpupsSparseKernelKey(pir::Operation* op) {`
	`44`	`+ return {phi::DataType::FLOAT32, phi::Backend::UNDEFINED};`
	`45`	`+}`
	`46`	`+`
`43`	`47`	`} // namespace paddle::dialect`
`44`	`48`
`45`	`49`	`IR_DEFINE_EXPLICIT_TYPE_ID(paddle::dialect::ParseKernelKeyInterface)`