[Mlir-commits] [mlir] e67cbbe - [mlir][python] Expose CallSiteLoc Python side

[Alex Zinenko]

Author: Jacques Pienaar
Date: 2021-10-13T10:25:40+02:00
New Revision: e67cbbef03392edc7f5a0e8bd1f64d798d2cfeef

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

LOG: [mlir][python] Expose CallSiteLoc Python side

This exposes creating a CallSiteLoc with a callee & list of frames for
callers. Follows the creation approach in C++ side where a list of
frames may be provided.

Reviewed By: ftynse

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

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 3226b7fe180f5..d53efd9c7bd16 100644
--- a/mlir/lib/Bindings/Python/IRCore.cpp
+++ b/mlir/lib/Bindings/Python/IRCore.cpp
@@ -17,6 +17,7 @@
 #include "mlir-c/Debug.h"
 #include "mlir-c/IR.h"
 #include "mlir-c/Registration.h"
+#include "llvm/ADT/ArrayRef.h"
 #include "llvm/ADT/SmallVector.h"
 #include <pybind11/stl.h>
 
@@ -40,6 +41,9 @@ Returns a Type object or raises a ValueError if the type cannot be parsed.
 See also: https://mlir.llvm.org/docs/LangRef/#type-system
 )";
 
+static const char kContextGetCallSiteLocationDocstring[] =
+    R"(Gets a Location representing a caller and callsite)";
+
 static const char kContextGetFileLocationDocstring[] =
     R"(Gets a Location representing a file, line and column)";
 
@@ -1962,6 +1966,21 @@ void mlir::python::populateIRCore(py::module &m) {
           },
           py::arg("context") = py::none(),
           "Gets a Location representing an unknown location")
+      .def_static(
+          "callsite",
+          [](PyLocation callee, const std::vector<PyLocation> &frames,
+             DefaultingPyMlirContext context) {
+            if (frames.empty())
+              throw py::value_error("No caller frames provided");
+            MlirLocation caller = frames.back().get();
+            for (PyLocation frame :
+                 llvm::reverse(llvm::makeArrayRef(frames).drop_back()))
+              caller = mlirLocationCallSiteGet(frame.get(), caller);
+            return PyLocation(context->getRef(),
+                              mlirLocationCallSiteGet(callee.get(), caller));
+          },
+          py::arg("callee"), py::arg("frames"), py::arg("context") = py::none(),
+          kContextGetCallSiteLocationDocstring)
       .def_static(
           "file",
           [](std::string filename, int line, int col,

diff  --git a/mlir/test/python/ir/location.py b/mlir/test/python/ir/location.py
index b17e5ec6be6a8..7bc7b28ae6a4a 100644
--- a/mlir/test/python/ir/location.py
+++ b/mlir/test/python/ir/location.py
@@ -58,6 +58,23 @@ def testName():
 run(testName)
 
 
+# CHECK-LABEL: TEST: testCallSite
+def testCallSite():
+  with Context() as ctx:
+    loc = Location.callsite(
+        Location.file("foo.text", 123, 45), [
+            Location.file("util.foo", 379, 21),
+            Location.file("main.foo", 100, 63)
+        ])
+  ctx = None
+  # CHECK: file str: loc(callsite("foo.text":123:45 at callsite("util.foo":379:21 at "main.foo":100:63))
+  print("file str:", str(loc))
+  # CHECK: file repr: loc(callsite("foo.text":123:45 at callsite("util.foo":379:21 at "main.foo":100:63))
+  print("file repr:", repr(loc))
+
+run(testCallSite)
+
+
 # CHECK-LABEL: TEST: testLocationCapsule
 def testLocationCapsule():
   with Context() as ctx: