MLIR技术杂谈

MLIR技术杂谈

MLIR: 编译器基础架构重定义

MLIR（多级中间表示）是语言（如 C）或库（如 TensorFlow）与编译器后端（如 LLVM）之间的中间表示 (IR) 系统。允许不同语言的不同编译器堆栈之间的代码重用以及其他性能和可用性优势。

MLIR 由Google开发为一个开源项目，主要是为了改进 TensorFlow 在不同后端的支持，但通常可用于任何语言。

Buddy-MLIR 详解

0x0. 前言

整个Buddy-MLIR项目给最大感觉就是，无论结果怎么样，都可以先 run 起来。虽然 MLIR 已经出现了几年并且也有一些明星项目比如 IREE 获得了成功。但相比于 TVM 的用例丰富度来说个人感觉还是有一点差距的，特别是在中文社区。这样就造成了一个问题，如果一个人对 MLIR 感兴趣或者要基于 MLIR 从事一些开发工作，就必须要啃 MLIR 的官方文档的 Toy Tutorial 来速成。不否认官方文档十分详尽并且结构组织也比较得当，但对于一个完全新手的用户来说，的确是不那么友好。有没有办法在对 MLIR 相关基础概念进行了解之后就快速进入到MLIR这个世界里提供的组件去构建一个真实的应用呢？

个人认为最近推出的Buddy-MLIR缓解了这一痛点，可以很轻易的跑起来基于 MLIR 做的一个应用然后一边学习 MLIR 的相关概念一边进行魔改来构建自己的应用。Buddy-MLIR 的另外一个亮点在于整个工程的组织结构和 LLVM/MLIR 项目本身一样十分清晰，使得把握整个工程的难度以及阅读相关代码的难度降低了很多。接下来，会从 Run 起来和工程结构解析两方面进行讲解。实际上这种组织结构在OneFlow仓库里的IR部分也是完全一样的，不过由于OneFlow 的计算图和 IR 进行了交互所以目前没有把 IR 部分独立出一个仓库，不然大家会看到 Buddy-MLIR 和 OneFlow-MLIR 的工程结构也是完全一样的。

参考文献链接

https://mp.weixin.qq.com/s/AM1hTcQsgbwG3hCzK6P_gQ

https://mp.weixin.qq.com/s/uE5VhU_s3NgndPk2X6zbAA

https://github.com/buddy-compiler/buddy-mlir

0x1. How to run？

怎么跑起来？这应该是拿到一个项目最重要的问题之一。实际上跟随 Buddy-MLIR 的 README 就可以，不过实际操作的时候还有一些细节需要注意。为了小白用户考虑，这里记录一下在一台Ubuntu20.04的完整编译和Run Buddy-MLIR 的流程。

Buddy-MLIR 项目是基于 LLVM/MLIR 项目扩展的，或者说 LLVM 是 Buddy-MLIR 的一个依赖，所以首先需要安装这个依赖。具体操作过程如下：

$ git 
clone
 git@github.com:buddy-compiler/buddy-mlir.git
$ 
cd
 buddy-mlir
$ git submodule update --init
$ 
cd
 buddy-mlir
$ mkdir llvm/build
$ 
cd
 llvm/build
$ cmake -G Ninja ../llvm \
    -DLLVM_ENABLE_PROJECTS=
"mlir"
 \
    -DLLVM_TARGETS_TO_BUILD=
"host;RISCV"
 \
    -DLLVM_ENABLE_ASSERTIONS=ON \
    -DCMAKE_BUILD_TYPE=RELEASE
$ ninja
$ ninja check-mlir

按照上面的命令操作就可以完成 LLVM 项目的编译，编译结果存放在 llvm/build 文件中。接下来就可以在 Buddy-MLIR 的工程目录下基于 LLVM 编译结果提供的库完成 Buddy-MLIR 本身的编译了。对 Buddy-MLIR 工程编译的如下：

$ cd buddy-mlir
$ mkdir build
$ cd build
$ cmake -G Ninja .. \
    -DMLIR_DIR=$PWD/../llvm/build/lib/cmake/mlir \
    -DLLVM_DIR=$PWD/../llvm/build/lib/cmake/llvm \
    -DLLVM_ENABLE_ASSERTIONS=ON \
    -DCMAKE_BUILD_TYPE=RELEASE
$ ninja check-buddy

编译完成后如果出现类似下面的输出，即 FileCheck 成功则可以证明 Buddy-MLIR 的构建流程已经成功了。

Testing Time: 0.06s
  Passed: 3

在 Buddy-MLIR 开源工程中目前有三种 Dialect 即 Bud Dialect，DIP Dialect 以及 RVV Dialect。从项目相关的介绍目前还不理解 RVV Dialect，所以本文只会涉及 Bud Dialect 以及 DIP Dialect。其中DIP Dialect是为数字图像处理（digital image processing ） 进行的抽象。由于 Buddy-MLIR C/C++ 前端依赖了 OpenCV 来做图片的编解码 ，所以 Buddy-MLIR 引入了 OpenCV 第三方库。如果你没有编译OpenCV 可以使用如下的命令进行编译：

$ sudo apt-get install libgtk2.0-dev pkg-config  libcanberra-gtk-module
$ git 
clone
 https://github.com/opencv/opencv.git
$ 
cd
 opencv && mkdir build && 
cd
 build
$ cmake -D CMAKE_BUILD_TYPE=RELEASE -D CMAKE_INSTALL_PREFIX=/usr/
local
 ..
$ make -j$(nproc)
$ sudo make install

这里可以把 /usr/local 换成任意自定义目录。后续在构建 DIP Dialect 相关的应用时需要指明 -DBUDDY_ENABLE_OPENCV=ON 这个选项启用OpenCV。

接下来看一下 Buddy-MLIR 中提供了哪些有趣的例子。

1. IR 级别的例子

