[mlir] Make classof substitution in interface use an instance

mlir/include/mlir/IR/Interfaces.td

@@ -114,7 +114,7 @@ class Interface<string name, list<Interface> baseInterfacesArg = []> {
   // be used to better enable "optional" interfaces, where an entity only
   // implements the interface if some dynamic characteristic holds.
   // `$_attr`/`$_op`/`$_type` may be used to refer to an instance of the
-  // entity being checked.
+  // interface instance being checked.
   code extraClassOf = "";
 
   // An optional set of base interfaces that this interface

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

@@ -134,4 +134,17 @@ class TestEffects<list<TestEffect> effects = []>
 
 def TestConcreteEffect : TestEffect<"TestEffects::Concrete">;
 
+def TestOptionallyImplementedOpInterface
+    : OpInterface<"TestOptionallyImplementedOpInterface"> {
+  let cppNamespace = "::mlir";
+
+  let methods = [
+    InterfaceMethod<"", "bool", "getImplementsInterface", (ins)>,
+  ];
+
+  let extraClassOf = [{
+    return $_op.getImplementsInterface();
+  }];
+}
+
 #endif // MLIR_TEST_DIALECT_TEST_INTERFACES

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

@@ -2952,4 +2952,10 @@ def TestStoreWithARegionTerminator : TEST_Op<"store_with_a_region_terminator",
   let assemblyFormat = "attr-dict";
 }
 
+def TestOpOptionallyImplementingInterface
+    : TEST_Op<"op_optionally_implementing_interface",
+        [TestOptionallyImplementedOpInterface]> {
+  let arguments = (ins BoolAttr:$implementsInterface);
+}
+
 #endif // TEST_OPS

mlir/test/mlir-tblgen/op-interface.td

@@ -11,9 +11,11 @@ def ExtraClassOfInterface : OpInterface<"ExtraClassOfInterface"> {
 
 // DECL: class ExtraClassOfInterface
 // DECL:   static bool classof(::mlir::Operation * base) {
-// DECL-NEXT:     if (!getInterfaceFor(base))
+// DECL-NEXT:     auto* concept = getInterfaceFor(base);
+// DECL-NEXT:     if (!concept)
 // DECL-NEXT:       return false;
-// DECL-NEXT:     return base->someOtherMethod();
+// DECL-NEXT:     ExtraClassOfInterface odsInterfaceInstance(base, concept);
+// DECL-NEXT:     return odsInterfaceInstance->someOtherMethod();
 // DECL-NEXT:   }
 
 def ExtraShardDeclsInterface : OpInterface<"ExtraShardDeclsInterface"> {

mlir/tools/mlir-tblgen/OpInterfacesGen.cpp

@@ -582,10 +582,12 @@ void InterfaceGenerator::emitInterfaceDecl(const Interface &interface) {
   // Emit classof code if necessary.
   if (std::optional<StringRef> extraClassOf = interface.getExtraClassOf()) {
     auto extraClassOfFmt = tblgen::FmtContext();
-    extraClassOfFmt.addSubst(substVar, "base");
+    extraClassOfFmt.addSubst(substVar, "odsInterfaceInstance");
     os << "  static bool classof(" << valueType << " base) {\n"
-       << "    if (!getInterfaceFor(base))\n"
+       << "    auto* concept = getInterfaceFor(base);\n"
+       << "    if (!concept)\n"
           "      return false;\n"
+          "    " << interfaceName << " odsInterfaceInstance(base, concept);\n"
        << "    " << tblgen::tgfmt(extraClassOf->trim(), &extraClassOfFmt)
        << "\n  }\n";
   }

mlir/unittests/IR/InterfaceTest.cpp

@@ -57,3 +57,15 @@ TEST(InterfaceTest, TypeInterfaceDenseMapKey) {
   EXPECT_TRUE(typeSet.contains(type2));
   EXPECT_FALSE(typeSet.contains(type3));
 }
+
+TEST(InterfaceTest, TestCustomClassOf) {
+  MLIRContext context;
+  context.loadDialect<test::TestDialect>();
+
+  OpBuilder builder(&context);
+  auto op = builder.create<TestOpOptionallyImplementingInterface>(
+      builder.getUnknownLoc(), /*implementsInterface=*/true);
+  EXPECT_TRUE(isa<TestOptionallyImplementedOpInterface>(*op));
+  op.setImplementsInterface(false);
+  EXPECT_FALSE(isa<TestOptionallyImplementedOpInterface>(*op));
+}