[mlir][Symbol] Change Symbol from a Trait into an OpInterface.

[rriddle]

Summary

This provides a much cleaner interface into Symbols, and allows for users to start injecting op-specific information. For example, derived op can now inject when a symbol can be discarded if use_empty. This would let us drop unused external functions, which generally have public visibility.

This revision also adds a new extraTraitClassDeclaration field to the ODS OpInterface class to allow for injecting declarations into the trait class that gets attached to the operations.

Depends On D78447

[bondhugula]

mlir/include/mlir/IR/SymbolInterfaces.td

26

Can an operation ever not have a name? Did you mean to say "A Symbol operation resides ..."?

106

"UseEmpty" seems to be an artifact of use_empty() - not perhaps a great fit here. canDiscardOnNoUses / canDiscardOnZeroUses?

123

Why is this still under a trait implementation namespace? On another note, I think verifySymbol is missing a zero region check and the code is assuming there is at least one region.

mlir/include/mlir/IR/SymbolTable.h

232–234

This information has I think gone missing from your operation interface doc comment.

mlir/lib/Transforms/SymbolDCE.cpp

81

Are you verifying somewhere on the op interface that a SymbolTableOp has at least one region?

[rriddle]

Resolve comments

[rriddle]

mlir/include/mlir/IR/SymbolInterfaces.td

106

This is intended to match the same verbiage as in LLVM.

123

Moved to mlir::detail. verifySymbolTable is the one that checks for regions, symbols don't need to have regions(e.g. global variables).

mlir/include/mlir/IR/SymbolTable.h

232–234

I added the link to SymbolsAndSymbolTables, which contains information on the constraints of symbols. Made an explicit mention of constraints.

mlir/lib/Transforms/SymbolDCE.cpp

81

Yes, this is checked as part of verifySymbolTable.

[rriddle]

Rebase

[mehdi_amini]

Can you document in the commit message the addition of extraTraitClassDeclaration? This seems like a new ODS feature independent of the symbol table change (someone using git blame in the future would end up in this commit)

[rriddle]

Rebase

flang/include/flang/Optimizer/Dialect/FIROps.td

 //===-- FIROps.td - FIR operation definitions --------------*- tablegen -*-===//
 //
 // 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
 //
 //===----------------------------------------------------------------------===//
 ///
 /// \file
 /// Definition of the FIR dialect operations
 ///
 //===----------------------------------------------------------------------===//
 
 #ifndef FIR_DIALECT_FIR_OPS
 #define FIR_DIALECT_FIR_OPS
 
+include "mlir/IR/SymbolInterfaces.td"
 include "mlir/Interfaces/ControlFlowInterfaces.td"
 include "mlir/Interfaces/SideEffects.td"
 
 def fir_Dialect : Dialect {
   let name = "fir";
 }
 
 // Types and predicates

mlir/include/mlir/Dialect/GPU/GPUOps.td

 //
 //===----------------------------------------------------------------------===//
 
 #ifndef GPU_OPS
 #define GPU_OPS
 
 include "mlir/Dialect/GPU/GPUBase.td"
 include "mlir/Dialect/LLVMIR/LLVMOpBase.td"
+include "mlir/IR/SymbolInterfaces.td"
 include "mlir/Interfaces/SideEffects.td"
 
 // Type constraint accepting standard integers, indices and wrapped LLVM integer
 // types.
 def IntLikeOrLLVMInt : TypeConstraint<
   Or<[AnySignlessInteger.predicate, Index.predicate, LLVMInt.predicate]>,
   "integer, index or LLVM dialect equivalent">;
 

mlir/include/mlir/Dialect/LLVMIR/LLVMOps.td

 //===-- LLVMOps.td - LLVM IR dialect op definition file ----*- tablegen -*-===//
 //
 // 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 is the LLVM IR operation definition file.
 //
 //===----------------------------------------------------------------------===//
 
 #ifndef LLVMIR_OPS
 #define LLVMIR_OPS
 
 include "mlir/Dialect/LLVMIR/LLVMOpBase.td"
+include "mlir/IR/SymbolInterfaces.td"
 include "mlir/Interfaces/ControlFlowInterfaces.td"
 include "mlir/Interfaces/SideEffects.td"
 
 class LLVM_Builder<string builder> {
   string llvmBuilder = builder;
 }
 
 def LLVM_OneResultOpBuilder : OpBuilder<

mlir/include/mlir/Dialect/SPIRV/SPIRVStructureOps.td

 // from the SPIR-V binary format in order to utilize MLIR mechanisms.
 //
 //===----------------------------------------------------------------------===//
 
 #ifndef SPIRV_STRUCTURE_OPS
 #define SPIRV_STRUCTURE_OPS
 
 include "mlir/Dialect/SPIRV/SPIRVBase.td"
