Triton是由OpenAI開發(fā)的一個(gè)開源編程語言和編譯器，旨在簡化高性能GPU內(nèi)核的編寫。它提供了類似Python的語法，并通過高級(jí)抽象降低了 GPU 編程的復(fù)雜性，同時(shí)保持了高性能。目前Pytorch已能做到100%替換CUDA，國內(nèi)也有智源研究院主導(dǎo)的FlagGems通用算子庫試圖構(gòu)建起不依賴CUDA的AI計(jì)算生態(tài)，截至今日，F(xiàn)lagGems已進(jìn)入Pytorch基金會(huì)生態(tài)項(xiàng)目體系。Triton生態(tài)內(nèi)少有CPU架構(gòu)的實(shí)踐，且多面向Host-Device的異構(gòu)方案，進(jìn)迭時(shí)空通過同構(gòu)融合RISC-V AI CPU技術(shù)，結(jié)合Triton輕量化的交互式編程模式，將構(gòu)建起比肩Triton GPGPU的AI高性能編程方案，從而推動(dòng)AI應(yīng)用的快速規(guī)模化落地。

為什么是Triton

▲

AI高性能編程模型趨于統(tǒng)一，多核并行的調(diào)度+Tile base的kernel基本成為固定范式。

▲

CUDA的話語權(quán)過高，為走出新AI架構(gòu)的路，需要有獨(dú)立的前端編程語言支撐，而Triton DSL的社區(qū)活躍度足夠高，也有相當(dāng)數(shù)量的大模型、CNN模型項(xiàng)目采用了Triton作為算子編程語言。

▲

Pytorch的成功表明，Python First讓更多開發(fā)者參與生態(tài)共建，降低介入門檻，也有利于新AI架構(gòu)輸出自己的性能優(yōu)化方案。

同構(gòu)融合AI

常見的Host-Device的異構(gòu)Triton方案，使得Triton算子編程的調(diào)試?yán)щy，內(nèi)存模型復(fù)雜，不利于開發(fā)者靈活的實(shí)現(xiàn)自己的想法，而搭建于傳統(tǒng)CPU之上的Triton-CPU方案，也缺乏在AI高性能計(jì)算上的硬件支持，例如核內(nèi)TensorCore、多核通信與訪存優(yōu)化、多卡互聯(lián)等。

進(jìn)迭時(shí)空踐行的同構(gòu)融合技術(shù)，創(chuàng)新性地在CPU內(nèi)集成TensorCore，以RISC-V指令集為統(tǒng)一的軟硬件接口，驅(qū)動(dòng)Scalar標(biāo)量算力、Vector向量算力和 Matrix AI算力，支持軟件和AI模型同時(shí)在RISC-V AI核上運(yùn)行，并通過程序正常跳轉(zhuǎn)實(shí)現(xiàn)軟件和AI模型之間的事件和數(shù)據(jù)交互，進(jìn)而完成整個(gè)AI應(yīng)用執(zhí)行。

基于同構(gòu)融合RISC-V AI CPU架構(gòu)的Triton方案，在編程調(diào)試視角看仍然類似于傳統(tǒng)CPU，并且消除了Host-Device的概念，采用統(tǒng)一內(nèi)存，調(diào)用側(cè)與執(zhí)行側(cè)是Linux軟件多線程的概念，這將極大的降低高性能算子的編程與調(diào)試難度。同時(shí)，在確保編程易用性的前提下，進(jìn)迭時(shí)空通過集成TensorCore、緊密耦合內(nèi)存、Core-to-Core coherence、Cluster-to-Cluster coherence、多核調(diào)度優(yōu)化、AI編譯器優(yōu)化等軟硬件創(chuàng)新，處理絕大部分性能優(yōu)化點(diǎn)，最終交給用戶一個(gè)上手即用的算子開發(fā)工具鏈。

RISC-V AI CPU Triton軟件棧

前端層面，支持Pytorch Triton Kernel以及第三方Triton Kernel，例如FlagGems，支持Triton DSL的全部語義。

中端層面，通過TTIR、TTSIR（Triton Shared）至標(biāo)準(zhǔn)Linag IR，不做任何Dialect擴(kuò)展。

后端層面，先驗(yàn)調(diào)優(yōu)的矩陣乘kernel與vector.contract并存，保證矩陣計(jì)算高效的同時(shí)，釋放更多vector codegen的可能性。

SpineTriton（即進(jìn)迭時(shí)空Triton解決方案）作為Triton的第三方后端，對Pytorch提供RISCV AI-CPU底層加速，兼容社區(qū)已有的Triton Kernel，充分融入現(xiàn)有基于Triton構(gòu)建的AI加速生態(tài)。同時(shí)，針對AI-CPU核內(nèi)擴(kuò)展指令、Core-to-Core高速緩存、異步訪存等特性，對tl.make_block_ptr進(jìn)行了專門特化，開發(fā)者在使用Triton DSL中的塊級(jí)訪存與計(jì)算時(shí)，獲得更大的優(yōu)化收益。

