【PIR Dist Op Reg No.19】 reg pull_box_sparse

paddle/fluid/ir_adaptor/translator/op_compat_gen.py

@@ -168,6 +168,9 @@ def insert_new_mutable_attributes(
     op_arg_name_mappings['push_sparse_v2'].update(
         {"out_grad_in": "Out@GRAD", "out_grad_out": "Out@GRAD"}
     )
+    op_arg_name_mappings['push_box_sparse'].update(
+        {"out_grad_in": "Out@GRAD", "out_grad_out": "Out@GRAD"}
+    )
 
     op_name_normalizer_template = env.get_template("op_compat_info.cc.j2")
     with open(output_source_file, 'wt') as f:

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

@@ -201,6 +201,10 @@
     'push_dense',
     'limit_by_capacity',
     'global_scatter',
+    'pull_box_sparse',
+    'pull_box_sparse_',
+    'push_box_sparse',
+    'push_box_sparse_',
 ]
 
 

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

@@ -1342,6 +1342,26 @@
     func : prune_gate_by_capacity
     data_type : gate_idx
 
+- op : pull_box_sparse
+  args : (Tensor w, Tensor[] ids, bool is_sparse = false, bool is_distributed = false, int size = 1)
+  output : Tensor[](out){ids.size()}
+  infer_meta :
+    func : PullBoxSparseInferMeta
+  kernel :
+    func : pull_box_sparse
+    data_type : ids
+
+- op : push_box_sparse
+  args : (Tensor[] ids, Tensor[] out_grad_in, bool is_sparse = false, bool is_distributed = false, int size = 1)
+  output : Tensor[](out_grad_out){out_grad_in.size()}
+  infer_meta :
+    func : UnchangedMultiInferMeta
+    param : [out_grad_in]
+  kernel :
+    func : push_box_sparse
+    data_type : out_grad_in
+  inplace : (out_grad_in -> out_grad_out)
+
 - op : push_dense
   args : (Tensor[] ids, int table_id = -1, float scale_data_norm = -1.0f, str[] input_names = {})
   output :

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

@@ -106,6 +106,8 @@ const std::unordered_set<std::string> LegacyOpList = {
     CReduceMinOp::name(),
     CReduceProdOp::name(),
     CScatterOp::name(),
+    PullBoxSparseOp::name(),
+    PushBoxSparseOp::name(),
     PushSparseV2Op::name(),
     PartialSendOp::name(),
     PartialRecvOp::name()};

paddle/phi/api/yaml/op_compat.yaml

@@ -2659,6 +2659,24 @@
   outputs :
     out : Out
 
+- op : pull_box_sparse
+  inputs :
+    { w : W, ids: Ids}
+  outputs :
+    out : Out
+  attrs :
+    sparse : is_sparse
+  extra :
+    attrs : [bool is_sparse = false, bool is_distributed = false, int size = 1]
+
+- op : push_box_sparse
+  inputs :
+    ids: Ids
+  outputs :
+    out : Out
+  attrs :
+    sparse : is_sparse
+
 - op : push_dense
   inputs :
     ids : Ids

paddle/phi/infermeta/binary.cc

@@ -2893,6 +2893,32 @@ void PruneGateByCapacityInferMeta(const MetaTensor& gate_idx,
   new_gate_idx->set_dtype(gate_idx.dtype());
 }
 
+void PullBoxSparseInferMeta(const MetaTensor& w,
+                            const std::vector<const MetaTensor*>& ids,
+                            bool is_sparse,
+                            bool is_distributed,
+                            int size,
+                            std::vector<MetaTensor*> out) {
+  auto hidden_size = static_cast<int64_t>(size);
+  const size_t n_ids = ids.size();
+  for (size_t i = 0; i < n_ids; ++i) {
+    MetaTensor* output = out[i];
+    auto ids_dims = ids[i]->dims();
+    int ids_rank = ids_dims.size();
+    PADDLE_ENFORCE_EQ(ids_dims[ids_rank - 1],
+                      1UL,
+                      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(hidden_size);
+    output->set_dims(common::make_ddim(out_dim));
+    output->share_lod(*ids[i]);
+  }
+}
+
 void RepeatInterleaveWithTensorIndexInferMeta(const MetaTensor& x,
                                               const MetaTensor& repeats,
                                               int dim,

paddle/phi/infermeta/binary.h

@@ -469,6 +469,13 @@ void PReluInferMeta(const MetaTensor& x,
                     MetaTensor* out,
                     MetaConfig config = MetaConfig());
 
+void PullBoxSparseInferMeta(const MetaTensor& w,
+                            const std::vector<const MetaTensor*>& ids,
+                            bool is_sparse,
+                            bool is_distributed,
+                            int size,
+                            std::vector<MetaTensor*> out);
+
 void RepeatInterleaveWithTensorIndexInferMeta(const MetaTensor& x,
                                               const MetaTensor& repeats,
                                               int dim,

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_box_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)

test/ir/pir/translator/test_pull_box_sparse_translator.py

@@ -0,0 +1,50 @@
+# 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 test_op_translator
+
+import paddle
+from paddle.base.layer_helper import LayerHelper
+
+
+class TestPullBoxSparseOpTranslator(
+    test_op_translator.TestOpWithBackwardTranslator
+):
+    def append_op(self):
+        self.forward_op_type = "pull_box_sparse"
+        self.backward_op_type = "push_box_sparse"
+        ids = paddle.ones(shape=(1, 1), dtype='float32')
+        w = paddle.ones(shape=(1, 1), dtype='float32')
+        out = paddle.ones(shape=(1, 1), dtype='float32')
+        attrs = {
+            'is_sparse': False,
+            'is_distributed': False,
+            'size': 1,
+        }
+        forward_helper = LayerHelper(self.forward_op_type)
+        forward_helper.append_op(
+            type=self.forward_op_type,
+            inputs={"W": w, "Ids": [ids]},
+            outputs={"Out": [out]},
+            attrs=attrs,
+        )
+
+    def test_translator(self):
+        self.check()
+
+
+if __name__ == "__main__":
+    unittest.main()