+include "mlir/IR/SymbolInterfaces.td"
 include "mlir/Interfaces/CallInterfaces.td"
 include "mlir/Interfaces/SideEffects.td"
 
 // -----
 
 def SPV_AddressOfOp : SPV_Op<"_address_of", [InFunctionScope, NoSideEffect]> {
   let summary = "Get the address of a global variable.";
 

mlir/include/mlir/IR/CMakeLists.txt

 set(LLVM_TARGET_DEFINITIONS OpAsmInterface.td)
 mlir_tablegen(OpAsmInterface.h.inc -gen-op-interface-decls)
 mlir_tablegen(OpAsmInterface.cpp.inc -gen-op-interface-defs)
 add_public_tablegen_target(MLIROpAsmInterfacesIncGen)
+
+set(LLVM_TARGET_DEFINITIONS SymbolInterfaces.td)
+mlir_tablegen(SymbolInterfaces.h.inc -gen-op-interface-decls)
+mlir_tablegen(SymbolInterfaces.cpp.inc -gen-op-interface-defs)
+add_public_tablegen_target(MLIRSymbolInterfacesIncGen)

mlir/include/mlir/IR/Function.h

 // Function Operation.
 //===--------------------------------------------------------------------===//
 
 /// FuncOp represents a function, or an operation containing one region that
 /// forms a CFG(Control Flow Graph). The region of a function is not allowed to
 /// implicitly capture global values, and all external references must use
 /// Function arguments or attributes that establish a symbolic connection(e.g.
 /// symbols referenced by name via a string attribute).
-class FuncOp
-    : public Op<FuncOp, OpTrait::ZeroOperands, OpTrait::ZeroResult,
-                OpTrait::IsIsolatedFromAbove, OpTrait::Symbol,
-                OpTrait::FunctionLike, OpTrait::AutomaticAllocationScope,
-                CallableOpInterface::Trait> {
+class FuncOp : public Op<FuncOp, OpTrait::ZeroOperands, OpTrait::ZeroResult,
+                         OpTrait::IsIsolatedFromAbove, OpTrait::FunctionLike,
+                         OpTrait::AutomaticAllocationScope,
+                         CallableOpInterface::Trait, SymbolOpInterface::Trait> {
 public:
   using Op::Op;
   using Op::print;
 
   static StringRef getOperationName() { return "func"; }
 
   static FuncOp create(Location location, StringRef name, FunctionType type,
                        ArrayRef<NamedAttribute> attrs = {});

mlir/include/mlir/IR/Module.h

 /// ModuleOp represents a module, or an operation containing one region with a
 /// single block containing opaque operations. The region of a module is not
 /// allowed to implicitly capture global values, and all external references
 /// must use symbolic references via attributes(e.g. via a string name).
 class ModuleOp
     : public Op<
           ModuleOp, OpTrait::ZeroOperands, OpTrait::ZeroResult,
           OpTrait::IsIsolatedFromAbove, OpTrait::SymbolTable,
-          OpTrait::SingleBlockImplicitTerminator<ModuleTerminatorOp>::Impl> {
+          OpTrait::SingleBlockImplicitTerminator<ModuleTerminatorOp>::Impl,
+          SymbolOpInterface::Trait> {
 public:
   using Op::Op;
   using Op::print;
 
   static StringRef getOperationName() { return "module"; }
 
   static void build(Builder *builder, OperationState &result,
                     Optional<StringRef> name = llvm::None);
void insert(Operation *insertPt, Operation *op) {
     insert(Block::iterator(insertPt), op);
   }
   void insert(Block::iterator insertPt, Operation *op) {
     auto *body = getBody();
     if (insertPt == body->end())
       insertPt = Block::iterator(body->getTerminator());
     body->getOperations().insert(insertPt, op);
   }
+
+  //===--------------------------------------------------------------------===//
+  // SymbolOpInterface Methods
+  //===--------------------------------------------------------------------===//
+
+  /// A ModuleOp may optionally define a symbol.
+  bool isOptionalSymbol() { return true; }
 };
 
 /// The ModuleTerminatorOp is a special terminator operation for the body of a
 /// ModuleOp, it has no semantic meaning beyond keeping the body of a ModuleOp
 /// well-formed.
 ///
 /// This operation does _not_ have a custom syntax. However, ModuleOp will omit
 /// the terminator in their custom syntax for brevity.

mlir/include/mlir/IR/OpBase.td

 def SameOperandsAndResultShape : NativeOpTrait<"SameOperandsAndResultShape">;
 // Op has the same operand and result type.
 def SameOperandsAndResultType : NativeOpTrait<"SameOperandsAndResultType">;
 // Op has the same element type (or type itself, if scalar) for all operands.
 def SameOperandsElementType : NativeOpTrait<"SameOperandsElementType">;
 // Op has the same operand and result element type (or type itself, if scalar).
 def SameOperandsAndResultElementType :
   NativeOpTrait<"SameOperandsAndResultElementType">;
-// Op is a symbol.
-def Symbol : NativeOpTrait<"Symbol">;
-// Op defines a symbol table.
-def SymbolTable : NativeOpTrait<"SymbolTable">;
 // Op is a terminator.
 def Terminator : NativeOpTrait<"IsTerminator">;
 
 // Op's regions have a single block with the specified terminator.
 class SingleBlockImplicitTerminator<string op>
     : ParamNativeOpTrait<"SingleBlockImplicitTerminator", op>;
 
 // Op's parent operation is the provided one.

 // C++. The purpose to wrap around C++ symbol string with this class is to make
 // interfaces specified for ops in TableGen less alien and more integrated.
 class OpInterfaceTrait<string name, code verifyBody = [{}]> : NativeOpTrait<""> {
   let trait = name # "::Trait";
 
   // Specify the body of the verification function. `$_op` will be replaced with
   // the operation being verified.
   code verify = verifyBody;
+
+  // An optional code block containing extra declarations to place in the
+  // interface trait declaration.
+  code extraTraitClassDeclaration = "";
 }
 
 // This class represents a single, optionally static, interface method.
 // Note: non-static interface methods have an implicit 'op' parameter
 // corresponding to an instance of the derived operation.
 class InterfaceMethod<string desc, string retTy, string methodName,
                       dag args = (ins), code methodBody = [{}],
                       code defaultImplementation = [{}]> {

mlir/include/mlir/IR/OpDefinition.h

 ///    };
 ///  };
 ///
 template <typename ConcreteType, typename Traits>
 class OpInterface : public Op<ConcreteType> {
 public:
   using Concept = typename Traits::Concept;
   template <typename T> using Model = typename Traits::template Model<T>;
+  using Base = OpInterface<ConcreteType, Traits>;
 
   OpInterface(Operation *op = nullptr)
       : Op<ConcreteType>(op), impl(op ? getInterfaceFor(op) : nullptr) {
     assert((!op || impl) &&
            "instantiating an interface with an unregistered operation");
   }
 
   /// Support 'classof' by checking if the given operation defines the concrete

mlir/include/mlir/IR/SymbolInterfaces.td

+//===- SymbolInterfaces.td - Interfaces for symbol ops -----*- tablegen -*-===//
+//
+// 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 contains a set of interfaces and traits that can be used to define
+// properties of symbol and symbol table operations.
+//
+//===----------------------------------------------------------------------===//
+
+#ifndef MLIR_IR_SYMBOLINTERFACES
+#define MLIR_IR_SYMBOLINTERFACES
+
+include "mlir/IR/OpBase.td"
+
+//===----------------------------------------------------------------------===//
+// SymbolOpInterface
+//===----------------------------------------------------------------------===//
+
+def Symbol : OpInterface<"SymbolOpInterface"> {
+  let description = [{
+    This interface describes an operation that may define a `Symbol`. A `Symbol`
+    is a named operation that resides immediately within a region that defines
+    a `SymbolTable`. See [Symbols and SymbolTables](SymbolsAndSymbolTables.md)
+    for more details.
+  }];
+
+  let methods = [
+    InterfaceMethod<"Returns the name of this symbol.",
+      "StringRef", "getName", (ins), [{
+        // Don't rely on the trait implementation as optional symbol operations
+        // may override this.
+        return mlir::SymbolTable::getSymbolName(op);
+      }], /*defaultImplementation=*/[{
+        return mlir::SymbolTable::getSymbolName(this->getOperation());
+      }]
+    >,
+    InterfaceMethod<"Sets the name of this symbol.",
+      "void", "setName", (ins "StringRef":$name), [{}],
+      /*defaultImplementation=*/[{
+        this->getOperation()->setAttr(
+            mlir::SymbolTable::getSymbolAttrName(),
+            StringAttr::get(name, this->getOperation()->getContext()));
+      }]
+    >,
+    InterfaceMethod<"Gets the visibility of this symbol.",
+      "mlir::SymbolTable::Visibility", "getVisibility", (ins), [{}],
+      /*defaultImplementation=*/[{
+        return mlir::SymbolTable::getSymbolVisibility(this->getOperation());
+      }]
+    >,
+    InterfaceMethod<"Sets the visibility of this symbol.",
+      "void", "setVisibility", (ins "mlir::SymbolTable::Visibility":$vis), [{}],
+      /*defaultImplementation=*/[{
+        mlir::SymbolTable::setSymbolVisibility(this->getOperation(), vis);
+      }]
+    >,
+    InterfaceMethod<[{
+        Get all of the uses of the current symbol that are nested within the
+        given operation 'from'.
+        Note: See mlir::SymbolTable::getSymbolUses for more details.
+      }],
+      "Optional<::mlir::SymbolTable::UseRange>", "getSymbolUses",
+      (ins "Operation *":$from), [{}],
+      /*defaultImplementation=*/[{
+        return ::mlir::SymbolTable::getSymbolUses(this->getOperation(), from);
+      }]
+    >,
+    InterfaceMethod<[{
+        Return if the current symbol is known to have no uses that are nested
+        within the given operation 'from'.
+        Note: See mlir::SymbolTable::symbolKnownUseEmpty for more details.
+      }],
+      "bool", "symbolKnownUseEmpty", (ins "Operation *":$from), [{}],
+      /*defaultImplementation=*/[{
+        return ::mlir::SymbolTable::symbolKnownUseEmpty(this->getOperation(),
+                                                        from);
+      }]
+    >,
+    InterfaceMethod<[{
+        Attempt to replace all uses of the current symbol with the provided
+        symbol 'newSymbol' that are nested within the given operation 'from'.
+        Note: See mlir::SymbolTable::replaceAllSymbolUses for more details.
+      }],
+      "LogicalResult", "replaceAllSymbolUses", (ins "StringRef":$newSymbol,
+                                                    "Operation *":$from), [{}],
+      /*defaultImplementation=*/[{
+        return ::mlir::SymbolTable::replaceAllSymbolUses(this->getOperation(),
+                                                         newSymbol, from);
+      }]
+    >,
+    InterfaceMethod<[{
+        Returns true if this operation optionally defines a symbol based on the
+        presence of the symbol name.
+      }],
+      "bool", "isOptionalSymbol", (ins), [{}],
+      /*defaultImplementation=*/[{ return false; }]
+    >,
+    InterfaceMethod<[{
+        Returns true if this operation can be discarded if it has no remaining
+        symbol uses.
+      }],
+      "bool", "canDiscardOnUseEmpty", (ins), [{}],
+      /*defaultImplementation=*/[{
+        // By default, base this on the visibility alone. A symbol can be
+        // discarded as long as it is not public. Only public symbols may be
+        // visible from outside of the IR.
+        return getVisibility() != ::mlir::SymbolTable::Visibility::Public;
+      }]
+    >,
+  ];
+
+  let verify = [{
+    // If this is an optional symbol, bail out early if possible.
+    auto concreteOp = cast<ConcreteOp>($_op);
+    if (concreteOp.isOptionalSymbol()) {
+      if(!concreteOp.getAttr(::mlir::SymbolTable::getSymbolAttrName()))
+        return success();
+    }
+    return ::mlir::OpTrait::impl::verifySymbol($_op);
+  }];
+
+  let extraClassDeclaration = [{
+    using Visibility = mlir::SymbolTable::Visibility;
+
+    /// Custom classof that handles the case where the symbol is optional.
+    static bool classof(Operation *op) {
+      return Base::classof(op)
+        && op->getAttr(::mlir::SymbolTable::getSymbolAttrName());
+    }
+
+    /// Returns true if this symbol has nested visibility.
+    bool isNested() { return getVisibility() == Visibility::Nested; }
+    /// Returns true if this symbol has private visibility.
+    bool isPrivate() { return getVisibility() == Visibility::Private; }
+    /// Returns true if this symbol has public visibility.
+    bool isPublic() { return getVisibility() == Visibility::Public; }
+  }];
+
+  let extraTraitClassDeclaration = [{
+    using Visibility = mlir::SymbolTable::Visibility;
+  }];
+}
+
+//===----------------------------------------------------------------------===//
+// Symbol Traits
+//===----------------------------------------------------------------------===//
+
+// Op defines a symbol table.
+def SymbolTable : NativeOpTrait<"SymbolTable">;
+
+#endif // MLIR_IR_SYMBOLINTERFACES

mlir/include/mlir/IR/SymbolTable.h

enum class Visibility {
 
     /// The symbol is visible to the current IR, which may include operations in
     /// symbol tables above the one that owns the current symbol. `Nested`
     /// visibility allows for referencing a symbol outside of its current symbol
     /// table, while retaining the ability to observe all uses.
     Nested,
   };
 
-  /// Returns true if the given operation defines a symbol.
-  static bool isSymbol(Operation *op);
-
   /// Returns the name of the given symbol operation.
   static StringRef getSymbolName(Operation *symbol);
   /// Sets the name of the given symbol operation.
   static void setSymbolName(Operation *symbol, StringRef name);
 
   /// Returns the visibility of the given symbol operation.
   static Visibility getSymbolVisibility(Operation *symbol);
   /// Sets the visibility of the given symbol operation.
public:
   Operation *lookupSymbol(StringRef name) {
     return mlir::SymbolTable::lookupSymbolIn(this->getOperation(), name);
   }
   template <typename T> T lookupSymbol(StringRef name) {
     return dyn_cast_or_null<T>(lookupSymbol(name));
   }
 };
 
-/// A trait used to define a symbol that can be used on operations within a
-/// symbol table. Operations using this trait must adhere to the following:
-///   * Have a StringAttr attribute named 'SymbolTable::getSymbolAttrName()'.
-template <typename ConcreteType>
-class Symbol : public TraitBase<ConcreteType, Symbol> {
-public:
-  using Visibility = mlir::SymbolTable::Visibility;
-
-  static LogicalResult verifyTrait(Operation *op) {
-    return impl::verifySymbol(op);
-  }
-
-  /// Returns the name of this symbol.
-  StringRef getName() {
-    return this->getOperation()
-        ->template getAttrOfType<StringAttr>(
-            mlir::SymbolTable::getSymbolAttrName())
-        .getValue();
-  }
-
-  /// Set the name of this symbol.
-  void setName(StringRef name) {
-    this->getOperation()->setAttr(
-        mlir::SymbolTable::getSymbolAttrName(),
-        StringAttr::get(name, this->getOperation()->getContext()));
-  }
-
-  /// Returns the visibility of the current symbol.
-  Visibility getVisibility() {
-    return mlir::SymbolTable::getSymbolVisibility(this->getOperation());
-  }
-
-  /// Sets the visibility of the current symbol.
-  void setVisibility(Visibility vis) {
-    mlir::SymbolTable::setSymbolVisibility(this->getOperation(), vis);
-  }
-
-  /// Get all of the uses of the current symbol that are nested within the given
-  /// operation 'from'.
-  /// Note: See mlir::SymbolTable::getSymbolUses for more details.
-  Optional<::mlir::SymbolTable::UseRange> getSymbolUses(Operation *from) {
-    return ::mlir::SymbolTable::getSymbolUses(this->getOperation(), from);
-  }
-
-  /// Return if the current symbol is known to have no uses that are nested
-  /// within the given operation 'from'.
-  /// Note: See mlir::SymbolTable::symbolKnownUseEmpty for more details.
-  bool symbolKnownUseEmpty(Operation *from) {
-    return ::mlir::SymbolTable::symbolKnownUseEmpty(this->getOperation(), from);
-  }
+} // end namespace OpTrait
 
-  /// Attempt to replace all uses of the current symbol with the provided symbol
-  /// 'newSymbol' that are nested within the given operation 'from'.
-  /// Note: See mlir::SymbolTable::replaceAllSymbolUses for more details.
-  LLVM_NODISCARD LogicalResult replaceAllSymbolUses(StringRef newSymbol,
-                                                    Operation *from) {
-    return ::mlir::SymbolTable::replaceAllSymbolUses(this->getOperation(),
-                                                     newSymbol, from);
-  }
-};
+/// Include the generated symbol interfaces.
+#include "mlir/IR/SymbolInterfaces.h.inc"
 
-} // end namespace OpTrait
 } // end namespace mlir
 
 #endif // MLIR_IR_SYMBOLTABLE_H

mlir/include/mlir/TableGen/OpInterfaces.h

public:
   ArrayRef<OpInterfaceMethod> getMethods() const;
 
   // Return the description of this method if it has one.
   llvm::Optional<StringRef> getDescription() const;
 
   // Return the interfaces extra class declaration code.
   llvm::Optional<StringRef> getExtraClassDeclaration() const;
 
+  // Return the traits extra class declaration code.
+  llvm::Optional<StringRef> getExtraTraitClassDeclaration() const;
+
   // Return the verify method body if it has one.
   llvm::Optional<StringRef> getVerify() const;
 
 private:
   // The TableGen definition of this interface.
   const llvm::Record *def;
 
   // The methods of this interface.
   SmallVector<OpInterfaceMethod, 8> methods;
 };
 
 } // end namespace tblgen
 } // end namespace mlir
 
 #endif // MLIR_TABLEGEN_OPINTERFACES_H_

