Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

[Hackathon 2 No.22 ]Add RFC for task 22 (add API paddle.index_add) #127

Closed
wants to merge 1 commit into from
Closed
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
264 changes: 264 additions & 0 deletions rfcs/APIs/20220509_api_design_for_index_add.md
Original file line number Diff line number Diff line change
@@ -0,0 +1,264 @@
# paddle.Tensor.index_add 设计文档



|API名称 | paddle.Tensor.index_add |

|---|------------------------------------------|

|提交作者<input type="checkbox" class="rowselector hidden"> | SmirnovKol |

|提交时间<input type="checkbox" class="rowselector hidden"> | 2022-05-09 |

|版本号 | V1.0 |

|依赖飞桨版本<input type="checkbox" class="rowselector hidden"> | develop |

|文件名 | 20220509_api_design_for_index_add.md<br> |



# 一、概述



## 1、相关背景

为了提升飞桨API丰富度,支持科学计算领域API,Paddle需要实现API`paddle.index_add功能需求。

## 2、功能目标

增加 API `paddle.index_add`, `paddle.index_add_`,`tensor.index_add`,`tensor.index_add_`, 在指定轴上, 通过按index中给定的顺序切片, 对每个切片的张量加上固定的常量。

## 3、意义

飞桨支持index_add算子进一步满足用户需求。



# 二、飞桨现状

目前paddle缺少相关功能实现, 且用其他API来组合实现也较为困难, 因为axis可以是任意轴而且index也不一定连续,`paddle.index_select`和`paddle.slice`也无法直接达到目的。

简单场景下计算逻辑如下:





```Python

import paddle
import numpy as np
np.random.seed(102)


x_np = np.random.rand(4, 3)
x_tensor = paddle.to_tensor(x_np)

result_np = x_np.copy()

added_value = 9.0

index = [0, 2]
axis = 0
result_np[index] += added_value

length = x_tensor.shape[axis]
for i in index:
if i < 0 or i >= length:
raise ValueError('index is wrong: {}'.format(index))
else:
result_tensor[i] += added_value

print(np.allclose(result_np, result_tensor.numpy()))

```



# 三、业内方案调研

## Numpy

### 实现方法

Numpy没有对该功能有特定的API进行支持,但是Numpy有非常完善的切片操作和广播机制,可以很好的实现。示例如下:



```Python

import numpy as np

axis = 2

index = [0, 2, 3]

added_value = 97

data = np.random.rand(4, 3, 7, 9)

data[:, :, index, :] += added_value


