窥探Trition的lower(一)

Triton的存在使得编写针对特定Target的算子的难度大大降低，相比于以往暴露更多硬件细节的编程方法，Triton编程提供了更加简洁的python接口去描述一个算子，并通过multi Level IR(MLIR)(也是多层Dialect)渐进式地lower，从设计上能够更灵活地支持抽象级别的扩展（在最后会对这）。

关于triton的宏观概念和流程，这里就不再展开了。这篇文章的意义在于当你了解了triton的作用（以及一些LLVM的知识），并成功安装后，想要通过实际操作来窥探到Triton内部的运行流程。由于Triton的很多内部函数通过pybind11交叉编译，跳转的过程有时不太直白。因此这篇文章希望通过一个简单的case，来帮助分析和跟踪它每一步在做什么，同时也是自己学习时的记录。此外，这篇文章重点关注编译的流程，因此和runtime/driver相关的部分会暂时忽略。如果您想复现文章中所用代码，需要准备编译好的Triton（https://github.com/triton-lan...）、LLVM（https://github.com/llvm/llvm-...）项目以及带有nvidia GPU的环境。

大体上讲，Triton的编译过程是从源语言生成 AST（Abstract Syntax Tree，抽象语法树），借助 dialect遍历 AST，产生 MLIR 表达式（此处可为多层IR通过 Lowering Pass 依次进行分析），最后经过 MLIR 分析器，生成目标硬件程序，简单来讲也就是source->AST->dialects，多条dialects组成ttir，ttir->ttgir->ttllir->backend。这篇文章主要关注前一个阶段，剩下的在后面阐述。关于dialect 的概念，我们可以暂时按字面理解为"方言"，各个Dialect分别对不同的层级概念进行建模。比如LLVM Dialect负责系统级别的转换，Linalg，Tensor，Vector等Dialect负责协同生成代码，而Affine，Math等Dialect用来描述底层计算。后面我们会通过代码具体了解Dialect的作用。

首先我们用python写一个最简单的case，代码中包含了一个kernel和几行编译代码，这个kernel由triton的jitFunction装饰，提供了被编译所需的必要接口。其中的操作就是一个普通的add和load store，同时为了代码和生成的ir简洁，这里设置了输入元素个数和block_size都为1。编译代码部分，指明了这个kernel的来源是AST（也可以直接写成IR），以及参数类型信息和目标Target信息，最后会打印出各个level的IR。

import tritonimport triton.language as tl@triton.jitdef addi_kernel(x_ptr,  # *Pointer* to first input vector.
               output_ptr,  # *Pointer* to output vector.
               n_elements,  # Size of the vector.
               BLOCK_SIZE: tl.constexpr,  # Number of elements each program should process.
               ):
    x = tl.load(x_ptr)
    output = x + 1
    tl.store(output_ptr, output)src = triton.compiler.ASTSource(fn=addi_kernel, signature="*i32,*i32", constants={"n_elements": 1,"BLOCK_SIZE": 1})ret = triton.compile(src, target=("cuda", 80))for k in ['ttir', 'ttgir', 'llir']:
    print(ret.asm[k])

这里的代码应该看上去足够清晰简洁了，运行成功的话会输出指定的IR，在下一篇文章会展示出它们的具体内容。在此之前，我们先从代码开始，跟踪一下triton是如何将代码逐步lower到多层IR的。

首先triton.compile函数开始。将上述kernel封装成一个ASTSource后，会进入代码的python/triton/compiler/compiler.py/compile函数，这里一些不太重要的功能用省略号跳过了，比如利用cache编译、解析选项等等。

def compile(src, target=None, options=None):
    ......
    backend = make_backend(target) # 根据指定的target获取一个backend，这里会返回CUDABackend
    ......
    backend.add_stages(stages, options) # 添加编译stage
    ......
    ir.load_dialects(context)  # 加载所需dialects
    backend.load_dialects(context)
    module = src.make_ir(options, context) # 创建IR
    ......
    for ext, compile_ir in list(stages.items())[first_stage:]:
        next_module = compile_ir(module, metadata) # 编译各个阶段IR
        ......
        module = next_module
    ......
    return CompiledKernel(src, metadata_group, hash)

这部分是编译的核心代码，上面的步骤可以分成五个主要阶段：make_backend、add_stage、load_dialects、make_ir、compile_ir。而我们这一篇文章会先分析前四个阶段（主要是load_dialects和make_ir），最后一个阶段放在下一章描述。

• make_backend & add_stage

前两个阶段的内容比较好理解，make_backend阶段我们会通过指定的Target名字获取一个CUDABackend，这其中还包括对应driver初始化的一些流程，这里暂时先不关注。然后调用它的add_stage成员函数，

