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)

453–481

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

503–511

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

453–481

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/include/mlir-c/BuiltinAttributes.h

 MLIR_CAPI_EXPORTED MlirAttribute mlirFlatSymbolRefAttrGet(MlirContext ctx,
                                                           MlirStringRef symbol);
 
 /// Returns the referenced symbol as a string reference. The data remains live
 /// as long as the context in which the attribute lives.
 MLIR_CAPI_EXPORTED MlirStringRef
 mlirFlatSymbolRefAttrGetValue(MlirAttribute attr);
 
-/// Returns the typeID of an FlatSymbolRef attribute.
-MLIR_CAPI_EXPORTED MlirTypeID mlirFlatSymbolRefAttrGetTypeID(void);
-
 //===----------------------------------------------------------------------===//
 // Type attribute.
 //===----------------------------------------------------------------------===//
 
 /// Checks whether the given attribute is a type attribute.
 MLIR_CAPI_EXPORTED bool mlirAttributeIsAType(MlirAttribute attr);
 
 /// Creates a type attribute wrapping the given type in the same context as the

mlir/lib/Bindings/Python/IRAttributes.cpp

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.");
+    MlirStringRef rootSymbol = toMlirStringRef(symbols[0]);
+    SmallVector<MlirAttribute, 3> referenceAttrs;
+    for (size_t i = 1; i < symbols.size(); ++i) {
+      referenceAttrs.push_back(
+          mlirFlatSymbolRefAttrGet(context.get(), toMlirStringRef(symbols[i])));
+    }
+    return mlirSymbolRefAttrGet(context.get(), rootSymbol,
+                                referenceAttrs.size(), referenceAttrs.data());
+  }
+
+  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 from a list of symbol names");
+    c.def_property_readonly(
+        "value",
+        [](PySymbolRefAttribute &self) {
+          std::vector<std::string> symbols = {
+              unwrap(mlirSymbolRefAttrGetRootReference(self)).str()};
+          for (int i = 0; i < mlirSymbolRefAttrGetNumNestedReferences(self);
+               ++i)
+            symbols.push_back(
+                unwrap(mlirSymbolRefAttrGetRootReference(
+                           mlirSymbolRefAttrGetNestedReference(self, i)))
+                    .str());
+          return symbols;
+        },
+        "Returns the value of the SymbolRef attribute as a list[str]");
+  }
+};
+
 class PyFlatSymbolRefAttribute
     : public PyConcreteAttribute<PyFlatSymbolRefAttribute> {
 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);
         },
         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);
         },
         "Returns the value of the FlatSymbolRef attribute as a string");
   }
 };
 
 class PyOpaqueAttribute : public PyConcreteAttribute<PyOpaqueAttribute> {
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/lib/Bindings/Python/IRCore.cpp

py::class_<PyAttribute>(m, "Attribute", py::module_local())
       // extend the backing context which owns the MlirAttribute.
       .def(py::init<PyAttribute &>(), py::arg("cast_from_type"),
            "Casts the passed attribute to the generic Attribute")
       .def_property_readonly(MLIR_PYTHON_CAPI_PTR_ATTR,
                              &PyAttribute::getCapsule)
       .def(MLIR_PYTHON_CAPI_FACTORY_ATTR, &PyAttribute::createFromCapsule)
       .def_static(
           "parse",
-          [](std::string attrSpec, DefaultingPyMlirContext context) {
+          [](const std::string &attrSpec, DefaultingPyMlirContext context) {
             PyMlirContext::ErrorCapture errors(context->getRef());
-            MlirAttribute type = mlirAttributeParseGet(
+            MlirAttribute attr = mlirAttributeParseGet(
                 context->get(), toMlirStringRef(attrSpec));
-            if (mlirAttributeIsNull(type))
+            if (mlirAttributeIsNull(attr))
               throw MLIRError("Unable to parse attribute", errors.take());
-            return type;
+            return attr;
           },
           py::arg("asm"), py::arg("context") = py::none(),
           "Parses an attribute from an assembly form. Raises an MLIRError on "
           "failure.")
       .def_property_readonly(
           "context",
           [](PyAttribute &self) { return self.getContext().getObject(); },
           "Context that owns the Attribute")

mlir/lib/CAPI/IR/BuiltinAttributes.cpp

 MlirAttribute mlirFlatSymbolRefAttrGet(MlirContext ctx, MlirStringRef symbol) {
   return wrap(FlatSymbolRefAttr::get(unwrap(ctx), unwrap(symbol)));
 }
 
 MlirStringRef mlirFlatSymbolRefAttrGetValue(MlirAttribute attr) {
   return wrap(llvm::cast<FlatSymbolRefAttr>(unwrap(attr)).getValue());
 }
 
-MlirTypeID mlirFlatSymbolRefAttrGetTypeID(void) {
-  return wrap(FlatSymbolRefAttr::getTypeID());
-}
-
 //===----------------------------------------------------------------------===//
 // Type attribute.
 //===----------------------------------------------------------------------===//
 
 bool mlirAttributeIsAType(MlirAttribute attr) {
   return llvm::isa<TypeAttr>(unwrap(attr));
 }
 

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):
+    if isinstance(x, list):
+        return SymbolRefAttr.get(x, context=context)
+    else:
+        return FlatSymbolRefAttr.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/dialects/ml_program.py

 # RUN: %PYTHON %s | FileCheck %s
 # This is just a smoke test that the dialect is functional.
 
 from mlir.ir import *
-from mlir.dialects import ml_program
+from mlir.dialects import ml_program, arith, builtin
 
 
 def constructAndPrintInModule(f):
     print("\nTEST:", f.__name__)
     with Context(), Location.unknown():
         module = Module.create()
         with InsertionPoint(module.body):
             f()
         print(module)
     return f
 
 
 # CHECK-LABEL: testFuncOp
 @constructAndPrintInModule
 def testFuncOp():
     # CHECK: ml_program.func @foobar(%arg0: si32) -> si32
     f = ml_program.FuncOp(
         name="foobar", type=([IntegerType.get_signed(32)], [IntegerType.get_signed(32)])
     )
     block = f.add_entry_block()
     with InsertionPoint(block):
         # CHECK: ml_program.return
         ml_program.ReturnOp([block.arguments[0]])
+
+
+# CHECK-LABEL: testGlobalStoreOp
+@constructAndPrintInModule
+def testGlobalStoreOp():
+    # CHECK: %cst = arith.constant 4.242000e+01 : f32
+    cst = arith.ConstantOp(value=42.42, result=F32Type.get())
+
+    m = builtin.ModuleOp()
+    m.sym_name = StringAttr.get("symbol1")
+    m.sym_visibility = StringAttr.get("public")
+    # CHECK: module @symbol1 attributes {sym_visibility = "public"} {
+    # CHECK:   ml_program.global public mutable @symbol2 : f32
+    # CHECK: }
+    with InsertionPoint(m.body):
+        ml_program.GlobalOp("symbol2", F32Type.get(), is_mutable=True)
+    # CHECK: ml_program.global_store @symbol1::@symbol2 = %cst : f32
+    ml_program.GlobalStoreOp(["symbol1", "symbol2"], cst)

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)
+
+        # 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)