Skip to content

Tutorial[CN]

Jump to bottom
HarmonyHu edited this page Nov 14, 2022 · 1 revision

简介

TPU-MLIR是算能智能AI芯片的TPU编译器工程。该工程提供了一套完整的工具链，其可以将不同框架下预训练的神经网络，转化为可以在算能TPU上高效运算的二进制文件 BModel 。代码已经开源到 GitHub 。TPU-MLIR提供了一种使用MLIR工程对DSA定制编译器的方式，该工程目前支持Sophgo的TPU芯片，包含1684X，1684以及18xx，也用于支持后续更多芯片。TPU-MLIR作为Sophgo的开源编译器工程，将长期用于支持Sophgo芯片的发展。

本文档用于描述TPU-MLIR工程代码结构并提供例子用于指导开发者入门。同时提供一写辅助信息用于帮助开发者理解整个编译结构和背后的逻辑。

工程结构

TPU-MLIR工程目录结构和编译流程图的关系如下所示（红色部分为前端，绿色部分为TOP层以及与之有关功能，蓝色部分为TPU层以及与之有关的功能）：

arch linked with code

更多详细信息可以参考整体设计。

TPU-MLIR工程编译需要使用Sophgo提供的docker镜像，更多关于如何编译TPU-MLIR工程信息可以参考开发环境配置。

TPU

Sophgo TPU的计算方式是标准的SIMD结构，同时是Load-Store结构。TPU的计算单元只能读写local memory（SRAM），整个计算流程为：数据先从外部存储（DDR）中Load到local memory中，然后EU完成计算，之后再将数据存回到DDR。

在1684X中，每个TPU有64个NPU，每个NPU上有相应的EU和local memory，一个NPU也称为一个Lane（后文中会混合使用这两个名词）。其中EU的数目会根据数据类型有变化，如：对于FP32/INT32 运算EU数为16，FP16/BF16/INT16 运算EU数为32，INT8 运算EU数为64；而每个Lane的local memory是一样的，都为256KB。所以单个1684X芯片的local memory为16MB。

和通用SIMD结构一样，1684X的每个NPU/Lane中的EU只能访问自己的local memory，且同一时刻所有的NPU执行的计算都是一样的。在计算的时候，需要先将数据切分，然后分散到不同的Lane上，计算完成后再将数据取回。

数据切分和取回可以使用高效且灵活的GDMA，GDMA起到了local memory、DDR、HAU L2 memory以及Static memory之间交换数据的功能，且与EU可以同时工作。

EU的一些计算命令会涉及到数据的broadcast（如convolution和matrix multiplication），此操作对用户不可见，在指令文档或是OK kernel编程中会有详细描述。

更多关于TPU的结构介绍以及编程模型（包含计算模型和内存模型）请参考Sophgo OKKernel文档。

内存结构，访问方式

不同于高级编程模型，如Tensor在TensorFlow、Pytorch以及Numpy中只是一个高纬数据，使用对应的函数（slice，tranpose，reshape等）可以获得对Tensor数据的抽取和查看，Tensor在具体硬件上需要考虑硬件的限制和结构。在高级编程中可线性访问（forward iterator）的数据在硬件中储存可能是非连续的。

Sophgo TPU的数据都是以4维Tensor[batch, channel, high, wide]为主，且默认会将channel分割到不同Lane的local memory中。粗略的一个例子：对于一个shape为[4, 23, 8, 8]的Tensor，其数据在TPU的每个local memory分布可能会是[4, 1, 8, 8]，且只有前23个NPU的local memory有有效的数据，后续的NPU中对应位置数据无效。更对关于内存分布和使用的例子可以参考Sophgo OKKernel Layouts。

数据分割和编排是使用SIMD硬件较难的部分，其中有很多窍门和陷阱^1。Stride是个非常有效且优雅的方式，用于描述Tensor在memory中的分布^2，如果感兴趣这部分，可以使用Numpy的ndarray^4来做验证和探索。如在上述例子中，如果认为64个NPU的local memory是连续编址的，那么[4, 23, 8, 8]Tensor在整个local memory的layout的stride为(84, 256K, 8, 1)。或是设置shape[23, 4, 1, 8, 8]，此时的Stride为[256K, 64, ?, 8, 1]。Stride给出一种将高纬数据与线性地址的映射，且这种映射一般是双射，即支持subscripts与 linear indices互相转化。

