MLIR-Tutorials-Toy Tutorial

玩具语言教程(Toy Tutorial)

第七章：向玩具语言增加一个复合类型(Chapter 7: Adding a Composite Type to Toy)

总结

  在前一章中，我们展示了从Toy语言的编译前端到LLVM IR的端到端编译流。在本章中，我们将扩展Toy语言以支持新的复合结构类型(struct)。

在Toy中定义一个struct(Defining a struct in Toy)

  首先，我们需要定义这种类型在Toy源语言中的接口。通常struct类型的语法如下：

# A struct is defined by using the `struct` keyword followed by a name.
struct MyStruct {
  # Inside of the struct is a list of variable declarations without initializers
  # or shapes, which may also be other previously defined structs.
  var a;
  var b;
}

  struct现在可以作为变量或者参数在函数中使用，方法是使用结构的名称而不是结构内的值？需要通过.符号来访问struct的成员。struct可以被一个复合初始化器(指同类型的其他struct变量？)初始化，也可以使用被使用逗号分隔的被{}包围的初始化列表初始化，一个简单的例子如下所示：

struct Struct {
  var a;
  var b;
}

# User defined generic function my operate on struct types as well.
def multiply_transpose(Struct value) {
  # We can access the elements of a struct via the '.' operator.
  return transpose(value.a) * transpose(value.b);
}

def main() {
  # We initialize struct values using a composite initializer.
  Struct value = {[[1,2,3],[4,5,6]],[[1,2,3],[4,5,6]]};

  # We pass these arguments to functions like we do with variables.
  var c = multiply_transpose(value);
  print(c);
}

在MLIR中定义一个struct类型(Defining a struct in MLIR)

  在MLIR中，我们也需要一个struct类型的表示。MLIR不会提供一个我们需要的完整的类型，所以我们得自己定义。我们将struct类型简单的定义为一个包含一组元素类型的无名容器。struct的名字和它的元素仅用于Toy语言编译器的AST(抽象语法树)，所以我们不需要将类型名字放进MLIR的表达式。

定义类型类(Defining the Type Class)

  就像是在第二章中提到的，MLIR中的类型对象是值类型的，并依靠一个内部存储对象来保存类型的实际数据。Type类本身作为内部TypeStorage对象的一个简单包装器，该对象在一个MLIRContext中只有一个实例。当构建了一个Type，我们只是在内部构造一个存储类的梳理并使其唯一化。
  定义包含参数数据的新类型时(例如，struct类型保存元素类型需要额外的信息)，我们需要提供派生存储类。单例(singleton)类型没有任何额外的数据(例如，index类型)，不需要存储类，仅使用默认的TypeStorage。
  类型存储对象(type storage objects)包含构建和唯一化一个类型实例的所有必要数据。派生存储类必须继承mlir::TypeStorage，并提供一组别名和钩子函数，这些将被MLIRContext用于唯一化。下面是我们的struct类型的存储实例的定义，每个必要的需求都内联了：

// 该类是Toy语言`StructType`类型的内部存储对象类
struct StructTypeStorage : public mlir::TypeStorage {
  /// `KeyTy`类型是必须的，它向存储实例提供了接口。这个类型用于类型存储实例的唯一化
  /// 对于我们提供的struct来说，我们要通过struct包含的元素来保证struct的唯一。
  /// 比方说一个包含两个int类型的struct，与一个包含两个float类型的struct，
  /// 就是通过所包含的成员的数量或类型不同来唯一化。
  using KeyTy = llvm::ArrayRef<mlir::Type>;
  /// StructTypeStorage的构造函数
  StructTypeStorage(llvm::ArrayRef<mlir::Type> elementTypes)
    : elementTypes(elementTypes) {}

  /// 定义比较函数用于KeyTy与当前实例的比较。
  /// 当我们构建了一个新实例时，这用于确保我们还没有这个key的唯一实例。
  bool operator==(const KeyTy &key) const { return key == elementTypes; }

  /// 为KeyTy定义一个哈希函数，用于唯一化该存储实例。
  /// 这个方法不是必须的，因为llvm::ArrayRef和mlir::Type都有可用的散列函数，所以我们可以完全省略它
  static llvm::hash_code hashKey(const KeyTy &key) {
    return llvm::hash_value(key);
  }

