TorchDynamo #FX Graph。是一个即时编译器,将一系列 PyTorch 操作转换成一个 FX 图 。
# 这个 FX 图是一个中间表示,可以被不同的后端编译器优化和执行。 TorchDynamo 本身并
# 不直接进行优化,而是将优化工作交给其他的后端,例如 Inductor。
|
|
v
Inductor (这是一个编译器后端 , 它接收 TorchDynamo 生成的 FX 图,并将它们编译成优化的 C++/Triton 内核)
# Inductor 将 TorchDynamo 提供的 FX 图转换成其自身的 IR,这个 IR 考虑了循环融合、内存访问优化等因素。
# 然后,Inductor 会根据目标硬件 (GPU 或 CPU) 和其他配置,将这个自身 IR 转换成优化的代码。
# 对于 GPU,它会使用 Triton 作为代码生成的后端。
(Triton Kernels) # @triton.jit, python function。 用户使用 Python 和 Triton API
# 定义 Triton kernel 的计算逻辑,Triton 编译器将这个 Python 定义的 kernel 转换成实际的 CUDA 或 ROCm 代码,
# 作为最终的 GPU kernel。
|
|
v
## Triton 里边
# FrontEnd 是 python 表达的计算-> Triton IR
TritonDialect, Triton IR [upstream]
# Pass
- Comnbine pass
- Braodcast reordering
- Tensor pointer rewriting
- other optimizers ...
||
||
vv
# Middle End
TritonGPU Dialect, TrtionGPU IR [upstream]
# Pass
- Coalescing
- Layout conversion removal
- Thread Locality optimization
- ...
||
||
vv
# TritonGPU Dialect 中最重要的是通过添加Layout 来改变一个tensor的表示形式,表达一个data在GPU的thread是
# 如何partition的
# .mlir 文件中的变量会通过 Blocked 和 Shared 这两中Layout 类型描述,blocked 表示数据切割方式是有 blocked 定义,
# 如:
# #blocked0 = #triton_gpu.blocked<{versionMajor = 3,
# versionMinor = 0,
# warpsPerCTA = [8, 1],
# CTAsPerCGA = [1, 1],
# CTASplitNum = [1, 1],
# CTAOrder = [1, 0],
# instrShape = [16, 256, 32]}>
# 它定义了数据的切割方式
# %cts : tensor<256xi1, #blocked0>, 逗号前是 mlir 定义的,后边是 TritonGPUDialect 自己定义的
[Shared Layout, Distributed layout (Block layout), Do operand laytout, MMA layout]
||
||
vv
## BackEnd
# Intel (&NV, AMD) 各家vendor自己的Dialect
Intel specified Dialect: TritonGEN, TritonIntelGPU 和 许多优化Pass。
At the same time re-use most of the Triton upstream infrastructure and optimizations
|||
|||
vvv
LLVM Dialect, LLVM IR
|||
|||
vvv
Intel: GenISA/GenX # IGC 编译器得到 SPIRV 中间表示, Intel 没有与 nvcc 对应的工具来生成asambly,
# 所以只能生成 SPIRV,然后使用官方的工具将 SPIRV 翻译为 LLVM
Nvidia: cubin # nvcc 得到 PTX 表示(nvcc 会生成 PTX(一种IR,与硬件无关)代码,并将其传递给 ptxas,
# 最终得到cubin,它是针对特定 GPU 架构编译的二进制可执行文件)
# IGC has 2 path to compile Triton kernels: SIMT & SIMD (Triton 走 SIMT,没有应用IMEX)
# SIMD: lowerTritonGPU IR to lowe level IR and maps to XeGPU Dialect (来自IMEX)
|||
|||
vvv
Runtime # IPEX, 目前是 Stock Pytorch
Tools SPIRV-LLVM-Translator
In-tree mode 在 mlir project 中使用
cmake 中通过使用 FetchContent(cmake, in-tree mode) 作为依赖构建它的编译, 详见 open source code。