[mlir] Add support for "promised" interfaces

[springerm]

Summary

Promised interfaces allow for a dialect to "promise" the implementation of an interface, i.e., declare that it supports an interface, but have the interface defined in an extension in a library separate from the dialect itself. A promised interface is powerful in that it alerts the user when the interface is attempted to be used (e.g. via cast/dyn_cast/etc.) and the implementation has not yet been provided. This makes the system much more robust against misconfiguration, and ensures that we do not lose the benefit we currently have of defining the interface in the dialect library.

All external model registrations in InitAllDialects.h are marked as "promised" in the respective dialects. An exception is BufferizableOpInterface: this interface cannot be marked as promised because it is in the Bufferization dialect and has dependencies on other dialects. Only interfaces located in mlir/Interfaces can be marked as "promised".

Note: This revision was copied and is a subset of D120368.

[springerm]

This is a re-upload of River's revision. Trying address a comment mentioned by @mehdi_amini in https://reviews.llvm.org/D145681?id=503721#inline-1408557. So that I can use this revision as a diffbase for my stack.

[mehdi_amini]

This patch will also not be enough I believe to handle OpInterfaces? The end-goal from the beginning was that all of the registerXXXXXOpInterfaceExternalModels() here: https://github.com/llvm/llvm-project/blob/main/mlir/include/mlir/InitAllDialects.h#L145-L168 should be forbidden unless they have been promised in the first place.

mlir/include/mlir/IR/Attributes.h

public:
   using Base = AttributeInterface<ConcreteType, Traits>;
   using InterfaceBase = detail::Interface<ConcreteType, Attribute, Traits,
                                           Attribute, AttributeTrait::TraitBase>;
   using InterfaceBase::InterfaceBase;
 
 private:
   /// Returns the impl interface instance for the given type.
   static typename InterfaceBase::Concept *getInterfaceFor(Attribute attr) {
+#ifndef NDEBUG
+    // Check that the current interface is not an unresolved promise for the
+    // given attribute.
+    dialect_extension_detail::handleUseOfUndefinedPromisedInterface(
+        attr.getDialect(), ConcreteType::getInterfaceID(),
+        llvm::getTypeName<ConcreteType>());
+#endif // NDEBUG
+
     return attr.getAbstractAttribute().getInterface<ConcreteType>();
   }
 
   /// Allow access to 'getInterfaceFor'.
   friend InterfaceBase;
 };
 
 //===----------------------------------------------------------------------===//

mlir/include/mlir/IR/Dialect.h

public:
 
   //===--------------------------------------------------------------------===//
   // Interfaces
   //===--------------------------------------------------------------------===//
 
   /// Lookup an interface for the given ID if one is registered, otherwise
   /// nullptr.
   DialectInterface *getRegisteredInterface(TypeID interfaceID) {
+#ifndef NDEBUG
+    handleUseOfUndefinedPromisedInterface(interfaceID);
+#endif // NDEBUG
+
     auto it = registeredInterfaces.find(interfaceID);
     return it != registeredInterfaces.end() ? it->getSecond().get() : nullptr;
   }
   template <typename InterfaceT>
   InterfaceT *getRegisteredInterface() {
+#ifndef NDEBUG
+    handleUseOfUndefinedPromisedInterface(InterfaceT::getInterfaceID(),
+                                          llvm::getTypeName<InterfaceT>());
+#endif // NDEBUG
+
     return static_cast<InterfaceT *>(
         getRegisteredInterface(InterfaceT::getInterfaceID()));
   }
 
   /// Lookup an op interface for the given ID if one is registered, otherwise
   /// nullptr.
   virtual void *getRegisteredInterfaceForOp(TypeID interfaceID,
                                             OperationName opName) {
#endif // NDEBUG
   }
   template <typename InterfaceT, typename... Args>
   InterfaceT &addInterface(Args &&...args) {
     InterfaceT *interface = new InterfaceT(this, std::forward<Args>(args)...);
     addInterface(std::unique_ptr<DialectInterface>(interface));
     return *interface;
   }
 