RISC-V AI CPU Triton實(shí)踐

前端

以一個(gè)矩陣乘的Triton Kernel為例，使用tl.make_block_ptr進(jìn)行訪存與計(jì)算。

pid_m = tl.program_id(0)pid_n = tl.program_id(1)# load matmul a and ba_block_ptr = tl.make_block_ptr( base=a_ptr, shape=[M, K], strides=[stride_am, stride_ak], offsets=[pid_m * BLOCK_SIZE_M, 0], block_shape=[BLOCK_SIZE_M, BLOCK_SIZE_K], order=[1, 0])b_block_ptr = ...accumulator = tl.zeros((BLOCK_SIZE_M, BLOCK_SIZE_N), dtype=tl.float32)for k in range(0, tl.cdiv(K, BLOCK_SIZE_K)): a = tl.load(a_block_ptr, boundary_check=(0, 1)) b = tl.load(b_block_ptr, boundary_check=(0, 1)) accumulator += tl.dot(a, b, allow_tf32=False) a_block_ptr = tl.advance(a_block_ptr, (0, BLOCK_SIZE_K)) b_block_ptr = tl.advance(b_block_ptr, (BLOCK_SIZE_K, 0))c = accumulator.to(dot_out_dtype)# maybe some epilogue for cc_block_ptr = tl.make_block_ptr( base=c_ptr, shape=[M, N], strides=[stride_cm, stride_cn], offsets=[pid_m * BLOCK_SIZE_M, pid_n * BLOCK_SIZE_N], block_shape=[BLOCK_SIZE_M, BLOCK_SIZE_N], order=[1, 0],)tl.store(c_block_ptr, c, boundary_check=(0, 1))

上文Triton Kernel描述的矩陣乘計(jì)算對應(yīng)于下圖計(jì)算過程，當(dāng)以一個(gè)Cluster進(jìn)行捆綁調(diào)度時(shí)，SPMD中的Single Program指向一個(gè)Cluster上的執(zhí)行程序，通過Program ID區(qū)分輸入與輸出數(shù)據(jù)位置。以開發(fā)者的視角看，Cluster上的編程是線性的，且不需要關(guān)心異步數(shù)據(jù)的訪問邏輯，后端編譯器將分析用戶代碼邏輯的潛在并行性，在Cluster內(nèi)完成并行化，以及使用高速緩存合并Cluster內(nèi)多核的訪存。

中后端

在矩陣乘內(nèi)部計(jì)算過程的轉(zhuǎn)換時(shí)，將完整的tl.dot即linalg.matmul進(jìn)行分塊分析，充分使用寄存器資源與近核緩存，在中端轉(zhuǎn)為linalg.mmt4d、linalg.pack、linalg.unpack及結(jié)構(gòu)化循環(huán)體的表示。linalg.mmt4d與手寫kernel直接映射并利用到Tensor算力，而其他的算子，則采用affine進(jìn)行向量化使用Vector算力。

由于采用了IME的方式擴(kuò)展AI指令（參考進(jìn)迭時(shí)空AI擴(kuò)展指令Spec，https://github.com/spacemit-com/riscv-ime-extension-spec），在linalg.mmt4d這樣的ukernel的轉(zhuǎn)換過程時(shí)，可以直接使用vector進(jìn)行交互，避免在延遲更高的存儲(chǔ)結(jié)構(gòu)上進(jìn)行交互，這是IME的一大優(yōu)點(diǎn)。

// load b// %acc: vector<16x32xf32>%0 = vector.load [...] : memref, vector<4x32xf32>// load a%1 = vector.load [...] : memref, vector<2x32xf32>// vfmadot -> 2x8x4 @ 4x8x4 => 2x4x8x8%2 = vector.contract {...} %1, %0, %acc : vector<2x32xf32>, vector<4x32xf32> into vector<16x32xf32>

在mlir-llvm的結(jié)合部分，通過vector.contract構(gòu)造了大量先驗(yàn)的手寫匯編序列，以確保最終性能的可靠性。

結(jié)束語

Triton目前仍然是一個(gè)GPGPU架構(gòu)主導(dǎo)的Python DSL及算子編譯器，在CPU架構(gòu)上發(fā)展緩慢，僅存在一些在x86架構(gòu)下的TritonCPU編程的社區(qū)工作，且不是最優(yōu)適配。RISC-V同構(gòu)融合AI算力的方式，利于打破算子內(nèi)多種計(jì)算模式（Scalar、Vector、Tensor）的隔閡，同時(shí)統(tǒng)一內(nèi)存、統(tǒng)一OS的軟硬件架構(gòu)，使得調(diào)試難度降低，系統(tǒng)內(nèi)多種軟硬件資源的交互難度降低。此外，未來也將逐步開源SpineTriton的軟件棧部分，共同建設(shè)RISCV Triton高性能編程社區(qū)。