自定义算子高性能开发

自定义算子高性能开发

在计图中,一共有三种方法来开发自定义的算子:

  1. 使用元算子进行组合。
  2. 使用Code算子开发自定义算子。
  3. 使用计图编译器编译自定义的模块和custom op。

其中,元算子开发是最为简单的, 但不免有些情况存在元算子表达能力不足。可以使用Code算子进行开发,Code算子在保持了开发的便捷性,还具有很高的可定制性和性能。和方法3相比,Code算子的开发更加简单,非常适合用户构建模型中的创新算子。

本文主要介绍Code算子,关于元算子和自定义模块,参考文档:

Code算子是一个基于高性能语言的动态编译算子,允许用户直接在Python中内联C++/CUDA代码,只需要寥寥数行代码,就可以完成高性能的自定义算子开发,降低用户开发自定义算子的难度。

Code 算子的输入参数

使用Python的help命令(help(jt.code)),可以看到文档如下:

@param[in]       shape   输出的形状, a integer array

@param[in]       dtype   输出的数据类型

@param[in]       inputs  一个计图变量数组

@param[in]       cpu_src CPU前向代码字符串,内建变量包括:

    *   in{x}, in{x}_shape{y}, in{x}_stride{y}, in{x}_type, in{x}_p, @in0(...)

    *   out{x}, out{x}_shape{y}, out{x}_stride{y}, out{x}_type, out{x}_p, @out0(...)

    *   out, out_shape{y}, out_stride{y}, out_type, out_p, @out(...)

@param[in]       cpu_header   CPU头文件字符串

@param[in]       cuda_src CUDA    前向代码字符串,和上述参数具有同样的内建变量。

@param[in]       cuda_header CUDA头文件字符串。

可以看到,用户需要提供Code算子的输入,输出的形状和类型,以及对应的代码。计图会通过编译缓存器,让相同的代码只编译一次。如果希望最大化Code算子的性能,尽量保证Code算子的代码不会出现过多变种。在Code算子的代码中,用户可以使用内建变量,访问计图的变量。下面将用若干个实例,来介绍Code算子的使用。

实例1:CPU算子以及导数

下面的实例中,首先生成了一个随机的长度为10的变量a,然后计算了$2a^2$ 和对应的导数$4a$,在这个例子中使用了@out, @in0,这种C++中没有的语法,这种语法目的是给用户提供方便的访问计图变量的接口。这种语法在后端会被翻译成C++可以识别的语法。

from jittor import Function

import jittor as jt

 

class Func(Function):

    def execute(self, x):

        self.save_vars = x

        return jt.code(x.shape, x.dtype, [x],

            cpu_src='''

                for (int i=0; i<in0_shape0; i++)

                    @out(i) = @in0(i)*@in0(i)*2;

            ''')

 

    def grad(self, grad_x):

        x = self.save_vars

        return jt.code(x.shape, x.dtype, [x, grad_x],

            cpu_src='''

                for (int i=0; i<in0_shape0; i++)

                    @out(i) = @in1(i)*@in0(i)*4;

            ''')

 

a = jt.random([10])

func = Func()

b = func(a)

print(b)

print(jt.grad(b,a))

实例2:使用stl和alias

下面的实例中,实现了一个简单的排序算法,演示了如何使用C++算法库中排序算法,以及使用别名alias来增加代码的可读性。

a = jt.array([3,2,1])

b = jt.code(a.shape, a.dtype, [a],

    cpu_header="""

        #include <algorithm>

        @alias(a, in0)

        @alias(b, out)

    """,

    cpu_src="""

        for (int i=0; i<a_shape0; i++)

            @b(i) = @a(i);

        std::sort(&@b(0), &@b(in0_shape0));

    """

)

assert (b.data==[1,2,3]).all()

实例3:多输出的Code算子

在某些情况下,算子可能有多个输出,在这个实例中,演示了如何设置多输出。该算子输入为一维向量,输出为两个长度为1的向量,分别是最小值和最大值。

同之前实例不同的地方是,原来传入单个shape和dtype,这里传入的是一个shape数组和dtype数组。同时还在这个实例中演示了如何使用cout。

