什么是Triton-shared?Triton-shared的安装和使用

经过前面几章关于triton在nv gpu上调优的讲解,我们这章开始来看看triton的一个third_party库,该库是为了让triton去支持更多其他的backend。该项目的地址如下所示,并且已经在triton的main分支中,作为third_party进行了官方支持,在clone triton的时候,只需要带上recursive的flag就可以完成对triton-shared的使用。
什么是triton-shared?
关于triton-shared的官方具体实现,如下github repo所示:
github - microsoft/triton-shared: shared middle-layer for triton compilationgithub.com/microsoft/triton-shared
如下所示为官方对triton-shared的解释:
a shared middle-layer for the triton compiler.currently the middle layer is not complete but has enough functionality to demonstrate how it can work. the general idea is that triton ir is lowered into an mlir core dialect to allow it to be both shared across triton targets as well as allow back-ends to be shared with other languages.the basic intended architecture looks like this:[triton ir] -> [middle layer] -> [hw specific ir]the middle-layer uses mlir's linalg and tenor dialects for operations on triton block values. operations on triton pointers use the memref dialect.  
triton-shared其实就是为了提供一个胶水一样的中间层,通过对middle-layer的设计来方便我们的编程语言或者编译器对接到下游不同的硬件生态,因为triton自身已经把nv和amd这两个比较常见的gpu后端实现了,如果第三方的厂商想通过复用triton的前端来对自己的芯片搞一套编译flow,那么triton-shared就起到了决定性的作用。下面这个图是triton的codebase所希望支持的一个愿景,可以看出来,中间这条垂直下来的分支就是triton所支持的nv gpu的优化路线,当用户写完的triton dsl会被翻译成python的ast,然后再从ast到对应的triton dialect,从这一步开始,也就正式将用户手写的成分转到了mlir这套生态,然后再从triton dialect进一步优化到triton gpu dialect,从trition gpu dialect开始,就走了比较标准的llvm代码生成,从llvm ir一路lower到ptx,再到sass,最终可以成功运行在nv的gpu上,这套codegen的路线相比tvm等其他编译框架来说更加的激进,直接越过了nvcc compiler,从而使得整个过程都变成了透明的,对于性能优化来说带来了更多的可能。
img
添加图片注释,不超过 140 字(可选)
triton-shared其实主要是用来cover最右边的分支,因为熟悉mlir的朋友都知道,在右边的分支中,linalg dialect是一个非常重要dialect,该dialect可以去承接很多不同的backend,在主流一些backend的编译优化环节,都会将linalg作为主要的dialect来进行上下游不同dialect之间的转换与对接。
triton-shared的安装
triton-shared的安装其实也很简单,只需要一开始通过recursive来clone整个triton的主分支,然后使用
export triton_codegen_triton_shared=1  
来指明,我们在build triton整个项目的过程中需要使用到triton-shared这个第三方的库。接下来的流程按照triton官方repo的readme一步一步进行即可,有关llvm我是使用的具体commit id下手动编译得到的llvm
llvm commit id: b1115f8ccefb380824a9d997622cc84fc0d84a89triton commit id: 1c2d2405bf04dca2de140bccd65480c3d02d995e  
为什么要选择如上两个固定的commit id,其实理由很简单,因为我前面做过一些关于triton和llvm的开发都是基于上面两个id做的,所以后面我的所有教程以及案例展示都是以这两个commit id为主进行。如果不知道怎么从0开始编译triton,可以参考我之前的教程:
科研败犬丶:openai/triton mlir 第零章: 源码编译70 赞同 · 7 评论文章
triton-shared的使用
讲解完了什么是triton-shared,以及triton-shared怎么安装,接下来,我们就来谈谈如何使用已经被编译好的triton-shared。当你按照我的上述流程编译好triton后,会在该路径下:
/triton/build/tools/triton-shared-opt  
看到一个triton-shared-opt的可执行文件,熟悉mlir的同学可能很快发现该方法其实就是mlir中最基本的opt,该二进制文件可以完成从一个dialect向另外一个dialect的lowering,那么我们使用--help来看看triton-shared-opt的所有功能。如果能在终端中输出如下所示的信息,说明你的triton-shared已经全部安装完毕了。
overview: triton-shared test driveravailable dialects: arith, builtin, cf, gpu, math, scf, triton_gpu, ttusage: triton-shared-opt [options] options:color options:  --color                                                    - use colors in output (default=autodetect)general options:  --abort-on-max-devirt-iterations-reached                   - abort when the max iterations for devirtualization cgscc repeat pass is reached  --allow-unregistered-dialect                               - allow operation with no registered dialects  compiler passes to run    passes:      --affine-data-copy-generate                            -   generate explicit copying for affine memory operations        --fast-mem-capacity=                          - set fast memory space capacity in kib (default: unlimited)        --fast-mem-space=                              - fast memory space identifier for copy generation (default: 1)        --generate-dma                                       - generate dma instead of point-wise copy        --min-dma-transfer=                             - minimum dma transfer size supported by the target in bytes        --skip-non-unit-stride-loops                         - testing purposes: avoid non-unit stride loop choice depths for copy placement        --slow-mem-space=                              - slow memory space identifier for copy generation (default: 0)        --tag-mem-space=                               - tag memory space identifier for copy generation (default: 0)      --affine-expand-index-ops                              -   lower affine operations operating on indices into more fundamental operations      --affine-loop-coalescing                               -   coalesce nested loops with independent bounds into a single loop      --affine-loop-fusion                                   -   fuse affine loop nests...  
这里先来展示
--triton-to-linalg                                     -   convert triton to linalg dialect  
这个pass的使用,因为triton-shared主要就是用来做该优化的。他表示的就是将triton dialect作为输入,然后经过triton-to-linalg这个pass,将其lowering到具有相同语义的linalg dialect上,那triton dialect从哪里来得到呢?不要慌,triton-shared的repo为我们提供了很多mlir格式的文件来方便我们使用该功能,具体路径如下:
/triton/third_party/triton_shared/test/conversion/tritontolinalg/*  
在该教程中,我们使用dot.mlir作为案例进行分析,具体代码如下所示:
// run: triton-shared-opt --triton-to-linalg %s | filecheck %smodule {  tt.func @kernel(    %arg0 : !tt.ptr,    %arg1 : !tt.ptr,    %arg2 : !tt.ptr  )  {    %0 = tt.make_range {end = 128 : i32, start = 0 : i32} : tensor    %c64 = arith.constant 128 : i32    %1 = tt.splat %c64 : (i32) -> tensor    %2 = arith.muli %0, %1 : tensor    %3 = tt.expand_dims %2 {axis = 1 : i32} : (tensor) -> tensor    %4 = tt.broadcast %3 : (tensor) -> tensor    %5 = tt.make_range {end = 64 : i32, start = 0 : i32} : tensor    %6 = tt.expand_dims %5 {axis = 0 : i32} : (tensor) -> tensor    %7 = tt.broadcast %6 : (tensor) -> tensor    %8 = arith.addi %4, %7 : tensor    %10 = tt.make_range {end = 256 : i32, start = 0 : i32} : tensor    %11 = tt.expand_dims %10 {axis = 1 : i32} : (tensor) -> tensor    %12 = tt.broadcast %11 : (tensor) -> tensor    %13 = tt.make_range {end = 64 : i32, start = 0 : i32} : tensor    %c256 = arith.constant 256 : i32    %14 = tt.splat %c256 : (i32) -> tensor    %15 = arith.muli %13, %14 : tensor    %16 = tt.expand_dims %15 {axis = 0 : i32} : (tensor) -> tensor    %17 = tt.broadcast %16 : (tensor) -> tensor    %18 = arith.addi %12, %17 : tensor    %20 = tt.splat %c256 : (i32) -> tensor    %21 = arith.muli %0, %20 : tensor    %22 = tt.expand_dims %21 {axis = 1 : i32} : (tensor) -> tensor    %23 = tt.broadcast %22 : (tensor) -> tensor    %24 = tt.expand_dims %10 {axis = 0 : i32} : (tensor) -> tensor    %25 = tt.broadcast %24 {axis = 0 : i32} : (tensor) -> tensor    %26 = arith.addi %23, %25 : tensor    %30 = tt.splat %arg0 : (!tt.ptr) -> tensor    %31 = tt.addptr %30, %8 : tensor, tensor    %32 = tt.load %31 {cache = 1 : i32, evict = 1 : i32, isvolatile = false}: tensor    %40 = tt.splat %arg1 : (!tt.ptr) -> tensor    %41 = tt.addptr %40, %18 : tensor, tensor    %42 = tt.load %41 {cache = 1 : i32, evict = 1 : i32, isvolatile = false}: tensor    %43 = tt.trans %42 : (tensor) -> tensor    %50 = tt.splat %arg2 : (!tt.ptr) -> tensor    %51 = tt.addptr %50, %26 : tensor, tensor    %52 = tt.load %51 {cache = 1 : i32, evict = 1 : i32, isvolatile = false}: tensor    %60 = tt.dot %32, %43, %52 {allowtf32 = false, maxnumimpreciseacc = 0 : i32} : tensor * tensor -> tensor    tt.store %51, %60 : tensor    tt.return  }}  
上述mlir其实很容易看懂,在%0->%10其实都是triton dialect的内容,该内容表示的就是从上层的triton dsl通过lower转换到对应的triton dialect的过程。其中tt就是表示的该mlir所处的dialect是triton dialect,然后tt.xxx则表示了该dialect所支持的所有operation,有关如何定义一个mlir dialect,我准备拿一个单独的教程来讲。
接下来,只需要在终端中输入
./triton-shared-opt --triton-to-linalg /triton/third_party/triton_shared/test/conversion/tritontolinalg/dot.mlir  
就可以得到从triton dialect转到linag dialect部分对应的内容
#map = affine_map (d0, d1)>module {  func.func @kernel(%arg0: memref, %arg1: memref, %arg2: memref, %arg3: i32, %arg4: i32, %arg5: i32, %arg6: i32, %arg7: i32, %arg8: i32) {    %c256 = arith.constant 256 : index    %c128 = arith.constant 128 : index    %reinterpret_cast = memref.reinterpret_cast %arg0 to offset: [0], sizes: [128, 64], strides: [%c128, 1] : memref to memref    %alloc = memref.alloc() : memref    memref.copy %reinterpret_cast, %alloc : memref to memref    %0 = bufferization.to_tensor %alloc restrict writable : memref    %reinterpret_cast_0 = memref.reinterpret_cast %arg1 to offset: [0], sizes: [256, 64], strides: [1, %c256] : memref to memref    %alloc_1 = memref.alloc() : memref    memref.copy %reinterpret_cast_0, %alloc_1 : memref to memref    %1 = bufferization.to_tensor %alloc_1 restrict writable : memref    %2 = tensor.empty() : tensor    %transposed = linalg.transpose ins(%1 : tensor) outs(%2 : tensor) permutation = [1, 0]     %reinterpret_cast_2 = memref.reinterpret_cast %arg2 to offset: [0], sizes: [128, 256], strides: [%c256, 1] : memref to memref    %alloc_3 = memref.alloc() : memref    memref.copy %reinterpret_cast_2, %alloc_3 : memref to memref    %3 = bufferization.to_tensor %alloc_3 restrict writable : memref    %4 = tensor.empty() : tensor    %5 = linalg.matmul ins(%0, %transposed : tensor, tensor) outs(%4 : tensor) -> tensor    %6 = linalg.generic {indexing_maps = [#map, #map, #map], iterator_types = [parallel, parallel]} ins(%5, %3 : tensor, tensor) outs(%5 : tensor) {    ^bb0(%in: bf16, %in_4: bf16, %out: bf16):      %7 = arith.addf %in, %in_4 : bf16      linalg.yield %7 : bf16    } -> tensor    memref.tensor_store %6, %reinterpret_cast_2 : memref    return  }}  
关于其他更加具体的operator,我们可以都按照上述流程来进行操作,一旦你的编译框架是基于mlir来开发的,那么如果能很好的转到linalg,那么就说明了后续在接入自己的backend以及适配一些isa的过程就会方便不少,这也从另外一个角度彰显了为什么现在的趋势都是将自己的compiler通过mlir进行重构。最重要的原因,其实就是以最小的开发成本方便的接入各种软件或者硬件的生态。
后记
对triton的研究已经有一段时间了,由于当时学triton也是基于源码一步一步硬吃过来的,并且triton也没有比较好的中文教程,所以后面会利用空闲时间将我目前对于使用triton来做codegen的各种优化方法(不同backend以及不同ir层面的pass)和细节(底层layout的设计)进行一个详细的梳理,来帮助更多想要使用triton来做codegen的同学。


智能制造代表工业机器人占据上风
中电数通深耕智慧城市建设 把安全带给千家万户
列车上PCB电路板如何用来抵抗干扰
AMD助力国宝、文物古迹数字化
未来一定是柔性的时代,但时代里是否有柔宇?
什么是Triton-shared?Triton-shared的安装和使用
谷歌和Facebook放弃美国通往中国香港的海底电缆计划
青灯有味是儿时 飞利浦晶亮LED台灯评测
光刻技术是迄今所能达到的最高精度的加工技术
制冷压缩机根据工作原理分几类
虚电路的概念及采用虚电路进行数据传输的过程
PLM、ERP、APS和MES几个信息系统的数据集成以及各自基本功能
企业可以怎样投资区块链
索尼准备和美国卡内基梅隆大学携手,染指家用机器人领域
无刷电机小车开发记录—PWM信号输入捕获驱动
FSH18频谱分析仪18GHz
apple watch2上市时间出来了!同时还支持防水和GPS
红米Pro2要来了,联发科双摄后置指纹
【节能学院】安科瑞电能管理系统在郑大体院的设计与应用
华为平板M3怎么样?华为平板M3测评:极致的影音体验和游戏体验,让你更好的享受生活