IR 级别示例展示了如何在上游 MLIR 和 Buddy-MLIR 中使用 pass，其中一些示例来自 MLIR 集成测试。 大多数情况可以直接使用 MLIR JIT 引擎 mlir-cpu-runner 运行。 递降管道和工具链配置在 makefile 目标中指定。 可以选择一个感兴趣的 Dialect 并到对应的目录下找到要运行的目标。Buddy-MLIR 中所有的示例都在 https://github.com/buddy-compiler/buddy-mlir/tree/main/examples 这个目录中:

 

 

 Buddy-MLIR 示例分类

点开任意一种 Dialect 示例的MakeFile，可以发现里面主要有三种测试：

  • <Dialect Name>-<Operation Name>-lower 。这个测试用来展示递降管道。会产生log.mlir文件。
  • <Dialect Name>-<Operation Name>-translate。这个测试用来展示从当前Dialect文件产生的LLVM IR。会生成一个log.ll文件。
  • <Dialect Name>-<Operation Name>-run。这个测试会使用MLIR JIT Engine 执行 LLVM IR 产生结果。

以 MemRef Dialect 里面的 memref.dim Op 为例，编译测试方法如下：

$ 
cd
 buddy-mlir/examples/MLIRMemRef
$ make memref-dim-lower
$ make memref-dim-translate
$ make memref-dim-run

原始的 memref.dim 长这样：

func.func @
main
() {
  %c0 = arith.constant 0 : index
  %c1 = arith.constant 1 : index
  %mem0 = memref.alloc() : memref<2x3xf32>
  %mem1 = memref.cast %mem0 : memref<2x3xf32> to memref<?x?xf32>
  %dim0 = memref.dim %mem0, %c0 : memref<2x3xf32>
  %dim1 = memref.dim %mem0, %c1 : memref<2x3xf32>
  %dim2 = memref.dim %mem1, %c0 : memref<?x?xf32>
  %dim3 = memref.dim %mem1, %c1 : memref<?x?xf32>
  vector.print %dim0 : index
  vector.print %dim1 : index
  vector.print %dim2 : index
  vector.print %dim3 : index  
  memref.dealloc %mem0 : memref<2x3xf32>
  func.return
}

使用 JIT Eagine 执行的输出：

2. Convolution Vectorization Examples

Buddy-MLIR 中提供了一个 2D 向量化卷积的 Pass conv-vectorization , 这个 Pass 实现了 Coefficients Broadcasting algorithm with Strip Mining 算法，然后 strip mining size 是可以配置的。这里将其配置为 256 进行演示：

$ 
cd
 buddy-mlir/build/bin
$ ./buddy-opt ../../examples/ConvOpt/conv2d.mlir -conv-vectorization=
"strip-mining=256"

原始的 conv2d.mlir 长这样：

func.func @conv_2d(%arg0: memref<?x?xf32>, %arg1: memref<?x?xf32>, %arg2: memref<?x?xf32>) {
  linalg.conv_2d ins (%arg0, %arg1: memref<?x?xf32>, memref<?x?xf32>)
                 outs (%arg2: memref<?x?xf32>)
  
return
}

经过上面的可执行命令后产生的 MLIR 文件结果如下：

#map0 = affine_map<(d0) -> (d0)>
#map1 = affine_map<(d0) -> (d0 ceildiv 256)>
module {
  func.func @conv_2d(%arg0: memref<?x?xf32>, %arg1: memref<?x?xf32>, %arg2: memref<?x?xf32>) {
    %c0 = arith.constant 0 : index
    %c1 = arith.constant 1 : index
    %c256 = arith.constant 256 : index
    %cst = arith.constant 0.000000e+00 : f32
    %0 = vector.splat %cst : vector<256xf32>
    %1 = memref.dim %arg1, %c0 : memref<?x?xf32>
    %2 = memref.dim %arg1, %c1 : memref<?x?xf32>
    %3 = memref.dim %arg2, %c0 : memref<?x?xf32>
    %4 = memref.dim %arg2, %c1 : memref<?x?xf32>
    affine.for %arg3 = 
#map0(%c0) to #map0(%3) {
      affine.for %arg4 = 
#map0(%c0) to #map0(%1) {
        affine.for %arg5 = 
#map0(%c0) to #map0(%2) {
          affine.for %arg6 = 
#map0(%c0) to #map1(%4) {
            // 
对应下面的步骤1
            %5 = affine.vector_load %arg1[%arg4, %arg5] : memref<?x?xf32>, vector<1xf32>
            %6 = vector.broadcast %5 : vector<1xf32> to vector<256xf32>
            %7 = arith.muli %arg6, %c256 : index
            %8 = arith.subi %4, %7 : index
            %9 = arith.cmpi sge, %8, %c256 : index
            scf.if %9 {
              // 
对应下面的步骤二
              %10 = affine.vector_load %arg0[%arg3 + %arg4, %arg5 + %arg6 * 256] : memref<?x?xf32>, vector<256xf32>
              // 
对应下面的步骤三
              %11 = affine.vector_load %arg2[%arg3, %arg6 * 256] : memref<?x?xf32>, vector<256xf32>
              // 
对应下面的步骤四
              %12 = vector.fma %10, %6, %11 : vector<256xf32>
              // 
对应下面的步骤五
              affine.vector_store %12, %arg2[%arg3, %arg6 * 256] : memref<?x?xf32>, vector<256xf32>
            } 
else
 {
              %10 = vector.create_mask %8 : vector<256xi1>
              %11 = arith.addi %arg3, %arg4 : index
              %12 = arith.muli %arg6, %c256 : index
              %13 = arith.addi %arg5, %12 : index
              %14 = vector.maskedload %arg0[%11, %13], %10, %0 : memref<?x?xf32>, vector<256xi1>, vector<256xf32> into vector<256xf32>
              %15 = vector.maskedload %arg2[%arg3, %12], %10, %0 : memref<?x?xf32>, vector<256xi1>, vector<256xf32> into vector<256xf32>
              %16 = vector.fma %14, %6, %15 : vector<256xf32>
              vector.maskedstore %arg2[%arg3, %12], %10, %16 : memref<?x?xf32>, vector<256xi1>, vector<256xf32>
            }
          }
        }
      }
    }
    
