深入解析 MLIR Toy Tutorial(Chapter 4):

概述

MLIR Toy Tutorial 的目标是通过构建一门编程语言编译器的完整过程(包括前端和后端技术),教授如何使用 MLIR 的各个组件来实现语言的解析、转换和代码生成等功能。

Chapter3 介绍了如何在 Canonicalizer pass 上应用自定义的 rewrite pattern 来重写优化 IR。回顾一下相关开发流程:

  • 自定义继承于 mlir::OpRewritePattern 的 RewritePattern,在 MatchAndRewrite() 中定义匹配和重写的规则。
  • 在 op 定义(Ops.td)中声明 let hasCanonicalizer = 1;,并实现 getCanonicalizationPatterns() 方法,用它来调用那些自定义的 RewritePattern。
  • 通过 mlir::createCanonicalizerPass() 创建 Canonicalizer pass,pass 会遍历所有 op,并调用它们的钩子(hook)函数 -- getCanonicalizationPatterns() 来重写 IR。

总之,Canonicalizer pass 通过提供 getCanonicalizationPatterns() 这个 hook(也就是回调函数)可以让各 op 灵活滴实现自定义操作。将这个设计进一步发扬光大就成了 Chapter4 要介绍的技术:接口(interfaces)。

这里说的接口和面向对象里的接口概念差不多,是一种将一组操作组织在一起的抽象机制,使得不同的实体可以共享相同的操作集合,从而实现代码的重用和灵活性。MLIR 为 pass、dialect 和 op 都提供有各种操作接口,用户只为这些组件实现相应的接口方法就可以应用这些操作。

Chapter4 以形状推断(shape inference)为例,介绍如何使用接口。

Shape Inference

形状推断是编译时的重要环节,它可以在运行时根据输入张量构建出一张静态计算图,用于后续如常量折叠、剪枝等图优化操作。它通常需要遍历所有的 op,根据 op 的输入张量的类型(包括形状和数据类型)来推断输出张量的类型,因此,每个 op 都要对外提供一个能做形状推断的方法,在Ch4,这个方法是 inferShapes()

Ch4 创建了一个 pass -- ShapeInferencePass 来遍历所有函数的 ops,找出会输出动态形状的 op,然后调用它们的 inferShapes() 来做形状推断。出于简单和教学的需要,Ch4 会先把所有函数中都内联(inlining)到 main 函数,这样 pass 只需要处理一个函数即可。

Inlining

内联函数是一个常用的操作,MLIR 自然也已经提供有相关功能模块,允许不同方言(dialect)之间的 ops 进行内联。Toy dialect 只要应用相关接口 -- DialectInlinerInterface,就能复用这些功能。

要应用 DialectInlinerInterface 接口,除了需要在 dialect 里声明使用该接口,还需要实现它的方法:isLegalToInline()、handleTerminator() 和 materializeCallConversion() 等:

/// Dialect initialization, the instance will be owned by the context. This is
/// the point of registration of types and operations for the dialect.
void ToyDialect::initialize() {
  ...
  addInterfaces<ToyInlinerInterface>();
}

/// This class defines the interface for handling inlining with Toy
/// operations.
struct ToyInlinerInterface : public DialectInlinerInterface {
    bool isLegalToInline(Operation *call, Operation *callable,
                       bool wouldBeCloned) const final {
    return true;
  }

  ...

  Operation *materializeCallConversion(OpBuilder &builder, Value input,
                                       Type resultType,
                                       Location conversionLoc) const final {
    return builder.create<CastOp>(conversionLoc, resultType, input);
  }
};
  • isLegalToInline() 用于判断两组 ops 做内联是否合法。
  • handleTerminator() 用来处理 toy.return op。
  • materializeCallConversion() 用来函数形参和实参的类型转换。

除此之外,由于内联函数需要确定被内联的对象(Callable)和 ops插入的位置(Caller),MLIR 也提供了相关接口 -- CallOpInterface/CallableOpInterface,来定义 ops 的角色。另外,被内联的函数存在形参和实参类型不匹配的情况,需要引入 CastOp 来对它们做转换。

内联部分涉及的内容较多,由于篇幅所限,这里不做过多展开,详情请看原文

ShapeInferenceOpInterface

现在我们已经内联了所有的函数,剩下的就只有主函数,它包含静态和动态形状的操作:

toy.func @main() {
  %0 = toy.constant dense<[[1.000000e+00, 2.000000e+00, 3.000000e+00], [4.000000e+00, 5.000000e+00, 6.000000e+00]]> : tensor<2x3xf64>
  %1 = toy.constant dense<[[1.000000e+00, 2.000000e+00, 3.000000e+00], [4.000000e+00, 5.000000e+00, 6.000000e+00]]> : tensor<2x3xf64>
  %2 = toy.cast %1 : tensor<2x3xf64> to tensor<*xf64>
  %3 = toy.cast %0 : tensor<2x3xf64> to tensor<*xf64>
  %4 = toy.transpose(%2 : tensor<*xf64>) to tensor<*xf64>
  %5 = toy.transpose(%3 : tensor<*xf64>) to tensor<*xf64>
  %6 = toy.mul %4, %5 : tensor<*xf64>
  toy.print %6 : tensor<*xf64>
  toy.return
}

