【从零开始学深度学习编译器】十八，MLIR中的Interfaces

最新推荐文章于 2025-08-30 17:55:41 发布

原创

最新推荐文章于 2025-08-30 17:55:41 发布 · 878 阅读

0 ·

CC 4.0 BY-SA版权

文章标签：

#深度学习 #人工智能

0x0. 前言

这篇文章用来了解一下MLIR中的Interfaces（接口）。MLIR是一个通用可扩展的框架，由不同层次的具有特定属性，Operation以及Type的Dialects构成。正是由于Dialects的分层设计, 使得MLIR可以表达多种语意和抽象级别的Operation。但这个分级设计也存在一个缺点，那就是在不同的Dialect层次进行Operation转换或者做变换（Pass）的时候我们需要明确每个Dialect下的每个Operation的具体语意，否则就可能会转换或变换失败。其实基于MLIR开发过的读者应该碰到过组合一些MLIR Pass对一个MLIR文件进行Lower的时候，有可能出现Op转换失败的情况。为了缓解这种情况，MLIR提出了Interfaces。实际上在【从零开始学深度学习编译器】十三，如何在MLIR里面写Pass？这里我们已经利用过Interfaces来实现内联以及形状推导Pass了。这一节就更深入的了解一下MLIR中的Interfaces，最后还结合了OneFlow IR中的UserOpCompatibleInterface例子来进一步加深了解。

本文提到的Operation和操作是一个东西，都是MLIR Dialect下的操作。

0x1. 动机

Interfaces可以翻译成接口，MLIR的Interfaces提供了和IR交互的通用方式。Interfaces的设计目标是可以不用侵入到具体某个Dialect下的特定Operation和Dialect的特定知识就达到可以转换和分析MLIR表达式。这样就可以将转换，分析和新增一个Dialect和对应的Operation 进行解耦，大大增强MLIR的可扩展性。

0x2. Dialect Interfaces定义（细看）

Dialect Interfaces一般用在想对一组属性，Operation，类型进行通用的转换（Pass）或者分析，这些属性，Operation，类型可以是由不同的Dialect定义的。这些Interfaces一般会广泛覆盖各个级别的Dialects，仅用于少数分析和变换。因此，我们要明确Interface并不是Operation的核心，而是一些通用变换的核心。在【从零开始学深度学习编译器】十三，如何在MLIR里面写Pass？这里有一个使用内联Interface实现内联Pass的例子。内联通常查询的是有关Dialect中Operation的高级信息，例如cost modeling和合法性，而这些信息通常不特定于某个Dialect下的某个Operation单独存在。

Dialect Interface可以通过继承一个CRTP基类DialectInterfaceBase::Base<>来进行定义。CRTP的介绍可以参考：https://zh.wikipedia.org/wiki/奇异递归模板模式，我理解静态多态（CRTP）是因为MLIR里面会存在很多Dialect Interface要从这个DialectInterfaceBase::Base<>基类派生出来，为了性能考虑用CRTP比较合适。这个基类提供了Dialect Interface注册必须的一些接口，方便将来引用它们。当Interface被定义之后，Dialects就可以使用Dialect特定信息去重写它。被一个Dialect定义的Interfaces通过addInterfaces<>进行注册，和属性，Operation，Type的注册机制类似。下面举一个栗子：

// 定义一个基础的内联Interface类以允许Dialect选择加入内联。 
class DialectInlinerInterface :
    public DialectInterface::Base<DialectInlinerInterface> {
   
   
public:
  /// 如果给定的区域 'src' 可以内联到该区域中，则返回 true。
  /// 'dest' 附加到注册到当前Dialect的Operation上。 
  /// 'valueMapping' 包含来自 'src' 区域内的任何重新映射的值。 
  /// 例如，这可用于检查哪些值将替换“src”区域中的条目参数。 
  virtual bool isLegalToInline(Region *dest, Region *src,
                               BlockAndValueMapping &valueMapping) const {
   
   
    return false;
  }
};

