[mlir] provide C API and Python bindings for symbol tables

Symbol tables are a largely useful top-level IR construct, for example, they
make it easy to access functions in a module by name instead of traversing the
list of module's operations to find the corresponding function.

Depends On D112886

Reviewed By: mehdi_amini

Differential Revision: https://reviews.llvm.org/D112821

Author: ftynse

mlir/include/mlir-c/IR.h

@@ -54,6 +54,7 @@ DEFINE_C_API_STRUCT(MlirOperation, void);
 DEFINE_C_API_STRUCT(MlirOpPrintingFlags, void);
 DEFINE_C_API_STRUCT(MlirBlock, void);
 DEFINE_C_API_STRUCT(MlirRegion, void);
+DEFINE_C_API_STRUCT(MlirSymbolTable, void);
 
 DEFINE_C_API_STRUCT(MlirAttribute, const void);
 DEFINE_C_API_STRUCT(MlirIdentifier, const void);
@@ -738,6 +739,47 @@ MLIR_CAPI_EXPORTED bool mlirTypeIDEqual(MlirTypeID typeID1, MlirTypeID typeID2);
 /// Returns the hash value of the type id.
 MLIR_CAPI_EXPORTED size_t mlirTypeIDHashValue(MlirTypeID typeID);
 
+//===----------------------------------------------------------------------===//
+// Symbol and SymbolTable API.
+//===----------------------------------------------------------------------===//
+
+/// Returns the name of the attribute used to store symbol names compatible with
+/// symbol tables.
+MLIR_CAPI_EXPORTED MlirStringRef mlirSymbolTableGetSymbolAttributeName();
+
+/// Creates a symbol table for the given operation. If the operation does not
+/// have the SymbolTable trait, returns a null symbol table.
+MLIR_CAPI_EXPORTED MlirSymbolTable
+mlirSymbolTableCreate(MlirOperation operation);
+
+/// Returns true if the symbol table is null.
+static inline bool mlirSymbolTableIsNull(MlirSymbolTable symbolTable) {
+  return !symbolTable.ptr;
+}
+
+/// Destroys the symbol table created with mlirSymbolTableCreate. This does not
+/// affect the operations in the table.
+MLIR_CAPI_EXPORTED void mlirSymbolTableDestroy(MlirSymbolTable symbolTable);
+
+/// Looks up a symbol with the given name in the given symbol table and returns
+/// the operation that corresponds to the symbol. If the symbol cannot be found,
+/// returns a null operation.
+MLIR_CAPI_EXPORTED MlirOperation
+mlirSymbolTableLookup(MlirSymbolTable symbolTable, MlirStringRef name);
+
+/// Inserts the given operation into the given symbol table. The operation must
+/// have the symbol trait. If the symbol table already has a symbol with the
+/// same name, renames the symbol being inserted to ensure name uniqueness. Note
+/// that this does not move the operation itself into the block of the symbol
+/// table operation, this should be done separately. Returns the name of the
+/// symbol after insertion.
+MLIR_CAPI_EXPORTED MlirAttribute
+mlirSymbolTableInsert(MlirSymbolTable symbolTable, MlirOperation operation);
+
+/// Removes the given operation from the symbol table and erases it.
+MLIR_CAPI_EXPORTED void mlirSymbolTableErase(MlirSymbolTable symbolTable,
+                                             MlirOperation operation);
+
 #ifdef __cplusplus
 }
 #endif

mlir/include/mlir-c/Support.h