def add_stages(self, stages, options):
        stages["ttir"] = lambda src, metadata: self.make_ttir(src, metadata, options)
        stages["ttgir"] = lambda src, metadata: self.make_ttgir(src, metadata, options, self.capability)
        stages["llir"] = lambda src, metadata: self.make_llir(src, metadata, options, self.capability)
        stages["ptx"] = lambda src, metadata: self.make_ptx(src, metadata, options, self.capability)
        stages["cubin"] = lambda src, metadata: self.make_cubin(src, metadata, options, self.capability)

这些stage会编译到不同级别的IR，后面两个阶段会借助llvm以及ptxas，前面三个stage会经过若干个pass对IR进行分析或优化。这些阶段在最后compile_ir步骤实际运行。

• load_dialects

这是我们需要着重follow的一个步骤，主要用来理解dialect。根据上面的代码我们发现这里有两个步骤：ir.load_dialects和backend.load_dialects，首先关注第一个。这里的ir对象是从libtriton.so文件import的，最初是在编译项目的时候由python/src/http://ir.cc（http://ir.cc/）通过pybind11交叉编译得到的，使得在python中直接调用，这里是pybind的语法

m.def("load_dialects", [](MLIRContext &context) {
    DialectRegistry registry;
    registry.insert<TritonDialect, ::mlir::triton::gpu::TritonGPUDialect,
                    math::MathDialect, arith::ArithDialect, index::IndexDialect,
                    scf::SCFDialect, ::mlir::gpu::GPUDialect,
                    cf::ControlFlowDialect, LLVM::LLVMDialect>();
    registerBuiltinDialectTranslation(registry);
    registerLLVMDialectTranslation(registry);
    context.appendDialectRegistry(registry);
    context.loadAllAvailableDialects();
  });

其中可以看到在一个注册器中注册了一些dialect，这些是triton的公共dialect，后面的一些register函数是在llvm/mlir中用来注册的辅助函数。那么这些dialect都是什么呢，我们选择其中一个arith::ArithDialect来分析，因为我们写的kernel中的加法操作最终会被ArithDialect的addi这个op表示。ArithDialect的定义是在llvm项目里的mlir/include/mlir/Dialect/Arith/IR/ArithBase.td，熟悉llvm的应该比较理解tablegen的写法

def Arith_Dialect : Dialect {
  let name = "arith";
  let cppNamespace = "::mlir::arith";
  let description = [{
    ......
  }];

  let hasConstantMaterializer = 1;
  let useDefaultAttributePrinterParser = 1;
}

这里包含了一些简单属性，在这个dialect中，还定义很多基础操作，用来表示我们源代码中的计算，这些定义存在在ArithOps.td中，我们选择其中的Arith_AddIOp来看

def Arith_AddIOp : Arith_IntBinaryOpWithOverflowFlags<"addi", [Commutative]> {
  let summary = "integer addition operation";
  let description = [{
    ......
  }];
  let hasFolder = 1;
  let hasCanonicalizer = 1;
}

这里指明了它的名字，是否满足交换律等等属性，最终继承自一个大类Op。刚看到这里的时候可能有些疑问，这里tablegen的定义和LLVM中的有些区别，通常在LLVM后端编写td文件时，对定义的操作都会指明一个对应的模式相匹配，这样在编译的时候才能顺利lower。因此，这里一个重要的事情是理解源代码中的加法操作是怎么变成Arith_AddIOp的，到后面我们会知道这个问题的答案。

除了ir.load_dialects步骤，还有backend.load_dialects，针对nvidia的后端这里会跳转到

  m.def("load_dialects", [](mlir::MLIRContext &context) {
    mlir::DialectRegistry registry;
    registry.insert<mlir::triton::nvidia_gpu::TritonNvidiaGPUDialect,
                    mlir::triton::nvgpu::NVGPUDialect>();
    mlir::registerNVVMDialectTranslation(registry);
    context.appendDialectRegistry(registry);
    context.loadAllAvailableDialects();
  });

nvidia的dialect这里注册了两个，一个是在triton中的include/triton/Dialect/TritonNvidiaGPU/IR/TritonNvidiaGPUDialect.td，一个是在llvm中mlir/include/mlir/Dialect/NVGPU/IR/NVGPU.td，这两者在编译到更低层次的ir时，会用到这些dialects，具体的在下一篇文章阐述。

• make_ir

在将需要的dialect load到上下文管理器之后，我们会开始创建一个初始ir。这里会根据输入文件形式的区别有所不同。triton支持源代码输入和MLIR输入。如果是MLIR，它的调用经过make_ir->parse_mlir_module->parseSourceFile->parseAsmSourceFile->TopLevelOperationParser::parse到达一个解析器，根据token的类型执行前端的解析。