+  /// Declare that the given interface will be implemented, but has a delayed
+  /// registration. The promised interface type can be an interface of any type
+  /// not just a dialect interface, i.e. it may also be an
+  /// AttributeInterface/OpInterface/TypeInterface/etc.
+  template <typename InterfaceT> void declarePromisedInterface() {
+    unresolvedPromisedInterfaces.insert(InterfaceT::getInterfaceID());
+  }
+
+  /// Checks if the given interface, which is attempting to be used, is a
+  /// promised interface of this dialect that has yet to be implemented. If so,
+  /// emits a fatal error. `interfaceName` is an optional string that contains a
+  /// more user readable name for the interface (such as the class name).
+  void handleUseOfUndefinedPromisedInterface(TypeID interfaceID,
+                                             StringRef interfaceName = "") {
+    if (unresolvedPromisedInterfaces.count(interfaceID)) {
+      llvm::report_fatal_error(
+          "checking for an interface (`" + interfaceName +
+          "`) that was promised by dialect '" + getNamespace() +
+          "' but never implemented. This is generally an indication "
+          "that the dialect extension implementing the interface was never "
+          "registered.");
+    }
+  }
+  /// Checks if the given interface, which is attempting to be attached to a
+  /// construct owned by this dialect, is a promised interface of this dialect
+  /// that has yet to be implemented. If so, it resolves the interface promise.
+  void handleAdditionOfUndefinedPromisedInterface(TypeID interfaceID) {
+    unresolvedPromisedInterfaces.erase(interfaceID);
+  }
+
 protected:
   /// The constructor takes a unique namespace for this dialect as well as the
   /// context to bind to.
   /// Note: The namespace must not contain '.' characters.
   /// Note: All operations belonging to this dialect must have names starting
   ///       with the namespace followed by '.'.
   /// Example:
   ///       - "tf" for the TensorFlow ops like "tf.add".
private:
   /// Flag that specifies whether this dialect allows unregistered types, i.e.
   /// types prefixed with the dialect namespace but not registered with addType.
   /// These types are represented with OpaqueType.
   bool unknownTypesAllowed = false;
 
   /// A collection of registered dialect interfaces.
   DenseMap<TypeID, std::unique_ptr<DialectInterface>> registeredInterfaces;
 
+  /// A set of interfaces that the dialect (or its constructs, i.e.
+  /// Attributes/Operations/Types/etc.) has promised to implement, but has yet
+  /// to provide an implementation for.
+  DenseSet<TypeID> unresolvedPromisedInterfaces;
+
   friend class DialectRegistry;
   friend void registerDialect();
   friend class MLIRContext;
 };
 
 } // namespace mlir
 
 namespace llvm {

mlir/include/mlir/IR/DialectRegistry.h

void apply(MLIRContext *context,
     unsigned dialectIdx = 0;
     auto derivedDialects = std::tuple<DialectsT *...>{
         static_cast<DialectsT *>(dialects[dialectIdx++])...};
     std::apply([&](DialectsT *...dialect) { apply(context, dialect...); },
                derivedDialects);
   }
 };
 
+namespace dialect_extension_detail {
+
+/// Checks if the given interface, which is attempting to be used, is a
+/// promised interface of this dialect that has yet to be implemented. If so,
+/// emits a fatal error.
+void handleUseOfUndefinedPromisedInterface(Dialect &dialect, TypeID interfaceID,
+                                           StringRef interfaceName);
+
+/// Checks if the given interface, which is attempting to be attached, is a
+/// promised interface of this dialect that has yet to be implemented. If so,
+/// the promised interface is marked as resolved.
+void handleAdditionOfUndefinedPromisedInterface(Dialect &dialect,
+                                                TypeID interfaceID);
+
+} // namespace dialect_extension_detail
+
 //===----------------------------------------------------------------------===//
 // DialectRegistry
 //===----------------------------------------------------------------------===//
 
 /// The DialectRegistry maps a dialect namespace to a constructor for the
 /// matching dialect. This allows for decoupling the list of dialects
 /// "available" from the dialects loaded in the Context. The parser in
 /// particular will lazily load dialects in the Context as operations are

mlir/include/mlir/IR/OpDefinition.h

