[mlir] Move SymbolOpInterfaces "classof" check to a proper "extraClassOf" interface field

[rriddle]

Summary

SymbolOpInterface overrides the base classof to provide support
for optionally implementing the interface. This is currently placed
in the extraClassDeclarations, but that is kind of awkard given that
it requires underlying knowledge of how the base classof is implemented.
This commit adds a proper "extraClassOf" field to allow interfaces to
implement this, which abstracts away the default classof logic.

[jpienaar]

mlir/include/mlir/IR/OpBase.td

2052

This is not arbitrary logic, this requires to end in a return that reports whether a given op implements this interface.

Here we rely on -> for accessing the operation being checked and so need nothing else like $_op ? it may be good to document this. It feels partially weird/unsafe to say its an optional interface which you check if implemented by calling functions on that interface.

mlir/include/mlir/IR/OpBase.td

class Interface<string name> {
 
   // An optional code block containing extra declarations to place in the
   // interface declaration.
   code extraClassDeclaration = "";
 
   // An optional code block containing extra declarations to place in both
   // the interface and trait declaration.
   code extraSharedClassDeclaration = "";
+
+  // An optional code block for adding additional "classof" logic. This can
+  // 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.
+  code extraClassOf = "";
 }
 
 // AttrInterface represents an interface registered to an attribute.
 class AttrInterface<string name> : Interface<name>, InterfaceTrait<name>,
   Attr<CPred<"$_self.isa<"
     # !if(!empty(cppNamespace),"", cppNamespace # "::") # name # ">()">,
       name # " instance">
 {

mlir/include/mlir/IR/SymbolInterfaces.td

let verify = [{
     if (::mlir::failed(::mlir::detail::verifySymbol($_op)))
       return ::mlir::failure();
     if (concreteOp.isDeclaration() && concreteOp.isPublic())
       return concreteOp.emitOpError("symbol declaration cannot have public "
              "visibility");
     return success();
   }];
 
-  let extraClassDeclaration = [{
-    /// Convenience version of `getNameAttr` that returns a StringRef.
-    StringRef getName() {
-      return getNameAttr().getValue();
-    }
-
-    /// Convenience version of `setName` that take a StringRef.
-    void setName(StringRef name) {
-      setName(StringAttr::get(this->getContext(), name));
-    }
-
-    /// Custom classof that handles the case where the symbol is optional.
-    static bool classof(Operation *op) {
-      auto *opConcept = getInterfaceFor(op);
-      if (!opConcept)
-        return false;
-      return !opConcept->isOptionalSymbol(opConcept, op) ||
-             op->getAttr(::mlir::SymbolTable::getSymbolAttrName());
-    }
-  }];
-
-  let extraTraitClassDeclaration = [{
+  let extraSharedClassDeclaration = [{
     using Visibility = mlir::SymbolTable::Visibility;
 
     /// Convenience version of `getNameAttr` that returns a StringRef.
     StringRef getName() {
       return getNameAttr().getValue();
     }
 
     /// Convenience version of `setName` that take a StringRef.
     void setName(StringRef name) {
       setName(StringAttr::get($_op->getContext(), name));
     }
   }];
+
+  // Add additional classof checks to properly handle "optional" symbols.
+  let extraClassOf = [{
+    return $_op->hasAttr(::mlir::SymbolTable::getSymbolAttrName());
+  }];
 }
 
 //===----------------------------------------------------------------------===//
 // SymbolUserOpInterface
 //===----------------------------------------------------------------------===//
 
 def SymbolUserOpInterface : OpInterface<"SymbolUserOpInterface"> {
   let description = [{

mlir/include/mlir/Support/InterfaceSupport.h

template <typename T,
             std::enable_if_t<std::is_base_of<Trait<T>, T>::value> * = nullptr>
   Interface(T t)
       : BaseType(t),
         conceptImpl(t ? ConcreteType::getInterfaceFor(t) : nullptr) {
     assert((!t || conceptImpl) &&
            "expected value to provide interface instance");
   }
 
+  /// Constructor for a known concept.
+  Interface(ValueT t, Concept *conceptImpl)
+      : BaseType(t), conceptImpl(conceptImpl) {
+    assert(!t || ConcreteType::getInterfaceFor(t) == conceptImpl);
+  }
+
   /// Constructor for DenseMapInfo's empty key and tombstone key.
   Interface(ValueT t, std::nullptr_t) : BaseType(t), conceptImpl(nullptr) {}
 
   /// Support 'classof' by checking if the given object defines the concrete
   /// interface.
   static bool classof(ValueT t) { return ConcreteType::getInterfaceFor(t); }
 
   /// Define an accessor for the ID of this interface.

mlir/include/mlir/TableGen/Format.h

 /// same unified setter.
 class FmtContext {
 public:
   // Placeholder kinds
   enum class PHKind : char {
     None,
     Custom,  // For custom placeholders
     Builder, // For the $_builder placeholder
-    Op,      // For the $_op placeholder
     Self,    // For the $_self placeholder
   };
 
   FmtContext() = default;
 
   // Create a format context with a list of substitutions.
   FmtContext(ArrayRef<std::pair<StringRef, StringRef>> subs);
 
   // Setter for custom placeholders
   FmtContext &addSubst(StringRef placeholder, const Twine &subst);
 
   // Setters for builtin placeholders
   FmtContext &withBuilder(Twine subst);
-  FmtContext &withOp(Twine subst);
   FmtContext &withSelf(Twine subst);
 
   std::optional<StringRef> getSubstFor(PHKind placeholder) const;
   std::optional<StringRef> getSubstFor(StringRef placeholder) const;
 
   static PHKind getPlaceHolderKind(StringRef str);
 
 private:

mlir/include/mlir/TableGen/Interfaces.h

public:
 
   // Return the traits extra class declaration code.
   std::optional<StringRef> getExtraTraitClassDeclaration() const;
 
   // Return the extra class declaration code shared between the interface and
   // trait classes.
   std::optional<StringRef> getExtraSharedClassDeclaration() const;
 
+  // Return the extra classof method code.
+  std::optional<StringRef> getExtraClassOf() const;
+
   // Return the verify method body if it has one.
   std::optional<StringRef> getVerify() const;
 
   // If there's a verify method, return if it needs to access the ops in the
   // regions.
   bool verifyWithRegions() const;
 
   // Returns the Tablegen definition this interface was constructed from.

mlir/lib/TableGen/CodeGenHelpers.cpp

static ::mlir::LogicalResult {0}(
   return ::mlir::success();
 }
 )";
 
 void StaticVerifierFunctionEmitter::emitConstraints(
     const ConstraintMap &constraints, StringRef selfName,
     const char *const codeTemplate) {
   FmtContext ctx;
-  ctx.withOp("*op").withSelf(selfName);
+  ctx.addSubst("_op", "*op").withSelf(selfName);
   for (auto &it : constraints) {
     os << formatv(codeTemplate, it.second,
                   tgfmt(it.first.getConditionTemplate(), &ctx),
                   escapeString(it.first.getSummary()));
   }
 }
 
 void StaticVerifierFunctionEmitter::emitTypeConstraints() {

 }
 
 void StaticVerifierFunctionEmitter::emitRegionConstraints() {
   emitConstraints(regionConstraints, "region", regionConstraintCode);
 }
 
 void StaticVerifierFunctionEmitter::emitPatternConstraints() {
   FmtContext ctx;
-  ctx.withOp("*op").withBuilder("rewriter").withSelf("type");
+  ctx.addSubst("_op", "*op").withBuilder("rewriter").withSelf("type");
   for (auto &it : typeConstraints) {
     os << formatv(patternAttrOrTypeConstraintCode, it.second,
                   tgfmt(it.first.getConditionTemplate(), &ctx),
                   escapeString(it.first.getSummary()), "Type type");
   }
   ctx.withSelf("attr");
   for (auto &it : attrConstraints) {
     os << formatv(patternAttrOrTypeConstraintCode, it.second,
                   tgfmt(it.first.getConditionTemplate(), &ctx),
                   escapeString(it.first.getSummary()), "Attribute attr");
   }
 }
 
 //===----------------------------------------------------------------------===//
 // Constraint Uniquing
 
 /// An attribute constraint that references anything other than itself and the
 /// current op cannot be generically extracted into a function. Most
 /// prohibitive are operands and results, which require calls to
 /// `getODSOperands` or `getODSResults`. Attribute references are tricky too
 /// because ops use cached identifiers.
 static bool canUniqueAttrConstraint(Attribute attr) {
   FmtContext ctx;
-  auto test =
-      tgfmt(attr.getConditionTemplate(), &ctx.withSelf("attr").withOp("*op"))
+  auto test = tgfmt(attr.getConditionTemplate(),
+                    &ctx.withSelf("attr").addSubst("_op", "*op"))
                   .str();
   return !StringRef(test).contains("<no-subst-found>");
 }
 
 std::string StaticVerifierFunctionEmitter::getUniqueName(StringRef kind,
                                                          unsigned index) {
   return ("__mlir_ods_local_" + kind + "_constraint_" + uniqueOutputLabel +
           Twine(index))
       .str();

mlir/lib/TableGen/Format.cpp

FmtContext &FmtContext::addSubst(StringRef placeholder, const Twine &subst) {
   return *this;
 }
 
 FmtContext &FmtContext::withBuilder(Twine subst) {
   builtinSubstMap[PHKind::Builder] = subst.str();
   return *this;
 }
 
-FmtContext &FmtContext::withOp(Twine subst) {
-  builtinSubstMap[PHKind::Op] = subst.str();
-  return *this;
-}
-
 FmtContext &FmtContext::withSelf(Twine subst) {
   builtinSubstMap[PHKind::Self] = subst.str();
   return *this;
 }
 
 std::optional<StringRef>
 FmtContext::getSubstFor(FmtContext::PHKind placeholder) const {
   if (placeholder == FmtContext::PHKind::None ||
std::optional<StringRef> FmtContext::getSubstFor(StringRef placeholder) const {
   if (it == customSubstMap.end())
     return {};
   return StringRef(it->second);
 }
 
 FmtContext::PHKind FmtContext::getPlaceHolderKind(StringRef str) {
   return StringSwitch<FmtContext::PHKind>(str)
       .Case("_builder", FmtContext::PHKind::Builder)
-      .Case("_op", FmtContext::PHKind::Op)
       .Case("_self", FmtContext::PHKind::Self)
       .Case("", FmtContext::PHKind::None)
       .Default(FmtContext::PHKind::Custom);
 }
 
 std::pair<FmtReplacement, StringRef>
 FmtObjectBase::splitFmtSegment(StringRef fmt) {
   size_t begin = fmt.find_first_of('$');

mlir/lib/TableGen/Interfaces.cpp

 }
 
 // Return the shared extra class declaration code.
 std::optional<StringRef> Interface::getExtraSharedClassDeclaration() const {
   auto value = def->getValueAsString("extraSharedClassDeclaration");
   return value.empty() ? std::optional<StringRef>() : value;
 }
 
+std::optional<StringRef> Interface::getExtraClassOf() const {
+  auto value = def->getValueAsString("extraClassOf");
+  return value.empty() ? std::optional<StringRef>() : value;
+}
+
 // Return the body for this method if it has one.
 std::optional<StringRef> Interface::getVerify() const {
   // Only OpInterface supports the verify method.
   if (!isa<OpInterface>(this))
     return std::nullopt;
   auto value = def->getValueAsString("verify");
   return value.empty() ? std::optional<StringRef>() : value;
 }

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

 // RUN: mlir-tblgen -gen-op-interface-decls -I %S/../../include %s | FileCheck %s --check-prefix=DECL
 // RUN: mlir-tblgen -gen-op-decls -I %S/../../include %s | FileCheck %s --check-prefix=OP_DECL
 // RUN: mlir-tblgen -gen-op-interface-docs -I %S/../../include %s | FileCheck %s --check-prefix=DOCS
 
 include "mlir/IR/OpBase.td"
 
+def ExtraClassOfInterface : OpInterface<"ExtraClassOfInterface"> {
+  let extraClassOf = "return $_op->someOtherMethod();";
+}
+
+// DECL: class ExtraClassOfInterface
+// DECL:   static bool classof(::mlir::Operation * base) {
+// DECL-NEXT:     if (!getInterfaceFor(base))
+// DECL-NEXT:       return false;
+// DECL-NEXT:     return base->someOtherMethod();
+// DECL-NEXT:   }
+
 def ExtraShardDeclsInterface : OpInterface<"ExtraShardDeclsInterface"> {
   let extraSharedClassDeclaration = [{
     bool sharedMethodDeclaration() {
       return $_op.someOtherMethod();
     }
   }];
 }
 

mlir/tools/mlir-tblgen/OpDefinitionsGen.cpp

 
 OpEmitter::OpEmitter(const Operator &op,
                      const StaticVerifierFunctionEmitter &staticVerifierEmitter)
     : def(op.getDef()), op(op),
       opClass(op.getCppClassName(), op.getExtraClassDeclaration(),
               formatExtraDefinitions(op)),
       staticVerifierEmitter(staticVerifierEmitter),
       emitHelper(op, /*emitForOp=*/true) {
-  verifyCtx.withOp("(*this->getOperation())");
+  verifyCtx.addSubst("_op", "(*this->getOperation())");
   verifyCtx.addSubst("_ctxt", "this->getOperation()->getContext()");
 
   genTraits();
 
   // Generate C++ code for various op methods. The order here determines the
   // methods in the generated file.
   genAttrNameGetters();
   genOpAsmInterface();

mlir/tools/mlir-tblgen/OpInterfacesGen.cpp

protected:
   // The stream to emit to.
   raw_ostream &os;
   /// The C++ value type of the interface, e.g. Operation*.
   StringRef valueType;
   /// The C++ base interface type.
   StringRef interfaceBaseType;
   /// The name of the typename for the value template.
   StringRef valueTemplate;
+  /// The name of the substituion variable for the value.
+  StringRef substVar;
   /// The format context to use for methods.
   tblgen::FmtContext nonStaticMethodFmt;
   tblgen::FmtContext traitMethodFmt;
   tblgen::FmtContext extraDeclsFmt;
 };
 
 /// A specialized generator for attribute interfaces.
 struct AttrInterfaceGenerator : public InterfaceGenerator {
   AttrInterfaceGenerator(const llvm::RecordKeeper &records, raw_ostream &os)
       : InterfaceGenerator(getAllInterfaceDefinitions(records, "Attr"), os) {
     valueType = "::mlir::Attribute";
     interfaceBaseType = "AttributeInterface";
     valueTemplate = "ConcreteAttr";
+    substVar = "_attr";
     StringRef castCode = "(tablegen_opaque_val.cast<ConcreteAttr>())";
-    nonStaticMethodFmt.addSubst("_attr", castCode).withSelf(castCode);
-    traitMethodFmt.addSubst("_attr",
+    nonStaticMethodFmt.addSubst(substVar, castCode).withSelf(castCode);
+    traitMethodFmt.addSubst(substVar,
                             "(*static_cast<const ConcreteAttr *>(this))");
-    extraDeclsFmt.addSubst("_attr", "(*this)");
+    extraDeclsFmt.addSubst(substVar, "(*this)");
   }
 };
 /// A specialized generator for operation interfaces.
 struct OpInterfaceGenerator : public InterfaceGenerator {
   OpInterfaceGenerator(const llvm::RecordKeeper &records, raw_ostream &os)
       : InterfaceGenerator(getAllInterfaceDefinitions(records, "Op"), os) {
     valueType = "::mlir::Operation *";
     interfaceBaseType = "OpInterface";
     valueTemplate = "ConcreteOp";
+    substVar = "_op";
     StringRef castCode = "(llvm::cast<ConcreteOp>(tablegen_opaque_val))";
     nonStaticMethodFmt.addSubst("_this", "impl")
-        .withOp(castCode)
+        .addSubst(substVar, castCode)
         .withSelf(castCode);
-    traitMethodFmt.withOp("(*static_cast<ConcreteOp *>(this))");
-    extraDeclsFmt.withOp("(*this)");
+    traitMethodFmt.addSubst(substVar, "(*static_cast<ConcreteOp *>(this))");
+    extraDeclsFmt.addSubst(substVar, "(*this)");
   }
 };
 /// A specialized generator for type interfaces.
 struct TypeInterfaceGenerator : public InterfaceGenerator {
   TypeInterfaceGenerator(const llvm::RecordKeeper &records, raw_ostream &os)
       : InterfaceGenerator(getAllInterfaceDefinitions(records, "Type"), os) {
     valueType = "::mlir::Type";
     interfaceBaseType = "TypeInterface";
     valueTemplate = "ConcreteType";
+    substVar = "_type";
     StringRef castCode = "(tablegen_opaque_val.cast<ConcreteType>())";
-    nonStaticMethodFmt.addSubst("_type", castCode).withSelf(castCode);
-    traitMethodFmt.addSubst("_type",
+    nonStaticMethodFmt.addSubst(substVar, castCode).withSelf(castCode);
+    traitMethodFmt.addSubst(substVar,
                             "(*static_cast<const ConcreteType *>(this))");
-    extraDeclsFmt.addSubst("_type", "(*this)");
+    extraDeclsFmt.addSubst(substVar, "(*this)");
   }
 };
 } // namespace
 
 //===----------------------------------------------------------------------===//
 // GEN: Interface definitions
 //===----------------------------------------------------------------------===//
 
for (auto &method : interface.getMethods()) {
     os << " {\n      " << tblgen::tgfmt(defaultImpl->trim(), &traitMethodFmt)
        << "\n    }\n";
   }
 
   if (auto verify = interface.getVerify()) {
     assert(isa<OpInterface>(interface) && "only OpInterface supports 'verify'");
 
     tblgen::FmtContext verifyCtx;
-    verifyCtx.withOp("op");
+    verifyCtx.addSubst("_op", "op");
     os << llvm::formatv(
               "    static ::mlir::LogicalResult {0}(::mlir::Operation *op) ",
               (interface.verifyWithRegions() ? "verifyRegionTrait"
                                              : "verifyTrait"))
        << "{\n      " << tblgen::tgfmt(verify->trim(), &verifyCtx)
        << "\n    }\n";
   }
   if (auto extraTraitDecls = interface.getExtraTraitClassDeclaration())
void InterfaceGenerator::emitInterfaceDecl(const Interface &interface) {
   // Emit any extra declarations.
   if (std::optional<StringRef> extraDecls =
           interface.getExtraClassDeclaration())
     os << *extraDecls << "\n";
   if (std::optional<StringRef> extraDecls =
           interface.getExtraSharedClassDeclaration())
     os << tblgen::tgfmt(*extraDecls, &extraDeclsFmt);
 
+  // Emit classof code if necessary.
+  if (std::optional<StringRef> extraClassOf = interface.getExtraClassOf()) {
+    auto extraClassOfFmt = tblgen::FmtContext();
+    extraClassOfFmt.addSubst(substVar, "base");
+    os << "  static bool classof(" << valueType << " base) {\n"
+       << "    if (!getInterfaceFor(base))\n"
+          "      return false;\n"
+       << "    " << tblgen::tgfmt(extraClassOf->trim(), &extraClassOfFmt)
+       << "\n  }\n";
+  }
+
   os << "};\n";
 
   os << "namespace detail {\n";
   emitTraitDecl(interface, interfaceName, interfaceTraitsName);
   os << "}// namespace detail\n";
 
   for (StringRef ns : llvm::reverse(namespaces))
     os << "} // namespace " << ns << "\n";

mlir/unittests/TableGen/FormatTest.cpp

 
 TEST(FormatTest, PlaceHolderFmtStrWithBuilder) {
   FmtContext ctx;
 
   std::string result = std::string(tgfmt("$_builder", &ctx.withBuilder("bbb")));
   EXPECT_THAT(result, StrEq("bbb"));
 }
 
-TEST(FormatTest, PlaceHolderFmtStrWithOp) {
-  FmtContext ctx;
-  std::string result = std::string(tgfmt("$_op", &ctx.withOp("ooo")));
-  EXPECT_THAT(result, StrEq("ooo"));
-}
-
 TEST(FormatTest, PlaceHolderMissingCtx) {
   std::string result = std::string(tgfmt("$_op", nullptr));
   EXPECT_THAT(result, StrEq("$_op<no-subst-found>"));
 }
 
 TEST(FormatTest, PlaceHolderMissingSubst) {
   FmtContext ctx;
   std::string result = std::string(tgfmt("$_op", &ctx.withBuilder("builder")));