Attribute and Constraint
Image credit: MLIROverview
在MLIR中,属性(Attribute)和约束(Constraint)是构建类型安全中间表示的两个核心概念。它们共同确保了操作语义的正确性,为编译器优化和硬件代码生成提供了坚实的基础。本文将深入探讨这两个概念的设计哲学、实现机制以及接着用上一个笔记中的Jiajun Dialect来展示实际应用。
Attribute的本质
Attribute(属性)是编译期就知道的Operation的常量。在MLIR的ODS语法中,TableGen会在C++属性类上提供属性包装器。在上一章中我们提到了Operation的具体定义声明以及实现方式(无论TableGen自动生成还是手写C++实现)。ODS允许在TableGen中使用这些属性来定义Operation,可能具有更细粒度的约束。比如StrAttr直接映射到StringAttr;F32Attr/F64Attr要求FloatAttr额外具有一定的位宽(这些在自定义硬件加速器中比较常见)。ODS属性被定义为具有存储类型(对应于存储属性的mlir::Attribute类),返回类型(对应于生成的getters帮助函数的C++返回类型)以及在内部存储类型和帮助函数进行互转的方法。
属性在MLIR中代表了编译时常量数据,用于参数化操作的行为和特征。与操作数不同,属性在编译期间就具有确定的值,不会随着程序的执行而改变。这种不可变性使得属性成为描述操作元数据的理想选择,比如卷积操作的步长、填充方式,或者常量操作的具体数值。
从实现角度来看,MLIR属性系统设计得十分丰富和灵活。基础属性类型包括整数属性、浮点数属性、字符串属性等,这些基本构建块可以组合成更复杂的结构。数组属性允许将多个相同类型的值组合在一起,这在描述形状或参数列表时特别有用。字典属性则提供了键值对的集合,能够表达更加结构化的配置信息。在AI加速器开发的语境下,属性的价值尤为突出。考虑一个专用的矩阵乘法操作,它可能需要描述data layout、精度模式(例如混合精度量化)等多个维度上的配置。通过属性系统,我们可以将这些编译时常量清晰地表达出来,使得后续的优化和代码生成阶段能够基于这些信息做出更加精准的决策。
这里强调一下Attribute和Type的区别
Type(类型):描述数据的形态和结构,如张量形状、元素类型等。Attribute(属性):描述编译时的常量值或元数据,如具体数值、字符串等。简单来说,Type 告诉你"是什么",Attribute 告诉你"是多少/怎么样"。两者协同工作,共同描述 MLIR 中的静态信息。对于程序员来说,Type应该是一个比较熟悉的概念,而Attribute会在接下来的描述中进行区别。
Attribute定义
让我们通过Jiajun Dialect中的具体例子来理解属性的定义和使用。假设我们需要为矩阵乘法操作定义一个精度控制属性:
def Jiajun_PrecisionAttr : EnumAttr<Jiajun_Dialect, "Precision", [
EnumAttrCase<"FP32", 0>,
EnumAttrCase<"FP16", 1>,
EnumAttrCase<"INT8", 2>,
EnumAttrCase<"INT4", 3>
]> {
let summary = "Compute precision for Jiajun operations";
let cppNamespace = "::jiajun";
}
这个枚举属性定义了四种不同的计算精度模式,每种模式对应一个整数值。在操作中使用这个属性时,我们可以确保只有预定义的精度模式被接受。
对于更复杂的配置场景,字典属性提供了更好的表达力。考虑一个卷积操作的完整配置:
def Jiajun_ConvConfigAttr : Jiajun_Attr<"ConvConfig"> {
let parameters = (ins
"int64_t":$kernel_size,
"int64_t":$stride,
"int64_t":$padding,
"std::string":$data_layout,
"bool":$use_winograd
);
let assemblyFormat = "`<` $kernel_size `,` $stride `,` $padding `,` $data_layout `,` $use_winograd `>`";
}
这种结构化的属性定义使得复杂的配置信息能够以类型安全的方式传递,同时保持文本表示的可读性。当然,这里举的两个例子都比较抽象,接下来将会把Attribute和之前定义过的operator结合起来,同ODS语法一起理解Attribute在tablegen文件中的使用方式。
Attribute 如何与自定义 Operation 紧密结合
让我们回到上一章的某一段代码:
def Jiajun_MulOp : Jiajun_Op<"mul"> {
let summary = "Multiplication operation for Jiajun accelerator";
let description = "This operation performs element-wise multiplication on two tensors";
let arguments = (ins
TensorOf<[F32]>:$lhs,
TensorOf<[F32]>:$rhs,
// 属性定义:一个可选的浮点属性,名为 ‘scale‘
OptionalAttr<F32Attr>:$scale
);
let results = (outs TensorOf<[F32]>:$result);
let assemblyFormat = ...
let builders = [
OpBuilder<(ins "Value":$lhs, "Value":$rhs)>,
OpBuilder<(ins "Value":$lhs, "Value":$rhs, "float":$scale)>
];
}
可以看到,我们在自定义的乘法operator中定义了一个可选的浮点属性。请注意,这里的scale是作为乘法运算中的一个操作属性,这清晰地表明 scale 是操作 jiajun.mul 的一个配置参数,而不是输入或者输出类型。由此也可以看出Attribute和Type的不同。
除此以外,我们还可以定义并使用自定义枚举属性。因为内置的基础属性类型(如整数、浮点数、字符串)虽然有用,但远不足以描述AI加速器中的复杂配置。枚举属性(EnumAttr)允许我们定义一组命名的、互斥的选项,是提高代码可读性和类型安全性的利器。这里就回到了我们最开始描述的代码:
def Jiajun_PrecisionAttr : EnumAttr<Jiajun_Dialect, “Precision”, [
EnumAttrCase<“FP32”, 32>,
EnumAttrCase<“FP16”, 16>,
EnumAttrCase<“INT8”, 8>,
EnumAttrCase<“INT4”, 4>
]> {
let summary = “Computational precision mode for Jiajun accelerator”;
let cppNamespace = “::jiajun”;
}
接着,在一个矩阵乘法操作 Jiajun_MatMulOp 中使用这个自定义的枚举属性:
def Jiajun_MatMulOp : Jiajun_Op<“matmul”> {
let arguments = (ins
TensorOf<[F32, F16]>:$A,
TensorOf<[F32, F16]>:$B,
// 使用自定义枚举属性,并设置默认值为FP32
Jiajun_PrecisionAttr:$precision = “FP32”
);
let results = (outs TensorOf<[F32, F16]>:$C);
let assemblyFormat = “$A `,` $B (`precision` `=` $precision^)? ”
“attr-dict `:` type($A) `,` type($B) `->` type($C)”;
}
现在,在MLIR文本中,我们可以清晰地指定矩阵乘法的精度:
%result = jiajun.matmul %matrix_a, %matrix_b precision = INT8 : tensor<128x128xf32>, tensor<128x128xf32> -> tensor<128x128xf16>
枚举属性将魔数(magic number)转换为了有意义的符号名称,使得IR更易于理解和维护。
对于更复杂的操作,单个标量或枚举属性可能不够。例如,一个卷积操作 Jiajun_ConvOp 需要内核大小、步长、填充模式、数据布局等多个参数。这时,字典属性(DictionaryAttr)或自定义的结构化属性就派上了用场,利用字典属性封装复杂配置:
我们可以定义一个专门的配置属性,将所有这些参数封装在一起:
def Jiajun_ConvConfigAttr : Jiajun_Attr<“ConvConfig”> {
let parameters = (ins
“int64_t”:$kernel_h,
“int64_t”:$kernel_w,
“int64_t”:$stride_h,
“int64_t”:$stride_w,
“std::string”:$padding_mode, // 如 “SAME”, “VALID”
“bool”:$use_winograd
);
let summary = “Structured configuration for convolution operation”;
}
然后,在卷积操作中引用它:
def Jiajun_ConvOp : Jiajun_Op<“conv”> {
let arguments = (ins
TensorOf<[F32]>:$input,
TensorOf<[F32]>:$filter,
// 使用结构化配置属性
Jiajun_ConvConfigAttr:$config
);
let results = (outs TensorOf<[F32]>:$output);
let assemblyFormat = “$input `,` $filter `config` `=` $config ”
“attr-dict `:` type($input) `,` type($filter) `->` type($output)”;
}
在C++代码中,我们可以方便地通过 getConfig() 方法访问这个结构化属性的所有字段,进行统一的验证和处理。这种设计不仅使Operation的定义更加简洁,也保证了相关配置参数的完整性和一致性。
关于Attribute,读者可以直接通过学习MLIR的C++源码,来进一步理解其定义与作用:https://github.com/llvm/llvm-project/blob/main/mlir/include/mlir/IR/Attributes.h
Constraint: 设计哲学
Constraint约束系统在MLIR中扮演着验证者的角色,它确保操作的使用符合设计者的预期。约束可以作用于多个层面:类型约束确保操作数和结果具有兼容的类型,属性约束验证属性值的有效性,而操作数数量约束则保证操作接收正确数量的输入。
类型约束是约束系统中最常见的形式之一。在深度学习计算中,我们经常需要确保两个张量具有相同的形状,或者某个操作数具有特定的元素类型。MLIR通过类型约束机制将这些要求显式地表达出来,而不是隐藏在操作的实现代码中。这种显式表达不仅提高了代码的可读性,还使得错误信息更加清晰和具体。
属性约束则关注于属性值的合理性检查。例如,卷积操作的步长必须为正整数,缩放因子不能为零等。这些约束在操作构建时就会进行验证,及早发现配置错误,避免将问题传递到后续的编译阶段。
约束系统的实现机制
约束在MLIR中的实现依赖于谓词系统,这些谓词在编译时或运行时对操作进行验证。让我们深入探讨几种常见的约束类型及其实现方式。
类型匹配约束是最基础的约束形式之一。在Jiajun Dialect中,我们可能要求矩阵乘法的两个输入矩阵在特定维度上具有兼容的形状:
def Jiajun_MatMulOp : Jiajun_Op<“matmul”> { let arguments = (ins TensorOf<[F32, F16]>:$A, TensorOf<[F32, F16]>:$B );
let constraints = [
TypesMatchWith<“A and B must have compatible inner dimensions”,
“A”, “B”, [{
auto typeA = $_self.cast
这个约束通过C++代码块实现了复杂的形状验证逻辑,确保第一个矩阵的列数与第二个矩阵的行数相匹配。
对于属性约束,我们可以定义更加精细的验证规则。考虑一个分块操作的块大小属性:
def PowerOfTwoAttr : AttrConstraint<CPred<"$_self.cast
def Jiajun_TileOp : Jiajun_Op<“tile”> { let arguments = (ins TensorOf<[AnyType]>:$input, I64Attr:$tile_size );
let verifier = [{ if (!getTileSize().getValue().isPowerOf2()) { return emitOpError(“tile size must be a power of two”); } return success(); }]; }
这里我们展示了两种实现方式:通过预定义的约束谓词和在验证器中手动实现。两种方法各有优势,约束谓词可以提供更早的错误检测,而验证器则支持更复杂的逻辑。
高级约束模式
在实际的AI加速器开发中,我们经常需要实现更加复杂的约束模式。区域约束确保操作包含特定结构的子图,特性约束验证操作具有某些语义特性,而互斥约束保证属性组合的合理性。
考虑一个融合操作,它需要验证内部区域的结构:
def Jiajun_FusedOp : Jiajun_Op<“fused”> { let regions = (region AnyRegion:$body);
let verifier = [{ auto &body = getBody(); if (body.empty()) { return emitOpError(“must have a non-empty body region”); }
// 验证区域只有一个块
if (std::next(body.begin()) != body.end()) {
return emitOpError("must have exactly one block");
}
// 验证特定操作序列
return verifyFusedPattern(body.front());
}]; }
这种区域级别的约束确保了融合操作内部结构的正确性,为后续的优化和代码生成提供了可靠的基础。
属性与约束的协同工作
属性和约束在实际使用中往往是紧密配合的。属性定义了操作的可配置参数,而约束则确保这些参数的合理性和一致性。这种协同工作模式在复杂的AI加速器操作中表现得尤为明显。
考虑一个支持多种数据布局的矩阵转置操作:
def Jiajun_LayoutAttr : EnumAttr<Jiajun_Dialect, “Layout”, [ EnumAttrCase<“NCHW”, 0>, EnumAttrCase<“NHWC”, 1>, EnumAttrCase<“CHWN”, 2> ]>;
def Jiajun_TransposeOp : Jiajun_Op<“transpose”> { let arguments = (ins TensorOf<[AnyType]>:$input, Jiajun_LayoutAttr:$src_layout, Jiajun_LayoutAttr:$dst_layout );
let constraints = [ AttrConstraintWith<“source and destination layouts must be different”, “src_layout”, “dst_layout”, [{ return $_self != $_other; }]> ];
let verifier = [{
auto inputType = getInput().getType().cast
// 根据布局验证输入形状的合理性
if (!isValidLayoutForShape(getSrcLayout(), inputType.getShape())) {
return emitOpError("input shape is incompatible with source layout");
}
return success();
}]; }
这个例子展示了属性定义、属性约束和复杂验证逻辑的有机结合。布局属性定义了操作的行为特征,约束确保配置的基本合理性,而验证器则执行更深层次的语义检查。
实际应用中的最佳实践
在AI加速器开发中有效使用属性和约束需要遵循一些最佳实践。首先,应该充分利用MLIR内置的约束类型,这些类型经过了充分的测试和优化。其次,对于复杂的验证逻辑,应该优先考虑在TableGen中表达,这样可以利用MLIR的自动验证基础设施。
自定义约束的命名应该具有描述性,这样在错误信息中能够提供清晰的指导。例如,SquareMatrixConstraint比简单的ShapeConstraint更能准确描述问题的本质。
在性能敏感的场景中,需要注意约束验证的开销。编译时常量检查应该尽量通过约束谓词实现,而不是推迟到运行时验证。对于复杂的形状推断和验证,考虑使用接口而不是硬编码的约束逻辑,这样可以提高代码的复用性。
调试与错误处理
当约束验证失败时,MLIR会提供详细的错误信息来帮助定位问题。理解这些错误信息的结构对于高效调试至关重要。典型的约束错误信息包括失败的位置、约束的描述以及相关的操作信息。
我们可以通过自定义错误信息来提升调试体验:
def CompatibleShapesConstraint : TypeConstraint< CPred<“areCompatibleShapes($_self, $_other)">, “compatible shapes with $0”>;
def Jiajun_BroadcastOp : Jiajun_Op<“broadcast”> { let arguments = (ins TensorOf<[AnyType]>:$input, TensorOf<[AnyType]>:$target );
let constraints = [ CompatibleShapesConstraint<“target”>, SameElementTypeConstraint ]; }
通过精心设计的约束描述,错误信息能够明确指出形状不兼容的具体维度,大大减少了调试时间。
总结
属性与约束系统是MLIR类型安全的核心保障,它们共同构建了一个既灵活又可靠的中间表示框架。在AI加速器开发中,合理运用这两个概念可以显著提高代码的健壮性和可维护性。
属性系统通过丰富的类型层次和组合机制,为操作参数化提供了强大的表达能力。约束系统则通过多层次的验证机制,确保操作使用符合设计意图。两者的协同工作使得MLIR能够在保持高度灵活性的同时,提供编译时的安全保障。
在实际项目中,建议从简单的属性定义和基础约束开始,逐步引入更复杂的验证逻辑。通过遵循最佳实践和充分利用MLIR的基础设施,可以构建出既强大又易于维护的方言系统,为高效的AI加速器开发奠定坚实基础。