a = jt.array([3,2,1])

b,c = jt.code([(1,), (1,)], [a.dtype, a.dtype], [a],

    cpu_header="""

        #include <iostream>

        using namespace std;

    """,

    cpu_src="""

        @alias(a, in0)

        @alias(b, out0)

        @alias(c, out1)

        @b(0) = @c(0) = @a(0);

        for (int i=0; i<a_shape0; i++) {

            @b(0) = std::min(@b(0), @a(i));

            @c(0) = std::max(@c(0), @a(i));

        }

        cout << "min:" << @b(0) << " max:" << @c(0) << endl;

    """

)

assert b.data == 1, b

assert c.data == 3, c

实例4:动态大小的输出

在某些情况下,算子的输出的大小可能是会变化的,比如把输入中大于0和小于等于0的数,分别紧密排列在两个向量中。下面的实例就实现了这样一个算子。

可以发现下面的数组的输出形状被设置成了负数,这是计图的特殊机制,传入负数代表这个数组的大小是不确定的,负数的绝对值,代表了这个维度最大上限。需要注意的是,动态大小只能在第一维度出现,而且在算法最后结束的时候,需要使用set_shape来设置确定的形状。

a = jt.array([5,-4,3,-2,1])

 

# negtive shape for max size of vary dimension

b,c = jt.code([(-5,), (-5,)], [a.dtype, a.dtype], [a],

    cpu_src="""

        @alias(a, in0)

        @alias(b, out0)

        @alias(c, out1)

        int num_b=0, num_c=0;

        for (int i=0; i<a_shape0; i++) {

            if (@a(i)>0)

                @b(num_b++) = @a(i);

            else

                @c(num_c++) = @a(i);

        }

        b->set_shape({num_b});

        c->set_shape({num_c});

    """

)

assert (b.data == [5,3,1]).all()

assert (c.data == [-4,-2]).all()

综合实例5:使用Code算子实现三维点云K近邻查找

下面的实例展示了如何使用code算子,使用数行代码实现三维点云中十分常用的K近邻查找。Code算子的设计和实现,让用户既可以享受到Python语言的便捷与易用性,又可以获得高性能语言的性能。

可以留意到,在计图的Code算子中,可以使用openmp实现自动并行化的,关于openmp的使用,可以参考openmp文档

a = jt.random((n,3))

b = jt.code([n, k], "int32", [a],

cpu_header="#include <algorithm>",

cpu_src="""

  using namespace std;

  auto n=out_shape0, k=out_shape1;

  // 使用openmp实现自动并行化

  #pragma omp parallel for

  for (int i=0; i<n; i++) {

    // 存储k近邻的距离和下标

    vector<pair<float,int>> id(n);

    for (int j=0; j<n; j++) {

      auto dx = @in0(i,0)-@in0(j,0);

      auto dy = @in0(i,1)-@in0(j,1);

      auto dz = @in0(i,2)-@in0(j,2);

      id[j] = {dx*dx+dy*dy+dz*dz, j};

    }

    // 使用c++算法库的nth_element排序

    nth_element(id.begin(),

      id.begin()+k, id.end());

    // 将下标输出到计图的变量中

    for (int j=0; j<k; j++)

      @out(i,j) = id[j].second;

  }"""

)

将计图使用code算子实现的K近邻查找,和PyTorch的算子用时进行比较,速度对比如下(k=10,点云数量n=[100,1000,10000]):

参数

n=100

n=1000

n=10000

PyTorch

433 µs

7.6 ms

623 ms

Jittor

68 µs

5.9 ms

484 ms

速度对比

6.4X

1.29X

1.29X

注:此处使用的K近邻算法为暴力算法,还存在更优的算法实现,由于文章篇幅有限,此处仅用于展示Code算子的使用。

实例6:使用CUDA进行加速

在这个实例中,使用CUDA实现了简单的两个2维向量相乘。并且反向传播对应的导数。

这个实例与之前的区别,定义了CUDA kernel,这需要用户有一定的CUDA基础。这里面的@ARGS_DEF,@ARGS分别是CUDA kernel函数的参数声明和参数传递,而@PRECALC包含了计图预处理内核的代码。除此之外,其他语法和CUDA保持高度一致。