mlir/lib/IR/CMakeLists.txt

 file(GLOB globbed *.c *.cpp)
 add_mlir_library(MLIRIR
   ${globbed}
 
   ADDITIONAL_HEADER_DIRS
   ${MLIR_MAIN_INCLUDE_DIR}/mlir/IR
 
   DEPENDS
   MLIRCallInterfacesIncGen
   MLIROpAsmInterfacesIncGen
+  MLIRSymbolInterfacesIncGen
   )
 target_link_libraries(MLIRIR
   PUBLIC
   MLIRSupport
   LLVMSupport
   )

mlir/lib/IR/SymbolTable.cpp

void SymbolTable::insert(Operation *symbol, Block::iterator insertPt) {
   do {
     nameBuffer.resize(originalLength);
     nameBuffer += '_';
     nameBuffer += std::to_string(uniquingCounter++);
   } while (!symbolTable.insert({nameBuffer, symbol}).second);
   setSymbolName(symbol, nameBuffer);
 }
 
-/// Returns true if the given operation defines a symbol.
-bool SymbolTable::isSymbol(Operation *op) {
-  return op->hasTrait<OpTrait::Symbol>() || getNameIfSymbol(op).hasValue();
-}
-
 /// Returns the name of the given symbol operation.
 StringRef SymbolTable::getSymbolName(Operation *symbol) {
   Optional<StringRef> name = getNameIfSymbol(symbol);
   assert(name && "expected valid symbol name");
   return *name;
 }
 /// Sets the name of the given symbol operation.
 void SymbolTable::setSymbolName(Operation *symbol, StringRef name) {
LogicalResult SymbolTable::replaceAllSymbolUses(StringRef oldSymbol,
                                                 Region *from) {
   return replaceAllSymbolUsesImpl(oldSymbol, newSymbol, from);
 }
 LogicalResult SymbolTable::replaceAllSymbolUses(Operation *oldSymbol,
                                                 StringRef newSymbol,
                                                 Region *from) {
   return replaceAllSymbolUsesImpl(oldSymbol, newSymbol, from);
 }
+
+//===----------------------------------------------------------------------===//
+// Symbol Interfaces
+//===----------------------------------------------------------------------===//
+
+/// Include the generated symbol interfaces.
+#include "mlir/IR/SymbolInterfaces.cpp.inc"

mlir/lib/TableGen/OpInterfaces.cpp

 }
 
 // Return the interfaces extra class declaration code.
 llvm::Optional<StringRef> OpInterface::getExtraClassDeclaration() const {
   auto value = def->getValueAsString("extraClassDeclaration");
   return value.empty() ? llvm::Optional<StringRef>() : value;
 }
 