@@ -79,6 +79,10 @@ inline static MlirStringRef mlirStringRefCreate(const char *str,
 MLIR_CAPI_EXPORTED MlirStringRef
 mlirStringRefCreateFromCString(const char *str);
 
+/// Returns true if two string references are equal, false otherwise.
+MLIR_CAPI_EXPORTED bool mlirStringRefEqual(MlirStringRef string,
+                                           MlirStringRef other);
+
 /// A callback for returning string references.
 ///
 /// This function is called back by the functions that need to return a

mlir/include/mlir/CAPI/IR.h

@@ -27,6 +27,7 @@ DEFINE_C_API_PTR_METHODS(MlirOperation, mlir::Operation)
 DEFINE_C_API_PTR_METHODS(MlirBlock, mlir::Block)
 DEFINE_C_API_PTR_METHODS(MlirOpPrintingFlags, mlir::OpPrintingFlags)
 DEFINE_C_API_PTR_METHODS(MlirRegion, mlir::Region)
+DEFINE_C_API_PTR_METHODS(MlirSymbolTable, mlir::SymbolTable);
 
 DEFINE_C_API_METHODS(MlirAttribute, mlir::Attribute)
 DEFINE_C_API_METHODS(MlirIdentifier, mlir::Identifier)

mlir/lib/Bindings/Python/IRCore.cpp

@@ -1530,6 +1530,57 @@ PyValue PyValue::createFromCapsule(pybind11::object capsule) {
   return PyValue(ownerRef, value);
 }
 
+//------------------------------------------------------------------------------
+// PySymbolTable.
+//------------------------------------------------------------------------------
+
+PySymbolTable::PySymbolTable(PyOperationBase &operation)
+    : operation(operation.getOperation().getRef()) {
+  symbolTable = mlirSymbolTableCreate(operation.getOperation().get());
+  if (mlirSymbolTableIsNull(symbolTable)) {
+    throw py::cast_error("Operation is not a Symbol Table.");
+  }
+}
+
+py::object PySymbolTable::dunderGetItem(const std::string &name) {
+  operation->checkValid();
+  MlirOperation symbol = mlirSymbolTableLookup(
+      symbolTable, mlirStringRefCreate(name.data(), name.length()));
+  if (mlirOperationIsNull(symbol))
+    throw py::key_error("Symbol '" + name + "' not in the symbol table.");
+
+  return PyOperation::forOperation(operation->getContext(), symbol,
+                                   operation.getObject())
+      ->createOpView();
+}
+
+void PySymbolTable::erase(PyOperationBase &symbol) {
+  operation->checkValid();
+  symbol.getOperation().checkValid();
+  mlirSymbolTableErase(symbolTable, symbol.getOperation().get());
+  // The operation is also erased, so we must invalidate it. There may be Python
+  // references to this operation so we don't want to delete it from the list of
+  // live operations here.
+  symbol.getOperation().valid = false;
+}
+
+void PySymbolTable::dunderDel(const std::string &name) {
+  py::object operation = dunderGetItem(name);
+  erase(py::cast<PyOperationBase &>(operation));
+}
+
+PyAttribute PySymbolTable::insert(PyOperationBase &symbol) {
+  operation->checkValid();
+  symbol.getOperation().checkValid();
+  MlirAttribute symbolAttr = mlirOperationGetAttributeByName(
+      symbol.getOperation().get(), mlirSymbolTableGetSymbolAttributeName());
+  if (mlirAttributeIsNull(symbolAttr))
+    throw py::value_error("Expected operation to have a symbol name.");
+  return PyAttribute(
+      symbol.getOperation().getContext(),
+      mlirSymbolTableInsert(symbolTable, symbol.getOperation().get()));
+}
+
 namespace {
 /// CRTP base class for Python MLIR values that subclass Value and should be
 /// castable from it. The value hierarchy is one level deep and is not supposed
@@ -2670,6 +2721,20 @@ void mlir::python::populateIRCore(py::module &m) {
   PyBlockArgument::bind(m);
   PyOpResult::bind(m);
 
+  //----------------------------------------------------------------------------
+  // Mapping of SymbolTable.
+  //----------------------------------------------------------------------------
+  py::class_<PySymbolTable>(m, "SymbolTable", py::module_local())
+      .def(py::init<PyOperationBase &>())
+      .def("__getitem__", &PySymbolTable::dunderGetItem)
+      .def("insert", &PySymbolTable::insert)
+      .def("erase", &PySymbolTable::erase)
+      .def("__delitem__", &PySymbolTable::dunderDel)
+      .def("__contains__", [](PySymbolTable &table, const std::string &name) {
+        return !mlirOperationIsNull(mlirSymbolTableLookup(
+            table, mlirStringRefCreate(name.data(), name.length())));
+      });
+
   // Container bindings.
   PyBlockArgumentList::bind(m);
   PyBlockIterator::bind(m);

mlir/lib/Bindings/Python/IRModule.h

@@ -32,6 +32,7 @@ class DefaultingPyMlirContext;
 class PyModule;
 class PyOperation;
 class PyType;
+class PySymbolTable;
 class PyValue;
 
 /// Template for a reference to a concrete type which captures a python
@@ -513,6 +514,7 @@ class PyOperation : public PyOperationBase, public BaseContextObject {
   bool valid = true;
 
   friend class PyOperationBase;
+  friend class PySymbolTable;
 };
 
 /// A PyOpView is equivalent to the C++ "Op" wrappers: these are the basis for
@@ -876,6 +878,38 @@ class PyIntegerSet : public BaseContextObject {
   MlirIntegerSet integerSet;
 };
 
+/// Bindings for MLIR symbol tables.
+class PySymbolTable {
+public:
+  /// Constructs a symbol table for the given operation.
+  explicit PySymbolTable(PyOperationBase &operation);
+
+  /// Destroys the symbol table.
+  ~PySymbolTable() { mlirSymbolTableDestroy(symbolTable); }
+
+  /// Returns the symbol (opview) with the given name, throws if there is no
+  /// such symbol in the table.
+  pybind11::object dunderGetItem(const std::string &name);
+
+  /// Removes the given operation from the symbol table and erases it.
+  void erase(PyOperationBase &symbol);
+
+  /// Removes the operation with the given name from the symbol table and erases
+  /// it, throws if there is no such symbol in the table.
+  void dunderDel(const std::string &name);
+
+  /// Inserts the given operation into the symbol table. The operation must have
+  /// the symbol trait.
+  PyAttribute insert(PyOperationBase &symbol);
+
+  /// Casts the bindings class into the C API structure.
+  operator MlirSymbolTable() { return symbolTable; }
+
+private:
+  PyOperationRef operation;
+  MlirSymbolTable symbolTable;
+};
+
 void populateIRAffine(pybind11::module &m);
 void populateIRAttributes(pybind11::module &m);
 void populateIRCore(pybind11::module &m);

mlir/lib/CAPI/IR/IR.cpp

@@ -763,3 +763,36 @@ bool mlirTypeIDEqual(MlirTypeID typeID1, MlirTypeID typeID2) {
 size_t mlirTypeIDHashValue(MlirTypeID typeID) {
   return hash_value(unwrap(typeID));
 }
+
+//===----------------------------------------------------------------------===//
+// Symbol and SymbolTable API.
+//===----------------------------------------------------------------------===//
+
+MlirStringRef mlirSymbolTableGetSymbolAttributeName() {
+  return wrap(SymbolTable::getSymbolAttrName());
+}
+
+MlirSymbolTable mlirSymbolTableCreate(MlirOperation operation) {
+  if (!unwrap(operation)->hasTrait<OpTrait::SymbolTable>())
+    return wrap(static_cast<SymbolTable *>(nullptr));
+  return wrap(new SymbolTable(unwrap(operation)));
+}
+
+void mlirSymbolTableDestroy(MlirSymbolTable symbolTable) {
+  delete unwrap(symbolTable);
+}
+
+MlirOperation mlirSymbolTableLookup(MlirSymbolTable symbolTable,
+                                    MlirStringRef name) {
+  return wrap(unwrap(symbolTable)->lookup(StringRef(name.data, name.length)));
+}
+
+MlirAttribute mlirSymbolTableInsert(MlirSymbolTable symbolTable,
+                                    MlirOperation operation) {
+  return wrap(unwrap(symbolTable)->insert(unwrap(operation)));
+}
+
+void mlirSymbolTableErase(MlirSymbolTable symbolTable,
+                          MlirOperation operation) {
+  unwrap(symbolTable)->erase(unwrap(operation));
+}

mlir/lib/CAPI/IR/Support.cpp

@@ -7,9 +7,15 @@
 //===----------------------------------------------------------------------===//
 
 #include "mlir-c/Support.h"
+#include "llvm/ADT/StringRef.h"
 
 #include <cstring>
 
 MlirStringRef mlirStringRefCreateFromCString(const char *str) {
   return mlirStringRefCreate(str, strlen(str));
 }
+
+bool mlirStringRefEqual(MlirStringRef string, MlirStringRef other) {
+  return llvm::StringRef(string.data, string.length) ==
+         llvm::StringRef(other.data, other.length);
+}