  /// 定义一个从参数集合到KeyTy的构造函数。当构建存储实例本身时，这些参数将会被提供。
  /// 详见后续的`StructType::get`方法
  /// 注意：该方法不是必须的，因为KeyTy可以用给定的参数直接构造。
  static KeyTy getKey(llvm::ArrayRef<mlir::Type> elementTypes) {
    return KeyTy(elementTypes);
  }

  /// 定义一个创建该存储的新实例的构造函数。
  /// 该方法需要一个存储分配器实例和一个`KeyTy`实例。
  /// 给定的存储分配器必须用于创建类型存储及其内部的*所有*必要的动态分配。
  static StructTypeStorage *construct(mlir::TypeStorageAllocator &allocator,
                                      const KeyTy &key) {
    // 从参数key中拷贝元素到分配器
    llvm::ArrayRef<mlir::Type> elementTypes = allocator.copyInfo(key);
 
    // 分配存储实例并构建它。
    return new (allocator.allocate<StructTypeStorage>())
      StructTypeStorage(elementTypes);
  }

  /// 下面的字段包含struct类型的元素类型。
  llvm::ArrayRef<mlir::Type> elementTypes;
}

  当存储类定义完毕，我们就可以增加用户可见的StructType类了。这些是我们将实际交互的类。

class StructType : public mlir::Type::TypeBase<StructType, mlir::Type,
                                               StructTypeStorage> {
public:
  /// 从'TypeBase'继承一些必要的构造函数
  using Base::Base;

  /// 根据给定的元素类型创建一个`StructType`实例。
  /// *必须*至少有一个元素类型
  static StructType get(llvm::ArrayRef<mlir::Type> elementTypes) {
    assert(!elementTypes.empty() && "expected at least 1 element type");

    // 调用'TypeBase'中的助手'get'方法来获得该类型的唯一实例。
    // 第一个参数是唯一的上下文。之后的参数将被发到存储实例。
    mlir::MLIRContext *ctx = elementTypes.front().getContext();
    return Base::get(ctx, elementTypes);
  }

  /// 返回该struct类型的所有元素类型
  llvm::ArrayRef<mlir::Type> getElementTypes() {
    // 'getImpl' 返回内在的存储实例指针
    return getImpl()->elementTypes;
  }
  
  /// 返回该struct持有的元素类型数量
  size_t getNumElementTypes() {return getElementTypes().size(); }
};

  我们在Toy方言初始化器中注册这个类型，方法与注册操作类似。

void ToyDialect::initialize() {
  addTypes<StructType>();
}

  需要注意的是，在注册类型是，存储类的dinginess必须是可见的。
  现在在Toy源语言生成MLIR的时候我们可以使用StrucType。详见 examples/toy/Ch7/mlir/MLIRGen.cpp

暴露到ODS？(Exposing to ODS)

  定义完一个新的类型之后，我们应该让ODS框架知道我们的新类型，这样我们才能在方言的Op定义和自动生成中使用它。一个简单的例子：

// 在ODS中提供一个Toy StructType的定义。
// 这允许类似于Tensor或MemRef的方式使用StructType
// 我们使用'DialectType'将structType类型划分为Toy方言之列。
def Toy_StructType :
  DialectType<Toy_Dialect, CPred<"$_self.isa<StructType>()">,
              "Toy strct type">；

// 提供Toy方言中使用的类型的定义。
def Toy_Type : AnyTypeOf<[F64Tensor, Toy_StructType]>;

解析与打印(Parsing and Printing)

  此时，我们可以在MLIR生成和转换期间使用StructType，但是我们还不能输出或解析.mlir文件。为了做到这点，我们需要支持StructType实例的解析与打印。可以通过在ToyDialect中重写parseType和printType方法完成。这些方法的声明在类型暴露给ODS时自动提供，如前一节所述。

class ToyDialect : public mlir::Dialect {
public:
  /// 解析一个注册在方言上的类型实例
  mlir::Type parseType(mlir::DialectAsmParser &parser) const override;

