自定義算子高性能開發
在計圖中，一共有三種方法來開發自定義的算子:

1. 使用元算子進行組合。
2. 使用Code算子開發自定義算子。
3. 使用計圖編譯器編譯自定義的模塊和custom op。
   其中，元算子開發是最為簡單的， 但不免有些情況存在元算子表達能力不足。可以使用Code算子進行開發，Code算子在保持了開發的便捷性，還具有很高的可定制性和性能。和方法3相比，Code算子的開發更加簡單，非常適合用戶構建模型中的創新算子。
   本文主要介紹Code算子，關于元算子和自定義模塊，參考文檔：
   ? 使用元算子開發卷積
   ? 使用計圖編譯器編譯自定義的模塊和custom op
   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_varsreturn 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
@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
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 “,
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] = {dxdx+dydy+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_varsreturn jt.code([a.shape, b.shape], [a.dtype, b.dtype], [a, b, grad],cuda_src='''__global__ static void kernel2(@ARGS_DEF) {@PRECALCfor (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?
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_varsreturn 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) {{@PRECALCint 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+padding2-kernel_size)//stride+1
w = (W+padding2-kernel_size)//stride+1

func = Func()
out = func(x)
print(out)
print(jt.grad(out, x))

總結

以上是生活随笔為你收集整理的自定义算子高性能开发的全部內容，希望文章能夠幫你解決所遇到的問題。

如果覺得生活随笔網站內容還不錯，歡迎將生活随笔推薦給好友。