Skip to content

Latest commit

 

History

History
249 lines (184 loc) · 8.37 KB

20220329_api_design_for_sparse_divide.md

File metadata and controls

249 lines (184 loc) · 8.37 KB

paddle.sparse.divide 设计文档

API名称 paddle.sparse.divide
提交作者 PeachML
提交时间 2022-03-29
版本号 V1.0
依赖飞桨版本 develop
文件名 20220329_api_design_for_sparse_divide.md

一、概述

1、相关背景

为了提升飞桨API丰富度,divide 是一个基础除法运算操作,目前 Paddle 中还没有 sparse 的除法算子。 本任务的目标是在 Paddle 中添加 sparse.divide 算子, 实现输入是两个 SparseCooTensor 或者两个 SparseCsrTensor 逐元素相除的功能。 Paddle需要扩充API,新增 sparse.divide API, 调用路径为:paddle.sparse.divide 实现稀疏Tensor相除的功能。

3、意义

支持稀疏tensor相除,提高空间利用效率,提升稀疏tensor的计算效率。

二、飞桨现状

目前paddle缺少相关功能实现。

三、业内方案调研

Pytorch

Pytorch中有APItorch.div(input, other, *, rounding_mode=None, out=None) , 在pytorch中,介绍为:

Divides each element of the input input by the corresponding element of other.

可以支持一般tensor和sparse tensor的相除

Scipy

Scipy中有csr类型的稀疏矩阵,可以支持相除操作,通过binary_op实现。对coo类型的稀疏矩阵,先转换成csr再相除。

实现方法

代码如下

/*
 * Compute C = A (binary_op) B for CSR matrices that are not
 * necessarily canonical CSR format.  Specifically, this method
 * works even when the input matrices have duplicate and/or
 * unsorted column indices within a given row.
 *
 * Refer to csr_binop_csr() for additional information
 *
 * Note:
 *   Output arrays Cp, Cj, and Cx must be preallocated
 *   If nnz(C) is not known a priori, a conservative bound is:
 *          nnz(C) <= nnz(A) + nnz(B)
 *
 * Note:
 *   Input:  A and B column indices are not assumed to be in sorted order
 *   Output: C column indices are not generally in sorted order
 *           C will not contain any duplicate entries or explicit zeros.
 *
 */
template <class I, class T, class T2, class binary_op>
void csr_binop_csr_general(const I n_row, const I n_col,
                           const I Ap[], const I Aj[], const T Ax[],
                           const I Bp[], const I Bj[], const T Bx[],
                                 I Cp[],       I Cj[],       T2 Cx[],
                           const binary_op& op)
{
    //Method that works for duplicate and/or unsorted indices

    std::vector<I>  next(n_col,-1);
    std::vector<T> A_row(n_col, 0);
    std::vector<T> B_row(n_col, 0);

    I nnz = 0;
    Cp[0] = 0;

    for(I i = 0; i < n_row; i++){
        I head   = -2;
        I length =  0;

        //add a row of A to A_row
        I i_start = Ap[i];
        I i_end   = Ap[i+1];
        for(I jj = i_start; jj < i_end; jj++){
            I j = Aj[jj];

            A_row[j] += Ax[jj];

            if(next[j] == -1){
                next[j] = head;
                head = j;
                length++;
            }
        }

        //add a row of B to B_row
        i_start = Bp[i];
        i_end   = Bp[i+1];
        for(I jj = i_start; jj < i_end; jj++){
            I j = Bj[jj];

            B_row[j] += Bx[jj];

            if(next[j] == -1){
                next[j] = head;
                head = j;
                length++;
            }
        }


        // scan through columns where A or B has
        // contributed a non-zero entry
        for(I jj = 0; jj < length; jj++){
            T result = op(A_row[head], B_row[head]);

            if(result != 0){
                Cj[nnz] = head;
                Cx[nnz] = result;
                nnz++;
            }

            I temp = head;
            head = next[head];

            next[temp]  = -1;
            A_row[temp] =  0;
            B_row[temp] =  0;
        }

        Cp[i + 1] = nnz;
    }
}

四、对比分析

torch设计结构复杂,为了适配paddle phi库的设计模式,故参考scipy的实现方式

五、方案设计

命名与参数设计

在paddle/phi/kernels/sparse/目录下, kernel设计为

void ElementWiseDivideCsrCPUKernel(const Context& dev_ctx,
                                   const SparseCsrTensor& x,
                                   const SparseCsrTensor& y,
                                   SparseCsrTensor* out) 
void ElementWiseDivideCooKernel(const Context& dev_ctx,
                                const SparseCooTensor& x,
                                const SparseCooTensor& y,
                                SparseCooTensor* out) 
template <typename T, typename Context>
void ElementWiseDivideCsrGradKernel(const Context& dev_ctx,
                                    const SparseCsrTensor& x,
                                    const SparseCsrTensor& y,
                                    const SparseCsrTensor& out,
                                    const SparseCsrTensor& dout,
                                    SparseCsrTensor* dx,
                                    SparseCsrTensor* dy);
void ElementWiseDivideCooGradKernel(const Context& dev_ctx,
                                    const SparseCooTensor& x,
                                    const SparseCooTensor& y,
                                    const SparseCooTensor& out,
                                    const SparseCooTensor& dout,
                                    SparseCooTensor* dx,
                                    SparseCooTensor* dy);

函数设计为

SparseCsrTensor ElementWiseDivideCsr(const Context& dev_ctx,
                                     const SparseCsrTensor& x,
                                     const SparseCsrTensor& y)

SparseCooTensor ElementWiseDivideCoo(const Context& dev_ctx,
                                     const SparseCooTensor& x,
                                     const SparseCooTensor& y)

底层OP设计

实现对应的 CPU Kernel,使用 Merge 两个有序数组的算法,然后使用已有op组合实现, 主要涉及SparseCooToCsrKernelSparseCsrToCooKernel

对于dense tensor,值连续的存储在一块内存中,二元运算需要处理每一个元素,即x[i][j] ∘ y[i][j],运算时间复杂度为 O(numel(x))numel(x)x中总素个数。

而sparse tensor以索引和值的模式存储一个多数元素为零的tensor,二元运算只需要处理两个输入不全为0的位置, 在sparse tensor构造时,索引按升序排序,可以采取merge有序数组的方式,若两输入索引相等,则计算x[i][j] ∘ y[i][j], 若不相等则说明该位置上的二元运算有一个元为0, x索引小时计算 x[i][j] ∘ 0y索引小时计算 0 ∘ y[i][j]。 对于除法而言,两个输入都为0的位置运算结果不为0,所以需要将y中索引未覆盖到的位置(原tensor中为0),在运算中通过指向一个零常量的方式填充为0, 以参与除法运算。 此时,时间复杂度为O(numel(x)),鉴于实际场景中动辄50% ~ 99%的稀疏性 [1]1 , 虽然时间复杂度与dense tensor一样,但是可以节省大量的存储空间。

API实现方案

对于SparseCsrTensor,将csr格式转换成coo格式再进行运算,然后转换回csr格式输出。

对于SparseCooTensor,直接进行运算。

六、测试和验收的考量

测试考虑的case如下:

  • 数值正确性
  • 反向
  • 不同 sparse_dim

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

方案主要依赖paddle现有op组合而成,并自行实现核心算法

八、影响面

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

名词解释

附件及参考资料

[1: 深度学习中的稀疏性]1