  /// 打印一个注册在方言上的类型实例
  void printType(mlir::Type type,
                 mlir::DialectAsmPrinter &printer) const override;
};

  这些方法需要高级别parser或printer的实例，可以方便的实现必要的功能。在开始实现之前，让我们考虑一下struct类型在输出IR上的语法(形状？)。根据MLIR language reference，方言类型的通用格式应形如：!dialect-namespace<type-data>，在某些情况下这样的写法能有一个漂亮的形式。Toy的parser与printer的职责就是提供格式中的type-data。我们将把StructType定义成如下格式：

struct-type ::= `struct` `<` type(`,` type)* `>`

解析(Parsing)

  一个parsing的实现如下：

/// 解析一个注册在方言上的类型
mlir::Type ToyDialect::parseType(mlir::DialectAsmParser &parser) const {
  // 以以下形式解析struct类型
  // struct-type ::= `struct` `<` type(`,` type)* `>`

  // 注意：所有的MLIR parser函数都返回ParseResult。
  // 这是LogicalResult的一个特化，在失败时自动转换为' true '布尔值以允许链接，
  // 但可以根据需要使用显式' mlir::failed/mlir::succeeded '
  // 没明白在失败时自动转为失败是为了什么？

  // Parse: `struct` `<`
  if (parser.parseKeyword("struct") || parser.parseLess())
    return Type();

  // Parse the element types of the struct
  SmallVector<mlir::Type, 1> elementTypes;
  do {
    // Parse the current element type.
    llvm::SMLoc typeLoc = parser.getCurrentLocation();
    mlir::Type elementType;
    if (parser.parseType(elementType))
      return nullptr;

    // Check that the type is either a TensorType or another StructType.
    if (!elementType.isa<mlir::TensorType, StructType>()) {
      parser.emitError(typeLoc, "element type for a struct must either "
                                "be a TensorType or a StructType, got: ")
        << elementType;
      return Type();
    }
    elementTypes.push_back(elementType);
  }while (succeeded(parser.parseOptionalComma()));

  // Pares: `>`
  if (parser.parseGreater())
    return Type();
  return StructType::get(elementTypes);

}

打印(Printing)

  一个printer的实现如下：

void ToyDialect::printType(mlir::Type type,
                           mlir::DialectAsmPrinter &printer) const {
  StructType structType = type.cast<StructType>;

  printer << "struct<";
  llvm::interleaveComma(structType.getElementTypes(), printer);
  printer << '>';
}

  在继续前，让我们看一下我们现有功能的示例：

struct Struct {
  var a;
  var b;
}

def multiply_transpose(Struct value) {
}

  生成如下：

module {
  func @multiply_transpose(%arg0: !toy.struct<tensor<*xf64>, tensor<*xf64>>) {
    toy.return
  }
}

在StructType上操作(Operationg on StructType)

  现在struct类型已经被定义，也可以在IR中往返(指的是，既可以打印到IR，也可以从IR中解析)。下一步要增加struct类型在Op中使用的支持。

对已存在的Op进行更新(Updating Existing Operations)

  现在我们有一些Op，例如：ReturnOp，需要被升级以处理Toy_StructType。

def ReturnOp : Toy_Op<"return", [Terminator, HasParent<"FuncOp">]> {
  ...
  let arguments = (ins Variadic<Toy_Type>:$input);
  ...
}

增加新的Toy Op(Adding New Toy Operations)

  除了现有的Op之外，我们还将添加一些新的Op，这些Op将提供对struct类型的更具体操作。
toy.struct_constant
  新的Op实现了struct类型的一个常量值。在我们当前的模型中，我们只需要用一个数组属性，该属性包含了一组代表struct元素类型的常量值。

%0 = toy.struct_constant [
  dense<[[1.0, 2.0, 3.0], [4.0, 5.0, 6.0]]> : tensor<2x3xf64>
] : !toy.struct<tensor<*xf64>>

toy.struct_access
  这个操作实现了访问struct第N个元素。

// 使用之前的%0
%1 = toy.struct_access %0[0] : !toy.struct<tensor<*xf64>> -> tensor<*xf64>

  有了这些操作，我们可以重新审视最初的示例:

struct Struct {
  var a;
  var b;
}