前面说过,ShapeInferencePass 会遍历 main 函数的所有含有动态形状(tensor<*xf64>)的 op,然后调用它们的 inferShapes 方法来推断出静态形状(如 tensor<2x3xf64>)。

Ch4 定义了 ShapeInferenceOpInterface 接口,这个接口会提供 inferShapes 方法,这样 Ops 只要应用该接口并实现 inferShapes() 即可:

def ShapeInferenceOpInterface : OpInterface<"ShapeInference"> {
  ...

  let methods = [
    InterfaceMethod<"Infer and set the output shape for the current operation.",
                    "void", "inferShapes">
  ];
}

def MulOp : Toy_Op<"mul",
    [..., DeclareOpInterfaceMethods<ShapeInferenceOpInterface>]> {
  ...
}

void MulOp::inferShapes() { getResult().setType(getLhs().getType()); }

这里,MulOp::inferShapes() 会将它的输出张量的数据类型和形状设置成和输入张量一样。

总结

Chapter4 以内联函数和形状推断为例,介绍了 MLIR 的接口技术的使用方法。MLIR 中的接口是一种抽象机制,用于定义和共享操作集合,促进代码的重用和模块化,从而提高代码的灵活性和可维护性。 MLIR pass、dialect 和 op 都可以定义和使用接口。

END

`

©著作权归作者所有,转载或内容合作请联系作者
  • 序言:七十年代末,一起剥皮案震惊了整个滨河市,随后出现的几起案子,更是在滨河造成了极大的恐慌,老刑警刘岩,带你破解...
    沈念sama阅读 203,547评论 6 477
  • 序言:滨河连续发生了三起死亡事件,死亡现场离奇诡异,居然都是意外死亡,警方通过查阅死者的电脑和手机,发现死者居然都...
    沈念sama阅读 85,399评论 2 381
  • 文/潘晓璐 我一进店门,熙熙楼的掌柜王于贵愁眉苦脸地迎上来,“玉大人,你说我怎么就摊上这事。” “怎么了?”我有些...
    开封第一讲书人阅读 150,428评论 0 337
  • 文/不坏的土叔 我叫张陵,是天一观的道长。 经常有香客问我,道长,这世上最难降的妖魔是什么? 我笑而不...
    开封第一讲书人阅读 54,599评论 1 274
  • 正文 为了忘掉前任,我火速办了婚礼,结果婚礼上,老公的妹妹穿的比我还像新娘。我一直安慰自己,他们只是感情好,可当我...
    茶点故事阅读 63,612评论 5 365
  • 文/花漫 我一把揭开白布。 她就那样静静地躺着,像睡着了一般。 火红的嫁衣衬着肌肤如雪。 梳的纹丝不乱的头发上,一...
    开封第一讲书人阅读 48,577评论 1 281
  • 那天,我揣着相机与录音,去河边找鬼。 笑死,一个胖子当着我的面吹牛,可吹牛的内容都是我干的。 我是一名探鬼主播,决...
    沈念sama阅读 37,941评论 3 395
  • 文/苍兰香墨 我猛地睁开眼,长吁一口气:“原来是场噩梦啊……” “哼!你这毒妇竟也来了?” 一声冷哼从身侧响起,我...
    开封第一讲书人阅读 36,603评论 0 258
  • 序言:老挝万荣一对情侣失踪,失踪者是张志新(化名)和其女友刘颖,没想到半个月后,有当地人在树林里发现了一具尸体,经...
    沈念sama阅读 40,852评论 1 297
  • 正文 独居荒郊野岭守林人离奇死亡,尸身上长有42处带血的脓包…… 初始之章·张勋 以下内容为张勋视角 年9月15日...
    茶点故事阅读 35,605评论 2 321
  • 正文 我和宋清朗相恋三年,在试婚纱的时候发现自己被绿了。 大学时的朋友给我发了我未婚夫和他白月光在一起吃饭的照片。...
    茶点故事阅读 37,693评论 1 329
  • 序言:一个原本活蹦乱跳的男人离奇死亡,死状恐怖,灵堂内的尸体忽然破棺而出,到底是诈尸还是另有隐情,我是刑警宁泽,带...
    沈念sama阅读 33,375评论 4 318
  • 正文 年R本政府宣布,位于F岛的核电站,受9级特大地震影响,放射性物质发生泄漏。R本人自食恶果不足惜,却给世界环境...
    茶点故事阅读 38,955评论 3 307
  • 文/蒙蒙 一、第九天 我趴在偏房一处隐蔽的房顶上张望。 院中可真热闹,春花似锦、人声如沸。这庄子的主人今日做“春日...
    开封第一讲书人阅读 29,936评论 0 19
  • 文/苍兰香墨 我抬头看了看天上的太阳。三九已至,却和暖如春,着一层夹袄步出监牢的瞬间,已是汗流浃背。 一阵脚步声响...
    开封第一讲书人阅读 31,172评论 1 259
  • 我被黑心中介骗来泰国打工, 没想到刚下飞机就差点儿被人妖公主榨干…… 1. 我叫王不留,地道东北人。 一个月前我还...
    沈念sama阅读 43,970评论 2 349
  • 正文 我出身青楼,却偏偏与公主长得像,于是被迫代替她去往敌国和亲。 传闻我的和亲对象是个残疾皇子,可洞房花烛夜当晚...
    茶点故事阅读 42,414评论 2 342

推荐阅读更多精彩内容