2. 程式人生 > 其它 >自定義運算元高效能開發

自定義運算元高效能開發

阿新 • • 發佈：2021-02-11

技術標籤：GPU視訊推理智慧推理

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

  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，然後計算了 2 a 2 2a^2 2a2 和對應的導數 4 a 4a 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
@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)[email protected](j,0);
auto dy = @in0(i,1)[email protected](j,1);
auto dz = @in0(i,2)[email protected](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_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©
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+padding2-kernel_size)//stride+1
w = (W+padding2-kernel_size)//stride+1

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

相關推薦

自定義運算元高效能開發

技術標籤：GPU視訊推理智慧推理 自定義運算元高效能開發 在計圖中，一共有三種方法來開發自定義的運算元:

自定義微信開發 spring-boot-starter

技術標籤：javaspringgit 自定義微信開發 spring-boot-starter 1.新建專案 ​ 推薦使用 xx-spring-boot-starter

ios開發 自定義btn_Flutter 混合開發

技術標籤：ios開發 自定義btn 目前大多數公司都有自己開發多年的專案，不可能直接用Flutter從頭開發一套，那樣不實現，除非是小專案，因此只能是在原有的基礎上用Flutter來開發新業務或重構舊業務，而這裡就需

C#--自定義控制元件-開發LED指示燈控制元件

以下是學習筆記： 參考：https://www.bilibili.com/video/BV1eQ4y1M7ZY?p=5 效果如下： 思考：實現以上效果要用到哪些屬性

如何在OpenVINO中實現自己的自定義運算元 – 第一篇：搞定模型轉換

技術標籤：程式設計技巧開源專案深度學習 前言 與主流的深度學習框架類似，OpenVINO也提供一套擴充套件機制，方便開發者可以實現自己的自定義運算元。自定義運算元有如下三種使用場景：

pytorch之C++實現自定義運算元

技術標籤：深度學習c++pythonjavalinux程式語言 自定義運算元 對於輸入 x，其輸出為

【pytorch->mindspore】1.自定義運算元遷移

要遷移的專案為影象壓縮演算法https://github.com/ywz978020607/HESIC 1.自定義運算元遷移--LowerBoundFunction類

Robot Framework自定義測試庫開發詳解

Robot Framework是一個開源自動化測試框架，主要特點是基於關鍵字驅動，本文介紹自定義測試庫及測試庫文件建立方法。

pytorch自定義運算元

參照官方教程，實現pytorch自定義運算元。主要分為以下幾步： 改寫運算元為torch C++版本

rxjs Observable 自定義 Operator 的開發技巧

操作符是 RxJS 庫和 Observables 的基礎塊。 它使我們能夠通過使用一些關鍵字（函式）來執行復雜的操作。 運算子只不過是獲取源 Observable、對其執行操作並返回新 Observable 的函式。

金蝶雲星空自定義WebApi介面開發和呼叫

1、金蝶雲星空自定義WebApi開發 外掛開發好之後，編譯生成DLL檔案，放到金蝶安裝目錄\\WebSite\\bin下，並重啟IIS

自定義MVC開發的Xcode模板

前言 Xcode分為系統模板和自定義模板。模板的好處如下： 1、節省重複程式碼手寫時間

python GUI庫圖形介面開發之PyQt5訊號與槽機制、自定義訊號基礎介紹

訊號和槽機制是 QT 的核心機制，要精通 QT 程式設計就必須對訊號和槽有所瞭解。訊號和槽是一種高階介面，應用於物件之間的通訊，它是 QT 的核心特性，也是 QT 區別於其它工具包的重要地方。

python GUI庫圖形介面開發之PyQt5訊號與槽的高階使用技巧(自定義訊號與槽)詳解與例項

PyQt5訊號與槽高階自定義訊號與槽 所謂高階自定義訊號與槽，指的就是我們可以以自己喜歡的方式定義訊號與槽函式，並傳遞引數，自定義訊號的一般流程如下

詳解使用python3.7配置開發釘釘群自定義機器人(2020年新版攻略)

最近疫情比較嚴重，很多公司依靠阿里旗下的辦公軟體釘釘來進行遠端辦公，當然了，釘釘這個產品真的是讓人一言難盡，要多難用有多難用，真的讓人覺得阿里的pm都是腦殘才會設計出這種腦殘產品，不過吐槽歸吐槽，該用還

Javaweb EL自定義函式開發及程式碼例項

1.什麼是EL自定義函式 EL自定義函式是在EL表示式中呼叫的某個java類的靜態方法，這個靜態方法需在web應用程式中進行配置才可以被EL表示式呼叫。

面向物件開發之自定義事件

class Event { constructor(){ this.handlers = {}; // 記錄所有的事件及處理函式 // { //click: [fn1, fn2],

Unity開發筆記-PlayerPrefs儲存bytes[]和自定義物件

PlayerPrefs沒有提供儲存bytes[]和其他自定義結構的介面，我們可以利用現有的SetString介面，將bytes[]轉化為string儲存，從而實現擴充套件

WPF介面開發技巧分享——如何實現自定義DateEdit並自動更正值

下載DevExpress v20.1完整版 通過DevExpress WPF Controls，您能建立有著強大互動功能的XAML基礎應用程式，這些應用程式專注於當代客戶的需求和構建未來新一代支援觸控的解決方案。

Winform介面開發：如何在程式碼中獲取自定義外觀元素屬性的值

下載DevExpress v20.1完整版 DevExpress Winforms Controls內建140多個UI控制元件和庫，完美構建流暢、美觀且易於使用的應用程式。