return
  }
}

初步看到这个变换可能还比较懵，结合这个算法和 Pass 实现进行理解一下。

Coefficients broadcasting（CB）算法是 2D 卷积的一种高效实现。Buddy-MLIR 基于 MLIR 基础设施完成了对这个算法的实现。实现这个算法涉及到的 MLIR Dialect 以及 Op 这里列一下：

  • affine.for ：执行指定次数循环体的操作。
  • affine.vector_load：从缓冲区切片中返回一个向量 （MLIR MemRef格式）。
  • affine.vector_store：将一个向量写到缓存区切片中（MLIR MemRef格式）。
  • vector.broadcast：将标量或向量值广播为 N-维 结果向量。
  • vector.fma：向量化类型的乘加混合指令。

然后 CB 算法的过程如下图所示：

 

 

 CB 算法流程

注意输入是一个通道数为 1 的图片或者特征图，然后 kernel 的通道数也是1。算法的执行流程大概为：

  • 首先将 kernel 的每个元素使用 vector_load 加载到缓冲区中 并使用 vector.broadcast 广播到 vector1 中。
  • 然后将特征图的元素使用 vector_load 加载到 vector2 中。
  • 第三步将输出特征图的元素使用 vector_load 加载到 vector3 中。
  • 然后使用 vector.fma 将 vector1 和 vector2 相乘并加到 vector3 上。
  • 最后使用 vector_store 将上述结果写回缓冲区中。

注意，经过 conv-vectorization Pass之后产生的 MLIR 文件中有2个部分。另外一个部分使用了vector.create_mask 和  vector.maskedstore ，这对应了上图中特征图在每一行最后加载的元素字节不够 fma 指令需要的 256Bit （这个256是通过 -conv-vectorization="strip-mining=256" 指定的），所以需要一个 Mask 来补齐然后进行计算。

  • Edge detection example

Buddy-MLIR 还提供了一个边缘检测示例来展示优化。 conv-vectorization pass 负责使用算法递降 linalg.conv_2d。 然后使用 mlir-translate 和 llc 工具生成目标文件。 最后，在 C++ 程序中调用这个 MLIR 卷积函数（在第二节会详细介绍这个过程）。再运行这个示例前需要保证 OpenCV 已经安装好了，安装方法上面介绍了。

这个例子还展示了 AutoConfig 机制的“魔力”，可以帮助末指定strip mining size、ISA SIMD/Vector extension 和 target triple。 您只需要启用 BUDDY_EXAMPLES 选项，无需担心工具链配置。操作命令如下：

$ 
cd
 buddy-mlir/build
$ cmake -G Ninja .. -DBUDDY_EXAMPLES=ON -DBUDDY_ENABLE_OPENCV=ON
$ ninja edge-detection

当然，也可以使用自己的配置值 -DBUDDY_CONV_OPT_STRIP_MINING（例如 64）和 -DBUDDY_OPT_ATTR（例如 avx2）。

仓库提供了一张图片，路径为 buddy-mlir/examples/ConvOpt/images/YuTu.png ，这是构成中国嫦娥三号任务一部分的机器人月球车。然后运行下面的命令对其进行边缘检测。

$ 
cd
 bin
$ ./edge-detection ../../examples/ConvOpt/images/YuTu.png result.png

 

 

 原图

 

 

 边缘检测后的图

3. Digital Image Processing Examples

Buddy-MLIR 还提供了 DIP Dialect 相关的展示例子，具体就是给一张图片进行 Constant Padding 或者 Replicate Padding 然后做卷积 。操作步骤和上面类似，这里就不再展示了。感兴趣的读者可以自行体验。链接：https://github.com/buddy-compiler/buddy-mlir/tree/main/examples#digital-image-processing-examples 。

0x2. How to Understand?

上面一节主要展示了在 Buddy-MLIR 中怎么把构建的应用跑起来，这一节将带大家从 Buddy-MLIR 的结构出发来理解这个工程。工程的整体结构可以总结为：

 

 

 Buddy-MLIR 工程结构

主要把目光放在 include 和  lib 两个文件夹上，其它的文档，测试以及工具类的源码读者可以有选择的进行查看。

2.1 Bud Dialect

从上图可以看出，Buddy-MLIR 主要有三种 Dialect ，即：Bud Dialect，DIP Dialect 还有 RVV Dialect。对于 Dialect 的定义遵循了和 LLVM 上游 Dialect 一样的文件结构和方法，因此这里不再赘述。如果你想了解更多细节可以查看 https://github.com/BBuf/tvm_mlir_learn 仓库中的 MLIR：摩尔定律终结的编译器基础结构 论文解读 文章。

这里主要关注一下 Bud Dialect 中定义了哪些操作。从 buddy-mlir/include/Dialect/Bud/BudOps.td 中可以看到 Bud Dialect 主要定义了 4 种类型的操作：

  • Bud_TestConstantOp。这个 Op 用于测试常量Op。
  • Bud_TestPrintOp 。这个Op 用于测试打印Op。
  • Bud_TestEnumAttrOp 。在 Op 中测试枚举属性。
  • Bud_TestArrayAttrOp 。在 Op 中测试数组属性。

构建了基础操作之后，需要为 Bud Dialect 注册一个递降的 Pipline，也即 lib/Conversion/LowerBud/LowerBudPass.cpp 中实现的 LowerBudPass 。

对 bud::TestConstantOp 的实现如下：