硬件文件格式

Sophgo TPU的可执行文件为自定义的二进制数据，称作BModel。BModel是一种面向算能TPU处理器的深度神经网络模型文件格式，其中包含目标网络的权重（weight）、TPU指令流等等。

该文件中存储形式使用了多段（section）式，包含数据头，数据描述段，代码段，数据段。数据头描述了数据描述段，数据描述段又描述了数据段和代码段，形成一种树形结构。其中数据描述段使用了flattenbuffer，相关的内容可以可以参考bmodel.fbs，BModel整体定义可以参考bmodel.hpp。

模型编译流程

整个编译流程涉及多个优化和验证，为了便于用户使用，TPU-MLIR定义了两个python工具来驱动整个编译过程：

关于CLI使用细节，可以参考用户界面。

添加后端

待开放

TOP/TPU推理

TPU-MLIR中提供两级抽象描述TOP和TPU层。

  • Top Dialect：与芯片无关层，包括图优化、量化、推理等等
  • Tpu Dialect：与芯片相关层，包括权重重排、算子切分、地址分配、推理等等

为了便于debug，TPU-MLIR支持在这两层上做推理运算（Inference）,通过比较计算的所有中间数据与前端（通常为onnxruntime）计算结果的相似性，来确认计算的正确性。同时对TOP层的推理也用于Calibration过程。

模型量化

模型量化是一个神经网络部署的特别领域，其中涉及到对特定加速硬件定制计算以及考虑网络结构和参数特性，调整原始网络，使其能在特定硬件上有最好的性能表现。由于神经网络有较多的冗余参数且有很好的抗噪声特性，这让量化方式可以有效地应用到神经网络中；同时神经网络的可解释性不强，这导致量化神经网络会有较多的edge cases，且这些cases较难解决。训练量化虽然可以表现出很好的性能，但是其设置流程较多，且需要的时间较长。训练后量化设置流程较少，且不需要label，整个处理流程时间很短，在快速部署上很常见。TPU-MLIR提供的量化方式为训练后量化，相关量化细节可以参考quantization。

训练后量化基本分为两个过程：

  1. 统计在少量样例下（inputs）每个activation /mutable Tensor的数据分布，根据数据分布，计算数据的范围min/max，此过程也称作calibration。相关代码位置calibration。
  2. 根据上述得到的数据范围min/max，以及整个神经网络和部署芯片的计算特点，将浮点计算调整为定点计算。相应的代码位于calibration_process。

调试

TPU-MLIR可以依赖MLIR工程提供的强大分析工具来debug。这里介绍一些TPU-MLIR额外的调试方式：

  1. 算子级测试：

    在test_onnx.py中可以加入算子级测试，将一个或多个算子组合进行全流程测试。如：单测 test_onnx.py Conv2d；全测 test_onnx.py。

  2. 模型截断：

    使用model_transform.py的时候指定--output_names。

  3. 对MLIR文件进行裁剪后编译。使用model_transform.py和model_deploy.py时工具都会打印实际执行的命令（多数是tpuc-opt），可以借助这些命令分步编译模型，缩小错误范围。

  4. 使用经典调试工具pdb和gdb。如使用pdb：

    import pdb
pdb.set_trace()

    或是gdb：

    gdb --args python /work/python/tools/model_runner.py --input resnet18_in_f32.npz --model resnet18_opt.mlir --output resnet18_out.npz

gdb --args tpuc-opt Resize_f32.mlir --weight-reorder --subnet-divide --layer-group --address-assign --save-weight --codegen="model_file=Resize_f32.bmodel" -o Resize_f32_final.mlir

代码规范

  1. 使用clang-format或是其他识别.clang-format工具格式化C++代码；

  2. 文件头添加license（本工程使用2-Clause BSD License）描述；

  3. 代码结构要符合TPU-MLIR设计哲学和结构。

