Remove DialectHooks and introduce a Dialect Interfaces instead

[mehdi_amini]

Summary

These hooks were introduced before the Interfaces mechanism was available.

DialectExtractElementHook is unused and entirely removed. The
DialectConstantFoldHook is used a fallback in the
operation fold() method, and is replaced by a DialectInterface.
The DialectConstantDecodeHook is used for interpreting OpaqueAttribute
and should be revamped, but is replaced with an interface in 1:1 fashion
for now.

[mehdi_amini]

Move the default implementation to the header / remove the .cpp file for the interface

[mehdi_amini]

Restore DecodeHook as an interface

[rriddle]

Thank you for cleaning this up!

mlir/include/mlir/Interfaces/ConstantFoldInterfaces.h

1 ↗

(On Diff #284167)

This is off.

21 ↗

(On Diff #284167)

Can you merge the OpFolderDialectInterface into this one, and remove all of the uses of "Constant" here. We've been using just "Fold" to refer to general folding.

[mehdi_amini]

Rename (drop the "Constant" from the name)

[mehdi_amini]

mlir/include/mlir/Interfaces/ConstantFoldInterfaces.h

21 ↗

(On Diff #284167)

Done in the child revision: https://reviews.llvm.org/D85823

[mehdi_amini]

use lower case for the method

[rriddle]

mlir/include/mlir/Interfaces/FoldInterfaces.h

20

nit: Remove the constant from comments here and below.

mlir/include/mlir/IR/Dialect.h

 
 namespace mlir {
 class DialectAsmParser;
 class DialectAsmPrinter;
 class DialectInterface;
 class OpBuilder;
 class Type;
 
-using DialectConstantDecodeHook =
-    std::function<bool(const OpaqueElementsAttr, ElementsAttr &)>;
-using DialectConstantFoldHook = std::function<LogicalResult(
-    Operation *, ArrayRef<Attribute>, SmallVectorImpl<Attribute> &)>;
-using DialectExtractElementHook =
-    std::function<Attribute(const OpaqueElementsAttr, ArrayRef<uint64_t>)>;
 using DialectAllocatorFunction = std::function<void(MLIRContext *)>;
 
 /// Dialects are groups of MLIR operations and behavior associated with the
 /// entire group.  For example, hooks into other systems for constant folding,
 /// default named types for asm printing, etc.
 ///
 /// Instances of the dialect object are global across all MLIRContext's that may
 /// be active in the process.
public:
   /// addOperation.
   bool allowsUnknownOperations() const { return unknownOpsAllowed; }
 
   /// Return true if this dialect allows for unregistered types, i.e., types
   /// prefixed with the dialect namespace but not registered with addType.
   /// These are represented with OpaqueType.
   bool allowsUnknownTypes() const { return unknownTypesAllowed; }
 
-  //===--------------------------------------------------------------------===//
-  // Constant Hooks
-  //===--------------------------------------------------------------------===//
-
-  /// Registered fallback constant fold hook for the dialect. Like the constant
-  /// fold hook of each operation, it attempts to constant fold the operation
-  /// with the specified constant operand values - the elements in "operands"
-  /// will correspond directly to the operands of the operation, but may be null
-  /// if non-constant.  If constant folding is successful, this fills in the
-  /// `results` vector.  If not, this returns failure and `results` is
-  /// unspecified.
-  DialectConstantFoldHook constantFoldHook =
-      [](Operation *op, ArrayRef<Attribute> operands,
-         SmallVectorImpl<Attribute> &results) { return failure(); };
-
-  /// Registered hook to decode opaque constants associated with this
-  /// dialect. The hook function attempts to decode an opaque constant tensor
-  /// into a tensor with non-opaque content. If decoding is successful, this
-  /// method returns false and sets 'output' attribute. If not, it returns true
-  /// and leaves 'output' unspecified. The default hook fails to decode.
-  DialectConstantDecodeHook decodeHook =
-      [](const OpaqueElementsAttr input, ElementsAttr &output) { return true; };
-
-  /// Registered hook to extract an element from an opaque constant associated
-  /// with this dialect. If element has been successfully extracted, this
-  /// method returns that element. If not, it returns an empty attribute.
-  /// The default hook fails to extract an element.
-  DialectExtractElementHook extractElementHook =
-      [](const OpaqueElementsAttr input, ArrayRef<uint64_t> index) {
-        return Attribute();
-      };
-
   /// Registered hook to materialize a single constant operation from a given
   /// attribute value with the desired resultant type. This method should use
   /// the provided builder to create the operation without changing the
   /// insertion position. The generated operation is expected to be constant
   /// like, i.e. single result, zero operands, non side-effecting, etc. On
   /// success, this hook should return the value generated to represent the
   /// constant value. Otherwise, it should return null on failure.
   virtual Operation *materializeConstant(OpBuilder &builder, Attribute value,

mlir/include/mlir/IR/DialectHooks.h

-//===- DialectHooks.h - MLIR DialectHooks mechanism -------------*- C++ -*-===//
-//
-// Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
-// See https://llvm.org/LICENSE.txt for license information.
-// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
-//
-//===----------------------------------------------------------------------===//
-//
-// This file defines abstraction and registration mechanism for dialect hooks.
-//
-//===----------------------------------------------------------------------===//
-
-#ifndef MLIR_IR_DIALECT_HOOKS_H
-#define MLIR_IR_DIALECT_HOOKS_H
-
-#include "mlir/IR/Dialect.h"
-#include "llvm/Support/raw_ostream.h"
-
-namespace mlir {
-using DialectHooksSetter = std::function<void(MLIRContext *)>;
-
-/// Dialect hooks allow external components to register their functions to
-/// be called for specific tasks specialized per dialect, such as decoding
-/// of opaque constants. To register concrete dialect hooks, one should
-/// define a DialectHooks subclass and use it as a template
-/// argument to DialectHooksRegistration. For example,
-///     class MyHooks : public DialectHooks {...};
-///     static DialectHooksRegistration<MyHooks, MyDialect> hooksReg;
-/// The subclass should override DialectHook methods for supported hooks.
-class DialectHooks {
-public:
-  // Returns hook to constant fold an operation.
-  DialectConstantFoldHook getConstantFoldHook() { return nullptr; }
-  // Returns hook to decode opaque constant tensor.
-  DialectConstantDecodeHook getDecodeHook() { return nullptr; }
-  // Returns hook to extract an element of an opaque constant tensor.
-  DialectExtractElementHook getExtractElementHook() { return nullptr; }
-
-private:
-  /// Registers a function that will set hooks in the registered dialects.
-  /// Registrations are deduplicated by dialect TypeID and only the first
-  /// registration will be used.
-  static void registerDialectHooksSetter(TypeID typeID,
-                                         const DialectHooksSetter &function);
-  template <typename ConcreteHooks>
-  friend void registerDialectHooks(StringRef dialectName);
-};
-
-void registerDialectHooksSetter(TypeID typeID,
-                                const DialectHooksSetter &function);
-
-/// Utility to register dialect hooks. Client can register their dialect hooks
-/// with the global registry by calling
-/// registerDialectHooks<MyHooks>("dialect_namespace");
-template <typename ConcreteHooks>
-void registerDialectHooks(StringRef dialectName) {
-  DialectHooks::registerDialectHooksSetter(
-      TypeID::get<ConcreteHooks>(), [dialectName](MLIRContext *ctx) {
-        Dialect *dialect = ctx->getRegisteredDialect(dialectName);
-        if (!dialect) {
-          llvm::errs() << "error: cannot register hooks for unknown dialect '"
-                       << dialectName << "'\n";
-          abort();
-        }
-        // Set hooks.
-        ConcreteHooks hooks;
-        if (auto h = hooks.getConstantFoldHook())
-          dialect->constantFoldHook = h;
-        if (auto h = hooks.getDecodeHook())
-          dialect->decodeHook = h;
-        if (auto h = hooks.getExtractElementHook())
-          dialect->extractElementHook = h;
-      });
-}
-
-/// DialectHooksRegistration provides a global initializer that registers
-/// a dialect hooks setter routine.
-/// Usage:
-///
-///   // At namespace scope.
-///   static DialectHooksRegistration<MyHooks> Unused("dialect_namespace");
-template <typename ConcreteHooks> struct DialectHooksRegistration {
-  DialectHooksRegistration(StringRef dialectName) {
-    registerDialectHooks<ConcreteHooks>(dialectName);
-  }
-};
-
-} // namespace mlir
-
-#endif

mlir/include/mlir/Interfaces/DecodeAttributesInterfaces.h

+//===- DecodeAttributesInterfaces.h - DecodeAttributes Interfaces -*- C++ -*-=//
+//
+// Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
+// See https://llvm.org/LICENSE.txt for license information.
+// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
+//
+//===----------------------------------------------------------------------===//
+#ifndef MLIR_INTERFACES_DECODEATTRIBUTESINTERFACES_H_
+#define MLIR_INTERFACES_DECODEATTRIBUTESINTERFACES_H_
+
+#include "mlir/IR/Attributes.h"
+#include "mlir/IR/DialectInterface.h"
+#include "mlir/Support/LogicalResult.h"
+
+namespace mlir {
+
+/// Define an interface to decode opaque constant tensor.
+class DialectDecodeAttributesInterface
+    : public DialectInterface::Base<DialectDecodeAttributesInterface> {
+public:
+  DialectDecodeAttributesInterface(Dialect *dialect) : Base(dialect) {}
+
+  /// Registered hook to decode opaque constants associated with this
+  /// dialect. The hook function attempts to decode an opaque constant tensor
+  /// into a tensor with non-opaque content. If decoding is successful, this
+  /// method returns success() and sets 'output' attribute. If not, it returns
+  /// failure() and leaves 'output' unspecified. The default hook fails to
+  /// decode.
+  virtual LogicalResult decode(OpaqueElementsAttr input,
+                               ElementsAttr &output) const {
+    return failure();
+  }
+};
+
+} // end namespace mlir
+
+#endif // MLIR_INTERFACES_DECODEATTRIBUTESINTERFACES_H_

mlir/include/mlir/Interfaces/FoldInterfaces.h

+//===- FoldInterfaces.h - Folding Interfaces --------------------*- C++ -*-===//
+//
+// Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
+// See https://llvm.org/LICENSE.txt for license information.
+// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
+//
+//===----------------------------------------------------------------------===//
+#ifndef MLIR_INTERFACES_FOLDINTERFACES_H_
+#define MLIR_INTERFACES_FOLDINTERFACES_H_
+
+#include "mlir/IR/DialectInterface.h"
+#include "mlir/Support/LogicalResult.h"
+#include "llvm/ADT/SmallVector.h"
+
+namespace mlir {
+class Attribute;
+class OpFoldResult;
+
+/// Define a fold interface to allow for dialects to opt-in specific
+/// folding for operations they define.
+class DialectFoldInterface
+    : public DialectInterface::Base<DialectFoldInterface> {
+public:
+  DialectFoldInterface(Dialect *dialect) : Base(dialect) {}
+
+  /// Registered fallback fold for the dialect. Like the fold hook of each
+  /// operation, it attempts to fold the operation with the specified constant
+  /// operand values - the elements in "operands" will correspond directly to
+  /// the operands of the operation, but may be null if non-constant.  If
+  /// folding is successful, this fills in the `results` vector.  If not, this
+  /// returns failure and `results` is unspecified.
+  virtual LogicalResult Fold(Operation *op, ArrayRef<Attribute> operands,
+                             SmallVectorImpl<OpFoldResult> &results) const {
+    return failure();
+  }
+};
+
+} // end namespace mlir
+
+#endif // MLIR_INTERFACES_FOLDINTERFACES_H_

mlir/lib/IR/Attributes.cpp

 //===- Attributes.cpp - MLIR Affine Expr Classes --------------------------===//
 //
 // Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
 // See https://llvm.org/LICENSE.txt for license information.
 // SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
 //
 //===----------------------------------------------------------------------===//
 
 #include "mlir/IR/Attributes.h"
 #include "AttributeDetail.h"
 #include "mlir/IR/AffineMap.h"
 #include "mlir/IR/Diagnostics.h"
 #include "mlir/IR/Dialect.h"
 #include "mlir/IR/Function.h"
 #include "mlir/IR/IntegerSet.h"
 #include "mlir/IR/Types.h"
+#include "mlir/Interfaces/DecodeAttributesInterfaces.h"
 #include "llvm/ADT/Sequence.h"
 #include "llvm/ADT/Twine.h"
 #include "llvm/Support/Endian.h"
 
 using namespace mlir;
 using namespace mlir::detail;
 
 //===----------------------------------------------------------------------===//

 }
 
 StringRef OpaqueElementsAttr::getValue() const { return getImpl()->bytes; }
 
 /// Return the value at the given index. If index does not refer to a valid
 /// element, then a null attribute is returned.
 Attribute OpaqueElementsAttr::getValue(ArrayRef<uint64_t> index) const {
   assert(isValidIndex(index) && "expected valid multi-dimensional index");
-  if (Dialect *dialect = getDialect())
-    return dialect->extractElementHook(*this, index);
   return Attribute();
 }
 
 Dialect *OpaqueElementsAttr::getDialect() const { return getImpl()->dialect; }
 
 bool OpaqueElementsAttr::decode(ElementsAttr &result) {
-  if (auto *d = getDialect())
-    return d->decodeHook(*this, result);
+  auto *d = getDialect();
+  if (!d)
     return true;
+  auto *interface =
+      d->getRegisteredInterface<DialectDecodeAttributesInterface>();
+  if (!interface)
+    return true;
+  return failed(interface->decode(*this, result));
 }
 
 //===----------------------------------------------------------------------===//
 // SparseElementsAttr
 //===----------------------------------------------------------------------===//
 
 SparseElementsAttr SparseElementsAttr::get(ShapedType type,
                                            DenseElementsAttr indices,

mlir/lib/IR/Dialect.cpp

 //===- Dialect.cpp - Dialect implementation -------------------------------===//
 //
 // Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
 // See https://llvm.org/LICENSE.txt for license information.
 // SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
 //
 //===----------------------------------------------------------------------===//
 
 #include "mlir/IR/Dialect.h"
 #include "mlir/IR/Diagnostics.h"
-#include "mlir/IR/DialectHooks.h"
 #include "mlir/IR/DialectImplementation.h"
 #include "mlir/IR/DialectInterface.h"
 #include "mlir/IR/MLIRContext.h"
 #include "mlir/IR/Operation.h"
 #include "llvm/ADT/MapVector.h"
 #include "llvm/ADT/Twine.h"
 #include "llvm/Support/ManagedStatic.h"
 #include "llvm/Support/Regex.h"
 
 using namespace mlir;
 using namespace detail;
 
 DialectAsmParser::~DialectAsmParser() {}
 
 //===----------------------------------------------------------------------===//
 // Dialect Registration
 //===----------------------------------------------------------------------===//
 
 /// Registry for all dialect allocation functions.
 static llvm::ManagedStatic<llvm::MapVector<TypeID, DialectAllocatorFunction>>
     dialectRegistry;
 
-/// Registry for functions that set dialect hooks.
-static llvm::ManagedStatic<llvm::MapVector<TypeID, DialectHooksSetter>>
-    dialectHooksRegistry;
-
 void Dialect::registerDialectAllocator(
     TypeID typeID, const DialectAllocatorFunction &function) {
   assert(function &&
          "Attempting to register an empty dialect initialize function");
   dialectRegistry->insert({typeID, function});
 }
 
-/// Registers a function to set specific hooks for a specific dialect, typically
-/// used through the DialectHooksRegistration template.
-void DialectHooks::registerDialectHooksSetter(
-    TypeID typeID, const DialectHooksSetter &function) {
-  assert(
-      function &&
-      "Attempting to register an empty dialect hooks initialization function");
-
-  dialectHooksRegistry->insert({typeID, function});
-}
-
 /// Registers all dialects and hooks from the global registries with the
 /// specified MLIRContext.
 void mlir::registerAllDialects(MLIRContext *context) {
   for (const auto &it : *dialectRegistry)
     it.second(context);
-  for (const auto &it : *dialectHooksRegistry)
-    it.second(context);
 }
 
 //===----------------------------------------------------------------------===//
 // Dialect
 //===----------------------------------------------------------------------===//
 
 Dialect::Dialect(StringRef name, MLIRContext *context, TypeID id)
     : name(name), dialectID(id), context(context) {

mlir/lib/IR/Operation.cpp

 //===- Operation.cpp - Operation support code -----------------------------===//
 //
 // Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
 // See https://llvm.org/LICENSE.txt for license information.
 // SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
 //
 //===----------------------------------------------------------------------===//
 
 #include "mlir/IR/Operation.h"
 #include "mlir/IR/BlockAndValueMapping.h"
 #include "mlir/IR/Dialect.h"
 #include "mlir/IR/OpImplementation.h"
 #include "mlir/IR/PatternMatch.h"
 #include "mlir/IR/StandardTypes.h"
 #include "mlir/IR/TypeUtilities.h"
+#include "mlir/Interfaces/FoldInterfaces.h"
 #include <numeric>
 
 using namespace mlir;
 
 OpAsmParser::~OpAsmParser() {}
 
 //===----------------------------------------------------------------------===//
 // OperationName
LogicalResult Operation::fold(ArrayRef<Attribute> operands,
   if (abstractOp && succeeded(abstractOp->foldHook(this, operands, results)))
     return success();
 
   // Otherwise, fall back on the dialect hook to handle it.
   Dialect *dialect = getDialect();
   if (!dialect)
     return failure();
 
-  SmallVector<Attribute, 8> constants;
-  if (failed(dialect->constantFoldHook(this, operands, constants)))
+  auto *interface = dialect->getRegisteredInterface<DialectFoldInterface>();
+  if (!interface)
     return failure();
-  results.assign(constants.begin(), constants.end());
-  return success();
+
+  return interface->Fold(this, operands, results);
 }
 
 /// Emit an error with the op name prefixed, like "'dim' op " which is
 /// convenient for verifiers.
 InFlightDiagnostic Operation::emitOpError(const Twine &message) {
   return emitError() << "'" << getName() << "' op " << message;
 }