Add a mechanism for Dialects to provide a fallback for OpInterface

[mehdi_amini]

Summary

This mechanism makes it possible for a dialect to not register all
operations but still answer interface-based queries.
This can useful for dialects that are "open" or connected to an external
system and still interoperate with the compiler. It can also open up the
possibility to have a more extensible compiler at runtime: the compiler
does not need a pre-registration for each operation and the dialect can
inject behavior dynamically.

[mehdi_amini]

Remove interfacemap

[mehdi_amini]

Typo fix

[jpienaar]

I forget, what was River's suggestion to make this more efficient?

mlir/include/mlir/IR/OpDefinition.h

1747

Nit:

if (auto *opIface = ...) return

to reduce scope (and not used in early return)

mlir/test/lib/Dialect/Test/TestDialect.cpp

177

If this was

return op->getName().getStringRef() == "test.unregistered_side_effect_op";

would that also restrict it to this op?

181

Why is ::mlir needed?

204

Instantiate

206

Is there a reason why not unique_ptr?

223

So for every op, one would have an entry here? Probably up to dialect to implement in efficient way for the fallback behavior they want.

[mehdi_amini]

I forget, what was River's suggestion to make this more efficient?

There wasn't an efficiency aspect left, the idea River had was about generating a different base class with TableGen so that the subclass implementation for the interface fallback could be friendlier to write.

[jpienaar]

In D93085#2538095, @mehdi_amini wrote:

I forget, what was River's suggestion to make this more efficient?

There wasn't an efficiency aspect left, the idea River had was about generating a different base class with TableGen so that the subclass implementation for the interface fallback could be friendlier to write.

So just convenience generation pending?

[mehdi_amini]

In D93085#2599166, @jpienaar wrote:

In D93085#2538095, @mehdi_amini wrote:

I forget, what was River's suggestion to make this more efficient?

There wasn't an efficiency aspect left, the idea River had was about generating a different base class with TableGen so that the subclass implementation for the interface fallback could be friendlier to write.

So just convenience generation pending?

Yes.
I plan to get back to this in order to prototype a new TF control dialect.

[mehdi_amini]

Rebase

[mehdi_amini]

Address comments

mlir/test/lib/Dialect/Test/TestDialect.cpp

177

Yes, I updated the code to make it more clear.

181

Nothing, copy/pasted from somewhere else.

Removed!

206

It is a void*, making it a unique_ptr would require to declare some private classes. I'd expect something not in the "TestDialect" to use another more structured solution.

223

Yes!

[jpienaar]

Nice, thanks

[rriddle]

Can you add documentation for this?

mlir/include/mlir/IR/OpDefinition.h

1750

nit: Please use complete sentences in comments.

1751–1752

mlir/test/lib/Dialect/Test/TestDialect.cpp

168

Functions should be static and top-level.

206

You only have to forward declare them. You could even put the class in the dialect declaration:

// in .td
  let extraClassDeclaration = [{
    ...
  private:
    class CustomFallbackInterface;

    // Storage for a custom fallback interface.
    std::unique_ptr<CustomFallbackInterface> fallbackEffectOpInterfaces;
  }];

// in .cpp
class TestDialect::CustomFallbackInterface : public TestEffectOpInterface::Model<TestOpEffectInterfaceFallback> {};

mlir/tools/mlir-tblgen/DialectGen.cpp

147

You need to add ::mlir:: to TypeID and Operation to not break out-of-tree projects.

[mehdi_amini]

Use the TableGen mechanism

[mehdi_amini]

Address River's comment

mlir/test/lib/Dialect/Test/TestDialect.cpp

206

It won't work because the constructor body is in the header...

[mehdi_amini]

Add some doc in Interfaces.md

[rriddle]

mlir/docs/Interfaces.md

237 ↗

