Triton編譯器的原理和性能

我們推出了一個(gè)新的系列，對(duì)PytorchConference2023 的博客進(jìn)行中文編譯，會(huì)陸續(xù)在公眾號(hào)發(fā)表。

Triton是一種用于編寫(xiě)高效自定義深度學(xué)習(xí)原語(yǔ)的語(yǔ)言和編譯器。Triton的目的是提供一個(gè)開(kāi)源環(huán)境，以比CUDA更高的生產(chǎn)力編寫(xiě)快速代碼，但也比其他現(xiàn)有DSL具有更大的靈活性。Triton已被采用為T(mén)orch inductor的基本組件，以合成針對(duì)GPU的高效內(nèi)核。與傳統(tǒng)庫(kù)使用相比，這具有多種優(yōu)勢(shì)。它允許創(chuàng)建各種各樣的融合，它可以獨(dú)立調(diào)整，并且它的內(nèi)存占用更小。本次演講將介紹Triton編譯器，并描述使其能夠以最少的用戶(hù)努力生成閃電般快速內(nèi)核的過(guò)程。

全文

今天我要和大家談?wù)劦氖荰riton。那么，我將要討論的大致內(nèi)容是Triton是什么？我們?yōu)槭裁匆獎(jiǎng)?chuàng)建這個(gè)工具？它可以用來(lái)做什么？然后，我將討論如何將其集成在ML編譯器堆棧中。最后，我將簡(jiǎn)要介紹其背后的原理以及編譯器是如何簡(jiǎn)化管理的。

Triton是一個(gè)Python DSL（領(lǐng)域特定語(yǔ)言），旨在用于編寫(xiě)機(jī)器學(xué)習(xí)內(nèi)核。 最初，它嚴(yán)格用于GPU內(nèi)核，但慢慢地?cái)U(kuò)展以支持用于機(jī)器學(xué)習(xí)的任何硬件，包括CPU、ASIC等。Triton的目標(biāo)是讓那些沒(méi)有GPU經(jīng)驗(yàn)的研究人員能夠編寫(xiě)高性能代碼。如果你看到幻燈片底部的圖表，那真的是Triton想要達(dá)到的地方。通過(guò)少量的開(kāi)發(fā)工作，你可以非常接近峰值性能。

簡(jiǎn)而言之，Triton是一個(gè)幫助研究人員輕松編寫(xiě)高性能機(jī)器學(xué)習(xí)內(nèi)核的工具，無(wú)論他們是否有GPU經(jīng)驗(yàn)。

當(dāng)然，總是會(huì)有像CUDA或匯編語(yǔ)言這樣的其他語(yǔ)言，它們能讓你獲得同樣或更高的性能，但通常你需要對(duì)硬件有更多的了解，并花費(fèi)更多的時(shí)間。為什么我們需要這種新的語(yǔ)言呢？如果你看看現(xiàn)有的選擇，例如在不同的硬件上編程機(jī)器學(xué)習(xí)，有PyTorch這樣的工具，它允許你輕松地將不同類(lèi)型的操作映射到硬件上，并且非常容易從中獲得高性能。

但問(wèn)題在于你對(duì)它的控制非常有限。如果現(xiàn)有的操作集中沒(méi)有你需要的東西，你就只能束手無(wú)策，唯一的解決辦法是走向另一個(gè)極端，例如編寫(xiě)CUDA或編寫(xiě)PTX，甚至直接編寫(xiě)匯編代碼。但問(wèn)題在于，要編寫(xiě)這些語(yǔ)言，你需要真正成為硬件方面的專(zhuān)家，并且用這些語(yǔ)言編寫(xiě)高效的內(nèi)核可能非常棘手 。所以Triton實(shí)際上是嘗試在這里找到一個(gè)中間地帶，它允許用戶(hù)編寫(xiě)高效的內(nèi)核，并有大量的控制權(quán)，但又不必關(guān)心那些微小的細(xì)節(jié)。

是的，硬件的細(xì)節(jié)以及如何在特定硬件上獲得性能。實(shí)際上，設(shè)計(jì)的難點(diǎn)在于找到這個(gè)最佳平衡點(diǎn)。Triton的設(shè)計(jì)方式就是找到這個(gè)抽象的平衡點(diǎn)，即你想向用戶(hù)暴露什么，以及你想讓編譯器做什么？

編譯器是生產(chǎn)力工具，真的……在這方面，Triton的目標(biāo)是讓編譯器為你完成你不想做的工作，但仍然讓你能夠控制算法、你想要用來(lái)進(jìn)行調(diào)整的任何tuning。Triton介于Cuda和Torch之間，因?yàn)槟闳匀豢梢跃帉?xiě)自己的算法，你仍然可以控制自己的類(lèi)型，你仍然需要決定是否需要以某種類(lèi)型來(lái)保存中間值，你控制所有的精度。你不必關(guān)心如何處理共享內(nèi)存、在目標(biāo)有張量核時(shí)使用張量核、如何很好地處理負(fù)載聚合，以便你有良好的內(nèi)存訪(fǎng)問(wèn)模式。 這些人們?cè)诰帉?xiě)GPU內(nèi)核時(shí)經(jīng)常要考慮的事情。你總是要擔(dān)心這些問(wèn)題，或者弄清楚我的中間數(shù)據(jù)的布局是什么等等。編譯器會(huì)為你完成這些工作。

讓我們來(lái)看一個(gè)例子。這是一個(gè)softmax內(nèi)核的示例。這是一個(gè)工作解決方案的復(fù)制品，它是有效的。