import jittor as jt

from jittor import Function

jt.flags.use_cuda = 1

 

class Func(Function):

    def execute(self, a, b):

        self.save_vars = a, b

        return jt.code(a.shape, a.dtype, [a,b],

            cuda_src='''

                __global__ static void kernel1(@ARGS_DEF) {

                    @PRECALC

                    for (int i=blockIdx.x; i<in0_shape0; i+=gridDim.x)

                    for (int j=threadIdx.x; j<in0_shape1; j+=blockDim.x)

                        @out(i,j) = @in0(i,j)*@in1(i,j);

                }

                kernel1<<<32, 32>>>(@ARGS);

            ''')

 

    def grad(self, grad):

        a, b = self.save_vars

        return jt.code([a.shape, b.shape], [a.dtype, b.dtype], [a, b, grad],

            cuda_src='''

                __global__ static void kernel2(@ARGS_DEF) {

                    @PRECALC

                    for (int i=blockIdx.x; i<in0_shape0; i+=gridDim.x)

                    for (int j=threadIdx.x; j<in0_shape1; j+=blockDim.x) {

                        @out0(i,j) = @in2(i,j)*@in1(i,j);

                        @out1(i,j) = @in2(i,j)*@in0(i,j);

                    }

                }

                kernel2<<<32, 32>>>(@ARGS);

            ''')

       

a = jt.random((100,100))

b = jt.random((100,100))

func = Func()

c = func(a,b)

print(c)

print(jt.grad(c, [a, b]))

综合实例7:实现可以同时在GPU和CPU上运行的Pool算法

注:计图内部已经实现了Pool,用户不需要自己实现

import jittor as jt

from jittor import Function

jt.flags.use_cuda = 1

 

