[Mlir-commits] [mlir] d5429a1 - [mlir][python] Add 'loc' property to ops

[Alex Zinenko]

Author: rkayaith
Date: 2021-10-18T16:01:12+02:00
New Revision: d5429a13da2b61b356f996a41ea314a1a1ce9076

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

LOG: [mlir][python] Add 'loc' property to ops

Add a read-only `loc` property to Operation and OpView

Reviewed By: ftynse

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

Added: 
    

Modified: 
    mlir/include/mlir-c/IR.h
    mlir/lib/Bindings/Python/IRCore.cpp
    mlir/lib/CAPI/IR/IR.cpp
    mlir/test/python/ir/operation.py

Removed: 
    


################################################################################
diff  --git a/mlir/include/mlir-c/IR.h b/mlir/include/mlir-c/IR.h
index 92697a248b71b..2fec0be4defb2 100644
--- a/mlir/include/mlir-c/IR.h
+++ b/mlir/include/mlir-c/IR.h
@@ -357,6 +357,9 @@ MLIR_CAPI_EXPORTED bool mlirOperationEqual(MlirOperation op,
 /// Gets the context this operation is associated with
 MLIR_CAPI_EXPORTED MlirContext mlirOperationGetContext(MlirOperation op);
 
+/// Gets the location of the operation.
+MLIR_CAPI_EXPORTED MlirLocation mlirOperationGetLocation(MlirOperation op);
+
 /// Gets the type id of the operation.
 /// Returns null if the operation does not have a registered operation
 /// description.

diff  --git a/mlir/lib/Bindings/Python/IRCore.cpp b/mlir/lib/Bindings/Python/IRCore.cpp
index 7b1c998297658..ed96bafe41919 100644
--- a/mlir/lib/Bindings/Python/IRCore.cpp
+++ b/mlir/lib/Bindings/Python/IRCore.cpp
@@ -2143,6 +2143,15 @@ void mlir::python::populateIRCore(py::module &m) {
           },
           "Shortcut to get an op result if it has only one (throws an error "
           "otherwise).")
+      .def_property_readonly(
+          "location",
+          [](PyOperationBase &self) {
+            PyOperation &operation = self.getOperation();
+            return PyLocation(operation.getContext(),
+                              mlirOperationGetLocation(operation.get()));
+          },
+          "Returns the source location the operation was defined or derived "
+          "from.")
       .def("__iter__",
            [](PyOperationBase &self) {
              return PyRegionIterator(self.getOperation().getRef());

diff  --git a/mlir/lib/CAPI/IR/IR.cpp b/mlir/lib/CAPI/IR/IR.cpp
index ee5a5551133c9..c738198f75b42 100644
--- a/mlir/lib/CAPI/IR/IR.cpp
+++ b/mlir/lib/CAPI/IR/IR.cpp
@@ -346,6 +346,10 @@ MlirContext mlirOperationGetContext(MlirOperation op) {
   return wrap(unwrap(op)->getContext());
 }
 
+MlirLocation mlirOperationGetLocation(MlirOperation op) {
+  return wrap(unwrap(op)->getLoc());
+}
+
 MlirTypeID mlirOperationGetTypeID(MlirOperation op) {
   if (const auto *abstractOp = unwrap(op)->getAbstractOperation()) {
     return wrap(abstractOp->typeID);

diff  --git a/mlir/test/python/ir/operation.py b/mlir/test/python/ir/operation.py
index 57a6d57a7f7bc..e53be97cf25d9 100644
--- a/mlir/test/python/ir/operation.py
+++ b/mlir/test/python/ir/operation.py
@@ -728,3 +728,15 @@ def testOperationErase():
 
       # Ensure we can create another operation
       Operation.create("custom.op2")
+
+
+# CHECK-LABEL: TEST: testOperationLoc
+ at run
+def testOperationLoc():
+  ctx = Context()
+  ctx.allow_unregistered_dialects = True
+  with ctx:
+    loc = Location.name("loc")
+    op = Operation.create("custom.op", loc=loc)
+    assert op.location == loc
+    assert op.operation.location == loc