+// Return the traits extra class declaration code.
+llvm::Optional<StringRef> OpInterface::getExtraTraitClassDeclaration() const {
+  auto value = def->getValueAsString("extraTraitClassDeclaration");
+  return value.empty() ? llvm::Optional<StringRef>() : value;
+}
+
 // Return the body for this method if it has one.
 llvm::Optional<StringRef> OpInterface::getVerify() const {
   auto value = def->getValueAsString("verify");
   return value.empty() ? llvm::Optional<StringRef>() : value;
 }

mlir/lib/Transforms/Inliner.cpp

 #define DEBUG_TYPE "inlining"
 
 using namespace mlir;
 
 //===----------------------------------------------------------------------===//
 // Symbol Use Tracking
 //===----------------------------------------------------------------------===//
 
-/// Returns true if this operation can be discarded if it is a symbol and has no
-/// uses. 'allUsesVisible' corresponds to if the parent symbol table is hidden
-/// from above.
-static bool canDiscardSymbolOnUseEmpty(Operation *op, bool allUsesVisible) {
-  if (!SymbolTable::isSymbol(op))
-    return false;
-
-  // TODO: This is essentially the same logic from SymbolDCE. Remove this when
-  // we have a 'Symbol' interface.
-  // Private symbols are always initially considered dead.
-  SymbolTable::Visibility visibility = SymbolTable::getSymbolVisibility(op);
-  if (visibility == mlir::SymbolTable::Visibility::Private)
-    return true;
-  // We only include nested visibility here if all uses are visible.
-  if (allUsesVisible && visibility == SymbolTable::Visibility::Nested)
-    return true;
-  // Otherwise, public symbols are never removable.
-  return false;
-}
-
 /// Walk all of the symbol table operations nested with 'op' along with a
 /// boolean signifying if the symbols within can be treated as if all uses are
 /// visible. The provided callback is invoked with the symbol table operation,
 /// and a boolean signaling if all of the uses within the symbol table are
 /// visible.
 static void walkSymbolTables(Operation *op, bool allSymUsesVisible,
                              function_ref<void(Operation *, bool)> callback) {
   if (op->hasTrait<OpTrait::SymbolTable>()) {
-    allSymUsesVisible = allSymUsesVisible || !SymbolTable::isSymbol(op) ||
-                        SymbolTable::getSymbolVisibility(op) ==
-                            SymbolTable::Visibility::Private;
+    SymbolOpInterface symbol = dyn_cast<SymbolOpInterface>(op);
+    allSymUsesVisible = allSymUsesVisible || !symbol || symbol.isPrivate();
     callback(op, allSymUsesVisible);
   } else {
     // Otherwise if 'op' is not a symbol table, any nested symbols are
     // guaranteed to be hidden.
     allSymUsesVisible = true;
   }
 
   for (Region &region : op->getRegions())
CGUseList::CGUseList(Operation *op, CallGraph &cg) {
 
   // Walk each of the symbol tables looking for discardable callgraph nodes.
   auto walkFn = [&](Operation *symbolTableOp, bool allUsesVisible) {
     for (Block &block : symbolTableOp->getRegion(0)) {
       for (Operation &op : block) {
         // If this is a callgraph operation, check to see if it is discardable.
         if (auto callable = dyn_cast<CallableOpInterface>(&op)) {
           if (auto *node = cg.lookupNode(callable.getCallableRegion())) {
-            if (canDiscardSymbolOnUseEmpty(&op, allUsesVisible))
+            SymbolOpInterface symbol = dyn_cast<SymbolOpInterface>(&op);
+            if (symbol && (allUsesVisible || symbol.isPrivate()) &&
+                symbol.canDiscardOnUseEmpty()) {
               discardableSymNodeUses.try_emplace(node, 0);
+            }
             continue;
           }
         }
         // Otherwise, check for any referenced nodes. These will be always-live.
         walkReferencedSymbolNodes(&op, cg, alwaysLiveNodes,
                                   [](CallGraphNode *, Operation *) {});
       }
     }
void CGUseList::eraseNode(CallGraphNode *node) {
   decrementDiscardableUses(useIt->getSecond());
   nodeUses.erase(useIt);
   discardableSymNodeUses.erase(node);
 }
 
 bool CGUseList::isDead(CallGraphNode *node) const {
   // If the parent operation isn't a symbol, simply check normal SSA deadness.
   Operation *nodeOp = node->getCallableRegion()->getParentOp();
-  if (!SymbolTable::isSymbol(nodeOp))
+  if (!isa<SymbolOpInterface>(nodeOp))
     return MemoryEffectOpInterface::hasNoEffect(nodeOp) && nodeOp->use_empty();
 
   // Otherwise, check the number of symbol uses.
   auto symbolIt = discardableSymNodeUses.find(node);
   return symbolIt != discardableSymNodeUses.end() && symbolIt->second == 0;
 }
 
 bool CGUseList::hasOneUseAndDiscardable(CallGraphNode *node) const {
   // If this isn't a symbol node, check for side-effects and SSA use count.
   Operation *nodeOp = node->getCallableRegion()->getParentOp();
-  if (!SymbolTable::isSymbol(nodeOp))
+  if (!isa<SymbolOpInterface>(nodeOp))
     return MemoryEffectOpInterface::hasNoEffect(nodeOp) && nodeOp->hasOneUse();
 
   // Otherwise, check the number of symbol uses.
   auto symbolIt = discardableSymNodeUses.find(node);
   return symbolIt != discardableSymNodeUses.end() && symbolIt->second == 1;
 }
 
 void CGUseList::recomputeUses(CallGraphNode *node, CallGraph &cg) {

mlir/lib/Transforms/SymbolDCE.cpp

symbolTableOp->emitOpError()
         << " was scheduled to run under SymbolDCE, but does not define a "
            "symbol table";
     return signalPassFailure();
   }
 
   // A flag that signals if the top level symbol table is hidden, i.e. not
   // accessible from parent scopes.
   bool symbolTableIsHidden = true;
-  if (symbolTableOp->getParentOp() && SymbolTable::isSymbol(symbolTableOp)) {
-    symbolTableIsHidden = SymbolTable::getSymbolVisibility(symbolTableOp) ==
-                          SymbolTable::Visibility::Private;
-  }
+  SymbolOpInterface symbol = dyn_cast<SymbolOpInterface>(symbolTableOp);
+  if (symbolTableOp->getParentOp() && symbol)
+    symbolTableIsHidden = symbol.isPrivate();
 
   // Compute the set of live symbols within the symbol table.
   DenseSet<Operation *> liveSymbols;
   if (failed(computeLiveness(symbolTableOp, symbolTableIsHidden, liveSymbols)))
     return signalPassFailure();
 
   // After computing the liveness, delete all of the symbols that were found to
   // be dead.
   symbolTableOp->walk([&](Operation *nestedSymbolTable) {
     if (!nestedSymbolTable->hasTrait<OpTrait::SymbolTable>())
       return;
     for (auto &block : nestedSymbolTable->getRegion(0)) {
       for (Operation &op :
            llvm::make_early_inc_range(block.without_terminator())) {
-        if (SymbolTable::isSymbol(&op) && !liveSymbols.count(&op))
+        if (isa<SymbolOpInterface>(&op) && !liveSymbols.count(&op))
           op.erase();
       }
     }
   });
 }
 
 /// Compute the liveness of the symbols within the given symbol table.
 /// `symbolTableIsHidden` is true if this symbol table is known to be
 /// unaccessible from operations in its parent regions.
 LogicalResult SymbolDCE::computeLiveness(Operation *symbolTableOp,
                                          bool symbolTableIsHidden,
                                          DenseSet<Operation *> &liveSymbols) {
   // A worklist of live operations to propagate uses from.
   SmallVector<Operation *, 16> worklist;
 
   // Walk the symbols within the current symbol table, marking the symbols that
   // are known to be live.
   for (auto &block : symbolTableOp->getRegion(0)) {
+    // Add all non-symbols or symbols that can't be discarded.
     for (Operation &op : block.without_terminator()) {
-      // Always add non symbol operations to the worklist.
-      if (!SymbolTable::isSymbol(&op)) {
+      SymbolOpInterface symbol = dyn_cast<SymbolOpInterface>(&op);
+      if (!symbol) {
         worklist.push_back(&op);
         continue;
       }
-
-      // Check the visibility to see if this symbol may be referenced
-      // externally.
-      SymbolTable::Visibility visibility =
-          SymbolTable::getSymbolVisibility(&op);
-
-      // Private symbols are always initially considered dead.
-      if (visibility == mlir::SymbolTable::Visibility::Private)
-        continue;
-      // We only include nested visibility here if the symbol table isn't
-      // hidden.
-      if (symbolTableIsHidden && visibility == SymbolTable::Visibility::Nested)
-        continue;
-
-      // TODO(riverriddle) Add hooks here to allow symbols to provide additional
-      // information, e.g. linkage can be used to drop some symbols that may
-      // otherwise be considered "live".
-      if (liveSymbols.insert(&op).second)
+      bool isDiscardable = (symbolTableIsHidden || symbol.isPrivate()) &&
+                           symbol.canDiscardOnUseEmpty();
+      if (!isDiscardable && liveSymbols.insert(&op).second)
         worklist.push_back(&op);
     }
   }
 
   // Process the set of symbols that were known to be live, adding new symbols
   // that are referenced within.
   while (!worklist.empty()) {
     Operation *op = worklist.pop_back_val();
 
     // If this is a symbol table, recursively compute its liveness.
     if (op->hasTrait<OpTrait::SymbolTable>()) {
       // The internal symbol table is hidden if the parent is, if its not a
       // symbol, or if it is a private symbol.
-      bool symbolIsHidden = symbolTableIsHidden || !SymbolTable::isSymbol(op) ||
-                            SymbolTable::getSymbolVisibility(op) ==
-                                SymbolTable::Visibility::Private;
-      if (failed(computeLiveness(op, symbolIsHidden, liveSymbols)))
+      SymbolOpInterface symbol = dyn_cast<SymbolOpInterface>(op);
+      bool symIsHidden = symbolTableIsHidden || !symbol || symbol.isPrivate();
+      if (failed(computeLiveness(op, symIsHidden, liveSymbols)))
         return failure();
     }
 
     // Collect the uses held by this operation.
     Optional<SymbolTable::UseRange> uses = SymbolTable::getSymbolUses(op);
     if (!uses) {
       return op->emitError()
              << "operation contains potentially unknown symbol table, "

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

 //===-- TestOps.td - Test dialect operation definitions ----*- tablegen -*-===//
 //
 // 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 TEST_OPS
 #define TEST_OPS
 
 include "mlir/IR/OpBase.td"
 include "mlir/IR/OpAsmInterface.td"
+include "mlir/IR/SymbolInterfaces.td"
 include "mlir/Interfaces/SideEffects.td"
 include "mlir/Interfaces/CallInterfaces.td"
 include "mlir/Interfaces/ControlFlowInterfaces.td"
 include "mlir/Interfaces/InferTypeOpInterface.td"
 include "mlir/Interfaces/SideEffects.td"
 
 def Test_Dialect : Dialect {
   let name = "test";

mlir/test/lib/IR/TestSymbolUses.cpp

struct SymbolUsesPass
   }
 
   void runOnOperation() override {
     auto module = getOperation();
 
     // Walk nested symbols.
     SmallVector<FuncOp, 4> deadFunctions;
     module.getBodyRegion().walk([&](Operation *nestedOp) {
-      if (SymbolTable::isSymbol(nestedOp))
+      if (isa<SymbolOpInterface>(nestedOp))
         return operateOnSymbol(nestedOp, module, deadFunctions);
       return WalkResult::advance();
     });
 
     SymbolTable table(module);
     for (Operation *op : deadFunctions) {
       // In order to test the SymbolTable::erase method, also erase completely
       // useless functions.

mlir/tools/mlir-tblgen/OpInterfacesGen.cpp

static void emitTraitDecl(OpInterface &interface, raw_ostream &os,
   }
 
   tblgen::FmtContext traitCtx;
   traitCtx.withOp("op");
   if (auto verify = interface.getVerify()) {
     os << "  static LogicalResult verifyTrait(Operation* op) {\n"
        << std::string(tblgen::tgfmt(*verify, &traitCtx)) << "\n  }\n";
   }
+  if (auto extraTraitDecls = interface.getExtraTraitClassDeclaration())
+    os << extraTraitDecls << "\n";
 
   os << "  };\n";
 }
 
 static void emitInterfaceDecl(OpInterface &interface, raw_ostream &os) {
   StringRef interfaceName = interface.getName();
   auto interfaceTraitsName = (interfaceName + "InterfaceTraits").str();