```



## TensorFlow

tensorflow目前也没有特定的API支持类似功能,

但在tensorflow里也可以通过tf.exprimental.numpy直接调用numpy函数



## Pytorch

Pytorch中有API`Tensor.index_add_(dim, index, source, *, alpha=1)`和`Tensor.index_add(dim, index, source, *, alpha=1)`,


[文档地址](https://pytorch.org/docs/stable/generated/torch.Tensor.index_add_.html), 介绍为:



# 四、对比分析

- Numpy基于切片操作和广播机制功能上更灵活更自由。

- Pytorch支只支持一个axis,不仅支持cpu还支持gpu。



# 五、方案设计

## 命名与参数设计

新增API设计为:

`paddle.index_add(x, axis, index, added_value)`

`paddle.index_add_(x, axis, index, added_value)`

`Tensor.index_add(axis, index, added_value)`

`Tensor.index_add_(axis, index, added_value)`



index_add_支持inplace方式修改输入张量。

axis是index索引选择的轴, 支持int类型。

index在指定轴上含索引下标的list of int, tuple of int 或者 1-D Tensor。

added_value是待相加的数据,参数类型支持bool, int, float。
Copy link

@iclementine iclementine May 11, 2022

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

我们希望 axis 能支持输入一个整数,或者一个 0D 的整形 Tensor,这样,如果 axis 是其他运算的结果,也能被支持。

add_value 我们希望支持 bool, int, float, complex 之类的 python scalar, 也能支持和 x 数据类型相同的 Tensor (只要它的形状可以 broadcast 到指定出来的形状)。而且当 add_value 是一个 Tensor 的时候,它的梯度也需要能正常回传。达到类似 numpy 切片 += 的效果。

import numpy as np
x = np.random.randn(2, 3)
x[:, [0: 2]] += 10 # ok, python scalalr
x[:, [0, 2]] += np.array(10) # ok, 0d array
x[:, [0, 2]] += np.array([1000, 1000]) # ok (2,) broadcast to (2, 2)

预期能实现这些

import paddle
x = paddle.randn([2, 3])
paddle.index_add(x, 1, [0, 2], 10) # ok add_value accepts python scalar
paddle.index_add(x, 1, [0, 2], paddle.to_tensor(10, dtype=paddle.float32)) # ok add_value accepts 0D Tensor
paddle.index_add(x, 1, paddle.to_tensor([0, 2]), paddle.to_tensor([100.0, 100.0])) # ok, index & add_value accept Tensors, add_value can be broadcast to the slice's shape

(反向的部分未写出示例。)

你可能需要了解一下,如何让一个参数既支持 python scalar, 也能支持 Tensor 以实现灵活的语义。Hint: 在 paddle 的 op 设计中存在 Input 和 Attribute 的区别。你可以参考

paddle/fluid/operators/scale_op.ccpaddle/phi/kernels/cpu/scale_kernel.ccpaddle/phi/ops/compat/scale_sig.cc 来实现。另外考虑到需要 add_value 也需要梯度回传,你可能需要实现两个 kernel, 然后在 python 接口根据 add_value 是不是一个 Tensor 来调用不同的 kernel.

如果有遇到问题,欢迎咨询。




## 底层OP设计

参考飞桨现有算子,分别实现cpu和cuda的算子kernel。



## API实现方案

在 python/paddle/tensor/manipulation.py 中增加index_add以及index_add_函数,分别通过_C_ops调用底层算子

计算正确的stride之后,参考index_select算子进行逻辑修改

输入tensor的所有元素梯度是1.0


## 代码实现文件路径



CPU中正向和反向计算:

paddle/phi/kernels/cpu/index_add_kernel.cc
paddle/phi/kernels/cpu/index_add_grad_kernel.cc

GPU中正向和反向计算:

paddle/phi/kernels/gpu/index_add_kernel.cu
paddle/phi/kernels/gpu/index_add_grad_kernel.cu



```c++

template <typename T, typename Context>
void IndexAddKernel(const Context& ctx,
const DenseTensor& x,
const DenseTensor& index,
int axis,
float added_value,
DenseTensor* output);


template <typename T, typename Context>
void IndexAddGradKernel(const Context& ctx,
const DenseTensor& out_grad,
int axis,
float added_value,
DenseTensor* x_grad);
```

算子注册路径:

paddle/fluid/operators/index_add_op.cc

函数API实现路径: python/paddle/tensor/manipulation.py

单元测试路径: python/paddle/fluid/tests/unittests/test_index_add_op.py



# 六、测试和验收的考量

测试考虑的case如下:



- 和numpy结果的数值的一致性, `paddle.index_add`和numpy切片操作结果是否一致;

- 参数`axis`校验参数类型int,判断axis合法,并进行边界检查;

- 校验参数`index`的正确性,索引边界检查,输出结果的正确性;

- 校验参数added_value的正确性, 是否是支持的数据类型

- 测试在进行反向梯度计算时结果的正确性;

- 错误检查:输入`x`不是Tensor时,能否正确抛出错误;



# 七、可行性分析及规划排期



方案实施难度可控,工期上可以满足在当前版本周期内开发完成。



# 八、影响面

为独立新增API,对其他模块没有影响



# 名词解释


# 附件及参考资料