[Mlir-commits] [mlir] 5c3861b - [MLIR][python binding] Add OpaqueAttribute to python binding.

[Alex Zinenko]

Author: Yun Long
Date: 2022-03-11T10:56:21+01:00
New Revision: 5c3861b2772c518645419667839febd23035e38f

URL: https://github.com/llvm/llvm-project/commit/5c3861b2772c518645419667839febd23035e38f
DIFF: https://github.com/llvm/llvm-project/commit/5c3861b2772c518645419667839febd23035e38f.diff

LOG: [MLIR][python binding] Add OpaqueAttribute to python binding.

Reviewed By: ftynse

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

Added: 
    

Modified: 
    mlir/lib/Bindings/Python/IRAttributes.cpp
    mlir/test/python/ir/attributes.py

Removed: 
    


################################################################################
diff  --git a/mlir/lib/Bindings/Python/IRAttributes.cpp b/mlir/lib/Bindings/Python/IRAttributes.cpp
index bef3b95a2487a..1093d50c88698 100644
--- a/mlir/lib/Bindings/Python/IRAttributes.cpp
+++ b/mlir/lib/Bindings/Python/IRAttributes.cpp
@@ -320,6 +320,43 @@ class PyFlatSymbolRefAttribute
   }
 };
 
+class PyOpaqueAttribute : public PyConcreteAttribute<PyOpaqueAttribute> {
+public:
+  static constexpr IsAFunctionTy isaFunction = mlirAttributeIsAOpaque;
+  static constexpr const char *pyClassName = "OpaqueAttr";
+  using PyConcreteAttribute::PyConcreteAttribute;
+
+  static void bindDerived(ClassTy &c) {
+    c.def_static(
+        "get",
+        [](std::string dialectNamespace, py::buffer buffer, PyType &type,
+           DefaultingPyMlirContext context) {
+          const py::buffer_info bufferInfo = buffer.request();
+          intptr_t bufferSize = bufferInfo.size;
+          MlirAttribute attr = mlirOpaqueAttrGet(
+              context->get(), toMlirStringRef(dialectNamespace), bufferSize,
+              static_cast<char *>(bufferInfo.ptr), type);
+          return PyOpaqueAttribute(context->getRef(), attr);
+        },
+        py::arg("dialect_namespace"), py::arg("buffer"), py::arg("type"),
+        py::arg("context") = py::none(), "Gets an Opaque attribute.");
+    c.def_property_readonly(
+        "dialect_namespace",
+        [](PyOpaqueAttribute &self) {
+          MlirStringRef stringRef = mlirOpaqueAttrGetDialectNamespace(self);
+          return py::str(stringRef.data, stringRef.length);
+        },
+        "Returns the dialect namespace for the Opaque attribute as a string");
+    c.def_property_readonly(
+        "data",
+        [](PyOpaqueAttribute &self) {
+          MlirStringRef stringRef = mlirOpaqueAttrGetData(self);
+          return py::str(stringRef.data, stringRef.length);
+        },
+        "Returns the data for the Opaqued attributes as a string");
+  }
+};
+
 class PyStringAttribute : public PyConcreteAttribute<PyStringAttribute> {
 public:
   static constexpr IsAFunctionTy isaFunction = mlirAttributeIsAString;
@@ -862,6 +899,7 @@ void mlir::python::populateIRAttributes(py::module &m) {
   PyDenseIntElementsAttribute::bind(m);
   PyDictAttribute::bind(m);
   PyFlatSymbolRefAttribute::bind(m);
+  PyOpaqueAttribute::bind(m);
   PyFloatAttribute::bind(m);
   PyIntegerAttribute::bind(m);
   PyStringAttribute::bind(m);

diff  --git a/mlir/test/python/ir/attributes.py b/mlir/test/python/ir/attributes.py
index 53d246b397528..f5efdf3880bb2 100644
--- a/mlir/test/python/ir/attributes.py
+++ b/mlir/test/python/ir/attributes.py
@@ -236,6 +236,24 @@ def testFlatSymbolRefAttr():
     print("default_get:", FlatSymbolRefAttr.get("foobar"))
 
 
+# CHECK-LABEL: TEST: testOpaqueAttr
+ at 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)
+    # CHECK: oattr value: dummyattr<>
+    print("oattr value:", oattr.data)
+
+    # Test factory methods.
+    # CHECK: default_get: #foobar<"123">
+    print(
+        "default_get:",
+        OpaqueAttr.get("foobar", bytes("123", "utf-8"), NoneType.get()))
+
+
 # CHECK-LABEL: TEST: testStringAttr
 @run
 def testStringAttr():