Add SymbolRefAttr to python bindings

[makslevental]

use MlirAttribute instead

[makslevental]

add AttrBuilder.get('SymbolNameAttr') test

[rkayaith]

mlir/lib/Bindings/Python/IRAttributes.cpp

452

should this be removed from the CAPI, since it's not a unique typeid? (similar to how BoolAttr doesn't have its typeid exposed)

460–488

not sure trying to parse @/::s here is a good idea, as far as I know they're valid characters in symbol names:

module @"@foo::bar" {}  // no complaints from the verifier

better to leave it to Attribute.parse to handle all that imo

510–518

Doesn't str() effectively do this already, since it prints the asm string? I would return List[str] here instead.

[makslevental]

mlir/lib/Bindings/Python/IRAttributes.cpp

460–488

fair enough

[makslevental]

don't parse

[makslevental]

handle degenerate case in python rather than cpp

[rkayaith]

mlir/test/python/ir/attributes.py

248–252

nit: default_get is used in both of these prints, can you change one of them

[makslevental]

rebase

[makslevental]

format

mlir/lib/Bindings/Python/IRAttributes.cpp

 //===- IRAttributes.cpp - Exports builtin and standard attributes ---------===//
 //
 // 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 <optional>
+#include <pybind11/detail/common.h>
+#include <stdexcept>
+#include <string>
 #include <utility>
+#include <vector>
 
 #include "IRModule.h"
 
 #include "PybindUtils.h"
 
 #include "mlir-c/BuiltinAttributes.h"
 #include "mlir-c/BuiltinTypes.h"
+#include "mlir-c/Support.h"
 #include "mlir/Bindings/Python/PybindAdaptors.h"
+#include "mlir/CAPI/Support.h"
+#include "llvm/ADT/StringRef.h"
 
 namespace py = pybind11;
 using namespace mlir;
 using namespace mlir::python;
 
 using llvm::SmallVector;
 
 //------------------------------------------------------------------------------
c.def_static(
         "Gets an uniqued bool attribute");
     c.def_property_readonly(
         "value",
         [](PyBoolAttribute &self) { return mlirBoolAttrGetValue(self); },
         "Returns the value of the bool attribute");
   }
 };
 
+class PySymbolRefAttribute : public PyConcreteAttribute<PySymbolRefAttribute> {
+public:
+  static constexpr IsAFunctionTy isaFunction = mlirAttributeIsASymbolRef;
+  static constexpr const char *pyClassName = "SymbolRefAttr";
+  using PyConcreteAttribute::PyConcreteAttribute;
+
+  static MlirAttribute fromList(const std::vector<std::string> &symbols,
+                                PyMlirContext &context) {
+    if (symbols.empty())
+      throw std::runtime_error("SymbolRefAttr must be composed of at least "
+                               "one symbol.");
+    std::string rootSym = symbols[0];
+    if (rootSym[0] == '@')
+      rootSym = rootSym.substr(1);
+    MlirStringRef rootSymbol = toMlirStringRef(rootSym);
+    SmallVector<MlirAttribute, 3> referenceAttrs;
+    for (size_t i = 1; i < symbols.size(); ++i) {
+      std::string sym = symbols[i];
+      if (sym[0] == '@')
+        sym = sym.substr(1);
+      referenceAttrs.push_back(
+          mlirFlatSymbolRefAttrGet(context.get(), toMlirStringRef(sym)));
+    }
+    return mlirSymbolRefAttrGet(context.get(), rootSymbol,
+                                referenceAttrs.size(), referenceAttrs.data());
+  }
+
+  static MlirAttribute fromStr(const std::string &symbols,
+                               DefaultingPyMlirContext context) {
+    std::vector<std::string> symbols_;
+    llvm::StringRef symbolsRef(symbols);
+    while (!symbolsRef.empty()) {
+      auto [head, tail] = symbolsRef.split("::");
+      symbols_.push_back(head.str());
+      symbolsRef = tail;
+    }
+    return PySymbolRefAttribute::fromList(symbols_, context.resolve());
+  }
+
+  static void bindDerived(ClassTy &c) {
+    c.def_static(
+        "get",
+        [](const std::vector<std::string> &symbols,
+           DefaultingPyMlirContext context) {
+          return PySymbolRefAttribute::fromList(symbols, context.resolve());
+        },
+        py::arg("symbols"), py::arg("context") = py::none(),
+        "Gets a uniqued SymbolRef attribute");
+    c.def_static(
+        "get",
+        [](const std::string &symbols, DefaultingPyMlirContext context) {
+          return PySymbolRefAttribute::fromStr(symbols, context.resolve());
+        },
+        py::arg("symbols"), py::arg("context") = py::none(),
+        "Gets a uniqued SymbolRef attribute");
+    c.def_property_readonly(
+        "value",
+        [](PySymbolRefAttribute &self) {
+          std::string symbol =
+              "@" + unwrap(mlirSymbolRefAttrGetRootReference(self)).str();
+          for (int i = 0; i < mlirSymbolRefAttrGetNumNestedReferences(self);
+               ++i)
+            symbol += "::@" +
+                      unwrap(mlirSymbolRefAttrGetRootReference(
+                                 mlirSymbolRefAttrGetNestedReference(self, i)))
+                          .str();
+          return symbol;
+        },
+        "Returns the value of the SymbolRef attribute as a string");
+  }
+};
+
 class PyFlatSymbolRefAttribute
-    : public PyConcreteAttribute<PyFlatSymbolRefAttribute> {
+    : public PyConcreteAttribute<PyFlatSymbolRefAttribute,
+                                 PySymbolRefAttribute> {
 public:
   static constexpr IsAFunctionTy isaFunction = mlirAttributeIsAFlatSymbolRef;
   static constexpr const char *pyClassName = "FlatSymbolRefAttr";
   using PyConcreteAttribute::PyConcreteAttribute;
-  static constexpr GetTypeIDFunctionTy getTypeIdFunction =
-      mlirFlatSymbolRefAttrGetTypeID;
 
   static void bindDerived(ClassTy &c) {
     c.def_static(
         "get",
-        [](std::string value, DefaultingPyMlirContext context) {
-          MlirAttribute attr =
-              mlirFlatSymbolRefAttrGet(context->get(), toMlirStringRef(value));
-          return PyFlatSymbolRefAttribute(context->getRef(), attr);
+        [](const std::string &value, DefaultingPyMlirContext context) {
+          return PySymbolRefAttribute::fromStr(value, context.resolve());
         },
         py::arg("value"), py::arg("context") = py::none(),
         "Gets a uniqued FlatSymbolRef attribute");
     c.def_property_readonly(
         "value",
         [](PyFlatSymbolRefAttribute &self) {
           MlirStringRef stringRef = mlirFlatSymbolRefAttrGetValue(self);
           return py::str(stringRef.data, stringRef.length);
py::object integerOrBoolAttributeCaster(PyAttribute &pyAttribute) {
   if (PyIntegerAttribute::isaFunction(pyAttribute))
     return py::cast(PyIntegerAttribute(pyAttribute));
   std::string msg =
       std::string("Can't cast unknown element type DenseArrayAttr (") +
       std::string(py::repr(py::cast(pyAttribute))) + ")";
   throw py::cast_error(msg);
 }
 
+py::object symbolRefOrFlatSymbolRefAttributeCaster(PyAttribute &pyAttribute) {
+  if (PyFlatSymbolRefAttribute::isaFunction(pyAttribute))
+    return py::cast(PyFlatSymbolRefAttribute(pyAttribute));
+  if (PySymbolRefAttribute::isaFunction(pyAttribute))
+    return py::cast(PySymbolRefAttribute(pyAttribute));
+  std::string msg = std::string("Can't cast unknown SymbolRef attribute (") +
+                    std::string(py::repr(py::cast(pyAttribute))) + ")";
+  throw py::cast_error(msg);
+}
+
 } // namespace
 
 void mlir::python::populateIRAttributes(py::module &m) {
   PyAffineMapAttribute::bind(m);
 
   PyDenseBoolArrayAttribute::bind(m);
   PyDenseBoolArrayAttribute::PyDenseArrayIterator::bind(m);
   PyDenseI8ArrayAttribute::bind(m);
void mlir::python::populateIRAttributes(py::module &m) {
   PyDenseElementsAttribute::bind(m);
   PyDenseFPElementsAttribute::bind(m);
   PyDenseIntElementsAttribute::bind(m);
   PyGlobals::get().registerTypeCaster(
       mlirDenseIntOrFPElementsAttrGetTypeID(),
       pybind11::cpp_function(denseIntOrFPElementsAttributeCaster));
 
   PyDictAttribute::bind(m);
+  PySymbolRefAttribute::bind(m);
+  PyGlobals::get().registerTypeCaster(
+      mlirSymbolRefAttrGetTypeID(),
+      pybind11::cpp_function(symbolRefOrFlatSymbolRefAttributeCaster));
+
   PyFlatSymbolRefAttribute::bind(m);
   PyOpaqueAttribute::bind(m);
   PyFloatAttribute::bind(m);
   PyIntegerAttribute::bind(m);
   PyStringAttribute::bind(m);
   PyTypeAttribute::bind(m);
   PyGlobals::get().registerTypeCaster(
       mlirIntegerAttrGetTypeID(),
       pybind11::cpp_function(integerOrBoolAttributeCaster));
   PyUnitAttribute::bind(m);
 
   PyStridedLayoutAttribute::bind(m);
 }

mlir/python/mlir/ir.py

 
 @register_attribute_builder("SymbolNameAttr")
 def _symbolNameAttr(x, context):
     return StringAttr.get(x, context=context)
 
 
 @register_attribute_builder("SymbolRefAttr")
 def _symbolRefAttr(x, context):
+    return SymbolRefAttr.get(x, context=context)
+
+
+@register_attribute_builder("FlatSymbolRefAttr")
+def _flatSymbolRefAttr(x, context):
     return FlatSymbolRefAttr.get(x, context=context)
 
 
 @register_attribute_builder("ArrayAttr")
 def _arrayAttr(x, context):
     return ArrayAttr.get(x, context=context)
 
 

 def _f64ArrayAttr(x, context):
     return ArrayAttr.get([_f64Attr(v, context) for v in x])
 
 
 @register_attribute_builder("DenseI64ArrayAttr")
 def _denseI64ArrayAttr(x, context):
     return DenseI64ArrayAttr.get(x, context=context)
 
+
 @register_attribute_builder("DenseBoolArrayAttr")
 def _denseBoolArrayAttr(x, context):
     return DenseBoolArrayAttr.get(x, context=context)
 
 
 @register_attribute_builder("TypeAttr")
 def _typeAttr(x, context):
     return TypeAttr.get(x, context=context)

mlir/test/python/ir/attributes.py

with Context() as ctx:
         # CHECK: default_get: true
         print("default_get:", BoolAttr.get(True))
 
 
 # CHECK-LABEL: TEST: testFlatSymbolRefAttr
 @run
 def testFlatSymbolRefAttr():
     with Context() as ctx:
-        sattr = FlatSymbolRefAttr(Attribute.parse("@symbol"))
+        sattr = Attribute.parse("@symbol")
         # CHECK: symattr value: symbol
         print("symattr value:", sattr.value)
 
         # Test factory methods.
         # CHECK: default_get: @foobar
         print("default_get:", FlatSymbolRefAttr.get("foobar"))
 
 
+# CHECK-LABEL: TEST: testSymbolRefAttr
+@run
+def testSymbolRefAttr():
+    with Context() as ctx:
+        sattr = Attribute.parse("@symbol1::@symbol2")
+        # CHECK: symattr value: @symbol1::@symbol2
+        print("symattr value:", sattr.value)
+
+        # Test factory methods.
+        # CHECK: default_get: @symbol1::@symbol2
+        print("default_get:", SymbolRefAttr.get("symbol1::symbol2"))
+
+        # CHECK: default_get: @symbol1::@symbol2
+        print("default_get:", SymbolRefAttr.get(["symbol1", "symbol2"]))
+
+        # CHECK: default_get: @symbol1::@symbol2
+        print("default_get:", SymbolRefAttr.get("@symbol1::@symbol2"))
+
+        # CHECK: default_get: @symbol1::@symbol2
+        print("default_get:", SymbolRefAttr.get(["@symbol1", "@symbol2"]))
+
+
 # CHECK-LABEL: TEST: testOpaqueAttr
 @run
 def testOpaqueAttr():
     with Context() as ctx:
         ctx.allow_unregistered_dialects = True
         oattr = OpaqueAttr(Attribute.parse("#pytest_dummy.dummyattr<>"))
         # CHECK: oattr value: pytest_dummy
         print("oattr value:", oattr.dialect_namespace)