[mlir][python] Expose CallSiteLoc Python side

[jpienaar]

Summary

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.

[jpienaar]

Rebase & update test

[ftynse]

mlir/lib/Bindings/Python/IRCore.cpp

1974

Nit: throw py::value_error

mlir/test/python/ir/location.py

70

Nit: run already does this.

mlir/lib/Bindings/Python/IRCore.cpp

 #include "PybindUtils.h"
 
 #include "mlir-c/Bindings/Python/Interop.h"
 #include "mlir-c/BuiltinAttributes.h"
 #include "mlir-c/BuiltinTypes.h"
 #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>
 
 namespace py = pybind11;
 using namespace mlir;
 using namespace mlir::python;
 
 using llvm::SmallVector;
 using llvm::StringRef;
 using llvm::Twine;
 
 //------------------------------------------------------------------------------
 // Docstrings (trivial, non-duplicated docstrings are included inline).
 //------------------------------------------------------------------------------
 
 static const char kContextParseTypeDocstring[] =
     R"(Parses the assembly form of a type.
 
 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)";
 
 static const char kContextGetNameLocationDocString[] =
     R"(Gets a Location representing a named location with optional child location)";
 
 static const char kModuleParseDocstring[] =
     R"(Parses a module's assembly format from a string.
py::class_<PyLocation>(m, "Location", py::module_local())
           "unknown",
           [](DefaultingPyMlirContext context) {
             return PyLocation(context->getRef(),
                               mlirLocationUnknownGet(context->get()));
           },
           py::arg("context") = py::none(),
           "Gets a Location representing an unknown location")
       .def_static(
+          "callsite",
+          [](PyLocation callee, std::vector<PyLocation> frames,
+             DefaultingPyMlirContext context) {
+            if (frames.empty())
+              throw SetPyError(PyExc_ValueError, "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,
              DefaultingPyMlirContext context) {
             return PyLocation(
                 context->getRef(),
                 mlirLocationFileLineColGet(
                     context->get(), toMlirStringRef(filename), line, col));
           },

mlir/test/python/ir/location.py

def testName():
   # CHECK: file str: loc("naam"("nombre"))
   print("file str:", str(locWithChildLoc))
   # CHECK: file repr: loc("naam"("nombre"))
   print("file repr:", repr(locWithChildLoc))
 
 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
+  gc.collect()
+  # 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:
     loc1 = Location.file("foo.txt", 123, 56)
   # CHECK: mlir.ir.Location._CAPIPtr
   loc_capsule = loc1._CAPIPtr
   print(loc_capsule)
   loc2 = Location._CAPICreate(loc_capsule)
   assert loc2 == loc1
   assert loc2.context is ctx
 
 run(testLocationCapsule)