MLIR入门教程-Toy扩展实践-1-添加一个新Op

hunterzju

发布于 2021-12-09 14:02:08

1.3K0

学习完了MLIR入门教程后，总觉得需要实践练练手才有感觉。之前折腾了把Toy教程从代码仓库里拿出来单独编译，就想着自己折腾下向toyDialect中添加新的op，熟悉下开发流程。这里项Toy Dialect中添加一个OrOp，支持按照Tensor元素执行逻辑或操作。

代码在github上管理：https://github.com/hunterzju/llvm-tutorial

def main() {
  # Define a variable `a` with shape <2, 3>, initialized with the literal value.
  # The shape is inferred from the supplied literal.
  var a = [[1, 2, 3], [4, 5, 6]];
  # b is identical to a, the literal array is implicitly reshaped: defining new
  # variables is the way to reshape arrays (element count in literal must match
  # the size of specified shape).
  var b<2, 3> = [1, 2, 3, 4, 5, 6];

  # add a new operation Or
  var c = a | b;
}

源码到AST

为toy语言新加入一个Op支持，首先要能够支持解析为AST；需要经过词法分析lexer和语法分析parser两个过程；

lexer支持:

lexer实现在mlir/mycode/Ch2/include/toy/Lexer.h中，新加入的操作符|在词法分析阶段被当作Identifier处理，并不需要新添加支持。

parser支持：

parser实现在mlir/mycode/Ch2/include/toy/Parser.h中，新加入的|操作需要被解析为BinaryOp, 该操作在parseBinOpRHS()中实现，需要在getTokPrecedence()中加入对|操作符的支持，参考c语言运算符优先级，|操作优先级低于+,-操作：

int getTokPrecedence() {
    // ...
    // 1 is lowest precedence.
    switch (static_cast<char>(lexer.getCurToken())) {
    case '|':
      return 10;
    case '-':
      return 20;
    case '+':
      return 20;
    case '*':
      return 40;
    default:
      return -1;
    }
  }

AST支持：

ast实现在mlir/mycode/Ch2/include/toy/AST.h，新加入的|操作需要被解析为AST中的二元操作符BinaryExprAST，新加入的|操作和+等二元操作符在仅是AST节点的操作符op有区别，该部分实现也不需要做改动。

MLIR Op构建

MLIR是一种图类型的IR表示，核心是节点Node和边Edge。在MLIR中节点是Operation，边是Values，这里的value可以是Operation的结果或者Block的参数（Block通常是多个不含分支的Operation构成）。新加的操作最终需要构建为MLIR中的一个Operation节点。关于MLIR更详细的内容可以参看官网的LanguageReference。

Image

MLIR提供了一套基于tablegen的Op实现框架ODS。由于MLIR支持各种自定义的Dialect，如果各种Dialect的构造器和参数等内容缺乏一些共识，会导致碎片化非常严重，Dialect转换成本会很高。采用表驱动（tablegen）的方式定义Dialect的核心内容，自动化生成接口代码可以减少工作量，同时达到一定程度“标准化”的效果。

Image

ODS框架中提供了一套自己的语法来定义Dialect和Op，可以通过tablegen生成对应类的C++代码。定义OP主要需要定义的内容包括arguments、results、verifier、builders、doc等。参考AddOp的实现，添加OrOp操作也很简单：

def OrOp : Toy_Op<"Or"> {
  let summary = "element-wise logic Or operation";
  let description = [{
    The "or" operation performs element-wise logic or between two tensors.
    The shapes of the tensor operands are expected to match.
  }];

  let arguments = (ins F64Tensor: $lhs, F64Tensor:$rhs);
  let results = (outs F64Tensor);

  // 构建一个MLIR Operation节点
  let builders = [
    OpBuilder<(ins "Value":$lhs, "Value":$rhs)>
  ];
}

一个Op的定义：通过arguments指定输入，通过results指定输出，然后可以通过builders来指定构建一个MLIR Operation节点的方法。这里以argument的构建为例，说明各个表项的定义：

let arguments = (ins
  <type-constraint>:$<operand-name>,
  ...
  <attr-constraint>:$<attr-name>,
  ...
);

上文中的F64Tensor对应<type-constraint>，

到此就可以实现将toy语言中的|操作符对应为MLIR中ToyDialect的一个Operation节点，可以编译验证一下：

# 编译
cmake --build . --target toyc-ch2
# 验证
./bin/toyc-ch2 ../../testcode/Ch2/ast.toy --emit=mlir

说明:AST解析支持了减法操作，但是目前toy对应的MLIR Op中并没有定义减法Op，感兴趣可以把支持减法当作练习。

本文参与腾讯云自媒体同步曝光计划，分享自作者个人站点/博客。

如有侵权请联系 cloudcommunity@tencent.com 删除

编程算法

github

git

开源

本文分享自作者个人站点/博客前往查看

如有侵权，请联系 cloudcommunity@tencent.com 删除。

本文参与腾讯云自媒体同步曝光计划，欢迎热爱写作的你一起参与！

登录后参与评论

0 条评论

热度