[mlir][python] Implement more SymbolTable methods.

[stellaraccident]

Summary

• set_symbol_name, get_symbol_name, set_visibility, get_visibility, replace_all_symbol_uses, walk_symbol_tables
• In integrations I've been doing, I've been reaching for all of these to do both general IR manipulation and module merging.
• I don't love the replace_all_symbol_uses underlying APIs since they necessitate SYMBOL_COUNT walks and have various sharp edges. I'm hoping that whatever emerges eventually for this can still retain this simple API as a one-shot.

[ftynse]

mlir/include/mlir-c/IR.h

797

Nit: this is already mentioned above.

809

Other similar callbacks, in particular in IR.h, take userData as the last argument. Let's do the same here for API consistency.

mlir/lib/Bindings/Python/IRCore.cpp

1599

Maybe this should return the attribute. I'm okay with both as long as it's the same for the entire SymbolTable API.

1601

Nit: MLIR code tends to expand auto whenever possible, e.g., here it's not clear if we are getting a PyOperation& or an Operation&.

1701–1703

Nit: throw std::runtime_error should work here.

1703

I don't think operator , works for Twine, did this mean + ?

[stellaraccident]

mlir/lib/Bindings/Python/IRCore.cpp

1703

Thanks for the catch.

[stellaraccident]

Comments and rebase.

[stellaraccident]

Add pyi bindings.

mlir/include/mlir-c/IR.h

 //===----------------------------------------------------------------------===//
 // Symbol and SymbolTable API.
 //===----------------------------------------------------------------------===//
 
 /// Returns the name of the attribute used to store symbol names compatible with
 /// symbol tables.
 MLIR_CAPI_EXPORTED MlirStringRef mlirSymbolTableGetSymbolAttributeName();
 
+/// Returns the name of the attribute used to store symbol visibility.
+MLIR_CAPI_EXPORTED MlirStringRef mlirSymbolTableGetVisibilityAttributeName();
+
 /// 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;

 /// 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);
 
+/// Attempt to replace all uses that are nested within the given operation
+/// of the given symbol 'oldSymbol' with the provided 'newSymbol'. This does
+/// not traverse into nested symbol tables. Will fail atomically if there are
+/// any unknown operations that may be potential symbol tables.
+MLIR_CAPI_EXPORTED MlirLogicalResult mlirSymbolTableReplaceAllSymbolUses(
+    MlirStringRef oldSymbol, MlirStringRef newSymbol, MlirOperation from);
+
+/// Walks all symbol table operations nested within, and including, `op`. For
+/// each symbol table operation, the provided callback is invoked with the op
+/// and a boolean signifying if the symbols within that symbol table can be
+/// treated as if all uses within the IR are visible to the caller.
+/// `allSymUsesVisible` identifies whether all of the symbol uses of symbols
+/// within `op` are visible.
+MLIR_CAPI_EXPORTED void mlirSymbolTableWalkSymbolTables(
+    MlirOperation from, bool allSymUsesVisible,
+    void (*callback)(MlirOperation, bool, void *userData), void *userData);
+
 #ifdef __cplusplus
 }
 #endif
 
 #endif // MLIR_C_IR_H

mlir/lib/Bindings/Python/IRCore.cpp

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()));
 }
 
+PyAttribute PySymbolTable::getSymbolName(PyOperationBase &symbol) {
+  // Op must already be a symbol.
+  PyOperation &operation = symbol.getOperation();
+  operation.checkValid();
+  MlirStringRef attrName = mlirSymbolTableGetSymbolAttributeName();
+  MlirAttribute existingNameAttr =
+      mlirOperationGetAttributeByName(operation.get(), attrName);
+  if (mlirAttributeIsNull(existingNameAttr))
+    throw py::value_error("Expected operation to have a symbol name.");
+  return PyAttribute(symbol.getOperation().getContext(), existingNameAttr);
+}
+
+void PySymbolTable::setSymbolName(PyOperationBase &symbol,
+                                  const std::string &name) {
+  // Op must already be a symbol.
+  PyOperation &operation = symbol.getOperation();
+  operation.checkValid();
+  MlirStringRef attrName = mlirSymbolTableGetSymbolAttributeName();
+  MlirAttribute existingNameAttr =
+      mlirOperationGetAttributeByName(operation.get(), attrName);
+  if (mlirAttributeIsNull(existingNameAttr))
+    throw py::value_error("Expected operation to have a symbol name.");
+  MlirAttribute newNameAttr =
+      mlirStringAttrGet(operation.getContext()->get(), toMlirStringRef(name));
+  mlirOperationSetAttributeByName(operation.get(), attrName, newNameAttr);
+}
+
+PyAttribute PySymbolTable::getVisibility(PyOperationBase &symbol) {
+  PyOperation &operation = symbol.getOperation();
+  operation.checkValid();
+  MlirStringRef attrName = mlirSymbolTableGetVisibilityAttributeName();
+  MlirAttribute existingVisAttr =
+      mlirOperationGetAttributeByName(operation.get(), attrName);
+  if (mlirAttributeIsNull(existingVisAttr))
+    throw py::value_error("Expected operation to have a symbol visibility.");
+  return PyAttribute(symbol.getOperation().getContext(), existingVisAttr);
+}
+
+void PySymbolTable::setVisibility(PyOperationBase &symbol,
+                                  const std::string &visibility) {
+  if (visibility != "public" && visibility != "private" &&
+      visibility != "nested")
+    throw py::value_error(
+        "Expected visibility to be 'public', 'private' or 'nested'");
+  PyOperation &operation = symbol.getOperation();
+  operation.checkValid();
+  MlirStringRef attrName = mlirSymbolTableGetVisibilityAttributeName();
+  MlirAttribute existingVisAttr =
+      mlirOperationGetAttributeByName(operation.get(), attrName);
+  if (mlirAttributeIsNull(existingVisAttr))
+    throw py::value_error("Expected operation to have a symbol visibility.");
+  MlirAttribute newVisAttr = mlirStringAttrGet(operation.getContext()->get(),
+                                               toMlirStringRef(visibility));
+  mlirOperationSetAttributeByName(operation.get(), attrName, newVisAttr);
+}
+
+void PySymbolTable::replaceAllSymbolUses(const std::string &oldSymbol,
+                                         const std::string &newSymbol,
+                                         PyOperationBase &from) {
+  PyOperation &fromOperation = from.getOperation();
+  fromOperation.checkValid();
+  if (mlirLogicalResultIsFailure(mlirSymbolTableReplaceAllSymbolUses(
+          toMlirStringRef(oldSymbol), toMlirStringRef(newSymbol),
+          from.getOperation())))
+
+    throw py::value_error("Symbol rename failed");
+}
+
+void PySymbolTable::walkSymbolTables(PyOperationBase &from,
+                                     bool allSymUsesVisible,
+                                     py::object callback) {
+  PyOperation &fromOperation = from.getOperation();
+  fromOperation.checkValid();
+  struct UserData {
+    PyMlirContextRef context;
+    py::object callback;
+    bool gotException;
+    std::string exceptionWhat;
+    py::object exceptionType;
+  };
+  UserData userData{
+      fromOperation.getContext(), std::move(callback), false, {}, {}};
+  mlirSymbolTableWalkSymbolTables(
+      fromOperation.get(), allSymUsesVisible,
+      [](MlirOperation foundOp, bool isVisible, void *calleeUserDataVoid) {
+        UserData *calleeUserData = static_cast<UserData *>(calleeUserDataVoid);
+        auto pyFoundOp =
+            PyOperation::forOperation(calleeUserData->context, foundOp);
+        if (calleeUserData->gotException)
+          return;
+        try {
+          calleeUserData->callback(pyFoundOp.getObject(), isVisible);
+        } catch (py::error_already_set &e) {
+          calleeUserData->gotException = true;
+          calleeUserData->exceptionWhat = e.what();
+          calleeUserData->exceptionType = e.type();
+        }
+      },
+      static_cast<void *>(&userData));
+  if (userData.gotException) {
+    std::string message("Exception raised in callback: ");
+    message.append(userData.exceptionWhat);
+    throw std::runtime_error(std::move(message));
+  }
+}
+
 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:
