[mlir][ods] Optimize FoldAdaptor constructor

[Mogball]

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.

[llvmbot]

@llvm/pr-subscribers-mlir

@llvm/pr-subscribers-mlir-core

Author: Jeff Niu (Mogball)

Changes

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.

---

Full diff: https://github.com/llvm/llvm-project/pull/93219.diff

3 Files Affected:

• (modified) mlir/test/mlir-tblgen/op-decl-and-defs.td (+7-12)
• (modified) mlir/test/mlir-tblgen/op-operand.td (-3)
• (modified) mlir/tools/mlir-tblgen/OpDefinitionsGen.cpp (+20-7)

diff --git a/mlir/test/mlir-tblgen/op-decl-and-defs.td b/mlir/test/mlir-tblgen/op-decl-and-defs.td
index 499e3ceecaf04..c615edeb014c6 100644
--- a/mlir/test/mlir-tblgen/op-decl-and-defs.td
+++ b/mlir/test/mlir-tblgen/op-decl-and-defs.td
@@ -58,7 +58,8 @@ def NS_AOp : NS_Op<"a_op", [IsolatedFromAbove, IsolatedFromAbove]> {
 // CHECK: namespace detail {
 // CHECK: class AOpGenericAdaptorBase {
 // CHECK: public:
-// CHECK:   AOpGenericAdaptorBase(AOp{{[[:space:]]}}
+// CHECK:   AOpGenericAdaptorBase(::mlir::DictionaryAttr attrs = {}, const ::mlir::EmptyProperties &properties = {}, ::mlir::RegionRange regions = {}) : odsAttrs(attrs), odsRegions(regions)
+// CHECK:   AOpGenericAdaptorBase(::mlir::Operation *op) : odsAttrs(op->getAttrDictionary()), odsOpName(op->getName()), odsRegions(op->getRegions()) {}
 // CHECK:   ::mlir::IntegerAttr getAttr1Attr();
 // CHECK:   uint32_t getAttr1();
 // CHECK:   ::mlir::FloatAttr getSomeAttr2Attr();
@@ -128,15 +129,8 @@ def NS_AOp : NS_Op<"a_op", [IsolatedFromAbove, IsolatedFromAbove]> {
 
 // DEFS-LABEL: NS::AOp definitions
 
-// DEFS: AOpGenericAdaptorBase::AOpGenericAdaptorBase(::mlir::DictionaryAttr attrs, const ::mlir::EmptyProperties &properties, ::mlir::RegionRange regions) : odsAttrs(attrs), odsRegions(regions)
-
 // Check that `getAttrDictionary()` is used when not using properties.
 
-// DEFS: AOpGenericAdaptorBase::AOpGenericAdaptorBase(AOp op)
-// DEFS-SAME: op->getAttrDictionary()
-// DEFS-SAME: p.getProperties()
-// DEFS-SAME: op->getRegions()
-
 // DECLS: ::mlir::RegionRange AOpGenericAdaptorBase::getSomeRegions()
 // DECLS-NEXT: return odsRegions.drop_front(1);
 // DECLS: ::mlir::RegionRange AOpGenericAdaptorBase::getRegions()
@@ -346,10 +340,11 @@ def NS_NOp : NS_Op<"op_with_properties", []> {
 
 // Check that `getDiscardableAttrDictionary()` is used with properties.
 
-// DEFS: NOpGenericAdaptorBase::NOpGenericAdaptorBase(NOp op) : NOpGenericAdaptorBase(
-// DEFS-SAME: op->getDiscardableAttrDictionary()
-// DEFS-SAME: op.getProperties()
-// DEFS-SAME: op->getRegions()
+// DEFS: NOpGenericAdaptorBase::NOpGenericAdaptorBase(NOp op) :
+// DEFS-SAME: odsAttrs(op->getDiscardableAttrDictionary())
+// DEFS-SAME: odsOpName(op->getName())
+// DEFS-SAME: properties(op.getProperties())
+// DEFS-SAME: odsRegions(op->getRegions())
 
 // Test that type defs have the proper namespaces when used as a constraint.
 // ---
diff --git a/mlir/test/mlir-tblgen/op-operand.td b/mlir/test/mlir-tblgen/op-operand.td
index a749708244798..a2fa1f7046a97 100644
--- a/mlir/test/mlir-tblgen/op-operand.td
+++ b/mlir/test/mlir-tblgen/op-operand.td
@@ -15,9 +15,6 @@ def OpA : NS_Op<"one_normal_operand_op", []> {
 
 // CHECK-LABEL: OpA definitions
 
-// CHECK:      OpAGenericAdaptorBase::OpAGenericAdaptorBase
-// CHECK-SAME: odsAttrs(attrs)
-
 // CHECK:      void OpA::build
 // CHECK:        ::mlir::Value input
 // CHECK:        odsState.addOperands(input);
diff --git a/mlir/tools/mlir-tblgen/OpDefinitionsGen.cpp b/mlir/tools/mlir-tblgen/OpDefinitionsGen.cpp
index e013ccac5dd0f..cddfd647129cb 100644
--- a/mlir/tools/mlir-tblgen/OpDefinitionsGen.cpp
+++ b/mlir/tools/mlir-tblgen/OpDefinitionsGen.cpp
@@ -4101,7 +4101,8 @@ OpOperandAdaptorEmitter::OpOperandAdaptorEmitter(
                              "{}");
     }
     paramList.emplace_back("::mlir::RegionRange", "regions", "{}");
-    auto *baseConstructor = genericAdaptorBase.addConstructor(paramList);
+    auto *baseConstructor =
+        genericAdaptorBase.addConstructor<Method::Inline>(paramList);
     baseConstructor->addMemberInitializer("odsAttrs", "attrs");
     if (useProperties)
       baseConstructor->addMemberInitializer("properties", "properties");
@@ -4163,14 +4164,26 @@ OpOperandAdaptorEmitter::OpOperandAdaptorEmitter(
   // and the value range from the parameter.
   {
     // Base class is in the cpp file and can simply access the members of the op
-    // class to initialize the template independent fields.
-    auto *constructor = genericAdaptorBase.addConstructor(
-        MethodParameter(op.getCppClassName(), "op"));
+    // class to initialize the template independent fields. If the op doesn't
+    // have properties, we can emit a generic constructor inline. Otherwise,
+    // emit it out-of-line because we need the op to be defined.
+    Constructor *constructor;
+    if (useProperties) {
+      constructor = genericAdaptorBase.addConstructor(
+          MethodParameter(op.getCppClassName(), "op"));
+    } else {
+      constructor = genericAdaptorBase.addConstructor<Method::Inline>(
+          MethodParameter("::mlir::Operation *", "op"));
+    }
     constructor->addMemberInitializer(
-        genericAdaptorBase.getClassName(),
+        "odsAttrs",
         llvm::Twine(!useProperties ? "op->getAttrDictionary()"
-                                   : "op->getDiscardableAttrDictionary()") +
-            ", op.getProperties(), op->getRegions()");
+                                   : "op->getDiscardableAttrDictionary()"));
+    // Retrieve the operation name from the op directly.
+    constructor->addMemberInitializer("odsOpName", "op->getName()");
+    if (useProperties)
+      constructor->addMemberInitializer("properties", "op.getProperties()");
+    constructor->addMemberInitializer("odsRegions", "op->getRegions()");
 
     // Generic adaptor is templated and therefore defined inline in the header.
     // We cannot use the Op class here as it is an incomplete type (we have a

[joker-eph]

Nice finding!

[jpienaar]

Thanks!