class Func(Function):

    def execute(self, x):

        out = jt.code([N,C,h,w], x.dtype, [x],

            cuda_src=f'''

                __global__ static void kernel1(@ARGS_DEF) {{

                    @PRECALC

                    int p3 = threadIdx.x;

                    int s3 = blockDim.x;

                    int p2 = threadIdx.y + blockIdx.x * blockDim.y;

                    int s2 = blockDim.y * gridDim.x;

                    int i1 = blockIdx.y;

                    int i0 = blockIdx.z;

                    for (int i3 = p3; i3 < out_shape3; i3 += s3)

                        for (int i2 = p2; i2 < out_shape2; i2 += s2) {{

                            int k3 = i3*{stride}-{padding};

                            int k2 = i2*{stride}-{padding};

                            int k3_ = min(k3 + {kernel_size}, in0_shape3);

                            int k2_ = min(k2 + {kernel_size}, in0_shape2);

                            k3 = max(0, k3);

                            k2 = max(0, k2);

                            @out(i0, i1, i2, i3) = @in0(i0, i1, k2, k3);

                            for (int p = k2; p < k2_; ++p)

                                for (int q = k3; q < k3_; ++q)

                                    @out(i0, i1, i2, i3) = {op}(@out(i0, i1, i2, i3), @in0(i0, i1, p, q));

                        }}

                }}

                int tx = min(1024, out_shape3);

                int ty = min(1024 / tx, out_shape2);

                int bx = (out_shape2 - 1) / ty + 1;

                int by = out_shape1;

                int bz = out_shape0;

                dim3 s1(bx, by, bz);

                dim3 s2(tx, ty);

                kernel1<<<s1, s2>>>(@ARGS);

            ''',

            cpu_src=f'''

                for (int i0=0; i0<out_shape0; i0++)

                for (int i1=0; i1<out_shape1; i1++)

                for (int i2=0; i2<out_shape2; i2++)

                for (int i3=0; i3<out_shape3; i3++) {{

                    int k2 = i2*{stride}-{padding};

                    int k3 = i3*{stride}-{padding};

                    int k2_ = std::min(k2 + {kernel_size}, in0_shape2);

                    int k3_ = std::min(k3 + {kernel_size}, in0_shape3);

                    k2 = std::max(0, k2);

                    k3 = std::max(0, k3);

                    @out(i0, i1, i2, i3) = @in0(i0, i1, k2, k3);

                    for (int p = k2; p < k2_; ++p)

                        for (int q = k3; q < k3_; ++q)

                            @out(i0, i1, i2, i3) = std::{op}(@out(i0, i1, i2, i3), @in0(i0, i1, p, q));

                }}

            ''')

        self.save_vars = x, out

        return out

 

    def grad(self, grad_x):

        x, pout = self.save_vars

        return jt.code(x.shape, x.dtype, [x, pout, grad_x],

            cuda_header=f'''

            @alias(pout, in1);

            ''',

            cuda_src=f'''

            __global__ static void kernel3(@ARGS_DEF) {{

                @PRECALC

                int p3 = threadIdx.x;

                int s3 = blockDim.x;

                int p2 = threadIdx.y + blockIdx.x * blockDim.y;

                int s2 = blockDim.y * gridDim.x;

                int i1 = blockIdx.y;

                int i0 = blockIdx.z;

                for (int i3 = p3; i3 < pout_shape3; i3 += s3)

                    for (int i2 = p2; i2 < pout_shape2; i2 += s2) {{

                        int k3 = i3*{stride}-{padding};

                        int k2 = i2*{stride}-{padding};

                        int k3_ = min(k3 + {kernel_size}, in0_shape3);

                        int k2_ = min(k2 + {kernel_size}, in0_shape2);

                        k3 = max(0, k3);

                        k2 = max(0, k2);

                        int bo=1;

                        for (int p = k2; p < k2_ && bo; ++p)

                            for (int q = k3; q < k3_ && bo; ++q) {{

                                if (@pout(i0,i1,i2,i3) == @in0(i0,i1,p,q)) {{

                                    atomicAdd(&@out(i0,i1,p,q), @in2(i0,i1,i2,i3));

                                    bo=0;

                                }}

                            }}

                    }}

            }}

            cudaMemsetAsync(out_p, 0, out->size);

            int tx = min(1024, pout_shape3);

            int ty = min(1024 / tx, pout_shape2);

            int bx = (pout_shape2 - 1) / ty + 1;

            int by = pout_shape1;

            int bz = pout_shape0;

            dim3 s1_(bx, by, bz);

            dim3 s2_(tx, ty);

            kernel3<<<s1_, s2_>>>(@ARGS);

            ''',

            cpu_src=f'''

                @alias(pout, in1);

                for (int i=0; i<out_shape0; i++)

                for (int j=0; j<out_shape1; j++)

                for (int k=0; k<out_shape2; k++)

                for (int l=0; l<out_shape3; l++) @out(i,j,k,l) = 0;

 

                for (int i0=0; i0<pout_shape0; i0++)

                for (int i1=0; i1<pout_shape1; i1++)

                for (int i2=0; i2<pout_shape2; i2++)

                for (int i3=0; i3<pout_shape3; i3++) {{

                    int k3 = i3*{stride}-{padding};

                    int k2 = i2*{stride}-{padding};

                    int k3_ = std::min(k3 + {kernel_size}, in0_shape3);

                    int k2_ = std::min(k2 + {kernel_size}, in0_shape2);

                    k3 = std::max(0, k3);

                    k2 = std::max(0, k2);

                    int bo=1;

                    for (int p = k2; p < k2_ && bo; ++p)

                        for (int q = k3; q < k3_ && bo; ++q) {{

                            if (@pout(i0,i1,i2,i3) == @in0(i0,i1,p,q)) {{

                                @out(i0,i1,p,q) += @in2(i0,i1,i2,i3);

                                bo=0;

                            }}

                        }}

                }}

            ''')

 

 

N,C,H,W = [2,10,100,100]

stride = 2

padding = 0

kernel_size = 3

op = "max"

 

x = jt.random((N,C,H,W))

h = (H+padding*2-kernel_size)//stride+1

w = (W+padding*2-kernel_size)//stride+1

 

func = Func()

out = func(x)

print(out)

print(jt.grad(out, x))

 

上一篇:K3忘记帐套管理员密码的解决办法


下一篇:【iOS面试粮食】Runtime—消息传递和转发机制、Method Swizzling