(On Diff #332061)

Do we need to pass in the operation here? Or can we just pass in the OperationName instead? We have several few places within the code base that may try to use the AbstractOperation to check for interfaces without having a concrete operation instead, it would be great not to lose that ability.

mlir/include/mlir/IR/OpDefinition.h

1755

Drop the getIdentifier here?

[rriddle]

(Marking as requested changes while discussing updated patch)

[mehdi_amini]

Fix doc: the fallback applies also to registered operations

[mehdi_amini]

Use an Identifier instead of an Operation* to retrieve the interface

mlir/docs/Interfaces.md

237 ↗

(On Diff #332061)

Yes the name should be enough here. The original version of the patched needed the operation here because it was captured in the returned instance IIRC.

[rriddle]

mlir/docs/Interfaces.md

237 ↗

(On Diff #332061)

Why Identifier and not OperationName?

[rriddle]

Only comment is on OperationName vs Identifier, but otherwise LGTM.

[mehdi_amini]

Use OperationName

[vinograd47]

Nice feature! I have one small question. Can this fallback OpInterface be represented as DialectInterface itself? In that case it can be registered in the Dialect via addInterfaces method without manual memory management (destructor overload or unique_ptr).

struct TestOpEffectInterfaceFallback
    : public DialectInterface::Base<TestEffectOpInterface::FallbackModel<TestOpEffectInterfaceFallback>> {
  ...
};


void TestDialect::initialize() {
  addInterfaces<TestOpEffectInterfaceFallback>();
}

void *TestDialect::getRegisteredInterfaceForOp(TypeID typeID,
                                               OperationName opName) {
  if (opName.getIdentifier() == "test.unregistered_side_effect_op" &&
      typeID == TypeID::get<TestEffectOpInterface>())
    return getRegisteredInterface<TestOpEffectInterfaceFallback>();
  return nullptr;
}

[mehdi_amini]

That's an excellent remark!
I intentionally did it this way because in my original mock-up I wanted to be able to initialize the interface instance with some state and keep it around. But it may not be warranted actually.

mlir/include/mlir/IR/Dialect.h

const DialectInterface *getRegisteredInterface(TypeID interfaceID) {
     auto it = registeredInterfaces.find(interfaceID);
     return it != registeredInterfaces.end() ? it->getSecond().get() : nullptr;
   }
   template <typename InterfaceT> const InterfaceT *getRegisteredInterface() {
     return static_cast<const InterfaceT *>(
         getRegisteredInterface(InterfaceT::getInterfaceID()));
   }
 
+  /// Lookup an op interface for the given ID if one is registered, otherwise
+  /// nullptr.
+  virtual void *getRegisteredInterfaceForOp(TypeID interfaceID, Operation *op) {
+    return nullptr;
+  }
+  template <typename InterfaceT>
+  typename InterfaceT::Concept *getRegisteredInterfaceForOp(Operation *op) {
+    return static_cast<typename InterfaceT::Concept *>(
+        getRegisteredInterfaceForOp(InterfaceT::getInterfaceID(), op));
+  }
+
 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".

mlir/include/mlir/IR/OpBase.td

class Dialect {
   bit hasOperationAttrVerify = 0;
 
   // If this dialect overrides the hook for verifying region argument
   // attributes.
   bit hasRegionArgAttrVerify = 0;
 
   // If this dialect overrides the hook for verifying region result attributes.
   bit hasRegionResultAttrVerify = 0;
+
+  // If this dialect overrides the hook for op interface fallback.
+  bit hasOperationInterfaceFallback = 0;
 }
 
 //===----------------------------------------------------------------------===//
 // Type definitions
 //===----------------------------------------------------------------------===//
 
 // A type, carries type constraints.
 class Type<Pred condition, string descr = ""> :

mlir/include/mlir/IR/OpDefinition.h

 // The purpose of these types are to allow light-weight implementation of
 // concrete ops (like DimOp) with very little boilerplate.
 //
 //===----------------------------------------------------------------------===//
 
 #ifndef MLIR_IR_OPDEFINITION_H
 #define MLIR_IR_OPDEFINITION_H
 
+#include "mlir/IR/Dialect.h"
 #include "mlir/IR/Operation.h"
 #include "llvm/Support/PointerLikeTypeTraits.h"
 
 #include <type_traits>
 
 namespace mlir {
 class Builder;
 class OpBuilder;
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) {
     // Access the raw interface from the abstract operation.
     auto *abstractOp = op->getAbstractOperation();
-    return abstractOp ? abstractOp->getInterface<ConcreteType>() : nullptr;
+    if (abstractOp) {
+      auto *opIface = abstractOp->getInterface<ConcreteType>();
+      if (opIface)
+        return opIface;
+    }
+    // Fallback to the dialect
+    if (op->getDialect())
+      return op->getDialect()->getRegisteredInterfaceForOp<ConcreteType>(op);
+    return nullptr;
   }
 
   /// Allow access to `getInterfaceFor`.
   friend InterfaceBase;
 };
 
 //===----------------------------------------------------------------------===//
 // Common Operation Folders/Parsers/Printers
 //===----------------------------------------------------------------------===//

mlir/include/mlir/Support/InterfaceSupport.h

public:
 
   /// Returns an instance of the concept object for the given interface if it
   /// was registered to this map, null otherwise.
   template <typename T> typename T::Concept *lookup() const {
     void *inst = interfaces ? interfaces->lookup(T::getInterfaceID()) : nullptr;
     return reinterpret_cast<typename T::Concept *>(inst);
   }
 
+  /// Returns an instance of the concept object for the given interface if it
+  /// was registered to this map, null otherwise.
+  void *lookup(TypeID typeID) const {
+    return interfaces ? interfaces->lookup(typeID) : nullptr;
+  }
+
 private:
   InterfaceMap() = default;
   InterfaceMap(MutableArrayRef<std::pair<TypeID, void *>> elements)
       : interfaces(std::make_unique<llvm::SmallDenseMap<TypeID, void *>>(
             elements.begin(), elements.end())) {}
 
   template <typename... Ts>
   static InterfaceMap getImpl(std::tuple<Ts...> *) {

mlir/include/mlir/TableGen/Dialect.h

public:
   bool hasOperationAttrVerify() const;
 
   /// Returns true if this dialect has a region argument attribute verifier.
   bool hasRegionArgAttrVerify() const;
 
   /// Returns true if this dialect has a region result attribute verifier.
   bool hasRegionResultAttrVerify() const;
 
+  /// Returns true if this dialect has fallback interfaces for its operations.
+  bool hasOperationInterfaceFallback() const;
+
   // Returns whether two dialects are equal by checking the equality of the
   // underlying record.
   bool operator==(const Dialect &other) const;
 
   bool operator!=(const Dialect &other) const { return !(*this == other); }
 
   // Compares two dialects by comparing the names of the dialects.
   bool operator<(const Dialect &other) const;

mlir/lib/TableGen/Dialect.cpp

 bool Dialect::hasRegionArgAttrVerify() const {
   return def->getValueAsBit("hasRegionArgAttrVerify");
 }
 
 bool Dialect::hasRegionResultAttrVerify() const {
   return def->getValueAsBit("hasRegionResultAttrVerify");
 }
 
+bool Dialect::hasOperationInterfaceFallback() const {
+  return def->getValueAsBit("hasOperationInterfaceFallback");
+}
+
 bool Dialect::operator==(const Dialect &other) const {
   return def == other.def;
 }
 
 bool Dialect::operator<(const Dialect &other) const {
   return getName() < other.getName();
 }

mlir/test/IR/test-side-effects.mlir

 ]} : () -> i32
 
 // No _memory_ effects, but a parametric test effect.
 // expected-remark@+2 {{operation has no memory effects}}
 // expected-remark@+1 {{found a parametric effect with affine_map<(d0, d1) -> (d1, d0)>}}
 %4 = "test.side_effect_op"() {
   effect_parameter = affine_map<(i, j) -> (j, i)>
 } : () -> i32
+
+// Check with this unregistered operation to test the fallback on the dialect.
+// expected-remark@+1 {{found a parametric effect with affine_map<(d0, d1) -> (d1, d0)>}}
+%5 = "test.unregistered_side_effect_op"() {
+  effect_parameter = affine_map<(i, j) -> (j, i)>
+} : () -> i32

mlir/test/lib/Dialect/Test/TestDialect.cpp

struct TestInlinerInterface : public DialectInlinerInterface {
   }
 };
 } // end anonymous namespace
 
 //===----------------------------------------------------------------------===//
 // TestDialect
 //===----------------------------------------------------------------------===//
 