/// 覆盖内联接口以添加对 AffineDialect 的支持以启用内联Affine Dialect的Operation。 
struct AffineInlinerInterface : public DialectInlinerInterface {
   
   
  /// Affine结构具有特定的内联约束。 
  bool isLegalToInline(Region *dest, Region *src,
                       BlockAndValueMapping &valueMapping) const final {
   
   
    ...
  }
};

/// 在Dialect下注册内联Interfaces
AffineDialect::AffineDialect(MLIRContext *context) ... {
   
   
  addInterfaces<AffineInlinerInterface>();
}

这些Interfaces被注册之后，在执行MLIR的变换和分析时就可以从Dialect中查到，并不需要确定特定的Dialect子类（如具体到某个Operation）。例如：

Dialect *dialect = ...;
if (DialectInlinerInterface *interface
      = dialect->getRegisteredInterface<DialectInlinerInterface>()) {
   
   
  // Dialect提供了这个Interface的实现
  ...
}

例如，在llvm/mlir/lib/IR/Dialect.cpp这个文件中的registerDelayedInterfaces函数就展示了上面这种用法，这个函数用于注册加载进来的Dialect的Interfaces：

void DialectRegistry::registerDelayedInterfaces(Dialect *dialect) const {
   
   
  auto it = interfaces.find(dialect->getTypeID());
  if (it == interfaces.end())
    return;

  // Add an interface if it is not already present.
  for (const auto &kvp : it->getSecond().dialectInterfaces) {
   
   
    if (dialect->getRegisteredInterface(kvp.first))
      continue;
    dialect->addInterface(kvp.second(dialect));
  }

  // Add attribute, operation and type interfaces.
  for (const auto &info : it->getSecond().objectInterfaces)
    std::get<2>(info)(dialect->getContext());
}

0x3. DialectInterfaceCollection（选看，我还没用过）

DialectInterfaceCollection提供了一个额外的实用程序。这个类允许收集已在 MLIRContext 实例中注册给定Interface的所有Dialect。这对于隐藏和优化已注册Dialect Interface的查找很有用。

class InlinerInterface : public
    DialectInterfaceCollection<DialectInlinerInterface> {
   
   
  // 此类的钩子是DialectInlinerInterface的钩子镜像，默认实现为调用给定Dialect上Interface上的钩子。 
  virtual bool isLegalToInline(Region *dest, Region *src,
                               BlockAndValueMapping &valueMapping) const {
   
   
    auto *handler = getInterfaceFor(dest->getContainingOp());
    return handler ? handler->isLegalToInline(dest, src, valueMapping) : false;
  }
};

MLIRContext *ctx = ...;
InlinerInterface interface(ctx);
if(!interface.isLegalToInline(...))
   ...

0x4. 属性，操作，类型 Interfaces（选看）

顾名思义，属性/操作/类型Interface是在特定属性/操作/类型级别注册的那些。这些Interface 通过提供必须实现的虚接口来提供对派生对象的访问。例如，许多分析和转换想要知道Operation的副作用以提高性能和正确性。 Operation的副作用通常与特定Operation的语义相关，例如 affine.load Operation具有读取效果（顾名思义）。

这些Interface是通过覆盖特定 IR 实体的 CRTP 类来定义的；分别是 AttrInterface、OpInterface 或 TypeInterface。这些类将定义Concept和Model类的 Traits 类作为模板参数。这些类提供了基于概念的多态性的实现，其中Concept定义了一组虚方法，这些方法被在具体实体类型上模板化的Model覆盖。需要注意的是，这些类应该是纯的，不应包含非静态数据成员或其他可变数据。为了将Interface附加到对象，基类提供了一个可以附加到该对象的特征列表的 Trait 类（跳过下面的示例代码就可以看到解释）。

最低0.47元/天解锁文章