void mlir::python::populateIRCore(py::module &m) {
   // Mapping of SymbolTable.
   //----------------------------------------------------------------------------
   py::class_<PySymbolTable>(m, "SymbolTable", py::module_local())
       .def(py::init<PyOperationBase &>())
       .def("__getitem__", &PySymbolTable::dunderGetItem)
       .def("insert", &PySymbolTable::insert, py::arg("operation"))
       .def("erase", &PySymbolTable::erase, py::arg("operation"))
       .def("__delitem__", &PySymbolTable::dunderDel)
-      .def("__contains__", [](PySymbolTable &table, const std::string &name) {
+      .def("__contains__",
+           [](PySymbolTable &table, const std::string &name) {
              return !mlirOperationIsNull(mlirSymbolTableLookup(
                  table, mlirStringRefCreate(name.data(), name.length())));
-      });
+           })
+      // Static helpers.
+      .def_static("set_symbol_name", &PySymbolTable::setSymbolName,
+                  py::arg("symbol"), py::arg("name"))
+      .def_static("get_symbol_name", &PySymbolTable::getSymbolName,
+                  py::arg("symbol"))
+      .def_static("get_visibility", &PySymbolTable::getVisibility,
+                  py::arg("symbol"))
+      .def_static("set_visibility", &PySymbolTable::setVisibility,
+                  py::arg("symbol"), py::arg("visibility"))
+      .def_static("replace_all_symbol_uses",
+                  &PySymbolTable::replaceAllSymbolUses, py::arg("old_symbol"),
+                  py::arg("new_symbol"), py::arg("from_op"))
+      .def_static("walk_symbol_tables", &PySymbolTable::walkSymbolTables,
+                  py::arg("from_op"), py::arg("all_sym_uses_visible"),
+                  py::arg("callback"));
 
   // 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

public:
   /// 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);
 
+  /// Gets and sets the name of a symbol op.
+  static PyAttribute getSymbolName(PyOperationBase &symbol);
+  static void setSymbolName(PyOperationBase &symbol, const std::string &name);
+
+  /// Gets and sets the visibility of a symbol op.
+  static PyAttribute getVisibility(PyOperationBase &symbol);
+  static void setVisibility(PyOperationBase &symbol,
+                            const std::string &visibility);
+
+  /// Replaces all symbol uses within an operation. See the API
+  /// mlirSymbolTableReplaceAllSymbolUses for all caveats.
+  static void replaceAllSymbolUses(const std::string &oldSymbol,
+                                   const std::string &newSymbol,
+                                   PyOperationBase &from);
+
+  /// Walks all symbol tables under and including 'from'.
+  static void walkSymbolTables(PyOperationBase &from, bool allSymUsesVisible,
+                               pybind11::object callback);
+
   /// Casts the bindings class into the C API structure.
   operator MlirSymbolTable() { return symbolTable; }
 
 private:
   PyOperationRef operation;
   MlirSymbolTable symbolTable;
 };
 

mlir/lib/CAPI/IR/IR.cpp

 //===----------------------------------------------------------------------===//
 // Symbol and SymbolTable API.
 //===----------------------------------------------------------------------===//
 
 MlirStringRef mlirSymbolTableGetSymbolAttributeName() {
   return wrap(SymbolTable::getSymbolAttrName());
 }
 
+MlirStringRef mlirSymbolTableGetVisibilityAttributeName() {
+  return wrap(SymbolTable::getVisibilityAttrName());
+}
+
 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((Attribute)unwrap(symbolTable)->insert(unwrap(operation)));
 }
 
 void mlirSymbolTableErase(MlirSymbolTable symbolTable,
                           MlirOperation operation) {
   unwrap(symbolTable)->erase(unwrap(operation));
 }