+namespace {
+void testSideEffectOpGetEffect(
+    Operation *op, SmallVectorImpl<::mlir::SideEffects::EffectInstance<
+                       ::mlir::TestEffects::Effect>> &effects);
+}
+
+// This is the implementation of a dialect fallback for `TestEffectOpInterface`.
+struct TestOpEffectInterfaceFallback {
+  TestOpEffectInterfaceFallback(Operation *op) : op(op) {}
+  static bool classof(Operation *op) {
+    assert(op->getName().getStringRef() == ==
+           "test.unregistered_side_effect_op");
+    return true;
+  }
+
+  void getEffects(SmallVectorImpl<::mlir::SideEffects::EffectInstance<
+                      ::mlir::TestEffects::Effect>> &effects) {
+    testSideEffectOpGetEffect(op, effects);
+  }
+
+private:
+  Operation *op;
+};
+
 void TestDialect::initialize() {
   addOperations<
 #define GET_OP_LIST
 #include "TestOps.cpp.inc"
       >();
   addInterfaces<TestOpAsmInterface, TestDialectFoldInterface,
                 TestInlinerInterface>();
   addTypes<TestType, TestRecursiveType,
 #define GET_TYPEDEF_LIST
 #include "TestTypeDefs.cpp.inc"
            >();
   allowUnknownOperations();
+
+  // Instanciate our fallback op interface that we'll use on specific
+  // unregistered op.
+  fallbackEffectOpInterfaces =
+      new TestEffectOpInterface::Model<TestOpEffectInterfaceFallback>();
+}
+TestDialect::~TestDialect() {
+  delete static_cast<
+      TestEffectOpInterface::Model<TestOpEffectInterfaceFallback> *>(
+      fallbackEffectOpInterfaces);
 }
 
 Operation *TestDialect::materializeConstant(OpBuilder &builder, Attribute value,
                                             Type type, Location loc) {
   return builder.create<TestOpConstant>(loc, type, value);
 }
 