ParseResult TopLevelOperationParser::parse(Block *topLevelBlock,
                                           Location parserLoc) {
  // Create a top-level operation to contain the parsed state.
  OwningOpRef<ModuleOp> topLevelOp(ModuleOp::create(parserLoc));
  OperationParser opParser(state, topLevelOp.get());
  while (true) {
    switch (getToken().getKind()) {
    default:
      // Parse a top-level operation.
      if (opParser.parseOperation())
        return failure();
      break;

    // If we got to the end of the file, then we're done.
    case Token::eof: {
      if (opParser.finalize())
        return failure();
     .......

  }
}

由于我们一般不采用这种形式，而且这部分一般不涉及到修改，所以就不深入了。

对于源代码的形式，make_ir的过程是make_ir->ast_to_ttir->ast.parse->generator.visit，在这一过程中ast.parse会把源代码解析成抽象语法树(AST)，其中树的节点包含了源代码的各种抽象模块，比如module、function、args、attr、name、op等等。而在generator.visit阶段，会对这些节点进行遍历lower。我们选择代码中的加法操作进行观察，首先在ast.parse阶段，"+"这个操作被映射成了一个ast.BinOp节点，在visit的过程中，会调用ast.py中NodeVisitor的'visit_' + node.__class__.__name__方法，调用到如下visit_BinOp函数

def visit_BinOp(self, node):
    lhs = self.visit(node.left)
    rhs = self.visit(node.right)
    method_name = self._method_name_for_bin_op.get(type(node.op))
    if method_name is None:
        raise self._unsupported(node,
                                "AST binary operator '{}' is not (currently) implemented.".format(node.op.__name__))
    return self._apply_binary_method(method_name, lhs, rhs)

首先对二元操作数的左右节点分别visit后（visit_name/constant），通过调试可知获取到的method_name为"__add__"，经过apply函数会调用如下函数python/triton/language/core.py

@builtin
def __add__(self, other, _builder=None):
    other = _to_tensor(other, _builder)
    return semantic.add(self, other, _builder)

进而调用python/triton/language/semantic.py/semantic.add：

def add(input: tl.tensor, other: tl.tensor, builder: ir.builder) -> tl.tensor:
    input, other = binary_op_type_checking_impl(input, other, builder, True, True)
    input_scalar_ty = input.type.scalar
    other_scalar_ty = other.type.scalar
    ......
    # int + int
    elif input_scalar_ty.is_int():
        return tl.tensor(builder.create_add(input.handle, other.handle), input.type)
    assert False

这里我们看到builder.create_add这个函数，再次跳转到python/src/http://ir.cc

.def("create_add",
       [](TritonOpBuilder &self, Value &lhs, Value &rhs) -> Value {
          return self.create<arith::AddIOp>(lhs, rhs);
       })

终于我们看到一个熟悉的东西，arith::AddIOp，也是就是我们上面分析的arithdialect中的op。因此目前我们可以得知，源代码中的"+"操作，先解析为ast的节点ast.BinOp(op.name="__add__")，进而利用dialect表示为arith::AddIOp。在这一阶段，代码已经被表示成了如下ir

#loc = loc("toy.py":28:0)
module {
  tt.func public @addi_kernel_01(%arg0: !tt.ptr<i32, 1> loc("toy.py":28:0), %arg1: !tt.ptr<i32, 1> loc("toy.py":28:0)) attributes {noinline = false} {
    %0 = tt.load %arg0 {cache = 1 : i32, evict = 1 : i32, isVolatile = false} : i32 loc(#loc1)
    %c1_i32 = arith.constant 1 : i32 loc(#loc2)
    %1 = arith.addi %0, %c1_i32 : i32 loc(#loc2)
    tt.store %arg1, %1 {cache = 1 : i32, evict = 1 : i32} : i32 loc(#loc3)
    tt.return loc(#loc4)
  } loc(#loc)
} loc(#loc)
#loc1 = loc("toy.py":38:16)
#loc2 = loc("toy.py":39:17)
#loc3 = loc("toy.py":40:25)
#loc4 = loc("toy.py":40:4)

可以看到此时加法的形式在ir中是用arith.addi的dialect表示的。

目前，整个编译流程只剩下compile_ir这一步了，但其实这一步包含更多内容，包括一些硬件特性信息的添加，当我们想在triton中支持一个新后端时，也更需要关注这一部分。下一篇会尽快完成，欢迎关注~

The End

作者：液态黑洞
来源：GiantPandaCV

推荐阅读

• MLIR_对自定义IR Dialect编写bufferization pass
• How to Do the Paper/Talk Reviews
• ​SIMD 指令集与数据并行程序
• 单处理器实现大规模计算能力的有效性

欢迎大家点赞留言，更多Arm技术文章动态请关注极术社区嵌入式AI专栏欢迎添加极术小姐姐微信（id:aijishu20)加入技术交流群，请备注研究方向。