From 0b6083930a39a0e4316eb60571e5060b9c6c642f Mon Sep 17 00:00:00 2001 From: Scotty Date: Fri, 1 Sep 2023 15:08:20 +0800 Subject: [PATCH] =?UTF-8?q?=E3=80=90Complex=20op=E3=80=91add=20complex=20s?= =?UTF-8?q?upport=20for=20index=5Fselect=20and=20index=5Fsample=20(#56457)?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit * support index_select op * index_sample in cpu * support index_sample in gpu * change data_transform * fix api gen and use skip_transform in yaml --- paddle/fluid/eager/grad_node_info.cc | 1 + paddle/phi/api/yaml/backward.yaml | 4 +++ paddle/phi/api/yaml/generator/api_base.py | 25 ++++++++++---- paddle/phi/api/yaml/ops.yaml | 4 +++ .../kernels/cpu/index_sample_grad_kernel.cc | 4 ++- paddle/phi/kernels/cpu/index_sample_kernel.cc | 4 ++- .../kernels/cpu/index_select_grad_kernel.cc | 2 ++ paddle/phi/kernels/cpu/index_select_kernel.cc | 2 ++ .../kernels/gpu/index_sample_grad_kernel.cu | 4 ++- paddle/phi/kernels/gpu/index_sample_kernel.cu | 4 ++- .../kernels/gpu/index_select_grad_kernel.cu | 2 ++ paddle/phi/kernels/gpu/index_select_kernel.cu | 2 ++ python/paddle/tensor/search.py | 26 ++++++++++++--- test/legacy_test/test_index_sample_op.py | 27 +++++++++++++++ test/legacy_test/test_index_select_op.py | 33 +++++++++++++++++-- 15 files changed, 127 insertions(+), 17 deletions(-) mode change 100755 => 100644 paddle/phi/kernels/gpu/index_sample_grad_kernel.cu mode change 100755 => 100644 paddle/phi/kernels/gpu/index_sample_kernel.cu diff --git a/paddle/fluid/eager/grad_node_info.cc b/paddle/fluid/eager/grad_node_info.cc index 718bc51b4f50a..6aba5f31b6d3f 100644 --- a/paddle/fluid/eager/grad_node_info.cc +++ b/paddle/fluid/eager/grad_node_info.cc @@ -548,6 +548,7 @@ void GradNodeBase::HandleComplexGradToRealGrad( for (size_t slot_id = 0; slot_id < out_grads->size(); slot_id++) { const std::vector& slot_out_grads = (*out_grads)[slot_id]; for (size_t rank_id = 0; rank_id < slot_out_grads.size(); rank_id++) { + if (bwd_out_meta_[slot_id].size() == 0) continue; const GradSlotMeta& slot_meta = bwd_out_meta_[slot_id][rank_id]; PADDLE_ENFORCE( diff --git a/paddle/phi/api/yaml/backward.yaml b/paddle/phi/api/yaml/backward.yaml index 4031ecf034be2..f82623c06f82c 100644 --- a/paddle/phi/api/yaml/backward.yaml +++ b/paddle/phi/api/yaml/backward.yaml @@ -1120,6 +1120,8 @@ func : index_sample_grad data_type : out_grad no_need_buffer : x + data_transform : + skip_transform : index - backward_op : index_select_grad forward : index_select(Tensor x, Tensor index, int axis) -> Tensor(out) @@ -1132,6 +1134,8 @@ func : index_select_grad data_type : out_grad no_need_buffer : x + data_transform : + skip_transform : index - backward_op : index_select_strided_grad forward : index_select_strided(Tensor x, int64_t index, int axis) -> Tensor(out) diff --git a/paddle/phi/api/yaml/generator/api_base.py b/paddle/phi/api/yaml/generator/api_base.py index e3aea13afe133..cbf4ed1dab837 100644 --- a/paddle/phi/api/yaml/generator/api_base.py +++ b/paddle/phi/api/yaml/generator/api_base.py @@ -14,11 +14,23 @@ import collections import re +from typing import List PREFIX_TENSOR_NAME = 'input_' PREFIX_META_TENSOR_NAME = 'meta_' +def parse_plain_list(s: str, sep=",") -> List[str]: + """Copy from `paddle/fluid/operators/generator/parse_utils.py`""" + if sep == ",": + patten = re.compile(r',(?![^{]*\})') # support "int[] a={1,2}" + items = re.split(patten, s.strip()) + items = [x.strip() for x in items] + return items + else: + return [item.strip() for item in s.strip().split(sep)] + + class BaseAPI: def __init__(self, api_item_yaml): self.api = self.get_api_name(api_item_yaml) @@ -367,14 +379,13 @@ def parse_data_transform(self, api_item_yaml): data_transform = {'skip_transform': [], 'support_trans_dtype': []} if 'data_transform' in api_item_yaml: if 'skip_transform' in api_item_yaml['data_transform']: - data_transform['skip_transform'] = api_item_yaml[ - 'data_transform' - ]['skip_transform'] + data_transform['skip_transform'] = parse_plain_list( + api_item_yaml['data_transform']['skip_transform'] + ) if 'support_trans_dtype' in api_item_yaml['data_transform']: - data_transform['support_trans_dtype'] = api_item_yaml[ - 'data_transform' - ]['support_trans_dtype'] - + data_transform['support_trans_dtype'] = parse_plain_list( + api_item_yaml['data_transform']['support_trans_dtype'] + ) return data_transform # Override by child class diff --git a/paddle/phi/api/yaml/ops.yaml b/paddle/phi/api/yaml/ops.yaml index fbc058ff64e78..4e67144ba8a89 100644 --- a/paddle/phi/api/yaml/ops.yaml +++ b/paddle/phi/api/yaml/ops.yaml @@ -1238,6 +1238,8 @@ func : index_sample data_type : x backward : index_sample_grad + data_transform : + skip_transform : index - op : index_select args : (Tensor x, Tensor index, int axis = 0) @@ -1248,6 +1250,8 @@ func : index_select data_type : x backward : index_select_grad + data_transform : + skip_transform : index - op : index_select_strided args : (Tensor x, int64_t index, int axis = 0) diff --git a/paddle/phi/kernels/cpu/index_sample_grad_kernel.cc b/paddle/phi/kernels/cpu/index_sample_grad_kernel.cc index fe5d769da4e8f..c49a4531aea7a 100644 --- a/paddle/phi/kernels/cpu/index_sample_grad_kernel.cc +++ b/paddle/phi/kernels/cpu/index_sample_grad_kernel.cc @@ -100,4 +100,6 @@ PD_REGISTER_KERNEL(index_sample_grad, float, double, int, - int64_t) {} + int64_t, + phi::dtype::complex, + phi::dtype::complex) {} diff --git a/paddle/phi/kernels/cpu/index_sample_kernel.cc b/paddle/phi/kernels/cpu/index_sample_kernel.cc index 1a9c495221fe9..02f3afcb67b6e 100644 --- a/paddle/phi/kernels/cpu/index_sample_kernel.cc +++ b/paddle/phi/kernels/cpu/index_sample_kernel.cc @@ -115,4 +115,6 @@ PD_REGISTER_KERNEL(index_sample, float, double, int, - int64_t) {} + int64_t, + phi::dtype::complex, + phi::dtype::complex) {} diff --git a/paddle/phi/kernels/cpu/index_select_grad_kernel.cc b/paddle/phi/kernels/cpu/index_select_grad_kernel.cc index fc2e88b2f9906..4e53056bd117f 100644 --- a/paddle/phi/kernels/cpu/index_select_grad_kernel.cc +++ b/paddle/phi/kernels/cpu/index_select_grad_kernel.cc @@ -60,5 +60,7 @@ PD_REGISTER_KERNEL(index_select_grad, float, double, phi::dtype::bfloat16, + phi::dtype::complex, + phi::dtype::complex, int, int64_t) {} diff --git a/paddle/phi/kernels/cpu/index_select_kernel.cc b/paddle/phi/kernels/cpu/index_select_kernel.cc index 5c905dd7cd161..a9e64d1d183b4 100644 --- a/paddle/phi/kernels/cpu/index_select_kernel.cc +++ b/paddle/phi/kernels/cpu/index_select_kernel.cc @@ -59,5 +59,7 @@ PD_REGISTER_KERNEL(index_select, float, double, phi::dtype::bfloat16, + phi::dtype::complex, + phi::dtype::complex, int, int64_t) {} diff --git a/paddle/phi/kernels/gpu/index_sample_grad_kernel.cu b/paddle/phi/kernels/gpu/index_sample_grad_kernel.cu old mode 100755 new mode 100644 index ee6602f0dacc2..ff23dc0c394b9 --- a/paddle/phi/kernels/gpu/index_sample_grad_kernel.cu +++ b/paddle/phi/kernels/gpu/index_sample_grad_kernel.cu @@ -135,4 +135,6 @@ PD_REGISTER_KERNEL(index_sample_grad, float, double, int, - int64_t) {} + int64_t, + phi::dtype::complex, + phi::dtype::complex) {} diff --git a/paddle/phi/kernels/gpu/index_sample_kernel.cu b/paddle/phi/kernels/gpu/index_sample_kernel.cu old mode 100755 new mode 100644 index 9175f783b8baf..c60abdb3372cb --- a/paddle/phi/kernels/gpu/index_sample_kernel.cu +++ b/paddle/phi/kernels/gpu/index_sample_kernel.cu @@ -108,4 +108,6 @@ PD_REGISTER_KERNEL(index_sample, float, double, int, - int64_t) {} + int64_t, + phi::dtype::complex, + phi::dtype::complex) {} diff --git a/paddle/phi/kernels/gpu/index_select_grad_kernel.cu b/paddle/phi/kernels/gpu/index_select_grad_kernel.cu index 03f74888fca0f..6d0ba9e5bd4ef 100644 --- a/paddle/phi/kernels/gpu/index_select_grad_kernel.cu +++ b/paddle/phi/kernels/gpu/index_select_grad_kernel.cu @@ -132,5 +132,7 @@ PD_REGISTER_KERNEL(index_select_grad, double, phi::dtype::float16, phi::dtype::bfloat16, + phi::dtype::complex, + phi::dtype::complex, int, int64_t) {} diff --git a/paddle/phi/kernels/gpu/index_select_kernel.cu b/paddle/phi/kernels/gpu/index_select_kernel.cu index c2f447dcb1bca..910015e00f010 100644 --- a/paddle/phi/kernels/gpu/index_select_kernel.cu +++ b/paddle/phi/kernels/gpu/index_select_kernel.cu @@ -85,5 +85,7 @@ PD_REGISTER_KERNEL(index_select, double, phi::dtype::float16, phi::dtype::bfloat16, + phi::dtype::complex, + phi::dtype::complex, int, int64_t) {} diff --git a/python/paddle/tensor/search.py b/python/paddle/tensor/search.py index 3906fb5aad1f9..b7fe99f477ee8 100755 --- a/python/paddle/tensor/search.py +++ b/python/paddle/tensor/search.py @@ -319,7 +319,7 @@ def index_select(x, index, axis=0, name=None): size as the length of ``index``; other dimensions have the same size as in the ``x`` tensor. Args: - x (Tensor): The input Tensor to be operated. The data of ``x`` can be one of float16, float32, float64, int32, int64. + x (Tensor): The input Tensor to be operated. The data of ``x`` can be one of float16, float32, float64, int32, int64, complex64 and complex128. index (Tensor): The 1-D Tensor containing the indices to index. The data type of ``index`` must be int32 or int64. axis (int, optional): The dimension in which we index. Default: if None, the ``axis`` is 0. name (str, optional): For details, please refer to :ref:`api_guide_Name`. Generally, no setting is required. Default: None. @@ -353,7 +353,16 @@ def index_select(x, index, axis=0, name=None): check_variable_and_dtype( x, 'x', - ['uint16', 'float16', 'float32', 'float64', 'int32', 'int64'], + [ + 'uint16', + 'float16', + 'float32', + 'float64', + 'int32', + 'int64', + 'complex64', + 'complex128', + ], 'paddle.tensor.search.index_select', ) check_variable_and_dtype( @@ -771,7 +780,7 @@ def index_sample(x, index): Args: x (Tensor): The source input tensor with 2-D shape. Supported data type is - int32, int64, bfloat16, float16, float32, float64. + int32, int64, bfloat16, float16, float32, float64, complex64, complex128. index (Tensor): The index input tensor with 2-D shape, first dimension should be same with X. Data type is int32 or int64. @@ -826,7 +835,16 @@ def index_sample(x, index): check_variable_and_dtype( x, 'x', - ['uint16', 'float16', 'float32', 'float64', 'int32', 'int64'], + [ + 'uint16', + 'float16', + 'float32', + 'float64', + 'int32', + 'int64', + 'complex64', + 'complex128', + ], 'paddle.tensor.search.index_sample', ) check_variable_and_dtype( diff --git a/test/legacy_test/test_index_sample_op.py b/test/legacy_test/test_index_sample_op.py index d9551e0bea8d8..b2fee3e3fd648 100755 --- a/test/legacy_test/test_index_sample_op.py +++ b/test/legacy_test/test_index_sample_op.py @@ -28,6 +28,11 @@ def setUp(self): self.python_api = paddle.index_sample self.config() xnp = np.random.random(self.x_shape).astype(self.x_type) + if self.x_type == np.complex64 or self.x_type == np.complex128: + xnp = ( + np.random.random(self.x_shape) + + 1j * np.random.random(self.x_shape) + ).astype(self.x_type) indexnp = np.random.randint( low=0, high=self.x_shape[1], size=self.index_shape ).astype(self.index_type) @@ -122,6 +127,28 @@ def config(self): self.index_type = "int64" +class TestIndexSampleComplex64(TestIndexSampleOp): + def config(self): + """ + For complex64 x type + """ + self.x_shape = (10, 128) + self.x_type = np.complex64 + self.index_shape = (10, 64) + self.index_type = "int64" + + +class TestIndexSampleComplex128(TestIndexSampleOp): + def config(self): + """ + For complex64 x type + """ + self.x_shape = (10, 128) + self.x_type = np.complex128 + self.index_shape = (10, 64) + self.index_type = "int64" + + @unittest.skipIf( not core.is_compiled_with_cuda() or not core.is_bfloat16_supported(core.CUDAPlace(0)), diff --git a/test/legacy_test/test_index_select_op.py b/test/legacy_test/test_index_select_op.py index ceb152a465fc6..7bd578dcf68aa 100644 --- a/test/legacy_test/test_index_select_op.py +++ b/test/legacy_test/test_index_select_op.py @@ -36,6 +36,11 @@ def setUp(self): low=0, high=self.x_shape[self.dim], size=self.index_size ) x_np = np.random.random(self.x_shape).astype(self.x_type) + if self.dtype == np.complex64 or self.dtype == np.complex128: + x_np = ( + np.random.random(self.x_shape) + + 1j * np.random.random(self.x_shape) + ).astype(self.x_type) self.inputs = {'X': x_np, 'Index': index_np} self.attrs = {'dim': self.dim} outer_loop = np.prod(self.x_shape[: self.dim]) @@ -60,10 +65,16 @@ def init_dtype_type(self): self.index_size = 100 def test_check_output(self): - self.check_output(check_prim=True) + if self.x_type == np.complex64 or self.x_type == np.complex128: + self.check_output(check_prim=False) + else: + self.check_output(check_prim=True) def test_check_grad_normal(self): - self.check_grad(['X'], 'Out', check_prim=True) + if self.x_type == np.complex64 or self.x_type == np.complex128: + self.check_grad(['X'], 'Out', check_prim=False) + else: + self.check_grad(['X'], 'Out', check_prim=True) class TestIndexSelectOpCase2(TestIndexSelectOp): @@ -146,6 +157,24 @@ def test_check_grad_normal(self): self.check_grad_with_place(place, ['X'], 'Out', check_prim=True) +class TestIndexSelectComplex64(TestIndexSelectOp): + def init_dtype_type(self): + self.x_type = np.complex64 + self.index_type = np.int32 + self.dim = -2 + self.x_shape = (10, 10, 4, 10) + self.index_size = 10 + + +class TestIndexSelectComplex128(TestIndexSelectOp): + def init_dtype_type(self): + self.x_type = np.complex128 + self.index_type = np.int32 + self.dim = -2 + self.x_shape = (10, 10, 4, 10) + self.index_size = 10 + + class TestIndexSelectAPI(unittest.TestCase): def input_data(self): self.data_x = np.array(