diff --git a/paddle/fluid/ir_adaptor/translator/op_compat_gen.py b/paddle/fluid/ir_adaptor/translator/op_compat_gen.py index 89e3f7b7583f64..2e6f870d5449d5 100644 --- a/paddle/fluid/ir_adaptor/translator/op_compat_gen.py +++ b/paddle/fluid/ir_adaptor/translator/op_compat_gen.py @@ -170,6 +170,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_arg_name_mappings['push_gpups_sparse'].update( {"out_grad": "Out@GRAD", "out_grad_grad": "Out@GRAD"} ) diff --git a/paddle/fluid/pir/dialect/op_generator/ops_api_gen.py b/paddle/fluid/pir/dialect/op_generator/ops_api_gen.py index 2647b579f2bc73..1e6799fccad854 100644 --- a/paddle/fluid/pir/dialect/op_generator/ops_api_gen.py +++ b/paddle/fluid/pir/dialect/op_generator/ops_api_gen.py @@ -206,6 +206,10 @@ 'push_dense', 'limit_by_capacity', 'global_scatter', + 'pull_box_sparse', + 'pull_box_sparse_', + 'push_box_sparse', + 'push_box_sparse_', ] diff --git a/paddle/fluid/pir/dialect/operator/ir/ops.yaml b/paddle/fluid/pir/dialect/operator/ir/ops.yaml index 7b24d6c4fec9a0..f19fcd340413dc 100644 --- a/paddle/fluid/pir/dialect/operator/ir/ops.yaml +++ b/paddle/fluid/pir/dialect/operator/ir/ops.yaml @@ -1351,6 +1351,15 @@ 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 : pull_gpups_sparse args : (Tensor w, Tensor[] ids, int[] size={}, bool is_sparse=false, bool is_distributed=false) output : Tensor[](out){ids.size()} diff --git a/paddle/fluid/pir/dialect/operator/ir/ops_backward.yaml b/paddle/fluid/pir/dialect/operator/ir/ops_backward.yaml index f407cf00c504f7..1b3b24153b34a6 100644 --- a/paddle/fluid/pir/dialect/operator/ir/ops_backward.yaml +++ b/paddle/fluid/pir/dialect/operator/ir/ops_backward.yaml @@ -659,6 +659,18 @@ func : prod_grad composite: prod_grad(x, out, out_grad, dims, keep_dim, reduce_all, x_grad) +- backward_op : push_box_sparse + forward : pull_box_sparse (Tensor w, Tensor[] ids, bool is_sparse = false, bool is_distributed = false, int size = 1) -> Tensor[](out){ids.size()} + 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) + - backward_op : rank_attention_grad forward : rank_attention (Tensor x, Tensor rank_offset, Tensor rank_param, int max_rank = 3, int max_size = 0) -> Tensor(input_help), Tensor(out), Tensor(ins_rank) args : (Tensor x, Tensor rank_offset, Tensor rank_param, Tensor input_help, Tensor ins_rank, Tensor out_grad, int max_rank = 3, int max_size = 0) diff --git a/paddle/fluid/pir/dialect/operator/utils/utils.cc b/paddle/fluid/pir/dialect/operator/utils/utils.cc index 551acf25fa16b5..711bd2f4e5f18a 100644 --- a/paddle/fluid/pir/dialect/operator/utils/utils.cc +++ b/paddle/fluid/pir/dialect/operator/utils/utils.cc @@ -105,6 +105,8 @@ const std::unordered_set LegacyOpList = { CReduceMinOp::name(), CReduceProdOp::name(), CScatterOp::name(), + PullBoxSparseOp::name(), + PushBoxSparseOp::name(), PushSparseV2Op::name(), PartialSendOp::name(), PartialRecvOp::name()}; diff --git a/paddle/phi/api/yaml/op_compat.yaml b/paddle/phi/api/yaml/op_compat.yaml index 883f6ec122f43d..8c08948524c8f3 100755 --- a/paddle/phi/api/yaml/op_compat.yaml +++ b/paddle/phi/api/yaml/op_compat.yaml @@ -2685,6 +2685,16 @@ 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 : pull_gpups_sparse backward : push_gpups_sparse inputs : @@ -2700,6 +2710,14 @@ extra : attrs : [int embedding_dim = 11, int table_id = 0, str accessor_class = "", str ctr_label_name = "", int padding_id = 0, bool scale_sparse_grad = true, 'str[] input_names = {}', bool is_distributed = true] +- op : push_box_sparse + inputs : + ids: Ids + outputs : + out : Out + attrs : + sparse : is_sparse + - op : push_dense inputs : ids : Ids diff --git a/paddle/phi/infermeta/binary.cc b/paddle/phi/infermeta/binary.cc index fe14db9ccb415e..421ec8357c1a20 100644 --- a/paddle/phi/infermeta/binary.cc +++ b/paddle/phi/infermeta/binary.cc @@ -2989,6 +2989,33 @@ void PruneGateByCapacityInferMeta(const MetaTensor& gate_idx, new_gate_idx->set_dtype(gate_idx.dtype()); } +void PullBoxSparseInferMeta(const MetaTensor& w, + const std::vector& ids, + bool is_sparse, + bool is_distributed, + int size, + std::vector out) { + auto hidden_size = static_cast(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]); + output->set_dtype(w.dtype()); + } +} + void RepeatInterleaveWithTensorIndexInferMeta(const MetaTensor& x, const MetaTensor& repeats, int dim, diff --git a/paddle/phi/infermeta/binary.h b/paddle/phi/infermeta/binary.h index c43b4e852a19ea..27f387eedb2aef 100644 --- a/paddle/phi/infermeta/binary.h +++ b/paddle/phi/infermeta/binary.h @@ -474,6 +474,13 @@ void PReluInferMeta(const MetaTensor& x, MetaTensor* out, MetaConfig config = MetaConfig()); +void PullBoxSparseInferMeta(const MetaTensor& w, + const std::vector& ids, + bool is_sparse, + bool is_distributed, + int size, + std::vector out); + void PullGpupsSparseInferMeta(const MetaTensor& w, const std::vector& ids, const std::vector& size, diff --git a/test/ir/pir/translator/test_pull_box_sparse_translator.py b/test/ir/pir/translator/test_pull_box_sparse_translator.py new file mode 100644 index 00000000000000..f691892adc4f41 --- /dev/null +++ b/test/ir/pir/translator/test_pull_box_sparse_translator.py @@ -0,0 +1,51 @@ +# 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, + ) + return out + + def test_translator(self): + self.check() + + +if __name__ == "__main__": + unittest.main()