class
 
BudTestConstantLowering
 : 
public
 OpRewritePattern<bud::TestConstantOp> {
public
:
  
using
 OpRewritePattern<bud::TestConstantOp>::OpRewritePattern;
  LogicalResult 
matchAndRewrite
(bud::TestConstantOp op,
                                PatternRewriter &rewriter)
 
const
 
override
 {
    
auto
 loc = op.getLoc();
    
// Get type from the origin operation.
    Type resultType = op.getResult().getType();
    
// Create constant operation.
    Attribute zeroAttr = rewriter.getZeroAttr(resultType);
    Value c0 = rewriter.create<mlir::arith::ConstantOp>(loc, resultType, zeroAttr);
    rewriter.replaceOp(op, c0);
    
return
 success();
  }
};

可以看到在匹配到 bud::TestConstantOp 后会将其重写为 mlir::arith::ConstantOp 。可以在buddy-mlir/examples/BudDialect 下执行 make bud-constant-lower 。得到的结果为：

module
 {
  %i0 = bud.test_constant : i32
}
=>
module
 {
  %c0_i32 = arith.constant 
0
 : i32
}

其它的几个操作类似，都是将 Bud Dialect 定义的几个操作 Lower 到指定的几个上游 Dialect 上。这个 LowerBudPass 的具体实现为：

namespace
 {
class
 
LowerBudPass
 : 
public
 PassWrapper<LowerBudPass, OperationPass<ModuleOp>> {
public
:
  MLIR_DEFINE_EXPLICIT_INTERNAL_INLINE_TYPE_ID(LowerBudPass)
  LowerBudPass() = 
default
;
  LowerBudPass(
const
 LowerBudPass &) {}
  StringRef 
getArgument
()
 
const
 
final
 { 
return
 
"lower-bud"
; }
  StringRef 
getDescription
()
 
const
 
final
 { 
return
 
"Lower Bud Dialect."
; }
  
void
 
runOnOperation
()
 
override
;
  
void
 
getDependentDialects
(DialectRegistry &registry)
 
const
 
override
 {
    
// clang-format off
    registry.insert<
        buddy::bud::BudDialect,
        func::FuncDialect,
        
vector
::VectorDialect,
        memref::MemRefDialect>();
    
// clang-format on
  }
};
} 
// end anonymous namespace.
void
 
LowerBudPass::runOnOperation
()
 {
  MLIRContext *context = &getContext();
  ModuleOp 
module
 = getOperation();
  ConversionTarget 
target
(*context)
;
  
// clang-format off
  target.addLegalDialect<
      arith::ArithmeticDialect,
      func::FuncDialect,
      
vector
::VectorDialect,
      memref::MemRefDialect>();
  
// clang-format on
  target.addLegalOp<ModuleOp, func::FuncOp, func::ReturnOp>();
  RewritePatternSet 
patterns
(context)
;
  populateLowerBudConversionPatterns(patterns);
  
if
 (failed(applyPartialConversion(
module
, target, 
std
::move(patterns))))
    signalPassFailure();
}

可以看到 Bud Dialect 的操作主要会被 Lower 到 arith::ArithmeticDialect, func::FuncDialect, vector::VectorDialect, memref::MemRefDialect 上。

从上面的介绍也可以看到，Buddy Dialect 实际上只是起一个演示作用，也许是教一下初学者怎么快速定义一个新的 Dialect 并接入 MLIR 的生态。

2.2 DIP Dialect

DIP Dialect 是数组图像处理的一个抽象。这里展示一下 DIP Dialect 目前定义的操作。

def DIP_ConstantPadding : I32EnumAttrCase<
"ConstantPadding"
, 0, 
"CONSTANT_PADDING"
>;
def DIP_ReplicatePadding : I32EnumAttrCase<
"ReplicatePadding"
, 1, 
"REPLICATE_PADDING"
>;
def DIP_BoundaryOption : I32EnumAttr<
"BoundaryOption"
,
    
"Specifies desired method of boundary extrapolation during image processing."
,
    [
      DIP_ConstantPadding,
      DIP_ReplicatePadding
    ]>{
  
let
 genSpecializedAttr = 0;
  
let
 cppNamespace = 
"::buddy::dip"
;
}
def DIP_BoundaryOptionAttr : EnumAttr<DIP_Dialect, DIP_BoundaryOption, 
"boundary_option"
>;

DIP Dialect 的 Corr2DOp

 

 DIP Dialect定义了唯一的一个操作 DIP_Corr2DOp ，这个 Op 在做 2D 卷积之前会先将输入进行 Padding 使得卷积后的输出特征图大小和输入一致。这里还涉及到了很多优化技巧，具体体现在 https://github.com/buddy-compiler/buddy-mlir/blob/main/docs/dip-opt.md 这篇文档以及 https://github.com/buddy-compiler/buddy-mlir/blob/main/lib/Conversion/LowerDIP/LowerDIPPass.cpp 这个 Pass 实现中。没有完全理清楚这个算法的逻辑，所以这里就不再讲解这部分，感兴趣的读者可以自行研究。

2.3 Interface

上面介绍了 Buddy-MLIR 项目中定义的两种 Dialect ，这一节需要解答这样一个问题。即，如基于 Buddy-MLIR 构建的算法在 C/C++ 前端中进行调用来实现一个完整的应用程序呢？

为了实现这一目的，Buddy-MLIR 实现了一个为 C/C++ 前端服务的数据结构 MemRef，https://github.com/buddy-compiler/buddy-mlir/blob/main/include/Interface/buddy/core/Container.h 。

// MemRef descriptor.
// - T represents the type of the elements.
// - N represents the number of dimensions.
// - The storage order is NCHW.
template
 <
typename
 T, 
size_t
 N> 
class
 
MemRef
 {
public
:
  
// Constructor from shape.
  MemRef(
intptr_t
 sizes[N], T init = T(
0
));
  
// Constructor from data.
  MemRef(
const
 T *data, 
intptr_t
 sizes[N], 
intptr_t
 offset = 
0
);
  
// Copy constructor.
  MemRef(
const
 MemRef<T, N> &other);
  
// Copy assignment operator.
  MemRef<T, N> &
operator
=(
const
 MemRef<T, N> &other);
  
// Move constructor.
  MemRef(MemRef<T, N> &&other) 
noexcept
;
  
// Move assignment operator.
  MemRef<T, N> &
operator
=(MemRef<T, N> &&other) 
noexcept
;
  
// Desctrutor.
  ~MemRef();
  
// Get the data pointer.
  T *
getData
()
;
  
// Get the sizes (shape).
  
const
 
intptr_t
 *
getSizes
()
 { 
return
 sizes; }
  
// Get the strides.
  
const
 
intptr_t
 *
getStrides
()
 { 
return
 strides; }
  
// Get the rank of the memref.
  
size_t
 
getRank
()
 
const
 { 
return
 N; }
  
// Get the size (number of elements).
  
size_t
 
getSize
()
 
const
 { 
return
 size; }
  
// Get the element at index.
  
const
 T &
operator
[](
size_t
 index) 
const
;
  T &
operator
[](
size_t
 index);
protected
:
  
// Default constructor.
  
// This constructor is desinged for derived domain-specific constructor.
  MemRef() {};
  
// Set the strides.
  
// Computes the strides of the transposed tensor for transpose=true.
  
void
 
setStrides
()
;
  
// Compute the product of array elements.
  
size_t
 
product
(
intptr_t
 sizes[N])
 
const
;
  
// Data.
  
// The `aligned` and `allocated` members point to the same address, `aligned`
  
// member is responsible for handling data, and `allocated` member is
  
// resposible for handling the memory space.
  T *allocated;
  T *aligned;
  
// Offset.
  
intptr_t
 offset = 
0
;
  
// Shape.
  
intptr_t
 sizes[N];
  
// Strides.
  
intptr_t
 strides[N];
  
// Number of elements.
  
size_t
 size;
};

具体的实现在：https://github.com/buddy-compiler/buddy-mlir/blob/main/lib/Interface/core/Container.cpp 。这里主要梳理一下这个自定义的 MemRef 类是如何服务于 C/C++ 前端的。这里以边缘检测为例。核心代码实现如下：

#include 
<iostream>
#include 
<opencv2/imgcodecs.hpp>
#include 
<time.h>
#include 
"Interface/buddy/core/ImageContainer.h"
#include 
"kernels.h"
using
 
namespace
 cv;
using
 
namespace
 
std
;
// Declare the conv2d C interface.
extern
 
"C"
 {
void
 _mlir_ciface_conv_2d(Img<
float
, 
2
> *input, MemRef<
float
, 
2
> *kernel,
                          MemRef<
float
, 
2
> *output);
}
int
 
main
(
int
 argc, 
char
 *argv[])
 {
  
printf
(
"Start processing...\n"
);
  
// Read as grayscale image.
  Mat image = imread(argv[
1
], IMREAD_GRAYSCALE);
  
if
 (image.empty()) {
    
cout
 << 
"Could not read the image: "
 << argv[
1
] << 
endl
;
    
return
 
1
;
  }
  Img<
float
, 2> 
input
(image)
;
  
// Define the kernel.
  
float
 *kernelAlign = laplacianKernelAlign;
  
int
 kernelRows = laplacianKernelRows;
  
int
 kernelCols = laplacianKernelCols;
  
intptr_t
 sizesKernel[
2
] = {kernelRows, kernelCols};
  MemRef<
float
, 2> 
kernel
(kernelAlign, sizesKernel)
;
  
// Define the output.
  
int
 outputRows = image.rows - kernelRows + 
1
;
  
int
 outputCols = image.cols - kernelCols + 
1
;
  
intptr_t
 sizesOutput[
2
] = {outputRows, outputCols};
  MemRef<
float
, 2> 
output
(sizesOutput)
;
  
// Run the convolution and record the time.
  
clock_t
 start, end;
  start = clock();
  
// Call the MLIR conv2d function.
  _mlir_ciface_conv_2d(&input, &kernel, &output);
  end = clock();
  
cout
 << 
"Execution time: "
 << (
double
)(end - start) / CLOCKS_PER_SEC << 
" s"
       << 
endl
;
  
// Define a cv::Mat with the output of the conv2d.
  Mat 
outputImage
(outputRows, outputCols, CV_32FC1, output.getData())
;
  
// Choose a PNG compression level
  
vector
<
int
> compression_params;
  compression_params.push_back(IMWRITE_PNG_COMPRESSION);
  compression_params.push_back(
9
);
  
// Write output to PNG.
  
bool
 result = 
false
;
  
try
 {
    result = imwrite(argv[
2
], outputImage, compression_params);
  } 
catch
 (
const
 cv::Exception &ex) {
    
fprintf
(
stderr
, 
"Exception converting image to PNG format: %s\n"
,
            ex.what());
  }
  
if
 (result)
    
cout
 << 
"Saved PNG file."
 << 
endl
;
  
else
    
cout
 << 
"ERROR: Can't save PNG file."
 << 
endl
;
  
return
 
0
;
}

注意这里的 Img 类的基类也是 MemRef类。

// Image container.
// - T represents the 
type
 of the elements.
// - N represents the number of dimensions.
template <typename T, size_t N> class Img : public MemRef<T, N> {
public:
  Img(cv::Mat image);
};

然后在上面的应用程序中定义了 conv2d Op 的 C 前端函数：

// Declare the conv2d C interface.
extern
 
