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

[ftynse]

Summary

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

[mehdi_amini]

mlir/include/mlir-c/IR.h

776

Why don't we return a MlirStringRef with the new name to match the C++ API?

mlir/test/CAPI/ir.c

1864

I think we're missing an example of renaming, if you add an operation in the second module with the same name as one in the first module, move it there, and get the new name back.

mlir/test/python/ir/operation.py

842

These two lines are stressing my point from the parent revision: there should be a single API call to move an operation instead of a detach/reattach sequence.

[ftynse]

Address review.

mlir/include/mlir-c/IR.h

776

Because the C++ API does not return the new name either https://github.com/llvm/llvm-project/blob/3be3c944a5bacfd208b56853941b0fa4dec3ddcc/mlir/lib/IR/SymbolTable.cpp#L155.

I am happy to change that if desired.

[mehdi_amini]

mlir/include/mlir-c/IR.h

776

Uh, I was sure we were returning the new name... Yeah it'd be good change to do so I think. But that's irrelevant to this patch.
(unless you want to minimize churn in the C API, so land the C++ change ahead of this one and rebase)

[ftynse]

Rebase and address review.

[ftynse]

Rebase.

[ftynse]

Another rebase.

mlir/include/mlir-c/IR.h

#define DEFINE_C_API_STRUCT(name, storage)                                     \
   typedef struct name name
 
 DEFINE_C_API_STRUCT(MlirContext, void);
 DEFINE_C_API_STRUCT(MlirDialect, void);
 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);
 DEFINE_C_API_STRUCT(MlirLocation, const void);
 DEFINE_C_API_STRUCT(MlirModule, const void);
 DEFINE_C_API_STRUCT(MlirType, const void);
 DEFINE_C_API_STRUCT(MlirTypeID, const void);
 DEFINE_C_API_STRUCT(MlirValue, const void);

 static inline bool mlirTypeIDIsNull(MlirTypeID typeID) { return !typeID.ptr; }
 
 /// Checks if two type ids are equal.
 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
 
 #endif // MLIR_C_IR_H

mlir/include/mlir-c/Support.h

inline static MlirStringRef mlirStringRefCreate(const char *str,
   return result;
 }
 
 /// Constructs a string reference from a null-terminated C string. Prefer
 /// mlirStringRefCreate if the length of the string is known.
 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
 /// reference to the portion of the string with the following arguments:
 ///  - an MlirStringRef representing the current portion of the string
 ///  - a pointer to user data forwarded from the printing call.
 typedef void (*MlirStringCallback)(MlirStringRef, void *);
 

mlir/include/mlir/CAPI/IR.h

 #include "mlir/IR/Operation.h"
 
 DEFINE_C_API_PTR_METHODS(MlirContext, mlir::MLIRContext)
 DEFINE_C_API_PTR_METHODS(MlirDialect, mlir::Dialect)
 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)
 DEFINE_C_API_METHODS(MlirLocation, mlir::Location)
 DEFINE_C_API_METHODS(MlirModule, mlir::ModuleOp)
 DEFINE_C_API_METHODS(MlirType, mlir::Type)
 DEFINE_C_API_METHODS(MlirTypeID, mlir::TypeID)
 DEFINE_C_API_METHODS(MlirValue, mlir::Value)
 
 #endif // MLIR_CAPI_IR_H

mlir/lib/Bindings/Python/IRCore.cpp

PyValue PyValue::createFromCapsule(pybind11::object capsule) {
   if (mlirOperationIsNull(owner))
     throw py::error_already_set();
   MlirContext ctx = mlirOperationGetContext(owner);
   PyOperationRef ownerRef =
       PyOperation::forOperation(PyMlirContext::forContext(ctx), owner);
   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
 /// to accommodate other levels unless core MLIR changes.
 template <typename DerivedTy>
 class PyConcreteValue : public PyValue {
 public:
   // Derived classes must define statics for:
py::class_<PyValue>(m, "Value", py::module_local())
           kValueDunderStrDocstring)
       .def_property_readonly("type", [](PyValue &self) {
         return PyType(self.getParentOperation()->getContext(),
                       mlirValueGetType(self.get()));
       });
   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);
   PyBlockList::bind(m);
   PyOperationIterator::bind(m);
   PyOperationList::bind(m);
   PyOpAttributeMap::bind(m);
   PyOpOperandList::bind(m);
   PyOpResultList::bind(m);
   PyRegionIterator::bind(m);
   PyRegionList::bind(m);
 
   // Debug bindings.
   PyGlobalDebugFlag::bind(m);
 }

mlir/lib/Bindings/Python/IRModule.h

 class PyInsertionPoint;
 class PyLocation;
 class DefaultingPyLocation;
 class PyMlirContext;
 class DefaultingPyMlirContext;
 class PyModule;
 class PyOperation;
 class PyType;
+class PySymbolTable;
 class PyValue;
 
 /// Template for a reference to a concrete type which captures a python
 /// reference to its underlying python object.
 template <typename T>
 class PyObjectRef {
 public:
   PyObjectRef(T *referrent, pybind11::object object)
private:
   // trivial module parent back-reference. Generalize this to also account for
   // transitions from detached to attached and address TODOs in the
   // ir_operation.py regarding testing corresponding lifetime guarantees.
   pybind11::object parentKeepAlive;
   bool attached = true;
   bool valid = true;
 
   friend class PyOperationBase;
+  friend class PySymbolTable;
 };
 
 /// A PyOpView is equivalent to the C++ "Op" wrappers: these are the basis for
 /// providing more instance-specific accessors and serve as the base class for
 /// custom ODS-style operation classes. Since this class is subclass on the
 /// python side, it must present an __init__ method that operates in pure
 /// python types.
 class PyOpView : public PyOperationBase {
public:
   /// Note that PyIntegerSet instances may be uniqued, so the returned object
   /// may be a pre-existing object. Integer sets are owned by the context.
   static PyIntegerSet createFromCapsule(pybind11::object capsule);
 
 private:
   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);
 void populateIRInterfaces(pybind11::module &m);
 void populateIRTypes(pybind11::module &m);
 
 } // namespace python
 } // namespace mlir

mlir/lib/CAPI/IR/IR.cpp

 
 bool mlirTypeIDEqual(MlirTypeID typeID1, MlirTypeID typeID2) {
   return unwrap(typeID1) == unwrap(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

 //===- Support.cpp - Helpers for C interface to MLIR API ------------------===//
 //
 // 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
 //
 //===----------------------------------------------------------------------===//
 
 #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);
+}

mlir/test/CAPI/ir.c

 }
 
 // Logs when the delete user data callback is called
 static void deleteUserData(void *userData) {
   fprintf(stderr, "deleting user data (userData: %" PRIdPTR ")\n",
           (intptr_t)userData);
 }
 
-void testDiagnostics() {
-  MlirContext ctx = mlirContextCreate();
-  MlirDiagnosticHandlerID id = mlirContextAttachDiagnosticHandler(
-      ctx, errorHandler, (void *)42, deleteUserData);
-  fprintf(stderr, "@test_diagnostics\n");
-  MlirLocation unknownLoc = mlirLocationUnknownGet(ctx);
-  mlirEmitError(unknownLoc, "test diagnostics");
-  MlirLocation fileLineColLoc = mlirLocationFileLineColGet(
-      ctx, mlirStringRefCreateFromCString("file.c"), 1, 2);
-  mlirEmitError(fileLineColLoc, "test diagnostics");
-  MlirLocation callSiteLoc = mlirLocationCallSiteGet(
-      mlirLocationFileLineColGet(
-          ctx, mlirStringRefCreateFromCString("other-file.c"), 2, 3),
-      fileLineColLoc);
-  mlirEmitError(callSiteLoc, "test diagnostics");
-  MlirLocation null = {0};
-  MlirLocation nameLoc =
-      mlirLocationNameGet(ctx, mlirStringRefCreateFromCString("named"), null);
-  mlirEmitError(nameLoc, "test diagnostics");
-  MlirLocation locs[2] = {nameLoc, callSiteLoc};
-  MlirAttribute nullAttr = {0};
-  MlirLocation fusedLoc = mlirLocationFusedGet(ctx, 2, locs, nullAttr);
-  mlirEmitError(fusedLoc, "test diagnostics");
-  mlirContextDetachDiagnosticHandler(ctx, id);
-  mlirEmitError(unknownLoc, "more test diagnostics");
-  // CHECK-LABEL: @test_diagnostics
-  // CHECK: processing diagnostic (userData: 42) <<
-  // CHECK:   test diagnostics
-  // CHECK:   loc(unknown)
-  // CHECK: >> end of diagnostic (userData: 42)
-  // CHECK: processing diagnostic (userData: 42) <<
-  // CHECK:   test diagnostics
-  // CHECK:   loc("file.c":1:2)
-  // CHECK: >> end of diagnostic (userData: 42)
-  // CHECK: processing diagnostic (userData: 42) <<
-  // CHECK:   test diagnostics
-  // CHECK:   loc(callsite("other-file.c":2:3 at "file.c":1:2))
-  // CHECK: >> end of diagnostic (userData: 42)
-  // CHECK: processing diagnostic (userData: 42) <<
-  // CHECK:   test diagnostics
-  // CHECK:   loc("named")
-  // CHECK: >> end of diagnostic (userData: 42)
-  // CHECK: processing diagnostic (userData: 42) <<
-  // CHECK:   test diagnostics
-  // CHECK:   loc(fused["named", callsite("other-file.c":2:3 at "file.c":1:2)])
-  // CHECK: deleting user data (userData: 42)
-  // CHECK-NOT: processing diagnostic
-  // CHECK:     more test diagnostics
-  mlirContextDestroy(ctx);
-}
-
 int testTypeID(MlirContext ctx) {
   fprintf(stderr, "@testTypeID\n");
 
   // Test getting and comparing type and attribute type ids.
   MlirType i32 = mlirIntegerTypeGet(ctx, 32);
   MlirTypeID i32ID = mlirTypeGetTypeID(i32);
   MlirType ui32 = mlirIntegerTypeUnsignedGet(ctx, 32);
   MlirTypeID ui32ID = mlirTypeGetTypeID(ui32);
int testTypeID(MlirContext ctx) {
   }
 
   mlirOperationDestroy(constZero);
   mlirOperationDestroy(unregisteredOp);
 
   return 0;
 }
 
+int testSymbolTable(MlirContext ctx) {
+  fprintf(stderr, "@testSymbolTable\n");
+
+  const char *moduleString = "func private @foo()"
+                             "func private @bar()";
+  const char *otherModuleString = "func private @qux()"
+                                  "func private @foo()";
+
+  MlirModule module =
+      mlirModuleCreateParse(ctx, mlirStringRefCreateFromCString(moduleString));
+  MlirModule otherModule = mlirModuleCreateParse(
+      ctx, mlirStringRefCreateFromCString(otherModuleString));
+
+  MlirSymbolTable symbolTable =
+      mlirSymbolTableCreate(mlirModuleGetOperation(module));
+
+  MlirOperation funcFoo =
+      mlirSymbolTableLookup(symbolTable, mlirStringRefCreateFromCString("foo"));
+  if (mlirOperationIsNull(funcFoo))
+    return 1;
+
+  MlirOperation funcBar =
+      mlirSymbolTableLookup(symbolTable, mlirStringRefCreateFromCString("bar"));
+  if (mlirOperationEqual(funcFoo, funcBar))
+    return 2;
+
+  MlirOperation missing =
+      mlirSymbolTableLookup(symbolTable, mlirStringRefCreateFromCString("qux"));
+  if (!mlirOperationIsNull(missing))
+    return 3;
+
+  MlirBlock moduleBody = mlirModuleGetBody(module);
+  MlirBlock otherModuleBody = mlirModuleGetBody(otherModule);
+  MlirOperation operation = mlirBlockGetFirstOperation(otherModuleBody);
+  mlirOperationRemoveFromParent(operation);
+  mlirBlockAppendOwnedOperation(moduleBody, operation);
+
+  // At this moment, the operation is still missing from the symbol table.
+  MlirOperation stillMissing =
+      mlirSymbolTableLookup(symbolTable, mlirStringRefCreateFromCString("qux"));
+  if (!mlirOperationIsNull(stillMissing))
+    return 4;
+
+  // After it is added to the symbol table, and not only the operation with
+  // which the table is associated, it can be looked up.
+  mlirSymbolTableInsert(symbolTable, operation);
+  MlirOperation funcQux =
+      mlirSymbolTableLookup(symbolTable, mlirStringRefCreateFromCString("qux"));
+  if (!mlirOperationEqual(operation, funcQux))
+    return 5;
+
+  // Erasing from the symbol table also removes the operation.
+  mlirSymbolTableErase(symbolTable, funcBar);
+  MlirOperation nowMissing =
+      mlirSymbolTableLookup(symbolTable, mlirStringRefCreateFromCString("bar"));
+  if (!mlirOperationIsNull(nowMissing))
+    return 6;
+
+  // Adding a symbol with the same name to the table should rename.
+  MlirOperation duplicateNameOp = mlirBlockGetFirstOperation(otherModuleBody);
+  mlirOperationRemoveFromParent(duplicateNameOp);
+  mlirBlockAppendOwnedOperation(moduleBody, duplicateNameOp);
+  MlirAttribute newName = mlirSymbolTableInsert(symbolTable, duplicateNameOp);
+  MlirStringRef newNameStr = mlirStringAttrGetValue(newName);
+  if (mlirStringRefEqual(newNameStr, mlirStringRefCreateFromCString("foo")))
+    return 7;
+  MlirAttribute updatedName = mlirOperationGetAttributeByName(
+      duplicateNameOp, mlirSymbolTableGetSymbolAttributeName());
+  if (!mlirAttributeEqual(updatedName, newName))
+    return 8;
+
+  mlirOperationDump(mlirModuleGetOperation(module));
+  mlirOperationDump(mlirModuleGetOperation(otherModule));
+  // clang-format off
+  // CHECK-LABEL: @testSymbolTable
+  // CHECK: module
+  // CHECK:   func private @foo
+  // CHECK:   func private @qux
+  // CHECK:   func private @foo{{.+}}
+  // CHECK: module
+  // CHECK-NOT: @qux
+  // CHECK-NOT: @foo
+  // clang-format on
+
+  mlirSymbolTableDestroy(symbolTable);
+  mlirModuleDestroy(module);
+  mlirModuleDestroy(otherModule);
+
+  return 0;
+}
+
+void testDiagnostics() {
+  MlirContext ctx = mlirContextCreate();
+  MlirDiagnosticHandlerID id = mlirContextAttachDiagnosticHandler(
+      ctx, errorHandler, (void *)42, deleteUserData);
+  fprintf(stderr, "@test_diagnostics\n");
+  MlirLocation unknownLoc = mlirLocationUnknownGet(ctx);
+  mlirEmitError(unknownLoc, "test diagnostics");
+  MlirLocation fileLineColLoc = mlirLocationFileLineColGet(
+      ctx, mlirStringRefCreateFromCString("file.c"), 1, 2);
+  mlirEmitError(fileLineColLoc, "test diagnostics");
+  MlirLocation callSiteLoc = mlirLocationCallSiteGet(
+      mlirLocationFileLineColGet(
+          ctx, mlirStringRefCreateFromCString("other-file.c"), 2, 3),
+      fileLineColLoc);
+  mlirEmitError(callSiteLoc, "test diagnostics");
+  MlirLocation null = {0};
+  MlirLocation nameLoc =
+      mlirLocationNameGet(ctx, mlirStringRefCreateFromCString("named"), null);
+  mlirEmitError(nameLoc, "test diagnostics");
+  MlirLocation locs[2] = {nameLoc, callSiteLoc};
+  MlirAttribute nullAttr = {0};
+  MlirLocation fusedLoc = mlirLocationFusedGet(ctx, 2, locs, nullAttr);
+  mlirEmitError(fusedLoc, "test diagnostics");
+  mlirContextDetachDiagnosticHandler(ctx, id);
+  mlirEmitError(unknownLoc, "more test diagnostics");
+  // CHECK-LABEL: @test_diagnostics
+  // CHECK: processing diagnostic (userData: 42) <<
+  // CHECK:   test diagnostics
+  // CHECK:   loc(unknown)
+  // CHECK: >> end of diagnostic (userData: 42)
+  // CHECK: processing diagnostic (userData: 42) <<
+  // CHECK:   test diagnostics
+  // CHECK:   loc("file.c":1:2)
+  // CHECK: >> end of diagnostic (userData: 42)
+  // CHECK: processing diagnostic (userData: 42) <<
+  // CHECK:   test diagnostics
+  // CHECK:   loc(callsite("other-file.c":2:3 at "file.c":1:2))
+  // CHECK: >> end of diagnostic (userData: 42)
+  // CHECK: processing diagnostic (userData: 42) <<
+  // CHECK:   test diagnostics
+  // CHECK:   loc("named")
+  // CHECK: >> end of diagnostic (userData: 42)
+  // CHECK: processing diagnostic (userData: 42) <<
+  // CHECK:   test diagnostics
+  // CHECK:   loc(fused["named", callsite("other-file.c":2:3 at "file.c":1:2)])
+  // CHECK: deleting user data (userData: 42)
+  // CHECK-NOT: processing diagnostic
+  // CHECK:     more test diagnostics
+  mlirContextDestroy(ctx);
+}
+
 int main() {
   MlirContext ctx = mlirContextCreate();
   mlirRegisterAllDialects(ctx);
   if (constructAndTraverseIr(ctx))
     return 1;
   buildWithInsertionsAndPrint(ctx);
   if (createOperationWithTypeInference(ctx))
     return 2;
int main() {
   if (registerOnlyStd())
     return 9;
   if (testBackreferences())
     return 10;
   if (testOperands())
     return 11;
   if (testClone())
     return 12;
-  if (testTypeID(ctx)) {
+  if (testTypeID(ctx))
     return 13;
-  }
+  if (testSymbolTable(ctx))
+    return 14;
 
   mlirContextDestroy(ctx);
 
   testDiagnostics();
   return 0;
 }

mlir/test/python/ir/operation.py

 
 
 # CHECK-LABEL: TEST: testAppendMoveFromAnotherBlock
 @run
 def testAppendMoveFromAnotherBlock():
   with Context():
     m1 = Module.parse("func private @foo()")
     m2 = Module.parse("func private @bar()")
-    func = m1.body.operations[0]  
+    func = m1.body.operations[0]
     m2.body.append(func)
 
     # CHECK: module
     # CHECK: func private @bar
     # CHECK: func private @foo
 
     print(m2)
     # CHECK: module {
with Context():
     except ValueError as e:
       if "has no parent" not in str(e):
         raise
     else:
       assert False, "expected ValueError when detaching a detached operation"
 
     print(m1)
     # CHECK-NOT: func private @foo
+
+
+# CHECK-LABEL: TEST: testSymbolTable
+@run
+def testSymbolTable():
+  with Context() as ctx:
+    ctx.allow_unregistered_dialects = True
+    m1 = Module.parse("""
+      func private @foo()
+      func private @bar()""")
+    m2 = Module.parse("""
+      func private @qux()
+      func private @foo()
+      "foo.bar"() : () -> ()""")
+
+    symbol_table = SymbolTable(m1.operation)
+
+    # CHECK: func private @foo
+    # CHECK: func private @bar
+    assert "foo" in symbol_table
+    print(symbol_table["foo"])
+    assert "bar" in symbol_table
+    bar = symbol_table["bar"]
+    print(symbol_table["bar"])
+
+    assert "qux" not in symbol_table
+
+    del symbol_table["bar"]
+    try:
+      symbol_table.erase(symbol_table["bar"])
+    except KeyError:
+      pass
+    else:
+      assert False, "expected KeyError"
+
+    # CHECK: module
+    # CHECK:   func private @foo()
+    print(m1)
+    assert "bar" not in symbol_table
+
+    try:
+      print(bar)
+    except RuntimeError as e:
+      if "the operation has been invalidated" not in str(e):
+        raise
+    else:
+      assert False, "expected RuntimeError due to invalidated operation"
+
+    qux = m2.body.operations[0]
+    m1.body.append(qux)
+    symbol_table.insert(qux)
+    assert "qux" in symbol_table
+
+    # Check that insertion actually renames this symbol in the symbol table.
+    foo2 = m2.body.operations[0]
+    m1.body.append(foo2)
+    updated_name = symbol_table.insert(foo2)
+    assert foo2.name.value != "foo"
+    assert foo2.name == updated_name
+
+    # CHECK: module
+    # CHECK:   func private @foo()
+    # CHECK:   func private @qux()
+    # CHECK:   func private @foo{{.*}}
+    print(m1)
+
+    try:
+      symbol_table.insert(m2.body.operations[0])
+    except ValueError as e:
+      if "Expected operation to have a symbol name" not in str(e):
+        raise
+    else:
+      assert False, "exepcted ValueError when adding a non-symbol"