标签：自定义 int in0 k3 k2 高性能 算子 out
自定义算子高性能开发
在计图中，一共有三种方法来开发自定义的算子:
- 使用元算子进行组合。
- 使用Code算子开发自定义算子。
- 使用计图编译器编译自定义的模块和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))
标签：自定义,int,in0,k3,k2,高性能,算子,out 来源： https://www.cnblogs.com/wujianming-110117/p/14394891.html
本站声明： 1. iCode9 技术分享网（下文简称本站）提供的所有内容，仅供技术学习、探讨和分享； 2. 关于本站的所有留言、评论、转载及引用，纯属内容发起人的个人观点，与本站观点和立场无关； 3. 关于本站的所有言论和文字，纯属内容发起人的个人观点，与本站观点和立场无关； 4. 本站文章均是网友提供，不完全保证技术分享内容的完整性、准确性、时效性、风险性和版权归属；如您发现该文章侵犯了您的权益，可联系我们第一时间进行删除； 5. 本站为非盈利性的个人网站，所有内容不会用来进行牟利，也不会利用任何形式的广告来间接获益，纯粹是为了广大技术爱好者提供技术内容和技术思想的分享性交流网站。