"C"
 {
void
 _mlir_ciface_conv_2d(Img<
float
, 
2
> *input, MemRef<
float
, 
2
> *kernel,
                          MemRef<
float
, 
2
> *output);
}

这个全局的 C 函数会在执行 buddy-opt 的过程中被翻译成 llvm.call 指令，即 CMakeLists.txt 中这一部分：

add_custom_command(OUTPUT conv2d.o
  COMMAND 
${BUDDY_BINARY_DIR}
/buddy-opt 
${BUDDY_EXAMPLES_DIR}
/ConvOpt/conv2d.mlir -conv-vectorization=
"strip-mining=
${SPLITING_SIZE}
"
 -lower-affine -convert-scf-to-cf -convert-vector-to-llvm -convert-memref-to-llvm -convert-func-to-llvm=
'emit-c-wrappers=1'
 -reconcile-unrealized-casts | 
          
${LLVM_MLIR_BINARY_DIR}
/mlir-translate --mlir-to-llvmir |
          
${LLVM_MLIR_BINARY_DIR}
/llc -mtriple=
${BUDDY_TARGET_TRIPLE}
 -mattr=
${BUDDY_OPT_ATTR}
 --filetype=obj -o 
${BUDDY_BINARY_DIR}
/../examples/ConvOpt/conv2d.o
  DEPENDS buddy-opt)

conv2d 操作的原始 MLIR 文件内容是：

func.func @conv_2d(%arg0: memref<?x?xf32>, %arg1: memref<?x?xf32>, %arg2: memref<?x?xf32>) {
  linalg.conv_2d 
ins
 
(%arg0, %arg1: memref<?x?xf32>, memref<?x?xf32>)
                 
outs
 
(%arg2: memref<?x?xf32>)
  
return
}

在执行 -convert-func-to-llvm='emit-c-wrappers=1' 这个 Pass 时会将上面的 Func Dialect 下的 conv2d 操作翻译为 LLVM IR 并将其包装为一个 llvm.call 指令。这里的详细交互过程在 buddy-mlir/llvm/mlir/docs/TargetLLVMIR.md 这个文档中可以看到，也即 MLIR 提供了一个 C/C++ 的前端接口功能， Buddy-MLIR 应用了这个前端接口功能完成了端到端的应用构建。

上面获得了 LLVM IR，然后从 cmake 的命令可以看到又调用了 LLVM llc命令编译LLVM源文件到用于指定的体系结构的汇编语言。然后，汇编语言输出可以通过本机汇编器和链接器传递，以生成本机可执行文件。这里可以指定执行架构以及一些优化参数等。

0x3. buddy-opt 和 buddy-translate

将上面介绍的实现的 Pass 添加到 MLIR 的上游 Pass 管理机制中就实现了 buddy-opt 工具。

而 buddy-translate 则只扩展了一个 Buddy Dialect 到 LLVMIR 翻译的一项功能。

总的来说，Buddy-MLIR 是入门 MLIR 或者以 MLIR 为基础设施构建自己应用一个比较好的示例。推荐有需要的读者或者开发者学习和探索更多可能。本文没有讲任何 RVV Dialect 的相关知识，因为目前也不太了解，希望后面洪滨可以讲讲这个 Dialect 的动机和细节。

MLIR - 一种新的IR表示和编译器框架

随着深度学习技术的发展，深度学习技术也逐渐从学术研究的方向转向了实践应用的方向，这不仅对深度模型的准确率有了较高的需求，也对深度模型的推理速度有了越来越高的需求。

目前深度模型的推理引擎按照实现方式大体分为两类：

  • 解释型推理引擎：一般包含一个模型解析器和一个模型解释器，一些推理引擎可能还包含一个模型优化器。模型解析器负责读取和解析模型文件，并将其转换为适用于解释器处理的内存格式；模型优化器负责将原始模型变换为等价的、但具有更快的推理速度的模型；模型解释器分析内存格式的模型并接受模型的输入数据，然后根据模型的结构依次执行相应的模型内部的算子，最后产生模型的输出。
  • 编译型推理引擎：一般包含一个模型解析器和一个模型编译器。模型解析器的作用与解释型推理引擎相同；模型编译器负责将模型编译为计算设备（CPU、GPU 等）可直接处理的机器码，并且可能在编译的过程中应用各种优化方法来提高生成的机器码的效率。由于机器码的模型可以直接被计算设备处理而无需额外的解释器的参与，其消除了解释器调度的开销。此外，相对于解释型推理引擎，由于生成机器码的过程更加靠底层，编译器有更多的优化机会以达到更高的执行效率。

由于现在业界对于推理引擎的执行速度有了更高的需求，编译型推理引擎也逐渐成为高速推理引擎的发展方向。目编译型推理引擎有 Apache TVM、oneDNN、PlaidML、TensorFlow XLA、TensorFlow Runtime 等。

为了便于优化，一般来说推理引擎会把模型转换为中间表示，然后对中间表示进行优化和变换，最终生成目标模型（对于解释型推理引擎）或目标机器码（对于编译型推理引擎）。此外，除了深度学习领域，在很早以前编程语言领域就引入了中间表示来做优化和变换。而新的编程语言层出不穷，因此就出现了各种各样的中间表示：

 

 

 

 

 

不同的推理引擎或者编译器都会有自己的中间表示和优化方案，而每种中间表示和优化方案可能都需要从头实现，最终可能会导致软件的碎片化和重复的开发工作。

MLIR 简介

MLIR（Multi-Level Intermediate Representation）是一种新型的用于构建可复用和可扩展的编译器的框架。MLIR 旨在解决软件碎片化、改善异构硬件的编译、降低构建领域特定编译器的成本，以及帮助将现有的编译器连接到一起。

