Skip to content

Commit 264aaa5

Browse files
author
Jeff Niu
authored
[mlir][ods] Optimize FoldAdaptor constructor (#93219)
FoldAdaptor is generated as a subclass of the operation's generic adaptor, which requires an OperationName instance. It called into the generic base constructor that constructed the OperationName from a string, requiring a StringMap lookup inside the MLIRContext. This makes constructing FoldAdaptors really slow, which is a shame because the `Operation *` is right there. This PR changes GenericAdaptor constructor from an operation to grab the OperationName directly from the `Operation *`. In addition, it generates the constructor inline if the operation doesn't have properties, since otherwise it requires the definition of the op.
1 parent 0f6c4d8 commit 264aaa5

File tree

3 files changed

+27
-24
lines changed

3 files changed

+27
-24
lines changed

mlir/test/mlir-tblgen/op-decl-and-defs.td

Lines changed: 7 additions & 12 deletions
Original file line numberDiff line numberDiff line change
@@ -58,7 +58,8 @@ def NS_AOp : NS_Op<"a_op", [IsolatedFromAbove, IsolatedFromAbove]> {
5858
// CHECK: namespace detail {
5959
// CHECK: class AOpGenericAdaptorBase {
6060
// CHECK: public:
61-
// CHECK: AOpGenericAdaptorBase(AOp{{[[:space:]]}}
61+
// CHECK: AOpGenericAdaptorBase(::mlir::DictionaryAttr attrs = {}, const ::mlir::EmptyProperties &properties = {}, ::mlir::RegionRange regions = {}) : odsAttrs(attrs), odsRegions(regions)
62+
// CHECK: AOpGenericAdaptorBase(::mlir::Operation *op) : odsAttrs(op->getRawDictionaryAttrs()), odsOpName(op->getName()), odsRegions(op->getRegions()) {}
6263
// CHECK: ::mlir::IntegerAttr getAttr1Attr();
6364
// CHECK: uint32_t getAttr1();
6465
// CHECK: ::mlir::FloatAttr getSomeAttr2Attr();
@@ -128,15 +129,8 @@ def NS_AOp : NS_Op<"a_op", [IsolatedFromAbove, IsolatedFromAbove]> {
128129

129130
// DEFS-LABEL: NS::AOp definitions
130131

131-
// DEFS: AOpGenericAdaptorBase::AOpGenericAdaptorBase(::mlir::DictionaryAttr attrs, const ::mlir::EmptyProperties &properties, ::mlir::RegionRange regions) : odsAttrs(attrs), odsRegions(regions)
132-
133132
// Check that `getAttrDictionary()` is used when not using properties.
134133

135-
// DEFS: AOpGenericAdaptorBase::AOpGenericAdaptorBase(AOp op)
136-
// DEFS-SAME: op->getAttrDictionary()
137-
// DEFS-SAME: p.getProperties()
138-
// DEFS-SAME: op->getRegions()
139-
140134
// DECLS: ::mlir::RegionRange AOpGenericAdaptorBase::getSomeRegions()
141135
// DECLS-NEXT: return odsRegions.drop_front(1);
142136
// DECLS: ::mlir::RegionRange AOpGenericAdaptorBase::getRegions()
@@ -346,10 +340,11 @@ def NS_NOp : NS_Op<"op_with_properties", []> {
346340

347341
// Check that `getDiscardableAttrDictionary()` is used with properties.
348342

349-
// DEFS: NOpGenericAdaptorBase::NOpGenericAdaptorBase(NOp op) : NOpGenericAdaptorBase(
350-
// DEFS-SAME: op->getDiscardableAttrDictionary()
351-
// DEFS-SAME: op.getProperties()
352-
// DEFS-SAME: op->getRegions()
343+
// DEFS: NOpGenericAdaptorBase::NOpGenericAdaptorBase(NOp op) :
344+
// DEFS-SAME: odsAttrs(op->getDiscardableAttrDictionary())
345+
// DEFS-SAME: odsOpName(op->getName())
346+
// DEFS-SAME: properties(op.getProperties())
347+
// DEFS-SAME: odsRegions(op->getRegions())
353348

354349
// Test that type defs have the proper namespaces when used as a constraint.
355350
// ---

mlir/test/mlir-tblgen/op-operand.td

Lines changed: 0 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -15,9 +15,6 @@ def OpA : NS_Op<"one_normal_operand_op", []> {
1515

1616
// CHECK-LABEL: OpA definitions
1717

18-
// CHECK: OpAGenericAdaptorBase::OpAGenericAdaptorBase
19-
// CHECK-SAME: odsAttrs(attrs)
20-
2118
// CHECK: void OpA::build
2219
// CHECK: ::mlir::Value input
2320
// CHECK: odsState.addOperands(input);

mlir/tools/mlir-tblgen/OpDefinitionsGen.cpp

Lines changed: 20 additions & 9 deletions
Original file line numberDiff line numberDiff line change
@@ -4101,7 +4101,8 @@ OpOperandAdaptorEmitter::OpOperandAdaptorEmitter(
41014101
"{}");
41024102
}
41034103
paramList.emplace_back("::mlir::RegionRange", "regions", "{}");
4104-
auto *baseConstructor = genericAdaptorBase.addConstructor(paramList);
4104+
auto *baseConstructor =
4105+
genericAdaptorBase.addConstructor<Method::Inline>(paramList);
41054106
baseConstructor->addMemberInitializer("odsAttrs", "attrs");
41064107
if (useProperties)
41074108
baseConstructor->addMemberInitializer("properties", "properties");
@@ -4163,14 +4164,24 @@ OpOperandAdaptorEmitter::OpOperandAdaptorEmitter(
41634164
// and the value range from the parameter.
41644165
{
41654166
// Base class is in the cpp file and can simply access the members of the op
4166-
// class to initialize the template independent fields.
4167-
auto *constructor = genericAdaptorBase.addConstructor(
4168-
MethodParameter(op.getCppClassName(), "op"));
4169-
constructor->addMemberInitializer(
4170-
genericAdaptorBase.getClassName(),
4171-
llvm::Twine(!useProperties ? "op->getAttrDictionary()"
4172-
: "op->getDiscardableAttrDictionary()") +
4173-
", op.getProperties(), op->getRegions()");
4167+
// class to initialize the template independent fields. If the op doesn't
4168+
// have properties, we can emit a generic constructor inline. Otherwise,
4169+
// emit it out-of-line because we need the op to be defined.
4170+
Constructor *constructor;
4171+
if (useProperties) {
4172+
constructor = genericAdaptorBase.addConstructor(
4173+
MethodParameter(op.getCppClassName(), "op"));
4174+
} else {
4175+
constructor = genericAdaptorBase.addConstructor<Method::Inline>(
4176+
MethodParameter("::mlir::Operation *", "op"));
4177+
}
4178+
constructor->addMemberInitializer("odsAttrs",
4179+
"op->getRawDictionaryAttrs()");
4180+
// Retrieve the operation name from the op directly.
4181+
constructor->addMemberInitializer("odsOpName", "op->getName()");
4182+
if (useProperties)
4183+
constructor->addMemberInitializer("properties", "op.getProperties()");
4184+
constructor->addMemberInitializer("odsRegions", "op->getRegions()");
41744185

41754186
// Generic adaptor is templated and therefore defined inline in the header.
41764187
// We cannot use the Op class here as it is an incomplete type (we have a

0 commit comments

Comments
 (0)