+
+MlirLogicalResult mlirSymbolTableReplaceAllSymbolUses(MlirStringRef oldSymbol,
+                                                      MlirStringRef newSymbol,
+                                                      MlirOperation from) {
+  auto cppFrom = unwrap(from);
+  auto *context = cppFrom->getContext();
+  auto oldSymbolAttr = StringAttr::get(unwrap(oldSymbol), context);
+  auto newSymbolAttr = StringAttr::get(unwrap(newSymbol), context);
+  return wrap(SymbolTable::replaceAllSymbolUses(oldSymbolAttr, newSymbolAttr,
+                                                unwrap(from)));
+}
+
+void mlirSymbolTableWalkSymbolTables(MlirOperation from, bool allSymUsesVisible,
+                                     void (*callback)(MlirOperation, bool,
+                                                      void *userData),
+                                     void *userData) {
+  SymbolTable::walkSymbolTables(unwrap(from), allSymUsesVisible,
+                                [&](Operation *foundOpCpp, bool isVisible) {
+                                  callback(wrap(foundOpCpp), isVisible,
+                                           userData);
+                                });
+}

mlir/test/python/ir/operation.py

except ValueError as 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"
-
-
 # CHECK-LABEL: TEST: testOperationHash
 @run
 def testOperationHash():
   ctx = Context()
   ctx.allow_unregistered_dialects = True
   with ctx, Location.unknown():
     op = Operation.create("custom.op1")
     assert hash(op) == hash(op.operation)

mlir/test/python/ir/symbol_table.py

+# RUN: %PYTHON %s | FileCheck %s
+
+import gc
+import io
+import itertools
+from mlir.ir import *
+
+
+def run(f):
+  print("\nTEST:", f.__name__)
+  f()
+  gc.collect()
+  assert Context._get_live_count() == 0
+  return f
+
+
+# CHECK-LABEL: TEST: testSymbolTableInsert
+@run
+def testSymbolTableInsert():
+  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"
+
+
+# CHECK-LABEL: testSymbolTableRAUW
+@run
+def testSymbolTableRAUW():
+  with Context() as ctx:
+    m = Module.parse("""
+      func private @foo() {
+        call @bar() : () -> ()
+        return
+      }
+      func private @bar()
+      """)
+    foo, bar = list(m.operation.regions[0].blocks[0].operations)[0:2]
+    SymbolTable.set_symbol_name(bar, "bam")
+    # Note that module.operation counts as a "nested symbol table" which won't
+    # be traversed into, so it is necessary to traverse its children.
+    SymbolTable.replace_all_symbol_uses("bar", "bam", foo)
+    # CHECK: call @bam()
+    # CHECK: func private @bam
+    print(m)
+    # CHECK: Foo symbol: "foo"
+    # CHECK: Bar symbol: "bam"
+    print(f"Foo symbol: {SymbolTable.get_symbol_name(foo)}")
+    print(f"Bar symbol: {SymbolTable.get_symbol_name(bar)}")
+
+
+# CHECK-LABEL: testSymbolTableVisibility
+@run
+def testSymbolTableVisibility():
+  with Context() as ctx:
+    m = Module.parse("""
+      func private @foo() {
+        return
+      }
+      """)
+    foo = m.operation.regions[0].blocks[0].operations[0]
+    # CHECK: Existing visibility: "private"
+    print(f"Existing visibility: {SymbolTable.get_visibility(foo)}")
+    SymbolTable.set_visibility(foo, "public")
+    # CHECK: func public @foo
+    print(m)
+
+
+# CHECK: testWalkSymbolTables
+@run
+def testWalkSymbolTables():
+  with Context() as ctx:
+    m = Module.parse("""
+      module @outer {
+        module @inner{
+        }
+      }
+      """)
+    def callback(symbol_table_op, uses_visible):
+      print(f"SYMBOL TABLE: {uses_visible}: {symbol_table_op}")
+    # CHECK: SYMBOL TABLE: True: module @inner
+    # CHECK: SYMBOL TABLE: True: module @outer
+    SymbolTable.walk_symbol_tables(m.operation, True, callback)
+
+    # Make sure exceptions in the callback are handled.
+    def error_callback(symbol_table_op, uses_visible):
+      assert False, "Raised from python"
+    try:
+      SymbolTable.walk_symbol_tables(m.operation, True, error_callback)
+    except RuntimeError as e:
+      # CHECK: GOT EXCEPTION: Exception raised in callback: AssertionError: Raised from python
+      print(f"GOT EXCEPTION: {e}")
+