[Mlir-commits] [mlir] [mlir] Share registered operation forwarding hooks (NFC) (PR #226319)
Mehdi Amini
llvmlistbot at llvm.org
Fri Sep 25 04:57:52 PDT 2026
https://github.com/joker-eph updated https://github.com/llvm/llvm-project/pull/226319
>From d9ad2c69821f30e8ad22c2e8c4b32c8ffc5ee6b1 Mon Sep 17 00:00:00 2001
From: Mehdi Amini <joker.eph at gmail.com>
Date: Wed, 23 Sep 2026 18:03:36 -0700
Subject: [PATCH] [mlir] Share registered operation forwarding hooks
Store eight operation-specific callbacks in a per-type table and forward
through common OperationName methods. This reduces repeated model method
instantiations when registering operations.
Dispatch all eight hooks directly from OperationName for static operation
models. Dynamic and unregistered operations keep their virtual fallback.
This covers folding, canonicalization, traits, parsing, printing, default
attributes, and both verifier hooks.
A Clang 22 Release microbenchmark with assertions off, pinned to CPU 2,
measured these medians over 15 interleaved runs (ns/call, parent -> final):
empty op verifyInvariants: 0.940 -> 0.942
empty op verify(op, false): 24.65 -> 24.45
empty ModuleOp invariants: 16.90 -> 16.12
empty ModuleOp full verify: 56.23 -> 55.91
Ten longer alternating pairs found no stable full ModuleOp regression:
56.22 -> 56.27 ns/call, with a median paired difference of +0.11 ns.
The Release MLIRIR and MLIRParser builds passed. Parsing, printing, and
calling the registered hooks passed a smoke test, as did the dynamic
operation verifier fallback.
Assisted-by: Codex
---
mlir/include/mlir/IR/OperationSupport.h | 112 +++++++++++++++++-------
mlir/lib/IR/OperationSupport.cpp | 35 ++++++++
2 files changed, 113 insertions(+), 34 deletions(-)
diff --git a/mlir/include/mlir/IR/OperationSupport.h b/mlir/include/mlir/IR/OperationSupport.h
index 991d9847672e2..c761e86a0d8d7 100644
--- a/mlir/include/mlir/IR/OperationSupport.h
+++ b/mlir/include/mlir/IR/OperationSupport.h
@@ -200,6 +200,33 @@ class OperationName {
: name(name), typeID(typeID), dialect(dialect),
interfaceMap(std::move(interfaceMap)) {}
+ /// Callbacks common to all registered operation models. Store one table
+ /// per operation type so the virtual forwarding methods can be shared.
+ struct HookFns {
+ LogicalResult (*fold)(Operation *, ArrayRef<Attribute>,
+ SmallVectorImpl<OpFoldResult> &);
+ void (*canonicalize)(RewritePatternSet &, MLIRContext *);
+ bool (*hasTrait)(TypeID);
+ ParseResult (*parse)(OpAsmParser &, OperationState &);
+ void (*populateDefaultAttrs)(const OperationName &, NamedAttrList &);
+ void (*print)(Operation *, OpAsmPrinter &, StringRef);
+ LogicalResult (*verify)(Operation *);
+ LogicalResult (*verifyRegions)(Operation *);
+ };
+
+ LogicalResult foldHook(Operation *op, ArrayRef<Attribute> attrs,
+ SmallVectorImpl<OpFoldResult> &results) override;
+ void getCanonicalizationPatterns(RewritePatternSet &set,
+ MLIRContext *context) override;
+ bool hasTrait(TypeID id) override;
+ OperationName::ParseAssemblyFn getParseAssemblyFn() override;
+ void populateDefaultAttrs(const OperationName &name,
+ NamedAttrList &attrs) override;
+ void printAssembly(Operation *op, OpAsmPrinter &printer,
+ StringRef name) override;
+ LogicalResult verifyInvariants(Operation *op) override;
+ LogicalResult verifyRegionInvariants(Operation *op) override;
+
/// Returns true if this is a registered operation.
bool isRegistered() const { return typeID != TypeID::get<void>(); }
detail::InterfaceMap &getInterfaceMap() { return interfaceMap; }
@@ -236,6 +263,11 @@ class OperationName {
/// The TypeID of the Properties struct for this operation.
TypeID propertiesTypeID;
+ /// The forwarding callbacks for a registered operation. Unregistered
+ /// operations override these hooks and leave this pointer null.
+ const HookFns *hookFns = nullptr;
+
+ friend class OperationName;
friend class RegisteredOperationName;
};
@@ -314,14 +346,20 @@ class OperationName {
/// generalized constant folding.
LogicalResult foldHook(Operation *op, ArrayRef<Attribute> operands,
SmallVectorImpl<OpFoldResult> &results) const {
- return getImpl()->foldHook(op, operands, results);
+ Impl *impl = getImpl();
+ if (const Impl::HookFns *hooks = impl->hookFns)
+ return hooks->fold(op, operands, results);
+ return impl->foldHook(op, operands, results);
}
/// This hook returns any canonicalization pattern rewrites that the
/// operation supports, for use by the canonicalization pass.
void getCanonicalizationPatterns(RewritePatternSet &results,
MLIRContext *context) const {
- return getImpl()->getCanonicalizationPatterns(results, context);
+ Impl *impl = getImpl();
+ if (const Impl::HookFns *hooks = impl->hookFns)
+ return hooks->canonicalize(results, context);
+ return impl->getCanonicalizationPatterns(results, context);
}
/// Returns true if the operation was registered with a particular trait, e.g.
@@ -331,7 +369,12 @@ class OperationName {
bool hasTrait() const {
return hasTrait(TypeID::get<Trait>());
}
- bool hasTrait(TypeID traitID) const { return getImpl()->hasTrait(traitID); }
+ bool hasTrait(TypeID traitID) const {
+ Impl *impl = getImpl();
+ if (const Impl::HookFns *hooks = impl->hookFns)
+ return hooks->hasTrait(traitID);
+ return impl->hasTrait(traitID);
+ }
/// Returns true if the operation *might* have the provided trait. This
/// means that either the operation is unregistered, or it was registered with
@@ -341,34 +384,49 @@ class OperationName {
return mightHaveTrait(TypeID::get<Trait>());
}
bool mightHaveTrait(TypeID traitID) const {
- return !isRegistered() || getImpl()->hasTrait(traitID);
+ return !isRegistered() || hasTrait(traitID);
}
/// Return the static hook for parsing this operation assembly.
ParseAssemblyFn getParseAssemblyFn() const {
- return getImpl()->getParseAssemblyFn();
+ Impl *impl = getImpl();
+ if (const Impl::HookFns *hooks = impl->hookFns)
+ return hooks->parse;
+ return impl->getParseAssemblyFn();
}
/// This hook implements the method to populate defaults attributes that are
/// unset.
void populateDefaultAttrs(NamedAttrList &attrs) const {
- getImpl()->populateDefaultAttrs(*this, attrs);
+ Impl *impl = getImpl();
+ if (const Impl::HookFns *hooks = impl->hookFns)
+ return hooks->populateDefaultAttrs(*this, attrs);
+ return impl->populateDefaultAttrs(*this, attrs);
}
/// This hook implements the AsmPrinter for this operation.
void printAssembly(Operation *op, OpAsmPrinter &p,
StringRef defaultDialect) const {
- return getImpl()->printAssembly(op, p, defaultDialect);
+ Impl *impl = getImpl();
+ if (const Impl::HookFns *hooks = impl->hookFns)
+ return hooks->print(op, p, defaultDialect);
+ return impl->printAssembly(op, p, defaultDialect);
}
/// These hooks implement the verifiers for this operation. It should emits
/// an error message and returns failure if a problem is detected, or
/// returns success if everything is ok.
LogicalResult verifyInvariants(Operation *op) const {
- return getImpl()->verifyInvariants(op);
+ Impl *impl = getImpl();
+ if (const Impl::HookFns *hooks = impl->hookFns)
+ return hooks->verify(op);
+ return impl->verifyInvariants(op);
}
LogicalResult verifyRegionInvariants(Operation *op) const {
- return getImpl()->verifyRegionInvariants(op);
+ Impl *impl = getImpl();
+ if (const Impl::HookFns *hooks = impl->hookFns)
+ return hooks->verifyRegions(op);
+ return impl->verifyRegionInvariants(op);
}
/// Return the list of cached attribute names registered to this operation.
@@ -599,32 +657,18 @@ class RegisteredOperationName : public OperationName {
: Impl(ConcreteOp::getOperationName(), dialect,
TypeID::get<ConcreteOp>(), ConcreteOp::getInterfaceMap()) {
propertiesTypeID = TypeID::get<Properties>();
+ hookFns = &getHookFns();
}
- LogicalResult foldHook(Operation *op, ArrayRef<Attribute> attrs,
- SmallVectorImpl<OpFoldResult> &results) final {
- return ConcreteOp::getFoldHookFn()(op, attrs, results);
- }
- void getCanonicalizationPatterns(RewritePatternSet &set,
- MLIRContext *context) final {
- ConcreteOp::getCanonicalizationPatterns(set, context);
- }
- bool hasTrait(TypeID id) final { return ConcreteOp::getHasTraitFn()(id); }
- OperationName::ParseAssemblyFn getParseAssemblyFn() final {
- return ConcreteOp::parse;
- }
- void populateDefaultAttrs(const OperationName &name,
- NamedAttrList &attrs) final {
- ConcreteOp::populateDefaultAttrs(name, attrs);
- }
- void printAssembly(Operation *op, OpAsmPrinter &printer,
- StringRef name) final {
- ConcreteOp::getPrintAssemblyFn()(op, printer, name);
- }
- LogicalResult verifyInvariants(Operation *op) final {
- return ConcreteOp::getVerifyInvariantsFn()(op);
- }
- LogicalResult verifyRegionInvariants(Operation *op) final {
- return ConcreteOp::getVerifyRegionInvariantsFn()(op);
+ static const HookFns &getHookFns() {
+ static const HookFns hooks = {ConcreteOp::getFoldHookFn(),
+ &ConcreteOp::getCanonicalizationPatterns,
+ ConcreteOp::getHasTraitFn(),
+ &ConcreteOp::parse,
+ &ConcreteOp::populateDefaultAttrs,
+ ConcreteOp::getPrintAssemblyFn(),
+ ConcreteOp::getVerifyInvariantsFn(),
+ ConcreteOp::getVerifyRegionInvariantsFn()};
+ return hooks;
}
/// Implementation for "Properties"
diff --git a/mlir/lib/IR/OperationSupport.cpp b/mlir/lib/IR/OperationSupport.cpp
index 545791df641fe..6fc30d413d4bf 100644
--- a/mlir/lib/IR/OperationSupport.cpp
+++ b/mlir/lib/IR/OperationSupport.cpp
@@ -30,6 +30,41 @@ void mlir::detail::appendAttributeProperty(
attrs.emplace_back(name, attr);
}
+LogicalResult
+OperationName::Impl::foldHook(Operation *op, ArrayRef<Attribute> attrs,
+ SmallVectorImpl<OpFoldResult> &results) {
+ return hookFns->fold(op, attrs, results);
+}
+
+void OperationName::Impl::getCanonicalizationPatterns(RewritePatternSet &set,
+ MLIRContext *context) {
+ hookFns->canonicalize(set, context);
+}
+
+bool OperationName::Impl::hasTrait(TypeID id) { return hookFns->hasTrait(id); }
+
+OperationName::ParseAssemblyFn OperationName::Impl::getParseAssemblyFn() {
+ return hookFns->parse;
+}
+
+void OperationName::Impl::populateDefaultAttrs(const OperationName &name,
+ NamedAttrList &attrs) {
+ hookFns->populateDefaultAttrs(name, attrs);
+}
+
+void OperationName::Impl::printAssembly(Operation *op, OpAsmPrinter &printer,
+ StringRef name) {
+ hookFns->print(op, printer, name);
+}
+
+LogicalResult OperationName::Impl::verifyInvariants(Operation *op) {
+ return hookFns->verify(op);
+}
+
+LogicalResult OperationName::Impl::verifyRegionInvariants(Operation *op) {
+ return hookFns->verifyRegions(op);
+}
+
ParseResult mlir::detail::parseOptionalOperandInto(
OpAsmParser &parser,
SmallVectorImpl<OpAsmParser::UnresolvedOperand> &operands) {
More information about the Mlir-commits
mailing list