实用工具

  1. TPU-MLIR已经添加了vscode的配置，使用vscode开发会比较容易；
  2. 推荐使用Netron来查看原始模型；
  3. MLIR代码有良好的语法定义，可以对vscode添加相关的语法高亮插件来提升效率；
  4. 如果有条件，可以安装LSP Server提升编写TableGen/MLIR文件的效率。

如何贡献

修改文档

TPU-MLIR的文档位于tpu-mlir/doc，并使用reStructuredText格式来编写，该格式和Markdown类似，都是文档标记语言，但是语法上有很大差别。关于Rst文档语法部分，可以参考Sphinx/Rest Memo。

Rst文档可以输出多种格式的最终文件，TPU-MLIR主要输出的格式为PDF和HTML。

贡献代码

添加算子

  1. 参考该算子的ONNX定义文档，然后在OnnxConverter中定义相应的函数，将该ONNX算子转换到TOP层对应的算子（如果没有对应的算子，也需要对TOP层添加对应算子），此过程称作为前端转换，详细文档可以参考前端转换。

  2. 经过步骤1，此时的网络已经从外部框架转化为MLIR格式。如果该算子可以优化合并，那么可以通过添加Canonicalize优化pattern，把图进一步简化，适合后续计算。添加Canonicalize优化需要在对应的TOP算子定义的时候设置 hasConstantMaterializer=1，细节可以参考Canonicalization。以下为一个TOP层Op的定义，对应的优化实现为Scale.cpp。

    def Top_ScaleOp: Top_Op<"Scale", [SupportFuseRelu, InOutSameShape]> {
  let summary = "Scale operator";
  let description = [{
    Y = X * S + B,
    where the shape of X/Y is [n, c, h, w] and the shape of S/B is [1, c, 1, 1].
  }];
  let arguments = (ins
    AnyTensor:$input,
    AnyTensor:$scale,
    AnyTensor:$bias,
    DefaultValuedAttr<BoolAttr, "false">:$do_relu,
    DefaultValuedAttr<F64Attr, "-1.0">:$relu_limit
  );
  let results = (outs AnyTensor:$output);
  let hasCanonicalizer = 1;
}

    • TPU-MLIR中大量使用Interface对Operation添加扩展功能。TOP dialect中有效的Interfeace有InferenceInterface，FlopsInterface。如果添加了对应的Interface，那么就需要实现相应的逻辑。在TOP dialect中，由于所有的Op都需要做前向运算，得到中间计算结果与原始模型的推理结果做比较（保证转化的正确性），所以TOP dialect的Op大部分都需要有InferenceInterface。FlopsInterface用于Op计算FLOPS，以便测量模型的计算量和优化性能。下边是TPU-MLIR定义的TOP Op的两个模版，Top_BaseOp用于不需要计算的Op，如Input，Weight和None。Top_Op适合需要计算的Op，如BatchNorm，Conv等。

      class Top_BaseOp<string mnemonic, list<Trait> traits = []> :
    Op<Top_Dialect, mnemonic, !listconcat(traits,[NoSideEffect])> ;

