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

This is off.

21

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

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

19 ↗

(On Diff #285026)

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/ConstantFoldInterfaces.h

+//===- ControlFlowInterfaces.h - ControlFlow 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_CONSTANTFOLDINTERFACES_H_
+#define MLIR_INTERFACES_CONSTANTFOLDINTERFACES_H_
+
+#include "mlir/IR/DialectInterface.h"
+#include "mlir/Support/LogicalResult.h"
+#include "llvm/ADT/SmallVector.h"
+
+namespace mlir {
+class Attribute;
+class OpFoldResult;
+
+/// Define a constant fold interface to allow for dialects to opt-in specific
+/// constant folding for operations they define.
+class DialectConstantFoldInterface
+    : public DialectInterface::Base<DialectConstantFoldInterface> {
+public:
+  /// Registered fallback constant fold 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.
+  virtual LogicalResult
+  constantFold(Operation *op, ArrayRef<Attribute> operands,
+               SmallVectorImpl<OpFoldResult> &results) const {
+    return failure();
+  }
+};
+
+} // end namespace mlir
+
+#endif // MLIR_INTERFACES_CONSTANTFOLDINTERFACES_H_

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:
+  /// 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/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/ConstantFoldInterfaces.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<DialectConstantFoldInterface>();
+  if (!interface)
     return failure();
-  results.assign(constants.begin(), constants.end());
-  return success();
+
+  return interface->constantFold(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;
 }