# User defined generic function my operate on struct types as well.
def multiply_transpose(Struct value) {
  # We can access the elements of a struct via the '.' operator.
  return transpose(value.a) * transpose(value.b);
}

def main() {
  # We initialize struct values using a composite initializer.
  Struct value = {[[1,2,3],[4,5,6]],[[1,2,3],[4,5,6]]};

  # We pass these arguments to functions like we do with variables.
  var c = multiply_transpose(value);
  print(c);
}

最终得到完整的MLIR模块:

module {
  func @multiply_transpose(%arg0: !toy.struct<tensor<*xf64>, tensor<*xf64>>) -> tensor<*xf64> {
    %0 = toy.struct_access %arg0[0] : !toy.struct<tensor<*xf64>, tensor<*xf64>> -> tensor<*xf64>
    %1 = toy.transpose(%0 : tensor<*xf64>) to tensor<*xf64>
    %2 = toy.struct_access %arg0[1] : !toy.struct<tensor<*xf64>, tensor<*xf64>> -> tensor<*xf64>
    %3 = toy.transpose(%2 : tensor<*xf64>) to tensor<*xf64>
    %4 = toy.mul %1, %3 : tensor<*xf64>
    toy.return %4 : tensor<*xf64>
  }
  func @main(){
    %0 = toy.struct_constant [
      dense<[[1.000000e+00, 2.000000e+00, 3.000000e+00], [4.000000e+00, 5.000000e+00, 6.000000e+00]]> : tensor<2x3xf64>,
      dense<[[1.000000e+00, 2.000000e+00, 3.000000e+00], [4.000000e+00, 5.000000e+00, 6.000000e+00]]> : tensor<2x3xf64>
      ] : !toy.struct<tensor<*xf64>, tensor<*xf64>>
      %1 = toy.generic_call @multiply_transpose(%0) : (!toy.struct<tensor<*xf64>, tensor<*xf64>>) -> tensor<*xf64>
      toy.print %1 : tensor<*xf64>
      toy.return
  }
}

优化StructType上的Op(Optimizing Operations on StructType)

  现在我们有几个能在StructType上操作的Op，我们也有了许多新的常量折叠机会。内联之后，前一章中的MLIR模块看起来如下所示：

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

  我们有几个toy.struct_accessOp用于访问toy.struct_constant。如第三章所述(FoldConstantReshape)，可以为这些toy操作增加文件夹，通过在Op定义上设置hasFolder位，并提供*Op::fold方法的定义。

/// Fold constants.
OpFoldResult ConstantOp::fold(ArrayRef<Attribute> operands) { return value(); }

/// Fold struct constants.
OpFoldResult StructConstantOp::fold(ArrayRef<Attribute> operands) {
  return value();
}

/// Fold simple struct access operations that access into a constant.
OpFoldResult StructAccessOp::fold(ArrayRef<Attribute> operands) {
  auto structAttr = operands.front().dyn_cast_or_null<mlir::ArrayAttr>();
  if (!structAttr)
    return nullptr;

  size_t elementIndex = index().getZExtValue();
  return structAttr[elementIndex];
}

  为了确保MLIR在折叠Toy Op时生成正确的常量操作，比如，TensorType的ConstantOp，StructType的StructConstant，我们需要重写方言钩子函数materializeConstant。允许通用MLIR Op在必要时创建Toy的常量。

mlir::Operation *ToyDialect::materializeConstant(mlir::OpBuilder &builder,
                                                 mlir::Attribute value,
                                                 mlir::Type type,
                                                 mlir::Location loc) {
  if (type.isa<StructType>())
    return builder.create<StructConstantOp>(loc, type,
                                            value.cast<mlir::ArrayAttr>());
  return builder.create<ConstantOp>(loc, type,
                                    value.cast<mlir::DenseElementsAttr>());
}

  有了这个，我们可以生成LLVM的代码，而不需要对管道进行任何更改。

module {
  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.transpose(%0 : tensor<2x3xf64>) to tensor<3x2xf64>
    %2 = toy.mul %1, %1 : tensor<3x2xf64>
    toy.print %2 : tensor<3x2xf64>
    toy.return
  }
}