+void *TestDialect::getRegisteredInterfaceForOp(mlir::TypeID typeID,
+                                               mlir::Operation *op) {
+  StringRef opName = op->getName().getStringRef();
+  if (opName == "test.unregistered_side_effect_op" &&
+      typeID == TypeID::get<TestEffectOpInterface>())
+    return fallbackEffectOpInterfaces;
+  return nullptr;
+}
+
 static Type parseTestType(MLIRContext *ctxt, DialectAsmParser &parser,
                           llvm::SetVector<Type> &stack) {
   StringRef typeTag;
   if (failed(parser.parseKeyword(&typeTag)))
     return Type();
 
   auto genType = generatedTypeParser(ctxt, parser, typeTag);
   if (genType != Type())

 // Test SideEffect interfaces
 //===----------------------------------------------------------------------===//
 
 namespace {
 /// A test resource for side effects.
 struct TestResource : public SideEffects::Resource::Base<TestResource> {
   StringRef getName() final { return "<Test>"; }
 };
+
+void testSideEffectOpGetEffect(
+    Operation *op, SmallVectorImpl<::mlir::SideEffects::EffectInstance<
+                       ::mlir::TestEffects::Effect>> &effects) {
+  auto effectsAttr = op->getAttrOfType<AffineMapAttr>("effect_parameter");
+  if (!effectsAttr)
+    return;
+
+  effects.emplace_back(TestEffects::Concrete::get(), effectsAttr);
+}
 } // end anonymous namespace
 
 void SideEffectOp::getEffects(
     SmallVectorImpl<MemoryEffects::EffectInstance> &effects) {
   // Check for an effects attribute on the op instance.
-  ArrayAttr effectsAttr = (*this)->getAttrOfType<ArrayAttr>("effects");
+  ArrayAttr effectsAttr = this->getAttrOfType<ArrayAttr>("effects");
   if (!effectsAttr)
     return;
 
   // If there is one, it is an array of dictionary attributes that hold
   // information on the effects of this operation.
   for (Attribute element : effectsAttr) {
     DictionaryAttr effectElement = element.cast<DictionaryAttr>();
 
else if (Attribute ref = effectElement.get("on_reference"))
       effects.emplace_back(effect, ref.cast<SymbolRefAttr>(), resource);
     else
       effects.emplace_back(effect, resource);
   }
 }
 
 void SideEffectOp::getEffects(
     SmallVectorImpl<TestEffects::EffectInstance> &effects) {
-  auto effectsAttr = (*this)->getAttrOfType<AffineMapAttr>("effect_parameter");
-  if (!effectsAttr)
-    return;
-
-  effects.emplace_back(TestEffects::Concrete::get(), effectsAttr);
+  testSideEffectOpGetEffect(getOperation(), effects);
 }
 
 //===----------------------------------------------------------------------===//
 // StringAttrPrettyNameOp
 //===----------------------------------------------------------------------===//
 
 // This op has fancy handling of its SSA result name.
 static ParseResult parseStringAttrPrettyNameOp(OpAsmParser &parser,

mlir/test/lib/Dialect/Test/TestOps.td

 
 def Test_Dialect : Dialect {
   let name = "test";
   let cppNamespace = "::mlir::test";
   let hasConstantMaterializer = 1;
   let hasOperationAttrVerify = 1;
   let hasRegionArgAttrVerify = 1;
   let hasRegionResultAttrVerify = 1;
+  let hasOperationInterfaceFallback = 1;
+  code extraClassDeclaration = [{
+    ~TestDialect();
+  private:
+    // Storage for a custom fallback interface.
+    void *fallbackEffectOpInterfaces;
+  }];
 }
 
 class TEST_Op<string mnemonic, list<OpTrait> traits = []> :
     Op<Test_Dialect, mnemonic, traits>;
 
 //===----------------------------------------------------------------------===//
 // Test Types
 //===----------------------------------------------------------------------===//

mlir/tools/mlir-tblgen/DialectGen.cpp

 static const char *const regionResultAttrVerifierDecl = R"(
     /// Provides a hook for verifying dialect attributes attached to the given
     /// op's region result.
     ::mlir::LogicalResult verifyRegionResultAttribute(
         ::mlir::Operation *op, unsigned regionIndex, unsigned resultIndex,
         ::mlir::NamedAttribute attribute) override;
 )";
 
+/// The code block for the op interface fallback hook.
+static const char *const operationInterfaceFallbackDecl = R"(
+    /// Provides a hook for op interface.
+    void *getRegisteredInterfaceForOp(TypeID interfaceID,
+                                      Operation *op) override;
+)";
+
 /// Generate the declaration for the given dialect class.
 static void emitDialectDecl(Dialect &dialect,
                             iterator_range<DialectFilterIterator> dialectAttrs,
                             iterator_range<DialectFilterIterator> dialectTypes,
                             raw_ostream &os) {
   /// Build the list of dependent dialects
   std::string dependentDialectRegistrations;
   {
static void emitDialectDecl(Dialect &dialect,
   if (dialect.hasConstantMaterializer())
     os << constantMaterializerDecl;
   if (dialect.hasOperationAttrVerify())
     os << opAttrVerifierDecl;
   if (dialect.hasRegionArgAttrVerify())
     os << regionArgAttrVerifierDecl;
   if (dialect.hasRegionResultAttrVerify())
     os << regionResultAttrVerifierDecl;
+  if (dialect.hasOperationInterfaceFallback())
+    os << operationInterfaceFallbackDecl;
   if (llvm::Optional<StringRef> extraDecl = dialect.getExtraClassDeclaration())
     os << *extraDecl;
 
   // End the dialect decl.
   os << "};\n";
 }
 
 static bool emitDialectDecls(const llvm::RecordKeeper &recordKeeper,