#https://github.com/openai/triton/blob/main/python/tutorials/02-fused-softmax.py
@triton.jit
defsoftmax_kernel(output_ptr,input_ptr,input_row_stride,output_row_stride,n_cols,BLOCK_SIZE:tl.constexpr):
#Therowsofthesoftmaxareindependent,soweparallelizeacrossthose
row_idx=tl.program_id(0)
#Thestriderepresentshowmuchweneedtoincreasethepointertoadvance1row
row_start_ptr=input_ptr+row_idx*input_row_stride
#Theblocksizeisthenextpoweroftwogreaterthann_cols,sowecanfiteach
#rowinasingleblock
col_offsets=tl.arange(0,BLOCK_SIZE)
input_ptrs=row_start_ptr+col_offsets
#LoadtherowintoSRAM,usingamasksinceBLOCK_SIZEmaybe>thann_cols
row=tl.load(input_ptrs,mask=col_offsets

	

	第一個(gè)有趣的事情是這段代碼相對(duì)較短。如果你用CUDA編寫(xiě)同樣的內(nèi)核，它實(shí)際需要更多的努力。我們可以注意到一些有趣的事情。例如，你可以控制如何在計(jì)算機(jī)上分配工作。多虧了這些編程思想。你可以看到，你仍然可以控制你的內(nèi)存訪(fǎng)問(wèn)，因?yàn)槟憧梢栽L(fǎng)問(wèn)指針。你可以基于一些原始指針加載一大塊數(shù)據(jù)。然后編譯器將在后臺(tái)決定將其映射到硬件的最佳方式，以及如何進(jìn)行聚合，如何處理所有事情，以便這個(gè)加載將是有效的，并將分布到你的GPU的不同線(xiàn)程和warp上。但你不必?fù)?dān)心這些。在底部，我們可以看到有一個(gè)歸約操作，通常它會(huì)隱式地使用共享內(nèi)存，但你不必?fù)?dān)心它。編譯器將確保你為其選擇最佳實(shí)現(xiàn)，并為你使用共享內(nèi)存。

	之后我將討論，如何在典型的設(shè)備上使用triton，除了內(nèi)核他還可以集成到完整的graph編譯器堆棧中：

	

	Triton為你提供了一個(gè)非常容易、非常自然的從graph表示直接到實(shí)現(xiàn)的lowering過(guò)程，并且它實(shí)際上允許更簡(jiǎn)單的graph表示實(shí)現(xiàn)，因?yàn)槟悴槐匾淮涡陨梢粋€(gè)完美的內(nèi)核。你可以只生成Triton部分，然后Triton編譯器將完成繁重的工作，找出如何有效地將其映射到硬件上。

	Triton可以被用作的另一個(gè)地方是它可以被用作自定義操作語(yǔ)言 。像PyTorch這樣的工具，因?yàn)槿绻阆萑肜Ь常鳳yTorch中沒(méi)有實(shí)現(xiàn)某些功能，添加自定義操作是你能夠完成你想要做的事情的唯一解決方案。

	讓我們稍微看一下編譯器架構(gòu)。這是一個(gè)非常高層次的查看Triton架構(gòu)的方式。

	

	Triton被構(gòu)建為一個(gè)老式編譯器，包括前端、中端和后端。這里有趣的部分是這兩個(gè)塊，Triton IR和Triton GPU IR，它們是Triton的中間IR，這里有很多魔法發(fā)生。你可以在這里看到的另一件有趣的事情是，Triton IR真的允許你針對(duì)不同的硬件進(jìn)行定位，因?yàn)門(mén)riton IR本身對(duì)于這硬件是完全無(wú)關(guān)的。如果我們放大這個(gè)有趣的部分，即基本上發(fā)生在Triton IR和最終的LLVM IR之間的事情，LLVM IR是最終的目標(biāo)。

	

	基本上，編譯器首先接收Triton IR，Triton IR與語(yǔ)言本身非常相似。然后，編譯器要做的第一件事是為描述張量如何分布到線(xiàn)程上的布局進(jìn)行關(guān)聯(lián)。這真的是編譯器的核心機(jī)制，因?yàn)榛谶@些布局，有多種路徑可以改變這些布局，并能夠生成一些能夠有效地映射到硬件上的東西。因此，我們會(huì)像進(jìn)行coalesce一樣，嘗試選擇一個(gè)布局，以便加載存儲(chǔ)聚合能夠高效進(jìn)行。

	如果機(jī)器有tensorcore，我們會(huì)嘗試使用非常適合tensorcore的布局。然后，我們會(huì)嘗試避免任何布局轉(zhuǎn)換，應(yīng)用一系列典型的編譯器傳遞，然后在此基礎(chǔ)上進(jìn)行轉(zhuǎn)換，基于分析轉(zhuǎn)到llvm ir。

	這是非常高層次的，但這就是編譯器的工作原理。嗯，這就是我想告訴你的全部?jī)?nèi)容。Triton正在完全開(kāi)源的情況下進(jìn)行開(kāi)發(fā)，非常歡迎貢獻(xiàn)者。我們每個(gè)月都會(huì)舉行社區(qū)會(huì)議。

	Triton IR本身對(duì)硬件無(wú)關(guān)。但是，如果你把一個(gè)在目標(biāo)上運(yùn)行良好的內(nèi)核拿過(guò)來(lái)，你可能需要重新調(diào)整它，以便在另一個(gè)目標(biāo)上運(yùn)行良好。

	審核編輯：湯梓紅