class Top_Op<string mnemonic, list<Trait> traits = []> :
    Top_BaseOp<mnemonic, !listconcat(traits,
       [DeclareOpInterfaceMethods<InferenceInterface>,
        DeclareOpInterfaceMethods<FlopsInterface>])> ;

    • InferenceInterface定义了init，deinit，inference三个函数接口，需要依据具体的算子实现相应的计算。如：Add.cpp。这三个函数都接受InferenceParameter参数，该结构中定义了输入数据，输出数据，以及用于执行计算实现的handle。

      struct InferenceParameter {
  std::vector<float *> inputs;
  std::vector<float *> outputs;
  void *handle = nullptr;
};

      TPU-MLIR使用OneDNN，作为X86计算加速库，所以handle多指向对OneDNN包装的对象。

      LogicalResult top::AddOp::init(InferenceParameter &p) {
  auto binary = new Binary();
  (*binary)
      .lhs(p.inputs[0], Module::getShape(inputs()[0]))
      .rhs(p.inputs[1], Module::getShape(inputs()[1]))
      .dst(p.outputs[0], Module::getShape(output()))
      .do_relu(do_relu())
      .relu_limit(relu_limit().convertToDouble())
      .algorithem(algorithm::binary_add)
      .setup();

  p.handle = (void *)binary;

  3. TOP层网络Lowering到TPU层。TOP dialect定义的是与神经网络框架无关且与硬件无关的抽象计算层，TPU定义与特定硬件有关层。TPU dialect定义的Op都是有对应后端实现的Op，且在TPU dialect还有内存优化，数据分块（tilting）和计算融合（fusion)（TPU-MLIR中称作Layergroup）。

    • Op convertion过程使用了MLIR的DialectConversion机制，同时提供了基础的OpRewritePattern TopLowering基础类，添加对应TOP Op到TPU Op需要在LoweringBM1684X （此处以转化到TPU 1684X为例）中添加对应的宏，然后在/lib/Conversion/TopToTpu/BM1684X添加具体的对应逻辑代码。以下为需要扩展的函数

      public:
  virtual void LoweringINT8(PatternRewriter &rewriter, OpTy opTy,
                            bool asymmetric) const {
    llvm_unreachable("Not Implemented");
  }
  virtual void LoweringBF16(PatternRewriter &rewriter, OpTy opTy) const {
    llvm_unreachable("Not Implemented");
  }
  virtual void LoweringF16(PatternRewriter &rewriter, OpTy opTy) const {
    llvm_unreachable("Not Implemented");
  }
  virtual void LoweringF32(PatternRewriter &rewriter, OpTy opTy) const {
    llvm_unreachable("Not Implemented");
  }
  virtual void LoweringQuantized(PatternRewriter &rewriter, OpTy opTy) const {
    llvm_unreachable("Not Implemented");
  }
};

      对应关系为：

      • LoweringQuantized将quant.calibration类型转化为quant.UniformQuantized等类型；
      • LoweringINT8将TOP Op转化为TPU Op，同时把quant.UniformQuantized转化为可计算的multiplier和shift计算；
      • LoweringBF16将TOP Op转化为TPU Op，同时把float类型转化为BF16类型；
      • LoweringF16 将TOP Op转化为TPU Op，同时把float类型转化为BF16类型；
      • LoweringF32将TOP Op转化为TPU Op，保留Tensor类型。

  4. 当计算使用TPU dialect表示后，就可以依据具体硬件完成内存分配和生成Op的schedule，之后调用后端，完成codegen。和TOP dialect一样，TPU 的Op也定义了一些Interface来完成不同方式的codegen。如：LocalGenInterface完成只使用local memory的计算，其假设输入数据已经存在local memory上，且会保证输出数据也储存在local memory；GlobalGenInterface假设输入数据存储在DDR上，且local memory没有需要保留的之前计算数据（即整个local memory可以随意读写，不用担心数据覆盖），同时输入也存储在DDR上；WeightReorderInterface用来表示相关的constant/coeff/weight Tensor需要针对硬件调整layout。根据对应Op支持的计算情况（是否支持仅使用local memory完成计算）以及图的特点，LayerGroup会把Op做fuse并对数据做title。对应的代码参见LayerGroup。最终的codegen代码可以参看lib/Dialect/Tpu/Transforms/BM168x/Codegen.cpp。整个TPU dialect到最终的codegen，TPU-MLIR会调用多个pass把整个优化过程串联起来，相关的pass定义可以参考include/tpu_mlir/Dialect/Tpu/Transforms/Passes.td。

测试/使用工程

如果你可以完成添加算子中步骤添加一个算子到TPU-MLIR中，那么可以构造一个单元测试，对该算子的正确行进行验证。TPU-MLIR的单元测试使用了python编写，相关代码参见python/test/test_onnx.py。如果发现测试case未覆盖一些边界问题，同样可以添加相应的单元测试。

TPU-MLIR提供了快速入门手册/开发手册来帮助用户快速入门和理解TPU-MLIR工程结构，及其工作原理。

反馈错误

在使用TPU-MLIR中遇到任何问题或是疑问都可以通过在Github上提issue，我们会尽快解答。

推广本工程

如果你觉得本工程有趣，请推荐给更多兴趣相同的同伴。

加入&联系我们

[email protected]

Clone this wiki locally