MLIR 旨在成为一种在统一的基础架构中支持多种不同需求的混合中间表示，例如：

  • 表示数据流图（例如在 TensorFlow 中）的能力，包括动态性状、用户可扩展的算子生态系统、TensorFlow 变量等。
  • 在这些图中进行优化和变换（例如在 Grappler 中）。
  • 适合优化的形式的机器学习算子内核的表示。
  • 能够承载跨内核的高性能计算风格的循环优化（融合、循环交换、分块等），并能够变换数据的内存布局。
  • 代码生成“下降”变换，例如 DMA 插入、显式缓存管理、内存分块，以及 1 维和 2 维寄存器架构的向量化。
  • 表示目标特定操作的能力，例如加速器特定的高层操作。
  • 在深度学习图中的做的量化和其他图变换。

MLIR 是一种支持硬件特定操作的通用中间表示。因此，对围绕 MLIR 的基础架构进行的任何投入（例如在编译器 pass 上的工作）都将产生良好的回报；许多目标都可以使用该基础架构，并从中受益。

尽管 MLIR 是一种强大的框架，也有一些非目标。MLIR 不试图去支持底层机器码生成算法（如寄存器分配和指令调度）。这些更适合于底层优化器（例如 LLVM）。此外，MLIR 也不意图成为最终用户写算子内核的源语言（类似于 CUDA 和 C++）。另一方面，MLIR 提供了用于表示此类领域特定语言并将其集成到生态系统中的支柱。

MLIR 在构建时受益于从构建其他中间表示（LLVM IR、XLA HLO 和 Swift SIL）的过程中获得的经验。MLIR 框架鼓励现存的最佳实践，例如：编写和维护中间表示规范、构建中间表示验证器、提供将 MLIR 文件转储和解析为文本的功能、使用 FileCheck 工具编写详尽的单元测试、以及以一组可以以新的方式组合的模块化库的形式构建基础框架。

其他的经验教训也已经整合到了设计中。例如，LLVM 有一个不明显的设计错误，其会阻止多线程编译器同时处理 LLVM 模块中的多个函数。MLIR 通过限制 SSA 作用域来减少使用-定义链，并用显式的符号引用代替跨函数引用来解决这些问题。

MLIR 方言（Dialect）

MLIR 通过“方言”来定义不同层次的中间表示，每一个方言都有自己唯一的名字空间。开发者可以创建自定义方言，并在方言内部定义操作、类型和属性，以及语义。MLIR 推荐使用方言来对 MLIR 进行扩展。有这样一个统一的中间表示框架降低了开发新的编译器的成本。除了可以使用 C++ 语言对方言进行定义之外，MLIR 也提供了一种声明式的方式来定义方言，即用户通过编写 TableGen 格式的文件来定义方言，然后使用 TableGen 工具来生成对应的 C++ 头文件和源文件，以及对应的文档。MLIR 也推荐使用这种声明式的的方式来定义方言。此外，MLIR 也提供了一个框架用于在方言之间或者方言内部进行转换。

为了方便开发，MLIR 也内置了一些方言可供直接使用：

  • acc
  • affine
  • async
  • avx512
  • gpu
  • linalg
  • llvm
  • nvvm
  • omp
  • pdl
  • pdl_interp
  • quant
  • rocdl
  • scf
  • shape
  • spv
  • std
  • vector

MLIR 使用“操作”来描述不同层次的抽象和计算。MLIR 中的操作也是可扩展的，用户可以创建自定义的操作并规定其语义。例如目标无关操作、仿射操作和目标特定操作。MLIR 也支持用户通过声明式的方式（TableGen）来创建自定义操作。

MLIR 中的每个值都有其对应的“类型”，MLIR 内置了一些原始类型（比如整数）和聚合类型（张量和内存缓冲区）。MLIR 的类型系统也允许用户对其进行扩展，创建自定义的类型以及规定其语义。

此外在 MLIR 中，用户可以通过指定操作的“属性”的值来控制操作的行为。操作可以定义自身的属性，比如卷积操作的 stride 属性等。

方言的变换

在 MLIR 中定义操作的时候可以定义其规范化的行为，比如将 x + 2 和 2 + x 统一规范化为 x + 2，以便后续的优化过程更为方便地进行。MLIR 以一种贪婪地策略，不断地应用规范化变换，直到中间表示收敛为止。

在 MLIR 中进行方言内部或方言之间的转换时，用户首先要定义一个转换目标。转换目标规定了生成的目标中可以出现哪些操作。然后用户需要指定一组重写模式，这些重写模式定义了操作之间的转换关系。最后框架根据用户指定的转换目标和重写模式执行转换。这个转换过程会自动检测转换方式，例如如果指定了 A → B 和 B → C 的重写模式，框架会自动完成 A → C 的转换过程。MLIR 也支持用户通过声明式的方式（TableGen）来创建自定义的重写模式。当转换的方言之间有着不同的类型系统，用户可以使用类型转换器来完成类型之间的转换。

MLIR 的用户

  • ONNX MLIR：将 ONNX 格式的深度学习网络模型转换为能在不同二进制执行的二进制格式。
  • PlaidML：一个开源的张量编译器，允许在不同的硬件平台上运行深度学习模型。
  • TensorFlow：TensorFlow 项目的 XLA 和 TensorFlow Lite 模型转换器用到了 MLIR。
  • TensorFlow Runtime：一种新的 TensorFlow 运行时。
  • Verona：一种新的研究型的编程语言，用于探索并发所有权。其提供了一个可以与所有权无缝集成新的并发模型。

结论

MLIR 是一种新型的编译器框架，其设计从已有的编译器的实现中吸取了经验和教训，包括了中间表示的定义、转换以及优化等功能，极大地方便了新的编译器的开发和调试工作。同时，MLIR 也包含了很多现成的工具可直接使用（batteries included）。MLIR 包揽了编译器设计中的通用部分，使得编译器的开发人员可以专注于核心的语义分析、中间表示的设计和变换，以此降低开发成本，提高开发效率和提高成品质量。

