【PIR Dist Op Reg No.15】 reg push_dense (#62505)

enkilee · web-flow · commit 11ba107a6611 · 2024-03-26T16:34:48.000+08:00
* fix

* fix

* fix

* fix

* fix

* fix

* fix
diff --git a/paddle/fluid/pir/dialect/op_generator/ops_api_gen.py b/paddle/fluid/pir/dialect/op_generator/ops_api_gen.py
@@ -192,6 +192,7 @@
     'partial_allgather_',
     'nop',
     'nop_',
+    'push_dense',
     'limit_by_capacity',
     'global_scatter',
 ]
diff --git a/paddle/fluid/pir/dialect/operator/ir/ops.yaml b/paddle/fluid/pir/dialect/operator/ir/ops.yaml
@@ -1305,6 +1305,16 @@
     func : prune_gate_by_capacity
     data_type : gate_idx
 
+- op : push_dense
+  args : (Tensor[] ids, int table_id = -1, float scale_data_norm = -1.0f, str[] input_names = {})
+  output :
+  infer_meta :
+    func : PushDenseInferMeta
+    param : [ids, table_id, scale_data_norm, input_names]
+  kernel :
+    func : push_dense
+    data_type : DataType::FLOAT32
+
 - 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()}
diff --git a/paddle/fluid/pir/dialect/operator/utils/utils.cc b/paddle/fluid/pir/dialect/operator/utils/utils.cc
@@ -64,6 +64,7 @@ const std::unordered_set<std::string> LegacyOpList = {
     CSoftmaxWithCrossEntropyOp::name(),
     CSoftmaxWithCrossEntropyGradOp::name(),
     CSplitOp::name(),
+    PushDenseOp::name(),
     SeedOp::name(),
     ShareDataOp::name(),
     SparseMomentumOp::name(),
diff --git a/paddle/fluid/primitive/codegen/gen.py b/paddle/fluid/primitive/codegen/gen.py
@@ -53,6 +53,7 @@
     "embedding_grad",
     "full",
     "partial_send",
+    "push_dense",
 ]
 
 # prim op with one input and one output, with no attribute
diff --git a/paddle/phi/api/yaml/op_compat.yaml b/paddle/phi/api/yaml/op_compat.yaml
@@ -2641,6 +2641,12 @@
   outputs :
     out : Out
 
+- op : push_dense
+  inputs :
+    ids : Ids
+  attrs :
+    {table_id : TableId, scale_data_norm : ScaleDataNorm, input_names: InputNames}
+
 - op : push_sparse_v2
   inputs :
     { x : Ids, W : w}
diff --git a/paddle/phi/infermeta/unary.cc b/paddle/phi/infermeta/unary.cc
@@ -3377,6 +3377,17 @@ void PoolInferMeta(const MetaTensor& x,
   out->set_dtype(x.dtype());
 }
 
+void PushDenseInferMeta(const std::vector<const MetaTensor*>& ids,
+                        int table_id,
+                        float scale_data_norm,
+                        const std::vector<std::string>& input_names) {
+  auto ids_num = ids.size();
+  PADDLE_ENFORCE_GE(ids_num,
+                    1UL,
+                    phi::errors::InvalidArgument(
+                        "Input(Ids) of PushDenseOp can not be null."));
+}
+
 void RealAndImagInferMeta(const MetaTensor& x, MetaTensor* out) {
   out->set_dims(x.dims());
   out->set_dtype(dtype::ToReal(x.dtype()));
diff --git a/paddle/phi/infermeta/unary.h b/paddle/phi/infermeta/unary.h
@@ -508,6 +508,11 @@ void PSendInferMeta(const MetaTensor& x, int peer);
 
 void PSendArrayInferMeta(const MetaTensor& x, int peer);
 
+void PushDenseInferMeta(const std::vector<const MetaTensor*>& ids,
+                        int table_id,
+                        float scale_data_norm,
+                        const std::vector<std::string>& input_names);
+
 void SendV2InferMeta(const int peer, const int ring_id);
 
 void QrInferMeta(const MetaTensor& x,
diff --git a/test/ir/pir/translator/CMakeLists.txt b/test/ir/pir/translator/CMakeLists.txt
@@ -22,6 +22,7 @@ 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_allgather_translator)
 list(APPEND DISTRIBUTED_OP_TRANSLATOR_TEST test_partial_send_translator)
+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)
diff --git a/test/ir/pir/translator/test_push_dense_translator.py b/test/ir/pir/translator/test_push_dense_translator.py
@@ -0,0 +1,45 @@
+# 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.layer_helper import LayerHelper
+
+
+class TestPushDenseOpTranslator(test_op_translator.TestOpTranslator):
+    def append_op(self):
+        self.op_type = "push_dense"
+        ids = paddle.ones(shape=(100, 2, 3), dtype='float32')
+        input_names = []
+        attrs = {
+            'TableId': 1,
+            'ScaleDataNorm': -1,
+            'InputNames': input_names,
+        }
+        helper = LayerHelper(self.op_type)
+        helper.append_op(
+            type=self.op_type,
+            inputs={"Ids": [ids]},
+            attrs=attrs,
+        )
+
+    def test_translator(self):
+        self.check()
+
+
+if __name__ == "__main__":
+    unittest.main()

Original file line number	Diff line number	Diff line change
`@@ -192,6 +192,7 @@`
`192`	`192`	`'partial_allgather_',`
`193`	`193`	`'nop',`
`194`	`194`	`'nop_',`
	`195`	`+ 'push_dense',`
`195`	`196`	`'limit_by_capacity',`
`196`	`197`	`'global_scatter',`
`197`	`198`	`]`
Original file line number	Diff line number	Diff line change
`@@ -53,6 +53,7 @@`
`53`	`53`	`"embedding_grad",`
`54`	`54`	`"full",`
`55`	`55`	`"partial_send",`
	`56`	`+ "push_dense",`
`56`	`57`	`]`
`57`	`58`
`58`	`59`	`# prim op with one input and one output, with no attribute`