public:
   /// Inherit the base class constructor.
   using InterfaceBase::InterfaceBase;
 
 protected:
   /// Returns the impl interface instance for the given operation.
   static typename InterfaceBase::Concept *getInterfaceFor(Operation *op) {
     OperationName name = op->getName();
 
+#ifndef NDEBUG
+    // Check that the current interface is not an unresolved promise for the
+    // given operation.
+    if (Dialect *dialect = name.getDialect()) {
+      dialect_extension_detail::handleUseOfUndefinedPromisedInterface(
+          *dialect, ConcreteType::getInterfaceID(),
+          llvm::getTypeName<ConcreteType>());
+    }
+#endif // NDEBUG
+
     // Access the raw interface from the operation info.
     if (std::optional<RegisteredOperationName> rInfo =
             name.getRegisteredInfo()) {
       if (auto *opIface = rInfo->getInterface<ConcreteType>())
         return opIface;
       // Fallback to the dialect to provide it with a chance to implement this
       // interface for this operation.
       return rInfo->getDialect().getRegisteredInterfaceForOp<ConcreteType>(

mlir/include/mlir/IR/OperationSupport.h

public:
   typename T::Concept *getInterface() const {
     return getImpl()->getInterfaceMap().lookup<T>();
   }
 
   /// Attach the given models as implementations of the corresponding
   /// interfaces for the concrete operation.
   template <typename... Models>
   void attachInterface() {
+    // Handle the case where the models resolve a promised interface.
+    (void)std::initializer_list<int>{
+        (dialect_extension_detail::handleAdditionOfUndefinedPromisedInterface(
+             *getDialect(), Models::Interface::getInterfaceID()),
+         0)...};
+
     getImpl()->getInterfaceMap().insertModels<Models...>();
   }
 
   /// Returns true if this operation has the given interface registered to it.
   template <typename T>
   bool hasInterface() const {
     return hasInterface(TypeID::get<T>());
   }

mlir/include/mlir/IR/StorageUniquerSupport.h

 //===- StorageUniquerSupport.h - MLIR Storage Uniquer Utilities -*- 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 utility classes for interfacing with StorageUniquer.
 //
 //===----------------------------------------------------------------------===//
 
 #ifndef MLIR_IR_STORAGEUNIQUERSUPPORT_H
 #define MLIR_IR_STORAGEUNIQUERSUPPORT_H
 
 #include "mlir/IR/AttrTypeSubElements.h"
+#include "mlir/IR/DialectRegistry.h"
 #include "mlir/Support/InterfaceSupport.h"
 #include "mlir/Support/LogicalResult.h"
 #include "mlir/Support/StorageUniquer.h"
 #include "mlir/Support/TypeID.h"
 #include "llvm/ADT/FunctionExtras.h"
 
 namespace mlir {
 class InFlightDiagnostic;
static void attachInterface(MLIRContext &context) {
     typename ConcreteT::AbstractTy *abstract =
         ConcreteT::AbstractTy::lookupMutable(TypeID::get<ConcreteT>(),
                                              &context);
     if (!abstract)
       llvm::report_fatal_error("Registering an interface for an attribute/type "
                                "that is not itself registered.");
 
     (checkInterfaceTarget<IfaceModels>(), ...);
+
+    // Handle the case where the models resolve a promised interface.
+    (void)std::initializer_list<int>{
+        (dialect_extension_detail::handleAdditionOfUndefinedPromisedInterface(
+             abstract->getDialect(), IfaceModels::Interface::getInterfaceID()),
+         0)...};
+
     abstract->interfaceMap.template insertModels<IfaceModels...>();
   }
 
   /// Get or create a new ConcreteT instance within the ctx. This
   /// function is guaranteed to return a non null object and will assert if
   /// the arguments provided are invalid.
   template <typename... Args>
   static ConcreteT get(MLIRContext *ctx, Args... args) {

mlir/include/mlir/IR/Types.h

public:
   using Base = TypeInterface<ConcreteType, Traits>;
   using InterfaceBase =
       detail::Interface<ConcreteType, Type, Traits, Type, TypeTrait::TraitBase>;
   using InterfaceBase::InterfaceBase;
 
 private:
   /// Returns the impl interface instance for the given type.
   static typename InterfaceBase::Concept *getInterfaceFor(Type type) {
+#ifndef NDEBUG
+    // Check that the current interface isn't an unresolved promise for the
+    // given type.
+    dialect_extension_detail::handleUseOfUndefinedPromisedInterface(
+        type.getDialect(), ConcreteType::getInterfaceID(),
+        llvm::getTypeName<ConcreteType>());
+#endif // NDEBUG
+
     return type.getAbstractType().getInterface<ConcreteType>();
   }
 
   /// Allow access to 'getInterfaceFor'.
   friend InterfaceBase;
 };
 
 //===----------------------------------------------------------------------===//

mlir/lib/Dialect/Linalg/IR/LinalgDialect.cpp

 #include "mlir/Dialect/Linalg/IR/Linalg.h"
 #include "mlir/Dialect/Math/IR/Math.h"
 #include "mlir/Dialect/MemRef/IR/MemRef.h"
 #include "mlir/Dialect/Tensor/IR/Tensor.h"
 #include "mlir/IR/BuiltinTypes.h"
 #include "mlir/IR/Dialect.h"
 #include "mlir/IR/DialectImplementation.h"
 #include "mlir/IR/FunctionInterfaces.h"
+#include "mlir/Interfaces/TilingInterface.h"
 #include "mlir/Parser/Parser.h"
 #include "mlir/Support/LLVM.h"
 #include "mlir/Transforms/InliningUtils.h"
 
 #include "llvm/ADT/StringExtras.h"
 #include "llvm/ADT/TypeSwitch.h"
 #include "llvm/Support/raw_ostream.h"
 
#include "mlir/Dialect/Linalg/IR/LinalgStructuredOps.cpp.inc"
 
   // Fill the Linalg-specific OpName to RegionBuilder map.
   addNamedOpBuilders<
 #define GET_OP_LIST
 #include "mlir/Dialect/Linalg/IR/LinalgStructuredOps.cpp.inc"
       >(namedStructuredOpRegionBuilders);
 
   addInterfaces<LinalgInlinerInterface>();
+
+  declarePromisedInterface<PartialReductionOpInterface>();
+  declarePromisedInterface<TilingInterface>();
 }
 
 LogicalResult LinalgDialect::verifyOperationAttribute(Operation *op,
                                                       NamedAttribute attr) {
   if (attr.getName() == LinalgDialect::kMemoizedIndexingMapsAttrName)
     return success();
   return op->emitError() << "attribute '" << attr.getName()
                          << "' not supported by the linalg dialect";
 }
 
 #include "mlir/Dialect/Linalg/IR/LinalgOpsEnums.cpp.inc"
 
 #define GET_ATTRDEF_CLASSES
 #include "mlir/Dialect/Linalg/IR/LinalgOpsAttrDefs.cpp.inc"
 
 #include "mlir/Dialect/Linalg/IR/LinalgOpsDialect.cpp.inc"

mlir/lib/Dialect/MemRef/IR/CMakeLists.txt

add_mlir_dialect_library(MLIRMemRefDialect
   LINK_LIBS PUBLIC
   MLIRArithDialect
   MLIRArithUtils
   MLIRControlFlowInterfaces
   MLIRDialect
   MLIRDialectUtils
   MLIRInferTypeOpInterface
   MLIRIR
+  MLIRRuntimeVerifiableOpInterface
   MLIRShapedOpInterfaces
   MLIRSideEffectInterfaces
   MLIRViewLikeInterface
 )

mlir/lib/Dialect/MemRef/IR/MemRefDialect.cpp

 //===----------------------------------------------------------------------===//
 //
 // 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/Dialect/Arith/IR/Arith.h"
 #include "mlir/Dialect/MemRef/IR/MemRef.h"
+#include "mlir/Interfaces/RuntimeVerifiableOpInterface.h"
 #include "mlir/Interfaces/SideEffectInterfaces.h"
 #include "mlir/Transforms/InliningUtils.h"
 #include <optional>
 
 using namespace mlir;
 using namespace mlir::memref;
 
 #include "mlir/Dialect/MemRef/IR/MemRefOpsDialect.cpp.inc"

 };
 } // namespace
 
 void mlir::memref::MemRefDialect::initialize() {
   addOperations<
 #define GET_OP_LIST
 #include "mlir/Dialect/MemRef/IR/MemRefOps.cpp.inc"
       >();
+
   addInterfaces<MemRefInlinerInterface>();
+  declarePromisedInterface<RuntimeVerifiableOpInterface>();
 }
 
 /// Finds the unique dealloc operation (if one exists) for `allocValue`.
 std::optional<Operation *> mlir::memref::findDealloc(Value allocValue) {
   Operation *dealloc = nullptr;
   for (Operation *user : allocValue.getUsers()) {
     if (!hasEffect<MemoryEffects::Free>(user, allocValue))
       continue;
     // If we found > 1 dealloc, return std::nullopt.
     if (dealloc)
       return std::nullopt;
     dealloc = user;
   }
   return dealloc;
 }

mlir/lib/Dialect/Tensor/IR/TensorDialect.cpp

 //===----------------------------------------------------------------------===//
 //
 // 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/Dialect/Affine/IR/AffineOps.h"
 #include "mlir/Dialect/Arith/IR/Arith.h"
 #include "mlir/Dialect/Complex/IR/Complex.h"
 #include "mlir/Dialect/Tensor/IR/Tensor.h"
+#include "mlir/Interfaces/InferTypeOpInterface.h"
+#include "mlir/Interfaces/TilingInterface.h"
 #include "mlir/Transforms/InliningUtils.h"
 
 using namespace mlir;
 using namespace mlir::tensor;
 
 #include "mlir/Dialect/Tensor/IR/TensorOpsDialect.cpp.inc"
 
 //===----------------------------------------------------------------------===//

 // TensorDialect Methods
 //===----------------------------------------------------------------------===//
 
 void TensorDialect::initialize() {
   addOperations<
 #define GET_OP_LIST
 #include "mlir/Dialect/Tensor/IR/TensorOps.cpp.inc"
       >();
+
   addInterfaces<TensorInlinerInterface>();
+
+  declarePromisedInterface<ReifyRankedShapedTypeOpInterface>();
+  declarePromisedInterface<TilingInterface>();
 }

mlir/lib/IR/Dialect.cpp

 /// namespace
 bool Dialect::isValidNamespace(StringRef str) {
   llvm::Regex dialectNameRegex("^[a-zA-Z_][a-zA-Z_0-9\\$]*$");
   return dialectNameRegex.match(str);
 }
 
 /// Register a set of dialect interfaces with this dialect instance.
 void Dialect::addInterface(std::unique_ptr<DialectInterface> interface) {
+  // Handle the case where the models resolve a promised interface.
+  handleAdditionOfUndefinedPromisedInterface(interface->getID());
+
   auto it = registeredInterfaces.try_emplace(interface->getID(),
                                              std::move(interface));
   (void)it;
   LLVM_DEBUG({
     if (!it.second) {
       llvm::dbgs() << "[" DEBUG_TYPE
                       "] repeated interface registration for dialect "
                    << getNamespace();

 }
 
 //===----------------------------------------------------------------------===//
 // DialectExtension
 //===----------------------------------------------------------------------===//
 
 DialectExtensionBase::~DialectExtensionBase() = default;
 
+void dialect_extension_detail::handleUseOfUndefinedPromisedInterface(
+    Dialect &dialect, TypeID interfaceID, StringRef interfaceName) {
+  dialect.handleUseOfUndefinedPromisedInterface(interfaceID, interfaceName);
+}
+
+void dialect_extension_detail::handleAdditionOfUndefinedPromisedInterface(
+    Dialect &dialect, TypeID interfaceID) {
+  dialect.handleAdditionOfUndefinedPromisedInterface(interfaceID);
+}
+
 //===----------------------------------------------------------------------===//
 // DialectRegistry
 //===----------------------------------------------------------------------===//
 
 DialectRegistry::DialectRegistry() { insert<BuiltinDialect>(); }
 
 DialectAllocatorFunctionRef
 DialectRegistry::getDialectAllocator(StringRef name) const {

utils/bazel/llvm-project-overlay/mlir/BUILD.bazel

deps = [
         ":ArithUtils",
         ":ControlFlowInterfaces",
         ":CopyOpInterface",
         ":DialectUtils",
         ":IR",
         ":InferTypeOpInterface",
         ":MemRefBaseIncGen",
         ":MemRefOpsIncGen",
+        ":RuntimeVerifiableOpInterface",
         ":ShapedOpInterfaces",
         ":ViewLikeInterface",
         "//llvm:Support",
         "//llvm:TargetParser",
     ],
 )
 
 gentbl_cc_library(