外部链接

  • MLIR 主页：https://mlir.llvm.org/。
  • MLIR 语言参考：https://mlir.llvm.org/docs/LangRef/。

附录：编译和安装 MLIR

下载 MLIR

MLIR 是 LLVM 项目的子项目，要编译 MLIR，首先获取 LLVM 的源代码。

LLVM 的源码可从 GitHub 获取：

git 
clone
 https://github.com/llvm/llvm-project.git

用户也可以直接下载源码包：https://github.com/llvm/llvm-project/releases。

假定 LLVM 的源码目录为 $LLVM_SRC。

编译 MLIR

首先用户需要指定一个路径用于存放编译中间产物，假定其路径为 $LLVM_BUILD。然后使用下列命令对 LLVM 进行配置：

cmake -S 
"
$LLVM_SRC
"
 -B 
"
$LLVM_BUILD
"
 -DLLVM_ENABLE_PROJECTS=mlir -DCMAKE_BUILD_TYPE=Release

默认情况下，LLVM 禁用了异常处理和运行时类型信息。如果应用程序需要依赖这些功能，可指定在配置时指定 LLVM_ENABLE_EH 和 LLVM_ENABLE_RTTI CMake 变量的值为 ON：

cmake -S 
"
$LLVM_SRC
"
 -B 
"
$LLVM_BUILD
"
 -DLLVM_ENABLE_PROJECTS=mlir -DLLVM_ENABLE_EH=ON -DLLVM_ENABLE_RTTI=ON -DCMAKE_BUILD_TYPE=Release

更多的 LLVM 配置参数参见 https://llvm.org/docs/CMake.html。

执行完配置过程后使用下列命令执行编译：

cmake --build 
"
$LLVM_BUILD
"

安装 MLIR

使用如下命令将 LLVM 安装到 /usr/local 目录：

cmake --install 
"
$LLVM_BUILD
"

如果想指定另外一个安装目录，例如 $INSTALL_DIR，可以使用 --prefix 命令行参数来指定：

cmake --install 
"
$LLVM_BUILD
"
 --prefix 
"
$INSTALL_DIR
"

在 CMake 项目中使用 MLIR

用户可以在 CMake 项目文件中使用下列语句添加查找 MLIR 依赖：

find_package(MLIR REQUIRED CONFIG)

如果 MLIR 被安装到了系统目录（比如 /、/usr、/usr/local 等），CMake 无需额外的配置就能找到 MLIR；如果 MLIR 被安装到了非系统目录，可以在 CMake 的配置过程通过 CMake 的 MLIR_DIR 变量来指定 MLIR 的安装位置：

cmake 
"
$MY_PROJECT_DIR
"
 -DMLIR_DIR=
"
$INSTALL_DIR
"

成功之后用户可以直接使用 MLIR 的库作为编译目标的依赖：

add_executable(my-executable main.cpp)
target_include_directories(my-executable SYSTEM PRIVATE ${MLIR_INCLUDE_DIRS})
target_link_libraries(my-executable PRIVATE MLIRIR)

其中 MLIR_INCLUDE_DIRS 是自动生成的变量，其指向 MLIR 的包含目录。

在使用 CMake 定义可执行文件目标时，如果 LLVM 禁用了运行时类型信息，那么依赖于 LLVM 的可执行文件目标也需要禁用运行时类型信息，否则可能会编译失败。LLVM 提供了一个 CMake 帮助函数 llvm_update_compile_flags 可以自动完成这个配置。这个函数定义在 LLVM 提供的 AddLLVM.cmake 文件中。用户可以使用下列语句导入 AddLLVM.cmake 文件：

list(APPEND CMAKE_MODULE_PATH "${LLVM_CMAKE_DIR}")
include(AddLLVM)

导入 AddLLVM.cmake 文件后就可以对编译目标进行配置了：

llvm_update_compile_flags(my-executable)

完整的 CMake 项目文件示例如下：

cmake_minimum_required(VERSION 3.15)
project(my-executable)
find_package(MLIR REQUIRED CONFIG)
list(APPEND CMAKE_MODULE_PATH "${LLVM_CMAKE_DIR}")
include(AddLLVM)
add_executable(my-executable main.cpp)
target_include_directories(my-executable SYSTEM PRIVATE ${MLIR_INCLUDE_DIRS})
target_link_libraries(my-executable PRIVATE MLIRIR)
llvm_update_compile_flags(my-executable)

 

参考文献链接

https://mp.weixin.qq.com/s/AM1hTcQsgbwG3hCzK6P_gQ

https://mp.weixin.qq.com/s/uE5VhU_s3NgndPk2X6zbAA

https://github.com/buddy-compiler/buddy-mlir

 

 

 

 

posted @   吴建明wujianming  阅读(734)  评论(0)  编辑  收藏  举报
相关博文：
· MLIR中间表示和编译器框架
· MLIR一些背景知识
· 从零开始教你写一个MLIR Pass
· 2、[MLIR] 转换流程详解(以Toy接入为例)
· 【源码研读】MLIR Dialect 分层设计
阅读排行：
· 全程不用写代码，我用AI程序员写了一个飞机大战
· DeepSeek 开源周回顾「GitHub 热点速览」
· 记一次.NET内存居高不下排查解决与启示
· MongoDB 8.0这个新功能碉堡了，比商业数据库还牛
· .NET10 - 预览版1新功能体验（一）
历史上的今天：
2021-06-08 GPU特征处理技术
2020-06-08 TOF摄像机可以替代Flash激光雷达吗？
2020-06-08 毫米波雷达分类和技术方案
2020-06-08 高精地图与自动驾驶（下）
2020-06-08 高精地图与自动驾驶（上）
2020-06-08 Linux内存技术分析（下）
2020-06-08 Linux内存技术分析（上）
点击右上角即可分享
微信分享提示