[Mlir-commits] [mlir] 792f3c8 - [MLIR] Add LocationAttr to the Python API

[Andrew Young]

Author: Andrew Young
Date: 2023-01-25T16:09:29-08:00
New Revision: 792f3c8141bc07140c08e5392edf58e43913d6b2

URL: https://github.com/llvm/llvm-project/commit/792f3c8141bc07140c08e5392edf58e43913d6b2
DIFF: https://github.com/llvm/llvm-project/commit/792f3c8141bc07140c08e5392edf58e43913d6b2.diff

LOG: [MLIR] Add LocationAttr to the Python API

This is a follow up to D142182, to expose LocationAttrs through Python.

Reviewed By: ftynse

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

Added: 
    

Modified: 
    mlir/lib/Bindings/Python/IRCore.cpp
    mlir/test/python/ir/location.py

Removed: 
    


################################################################################
diff  --git a/mlir/lib/Bindings/Python/IRCore.cpp b/mlir/lib/Bindings/Python/IRCore.cpp
index eb7d18f9842c3..2ecfc36d41a55 100644
--- a/mlir/lib/Bindings/Python/IRCore.cpp
+++ b/mlir/lib/Bindings/Python/IRCore.cpp
@@ -2546,10 +2546,25 @@ void mlir::python::populateIRCore(py::module &m) {
           },
           py::arg("name"), py::arg("childLoc") = py::none(),
           py::arg("context") = py::none(), kContextGetNameLocationDocString)
+      .def_static(
+          "from_attr",
+          [](PyAttribute &attribute, DefaultingPyMlirContext context) {
+            return PyLocation(context->getRef(),
+                              mlirLocationFromAttribute(attribute));
+          },
+          py::arg("attribute"), py::arg("context") = py::none(),
+          "Gets a Location from a LocationAttr")
       .def_property_readonly(
           "context",
           [](PyLocation &self) { return self.getContext().getObject(); },
           "Context that owns the Location")
+      .def_property_readonly(
+          "attr",
+          [](PyLocation &self) {
+            return PyAttribute(self.getContext(),
+                               mlirLocationGetAttribute(self));
+          },
+          "Get the underlying LocationAttr")
       .def(
           "emit_error",
           [](PyLocation &self, std::string message) {

diff  --git a/mlir/test/python/ir/location.py b/mlir/test/python/ir/location.py
index ecdd02efb0aee..c1c96d0297c17 100644
--- a/mlir/test/python/ir/location.py
+++ b/mlir/test/python/ir/location.py
@@ -25,6 +25,21 @@ def testUnknown():
 run(testUnknown)
 
 
+# CHECK-LABEL: TEST: testLocationAttr
+def testLocationAttr():
+  with Context() as ctxt:
+    loc = Location.unknown()
+    attr = loc.get_attr()
+    clone = Location.from_attr(attr)
+  gc.collect()
+  # CHECK: loc: loc(unknown)
+  print("loc:", str(loc))
+  # CHECK: clone: loc(unknown)
+  print("clone:", str(clone))
+  assert loc == clone
+
+run(testLocationAttr)
+
 # CHECK-LABEL: TEST: testFileLineCol
 def testFileLineCol():
   with Context() as ctx: