diff --git a/mlir/lib/Bindings/Python/IRCore.cpp b/mlir/lib/Bindings/Python/IRCore.cpp --- 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 @@ -40,6 +41,9 @@ 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 @@ }, py::arg("context") = py::none(), "Gets a Location representing an unknown location") + .def_static( + "callsite", + [](PyLocation callee, const std::vector &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 --- a/mlir/test/python/ir/location.py +++ b/mlir/test/python/ir/location.py @@ -58,6 +58,23 @@ 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: